From 42e22c00cfe21f8ab86a3d609a0c420eb452aa0d Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Tue, 29 Sep 2026 20:53:01 +0000 Subject: [PATCH 01/58] tmp Signed-off-by: Kaining Zhong --- .../cast/mxfp8/group_quantize_mxfp8.py | 1136 +++++++++++++++++ .../common/cast/dispatch/quantize.cuh | 15 +- .../mxfp8/group_quantize_mxfp8_cutedsl.cuh | 312 +++++ 3 files changed, 1460 insertions(+), 3 deletions(-) create mode 100644 transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8.py create mode 100644 transformer_engine/common/cast/mxfp8/group_quantize_mxfp8_cutedsl.cuh diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8.py new file mode 100644 index 00000000000..11c288a7d36 --- /dev/null +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8.py @@ -0,0 +1,1136 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Grouped MXFP8 quantization kernel implemented in CuTeDSL. + +Strategy-aligned port of group_quantize_mxfp8.cuh. The scheduling, descriptor +management and per-tensor scale addressing mirror the CUDA kernel one-for-one: + + * `is_single_tensor` reps (SAME_BOTH_DIMS, VARYING_FIRST_DIM) launch ONE CTA per + 128x128 chunk and address the group through ONE static TMA descriptor with + global block offsets -- the CUDA `tensor_map_*_static` "direct mapper" path. + For SAME_BOTH_DIMS the CUDA grid is linearized per tensor (X, Y-in-tensor, + tensor) while this kernel linearizes it flat over the stacked rows; both + require every member's row count to be a multiple of CHUNK_DIM_Y, and under + that precondition the two decode to the identical (block_offset_Y, block_id_X) + for every block index. + * the other reps launch grid=(workers_per_tensor, num_tensors) and bind + tensor_id to blockIdx.y, so a CTA grid-strides only within its own tensor and + never re-resolves which tensor a chunk belongs to. They get per-tensor + descriptors written by a prologue kernel (the CuTeDSL analog of + update_tma_descriptors filling g_tensor_maps) and acquired with a tensormap + proxy fence. + * per-tensor scale bases/strides follow the CUDA formulas: + scales_* += is_single_tensor ? 0 : tensor_base / 32 + stride_rowwise = roundup(cols/32, 4) stride_colwise = roundup(cols, 128) + +Both kernels dropped the older flat persistent grid that strided across tensor +boundaries (CUDA's decode_job / advance_to_next_job, removed in #3483); the +per-tensor grid above is what replaced it on both sides. + +Mechanics that provably yield the same bytes may differ: the mbarrier pipeline is +expressed with PipelineTmaAsync instead of hand-rolled mbarriers, and +out-of-bounds scale padding is skipped rather than explicitly zeroed (the CUDA +kernel writes 0 there; downstream only consumes the meaningful region). + +Scope: cast-only (no dbias / activation / dact / amax), compact (non-swizzled) +scales, rowwise and/or colwise. VARYING_BOTH_DIMS is not handled here (its +logical shape [1, total] is not tileable); the C++ bridge falls back to CUDA. + +Known gap vs CUDA: the CUDA kernel validates on device that every member's first +dim is a multiple of 128 (get_tensor_rows_num -> NVTE_DEVICE_ERROR). This kernel +shares the precondition but cannot check it -- CuTeDSL has no device-side assert, +and for VARYING_FIRST_DIM the per-tensor extents live in device memory, so the +C++ bridge cannot check them either without a sync. A violating group mis-tiles +silently here where CUDA raises. +""" + +""" +Measured, deliberately NOT changed: + - sO_row and sO_col are both allocated unconditionally, where CUDA sizes only the + direction in use. Sizing them conditionally does work -- ncu confirms the shared-memory + occupancy limit goes 6 -> 9 CTAs/SM for a single-direction bf16 config -- but it is a + small LOSS on GB200, not a win: rep_med_sbd bf16 colwise 54.9 -> 56.1 us, rowwise + 56.5 -> 57.0 us (fp32 rowwise gains ~1%). The kernel is DRAM-bandwidth-bound at + ~6.3 TB/s, so extra resident CTAs only add contention. Verified by a control that kept + the conditional code but padded SMEM back to the old size: timings returned exactly to + the unconditional numbers, so the effect is the occupancy, not codegen. Revisit if a + future variant (dbias / activation) makes this kernel latency- rather than + bandwidth-bound. +""" + +# pylint: disable=missing-class-docstring + +import logging +import os +from typing import Type + +import cutlass +from cutlass import cute +from cutlass import pipeline +from cutlass import Boolean, Int32, Int64, Float8E8M0FNU +from cutlass.cute.nvgpu import cpasync +from cutlass.tensor_utils import TensorMapManager, TensorMapUpdateMode +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.cast.mxfp8.quantize_mxfp8 import ( + MXFP8_BLOCK_SCALING_SIZE, + FP8E4M3_MAX_NORM_RCP, + FP8E5M2_MAX_NORM_RCP, + quantize_rowwise_mxfp8, + quantize_colwise_mxfp8, +) + +CUTEDSL_DEBUG_LOGGING = os.environ.get("CUTEDSL_DEBUG_LOGGING", "0") == "1" +logger = logging.getLogger("transformer_engine.cutedsl.mxfp8") + +THREADS_PER_WARP = 32 +BYTES_PER_TENSORMAP = 128 +# Descriptor slots per tensor: input, rowwise output, colwise output. +NUM_TENSORMAPS = 3 +# One extra slot holds per-tensor (rows, cols, base_elts), so the main kernel never has to +# binary-search the offsets array. Mirrors TensorMapStorage::rows/cols/offsets upstream. +META_SLOT = NUM_TENSORMAPS +NUM_WORKSPACE_SLOTS = NUM_TENSORMAPS + 1 + +# Shape representations, mirroring ShapeRepresentation in common/utils.cuh. +SAME_BOTH_DIMS = "same_both_dims" +VARYING_FIRST_DIM = "varying_first_dim" +VARYING_LAST_DIM = "varying_last_dim" +SUPPORTED_SHAPE_REPS = (SAME_BOTH_DIMS, VARYING_FIRST_DIM, VARYING_LAST_DIM) + + +class MXFP8GroupQuantizeConfig: + """Compile-time config for the grouped MXFP8 quantize kernel.""" + + def __init__(self, dtype: str, fp8_dtype: str, rowwise: bool, colwise: bool, shape_rep: str): + if dtype not in ("fp32", "fp16", "bf16"): + raise ValueError(f"unknown input dtype {dtype!r}; expected fp32|fp16|bf16") + self.DTYPE = str_to_cutlass_dtype(dtype) + self.DTYPE_STR = dtype + if fp8_dtype not in ("fp8_e4m3fn", "fp8_e5m2"): + raise ValueError(f"unknown FP8 dtype {fp8_dtype!r}; expected fp8_e4m3fn|fp8_e5m2") + self.FP8_DTYPE = str_to_cutlass_dtype(fp8_dtype) + self.FP8_DTYPE_STR = fp8_dtype + if not (rowwise or colwise): + raise ValueError("at least one of rowwise or colwise must be true") + self.ROWWISE = rowwise + self.COLWISE = colwise + if shape_rep not in SUPPORTED_SHAPE_REPS: + raise ValueError( + f"unsupported shape representation {shape_rep!r}; expected one of" + f" {SUPPORTED_SHAPE_REPS}" + ) + self.SHAPE_REP = shape_rep + # Mirrors `is_single_tensor` in group_quantize_mxfp8.cuh. + self.IS_SINGLE_TENSOR = shape_rep in (SAME_BOTH_DIMS, VARYING_FIRST_DIM) + self.MAX_NORM_RCP = ( + FP8E4M3_MAX_NORM_RCP if fp8_dtype == "fp8_e4m3fn" else FP8E5M2_MAX_NORM_RCP + ) + + def __str__(self): + return ( + f"MXFP8GroupQuantizeConfig(dtype={self.DTYPE_STR}, fp8_dtype={self.FP8_DTYPE_STR}, " + f"rowwise={self.ROWWISE}, colwise={self.COLWISE}, shape_rep={self.SHAPE_REP})" + ) + + __repr__ = __str__ + + +class MXFP8GroupQuantizeKernel: + """Grouped MXFP8 quantize mirroring group_quantize_mxfp8_kernel's strategy.""" + + # TunableConfig / derived constants from group_quantize_mxfp8.cuh. + CHUNK_DIM_Y = 128 + CHUNK_DIM_X = 128 + THREADS_PER_CHUNK = 128 + STATIC_PERSISTENT_BLOCKS_PER_SM = 24 + ELTS_PER_CHUNK = CHUNK_DIM_Y * CHUNK_DIM_X + THREADS_X = CHUNK_DIM_X // MXFP8_BLOCK_SCALING_SIZE # 4 + THREADS_Y = THREADS_PER_CHUNK // THREADS_X # 32 + BUFF_DIM_Y = THREADS_Y # 32 + BUFF_DIM_X = CHUNK_DIM_X # 128 + # Each block of (CHUNK_DIM_Y, CHUNK_DIM_X) consists of STAGES tiles of (BUFF_DIM_X, BUFF_DIM_Y) stacked vertically + STAGES = CHUNK_DIM_Y // BUFF_DIM_Y # 4 + PIPELINE_DEPTH = 2 # PREFETCH_STAGES(1) + 1 + NUM_WARPS = THREADS_PER_CHUNK // THREADS_PER_WARP # 4 + + # Rowwise vectorization constants (mirror MXFP8QuantizeKernel / CUDA PACK_SIZE). + PACK_SIZE = 4 + WAVES = MXFP8_BLOCK_SCALING_SIZE // PACK_SIZE # 8 + THREADS_PER_BANK = (32 * 4) // MXFP8_BLOCK_SCALING_SIZE # 4 + + def __init__(self, cfg: MXFP8GroupQuantizeConfig, SM_COUNT: int): + self.cfg = cfg + self.SM_COUNT = SM_COUNT + + # ---------------------------------------------------------------- helpers + @cute.jit + def _tensor_rows_cols( + self, tensor_id, mFirstDims, mLastDims, first_logical_dim, last_logical_dim + ): + """Get the shape (rows, cols) of the tensor by tensor_id.""" + cfg = self.cfg + if cutlass.const_expr(cfg.SHAPE_REP == VARYING_FIRST_DIM): + rows = Int32(mFirstDims[tensor_id]) + else: + rows = Int32(first_logical_dim) + if cutlass.const_expr(cfg.SHAPE_REP == VARYING_LAST_DIM): + cols = Int32(mLastDims[tensor_id]) + else: + cols = Int32(last_logical_dim) + return rows, cols + + # ------------------------------------------------------------ entry point + @cute.jit + def __call__( + self, + mX: cute.Tensor, + mO_row: cute.Tensor, + mO_col: cute.Tensor, + mS_row: cute.Tensor, + mS_col: cute.Tensor, + mOffsets: cute.Tensor, # int64[num_tensors + 1], CSR element offsets + mFirstDims: cute.Tensor, # int64[num_tensors] (VARYING_FIRST_DIM) + mLastDims: cute.Tensor, # int64[num_tensors] (VARYING_LAST_DIM) + mTensormaps: cute.Tensor, # int64[num_tensors, NUM_TENSORMAPS, 16] + stream: CUstream, + ): + if cutlass.const_expr(CUTEDSL_DEBUG_LOGGING): + cute.printf(f"[CuTeDSL] MXFP8GroupQuantizeKernel.__call__() cfg: {self.cfg}\n") + + cfg = self.cfg + first_logical_dim = mX.shape[0] + last_logical_dim = mX.shape[1] + # Number of group members. Do NOT derive this from mOffsets: a caller is free to + # pass a length-num_tensors stub for SAME_BOTH_DIMS (where the offsets array is + # unused), which would make `mOffsets.shape[0] - 1` read one too few and divide by + # zero at num_tensors == 1. The per-tensor descriptor workspace is num_tensors long + # by construction -- one slot set per member -- so it is the reliable source, and it + # is only consulted on the multi-tensor path that actually needs the workspace. + num_tensors = mTensormaps.shape[0] + + smem_tile_layout = cute.make_ordered_layout( + (self.BUFF_DIM_Y, self.BUFF_DIM_X), order=(1, 0) + ) + cta_tiler = (self.BUFF_DIM_Y, self.BUFF_DIM_X) + print(f"mx={mX}, smem_tile_layout={smem_tile_layout}, cta_tiler={cta_tiler}\n") + + op_load = cpasync.CopyBulkTensorTileG2SOp() + tma_atom_x, tma_src = cpasync.make_tiled_tma_atom( + op_load, mX, smem_tile_layout, cta_tiler, num_multicast=1 + ) + print(f"tma_atom_x={tma_atom_x}\n") + print(f"tma_src={tma_src}\n") + op_store = cpasync.CopyBulkTensorTileS2GOp() + tma_atom_out_row, tma_dst_out_row = cpasync.make_tiled_tma_atom( + op_store, mO_row, smem_tile_layout, cta_tiler, num_multicast=1 + ) + tma_atom_out_col, tma_dst_out_col = cpasync.make_tiled_tma_atom( + op_store, mO_col, smem_tile_layout, cta_tiler, num_multicast=1 + ) + + if cutlass.const_expr(cfg.IS_SINGLE_TENSOR): + # How many blocks does the grouped tensor have in both directions + work_blocks_X = cute.ceil_div(Int32(last_logical_dim), self.CHUNK_DIM_X) + work_blocks_Y = cute.ceil_div(Int32(first_logical_dim), self.CHUNK_DIM_Y) + # Each CTA handles a block from one individual tensor + grid = [work_blocks_X * work_blocks_Y, 1, 1] + else: + # The work-block grid is per-tensor here, not global: each CTA derives its own + # block range from its tensor's extents (see the kernel's `else` branch), so + # work_blocks_X/Y are dead on this path -- the kernel reads them only under + # IS_SINGLE_TENSOR. Pass dummies rather than computing first*last, which is an + # ELEMENT count and would wrap Int32 past 2^31 elements. + work_blocks_Y = Int32(1) + work_blocks_X = Int32(1) + # Persistent worker count, mirroring get_launch_config() in + # group_quantize_mxfp8.cuh. There are SM_COUNT * STATIC_PERSISTENT_BLOCKS_PER_SM + # workers in total, split evenly across tensors -- but clamped to the average + # number of chunks a tensor actually holds, so a group of small tensors does not + # launch CTAs that only reach the `first_block_id >= blocks_in_tensor` early-out. + n_tensors = cutlass.max(Int32(num_tensors), Int32(1)) # never divide by zero + # CUDA's DIVUP(elts_total, CHUNK_DIM_Y * TILE_DIM_X) / STAGES_X, where TILE_DIM_X + # is BUFF_DIM_X here and STAGES_X is 1 for every rep this kernel supports. The + # element product would wrap Int32 past 2^31 elements, so divide the first extent + # by CHUNK_DIM_Y up front -- it is exact, the kernel is compiled with + # sym_int32(divisibility=CHUNK_DIM_Y) on that extent -- and keep the whole + # estimate in Int32. Feeding an Int64 into the grid poisons the tile arithmetic + # downstream ('cute.make_tile' expects width=32). + estimated_work_blocks = cute.ceil_div( + (Int32(first_logical_dim) // self.CHUNK_DIM_Y) * Int32(last_logical_dim), + self.ELTS_PER_CHUNK // self.CHUNK_DIM_Y, + ) + requested_workers_per_tensor = cutlass.max( + Int32(1), + Int32(self.SM_COUNT * self.STATIC_PERSISTENT_BLOCKS_PER_SM) // n_tensors, + ) + average_work_blocks_per_tensor = cutlass.max( + Int32(1), cute.ceil_div(estimated_work_blocks, n_tensors) + ) + workers_per_tensor = cutlass.min( + requested_workers_per_tensor, average_work_blocks_per_tensor + ) + # Each group of workers_per_tensor CTAs serves one tensor, and we launch + # num_tensors such groups to cover the whole group. + grid = [workers_per_tensor, Int32(num_tensors), 1] + + # Only manually create descriptors for the non-single-tensor case because we will need to manually + # overwrite the descriptors as we visit different groups + if cutlass.const_expr(not cfg.IS_SINGLE_TENSOR): + self.update_descriptors_kernel( + mX, + mO_row, + mO_col, + mOffsets, + mFirstDims, + mLastDims, + mTensormaps, + first_logical_dim, + last_logical_dim, + mX.element_type, + tma_atom_x, + tma_atom_out_row, + tma_atom_out_col, + ).launch(grid=[num_tensors, 1, 1], block=[THREADS_PER_WARP, 1, 1], stream=stream) + + self.kernel( + mS_row, + mS_col, + mOffsets, + mFirstDims, + mLastDims, + mTensormaps, + first_logical_dim, + last_logical_dim, + num_tensors, + work_blocks_X, + work_blocks_Y, + mX.element_type, + tma_atom_x, + tma_src, + tma_atom_out_row, + tma_dst_out_row, + tma_atom_out_col, + tma_dst_out_col, + ).launch( + grid=grid, + block=[self.THREADS_PER_CHUNK, 1, 1], + stream=stream, + ) + + # ------------------------------------------------- descriptor prologue + @cute.kernel + def update_descriptors_kernel( + self, + mX, + mO_row, + mO_col, + mOffsets, + mFirstDims, + mLastDims, + mTensormaps, + first_logical_dim, + last_logical_dim, + dtype: cutlass.Constexpr[Type[cutlass.Numeric]], + tma_atom_x, + tma_atom_orow, + tma_atom_ocol, + ): + """One CTA per tensor: point that tensor's TMA descriptors at its own block. + + CuTeDSL analog of common::update_tma_descriptors writing g_tensor_maps[]. + """ + cfg = self.cfg + tensor_id, _, _ = cute.arch.block_idx() + rows, cols = self._tensor_rows_cols( + tensor_id, mFirstDims, mLastDims, first_logical_dim, last_logical_dim + ) + base_elts = Int64(mOffsets[tensor_id]) + + # Publish this tensor's geometry for the main kernel (written even when empty). + meta = mTensormaps[(tensor_id, META_SLOT, None)] + meta[0] = Int64(rows) + meta[1] = Int64(cols) + meta[2] = base_elts + + tmap = TensorMapManager(TensorMapUpdateMode.GMEM, BYTES_PER_TENSORMAP) + desc_x = tmap.get_tensormap_ptr(mTensormaps[(tensor_id, 0, None)].iterator) + desc_orow = tmap.get_tensormap_ptr(mTensormaps[(tensor_id, 1, None)].iterator) + desc_ocol = tmap.get_tensormap_ptr(mTensormaps[(tensor_id, 2, None)].iterator) + + # Zero-sized groups: creating a descriptor with a zero extent is invalid, + # so skip (the main kernel skips these jobs via job_has_work). + if rows > 0 and cols > 0: + gX = cute.make_tensor( + cute.make_ptr( + dtype, + mX.iterator.toint() + base_elts * (dtype.width // 8), + cute.AddressSpace.gmem, + assumed_align=16, + ), + cute.make_layout((rows, cols), stride=(cols, 1)), + ) + gO_row = cute.make_tensor( + cute.make_ptr( + cfg.FP8_DTYPE, + mO_row.iterator.toint() + base_elts, + cute.AddressSpace.gmem, + assumed_align=16, + ), + cute.make_layout((rows, cols), stride=(cols, 1)), + ) + gO_col = cute.make_tensor( + cute.make_ptr( + cfg.FP8_DTYPE, + mO_col.iterator.toint() + base_elts, + cute.AddressSpace.gmem, + assumed_align=16, + ), + cute.make_layout((rows, cols), stride=(cols, 1)), + ) + + tmap.init_tensormap_from_atom(tma_atom_x, desc_x, 0) + if cutlass.const_expr(cfg.ROWWISE): + tmap.init_tensormap_from_atom(tma_atom_orow, desc_orow, 0) + if cutlass.const_expr(cfg.COLWISE): + tmap.init_tensormap_from_atom(tma_atom_ocol, desc_ocol, 0) + tmap.fence_tensormap_initialization() + + if cutlass.const_expr(cfg.ROWWISE and cfg.COLWISE): + tmap.update_tensormap( + (gX, gO_row, gO_col), + (tma_atom_x, tma_atom_orow, tma_atom_ocol), + (desc_x, desc_orow, desc_ocol), + 0, + (), # smem staging is unused in GMEM update mode + ) + elif cutlass.const_expr(cfg.ROWWISE): + tmap.update_tensormap( + (gX, gO_row), + (tma_atom_x, tma_atom_orow), + (desc_x, desc_orow), + 0, + (), # smem staging is unused in GMEM update mode + ) + else: + tmap.update_tensormap( + (gX, gO_col), + (tma_atom_x, tma_atom_ocol), + (desc_x, desc_ocol), + 0, + (), # smem staging is unused in GMEM update mode + ) + + # ------------------------------------------------------------ main kernel + @cute.kernel + def kernel( + self, + mS_row, + mS_col, + mOffsets, + mFirstDims, + mLastDims, + mTensormaps, + first_logical_dim, + last_logical_dim, + num_tensors, + work_blocks_X, + work_blocks_Y, + dtype: cutlass.Constexpr[Type[cutlass.Numeric]], + tma_atom_x, + tma_src, + tma_atom_out_row, + tma_dst_out_row, + tma_atom_out_col, + tma_dst_out_col, + ): + cfg = self.cfg + FP8_DTYPE = cfg.FP8_DTYPE + tidx, _, _ = cute.arch.thread_idx() + bidx, bidy, _ = cute.arch.block_idx() + gdx, _, _ = cute.arch.grid_dim() + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + + # --- shared memory (allocated once, reused across jobs) --- + @cute.struct + class SharedStorage: + mbar: cute.struct.MemRange[cute.Int64, 2 * self.PIPELINE_DEPTH] + sX: cute.struct.Align[ + cute.struct.MemRange[ + dtype, self.BUFF_DIM_Y * self.BUFF_DIM_X * self.PIPELINE_DEPTH + ], + 128, + ] + sO_row: cute.struct.Align[ + cute.struct.MemRange[ + FP8_DTYPE, self.BUFF_DIM_Y * self.BUFF_DIM_X * self.PIPELINE_DEPTH + ], + 128, + ] + sO_col: cute.struct.Align[ + cute.struct.MemRange[ + FP8_DTYPE, self.BUFF_DIM_Y * self.BUFF_DIM_X * self.PIPELINE_DEPTH + ], + 128, + ] + + storage = cutlass.utils.SmemAllocator().allocate(SharedStorage) + tile_layout = cute.make_layout( + ((self.BUFF_DIM_Y, self.BUFF_DIM_X), self.PIPELINE_DEPTH), + stride=((self.BUFF_DIM_X, 1), self.BUFF_DIM_Y * self.BUFF_DIM_X), + ) + sX = storage.sX.get_tensor(tile_layout) + print(f"sX={sX}\n") + sO_row = storage.sO_row.get_tensor(tile_layout) + sO_col = storage.sO_col.get_tensor(tile_layout) + + mainloop_pipeline = pipeline.PipelineTmaAsync.create( + barrier_storage=storage.mbar.data_ptr(), + num_stages=self.PIPELINE_DEPTH, + producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), + consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, self.NUM_WARPS), + tx_count=self.BUFF_DIM_Y * self.BUFF_DIM_X * dtype.width // 8, + cta_layout_vmnk=None, + ) + prod_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.PIPELINE_DEPTH + ) + cons_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.PIPELINE_DEPTH + ) + + # TMA partitions built from the representative views. For multi-tensor reps + # the descriptor is swapped per tensor and the tile coords are tensor-local. + gX_tiled = cute.zipped_divide(tma_src, (self.BUFF_DIM_Y, self.BUFF_DIM_X)) + tXsX, tXgX = cpasync.tma_partition(tma_atom_x, 0, cute.make_layout(1), sX, gX_tiled) + gO_row_tiled = cute.zipped_divide(tma_dst_out_row, (self.BUFF_DIM_Y, self.BUFF_DIM_X)) + tXsO_row, tXgO_row = 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.BUFF_DIM_Y, self.BUFF_DIM_X)) + tXsO_col, tXgO_col = cpasync.tma_partition( + tma_atom_out_col, 0, cute.make_layout(1), sO_col, gO_col_tiled + ) + print(f"tma_atom_x={tma_atom_x}\n") + print(f"tma_src={tma_src}\n") + print(f"gX_tiled={gX_tiled}\n") + print(f"tXsX={tXsX}\n") + print(f"tXgX={tXgX}\n") + + tmap = TensorMapManager(TensorMapUpdateMode.GMEM, BYTES_PER_TENSORMAP) + cute.arch.sync_threads() + + # If the CTA has work to do + has_work = Boolean(True) + # Metadata of the tensor that owns this block + tensor_rows = Int32(0) + tensor_cols = Int32(0) + # Element offset of this tensor within the group: Int64 (CUDA uses size_t), since a + # group can exceed 2^31 elements even when every individual extent is small. + # Element offset of this tensor within the group: Int64 (CUDA uses size_t), since a + # group can exceed 2^31 elements even when every individual extent is small. + tensor_base = Int64(0) + # Block's offset and id in this individual tensor / global single tensor + block_offset_Y = Int32(0) + block_id_X = Int32(0) + + first_block_id = Int32(0) + blocks_in_tensor = Int32(1) + block_stride = Int32(1) + block_columns_in_tensor = Int32(1) + + if cutlass.const_expr(cfg.IS_SINGLE_TENSOR): + # grid = [work_blocks_X * work_blocks_Y, 1, 1] + block_id_Y = Int32(bidx) // work_blocks_X + block_id_X = Int32(bidx) % work_blocks_X + # View the grouped tensor as a single tensor of shape (first_logical_dim, last_logical_dim) + tensor_rows = Int32(first_logical_dim) + tensor_cols = Int32(last_logical_dim) + # Which row does this block start from + block_offset_Y = block_id_Y * self.CHUNK_DIM_Y + if cutlass.const_expr(cfg.SHAPE_REP == VARYING_FIRST_DIM): + # Check if this block contains any valid tokens since the last offset <= logical_first_dim * logical_last_dim + # Last CSR offset == the group's total element count. Indexed off the + # array's own length so it does not depend on num_tensors. + total_elts = Int64(mOffsets[mOffsets.shape[0] - 1]) + if Int64(block_offset_Y) * Int64(last_logical_dim) >= total_elts: + has_work = Boolean(False) + else: + # grid = [workers_per_tensor, Int32(num_tensors), 1] + tensor_id = Int32(bidy) + # Extract tensor's metadata + meta = mTensormaps[(tensor_id, META_SLOT, None)] + tensor_rows = Int32(meta[0]) + tensor_cols = Int32(meta[1]) + tensor_base = Int64(meta[2]) + if tensor_rows > 0 and tensor_cols > 0: + # How many blocks does this tensor have in both directions + block_columns_in_tensor = cute.ceil_div(tensor_cols, self.CHUNK_DIM_X) + block_rows_in_tensor = cute.ceil_div(tensor_rows, self.CHUNK_DIM_Y) + # How many blocks does this tensor have + blocks_in_tensor = block_columns_in_tensor * block_rows_in_tensor + # Which block (1D index) does this CTA start from + first_block_id = Int32(bidx) + # gdx is workers_per_tensor (how many CTAs are assigned to this tensor) + block_stride = Int32(gdx) + # If my first block_id is already beyond the tensor's last block, I have no work to do + if first_block_id >= blocks_in_tensor: + has_work = Boolean(False) + else: + # This tensor is empty, so this CTA has no work to do + has_work = Boolean(False) + + if has_work: + if cutlass.const_expr(cfg.IS_SINGLE_TENSOR): + # For single tensor case we don't use tensor descriptors + desc_x = desc_out_row = desc_out_col = None + + cute.arch.sync_threads() + + self._process_block( + block_offset_Y, + block_id_X, + tensor_rows, + tensor_cols, + tensor_base, + desc_x, + desc_out_row, + desc_out_col, + mS_row, + mS_col, + first_logical_dim, + tmap, + warp_idx, + tidx, + sX, + sO_row, + sO_col, + tXsX, + tXgX, + tXsO_row, + tXgO_row, + tXsO_col, + tXgO_col, + tma_atom_x, + tma_atom_out_row, + tma_atom_out_col, + mainloop_pipeline, + prod_state, + cons_state, + ) + else: + # For multi-tensor case, retrieve tensor descriptors we processed early in the prologue kernel + desc_x = tmap.get_tensormap_ptr(mTensormaps[(tensor_id, 0, None)].iterator) + desc_out_row = tmap.get_tensormap_ptr(mTensormaps[(tensor_id, 1, None)].iterator) + desc_out_col = tmap.get_tensormap_ptr(mTensormaps[(tensor_id, 2, None)].iterator) + # Acquire the descriptors on ONE thread, as the CUDA kernel does + # (`leading_thread` in group_quantize_mxfp8.cuh); the sync_threads below + # publishes it CTA-wide. Running the tensormap acquire fence on all 128 + # threads is correct but very expensive -- it more than doubles the + # kernel time on the multi-tensor path (4096x14336 bidirectional: + # 133 us -> 58 us), since the cost scales with threads x descriptors. + if tidx == 0: + tmap.fence_tensormap_update(desc_x) + if cutlass.const_expr(cfg.ROWWISE): + tmap.fence_tensormap_update(desc_out_row) + if cutlass.const_expr(cfg.COLWISE): + tmap.fence_tensormap_update(desc_out_col) + + cute.arch.sync_threads() + + # Grid-stride over this tensor's own chunks; the descriptors never change. + block_id = first_block_id + job_finished = Boolean(False) + while not job_finished: + block_id_Y_in_tensor = block_id // block_columns_in_tensor + block_id_X_in_tensor = block_id % block_columns_in_tensor + self._process_block( + block_id_Y_in_tensor * self.CHUNK_DIM_Y, + block_id_X_in_tensor, + tensor_rows, + tensor_cols, + tensor_base, + desc_x, + desc_out_row, + desc_out_col, + mS_row, + mS_col, + first_logical_dim, + tmap, + warp_idx, + tidx, + sX, + sO_row, + sO_col, + tXsX, + tXgX, + tXsO_row, + tXgO_row, + tXsO_col, + tXgO_col, + tma_atom_x, + tma_atom_out_row, + tma_atom_out_col, + mainloop_pipeline, + prod_state, + cons_state, + ) + # Find the next block to process + block_id = block_id + block_stride + if block_id >= blocks_in_tensor: + job_finished = Boolean(True) + + # Drain every TMA store before the CTA releases its shared-memory source buffers. + if warp_idx == 0: + cute.arch.cp_async_bulk_wait_group(0, read=False) + cute.arch.sync_threads() + + def _issue_load( + self, + pipeline_obj, + prod_state, + tile_y, + tile_x, + tma_atom_x, + tXgX, + tXsX, + tmap, + desc_x, + ): + """Emit one 32x128 TMA load into the current pipeline buffer. + + Caller gates this on warp 0 and advances `prod_state` afterwards -- the advance + must happen outside the gate or the mutated SSA values stay trapped in the scf.if. + """ + # Wait for the consumer to finish using this SMEM buffer + pipeline_obj.producer_acquire(prod_state) + if cutlass.const_expr(self.cfg.IS_SINGLE_TENSOR): + cute.copy( + tma_atom_x, + tXgX[(None, (tile_y, tile_x))], + tXsX[(None, prod_state.index)], + tma_bar_ptr=pipeline_obj.producer_get_barrier(prod_state), + ) + else: + # Every member shares tXgX's tile-coordinate arithmetic (the coefficients are + # just the tile size); tma_desc_ptr supplies this member's geometry. + cute.copy( + tma_atom_x, + tXgX[(None, (tile_y, tile_x))], + tXsX[(None, prod_state.index)], + tma_bar_ptr=pipeline_obj.producer_get_barrier(prod_state), + tma_desc_ptr=tmap.get_tensormap_ptr(desc_x, cute.AddressSpace.generic), + ) + # Notify the consumer that this SMEM buffer is ready for consumption + pipeline_obj.producer_commit(prod_state) + + @cute.jit + def _process_block( + self, + block_offset_Y_in_tensor, # Row offset of this chunk (global for single-tensor, else tensor-local) + block_id_X_in_tensor, # Column-chunk index within the tensor + rows, # Number of rows in this tensor (logical shape) + cols, # Number of columns in this tensor (logical shape) + tensor_base, # Int64 element offset of this tensor in the group (0 for single-tensor) + desc_x, # Per-tensor descriptors, already acquired by the caller (None if single-tensor) + desc_out_row, + desc_out_col, + mS_row, # Grouped rowwise scales + mS_col, # Grouped colwise scales + first_logical_dim, # First dimension of the grouped tensor + tmap, # TensorMapManager for managing TMA descriptors + warp_idx, + tidx, + sX, # SMEM input for this block + sO_row, # SMEM rowwise output for this block + sO_col, # SMEM colwise output for this block + tXsX, # Tiled sX for TMA + tXgX, # Tiled gX for TMA + tXsO_row, # Tiled sO_row for TMA + tXgO_row, # Tiled gO_row for TMA + tXsO_col, # Tiled sO_col for TMA + tXgO_col, # Tiled gO_col for TMA + tma_atom_x, # TMA atom for input + tma_atom_out_row, # TMA atom for rowwise output + tma_atom_out_col, # TMA atom for colwise output + mainloop_pipeline: cutlass.pipeline.PipelineTmaAsync, + prod_state, + cons_state, + ): + """Quantize one 128x128 chunk in STAGES slices of BUFF_DIM_Y rows.""" + cfg = self.cfg + block_offset_X = block_id_X_in_tensor * self.CHUNK_DIM_X + + # This tensor's rowwise scales + scale_rows = Int32(first_logical_dim) if cutlass.const_expr(cfg.IS_SINGLE_TENSOR) else rows + scale_base = ( + Int64(0) + if cutlass.const_expr(cfg.IS_SINGLE_TENSOR) + else tensor_base // Int64(MXFP8_BLOCK_SCALING_SIZE) + ) + + if cutlass.const_expr(cfg.ROWWISE): + # Rowwise scale's divisibility guarantee: (128, 4) + scale_row_stride = cute.round_up(cute.ceil_div(cols, MXFP8_BLOCK_SCALING_SIZE), 4) + # Advance to this tensor's rowwise scales + mS_row_t = cute.make_tensor( + cute.make_ptr( + Float8E8M0FNU, + mS_row.iterator.toint() + scale_base, + cute.AddressSpace.gmem, + assumed_align=4, + ), + cute.make_layout((scale_rows, scale_row_stride), stride=(scale_row_stride, 1)), + ) + mS_row_tiled = cute.zipped_divide( + mS_row_t, (self.BUFF_DIM_Y, self.CHUNK_DIM_X // MXFP8_BLOCK_SCALING_SIZE) + ) + + if cutlass.const_expr(cfg.COLWISE): + # Colwise scale's divisibility guarantee: (4, 128) + scale_col_stride = cute.round_up(cols, 128) + # Advance to this tensor's colwise scales + mS_col_t = cute.make_tensor( + cute.make_ptr( + Float8E8M0FNU, + mS_col.iterator.toint() + scale_base, + cute.AddressSpace.gmem, + assumed_align=4, + ), + cute.make_layout( + (scale_rows // MXFP8_BLOCK_SCALING_SIZE, scale_col_stride), + stride=(scale_col_stride, 1), + ), + ) + mS_col_tiled = cute.zipped_divide( + mS_col_t, (self.BUFF_DIM_Y // MXFP8_BLOCK_SCALING_SIZE, self.CHUNK_DIM_X) + ) + + cute.arch.sync_threads() + + # This chunk's coordinates in the tile grid (32x128 TMA boxes, not elements). + tile_id_Y = block_offset_Y_in_tensor // self.BUFF_DIM_Y + tile_id_X = block_id_X_in_tensor + + # Fill every buffer up front, then issue one more each time a stage is consumed. + for prologue_stage in cutlass.range_constexpr(self.PIPELINE_DEPTH): + if warp_idx == 0: + self._issue_load( + mainloop_pipeline, + prod_state, + tile_id_Y + prologue_stage, + tile_id_X, + tma_atom_x, + tXgX, + tXsX, + tmap, + desc_x, + ) + prod_state.advance() + + for stage in cutlass.range_constexpr(self.STAGES): + # Wait for at most DEPTH-1 iters on the fly, which means the the last DEPTH iter has finished + # so we can reuse its SMEM output buffer + # (input buffer is managed by the producer and consumer pipeline states) + if warp_idx == 0: + cute.arch.cp_async_bulk_wait_group(self.PIPELINE_DEPTH - 1, read=True) + # Wait for this stage's input buffer to be filled by the producer + mainloop_pipeline.consumer_wait(cons_state) + cute.arch.sync_threads() + sX_tile = sX[(None, cons_state.index)] + row_tile = tile_id_Y + stage + tile_row_start = block_offset_Y_in_tensor + stage * self.BUFF_DIM_Y + + if cutlass.const_expr(cfg.COLWISE): + quantize_colwise_mxfp8( + sX_tile, + None, + sO_col[(None, cons_state.index)], + cute.flatten(mS_col_tiled[(None, (row_tile, tile_id_X))]), + cfg.MAX_NORM_RCP, + tile_row_start, + block_offset_X, + scale_rows, + cols, + ACTIVATION=None, + DTYPE=cfg.DTYPE, + FP8_DTYPE=cfg.FP8_DTYPE, + SWIZZLE=False, + TILE_X=self.BUFF_DIM_X, + TILE_Y=self.BUFF_DIM_Y, + SKIP_MASKING=False, + ) + if cutlass.const_expr(cfg.ROWWISE): + quantize_rowwise_mxfp8( + sX_tile, + None, + sO_row[(None, cons_state.index)], + cute.flatten(mS_row_tiled[(None, (row_tile, tile_id_X))]), + cfg.MAX_NORM_RCP, + tile_row_start, + block_offset_X, + scale_rows, + cols, + ACTIVATION=None, + DTYPE=cfg.DTYPE, + FP8_DTYPE=cfg.FP8_DTYPE, + TILE_X=self.BUFF_DIM_X, + TILE_Y=self.BUFF_DIM_Y, + WAVES=self.WAVES, + THREADS_PER_BANK=self.THREADS_PER_BANK, + PACK_SIZE=self.PACK_SIZE, + SKIP_MASKING=False, + ) + + # Force consumer's write to SMEM to be visible to TMA stores later + cute.arch.fence_proxy("async.shared", space="cta") + # Only after everyone finishes computation then this stage can be considered as "consumed" + cute.arch.sync_threads() + # I'm done with my input SMEM buffer, so the producer can write the next stage's data into it + mainloop_pipeline.consumer_release(cons_state) + + # I just freed my input SMEM buffer (stage), so the producer now can use it for writing + # (stage+DEPTH) stage's data if that stage exists + if cutlass.const_expr(stage + self.PIPELINE_DEPTH < self.STAGES): + if warp_idx == 0: + self._issue_load( + mainloop_pipeline, + prod_state, + tile_id_Y + stage + self.PIPELINE_DEPTH, + tile_id_X, + tma_atom_x, + tXgX, + tXsX, + tmap, + desc_x, + ) + prod_state.advance() + + # Write result to GMEM via TMA + if warp_idx == 0: + if cutlass.const_expr(cfg.ROWWISE): + if cutlass.const_expr(cfg.IS_SINGLE_TENSOR): + cute.copy( + tma_atom_out_row, + tXsO_row[(None, cons_state.index)], + tXgO_row[(None, (row_tile, tile_id_X))], + ) + else: + cute.copy( + tma_atom_out_row, + tXsO_row[(None, cons_state.index)], + tXgO_row[(None, (row_tile, tile_id_X))], + tma_desc_ptr=tmap.get_tensormap_ptr( + desc_out_row, cute.AddressSpace.generic + ), + ) + if cutlass.const_expr(cfg.COLWISE): + if cutlass.const_expr(cfg.IS_SINGLE_TENSOR): + cute.copy( + tma_atom_out_col, + tXsO_col[(None, cons_state.index)], + tXgO_col[(None, (row_tile, tile_id_X))], + ) + else: + cute.copy( + tma_atom_out_col, + tXsO_col[(None, cons_state.index)], + tXgO_col[(None, (row_tile, tile_id_X))], + tma_desc_ptr=tmap.get_tensormap_ptr( + desc_out_col, cute.AddressSpace.generic + ), + ) + # Commit all TMA operations of this iteration + cute.arch.cp_async_bulk_commit_group() + + cons_state.advance() + + +def compile_cutedsl_function_from_cfg(cfg: MXFP8GroupQuantizeConfig): + """Return the compiled CuTeDSL function object for the given grouped config.""" + # CUDA requires the group's first logical dim to be a multiple of 128 (and each + # tensor's rows/cols likewise); MXFP8 needs the last dim divisible by 32. + sym_M = cute.sym_int32(divisibility=128) + sym_N = cute.sym_int32(divisibility=MXFP8_BLOCK_SCALING_SIZE) + logical_shape = (sym_M, sym_N) + + out_dtype = cfg.FP8_DTYPE + scale_dtype = cutlass.Float8E8M0FNU + + def g2d(dtype, align=16): + return cute.runtime.make_fake_compact_tensor( + dtype, + logical_shape, + stride_order=(1, 0), + memspace=cute.AddressSpace.gmem, + assumed_align=align, + ) + + def g1d(dtype, align=4): + return cute.runtime.make_fake_compact_tensor( + dtype, + (cute.sym_int32(),), + stride_order=(0,), + memspace=cute.AddressSpace.gmem, + assumed_align=align, + ) + + # The kernel only takes the base address of the scale buffers (per-tensor strides + # are derived from cols), so their fake shape is a flat 1D byte run. + in_fake = g2d(cfg.DTYPE) + out_row_fake = g2d(out_dtype) + out_col_fake = g2d(out_dtype) + scale_row_fake = g1d(scale_dtype) + scale_col_fake = g1d(scale_dtype) + offsets_fake = g1d(cutlass.Int64, align=8) + first_dims_fake = g1d(cutlass.Int64, align=8) + last_dims_fake = g1d(cutlass.Int64, align=8) + tensormaps_fake = cute.runtime.make_fake_compact_tensor( + cutlass.Int64, + (cute.sym_int32(), NUM_WORKSPACE_SLOTS, BYTES_PER_TENSORMAP // 8), + stride_order=(2, 1, 0), + memspace=cute.AddressSpace.gmem, + assumed_align=128, + ) + + from cutlass.utils import HardwareInfo # pylint: disable=import-outside-toplevel + + sm_count = HardwareInfo().get_device_multiprocessor_count() + kernel_obj = MXFP8GroupQuantizeKernel(cfg, sm_count) + return cute.compile( + kernel_obj, + in_fake, + out_row_fake, + out_col_fake, + scale_row_fake, + scale_col_fake, + offsets_fake, + first_dims_fake, + last_dims_fake, + tensormaps_fake, + cute.runtime.make_fake_stream(), + options="--enable-tvm-ffi", + ) + + +# TEMPORARY (demo only): same as compile_cutedsl_function_from_cfg but with every +# extent pinned to a compile-time constant, so traced layouts print concrete numbers +# (e.g. `(384,256):(1@1,1@0)`) instead of `?{div=128}`. Delete once done inspecting. +def compile_cutedsl_function_from_cfg_static( + cfg: MXFP8GroupQuantizeConfig, + M_total: int, + N: int, + num_tensors: int, + scale_row_numel: int, + scale_col_numel: int, +): + """Compile with fully static shapes. Only accepts inputs of exactly these extents.""" + logical_shape = (M_total, N) + + def g2d(dtype, align=16): + return cute.runtime.make_fake_compact_tensor( + dtype, + logical_shape, + stride_order=(1, 0), + memspace=cute.AddressSpace.gmem, + assumed_align=align, + ) + + def g1d(dtype, numel, align=4): + return cute.runtime.make_fake_compact_tensor( + dtype, + (numel,), + stride_order=(0,), + memspace=cute.AddressSpace.gmem, + assumed_align=align, + ) + + scale_dtype = cutlass.Float8E8M0FNU + tensormaps_fake = cute.runtime.make_fake_compact_tensor( + cutlass.Int64, + (num_tensors, NUM_WORKSPACE_SLOTS, BYTES_PER_TENSORMAP // 8), + stride_order=(2, 1, 0), + memspace=cute.AddressSpace.gmem, + assumed_align=128, + ) + + from cutlass.utils import HardwareInfo # pylint: disable=import-outside-toplevel + + sm_count = HardwareInfo().get_device_multiprocessor_count() + return cute.compile( + MXFP8GroupQuantizeKernel(cfg, sm_count), + g2d(cfg.DTYPE), + g2d(cfg.FP8_DTYPE), + g2d(cfg.FP8_DTYPE), + g1d(scale_dtype, scale_row_numel), + g1d(scale_dtype, scale_col_numel), + g1d(cutlass.Int64, num_tensors + 1, align=8), + g1d(cutlass.Int64, num_tensors, align=8), + g1d(cutlass.Int64, num_tensors, align=8), + tensormaps_fake, + cute.runtime.make_fake_stream(), + options="--enable-tvm-ffi", + ) + + +def get_mxfp8_group_quantization_function( + fn_name: str, + dtype: str, + fp8_dtype: str, + rowwise: bool, + colwise: bool, + shape_rep: str, +) -> bool: + """Compile the grouped MXFP8 quantize kernel for this config and register it in the + TVM-FFI global registry under EXACTLY `fn_name`. Returns True on success (the C++ + dispatcher then fetches it with GetGlobal(fn_name)); False if unsupported, so the + caller falls back to the CUDA C++ grouped kernel. + """ + 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 = MXFP8GroupQuantizeConfig( + dtype=dtype, + fp8_dtype=fp8_dtype, + rowwise=rowwise, + colwise=colwise, + shape_rep=shape_rep, + ) + except ValueError as e: + logger.warning( + "CuTeDSL grouped MXFP8 backend does not support this config, " + "falling back to the CUDA C++ kernel: %s", + e, + ) + return False + + logger.debug("Compiling CuTeDSL grouped MXFP8 quantization kernel for %s", cfg) + try: + compiled = compile_cutedsl_function_from_cfg(cfg) + except Exception as e: # pylint: disable=broad-exception-caught + logger.error( + "CuTeDSL grouped MXFP8 kernel compilation failed, " + "falling back to the CUDA C++ kernel: %s", + e, + ) + return False + tvm_ffi.register_global_func(fn_name, compiled, override=True) + return True diff --git a/transformer_engine/common/cast/dispatch/quantize.cuh b/transformer_engine/common/cast/dispatch/quantize.cuh index aa1d16fd894..ae68069b049 100644 --- a/transformer_engine/common/cast/dispatch/quantize.cuh +++ b/transformer_engine/common/cast/dispatch/quantize.cuh @@ -541,9 +541,18 @@ void group_quantize_fwd_helper(const NVTEGroupedTensor input, NVTEGroupedTensor break; } case NVTE_MXFP8_1D_SCALING: { - mxfp8::group_quantize( - input_tensor, activations_tensor, noop_tensor, output_tensor, dbias_tensor, - workspace_tensor, &quant_config_cpp, stream); + #ifdef NVTE_WITH_CUTEDSL + quantized_with_cutedsl = + cutedsl_backend::mxfp8_group_quantize_cutedsl( + input_tensor, noop_tensor, output_tensor, quant_config_cpp.mxfp8_2d_quantization, + stream); +#endif + if (!quantized_with_cutedsl) { + mxfp8::group_quantize( + input_tensor, activations_tensor, noop_tensor, output_tensor, dbias_tensor, + workspace_tensor, &quant_config_cpp, stream); + } break; } case NVTE_BLOCK_SCALING_1D: { diff --git a/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8_cutedsl.cuh b/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8_cutedsl.cuh new file mode 100644 index 00000000000..ca33f0b25ef --- /dev/null +++ b/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8_cutedsl.cuh @@ -0,0 +1,312 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#ifndef TRANSFORMER_ENGINE_COMMON_CAST_MXFP8_GROUP_QUANTIZE_MXFP8_CUTEDSL_CUH_ +#define TRANSFORMER_ENGINE_COMMON_CAST_MXFP8_GROUP_QUANTIZE_MXFP8_CUTEDSL_CUH_ + +#include +#include + +#include +#include +#include +#include +#include + +#include "../../common.h" +#include "../../tvm_ffi_bridge.h" +#include "../../utils.cuh" // ShapeRepresentation +#include "../core/common.cuh" +#include "group_quantize_mxfp8.cuh" + +namespace transformer_engine { +namespace cutedsl_backend { + +// te_dtype_to_str, DLTensorWrapper, TVMFFICentral all live in +// transformer_engine::tvm_ffi_bridge (tvm_ffi_bridge.h). +using namespace tvm_ffi_bridge; + +inline const char *shape_rep_to_str(ShapeRepresentation shape_rep) { + switch (shape_rep) { + case ShapeRepresentation::SAME_BOTH_DIMS: + return "same_both_dims"; + case ShapeRepresentation::VARYING_FIRST_DIM: + return "varying_first_dim"; + case ShapeRepresentation::VARYING_LAST_DIM: + return "varying_last_dim"; + default: + return "varying_both_dims"; + } +} + +struct MXFP8GroupQuantConfig { + static constexpr const char *kEntrypointName = "get_mxfp8_group_quantization_function"; + + DType dtype; // The input format + DType fp8_dtype; // The fp8 output format + bool rowwise; // If quantize rowwisely + bool colwise; // If quantize columnwisely + ShapeRepresentation shape_rep; // How the member shapes vary across the group + + constexpr uint32_t to_id() const { + static_assert(static_cast(DType::kNumTypes) <= 256, + "DType no longer fits in the 8 bits to_id() gives it."); + return static_cast(dtype) | (static_cast(fp8_dtype) << 8) | + (static_cast(rowwise) << 16) | (static_cast(colwise) << 17) | + (static_cast(shape_rep) << 18); + } + + std::optional get_kernel() const { + static TVMFFIConfigCache &cache = TVMFFIConfigCache::create(); + return cache.get_or_load(*this); + } + + // Globally unique TVM-FFI registry key used when the CuTeDSL function is + // compiled and registered on a cache miss. + std::string to_key() const { + std::string key; + key.reserve(64); // longest: cutedsl_group_mxfp8_bf16_fp8_e4m3fn_1_1_varying_first_dim + key.append("cutedsl_group_mxfp8_") + .append(te_dtype_to_str(dtype)) + .append("_") + .append(te_dtype_to_str(fp8_dtype)) + .append("_") + .append(rowwise ? "1" : "0") + .append("_") + .append(colwise ? "1" : "0") + .append("_") + .append(shape_rep_to_str(shape_rep)); + return key; + } + + bool retrieve_func_from_python(const std::string &fn_name) const { + auto entrypoint = tvm::ffi::Function::GetGlobal(kEntrypointName); + if (!entrypoint.has_value()) { + return false; + } + tvm::ffi::Any result = + (*entrypoint)(tvm::ffi::String(fn_name), tvm::ffi::String(te_dtype_to_str(dtype)), + tvm::ffi::String(te_dtype_to_str(fp8_dtype)), rowwise, colwise, + tvm::ffi::String(shape_rep_to_str(shape_rep))); + return result.try_cast().value_or(false); + } +}; + +// Descriptor slots per group member: input, rowwise output, colwise output, plus one +// carrying (rows, cols, base_elts). Mirrors NUM_WORKSPACE_SLOTS / BYTES_PER_TENSORMAP in +// CuTeDSL/cast/mxfp8/group_quantize_mxfp8.py. +constexpr size_t kGroupTensorMapSlots = 4; +constexpr size_t kInt64PerTensorMap = 128 / sizeof(int64_t); + +struct alignas(128) GroupDescriptorWorkspace { + alignas(128) int64_t tensor_maps[dispatch::common::MAX_SUPPORTED_TENSOR_DESCRIPTORS] + [kGroupTensorMapSlots][kInt64PerTensorMap]; + // Stand-in for the offsets / first_dims / last_dims arrays a given shape representation + // does not carry: the kernel takes all three unconditionally but only dereferences the + // ones its representation uses, so the contents are never read. Sized num_tensors + 1 + // for the CSR offsets array, the longest of the three. + int64_t unused_dims[dispatch::common::MAX_SUPPORTED_TENSOR_DESCRIPTORS + 1]; +}; + +// One workspace per translation unit, mirroring `g_tensor_maps` on the CUDA path -- and +// sharing its caveat that two grouped quantize calls in flight on different streams would +// overwrite each other's descriptors. +static __device__ GroupDescriptorWorkspace g_group_descriptor_workspace; + +inline GroupDescriptorWorkspace *group_descriptor_workspace_ptr() { + static GroupDescriptorWorkspace *const ptr = [] { + void *p = nullptr; + NVTE_CHECK_CUDA(cudaGetSymbolAddress(&p, g_group_descriptor_workspace)); + return static_cast(p); + }(); + return ptr; +} + +inline NVTEBasicTensor make_basic_tensor(void *dptr, DType dtype, + const std::vector &shape) { + return NVTEBasicTensor{dptr, static_cast(dtype), + nvte_make_shape(shape.data(), shape.size())}; +} + +// Signature mirrors mxfp8::group_quantize (input, output, stream) for the subset the +// CuTeDSL kernel covers. Returns false to fall back to the CUDA kernel. +inline bool mxfp8_group_quantize_cutedsl(const MXFP8GroupQuantConfig &config, + const GroupedTensor *input_tensor, + GroupedTensor *output_tensor, cudaStream_t stream) { + using namespace dispatch::mxfp8::group_quantize_kernel; + + const size_t num_tensors = input_tensor->num_tensors; + const size_t first_logical_dim = input_tensor->logical_shape.data[0]; + const size_t last_logical_dim = input_tensor->logical_shape.data[1]; + + // The kernel is compiled with cute.sym_int32(divisibility=...) on both logical extents, + // so a violating shape would silently mis-tile rather than fail. These mirror sym_M / + // sym_N in CuTeDSL/cast/mxfp8/group_quantize_mxfp8.py -- the DSL kernel's own chunk + // height and MXFP8 block size, which it tiles independently of the CUDA kernel's + // CastTraits. + constexpr size_t kChunkDimY = 128; + constexpr size_t kScaleDimX = 32; + if (first_logical_dim % kChunkDimY != 0 || last_logical_dim % kScaleDimX != 0) { + maybe_warn_cutedsl_not_chosen("the grouped logical shape is not a multiple of (", kChunkDimY, + ", ", kScaleDimX, ")."); + return false; + } + + std::optional group_quant_func_opt = config.get_kernel(); + if (!group_quant_func_opt.has_value()) { + return false; + } + + GroupDescriptorWorkspace *const workspace = group_descriptor_workspace_ptr(); + + // Both output directions are handed to the kernel unconditionally: the compiled + // signature has no optional tensors, and building a TMA descriptor needs a real + // address for each. The disabled direction is never read or written, so it points at + // the enabled one instead of at a buffer that would have to be allocated. + const SimpleTensor &data_row = + config.rowwise ? output_tensor->data : output_tensor->columnwise_data; + const SimpleTensor &data_col = + config.colwise ? output_tensor->columnwise_data : output_tensor->data; + const SimpleTensor &scale_row = + config.rowwise ? output_tensor->scale_inv : output_tensor->columnwise_scale_inv; + const SimpleTensor &scale_col = + config.colwise ? output_tensor->columnwise_scale_inv : output_tensor->scale_inv; + + // The group's payload is stored flat; the kernel wants it as the logical 2D view. + const std::vector logical_shape{first_logical_dim, last_logical_dim}; + const NVTEBasicTensor x_bt = + make_basic_tensor(input_tensor->data.dptr, input_tensor->dtype(), logical_shape); + const NVTEBasicTensor o_row_bt = make_basic_tensor(data_row.dptr, data_row.dtype, logical_shape); + const NVTEBasicTensor o_col_bt = make_basic_tensor(data_col.dptr, data_col.dtype, logical_shape); + DLTensorWrapper mX(x_bt), mO_row(o_row_bt), mO_col(o_col_bt); + + // The kernel only takes the base address of the scale buffers (per-tensor bases and + // strides are derived from the member shapes), so these stay 1D. + const NVTEBasicTensor s_row_bt = + make_basic_tensor(scale_row.dptr, scale_row.dtype, {scale_row.numel()}); + const NVTEBasicTensor s_col_bt = + make_basic_tensor(scale_col.dptr, scale_col.dtype, {scale_col.numel()}); + DLTensorWrapper mS_row(s_row_bt, /*flatten_2D=*/false), mS_col(s_col_bt, /*flatten_2D=*/false); + + const SimpleTensor &offsets = output_tensor->tensor_offsets; + const SimpleTensor &first_dims = output_tensor->first_dims; + const SimpleTensor &last_dims = output_tensor->last_dims; + const NVTEBasicTensor offsets_bt = make_basic_tensor( + offsets.has_data() ? offsets.dptr : static_cast(workspace->unused_dims), + DType::kInt64, {num_tensors + 1}); + const NVTEBasicTensor first_dims_bt = make_basic_tensor( + first_dims.has_data() ? first_dims.dptr : static_cast(workspace->unused_dims), + DType::kInt64, {num_tensors}); + const NVTEBasicTensor last_dims_bt = make_basic_tensor( + last_dims.has_data() ? last_dims.dptr : static_cast(workspace->unused_dims), + DType::kInt64, {num_tensors}); + DLTensorWrapper mOffsets(offsets_bt, /*flatten_2D=*/false), + mFirstDims(first_dims_bt, /*flatten_2D=*/false), + mLastDims(last_dims_bt, /*flatten_2D=*/false); + + // The kernel reads num_tensors off this tensor's leading extent, so it must be exactly + // the group size even on the single-tensor path that leaves the descriptors untouched. + const NVTEBasicTensor tensormaps_bt = + make_basic_tensor(static_cast(workspace->tensor_maps), DType::kInt64, + {num_tensors, kGroupTensorMapSlots, kInt64PerTensorMap}); + DLTensorWrapper mTensormaps(tensormaps_bt, /*flatten_2D=*/false); + + // stream is a tvm-ffi opaque "handle"; pass it as void*. + (*group_quant_func_opt)(&mX, &mO_row, &mO_col, &mS_row, &mS_col, &mOffsets, &mFirstDims, + &mLastDims, &mTensormaps, static_cast(stream)); + return true; +} + +template +bool mxfp8_group_quantize_cutedsl(const GroupedTensor *input_tensor, const Tensor *noop_tensor, + GroupedTensor *output_tensor, const bool use_2d_quantization, + cudaStream_t stream) { + if (!tvm_ffi_bridge::TVMFFICentral::getInstance().get_cutedsl_backend_enabled()) { + maybe_warn_cutedsl_not_chosen("the CuTeDSL backend is disabled."); + return false; + } + // The CuTeDSL grouped kernel is cast-only: no dbias, no fused (derivative) activation. + if constexpr (IS_DBIAS || IS_DACT || IS_ACT || OP != nullptr) { + maybe_warn_cutedsl_not_chosen( + "grouped quantization with dbias or a fused activation is not supported."); + return false; + } else { + // TODO(kainingz): port 2D quantization to CuTeDSL + if (use_2d_quantization) { + maybe_warn_cutedsl_not_chosen("2D quantization is not supported."); + return false; + } + // The kernel takes no noop flag, no amax accumulator, and writes compact scales only. + if (noop_tensor != nullptr && noop_tensor->data.dptr != nullptr) { + maybe_warn_cutedsl_not_chosen("the cast-noop flag is not supported."); + return false; + } + if (output_tensor->amax.dptr != nullptr) { + maybe_warn_cutedsl_not_chosen("amax computation is not supported."); + return false; + } + if (output_tensor->with_gemm_swizzled_scales) { + maybe_warn_cutedsl_not_chosen("GEMM-swizzled scales are not supported."); + return false; + } + + // Mirrors the shape-representation selection in mxfp8::group_quantize. + ShapeRepresentation shape_rep = ShapeRepresentation::SAME_BOTH_DIMS; + if (output_tensor->all_same_shape()) { + shape_rep = ShapeRepresentation::SAME_BOTH_DIMS; + } else if (output_tensor->all_same_first_dim()) { + shape_rep = ShapeRepresentation::VARYING_LAST_DIM; + } else if (output_tensor->all_same_last_dim()) { + shape_rep = ShapeRepresentation::VARYING_FIRST_DIM; + } else { + // VARYING_BOTH_DIMS: the logical shape is [1, total], which is not tileable. + maybe_warn_cutedsl_not_chosen("groups with both dimensions varying are not supported."); + return false; + } + // Every member gets a descriptor slot in the fixed-size workspace, so the CUDA + // kernel's descriptor limit applies to the single-tensor representations here too. + if (input_tensor->num_tensors > dispatch::common::MAX_SUPPORTED_TENSOR_DESCRIPTORS) { + maybe_warn_cutedsl_not_chosen("the group has more than ", + dispatch::common::MAX_SUPPORTED_TENSOR_DESCRIPTORS, + " tensors."); + return false; + } + + const bool rowwise = output_tensor->has_data(); + const bool colwise = output_tensor->has_columnwise_data(); + if (!rowwise && !colwise) { + // mxfp8::group_quantize raises a proper error for this. + return false; + } + + checkCuDriverContext(stream); + // Sanity checks, mirroring mxfp8::group_quantize + if (rowwise) { + NVTE_CHECK(output_tensor->scale_inv.dptr != nullptr, "Scaling tensor must be allocated"); + } + if (colwise) { + NVTE_CHECK(output_tensor->columnwise_scale_inv.dptr != nullptr, + "Columnwise scaling tensor must be allocated"); + } + NVTE_CHECK(input_tensor->num_tensors == output_tensor->num_tensors, + "Number of input and output tensors must be same."); + NVTE_CHECK(input_tensor->has_data(), "Cannot quantize tensor without rowwise data."); + NVTE_CHECK(is_fp8_dtype(output_tensor->dtype()), "Output must have FP8 type."); + + const MXFP8GroupQuantConfig config{/*dtype=*/input_tensor->dtype(), + /*fp8_dtype=*/output_tensor->dtype(), + /*rowwise=*/rowwise, + /*colwise=*/colwise, + /*shape_rep=*/shape_rep}; + return mxfp8_group_quantize_cutedsl(config, input_tensor, output_tensor, stream); + } +} + +} // namespace cutedsl_backend +} // namespace transformer_engine + +#endif // TRANSFORMER_ENGINE_COMMON_CAST_MXFP8_GROUP_QUANTIZE_MXFP8_CUTEDSL_CUH_ From 79d521712ad118db9944b871a231dfab0db20fe0 Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Tue, 29 Sep 2026 21:15:15 +0000 Subject: [PATCH 02/58] fix Signed-off-by: Kaining Zhong --- transformer_engine/common/CuTeDSL/__init__.py | 8 + .../common/CuTeDSL/cast/mxfp8/__init__.py | 1 + .../cast/mxfp8/group_quantize_mxfp8.py | 155 ++++++----------- .../common/cast/dispatch/quantize.cuh | 4 +- .../mxfp8/group_quantize_mxfp8_cutedsl.cuh | 158 ++++++++++-------- 5 files changed, 155 insertions(+), 171 deletions(-) diff --git a/transformer_engine/common/CuTeDSL/__init__.py b/transformer_engine/common/CuTeDSL/__init__.py index 50d536852c2..6698e22965b 100644 --- a/transformer_engine/common/CuTeDSL/__init__.py +++ b/transformer_engine/common/CuTeDSL/__init__.py @@ -13,6 +13,9 @@ from transformer_engine.common.CuTeDSL.cast.mxfp8.quantize_mxfp8 import ( get_mxfp8_quantization_function, ) +from transformer_engine.common.CuTeDSL.cast.mxfp8.group_quantize_mxfp8 import ( + get_mxfp8_group_quantization_function, +) def register_cutedsl_backends(): @@ -22,3 +25,8 @@ def register_cutedsl_backends(): tvm_ffi.register_global_func( "get_mxfp8_quantization_function", get_mxfp8_quantization_function, override=True ) + tvm_ffi.register_global_func( + "get_mxfp8_group_quantization_function", + get_mxfp8_group_quantization_function, + override=True, + ) diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/__init__.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/__init__.py index 6491573690f..97cc599ce29 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/__init__.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/__init__.py @@ -5,3 +5,4 @@ """CuTeDSL MXFP8 quantization kernels.""" from . import quantize_mxfp8 +from . import group_quantize_mxfp8 diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8.py index 11c288a7d36..116b07200f6 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8.py @@ -110,12 +110,14 @@ class MXFP8GroupQuantizeConfig: """Compile-time config for the grouped MXFP8 quantize kernel.""" def __init__(self, dtype: str, fp8_dtype: str, rowwise: bool, colwise: bool, shape_rep: str): - if dtype not in ("fp32", "fp16", "bf16"): - raise ValueError(f"unknown input dtype {dtype!r}; expected fp32|fp16|bf16") + if 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 - if fp8_dtype not in ("fp8_e4m3fn", "fp8_e5m2"): - raise ValueError(f"unknown FP8 dtype {fp8_dtype!r}; expected fp8_e4m3fn|fp8_e5m2") + 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 if not (rowwise or colwise): @@ -131,7 +133,7 @@ def __init__(self, dtype: str, fp8_dtype: str, rowwise: bool, colwise: bool, sha # Mirrors `is_single_tensor` in group_quantize_mxfp8.cuh. self.IS_SINGLE_TENSOR = shape_rep in (SAME_BOTH_DIMS, VARYING_FIRST_DIM) self.MAX_NORM_RCP = ( - FP8E4M3_MAX_NORM_RCP if fp8_dtype == "fp8_e4m3fn" else FP8E5M2_MAX_NORM_RCP + FP8E4M3_MAX_NORM_RCP if fp8_dtype == "Float8E4M3" else FP8E5M2_MAX_NORM_RCP ) def __str__(self): @@ -1020,66 +1022,6 @@ def g1d(dtype, align=4): ) -# TEMPORARY (demo only): same as compile_cutedsl_function_from_cfg but with every -# extent pinned to a compile-time constant, so traced layouts print concrete numbers -# (e.g. `(384,256):(1@1,1@0)`) instead of `?{div=128}`. Delete once done inspecting. -def compile_cutedsl_function_from_cfg_static( - cfg: MXFP8GroupQuantizeConfig, - M_total: int, - N: int, - num_tensors: int, - scale_row_numel: int, - scale_col_numel: int, -): - """Compile with fully static shapes. Only accepts inputs of exactly these extents.""" - logical_shape = (M_total, N) - - def g2d(dtype, align=16): - return cute.runtime.make_fake_compact_tensor( - dtype, - logical_shape, - stride_order=(1, 0), - memspace=cute.AddressSpace.gmem, - assumed_align=align, - ) - - def g1d(dtype, numel, align=4): - return cute.runtime.make_fake_compact_tensor( - dtype, - (numel,), - stride_order=(0,), - memspace=cute.AddressSpace.gmem, - assumed_align=align, - ) - - scale_dtype = cutlass.Float8E8M0FNU - tensormaps_fake = cute.runtime.make_fake_compact_tensor( - cutlass.Int64, - (num_tensors, NUM_WORKSPACE_SLOTS, BYTES_PER_TENSORMAP // 8), - stride_order=(2, 1, 0), - memspace=cute.AddressSpace.gmem, - assumed_align=128, - ) - - from cutlass.utils import HardwareInfo # pylint: disable=import-outside-toplevel - - sm_count = HardwareInfo().get_device_multiprocessor_count() - return cute.compile( - MXFP8GroupQuantizeKernel(cfg, sm_count), - g2d(cfg.DTYPE), - g2d(cfg.FP8_DTYPE), - g2d(cfg.FP8_DTYPE), - g1d(scale_dtype, scale_row_numel), - g1d(scale_dtype, scale_col_numel), - g1d(cutlass.Int64, num_tensors + 1, align=8), - g1d(cutlass.Int64, num_tensors, align=8), - g1d(cutlass.Int64, num_tensors, align=8), - tensormaps_fake, - cute.runtime.make_fake_stream(), - options="--enable-tvm-ffi", - ) - - def get_mxfp8_group_quantization_function( fn_name: str, dtype: str, @@ -1088,49 +1030,58 @@ def get_mxfp8_group_quantization_function( colwise: bool, shape_rep: str, ) -> bool: - """Compile the grouped MXFP8 quantize kernel for this config and register it in the - TVM-FFI global registry under EXACTLY `fn_name`. Returns True on success (the C++ - dispatcher then fetches it with GetGlobal(fn_name)); False if unsupported, so the - caller falls back to the CUDA C++ grouped kernel. + """Compile the grouped 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++ grouped kernel. """ - 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 = MXFP8GroupQuantizeConfig( - dtype=dtype, - fp8_dtype=fp8_dtype, - rowwise=rowwise, - colwise=colwise, - shape_rep=shape_rep, - ) - except ValueError as e: - logger.warning( - "CuTeDSL grouped MXFP8 backend does not support this config, " - "falling back to the CUDA C++ kernel: %s", - e, - ) - return False + # 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 = MXFP8GroupQuantizeConfig( + dtype=dtype, + fp8_dtype=fp8_dtype, + rowwise=rowwise, + colwise=colwise, + shape_rep=shape_rep, + ) + except ValueError as e: + logger.warning( + "CuTeDSL grouped MXFP8 backend does not support this config, " + "falling back to the CUDA C++ kernel: %s", + e, + ) + return False - logger.debug("Compiling CuTeDSL grouped MXFP8 quantization kernel for %s", cfg) - try: + logger.debug("Compiling CuTeDSL grouped MXFP8 quantization kernel for %s", cfg) compiled = compile_cutedsl_function_from_cfg(cfg) + # Register the native TVM-FFI function rather than its Python argument-parsing wrapper; + # see get_mxfp8_quantization_function for why. + 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 grouped MXFP8 kernel compilation failed, " - "falling back to the CUDA C++ kernel: %s", + "CuTeDSL grouped 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 - tvm_ffi.register_global_func(fn_name, compiled, override=True) - return True diff --git a/transformer_engine/common/cast/dispatch/quantize.cuh b/transformer_engine/common/cast/dispatch/quantize.cuh index ae68069b049..ebed509813f 100644 --- a/transformer_engine/common/cast/dispatch/quantize.cuh +++ b/transformer_engine/common/cast/dispatch/quantize.cuh @@ -32,6 +32,7 @@ #include "../nvfp4/quantize_transpose_nvfp4.cuh" #ifdef NVTE_WITH_CUTEDSL +#include "../mxfp8/group_quantize_mxfp8_cutedsl.cuh" #include "../mxfp8/quantize_mxfp8_cutedsl.cuh" #endif @@ -541,7 +542,8 @@ void group_quantize_fwd_helper(const NVTEGroupedTensor input, NVTEGroupedTensor break; } case NVTE_MXFP8_1D_SCALING: { - #ifdef NVTE_WITH_CUTEDSL + bool quantized_with_cutedsl = false; +#ifdef NVTE_WITH_CUTEDSL quantized_with_cutedsl = cutedsl_backend::mxfp8_group_quantize_cutedsl( diff --git a/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8_cutedsl.cuh b/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8_cutedsl.cuh index ca33f0b25ef..b5b72ad20fc 100644 --- a/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8_cutedsl.cuh +++ b/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8_cutedsl.cuh @@ -12,21 +12,21 @@ #include #include +#include #include #include #include #include "../../common.h" #include "../../tvm_ffi_bridge.h" -#include "../../utils.cuh" // ShapeRepresentation -#include "../core/common.cuh" -#include "group_quantize_mxfp8.cuh" +#include "../../util/cuda_runtime.h" +#include "../../utils.cuh" // ShapeRepresentation +#include "../core/grouped_tma.cuh" // dispatch::common::MAX_SUPPORTED_TENSOR_DESCRIPTORS namespace transformer_engine { namespace cutedsl_backend { -// te_dtype_to_str, DLTensorWrapper, TVMFFICentral all live in -// transformer_engine::tvm_ffi_bridge (tvm_ffi_bridge.h). +// DLTensorWrapper and TVMFFICentral live in transformer_engine::tvm_ffi_bridge. using namespace tvm_ffi_bridge; inline const char *shape_rep_to_str(ShapeRepresentation shape_rep) { @@ -50,16 +50,23 @@ struct MXFP8GroupQuantConfig { bool rowwise; // If quantize rowwisely bool colwise; // If quantize columnwisely ShapeRepresentation shape_rep; // How the member shapes vary across the group - - constexpr uint32_t to_id() const { - static_assert(static_cast(DType::kNumTypes) <= 256, - "DType no longer fits in the 8 bits to_id() gives it."); - return static_cast(dtype) | (static_cast(fp8_dtype) << 8) | - (static_cast(rowwise) << 16) | (static_cast(colwise) << 17) | - (static_cast(shape_rep) << 18); + uint32_t sm_arch = static_cast(cuda::sm_arch()); + + // Bit layout: dtype [3:0] (4 used/reserved), fp8_dtype [7:4] (4 used/reserved), + // flags [9:8] (2 used), shape_rep [11:10] (2 used), and SM architecture [20:12] + // (9 used/reserved). Bits [31:21] are unused. + uint32_t to_id() const { + static_assert(static_cast(DType::kNumTypes) <= 16, + "DType no longer fits in the 4 bits to_id() gives it."); + static_assert(ShapeRepresentation::VARYING_BOTH_DIMS < 4, + "ShapeRepresentation no longer fits in the 2 bits to_id() gives it."); + NVTE_CHECK(sm_arch < 512, "SM architecture no longer fits in the 9 bits to_id() gives it."); + return static_cast(dtype) | (static_cast(fp8_dtype) << 4) | + (static_cast(rowwise) << 8) | (static_cast(colwise) << 9) | + (static_cast(shape_rep) << 10) | (sm_arch << 12); } - std::optional get_kernel() const { + std::optional get_kernel() const { static TVMFFIConfigCache &cache = TVMFFIConfigCache::create(); return cache.get_or_load(*this); } @@ -68,11 +75,14 @@ struct MXFP8GroupQuantConfig { // compiled and registered on a cache miss. std::string to_key() const { std::string key; - key.reserve(64); // longest: cutedsl_group_mxfp8_bf16_fp8_e4m3fn_1_1_varying_first_dim - key.append("cutedsl_group_mxfp8_") - .append(te_dtype_to_str(dtype)) + key.reserve( + 80); // longest: cutedsl_group_mxfp8_smXXX_BFloat16_Float8E4M3_1_1_varying_first_dim + key.append("cutedsl_group_mxfp8_sm") + .append(std::to_string(sm_arch)) + .append("_") + .append(to_string(dtype)) .append("_") - .append(te_dtype_to_str(fp8_dtype)) + .append(to_string(fp8_dtype)) .append("_") .append(rowwise ? "1" : "0") .append("_") @@ -88,8 +98,8 @@ struct MXFP8GroupQuantConfig { return false; } tvm::ffi::Any result = - (*entrypoint)(tvm::ffi::String(fn_name), tvm::ffi::String(te_dtype_to_str(dtype)), - tvm::ffi::String(te_dtype_to_str(fp8_dtype)), rowwise, colwise, + (*entrypoint)(tvm::ffi::String(fn_name), tvm::ffi::String(to_string(dtype)), + tvm::ffi::String(to_string(fp8_dtype)), rowwise, colwise, tvm::ffi::String(shape_rep_to_str(shape_rep))); return result.try_cast().value_or(false); } @@ -100,29 +110,39 @@ struct MXFP8GroupQuantConfig { // CuTeDSL/cast/mxfp8/group_quantize_mxfp8.py. constexpr size_t kGroupTensorMapSlots = 4; constexpr size_t kInt64PerTensorMap = 128 / sizeof(int64_t); +constexpr size_t kMaxGroupTensors = + static_cast(dispatch::common::MAX_SUPPORTED_TENSOR_DESCRIPTORS); struct alignas(128) GroupDescriptorWorkspace { - alignas(128) int64_t tensor_maps[dispatch::common::MAX_SUPPORTED_TENSOR_DESCRIPTORS] - [kGroupTensorMapSlots][kInt64PerTensorMap]; + alignas(128) int64_t tensor_maps[kMaxGroupTensors][kGroupTensorMapSlots][kInt64PerTensorMap]; // Stand-in for the offsets / first_dims / last_dims arrays a given shape representation // does not carry: the kernel takes all three unconditionally but only dereferences the // ones its representation uses, so the contents are never read. Sized num_tensors + 1 // for the CSR offsets array, the longest of the three. - int64_t unused_dims[dispatch::common::MAX_SUPPORTED_TENSOR_DESCRIPTORS + 1]; + int64_t unused_dims[kMaxGroupTensors + 1]; }; -// One workspace per translation unit, mirroring `g_tensor_maps` on the CUDA path -- and -// sharing its caveat that two grouped quantize calls in flight on different streams would -// overwrite each other's descriptors. +// Like `g_tensor_maps` on the CUDA path, this has internal linkage, so every translation +// unit including this header gets its own copy. It shares that path's caveat that two +// grouped quantize calls in flight on different streams would overwrite each other's +// descriptors. static __device__ GroupDescriptorWorkspace g_group_descriptor_workspace; -inline GroupDescriptorWorkspace *group_descriptor_workspace_ptr() { - static GroupDescriptorWorkspace *const ptr = [] { +// Device address of this translation unit's workspace on the current device. The address +// is per device (each device context loads its own copy of the module), so it is cached +// per device. `static` rather than `inline`: it refers to the internal-linkage symbol above, +// so each translation unit needs its own definition and cache. +static GroupDescriptorWorkspace *group_descriptor_workspace_ptr() { + static std::vector cache(cuda::num_devices(), nullptr); + static std::vector flags(cuda::num_devices()); + const int device_id = cuda::current_device(); + NVTE_CHECK(0 <= device_id && device_id < cuda::num_devices(), "invalid CUDA device ID"); + std::call_once(flags[device_id], [&]() { void *p = nullptr; NVTE_CHECK_CUDA(cudaGetSymbolAddress(&p, g_group_descriptor_workspace)); - return static_cast(p); - }(); - return ptr; + cache[device_id] = static_cast(p); + }); + return cache[device_id]; } inline NVTEBasicTensor make_basic_tensor(void *dptr, DType dtype, @@ -136,8 +156,6 @@ inline NVTEBasicTensor make_basic_tensor(void *dptr, DType dtype, inline bool mxfp8_group_quantize_cutedsl(const MXFP8GroupQuantConfig &config, const GroupedTensor *input_tensor, GroupedTensor *output_tensor, cudaStream_t stream) { - using namespace dispatch::mxfp8::group_quantize_kernel; - const size_t num_tensors = input_tensor->num_tensors; const size_t first_logical_dim = input_tensor->logical_shape.data[0]; const size_t last_logical_dim = input_tensor->logical_shape.data[1]; @@ -154,12 +172,19 @@ inline bool mxfp8_group_quantize_cutedsl(const MXFP8GroupQuantConfig &config, ", ", kScaleDimX, ")."); return false; } + // The same extents are sym_int32 in the compiled kernel. + if (first_logical_dim > static_cast(INT32_MAX) || + last_logical_dim > static_cast(INT32_MAX)) { + maybe_warn_cutedsl_not_chosen("the grouped logical shape does not fit in int32."); + return false; + } - std::optional group_quant_func_opt = config.get_kernel(); + std::optional group_quant_func_opt = config.get_kernel(); if (!group_quant_func_opt.has_value()) { return false; } + const int32_t device_index = transformer_engine::cuda::current_device(); GroupDescriptorWorkspace *const workspace = group_descriptor_workspace_ptr(); // Both output directions are handed to the kernel unconditionally: the compiled @@ -177,42 +202,36 @@ inline bool mxfp8_group_quantize_cutedsl(const MXFP8GroupQuantConfig &config, // The group's payload is stored flat; the kernel wants it as the logical 2D view. const std::vector logical_shape{first_logical_dim, last_logical_dim}; - const NVTEBasicTensor x_bt = - make_basic_tensor(input_tensor->data.dptr, input_tensor->dtype(), logical_shape); - const NVTEBasicTensor o_row_bt = make_basic_tensor(data_row.dptr, data_row.dtype, logical_shape); - const NVTEBasicTensor o_col_bt = make_basic_tensor(data_col.dptr, data_col.dtype, logical_shape); - DLTensorWrapper mX(x_bt), mO_row(o_row_bt), mO_col(o_col_bt); + DLTensorWrapper mX( + make_basic_tensor(input_tensor->data.dptr, input_tensor->dtype(), logical_shape), true, + device_index); + DLTensorWrapper mO_row(make_basic_tensor(data_row.dptr, data_row.dtype, logical_shape), true, + device_index); + DLTensorWrapper mO_col(make_basic_tensor(data_col.dptr, data_col.dtype, logical_shape), true, + device_index); // The kernel only takes the base address of the scale buffers (per-tensor bases and // strides are derived from the member shapes), so these stay 1D. - const NVTEBasicTensor s_row_bt = - make_basic_tensor(scale_row.dptr, scale_row.dtype, {scale_row.numel()}); - const NVTEBasicTensor s_col_bt = - make_basic_tensor(scale_col.dptr, scale_col.dtype, {scale_col.numel()}); - DLTensorWrapper mS_row(s_row_bt, /*flatten_2D=*/false), mS_col(s_col_bt, /*flatten_2D=*/false); - - const SimpleTensor &offsets = output_tensor->tensor_offsets; - const SimpleTensor &first_dims = output_tensor->first_dims; - const SimpleTensor &last_dims = output_tensor->last_dims; - const NVTEBasicTensor offsets_bt = make_basic_tensor( - offsets.has_data() ? offsets.dptr : static_cast(workspace->unused_dims), - DType::kInt64, {num_tensors + 1}); - const NVTEBasicTensor first_dims_bt = make_basic_tensor( - first_dims.has_data() ? first_dims.dptr : static_cast(workspace->unused_dims), - DType::kInt64, {num_tensors}); - const NVTEBasicTensor last_dims_bt = make_basic_tensor( - last_dims.has_data() ? last_dims.dptr : static_cast(workspace->unused_dims), - DType::kInt64, {num_tensors}); - DLTensorWrapper mOffsets(offsets_bt, /*flatten_2D=*/false), - mFirstDims(first_dims_bt, /*flatten_2D=*/false), - mLastDims(last_dims_bt, /*flatten_2D=*/false); + DLTensorWrapper mS_row(make_basic_tensor(scale_row.dptr, scale_row.dtype, {scale_row.numel()}), + false, device_index); + DLTensorWrapper mS_col(make_basic_tensor(scale_col.dptr, scale_col.dtype, {scale_col.numel()}), + false, device_index); + + // Offsets and member dims are read from the output, as in mxfp8::group_quantize. + auto dims_or_unused = [&](const SimpleTensor &t, size_t numel) { + void *dptr = t.has_data() ? t.dptr : static_cast(workspace->unused_dims); + return DLTensorWrapper(make_basic_tensor(dptr, DType::kInt64, {numel}), false, device_index); + }; + DLTensorWrapper mOffsets = dims_or_unused(output_tensor->tensor_offsets, num_tensors + 1); + DLTensorWrapper mFirstDims = dims_or_unused(output_tensor->first_dims, num_tensors); + DLTensorWrapper mLastDims = dims_or_unused(output_tensor->last_dims, num_tensors); // The kernel reads num_tensors off this tensor's leading extent, so it must be exactly // the group size even on the single-tensor path that leaves the descriptors untouched. - const NVTEBasicTensor tensormaps_bt = + DLTensorWrapper mTensormaps( make_basic_tensor(static_cast(workspace->tensor_maps), DType::kInt64, - {num_tensors, kGroupTensorMapSlots, kInt64PerTensorMap}); - DLTensorWrapper mTensormaps(tensormaps_bt, /*flatten_2D=*/false); + {num_tensors, kGroupTensorMapSlots, kInt64PerTensorMap}), + false, device_index); // stream is a tvm-ffi opaque "handle"; pass it as void*. (*group_quant_func_opt)(&mX, &mO_row, &mO_col, &mS_row, &mS_col, &mOffsets, &mFirstDims, @@ -262,17 +281,20 @@ bool mxfp8_group_quantize_cutedsl(const GroupedTensor *input_tensor, const Tenso shape_rep = ShapeRepresentation::VARYING_LAST_DIM; } else if (output_tensor->all_same_last_dim()) { shape_rep = ShapeRepresentation::VARYING_FIRST_DIM; - } else { - // VARYING_BOTH_DIMS: the logical shape is [1, total], which is not tileable. + } else if (output_tensor->varying_both_dims()) { + shape_rep = ShapeRepresentation::VARYING_BOTH_DIMS; + } + if (shape_rep == ShapeRepresentation::VARYING_BOTH_DIMS) { + // The logical shape is [1, total], which is not tileable. maybe_warn_cutedsl_not_chosen("groups with both dimensions varying are not supported."); return false; } + // Leave invalid group sizes to mxfp8::group_quantize, which raises a proper error. // Every member gets a descriptor slot in the fixed-size workspace, so the CUDA // kernel's descriptor limit applies to the single-tensor representations here too. - if (input_tensor->num_tensors > dispatch::common::MAX_SUPPORTED_TENSOR_DESCRIPTORS) { - maybe_warn_cutedsl_not_chosen("the group has more than ", - dispatch::common::MAX_SUPPORTED_TENSOR_DESCRIPTORS, - " tensors."); + if (input_tensor->num_tensors == 0 || input_tensor->num_tensors > kMaxGroupTensors) { + maybe_warn_cutedsl_not_chosen("the group size ", input_tensor->num_tensors, + " is not between 1 and ", kMaxGroupTensors, "."); return false; } From 0ff3855d09367a5ee4f6132586af82a94cc2b6b7 Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Tue, 29 Sep 2026 21:41:49 +0000 Subject: [PATCH 03/58] fix Signed-off-by: Kaining Zhong --- .../cast/mxfp8/group_quantize_mxfp8.py | 24 ++++--------------- .../CuTeDSL/cast/mxfp8/quantize_mxfp8.py | 16 +++++++++++++ .../mxfp8/group_quantize_mxfp8_cutedsl.cuh | 21 +++++++++++++--- 3 files changed, 39 insertions(+), 22 deletions(-) diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8.py index 116b07200f6..c0d03d37bc7 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8.py @@ -30,9 +30,8 @@ per-tensor grid above is what replaced it on both sides. Mechanics that provably yield the same bytes may differ: the mbarrier pipeline is -expressed with PipelineTmaAsync instead of hand-rolled mbarriers, and -out-of-bounds scale padding is skipped rather than explicitly zeroed (the CUDA -kernel writes 0 there; downstream only consumes the meaningful region). +expressed with PipelineTmaAsync instead of hand-rolled mbarriers. As in CUDA, the +scales of out-of-bounds columns in a chunk (the scale-row padding) are written as 0. Scope: cast-only (no dbias / activation / dact / amax), compact (non-swizzled) scales, rowwise and/or colwise. VARYING_BOTH_DIMS is not handled here (its @@ -44,9 +43,7 @@ and for VARYING_FIRST_DIM the per-tensor extents live in device memory, so the C++ bridge cannot check them either without a sync. A violating group mis-tiles silently here where CUDA raises. -""" -""" Measured, deliberately NOT changed: - sO_row and sO_col are both allocated unconditionally, where CUDA sizes only the direction in use. Sizing them conditionally does work -- ncu confirms the shared-memory @@ -71,7 +68,7 @@ from cutlass import pipeline from cutlass import Boolean, Int32, Int64, Float8E8M0FNU from cutlass.cute.nvgpu import cpasync -from cutlass.tensor_utils import TensorMapManager, TensorMapUpdateMode +from cutlass.utils import TensorMapManager, TensorMapUpdateMode from cuda.bindings.driver import CUstream # pylint: disable=no-name-in-module import tvm_ffi @@ -222,14 +219,11 @@ def __call__( (self.BUFF_DIM_Y, self.BUFF_DIM_X), order=(1, 0) ) cta_tiler = (self.BUFF_DIM_Y, self.BUFF_DIM_X) - print(f"mx={mX}, smem_tile_layout={smem_tile_layout}, cta_tiler={cta_tiler}\n") op_load = cpasync.CopyBulkTensorTileG2SOp() tma_atom_x, tma_src = cpasync.make_tiled_tma_atom( op_load, mX, smem_tile_layout, cta_tiler, num_multicast=1 ) - print(f"tma_atom_x={tma_atom_x}\n") - print(f"tma_src={tma_src}\n") op_store = cpasync.CopyBulkTensorTileS2GOp() tma_atom_out_row, tma_dst_out_row = cpasync.make_tiled_tma_atom( op_store, mO_row, smem_tile_layout, cta_tiler, num_multicast=1 @@ -489,7 +483,6 @@ class SharedStorage: stride=((self.BUFF_DIM_X, 1), self.BUFF_DIM_Y * self.BUFF_DIM_X), ) sX = storage.sX.get_tensor(tile_layout) - print(f"sX={sX}\n") sO_row = storage.sO_row.get_tensor(tile_layout) sO_col = storage.sO_col.get_tensor(tile_layout) @@ -520,11 +513,6 @@ class SharedStorage: tXsO_col, tXgO_col = cpasync.tma_partition( tma_atom_out_col, 0, cute.make_layout(1), sO_col, gO_col_tiled ) - print(f"tma_atom_x={tma_atom_x}\n") - print(f"tma_src={tma_src}\n") - print(f"gX_tiled={gX_tiled}\n") - print(f"tXsX={tXsX}\n") - print(f"tXgX={tXgX}\n") tmap = TensorMapManager(TensorMapUpdateMode.GMEM, BYTES_PER_TENSORMAP) cute.arch.sync_threads() @@ -536,8 +524,6 @@ class SharedStorage: tensor_cols = Int32(0) # Element offset of this tensor within the group: Int64 (CUDA uses size_t), since a # group can exceed 2^31 elements even when every individual extent is small. - # Element offset of this tensor within the group: Int64 (CUDA uses size_t), since a - # group can exceed 2^31 elements even when every individual extent is small. tensor_base = Int64(0) # Block's offset and id in this individual tensor / global single tensor block_offset_Y = Int32(0) @@ -867,7 +853,7 @@ def _process_block( SWIZZLE=False, TILE_X=self.BUFF_DIM_X, TILE_Y=self.BUFF_DIM_Y, - SKIP_MASKING=False, + ZERO_OOB_SCALES=True, ) if cutlass.const_expr(cfg.ROWWISE): quantize_rowwise_mxfp8( @@ -888,7 +874,7 @@ def _process_block( WAVES=self.WAVES, THREADS_PER_BANK=self.THREADS_PER_BANK, PACK_SIZE=self.PACK_SIZE, - SKIP_MASKING=False, + ZERO_OOB_SCALES=True, ) # Force consumer's write to SMEM to be visible to TMA stores later diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8.py index 824ed18aa59..a59b9745a24 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8.py @@ -175,6 +175,9 @@ def quantize_rowwise_mxfp8( 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 + # Write 0 to the scales of col-blocks past N instead of skipping them, so the padding of the + # scale row is zeroed (mirrors group_quantize_mxfp8.cuh). + ZERO_OOB_SCALES: cutlass.Constexpr[bool] = False, ): """Quantize one SMEM tile rowwise to MXFP8 (per-row 32-elt block scales); returns the tile amax.""" tidx, _, _ = cute.arch.thread_idx() @@ -375,6 +378,11 @@ def quantize_rowwise_mxfp8( 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 + if cutlass.const_expr(ZERO_OOB_SCALES): + if scale_row < M and scale_col_first_elt >= N: + mS_row_stage[(tidx // CTA_THREADS_X, tidx % CTA_THREADS_X)] = Uint8(0).bitcast( + Float8E8M0FNU + ) inv_scale_r = exp2f_rcp(biased_exp_r) # f32 reciprocal of the scale scale_2x = pack_f32x2(inv_scale_r, inv_scale_r) @@ -421,6 +429,9 @@ def quantize_colwise_mxfp8( 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 + # Write 0 to the scales of columns past N instead of skipping them, so the padding of the + # scale row is zeroed (mirrors group_quantize_mxfp8.cuh). + ZERO_OOB_SCALES: cutlass.Constexpr[bool] = False, ): """Quantize one SMEM tile colwise to MXFP8 (per-column 32-elt block scales); returns (amax, dbias_partial).""" tidx, _, _ = cute.arch.thread_idx() @@ -541,6 +552,11 @@ def quantize_colwise_mxfp8( mS_col_stage[(0, tidx % 32, tidx // 32)] = biased_exp_c else: mS_col_stage[(0, tidx)] = biased_exp_c + if cutlass.const_expr(ZERO_OOB_SCALES): + # Swizzled layouts pad through derive_swizzled_scale_layout instead. + assert not SWIZZLE + if tile_row_start < M and scale_col >= N: + mS_col_stage[(0, tidx)] = Uint8(0).bitcast(Float8E8M0FNU) inv_scale_c = exp2f_rcp(biased_exp_c) # cvt.rn.satfinite can be vectorized to convert 2 f32 to 2 fp8 in one instruction diff --git a/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8_cutedsl.cuh b/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8_cutedsl.cuh index b5b72ad20fc..41d4b23cd17 100644 --- a/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8_cutedsl.cuh +++ b/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8_cutedsl.cuh @@ -292,9 +292,24 @@ bool mxfp8_group_quantize_cutedsl(const GroupedTensor *input_tensor, const Tenso // Leave invalid group sizes to mxfp8::group_quantize, which raises a proper error. // Every member gets a descriptor slot in the fixed-size workspace, so the CUDA // kernel's descriptor limit applies to the single-tensor representations here too. - if (input_tensor->num_tensors == 0 || input_tensor->num_tensors > kMaxGroupTensors) { - maybe_warn_cutedsl_not_chosen("the group size ", input_tensor->num_tensors, - " is not between 1 and ", kMaxGroupTensors, "."); + const size_t num_tensors = input_tensor->num_tensors; + if (num_tensors == 0 || num_tensors > kMaxGroupTensors) { + maybe_warn_cutedsl_not_chosen("the group size ", num_tensors, " is not between 1 and ", + kMaxGroupTensors, "."); + return false; + } + if (shape_rep == ShapeRepresentation::SAME_BOTH_DIMS) { + // The kernel tiles the stacked rows without tensor boundaries, which matches the CUDA + // kernel's per-tensor tiling only when every member's rows are a multiple of its + // 128-row chunk. mxfp8::group_quantize raises for a non-integral row count. + const size_t first_logical_dim = input_tensor->logical_shape.data[0]; + if (first_logical_dim % num_tensors != 0 || (first_logical_dim / num_tensors) % 128 != 0) { + maybe_warn_cutedsl_not_chosen("the rows of each group member are not a multiple of 128."); + return false; + } + } else if (!output_tensor->tensor_offsets.has_data()) { + // The varying representations read per-member offsets, as the CUDA kernel does. + maybe_warn_cutedsl_not_chosen("the grouped tensor has no tensor offsets."); return false; } From 1dc64bd2d0c1e58528702d48d5682f13e16251bd Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Wed, 30 Sep 2026 18:47:15 +0000 Subject: [PATCH 04/58] fix Signed-off-by: Kaining Zhong --- .../mxfp8/test_mxfp8_cutedsl_backend.py | 219 ++++ .../cast/mxfp8/group_quantize_mxfp8.py | 1015 +++++++++++------ .../CuTeDSL/cast/mxfp8/quantize_mxfp8.py | 12 +- .../common/cast/dispatch/quantize.cuh | 20 +- .../mxfp8/group_quantize_mxfp8_cutedsl.cuh | 199 +++- 5 files changed, 1075 insertions(+), 390 deletions(-) diff --git a/tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py b/tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py index f409f76993d..c3ab0bfa8c2 100644 --- a/tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py +++ b/tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py @@ -370,3 +370,222 @@ 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) + + +# Grouped quantization (nvte_group_quantize / nvte_group_quantize_dbias). Every member's first +# dim is a multiple of 128, which both grouped kernels require. VARYING_LAST_DIM and +# VARYING_BOTH_DIMS members also keep their last dims 128-aligned, as both kernels' per-member +# scale layout assumes. The fused activations have no PyTorch binding; the CuTeDSL path for +# them is covered by tests/cpp/operator/test_cast_mxfp8_grouped.cu run with +# NVTE_ENABLE_CUTEDSL_BACKEND=1. +# (name, shape representation in the config key, per-member shapes) +GROUP_CASES = [ + ("same_both", "same_both_dims", [(256, 512)] * 3), + # N is 32- but not 128-divisible, so the rowwise and colwise scales carry zeroed padding. + ("same_both_n96", "same_both_dims", [(128, 96)] * 2), + ("varying_first", "varying_first_dim", [(128, 256), (384, 256), (256, 256)]), + ("varying_first_n160", "varying_first_dim", [(128, 160), (256, 160)]), + # N ends in a partial 32-element scale block. + ("varying_first_n144", "varying_first_dim", [(128, 144), (384, 144)]), + ("varying_last", "varying_last_dim", [(256, 128), (256, 384), (256, 256)]), + ("varying_both", "varying_both_dims", [(128, 256), (256, 128), (384, 512)]), +] +SINGLE_TENSOR_GROUP_CASES = [ + c for c in GROUP_CASES if c[1] in ("same_both_dims", "varying_first_dim") +] +get_group_case_id = lambda c: c[0] + + +def group_dims(members): + """(logical shape, first_dims, last_dims) of the grouped tensor made of `members`.""" + first_dims = [m for m, _ in members] + last_dims = [n for _, n in members] + same_first = len(set(first_dims)) == 1 + same_last = len(set(last_dims)) == 1 + if same_last: + logical_shape = (sum(first_dims), last_dims[0]) + elif same_first: + logical_shape = (first_dims[0], sum(last_dims)) + else: + logical_shape = (1, sum(m * n for m, n in members)) + to_tensor = lambda dims: torch.tensor(dims, dtype=torch.int64, device="cuda") + return ( + logical_shape, + None if same_first else to_tensor(first_dims), + None if same_last else to_tensor(last_dims), + ) + + +def run_group_quantize(members, in_dtype, fp8_dtype, rowwise, columnwise, swizzled, dbias): + """Quantize the concatenated members; returns (grouped output, dbias or None).""" + logical_shape, first_dims, last_dims = group_dims(members) + x, _ = generate_inputs(*logical_shape, in_dtype) + q = MXFP8Quantizer(fp8_dtype=fp8_dtype, rowwise=rowwise, columnwise=columnwise) + q.optimize_for_gemm = swizzled + if dbias: + return tex.bgrad_group_quantize(x, q, len(members), first_dims, last_dims) + return tex.group_quantize(x, q, len(members), first_dims, last_dims), None + + +def extract_group_quantized_output(out, rowwise, columnwise): + """Extract the bytes to compare between backends. + + Unlike the single-tensor kernels, both grouped backends zero the scale padding, so the + whole data and scale buffers are compared. + """ + parts = {} + if rowwise: + parts["rowwise data"] = out.rowwise_data.view(torch.uint8).clone() + parts["rowwise scales"] = out.scale_inv.view(torch.uint8).clone() + if columnwise: + parts["colwise data"] = out.columnwise_data.view(torch.uint8).clone() + parts["colwise scales"] = out.columnwise_scale_inv.view(torch.uint8).clone() + return parts + + +def get_group_cfg_key(shape_rep, in_dtype, fp8_dtype, rowwise, colwise, swizzled, dbias): + """Mirror of MXFP8GroupQuantConfig::to_key (group_quantize_mxfp8_cutedsl.cuh) for the + configs reachable from PyTorch (no fused activation).""" + major, minor = device_compute_capability() + flags = (swizzled, dbias, False, False) # swizzled, with_dbias, with_dact, with_act + return ( + f"cutedsl_group_mxfp8_sm{major * 10 + minor}_{DTYPE_TO_STR[in_dtype]}_" + f"{FP8_TO_KEY[fp8_dtype]}_{int(rowwise)}_{int(colwise)}_{shape_rep}_" + + "_".join("1" if f else "0" for f in flags) + + "_none" + ) + + +def assert_group_cutedsl_registered(*key_args): + """Guard against a silent CUDA fallback; see run_test_case.""" + key = get_group_cfg_key(*key_args) + assert tvm_ffi.get_global_func(key, allow_missing=True) is not None, ( + f"CuTeDSL kernel not registered for {key}; the CuTeDSL backend fell back " + "to CUDA and this case compared CUDA against itself" + ) + + +def run_group_test_case( + members, shape_rep, block_size, in_dtype, fp8_dtype, swizzled=False, dbias=False +): + """Assert the CuTeDSL and CUDA grouped backends produce bit-identical outputs, including + dbias, which both accumulate in the same order.""" + rowwise = block_size[1] != 1 + columnwise = block_size[0] != 1 + args = (members, in_dtype, fp8_dtype, rowwise, columnwise, swizzled, dbias) + + set_cutedsl_backend(False) + out_cuda, dbias_cuda = run_group_quantize(*args) + cuda_output = extract_group_quantized_output(out_cuda, rowwise, columnwise) + + set_cutedsl_backend(True) + try: + out_cutedsl, dbias_cutedsl = run_group_quantize(*args) + cutedsl_output = extract_group_quantized_output(out_cutedsl, rowwise, columnwise) + finally: + set_cutedsl_backend(False) + + assert_group_cutedsl_registered( + shape_rep, in_dtype, fp8_dtype, rowwise, columnwise, swizzled, dbias + ) + tag = f"group/{members}/{DTYPE_TO_STR[in_dtype]}/{FP8_TO_STR[fp8_dtype]}" + for name, cuda_bytes in cuda_output.items(): + assert torch.equal( + cutedsl_output[name], cuda_bytes + ), f"{tag}: {name} differ between backends" + if dbias: + assert torch.equal(dbias_cutedsl, dbias_cuda), f"{tag}: dbias differs between backends" + + +@pytest.mark.parametrize("case", GROUP_CASES, ids=get_group_case_id) +@pytest.mark.parametrize("block_size", BLOCK_SIZES, ids=get_block_id) +@pytest.mark.parametrize("in_dtype", IN_DTYPES, ids=get_dtype_id) +@pytest.mark.parametrize("fp8_dtype", FP8_DTYPES, ids=get_fp8_id) +def test_group_cast_only(fp8_dtype, in_dtype, block_size, case): + _, shape_rep, members = case + run_group_test_case(members, shape_rep, block_size, in_dtype, fp8_dtype) + + +@pytest.mark.parametrize("case", GROUP_CASES, ids=get_group_case_id) +@pytest.mark.parametrize("block_size", BLOCK_SIZES, ids=get_block_id) +def test_group_swizzled(block_size, case): + _, shape_rep, members = case + if block_size[0] != 1 and shape_rep in ("varying_last_dim", "varying_both_dims"): + pytest.skip( + "The CUDA kernel's GEMM-swizzled colwise scale index double-counts the tensor base" + " for varying last dims; the CuTeDSL backend leaves these configs to it." + ) + run_group_test_case( + members, shape_rep, block_size, torch.bfloat16, tex.DType.kFloat8E4M3, swizzled=True + ) + + +@pytest.mark.parametrize("case", SINGLE_TENSOR_GROUP_CASES, ids=get_group_case_id) +@pytest.mark.parametrize("block_size", BLOCK_SIZES, ids=get_block_id) +@pytest.mark.parametrize("in_dtype", [torch.bfloat16, torch.float32], ids=get_dtype_id) +@pytest.mark.parametrize("swizzled", SWIZZLE_MODES, ids=get_swizzle_id) +def test_group_dbias(swizzled, in_dtype, block_size, case): + _, shape_rep, members = case + run_group_test_case( + members, + shape_rep, + block_size, + in_dtype, + tex.DType.kFloat8E4M3, + swizzled=swizzled, + dbias=True, + ) + + +@pytest.mark.parametrize("cutedsl", [False, True], ids=["cuda", "cutedsl"]) +def test_group_noop(cutedsl): + """A set cast-noop flag leaves a reused grouped output untouched; a clear one does not.""" + members = [(256, 512)] * 3 + logical_shape, _, _ = group_dims(members) + x1, x2 = generate_inputs(*logical_shape, torch.bfloat16) + q = MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3, rowwise=True, columnwise=True) + set_cutedsl_backend(cutedsl) + try: + out = tex.group_quantize(x1, q, len(members), None) + first = extract_group_quantized_output(out, True, True) + noop = torch.ones(1, dtype=torch.float32, device="cuda") + tex.group_quantize(x2, q, len(members), None, noop_flag=noop, output=out) + skipped = extract_group_quantized_output(out, True, True) + noop.zero_() + tex.group_quantize(x2, q, len(members), None, noop_flag=noop, output=out) + quantized = extract_group_quantized_output(out, True, True) + expected = extract_group_quantized_output( + tex.group_quantize(x2, q, len(members), None), True, True + ) + finally: + set_cutedsl_backend(False) + if cutedsl: + assert_group_cutedsl_registered( + "same_both_dims", torch.bfloat16, tex.DType.kFloat8E4M3, True, True, False, False + ) + for name, first_bytes in first.items(): + assert torch.equal(skipped[name], first_bytes), f"{name} changed under a set noop flag" + assert torch.equal(quantized[name], expected[name]), f"{name} wrong under a clear flag" + + +def test_group_2d_quantization_fallback(): + """2D block scaling is left to the CUDA kernel, which must still be what runs.""" + members = [(128, 256), (384, 256), (256, 256)] + logical_shape, first_dims, _ = group_dims(members) + x, _ = generate_inputs(*logical_shape, torch.bfloat16) + outputs = [] + for cutedsl in (False, True): + q = MXFP8Quantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=True, + columnwise=True, + with_2d_quantization=True, + ) + set_cutedsl_backend(cutedsl) + try: + out = tex.group_quantize(x, q, len(members), first_dims) + outputs.append(extract_group_quantized_output(out, True, True)) + finally: + set_cutedsl_backend(False) + for name, cuda_bytes in outputs[0].items(): + assert torch.equal(outputs[1][name], cuda_bytes), f"{name} differ between backends" diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8.py index c0d03d37bc7..26e36db8100 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8.py @@ -33,16 +33,23 @@ expressed with PipelineTmaAsync instead of hand-rolled mbarriers. As in CUDA, the scales of out-of-bounds columns in a chunk (the scale-row padding) are written as 0. -Scope: cast-only (no dbias / activation / dact / amax), compact (non-swizzled) -scales, rowwise and/or colwise. VARYING_BOTH_DIMS is not handled here (its -logical shape [1, total] is not tileable); the C++ bridge falls back to CUDA. - -Known gap vs CUDA: the CUDA kernel validates on device that every member's first -dim is a multiple of 128 (get_tensor_rows_num -> NVTE_DEVICE_ERROR). This kernel -shares the precondition but cannot check it -- CuTeDSL has no device-side assert, -and for VARYING_FIRST_DIM the per-tensor extents live in device memory, so the -C++ bridge cannot check them either without a sync. A violating group mis-tiles -silently here where CUDA raises. +Scope: everything group_quantize_mxfp8.cuh covers except 2D block scaling -- the +cast-noop flag, fused activation (IS_ACT) and activation derivative (IS_DACT), dbias, +compact and GEMM-swizzled scales, rowwise and/or colwise, and all four shape +representations. Differences from CUDA: + * the grouped amax pointer is accepted and left untouched, as the CUDA kernel does; + * GEMM-swizzled colwise scales are only produced for the single-tensor reps. For + VARYING_LAST_DIM / VARYING_BOTH_DIMS the CUDA kernel adds the tensor base to the + colwise swizzled index twice, which is only in bounds for the first member, so the + C++ bridge falls back to CUDA instead of reproducing it; + * dbias partial sums are accumulated in the CUDA kernel's order (a running column sum + with colwise output, otherwise per-thread sums reduced across the CTA), and the C++ + bridge reduces the workspace with the same grouped_reduce_dbias. + +Like the CUDA kernel, every member's first dim must be a multiple of 128 (and, for the +varying-last reps, its last dim too). The kernel prints the same diagnostics as +get_tensor_rows_num / get_tensor_cols_num when a group violates this and, like +NVTE_DEVICE_ERROR in a release build, carries on. Measured, deliberately NOT changed: - sO_row and sO_col are both allocated unconditionally, where CUDA sizes only the @@ -61,12 +68,12 @@ import logging import os -from typing import Type +from typing import Optional, Type import cutlass from cutlass import cute from cutlass import pipeline -from cutlass import Boolean, Int32, Int64, Float8E8M0FNU +from cutlass import Boolean, Float32, Int32, Int64, Float8E8M0FNU from cutlass.cute.nvgpu import cpasync from cutlass.utils import TensorMapManager, TensorMapUpdateMode from cuda.bindings.driver import CUstream # pylint: disable=no-name-in-module @@ -78,8 +85,13 @@ ) from transformer_engine.common.CuTeDSL.cast.mxfp8.quantize_mxfp8 import ( MXFP8_BLOCK_SCALING_SIZE, + SYM_N_DIVISIBILITY, FP8E4M3_MAX_NORM_RCP, FP8E5M2_MAX_NORM_RCP, + SUPPORTED_ACTIVATIONS, + SUPPORTED_DACTIVATIONS, + derive_swizzled_scale_layout, + noop_flag_is_set, quantize_rowwise_mxfp8, quantize_colwise_mxfp8, ) @@ -89,8 +101,9 @@ THREADS_PER_WARP = 32 BYTES_PER_TENSORMAP = 128 -# Descriptor slots per tensor: input, rowwise output, colwise output. -NUM_TENSORMAPS = 3 +# Descriptor slots per tensor: input, rowwise output, colwise output, activation input. +NUM_TENSORMAPS = 4 +ACT_INPUT_SLOT = 3 # One extra slot holds per-tensor (rows, cols, base_elts), so the main kernel never has to # binary-search the offsets array. Mirrors TensorMapStorage::rows/cols/offsets upstream. META_SLOT = NUM_TENSORMAPS @@ -100,13 +113,30 @@ SAME_BOTH_DIMS = "same_both_dims" VARYING_FIRST_DIM = "varying_first_dim" VARYING_LAST_DIM = "varying_last_dim" -SUPPORTED_SHAPE_REPS = (SAME_BOTH_DIMS, VARYING_FIRST_DIM, VARYING_LAST_DIM) +VARYING_BOTH_DIMS = "varying_both_dims" +SUPPORTED_SHAPE_REPS = (SAME_BOTH_DIMS, VARYING_FIRST_DIM, VARYING_LAST_DIM, VARYING_BOTH_DIMS) + +# Upper bound on the group size (MAX_SUPPORTED_TENSOR_DESCRIPTORS in grouped_tma.cuh); sizes +# the fixed binary search over the offsets. +MAX_SUPPORTED_TENSORS = 64 class MXFP8GroupQuantizeConfig: """Compile-time config for the grouped MXFP8 quantize kernel.""" - def __init__(self, dtype: str, fp8_dtype: str, rowwise: bool, colwise: bool, shape_rep: str): + def __init__( + self, + dtype: str, + fp8_dtype: str, + rowwise: bool, + colwise: bool, + shape_rep: str, + with_gemm_swizzled_scales: bool = False, + with_dbias: bool = False, + with_dact: bool = False, + with_act: bool = False, + activation: str = "none", + ): if 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) @@ -133,10 +163,48 @@ def __init__(self, dtype: str, fp8_dtype: str, rowwise: bool, colwise: bool, sha FP8E4M3_MAX_NORM_RCP if fp8_dtype == "Float8E4M3" else FP8E5M2_MAX_NORM_RCP ) + self.WITH_GEMM_SWIZZLED_SCALES = with_gemm_swizzled_scales + if with_gemm_swizzled_scales and colwise and not self.IS_SINGLE_TENSOR: + # The CUDA kernel offsets the colwise swizzled scales of these representations by + # the tensor base twice, which only lands inside the buffer for the first member. + raise ValueError( + "GEMM-swizzled colwise scales are only supported for single-tensor representations" + ) + if with_dbias and not self.IS_SINGLE_TENSOR: + # mxfp8::group_quantize raises for this. + raise ValueError("dbias is only supported for tensors with a common last dimension") + self.WITH_DBIAS = with_dbias + + if with_dact and with_act: + raise ValueError("with_dact and with_act cannot both be set") + if with_dact: + if activation not in SUPPORTED_DACTIVATIONS: + raise ValueError( + f"unknown activation {activation!r} for with_dact=True; expected one of" + f" {sorted(SUPPORTED_DACTIVATIONS)}" + ) + self.ACTIVATION = activation + elif with_act: + if activation not in SUPPORTED_ACTIVATIONS: + raise ValueError( + f"unknown activation {activation!r} for with_act=True; expected one of" + f" {sorted(SUPPORTED_ACTIVATIONS)}" + ) + self.ACTIVATION = activation + else: + if activation != "none": + raise ValueError("activation must be none when with_dact and with_act are False") + self.ACTIVATION = None + self.WITH_DACT = with_dact + self.WITH_ACT = with_act + def __str__(self): return ( f"MXFP8GroupQuantizeConfig(dtype={self.DTYPE_STR}, fp8_dtype={self.FP8_DTYPE_STR}, " - f"rowwise={self.ROWWISE}, colwise={self.COLWISE}, shape_rep={self.SHAPE_REP})" + f"rowwise={self.ROWWISE}, colwise={self.COLWISE}, shape_rep={self.SHAPE_REP}, " + f"swizzled={self.WITH_GEMM_SWIZZLED_SCALES}, with_dbias={self.WITH_DBIAS}, " + f"with_dact={self.WITH_DACT}, with_act={self.WITH_ACT}, " + f"activation={self.ACTIVATION})" ) __repr__ = __str__ @@ -168,6 +236,24 @@ class MXFP8GroupQuantizeKernel: def __init__(self, cfg: MXFP8GroupQuantizeConfig, SM_COUNT: int): self.cfg = cfg self.SM_COUNT = SM_COUNT + # CastConfig widens the chunk to 128x256, i.e. STAGES_X = 2 tiles of + # BUFF_DIM_X columns, each traversed in STAGES row stages before moving right. + self.STAGES_X = 2 if cfg.SHAPE_REP == VARYING_BOTH_DIMS else 1 + self.CHUNK_WIDTH = self.CHUNK_DIM_X * self.STAGES_X + # The CUDA kernel honors the noop flag only without fused activations or dbias. + self.CHECK_NOOP_FLAG = not (cfg.WITH_ACT or cfg.WITH_DACT or cfg.WITH_DBIAS) + # Like IS_CACHED_ACT_OP in CUDA: with both directions, the colwise pass caches the + # activation in the input tile for the rowwise pass. ReLU is fused into the conversion + # instead, which yields the same bytes. + self.CACHE_ACTIVATION = ( + (cfg.WITH_ACT or cfg.WITH_DACT) + and cfg.ROWWISE + and cfg.COLWISE + and cfg.ACTIVATION != "relu" + ) + # CUDA reduces dbias in the colwise pass when there is one, else in the rowwise pass. + self.DBIAS_IN_COLWISE = cfg.WITH_DBIAS and cfg.COLWISE + self.DBIAS_IN_ROWWISE = cfg.WITH_DBIAS and not cfg.COLWISE # ---------------------------------------------------------------- helpers @cute.jit @@ -176,16 +262,80 @@ def _tensor_rows_cols( ): """Get the shape (rows, cols) of the tensor by tensor_id.""" cfg = self.cfg - if cutlass.const_expr(cfg.SHAPE_REP == VARYING_FIRST_DIM): + if cutlass.const_expr(cfg.SHAPE_REP in (VARYING_FIRST_DIM, VARYING_BOTH_DIMS)): rows = Int32(mFirstDims[tensor_id]) else: rows = Int32(first_logical_dim) - if cutlass.const_expr(cfg.SHAPE_REP == VARYING_LAST_DIM): + if cutlass.const_expr(cfg.SHAPE_REP in (VARYING_LAST_DIM, VARYING_BOTH_DIMS)): cols = Int32(mLastDims[tensor_id]) else: cols = Int32(last_logical_dim) return rows, cols + @cute.jit + def _find_tensor_from_offsets(self, mOffsets, num_tensors, offset: Int64): + """Index of the tensor whose element range holds `offset` (find_tensor_from_offsets).""" + low = Int32(1) + hi = Int32(num_tensors) + # Enough bisection steps for any group of up to MAX_SUPPORTED_TENSORS members. + for _ in cutlass.range_constexpr(MAX_SUPPORTED_TENSORS.bit_length()): + if low < hi: + mid = low + (hi - low) // 2 + if Int64(mOffsets[mid]) <= offset: + low = mid + 1 + else: + hi = mid + return low - 1 + + @cute.jit + def _scale_tensor(self, mS, base: Int64, layout): + """View the scale buffer from element `base` on with `layout`.""" + return cute.make_tensor( + cute.make_ptr( + Float8E8M0FNU, + mS.iterator.toint() + base, + cute.AddressSpace.gmem, + assumed_align=4, + ), + layout, + ) + + @cute.jit + def _rowwise_scales(self, mS_row, base: Int64, rows, cols): + """Rowwise scales of a (rows, cols) tensor at `base`, tiled per 32x128 stage.""" + if cutlass.const_expr(self.cfg.WITH_GEMM_SWIZZLED_SCALES): + mS_t, _ = derive_swizzled_scale_layout( + rows, cols, True, False, self._scale_tensor(mS_row, base, cute.make_layout(1)), None + ) + else: + # Rowwise scale's divisibility guarantee: (128, 4) + stride = cute.round_up(cute.ceil_div(cols, MXFP8_BLOCK_SCALING_SIZE), 4) + mS_t = self._scale_tensor( + mS_row, base, cute.make_layout((rows, stride), stride=(stride, 1)) + ) + return cute.zipped_divide( + mS_t, (self.BUFF_DIM_Y, self.BUFF_DIM_X // MXFP8_BLOCK_SCALING_SIZE) + ) + + @cute.jit + def _colwise_scales(self, mS_col, base: Int64, rows, cols): + """Colwise scales of a (rows, cols) tensor at `base`, tiled per 32x128 stage.""" + if cutlass.const_expr(self.cfg.WITH_GEMM_SWIZZLED_SCALES): + _, mS_t = derive_swizzled_scale_layout( + rows, cols, False, True, None, self._scale_tensor(mS_col, base, cute.make_layout(1)) + ) + else: + # Colwise scale's divisibility guarantee: (4, 128) + stride = cute.round_up(cols, 128) + mS_t = self._scale_tensor( + mS_col, + base, + cute.make_layout((rows // MXFP8_BLOCK_SCALING_SIZE, stride), stride=(stride, 1)), + ) + return cute.zipped_divide( + mS_t, (self.BUFF_DIM_Y // MXFP8_BLOCK_SCALING_SIZE, self.BUFF_DIM_X) + ) + # ------------------------------------------------------------ entry point @cute.jit def __call__( @@ -196,9 +346,12 @@ def __call__( mS_row: cute.Tensor, mS_col: cute.Tensor, mOffsets: cute.Tensor, # int64[num_tensors + 1], CSR element offsets - mFirstDims: cute.Tensor, # int64[num_tensors] (VARYING_FIRST_DIM) - mLastDims: cute.Tensor, # int64[num_tensors] (VARYING_LAST_DIM) - mTensormaps: cute.Tensor, # int64[num_tensors, NUM_TENSORMAPS, 16] + mFirstDims: cute.Tensor, # int64[num_tensors] (VARYING_FIRST_DIM / VARYING_BOTH_DIMS) + mLastDims: cute.Tensor, # int64[num_tensors] (VARYING_LAST_DIM / VARYING_BOTH_DIMS) + mTensormaps: cute.Tensor, # int64[num_tensors, NUM_WORKSPACE_SLOTS, 16] + mNoop: cute.Pointer, # f32 cast_noop flag; may be null, checked on device + mActInput: Optional[cute.Tensor], # activation input, only with WITH_DACT + mWorkspace: Optional[cute.Tensor], # f32 partial dbias, only with WITH_DBIAS stream: CUstream, ): if cutlass.const_expr(CUTEDSL_DEBUG_LOGGING): @@ -211,8 +364,7 @@ def __call__( # pass a length-num_tensors stub for SAME_BOTH_DIMS (where the offsets array is # unused), which would make `mOffsets.shape[0] - 1` read one too few and divide by # zero at num_tensors == 1. The per-tensor descriptor workspace is num_tensors long - # by construction -- one slot set per member -- so it is the reliable source, and it - # is only consulted on the multi-tensor path that actually needs the workspace. + # by construction -- one slot set per member -- so it is the reliable source. num_tensors = mTensormaps.shape[0] smem_tile_layout = cute.make_ordered_layout( @@ -224,6 +376,12 @@ def __call__( tma_atom_x, tma_src = cpasync.make_tiled_tma_atom( op_load, mX, smem_tile_layout, cta_tiler, num_multicast=1 ) + tma_atom_act = None + tma_src_act = None + if cutlass.const_expr(cfg.WITH_DACT): + tma_atom_act, tma_src_act = cpasync.make_tiled_tma_atom( + op_load, mActInput, smem_tile_layout, cta_tiler, num_multicast=1 + ) op_store = cpasync.CopyBulkTensorTileS2GOp() tma_atom_out_row, tma_dst_out_row = cpasync.make_tiled_tma_atom( op_store, mO_row, smem_tile_layout, cta_tiler, num_multicast=1 @@ -236,33 +394,30 @@ def __call__( # How many blocks does the grouped tensor have in both directions work_blocks_X = cute.ceil_div(Int32(last_logical_dim), self.CHUNK_DIM_X) work_blocks_Y = cute.ceil_div(Int32(first_logical_dim), self.CHUNK_DIM_Y) - # Each CTA handles a block from one individual tensor + # Each CTA handles one chunk. With every member's rows a multiple of CHUNK_DIM_Y, + # this linear order is the CUDA kernel's (X, Y-in-tensor, tensor) order. grid = [work_blocks_X * work_blocks_Y, 1, 1] else: - # The work-block grid is per-tensor here, not global: each CTA derives its own - # block range from its tensor's extents (see the kernel's `else` branch), so - # work_blocks_X/Y are dead on this path -- the kernel reads them only under - # IS_SINGLE_TENSOR. Pass dummies rather than computing first*last, which is an - # ELEMENT count and would wrap Int32 past 2^31 elements. + # The work-block grid is per-tensor here: each CTA derives its own block range + # from its tensor's extents, so work_blocks_X/Y are unused on this path. work_blocks_Y = Int32(1) work_blocks_X = Int32(1) # Persistent worker count, mirroring get_launch_config() in - # group_quantize_mxfp8.cuh. There are SM_COUNT * STATIC_PERSISTENT_BLOCKS_PER_SM - # workers in total, split evenly across tensors -- but clamped to the average - # number of chunks a tensor actually holds, so a group of small tensors does not - # launch CTAs that only reach the `first_block_id >= blocks_in_tensor` early-out. + # group_quantize_mxfp8.cuh: SM_COUNT * STATIC_PERSISTENT_BLOCKS_PER_SM workers + # split evenly across tensors, clamped to the average number of chunks a tensor + # holds. The element count would wrap Int32, so the estimate + # DIVUP(elts_total, CHUNK_DIM_Y * TILE_DIM_X) is formed without it. n_tensors = cutlass.max(Int32(num_tensors), Int32(1)) # never divide by zero - # CUDA's DIVUP(elts_total, CHUNK_DIM_Y * TILE_DIM_X) / STAGES_X, where TILE_DIM_X - # is BUFF_DIM_X here and STAGES_X is 1 for every rep this kernel supports. The - # element product would wrap Int32 past 2^31 elements, so divide the first extent - # by CHUNK_DIM_Y up front -- it is exact, the kernel is compiled with - # sym_int32(divisibility=CHUNK_DIM_Y) on that extent -- and keep the whole - # estimate in Int32. Feeding an Int64 into the grid poisons the tile arithmetic - # downstream ('cute.make_tile' expects width=32). - estimated_work_blocks = cute.ceil_div( - (Int32(first_logical_dim) // self.CHUNK_DIM_Y) * Int32(last_logical_dim), - self.ELTS_PER_CHUNK // self.CHUNK_DIM_Y, - ) + if cutlass.const_expr(cfg.SHAPE_REP == VARYING_BOTH_DIMS): + # The logical shape is [1, total]. + estimated_work_blocks = cute.ceil_div(Int32(last_logical_dim), self.ELTS_PER_CHUNK) + else: + # Exact: the first extent is compiled with divisibility CHUNK_DIM_Y. + estimated_work_blocks = cute.ceil_div( + (Int32(first_logical_dim) // self.CHUNK_DIM_Y) * Int32(last_logical_dim), + self.ELTS_PER_CHUNK // self.CHUNK_DIM_Y, + ) + estimated_work_blocks = cute.ceil_div(estimated_work_blocks, self.STAGES_X) requested_workers_per_tensor = cutlass.max( Int32(1), Int32(self.SM_COUNT * self.STATIC_PERSISTENT_BLOCKS_PER_SM) // n_tensors, @@ -273,17 +428,15 @@ def __call__( workers_per_tensor = cutlass.min( requested_workers_per_tensor, average_work_blocks_per_tensor ) - # Each group of workers_per_tensor CTAs serves one tensor, and we launch - # num_tensors such groups to cover the whole group. grid = [workers_per_tensor, Int32(num_tensors), 1] - # Only manually create descriptors for the non-single-tensor case because we will need to manually - # overwrite the descriptors as we visit different groups + # Only the multi-tensor representations need per-tensor descriptors. if cutlass.const_expr(not cfg.IS_SINGLE_TENSOR): self.update_descriptors_kernel( mX, mO_row, mO_col, + mActInput, mOffsets, mFirstDims, mLastDims, @@ -294,6 +447,7 @@ def __call__( tma_atom_x, tma_atom_out_row, tma_atom_out_col, + tma_atom_act, ).launch(grid=[num_tensors, 1, 1], block=[THREADS_PER_WARP, 1, 1], stream=stream) self.kernel( @@ -301,16 +455,18 @@ def __call__( mS_col, mOffsets, mFirstDims, - mLastDims, mTensormaps, + mNoop, + mWorkspace, first_logical_dim, last_logical_dim, num_tensors, work_blocks_X, - work_blocks_Y, mX.element_type, tma_atom_x, tma_src, + tma_atom_act, + tma_src_act, tma_atom_out_row, tma_dst_out_row, tma_atom_out_col, @@ -328,6 +484,7 @@ def update_descriptors_kernel( mX, mO_row, mO_col, + mActInput, mOffsets, mFirstDims, mLastDims, @@ -338,6 +495,7 @@ def update_descriptors_kernel( tma_atom_x, tma_atom_orow, tma_atom_ocol, + tma_atom_act, ): """One CTA per tensor: point that tensor's TMA descriptors at its own block. @@ -345,11 +503,28 @@ def update_descriptors_kernel( """ cfg = self.cfg tensor_id, _, _ = cute.arch.block_idx() + tidx, _, _ = cute.arch.thread_idx() rows, cols = self._tensor_rows_cols( tensor_id, mFirstDims, mLastDims, first_logical_dim, last_logical_dim ) base_elts = Int64(mOffsets[tensor_id]) + # Same diagnostics as get_tensor_rows_num / get_tensor_cols_num. Like NVTE_DEVICE_ERROR + # in a release build, they only print. + if tidx == 0: + if rows % 128 != 0: + cute.printf( + "tensor %d: First dimension of each tensor in a group must be divisible" + " by 128.\n", + tensor_id, + ) + if cols % 128 != 0: + cute.printf( + "tensor %d: For varying last dimensions support, the last dimension of each" + " tensor in a group must be divisible by 128.\n", + tensor_id, + ) + # Publish this tensor's geometry for the main kernel (written even when empty). meta = mTensormaps[(tensor_id, META_SLOT, None)] meta[0] = Int64(rows) @@ -360,69 +535,51 @@ def update_descriptors_kernel( desc_x = tmap.get_tensormap_ptr(mTensormaps[(tensor_id, 0, None)].iterator) desc_orow = tmap.get_tensormap_ptr(mTensormaps[(tensor_id, 1, None)].iterator) desc_ocol = tmap.get_tensormap_ptr(mTensormaps[(tensor_id, 2, None)].iterator) + desc_act = tmap.get_tensormap_ptr(mTensormaps[(tensor_id, ACT_INPUT_SLOT, None)].iterator) # Zero-sized groups: creating a descriptor with a zero extent is invalid, - # so skip (the main kernel skips these jobs via job_has_work). + # so skip (the main kernel skips these tensors as well). if rows > 0 and cols > 0: - gX = cute.make_tensor( - cute.make_ptr( - dtype, - mX.iterator.toint() + base_elts * (dtype.width // 8), - cute.AddressSpace.gmem, - assumed_align=16, - ), - cute.make_layout((rows, cols), stride=(cols, 1)), - ) - gO_row = cute.make_tensor( - cute.make_ptr( - cfg.FP8_DTYPE, - mO_row.iterator.toint() + base_elts, - cute.AddressSpace.gmem, - assumed_align=16, - ), - cute.make_layout((rows, cols), stride=(cols, 1)), - ) - gO_col = cute.make_tensor( - cute.make_ptr( - cfg.FP8_DTYPE, - mO_col.iterator.toint() + base_elts, - cute.AddressSpace.gmem, - assumed_align=16, - ), - cute.make_layout((rows, cols), stride=(cols, 1)), - ) + member_layout = cute.make_layout((rows, cols), stride=(cols, 1)) + + def member_view(tensor, elt_dtype): + return cute.make_tensor( + cute.make_ptr( + elt_dtype, + tensor.iterator.toint() + base_elts * (elt_dtype.width // 8), + cute.AddressSpace.gmem, + assumed_align=16, + ), + member_layout, + ) + views = [member_view(mX, dtype)] + atoms = [tma_atom_x] + descs = [desc_x] tmap.init_tensormap_from_atom(tma_atom_x, desc_x, 0) if cutlass.const_expr(cfg.ROWWISE): + views.append(member_view(mO_row, cfg.FP8_DTYPE)) + atoms.append(tma_atom_orow) + descs.append(desc_orow) tmap.init_tensormap_from_atom(tma_atom_orow, desc_orow, 0) if cutlass.const_expr(cfg.COLWISE): + views.append(member_view(mO_col, cfg.FP8_DTYPE)) + atoms.append(tma_atom_ocol) + descs.append(desc_ocol) tmap.init_tensormap_from_atom(tma_atom_ocol, desc_ocol, 0) + if cutlass.const_expr(cfg.WITH_DACT): + views.append(member_view(mActInput, dtype)) + atoms.append(tma_atom_act) + descs.append(desc_act) + tmap.init_tensormap_from_atom(tma_atom_act, desc_act, 0) tmap.fence_tensormap_initialization() - - if cutlass.const_expr(cfg.ROWWISE and cfg.COLWISE): - tmap.update_tensormap( - (gX, gO_row, gO_col), - (tma_atom_x, tma_atom_orow, tma_atom_ocol), - (desc_x, desc_orow, desc_ocol), - 0, - (), # smem staging is unused in GMEM update mode - ) - elif cutlass.const_expr(cfg.ROWWISE): - tmap.update_tensormap( - (gX, gO_row), - (tma_atom_x, tma_atom_orow), - (desc_x, desc_orow), - 0, - (), # smem staging is unused in GMEM update mode - ) - else: - tmap.update_tensormap( - (gX, gO_col), - (tma_atom_x, tma_atom_ocol), - (desc_x, desc_ocol), - 0, - (), # smem staging is unused in GMEM update mode - ) + tmap.update_tensormap( + tuple(views), + tuple(atoms), + tuple(descs), + 0, + (), # smem staging is unused in GMEM update mode + ) # ------------------------------------------------------------ main kernel @cute.kernel @@ -432,16 +589,68 @@ def kernel( mS_col, mOffsets, mFirstDims, - mLastDims, mTensormaps, + mNoop, + mWorkspace, + first_logical_dim, + last_logical_dim, + num_tensors, + work_blocks_X, + dtype: cutlass.Constexpr[Type[cutlass.Numeric]], + tma_atom_x, + tma_src, + tma_atom_act, + tma_src_act, + tma_atom_out_row, + tma_dst_out_row, + tma_atom_out_col, + tma_dst_out_col, + ): + """No-op the CTA when the noop flag is set, else run the quantize main loop.""" + skip_execution = Boolean(False) + if cutlass.const_expr(self.CHECK_NOOP_FLAG): + skip_execution = noop_flag_is_set(mNoop) + if not skip_execution: + self._kernel_main( + mS_row, + mS_col, + mOffsets, + mFirstDims, + mTensormaps, + mWorkspace, + first_logical_dim, + last_logical_dim, + num_tensors, + work_blocks_X, + dtype, + tma_atom_x, + tma_src, + tma_atom_act, + tma_src_act, + tma_atom_out_row, + tma_dst_out_row, + tma_atom_out_col, + tma_dst_out_col, + ) + + @cute.jit + def _kernel_main( + self, + mS_row, + mS_col, + mOffsets, + mFirstDims, + mTensormaps, + mWorkspace, first_logical_dim, last_logical_dim, num_tensors, work_blocks_X, - work_blocks_Y, dtype: cutlass.Constexpr[Type[cutlass.Numeric]], tma_atom_x, tma_src, + tma_atom_act, + tma_src_act, tma_atom_out_row, tma_dst_out_row, tma_atom_out_col, @@ -454,6 +663,18 @@ def kernel( gdx, _, _ = cute.arch.grid_dim() warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + if cutlass.const_expr(cfg.SHAPE_REP == VARYING_FIRST_DIM): + # The first CTA validates every member's rows, as the CUDA kernel does. Like + # NVTE_DEVICE_ERROR in a release build, this only prints. + if bidx == 0: + if tidx < num_tensors: + if Int64(mFirstDims[tidx]) % 128 != 0: + cute.printf( + "tensor %d: First dimension of each tensor in a group must be" + " divisible by 128.\n", + tidx, + ) + # --- shared memory (allocated once, reused across jobs) --- @cute.struct class SharedStorage: @@ -477,7 +698,8 @@ class SharedStorage: 128, ] - storage = cutlass.utils.SmemAllocator().allocate(SharedStorage) + smem = cutlass.utils.SmemAllocator() + storage = smem.allocate(SharedStorage) tile_layout = cute.make_layout( ((self.BUFF_DIM_Y, self.BUFF_DIM_X), self.PIPELINE_DEPTH), stride=((self.BUFF_DIM_X, 1), self.BUFF_DIM_Y * self.BUFF_DIM_X), @@ -486,12 +708,43 @@ class SharedStorage: sO_row = storage.sO_row.get_tensor(tile_layout) sO_col = storage.sO_col.get_tensor(tile_layout) + sActInput = None + if cutlass.const_expr(cfg.WITH_DACT): + + @cute.struct + class DactStorage: + sActInput: cute.struct.Align[ + cute.struct.MemRange[ + dtype, self.BUFF_DIM_Y * self.BUFF_DIM_X * self.PIPELINE_DEPTH + ], + 128, + ] + + sActInput = smem.allocate(DactStorage).sActInput.get_tensor(tile_layout) + + sDbias = None + if cutlass.const_expr(self.DBIAS_IN_ROWWISE): + # Padded like the CUDA kernel's partial_dbias_rowwise to avoid bank conflicts. + DBIAS_BUFF_WIDTH = self.THREADS_X * (MXFP8_BLOCK_SCALING_SIZE + 1) + + @cute.struct + class DbiasStorage: + sDbias: cute.struct.MemRange[Float32, self.THREADS_Y * DBIAS_BUFF_WIDTH] + + sDbias = smem.allocate(DbiasStorage).sDbias.get_tensor( + cute.make_layout((self.THREADS_Y, self.BUFF_DIM_X), stride=(DBIAS_BUFF_WIDTH, 1)) + ) + + # Grad and activation input share each stage's barrier. + tx_count = self.BUFF_DIM_Y * self.BUFF_DIM_X * dtype.width // 8 + if cutlass.const_expr(cfg.WITH_DACT): + tx_count *= 2 mainloop_pipeline = pipeline.PipelineTmaAsync.create( barrier_storage=storage.mbar.data_ptr(), num_stages=self.PIPELINE_DEPTH, producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, self.NUM_WARPS), - tx_count=self.BUFF_DIM_Y * self.BUFF_DIM_X * dtype.width // 8, + tx_count=tx_count, cta_layout_vmnk=None, ) prod_state = pipeline.make_pipeline_state( @@ -505,6 +758,13 @@ class SharedStorage: # the descriptor is swapped per tensor and the tile coords are tensor-local. gX_tiled = cute.zipped_divide(tma_src, (self.BUFF_DIM_Y, self.BUFF_DIM_X)) tXsX, tXgX = cpasync.tma_partition(tma_atom_x, 0, cute.make_layout(1), sX, gX_tiled) + tXsA = None + tXgA = None + if cutlass.const_expr(cfg.WITH_DACT): + gA_tiled = cute.zipped_divide(tma_src_act, (self.BUFF_DIM_Y, self.BUFF_DIM_X)) + tXsA, tXgA = cpasync.tma_partition( + tma_atom_act, 0, cute.make_layout(1), sActInput, gA_tiled + ) gO_row_tiled = cute.zipped_divide(tma_dst_out_row, (self.BUFF_DIM_Y, self.BUFF_DIM_X)) tXsO_row, tXgO_row = cpasync.tma_partition( tma_atom_out_row, 0, cute.make_layout(1), sO_row, gO_row_tiled @@ -544,9 +804,8 @@ class SharedStorage: # Which row does this block start from block_offset_Y = block_id_Y * self.CHUNK_DIM_Y if cutlass.const_expr(cfg.SHAPE_REP == VARYING_FIRST_DIM): - # Check if this block contains any valid tokens since the last offset <= logical_first_dim * logical_last_dim - # Last CSR offset == the group's total element count. Indexed off the - # array's own length so it does not depend on num_tensors. + # logical_shape may describe graph-safe capacity beyond the active tensors, + # whose total element count is the last CSR offset. total_elts = Int64(mOffsets[mOffsets.shape[0] - 1]) if Int64(block_offset_Y) * Int64(last_logical_dim) >= total_elts: has_work = Boolean(False) @@ -560,7 +819,7 @@ class SharedStorage: tensor_base = Int64(meta[2]) if tensor_rows > 0 and tensor_cols > 0: # How many blocks does this tensor have in both directions - block_columns_in_tensor = cute.ceil_div(tensor_cols, self.CHUNK_DIM_X) + block_columns_in_tensor = cute.ceil_div(tensor_cols, self.CHUNK_WIDTH) block_rows_in_tensor = cute.ceil_div(tensor_rows, self.CHUNK_DIM_Y) # How many blocks does this tensor have blocks_in_tensor = block_columns_in_tensor * block_rows_in_tensor @@ -575,10 +834,51 @@ class SharedStorage: # This tensor is empty, so this CTA has no work to do has_work = Boolean(False) + partitions = (tXsX, tXgX, tXsA, tXgA, tXsO_row, tXgO_row, tXsO_col, tXgO_col) + atoms = (tma_atom_x, tma_atom_act, tma_atom_out_row, tma_atom_out_col) + if has_work: if cutlass.const_expr(cfg.IS_SINGLE_TENSOR): # For single tensor case we don't use tensor descriptors - desc_x = desc_out_row = desc_out_col = None + descs = (None, None, None, None) + + # Rowwise scales span the whole group: members are stacked on 128-row + # boundaries, so each member's swizzled tiles follow the previous member's. + row_scales = None + if cutlass.const_expr(cfg.ROWWISE): + row_scales = self._rowwise_scales(mS_row, Int64(0), tensor_rows, tensor_cols) + col_scales = None + col_scale_row0 = block_offset_Y + col_scale_rows = tensor_rows + if cutlass.const_expr(cfg.COLWISE): + col_scale_base = Int64(0) + if cutlass.const_expr(cfg.WITH_GEMM_SWIZZLED_SCALES): + # Colwise swizzled scale indices restart at each member and depend + # on its rows (process_colwise_stage), so address the member that + # owns this chunk. + member_rows = Int32(0) + member_row0 = Int32(0) + if cutlass.const_expr(cfg.SHAPE_REP == SAME_BOTH_DIMS): + member_rows = tensor_rows // Int32(num_tensors) + member_row0 = block_offset_Y // member_rows * member_rows + else: + member_id = self._find_tensor_from_offsets( + mOffsets, + num_tensors, + Int64(block_offset_Y) * Int64(tensor_cols), + ) + member_rows = Int32(mFirstDims[member_id]) + member_row0 = Int32(Int64(mOffsets[member_id]) // Int64(tensor_cols)) + col_scale_base = ( + Int64(member_row0) + * Int64(cute.round_up(tensor_cols, 128)) + // MXFP8_BLOCK_SCALING_SIZE + ) + col_scale_row0 = block_offset_Y - member_row0 + col_scale_rows = member_rows + col_scales = self._colwise_scales( + mS_col, col_scale_base, col_scale_rows, tensor_cols + ) cute.arch.sync_threads() @@ -587,28 +887,23 @@ class SharedStorage: block_id_X, tensor_rows, tensor_cols, - tensor_base, - desc_x, - desc_out_row, - desc_out_col, - mS_row, - mS_col, - first_logical_dim, + row_scales, + col_scales, + col_scale_row0, + col_scale_rows, + block_offset_Y // self.CHUNK_DIM_Y, + mWorkspace, + sDbias, + descs, tmap, warp_idx, tidx, sX, + sActInput, sO_row, sO_col, - tXsX, - tXgX, - tXsO_row, - tXgO_row, - tXsO_col, - tXgO_col, - tma_atom_x, - tma_atom_out_row, - tma_atom_out_col, + partitions, + atoms, mainloop_pipeline, prod_state, cons_state, @@ -618,6 +913,9 @@ class SharedStorage: desc_x = tmap.get_tensormap_ptr(mTensormaps[(tensor_id, 0, None)].iterator) desc_out_row = tmap.get_tensormap_ptr(mTensormaps[(tensor_id, 1, None)].iterator) desc_out_col = tmap.get_tensormap_ptr(mTensormaps[(tensor_id, 2, None)].iterator) + desc_act = tmap.get_tensormap_ptr( + mTensormaps[(tensor_id, ACT_INPUT_SLOT, None)].iterator + ) # Acquire the descriptors on ONE thread, as the CUDA kernel does # (`leading_thread` in group_quantize_mxfp8.cuh); the sync_threads below # publishes it CTA-wide. Running the tensormap acquire fence on all 128 @@ -626,10 +924,22 @@ class SharedStorage: # 133 us -> 58 us), since the cost scales with threads x descriptors. if tidx == 0: tmap.fence_tensormap_update(desc_x) + if cutlass.const_expr(cfg.WITH_DACT): + tmap.fence_tensormap_update(desc_act) if cutlass.const_expr(cfg.ROWWISE): tmap.fence_tensormap_update(desc_out_row) if cutlass.const_expr(cfg.COLWISE): tmap.fence_tensormap_update(desc_out_col) + descs = (desc_x, desc_act, desc_out_row, desc_out_col) + + # This tensor's scales start at tensor_base / 32 in both directions. + scale_base = tensor_base // Int64(MXFP8_BLOCK_SCALING_SIZE) + row_scales = None + if cutlass.const_expr(cfg.ROWWISE): + row_scales = self._rowwise_scales(mS_row, scale_base, tensor_rows, tensor_cols) + col_scales = None + if cutlass.const_expr(cfg.COLWISE): + col_scales = self._colwise_scales(mS_col, scale_base, tensor_rows, tensor_cols) cute.arch.sync_threads() @@ -639,37 +949,69 @@ class SharedStorage: while not job_finished: block_id_Y_in_tensor = block_id // block_columns_in_tensor block_id_X_in_tensor = block_id % block_columns_in_tensor - self._process_block( - block_id_Y_in_tensor * self.CHUNK_DIM_Y, - block_id_X_in_tensor, - tensor_rows, - tensor_cols, - tensor_base, - desc_x, - desc_out_row, - desc_out_col, - mS_row, - mS_col, - first_logical_dim, - tmap, - warp_idx, - tidx, - sX, - sO_row, - sO_col, - tXsX, - tXgX, - tXsO_row, - tXgO_row, - tXsO_col, - tXgO_col, - tma_atom_x, - tma_atom_out_row, - tma_atom_out_col, - mainloop_pipeline, - prod_state, - cons_state, - ) + block_offset_Y_in_tensor = block_id_Y_in_tensor * self.CHUNK_DIM_Y + if cutlass.const_expr(self.STAGES_X == 1): + self._process_block( + block_offset_Y_in_tensor, + block_id_X_in_tensor, + tensor_rows, + tensor_cols, + row_scales, + col_scales, + block_offset_Y_in_tensor, + tensor_rows, + Int32(0), # dbias is only supported for single-tensor reps + mWorkspace, + sDbias, + descs, + tmap, + warp_idx, + tidx, + sX, + sActInput, + sO_row, + sO_col, + partitions, + atoms, + mainloop_pipeline, + prod_state, + cons_state, + ) + else: + # The chunk's column tiles in order, stopping at the tensor's last + # column like the CUDA kernel's stages_X = DIVUP(chunk_cols, TILE_DIM_X). + chunk_col0 = block_id_X_in_tensor * self.CHUNK_WIDTH + tiles_X = cutlass.min( + Int32(self.STAGES_X), + cute.ceil_div(tensor_cols - chunk_col0, self.BUFF_DIM_X), + ) + for stage_X in cutlass.range(tiles_X, unroll=1): + self._process_block( + block_offset_Y_in_tensor, + block_id_X_in_tensor * self.STAGES_X + stage_X, + tensor_rows, + tensor_cols, + row_scales, + col_scales, + block_offset_Y_in_tensor, + tensor_rows, + Int32(0), # dbias is only supported for single-tensor reps + mWorkspace, + sDbias, + descs, + tmap, + warp_idx, + tidx, + sX, + sActInput, + sO_row, + sO_col, + partitions, + atoms, + mainloop_pipeline, + prod_state, + cons_state, + ) # Find the next block to process block_id = block_id + block_stride if block_id >= blocks_in_tensor: @@ -686,126 +1028,98 @@ def _issue_load( prod_state, tile_y, tile_x, - tma_atom_x, - tXgX, - tXsX, + atoms, + partitions, tmap, - desc_x, + descs, ): - """Emit one 32x128 TMA load into the current pipeline buffer. + """Emit the 32x128 TMA load(s) of one stage into the current pipeline buffer. Caller gates this on warp 0 and advances `prod_state` afterwards -- the advance must happen outside the gate or the mutated SSA values stay trapped in the scf.if. """ + tma_atom_x, tma_atom_act, _, _ = atoms + tXsX, tXgX, tXsA, tXgA, _, _, _, _ = partitions + desc_x, desc_act, _, _ = descs # Wait for the consumer to finish using this SMEM buffer pipeline_obj.producer_acquire(prod_state) - if cutlass.const_expr(self.cfg.IS_SINGLE_TENSOR): - cute.copy( - tma_atom_x, - tXgX[(None, (tile_y, tile_x))], - tXsX[(None, prod_state.index)], - tma_bar_ptr=pipeline_obj.producer_get_barrier(prod_state), - ) - else: - # Every member shares tXgX's tile-coordinate arithmetic (the coefficients are - # just the tile size); tma_desc_ptr supplies this member's geometry. - cute.copy( - tma_atom_x, - tXgX[(None, (tile_y, tile_x))], - tXsX[(None, prod_state.index)], - tma_bar_ptr=pipeline_obj.producer_get_barrier(prod_state), - tma_desc_ptr=tmap.get_tensormap_ptr(desc_x, cute.AddressSpace.generic), - ) + barrier = pipeline_obj.producer_get_barrier(prod_state) + loads = [(tma_atom_x, tXgX, tXsX, desc_x)] + if cutlass.const_expr(self.cfg.WITH_DACT): + loads.append((tma_atom_act, tXgA, tXsA, desc_act)) + for atom, tXg, tXs, desc in loads: + if cutlass.const_expr(self.cfg.IS_SINGLE_TENSOR): + cute.copy( + atom, + tXg[(None, (tile_y, tile_x))], + tXs[(None, prod_state.index)], + tma_bar_ptr=barrier, + ) + else: + # Every member shares tXg's tile-coordinate arithmetic (the coefficients are + # just the tile size); tma_desc_ptr supplies this member's geometry. + cute.copy( + atom, + tXg[(None, (tile_y, tile_x))], + tXs[(None, prod_state.index)], + tma_bar_ptr=barrier, + tma_desc_ptr=tmap.get_tensormap_ptr(desc, cute.AddressSpace.generic), + ) # Notify the consumer that this SMEM buffer is ready for consumption pipeline_obj.producer_commit(prod_state) @cute.jit def _process_block( self, - block_offset_Y_in_tensor, # Row offset of this chunk (global for single-tensor, else tensor-local) - block_id_X_in_tensor, # Column-chunk index within the tensor - rows, # Number of rows in this tensor (logical shape) - cols, # Number of columns in this tensor (logical shape) - tensor_base, # Int64 element offset of this tensor in the group (0 for single-tensor) - desc_x, # Per-tensor descriptors, already acquired by the caller (None if single-tensor) - desc_out_row, - desc_out_col, - mS_row, # Grouped rowwise scales - mS_col, # Grouped colwise scales - first_logical_dim, # First dimension of the grouped tensor + block_offset_Y, # Row offset of this chunk (global for single-tensor, else tensor-local) + block_id_X, # Column-chunk index within the tensor + rows, # Rows of the rowwise-scale view (the group for single-tensor, else the tensor) + cols, # Number of columns in this tensor + row_scales, # Rowwise scales tiled per stage, rows counted like block_offset_Y + col_scales, # Colwise scales tiled per stage + col_scale_row0, # Row of this chunk in the colwise-scale view + col_scale_rows, # Rows of the colwise-scale view + dbias_row, # Row of the dbias workspace this chunk reduces into + mWorkspace, # f32 partial dbias workspace (WITH_DBIAS) + sDbias, # SMEM buffer for the rowwise dbias reduction (rowwise-only dbias) + descs, # Per-tensor descriptors (x, act, out_row, out_col), None if single-tensor tmap, # TensorMapManager for managing TMA descriptors warp_idx, tidx, - sX, # SMEM input for this block - sO_row, # SMEM rowwise output for this block - sO_col, # SMEM colwise output for this block - tXsX, # Tiled sX for TMA - tXgX, # Tiled gX for TMA - tXsO_row, # Tiled sO_row for TMA - tXgO_row, # Tiled gO_row for TMA - tXsO_col, # Tiled sO_col for TMA - tXgO_col, # Tiled gO_col for TMA - tma_atom_x, # TMA atom for input - tma_atom_out_row, # TMA atom for rowwise output - tma_atom_out_col, # TMA atom for colwise output + sX, # SMEM input ring + sActInput, # SMEM activation input ring (WITH_DACT) + sO_row, # SMEM rowwise output ring + sO_col, # SMEM colwise output ring + partitions, # TMA partitions (x, act, out_row, out_col) + atoms, # TMA atoms (x, act, out_row, out_col) mainloop_pipeline: cutlass.pipeline.PipelineTmaAsync, prod_state, cons_state, ): - """Quantize one 128x128 chunk in STAGES slices of BUFF_DIM_Y rows.""" + """Quantize one 128x128 tile of a chunk in STAGES slices of BUFF_DIM_Y rows.""" cfg = self.cfg - block_offset_X = block_id_X_in_tensor * self.CHUNK_DIM_X - - # This tensor's rowwise scales - scale_rows = Int32(first_logical_dim) if cutlass.const_expr(cfg.IS_SINGLE_TENSOR) else rows - scale_base = ( - Int64(0) - if cutlass.const_expr(cfg.IS_SINGLE_TENSOR) - else tensor_base // Int64(MXFP8_BLOCK_SCALING_SIZE) - ) - - if cutlass.const_expr(cfg.ROWWISE): - # Rowwise scale's divisibility guarantee: (128, 4) - scale_row_stride = cute.round_up(cute.ceil_div(cols, MXFP8_BLOCK_SCALING_SIZE), 4) - # Advance to this tensor's rowwise scales - mS_row_t = cute.make_tensor( - cute.make_ptr( - Float8E8M0FNU, - mS_row.iterator.toint() + scale_base, - cute.AddressSpace.gmem, - assumed_align=4, - ), - cute.make_layout((scale_rows, scale_row_stride), stride=(scale_row_stride, 1)), - ) - mS_row_tiled = cute.zipped_divide( - mS_row_t, (self.BUFF_DIM_Y, self.CHUNK_DIM_X // MXFP8_BLOCK_SCALING_SIZE) - ) - - if cutlass.const_expr(cfg.COLWISE): - # Colwise scale's divisibility guarantee: (4, 128) - scale_col_stride = cute.round_up(cols, 128) - # Advance to this tensor's colwise scales - mS_col_t = cute.make_tensor( - cute.make_ptr( - Float8E8M0FNU, - mS_col.iterator.toint() + scale_base, - cute.AddressSpace.gmem, - assumed_align=4, - ), - cute.make_layout( - (scale_rows // MXFP8_BLOCK_SCALING_SIZE, scale_col_stride), - stride=(scale_col_stride, 1), - ), - ) - mS_col_tiled = cute.zipped_divide( - mS_col_t, (self.BUFF_DIM_Y // MXFP8_BLOCK_SCALING_SIZE, self.CHUNK_DIM_X) - ) - - cute.arch.sync_threads() + _, _, tma_atom_out_row, tma_atom_out_col = atoms + _, _, _, _, tXsO_row, tXgO_row, tXsO_col, tXgO_col = partitions + _, _, desc_out_row, desc_out_col = descs + block_offset_X = block_id_X * self.CHUNK_DIM_X # This chunk's coordinates in the tile grid (32x128 TMA boxes, not elements). - tile_id_Y = block_offset_Y_in_tensor // self.BUFF_DIM_Y - tile_id_X = block_id_X_in_tensor + tile_id_Y = block_offset_Y // self.BUFF_DIM_Y + tile_id_X = block_id_X + col_scale_tile_Y = col_scale_row0 // self.BUFF_DIM_Y + + # Per-chunk dbias accumulators, in the CUDA kernel's summation order: a running + # column sum over the chunk's rows (colwise), or per-thread partial sums over its + # stages that the whole CTA reduces afterwards (rowwise-only). + dbias_col = Float32(0.0) + dbias_row_acc = None + if cutlass.const_expr(self.DBIAS_IN_ROWWISE): + dbias_row_acc = cute.make_rmem_tensor( + layout_or_shape=cute.make_layout((MXFP8_BLOCK_SCALING_SIZE,), stride=(1,)), + dtype=Float32, + ) + for c in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE): + dbias_row_acc[c] = Float32(0.0) # Fill every buffer up front, then issue one more each time a stage is consumed. for prologue_stage in cutlass.range_constexpr(self.PIPELINE_DEPTH): @@ -815,11 +1129,10 @@ def _process_block( prod_state, tile_id_Y + prologue_stage, tile_id_X, - tma_atom_x, - tXgX, - tXsX, + atoms, + partitions, tmap, - desc_x, + descs, ) prod_state.advance() @@ -833,40 +1146,50 @@ def _process_block( mainloop_pipeline.consumer_wait(cons_state) cute.arch.sync_threads() sX_tile = sX[(None, cons_state.index)] + sAct_tile = None + if cutlass.const_expr(cfg.WITH_DACT): + sAct_tile = sActInput[(None, cons_state.index)] row_tile = tile_id_Y + stage - tile_row_start = block_offset_Y_in_tensor + stage * self.BUFF_DIM_Y if cutlass.const_expr(cfg.COLWISE): - quantize_colwise_mxfp8( + _, dbias_col = quantize_colwise_mxfp8( sX_tile, - None, + sAct_tile, sO_col[(None, cons_state.index)], - cute.flatten(mS_col_tiled[(None, (row_tile, tile_id_X))]), + cute.flatten(col_scales[(None, (col_scale_tile_Y + stage, tile_id_X))]), cfg.MAX_NORM_RCP, - tile_row_start, + (col_scale_tile_Y + stage) * self.BUFF_DIM_Y, block_offset_X, - scale_rows, + col_scale_rows, cols, - ACTIVATION=None, + ACTIVATION=cfg.ACTIVATION, DTYPE=cfg.DTYPE, FP8_DTYPE=cfg.FP8_DTYPE, - SWIZZLE=False, + SWIZZLE=cfg.WITH_GEMM_SWIZZLED_SCALES, TILE_X=self.BUFF_DIM_X, TILE_Y=self.BUFF_DIM_Y, + WITH_ACT=cfg.WITH_ACT, + WITH_DACT=cfg.WITH_DACT, + WITH_DBIAS=self.DBIAS_IN_COLWISE, + CACHE_ACTIVATION=self.CACHE_ACTIVATION, ZERO_OOB_SCALES=True, + dbias_init=dbias_col, ) + if cutlass.const_expr(self.CACHE_ACTIVATION): + # The rowwise pass reads the activation the colwise pass cached in sX. + cute.arch.sync_threads() if cutlass.const_expr(cfg.ROWWISE): quantize_rowwise_mxfp8( sX_tile, - None, + None if self.CACHE_ACTIVATION else sAct_tile, sO_row[(None, cons_state.index)], - cute.flatten(mS_row_tiled[(None, (row_tile, tile_id_X))]), + cute.flatten(row_scales[(None, (row_tile, tile_id_X))]), cfg.MAX_NORM_RCP, - tile_row_start, + row_tile * self.BUFF_DIM_Y, block_offset_X, - scale_rows, + rows, cols, - ACTIVATION=None, + ACTIVATION=None if self.CACHE_ACTIVATION else cfg.ACTIVATION, DTYPE=cfg.DTYPE, FP8_DTYPE=cfg.FP8_DTYPE, TILE_X=self.BUFF_DIM_X, @@ -874,6 +1197,10 @@ def _process_block( WAVES=self.WAVES, THREADS_PER_BANK=self.THREADS_PER_BANK, PACK_SIZE=self.PACK_SIZE, + WITH_ACT=cfg.WITH_ACT and not self.CACHE_ACTIVATION, + WITH_DACT=cfg.WITH_DACT and not self.CACHE_ACTIVATION, + WITH_DBIAS=self.DBIAS_IN_ROWWISE, + dbias_acc=dbias_row_acc, ZERO_OOB_SCALES=True, ) @@ -893,69 +1220,97 @@ def _process_block( prod_state, tile_id_Y + stage + self.PIPELINE_DEPTH, tile_id_X, - tma_atom_x, - tXgX, - tXsX, + atoms, + partitions, tmap, - desc_x, + descs, ) prod_state.advance() # Write result to GMEM via TMA if warp_idx == 0: + stores = [] if cutlass.const_expr(cfg.ROWWISE): - if cutlass.const_expr(cfg.IS_SINGLE_TENSOR): - cute.copy( - tma_atom_out_row, - tXsO_row[(None, cons_state.index)], - tXgO_row[(None, (row_tile, tile_id_X))], - ) - else: - cute.copy( - tma_atom_out_row, - tXsO_row[(None, cons_state.index)], - tXgO_row[(None, (row_tile, tile_id_X))], - tma_desc_ptr=tmap.get_tensormap_ptr( - desc_out_row, cute.AddressSpace.generic - ), - ) + stores.append((tma_atom_out_row, tXsO_row, tXgO_row, desc_out_row)) if cutlass.const_expr(cfg.COLWISE): + stores.append((tma_atom_out_col, tXsO_col, tXgO_col, desc_out_col)) + for atom, tXs, tXg, desc in stores: if cutlass.const_expr(cfg.IS_SINGLE_TENSOR): cute.copy( - tma_atom_out_col, - tXsO_col[(None, cons_state.index)], - tXgO_col[(None, (row_tile, tile_id_X))], + atom, + tXs[(None, cons_state.index)], + tXg[(None, (row_tile, tile_id_X))], ) else: cute.copy( - tma_atom_out_col, - tXsO_col[(None, cons_state.index)], - tXgO_col[(None, (row_tile, tile_id_X))], - tma_desc_ptr=tmap.get_tensormap_ptr( - desc_out_col, cute.AddressSpace.generic - ), + atom, + tXs[(None, cons_state.index)], + tXg[(None, (row_tile, tile_id_X))], + tma_desc_ptr=tmap.get_tensormap_ptr(desc, cute.AddressSpace.generic), ) # Commit all TMA operations of this iteration cute.arch.cp_async_bulk_commit_group() cons_state.advance() + if cutlass.const_expr(cfg.WITH_DBIAS): + if cutlass.const_expr(self.DBIAS_IN_ROWWISE): + dbias_col = self._reduce_rowwise_dbias(sDbias, tidx, dbias_row_acc) + # One partial-dbias row per chunk, as in the CUDA kernel's dbias_workspace. + dbias_x = block_offset_X + tidx + if dbias_x < cols: + mWorkspace[(dbias_row, dbias_x)] = dbias_col + + @cute.jit + def _reduce_rowwise_dbias(self, sDbias, tidx, dbias_row_acc): + """Reduce the per-thread rowwise partial sums to one sum per column, in the order of the + CUDA kernel's partial_dbias_rowwise reduction.""" + _, tv_write = cute.make_layout_tv( + thr_layout=cute.make_layout( + (self.THREADS_Y, self.THREADS_X), stride=(self.THREADS_X, 1) + ), + val_layout=cute.make_layout( + (1, MXFP8_BLOCK_SCALING_SIZE), stride=(MXFP8_BLOCK_SCALING_SIZE, 1) + ), + ) + sDbias_write = cute.composition(sDbias, tv_write) + bank_group = (tidx % THREADS_PER_WARP) // self.THREADS_PER_BANK + offset = bank_group * self.PACK_SIZE + for w in cutlass.range_constexpr(self.WAVES): + # Undo the bank-conflict rotation quantize_rowwise_mxfp8 accumulated in. + start = (w * self.PACK_SIZE + offset) % MXFP8_BLOCK_SCALING_SIZE + for i in cutlass.range_constexpr(self.PACK_SIZE): + sDbias_write[(tidx, start + i)] = dbias_row_acc[w * self.PACK_SIZE + i] + cute.arch.sync_threads() + # Thread tidx sums column tidx over the THREADS_Y partial rows. + dbias = Float32(0.0) + for i in cutlass.range_constexpr(self.THREADS_Y): + dbias += sDbias[(i, tidx)] + # The buffer is rewritten by the next chunk. + cute.arch.sync_threads() + return dbias + def compile_cutedsl_function_from_cfg(cfg: MXFP8GroupQuantizeConfig): """Return the compiled CuTeDSL function object for the given grouped config.""" # CUDA requires the group's first logical dim to be a multiple of 128 (and each - # tensor's rows/cols likewise); MXFP8 needs the last dim divisible by 32. - sym_M = cute.sym_int32(divisibility=128) - sym_N = cute.sym_int32(divisibility=MXFP8_BLOCK_SCALING_SIZE) + # tensor's rows likewise). VARYING_BOTH_DIMS is the exception: its logical shape is + # [1, total]. The last dim only needs the 16-byte TMA row alignment; a partial 32-element + # scale block at the end of a row is zero-filled by TMA, as in the CUDA kernel. + if cfg.SHAPE_REP == VARYING_BOTH_DIMS: + sym_M = cute.sym_int32() + else: + sym_M = cute.sym_int32(divisibility=128) + sym_N = cute.sym_int32(divisibility=SYM_N_DIVISIBILITY) logical_shape = (sym_M, sym_N) out_dtype = cfg.FP8_DTYPE scale_dtype = cutlass.Float8E8M0FNU - def g2d(dtype, align=16): + def g2d(dtype, shape=logical_shape, align=16): return cute.runtime.make_fake_compact_tensor( dtype, - logical_shape, + shape, stride_order=(1, 0), memspace=cute.AddressSpace.gmem, assumed_align=align, @@ -972,14 +1327,6 @@ def g1d(dtype, align=4): # The kernel only takes the base address of the scale buffers (per-tensor strides # are derived from cols), so their fake shape is a flat 1D byte run. - in_fake = g2d(cfg.DTYPE) - out_row_fake = g2d(out_dtype) - out_col_fake = g2d(out_dtype) - scale_row_fake = g1d(scale_dtype) - scale_col_fake = g1d(scale_dtype) - offsets_fake = g1d(cutlass.Int64, align=8) - first_dims_fake = g1d(cutlass.Int64, align=8) - last_dims_fake = g1d(cutlass.Int64, align=8) tensormaps_fake = cute.runtime.make_fake_compact_tensor( cutlass.Int64, (cute.sym_int32(), NUM_WORKSPACE_SLOTS, BYTES_PER_TENSORMAP // 8), @@ -987,6 +1334,15 @@ def g1d(dtype, align=4): memspace=cute.AddressSpace.gmem, assumed_align=128, ) + # The cast-noop flag is an always-present f32 pointer instead of an optional tensor, so + # that one compiled kernel serves both an absent and a present flag (noop_flag_is_set). + noop_fake = cute.runtime.nullptr(Float32, mem_space=cute.AddressSpace.gmem, assumed_align=4) + act_input_fake = g2d(cfg.DTYPE) if cfg.WITH_DACT else None + workspace_fake = ( + g2d(Float32, shape=(cute.sym_int32(), cute.sym_int32()), align=4) + if cfg.WITH_DBIAS + else None + ) from cutlass.utils import HardwareInfo # pylint: disable=import-outside-toplevel @@ -994,15 +1350,18 @@ def g1d(dtype, align=4): kernel_obj = MXFP8GroupQuantizeKernel(cfg, sm_count) return cute.compile( kernel_obj, - in_fake, - out_row_fake, - out_col_fake, - scale_row_fake, - scale_col_fake, - offsets_fake, - first_dims_fake, - last_dims_fake, - tensormaps_fake, + g2d(cfg.DTYPE), # mX + g2d(out_dtype), # mO_row + g2d(out_dtype), # mO_col + g1d(scale_dtype), # mS_row + g1d(scale_dtype), # mS_col + g1d(cutlass.Int64, align=8), # mOffsets + g1d(cutlass.Int64, align=8), # mFirstDims + g1d(cutlass.Int64, align=8), # mLastDims + tensormaps_fake, # mTensormaps + noop_fake, # mNoop + act_input_fake, # mActInput + workspace_fake, # mWorkspace cute.runtime.make_fake_stream(), options="--enable-tvm-ffi", ) @@ -1015,6 +1374,11 @@ def get_mxfp8_group_quantization_function( rowwise: bool, colwise: bool, shape_rep: str, + with_gemm_swizzled_scales: bool, + with_dbias: bool, + with_dact: bool, + with_act: bool, + activation: str, ) -> bool: """Compile the grouped 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 @@ -1044,6 +1408,11 @@ def get_mxfp8_group_quantization_function( rowwise=rowwise, colwise=colwise, shape_rep=shape_rep, + with_gemm_swizzled_scales=with_gemm_swizzled_scales, + with_dbias=with_dbias, + with_dact=with_dact, + with_act=with_act, + activation=activation, ) except ValueError as e: logger.warning( diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8.py index a59b9745a24..9880276d396 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8.py @@ -432,6 +432,9 @@ def quantize_colwise_mxfp8( # Write 0 to the scales of columns past N instead of skipping them, so the padding of the # scale row is zeroed (mirrors group_quantize_mxfp8.cuh). ZERO_OOB_SCALES: cutlass.Constexpr[bool] = False, + # Running column sum to continue the dbias accumulation from, so that consecutive tiles add + # their elements in row order as the CUDA kernel does. Starts from 0 when None. + dbias_init: Optional[Float32] = None, ): """Quantize one SMEM tile colwise to MXFP8 (per-column 32-elt block scales); returns (amax, dbias_partial).""" tidx, _, _ = cute.arch.thread_idx() @@ -452,7 +455,7 @@ def quantize_colwise_mxfp8( 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) + dbias_partial = Float32(0.0) if cutlass.const_expr(dbias_init is None) else dbias_init if cutlass.const_expr(USE_HALF_PRECISION): max_scalar = max_scalar_f16 if DTYPE is cutlass.Float16 else max_scalar_bf16 @@ -553,10 +556,11 @@ def quantize_colwise_mxfp8( else: mS_col_stage[(0, tidx)] = biased_exp_c if cutlass.const_expr(ZERO_OOB_SCALES): - # Swizzled layouts pad through derive_swizzled_scale_layout instead. - assert not SWIZZLE if tile_row_start < M and scale_col >= N: - mS_col_stage[(0, tidx)] = Uint8(0).bitcast(Float8E8M0FNU) + if cutlass.const_expr(SWIZZLE): + mS_col_stage[(0, tidx % 32, tidx // 32)] = Uint8(0).bitcast(Float8E8M0FNU) + else: + mS_col_stage[(0, tidx)] = Uint8(0).bitcast(Float8E8M0FNU) inv_scale_c = exp2f_rcp(biased_exp_c) # cvt.rn.satfinite can be vectorized to convert 2 f32 to 2 fp8 in one instruction diff --git a/transformer_engine/common/cast/dispatch/quantize.cuh b/transformer_engine/common/cast/dispatch/quantize.cuh index ebed509813f..985c0f625d2 100644 --- a/transformer_engine/common/cast/dispatch/quantize.cuh +++ b/transformer_engine/common/cast/dispatch/quantize.cuh @@ -547,8 +547,8 @@ void group_quantize_fwd_helper(const NVTEGroupedTensor input, NVTEGroupedTensor quantized_with_cutedsl = cutedsl_backend::mxfp8_group_quantize_cutedsl( - input_tensor, noop_tensor, output_tensor, quant_config_cpp.mxfp8_2d_quantization, - stream); + input_tensor, activations_tensor, noop_tensor, output_tensor, dbias_tensor, + workspace_tensor, quant_config_cpp.mxfp8_2d_quantization, stream); #endif if (!quantized_with_cutedsl) { mxfp8::group_quantize( @@ -647,9 +647,19 @@ void group_quantize_bwd_helper(const NVTEGroupedTensor grad, const NVTEGroupedTe // Dispatch to quantization kernel depending on data format switch (scaling_mode) { case NVTE_MXFP8_1D_SCALING: { - mxfp8::group_quantize( - grad_tensor, input_tensor, noop_tensor, output_tensor, dbias_tensor, workspace_tensor, - &quant_config_cpp, stream); + bool quantized_with_cutedsl = false; +#ifdef NVTE_WITH_CUTEDSL + quantized_with_cutedsl = + cutedsl_backend::mxfp8_group_quantize_cutedsl( + grad_tensor, input_tensor, noop_tensor, output_tensor, dbias_tensor, workspace_tensor, + quant_config_cpp.mxfp8_2d_quantization, stream); +#endif + if (!quantized_with_cutedsl) { + mxfp8::group_quantize( + grad_tensor, input_tensor, noop_tensor, output_tensor, dbias_tensor, workspace_tensor, + &quant_config_cpp, stream); + } break; } case NVTE_BLOCK_SCALING_1D: diff --git a/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8_cutedsl.cuh b/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8_cutedsl.cuh index 41d4b23cd17..c1b1aae0edd 100644 --- a/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8_cutedsl.cuh +++ b/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8_cutedsl.cuh @@ -20,8 +20,9 @@ #include "../../common.h" #include "../../tvm_ffi_bridge.h" #include "../../util/cuda_runtime.h" -#include "../../utils.cuh" // ShapeRepresentation -#include "../core/grouped_tma.cuh" // dispatch::common::MAX_SUPPORTED_TENSOR_DESCRIPTORS +#include "../../util/cutedsl_utils.h" +#include "../../utils.cuh" // ShapeRepresentation +#include "../core/common.cuh" // MAX_SUPPORTED_TENSOR_DESCRIPTORS, grouped_reduce_dbias namespace transformer_engine { namespace cutedsl_backend { @@ -50,20 +51,30 @@ struct MXFP8GroupQuantConfig { bool rowwise; // If quantize rowwisely bool colwise; // If quantize columnwisely ShapeRepresentation shape_rep; // How the member shapes vary across the group + bool swizzled; // If the scales are written in the GEMM-swizzled layout + bool with_dbias; // If the partial dbias is computed (via the workspace tensor) + bool with_dact; // If an activation derivative operation is fused + bool with_act; // If an activation operation is fused + Activation activation = Activation::kNone; uint32_t sm_arch = static_cast(cuda::sm_arch()); // Bit layout: dtype [3:0] (4 used/reserved), fp8_dtype [7:4] (4 used/reserved), - // flags [9:8] (2 used), shape_rep [11:10] (2 used), and SM architecture [20:12] - // (9 used/reserved). Bits [31:21] are unused. + // flags [13:8] (6 used), shape_rep [15:14] (2 used), activation [21:16] (6 used/reserved), + // and SM architecture [30:22] (9 used/reserved). Bit 31 is unused. uint32_t to_id() const { static_assert(static_cast(DType::kNumTypes) <= 16, "DType no longer fits in the 4 bits to_id() gives it."); static_assert(ShapeRepresentation::VARYING_BOTH_DIMS < 4, "ShapeRepresentation no longer fits in the 2 bits to_id() gives it."); + static_assert(static_cast(Activation::kNumTypes) <= 64, + "Activation no longer fits in the 6 bits to_id() gives it."); NVTE_CHECK(sm_arch < 512, "SM architecture no longer fits in the 9 bits to_id() gives it."); return static_cast(dtype) | (static_cast(fp8_dtype) << 4) | (static_cast(rowwise) << 8) | (static_cast(colwise) << 9) | - (static_cast(shape_rep) << 10) | (sm_arch << 12); + (static_cast(swizzled) << 10) | (static_cast(with_dbias) << 11) | + (static_cast(with_dact) << 12) | (static_cast(with_act) << 13) | + (static_cast(shape_rep) << 14) | (static_cast(activation) << 16) | + (sm_arch << 22); } std::optional get_kernel() const { @@ -75,8 +86,8 @@ struct MXFP8GroupQuantConfig { // compiled and registered on a cache miss. std::string to_key() const { std::string key; - key.reserve( - 80); // longest: cutedsl_group_mxfp8_smXXX_BFloat16_Float8E4M3_1_1_varying_first_dim + // longest: cutedsl_group_mxfp8_smXXX_BFloat16_Float8E4M3_1_1_varying_first_dim_1_1_1_0_dqgelu + key.reserve(96); key.append("cutedsl_group_mxfp8_sm") .append(std::to_string(sm_arch)) .append("_") @@ -88,7 +99,17 @@ struct MXFP8GroupQuantConfig { .append("_") .append(colwise ? "1" : "0") .append("_") - .append(shape_rep_to_str(shape_rep)); + .append(shape_rep_to_str(shape_rep)) + .append("_") + .append(swizzled ? "1" : "0") + .append("_") + .append(with_dbias ? "1" : "0") + .append("_") + .append(with_dact ? "1" : "0") + .append("_") + .append(with_act ? "1" : "0") + .append("_") + .append(activation_to_str(activation)); return key; } @@ -100,15 +121,16 @@ struct MXFP8GroupQuantConfig { tvm::ffi::Any result = (*entrypoint)(tvm::ffi::String(fn_name), tvm::ffi::String(to_string(dtype)), tvm::ffi::String(to_string(fp8_dtype)), rowwise, colwise, - tvm::ffi::String(shape_rep_to_str(shape_rep))); + tvm::ffi::String(shape_rep_to_str(shape_rep)), swizzled, with_dbias, + with_dact, with_act, tvm::ffi::String(activation_to_str(activation))); return result.try_cast().value_or(false); } }; -// Descriptor slots per group member: input, rowwise output, colwise output, plus one -// carrying (rows, cols, base_elts). Mirrors NUM_WORKSPACE_SLOTS / BYTES_PER_TENSORMAP in -// CuTeDSL/cast/mxfp8/group_quantize_mxfp8.py. -constexpr size_t kGroupTensorMapSlots = 4; +// Descriptor slots per group member: input, rowwise output, colwise output, activation input, +// plus one carrying (rows, cols, base_elts). Mirrors NUM_WORKSPACE_SLOTS / BYTES_PER_TENSORMAP +// in CuTeDSL/cast/mxfp8/group_quantize_mxfp8.py. +constexpr size_t kGroupTensorMapSlots = 5; constexpr size_t kInt64PerTensorMap = 128 / sizeof(int64_t); constexpr size_t kMaxGroupTensors = static_cast(dispatch::common::MAX_SUPPORTED_TENSOR_DESCRIPTORS); @@ -151,25 +173,29 @@ inline NVTEBasicTensor make_basic_tensor(void *dptr, DType dtype, nvte_make_shape(shape.data(), shape.size())}; } -// Signature mirrors mxfp8::group_quantize (input, output, stream) for the subset the -// CuTeDSL kernel covers. Returns false to fall back to the CUDA kernel. +// Signature mirrors mxfp8::group_quantize (input, act_input, noop, output, dbias, workspace, +// stream). Returns false to fall back to the CUDA kernel. inline bool mxfp8_group_quantize_cutedsl(const MXFP8GroupQuantConfig &config, const GroupedTensor *input_tensor, - GroupedTensor *output_tensor, cudaStream_t stream) { + const GroupedTensor *act_input_tensor, + const Tensor *noop_tensor, GroupedTensor *output_tensor, + GroupedTensor *dbias_tensor, Tensor *workspace_tensor, + cudaStream_t stream) { const size_t num_tensors = input_tensor->num_tensors; const size_t first_logical_dim = input_tensor->logical_shape.data[0]; const size_t last_logical_dim = input_tensor->logical_shape.data[1]; // The kernel is compiled with cute.sym_int32(divisibility=...) on both logical extents, // so a violating shape would silently mis-tile rather than fail. These mirror sym_M / - // sym_N in CuTeDSL/cast/mxfp8/group_quantize_mxfp8.py -- the DSL kernel's own chunk - // height and MXFP8 block size, which it tiles independently of the CUDA kernel's - // CastTraits. + // sym_N in CuTeDSL/cast/mxfp8/group_quantize_mxfp8.py -- the DSL kernel's chunk height and + // the 16-byte TMA row alignment. The logical shape of VARYING_BOTH_DIMS is [1, total]. constexpr size_t kChunkDimY = 128; - constexpr size_t kScaleDimX = 32; - if (first_logical_dim % kChunkDimY != 0 || last_logical_dim % kScaleDimX != 0) { + constexpr size_t kLastDimAlignment = 16; + const bool first_dim_tiles = config.shape_rep == ShapeRepresentation::VARYING_BOTH_DIMS || + first_logical_dim % kChunkDimY == 0; + if (!first_dim_tiles || last_logical_dim % kLastDimAlignment != 0) { maybe_warn_cutedsl_not_chosen("the grouped logical shape is not a multiple of (", kChunkDimY, - ", ", kScaleDimX, ")."); + ", ", kLastDimAlignment, ")."); return false; } // The same extents are sym_int32 in the compiled kernel. @@ -179,6 +205,15 @@ inline bool mxfp8_group_quantize_cutedsl(const MXFP8GroupQuantConfig &config, return false; } + // dbias workspace-size query, mirroring mxfp8::group_quantize: the framework first calls + // with an unallocated workspace to learn its shape, allocates it, then calls again to run. + // The kernel writes one partial-dbias row per 128-row chunk. + if (config.with_dbias && workspace_tensor->data.dptr == nullptr) { + workspace_tensor->data.shape = {DIVUP(first_logical_dim, kChunkDimY), last_logical_dim}; + workspace_tensor->data.dtype = DType::kFloat32; + return true; + } + std::optional group_quant_func_opt = config.get_kernel(); if (!group_quant_func_opt.has_value()) { return false; @@ -188,7 +223,7 @@ inline bool mxfp8_group_quantize_cutedsl(const MXFP8GroupQuantConfig &config, GroupDescriptorWorkspace *const workspace = group_descriptor_workspace_ptr(); // Both output directions are handed to the kernel unconditionally: the compiled - // signature has no optional tensors, and building a TMA descriptor needs a real + // signature has no optional outputs, and building a TMA descriptor needs a real // address for each. The disabled direction is never read or written, so it points at // the enabled one instead of at a buffer that would have to be allocated. const SimpleTensor &data_row = @@ -233,46 +268,62 @@ inline bool mxfp8_group_quantize_cutedsl(const MXFP8GroupQuantConfig &config, {num_tensors, kGroupTensorMapSlots, kInt64PerTensorMap}), false, device_index); - // stream is a tvm-ffi opaque "handle"; pass it as void*. + // Optional inputs: a wrapper over a null buffer packs as TVM-FFI None. + DLTensorWrapper mActInput, mWorkspace; + if (config.with_dact) { + mActInput = DLTensorWrapper( + make_basic_tensor(act_input_tensor->data.dptr, act_input_tensor->dtype(), logical_shape), + true, device_index); + } + if (config.with_dbias) { + mWorkspace = DLTensorWrapper(workspace_tensor->data, true, device_index); + } + + // The cast-noop flag travels as a raw device pointer (not a tensor): it may be null, and the + // kernel null-checks it on device, so one compiled kernel serves both cases. + void *noop_ptr = (noop_tensor != nullptr) ? noop_tensor->data.dptr : nullptr; + + // noop and stream are tvm-ffi opaque "handles"; pass them as void*. (*group_quant_func_opt)(&mX, &mO_row, &mO_col, &mS_row, &mS_col, &mOffsets, &mFirstDims, - &mLastDims, &mTensormaps, static_cast(stream)); + &mLastDims, &mTensormaps, noop_ptr, &mActInput, &mWorkspace, + static_cast(stream)); + + // Reduce the per-chunk partial dbias per member with the CUDA kernel's reduction. + if (config.with_dbias) { + const float *workspace_ptr = reinterpret_cast(workspace_tensor->data.dptr); + TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY( + input_tensor->dtype(), IType, + dispatch::common::grouped_reduce_dbias( + config.shape_rep, num_tensors, first_logical_dim, last_logical_dim, + reinterpret_cast(output_tensor->tensor_offsets.dptr), + reinterpret_cast(output_tensor->first_dims.dptr), + reinterpret_cast(output_tensor->last_dims.dptr), dbias_tensor, + workspace_ptr, kChunkDimY, stream);) // NOLINT(*) + } return true; } template -bool mxfp8_group_quantize_cutedsl(const GroupedTensor *input_tensor, const Tensor *noop_tensor, - GroupedTensor *output_tensor, const bool use_2d_quantization, +bool mxfp8_group_quantize_cutedsl(const GroupedTensor *input_tensor, + const GroupedTensor *act_input_tensor, const Tensor *noop_tensor, + GroupedTensor *output_tensor, GroupedTensor *dbias_tensor, + Tensor *workspace_tensor, const bool use_2d_quantization, cudaStream_t stream) { if (!tvm_ffi_bridge::TVMFFICentral::getInstance().get_cutedsl_backend_enabled()) { maybe_warn_cutedsl_not_chosen("the CuTeDSL backend is disabled."); return false; } - // The CuTeDSL grouped kernel is cast-only: no dbias, no fused (derivative) activation. - if constexpr (IS_DBIAS || IS_DACT || IS_ACT || OP != nullptr) { - maybe_warn_cutedsl_not_chosen( - "grouped quantization with dbias or a fused activation is not supported."); + // TODO(kainingz): port 2D quantization to CuTeDSL + if (use_2d_quantization) { + maybe_warn_cutedsl_not_chosen("2D quantization is not supported."); + return false; + } + constexpr Activation activation = activation_func_to_enum(); + if constexpr (activation == Activation::kUnsupported) { + maybe_warn_cutedsl_not_chosen("the fused activation is not supported."); return false; } else { - // TODO(kainingz): port 2D quantization to CuTeDSL - if (use_2d_quantization) { - maybe_warn_cutedsl_not_chosen("2D quantization is not supported."); - return false; - } - // The kernel takes no noop flag, no amax accumulator, and writes compact scales only. - if (noop_tensor != nullptr && noop_tensor->data.dptr != nullptr) { - maybe_warn_cutedsl_not_chosen("the cast-noop flag is not supported."); - return false; - } - if (output_tensor->amax.dptr != nullptr) { - maybe_warn_cutedsl_not_chosen("amax computation is not supported."); - return false; - } - if (output_tensor->with_gemm_swizzled_scales) { - maybe_warn_cutedsl_not_chosen("GEMM-swizzled scales are not supported."); - return false; - } - // Mirrors the shape-representation selection in mxfp8::group_quantize. ShapeRepresentation shape_rep = ShapeRepresentation::SAME_BOTH_DIMS; if (output_tensor->all_same_shape()) { @@ -284,11 +335,9 @@ bool mxfp8_group_quantize_cutedsl(const GroupedTensor *input_tensor, const Tenso } else if (output_tensor->varying_both_dims()) { shape_rep = ShapeRepresentation::VARYING_BOTH_DIMS; } - if (shape_rep == ShapeRepresentation::VARYING_BOTH_DIMS) { - // The logical shape is [1, total], which is not tileable. - maybe_warn_cutedsl_not_chosen("groups with both dimensions varying are not supported."); - return false; - } + const bool is_single_tensor = shape_rep == ShapeRepresentation::SAME_BOTH_DIMS || + shape_rep == ShapeRepresentation::VARYING_FIRST_DIM; + // Leave invalid group sizes to mxfp8::group_quantize, which raises a proper error. // Every member gets a descriptor slot in the fixed-size workspace, so the CUDA // kernel's descriptor limit applies to the single-tensor representations here too. @@ -312,6 +361,11 @@ bool mxfp8_group_quantize_cutedsl(const GroupedTensor *input_tensor, const Tenso maybe_warn_cutedsl_not_chosen("the grouped tensor has no tensor offsets."); return false; } + if (IS_DBIAS && !is_single_tensor) { + // mxfp8::group_quantize raises a proper error for this. + maybe_warn_cutedsl_not_chosen("dbias is only supported for a common last dimension."); + return false; + } const bool rowwise = output_tensor->has_data(); const bool colwise = output_tensor->has_columnwise_data(); @@ -319,9 +373,18 @@ bool mxfp8_group_quantize_cutedsl(const GroupedTensor *input_tensor, const Tenso // mxfp8::group_quantize raises a proper error for this. return false; } + const bool swizzled = output_tensor->with_gemm_swizzled_scales; + if (swizzled && colwise && !is_single_tensor) { + // For these representations the CUDA kernel adds the tensor base to the colwise + // swizzled scale index twice, so leave them to it rather than reproduce that. + maybe_warn_cutedsl_not_chosen( + "GEMM-swizzled colwise scales are only supported for a common last dimension."); + return false; + } - checkCuDriverContext(stream); // Sanity checks, mirroring mxfp8::group_quantize + checkCuDriverContext(stream); + CheckNoopTensor(*noop_tensor, "cast_noop"); if (rowwise) { NVTE_CHECK(output_tensor->scale_inv.dptr != nullptr, "Scaling tensor must be allocated"); } @@ -333,13 +396,33 @@ bool mxfp8_group_quantize_cutedsl(const GroupedTensor *input_tensor, const Tenso "Number of input and output tensors must be same."); NVTE_CHECK(input_tensor->has_data(), "Cannot quantize tensor without rowwise data."); NVTE_CHECK(is_fp8_dtype(output_tensor->dtype()), "Output must have FP8 type."); + if constexpr (IS_DACT) { + NVTE_CHECK(act_input_tensor->has_data(), "Activations tensor must have data."); + NVTE_CHECK(input_tensor->num_tensors == act_input_tensor->num_tensors, + "Number of grad and activations tensors must be same."); + NVTE_CHECK(input_tensor->dtype() == act_input_tensor->dtype(), + "Grad and activations tensors must have the same type."); + } + if constexpr (IS_DBIAS) { + NVTE_CHECK(dbias_tensor->data.dtype == input_tensor->dtype(), + "DBias must have the same type as input_tensor."); + const Shape expected_shape_dbias_tensor = {num_tensors, input_tensor->logical_shape.data[1]}; + NVTE_CHECK(dbias_tensor->data.shape == expected_shape_dbias_tensor, "Wrong shape of DBias."); + NVTE_CHECK(workspace_tensor != nullptr, "Workspace must be a tensor."); + } const MXFP8GroupQuantConfig config{/*dtype=*/input_tensor->dtype(), /*fp8_dtype=*/output_tensor->dtype(), /*rowwise=*/rowwise, /*colwise=*/colwise, - /*shape_rep=*/shape_rep}; - return mxfp8_group_quantize_cutedsl(config, input_tensor, output_tensor, stream); + /*shape_rep=*/shape_rep, + /*swizzled=*/swizzled, + /*with_dbias=*/IS_DBIAS, + /*with_dact=*/IS_DACT, + /*with_act=*/IS_ACT, + /*activation=*/activation}; + return mxfp8_group_quantize_cutedsl(config, input_tensor, act_input_tensor, noop_tensor, + output_tensor, dbias_tensor, workspace_tensor, stream); } } From 2a0d2581aa7aa81f935e19c6b57d71e8526be3c9 Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Fri, 2 Oct 2026 03:16:22 +0000 Subject: [PATCH 05/58] tmp Signed-off-by: Kaining Zhong --- .../cast/mxfp8/group_quantize_mxfp8.py | 713 +++++++++++------- 1 file changed, 430 insertions(+), 283 deletions(-) diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8.py index 26e36db8100..b86476c5824 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8.py @@ -12,8 +12,8 @@ global block offsets -- the CUDA `tensor_map_*_static` "direct mapper" path. For SAME_BOTH_DIMS the CUDA grid is linearized per tensor (X, Y-in-tensor, tensor) while this kernel linearizes it flat over the stacked rows; both - require every member's row count to be a multiple of CHUNK_DIM_Y, and under - that precondition the two decode to the identical (block_offset_Y, block_id_X) + require every member's row count to be a multiple of 128, and under + that precondition the two decode to the identical (block_start_row, block_id_X) for every block index. * the other reps launch grid=(workers_per_tensor, num_tensors) and bind tensor_id to blockIdx.y, so a CTA grid-strides only within its own tensor and @@ -38,10 +38,6 @@ compact and GEMM-swizzled scales, rowwise and/or colwise, and all four shape representations. Differences from CUDA: * the grouped amax pointer is accepted and left untouched, as the CUDA kernel does; - * GEMM-swizzled colwise scales are only produced for the single-tensor reps. For - VARYING_LAST_DIM / VARYING_BOTH_DIMS the CUDA kernel adds the tensor base to the - colwise swizzled index twice, which is only in bounds for the first member, so the - C++ bridge falls back to CUDA instead of reproducing it; * dbias partial sums are accumulated in the CUDA kernel's order (a running column sum with colwise output, otherwise per-thread sums reduced across the CTA), and the C++ bridge reduces the workspace with the same grouped_reduce_dbias. @@ -75,7 +71,9 @@ from cutlass import pipeline from cutlass import Boolean, Float32, Int32, Int64, Float8E8M0FNU from cutlass.cute.nvgpu import cpasync +from cutlass.cute.testing import assert_ as runtime_assert from cutlass.utils import TensorMapManager, TensorMapUpdateMode +from cutlass.utils import HardwareInfo from cuda.bindings.driver import CUstream # pylint: disable=no-name-in-module import tvm_ffi @@ -157,19 +155,13 @@ def __init__( f" {SUPPORTED_SHAPE_REPS}" ) self.SHAPE_REP = shape_rep - # Mirrors `is_single_tensor` in group_quantize_mxfp8.cuh. + # Only can view the grouped tensor as one tensor when the last dim is same across all tensors self.IS_SINGLE_TENSOR = shape_rep in (SAME_BOTH_DIMS, VARYING_FIRST_DIM) self.MAX_NORM_RCP = ( FP8E4M3_MAX_NORM_RCP if fp8_dtype == "Float8E4M3" else FP8E5M2_MAX_NORM_RCP ) self.WITH_GEMM_SWIZZLED_SCALES = with_gemm_swizzled_scales - if with_gemm_swizzled_scales and colwise and not self.IS_SINGLE_TENSOR: - # The CUDA kernel offsets the colwise swizzled scales of these representations by - # the tensor base twice, which only lands inside the buffer for the first member. - raise ValueError( - "GEMM-swizzled colwise scales are only supported for single-tensor representations" - ) if with_dbias and not self.IS_SINGLE_TENSOR: # mxfp8::group_quantize raises for this. raise ValueError("dbias is only supported for tensors with a common last dimension") @@ -213,65 +205,48 @@ def __str__(self): class MXFP8GroupQuantizeKernel: """Grouped MXFP8 quantize mirroring group_quantize_mxfp8_kernel's strategy.""" - # TunableConfig / derived constants from group_quantize_mxfp8.cuh. - CHUNK_DIM_Y = 128 - CHUNK_DIM_X = 128 - THREADS_PER_CHUNK = 128 - STATIC_PERSISTENT_BLOCKS_PER_SM = 24 - ELTS_PER_CHUNK = CHUNK_DIM_Y * CHUNK_DIM_X - THREADS_X = CHUNK_DIM_X // MXFP8_BLOCK_SCALING_SIZE # 4 - THREADS_Y = THREADS_PER_CHUNK // THREADS_X # 32 - BUFF_DIM_Y = THREADS_Y # 32 - BUFF_DIM_X = CHUNK_DIM_X # 128 - # Each block of (CHUNK_DIM_Y, CHUNK_DIM_X) consists of STAGES tiles of (BUFF_DIM_X, BUFF_DIM_Y) stacked vertically - STAGES = CHUNK_DIM_Y // BUFF_DIM_Y # 4 - PIPELINE_DEPTH = 2 # PREFETCH_STAGES(1) + 1 - NUM_WARPS = THREADS_PER_CHUNK // THREADS_PER_WARP # 4 - - # Rowwise vectorization constants (mirror MXFP8QuantizeKernel / CUDA PACK_SIZE). + # Target persistent CTA count per SM for sizing the grid + STATIC_PERSISTENT_WORK_UNITS_PER_SM = 24 + # The shape of one pipeline stage processed by a CTA + TILE_ROWS = 32 + TILE_COLS = 128 + PIPELINE_DEPTH = 2 + # The number of elements in a tile + ELTS_PER_TILE = TILE_ROWS * TILE_COLS + # CTA shape + THREADS_PER_CTA = 128 + NUM_WARPS = THREADS_PER_CTA // THREADS_PER_WARP # 4 + THREADS_X = TILE_COLS // MXFP8_BLOCK_SCALING_SIZE # 4 + THREADS_Y = THREADS_PER_CTA // THREADS_X # 32 + # How many elements a thread handles in a wave PACK_SIZE = 4 + # How many waves needed to handle a MXFP8 block WAVES = MXFP8_BLOCK_SCALING_SIZE // PACK_SIZE # 8 + # How many threads per bank -- for avoiding bank conflicts THREADS_PER_BANK = (32 * 4) // MXFP8_BLOCK_SCALING_SIZE # 4 def __init__(self, cfg: MXFP8GroupQuantizeConfig, SM_COUNT: int): self.cfg = cfg self.SM_COUNT = SM_COUNT - # CastConfig widens the chunk to 128x256, i.e. STAGES_X = 2 tiles of - # BUFF_DIM_X columns, each traversed in STAGES row stages before moving right. - self.STAGES_X = 2 if cfg.SHAPE_REP == VARYING_BOTH_DIMS else 1 - self.CHUNK_WIDTH = self.CHUNK_DIM_X * self.STAGES_X + # A CTA processes (NUM_TILES_Y, NUM_TILES_X) tiles, NUM_STAGES tiles in total + self.NUM_TILES_X = 2 if cfg.SHAPE_REP == VARYING_BOTH_DIMS else 1 + self.NUM_TILES_Y = 4 + self.NUM_STAGES = self.NUM_TILES_X * self.NUM_TILES_Y + self.ELTS_PER_CTA = self.ELTS_PER_TILE * self.NUM_STAGES # The CUDA kernel honors the noop flag only without fused activations or dbias. self.CHECK_NOOP_FLAG = not (cfg.WITH_ACT or cfg.WITH_DACT or cfg.WITH_DBIAS) - # Like IS_CACHED_ACT_OP in CUDA: with both directions, the colwise pass caches the - # activation in the input tile for the rowwise pass. ReLU is fused into the conversion - # instead, which yields the same bytes. + # The colwise pass caches the activation in the input tile for the rowwise pass. + # Note: ReLU is fused into the conversion instead, which yields the same bytes. self.CACHE_ACTIVATION = ( (cfg.WITH_ACT or cfg.WITH_DACT) and cfg.ROWWISE and cfg.COLWISE and cfg.ACTIVATION != "relu" ) - # CUDA reduces dbias in the colwise pass when there is one, else in the rowwise pass. + # Prefer to reduce dbias in the colwise pass if quantized in columnwise, else in the rowwise pass. self.DBIAS_IN_COLWISE = cfg.WITH_DBIAS and cfg.COLWISE self.DBIAS_IN_ROWWISE = cfg.WITH_DBIAS and not cfg.COLWISE - # ---------------------------------------------------------------- helpers - @cute.jit - def _tensor_rows_cols( - self, tensor_id, mFirstDims, mLastDims, first_logical_dim, last_logical_dim - ): - """Get the shape (rows, cols) of the tensor by tensor_id.""" - cfg = self.cfg - if cutlass.const_expr(cfg.SHAPE_REP in (VARYING_FIRST_DIM, VARYING_BOTH_DIMS)): - rows = Int32(mFirstDims[tensor_id]) - else: - rows = Int32(first_logical_dim) - if cutlass.const_expr(cfg.SHAPE_REP in (VARYING_LAST_DIM, VARYING_BOTH_DIMS)): - cols = Int32(mLastDims[tensor_id]) - else: - cols = Int32(last_logical_dim) - return rows, cols - @cute.jit def _find_tensor_from_offsets(self, mOffsets, num_tensors, offset: Int64): """Index of the tensor whose element range holds `offset` (find_tensor_from_offsets).""" @@ -314,7 +289,7 @@ def _rowwise_scales(self, mS_row, base: Int64, rows, cols): mS_row, base, cute.make_layout((rows, stride), stride=(stride, 1)) ) return cute.zipped_divide( - mS_t, (self.BUFF_DIM_Y, self.BUFF_DIM_X // MXFP8_BLOCK_SCALING_SIZE) + mS_t, (self.TILE_ROWS, self.TILE_COLS // MXFP8_BLOCK_SCALING_SIZE) ) @cute.jit @@ -333,21 +308,24 @@ def _colwise_scales(self, mS_col, base: Int64, rows, cols): cute.make_layout((rows // MXFP8_BLOCK_SCALING_SIZE, stride), stride=(stride, 1)), ) return cute.zipped_divide( - mS_t, (self.BUFF_DIM_Y // MXFP8_BLOCK_SCALING_SIZE, self.BUFF_DIM_X) + mS_t, (self.TILE_ROWS // MXFP8_BLOCK_SCALING_SIZE, self.TILE_COLS) ) - # ------------------------------------------------------------ entry point @cute.jit def __call__( self, mX: cute.Tensor, - mO_row: cute.Tensor, - mO_col: cute.Tensor, + mO_row: Optional[cute.Tensor], + mO_col: Optional[cute.Tensor], mS_row: cute.Tensor, mS_col: cute.Tensor, mOffsets: cute.Tensor, # int64[num_tensors + 1], CSR element offsets - mFirstDims: cute.Tensor, # int64[num_tensors] (VARYING_FIRST_DIM / VARYING_BOTH_DIMS) - mLastDims: cute.Tensor, # int64[num_tensors] (VARYING_LAST_DIM / VARYING_BOTH_DIMS) + mFirstDims: Optional[ + cute.Tensor + ], # int64[num_tensors] (VARYING_FIRST_DIM / VARYING_BOTH_DIMS) + mLastDims: Optional[ + cute.Tensor + ], # int64[num_tensors] (VARYING_LAST_DIM / VARYING_BOTH_DIMS) mTensormaps: cute.Tensor, # int64[num_tensors, NUM_WORKSPACE_SLOTS, 16] mNoop: cute.Pointer, # f32 cast_noop flag; may be null, checked on device mActInput: Optional[cute.Tensor], # activation input, only with WITH_DACT @@ -360,18 +338,19 @@ def __call__( cfg = self.cfg first_logical_dim = mX.shape[0] last_logical_dim = mX.shape[1] - # Number of group members. Do NOT derive this from mOffsets: a caller is free to - # pass a length-num_tensors stub for SAME_BOTH_DIMS (where the offsets array is - # unused), which would make `mOffsets.shape[0] - 1` read one too few and divide by - # zero at num_tensors == 1. The per-tensor descriptor workspace is num_tensors long - # by construction -- one slot set per member -- so it is the reliable source. + if cutlass.const_expr(cfg.SHAPE_REP == VARYING_BOTH_DIMS): + runtime_assert( + first_logical_dim == 1, "VARYING_BOTH_DIMS requires logical shape [1, total]" + ) + num_tensors = mTensormaps.shape[0] + runtime_assert(num_tensors > 0, "Grouped quantization requires at least one tensor") - smem_tile_layout = cute.make_ordered_layout( - (self.BUFF_DIM_Y, self.BUFF_DIM_X), order=(1, 0) - ) - cta_tiler = (self.BUFF_DIM_Y, self.BUFF_DIM_X) + # A TMA atom copies a TILE at a time + smem_tile_layout = cute.make_ordered_layout((self.TILE_ROWS, self.TILE_COLS), order=(1, 0)) + cta_tiler = (self.TILE_ROWS, self.TILE_COLS) + # TMA atom for loading the input tensor and the activation input (if WITH_DACT) op_load = cpasync.CopyBulkTensorTileG2SOp() tma_atom_x, tma_src = cpasync.make_tiled_tma_atom( op_load, mX, smem_tile_layout, cta_tiler, num_multicast=1 @@ -382,57 +361,72 @@ def __call__( tma_atom_act, tma_src_act = cpasync.make_tiled_tma_atom( op_load, mActInput, smem_tile_layout, cta_tiler, num_multicast=1 ) + + # TMA atom for storing the rowwise and colwise outputs (if enabled) op_store = cpasync.CopyBulkTensorTileS2GOp() - tma_atom_out_row, tma_dst_out_row = cpasync.make_tiled_tma_atom( - op_store, mO_row, smem_tile_layout, cta_tiler, num_multicast=1 - ) - tma_atom_out_col, tma_dst_out_col = cpasync.make_tiled_tma_atom( - op_store, mO_col, smem_tile_layout, cta_tiler, num_multicast=1 - ) + tma_atom_out_row = None + tma_dst_out_row = None + if cutlass.const_expr(cfg.ROWWISE): + tma_atom_out_row, tma_dst_out_row = cpasync.make_tiled_tma_atom( + op_store, mO_row, smem_tile_layout, cta_tiler, num_multicast=1 + ) + tma_atom_out_col = None + tma_dst_out_col = None + if cutlass.const_expr(cfg.COLWISE): + tma_atom_out_col, tma_dst_out_col = cpasync.make_tiled_tma_atom( + op_store, mO_col, smem_tile_layout, cta_tiler, num_multicast=1 + ) if cutlass.const_expr(cfg.IS_SINGLE_TENSOR): - # How many blocks does the grouped tensor have in both directions - work_blocks_X = cute.ceil_div(Int32(last_logical_dim), self.CHUNK_DIM_X) - work_blocks_Y = cute.ceil_div(Int32(first_logical_dim), self.CHUNK_DIM_Y) - # Each CTA handles one chunk. With every member's rows a multiple of CHUNK_DIM_Y, - # this linear order is the CUDA kernel's (X, Y-in-tensor, tensor) order. + # How many CTAs does the grouped tensor have in both directions + work_blocks_Y = cute.ceil_div( + Int32(first_logical_dim), (self.TILE_ROWS * self.NUM_TILES_Y) + ) + work_blocks_X = cute.ceil_div( + Int32(last_logical_dim), self.TILE_COLS * self.NUM_TILES_X + ) + # Flatten it to an 1D grid grid = [work_blocks_X * work_blocks_Y, 1, 1] else: - # The work-block grid is per-tensor here: each CTA derives its own block range - # from its tensor's extents, so work_blocks_X/Y are unused on this path. - work_blocks_Y = Int32(1) - work_blocks_X = Int32(1) - # Persistent worker count, mirroring get_launch_config() in - # group_quantize_mxfp8.cuh: SM_COUNT * STATIC_PERSISTENT_BLOCKS_PER_SM workers - # split evenly across tensors, clamped to the average number of chunks a tensor - # holds. The element count would wrap Int32, so the estimate - # DIVUP(elts_total, CHUNK_DIM_Y * TILE_DIM_X) is formed without it. - n_tensors = cutlass.max(Int32(num_tensors), Int32(1)) # never divide by zero + # A placeholder for the kernel signature only; we won't use it in non-single tensor cases + work_blocks_X = None + # Estimate how many CTA work unit we need to do, where one work unit is NUM_TILES_X * NUM_TILES_Y tiles if cutlass.const_expr(cfg.SHAPE_REP == VARYING_BOTH_DIMS): - # The logical shape is [1, total]. - estimated_work_blocks = cute.ceil_div(Int32(last_logical_dim), self.ELTS_PER_CHUNK) - else: - # Exact: the first extent is compiled with divisibility CHUNK_DIM_Y. - estimated_work_blocks = cute.ceil_div( - (Int32(first_logical_dim) // self.CHUNK_DIM_Y) * Int32(last_logical_dim), - self.ELTS_PER_CHUNK // self.CHUNK_DIM_Y, + # Note: when VARYING_BOTH_DIMS, the first_logical_dim must be 1 + estimated_work_units = cute.ceil_div( + Int32(first_logical_dim) * Int32(last_logical_dim), self.ELTS_PER_CTA ) - estimated_work_blocks = cute.ceil_div(estimated_work_blocks, self.STAGES_X) - requested_workers_per_tensor = cutlass.max( + elif cutlass.const_expr(cfg.SHAPE_REP == VARYING_LAST_DIM): + # Same as VARYING_BOTH_DIMS but we divide 128 before multiplying to avoid overflowing Int32 + # because the first logical dimension is always 128-aligned when not VARYING_BOTH_DIMS + # so they are equivalent + estimated_work_units = cute.ceil_div( + (Int32(first_logical_dim) // 128) * Int32(last_logical_dim), + self.ELTS_PER_CTA // 128, + ) + else: + raise ValueError(f"unexpected shape representation {cfg.SHAPE_REP!r}") + + # All SMs can do SM_COUNT * STATIC_PERSISTENT_WORK_UNITS_PER_SM units of work, which are distributed evenly + # to all tensors in the group. + # If there are more tensors than work units, we will launch one CTA per tensor since we have enough parallelism here + requested_CTAs_per_tensor = cutlass.max( Int32(1), - Int32(self.SM_COUNT * self.STATIC_PERSISTENT_BLOCKS_PER_SM) // n_tensors, - ) - average_work_blocks_per_tensor = cutlass.max( - Int32(1), cute.ceil_div(estimated_work_blocks, n_tensors) + Int32(self.SM_COUNT * self.STATIC_PERSISTENT_WORK_UNITS_PER_SM) + // Int32(num_tensors), ) - workers_per_tensor = cutlass.min( - requested_workers_per_tensor, average_work_blocks_per_tensor + # In average how many work units per tensor (only an average, the actual work units per tensor may vary) + average_work_units_per_tensor = cutlass.max( + Int32(1), cute.ceil_div(estimated_work_units, Int32(num_tensors)) ) - grid = [workers_per_tensor, Int32(num_tensors), 1] + # Don't launch more CTAs than the average work units per tensor in case + # STATIC_PERSISTENT_WORK_UNITS_PER_SM causes redundancy + CTAs_per_tensor = cutlass.min(requested_CTAs_per_tensor, average_work_units_per_tensor) + grid = [CTAs_per_tensor, Int32(num_tensors), 1] # Only the multi-tensor representations need per-tensor descriptors. if cutlass.const_expr(not cfg.IS_SINGLE_TENSOR): - self.update_descriptors_kernel( + self._update_descriptors_kernel( mX, mO_row, mO_col, @@ -473,13 +467,12 @@ def __call__( tma_dst_out_col, ).launch( grid=grid, - block=[self.THREADS_PER_CHUNK, 1, 1], + block=[self.THREADS_PER_CTA, 1, 1], stream=stream, ) - # ------------------------------------------------- descriptor prologue @cute.kernel - def update_descriptors_kernel( + def _update_descriptors_kernel( self, mX, mO_row, @@ -497,20 +490,33 @@ def update_descriptors_kernel( tma_atom_ocol, tma_atom_act, ): - """One CTA per tensor: point that tensor's TMA descriptors at its own block. - - CuTeDSL analog of common::update_tma_descriptors writing g_tensor_maps[]. + """Update the per-tensor TMA descriptors for the group quantization kernel. + + mTensormaps: int64[num_tensors, NUM_WORKSPACE_SLOTS, 16], where the slots are: + - 0: input tensor + - 1: rowwise output tensor + - 2: colwise output tensor + - 3: activation input tensor (only with WITH_DACT) + - 4: metadata (rows, cols, base_offset) """ cfg = self.cfg + + # For the descriptor update kernel, grid=[num_tensors, 1, 1] tensor_id, _, _ = cute.arch.block_idx() - tidx, _, _ = cute.arch.thread_idx() - rows, cols = self._tensor_rows_cols( - tensor_id, mFirstDims, mLastDims, first_logical_dim, last_logical_dim - ) - base_elts = Int64(mOffsets[tensor_id]) + # Figure out how many rows and columns this tensor has, and where its first element is in the group. + if cutlass.const_expr(cfg.SHAPE_REP in (VARYING_FIRST_DIM, VARYING_BOTH_DIMS)): + rows = Int32(mFirstDims[tensor_id]) + else: + rows = Int32(first_logical_dim) + if cutlass.const_expr(cfg.SHAPE_REP in (VARYING_LAST_DIM, VARYING_BOTH_DIMS)): + cols = Int32(mLastDims[tensor_id]) + else: + cols = Int32(last_logical_dim) + base_offset = Int64(mOffsets[tensor_id]) # Same diagnostics as get_tensor_rows_num / get_tensor_cols_num. Like NVTE_DEVICE_ERROR # in a release build, they only print. + tidx, _, _ = cute.arch.thread_idx() if tidx == 0: if rows % 128 != 0: cute.printf( @@ -525,13 +531,13 @@ def update_descriptors_kernel( tensor_id, ) - # Publish this tensor's geometry for the main kernel (written even when empty). meta = mTensormaps[(tensor_id, META_SLOT, None)] meta[0] = Int64(rows) meta[1] = Int64(cols) - meta[2] = base_elts + meta[2] = base_offset tmap = TensorMapManager(TensorMapUpdateMode.GMEM, BYTES_PER_TENSORMAP) + # Obtain the pointers of these descriptors desc_x = tmap.get_tensormap_ptr(mTensormaps[(tensor_id, 0, None)].iterator) desc_orow = tmap.get_tensormap_ptr(mTensormaps[(tensor_id, 1, None)].iterator) desc_ocol = tmap.get_tensormap_ptr(mTensormaps[(tensor_id, 2, None)].iterator) @@ -546,7 +552,7 @@ def member_view(tensor, elt_dtype): return cute.make_tensor( cute.make_ptr( elt_dtype, - tensor.iterator.toint() + base_elts * (elt_dtype.width // 8), + tensor.iterator.toint() + base_offset * (elt_dtype.width // 8), cute.AddressSpace.gmem, assumed_align=16, ), @@ -560,19 +566,20 @@ def member_view(tensor, elt_dtype): if cutlass.const_expr(cfg.ROWWISE): views.append(member_view(mO_row, cfg.FP8_DTYPE)) atoms.append(tma_atom_orow) - descs.append(desc_orow) tmap.init_tensormap_from_atom(tma_atom_orow, desc_orow, 0) + descs.append(desc_orow) if cutlass.const_expr(cfg.COLWISE): views.append(member_view(mO_col, cfg.FP8_DTYPE)) atoms.append(tma_atom_ocol) - descs.append(desc_ocol) tmap.init_tensormap_from_atom(tma_atom_ocol, desc_ocol, 0) + descs.append(desc_ocol) if cutlass.const_expr(cfg.WITH_DACT): views.append(member_view(mActInput, dtype)) atoms.append(tma_atom_act) - descs.append(desc_act) tmap.init_tensormap_from_atom(tma_atom_act, desc_act, 0) + descs.append(desc_act) tmap.fence_tensormap_initialization() + # Update descriptors in global memory with views and atoms tmap.update_tensormap( tuple(views), tuple(atoms), @@ -581,7 +588,125 @@ def member_view(tensor, elt_dtype): (), # smem staging is unused in GMEM update mode ) - # ------------------------------------------------------------ main kernel + @cute.jit + def _make_shared_storage( + self, + smem: cutlass.Constexpr, + dtype: cutlass.Constexpr[Type[cutlass.Numeric]], + ): + """Allocate pipeline buffers and optional activation input and dbias storage.""" + FP8_DTYPE = self.cfg.FP8_DTYPE + tile_layout = cute.make_layout( + ((self.TILE_ROWS, self.TILE_COLS), self.PIPELINE_DEPTH), + stride=((self.TILE_COLS, 1), self.TILE_ROWS * self.TILE_COLS), + ) + + sX = None + sO_row = None + sO_col = None + sActInput = None + + if cutlass.const_expr(self.cfg.ROWWISE and self.cfg.COLWISE): + + @cute.struct + class SharedStorage: + mbar: cute.struct.MemRange[cute.Int64, 2 * self.PIPELINE_DEPTH] + sX: cute.struct.Align[ + cute.struct.MemRange[ + dtype, self.TILE_ROWS * self.TILE_COLS * self.PIPELINE_DEPTH + ], + 128, + ] + sO_row: cute.struct.Align[ + cute.struct.MemRange[ + FP8_DTYPE, self.TILE_ROWS * self.TILE_COLS * self.PIPELINE_DEPTH + ], + 128, + ] + sO_col: cute.struct.Align[ + cute.struct.MemRange[ + FP8_DTYPE, self.TILE_ROWS * self.TILE_COLS * self.PIPELINE_DEPTH + ], + 128, + ] + + storage = smem.allocate(SharedStorage) + sX = storage.sX.get_tensor(tile_layout) + sO_row = storage.sO_row.get_tensor(tile_layout) + sO_col = storage.sO_col.get_tensor(tile_layout) + + elif cutlass.const_expr(self.cfg.ROWWISE): + + @cute.struct + class SharedStorage: + mbar: cute.struct.MemRange[cute.Int64, 2 * self.PIPELINE_DEPTH] + sX: cute.struct.Align[ + cute.struct.MemRange[ + dtype, self.TILE_ROWS * self.TILE_COLS * self.PIPELINE_DEPTH + ], + 128, + ] + sO_row: cute.struct.Align[ + cute.struct.MemRange[ + FP8_DTYPE, self.TILE_ROWS * self.TILE_COLS * self.PIPELINE_DEPTH + ], + 128, + ] + + storage = smem.allocate(SharedStorage) + sX = storage.sX.get_tensor(tile_layout) + sO_row = storage.sO_row.get_tensor(tile_layout) + + elif cutlass.const_expr(self.cfg.COLWISE): + + @cute.struct + class SharedStorage: + mbar: cute.struct.MemRange[cute.Int64, 2 * self.PIPELINE_DEPTH] + sX: cute.struct.Align[ + cute.struct.MemRange[ + dtype, self.TILE_ROWS * self.TILE_COLS * self.PIPELINE_DEPTH + ], + 128, + ] + sO_col: cute.struct.Align[ + cute.struct.MemRange[ + FP8_DTYPE, self.TILE_ROWS * self.TILE_COLS * self.PIPELINE_DEPTH + ], + 128, + ] + + storage = smem.allocate(SharedStorage) + sX = storage.sX.get_tensor(tile_layout) + sO_col = storage.sO_col.get_tensor(tile_layout) + + if cutlass.const_expr(self.cfg.WITH_DACT): + + @cute.struct + class DactStorage: + sActInput: cute.struct.Align[ + cute.struct.MemRange[ + dtype, self.TILE_ROWS * self.TILE_COLS * self.PIPELINE_DEPTH + ], + 128, + ] + + sActInput = smem.allocate(DactStorage).sActInput.get_tensor(tile_layout) + + sDbias = None + if cutlass.const_expr(self.DBIAS_IN_ROWWISE): + # Padded like the CUDA kernel's partial_dbias_rowwise to avoid bank conflicts. + DBIAS_BUFF_WIDTH = self.THREADS_X * (MXFP8_BLOCK_SCALING_SIZE + 1) + + @cute.struct + class DbiasStorage: + sDbias: cute.struct.MemRange[Float32, self.THREADS_Y * DBIAS_BUFF_WIDTH] + + sDbias = smem.allocate(DbiasStorage).sDbias.get_tensor( + cute.make_layout((self.THREADS_Y, self.TILE_COLS), stride=(DBIAS_BUFF_WIDTH, 1)) + ) + + return storage.mbar.data_ptr(), sX, sO_row, sO_col, sActInput, sDbias + @cute.kernel def kernel( self, @@ -645,7 +770,7 @@ def _kernel_main( first_logical_dim, last_logical_dim, num_tensors, - work_blocks_X, + work_units_X, dtype: cutlass.Constexpr[Type[cutlass.Numeric]], tma_atom_x, tma_src, @@ -675,72 +800,18 @@ def _kernel_main( tidx, ) - # --- shared memory (allocated once, reused across jobs) --- - @cute.struct - class SharedStorage: - mbar: cute.struct.MemRange[cute.Int64, 2 * self.PIPELINE_DEPTH] - sX: cute.struct.Align[ - cute.struct.MemRange[ - dtype, self.BUFF_DIM_Y * self.BUFF_DIM_X * self.PIPELINE_DEPTH - ], - 128, - ] - sO_row: cute.struct.Align[ - cute.struct.MemRange[ - FP8_DTYPE, self.BUFF_DIM_Y * self.BUFF_DIM_X * self.PIPELINE_DEPTH - ], - 128, - ] - sO_col: cute.struct.Align[ - cute.struct.MemRange[ - FP8_DTYPE, self.BUFF_DIM_Y * self.BUFF_DIM_X * self.PIPELINE_DEPTH - ], - 128, - ] - smem = cutlass.utils.SmemAllocator() - storage = smem.allocate(SharedStorage) - tile_layout = cute.make_layout( - ((self.BUFF_DIM_Y, self.BUFF_DIM_X), self.PIPELINE_DEPTH), - stride=((self.BUFF_DIM_X, 1), self.BUFF_DIM_Y * self.BUFF_DIM_X), - ) - sX = storage.sX.get_tensor(tile_layout) - sO_row = storage.sO_row.get_tensor(tile_layout) - sO_col = storage.sO_col.get_tensor(tile_layout) - - sActInput = None - if cutlass.const_expr(cfg.WITH_DACT): - - @cute.struct - class DactStorage: - sActInput: cute.struct.Align[ - cute.struct.MemRange[ - dtype, self.BUFF_DIM_Y * self.BUFF_DIM_X * self.PIPELINE_DEPTH - ], - 128, - ] - - sActInput = smem.allocate(DactStorage).sActInput.get_tensor(tile_layout) - - sDbias = None - if cutlass.const_expr(self.DBIAS_IN_ROWWISE): - # Padded like the CUDA kernel's partial_dbias_rowwise to avoid bank conflicts. - DBIAS_BUFF_WIDTH = self.THREADS_X * (MXFP8_BLOCK_SCALING_SIZE + 1) - - @cute.struct - class DbiasStorage: - sDbias: cute.struct.MemRange[Float32, self.THREADS_Y * DBIAS_BUFF_WIDTH] - - sDbias = smem.allocate(DbiasStorage).sDbias.get_tensor( - cute.make_layout((self.THREADS_Y, self.BUFF_DIM_X), stride=(DBIAS_BUFF_WIDTH, 1)) - ) + # Layouts for sX, sO_row, sO_col, sActInput: + # ((TILE_ROWS, TILE_COLS), PIPELINE_DEPTH):((TILE_COLS, 1), TILE_ROWS * TILE_COLS) + # Layout for sDbias: + # (THREADS_Y, TILE_COLS):(THREADS_X * (MXFP8_BLOCK_SCALING_SIZE + 1), 1) + mbar_ptr, sX, sO_row, sO_col, sActInput, sDbias = self._make_shared_storage(smem, dtype) - # Grad and activation input share each stage's barrier. - tx_count = self.BUFF_DIM_Y * self.BUFF_DIM_X * dtype.width // 8 + tx_count = self.TILE_ROWS * self.TILE_COLS * dtype.width // 8 if cutlass.const_expr(cfg.WITH_DACT): tx_count *= 2 mainloop_pipeline = pipeline.PipelineTmaAsync.create( - barrier_storage=storage.mbar.data_ptr(), + barrier_storage=mbar_ptr, num_stages=self.PIPELINE_DEPTH, producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, self.NUM_WARPS), @@ -754,25 +825,32 @@ class DbiasStorage: pipeline.PipelineUserType.Consumer, self.PIPELINE_DEPTH ) - # TMA partitions built from the representative views. For multi-tensor reps - # the descriptor is swapped per tensor and the tile coords are tensor-local. - gX_tiled = cute.zipped_divide(tma_src, (self.BUFF_DIM_Y, self.BUFF_DIM_X)) + gX_tiled = cute.zipped_divide(tma_src, (self.TILE_ROWS, self.TILE_COLS)) tXsX, tXgX = cpasync.tma_partition(tma_atom_x, 0, cute.make_layout(1), sX, gX_tiled) + tXsA = None tXgA = None if cutlass.const_expr(cfg.WITH_DACT): - gA_tiled = cute.zipped_divide(tma_src_act, (self.BUFF_DIM_Y, self.BUFF_DIM_X)) + gA_tiled = cute.zipped_divide(tma_src_act, (self.TILE_ROWS, self.TILE_COLS)) tXsA, tXgA = cpasync.tma_partition( tma_atom_act, 0, cute.make_layout(1), sActInput, gA_tiled ) - gO_row_tiled = cute.zipped_divide(tma_dst_out_row, (self.BUFF_DIM_Y, self.BUFF_DIM_X)) - tXsO_row, tXgO_row = 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.BUFF_DIM_Y, self.BUFF_DIM_X)) - tXsO_col, tXgO_col = cpasync.tma_partition( - tma_atom_out_col, 0, cute.make_layout(1), sO_col, gO_col_tiled - ) + + tXsO_row = None + tXgO_row = None + 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 = cpasync.tma_partition( + tma_atom_out_row, 0, cute.make_layout(1), sO_row, gO_row_tiled + ) + + tXsO_col = None + tXgO_col = None + 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 = cpasync.tma_partition( + tma_atom_out_col, 0, cute.make_layout(1), sO_col, gO_col_tiled + ) tmap = TensorMapManager(TensorMapUpdateMode.GMEM, BYTES_PER_TENSORMAP) cute.arch.sync_threads() @@ -786,29 +864,30 @@ class DbiasStorage: # group can exceed 2^31 elements even when every individual extent is small. tensor_base = Int64(0) # Block's offset and id in this individual tensor / global single tensor - block_offset_Y = Int32(0) + block_start_row = Int32(0) block_id_X = Int32(0) first_block_id = Int32(0) - blocks_in_tensor = Int32(1) + units_in_tensor = Int32(1) block_stride = Int32(1) - block_columns_in_tensor = Int32(1) + units_X_in_tensor = Int32(1) if cutlass.const_expr(cfg.IS_SINGLE_TENSOR): - # grid = [work_blocks_X * work_blocks_Y, 1, 1] - block_id_Y = Int32(bidx) // work_blocks_X - block_id_X = Int32(bidx) % work_blocks_X + # grid = [work_units_X * work_units_Y, 1, 1] + block_id_Y = Int32(bidx) // work_units_X + block_id_X = Int32(bidx) % work_units_X # View the grouped tensor as a single tensor of shape (first_logical_dim, last_logical_dim) tensor_rows = Int32(first_logical_dim) tensor_cols = Int32(last_logical_dim) # Which row does this block start from - block_offset_Y = block_id_Y * self.CHUNK_DIM_Y + block_start_row = block_id_Y * (self.TILE_ROWS * self.NUM_TILES_Y) if cutlass.const_expr(cfg.SHAPE_REP == VARYING_FIRST_DIM): - # logical_shape may describe graph-safe capacity beyond the active tensors, - # whose total element count is the last CSR offset. total_elts = Int64(mOffsets[mOffsets.shape[0] - 1]) - if Int64(block_offset_Y) * Int64(last_logical_dim) >= total_elts: + # If the starting element is already beyond the last element of the group, this CTA can stop + if Int64(block_start_row) * Int64(last_logical_dim) >= total_elts: has_work = Boolean(False) + # When SAME_BOTH_DIM, M is always divisible by 128, which is exactly TILE_ROWS * NUM_TILES_Y, + # so no need to check for the last block's starting row being beyond the last row of the tensor. else: # grid = [workers_per_tensor, Int32(num_tensors), 1] tensor_id = Int32(bidy) @@ -818,17 +897,17 @@ class DbiasStorage: tensor_cols = Int32(meta[1]) tensor_base = Int64(meta[2]) if tensor_rows > 0 and tensor_cols > 0: - # How many blocks does this tensor have in both directions - block_columns_in_tensor = cute.ceil_div(tensor_cols, self.CHUNK_WIDTH) - block_rows_in_tensor = cute.ceil_div(tensor_rows, self.CHUNK_DIM_Y) - # How many blocks does this tensor have - blocks_in_tensor = block_columns_in_tensor * block_rows_in_tensor - # Which block (1D index) does this CTA start from + # How many work units does this tensor have in both directions + units_X_in_tensor = cute.ceil_div(tensor_cols, (self.TILE_COLS * self.NUM_TILES_X)) + units_Y_in_tensor = cute.ceil_div(tensor_rows, (self.TILE_ROWS * self.NUM_TILES_Y)) + # How many work units does this tensor have + units_in_tensor = units_X_in_tensor * units_Y_in_tensor + # Which work unit (1D index) does this CTA start from, which is also their worker ID first_block_id = Int32(bidx) - # gdx is workers_per_tensor (how many CTAs are assigned to this tensor) + # gdx is how many CTAs are assigned to this tensor block_stride = Int32(gdx) - # If my first block_id is already beyond the tensor's last block, I have no work to do - if first_block_id >= blocks_in_tensor: + # If my first work unit is already beyond the tensor's last work unit, I have no work to do + if first_block_id >= units_in_tensor: has_work = Boolean(False) else: # This tensor is empty, so this CTA has no work to do @@ -838,6 +917,7 @@ class DbiasStorage: atoms = (tma_atom_x, tma_atom_act, tma_atom_out_row, tma_atom_out_col) if has_work: + # For single tensor case, each CTA only processes one work unit if cutlass.const_expr(cfg.IS_SINGLE_TENSOR): # For single tensor case we don't use tensor descriptors descs = (None, None, None, None) @@ -847,8 +927,9 @@ class DbiasStorage: row_scales = None if cutlass.const_expr(cfg.ROWWISE): row_scales = self._rowwise_scales(mS_row, Int64(0), tensor_rows, tensor_cols) + col_scales = None - col_scale_row0 = block_offset_Y + col_scale_row_start = block_start_row col_scale_rows = tensor_rows if cutlass.const_expr(cfg.COLWISE): col_scale_base = Int64(0) @@ -860,12 +941,12 @@ class DbiasStorage: member_row0 = Int32(0) if cutlass.const_expr(cfg.SHAPE_REP == SAME_BOTH_DIMS): member_rows = tensor_rows // Int32(num_tensors) - member_row0 = block_offset_Y // member_rows * member_rows + member_row0 = block_start_row // member_rows * member_rows else: member_id = self._find_tensor_from_offsets( mOffsets, num_tensors, - Int64(block_offset_Y) * Int64(tensor_cols), + Int64(block_start_row) * Int64(tensor_cols), ) member_rows = Int32(mFirstDims[member_id]) member_row0 = Int32(Int64(mOffsets[member_id]) // Int64(tensor_cols)) @@ -874,7 +955,7 @@ class DbiasStorage: * Int64(cute.round_up(tensor_cols, 128)) // MXFP8_BLOCK_SCALING_SIZE ) - col_scale_row0 = block_offset_Y - member_row0 + col_scale_row_start = block_start_row - member_row0 col_scale_rows = member_rows col_scales = self._colwise_scales( mS_col, col_scale_base, col_scale_rows, tensor_cols @@ -883,15 +964,15 @@ class DbiasStorage: cute.arch.sync_threads() self._process_block( - block_offset_Y, + block_start_row, block_id_X, tensor_rows, tensor_cols, row_scales, col_scales, - col_scale_row0, + col_scale_row_start, col_scale_rows, - block_offset_Y // self.CHUNK_DIM_Y, + block_start_row // (self.TILE_ROWS * self.NUM_TILES_Y), mWorkspace, sDbias, descs, @@ -909,19 +990,14 @@ class DbiasStorage: cons_state, ) else: - # For multi-tensor case, retrieve tensor descriptors we processed early in the prologue kernel + # For non-single tensor case, we use persistent kernel so each CTA keeps processing work units until none is left desc_x = tmap.get_tensormap_ptr(mTensormaps[(tensor_id, 0, None)].iterator) desc_out_row = tmap.get_tensormap_ptr(mTensormaps[(tensor_id, 1, None)].iterator) desc_out_col = tmap.get_tensormap_ptr(mTensormaps[(tensor_id, 2, None)].iterator) desc_act = tmap.get_tensormap_ptr( mTensormaps[(tensor_id, ACT_INPUT_SLOT, None)].iterator ) - # Acquire the descriptors on ONE thread, as the CUDA kernel does - # (`leading_thread` in group_quantize_mxfp8.cuh); the sync_threads below - # publishes it CTA-wide. Running the tensormap acquire fence on all 128 - # threads is correct but very expensive -- it more than doubles the - # kernel time on the multi-tensor path (4096x14336 bidirectional: - # 133 us -> 58 us), since the cost scales with threads x descriptors. + if tidx == 0: tmap.fence_tensormap_update(desc_x) if cutlass.const_expr(cfg.WITH_DACT): @@ -941,24 +1017,27 @@ class DbiasStorage: if cutlass.const_expr(cfg.COLWISE): col_scales = self._colwise_scales(mS_col, scale_base, tensor_rows, tensor_cols) + # Make sure all threads see the updated descriptors and scales before processing any blocks. cute.arch.sync_threads() # Grid-stride over this tensor's own chunks; the descriptors never change. block_id = first_block_id job_finished = Boolean(False) while not job_finished: - block_id_Y_in_tensor = block_id // block_columns_in_tensor - block_id_X_in_tensor = block_id % block_columns_in_tensor - block_offset_Y_in_tensor = block_id_Y_in_tensor * self.CHUNK_DIM_Y - if cutlass.const_expr(self.STAGES_X == 1): + block_id_Y_in_tensor = block_id // units_X_in_tensor + block_id_X_in_tensor = block_id % units_X_in_tensor + block_start_row_in_tensor = block_id_Y_in_tensor * ( + self.TILE_ROWS * self.NUM_TILES_Y + ) + if cutlass.const_expr(self.NUM_TILES_X == 1): self._process_block( - block_offset_Y_in_tensor, + block_start_row_in_tensor, block_id_X_in_tensor, tensor_rows, tensor_cols, row_scales, col_scales, - block_offset_Y_in_tensor, + block_start_row_in_tensor, tensor_rows, Int32(0), # dbias is only supported for single-tensor reps mWorkspace, @@ -980,20 +1059,20 @@ class DbiasStorage: else: # The chunk's column tiles in order, stopping at the tensor's last # column like the CUDA kernel's stages_X = DIVUP(chunk_cols, TILE_DIM_X). - chunk_col0 = block_id_X_in_tensor * self.CHUNK_WIDTH + chunk_col0 = block_id_X_in_tensor * (self.TILE_COLS * self.NUM_TILES_X) tiles_X = cutlass.min( - Int32(self.STAGES_X), - cute.ceil_div(tensor_cols - chunk_col0, self.BUFF_DIM_X), + Int32(self.NUM_TILES_X), + cute.ceil_div(tensor_cols - chunk_col0, self.TILE_COLS), ) for stage_X in cutlass.range(tiles_X, unroll=1): self._process_block( - block_offset_Y_in_tensor, - block_id_X_in_tensor * self.STAGES_X + stage_X, + block_start_row_in_tensor, + block_id_X_in_tensor * self.NUM_TILES_X + stage_X, tensor_rows, tensor_cols, row_scales, col_scales, - block_offset_Y_in_tensor, + block_start_row_in_tensor, tensor_rows, Int32(0), # dbias is only supported for single-tensor reps mWorkspace, @@ -1014,7 +1093,7 @@ class DbiasStorage: ) # Find the next block to process block_id = block_id + block_stride - if block_id >= blocks_in_tensor: + if block_id >= units_in_tensor: job_finished = Boolean(True) # Drain every TMA store before the CTA releases its shared-memory source buffers. @@ -1071,13 +1150,13 @@ def _issue_load( @cute.jit def _process_block( self, - block_offset_Y, # Row offset of this chunk (global for single-tensor, else tensor-local) + block_start_row, # Row offset of this chunk (global for single-tensor, else tensor-local) block_id_X, # Column-chunk index within the tensor rows, # Rows of the rowwise-scale view (the group for single-tensor, else the tensor) cols, # Number of columns in this tensor - row_scales, # Rowwise scales tiled per stage, rows counted like block_offset_Y + row_scales, # Rowwise scales tiled per stage, rows counted like block_start_row col_scales, # Colwise scales tiled per stage - col_scale_row0, # Row of this chunk in the colwise-scale view + col_scale_row_start, # Row of this chunk in the colwise-scale view col_scale_rows, # Rows of the colwise-scale view dbias_row, # Row of the dbias workspace this chunk reduces into mWorkspace, # f32 partial dbias workspace (WITH_DBIAS) @@ -1096,17 +1175,17 @@ def _process_block( prod_state, cons_state, ): - """Quantize one 128x128 tile of a chunk in STAGES slices of BUFF_DIM_Y rows.""" + """Quantize NUM_TILES_Y vertically stacked tiles in one column strip of a work item.""" cfg = self.cfg _, _, tma_atom_out_row, tma_atom_out_col = atoms _, _, _, _, tXsO_row, tXgO_row, tXsO_col, tXgO_col = partitions _, _, desc_out_row, desc_out_col = descs - block_offset_X = block_id_X * self.CHUNK_DIM_X + block_offset_X = block_id_X * self.TILE_COLS # This chunk's coordinates in the tile grid (32x128 TMA boxes, not elements). - tile_id_Y = block_offset_Y // self.BUFF_DIM_Y + tile_id_Y = block_start_row // self.TILE_ROWS tile_id_X = block_id_X - col_scale_tile_Y = col_scale_row0 // self.BUFF_DIM_Y + col_scale_tile_Y = col_scale_row_start // self.TILE_ROWS # Per-chunk dbias accumulators, in the CUDA kernel's summation order: a running # column sum over the chunk's rows (colwise), or per-thread partial sums over its @@ -1136,7 +1215,7 @@ def _process_block( ) prod_state.advance() - for stage in cutlass.range_constexpr(self.STAGES): + for stage in cutlass.range_constexpr(self.NUM_TILES_Y): # Wait for at most DEPTH-1 iters on the fly, which means the the last DEPTH iter has finished # so we can reuse its SMEM output buffer # (input buffer is managed by the producer and consumer pipeline states) @@ -1158,7 +1237,7 @@ def _process_block( sO_col[(None, cons_state.index)], cute.flatten(col_scales[(None, (col_scale_tile_Y + stage, tile_id_X))]), cfg.MAX_NORM_RCP, - (col_scale_tile_Y + stage) * self.BUFF_DIM_Y, + (col_scale_tile_Y + stage) * self.TILE_ROWS, block_offset_X, col_scale_rows, cols, @@ -1166,8 +1245,8 @@ def _process_block( DTYPE=cfg.DTYPE, FP8_DTYPE=cfg.FP8_DTYPE, SWIZZLE=cfg.WITH_GEMM_SWIZZLED_SCALES, - TILE_X=self.BUFF_DIM_X, - TILE_Y=self.BUFF_DIM_Y, + TILE_X=self.TILE_COLS, + TILE_Y=self.TILE_ROWS, WITH_ACT=cfg.WITH_ACT, WITH_DACT=cfg.WITH_DACT, WITH_DBIAS=self.DBIAS_IN_COLWISE, @@ -1185,15 +1264,15 @@ def _process_block( sO_row[(None, cons_state.index)], cute.flatten(row_scales[(None, (row_tile, tile_id_X))]), cfg.MAX_NORM_RCP, - row_tile * self.BUFF_DIM_Y, + row_tile * self.TILE_ROWS, block_offset_X, rows, cols, ACTIVATION=None if self.CACHE_ACTIVATION else cfg.ACTIVATION, DTYPE=cfg.DTYPE, FP8_DTYPE=cfg.FP8_DTYPE, - TILE_X=self.BUFF_DIM_X, - TILE_Y=self.BUFF_DIM_Y, + 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, @@ -1213,7 +1292,7 @@ def _process_block( # I just freed my input SMEM buffer (stage), so the producer now can use it for writing # (stage+DEPTH) stage's data if that stage exists - if cutlass.const_expr(stage + self.PIPELINE_DEPTH < self.STAGES): + if cutlass.const_expr(stage + self.PIPELINE_DEPTH < self.NUM_TILES_Y): if warp_idx == 0: self._issue_load( mainloop_pipeline, @@ -1307,23 +1386,29 @@ def compile_cutedsl_function_from_cfg(cfg: MXFP8GroupQuantizeConfig): out_dtype = cfg.FP8_DTYPE scale_dtype = cutlass.Float8E8M0FNU - def g2d(dtype, shape=logical_shape, align=16): - return cute.runtime.make_fake_compact_tensor( - dtype, - shape, + out_col_fake = ( + cute.runtime.make_fake_compact_tensor( + out_dtype, + logical_shape, stride_order=(1, 0), memspace=cute.AddressSpace.gmem, - assumed_align=align, + assumed_align=16, ) + if cfg.COLWISE + else None + ) - def g1d(dtype, align=4): - return cute.runtime.make_fake_compact_tensor( - dtype, - (cute.sym_int32(),), - stride_order=(0,), + out_row_fake = ( + cute.runtime.make_fake_compact_tensor( + out_dtype, + logical_shape, + stride_order=(1, 0), memspace=cute.AddressSpace.gmem, - assumed_align=align, + assumed_align=16, ) + if cfg.ROWWISE + else None + ) # The kernel only takes the base address of the scale buffers (per-tensor strides # are derived from cols), so their fake shape is a flat 1D byte run. @@ -1337,27 +1422,89 @@ def g1d(dtype, align=4): # The cast-noop flag is an always-present f32 pointer instead of an optional tensor, so # that one compiled kernel serves both an absent and a present flag (noop_flag_is_set). noop_fake = cute.runtime.nullptr(Float32, mem_space=cute.AddressSpace.gmem, assumed_align=4) - act_input_fake = g2d(cfg.DTYPE) if cfg.WITH_DACT else None + act_input_fake = ( + cute.runtime.make_fake_compact_tensor( + cfg.DTYPE, + logical_shape, + stride_order=(1, 0), + memspace=cute.AddressSpace.gmem, + assumed_align=16, + ) + if cfg.WITH_DACT + else None + ) workspace_fake = ( - g2d(Float32, shape=(cute.sym_int32(), cute.sym_int32()), align=4) + cute.runtime.make_fake_compact_tensor( + Float32, + (cute.sym_int32(), cute.sym_int32()), + stride_order=(1, 0), + memspace=cute.AddressSpace.gmem, + assumed_align=4, + ) if cfg.WITH_DBIAS else None ) - from cutlass.utils import HardwareInfo # pylint: disable=import-outside-toplevel + first_dims_fake = ( + cute.runtime.make_fake_compact_tensor( + cutlass.Int64, + (cute.sym_int32(),), + stride_order=(0,), + memspace=cute.AddressSpace.gmem, + assumed_align=8, + ) + if cfg.SHAPE_REP in (VARYING_FIRST_DIM, VARYING_BOTH_DIMS) + else None + ) + + last_dims_fake = ( + cute.runtime.make_fake_compact_tensor( + cutlass.Int64, + (cute.sym_int32(),), + stride_order=(0,), + memspace=cute.AddressSpace.gmem, + assumed_align=8, + ) + if cfg.SHAPE_REP in (VARYING_LAST_DIM, VARYING_BOTH_DIMS) + else None + ) sm_count = HardwareInfo().get_device_multiprocessor_count() kernel_obj = MXFP8GroupQuantizeKernel(cfg, sm_count) return cute.compile( kernel_obj, - g2d(cfg.DTYPE), # mX - g2d(out_dtype), # mO_row - g2d(out_dtype), # mO_col - g1d(scale_dtype), # mS_row - g1d(scale_dtype), # mS_col - g1d(cutlass.Int64, align=8), # mOffsets - g1d(cutlass.Int64, align=8), # mFirstDims - g1d(cutlass.Int64, align=8), # mLastDims + cute.runtime.make_fake_compact_tensor( # mX + cfg.DTYPE, + logical_shape, + stride_order=(1, 0), + memspace=cute.AddressSpace.gmem, + assumed_align=16, + ), + out_row_fake, # mO_row + out_col_fake, # mO_col + cute.runtime.make_fake_compact_tensor( # mS_row + scale_dtype, + (cute.sym_int32(),), + stride_order=(0,), + memspace=cute.AddressSpace.gmem, + assumed_align=4, + ), + cute.runtime.make_fake_compact_tensor( # mS_col + scale_dtype, + (cute.sym_int32(),), + stride_order=(0,), + memspace=cute.AddressSpace.gmem, + assumed_align=4, + ), + cute.runtime.make_fake_compact_tensor( # mOffsets + cutlass.Int64, + (cute.sym_int32(),), + stride_order=(0,), + memspace=cute.AddressSpace.gmem, + assumed_align=8, + ), + first_dims_fake, # mFirstDims + last_dims_fake, # mLastDims tensormaps_fake, # mTensormaps noop_fake, # mNoop act_input_fake, # mActInput From bce23041794ed25434e1108dd9acb259f95bffd3 Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Fri, 2 Oct 2026 18:10:49 +0000 Subject: [PATCH 06/58] add tests Signed-off-by: Kaining Zhong --- tests/cpp/CMakeLists.txt | 1 + tests/cpp/operator/test_cast_mxfp8_grouped.cu | 59 ++++++++++++++++ tests/jax/test_mxfp8_cutedsl_backend.py | 67 +++++++++++++++++++ .../mxfp8/test_mxfp8_cutedsl_backend.py | 20 +++--- 4 files changed, 138 insertions(+), 9 deletions(-) diff --git a/tests/cpp/CMakeLists.txt b/tests/cpp/CMakeLists.txt index af074be1c78..9a1e0bae18f 100644 --- a/tests/cpp/CMakeLists.txt +++ b/tests/cpp/CMakeLists.txt @@ -57,6 +57,7 @@ if(NVTE_WITH_CUTEDSL) # Find the interpreter too so older CMake versions use its version to locate libpython. find_package(Python COMPONENTS Interpreter Development.Embed REQUIRED) foreach(test_target test_operator test_util) + target_compile_definitions(${test_target} PRIVATE NVTE_WITH_CUTEDSL) target_link_libraries(${test_target} PRIVATE "-Wl,--no-as-needed" Python::Python diff --git a/tests/cpp/operator/test_cast_mxfp8_grouped.cu b/tests/cpp/operator/test_cast_mxfp8_grouped.cu index c2309f7aa4b..13c74491a11 100644 --- a/tests/cpp/operator/test_cast_mxfp8_grouped.cu +++ b/tests/cpp/operator/test_cast_mxfp8_grouped.cu @@ -9,6 +9,11 @@ #include #include +#ifdef NVTE_WITH_CUTEDSL +#include +#include +#endif + #include #include #include "../test_common.h" @@ -1030,3 +1035,57 @@ INSTANTIATE_TEST_SUITE_P( ::testing::Values(DType::kBFloat16), ::testing::Values(DType::kFloat8E4M3)), MakeGroupedFusedCastMXFP8TestName); + +// Exercise the grouped C API with NVTE_ENABLE_CUTEDSL_BACKEND=1 as well as CUDA. +// These cases cover both column strips, a final half-width chunk, and every fused +// activation on all four shape representations against the independent CPU reference. +INSTANTIATE_TEST_SUITE_P( + OperatorTest_GroupedFusedCastMXFP8_MultiChunkActivations, + GroupedFusedCastMXFP8TestSuite, + ::testing::Combine( + ::testing::Values(ProcessingMethod::CAST_ACT, ProcessingMethod::CAST_DACT, + ProcessingMethod::CAST_DBIAS_DACT), + ::testing::Values(ActivationKind::GeLU, ActivationKind::SiLU, ActivationKind::ReLU, + ActivationKind::QGeLU, ActivationKind::SReLU), + ::testing::ValuesIn(scaling_directions), + ::testing::ValuesIn(input_config_multichunk), + ::testing::Values(DType::kBFloat16), + ::testing::Values(DType::kFloat8E4M3)), + MakeGroupedFusedCastMXFP8TestName); + +#ifdef NVTE_WITH_CUTEDSL +TEST(OperatorTest_GroupedFusedCastMXFP8, TestCuTeDSLRegistration) { + const char* enabled = std::getenv("NVTE_ENABLE_CUTEDSL_BACKEND"); + if (enabled == nullptr || enabled[0] == '0' || + getDeviceComputeCapability() < blackwellComputeCapability) { + GTEST_SKIP() << "Requires Blackwell and NVTE_ENABLE_CUTEDSL_BACKEND=1"; + } + + // Cover both column strips and a final half-width chunk. Verify registration + // after the C API call so a silent CUDA fallback cannot satisfy this test. + const std::vector first_dims = {128, 256}; + const std::vector last_dims = {128, 384}; + const std::vector offsets = {0, 128 * 128, 128 * 128 + 256 * 384}; + performTest(CAST_ONLY, &identity, VARYING_BOTH_DIMS, 2, + {1, offsets.back()}, first_dims, last_dims, offsets, + /*rowwise=*/true, /*colwise=*/true); + + ASSERT_TRUE(Py_IsInitialized()) << "CuTeDSL did not initialize embedded Python"; + const std::string key = + "cutedsl_group_mxfp8_sm" + std::to_string(getDeviceComputeCapability()) + + "_BFloat16_Float8E4M3_1_1_varying_both_dims_0_0_0_0_none"; + const PyGILState_STATE gil = PyGILState_Ensure(); + PyObject* module = PyImport_ImportModule("tvm_ffi"); + PyObject* kernel = module == nullptr ? nullptr : + PyObject_CallMethod(module, "get_global_func", "s", key.c_str()); + const bool registered = kernel != nullptr && kernel != Py_None; + if (PyErr_Occurred() != nullptr) { + PyErr_Print(); + } + Py_XDECREF(kernel); + Py_XDECREF(module); + PyGILState_Release(gil); + EXPECT_TRUE(registered) << "CuTeDSL kernel not registered for " << key + << "; the grouped C API fell back to CUDA"; +} +#endif diff --git a/tests/jax/test_mxfp8_cutedsl_backend.py b/tests/jax/test_mxfp8_cutedsl_backend.py index 4671ea06869..a74baf93ec1 100644 --- a/tests/jax/test_mxfp8_cutedsl_backend.py +++ b/tests/jax/test_mxfp8_cutedsl_backend.py @@ -283,3 +283,70 @@ def test_dtypes(method, act_type, act_desc, fp8_dtype, in_dtype): in_dtype, fp8_dtype, ) + + +# JAX's V2 grouped quantize emits GEMM-swizzled scales and supports common last +# dimensions, represented by device row counts even for equal-size groups. +# Varying last dimensions and grouped fused activations are covered +# through the C API / PyTorch, since JAX has no wrappers for those operations. +# (name, shape representation, input shape, optional device-resident row counts) +GROUP_CASES = [ + ("single_member", "varying_first_dim", (1, 128, 128), None), + ("same_both_multichunk", "varying_first_dim", (3, 384, 384), None), + ("varying_first", "varying_first_dim", (768, 256), (128, 384, 256)), + ("varying_first_multichunk", "varying_first_dim", (1024, 384), (128, 256, 384, 256)), + ("varying_first_empty", "varying_first_dim", (512, 256), (128, 0, 384)), +] +GROUP_Q_LAYOUTS = [QuantizeLayout.ROWWISE, QuantizeLayout.COLWISE, QuantizeLayout.ROWWISE_COLWISE] + + +def get_group_cfg_key(shape_rep, in_dtype, fp8_dtype, q_layout): + """Registry key for the grouped V2 kernel (swizzled, without dbias or activation).""" + major, minor = device_compute_capability() + return ( + f"cutedsl_group_mxfp8_sm{major * 10 + minor}_{DTYPE_TO_STR[in_dtype]}_" + f"{FP8_TO_KEY[fp8_dtype]}_{int(q_layout.has_rowwise)}_{int(q_layout.has_colwise)}_" + f"{shape_rep}_1_0_0_0_none" + ) + + +@pytest.mark.parametrize("case", GROUP_CASES, ids=lambda c: c[0]) +@pytest.mark.parametrize("q_layout", GROUP_Q_LAYOUTS, ids=lambda l: l.name.lower()) +@pytest.mark.parametrize("in_dtype", IN_DTYPES, ids=get_dtype_id) +@pytest.mark.parametrize("fp8_dtype", FP8_DTYPES, ids=get_fp8_id) +def test_group_cast_only(case, q_layout, in_dtype, fp8_dtype): + """Compare every grouped data/scale byte, including swizzled scales, across backends.""" + name, shape_rep, shape, row_counts = case + x, _ = generate_inputs(int(np.prod(shape[:-1])), shape[-1], in_dtype) + x = x.reshape(shape) + group_sizes = None if row_counts is None else jnp.asarray(row_counts, dtype=jnp.int32) + n_groups = shape[0] if row_counts is None else len(row_counts) + quantizer = QuantizerFactory.create( + scaling_mode=ScalingMode.MXFP8_1D_SCALING, + q_dtype=fp8_dtype, + q_layout=q_layout, + n_groups=n_groups, + ) + + def run(): + out = tex.grouped_quantize(x, quantizer=quantizer, group_sizes=group_sizes) + # Host materialization completes the FFI call before changing the backend flag. + return extract_quantized_output(out, None)[0] + + set_cutedsl_backend(False) + cuda_output = run() + set_cutedsl_backend(True) + try: + cutedsl_output = run() + finally: + set_cutedsl_backend(False) + + key = get_group_cfg_key(shape_rep, in_dtype, fp8_dtype, q_layout) + assert ( + tvm_ffi.get_global_func(key, allow_missing=True) is not None + ), f"CuTeDSL kernel not registered for {key}; grouped quantization fell back to CUDA" + for part, cuda_bytes in cuda_output.items(): + assert np.array_equal(cutedsl_output[part], cuda_bytes), ( + f"group/{name}/{get_dtype_id(in_dtype)}/{get_fp8_id(fp8_dtype)}/{q_layout.name}: " + f"{part} differ between backends" + ) diff --git a/tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py b/tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py index c3ab0bfa8c2..0847421383e 100644 --- a/tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py +++ b/tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py @@ -381,14 +381,19 @@ def test_dtypes(swizzled, method, act, fp8_dtype, in_dtype): # (name, shape representation in the config key, per-member shapes) GROUP_CASES = [ ("same_both", "same_both_dims", [(256, 512)] * 3), + ("single_member", "same_both_dims", [(384, 384)]), # N is 32- but not 128-divisible, so the rowwise and colwise scales carry zeroed padding. ("same_both_n96", "same_both_dims", [(128, 96)] * 2), ("varying_first", "varying_first_dim", [(128, 256), (384, 256), (256, 256)]), ("varying_first_n160", "varying_first_dim", [(128, 160), (256, 160)]), - # N ends in a partial 32-element scale block. - ("varying_first_n144", "varying_first_dim", [(128, 144), (384, 144)]), + # Partial 32-element blocks (e.g. N=144) are covered by the C++ grouped tests; + # the PyTorch MXFP8 quantizer requires dimensions divisible by 32. ("varying_last", "varying_last_dim", [(256, 128), (256, 384), (256, 256)]), ("varying_both", "varying_both_dims", [(128, 256), (256, 128), (384, 512)]), + # Multiple chunks in both dimensions, including a final half-width 128-column strip. + ("varying_both_multichunk", "varying_both_dims", [(128, 128), (256, 384), (384, 640)]), + ("varying_first_empty", "varying_first_dim", [(128, 256), (0, 256), (384, 256)]), + ("varying_both_empty", "varying_both_dims", [(128, 128), (0, 256), (256, 384)]), ] SINGLE_TENSOR_GROUP_CASES = [ c for c in GROUP_CASES if c[1] in ("same_both_dims", "varying_first_dim") @@ -494,7 +499,9 @@ def run_group_test_case( cutedsl_output[name], cuda_bytes ), f"{tag}: {name} differ between backends" if dbias: - assert torch.equal(dbias_cutedsl, dbias_cuda), f"{tag}: dbias differs between backends" + assert torch.equal( + dbias_cutedsl.view(torch.uint8), dbias_cuda.view(torch.uint8) + ), f"{tag}: dbias differs between backends" @pytest.mark.parametrize("case", GROUP_CASES, ids=get_group_case_id) @@ -510,11 +517,6 @@ def test_group_cast_only(fp8_dtype, in_dtype, block_size, case): @pytest.mark.parametrize("block_size", BLOCK_SIZES, ids=get_block_id) def test_group_swizzled(block_size, case): _, shape_rep, members = case - if block_size[0] != 1 and shape_rep in ("varying_last_dim", "varying_both_dims"): - pytest.skip( - "The CUDA kernel's GEMM-swizzled colwise scale index double-counts the tensor base" - " for varying last dims; the CuTeDSL backend leaves these configs to it." - ) run_group_test_case( members, shape_rep, block_size, torch.bfloat16, tex.DType.kFloat8E4M3, swizzled=True ) @@ -522,7 +524,7 @@ def test_group_swizzled(block_size, case): @pytest.mark.parametrize("case", SINGLE_TENSOR_GROUP_CASES, ids=get_group_case_id) @pytest.mark.parametrize("block_size", BLOCK_SIZES, ids=get_block_id) -@pytest.mark.parametrize("in_dtype", [torch.bfloat16, torch.float32], ids=get_dtype_id) +@pytest.mark.parametrize("in_dtype", IN_DTYPES, ids=get_dtype_id) @pytest.mark.parametrize("swizzled", SWIZZLE_MODES, ids=get_swizzle_id) def test_group_dbias(swizzled, in_dtype, block_size, case): _, shape_rep, members = case From d0fce6b17cab631df3d3b1f57e9d5156ba7cb04b Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Fri, 2 Oct 2026 18:14:29 +0000 Subject: [PATCH 07/58] fix Signed-off-by: Kaining Zhong --- .../cast/mxfp8/group_quantize_mxfp8.py | 216 +++++++++--------- .../cast/mxfp8/group_quantize_mxfp8.cuh | 3 +- .../mxfp8/group_quantize_mxfp8_cutedsl.cuh | 52 ++--- 3 files changed, 134 insertions(+), 137 deletions(-) diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8.py index b86476c5824..d4c61f5730e 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8.py @@ -8,16 +8,16 @@ management and per-tensor scale addressing mirror the CUDA kernel one-for-one: * `is_single_tensor` reps (SAME_BOTH_DIMS, VARYING_FIRST_DIM) launch ONE CTA per - 128x128 chunk and address the group through ONE static TMA descriptor with - global block offsets -- the CUDA `tensor_map_*_static` "direct mapper" path. + 128x128 job and address the group through ONE static TMA descriptor with + global job offsets -- the CUDA `tensor_map_*_static` "direct mapper" path. For SAME_BOTH_DIMS the CUDA grid is linearized per tensor (X, Y-in-tensor, tensor) while this kernel linearizes it flat over the stacked rows; both require every member's row count to be a multiple of 128, and under - that precondition the two decode to the identical (block_start_row, block_id_X) - for every block index. + that precondition the two decode to the identical (job_start_row, job_id_X) + for every job index. * the other reps launch grid=(workers_per_tensor, num_tensors) and bind tensor_id to blockIdx.y, so a CTA grid-strides only within its own tensor and - never re-resolves which tensor a chunk belongs to. They get per-tensor + never re-resolves which tensor a job belongs to. They get per-tensor descriptors written by a prologue kernel (the CuTeDSL analog of update_tma_descriptors filling g_tensor_maps) and acquired with a tensormap proxy fence. @@ -31,7 +31,7 @@ Mechanics that provably yield the same bytes may differ: the mbarrier pipeline is expressed with PipelineTmaAsync instead of hand-rolled mbarriers. As in CUDA, the -scales of out-of-bounds columns in a chunk (the scale-row padding) are written as 0. +scales of out-of-bounds columns in a job (the scale-row padding) are written as 0. Scope: everything group_quantize_mxfp8.cuh covers except 2D block scaling -- the cast-noop flag, fused activation (IS_ACT) and activation derivative (IS_DACT), dbias, @@ -206,7 +206,7 @@ class MXFP8GroupQuantizeKernel: """Grouped MXFP8 quantize mirroring group_quantize_mxfp8_kernel's strategy.""" # Target persistent CTA count per SM for sizing the grid - STATIC_PERSISTENT_WORK_UNITS_PER_SM = 24 + STATIC_PERSISTENT_WORKERS_PER_SM = 24 # The shape of one pipeline stage processed by a CTA TILE_ROWS = 32 TILE_COLS = 128 @@ -379,49 +379,43 @@ def __call__( if cutlass.const_expr(cfg.IS_SINGLE_TENSOR): # How many CTAs does the grouped tensor have in both directions - work_blocks_Y = cute.ceil_div( - Int32(first_logical_dim), (self.TILE_ROWS * self.NUM_TILES_Y) - ) - work_blocks_X = cute.ceil_div( - Int32(last_logical_dim), self.TILE_COLS * self.NUM_TILES_X - ) + jobs_Y = cute.ceil_div(Int32(first_logical_dim), (self.TILE_ROWS * self.NUM_TILES_Y)) + jobs_X = cute.ceil_div(Int32(last_logical_dim), self.TILE_COLS * self.NUM_TILES_X) # Flatten it to an 1D grid - grid = [work_blocks_X * work_blocks_Y, 1, 1] + grid = [jobs_X * jobs_Y, 1, 1] else: # A placeholder for the kernel signature only; we won't use it in non-single tensor cases - work_blocks_X = None - # Estimate how many CTA work unit we need to do, where one work unit is NUM_TILES_X * NUM_TILES_Y tiles + jobs_X = None + # Estimate the total jobs across the group; each job has NUM_TILES_X * NUM_TILES_Y tiles if cutlass.const_expr(cfg.SHAPE_REP == VARYING_BOTH_DIMS): # Note: when VARYING_BOTH_DIMS, the first_logical_dim must be 1 - estimated_work_units = cute.ceil_div( + estimated_jobs = cute.ceil_div( Int32(first_logical_dim) * Int32(last_logical_dim), self.ELTS_PER_CTA ) elif cutlass.const_expr(cfg.SHAPE_REP == VARYING_LAST_DIM): # Same as VARYING_BOTH_DIMS but we divide 128 before multiplying to avoid overflowing Int32 # because the first logical dimension is always 128-aligned when not VARYING_BOTH_DIMS # so they are equivalent - estimated_work_units = cute.ceil_div( + estimated_jobs = cute.ceil_div( (Int32(first_logical_dim) // 128) * Int32(last_logical_dim), self.ELTS_PER_CTA // 128, ) else: raise ValueError(f"unexpected shape representation {cfg.SHAPE_REP!r}") - # All SMs can do SM_COUNT * STATIC_PERSISTENT_WORK_UNITS_PER_SM units of work, which are distributed evenly - # to all tensors in the group. - # If there are more tensors than work units, we will launch one CTA per tensor since we have enough parallelism here + # Divide the persistent worker budget evenly across tensors, with at least + # one worker per tensor. Each worker may process several jobs. requested_CTAs_per_tensor = cutlass.max( Int32(1), - Int32(self.SM_COUNT * self.STATIC_PERSISTENT_WORK_UNITS_PER_SM) - // Int32(num_tensors), + Int32(self.SM_COUNT * self.STATIC_PERSISTENT_WORKERS_PER_SM) // Int32(num_tensors), ) - # In average how many work units per tensor (only an average, the actual work units per tensor may vary) - average_work_units_per_tensor = cutlass.max( - Int32(1), cute.ceil_div(estimated_work_units, Int32(num_tensors)) + # In average how many jobs per tensor (only an average, the actual jobs per tensor may vary) + average_jobs_per_tensor = cutlass.max( + Int32(1), cute.ceil_div(estimated_jobs, Int32(num_tensors)) ) - # Don't launch more CTAs than the average work units per tensor in case - # STATIC_PERSISTENT_WORK_UNITS_PER_SM causes redundancy - CTAs_per_tensor = cutlass.min(requested_CTAs_per_tensor, average_work_units_per_tensor) + # Don't launch more CTAs than the average jobs per tensor in case + # STATIC_PERSISTENT_WORKERS_PER_SM causes redundancy + CTAs_per_tensor = cutlass.min(requested_CTAs_per_tensor, average_jobs_per_tensor) grid = [CTAs_per_tensor, Int32(num_tensors), 1] # Only the multi-tensor representations need per-tensor descriptors. @@ -455,7 +449,7 @@ def __call__( first_logical_dim, last_logical_dim, num_tensors, - work_blocks_X, + jobs_X, mX.element_type, tma_atom_x, tma_src, @@ -720,7 +714,7 @@ def kernel( first_logical_dim, last_logical_dim, num_tensors, - work_blocks_X, + jobs_X, dtype: cutlass.Constexpr[Type[cutlass.Numeric]], tma_atom_x, tma_src, @@ -746,7 +740,7 @@ def kernel( first_logical_dim, last_logical_dim, num_tensors, - work_blocks_X, + jobs_X, dtype, tma_atom_x, tma_src, @@ -770,7 +764,7 @@ def _kernel_main( first_logical_dim, last_logical_dim, num_tensors, - work_units_X, + jobs_X, dtype: cutlass.Constexpr[Type[cutlass.Numeric]], tma_atom_x, tma_src, @@ -857,37 +851,37 @@ def _kernel_main( # If the CTA has work to do has_work = Boolean(True) - # Metadata of the tensor that owns this block + # Metadata of the tensor that owns this job tensor_rows = Int32(0) tensor_cols = Int32(0) # Element offset of this tensor within the group: Int64 (CUDA uses size_t), since a # group can exceed 2^31 elements even when every individual extent is small. tensor_base = Int64(0) - # Block's offset and id in this individual tensor / global single tensor - block_start_row = Int32(0) - block_id_X = Int32(0) + # Job's starting row and id in this individual tensor / global single tensor + job_start_row = Int32(0) + job_id_X = Int32(0) - first_block_id = Int32(0) - units_in_tensor = Int32(1) - block_stride = Int32(1) - units_X_in_tensor = Int32(1) + first_job_id = Int32(0) + jobs_in_tensor = Int32(1) + job_stride = Int32(1) + jobs_X_in_tensor = Int32(1) if cutlass.const_expr(cfg.IS_SINGLE_TENSOR): - # grid = [work_units_X * work_units_Y, 1, 1] - block_id_Y = Int32(bidx) // work_units_X - block_id_X = Int32(bidx) % work_units_X + # grid = [jobs_X * jobs_Y, 1, 1] + job_id_Y = Int32(bidx) // jobs_X + job_id_X = Int32(bidx) % jobs_X # View the grouped tensor as a single tensor of shape (first_logical_dim, last_logical_dim) tensor_rows = Int32(first_logical_dim) tensor_cols = Int32(last_logical_dim) - # Which row does this block start from - block_start_row = block_id_Y * (self.TILE_ROWS * self.NUM_TILES_Y) + # Which row does this job start from + job_start_row = job_id_Y * (self.TILE_ROWS * self.NUM_TILES_Y) if cutlass.const_expr(cfg.SHAPE_REP == VARYING_FIRST_DIM): total_elts = Int64(mOffsets[mOffsets.shape[0] - 1]) # If the starting element is already beyond the last element of the group, this CTA can stop - if Int64(block_start_row) * Int64(last_logical_dim) >= total_elts: + if Int64(job_start_row) * Int64(last_logical_dim) >= total_elts: has_work = Boolean(False) # When SAME_BOTH_DIM, M is always divisible by 128, which is exactly TILE_ROWS * NUM_TILES_Y, - # so no need to check for the last block's starting row being beyond the last row of the tensor. + # so no need to check for the last job's starting row being beyond the last row of the tensor. else: # grid = [workers_per_tensor, Int32(num_tensors), 1] tensor_id = Int32(bidy) @@ -897,17 +891,17 @@ def _kernel_main( tensor_cols = Int32(meta[1]) tensor_base = Int64(meta[2]) if tensor_rows > 0 and tensor_cols > 0: - # How many work units does this tensor have in both directions - units_X_in_tensor = cute.ceil_div(tensor_cols, (self.TILE_COLS * self.NUM_TILES_X)) - units_Y_in_tensor = cute.ceil_div(tensor_rows, (self.TILE_ROWS * self.NUM_TILES_Y)) - # How many work units does this tensor have - units_in_tensor = units_X_in_tensor * units_Y_in_tensor - # Which work unit (1D index) does this CTA start from, which is also their worker ID - first_block_id = Int32(bidx) + # How many jobs does this tensor have in both directions + jobs_X_in_tensor = cute.ceil_div(tensor_cols, (self.TILE_COLS * self.NUM_TILES_X)) + jobs_Y_in_tensor = cute.ceil_div(tensor_rows, (self.TILE_ROWS * self.NUM_TILES_Y)) + # How many jobs does this tensor have + jobs_in_tensor = jobs_X_in_tensor * jobs_Y_in_tensor + # Which job (1D index) does this CTA start from, which is also their worker ID + first_job_id = Int32(bidx) # gdx is how many CTAs are assigned to this tensor - block_stride = Int32(gdx) - # If my first work unit is already beyond the tensor's last work unit, I have no work to do - if first_block_id >= units_in_tensor: + job_stride = Int32(gdx) + # If my first job is already beyond the tensor's last job, I have no work to do + if first_job_id >= jobs_in_tensor: has_work = Boolean(False) else: # This tensor is empty, so this CTA has no work to do @@ -917,7 +911,7 @@ def _kernel_main( atoms = (tma_atom_x, tma_atom_act, tma_atom_out_row, tma_atom_out_col) if has_work: - # For single tensor case, each CTA only processes one work unit + # For single tensor case, each CTA only processes one job if cutlass.const_expr(cfg.IS_SINGLE_TENSOR): # For single tensor case we don't use tensor descriptors descs = (None, None, None, None) @@ -929,33 +923,35 @@ def _kernel_main( row_scales = self._rowwise_scales(mS_row, Int64(0), tensor_rows, tensor_cols) col_scales = None - col_scale_row_start = block_start_row + col_scale_row_start = job_start_row col_scale_rows = tensor_rows if cutlass.const_expr(cfg.COLWISE): col_scale_base = Int64(0) if cutlass.const_expr(cfg.WITH_GEMM_SWIZZLED_SCALES): # Colwise swizzled scale indices restart at each member and depend # on its rows (process_colwise_stage), so address the member that - # owns this chunk. + # owns this job. member_rows = Int32(0) - member_row0 = Int32(0) + member_row_start = Int32(0) if cutlass.const_expr(cfg.SHAPE_REP == SAME_BOTH_DIMS): member_rows = tensor_rows // Int32(num_tensors) - member_row0 = block_start_row // member_rows * member_rows + member_row_start = job_start_row // member_rows * member_rows else: member_id = self._find_tensor_from_offsets( mOffsets, num_tensors, - Int64(block_start_row) * Int64(tensor_cols), + Int64(job_start_row) * Int64(tensor_cols), ) member_rows = Int32(mFirstDims[member_id]) - member_row0 = Int32(Int64(mOffsets[member_id]) // Int64(tensor_cols)) + member_row_start = Int32( + Int64(mOffsets[member_id]) // Int64(tensor_cols) + ) col_scale_base = ( - Int64(member_row0) + Int64(member_row_start) * Int64(cute.round_up(tensor_cols, 128)) // MXFP8_BLOCK_SCALING_SIZE ) - col_scale_row_start = block_start_row - member_row0 + col_scale_row_start = job_start_row - member_row_start col_scale_rows = member_rows col_scales = self._colwise_scales( mS_col, col_scale_base, col_scale_rows, tensor_cols @@ -963,16 +959,16 @@ def _kernel_main( cute.arch.sync_threads() - self._process_block( - block_start_row, - block_id_X, + self._process_job_strip( + job_start_row, + job_id_X, tensor_rows, tensor_cols, row_scales, col_scales, col_scale_row_start, col_scale_rows, - block_start_row // (self.TILE_ROWS * self.NUM_TILES_Y), + job_start_row // (self.TILE_ROWS * self.NUM_TILES_Y), mWorkspace, sDbias, descs, @@ -990,7 +986,7 @@ def _kernel_main( cons_state, ) else: - # For non-single tensor case, we use persistent kernel so each CTA keeps processing work units until none is left + # For non-single tensor case, we use persistent kernel so each CTA keeps processing jobs until none is left desc_x = tmap.get_tensormap_ptr(mTensormaps[(tensor_id, 0, None)].iterator) desc_out_row = tmap.get_tensormap_ptr(mTensormaps[(tensor_id, 1, None)].iterator) desc_out_col = tmap.get_tensormap_ptr(mTensormaps[(tensor_id, 2, None)].iterator) @@ -1017,27 +1013,27 @@ def _kernel_main( if cutlass.const_expr(cfg.COLWISE): col_scales = self._colwise_scales(mS_col, scale_base, tensor_rows, tensor_cols) - # Make sure all threads see the updated descriptors and scales before processing any blocks. + # Make sure all threads see the updated descriptors and scales before processing any jobs. cute.arch.sync_threads() - # Grid-stride over this tensor's own chunks; the descriptors never change. - block_id = first_block_id + # Grid-stride over this tensor's own jobs; the descriptors never change. + job_id = first_job_id job_finished = Boolean(False) while not job_finished: - block_id_Y_in_tensor = block_id // units_X_in_tensor - block_id_X_in_tensor = block_id % units_X_in_tensor - block_start_row_in_tensor = block_id_Y_in_tensor * ( + job_id_Y_in_tensor = job_id // jobs_X_in_tensor + job_id_X_in_tensor = job_id % jobs_X_in_tensor + job_start_row_in_tensor = job_id_Y_in_tensor * ( self.TILE_ROWS * self.NUM_TILES_Y ) if cutlass.const_expr(self.NUM_TILES_X == 1): - self._process_block( - block_start_row_in_tensor, - block_id_X_in_tensor, + self._process_job_strip( + job_start_row_in_tensor, + job_id_X_in_tensor, tensor_rows, tensor_cols, row_scales, col_scales, - block_start_row_in_tensor, + job_start_row_in_tensor, tensor_rows, Int32(0), # dbias is only supported for single-tensor reps mWorkspace, @@ -1057,22 +1053,22 @@ def _kernel_main( cons_state, ) else: - # The chunk's column tiles in order, stopping at the tensor's last + # The job's column tiles in order, stopping at the tensor's last # column like the CUDA kernel's stages_X = DIVUP(chunk_cols, TILE_DIM_X). - chunk_col0 = block_id_X_in_tensor * (self.TILE_COLS * self.NUM_TILES_X) + job_start_col = job_id_X_in_tensor * (self.TILE_COLS * self.NUM_TILES_X) tiles_X = cutlass.min( Int32(self.NUM_TILES_X), - cute.ceil_div(tensor_cols - chunk_col0, self.TILE_COLS), + cute.ceil_div(tensor_cols - job_start_col, self.TILE_COLS), ) for stage_X in cutlass.range(tiles_X, unroll=1): - self._process_block( - block_start_row_in_tensor, - block_id_X_in_tensor * self.NUM_TILES_X + stage_X, + self._process_job_strip( + job_start_row_in_tensor, + job_id_X_in_tensor * self.NUM_TILES_X + stage_X, tensor_rows, tensor_cols, row_scales, col_scales, - block_start_row_in_tensor, + job_start_row_in_tensor, tensor_rows, Int32(0), # dbias is only supported for single-tensor reps mWorkspace, @@ -1091,9 +1087,9 @@ def _kernel_main( prod_state, cons_state, ) - # Find the next block to process - block_id = block_id + block_stride - if block_id >= units_in_tensor: + # Find the next job to process + job_id = job_id + job_stride + if job_id >= jobs_in_tensor: job_finished = Boolean(True) # Drain every TMA store before the CTA releases its shared-memory source buffers. @@ -1148,17 +1144,17 @@ def _issue_load( pipeline_obj.producer_commit(prod_state) @cute.jit - def _process_block( + def _process_job_strip( self, - block_start_row, # Row offset of this chunk (global for single-tensor, else tensor-local) - block_id_X, # Column-chunk index within the tensor + job_start_row, # Row offset of this job (global for single-tensor, else tensor-local) + column_tile_id, # Column-tile index within the tensor rows, # Rows of the rowwise-scale view (the group for single-tensor, else the tensor) cols, # Number of columns in this tensor - row_scales, # Rowwise scales tiled per stage, rows counted like block_start_row + row_scales, # Rowwise scales tiled per stage, rows counted like job_start_row col_scales, # Colwise scales tiled per stage - col_scale_row_start, # Row of this chunk in the colwise-scale view + col_scale_row_start, # Row of this job in the colwise-scale view col_scale_rows, # Rows of the colwise-scale view - dbias_row, # Row of the dbias workspace this chunk reduces into + dbias_row, # Row of the dbias workspace this job reduces into mWorkspace, # f32 partial dbias workspace (WITH_DBIAS) sDbias, # SMEM buffer for the rowwise dbias reduction (rowwise-only dbias) descs, # Per-tensor descriptors (x, act, out_row, out_col), None if single-tensor @@ -1175,20 +1171,20 @@ def _process_block( prod_state, cons_state, ): - """Quantize NUM_TILES_Y vertically stacked tiles in one column strip of a work item.""" + """Quantize NUM_TILES_Y vertically stacked tiles in one column strip of a job.""" cfg = self.cfg _, _, tma_atom_out_row, tma_atom_out_col = atoms _, _, _, _, tXsO_row, tXgO_row, tXsO_col, tXgO_col = partitions _, _, desc_out_row, desc_out_col = descs - block_offset_X = block_id_X * self.TILE_COLS + job_start_col = column_tile_id * self.TILE_COLS - # This chunk's coordinates in the tile grid (32x128 TMA boxes, not elements). - tile_id_Y = block_start_row // self.TILE_ROWS - tile_id_X = block_id_X + # This job's coordinates in the tile grid (32x128 TMA boxes, not elements). + tile_id_Y = job_start_row // self.TILE_ROWS + tile_id_X = column_tile_id col_scale_tile_Y = col_scale_row_start // self.TILE_ROWS - # Per-chunk dbias accumulators, in the CUDA kernel's summation order: a running - # column sum over the chunk's rows (colwise), or per-thread partial sums over its + # Per-job dbias accumulators, in the CUDA kernel's summation order: a running + # column sum over the job's rows (colwise), or per-thread partial sums over its # stages that the whole CTA reduces afterwards (rowwise-only). dbias_col = Float32(0.0) dbias_row_acc = None @@ -1238,7 +1234,7 @@ def _process_block( cute.flatten(col_scales[(None, (col_scale_tile_Y + stage, tile_id_X))]), cfg.MAX_NORM_RCP, (col_scale_tile_Y + stage) * self.TILE_ROWS, - block_offset_X, + job_start_col, col_scale_rows, cols, ACTIVATION=cfg.ACTIVATION, @@ -1265,7 +1261,7 @@ def _process_block( cute.flatten(row_scales[(None, (row_tile, tile_id_X))]), cfg.MAX_NORM_RCP, row_tile * self.TILE_ROWS, - block_offset_X, + job_start_col, rows, cols, ACTIVATION=None if self.CACHE_ACTIVATION else cfg.ACTIVATION, @@ -1335,8 +1331,8 @@ def _process_block( if cutlass.const_expr(cfg.WITH_DBIAS): if cutlass.const_expr(self.DBIAS_IN_ROWWISE): dbias_col = self._reduce_rowwise_dbias(sDbias, tidx, dbias_row_acc) - # One partial-dbias row per chunk, as in the CUDA kernel's dbias_workspace. - dbias_x = block_offset_X + tidx + # One partial-dbias row per job, as in the CUDA kernel's dbias_workspace. + dbias_x = job_start_col + tidx if dbias_x < cols: mWorkspace[(dbias_row, dbias_x)] = dbias_col @@ -1365,7 +1361,7 @@ def _reduce_rowwise_dbias(self, sDbias, tidx, dbias_row_acc): dbias = Float32(0.0) for i in cutlass.range_constexpr(self.THREADS_Y): dbias += sDbias[(i, tidx)] - # The buffer is rewritten by the next chunk. + # The buffer is rewritten by the next job. cute.arch.sync_threads() return dbias diff --git a/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh index bd8c3052f60..c39624889ca 100644 --- a/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh @@ -882,7 +882,8 @@ __global__ void __launch_bounds__(CastTraits::THREADS_PER_CHUNK) group_quantize_ DIVUP_TO_MULTIPLE(DIVUP(cols, static_cast(SCALE_DIM_X)), scale_alignment_X_rowwise); const size_t scale_stride_colwise = DIVUP_TO_MULTIPLE(cols, scale_alignment_X_colwise); - const size_t tensor_base_for_scales = is_single_tensor ? tensor_start_offset : tensor_base; + // Non-single-tensor scale pointers already include the member offset. + const size_t tensor_base_for_scales = is_single_tensor ? tensor_start_offset : 0; e8m0_t *const scales_rowwise = scales_rowwise_ptr + (is_single_tensor ? 0 : tensor_base / SCALE_DIM_X); e8m0_t *const scales_colwise = diff --git a/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8_cutedsl.cuh b/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8_cutedsl.cuh index c1b1aae0edd..71f660f3748 100644 --- a/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8_cutedsl.cuh +++ b/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8_cutedsl.cuh @@ -137,10 +137,8 @@ constexpr size_t kMaxGroupTensors = struct alignas(128) GroupDescriptorWorkspace { alignas(128) int64_t tensor_maps[kMaxGroupTensors][kGroupTensorMapSlots][kInt64PerTensorMap]; - // Stand-in for the offsets / first_dims / last_dims arrays a given shape representation - // does not carry: the kernel takes all three unconditionally but only dereferences the - // ones its representation uses, so the contents are never read. Sized num_tensors + 1 - // for the CSR offsets array, the longest of the three. + // Stand-in for the unused offsets array in SAME_BOTH_DIMS. The kernel takes + // offsets unconditionally but does not read them for this representation. int64_t unused_dims[kMaxGroupTensors + 1]; }; @@ -222,14 +220,6 @@ inline bool mxfp8_group_quantize_cutedsl(const MXFP8GroupQuantConfig &config, const int32_t device_index = transformer_engine::cuda::current_device(); GroupDescriptorWorkspace *const workspace = group_descriptor_workspace_ptr(); - // Both output directions are handed to the kernel unconditionally: the compiled - // signature has no optional outputs, and building a TMA descriptor needs a real - // address for each. The disabled direction is never read or written, so it points at - // the enabled one instead of at a buffer that would have to be allocated. - const SimpleTensor &data_row = - config.rowwise ? output_tensor->data : output_tensor->columnwise_data; - const SimpleTensor &data_col = - config.colwise ? output_tensor->columnwise_data : output_tensor->data; const SimpleTensor &scale_row = config.rowwise ? output_tensor->scale_inv : output_tensor->columnwise_scale_inv; const SimpleTensor &scale_col = @@ -240,10 +230,17 @@ inline bool mxfp8_group_quantize_cutedsl(const MXFP8GroupQuantConfig &config, DLTensorWrapper mX( make_basic_tensor(input_tensor->data.dptr, input_tensor->dtype(), logical_shape), true, device_index); - DLTensorWrapper mO_row(make_basic_tensor(data_row.dptr, data_row.dtype, logical_shape), true, - device_index); - DLTensorWrapper mO_col(make_basic_tensor(data_col.dptr, data_col.dtype, logical_shape), true, - device_index); + DLTensorWrapper mO_row, mO_col; + if (config.rowwise) { + mO_row = DLTensorWrapper( + make_basic_tensor(output_tensor->data.dptr, output_tensor->data.dtype, logical_shape), true, + device_index); + } + if (config.colwise) { + mO_col = DLTensorWrapper(make_basic_tensor(output_tensor->columnwise_data.dptr, + output_tensor->columnwise_data.dtype, logical_shape), + true, device_index); + } // The kernel only takes the base address of the scale buffers (per-tensor bases and // strides are derived from the member shapes), so these stay 1D. @@ -258,8 +255,19 @@ inline bool mxfp8_group_quantize_cutedsl(const MXFP8GroupQuantConfig &config, return DLTensorWrapper(make_basic_tensor(dptr, DType::kInt64, {numel}), false, device_index); }; DLTensorWrapper mOffsets = dims_or_unused(output_tensor->tensor_offsets, num_tensors + 1); - DLTensorWrapper mFirstDims = dims_or_unused(output_tensor->first_dims, num_tensors); - DLTensorWrapper mLastDims = dims_or_unused(output_tensor->last_dims, num_tensors); + DLTensorWrapper mFirstDims, mLastDims; + if (config.shape_rep == ShapeRepresentation::VARYING_FIRST_DIM || + config.shape_rep == ShapeRepresentation::VARYING_BOTH_DIMS) { + mFirstDims = DLTensorWrapper( + make_basic_tensor(output_tensor->first_dims.dptr, DType::kInt64, {num_tensors}), + false, device_index); + } + if (config.shape_rep == ShapeRepresentation::VARYING_LAST_DIM || + config.shape_rep == ShapeRepresentation::VARYING_BOTH_DIMS) { + mLastDims = DLTensorWrapper( + make_basic_tensor(output_tensor->last_dims.dptr, DType::kInt64, {num_tensors}), + false, device_index); + } // The kernel reads num_tensors off this tensor's leading extent, so it must be exactly // the group size even on the single-tensor path that leaves the descriptors untouched. @@ -374,14 +382,6 @@ bool mxfp8_group_quantize_cutedsl(const GroupedTensor *input_tensor, return false; } const bool swizzled = output_tensor->with_gemm_swizzled_scales; - if (swizzled && colwise && !is_single_tensor) { - // For these representations the CUDA kernel adds the tensor base to the colwise - // swizzled scale index twice, so leave them to it rather than reproduce that. - maybe_warn_cutedsl_not_chosen( - "GEMM-swizzled colwise scales are only supported for a common last dimension."); - return false; - } - // Sanity checks, mirroring mxfp8::group_quantize checkCuDriverContext(stream); CheckNoopTensor(*noop_tensor, "cast_noop"); From 736efe5eaeb620239f5698562b503532cdbffbe0 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 2 Oct 2026 18:16:10 +0000 Subject: [PATCH 08/58] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../common/cast/mxfp8/group_quantize_mxfp8_cutedsl.cuh | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8_cutedsl.cuh b/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8_cutedsl.cuh index 71f660f3748..11d6500ff71 100644 --- a/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8_cutedsl.cuh +++ b/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8_cutedsl.cuh @@ -259,14 +259,14 @@ inline bool mxfp8_group_quantize_cutedsl(const MXFP8GroupQuantConfig &config, if (config.shape_rep == ShapeRepresentation::VARYING_FIRST_DIM || config.shape_rep == ShapeRepresentation::VARYING_BOTH_DIMS) { mFirstDims = DLTensorWrapper( - make_basic_tensor(output_tensor->first_dims.dptr, DType::kInt64, {num_tensors}), - false, device_index); + make_basic_tensor(output_tensor->first_dims.dptr, DType::kInt64, {num_tensors}), false, + device_index); } if (config.shape_rep == ShapeRepresentation::VARYING_LAST_DIM || config.shape_rep == ShapeRepresentation::VARYING_BOTH_DIMS) { mLastDims = DLTensorWrapper( - make_basic_tensor(output_tensor->last_dims.dptr, DType::kInt64, {num_tensors}), - false, device_index); + make_basic_tensor(output_tensor->last_dims.dptr, DType::kInt64, {num_tensors}), false, + device_index); } // The kernel reads num_tensors off this tensor's leading extent, so it must be exactly From 3d1a95cee143fefd29a662465f9bd0e95b4de0a4 Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Fri, 2 Oct 2026 18:24:51 +0000 Subject: [PATCH 09/58] nit Signed-off-by: Kaining Zhong --- .../cast/mxfp8/group_quantize_mxfp8_cutedsl.cuh | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8_cutedsl.cuh b/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8_cutedsl.cuh index 11d6500ff71..fee0dc8a7d9 100644 --- a/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8_cutedsl.cuh +++ b/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8_cutedsl.cuh @@ -139,7 +139,7 @@ struct alignas(128) GroupDescriptorWorkspace { alignas(128) int64_t tensor_maps[kMaxGroupTensors][kGroupTensorMapSlots][kInt64PerTensorMap]; // Stand-in for the unused offsets array in SAME_BOTH_DIMS. The kernel takes // offsets unconditionally but does not read them for this representation. - int64_t unused_dims[kMaxGroupTensors + 1]; + int64_t unused_offsets[kMaxGroupTensors + 1]; }; // Like `g_tensor_maps` on the CUDA path, this has internal linkage, so every translation @@ -250,20 +250,24 @@ inline bool mxfp8_group_quantize_cutedsl(const MXFP8GroupQuantConfig &config, false, device_index); // Offsets and member dims are read from the output, as in mxfp8::group_quantize. - auto dims_or_unused = [&](const SimpleTensor &t, size_t numel) { - void *dptr = t.has_data() ? t.dptr : static_cast(workspace->unused_dims); + auto offsets_or_unused = [&](const SimpleTensor &t, size_t numel) { + void *dptr = t.has_data() ? t.dptr : static_cast(workspace->unused_offsets); return DLTensorWrapper(make_basic_tensor(dptr, DType::kInt64, {numel}), false, device_index); }; - DLTensorWrapper mOffsets = dims_or_unused(output_tensor->tensor_offsets, num_tensors + 1); + DLTensorWrapper mOffsets = offsets_or_unused(output_tensor->tensor_offsets, num_tensors + 1); DLTensorWrapper mFirstDims, mLastDims; if (config.shape_rep == ShapeRepresentation::VARYING_FIRST_DIM || config.shape_rep == ShapeRepresentation::VARYING_BOTH_DIMS) { + NVTE_CHECK(output_tensor->first_dims.has_data(), "Grouped MXFP8 quantization with ", + shape_rep_to_str(config.shape_rep), " requires an allocated first_dims buffer."); mFirstDims = DLTensorWrapper( make_basic_tensor(output_tensor->first_dims.dptr, DType::kInt64, {num_tensors}), false, device_index); } if (config.shape_rep == ShapeRepresentation::VARYING_LAST_DIM || config.shape_rep == ShapeRepresentation::VARYING_BOTH_DIMS) { + NVTE_CHECK(output_tensor->last_dims.has_data(), "Grouped MXFP8 quantization with ", + shape_rep_to_str(config.shape_rep), " requires an allocated last_dims buffer."); mLastDims = DLTensorWrapper( make_basic_tensor(output_tensor->last_dims.dptr, DType::kInt64, {num_tensors}), false, device_index); From 982b900a7a276353ff69e99939daf30cb705ce96 Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Fri, 2 Oct 2026 22:11:12 +0000 Subject: [PATCH 10/58] nit Signed-off-by: Kaining Zhong --- .../cast/mxfp8/group_quantize_mxfp8.py | 419 ++++++++---------- .../CuTeDSL/cast/mxfp8/quantize_mxfp8.py | 86 ++-- 2 files changed, 238 insertions(+), 267 deletions(-) diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8.py index d4c61f5730e..fc65638faf2 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8.py @@ -90,6 +90,7 @@ SUPPORTED_DACTIVATIONS, derive_swizzled_scale_layout, noop_flag_is_set, + reduce_rowwise_dbias, quantize_rowwise_mxfp8, quantize_colwise_mxfp8, ) @@ -922,6 +923,7 @@ def _kernel_main( if cutlass.const_expr(cfg.ROWWISE): row_scales = self._rowwise_scales(mS_row, Int64(0), tensor_rows, tensor_cols) + # Colwise scales require special handling for swizzled scales col_scales = None col_scale_row_start = job_start_row col_scale_rows = tensor_rows @@ -959,16 +961,15 @@ def _kernel_main( cute.arch.sync_threads() - self._process_job_strip( + self._process_job( job_start_row, - job_id_X, + job_id_X * self.TILE_COLS, tensor_rows, tensor_cols, row_scales, col_scales, col_scale_row_start, col_scale_rows, - job_start_row // (self.TILE_ROWS * self.NUM_TILES_Y), mWorkspace, sDbias, descs, @@ -1025,68 +1026,32 @@ def _kernel_main( job_start_row_in_tensor = job_id_Y_in_tensor * ( self.TILE_ROWS * self.NUM_TILES_Y ) - if cutlass.const_expr(self.NUM_TILES_X == 1): - self._process_job_strip( - job_start_row_in_tensor, - job_id_X_in_tensor, - tensor_rows, - tensor_cols, - row_scales, - col_scales, - job_start_row_in_tensor, - tensor_rows, - Int32(0), # dbias is only supported for single-tensor reps - mWorkspace, - sDbias, - descs, - tmap, - warp_idx, - tidx, - sX, - sActInput, - sO_row, - sO_col, - partitions, - atoms, - mainloop_pipeline, - prod_state, - cons_state, - ) - else: - # The job's column tiles in order, stopping at the tensor's last - # column like the CUDA kernel's stages_X = DIVUP(chunk_cols, TILE_DIM_X). - job_start_col = job_id_X_in_tensor * (self.TILE_COLS * self.NUM_TILES_X) - tiles_X = cutlass.min( - Int32(self.NUM_TILES_X), - cute.ceil_div(tensor_cols - job_start_col, self.TILE_COLS), - ) - for stage_X in cutlass.range(tiles_X, unroll=1): - self._process_job_strip( - job_start_row_in_tensor, - job_id_X_in_tensor * self.NUM_TILES_X + stage_X, - tensor_rows, - tensor_cols, - row_scales, - col_scales, - job_start_row_in_tensor, - tensor_rows, - Int32(0), # dbias is only supported for single-tensor reps - mWorkspace, - sDbias, - descs, - tmap, - warp_idx, - tidx, - sX, - sActInput, - sO_row, - sO_col, - partitions, - atoms, - mainloop_pipeline, - prod_state, - cons_state, - ) + job_start_col = job_id_X_in_tensor * (self.TILE_COLS * self.NUM_TILES_X) + self._process_job( + job_start_row_in_tensor, + job_start_col, + tensor_rows, + tensor_cols, + row_scales, + col_scales, + job_start_row_in_tensor, + tensor_rows, + mWorkspace, + sDbias, + descs, + tmap, + warp_idx, + tidx, + sX, + sActInput, + sO_row, + sO_col, + partitions, + atoms, + mainloop_pipeline, + prod_state, + cons_state, + ) # Find the next job to process job_id = job_id + job_stride if job_id >= jobs_in_tensor: @@ -1144,17 +1109,16 @@ def _issue_load( pipeline_obj.producer_commit(prod_state) @cute.jit - def _process_job_strip( + def _process_job( self, job_start_row, # Row offset of this job (global for single-tensor, else tensor-local) - column_tile_id, # Column-tile index within the tensor + job_start_col, # Column offset of this job within the tensor rows, # Rows of the rowwise-scale view (the group for single-tensor, else the tensor) cols, # Number of columns in this tensor row_scales, # Rowwise scales tiled per stage, rows counted like job_start_row col_scales, # Colwise scales tiled per stage col_scale_row_start, # Row of this job in the colwise-scale view col_scale_rows, # Rows of the colwise-scale view - dbias_row, # Row of the dbias workspace this job reduces into mWorkspace, # f32 partial dbias workspace (WITH_DBIAS) sDbias, # SMEM buffer for the rowwise dbias reduction (rowwise-only dbias) descs, # Per-tensor descriptors (x, act, out_row, out_col), None if single-tensor @@ -1171,30 +1135,19 @@ def _process_job_strip( prod_state, cons_state, ): - """Quantize NUM_TILES_Y vertically stacked tiles in one column strip of a job.""" + """Quantize a job with one continuous pipeline across its column and row tiles.""" cfg = self.cfg _, _, tma_atom_out_row, tma_atom_out_col = atoms _, _, _, _, tXsO_row, tXgO_row, tXsO_col, tXgO_col = partitions _, _, desc_out_row, desc_out_col = descs - job_start_col = column_tile_id * self.TILE_COLS - # This job's coordinates in the tile grid (32x128 TMA boxes, not elements). - tile_id_Y = job_start_row // self.TILE_ROWS - tile_id_X = column_tile_id + job_tile_Y = job_start_row // self.TILE_ROWS + job_tile_X = job_start_col // self.TILE_COLS col_scale_tile_Y = col_scale_row_start // self.TILE_ROWS - - # Per-job dbias accumulators, in the CUDA kernel's summation order: a running - # column sum over the job's rows (colwise), or per-thread partial sums over its - # stages that the whole CTA reduces afterwards (rowwise-only). - dbias_col = Float32(0.0) - dbias_row_acc = None - if cutlass.const_expr(self.DBIAS_IN_ROWWISE): - dbias_row_acc = cute.make_rmem_tensor( - layout_or_shape=cute.make_layout((MXFP8_BLOCK_SCALING_SIZE,), stride=(1,)), - dtype=Float32, - ) - for c in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE): - dbias_row_acc[c] = Float32(0.0) + tiles_X = cutlass.min( + Int32(self.NUM_TILES_X), cute.ceil_div(cols - job_start_col, self.TILE_COLS) + ) + num_tiles = tiles_X * self.NUM_TILES_Y # Fill every buffer up front, then issue one more each time a stage is consumed. for prologue_stage in cutlass.range_constexpr(self.PIPELINE_DEPTH): @@ -1202,8 +1155,8 @@ def _process_job_strip( self._issue_load( mainloop_pipeline, prod_state, - tile_id_Y + prologue_stage, - tile_id_X, + job_tile_Y + prologue_stage % self.NUM_TILES_Y, + job_tile_X + prologue_stage // self.NUM_TILES_Y, atoms, partitions, tmap, @@ -1211,159 +1164,159 @@ def _process_job_strip( ) prod_state.advance() - for stage in cutlass.range_constexpr(self.NUM_TILES_Y): - # Wait for at most DEPTH-1 iters on the fly, which means the the last DEPTH iter has finished - # so we can reuse its SMEM output buffer - # (input buffer is managed by the producer and consumer pipeline states) - if warp_idx == 0: - cute.arch.cp_async_bulk_wait_group(self.PIPELINE_DEPTH - 1, read=True) - # Wait for this stage's input buffer to be filled by the producer - mainloop_pipeline.consumer_wait(cons_state) - cute.arch.sync_threads() - sX_tile = sX[(None, cons_state.index)] - sAct_tile = None - if cutlass.const_expr(cfg.WITH_DACT): - sAct_tile = sActInput[(None, cons_state.index)] - row_tile = tile_id_Y + stage - - if cutlass.const_expr(cfg.COLWISE): - _, dbias_col = quantize_colwise_mxfp8( - sX_tile, - sAct_tile, - sO_col[(None, cons_state.index)], - cute.flatten(col_scales[(None, (col_scale_tile_Y + stage, tile_id_X))]), - cfg.MAX_NORM_RCP, - (col_scale_tile_Y + stage) * self.TILE_ROWS, - job_start_col, - col_scale_rows, - cols, - 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, - WITH_ACT=cfg.WITH_ACT, - WITH_DACT=cfg.WITH_DACT, - WITH_DBIAS=self.DBIAS_IN_COLWISE, - CACHE_ACTIVATION=self.CACHE_ACTIVATION, - ZERO_OOB_SCALES=True, - dbias_init=dbias_col, + for tile_X in cutlass.range(tiles_X, unroll=1): + tile_id_X = job_tile_X + tile_X + tile_start_col = job_start_col + tile_X * self.TILE_COLS + partial_dbias = Float32(0.0) + dbias_row_acc = None + if cutlass.const_expr(self.DBIAS_IN_ROWWISE): + dbias_row_acc = cute.make_rmem_tensor( + layout_or_shape=cute.make_layout((MXFP8_BLOCK_SCALING_SIZE,), stride=(1,)), + dtype=Float32, ) - if cutlass.const_expr(self.CACHE_ACTIVATION): - # The rowwise pass reads the activation the colwise pass cached in sX. + for c in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE): + dbias_row_acc[c] = Float32(0.0) + + for tile_Y in cutlass.range_constexpr(self.NUM_TILES_Y): + stage = tile_X * self.NUM_TILES_Y + tile_Y + # Wait for at most DEPTH-1 iters on the fly, which means the the last DEPTH iter has finished + # so we can reuse its SMEM output buffer + # (input buffer is managed by the producer and consumer pipeline states) + if warp_idx == 0: + cute.arch.cp_async_bulk_wait_group(self.PIPELINE_DEPTH - 1, read=True) + # Wait for this stage's input buffer to be filled by the producer + mainloop_pipeline.consumer_wait(cons_state) cute.arch.sync_threads() - if cutlass.const_expr(cfg.ROWWISE): - quantize_rowwise_mxfp8( - sX_tile, - None if self.CACHE_ACTIVATION else sAct_tile, - sO_row[(None, cons_state.index)], - cute.flatten(row_scales[(None, (row_tile, tile_id_X))]), - cfg.MAX_NORM_RCP, - row_tile * self.TILE_ROWS, - job_start_col, - rows, - cols, - 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, - WITH_ACT=cfg.WITH_ACT and not self.CACHE_ACTIVATION, - WITH_DACT=cfg.WITH_DACT and not self.CACHE_ACTIVATION, - WITH_DBIAS=self.DBIAS_IN_ROWWISE, - dbias_acc=dbias_row_acc, - ZERO_OOB_SCALES=True, - ) + sX_tile = sX[(None, cons_state.index)] + sAct_tile = None + if cutlass.const_expr(cfg.WITH_DACT): + sAct_tile = sActInput[(None, cons_state.index)] + row_tile = job_tile_Y + tile_Y - # Force consumer's write to SMEM to be visible to TMA stores later - cute.arch.fence_proxy("async.shared", space="cta") - # Only after everyone finishes computation then this stage can be considered as "consumed" - cute.arch.sync_threads() - # I'm done with my input SMEM buffer, so the producer can write the next stage's data into it - mainloop_pipeline.consumer_release(cons_state) - - # I just freed my input SMEM buffer (stage), so the producer now can use it for writing - # (stage+DEPTH) stage's data if that stage exists - if cutlass.const_expr(stage + self.PIPELINE_DEPTH < self.NUM_TILES_Y): - if warp_idx == 0: - self._issue_load( - mainloop_pipeline, - prod_state, - tile_id_Y + stage + self.PIPELINE_DEPTH, - tile_id_X, - atoms, - partitions, - tmap, - descs, + if cutlass.const_expr(cfg.COLWISE): + _, partial_dbias = quantize_colwise_mxfp8( + sX_tile, + sAct_tile, + sO_col[(None, cons_state.index)], + cute.flatten(col_scales[(None, (col_scale_tile_Y + tile_Y, tile_id_X))]), + cfg.MAX_NORM_RCP, + (col_scale_tile_Y + tile_Y) * self.TILE_ROWS, + tile_start_col, + col_scale_rows, + cols, + 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, + WITH_ACT=cfg.WITH_ACT, + WITH_DACT=cfg.WITH_DACT, + WITH_DBIAS=self.DBIAS_IN_COLWISE, + CACHE_ACTIVATION=self.CACHE_ACTIVATION, + ZERO_OOB_SCALES=True, + dbias_init=partial_dbias, ) - prod_state.advance() + if cutlass.const_expr(self.CACHE_ACTIVATION): + # The rowwise pass reads the activation the colwise pass cached in sX. + cute.arch.sync_threads() - # Write result to GMEM via TMA - if warp_idx == 0: - stores = [] if cutlass.const_expr(cfg.ROWWISE): - stores.append((tma_atom_out_row, tXsO_row, tXgO_row, desc_out_row)) - if cutlass.const_expr(cfg.COLWISE): - stores.append((tma_atom_out_col, tXsO_col, tXgO_col, desc_out_col)) - for atom, tXs, tXg, desc in stores: - if cutlass.const_expr(cfg.IS_SINGLE_TENSOR): - cute.copy( - atom, - tXs[(None, cons_state.index)], - tXg[(None, (row_tile, tile_id_X))], - ) - else: - cute.copy( - atom, - tXs[(None, cons_state.index)], - tXg[(None, (row_tile, tile_id_X))], - tma_desc_ptr=tmap.get_tensormap_ptr(desc, cute.AddressSpace.generic), - ) - # Commit all TMA operations of this iteration - cute.arch.cp_async_bulk_commit_group() - - cons_state.advance() + quantize_rowwise_mxfp8( + sX_tile, + None if self.CACHE_ACTIVATION else sAct_tile, + sO_row[(None, cons_state.index)], + cute.flatten(row_scales[(None, (row_tile, tile_id_X))]), + cfg.MAX_NORM_RCP, + row_tile * self.TILE_ROWS, + tile_start_col, + rows, + cols, + 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, + WITH_ACT=cfg.WITH_ACT and not self.CACHE_ACTIVATION, + WITH_DACT=cfg.WITH_DACT and not self.CACHE_ACTIVATION, + WITH_DBIAS=self.DBIAS_IN_ROWWISE, + dbias_acc=dbias_row_acc, + ZERO_OOB_SCALES=True, + ) - if cutlass.const_expr(cfg.WITH_DBIAS): - if cutlass.const_expr(self.DBIAS_IN_ROWWISE): - dbias_col = self._reduce_rowwise_dbias(sDbias, tidx, dbias_row_acc) - # One partial-dbias row per job, as in the CUDA kernel's dbias_workspace. - dbias_x = job_start_col + tidx - if dbias_x < cols: - mWorkspace[(dbias_row, dbias_x)] = dbias_col + # Force consumer's write to SMEM to be visible to TMA stores later + cute.arch.fence_proxy("async.shared", space="cta") + # Only after everyone finishes computation then this stage can be considered as "consumed" + cute.arch.sync_threads() + # I'm done with my input SMEM buffer, so the producer can write the next stage's data into it + mainloop_pipeline.consumer_release(cons_state) + + # I just freed my input SMEM buffer (stage), so the producer now can use it for writing + # (stage+DEPTH) stage's data if that stage exists + if stage + self.PIPELINE_DEPTH < num_tiles: + if warp_idx == 0: + self._issue_load( + mainloop_pipeline, + prod_state, + job_tile_Y + (stage + self.PIPELINE_DEPTH) % self.NUM_TILES_Y, + job_tile_X + (stage + self.PIPELINE_DEPTH) // self.NUM_TILES_Y, + atoms, + partitions, + tmap, + descs, + ) + prod_state.advance() - @cute.jit - def _reduce_rowwise_dbias(self, sDbias, tidx, dbias_row_acc): - """Reduce the per-thread rowwise partial sums to one sum per column, in the order of the - CUDA kernel's partial_dbias_rowwise reduction.""" - _, tv_write = cute.make_layout_tv( - thr_layout=cute.make_layout( - (self.THREADS_Y, self.THREADS_X), stride=(self.THREADS_X, 1) - ), - val_layout=cute.make_layout( - (1, MXFP8_BLOCK_SCALING_SIZE), stride=(MXFP8_BLOCK_SCALING_SIZE, 1) - ), - ) - sDbias_write = cute.composition(sDbias, tv_write) - bank_group = (tidx % THREADS_PER_WARP) // self.THREADS_PER_BANK - offset = bank_group * self.PACK_SIZE - for w in cutlass.range_constexpr(self.WAVES): - # Undo the bank-conflict rotation quantize_rowwise_mxfp8 accumulated in. - start = (w * self.PACK_SIZE + offset) % MXFP8_BLOCK_SCALING_SIZE - for i in cutlass.range_constexpr(self.PACK_SIZE): - sDbias_write[(tidx, start + i)] = dbias_row_acc[w * self.PACK_SIZE + i] - cute.arch.sync_threads() - # Thread tidx sums column tidx over the THREADS_Y partial rows. - dbias = Float32(0.0) - for i in cutlass.range_constexpr(self.THREADS_Y): - dbias += sDbias[(i, tidx)] - # The buffer is rewritten by the next job. - cute.arch.sync_threads() - return dbias + # Write result to GMEM via TMA + if warp_idx == 0: + stores = [] + if cutlass.const_expr(cfg.ROWWISE): + stores.append((tma_atom_out_row, tXsO_row, tXgO_row, desc_out_row)) + if cutlass.const_expr(cfg.COLWISE): + stores.append((tma_atom_out_col, tXsO_col, tXgO_col, desc_out_col)) + for atom, tXs, tXg, desc in stores: + if cutlass.const_expr(cfg.IS_SINGLE_TENSOR): + cute.copy( + atom, + tXs[(None, cons_state.index)], + tXg[(None, (row_tile, tile_id_X))], + ) + else: + cute.copy( + atom, + tXs[(None, cons_state.index)], + tXg[(None, (row_tile, tile_id_X))], + tma_desc_ptr=tmap.get_tensormap_ptr( + desc, cute.AddressSpace.generic + ), + ) + # Commit all TMA operations of this iteration + cute.arch.cp_async_bulk_commit_group() + + cons_state.advance() + + if cutlass.const_expr(cfg.WITH_DBIAS): + if cutlass.const_expr(self.DBIAS_IN_ROWWISE): + partial_dbias = reduce_rowwise_dbias( + sDbias, + tidx, + dbias_row_acc, + self.TILE_ROWS, + self.TILE_COLS, + self.PACK_SIZE, + self.THREADS_PER_BANK, + ) + # All threads must finish reading before the next job rewrites the buffer. + cute.arch.sync_threads() + + # A job has TILE_ROWS * NUM_TILES_Y rows, and dbias reduces them to one row + dbias_row = job_start_row // (self.TILE_ROWS * self.NUM_TILES_Y) + dbias_col = tile_start_col + tidx + if dbias_col < cols: + mWorkspace[(dbias_row, dbias_col)] = partial_dbias def compile_cutedsl_function_from_cfg(cfg: MXFP8GroupQuantizeConfig): diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8.py index 9880276d396..c267638bbd7 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8.py @@ -815,6 +815,50 @@ def quantize_bidimensional_mxfp8_swizzled( cute.autovec_copy(rO_col, tXsO_col) +@cute.jit +def reduce_rowwise_dbias( + sDbias: cute.Tensor, + tidx: Int32, + rowwise_dbias_acc: cute.Tensor, + TILE_ROWS: cutlass.Constexpr[int], + TILE_COLS: cutlass.Constexpr[int], + PACK_SIZE: cutlass.Constexpr[int], + THREADS_PER_BANK: cutlass.Constexpr[int], +): + """Reduce per-thread rowwise partial sums to one sum per column using shared memory.""" + _, tv_layout_dbias_write = cute.make_layout_tv( + thr_layout=cute.make_layout( + (TILE_ROWS, TILE_COLS // MXFP8_BLOCK_SCALING_SIZE), + stride=(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) + # Undo the bank-conflict rotation used when accumulating the per-thread sums. + bank_group = (tidx % THREADS_PER_WARP) // THREADS_PER_BANK + offset = bank_group * PACK_SIZE + for w in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE // PACK_SIZE): + start = (w * PACK_SIZE + offset) % MXFP8_BLOCK_SCALING_SIZE + for i in cutlass.range_constexpr(PACK_SIZE): + # All threads write their per-thread partial sum results to the shared buffer. + sDbias_write[(tidx, start + i)] = rowwise_dbias_acc[w * 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, TILE_COLS), stride=(TILE_COLS, 1)), + val_layout=cute.make_layout((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(TILE_ROWS): + block_dbias += sDbias_reduce[tidx, i] + return block_dbias + + @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. @@ -1645,41 +1689,15 @@ class 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) - ), + return reduce_rowwise_dbias( + sDbias, + tidx, + rowwise_dbias_acc, + self._TILE_ROWS, + self._TILE_COLS, + self._PACK_SIZE, + self._THREADS_PER_BANK, ) - 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( From 436a479c0668a8bec00a58a747106e946cca7ce6 Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Fri, 2 Oct 2026 23:17:46 +0000 Subject: [PATCH 11/58] type annotation Signed-off-by: Kaining Zhong --- .../cast/mxfp8/group_quantize_mxfp8.py | 271 ++++++++---------- 1 file changed, 117 insertions(+), 154 deletions(-) diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8.py index fc65638faf2..7d90d56da93 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8.py @@ -2,63 +2,7 @@ # # See LICENSE for license information. -"""Grouped MXFP8 quantization kernel implemented in CuTeDSL. - -Strategy-aligned port of group_quantize_mxfp8.cuh. The scheduling, descriptor -management and per-tensor scale addressing mirror the CUDA kernel one-for-one: - - * `is_single_tensor` reps (SAME_BOTH_DIMS, VARYING_FIRST_DIM) launch ONE CTA per - 128x128 job and address the group through ONE static TMA descriptor with - global job offsets -- the CUDA `tensor_map_*_static` "direct mapper" path. - For SAME_BOTH_DIMS the CUDA grid is linearized per tensor (X, Y-in-tensor, - tensor) while this kernel linearizes it flat over the stacked rows; both - require every member's row count to be a multiple of 128, and under - that precondition the two decode to the identical (job_start_row, job_id_X) - for every job index. - * the other reps launch grid=(workers_per_tensor, num_tensors) and bind - tensor_id to blockIdx.y, so a CTA grid-strides only within its own tensor and - never re-resolves which tensor a job belongs to. They get per-tensor - descriptors written by a prologue kernel (the CuTeDSL analog of - update_tma_descriptors filling g_tensor_maps) and acquired with a tensormap - proxy fence. - * per-tensor scale bases/strides follow the CUDA formulas: - scales_* += is_single_tensor ? 0 : tensor_base / 32 - stride_rowwise = roundup(cols/32, 4) stride_colwise = roundup(cols, 128) - -Both kernels dropped the older flat persistent grid that strided across tensor -boundaries (CUDA's decode_job / advance_to_next_job, removed in #3483); the -per-tensor grid above is what replaced it on both sides. - -Mechanics that provably yield the same bytes may differ: the mbarrier pipeline is -expressed with PipelineTmaAsync instead of hand-rolled mbarriers. As in CUDA, the -scales of out-of-bounds columns in a job (the scale-row padding) are written as 0. - -Scope: everything group_quantize_mxfp8.cuh covers except 2D block scaling -- the -cast-noop flag, fused activation (IS_ACT) and activation derivative (IS_DACT), dbias, -compact and GEMM-swizzled scales, rowwise and/or colwise, and all four shape -representations. Differences from CUDA: - * the grouped amax pointer is accepted and left untouched, as the CUDA kernel does; - * dbias partial sums are accumulated in the CUDA kernel's order (a running column sum - with colwise output, otherwise per-thread sums reduced across the CTA), and the C++ - bridge reduces the workspace with the same grouped_reduce_dbias. - -Like the CUDA kernel, every member's first dim must be a multiple of 128 (and, for the -varying-last reps, its last dim too). The kernel prints the same diagnostics as -get_tensor_rows_num / get_tensor_cols_num when a group violates this and, like -NVTE_DEVICE_ERROR in a release build, carries on. - -Measured, deliberately NOT changed: - - sO_row and sO_col are both allocated unconditionally, where CUDA sizes only the - direction in use. Sizing them conditionally does work -- ncu confirms the shared-memory - occupancy limit goes 6 -> 9 CTAs/SM for a single-direction bf16 config -- but it is a - small LOSS on GB200, not a win: rep_med_sbd bf16 colwise 54.9 -> 56.1 us, rowwise - 56.5 -> 57.0 us (fp32 rowwise gains ~1%). The kernel is DRAM-bandwidth-bound at - ~6.3 TB/s, so extra resident CTAs only add contention. Verified by a control that kept - the conditional code but padded SMEM back to the old size: timings returned exactly to - the unconditional numbers, so the effect is the occupancy, not codegen. Revisit if a - future variant (dbias / activation) makes this kernel latency- rather than - bandwidth-bound. -""" +"""Grouped MXFP8 quantization kernel implemented in CuTeDSL.""" # pylint: disable=missing-class-docstring @@ -72,8 +16,8 @@ from cutlass import Boolean, Float32, Int32, Int64, Float8E8M0FNU from cutlass.cute.nvgpu import cpasync from cutlass.cute.testing import assert_ as runtime_assert -from cutlass.utils import TensorMapManager, TensorMapUpdateMode -from cutlass.utils import HardwareInfo +from cutlass.utils import TensorMapUpdateMode, HardwareInfo +from cutlass.tensor_utils import TensorMapManager from cuda.bindings.driver import CUstream # pylint: disable=no-name-in-module import tvm_ffi @@ -226,7 +170,7 @@ class MXFP8GroupQuantizeKernel: # How many threads per bank -- for avoiding bank conflicts THREADS_PER_BANK = (32 * 4) // MXFP8_BLOCK_SCALING_SIZE # 4 - def __init__(self, cfg: MXFP8GroupQuantizeConfig, SM_COUNT: int): + def __init__(self, cfg: MXFP8GroupQuantizeConfig, SM_COUNT: int) -> None: self.cfg = cfg self.SM_COUNT = SM_COUNT # A CTA processes (NUM_TILES_Y, NUM_TILES_X) tiles, NUM_STAGES tiles in total @@ -249,7 +193,9 @@ def __init__(self, cfg: MXFP8GroupQuantizeConfig, SM_COUNT: int): self.DBIAS_IN_ROWWISE = cfg.WITH_DBIAS and not cfg.COLWISE @cute.jit - def _find_tensor_from_offsets(self, mOffsets, num_tensors, offset: Int64): + def _find_tensor_from_offsets( + self, mOffsets: cute.Tensor, num_tensors: Int32, offset: Int64 + ) -> Int32: """Index of the tensor whose element range holds `offset` (find_tensor_from_offsets).""" low = Int32(1) hi = Int32(num_tensors) @@ -264,7 +210,7 @@ def _find_tensor_from_offsets(self, mOffsets, num_tensors, offset: Int64): return low - 1 @cute.jit - def _scale_tensor(self, mS, base: Int64, layout): + def _scale_tensor(self, mS: cute.Tensor, base: Int64, layout: cute.Layout) -> cute.Tensor: """View the scale buffer from element `base` on with `layout`.""" return cute.make_tensor( cute.make_ptr( @@ -277,7 +223,9 @@ def _scale_tensor(self, mS, base: Int64, layout): ) @cute.jit - def _rowwise_scales(self, mS_row, base: Int64, rows, cols): + def _rowwise_scales( + self, mS_row: cute.Tensor, base: Int64, rows: Int32, cols: Int32 + ) -> cute.Tensor: """Rowwise scales of a (rows, cols) tensor at `base`, tiled per 32x128 stage.""" if cutlass.const_expr(self.cfg.WITH_GEMM_SWIZZLED_SCALES): mS_t, _ = derive_swizzled_scale_layout( @@ -294,7 +242,9 @@ def _rowwise_scales(self, mS_row, base: Int64, rows, cols): ) @cute.jit - def _colwise_scales(self, mS_col, base: Int64, rows, cols): + def _colwise_scales( + self, mS_col: cute.Tensor, base: Int64, rows: Int32, cols: Int32 + ) -> cute.Tensor: """Colwise scales of a (rows, cols) tensor at `base`, tiled per 32x128 stage.""" if cutlass.const_expr(self.cfg.WITH_GEMM_SWIZZLED_SCALES): _, mS_t = derive_swizzled_scale_layout( @@ -332,7 +282,7 @@ def __call__( mActInput: Optional[cute.Tensor], # activation input, only with WITH_DACT mWorkspace: Optional[cute.Tensor], # f32 partial dbias, only with WITH_DBIAS stream: CUstream, - ): + ) -> None: if cutlass.const_expr(CUTEDSL_DEBUG_LOGGING): cute.printf(f"[CuTeDSL] MXFP8GroupQuantizeKernel.__call__() cfg: {self.cfg}\n") @@ -469,22 +419,22 @@ def __call__( @cute.kernel def _update_descriptors_kernel( self, - mX, - mO_row, - mO_col, - mActInput, - mOffsets, - mFirstDims, - mLastDims, - mTensormaps, - first_logical_dim, - last_logical_dim, + mX: cute.Tensor, + mO_row: Optional[cute.Tensor], + mO_col: Optional[cute.Tensor], + mActInput: Optional[cute.Tensor], + mOffsets: cute.Tensor, + mFirstDims: Optional[cute.Tensor], + mLastDims: Optional[cute.Tensor], + mTensormaps: cute.Tensor, + first_logical_dim: Int32, + last_logical_dim: Int32, dtype: cutlass.Constexpr[Type[cutlass.Numeric]], - tma_atom_x, - tma_atom_orow, - tma_atom_ocol, - tma_atom_act, - ): + tma_atom_x: cute.CopyAtom, + tma_atom_orow: Optional[cute.CopyAtom], + tma_atom_ocol: Optional[cute.CopyAtom], + tma_atom_act: Optional[cute.CopyAtom], + ) -> None: """Update the per-tensor TMA descriptors for the group quantization kernel. mTensormaps: int64[num_tensors, NUM_WORKSPACE_SLOTS, 16], where the slots are: @@ -543,7 +493,7 @@ def _update_descriptors_kernel( if rows > 0 and cols > 0: member_layout = cute.make_layout((rows, cols), stride=(cols, 1)) - def member_view(tensor, elt_dtype): + def member_view(tensor: cute.Tensor, elt_dtype: Type[cutlass.Numeric]) -> cute.Tensor: return cute.make_tensor( cute.make_ptr( elt_dtype, @@ -586,9 +536,16 @@ def member_view(tensor, elt_dtype): @cute.jit def _make_shared_storage( self, - smem: cutlass.Constexpr, + smem: cutlass.Constexpr[cutlass.utils.SmemAllocator], dtype: cutlass.Constexpr[Type[cutlass.Numeric]], - ): + ) -> tuple[ + cute.Pointer, + cute.Tensor, + Optional[cute.Tensor], + Optional[cute.Tensor], + Optional[cute.Tensor], + Optional[cute.Tensor], + ]: """Allocate pipeline buffers and optional activation input and dbias storage.""" FP8_DTYPE = self.cfg.FP8_DTYPE tile_layout = cute.make_layout( @@ -705,27 +662,27 @@ class DbiasStorage: @cute.kernel def kernel( self, - mS_row, - mS_col, - mOffsets, - mFirstDims, - mTensormaps, - mNoop, - mWorkspace, - first_logical_dim, - last_logical_dim, - num_tensors, - jobs_X, + mS_row: cute.Tensor, + mS_col: cute.Tensor, + mOffsets: cute.Tensor, + mFirstDims: Optional[cute.Tensor], + mTensormaps: cute.Tensor, + mNoop: cute.Pointer, + mWorkspace: Optional[cute.Tensor], + first_logical_dim: Int32, + last_logical_dim: Int32, + num_tensors: Int32, + jobs_X: Optional[Int32], dtype: cutlass.Constexpr[Type[cutlass.Numeric]], - tma_atom_x, - tma_src, - tma_atom_act, - tma_src_act, - tma_atom_out_row, - tma_dst_out_row, - tma_atom_out_col, - tma_dst_out_col, - ): + tma_atom_x: cute.CopyAtom, + tma_src: cute.Tensor, + tma_atom_act: Optional[cute.CopyAtom], + tma_src_act: Optional[cute.Tensor], + tma_atom_out_row: Optional[cute.CopyAtom], + tma_dst_out_row: Optional[cute.Tensor], + tma_atom_out_col: Optional[cute.CopyAtom], + tma_dst_out_col: Optional[cute.Tensor], + ) -> None: """No-op the CTA when the noop flag is set, else run the quantize main loop.""" skip_execution = Boolean(False) if cutlass.const_expr(self.CHECK_NOOP_FLAG): @@ -756,26 +713,26 @@ def kernel( @cute.jit def _kernel_main( self, - mS_row, - mS_col, - mOffsets, - mFirstDims, - mTensormaps, - mWorkspace, - first_logical_dim, - last_logical_dim, - num_tensors, - jobs_X, + mS_row: cute.Tensor, + mS_col: cute.Tensor, + mOffsets: cute.Tensor, + mFirstDims: Optional[cute.Tensor], + mTensormaps: cute.Tensor, + mWorkspace: Optional[cute.Tensor], + first_logical_dim: Int32, + last_logical_dim: Int32, + num_tensors: Int32, + jobs_X: Optional[Int32], dtype: cutlass.Constexpr[Type[cutlass.Numeric]], - tma_atom_x, - tma_src, - tma_atom_act, - tma_src_act, - tma_atom_out_row, - tma_dst_out_row, - tma_atom_out_col, - tma_dst_out_col, - ): + tma_atom_x: cute.CopyAtom, + tma_src: cute.Tensor, + tma_atom_act: Optional[cute.CopyAtom], + tma_src_act: Optional[cute.Tensor], + tma_atom_out_row: Optional[cute.CopyAtom], + tma_dst_out_row: Optional[cute.Tensor], + tma_atom_out_col: Optional[cute.CopyAtom], + tma_dst_out_col: Optional[cute.Tensor], + ) -> None: cfg = self.cfg FP8_DTYPE = cfg.FP8_DTYPE tidx, _, _ = cute.arch.thread_idx() @@ -1064,15 +1021,15 @@ def _kernel_main( def _issue_load( self, - pipeline_obj, - prod_state, - tile_y, - tile_x, - atoms, - partitions, - tmap, - descs, - ): + pipeline_obj: pipeline.PipelineTmaAsync, + prod_state: pipeline.PipelineState, + tile_y: Int32, + tile_x: Int32, + atoms: tuple[Optional[cute.CopyAtom], ...], + partitions: tuple[Optional[cute.Tensor], ...], + tmap: TensorMapManager, + descs: tuple[Optional[cute.Pointer], ...], + ) -> None: """Emit the 32x128 TMA load(s) of one stage into the current pipeline buffer. Caller gates this on warp 0 and advances `prod_state` afterwards -- the advance @@ -1111,30 +1068,36 @@ def _issue_load( @cute.jit def _process_job( self, - job_start_row, # Row offset of this job (global for single-tensor, else tensor-local) - job_start_col, # Column offset of this job within the tensor - rows, # Rows of the rowwise-scale view (the group for single-tensor, else the tensor) - cols, # Number of columns in this tensor - row_scales, # Rowwise scales tiled per stage, rows counted like job_start_row - col_scales, # Colwise scales tiled per stage - col_scale_row_start, # Row of this job in the colwise-scale view - col_scale_rows, # Rows of the colwise-scale view - mWorkspace, # f32 partial dbias workspace (WITH_DBIAS) - sDbias, # SMEM buffer for the rowwise dbias reduction (rowwise-only dbias) - descs, # Per-tensor descriptors (x, act, out_row, out_col), None if single-tensor - tmap, # TensorMapManager for managing TMA descriptors - warp_idx, - tidx, - sX, # SMEM input ring - sActInput, # SMEM activation input ring (WITH_DACT) - sO_row, # SMEM rowwise output ring - sO_col, # SMEM colwise output ring - partitions, # TMA partitions (x, act, out_row, out_col) - atoms, # TMA atoms (x, act, out_row, out_col) + job_start_row: Int32, # Row offset of this job (global for single-tensor, else tensor-local) + job_start_col: Int32, # Column offset of this job within the tensor + rows: Int32, # Rows of the rowwise-scale view (the group for single-tensor, else the tensor) + cols: Int32, # Number of columns in this tensor + row_scales: Optional[ + cute.Tensor + ], # Rowwise scales tiled per stage, rows counted like job_start_row + col_scales: Optional[cute.Tensor], # Colwise scales tiled per stage + col_scale_row_start: Int32, # Row of this job in the colwise-scale view + col_scale_rows: Int32, # Rows of the colwise-scale view + mWorkspace: Optional[cute.Tensor], # f32 partial dbias workspace (WITH_DBIAS) + sDbias: Optional[ + cute.Tensor + ], # SMEM buffer for the rowwise dbias reduction (rowwise-only dbias) + descs: tuple[ + Optional[cute.Pointer], ... + ], # Per-tensor descriptors (x, act, out_row, out_col), None if single-tensor + tmap: TensorMapManager, # TensorMapManager for managing TMA descriptors + warp_idx: Int32, + tidx: Int32, + sX: cute.Tensor, # SMEM input ring + sActInput: Optional[cute.Tensor], # SMEM activation input ring (WITH_DACT) + sO_row: Optional[cute.Tensor], # SMEM rowwise output ring + sO_col: Optional[cute.Tensor], # SMEM colwise output ring + partitions: tuple[Optional[cute.Tensor], ...], # TMA partitions (x, act, out_row, out_col) + atoms: tuple[Optional[cute.CopyAtom], ...], # TMA atoms (x, act, out_row, out_col) mainloop_pipeline: cutlass.pipeline.PipelineTmaAsync, - prod_state, - cons_state, - ): + prod_state: pipeline.PipelineState, + cons_state: pipeline.PipelineState, + ) -> None: """Quantize a job with one continuous pipeline across its column and row tiles.""" cfg = self.cfg _, _, tma_atom_out_row, tma_atom_out_col = atoms From a87f59fde3d1dedef398d61be90a8bb344c45337 Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Fri, 2 Oct 2026 23:21:33 +0000 Subject: [PATCH 12/58] Remove unrelated C++ test additions from CuTeDSL changes Signed-off-by: Kaining Zhong --- tests/cpp/CMakeLists.txt | 1 - tests/cpp/operator/test_cast_mxfp8_grouped.cu | 59 ------------------- 2 files changed, 60 deletions(-) diff --git a/tests/cpp/CMakeLists.txt b/tests/cpp/CMakeLists.txt index 9a1e0bae18f..af074be1c78 100644 --- a/tests/cpp/CMakeLists.txt +++ b/tests/cpp/CMakeLists.txt @@ -57,7 +57,6 @@ if(NVTE_WITH_CUTEDSL) # Find the interpreter too so older CMake versions use its version to locate libpython. find_package(Python COMPONENTS Interpreter Development.Embed REQUIRED) foreach(test_target test_operator test_util) - target_compile_definitions(${test_target} PRIVATE NVTE_WITH_CUTEDSL) target_link_libraries(${test_target} PRIVATE "-Wl,--no-as-needed" Python::Python diff --git a/tests/cpp/operator/test_cast_mxfp8_grouped.cu b/tests/cpp/operator/test_cast_mxfp8_grouped.cu index 13c74491a11..c2309f7aa4b 100644 --- a/tests/cpp/operator/test_cast_mxfp8_grouped.cu +++ b/tests/cpp/operator/test_cast_mxfp8_grouped.cu @@ -9,11 +9,6 @@ #include #include -#ifdef NVTE_WITH_CUTEDSL -#include -#include -#endif - #include #include #include "../test_common.h" @@ -1035,57 +1030,3 @@ INSTANTIATE_TEST_SUITE_P( ::testing::Values(DType::kBFloat16), ::testing::Values(DType::kFloat8E4M3)), MakeGroupedFusedCastMXFP8TestName); - -// Exercise the grouped C API with NVTE_ENABLE_CUTEDSL_BACKEND=1 as well as CUDA. -// These cases cover both column strips, a final half-width chunk, and every fused -// activation on all four shape representations against the independent CPU reference. -INSTANTIATE_TEST_SUITE_P( - OperatorTest_GroupedFusedCastMXFP8_MultiChunkActivations, - GroupedFusedCastMXFP8TestSuite, - ::testing::Combine( - ::testing::Values(ProcessingMethod::CAST_ACT, ProcessingMethod::CAST_DACT, - ProcessingMethod::CAST_DBIAS_DACT), - ::testing::Values(ActivationKind::GeLU, ActivationKind::SiLU, ActivationKind::ReLU, - ActivationKind::QGeLU, ActivationKind::SReLU), - ::testing::ValuesIn(scaling_directions), - ::testing::ValuesIn(input_config_multichunk), - ::testing::Values(DType::kBFloat16), - ::testing::Values(DType::kFloat8E4M3)), - MakeGroupedFusedCastMXFP8TestName); - -#ifdef NVTE_WITH_CUTEDSL -TEST(OperatorTest_GroupedFusedCastMXFP8, TestCuTeDSLRegistration) { - const char* enabled = std::getenv("NVTE_ENABLE_CUTEDSL_BACKEND"); - if (enabled == nullptr || enabled[0] == '0' || - getDeviceComputeCapability() < blackwellComputeCapability) { - GTEST_SKIP() << "Requires Blackwell and NVTE_ENABLE_CUTEDSL_BACKEND=1"; - } - - // Cover both column strips and a final half-width chunk. Verify registration - // after the C API call so a silent CUDA fallback cannot satisfy this test. - const std::vector first_dims = {128, 256}; - const std::vector last_dims = {128, 384}; - const std::vector offsets = {0, 128 * 128, 128 * 128 + 256 * 384}; - performTest(CAST_ONLY, &identity, VARYING_BOTH_DIMS, 2, - {1, offsets.back()}, first_dims, last_dims, offsets, - /*rowwise=*/true, /*colwise=*/true); - - ASSERT_TRUE(Py_IsInitialized()) << "CuTeDSL did not initialize embedded Python"; - const std::string key = - "cutedsl_group_mxfp8_sm" + std::to_string(getDeviceComputeCapability()) + - "_BFloat16_Float8E4M3_1_1_varying_both_dims_0_0_0_0_none"; - const PyGILState_STATE gil = PyGILState_Ensure(); - PyObject* module = PyImport_ImportModule("tvm_ffi"); - PyObject* kernel = module == nullptr ? nullptr : - PyObject_CallMethod(module, "get_global_func", "s", key.c_str()); - const bool registered = kernel != nullptr && kernel != Py_None; - if (PyErr_Occurred() != nullptr) { - PyErr_Print(); - } - Py_XDECREF(kernel); - Py_XDECREF(module); - PyGILState_Release(gil); - EXPECT_TRUE(registered) << "CuTeDSL kernel not registered for " << key - << "; the grouped C API fell back to CUDA"; -} -#endif From a04ca31c4ddf0b44ce88621f2b4bbd7046080dcd Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Fri, 2 Oct 2026 23:57:58 +0000 Subject: [PATCH 13/58] nit Signed-off-by: Kaining Zhong --- transformer_engine/common/cast/dispatch/quantize.cuh | 4 ++-- .../common/cast/mxfp8/group_quantize_mxfp8_cutedsl.cuh | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/transformer_engine/common/cast/dispatch/quantize.cuh b/transformer_engine/common/cast/dispatch/quantize.cuh index 985c0f625d2..2e0d681d1a9 100644 --- a/transformer_engine/common/cast/dispatch/quantize.cuh +++ b/transformer_engine/common/cast/dispatch/quantize.cuh @@ -548,7 +548,7 @@ void group_quantize_fwd_helper(const NVTEGroupedTensor input, NVTEGroupedTensor cutedsl_backend::mxfp8_group_quantize_cutedsl( input_tensor, activations_tensor, noop_tensor, output_tensor, dbias_tensor, - workspace_tensor, quant_config_cpp.mxfp8_2d_quantization, stream); + workspace_tensor, &quant_config_cpp, stream); #endif if (!quantized_with_cutedsl) { mxfp8::group_quantize( @@ -653,7 +653,7 @@ void group_quantize_bwd_helper(const NVTEGroupedTensor grad, const NVTEGroupedTe cutedsl_backend::mxfp8_group_quantize_cutedsl( grad_tensor, input_tensor, noop_tensor, output_tensor, dbias_tensor, workspace_tensor, - quant_config_cpp.mxfp8_2d_quantization, stream); + &quant_config_cpp, stream); #endif if (!quantized_with_cutedsl) { mxfp8::group_quantize( diff --git a/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8_cutedsl.cuh b/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8_cutedsl.cuh index fee0dc8a7d9..edb674a09c0 100644 --- a/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8_cutedsl.cuh +++ b/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8_cutedsl.cuh @@ -320,14 +320,14 @@ template mxfp8_2d_quantization) { maybe_warn_cutedsl_not_chosen("2D quantization is not supported."); return false; } From beaded62e19e252ccbe3643309d3ea34fd95dd20 Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Fri, 2 Oct 2026 23:59:49 +0000 Subject: [PATCH 14/58] nit Signed-off-by: Kaining Zhong --- transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh index c39624889ca..bd8c3052f60 100644 --- a/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh @@ -882,8 +882,7 @@ __global__ void __launch_bounds__(CastTraits::THREADS_PER_CHUNK) group_quantize_ DIVUP_TO_MULTIPLE(DIVUP(cols, static_cast(SCALE_DIM_X)), scale_alignment_X_rowwise); const size_t scale_stride_colwise = DIVUP_TO_MULTIPLE(cols, scale_alignment_X_colwise); - // Non-single-tensor scale pointers already include the member offset. - const size_t tensor_base_for_scales = is_single_tensor ? tensor_start_offset : 0; + const size_t tensor_base_for_scales = is_single_tensor ? tensor_start_offset : tensor_base; e8m0_t *const scales_rowwise = scales_rowwise_ptr + (is_single_tensor ? 0 : tensor_base / SCALE_DIM_X); e8m0_t *const scales_colwise = From 8f88f4f7ba52e45cb537a79edf9324c1bda0c5f8 Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Sat, 3 Oct 2026 00:05:34 +0000 Subject: [PATCH 15/58] remove first launch protection Signed-off-by: Kaining Zhong --- .../cast/mxfp8/quantize_mxfp8_cutedsl.cuh | 4 +- transformer_engine/common/tvm_ffi_bridge.h | 74 +++---------------- 2 files changed, 14 insertions(+), 64 deletions(-) diff --git a/transformer_engine/common/cast/mxfp8/quantize_mxfp8_cutedsl.cuh b/transformer_engine/common/cast/mxfp8/quantize_mxfp8_cutedsl.cuh index dffeff6ed4f..431812db143 100644 --- a/transformer_engine/common/cast/mxfp8/quantize_mxfp8_cutedsl.cuh +++ b/transformer_engine/common/cast/mxfp8/quantize_mxfp8_cutedsl.cuh @@ -61,7 +61,7 @@ struct MXFP8QuantConfig { (static_cast(activation) << 16) | (sm_arch << 22); } - std::optional get_kernel() const { + std::optional get_kernel() const { static TVMFFIConfigCache &cache = TVMFFIConfigCache::create(); return cache.get_or_load(*this); } @@ -142,7 +142,7 @@ inline bool mxfp8_quantize_cutedsl(const MXFP8QuantConfig &config, const Tensor return true; } - std::optional mxfp8_quant_func_opt = config.get_kernel(); + std::optional mxfp8_quant_func_opt = config.get_kernel(); if (!mxfp8_quant_func_opt.has_value()) { return false; } diff --git a/transformer_engine/common/tvm_ffi_bridge.h b/transformer_engine/common/tvm_ffi_bridge.h index a231b452813..d57fd8f03da 100644 --- a/transformer_engine/common/tvm_ffi_bridge.h +++ b/transformer_engine/common/tvm_ffi_bridge.h @@ -16,7 +16,6 @@ #include #include #include -#include #include #include #include @@ -34,49 +33,6 @@ namespace transformer_engine { namespace tvm_ffi_bridge { -// All CuTeDSL kernels share this lock while their TVM-FFI entrypoints are first used. -inline std::mutex first_cutedsl_launch_mutex; - -// Cached copies of the lambda share this state for one compiled kernel. -struct TVMFFIKernelState { - static constexpr int kMaxDevices = 32; - explicit TVMFFIKernelState(tvm::ffi::Function function) - : function(std::move(function)), num_devices(cuda::num_devices()) { - NVTE_CHECK(num_devices <= kMaxDevices, - "Too many visible CUDA devices for CuTeDSL: ", num_devices); - // All devices start out unlaunched. - for (int device = 0; device < num_devices; ++device) { - launched[device].store(false, std::memory_order_relaxed); - } - } - - tvm::ffi::Function function; - int num_devices; - std::atomic launched[kMaxDevices]; -}; - -inline auto make_tvm_ffi_kernel(tvm::ffi::Function function) { - auto state = std::make_shared(std::move(function)); - return [state = std::move(state)](auto &&...args) -> tvm::ffi::Any { - const int device = cuda::current_device(); - NVTE_CHECK(device >= 0 && device < state->num_devices, "Invalid CUDA device index: ", device); - // If we never launch the kernel on this device, we need to make its initial launch serialized - if (!state->launched[device].load(std::memory_order_acquire)) { - std::lock_guard lock(first_cutedsl_launch_mutex); - // Check again in case another thread already launched it while we were waiting for the lock - if (!state->launched[device].load(std::memory_order_relaxed)) { - // Launch the kernel while holding the global mutex and mark it as launched for this device after we're done. - tvm::ffi::Any result = state->function(std::forward(args)...); - state->launched[device].store(true, std::memory_order_release); - return result; - } - } - return state->function(std::forward(args)...); - }; -} - -using TVMFFIKernel = decltype(make_tvm_ffi_kernel(std::declval())); - bool initialize_python_cutedsl_backend(); inline DLDataType convert_to_dltype(NVTEDType type) { @@ -236,12 +192,12 @@ namespace tvm_ffi_bridge { // This compiles + globally registers the kernel under `key`; returns whether // a kernel is now registered / the config is supported // -// - std::optional get_kernel() const -// Retrieves the possibly cached callable. If the config provides a +// - std::optional get_kernel() const +// Retrieves the possibly cached function. If the config provides a // `uint32_t to_id() const` method, it can reuse TVMFFIConfigCache with the // canonical implementation: // ``` -// std::optional get_kernel() const { +// std::optional get_kernel() const { // static TVMFFIConfigCache &cache = TVMFFIConfigCache::create(); // return cache.get_or_load(*this); // } @@ -251,7 +207,7 @@ namespace tvm_ffi_bridge { // policy may implement get_kernel() themselves. // // Note: TVMFFIConfigCache::create() intentionally gives its cache process lifetime. -// This prevents cached TVM-FFI function handles from being destroyed during +// This prevents cached tvm::ffi::Function handles from being destroyed during // static teardown, when Python or TVM-FFI runtime state may already have been // finalized. The OS reclaims the allocation when the process exits. class TVMFFIConfigCache; @@ -266,7 +222,7 @@ struct is_lazyloadable_config< std::enable_if_t::value>, std::enable_if_t (T::*)() const>::value>>> + std::optional (T::*)() const>::value>>> : std::true_type {}; } // namespace detail @@ -282,7 +238,7 @@ class TVMFFICentral { static_assert(detail::is_lazyloadable_config::value, "Config must define `std::string to_key() const`, " "`bool retrieve_func_from_python(const std::string&) const`, " - "and `std::optional get_kernel() const`."); + "and `std::optional get_kernel() const`."); if (!cutedsl_backend_enabled_.load(std::memory_order_relaxed)) { maybe_warn_not_chosen( "the CuTeDSL backend is disabled, so no kernel is available for " @@ -431,14 +387,12 @@ class TVMFFIConfigCache { // This requires the config to have a `uint32_t to_id() const` method, which returns an unique // identifier among all configs of this type. template - std::optional get_or_load(const Config &cfg) { + std::optional get_or_load(const Config &cfg) { TVMFFICentral ¢ral = TVMFFICentral::getInstance(); // Checked ahead of the cache so that toggling the backend off still disables - // already-resolved configs. + // already-resolved configs; load_tvm_ffi_function emits the "disabled" warning. if (!central.get_cutedsl_backend_enabled()) { - // Preserve the "disabled" warning from load_tvm_ffi_function. - central.load_tvm_ffi_function(cfg); - return std::nullopt; + return central.load_tvm_ffi_function(cfg); } // Otherwise try the cache first, and ask Python to compile/register the kernel if not found. const uint32_t id = cfg.to_id(); @@ -460,12 +414,8 @@ class TVMFFIConfigCache { // No other thread has loaded it, and none can load it now while I hold the write lock. std::optional fn = central.load_tvm_ffi_function(cfg); - std::optional kernel; - if (fn) { - kernel.emplace(make_tvm_ffi_kernel(std::move(*fn))); - } - map_.emplace(id, kernel); - return kernel; + map_.emplace(id, fn); + return fn; } private: @@ -473,7 +423,7 @@ class TVMFFIConfigCache { ~TVMFFIConfigCache() = default; std::shared_mutex mutex_; - std::unordered_map> map_; + std::unordered_map> map_; }; // Optionally emit a warning explaining why the CuTeDSL backend was not chosen for this config. From c8d0016d1801334cab9428ce74da89156cbad403 Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Sat, 3 Oct 2026 00:15:06 +0000 Subject: [PATCH 16/58] remove TVMFFIKernel Signed-off-by: Kaining Zhong --- .../common/cast/mxfp8/group_quantize_mxfp8_cutedsl.cuh | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8_cutedsl.cuh b/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8_cutedsl.cuh index edb674a09c0..92bfa47a88b 100644 --- a/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8_cutedsl.cuh +++ b/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8_cutedsl.cuh @@ -77,7 +77,7 @@ struct MXFP8GroupQuantConfig { (sm_arch << 22); } - std::optional get_kernel() const { + std::optional get_kernel() const { static TVMFFIConfigCache &cache = TVMFFIConfigCache::create(); return cache.get_or_load(*this); } @@ -212,7 +212,7 @@ inline bool mxfp8_group_quantize_cutedsl(const MXFP8GroupQuantConfig &config, return true; } - std::optional group_quant_func_opt = config.get_kernel(); + std::optional group_quant_func_opt = config.get_kernel(); if (!group_quant_func_opt.has_value()) { return false; } From 62e4fd7fa911138c895db33364905903c636f74d Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Mon, 5 Oct 2026 21:41:48 +0000 Subject: [PATCH 17/58] refactor Signed-off-by: Kaining Zhong --- tests/jax/test_mxfp8_cutedsl_backend.py | 10 +- .../mxfp8/test_mxfp8_cutedsl_backend.py | 32 +++--- .../cast/mxfp8/group_quantize_mxfp8.py | 34 +++--- .../mxfp8/group_quantize_mxfp8_cutedsl.cuh | 106 ++++++++---------- transformer_engine/common/tvm_ffi_bridge.h | 6 + 5 files changed, 90 insertions(+), 98 deletions(-) diff --git a/tests/jax/test_mxfp8_cutedsl_backend.py b/tests/jax/test_mxfp8_cutedsl_backend.py index a74baf93ec1..e9cfa0ad1c3 100644 --- a/tests/jax/test_mxfp8_cutedsl_backend.py +++ b/tests/jax/test_mxfp8_cutedsl_backend.py @@ -291,11 +291,11 @@ def test_dtypes(method, act_type, act_desc, fp8_dtype, in_dtype): # through the C API / PyTorch, since JAX has no wrappers for those operations. # (name, shape representation, input shape, optional device-resident row counts) GROUP_CASES = [ - ("single_member", "varying_first_dim", (1, 128, 128), None), - ("same_both_multichunk", "varying_first_dim", (3, 384, 384), None), - ("varying_first", "varying_first_dim", (768, 256), (128, 384, 256)), - ("varying_first_multichunk", "varying_first_dim", (1024, 384), (128, 256, 384, 256)), - ("varying_first_empty", "varying_first_dim", (512, 256), (128, 0, 384)), + ("single_member", "vfd", (1, 128, 128), None), + ("same_both_multichunk", "vfd", (3, 384, 384), None), + ("varying_first", "vfd", (768, 256), (128, 384, 256)), + ("varying_first_multichunk", "vfd", (1024, 384), (128, 256, 384, 256)), + ("varying_first_empty", "vfd", (512, 256), (128, 0, 384)), ] GROUP_Q_LAYOUTS = [QuantizeLayout.ROWWISE, QuantizeLayout.COLWISE, QuantizeLayout.ROWWISE_COLWISE] diff --git a/tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py b/tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py index 0847421383e..a9479a53cda 100644 --- a/tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py +++ b/tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py @@ -138,8 +138,7 @@ def set_cutedsl_backend(enabled): def test_enable_cutedsl_backend_after_import(): """Test if we can manually enable the CuTeDSL backend without enabling CuTeDSL backend in the beginning.""" - script = textwrap.dedent( - """ + script = textwrap.dedent(""" import ctypes import transformer_engine.pytorch # pylint: disable=unused-import @@ -157,8 +156,7 @@ def test_enable_cutedsl_backend_after_import(): setter(1) assert tvm_ffi.get_global_func(entrypoint, allow_missing=True) is not None - """ - ) + """) env = os.environ.copy() env["NVTE_ENABLE_CUTEDSL_BACKEND"] = "0" result = subprocess.run( @@ -380,24 +378,22 @@ def test_dtypes(swizzled, method, act, fp8_dtype, in_dtype): # NVTE_ENABLE_CUTEDSL_BACKEND=1. # (name, shape representation in the config key, per-member shapes) GROUP_CASES = [ - ("same_both", "same_both_dims", [(256, 512)] * 3), - ("single_member", "same_both_dims", [(384, 384)]), + ("same_both", "sbd", [(256, 512)] * 3), + ("single_member", "sbd", [(384, 384)]), # N is 32- but not 128-divisible, so the rowwise and colwise scales carry zeroed padding. - ("same_both_n96", "same_both_dims", [(128, 96)] * 2), - ("varying_first", "varying_first_dim", [(128, 256), (384, 256), (256, 256)]), - ("varying_first_n160", "varying_first_dim", [(128, 160), (256, 160)]), + ("same_both_n96", "sbd", [(128, 96)] * 2), + ("varying_first", "vfd", [(128, 256), (384, 256), (256, 256)]), + ("varying_first_n160", "vfd", [(128, 160), (256, 160)]), # Partial 32-element blocks (e.g. N=144) are covered by the C++ grouped tests; # the PyTorch MXFP8 quantizer requires dimensions divisible by 32. - ("varying_last", "varying_last_dim", [(256, 128), (256, 384), (256, 256)]), - ("varying_both", "varying_both_dims", [(128, 256), (256, 128), (384, 512)]), + ("varying_last", "vld", [(256, 128), (256, 384), (256, 256)]), + ("varying_both", "vbd", [(128, 256), (256, 128), (384, 512)]), # Multiple chunks in both dimensions, including a final half-width 128-column strip. - ("varying_both_multichunk", "varying_both_dims", [(128, 128), (256, 384), (384, 640)]), - ("varying_first_empty", "varying_first_dim", [(128, 256), (0, 256), (384, 256)]), - ("varying_both_empty", "varying_both_dims", [(128, 128), (0, 256), (256, 384)]), -] -SINGLE_TENSOR_GROUP_CASES = [ - c for c in GROUP_CASES if c[1] in ("same_both_dims", "varying_first_dim") + ("varying_both_multichunk", "vbd", [(128, 128), (256, 384), (384, 640)]), + ("varying_first_empty", "vfd", [(128, 256), (0, 256), (384, 256)]), + ("varying_both_empty", "vbd", [(128, 128), (0, 256), (256, 384)]), ] +SINGLE_TENSOR_GROUP_CASES = [c for c in GROUP_CASES if c[1] in ("sbd", "vfd")] get_group_case_id = lambda c: c[0] @@ -563,7 +559,7 @@ def test_group_noop(cutedsl): set_cutedsl_backend(False) if cutedsl: assert_group_cutedsl_registered( - "same_both_dims", torch.bfloat16, tex.DType.kFloat8E4M3, True, True, False, False + "sbd", torch.bfloat16, tex.DType.kFloat8E4M3, True, True, False, False ) for name, first_bytes in first.items(): assert torch.equal(skipped[name], first_bytes), f"{name} changed under a set noop flag" diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8.py index 7d90d56da93..a6d6b82d1d8 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8.py @@ -53,10 +53,10 @@ NUM_WORKSPACE_SLOTS = NUM_TENSORMAPS + 1 # Shape representations, mirroring ShapeRepresentation in common/utils.cuh. -SAME_BOTH_DIMS = "same_both_dims" -VARYING_FIRST_DIM = "varying_first_dim" -VARYING_LAST_DIM = "varying_last_dim" -VARYING_BOTH_DIMS = "varying_both_dims" +SAME_BOTH_DIMS = "sbd" +VARYING_FIRST_DIM = "vfd" +VARYING_LAST_DIM = "vld" +VARYING_BOTH_DIMS = "vbd" SUPPORTED_SHAPE_REPS = (SAME_BOTH_DIMS, VARYING_FIRST_DIM, VARYING_LAST_DIM, VARYING_BOTH_DIMS) # Upper bound on the group size (MAX_SUPPORTED_TENSOR_DESCRIPTORS in grouped_tma.cuh); sizes @@ -270,7 +270,7 @@ def __call__( mO_col: Optional[cute.Tensor], mS_row: cute.Tensor, mS_col: cute.Tensor, - mOffsets: cute.Tensor, # int64[num_tensors + 1], CSR element offsets + mOffsets: Optional[cute.Tensor], # int64[num_tensors + 1], CSR element offsets mFirstDims: Optional[ cute.Tensor ], # int64[num_tensors] (VARYING_FIRST_DIM / VARYING_BOTH_DIMS) @@ -664,7 +664,7 @@ def kernel( self, mS_row: cute.Tensor, mS_col: cute.Tensor, - mOffsets: cute.Tensor, + mOffsets: Optional[cute.Tensor], mFirstDims: Optional[cute.Tensor], mTensormaps: cute.Tensor, mNoop: cute.Pointer, @@ -715,7 +715,7 @@ def _kernel_main( self, mS_row: cute.Tensor, mS_col: cute.Tensor, - mOffsets: cute.Tensor, + mOffsets: Optional[cute.Tensor], mFirstDims: Optional[cute.Tensor], mTensormaps: cute.Tensor, mWorkspace: Optional[cute.Tensor], @@ -1381,6 +1381,18 @@ def compile_cutedsl_function_from_cfg(cfg: MXFP8GroupQuantizeConfig): else None ) + offsets_fake = ( + cute.runtime.make_fake_compact_tensor( + cutlass.Int64, + (cute.sym_int32(),), + stride_order=(0,), + memspace=cute.AddressSpace.gmem, + assumed_align=8, + ) + if cfg.SHAPE_REP != SAME_BOTH_DIMS + else None + ) + sm_count = HardwareInfo().get_device_multiprocessor_count() kernel_obj = MXFP8GroupQuantizeKernel(cfg, sm_count) return cute.compile( @@ -1408,13 +1420,7 @@ def compile_cutedsl_function_from_cfg(cfg: MXFP8GroupQuantizeConfig): memspace=cute.AddressSpace.gmem, assumed_align=4, ), - cute.runtime.make_fake_compact_tensor( # mOffsets - cutlass.Int64, - (cute.sym_int32(),), - stride_order=(0,), - memspace=cute.AddressSpace.gmem, - assumed_align=8, - ), + offsets_fake, # mOffsets first_dims_fake, # mFirstDims last_dims_fake, # mLastDims tensormaps_fake, # mTensormaps diff --git a/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8_cutedsl.cuh b/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8_cutedsl.cuh index 92bfa47a88b..96c51b0c241 100644 --- a/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8_cutedsl.cuh +++ b/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8_cutedsl.cuh @@ -33,13 +33,13 @@ using namespace tvm_ffi_bridge; inline const char *shape_rep_to_str(ShapeRepresentation shape_rep) { switch (shape_rep) { case ShapeRepresentation::SAME_BOTH_DIMS: - return "same_both_dims"; + return "sbd"; case ShapeRepresentation::VARYING_FIRST_DIM: - return "varying_first_dim"; + return "vfd"; case ShapeRepresentation::VARYING_LAST_DIM: - return "varying_last_dim"; + return "vld"; default: - return "varying_both_dims"; + return "vbd"; } } @@ -86,7 +86,7 @@ struct MXFP8GroupQuantConfig { // compiled and registered on a cache miss. std::string to_key() const { std::string key; - // longest: cutedsl_group_mxfp8_smXXX_BFloat16_Float8E4M3_1_1_varying_first_dim_1_1_1_0_dqgelu + // longest: cutedsl_group_mxfp8_smXXX_BFloat16_Float8E4M3_1_1_vfd_1_1_1_0_dqgelu key.reserve(96); key.append("cutedsl_group_mxfp8_sm") .append(std::to_string(sm_arch)) @@ -127,50 +127,35 @@ struct MXFP8GroupQuantConfig { } }; -// Descriptor slots per group member: input, rowwise output, colwise output, activation input, -// plus one carrying (rows, cols, base_elts). Mirrors NUM_WORKSPACE_SLOTS / BYTES_PER_TENSORMAP -// in CuTeDSL/cast/mxfp8/group_quantize_mxfp8.py. +// kGroupTensorMapSlots is 5 slots for: input, rowwise output, colwise output, activation input, +// plus one carrying the metadata of tensor -- (rows, cols, base_elts) constexpr size_t kGroupTensorMapSlots = 5; constexpr size_t kInt64PerTensorMap = 128 / sizeof(int64_t); constexpr size_t kMaxGroupTensors = static_cast(dispatch::common::MAX_SUPPORTED_TENSOR_DESCRIPTORS); -struct alignas(128) GroupDescriptorWorkspace { +// We need to use int64_t here instead of CUtensorMap so we can pass this through tvm-ffi boundary +struct alignas(128) TensorMapStorage { alignas(128) int64_t tensor_maps[kMaxGroupTensors][kGroupTensorMapSlots][kInt64PerTensorMap]; - // Stand-in for the unused offsets array in SAME_BOTH_DIMS. The kernel takes - // offsets unconditionally but does not read them for this representation. - int64_t unused_offsets[kMaxGroupTensors + 1]; }; +static __device__ TensorMapStorage g_group_descriptor_workspace; -// Like `g_tensor_maps` on the CUDA path, this has internal linkage, so every translation -// unit including this header gets its own copy. It shares that path's caveat that two -// grouped quantize calls in flight on different streams would overwrite each other's -// descriptors. -static __device__ GroupDescriptorWorkspace g_group_descriptor_workspace; - -// Device address of this translation unit's workspace on the current device. The address -// is per device (each device context loads its own copy of the module), so it is cached -// per device. `static` rather than `inline`: it refers to the internal-linkage symbol above, -// so each translation unit needs its own definition and cache. -static GroupDescriptorWorkspace *group_descriptor_workspace_ptr() { - static std::vector cache(cuda::num_devices(), nullptr); +static TensorMapStorage *group_descriptor_workspace_ptr() { + // Each device has its own workspace for descriptors + static std::vector cache(cuda::num_devices(), nullptr); static std::vector flags(cuda::num_devices()); const int device_id = cuda::current_device(); NVTE_CHECK(0 <= device_id && device_id < cuda::num_devices(), "invalid CUDA device ID"); + // Copy the device symbol address on the current device into the cache on its first use only std::call_once(flags[device_id], [&]() { void *p = nullptr; NVTE_CHECK_CUDA(cudaGetSymbolAddress(&p, g_group_descriptor_workspace)); - cache[device_id] = static_cast(p); + cache[device_id] = static_cast(p); }); + // Return the cached device pointer for the current device to the host return cache[device_id]; } -inline NVTEBasicTensor make_basic_tensor(void *dptr, DType dtype, - const std::vector &shape) { - return NVTEBasicTensor{dptr, static_cast(dtype), - nvte_make_shape(shape.data(), shape.size())}; -} - // Signature mirrors mxfp8::group_quantize (input, act_input, noop, output, dbias, workspace, // stream). Returns false to fall back to the CUDA kernel. inline bool mxfp8_group_quantize_cutedsl(const MXFP8GroupQuantConfig &config, @@ -183,19 +168,16 @@ inline bool mxfp8_group_quantize_cutedsl(const MXFP8GroupQuantConfig &config, const size_t first_logical_dim = input_tensor->logical_shape.data[0]; const size_t last_logical_dim = input_tensor->logical_shape.data[1]; - // The kernel is compiled with cute.sym_int32(divisibility=...) on both logical extents, - // so a violating shape would silently mis-tile rather than fail. These mirror sym_M / - // sym_N in CuTeDSL/cast/mxfp8/group_quantize_mxfp8.py -- the DSL kernel's chunk height and - // the 16-byte TMA row alignment. The logical shape of VARYING_BOTH_DIMS is [1, total]. - constexpr size_t kChunkDimY = 128; - constexpr size_t kLastDimAlignment = 16; - const bool first_dim_tiles = config.shape_rep == ShapeRepresentation::VARYING_BOTH_DIMS || - first_logical_dim % kChunkDimY == 0; - if (!first_dim_tiles || last_logical_dim % kLastDimAlignment != 0) { - maybe_warn_cutedsl_not_chosen("the grouped logical shape is not a multiple of (", kChunkDimY, - ", ", kLastDimAlignment, ")."); + // Match the symbolic divisibility checks in group_quantize_mxfp8.py. + if (config.shape_rep != ShapeRepresentation::VARYING_BOTH_DIMS && first_logical_dim % 128 != 0) { + maybe_warn_cutedsl_not_chosen("the first logical dimension is not divisible by 128 for a non-varying both dimensions tensor."); + return false; + } + if (last_logical_dim % 16 != 0) { + maybe_warn_cutedsl_not_chosen("the last logical dimension is not divisible by 16."); return false; } + // The same extents are sym_int32 in the compiled kernel. if (first_logical_dim > static_cast(INT32_MAX) || last_logical_dim > static_cast(INT32_MAX)) { @@ -203,11 +185,11 @@ inline bool mxfp8_group_quantize_cutedsl(const MXFP8GroupQuantConfig &config, return false; } - // dbias workspace-size query, mirroring mxfp8::group_quantize: the framework first calls - // with an unallocated workspace to learn its shape, allocates it, then calls again to run. - // The kernel writes one partial-dbias row per 128-row chunk. + // How many rows a job processes (see the CuTeDSL kernel) + constexpr size_t kRowsPerJob = 128; + // For dbias workspace-size query if (config.with_dbias && workspace_tensor->data.dptr == nullptr) { - workspace_tensor->data.shape = {DIVUP(first_logical_dim, kChunkDimY), last_logical_dim}; + workspace_tensor->data.shape = {DIVUP(first_logical_dim, kRowsPerJob), last_logical_dim}; workspace_tensor->data.dtype = DType::kFloat32; return true; } @@ -218,7 +200,7 @@ inline bool mxfp8_group_quantize_cutedsl(const MXFP8GroupQuantConfig &config, } const int32_t device_index = transformer_engine::cuda::current_device(); - GroupDescriptorWorkspace *const workspace = group_descriptor_workspace_ptr(); + TensorMapStorage *const workspace = group_descriptor_workspace_ptr(); const SimpleTensor &scale_row = config.rowwise ? output_tensor->scale_inv : output_tensor->columnwise_scale_inv; @@ -230,6 +212,7 @@ inline bool mxfp8_group_quantize_cutedsl(const MXFP8GroupQuantConfig &config, DLTensorWrapper mX( make_basic_tensor(input_tensor->data.dptr, input_tensor->dtype(), logical_shape), true, device_index); + DLTensorWrapper mO_row, mO_col; if (config.rowwise) { mO_row = DLTensorWrapper( @@ -250,12 +233,14 @@ inline bool mxfp8_group_quantize_cutedsl(const MXFP8GroupQuantConfig &config, false, device_index); // Offsets and member dims are read from the output, as in mxfp8::group_quantize. - auto offsets_or_unused = [&](const SimpleTensor &t, size_t numel) { - void *dptr = t.has_data() ? t.dptr : static_cast(workspace->unused_offsets); - return DLTensorWrapper(make_basic_tensor(dptr, DType::kInt64, {numel}), false, device_index); - }; - DLTensorWrapper mOffsets = offsets_or_unused(output_tensor->tensor_offsets, num_tensors + 1); - DLTensorWrapper mFirstDims, mLastDims; + DLTensorWrapper mOffsets, mFirstDims, mLastDims; + if (config.shape_rep != ShapeRepresentation::SAME_BOTH_DIMS) { + NVTE_CHECK(output_tensor->tensor_offsets.has_data(), "Grouped MXFP8 quantization with ", + shape_rep_to_str(config.shape_rep), " requires an allocated tensor_offsets buffer."); + mOffsets = DLTensorWrapper( + make_basic_tensor(output_tensor->tensor_offsets.dptr, DType::kInt64, {num_tensors + 1}), + false, device_index); + } if (config.shape_rep == ShapeRepresentation::VARYING_FIRST_DIM || config.shape_rep == ShapeRepresentation::VARYING_BOTH_DIMS) { NVTE_CHECK(output_tensor->first_dims.has_data(), "Grouped MXFP8 quantization with ", @@ -273,8 +258,7 @@ inline bool mxfp8_group_quantize_cutedsl(const MXFP8GroupQuantConfig &config, device_index); } - // The kernel reads num_tensors off this tensor's leading extent, so it must be exactly - // the group size even on the single-tensor path that leaves the descriptors untouched. + // Pass tensormaps as a 3D tensor of int64_t DLTensorWrapper mTensormaps( make_basic_tensor(static_cast(workspace->tensor_maps), DType::kInt64, {num_tensors, kGroupTensorMapSlots, kInt64PerTensorMap}), @@ -291,11 +275,8 @@ inline bool mxfp8_group_quantize_cutedsl(const MXFP8GroupQuantConfig &config, mWorkspace = DLTensorWrapper(workspace_tensor->data, true, device_index); } - // The cast-noop flag travels as a raw device pointer (not a tensor): it may be null, and the - // kernel null-checks it on device, so one compiled kernel serves both cases. void *noop_ptr = (noop_tensor != nullptr) ? noop_tensor->data.dptr : nullptr; - // noop and stream are tvm-ffi opaque "handles"; pass them as void*. (*group_quant_func_opt)(&mX, &mO_row, &mO_col, &mS_row, &mS_col, &mOffsets, &mFirstDims, &mLastDims, &mTensormaps, noop_ptr, &mActInput, &mWorkspace, static_cast(stream)); @@ -359,6 +340,7 @@ bool mxfp8_group_quantize_cutedsl(const GroupedTensor *input_tensor, kMaxGroupTensors, "."); return false; } + if (shape_rep == ShapeRepresentation::SAME_BOTH_DIMS) { // The kernel tiles the stacked rows without tensor boundaries, which matches the CUDA // kernel's per-tensor tiling only when every member's rows are a multiple of its @@ -373,6 +355,7 @@ bool mxfp8_group_quantize_cutedsl(const GroupedTensor *input_tensor, maybe_warn_cutedsl_not_chosen("the grouped tensor has no tensor offsets."); return false; } + if (IS_DBIAS && !is_single_tensor) { // mxfp8::group_quantize raises a proper error for this. maybe_warn_cutedsl_not_chosen("dbias is only supported for a common last dimension."); @@ -381,12 +364,13 @@ bool mxfp8_group_quantize_cutedsl(const GroupedTensor *input_tensor, const bool rowwise = output_tensor->has_data(); const bool colwise = output_tensor->has_columnwise_data(); + const bool swizzled = output_tensor->with_gemm_swizzled_scales; + + // Some Sanity checks if (!rowwise && !colwise) { - // mxfp8::group_quantize raises a proper error for this. + maybe_warn_cutedsl_not_chosen("the grouped tensor has neither rowwise nor columnwise data."); return false; } - const bool swizzled = output_tensor->with_gemm_swizzled_scales; - // Sanity checks, mirroring mxfp8::group_quantize checkCuDriverContext(stream); CheckNoopTensor(*noop_tensor, "cast_noop"); if (rowwise) { diff --git a/transformer_engine/common/tvm_ffi_bridge.h b/transformer_engine/common/tvm_ffi_bridge.h index d57fd8f03da..d2f24925844 100644 --- a/transformer_engine/common/tvm_ffi_bridge.h +++ b/transformer_engine/common/tvm_ffi_bridge.h @@ -64,6 +64,12 @@ inline DLDataType convert_to_dltype(NVTEDType type) { } } +inline NVTEBasicTensor make_basic_tensor(void *dptr, DType dtype, + const std::vector &shape) { + return NVTEBasicTensor{dptr, static_cast(dtype), + nvte_make_shape(shape.data(), shape.size())}; +} + class DLTensorWrapper : public DLTensor { public: static constexpr int kMaxNDim = sizeof(NVTEShape::data) / sizeof(NVTEShape::data[0]); From 6720ae2af0b226078f1b61aa80f20e3ce8848963 Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Tue, 6 Oct 2026 01:07:08 +0000 Subject: [PATCH 18/58] reuse grouped_tma.cuh's TensorMap struct Signed-off-by: Kaining Zhong --- .../cast/mxfp8/group_quantize_mxfp8.py | 237 ++++++++++-------- .../CuTeDSL/cast/mxfp8/quantize_mxfp8.py | 20 +- .../mxfp8/group_quantize_mxfp8_cutedsl.cuh | 92 ++++--- 3 files changed, 189 insertions(+), 160 deletions(-) diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8.py index a6d6b82d1d8..ed83e7a4b57 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8.py @@ -43,14 +43,12 @@ logger = logging.getLogger("transformer_engine.cutedsl.mxfp8") THREADS_PER_WARP = 32 -BYTES_PER_TENSORMAP = 128 -# Descriptor slots per tensor: input, rowwise output, colwise output, activation input. +# CUDA stores four descriptor arrays: input, act_input, output_rowwise, output_colwise. NUM_TENSORMAPS = 4 -ACT_INPUT_SLOT = 3 -# One extra slot holds per-tensor (rows, cols, base_elts), so the main kernel never has to -# binary-search the offsets array. Mirrors TensorMapStorage::rows/cols/offsets upstream. -META_SLOT = NUM_TENSORMAPS -NUM_WORKSPACE_SLOTS = NUM_TENSORMAPS + 1 +INPUT_SLOT = 0 +ACT_INPUT_SLOT = 1 +ROWWISE_OUTPUT_SLOT = 2 +COLWISE_OUTPUT_SLOT = 3 # Shape representations, mirroring ShapeRepresentation in common/utils.cuh. SAME_BOTH_DIMS = "sbd" @@ -62,6 +60,15 @@ # Upper bound on the group size (MAX_SUPPORTED_TENSOR_DESCRIPTORS in grouped_tma.cuh); sizes # the fixed binary search over the offsets. MAX_SUPPORTED_TENSORS = 64 +# How many bytes is a TensorMap +BYTES_PER_TENSORMAP = 128 +INT64_PER_TENSORMAP = BYTES_PER_TENSORMAP // 8 +# The byte offset of fields in the TensorMap struct +ROWS_OFFSET = NUM_TENSORMAPS * MAX_SUPPORTED_TENSORS * INT64_PER_TENSORMAP +COLS_OFFSET = ROWS_OFFSET + MAX_SUPPORTED_TENSORS +OFFSETS_OFFSET = COLS_OFFSET + MAX_SUPPORTED_TENSORS +# Size of the TensorMap struct in int64 +TENSOR_MAP_STORAGE_INT64_COUNT = OFFSETS_OFFSET + MAX_SUPPORTED_TENSORS class MXFP8GroupQuantizeConfig: @@ -224,7 +231,7 @@ def _scale_tensor(self, mS: cute.Tensor, base: Int64, layout: cute.Layout) -> cu @cute.jit def _rowwise_scales( - self, mS_row: cute.Tensor, base: Int64, rows: Int32, cols: Int32 + self, mS_row: cute.Tensor, base: Int64, rows: Int64, cols: Int64 ) -> cute.Tensor: """Rowwise scales of a (rows, cols) tensor at `base`, tiled per 32x128 stage.""" if cutlass.const_expr(self.cfg.WITH_GEMM_SWIZZLED_SCALES): @@ -243,7 +250,7 @@ def _rowwise_scales( @cute.jit def _colwise_scales( - self, mS_col: cute.Tensor, base: Int64, rows: Int32, cols: Int32 + self, mS_col: cute.Tensor, base: Int64, rows: Int64, cols: Int64 ) -> cute.Tensor: """Colwise scales of a (rows, cols) tensor at `base`, tiled per 32x128 stage.""" if cutlass.const_expr(self.cfg.WITH_GEMM_SWIZZLED_SCALES): @@ -277,7 +284,8 @@ def __call__( mLastDims: Optional[ cute.Tensor ], # int64[num_tensors] (VARYING_LAST_DIM / VARYING_BOTH_DIMS) - mTensormaps: cute.Tensor, # int64[num_tensors, NUM_WORKSPACE_SLOTS, 16] + mTensormaps: Optional[cute.Tensor], # flat int64 view of CUDA TensorMapStorage + num_tensors: Int32, mNoop: cute.Pointer, # f32 cast_noop flag; may be null, checked on device mActInput: Optional[cute.Tensor], # activation input, only with WITH_DACT mWorkspace: Optional[cute.Tensor], # f32 partial dbias, only with WITH_DBIAS @@ -294,7 +302,6 @@ def __call__( first_logical_dim == 1, "VARYING_BOTH_DIMS requires logical shape [1, total]" ) - num_tensors = mTensormaps.shape[0] runtime_assert(num_tensors > 0, "Grouped quantization requires at least one tensor") # A TMA atom copies a TILE at a time @@ -330,44 +337,32 @@ def __call__( if cutlass.const_expr(cfg.IS_SINGLE_TENSOR): # How many CTAs does the grouped tensor have in both directions - jobs_Y = cute.ceil_div(Int32(first_logical_dim), (self.TILE_ROWS * self.NUM_TILES_Y)) - jobs_X = cute.ceil_div(Int32(last_logical_dim), self.TILE_COLS * self.NUM_TILES_X) + jobs_Y = cute.ceil_div(Int64(first_logical_dim), (self.TILE_ROWS * self.NUM_TILES_Y)) + jobs_X = cute.ceil_div(Int64(last_logical_dim), self.TILE_COLS * self.NUM_TILES_X) # Flatten it to an 1D grid grid = [jobs_X * jobs_Y, 1, 1] else: # A placeholder for the kernel signature only; we won't use it in non-single tensor cases jobs_X = None # Estimate the total jobs across the group; each job has NUM_TILES_X * NUM_TILES_Y tiles - if cutlass.const_expr(cfg.SHAPE_REP == VARYING_BOTH_DIMS): - # Note: when VARYING_BOTH_DIMS, the first_logical_dim must be 1 - estimated_jobs = cute.ceil_div( - Int32(first_logical_dim) * Int32(last_logical_dim), self.ELTS_PER_CTA - ) - elif cutlass.const_expr(cfg.SHAPE_REP == VARYING_LAST_DIM): - # Same as VARYING_BOTH_DIMS but we divide 128 before multiplying to avoid overflowing Int32 - # because the first logical dimension is always 128-aligned when not VARYING_BOTH_DIMS - # so they are equivalent - estimated_jobs = cute.ceil_div( - (Int32(first_logical_dim) // 128) * Int32(last_logical_dim), - self.ELTS_PER_CTA // 128, - ) - else: - raise ValueError(f"unexpected shape representation {cfg.SHAPE_REP!r}") + estimated_jobs = cute.ceil_div( + Int64(first_logical_dim) * Int64(last_logical_dim), self.ELTS_PER_CTA + ) # Divide the persistent worker budget evenly across tensors, with at least # one worker per tensor. Each worker may process several jobs. requested_CTAs_per_tensor = cutlass.max( - Int32(1), - Int32(self.SM_COUNT * self.STATIC_PERSISTENT_WORKERS_PER_SM) // Int32(num_tensors), + Int64(1), + Int64(self.SM_COUNT * self.STATIC_PERSISTENT_WORKERS_PER_SM) // Int64(num_tensors), ) # In average how many jobs per tensor (only an average, the actual jobs per tensor may vary) average_jobs_per_tensor = cutlass.max( - Int32(1), cute.ceil_div(estimated_jobs, Int32(num_tensors)) + Int64(1), cute.ceil_div(estimated_jobs, Int64(num_tensors)) ) # Don't launch more CTAs than the average jobs per tensor in case # STATIC_PERSISTENT_WORKERS_PER_SM causes redundancy CTAs_per_tensor = cutlass.min(requested_CTAs_per_tensor, average_jobs_per_tensor) - grid = [CTAs_per_tensor, Int32(num_tensors), 1] + grid = [Int32(CTAs_per_tensor), num_tensors, 1] # Only the multi-tensor representations need per-tensor descriptors. if cutlass.const_expr(not cfg.IS_SINGLE_TENSOR): @@ -427,8 +422,8 @@ def _update_descriptors_kernel( mFirstDims: Optional[cute.Tensor], mLastDims: Optional[cute.Tensor], mTensormaps: cute.Tensor, - first_logical_dim: Int32, - last_logical_dim: Int32, + first_logical_dim: Int64, + last_logical_dim: Int64, dtype: cutlass.Constexpr[Type[cutlass.Numeric]], tma_atom_x: cute.CopyAtom, tma_atom_orow: Optional[cute.CopyAtom], @@ -437,12 +432,8 @@ def _update_descriptors_kernel( ) -> None: """Update the per-tensor TMA descriptors for the group quantization kernel. - mTensormaps: int64[num_tensors, NUM_WORKSPACE_SLOTS, 16], where the slots are: - - 0: input tensor - - 1: rowwise output tensor - - 2: colwise output tensor - - 3: activation input tensor (only with WITH_DACT) - - 4: metadata (rows, cols, base_offset) + mTensormaps is a flat int64 view of CUDA's TensorMapStorage: four descriptor + arrays followed by the rows, cols, and offsets arrays. """ cfg = self.cfg @@ -450,13 +441,13 @@ def _update_descriptors_kernel( tensor_id, _, _ = cute.arch.block_idx() # Figure out how many rows and columns this tensor has, and where its first element is in the group. if cutlass.const_expr(cfg.SHAPE_REP in (VARYING_FIRST_DIM, VARYING_BOTH_DIMS)): - rows = Int32(mFirstDims[tensor_id]) + rows = Int64(mFirstDims[tensor_id]) else: - rows = Int32(first_logical_dim) + rows = Int64(first_logical_dim) if cutlass.const_expr(cfg.SHAPE_REP in (VARYING_LAST_DIM, VARYING_BOTH_DIMS)): - cols = Int32(mLastDims[tensor_id]) + cols = Int64(mLastDims[tensor_id]) else: - cols = Int32(last_logical_dim) + cols = Int64(last_logical_dim) base_offset = Int64(mOffsets[tensor_id]) # Same diagnostics as get_tensor_rows_num / get_tensor_cols_num. Like NVTE_DEVICE_ERROR @@ -476,17 +467,28 @@ def _update_descriptors_kernel( tensor_id, ) - meta = mTensormaps[(tensor_id, META_SLOT, None)] - meta[0] = Int64(rows) - meta[1] = Int64(cols) - meta[2] = base_offset + mTensormaps[ROWS_OFFSET + tensor_id] = Int64(rows) + mTensormaps[COLS_OFFSET + tensor_id] = Int64(cols) + mTensormaps[OFFSETS_OFFSET + tensor_id] = base_offset tmap = TensorMapManager(TensorMapUpdateMode.GMEM, BYTES_PER_TENSORMAP) # Obtain the pointers of these descriptors - desc_x = tmap.get_tensormap_ptr(mTensormaps[(tensor_id, 0, None)].iterator) - desc_orow = tmap.get_tensormap_ptr(mTensormaps[(tensor_id, 1, None)].iterator) - desc_ocol = tmap.get_tensormap_ptr(mTensormaps[(tensor_id, 2, None)].iterator) - desc_act = tmap.get_tensormap_ptr(mTensormaps[(tensor_id, ACT_INPUT_SLOT, None)].iterator) + desc_x = tmap.get_tensormap_ptr( + mTensormaps.iterator + + (INPUT_SLOT * MAX_SUPPORTED_TENSORS + tensor_id) * INT64_PER_TENSORMAP + ) + desc_orow = tmap.get_tensormap_ptr( + mTensormaps.iterator + + (ROWWISE_OUTPUT_SLOT * MAX_SUPPORTED_TENSORS + tensor_id) * INT64_PER_TENSORMAP + ) + desc_ocol = tmap.get_tensormap_ptr( + mTensormaps.iterator + + (COLWISE_OUTPUT_SLOT * MAX_SUPPORTED_TENSORS + tensor_id) * INT64_PER_TENSORMAP + ) + desc_act = tmap.get_tensormap_ptr( + mTensormaps.iterator + + (ACT_INPUT_SLOT * MAX_SUPPORTED_TENSORS + tensor_id) * INT64_PER_TENSORMAP + ) # Zero-sized groups: creating a descriptor with a zero extent is invalid, # so skip (the main kernel skips these tensors as well). @@ -666,13 +668,13 @@ def kernel( mS_col: cute.Tensor, mOffsets: Optional[cute.Tensor], mFirstDims: Optional[cute.Tensor], - mTensormaps: cute.Tensor, + mTensormaps: Optional[cute.Tensor], mNoop: cute.Pointer, mWorkspace: Optional[cute.Tensor], - first_logical_dim: Int32, - last_logical_dim: Int32, + first_logical_dim: Int64, + last_logical_dim: Int64, num_tensors: Int32, - jobs_X: Optional[Int32], + jobs_X: Optional[Int64], dtype: cutlass.Constexpr[Type[cutlass.Numeric]], tma_atom_x: cute.CopyAtom, tma_src: cute.Tensor, @@ -717,12 +719,12 @@ def _kernel_main( mS_col: cute.Tensor, mOffsets: Optional[cute.Tensor], mFirstDims: Optional[cute.Tensor], - mTensormaps: cute.Tensor, + mTensormaps: Optional[cute.Tensor], mWorkspace: Optional[cute.Tensor], - first_logical_dim: Int32, - last_logical_dim: Int32, + first_logical_dim: Int64, + last_logical_dim: Int64, num_tensors: Int32, - jobs_X: Optional[Int32], + jobs_X: Optional[Int64], dtype: cutlass.Constexpr[Type[cutlass.Numeric]], tma_atom_x: cute.CopyAtom, tma_src: cute.Tensor, @@ -810,27 +812,27 @@ def _kernel_main( # If the CTA has work to do has_work = Boolean(True) # Metadata of the tensor that owns this job - tensor_rows = Int32(0) - tensor_cols = Int32(0) + tensor_rows = Int64(0) + tensor_cols = Int64(0) # Element offset of this tensor within the group: Int64 (CUDA uses size_t), since a # group can exceed 2^31 elements even when every individual extent is small. tensor_base = Int64(0) # Job's starting row and id in this individual tensor / global single tensor - job_start_row = Int32(0) - job_id_X = Int32(0) + job_start_row = Int64(0) + job_id_X = Int64(0) - first_job_id = Int32(0) - jobs_in_tensor = Int32(1) - job_stride = Int32(1) - jobs_X_in_tensor = Int32(1) + first_job_id = Int64(0) + jobs_in_tensor = Int64(1) + job_stride = Int64(1) + jobs_X_in_tensor = Int64(1) if cutlass.const_expr(cfg.IS_SINGLE_TENSOR): # grid = [jobs_X * jobs_Y, 1, 1] - job_id_Y = Int32(bidx) // jobs_X - job_id_X = Int32(bidx) % jobs_X + job_id_Y = Int64(bidx) // jobs_X + job_id_X = Int64(bidx) % jobs_X # View the grouped tensor as a single tensor of shape (first_logical_dim, last_logical_dim) - tensor_rows = Int32(first_logical_dim) - tensor_cols = Int32(last_logical_dim) + tensor_rows = Int64(first_logical_dim) + tensor_cols = Int64(last_logical_dim) # Which row does this job start from job_start_row = job_id_Y * (self.TILE_ROWS * self.NUM_TILES_Y) if cutlass.const_expr(cfg.SHAPE_REP == VARYING_FIRST_DIM): @@ -841,13 +843,12 @@ def _kernel_main( # When SAME_BOTH_DIM, M is always divisible by 128, which is exactly TILE_ROWS * NUM_TILES_Y, # so no need to check for the last job's starting row being beyond the last row of the tensor. else: - # grid = [workers_per_tensor, Int32(num_tensors), 1] + # grid = [workers_per_tensor, num_tensors, 1] tensor_id = Int32(bidy) # Extract tensor's metadata - meta = mTensormaps[(tensor_id, META_SLOT, None)] - tensor_rows = Int32(meta[0]) - tensor_cols = Int32(meta[1]) - tensor_base = Int64(meta[2]) + tensor_rows = Int64(mTensormaps[ROWS_OFFSET + tensor_id]) + tensor_cols = Int64(mTensormaps[COLS_OFFSET + tensor_id]) + tensor_base = Int64(mTensormaps[OFFSETS_OFFSET + tensor_id]) if tensor_rows > 0 and tensor_cols > 0: # How many jobs does this tensor have in both directions jobs_X_in_tensor = cute.ceil_div(tensor_cols, (self.TILE_COLS * self.NUM_TILES_X)) @@ -855,9 +856,9 @@ def _kernel_main( # How many jobs does this tensor have jobs_in_tensor = jobs_X_in_tensor * jobs_Y_in_tensor # Which job (1D index) does this CTA start from, which is also their worker ID - first_job_id = Int32(bidx) + first_job_id = Int64(bidx) # gdx is how many CTAs are assigned to this tensor - job_stride = Int32(gdx) + job_stride = Int64(gdx) # If my first job is already beyond the tensor's last job, I have no work to do if first_job_id >= jobs_in_tensor: has_work = Boolean(False) @@ -890,10 +891,10 @@ def _kernel_main( # Colwise swizzled scale indices restart at each member and depend # on its rows (process_colwise_stage), so address the member that # owns this job. - member_rows = Int32(0) - member_row_start = Int32(0) + member_rows = Int64(0) + member_row_start = Int64(0) if cutlass.const_expr(cfg.SHAPE_REP == SAME_BOTH_DIMS): - member_rows = tensor_rows // Int32(num_tensors) + member_rows = tensor_rows // Int64(num_tensors) member_row_start = job_start_row // member_rows * member_rows else: member_id = self._find_tensor_from_offsets( @@ -901,8 +902,8 @@ def _kernel_main( num_tensors, Int64(job_start_row) * Int64(tensor_cols), ) - member_rows = Int32(mFirstDims[member_id]) - member_row_start = Int32( + member_rows = Int64(mFirstDims[member_id]) + member_row_start = Int64( Int64(mOffsets[member_id]) // Int64(tensor_cols) ) col_scale_base = ( @@ -945,11 +946,23 @@ def _kernel_main( ) else: # For non-single tensor case, we use persistent kernel so each CTA keeps processing jobs until none is left - desc_x = tmap.get_tensormap_ptr(mTensormaps[(tensor_id, 0, None)].iterator) - desc_out_row = tmap.get_tensormap_ptr(mTensormaps[(tensor_id, 1, None)].iterator) - desc_out_col = tmap.get_tensormap_ptr(mTensormaps[(tensor_id, 2, None)].iterator) + desc_x = tmap.get_tensormap_ptr( + mTensormaps.iterator + + (INPUT_SLOT * MAX_SUPPORTED_TENSORS + tensor_id) * INT64_PER_TENSORMAP + ) + desc_out_row = tmap.get_tensormap_ptr( + mTensormaps.iterator + + (ROWWISE_OUTPUT_SLOT * MAX_SUPPORTED_TENSORS + tensor_id) + * INT64_PER_TENSORMAP + ) + desc_out_col = tmap.get_tensormap_ptr( + mTensormaps.iterator + + (COLWISE_OUTPUT_SLOT * MAX_SUPPORTED_TENSORS + tensor_id) + * INT64_PER_TENSORMAP + ) desc_act = tmap.get_tensormap_ptr( - mTensormaps[(tensor_id, ACT_INPUT_SLOT, None)].iterator + mTensormaps.iterator + + (ACT_INPUT_SLOT * MAX_SUPPORTED_TENSORS + tensor_id) * INT64_PER_TENSORMAP ) if tidx == 0: @@ -1023,8 +1036,8 @@ def _issue_load( self, pipeline_obj: pipeline.PipelineTmaAsync, prod_state: pipeline.PipelineState, - tile_y: Int32, - tile_x: Int32, + tile_y: Int64, + tile_x: Int64, atoms: tuple[Optional[cute.CopyAtom], ...], partitions: tuple[Optional[cute.Tensor], ...], tmap: TensorMapManager, @@ -1068,16 +1081,16 @@ def _issue_load( @cute.jit def _process_job( self, - job_start_row: Int32, # Row offset of this job (global for single-tensor, else tensor-local) - job_start_col: Int32, # Column offset of this job within the tensor - rows: Int32, # Rows of the rowwise-scale view (the group for single-tensor, else the tensor) - cols: Int32, # Number of columns in this tensor + job_start_row: Int64, # Row offset of this job (global for single-tensor, else tensor-local) + job_start_col: Int64, # Column offset of this job within the tensor + rows: Int64, # Rows of the rowwise-scale view (the group for single-tensor, else the tensor) + cols: Int64, # Number of columns in this tensor row_scales: Optional[ cute.Tensor ], # Rowwise scales tiled per stage, rows counted like job_start_row col_scales: Optional[cute.Tensor], # Colwise scales tiled per stage - col_scale_row_start: Int32, # Row of this job in the colwise-scale view - col_scale_rows: Int32, # Rows of the colwise-scale view + col_scale_row_start: Int64, # Row of this job in the colwise-scale view + col_scale_rows: Int64, # Rows of the colwise-scale view mWorkspace: Optional[cute.Tensor], # f32 partial dbias workspace (WITH_DBIAS) sDbias: Optional[ cute.Tensor @@ -1108,7 +1121,7 @@ def _process_job( job_tile_X = job_start_col // self.TILE_COLS col_scale_tile_Y = col_scale_row_start // self.TILE_ROWS tiles_X = cutlass.min( - Int32(self.NUM_TILES_X), cute.ceil_div(cols - job_start_col, self.TILE_COLS) + Int64(self.NUM_TILES_X), cute.ceil_div(cols - job_start_col, self.TILE_COLS) ) num_tiles = tiles_X * self.NUM_TILES_Y @@ -1289,10 +1302,10 @@ def compile_cutedsl_function_from_cfg(cfg: MXFP8GroupQuantizeConfig): # [1, total]. The last dim only needs the 16-byte TMA row alignment; a partial 32-element # scale block at the end of a row is zero-filled by TMA, as in the CUDA kernel. if cfg.SHAPE_REP == VARYING_BOTH_DIMS: - sym_M = cute.sym_int32() + sym_M = cute.sym_int64() else: - sym_M = cute.sym_int32(divisibility=128) - sym_N = cute.sym_int32(divisibility=SYM_N_DIVISIBILITY) + sym_M = cute.sym_int64(divisibility=128) + sym_N = cute.sym_int64(divisibility=SYM_N_DIVISIBILITY) logical_shape = (sym_M, sym_N) out_dtype = cfg.FP8_DTYPE @@ -1324,12 +1337,17 @@ def compile_cutedsl_function_from_cfg(cfg: MXFP8GroupQuantizeConfig): # The kernel only takes the base address of the scale buffers (per-tensor strides # are derived from cols), so their fake shape is a flat 1D byte run. - tensormaps_fake = cute.runtime.make_fake_compact_tensor( - cutlass.Int64, - (cute.sym_int32(), NUM_WORKSPACE_SLOTS, BYTES_PER_TENSORMAP // 8), - stride_order=(2, 1, 0), - memspace=cute.AddressSpace.gmem, - assumed_align=128, + # One compact view covers CUDA's descriptor and metadata arrays in their native order. + tensormaps_fake = ( + cute.runtime.make_fake_compact_tensor( + cutlass.Int64, + (TENSOR_MAP_STORAGE_INT64_COUNT,), + stride_order=(0,), + memspace=cute.AddressSpace.gmem, + assumed_align=128, + ) + if not cfg.IS_SINGLE_TENSOR + else None ) # The cast-noop flag is an always-present f32 pointer instead of an optional tensor, so # that one compiled kernel serves both an absent and a present flag (noop_flag_is_set). @@ -1348,7 +1366,7 @@ def compile_cutedsl_function_from_cfg(cfg: MXFP8GroupQuantizeConfig): workspace_fake = ( cute.runtime.make_fake_compact_tensor( Float32, - (cute.sym_int32(), cute.sym_int32()), + (cute.sym_int64(), cute.sym_int64()), stride_order=(1, 0), memspace=cute.AddressSpace.gmem, assumed_align=4, @@ -1408,14 +1426,14 @@ def compile_cutedsl_function_from_cfg(cfg: MXFP8GroupQuantizeConfig): out_col_fake, # mO_col cute.runtime.make_fake_compact_tensor( # mS_row scale_dtype, - (cute.sym_int32(),), + (cute.sym_int64(),), stride_order=(0,), memspace=cute.AddressSpace.gmem, assumed_align=4, ), cute.runtime.make_fake_compact_tensor( # mS_col scale_dtype, - (cute.sym_int32(),), + (cute.sym_int64(),), stride_order=(0,), memspace=cute.AddressSpace.gmem, assumed_align=4, @@ -1424,6 +1442,7 @@ def compile_cutedsl_function_from_cfg(cfg: MXFP8GroupQuantizeConfig): first_dims_fake, # mFirstDims last_dims_fake, # mLastDims tensormaps_fake, # mTensormaps + Int32(1), # num_tensors, supplied at runtime noop_fake, # mNoop act_input_fake, # mActInput workspace_fake, # mWorkspace diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8.py index c267638bbd7..34eac5e440e 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8.py @@ -115,8 +115,8 @@ @cute.jit def derive_swizzled_scale_layout( - M: Int32, - N: Int32, + M: Int32 | Int64, + N: Int32 | Int64, ROWWISE: cutlass.Constexpr[bool], COLWISE: cutlass.Constexpr[bool], mS_row: Optional[cute.Tensor], @@ -157,10 +157,10 @@ def quantize_rowwise_mxfp8( 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, + tile_row_start: Int32 | Int64, # global row index of this stage's row 0 + tile_col_start: Int32 | Int64, # global col index of this CTA's col 0 + M: Int32 | Int64, + N: Int32 | Int64, ACTIVATION: cutlass.Constexpr[str | None], DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], FP8_DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], @@ -413,10 +413,10 @@ def quantize_colwise_mxfp8( 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, + tile_row_start: Int32 | Int64, # global row index of this stage's row 0 + tile_col_start: Int32 | Int64, # global col index of this CTA's col 0 + M: Int32 | Int64, + N: Int32 | Int64, ACTIVATION: cutlass.Constexpr[str | None], DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], FP8_DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], diff --git a/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8_cutedsl.cuh b/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8_cutedsl.cuh index 96c51b0c241..4722935a1af 100644 --- a/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8_cutedsl.cuh +++ b/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8_cutedsl.cuh @@ -21,8 +21,9 @@ #include "../../tvm_ffi_bridge.h" #include "../../util/cuda_runtime.h" #include "../../util/cutedsl_utils.h" -#include "../../utils.cuh" // ShapeRepresentation -#include "../core/common.cuh" // MAX_SUPPORTED_TENSOR_DESCRIPTORS, grouped_reduce_dbias +#include "../../utils.cuh" // ShapeRepresentation +#include "../core/common.cuh" // grouped_reduce_dbias +#include "../core/grouped_tma.cuh" // TensorMapStorage namespace transformer_engine { namespace cutedsl_backend { @@ -127,32 +128,41 @@ struct MXFP8GroupQuantConfig { } }; -// kGroupTensorMapSlots is 5 slots for: input, rowwise output, colwise output, activation input, -// plus one carrying the metadata of tensor -- (rows, cols, base_elts) -constexpr size_t kGroupTensorMapSlots = 5; -constexpr size_t kInt64PerTensorMap = 128 / sizeof(int64_t); -constexpr size_t kMaxGroupTensors = - static_cast(dispatch::common::MAX_SUPPORTED_TENSOR_DESCRIPTORS); - -// We need to use int64_t here instead of CUtensorMap so we can pass this through tvm-ffi boundary -struct alignas(128) TensorMapStorage { - alignas(128) int64_t tensor_maps[kMaxGroupTensors][kGroupTensorMapSlots][kInt64PerTensorMap]; -}; -static __device__ TensorMapStorage g_group_descriptor_workspace; - -static TensorMapStorage *group_descriptor_workspace_ptr() { - // Each device has its own workspace for descriptors +using dispatch::common::TensorMapStorage; +constexpr size_t kGroupTensorMapSlots = 4; +constexpr size_t kMaxGroupTensors = dispatch::common::MAX_SUPPORTED_TENSOR_DESCRIPTORS; + +// Make sure TensorMapStorage from grouped_tma.cuh still works with CuTeDSL. +// If someone modified that, they need to fix how the CuTeDSL path uses it too. +static_assert(sizeof(CUtensorMap) == 128); +static_assert(sizeof(size_t) == sizeof(int64_t)); +static_assert(offsetof(TensorMapStorage, act_input) == kMaxGroupTensors * sizeof(CUtensorMap)); +static_assert(offsetof(TensorMapStorage, output_rowwise) == + 2 * kMaxGroupTensors * sizeof(CUtensorMap)); +static_assert(offsetof(TensorMapStorage, output_colwise) == + 3 * kMaxGroupTensors * sizeof(CUtensorMap)); +static_assert(offsetof(TensorMapStorage, rows) == + kGroupTensorMapSlots * kMaxGroupTensors * sizeof(CUtensorMap)); +static_assert(sizeof(TensorMapStorage) == + kGroupTensorMapSlots * kMaxGroupTensors * sizeof(CUtensorMap) + + 3 * kMaxGroupTensors * sizeof(size_t)); +static_assert(offsetof(TensorMapStorage, cols) - offsetof(TensorMapStorage, rows) == + kMaxGroupTensors * sizeof(size_t)); +static_assert(offsetof(TensorMapStorage, offsets) - offsetof(TensorMapStorage, rows) == + 2 * kMaxGroupTensors * sizeof(size_t)); + +// Cache the address of CUDA's device workspace once per device. Internal linkage +// keeps this cache paired with the translation unit's static g_tensor_maps symbol. +static TensorMapStorage *group_tensor_map_storage_ptr() { static std::vector cache(cuda::num_devices(), nullptr); static std::vector flags(cuda::num_devices()); const int device_id = cuda::current_device(); NVTE_CHECK(0 <= device_id && device_id < cuda::num_devices(), "invalid CUDA device ID"); - // Copy the device symbol address on the current device into the cache on its first use only std::call_once(flags[device_id], [&]() { - void *p = nullptr; - NVTE_CHECK_CUDA(cudaGetSymbolAddress(&p, g_group_descriptor_workspace)); - cache[device_id] = static_cast(p); + void *ptr = nullptr; + NVTE_CHECK_CUDA(cudaGetSymbolAddress(&ptr, dispatch::common::g_tensor_maps)); + cache[device_id] = static_cast(ptr); }); - // Return the cached device pointer for the current device to the host return cache[device_id]; } @@ -170,7 +180,9 @@ inline bool mxfp8_group_quantize_cutedsl(const MXFP8GroupQuantConfig &config, // Match the symbolic divisibility checks in group_quantize_mxfp8.py. if (config.shape_rep != ShapeRepresentation::VARYING_BOTH_DIMS && first_logical_dim % 128 != 0) { - maybe_warn_cutedsl_not_chosen("the first logical dimension is not divisible by 128 for a non-varying both dimensions tensor."); + maybe_warn_cutedsl_not_chosen( + "the first logical dimension is not divisible by 128 for a non-varying both dimensions " + "tensor."); return false; } if (last_logical_dim % 16 != 0) { @@ -178,13 +190,6 @@ inline bool mxfp8_group_quantize_cutedsl(const MXFP8GroupQuantConfig &config, return false; } - // The same extents are sym_int32 in the compiled kernel. - if (first_logical_dim > static_cast(INT32_MAX) || - last_logical_dim > static_cast(INT32_MAX)) { - maybe_warn_cutedsl_not_chosen("the grouped logical shape does not fit in int32."); - return false; - } - // How many rows a job processes (see the CuTeDSL kernel) constexpr size_t kRowsPerJob = 128; // For dbias workspace-size query @@ -200,7 +205,6 @@ inline bool mxfp8_group_quantize_cutedsl(const MXFP8GroupQuantConfig &config, } const int32_t device_index = transformer_engine::cuda::current_device(); - TensorMapStorage *const workspace = group_descriptor_workspace_ptr(); const SimpleTensor &scale_row = config.rowwise ? output_tensor->scale_inv : output_tensor->columnwise_scale_inv; @@ -258,11 +262,18 @@ inline bool mxfp8_group_quantize_cutedsl(const MXFP8GroupQuantConfig &config, device_index); } - // Pass tensormaps as a 3D tensor of int64_t - DLTensorWrapper mTensormaps( - make_basic_tensor(static_cast(workspace->tensor_maps), DType::kInt64, - {num_tensors, kGroupTensorMapSlots, kInt64PerTensorMap}), - false, device_index); + DLTensorWrapper mTensormaps; + const bool is_single_tensor = config.shape_rep == ShapeRepresentation::SAME_BOTH_DIMS || + config.shape_rep == ShapeRepresentation::VARYING_FIRST_DIM; + if (!is_single_tensor) { + TensorMapStorage *const ptr = group_tensor_map_storage_ptr(); + + // Keep CUDA's native layout: input, activation input, rowwise output, colwise output, + // then rows, cols, and offsets. CuTeDSL accesses these arrays by their int64 offsets. + mTensormaps = DLTensorWrapper(make_basic_tensor(static_cast(ptr), DType::kInt64, + {sizeof(TensorMapStorage) / sizeof(int64_t)}), + false, device_index); + } // Optional inputs: a wrapper over a null buffer packs as TVM-FFI None. DLTensorWrapper mActInput, mWorkspace; @@ -278,8 +289,8 @@ inline bool mxfp8_group_quantize_cutedsl(const MXFP8GroupQuantConfig &config, void *noop_ptr = (noop_tensor != nullptr) ? noop_tensor->data.dptr : nullptr; (*group_quant_func_opt)(&mX, &mO_row, &mO_col, &mS_row, &mS_col, &mOffsets, &mFirstDims, - &mLastDims, &mTensormaps, noop_ptr, &mActInput, &mWorkspace, - static_cast(stream)); + &mLastDims, &mTensormaps, static_cast(num_tensors), noop_ptr, + &mActInput, &mWorkspace, static_cast(stream)); // Reduce the per-chunk partial dbias per member with the CUDA kernel's reduction. if (config.with_dbias) { @@ -291,7 +302,7 @@ inline bool mxfp8_group_quantize_cutedsl(const MXFP8GroupQuantConfig &config, reinterpret_cast(output_tensor->tensor_offsets.dptr), reinterpret_cast(output_tensor->first_dims.dptr), reinterpret_cast(output_tensor->last_dims.dptr), dbias_tensor, - workspace_ptr, kChunkDimY, stream);) // NOLINT(*) + workspace_ptr, kRowsPerJob, stream);) // NOLINT(*) } return true; } @@ -332,8 +343,7 @@ bool mxfp8_group_quantize_cutedsl(const GroupedTensor *input_tensor, shape_rep == ShapeRepresentation::VARYING_FIRST_DIM; // Leave invalid group sizes to mxfp8::group_quantize, which raises a proper error. - // Every member gets a descriptor slot in the fixed-size workspace, so the CUDA - // kernel's descriptor limit applies to the single-tensor representations here too. + // Keep the same group-size limit as the CUDA implementation. const size_t num_tensors = input_tensor->num_tensors; if (num_tensors == 0 || num_tensors > kMaxGroupTensors) { maybe_warn_cutedsl_not_chosen("the group size ", num_tensors, " is not between 1 and ", From 4578caa9a22e745b48b5496a40ba6a719d0fbdf5 Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Tue, 6 Oct 2026 23:57:27 +0000 Subject: [PATCH 19/58] refactor Signed-off-by: Kaining Zhong --- .../cast/mxfp8/group_quantize_mxfp8.py | 172 +++++++++--------- 1 file changed, 84 insertions(+), 88 deletions(-) diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8.py index ed83e7a4b57..fbf2c6025da 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8.py @@ -344,11 +344,13 @@ def __call__( else: # A placeholder for the kernel signature only; we won't use it in non-single tensor cases jobs_X = None + + # Below is just some heuristic to decide how many CTAs we should use per tensor + # Estimate the total jobs across the group; each job has NUM_TILES_X * NUM_TILES_Y tiles estimated_jobs = cute.ceil_div( Int64(first_logical_dim) * Int64(last_logical_dim), self.ELTS_PER_CTA ) - # Divide the persistent worker budget evenly across tensors, with at least # one worker per tensor. Each worker may process several jobs. requested_CTAs_per_tensor = cutlass.max( @@ -812,11 +814,11 @@ def _kernel_main( # If the CTA has work to do has_work = Boolean(True) # Metadata of the tensor that owns this job - tensor_rows = Int64(0) - tensor_cols = Int64(0) + member_rows = Int64(0) + member_cols = Int64(0) # Element offset of this tensor within the group: Int64 (CUDA uses size_t), since a # group can exceed 2^31 elements even when every individual extent is small. - tensor_base = Int64(0) + member_base = Int64(0) # Job's starting row and id in this individual tensor / global single tensor job_start_row = Int64(0) job_id_X = Int64(0) @@ -830,9 +832,8 @@ def _kernel_main( # grid = [jobs_X * jobs_Y, 1, 1] job_id_Y = Int64(bidx) // jobs_X job_id_X = Int64(bidx) % jobs_X - # View the grouped tensor as a single tensor of shape (first_logical_dim, last_logical_dim) - tensor_rows = Int64(first_logical_dim) - tensor_cols = Int64(last_logical_dim) + # All members share the stacked tensor's column count. + member_cols = Int64(last_logical_dim) # Which row does this job start from job_start_row = job_id_Y * (self.TILE_ROWS * self.NUM_TILES_Y) if cutlass.const_expr(cfg.SHAPE_REP == VARYING_FIRST_DIM): @@ -846,13 +847,13 @@ def _kernel_main( # grid = [workers_per_tensor, num_tensors, 1] tensor_id = Int32(bidy) # Extract tensor's metadata - tensor_rows = Int64(mTensormaps[ROWS_OFFSET + tensor_id]) - tensor_cols = Int64(mTensormaps[COLS_OFFSET + tensor_id]) - tensor_base = Int64(mTensormaps[OFFSETS_OFFSET + tensor_id]) - if tensor_rows > 0 and tensor_cols > 0: + member_rows = Int64(mTensormaps[ROWS_OFFSET + tensor_id]) + member_cols = Int64(mTensormaps[COLS_OFFSET + tensor_id]) + member_base = Int64(mTensormaps[OFFSETS_OFFSET + tensor_id]) + if member_rows > 0 and member_cols > 0: # How many jobs does this tensor have in both directions - jobs_X_in_tensor = cute.ceil_div(tensor_cols, (self.TILE_COLS * self.NUM_TILES_X)) - jobs_Y_in_tensor = cute.ceil_div(tensor_rows, (self.TILE_ROWS * self.NUM_TILES_Y)) + jobs_X_in_tensor = cute.ceil_div(member_cols, (self.TILE_COLS * self.NUM_TILES_X)) + jobs_Y_in_tensor = cute.ceil_div(member_rows, (self.TILE_ROWS * self.NUM_TILES_Y)) # How many jobs does this tensor have jobs_in_tensor = jobs_X_in_tensor * jobs_Y_in_tensor # Which job (1D index) does this CTA start from, which is also their worker ID @@ -875,59 +876,50 @@ def _kernel_main( # For single tensor case we don't use tensor descriptors descs = (None, None, None, None) - # Rowwise scales span the whole group: members are stacked on 128-row - # boundaries, so each member's swizzled tiles follow the previous member's. + # Find how many rows this tensor member has and which row the member starts from + if cutlass.const_expr(cfg.SHAPE_REP == SAME_BOTH_DIMS): + member_rows = Int64(first_logical_dim) // Int64(num_tensors) + member_start_row = job_start_row - job_start_row % member_rows + elif cutlass.const_expr(cfg.SHAPE_REP == VARYING_FIRST_DIM): + member_id = self._find_tensor_from_offsets( + mOffsets, num_tensors, job_start_row * member_cols + ) + member_rows = Int64(mFirstDims[member_id]) + member_start_row = Int64(mOffsets[member_id]) // member_cols + else: + raise RuntimeError( + "is_single_tensor requires SAME_BOTH_DIMS or VARYING_FIRST_DIM" + ) + + # The starting offset for this tensor member's scale + member_scale_base = ( + member_start_row * cute.round_up(member_cols, 128) // MXFP8_BLOCK_SCALING_SIZE + ) + row_scales = None + # Starting data row of this job relative to the member + job_start_row_from_member = job_start_row - member_start_row + job_start_col_from_member = job_id_X * self.TILE_COLS if cutlass.const_expr(cfg.ROWWISE): - row_scales = self._rowwise_scales(mS_row, Int64(0), tensor_rows, tensor_cols) + row_scales = self._rowwise_scales( + mS_row, member_scale_base, member_rows, member_cols + ) - # Colwise scales require special handling for swizzled scales col_scales = None - col_scale_row_start = job_start_row - col_scale_rows = tensor_rows if cutlass.const_expr(cfg.COLWISE): - col_scale_base = Int64(0) - if cutlass.const_expr(cfg.WITH_GEMM_SWIZZLED_SCALES): - # Colwise swizzled scale indices restart at each member and depend - # on its rows (process_colwise_stage), so address the member that - # owns this job. - member_rows = Int64(0) - member_row_start = Int64(0) - if cutlass.const_expr(cfg.SHAPE_REP == SAME_BOTH_DIMS): - member_rows = tensor_rows // Int64(num_tensors) - member_row_start = job_start_row // member_rows * member_rows - else: - member_id = self._find_tensor_from_offsets( - mOffsets, - num_tensors, - Int64(job_start_row) * Int64(tensor_cols), - ) - member_rows = Int64(mFirstDims[member_id]) - member_row_start = Int64( - Int64(mOffsets[member_id]) // Int64(tensor_cols) - ) - col_scale_base = ( - Int64(member_row_start) - * Int64(cute.round_up(tensor_cols, 128)) - // MXFP8_BLOCK_SCALING_SIZE - ) - col_scale_row_start = job_start_row - member_row_start - col_scale_rows = member_rows col_scales = self._colwise_scales( - mS_col, col_scale_base, col_scale_rows, tensor_cols + mS_col, member_scale_base, member_rows, member_cols ) - cute.arch.sync_threads() self._process_job( - job_start_row, - job_id_X * self.TILE_COLS, - tensor_rows, - tensor_cols, + member_start_row, + job_start_row_from_member, + job_start_col_from_member, + member_rows, + member_cols, # all tensors have the same last dimension in single tensor case row_scales, col_scales, - col_scale_row_start, - col_scale_rows, mWorkspace, sDbias, descs, @@ -975,14 +967,18 @@ def _kernel_main( tmap.fence_tensormap_update(desc_out_col) descs = (desc_x, desc_act, desc_out_row, desc_out_col) - # This tensor's scales start at tensor_base / 32 in both directions. - scale_base = tensor_base // Int64(MXFP8_BLOCK_SCALING_SIZE) + # This member's scales start at member_base / 32 in both directions. + member_scale_base = member_base // Int64(MXFP8_BLOCK_SCALING_SIZE) row_scales = None if cutlass.const_expr(cfg.ROWWISE): - row_scales = self._rowwise_scales(mS_row, scale_base, tensor_rows, tensor_cols) + row_scales = self._rowwise_scales( + mS_row, member_scale_base, member_rows, member_cols + ) col_scales = None if cutlass.const_expr(cfg.COLWISE): - col_scales = self._colwise_scales(mS_col, scale_base, tensor_rows, tensor_cols) + col_scales = self._colwise_scales( + mS_col, member_scale_base, member_rows, member_cols + ) # Make sure all threads see the updated descriptors and scales before processing any jobs. cute.arch.sync_threads() @@ -993,19 +989,20 @@ def _kernel_main( while not job_finished: job_id_Y_in_tensor = job_id // jobs_X_in_tensor job_id_X_in_tensor = job_id % jobs_X_in_tensor - job_start_row_in_tensor = job_id_Y_in_tensor * ( + job_start_row_from_member = job_id_Y_in_tensor * ( self.TILE_ROWS * self.NUM_TILES_Y ) - job_start_col = job_id_X_in_tensor * (self.TILE_COLS * self.NUM_TILES_X) + job_start_col_from_member = job_id_X_in_tensor * ( + self.TILE_COLS * self.NUM_TILES_X + ) self._process_job( - job_start_row_in_tensor, - job_start_col, - tensor_rows, - tensor_cols, + Int64(0), # For non-single tensor case, each TMA loads one member + job_start_row_from_member, + job_start_col_from_member, + member_rows, + member_cols, row_scales, col_scales, - job_start_row_in_tensor, - tensor_rows, mWorkspace, sDbias, descs, @@ -1081,16 +1078,13 @@ def _issue_load( @cute.jit def _process_job( self, - job_start_row: Int64, # Row offset of this job (global for single-tensor, else tensor-local) - job_start_col: Int64, # Column offset of this job within the tensor - rows: Int64, # Rows of the rowwise-scale view (the group for single-tensor, else the tensor) - cols: Int64, # Number of columns in this tensor - row_scales: Optional[ - cute.Tensor - ], # Rowwise scales tiled per stage, rows counted like job_start_row + member_start_row: Int64, # Which row the member starts from the tensor loaded by TMA + job_start_row_from_member: Int64, # Rows between this job and the start of this member + job_start_col_from_member: Int64, # Columns between this job and the start of this member + member_rows: Int64, # Number of data rows in this member + member_cols: Int64, # Number of columns in this member + row_scales: Optional[cute.Tensor], # Member-local rowwise scales tiled per stage col_scales: Optional[cute.Tensor], # Colwise scales tiled per stage - col_scale_row_start: Int64, # Row of this job in the colwise-scale view - col_scale_rows: Int64, # Rows of the colwise-scale view mWorkspace: Optional[cute.Tensor], # f32 partial dbias workspace (WITH_DBIAS) sDbias: Optional[ cute.Tensor @@ -1117,11 +1111,13 @@ def _process_job( _, _, _, _, tXsO_row, tXgO_row, tXsO_col, tXgO_col = partitions _, _, desc_out_row, desc_out_col = descs + job_start_row = member_start_row + job_start_row_from_member job_tile_Y = job_start_row // self.TILE_ROWS - job_tile_X = job_start_col // self.TILE_COLS - col_scale_tile_Y = col_scale_row_start // self.TILE_ROWS + job_tile_X = job_start_col_from_member // self.TILE_COLS + member_tile_Y = job_start_row_from_member // self.TILE_ROWS tiles_X = cutlass.min( - Int64(self.NUM_TILES_X), cute.ceil_div(cols - job_start_col, self.TILE_COLS) + Int64(self.NUM_TILES_X), + cute.ceil_div(member_cols - job_start_col_from_member, self.TILE_COLS), ) num_tiles = tiles_X * self.NUM_TILES_Y @@ -1142,7 +1138,7 @@ def _process_job( for tile_X in cutlass.range(tiles_X, unroll=1): tile_id_X = job_tile_X + tile_X - tile_start_col = job_start_col + tile_X * self.TILE_COLS + tile_start_col = job_start_col_from_member + tile_X * self.TILE_COLS partial_dbias = Float32(0.0) dbias_row_acc = None if cutlass.const_expr(self.DBIAS_IN_ROWWISE): @@ -1174,12 +1170,12 @@ def _process_job( sX_tile, sAct_tile, sO_col[(None, cons_state.index)], - cute.flatten(col_scales[(None, (col_scale_tile_Y + tile_Y, tile_id_X))]), + cute.flatten(col_scales[(None, (member_tile_Y + tile_Y, tile_id_X))]), cfg.MAX_NORM_RCP, - (col_scale_tile_Y + tile_Y) * self.TILE_ROWS, + (member_tile_Y + tile_Y) * self.TILE_ROWS, tile_start_col, - col_scale_rows, - cols, + member_rows, + member_cols, ACTIVATION=cfg.ACTIVATION, DTYPE=cfg.DTYPE, FP8_DTYPE=cfg.FP8_DTYPE, @@ -1202,12 +1198,12 @@ def _process_job( sX_tile, None if self.CACHE_ACTIVATION else sAct_tile, sO_row[(None, cons_state.index)], - cute.flatten(row_scales[(None, (row_tile, tile_id_X))]), + cute.flatten(row_scales[(None, (member_tile_Y + tile_Y, tile_id_X))]), cfg.MAX_NORM_RCP, - row_tile * self.TILE_ROWS, + (member_tile_Y + tile_Y) * self.TILE_ROWS, tile_start_col, - rows, - cols, + member_rows, + member_cols, ACTIVATION=None if self.CACHE_ACTIVATION else cfg.ACTIVATION, DTYPE=cfg.DTYPE, FP8_DTYPE=cfg.FP8_DTYPE, @@ -1291,7 +1287,7 @@ def _process_job( # A job has TILE_ROWS * NUM_TILES_Y rows, and dbias reduces them to one row dbias_row = job_start_row // (self.TILE_ROWS * self.NUM_TILES_Y) dbias_col = tile_start_col + tidx - if dbias_col < cols: + if dbias_col < member_cols: mWorkspace[(dbias_row, dbias_col)] = partial_dbias From 00dfd90634c1dff32aed2b20ea75c745ef5ee1fd Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 7 Oct 2026 01:43:50 +0000 Subject: [PATCH 20/58] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py b/tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py index a9479a53cda..f984ad4be36 100644 --- a/tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py +++ b/tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py @@ -138,7 +138,8 @@ def set_cutedsl_backend(enabled): def test_enable_cutedsl_backend_after_import(): """Test if we can manually enable the CuTeDSL backend without enabling CuTeDSL backend in the beginning.""" - script = textwrap.dedent(""" + script = textwrap.dedent( + """ import ctypes import transformer_engine.pytorch # pylint: disable=unused-import @@ -156,7 +157,8 @@ def test_enable_cutedsl_backend_after_import(): setter(1) assert tvm_ffi.get_global_func(entrypoint, allow_missing=True) is not None - """) + """ + ) env = os.environ.copy() env["NVTE_ENABLE_CUTEDSL_BACKEND"] = "0" result = subprocess.run( From 7c3abf6bdc9442b8645e4d639c603c1b8fc98cdc Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Wed, 7 Oct 2026 07:22:07 +0000 Subject: [PATCH 21/58] [Common] More CuTeDSL MXFP8 special kernels Signed-off-by: Kaining Zhong --- transformer_engine/common/CuTeDSL/__init__.py | 2 +- .../common/CuTeDSL/cast/mxfp8/__init__.py | 2 +- .../cast/mxfp8/quantize_mxfp8_common.py | 685 ++++++++ .../mxfp8/quantize_mxfp8_g2r_bidimensional.py | 490 ++++++ .../mxfp8/quantize_mxfp8_g2r_rowwise_1lane.py | 265 +++ .../mxfp8/quantize_mxfp8_g2r_rowwise_2lane.py | 368 ++++ .../mxfp8/quantize_mxfp8_g2s_bidimensional.py | 712 ++++++++ ...ize_mxfp8.py => quantize_mxfp8_general.py} | 1498 +---------------- transformer_engine/common/CuTeDSL/utils.py | 20 + .../common/CuTeDSL/utils_fp8.py | 43 + 10 files changed, 2603 insertions(+), 1482 deletions(-) create mode 100644 transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_common.py create mode 100644 transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py create mode 100644 transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_1lane.py create mode 100644 transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_2lane.py create mode 100644 transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2s_bidimensional.py rename transformer_engine/common/CuTeDSL/cast/mxfp8/{quantize_mxfp8.py => quantize_mxfp8_general.py} (50%) diff --git a/transformer_engine/common/CuTeDSL/__init__.py b/transformer_engine/common/CuTeDSL/__init__.py index 50d536852c2..e3920b13486 100644 --- a/transformer_engine/common/CuTeDSL/__init__.py +++ b/transformer_engine/common/CuTeDSL/__init__.py @@ -10,7 +10,7 @@ import tvm_ffi -from transformer_engine.common.CuTeDSL.cast.mxfp8.quantize_mxfp8 import ( +from transformer_engine.common.CuTeDSL.cast.mxfp8.quantize_mxfp8_common import ( get_mxfp8_quantization_function, ) diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/__init__.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/__init__.py index 6491573690f..444d85234ec 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/__init__.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/__init__.py @@ -4,4 +4,4 @@ """CuTeDSL MXFP8 quantization kernels.""" -from . import quantize_mxfp8 +from . import quantize_mxfp8_common diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_common.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_common.py new file mode 100644 index 00000000000..f281937ffca --- /dev/null +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_common.py @@ -0,0 +1,685 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Shared definitions, kernel dispatch and TVM-FFI registration for CuTeDSL MXFP8.""" + +import abc +import logging +import os +from typing import Optional + +import cutlass +from cutlass import cute +from cutlass import Boolean, Float32, Int32, Int64, Uint32 +from cuda.bindings.driver import CUstream # pylint: disable=no-name-in-module +import tvm_ffi + +from transformer_engine.common.CuTeDSL.utils import ( + str_to_cutlass_dtype, + abs_max_x2_bf16, + device_compute_capability, + is_packed16, +) +from transformer_engine.common.CuTeDSL.activations import ( + act_relu, + act_gelu, + act_silu, + act_qgelu, + act_srelu, + dact_drelu, + dact_dsrelu, + dact_dsilu, + dact_dqgelu, + dact_dgelu, +) + +CUTEDSL_DEBUG_LOGGING = os.environ.get("CUTEDSL_DEBUG_LOGGING", "0") == "1" + +logger = logging.getLogger("transformer_engine.cutedsl.mxfp8") + +# Number of elements per MXFP8 scale block. They will share the same E8M0 scale factor +MXFP8_BLOCK_SCALING_SIZE = 32 +# How many threads are in one warp +THREADS_PER_WARP = 32 + +# Prefer the current allocator API while retaining compatibility with older DSL releases. +_SmemAllocator = getattr(cutlass, "memory", cutlass.utils).SmemAllocator + +# FP8E4M3 max representable value +FP8E4M3_MAX_NORM = 448.0 +FP8E4M3_MAX_NORM_RCP = 1.0 / FP8E4M3_MAX_NORM +FP8E5M2_MAX_NORM = 57344.0 +FP8E5M2_MAX_NORM_RCP = 1.0 / FP8E5M2_MAX_NORM + +# If N is not divisible by 16, writing quantized FP8 output with N bytes per row will fail because +# TMA requires 16-byte alignment for each row +SYM_N_DIVISIBILITY = 16 + +SUPPORTED_ACTIVATIONS = { + "relu": act_relu, + "gelu": act_gelu, + "silu": act_silu, + "qgelu": act_qgelu, + "srelu": act_srelu, +} + +SUPPORTED_DACTIVATIONS = { + "drelu": dact_drelu, + "dgelu": dact_dgelu, + "dsilu": dact_dsilu, + "dqgelu": dact_dqgelu, + "dsrelu": dact_dsrelu, +} + + +@cute.jit +def derive_swizzled_scale_layout( + M: Int32, + N: Int32, + ROWWISE: cutlass.Constexpr[bool], + COLWISE: cutlass.Constexpr[bool], + mS_row: Optional[cute.Tensor], + mS_col: Optional[cute.Tensor], +): + """Derive the swizzled layout for the rowwise and colwise scale tensors.""" + num_scale_cols = cute.ceil_div(N, MXFP8_BLOCK_SCALING_SIZE) + num_scale_rows = cute.ceil_div(M, MXFP8_BLOCK_SCALING_SIZE) + + num_tiles_M = cute.ceil_div(M, 128) + num_tiles_SC = cute.ceil_div(num_scale_cols, 4) + num_tiles_SR = cute.ceil_div(num_scale_rows, 4) + num_tiles_N = cute.ceil_div(N, 128) + + if cutlass.const_expr(ROWWISE): + mS_row = cute.make_tensor( + mS_row.iterator, + cute.make_layout( + ((32, 4, num_tiles_M), (4, num_tiles_SC)), + stride=((16, 4, num_tiles_SC * 512), (1, 512)), + ), + ) + if cutlass.const_expr(COLWISE): + mS_col = cute.make_tensor( + mS_col.iterator, + cute.make_layout( + ((4, num_tiles_SR), (32, 4, num_tiles_N)), + stride=((1, 512), (16, 4, num_tiles_SR * 512)), + ), + ) + return mS_row, mS_col + + +@cute.jit +def noop_flag_is_set(mNoop: cute.Pointer) -> Boolean: + """Whether the cast_noop flag says this quantization is a no-op and must be skipped. + + mNoop is a pointer rather than a tensor so that one compiled kernel serves both a present and + an absent flag, hence the address is checked before it is dereferenced, exactly like the CUDA + C++ kernel's `noop != nullptr && noop[0] == 1.0f`. The two checks cannot be joined with `and`, + which the DSL lowers to a non-short-circuiting op that would load from the null pointer. + """ + flag_is_set = Boolean(False) + if mNoop.toint() != Int64(0): + flag_is_set = cute.make_tensor(mNoop, cute.make_layout((1,)))[0] == Float32(1.0) + return flag_is_set + + +class MXFP8QuantizeConfig: + """Configs for the compiled CuTeDSL kernel. These will be fixed once the kernel is compiled and + they will behave as const expressions. + """ + + def __init__( + self, + dtype: str, + fp8_dtype: str, + rowwise: bool, + colwise: bool, + with_gemm_swizzled_scales: bool, + with_amax: bool, + with_dbias: bool = False, + with_dact: bool = False, + with_act: bool = False, + use_2d_quantization: bool = False, + activation: Optional[str] = None, + ): + if use_2d_quantization: + raise ValueError("2D block scaling is not implemented by the CuTeDSL MXFP8 kernels") + if dtype is None or dtype not in ("Float32", "Float16", "BFloat16"): + raise ValueError(f"unknown input dtype {dtype!r}; expected Float32|Float16|BFloat16") + self.DTYPE = str_to_cutlass_dtype(dtype) + self.DTYPE_STR = dtype # readable input-dtype token, for __str__ + if fp8_dtype not in ("Float8E4M3", "Float8E5M2"): + raise ValueError( + f"unknown FP8 dtype {fp8_dtype!r}; expected 'Float8E4M3' or 'Float8E5M2'" + ) + self.FP8_DTYPE = str_to_cutlass_dtype(fp8_dtype) + self.FP8_DTYPE_STR = fp8_dtype # readable token, for __str__ + self.ROWWISE = rowwise + self.COLWISE = colwise + if not (rowwise or colwise): + raise ValueError("at least one of rowwise or colwise must be true") + self.WITH_GEMM_SWIZZLED_SCALES = with_gemm_swizzled_scales + self.WITH_AMAX = with_amax + if not with_dact and not with_act: + if activation == "none": + self.ACTIVATION = None + else: + raise ValueError( + "activation must be none when with_dact and with_act are both False" + ) + else: + if with_dact and with_act: + raise ValueError( + "with_dact and with_act cannot be true at the same time since they are used for" + " different paths (bwd vs fwd)" + ) + if with_dact: + if activation in SUPPORTED_DACTIVATIONS: + self.ACTIVATION = activation + else: + raise ValueError( + f"unknown activation {activation!r} for with_dact=True; expected one of" + f" {sorted(SUPPORTED_DACTIVATIONS)}" + ) + elif with_act: + if activation in SUPPORTED_ACTIVATIONS: + self.ACTIVATION = activation + else: + raise ValueError( + f"unknown activation {activation!r} for with_act=True; expected one of" + f" {sorted(SUPPORTED_ACTIVATIONS)}" + ) + self.WITH_DACT = with_dact + self.WITH_ACT = with_act + # dbias is the column reduction of the (post-act/dact) element. With colwise + # output each thread owns a full column (trivial reduction); rowwise-only + # uses a cross-thread smem reduction over THREADS_Y. Both mirror the CUDA + # kernel's COLWISE_SCALING / rowwise dbias branches. + self.WITH_DBIAS = with_dbias + self.MAX_NORM_RCP = ( + FP8E4M3_MAX_NORM_RCP if fp8_dtype == "Float8E4M3" else FP8E5M2_MAX_NORM_RCP + ) + + def __str__(self): + return ( + f"MXFP8QuantizeConfig(dtype={self.DTYPE_STR}, fp8_dtype={self.FP8_DTYPE_STR}, " + f"rowwise={self.ROWWISE}, colwise={self.COLWISE}, " + f"swizzled={self.WITH_GEMM_SWIZZLED_SCALES}, with_amax={self.WITH_AMAX}, " + f"with_dbias={self.WITH_DBIAS}, with_dact={self.WITH_DACT}, " + f"with_act={self.WITH_ACT}, activation={self.ACTIVATION})" + ) + + __repr__ = __str__ + + +class MXFP8QuantizeKernelBase(abc.ABC): + """Base class for MXFP8 quantize kernels.""" + + @abc.abstractmethod + def __call__( + self, + mX: cute.Tensor, + mO_row: Optional[cute.Tensor], + mS_row: Optional[cute.Tensor], + mO_col: Optional[cute.Tensor], + mS_col: Optional[cute.Tensor], + mAmax: Optional[cute.Tensor], + mNoop: cute.Pointer, + mDActInput: Optional[cute.Tensor], + mWorkspace: Optional[cute.Tensor], + stream: CUstream, + ): + """ + Compiled kernel entrypoint (decorate with @cute.jit). + All MXFP8 quantize kernels must implement this interface because this is our C++ call site's contract in `quantize_mxfp8_cutedsl.cuh`. + C++ call site will pass the arguments in this exact order via tvm-ffi and the kernel must accept them even if they are not used. + """ + + +@cute.jit +def bf16_pair_magnitude(pair: Int32) -> Uint32: + """Fold both BF16 halves and clear the sign set by max.xorsign.abs.""" + swapped = cute.arch.inline_ptx( + "prmt.b32 {$w0}, {$r0}, {$r0}, 0x1032;", + write_only_types=[Int32], + read_only_args=[pair], + ) + return Uint32(abs_max_x2_bf16(pair, swapped)) & Uint32(0x7FFF) + + +class MXFP8QuantizeEntry(MXFP8QuantizeKernelBase): + """Select the appropriate MXFP8 quantization kernel based on the configuration and runtime shapes.""" + + def __init__(self, cfg: MXFP8QuantizeConfig): + # Kernels import shared definitions from this module. Resolve their + # implementations only when constructing the compile-time entry object. + # pylint: disable=import-outside-toplevel + from .quantize_mxfp8_general import MXFP8QuantizeKernel + from .quantize_mxfp8_g2r_rowwise_1lane import MXFP8QuantizeG2RRowwise1LaneKernel + from .quantize_mxfp8_g2s_bidimensional import MXFP8QuantizeSpecializedBidimensionalKernel + from .quantize_mxfp8_g2r_rowwise_2lane import MXFP8QuantizeG2RRowwise2LaneKernel + from .quantize_mxfp8_g2r_bidimensional import MXFP8QuantizeRegisterBidimensionalKernel + + # pylint: enable=import-outside-toplevel + + self.cfg = cfg + # These activation functions satisfy f(0) = 0 so we don't need to mask with OOB regions + # (zeros filled by TMA are still zeros without applying activation to them) + # Note that as the time of writing, all activations supported by MXFP8 satisfy this property, so !ACT_NEED_MASKING is never reachable. + # This is only a precaution for future use cases if someone wants to add a new activation function. + self.ACT_NEED_MASKING = cfg.WITH_ACT and cfg.ACTIVATION not in ( + "relu", + "gelu", + "silu", + "qgelu", + "srelu", + ) + # Instantiate all possible kernels at compile time, + # and we will pick the right one at runtime based on the input shape and config. + self.general_skip_input_masking_kernel = MXFP8QuantizeKernel(cfg, SKIP_INPUT_MASKING=True) + self.general_skip_input_masking_and_scale_bounds_kernel = MXFP8QuantizeKernel( + cfg, SKIP_INPUT_MASKING=True, SKIP_SCALE_BOUNDS=True + ) + # We only need to mask when WITH_ACT is enabled because with activations applied zeros filled by TMA affect block statistics + self.general_kernel = MXFP8QuantizeKernel(cfg) if self.ACT_NEED_MASKING else None + self.g2r_rowwise_1lane = MXFP8QuantizeG2RRowwise1LaneKernel(cfg) if cfg.ROWWISE else None + self.g2s_bidim = ( + MXFP8QuantizeSpecializedBidimensionalKernel(cfg) + if cfg.ROWWISE and cfg.COLWISE + else None + ) + self.g2r_rowwise_2lane = MXFP8QuantizeG2RRowwise2LaneKernel(cfg) if cfg.ROWWISE else None + self.g2r_bidim = ( + MXFP8QuantizeRegisterBidimensionalKernel(cfg) if cfg.ROWWISE and cfg.COLWISE else None + ) + + @cute.jit + def __call__( + self, + mX: cute.Tensor, + mO_row: Optional[cute.Tensor], + mS_row: Optional[cute.Tensor], + mO_col: Optional[cute.Tensor], + mS_col: Optional[cute.Tensor], + mAmax: Optional[cute.Tensor], + mNoop: cute.Pointer, + mDActInput: Optional[cute.Tensor], + mWorkspace: Optional[cute.Tensor], + stream: CUstream, + ): + M = mX.shape[0] + N = mX.shape[1] + # Select specialized kernels if possible + plain_cast_only = ( + not self.cfg.WITH_AMAX + and not self.cfg.WITH_DBIAS + and not self.cfg.WITH_DACT + and not self.cfg.WITH_ACT + ) + dispatched_to_specialized = False + # The two-lane rowwise and g2r bidimensional kernels use BF16 packed math and 32-byte loads. + # Keep the existing kernels for FP16, fusions, ragged bidimensional tiles, + # misaligned input views and compilers predating 256-bit global loads. + if cutlass.const_expr( + plain_cast_only + and self.cfg.DTYPE is cutlass.BFloat16 + and cutlass.target_version(min_version="12.9") + ): + input_is_aligned = mX.iterator.toint() % 32 == 0 + if cutlass.const_expr(self.cfg.ROWWISE and not self.cfg.COLWISE): + scale_stride_matches = Boolean(True) + if cutlass.const_expr(not self.cfg.WITH_GEMM_SWIZZLED_SCALES): + scale_stride_matches = mS_row.shape[1] == N // 32 + # Column tiles use CUDA's grid.y axis. Inputs this wide always use two tiles. + grid_cols_fit = ( + N + <= 65535 + * self.g2r_rowwise_2lane._CTA_THREADS_X + * self.g2r_rowwise_2lane._ELEMENTS_PER_THREAD + * 2 + ) + if N % 128 == 0 and input_is_aligned and scale_stride_matches and grid_cols_fit: + dispatched_to_specialized = True + self.g2r_rowwise_2lane( + mX, + mO_row, + mS_row, + mO_col, + mS_col, + mAmax, + mNoop, + mDActInput, + mWorkspace, + stream, + ) + if cutlass.const_expr( + self.cfg.ROWWISE and self.cfg.COLWISE and not self.cfg.WITH_GEMM_SWIZZLED_SCALES + ): + columnwise_scales_are_aligned = mS_col.iterator.toint() % 16 == 0 + if ( + M % 32 == 0 + and N % 256 == 0 + and M // 32 <= 65535 + and input_is_aligned + and columnwise_scales_are_aligned + ): + dispatched_to_specialized = True + self.g2r_bidim( + mX, + mO_row, + mS_row, + mO_col, + mS_col, + mAmax, + mNoop, + mDActInput, + mWorkspace, + stream, + ) + # The one-lane rowwise and g2s bidimensional kernels support packed16 types (bf16/fp16). + if cutlass.const_expr(plain_cast_only and is_packed16(self.cfg.DTYPE)): + if cutlass.const_expr(self.cfg.ROWWISE and not self.cfg.COLWISE): + g2r_rowwise_1lane_grid_fits = ( + cute.ceil_div(M, self.g2r_rowwise_1lane._TILE_ROWS) <= 65535 + ) + # The one-lane rowwise kernel requires N divisible by 128 for vectorized stores. + if not dispatched_to_specialized and N % 128 == 0 and g2r_rowwise_1lane_grid_fits: + dispatched_to_specialized = True + self.g2r_rowwise_1lane( + mX, + mO_row, + mS_row, + mO_col, + mS_col, + mAmax, + mNoop, + mDActInput, + mWorkspace, + stream, + ) + # The g2s bidimensional kernel supports swizzled scales but we don't dispatch to it for now + if cutlass.const_expr( + self.cfg.ROWWISE and self.cfg.COLWISE and not self.cfg.WITH_GEMM_SWIZZLED_SCALES + ): + g2s_bidim_grid_fits = ( + cute.ceil_div( + M, + self.g2s_bidim._TILE_ROWS * self.g2s_bidim._NUM_TILES_Y, + ) + <= 65535 + ) + if not dispatched_to_specialized and g2s_bidim_grid_fits: + dispatched_to_specialized = True + self.g2s_bidim( + mX, + mO_row, + mS_row, + mO_col, + mS_col, + mAmax, + mNoop, + mDActInput, + mWorkspace, + stream, + ) + # If not using a specialized kernel, fall back to the general kernel + if not dispatched_to_specialized: + # If the input shape can be perfectly tiled by the general kernel's tile size, we can skip some boundary check because + # we know we will not touch any out-of-bounds region. + shape_is_divisible = ( + mX.shape[0] % self.general_skip_input_masking_kernel._TILE_ROWS == 0 + and mX.shape[1] % self.general_skip_input_masking_kernel._TILE_COLS == 0 + ) + # If shape is divisible, then we won't access OOB regions regardless whatever + if shape_is_divisible: + self.general_skip_input_masking_and_scale_bounds_kernel( + mX, + mO_row, + mS_row, + mO_col, + mS_col, + mAmax, + mNoop, + mDActInput, + mWorkspace, + stream, + ) + else: + # Otherwise we can't skip scale bound check, but we may still skip masking activation inputs with zero + if cutlass.const_expr(self.ACT_NEED_MASKING): + # Masking and the scale guards now share one condition, so this is a plain + # two-way choice: either the shape tiles exactly (drop both) or it does not. + self.general_kernel( + mX, + mO_row, + mS_row, + mO_col, + mS_col, + mAmax, + mNoop, + mDActInput, + mWorkspace, + stream, + ) + else: + # We still need to check the scale bounds, but we can skip some activations masking because their output is 0 + # for OOB regions where TMA fills with zeros + self.general_skip_input_masking_kernel( + mX, + mO_row, + mS_row, + mO_col, + mS_col, + mAmax, + mNoop, + mDActInput, + mWorkspace, + stream, + ) + + +def compile_cutedsl_function_from_cfg(cfg): + """ + Return the compiled CuTeDSL function object for the given MXFP8 quantization config. + """ + + kernel_obj = MXFP8QuantizeEntry(cfg) + sym_M = cute.sym_int32() + sym_N = cute.sym_int32(divisibility=SYM_N_DIVISIBILITY) + in_shape = out_shape = (sym_M, sym_N) + # TE allocates scale tensors at a padded shape (see + # MXFP8Quantizer::get_scale_shape in transformer_engine/pytorch/csrc): + # rowwise: (roundup(M, 128), roundup(ceildiv(N, 32), 4)) + # columnwise: (roundup(ceildiv(M, 32), 4), roundup(N, 128)) + # These padded extents are NOT M/N (and SymInt has no `//`/`+`), so give the + # scales their own fresh syms carrying the divisibility the padding + # guarantees (rowwise: 128 x 4; colwise: 4 x 128). + scale_rowwise_shape = (cute.sym_int32(divisibility=128), cute.sym_int32(divisibility=4)) + scale_colwise_shape = (cute.sym_int32(divisibility=4), cute.sym_int32(divisibility=128)) + ws_shape = (cute.sym_int32(), sym_N) # (blocks_Y, N); N ties to input N + # Native FP8/E8M0 dtypes at the FFI boundary (matches the DLPack dtype the C++ + # bridge sends). + out_dtype = cfg.FP8_DTYPE + scale_dtype = cutlass.Float8E8M0FNU + + in_fake = cute.runtime.make_fake_compact_tensor( + cfg.DTYPE, in_shape, stride_order=(1, 0), memspace=cute.AddressSpace.gmem, assumed_align=16 + ) + out_row_fake = ( + cute.runtime.make_fake_compact_tensor( + out_dtype, + out_shape, + stride_order=(1, 0), + memspace=cute.AddressSpace.gmem, + assumed_align=16, + ) + if cfg.ROWWISE + else None + ) + scale_row_fake = ( + cute.runtime.make_fake_compact_tensor( + scale_dtype, + scale_rowwise_shape, + stride_order=(1, 0), + memspace=cute.AddressSpace.gmem, + assumed_align=4, + ) + if cfg.ROWWISE + else None + ) + out_col_fake = ( + cute.runtime.make_fake_compact_tensor( + out_dtype, + out_shape, + stride_order=(1, 0), + memspace=cute.AddressSpace.gmem, + assumed_align=16, + ) + if cfg.COLWISE + else None + ) + scale_col_fake = ( + cute.runtime.make_fake_compact_tensor( + scale_dtype, + scale_colwise_shape, + stride_order=(1, 0), + memspace=cute.AddressSpace.gmem, + assumed_align=4, + ) + if cfg.COLWISE + else None + ) + amax_fake = ( + cute.runtime.make_fake_compact_tensor( + Float32, (1,), stride_order=(0,), memspace=cute.AddressSpace.gmem, assumed_align=4 + ) + if cfg.WITH_AMAX + else None + ) + # The cast-noop flag is an always-present f32 pointer instead of an optional tensor, so that + # the same compiled kernel serves both an absent and a present flag (see noop_flag_is_set). + # The null address here is only a compile-time placeholder for the pointer the C++ dispatcher + # passes at launch. + noop_fake = cute.runtime.nullptr(Float32, mem_space=cute.AddressSpace.gmem, assumed_align=4) + act_input_fake = ( + cute.runtime.make_fake_compact_tensor( + cfg.DTYPE, + in_shape, + stride_order=(1, 0), + memspace=cute.AddressSpace.gmem, + assumed_align=16, + ) + if cfg.WITH_DACT + else None + ) + workspace_fake = ( + cute.runtime.make_fake_compact_tensor( + Float32, ws_shape, stride_order=(1, 0), memspace=cute.AddressSpace.gmem, assumed_align=4 + ) + if cfg.WITH_DBIAS + else None + ) + + compiled = cute.compile( + kernel_obj, + in_fake, # mX + out_row_fake, + scale_row_fake, # mO_row, mS_row + out_col_fake, + scale_col_fake, # mO_col, mS_col + amax_fake, # mAmax + noop_fake, # mNoop (pointer to the cast_noop flag, may be null at launch) + act_input_fake, # mDActInput (backward slot, unused) + workspace_fake, # mWorkspace(backward slot, unused) + cute.runtime.make_fake_stream(), # stream (compiled as an explicit tvm-ffi + # "handle" arg; C++ passes the CUDA stream + # as void*) + options="--enable-tvm-ffi", + ) + return compiled + + +def get_mxfp8_quantization_function( + fn_name: str, + dtype: str, + fp8_dtype: str, + rowwise: bool, + colwise: bool, + with_gemm_swizzled_scales: bool, + with_amax: bool, + with_dbias: bool, + with_dact: bool, + with_act: bool, + use_2d_quantization: bool, + activation: str, +) -> bool: + """Compile the MXFP8 quantize kernel for this config and register it in the TVM-FFI global registry + under EXACTLY `fn_name` (the key the C++ dispatcher built; Python treats it as an opaque name). + Returns True if a kernel is successfully registered under `fn_name` (the C++ side then fetches it with GetGlobal(fn_name)); + False if the config is unsupported, so the caller caches the negative result and falls back to the CUDA C++ kernel. + """ + try: + # Already registered (e.g. by a prior call) -> supported. + if tvm_ffi.get_global_func(fn_name, allow_missing=True) is not None: + return True + + major, minor = device_compute_capability() + if major < 10: + logger.warning( + "CuTeDSL MXFP8 backend requires compute capability >= 10.0 (Blackwell), " + "but detected %d.%d; falling back to the CUDA C++ kernel.", + major, + minor, + ) + return False + + try: + cfg = MXFP8QuantizeConfig( + dtype=dtype, + fp8_dtype=fp8_dtype, + rowwise=rowwise, + colwise=colwise, + with_gemm_swizzled_scales=with_gemm_swizzled_scales, + with_amax=with_amax, + with_dbias=with_dbias, + with_dact=with_dact, + with_act=with_act, + use_2d_quantization=use_2d_quantization, + activation=activation, + ) + except ValueError as e: + # The exception message states exactly why the config is unsupported + # (unknown dtype/activation, dbias not implemented, ...). Surfacing it as a + # warning lets the C++ dispatcher's CUDA fallback be recognized as expected. + logger.warning( + "CuTeDSL MXFP8 backend does not support this config, " + "falling back to the CUDA C++ kernel: %s", + e, + ) + return False + + logger.debug("Compiling CuTeDSL MXFP8 quantization kernel for %s", cfg) + compiled = compile_cutedsl_function_from_cfg(cfg) + # The returned compiled object is not neccessarily the compiled function itself. It could be a Python + # wrapper that parses arguments and calls the underlying function. This is needed in the general case + # as TVM-FFI is positional-only, however, for us this is unnecessary launch overhead as we pass the + # arguments in the correct order on the C++ side. + # Relevant issue: https://github.com/NVIDIA/cutlass/issues/3527 + # https://github.com/NVIDIA/cutlass/pull/3589 recommends using an explicit positional only marker (/) + # in the function signature, however it is not merged and it is something easy to miss and silently + # introduce unnecessarry overhead. + native = getattr(compiled, "__tvm_ffi_object__", lambda: None)() + tvm_ffi.register_global_func( + fn_name, native if native is not None else compiled, override=True + ) + return True + except Exception as e: # pylint: disable=broad-exception-caught + logger.error( + "CuTeDSL MXFP8 kernel compilation & registration failed, falling back to the CUDA" + " C++ kernel: %s", + e, + ) + # Unconditionally fallback to CUDA path because we can't tell if this exception is transient or permanent. + return False diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py new file mode 100644 index 00000000000..0713d68f0fd --- /dev/null +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py @@ -0,0 +1,490 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Register-resident CuTeDSL bidimensional MXFP8 cast, ported from cast_bidim.cuh.""" + +from typing import Optional + +import cutlass +from cutlass import cute +from cutlass import Int32, Int64, Uint32, Uint8 +from cuda.bindings.driver import CUstream # pylint: disable=no-name-in-module + +from transformer_engine.common.CuTeDSL.utils import ( + abs_max_x2_bf16, + create_l2_policy, +) +from transformer_engine.common.CuTeDSL.utils_fp8 import mul_bf16x2x2_cvt_bf16x4_to_fp8x4 +from .quantize_mxfp8_common import ( + MXFP8QuantizeConfig, + MXFP8QuantizeKernelBase, + _SmemAllocator, + bf16_pair_magnitude, + noop_flag_is_set, +) + + +@cute.jit +def bf16_mx_scale_reciprocal(magnitude: Uint32, MAX_EXPONENT: cutlass.Constexpr[int]) -> Uint32: + """Exact BF16 reciprocal scale bits, including E8M0 zero, Inf and NaN cases.""" + rounded = (magnitude + Uint32(31)) & Uint32(0x7F80) + if rounded < Uint32(MAX_EXPONENT << 7): # pylint: disable=consider-using-max-builtin + rounded = Uint32(MAX_EXPONENT << 7) + reciprocal = Uint32((254 + MAX_EXPONENT) << 7) - rounded + if magnitude >= Uint32(0x7F80): + reciprocal = Uint32(0x0040) + if magnitude != Uint32(0x7F80): + reciprocal = Uint32(0x7FFF) + return reciprocal + + +@cute.jit +def bf16_mx_scale_reciprocal_pair(pair: Int32, MAX_EXPONENT: cutlass.Constexpr[int]) -> Uint32: + """Compute two reciprocal scales without borrowing across packed BF16 halves.""" + magnitudes = Uint32(pair) & Uint32(0x7FFF7FFF) + rounded = (magnitudes + Uint32(0x001F001F)) & Uint32(0xFF80FF80) + floor = (MAX_EXPONENT << 23) | (MAX_EXPONENT << 7) + # Packed BF16 max accepts registers, so materialize the constant inside PTX. + clamped = cute.arch.inline_ptx( + "{ .reg.b32 floor_pair; " + f"mov.b32 floor_pair, {floor}; " + "max.xorsign.abs.bf16x2 {$w0}, {$r0}, floor_pair; }", + write_only_types=[Int32], + read_only_args=[rounded.bitcast(Int32)], + ) + bias = ((254 + MAX_EXPONENT) << 23) | ((254 + MAX_EXPONENT) << 7) + reciprocal = Uint32(bias) - Uint32(clamped) + if ((magnitudes + Uint32(0x00800080)) & Uint32(0x80008000)) != Uint32(0): + lo = bf16_mx_scale_reciprocal(magnitudes & Uint32(0x7FFF), MAX_EXPONENT) + hi = bf16_mx_scale_reciprocal(magnitudes >> Uint32(16), MAX_EXPONENT) + reciprocal = lo | (hi << Uint32(16)) + return reciprocal + + +@cute.jit +def bf16_mx_scale_bytes_pair(reciprocal: Uint32) -> Uint32: + """Gather the E8M0 bytes of two packed BF16 reciprocals into the low half.""" + bytes_pair = (Uint32(0x01FE01FE) - ((reciprocal >> Uint32(7)) & Uint32(0x00FF00FF))) & Uint32( + 0x00FF00FF + ) + return (bytes_pair & Uint32(0xFF)) | ((bytes_pair >> Uint32(8)) & Uint32(0xFF00)) + + +@cute.jit +def quantize_bidimensional_mxfp8_g2r( + gX_tile: cute.Tensor, + gO_row_tile: cute.Tensor, + gS_row_tile: cute.Tensor, + gO_col_tile: cute.Tensor, + gS_col_tile: cute.Tensor, + cfg: cutlass.Constexpr[MXFP8QuantizeConfig], + CTA_THREADS_Y: cutlass.Constexpr[int], + CTA_THREADS_X: cutlass.Constexpr[int], + ROWS_PER_THREAD: cutlass.Constexpr[int], + ELEMENTS_PER_THREAD: cutlass.Constexpr[int], + cache_policy: Int64, +): + """Quantize one (TILE_ROWS, TILE_COLS) GMEM tile in both directions. + + TILE_ROWS = CTA_THREADS_Y * ROWS_PER_THREAD = 32, so every columnwise + MXFP8 block is contained in the tile. Each warp owns ROWS_PER_THREAD + consecutive rows, and each lane owns ELEMENTS_PER_THREAD adjacent columns. + Load the tile once, then reuse its registers for both quantization passes. + """ + TILE_ROWS = CTA_THREADS_Y * ROWS_PER_THREAD + TILE_COLS = CTA_THREADS_X * ELEMENTS_PER_THREAD + LANES_PER_MX_BLOCK = 32 // ELEMENTS_PER_THREAD + COLUMN_PAIRS = TILE_COLS // 2 + COLUMN_PAIRS_PER_THREAD = ELEMENTS_PER_THREAD // 2 + PAIRS_PER_HALF = COLUMN_PAIRS_PER_THREAD // 2 + MAX_EXPONENT = 8 if cfg.FP8_DTYPE is cutlass.Float8E4M3FN else 15 + SCALE_DTYPE = cutlass.Float8E8M0FNU + tidx, _, _ = cute.arch.thread_idx() + warp, lane = tidx // CTA_THREADS_X, tidx % CTA_THREADS_X + + _, tv_payload = cute.make_layout_tv( + thr_layout=cute.make_layout((CTA_THREADS_Y, CTA_THREADS_X), stride=(CTA_THREADS_X, 1)), + val_layout=cute.make_layout( + (ROWS_PER_THREAD, ELEMENTS_PER_THREAD), stride=(ELEMENTS_PER_THREAD, 1) + ), + ) + # Reshape each thread's values into (columns, rows), with contiguous columns. + fragment_layout = cute.make_layout( + (ELEMENTS_PER_THREAD, ROWS_PER_THREAD), stride=(1, ELEMENTS_PER_THREAD) + ) + gX_thread = cute.composition(cute.composition(gX_tile, tv_payload)[tidx, None], fragment_layout) + gO_row_thread = cute.composition( + cute.composition(gO_row_tile, tv_payload)[tidx, None], fragment_layout + ) + gO_col_thread = cute.composition( + cute.composition(gO_col_tile, tv_payload)[tidx, None], fragment_layout + ) + # Each group of lanes shares one rowwise scale for each of its rows. + tv_row_scale = cute.make_layout( + ( + (LANES_PER_MX_BLOCK, (CTA_THREADS_X // LANES_PER_MX_BLOCK, CTA_THREADS_Y)), + ROWS_PER_THREAD, + ), + stride=((0, (TILE_ROWS, ROWS_PER_THREAD)), 1), + ) + gS_row_thread = cute.composition(gS_row_tile, tv_row_scale)[tidx, None] + # Warp 0 writes one scale per column, in each lane's original column order. + gS_col_thread = cute.local_tile(gS_col_tile, (1, ELEMENTS_PER_THREAD), (0, lane))[0, None] + + # Keep the complete input fragment; reuse one row's output registers. + rX = cute.make_rmem_tensor((ELEMENTS_PER_THREAD, ROWS_PER_THREAD), cfg.DTYPE) + rO = cute.make_rmem_tensor((ELEMENTS_PER_THREAD,), cfg.FP8_DTYPE) + rS_row = cute.make_rmem_tensor((1,), SCALE_DTYPE) + rS_col = cute.make_rmem_tensor((ELEMENTS_PER_THREAD,), SCALE_DTYPE) + load_atom = cute.make_copy_atom( + cute.nvgpu.CopyG2ROp(), + cfg.DTYPE, + num_bits_per_copy=ELEMENTS_PER_THREAD * cfg.DTYPE.width, + invariant=True, + ) + store_atom = cute.make_copy_atom( + cute.nvgpu.CopyR2GOp(), + cfg.FP8_DTYPE, + num_bits_per_copy=ELEMENTS_PER_THREAD * cfg.FP8_DTYPE.width, + ) + row_scale_atom = cute.make_copy_atom(cute.nvgpu.CopyR2GOp(), SCALE_DTYPE, num_bits_per_copy=8) + col_scale_atom = cute.make_copy_atom( + cute.nvgpu.CopyR2GOp(), SCALE_DTYPE, num_bits_per_copy=ELEMENTS_PER_THREAD * 8 + ) + scale_and_pack = mul_bf16x2x2_cvt_bf16x4_to_fp8x4(cfg.FP8_DTYPE) + rO_u32 = cute.recast_tensor(rO, Uint32) + + for row_idx in cutlass.range_constexpr(ROWS_PER_THREAD): + cute.copy(load_atom, gX_thread[None, row_idx], rX[None, row_idx], cache_policy=cache_policy) + + @cute.jit + def extract_row_amax(values_i32: cute.Tensor) -> Uint32: + """Reduce a 32-column MXFP8 block across its two or four lane owners.""" + amax = values_i32[0] + for pair_idx in cutlass.range_constexpr(1, cute.size(values_i32)): + amax = abs_max_x2_bf16(amax, values_i32[pair_idx]) + for stage in cutlass.range_constexpr(LANES_PER_MX_BLOCK.bit_length() - 1): + amax = abs_max_x2_bf16(amax, cute.arch.shuffle_sync_bfly(amax, 1 << stage)) + return bf16_pair_magnitude(amax) + + # Rowwise pass: lane shuffles close each 32-column reduction. + for row_idx in cutlass.range_constexpr(ROWS_PER_THREAD): + rX_row_i32 = cute.recast_tensor(rX[None, row_idx], Int32) + reciprocal = bf16_mx_scale_reciprocal(extract_row_amax(rX_row_i32), MAX_EXPONENT) + reciprocal_packed = reciprocal | (reciprocal << Uint32(16)) + for pack_idx in cutlass.range_constexpr(cute.size(rO_u32)): + rO_u32[pack_idx] = scale_and_pack( + rX_row_i32[2 * pack_idx], + reciprocal_packed, + rX_row_i32[2 * pack_idx + 1], + reciprocal_packed, + ) + cute.copy(store_atom, rO, gO_row_thread[None, row_idx], cache_policy=cache_policy) + if lane % LANES_PER_MX_BLOCK == 0: + rS_row[0] = Uint8(Uint32(254) - (reciprocal >> Uint32(7))).bitcast(SCALE_DTYPE) + cute.copy( + row_scale_atom, + rS_row, + cute.local_tile(gS_row_thread, (1,), (row_idx,)), + cache_policy=cache_policy, + ) + + # Columnwise pass: first reduce each thread's rows into one maximum per column. + rColumnAmax = cute.make_rmem_tensor((ELEMENTS_PER_THREAD,), cfg.DTYPE) + rColumnScale = cute.make_rmem_tensor((ELEMENTS_PER_THREAD,), cfg.DTYPE) + rColumnAmax_i32 = cute.recast_tensor(rColumnAmax, Int32) + rColumnScale_u32 = cute.recast_tensor(rColumnScale, Uint32) + rX_i32 = cute.recast_tensor(rX, Int32) + for pair_idx in cutlass.range_constexpr(COLUMN_PAIRS_PER_THREAD): + amax = rX_i32[pair_idx, 0] + for row_idx in cutlass.range_constexpr(1, ROWS_PER_THREAD): + amax = abs_max_x2_bf16(amax, rX_i32[pair_idx, row_idx]) + rColumnAmax_i32[pair_idx] = amax + + allocator = _SmemAllocator() + sColumnAmax = allocator.allocate_tensor( + cfg.DTYPE, + cute.make_layout((CTA_THREADS_Y, TILE_COLS), stride=(TILE_COLS, 1)), + byte_alignment=16, + ) + sColumnScale = allocator.allocate_tensor( + cfg.DTYPE, cute.make_layout((TILE_COLS,)), byte_alignment=16 + ) + sColumnAmax_i32 = cute.recast_tensor(sColumnAmax, Int32) + sColumnScale_u32 = cute.recast_tensor(sColumnScale, Uint32) + # Group each lane's column pairs into two halves. Placing the first halves + # of all lanes before the second halves preserves CUDA's conflict-free + # vectorized shared-memory accesses, without calculating pointer offsets. + column_pair_layout = cute.make_layout( + ((PAIRS_PER_HALF, 2), CTA_THREADS_X), + stride=((1, COLUMN_PAIRS // 2), PAIRS_PER_HALF), + ) + sColumnAmax_thread = cute.composition(sColumnAmax_i32, (None, column_pair_layout))[ + warp, (None, lane) + ] + sColumnScale_thread = cute.composition(sColumnScale_u32, column_pair_layout)[None, lane] + for half_idx in cutlass.range_constexpr(2): + cute.autovec_copy( + cute.local_tile(rColumnAmax_i32, (PAIRS_PER_HALF,), (half_idx,)), + cute.local_tile(sColumnAmax_thread, (PAIRS_PER_HALF,), (half_idx,)), + ) + cute.arch.sync_threads() + + # One thread reduces each column pair across all warps (all 32 tile rows). + if tidx < COLUMN_PAIRS: + amax = sColumnAmax_i32[0, tidx] + for warp_idx in cutlass.range_constexpr(1, CTA_THREADS_Y): + amax = abs_max_x2_bf16(amax, sColumnAmax_i32[warp_idx, tidx]) + sColumnScale_u32[tidx] = bf16_mx_scale_reciprocal_pair(amax, MAX_EXPONENT) + cute.arch.sync_threads() + + # The same layout returns the reciprocals to each lane's original columns. + for half_idx in cutlass.range_constexpr(2): + cute.autovec_copy( + cute.local_tile(sColumnScale_thread, (PAIRS_PER_HALF,), (half_idx,)), + cute.local_tile(rColumnScale_u32, (PAIRS_PER_HALF,), (half_idx,)), + ) + if warp == 0: + rS_col_u32 = cute.recast_tensor(rS_col, Uint32) + for pack_idx in cutlass.range_constexpr(cute.size(rS_col_u32)): + rS_col_u32[pack_idx] = bf16_mx_scale_bytes_pair(rColumnScale_u32[2 * pack_idx]) | ( + bf16_mx_scale_bytes_pair(rColumnScale_u32[2 * pack_idx + 1]) << Uint32(16) + ) + cute.copy(col_scale_atom, rS_col, gS_col_thread, cache_policy=cache_policy) + + for row_idx in cutlass.range_constexpr(ROWS_PER_THREAD): + rX_row_i32 = cute.recast_tensor(rX[None, row_idx], Int32) + for pack_idx in cutlass.range_constexpr(cute.size(rO_u32)): + rO_u32[pack_idx] = scale_and_pack( + rX_row_i32[2 * pack_idx], + rColumnScale_u32[2 * pack_idx], + rX_row_i32[2 * pack_idx + 1], + rColumnScale_u32[2 * pack_idx + 1], + ) + cute.copy(store_atom, rO, gO_col_thread[None, row_idx], cache_policy=cache_policy) + + +class MXFP8QuantizeRegisterBidimensionalKernel(MXFP8QuantizeKernelBase): + """Load one CTA tile into registers and quantize it along both axes. + + TILE_ROWS = _CTA_THREADS_Y * _ROWS_PER_THREAD = 32. + TILE_COLS = _CTA_THREADS_X * ELEMENTS_PER_THREAD = 256 or 512. + The 32-row tile contains complete columnwise MXFP8 blocks; shared memory + combines the partial column maxima from the eight warps. Requires aligned + BF16 input, M divisible by 32, N divisible by 256 and plain scale layouts. + """ + + _CTA_THREADS_Y = 8 + _CTA_THREADS_X = 32 + _ROWS_PER_THREAD = 4 + _ELEMENTS_PER_THREAD = 16 + + def __init__(self, cfg: MXFP8QuantizeConfig): + self.cfg = cfg + + @cute.jit + def __call__( + self, + mX: cute.Tensor, + mO_row: Optional[cute.Tensor], + mS_row: Optional[cute.Tensor], + mO_col: Optional[cute.Tensor], + mS_col: Optional[cute.Tensor], + mAmax: Optional[cute.Tensor], + mNoop: cute.Pointer, + mDActInput: Optional[cute.Tensor], + mWorkspace: Optional[cute.Tensor], + stream: CUstream, + ): + M, N = mX.shape + TILE_ROWS = self._CTA_THREADS_Y * self._ROWS_PER_THREAD + WIDE_TILE_COLS = self._CTA_THREADS_X * self._ELEMENTS_PER_THREAD + if N % WIDE_TILE_COLS == 0: + if Int64(M // TILE_ROWS) * (N // WIDE_TILE_COLS) >= 4096: + self._launch( + mX, + mO_row, + mS_row, + mO_col, + mS_col, + mNoop, + stream, + CTA_THREADS_Y=self._CTA_THREADS_Y, + CTA_THREADS_X=self._CTA_THREADS_X, + ROWS_PER_THREAD=self._ROWS_PER_THREAD, + ELEMENTS_PER_THREAD=self._ELEMENTS_PER_THREAD, + CLUSTER_WIDTH=4, + ) + else: + self._launch( + mX, + mO_row, + mS_row, + mO_col, + mS_col, + mNoop, + stream, + CTA_THREADS_Y=self._CTA_THREADS_Y, + CTA_THREADS_X=self._CTA_THREADS_X, + ROWS_PER_THREAD=self._ROWS_PER_THREAD, + ELEMENTS_PER_THREAD=self._ELEMENTS_PER_THREAD, + CLUSTER_WIDTH=1, + ) + else: + self._launch( + mX, + mO_row, + mS_row, + mO_col, + mS_col, + mNoop, + stream, + CTA_THREADS_Y=self._CTA_THREADS_Y, + CTA_THREADS_X=self._CTA_THREADS_X, + ROWS_PER_THREAD=self._ROWS_PER_THREAD, + ELEMENTS_PER_THREAD=self._ELEMENTS_PER_THREAD // 2, + CLUSTER_WIDTH=8, + ) + + @cute.jit + def _launch( + self, + mX: cute.Tensor, + mO_row: cute.Tensor, + mS_row: cute.Tensor, + mO_col: cute.Tensor, + mS_col: cute.Tensor, + mNoop: cute.Pointer, + stream: CUstream, + CTA_THREADS_Y: cutlass.Constexpr[int], + CTA_THREADS_X: cutlass.Constexpr[int], + ROWS_PER_THREAD: cutlass.Constexpr[int], + ELEMENTS_PER_THREAD: cutlass.Constexpr[int], + CLUSTER_WIDTH: cutlass.Constexpr[int], + ): + assert CTA_THREADS_X == 32 + assert CTA_THREADS_Y * ROWS_PER_THREAD == 32 + assert ELEMENTS_PER_THREAD in (8, 16) + M, N = mX.shape + TILE_ROWS = CTA_THREADS_Y * ROWS_PER_THREAD + TILE_COLS = CTA_THREADS_X * ELEMENTS_PER_THREAD + THREADS_PER_CTA = CTA_THREADS_Y * CTA_THREADS_X + grid_cols = N // TILE_COLS + grid_rows = M // TILE_ROWS + cluster_cols = Int32(CLUSTER_WIDTH) + while cluster_cols > 1 and grid_cols % cluster_cols != 0: + cluster_cols -= 1 + # Keep column CTAs on grid.x so clusters read adjacent column tiles. + grid = [grid_cols, grid_rows] + block = [THREADS_PER_CTA] + cluster = [cluster_cols, 1, 1] + # Preserve CUDA's compile-time strides for the two common column counts. + if N == 2048: + self.kernel( + mX, + mO_row, + mS_row, + mO_col, + mS_col, + mNoop, + CTA_THREADS_Y=CTA_THREADS_Y, + CTA_THREADS_X=CTA_THREADS_X, + ROWS_PER_THREAD=ROWS_PER_THREAD, + ELEMENTS_PER_THREAD=ELEMENTS_PER_THREAD, + N_CONST=2048, + ).launch(grid=grid, block=block, cluster=cluster, min_blocks_per_mp=4, stream=stream) + elif N == 4096: + self.kernel( + mX, + mO_row, + mS_row, + mO_col, + mS_col, + mNoop, + CTA_THREADS_Y=CTA_THREADS_Y, + CTA_THREADS_X=CTA_THREADS_X, + ROWS_PER_THREAD=ROWS_PER_THREAD, + ELEMENTS_PER_THREAD=ELEMENTS_PER_THREAD, + N_CONST=4096, + ).launch(grid=grid, block=block, cluster=cluster, min_blocks_per_mp=4, stream=stream) + else: + self.kernel( + mX, + mO_row, + mS_row, + mO_col, + mS_col, + mNoop, + CTA_THREADS_Y=CTA_THREADS_Y, + CTA_THREADS_X=CTA_THREADS_X, + ROWS_PER_THREAD=ROWS_PER_THREAD, + ELEMENTS_PER_THREAD=ELEMENTS_PER_THREAD, + N_CONST=0, + ).launch(grid=grid, block=block, cluster=cluster, min_blocks_per_mp=4, stream=stream) + + @cute.kernel + def kernel( + self, + mX: cute.Tensor, + mO_row: cute.Tensor, + mS_row: cute.Tensor, + mO_col: cute.Tensor, + mS_col: cute.Tensor, + mNoop: cute.Pointer, + CTA_THREADS_Y: cutlass.Constexpr[int], + CTA_THREADS_X: cutlass.Constexpr[int], + ROWS_PER_THREAD: cutlass.Constexpr[int], + ELEMENTS_PER_THREAD: cutlass.Constexpr[int], + N_CONST: cutlass.Constexpr[int], + ): + """Check noop before constructing or accessing any CTA tile.""" + if not noop_flag_is_set(mNoop): + M, N = mX.shape + if cutlass.const_expr(N_CONST != 0): + N = N_CONST + payload_layout = cute.make_layout((M, N), stride=(N, 1)) + gX = cute.make_tensor(mX.iterator.align(32), payload_layout) + gO_row = cute.make_tensor(mO_row.iterator, payload_layout) + gO_col = cute.make_tensor(mO_col.iterator, payload_layout) + gS_row = cute.make_tensor( + mS_row.iterator, + cute.make_layout((cute.size(mS_row, mode=[0]), N // 32), stride=(N // 32, 1)), + ) + gS_col = cute.make_tensor( + mS_col.iterator.align(16), + cute.make_layout((cute.size(mS_col, mode=[0]), N), stride=(N, 1)), + ) + else: + # The dispatcher checks these alignments before choosing this kernel. + gX = cute.make_tensor(mX.iterator.align(32), mX.layout) + gO_row = mO_row + gO_col = mO_col + gS_row = mS_row + gS_col = cute.make_tensor(mS_col.iterator.align(16), mS_col.layout) + cta_col, cta_row, _ = cute.arch.block_idx() + cta_coord = (cta_row, cta_col) + TILE_ROWS = CTA_THREADS_Y * ROWS_PER_THREAD + TILE_COLS = CTA_THREADS_X * ELEMENTS_PER_THREAD + TILER = (TILE_ROWS, TILE_COLS) + ROW_SCALE_TILER = (TILE_ROWS, TILE_COLS // 32) + COL_SCALE_TILER = (TILE_ROWS // 32, TILE_COLS) + gX_tile = cute.local_tile(gX, TILER, cta_coord) + gO_row_tile = cute.local_tile(gO_row, TILER, cta_coord) + gO_col_tile = cute.local_tile(gO_col, TILER, cta_coord) + gS_row_tile = cute.local_tile(gS_row, ROW_SCALE_TILER, cta_coord) + gS_col_tile = cute.local_tile(gS_col, COL_SCALE_TILER, cta_coord) + cache_policy = create_l2_policy(True) + quantize_bidimensional_mxfp8_g2r( + gX_tile, + gO_row_tile, + gS_row_tile, + gO_col_tile, + gS_col_tile, + self.cfg, + CTA_THREADS_Y=CTA_THREADS_Y, + CTA_THREADS_X=CTA_THREADS_X, + ROWS_PER_THREAD=ROWS_PER_THREAD, + ELEMENTS_PER_THREAD=ELEMENTS_PER_THREAD, + cache_policy=cache_policy, + ) diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_1lane.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_1lane.py new file mode 100644 index 00000000000..bbdf9a4b510 --- /dev/null +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_1lane.py @@ -0,0 +1,265 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""One-lane GMEM-to-RMEM rowwise MXFP8 cast with vectorized global I/O.""" + +# Local @cute.struct classes are shared-memory storage descriptors. +# pylint: disable=missing-class-docstring + +from typing import Optional, Type + +import cutlass +from cutlass import cute +from cutlass import Float8E8M0FNU, Int32, Uint32, Uint8 +from cuda.bindings.driver import CUstream # pylint: disable=no-name-in-module + +from transformer_engine.common.CuTeDSL.utils import ( + abs_max_x2_bf16, + abs_max_x2_f16, + x2_lo_to_f32_bf16, + x2_lo_to_f32_f16, + x2_hi_to_f32_bf16, + x2_hi_to_f32_f16, + fabs_f32, + exp2f_rcp, + pack_f32x2, +) +from transformer_engine.common.CuTeDSL.utils_fp8 import ( + cvt_f32_to_fp8e8m0fnu, + mul_f32x2_cvt_packed16x4_to_fp8x4, +) +from .quantize_mxfp8_common import ( + CUTEDSL_DEBUG_LOGGING, + MXFP8QuantizeConfig, + MXFP8QuantizeKernelBase, + MXFP8_BLOCK_SCALING_SIZE, + derive_swizzled_scale_layout, + noop_flag_is_set, +) + + +class MXFP8QuantizeG2RRowwise1LaneKernel(MXFP8QuantizeKernelBase): + """Specialized cast-only ROWWISE-only MXFP8 kernel. + + Requires N % 128 == 0 (full vectorizable column chunks). + + Plain rowwise-only quantize. Each thread owns one 32-element MXFP8 chunk and + uses vectorized global loads/stores (no TMA used).""" + + _TILE_ROWS = 4 + _TILE_COLS = 1024 + _THREADS_PER_CTA = 128 + + def __init__(self, cfg: MXFP8QuantizeConfig): + self.cfg = cfg + # If True, then this kernel will first write each thread's scale byte to a shared memory buffer, + # then utilize vectorized store to flush the buffer to global memory. + self._STASH_SCALE_TO_SMEM = True # Hardcode to true for now + + @cute.jit + def __call__( + self, + mX: cute.Tensor, + mO_row: Optional[cute.Tensor], + mS_row: Optional[cute.Tensor], + mO_col: Optional[cute.Tensor], + mS_col: Optional[cute.Tensor], # Unused, kept for API compatibility + mAmax: Optional[cute.Tensor], # Unused, kept for API compatibility + mNoop: cute.Pointer, # f32 cast_noop flag; may be null, checked on device + mDActInput: Optional[cute.Tensor], # Unused, kept for API compatibility + mWorkspace: Optional[cute.Tensor], # Unused, kept for API compatibility + stream: CUstream, + ): + if cutlass.const_expr(CUTEDSL_DEBUG_LOGGING): + cute.printf( + f"[CuTeDSL] MXFP8QuantizeG2RRowwise1LaneKernel.__call__() with config: {self.cfg}\n" + ) + + M = mX.shape[0] + N = mX.shape[1] + + if cutlass.const_expr(self.cfg.WITH_GEMM_SWIZZLED_SCALES): + mS_row, _ = derive_swizzled_scale_layout(M, N, True, False, mS_row, None) + + grid = [ + cute.ceil_div(Int32(N), self._TILE_COLS), + cute.ceil_div(M, self._TILE_ROWS), + ] + block = [self._THREADS_PER_CTA] + + self.kernel( + mX, + mO_row, + mS_row, + mNoop, + mX.element_type, + ).launch(grid=grid, block=block, stream=stream) + + @cute.kernel + def kernel( + self, + mX: cute.Tensor, + mO_row: cute.Tensor, + mS_row: cute.Tensor, + mNoop: cute.Pointer, + DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], + ): + """Device entry: no-op the CTA when the noop flag is set, else run the quantize main loop.""" + if not noop_flag_is_set(mNoop): + self._kernel_main(mX, mO_row, mS_row, DTYPE) + + @cute.jit + def _kernel_main( + self, + mX: cute.Tensor, + mO_row: cute.Tensor, + mS_row: cute.Tensor, + DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], + ): + """Device entry for the specialized rowwise-only cast kernel (vectorized global loads/stores, no TMA).""" + tidx, _, _ = cute.arch.thread_idx() + bidx, bidy, _ = cute.arch.block_idx() + M = mX.shape[0] + N = mX.shape[1] + + # Each thread handles one 32-element MXFP8 chunk (= one scale block). + # The 128 threads in the CTA are grouped as (4, 32), so they cover a + # (4, 1024) input tile and the matching (4, 32) scale tile. + CTA_Y = self._TILE_ROWS + CTA_X = self._TILE_COLS // MXFP8_BLOCK_SCALING_SIZE + tiler, tv_layout = cute.make_layout_tv( + thr_layout=cute.make_layout((CTA_Y, CTA_X), stride=(CTA_X, 1)), + val_layout=cute.make_layout( + (1, MXFP8_BLOCK_SCALING_SIZE), stride=(MXFP8_BLOCK_SCALING_SIZE, 1) + ), + ) + tiler_scale, tv_layout_scale = cute.make_layout_tv( + thr_layout=cute.make_layout((CTA_Y, CTA_X), stride=(CTA_X, 1)), + val_layout=cute.make_layout((1, 1), stride=(1, 1)), + ) + + # Select the tile that belongs to this CTA, then the fragment per thread. + mX_tile = cute.local_tile(mX, tiler, (bidy, bidx)) + mO_tile = cute.local_tile(mO_row, tiler, (bidy, bidx)) + mS_tile = cute.local_tile(mS_row, tiler_scale, (bidy, bidx)) + mX_thread = cute.composition(mX_tile, tv_layout)[tidx, None] + mO_thread = cute.composition(mO_tile, tv_layout)[tidx, None] + mS_thread = cute.composition(mS_tile, tv_layout_scale)[tidx, None] + + rX_thread = cute.make_rmem_tensor( + cute.make_layout((1, MXFP8_BLOCK_SCALING_SIZE), stride=(MXFP8_BLOCK_SCALING_SIZE, 1)), + dtype=DTYPE, + ) + abs_max_x2 = abs_max_x2_f16 if DTYPE is cutlass.Float16 else abs_max_x2_bf16 + x2_lo_to_f32 = x2_lo_to_f32_f16 if DTYPE is cutlass.Float16 else x2_lo_to_f32_bf16 + x2_hi_to_f32 = x2_hi_to_f32_f16 if DTYPE is cutlass.Float16 else x2_hi_to_f32_bf16 + mul_cvt4 = mul_f32x2_cvt_packed16x4_to_fp8x4(DTYPE, self.cfg.FP8_DTYPE) + rX_i32 = cute.make_tensor( + cute.recast_ptr(rX_thread.iterator, dtype=Int32), + cute.make_layout((MXFP8_BLOCK_SCALING_SIZE // 2,), stride=(1,)), + ) + rO_thread = cute.make_rmem_tensor( + cute.make_layout((1, MXFP8_BLOCK_SCALING_SIZE), stride=(MXFP8_BLOCK_SCALING_SIZE, 1)), + dtype=self.cfg.FP8_DTYPE, + ) + rO_u32 = cute.make_tensor( + cute.recast_ptr(rO_thread.iterator, dtype=Uint32), + cute.make_layout( + (MXFP8_BLOCK_SCALING_SIZE // 4,), stride=(1,) + ), # Unit is Uint32, divide by 4 here + ) + + sS_thread = None + if cutlass.const_expr(self._STASH_SCALE_TO_SMEM): + + @cute.struct + class SharedStorage: + buf: cute.struct.Align[cute.struct.MemRange[Float8E8M0FNU, CTA_Y * CTA_X], 16] + + storage = cutlass.utils.SmemAllocator().allocate(SharedStorage) + sScale = storage.buf.get_tensor(cute.make_layout((CTA_Y, CTA_X), stride=(CTA_X, 1))) + # sScale is (CTA_Y, CTA_X):(CTA_X, 1), which is the same layout as tv_layout_scale + # so sS_thread is really just a 1-element e8m0 buffer for this thread's scale. + sS_thread = cute.composition(sScale, tv_layout_scale)[tidx, None] + # Zero first so padding columns (cols past N/32 in the padded scale + # matrix) flush as 0 and we never read uninitialized smem. + # Raw 0x00: e8m0 has no zero value to convert from. + sS_thread[0] = Uint8(0).bitcast(Float8E8M0FNU) + cute.arch.sync_threads() + + row = bidy * self._TILE_ROWS + tidx // CTA_X + col = bidx * self._TILE_COLS + (tidx % CTA_X) * MXFP8_BLOCK_SCALING_SIZE + if row < M and col < N: + cute.autovec_copy(mX_thread, rX_thread) + amax_2x = rX_i32[0] + # Skip the first iteration because we assigned rX_i32[0] to amax_2x already + for i in cutlass.range_constexpr(1, MXFP8_BLOCK_SCALING_SIZE // 2): + amax_2x = abs_max_x2(amax_2x, rX_i32[i]) + amax = cute.arch.fmax(fabs_f32(x2_lo_to_f32(amax_2x)), fabs_f32(x2_hi_to_f32(amax_2x))) + + biased_exp = cvt_f32_to_fp8e8m0fnu(amax * self.cfg.MAX_NORM_RCP) + if cutlass.const_expr(self._STASH_SCALE_TO_SMEM): + sS_thread[0] = biased_exp + else: + mS_thread[0] = biased_exp + + # Rescale + FP8 cast, 4 elements per fused mul_cvt (one uint32 out), + # then a vectorized store. Mirrors CUDA's _use_cvt_4x path. + inv_scale = exp2f_rcp(biased_exp) + scale_2x = pack_f32x2(inv_scale, inv_scale) + for i in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE // 4): + rO_u32[i] = mul_cvt4(rX_i32[2 * i], rX_i32[2 * i + 1], scale_2x) + cute.autovec_copy(rO_thread, mO_thread) + + # Cooperative wide flush of the staged scales where padding columns flush as 0. + if cutlass.const_expr(self._STASH_SCALE_TO_SMEM): + cute.arch.sync_threads() + # Use cute.size instead of .shape[1] because under the swizzled layout mode 1 is + # a nested tuple ((4, num_tiles_SC)), not a plain int. + padded_cols = cute.size(mS_row, mode=[1]) + if cutlass.const_expr(self.cfg.WITH_GEMM_SWIZZLED_SCALES): + # Swizzled rowwise scale layout is ((32, 4, num_tiles_M), (4, num_tiles_SC)):((16, 4, num_tiles_SC * 512), (1, 512)) + # so we have at most 4 elements continuous which is our vectorized store width + self._flush_scales_to_gmem(sScale, mS_tile, tidx, bidx, bidy, M, padded_cols, 4) + elif padded_cols % 16 == 0: + # If columns is divisible by 16, use 16 bytes as the vectorized store width + self._flush_scales_to_gmem(sScale, mS_tile, tidx, bidx, bidy, M, padded_cols, 16) + else: + # Otherwise use 4 bytes as the vectorized store width. + # Note our fake tensor requires 4 divisibility so this is enforced as long as you can get here + self._flush_scales_to_gmem(sScale, mS_tile, tidx, bidx, bidy, M, padded_cols, 4) + + @cute.jit + def _flush_scales_to_gmem( + self, + sScale: cute.Tensor, + mS_tile: cute.Tensor, + tidx: Int32, + bidx: Int32, + bidy: Int32, + M: Int32, + padded_cols: Int32, + width: cutlass.Constexpr[int], + ): + """Flush the staged (CTA_Y, CTA_X) scale tile to gmem with vectorized stores.""" + CTA_Y = self._TILE_ROWS + CTA_X = self._TILE_COLS // MXFP8_BLOCK_SCALING_SIZE + # Previously each threads has 1 byte, but now we are doing vectorized store, + # which means only a subset of threads will need to issue the store while other threads are not used. + active_threads = CTA_X // width + _, tv_flush = cute.make_layout_tv( + thr_layout=cute.make_layout((CTA_Y, active_threads), stride=(active_threads, 1)), + val_layout=cute.make_layout((1, width), stride=(width, 1)), + ) + # We only need to use a subset of threads with shape (CTA_Y, active_threads) to write + # so if the thread is outside of this subset, it will remain inactive + if tidx < CTA_Y * active_threads: + # Absolute position of the scale vector to write in the GMEM buffer + thread_y = bidy * CTA_Y + tidx // active_threads + thread_x = bidx * CTA_X + (tidx % active_threads) * width + if thread_y < M and thread_x < padded_cols: + cute.autovec_copy( + cute.composition(sScale, tv_flush)[tidx, None], + cute.composition(mS_tile, tv_flush)[tidx, None], + ) diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_2lane.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_2lane.py new file mode 100644 index 00000000000..95882a4161f --- /dev/null +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_2lane.py @@ -0,0 +1,368 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Two-lane GMEM-to-RMEM rowwise MXFP8 cast, ported from cast_rowwise.cuh.""" + +from typing import Optional + +import cutlass +from cutlass import cute +from cutlass import Boolean, Float32, Int32, Int64, Uint32, Uint8 +from cuda.bindings.driver import CUstream + +from transformer_engine.common.CuTeDSL.utils import ( + abs_max_x2_bf16, + create_l2_policy, + exp2_bf16x2_rcp, +) +from transformer_engine.common.CuTeDSL.utils_fp8 import ( + cvt_f32_to_fp8e8m0fnu, + mul_bf16x2x2_cvt_bf16x4_to_fp8x4, +) +from .quantize_mxfp8_common import ( + MXFP8QuantizeConfig, + MXFP8QuantizeKernelBase, + bf16_pair_magnitude, + derive_swizzled_scale_layout, + noop_flag_is_set, +) + + +@cute.jit +def quantize_rowwise_mxfp8_g2r( + gX_tiles: cute.Tensor, + gO_tiles: cute.Tensor, + mS_tiles: cute.Tensor, + cfg: cutlass.Constexpr[MXFP8QuantizeConfig], + M: Int32, + N: Int32, + CTA_THREADS_Y: cutlass.Constexpr[int], + CTA_THREADS_X: cutlass.Constexpr[int], + ELEMENTS_PER_THREAD: cutlass.Constexpr[int], + NUM_TILES: cutlass.Constexpr[int], + CHECK_BOUNDS: cutlass.Constexpr[bool], + input_policy: Int64, + output_policy: Int64, +): + """Quantize a CTA's (TILE_ROWS, TILE_COLS, NUM_TILES) GMEM tiles. + + CTA_THREADS_X threads own each tile row. Each thread loads 16 BF16 values, and + two adjacent lanes share the maximum and scale for a 32-value MXFP8 block. + Load all tiles before reducing to preserve independent memory requests. + """ + TILE_ROWS = CTA_THREADS_Y + TILE_COLS = CTA_THREADS_X * ELEMENTS_PER_THREAD + tidx, _, _ = cute.arch.thread_idx() + + _, tv_payload = cute.make_layout_tv( + thr_layout=cute.make_layout((CTA_THREADS_Y, CTA_THREADS_X), stride=(CTA_THREADS_X, 1)), + val_layout=cute.make_layout((1, ELEMENTS_PER_THREAD), stride=(ELEMENTS_PER_THREAD, 1)), + ) + # Mode 0.0: two adjacent threads share one scale (stride 0). + # Mode 0.1: thread pairs are arranged as (scale columns, tile rows). + # Mode 1: each thread addresses one scale element. + tv_scale = cute.make_layout( + ((2, (CTA_THREADS_X // 2, CTA_THREADS_Y)), 1), stride=((0, (TILE_ROWS, 1)), 0) + ) + + # We load fragments from all GMEM tiles at once + rX = cute.make_rmem_tensor((ELEMENTS_PER_THREAD, NUM_TILES), cfg.DTYPE) + # We store a fragment of one tile at a time + rO = cute.make_rmem_tensor((ELEMENTS_PER_THREAD,), cfg.FP8_DTYPE) + # RMEM flags to indicate whether each tile is out of bounds or not + # With CHECK_BOUNDS=False DCE will remove it so we don't pay extra cost when we don't use it + tile_valid_flags = cute.make_rmem_tensor((NUM_TILES,), Boolean) + + # 16 * 2 bytes = 32 bytes which is already the widest + LOAD_WIDTH = ELEMENTS_PER_THREAD * cfg.DTYPE.width + load_atom = cute.make_copy_atom( + cute.nvgpu.CopyG2ROp(), + cfg.DTYPE, + num_bits_per_copy=LOAD_WIDTH, + invariant=True, # hint that data being read stays unchanged during the kernel + ) + store_atom = cute.make_copy_atom( + cute.nvgpu.CopyR2GOp(), + cfg.FP8_DTYPE, + num_bits_per_copy=ELEMENTS_PER_THREAD * cfg.FP8_DTYPE.width, + ) + + scale_and_pack = mul_bf16x2x2_cvt_bf16x4_to_fp8x4(cfg.FP8_DTYPE) + + # Copy all tiles from GMEM to RMEM and potentially check their bounds + for tile_idx in cutlass.range_constexpr(NUM_TILES): + # gX_tiles is (TILE_ROWS, TILE_COLS, NUM_TILES) + gX_tile = gX_tiles[None, None, tile_idx] + gX_thread = cute.composition(gX_tile, tv_payload)[tidx, None] + valid = Boolean(True) + if cutlass.const_expr(CHECK_BOUNDS): + cta_row, cta_col, _ = cute.arch.block_idx() + row = cta_row * TILE_ROWS + tidx // CTA_THREADS_X + col = (cta_col * NUM_TILES + tile_idx) * TILE_COLS + ( + tidx % CTA_THREADS_X + ) * ELEMENTS_PER_THREAD + valid = row < M and col < N + tile_valid_flags[tile_idx] = valid + rX[None, tile_idx].fill(0) + if valid: + cute.copy(load_atom, gX_thread, rX[None, tile_idx], cache_policy=input_policy) + + @cute.jit + def extract_tile_amax(values_i32: cute.Tensor) -> Float32: + """Return the 32-value MXFP8 block amax shared by two adjacent threads. + + Balanced eight-pair tree, followed by one butterfly between lane partners. + """ + level = [ + abs_max_x2_bf16(values_i32[2 * i], values_i32[2 * i + 1]) + for i in range(cute.size(values_i32) // 2) + ] + amax_pair = abs_max_x2_bf16( + abs_max_x2_bf16(level[0], level[2]), abs_max_x2_bf16(level[1], level[3]) + ) + # Compare with the neighbor thread's amax (2 threads cooperate to process a MXFP8 block). + partner = cute.arch.shuffle_sync_bfly(amax_pair, 1) + amax_pair = abs_max_x2_bf16(amax_pair, partner) + return (bf16_pair_magnitude(amax_pair) << Uint32(16)).bitcast(Float32) + + for tile_idx in cutlass.range_constexpr(NUM_TILES): + gO_tile = gO_tiles[None, None, tile_idx] + mS_tile = mS_tiles[None, None, tile_idx] + gO_thread = cute.composition(gO_tile, tv_payload)[tidx, None] + mS_thread = cute.composition(mS_tile, tv_scale)[tidx, None] + # Reinterpret the native tensors only for the packed BF16/FP8 instructions. + rX_tile_i32 = cute.recast_tensor(rX[None, tile_idx], Int32) + rO_u32 = cute.recast_tensor(rO, Uint32) + amax = extract_tile_amax(rX_tile_i32) + exponent = cvt_f32_to_fp8e8m0fnu(amax * cfg.MAX_NORM_RCP) + + # Collect 4 scale bytes and write them using one thread + byte = Uint32(exponent.bitcast(Uint8)) + c1 = cute.arch.shuffle_sync_down(byte, 2) + c2 = cute.arch.shuffle_sync_down(byte, 4) + c3 = cute.arch.shuffle_sync_down(byte, 6) + if tidx % 8 == 0 and tile_valid_flags[tile_idx]: + scale_word = cute.make_tensor( + cute.recast_ptr(mS_thread.iterator, dtype=Uint32), cute.make_layout(1) + ) + scale_word[0] = byte | (c1 << Uint32(8)) | (c2 << Uint32(16)) | (c3 << Uint32(24)) + + # Compute and write the quantized MXFP8 values to GMEM + reciprocal_packed = exp2_bf16x2_rcp(exponent) + for i in cutlass.range_constexpr(cute.size(rO_u32)): + rO_u32[i] = scale_and_pack( + rX_tile_i32[2 * i], reciprocal_packed, rX_tile_i32[2 * i + 1], reciprocal_packed + ) + if tile_valid_flags[tile_idx]: + cute.copy(store_atom, rO, gO_thread, cache_policy=output_policy) + + +class MXFP8QuantizeG2RRowwise2LaneKernel(MXFP8QuantizeKernelBase): + """Two lanes per MXFP8 block, with a CTA sweeping NUM_TILES horizontally. + + One CTA tile has TILE_ROWS = CTA_THREADS_Y and + TILE_COLS = CTA_THREADS_X * _ELEMENTS_PER_THREAD. The CTA + processes NUM_TILES = 1 or 2 in the original input matrix. Both scale + formats use the same payload tiling; a composed layout maps swizzled scales. + _L2_CACHED_CTA_PERCENT tunes the share of CTAs whose input loads use normal + L2 retention: 0 streams all input and 100 caches all input normally. + """ + + _CTA_THREADS_X = 32 + _ELEMENTS_PER_THREAD = 16 + _L2_CACHED_CTA_PERCENT = 40 + + def __init__(self, cfg: MXFP8QuantizeConfig): + self.cfg = cfg + + @cute.jit + def __call__( + self, + mX: cute.Tensor, + mO_row: Optional[cute.Tensor], + mS_row: Optional[cute.Tensor], + mO_col: Optional[cute.Tensor], + mS_col: Optional[cute.Tensor], + mAmax: Optional[cute.Tensor], + mNoop: cute.Pointer, + mDActInput: Optional[cute.Tensor], + mWorkspace: Optional[cute.Tensor], + stream: CUstream, + ): + output_bytes = Int64(mX.shape[0]) * mX.shape[1] + if output_bytes <= 12 << 20: + self._launch( + mX, + mO_row, + mS_row, + mNoop, + stream, + CTA_THREADS_Y=8, + CTA_THREADS_X=self._CTA_THREADS_X, + NUM_TILES=1, + ) + elif output_bytes <= 48 << 20: + self._launch( + mX, + mO_row, + mS_row, + mNoop, + stream, + CTA_THREADS_Y=8, + CTA_THREADS_X=self._CTA_THREADS_X, + NUM_TILES=2, + ) + elif output_bytes <= 96 << 20: + self._launch( + mX, + mO_row, + mS_row, + mNoop, + stream, + CTA_THREADS_Y=4, + CTA_THREADS_X=self._CTA_THREADS_X, + NUM_TILES=2, + ) + else: + self._launch( + mX, + mO_row, + mS_row, + mNoop, + stream, + CTA_THREADS_Y=8, + CTA_THREADS_X=self._CTA_THREADS_X, + NUM_TILES=2, + ) + + @cute.jit + def _launch( + self, + mX: cute.Tensor, + mO: cute.Tensor, + mS: cute.Tensor, + mNoop: cute.Pointer, + stream: CUstream, + CTA_THREADS_Y: cutlass.Constexpr[int], + CTA_THREADS_X: cutlass.Constexpr[int], + NUM_TILES: cutlass.Constexpr[int], + ): + assert 0 <= self._L2_CACHED_CTA_PERCENT <= 100 + assert self._ELEMENTS_PER_THREAD == 16 + assert CTA_THREADS_X % 8 == 0 + assert CTA_THREADS_Y * CTA_THREADS_X % 32 == 0 + M, N = mX.shape + # Each CTA processes NUM_TILES horizontal tiles of (TILE_ROWS, TILE_COLS) values. + TILE_ROWS = CTA_THREADS_Y + TILE_COLS = CTA_THREADS_X * self._ELEMENTS_PER_THREAD + THREADS_PER_CTA = CTA_THREADS_Y * CTA_THREADS_X + # Total columns covered by a CTA; its row count stays TILE_ROWS. + CTA_ROWS = TILE_ROWS + CTA_COLS = TILE_COLS * NUM_TILES + + grid_rows = cute.ceil_div(M, CTA_ROWS) + grid_cols = cute.ceil_div(N, CTA_COLS) + check_bounds = M % CTA_ROWS != 0 or N % CTA_COLS != 0 + grid = Int64(grid_rows) * grid_cols + first_streaming_cta = grid * self._L2_CACHED_CTA_PERCENT // 100 + + if grid > 0: + if check_bounds: + self.kernel( + mX, + mO, + mS, + mNoop, + first_streaming_cta, + CTA_THREADS_Y=CTA_THREADS_Y, + CTA_THREADS_X=CTA_THREADS_X, + NUM_TILES=NUM_TILES, + CHECK_BOUNDS=True, + ).launch(grid=[grid_rows, grid_cols], block=[THREADS_PER_CTA], stream=stream) + else: + self.kernel( + mX, + mO, + mS, + mNoop, + first_streaming_cta, + CTA_THREADS_Y=CTA_THREADS_Y, + CTA_THREADS_X=CTA_THREADS_X, + NUM_TILES=NUM_TILES, + CHECK_BOUNDS=False, + ).launch(grid=[grid_rows, grid_cols], block=[THREADS_PER_CTA], stream=stream) + + @cute.kernel + def kernel( + self, + mX: cute.Tensor, + mO: cute.Tensor, + mS: cute.Tensor, + mNoop: cute.Pointer, + first_streaming_cta: Int64, + CTA_THREADS_Y: cutlass.Constexpr[int], + CTA_THREADS_X: cutlass.Constexpr[int], + NUM_TILES: cutlass.Constexpr[int], + CHECK_BOUNDS: cutlass.Constexpr[bool], + ): + """Check the noop flag before selecting or accessing any GMEM tile.""" + if not noop_flag_is_set(mNoop): + # Dispatch checks 32-byte input alignment before selecting this kernel. + gX = cute.make_tensor(mX.iterator.align(32), mX.layout) + # CUDA's x/y axes hold row/column tiles in the same order as local_tile. + cta_coord = cute.arch.block_idx()[:2] + M, N = mX.shape + TILE_ROWS = CTA_THREADS_Y + TILE_COLS = CTA_THREADS_X * self._ELEMENTS_PER_THREAD + CTA_ROWS = TILE_ROWS + CTA_COLS = TILE_COLS * NUM_TILES + + gS = mS + # Apply the swizzled scale layout if requested + if cutlass.const_expr(self.cfg.WITH_GEMM_SWIZZLED_SCALES): + mS_swizzled, _ = derive_swizzled_scale_layout(M, N, True, False, mS, None) + gS = cute.composition(mS_swizzled, (cute.make_layout(M), cute.make_layout(N // 32))) + + # CTA_TILER is the block that a CTA processes in total, which consists of NUM_TILES horizontal tiles + CTA_TILER = (CTA_ROWS, CTA_COLS) + CTA_SCALE_TILER = (CTA_ROWS, CTA_COLS // 32) + gX_cta = cute.local_tile(gX, CTA_TILER, cta_coord) + gO_cta = cute.local_tile(mO, CTA_TILER, cta_coord) + gS_cta = cute.local_tile(gS, CTA_SCALE_TILER, cta_coord) + + fraction = Float32(1.0) + if cutlass.const_expr(self._L2_CACHED_CTA_PERCENT > 0): + fraction = Float32(0.0) + # A linear index is needed only for the L2 policy split. + cta_id = cute.crd2idx( + (Int64(cta_coord[0]), Int64(cta_coord[1])), cute.arch.grid_dim()[:2] + ) + if cta_id >= first_streaming_cta: + fraction = Float32(1.0) + input_policy = create_l2_policy(False, fraction) + output_policy = create_l2_policy(True) + + # Divide CTA_TILER to multiple TILER shape TILES + TILER = (TILE_ROWS, TILE_COLS) + SCALE_TILER = (TILE_ROWS, TILE_COLS // 32) + # Mode 1's first index is 0 because CTA_ROWS == TILE_ROWS (there is only 1 tile on the row dimension). + gX_tiles = cute.local_tile(gX_cta, TILER, (0, None)) + gO_tiles = cute.local_tile(gO_cta, TILER, (0, None)) + mS_tiles = cute.local_tile(gS_cta, SCALE_TILER, (0, None)) + + quantize_rowwise_mxfp8_g2r( + gX_tiles, + gO_tiles, + mS_tiles, + self.cfg, + M, + N, + CTA_THREADS_Y=CTA_THREADS_Y, + CTA_THREADS_X=CTA_THREADS_X, + ELEMENTS_PER_THREAD=self._ELEMENTS_PER_THREAD, + NUM_TILES=NUM_TILES, + CHECK_BOUNDS=CHECK_BOUNDS, + input_policy=input_policy, + output_policy=output_policy, + ) diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2s_bidimensional.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2s_bidimensional.py new file mode 100644 index 00000000000..81a94382bc9 --- /dev/null +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2s_bidimensional.py @@ -0,0 +1,712 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Specialized CuTeDSL bidimensional MXFP8 cast with staged TMA I/O.""" + +# Local @cute.struct classes are shared-memory storage descriptors. +# pylint: disable=missing-class-docstring + +from typing import Optional, Type + +import cutlass +from cutlass import cute +from cutlass import pipeline +from cutlass import Float32, Float8E8M0FNU, Int32, Int64, Uint32 +from cuda.bindings.driver import CUstream # pylint: disable=no-name-in-module + +from transformer_engine.common.CuTeDSL.utils import ( + is_packed16, + abs_max_x2_bf16, + abs_max_x2_f16, + x2_lo_to_f32_bf16, + x2_lo_to_f32_f16, + x2_hi_to_f32_bf16, + x2_hi_to_f32_f16, + fabs_f32, + exp2f_rcp, + pack_f32x2, +) +from transformer_engine.common.CuTeDSL.utils_fp8 import ( + cvt_f32_to_fp8e8m0fnu, + mul_f32x2_cvt_f32x4_to_fp8x4, + mul_f32x4_cvt_f32x4_to_fp8x4, +) +from .quantize_mxfp8_common import ( + CUTEDSL_DEBUG_LOGGING, + MXFP8QuantizeConfig, + MXFP8QuantizeKernelBase, + MXFP8_BLOCK_SCALING_SIZE, + THREADS_PER_WARP, + derive_swizzled_scale_layout, + noop_flag_is_set, +) + + +@cute.jit +def quantize_bidimensional_mxfp8_swizzled( + sX_tile: cute.Tensor, # 32x64 input tile (already sliced to this stage), Sw<3,4,3> swizzled + sO_row_tile: cute.Tensor, # 32x64 rowwise-output tile + sO_col_tile: cute.Tensor, # 32x64 colwise-output tile + sS_row_tile: cute.Tensor, # (32, 2) smem rowwise-scale staging tile (flushed in the epilogue) + sS_col_tile: cute.Tensor, # (1, 64) smem colwise-scale staging tile (flushed in the epilogue) + sColReduce_warp: cute.Tensor, # (32,) SMEM fp32 columnwise scale reduction buffer for this warp + WARPS_PER_CTA: cutlass.Constexpr[int], + MAX_NORM_RCP: cutlass.Constexpr[float], + DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], + FP8_DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], +): + """Quantize a pre-sliced 32x64 tile -> both rowwise and colwise MXFP8. Elements are + addressed through TV layouts: a warp = 32 lanes = 32 rows (lane == row), and each + thread owns one 32-col row segment (its 32-col block) and one output column. + + No bounds check: the caller only loops over valid tiles (`num_tiles`), and + M % 32 == 0 means a tile can only be partial along N. The partial tile's + out-of-bounds warp still runs, but harmlessly -- it reads TMA zero-filled smem, + its output columns are masked by the caller's TMA store, and its scale writes + (rowwise and colwise alike) go to staging slots whose flush targets are past-N + padding columns of the respective scale tensors. + """ + mul_cvt4 = mul_f32x2_cvt_f32x4_to_fp8x4(FP8_DTYPE) + mul_cvt4_elemwise = mul_f32x4_cvt_f32x4_to_fp8x4(FP8_DTYPE) + + _, tv_layout = cute.make_layout_tv( + thr_layout=cute.make_layout(((MXFP8_BLOCK_SCALING_SIZE, 1), WARPS_PER_CTA)), # ((32, 1) 2) + val_layout=cute.make_layout((1, MXFP8_BLOCK_SCALING_SIZE)), + ) + _, tv_layout_rowwise_scale = cute.make_layout_tv( + thr_layout=cute.make_layout(((MXFP8_BLOCK_SCALING_SIZE, 1), WARPS_PER_CTA)), # ((32, 1) 2) + val_layout=cute.make_layout((1, 1)), + ) + _, tv_layout_colwise_scale = cute.make_layout_tv( + thr_layout=cute.make_layout((1, (MXFP8_BLOCK_SCALING_SIZE, WARPS_PER_CTA))), # (1, (32, 2)) + val_layout=cute.make_layout((1, 1)), + ) + + tidx, _, _ = cute.arch.thread_idx() + lane = tidx % MXFP8_BLOCK_SCALING_SIZE + + # Each composed [tidx, None] slice is 1-D (the value mode is flattened): the data + # slices are size 32 (this thread's row segment), the scale slices size 1. + tXsX = cute.composition(sX_tile, tv_layout)[tidx, None] # (32,) input row segment + tXsO_row = cute.composition(sO_row_tile, tv_layout)[tidx, None] # (32,) rowwise out + tXsO_col = cute.composition(sO_col_tile, tv_layout)[tidx, None] # (32,) colwise out + tSsS_row_tile = cute.composition(sS_row_tile, tv_layout_rowwise_scale)[tidx, None] # (1,) + tSsS_col_tile = cute.composition(sS_col_tile, tv_layout_colwise_scale)[tidx, None] # (1,) + + rO_row = cute.make_rmem_tensor(MXFP8_BLOCK_SCALING_SIZE, FP8_DTYPE) + rO_col = cute.make_rmem_tensor(MXFP8_BLOCK_SCALING_SIZE, FP8_DTYPE) + rO_row_u32 = cute.make_tensor( + cute.recast_ptr(rO_row.iterator, dtype=Uint32), + cute.make_layout((MXFP8_BLOCK_SCALING_SIZE // 4,), stride=(1,)), + ) + rO_col_u32 = cute.make_tensor( + cute.recast_ptr(rO_col.iterator, dtype=Uint32), + cute.make_layout((MXFP8_BLOCK_SCALING_SIZE // 4,), stride=(1,)), + ) + sColReduce = cute.make_tensor( + sColReduce_warp.iterator, + cute.make_layout((MXFP8_BLOCK_SCALING_SIZE // 4, 4), stride=(4, 1)), + ) + + if cutlass.const_expr(is_packed16(DTYPE)): + # If the input is bf16 / fp16, take this fast path and process 2 elements at a time in a packed i32 + abs_max_x2 = abs_max_x2_f16 if DTYPE is cutlass.Float16 else abs_max_x2_bf16 + x2_lo_to_f32 = x2_lo_to_f32_f16 if DTYPE is cutlass.Float16 else x2_lo_to_f32_bf16 + x2_hi_to_f32 = x2_hi_to_f32_f16 if DTYPE is cutlass.Float16 else x2_hi_to_f32_bf16 + rX = cute.make_rmem_tensor(MXFP8_BLOCK_SCALING_SIZE, DTYPE) + # Do a vectorized load from SMEM to RMEM and unswizzle in the meantime. + cute.autovec_copy(tXsX, rX) + rX_2x = cute.make_tensor( + cute.recast_ptr(rX.iterator, dtype=Int32), + cute.make_layout((MXFP8_BLOCK_SCALING_SIZE // 2,), stride=(1,)), + ) + sColReduce_2x = cute.make_tensor( + cute.recast_ptr(sColReduce.iterator, dtype=Int64), + cute.make_layout((MXFP8_BLOCK_SCALING_SIZE // 2,), stride=(1,)), + ) + + row_amax2 = rX_2x[0] + for i in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE // 2): + pair = rX_2x[i] + # Skip the first iteration because we assigned rX_2x[0] to row_amax2 already + if cutlass.const_expr(i > 0): + row_amax2 = abs_max_x2(row_amax2, pair) + a_lo = fabs_f32(x2_lo_to_f32(pair)) + a_hi = fabs_f32(x2_hi_to_f32(pair)) + col_lo = cute.arch.warp_redux_sync(a_lo, kind="fmax") + col_hi = cute.arch.warp_redux_sync(a_hi, kind="fmax") + with cute.arch.elect_one(): + sColReduce_2x[i] = pack_f32x2(col_lo, col_hi) + + # Compute the rowwise scale factor + row_amax = cute.arch.fmax( + fabs_f32(x2_lo_to_f32(row_amax2)), fabs_f32(x2_hi_to_f32(row_amax2)) + ) + row_exp = cvt_f32_to_fp8e8m0fnu(row_amax * MAX_NORM_RCP) + row_inv = exp2f_rcp(row_exp) + tSsS_row_tile[0] = row_exp + cute.arch.sync_warp() + + # Compute the colwise scale factor (only handle the one that belongs to this thread / lane) + col_exp = cvt_f32_to_fp8e8m0fnu(sColReduce_warp[lane] * MAX_NORM_RCP) + tSsS_col_tile[0] = col_exp + sColReduce_warp[lane] = exp2f_rcp(col_exp) + cute.arch.sync_warp() + + row_scale_2x = pack_f32x2(row_inv, row_inv) + # Vectorized multiply-and-convert: 4 f32 → 4 fp8 + for j in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE // 4): + # Vectorized load for 4 columnwise scale factors from SMEM to RMEM + col_inv4 = cute.make_rmem_tensor(4, Float32) + cute.autovec_copy(sColReduce[j, None], col_inv4) # LDS.128, warp-broadcast + p01 = rX_2x[2 * j] + p23 = rX_2x[2 * j + 1] + f0 = x2_lo_to_f32(p01) + f1 = x2_hi_to_f32(p01) + f2 = x2_lo_to_f32(p23) + f3 = x2_hi_to_f32(p23) + # For rowwise quantized values, they use the same scale for all 4 elements, + # so we can just pass two row_inv to mul_cvt4 to apply it to all 4 elements at once. + rO_row_u32[j] = mul_cvt4(f0, f1, f2, f3, row_scale_2x) + # For columnwise quantized values, each element has its own scale; the + # elementwise variant fuses the per-element multiply into the cvt sequence. + rO_col_u32[j] = mul_cvt4_elemwise( + f0, f1, f2, f3, col_inv4[0], col_inv4[1], col_inv4[2], col_inv4[3] + ) + cute.autovec_copy(rO_row, tXsO_row) + cute.autovec_copy(rO_col, tXsO_col) + else: + # If input is fp32, take this slow path and process it normally without packing + rX = cute.make_rmem_tensor(MXFP8_BLOCK_SCALING_SIZE, Float32) + cute.autovec_copy(tXsX, rX) + + row_amax = Float32(0.0) + for c in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE): + a = fabs_f32(rX[c]) + row_amax = cute.arch.fmax(row_amax, a) + col_amax = cute.arch.warp_redux_sync(a, kind="fmax") + with cute.arch.elect_one(): + sColReduce_warp[c] = col_amax + + # Compute the rowwise scale factor + row_exp = cvt_f32_to_fp8e8m0fnu(row_amax * MAX_NORM_RCP) + row_inv = exp2f_rcp(row_exp) + tSsS_row_tile[0] = row_exp # rowwise scale (this row-block, staged in smem) + cute.arch.sync_warp() + + # Compute the colwise scale factor (only handle the one that belongs to this thread / lane) + col_exp = cvt_f32_to_fp8e8m0fnu(sColReduce_warp[lane] * MAX_NORM_RCP) + tSsS_col_tile[0] = col_exp # colwise scale (this thread's column) + sColReduce_warp[lane] = exp2f_rcp(col_exp) + cute.arch.sync_warp() + + row_scale_2x = pack_f32x2(row_inv, row_inv) + # Vectorized multiply-and-convert: 4 f32 → 4 fp8 + for j in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE // 4): + # Vectorized load for 4 columnwise scale factors from SMEM to RMEM + col_inv4 = cute.make_rmem_tensor(4, Float32) + cute.autovec_copy(sColReduce[j, None], col_inv4) + offset = 4 * j + # For rowwise quantized values, they use the same scale for all 4 elements, + # so we can just pass two row_inv to mul_cvt4 to apply it to all 4 elements at once. + rO_row_u32[j] = mul_cvt4( + rX[offset], rX[offset + 1], rX[offset + 2], rX[offset + 3], row_scale_2x + ) + # For columnwise quantized values, each element has its own scale; the + # elementwise variant fuses the per-element multiply into the cvt sequence. + rO_col_u32[j] = mul_cvt4_elemwise( + rX[offset], + rX[offset + 1], + rX[offset + 2], + rX[offset + 3], + col_inv4[0], + col_inv4[1], + col_inv4[2], + col_inv4[3], + ) + cute.autovec_copy(rO_row, tXsO_row) + cute.autovec_copy(rO_col, tXsO_col) + + +class MXFP8QuantizeSpecializedBidimensionalKernel(MXFP8QuantizeKernelBase): + """Specialized cast-only rowwise+colwise MXFP8 kernel (swizzled TMA, one 32x32 tile per warp).""" + + _WARPS_PER_CTA = 2 + _THREADS_PER_CTA = _WARPS_PER_CTA * THREADS_PER_WARP + # A warp handles a 32x32 tile + _WARP_ROWS = MXFP8_BLOCK_SCALING_SIZE + _WARP_COLS = MXFP8_BLOCK_SCALING_SIZE + # A CTA handles a tile consisted two 32x32 warp subtiles side-by-side at a time, each handled by a warp + _TILE_ROWS = _WARP_ROWS + _TILE_COLS = _WARP_COLS * _WARPS_PER_CTA + _NUM_STAGES = 2 + _NUM_TILES_X = 4 + _NUM_TILES_Y = 1 + _NUM_TILES = _NUM_TILES_X * _NUM_TILES_Y + + # Rows and columns for the rowwise / columnwise scale factor tensor + _SCALE_ROWS = _TILE_ROWS // MXFP8_BLOCK_SCALING_SIZE + _SCALE_COLS = _TILE_COLS // MXFP8_BLOCK_SCALING_SIZE + + def __init__(self, cfg: MXFP8QuantizeConfig): + self.cfg = cfg + if cfg.WITH_GEMM_SWIZZLED_SCALES: + self._NUM_TILES_X = 2 + self._NUM_TILES_Y = 2 + self._NUM_TILES = self._NUM_TILES_X * self._NUM_TILES_Y + + @cute.jit + def __call__( + self, + mX: cute.Tensor, + mO_row: Optional[cute.Tensor], + mS_row: Optional[cute.Tensor], + mO_col: Optional[cute.Tensor], + mS_col: Optional[cute.Tensor], + mAmax: Optional[cute.Tensor], + mNoop: cute.Pointer, # f32 cast_noop flag; may be null, checked on device + mDActInput: Optional[cute.Tensor], + mWorkspace: Optional[cute.Tensor], + stream: CUstream, + ): + if cutlass.const_expr(CUTEDSL_DEBUG_LOGGING): + cute.printf( + "[CuTeDSL] MXFP8QuantizeSpecializedBidimensionalKernel.__call__() with config:" + f" {self.cfg}\n" + ) + M = mX.shape[0] + N = mX.shape[1] + + if cutlass.const_expr(self.cfg.WITH_GEMM_SWIZZLED_SCALES): + mS_row, mS_col = derive_swizzled_scale_layout( + M, N, self.cfg.ROWWISE, self.cfg.COLWISE, mS_row, mS_col + ) + + smem_tile_layout = cute.make_ordered_layout( + (self._TILE_ROWS, self._TILE_COLS), order=(1, 0) + ) + cta_tiler = (self._TILE_ROWS, self._TILE_COLS) + # Apply 128B input Swizzle<3,4,3> to input tiles + in_smem_layout = cute.make_composed_layout(cute.make_swizzle(3, 4, 3), 0, smem_tile_layout) + # Apply 64B output swizzle<2,4,3> for output tiles + out_smem_layout = cute.make_composed_layout(cute.make_swizzle(2, 4, 3), 0, smem_tile_layout) + op_load = cute.nvgpu.cpasync.CopyBulkTensorTileG2SOp() + op_store = cute.nvgpu.cpasync.CopyBulkTensorTileS2GOp() + + # Input TMA atom + tma_atom, tma_src = cute.nvgpu.cpasync.make_tiled_tma_atom( + op_load, + mX, + in_smem_layout, + cta_tiler, + num_multicast=1, + ) + # Rowwise output TMA atoms + tma_atom_out_row, tma_dst_out_row = cute.nvgpu.cpasync.make_tiled_tma_atom( + op_store, + mO_row, + out_smem_layout, + cta_tiler, + num_multicast=1, + ) + # Colwise output TMA atoms + tma_atom_out_col, tma_dst_out_col = cute.nvgpu.cpasync.make_tiled_tma_atom( + op_store, + mO_col, + out_smem_layout, + cta_tiler, + num_multicast=1, + ) + + grid = [ + # Each CTA covers a (TILE_ROWS * NUM_TILES_Y, TILE_COLS * NUM_TILES_X) area. + cute.ceil_div(Int32(N), self._TILE_COLS * self._NUM_TILES_X), + cute.ceil_div(M, self._TILE_ROWS * self._NUM_TILES_Y), + ] + block = [self._THREADS_PER_CTA] + self.kernel( + mX, + mS_row, + mS_col, + mNoop, + mX.element_type, + tma_atom, + tma_src, + tma_atom_out_row, + tma_dst_out_row, + tma_atom_out_col, + tma_dst_out_col, + ).launch(grid=grid, block=block, stream=stream) + + @cute.kernel + def kernel( + self, + mX: cute.Tensor, + mS_row: cute.Tensor, + mS_col: cute.Tensor, + mNoop: cute.Pointer, + dtype: cutlass.Constexpr[Type[cutlass.Numeric]], + tma_atom: cute.CopyAtom, + tma_src: cute.Tensor, + tma_atom_out_row: cute.CopyAtom, + tma_dst_out_row: cute.Tensor, + tma_atom_out_col: cute.CopyAtom, + tma_dst_out_col: cute.Tensor, + ): + """Device entry: no-op the CTA when the noop flag is set, else run the quantize main loop.""" + if not noop_flag_is_set(mNoop): + self._kernel_main( + mX, + mS_row, + mS_col, + dtype, + tma_atom, + tma_src, + tma_atom_out_row, + tma_dst_out_row, + tma_atom_out_col, + tma_dst_out_col, + ) + + @cute.jit + def _kernel_main( + self, + mX: cute.Tensor, + mS_row: cute.Tensor, + mS_col: cute.Tensor, + dtype: cutlass.Constexpr[Type[cutlass.Numeric]], + tma_atom: cute.CopyAtom, + tma_src: cute.Tensor, + tma_atom_out_row: cute.CopyAtom, + tma_dst_out_row: cute.Tensor, + tma_atom_out_col: cute.CopyAtom, + tma_dst_out_col: cute.Tensor, + ): + """Device entry for the specialized bidimensional (rowwise+colwise) cast kernel.""" + + bidx, bidy, _ = cute.arch.block_idx() + + M = mX.shape[0] + N = mX.shape[1] + + # The CTA owns a NUM_TILES_Y x NUM_TILES_X arrangement of tiles, starting here. + tile_x_base = bidx * self._NUM_TILES_X + tile_y_base = bidy * self._NUM_TILES_Y + + num_tiles_x = cutlass.min( + self._NUM_TILES_X, + # Valid 64-col tiles remaining from this CTA's start column (bidx * span). + cute.ceil_div(Int32(N) - bidx * self._TILE_COLS * self._NUM_TILES_X, self._TILE_COLS), + ) + num_tiles_y = cutlass.min( + self._NUM_TILES_Y, + # Valid 32-row tiles remaining from this CTA's start row (bidy * span). + cute.ceil_div(Int32(M) - bidy * self._TILE_ROWS * self._NUM_TILES_Y, self._TILE_ROWS), + ) + num_tiles = num_tiles_x * num_tiles_y + + FP8_DTYPE = self.cfg.FP8_DTYPE + + @cute.struct + class SharedStorage: + mbar_storage: cute.struct.MemRange[ + cute.Int64, 2 * self._NUM_STAGES + ] # (full,empty) per stage + sX: cute.struct.Align[ + cute.struct.MemRange[dtype, self._TILE_ROWS * self._TILE_COLS * self._NUM_STAGES], + 128, + ] + sO_row: cute.struct.Align[ + cute.struct.MemRange[ + FP8_DTYPE, self._TILE_ROWS * self._TILE_COLS * self._NUM_STAGES + ], + 128, + ] + sO_col: cute.struct.Align[ + cute.struct.MemRange[ + FP8_DTYPE, self._TILE_ROWS * self._TILE_COLS * self._NUM_STAGES + ], + 128, + ] + # Per-warp colwise-reduce scratchpad + sColReduce: cute.struct.Align[ + cute.struct.MemRange[Float32, THREADS_PER_WARP * self._WARPS_PER_CTA], 16 + ] + # Staged rowwise scales for the whole CTA + sScaleRow: cute.struct.Align[ + cute.struct.MemRange[ + Float8E8M0FNU, + self._TILE_ROWS * self._NUM_TILES * self._TILE_COLS // MXFP8_BLOCK_SCALING_SIZE, + ], + 16, + ] + # Staged colwise scales for the whole CTA + sScaleCol: cute.struct.Align[ + cute.struct.MemRange[Float8E8M0FNU, self._NUM_TILES * self._TILE_COLS], 16 + ] + + smem = cutlass.utils.SmemAllocator() + storage = smem.allocate(SharedStorage) + + tile_layout = cute.make_layout( + ((self._TILE_ROWS, self._TILE_COLS), self._NUM_STAGES), + stride=((self._TILE_COLS, 1), self._TILE_ROWS * self._TILE_COLS), + ) + # SMEM input tile should have the same Swizzle<3,4,3> as the input TMA atom + sX = storage.sX.get_tensor(tile_layout, swizzle=cute.make_swizzle(3, 4, 3)) + # SMEM output tiles should have the same Swizzle<2,4,3> as the output TMA atom + sO_row = storage.sO_row.get_tensor(tile_layout, swizzle=cute.make_swizzle(2, 4, 3)) + sO_col = storage.sO_col.get_tensor(tile_layout, swizzle=cute.make_swizzle(2, 4, 3)) + # Reshape the per-warp colwise-reduce scratchpad to a (threads, warps) layout for easier indexing. + sColReduce = storage.sColReduce.get_tensor( + cute.make_layout((THREADS_PER_WARP, self._WARPS_PER_CTA), stride=(1, THREADS_PER_WARP)) + ) + # The whole rowwise scale factor SMEM tensor covered by a CTA + sScaleRow2D = storage.sScaleRow.get_tensor( + cute.make_ordered_layout( + ( + self._TILE_ROWS * self._NUM_TILES_Y, + self._SCALE_COLS * self._NUM_TILES_X, + ), + order=(1, 0), + ) + ) + # Divide by a single tile's shape, so it becomes ((TILE_ROWS, SCALE_COLS), TILES) for easier indexing + # where TILES = NUM_TILES_X * NUM_TILES_Y + sScaleRow = cute.zipped_divide(sScaleRow2D, (self._TILE_ROWS, self._SCALE_COLS)) + # The whole colwise scale factor SMEM tensor covered by a CTA + sScaleCol2D = storage.sScaleCol.get_tensor( + cute.make_ordered_layout( + ( + self._SCALE_ROWS * self._NUM_TILES_Y, + self._TILE_COLS * self._NUM_TILES_X, + ), + order=(1, 0), + ) + ) + # Divide by a single tile's shape, so it becomes ((SCALE_ROWS, TILE_COLS), TILES) for easier indexing + # where TILES = NUM_TILES_X * NUM_TILES_Y + sScaleCol = cute.zipped_divide(sScaleCol2D, (self._SCALE_ROWS, self._TILE_COLS)) + + # Zero the scale SMEM buffers so partial tiles / padding columns flush as 0. + tidx, _, _ = cute.arch.thread_idx() + # View rowwise staging buffers as flat uint32 and stride the CTA over them. + _ROW_SCALE_WORDS = self._TILE_ROWS * self._SCALE_COLS * self._NUM_TILES // 4 + sScaleRow_u32 = cute.make_tensor( + cute.recast_ptr(sScaleRow2D.iterator, dtype=Uint32), + cute.make_layout((_ROW_SCALE_WORDS,), stride=(1,)), + ) + for i in cutlass.range_constexpr(cute.ceil_div(_ROW_SCALE_WORDS, self._THREADS_PER_CTA)): + slot = i * self._THREADS_PER_CTA + tidx + if slot < _ROW_SCALE_WORDS: + sScaleRow_u32[slot] = Uint32(0) + # View colwise staging buffers as flat uint32 and stride the CTA over them. + _COL_SCALE_WORDS = self._SCALE_ROWS * self._TILE_COLS * self._NUM_TILES // 4 + sScaleCol_u32 = cute.make_tensor( + cute.recast_ptr(sScaleCol2D.iterator, dtype=Uint32), + cute.make_layout((_COL_SCALE_WORDS,), stride=(1,)), + ) + for i in cutlass.range_constexpr(cute.ceil_div(_COL_SCALE_WORDS, self._THREADS_PER_CTA)): + slot = i * self._THREADS_PER_CTA + tidx + if slot < _COL_SCALE_WORDS: + sScaleCol_u32[slot] = Uint32(0) + + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + if warp_idx == 0: + cute.nvgpu.cpasync.prefetch_descriptor(tma_atom) + + # Only warp 0 is the producer (issues TMA) + producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, 1) + # Every warp is the consumer + consumer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, self._WARPS_PER_CTA) + # One TMA loads a tile + tx_count = self._TILE_ROWS * self._TILE_COLS * dtype.width // 8 + + mainloop_pipeline = pipeline.PipelineTmaAsync.create( + barrier_storage=storage.mbar_storage.data_ptr(), + num_stages=self._NUM_STAGES, + producer_group=producer_group, + consumer_group=consumer_group, + tx_count=tx_count, + cta_layout_vmnk=None, + ) + + prod_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self._NUM_STAGES + ) + cons_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self._NUM_STAGES + ) + + gX_tiled = cute.zipped_divide(tma_src, (self._TILE_ROWS, self._TILE_COLS)) + tXsX, tXgX = cute.nvgpu.cpasync.tma_partition( + tma_atom, 0, cute.make_layout(1), sX, gX_tiled + ) + gO_row_tiled = cute.zipped_divide(tma_dst_out_row, (self._TILE_ROWS, self._TILE_COLS)) + tXsO_row, tXgO_row = cute.nvgpu.cpasync.tma_partition( + tma_atom_out_row, 0, cute.make_layout(1), sO_row, gO_row_tiled + ) + gO_col_tiled = cute.zipped_divide(tma_dst_out_col, (self._TILE_ROWS, self._TILE_COLS)) + tXsO_col, tXgO_col = cute.nvgpu.cpasync.tma_partition( + tma_atom_out_col, 0, cute.make_layout(1), sO_col, gO_col_tiled + ) + + cute.arch.sync_threads() + + # Prologue: warp 0 prefetches up to NUM_STAGES tiles to fully fill the pipeline + if warp_idx == 0: + for s in cutlass.range_constexpr(self._NUM_STAGES): + if s < num_tiles: + # This tile's position in the CTA's NUM_TILES_Y x NUM_TILES_X arrangement. + tile_y = s // num_tiles_x + tile_x = s % num_tiles_x + mainloop_pipeline.producer_acquire(prod_state) + cute.copy( + tma_atom, + tXgX[(None, (tile_y_base + tile_y, tile_x_base + tile_x))], + tXsX[(None, prod_state.index)], + tma_bar_ptr=mainloop_pipeline.producer_get_barrier(prod_state), + ) + mainloop_pipeline.producer_commit(prod_state) + prod_state.advance() + + # Consumer: all warps fetch from the pipeline, process its tile, and issue a new load to the tile buffer it just consumed + for tile_idx in cutlass.range(num_tiles, unroll=1): + mainloop_pipeline.consumer_wait(cons_state) + # Only allow at most _NUM_STAGES-1 stages to be in-flight, because this iteration will reuse the ring buffer + # that is read _NUM_STAGES iterations ago so we must wait for whoever is reading that buffer to finish + if warp_idx == 0: + cute.arch.cp_async_bulk_wait_group(self._NUM_STAGES - 1, read=True) + cute.arch.sync_threads() + # The current pipeline stage index, which is the tile index modulo the number of stages. + # This is used to index into the shared memory ring buffers that are wrapped around the number of stages. + stage_idx = cons_state.index + + # This tile's position in the CTA's NUM_TILES_Y x NUM_TILES_X arrangement. + tile_y = tile_idx // num_tiles_x + tile_x = tile_idx % num_tiles_x + + # Process the 32x64 SMEM tile for this stage + quantize_bidimensional_mxfp8_swizzled( + sX[None, stage_idx], + sO_row[None, stage_idx], + sO_col[None, stage_idx], + sScaleRow[(None, (tile_y, tile_x))], + sScaleCol[(None, (tile_y, tile_x))], + sColReduce[ + None, warp_idx + ], # Pick the per-warp colwise-reduce scratchpad for this warp + self._WARPS_PER_CTA, + self.cfg.MAX_NORM_RCP, + dtype, + self.cfg.FP8_DTYPE, + ) + + # Make the smem output writes visible to the TMA async proxy, then store. + cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.sync_threads() + + # We are done with this input pipeline SMEM buffer, signal the producer that it can write to this buffer + mainloop_pipeline.consumer_release(cons_state) + + # Issue TMA and write the output to GMEM + if warp_idx == 0: + tile_coord = (tile_y_base + tile_y, tile_x_base + tile_x) + cute.copy( + tma_atom_out_row, tXsO_row[(None, stage_idx)], tXgO_row[(None, tile_coord)] + ) + cute.copy( + tma_atom_out_col, tXsO_col[(None, stage_idx)], tXgO_col[(None, tile_coord)] + ) + cute.arch.cp_async_bulk_commit_group() + + cons_state.advance() + + # Producer: refill the buffer the consumer just freed with the tile + # NUM_STAGES ahead. + next_tile_idx = tile_idx + self._NUM_STAGES + if next_tile_idx < num_tiles: + if warp_idx == 0: + next_tile_y = next_tile_idx // num_tiles_x + next_tile_x = next_tile_idx % num_tiles_x + mainloop_pipeline.producer_acquire(prod_state) + cute.copy( + tma_atom, + tXgX[(None, (tile_y_base + next_tile_y, tile_x_base + next_tile_x))], + tXsX[(None, prod_state.index)], + tma_bar_ptr=mainloop_pipeline.producer_get_barrier(prod_state), + ) + mainloop_pipeline.producer_commit(prod_state) + prod_state.advance() + + # TODO(kainingz): the rowwise WIDTH is hardcoded to 4. Consider using other width if we can convince + # cute.autovec_copy that wider stores are safe (e.g. 16B) and the runtime row pitch allows it. This might + # require us to change `scale_rowwise_shape`'s divisibility when we compile with fake tensors + self._flush_scales_to_gmem( + sScaleRow2D, + mS_row, + tidx, + bidx, + bidy, + ROWS=self._TILE_ROWS * self._NUM_TILES_Y, + COLS=self._SCALE_COLS * self._NUM_TILES_X, + WIDTH=4, + ) + self._flush_scales_to_gmem( + sScaleCol2D, + mS_col, + tidx, + bidx, + bidy, + ROWS=self._SCALE_ROWS * self._NUM_TILES_Y, + COLS=self._TILE_COLS * self._NUM_TILES_X, + WIDTH=1 if self.cfg.WITH_GEMM_SWIZZLED_SCALES else 16, + ) + + # Drain the final stores (their gmem writes must complete before the CTA exits). + cute.arch.cp_async_bulk_wait_group(0, read=False) + + @cute.jit + def _flush_scales_to_gmem( + self, + sScale2D: cute.Tensor, + mS: cute.Tensor, + tidx: Int32, + bidx: Int32, + bidy: Int32, + ROWS: cutlass.Constexpr[int], + COLS: cutlass.Constexpr[int], + WIDTH: cutlass.Constexpr[int], + ): + """Flush a staged (ROWS, COLS) SMEM scale block (a plain row-major 2D tensor) to its (bidy, bidx) slice of the gmem scale tensor, + where each GMEM slice has the same shape as this SMEM tile. Use `WIDTH` bytes per vectorized store. + """ + # Use cute.size instead of cute.shape because under the swizzled layout each mode is a nested tuple ((32, 4, num_tiles_M) etc.) + mS_M = cute.size(mS, mode=[0]) + mS_N = cute.size(mS, mode=[1]) + # Obtain the GMEM slice for the output scale factor block of this CTA + mS_tile = cute.local_tile(mS, (ROWS, COLS), (bidy, bidx)) + + ACTIVE_THREAD_COLS = COLS // WIDTH + _, tv_flush_layout = cute.make_layout_tv( + thr_layout=cute.make_layout((ROWS, ACTIVE_THREAD_COLS), stride=(ACTIVE_THREAD_COLS, 1)), + val_layout=cute.make_layout((1, WIDTH), stride=(WIDTH, 1)), + ) + TOTAL_ACTIVE = ROWS * ACTIVE_THREAD_COLS + + # We may have more total active threads than threads in a CTA, so we do multiple waves of stores to flush the whole buffer + for wave in cutlass.range_constexpr(cute.ceil_div(TOTAL_ACTIVE, self._THREADS_PER_CTA)): + thread_idx = wave * self._THREADS_PER_CTA + tidx + if thread_idx < TOTAL_ACTIVE: + # Find the position for this slot's vectorized store in the GMEM scale factor buffer + thread_y = bidy * ROWS + thread_idx // ACTIVE_THREAD_COLS + thread_x = bidx * COLS + (thread_idx % ACTIVE_THREAD_COLS) * WIDTH + # For rowwise we have N divisible by 4 and WIDTH=4, and for colwise we have N divisible by 128 and WIDTH=16, + # so `thread_x < mS_N` with vectorized store is safe here. + # A thread only writes to a single row so `thread_y < mS_M` is also safe here. + if thread_y < mS_M and thread_x < mS_N: + cute.autovec_copy( + cute.composition(sScale2D, tv_flush_layout)[thread_idx, None], + cute.composition(mS_tile, tv_flush_layout)[thread_idx, None], + ) diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_general.py similarity index 50% rename from transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8.py rename to transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_general.py index 824ed18aa59..26ff4c10b2d 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_general.py @@ -7,6 +7,11 @@ Replicates the core logic of quantize_mxfp8.cuh: given a 2D tensor of BF16/FP16 values, quantize to MXFP8 format (FP8E4M3 data + E8M0 per-block scales). +BF16 cast-only calls with aligned input use the register-resident ports of +specialized/cast_rowwise.cuh and specialized/cast_bidim.cuh. Other configurations +use the one-lane rowwise, staged bidimensional or general kernels. All paths share the +TVM-FFI entrypoint. + Matches the C++ kernel's tile dimensions and thread layout: CHUNK_DIM_Y = 64, CHUNK_DIM_X = 64, THREADS_PER_CTA = 64 BUFF_DIM_Y = 32, BUFF_DIM_X = 64, STAGES = 2 @@ -20,21 +25,15 @@ # Local @cute.struct classes are SMEM-layout descriptors that need no docstrings. # pylint: disable=missing-class-docstring -import abc -import logging -import os from typing import Optional, Type import cutlass from cutlass import cute from cutlass import pipeline -from cutlass import Boolean, Float32, Float8E8M0FNU, Int32, Int64, Uint16, Uint32, Uint8 +from cutlass import Float32, Int32, Int64, Uint16, Uint32, Uint8 from cuda.bindings.driver import CUstream # pylint: disable=no-name-in-module -import tvm_ffi from transformer_engine.common.CuTeDSL.utils import ( - device_compute_capability, - str_to_cutlass_dtype, is_packed16, abs_max_x2_bf16, abs_max_x2_f16, @@ -57,97 +56,24 @@ pack_f32x2, unpack_i64_to_i32x2, ) -from transformer_engine.common.CuTeDSL.activations import ( - act_relu, - act_gelu, - act_silu, - act_qgelu, - act_srelu, - dact_drelu, - dact_dsrelu, - dact_dsilu, - dact_dqgelu, - dact_dgelu, -) from transformer_engine.common.CuTeDSL.utils_fp8 import ( get_cvt_f32x2_to_fp8x2_func, cvt_f32_to_fp8e8m0fnu, mul_f32x2_cvt_f32x4_to_fp8x4, - mul_f32x4_cvt_f32x4_to_fp8x4, mul_f32x2_cvt_packed16x4_to_fp8x4, ) - -CUTEDSL_DEBUG_LOGGING = os.environ.get("CUTEDSL_DEBUG_LOGGING", "0") == "1" - -logger = logging.getLogger("transformer_engine.cutedsl.mxfp8") - -# Number of elements per MXFP8 scale block. They will share the same E8M0 scale factor -MXFP8_BLOCK_SCALING_SIZE = 32 -# How many threads are in one warp -THREADS_PER_WARP = 32 - -# FP8E4M3 max representable value -FP8E4M3_MAX_NORM = 448.0 -FP8E4M3_MAX_NORM_RCP = 1.0 / FP8E4M3_MAX_NORM -FP8E5M2_MAX_NORM = 57344.0 -FP8E5M2_MAX_NORM_RCP = 1.0 / FP8E5M2_MAX_NORM - -# If N is not divisible by 16, writing quantized FP8 output with N bytes per row will fail because -# TMA requires 16-byte alignment for each row -SYM_N_DIVISIBILITY = 16 - -SUPPORTED_ACTIVATIONS = { - "relu": act_relu, - "gelu": act_gelu, - "silu": act_silu, - "qgelu": act_qgelu, - "srelu": act_srelu, -} - -SUPPORTED_DACTIVATIONS = { - "drelu": dact_drelu, - "dgelu": dact_dgelu, - "dsilu": dact_dsilu, - "dqgelu": dact_dqgelu, - "dsrelu": dact_dsrelu, -} - - -@cute.jit -def derive_swizzled_scale_layout( - M: Int32, - N: Int32, - ROWWISE: cutlass.Constexpr[bool], - COLWISE: cutlass.Constexpr[bool], - mS_row: Optional[cute.Tensor], - mS_col: Optional[cute.Tensor], -): - """Derive the swizzled layout for the rowwise and colwise scale tensors.""" - num_scale_cols = cute.ceil_div(N, MXFP8_BLOCK_SCALING_SIZE) - num_scale_rows = cute.ceil_div(M, MXFP8_BLOCK_SCALING_SIZE) - - num_tiles_M = cute.ceil_div(M, 128) - num_tiles_SC = cute.ceil_div(num_scale_cols, 4) - num_tiles_SR = cute.ceil_div(num_scale_rows, 4) - num_tiles_N = cute.ceil_div(N, 128) - - if cutlass.const_expr(ROWWISE): - mS_row = cute.make_tensor( - mS_row.iterator, - cute.make_layout( - ((32, 4, num_tiles_M), (4, num_tiles_SC)), - stride=((16, 4, num_tiles_SC * 512), (1, 512)), - ), - ) - if cutlass.const_expr(COLWISE): - mS_col = cute.make_tensor( - mS_col.iterator, - cute.make_layout( - ((4, num_tiles_SR), (32, 4, num_tiles_N)), - stride=((1, 512), (16, 4, num_tiles_SR * 512)), - ), - ) - return mS_row, mS_col +from .quantize_mxfp8_common import ( + CUTEDSL_DEBUG_LOGGING, + MXFP8QuantizeConfig, + MXFP8QuantizeKernelBase, + MXFP8_BLOCK_SCALING_SIZE, + SUPPORTED_ACTIVATIONS, + SUPPORTED_DACTIVATIONS, + SYM_N_DIVISIBILITY, + THREADS_PER_WARP, + derive_swizzled_scale_layout, + noop_flag_is_set, +) @cute.jit @@ -609,320 +535,6 @@ def dbias_reduction_colwise( return dbias_partial -@cute.jit -def quantize_bidimensional_mxfp8_swizzled( - sX_tile: cute.Tensor, # 32x64 input tile (already sliced to this stage), Sw<3,4,3> swizzled - sO_row_tile: cute.Tensor, # 32x64 rowwise-output tile - sO_col_tile: cute.Tensor, # 32x64 colwise-output tile - sS_row_tile: cute.Tensor, # (32, 2) smem rowwise-scale staging tile (flushed in the epilogue) - sS_col_tile: cute.Tensor, # (1, 64) smem colwise-scale staging tile (flushed in the epilogue) - sColReduce_warp: cute.Tensor, # (32,) SMEM fp32 columnwise scale reduction buffer for this warp - WARPS_PER_CTA: cutlass.Constexpr[int], - MAX_NORM_RCP: cutlass.Constexpr[float], - DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], - FP8_DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], -): - """Quantize a pre-sliced 32x64 tile -> both rowwise and colwise MXFP8. Elements are - addressed through TV layouts: a warp = 32 lanes = 32 rows (lane == row), and each - thread owns one 32-col row segment (its 32-col block) and one output column. - - No bounds check: the caller only loops over valid tiles (`num_tiles`), and - M % 32 == 0 means a tile can only be partial along N. The partial tile's - out-of-bounds warp still runs, but harmlessly -- it reads TMA zero-filled smem, - its output columns are masked by the caller's TMA store, and its scale writes - (rowwise and colwise alike) go to staging slots whose flush targets are past-N - padding columns of the respective scale tensors. - """ - mul_cvt4 = mul_f32x2_cvt_f32x4_to_fp8x4(FP8_DTYPE) - mul_cvt4_elemwise = mul_f32x4_cvt_f32x4_to_fp8x4(FP8_DTYPE) - - _, tv_layout = cute.make_layout_tv( - thr_layout=cute.make_layout(((MXFP8_BLOCK_SCALING_SIZE, 1), WARPS_PER_CTA)), # ((32, 1) 2) - val_layout=cute.make_layout((1, MXFP8_BLOCK_SCALING_SIZE)), - ) - _, tv_layout_rowwise_scale = cute.make_layout_tv( - thr_layout=cute.make_layout(((MXFP8_BLOCK_SCALING_SIZE, 1), WARPS_PER_CTA)), # ((32, 1) 2) - val_layout=cute.make_layout((1, 1)), - ) - _, tv_layout_colwise_scale = cute.make_layout_tv( - thr_layout=cute.make_layout((1, (MXFP8_BLOCK_SCALING_SIZE, WARPS_PER_CTA))), # (1, (32, 2)) - val_layout=cute.make_layout((1, 1)), - ) - - tidx, _, _ = cute.arch.thread_idx() - lane = tidx % MXFP8_BLOCK_SCALING_SIZE - - # Each composed [tidx, None] slice is 1-D (the value mode is flattened): the data - # slices are size 32 (this thread's row segment), the scale slices size 1. - tXsX = cute.composition(sX_tile, tv_layout)[tidx, None] # (32,) input row segment - tXsO_row = cute.composition(sO_row_tile, tv_layout)[tidx, None] # (32,) rowwise out - tXsO_col = cute.composition(sO_col_tile, tv_layout)[tidx, None] # (32,) colwise out - tSsS_row_tile = cute.composition(sS_row_tile, tv_layout_rowwise_scale)[tidx, None] # (1,) - tSsS_col_tile = cute.composition(sS_col_tile, tv_layout_colwise_scale)[tidx, None] # (1,) - - rO_row = cute.make_rmem_tensor(MXFP8_BLOCK_SCALING_SIZE, FP8_DTYPE) - rO_col = cute.make_rmem_tensor(MXFP8_BLOCK_SCALING_SIZE, FP8_DTYPE) - rO_row_u32 = cute.make_tensor( - cute.recast_ptr(rO_row.iterator, dtype=Uint32), - cute.make_layout((MXFP8_BLOCK_SCALING_SIZE // 4,), stride=(1,)), - ) - rO_col_u32 = cute.make_tensor( - cute.recast_ptr(rO_col.iterator, dtype=Uint32), - cute.make_layout((MXFP8_BLOCK_SCALING_SIZE // 4,), stride=(1,)), - ) - sColReduce = cute.make_tensor( - sColReduce_warp.iterator, - cute.make_layout((MXFP8_BLOCK_SCALING_SIZE // 4, 4), stride=(4, 1)), - ) - - if cutlass.const_expr(is_packed16(DTYPE)): - # If the input is bf16 / fp16, take this fast path and process 2 elements at a time in a packed i32 - abs_max_x2 = abs_max_x2_f16 if DTYPE is cutlass.Float16 else abs_max_x2_bf16 - x2_lo_to_f32 = x2_lo_to_f32_f16 if DTYPE is cutlass.Float16 else x2_lo_to_f32_bf16 - x2_hi_to_f32 = x2_hi_to_f32_f16 if DTYPE is cutlass.Float16 else x2_hi_to_f32_bf16 - rX = cute.make_rmem_tensor(MXFP8_BLOCK_SCALING_SIZE, DTYPE) - # Do a vectorized load from SMEM to RMEM and unswizzle in the meantime. - cute.autovec_copy(tXsX, rX) - rX_2x = cute.make_tensor( - cute.recast_ptr(rX.iterator, dtype=Int32), - cute.make_layout((MXFP8_BLOCK_SCALING_SIZE // 2,), stride=(1,)), - ) - sColReduce_2x = cute.make_tensor( - cute.recast_ptr(sColReduce.iterator, dtype=Int64), - cute.make_layout((MXFP8_BLOCK_SCALING_SIZE // 2,), stride=(1,)), - ) - - row_amax2 = rX_2x[0] - for i in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE // 2): - pair = rX_2x[i] - # Skip the first iteration because we assigned rX_2x[0] to row_amax2 already - if cutlass.const_expr(i > 0): - row_amax2 = abs_max_x2(row_amax2, pair) - a_lo = fabs_f32(x2_lo_to_f32(pair)) - a_hi = fabs_f32(x2_hi_to_f32(pair)) - col_lo = cute.arch.warp_redux_sync(a_lo, kind="fmax") - col_hi = cute.arch.warp_redux_sync(a_hi, kind="fmax") - with cute.arch.elect_one(): - sColReduce_2x[i] = pack_f32x2(col_lo, col_hi) - - # Compute the rowwise scale factor - row_amax = cute.arch.fmax( - fabs_f32(x2_lo_to_f32(row_amax2)), fabs_f32(x2_hi_to_f32(row_amax2)) - ) - row_exp = cvt_f32_to_fp8e8m0fnu(row_amax * MAX_NORM_RCP) - row_inv = exp2f_rcp(row_exp) - tSsS_row_tile[0] = row_exp - cute.arch.sync_warp() - - # Compute the colwise scale factor (only handle the one that belongs to this thread / lane) - col_exp = cvt_f32_to_fp8e8m0fnu(sColReduce_warp[lane] * MAX_NORM_RCP) - tSsS_col_tile[0] = col_exp - sColReduce_warp[lane] = exp2f_rcp(col_exp) - cute.arch.sync_warp() - - row_scale_2x = pack_f32x2(row_inv, row_inv) - # Vectorized multiply-and-convert: 4 f32 → 4 fp8 - for j in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE // 4): - # Vectorized load for 4 columnwise scale factors from SMEM to RMEM - col_inv4 = cute.make_rmem_tensor(4, Float32) - cute.autovec_copy(sColReduce[j, None], col_inv4) # LDS.128, warp-broadcast - p01 = rX_2x[2 * j] - p23 = rX_2x[2 * j + 1] - f0 = x2_lo_to_f32(p01) - f1 = x2_hi_to_f32(p01) - f2 = x2_lo_to_f32(p23) - f3 = x2_hi_to_f32(p23) - # For rowwise quantized values, they use the same scale for all 4 elements, - # so we can just pass two row_inv to mul_cvt4 to apply it to all 4 elements at once. - rO_row_u32[j] = mul_cvt4(f0, f1, f2, f3, row_scale_2x) - # For columnwise quantized values, each element has its own scale; the - # elementwise variant fuses the per-element multiply into the cvt sequence. - rO_col_u32[j] = mul_cvt4_elemwise( - f0, f1, f2, f3, col_inv4[0], col_inv4[1], col_inv4[2], col_inv4[3] - ) - cute.autovec_copy(rO_row, tXsO_row) - cute.autovec_copy(rO_col, tXsO_col) - else: - # If input is fp32, take this slow path and process it normally without packing - rX = cute.make_rmem_tensor(MXFP8_BLOCK_SCALING_SIZE, Float32) - cute.autovec_copy(tXsX, rX) - - row_amax = Float32(0.0) - for c in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE): - a = fabs_f32(rX[c]) - row_amax = cute.arch.fmax(row_amax, a) - col_amax = cute.arch.warp_redux_sync(a, kind="fmax") - with cute.arch.elect_one(): - sColReduce_warp[c] = col_amax - - # Compute the rowwise scale factor - row_exp = cvt_f32_to_fp8e8m0fnu(row_amax * MAX_NORM_RCP) - row_inv = exp2f_rcp(row_exp) - tSsS_row_tile[0] = row_exp # rowwise scale (this row-block, staged in smem) - cute.arch.sync_warp() - - # Compute the colwise scale factor (only handle the one that belongs to this thread / lane) - col_exp = cvt_f32_to_fp8e8m0fnu(sColReduce_warp[lane] * MAX_NORM_RCP) - tSsS_col_tile[0] = col_exp # colwise scale (this thread's column) - sColReduce_warp[lane] = exp2f_rcp(col_exp) - cute.arch.sync_warp() - - row_scale_2x = pack_f32x2(row_inv, row_inv) - # Vectorized multiply-and-convert: 4 f32 → 4 fp8 - for j in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE // 4): - # Vectorized load for 4 columnwise scale factors from SMEM to RMEM - col_inv4 = cute.make_rmem_tensor(4, Float32) - cute.autovec_copy(sColReduce[j, None], col_inv4) - offset = 4 * j - # For rowwise quantized values, they use the same scale for all 4 elements, - # so we can just pass two row_inv to mul_cvt4 to apply it to all 4 elements at once. - rO_row_u32[j] = mul_cvt4( - rX[offset], rX[offset + 1], rX[offset + 2], rX[offset + 3], row_scale_2x - ) - # For columnwise quantized values, each element has its own scale; the - # elementwise variant fuses the per-element multiply into the cvt sequence. - rO_col_u32[j] = mul_cvt4_elemwise( - rX[offset], - rX[offset + 1], - rX[offset + 2], - rX[offset + 3], - col_inv4[0], - col_inv4[1], - col_inv4[2], - col_inv4[3], - ) - cute.autovec_copy(rO_row, tXsO_row) - cute.autovec_copy(rO_col, tXsO_col) - - -@cute.jit -def noop_flag_is_set(mNoop: cute.Pointer) -> Boolean: - """Whether the cast_noop flag says this quantization is a no-op and must be skipped. - - mNoop is a pointer rather than a tensor so that one compiled kernel serves both a present and - an absent flag, hence the address is checked before it is dereferenced, exactly like the CUDA - C++ kernel's `noop != nullptr && noop[0] == 1.0f`. The two checks cannot be joined with `and`, - which the DSL lowers to a non-short-circuiting op that would load from the null pointer. - """ - flag_is_set = Boolean(False) - if mNoop.toint() != Int64(0): - flag_is_set = cute.make_tensor(mNoop, cute.make_layout((1,)))[0] == Float32(1.0) - return flag_is_set - - -class MXFP8QuantizeConfig: - """Configs for the compiled CuTeDSL kernel. These will be fixed once the kernel is compiled and - they will behave as const expressions. - """ - - def __init__( - self, - dtype: str, - fp8_dtype: str, - rowwise: bool, - colwise: bool, - with_gemm_swizzled_scales: bool, - with_amax: bool, - with_dbias: bool = False, - with_dact: bool = False, - with_act: bool = False, - use_2d_quantization: bool = False, - activation: Optional[str] = None, - ): - if use_2d_quantization: - raise ValueError("2D block scaling is not implemented by the CuTeDSL MXFP8 kernels") - if dtype is None or dtype not in ("Float32", "Float16", "BFloat16"): - raise ValueError(f"unknown input dtype {dtype!r}; expected Float32|Float16|BFloat16") - self.DTYPE = str_to_cutlass_dtype(dtype) - self.DTYPE_STR = dtype # readable input-dtype token, for __str__ - if fp8_dtype not in ("Float8E4M3", "Float8E5M2"): - raise ValueError( - f"unknown FP8 dtype {fp8_dtype!r}; expected 'Float8E4M3' or 'Float8E5M2'" - ) - self.FP8_DTYPE = str_to_cutlass_dtype(fp8_dtype) - self.FP8_DTYPE_STR = fp8_dtype # readable token, for __str__ - self.ROWWISE = rowwise - self.COLWISE = colwise - if not (rowwise or colwise): - raise ValueError("at least one of rowwise or colwise must be true") - self.WITH_GEMM_SWIZZLED_SCALES = with_gemm_swizzled_scales - self.WITH_AMAX = with_amax - if not with_dact and not with_act: - if activation == "none": - self.ACTIVATION = None - else: - raise ValueError( - "activation must be none when with_dact and with_act are both False" - ) - else: - if with_dact and with_act: - raise ValueError( - "with_dact and with_act cannot be true at the same time since they are used for" - " different paths (bwd vs fwd)" - ) - if with_dact: - if activation in SUPPORTED_DACTIVATIONS: - self.ACTIVATION = activation - else: - raise ValueError( - f"unknown activation {activation!r} for with_dact=True; expected one of" - f" {sorted(SUPPORTED_DACTIVATIONS)}" - ) - elif with_act: - if activation in SUPPORTED_ACTIVATIONS: - self.ACTIVATION = activation - else: - raise ValueError( - f"unknown activation {activation!r} for with_act=True; expected one of" - f" {sorted(SUPPORTED_ACTIVATIONS)}" - ) - self.WITH_DACT = with_dact - self.WITH_ACT = with_act - # dbias is the column reduction of the (post-act/dact) element. With colwise - # output each thread owns a full column (trivial reduction); rowwise-only - # uses a cross-thread smem reduction over THREADS_Y. Both mirror the CUDA - # kernel's COLWISE_SCALING / rowwise dbias branches. - self.WITH_DBIAS = with_dbias - self.MAX_NORM_RCP = ( - FP8E4M3_MAX_NORM_RCP if fp8_dtype == "Float8E4M3" else FP8E5M2_MAX_NORM_RCP - ) - - def __str__(self): - return ( - f"MXFP8QuantizeConfig(dtype={self.DTYPE_STR}, fp8_dtype={self.FP8_DTYPE_STR}, " - f"rowwise={self.ROWWISE}, colwise={self.COLWISE}, " - f"swizzled={self.WITH_GEMM_SWIZZLED_SCALES}, with_amax={self.WITH_AMAX}, " - f"with_dbias={self.WITH_DBIAS}, with_dact={self.WITH_DACT}, " - f"with_act={self.WITH_ACT}, activation={self.ACTIVATION})" - ) - - __repr__ = __str__ - - -class MXFP8QuantizeKernelBase(abc.ABC): - """Base class for MXFP8 quantize kernels.""" - - @abc.abstractmethod - def __call__( - self, - mX: cute.Tensor, - mO_row: Optional[cute.Tensor], - mS_row: Optional[cute.Tensor], - mO_col: Optional[cute.Tensor], - mS_col: Optional[cute.Tensor], - mAmax: Optional[cute.Tensor], - mNoop: cute.Pointer, - mDActInput: Optional[cute.Tensor], - mWorkspace: Optional[cute.Tensor], - stream: CUstream, - ): - """ - Compiled kernel entrypoint (decorate with @cute.jit). - All MXFP8 quantize kernels must implement this interface because this is our C++ call site's contract in `quantize_mxfp8_cutedsl.cuh`. - C++ call site will pass the arguments in this exact order via tvm-ffi and the kernel must accept them even if they are not used. - """ - - class MXFP8QuantizeKernel(MXFP8QuantizeKernelBase): """The MXFP8 quantization kernel that mirrors the standard (non-specialized) MXFP8 CUDA C++ quantization kernel with multiple fusions (activation, dbias, etc.). @@ -1770,1077 +1382,3 @@ def _process_colwise( WITH_DBIAS=self.DBIAS_REDUCTION_IN_COLWISE, CACHE_ACTIVATION=self.CACHE_ACTIVATION, ) - - -class MXFP8QuantizeSpecializedRowwiseKernel(MXFP8QuantizeKernelBase): - """Specialized cast-only ROWWISE-only MXFP8 kernel. - - Requires N % 128 == 0 (full vectorizable column chunks). - - Plain rowwise-only quantize. Each thread owns one 32-element MXFP8 chunk and - uses vectorized global loads/stores (no TMA used).""" - - _TILE_ROWS = 4 - _TILE_COLS = 1024 - _THREADS_PER_CTA = 128 - - def __init__(self, cfg: MXFP8QuantizeConfig): - self.cfg = cfg - # If True, then this kernel will first write each thread's scale byte to a shared memory buffer, - # then utilize vectorized store to flush the buffer to global memory. - self._STASH_SCALE_TO_SMEM = True # Hardcode to true for now - - @cute.jit - def __call__( - self, - mX: cute.Tensor, - mO_row: Optional[cute.Tensor], - mS_row: Optional[cute.Tensor], - mO_col: Optional[cute.Tensor], - mS_col: Optional[cute.Tensor], # Unused, kept for API compatibility - mAmax: Optional[cute.Tensor], # Unused, kept for API compatibility - mNoop: cute.Pointer, # f32 cast_noop flag; may be null, checked on device - mDActInput: Optional[cute.Tensor], # Unused, kept for API compatibility - mWorkspace: Optional[cute.Tensor], # Unused, kept for API compatibility - stream: CUstream, - ): - if cutlass.const_expr(CUTEDSL_DEBUG_LOGGING): - cute.printf( - "[CuTeDSL] MXFP8QuantizeSpecializedRowwiseKernel.__call__() with config:" - f" {self.cfg}\n" - ) - - M = mX.shape[0] - N = mX.shape[1] - - if cutlass.const_expr(self.cfg.WITH_GEMM_SWIZZLED_SCALES): - mS_row, _ = derive_swizzled_scale_layout(M, N, True, False, mS_row, None) - - grid = [ - cute.ceil_div(Int32(N), self._TILE_COLS), - cute.ceil_div(M, self._TILE_ROWS), - ] - block = [self._THREADS_PER_CTA] - - self.kernel( - mX, - mO_row, - mS_row, - mNoop, - mX.element_type, - ).launch(grid=grid, block=block, stream=stream) - - @cute.kernel - def kernel( - self, - mX: cute.Tensor, - mO_row: cute.Tensor, - mS_row: cute.Tensor, - mNoop: cute.Pointer, - DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], - ): - """Device entry: no-op the CTA when the noop flag is set, else run the quantize main loop.""" - if not noop_flag_is_set(mNoop): - self._kernel_main(mX, mO_row, mS_row, DTYPE) - - @cute.jit - def _kernel_main( - self, - mX: cute.Tensor, - mO_row: cute.Tensor, - mS_row: cute.Tensor, - DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], - ): - """Device entry for the specialized rowwise-only cast kernel (vectorized global loads/stores, no TMA).""" - tidx, _, _ = cute.arch.thread_idx() - bidx, bidy, _ = cute.arch.block_idx() - M = mX.shape[0] - N = mX.shape[1] - - # Each thread handles one 32-element MXFP8 chunk (= one scale block). - # The 128 threads in the CTA are grouped as (4, 32), so they cover a - # (4, 1024) input tile and the matching (4, 32) scale tile. - CTA_Y = self._TILE_ROWS - CTA_X = self._TILE_COLS // MXFP8_BLOCK_SCALING_SIZE - tiler, tv_layout = cute.make_layout_tv( - thr_layout=cute.make_layout((CTA_Y, CTA_X), stride=(CTA_X, 1)), - val_layout=cute.make_layout( - (1, MXFP8_BLOCK_SCALING_SIZE), stride=(MXFP8_BLOCK_SCALING_SIZE, 1) - ), - ) - tiler_scale, tv_layout_scale = cute.make_layout_tv( - thr_layout=cute.make_layout((CTA_Y, CTA_X), stride=(CTA_X, 1)), - val_layout=cute.make_layout((1, 1), stride=(1, 1)), - ) - - # Select the tile that belongs to this CTA, then the fragment per thread. - mX_tile = cute.local_tile(mX, tiler, (bidy, bidx)) - mO_tile = cute.local_tile(mO_row, tiler, (bidy, bidx)) - mS_tile = cute.local_tile(mS_row, tiler_scale, (bidy, bidx)) - mX_thread = cute.composition(mX_tile, tv_layout)[tidx, None] - mO_thread = cute.composition(mO_tile, tv_layout)[tidx, None] - mS_thread = cute.composition(mS_tile, tv_layout_scale)[tidx, None] - - rX_thread = cute.make_rmem_tensor( - cute.make_layout((1, MXFP8_BLOCK_SCALING_SIZE), stride=(MXFP8_BLOCK_SCALING_SIZE, 1)), - dtype=DTYPE, - ) - abs_max_x2 = abs_max_x2_f16 if DTYPE is cutlass.Float16 else abs_max_x2_bf16 - x2_lo_to_f32 = x2_lo_to_f32_f16 if DTYPE is cutlass.Float16 else x2_lo_to_f32_bf16 - x2_hi_to_f32 = x2_hi_to_f32_f16 if DTYPE is cutlass.Float16 else x2_hi_to_f32_bf16 - mul_cvt4 = mul_f32x2_cvt_packed16x4_to_fp8x4(DTYPE, self.cfg.FP8_DTYPE) - rX_i32 = cute.make_tensor( - cute.recast_ptr(rX_thread.iterator, dtype=Int32), - cute.make_layout((MXFP8_BLOCK_SCALING_SIZE // 2,), stride=(1,)), - ) - rO_thread = cute.make_rmem_tensor( - cute.make_layout((1, MXFP8_BLOCK_SCALING_SIZE), stride=(MXFP8_BLOCK_SCALING_SIZE, 1)), - dtype=self.cfg.FP8_DTYPE, - ) - rO_u32 = cute.make_tensor( - cute.recast_ptr(rO_thread.iterator, dtype=Uint32), - cute.make_layout( - (MXFP8_BLOCK_SCALING_SIZE // 4,), stride=(1,) - ), # Unit is Uint32, divide by 4 here - ) - - sS_thread = None - if cutlass.const_expr(self._STASH_SCALE_TO_SMEM): - - @cute.struct - class SharedStorage: - buf: cute.struct.Align[cute.struct.MemRange[Float8E8M0FNU, CTA_Y * CTA_X], 16] - - storage = cutlass.utils.SmemAllocator().allocate(SharedStorage) - sScale = storage.buf.get_tensor(cute.make_layout((CTA_Y, CTA_X), stride=(CTA_X, 1))) - # sScale is (CTA_Y, CTA_X):(CTA_X, 1), which is the same layout as tv_layout_scale - # so sS_thread is really just a 1-element e8m0 buffer for this thread's scale. - sS_thread = cute.composition(sScale, tv_layout_scale)[tidx, None] - # Zero first so padding columns (cols past N/32 in the padded scale - # matrix) flush as 0 and we never read uninitialized smem. - # Raw 0x00: e8m0 has no zero value to convert from. - sS_thread[0] = Uint8(0).bitcast(Float8E8M0FNU) - cute.arch.sync_threads() - - row = bidy * self._TILE_ROWS + tidx // CTA_X - col = bidx * self._TILE_COLS + (tidx % CTA_X) * MXFP8_BLOCK_SCALING_SIZE - if row < M and col < N: - cute.autovec_copy(mX_thread, rX_thread) - amax_2x = rX_i32[0] - # Skip the first iteration because we assigned rX_i32[0] to amax_2x already - for i in cutlass.range_constexpr(1, MXFP8_BLOCK_SCALING_SIZE // 2): - amax_2x = abs_max_x2(amax_2x, rX_i32[i]) - amax = cute.arch.fmax(fabs_f32(x2_lo_to_f32(amax_2x)), fabs_f32(x2_hi_to_f32(amax_2x))) - - biased_exp = cvt_f32_to_fp8e8m0fnu(amax * self.cfg.MAX_NORM_RCP) - if cutlass.const_expr(self._STASH_SCALE_TO_SMEM): - sS_thread[0] = biased_exp - else: - mS_thread[0] = biased_exp - - # Rescale + FP8 cast, 4 elements per fused mul_cvt (one uint32 out), - # then a vectorized store. Mirrors CUDA's _use_cvt_4x path. - inv_scale = exp2f_rcp(biased_exp) - scale_2x = pack_f32x2(inv_scale, inv_scale) - for i in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE // 4): - rO_u32[i] = mul_cvt4(rX_i32[2 * i], rX_i32[2 * i + 1], scale_2x) - cute.autovec_copy(rO_thread, mO_thread) - - # Cooperative wide flush of the staged scales where padding columns flush as 0. - if cutlass.const_expr(self._STASH_SCALE_TO_SMEM): - cute.arch.sync_threads() - # Use cute.size instead of .shape[1] because under the swizzled layout mode 1 is - # a nested tuple ((4, num_tiles_SC)), not a plain int. - padded_cols = cute.size(mS_row, mode=[1]) - if cutlass.const_expr(self.cfg.WITH_GEMM_SWIZZLED_SCALES): - # Swizzled rowwise scale layout is ((32, 4, num_tiles_M), (4, num_tiles_SC)):((16, 4, num_tiles_SC * 512), (1, 512)) - # so we have at most 4 elements continuous which is our vectorized store width - self._flush_scales_to_gmem(sScale, mS_tile, tidx, bidx, bidy, M, padded_cols, 4) - elif padded_cols % 16 == 0: - # If columns is divisible by 16, use 16 bytes as the vectorized store width - self._flush_scales_to_gmem(sScale, mS_tile, tidx, bidx, bidy, M, padded_cols, 16) - else: - # Otherwise use 4 bytes as the vectorized store width. - # Note our fake tensor requires 4 divisibility so this is enforced as long as you can get here - self._flush_scales_to_gmem(sScale, mS_tile, tidx, bidx, bidy, M, padded_cols, 4) - - @cute.jit - def _flush_scales_to_gmem( - self, - sScale: cute.Tensor, - mS_tile: cute.Tensor, - tidx: Int32, - bidx: Int32, - bidy: Int32, - M: Int32, - padded_cols: Int32, - width: cutlass.Constexpr[int], - ): - """Flush the staged (CTA_Y, CTA_X) scale tile to gmem with vectorized stores.""" - CTA_Y = self._TILE_ROWS - CTA_X = self._TILE_COLS // MXFP8_BLOCK_SCALING_SIZE - # Previously each threads has 1 byte, but now we are doing vectorized store, - # which means only a subset of threads will need to issue the store while other threads are not used. - active_threads = CTA_X // width - _, tv_flush = cute.make_layout_tv( - thr_layout=cute.make_layout((CTA_Y, active_threads), stride=(active_threads, 1)), - val_layout=cute.make_layout((1, width), stride=(width, 1)), - ) - # We only need to use a subset of threads with shape (CTA_Y, active_threads) to write - # so if the thread is outside of this subset, it will remain inactive - if tidx < CTA_Y * active_threads: - # Absolute position of the scale vector to write in the GMEM buffer - thread_y = bidy * CTA_Y + tidx // active_threads - thread_x = bidx * CTA_X + (tidx % active_threads) * width - if thread_y < M and thread_x < padded_cols: - cute.autovec_copy( - cute.composition(sScale, tv_flush)[tidx, None], - cute.composition(mS_tile, tv_flush)[tidx, None], - ) - - -class MXFP8QuantizeSpecializedBidimensionalKernel(MXFP8QuantizeKernelBase): - """Specialized cast-only rowwise+colwise MXFP8 kernel (swizzled TMA, one 32x32 tile per warp).""" - - _WARPS_PER_CTA = 2 - _THREADS_PER_CTA = _WARPS_PER_CTA * THREADS_PER_WARP - # A warp handles a 32x32 tile - _WARP_ROWS = MXFP8_BLOCK_SCALING_SIZE - _WARP_COLS = MXFP8_BLOCK_SCALING_SIZE - # A CTA handles a tile consisted two 32x32 warp subtiles side-by-side at a time, each handled by a warp - _TILE_ROWS = _WARP_ROWS - _TILE_COLS = _WARP_COLS * _WARPS_PER_CTA - _NUM_STAGES = 2 - _NUM_TILES_X = 4 - _NUM_TILES_Y = 1 - _NUM_TILES = _NUM_TILES_X * _NUM_TILES_Y - - # Rows and columns for the rowwise / columnwise scale factor tensor - _SCALE_ROWS = _TILE_ROWS // MXFP8_BLOCK_SCALING_SIZE - _SCALE_COLS = _TILE_COLS // MXFP8_BLOCK_SCALING_SIZE - - def __init__(self, cfg: MXFP8QuantizeConfig): - self.cfg = cfg - if cfg.WITH_GEMM_SWIZZLED_SCALES: - self._NUM_TILES_X = 2 - self._NUM_TILES_Y = 2 - self._NUM_TILES = self._NUM_TILES_X * self._NUM_TILES_Y - - @cute.jit - def __call__( - self, - mX: cute.Tensor, - mO_row: Optional[cute.Tensor], - mS_row: Optional[cute.Tensor], - mO_col: Optional[cute.Tensor], - mS_col: Optional[cute.Tensor], - mAmax: Optional[cute.Tensor], - mNoop: cute.Pointer, # f32 cast_noop flag; may be null, checked on device - mDActInput: Optional[cute.Tensor], - mWorkspace: Optional[cute.Tensor], - stream: CUstream, - ): - if cutlass.const_expr(CUTEDSL_DEBUG_LOGGING): - cute.printf( - "[CuTeDSL] MXFP8QuantizeSpecializedBidimensionalKernel.__call__() with config:" - f" {self.cfg}\n" - ) - M = mX.shape[0] - N = mX.shape[1] - - if cutlass.const_expr(self.cfg.WITH_GEMM_SWIZZLED_SCALES): - mS_row, mS_col = derive_swizzled_scale_layout( - M, N, self.cfg.ROWWISE, self.cfg.COLWISE, mS_row, mS_col - ) - - smem_tile_layout = cute.make_ordered_layout( - (self._TILE_ROWS, self._TILE_COLS), order=(1, 0) - ) - cta_tiler = (self._TILE_ROWS, self._TILE_COLS) - # Apply 128B input Swizzle<3,4,3> to input tiles - in_smem_layout = cute.make_composed_layout(cute.make_swizzle(3, 4, 3), 0, smem_tile_layout) - # Apply 64B output swizzle<2,4,3> for output tiles - out_smem_layout = cute.make_composed_layout(cute.make_swizzle(2, 4, 3), 0, smem_tile_layout) - op_load = cute.nvgpu.cpasync.CopyBulkTensorTileG2SOp() - op_store = cute.nvgpu.cpasync.CopyBulkTensorTileS2GOp() - - # Input TMA atom - tma_atom, tma_src = cute.nvgpu.cpasync.make_tiled_tma_atom( - op_load, - mX, - in_smem_layout, - cta_tiler, - num_multicast=1, - ) - # Rowwise output TMA atoms - tma_atom_out_row, tma_dst_out_row = cute.nvgpu.cpasync.make_tiled_tma_atom( - op_store, - mO_row, - out_smem_layout, - cta_tiler, - num_multicast=1, - ) - # Colwise output TMA atoms - tma_atom_out_col, tma_dst_out_col = cute.nvgpu.cpasync.make_tiled_tma_atom( - op_store, - mO_col, - out_smem_layout, - cta_tiler, - num_multicast=1, - ) - - grid = [ - # Each CTA covers a (TILE_ROWS * NUM_TILES_Y, TILE_COLS * NUM_TILES_X) area. - cute.ceil_div(Int32(N), self._TILE_COLS * self._NUM_TILES_X), - cute.ceil_div(M, self._TILE_ROWS * self._NUM_TILES_Y), - ] - block = [self._THREADS_PER_CTA] - self.kernel( - mX, - mS_row, - mS_col, - mNoop, - mX.element_type, - tma_atom, - tma_src, - tma_atom_out_row, - tma_dst_out_row, - tma_atom_out_col, - tma_dst_out_col, - ).launch(grid=grid, block=block, stream=stream) - - @cute.kernel - def kernel( - self, - mX: cute.Tensor, - mS_row: cute.Tensor, - mS_col: cute.Tensor, - mNoop: cute.Pointer, - dtype: cutlass.Constexpr[Type[cutlass.Numeric]], - tma_atom: cute.CopyAtom, - tma_src: cute.Tensor, - tma_atom_out_row: cute.CopyAtom, - tma_dst_out_row: cute.Tensor, - tma_atom_out_col: cute.CopyAtom, - tma_dst_out_col: cute.Tensor, - ): - """Device entry: no-op the CTA when the noop flag is set, else run the quantize main loop.""" - if not noop_flag_is_set(mNoop): - self._kernel_main( - mX, - mS_row, - mS_col, - dtype, - tma_atom, - tma_src, - tma_atom_out_row, - tma_dst_out_row, - tma_atom_out_col, - tma_dst_out_col, - ) - - @cute.jit - def _kernel_main( - self, - mX: cute.Tensor, - mS_row: cute.Tensor, - mS_col: cute.Tensor, - dtype: cutlass.Constexpr[Type[cutlass.Numeric]], - tma_atom: cute.CopyAtom, - tma_src: cute.Tensor, - tma_atom_out_row: cute.CopyAtom, - tma_dst_out_row: cute.Tensor, - tma_atom_out_col: cute.CopyAtom, - tma_dst_out_col: cute.Tensor, - ): - """Device entry for the specialized bidimensional (rowwise+colwise) cast kernel.""" - - bidx, bidy, _ = cute.arch.block_idx() - - M = mX.shape[0] - N = mX.shape[1] - - # The CTA owns a NUM_TILES_Y x NUM_TILES_X arrangement of tiles, starting here. - tile_x_base = bidx * self._NUM_TILES_X - tile_y_base = bidy * self._NUM_TILES_Y - - num_tiles_x = cutlass.min( - self._NUM_TILES_X, - # Valid 64-col tiles remaining from this CTA's start column (bidx * span). - cute.ceil_div(Int32(N) - bidx * self._TILE_COLS * self._NUM_TILES_X, self._TILE_COLS), - ) - num_tiles_y = cutlass.min( - self._NUM_TILES_Y, - # Valid 32-row tiles remaining from this CTA's start row (bidy * span). - cute.ceil_div(Int32(M) - bidy * self._TILE_ROWS * self._NUM_TILES_Y, self._TILE_ROWS), - ) - num_tiles = num_tiles_x * num_tiles_y - - FP8_DTYPE = self.cfg.FP8_DTYPE - - @cute.struct - class SharedStorage: - mbar_storage: cute.struct.MemRange[ - cute.Int64, 2 * self._NUM_STAGES - ] # (full,empty) per stage - sX: cute.struct.Align[ - cute.struct.MemRange[dtype, self._TILE_ROWS * self._TILE_COLS * self._NUM_STAGES], - 128, - ] - sO_row: cute.struct.Align[ - cute.struct.MemRange[ - FP8_DTYPE, self._TILE_ROWS * self._TILE_COLS * self._NUM_STAGES - ], - 128, - ] - sO_col: cute.struct.Align[ - cute.struct.MemRange[ - FP8_DTYPE, self._TILE_ROWS * self._TILE_COLS * self._NUM_STAGES - ], - 128, - ] - # Per-warp colwise-reduce scratchpad - sColReduce: cute.struct.Align[ - cute.struct.MemRange[Float32, THREADS_PER_WARP * self._WARPS_PER_CTA], 16 - ] - # Staged rowwise scales for the whole CTA - sScaleRow: cute.struct.Align[ - cute.struct.MemRange[ - Float8E8M0FNU, - self._TILE_ROWS * self._NUM_TILES * self._TILE_COLS // MXFP8_BLOCK_SCALING_SIZE, - ], - 16, - ] - # Staged colwise scales for the whole CTA - sScaleCol: cute.struct.Align[ - cute.struct.MemRange[Float8E8M0FNU, self._NUM_TILES * self._TILE_COLS], 16 - ] - - smem = cutlass.utils.SmemAllocator() - storage = smem.allocate(SharedStorage) - - tile_layout = cute.make_layout( - ((self._TILE_ROWS, self._TILE_COLS), self._NUM_STAGES), - stride=((self._TILE_COLS, 1), self._TILE_ROWS * self._TILE_COLS), - ) - # SMEM input tile should have the same Swizzle<3,4,3> as the input TMA atom - sX = storage.sX.get_tensor(tile_layout, swizzle=cute.make_swizzle(3, 4, 3)) - # SMEM output tiles should have the same Swizzle<2,4,3> as the output TMA atom - sO_row = storage.sO_row.get_tensor(tile_layout, swizzle=cute.make_swizzle(2, 4, 3)) - sO_col = storage.sO_col.get_tensor(tile_layout, swizzle=cute.make_swizzle(2, 4, 3)) - # Reshape the per-warp colwise-reduce scratchpad to a (threads, warps) layout for easier indexing. - sColReduce = storage.sColReduce.get_tensor( - cute.make_layout((THREADS_PER_WARP, self._WARPS_PER_CTA), stride=(1, THREADS_PER_WARP)) - ) - # The whole rowwise scale factor SMEM tensor covered by a CTA - sScaleRow2D = storage.sScaleRow.get_tensor( - cute.make_ordered_layout( - ( - self._TILE_ROWS * self._NUM_TILES_Y, - self._SCALE_COLS * self._NUM_TILES_X, - ), - order=(1, 0), - ) - ) - # Divide by a single tile's shape, so it becomes ((TILE_ROWS, SCALE_COLS), TILES) for easier indexing - # where TILES = NUM_TILES_X * NUM_TILES_Y - sScaleRow = cute.zipped_divide(sScaleRow2D, (self._TILE_ROWS, self._SCALE_COLS)) - # The whole colwise scale factor SMEM tensor covered by a CTA - sScaleCol2D = storage.sScaleCol.get_tensor( - cute.make_ordered_layout( - ( - self._SCALE_ROWS * self._NUM_TILES_Y, - self._TILE_COLS * self._NUM_TILES_X, - ), - order=(1, 0), - ) - ) - # Divide by a single tile's shape, so it becomes ((SCALE_ROWS, TILE_COLS), TILES) for easier indexing - # where TILES = NUM_TILES_X * NUM_TILES_Y - sScaleCol = cute.zipped_divide(sScaleCol2D, (self._SCALE_ROWS, self._TILE_COLS)) - - # Zero the scale SMEM buffers so partial tiles / padding columns flush as 0. - tidx, _, _ = cute.arch.thread_idx() - # View rowwise staging buffers as flat uint32 and stride the CTA over them. - _ROW_SCALE_WORDS = self._TILE_ROWS * self._SCALE_COLS * self._NUM_TILES // 4 - sScaleRow_u32 = cute.make_tensor( - cute.recast_ptr(sScaleRow2D.iterator, dtype=Uint32), - cute.make_layout((_ROW_SCALE_WORDS,), stride=(1,)), - ) - for i in cutlass.range_constexpr(cute.ceil_div(_ROW_SCALE_WORDS, self._THREADS_PER_CTA)): - slot = i * self._THREADS_PER_CTA + tidx - if slot < _ROW_SCALE_WORDS: - sScaleRow_u32[slot] = Uint32(0) - # View colwise staging buffers as flat uint32 and stride the CTA over them. - _COL_SCALE_WORDS = self._SCALE_ROWS * self._TILE_COLS * self._NUM_TILES // 4 - sScaleCol_u32 = cute.make_tensor( - cute.recast_ptr(sScaleCol2D.iterator, dtype=Uint32), - cute.make_layout((_COL_SCALE_WORDS,), stride=(1,)), - ) - for i in cutlass.range_constexpr(cute.ceil_div(_COL_SCALE_WORDS, self._THREADS_PER_CTA)): - slot = i * self._THREADS_PER_CTA + tidx - if slot < _COL_SCALE_WORDS: - sScaleCol_u32[slot] = Uint32(0) - - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - if warp_idx == 0: - cute.nvgpu.cpasync.prefetch_descriptor(tma_atom) - - # Only warp 0 is the producer (issues TMA) - producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, 1) - # Every warp is the consumer - consumer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, self._WARPS_PER_CTA) - # One TMA loads a tile - tx_count = self._TILE_ROWS * self._TILE_COLS * dtype.width // 8 - - mainloop_pipeline = pipeline.PipelineTmaAsync.create( - barrier_storage=storage.mbar_storage.data_ptr(), - num_stages=self._NUM_STAGES, - producer_group=producer_group, - consumer_group=consumer_group, - tx_count=tx_count, - cta_layout_vmnk=None, - ) - - prod_state = pipeline.make_pipeline_state( - pipeline.PipelineUserType.Producer, self._NUM_STAGES - ) - cons_state = pipeline.make_pipeline_state( - pipeline.PipelineUserType.Consumer, self._NUM_STAGES - ) - - gX_tiled = cute.zipped_divide(tma_src, (self._TILE_ROWS, self._TILE_COLS)) - tXsX, tXgX = cute.nvgpu.cpasync.tma_partition( - tma_atom, 0, cute.make_layout(1), sX, gX_tiled - ) - gO_row_tiled = cute.zipped_divide(tma_dst_out_row, (self._TILE_ROWS, self._TILE_COLS)) - tXsO_row, tXgO_row = cute.nvgpu.cpasync.tma_partition( - tma_atom_out_row, 0, cute.make_layout(1), sO_row, gO_row_tiled - ) - gO_col_tiled = cute.zipped_divide(tma_dst_out_col, (self._TILE_ROWS, self._TILE_COLS)) - tXsO_col, tXgO_col = cute.nvgpu.cpasync.tma_partition( - tma_atom_out_col, 0, cute.make_layout(1), sO_col, gO_col_tiled - ) - - cute.arch.sync_threads() - - # Prologue: warp 0 prefetches up to NUM_STAGES tiles to fully fill the pipeline - if warp_idx == 0: - for s in cutlass.range_constexpr(self._NUM_STAGES): - if s < num_tiles: - # This tile's position in the CTA's NUM_TILES_Y x NUM_TILES_X arrangement. - tile_y = s // num_tiles_x - tile_x = s % num_tiles_x - mainloop_pipeline.producer_acquire(prod_state) - cute.copy( - tma_atom, - tXgX[(None, (tile_y_base + tile_y, tile_x_base + tile_x))], - tXsX[(None, prod_state.index)], - tma_bar_ptr=mainloop_pipeline.producer_get_barrier(prod_state), - ) - mainloop_pipeline.producer_commit(prod_state) - prod_state.advance() - - # Consumer: all warps fetch from the pipeline, process its tile, and issue a new load to the tile buffer it just consumed - for tile_idx in cutlass.range(num_tiles, unroll=1): - mainloop_pipeline.consumer_wait(cons_state) - # Only allow at most _NUM_STAGES-1 stages to be in-flight, because this iteration will reuse the ring buffer - # that is read _NUM_STAGES iterations ago so we must wait for whoever is reading that buffer to finish - if warp_idx == 0: - cute.arch.cp_async_bulk_wait_group(self._NUM_STAGES - 1, read=True) - cute.arch.sync_threads() - # The current pipeline stage index, which is the tile index modulo the number of stages. - # This is used to index into the shared memory ring buffers that are wrapped around the number of stages. - stage_idx = cons_state.index - - # This tile's position in the CTA's NUM_TILES_Y x NUM_TILES_X arrangement. - tile_y = tile_idx // num_tiles_x - tile_x = tile_idx % num_tiles_x - - # Process the 32x64 SMEM tile for this stage - quantize_bidimensional_mxfp8_swizzled( - sX[None, stage_idx], - sO_row[None, stage_idx], - sO_col[None, stage_idx], - sScaleRow[(None, (tile_y, tile_x))], - sScaleCol[(None, (tile_y, tile_x))], - sColReduce[ - None, warp_idx - ], # Pick the per-warp colwise-reduce scratchpad for this warp - self._WARPS_PER_CTA, - self.cfg.MAX_NORM_RCP, - dtype, - self.cfg.FP8_DTYPE, - ) - - # Make the smem output writes visible to the TMA async proxy, then store. - cute.arch.fence_proxy("async.shared", space="cta") - cute.arch.sync_threads() - - # We are done with this input pipeline SMEM buffer, signal the producer that it can write to this buffer - mainloop_pipeline.consumer_release(cons_state) - - # Issue TMA and write the output to GMEM - if warp_idx == 0: - tile_coord = (tile_y_base + tile_y, tile_x_base + tile_x) - cute.copy( - tma_atom_out_row, tXsO_row[(None, stage_idx)], tXgO_row[(None, tile_coord)] - ) - cute.copy( - tma_atom_out_col, tXsO_col[(None, stage_idx)], tXgO_col[(None, tile_coord)] - ) - cute.arch.cp_async_bulk_commit_group() - - cons_state.advance() - - # Producer: refill the buffer the consumer just freed with the tile - # NUM_STAGES ahead. - next_tile_idx = tile_idx + self._NUM_STAGES - if next_tile_idx < num_tiles: - if warp_idx == 0: - next_tile_y = next_tile_idx // num_tiles_x - next_tile_x = next_tile_idx % num_tiles_x - mainloop_pipeline.producer_acquire(prod_state) - cute.copy( - tma_atom, - tXgX[(None, (tile_y_base + next_tile_y, tile_x_base + next_tile_x))], - tXsX[(None, prod_state.index)], - tma_bar_ptr=mainloop_pipeline.producer_get_barrier(prod_state), - ) - mainloop_pipeline.producer_commit(prod_state) - prod_state.advance() - - # TODO(kainingz): the rowwise WIDTH is hardcoded to 4. Consider using other width if we can convince - # cute.autovec_copy that wider stores are safe (e.g. 16B) and the runtime row pitch allows it. This might - # require us to change `scale_rowwise_shape`'s divisibility when we compile with fake tensors - self._flush_scales_to_gmem( - sScaleRow2D, - mS_row, - tidx, - bidx, - bidy, - ROWS=self._TILE_ROWS * self._NUM_TILES_Y, - COLS=self._SCALE_COLS * self._NUM_TILES_X, - WIDTH=4, - ) - self._flush_scales_to_gmem( - sScaleCol2D, - mS_col, - tidx, - bidx, - bidy, - ROWS=self._SCALE_ROWS * self._NUM_TILES_Y, - COLS=self._TILE_COLS * self._NUM_TILES_X, - WIDTH=1 if self.cfg.WITH_GEMM_SWIZZLED_SCALES else 16, - ) - - # Drain the final stores (their gmem writes must complete before the CTA exits). - cute.arch.cp_async_bulk_wait_group(0, read=False) - - @cute.jit - def _flush_scales_to_gmem( - self, - sScale2D: cute.Tensor, - mS: cute.Tensor, - tidx: Int32, - bidx: Int32, - bidy: Int32, - ROWS: cutlass.Constexpr[int], - COLS: cutlass.Constexpr[int], - WIDTH: cutlass.Constexpr[int], - ): - """Flush a staged (ROWS, COLS) SMEM scale block (a plain row-major 2D tensor) to its (bidy, bidx) slice of the gmem scale tensor, - where each GMEM slice has the same shape as this SMEM tile. Use `WIDTH` bytes per vectorized store. - """ - # Use cute.size instead of cute.shape because under the swizzled layout each mode is a nested tuple ((32, 4, num_tiles_M) etc.) - mS_M = cute.size(mS, mode=[0]) - mS_N = cute.size(mS, mode=[1]) - # Obtain the GMEM slice for the output scale factor block of this CTA - mS_tile = cute.local_tile(mS, (ROWS, COLS), (bidy, bidx)) - - ACTIVE_THREAD_COLS = COLS // WIDTH - _, tv_flush_layout = cute.make_layout_tv( - thr_layout=cute.make_layout((ROWS, ACTIVE_THREAD_COLS), stride=(ACTIVE_THREAD_COLS, 1)), - val_layout=cute.make_layout((1, WIDTH), stride=(WIDTH, 1)), - ) - TOTAL_ACTIVE = ROWS * ACTIVE_THREAD_COLS - - # We may have more total active threads than threads in a CTA, so we do multiple waves of stores to flush the whole buffer - for wave in cutlass.range_constexpr(cute.ceil_div(TOTAL_ACTIVE, self._THREADS_PER_CTA)): - thread_idx = wave * self._THREADS_PER_CTA + tidx - if thread_idx < TOTAL_ACTIVE: - # Find the position for this slot's vectorized store in the GMEM scale factor buffer - thread_y = bidy * ROWS + thread_idx // ACTIVE_THREAD_COLS - thread_x = bidx * COLS + (thread_idx % ACTIVE_THREAD_COLS) * WIDTH - # For rowwise we have N divisible by 4 and WIDTH=4, and for colwise we have N divisible by 128 and WIDTH=16, - # so `thread_x < mS_N` with vectorized store is safe here. - # A thread only writes to a single row so `thread_y < mS_M` is also safe here. - if thread_y < mS_M and thread_x < mS_N: - cute.autovec_copy( - cute.composition(sScale2D, tv_flush_layout)[thread_idx, None], - cute.composition(mS_tile, tv_flush_layout)[thread_idx, None], - ) - - -class MXFP8QuantizeEntry(MXFP8QuantizeKernelBase): - """Select the appropriate MXFP8 quantization kernel based on the configuration and runtime shapes.""" - - def __init__(self, cfg: MXFP8QuantizeConfig): - self.cfg = cfg - # These activation functions satisfy f(0) = 0 so we don't need to mask with OOB regions - # (zeros filled by TMA are still zeros without applying activation to them) - # Note that as the time of writing, all activations supported by MXFP8 satisfy this property, so !ACT_NEED_MASKING is never reachable. - # This is only a precaution for future use cases if someone wants to add a new activation function. - self.ACT_NEED_MASKING = cfg.WITH_ACT and cfg.ACTIVATION not in ( - "relu", - "gelu", - "silu", - "qgelu", - "srelu", - ) - # Instantiate all possible kernels at compile time, - # and we will pick the right one at runtime based on the input shape and config. - self.general_skip_input_masking_kernel = MXFP8QuantizeKernel(cfg, SKIP_INPUT_MASKING=True) - self.general_skip_input_masking_and_scale_bounds_kernel = MXFP8QuantizeKernel( - cfg, SKIP_INPUT_MASKING=True, SKIP_SCALE_BOUNDS=True - ) - # We only need to mask when WITH_ACT is enabled because with activations applied zeros filled by TMA affect block statistics - self.general_kernel = MXFP8QuantizeKernel(cfg) if self.ACT_NEED_MASKING else None - self.specialized_rowwise = ( - MXFP8QuantizeSpecializedRowwiseKernel(cfg) if cfg.ROWWISE else None - ) - self.specialized_bidim = ( - MXFP8QuantizeSpecializedBidimensionalKernel(cfg) - if cfg.ROWWISE and cfg.COLWISE - else None - ) - - @cute.jit - def __call__( - self, - mX: cute.Tensor, - mO_row: Optional[cute.Tensor], - mS_row: Optional[cute.Tensor], - mO_col: Optional[cute.Tensor], - mS_col: Optional[cute.Tensor], - mAmax: Optional[cute.Tensor], - mNoop: cute.Pointer, - mDActInput: Optional[cute.Tensor], - mWorkspace: Optional[cute.Tensor], - stream: CUstream, - ): - M = mX.shape[0] - N = mX.shape[1] - # Select specialized kernels if possible - plain_cast_only = ( - not self.cfg.WITH_AMAX - and not self.cfg.WITH_DBIAS - and not self.cfg.WITH_DACT - and not self.cfg.WITH_ACT - ) - dispatched_to_specialized = False - # Only dispatch to the specialized kernels for packed16 types (bf16/fp16) - if cutlass.const_expr(plain_cast_only and is_packed16(self.cfg.DTYPE)): - if cutlass.const_expr(self.cfg.ROWWISE and not self.cfg.COLWISE): - rowwise_specialized_grid_fits = ( - cute.ceil_div(M, self.specialized_rowwise._TILE_ROWS) <= 65535 - ) - # The rowwise specialized kernel requires N divisible by 128 for vectorized stores - if N % 128 == 0 and rowwise_specialized_grid_fits: - dispatched_to_specialized = True - self.specialized_rowwise( - mX, - mO_row, - mS_row, - mO_col, - mS_col, - mAmax, - mNoop, - mDActInput, - mWorkspace, - stream, - ) - # The bidimensional specialized kernel supports swizzled scales but we don't dispatch to it for now - if cutlass.const_expr( - self.cfg.ROWWISE and self.cfg.COLWISE and not self.cfg.WITH_GEMM_SWIZZLED_SCALES - ): - bidim_specialized_grid_fits = ( - cute.ceil_div( - M, - self.specialized_bidim._TILE_ROWS * self.specialized_bidim._NUM_TILES_Y, - ) - <= 65535 - ) - if bidim_specialized_grid_fits: - dispatched_to_specialized = True - self.specialized_bidim( - mX, - mO_row, - mS_row, - mO_col, - mS_col, - mAmax, - mNoop, - mDActInput, - mWorkspace, - stream, - ) - # If not using a specialized kernel, fall back to the general kernel - if not dispatched_to_specialized: - # If the input shape can be perfectly tiled by the general kernel's tile size, we can skip some boundary check because - # we know we will not touch any out-of-bounds region. - shape_is_divisible = ( - mX.shape[0] % self.general_skip_input_masking_kernel._TILE_ROWS == 0 - and mX.shape[1] % self.general_skip_input_masking_kernel._TILE_COLS == 0 - ) - # If shape is divisible, then we won't access OOB regions regardless whatever - if shape_is_divisible: - self.general_skip_input_masking_and_scale_bounds_kernel( - mX, - mO_row, - mS_row, - mO_col, - mS_col, - mAmax, - mNoop, - mDActInput, - mWorkspace, - stream, - ) - else: - # Otherwise we can't skip scale bound check, but we may still skip masking activation inputs with zero - if cutlass.const_expr(self.ACT_NEED_MASKING): - # Masking and the scale guards now share one condition, so this is a plain - # two-way choice: either the shape tiles exactly (drop both) or it does not. - self.general_kernel( - mX, - mO_row, - mS_row, - mO_col, - mS_col, - mAmax, - mNoop, - mDActInput, - mWorkspace, - stream, - ) - else: - # We still need to check the scale bounds, but we can skip some activations masking because their output is 0 - # for OOB regions where TMA fills with zeros - self.general_skip_input_masking_kernel( - mX, - mO_row, - mS_row, - mO_col, - mS_col, - mAmax, - mNoop, - mDActInput, - mWorkspace, - stream, - ) - - -def compile_cutedsl_function_from_cfg(cfg): - """ - Return the compiled CuTeDSL function object for the given MXFP8 quantization config. - """ - - kernel_obj = MXFP8QuantizeEntry(cfg) - sym_M = cute.sym_int32() - sym_N = cute.sym_int32(divisibility=SYM_N_DIVISIBILITY) - in_shape = out_shape = (sym_M, sym_N) - # TE allocates scale tensors at a padded shape (see - # MXFP8Quantizer::get_scale_shape in transformer_engine/pytorch/csrc): - # rowwise: (roundup(M, 128), roundup(ceildiv(N, 32), 4)) - # columnwise: (roundup(ceildiv(M, 32), 4), roundup(N, 128)) - # These padded extents are NOT M/N (and SymInt has no `//`/`+`), so give the - # scales their own fresh syms carrying the divisibility the padding - # guarantees (rowwise: 128 x 4; colwise: 4 x 128). - scale_rowwise_shape = (cute.sym_int32(divisibility=128), cute.sym_int32(divisibility=4)) - scale_colwise_shape = (cute.sym_int32(divisibility=4), cute.sym_int32(divisibility=128)) - ws_shape = (cute.sym_int32(), sym_N) # (blocks_Y, N); N ties to input N - # Native FP8/E8M0 dtypes at the FFI boundary (matches the DLPack dtype the C++ - # bridge sends). - out_dtype = cfg.FP8_DTYPE - scale_dtype = cutlass.Float8E8M0FNU - - in_fake = cute.runtime.make_fake_compact_tensor( - cfg.DTYPE, in_shape, stride_order=(1, 0), memspace=cute.AddressSpace.gmem, assumed_align=16 - ) - out_row_fake = ( - cute.runtime.make_fake_compact_tensor( - out_dtype, - out_shape, - stride_order=(1, 0), - memspace=cute.AddressSpace.gmem, - assumed_align=16, - ) - if cfg.ROWWISE - else None - ) - scale_row_fake = ( - cute.runtime.make_fake_compact_tensor( - scale_dtype, - scale_rowwise_shape, - stride_order=(1, 0), - memspace=cute.AddressSpace.gmem, - assumed_align=4, - ) - if cfg.ROWWISE - else None - ) - out_col_fake = ( - cute.runtime.make_fake_compact_tensor( - out_dtype, - out_shape, - stride_order=(1, 0), - memspace=cute.AddressSpace.gmem, - assumed_align=16, - ) - if cfg.COLWISE - else None - ) - scale_col_fake = ( - cute.runtime.make_fake_compact_tensor( - scale_dtype, - scale_colwise_shape, - stride_order=(1, 0), - memspace=cute.AddressSpace.gmem, - assumed_align=4, - ) - if cfg.COLWISE - else None - ) - amax_fake = ( - cute.runtime.make_fake_compact_tensor( - Float32, (1,), stride_order=(0,), memspace=cute.AddressSpace.gmem, assumed_align=4 - ) - if cfg.WITH_AMAX - else None - ) - # The cast-noop flag is an always-present f32 pointer instead of an optional tensor, so that - # the same compiled kernel serves both an absent and a present flag (see noop_flag_is_set). - # The null address here is only a compile-time placeholder for the pointer the C++ dispatcher - # passes at launch. - noop_fake = cute.runtime.nullptr(Float32, mem_space=cute.AddressSpace.gmem, assumed_align=4) - act_input_fake = ( - cute.runtime.make_fake_compact_tensor( - cfg.DTYPE, - in_shape, - stride_order=(1, 0), - memspace=cute.AddressSpace.gmem, - assumed_align=16, - ) - if cfg.WITH_DACT - else None - ) - workspace_fake = ( - cute.runtime.make_fake_compact_tensor( - Float32, ws_shape, stride_order=(1, 0), memspace=cute.AddressSpace.gmem, assumed_align=4 - ) - if cfg.WITH_DBIAS - else None - ) - - compiled = cute.compile( - kernel_obj, - in_fake, # mX - out_row_fake, - scale_row_fake, # mO_row, mS_row - out_col_fake, - scale_col_fake, # mO_col, mS_col - amax_fake, # mAmax - noop_fake, # mNoop (pointer to the cast_noop flag, may be null at launch) - act_input_fake, # mDActInput (backward slot, unused) - workspace_fake, # mWorkspace(backward slot, unused) - cute.runtime.make_fake_stream(), # stream (compiled as an explicit tvm-ffi - # "handle" arg; C++ passes the CUDA stream - # as void*) - options="--enable-tvm-ffi", - ) - return compiled - - -def get_mxfp8_quantization_function( - fn_name: str, - dtype: str, - fp8_dtype: str, - rowwise: bool, - colwise: bool, - with_gemm_swizzled_scales: bool, - with_amax: bool, - with_dbias: bool, - with_dact: bool, - with_act: bool, - use_2d_quantization: bool, - activation: str, -) -> bool: - """Compile the MXFP8 quantize kernel for this config and register it in the TVM-FFI global registry - under EXACTLY `fn_name` (the key the C++ dispatcher built; Python treats it as an opaque name). - Returns True if a kernel is successfully registered under `fn_name` (the C++ side then fetches it with GetGlobal(fn_name)); - False if the config is unsupported, so the caller caches the negative result and falls back to the CUDA C++ kernel. - """ - try: - # Already registered (e.g. by a prior call) -> supported. - if tvm_ffi.get_global_func(fn_name, allow_missing=True) is not None: - return True - - major, minor = device_compute_capability() - if major < 10: - logger.warning( - "CuTeDSL MXFP8 backend requires compute capability >= 10.0 (Blackwell), " - "but detected %d.%d; falling back to the CUDA C++ kernel.", - major, - minor, - ) - return False - - try: - cfg = MXFP8QuantizeConfig( - dtype=dtype, - fp8_dtype=fp8_dtype, - rowwise=rowwise, - colwise=colwise, - with_gemm_swizzled_scales=with_gemm_swizzled_scales, - with_amax=with_amax, - with_dbias=with_dbias, - with_dact=with_dact, - with_act=with_act, - use_2d_quantization=use_2d_quantization, - activation=activation, - ) - except ValueError as e: - # The exception message states exactly why the config is unsupported - # (unknown dtype/activation, dbias not implemented, ...). Surfacing it as a - # warning lets the C++ dispatcher's CUDA fallback be recognized as expected. - logger.warning( - "CuTeDSL MXFP8 backend does not support this config, " - "falling back to the CUDA C++ kernel: %s", - e, - ) - return False - - logger.debug("Compiling CuTeDSL MXFP8 quantization kernel for %s", cfg) - compiled = compile_cutedsl_function_from_cfg(cfg) - # The returned compiled object is not neccessarily the compiled function itself. It could be a Python - # wrapper that parses arguments and calls the underlying function. This is needed in the general case - # as TVM-FFI is positional-only, however, for us this is unnecessary launch overhead as we pass the - # arguments in the correct order on the C++ side. - # Relevant issue: https://github.com/NVIDIA/cutlass/issues/3527 - # https://github.com/NVIDIA/cutlass/pull/3589 recommends using an explicit positional only marker (/) - # in the function signature, however it is not merged and it is something easy to miss and silently - # introduce unnecessarry overhead. - native = getattr(compiled, "__tvm_ffi_object__", lambda: None)() - tvm_ffi.register_global_func( - fn_name, native if native is not None else compiled, override=True - ) - return True - except Exception as e: # pylint: disable=broad-exception-caught - logger.error( - "CuTeDSL MXFP8 kernel compilation & registration failed, falling back to the CUDA" - " C++ kernel: %s", - e, - ) - # Unconditionally fallback to CUDA path because we can't tell if this exception is transient or permanent. - return False diff --git a/transformer_engine/common/CuTeDSL/utils.py b/transformer_engine/common/CuTeDSL/utils.py index ed05aa1c3c6..80a90a85ab8 100644 --- a/transformer_engine/common/CuTeDSL/utils.py +++ b/transformer_engine/common/CuTeDSL/utils.py @@ -83,6 +83,15 @@ def exp2f_rcp(scale_e8m0) -> Float32: return result +@cute.jit +def exp2_bf16x2_rcp(scale_e8m0) -> Uint32: + """Pack two BF16 copies of the E8M0 reciprocal computed by exp2f_rcp.""" + # These reciprocal powers of two are exact in BF16. Taking the upper bits + # also preserves exp2f_rcp's chosen subnormal and NaN encodings. + bf16_bits = exp2f_rcp(scale_e8m0).bitcast(Uint32) >> Uint32(16) + return bf16_bits | (bf16_bits << Uint32(16)) + + @cute.jit def pack_f32x2(lo: Float32, hi: Float32) -> Int64: """Pack two f32 scalars into a single 64-bit register (`floatx2` layout). @@ -258,3 +267,14 @@ def truncate_f32_f16(val: Float32) -> Float32: def is_packed16(dtype) -> bool: """True if `dtype` is one of the 16-bit packed input formats.""" return dtype is cutlass.BFloat16 or dtype is cutlass.Float16 + + +@cute.jit +def create_l2_policy(evict_last: cutlass.Constexpr[bool], fraction: Float32 = 1.0) -> Int64: + """Create the L2 eviction policy used by register-resident casts.""" + priority = "evict_last" if evict_last else "evict_first" + return cute.arch.inline_ptx( + f"createpolicy.fractional.L2::{priority}.b64 {{$w0}}, {{$r0}};", + write_only_types=[Int64], + read_only_args=[Float32(fraction)], + ) diff --git a/transformer_engine/common/CuTeDSL/utils_fp8.py b/transformer_engine/common/CuTeDSL/utils_fp8.py index b3b73e7998d..9a972e39469 100644 --- a/transformer_engine/common/CuTeDSL/utils_fp8.py +++ b/transformer_engine/common/CuTeDSL/utils_fp8.py @@ -180,6 +180,49 @@ def mul_f32x4_cvt_f32x4_to_fp8x4(fp8_dtype, relu: bool = False) -> Callable[..., return _build_mul_f32x4_cvt_f32x4_to_fp8x4(fp8_dtype, relu) +def mul_bf16x2x2_cvt_bf16x4_to_fp8x4(fp8_dtype) -> Callable[..., Uint32]: + """Return a multiply/cast for two BF16 pairs, producing four packed FP8 bytes. + + CUDA 13.1+ supports direct BF16-to-FP8 conversion. Older compilers use + FP32 fused multiply-add, matching the corresponding ptx::mul_cvt_4x overload. + """ + out_op = "e5m2x2" if fp8_dtype is Float8E5M2 else "e4m3x2" + if cutlass.target_version(min_version="13.1"): + body = ( + ".reg.b32 y0, y1;\n\t" + "mul.rn.bf16x2 y0, {$r0}, {$r1};\n\t" + "mul.rn.bf16x2 y1, {$r2}, {$r3};\n\t" + f"cvt.rn.satfinite.{out_op}.bf16x2 z0, y0;\n\t" + f"cvt.rn.satfinite.{out_op}.bf16x2 z1, y1;\n\t" + ) + else: + body = ( + ".reg.b16 x0,x1,x2,x3,s0,s1,s2,s3;\n\t" + ".reg.f32 y0,y1,y2,y3;\n\t" + "mov.b32 {x0,x1}, {$r0};\n\t" + "mov.b32 {s0,s1}, {$r1};\n\t" + "mov.b32 {x2,x3}, {$r2};\n\t" + "mov.b32 {s2,s3}, {$r3};\n\t" + "fma.rn.f32.bf16 y0, x0, s0, 0f00000000;\n\t" + "fma.rn.f32.bf16 y1, x1, s1, 0f00000000;\n\t" + "fma.rn.f32.bf16 y2, x2, s2, 0f00000000;\n\t" + "fma.rn.f32.bf16 y3, x3, s3, 0f00000000;\n\t" + f"cvt.rn.satfinite.{out_op}.f32 z0, y1, y0;\n\t" + f"cvt.rn.satfinite.{out_op}.f32 z1, y3, y2;\n\t" + ) + asm = "{\n.reg.b16 z0,z1;\n\t" + body + "mov.b32 {$w0}, {z0,z1};\n}" + + @cute.jit + def fn(in0: Int32, scale0: Uint32, in1: Int32, scale1: Uint32) -> Uint32: + return cute.arch.inline_ptx( + asm, + write_only_types=[Uint32], + read_only_args=[in0, scale0, in1, scale1], + ) + + return fn + + def _build_mul_f32x2_cvt_packed16x4_to_fp8x4( in_dtype, fp8_dtype, relu: bool = False ) -> Callable[..., Uint32]: From 1be6447b126772e48d575508a83f1fde41d837c1 Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Thu, 8 Oct 2026 21:42:05 +0000 Subject: [PATCH 22/58] wip Signed-off-by: Kaining Zhong --- docs/envvars.rst | 2 +- .../cast/mxfp8/quantize_mxfp8_common.py | 3 - .../mxfp8/quantize_mxfp8_g2r_bidimensional.py | 541 +++++++++--------- .../mxfp8/quantize_mxfp8_g2r_rowwise_1lane.py | 2 +- .../mxfp8/quantize_mxfp8_g2r_rowwise_2lane.py | 136 +++-- .../mxfp8/quantize_mxfp8_g2s_bidimensional.py | 2 +- .../cast/mxfp8/quantize_mxfp8_general.py | 4 +- transformer_engine/common/CuTeDSL/utils.py | 47 ++ 8 files changed, 407 insertions(+), 330 deletions(-) diff --git a/docs/envvars.rst b/docs/envvars.rst index 46b70bbe46a..14f562f1069 100644 --- a/docs/envvars.rst +++ b/docs/envvars.rst @@ -150,7 +150,7 @@ General :Type: ``int`` (0 or 1) :Default: ``0`` :Description: Use CuTeDSL kernel implementations when available. This requires that Transformer Engine was built with ``NVTE_WITH_CUTEDSL=1``. - It also requires ``apache-tvm-ffi``, ``nvidia-cutlass-dsl`` to be installed and you are using CPython 3.10 or later. + It also requires ``apache-tvm-ffi`` and ``nvidia-cutlass-dsl>=4.8.0`` to be installed and you are using CPython 3.10 or later. If these requirements are not met, Transformer Engine will fall back to using CUDA C++ kernels if possible. .. envvar:: NVTE_WARN_IF_CUTEDSL_BACKEND_NOT_CHOSEN diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_common.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_common.py index f281937ffca..249e178abd6 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_common.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_common.py @@ -43,9 +43,6 @@ # How many threads are in one warp THREADS_PER_WARP = 32 -# Prefer the current allocator API while retaining compatibility with older DSL releases. -_SmemAllocator = getattr(cutlass, "memory", cutlass.utils).SmemAllocator - # FP8E4M3 max representable value FP8E4M3_MAX_NORM = 448.0 FP8E4M3_MAX_NORM_RCP = 1.0 / FP8E4M3_MAX_NORM diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py index 0713d68f0fd..4a1692df68f 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py @@ -13,64 +13,22 @@ from transformer_engine.common.CuTeDSL.utils import ( abs_max_x2_bf16, + bf16_mx_scale_reciprocal, + bf16_mx_scale_reciprocal_pair, + bf16_mx_scale_bytes_pair, create_l2_policy, ) from transformer_engine.common.CuTeDSL.utils_fp8 import mul_bf16x2x2_cvt_bf16x4_to_fp8x4 from .quantize_mxfp8_common import ( + MXFP8_BLOCK_SCALING_SIZE, + THREADS_PER_WARP, MXFP8QuantizeConfig, MXFP8QuantizeKernelBase, - _SmemAllocator, bf16_pair_magnitude, noop_flag_is_set, ) -@cute.jit -def bf16_mx_scale_reciprocal(magnitude: Uint32, MAX_EXPONENT: cutlass.Constexpr[int]) -> Uint32: - """Exact BF16 reciprocal scale bits, including E8M0 zero, Inf and NaN cases.""" - rounded = (magnitude + Uint32(31)) & Uint32(0x7F80) - if rounded < Uint32(MAX_EXPONENT << 7): # pylint: disable=consider-using-max-builtin - rounded = Uint32(MAX_EXPONENT << 7) - reciprocal = Uint32((254 + MAX_EXPONENT) << 7) - rounded - if magnitude >= Uint32(0x7F80): - reciprocal = Uint32(0x0040) - if magnitude != Uint32(0x7F80): - reciprocal = Uint32(0x7FFF) - return reciprocal - - -@cute.jit -def bf16_mx_scale_reciprocal_pair(pair: Int32, MAX_EXPONENT: cutlass.Constexpr[int]) -> Uint32: - """Compute two reciprocal scales without borrowing across packed BF16 halves.""" - magnitudes = Uint32(pair) & Uint32(0x7FFF7FFF) - rounded = (magnitudes + Uint32(0x001F001F)) & Uint32(0xFF80FF80) - floor = (MAX_EXPONENT << 23) | (MAX_EXPONENT << 7) - # Packed BF16 max accepts registers, so materialize the constant inside PTX. - clamped = cute.arch.inline_ptx( - "{ .reg.b32 floor_pair; " - f"mov.b32 floor_pair, {floor}; " - "max.xorsign.abs.bf16x2 {$w0}, {$r0}, floor_pair; }", - write_only_types=[Int32], - read_only_args=[rounded.bitcast(Int32)], - ) - bias = ((254 + MAX_EXPONENT) << 23) | ((254 + MAX_EXPONENT) << 7) - reciprocal = Uint32(bias) - Uint32(clamped) - if ((magnitudes + Uint32(0x00800080)) & Uint32(0x80008000)) != Uint32(0): - lo = bf16_mx_scale_reciprocal(magnitudes & Uint32(0x7FFF), MAX_EXPONENT) - hi = bf16_mx_scale_reciprocal(magnitudes >> Uint32(16), MAX_EXPONENT) - reciprocal = lo | (hi << Uint32(16)) - return reciprocal - - -@cute.jit -def bf16_mx_scale_bytes_pair(reciprocal: Uint32) -> Uint32: - """Gather the E8M0 bytes of two packed BF16 reciprocals into the low half.""" - bytes_pair = (Uint32(0x01FE01FE) - ((reciprocal >> Uint32(7)) & Uint32(0x00FF00FF))) & Uint32( - 0x00FF00FF - ) - return (bytes_pair & Uint32(0xFF)) | ((bytes_pair >> Uint32(8)) & Uint32(0xFF00)) - - @cute.jit def quantize_bidimensional_mxfp8_g2r( gX_tile: cute.Tensor, @@ -81,97 +39,117 @@ def quantize_bidimensional_mxfp8_g2r( cfg: cutlass.Constexpr[MXFP8QuantizeConfig], CTA_THREADS_Y: cutlass.Constexpr[int], CTA_THREADS_X: cutlass.Constexpr[int], - ROWS_PER_THREAD: cutlass.Constexpr[int], - ELEMENTS_PER_THREAD: cutlass.Constexpr[int], + ELEMENTS_Y_PER_THREAD: cutlass.Constexpr[int], + ELEMENTS_X_PER_THREAD: cutlass.Constexpr[int], cache_policy: Int64, ): """Quantize one (TILE_ROWS, TILE_COLS) GMEM tile in both directions. - TILE_ROWS = CTA_THREADS_Y * ROWS_PER_THREAD = 32, so every columnwise - MXFP8 block is contained in the tile. Each warp owns ROWS_PER_THREAD - consecutive rows, and each lane owns ELEMENTS_PER_THREAD adjacent columns. + TILE_ROWS = CTA_THREADS_Y * ELEMENTS_Y_PER_THREAD = 32, so every columnwise + MXFP8 block is contained in the tile. Each warp owns ELEMENTS_Y_PER_THREAD + consecutive rows, and each lane owns ELEMENTS_X_PER_THREAD adjacent columns. Load the tile once, then reuse its registers for both quantization passes. """ - TILE_ROWS = CTA_THREADS_Y * ROWS_PER_THREAD - TILE_COLS = CTA_THREADS_X * ELEMENTS_PER_THREAD - LANES_PER_MX_BLOCK = 32 // ELEMENTS_PER_THREAD - COLUMN_PAIRS = TILE_COLS // 2 - COLUMN_PAIRS_PER_THREAD = ELEMENTS_PER_THREAD // 2 - PAIRS_PER_HALF = COLUMN_PAIRS_PER_THREAD // 2 + # Shape of the tile + TILE_ROWS = CTA_THREADS_Y * ELEMENTS_Y_PER_THREAD + TILE_COLS = CTA_THREADS_X * ELEMENTS_X_PER_THREAD + + # How many threads cooperate to process a single MXFP8 block (32 elements) for both quantized directions + THREADS_X_PER_MX_BLOCK = MXFP8_BLOCK_SCALING_SIZE // ELEMENTS_X_PER_THREAD + THREADS_Y_PER_MX_BLOCK = MXFP8_BLOCK_SCALING_SIZE // ELEMENTS_Y_PER_THREAD + MAX_EXPONENT = 8 if cfg.FP8_DTYPE is cutlass.Float8E4M3FN else 15 SCALE_DTYPE = cutlass.Float8E8M0FNU + tidx, _, _ = cute.arch.thread_idx() - warp, lane = tidx // CTA_THREADS_X, tidx % CTA_THREADS_X - _, tv_payload = cute.make_layout_tv( - thr_layout=cute.make_layout((CTA_THREADS_Y, CTA_THREADS_X), stride=(CTA_THREADS_X, 1)), + thread_layout = cute.make_layout((CTA_THREADS_Y, CTA_THREADS_X), stride=(CTA_THREADS_X, 1)) + _, tv_layout = cute.make_layout_tv( + thr_layout=thread_layout, val_layout=cute.make_layout( - (ROWS_PER_THREAD, ELEMENTS_PER_THREAD), stride=(ELEMENTS_PER_THREAD, 1) + (ELEMENTS_Y_PER_THREAD, ELEMENTS_X_PER_THREAD), stride=(ELEMENTS_X_PER_THREAD, 1) ), ) - # Reshape each thread's values into (columns, rows), with contiguous columns. - fragment_layout = cute.make_layout( - (ELEMENTS_PER_THREAD, ROWS_PER_THREAD), stride=(1, ELEMENTS_PER_THREAD) - ) - gX_thread = cute.composition(cute.composition(gX_tile, tv_payload)[tidx, None], fragment_layout) - gO_row_thread = cute.composition( - cute.composition(gO_row_tile, tv_payload)[tidx, None], fragment_layout - ) - gO_col_thread = cute.composition( - cute.composition(gO_col_tile, tv_payload)[tidx, None], fragment_layout - ) - # Each group of lanes shares one rowwise scale for each of its rows. + + gX_thread = cute.composition(gX_tile, tv_layout)[tidx, None] + gO_row_thread = cute.composition(gO_row_tile, tv_layout)[tidx, None] + gO_col_thread = cute.composition(gO_col_tile, tv_layout)[tidx, None] + + # Group CTA threads by THREADS_X_PER_MX_BLOCK since they share the same rowwise scale for a MXFP8 block. + # And assign stride 0 to that submode tv_row_scale = cute.make_layout( ( - (LANES_PER_MX_BLOCK, (CTA_THREADS_X // LANES_PER_MX_BLOCK, CTA_THREADS_Y)), - ROWS_PER_THREAD, + (THREADS_X_PER_MX_BLOCK, (CTA_THREADS_X // THREADS_X_PER_MX_BLOCK, CTA_THREADS_Y)), + ELEMENTS_Y_PER_THREAD, + ), + stride=( + (0, (TILE_ROWS, ELEMENTS_Y_PER_THREAD)), + 1, ), - stride=((0, (TILE_ROWS, ROWS_PER_THREAD)), 1), ) gS_row_thread = cute.composition(gS_row_tile, tv_row_scale)[tidx, None] - # Warp 0 writes one scale per column, in each lane's original column order. - gS_col_thread = cute.local_tile(gS_col_tile, (1, ELEMENTS_PER_THREAD), (0, lane))[0, None] - # Keep the complete input fragment; reuse one row's output registers. - rX = cute.make_rmem_tensor((ELEMENTS_PER_THREAD, ROWS_PER_THREAD), cfg.DTYPE) - rO = cute.make_rmem_tensor((ELEMENTS_PER_THREAD,), cfg.FP8_DTYPE) + # Group CTA threads by THREADS_Y_PER_MX_BLOCK since they share the same columnwise scale for a MXFP8 block. + # And assign stride 0 to that submode + tv_col_scale = cute.make_layout( + ( + (CTA_THREADS_X, (THREADS_Y_PER_MX_BLOCK, CTA_THREADS_Y // THREADS_Y_PER_MX_BLOCK)), + ELEMENTS_X_PER_THREAD, + ), + stride=( + (ELEMENTS_X_PER_THREAD, (0, 1)), + 1, + ), + ) + gS_col_thread = cute.composition(gS_col_tile, tv_col_scale)[tidx, None] + + # Input fragment contains all rows it handles; output fragment contains only one row and we write it back row by row + rX = cute.make_rmem_tensor((ELEMENTS_X_PER_THREAD, ELEMENTS_Y_PER_THREAD), cfg.DTYPE) + rO = cute.make_rmem_tensor((ELEMENTS_X_PER_THREAD,), cfg.FP8_DTYPE) + + # For rowwise scales, we write one row at a time rS_row = cute.make_rmem_tensor((1,), SCALE_DTYPE) - rS_col = cute.make_rmem_tensor((ELEMENTS_PER_THREAD,), SCALE_DTYPE) + load_atom = cute.make_copy_atom( cute.nvgpu.CopyG2ROp(), cfg.DTYPE, - num_bits_per_copy=ELEMENTS_PER_THREAD * cfg.DTYPE.width, + num_bits_per_copy=ELEMENTS_X_PER_THREAD * cfg.DTYPE.width, invariant=True, ) store_atom = cute.make_copy_atom( cute.nvgpu.CopyR2GOp(), cfg.FP8_DTYPE, - num_bits_per_copy=ELEMENTS_PER_THREAD * cfg.FP8_DTYPE.width, + num_bits_per_copy=ELEMENTS_X_PER_THREAD * cfg.FP8_DTYPE.width, ) + row_scale_atom = cute.make_copy_atom(cute.nvgpu.CopyR2GOp(), SCALE_DTYPE, num_bits_per_copy=8) col_scale_atom = cute.make_copy_atom( - cute.nvgpu.CopyR2GOp(), SCALE_DTYPE, num_bits_per_copy=ELEMENTS_PER_THREAD * 8 + cute.nvgpu.CopyR2GOp(), SCALE_DTYPE, num_bits_per_copy=ELEMENTS_X_PER_THREAD * 8 ) + scale_and_pack = mul_bf16x2x2_cvt_bf16x4_to_fp8x4(cfg.FP8_DTYPE) - rO_u32 = cute.recast_tensor(rO, Uint32) - for row_idx in cutlass.range_constexpr(ROWS_PER_THREAD): + # Copy all rows a thread sweeps into the RMEM fragment at once + for row_idx in cutlass.range_constexpr(ELEMENTS_Y_PER_THREAD): cute.copy(load_atom, gX_thread[None, row_idx], rX[None, row_idx], cache_policy=cache_policy) @cute.jit - def extract_row_amax(values_i32: cute.Tensor) -> Uint32: + def extract_mx_block_amax(values_i32: cute.Tensor) -> Uint32: """Reduce a 32-column MXFP8 block across its two or four lane owners.""" amax = values_i32[0] - for pair_idx in cutlass.range_constexpr(1, cute.size(values_i32)): - amax = abs_max_x2_bf16(amax, values_i32[pair_idx]) - for stage in cutlass.range_constexpr(LANES_PER_MX_BLOCK.bit_length() - 1): + for bf16x2_idx in cutlass.range_constexpr(1, cute.size(values_i32)): + amax = abs_max_x2_bf16(amax, values_i32[bf16x2_idx]) + for stage in cutlass.range_constexpr(THREADS_X_PER_MX_BLOCK.bit_length() - 1): amax = abs_max_x2_bf16(amax, cute.arch.shuffle_sync_bfly(amax, 1 << stage)) return bf16_pair_magnitude(amax) - # Rowwise pass: lane shuffles close each 32-column reduction. - for row_idx in cutlass.range_constexpr(ROWS_PER_THREAD): + # Rowwise pass: produce each row's rowwise scales and quantized outputs, and write them back to GMEM immediately. + rO_u32 = cute.recast_tensor(rO, Uint32) + # Only one thread from threads that share the MXFP8 block writes the rowwise scale + is_writer_thread = tidx % THREADS_X_PER_MX_BLOCK == 0 + for row_idx in cutlass.range_constexpr(ELEMENTS_Y_PER_THREAD): rX_row_i32 = cute.recast_tensor(rX[None, row_idx], Int32) - reciprocal = bf16_mx_scale_reciprocal(extract_row_amax(rX_row_i32), MAX_EXPONENT) + reciprocal = bf16_mx_scale_reciprocal(extract_mx_block_amax(rX_row_i32), MAX_EXPONENT) reciprocal_packed = reciprocal | (reciprocal << Uint32(16)) for pack_idx in cutlass.range_constexpr(cute.size(rO_u32)): rO_u32[pack_idx] = scale_and_pack( @@ -180,9 +158,10 @@ def extract_row_amax(values_i32: cute.Tensor) -> Uint32: rX_row_i32[2 * pack_idx + 1], reciprocal_packed, ) + # rO_u32 and rO share the same memory cute.copy(store_atom, rO, gO_row_thread[None, row_idx], cache_policy=cache_policy) - if lane % LANES_PER_MX_BLOCK == 0: - rS_row[0] = Uint8(Uint32(254) - (reciprocal >> Uint32(7))).bitcast(SCALE_DTYPE) + if is_writer_thread: + rS_row[0] = Uint8(bf16_mx_scale_bytes_pair(reciprocal_packed)).bitcast(SCALE_DTYPE) cute.copy( row_scale_atom, rS_row, @@ -190,86 +169,113 @@ def extract_row_amax(values_i32: cute.Tensor) -> Uint32: cache_policy=cache_policy, ) - # Columnwise pass: first reduce each thread's rows into one maximum per column. - rColumnAmax = cute.make_rmem_tensor((ELEMENTS_PER_THREAD,), cfg.DTYPE) - rColumnScale = cute.make_rmem_tensor((ELEMENTS_PER_THREAD,), cfg.DTYPE) - rColumnAmax_i32 = cute.recast_tensor(rColumnAmax, Int32) - rColumnScale_u32 = cute.recast_tensor(rColumnScale, Uint32) + # Columnwise pass: only produce the colwise scale. + # We first reduce each thread's rows into one maximum per column before reduce across different threads in the CTA + rAmax_col = cute.make_rmem_tensor((ELEMENTS_X_PER_THREAD,), cfg.DTYPE) + + # Gather the amax for each column that this thread sweeps rX_i32 = cute.recast_tensor(rX, Int32) - for pair_idx in cutlass.range_constexpr(COLUMN_PAIRS_PER_THREAD): - amax = rX_i32[pair_idx, 0] - for row_idx in cutlass.range_constexpr(1, ROWS_PER_THREAD): - amax = abs_max_x2_bf16(amax, rX_i32[pair_idx, row_idx]) - rColumnAmax_i32[pair_idx] = amax - - allocator = _SmemAllocator() - sColumnAmax = allocator.allocate_tensor( - cfg.DTYPE, - cute.make_layout((CTA_THREADS_Y, TILE_COLS), stride=(TILE_COLS, 1)), - byte_alignment=16, + rAmax_col_i32 = cute.recast_tensor(rAmax_col, Int32) + for bf16x2_idx in cutlass.range_constexpr(cute.size(rX_i32, mode=[0])): + amax = rX_i32[bf16x2_idx, 0] + for row_idx in cutlass.range_constexpr(1, ELEMENTS_Y_PER_THREAD): + amax = abs_max_x2_bf16(amax, rX_i32[bf16x2_idx, row_idx]) + rAmax_col_i32[bf16x2_idx] = amax + + # Threads that process a columnwise MXFP8 block don't belong to the same warp so we have to use SMEM to reduce the scale + allocator = cutlass.memory.SmemAllocator() + sAmaxs_col_scratch_layout = cute.zipped_product( + rAmax_col.layout, + thread_layout, ) - sColumnScale = allocator.allocate_tensor( - cfg.DTYPE, cute.make_layout((TILE_COLS,)), byte_alignment=16 + sAmaxs_col_scratch = allocator.allocate_tensor( + cfg.DTYPE, sAmaxs_col_scratch_layout, byte_alignment=16 ) - sColumnAmax_i32 = cute.recast_tensor(sColumnAmax, Int32) - sColumnScale_u32 = cute.recast_tensor(sColumnScale, Uint32) - # Group each lane's column pairs into two halves. Placing the first halves - # of all lanes before the second halves preserves CUDA's conflict-free - # vectorized shared-memory accesses, without calculating pointer offsets. - column_pair_layout = cute.make_layout( - ((PAIRS_PER_HALF, 2), CTA_THREADS_X), - stride=((1, COLUMN_PAIRS // 2), PAIRS_PER_HALF), + sAmaxs_col_scratch_i32 = cute.recast_tensor(sAmaxs_col_scratch, Int32) + sS_col_reduction_scratch = allocator.allocate_tensor( + cfg.DTYPE, (CTA_THREADS_Y, CTA_THREADS_X), byte_alignment=16 ) - sColumnAmax_thread = cute.composition(sColumnAmax_i32, (None, column_pair_layout))[ - warp, (None, lane) - ] - sColumnScale_thread = cute.composition(sColumnScale_u32, column_pair_layout)[None, lane] - for half_idx in cutlass.range_constexpr(2): - cute.autovec_copy( - cute.local_tile(rColumnAmax_i32, (PAIRS_PER_HALF,), (half_idx,)), - cute.local_tile(sColumnAmax_thread, (PAIRS_PER_HALF,), (half_idx,)), - ) - cute.arch.sync_threads() + sS_col_reduction_scratch_u32 = cute.recast_tensor(sS_col_reduction_scratch, Uint32) - # One thread reduces each column pair across all warps (all 32 tile rows). - if tidx < COLUMN_PAIRS: - amax = sColumnAmax_i32[0, tidx] - for warp_idx in cutlass.range_constexpr(1, CTA_THREADS_Y): - amax = abs_max_x2_bf16(amax, sColumnAmax_i32[warp_idx, tidx]) - sColumnScale_u32[tidx] = bf16_mx_scale_reciprocal_pair(amax, MAX_EXPONENT) + # Every thread copies the RMEM fragment to SMEM scratch space for reduction across warps later + cute.autovec_copy( + rAmax_col, + sAmaxs_col_scratch[None, tidx], + ) + + # Reduce thread read amaxes from SMEM then calculate & write the columnwise scale to SMEM scratch space + is_reduce_thread = tidx < CTA_THREADS_X + if is_reduce_thread: + amaxX2 = sAmaxs_col_scratch_i32[0, tidx] + for scratch_row in cutlass.range_constexpr(CTA_THREADS_Y): + amaxX2 = abs_max_x2_bf16(amaxX2, sAmaxs_col_scratch_i32[scratch_row, tidx]) + sS_col_reduction_scratch_u32[tidx] = bf16_mx_scale_reciprocal(amaxX2, MAX_EXPONENT) cute.arch.sync_threads() - # The same layout returns the reciprocals to each lane's original columns. - for half_idx in cutlass.range_constexpr(2): - cute.autovec_copy( - cute.local_tile(sColumnScale_thread, (PAIRS_PER_HALF,), (half_idx,)), - cute.local_tile(rColumnScale_u32, (PAIRS_PER_HALF,), (half_idx,)), - ) - if warp == 0: + # Let the reduction threads write the RMEM columnwise scales back to GMEM + if is_reduce_thread: rS_col_u32 = cute.recast_tensor(rS_col, Uint32) - for pack_idx in cutlass.range_constexpr(cute.size(rS_col_u32)): - rS_col_u32[pack_idx] = bf16_mx_scale_bytes_pair(rColumnScale_u32[2 * pack_idx]) | ( - bf16_mx_scale_bytes_pair(rColumnScale_u32[2 * pack_idx + 1]) << Uint32(16) + for packed_idx in cutlass.range_constexpr(cute.size(rS_col_u32)): + rS_col_u32[packed_idx] = bf16_mx_scale_bytes_pair(rS_col_u32[2 * packed_idx]) | ( + bf16_mx_scale_bytes_pair(rS_col_u32[2 * packed_idx + 1]) << Uint32(16) ) cute.copy(col_scale_atom, rS_col, gS_col_thread, cache_policy=cache_policy) - for row_idx in cutlass.range_constexpr(ROWS_PER_THREAD): - rX_row_i32 = cute.recast_tensor(rX[None, row_idx], Int32) - for pack_idx in cutlass.range_constexpr(cute.size(rO_u32)): - rO_u32[pack_idx] = scale_and_pack( - rX_row_i32[2 * pack_idx], - rColumnScale_u32[2 * pack_idx], - rX_row_i32[2 * pack_idx + 1], - rColumnScale_u32[2 * pack_idx + 1], - ) - cute.copy(store_atom, rO, gO_col_thread[None, row_idx], cache_policy=cache_policy) + # Every thread reads the reduced columnwise scales from SMEM to RMEM + rS_col = cute.make_rmem_tensor((ELEMENTS_X_PER_THREAD,), cfg.DTYPE) + cute.autovec_copy( + sS_col_reduction_scratch[tidx], + rS_col + ) + + # And then every thread uses the columnwise scales to calculate quantized outputs and writes it to GMEM + + + + + # # One thread reduces each column pair across all warps (all 32 tile rows). + # if tidx < cute.size(sColumnScales) // 2: + # sColumnAmaxs_i32 = cute.recast_tensor(sColumnAmaxs, Int32) + # sColumnScales_u32 = cute.recast_tensor(sColumnScales, Uint32) + # amax = sColumnAmaxs_i32[0, tidx] + # for warp_idx in cutlass.range_constexpr(1, CTA_THREADS_Y): + # amax = abs_max_x2_bf16(amax, sColumnAmaxs_i32[warp_idx, tidx]) + # sColumnScales_u32[tidx] = bf16_mx_scale_reciprocal_pair(amax, MAX_EXPONENT) + # cute.arch.sync_threads() + + # # The same layout returns the reciprocals to each lane's original columns. + # for half_idx in cutlass.range_constexpr(2): + # cute.autovec_copy( + # cute.local_tile(sColumnScale_thread, (cute.size(rColumnScales) // 2,), (half_idx,)), + # cute.local_tile(rColumnScales, (cute.size(rColumnScales) // 2,), (half_idx,)), + # ) + + # rColumnScales_u32 = cute.recast_tensor(rColumnScales, Uint32) + # if warp == 0: + # rS_col_u32 = cute.recast_tensor(rS_col, Uint32) + # for pack_idx in cutlass.range_constexpr(cute.size(rS_col_u32)): + # rS_col_u32[pack_idx] = bf16_mx_scale_bytes_pair(rColumnScales_u32[2 * pack_idx]) | ( + # bf16_mx_scale_bytes_pair(rColumnScales_u32[2 * pack_idx + 1]) << Uint32(16) + # ) + # cute.copy(col_scale_atom, rS_col, gS_col_thread, cache_policy=cache_policy) + + # for row_idx in cutlass.range_constexpr(ELEMENTS_Y_PER_THREAD): + # rX_row_i32 = cute.recast_tensor(rX[None, row_idx], Int32) + # for pack_idx in cutlass.range_constexpr(cute.size(rO_u32)): + # rO_u32[pack_idx] = scale_and_pack( + # rX_row_i32[2 * pack_idx], + # rColumnScales_u32[2 * pack_idx], + # rX_row_i32[2 * pack_idx + 1], + # rColumnScales_u32[2 * pack_idx + 1], + # ) + # cute.copy(store_atom, rO, gO_col_thread[None, row_idx], cache_policy=cache_policy) class MXFP8QuantizeRegisterBidimensionalKernel(MXFP8QuantizeKernelBase): """Load one CTA tile into registers and quantize it along both axes. - TILE_ROWS = _CTA_THREADS_Y * _ROWS_PER_THREAD = 32. - TILE_COLS = _CTA_THREADS_X * ELEMENTS_PER_THREAD = 256 or 512. + TILE_ROWS = _CTA_THREADS_Y * _ELEMENTS_Y_PER_THREAD = 32. + TILE_COLS = _CTA_THREADS_X * ELEMENTS_X_PER_THREAD = 256 or 512. The 32-row tile contains complete columnwise MXFP8 blocks; shared memory combines the partial column maxima from the eight warps. Requires aligned BF16 input, M divisible by 32, N divisible by 256 and plain scale layouts. @@ -277,11 +283,10 @@ class MXFP8QuantizeRegisterBidimensionalKernel(MXFP8QuantizeKernelBase): _CTA_THREADS_Y = 8 _CTA_THREADS_X = 32 - _ROWS_PER_THREAD = 4 - _ELEMENTS_PER_THREAD = 16 def __init__(self, cfg: MXFP8QuantizeConfig): self.cfg = cfg + self._ELEMENTS_Y_PER_THREAD = MXFP8_BLOCK_SCALING_SIZE // self._CTA_THREADS_Y # 4 @cute.jit def __call__( @@ -298,8 +303,8 @@ def __call__( stream: CUstream, ): M, N = mX.shape - TILE_ROWS = self._CTA_THREADS_Y * self._ROWS_PER_THREAD - WIDE_TILE_COLS = self._CTA_THREADS_X * self._ELEMENTS_PER_THREAD + TILE_ROWS = self._CTA_THREADS_Y * self._ELEMENTS_Y_PER_THREAD + WIDE_TILE_COLS = self._CTA_THREADS_X * 16 if N % WIDE_TILE_COLS == 0: if Int64(M // TILE_ROWS) * (N // WIDE_TILE_COLS) >= 4096: self._launch( @@ -310,10 +315,7 @@ def __call__( mS_col, mNoop, stream, - CTA_THREADS_Y=self._CTA_THREADS_Y, - CTA_THREADS_X=self._CTA_THREADS_X, - ROWS_PER_THREAD=self._ROWS_PER_THREAD, - ELEMENTS_PER_THREAD=self._ELEMENTS_PER_THREAD, + ELEMENTS_X_PER_THREAD=16, CLUSTER_WIDTH=4, ) else: @@ -325,10 +327,7 @@ def __call__( mS_col, mNoop, stream, - CTA_THREADS_Y=self._CTA_THREADS_Y, - CTA_THREADS_X=self._CTA_THREADS_X, - ROWS_PER_THREAD=self._ROWS_PER_THREAD, - ELEMENTS_PER_THREAD=self._ELEMENTS_PER_THREAD, + ELEMENTS_X_PER_THREAD=16, CLUSTER_WIDTH=1, ) else: @@ -340,10 +339,7 @@ def __call__( mS_col, mNoop, stream, - CTA_THREADS_Y=self._CTA_THREADS_Y, - CTA_THREADS_X=self._CTA_THREADS_X, - ROWS_PER_THREAD=self._ROWS_PER_THREAD, - ELEMENTS_PER_THREAD=self._ELEMENTS_PER_THREAD // 2, + ELEMENTS_X_PER_THREAD=8, CLUSTER_WIDTH=8, ) @@ -357,56 +353,74 @@ def _launch( mS_col: cute.Tensor, mNoop: cute.Pointer, stream: CUstream, - CTA_THREADS_Y: cutlass.Constexpr[int], - CTA_THREADS_X: cutlass.Constexpr[int], - ROWS_PER_THREAD: cutlass.Constexpr[int], - ELEMENTS_PER_THREAD: cutlass.Constexpr[int], + ELEMENTS_X_PER_THREAD: cutlass.Constexpr[int], CLUSTER_WIDTH: cutlass.Constexpr[int], ): - assert CTA_THREADS_X == 32 - assert CTA_THREADS_Y * ROWS_PER_THREAD == 32 - assert ELEMENTS_PER_THREAD in (8, 16) + assert ELEMENTS_X_PER_THREAD in (8, 16) + + # Dispatch checks these alignments before selecting this kernel. + mX = cute.make_tensor(mX.iterator.align(32), mX.layout) + mS_col = cute.make_tensor(mS_col.iterator.align(16), mS_col.layout) M, N = mX.shape - TILE_ROWS = CTA_THREADS_Y * ROWS_PER_THREAD - TILE_COLS = CTA_THREADS_X * ELEMENTS_PER_THREAD - THREADS_PER_CTA = CTA_THREADS_Y * CTA_THREADS_X + TILE_ROWS = self._CTA_THREADS_Y * self._ELEMENTS_Y_PER_THREAD + TILE_COLS = self._CTA_THREADS_X * ELEMENTS_X_PER_THREAD + THREADS_PER_CTA = self._CTA_THREADS_Y * self._CTA_THREADS_X + grid_cols = N // TILE_COLS grid_rows = M // TILE_ROWS cluster_cols = Int32(CLUSTER_WIDTH) + # Adjust cluster size so that the CTAs in the column direction can be evenly divided into clusters. while cluster_cols > 1 and grid_cols % cluster_cols != 0: cluster_cols -= 1 # Keep column CTAs on grid.x so clusters read adjacent column tiles. grid = [grid_cols, grid_rows] block = [THREADS_PER_CTA] cluster = [cluster_cols, 1, 1] - # Preserve CUDA's compile-time strides for the two common column counts. + + # Hardcode the layout using static N for some of the commonly used shapes if N == 2048: + payload_layout = cute.make_layout((M, 2048), stride=(2048, 1)) + gX = cute.make_tensor(mX.iterator, payload_layout) + gO_row = cute.make_tensor(mO_row.iterator, payload_layout) + gO_col = cute.make_tensor(mO_col.iterator, payload_layout) + gS_row = cute.make_tensor( + mS_row.iterator, + cute.make_layout((cute.size(mS_row, mode=[0]), 64), stride=(64, 1)), + ) + gS_col = cute.make_tensor( + mS_col.iterator, + cute.make_layout((cute.size(mS_col, mode=[0]), 2048), stride=(2048, 1)), + ) self.kernel( - mX, - mO_row, - mS_row, - mO_col, - mS_col, + gX, + gO_row, + gS_row, + gO_col, + gS_col, mNoop, - CTA_THREADS_Y=CTA_THREADS_Y, - CTA_THREADS_X=CTA_THREADS_X, - ROWS_PER_THREAD=ROWS_PER_THREAD, - ELEMENTS_PER_THREAD=ELEMENTS_PER_THREAD, - N_CONST=2048, + ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, ).launch(grid=grid, block=block, cluster=cluster, min_blocks_per_mp=4, stream=stream) elif N == 4096: + payload_layout = cute.make_layout((M, 4096), stride=(4096, 1)) + gX = cute.make_tensor(mX.iterator, payload_layout) + gO_row = cute.make_tensor(mO_row.iterator, payload_layout) + gO_col = cute.make_tensor(mO_col.iterator, payload_layout) + gS_row = cute.make_tensor( + mS_row.iterator, + cute.make_layout((cute.size(mS_row, mode=[0]), 128), stride=(128, 1)), + ) + gS_col = cute.make_tensor( + mS_col.iterator, + cute.make_layout((cute.size(mS_col, mode=[0]), 4096), stride=(4096, 1)), + ) self.kernel( - mX, - mO_row, - mS_row, - mO_col, - mS_col, + gX, + gO_row, + gS_row, + gO_col, + gS_col, mNoop, - CTA_THREADS_Y=CTA_THREADS_Y, - CTA_THREADS_X=CTA_THREADS_X, - ROWS_PER_THREAD=ROWS_PER_THREAD, - ELEMENTS_PER_THREAD=ELEMENTS_PER_THREAD, - N_CONST=4096, + ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, ).launch(grid=grid, block=block, cluster=cluster, min_blocks_per_mp=4, stream=stream) else: self.kernel( @@ -416,11 +430,7 @@ def _launch( mO_col, mS_col, mNoop, - CTA_THREADS_Y=CTA_THREADS_Y, - CTA_THREADS_X=CTA_THREADS_X, - ROWS_PER_THREAD=ROWS_PER_THREAD, - ELEMENTS_PER_THREAD=ELEMENTS_PER_THREAD, - N_CONST=0, + ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, ).launch(grid=grid, block=block, cluster=cluster, min_blocks_per_mp=4, stream=stream) @cute.kernel @@ -432,59 +442,58 @@ def kernel( mO_col: cute.Tensor, mS_col: cute.Tensor, mNoop: cute.Pointer, - CTA_THREADS_Y: cutlass.Constexpr[int], - CTA_THREADS_X: cutlass.Constexpr[int], - ROWS_PER_THREAD: cutlass.Constexpr[int], - ELEMENTS_PER_THREAD: cutlass.Constexpr[int], - N_CONST: cutlass.Constexpr[int], + ELEMENTS_X_PER_THREAD: cutlass.Constexpr[int], ): - """Check noop before constructing or accessing any CTA tile.""" + """Skip the CTA when the noop flag is set; otherwise run the kernel body.""" if not noop_flag_is_set(mNoop): - M, N = mX.shape - if cutlass.const_expr(N_CONST != 0): - N = N_CONST - payload_layout = cute.make_layout((M, N), stride=(N, 1)) - gX = cute.make_tensor(mX.iterator.align(32), payload_layout) - gO_row = cute.make_tensor(mO_row.iterator, payload_layout) - gO_col = cute.make_tensor(mO_col.iterator, payload_layout) - gS_row = cute.make_tensor( - mS_row.iterator, - cute.make_layout((cute.size(mS_row, mode=[0]), N // 32), stride=(N // 32, 1)), - ) - gS_col = cute.make_tensor( - mS_col.iterator.align(16), - cute.make_layout((cute.size(mS_col, mode=[0]), N), stride=(N, 1)), - ) - else: - # The dispatcher checks these alignments before choosing this kernel. - gX = cute.make_tensor(mX.iterator.align(32), mX.layout) - gO_row = mO_row - gO_col = mO_col - gS_row = mS_row - gS_col = cute.make_tensor(mS_col.iterator.align(16), mS_col.layout) - cta_col, cta_row, _ = cute.arch.block_idx() - cta_coord = (cta_row, cta_col) - TILE_ROWS = CTA_THREADS_Y * ROWS_PER_THREAD - TILE_COLS = CTA_THREADS_X * ELEMENTS_PER_THREAD - TILER = (TILE_ROWS, TILE_COLS) - ROW_SCALE_TILER = (TILE_ROWS, TILE_COLS // 32) - COL_SCALE_TILER = (TILE_ROWS // 32, TILE_COLS) - gX_tile = cute.local_tile(gX, TILER, cta_coord) - gO_row_tile = cute.local_tile(gO_row, TILER, cta_coord) - gO_col_tile = cute.local_tile(gO_col, TILER, cta_coord) - gS_row_tile = cute.local_tile(gS_row, ROW_SCALE_TILER, cta_coord) - gS_col_tile = cute.local_tile(gS_col, COL_SCALE_TILER, cta_coord) - cache_policy = create_l2_policy(True) - quantize_bidimensional_mxfp8_g2r( - gX_tile, - gO_row_tile, - gS_row_tile, - gO_col_tile, - gS_col_tile, - self.cfg, - CTA_THREADS_Y=CTA_THREADS_Y, - CTA_THREADS_X=CTA_THREADS_X, - ROWS_PER_THREAD=ROWS_PER_THREAD, - ELEMENTS_PER_THREAD=ELEMENTS_PER_THREAD, - cache_policy=cache_policy, + self._kernel_main( + mX, + mO_row, + mS_row, + mO_col, + mS_col, + ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, ) + + @cute.jit + def _kernel_main( + self, + mX: cute.Tensor, + mO_row: cute.Tensor, + mS_row: cute.Tensor, + mO_col: cute.Tensor, + mS_col: cute.Tensor, + ELEMENTS_X_PER_THREAD: cutlass.Constexpr[int], + ): + """Construct CTA tile views and run bidimensional quantization.""" + cta_col, cta_row, _ = cute.arch.block_idx() + cta_coord = (cta_row, cta_col) + + TILE_ROWS = self._CTA_THREADS_Y * self._ELEMENTS_Y_PER_THREAD + TILE_COLS = self._CTA_THREADS_X * ELEMENTS_X_PER_THREAD + + TILER = (TILE_ROWS, TILE_COLS) + ROW_SCALE_TILER = (TILE_ROWS, TILE_COLS // 32) + COL_SCALE_TILER = (TILE_ROWS // 32, TILE_COLS) + + gX_tile = cute.local_tile(mX, TILER, cta_coord) + gO_row_tile = cute.local_tile(mO_row, TILER, cta_coord) + gO_col_tile = cute.local_tile(mO_col, TILER, cta_coord) + gS_row_tile = cute.local_tile(mS_row, ROW_SCALE_TILER, cta_coord) + gS_col_tile = cute.local_tile(mS_col, COL_SCALE_TILER, cta_coord) + + cache_policy = create_l2_policy(True) + + quantize_bidimensional_mxfp8_g2r( + gX_tile, + gO_row_tile, + gS_row_tile, + gO_col_tile, + gS_col_tile, + self.cfg, + CTA_THREADS_Y=self._CTA_THREADS_Y, + CTA_THREADS_X=self._CTA_THREADS_X, + ELEMENTS_Y_PER_THREAD=self._ELEMENTS_Y_PER_THREAD, + ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, + cache_policy=cache_policy, + ) diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_1lane.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_1lane.py index bbdf9a4b510..552b6429465 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_1lane.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_1lane.py @@ -177,7 +177,7 @@ def _kernel_main( class SharedStorage: buf: cute.struct.Align[cute.struct.MemRange[Float8E8M0FNU, CTA_Y * CTA_X], 16] - storage = cutlass.utils.SmemAllocator().allocate(SharedStorage) + storage = cutlass.memory.SmemAllocator().allocate(SharedStorage) sScale = storage.buf.get_tensor(cute.make_layout((CTA_Y, CTA_X), stride=(CTA_X, 1))) # sScale is (CTA_Y, CTA_X):(CTA_X, 1), which is the same layout as tv_layout_scale # so sS_thread is really just a 1-element e8m0 buffer for this thread's scale. diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_2lane.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_2lane.py index 95882a4161f..edc6cf850f5 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_2lane.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_2lane.py @@ -109,7 +109,7 @@ def quantize_rowwise_mxfp8_g2r( cute.copy(load_atom, gX_thread, rX[None, tile_idx], cache_policy=input_policy) @cute.jit - def extract_tile_amax(values_i32: cute.Tensor) -> Float32: + def extract_mx_block_amax(values_i32: cute.Tensor) -> Float32: """Return the 32-value MXFP8 block amax shared by two adjacent threads. Balanced eight-pair tree, followed by one butterfly between lane partners. @@ -134,7 +134,7 @@ def extract_tile_amax(values_i32: cute.Tensor) -> Float32: # Reinterpret the native tensors only for the packed BF16/FP8 instructions. rX_tile_i32 = cute.recast_tensor(rX[None, tile_idx], Int32) rO_u32 = cute.recast_tensor(rO, Uint32) - amax = extract_tile_amax(rX_tile_i32) + amax = extract_mx_block_amax(rX_tile_i32) exponent = cvt_f32_to_fp8e8m0fnu(amax * cfg.MAX_NORM_RCP) # Collect 4 scale bytes and write them using one thread @@ -306,63 +306,87 @@ def kernel( NUM_TILES: cutlass.Constexpr[int], CHECK_BOUNDS: cutlass.Constexpr[bool], ): - """Check the noop flag before selecting or accessing any GMEM tile.""" + """Skip the CTA when the noop flag is set; otherwise run the kernel body.""" if not noop_flag_is_set(mNoop): - # Dispatch checks 32-byte input alignment before selecting this kernel. - gX = cute.make_tensor(mX.iterator.align(32), mX.layout) - # CUDA's x/y axes hold row/column tiles in the same order as local_tile. - cta_coord = cute.arch.block_idx()[:2] - M, N = mX.shape - TILE_ROWS = CTA_THREADS_Y - TILE_COLS = CTA_THREADS_X * self._ELEMENTS_PER_THREAD - CTA_ROWS = TILE_ROWS - CTA_COLS = TILE_COLS * NUM_TILES - - gS = mS - # Apply the swizzled scale layout if requested - if cutlass.const_expr(self.cfg.WITH_GEMM_SWIZZLED_SCALES): - mS_swizzled, _ = derive_swizzled_scale_layout(M, N, True, False, mS, None) - gS = cute.composition(mS_swizzled, (cute.make_layout(M), cute.make_layout(N // 32))) - - # CTA_TILER is the block that a CTA processes in total, which consists of NUM_TILES horizontal tiles - CTA_TILER = (CTA_ROWS, CTA_COLS) - CTA_SCALE_TILER = (CTA_ROWS, CTA_COLS // 32) - gX_cta = cute.local_tile(gX, CTA_TILER, cta_coord) - gO_cta = cute.local_tile(mO, CTA_TILER, cta_coord) - gS_cta = cute.local_tile(gS, CTA_SCALE_TILER, cta_coord) - - fraction = Float32(1.0) - if cutlass.const_expr(self._L2_CACHED_CTA_PERCENT > 0): - fraction = Float32(0.0) - # A linear index is needed only for the L2 policy split. - cta_id = cute.crd2idx( - (Int64(cta_coord[0]), Int64(cta_coord[1])), cute.arch.grid_dim()[:2] - ) - if cta_id >= first_streaming_cta: - fraction = Float32(1.0) - input_policy = create_l2_policy(False, fraction) - output_policy = create_l2_policy(True) - - # Divide CTA_TILER to multiple TILER shape TILES - TILER = (TILE_ROWS, TILE_COLS) - SCALE_TILER = (TILE_ROWS, TILE_COLS // 32) - # Mode 1's first index is 0 because CTA_ROWS == TILE_ROWS (there is only 1 tile on the row dimension). - gX_tiles = cute.local_tile(gX_cta, TILER, (0, None)) - gO_tiles = cute.local_tile(gO_cta, TILER, (0, None)) - mS_tiles = cute.local_tile(gS_cta, SCALE_TILER, (0, None)) - - quantize_rowwise_mxfp8_g2r( - gX_tiles, - gO_tiles, - mS_tiles, - self.cfg, - M, - N, + self._kernel_main( + mX, + mO, + mS, + first_streaming_cta, CTA_THREADS_Y=CTA_THREADS_Y, CTA_THREADS_X=CTA_THREADS_X, - ELEMENTS_PER_THREAD=self._ELEMENTS_PER_THREAD, NUM_TILES=NUM_TILES, CHECK_BOUNDS=CHECK_BOUNDS, - input_policy=input_policy, - output_policy=output_policy, ) + + @cute.jit + def _kernel_main( + self, + mX: cute.Tensor, + mO: cute.Tensor, + mS: cute.Tensor, + first_streaming_cta: Int64, + CTA_THREADS_Y: cutlass.Constexpr[int], + CTA_THREADS_X: cutlass.Constexpr[int], + NUM_TILES: cutlass.Constexpr[int], + CHECK_BOUNDS: cutlass.Constexpr[bool], + ): + """Construct CTA tile views and run the specialized quantization.""" + # Dispatch checks 32-byte input alignment before selecting this kernel. + gX = cute.make_tensor(mX.iterator.align(32), mX.layout) + # CUDA's x/y axes hold row/column tiles in the same order as local_tile. + cta_coord = cute.arch.block_idx()[:2] + M, N = mX.shape + TILE_ROWS = CTA_THREADS_Y + TILE_COLS = CTA_THREADS_X * self._ELEMENTS_PER_THREAD + CTA_ROWS = TILE_ROWS + CTA_COLS = TILE_COLS * NUM_TILES + + gS = mS + # Apply the swizzled scale layout if requested + if cutlass.const_expr(self.cfg.WITH_GEMM_SWIZZLED_SCALES): + mS_swizzled, _ = derive_swizzled_scale_layout(M, N, True, False, mS, None) + gS = cute.composition(mS_swizzled, (cute.make_layout(M), cute.make_layout(N // 32))) + + # CTA_TILER is the block that a CTA processes in total, which consists of NUM_TILES horizontal tiles + CTA_TILER = (CTA_ROWS, CTA_COLS) + CTA_SCALE_TILER = (CTA_ROWS, CTA_COLS // 32) + gX_cta = cute.local_tile(gX, CTA_TILER, cta_coord) + gO_cta = cute.local_tile(mO, CTA_TILER, cta_coord) + gS_cta = cute.local_tile(gS, CTA_SCALE_TILER, cta_coord) + + fraction = Float32(1.0) + if cutlass.const_expr(self._L2_CACHED_CTA_PERCENT > 0): + fraction = Float32(0.0) + # A linear index is needed only for the L2 policy split. + cta_id = cute.crd2idx( + (Int64(cta_coord[0]), Int64(cta_coord[1])), cute.arch.grid_dim()[:2] + ) + if cta_id >= first_streaming_cta: + fraction = Float32(1.0) + input_policy = create_l2_policy(False, fraction) + output_policy = create_l2_policy(True) + + # Divide CTA_TILER to multiple TILER shape TILES + TILER = (TILE_ROWS, TILE_COLS) + SCALE_TILER = (TILE_ROWS, TILE_COLS // 32) + # Mode 1's first index is 0 because CTA_ROWS == TILE_ROWS (there is only 1 tile on the row dimension). + gX_tiles = cute.local_tile(gX_cta, TILER, (0, None)) + gO_tiles = cute.local_tile(gO_cta, TILER, (0, None)) + mS_tiles = cute.local_tile(gS_cta, SCALE_TILER, (0, None)) + + quantize_rowwise_mxfp8_g2r( + gX_tiles, + gO_tiles, + mS_tiles, + self.cfg, + M, + N, + CTA_THREADS_Y=CTA_THREADS_Y, + CTA_THREADS_X=CTA_THREADS_X, + ELEMENTS_PER_THREAD=self._ELEMENTS_PER_THREAD, + NUM_TILES=NUM_TILES, + CHECK_BOUNDS=CHECK_BOUNDS, + input_policy=input_policy, + output_policy=output_policy, + ) diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2s_bidimensional.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2s_bidimensional.py index 81a94382bc9..89cfff90093 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2s_bidimensional.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2s_bidimensional.py @@ -446,7 +446,7 @@ class SharedStorage: cute.struct.MemRange[Float8E8M0FNU, self._NUM_TILES * self._TILE_COLS], 16 ] - smem = cutlass.utils.SmemAllocator() + smem = cutlass.memory.SmemAllocator() storage = smem.allocate(SharedStorage) tile_layout = cute.make_layout( diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_general.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_general.py index 26ff4c10b2d..a2b4f5e55ff 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_general.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_general.py @@ -867,7 +867,7 @@ class SharedStorage: ] sAmax: cute.struct.MemRange[Float32, self._NUM_WARPS] - smem = cutlass.utils.SmemAllocator() + smem = cutlass.memory.SmemAllocator() storage = smem.allocate(SharedStorage) # Apply the layout to the allocated shared memory buffers so the first rank is the tile (nested layout) # and the second rank is the pipeline stage @@ -1220,7 +1220,7 @@ def _dbias_only_colwise( @cute.jit def _dbias_reduction_rowwise_epilouge( - self, smem: cutlass.utils.SmemAllocator, tidx: Int32, rowwise_dbias_acc: cute.Tensor + self, smem: cutlass.memory.SmemAllocator, tidx: Int32, rowwise_dbias_acc: cute.Tensor ): # Pad the buffer to avoid bank conflicts. The logical shape is still the same. Only the stride is different. DBIAS_BUFF_WIDTH = ( diff --git a/transformer_engine/common/CuTeDSL/utils.py b/transformer_engine/common/CuTeDSL/utils.py index 80a90a85ab8..48f3a956544 100644 --- a/transformer_engine/common/CuTeDSL/utils.py +++ b/transformer_engine/common/CuTeDSL/utils.py @@ -92,6 +92,52 @@ def exp2_bf16x2_rcp(scale_e8m0) -> Uint32: return bf16_bits | (bf16_bits << Uint32(16)) +@cute.jit +def bf16_mx_scale_reciprocal(magnitude: Uint32, MAX_EXPONENT: cutlass.Constexpr[int]) -> Uint32: + """Exact BF16 reciprocal scale bits, including E8M0 zero, Inf and NaN cases.""" + rounded = (magnitude + Uint32(31)) & Uint32(0x7F80) + if rounded < Uint32(MAX_EXPONENT << 7): # pylint: disable=consider-using-max-builtin + rounded = Uint32(MAX_EXPONENT << 7) + reciprocal = Uint32((254 + MAX_EXPONENT) << 7) - rounded + if magnitude >= Uint32(0x7F80): + reciprocal = Uint32(0x0040) + if magnitude != Uint32(0x7F80): + reciprocal = Uint32(0x7FFF) + return reciprocal + + +@cute.jit +def bf16_mx_scale_reciprocal_pair(pair: Int32, MAX_EXPONENT: cutlass.Constexpr[int]) -> Uint32: + """Compute two reciprocal scales without borrowing across packed BF16 halves.""" + magnitudes = Uint32(pair) & Uint32(0x7FFF7FFF) + rounded = (magnitudes + Uint32(0x001F001F)) & Uint32(0xFF80FF80) + floor = (MAX_EXPONENT << 23) | (MAX_EXPONENT << 7) + # Packed BF16 max accepts registers, so materialize the constant inside PTX. + clamped = cute.arch.inline_ptx( + "{ .reg.b32 floor_pair; " + f"mov.b32 floor_pair, {floor}; " + "max.xorsign.abs.bf16x2 {$w0}, {$r0}, floor_pair; }", + write_only_types=[Int32], + read_only_args=[rounded.bitcast(Int32)], + ) + bias = ((254 + MAX_EXPONENT) << 23) | ((254 + MAX_EXPONENT) << 7) + reciprocal = Uint32(bias) - Uint32(clamped) + if ((magnitudes + Uint32(0x00800080)) & Uint32(0x80008000)) != Uint32(0): + lo = bf16_mx_scale_reciprocal(magnitudes & Uint32(0x7FFF), MAX_EXPONENT) + hi = bf16_mx_scale_reciprocal(magnitudes >> Uint32(16), MAX_EXPONENT) + reciprocal = lo | (hi << Uint32(16)) + return reciprocal + + +@cute.jit +def bf16_mx_scale_bytes_pair(reciprocal: Uint32) -> Uint32: + """Gather the E8M0 bytes of two packed BF16 reciprocals into the low half.""" + bytes_pair = (Uint32(0x01FE01FE) - ((reciprocal >> Uint32(7)) & Uint32(0x00FF00FF))) & Uint32( + 0x00FF00FF + ) + return (bytes_pair & Uint32(0xFF)) | ((bytes_pair >> Uint32(8)) & Uint32(0xFF00)) + + @cute.jit def pack_f32x2(lo: Float32, hi: Float32) -> Int64: """Pack two f32 scalars into a single 64-bit register (`floatx2` layout). @@ -269,6 +315,7 @@ def is_packed16(dtype) -> bool: return dtype is cutlass.BFloat16 or dtype is cutlass.Float16 +# TODO(kainingz): remove this once https://github.com/NVIDIA/cutlass/pull/3717 is merged and released @cute.jit def create_l2_policy(evict_last: cutlass.Constexpr[bool], fraction: Float32 = 1.0) -> Int64: """Create the L2 eviction policy used by register-resident casts.""" From 1995fea3cf5f32115a20a6d2880b89d466d93a0d Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Thu, 8 Oct 2026 23:39:53 +0000 Subject: [PATCH 23/58] wip Signed-off-by: Kaining Zhong --- .../mxfp8/quantize_mxfp8_g2r_bidimensional.py | 170 +++++++++--------- 1 file changed, 90 insertions(+), 80 deletions(-) diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py index 4a1692df68f..c88dab0d4c4 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py @@ -21,7 +21,6 @@ from transformer_engine.common.CuTeDSL.utils_fp8 import mul_bf16x2x2_cvt_bf16x4_to_fp8x4 from .quantize_mxfp8_common import ( MXFP8_BLOCK_SCALING_SIZE, - THREADS_PER_WARP, MXFP8QuantizeConfig, MXFP8QuantizeKernelBase, bf16_pair_magnitude, @@ -101,14 +100,13 @@ def quantize_bidimensional_mxfp8_g2r( 1, ), ) + # If CTA_THREADS_Y // THREADS_Y_PER_MX_BLOCK == 1, this will help compiler derive alignment + # Otherwise coalesce will not do anything which is still fine + tv_col_scale = cute.coalesce(tv_col_scale, target_profile=(1, 1)) gS_col_thread = cute.composition(gS_col_tile, tv_col_scale)[tidx, None] # Input fragment contains all rows it handles; output fragment contains only one row and we write it back row by row rX = cute.make_rmem_tensor((ELEMENTS_X_PER_THREAD, ELEMENTS_Y_PER_THREAD), cfg.DTYPE) - rO = cute.make_rmem_tensor((ELEMENTS_X_PER_THREAD,), cfg.FP8_DTYPE) - - # For rowwise scales, we write one row at a time - rS_row = cute.make_rmem_tensor((1,), SCALE_DTYPE) load_atom = cute.make_copy_atom( cute.nvgpu.CopyG2ROp(), @@ -131,7 +129,9 @@ def quantize_bidimensional_mxfp8_g2r( # Copy all rows a thread sweeps into the RMEM fragment at once for row_idx in cutlass.range_constexpr(ELEMENTS_Y_PER_THREAD): - cute.copy(load_atom, gX_thread[None, row_idx], rX[None, row_idx], cache_policy=cache_policy) + cute.copy( + load_atom, gX_thread[(None, row_idx),], rX[None, row_idx], cache_policy=cache_policy + ) @cute.jit def extract_mx_block_amax(values_i32: cute.Tensor) -> Uint32: @@ -143,23 +143,25 @@ def extract_mx_block_amax(values_i32: cute.Tensor) -> Uint32: amax = abs_max_x2_bf16(amax, cute.arch.shuffle_sync_bfly(amax, 1 << stage)) return bf16_pair_magnitude(amax) - # Rowwise pass: produce each row's rowwise scales and quantized outputs, and write them back to GMEM immediately. - rO_u32 = cute.recast_tensor(rO, Uint32) + rS_row = cute.make_rmem_tensor((1,), SCALE_DTYPE) + rO_row = cute.make_rmem_tensor((ELEMENTS_X_PER_THREAD,), cfg.FP8_DTYPE) + rO_row_u32 = cute.recast_tensor(rO_row, Uint32) + # Only one thread from threads that share the MXFP8 block writes the rowwise scale is_writer_thread = tidx % THREADS_X_PER_MX_BLOCK == 0 for row_idx in cutlass.range_constexpr(ELEMENTS_Y_PER_THREAD): rX_row_i32 = cute.recast_tensor(rX[None, row_idx], Int32) reciprocal = bf16_mx_scale_reciprocal(extract_mx_block_amax(rX_row_i32), MAX_EXPONENT) reciprocal_packed = reciprocal | (reciprocal << Uint32(16)) - for pack_idx in cutlass.range_constexpr(cute.size(rO_u32)): - rO_u32[pack_idx] = scale_and_pack( + for pack_idx in cutlass.range_constexpr(cute.size(rO_row_u32)): + rO_row_u32[pack_idx] = scale_and_pack( rX_row_i32[2 * pack_idx], reciprocal_packed, rX_row_i32[2 * pack_idx + 1], reciprocal_packed, ) - # rO_u32 and rO share the same memory - cute.copy(store_atom, rO, gO_row_thread[None, row_idx], cache_policy=cache_policy) + # rO_row_u32 and rO_row share the same memory + cute.copy(store_atom, rO_row, gO_row_thread[(None, row_idx),], cache_policy=cache_policy) if is_writer_thread: rS_row[0] = Uint8(bf16_mx_scale_bytes_pair(reciprocal_packed)).bitcast(SCALE_DTYPE) cute.copy( @@ -169,7 +171,6 @@ def extract_mx_block_amax(values_i32: cute.Tensor) -> Uint32: cache_policy=cache_policy, ) - # Columnwise pass: only produce the colwise scale. # We first reduce each thread's rows into one maximum per column before reduce across different threads in the CTA rAmax_col = cute.make_rmem_tensor((ELEMENTS_X_PER_THREAD,), cfg.DTYPE) @@ -184,91 +185,100 @@ def extract_mx_block_amax(values_i32: cute.Tensor) -> Uint32: # Threads that process a columnwise MXFP8 block don't belong to the same warp so we have to use SMEM to reduce the scale allocator = cutlass.memory.SmemAllocator() + + if cutlass.const_expr(ELEMENTS_X_PER_THREAD == 16): + scratch_swizzle = cute.make_swizzle(3, 1, 5) + else: + scratch_swizzle = cute.make_swizzle(2, 1, 5) + sAmaxs_col_scratch_layout = cute.zipped_product( rAmax_col.layout, - thread_layout, + # Store fragments in flat thread-ID order while keeping the same shape. + cute.make_layout(thread_layout.shape), + ) + sAmaxs_col_scratch_layout = cute.make_composed_layout( + scratch_swizzle, 0, sAmaxs_col_scratch_layout ) sAmaxs_col_scratch = allocator.allocate_tensor( cfg.DTYPE, sAmaxs_col_scratch_layout, byte_alignment=16 ) - sAmaxs_col_scratch_i32 = cute.recast_tensor(sAmaxs_col_scratch, Int32) - sS_col_reduction_scratch = allocator.allocate_tensor( - cfg.DTYPE, (CTA_THREADS_Y, CTA_THREADS_X), byte_alignment=16 - ) - sS_col_reduction_scratch_u32 = cute.recast_tensor(sS_col_reduction_scratch, Uint32) - # Every thread copies the RMEM fragment to SMEM scratch space for reduction across warps later - cute.autovec_copy( - rAmax_col, - sAmaxs_col_scratch[None, tidx], + # RMEM buffer for the reciprocal of the columnwise scale to calculate the quantized output + rS_col_rcp = cute.make_rmem_tensor((ELEMENTS_X_PER_THREAD,), cfg.DTYPE) + # SMEM buffer for colwise scale reduction + sS_col_reduction_scratch_layout = cute.zipped_product( + rS_col_rcp.layout, + cute.make_layout(CTA_THREADS_X), + ) + sS_col_reduction_scratch_layout = cute.make_composed_layout( + scratch_swizzle, 0, sS_col_reduction_scratch_layout + ) + sRcp_col_reduction_scratch = allocator.allocate_tensor( + cfg.DTYPE, sS_col_reduction_scratch_layout, byte_alignment=16 ) + sRcp_col_thread = cute.composition(sRcp_col_reduction_scratch, tv_col_scale)[tidx, None] - # Reduce thread read amaxes from SMEM then calculate & write the columnwise scale to SMEM scratch space - is_reduce_thread = tidx < CTA_THREADS_X - if is_reduce_thread: - amaxX2 = sAmaxs_col_scratch_i32[0, tidx] - for scratch_row in cutlass.range_constexpr(CTA_THREADS_Y): - amaxX2 = abs_max_x2_bf16(amaxX2, sAmaxs_col_scratch_i32[scratch_row, tidx]) - sS_col_reduction_scratch_u32[tidx] = bf16_mx_scale_reciprocal(amaxX2, MAX_EXPONENT) + # Every thread writes its partial maxima using its flat thread ID. + cute.autovec_copy(rAmax_col, sAmaxs_col_scratch[None, tidx]) cute.arch.sync_threads() - # Let the reduction threads write the RMEM columnwise scales back to GMEM - if is_reduce_thread: + sAmaxs_col_scratch_i32 = cute.recast_tensor(sAmaxs_col_scratch, Int32) + + # The layout of sAmaxs_col_scratch_i32 is a zipped product of thread layout (CTA_THREADS_Y, CTA_THREADS_X) + # and each thread's value layout from rAmax_col (ELEMENTS_X_PER_THREAD,) + # And we let each thread here reduce two full column (2 bf16 as a packed i32) of CTA_THREADS_Y elements, + # so we have CTA_THREADS_X * ELEMENT_X_PER_THREAD reduction threads + reduction_threads_num = CTA_THREADS_X * ELEMENTS_X_PER_THREAD // 2 + is_reduction_thread = tidx < reduction_threads_num + + # And apply the colwise amax reduction layout to the SMEM buffer + tv_col_reduce_layout = cute.make_layout_tv( + thr_layout=cute.make_layout((reduction_threads_num,), stride=(1,)), + val_layout=cute.make_layout((CTA_THREADS_Y,), stride=(1,)), + ) + + rS_col_rcp_u32 = cute.recast_tensor(rS_col_rcp, Uint32) + if is_reduction_thread: + rAmaxs_col_i32 = cute.composition(sAmaxs_col_scratch_i32, tv_col_reduce_layout)[tidx, None] + # Find the amax for all columns by comparing each row's partial amax + for bf16x2_idx in cutlass.range_constexpr(cute.size(rS_col_rcp_u32)): + amaxX2 = rAmaxs_col_i32[bf16x2_idx] + for scratch_row in cutlass.range_constexpr(1, CTA_THREADS_Y): + peer_thread = thread_layout((scratch_row, tidx)) + amaxX2 = abs_max_x2_bf16(amaxX2, rAmaxs_col_i32[bf16x2_idx, peer_thread]) + # Get the reciprocal directly from the amax + rS_col_rcp_u32[bf16x2_idx] = bf16_mx_scale_reciprocal_pair(amaxX2, MAX_EXPONENT) + # Write the reciprocal to the SMEM scratch buffer so that all threads can use it to quantize their columns + cute.autovec_copy(rS_col_rcp, sRcp_col_thread) + + rS_col = cute.make_rmem_tensor((ELEMENTS_X_PER_THREAD,), SCALE_DTYPE) rS_col_u32 = cute.recast_tensor(rS_col, Uint32) + # Get the colwise scale from the reciprocal for packed_idx in cutlass.range_constexpr(cute.size(rS_col_u32)): - rS_col_u32[packed_idx] = bf16_mx_scale_bytes_pair(rS_col_u32[2 * packed_idx]) | ( - bf16_mx_scale_bytes_pair(rS_col_u32[2 * packed_idx + 1]) << Uint32(16) + rS_col_u32[packed_idx] = bf16_mx_scale_bytes_pair(rS_col_rcp_u32[2 * packed_idx]) | ( + bf16_mx_scale_bytes_pair(rS_col_rcp_u32[2 * packed_idx + 1]) << Uint32(16) ) + # Use RMEM buffer to stash the columnwise scale in packed bytes and flush it to GMEM using vectorized store at once cute.copy(col_scale_atom, rS_col, gS_col_thread, cache_policy=cache_policy) + cute.arch.sync_threads() - # Every thread reads the reduced columnwise scales from SMEM to RMEM - rS_col = cute.make_rmem_tensor((ELEMENTS_X_PER_THREAD,), cfg.DTYPE) - cute.autovec_copy( - sS_col_reduction_scratch[tidx], - rS_col - ) - - # And then every thread uses the columnwise scales to calculate quantized outputs and writes it to GMEM - - + # All warps read the same reciprocals for their columns through the overlapping TV layout. + cute.autovec_copy(sRcp_col_thread, rS_col_rcp) + rO_col = cute.make_rmem_tensor((ELEMENTS_X_PER_THREAD,), cfg.FP8_DTYPE) + rO_col_u32 = cute.recast_tensor(rO_col, Uint32) - # # One thread reduces each column pair across all warps (all 32 tile rows). - # if tidx < cute.size(sColumnScales) // 2: - # sColumnAmaxs_i32 = cute.recast_tensor(sColumnAmaxs, Int32) - # sColumnScales_u32 = cute.recast_tensor(sColumnScales, Uint32) - # amax = sColumnAmaxs_i32[0, tidx] - # for warp_idx in cutlass.range_constexpr(1, CTA_THREADS_Y): - # amax = abs_max_x2_bf16(amax, sColumnAmaxs_i32[warp_idx, tidx]) - # sColumnScales_u32[tidx] = bf16_mx_scale_reciprocal_pair(amax, MAX_EXPONENT) - # cute.arch.sync_threads() - - # # The same layout returns the reciprocals to each lane's original columns. - # for half_idx in cutlass.range_constexpr(2): - # cute.autovec_copy( - # cute.local_tile(sColumnScale_thread, (cute.size(rColumnScales) // 2,), (half_idx,)), - # cute.local_tile(rColumnScales, (cute.size(rColumnScales) // 2,), (half_idx,)), - # ) - - # rColumnScales_u32 = cute.recast_tensor(rColumnScales, Uint32) - # if warp == 0: - # rS_col_u32 = cute.recast_tensor(rS_col, Uint32) - # for pack_idx in cutlass.range_constexpr(cute.size(rS_col_u32)): - # rS_col_u32[pack_idx] = bf16_mx_scale_bytes_pair(rColumnScales_u32[2 * pack_idx]) | ( - # bf16_mx_scale_bytes_pair(rColumnScales_u32[2 * pack_idx + 1]) << Uint32(16) - # ) - # cute.copy(col_scale_atom, rS_col, gS_col_thread, cache_policy=cache_policy) - - # for row_idx in cutlass.range_constexpr(ELEMENTS_Y_PER_THREAD): - # rX_row_i32 = cute.recast_tensor(rX[None, row_idx], Int32) - # for pack_idx in cutlass.range_constexpr(cute.size(rO_u32)): - # rO_u32[pack_idx] = scale_and_pack( - # rX_row_i32[2 * pack_idx], - # rColumnScales_u32[2 * pack_idx], - # rX_row_i32[2 * pack_idx + 1], - # rColumnScales_u32[2 * pack_idx + 1], - # ) - # cute.copy(store_atom, rO, gO_col_thread[None, row_idx], cache_policy=cache_policy) + # Each thread quantizes its own rows with the columnwise reciprocals and write its output to GMEM. + for row_idx in cutlass.range_constexpr(ELEMENTS_Y_PER_THREAD): + rX_row_i32 = cute.recast_tensor(rX[None, row_idx], Int32) + for packed_idx in cutlass.range_constexpr(cute.size(rO_col_u32)): + rO_col_u32[packed_idx] = scale_and_pack( + rX_row_i32[2 * packed_idx], + rS_col_rcp_u32[2 * packed_idx], + rX_row_i32[2 * packed_idx + 1], + rS_col_rcp_u32[2 * packed_idx + 1], + ) + cute.copy(store_atom, rO_col, gO_col_thread[(None, row_idx),], cache_policy=cache_policy) class MXFP8QuantizeRegisterBidimensionalKernel(MXFP8QuantizeKernelBase): From 3e7a653962a1f358eac82a44d55a2b8b5399c93c Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Fri, 9 Oct 2026 00:56:06 +0000 Subject: [PATCH 24/58] refactor Signed-off-by: Kaining Zhong --- .../mxfp8/quantize_mxfp8_g2r_bidimensional.py | 424 ++++++++++++------ 1 file changed, 275 insertions(+), 149 deletions(-) diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py index c88dab0d4c4..c056a104979 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py @@ -4,7 +4,7 @@ """Register-resident CuTeDSL bidimensional MXFP8 cast, ported from cast_bidim.cuh.""" -from typing import Optional +from typing import Callable, Optional import cutlass from cutlass import cute @@ -28,6 +28,280 @@ ) +@cute.jit +def _quantize_rowwise( + rX: cute.Tensor, + gO_row_thread: cute.Tensor, + gS_row_thread: cute.Tensor, + # Atoms + output_store_atom: cute.CopyAtom, + row_scale_store_atom: cute.CopyAtom, + # Miscellaneous + cache_policy: Int64, + scale_and_pack: Callable[..., Uint32], + # Constexprs + THREADS_X_PER_MX_BLOCK: cutlass.Constexpr[int], + ELEMENTS_Y_PER_THREAD: cutlass.Constexpr[int], + ELEMENTS_X_PER_THREAD: cutlass.Constexpr[int], + SCALE_DTYPE: cutlass.Constexpr[cutlass.DataType], + FP8_DTYPE: cutlass.Constexpr[cutlass.DataType], +): + tidx, _, _ = cute.arch.thread_idx() + MAX_EXPONENT = 8 if FP8_DTYPE is cutlass.Float8E4M3FN else 15 + + @cute.jit + def extract_mx_block_amax(values_i32: cute.Tensor) -> Uint32: + """Reduce a 32-column MXFP8 block across its two or four lane owners.""" + amax = values_i32[0] + for bf16x2_idx in cutlass.range_constexpr(1, cute.size(values_i32)): + amax = abs_max_x2_bf16(amax, values_i32[bf16x2_idx]) + for stage in cutlass.range_constexpr(THREADS_X_PER_MX_BLOCK.bit_length() - 1): + amax = abs_max_x2_bf16(amax, cute.arch.shuffle_sync_bfly(amax, 1 << stage)) + return bf16_pair_magnitude(amax) + + rS_row = cute.make_rmem_tensor((1,), SCALE_DTYPE) + rO_row = cute.make_rmem_tensor((ELEMENTS_X_PER_THREAD,), FP8_DTYPE) + rO_row_u32 = cute.recast_tensor(rO_row, Uint32) + + # Only one thread from threads that share the MXFP8 block writes the rowwise scale + is_writer_thread = tidx % THREADS_X_PER_MX_BLOCK == 0 + for row_idx in cutlass.range_constexpr(ELEMENTS_Y_PER_THREAD): + rX_row_i32 = cute.recast_tensor(rX[None, row_idx], Int32) + reciprocal = bf16_mx_scale_reciprocal(extract_mx_block_amax(rX_row_i32), MAX_EXPONENT) + reciprocal_packed = reciprocal | (reciprocal << Uint32(16)) + for pack_idx in cutlass.range_constexpr(cute.size(rO_row_u32)): + rO_row_u32[pack_idx] = scale_and_pack( + rX_row_i32[2 * pack_idx], + reciprocal_packed, + rX_row_i32[2 * pack_idx + 1], + reciprocal_packed, + ) + # rO_row_u32 and rO_row share the same memory + cute.copy(output_store_atom, rO_row, gO_row_thread[(None, row_idx),], cache_policy=cache_policy) + if is_writer_thread: + rS_row[0] = Uint8(bf16_mx_scale_bytes_pair(reciprocal_packed)).bitcast(SCALE_DTYPE) + cute.copy( + row_scale_store_atom, + rS_row, + cute.local_tile(gS_row_thread, (1,), (row_idx,)), + cache_policy=cache_policy, + ) + +@cute.jit +def _gather_partial_colwise_amaxes( + rX: cute.Tensor, + # Constexprs + CTA_THREADS_Y: cutlass.Constexpr[int], + CTA_THREADS_X: cutlass.Constexpr[int], + ELEMENTS_Y_PER_THREAD: cutlass.Constexpr[int], + ELEMENTS_X_PER_THREAD: cutlass.Constexpr[int], + DTYPE: cutlass.Constexpr[cutlass.DataType], +) -> cute.Tensor: + """ + rX is a RMEM fragment of layout (ELEMENTS_X_PER_THREAD, ELEMENTS_Y_PER_THREAD) that contains this thread's data + + It returns a SMEM buffer of layout (ELEMENTS_X_PER_THREAD, (CTA_THREADS_Y, CTA_THREADS_X)), where each thread + reduces their partial columnwise amax of shape (ELEMENTS_X_PER_THREAD,) and contributes it to the SMEM buffer + """ + tidx, _, _ = cute.arch.thread_idx() + + # We first reduce each thread's rows into one maximum per column before reduce across different threads in the CTA + rAmax_col = cute.make_rmem_tensor((ELEMENTS_X_PER_THREAD,), DTYPE) + + # Gather the amax for each column that this thread sweeps + rX_i32 = cute.recast_tensor(rX, Int32) + rAmax_col_i32 = cute.recast_tensor(rAmax_col, Int32) + for bf16x2_idx in cutlass.range_constexpr(cute.size(rX_i32, mode=[0])): + amax = rX_i32[bf16x2_idx, 0] + for row_idx in cutlass.range_constexpr(1, ELEMENTS_Y_PER_THREAD): + amax = abs_max_x2_bf16(amax, rX_i32[bf16x2_idx, row_idx]) + rAmax_col_i32[bf16x2_idx] = amax + + if cutlass.const_expr(ELEMENTS_X_PER_THREAD == 16): + scratch_swizzle = cute.make_swizzle(3, 1, 5) + else: + scratch_swizzle = cute.make_swizzle(2, 1, 5) + + # Prepare amax reduction scratch space in SMEM. + sAmaxs_col_scratch_layout = cute.make_composed_layout( + scratch_swizzle, 0, cute.zipped_product( + rAmax_col.layout, # (ELEMENTS_X_PER_THREAD,) + # Store fragments in flat thread-ID order while keeping the same shape. + cute.make_layout((CTA_THREADS_Y, CTA_THREADS_X), stride=(CTA_THREADS_X, 1)) + ) + ) + allocator = cutlass.memory.SmemAllocator() + sAmaxs_col_scratch = allocator.allocate_tensor( + DTYPE, sAmaxs_col_scratch_layout, byte_alignment=16 + ) + + # Every thread writes its partial maxima using its flat thread ID. + cute.autovec_copy(rAmax_col, sAmaxs_col_scratch[None, tidx]) + cute.arch.sync_threads() + + return sAmaxs_col_scratch + +@cute.jit +def _reduce_partial_colwise_amaxes( + rS_col_rcp: cute.Tensor, + sAmaxs_col_scratch: cute.Tensor, + # Constexprs + CTA_THREADS_Y: cutlass.Constexpr[int], + CTA_THREADS_X: cutlass.Constexpr[int], + ELEMENTS_Y_PER_THREAD: cutlass.Constexpr[int], + ELEMENTS_X_PER_THREAD: cutlass.Constexpr[int], + DTYPE: cutlass.Constexpr[cutlass.DataType], + FP8_DTYPE: cutlass.Constexpr[cutlass.DataType], +) -> cute.Tensor: + """ + sAmaxs_col_scratch: the SMEM buffer of layout (ELEMENTS_X_PER_THREAD, (CTA_THREADS_Y, CTA_THREADS_X)) + that contains each threads' own columnwise amax of its ELEMENTS_X_PER_THREAD columns + rS_col_rcp: the RMEM fragment of this thread's columnwise scale reciprocals with shape (ELEMENTS_X_PER_THREAD,) + + This function reduces the partial amax buffer of layout (ELEMENTS_X_PER_THREAD, (CTA_THREADS_Y, CTA_THREADS_X)) + along the Y axis, to a single columnwise scale SMEM buffer of layout (ELEMENTS_X_PER_THREAD * CTA_THREADS_X), + and write the columnwise scale reciprocals belonging to this thread to its RMEM fragment + """ + tidx, _, _ = cute.arch.thread_idx() + + TILE_COLS = CTA_THREADS_X * ELEMENTS_X_PER_THREAD + + # The layout of sAmaxs_col_scratch_i32 is a zipped product of thread layout (CTA_THREADS_Y, CTA_THREADS_X) + # and each thread's value layout from rAmax_col (ELEMENTS_X_PER_THREAD,) + # And we let each thread here reduce two full column (2 bf16 as a packed i32) of CTA_THREADS_Y elements, + # so we have CTA_THREADS_X * ELEMENT_X_PER_THREAD reduction threads + REDUCTION_THREADS_NUM = CTA_THREADS_X * ELEMENTS_X_PER_THREAD // 2 + is_reduction_thread = tidx < REDUCTION_THREADS_NUM + # How many bf16x2 pairs we need to reduce + REDUCTION_i32_PACKS = TILE_COLS // REDUCTION_THREADS_NUM // 2 + + _, tv_col_reduce_layout = cute.make_layout_tv( + thr_layout=cute.make_layout((REDUCTION_THREADS_NUM, 1), stride=(1, REDUCTION_THREADS_NUM)), + val_layout=cute.make_layout((CTA_THREADS_Y, REDUCTION_i32_PACKS), stride=(REDUCTION_i32_PACKS, 1)), + ) + + # SMEM buffer for colwise scale reduction where mode 0 is how many reciprocals a thread has, + # and mode 1 is how many threads we have in X axis, and the layout size is just TILE_COLS + if cutlass.const_expr(ELEMENTS_X_PER_THREAD == 16): + scratch_swizzle = cute.make_swizzle(3, 1, 5) + else: + scratch_swizzle = cute.make_swizzle(2, 1, 5) + sS_col_layout = cute.make_composed_layout( + scratch_swizzle, 0, cute.zipped_product( + cute.make_layout((ELEMENTS_X_PER_THREAD,)), + cute.make_layout((CTA_THREADS_X,)), + ) + ) + allocator = cutlass.memory.SmemAllocator() + sS_col = allocator.allocate_tensor( + DTYPE, sS_col_layout, byte_alignment=16 + ) + + # View sS_col using the TV layout from reduction threads. + # REDUCTION_THREADS_NUM * ELEMENTS_X_PER_THREAD = TILE_COLS = CTA_THREADS_X * ELEMENTS_X_PER_THREAD + sS_col_i32 = cute.recast_tensor(sS_col, Int32) + sS_col_i32_reduction_view = cute.composition(sS_col_i32, cute.make_tv_layout( + thr_layout=cute.make_layout((1, REDUCTION_THREADS_NUM), stride=(1, 1)), + val_layout=cute.make_layout((1, ELEMENTS_X_PER_THREAD,), stride=(1, 1)), + )) + + # Each reduction process `REDUCTION_i32_PACKS` bf16x2 packs by visiting all rows (CTA_THREADS_Y) + # to extract the amax of that column + sAmaxs_col_scratch_i32 = cute.recast_tensor(sAmaxs_col_scratch, Int32) + MAX_EXPONENT = 8 if FP8_DTYPE is cutlass.Float8E4M3FN else 15 + if is_reduction_thread: + sRcp_col_i32_reduction_frag = cute.composition(sS_col_i32, sS_col_i32_reduction_view)[tidx, None] + sAmaxs_col_frag_i32 = cute.composition(sAmaxs_col_scratch_i32, tv_col_reduce_layout)[tidx, None] + for pack_idx in cutlass.range_constexpr(REDUCTION_i32_PACKS): + sAmaxs_col_frag_i32_rows = sAmaxs_col_frag_i32[None, pack_idx] + amaxX2 = sAmaxs_col_frag_i32_rows[0] + for scratch_row_idx in cutlass.range_constexpr(1, CTA_THREADS_Y): + amaxX2 = abs_max_x2_bf16(amaxX2, sAmaxs_col_frag_i32_rows[scratch_row_idx]) + sRcp_col_i32_reduction_frag[pack_idx] = bf16_mx_scale_reciprocal_pair(amaxX2, MAX_EXPONENT) + cute.arch.sync_threads() + + cute.autovec_copy(sS_col[None, tidx], rS_col_rcp) + + +@cute.jit +def _write_colwise_scale( + rS_col_rcp: cute.Tensor, + gS_col_thread: cute.Tensor, + # Atoms + col_scale_atom: cute.CopyAtom, + # Miscellaneous + cache_policy: Int64, + # Constexprs + CTA_THREADS_Y: cutlass.Constexpr[int], + CTA_THREADS_X: cutlass.Constexpr[int], + ELEMENTS_Y_PER_THREAD: cutlass.Constexpr[int], + ELEMENTS_X_PER_THREAD: cutlass.Constexpr[int], + DTYPE: cutlass.Constexpr[cutlass.DataType], + FP8_DTYPE: cutlass.Constexpr[cutlass.DataType], + SCALE_DTYPE: cutlass.Constexpr[cutlass.DataType], +): + """ + rS_col_rcp: RMEM fragment of this thread's columnwise scale reciprocals with shape (ELEMENTS_X_PER_THREAD,) + """ + tidx, _, _ = cute.arch.thread_idx() + + rS_col_rcp = cute.make_rmem_tensor((ELEMENTS_X_PER_THREAD,), DTYPE) + rS_col_rcp_u32 = cute.recast_tensor(rS_col_rcp, Uint32) + + # Only threads from CTA's the first row writes the scale fragments, retaining vectorized GMEM stores. + if tidx < CTA_THREADS_X: + rS_col = cute.make_rmem_tensor((ELEMENTS_X_PER_THREAD,), SCALE_DTYPE) + rS_col_u32 = cute.recast_tensor(rS_col, Uint32) + # Get the colwise scale from the reciprocal + for packed_idx in cutlass.range_constexpr(cute.size(rS_col_u32)): + rS_col_u32[packed_idx] = bf16_mx_scale_bytes_pair(rS_col_rcp_u32[2 * packed_idx]) | ( + bf16_mx_scale_bytes_pair(rS_col_rcp_u32[2 * packed_idx + 1]) << Uint32(16) + ) + # Use RMEM buffer to stash the columnwise scale in packed bytes and flush it to GMEM using vectorized store at once + cute.copy(col_scale_atom, rS_col, gS_col_thread, cache_policy=cache_policy) + + +@cute.jit +def _quantize_colwise( + rX: cute.Tensor, + rS_col_rcp: cute.Tensor, + gO_col_thread: cute.Tensor, + # Atoms + store_atom: cute.CopyAtom, + # Miscellaneous + cache_policy: Int64, + scale_and_pack: Callable[..., Uint32], + # Constexprs + ELEMENTS_Y_PER_THREAD: cutlass.Constexpr[int], + ELEMENTS_X_PER_THREAD: cutlass.Constexpr[int], + FP8_DTYPE: cutlass.Constexpr[cutlass.DataType], +): + """ + rX: RMEM fragment of this thread's rows with shape (ELEMENTS_X_PER_THREAD, ELEMENTS_Y_PER_THREAD) + rS_col_rcp: RMEM fragment of this thread's columnwise scale reciprocals with shape (ELEMENTS_X_PER_THREAD,) + gO_col_thread: GMEM fragment of this thread's columnwise quantized output with shape + (ELEMENTS_X_PER_THREAD, ELEMENTS_Y_PER_THREAD) + + This function quantizes its fragments using the input and columnwise scales, and write to its output fragment + """ + rS_col_rcp_u32 = cute.recast_tensor(rS_col_rcp, Uint32) + + rO_col = cute.make_rmem_tensor((ELEMENTS_X_PER_THREAD,), FP8_DTYPE) + rO_col_u32 = cute.recast_tensor(rO_col, Uint32) + + # Each thread quantizes its own rows with the columnwise reciprocals and write its output to GMEM. + for row_idx in cutlass.range_constexpr(ELEMENTS_Y_PER_THREAD): + rX_row_i32 = cute.recast_tensor(rX[None, row_idx], Int32) + for packed_idx in cutlass.range_constexpr(ELEMENTS_X_PER_THREAD // 2): # packed in 2 + rO_col_u32[packed_idx] = scale_and_pack( + rX_row_i32[2 * packed_idx], + rS_col_rcp_u32[2 * packed_idx], + rX_row_i32[2 * packed_idx + 1], + rS_col_rcp_u32[2 * packed_idx + 1], + ) + cute.copy(store_atom, rO_col, gO_col_thread[(None, row_idx),], cache_policy=cache_policy) + + @cute.jit def quantize_bidimensional_mxfp8_g2r( gX_tile: cute.Tensor, @@ -51,13 +325,11 @@ def quantize_bidimensional_mxfp8_g2r( """ # Shape of the tile TILE_ROWS = CTA_THREADS_Y * ELEMENTS_Y_PER_THREAD - TILE_COLS = CTA_THREADS_X * ELEMENTS_X_PER_THREAD # How many threads cooperate to process a single MXFP8 block (32 elements) for both quantized directions THREADS_X_PER_MX_BLOCK = MXFP8_BLOCK_SCALING_SIZE // ELEMENTS_X_PER_THREAD THREADS_Y_PER_MX_BLOCK = MXFP8_BLOCK_SCALING_SIZE // ELEMENTS_Y_PER_THREAD - MAX_EXPONENT = 8 if cfg.FP8_DTYPE is cutlass.Float8E4M3FN else 15 SCALE_DTYPE = cutlass.Float8E8M0FNU tidx, _, _ = cute.arch.thread_idx() @@ -133,152 +405,6 @@ def quantize_bidimensional_mxfp8_g2r( load_atom, gX_thread[(None, row_idx),], rX[None, row_idx], cache_policy=cache_policy ) - @cute.jit - def extract_mx_block_amax(values_i32: cute.Tensor) -> Uint32: - """Reduce a 32-column MXFP8 block across its two or four lane owners.""" - amax = values_i32[0] - for bf16x2_idx in cutlass.range_constexpr(1, cute.size(values_i32)): - amax = abs_max_x2_bf16(amax, values_i32[bf16x2_idx]) - for stage in cutlass.range_constexpr(THREADS_X_PER_MX_BLOCK.bit_length() - 1): - amax = abs_max_x2_bf16(amax, cute.arch.shuffle_sync_bfly(amax, 1 << stage)) - return bf16_pair_magnitude(amax) - - rS_row = cute.make_rmem_tensor((1,), SCALE_DTYPE) - rO_row = cute.make_rmem_tensor((ELEMENTS_X_PER_THREAD,), cfg.FP8_DTYPE) - rO_row_u32 = cute.recast_tensor(rO_row, Uint32) - - # Only one thread from threads that share the MXFP8 block writes the rowwise scale - is_writer_thread = tidx % THREADS_X_PER_MX_BLOCK == 0 - for row_idx in cutlass.range_constexpr(ELEMENTS_Y_PER_THREAD): - rX_row_i32 = cute.recast_tensor(rX[None, row_idx], Int32) - reciprocal = bf16_mx_scale_reciprocal(extract_mx_block_amax(rX_row_i32), MAX_EXPONENT) - reciprocal_packed = reciprocal | (reciprocal << Uint32(16)) - for pack_idx in cutlass.range_constexpr(cute.size(rO_row_u32)): - rO_row_u32[pack_idx] = scale_and_pack( - rX_row_i32[2 * pack_idx], - reciprocal_packed, - rX_row_i32[2 * pack_idx + 1], - reciprocal_packed, - ) - # rO_row_u32 and rO_row share the same memory - cute.copy(store_atom, rO_row, gO_row_thread[(None, row_idx),], cache_policy=cache_policy) - if is_writer_thread: - rS_row[0] = Uint8(bf16_mx_scale_bytes_pair(reciprocal_packed)).bitcast(SCALE_DTYPE) - cute.copy( - row_scale_atom, - rS_row, - cute.local_tile(gS_row_thread, (1,), (row_idx,)), - cache_policy=cache_policy, - ) - - # We first reduce each thread's rows into one maximum per column before reduce across different threads in the CTA - rAmax_col = cute.make_rmem_tensor((ELEMENTS_X_PER_THREAD,), cfg.DTYPE) - - # Gather the amax for each column that this thread sweeps - rX_i32 = cute.recast_tensor(rX, Int32) - rAmax_col_i32 = cute.recast_tensor(rAmax_col, Int32) - for bf16x2_idx in cutlass.range_constexpr(cute.size(rX_i32, mode=[0])): - amax = rX_i32[bf16x2_idx, 0] - for row_idx in cutlass.range_constexpr(1, ELEMENTS_Y_PER_THREAD): - amax = abs_max_x2_bf16(amax, rX_i32[bf16x2_idx, row_idx]) - rAmax_col_i32[bf16x2_idx] = amax - - # Threads that process a columnwise MXFP8 block don't belong to the same warp so we have to use SMEM to reduce the scale - allocator = cutlass.memory.SmemAllocator() - - if cutlass.const_expr(ELEMENTS_X_PER_THREAD == 16): - scratch_swizzle = cute.make_swizzle(3, 1, 5) - else: - scratch_swizzle = cute.make_swizzle(2, 1, 5) - - sAmaxs_col_scratch_layout = cute.zipped_product( - rAmax_col.layout, - # Store fragments in flat thread-ID order while keeping the same shape. - cute.make_layout(thread_layout.shape), - ) - sAmaxs_col_scratch_layout = cute.make_composed_layout( - scratch_swizzle, 0, sAmaxs_col_scratch_layout - ) - sAmaxs_col_scratch = allocator.allocate_tensor( - cfg.DTYPE, sAmaxs_col_scratch_layout, byte_alignment=16 - ) - - # RMEM buffer for the reciprocal of the columnwise scale to calculate the quantized output - rS_col_rcp = cute.make_rmem_tensor((ELEMENTS_X_PER_THREAD,), cfg.DTYPE) - # SMEM buffer for colwise scale reduction - sS_col_reduction_scratch_layout = cute.zipped_product( - rS_col_rcp.layout, - cute.make_layout(CTA_THREADS_X), - ) - sS_col_reduction_scratch_layout = cute.make_composed_layout( - scratch_swizzle, 0, sS_col_reduction_scratch_layout - ) - sRcp_col_reduction_scratch = allocator.allocate_tensor( - cfg.DTYPE, sS_col_reduction_scratch_layout, byte_alignment=16 - ) - sRcp_col_thread = cute.composition(sRcp_col_reduction_scratch, tv_col_scale)[tidx, None] - - # Every thread writes its partial maxima using its flat thread ID. - cute.autovec_copy(rAmax_col, sAmaxs_col_scratch[None, tidx]) - cute.arch.sync_threads() - - sAmaxs_col_scratch_i32 = cute.recast_tensor(sAmaxs_col_scratch, Int32) - - # The layout of sAmaxs_col_scratch_i32 is a zipped product of thread layout (CTA_THREADS_Y, CTA_THREADS_X) - # and each thread's value layout from rAmax_col (ELEMENTS_X_PER_THREAD,) - # And we let each thread here reduce two full column (2 bf16 as a packed i32) of CTA_THREADS_Y elements, - # so we have CTA_THREADS_X * ELEMENT_X_PER_THREAD reduction threads - reduction_threads_num = CTA_THREADS_X * ELEMENTS_X_PER_THREAD // 2 - is_reduction_thread = tidx < reduction_threads_num - - # And apply the colwise amax reduction layout to the SMEM buffer - tv_col_reduce_layout = cute.make_layout_tv( - thr_layout=cute.make_layout((reduction_threads_num,), stride=(1,)), - val_layout=cute.make_layout((CTA_THREADS_Y,), stride=(1,)), - ) - - rS_col_rcp_u32 = cute.recast_tensor(rS_col_rcp, Uint32) - if is_reduction_thread: - rAmaxs_col_i32 = cute.composition(sAmaxs_col_scratch_i32, tv_col_reduce_layout)[tidx, None] - # Find the amax for all columns by comparing each row's partial amax - for bf16x2_idx in cutlass.range_constexpr(cute.size(rS_col_rcp_u32)): - amaxX2 = rAmaxs_col_i32[bf16x2_idx] - for scratch_row in cutlass.range_constexpr(1, CTA_THREADS_Y): - peer_thread = thread_layout((scratch_row, tidx)) - amaxX2 = abs_max_x2_bf16(amaxX2, rAmaxs_col_i32[bf16x2_idx, peer_thread]) - # Get the reciprocal directly from the amax - rS_col_rcp_u32[bf16x2_idx] = bf16_mx_scale_reciprocal_pair(amaxX2, MAX_EXPONENT) - # Write the reciprocal to the SMEM scratch buffer so that all threads can use it to quantize their columns - cute.autovec_copy(rS_col_rcp, sRcp_col_thread) - - rS_col = cute.make_rmem_tensor((ELEMENTS_X_PER_THREAD,), SCALE_DTYPE) - rS_col_u32 = cute.recast_tensor(rS_col, Uint32) - # Get the colwise scale from the reciprocal - for packed_idx in cutlass.range_constexpr(cute.size(rS_col_u32)): - rS_col_u32[packed_idx] = bf16_mx_scale_bytes_pair(rS_col_rcp_u32[2 * packed_idx]) | ( - bf16_mx_scale_bytes_pair(rS_col_rcp_u32[2 * packed_idx + 1]) << Uint32(16) - ) - # Use RMEM buffer to stash the columnwise scale in packed bytes and flush it to GMEM using vectorized store at once - cute.copy(col_scale_atom, rS_col, gS_col_thread, cache_policy=cache_policy) - cute.arch.sync_threads() - - # All warps read the same reciprocals for their columns through the overlapping TV layout. - cute.autovec_copy(sRcp_col_thread, rS_col_rcp) - - rO_col = cute.make_rmem_tensor((ELEMENTS_X_PER_THREAD,), cfg.FP8_DTYPE) - rO_col_u32 = cute.recast_tensor(rO_col, Uint32) - - # Each thread quantizes its own rows with the columnwise reciprocals and write its output to GMEM. - for row_idx in cutlass.range_constexpr(ELEMENTS_Y_PER_THREAD): - rX_row_i32 = cute.recast_tensor(rX[None, row_idx], Int32) - for packed_idx in cutlass.range_constexpr(cute.size(rO_col_u32)): - rO_col_u32[packed_idx] = scale_and_pack( - rX_row_i32[2 * packed_idx], - rS_col_rcp_u32[2 * packed_idx], - rX_row_i32[2 * packed_idx + 1], - rS_col_rcp_u32[2 * packed_idx + 1], - ) - cute.copy(store_atom, rO_col, gO_col_thread[(None, row_idx),], cache_policy=cache_policy) class MXFP8QuantizeRegisterBidimensionalKernel(MXFP8QuantizeKernelBase): From 79c5d5325104ccc209a997a0207a70edb4e45c38 Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Fri, 9 Oct 2026 01:28:20 +0000 Subject: [PATCH 25/58] refactor Signed-off-by: Kaining Zhong --- .../mxfp8/quantize_mxfp8_g2r_bidimensional.py | 197 +++++++++++++----- 1 file changed, 142 insertions(+), 55 deletions(-) diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py index c056a104979..9db22913ddd 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py @@ -4,7 +4,7 @@ """Register-resident CuTeDSL bidimensional MXFP8 cast, ported from cast_bidim.cuh.""" -from typing import Callable, Optional +from typing import Callable, Optional, Type import cutlass from cutlass import cute @@ -38,13 +38,13 @@ def _quantize_rowwise( row_scale_store_atom: cute.CopyAtom, # Miscellaneous cache_policy: Int64, - scale_and_pack: Callable[..., Uint32], + scale_and_pack: cutlass.Constexpr[Callable[..., Uint32]], # Constexprs THREADS_X_PER_MX_BLOCK: cutlass.Constexpr[int], ELEMENTS_Y_PER_THREAD: cutlass.Constexpr[int], ELEMENTS_X_PER_THREAD: cutlass.Constexpr[int], - SCALE_DTYPE: cutlass.Constexpr[cutlass.DataType], - FP8_DTYPE: cutlass.Constexpr[cutlass.DataType], + SCALE_DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], + FP8_DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], ): tidx, _, _ = cute.arch.thread_idx() MAX_EXPONENT = 8 if FP8_DTYPE is cutlass.Float8E4M3FN else 15 @@ -77,7 +77,9 @@ def extract_mx_block_amax(values_i32: cute.Tensor) -> Uint32: reciprocal_packed, ) # rO_row_u32 and rO_row share the same memory - cute.copy(output_store_atom, rO_row, gO_row_thread[(None, row_idx),], cache_policy=cache_policy) + cute.copy( + output_store_atom, rO_row, gO_row_thread[(None, row_idx),], cache_policy=cache_policy + ) if is_writer_thread: rS_row[0] = Uint8(bf16_mx_scale_bytes_pair(reciprocal_packed)).bitcast(SCALE_DTYPE) cute.copy( @@ -87,6 +89,7 @@ def extract_mx_block_amax(values_i32: cute.Tensor) -> Uint32: cache_policy=cache_policy, ) + @cute.jit def _gather_partial_colwise_amaxes( rX: cute.Tensor, @@ -95,7 +98,7 @@ def _gather_partial_colwise_amaxes( CTA_THREADS_X: cutlass.Constexpr[int], ELEMENTS_Y_PER_THREAD: cutlass.Constexpr[int], ELEMENTS_X_PER_THREAD: cutlass.Constexpr[int], - DTYPE: cutlass.Constexpr[cutlass.DataType], + DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], ) -> cute.Tensor: """ rX is a RMEM fragment of layout (ELEMENTS_X_PER_THREAD, ELEMENTS_Y_PER_THREAD) that contains this thread's data @@ -124,11 +127,12 @@ def _gather_partial_colwise_amaxes( # Prepare amax reduction scratch space in SMEM. sAmaxs_col_scratch_layout = cute.make_composed_layout( - scratch_swizzle, 0, cute.zipped_product( - rAmax_col.layout, # (ELEMENTS_X_PER_THREAD,) - # Store fragments in flat thread-ID order while keeping the same shape. - cute.make_layout((CTA_THREADS_Y, CTA_THREADS_X), stride=(CTA_THREADS_X, 1)) - ) + scratch_swizzle, + 0, + cute.zipped_product( + rAmax_col.layout, # Each thread contributes (ELEMENTS_X_PER_THREAD,) + cute.make_layout((CTA_THREADS_Y, CTA_THREADS_X)), # Total threads in the CTA + ), ) allocator = cutlass.memory.SmemAllocator() sAmaxs_col_scratch = allocator.allocate_tensor( @@ -141,86 +145,103 @@ def _gather_partial_colwise_amaxes( return sAmaxs_col_scratch + @cute.jit def _reduce_partial_colwise_amaxes( - rS_col_rcp: cute.Tensor, sAmaxs_col_scratch: cute.Tensor, # Constexprs CTA_THREADS_Y: cutlass.Constexpr[int], CTA_THREADS_X: cutlass.Constexpr[int], ELEMENTS_Y_PER_THREAD: cutlass.Constexpr[int], ELEMENTS_X_PER_THREAD: cutlass.Constexpr[int], - DTYPE: cutlass.Constexpr[cutlass.DataType], - FP8_DTYPE: cutlass.Constexpr[cutlass.DataType], -) -> cute.Tensor: + DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], + FP8_DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], +) -> None: """ - sAmaxs_col_scratch: the SMEM buffer of layout (ELEMENTS_X_PER_THREAD, (CTA_THREADS_Y, CTA_THREADS_X)) + sAmaxs_col_scratch: the SMEM buffer of layout (ELEMENTS_X_PER_THREAD, (CTA_THREADS_Y, CTA_THREADS_X)) that contains each threads' own columnwise amax of its ELEMENTS_X_PER_THREAD columns rS_col_rcp: the RMEM fragment of this thread's columnwise scale reciprocals with shape (ELEMENTS_X_PER_THREAD,) - This function reduces the partial amax buffer of layout (ELEMENTS_X_PER_THREAD, (CTA_THREADS_Y, CTA_THREADS_X)) + This function reduces the partial amax buffer of layout (ELEMENTS_X_PER_THREAD, (CTA_THREADS_Y, CTA_THREADS_X)) along the Y axis, to a single columnwise scale SMEM buffer of layout (ELEMENTS_X_PER_THREAD * CTA_THREADS_X), - and write the columnwise scale reciprocals belonging to this thread to its RMEM fragment + which is the return value """ tidx, _, _ = cute.arch.thread_idx() TILE_COLS = CTA_THREADS_X * ELEMENTS_X_PER_THREAD - # The layout of sAmaxs_col_scratch_i32 is a zipped product of thread layout (CTA_THREADS_Y, CTA_THREADS_X) - # and each thread's value layout from rAmax_col (ELEMENTS_X_PER_THREAD,) - # And we let each thread here reduce two full column (2 bf16 as a packed i32) of CTA_THREADS_Y elements, - # so we have CTA_THREADS_X * ELEMENT_X_PER_THREAD reduction threads + # Each reducer owns one column pair and visits its partial maximum from + # every warp. Column pairs precede warp rows in the scratch's logical order. REDUCTION_THREADS_NUM = CTA_THREADS_X * ELEMENTS_X_PER_THREAD // 2 is_reduction_thread = tidx < REDUCTION_THREADS_NUM # How many bf16x2 pairs we need to reduce REDUCTION_i32_PACKS = TILE_COLS // REDUCTION_THREADS_NUM // 2 - _, tv_col_reduce_layout = cute.make_layout_tv( - thr_layout=cute.make_layout((REDUCTION_THREADS_NUM, 1), stride=(1, REDUCTION_THREADS_NUM)), - val_layout=cute.make_layout((CTA_THREADS_Y, REDUCTION_i32_PACKS), stride=(REDUCTION_i32_PACKS, 1)), - ) - - # SMEM buffer for colwise scale reduction where mode 0 is how many reciprocals a thread has, - # and mode 1 is how many threads we have in X axis, and the layout size is just TILE_COLS + # Allocate colwise scale SMEM buffer with swizzling if cutlass.const_expr(ELEMENTS_X_PER_THREAD == 16): scratch_swizzle = cute.make_swizzle(3, 1, 5) else: scratch_swizzle = cute.make_swizzle(2, 1, 5) + # SMEM buffer for colwise scale reduction where mode 0 is how many reciprocals a thread has, + # and mode 1 is how many threads we have in X axis, and the layout size is just TILE_COLS sS_col_layout = cute.make_composed_layout( - scratch_swizzle, 0, cute.zipped_product( + scratch_swizzle, + 0, + cute.zipped_product( cute.make_layout((ELEMENTS_X_PER_THREAD,)), cute.make_layout((CTA_THREADS_X,)), - ) + ), ) allocator = cutlass.memory.SmemAllocator() - sS_col = allocator.allocate_tensor( - DTYPE, sS_col_layout, byte_alignment=16 - ) - - # View sS_col using the TV layout from reduction threads. - # REDUCTION_THREADS_NUM * ELEMENTS_X_PER_THREAD = TILE_COLS = CTA_THREADS_X * ELEMENTS_X_PER_THREAD + sS_col = allocator.allocate_tensor(DTYPE, sS_col_layout, byte_alignment=16) sS_col_i32 = cute.recast_tensor(sS_col, Int32) - sS_col_i32_reduction_view = cute.composition(sS_col_i32, cute.make_tv_layout( - thr_layout=cute.make_layout((1, REDUCTION_THREADS_NUM), stride=(1, 1)), - val_layout=cute.make_layout((1, ELEMENTS_X_PER_THREAD,), stride=(1, 1)), - )) + + # Colwise amaxes layout: each reduction thread owns CTA_THREADS_Y rows, with REDUCTION_i32_PACKS column packs + _, tv_col_amaxes_layout = cute.make_layout_tv( + thr_layout=cute.make_layout((REDUCTION_THREADS_NUM, 1), stride=(1, REDUCTION_THREADS_NUM)), + val_layout=cute.make_layout( + (REDUCTION_i32_PACKS, CTA_THREADS_Y), stride=(1, REDUCTION_i32_PACKS) + ), + ) + # Colwise rcp layout: each reduction thread owns only 1 row, with REDUCTION_i32_PACKS column packs + _, tv_col_rcp_layout = cute.make_layout_tv( + thr_layout=cute.make_layout((REDUCTION_THREADS_NUM, 1), stride=(1, REDUCTION_THREADS_NUM)), + val_layout=cute.make_layout((REDUCTION_i32_PACKS, 1), stride=(1, REDUCTION_i32_PACKS)), + ) + + sS_col_i32 = cute.composition(sS_col_i32, tv_col_rcp_layout) # Each reduction process `REDUCTION_i32_PACKS` bf16x2 packs by visiting all rows (CTA_THREADS_Y) # to extract the amax of that column sAmaxs_col_scratch_i32 = cute.recast_tensor(sAmaxs_col_scratch, Int32) MAX_EXPONENT = 8 if FP8_DTYPE is cutlass.Float8E4M3FN else 15 if is_reduction_thread: - sRcp_col_i32_reduction_frag = cute.composition(sS_col_i32, sS_col_i32_reduction_view)[tidx, None] - sAmaxs_col_frag_i32 = cute.composition(sAmaxs_col_scratch_i32, tv_col_reduce_layout)[tidx, None] + sAmaxs_col_frag_i32 = cute.composition(sAmaxs_col_scratch_i32, tv_col_amaxes_layout)[ + tidx, None + ] + # Expose the thread's value mode as (partial rows, column pairs). + sAmaxs_col_frag_i32 = cute.composition( + sAmaxs_col_frag_i32, + cute.make_layout((CTA_THREADS_Y, REDUCTION_i32_PACKS), stride=(REDUCTION_i32_PACKS, 1)), + ) + + # Iterative over all packed columns for pack_idx in cutlass.range_constexpr(REDUCTION_i32_PACKS): + # Select the packed column sAmaxs_col_frag_i32_rows = sAmaxs_col_frag_i32[None, pack_idx] + + # Reduce across all rows' partial amaxes to obtain the columnwise amax for this column pack amaxX2 = sAmaxs_col_frag_i32_rows[0] for scratch_row_idx in cutlass.range_constexpr(1, CTA_THREADS_Y): amaxX2 = abs_max_x2_bf16(amaxX2, sAmaxs_col_frag_i32_rows[scratch_row_idx]) - sRcp_col_i32_reduction_frag[pack_idx] = bf16_mx_scale_reciprocal_pair(amaxX2, MAX_EXPONENT) + + # Obtain the columnwise scale reciprocal for this column pack + sS_col_i32[tidx, pack_idx] = bf16_mx_scale_reciprocal_pair( + amaxX2, MAX_EXPONENT + ) cute.arch.sync_threads() - cute.autovec_copy(sS_col[None, tidx], rS_col_rcp) + return sS_col @cute.jit @@ -236,19 +257,18 @@ def _write_colwise_scale( CTA_THREADS_X: cutlass.Constexpr[int], ELEMENTS_Y_PER_THREAD: cutlass.Constexpr[int], ELEMENTS_X_PER_THREAD: cutlass.Constexpr[int], - DTYPE: cutlass.Constexpr[cutlass.DataType], - FP8_DTYPE: cutlass.Constexpr[cutlass.DataType], - SCALE_DTYPE: cutlass.Constexpr[cutlass.DataType], + DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], + FP8_DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], + SCALE_DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], ): """ rS_col_rcp: RMEM fragment of this thread's columnwise scale reciprocals with shape (ELEMENTS_X_PER_THREAD,) """ tidx, _, _ = cute.arch.thread_idx() - rS_col_rcp = cute.make_rmem_tensor((ELEMENTS_X_PER_THREAD,), DTYPE) rS_col_rcp_u32 = cute.recast_tensor(rS_col_rcp, Uint32) - # Only threads from CTA's the first row writes the scale fragments, retaining vectorized GMEM stores. + # Only threads from CTA's the first row writes the scale fragments if tidx < CTA_THREADS_X: rS_col = cute.make_rmem_tensor((ELEMENTS_X_PER_THREAD,), SCALE_DTYPE) rS_col_u32 = cute.recast_tensor(rS_col, Uint32) @@ -257,7 +277,8 @@ def _write_colwise_scale( rS_col_u32[packed_idx] = bf16_mx_scale_bytes_pair(rS_col_rcp_u32[2 * packed_idx]) | ( bf16_mx_scale_bytes_pair(rS_col_rcp_u32[2 * packed_idx + 1]) << Uint32(16) ) - # Use RMEM buffer to stash the columnwise scale in packed bytes and flush it to GMEM using vectorized store at once + # Use RMEM buffer to stash the columnwise scale in packed bytes and flush it to GMEM + # using vectorized store at once cute.copy(col_scale_atom, rS_col, gS_col_thread, cache_policy=cache_policy) @@ -270,18 +291,18 @@ def _quantize_colwise( store_atom: cute.CopyAtom, # Miscellaneous cache_policy: Int64, - scale_and_pack: Callable[..., Uint32], + scale_and_pack: cutlass.Constexpr[Callable[..., Uint32]], # Constexprs ELEMENTS_Y_PER_THREAD: cutlass.Constexpr[int], ELEMENTS_X_PER_THREAD: cutlass.Constexpr[int], - FP8_DTYPE: cutlass.Constexpr[cutlass.DataType], + FP8_DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], ): """ rX: RMEM fragment of this thread's rows with shape (ELEMENTS_X_PER_THREAD, ELEMENTS_Y_PER_THREAD) rS_col_rcp: RMEM fragment of this thread's columnwise scale reciprocals with shape (ELEMENTS_X_PER_THREAD,) gO_col_thread: GMEM fragment of this thread's columnwise quantized output with shape (ELEMENTS_X_PER_THREAD, ELEMENTS_Y_PER_THREAD) - + This function quantizes its fragments using the input and columnwise scales, and write to its output fragment """ rS_col_rcp_u32 = cute.recast_tensor(rS_col_rcp, Uint32) @@ -292,7 +313,7 @@ def _quantize_colwise( # Each thread quantizes its own rows with the columnwise reciprocals and write its output to GMEM. for row_idx in cutlass.range_constexpr(ELEMENTS_Y_PER_THREAD): rX_row_i32 = cute.recast_tensor(rX[None, row_idx], Int32) - for packed_idx in cutlass.range_constexpr(ELEMENTS_X_PER_THREAD // 2): # packed in 2 + for packed_idx in cutlass.range_constexpr(cute.size(rO_col_u32)): # 4 values in one pack rO_col_u32[packed_idx] = scale_and_pack( rX_row_i32[2 * packed_idx], rS_col_rcp_u32[2 * packed_idx], @@ -405,6 +426,71 @@ def quantize_bidimensional_mxfp8_g2r( load_atom, gX_thread[(None, row_idx),], rX[None, row_idx], cache_policy=cache_policy ) + _quantize_rowwise( + rX, + gO_row_thread, + gS_row_thread, + output_store_atom=store_atom, + row_scale_store_atom=row_scale_atom, + cache_policy=cache_policy, + scale_and_pack=scale_and_pack, + THREADS_X_PER_MX_BLOCK=THREADS_X_PER_MX_BLOCK, + ELEMENTS_Y_PER_THREAD=ELEMENTS_Y_PER_THREAD, + ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, + SCALE_DTYPE=SCALE_DTYPE, + FP8_DTYPE=cfg.FP8_DTYPE, + ) + + sAmaxs_col_scratch = _gather_partial_colwise_amaxes( + rX, + CTA_THREADS_Y=CTA_THREADS_Y, + CTA_THREADS_X=CTA_THREADS_X, + ELEMENTS_Y_PER_THREAD=ELEMENTS_Y_PER_THREAD, + ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, + DTYPE=cfg.DTYPE, + ) + + sS_col = _reduce_partial_colwise_amaxes( + rS_col_rcp, + sAmaxs_col_scratch, + CTA_THREADS_Y=CTA_THREADS_Y, + CTA_THREADS_X=CTA_THREADS_X, + ELEMENTS_Y_PER_THREAD=ELEMENTS_Y_PER_THREAD, + ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, + DTYPE=cfg.DTYPE, + FP8_DTYPE=cfg.FP8_DTYPE, + ) + + # Copy the reduced columnwise scale reciprocals to RMEM. + # Threads that quantize different rows share the same columnwise scale + rS_col_rcp = cute.make_rmem_tensor((ELEMENTS_X_PER_THREAD,), cfg.DTYPE) + cute.autovec_copy(sS_col[None, tidx % CTA_THREADS_X], rS_col_rcp) + + _write_colwise_scale( + rS_col_rcp, + gS_col_thread, + col_scale_atom=col_scale_atom, + cache_policy=cache_policy, + CTA_THREADS_Y=CTA_THREADS_Y, + CTA_THREADS_X=CTA_THREADS_X, + ELEMENTS_Y_PER_THREAD=ELEMENTS_Y_PER_THREAD, + ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, + DTYPE=cfg.DTYPE, + FP8_DTYPE=cfg.FP8_DTYPE, + SCALE_DTYPE=SCALE_DTYPE, + ) + + _quantize_colwise( + rX, + rS_col_rcp, + gO_col_thread, + store_atom=store_atom, + cache_policy=cache_policy, + scale_and_pack=scale_and_pack, + ELEMENTS_Y_PER_THREAD=ELEMENTS_Y_PER_THREAD, + ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, + FP8_DTYPE=cfg.FP8_DTYPE, + ) class MXFP8QuantizeRegisterBidimensionalKernel(MXFP8QuantizeKernelBase): @@ -422,6 +508,7 @@ class MXFP8QuantizeRegisterBidimensionalKernel(MXFP8QuantizeKernelBase): def __init__(self, cfg: MXFP8QuantizeConfig): self.cfg = cfg + # One CTA only quantize one columnwise MXFP8 block at a time self._ELEMENTS_Y_PER_THREAD = MXFP8_BLOCK_SCALING_SIZE // self._CTA_THREADS_Y # 4 @cute.jit From c51a3b8e8b408db1440336905a0d2468a5a63d65 Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Fri, 9 Oct 2026 04:50:18 +0000 Subject: [PATCH 26/58] refactor Signed-off-by: Kaining Zhong --- .../mxfp8/quantize_mxfp8_g2r_bidimensional.py | 84 +++++++++---------- 1 file changed, 38 insertions(+), 46 deletions(-) diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py index 9db22913ddd..5e2d1a6a8aa 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py @@ -68,20 +68,20 @@ def extract_mx_block_amax(values_i32: cute.Tensor) -> Uint32: for row_idx in cutlass.range_constexpr(ELEMENTS_Y_PER_THREAD): rX_row_i32 = cute.recast_tensor(rX[None, row_idx], Int32) reciprocal = bf16_mx_scale_reciprocal(extract_mx_block_amax(rX_row_i32), MAX_EXPONENT) - reciprocal_packed = reciprocal | (reciprocal << Uint32(16)) - for pack_idx in cutlass.range_constexpr(cute.size(rO_row_u32)): + reciprocal_pair = reciprocal | (reciprocal << Uint32(16)) + for pack_idx in cutlass.range_constexpr(cute.size(rO_row_u32)): # 4 output bytes in one pack rO_row_u32[pack_idx] = scale_and_pack( rX_row_i32[2 * pack_idx], - reciprocal_packed, + reciprocal_pair, rX_row_i32[2 * pack_idx + 1], - reciprocal_packed, + reciprocal_pair, ) # rO_row_u32 and rO_row share the same memory cute.copy( output_store_atom, rO_row, gO_row_thread[(None, row_idx),], cache_policy=cache_policy ) if is_writer_thread: - rS_row[0] = Uint8(bf16_mx_scale_bytes_pair(reciprocal_packed)).bitcast(SCALE_DTYPE) + rS_row[0] = Uint8(bf16_mx_scale_bytes_pair(reciprocal_pair)).bitcast(SCALE_DTYPE) cute.copy( row_scale_store_atom, rS_row, @@ -120,10 +120,9 @@ def _gather_partial_colwise_amaxes( amax = abs_max_x2_bf16(amax, rX_i32[bf16x2_idx, row_idx]) rAmax_col_i32[bf16x2_idx] = amax - if cutlass.const_expr(ELEMENTS_X_PER_THREAD == 16): - scratch_swizzle = cute.make_swizzle(3, 1, 5) - else: - scratch_swizzle = cute.make_swizzle(2, 1, 5) + # Keep eight BF16 values (16 bytes) contiguous. XOR BF16 offset bit 6 into + # bit 3 so each 128-bit warp transaction uses all 32 banks without overlap. + scratch_swizzle = cute.make_swizzle(1, 3, 3) # Prepare amax reduction scratch space in SMEM. sAmaxs_col_scratch_layout = cute.make_composed_layout( @@ -131,7 +130,7 @@ def _gather_partial_colwise_amaxes( 0, cute.zipped_product( rAmax_col.layout, # Each thread contributes (ELEMENTS_X_PER_THREAD,) - cute.make_layout((CTA_THREADS_Y, CTA_THREADS_X)), # Total threads in the CTA + cute.make_layout((CTA_THREADS_Y, CTA_THREADS_X)), # Total threads in the CTA ), ) allocator = cutlass.memory.SmemAllocator() @@ -156,11 +155,10 @@ def _reduce_partial_colwise_amaxes( ELEMENTS_X_PER_THREAD: cutlass.Constexpr[int], DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], FP8_DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], -) -> None: +) -> cute.Tensor: """ sAmaxs_col_scratch: the SMEM buffer of layout (ELEMENTS_X_PER_THREAD, (CTA_THREADS_Y, CTA_THREADS_X)) that contains each threads' own columnwise amax of its ELEMENTS_X_PER_THREAD columns - rS_col_rcp: the RMEM fragment of this thread's columnwise scale reciprocals with shape (ELEMENTS_X_PER_THREAD,) This function reduces the partial amax buffer of layout (ELEMENTS_X_PER_THREAD, (CTA_THREADS_Y, CTA_THREADS_X)) along the Y axis, to a single columnwise scale SMEM buffer of layout (ELEMENTS_X_PER_THREAD * CTA_THREADS_X), @@ -175,13 +173,11 @@ def _reduce_partial_colwise_amaxes( REDUCTION_THREADS_NUM = CTA_THREADS_X * ELEMENTS_X_PER_THREAD // 2 is_reduction_thread = tidx < REDUCTION_THREADS_NUM # How many bf16x2 pairs we need to reduce - REDUCTION_i32_PACKS = TILE_COLS // REDUCTION_THREADS_NUM // 2 + REDUCTION_i32_PAIRS = TILE_COLS // REDUCTION_THREADS_NUM // 2 - # Allocate colwise scale SMEM buffer with swizzling - if cutlass.const_expr(ELEMENTS_X_PER_THREAD == 16): - scratch_swizzle = cute.make_swizzle(3, 1, 5) - else: - scratch_swizzle = cute.make_swizzle(2, 1, 5) + # Use the same 16-byte swizzle as the partial-max scratch. Consecutive + # column-pair reducers still access distinct banks through this layout. + scratch_swizzle = cute.make_swizzle(1, 3, 3) # SMEM buffer for colwise scale reduction where mode 0 is how many reciprocals a thread has, # and mode 1 is how many threads we have in X axis, and the layout size is just TILE_COLS sS_col_layout = cute.make_composed_layout( @@ -196,22 +192,22 @@ def _reduce_partial_colwise_amaxes( sS_col = allocator.allocate_tensor(DTYPE, sS_col_layout, byte_alignment=16) sS_col_i32 = cute.recast_tensor(sS_col, Int32) - # Colwise amaxes layout: each reduction thread owns CTA_THREADS_Y rows, with REDUCTION_i32_PACKS column packs + # Colwise amaxes layout: each reduction thread owns CTA_THREADS_Y rows, with REDUCTION_i32_PAIRS column pairs _, tv_col_amaxes_layout = cute.make_layout_tv( thr_layout=cute.make_layout((REDUCTION_THREADS_NUM, 1), stride=(1, REDUCTION_THREADS_NUM)), val_layout=cute.make_layout( - (REDUCTION_i32_PACKS, CTA_THREADS_Y), stride=(1, REDUCTION_i32_PACKS) + (REDUCTION_i32_PAIRS, CTA_THREADS_Y), stride=(1, REDUCTION_i32_PAIRS) ), ) - # Colwise rcp layout: each reduction thread owns only 1 row, with REDUCTION_i32_PACKS column packs + # Colwise rcp layout: each reduction thread owns only 1 row, with REDUCTION_i32_PAIRS column pairs _, tv_col_rcp_layout = cute.make_layout_tv( thr_layout=cute.make_layout((REDUCTION_THREADS_NUM, 1), stride=(1, REDUCTION_THREADS_NUM)), - val_layout=cute.make_layout((REDUCTION_i32_PACKS, 1), stride=(1, REDUCTION_i32_PACKS)), + val_layout=cute.make_layout((REDUCTION_i32_PAIRS, 1), stride=(1, REDUCTION_i32_PAIRS)), ) - + sS_col_i32 = cute.composition(sS_col_i32, tv_col_rcp_layout) - # Each reduction process `REDUCTION_i32_PACKS` bf16x2 packs by visiting all rows (CTA_THREADS_Y) + # Each reduction process `REDUCTION_i32_PAIRS` bf16x2 pairs by visiting all rows (CTA_THREADS_Y) # to extract the amax of that column sAmaxs_col_scratch_i32 = cute.recast_tensor(sAmaxs_col_scratch, Int32) MAX_EXPONENT = 8 if FP8_DTYPE is cutlass.Float8E4M3FN else 15 @@ -222,23 +218,21 @@ def _reduce_partial_colwise_amaxes( # Expose the thread's value mode as (partial rows, column pairs). sAmaxs_col_frag_i32 = cute.composition( sAmaxs_col_frag_i32, - cute.make_layout((CTA_THREADS_Y, REDUCTION_i32_PACKS), stride=(REDUCTION_i32_PACKS, 1)), + cute.make_layout((CTA_THREADS_Y, REDUCTION_i32_PAIRS), stride=(REDUCTION_i32_PAIRS, 1)), ) - # Iterative over all packed columns - for pack_idx in cutlass.range_constexpr(REDUCTION_i32_PACKS): - # Select the packed column - sAmaxs_col_frag_i32_rows = sAmaxs_col_frag_i32[None, pack_idx] + # Iterative over all column pairs + for pair_idx in cutlass.range_constexpr(REDUCTION_i32_PAIRS): + # Select the pair column + sAmaxs_col_frag_i32_rows = sAmaxs_col_frag_i32[None, pair_idx] - # Reduce across all rows' partial amaxes to obtain the columnwise amax for this column pack + # Reduce across all rows' partial amaxes to obtain the columnwise amax for this column pair amaxX2 = sAmaxs_col_frag_i32_rows[0] for scratch_row_idx in cutlass.range_constexpr(1, CTA_THREADS_Y): amaxX2 = abs_max_x2_bf16(amaxX2, sAmaxs_col_frag_i32_rows[scratch_row_idx]) - # Obtain the columnwise scale reciprocal for this column pack - sS_col_i32[tidx, pack_idx] = bf16_mx_scale_reciprocal_pair( - amaxX2, MAX_EXPONENT - ) + # Obtain the columnwise scale reciprocal for this column pair + sS_col_i32[tidx, pair_idx] = bf16_mx_scale_reciprocal_pair(amaxX2, MAX_EXPONENT) cute.arch.sync_threads() return sS_col @@ -273,12 +267,11 @@ def _write_colwise_scale( rS_col = cute.make_rmem_tensor((ELEMENTS_X_PER_THREAD,), SCALE_DTYPE) rS_col_u32 = cute.recast_tensor(rS_col, Uint32) # Get the colwise scale from the reciprocal - for packed_idx in cutlass.range_constexpr(cute.size(rS_col_u32)): - rS_col_u32[packed_idx] = bf16_mx_scale_bytes_pair(rS_col_rcp_u32[2 * packed_idx]) | ( - bf16_mx_scale_bytes_pair(rS_col_rcp_u32[2 * packed_idx + 1]) << Uint32(16) + for pack_idx in cutlass.range_constexpr(cute.size(rS_col_u32)): # 4 output bytes in one pack + rS_col_u32[pack_idx] = bf16_mx_scale_bytes_pair(rS_col_rcp_u32[2 * pack_idx]) | ( + bf16_mx_scale_bytes_pair(rS_col_rcp_u32[2 * pack_idx + 1]) << Uint32(16) ) - # Use RMEM buffer to stash the columnwise scale in packed bytes and flush it to GMEM - # using vectorized store at once + # Use RMEM buffer to stash the columnwise scale and flush it to GMEM using vectorized store at once cute.copy(col_scale_atom, rS_col, gS_col_thread, cache_policy=cache_policy) @@ -313,12 +306,12 @@ def _quantize_colwise( # Each thread quantizes its own rows with the columnwise reciprocals and write its output to GMEM. for row_idx in cutlass.range_constexpr(ELEMENTS_Y_PER_THREAD): rX_row_i32 = cute.recast_tensor(rX[None, row_idx], Int32) - for packed_idx in cutlass.range_constexpr(cute.size(rO_col_u32)): # 4 values in one pack - rO_col_u32[packed_idx] = scale_and_pack( - rX_row_i32[2 * packed_idx], - rS_col_rcp_u32[2 * packed_idx], - rX_row_i32[2 * packed_idx + 1], - rS_col_rcp_u32[2 * packed_idx + 1], + for pack_idx in cutlass.range_constexpr(cute.size(rO_col_u32)): # 4 output bytes in one pack + rO_col_u32[pack_idx] = scale_and_pack( + rX_row_i32[2 * pack_idx], + rS_col_rcp_u32[2 * pack_idx], + rX_row_i32[2 * pack_idx + 1], + rS_col_rcp_u32[2 * pack_idx + 1], ) cute.copy(store_atom, rO_col, gO_col_thread[(None, row_idx),], cache_policy=cache_policy) @@ -451,7 +444,6 @@ def quantize_bidimensional_mxfp8_g2r( ) sS_col = _reduce_partial_colwise_amaxes( - rS_col_rcp, sAmaxs_col_scratch, CTA_THREADS_Y=CTA_THREADS_Y, CTA_THREADS_X=CTA_THREADS_X, From f14e433781ba19ca84105e023fdc796ac12dc4fd Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 9 Oct 2026 04:52:08 +0000 Subject: [PATCH 27/58] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py index 5e2d1a6a8aa..886b9eb27bd 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py @@ -69,7 +69,9 @@ def extract_mx_block_amax(values_i32: cute.Tensor) -> Uint32: rX_row_i32 = cute.recast_tensor(rX[None, row_idx], Int32) reciprocal = bf16_mx_scale_reciprocal(extract_mx_block_amax(rX_row_i32), MAX_EXPONENT) reciprocal_pair = reciprocal | (reciprocal << Uint32(16)) - for pack_idx in cutlass.range_constexpr(cute.size(rO_row_u32)): # 4 output bytes in one pack + for pack_idx in cutlass.range_constexpr( + cute.size(rO_row_u32) + ): # 4 output bytes in one pack rO_row_u32[pack_idx] = scale_and_pack( rX_row_i32[2 * pack_idx], reciprocal_pair, @@ -267,7 +269,9 @@ def _write_colwise_scale( rS_col = cute.make_rmem_tensor((ELEMENTS_X_PER_THREAD,), SCALE_DTYPE) rS_col_u32 = cute.recast_tensor(rS_col, Uint32) # Get the colwise scale from the reciprocal - for pack_idx in cutlass.range_constexpr(cute.size(rS_col_u32)): # 4 output bytes in one pack + for pack_idx in cutlass.range_constexpr( + cute.size(rS_col_u32) + ): # 4 output bytes in one pack rS_col_u32[pack_idx] = bf16_mx_scale_bytes_pair(rS_col_rcp_u32[2 * pack_idx]) | ( bf16_mx_scale_bytes_pair(rS_col_rcp_u32[2 * pack_idx + 1]) << Uint32(16) ) @@ -306,7 +310,9 @@ def _quantize_colwise( # Each thread quantizes its own rows with the columnwise reciprocals and write its output to GMEM. for row_idx in cutlass.range_constexpr(ELEMENTS_Y_PER_THREAD): rX_row_i32 = cute.recast_tensor(rX[None, row_idx], Int32) - for pack_idx in cutlass.range_constexpr(cute.size(rO_col_u32)): # 4 output bytes in one pack + for pack_idx in cutlass.range_constexpr( + cute.size(rO_col_u32) + ): # 4 output bytes in one pack rO_col_u32[pack_idx] = scale_and_pack( rX_row_i32[2 * pack_idx], rS_col_rcp_u32[2 * pack_idx], From 5f9ece21f3bc3241bd1817493a489a96ccd07133 Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Fri, 9 Oct 2026 13:57:11 -0700 Subject: [PATCH 28/58] fix Signed-off-by: Kaining Zhong --- .../cast/mxfp8/quantize_mxfp8_common.py | 12 +- .../mxfp8/quantize_mxfp8_g2r_rowwise_2lane.py | 112 +++++++++++------- 2 files changed, 72 insertions(+), 52 deletions(-) diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_common.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_common.py index 249e178abd6..14d3b88ca66 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_common.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_common.py @@ -329,15 +329,9 @@ def __call__( scale_stride_matches = Boolean(True) if cutlass.const_expr(not self.cfg.WITH_GEMM_SWIZZLED_SCALES): scale_stride_matches = mS_row.shape[1] == N // 32 - # Column tiles use CUDA's grid.y axis. Inputs this wide always use two tiles. - grid_cols_fit = ( - N - <= 65535 - * self.g2r_rowwise_2lane._CTA_THREADS_X - * self.g2r_rowwise_2lane._ELEMENTS_PER_THREAD - * 2 - ) - if N % 128 == 0 and input_is_aligned and scale_stride_matches and grid_cols_fit: + # Column tiles use grid.x and row tiles spill from grid.y into grid.z, so + # no shape exceeds the grid limits. + if N % 128 == 0 and input_is_aligned and scale_stride_matches: dispatched_to_specialized = True self.g2r_rowwise_2lane( mX, diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_2lane.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_2lane.py index edc6cf850f5..a264b578575 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_2lane.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_2lane.py @@ -8,7 +8,7 @@ import cutlass from cutlass import cute -from cutlass import Boolean, Float32, Int32, Int64, Uint32, Uint8 +from cutlass import Boolean, Float32, Int32, Int64, Uint32 from cuda.bindings.driver import CUstream from transformer_engine.common.CuTeDSL.utils import ( @@ -44,6 +44,8 @@ def quantize_rowwise_mxfp8_g2r( CHECK_BOUNDS: cutlass.Constexpr[bool], input_policy: Int64, output_policy: Int64, + cta_row: Int32, + cta_col: Int32, ): """Quantize a CTA's (TILE_ROWS, TILE_COLS, NUM_TILES) GMEM tiles. @@ -97,7 +99,6 @@ def quantize_rowwise_mxfp8_g2r( gX_thread = cute.composition(gX_tile, tv_payload)[tidx, None] valid = Boolean(True) if cutlass.const_expr(CHECK_BOUNDS): - cta_row, cta_col, _ = cute.arch.block_idx() row = cta_row * TILE_ROWS + tidx // CTA_THREADS_X col = (cta_col * NUM_TILES + tile_idx) * TILE_COLS + ( tidx % CTA_THREADS_X @@ -137,16 +138,14 @@ def extract_mx_block_amax(values_i32: cute.Tensor) -> Float32: amax = extract_mx_block_amax(rX_tile_i32) exponent = cvt_f32_to_fp8e8m0fnu(amax * cfg.MAX_NORM_RCP) - # Collect 4 scale bytes and write them using one thread - byte = Uint32(exponent.bitcast(Uint8)) - c1 = cute.arch.shuffle_sync_down(byte, 2) - c2 = cute.arch.shuffle_sync_down(byte, 4) - c3 = cute.arch.shuffle_sync_down(byte, 6) - if tidx % 8 == 0 and tile_valid_flags[tile_idx]: - scale_word = cute.make_tensor( - cute.recast_ptr(mS_thread.iterator, dtype=Uint32), cute.make_layout(1) - ) - scale_word[0] = byte | (c1 << Uint32(8)) | (c2 << Uint32(16)) | (c3 << Uint32(24)) + # The even lane of each pair stores its block's scale byte, as cast_rowwise.cuh does. + # A warp's stores coalesce in either layout: one 16-byte run when packed, four + # 4-byte runs when GEMM-swizzled. Gathering the bytes into 32-bit words first + # touches the same sectors but adds shuffles and a branch (+12% instructions). + # TODO(kainingz): not using vectorized word store seems to be faster in CuTeDSL. + # Check if we should port this to CUDA + if tidx % 2 == 0 and tile_valid_flags[tile_idx]: + mS_thread[0] = exponent # Compute and write the quantized MXFP8 values to GMEM reciprocal_packed = exp2_bf16x2_rcp(exponent) @@ -164,14 +163,17 @@ class MXFP8QuantizeG2RRowwise2LaneKernel(MXFP8QuantizeKernelBase): One CTA tile has TILE_ROWS = CTA_THREADS_Y and TILE_COLS = CTA_THREADS_X * _ELEMENTS_PER_THREAD. The CTA processes NUM_TILES = 1 or 2 in the original input matrix. Both scale - formats use the same payload tiling; a composed layout maps swizzled scales. + formats use the same payload tiling; swizzled scale offsets are computed directly. _L2_CACHED_CTA_PERCENT tunes the share of CTAs whose input loads use normal L2 retention: 0 streams all input and 100 caches all input normally. """ _CTA_THREADS_X = 32 _ELEMENTS_PER_THREAD = 16 + # cast_rowwise.cuh's tier table: the two smaller tiers stream all of their input. _L2_CACHED_CTA_PERCENT = 40 + _SMALL_TIER_CACHED_PERCENT = 0 + _MAX_GRID_Y = 65535 def __init__(self, cfg: MXFP8QuantizeConfig): self.cfg = cfg @@ -201,6 +203,7 @@ def __call__( CTA_THREADS_Y=8, CTA_THREADS_X=self._CTA_THREADS_X, NUM_TILES=1, + L2_CACHED_CTA_PERCENT=self._SMALL_TIER_CACHED_PERCENT, ) elif output_bytes <= 48 << 20: self._launch( @@ -212,6 +215,7 @@ def __call__( CTA_THREADS_Y=8, CTA_THREADS_X=self._CTA_THREADS_X, NUM_TILES=2, + L2_CACHED_CTA_PERCENT=self._SMALL_TIER_CACHED_PERCENT, ) elif output_bytes <= 96 << 20: self._launch( @@ -223,6 +227,7 @@ def __call__( CTA_THREADS_Y=4, CTA_THREADS_X=self._CTA_THREADS_X, NUM_TILES=2, + L2_CACHED_CTA_PERCENT=self._L2_CACHED_CTA_PERCENT, ) else: self._launch( @@ -234,6 +239,7 @@ def __call__( CTA_THREADS_Y=8, CTA_THREADS_X=self._CTA_THREADS_X, NUM_TILES=2, + L2_CACHED_CTA_PERCENT=self._L2_CACHED_CTA_PERCENT, ) @cute.jit @@ -247,8 +253,9 @@ def _launch( CTA_THREADS_Y: cutlass.Constexpr[int], CTA_THREADS_X: cutlass.Constexpr[int], NUM_TILES: cutlass.Constexpr[int], + L2_CACHED_CTA_PERCENT: cutlass.Constexpr[int], ): - assert 0 <= self._L2_CACHED_CTA_PERCENT <= 100 + assert 0 <= L2_CACHED_CTA_PERCENT <= 100 assert self._ELEMENTS_PER_THREAD == 16 assert CTA_THREADS_X % 8 == 0 assert CTA_THREADS_Y * CTA_THREADS_X % 32 == 0 @@ -265,8 +272,15 @@ def _launch( grid_cols = cute.ceil_div(N, CTA_COLS) check_bounds = M % CTA_ROWS != 0 or N % CTA_COLS != 0 grid = Int64(grid_rows) * grid_cols - first_streaming_cta = grid * self._L2_CACHED_CTA_PERCENT // 100 - + first_streaming_cta = grid * L2_CACHED_CTA_PERCENT // 100 + + # Column tiles vary fastest, so consecutive CTAs read consecutive memory along a + # row, as cast_rowwise.cuh's flat grid does. Putting row tiles on x instead + # sends concurrent CTAs down a column strip, one 2 KiB piece per row, and costs + # up to 9% at large shapes. Row tiles past grid.y's limit spill into grid.z. + grid_y = cutlass.min(grid_rows, self._MAX_GRID_Y) + grid_z = cute.ceil_div(grid_rows, self._MAX_GRID_Y) + launch_grid = [grid_cols, grid_y, grid_z] if grid > 0: if check_bounds: self.kernel( @@ -278,8 +292,9 @@ def _launch( CTA_THREADS_Y=CTA_THREADS_Y, CTA_THREADS_X=CTA_THREADS_X, NUM_TILES=NUM_TILES, + L2_CACHED_CTA_PERCENT=L2_CACHED_CTA_PERCENT, CHECK_BOUNDS=True, - ).launch(grid=[grid_rows, grid_cols], block=[THREADS_PER_CTA], stream=stream) + ).launch(grid=launch_grid, block=[THREADS_PER_CTA], stream=stream) else: self.kernel( mX, @@ -290,8 +305,9 @@ def _launch( CTA_THREADS_Y=CTA_THREADS_Y, CTA_THREADS_X=CTA_THREADS_X, NUM_TILES=NUM_TILES, + L2_CACHED_CTA_PERCENT=L2_CACHED_CTA_PERCENT, CHECK_BOUNDS=False, - ).launch(grid=[grid_rows, grid_cols], block=[THREADS_PER_CTA], stream=stream) + ).launch(grid=launch_grid, block=[THREADS_PER_CTA], stream=stream) @cute.kernel def kernel( @@ -304,6 +320,7 @@ def kernel( CTA_THREADS_Y: cutlass.Constexpr[int], CTA_THREADS_X: cutlass.Constexpr[int], NUM_TILES: cutlass.Constexpr[int], + L2_CACHED_CTA_PERCENT: cutlass.Constexpr[int], CHECK_BOUNDS: cutlass.Constexpr[bool], ): """Skip the CTA when the noop flag is set; otherwise run the kernel body.""" @@ -316,6 +333,7 @@ def kernel( CTA_THREADS_Y=CTA_THREADS_Y, CTA_THREADS_X=CTA_THREADS_X, NUM_TILES=NUM_TILES, + L2_CACHED_CTA_PERCENT=L2_CACHED_CTA_PERCENT, CHECK_BOUNDS=CHECK_BOUNDS, ) @@ -329,13 +347,16 @@ def _kernel_main( CTA_THREADS_Y: cutlass.Constexpr[int], CTA_THREADS_X: cutlass.Constexpr[int], NUM_TILES: cutlass.Constexpr[int], + L2_CACHED_CTA_PERCENT: cutlass.Constexpr[int], CHECK_BOUNDS: cutlass.Constexpr[bool], ): """Construct CTA tile views and run the specialized quantization.""" # Dispatch checks 32-byte input alignment before selecting this kernel. gX = cute.make_tensor(mX.iterator.align(32), mX.layout) - # CUDA's x/y axes hold row/column tiles in the same order as local_tile. - cta_coord = cute.arch.block_idx()[:2] + # grid.x walks column tiles; grid.y and grid.z together walk row tiles. + bidx, bidy, bidz = cute.arch.block_idx() + cta_row = bidz * self._MAX_GRID_Y + bidy + cta_coord = (cta_row, bidx) M, N = mX.shape TILE_ROWS = CTA_THREADS_Y TILE_COLS = CTA_THREADS_X * self._ELEMENTS_PER_THREAD @@ -343,10 +364,11 @@ def _kernel_main( CTA_COLS = TILE_COLS * NUM_TILES gS = mS - # Apply the swizzled scale layout if requested + # Use the swizzled layout as is. Composing it with (M, N // 32) layouts to flatten + # its shape made CuTe emit runtime integer div/mod on every scale address, + # about 7x this kernel's instruction count. if cutlass.const_expr(self.cfg.WITH_GEMM_SWIZZLED_SCALES): - mS_swizzled, _ = derive_swizzled_scale_layout(M, N, True, False, mS, None) - gS = cute.composition(mS_swizzled, (cute.make_layout(M), cute.make_layout(N // 32))) + gS, _ = derive_swizzled_scale_layout(M, N, True, False, mS, None) # CTA_TILER is the block that a CTA processes in total, which consists of NUM_TILES horizontal tiles CTA_TILER = (CTA_ROWS, CTA_COLS) @@ -356,12 +378,12 @@ def _kernel_main( gS_cta = cute.local_tile(gS, CTA_SCALE_TILER, cta_coord) fraction = Float32(1.0) - if cutlass.const_expr(self._L2_CACHED_CTA_PERCENT > 0): + if cutlass.const_expr(L2_CACHED_CTA_PERCENT > 0): fraction = Float32(0.0) - # A linear index is needed only for the L2 policy split. - cta_id = cute.crd2idx( - (Int64(cta_coord[0]), Int64(cta_coord[1])), cute.arch.grid_dim()[:2] - ) + # A linear index is needed only for the L2 policy split. It follows launch + # order, so the cached CTAs are the first ones scheduled. + gdx, _, _ = cute.arch.grid_dim() + cta_id = Int64(cta_row) * gdx + bidx if cta_id >= first_streaming_cta: fraction = Float32(1.0) input_policy = create_l2_policy(False, fraction) @@ -375,18 +397,22 @@ def _kernel_main( gO_tiles = cute.local_tile(gO_cta, TILER, (0, None)) mS_tiles = cute.local_tile(gS_cta, SCALE_TILER, (0, None)) - quantize_rowwise_mxfp8_g2r( - gX_tiles, - gO_tiles, - mS_tiles, - self.cfg, - M, - N, - CTA_THREADS_Y=CTA_THREADS_Y, - CTA_THREADS_X=CTA_THREADS_X, - ELEMENTS_PER_THREAD=self._ELEMENTS_PER_THREAD, - NUM_TILES=NUM_TILES, - CHECK_BOUNDS=CHECK_BOUNDS, - input_policy=input_policy, - output_policy=output_policy, - ) + # The last grid.z slice can overshoot the row tiles; those CTAs have no work. + if cta_row * CTA_ROWS < M: + quantize_rowwise_mxfp8_g2r( + gX_tiles, + gO_tiles, + mS_tiles, + self.cfg, + M, + N, + CTA_THREADS_Y=CTA_THREADS_Y, + CTA_THREADS_X=CTA_THREADS_X, + ELEMENTS_PER_THREAD=self._ELEMENTS_PER_THREAD, + NUM_TILES=NUM_TILES, + CHECK_BOUNDS=CHECK_BOUNDS, + input_policy=input_policy, + output_policy=output_policy, + cta_row=cta_coord[0], + cta_col=cta_coord[1], + ) From 870b23cbe23e3073c0d3d5898cccba015d378595 Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Fri, 9 Oct 2026 15:22:49 -0700 Subject: [PATCH 29/58] disable invariant temporarily Signed-off-by: Kaining Zhong --- .../CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py | 4 +++- .../CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_2lane.py | 4 +++- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py index 886b9eb27bd..2f194e804cf 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py @@ -404,7 +404,9 @@ def quantize_bidimensional_mxfp8_g2r( cute.nvgpu.CopyG2ROp(), cfg.DTYPE, num_bits_per_copy=ELEMENTS_X_PER_THREAD * cfg.DTYPE.width, - invariant=True, + # TODO(kainingz): due to https://github.com/NVIDIA/cutlass/issues/3726 using + # invariant=True here will drop the cache policy + invariant=False, ) store_atom = cute.make_copy_atom( cute.nvgpu.CopyR2GOp(), diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_2lane.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_2lane.py index a264b578575..dc39ecfb02d 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_2lane.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_2lane.py @@ -82,7 +82,9 @@ def quantize_rowwise_mxfp8_g2r( cute.nvgpu.CopyG2ROp(), cfg.DTYPE, num_bits_per_copy=LOAD_WIDTH, - invariant=True, # hint that data being read stays unchanged during the kernel + # TODO(kainingz): due to https://github.com/NVIDIA/cutlass/issues/3726 using + # invariant=True here will drop the cache policy + invariant=False, ) store_atom = cute.make_copy_atom( cute.nvgpu.CopyR2GOp(), From 4ef3404dfe6eb1efc3cbb51df83790afae6af9c7 Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Fri, 9 Oct 2026 17:02:12 -0700 Subject: [PATCH 30/58] fix Signed-off-by: Kaining Zhong --- .../mxfp8/quantize_mxfp8_g2r_rowwise_2lane.py | 54 +++++++++++++++++-- 1 file changed, 51 insertions(+), 3 deletions(-) diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_2lane.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_2lane.py index dc39ecfb02d..7f285677357 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_2lane.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_2lane.py @@ -46,6 +46,7 @@ def quantize_rowwise_mxfp8_g2r( output_policy: Int64, cta_row: Int32, cta_col: Int32, + num_elems: Int64, ): """Quantize a CTA's (TILE_ROWS, TILE_COLS, NUM_TILES) GMEM tiles. @@ -105,7 +106,10 @@ def quantize_rowwise_mxfp8_g2r( col = (cta_col * NUM_TILES + tile_idx) * TILE_COLS + ( tidx % CTA_THREADS_X ) * ELEMENTS_PER_THREAD - valid = row < M and col < N + # When we apply the (-1, 1024) manually constructed layout to the input tensor, + # the last row might not be a full row so we need to check OOB with element number + # instead of just M and N + valid = col < N and Int64(row) * N + col < num_elems tile_valid_flags[tile_idx] = valid rX[None, tile_idx].fill(0) if valid: @@ -194,7 +198,40 @@ def __call__( mWorkspace: Optional[cute.Tensor], stream: CUstream, ): - output_bytes = Int64(mX.shape[0]) * mX.shape[1] + M, N = mX.shape + num_elems = Int64(M) * N + # View the input tensor as a layout with last dimension = 1024 when not swizzled and the number of elements can be expressed using i32 + # However, M*N might not be divisible by 1024, so the last row of this layout might be a partial row with fewer than 1024 elements, + # which is why we need to pass num_elems to the kernel so OOB checks will identify these regions as invalid offsets correctly + if cutlass.const_expr(not self.cfg.WITH_GEMM_SWIZZLED_SCALES): + if num_elems <= 2**31 - 1: + rows = cute.ceil_div(num_elems, 1024) + flat = cute.make_layout((rows, 1024), stride=(1024, 1)) + self._launch_tiers( + cute.make_tensor(mX.iterator, flat), + cute.make_tensor(mO_row.iterator, flat), + cute.make_tensor(mS_row.iterator, cute.make_layout((rows, 32), stride=(32, 1))), + mNoop, + stream, + num_elems, + ) + else: + self._launch_tiers(mX, mO_row, mS_row, mNoop, stream, num_elems) + else: + self._launch_tiers(mX, mO_row, mS_row, mNoop, stream, num_elems) + + @cute.jit + def _launch_tiers( + self, + mX: cute.Tensor, + mO_row: cute.Tensor, + mS_row: cute.Tensor, + mNoop: cute.Pointer, + stream: CUstream, + num_elems: Int64, + ): + """Pick the launch configuration for this many elements (cast_rowwise.cuh's tiers).""" + output_bytes = num_elems if output_bytes <= 12 << 20: self._launch( mX, @@ -202,6 +239,7 @@ def __call__( mS_row, mNoop, stream, + num_elems, CTA_THREADS_Y=8, CTA_THREADS_X=self._CTA_THREADS_X, NUM_TILES=1, @@ -214,6 +252,7 @@ def __call__( mS_row, mNoop, stream, + num_elems, CTA_THREADS_Y=8, CTA_THREADS_X=self._CTA_THREADS_X, NUM_TILES=2, @@ -226,6 +265,7 @@ def __call__( mS_row, mNoop, stream, + num_elems, CTA_THREADS_Y=4, CTA_THREADS_X=self._CTA_THREADS_X, NUM_TILES=2, @@ -238,6 +278,7 @@ def __call__( mS_row, mNoop, stream, + num_elems, CTA_THREADS_Y=8, CTA_THREADS_X=self._CTA_THREADS_X, NUM_TILES=2, @@ -252,6 +293,7 @@ def _launch( mS: cute.Tensor, mNoop: cute.Pointer, stream: CUstream, + num_elems: Int64, CTA_THREADS_Y: cutlass.Constexpr[int], CTA_THREADS_X: cutlass.Constexpr[int], NUM_TILES: cutlass.Constexpr[int], @@ -272,7 +314,7 @@ def _launch( grid_rows = cute.ceil_div(M, CTA_ROWS) grid_cols = cute.ceil_div(N, CTA_COLS) - check_bounds = M % CTA_ROWS != 0 or N % CTA_COLS != 0 + check_bounds = M % CTA_ROWS != 0 or N % CTA_COLS != 0 or num_elems != Int64(M) * N grid = Int64(grid_rows) * grid_cols first_streaming_cta = grid * L2_CACHED_CTA_PERCENT // 100 @@ -291,6 +333,7 @@ def _launch( mS, mNoop, first_streaming_cta, + num_elems, CTA_THREADS_Y=CTA_THREADS_Y, CTA_THREADS_X=CTA_THREADS_X, NUM_TILES=NUM_TILES, @@ -304,6 +347,7 @@ def _launch( mS, mNoop, first_streaming_cta, + num_elems, CTA_THREADS_Y=CTA_THREADS_Y, CTA_THREADS_X=CTA_THREADS_X, NUM_TILES=NUM_TILES, @@ -319,6 +363,7 @@ def kernel( mS: cute.Tensor, mNoop: cute.Pointer, first_streaming_cta: Int64, + num_elems: Int64, CTA_THREADS_Y: cutlass.Constexpr[int], CTA_THREADS_X: cutlass.Constexpr[int], NUM_TILES: cutlass.Constexpr[int], @@ -332,6 +377,7 @@ def kernel( mO, mS, first_streaming_cta, + num_elems, CTA_THREADS_Y=CTA_THREADS_Y, CTA_THREADS_X=CTA_THREADS_X, NUM_TILES=NUM_TILES, @@ -346,6 +392,7 @@ def _kernel_main( mO: cute.Tensor, mS: cute.Tensor, first_streaming_cta: Int64, + num_elems: Int64, CTA_THREADS_Y: cutlass.Constexpr[int], CTA_THREADS_X: cutlass.Constexpr[int], NUM_TILES: cutlass.Constexpr[int], @@ -417,4 +464,5 @@ def _kernel_main( output_policy=output_policy, cta_row=cta_coord[0], cta_col=cta_coord[1], + num_elems=num_elems, ) From ef9877df4a10d5efa09b1461688fb8612b4d3704 Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Mon, 12 Oct 2026 18:45:39 +0000 Subject: [PATCH 31/58] refactor Signed-off-by: Kaining Zhong --- .../cast/mxfp8/quantize_mxfp8_common.py | 146 +++++------------- .../mxfp8/quantize_mxfp8_g2r_bidimensional.py | 69 ++++++--- .../mxfp8/quantize_mxfp8_g2r_rowwise_1lane.py | 37 ++++- .../mxfp8/quantize_mxfp8_g2r_rowwise_2lane.py | 85 ++++++---- .../mxfp8/quantize_mxfp8_g2s_bidimensional.py | 61 +++++--- 5 files changed, 212 insertions(+), 186 deletions(-) diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_common.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_common.py index 14d3b88ca66..b8859210d0b 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_common.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_common.py @@ -19,7 +19,6 @@ str_to_cutlass_dtype, abs_max_x2_bf16, device_compute_capability, - is_packed16, ) from transformer_engine.common.CuTeDSL.activations import ( act_relu, @@ -211,9 +210,32 @@ def __str__(self): __repr__ = __str__ + def is_plain_cast(self) -> bool: + """Whether cfg is a plain cast: no amax, dbias or (d)activation fused in.""" + return not (self.WITH_AMAX or self.WITH_DBIAS or self.WITH_DACT or self.WITH_ACT) + + class MXFP8QuantizeKernelBase(abc.ABC): """Base class for MXFP8 quantize kernels.""" + @staticmethod + def supports(cfg: MXFP8QuantizeConfig) -> bool: # pylint: disable=unused-argument + """Compile time check: whether this specialized kernel handles cfg at all.""" + return False + + @staticmethod + def eligible_to_dispatch( + cfg: MXFP8QuantizeConfig, + mX: cute.Tensor, + mO_row: Optional[cute.Tensor], + mS_row: Optional[cute.Tensor], + mO_col: Optional[cute.Tensor], + mS_col: Optional[cute.Tensor], + mNoop: cute.Pointer, + ) -> Boolean: + """Runtime check: whether we can dispatch to the instantiated kernel with this call's shapes and noops.""" + raise NotImplementedError + @abc.abstractmethod def __call__( self, @@ -281,16 +303,18 @@ def __init__(self, cfg: MXFP8QuantizeConfig): ) # We only need to mask when WITH_ACT is enabled because with activations applied zeros filled by TMA affect block statistics self.general_kernel = MXFP8QuantizeKernel(cfg) if self.ACT_NEED_MASKING else None - self.g2r_rowwise_1lane = MXFP8QuantizeG2RRowwise1LaneKernel(cfg) if cfg.ROWWISE else None - self.g2s_bidim = ( - MXFP8QuantizeSpecializedBidimensionalKernel(cfg) - if cfg.ROWWISE and cfg.COLWISE - else None - ) - self.g2r_rowwise_2lane = MXFP8QuantizeG2RRowwise2LaneKernel(cfg) if cfg.ROWWISE else None - self.g2r_bidim = ( - MXFP8QuantizeRegisterBidimensionalKernel(cfg) if cfg.ROWWISE and cfg.COLWISE else None - ) + + # Specialized cast-only kernels in a specific dispatch priority order so we prefer to dispatch to kernels that come first in the list. + self.specialized_kernels = [ + kernel_cls(cfg) + for kernel_cls in ( + MXFP8QuantizeG2RRowwise2LaneKernel, + MXFP8QuantizeRegisterBidimensionalKernel, + MXFP8QuantizeG2RRowwise1LaneKernel, + MXFP8QuantizeSpecializedBidimensionalKernel, + ) + if kernel_cls.supports(cfg) + ] @cute.jit def __call__( @@ -306,104 +330,12 @@ def __call__( mWorkspace: Optional[cute.Tensor], stream: CUstream, ): - M = mX.shape[0] - N = mX.shape[1] - # Select specialized kernels if possible - plain_cast_only = ( - not self.cfg.WITH_AMAX - and not self.cfg.WITH_DBIAS - and not self.cfg.WITH_DACT - and not self.cfg.WITH_ACT - ) dispatched_to_specialized = False - # The two-lane rowwise and g2r bidimensional kernels use BF16 packed math and 32-byte loads. - # Keep the existing kernels for FP16, fusions, ragged bidimensional tiles, - # misaligned input views and compilers predating 256-bit global loads. - if cutlass.const_expr( - plain_cast_only - and self.cfg.DTYPE is cutlass.BFloat16 - and cutlass.target_version(min_version="12.9") - ): - input_is_aligned = mX.iterator.toint() % 32 == 0 - if cutlass.const_expr(self.cfg.ROWWISE and not self.cfg.COLWISE): - scale_stride_matches = Boolean(True) - if cutlass.const_expr(not self.cfg.WITH_GEMM_SWIZZLED_SCALES): - scale_stride_matches = mS_row.shape[1] == N // 32 - # Column tiles use grid.x and row tiles spill from grid.y into grid.z, so - # no shape exceeds the grid limits. - if N % 128 == 0 and input_is_aligned and scale_stride_matches: - dispatched_to_specialized = True - self.g2r_rowwise_2lane( - mX, - mO_row, - mS_row, - mO_col, - mS_col, - mAmax, - mNoop, - mDActInput, - mWorkspace, - stream, - ) - if cutlass.const_expr( - self.cfg.ROWWISE and self.cfg.COLWISE and not self.cfg.WITH_GEMM_SWIZZLED_SCALES - ): - columnwise_scales_are_aligned = mS_col.iterator.toint() % 16 == 0 - if ( - M % 32 == 0 - and N % 256 == 0 - and M // 32 <= 65535 - and input_is_aligned - and columnwise_scales_are_aligned - ): - dispatched_to_specialized = True - self.g2r_bidim( - mX, - mO_row, - mS_row, - mO_col, - mS_col, - mAmax, - mNoop, - mDActInput, - mWorkspace, - stream, - ) - # The one-lane rowwise and g2s bidimensional kernels support packed16 types (bf16/fp16). - if cutlass.const_expr(plain_cast_only and is_packed16(self.cfg.DTYPE)): - if cutlass.const_expr(self.cfg.ROWWISE and not self.cfg.COLWISE): - g2r_rowwise_1lane_grid_fits = ( - cute.ceil_div(M, self.g2r_rowwise_1lane._TILE_ROWS) <= 65535 - ) - # The one-lane rowwise kernel requires N divisible by 128 for vectorized stores. - if not dispatched_to_specialized and N % 128 == 0 and g2r_rowwise_1lane_grid_fits: - dispatched_to_specialized = True - self.g2r_rowwise_1lane( - mX, - mO_row, - mS_row, - mO_col, - mS_col, - mAmax, - mNoop, - mDActInput, - mWorkspace, - stream, - ) - # The g2s bidimensional kernel supports swizzled scales but we don't dispatch to it for now - if cutlass.const_expr( - self.cfg.ROWWISE and self.cfg.COLWISE and not self.cfg.WITH_GEMM_SWIZZLED_SCALES - ): - g2s_bidim_grid_fits = ( - cute.ceil_div( - M, - self.g2s_bidim._TILE_ROWS * self.g2s_bidim._NUM_TILES_Y, - ) - <= 65535 - ) - if not dispatched_to_specialized and g2s_bidim_grid_fits: + for kernel in self.specialized_kernels: + if not dispatched_to_specialized: + if kernel.eligible_to_dispatch(self.cfg, mX, mO_row, mS_row, mO_col, mS_col, mNoop): dispatched_to_specialized = True - self.g2s_bidim( + kernel( mX, mO_row, mS_row, diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py index 2f194e804cf..27e43988450 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py @@ -8,7 +8,7 @@ import cutlass from cutlass import cute -from cutlass import Int32, Int64, Uint32, Uint8 +from cutlass import Boolean, Int32, Int64, Uint32, Uint8 from cuda.bindings.driver import CUstream # pylint: disable=no-name-in-module from transformer_engine.common.CuTeDSL.utils import ( @@ -20,11 +20,11 @@ ) from transformer_engine.common.CuTeDSL.utils_fp8 import mul_bf16x2x2_cvt_bf16x4_to_fp8x4 from .quantize_mxfp8_common import ( + is_plain_cast, MXFP8_BLOCK_SCALING_SIZE, MXFP8QuantizeConfig, MXFP8QuantizeKernelBase, bf16_pair_magnitude, - noop_flag_is_set, ) @@ -506,6 +506,41 @@ class MXFP8QuantizeRegisterBidimensionalKernel(MXFP8QuantizeKernelBase): _CTA_THREADS_Y = 8 _CTA_THREADS_X = 32 + @staticmethod + def supports(cfg: MXFP8QuantizeConfig) -> bool: + # BF16 packed math and 256-bit loads; packed scales only (no colwise GEMM swizzle yet). + return ( + cfg.is_plain_cast() + and cfg.DTYPE is cutlass.BFloat16 + and cfg.ROWWISE + and cfg.COLWISE + and not cfg.WITH_GEMM_SWIZZLED_SCALES + and cutlass.target_version(min_version="12.9") + ) + + @staticmethod + @cute.jit + def eligible_to_dispatch( + cfg: cutlass.Constexpr[MXFP8QuantizeConfig], + mX: cute.Tensor, + mO_row: Optional[cute.Tensor], + mS_row: Optional[cute.Tensor], + mO_col: Optional[cute.Tensor], + mS_col: Optional[cute.Tensor], + mNoop: cute.Pointer, + ) -> Boolean: + M, N = mX.shape + # Whole 32-row by 256-column tiles, grid.y within limits, and aligned input and + # colwise-scale pointers. + return ( + M % 32 == 0 + and N % 256 == 0 + and M // 32 <= 65535 + and mX.iterator.toint() % 32 == 0 + and mS_col.iterator.toint() % 16 == 0 + and mNoop.toint() == 0 + ) + def __init__(self, cfg: MXFP8QuantizeConfig): self.cfg = cfg # One CTA only quantize one columnwise MXFP8 block at a time @@ -520,7 +555,7 @@ def __call__( mO_col: Optional[cute.Tensor], mS_col: Optional[cute.Tensor], mAmax: Optional[cute.Tensor], - mNoop: cute.Pointer, + mNoop: cute.Pointer, # Unused: the dispatcher routes noop calls to the general kernel mDActInput: Optional[cute.Tensor], mWorkspace: Optional[cute.Tensor], stream: CUstream, @@ -536,7 +571,6 @@ def __call__( mS_row, mO_col, mS_col, - mNoop, stream, ELEMENTS_X_PER_THREAD=16, CLUSTER_WIDTH=4, @@ -548,7 +582,6 @@ def __call__( mS_row, mO_col, mS_col, - mNoop, stream, ELEMENTS_X_PER_THREAD=16, CLUSTER_WIDTH=1, @@ -560,7 +593,6 @@ def __call__( mS_row, mO_col, mS_col, - mNoop, stream, ELEMENTS_X_PER_THREAD=8, CLUSTER_WIDTH=8, @@ -574,7 +606,6 @@ def _launch( mS_row: cute.Tensor, mO_col: cute.Tensor, mS_col: cute.Tensor, - mNoop: cute.Pointer, stream: CUstream, ELEMENTS_X_PER_THREAD: cutlass.Constexpr[int], CLUSTER_WIDTH: cutlass.Constexpr[int], @@ -620,7 +651,6 @@ def _launch( gS_row, gO_col, gS_col, - mNoop, ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, ).launch(grid=grid, block=block, cluster=cluster, min_blocks_per_mp=4, stream=stream) elif N == 4096: @@ -642,7 +672,6 @@ def _launch( gS_row, gO_col, gS_col, - mNoop, ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, ).launch(grid=grid, block=block, cluster=cluster, min_blocks_per_mp=4, stream=stream) else: @@ -652,7 +681,6 @@ def _launch( mS_row, mO_col, mS_col, - mNoop, ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, ).launch(grid=grid, block=block, cluster=cluster, min_blocks_per_mp=4, stream=stream) @@ -664,19 +692,18 @@ def kernel( mS_row: cute.Tensor, mO_col: cute.Tensor, mS_col: cute.Tensor, - mNoop: cute.Pointer, ELEMENTS_X_PER_THREAD: cutlass.Constexpr[int], ): - """Skip the CTA when the noop flag is set; otherwise run the kernel body.""" - if not noop_flag_is_set(mNoop): - self._kernel_main( - mX, - mO_row, - mS_row, - mO_col, - mS_col, - ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, - ) + """Run the kernel body. Like the CUDA specialized cast kernels it does not consume the + cast_noop flag: the dispatcher routes calls with one to the general kernel.""" + self._kernel_main( + mX, + mO_row, + mS_row, + mO_col, + mS_col, + ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, + ) @cute.jit def _kernel_main( diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_1lane.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_1lane.py index 552b6429465..fd76e039fa6 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_1lane.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_1lane.py @@ -11,10 +11,11 @@ import cutlass from cutlass import cute -from cutlass import Float8E8M0FNU, Int32, Uint32, Uint8 +from cutlass import Boolean, Float8E8M0FNU, Int32, Uint32, Uint8 from cuda.bindings.driver import CUstream # pylint: disable=no-name-in-module from transformer_engine.common.CuTeDSL.utils import ( + is_packed16, abs_max_x2_bf16, abs_max_x2_f16, x2_lo_to_f32_bf16, @@ -35,7 +36,6 @@ MXFP8QuantizeKernelBase, MXFP8_BLOCK_SCALING_SIZE, derive_swizzled_scale_layout, - noop_flag_is_set, ) @@ -51,6 +51,29 @@ class MXFP8QuantizeG2RRowwise1LaneKernel(MXFP8QuantizeKernelBase): _TILE_COLS = 1024 _THREADS_PER_CTA = 128 + @staticmethod + def supports(cfg: MXFP8QuantizeConfig) -> bool: + return cfg.is_plain_cast() and is_packed16(cfg.DTYPE) and cfg.ROWWISE and not cfg.COLWISE + + @staticmethod + @cute.jit + def eligible_to_dispatch( + cfg: cutlass.Constexpr[MXFP8QuantizeConfig], + mX: cute.Tensor, + mO_row: Optional[cute.Tensor], + mS_row: Optional[cute.Tensor], + mO_col: Optional[cute.Tensor], + mS_col: Optional[cute.Tensor], + mNoop: cute.Pointer, + ) -> Boolean: + M, N = mX.shape + # N divisible by 128 for vectorized stores, and grid.y within limits. + return ( + N % 128 == 0 + and cute.ceil_div(M, MXFP8QuantizeG2RRowwise1LaneKernel._TILE_ROWS) <= 65535 + and mNoop.toint() == 0 + ) + def __init__(self, cfg: MXFP8QuantizeConfig): self.cfg = cfg # If True, then this kernel will first write each thread's scale byte to a shared memory buffer, @@ -66,7 +89,7 @@ def __call__( mO_col: Optional[cute.Tensor], mS_col: Optional[cute.Tensor], # Unused, kept for API compatibility mAmax: Optional[cute.Tensor], # Unused, kept for API compatibility - mNoop: cute.Pointer, # f32 cast_noop flag; may be null, checked on device + mNoop: cute.Pointer, # Unused: the dispatcher routes noop calls to the general kernel mDActInput: Optional[cute.Tensor], # Unused, kept for API compatibility mWorkspace: Optional[cute.Tensor], # Unused, kept for API compatibility stream: CUstream, @@ -92,7 +115,6 @@ def __call__( mX, mO_row, mS_row, - mNoop, mX.element_type, ).launch(grid=grid, block=block, stream=stream) @@ -102,12 +124,11 @@ def kernel( mX: cute.Tensor, mO_row: cute.Tensor, mS_row: cute.Tensor, - mNoop: cute.Pointer, DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], ): - """Device entry: no-op the CTA when the noop flag is set, else run the quantize main loop.""" - if not noop_flag_is_set(mNoop): - self._kernel_main(mX, mO_row, mS_row, DTYPE) + """Device entry. Like the CUDA specialized cast kernels it does not consume the + cast_noop flag: the dispatcher routes calls with one to the general kernel.""" + self._kernel_main(mX, mO_row, mS_row, DTYPE) @cute.jit def _kernel_main( diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_2lane.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_2lane.py index 7f285677357..1304a47dd00 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_2lane.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_2lane.py @@ -9,7 +9,7 @@ import cutlass from cutlass import cute from cutlass import Boolean, Float32, Int32, Int64, Uint32 -from cuda.bindings.driver import CUstream +from cuda.bindings.driver import CUstream # pylint: disable=no-name-in-module from transformer_engine.common.CuTeDSL.utils import ( abs_max_x2_bf16, @@ -25,7 +25,6 @@ MXFP8QuantizeKernelBase, bf16_pair_magnitude, derive_swizzled_scale_layout, - noop_flag_is_set, ) @@ -35,7 +34,6 @@ def quantize_rowwise_mxfp8_g2r( gO_tiles: cute.Tensor, mS_tiles: cute.Tensor, cfg: cutlass.Constexpr[MXFP8QuantizeConfig], - M: Int32, N: Int32, CTA_THREADS_Y: cutlass.Constexpr[int], CTA_THREADS_X: cutlass.Constexpr[int], @@ -181,6 +179,42 @@ class MXFP8QuantizeG2RRowwise2LaneKernel(MXFP8QuantizeKernelBase): _SMALL_TIER_CACHED_PERCENT = 0 _MAX_GRID_Y = 65535 + @staticmethod + def supports(cfg: MXFP8QuantizeConfig) -> bool: + # BF16 packed math and 256-bit loads; compilers predating those keep the other kernels. + return ( + cfg.is_plain_cast() + and cfg.DTYPE is cutlass.BFloat16 + and cfg.ROWWISE + and not cfg.COLWISE + and cutlass.target_version(min_version="12.9") + ) + + @staticmethod + @cute.jit + def eligible_to_dispatch( + cfg: cutlass.Constexpr[MXFP8QuantizeConfig], + mX: cute.Tensor, + mO_row: Optional[cute.Tensor], + mS_row: Optional[cute.Tensor], + mO_col: Optional[cute.Tensor], + mS_col: Optional[cute.Tensor], + mNoop: cute.Pointer, + ) -> Boolean: + N = mX.shape[1] + scale_stride_matches = Boolean(True) + if cutlass.const_expr(not cfg.WITH_GEMM_SWIZZLED_SCALES): + # The flat view needs the packed scales to be contiguous. + scale_stride_matches = mS_row.shape[1] == N // 32 + # Column tiles use grid.x and row tiles spill from grid.y into grid.z, so no shape + # exceeds the grid limits. + return ( + N % 128 == 0 + and mX.iterator.toint() % 32 == 0 + and scale_stride_matches + and mNoop.toint() == 0 + ) + def __init__(self, cfg: MXFP8QuantizeConfig): self.cfg = cfg @@ -193,7 +227,7 @@ def __call__( mO_col: Optional[cute.Tensor], mS_col: Optional[cute.Tensor], mAmax: Optional[cute.Tensor], - mNoop: cute.Pointer, + mNoop: cute.Pointer, # Unused: the dispatcher routes noop calls to the general kernel mDActInput: Optional[cute.Tensor], mWorkspace: Optional[cute.Tensor], stream: CUstream, @@ -211,14 +245,13 @@ def __call__( cute.make_tensor(mX.iterator, flat), cute.make_tensor(mO_row.iterator, flat), cute.make_tensor(mS_row.iterator, cute.make_layout((rows, 32), stride=(32, 1))), - mNoop, stream, num_elems, ) else: - self._launch_tiers(mX, mO_row, mS_row, mNoop, stream, num_elems) + self._launch_tiers(mX, mO_row, mS_row, stream, num_elems) else: - self._launch_tiers(mX, mO_row, mS_row, mNoop, stream, num_elems) + self._launch_tiers(mX, mO_row, mS_row, stream, num_elems) @cute.jit def _launch_tiers( @@ -226,7 +259,6 @@ def _launch_tiers( mX: cute.Tensor, mO_row: cute.Tensor, mS_row: cute.Tensor, - mNoop: cute.Pointer, stream: CUstream, num_elems: Int64, ): @@ -237,7 +269,6 @@ def _launch_tiers( mX, mO_row, mS_row, - mNoop, stream, num_elems, CTA_THREADS_Y=8, @@ -250,7 +281,6 @@ def _launch_tiers( mX, mO_row, mS_row, - mNoop, stream, num_elems, CTA_THREADS_Y=8, @@ -263,7 +293,6 @@ def _launch_tiers( mX, mO_row, mS_row, - mNoop, stream, num_elems, CTA_THREADS_Y=4, @@ -276,7 +305,6 @@ def _launch_tiers( mX, mO_row, mS_row, - mNoop, stream, num_elems, CTA_THREADS_Y=8, @@ -291,7 +319,6 @@ def _launch( mX: cute.Tensor, mO: cute.Tensor, mS: cute.Tensor, - mNoop: cute.Pointer, stream: CUstream, num_elems: Int64, CTA_THREADS_Y: cutlass.Constexpr[int], @@ -331,7 +358,6 @@ def _launch( mX, mO, mS, - mNoop, first_streaming_cta, num_elems, CTA_THREADS_Y=CTA_THREADS_Y, @@ -345,7 +371,6 @@ def _launch( mX, mO, mS, - mNoop, first_streaming_cta, num_elems, CTA_THREADS_Y=CTA_THREADS_Y, @@ -361,7 +386,6 @@ def kernel( mX: cute.Tensor, mO: cute.Tensor, mS: cute.Tensor, - mNoop: cute.Pointer, first_streaming_cta: Int64, num_elems: Int64, CTA_THREADS_Y: cutlass.Constexpr[int], @@ -370,20 +394,20 @@ def kernel( L2_CACHED_CTA_PERCENT: cutlass.Constexpr[int], CHECK_BOUNDS: cutlass.Constexpr[bool], ): - """Skip the CTA when the noop flag is set; otherwise run the kernel body.""" - if not noop_flag_is_set(mNoop): - self._kernel_main( - mX, - mO, - mS, - first_streaming_cta, - num_elems, - CTA_THREADS_Y=CTA_THREADS_Y, - CTA_THREADS_X=CTA_THREADS_X, - NUM_TILES=NUM_TILES, - L2_CACHED_CTA_PERCENT=L2_CACHED_CTA_PERCENT, - CHECK_BOUNDS=CHECK_BOUNDS, - ) + """Run the kernel body. Like cast_rowwise.cuh, there is no noop handling: the + dispatcher only selects this kernel when no cast_noop flag is passed.""" + self._kernel_main( + mX, + mO, + mS, + first_streaming_cta, + num_elems, + CTA_THREADS_Y=CTA_THREADS_Y, + CTA_THREADS_X=CTA_THREADS_X, + NUM_TILES=NUM_TILES, + L2_CACHED_CTA_PERCENT=L2_CACHED_CTA_PERCENT, + CHECK_BOUNDS=CHECK_BOUNDS, + ) @cute.jit def _kernel_main( @@ -453,7 +477,6 @@ def _kernel_main( gO_tiles, mS_tiles, self.cfg, - M, N, CTA_THREADS_Y=CTA_THREADS_Y, CTA_THREADS_X=CTA_THREADS_X, diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2s_bidimensional.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2s_bidimensional.py index 89cfff90093..f6c6d391d99 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2s_bidimensional.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2s_bidimensional.py @@ -12,7 +12,7 @@ import cutlass from cutlass import cute from cutlass import pipeline -from cutlass import Float32, Float8E8M0FNU, Int32, Int64, Uint32 +from cutlass import Boolean, Float32, Float8E8M0FNU, Int32, Int64, Uint32 from cuda.bindings.driver import CUstream # pylint: disable=no-name-in-module from transformer_engine.common.CuTeDSL.utils import ( @@ -39,7 +39,6 @@ MXFP8_BLOCK_SCALING_SIZE, THREADS_PER_WARP, derive_swizzled_scale_layout, - noop_flag_is_set, ) @@ -249,6 +248,32 @@ class MXFP8QuantizeSpecializedBidimensionalKernel(MXFP8QuantizeKernelBase): _SCALE_ROWS = _TILE_ROWS // MXFP8_BLOCK_SCALING_SIZE _SCALE_COLS = _TILE_COLS // MXFP8_BLOCK_SCALING_SIZE + @staticmethod + def supports(cfg: MXFP8QuantizeConfig) -> bool: + # This kernel supports swizzled scales but we don't dispatch to it for now. + return ( + cfg.is_plain_cast() + and is_packed16(cfg.DTYPE) + and cfg.ROWWISE + and cfg.COLWISE + and not cfg.WITH_GEMM_SWIZZLED_SCALES + ) + + @staticmethod + @cute.jit + def eligible_to_dispatch( + cfg: cutlass.Constexpr[MXFP8QuantizeConfig], + mX: cute.Tensor, + mO_row: Optional[cute.Tensor], + mS_row: Optional[cute.Tensor], + mO_col: Optional[cute.Tensor], + mS_col: Optional[cute.Tensor], + mNoop: cute.Pointer, + ) -> Boolean: + cls = MXFP8QuantizeSpecializedBidimensionalKernel + rows_per_cta = cls._TILE_ROWS * cls._NUM_TILES_Y # packed scales: one tile row + return cute.ceil_div(mX.shape[0], rows_per_cta) <= 65535 and mNoop.toint() == 0 + def __init__(self, cfg: MXFP8QuantizeConfig): self.cfg = cfg if cfg.WITH_GEMM_SWIZZLED_SCALES: @@ -265,7 +290,7 @@ def __call__( mO_col: Optional[cute.Tensor], mS_col: Optional[cute.Tensor], mAmax: Optional[cute.Tensor], - mNoop: cute.Pointer, # f32 cast_noop flag; may be null, checked on device + mNoop: cute.Pointer, # Unused: the dispatcher routes noop calls to the general kernel mDActInput: Optional[cute.Tensor], mWorkspace: Optional[cute.Tensor], stream: CUstream, @@ -329,7 +354,6 @@ def __call__( mX, mS_row, mS_col, - mNoop, mX.element_type, tma_atom, tma_src, @@ -345,7 +369,6 @@ def kernel( mX: cute.Tensor, mS_row: cute.Tensor, mS_col: cute.Tensor, - mNoop: cute.Pointer, dtype: cutlass.Constexpr[Type[cutlass.Numeric]], tma_atom: cute.CopyAtom, tma_src: cute.Tensor, @@ -354,20 +377,20 @@ def kernel( tma_atom_out_col: cute.CopyAtom, tma_dst_out_col: cute.Tensor, ): - """Device entry: no-op the CTA when the noop flag is set, else run the quantize main loop.""" - if not noop_flag_is_set(mNoop): - self._kernel_main( - mX, - mS_row, - mS_col, - dtype, - tma_atom, - tma_src, - tma_atom_out_row, - tma_dst_out_row, - tma_atom_out_col, - tma_dst_out_col, - ) + """Device entry. Like the CUDA specialized cast kernels it does not consume the + cast_noop flag: the dispatcher routes calls with one to the general kernel.""" + self._kernel_main( + mX, + mS_row, + mS_col, + dtype, + tma_atom, + tma_src, + tma_atom_out_row, + tma_dst_out_row, + tma_atom_out_col, + tma_dst_out_col, + ) @cute.jit def _kernel_main( From d6bb25862eb793ce1d99455cc8363ac1b9343806 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 12 Oct 2026 18:47:00 +0000 Subject: [PATCH 32/58] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../common/CuTeDSL/cast/mxfp8/quantize_mxfp8_common.py | 1 - 1 file changed, 1 deletion(-) diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_common.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_common.py index b8859210d0b..748f19756f4 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_common.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_common.py @@ -209,7 +209,6 @@ def __str__(self): __repr__ = __str__ - def is_plain_cast(self) -> bool: """Whether cfg is a plain cast: no amax, dbias or (d)activation fused in.""" return not (self.WITH_AMAX or self.WITH_DBIAS or self.WITH_DACT or self.WITH_ACT) From fffc87490c028869d45f13a3994c1b383b713033 Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Mon, 12 Oct 2026 18:48:29 +0000 Subject: [PATCH 33/58] fix Signed-off-by: Kaining Zhong --- .../CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py | 1 - 1 file changed, 1 deletion(-) diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py index 27e43988450..fa1c2d12807 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py @@ -20,7 +20,6 @@ ) from transformer_engine.common.CuTeDSL.utils_fp8 import mul_bf16x2x2_cvt_bf16x4_to_fp8x4 from .quantize_mxfp8_common import ( - is_plain_cast, MXFP8_BLOCK_SCALING_SIZE, MXFP8QuantizeConfig, MXFP8QuantizeKernelBase, From 2fd0aca375bda99e0bc11757550ce5842439d06a Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Mon, 12 Oct 2026 11:59:00 -0700 Subject: [PATCH 34/58] fix Signed-off-by: Kaining Zhong --- .../mxfp8/quantize_mxfp8_g2r_bidimensional.py | 18 ++++++++++++++--- .../mxfp8/quantize_mxfp8_g2r_rowwise_2lane.py | 20 +++++++++++++++---- 2 files changed, 31 insertions(+), 7 deletions(-) diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py index fa1c2d12807..e0df0abf804 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py @@ -402,7 +402,10 @@ def quantize_bidimensional_mxfp8_g2r( load_atom = cute.make_copy_atom( cute.nvgpu.CopyG2ROp(), cfg.DTYPE, - num_bits_per_copy=ELEMENTS_X_PER_THREAD * cfg.DTYPE.width, + num_bits_per_copy=min( + ELEMENTS_X_PER_THREAD * cfg.DTYPE.width, + MXFP8QuantizeRegisterBidimensionalKernel.max_num_bits_per_copy(), + ), # TODO(kainingz): due to https://github.com/NVIDIA/cutlass/issues/3726 using # invariant=True here will drop the cache policy invariant=False, @@ -505,16 +508,25 @@ class MXFP8QuantizeRegisterBidimensionalKernel(MXFP8QuantizeKernelBase): _CTA_THREADS_Y = 8 _CTA_THREADS_X = 32 + @staticmethod + @cute.jit + def max_num_bits_per_copy() -> int: + """Widest global load to issue: 256-bit loads (`ld.global.v8.b32`) need PTX ISA 8.8, + i.e. CUDA 12.9. Older toolchains split each one into two 128-bit loads, as + cast_rowwise.cuh / cast_bidim.cuh do via util/ptx.cuh's ld_global_nc_b32x8.""" + if cutlass.const_expr(cutlass.target_version(min_version="12.9")): + return 256 + return 128 + @staticmethod def supports(cfg: MXFP8QuantizeConfig) -> bool: - # BF16 packed math and 256-bit loads; packed scales only (no colwise GEMM swizzle yet). + # BF16 packed math; packed scales only (no colwise GEMM swizzle yet). return ( cfg.is_plain_cast() and cfg.DTYPE is cutlass.BFloat16 and cfg.ROWWISE and cfg.COLWISE and not cfg.WITH_GEMM_SWIZZLED_SCALES - and cutlass.target_version(min_version="12.9") ) @staticmethod diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_2lane.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_2lane.py index 1304a47dd00..8df894de373 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_2lane.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_2lane.py @@ -75,8 +75,11 @@ def quantize_rowwise_mxfp8_g2r( # With CHECK_BOUNDS=False DCE will remove it so we don't pay extra cost when we don't use it tile_valid_flags = cute.make_rmem_tensor((NUM_TILES,), Boolean) - # 16 * 2 bytes = 32 bytes which is already the widest - LOAD_WIDTH = ELEMENTS_PER_THREAD * cfg.DTYPE.width + # 16 * 2 bytes = 32 bytes, one 256-bit load where the toolchain supports it + LOAD_WIDTH = min( + ELEMENTS_PER_THREAD * cfg.DTYPE.width, + MXFP8QuantizeG2RRowwise2LaneKernel.max_num_bits_per_copy(), + ) load_atom = cute.make_copy_atom( cute.nvgpu.CopyG2ROp(), cfg.DTYPE, @@ -179,15 +182,24 @@ class MXFP8QuantizeG2RRowwise2LaneKernel(MXFP8QuantizeKernelBase): _SMALL_TIER_CACHED_PERCENT = 0 _MAX_GRID_Y = 65535 + @staticmethod + @cute.jit + def max_num_bits_per_copy() -> int: + """Widest global load to issue: 256-bit loads (`ld.global.v8.b32`) need PTX ISA 8.8, + i.e. CUDA 12.9. Older toolchains split each one into two 128-bit loads, as + cast_rowwise.cuh / cast_bidim.cuh do via util/ptx.cuh's ld_global_nc_b32x8.""" + if cutlass.const_expr(cutlass.target_version(min_version="12.9")): + return 256 + return 128 + @staticmethod def supports(cfg: MXFP8QuantizeConfig) -> bool: - # BF16 packed math and 256-bit loads; compilers predating those keep the other kernels. + # BF16 packed math. return ( cfg.is_plain_cast() and cfg.DTYPE is cutlass.BFloat16 and cfg.ROWWISE and not cfg.COLWISE - and cutlass.target_version(min_version="12.9") ) @staticmethod From 8afd230a0d7455903ac7f24d3e3851fb2b32d5d5 Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Mon, 12 Oct 2026 13:11:25 -0700 Subject: [PATCH 35/58] also support swizzled scale for g2r bidim Signed-off-by: Kaining Zhong --- .../mxfp8/quantize_mxfp8_g2r_bidimensional.py | 63 +++++++++++-------- 1 file changed, 36 insertions(+), 27 deletions(-) diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py index e0df0abf804..1b9f6ee4efb 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py @@ -24,6 +24,7 @@ MXFP8QuantizeConfig, MXFP8QuantizeKernelBase, bf16_pair_magnitude, + derive_swizzled_scale_layout, ) @@ -417,8 +418,12 @@ def quantize_bidimensional_mxfp8_g2r( ) row_scale_atom = cute.make_copy_atom(cute.nvgpu.CopyR2GOp(), SCALE_DTYPE, num_bits_per_copy=8) + # GEMM-swizzled colwise scales are not contiguous (adjacent columns land 16 bytes apart), + # so they are stored a byte at a time, as the general kernel does. col_scale_atom = cute.make_copy_atom( - cute.nvgpu.CopyR2GOp(), SCALE_DTYPE, num_bits_per_copy=ELEMENTS_X_PER_THREAD * 8 + cute.nvgpu.CopyR2GOp(), + SCALE_DTYPE, + num_bits_per_copy=8 if cfg.WITH_GEMM_SWIZZLED_SCALES else ELEMENTS_X_PER_THREAD * 8, ) scale_and_pack = mul_bf16x2x2_cvt_bf16x4_to_fp8x4(cfg.FP8_DTYPE) @@ -520,14 +525,8 @@ def max_num_bits_per_copy() -> int: @staticmethod def supports(cfg: MXFP8QuantizeConfig) -> bool: - # BF16 packed math; packed scales only (no colwise GEMM swizzle yet). - return ( - cfg.is_plain_cast() - and cfg.DTYPE is cutlass.BFloat16 - and cfg.ROWWISE - and cfg.COLWISE - and not cfg.WITH_GEMM_SWIZZLED_SCALES - ) + # BF16 packed math. + return cfg.is_plain_cast() and cfg.DTYPE is cutlass.BFloat16 and cfg.ROWWISE and cfg.COLWISE @staticmethod @cute.jit @@ -648,14 +647,17 @@ def _launch( gX = cute.make_tensor(mX.iterator, payload_layout) gO_row = cute.make_tensor(mO_row.iterator, payload_layout) gO_col = cute.make_tensor(mO_col.iterator, payload_layout) - gS_row = cute.make_tensor( - mS_row.iterator, - cute.make_layout((cute.size(mS_row, mode=[0]), 64), stride=(64, 1)), - ) - gS_col = cute.make_tensor( - mS_col.iterator, - cute.make_layout((cute.size(mS_col, mode=[0]), 2048), stride=(2048, 1)), - ) + if cutlass.const_expr(self.cfg.WITH_GEMM_SWIZZLED_SCALES): + gS_row, gS_col = derive_swizzled_scale_layout(M, 2048, True, True, mS_row, mS_col) + else: + gS_row = cute.make_tensor( + mS_row.iterator, + cute.make_layout((cute.size(mS_row, mode=[0]), 64), stride=(64, 1)), + ) + gS_col = cute.make_tensor( + mS_col.iterator, + cute.make_layout((cute.size(mS_col, mode=[0]), 2048), stride=(2048, 1)), + ) self.kernel( gX, gO_row, @@ -669,14 +671,17 @@ def _launch( gX = cute.make_tensor(mX.iterator, payload_layout) gO_row = cute.make_tensor(mO_row.iterator, payload_layout) gO_col = cute.make_tensor(mO_col.iterator, payload_layout) - gS_row = cute.make_tensor( - mS_row.iterator, - cute.make_layout((cute.size(mS_row, mode=[0]), 128), stride=(128, 1)), - ) - gS_col = cute.make_tensor( - mS_col.iterator, - cute.make_layout((cute.size(mS_col, mode=[0]), 4096), stride=(4096, 1)), - ) + if cutlass.const_expr(self.cfg.WITH_GEMM_SWIZZLED_SCALES): + gS_row, gS_col = derive_swizzled_scale_layout(M, 4096, True, True, mS_row, mS_col) + else: + gS_row = cute.make_tensor( + mS_row.iterator, + cute.make_layout((cute.size(mS_row, mode=[0]), 128), stride=(128, 1)), + ) + gS_col = cute.make_tensor( + mS_col.iterator, + cute.make_layout((cute.size(mS_col, mode=[0]), 4096), stride=(4096, 1)), + ) self.kernel( gX, gO_row, @@ -686,12 +691,16 @@ def _launch( ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, ).launch(grid=grid, block=block, cluster=cluster, min_blocks_per_mp=4, stream=stream) else: + if cutlass.const_expr(self.cfg.WITH_GEMM_SWIZZLED_SCALES): + gS_row, gS_col = derive_swizzled_scale_layout(M, N, True, True, mS_row, mS_col) + else: + gS_row, gS_col = mS_row, mS_col self.kernel( mX, mO_row, - mS_row, + gS_row, mO_col, - mS_col, + gS_col, ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, ).launch(grid=grid, block=block, cluster=cluster, min_blocks_per_mp=4, stream=stream) From 87ab24e659e8e6c73e7b60558e6dc95628f37c44 Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Mon, 12 Oct 2026 13:20:47 -0700 Subject: [PATCH 36/58] lint Signed-off-by: Kaining Zhong --- .../cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py | 10 ---------- 1 file changed, 10 deletions(-) diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py index 1b9f6ee4efb..e6794283a03 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py @@ -153,7 +153,6 @@ def _reduce_partial_colwise_amaxes( # Constexprs CTA_THREADS_Y: cutlass.Constexpr[int], CTA_THREADS_X: cutlass.Constexpr[int], - ELEMENTS_Y_PER_THREAD: cutlass.Constexpr[int], ELEMENTS_X_PER_THREAD: cutlass.Constexpr[int], DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], FP8_DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], @@ -249,12 +248,8 @@ def _write_colwise_scale( # Miscellaneous cache_policy: Int64, # Constexprs - CTA_THREADS_Y: cutlass.Constexpr[int], CTA_THREADS_X: cutlass.Constexpr[int], - ELEMENTS_Y_PER_THREAD: cutlass.Constexpr[int], ELEMENTS_X_PER_THREAD: cutlass.Constexpr[int], - DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], - FP8_DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], SCALE_DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], ): """ @@ -462,7 +457,6 @@ def quantize_bidimensional_mxfp8_g2r( sAmaxs_col_scratch, CTA_THREADS_Y=CTA_THREADS_Y, CTA_THREADS_X=CTA_THREADS_X, - ELEMENTS_Y_PER_THREAD=ELEMENTS_Y_PER_THREAD, ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, DTYPE=cfg.DTYPE, FP8_DTYPE=cfg.FP8_DTYPE, @@ -478,12 +472,8 @@ def quantize_bidimensional_mxfp8_g2r( gS_col_thread, col_scale_atom=col_scale_atom, cache_policy=cache_policy, - CTA_THREADS_Y=CTA_THREADS_Y, CTA_THREADS_X=CTA_THREADS_X, - ELEMENTS_Y_PER_THREAD=ELEMENTS_Y_PER_THREAD, ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, - DTYPE=cfg.DTYPE, - FP8_DTYPE=cfg.FP8_DTYPE, SCALE_DTYPE=SCALE_DTYPE, ) From cae94f577d812a6ae42ea890dca23d9f9956084a Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Mon, 12 Oct 2026 21:05:15 +0000 Subject: [PATCH 37/58] add tests Signed-off-by: Kaining Zhong --- tests/jax/test_mxfp8_cutedsl_backend.py | 31 ++++++++++++ .../mxfp8/test_mxfp8_cutedsl_backend.py | 50 ++++++++++++++++++- 2 files changed, 80 insertions(+), 1 deletion(-) diff --git a/tests/jax/test_mxfp8_cutedsl_backend.py b/tests/jax/test_mxfp8_cutedsl_backend.py index 4671ea06869..10f242fb96c 100644 --- a/tests/jax/test_mxfp8_cutedsl_backend.py +++ b/tests/jax/test_mxfp8_cutedsl_backend.py @@ -283,3 +283,34 @@ def test_dtypes(method, act_type, act_desc, fp8_dtype, in_dtype): in_dtype, fp8_dtype, ) + + +# Cover the specialized kernels' shapes +# Note that JAX only emits linear scales and passes no noop tensor +SPECIALIZED_CASES = { + # G2RRowwise2Lane: its four launch tiers (by output size). 4064x4480 leaves a partial last + # row in the flat 1024-column view. + "g2r_rowwise2lane_tier0": ((1024, 4096), QuantizeLayout.ROWWISE, jnp.bfloat16), + "g2r_rowwise2lane_tier1_flat_tail": ((4064, 4480), QuantizeLayout.ROWWISE, jnp.bfloat16), + "g2r_rowwise2lane_tier2": ((8192, 8192), QuantizeLayout.ROWWISE, jnp.bfloat16), + "g2r_rowwise2lane_tier3": ((16384, 8192), QuantizeLayout.ROWWISE, jnp.bfloat16), + # G2RRowwise1Lane + "g2r_rowwise1lane_fp16": ((1024, 4096), QuantizeLayout.ROWWISE, jnp.float16), + # RegisterBidimensional: static N == 2048 / 4096, the wide tile with and without the + # deep cluster, and the narrow tile (N % 512 != 0). + "g2r_bidim_n2048": ((256, 2048), QuantizeLayout.ROWWISE_COLWISE, jnp.bfloat16), + "g2r_bidim_n4096": ((256, 4096), QuantizeLayout.ROWWISE_COLWISE, jnp.bfloat16), + "g2r_bidim_wide": ((256, 6144), QuantizeLayout.ROWWISE_COLWISE, jnp.bfloat16), + "g2r_bidim_wide_cluster": ((4096, 16384), QuantizeLayout.ROWWISE_COLWISE, jnp.bfloat16), + "g2r_bidim_narrow": ((160, 2304), QuantizeLayout.ROWWISE_COLWISE, jnp.bfloat16), + # SpecializedBidimensional (g2s): fp16, or bf16 with N % 256 != 0. + "g2s_bidim_fp16": ((256, 1024), QuantizeLayout.ROWWISE_COLWISE, jnp.float16), + "g2s_bidim_bf16_n384": ((256, 384), QuantizeLayout.ROWWISE_COLWISE, jnp.bfloat16), +} + + +@pytest.mark.parametrize( + "shape,q_layout,in_dtype", SPECIALIZED_CASES.values(), ids=SPECIALIZED_CASES.keys() +) +def test_specialized_kernels(shape, q_layout, in_dtype): + run_test_case("CAST_ONLY", *NO_ACT, shape, q_layout, in_dtype, jnp.float8_e4m3fn) diff --git a/tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py b/tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py index f409f76993d..84d70563501 100644 --- a/tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py +++ b/tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py @@ -9,6 +9,7 @@ import ctypes import os +import re import subprocess import sys import textwrap @@ -59,7 +60,8 @@ (96, 128), ] # (block_rows, block_cols): (1,32)=rowwise, (32,1)=colwise, (32,32)=both. -BLOCK_SIZES = [(1, 32), (32, 1), (32, 32)] +ROW, COL, BOTH = (1, 32), (32, 1), (32, 32) +BLOCK_SIZES = [ROW, COL, BOTH] # Every activation the CuTeDSL backend implements; "desc" is the name it goes by in the # config key (Activation in common/util/cutedsl_utils.h). Gated variants are absent because @@ -370,3 +372,49 @@ def test_sizes(swizzled, method, act, block_size, shape): @pytest.mark.parametrize("swizzled", SWIZZLE_MODES, ids=get_swizzle_id) def test_dtypes(swizzled, method, act, fp8_dtype, in_dtype): run_test_case(method, act, (256, 384), (32, 32), in_dtype, fp8_dtype, swizzled) + + +# Cover the specialized kernels' shapes + +SPECIALIZED_PATH_CASES = [ + # G2RRowwise2Lane: the four launch tiers (by output size) on the flat packed-scale view. + ("g2r_rowwise2lane_tier0", (1024, 4096), ROW, torch.bfloat16, False, "G2RRowwise2Lane", 0, False), + # M*N % 1024 != 0: the flat view's last row is partial. + ("g2r_rowwise2lane_tier1_flat_tail", (4064, 4480), ROW, torch.bfloat16, False, "G2RRowwise2Lane", 0, False), + ("g2r_rowwise2lane_tier2", (8192, 8192), ROW, torch.bfloat16, False, "G2RRowwise2Lane", 0, False), + ("g2r_rowwise2lane_tier3", (16384, 8192), ROW, torch.bfloat16, False, "G2RRowwise2Lane", 0, False), + # Swizzled scales use the 2-D grid instead of the flat view. + ("g2r_rowwise2lane_swizzled", (4096, 4096), ROW, torch.bfloat16, True, "G2RRowwise2Lane", 0, False), + # Partial row and column tiles, and scale padding. + ("g2r_rowwise2lane_swizzled_partial", (160, 384), ROW, torch.bfloat16, True, "G2RRowwise2Lane", 0, False), + # Row tiles past grid.y's limit spill into grid.z. + ("g2r_rowwise2lane_swizzled_grid_z", (524352, 128), ROW, torch.bfloat16, True, "G2RRowwise2Lane", 0, False), + # G2RRowwise1Lane: input only 16-byte aligned, or fp16. + ("g2r_rowwise1lane_misaligned_bf16", (1024, 4096), ROW, torch.bfloat16, False, "G2RRowwise1Lane", 8, False), + ("g2r_rowwise1lane_fp16", (1024, 4096), ROW, torch.float16, False, "G2RRowwise1Lane", 0, False), + # G2RBidim: static N == 2048 / 4096, the wide tile with and without the + # deep cluster, and the narrow tile (N % 512 != 0). + ("g2r_bidim_n2048", (256, 2048), BOTH, torch.bfloat16, False, "G2RBidim", 0, False), + ("g2r_bidim_n2048_swizzled", (256, 2048), BOTH, torch.bfloat16, True, "G2RBidim", 0, False), + ("g2r_bidim_n4096_swizzled", (256, 4096), BOTH, torch.bfloat16, True, "G2RBidim", 0, False), + ("g2r_bidim_wide", (256, 6144), BOTH, torch.bfloat16, False, "G2RBidim", 0, False), + ("g2r_bidim_wide_cluster", (4096, 16384), BOTH, torch.bfloat16, False, "G2RBidim", 0, False), + ("g2r_bidim_narrow", (160, 2304), BOTH, torch.bfloat16, False, "G2RBidim", 0, False), + ("g2r_bidim_narrow_swizzled", (160, 2304), BOTH, torch.bfloat16, True, "G2RBidim", 0, False), + # M % 128 != 0: swizzled scale padding. + ("g2r_bidim_wide_swizzled_pad", (4128, 6144), BOTH, torch.bfloat16, True, "G2RBidim", 0, False), + # G2SBidim: fp16, or bf16 with N % 256 != 0. + ("g2s_bidim_fp16", (256, 1024), BOTH, torch.float16, False, "G2SBidim", 0, False), + ("g2s_bidim_bf16_n384", (256, 384), BOTH, torch.bfloat16, False, "G2SBidim", 0, False), +] + + +@pytest.mark.parametrize( + "shape,block_size,in_dtype,swizzled,expected_kernel,input_offset,noop", + [case[1:] for case in SPECIALIZED_PATH_CASES], + ids=[case[0] for case in SPECIALIZED_PATH_CASES], +) +def test_specialized_kernels( + shape, block_size, in_dtype, swizzled +): + run_test_case("CAST_ONLY", IDENTITY, shape, block_size, in_dtype, tex.DType.kFloat8E4M3,swizzled) From f67b1e5945d194d3e90c5209673166a8b78761af Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Mon, 12 Oct 2026 14:10:58 -0700 Subject: [PATCH 38/58] fix Signed-off-by: Kaining Zhong --- .../mxfp8/test_mxfp8_cutedsl_backend.py | 66 +++++++++++-------- 1 file changed, 39 insertions(+), 27 deletions(-) diff --git a/tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py b/tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py index 84d70563501..8d99ccd7dfa 100644 --- a/tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py +++ b/tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py @@ -9,7 +9,6 @@ import ctypes import os -import re import subprocess import sys import textwrap @@ -269,14 +268,20 @@ def extract_quantized_output(out, rowwise, columnwise, swizzled): return parts -def run_test_case(method, act, shape, block_size, in_dtype, fp8_dtype, swizzled=False): +def run_test_case( + method, act, shape, block_size, in_dtype, fp8_dtype, swizzled=False, input_offset=0 +): """Assert the CuTeDSL and CUDA backends produce bit-identical outputs for the - same input and config. + same input and config. input_offset (in elements) shifts x's data pointer past the + start of its allocation, to test misaligned inputs. """ M, N = shape rowwise = block_size[1] != 1 columnwise = block_size[0] != 1 x, act_input = generate_inputs(M, N, in_dtype) + if input_offset: + buf = torch.empty(input_offset + M * N, dtype=in_dtype, device="cuda") + x = buf[input_offset:].view(M, N).copy_(x) set_cutedsl_backend(False) out_cuda, dbias_cuda = run_quantize( @@ -378,43 +383,50 @@ def test_dtypes(swizzled, method, act, fp8_dtype, in_dtype): SPECIALIZED_PATH_CASES = [ # G2RRowwise2Lane: the four launch tiers (by output size) on the flat packed-scale view. - ("g2r_rowwise2lane_tier0", (1024, 4096), ROW, torch.bfloat16, False, "G2RRowwise2Lane", 0, False), + ("g2r_rowwise2lane_tier0", (1024, 4096), ROW, torch.bfloat16, False, 0), # M*N % 1024 != 0: the flat view's last row is partial. - ("g2r_rowwise2lane_tier1_flat_tail", (4064, 4480), ROW, torch.bfloat16, False, "G2RRowwise2Lane", 0, False), - ("g2r_rowwise2lane_tier2", (8192, 8192), ROW, torch.bfloat16, False, "G2RRowwise2Lane", 0, False), - ("g2r_rowwise2lane_tier3", (16384, 8192), ROW, torch.bfloat16, False, "G2RRowwise2Lane", 0, False), + ("g2r_rowwise2lane_tier1_flat_tail", (4064, 4480), ROW, torch.bfloat16, False, 0), + ("g2r_rowwise2lane_tier2", (8192, 8192), ROW, torch.bfloat16, False, 0), + ("g2r_rowwise2lane_tier3", (16384, 8192), ROW, torch.bfloat16, False, 0), # Swizzled scales use the 2-D grid instead of the flat view. - ("g2r_rowwise2lane_swizzled", (4096, 4096), ROW, torch.bfloat16, True, "G2RRowwise2Lane", 0, False), + ("g2r_rowwise2lane_swizzled", (4096, 4096), ROW, torch.bfloat16, True, 0), # Partial row and column tiles, and scale padding. - ("g2r_rowwise2lane_swizzled_partial", (160, 384), ROW, torch.bfloat16, True, "G2RRowwise2Lane", 0, False), + ("g2r_rowwise2lane_swizzled_partial", (160, 384), ROW, torch.bfloat16, True, 0), # Row tiles past grid.y's limit spill into grid.z. - ("g2r_rowwise2lane_swizzled_grid_z", (524352, 128), ROW, torch.bfloat16, True, "G2RRowwise2Lane", 0, False), + ("g2r_rowwise2lane_swizzled_grid_z", (524352, 128), ROW, torch.bfloat16, True, 0), # G2RRowwise1Lane: input only 16-byte aligned, or fp16. - ("g2r_rowwise1lane_misaligned_bf16", (1024, 4096), ROW, torch.bfloat16, False, "G2RRowwise1Lane", 8, False), - ("g2r_rowwise1lane_fp16", (1024, 4096), ROW, torch.float16, False, "G2RRowwise1Lane", 0, False), + ("g2r_rowwise1lane_misaligned_bf16", (1024, 4096), ROW, torch.bfloat16, False, 8), + ("g2r_rowwise1lane_fp16", (1024, 4096), ROW, torch.float16, False, 0), # G2RBidim: static N == 2048 / 4096, the wide tile with and without the # deep cluster, and the narrow tile (N % 512 != 0). - ("g2r_bidim_n2048", (256, 2048), BOTH, torch.bfloat16, False, "G2RBidim", 0, False), - ("g2r_bidim_n2048_swizzled", (256, 2048), BOTH, torch.bfloat16, True, "G2RBidim", 0, False), - ("g2r_bidim_n4096_swizzled", (256, 4096), BOTH, torch.bfloat16, True, "G2RBidim", 0, False), - ("g2r_bidim_wide", (256, 6144), BOTH, torch.bfloat16, False, "G2RBidim", 0, False), - ("g2r_bidim_wide_cluster", (4096, 16384), BOTH, torch.bfloat16, False, "G2RBidim", 0, False), - ("g2r_bidim_narrow", (160, 2304), BOTH, torch.bfloat16, False, "G2RBidim", 0, False), - ("g2r_bidim_narrow_swizzled", (160, 2304), BOTH, torch.bfloat16, True, "G2RBidim", 0, False), + ("g2r_bidim_n2048", (256, 2048), BOTH, torch.bfloat16, False, 0), + ("g2r_bidim_n2048_swizzled", (256, 2048), BOTH, torch.bfloat16, True, 0), + ("g2r_bidim_n4096_swizzled", (256, 4096), BOTH, torch.bfloat16, True, 0), + ("g2r_bidim_wide", (256, 6144), BOTH, torch.bfloat16, False, 0), + ("g2r_bidim_wide_cluster", (4096, 16384), BOTH, torch.bfloat16, False, 0), + ("g2r_bidim_narrow", (160, 2304), BOTH, torch.bfloat16, False, 0), + ("g2r_bidim_narrow_swizzled", (160, 2304), BOTH, torch.bfloat16, True, 0), # M % 128 != 0: swizzled scale padding. - ("g2r_bidim_wide_swizzled_pad", (4128, 6144), BOTH, torch.bfloat16, True, "G2RBidim", 0, False), + ("g2r_bidim_wide_swizzled_pad", (4128, 6144), BOTH, torch.bfloat16, True, 0), # G2SBidim: fp16, or bf16 with N % 256 != 0. - ("g2s_bidim_fp16", (256, 1024), BOTH, torch.float16, False, "G2SBidim", 0, False), - ("g2s_bidim_bf16_n384", (256, 384), BOTH, torch.bfloat16, False, "G2SBidim", 0, False), + ("g2s_bidim_fp16", (256, 1024), BOTH, torch.float16, False, 0), + ("g2s_bidim_bf16_n384", (256, 384), BOTH, torch.bfloat16, False, 0), ] @pytest.mark.parametrize( - "shape,block_size,in_dtype,swizzled,expected_kernel,input_offset,noop", + "shape,block_size,in_dtype,swizzled,input_offset", [case[1:] for case in SPECIALIZED_PATH_CASES], ids=[case[0] for case in SPECIALIZED_PATH_CASES], ) -def test_specialized_kernels( - shape, block_size, in_dtype, swizzled -): - run_test_case("CAST_ONLY", IDENTITY, shape, block_size, in_dtype, tex.DType.kFloat8E4M3,swizzled) +def test_specialized_kernels(shape, block_size, in_dtype, swizzled, input_offset): + run_test_case( + "CAST_ONLY", + IDENTITY, + shape, + block_size, + in_dtype, + tex.DType.kFloat8E4M3, + swizzled, + input_offset, + ) From 177647e1ca58620db7dfd2a50bda4f0171ed0bf9 Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Mon, 12 Oct 2026 15:24:59 -0700 Subject: [PATCH 39/58] nit Signed-off-by: Kaining Zhong --- tests/cpp/operator/test_cast_mxfp8.cu | 3 +++ tests/jax/test_mxfp8_cutedsl_backend.py | 5 ++--- .../pytorch/mxfp8/test_mxfp8_cutedsl_backend.py | 3 +-- .../mxfp8/quantize_mxfp8_g2r_bidimensional.py | 16 ++++------------ .../mxfp8/quantize_mxfp8_g2r_rowwise_2lane.py | 13 ++----------- .../CuTeDSL/cast/mxfp8/quantize_mxfp8_general.py | 7 +++---- transformer_engine/common/CuTeDSL/utils.py | 10 ++++++++++ 7 files changed, 25 insertions(+), 32 deletions(-) diff --git a/tests/cpp/operator/test_cast_mxfp8.cu b/tests/cpp/operator/test_cast_mxfp8.cu index 9368d4030b5..bc62d333b5f 100644 --- a/tests/cpp/operator/test_cast_mxfp8.cu +++ b/tests/cpp/operator/test_cast_mxfp8.cu @@ -699,6 +699,9 @@ std::vector> matrix_sizes = { {1, 16}, {16, 48}, {128, 128}, + // This case will cover cast_rowwise's last row being partial case, in which + // the kernel will reshape the input tensor with last dim being 1024 so the last row + // is not a full row due to 993*512 % 1024 != 0 {993, 512}, {1024}, {8, 32, 1024}, diff --git a/tests/jax/test_mxfp8_cutedsl_backend.py b/tests/jax/test_mxfp8_cutedsl_backend.py index 10f242fb96c..bd492c967a1 100644 --- a/tests/jax/test_mxfp8_cutedsl_backend.py +++ b/tests/jax/test_mxfp8_cutedsl_backend.py @@ -288,10 +288,9 @@ def test_dtypes(method, act_type, act_desc, fp8_dtype, in_dtype): # Cover the specialized kernels' shapes # Note that JAX only emits linear scales and passes no noop tensor SPECIALIZED_CASES = { - # G2RRowwise2Lane: its four launch tiers (by output size). 4064x4480 leaves a partial last - # row in the flat 1024-column view. + # G2RRowwise2Lane: its four launch tiers (by output size). "g2r_rowwise2lane_tier0": ((1024, 4096), QuantizeLayout.ROWWISE, jnp.bfloat16), - "g2r_rowwise2lane_tier1_flat_tail": ((4064, 4480), QuantizeLayout.ROWWISE, jnp.bfloat16), + "g2r_rowwise2lane_tier1": ((4064, 4480), QuantizeLayout.ROWWISE, jnp.bfloat16), "g2r_rowwise2lane_tier2": ((8192, 8192), QuantizeLayout.ROWWISE, jnp.bfloat16), "g2r_rowwise2lane_tier3": ((16384, 8192), QuantizeLayout.ROWWISE, jnp.bfloat16), # G2RRowwise1Lane diff --git a/tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py b/tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py index 8d99ccd7dfa..a886b63c088 100644 --- a/tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py +++ b/tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py @@ -384,8 +384,7 @@ def test_dtypes(swizzled, method, act, fp8_dtype, in_dtype): SPECIALIZED_PATH_CASES = [ # G2RRowwise2Lane: the four launch tiers (by output size) on the flat packed-scale view. ("g2r_rowwise2lane_tier0", (1024, 4096), ROW, torch.bfloat16, False, 0), - # M*N % 1024 != 0: the flat view's last row is partial. - ("g2r_rowwise2lane_tier1_flat_tail", (4064, 4480), ROW, torch.bfloat16, False, 0), + ("g2r_rowwise2lane_tier1", (4064, 4480), ROW, torch.bfloat16, False, 0), ("g2r_rowwise2lane_tier2", (8192, 8192), ROW, torch.bfloat16, False, 0), ("g2r_rowwise2lane_tier3", (16384, 8192), ROW, torch.bfloat16, False, 0), # Swizzled scales use the 2-D grid instead of the flat view. diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py index e6794283a03..b636d04edfb 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py @@ -17,6 +17,7 @@ bf16_mx_scale_reciprocal_pair, bf16_mx_scale_bytes_pair, create_l2_policy, + max_num_bits_per_copy_ld_global, ) from transformer_engine.common.CuTeDSL.utils_fp8 import mul_bf16x2x2_cvt_bf16x4_to_fp8x4 from .quantize_mxfp8_common import ( @@ -400,7 +401,7 @@ def quantize_bidimensional_mxfp8_g2r( cfg.DTYPE, num_bits_per_copy=min( ELEMENTS_X_PER_THREAD * cfg.DTYPE.width, - MXFP8QuantizeRegisterBidimensionalKernel.max_num_bits_per_copy(), + max_num_bits_per_copy_ld_global(), ), # TODO(kainingz): due to https://github.com/NVIDIA/cutlass/issues/3726 using # invariant=True here will drop the cache policy @@ -497,22 +498,13 @@ class MXFP8QuantizeRegisterBidimensionalKernel(MXFP8QuantizeKernelBase): TILE_COLS = _CTA_THREADS_X * ELEMENTS_X_PER_THREAD = 256 or 512. The 32-row tile contains complete columnwise MXFP8 blocks; shared memory combines the partial column maxima from the eight warps. Requires aligned - BF16 input, M divisible by 32, N divisible by 256 and plain scale layouts. + BF16 input, M divisible by 32 and N divisible by 256; writes packed or GEMM-swizzled + scales. """ _CTA_THREADS_Y = 8 _CTA_THREADS_X = 32 - @staticmethod - @cute.jit - def max_num_bits_per_copy() -> int: - """Widest global load to issue: 256-bit loads (`ld.global.v8.b32`) need PTX ISA 8.8, - i.e. CUDA 12.9. Older toolchains split each one into two 128-bit loads, as - cast_rowwise.cuh / cast_bidim.cuh do via util/ptx.cuh's ld_global_nc_b32x8.""" - if cutlass.const_expr(cutlass.target_version(min_version="12.9")): - return 256 - return 128 - @staticmethod def supports(cfg: MXFP8QuantizeConfig) -> bool: # BF16 packed math. diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_2lane.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_2lane.py index 8df894de373..68c69040cfc 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_2lane.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_2lane.py @@ -15,6 +15,7 @@ abs_max_x2_bf16, create_l2_policy, exp2_bf16x2_rcp, + max_num_bits_per_copy_ld_global, ) from transformer_engine.common.CuTeDSL.utils_fp8 import ( cvt_f32_to_fp8e8m0fnu, @@ -78,7 +79,7 @@ def quantize_rowwise_mxfp8_g2r( # 16 * 2 bytes = 32 bytes, one 256-bit load where the toolchain supports it LOAD_WIDTH = min( ELEMENTS_PER_THREAD * cfg.DTYPE.width, - MXFP8QuantizeG2RRowwise2LaneKernel.max_num_bits_per_copy(), + max_num_bits_per_copy_ld_global(), ) load_atom = cute.make_copy_atom( cute.nvgpu.CopyG2ROp(), @@ -182,16 +183,6 @@ class MXFP8QuantizeG2RRowwise2LaneKernel(MXFP8QuantizeKernelBase): _SMALL_TIER_CACHED_PERCENT = 0 _MAX_GRID_Y = 65535 - @staticmethod - @cute.jit - def max_num_bits_per_copy() -> int: - """Widest global load to issue: 256-bit loads (`ld.global.v8.b32`) need PTX ISA 8.8, - i.e. CUDA 12.9. Older toolchains split each one into two 128-bit loads, as - cast_rowwise.cuh / cast_bidim.cuh do via util/ptx.cuh's ld_global_nc_b32x8.""" - if cutlass.const_expr(cutlass.target_version(min_version="12.9")): - return 256 - return 128 - @staticmethod def supports(cfg: MXFP8QuantizeConfig) -> bool: # BF16 packed math. diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_general.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_general.py index a2b4f5e55ff..aca02a253b1 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_general.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_general.py @@ -7,10 +7,9 @@ Replicates the core logic of quantize_mxfp8.cuh: given a 2D tensor of BF16/FP16 values, quantize to MXFP8 format (FP8E4M3 data + E8M0 per-block scales). -BF16 cast-only calls with aligned input use the register-resident ports of -specialized/cast_rowwise.cuh and specialized/cast_bidim.cuh. Other configurations -use the one-lane rowwise, staged bidimensional or general kernels. All paths share the -TVM-FFI entrypoint. +This is the general kernel: it handles every config. MXFP8QuantizeEntry in +quantize_mxfp8_common.py dispatches plain casts to the specialized kernels when they +apply and falls back to this kernel otherwise. Matches the C++ kernel's tile dimensions and thread layout: CHUNK_DIM_Y = 64, CHUNK_DIM_X = 64, THREADS_PER_CTA = 64 diff --git a/transformer_engine/common/CuTeDSL/utils.py b/transformer_engine/common/CuTeDSL/utils.py index 48f3a956544..d5ffd4abfe9 100644 --- a/transformer_engine/common/CuTeDSL/utils.py +++ b/transformer_engine/common/CuTeDSL/utils.py @@ -27,6 +27,16 @@ logger = logging.getLogger("transformer_engine.cutedsl.utils") +@cute.jit +def max_num_bits_per_copy_ld_global() -> int: + """Widest global load to issue: 256-bit loads (`ld.global.v8.b32`) need PTX ISA 8.8, + i.e. CUDA 12.9. Older toolchains split each one into two 128-bit loads, as + cast_rowwise.cuh / cast_bidim.cuh do via util/ptx.cuh's ld_global_nc_b32x8.""" + if cutlass.const_expr(cutlass.target_version(min_version="12.9")): + return 256 + return 128 + + def device_compute_capability(device_id: Optional[int] = None) -> tuple: """(major, minor) compute capability of a CUDA device (current by default), or (0, 0) if it can't be queried.""" from cuda.core import Device # pylint: disable=no-name-in-module From 0eb539d9ffbda87a27d786d6f710db4291638c42 Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Mon, 12 Oct 2026 23:01:27 +0000 Subject: [PATCH 40/58] fix merge Signed-off-by: Kaining Zhong --- .../cast/mxfp8/group_quantize_mxfp8.py | 4 +- .../cast/mxfp8/quantize_mxfp8_common.py | 4 +- .../cast/mxfp8/quantize_mxfp8_general.py | 74 +------------------ 3 files changed, 6 insertions(+), 76 deletions(-) diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8.py index fbf2c6025da..2a556ca3d7e 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8.py @@ -25,7 +25,7 @@ str_to_cutlass_dtype, device_compute_capability, ) -from transformer_engine.common.CuTeDSL.cast.mxfp8.quantize_mxfp8 import ( +from .quantize_mxfp8_common import ( MXFP8_BLOCK_SCALING_SIZE, SYM_N_DIVISIBILITY, FP8E4M3_MAX_NORM_RCP, @@ -34,6 +34,8 @@ SUPPORTED_DACTIVATIONS, derive_swizzled_scale_layout, noop_flag_is_set, +) +from .quantize_mxfp8_general import ( reduce_rowwise_dbias, quantize_rowwise_mxfp8, quantize_colwise_mxfp8, diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_common.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_common.py index 748f19756f4..467bf38688f 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_common.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_common.py @@ -71,8 +71,8 @@ @cute.jit def derive_swizzled_scale_layout( - M: Int32, - N: Int32, + M: Int32 | Int64, + N: Int32 | Int64, ROWWISE: cutlass.Constexpr[bool], COLWISE: cutlass.Constexpr[bool], mS_row: Optional[cute.Tensor], diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_general.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_general.py index b964f8ed57a..7fe0c383bfb 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_general.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_general.py @@ -29,7 +29,7 @@ import cutlass from cutlass import cute from cutlass import pipeline -from cutlass import Float32, Int32, Int64, Uint16, Uint32, Uint8 +from cutlass import Float32, Float8E8M0FNU, Int32, Int64, Uint16, Uint32, Uint8 from cuda.bindings.driver import CUstream # pylint: disable=no-name-in-module from transformer_engine.common.CuTeDSL.utils import ( @@ -62,77 +62,6 @@ 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 | Int64, - N: Int32 | Int64, - ROWWISE: cutlass.Constexpr[bool], - COLWISE: cutlass.Constexpr[bool], - mS_row: Optional[cute.Tensor], - mS_col: Optional[cute.Tensor], -): - """Derive the swizzled layout for the rowwise and colwise scale tensors.""" - num_scale_cols = cute.ceil_div(N, MXFP8_BLOCK_SCALING_SIZE) - num_scale_rows = cute.ceil_div(M, MXFP8_BLOCK_SCALING_SIZE) - - num_tiles_M = cute.ceil_div(M, 128) - num_tiles_SC = cute.ceil_div(num_scale_cols, 4) - num_tiles_SR = cute.ceil_div(num_scale_rows, 4) - num_tiles_N = cute.ceil_div(N, 128) - - if cutlass.const_expr(ROWWISE): - mS_row = cute.make_tensor( - mS_row.iterator, - cute.make_layout( - ((32, 4, num_tiles_M), (4, num_tiles_SC)), - stride=((16, 4, num_tiles_SC * 512), (1, 512)), - ), - ) - if cutlass.const_expr(COLWISE): - mS_col = cute.make_tensor( - mS_col.iterator, - cute.make_layout( - ((4, num_tiles_SR), (32, 4, num_tiles_N)), - stride=((1, 512), (16, 4, num_tiles_SR * 512)), - ), - ) - return mS_row, mS_col from .quantize_mxfp8_common import ( CUTEDSL_DEBUG_LOGGING, MXFP8QuantizeConfig, @@ -191,7 +120,6 @@ def reduce_rowwise_dbias( return block_dbias - @cute.jit def quantize_rowwise_mxfp8( sX_tile: cute.Tensor, # (TILE_Y, TILE_X) bf16/fp16 smem view, post-TMA From c14c2e3b56f5068c2a3205ce8613b6a01179514a Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Mon, 12 Oct 2026 17:20:46 -0700 Subject: [PATCH 41/58] fix Signed-off-by: Kaining Zhong --- .../cast/mxfp8/group_quantize_mxfp8.py | 35 ++++++++++++------- 1 file changed, 22 insertions(+), 13 deletions(-) diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8.py index 2a556ca3d7e..5db65bf6c3f 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8.py @@ -361,7 +361,7 @@ def __call__( ) # In average how many jobs per tensor (only an average, the actual jobs per tensor may vary) average_jobs_per_tensor = cutlass.max( - Int64(1), cute.ceil_div(estimated_jobs, Int64(num_tensors)) + Int64(1), (estimated_jobs + Int64(num_tensors) - 1) // Int64(num_tensors) ) # Don't launch more CTAs than the average jobs per tensor in case # STATIC_PERSISTENT_WORKERS_PER_SM causes redundancy @@ -879,19 +879,28 @@ def _kernel_main( descs = (None, None, None, None) # Find how many rows this tensor member has and which row the member starts from - if cutlass.const_expr(cfg.SHAPE_REP == SAME_BOTH_DIMS): - member_rows = Int64(first_logical_dim) // Int64(num_tensors) - member_start_row = job_start_row - job_start_row % member_rows - elif cutlass.const_expr(cfg.SHAPE_REP == VARYING_FIRST_DIM): - member_id = self._find_tensor_from_offsets( - mOffsets, num_tensors, job_start_row * member_cols - ) - member_rows = Int64(mFirstDims[member_id]) - member_start_row = Int64(mOffsets[member_id]) // member_cols + if cutlass.const_expr(cfg.COLWISE and cfg.WITH_GEMM_SWIZZLED_SCALES): + # If colwise scale is swizzled, we can't view the grouped tensor as one tensor as a whole + # in which case where the colwise is written to depends on the tensor member's first dimension, + # so we need to actually find out the start of the tensor member the job belongs to + if cutlass.const_expr(cfg.SHAPE_REP == SAME_BOTH_DIMS): + # If SAME_BOTH_DIMS, the tensor member's start offset and rows can be obtained with division + member_rows = Int64(first_logical_dim) // Int64(num_tensors) + member_start_row = job_start_row - job_start_row % member_rows + else: # VARYING_FIRST_DIM + # Otherwise if VARYING_FIRST_DIM, we need to use find these tensor member metadata + # from mFirstDims & mOffsets + member_id = self._find_tensor_from_offsets( + mOffsets, num_tensors, job_start_row * member_cols + ) + member_rows = Int64(mFirstDims[member_id]) + member_start_row = Int64(mOffsets[member_id]) // member_cols else: - raise RuntimeError( - "is_single_tensor requires SAME_BOTH_DIMS or VARYING_FIRST_DIM" - ) + # If no swizzled colwise scale, quantizing a grouped tensor is equivalent to + # quantizing a single tensor as a whole as long as M % 128 == 0 holds, + # so we just pretend the whole tensor is the tensor member this CTA is working on + member_rows = Int64(first_logical_dim) + member_start_row = Int64(0) # The starting offset for this tensor member's scale member_scale_base = ( From 7f5ac46017bfa5dc0b307c991f31e19df553aa7e Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Tue, 13 Oct 2026 14:12:50 -0700 Subject: [PATCH 42/58] add grouped g2r kernels Signed-off-by: Kaining Zhong --- transformer_engine/common/CuTeDSL/__init__.py | 2 +- .../common/CuTeDSL/cast/mxfp8/__init__.py | 2 +- .../cast/mxfp8/group_quantize_mxfp8_common.py | 571 ++++++++++++++++++ .../cast/mxfp8/group_quantize_mxfp8_g2r.py | 209 +++++++ ...fp8.py => group_quantize_mxfp8_general.py} | 384 +----------- .../mxfp8/quantize_mxfp8_g2r_bidimensional.py | 72 ++- .../mxfp8/group_quantize_mxfp8_cutedsl.cuh | 2 +- 7 files changed, 861 insertions(+), 381 deletions(-) create mode 100644 transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8_common.py create mode 100644 transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8_g2r.py rename transformer_engine/common/CuTeDSL/cast/mxfp8/{group_quantize_mxfp8.py => group_quantize_mxfp8_general.py} (79%) diff --git a/transformer_engine/common/CuTeDSL/__init__.py b/transformer_engine/common/CuTeDSL/__init__.py index 3b2a4255666..bca6bec9edd 100644 --- a/transformer_engine/common/CuTeDSL/__init__.py +++ b/transformer_engine/common/CuTeDSL/__init__.py @@ -13,7 +13,7 @@ from transformer_engine.common.CuTeDSL.cast.mxfp8.quantize_mxfp8_common import ( get_mxfp8_quantization_function, ) -from transformer_engine.common.CuTeDSL.cast.mxfp8.group_quantize_mxfp8 import ( +from transformer_engine.common.CuTeDSL.cast.mxfp8.group_quantize_mxfp8_common import ( get_mxfp8_group_quantization_function, ) diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/__init__.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/__init__.py index 54ddad2baab..e8804b071cf 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/__init__.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/__init__.py @@ -5,4 +5,4 @@ """CuTeDSL MXFP8 quantization kernels.""" from . import quantize_mxfp8_common -from . import group_quantize_mxfp8 +from . import group_quantize_mxfp8_common diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8_common.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8_common.py new file mode 100644 index 00000000000..1149d0b75ca --- /dev/null +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8_common.py @@ -0,0 +1,571 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Grouped MXFP8 quantization in CuTeDSL: shared definitions and kernel dispatch. + +Holds the grouped config, the constants the grouped kernels share, and +MXFP8GroupQuantizeEntry, which picks the kernel for a config: the single-tensor entry +(quantize_mxfp8_common.py) on a stacked or flat view of the group, the GMEM-to-RMEM grouped +kernel (group_quantize_mxfp8_g2r.py), or the generic grouped kernel +(group_quantize_mxfp8_general.py). Also compiles and registers the TVM-FFI function the C++ +bridge (group_quantize_mxfp8_cutedsl.cuh) calls. +""" + +# pylint: disable=missing-class-docstring + +import logging +import os +from typing import Optional + +import cutlass +from cutlass import cute +from cutlass import Float32, Int32, Int64 +from cutlass.utils import HardwareInfo +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 .quantize_mxfp8_common import ( + SYM_N_DIVISIBILITY, + FP8E4M3_MAX_NORM_RCP, + FP8E5M2_MAX_NORM_RCP, + SUPPORTED_ACTIVATIONS, + SUPPORTED_DACTIVATIONS, +) + +CUTEDSL_DEBUG_LOGGING = os.environ.get("CUTEDSL_DEBUG_LOGGING", "0") == "1" +logger = logging.getLogger("transformer_engine.cutedsl.mxfp8") + +THREADS_PER_WARP = 32 +# CUDA stores four descriptor arrays: input, act_input, output_rowwise, output_colwise. +NUM_TENSORMAPS = 4 +INPUT_SLOT = 0 +ACT_INPUT_SLOT = 1 +ROWWISE_OUTPUT_SLOT = 2 +COLWISE_OUTPUT_SLOT = 3 + +# Shape representations, mirroring ShapeRepresentation in common/utils.cuh. +SAME_BOTH_DIMS = "sbd" +VARYING_FIRST_DIM = "vfd" +VARYING_LAST_DIM = "vld" +VARYING_BOTH_DIMS = "vbd" +SUPPORTED_SHAPE_REPS = (SAME_BOTH_DIMS, VARYING_FIRST_DIM, VARYING_LAST_DIM, VARYING_BOTH_DIMS) + +# Upper bound on the group size (MAX_SUPPORTED_TENSOR_DESCRIPTORS in grouped_tma.cuh); sizes +# the fixed binary search over the offsets. +MAX_SUPPORTED_TENSORS = 64 +# How many bytes is a TensorMap +BYTES_PER_TENSORMAP = 128 +INT64_PER_TENSORMAP = BYTES_PER_TENSORMAP // 8 +# The byte offset of fields in the TensorMap struct +ROWS_OFFSET = NUM_TENSORMAPS * MAX_SUPPORTED_TENSORS * INT64_PER_TENSORMAP +COLS_OFFSET = ROWS_OFFSET + MAX_SUPPORTED_TENSORS +OFFSETS_OFFSET = COLS_OFFSET + MAX_SUPPORTED_TENSORS +# Size of the TensorMap struct in int64 +TENSOR_MAP_STORAGE_INT64_COUNT = OFFSETS_OFFSET + MAX_SUPPORTED_TENSORS + + +class MXFP8GroupQuantizeConfig: + """Compile-time config for the grouped MXFP8 quantize kernel.""" + + def __init__( + self, + dtype: str, + fp8_dtype: str, + rowwise: bool, + colwise: bool, + shape_rep: str, + with_gemm_swizzled_scales: bool = False, + with_dbias: bool = False, + with_dact: bool = False, + with_act: bool = False, + activation: str = "none", + ): + if 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 + 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 + if not (rowwise or colwise): + raise ValueError("at least one of rowwise or colwise must be true") + self.ROWWISE = rowwise + self.COLWISE = colwise + if shape_rep not in SUPPORTED_SHAPE_REPS: + raise ValueError( + f"unsupported shape representation {shape_rep!r}; expected one of" + f" {SUPPORTED_SHAPE_REPS}" + ) + self.SHAPE_REP = shape_rep + # Only can view the grouped tensor as one tensor when the last dim is same across all tensors + self.IS_SINGLE_TENSOR = shape_rep in (SAME_BOTH_DIMS, VARYING_FIRST_DIM) + self.MAX_NORM_RCP = ( + FP8E4M3_MAX_NORM_RCP if fp8_dtype == "Float8E4M3" else FP8E5M2_MAX_NORM_RCP + ) + + self.WITH_GEMM_SWIZZLED_SCALES = with_gemm_swizzled_scales + if with_dbias and not self.IS_SINGLE_TENSOR: + # mxfp8::group_quantize raises for this. + raise ValueError("dbias is only supported for tensors with a common last dimension") + self.WITH_DBIAS = with_dbias + + if with_dact and with_act: + raise ValueError("with_dact and with_act cannot both be set") + if with_dact: + if activation not in SUPPORTED_DACTIVATIONS: + raise ValueError( + f"unknown activation {activation!r} for with_dact=True; expected one of" + f" {sorted(SUPPORTED_DACTIVATIONS)}" + ) + self.ACTIVATION = activation + elif with_act: + if activation not in SUPPORTED_ACTIVATIONS: + raise ValueError( + f"unknown activation {activation!r} for with_act=True; expected one of" + f" {sorted(SUPPORTED_ACTIVATIONS)}" + ) + self.ACTIVATION = activation + else: + if activation != "none": + raise ValueError("activation must be none when with_dact and with_act are False") + self.ACTIVATION = None + self.WITH_DACT = with_dact + self.WITH_ACT = with_act + + def __str__(self): + return ( + f"MXFP8GroupQuantizeConfig(dtype={self.DTYPE_STR}, fp8_dtype={self.FP8_DTYPE_STR}, " + f"rowwise={self.ROWWISE}, colwise={self.COLWISE}, shape_rep={self.SHAPE_REP}, " + f"swizzled={self.WITH_GEMM_SWIZZLED_SCALES}, with_dbias={self.WITH_DBIAS}, " + f"with_dact={self.WITH_DACT}, with_act={self.WITH_ACT}, " + f"activation={self.ACTIVATION})" + ) + + __repr__ = __str__ + + +class MXFP8GroupQuantizeEntry: + """Pick the kernel for a grouped config at compile time. + + Plain casts (no dbias/act/dact) whose scales need no per-member addressing are run as a + single-tensor cast by the single-tensor entry, which dispatches to the GMEM-to-RMEM + specialized kernels: + + * "stacked": SAME_BOTH_DIMS / VARYING_FIRST_DIM members start on 128-row boundaries and + share the column count, so unless colwise scales are GEMM-swizzled (whose indices + restart per member) the group is byte-for-byte the stacked (sum(rows), cols) tensor. + * "flat_rowwise": with packed (non-swizzled) rowwise scales only, scale i covers elements + 32i..32i+31 of the flat payload whatever the member shapes, as long as every member's + cols are a multiple of 128 (which both grouped kernels require for VARYING_LAST_DIM / + VARYING_BOTH_DIMS). Any group is then a cast of a (total / 128, 128) tensor. + + * "grouped_g2r": the remaining BF16 plain casts (colwise / both directions of + VARYING_LAST_DIM / VARYING_BOTH_DIMS, and swizzled colwise of every representation) run + MXFP8GroupQuantizeG2RKernel, which addresses each member on its own. + + Everything else, and any call these paths cannot take at runtime (a cast-noop flag, + 2^31+ elements, cols not a multiple of 128, unaligned buffers), runs + MXFP8GroupQuantizeKernel. + """ + + def __init__(self, cfg: MXFP8GroupQuantizeConfig, sm_count: int) -> None: + # pylint: disable=import-outside-toplevel + from .group_quantize_mxfp8_general import MXFP8GroupQuantizeKernel + + self.cfg = cfg + self.grouped_kernel = MXFP8GroupQuantizeKernel(cfg, sm_count) + plain_cast = not (cfg.WITH_DBIAS or cfg.WITH_ACT or cfg.WITH_DACT) + self.SINGLE_TENSOR_MODE = None + if plain_cast and cfg.IS_SINGLE_TENSOR: + if not (cfg.COLWISE and cfg.WITH_GEMM_SWIZZLED_SCALES): + self.SINGLE_TENSOR_MODE = "stacked" + elif plain_cast and cfg.ROWWISE and not cfg.COLWISE and not cfg.WITH_GEMM_SWIZZLED_SCALES: + self.SINGLE_TENSOR_MODE = "flat_rowwise" + self.single_tensor_entry = None + self.g2r_kernel = None + if self.SINGLE_TENSOR_MODE is None and plain_cast and cfg.DTYPE is cutlass.BFloat16: + # pylint: disable=import-outside-toplevel + from .quantize_mxfp8_common import MXFP8QuantizeConfig + from .group_quantize_mxfp8_g2r import MXFP8GroupQuantizeG2RKernel + + self.g2r_kernel = MXFP8GroupQuantizeG2RKernel( + cfg.SHAPE_REP, + MXFP8QuantizeConfig( + dtype=cfg.DTYPE_STR, + fp8_dtype=cfg.FP8_DTYPE_STR, + rowwise=cfg.ROWWISE, + colwise=cfg.COLWISE, + with_gemm_swizzled_scales=cfg.WITH_GEMM_SWIZZLED_SCALES, + with_amax=False, + activation="none", + ), + sm_count, + ) + if self.SINGLE_TENSOR_MODE is not None: + # pylint: disable=import-outside-toplevel + from .quantize_mxfp8_common import MXFP8QuantizeConfig, MXFP8QuantizeEntry + + self.single_tensor_entry = MXFP8QuantizeEntry( + MXFP8QuantizeConfig( + dtype=cfg.DTYPE_STR, + fp8_dtype=cfg.FP8_DTYPE_STR, + rowwise=cfg.ROWWISE, + colwise=cfg.COLWISE, + with_gemm_swizzled_scales=cfg.WITH_GEMM_SWIZZLED_SCALES, + with_amax=False, + activation="none", + ) + ) + + @cute.jit + def __call__( + self, + mX: cute.Tensor, + mO_row: Optional[cute.Tensor], + mO_col: Optional[cute.Tensor], + mS_row: cute.Tensor, + mS_col: cute.Tensor, + mOffsets: Optional[cute.Tensor], + mFirstDims: Optional[cute.Tensor], + mLastDims: Optional[cute.Tensor], + mTensormaps: Optional[cute.Tensor], + num_tensors: Int32, + mNoop: cute.Pointer, + mActInput: Optional[cute.Tensor], + mWorkspace: Optional[cute.Tensor], + stream: CUstream, + ) -> None: + grouped_args = ( + mX, + mO_row, + mO_col, + mS_row, + mS_col, + mOffsets, + mFirstDims, + mLastDims, + mTensormaps, + num_tensors, + mNoop, + mActInput, + mWorkspace, + stream, + ) + # The single-tensor kernels index with Int32, so groups of 2^31 or more elements + # stay on the grouped kernel. + num_elems = Int64(mX.shape[0]) * Int64(mX.shape[1]) + if cutlass.const_expr(self.SINGLE_TENSOR_MODE == "stacked"): + # The grouped kernels zero the scale padding and the single-tensor kernels leave it + # untouched, so stack only when there is none: rows are already 128-aligned, and + # N % 128 == 0 leaves no padded rowwise (N / 32 % 4) or colwise (N % 128) columns. + if mX.shape[1] % 128 == 0 and num_elems < Int64(2**31): + # Restate the divisibility the Int32 casts drop -- M by the compiled sym, N by + # the check above -- so the vectorized GMEM copies can prove their alignment. + M = cute.assume(Int32(mX.shape[0]), divby=128) + N = cute.assume(Int32(mX.shape[1]), divby=128) + self._run_single_tensor(mX, mO_row, mO_col, mS_row, mS_col, mNoop, stream, M, N) + else: + self.grouped_kernel(*grouped_args) + elif cutlass.const_expr(self.SINGLE_TENSOR_MODE == "flat_rowwise"): + if num_elems < Int64(2**31): + M = Int32(num_elems // 128) + self._run_single_tensor(mX, mO_row, mO_col, mS_row, mS_col, mNoop, stream, M, 128) + else: + self.grouped_kernel(*grouped_args) + elif cutlass.const_expr(self.g2r_kernel is not None): + eligible = ( + num_elems < Int64(2**31) + and mNoop.toint() == 0 + and mX.iterator.toint() % 32 == 0 + ) + if cutlass.const_expr(self.cfg.COLWISE): + eligible = eligible and mS_col.iterator.toint() % 16 == 0 + if cutlass.const_expr(self.cfg.IS_SINGLE_TENSOR): + # VARYING_LAST_DIM / VARYING_BOTH_DIMS members have 128-aligned cols already. + eligible = eligible and mX.shape[1] % 128 == 0 + if eligible: + self.g2r_kernel( + mX, + mO_row, + mO_col, + mS_row, + mS_col, + mOffsets, + mFirstDims, + mLastDims, + num_tensors, + stream, + ) + else: + self.grouped_kernel(*grouped_args) + else: + self.grouped_kernel(*grouped_args) + + @cute.jit + def _run_single_tensor( + self, + mX: cute.Tensor, + mO_row: Optional[cute.Tensor], + mO_col: Optional[cute.Tensor], + mS_row: cute.Tensor, + mS_col: cute.Tensor, + mNoop: cute.Pointer, + stream: CUstream, + M: Int32, + N: Int32, + ) -> None: + """Quantize the group's flat payload as one (M, N) tensor with the single-tensor entry. + + The single-tensor kernels are compiled for Int32 extents. The scales arrive as flat + buffers; view them at the 2D shapes TE allocates for a single (M, N) tensor, which + carry no padding here (M % 128 == N % 128 == 0). + """ + data_layout = cute.make_layout((M, N), stride=(N, 1)) + x = cute.make_tensor(mX.iterator, data_layout) + o_row = s_row = o_col = s_col = None + if cutlass.const_expr(self.cfg.ROWWISE): + o_row = cute.make_tensor(mO_row.iterator, data_layout) + s_row = cute.make_tensor(mS_row.iterator, cute.make_layout((M, N // 32), stride=(N // 32, 1))) + if cutlass.const_expr(self.cfg.COLWISE): + o_col = cute.make_tensor(mO_col.iterator, data_layout) + s_col = cute.make_tensor(mS_col.iterator, cute.make_layout((M // 32, N), stride=(N, 1))) + self.single_tensor_entry(x, o_row, s_row, o_col, s_col, None, mNoop, None, None, stream) + + +def compile_cutedsl_function_from_cfg(cfg: MXFP8GroupQuantizeConfig): + """Return the compiled CuTeDSL function object for the given grouped config.""" + # CUDA requires the group's first logical dim to be a multiple of 128 (and each + # tensor's rows likewise). VARYING_BOTH_DIMS is the exception: its logical shape is + # [1, total]. The last dim only needs the 16-byte TMA row alignment; a partial 32-element + # scale block at the end of a row is zero-filled by TMA, as in the CUDA kernel. + if cfg.SHAPE_REP == VARYING_BOTH_DIMS: + sym_M = cute.sym_int64() + else: + sym_M = cute.sym_int64(divisibility=128) + sym_N = cute.sym_int64(divisibility=SYM_N_DIVISIBILITY) + logical_shape = (sym_M, sym_N) + + out_dtype = cfg.FP8_DTYPE + scale_dtype = cutlass.Float8E8M0FNU + + out_col_fake = ( + cute.runtime.make_fake_compact_tensor( + out_dtype, + logical_shape, + stride_order=(1, 0), + memspace=cute.AddressSpace.gmem, + assumed_align=16, + ) + if cfg.COLWISE + else None + ) + + out_row_fake = ( + cute.runtime.make_fake_compact_tensor( + out_dtype, + logical_shape, + stride_order=(1, 0), + memspace=cute.AddressSpace.gmem, + assumed_align=16, + ) + if cfg.ROWWISE + else None + ) + + # The kernel only takes the base address of the scale buffers (per-tensor strides + # are derived from cols), so their fake shape is a flat 1D byte run. + # One compact view covers CUDA's descriptor and metadata arrays in their native order. + tensormaps_fake = ( + cute.runtime.make_fake_compact_tensor( + cutlass.Int64, + (TENSOR_MAP_STORAGE_INT64_COUNT,), + stride_order=(0,), + memspace=cute.AddressSpace.gmem, + assumed_align=128, + ) + if not cfg.IS_SINGLE_TENSOR + else None + ) + # The cast-noop flag is an always-present f32 pointer instead of an optional tensor, so + # that one compiled kernel serves both an absent and a present flag (noop_flag_is_set). + 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, + logical_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, + (cute.sym_int64(), cute.sym_int64()), + stride_order=(1, 0), + memspace=cute.AddressSpace.gmem, + assumed_align=4, + ) + if cfg.WITH_DBIAS + else None + ) + + first_dims_fake = ( + cute.runtime.make_fake_compact_tensor( + cutlass.Int64, + (cute.sym_int32(),), + stride_order=(0,), + memspace=cute.AddressSpace.gmem, + assumed_align=8, + ) + if cfg.SHAPE_REP in (VARYING_FIRST_DIM, VARYING_BOTH_DIMS) + else None + ) + + last_dims_fake = ( + cute.runtime.make_fake_compact_tensor( + cutlass.Int64, + (cute.sym_int32(),), + stride_order=(0,), + memspace=cute.AddressSpace.gmem, + assumed_align=8, + ) + if cfg.SHAPE_REP in (VARYING_LAST_DIM, VARYING_BOTH_DIMS) + else None + ) + + offsets_fake = ( + cute.runtime.make_fake_compact_tensor( + cutlass.Int64, + (cute.sym_int32(),), + stride_order=(0,), + memspace=cute.AddressSpace.gmem, + assumed_align=8, + ) + if cfg.SHAPE_REP != SAME_BOTH_DIMS + else None + ) + + sm_count = HardwareInfo().get_device_multiprocessor_count() + kernel_obj = MXFP8GroupQuantizeEntry(cfg, sm_count) + return cute.compile( + kernel_obj, + cute.runtime.make_fake_compact_tensor( # mX + cfg.DTYPE, + logical_shape, + stride_order=(1, 0), + memspace=cute.AddressSpace.gmem, + assumed_align=16, + ), + out_row_fake, # mO_row + out_col_fake, # mO_col + cute.runtime.make_fake_compact_tensor( # mS_row + scale_dtype, + (cute.sym_int64(),), + stride_order=(0,), + memspace=cute.AddressSpace.gmem, + assumed_align=4, + ), + cute.runtime.make_fake_compact_tensor( # mS_col + scale_dtype, + (cute.sym_int64(),), + stride_order=(0,), + memspace=cute.AddressSpace.gmem, + assumed_align=4, + ), + offsets_fake, # mOffsets + first_dims_fake, # mFirstDims + last_dims_fake, # mLastDims + tensormaps_fake, # mTensormaps + Int32(1), # num_tensors, supplied at runtime + noop_fake, # mNoop + act_input_fake, # mActInput + workspace_fake, # mWorkspace + cute.runtime.make_fake_stream(), + options="--enable-tvm-ffi", + ) + + +def get_mxfp8_group_quantization_function( + fn_name: str, + dtype: str, + fp8_dtype: str, + rowwise: bool, + colwise: bool, + shape_rep: str, + with_gemm_swizzled_scales: bool, + with_dbias: bool, + with_dact: bool, + with_act: bool, + activation: str, +) -> bool: + """Compile the grouped 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++ grouped 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 = MXFP8GroupQuantizeConfig( + dtype=dtype, + fp8_dtype=fp8_dtype, + rowwise=rowwise, + colwise=colwise, + shape_rep=shape_rep, + with_gemm_swizzled_scales=with_gemm_swizzled_scales, + with_dbias=with_dbias, + with_dact=with_dact, + with_act=with_act, + activation=activation, + ) + except ValueError as e: + logger.warning( + "CuTeDSL grouped MXFP8 backend does not support this config, " + "falling back to the CUDA C++ kernel: %s", + e, + ) + return False + + logger.debug("Compiling CuTeDSL grouped MXFP8 quantization kernel for %s", cfg) + compiled = compile_cutedsl_function_from_cfg(cfg) + # Register the native TVM-FFI function rather than its Python argument-parsing wrapper; + # see get_mxfp8_quantization_function for why. + 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 grouped 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/group_quantize_mxfp8_g2r.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8_g2r.py new file mode 100644 index 00000000000..70ccda6a7b0 --- /dev/null +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8_g2r.py @@ -0,0 +1,209 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Grouped GMEM-to-RMEM MXFP8 cast for BF16 input. + +Each CTA serves one group member (blockIdx.y), reading the member's rows, columns and element +offset directly -- no search, and, since the tiles are loaded with plain global loads rather +than TMA, no per-member descriptors and no descriptor-update prologue kernel. A member is +viewed as its own (rows, cols) tensor whose scales start at offset / 32, packed or +GEMM-swizzled, and its 32 x 128 tiles are quantized by the register-resident tile function of +the single-tensor bidimensional kernel (quantize_bidimensional_mxfp8_g2r), which runs the +rowwise and/or colwise pass the config asks for. + +Requires every member's rows and cols to be multiples of 128 (both grouped kernels require +this for VARYING_LAST_DIM / VARYING_BOTH_DIMS; for the single-tensor representations the +dispatcher checks cols), and fewer than 2^31 elements per member. +""" + +from typing import Optional + +import cutlass +from cutlass import cute +from cutlass import Int32, Int64 +from cuda.bindings.driver import CUstream # pylint: disable=no-name-in-module + +from transformer_engine.common.CuTeDSL.utils import create_l2_policy +from .quantize_mxfp8_common import ( + MXFP8_BLOCK_SCALING_SIZE, + MXFP8QuantizeConfig, + derive_swizzled_scale_layout, +) +from .quantize_mxfp8_g2r_bidimensional import quantize_bidimensional_mxfp8_g2r +from .group_quantize_mxfp8_common import ( + SAME_BOTH_DIMS, + VARYING_FIRST_DIM, + VARYING_LAST_DIM, +) + + +class MXFP8GroupQuantizeG2RKernel: + """Grid of (CTAs per member, num_tensors); each CTA grid-strides its member's tiles.""" + + _CTA_THREADS_Y = 8 + # Eight lanes of 16 BF16 values (one 256-bit load each) per tile row: 1-8% faster than + # sixteen lanes of 8 values (128-bit loads) on GB200. + _CTA_THREADS_X = 8 + _ELEMENTS_X_PER_THREAD = 16 + _ELEMENTS_Y_PER_THREAD = MXFP8_BLOCK_SCALING_SIZE // _CTA_THREADS_Y # 4 + TILE_ROWS = _CTA_THREADS_Y * _ELEMENTS_Y_PER_THREAD # 32 + TILE_COLS = _CTA_THREADS_X * _ELEMENTS_X_PER_THREAD # 128 + THREADS_PER_CTA = _CTA_THREADS_Y * _CTA_THREADS_X # 64 + # CTAs per SM the grid is sized for, split evenly across the members. Far more than are + # resident: on GB200 the loop's latency is better hidden by more CTAs in flight than by + # each CTA walking more tiles (16 -> 256 per SM took colwise from ~107 to ~92 us at + # 32768x7168). + WORKERS_PER_SM = 256 + + def __init__(self, shape_rep: str, cfg: MXFP8QuantizeConfig, sm_count: int) -> None: + self.cfg = cfg + self.SM_COUNT = sm_count + self.ROWS_VARY = shape_rep not in (SAME_BOTH_DIMS, VARYING_LAST_DIM) + self.COLS_VARY = shape_rep not in (SAME_BOTH_DIMS, VARYING_FIRST_DIM) + self.SAME_BOTH_DIMS = shape_rep == SAME_BOTH_DIMS + + @staticmethod + def supports(cfg: MXFP8QuantizeConfig) -> bool: + return cfg.is_plain_cast() and cfg.DTYPE is cutlass.BFloat16 + + @cute.jit + def __call__( + self, + mX: cute.Tensor, + mO_row: Optional[cute.Tensor], + mO_col: Optional[cute.Tensor], + mS_row: cute.Tensor, + mS_col: cute.Tensor, + mOffsets: Optional[cute.Tensor], + mFirstDims: Optional[cute.Tensor], + mLastDims: Optional[cute.Tensor], + num_tensors: Int32, + stream: CUstream, + ) -> None: + # Spread the persistent worker budget across the members, but never launch more CTAs + # per member than the average member has tiles. + num_elems = Int64(mX.shape[0]) * Int64(mX.shape[1]) + tile_elems = self.TILE_ROWS * self.TILE_COLS + tiles = (num_elems + tile_elems - 1) // tile_elems + avg_tiles = (tiles + Int64(num_tensors) - 1) // Int64(num_tensors) + budget = Int64(self.SM_COUNT * self.WORKERS_PER_SM) // Int64(num_tensors) + ctas_per_member = cutlass.max(Int64(1), cutlass.min(budget, avg_tiles)) + self.kernel( + mX, + mO_row, + mO_col, + mS_row, + mS_col, + mOffsets, + mFirstDims, + mLastDims, + Int64(mX.shape[0]), + Int64(mX.shape[1]), + num_tensors, + ).launch( + grid=[Int32(ctas_per_member), num_tensors, 1], + block=[self.THREADS_PER_CTA, 1, 1], + stream=stream, + ) + + @cute.kernel + def kernel( + self, + mX: cute.Tensor, + mO_row: Optional[cute.Tensor], + mO_col: Optional[cute.Tensor], + mS_row: cute.Tensor, + mS_col: cute.Tensor, + mOffsets: Optional[cute.Tensor], + mFirstDims: Optional[cute.Tensor], + mLastDims: Optional[cute.Tensor], + first_logical_dim: Int64, + last_logical_dim: Int64, + num_tensors: Int32, + ): + cfg = self.cfg + bidx, member_id, _ = cute.arch.block_idx() + gdx, _, _ = cute.arch.grid_dim() + + # This member's extents and element offset in the flat payload. + rows = Int64(0) + cols = Int64(last_logical_dim) + base = Int64(0) + if cutlass.const_expr(self.SAME_BOTH_DIMS): + rows = first_logical_dim // Int64(num_tensors) + base = Int64(member_id) * rows * cols + else: + base = Int64(mOffsets[member_id]) + if cutlass.const_expr(self.ROWS_VARY): + rows = Int64(mFirstDims[member_id]) + else: + rows = first_logical_dim + if cutlass.const_expr(self.COLS_VARY): + cols = Int64(mLastDims[member_id]) + + # Member-local extents fit Int32; the 128-divisibility lets the vectorized copies + # prove their alignment. Members start on 128 x 128 element boundaries and the + # dispatcher checks the payload is 32-byte aligned, as the 256-bit loads need. + M = cute.assume(Int32(rows), divby=128) + N = cute.assume(Int32(cols), divby=128) + scale_base = base // MXFP8_BLOCK_SCALING_SIZE + data_layout = cute.make_layout((M, N), stride=(N, 1)) + gX = cute.make_tensor((mX.iterator + base).align(32), data_layout) + gO_row = gS_row = gO_col = gS_col = None + if cutlass.const_expr(cfg.ROWWISE): + gO_row = cute.make_tensor((mO_row.iterator + base).align(16), data_layout) + gS_row = cute.make_tensor( + mS_row.iterator + scale_base, cute.make_layout((M, N // 32), stride=(N // 32, 1)) + ) + if cutlass.const_expr(cfg.COLWISE): + gO_col = cute.make_tensor((mO_col.iterator + base).align(16), data_layout) + # Member scale offsets are multiples of 512 (members start on 128 x 128 element + # boundaries) and the dispatcher checks the buffer is 16-byte aligned. + gS_col = cute.make_tensor( + (mS_col.iterator + scale_base).align(16), + cute.make_layout((M // 32, N), stride=(N, 1)), + ) + if cutlass.const_expr(cfg.WITH_GEMM_SWIZZLED_SCALES): + gS_row, gS_col = derive_swizzled_scale_layout( + M, N, cfg.ROWWISE, cfg.COLWISE, gS_row, gS_col + ) + + cache_policy = create_l2_policy(True) + tiles_X = N // self.TILE_COLS + num_tiles = (M // self.TILE_ROWS) * tiles_X + # Grid-stride over this member's tiles (column tiles fastest, as the single-tensor + # kernels launch them). An empty member has no tiles. + tile = Int32(bidx) + while tile < num_tiles: + coord = (tile // tiles_X, tile % tiles_X) + self._quantize_tile(gX, gO_row, gS_row, gO_col, gS_col, coord, cache_policy) + tile += Int32(gdx) + + @cute.jit + def _quantize_tile(self, gX, gO_row, gS_row, gO_col, gS_col, coord, cache_policy): + """Quantize the (TILE_ROWS, TILE_COLS) tile at `coord` of one member.""" + cfg = self.cfg + TILER = (self.TILE_ROWS, self.TILE_COLS) + ROW_SCALE_TILER = (self.TILE_ROWS, self.TILE_COLS // MXFP8_BLOCK_SCALING_SIZE) + COL_SCALE_TILER = (self.TILE_ROWS // MXFP8_BLOCK_SCALING_SIZE, self.TILE_COLS) + gO_row_tile = gS_row_tile = gO_col_tile = gS_col_tile = None + if cutlass.const_expr(cfg.ROWWISE): + gO_row_tile = cute.local_tile(gO_row, TILER, coord) + gS_row_tile = cute.local_tile(gS_row, ROW_SCALE_TILER, coord) + if cutlass.const_expr(cfg.COLWISE): + gO_col_tile = cute.local_tile(gO_col, TILER, coord) + gS_col_tile = cute.local_tile(gS_col, COL_SCALE_TILER, coord) + quantize_bidimensional_mxfp8_g2r( + cute.local_tile(gX, TILER, coord), + gO_row_tile, + gS_row_tile, + gO_col_tile, + gS_col_tile, + cfg, + CTA_THREADS_Y=self._CTA_THREADS_Y, + CTA_THREADS_X=self._CTA_THREADS_X, + ELEMENTS_Y_PER_THREAD=self._ELEMENTS_Y_PER_THREAD, + ELEMENTS_X_PER_THREAD=self._ELEMENTS_X_PER_THREAD, + cache_policy=cache_policy, + ) diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8_general.py similarity index 79% rename from transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8.py rename to transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8_general.py index 5db65bf6c3f..8401dca2e73 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8_general.py @@ -2,12 +2,14 @@ # # See LICENSE for license information. -"""Grouped MXFP8 quantization kernel implemented in CuTeDSL.""" +"""Generic grouped MXFP8 quantization kernel implemented in CuTeDSL (TMA -> SMEM). + +Handles every grouped config. MXFP8GroupQuantizeEntry in group_quantize_mxfp8_common.py +dispatches plain casts to faster kernels when they apply and falls back to this one otherwise. +""" # pylint: disable=missing-class-docstring -import logging -import os from typing import Optional, Type import cutlass @@ -16,22 +18,12 @@ from cutlass import Boolean, Float32, Int32, Int64, Float8E8M0FNU from cutlass.cute.nvgpu import cpasync from cutlass.cute.testing import assert_ as runtime_assert -from cutlass.utils import TensorMapUpdateMode, HardwareInfo +from cutlass.utils import TensorMapUpdateMode from cutlass.tensor_utils import TensorMapManager 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 .quantize_mxfp8_common import ( MXFP8_BLOCK_SCALING_SIZE, - SYM_N_DIVISIBILITY, - FP8E4M3_MAX_NORM_RCP, - FP8E5M2_MAX_NORM_RCP, - SUPPORTED_ACTIVATIONS, - SUPPORTED_DACTIVATIONS, derive_swizzled_scale_layout, noop_flag_is_set, ) @@ -40,120 +32,25 @@ quantize_rowwise_mxfp8, quantize_colwise_mxfp8, ) - -CUTEDSL_DEBUG_LOGGING = os.environ.get("CUTEDSL_DEBUG_LOGGING", "0") == "1" -logger = logging.getLogger("transformer_engine.cutedsl.mxfp8") - -THREADS_PER_WARP = 32 -# CUDA stores four descriptor arrays: input, act_input, output_rowwise, output_colwise. -NUM_TENSORMAPS = 4 -INPUT_SLOT = 0 -ACT_INPUT_SLOT = 1 -ROWWISE_OUTPUT_SLOT = 2 -COLWISE_OUTPUT_SLOT = 3 - -# Shape representations, mirroring ShapeRepresentation in common/utils.cuh. -SAME_BOTH_DIMS = "sbd" -VARYING_FIRST_DIM = "vfd" -VARYING_LAST_DIM = "vld" -VARYING_BOTH_DIMS = "vbd" -SUPPORTED_SHAPE_REPS = (SAME_BOTH_DIMS, VARYING_FIRST_DIM, VARYING_LAST_DIM, VARYING_BOTH_DIMS) - -# Upper bound on the group size (MAX_SUPPORTED_TENSOR_DESCRIPTORS in grouped_tma.cuh); sizes -# the fixed binary search over the offsets. -MAX_SUPPORTED_TENSORS = 64 -# How many bytes is a TensorMap -BYTES_PER_TENSORMAP = 128 -INT64_PER_TENSORMAP = BYTES_PER_TENSORMAP // 8 -# The byte offset of fields in the TensorMap struct -ROWS_OFFSET = NUM_TENSORMAPS * MAX_SUPPORTED_TENSORS * INT64_PER_TENSORMAP -COLS_OFFSET = ROWS_OFFSET + MAX_SUPPORTED_TENSORS -OFFSETS_OFFSET = COLS_OFFSET + MAX_SUPPORTED_TENSORS -# Size of the TensorMap struct in int64 -TENSOR_MAP_STORAGE_INT64_COUNT = OFFSETS_OFFSET + MAX_SUPPORTED_TENSORS - - -class MXFP8GroupQuantizeConfig: - """Compile-time config for the grouped MXFP8 quantize kernel.""" - - def __init__( - self, - dtype: str, - fp8_dtype: str, - rowwise: bool, - colwise: bool, - shape_rep: str, - with_gemm_swizzled_scales: bool = False, - with_dbias: bool = False, - with_dact: bool = False, - with_act: bool = False, - activation: str = "none", - ): - if 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 - 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 - if not (rowwise or colwise): - raise ValueError("at least one of rowwise or colwise must be true") - self.ROWWISE = rowwise - self.COLWISE = colwise - if shape_rep not in SUPPORTED_SHAPE_REPS: - raise ValueError( - f"unsupported shape representation {shape_rep!r}; expected one of" - f" {SUPPORTED_SHAPE_REPS}" - ) - self.SHAPE_REP = shape_rep - # Only can view the grouped tensor as one tensor when the last dim is same across all tensors - self.IS_SINGLE_TENSOR = shape_rep in (SAME_BOTH_DIMS, VARYING_FIRST_DIM) - self.MAX_NORM_RCP = ( - FP8E4M3_MAX_NORM_RCP if fp8_dtype == "Float8E4M3" else FP8E5M2_MAX_NORM_RCP - ) - - self.WITH_GEMM_SWIZZLED_SCALES = with_gemm_swizzled_scales - if with_dbias and not self.IS_SINGLE_TENSOR: - # mxfp8::group_quantize raises for this. - raise ValueError("dbias is only supported for tensors with a common last dimension") - self.WITH_DBIAS = with_dbias - - if with_dact and with_act: - raise ValueError("with_dact and with_act cannot both be set") - if with_dact: - if activation not in SUPPORTED_DACTIVATIONS: - raise ValueError( - f"unknown activation {activation!r} for with_dact=True; expected one of" - f" {sorted(SUPPORTED_DACTIVATIONS)}" - ) - self.ACTIVATION = activation - elif with_act: - if activation not in SUPPORTED_ACTIVATIONS: - raise ValueError( - f"unknown activation {activation!r} for with_act=True; expected one of" - f" {sorted(SUPPORTED_ACTIVATIONS)}" - ) - self.ACTIVATION = activation - else: - if activation != "none": - raise ValueError("activation must be none when with_dact and with_act are False") - self.ACTIVATION = None - self.WITH_DACT = with_dact - self.WITH_ACT = with_act - - def __str__(self): - return ( - f"MXFP8GroupQuantizeConfig(dtype={self.DTYPE_STR}, fp8_dtype={self.FP8_DTYPE_STR}, " - f"rowwise={self.ROWWISE}, colwise={self.COLWISE}, shape_rep={self.SHAPE_REP}, " - f"swizzled={self.WITH_GEMM_SWIZZLED_SCALES}, with_dbias={self.WITH_DBIAS}, " - f"with_dact={self.WITH_DACT}, with_act={self.WITH_ACT}, " - f"activation={self.ACTIVATION})" - ) - - __repr__ = __str__ +from .group_quantize_mxfp8_common import ( + CUTEDSL_DEBUG_LOGGING, + THREADS_PER_WARP, + INPUT_SLOT, + ACT_INPUT_SLOT, + ROWWISE_OUTPUT_SLOT, + COLWISE_OUTPUT_SLOT, + SAME_BOTH_DIMS, + VARYING_FIRST_DIM, + VARYING_LAST_DIM, + VARYING_BOTH_DIMS, + MAX_SUPPORTED_TENSORS, + BYTES_PER_TENSORMAP, + INT64_PER_TENSORMAP, + ROWS_OFFSET, + COLS_OFFSET, + OFFSETS_OFFSET, + MXFP8GroupQuantizeConfig, +) class MXFP8GroupQuantizeKernel: @@ -1300,234 +1197,3 @@ def _process_job( dbias_col = tile_start_col + tidx if dbias_col < member_cols: mWorkspace[(dbias_row, dbias_col)] = partial_dbias - - -def compile_cutedsl_function_from_cfg(cfg: MXFP8GroupQuantizeConfig): - """Return the compiled CuTeDSL function object for the given grouped config.""" - # CUDA requires the group's first logical dim to be a multiple of 128 (and each - # tensor's rows likewise). VARYING_BOTH_DIMS is the exception: its logical shape is - # [1, total]. The last dim only needs the 16-byte TMA row alignment; a partial 32-element - # scale block at the end of a row is zero-filled by TMA, as in the CUDA kernel. - if cfg.SHAPE_REP == VARYING_BOTH_DIMS: - sym_M = cute.sym_int64() - else: - sym_M = cute.sym_int64(divisibility=128) - sym_N = cute.sym_int64(divisibility=SYM_N_DIVISIBILITY) - logical_shape = (sym_M, sym_N) - - out_dtype = cfg.FP8_DTYPE - scale_dtype = cutlass.Float8E8M0FNU - - out_col_fake = ( - cute.runtime.make_fake_compact_tensor( - out_dtype, - logical_shape, - stride_order=(1, 0), - memspace=cute.AddressSpace.gmem, - assumed_align=16, - ) - if cfg.COLWISE - else None - ) - - out_row_fake = ( - cute.runtime.make_fake_compact_tensor( - out_dtype, - logical_shape, - stride_order=(1, 0), - memspace=cute.AddressSpace.gmem, - assumed_align=16, - ) - if cfg.ROWWISE - else None - ) - - # The kernel only takes the base address of the scale buffers (per-tensor strides - # are derived from cols), so their fake shape is a flat 1D byte run. - # One compact view covers CUDA's descriptor and metadata arrays in their native order. - tensormaps_fake = ( - cute.runtime.make_fake_compact_tensor( - cutlass.Int64, - (TENSOR_MAP_STORAGE_INT64_COUNT,), - stride_order=(0,), - memspace=cute.AddressSpace.gmem, - assumed_align=128, - ) - if not cfg.IS_SINGLE_TENSOR - else None - ) - # The cast-noop flag is an always-present f32 pointer instead of an optional tensor, so - # that one compiled kernel serves both an absent and a present flag (noop_flag_is_set). - 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, - logical_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, - (cute.sym_int64(), cute.sym_int64()), - stride_order=(1, 0), - memspace=cute.AddressSpace.gmem, - assumed_align=4, - ) - if cfg.WITH_DBIAS - else None - ) - - first_dims_fake = ( - cute.runtime.make_fake_compact_tensor( - cutlass.Int64, - (cute.sym_int32(),), - stride_order=(0,), - memspace=cute.AddressSpace.gmem, - assumed_align=8, - ) - if cfg.SHAPE_REP in (VARYING_FIRST_DIM, VARYING_BOTH_DIMS) - else None - ) - - last_dims_fake = ( - cute.runtime.make_fake_compact_tensor( - cutlass.Int64, - (cute.sym_int32(),), - stride_order=(0,), - memspace=cute.AddressSpace.gmem, - assumed_align=8, - ) - if cfg.SHAPE_REP in (VARYING_LAST_DIM, VARYING_BOTH_DIMS) - else None - ) - - offsets_fake = ( - cute.runtime.make_fake_compact_tensor( - cutlass.Int64, - (cute.sym_int32(),), - stride_order=(0,), - memspace=cute.AddressSpace.gmem, - assumed_align=8, - ) - if cfg.SHAPE_REP != SAME_BOTH_DIMS - else None - ) - - sm_count = HardwareInfo().get_device_multiprocessor_count() - kernel_obj = MXFP8GroupQuantizeKernel(cfg, sm_count) - return cute.compile( - kernel_obj, - cute.runtime.make_fake_compact_tensor( # mX - cfg.DTYPE, - logical_shape, - stride_order=(1, 0), - memspace=cute.AddressSpace.gmem, - assumed_align=16, - ), - out_row_fake, # mO_row - out_col_fake, # mO_col - cute.runtime.make_fake_compact_tensor( # mS_row - scale_dtype, - (cute.sym_int64(),), - stride_order=(0,), - memspace=cute.AddressSpace.gmem, - assumed_align=4, - ), - cute.runtime.make_fake_compact_tensor( # mS_col - scale_dtype, - (cute.sym_int64(),), - stride_order=(0,), - memspace=cute.AddressSpace.gmem, - assumed_align=4, - ), - offsets_fake, # mOffsets - first_dims_fake, # mFirstDims - last_dims_fake, # mLastDims - tensormaps_fake, # mTensormaps - Int32(1), # num_tensors, supplied at runtime - noop_fake, # mNoop - act_input_fake, # mActInput - workspace_fake, # mWorkspace - cute.runtime.make_fake_stream(), - options="--enable-tvm-ffi", - ) - - -def get_mxfp8_group_quantization_function( - fn_name: str, - dtype: str, - fp8_dtype: str, - rowwise: bool, - colwise: bool, - shape_rep: str, - with_gemm_swizzled_scales: bool, - with_dbias: bool, - with_dact: bool, - with_act: bool, - activation: str, -) -> bool: - """Compile the grouped 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++ grouped 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 = MXFP8GroupQuantizeConfig( - dtype=dtype, - fp8_dtype=fp8_dtype, - rowwise=rowwise, - colwise=colwise, - shape_rep=shape_rep, - with_gemm_swizzled_scales=with_gemm_swizzled_scales, - with_dbias=with_dbias, - with_dact=with_dact, - with_act=with_act, - activation=activation, - ) - except ValueError as e: - logger.warning( - "CuTeDSL grouped MXFP8 backend does not support this config, " - "falling back to the CUDA C++ kernel: %s", - e, - ) - return False - - logger.debug("Compiling CuTeDSL grouped MXFP8 quantization kernel for %s", cfg) - compiled = compile_cutedsl_function_from_cfg(cfg) - # Register the native TVM-FFI function rather than its Python argument-parsing wrapper; - # see get_mxfp8_quantization_function for why. - 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 grouped MXFP8 kernel compilation & registration failed, falling back to the" - " CUDA C++ kernel: %s", - e, - ) - # Unconditionally fallback to CUDA path because we can't tell if this exception is - # transient or permanent. - return False diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py index b636d04edfb..efd0c4c4879 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py @@ -332,12 +332,14 @@ def quantize_bidimensional_mxfp8_g2r( ELEMENTS_X_PER_THREAD: cutlass.Constexpr[int], cache_policy: Int64, ): - """Quantize one (TILE_ROWS, TILE_COLS) GMEM tile in both directions. + """Quantize one (TILE_ROWS, TILE_COLS) GMEM tile in cfg's directions. TILE_ROWS = CTA_THREADS_Y * ELEMENTS_Y_PER_THREAD = 32, so every columnwise MXFP8 block is contained in the tile. Each warp owns ELEMENTS_Y_PER_THREAD consecutive rows, and each lane owns ELEMENTS_X_PER_THREAD adjacent columns. Load the tile once, then reuse its registers for both quantization passes. + cfg.ROWWISE / cfg.COLWISE select the passes; the tensors of a disabled + direction are unused and may be None. """ # Shape of the tile TILE_ROWS = CTA_THREADS_Y * ELEMENTS_Y_PER_THREAD @@ -359,8 +361,6 @@ def quantize_bidimensional_mxfp8_g2r( ) gX_thread = cute.composition(gX_tile, tv_layout)[tidx, None] - gO_row_thread = cute.composition(gO_row_tile, tv_layout)[tidx, None] - gO_col_thread = cute.composition(gO_col_tile, tv_layout)[tidx, None] # Group CTA threads by THREADS_X_PER_MX_BLOCK since they share the same rowwise scale for a MXFP8 block. # And assign stride 0 to that submode @@ -374,7 +374,6 @@ def quantize_bidimensional_mxfp8_g2r( 1, ), ) - gS_row_thread = cute.composition(gS_row_tile, tv_row_scale)[tidx, None] # Group CTA threads by THREADS_Y_PER_MX_BLOCK since they share the same columnwise scale for a MXFP8 block. # And assign stride 0 to that submode @@ -391,7 +390,6 @@ def quantize_bidimensional_mxfp8_g2r( # If CTA_THREADS_Y // THREADS_Y_PER_MX_BLOCK == 1, this will help compiler derive alignment # Otherwise coalesce will not do anything which is still fine tv_col_scale = cute.coalesce(tv_col_scale, target_profile=(1, 1)) - gS_col_thread = cute.composition(gS_col_tile, tv_col_scale)[tidx, None] # Input fragment contains all rows it handles; output fragment contains only one row and we write it back row by row rX = cute.make_rmem_tensor((ELEMENTS_X_PER_THREAD, ELEMENTS_Y_PER_THREAD), cfg.DTYPE) @@ -430,21 +428,57 @@ def quantize_bidimensional_mxfp8_g2r( load_atom, gX_thread[(None, row_idx),], rX[None, row_idx], cache_policy=cache_policy ) - _quantize_rowwise( - rX, - gO_row_thread, - gS_row_thread, - output_store_atom=store_atom, - row_scale_store_atom=row_scale_atom, - cache_policy=cache_policy, - scale_and_pack=scale_and_pack, - THREADS_X_PER_MX_BLOCK=THREADS_X_PER_MX_BLOCK, - ELEMENTS_Y_PER_THREAD=ELEMENTS_Y_PER_THREAD, - ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, - SCALE_DTYPE=SCALE_DTYPE, - FP8_DTYPE=cfg.FP8_DTYPE, - ) + if cutlass.const_expr(cfg.ROWWISE): + _quantize_rowwise( + rX, + cute.composition(gO_row_tile, tv_layout)[tidx, None], + cute.composition(gS_row_tile, tv_row_scale)[tidx, None], + output_store_atom=store_atom, + row_scale_store_atom=row_scale_atom, + cache_policy=cache_policy, + scale_and_pack=scale_and_pack, + THREADS_X_PER_MX_BLOCK=THREADS_X_PER_MX_BLOCK, + ELEMENTS_Y_PER_THREAD=ELEMENTS_Y_PER_THREAD, + ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, + SCALE_DTYPE=SCALE_DTYPE, + FP8_DTYPE=cfg.FP8_DTYPE, + ) + if cutlass.const_expr(cfg.COLWISE): + _quantize_colwise_tile( + rX, + cute.composition(gO_col_tile, tv_layout)[tidx, None], + cute.composition(gS_col_tile, tv_col_scale)[tidx, None], + store_atom=store_atom, + col_scale_atom=col_scale_atom, + cache_policy=cache_policy, + scale_and_pack=scale_and_pack, + cfg=cfg, + 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, + ) + +@cute.jit +def _quantize_colwise_tile( + rX: cute.Tensor, + gO_col_thread: cute.Tensor, + gS_col_thread: cute.Tensor, + store_atom: cute.CopyAtom, + col_scale_atom: cute.CopyAtom, + cache_policy: Int64, + scale_and_pack: cutlass.Constexpr[Callable[..., Uint32]], + cfg: cutlass.Constexpr[MXFP8QuantizeConfig], + 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]], +): + """Columnwise pass of quantize_bidimensional_mxfp8_g2r over the registers in rX.""" + tidx, _, _ = cute.arch.thread_idx() sAmaxs_col_scratch = _gather_partial_colwise_amaxes( rX, CTA_THREADS_Y=CTA_THREADS_Y, diff --git a/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8_cutedsl.cuh b/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8_cutedsl.cuh index 4722935a1af..f3d77377a8b 100644 --- a/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8_cutedsl.cuh +++ b/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8_cutedsl.cuh @@ -178,7 +178,7 @@ inline bool mxfp8_group_quantize_cutedsl(const MXFP8GroupQuantConfig &config, const size_t first_logical_dim = input_tensor->logical_shape.data[0]; const size_t last_logical_dim = input_tensor->logical_shape.data[1]; - // Match the symbolic divisibility checks in group_quantize_mxfp8.py. + // Match the symbolic divisibility checks in group_quantize_mxfp8_common.py. if (config.shape_rep != ShapeRepresentation::VARYING_BOTH_DIMS && first_logical_dim % 128 != 0) { maybe_warn_cutedsl_not_chosen( "the first logical dimension is not divisible by 128 for a non-varying both dimensions " From 59c32deb5d349a132cf729fc6dc7103c11d09bd5 Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Tue, 13 Oct 2026 14:23:22 -0700 Subject: [PATCH 43/58] fix Signed-off-by: Kaining Zhong --- .../cast/mxfp8/group_quantize_mxfp8_common.py | 18 ++---------- .../cast/mxfp8/group_quantize_mxfp8_g2r.py | 29 ++++++++++++++----- 2 files changed, 25 insertions(+), 22 deletions(-) diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8_common.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8_common.py index 1149d0b75ca..493ca566d7d 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8_common.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8_common.py @@ -191,24 +191,12 @@ def __init__(self, cfg: MXFP8GroupQuantizeConfig, sm_count: int) -> None: self.SINGLE_TENSOR_MODE = "flat_rowwise" self.single_tensor_entry = None self.g2r_kernel = None - if self.SINGLE_TENSOR_MODE is None and plain_cast and cfg.DTYPE is cutlass.BFloat16: + if self.SINGLE_TENSOR_MODE is None: # pylint: disable=import-outside-toplevel - from .quantize_mxfp8_common import MXFP8QuantizeConfig from .group_quantize_mxfp8_g2r import MXFP8GroupQuantizeG2RKernel - self.g2r_kernel = MXFP8GroupQuantizeG2RKernel( - cfg.SHAPE_REP, - MXFP8QuantizeConfig( - dtype=cfg.DTYPE_STR, - fp8_dtype=cfg.FP8_DTYPE_STR, - rowwise=cfg.ROWWISE, - colwise=cfg.COLWISE, - with_gemm_swizzled_scales=cfg.WITH_GEMM_SWIZZLED_SCALES, - with_amax=False, - activation="none", - ), - sm_count, - ) + if MXFP8GroupQuantizeG2RKernel.supports(cfg): + self.g2r_kernel = MXFP8GroupQuantizeG2RKernel(cfg, sm_count) if self.SINGLE_TENSOR_MODE is not None: # pylint: disable=import-outside-toplevel from .quantize_mxfp8_common import MXFP8QuantizeConfig, MXFP8QuantizeEntry diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8_g2r.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8_g2r.py index 70ccda6a7b0..b38c9ac9d9c 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8_g2r.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8_g2r.py @@ -32,6 +32,7 @@ ) from .quantize_mxfp8_g2r_bidimensional import quantize_bidimensional_mxfp8_g2r from .group_quantize_mxfp8_common import ( + MXFP8GroupQuantizeConfig, SAME_BOTH_DIMS, VARYING_FIRST_DIM, VARYING_LAST_DIM, @@ -56,16 +57,30 @@ class MXFP8GroupQuantizeG2RKernel: # 32768x7168). WORKERS_PER_SM = 256 - def __init__(self, shape_rep: str, cfg: MXFP8QuantizeConfig, sm_count: int) -> None: + def __init__(self, cfg: MXFP8GroupQuantizeConfig, sm_count: int) -> None: + assert self.supports(cfg), f"unsupported config {cfg}" self.cfg = cfg self.SM_COUNT = sm_count - self.ROWS_VARY = shape_rep not in (SAME_BOTH_DIMS, VARYING_LAST_DIM) - self.COLS_VARY = shape_rep not in (SAME_BOTH_DIMS, VARYING_FIRST_DIM) - self.SAME_BOTH_DIMS = shape_rep == SAME_BOTH_DIMS + self.ROWS_VARY = cfg.SHAPE_REP not in (SAME_BOTH_DIMS, VARYING_LAST_DIM) + self.COLS_VARY = cfg.SHAPE_REP not in (SAME_BOTH_DIMS, VARYING_FIRST_DIM) + self.SAME_BOTH_DIMS = cfg.SHAPE_REP == SAME_BOTH_DIMS + # The tile function is shared with the single-tensor bidimensional kernel and takes + # its config type. + self.tile_cfg = MXFP8QuantizeConfig( + dtype=cfg.DTYPE_STR, + fp8_dtype=cfg.FP8_DTYPE_STR, + rowwise=cfg.ROWWISE, + colwise=cfg.COLWISE, + with_gemm_swizzled_scales=cfg.WITH_GEMM_SWIZZLED_SCALES, + with_amax=False, + activation="none", + ) @staticmethod - def supports(cfg: MXFP8QuantizeConfig) -> bool: - return cfg.is_plain_cast() and cfg.DTYPE is cutlass.BFloat16 + def supports(cfg: MXFP8GroupQuantizeConfig) -> bool: + """BF16 plain casts: no dbias and no fused (d)activation.""" + plain_cast = not (cfg.WITH_DBIAS or cfg.WITH_ACT or cfg.WITH_DACT) + return plain_cast and cfg.DTYPE is cutlass.BFloat16 @cute.jit def __call__( @@ -200,7 +215,7 @@ def _quantize_tile(self, gX, gO_row, gS_row, gO_col, gS_col, coord, cache_policy gS_row_tile, gO_col_tile, gS_col_tile, - cfg, + self.tile_cfg, CTA_THREADS_Y=self._CTA_THREADS_Y, CTA_THREADS_X=self._CTA_THREADS_X, ELEMENTS_Y_PER_THREAD=self._ELEMENTS_Y_PER_THREAD, From 8f4c54d6b124fa863a3282d28496c709154191fd Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Tue, 13 Oct 2026 22:47:11 +0000 Subject: [PATCH 44/58] refactor Signed-off-by: Kaining Zhong --- .../mxfp8/quantize_mxfp8_g2r_bidimensional.py | 281 +++++++++------ .../mxfp8/quantize_mxfp8_g2r_rowwise_1lane.py | 329 +++++++++-------- .../cast/mxfp8/quantize_mxfp8_general.py | 335 +++++++++--------- 3 files changed, 532 insertions(+), 413 deletions(-) diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py index b636d04edfb..75ad5799cbf 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py @@ -276,7 +276,7 @@ def _write_colwise_scale( @cute.jit -def _quantize_colwise( +def _apply_colwise_scales( rX: cute.Tensor, rS_col_rcp: cute.Tensor, gO_col_thread: cute.Tensor, @@ -318,26 +318,100 @@ def _quantize_colwise( cute.copy(store_atom, rO_col, gO_col_thread[(None, row_idx),], cache_policy=cache_policy) +@cute.jit +def _quantize_colwise( + rX: cute.Tensor, + gO_col_thread: cute.Tensor, + gS_col_thread: cute.Tensor, + # Atoms + output_store_atom: cute.CopyAtom, + col_scale_store_atom: cute.CopyAtom, + # Miscellaneous + cache_policy: Int64, + scale_and_pack: cutlass.Constexpr[Callable[..., Uint32]], + # Constexprs + CTA_THREADS_Y: cutlass.Constexpr[int], + CTA_THREADS_X: cutlass.Constexpr[int], + ELEMENTS_Y_PER_THREAD: cutlass.Constexpr[int], + ELEMENTS_X_PER_THREAD: cutlass.Constexpr[int], + SCALE_DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], + DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], + FP8_DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], +): + """Reduce columnwise amaxes across the CTA, write the columnwise scales, and quantize rX with them.""" + tidx, _, _ = cute.arch.thread_idx() + + sAmaxs_col_scratch = _gather_partial_colwise_amaxes( + rX, + CTA_THREADS_Y=CTA_THREADS_Y, + CTA_THREADS_X=CTA_THREADS_X, + ELEMENTS_Y_PER_THREAD=ELEMENTS_Y_PER_THREAD, + ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, + DTYPE=DTYPE, + ) + + sS_col = _reduce_partial_colwise_amaxes( + sAmaxs_col_scratch, + CTA_THREADS_Y=CTA_THREADS_Y, + CTA_THREADS_X=CTA_THREADS_X, + ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, + DTYPE=DTYPE, + FP8_DTYPE=FP8_DTYPE, + ) + + # Copy the reduced columnwise scale reciprocals to RMEM. + # Threads that quantize different rows share the same columnwise scale + rS_col_rcp = cute.make_rmem_tensor((ELEMENTS_X_PER_THREAD,), DTYPE) + cute.autovec_copy(sS_col[None, tidx % CTA_THREADS_X], rS_col_rcp) + + _write_colwise_scale( + rS_col_rcp, + gS_col_thread, + col_scale_atom=col_scale_store_atom, + cache_policy=cache_policy, + CTA_THREADS_X=CTA_THREADS_X, + ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, + SCALE_DTYPE=SCALE_DTYPE, + ) + + _apply_colwise_scales( + rX, + rS_col_rcp, + gO_col_thread, + store_atom=output_store_atom, + cache_policy=cache_policy, + scale_and_pack=scale_and_pack, + ELEMENTS_Y_PER_THREAD=ELEMENTS_Y_PER_THREAD, + ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, + FP8_DTYPE=FP8_DTYPE, + ) + + @cute.jit def quantize_bidimensional_mxfp8_g2r( gX_tile: cute.Tensor, - gO_row_tile: cute.Tensor, - gS_row_tile: cute.Tensor, - gO_col_tile: cute.Tensor, - gS_col_tile: cute.Tensor, - cfg: cutlass.Constexpr[MXFP8QuantizeConfig], + gO_row_tile: Optional[cute.Tensor], + gS_row_tile: Optional[cute.Tensor], + gO_col_tile: Optional[cute.Tensor], + gS_col_tile: Optional[cute.Tensor], + ROWWISE: cutlass.Constexpr[bool], + COLWISE: cutlass.Constexpr[bool], + DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], + FP8_DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], + WITH_GEMM_SWIZZLED_SCALES: cutlass.Constexpr[bool], CTA_THREADS_Y: cutlass.Constexpr[int], CTA_THREADS_X: cutlass.Constexpr[int], ELEMENTS_Y_PER_THREAD: cutlass.Constexpr[int], ELEMENTS_X_PER_THREAD: cutlass.Constexpr[int], cache_policy: Int64, ): - """Quantize one (TILE_ROWS, TILE_COLS) GMEM tile in both directions. + """Quantize one (TILE_ROWS, TILE_COLS) GMEM tile in the directions enabled by ROWWISE / COLWISE. TILE_ROWS = CTA_THREADS_Y * ELEMENTS_Y_PER_THREAD = 32, so every columnwise MXFP8 block is contained in the tile. Each warp owns ELEMENTS_Y_PER_THREAD consecutive rows, and each lane owns ELEMENTS_X_PER_THREAD adjacent columns. - Load the tile once, then reuse its registers for both quantization passes. + Load the tile once, then reuse its registers for each enabled quantization pass. + The tensors of a disabled direction are unused and may be None. """ # Shape of the tile TILE_ROWS = CTA_THREADS_Y * ELEMENTS_Y_PER_THREAD @@ -359,48 +433,15 @@ def quantize_bidimensional_mxfp8_g2r( ) gX_thread = cute.composition(gX_tile, tv_layout)[tidx, None] - gO_row_thread = cute.composition(gO_row_tile, tv_layout)[tidx, None] - gO_col_thread = cute.composition(gO_col_tile, tv_layout)[tidx, None] - - # Group CTA threads by THREADS_X_PER_MX_BLOCK since they share the same rowwise scale for a MXFP8 block. - # And assign stride 0 to that submode - tv_row_scale = cute.make_layout( - ( - (THREADS_X_PER_MX_BLOCK, (CTA_THREADS_X // THREADS_X_PER_MX_BLOCK, CTA_THREADS_Y)), - ELEMENTS_Y_PER_THREAD, - ), - stride=( - (0, (TILE_ROWS, ELEMENTS_Y_PER_THREAD)), - 1, - ), - ) - gS_row_thread = cute.composition(gS_row_tile, tv_row_scale)[tidx, None] - - # Group CTA threads by THREADS_Y_PER_MX_BLOCK since they share the same columnwise scale for a MXFP8 block. - # And assign stride 0 to that submode - tv_col_scale = cute.make_layout( - ( - (CTA_THREADS_X, (THREADS_Y_PER_MX_BLOCK, CTA_THREADS_Y // THREADS_Y_PER_MX_BLOCK)), - ELEMENTS_X_PER_THREAD, - ), - stride=( - (ELEMENTS_X_PER_THREAD, (0, 1)), - 1, - ), - ) - # If CTA_THREADS_Y // THREADS_Y_PER_MX_BLOCK == 1, this will help compiler derive alignment - # Otherwise coalesce will not do anything which is still fine - tv_col_scale = cute.coalesce(tv_col_scale, target_profile=(1, 1)) - gS_col_thread = cute.composition(gS_col_tile, tv_col_scale)[tidx, None] # Input fragment contains all rows it handles; output fragment contains only one row and we write it back row by row - rX = cute.make_rmem_tensor((ELEMENTS_X_PER_THREAD, ELEMENTS_Y_PER_THREAD), cfg.DTYPE) + rX = cute.make_rmem_tensor((ELEMENTS_X_PER_THREAD, ELEMENTS_Y_PER_THREAD), DTYPE) load_atom = cute.make_copy_atom( cute.nvgpu.CopyG2ROp(), - cfg.DTYPE, + DTYPE, num_bits_per_copy=min( - ELEMENTS_X_PER_THREAD * cfg.DTYPE.width, + ELEMENTS_X_PER_THREAD * DTYPE.width, max_num_bits_per_copy_ld_global(), ), # TODO(kainingz): due to https://github.com/NVIDIA/cutlass/issues/3726 using @@ -409,20 +450,11 @@ def quantize_bidimensional_mxfp8_g2r( ) store_atom = cute.make_copy_atom( cute.nvgpu.CopyR2GOp(), - cfg.FP8_DTYPE, - num_bits_per_copy=ELEMENTS_X_PER_THREAD * cfg.FP8_DTYPE.width, + FP8_DTYPE, + num_bits_per_copy=ELEMENTS_X_PER_THREAD * FP8_DTYPE.width, ) - row_scale_atom = cute.make_copy_atom(cute.nvgpu.CopyR2GOp(), SCALE_DTYPE, num_bits_per_copy=8) - # GEMM-swizzled colwise scales are not contiguous (adjacent columns land 16 bytes apart), - # so they are stored a byte at a time, as the general kernel does. - col_scale_atom = cute.make_copy_atom( - cute.nvgpu.CopyR2GOp(), - SCALE_DTYPE, - num_bits_per_copy=8 if cfg.WITH_GEMM_SWIZZLED_SCALES else ELEMENTS_X_PER_THREAD * 8, - ) - - scale_and_pack = mul_bf16x2x2_cvt_bf16x4_to_fp8x4(cfg.FP8_DTYPE) + scale_and_pack = mul_bf16x2x2_cvt_bf16x4_to_fp8x4(FP8_DTYPE) # Copy all rows a thread sweeps into the RMEM fragment at once for row_idx in cutlass.range_constexpr(ELEMENTS_Y_PER_THREAD): @@ -430,65 +462,86 @@ def quantize_bidimensional_mxfp8_g2r( load_atom, gX_thread[(None, row_idx),], rX[None, row_idx], cache_policy=cache_policy ) - _quantize_rowwise( - rX, - gO_row_thread, - gS_row_thread, - output_store_atom=store_atom, - row_scale_store_atom=row_scale_atom, - cache_policy=cache_policy, - scale_and_pack=scale_and_pack, - THREADS_X_PER_MX_BLOCK=THREADS_X_PER_MX_BLOCK, - ELEMENTS_Y_PER_THREAD=ELEMENTS_Y_PER_THREAD, - ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, - SCALE_DTYPE=SCALE_DTYPE, - FP8_DTYPE=cfg.FP8_DTYPE, - ) - - sAmaxs_col_scratch = _gather_partial_colwise_amaxes( - rX, - CTA_THREADS_Y=CTA_THREADS_Y, - CTA_THREADS_X=CTA_THREADS_X, - ELEMENTS_Y_PER_THREAD=ELEMENTS_Y_PER_THREAD, - ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, - DTYPE=cfg.DTYPE, - ) + if cutlass.const_expr(ROWWISE): + gO_row_thread = cute.composition(gO_row_tile, tv_layout)[tidx, None] + + # Group CTA threads by THREADS_X_PER_MX_BLOCK since they share the same rowwise scale for a MXFP8 block. + # And assign stride 0 to that submode + tv_row_scale = cute.make_layout( + ( + (THREADS_X_PER_MX_BLOCK, (CTA_THREADS_X // THREADS_X_PER_MX_BLOCK, CTA_THREADS_Y)), + ELEMENTS_Y_PER_THREAD, + ), + stride=( + (0, (TILE_ROWS, ELEMENTS_Y_PER_THREAD)), + 1, + ), + ) + gS_row_thread = cute.composition(gS_row_tile, tv_row_scale)[tidx, None] - sS_col = _reduce_partial_colwise_amaxes( - sAmaxs_col_scratch, - CTA_THREADS_Y=CTA_THREADS_Y, - CTA_THREADS_X=CTA_THREADS_X, - ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, - DTYPE=cfg.DTYPE, - FP8_DTYPE=cfg.FP8_DTYPE, - ) + row_scale_atom = cute.make_copy_atom( + cute.nvgpu.CopyR2GOp(), SCALE_DTYPE, num_bits_per_copy=8 + ) - # Copy the reduced columnwise scale reciprocals to RMEM. - # Threads that quantize different rows share the same columnwise scale - rS_col_rcp = cute.make_rmem_tensor((ELEMENTS_X_PER_THREAD,), cfg.DTYPE) - cute.autovec_copy(sS_col[None, tidx % CTA_THREADS_X], rS_col_rcp) + _quantize_rowwise( + rX, + gO_row_thread, + gS_row_thread, + output_store_atom=store_atom, + row_scale_store_atom=row_scale_atom, + cache_policy=cache_policy, + scale_and_pack=scale_and_pack, + THREADS_X_PER_MX_BLOCK=THREADS_X_PER_MX_BLOCK, + ELEMENTS_Y_PER_THREAD=ELEMENTS_Y_PER_THREAD, + ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, + SCALE_DTYPE=SCALE_DTYPE, + FP8_DTYPE=FP8_DTYPE, + ) - _write_colwise_scale( - rS_col_rcp, - gS_col_thread, - col_scale_atom=col_scale_atom, - cache_policy=cache_policy, - CTA_THREADS_X=CTA_THREADS_X, - ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, - SCALE_DTYPE=SCALE_DTYPE, - ) + if cutlass.const_expr(COLWISE): + gO_col_thread = cute.composition(gO_col_tile, tv_layout)[tidx, None] + + # Group CTA threads by THREADS_Y_PER_MX_BLOCK since they share the same columnwise scale for a MXFP8 block. + # And assign stride 0 to that submode + tv_col_scale = cute.make_layout( + ( + (CTA_THREADS_X, (THREADS_Y_PER_MX_BLOCK, CTA_THREADS_Y // THREADS_Y_PER_MX_BLOCK)), + ELEMENTS_X_PER_THREAD, + ), + stride=( + (ELEMENTS_X_PER_THREAD, (0, 1)), + 1, + ), + ) + # If CTA_THREADS_Y // THREADS_Y_PER_MX_BLOCK == 1, this will help compiler derive alignment + # Otherwise coalesce will not do anything which is still fine + tv_col_scale = cute.coalesce(tv_col_scale, target_profile=(1, 1)) + gS_col_thread = cute.composition(gS_col_tile, tv_col_scale)[tidx, None] + + # GEMM-swizzled colwise scales are not contiguous (adjacent columns land 16 bytes apart), + # so they are stored a byte at a time, as the general kernel does. + col_scale_atom = cute.make_copy_atom( + cute.nvgpu.CopyR2GOp(), + SCALE_DTYPE, + num_bits_per_copy=8 if WITH_GEMM_SWIZZLED_SCALES else ELEMENTS_X_PER_THREAD * 8, + ) - _quantize_colwise( - rX, - rS_col_rcp, - gO_col_thread, - store_atom=store_atom, - cache_policy=cache_policy, - scale_and_pack=scale_and_pack, - ELEMENTS_Y_PER_THREAD=ELEMENTS_Y_PER_THREAD, - ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, - FP8_DTYPE=cfg.FP8_DTYPE, - ) + _quantize_colwise( + rX, + gO_col_thread, + gS_col_thread, + output_store_atom=store_atom, + col_scale_store_atom=col_scale_atom, + cache_policy=cache_policy, + scale_and_pack=scale_and_pack, + CTA_THREADS_Y=CTA_THREADS_Y, + CTA_THREADS_X=CTA_THREADS_X, + ELEMENTS_Y_PER_THREAD=ELEMENTS_Y_PER_THREAD, + ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, + SCALE_DTYPE=SCALE_DTYPE, + DTYPE=DTYPE, + FP8_DTYPE=FP8_DTYPE, + ) class MXFP8QuantizeRegisterBidimensionalKernel(MXFP8QuantizeKernelBase): @@ -742,7 +795,11 @@ def _kernel_main( gS_row_tile, gO_col_tile, gS_col_tile, - self.cfg, + ROWWISE=self.cfg.ROWWISE, + COLWISE=self.cfg.COLWISE, + DTYPE=self.cfg.DTYPE, + FP8_DTYPE=self.cfg.FP8_DTYPE, + WITH_GEMM_SWIZZLED_SCALES=self.cfg.WITH_GEMM_SWIZZLED_SCALES, CTA_THREADS_Y=self._CTA_THREADS_Y, CTA_THREADS_X=self._CTA_THREADS_X, ELEMENTS_Y_PER_THREAD=self._ELEMENTS_Y_PER_THREAD, diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_1lane.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_1lane.py index fd76e039fa6..1c005c22fdb 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_1lane.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_1lane.py @@ -39,6 +39,172 @@ ) +@cute.jit +def _flush_scales_to_gmem( + sScale: cute.Tensor, + gS_tile: cute.Tensor, + cta_row: Int32, + cta_col: Int32, + M: Int32, + padded_cols: Int32, + CTA_Y: cutlass.Constexpr[int], + CTA_X: cutlass.Constexpr[int], + width: cutlass.Constexpr[int], +): + """Flush the staged (CTA_Y, CTA_X) scale tile to gmem with vectorized stores.""" + tidx, _, _ = cute.arch.thread_idx() + # Previously each threads has 1 byte, but now we are doing vectorized store, + # which means only a subset of threads will need to issue the store while other threads are not used. + active_threads = CTA_X // width + _, tv_flush = cute.make_layout_tv( + thr_layout=cute.make_layout((CTA_Y, active_threads), stride=(active_threads, 1)), + val_layout=cute.make_layout((1, width), stride=(width, 1)), + ) + # We only need to use a subset of threads with shape (CTA_Y, active_threads) to write + # so if the thread is outside of this subset, it will remain inactive + if tidx < CTA_Y * active_threads: + # Absolute position of the scale vector to write in the GMEM buffer + thread_y = cta_row * CTA_Y + tidx // active_threads + thread_x = cta_col * CTA_X + (tidx % active_threads) * width + if thread_y < M and thread_x < padded_cols: + cute.autovec_copy( + cute.composition(sScale, tv_flush)[tidx, None], + cute.composition(gS_tile, tv_flush)[tidx, None], + ) + + +@cute.jit +def quantize_rowwise_mxfp8_g2r_1lane( + gX_tile: cute.Tensor, + gO_tile: cute.Tensor, + gS_tile: cute.Tensor, + M: Int32, + N: Int32, + padded_cols: Int32, + cta_row: Int32, + cta_col: Int32, + DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], + FP8_DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], + MAX_NORM_RCP: cutlass.Constexpr[float], + WITH_GEMM_SWIZZLED_SCALES: cutlass.Constexpr[bool], + TILE_ROWS: cutlass.Constexpr[int], + TILE_COLS: cutlass.Constexpr[int], + STASH_SCALE_TO_SMEM: cutlass.Constexpr[bool], +): + """Quantize one (TILE_ROWS, TILE_COLS) GMEM tile rowwise with TILE_ROWS * TILE_COLS / 32 threads. + + Each thread handles one 32-element MXFP8 chunk (= one scale block). (cta_row, cta_col) is the + tile coordinate, used with M and N for bounds checks; padded_cols is the column count of the + (padded) scale tensor that gS_tile is a tile of. + """ + tidx, _, _ = cute.arch.thread_idx() + + # The CTA threads are grouped as (CTA_Y, CTA_X), so they cover a (TILE_ROWS, TILE_COLS) + # input tile and the matching (CTA_Y, CTA_X) scale tile. + CTA_Y = TILE_ROWS + CTA_X = TILE_COLS // MXFP8_BLOCK_SCALING_SIZE + _, tv_layout = cute.make_layout_tv( + thr_layout=cute.make_layout((CTA_Y, CTA_X), stride=(CTA_X, 1)), + val_layout=cute.make_layout( + (1, MXFP8_BLOCK_SCALING_SIZE), stride=(MXFP8_BLOCK_SCALING_SIZE, 1) + ), + ) + _, tv_layout_scale = cute.make_layout_tv( + thr_layout=cute.make_layout((CTA_Y, CTA_X), stride=(CTA_X, 1)), + val_layout=cute.make_layout((1, 1), stride=(1, 1)), + ) + + gX_thread = cute.composition(gX_tile, tv_layout)[tidx, None] + gO_thread = cute.composition(gO_tile, tv_layout)[tidx, None] + gS_thread = cute.composition(gS_tile, tv_layout_scale)[tidx, None] + + rX_thread = cute.make_rmem_tensor( + cute.make_layout((1, MXFP8_BLOCK_SCALING_SIZE), stride=(MXFP8_BLOCK_SCALING_SIZE, 1)), + dtype=DTYPE, + ) + abs_max_x2 = abs_max_x2_f16 if DTYPE is cutlass.Float16 else abs_max_x2_bf16 + x2_lo_to_f32 = x2_lo_to_f32_f16 if DTYPE is cutlass.Float16 else x2_lo_to_f32_bf16 + x2_hi_to_f32 = x2_hi_to_f32_f16 if DTYPE is cutlass.Float16 else x2_hi_to_f32_bf16 + mul_cvt4 = mul_f32x2_cvt_packed16x4_to_fp8x4(DTYPE, FP8_DTYPE) + rX_i32 = cute.make_tensor( + cute.recast_ptr(rX_thread.iterator, dtype=Int32), + cute.make_layout((MXFP8_BLOCK_SCALING_SIZE // 2,), stride=(1,)), + ) + rO_thread = cute.make_rmem_tensor( + cute.make_layout((1, MXFP8_BLOCK_SCALING_SIZE), stride=(MXFP8_BLOCK_SCALING_SIZE, 1)), + dtype=FP8_DTYPE, + ) + rO_u32 = cute.make_tensor( + cute.recast_ptr(rO_thread.iterator, dtype=Uint32), + cute.make_layout( + (MXFP8_BLOCK_SCALING_SIZE // 4,), stride=(1,) + ), # Unit is Uint32, divide by 4 here + ) + + sS_thread = None + if cutlass.const_expr(STASH_SCALE_TO_SMEM): + + @cute.struct + class SharedStorage: + buf: cute.struct.Align[cute.struct.MemRange[Float8E8M0FNU, CTA_Y * CTA_X], 16] + + storage = cutlass.memory.SmemAllocator().allocate(SharedStorage) + sScale = storage.buf.get_tensor(cute.make_layout((CTA_Y, CTA_X), stride=(CTA_X, 1))) + # sScale is (CTA_Y, CTA_X):(CTA_X, 1), which is the same layout as tv_layout_scale + # so sS_thread is really just a 1-element e8m0 buffer for this thread's scale. + sS_thread = cute.composition(sScale, tv_layout_scale)[tidx, None] + # Zero first so padding columns (cols past N/32 in the padded scale + # matrix) flush as 0 and we never read uninitialized smem. + # Raw 0x00: e8m0 has no zero value to convert from. + sS_thread[0] = Uint8(0).bitcast(Float8E8M0FNU) + cute.arch.sync_threads() + + row = cta_row * TILE_ROWS + tidx // CTA_X + col = cta_col * TILE_COLS + (tidx % CTA_X) * MXFP8_BLOCK_SCALING_SIZE + if row < M and col < N: + cute.autovec_copy(gX_thread, rX_thread) + amax_2x = rX_i32[0] + # Skip the first iteration because we assigned rX_i32[0] to amax_2x already + for i in cutlass.range_constexpr(1, MXFP8_BLOCK_SCALING_SIZE // 2): + amax_2x = abs_max_x2(amax_2x, rX_i32[i]) + amax = cute.arch.fmax(fabs_f32(x2_lo_to_f32(amax_2x)), fabs_f32(x2_hi_to_f32(amax_2x))) + + biased_exp = cvt_f32_to_fp8e8m0fnu(amax * MAX_NORM_RCP) + if cutlass.const_expr(STASH_SCALE_TO_SMEM): + sS_thread[0] = biased_exp + else: + gS_thread[0] = biased_exp + + # Rescale + FP8 cast, 4 elements per fused mul_cvt (one uint32 out), + # then a vectorized store. Mirrors CUDA's _use_cvt_4x path. + inv_scale = exp2f_rcp(biased_exp) + scale_2x = pack_f32x2(inv_scale, inv_scale) + for i in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE // 4): + rO_u32[i] = mul_cvt4(rX_i32[2 * i], rX_i32[2 * i + 1], scale_2x) + cute.autovec_copy(rO_thread, gO_thread) + + # Cooperative wide flush of the staged scales where padding columns flush as 0. + if cutlass.const_expr(STASH_SCALE_TO_SMEM): + cute.arch.sync_threads() + if cutlass.const_expr(WITH_GEMM_SWIZZLED_SCALES): + # Swizzled rowwise scale layout is ((32, 4, num_tiles_M), (4, num_tiles_SC)):((16, 4, num_tiles_SC * 512), (1, 512)) + # so we have at most 4 elements continuous which is our vectorized store width + _flush_scales_to_gmem( + sScale, gS_tile, cta_row, cta_col, M, padded_cols, CTA_Y, CTA_X, 4 + ) + elif padded_cols % 16 == 0: + # If columns is divisible by 16, use 16 bytes as the vectorized store width + _flush_scales_to_gmem( + sScale, gS_tile, cta_row, cta_col, M, padded_cols, CTA_Y, CTA_X, 16 + ) + else: + # Otherwise use 4 bytes as the vectorized store width. + # Note our fake tensor requires 4 divisibility so this is enforced as long as you can get here + _flush_scales_to_gmem( + sScale, gS_tile, cta_row, cta_col, M, padded_cols, CTA_Y, CTA_X, 4 + ) + + class MXFP8QuantizeG2RRowwise1LaneKernel(MXFP8QuantizeKernelBase): """Specialized cast-only ROWWISE-only MXFP8 kernel. @@ -115,7 +281,6 @@ def __call__( mX, mO_row, mS_row, - mX.element_type, ).launch(grid=grid, block=block, stream=stream) @cute.kernel @@ -124,11 +289,10 @@ def kernel( mX: cute.Tensor, mO_row: cute.Tensor, mS_row: cute.Tensor, - DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], ): """Device entry. Like the CUDA specialized cast kernels it does not consume the cast_noop flag: the dispatcher routes calls with one to the general kernel.""" - self._kernel_main(mX, mO_row, mS_row, DTYPE) + self._kernel_main(mX, mO_row, mS_row) @cute.jit def _kernel_main( @@ -136,151 +300,34 @@ def _kernel_main( mX: cute.Tensor, mO_row: cute.Tensor, mS_row: cute.Tensor, - DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], ): """Device entry for the specialized rowwise-only cast kernel (vectorized global loads/stores, no TMA).""" - tidx, _, _ = cute.arch.thread_idx() bidx, bidy, _ = cute.arch.block_idx() M = mX.shape[0] N = mX.shape[1] - # Each thread handles one 32-element MXFP8 chunk (= one scale block). - # The 128 threads in the CTA are grouped as (4, 32), so they cover a - # (4, 1024) input tile and the matching (4, 32) scale tile. - CTA_Y = self._TILE_ROWS - CTA_X = self._TILE_COLS // MXFP8_BLOCK_SCALING_SIZE - tiler, tv_layout = cute.make_layout_tv( - thr_layout=cute.make_layout((CTA_Y, CTA_X), stride=(CTA_X, 1)), - val_layout=cute.make_layout( - (1, MXFP8_BLOCK_SCALING_SIZE), stride=(MXFP8_BLOCK_SCALING_SIZE, 1) - ), - ) - tiler_scale, tv_layout_scale = cute.make_layout_tv( - thr_layout=cute.make_layout((CTA_Y, CTA_X), stride=(CTA_X, 1)), - val_layout=cute.make_layout((1, 1), stride=(1, 1)), - ) + TILER = (self._TILE_ROWS, self._TILE_COLS) + SCALE_TILER = (self._TILE_ROWS, self._TILE_COLS // MXFP8_BLOCK_SCALING_SIZE) + gX_tile = cute.local_tile(mX, TILER, (bidy, bidx)) + gO_tile = cute.local_tile(mO_row, TILER, (bidy, bidx)) + gS_tile = cute.local_tile(mS_row, SCALE_TILER, (bidy, bidx)) - # Select the tile that belongs to this CTA, then the fragment per thread. - mX_tile = cute.local_tile(mX, tiler, (bidy, bidx)) - mO_tile = cute.local_tile(mO_row, tiler, (bidy, bidx)) - mS_tile = cute.local_tile(mS_row, tiler_scale, (bidy, bidx)) - mX_thread = cute.composition(mX_tile, tv_layout)[tidx, None] - mO_thread = cute.composition(mO_tile, tv_layout)[tidx, None] - mS_thread = cute.composition(mS_tile, tv_layout_scale)[tidx, None] - - rX_thread = cute.make_rmem_tensor( - cute.make_layout((1, MXFP8_BLOCK_SCALING_SIZE), stride=(MXFP8_BLOCK_SCALING_SIZE, 1)), - dtype=DTYPE, - ) - abs_max_x2 = abs_max_x2_f16 if DTYPE is cutlass.Float16 else abs_max_x2_bf16 - x2_lo_to_f32 = x2_lo_to_f32_f16 if DTYPE is cutlass.Float16 else x2_lo_to_f32_bf16 - x2_hi_to_f32 = x2_hi_to_f32_f16 if DTYPE is cutlass.Float16 else x2_hi_to_f32_bf16 - mul_cvt4 = mul_f32x2_cvt_packed16x4_to_fp8x4(DTYPE, self.cfg.FP8_DTYPE) - rX_i32 = cute.make_tensor( - cute.recast_ptr(rX_thread.iterator, dtype=Int32), - cute.make_layout((MXFP8_BLOCK_SCALING_SIZE // 2,), stride=(1,)), - ) - rO_thread = cute.make_rmem_tensor( - cute.make_layout((1, MXFP8_BLOCK_SCALING_SIZE), stride=(MXFP8_BLOCK_SCALING_SIZE, 1)), - dtype=self.cfg.FP8_DTYPE, - ) - rO_u32 = cute.make_tensor( - cute.recast_ptr(rO_thread.iterator, dtype=Uint32), - cute.make_layout( - (MXFP8_BLOCK_SCALING_SIZE // 4,), stride=(1,) - ), # Unit is Uint32, divide by 4 here - ) - - sS_thread = None - if cutlass.const_expr(self._STASH_SCALE_TO_SMEM): - - @cute.struct - class SharedStorage: - buf: cute.struct.Align[cute.struct.MemRange[Float8E8M0FNU, CTA_Y * CTA_X], 16] - - storage = cutlass.memory.SmemAllocator().allocate(SharedStorage) - sScale = storage.buf.get_tensor(cute.make_layout((CTA_Y, CTA_X), stride=(CTA_X, 1))) - # sScale is (CTA_Y, CTA_X):(CTA_X, 1), which is the same layout as tv_layout_scale - # so sS_thread is really just a 1-element e8m0 buffer for this thread's scale. - sS_thread = cute.composition(sScale, tv_layout_scale)[tidx, None] - # Zero first so padding columns (cols past N/32 in the padded scale - # matrix) flush as 0 and we never read uninitialized smem. - # Raw 0x00: e8m0 has no zero value to convert from. - sS_thread[0] = Uint8(0).bitcast(Float8E8M0FNU) - cute.arch.sync_threads() - - row = bidy * self._TILE_ROWS + tidx // CTA_X - col = bidx * self._TILE_COLS + (tidx % CTA_X) * MXFP8_BLOCK_SCALING_SIZE - if row < M and col < N: - cute.autovec_copy(mX_thread, rX_thread) - amax_2x = rX_i32[0] - # Skip the first iteration because we assigned rX_i32[0] to amax_2x already - for i in cutlass.range_constexpr(1, MXFP8_BLOCK_SCALING_SIZE // 2): - amax_2x = abs_max_x2(amax_2x, rX_i32[i]) - amax = cute.arch.fmax(fabs_f32(x2_lo_to_f32(amax_2x)), fabs_f32(x2_hi_to_f32(amax_2x))) - - biased_exp = cvt_f32_to_fp8e8m0fnu(amax * self.cfg.MAX_NORM_RCP) - if cutlass.const_expr(self._STASH_SCALE_TO_SMEM): - sS_thread[0] = biased_exp - else: - mS_thread[0] = biased_exp - - # Rescale + FP8 cast, 4 elements per fused mul_cvt (one uint32 out), - # then a vectorized store. Mirrors CUDA's _use_cvt_4x path. - inv_scale = exp2f_rcp(biased_exp) - scale_2x = pack_f32x2(inv_scale, inv_scale) - for i in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE // 4): - rO_u32[i] = mul_cvt4(rX_i32[2 * i], rX_i32[2 * i + 1], scale_2x) - cute.autovec_copy(rO_thread, mO_thread) - - # Cooperative wide flush of the staged scales where padding columns flush as 0. - if cutlass.const_expr(self._STASH_SCALE_TO_SMEM): - cute.arch.sync_threads() + quantize_rowwise_mxfp8_g2r_1lane( + gX_tile, + gO_tile, + gS_tile, + M, + N, # Use cute.size instead of .shape[1] because under the swizzled layout mode 1 is # a nested tuple ((4, num_tiles_SC)), not a plain int. - padded_cols = cute.size(mS_row, mode=[1]) - if cutlass.const_expr(self.cfg.WITH_GEMM_SWIZZLED_SCALES): - # Swizzled rowwise scale layout is ((32, 4, num_tiles_M), (4, num_tiles_SC)):((16, 4, num_tiles_SC * 512), (1, 512)) - # so we have at most 4 elements continuous which is our vectorized store width - self._flush_scales_to_gmem(sScale, mS_tile, tidx, bidx, bidy, M, padded_cols, 4) - elif padded_cols % 16 == 0: - # If columns is divisible by 16, use 16 bytes as the vectorized store width - self._flush_scales_to_gmem(sScale, mS_tile, tidx, bidx, bidy, M, padded_cols, 16) - else: - # Otherwise use 4 bytes as the vectorized store width. - # Note our fake tensor requires 4 divisibility so this is enforced as long as you can get here - self._flush_scales_to_gmem(sScale, mS_tile, tidx, bidx, bidy, M, padded_cols, 4) - - @cute.jit - def _flush_scales_to_gmem( - self, - sScale: cute.Tensor, - mS_tile: cute.Tensor, - tidx: Int32, - bidx: Int32, - bidy: Int32, - M: Int32, - padded_cols: Int32, - width: cutlass.Constexpr[int], - ): - """Flush the staged (CTA_Y, CTA_X) scale tile to gmem with vectorized stores.""" - CTA_Y = self._TILE_ROWS - CTA_X = self._TILE_COLS // MXFP8_BLOCK_SCALING_SIZE - # Previously each threads has 1 byte, but now we are doing vectorized store, - # which means only a subset of threads will need to issue the store while other threads are not used. - active_threads = CTA_X // width - _, tv_flush = cute.make_layout_tv( - thr_layout=cute.make_layout((CTA_Y, active_threads), stride=(active_threads, 1)), - val_layout=cute.make_layout((1, width), stride=(width, 1)), + cute.size(mS_row, mode=[1]), + bidy, + bidx, + DTYPE=self.cfg.DTYPE, + FP8_DTYPE=self.cfg.FP8_DTYPE, + MAX_NORM_RCP=self.cfg.MAX_NORM_RCP, + WITH_GEMM_SWIZZLED_SCALES=self.cfg.WITH_GEMM_SWIZZLED_SCALES, + TILE_ROWS=self._TILE_ROWS, + TILE_COLS=self._TILE_COLS, + STASH_SCALE_TO_SMEM=self._STASH_SCALE_TO_SMEM, ) - # We only need to use a subset of threads with shape (CTA_Y, active_threads) to write - # so if the thread is outside of this subset, it will remain inactive - if tidx < CTA_Y * active_threads: - # Absolute position of the scale vector to write in the GMEM buffer - thread_y = bidy * CTA_Y + tidx // active_threads - thread_x = bidx * CTA_X + (tidx % active_threads) * width - if thread_y < M and thread_x < padded_cols: - cute.autovec_copy( - cute.composition(sScale, tv_flush)[tidx, None], - cute.composition(mS_tile, tv_flush)[tidx, None], - ) diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_general.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_general.py index aca02a253b1..d74d4ca824b 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_general.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_general.py @@ -91,9 +91,6 @@ def quantize_rowwise_mxfp8( FP8_DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], TILE_X: cutlass.Constexpr[int], TILE_Y: cutlass.Constexpr[int], - WAVES: cutlass.Constexpr[int], - THREADS_PER_BANK: cutlass.Constexpr[int], - PACK_SIZE: cutlass.Constexpr[int], SKIP_INPUT_MASKING: cutlass.Constexpr[bool] = False, SKIP_SCALE_BOUNDS: cutlass.Constexpr[bool] = False, WITH_ACT: cutlass.Constexpr[bool] = False, @@ -104,6 +101,14 @@ def quantize_rowwise_mxfp8( """Quantize one SMEM tile rowwise to MXFP8 (per-row 32-elt block scales); returns the tile amax.""" tidx, _, _ = cute.arch.thread_idx() + # Vectorised SMEM access constants for bank-conflict avoidance. The loads and the fused 4-wide + # mul_cvt below assume PACK_SIZE == 4, so these are not tunable. + PACK_SIZE = 4 # Elements per vector load + # Each thread reads 8 waves with each wave reads 4 packed bf16, so it reads a whole MXFP8 block in total + WAVES = MXFP8_BLOCK_SCALING_SIZE // PACK_SIZE + TOTAL_BANKS_WIDTH = 32 * 4 # 32 banks × 4 bytes, in bytes + THREADS_PER_BANK = TOTAL_BANKS_WIDTH // MXFP8_BLOCK_SCALING_SIZE # 4 threads per bank + CTA_THREADS_Y = TILE_Y # threads per column (rows per tile) CTA_THREADS_X = TILE_X // MXFP8_BLOCK_SCALING_SIZE # threads per row (chunks per row) @@ -534,18 +539,125 @@ def dbias_reduction_colwise( return dbias_partial +@cute.jit +def quantize_mxfp8_tile( + sX_tile: cute.Tensor, # (TILE_Y, TILE_X) bf16/fp16 smem view, post-TMA + sActInput_tile: Optional[cute.Tensor], # (TILE_Y, TILE_X) act-input smem tile (dact only) + sO_row_tile: Optional[cute.Tensor], # (TILE_Y, TILE_X) fp8 smem view (rowwise FP8 output) + mS_row_stage: Optional[cute.Tensor], # rowwise scale tensor (1D swizzled, or 2D linear) + sO_col_tile: Optional[cute.Tensor], # (TILE_Y, TILE_X) fp8 smem view (colwise FP8 output) + mS_col_stage: Optional[cute.Tensor], # colwise scale tensor (1D swizzled, or 2D linear) + rowwise_dbias_acc: Optional[cute.Tensor], # rmem Float32[32] (DBIAS_REDUCTION_IN_ROWWISE only) + tile_row_start: Int32, # global row index of this tile's row 0 + tile_col_start: Int32, # global col index of this tile's col 0 + M: Int32, + N: Int32, + ROWWISE: cutlass.Constexpr[bool], + COLWISE: cutlass.Constexpr[bool], + WITH_AMAX: cutlass.Constexpr[bool], + WITH_ACT: cutlass.Constexpr[bool], + WITH_DACT: cutlass.Constexpr[bool], + ACTIVATION: cutlass.Constexpr[str | None], + MAX_NORM_RCP: cutlass.Constexpr[float], + DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], + FP8_DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], + WITH_GEMM_SWIZZLED_SCALES: cutlass.Constexpr[bool], + TILE_X: cutlass.Constexpr[int], + TILE_Y: cutlass.Constexpr[int], + SKIP_INPUT_MASKING: cutlass.Constexpr[bool], + SKIP_SCALE_BOUNDS: cutlass.Constexpr[bool], + CACHE_ACTIVATION: cutlass.Constexpr[bool], + DBIAS_REDUCTION_IN_COLWISE: cutlass.Constexpr[bool], + COLWISE_DBIAS_REDUCTION_ONLY: cutlass.Constexpr[bool], + DBIAS_REDUCTION_IN_ROWWISE: cutlass.Constexpr[bool], +): + """Quantize one SMEM tile in the directions enabled by ROWWISE / COLWISE; returns (amax, dbias). + + amax is this thread's amax over the tile (None unless WITH_AMAX). dbias is this thread's + column partial sum (None unless DBIAS_REDUCTION_IN_COLWISE or COLWISE_DBIAS_REDUCTION_ONLY); + with DBIAS_REDUCTION_IN_ROWWISE the partial sums are accumulated into rowwise_dbias_acc instead. + """ + amax = None + dbias = None + if cutlass.const_expr(COLWISE): + amax_c, dbias_c = quantize_colwise_mxfp8( + sX_tile, + sActInput_tile, + sO_col_tile, + mS_col_stage, + MAX_NORM_RCP, + tile_row_start, + tile_col_start, + M, + N, + ACTIVATION=ACTIVATION, + DTYPE=DTYPE, + FP8_DTYPE=FP8_DTYPE, + SWIZZLE=WITH_GEMM_SWIZZLED_SCALES, + TILE_X=TILE_X, + TILE_Y=TILE_Y, + SKIP_INPUT_MASKING=SKIP_INPUT_MASKING, + SKIP_SCALE_BOUNDS=SKIP_SCALE_BOUNDS, + WITH_ACT=WITH_ACT, + WITH_DACT=WITH_DACT, + WITH_DBIAS=DBIAS_REDUCTION_IN_COLWISE, + CACHE_ACTIVATION=CACHE_ACTIVATION, + ) + # Either direction's amax already spans every element of the tile, so favor colwise + # and skip the redundant comparison in the rowwise pass. + if cutlass.const_expr(WITH_AMAX): + amax = amax_c + if cutlass.const_expr(DBIAS_REDUCTION_IN_COLWISE): + dbias = dbias_c + # If we don't quantize columnwise but we fuse dbias, do a dbias reduction only pass in the columnwise direction without quantization + if cutlass.const_expr(COLWISE_DBIAS_REDUCTION_ONLY): + dbias = dbias_reduction_colwise( + sX_tile, + sActInput_tile, + ACTIVATION=ACTIVATION, + DTYPE=DTYPE, + TILE_X=TILE_X, + WITH_DACT=WITH_DACT, + CACHE_ACTIVATION=CACHE_ACTIVATION, + ) + # If we cache the activation in shared memory, we need to ensure that all threads have finished writing to the shared memory + # from the columnwise pass before any thread reads from it in the rowwise pass. + if cutlass.const_expr(CACHE_ACTIVATION): + cute.arch.sync_threads() + if cutlass.const_expr(ROWWISE): + amax_r = quantize_rowwise_mxfp8( + sX_tile, + None if CACHE_ACTIVATION else sActInput_tile, + sO_row_tile, + mS_row_stage, + MAX_NORM_RCP, + tile_row_start, + tile_col_start, + M, + N, + ACTIVATION=None if CACHE_ACTIVATION else ACTIVATION, + DTYPE=DTYPE, + FP8_DTYPE=FP8_DTYPE, + TILE_X=TILE_X, + TILE_Y=TILE_Y, + SKIP_INPUT_MASKING=SKIP_INPUT_MASKING, + SKIP_SCALE_BOUNDS=SKIP_SCALE_BOUNDS, + WITH_ACT=WITH_ACT and not CACHE_ACTIVATION, + WITH_DACT=WITH_DACT and not CACHE_ACTIVATION, + WITH_DBIAS=DBIAS_REDUCTION_IN_ROWWISE, + dbias_acc=rowwise_dbias_acc, + ) + if cutlass.const_expr(WITH_AMAX and not COLWISE): + amax = amax_r + return amax, dbias + + class MXFP8QuantizeKernel(MXFP8QuantizeKernelBase): """The MXFP8 quantization kernel that mirrors the standard (non-specialized) MXFP8 CUDA C++ quantization kernel with multiple fusions (activation, dbias, etc.). `__call__` method is the entrypoint which is AOT compiled. `self` will be captured so it's fixed per compiled kernel """ - # Vectorised access constants for bank-conflict avoidance (rowwise pass) - _PACK_SIZE = 4 # Elements per vector load - # Each thread reads 8 waves with each wave reads 4 packed bf16, so it reads a whole MXFP8 block in total - _WAVES = MXFP8_BLOCK_SCALING_SIZE // _PACK_SIZE - _TOTAL_BANKS_WIDTH = (32 * 4) // 1 # 32 banks × 4 bytes, in bytes - _THREADS_PER_BANK = _TOTAL_BANKS_WIDTH // MXFP8_BLOCK_SCALING_SIZE # 4 threads per bank _NUM_STAGES = 2 # The pipeline depth is always 2 def __init__( @@ -602,13 +714,6 @@ def __init__( and (cfg.COLWISE or self.COLWISE_DBIAS_REDUCTION_ONLY) and cfg.ACTIVATION != "relu" ) - # The global tensor amax (mAmax) is the max over ALL elements. Each direction's - # per-block amaxes already span every element, so when both passes run we only - # fold the global amax from one of them — favor colwise (matches the flags - # above). The per-block *scale* amax is still computed in each pass for its own - # scale; this only skips the redundant global comparison in the other pass. - self.AMAX_FROM_COLWISE = cfg.WITH_AMAX and cfg.COLWISE - self.AMAX_FROM_ROWWISE = cfg.WITH_AMAX and not cfg.COLWISE @cute.jit def __call__( @@ -1074,34 +1179,8 @@ class DactStorage: # and tile_idx_y is the tile that this stage handles out of the `NUM_TILES` tiles tile_idx_x = bidx tile_idx_y = bidy * self._NUM_TILES + tile_idx - # Process rowwise and colwise quantization separately - if cutlass.const_expr(cfg.COLWISE): - # The first row that belongs to this CTA. Each CTA handles NUM_TILES of (TILE_Y, TILE_X) tiles stacked vertically, - # and each stage handles one of them. - sO_col_tile = sO_col[(None, stage_idx)] - mS_col_stage = cute.flatten(mS_col[(None, (tile_idx_y, tile_idx_x))]) - - amax_c, dbias_c = self._process_colwise( - sX_tile, - sO_col_tile, - mS_col_stage, - tile_idx_y * self._TILE_ROWS, - bidx * self._TILE_COLS, - M, - N, - sActInput_tile, - ) - if cutlass.const_expr(self.AMAX_FROM_COLWISE): - per_thread_amax = cute.arch.fmax(per_thread_amax, amax_c) - if cutlass.const_expr(self.DBIAS_REDUCTION_IN_COLWISE): - block_dbias += dbias_c - # If we don't quantize columnwise but we fuse dbias, do a dbias reduction only pass in the columnwise direction without quantization - if cutlass.const_expr(self.COLWISE_DBIAS_REDUCTION_ONLY): - block_dbias += self._dbias_only_colwise(sX_tile, sActInput_tile) - # If we cache the activation in shared memory, we need to ensure that all threads have finished writing to the shared memory - # from the columnwise pass before any thread reads from it in the rowwise pass. - if cutlass.const_expr(self.CACHE_ACTIVATION): - cute.arch.sync_threads() + sO_row_tile = None + mS_row_stage = None if cutlass.const_expr(cfg.ROWWISE): sO_row_tile = sO_row[(None, stage_idx)] # mS_row is ((SCALE_TILE), (GRID)) where SCALE_TILE = (32, 2). @@ -1111,20 +1190,49 @@ class DactStorage: # (i_hi, tile_Y) tile-major order for swizzled, and the plain # row-tile order for compact. Same source, both layouts correct. mS_row_stage = cute.flatten(mS_row[(None, (tile_idx_y, tile_idx_x))]) - amax_r = self._process_rowwise( - sX_tile, - sO_row_tile, - mS_row_stage, - tile_idx_y * self._TILE_ROWS, - bidx * self._TILE_COLS, - M, - N, - sActInput_tile, - rowwise_dbias_acc, - ) + sO_col_tile = None + mS_col_stage = None + if cutlass.const_expr(cfg.COLWISE): + sO_col_tile = sO_col[(None, stage_idx)] + mS_col_stage = cute.flatten(mS_col[(None, (tile_idx_y, tile_idx_x))]) - if cutlass.const_expr(self.AMAX_FROM_ROWWISE): - per_thread_amax = cute.arch.fmax(per_thread_amax, amax_r) + amax, dbias = quantize_mxfp8_tile( + sX_tile, + sActInput_tile, + sO_row_tile, + mS_row_stage, + sO_col_tile, + mS_col_stage, + rowwise_dbias_acc, + tile_idx_y * self._TILE_ROWS, + tile_idx_x * self._TILE_COLS, + M, + N, + ROWWISE=self.cfg.ROWWISE, + COLWISE=self.cfg.COLWISE, + WITH_AMAX=self.cfg.WITH_AMAX, + WITH_ACT=self.cfg.WITH_ACT, + WITH_DACT=self.cfg.WITH_DACT, + ACTIVATION=self.cfg.ACTIVATION, + MAX_NORM_RCP=self.cfg.MAX_NORM_RCP, + DTYPE=self.cfg.DTYPE, + FP8_DTYPE=self.cfg.FP8_DTYPE, + WITH_GEMM_SWIZZLED_SCALES=self.cfg.WITH_GEMM_SWIZZLED_SCALES, + TILE_X=self._TILE_COLS, + TILE_Y=self._TILE_ROWS, + SKIP_INPUT_MASKING=self.SKIP_INPUT_MASKING, + SKIP_SCALE_BOUNDS=self.SKIP_SCALE_BOUNDS, + CACHE_ACTIVATION=self.CACHE_ACTIVATION, + DBIAS_REDUCTION_IN_COLWISE=self.DBIAS_REDUCTION_IN_COLWISE, + COLWISE_DBIAS_REDUCTION_ONLY=self.COLWISE_DBIAS_REDUCTION_ONLY, + DBIAS_REDUCTION_IN_ROWWISE=self.DBIAS_REDUCTION_IN_ROWWISE, + ) + if cutlass.const_expr(cfg.WITH_AMAX): + per_thread_amax = cute.arch.fmax(per_thread_amax, amax) + if cutlass.const_expr( + self.DBIAS_REDUCTION_IN_COLWISE or self.COLWISE_DBIAS_REDUCTION_ONLY + ): + block_dbias += dbias # Make the shared-memory writes visible to the TMA's async proxy before the TMA reads them. cute.arch.fence_proxy( @@ -1199,24 +1307,6 @@ class DactStorage: # before the kernel returns. cute.arch.cp_async_bulk_wait_group(0, read=False) - @cute.jit - def _dbias_only_colwise( - self, - sX_tile: cute.Tensor, # (TILE_Y, TILE_X) bf16/fp16 smem view, post-TMA - sActInput_tile: Optional[ - cute.Tensor - ] = None, # (TILE_Y, TILE_X) act_input tile (dact only), - ): - return dbias_reduction_colwise( - sX_tile, - sActInput_tile, - ACTIVATION=self.cfg.ACTIVATION, - DTYPE=self.cfg.DTYPE, - TILE_X=self._TILE_COLS, - WITH_DACT=self.cfg.WITH_DACT, - CACHE_ACTIVATION=self.CACHE_ACTIVATION, - ) - @cute.jit def _dbias_reduction_rowwise_epilouge( self, smem: cutlass.memory.SmemAllocator, tidx: Int32, rowwise_dbias_acc: cute.Tensor @@ -1246,18 +1336,22 @@ class DbiasStorage: ), ) sDbias_write = cute.composition(sDbias, tv_layout_dbias_write) + # rowwise_dbias_acc is in quantize_rowwise_mxfp8's rotated wave order; these must match its constants. + PACK_SIZE = 4 + WAVES = MXFP8_BLOCK_SCALING_SIZE // PACK_SIZE + THREADS_PER_BANK = 32 * 4 // MXFP8_BLOCK_SCALING_SIZE # Each thread start reading from the specfic bank based on its thread ID so they can do their best to access different banks # to avoid bank conflict. - bank_group = (tidx % THREADS_PER_WARP) // self._THREADS_PER_BANK + bank_group = (tidx % THREADS_PER_WARP) // THREADS_PER_BANK # The offset this thread should start reading from based on what's its first bank to access. - offset = bank_group * self._PACK_SIZE + offset = bank_group * PACK_SIZE for w in cutlass.range_constexpr( - self._WAVES + WAVES ): # Each thread starts from this offset when writing into SMEM to avoid bank conflict - start = (w * self._PACK_SIZE + offset) % MXFP8_BLOCK_SCALING_SIZE - for i in cutlass.range_constexpr(self._PACK_SIZE): + start = (w * PACK_SIZE + offset) % MXFP8_BLOCK_SCALING_SIZE + for i in cutlass.range_constexpr(PACK_SIZE): # All threads write their per-thread partial sum results to the shared buffer. - sDbias_write[(tidx, start + i)] = rowwise_dbias_acc[w * self._PACK_SIZE + i] + sDbias_write[(tidx, start + i)] = rowwise_dbias_acc[w * PACK_SIZE + i] cute.arch.sync_threads() # All threads reduce the cross-thread partial sums to the per-block partial sum. _, tv_layout_dbias_reduce = cute.make_layout_tv( @@ -1302,82 +1396,3 @@ def _amax_epilogue( amax_i32.iterator, cta_amax.bitcast(Int32), ) - - @cute.jit - def _process_rowwise( - self, - sX_tile: cute.Tensor, # (TILE_Y, TILE_X) bf16/fp16 smem view, post-TMA - sO_row_tile: cute.Tensor, # (TILE_Y, TILE_X) fp8 smem view (rowwise FP8 output) - mS_row_stage: cute.Tensor, # rowwise scale tensor (1D swizzled, or 2D linear) - tile_row_start: Int32, # global row of this stage's row 0 - tile_col_start: Int32, # global col of this CTA's col 0 - M: Int32, - N: Int32, # full input extents, for OOB masking - sActInput_tile: Optional[cute.Tensor] = None, # (TILE_Y, TILE_X) act_input tile (dact only) - dbias_acc: Optional[ - cute.Tensor - ] = None, # rmem Float32[32] dbias accumulator (rowwise-only dbias) - ): - cfg = self.cfg - return quantize_rowwise_mxfp8( - sX_tile, - None if self.CACHE_ACTIVATION else sActInput_tile, - sO_row_tile, - mS_row_stage, - self.cfg.MAX_NORM_RCP, - tile_row_start, - tile_col_start, - M, - N, - ACTIVATION=None if self.CACHE_ACTIVATION else cfg.ACTIVATION, - DTYPE=cfg.DTYPE, - FP8_DTYPE=cfg.FP8_DTYPE, - TILE_X=self._TILE_COLS, - TILE_Y=self._TILE_ROWS, - WAVES=self._WAVES, - THREADS_PER_BANK=self._THREADS_PER_BANK, - PACK_SIZE=self._PACK_SIZE, - SKIP_INPUT_MASKING=self.SKIP_INPUT_MASKING, - SKIP_SCALE_BOUNDS=self.SKIP_SCALE_BOUNDS, - WITH_ACT=cfg.WITH_ACT and not self.CACHE_ACTIVATION, - WITH_DACT=cfg.WITH_DACT and not self.CACHE_ACTIVATION, - WITH_DBIAS=self.DBIAS_REDUCTION_IN_ROWWISE, - dbias_acc=dbias_acc, - ) - - @cute.jit - def _process_colwise( - self, - sX_tile: cute.Tensor, # (TILE_Y, TILE_X) bf16/fp16 smem view, post-TMA - sO_col_tile: cute.Tensor, # (TILE_Y, TILE_X) fp8 smem view (colwise FP8 output) - mS_col_stage: cute.Tensor, # colwise scale tensor (1D swizzled, or 2D linear) - tile_row_start: Int32, # global row of this stage's row 0 - tile_col_start: Int32, # global col of this CTA's col 0 - M: Int32, - N: Int32, # full input extents, for OOB masking - sActInput_tile: Optional[cute.Tensor] = None, # (TILE_Y, TILE_X) act_input tile (dact only) - ): - cfg = self.cfg - return quantize_colwise_mxfp8( - sX_tile, - sActInput_tile, - sO_col_tile, - mS_col_stage, - self.cfg.MAX_NORM_RCP, - tile_row_start, - tile_col_start, - M, - N, - ACTIVATION=cfg.ACTIVATION, - DTYPE=cfg.DTYPE, - FP8_DTYPE=cfg.FP8_DTYPE, - SWIZZLE=cfg.WITH_GEMM_SWIZZLED_SCALES, - TILE_X=self._TILE_COLS, - TILE_Y=self._TILE_ROWS, - SKIP_INPUT_MASKING=self.SKIP_INPUT_MASKING, - SKIP_SCALE_BOUNDS=self.SKIP_SCALE_BOUNDS, - WITH_ACT=cfg.WITH_ACT, - WITH_DACT=cfg.WITH_DACT, - WITH_DBIAS=self.DBIAS_REDUCTION_IN_COLWISE, - CACHE_ACTIVATION=self.CACHE_ACTIVATION, - ) From fddd076e98610520213d46b2399d783a7d6db53a Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Wed, 14 Oct 2026 00:09:31 +0000 Subject: [PATCH 45/58] refactor Signed-off-by: Kaining Zhong --- .../mxfp8/quantize_mxfp8_g2r_rowwise_2lane.py | 26 +++++++++++-------- .../mxfp8/quantize_mxfp8_g2s_bidimensional.py | 7 ++--- 2 files changed, 19 insertions(+), 14 deletions(-) diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_2lane.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_2lane.py index 68c69040cfc..9db77f4d36e 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_2lane.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_2lane.py @@ -4,7 +4,7 @@ """Two-lane GMEM-to-RMEM rowwise MXFP8 cast, ported from cast_rowwise.cuh.""" -from typing import Optional +from typing import Optional, Type import cutlass from cutlass import cute @@ -34,8 +34,10 @@ def quantize_rowwise_mxfp8_g2r( gX_tiles: cute.Tensor, gO_tiles: cute.Tensor, mS_tiles: cute.Tensor, - cfg: cutlass.Constexpr[MXFP8QuantizeConfig], N: Int32, + DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], + FP8_DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], + MAX_NORM_RCP: cutlass.Constexpr[float], CTA_THREADS_Y: cutlass.Constexpr[int], CTA_THREADS_X: cutlass.Constexpr[int], ELEMENTS_PER_THREAD: cutlass.Constexpr[int], @@ -69,21 +71,21 @@ def quantize_rowwise_mxfp8_g2r( ) # We load fragments from all GMEM tiles at once - rX = cute.make_rmem_tensor((ELEMENTS_PER_THREAD, NUM_TILES), cfg.DTYPE) + rX = cute.make_rmem_tensor((ELEMENTS_PER_THREAD, NUM_TILES), DTYPE) # We store a fragment of one tile at a time - rO = cute.make_rmem_tensor((ELEMENTS_PER_THREAD,), cfg.FP8_DTYPE) + rO = cute.make_rmem_tensor((ELEMENTS_PER_THREAD,), FP8_DTYPE) # RMEM flags to indicate whether each tile is out of bounds or not # With CHECK_BOUNDS=False DCE will remove it so we don't pay extra cost when we don't use it tile_valid_flags = cute.make_rmem_tensor((NUM_TILES,), Boolean) # 16 * 2 bytes = 32 bytes, one 256-bit load where the toolchain supports it LOAD_WIDTH = min( - ELEMENTS_PER_THREAD * cfg.DTYPE.width, + ELEMENTS_PER_THREAD * DTYPE.width, max_num_bits_per_copy_ld_global(), ) load_atom = cute.make_copy_atom( cute.nvgpu.CopyG2ROp(), - cfg.DTYPE, + DTYPE, num_bits_per_copy=LOAD_WIDTH, # TODO(kainingz): due to https://github.com/NVIDIA/cutlass/issues/3726 using # invariant=True here will drop the cache policy @@ -91,11 +93,11 @@ def quantize_rowwise_mxfp8_g2r( ) store_atom = cute.make_copy_atom( cute.nvgpu.CopyR2GOp(), - cfg.FP8_DTYPE, - num_bits_per_copy=ELEMENTS_PER_THREAD * cfg.FP8_DTYPE.width, + FP8_DTYPE, + num_bits_per_copy=ELEMENTS_PER_THREAD * FP8_DTYPE.width, ) - scale_and_pack = mul_bf16x2x2_cvt_bf16x4_to_fp8x4(cfg.FP8_DTYPE) + scale_and_pack = mul_bf16x2x2_cvt_bf16x4_to_fp8x4(FP8_DTYPE) # Copy all tiles from GMEM to RMEM and potentially check their bounds for tile_idx in cutlass.range_constexpr(NUM_TILES): @@ -144,7 +146,7 @@ def extract_mx_block_amax(values_i32: cute.Tensor) -> Float32: rX_tile_i32 = cute.recast_tensor(rX[None, tile_idx], Int32) rO_u32 = cute.recast_tensor(rO, Uint32) amax = extract_mx_block_amax(rX_tile_i32) - exponent = cvt_f32_to_fp8e8m0fnu(amax * cfg.MAX_NORM_RCP) + exponent = cvt_f32_to_fp8e8m0fnu(amax * MAX_NORM_RCP) # The even lane of each pair stores its block's scale byte, as cast_rowwise.cuh does. # A warp's stores coalesce in either layout: one 16-byte run when packed, four @@ -479,8 +481,10 @@ def _kernel_main( gX_tiles, gO_tiles, mS_tiles, - self.cfg, N, + DTYPE=self.cfg.DTYPE, + FP8_DTYPE=self.cfg.FP8_DTYPE, + MAX_NORM_RCP=self.cfg.MAX_NORM_RCP, CTA_THREADS_Y=CTA_THREADS_Y, CTA_THREADS_X=CTA_THREADS_X, ELEMENTS_PER_THREAD=self._ELEMENTS_PER_THREAD, diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2s_bidimensional.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2s_bidimensional.py index f6c6d391d99..004d2a64a99 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2s_bidimensional.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2s_bidimensional.py @@ -43,8 +43,9 @@ @cute.jit -def quantize_bidimensional_mxfp8_swizzled( - sX_tile: cute.Tensor, # 32x64 input tile (already sliced to this stage), Sw<3,4,3> swizzled +def quantize_bidimensional_mxfp8( + sX_tile: cute.Tensor, # 32x64 input tile (already sliced to this stage); any SMEM layout, + # Sw<3,4,3> avoids bank conflicts sO_row_tile: cute.Tensor, # 32x64 rowwise-output tile sO_col_tile: cute.Tensor, # 32x64 colwise-output tile sS_row_tile: cute.Tensor, # (32, 2) smem rowwise-scale staging tile (flushed in the epilogue) @@ -611,7 +612,7 @@ class SharedStorage: tile_x = tile_idx % num_tiles_x # Process the 32x64 SMEM tile for this stage - quantize_bidimensional_mxfp8_swizzled( + quantize_bidimensional_mxfp8( sX[None, stage_idx], sO_row[None, stage_idx], sO_col[None, stage_idx], From 3a8a8825f9154df770e2bd3d4f218bf961f9b801 Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Wed, 14 Oct 2026 20:03:01 +0000 Subject: [PATCH 46/58] more refactor Signed-off-by: Kaining Zhong --- .../cast/mxfp8/quantize_mxfp8_common.py | 40 +++++++------------ ...nsional.py => quantize_mxfp8_g2r_bidim.py} | 8 ++-- .../mxfp8/quantize_mxfp8_g2r_rowwise_1lane.py | 10 ++--- .../mxfp8/quantize_mxfp8_g2r_rowwise_2lane.py | 8 ++-- ...nsional.py => quantize_mxfp8_g2s_bidim.py} | 12 +++--- ...neral.py => quantize_mxfp8_g2s_general.py} | 26 ++++++------ transformer_engine/common/CuTeDSL/utils.py | 11 +++++ 7 files changed, 59 insertions(+), 56 deletions(-) rename transformer_engine/common/CuTeDSL/cast/mxfp8/{quantize_mxfp8_g2r_bidimensional.py => quantize_mxfp8_g2r_bidim.py} (99%) rename transformer_engine/common/CuTeDSL/cast/mxfp8/{quantize_mxfp8_g2s_bidimensional.py => quantize_mxfp8_g2s_bidim.py} (98%) rename transformer_engine/common/CuTeDSL/cast/mxfp8/{quantize_mxfp8_general.py => quantize_mxfp8_g2s_general.py} (98%) diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_common.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_common.py index 748f19756f4..133e7232f48 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_common.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_common.py @@ -11,13 +11,12 @@ import cutlass from cutlass import cute -from cutlass import Boolean, Float32, Int32, Int64, Uint32 +from cutlass import Boolean, Float32, Int32, Int64 from cuda.bindings.driver import CUstream # pylint: disable=no-name-in-module import tvm_ffi from transformer_engine.common.CuTeDSL.utils import ( str_to_cutlass_dtype, - abs_max_x2_bf16, device_compute_capability, ) from transformer_engine.common.CuTeDSL.activations import ( @@ -256,17 +255,6 @@ def __call__( """ -@cute.jit -def bf16_pair_magnitude(pair: Int32) -> Uint32: - """Fold both BF16 halves and clear the sign set by max.xorsign.abs.""" - swapped = cute.arch.inline_ptx( - "prmt.b32 {$w0}, {$r0}, {$r0}, 0x1032;", - write_only_types=[Int32], - read_only_args=[pair], - ) - return Uint32(abs_max_x2_bf16(pair, swapped)) & Uint32(0x7FFF) - - class MXFP8QuantizeEntry(MXFP8QuantizeKernelBase): """Select the appropriate MXFP8 quantization kernel based on the configuration and runtime shapes.""" @@ -274,11 +262,11 @@ def __init__(self, cfg: MXFP8QuantizeConfig): # Kernels import shared definitions from this module. Resolve their # implementations only when constructing the compile-time entry object. # pylint: disable=import-outside-toplevel - from .quantize_mxfp8_general import MXFP8QuantizeKernel - from .quantize_mxfp8_g2r_rowwise_1lane import MXFP8QuantizeG2RRowwise1LaneKernel - from .quantize_mxfp8_g2s_bidimensional import MXFP8QuantizeSpecializedBidimensionalKernel - from .quantize_mxfp8_g2r_rowwise_2lane import MXFP8QuantizeG2RRowwise2LaneKernel - from .quantize_mxfp8_g2r_bidimensional import MXFP8QuantizeRegisterBidimensionalKernel + from .quantize_mxfp8_g2s_general import MXFP8QuantizeGeneralG2SKernel + from .quantize_mxfp8_g2r_rowwise_1lane import MXFP8QuantizeRowwise1LaneG2RKernel + from .quantize_mxfp8_g2s_bidim import MXFP8QuantizeBidimensionalG2SKernel + from .quantize_mxfp8_g2r_rowwise_2lane import MXFP8QuantizeRowwise2LaneG2RKernel + from .quantize_mxfp8_g2r_bidim import MXFP8QuantizeBidimensionalG2RKernel # pylint: enable=import-outside-toplevel @@ -296,21 +284,23 @@ def __init__(self, cfg: MXFP8QuantizeConfig): ) # Instantiate all possible kernels at compile time, # and we will pick the right one at runtime based on the input shape and config. - self.general_skip_input_masking_kernel = MXFP8QuantizeKernel(cfg, SKIP_INPUT_MASKING=True) - self.general_skip_input_masking_and_scale_bounds_kernel = MXFP8QuantizeKernel( + self.general_skip_input_masking_kernel = MXFP8QuantizeGeneralG2SKernel( + cfg, SKIP_INPUT_MASKING=True + ) + self.general_skip_input_masking_and_scale_bounds_kernel = MXFP8QuantizeGeneralG2SKernel( cfg, SKIP_INPUT_MASKING=True, SKIP_SCALE_BOUNDS=True ) # We only need to mask when WITH_ACT is enabled because with activations applied zeros filled by TMA affect block statistics - self.general_kernel = MXFP8QuantizeKernel(cfg) if self.ACT_NEED_MASKING else None + self.general_kernel = MXFP8QuantizeGeneralG2SKernel(cfg) if self.ACT_NEED_MASKING else None # Specialized cast-only kernels in a specific dispatch priority order so we prefer to dispatch to kernels that come first in the list. self.specialized_kernels = [ kernel_cls(cfg) for kernel_cls in ( - MXFP8QuantizeG2RRowwise2LaneKernel, - MXFP8QuantizeRegisterBidimensionalKernel, - MXFP8QuantizeG2RRowwise1LaneKernel, - MXFP8QuantizeSpecializedBidimensionalKernel, + MXFP8QuantizeRowwise2LaneG2RKernel, + MXFP8QuantizeBidimensionalG2RKernel, + MXFP8QuantizeRowwise1LaneG2RKernel, + MXFP8QuantizeBidimensionalG2SKernel, ) if kernel_cls.supports(cfg) ] diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidim.py similarity index 99% rename from transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py rename to transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidim.py index 75ad5799cbf..e1d828147c9 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidim.py @@ -13,6 +13,7 @@ from transformer_engine.common.CuTeDSL.utils import ( abs_max_x2_bf16, + bf16_pair_magnitude, bf16_mx_scale_reciprocal, bf16_mx_scale_reciprocal_pair, bf16_mx_scale_bytes_pair, @@ -24,7 +25,6 @@ MXFP8_BLOCK_SCALING_SIZE, MXFP8QuantizeConfig, MXFP8QuantizeKernelBase, - bf16_pair_magnitude, derive_swizzled_scale_layout, ) @@ -388,7 +388,7 @@ def _quantize_colwise( @cute.jit -def quantize_bidimensional_mxfp8_g2r( +def quantize_mxfp8_gmem_tile_bidimensional( gX_tile: cute.Tensor, gO_row_tile: Optional[cute.Tensor], gS_row_tile: Optional[cute.Tensor], @@ -544,7 +544,7 @@ def quantize_bidimensional_mxfp8_g2r( ) -class MXFP8QuantizeRegisterBidimensionalKernel(MXFP8QuantizeKernelBase): +class MXFP8QuantizeBidimensionalG2RKernel(MXFP8QuantizeKernelBase): """Load one CTA tile into registers and quantize it along both axes. TILE_ROWS = _CTA_THREADS_Y * _ELEMENTS_Y_PER_THREAD = 32. @@ -789,7 +789,7 @@ def _kernel_main( cache_policy = create_l2_policy(True) - quantize_bidimensional_mxfp8_g2r( + quantize_mxfp8_gmem_tile_bidimensional( gX_tile, gO_row_tile, gS_row_tile, diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_1lane.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_1lane.py index 1c005c22fdb..97ea8c0f5f4 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_1lane.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_1lane.py @@ -74,7 +74,7 @@ def _flush_scales_to_gmem( @cute.jit -def quantize_rowwise_mxfp8_g2r_1lane( +def quantize_mxfp8_gmem_tile_rowwise_1lane( gX_tile: cute.Tensor, gO_tile: cute.Tensor, gS_tile: cute.Tensor, @@ -205,7 +205,7 @@ class SharedStorage: ) -class MXFP8QuantizeG2RRowwise1LaneKernel(MXFP8QuantizeKernelBase): +class MXFP8QuantizeRowwise1LaneG2RKernel(MXFP8QuantizeKernelBase): """Specialized cast-only ROWWISE-only MXFP8 kernel. Requires N % 128 == 0 (full vectorizable column chunks). @@ -236,7 +236,7 @@ def eligible_to_dispatch( # N divisible by 128 for vectorized stores, and grid.y within limits. return ( N % 128 == 0 - and cute.ceil_div(M, MXFP8QuantizeG2RRowwise1LaneKernel._TILE_ROWS) <= 65535 + and cute.ceil_div(M, MXFP8QuantizeRowwise1LaneG2RKernel._TILE_ROWS) <= 65535 and mNoop.toint() == 0 ) @@ -262,7 +262,7 @@ def __call__( ): if cutlass.const_expr(CUTEDSL_DEBUG_LOGGING): cute.printf( - f"[CuTeDSL] MXFP8QuantizeG2RRowwise1LaneKernel.__call__() with config: {self.cfg}\n" + f"[CuTeDSL] MXFP8QuantizeRowwise1LaneG2RKernel.__call__() with config: {self.cfg}\n" ) M = mX.shape[0] @@ -312,7 +312,7 @@ def _kernel_main( gO_tile = cute.local_tile(mO_row, TILER, (bidy, bidx)) gS_tile = cute.local_tile(mS_row, SCALE_TILER, (bidy, bidx)) - quantize_rowwise_mxfp8_g2r_1lane( + quantize_mxfp8_gmem_tile_rowwise_1lane( gX_tile, gO_tile, gS_tile, diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_2lane.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_2lane.py index 9db77f4d36e..d1eb2cd0753 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_2lane.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_2lane.py @@ -13,6 +13,7 @@ from transformer_engine.common.CuTeDSL.utils import ( abs_max_x2_bf16, + bf16_pair_magnitude, create_l2_policy, exp2_bf16x2_rcp, max_num_bits_per_copy_ld_global, @@ -24,13 +25,12 @@ from .quantize_mxfp8_common import ( MXFP8QuantizeConfig, MXFP8QuantizeKernelBase, - bf16_pair_magnitude, derive_swizzled_scale_layout, ) @cute.jit -def quantize_rowwise_mxfp8_g2r( +def quantize_mxfp8_gmem_tile_rowwise_2lane( gX_tiles: cute.Tensor, gO_tiles: cute.Tensor, mS_tiles: cute.Tensor, @@ -167,7 +167,7 @@ def extract_mx_block_amax(values_i32: cute.Tensor) -> Float32: cute.copy(store_atom, rO, gO_thread, cache_policy=output_policy) -class MXFP8QuantizeG2RRowwise2LaneKernel(MXFP8QuantizeKernelBase): +class MXFP8QuantizeRowwise2LaneG2RKernel(MXFP8QuantizeKernelBase): """Two lanes per MXFP8 block, with a CTA sweeping NUM_TILES horizontally. One CTA tile has TILE_ROWS = CTA_THREADS_Y and @@ -477,7 +477,7 @@ def _kernel_main( # The last grid.z slice can overshoot the row tiles; those CTAs have no work. if cta_row * CTA_ROWS < M: - quantize_rowwise_mxfp8_g2r( + quantize_mxfp8_gmem_tile_rowwise_2lane( gX_tiles, gO_tiles, mS_tiles, diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2s_bidimensional.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2s_bidim.py similarity index 98% rename from transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2s_bidimensional.py rename to transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2s_bidim.py index 004d2a64a99..65bf548022a 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2s_bidimensional.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2s_bidim.py @@ -43,7 +43,7 @@ @cute.jit -def quantize_bidimensional_mxfp8( +def quantize_mxfp8_smem_tile_bidimensional( sX_tile: cute.Tensor, # 32x64 input tile (already sliced to this stage); any SMEM layout, # Sw<3,4,3> avoids bank conflicts sO_row_tile: cute.Tensor, # 32x64 rowwise-output tile @@ -229,7 +229,7 @@ def quantize_bidimensional_mxfp8( cute.autovec_copy(rO_col, tXsO_col) -class MXFP8QuantizeSpecializedBidimensionalKernel(MXFP8QuantizeKernelBase): +class MXFP8QuantizeBidimensionalG2SKernel(MXFP8QuantizeKernelBase): """Specialized cast-only rowwise+colwise MXFP8 kernel (swizzled TMA, one 32x32 tile per warp).""" _WARPS_PER_CTA = 2 @@ -271,7 +271,7 @@ def eligible_to_dispatch( mS_col: Optional[cute.Tensor], mNoop: cute.Pointer, ) -> Boolean: - cls = MXFP8QuantizeSpecializedBidimensionalKernel + cls = MXFP8QuantizeBidimensionalG2SKernel rows_per_cta = cls._TILE_ROWS * cls._NUM_TILES_Y # packed scales: one tile row return cute.ceil_div(mX.shape[0], rows_per_cta) <= 65535 and mNoop.toint() == 0 @@ -298,7 +298,7 @@ def __call__( ): if cutlass.const_expr(CUTEDSL_DEBUG_LOGGING): cute.printf( - "[CuTeDSL] MXFP8QuantizeSpecializedBidimensionalKernel.__call__() with config:" + "[CuTeDSL] MXFP8QuantizeBidimensionalG2SKernel.__call__() with config:" f" {self.cfg}\n" ) M = mX.shape[0] @@ -611,8 +611,8 @@ class SharedStorage: tile_y = tile_idx // num_tiles_x tile_x = tile_idx % num_tiles_x - # Process the 32x64 SMEM tile for this stage - quantize_bidimensional_mxfp8( + # Process the 32x64 SMEM tile for thiHs stage + quantize_mxfp8_smem_tile_bidimensional( sX[None, stage_idx], sO_row[None, stage_idx], sO_col[None, stage_idx], diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_general.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2s_general.py similarity index 98% rename from transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_general.py rename to transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2s_general.py index d74d4ca824b..9801d82a259 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_general.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2s_general.py @@ -76,7 +76,7 @@ @cute.jit -def quantize_rowwise_mxfp8( +def _quantize_rowwise_mxfp8( sX_tile: cute.Tensor, # (TILE_Y, TILE_X) bf16/fp16 smem view, post-TMA sActInput_tile: Optional[cute.Tensor], # (TILE_Y, TILE_X) act-input smem tile (dact only) sO_row_tile: cute.Tensor, # (TILE_Y, TILE_X) fp8 smem view (rowwise FP8 output) @@ -329,7 +329,7 @@ def quantize_rowwise_mxfp8( @cute.jit -def quantize_colwise_mxfp8( +def _quantize_colwise_mxfp8( sX_tile: cute.Tensor, # (TILE_Y, TILE_X) bf16/fp16 smem view, post-TMA sActInput_tile: Optional[cute.Tensor], # (TILE_Y, TILE_X) act-input smem tile (dact only) sO_col_tile: cute.Tensor, # (TILE_Y, TILE_X) fp8 smem view (colwise FP8 output) @@ -505,7 +505,7 @@ def quantize_colwise_mxfp8( @cute.jit -def dbias_reduction_colwise( +def _dbias_reduction_colwise( sX_tile: cute.Tensor, # (TILE_Y, TILE_X) bf16/fp16/fp32 smem view, post-TMA sA_tile: Optional[cute.Tensor], # (TILE_Y, TILE_X) activation-input smem tile (dact only) ACTIVATION: cutlass.Constexpr[str | None], @@ -517,7 +517,7 @@ def dbias_reduction_colwise( """Reduce one SMEM tile along rows (the dbias direction) and possibly cache dact values, with no quantization.""" tidx, _, _ = cute.arch.thread_idx() - # Same TV layout as quantize_colwise_mxfp8: thread tidx owns column tidx, all TILE_Y rows. + # Same TV layout as _quantize_colwise_mxfp8: thread tidx owns column tidx, all TILE_Y rows. _, tv_layout = cute.make_layout_tv( thr_layout=cute.make_layout((1, TILE_X), stride=(TILE_X, 1)), val_layout=cute.make_layout((MXFP8_BLOCK_SCALING_SIZE, 1), stride=(1, 1)), @@ -540,7 +540,7 @@ def dbias_reduction_colwise( @cute.jit -def quantize_mxfp8_tile( +def quantize_mxfp8_smem_tile_general( sX_tile: cute.Tensor, # (TILE_Y, TILE_X) bf16/fp16 smem view, post-TMA sActInput_tile: Optional[cute.Tensor], # (TILE_Y, TILE_X) act-input smem tile (dact only) sO_row_tile: Optional[cute.Tensor], # (TILE_Y, TILE_X) fp8 smem view (rowwise FP8 output) @@ -580,7 +580,7 @@ def quantize_mxfp8_tile( amax = None dbias = None if cutlass.const_expr(COLWISE): - amax_c, dbias_c = quantize_colwise_mxfp8( + amax_c, dbias_c = _quantize_colwise_mxfp8( sX_tile, sActInput_tile, sO_col_tile, @@ -611,7 +611,7 @@ def quantize_mxfp8_tile( dbias = dbias_c # If we don't quantize columnwise but we fuse dbias, do a dbias reduction only pass in the columnwise direction without quantization if cutlass.const_expr(COLWISE_DBIAS_REDUCTION_ONLY): - dbias = dbias_reduction_colwise( + dbias = _dbias_reduction_colwise( sX_tile, sActInput_tile, ACTIVATION=ACTIVATION, @@ -625,7 +625,7 @@ def quantize_mxfp8_tile( if cutlass.const_expr(CACHE_ACTIVATION): cute.arch.sync_threads() if cutlass.const_expr(ROWWISE): - amax_r = quantize_rowwise_mxfp8( + amax_r = _quantize_rowwise_mxfp8( sX_tile, None if CACHE_ACTIVATION else sActInput_tile, sO_row_tile, @@ -652,7 +652,7 @@ def quantize_mxfp8_tile( return amax, dbias -class MXFP8QuantizeKernel(MXFP8QuantizeKernelBase): +class MXFP8QuantizeGeneralG2SKernel(MXFP8QuantizeKernelBase): """The MXFP8 quantization kernel that mirrors the standard (non-specialized) MXFP8 CUDA C++ quantization kernel with multiple fusions (activation, dbias, etc.). `__call__` method is the entrypoint which is AOT compiled. `self` will be captured so it's fixed per compiled kernel @@ -734,7 +734,9 @@ def __call__( stream: CUstream, ): if cutlass.const_expr(CUTEDSL_DEBUG_LOGGING): - cute.printf(f"[CuTeDSL] MXFP8QuantizeKernel.__call__() with config: {self.cfg}\n") + cute.printf( + f"[CuTeDSL] MXFP8QuantizeGeneralG2SKernel.__call__() with config: {self.cfg}\n" + ) M = mX.shape[0] N = mX.shape[1] @@ -1196,7 +1198,7 @@ class DactStorage: sO_col_tile = sO_col[(None, stage_idx)] mS_col_stage = cute.flatten(mS_col[(None, (tile_idx_y, tile_idx_x))]) - amax, dbias = quantize_mxfp8_tile( + amax, dbias = quantize_mxfp8_smem_tile_general( sX_tile, sActInput_tile, sO_row_tile, @@ -1336,7 +1338,7 @@ class DbiasStorage: ), ) sDbias_write = cute.composition(sDbias, tv_layout_dbias_write) - # rowwise_dbias_acc is in quantize_rowwise_mxfp8's rotated wave order; these must match its constants. + # rowwise_dbias_acc is in _quantize_rowwise_mxfp8's rotated wave order; these must match its constants. PACK_SIZE = 4 WAVES = MXFP8_BLOCK_SCALING_SIZE // PACK_SIZE THREADS_PER_BANK = 32 * 4 // MXFP8_BLOCK_SCALING_SIZE diff --git a/transformer_engine/common/CuTeDSL/utils.py b/transformer_engine/common/CuTeDSL/utils.py index d5ffd4abfe9..caaf8339524 100644 --- a/transformer_engine/common/CuTeDSL/utils.py +++ b/transformer_engine/common/CuTeDSL/utils.py @@ -102,6 +102,17 @@ def exp2_bf16x2_rcp(scale_e8m0) -> Uint32: return bf16_bits | (bf16_bits << Uint32(16)) +@cute.jit +def bf16_pair_magnitude(pair: Int32) -> Uint32: + """Fold both BF16 halves and clear the sign set by max.xorsign.abs.""" + swapped = cute.arch.inline_ptx( + "prmt.b32 {$w0}, {$r0}, {$r0}, 0x1032;", + write_only_types=[Int32], + read_only_args=[pair], + ) + return Uint32(abs_max_x2_bf16(pair, swapped)) & Uint32(0x7FFF) + + @cute.jit def bf16_mx_scale_reciprocal(magnitude: Uint32, MAX_EXPONENT: cutlass.Constexpr[int]) -> Uint32: """Exact BF16 reciprocal scale bits, including E8M0 zero, Inf and NaN cases.""" From 30c92ab3bc274e01fecb36554cfea5c9a46705b4 Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Wed, 14 Oct 2026 14:16:13 -0700 Subject: [PATCH 47/58] g2r bidim can also support colwise only Signed-off-by: Kaining Zhong --- .../cast/mxfp8/quantize_mxfp8_g2r_bidim.py | 53 ++++++++++++++----- 1 file changed, 39 insertions(+), 14 deletions(-) diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidim.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidim.py index e1d828147c9..5e8d724504e 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidim.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidim.py @@ -561,7 +561,7 @@ class MXFP8QuantizeBidimensionalG2RKernel(MXFP8QuantizeKernelBase): @staticmethod def supports(cfg: MXFP8QuantizeConfig) -> bool: # BF16 packed math. - return cfg.is_plain_cast() and cfg.DTYPE is cutlass.BFloat16 and cfg.ROWWISE and cfg.COLWISE + return cfg.is_plain_cast() and cfg.DTYPE is cutlass.BFloat16 and cfg.COLWISE @staticmethod @cute.jit @@ -680,14 +680,24 @@ def _launch( if N == 2048: payload_layout = cute.make_layout((M, 2048), stride=(2048, 1)) gX = cute.make_tensor(mX.iterator, payload_layout) - gO_row = cute.make_tensor(mO_row.iterator, payload_layout) + gO_row = ( + cute.make_tensor(mO_row.iterator, payload_layout) + if cutlass.const_expr(self.cfg.ROWWISE) + else None + ) gO_col = cute.make_tensor(mO_col.iterator, payload_layout) if cutlass.const_expr(self.cfg.WITH_GEMM_SWIZZLED_SCALES): - gS_row, gS_col = derive_swizzled_scale_layout(M, 2048, True, True, mS_row, mS_col) + gS_row, gS_col = derive_swizzled_scale_layout( + M, 2048, self.cfg.ROWWISE, True, mS_row, mS_col + ) else: - gS_row = cute.make_tensor( - mS_row.iterator, - cute.make_layout((cute.size(mS_row, mode=[0]), 64), stride=(64, 1)), + gS_row = ( + cute.make_tensor( + mS_row.iterator, + cute.make_layout((cute.size(mS_row, mode=[0]), 64), stride=(64, 1)), + ) + if cutlass.const_expr(self.cfg.ROWWISE) + else None ) gS_col = cute.make_tensor( mS_col.iterator, @@ -704,14 +714,24 @@ def _launch( elif N == 4096: payload_layout = cute.make_layout((M, 4096), stride=(4096, 1)) gX = cute.make_tensor(mX.iterator, payload_layout) - gO_row = cute.make_tensor(mO_row.iterator, payload_layout) + gO_row = ( + cute.make_tensor(mO_row.iterator, payload_layout) + if cutlass.const_expr(self.cfg.ROWWISE) + else None + ) gO_col = cute.make_tensor(mO_col.iterator, payload_layout) if cutlass.const_expr(self.cfg.WITH_GEMM_SWIZZLED_SCALES): - gS_row, gS_col = derive_swizzled_scale_layout(M, 4096, True, True, mS_row, mS_col) + gS_row, gS_col = derive_swizzled_scale_layout( + M, 4096, self.cfg.ROWWISE, True, mS_row, mS_col + ) else: - gS_row = cute.make_tensor( - mS_row.iterator, - cute.make_layout((cute.size(mS_row, mode=[0]), 128), stride=(128, 1)), + gS_row = ( + cute.make_tensor( + mS_row.iterator, + cute.make_layout((cute.size(mS_row, mode=[0]), 128), stride=(128, 1)), + ) + if cutlass.const_expr(self.cfg.ROWWISE) + else None ) gS_col = cute.make_tensor( mS_col.iterator, @@ -727,7 +747,9 @@ def _launch( ).launch(grid=grid, block=block, cluster=cluster, min_blocks_per_mp=4, stream=stream) else: if cutlass.const_expr(self.cfg.WITH_GEMM_SWIZZLED_SCALES): - gS_row, gS_col = derive_swizzled_scale_layout(M, N, True, True, mS_row, mS_col) + gS_row, gS_col = derive_swizzled_scale_layout( + M, N, self.cfg.ROWWISE, True, mS_row, mS_col + ) else: gS_row, gS_col = mS_row, mS_col self.kernel( @@ -782,9 +804,12 @@ def _kernel_main( COL_SCALE_TILER = (TILE_ROWS // 32, TILE_COLS) gX_tile = cute.local_tile(mX, TILER, cta_coord) - gO_row_tile = cute.local_tile(mO_row, TILER, cta_coord) + gO_row_tile = None + gS_row_tile = None + if cutlass.const_expr(self.cfg.ROWWISE): + gO_row_tile = cute.local_tile(mO_row, TILER, cta_coord) + gS_row_tile = cute.local_tile(mS_row, ROW_SCALE_TILER, cta_coord) gO_col_tile = cute.local_tile(mO_col, TILER, cta_coord) - gS_row_tile = cute.local_tile(mS_row, ROW_SCALE_TILER, cta_coord) gS_col_tile = cute.local_tile(mS_col, COL_SCALE_TILER, cta_coord) cache_policy = create_l2_policy(True) From 5c3465a56103e1805e4e737ed0334b1139358f42 Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Thu, 15 Oct 2026 20:29:21 +0000 Subject: [PATCH 48/58] make g2r support activations Signed-off-by: Kaining Zhong --- .../cast/mxfp8/quantize_mxfp8_common.py | 7 +- .../cast/mxfp8/quantize_mxfp8_g2r_bidim.py | 93 +++++++++++++- .../mxfp8/quantize_mxfp8_g2r_rowwise_1lane.py | 1 + .../mxfp8/quantize_mxfp8_g2r_rowwise_2lane.py | 114 +++++++++++++++++- .../cast/mxfp8/quantize_mxfp8_g2s_bidim.py | 1 + transformer_engine/common/CuTeDSL/utils.py | 10 ++ 6 files changed, 219 insertions(+), 7 deletions(-) diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_common.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_common.py index 133e7232f48..0df60e0714e 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_common.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_common.py @@ -230,6 +230,7 @@ def eligible_to_dispatch( mO_col: Optional[cute.Tensor], mS_col: Optional[cute.Tensor], mNoop: cute.Pointer, + mDActInput: Optional[cute.Tensor], ) -> Boolean: """Runtime check: whether we can dispatch to the instantiated kernel with this call's shapes and noops.""" raise NotImplementedError @@ -322,7 +323,9 @@ def __call__( dispatched_to_specialized = False for kernel in self.specialized_kernels: if not dispatched_to_specialized: - if kernel.eligible_to_dispatch(self.cfg, mX, mO_row, mS_row, mO_col, mS_col, mNoop): + if kernel.eligible_to_dispatch( + self.cfg, mX, mO_row, mS_row, mO_col, mS_col, mNoop, mDActInput + ): dispatched_to_specialized = True kernel( mX, @@ -503,7 +506,7 @@ def compile_cutedsl_function_from_cfg(cfg): scale_col_fake, # mO_col, mS_col amax_fake, # mAmax noop_fake, # mNoop (pointer to the cast_noop flag, may be null at launch) - act_input_fake, # mDActInput (backward slot, unused) + act_input_fake, # mDActInput (activation input, dact only) workspace_fake, # mWorkspace(backward slot, unused) cute.runtime.make_fake_stream(), # stream (compiled as an explicit tvm-ffi # "handle" arg; C++ passes the CUDA stream diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidim.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidim.py index 5e8d724504e..0e8efdb3bd3 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidim.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidim.py @@ -8,7 +8,7 @@ import cutlass from cutlass import cute -from cutlass import Boolean, Int32, Int64, Uint32, Uint8 +from cutlass import Boolean, Float32, Int32, Int64, Uint32, Uint8 from cuda.bindings.driver import CUstream # pylint: disable=no-name-in-module from transformer_engine.common.CuTeDSL.utils import ( @@ -19,12 +19,17 @@ bf16_mx_scale_bytes_pair, create_l2_policy, max_num_bits_per_copy_ld_global, + x2_hi_to_f32_bf16, + x2_lo_to_f32_bf16, + cvt_f32x2_to_bf16x2, ) from transformer_engine.common.CuTeDSL.utils_fp8 import mul_bf16x2x2_cvt_bf16x4_to_fp8x4 from .quantize_mxfp8_common import ( MXFP8_BLOCK_SCALING_SIZE, MXFP8QuantizeConfig, MXFP8QuantizeKernelBase, + SUPPORTED_ACTIVATIONS, + SUPPORTED_DACTIVATIONS, derive_swizzled_scale_layout, ) @@ -390,12 +395,16 @@ def _quantize_colwise( @cute.jit def quantize_mxfp8_gmem_tile_bidimensional( gX_tile: cute.Tensor, + gA_tile: Optional[cute.Tensor], gO_row_tile: Optional[cute.Tensor], gS_row_tile: Optional[cute.Tensor], gO_col_tile: Optional[cute.Tensor], gS_col_tile: Optional[cute.Tensor], ROWWISE: cutlass.Constexpr[bool], COLWISE: cutlass.Constexpr[bool], + ACTIVATION: cutlass.Constexpr[Optional[str]], + WITH_ACT: cutlass.Constexpr[bool], + WITH_DACT: cutlass.Constexpr[bool], DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], FP8_DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], WITH_GEMM_SWIZZLED_SCALES: cutlass.Constexpr[bool], @@ -462,6 +471,44 @@ def quantize_mxfp8_gmem_tile_bidimensional( load_atom, gX_thread[(None, row_idx),], rX[None, row_idx], cache_policy=cache_policy ) + if cutlass.const_expr(WITH_ACT or WITH_DACT): + rA = cute.make_rmem_tensor((ELEMENTS_X_PER_THREAD, ELEMENTS_Y_PER_THREAD), DTYPE) + if cutlass.const_expr(WITH_DACT): + act_load_atom = cute.make_copy_atom( + cute.nvgpu.CopyG2ROp(), + DTYPE, + num_bits_per_copy=min( + ELEMENTS_X_PER_THREAD * DTYPE.width, + max_num_bits_per_copy_ld_global(), + ), + invariant=False, + ) + # Activation tile has the same layout as the input tile so we can reuse the tv layout + gA_thread = cute.composition(gA_tile, tv_layout)[tidx, None] + for row_idx in cutlass.range_constexpr(ELEMENTS_Y_PER_THREAD): + cute.copy( + act_load_atom, + gA_thread[(None, row_idx),], + rA[None, row_idx], + cache_policy=cache_policy, + ) + # Apply activation row by row + for row_idx in cutlass.range_constexpr(ELEMENTS_Y_PER_THREAD): + rX_i32 = cute.recast_tensor(rX[None, row_idx], Int32) + rA_i32 = cute.recast_tensor(rA[None, row_idx], Int32) + for i in cutlass.range_constexpr(ELEMENTS_X_PER_THREAD // 2): + lo = x2_lo_to_f32_bf16(rX_i32[i]) + hi = x2_hi_to_f32_bf16(rX_i32[i]) + if cutlass.const_expr(WITH_DACT): + dop = SUPPORTED_DACTIVATIONS[ACTIVATION] + lo = lo * dop(x2_lo_to_f32_bf16(rA_i32[i])) + hi = hi * dop(x2_hi_to_f32_bf16(rA_i32[i])) + else: + op = SUPPORTED_ACTIVATIONS[ACTIVATION] + lo = op(lo) + hi = op(hi) + rX_i32[i] = cvt_f32x2_to_bf16x2(lo, hi) + if cutlass.const_expr(ROWWISE): gO_row_thread = cute.composition(gO_row_tile, tv_layout)[tidx, None] @@ -561,7 +608,12 @@ class MXFP8QuantizeBidimensionalG2RKernel(MXFP8QuantizeKernelBase): @staticmethod def supports(cfg: MXFP8QuantizeConfig) -> bool: # BF16 packed math. - return cfg.is_plain_cast() and cfg.DTYPE is cutlass.BFloat16 and cfg.COLWISE + return ( + not cfg.WITH_AMAX + and not cfg.WITH_DBIAS + and cfg.DTYPE is cutlass.BFloat16 + and cfg.COLWISE + ) @staticmethod @cute.jit @@ -573,8 +625,12 @@ def eligible_to_dispatch( mO_col: Optional[cute.Tensor], mS_col: Optional[cute.Tensor], mNoop: cute.Pointer, + mDActInput: Optional[cute.Tensor], ) -> Boolean: M, N = mX.shape + act_input_aligned = Boolean(True) + if cutlass.const_expr(mDActInput is not None): + act_input_aligned = mDActInput.iterator.toint() % 32 == 0 # Whole 32-row by 256-column tiles, grid.y within limits, and aligned input and # colwise-scale pointers. return ( @@ -582,6 +638,7 @@ def eligible_to_dispatch( and N % 256 == 0 and M // 32 <= 65535 and mX.iterator.toint() % 32 == 0 + and act_input_aligned and mS_col.iterator.toint() % 16 == 0 and mNoop.toint() == 0 ) @@ -612,6 +669,7 @@ def __call__( if Int64(M // TILE_ROWS) * (N // WIDE_TILE_COLS) >= 4096: self._launch( mX, + mDActInput, mO_row, mS_row, mO_col, @@ -623,6 +681,7 @@ def __call__( else: self._launch( mX, + mDActInput, mO_row, mS_row, mO_col, @@ -634,6 +693,7 @@ def __call__( else: self._launch( mX, + mDActInput, mO_row, mS_row, mO_col, @@ -647,6 +707,7 @@ def __call__( def _launch( self, mX: cute.Tensor, + mA: Optional[cute.Tensor], mO_row: cute.Tensor, mS_row: cute.Tensor, mO_col: cute.Tensor, @@ -660,6 +721,8 @@ def _launch( # Dispatch checks these alignments before selecting this kernel. mX = cute.make_tensor(mX.iterator.align(32), mX.layout) mS_col = cute.make_tensor(mS_col.iterator.align(16), mS_col.layout) + if cutlass.const_expr(self.cfg.WITH_DACT): + mA = cute.make_tensor(mA.iterator.align(32), mA.layout) M, N = mX.shape TILE_ROWS = self._CTA_THREADS_Y * self._ELEMENTS_Y_PER_THREAD TILE_COLS = self._CTA_THREADS_X * ELEMENTS_X_PER_THREAD @@ -676,10 +739,17 @@ def _launch( block = [THREADS_PER_CTA] cluster = [cluster_cols, 1, 1] + gA = None # Hardcode the layout using static N for some of the commonly used shapes if N == 2048: payload_layout = cute.make_layout((M, 2048), stride=(2048, 1)) gX = cute.make_tensor(mX.iterator, payload_layout) + if cutlass.const_expr(self.cfg.WITH_DACT): + gA = ( + cute.make_tensor(mA.iterator, payload_layout) + if cutlass.const_expr(self.cfg.WITH_DACT) + else None + ) gO_row = ( cute.make_tensor(mO_row.iterator, payload_layout) if cutlass.const_expr(self.cfg.ROWWISE) @@ -705,6 +775,7 @@ def _launch( ) self.kernel( gX, + gA, gO_row, gS_row, gO_col, @@ -714,6 +785,12 @@ def _launch( elif N == 4096: payload_layout = cute.make_layout((M, 4096), stride=(4096, 1)) gX = cute.make_tensor(mX.iterator, payload_layout) + if cutlass.const_expr(self.cfg.WITH_DACT): + gA = ( + cute.make_tensor(mA.iterator, payload_layout) + if cutlass.const_expr(self.cfg.WITH_DACT) + else None + ) gO_row = ( cute.make_tensor(mO_row.iterator, payload_layout) if cutlass.const_expr(self.cfg.ROWWISE) @@ -739,6 +816,7 @@ def _launch( ) self.kernel( gX, + gA, gO_row, gS_row, gO_col, @@ -754,6 +832,7 @@ def _launch( gS_row, gS_col = mS_row, mS_col self.kernel( mX, + mA, mO_row, gS_row, mO_col, @@ -765,6 +844,7 @@ def _launch( def kernel( self, mX: cute.Tensor, + mA: Optional[cute.Tensor], mO_row: cute.Tensor, mS_row: cute.Tensor, mO_col: cute.Tensor, @@ -775,6 +855,7 @@ def kernel( cast_noop flag: the dispatcher routes calls with one to the general kernel.""" self._kernel_main( mX, + mA, mO_row, mS_row, mO_col, @@ -786,6 +867,7 @@ def kernel( def _kernel_main( self, mX: cute.Tensor, + mA: Optional[cute.Tensor], mO_row: cute.Tensor, mS_row: cute.Tensor, mO_col: cute.Tensor, @@ -804,6 +886,9 @@ def _kernel_main( COL_SCALE_TILER = (TILE_ROWS // 32, TILE_COLS) gX_tile = cute.local_tile(mX, TILER, cta_coord) + gA_tile = None + if cutlass.const_expr(self.cfg.WITH_DACT): + gA_tile = cute.local_tile(mA, TILER, cta_coord) gO_row_tile = None gS_row_tile = None if cutlass.const_expr(self.cfg.ROWWISE): @@ -816,12 +901,16 @@ def _kernel_main( quantize_mxfp8_gmem_tile_bidimensional( gX_tile, + gA_tile, gO_row_tile, gS_row_tile, gO_col_tile, gS_col_tile, ROWWISE=self.cfg.ROWWISE, COLWISE=self.cfg.COLWISE, + ACTIVATION=self.cfg.ACTIVATION, + WITH_ACT=self.cfg.WITH_ACT, + WITH_DACT=self.cfg.WITH_DACT, DTYPE=self.cfg.DTYPE, FP8_DTYPE=self.cfg.FP8_DTYPE, WITH_GEMM_SWIZZLED_SCALES=self.cfg.WITH_GEMM_SWIZZLED_SCALES, diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_1lane.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_1lane.py index 97ea8c0f5f4..682ba02c611 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_1lane.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_1lane.py @@ -231,6 +231,7 @@ def eligible_to_dispatch( mO_col: Optional[cute.Tensor], mS_col: Optional[cute.Tensor], mNoop: cute.Pointer, + mDActInput: Optional[cute.Tensor], ) -> Boolean: M, N = mX.shape # N divisible by 128 for vectorized stores, and grid.y within limits. diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_2lane.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_2lane.py index d1eb2cd0753..ee2f53c1ccb 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_2lane.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_2lane.py @@ -17,6 +17,9 @@ create_l2_policy, exp2_bf16x2_rcp, max_num_bits_per_copy_ld_global, + x2_hi_to_f32_bf16, + x2_lo_to_f32_bf16, + cvt_f32x2_to_bf16x2, ) from transformer_engine.common.CuTeDSL.utils_fp8 import ( cvt_f32_to_fp8e8m0fnu, @@ -25,6 +28,8 @@ from .quantize_mxfp8_common import ( MXFP8QuantizeConfig, MXFP8QuantizeKernelBase, + SUPPORTED_ACTIVATIONS, + SUPPORTED_DACTIVATIONS, derive_swizzled_scale_layout, ) @@ -32,9 +37,13 @@ @cute.jit def quantize_mxfp8_gmem_tile_rowwise_2lane( gX_tiles: cute.Tensor, + gA_tiles: Optional[cute.Tensor], gO_tiles: cute.Tensor, mS_tiles: cute.Tensor, N: Int32, + ACTIVATION: cutlass.Constexpr[Optional[str]], + WITH_ACT: cutlass.Constexpr[bool], + WITH_DACT: cutlass.Constexpr[bool], DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], FP8_DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], MAX_NORM_RCP: cutlass.Constexpr[float], @@ -72,6 +81,8 @@ def quantize_mxfp8_gmem_tile_rowwise_2lane( # We load fragments from all GMEM tiles at once rX = cute.make_rmem_tensor((ELEMENTS_PER_THREAD, NUM_TILES), DTYPE) + # Backward: the activation input, loaded alongside the incoming gradient in rX + rA = cute.make_rmem_tensor((ELEMENTS_PER_THREAD, NUM_TILES), DTYPE) # We store a fragment of one tile at a time rO = cute.make_rmem_tensor((ELEMENTS_PER_THREAD,), FP8_DTYPE) # RMEM flags to indicate whether each tile is out of bounds or not @@ -91,6 +102,12 @@ def quantize_mxfp8_gmem_tile_rowwise_2lane( # invariant=True here will drop the cache policy invariant=False, ) + act_load_atom = cute.make_copy_atom( + cute.nvgpu.CopyG2ROp(), + DTYPE, + num_bits_per_copy=LOAD_WIDTH, + invariant=False, + ) store_atom = cute.make_copy_atom( cute.nvgpu.CopyR2GOp(), FP8_DTYPE, @@ -116,8 +133,33 @@ def quantize_mxfp8_gmem_tile_rowwise_2lane( valid = col < N and Int64(row) * N + col < num_elems tile_valid_flags[tile_idx] = valid rX[None, tile_idx].fill(0) + if cutlass.const_expr(WITH_DACT): + rA[None, tile_idx].fill(0) if valid: cute.copy(load_atom, gX_thread, rX[None, tile_idx], cache_policy=input_policy) + if cutlass.const_expr(WITH_DACT): + gA_thread = cute.composition(gA_tiles[None, None, tile_idx], tv_payload)[tidx, None] + cute.copy(act_load_atom, gA_thread, rA[None, tile_idx], cache_policy=input_policy) + + # Fused activation: widen to FP32, apply it, and round back to BF16 as the general kernel + # does (it truncates the activated value to the input type before the amax). The values are + # then exactly BF16, so the packed BF16 amax and scaling below apply unchanged. + if cutlass.const_expr(WITH_ACT or WITH_DACT): + for tile_idx in cutlass.range_constexpr(NUM_TILES): + rX_i32 = cute.recast_tensor(rX[None, tile_idx], Int32) + rA_i32 = cute.recast_tensor(rA[None, tile_idx], Int32) + for i in cutlass.range_constexpr(ELEMENTS_PER_THREAD // 2): + lo = x2_lo_to_f32_bf16(rX_i32[i]) + hi = x2_hi_to_f32_bf16(rX_i32[i]) + if cutlass.const_expr(WITH_DACT): + dop = SUPPORTED_DACTIVATIONS[ACTIVATION] + lo = lo * dop(x2_lo_to_f32_bf16(rA_i32[i])) + hi = hi * dop(x2_hi_to_f32_bf16(rA_i32[i])) + else: + op = SUPPORTED_ACTIVATIONS[ACTIVATION] + lo = op(lo) + hi = op(hi) + rX_i32[i] = cvt_f32x2_to_bf16x2(lo, hi) @cute.jit def extract_mx_block_amax(values_i32: cute.Tensor) -> Float32: @@ -189,7 +231,8 @@ class MXFP8QuantizeRowwise2LaneG2RKernel(MXFP8QuantizeKernelBase): def supports(cfg: MXFP8QuantizeConfig) -> bool: # BF16 packed math. return ( - cfg.is_plain_cast() + not cfg.WITH_AMAX + and not cfg.WITH_DBIAS and cfg.DTYPE is cutlass.BFloat16 and cfg.ROWWISE and not cfg.COLWISE @@ -205,8 +248,12 @@ def eligible_to_dispatch( mO_col: Optional[cute.Tensor], mS_col: Optional[cute.Tensor], mNoop: cute.Pointer, + mDActInput: Optional[cute.Tensor], ) -> Boolean: N = mX.shape[1] + act_input_aligned = Boolean(True) + if cutlass.const_expr(mDActInput is not None): + act_input_aligned = mDActInput.iterator.toint() % 32 == 0 scale_stride_matches = Boolean(True) if cutlass.const_expr(not cfg.WITH_GEMM_SWIZZLED_SCALES): # The flat view needs the packed scales to be contiguous. @@ -216,6 +263,7 @@ def eligible_to_dispatch( return ( N % 128 == 0 and mX.iterator.toint() % 32 == 0 + and act_input_aligned and scale_stride_matches and mNoop.toint() == 0 ) @@ -246,22 +294,63 @@ def __call__( if num_elems <= 2**31 - 1: rows = cute.ceil_div(num_elems, 1024) flat = cute.make_layout((rows, 1024), stride=(1024, 1)) + mA = None + if cutlass.const_expr(self.cfg.WITH_DACT): + mA = cute.make_tensor(mDActInput.iterator, flat) self._launch_tiers( cute.make_tensor(mX.iterator, flat), + mA, cute.make_tensor(mO_row.iterator, flat), cute.make_tensor(mS_row.iterator, cute.make_layout((rows, 32), stride=(32, 1))), stream, num_elems, ) else: - self._launch_tiers(mX, mO_row, mS_row, stream, num_elems) + self._launch_tiers(mX, mDActInput, mO_row, mS_row, stream, num_elems) else: - self._launch_tiers(mX, mO_row, mS_row, stream, num_elems) + self._launch_tiers(mX, mDActInput, mO_row, mS_row, stream, num_elems) @cute.jit def _launch_tiers( self, mX: cute.Tensor, + mA: Optional[cute.Tensor], + mO_row: cute.Tensor, + mS_row: cute.Tensor, + stream: CUstream, + num_elems: Int64, + ): + """Pick the launch configuration: the size tiers, except for swizzled scales with a fused + activation when N does not fill the tiers' 1024-column CTAs.""" + if cutlass.const_expr( + self.cfg.WITH_GEMM_SWIZZLED_SCALES and (self.cfg.WITH_ACT or self.cfg.WITH_DACT) + ): + if mX.shape[1] % 1024 != 0: + # The swizzled path tiles the real (M, N) shape, so a 1024-column CTA overhangs + # N and its idle lanes still pay for the activation math. 256-column CTAs divide + # every eligible N (N % 128 == 0 and the tile is 2 x 128 columns). + self._launch( + mX, + mA, + mO_row, + mS_row, + stream, + num_elems, + CTA_THREADS_Y=16, + CTA_THREADS_X=8, + NUM_TILES=2, + L2_CACHED_CTA_PERCENT=0, + ) + else: + self._launch_size_tiers(mX, mA, mO_row, mS_row, stream, num_elems) + else: + self._launch_size_tiers(mX, mA, mO_row, mS_row, stream, num_elems) + + @cute.jit + def _launch_size_tiers( + self, + mX: cute.Tensor, + mA: Optional[cute.Tensor], mO_row: cute.Tensor, mS_row: cute.Tensor, stream: CUstream, @@ -272,6 +361,7 @@ def _launch_tiers( if output_bytes <= 12 << 20: self._launch( mX, + mA, mO_row, mS_row, stream, @@ -284,6 +374,7 @@ def _launch_tiers( elif output_bytes <= 48 << 20: self._launch( mX, + mA, mO_row, mS_row, stream, @@ -296,6 +387,7 @@ def _launch_tiers( elif output_bytes <= 96 << 20: self._launch( mX, + mA, mO_row, mS_row, stream, @@ -308,6 +400,7 @@ def _launch_tiers( else: self._launch( mX, + mA, mO_row, mS_row, stream, @@ -322,6 +415,7 @@ def _launch_tiers( def _launch( self, mX: cute.Tensor, + mA: Optional[cute.Tensor], mO: cute.Tensor, mS: cute.Tensor, stream: CUstream, @@ -361,6 +455,7 @@ def _launch( if check_bounds: self.kernel( mX, + mA, mO, mS, first_streaming_cta, @@ -374,6 +469,7 @@ def _launch( else: self.kernel( mX, + mA, mO, mS, first_streaming_cta, @@ -389,6 +485,7 @@ def _launch( def kernel( self, mX: cute.Tensor, + mA: Optional[cute.Tensor], mO: cute.Tensor, mS: cute.Tensor, first_streaming_cta: Int64, @@ -403,6 +500,7 @@ def kernel( dispatcher only selects this kernel when no cast_noop flag is passed.""" self._kernel_main( mX, + mA, mO, mS, first_streaming_cta, @@ -418,6 +516,7 @@ def kernel( def _kernel_main( self, mX: cute.Tensor, + mA: Optional[cute.Tensor], mO: cute.Tensor, mS: cute.Tensor, first_streaming_cta: Int64, @@ -431,6 +530,8 @@ def _kernel_main( """Construct CTA tile views and run the specialized quantization.""" # Dispatch checks 32-byte input alignment before selecting this kernel. gX = cute.make_tensor(mX.iterator.align(32), mX.layout) + if cutlass.const_expr(self.cfg.WITH_DACT): + mA = cute.make_tensor(mA.iterator.align(32), mA.layout) # grid.x walks column tiles; grid.y and grid.z together walk row tiles. bidx, bidy, bidz = cute.arch.block_idx() cta_row = bidz * self._MAX_GRID_Y + bidy @@ -472,6 +573,9 @@ def _kernel_main( SCALE_TILER = (TILE_ROWS, TILE_COLS // 32) # Mode 1's first index is 0 because CTA_ROWS == TILE_ROWS (there is only 1 tile on the row dimension). gX_tiles = cute.local_tile(gX_cta, TILER, (0, None)) + gA_tiles = None + if cutlass.const_expr(self.cfg.WITH_DACT): + gA_tiles = cute.local_tile(cute.local_tile(mA, CTA_TILER, cta_coord), TILER, (0, None)) gO_tiles = cute.local_tile(gO_cta, TILER, (0, None)) mS_tiles = cute.local_tile(gS_cta, SCALE_TILER, (0, None)) @@ -479,9 +583,13 @@ def _kernel_main( if cta_row * CTA_ROWS < M: quantize_mxfp8_gmem_tile_rowwise_2lane( gX_tiles, + gA_tiles, gO_tiles, mS_tiles, N, + ACTIVATION=self.cfg.ACTIVATION, + WITH_ACT=self.cfg.WITH_ACT, + WITH_DACT=self.cfg.WITH_DACT, DTYPE=self.cfg.DTYPE, FP8_DTYPE=self.cfg.FP8_DTYPE, MAX_NORM_RCP=self.cfg.MAX_NORM_RCP, diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2s_bidim.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2s_bidim.py index 65bf548022a..1192d98cb3a 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2s_bidim.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2s_bidim.py @@ -270,6 +270,7 @@ def eligible_to_dispatch( mO_col: Optional[cute.Tensor], mS_col: Optional[cute.Tensor], mNoop: cute.Pointer, + mDActInput: Optional[cute.Tensor], ) -> Boolean: cls = MXFP8QuantizeBidimensionalG2SKernel rows_per_cta = cls._TILE_ROWS * cls._NUM_TILES_Y # packed scales: one tile row diff --git a/transformer_engine/common/CuTeDSL/utils.py b/transformer_engine/common/CuTeDSL/utils.py index caaf8339524..90c61d03052 100644 --- a/transformer_engine/common/CuTeDSL/utils.py +++ b/transformer_engine/common/CuTeDSL/utils.py @@ -286,6 +286,16 @@ def x2_hi_to_f32_bf16(bits: Int32) -> Float32: return ((bits >> Int32(16)) << Int32(16)).bitcast(Float32) +@cute.jit +def cvt_f32x2_to_bf16x2(lo: Float32, hi: Float32) -> Int32: + """Round two FP32 values to BF16 (round-to-nearest-even) and pack them, lo in the low half.""" + return cute.arch.inline_ptx( + "cvt.rn.bf16x2.f32 {$w0}, {$r0}, {$r1};", + write_only_types=[Int32], + read_only_args=[hi, lo], + ) + + @cute.jit def truncate_f32_bf16(val: Float32) -> Float32: """Round FP32 to BF16 precision (round-to-nearest-even), keeping FP32 storage.""" From b1db277941e6221817f3d86d718206d9e78ab00b Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Thu, 15 Oct 2026 13:44:36 -0700 Subject: [PATCH 49/58] fix Signed-off-by: Kaining Zhong --- .../cast/mxfp8/quantize_mxfp8_g2r_bidim.py | 23 ++++++++----------- 1 file changed, 10 insertions(+), 13 deletions(-) diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidim.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidim.py index 0e8efdb3bd3..776fb1db21e 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidim.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidim.py @@ -739,17 +739,15 @@ def _launch( block = [THREADS_PER_CTA] cluster = [cluster_cols, 1, 1] - gA = None # Hardcode the layout using static N for some of the commonly used shapes if N == 2048: payload_layout = cute.make_layout((M, 2048), stride=(2048, 1)) gX = cute.make_tensor(mX.iterator, payload_layout) - if cutlass.const_expr(self.cfg.WITH_DACT): - gA = ( - cute.make_tensor(mA.iterator, payload_layout) - if cutlass.const_expr(self.cfg.WITH_DACT) - else None - ) + gA = ( + cute.make_tensor(mA.iterator, payload_layout) + if cutlass.const_expr(self.cfg.WITH_DACT) + else None + ) gO_row = ( cute.make_tensor(mO_row.iterator, payload_layout) if cutlass.const_expr(self.cfg.ROWWISE) @@ -785,12 +783,11 @@ def _launch( elif N == 4096: payload_layout = cute.make_layout((M, 4096), stride=(4096, 1)) gX = cute.make_tensor(mX.iterator, payload_layout) - if cutlass.const_expr(self.cfg.WITH_DACT): - gA = ( - cute.make_tensor(mA.iterator, payload_layout) - if cutlass.const_expr(self.cfg.WITH_DACT) - else None - ) + gA = ( + cute.make_tensor(mA.iterator, payload_layout) + if cutlass.const_expr(self.cfg.WITH_DACT) + else None + ) gO_row = ( cute.make_tensor(mO_row.iterator, payload_layout) if cutlass.const_expr(self.cfg.ROWWISE) From dcd8b08ed7491d7772863c8919e27218f0d23cfe Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Thu, 15 Oct 2026 14:51:01 -0700 Subject: [PATCH 50/58] fix Signed-off-by: Kaining Zhong --- transformer_engine/common/CuTeDSL/utils_fp8.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/transformer_engine/common/CuTeDSL/utils_fp8.py b/transformer_engine/common/CuTeDSL/utils_fp8.py index 9a972e39469..61355cca1f2 100644 --- a/transformer_engine/common/CuTeDSL/utils_fp8.py +++ b/transformer_engine/common/CuTeDSL/utils_fp8.py @@ -183,11 +183,11 @@ def mul_f32x4_cvt_f32x4_to_fp8x4(fp8_dtype, relu: bool = False) -> Callable[..., def mul_bf16x2x2_cvt_bf16x4_to_fp8x4(fp8_dtype) -> Callable[..., Uint32]: """Return a multiply/cast for two BF16 pairs, producing four packed FP8 bytes. - CUDA 13.1+ supports direct BF16-to-FP8 conversion. Older compilers use - FP32 fused multiply-add, matching the corresponding ptx::mul_cvt_4x overload. + Uses a direct BF16-to-FP8 conversion on CUDA 13.1+ and an FP32 detour before that. """ out_op = "e5m2x2" if fp8_dtype is Float8E5M2 else "e4m3x2" if cutlass.target_version(min_version="13.1"): + # Multiply in BF16, then convert BF16 -> FP8. body = ( ".reg.b32 y0, y1;\n\t" "mul.rn.bf16x2 y0, {$r0}, {$r1};\n\t" @@ -196,6 +196,8 @@ def mul_bf16x2x2_cvt_bf16x4_to_fp8x4(fp8_dtype) -> Callable[..., Uint32]: f"cvt.rn.satfinite.{out_op}.bf16x2 z1, y1;\n\t" ) else: + # No BF16 -> FP8 conversion before CUDA 13.1: fma.rn.f32.bf16 multiplies into FP32 + # (a*b + -0.0 is exact and keeps -0.0), then convert FP32 -> FP8. body = ( ".reg.b16 x0,x1,x2,x3,s0,s1,s2,s3;\n\t" ".reg.f32 y0,y1,y2,y3;\n\t" @@ -203,10 +205,10 @@ def mul_bf16x2x2_cvt_bf16x4_to_fp8x4(fp8_dtype) -> Callable[..., Uint32]: "mov.b32 {s0,s1}, {$r1};\n\t" "mov.b32 {x2,x3}, {$r2};\n\t" "mov.b32 {s2,s3}, {$r3};\n\t" - "fma.rn.f32.bf16 y0, x0, s0, 0f00000000;\n\t" - "fma.rn.f32.bf16 y1, x1, s1, 0f00000000;\n\t" - "fma.rn.f32.bf16 y2, x2, s2, 0f00000000;\n\t" - "fma.rn.f32.bf16 y3, x3, s3, 0f00000000;\n\t" + "fma.rn.f32.bf16 y0, x0, s0, 0f80000000;\n\t" + "fma.rn.f32.bf16 y1, x1, s1, 0f80000000;\n\t" + "fma.rn.f32.bf16 y2, x2, s2, 0f80000000;\n\t" + "fma.rn.f32.bf16 y3, x3, s3, 0f80000000;\n\t" f"cvt.rn.satfinite.{out_op}.f32 z0, y1, y0;\n\t" f"cvt.rn.satfinite.{out_op}.f32 z1, y3, y2;\n\t" ) From 86233c813ea4f8b572b23e501cb1639198319cf9 Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Thu, 15 Oct 2026 16:51:00 -0700 Subject: [PATCH 51/58] +0 and -0 should both legal Signed-off-by: Kaining Zhong --- tests/jax/test_mxfp8_cutedsl_backend.py | 5 ++++- tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py | 6 ++++-- 2 files changed, 8 insertions(+), 3 deletions(-) diff --git a/tests/jax/test_mxfp8_cutedsl_backend.py b/tests/jax/test_mxfp8_cutedsl_backend.py index bd492c967a1..2ee82f8c5d5 100644 --- a/tests/jax/test_mxfp8_cutedsl_backend.py +++ b/tests/jax/test_mxfp8_cutedsl_backend.py @@ -197,12 +197,15 @@ def extract_quantized_output(out, dbias): ScaledTensor1x carries scale_inv already trimmed to the unpadded shape (see its __post_init__), so there is no uninitialized scale padding to exclude. + FP8 -0.0 (0x80) is folded into +0.0: both backends round to zero there and may differ only + in its sign, which training does not depend on. """ tensors = [out] if isinstance(out, ScaledTensor1x) else [out.rowwise_tensor, out.colwise_tensor] parts = {} for t in tensors: name = "colwise" if t.is_colwise else "rowwise" - parts[f"{name} data"] = np.asarray(t.data.view(jnp.uint8)) + data = np.asarray(t.data.view(jnp.uint8)) + parts[f"{name} data"] = np.where(data == 0x80, 0, data).astype(np.uint8) parts[f"{name} scales"] = np.asarray(t.scale_inv.view(jnp.uint8)) return parts, None if dbias is None else np.asarray(dbias) diff --git a/tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py b/tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py index a886b63c088..5d7feeb6ce6 100644 --- a/tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py +++ b/tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py @@ -251,18 +251,20 @@ def extract_quantized_output(out, rowwise, columnwise, swizzled): compared. Swizzled layout: the meaningful scales are scattered by the swizzle, so the top-left slice is meaningless; both backends zero the padding (see zero_scales_kernel), so compare the whole buffer instead -- which also covers the padding-zeroing itself. + FP8 -0.0 (0x80) is folded into +0.0: both backends round to zero there and may differ only + in its sign, which training does not depend on. """ parts = {} if rowwise: d = out._rowwise_data.view(torch.uint8) M, N = d.shape - parts["rowwise data"] = d.clone() + parts["rowwise data"] = d.masked_fill(d == 0x80, 0) s = out._rowwise_scale_inv parts["rowwise scales"] = (s if swizzled else s[:M, : (N + 31) // 32]).clone() if columnwise: d = out._columnwise_data.view(torch.uint8) M, N = d.shape - parts["colwise data"] = d.clone() + parts["colwise data"] = d.masked_fill(d == 0x80, 0) s = out._columnwise_scale_inv parts["colwise scales"] = (s if swizzled else s[: (M + 31) // 32, :N]).clone() return parts From dba8dbf3f408053cbc627873be096c4ba2aee513 Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Thu, 15 Oct 2026 16:57:58 -0700 Subject: [PATCH 52/58] nit Signed-off-by: Kaining Zhong --- .../CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidim.py | 2 +- transformer_engine/common/CuTeDSL/utils_fp8.py | 14 ++++++-------- 2 files changed, 7 insertions(+), 9 deletions(-) diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidim.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidim.py index 776fb1db21e..61598a7d37b 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidim.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidim.py @@ -8,7 +8,7 @@ import cutlass from cutlass import cute -from cutlass import Boolean, Float32, Int32, Int64, Uint32, Uint8 +from cutlass import Boolean, Int32, Int64, Uint32, Uint8 from cuda.bindings.driver import CUstream # pylint: disable=no-name-in-module from transformer_engine.common.CuTeDSL.utils import ( diff --git a/transformer_engine/common/CuTeDSL/utils_fp8.py b/transformer_engine/common/CuTeDSL/utils_fp8.py index 61355cca1f2..9a972e39469 100644 --- a/transformer_engine/common/CuTeDSL/utils_fp8.py +++ b/transformer_engine/common/CuTeDSL/utils_fp8.py @@ -183,11 +183,11 @@ def mul_f32x4_cvt_f32x4_to_fp8x4(fp8_dtype, relu: bool = False) -> Callable[..., def mul_bf16x2x2_cvt_bf16x4_to_fp8x4(fp8_dtype) -> Callable[..., Uint32]: """Return a multiply/cast for two BF16 pairs, producing four packed FP8 bytes. - Uses a direct BF16-to-FP8 conversion on CUDA 13.1+ and an FP32 detour before that. + CUDA 13.1+ supports direct BF16-to-FP8 conversion. Older compilers use + FP32 fused multiply-add, matching the corresponding ptx::mul_cvt_4x overload. """ out_op = "e5m2x2" if fp8_dtype is Float8E5M2 else "e4m3x2" if cutlass.target_version(min_version="13.1"): - # Multiply in BF16, then convert BF16 -> FP8. body = ( ".reg.b32 y0, y1;\n\t" "mul.rn.bf16x2 y0, {$r0}, {$r1};\n\t" @@ -196,8 +196,6 @@ def mul_bf16x2x2_cvt_bf16x4_to_fp8x4(fp8_dtype) -> Callable[..., Uint32]: f"cvt.rn.satfinite.{out_op}.bf16x2 z1, y1;\n\t" ) else: - # No BF16 -> FP8 conversion before CUDA 13.1: fma.rn.f32.bf16 multiplies into FP32 - # (a*b + -0.0 is exact and keeps -0.0), then convert FP32 -> FP8. body = ( ".reg.b16 x0,x1,x2,x3,s0,s1,s2,s3;\n\t" ".reg.f32 y0,y1,y2,y3;\n\t" @@ -205,10 +203,10 @@ def mul_bf16x2x2_cvt_bf16x4_to_fp8x4(fp8_dtype) -> Callable[..., Uint32]: "mov.b32 {s0,s1}, {$r1};\n\t" "mov.b32 {x2,x3}, {$r2};\n\t" "mov.b32 {s2,s3}, {$r3};\n\t" - "fma.rn.f32.bf16 y0, x0, s0, 0f80000000;\n\t" - "fma.rn.f32.bf16 y1, x1, s1, 0f80000000;\n\t" - "fma.rn.f32.bf16 y2, x2, s2, 0f80000000;\n\t" - "fma.rn.f32.bf16 y3, x3, s3, 0f80000000;\n\t" + "fma.rn.f32.bf16 y0, x0, s0, 0f00000000;\n\t" + "fma.rn.f32.bf16 y1, x1, s1, 0f00000000;\n\t" + "fma.rn.f32.bf16 y2, x2, s2, 0f00000000;\n\t" + "fma.rn.f32.bf16 y3, x3, s3, 0f00000000;\n\t" f"cvt.rn.satfinite.{out_op}.f32 z0, y1, y0;\n\t" f"cvt.rn.satfinite.{out_op}.f32 z1, y3, y2;\n\t" ) From 444bc94bdde932c84a4ee66318e3e9066ba9364e Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Thu, 15 Oct 2026 18:05:26 -0700 Subject: [PATCH 53/58] fix after merge Signed-off-by: Kaining Zhong --- .../cast/mxfp8/group_quantize_mxfp8_common.py | 223 ++++++++---------- .../cast/mxfp8/group_quantize_mxfp8_g2r.py | 103 ++++++-- .../mxfp8/group_quantize_mxfp8_general.py | 25 +- ...group_quantize_mxfp8_single_tensor_view.py | 160 +++++++++++++ .../cast/mxfp8/quantize_mxfp8_g2r_bidim.py | 9 +- 5 files changed, 375 insertions(+), 145 deletions(-) create mode 100644 transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8_single_tensor_view.py diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8_common.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8_common.py index 493ca566d7d..e43948f3c31 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8_common.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8_common.py @@ -14,13 +14,14 @@ # pylint: disable=missing-class-docstring +import abc import logging import os from typing import Optional import cutlass from cutlass import cute -from cutlass import Float32, Int32, Int64 +from cutlass import Boolean, Float32, Int32 from cutlass.utils import HardwareInfo from cuda.bindings.driver import CUstream # pylint: disable=no-name-in-module import tvm_ffi @@ -152,66 +153,102 @@ def __str__(self): __repr__ = __str__ +class MXFP8GroupQuantizeKernelBase(abc.ABC): + """Base class for grouped MXFP8 quantize kernels dispatched by MXFP8GroupQuantizeEntry.""" + + @staticmethod + def supports(cfg: MXFP8GroupQuantizeConfig) -> bool: # pylint: disable=unused-argument + """Compile time check: whether this kernel handles cfg at all.""" + return False + + def eligible_to_dispatch( + self, + mX: cute.Tensor, + mO_row: Optional[cute.Tensor], + mO_col: Optional[cute.Tensor], + mS_row: cute.Tensor, + mS_col: cute.Tensor, + mOffsets: Optional[cute.Tensor], + mFirstDims: Optional[cute.Tensor], + mLastDims: Optional[cute.Tensor], + num_tensors: Int32, + mNoop: cute.Pointer, + ) -> Boolean: + """Runtime check: whether this call's shapes, alignments and cast-noop flag allow + dispatching to this kernel (decorate with @cute.jit).""" + raise NotImplementedError + + @abc.abstractmethod + def __call__( + self, + mX: cute.Tensor, + mO_row: Optional[cute.Tensor], + mO_col: Optional[cute.Tensor], + mS_row: cute.Tensor, + mS_col: cute.Tensor, + mOffsets: Optional[cute.Tensor], + mFirstDims: Optional[cute.Tensor], + mLastDims: Optional[cute.Tensor], + mTensormaps: Optional[cute.Tensor], + num_tensors: Int32, + mNoop: cute.Pointer, + mActInput: Optional[cute.Tensor], + mWorkspace: Optional[cute.Tensor], + stream: CUstream, + ): + """Compiled kernel entrypoint (decorate with @cute.jit), with the argument order the C++ + bridge (group_quantize_mxfp8_cutedsl.cuh) calls the compiled function with.""" + + class MXFP8GroupQuantizeEntry: - """Pick the kernel for a grouped config at compile time. - - Plain casts (no dbias/act/dact) whose scales need no per-member addressing are run as a - single-tensor cast by the single-tensor entry, which dispatches to the GMEM-to-RMEM - specialized kernels: - - * "stacked": SAME_BOTH_DIMS / VARYING_FIRST_DIM members start on 128-row boundaries and - share the column count, so unless colwise scales are GEMM-swizzled (whose indices - restart per member) the group is byte-for-byte the stacked (sum(rows), cols) tensor. - * "flat_rowwise": with packed (non-swizzled) rowwise scales only, scale i covers elements - 32i..32i+31 of the flat payload whatever the member shapes, as long as every member's - cols are a multiple of 128 (which both grouped kernels require for VARYING_LAST_DIM / - VARYING_BOTH_DIMS). Any group is then a cast of a (total / 128, 128) tensor. - - * "grouped_g2r": the remaining BF16 plain casts (colwise / both directions of - VARYING_LAST_DIM / VARYING_BOTH_DIMS, and swizzled colwise of every representation) run - MXFP8GroupQuantizeG2RKernel, which addresses each member on its own. - - Everything else, and any call these paths cannot take at runtime (a cast-noop flag, - 2^31+ elements, cols not a multiple of 128, unaligned buffers), runs - MXFP8GroupQuantizeKernel. + """Select the grouped MXFP8 quantization kernel from the config and the runtime arguments. + + As MXFP8QuantizeEntry does for single tensors: every kernel that supports the config is + instantiated at compile time, and each call runs the first one, in this priority order, + whose runtime check passes: + + 1. A non-grouped GMEM-to-RMEM kernel on a single-tensor view of the group + (MXFP8GroupQuantizeSingleTensorViewKernel; 2-lane rowwise, register bidimensional, then + 1-lane rowwise), where the group's data and scales are laid out as one tensor's. + 2. The grouped GMEM-to-RMEM kernel (MXFP8GroupQuantizeG2RKernel), which addresses each + member on its own: BF16 plain casts with no single-tensor view (GEMM-swizzled colwise + scales, colwise passes of VARYING_LAST_DIM / VARYING_BOTH_DIMS, ...) or whose + non-grouped kernels cannot take the call (colwise only). + 3. Otherwise the generic grouped kernel (MXFP8GroupQuantizeKernel), which handles every + config: dbias, fused (d)activations, FP16/FP32 cases the views do not cover, a set + cast-noop flag, 2^31+ elements, and shapes or alignments the others reject. """ def __init__(self, cfg: MXFP8GroupQuantizeConfig, sm_count: int) -> None: + # 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_g2r_rowwise_2lane import MXFP8QuantizeRowwise2LaneG2RKernel + from .quantize_mxfp8_g2r_bidim import MXFP8QuantizeBidimensionalG2RKernel + from .quantize_mxfp8_g2r_rowwise_1lane import MXFP8QuantizeRowwise1LaneG2RKernel + from .group_quantize_mxfp8_single_tensor_view import ( + MXFP8GroupQuantizeSingleTensorViewKernel, + ) + from .group_quantize_mxfp8_g2r import MXFP8GroupQuantizeG2RKernel from .group_quantize_mxfp8_general import MXFP8GroupQuantizeKernel + # pylint: enable=import-outside-toplevel + self.cfg = cfg - self.grouped_kernel = MXFP8GroupQuantizeKernel(cfg, sm_count) - plain_cast = not (cfg.WITH_DBIAS or cfg.WITH_ACT or cfg.WITH_DACT) - self.SINGLE_TENSOR_MODE = None - if plain_cast and cfg.IS_SINGLE_TENSOR: - if not (cfg.COLWISE and cfg.WITH_GEMM_SWIZZLED_SCALES): - self.SINGLE_TENSOR_MODE = "stacked" - elif plain_cast and cfg.ROWWISE and not cfg.COLWISE and not cfg.WITH_GEMM_SWIZZLED_SCALES: - self.SINGLE_TENSOR_MODE = "flat_rowwise" - self.single_tensor_entry = None - self.g2r_kernel = None - if self.SINGLE_TENSOR_MODE is None: - # pylint: disable=import-outside-toplevel - from .group_quantize_mxfp8_g2r import MXFP8GroupQuantizeG2RKernel - - if MXFP8GroupQuantizeG2RKernel.supports(cfg): - self.g2r_kernel = MXFP8GroupQuantizeG2RKernel(cfg, sm_count) - if self.SINGLE_TENSOR_MODE is not None: - # pylint: disable=import-outside-toplevel - from .quantize_mxfp8_common import MXFP8QuantizeConfig, MXFP8QuantizeEntry - - self.single_tensor_entry = MXFP8QuantizeEntry( - MXFP8QuantizeConfig( - dtype=cfg.DTYPE_STR, - fp8_dtype=cfg.FP8_DTYPE_STR, - rowwise=cfg.ROWWISE, - colwise=cfg.COLWISE, - with_gemm_swizzled_scales=cfg.WITH_GEMM_SWIZZLED_SCALES, - with_amax=False, - activation="none", - ) + # Candidate kernels in dispatch priority order; the first eligible one runs. + self.kernels = [ + MXFP8GroupQuantizeSingleTensorViewKernel(cfg, kernel_cls) + for kernel_cls in ( + MXFP8QuantizeRowwise2LaneG2RKernel, + MXFP8QuantizeBidimensionalG2RKernel, + MXFP8QuantizeRowwise1LaneG2RKernel, ) + if MXFP8GroupQuantizeSingleTensorViewKernel.supports(cfg, kernel_cls) + ] + if MXFP8GroupQuantizeG2RKernel.supports(cfg): + self.kernels.append(MXFP8GroupQuantizeG2RKernel(cfg, sm_count)) + # Handles every config and call; runs when no candidate above is eligible. + self.general_kernel = MXFP8GroupQuantizeKernel(cfg, sm_count) @cute.jit def __call__( @@ -231,7 +268,7 @@ def __call__( mWorkspace: Optional[cute.Tensor], stream: CUstream, ) -> None: - grouped_args = ( + args = ( mX, mO_row, mO_col, @@ -247,40 +284,10 @@ def __call__( mWorkspace, stream, ) - # The single-tensor kernels index with Int32, so groups of 2^31 or more elements - # stay on the grouped kernel. - num_elems = Int64(mX.shape[0]) * Int64(mX.shape[1]) - if cutlass.const_expr(self.SINGLE_TENSOR_MODE == "stacked"): - # The grouped kernels zero the scale padding and the single-tensor kernels leave it - # untouched, so stack only when there is none: rows are already 128-aligned, and - # N % 128 == 0 leaves no padded rowwise (N / 32 % 4) or colwise (N % 128) columns. - if mX.shape[1] % 128 == 0 and num_elems < Int64(2**31): - # Restate the divisibility the Int32 casts drop -- M by the compiled sym, N by - # the check above -- so the vectorized GMEM copies can prove their alignment. - M = cute.assume(Int32(mX.shape[0]), divby=128) - N = cute.assume(Int32(mX.shape[1]), divby=128) - self._run_single_tensor(mX, mO_row, mO_col, mS_row, mS_col, mNoop, stream, M, N) - else: - self.grouped_kernel(*grouped_args) - elif cutlass.const_expr(self.SINGLE_TENSOR_MODE == "flat_rowwise"): - if num_elems < Int64(2**31): - M = Int32(num_elems // 128) - self._run_single_tensor(mX, mO_row, mO_col, mS_row, mS_col, mNoop, stream, M, 128) - else: - self.grouped_kernel(*grouped_args) - elif cutlass.const_expr(self.g2r_kernel is not None): - eligible = ( - num_elems < Int64(2**31) - and mNoop.toint() == 0 - and mX.iterator.toint() % 32 == 0 - ) - if cutlass.const_expr(self.cfg.COLWISE): - eligible = eligible and mS_col.iterator.toint() % 16 == 0 - if cutlass.const_expr(self.cfg.IS_SINGLE_TENSOR): - # VARYING_LAST_DIM / VARYING_BOTH_DIMS members have 128-aligned cols already. - eligible = eligible and mX.shape[1] % 128 == 0 - if eligible: - self.g2r_kernel( + dispatched = False + for kernel in self.kernels: + if not dispatched: + if kernel.eligible_to_dispatch( mX, mO_row, mO_col, @@ -290,42 +297,12 @@ def __call__( mFirstDims, mLastDims, num_tensors, - stream, - ) - else: - self.grouped_kernel(*grouped_args) - else: - self.grouped_kernel(*grouped_args) - - @cute.jit - def _run_single_tensor( - self, - mX: cute.Tensor, - mO_row: Optional[cute.Tensor], - mO_col: Optional[cute.Tensor], - mS_row: cute.Tensor, - mS_col: cute.Tensor, - mNoop: cute.Pointer, - stream: CUstream, - M: Int32, - N: Int32, - ) -> None: - """Quantize the group's flat payload as one (M, N) tensor with the single-tensor entry. - - The single-tensor kernels are compiled for Int32 extents. The scales arrive as flat - buffers; view them at the 2D shapes TE allocates for a single (M, N) tensor, which - carry no padding here (M % 128 == N % 128 == 0). - """ - data_layout = cute.make_layout((M, N), stride=(N, 1)) - x = cute.make_tensor(mX.iterator, data_layout) - o_row = s_row = o_col = s_col = None - if cutlass.const_expr(self.cfg.ROWWISE): - o_row = cute.make_tensor(mO_row.iterator, data_layout) - s_row = cute.make_tensor(mS_row.iterator, cute.make_layout((M, N // 32), stride=(N // 32, 1))) - if cutlass.const_expr(self.cfg.COLWISE): - o_col = cute.make_tensor(mO_col.iterator, data_layout) - s_col = cute.make_tensor(mS_col.iterator, cute.make_layout((M // 32, N), stride=(N, 1))) - self.single_tensor_entry(x, o_row, s_row, o_col, s_col, None, mNoop, None, None, stream) + mNoop, + ): + dispatched = True + kernel(*args) + if not dispatched: + self.general_kernel(*args) def compile_cutedsl_function_from_cfg(cfg: MXFP8GroupQuantizeConfig): diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8_g2r.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8_g2r.py index d8941f1f13b..e776c1bc55b 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8_g2r.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8_g2r.py @@ -21,7 +21,7 @@ import cutlass from cutlass import cute -from cutlass import Int32, Int64 +from cutlass import Boolean, Int32, Int64 from cuda.bindings.driver import CUstream # pylint: disable=no-name-in-module from transformer_engine.common.CuTeDSL.utils import create_l2_policy @@ -32,37 +32,60 @@ from .quantize_mxfp8_g2r_bidim import quantize_mxfp8_gmem_tile_bidimensional from .group_quantize_mxfp8_common import ( MXFP8GroupQuantizeConfig, + MXFP8GroupQuantizeKernelBase, SAME_BOTH_DIMS, VARYING_FIRST_DIM, VARYING_LAST_DIM, ) -class MXFP8GroupQuantizeG2RKernel: +class MXFP8GroupQuantizeG2RKernel(MXFP8GroupQuantizeKernelBase): """Grid of (CTAs per member, num_tensors); each CTA grid-strides its member's tiles.""" - _CTA_THREADS_Y = 8 - # Eight lanes of 16 BF16 values (one 256-bit load each) per tile row: 1-8% faster than - # sixteen lanes of 8 values (128-bit loads) on GB200. - _CTA_THREADS_X = 8 - _ELEMENTS_X_PER_THREAD = 16 - _ELEMENTS_Y_PER_THREAD = MXFP8_BLOCK_SCALING_SIZE // _CTA_THREADS_Y # 4 - TILE_ROWS = _CTA_THREADS_Y * _ELEMENTS_Y_PER_THREAD # 32 - TILE_COLS = _CTA_THREADS_X * _ELEMENTS_X_PER_THREAD # 128 - THREADS_PER_CTA = _CTA_THREADS_Y * _CTA_THREADS_X # 64 - # CTAs per SM the grid is sized for, split evenly across the members. Far more than are - # resident: on GB200 the loop's latency is better hidden by more CTAs in flight than by - # each CTA walking more tiles (16 -> 256 per SM took colwise from ~107 to ~92 us at - # 32768x7168). - WORKERS_PER_SM = 256 - - def __init__(self, cfg: MXFP8GroupQuantizeConfig, sm_count: int) -> None: + def __init__( + self, + cfg: MXFP8GroupQuantizeConfig, + sm_count: int, + CTA_THREADS_Y: int = 8, + CTA_THREADS_X: int = 8, + ELEMENTS_X_PER_THREAD: int = 16, + WORKERS_PER_SM: int = 256, + INPUT_EVICT_LAST: bool = True, + ) -> None: + """Tunable launch parameters (defaults are the current choice): + + CTA_THREADS_Y x CTA_THREADS_X threads per CTA; each thread covers + 32 / CTA_THREADS_Y rows of ELEMENTS_X_PER_THREAD columns, so a tile is 32 rows by + CTA_THREADS_X * ELEMENTS_X_PER_THREAD columns (TILE_COLS; every member's cols must be + a multiple of it). Eight lanes of 16 BF16 values (one 256-bit load each) were 1-8% + faster than sixteen lanes of 8 on GB200. + + WORKERS_PER_SM: CTAs per SM the grid is sized for, split evenly across the members. + Far more than are resident: on GB200 the loop's latency is better hidden by more + CTAs in flight than by each CTA walking more tiles (16 -> 256 per SM took colwise + from ~107 to ~92 us at 32768x7168). + + INPUT_EVICT_LAST: L2 policy of the input loads (outputs always use evict_last). + """ assert self.supports(cfg), f"unsupported config {cfg}" + assert MXFP8_BLOCK_SCALING_SIZE % CTA_THREADS_Y == 0 + assert MXFP8_BLOCK_SCALING_SIZE % ELEMENTS_X_PER_THREAD == 0 self.cfg = cfg self.SM_COUNT = sm_count + self._CTA_THREADS_Y = CTA_THREADS_Y + self._CTA_THREADS_X = CTA_THREADS_X + self._ELEMENTS_X_PER_THREAD = ELEMENTS_X_PER_THREAD + self._ELEMENTS_Y_PER_THREAD = MXFP8_BLOCK_SCALING_SIZE // CTA_THREADS_Y + self.TILE_ROWS = CTA_THREADS_Y * self._ELEMENTS_Y_PER_THREAD # 32 + self.TILE_COLS = CTA_THREADS_X * ELEMENTS_X_PER_THREAD + self.THREADS_PER_CTA = CTA_THREADS_Y * CTA_THREADS_X + self.WORKERS_PER_SM = WORKERS_PER_SM + self.INPUT_EVICT_LAST = INPUT_EVICT_LAST self.ROWS_VARY = cfg.SHAPE_REP not in (SAME_BOTH_DIMS, VARYING_LAST_DIM) self.COLS_VARY = cfg.SHAPE_REP not in (SAME_BOTH_DIMS, VARYING_FIRST_DIM) self.SAME_BOTH_DIMS = cfg.SHAPE_REP == SAME_BOTH_DIMS + # Members of VARYING_LAST_DIM / VARYING_BOTH_DIMS only have 128-aligned cols. + assert not self.COLS_VARY or self.TILE_COLS == 128, "varying cols need 128-wide tiles" @staticmethod def supports(cfg: MXFP8GroupQuantizeConfig) -> bool: @@ -70,6 +93,36 @@ def supports(cfg: MXFP8GroupQuantizeConfig) -> bool: plain_cast = not (cfg.WITH_DBIAS or cfg.WITH_ACT or cfg.WITH_DACT) return plain_cast and cfg.DTYPE is cutlass.BFloat16 + @cute.jit + def eligible_to_dispatch( + self, + mX: cute.Tensor, + mO_row: Optional[cute.Tensor], + mO_col: Optional[cute.Tensor], + mS_row: cute.Tensor, + mS_col: cute.Tensor, + mOffsets: Optional[cute.Tensor], + mFirstDims: Optional[cute.Tensor], + mLastDims: Optional[cute.Tensor], + num_tensors: Int32, + mNoop: cute.Pointer, + ) -> Boolean: + # No cast-noop support (as the non-grouped GMEM-to-RMEM kernels); member-local Int32 + # indexing; 256-bit input loads need a 32-byte aligned payload and the vectorized + # colwise scale stores a 16-byte aligned buffer. + eligible = ( + Int64(mX.shape[0]) * Int64(mX.shape[1]) < Int64(2**31) + and mNoop.toint() == 0 + and mX.iterator.toint() % 32 == 0 + ) + if cutlass.const_expr(self.cfg.COLWISE): + eligible = eligible and mS_col.iterator.toint() % 16 == 0 + if cutlass.const_expr(not self.COLS_VARY): + # The common cols must tile exactly; VARYING_LAST_DIM / VARYING_BOTH_DIMS members + # are 128-aligned by the grouped kernels' contract and use 128-wide tiles. + eligible = eligible and mX.shape[1] % self.TILE_COLS == 0 + return eligible + @cute.jit def __call__( self, @@ -81,7 +134,11 @@ def __call__( mOffsets: Optional[cute.Tensor], mFirstDims: Optional[cute.Tensor], mLastDims: Optional[cute.Tensor], + mTensormaps: Optional[cute.Tensor], num_tensors: Int32, + mNoop: cute.Pointer, + mActInput: Optional[cute.Tensor], + mWorkspace: Optional[cute.Tensor], stream: CUstream, ) -> None: # Spread the persistent worker budget across the members, but never launch more CTAs @@ -173,6 +230,7 @@ def kernel( ) cache_policy = create_l2_policy(True) + input_policy = create_l2_policy(self.INPUT_EVICT_LAST) tiles_X = N // self.TILE_COLS num_tiles = (M // self.TILE_ROWS) * tiles_X # Grid-stride over this member's tiles (column tiles fastest, as the single-tensor @@ -180,11 +238,15 @@ def kernel( tile = Int32(bidx) while tile < num_tiles: coord = (tile // tiles_X, tile % tiles_X) - self._quantize_tile(gX, gO_row, gS_row, gO_col, gS_col, coord, cache_policy) + self._quantize_tile( + gX, gO_row, gS_row, gO_col, gS_col, coord, input_policy, cache_policy + ) tile += Int32(gdx) @cute.jit - def _quantize_tile(self, gX, gO_row, gS_row, gO_col, gS_col, coord, cache_policy): + def _quantize_tile( + self, gX, gO_row, gS_row, gO_col, gS_col, coord, input_policy, cache_policy + ): """Quantize the (TILE_ROWS, TILE_COLS) tile at `coord` of one member.""" cfg = self.cfg TILER = (self.TILE_ROWS, self.TILE_COLS) @@ -217,4 +279,5 @@ def _quantize_tile(self, gX, gO_row, gS_row, gO_col, gS_col, coord, cache_policy ELEMENTS_Y_PER_THREAD=self._ELEMENTS_Y_PER_THREAD, ELEMENTS_X_PER_THREAD=self._ELEMENTS_X_PER_THREAD, cache_policy=cache_policy, + load_cache_policy=input_policy, ) diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8_general.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8_general.py index 0bfc0f9688d..8604f6b4318 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8_general.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8_general.py @@ -50,10 +50,11 @@ COLS_OFFSET, OFFSETS_OFFSET, MXFP8GroupQuantizeConfig, + MXFP8GroupQuantizeKernelBase, ) -class MXFP8GroupQuantizeKernel: +class MXFP8GroupQuantizeKernel(MXFP8GroupQuantizeKernelBase): """Grouped MXFP8 quantize mirroring group_quantize_mxfp8_kernel's strategy.""" # Target persistent CTA count per SM for sizing the grid @@ -70,6 +71,28 @@ class MXFP8GroupQuantizeKernel: THREADS_X = TILE_COLS // MXFP8_BLOCK_SCALING_SIZE # 4 THREADS_Y = THREADS_PER_CTA // THREADS_X # 32 + @staticmethod + def supports(cfg: MXFP8GroupQuantizeConfig) -> bool: # pylint: disable=unused-argument + """Every grouped config.""" + return True + + @cute.jit + def eligible_to_dispatch( + self, + mX: cute.Tensor, + mO_row: Optional[cute.Tensor], + mO_col: Optional[cute.Tensor], + mS_row: cute.Tensor, + mS_col: cute.Tensor, + mOffsets: Optional[cute.Tensor], + mFirstDims: Optional[cute.Tensor], + mLastDims: Optional[cute.Tensor], + num_tensors: Int32, + mNoop: cute.Pointer, + ) -> Boolean: + """Every call.""" + return Boolean(True) + def __init__(self, cfg: MXFP8GroupQuantizeConfig, SM_COUNT: int) -> None: self.cfg = cfg self.SM_COUNT = SM_COUNT diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8_single_tensor_view.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8_single_tensor_view.py new file mode 100644 index 00000000000..1121b49af42 --- /dev/null +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8_single_tensor_view.py @@ -0,0 +1,160 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Grouped MXFP8 quantization by a non-grouped kernel on a single-tensor view of the group. + +When the group's data and scales are laid out exactly as one tensor's, a plain cast of the +group is byte-for-byte a single-tensor cast, so the non-grouped kernels apply unchanged: + +* "stacked": SAME_BOTH_DIMS / VARYING_FIRST_DIM members start on 128-row boundaries and share + the column count, so unless colwise scales are GEMM-swizzled (whose indices restart per + member) the group is the stacked (sum(rows), cols) tensor. +* "flat": with packed (non-swizzled) rowwise scales only, scale i covers elements 32i..32i+31 + of the flat payload whatever the member shapes, as long as every member's cols are a + multiple of 128 (which both grouped kernels require for VARYING_LAST_DIM / + VARYING_BOTH_DIMS). Any group is then a cast of a (total / 128, 128) tensor. + +The grouped kernels zero the scale padding and the non-grouped ones leave it untouched, so the +views are only taken when there is none. +""" + +from typing import Optional, Type + +import cutlass +from cutlass import cute +from cutlass import Boolean, Int32, Int64 +from cuda.bindings.driver import CUstream # pylint: disable=no-name-in-module + +from .quantize_mxfp8_common import MXFP8QuantizeConfig, MXFP8QuantizeKernelBase +from .group_quantize_mxfp8_common import MXFP8GroupQuantizeConfig, MXFP8GroupQuantizeKernelBase + +STACKED = "stacked" +FLAT = "flat" + + +class MXFP8GroupQuantizeSingleTensorViewKernel(MXFP8GroupQuantizeKernelBase): + """Run one non-grouped MXFP8 quantize kernel on the stacked or flat view of the group.""" + + def __init__( + self, cfg: MXFP8GroupQuantizeConfig, kernel_cls: Type[MXFP8QuantizeKernelBase] + ) -> None: + assert self.supports(cfg, kernel_cls), f"{kernel_cls.__name__} cannot run {cfg}" + self.cfg = cfg + self.VIEW = self.view(cfg) + self.single_cfg = self.single_tensor_cfg(cfg) + self.kernel = kernel_cls(self.single_cfg) + + @staticmethod + def view(cfg: MXFP8GroupQuantizeConfig) -> Optional[str]: + """The single-tensor view the group's plain cast has, if any.""" + if cfg.WITH_DBIAS or cfg.WITH_ACT or cfg.WITH_DACT: + return None + if cfg.IS_SINGLE_TENSOR and not (cfg.COLWISE and cfg.WITH_GEMM_SWIZZLED_SCALES): + return STACKED + if cfg.ROWWISE and not cfg.COLWISE and not cfg.WITH_GEMM_SWIZZLED_SCALES: + return FLAT + return None + + @staticmethod + def single_tensor_cfg(cfg: MXFP8GroupQuantizeConfig) -> MXFP8QuantizeConfig: + """The non-grouped config of the same plain cast.""" + return MXFP8QuantizeConfig( + dtype=cfg.DTYPE_STR, + fp8_dtype=cfg.FP8_DTYPE_STR, + rowwise=cfg.ROWWISE, + colwise=cfg.COLWISE, + with_gemm_swizzled_scales=cfg.WITH_GEMM_SWIZZLED_SCALES, + with_amax=False, + activation="none", + ) + + # pylint: disable-next=arguments-differ + @staticmethod + def supports( + cfg: MXFP8GroupQuantizeConfig, kernel_cls: Type[MXFP8QuantizeKernelBase] + ) -> bool: + view = MXFP8GroupQuantizeSingleTensorViewKernel.view(cfg) + return view is not None and kernel_cls.supports( + MXFP8GroupQuantizeSingleTensorViewKernel.single_tensor_cfg(cfg) + ) + + @cute.jit + def _views( + self, + mX: cute.Tensor, + mO_row: Optional[cute.Tensor], + mO_col: Optional[cute.Tensor], + mS_row: cute.Tensor, + mS_col: cute.Tensor, + ): + """(x, o_row, s_row, o_col, s_col) of the single-tensor view, in the non-grouped + kernels' argument order. They are compiled for Int32 extents; the scales arrive as flat + buffers and are viewed at the 2D shapes TE allocates for a single (M, N) tensor.""" + if cutlass.const_expr(self.VIEW == STACKED): + # Restate the divisibility the Int32 casts drop -- M by the compiled sym, N by the + # eligibility check -- so the vectorized GMEM copies can prove their alignment. + M = cute.assume(Int32(mX.shape[0]), divby=128) + N = cute.assume(Int32(mX.shape[1]), divby=128) + else: + M = Int32(Int64(mX.shape[0]) * Int64(mX.shape[1]) // 128) + N = 128 + data_layout = cute.make_layout((M, N), stride=(N, 1)) + x = cute.make_tensor(mX.iterator, data_layout) + o_row = s_row = o_col = s_col = None + if cutlass.const_expr(self.cfg.ROWWISE): + o_row = cute.make_tensor(mO_row.iterator, data_layout) + s_row = cute.make_tensor( + mS_row.iterator, cute.make_layout((M, N // 32), stride=(N // 32, 1)) + ) + if cutlass.const_expr(self.cfg.COLWISE): + o_col = cute.make_tensor(mO_col.iterator, data_layout) + s_col = cute.make_tensor(mS_col.iterator, cute.make_layout((M // 32, N), stride=(N, 1))) + return x, o_row, s_row, o_col, s_col + + @cute.jit + def eligible_to_dispatch( + self, + mX: cute.Tensor, + mO_row: Optional[cute.Tensor], + mO_col: Optional[cute.Tensor], + mS_row: cute.Tensor, + mS_col: cute.Tensor, + mOffsets: Optional[cute.Tensor], + mFirstDims: Optional[cute.Tensor], + mLastDims: Optional[cute.Tensor], + num_tensors: Int32, + mNoop: cute.Pointer, + ) -> Boolean: + # The non-grouped kernels index with Int32, so groups of 2^31 or more elements cannot + # be viewed as one tensor. + view_ok = Int64(mX.shape[0]) * Int64(mX.shape[1]) < Int64(2**31) + if cutlass.const_expr(self.VIEW == STACKED): + # No scale padding: rows are already 128-aligned, and N % 128 == 0 leaves no padded + # rowwise (N / 32 % 4) or colwise (N % 128) columns. + view_ok = view_ok and mX.shape[1] % 128 == 0 + x, o_row, s_row, o_col, s_col = self._views(mX, mO_row, mO_col, mS_row, mS_col) + return view_ok and self.kernel.eligible_to_dispatch( + self.single_cfg, x, o_row, s_row, o_col, s_col, mNoop, None + ) + + @cute.jit + def __call__( + self, + mX: cute.Tensor, + mO_row: Optional[cute.Tensor], + mO_col: Optional[cute.Tensor], + mS_row: cute.Tensor, + mS_col: cute.Tensor, + mOffsets: Optional[cute.Tensor], + mFirstDims: Optional[cute.Tensor], + mLastDims: Optional[cute.Tensor], + mTensormaps: Optional[cute.Tensor], + num_tensors: Int32, + mNoop: cute.Pointer, + mActInput: Optional[cute.Tensor], + mWorkspace: Optional[cute.Tensor], + stream: CUstream, + ) -> None: + x, o_row, s_row, o_col, s_col = self._views(mX, mO_row, mO_col, mS_row, mS_col) + self.kernel(x, o_row, s_row, o_col, s_col, None, mNoop, None, None, stream) diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidim.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidim.py index 61598a7d37b..855a93e2ab1 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidim.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidim.py @@ -413,9 +413,13 @@ def quantize_mxfp8_gmem_tile_bidimensional( ELEMENTS_Y_PER_THREAD: cutlass.Constexpr[int], ELEMENTS_X_PER_THREAD: cutlass.Constexpr[int], cache_policy: Int64, + load_cache_policy: Optional[Int64] = None, ): """Quantize one (TILE_ROWS, TILE_COLS) GMEM tile in the directions enabled by ROWWISE / COLWISE. + cache_policy is the L2 policy of the stores, and of the loads unless load_cache_policy is + given. + 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. @@ -468,7 +472,10 @@ def quantize_mxfp8_gmem_tile_bidimensional( # 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 + load_atom, + gX_thread[(None, row_idx),], + rX[None, row_idx], + cache_policy=cache_policy if load_cache_policy is None else load_cache_policy, ) if cutlass.const_expr(WITH_ACT or WITH_DACT): From 253928908aa100af9072bd22351cfac50a940d67 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 9 Nov 2026 21:46:45 +0000 Subject: [PATCH 54/58] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8_g2r.py | 4 +--- .../common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8_general.py | 2 +- .../cast/mxfp8/group_quantize_mxfp8_single_tensor_view.py | 4 +--- 3 files changed, 3 insertions(+), 7 deletions(-) diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8_g2r.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8_g2r.py index e776c1bc55b..b55594e3e17 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8_g2r.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8_g2r.py @@ -244,9 +244,7 @@ def kernel( tile += Int32(gdx) @cute.jit - def _quantize_tile( - self, gX, gO_row, gS_row, gO_col, gS_col, coord, input_policy, cache_policy - ): + def _quantize_tile(self, gX, gO_row, gS_row, gO_col, gS_col, coord, input_policy, cache_policy): """Quantize the (TILE_ROWS, TILE_COLS) tile at `coord` of one member.""" cfg = self.cfg TILER = (self.TILE_ROWS, self.TILE_COLS) diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8_general.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8_general.py index 8604f6b4318..569e4023522 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8_general.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8_general.py @@ -801,7 +801,7 @@ def _kernel_main( # If SAME_BOTH_DIMS, the tensor member's start offset and rows can be obtained with division member_rows = Int64(first_logical_dim) // Int64(num_tensors) member_start_row = job_start_row - job_start_row % member_rows - else: # VARYING_FIRST_DIM + else: # VARYING_FIRST_DIM # Otherwise if VARYING_FIRST_DIM, we need to use find these tensor member metadata # from mFirstDims & mOffsets member_id = self._find_tensor_from_offsets( diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8_single_tensor_view.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8_single_tensor_view.py index 1121b49af42..dffe8daee4f 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8_single_tensor_view.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8_single_tensor_view.py @@ -71,9 +71,7 @@ def single_tensor_cfg(cfg: MXFP8GroupQuantizeConfig) -> MXFP8QuantizeConfig: # pylint: disable-next=arguments-differ @staticmethod - def supports( - cfg: MXFP8GroupQuantizeConfig, kernel_cls: Type[MXFP8QuantizeKernelBase] - ) -> bool: + def supports(cfg: MXFP8GroupQuantizeConfig, kernel_cls: Type[MXFP8QuantizeKernelBase]) -> bool: view = MXFP8GroupQuantizeSingleTensorViewKernel.view(cfg) return view is not None and kernel_cls.supports( MXFP8GroupQuantizeSingleTensorViewKernel.single_tensor_cfg(cfg) From 303ef6e22d3a2f8b38d912c14300ffe1ff9f52e1 Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Tue, 10 Nov 2026 00:55:05 +0000 Subject: [PATCH 55/58] fix zeros Signed-off-by: Kaining Zhong --- tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py b/tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py index 45381ed6652..c5aa31ab5fe 100644 --- a/tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py +++ b/tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py @@ -443,14 +443,18 @@ def extract_group_quantized_output(out, rowwise, columnwise): """Extract the bytes to compare between backends. Unlike the single-tensor kernels, both grouped backends zero the scale padding, so the - whole data and scale buffers are compared. + whole data and scale buffers are compared. + Note: we mask the data to fold -0.0 (0x80) into +0.0, because the CuTeDSL implementation + differs from the CUDA implementation in some cases, but both should be correct. """ parts = {} if rowwise: - parts["rowwise data"] = out.rowwise_data.view(torch.uint8).clone() + d = out.rowwise_data.view(torch.uint8) + parts["rowwise data"] = d.masked_fill(d == 0x80, 0) parts["rowwise scales"] = out.scale_inv.view(torch.uint8).clone() if columnwise: - parts["colwise data"] = out.columnwise_data.view(torch.uint8).clone() + d = out.columnwise_data.view(torch.uint8) + parts["colwise data"] = d.masked_fill(d == 0x80, 0) parts["colwise scales"] = out.columnwise_scale_inv.view(torch.uint8).clone() return parts From cd185b8272f67071e9159be1e4edea6d9cece120 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 10 Nov 2026 00:57:29 +0000 Subject: [PATCH 56/58] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py b/tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py index c5aa31ab5fe..05103640285 100644 --- a/tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py +++ b/tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py @@ -443,7 +443,7 @@ def extract_group_quantized_output(out, rowwise, columnwise): """Extract the bytes to compare between backends. Unlike the single-tensor kernels, both grouped backends zero the scale padding, so the - whole data and scale buffers are compared. + whole data and scale buffers are compared. Note: we mask the data to fold -0.0 (0x80) into +0.0, because the CuTeDSL implementation differs from the CUDA implementation in some cases, but both should be correct. """ From 2cc9417b2e8d75daecc280b4e9b362013ddd62fe Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Tue, 10 Nov 2026 00:57:34 +0000 Subject: [PATCH 57/58] nit Signed-off-by: Kaining Zhong --- .../common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8_general.py | 1 - 1 file changed, 1 deletion(-) diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8_general.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8_general.py index 569e4023522..3374400662e 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8_general.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8_general.py @@ -654,7 +654,6 @@ def _kernel_main( tma_dst_out_col: Optional[cute.Tensor], ) -> None: cfg = self.cfg - FP8_DTYPE = cfg.FP8_DTYPE tidx, _, _ = cute.arch.thread_idx() bidx, bidy, _ = cute.arch.block_idx() gdx, _, _ = cute.arch.grid_dim() From bafb9b4bdb4750c19440d34e3d6c93e4dcfeed9e Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Tue, 10 Nov 2026 01:10:58 +0000 Subject: [PATCH 58/58] nit Signed-off-by: Kaining Zhong --- .../common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8_g2r.py | 1 + 1 file changed, 1 insertion(+) diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8_g2r.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8_g2r.py index b55594e3e17..6e89109b986 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8_g2r.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/group_quantize_mxfp8_g2r.py @@ -182,6 +182,7 @@ def kernel( last_logical_dim: Int64, num_tensors: Int32, ): + """Quantize the tiles of group member blockIdx.y, grid-striding over blockIdx.x.""" cfg = self.cfg bidx, member_id, _ = cute.arch.block_idx() gdx, _, _ = cute.arch.grid_dim()