diff --git a/pyproject.toml b/pyproject.toml index 2c9f224c14..924ca49a28 100755 --- a/pyproject.toml +++ b/pyproject.toml @@ -3,7 +3,7 @@ # See LICENSE for license information. [build-system] -requires = ["setuptools>=61.0", "cmake>=3.21", "wheel", "pybind11[global]", "ninja", "pip", "torch>=2.1", "jax>=0.5.0", "flax>=0.7.1", "nvidia-cudnn-frontend>=1.25.0"] +requires = ["setuptools>=61.0", "cmake>=3.21", "wheel", "pybind11[global]", "ninja", "pip", "torch>=2.1", "jax>=0.5.0", "flax>=0.7.1", "nvidia-cudnn-frontend>=1.25.0", "apache-tvm-ffi>=0.1.12"] # Use legacy backend to import local packages in setup.py build-backend = "setuptools.build_meta:__legacy__" diff --git a/setup.py b/setup.py index 150d92969c..cdf1acbf35 100644 --- a/setup.py +++ b/setup.py @@ -122,6 +122,8 @@ def setup_requirements() -> Tuple[List[str], List[str]]: "pydantic", "importlib-metadata>=1.0", "packaging", + "apache-tvm-ffi>=0.1.12", + "nvidia-cutlass-dsl>=4.4.2", ] test_reqs: List[str] = ["pytest>=8.2.1"] diff --git a/tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py b/tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py new file mode 100644 index 0000000000..df048661bb --- /dev/null +++ b/tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py @@ -0,0 +1,237 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Cross-backend bit-exactness tests for the CuTeDSL MXFP8 quantize kernels.""" + +import ctypes +import os +from typing import Callable, NamedTuple, Optional + +import pytest +import torch + +import transformer_engine.pytorch as te +import transformer_engine_torch as tex +import tvm_ffi +from transformer_engine.common import _get_shared_object_file +from transformer_engine.pytorch import MXFP8Quantizer + +recipe_available, reason_for_no_recipe = te.is_mxfp8_available(return_reason=True) + +# The already-loaded core lib (dlopen refcounts: this returns the same handle, +# so the call mutates the same dispatcher singleton the quantize ops read). +CORE_LIB = ctypes.CDLL(str(_get_shared_object_file("core"))) +if not hasattr(CORE_LIB, "nvte_set_cutedsl_quant_backend"): + raise RuntimeError( + "libtransformer_engine.so lacks nvte_set_cutedsl_quant_backend -- rebuild the " + "Transformer Engine core library." + ) + +pytestmark = pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) + + +class Fusion(NamedTuple): + """An ActivationType from the C++ test: its tex ops per ProcessingMethod and + the activation desc used in the CuTeDSL config key.""" + + name: str + act: Optional[Callable] # CAST_ACT: act(x, quantizer) + dact: Optional[Callable] # CAST_DACT: dact(grad, act_input, quantizer) + dbias_dact: Optional[Callable] # CAST_DBIAS_DACT: dbias_dact(grad, act_input, quantizer) + desc: str + + +# --- Case matrix, mirroring test_cast_mxfp8.cu --- +# The C++ multi-dim sizes flattened to the 2D (rows, cols) the kernels see: +# {8,32,1024} -> (256, 1024), {16,8,4,512} -> (512, 512). The C++ list also has +# non-32-divisible shapes ({1,16}, {16,48}, {993,512}, {1024}) that exercise the +# CUDA kernels' partial-block edges; the CuTeDSL backend's contract is +# 32-divisible flat dims (the dispatcher falls back to CUDA otherwise), so those +# cases are omitted here rather than vacuously comparing CUDA against CUDA. +MATRIX_SIZES = [ + (128, 128), + (256, 1024), + (512, 512), + (8192, 7168), +] +# (block_rows, block_cols): (1,32)=rowwise, (32,1)=colwise, (32,32)=both. +BLOCK_SIZES = [(1, 32), (32, 1), (32, 32)] +# Only GeLU activation tests are used (SiLU/ReLU/QGeLU/SReLU commented out +# in the C++ test as well). +IDENTITY = Fusion("Identity", None, None, None, "none") +GELU = Fusion("GeLU", tex.gelu, tex.dgelu, tex.dbias_dgelu, "gelu") +# SILU = Fusion("SiLU", tex.silu, tex.dsilu, tex.dbias_dsilu, "silu") +# RELU = Fusion("ReLU", tex.relu, tex.drelu, tex.dbias_drelu, "relu") +# QGELU = Fusion("QGeLU", tex.qgelu, tex.dqgelu, tex.dbias_dqgelu, "qgelu") +# SRELU = Fusion("SReLU", tex.srelu, tex.dsrelu, tex.dbias_dsrelu, "srelu") + +# Valid (ProcessingMethod, ActivationType) pairs. The C++ test crosses the two +# axes and GTEST_SKIPs the mismatched half; only the meaningful pairs are +# generated here. A newly enabled activation adds its ACT/DACT/DBIAS_DACT pairs. +METHOD_FUSION_CASES = [ + ("CAST_ONLY", IDENTITY), + ("CAST_DBIAS", IDENTITY), + ("CAST_ACT", GELU), + ("CAST_DACT", GELU), + ("CAST_DBIAS_DACT", GELU), +] +METHOD_FUSION_IDS = [f"{m}X{f.name}" for m, f in METHOD_FUSION_CASES] +IN_DTYPES = [torch.float32, torch.bfloat16, torch.float16] +FP8_DTYPES = [tex.DType.kFloat8E4M3, tex.DType.kFloat8E5M2] + +# Description strings for pytest case ids and the CuTeDSL config key. +DTYPE_TO_STR = {torch.float32: "fp32", torch.bfloat16: "bf16", torch.float16: "fp16"} +FP8_TO_STR = {tex.DType.kFloat8E4M3: "e4m3", tex.DType.kFloat8E5M2: "e5m2"} + +get_shape_id = lambda s: f"{s[0]}x{s[1]}" +get_block_id = lambda b: f"{b[0]}x{b[1]}" +get_dtype_id = DTYPE_TO_STR.get +get_fp8_id = FP8_TO_STR.get + + +def set_cutedsl_backend(enabled): + CORE_LIB.nvte_set_cutedsl_quant_backend(1 if enabled else 0) + + +@pytest.fixture(scope="module", autouse=True) +def _restore_backend_choice_from_env(): + """Restore the flag that decides the CuTeDSL / CUDA backend choice when this pytest module is done.""" + yield + flag = os.getenv("NVTE_ENABLE_CUTEDSL_QUANT_BACKEND") + set_cutedsl_backend(flag is not None and not flag.startswith("0")) + + +def generate_inputs(M, N, in_dtype, seed=0): + g = torch.Generator(device="cuda").manual_seed(seed) + x = torch.empty(M, N, dtype=in_dtype, device="cuda").uniform_(-2.0, 1.0, generator=g) + ain = torch.empty(M, N, dtype=in_dtype, device="cuda").uniform_(-2.0, 1.0, generator=g) + return x, ain + + +def run_quantize(method, act, x, ain, rowwise, columnwise, fp8_dtype): + """Quantize via the public dispatch; returns (mxfp8_tensor, dbias_or_None).""" + q = MXFP8Quantizer(fp8_dtype=fp8_dtype, rowwise=rowwise, columnwise=columnwise) + if method == "CAST_ONLY": + return q(x), None + if method == "CAST_DBIAS": + db, out = tex.bgrad_quantize(x, q) + return out, db + if method == "CAST_ACT": + return act.act(x, q), None + if method == "CAST_DACT": + return act.dact(x, ain, q), None + if method == "CAST_DBIAS_DACT": + db, out = act.dbias_dact(x, ain, q) + return out, db + raise ValueError(f"unknown method {method!r}") + + +def get_cfg_key(method, act, in_dtype, fp8_dtype, rowwise, colwise): + """Mirror of MXFP8QuantConfig::to_key (quantize_mxfp8_cutedsl.cuh): the name the CuTeDSL backend registers its compiled kernel under for this config. + Used to check if the CuTeDSL implmentation is registered + """ + with_dbias = method in ("CAST_DBIAS", "CAST_DBIAS_DACT") + with_dact = method in ("CAST_DACT", "CAST_DBIAS_DACT") + with_act = method == "CAST_ACT" + desc = "none" + if with_act: + desc = act.desc + elif with_dact: + desc = f"d{act.desc}" + flags = (rowwise, colwise, False, False, with_dbias, with_dact, with_act, False) + return ( + "cutedsl_mxfp8_" + + DTYPE_TO_STR[in_dtype] + + "_" + + FP8_TO_STR[fp8_dtype] + + "_" + + "_".join("1" if f else "0" for f in flags) + + "_" + + desc + ) + + +def extract_quantized_output(out, rowwise, columnwise): + """Extract the meaningful bytes from the MXFP8Quantizer output for comparison between backends. The scale padding is uninitialized, so only the meaningful + region is compared. + """ + parts = {} + if rowwise: + d = out._rowwise_data.view(torch.uint8) + M, N = d.shape + parts["rowwise data"] = d.clone() + parts["rowwise scales"] = out._rowwise_scale_inv[: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 scales"] = out._columnwise_scale_inv[: (M + 31) // 32, :N].clone() + return parts + + +def run_test_case(method, act, shape, block_size, in_dtype, fp8_dtype): + """Assert the CuTeDSL and CUDA backends produce bit-identical outputs for the + same input and config. + """ + M, N = shape + rowwise = block_size[1] != 1 + columnwise = block_size[0] != 1 + x, act_input = generate_inputs(M, N, in_dtype) + + set_cutedsl_backend(False) + out_cuda, dbias_cuda = run_quantize(method, act, x, act_input, rowwise, columnwise, fp8_dtype) + cuda_output = extract_quantized_output(out_cuda, rowwise, columnwise) + + set_cutedsl_backend(True) + try: + out_cutedsl, dbias_cutedsl = run_quantize( + method, act, x, act_input, rowwise, columnwise, fp8_dtype + ) + cutedsl_output = extract_quantized_output(out_cutedsl, rowwise, columnwise) + finally: + set_cutedsl_backend(False) + + # Guard against a silent CUDA fallback: every config in the matrix is one the + # CuTeDSL backend supports, so its kernel must have been registered under the + # config key. If not, the backend rejected or missed the config and the + # comparison above was CUDA vs CUDA. + key = get_cfg_key(method, act, in_dtype, fp8_dtype, rowwise, columnwise) + assert tvm_ffi.get_global_func(key, allow_missing=True) is not None, ( + f"CuTeDSL kernel not registered for {key}; the CuTeDSL backend fell back " + "to CUDA and this case compared CUDA against itself" + ) + + tag = f"{method}/{act.name}/{M}x{N}/{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_cuda is not None: + torch.testing.assert_close(dbias_cutedsl, dbias_cuda) + + +# Test cases with only cast kernels (mirrors C++ test's OperatorTest_FusedCastMXFP8_CastOnly). +@pytest.mark.parametrize("shape", MATRIX_SIZES, ids=get_shape_id) +@pytest.mark.parametrize("block_size", BLOCK_SIZES, ids=get_block_id) +@pytest.mark.parametrize("in_dtype", IN_DTYPES, ids=get_dtype_id) +@pytest.mark.parametrize("fp8_dtype", FP8_DTYPES, ids=get_fp8_id) +def test_cast_only(fp8_dtype, in_dtype, block_size, shape): + run_test_case("CAST_ONLY", IDENTITY, shape, block_size, in_dtype, fp8_dtype) + + +# Test cases with varying matrix shapes and block shapes +# (OperatorTest_FusedCastMXFP8_Sizes). +@pytest.mark.parametrize("shape", MATRIX_SIZES, ids=get_shape_id) +@pytest.mark.parametrize("block_size", BLOCK_SIZES, ids=get_block_id) +@pytest.mark.parametrize("method,act", METHOD_FUSION_CASES, ids=METHOD_FUSION_IDS) +def test_sizes(method, act, block_size, shape): + run_test_case(method, act, shape, block_size, torch.bfloat16, tex.DType.kFloat8E4M3) + + +# Test cases with varying dtypes (OperatorTest_FusedCastMXFP8_Dtypes). +@pytest.mark.parametrize("in_dtype", IN_DTYPES, ids=get_dtype_id) +@pytest.mark.parametrize("fp8_dtype", FP8_DTYPES, ids=get_fp8_id) +@pytest.mark.parametrize("method,act", METHOD_FUSION_CASES, ids=METHOD_FUSION_IDS) +def test_dtypes(method, act, fp8_dtype, in_dtype): + run_test_case(method, act, (256, 384), (32, 32), in_dtype, fp8_dtype) diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index be64fcb2be..019749e09b 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -103,6 +103,23 @@ set(CUTLASS_TOOLS_INCLUDE_DIR # Python find_package(Python COMPONENTS Interpreter Development.Module REQUIRED) +# Find tvm_ffi pip package's path in build time. We need to retrieve the include path to obtain its headers +# in order to compile the CuTeDSL quantize backend bridge (common/tvm_ffi_bridge.h). +# At runtime, we will attempt to find tvm_ffi's shared library path (libtvmffi.so) via python import in common/__init__.py. +# If no libtvmffi.so is found, the CuTeDSL backend will be disabled and fall back to the default TE CUDA C++ kernels. +execute_process( + COMMAND ${Python_EXECUTABLE} -c "import tvm_ffi.libinfo as li; print(li.find_include_path())" + OUTPUT_VARIABLE TVM_FFI_INCLUDE_DIR + OUTPUT_STRIP_TRAILING_WHITESPACE + RESULT_VARIABLE TVM_FFI_INCLUDE_QUERY) +if(NOT TVM_FFI_INCLUDE_QUERY EQUAL 0) + message(FATAL_ERROR + "Could not import the tvm_ffi Python package (with '${Python_EXECUTABLE}'), " + "whose headers Transformer Engine needs to compile the CuTeDSL quantize " + "backend bridge (common/tvm_ffi_bridge.h). Install it into this Python " + "environment: `pip install apache-tvm-ffi`.") +endif() + function(find_nccl_version OUT_VERSION OUT_INCLUDE_DIR) find_path(_nvte_nccl_include_dir NAMES nccl.h @@ -357,6 +374,10 @@ target_link_libraries(transformer_engine PUBLIC CUDA::cudart CUDNN::cudnn_all) +# Include tvm-ffi headers to compile tvm_ffi_bridge.h. The tvm-ffi library is loaded at runtime via dlopen. +# If libtvmffi.so is not successfully loaded, the CuTeDSL backend is disabled and will fall back to the default TE CUDA C++ kernels. +target_include_directories(transformer_engine PRIVATE ${TVM_FFI_INCLUDE_DIR}) + target_include_directories(transformer_engine PRIVATE ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES}) target_include_directories(transformer_engine SYSTEM PRIVATE diff --git a/transformer_engine/common/CuTeDSL/__init__.py b/transformer_engine/common/CuTeDSL/__init__.py new file mode 100644 index 0000000000..50d536852c --- /dev/null +++ b/transformer_engine/common/CuTeDSL/__init__.py @@ -0,0 +1,24 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""CuTeDSL kernels for Transformer Engine. + +To expose CuTeDSL kernels to C++, they should be registered via `register_cutedsl_backends()`. +They should provide a string function name which can be used to retrieve the corresponding CuTeDSL kernel function via TVM-FFI. +""" + +import tvm_ffi + +from transformer_engine.common.CuTeDSL.cast.mxfp8.quantize_mxfp8 import ( + get_mxfp8_quantization_function, +) + + +def register_cutedsl_backends(): + """Register all available CuTeDSL backends for on-demand compilation via TVM-FFI. + The C++ dispatcher retrieves them by the names defined here. + """ + tvm_ffi.register_global_func( + "get_mxfp8_quantization_function", get_mxfp8_quantization_function, override=True + ) diff --git a/transformer_engine/common/CuTeDSL/activations.py b/transformer_engine/common/CuTeDSL/activations.py new file mode 100644 index 0000000000..d6e0e6ce87 --- /dev/null +++ b/transformer_engine/common/CuTeDSL/activations.py @@ -0,0 +1,97 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Elementwise activations and their derivatives (f32) +A CuTeDSL port of the CUDA C++ implementation in common/util/math.h.""" + +# Each function is a 1:1 port of its math.h counterpart (see module docstring). +# pylint: disable=missing-function-docstring +import os + +from cutlass import cute +from cutlass import Float32 +from cutlass._mlir.dialects import arith as mlir_arith +from cutlass.cutlass_dsl import dsl_user_op + +from transformer_engine.common.CuTeDSL.utils import fma_f32 + + +USE_FAST_MATH = os.environ.get("NVTE_BUILD_ACTIVATION_WITH_FAST_MATH", "0") == "1" + + +def act_relu(x: Float32) -> Float32: + return cute.arch.fmax(x, Float32(0.0)) + + +def act_gelu(x: Float32) -> Float32: + A = Float32(0.79788456) # sqrt(2/π) truncated to TE's 8-digit literal + B = Float32(0.03567741) # = sqrt(2/π) · 0.044715, same truncation + return x * ( + Float32(0.5) + Float32(0.5) * cute.math.tanh(x * (A + B * x * x), fastmath=USE_FAST_MATH) + ) + + +def act_silu(x: Float32) -> Float32: + return x / (Float32(1.0) + cute.math.exp(-x, fastmath=USE_FAST_MATH)) + + +def act_qgelu(x: Float32) -> Float32: + return x * sigmoid(Float32(1.702) * x) + + +def act_srelu(x: Float32) -> Float32: + r = cute.arch.fmax(x, Float32(0.0)) + return r * r + + +@dsl_user_op +def dact_drelu(x: Float32, *, loc=None, ip=None) -> Float32: + cond = mlir_arith.cmpf( + mlir_arith.CmpFPredicate.OGT, + x.ir_value(loc=loc, ip=ip), + Float32(0.0).ir_value(loc=loc, ip=ip), + loc=loc, + ip=ip, + ) + return Float32( + mlir_arith.select( + cond, + Float32(1.0).ir_value(loc=loc, ip=ip), + Float32(0.0).ir_value(loc=loc, ip=ip), + loc=loc, + ip=ip, + ) + ) + + +def dact_dsrelu(x: Float32) -> Float32: + return cute.arch.fmax(Float32(2.0) * x, Float32(0.0)) + + +def sigmoid(x: Float32) -> Float32: + return Float32(1.0) / (Float32(1.0) + cute.math.exp(-x, fastmath=USE_FAST_MATH)) + + +def dsigmoid(x: Float32) -> Float32: + s = sigmoid(x) + return s * (Float32(1.0) - s) + + +def dact_dsilu(x: Float32) -> Float32: + return fma_f32(x, dsigmoid(x), sigmoid(x)) + + +def dact_dqgelu(x: Float32) -> Float32: + ax = Float32(1.702) * x + return ax * dsigmoid(ax) + sigmoid(ax) + + +def dact_dgelu(x: Float32) -> Float32: + t = cute.math.tanh( + Float32(0.79788456) * x * (Float32(1.0) + Float32(0.044715) * x * x), + fastmath=USE_FAST_MATH, + ) + return Float32(0.5) * x * ( + (Float32(1.0) - t * t) * (Float32(0.79788456) + Float32(0.1070322243) * x * x) + ) + Float32(0.5) * (Float32(1.0) + t) diff --git a/transformer_engine/common/CuTeDSL/cast/__init__.py b/transformer_engine/common/CuTeDSL/cast/__init__.py new file mode 100644 index 0000000000..aefd0b9771 --- /dev/null +++ b/transformer_engine/common/CuTeDSL/cast/__init__.py @@ -0,0 +1,7 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""CuTeDSL cast/quantize kernels.""" + +from . import mxfp8 diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/__init__.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/__init__.py new file mode 100644 index 0000000000..6491573690 --- /dev/null +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/__init__.py @@ -0,0 +1,7 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""CuTeDSL MXFP8 quantization kernels.""" + +from . import quantize_mxfp8 diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8.py new file mode 100644 index 0000000000..8aad84de10 --- /dev/null +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8.py @@ -0,0 +1,2366 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""MXFP8 quantization kernel implemented in CuTeDSL. + +Replicates the core logic of quantize_mxfp8.cuh: given a 2D tensor of BF16/FP16 +values, quantize to MXFP8 format (FP8E4M3 data + E8M0 per-block scales). + +Matches the C++ kernel's tile dimensions and thread layout: + CHUNK_DIM_Y = 64, CHUNK_DIM_X = 64, THREADS_PER_CTA = 64 + BUFF_DIM_Y = 32, BUFF_DIM_X = 64, STAGES = 2 + MXFP8_BLOCK_SCALING_SIZE = 32 (elements per MXFP8 scaling block) + +Grid: (ceil(N / 64), ceil(M / 64)) +Each block processes a 64x64 chunk in 2 stages of 32x64 tiles loaded into +shared memory. +""" + +# Local @cute.struct classes are SMEM-layout descriptors that need no docstrings. +# pylint: disable=missing-class-docstring + +import logging +import os +from typing import Optional, Type + +import cutlass +from cutlass import cute +from cutlass import pipeline +from cutlass import Float32, Int16, Int32, Int64, Uint32, Uint8 +from cuda.bindings.driver import CUstream # pylint: disable=no-name-in-module +import tvm_ffi + +from transformer_engine.common.CuTeDSL.utils import ( + _bitcast_f32_to_i32, + device_compute_capability, + str_to_cutlass_dtype, + is_packed16, + packed16_kit, + fabs_f32, + exp2f_rcp, + pack_f32x2, + unpack_i64_to_i32x2, +) +from transformer_engine.common.CuTeDSL.activations import ( + act_relu, + act_gelu, + act_silu, + act_qgelu, + act_srelu, + dact_drelu, + dact_dsrelu, + dact_dsilu, + dact_dqgelu, + dact_dgelu, +) +from transformer_engine.common.CuTeDSL.utils_fp8 import ( + as_byte_tensor, + get_cvt_f32_to_fp8_func, + cvt_f32_to_fp8e8m0, + mul_i64_cvt_f32x4_to_fp8x4, + mul_f32x4_cvt_f32x4_to_fp8x4, + mul_i64_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 + + +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 quantize_rowwise_mxfp8( + sX_tile, # (TILE_Y, TILE_X) bf16/fp16 smem view, post-TMA + sA_tile, # (TILE_Y, TILE_X) activation-input smem tile (dact only) + sO_row_tile, # (TILE_Y, TILE_X) uint8 smem view (rowwise FP8 output) + mS_row_stage, # rowwise scale tensor (1D swizzled, or 2D linear) + max_norm_rcp, + tile_row_start, # Int32 — global row index of this stage's row 0 + # (= tile_idx_y * TILE_Y). Used to mask OOB scale stores + # for irregular shapes. + tile_col_start, # Int32 — global col index of this CTA's col 0 + # (= bidx * TILE_X). Same purpose. + M, + N, # Int32 — full tensor extents; OOB threads skip their + # scale store. + ACTIVATION, + DTYPE, + FP8_DTYPE, + TILE_X, + TILE_Y, + WAVES, + THREADS_PER_BANK, + PACK_SIZE, + WITH_ACT=False, + WITH_DACT=False, + WITH_DBIAS=False, # rowwise-only dbias: accumulate per-column partials + dbias_acc=None, # only needed when WITH_DBIAS is True +): + """Quantize one SMEM tile rowwise to MXFP8 (per-row 32-elt block scales); returns the tile amax.""" + tidx, _, _ = cute.arch.thread_idx() + + CTA_THREADS_Y = TILE_Y # threads per column (rows per tile) + CTA_THREADS_X = TILE_X // MXFP8_BLOCK_SCALING_SIZE # threads per row (chunks per row) + + _, tv_layout = cute.make_layout_tv( + thr_layout=cute.make_layout((CTA_THREADS_Y, CTA_THREADS_X), stride=(CTA_THREADS_X, 1)), + val_layout=cute.make_layout((1, MXFP8_BLOCK_SCALING_SIZE), stride=(0, 1)), + ) + + sX_tv = cute.composition(sX_tile, tv_layout) + sO_tv = cute.composition(sO_row_tile, tv_layout) + + # I/O Elements that belong to this thread + sX_thread = sX_tv[tidx, None] # shape (32,) bf16 + sO_thread = sO_tv[tidx, None] # shape (32,) uint8 + + sO_thread_u32_ptr = cute.recast_ptr(sO_thread.iterator, dtype=Uint32) + # Each wave it writes 32 bytes = 8 uint32s, so in 4 waves we write all 32 quantized elements. + sO_thread_u32 = cute.make_tensor( + sO_thread_u32_ptr, + cute.make_layout( + (MXFP8_BLOCK_SCALING_SIZE // 4,), stride=(1,) + ), # 1 uint32 is 4 fp8 elements + ) + + # PTX allows to fuse relu activation in `cvt.rn.satfinite` + FUSE_RELU = cutlass.const_expr(ACTIVATION == "relu") + # For this fast path we can read in pack of 2 instead of reading individual f16 / bf16 element. + # dbias needs the per-element fp32 values to accumulate, so it forces the slow path. + _row_fast = is_packed16(DTYPE) and (ACTIVATION is None or FUSE_RELU) and not WITH_DBIAS + + amax_r = Float32(0.0) + + # Each thread start reading from the specfic bank based on its thread ID so they can do their best to access different banks + # to avoid bank conflict. + bank_group = (tidx % THREADS_PER_WARP) // THREADS_PER_BANK + # The offset this thread should start reading from based on what's its first bank to access. + offset = bank_group * PACK_SIZE + if cutlass.const_expr(_row_fast): + # If no activation, f16 / bf16 and rowwise quantization, we can read 2 f16 / bf16 at once in a pack + # and use max.xorsign.abs.f16x2 / max.xorsign.abs.bf16x2 to compute + kit = packed16_kit(DTYPE) + sX_thread_rw_i64 = cute.make_tensor( + cute.recast_ptr(sX_thread.iterator, dtype=Int64), + cute.make_layout( + (1, MXFP8_BLOCK_SCALING_SIZE // 4), stride=(0, 1) + ), # 1 int64 is 4 fp16/bf16 elements + ) + # Each wave reads its 4 elements (PACK_SIZE) as one 8-byte vectorized load + in_r = [[None, None] for _ in range(WAVES)] + for w in cutlass.range_constexpr(WAVES): + idx = (w + offset // 4) % (MXFP8_BLOCK_SCALING_SIZE // 4) + in_r[w][0], in_r[w][1] = unpack_i64_to_i32x2(sX_thread_rw_i64[0, idx]) + + amax_2x = Int32(0) + # Each wave will use max.xorsign.abs.f16x2 or max.xorsign.abs.bf16x2 to compare 2 packed elements in parallel + for w in cutlass.range_constexpr(WAVES): + if cutlass.const_expr(FUSE_RELU): + # If we fuse relu then we don't want to do abs since negative value will be set to 0 and they will lose comparison automatically + amax_2x = kit.max_x2(amax_2x, in_r[w][0]) + amax_2x = kit.max_x2(amax_2x, in_r[w][1]) + else: + amax_2x = kit.abs_max_x2(amax_2x, in_r[w][0]) + amax_2x = kit.abs_max_x2(amax_2x, in_r[w][1]) + if cutlass.const_expr(FUSE_RELU): + # Compare the 2 packed max without abs + amax_r = cute.arch.fmax( + kit.x2_lo_to_f32(amax_2x), + kit.x2_hi_to_f32(amax_2x), + ) + # For relu the max is at least 0 + amax_r = cute.arch.fmax(amax_r, Float32(0.0)) + else: + # Compare the 2 packed abs max + amax_r = cute.arch.fmax( + fabs_f32(kit.x2_lo_to_f32(amax_2x)), + fabs_f32(kit.x2_hi_to_f32(amax_2x)), + ) + else: + # Since we need to do computation on individual f16 / bf16 elements, we can't read in pack + sX_thread_rw = cute.make_tensor( + sX_thread.iterator, + cute.make_layout((1, MXFP8_BLOCK_SCALING_SIZE), stride=(0, 1)), + ) + + if cutlass.const_expr(WITH_DACT): + # Backward: out = grad · act'(act_input). sX is grad, sA is act_input. + dop = SUPPORTED_DACTIVATIONS[ACTIVATION] + sA_thread = cute.composition(sA_tile, tv_layout)[tidx, None] + sA_thread_rw = cute.make_tensor( + sA_thread.iterator, + cute.make_layout((1, MXFP8_BLOCK_SCALING_SIZE), stride=(0, 1)), + ) + elif cutlass.const_expr(WITH_ACT): + op = SUPPORTED_ACTIVATIONS[ACTIVATION] + + if cutlass.const_expr(is_packed16(DTYPE) and ACTIVATION is not None): + kit_act = packed16_kit(DTYPE) + + # Each wave we read PACK_SIZE elements, and we have WAVES waves, so we read WAVES * PACK_SIZE (= MXFP8_BLOCK_SCALING_SIZE) elements in total. + in_r = [[None] * PACK_SIZE for _ in range(WAVES)] + for w in cutlass.range_constexpr(WAVES): + start = (w * PACK_SIZE + offset) % MXFP8_BLOCK_SCALING_SIZE + for i in cutlass.range_constexpr(PACK_SIZE): + x = Float32(sX_thread_rw[0, start + i]) + if cutlass.const_expr(WITH_DACT): + # out = grad · act'(act_input) + x = x * dop(Float32(sA_thread_rw[0, start + i])) + # If IS_ACT, apply activation function to x in f32 + elif cutlass.const_expr(WITH_ACT): + # If it's relu, we can handle it later + if not cutlass.const_expr(FUSE_RELU): + x = op(x) + # Accumulate to the per-thread dbias register buffer for this tile if WITH_DBIAS + if cutlass.const_expr(WITH_DBIAS): + # dbias_acc is register buffer so we can just write without bank conflict + dbias_acc[w * PACK_SIZE + i] += x + # If 16-bit input with activation, truncate to IType + if cutlass.const_expr(is_packed16(DTYPE) and ACTIVATION is not None): + x = kit_act.truncate_f32(x) + in_r[w][i] = x + if cutlass.const_expr(FUSE_RELU): + amax_r = cute.arch.fmax( + amax_r, x + ) # For relu cases, we don't need abs since negative values will be 0 so they lose comparison automatically + else: + amax_r = cute.arch.fmax(amax_r, fabs_f32(x)) + if cutlass.const_expr(FUSE_RELU): + amax_r = cute.arch.fmax(amax_r, Float32(0.0)) # If relu, the amax is at least 0 + + biased_exp_r = cvt_f32_to_fp8e8m0(amax_r * max_norm_rcp) + + # mS_row_stage has logical shape (32, 2) and we have 64 threads where each is mapped to one scale factor + # The TV layout is equivalent to TV layout with thr_layout=(32, 2):(2, 1), val_layout=(1,) + # but it's too trival so let's just index it directly without using layout + # Note this is the logical layout, which is on top of the swizzled / non-swizzled scale factor layout + # that mappes the logical index to the physical offset + + # For irregular shapes, skip the scale store if this thread's logical row / col-block lies past the input's actual extents. + # TMA already zero-fills OOB input reads and drops OOB output writes; only the direct scale-byte gmem store needs an explicit guard. + scale_row = tile_row_start + tidx // CTA_THREADS_X + scale_col_first_elt = tile_col_start + (tidx % CTA_THREADS_X) * MXFP8_BLOCK_SCALING_SIZE + if scale_row < M and scale_col_first_elt < N: + mS_row_stage[(tidx // CTA_THREADS_X, tidx % CTA_THREADS_X)] = Uint8(biased_exp_r) + + inv_scale_r = exp2f_rcp(biased_exp_r) # f32 reciprocal of the scale + scale_2x = pack_f32x2(inv_scale_r, inv_scale_r) + if cutlass.const_expr(_row_fast): + mul_cvt_x4_func = mul_i64_cvt_packed16x4_to_fp8x4(DTYPE, FP8_DTYPE, FUSE_RELU) + else: + mul_cvt_x4_func = mul_i64_cvt_f32x4_to_fp8x4(FP8_DTYPE, FUSE_RELU) + + for w in cutlass.range_constexpr(WAVES): + idx = (w * 4 + offset) % MXFP8_BLOCK_SCALING_SIZE + idx = idx // 4 + if cutlass.const_expr(_row_fast): + # Convert 2 packed f16/bf16 pairs to 4 fp8 in one fused op + sO_thread_u32[idx] = mul_cvt_x4_func(in_r[w][0], in_r[w][1], scale_2x) + else: + # Convert 4 f32 to 4 fp8 in one fused op + sO_thread_u32[idx] = mul_cvt_x4_func( + in_r[w][0], in_r[w][1], in_r[w][2], in_r[w][3], scale_2x + ) + + return amax_r + + +@cute.jit +def quantize_colwise_mxfp8( + sX_tile, # (TILE_Y, TILE_X) bf16/fp16 smem view, post-TMA + sO_col_tile, # (TILE_Y, TILE_X) uint8 smem view (colwise FP8 output) + mS_col_stage, # colwise scale tensor (1D swizzled, or 2D linear) + max_norm_rcp, + tile_row_start, # Int32 — global row index of this stage's row 0 + # (= tile_idx_y * TILE_Y). Used to mask OOB scale stores + # for irregular shapes. + tile_col_start, # Int32 — global col index of this CTA's col 0 + # (= bidx * TILE_X). + M, + N, # Int32 — full tensor extents. + ACTIVATION, + DTYPE, + FP8_DTYPE, + SWIZZLE, + TILE_X, + TILE_Y, # pylint: disable=unused-argument # kept for API symmetry with the rowwise path + WITH_ACT=False, # forward: apply activation to the element + WITH_DACT=False, # backward: out = grad · act'(act_input) + sA_tile=None, # (TILE_Y, TILE_X) activation-input smem tile (dact only) + WITH_DBIAS=False, # also return this thread's column sum (pre-truncate) + CACHE_ACTIVATION=False, # overwrite sX_tile in place with the post-activation + # (IType-truncated) values, so the rowwise pass can read + # them instead of recomputing op +): + """Quantize one SMEM tile colwise to MXFP8 (per-column 32-elt block scales); returns (amax, dbias_partial).""" + tidx, _, _ = cute.arch.thread_idx() + + _, tv_layout = cute.make_layout_tv( + thr_layout=cute.make_layout((1, TILE_X), stride=(TILE_X, 1)), + val_layout=cute.make_layout((MXFP8_BLOCK_SCALING_SIZE, 1), stride=(1, 1)), + ) + + sX_tv = cute.composition(sX_tile, tv_layout) + sO_tv = cute.composition(sO_col_tile, tv_layout) + + # I/O Elements that belong to this thread + sX_thread = sX_tv[tidx, None] + sO_thread = sO_tv[tidx, None] + + USE_HALF_PRECISION = is_packed16(DTYPE) and ACTIVATION is None + dbias_partial = Float32(0.0) + + if cutlass.const_expr(USE_HALF_PRECISION): + kit = packed16_kit(DTYPE) + # If we can use the half precision format, then use the input tile directly since there is no need to upcast + sX_thread_i16 = cute.make_tensor( + cute.recast_ptr(sX_thread.iterator, dtype=Int16), + cute.make_layout((MXFP8_BLOCK_SCALING_SIZE,), stride=(TILE_X,)), + ) + # Stash the strided column reads in registers (CUDA's in_colwise_IType): + # the cvt loop below reuses them instead of re-reading smem. + in_c = [None] * MXFP8_BLOCK_SCALING_SIZE + amax_bits = Int16(0) + for i in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE): + in_c[i] = sX_thread_i16[i] + amax_bits = kit.abs_max_scalar(amax_bits, in_c[i]) + amax_c = fabs_f32(kit.bits_to_f32(amax_bits)) + else: + # Otherwise we need to case input values to fp32. Allocate the register tensor and load from SMEM input tiles. + sX_thread_f32 = cute.make_rmem_tensor( + layout_or_shape=cute.make_layout((MXFP8_BLOCK_SCALING_SIZE,), stride=(1,)), + dtype=Float32, + ) + for i in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE): + sX_thread_f32[i] = Float32(sX_thread[i]) + # Apply activation (fwd) or grad·act'(act_input) (bwd dact) in f32. + if cutlass.const_expr(WITH_DACT): + dop = SUPPORTED_DACTIVATIONS[ACTIVATION] + sA_thread = cute.composition(sA_tile, tv_layout)[tidx, None] + for i in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE): + sX_thread_f32[i] = sX_thread_f32[i] * dop(Float32(sA_thread[i])) + elif cutlass.const_expr(WITH_ACT): + op = SUPPORTED_ACTIVATIONS[ACTIVATION] + for i in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE): + sX_thread_f32[i] = op(sX_thread_f32[i]) + # Truncate the activation (after we apply op) back to the half precision type if input is also half precision. + if cutlass.const_expr(is_packed16(DTYPE) and ACTIVATION is not None): + kit_act = packed16_kit(DTYPE) + for i in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE): + sX_thread_f32[i] = kit_act.truncate_f32(sX_thread_f32[i]) + # Columnwise is the preferred direction so it runs first. If it needs to cache the activation in the input tile + # to let the rowwise pass read it, we need to cast and overwrite the input data in-place here + if cutlass.const_expr(CACHE_ACTIVATION): + for i in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE): + sX_thread[i] = DTYPE(sX_thread_f32[i]) + amax_c = Float32(0.0) + for i in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE): + amax_c = cute.arch.fmax(amax_c, fabs_f32(sX_thread_f32[i])) + + # Irregular shapes: skip when this stage's row range or this thread's + # column lies past the input extents. TILE_Y == MXFP8_BLOCK_SCALING_SIZE so each stage + # is exactly one scale-row; valid iff `tile_row_start < M`. + biased_exp_c = cvt_f32_to_fp8e8m0(amax_c * max_norm_rcp) + scale_col = tile_col_start + tidx + if tile_row_start < M and scale_col < N: + if cutlass.const_expr(SWIZZLE): + mS_col_stage[(0, tidx % 32, tidx // 32)] = Uint8(biased_exp_c) + else: + mS_col_stage[(0, tidx)] = Uint8(biased_exp_c) + + inv_scale_c = exp2f_rcp(biased_exp_c) + cvt_to_fp8_func = get_cvt_f32_to_fp8_func(FP8_DTYPE) + if cutlass.const_expr(USE_HALF_PRECISION): + kit_cast = packed16_kit(DTYPE) + for i in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE): + v_f32 = kit_cast.bits_to_f32(in_c[i]) + if cutlass.const_expr(WITH_DBIAS): + dbias_partial += v_f32 + sO_thread[i] = Uint8(cvt_to_fp8_func(v_f32 * inv_scale_c)) + else: + for i in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE): + # Accumulate the per-thread column partial for dbias if WITH_DBIAS. + if cutlass.const_expr(WITH_DBIAS): + dbias_partial += sX_thread_f32[i] + sO_thread[i] = Uint8(cvt_to_fp8_func(sX_thread_f32[i] * inv_scale_c)) + + # Return this stage's per-column partial alongside amax; the caller accumulates + # it across stages (a scalar can't be updated in-place through the arg). + return amax_c, dbias_partial + + +@cute.jit +def 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, + max_norm_rcp, + DTYPE: cutlass.Constexpr, + fp8_dtype: cutlass.Constexpr, +): + """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_i64_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, Uint8) + rO_col = cute.make_rmem_tensor(MXFP8_BLOCK_SCALING_SIZE, Uint8) + 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 + kit = packed16_kit(DTYPE) + 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 = Int32(0) + for i in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE // 2): + pair = rX_2x[i] + row_amax2 = kit.abs_max_x2(row_amax2, pair) + a_lo = fabs_f32(kit.x2_lo_to_f32(pair)) + a_hi = fabs_f32(kit.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(kit.x2_lo_to_f32(row_amax2)), fabs_f32(kit.x2_hi_to_f32(row_amax2)) + ) + row_exp = cvt_f32_to_fp8e8m0(row_amax * max_norm_rcp) + row_inv = exp2f_rcp(row_exp) + tSsS_row_tile[0] = Uint8(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_fp8e8m0(sColReduce_warp[lane] * max_norm_rcp) + tSsS_col_tile[0] = Uint8(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 = kit.x2_lo_to_f32(p01) + f1 = kit.x2_hi_to_f32(p01) + f2 = kit.x2_lo_to_f32(p23) + f3 = kit.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_fp8e8m0(row_amax * max_norm_rcp) + row_inv = exp2f_rcp(row_exp) + tSsS_row_tile[0] = Uint8(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_fp8e8m0(sColReduce_warp[lane] * max_norm_rcp) + tSsS_col_tile[0] = Uint8(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 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, + with_noop: bool = False, + activation: Optional[str] = None, + ): + if dtype is None or 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 # readable input-dtype token, for __str__ + if fp8_dtype not in ("e4m3", "e5m2"): + raise ValueError(f"unknown FP8 dtype {fp8_dtype!r}; expected 'e4m3' or 'e5m2'") + self.FP8_DTYPE = fp8_dtype + 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.WITH_NOOP = with_noop + self.MAX_NORM_RCP = FP8E4M3_MAX_NORM_RCP if fp8_dtype == "e4m3" else FP8E5M2_MAX_NORM_RCP + + def __str__(self): + return ( + f"MXFP8QuantizeConfig(dtype={self.DTYPE_STR}, fp8_dtype={self.FP8_DTYPE}, " + 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}, with_noop={self.WITH_NOOP}, " + f"activation={self.ACTIVATION})" + ) + + __repr__ = __str__ + + +class MXFP8QuantizeKernel: + """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 + _WAVES = ( + MXFP8_BLOCK_SCALING_SIZE // _PACK_SIZE + ) # Each thread reads 8 waves with each wave reads 4 packed bf16, so it reads a whole MXFP8 block in total + _TOTAL_BANKS_WIDTH = (32 * 4) // 1 # 32 banks × 4 bytes, in bytes (uint8 stride) + _THREADS_PER_BANK = _TOTAL_BANKS_WIDTH // MXFP8_BLOCK_SCALING_SIZE # 4 threads per bank + _NUM_STAGES = 2 # The pipeline depth is always 2 + + def __init__(self, cfg): + self.cfg = cfg + # Cast + dbias with no activation gets the larger tile (CUDA CAST_DBIAS_ONLY). + cast_dbias_only = cfg.WITH_DBIAS and not cfg.WITH_DACT and not cfg.WITH_ACT + # Use a different tile size for dbias only config + # No matter what tile size we use, each thread always handles a (1, MXFP8_BLOCK_SCALING_SIZE) chunk + if cast_dbias_only: + self._NUM_TILES = 4 # Each CTA handles 4 tiles stacked vertically + self._THREADS_PER_CTA = 128 + self._TILE_COLS = 128 + self._TILE_ROWS = 32 + else: + self._NUM_TILES = 2 # Each CTA handles 2 tiles stacked vertically + self._THREADS_PER_CTA = 64 + self._TILE_COLS = 64 + self._TILE_ROWS = 32 + self._NUM_WARPS = self._THREADS_PER_CTA // 32 + # We prefer to do dbias reduction in colwise which is easier (no cross-thread reduction needed). + # Only do rowwise reduction when we don't quantize columnwisely when WITH_DBIAS is True. + self.DBIAS_REDUCTION_COLWISE = cfg.WITH_DBIAS and cfg.COLWISE + self.DBIAS_REDUCTION_ROWWISE = cfg.WITH_DBIAS and not cfg.COLWISE + # Cache activation in-place in the SMEM input tile when we process both rowwise and colwise passes + # so the activation is only computed once in the direction we favor (columnwise) and the other direction (rowwise) + # reads the cached value instead of recomputing it. + # Note: if activation is relu, there is no standalong relu applied because it's already fused into `cvt.rn.satfinite` + # so it should be treated as "no activation" + self.CACHE_ACTIVATION = ( + (cfg.WITH_ACT or cfg.WITH_DACT) + and cfg.ROWWISE + and cfg.COLWISE + and cfg.ACTIVATION != "relu" + ) + # The global tensor amax (mAmax) is the max over ALL elements. Each direction's + # per-block amaxes already span every element, so when both passes run we only + # fold the global amax from one of them — favor colwise (matches the flags + # above). The per-block *scale* amax is still computed in each pass for its own + # scale; this only skips the redundant global comparison in the other pass. + self.AMAX_FROM_COLWISE = cfg.WITH_AMAX and cfg.COLWISE + self.AMAX_FROM_ROWWISE = cfg.WITH_AMAX and not cfg.COLWISE + + @cute.jit + def __call__( + self, + mX: cute.Tensor, # Input tensor to quantize + mO_row: Optional[cute.Tensor], + mS_row: Optional[cute.Tensor], # Rowwise output and scale tensors + mO_col: Optional[cute.Tensor], + mS_col: Optional[cute.Tensor], # Colwise output and scale tensors + mAmax: Optional[cute.Tensor], # Global amax accumulator, only used when WITH_AMAX is True + mNoop: Optional[cute.Tensor], # 1-element cast_noop flag, only used when WITH_NOOP is True + mDActInput: Optional[ + cute.Tensor + ], # Activation input for activation derivative fusion, only used when WITH_DACT is True + mWorkspace: Optional[ + cute.Tensor + ], # Workspace for the dbias reduction, only used when WITH_DBIAS is True + stream: CUstream, + ): + if cutlass.const_expr(CUTEDSL_DEBUG_LOGGING): + cute.printf(f"[CuTeDSL] MXFP8QuantizeKernel.__call__() with config: {self.cfg}\n") + + M = mX.shape[0] + N = mX.shape[1] + cfg = self.cfg + max_norm_rcp = cfg.MAX_NORM_RCP + num_scale_cols = N // MXFP8_BLOCK_SCALING_SIZE + num_scale_rows = M // MXFP8_BLOCK_SCALING_SIZE + + # The FFI boundary carries native FP8/E8M0 dtypes; the kernel works on bytes. + if cutlass.const_expr(cfg.ROWWISE): + mO_row = as_byte_tensor(mO_row) + mS_row = as_byte_tensor(mS_row) + if cutlass.const_expr(cfg.COLWISE): + mO_col = as_byte_tensor(mO_col) + mS_col = as_byte_tensor(mS_col) + + # If WITH_GEMM_SWIZZLED_SCALES is enabled, the output must satisfy cublas's swizzled layout + # This is expressed as a CuTe layout applied to the output tensor so it can be transparent throughout the kernel implementation. + # See https://docs.nvidia.com/cuda/cublas/#d-block-scaling-factors-layout for more details. + if cutlass.const_expr(cfg.WITH_GEMM_SWIZZLED_SCALES): + num_tiles_M = (M + 127) // 128 + num_tiles_SC = (num_scale_cols + 3) // 4 + num_tiles_SR = (num_scale_rows + 3) // 4 + num_tiles_N = (N + 127) // 128 + + if cutlass.const_expr(cfg.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(cfg.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)), + ), + ) + + # We have 2 stages in our pipeline where each stage loads / computes a (TILE_Y, TILE_X) tile + smem_tile_layout = cute.make_ordered_layout( + (self._TILE_ROWS, self._TILE_COLS), order=(1, 0) + ) + cta_tiler = (self._TILE_ROWS, self._TILE_COLS) + + # Input TMA atoms + op_load = cute.nvgpu.cpasync.CopyBulkTensorTileG2SOp() + tma_atom, tma_src = cute.nvgpu.cpasync.make_tiled_tma_atom( + op_load, + mX, + smem_tile_layout, + cta_tiler, + num_multicast=1, + ) + + # Activation input TMA atoms for activation derivative fusion + tma_atom_act = None + tma_src_act = None + if cutlass.const_expr(cfg.WITH_DACT): + tma_atom_act, tma_src_act = cute.nvgpu.cpasync.make_tiled_tma_atom( + op_load, + mDActInput, + smem_tile_layout, + cta_tiler, + num_multicast=1, + ) + + # Output TMA atoms + op_store = cute.nvgpu.cpasync.CopyBulkTensorTileS2GOp() + out_smem_layout = cute.make_ordered_layout((self._TILE_ROWS, self._TILE_COLS), order=(1, 0)) + tma_atom_out_row = None + tma_dst_out_row = None + tma_atom_out_col = None + tma_dst_out_col = None + if cutlass.const_expr(cfg.ROWWISE): + tma_atom_out_row, tma_dst_out_row = cute.nvgpu.cpasync.make_tiled_tma_atom( + op_store, + mO_row, + out_smem_layout, + cta_tiler, + num_multicast=1, + ) + if cutlass.const_expr(cfg.COLWISE): + tma_atom_out_col, tma_dst_out_col = cute.nvgpu.cpasync.make_tiled_tma_atom( + op_store, + mO_col, + out_smem_layout, + cta_tiler, + num_multicast=1, + ) + + grid = [ + cute.ceil_div(Int32(N), self._TILE_COLS), + cute.ceil_div(M, self._TILE_ROWS * self._NUM_TILES), + ] + block = [ + self._THREADS_PER_CTA, + ] + + self.kernel( + mX, + mS_row, + mS_col, + mAmax, + mNoop, + mWorkspace, + max_norm_rcp, + mX.element_type, + tma_atom, + tma_src, + tma_atom_out_row, + tma_dst_out_row, + tma_atom_out_col, + tma_dst_out_col, + tma_atom_act, + tma_src_act, + ).launch( + grid=grid, + block=block, + stream=stream, + ) + + @cute.kernel + def kernel( + self, + mX, + mS_row, + mS_col, + mAmax, + mNoop, + mWorkspace, + max_norm_rcp, + dtype: cutlass.Constexpr[Type[cutlass.Numeric]], + tma_atom, + tma_src, # Input TMA atoms + tma_atom_out_row, + tma_dst_out_row, # Rowwise output TMA atoms + tma_atom_out_col, + tma_dst_out_col, # Colwise output TMA atoms + tma_atom_act, + tma_src_act, # Activation derivative TMA atoms, or None if WITH_DACT is False + ): + """Device entry: no-op the CTA when the noop flag is set, else run the quantize main loop.""" + cfg = self.cfg + # If WITH_NOOP is True and no ACT / DACT fusion is involved, the kernel becomes a no-op + noop = cfg.WITH_NOOP and not cfg.WITH_ACT and not cfg.WITH_DACT + if not cutlass.const_expr(noop) or mNoop[0] != Float32(1.0): + self._kernel_main( + mX, + mS_row, + mS_col, + mAmax, + mWorkspace, + max_norm_rcp, + dtype, + tma_atom, + tma_src, + tma_atom_out_row, + tma_dst_out_row, + tma_atom_out_col, + tma_dst_out_col, + tma_atom_act, + tma_src_act, + ) + + @cute.jit + def _kernel_main( + self, + mX, + mS_row, + mS_col, + mAmax, + mWorkspace, + max_norm_rcp, + dtype: cutlass.Constexpr[Type[cutlass.Numeric]], + tma_atom, + tma_src, # Input TMA atoms + tma_atom_out_row, + tma_dst_out_row, # Rowwise output TMA atoms + tma_atom_out_col, + tma_dst_out_col, # Colwise output TMA atoms + tma_atom_act, + tma_src_act, # Activation derivative TMA atoms, or None if WITH_DACT is False + ): + cfg = self.cfg + + if cutlass.const_expr(cfg.ROWWISE): + mS_row = cute.zipped_divide( + mS_row, (self._TILE_ROWS, self._TILE_COLS // MXFP8_BLOCK_SCALING_SIZE) + ) + if cutlass.const_expr(cfg.COLWISE): + mS_col = cute.zipped_divide( + mS_col, (self._TILE_ROWS // MXFP8_BLOCK_SCALING_SIZE, self._TILE_COLS) + ) + + # Allocate shared memory for the input and rowwise / columnwise outputs + if cutlass.const_expr(cfg.ROWWISE and cfg.COLWISE): + + @cute.struct + class SharedStorage: + mbar_storage: cute.struct.MemRange[cute.Int64, 2 * self._NUM_STAGES] + sX: cute.struct.Align[ + cute.struct.MemRange[ + dtype, self._TILE_ROWS * self._TILE_COLS * self._NUM_STAGES + ], + 128, + ] + sO_row: cute.struct.Align[ + cute.struct.MemRange[ + Uint8, self._TILE_ROWS * self._TILE_COLS * self._NUM_STAGES + ], + 128, + ] + sO_col: cute.struct.Align[ + cute.struct.MemRange[ + Uint8, self._TILE_ROWS * self._TILE_COLS * self._NUM_STAGES + ], + 128, + ] + sAmax: cute.struct.MemRange[Float32, self._NUM_WARPS] + + elif cutlass.const_expr(cfg.ROWWISE and not cfg.COLWISE): + + @cute.struct + class SharedStorage: + mbar_storage: cute.struct.MemRange[cute.Int64, 2 * self._NUM_STAGES] + sX: cute.struct.Align[ + cute.struct.MemRange[ + dtype, self._TILE_ROWS * self._TILE_COLS * self._NUM_STAGES + ], + 128, + ] + sO_row: cute.struct.Align[ + cute.struct.MemRange[ + Uint8, self._TILE_ROWS * self._TILE_COLS * self._NUM_STAGES + ], + 128, + ] + sAmax: cute.struct.MemRange[Float32, self._NUM_WARPS] + + elif cutlass.const_expr(cfg.ROWWISE): + + @cute.struct + class SharedStorage: + mbar_storage: cute.struct.MemRange[cute.Int64, 2 * self._NUM_STAGES] + sX: cute.struct.Align[ + cute.struct.MemRange[ + dtype, self._TILE_ROWS * self._TILE_COLS * self._NUM_STAGES + ], + 128, + ] + sO_row: cute.struct.Align[ + cute.struct.MemRange[ + Uint8, self._TILE_ROWS * self._TILE_COLS * self._NUM_STAGES + ], + 128, + ] + sAmax: cute.struct.MemRange[Float32, self._NUM_WARPS] + + else: + + @cute.struct + class SharedStorage: + mbar_storage: cute.struct.MemRange[cute.Int64, 2 * self._NUM_STAGES] + sX: cute.struct.Align[ + cute.struct.MemRange[ + dtype, self._TILE_ROWS * self._TILE_COLS * self._NUM_STAGES + ], + 128, + ] + sO_col: cute.struct.Align[ + cute.struct.MemRange[ + Uint8, self._TILE_ROWS * self._TILE_COLS * self._NUM_STAGES + ], + 128, + ] + sAmax: cute.struct.MemRange[Float32, self._NUM_WARPS] + + smem = cutlass.utils.SmemAllocator() + storage = smem.allocate(SharedStorage) + # Apply the layout to the allocated shared memory buffers so the first rank is the tile (nested layout) + # and the second rank is the pipeline stage + sX = storage.sX.get_tensor( + cute.make_layout( + ((self._TILE_ROWS, self._TILE_COLS), self._NUM_STAGES), + stride=((self._TILE_COLS, 1), self._TILE_ROWS * self._TILE_COLS), + ) + ) + if cutlass.const_expr(cfg.ROWWISE): + sO_row = storage.sO_row.get_tensor( + cute.make_layout( + ((self._TILE_ROWS, self._TILE_COLS), self._NUM_STAGES), + stride=((self._TILE_COLS, 1), self._TILE_ROWS * self._TILE_COLS), + ) + ) + if cutlass.const_expr(cfg.COLWISE): + sO_col = storage.sO_col.get_tensor( + cute.make_layout( + ((self._TILE_ROWS, self._TILE_COLS), self._NUM_STAGES), + stride=((self._TILE_COLS, 1), self._TILE_ROWS * self._TILE_COLS), + ) + ) + + # Allocate shared memory for the activation input used for the activation derivative fusion. + if cutlass.const_expr(cfg.WITH_DACT): + + @cute.struct + class DactStorage: + sActInput: cute.struct.Align[ + cute.struct.MemRange[ + dtype, self._TILE_ROWS * self._TILE_COLS * self._NUM_STAGES + ], + 128, + ] + + dact_storage = smem.allocate(DactStorage) + # Apply the same layout as the input + sActInput = dact_storage.sActInput.get_tensor( + cute.make_layout( + ((self._TILE_ROWS, self._TILE_COLS), self._NUM_STAGES), + stride=((self._TILE_COLS, 1), self._TILE_ROWS * self._TILE_COLS), + ) + ) + + warp_idx = cute.arch.warp_idx() + warp_idx = cute.arch.make_warp_uniform(warp_idx) + + # Prefetch TMA descriptors + if warp_idx == 0: + cute.nvgpu.cpasync.prefetch_descriptor(tma_atom) + if cutlass.const_expr(cfg.WITH_DACT): + cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_act) + + tidx, _, _ = cute.arch.thread_idx() + bidx, bidy, _ = cute.arch.block_idx() + + # Only warp 0 is the producer (issues TMA) + producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, 1) + # Every warp is the consumer (reads the data loaded by TMA) + consumer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, self._NUM_WARPS) + + # Bytes transferred per TMA copy: one (TILE_Y, TILE_X) tile of dtype. + tx_count = self._TILE_ROWS * self._TILE_COLS * dtype.width // 8 + # dact loads two tiles (grad + act_input) under the same per-stage barrier, + # so the barrier must expect both copies' bytes. + if cutlass.const_expr(cfg.WITH_DACT): + tx_count *= 2 + + mainloop_pipeline = pipeline.PipelineTmaAsync.create( + barrier_storage=storage.mbar_storage.data_ptr(), + num_stages=self._NUM_STAGES, + producer_group=producer_group, + consumer_group=consumer_group, + tx_count=tx_count, + cta_layout_vmnk=None, # single-CTA, no cluster/multicast + ) + + prod_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self._NUM_STAGES + ) + cons_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self._NUM_STAGES + ) + + M = mX.shape[0] + N = mX.shape[1] + + num_tiles = cutlass.min( + self._NUM_TILES, + cute.ceil_div(M - bidy * self._TILE_ROWS * self._NUM_TILES, self._TILE_ROWS), + ) + + # Tile the TMA gmem view: ((TILE_Y, TILE_X), (M/TILE_Y, N/TILE_X)). + gX_tiled = cute.zipped_divide(tma_src, (self._TILE_ROWS, self._TILE_COLS)) + + # Partition sX/gX for the TMA atom (single-CTA, no cluster/multicast). + tXsX, tXgX = cute.nvgpu.cpasync.tma_partition( + tma_atom, + 0, # Use the only CTA to do the TMA copy + cute.make_layout(1), # This cluster only has 1 CTAs + sX, + gX_tiled, + ) + + # If WITH_DACT, partition the activation input for TMA as well in the same way + if cutlass.const_expr(cfg.WITH_DACT): + gA_tiled = cute.zipped_divide(tma_src_act, (self._TILE_ROWS, self._TILE_COLS)) + tXsA, tXgA = cute.nvgpu.cpasync.tma_partition( + tma_atom_act, + 0, + cute.make_layout(1), + sActInput, + gA_tiled, + ) + + # Partitioning for rowwise / columnwise outputs + if cutlass.const_expr(cfg.ROWWISE): + gO_row_tiled = cute.zipped_divide(tma_dst_out_row, (self._TILE_ROWS, self._TILE_COLS)) + tXsO_row, tXgO_row = cute.nvgpu.cpasync.tma_partition( + tma_atom_out_row, + 0, + cute.make_layout(1), + sO_row, + gO_row_tiled, + ) + if cutlass.const_expr(cfg.COLWISE): + gO_col_tiled = cute.zipped_divide(tma_dst_out_col, (self._TILE_ROWS, self._TILE_COLS)) + tXsO_col, tXgO_col = cute.nvgpu.cpasync.tma_partition( + tma_atom_out_col, + 0, + cute.make_layout(1), + sO_col, + gO_col_tiled, + ) + + # Ensure barrier init is visible to all threads before the pipeline is used. + cute.arch.sync_threads() + + # Prologue: warp 0 prefetches up to NUM_STAGES tiles to fully fill the pipeline + if warp_idx == 0: + for s in cutlass.range_constexpr(self._NUM_STAGES): + if s < num_tiles: + mainloop_pipeline.producer_acquire(prod_state) + tile_y = bidy * self._NUM_TILES + s + cute.copy( + tma_atom, + tXgX[(None, (tile_y, bidx))], + tXsX[(None, prod_state.index)], + tma_bar_ptr=mainloop_pipeline.producer_get_barrier(prod_state), + ) + if cutlass.const_expr(cfg.WITH_DACT): + cute.copy( + tma_atom_act, + tXgA[(None, (tile_y, bidx))], + tXsA[(None, prod_state.index)], + tma_bar_ptr=mainloop_pipeline.producer_get_barrier(prod_state), + ) + mainloop_pipeline.producer_commit(prod_state) + prod_state.advance() + + # Per-thread amax accumulator + if cutlass.const_expr(cfg.WITH_AMAX): + per_thread_amax = Float32(0.0) + + # Prepare thread-level register accumulators for rowwise dbias reduction. + # Each thread will process two (1, MXFP8_BLOCK_SCALING_SIZE) rows in two stages, and in each stage the thread will add the + # (after dact applied) value to this register array with the same shape so it carries the the two stages' partial sum. + # Then it will be written to a SMEM buffer to let the whole CTA do the reduction separately to yield + # the final (1, TILE_X) dbias workspace output. + rowwise_dbias_acc = None + if cutlass.const_expr(self.DBIAS_REDUCTION_ROWWISE): + rowwise_dbias_acc = cute.make_rmem_tensor( + layout_or_shape=cute.make_layout((MXFP8_BLOCK_SCALING_SIZE,), stride=(1,)), + dtype=Float32, + ) + # Zero the accumulator registers. + for c in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE): + rowwise_dbias_acc[c] = Float32(0.0) + block_dbias = Float32(0.0) + # Prepare thread-level register for columnwise dbias reduction. + # Each thread will process two (MXFP8_BLOCK_SCALING_SIZE, 1) columns in two stages, and in each stage the thread will reduce the + # (after dact applied) column to (1,) and add to this register. + # Then this partial sum scalar will be written to the GMEM workspace buffer directly. + if cutlass.const_expr(self.DBIAS_REDUCTION_COLWISE): + block_dbias = Float32(0.0) + + # Consumer: all warps fetch from the pipeline, process its tile, and issue a new load to the tile buffer it just consumed + for tile_idx in cutlass.range(num_tiles, unroll=1): + mainloop_pipeline.consumer_wait(cons_state) + # Only allow at most _NUM_STAGES-1 stages to be in-flight, because this iteration will reuse the ring buffer + # that is read _NUM_STAGES iterations ago so we must wait for whoever is reading that buffer to finish + if warp_idx == 0: + cute.arch.cp_async_bulk_wait_group(self._NUM_STAGES - 1, read=True) + cute.arch.sync_threads() + # The current pipeline stage index, which is the tile index modulo the number of stages. + # This is used to index into the shared memory ring buffers that are wrapped around the number of stages. + stage_idx = cons_state.index + sX_tile = sX[(None, stage_idx)] + # Also fetch the activation input if WITH_DACT + sActInput_tile = None + if cutlass.const_expr(cfg.WITH_DACT): + sActInput_tile = sActInput[(None, stage_idx)] + # Each CTA handles `NUM_TILES` tiles stacked vertically, so tile_idx_x is just the block index along X dimension + # and tile_idx_y is the tile that this stage handles out of the `NUM_TILES` tiles + tile_idx_x = bidx + tile_idx_y = bidy * self._NUM_TILES + tile_idx + # Process rowwise and colwise quantization separately + if cutlass.const_expr(cfg.COLWISE): + # The first row that belongs to this CTA. Each CTA handles NUM_TILES of (TILE_Y, TILE_X) tiles stacked vertically, + # and each stage handles one of them. + sO_col_tile = sO_col[(None, stage_idx)] + mS_col_stage = cute.flatten(mS_col[(None, (tile_idx_y, tile_idx_x))]) + + amax_c, dbias_c = self._process_colwise( + sX_tile, + sO_col_tile, + mS_col_stage, + max_norm_rcp, + 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_COLWISE): + block_dbias += dbias_c + # If we cache the activation in shared memory, we need to ensure that all threads have finished writing to the shared memory + # from the columnwise pass before any thread reads from it in the rowwise pass. + if cutlass.const_expr(self.CACHE_ACTIVATION): + cute.arch.sync_threads() + if cutlass.const_expr(cfg.ROWWISE): + sO_row_tile = sO_row[(None, stage_idx)] + # mS_row is ((SCALE_TILE), (GRID)) where SCALE_TILE = (32, 2). + # Each CTA owns NUM_TILES consecutive row-tiles of GRID. cute + # auto-decomposes the flat row coord `bidy * NUM_TILES + tile_idx` + # onto GRID's hierarchical row modes — which is the + # (i_hi, tile_Y) tile-major order for swizzled, and the plain + # row-tile order for compact. Same source, both layouts correct. + mS_row_stage = cute.flatten(mS_row[(None, (tile_idx_y, tile_idx_x))]) + amax_r = self._process_rowwise( + sX_tile, + sO_row_tile, + mS_row_stage, + max_norm_rcp, + tile_idx_y * self._TILE_ROWS, + bidx * self._TILE_COLS, + M, + N, + sActInput_tile, + rowwise_dbias_acc, + ) + + if cutlass.const_expr(self.AMAX_FROM_ROWWISE): + per_thread_amax = cute.arch.fmax(per_thread_amax, amax_r) + + # Make the shared-memory writes visible to the TMA's async proxy before the TMA reads them. + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + cute.arch.sync_threads() + + # We are done with this input pipeline SMEM buffer, signal the producer that it can write to this buffer + mainloop_pipeline.consumer_release(cons_state) + + # Warp 0 issues TMA copy to write the quantized output tile from shared memory to global memory and then commits + if warp_idx == 0: + tile_y = bidy * self._NUM_TILES + tile_idx + if cutlass.const_expr(cfg.ROWWISE): + cute.copy( + tma_atom_out_row, + tXsO_row[(None, stage_idx)], + tXgO_row[(None, (tile_y, bidx))], + ) + if cutlass.const_expr(cfg.COLWISE): + cute.copy( + tma_atom_out_col, + tXsO_col[(None, stage_idx)], + tXgO_col[(None, (tile_y, bidx))], + ) + cute.arch.cp_async_bulk_commit_group() + + cons_state.advance() + + # The pipeline is no longer fully filled after we consume this tile, so we fetch a new tile to fill the pipeline. + # The next _NUM_STAGES-1 tiles are already in-flight, so the next tile to fetch is after _NUM_STAGES tiles. + if warp_idx == 0: + next_tile_idx = tile_idx + self._NUM_STAGES + if next_tile_idx < num_tiles: + mainloop_pipeline.producer_acquire(prod_state) + tile_y = bidy * self._NUM_TILES + next_tile_idx + cute.copy( + tma_atom, + tXgX[(None, (tile_y, bidx))], + tXsX[(None, prod_state.index)], + tma_bar_ptr=mainloop_pipeline.producer_get_barrier(prod_state), + ) + if cutlass.const_expr(cfg.WITH_DACT): + cute.copy( + tma_atom_act, + tXgA[(None, (tile_y, bidx))], + tXsA[(None, prod_state.index)], + tma_bar_ptr=mainloop_pipeline.producer_get_barrier(prod_state), + ) + mainloop_pipeline.producer_commit(prod_state) + prod_state.advance() + # End of the main pipeline loop + + # Complete the cross-thread dbias reduction after each thread has its own per-thread partial sum after the rowwise quantization. + if cutlass.const_expr(self.DBIAS_REDUCTION_ROWWISE): + # If we do the dbias reduction in the rowwise pass, each thread will have a (1, MXFP8_BLOCK_SCALING_SIZE) partial sum + # and we need to write these to a SMEM buffer and let each thread reduce it in the columnwise direction + block_dbias = self._dbias_reduction_rowwise_epilouge(smem, tidx, rowwise_dbias_acc) + + # Write the per-tile reduced dbias to the global workspace. + if cutlass.const_expr(cfg.WITH_DBIAS): + dbias_col = bidx * self._TILE_COLS + tidx + if dbias_col < N: + mWorkspace[(bidy, dbias_col)] = block_dbias + + if cutlass.const_expr(cfg.WITH_AMAX): + sAmax = storage.sAmax.get_tensor(cute.make_layout(self._NUM_WARPS)) + self._amax_epilogue(sAmax, mAmax, tidx, warp_idx, per_thread_amax) + + # Wait for in-flight TMA stores so data is visible to the host + # before the kernel returns. + cute.arch.cp_async_bulk_wait_group(0, read=False) + + @cute.jit + def _dbias_reduction_rowwise_epilouge(self, smem, tidx, rowwise_dbias_acc): + # Pad the buffer to avoid bank conflicts. The logical shape is still the same. Only the stride is different. + DBIAS_BUFF_WIDTH = ( + self._TILE_COLS // MXFP8_BLOCK_SCALING_SIZE * (MXFP8_BLOCK_SCALING_SIZE + 1) + ) + + # Allocate the SMEM buffer that all threads use to reduce the two-stage partial sum (per thread) to the + # partial sum (per block). + @cute.struct + class DbiasStorage: + sDbias: cute.struct.MemRange[Float32, self._TILE_ROWS * DBIAS_BUFF_WIDTH] + + dbias_storage = smem.allocate(DbiasStorage) + sDbias = dbias_storage.sDbias.get_tensor( + cute.make_layout((self._TILE_ROWS, self._TILE_COLS), stride=(DBIAS_BUFF_WIDTH, 1)), + ) + _, tv_layout_dbias_write = cute.make_layout_tv( + thr_layout=cute.make_layout( + (self._TILE_ROWS, self._TILE_COLS // MXFP8_BLOCK_SCALING_SIZE), + stride=(self._TILE_COLS // MXFP8_BLOCK_SCALING_SIZE, 1), + ), + val_layout=cute.make_layout( + (1, MXFP8_BLOCK_SCALING_SIZE), stride=(MXFP8_BLOCK_SCALING_SIZE, 1) + ), + ) + sDbias_write = cute.composition(sDbias, tv_layout_dbias_write) + # Each thread start reading from the specfic bank based on its thread ID so they can do their best to access different banks + # to avoid bank conflict. + bank_group = (tidx % THREADS_PER_WARP) // self._THREADS_PER_BANK + # The offset this thread should start reading from based on what's its first bank to access. + offset = bank_group * self._PACK_SIZE + for w in cutlass.range_constexpr( + self._WAVES + ): # Each thread starts from this offset when writing into SMEM to avoid bank conflict + start = (w * self._PACK_SIZE + offset) % MXFP8_BLOCK_SCALING_SIZE + for i in cutlass.range_constexpr(self._PACK_SIZE): + # All threads write their per-thread partial sum results to the shared buffer. + sDbias_write[(tidx, start + i)] = rowwise_dbias_acc[w * self._PACK_SIZE + i] + cute.arch.sync_threads() + # All threads reduce the cross-thread partial sums to the per-block partial sum. + _, tv_layout_dbias_reduce = cute.make_layout_tv( + thr_layout=cute.make_layout((1, self._TILE_COLS), stride=(self._TILE_COLS, 1)), + val_layout=cute.make_layout((self._TILE_ROWS, 1), stride=(1, 1)), + ) + sDbias_reduce = cute.composition(sDbias, tv_layout_dbias_reduce) + # make_layout_tv yields a (thread, value) layout: thread=tidx -> column tidx, + # value=i -> row i. So index [tidx, i] (thread first), summing the column's rows. + block_dbias = Float32(0.0) + for i in cutlass.range_constexpr(self._TILE_ROWS): + block_dbias += sDbias_reduce[tidx, i] + return block_dbias + + @cute.jit + def _amax_epilogue(self, sAmax, mAmax, tidx, warp_idx, per_thread_amax): + # Reduce and get the per-warp amax. + warp_amax = cute.arch.warp_redux_sync(per_thread_amax, kind="fmax") + # Write the per-warp amax to shared memory + lane_idx = tidx % 32 + if lane_idx == 0: + sAmax[warp_idx] = warp_amax + cute.arch.sync_threads() + if tidx == 0: + cta_amax = Float32(0.0) + # The first thread reduces all the per-warp amax to the per-CTA amax + for w in cutlass.range_constexpr(self._NUM_WARPS): + cta_amax = cute.arch.fmax(cta_amax, sAmax[w]) + amax_i32 = cute.make_tensor( + cute.recast_ptr(mAmax.iterator, dtype=Int32), + cute.make_layout(1), + ) + # The first thread updates the global amax with an atomic max on the bitcasted float value + cute.arch.atomic_max( + amax_i32.iterator, + _bitcast_f32_to_i32(cta_amax), + ) + + @cute.jit + def _process_rowwise( + self, + sX_tile, # (TILE_Y, TILE_X) bf16/fp16 smem view, post-TMA + sO_row_tile, # (TILE_Y, TILE_X) uint8 smem view (rowwise FP8 output) + mS_row_stage, # rowwise scale tensor (1D swizzled, or 2D linear) + max_norm_rcp, + 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, + N, # Int32 — full input extents, for OOB masking + sActInput_tile=None, # (TILE_Y, TILE_X) act_input tile (dact only) + dbias_acc=None, # rmem Float32[MXFP8_BLOCK_SCALING_SIZE] 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, + 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, + 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_ROWWISE, + dbias_acc=dbias_acc, + ) + + @cute.jit + def _process_colwise( + self, + sX_tile, # (TILE_Y, TILE_X) bf16/fp16 smem view, post-TMA + sO_col_tile, # (TILE_Y, TILE_X) uint8 smem view (colwise FP8 output) + mS_col_stage, # colwise scale tensor (1D swizzled, or 2D linear) + max_norm_rcp, + 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, + N, # Int32 — full input extents, for OOB masking + sActInput_tile=None, # (TILE_Y, TILE_X) act_input tile (dact only) + ): + cfg = self.cfg + return quantize_colwise_mxfp8( + sX_tile, + sO_col_tile, + mS_col_stage, + 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, + WITH_ACT=cfg.WITH_ACT, + WITH_DACT=cfg.WITH_DACT, + sA_tile=sActInput_tile, + WITH_DBIAS=self.DBIAS_REDUCTION_COLWISE, + CACHE_ACTIVATION=self.CACHE_ACTIVATION, + ) + + +class MXFP8QuantizeSpecializedRowwiseKernel: + """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): + 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: Optional[cute.Tensor], # Unused, kept for API compatibility + 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] + + # The FFI boundary carries native FP8/E8M0 dtypes; the kernel works on bytes. + mO_row = as_byte_tensor(mO_row) + mS_row = as_byte_tensor(mS_row) + + 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, + self.cfg.MAX_NORM_RCP, + mX.element_type, + ).launch(grid=grid, block=block, stream=stream) + + @cute.kernel + def kernel(self, mX, mO_row, mS_row, max_norm_rcp, DTYPE): + """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, + ) + # Inputs widened to f32 once (reused by amax and the fused cvt). The FP8 + # output stays a uint8 fragment; we write it through a uint32 view so the + # 4-wide mul_cvt drops one packed word per call (see the cvt loop). + rX_f32 = cute.make_rmem_tensor( + cute.make_layout((1, MXFP8_BLOCK_SCALING_SIZE), stride=(MXFP8_BLOCK_SCALING_SIZE, 1)), + dtype=Float32, + ) + rO_thread = cute.make_rmem_tensor( + cute.make_layout((1, MXFP8_BLOCK_SCALING_SIZE), stride=(MXFP8_BLOCK_SCALING_SIZE, 1)), + dtype=Uint8, + ) + 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[Uint8, 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 an 1 Uint8 buffer for this thread's scale byte. + 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. + sS_thread[0] = Uint8(0) + 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) + + # Widen once and reduce. bf16/fp16 -> f32 widening is exact, so the + # amax matches the CUDA 16-bit abs_max path bit-for-bit. + amax = Float32(0.0) + for i in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE): + rX_f32[0, i] = Float32(rX_thread[0, i]) + amax = cute.arch.fmax(amax, fabs_f32(rX_f32[0, i])) + + biased_exp = cvt_f32_to_fp8e8m0(amax * max_norm_rcp) + if cutlass.const_expr(self._STASH_SCALE_TO_SMEM): + sS_thread[0] = Uint8(biased_exp) + else: + mS_thread[0] = Uint8(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) + mul_cvt4 = mul_i64_cvt_f32x4_to_fp8x4(self.cfg.FP8_DTYPE) + for i in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE // 4): + offset = 4 * i + rO_u32[i] = mul_cvt4( + rX_f32[0, offset], + rX_f32[0, offset + 1], + rX_f32[0, offset + 2], + rX_f32[0, offset + 3], + scale_2x, + ) + cute.autovec_copy(rO_thread, mO_thread) + + # Cooperative wide flush of the staged scales: the first CTA_Y*(CTA_X/G) + # threads each store one G-wide group, the rest idle. Pick the widest store + # the runtime row pitch allows (mirrors CUDA's PreferredDataType + runtime + # check): uint4 (16 scale bytes) when padded_cols % 16 == 0, else uint32 + # (4 bytes, always safe since padded_cols is a multiple of 4). A group never + # straddles the allocation boundary (base and padded_cols are both multiples + # of G), so flush a group iff its first column is in-allocation. Zeroed + # padding columns flush as 0. + if cutlass.const_expr(self._STASH_SCALE_TO_SMEM): + cute.arch.sync_threads() + padded_cols = mS_row.shape[1] + if 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, mS_tile, tidx, bidx, bidy, M, padded_cols, width): + """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: + """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_TILES = 4 # Each CTA handles 4 tiles in total side by side horizontally + _NUM_STAGES = 2 + # 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): + 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: Optional[cute.Tensor], + 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] + + # The FFI boundary carries native FP8/E8M0 dtypes; the kernel works on bytes. + mO_row = as_byte_tensor(mO_row) + mS_row = as_byte_tensor(mS_row) + mO_col = as_byte_tensor(mO_col) + mS_col = as_byte_tensor(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 NUM_TILES consecutive 64-col tiles = TILE_COLS*NUM_TILES cols. + cute.ceil_div(Int32(N), self._TILE_COLS * self._NUM_TILES), + cute.ceil_div(M, self._TILE_ROWS), + ] + block = [self._THREADS_PER_CTA] + self.kernel( + mX, + mS_row, + mS_col, + self.cfg.MAX_NORM_RCP, + 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, + mS_row, + mS_col, + max_norm_rcp, + dtype: cutlass.Constexpr[Type[cutlass.Numeric]], + tma_atom, + tma_src, + tma_atom_out_row, + tma_dst_out_row, + tma_atom_out_col, + tma_dst_out_col, + ): + """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 NUM_TILES consecutive 64-col tiles along N, starting here. + tile_n_base = bidx * self._NUM_TILES + + num_tiles = cutlass.min( + self._NUM_TILES, + # 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, self._TILE_COLS), + ) + + @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[Uint8, self._TILE_ROWS * self._TILE_COLS * self._NUM_STAGES], + 128, + ] + sO_col: cute.struct.Align[ + cute.struct.MemRange[Uint8, 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[ + Uint8, + 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[Uint8, 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)) + ) + # Reshape the rowwise scale SMEM buffer + sScaleRow = storage.sScaleRow.get_tensor( + cute.make_layout( + (self._TILE_ROWS, (self._SCALE_COLS, self._NUM_TILES)), + stride=(self._SCALE_COLS * self._NUM_TILES, (1, self._SCALE_COLS)), + ) + ) + # Reshape the colwise scale SMEM buffer + sScaleCol = storage.sScaleCol.get_tensor( + cute.make_layout( + ((self._SCALE_ROWS, self._TILE_COLS), self._NUM_TILES), + stride=((0, 1), 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(sScaleRow.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(sScaleCol.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: + mainloop_pipeline.producer_acquire(prod_state) + cute.copy( + tma_atom, + tXgX[(None, (bidy, tile_n_base + s))], + 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 + + # 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, (None, tile_idx)], + sScaleCol[(None, tile_idx)], + sColReduce[ + None, warp_idx + ], # Pick the per-warp colwise-reduce scratchpad for this warp + self._WARPS_PER_CTA, + 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_x = bidx * self._NUM_TILES + tile_idx + cute.copy( + tma_atom_out_row, tXsO_row[(None, stage_idx)], tXgO_row[(None, (bidy, tile_x))] + ) + cute.copy( + tma_atom_out_col, tXsO_col[(None, stage_idx)], tXgO_col[(None, (bidy, tile_x))] + ) + 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: + mainloop_pipeline.producer_acquire(prod_state) + cute.copy( + tma_atom, + tXgX[(None, (bidy, tile_n_base + next_tile_idx))], + tXsX[(None, prod_state.index)], + tma_bar_ptr=mainloop_pipeline.producer_get_barrier(prod_state), + ) + mainloop_pipeline.producer_commit(prod_state) + prod_state.advance() + + mS_row_N = mS_row.shape[1] + if mS_row_N % 8 == 0: + self._flush_rowwise_scales_to_gmem(sScaleRow, mS_row, tidx, bidx, bidy, M, mS_row_N, 8) + else: + self._flush_rowwise_scales_to_gmem(sScaleRow, mS_row, tidx, bidx, bidy, M, mS_row_N, 4) + self._flush_colwise_scales_to_gmem(sScaleCol, mS_col, tidx, bidx, bidy) + + # 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_rowwise_scales_to_gmem( + self, sScaleRow, mS_row, tidx, bidx, bidy, M, mS_row_N, WIDTH: cutlass.Constexpr + ): + """Flush the staged (TILE_ROWS, TILE_COLS / MXFP8_BLOCK_SCALING_SIZE * NUM_TILES) SMEM rowwise scale block to gmem + using `width` bytes per store""" + + ROWS = self._TILE_ROWS + COLS = self._NUM_TILES * self._SCALE_COLS + 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)), + ) + # View sScaleRow as a 2D tensor: + # (_TILE_ROWS, (_SCALE_COLS, _NUM_TILES)):(_NUM_TILES * _SCALE_COLS, (1, _SCALE_COLS)) -> + # (_TILE_ROWS, _SCALE_COLS * _NUM_TILES):(_NUM_TILES * _SCALE_COLS, 1) + sScaleRow2D = cute.make_tensor( + sScaleRow.iterator, + cute.make_layout((ROWS, COLS), stride=(COLS, 1)), + ) + # Obtain the GMEM slice for the output rowwise scale factor for this CTA + mS_row_tile = cute.local_tile(mS_row, (ROWS, COLS), (bidy, bidx)) + # Each flush slot stores one `WIDTH`-byte group; the CTA strides over the + # slots when there are more slots than threads. + 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 rowwise scale factor buffer + thread_y = bidy * ROWS + thread_idx // ACTIVE_THREAD_COLS + thread_x = bidx * COLS + (thread_idx % ACTIVE_THREAD_COLS) * WIDTH + # Only copy if the slot's position is within the valid GMEM rowwise scale factor buffer + if thread_y < M and thread_x < mS_row_N: + cute.autovec_copy( + cute.composition(sScaleRow2D, tv_flush_layout)[thread_idx, None], + cute.composition(mS_row_tile, tv_flush_layout)[thread_idx, None], + ) + + @cute.jit + def _flush_colwise_scales_to_gmem(self, sScaleCol, mS_col, tidx, bidx, bidy): + """Flush the staged (TILE_ROWS / MXFP8_BLOCK_SCALING_SIZE, TILE_COLS*NUM_TILES) SMEM colwise scale block to gmem + using with 16-byte stores.""" + + ROWS = self._SCALE_ROWS + COLS = self._NUM_TILES * self._TILE_COLS + width = 16 # span and the gmem row pitch (mult of 128) are both 16B-aligned + mS_col_N = mS_col.shape[1] + # View sScaleCol as a 2D tensor: + # ((_SCALE_ROWS, TILE_COLS), NUM_TILES):((0, 1), TILE_COLS) -> + # (_SCALE_ROWS, TILE_COLS * NUM_TILES):(TILE_COLS * NUM_TILES, 1) + sScaleCol2D = cute.make_tensor( + cute.recast_ptr(sScaleCol.iterator, dtype=Uint8), + cute.make_layout((ROWS, COLS), stride=(0, 1)), + ) + mS_col_tile = cute.local_tile(mS_col, (1, COLS), (bidy, bidx)) + ACTIVE_THREAD_COLS = COLS // width + _, tv_flush = cute.make_layout_tv( + thr_layout=cute.make_layout((1, 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: + col = bidx * COLS + thread_idx * width + if col < mS_col_N: + cute.autovec_copy( + cute.composition(sScaleCol2D, tv_flush)[thread_idx, None], + cute.composition(mS_col_tile, tv_flush)[thread_idx, None], + ) + + +def get_kernel_class(cfg): + """If no fusion is involved and the kernel only quantizes, dispatch to the specialized kernel for better performance.""" + plain_cast_only = ( + not cfg.WITH_GEMM_SWIZZLED_SCALES + and not cfg.WITH_AMAX + and not cfg.WITH_DBIAS + and not cfg.WITH_DACT + and not cfg.WITH_ACT + and not cfg.WITH_NOOP + ) + # Only dispatch to the specialized kernels for packed16 types (bf16/fp16) + if plain_cast_only and is_packed16(cfg.DTYPE): + if cfg.ROWWISE and not cfg.COLWISE: + return MXFP8QuantizeSpecializedRowwiseKernel + if cfg.ROWWISE and cfg.COLWISE: + return MXFP8QuantizeSpecializedBidimensionalKernel + return MXFP8QuantizeKernel + + +def compile_cutedsl_function_from_cfg(cfg): + """ + Return the compiled CuTeDSL function object for the given MXFP8 quantization config. + """ + + # Route plain cast-only configs to the matching specialized kernel (mirrors the + # CUDA dispatcher); everything else uses the general standard kernel. + kernel_class = get_kernel_class(cfg) + kernel_obj = kernel_class(cfg) + # M, N must be divisible by the MXFP8 scale-block size (MXFP8_BLOCK_SCALING_SIZE = 32) — the + # same alignment the CUDA C++ kernel requires. + sym_M = cute.sym_int32(divisibility=MXFP8_BLOCK_SCALING_SIZE) + sym_N = cute.sym_int32(divisibility=MXFP8_BLOCK_SCALING_SIZE) + 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(N // 32, 4)) + # columnwise: (roundup(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); the kernels view these buffers as raw bytes internally. + out_dtype = cutlass.Float8E4M3FN if cfg.FP8_DTYPE == "e4m3" else cutlass.Float8E5M2 + 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 + ) + noop_fake = ( + cute.runtime.make_fake_compact_tensor( + Float32, (1,), stride_order=(0,), memspace=cute.AddressSpace.gmem, assumed_align=4 + ) + if cfg.WITH_NOOP + else None + ) + 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 (1-element cast_noop flag) + 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, + with_noop: 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. + """ + # 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 + + 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, + with_noop=with_noop, + 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) + try: + compiled = compile_cutedsl_function_from_cfg(cfg) + except Exception as e: # pylint: disable=broad-exception-caught + logger.warning( + "CuTeDSL 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/CuTeDSL/utils.py b/transformer_engine/common/CuTeDSL/utils.py new file mode 100644 index 0000000000..001c862e92 --- /dev/null +++ b/transformer_engine/common/CuTeDSL/utils.py @@ -0,0 +1,331 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Low-level CuTeDSL helpers: bitcast/fma/exp2 intrinsics, f32 packing, and the 16-bit (bf16/fp16) packed-op kit.""" + +import functools +import logging +from types import SimpleNamespace + +import cutlass +from cutlass import Float32, Int64, Int32, Int16 +from cutlass._mlir.dialects import arith as mlir_arith +from cutlass._mlir.dialects import llvm +from cutlass.cutlass_dsl import T, dsl_user_op + +_CUTLASS_DTYPE_FROM_STR = { + "fp32": cutlass.Float32, + "fp16": cutlass.Float16, + "bf16": cutlass.BFloat16, +} +_STR_FROM_CUTLASS_DTYPE = {v: k for k, v in _CUTLASS_DTYPE_FROM_STR.items()} + +logger = logging.getLogger("transformer_engine.cutedsl.utils") + + +@functools.lru_cache(maxsize=None) +def device_compute_capability() -> tuple: + """(major, minor) of CUDA device 0, or (0, 0) if it can't be queried. + + Used to pick per-arch kernel shapes at compile time. Device 0 stands in for + the whole node (homogeneous-GPU assumption, same as the CUTE_DSL_ARCH env).""" + from cuda.core import Device # pylint: disable=no-name-in-module + + major_minor = Device().arch # compute capability as digits, e.g. "120" + return (int(major_minor[:-1]), int(major_minor[-1])) if major_minor else (0, 0) + + +@functools.lru_cache(maxsize=None) +def device_is_blackwell() -> bool: + """Return True for the Blackwell family (SM 10.0 / 11.0 / 12.0) + This is a run-time check, not a compile-time check. It check if the current device is Blackwell architecture. + """ + major, minor = device_compute_capability() + return ( + (major == 10 and minor == 0) or (major == 11 and minor == 0) or (major == 12 and minor == 0) + ) + + +def str_to_cutlass_dtype(dtype_str: str): + """Convert a string dtype to a cutlass dtype, or None if unknown.""" + return _CUTLASS_DTYPE_FROM_STR.get(dtype_str, None) + + +def cutlass_dtype_to_str(dtype): + """Convert a cutlass dtype back to its protocol string, or None if unknown.""" + return _STR_FROM_CUTLASS_DTYPE.get(dtype, None) + + +FP32_MANTISSA_BITS = 23 + + +@dsl_user_op +def _bitcast_f32_to_i32(val: Float32, *, loc=None, ip=None) -> Int32: + """Bitcast a float32 value to int32 without changing the bit pattern.""" + return Int32(mlir_arith.bitcast(T.i32(), val.ir_value(loc=loc, ip=ip), loc=loc, ip=ip)) + + +@dsl_user_op +def _bitcast_i32_to_f32(val: Int32, *, loc=None, ip=None) -> Float32: + """Bitcast an int32 value to float32 without changing the bit pattern.""" + return Float32(mlir_arith.bitcast(T.f32(), val.ir_value(loc=loc, ip=ip), loc=loc, ip=ip)) + + +@dsl_user_op +def fabs_f32(val: Float32, *, loc=None, ip=None) -> Float32: + """Compute the absolute value of a float32.""" + val_i32 = _bitcast_f32_to_i32(val, loc=loc, ip=ip) + abs_i32 = val_i32 & Int32(0x7FFFFFFF) + return _bitcast_i32_to_f32(abs_i32, loc=loc, ip=ip) + + +@dsl_user_op +def fma_f32(a: Float32, b: Float32, c: Float32, *, loc=None, ip=None) -> Float32: + """Compute the fused multiply-add of three float32 values: a * b + c.""" + return Float32( + llvm.inline_asm( + T.f32(), + [a.ir_value(loc=loc, ip=ip), b.ir_value(loc=loc, ip=ip), c.ir_value(loc=loc, ip=ip)], + "fma.rn.f32 $0, $1, $2, $3;", + "=f,f,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + +@dsl_user_op +def exp2f_rcp(biased_exp: Int32, *, loc=None, ip=None) -> Float32: + """2^(127 - biased_exp) with special-case handling.""" + new_exp = (Int32(254) - biased_exp) << Int32(FP32_MANTISSA_BITS) + result = _bitcast_i32_to_f32(new_exp, loc=loc, ip=ip) + for cmp_val, repl_bits in [(255, 0x7FFFFFFF), (254, 0x00400000), (0, 0x7F000000)]: + cond = mlir_arith.cmpi( + mlir_arith.CmpIPredicate.eq, + biased_exp.ir_value(loc=loc, ip=ip), + Int32(cmp_val).ir_value(loc=loc, ip=ip), + loc=loc, + ip=ip, + ) + alt = _bitcast_i32_to_f32(Int32(repl_bits), loc=loc, ip=ip) + result = Float32( + mlir_arith.select( + cond, alt.ir_value(loc=loc, ip=ip), result.ir_value(loc=loc, ip=ip), loc=loc, ip=ip + ) + ) + return result + + +@dsl_user_op +def pack_f32x2(lo: Float32, hi: Float32, *, loc=None, ip=None) -> Int64: + """Pack two f32 scalars into a single 64-bit register (`floatx2` layout). + + Low 32 bits = `lo`, high 32 bits = `hi`. Uses `mov.b64 %dst, {%lo, %hi};` + which lowers to a single register move — no actual memory traffic. + """ + return Int64( + llvm.inline_asm( + T.i64(), + [lo.ir_value(loc=loc, ip=ip), hi.ir_value(loc=loc, ip=ip)], + "mov.b64 $0, {$1, $2};", + "=l,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + +@dsl_user_op +def unpack_i64_to_i32x2(v: Int64, *, loc=None, ip=None): + """Split a 64-bit value into (lo, hi) 32-bit halves. + + Inverse of pack_f32x2's register-pair layout. Lowers to register-pair + aliasing in SASS (no real instructions), so an 8-byte smem load + this + split costs one LDS.64 total.""" + lo = Int32(mlir_arith.trunci(T.i32(), v.ir_value(loc=loc, ip=ip), loc=loc, ip=ip)) + hi_64 = mlir_arith.shrui( + v.ir_value(loc=loc, ip=ip), Int64(32).ir_value(loc=loc, ip=ip), loc=loc, ip=ip + ) + hi = Int32(mlir_arith.trunci(T.i32(), hi_64, loc=loc, ip=ip)) + return lo, hi + + +def _build_packed16_kit(in_fmt: str): + """Build a kit of PTX wrappers for a 16-bit input format so we don't have to repeat + the same inline asm boilerplate code for FP16 and BF16 dtypes. + + `in_fmt` is the PTX format string ('bf16' or 'f16'). Returns a namespace + with the per-format ops the rowwise/colwise inner loops need: + + abs_max_x2(Int32, Int32) -> Int32 # `max.xorsign.abs.x2` + abs_max_scalar(Int16, Int16) -> Int16 # `max.xorsign.abs.` + bits_to_f32(Int16) -> Float32 # widen one 16-bit element + x2_lo_to_f32(Int32) -> Float32 # extract+widen low half + x2_hi_to_f32(Int32) -> Float32 # extract+widen high half + """ + + @dsl_user_op + def abs_max_x2(a: Int32, b: Int32, *, loc=None, ip=None) -> Int32: + return Int32( + llvm.inline_asm( + T.i32(), + [a.ir_value(loc=loc, ip=ip), b.ir_value(loc=loc, ip=ip)], + f"max.xorsign.abs.{in_fmt}x2 $0, $1, $2;", + "=r,r,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + @dsl_user_op + def max_x2(a: Int32, b: Int32, *, loc=None, ip=None) -> Int32: + return Int32( + llvm.inline_asm( + T.i32(), + [a.ir_value(loc=loc, ip=ip), b.ir_value(loc=loc, ip=ip)], + f"max.{in_fmt}x2 $0, $1, $2;", + "=r,r,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + @dsl_user_op + def abs_max_scalar(a: Int16, b: Int16, *, loc=None, ip=None) -> Int16: + return Int16( + llvm.inline_asm( + T.i16(), + [a.ir_value(loc=loc, ip=ip), b.ir_value(loc=loc, ip=ip)], + f"max.xorsign.abs.{in_fmt} $0, $1, $2;", + "=h,h,h", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + if in_fmt == "bf16": + # bf16 == top 16 bits of f32 — widening is a free bit-shift. + @dsl_user_op + def bits_to_f32(bits: Int16, *, loc=None, ip=None) -> Float32: + i32 = Int32(mlir_arith.extui(T.i32(), bits.ir_value(loc=loc, ip=ip), loc=loc, ip=ip)) + return _bitcast_i32_to_f32(i32 << Int32(16), loc=loc, ip=ip) + + @dsl_user_op + def x2_lo_to_f32(bits: Int32, *, loc=None, ip=None) -> Float32: + return _bitcast_i32_to_f32((bits & Int32(0xFFFF)) << Int32(16), loc=loc, ip=ip) + + @dsl_user_op + def x2_hi_to_f32(bits: Int32, *, loc=None, ip=None) -> Float32: + # `(x >> 16) << 16` ≡ `x & 0xFFFF0000`, sidestepping signed-literal + # issues. Sign bits from the arith-right shift get zeroed by the + # left shift. + return _bitcast_i32_to_f32((bits >> Int32(16)) << Int32(16), loc=loc, ip=ip) + + @dsl_user_op + def truncate_f32(val: Float32, *, loc=None, ip=None) -> Float32: + """Round f32 to bf16 precision (round-to-nearest-even), keep f32. + Matches C++'s `static_cast(static_cast(elt))`.""" + bf16_bits = Int16( + llvm.inline_asm( + T.i16(), + [val.ir_value(loc=loc, ip=ip)], + "cvt.rn.bf16.f32 $0, $1;", + "=h,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + i32 = Int32( + mlir_arith.extui(T.i32(), bf16_bits.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + ) + return _bitcast_i32_to_f32(i32 << Int32(16), loc=loc, ip=ip) + + else: + # f16 has its own bit layout; widening requires `cvt.f32.f16`. + @dsl_user_op + def bits_to_f32(bits: Int16, *, loc=None, ip=None) -> Float32: + return Float32( + llvm.inline_asm( + T.f32(), + [bits.ir_value(loc=loc, ip=ip)], + "cvt.f32.f16 $0, $1;", + "=f,h", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + @dsl_user_op + def x2_lo_to_f32(bits: Int32, *, loc=None, ip=None) -> Float32: + lo_i16 = Int16( + mlir_arith.trunci(T.i16(), bits.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + ) + return bits_to_f32(lo_i16, loc=loc, ip=ip) + + @dsl_user_op + def x2_hi_to_f32(bits: Int32, *, loc=None, ip=None) -> Float32: + hi_shifted = bits >> Int32(16) + hi_i16 = Int16( + mlir_arith.trunci(T.i16(), hi_shifted.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + ) + return bits_to_f32(hi_i16, loc=loc, ip=ip) + + @dsl_user_op + def truncate_f32(val: Float32, *, loc=None, ip=None) -> Float32: + """Round f32 to f16 precision, keep f32.""" + f16_bits = Int16( + llvm.inline_asm( + T.i16(), + [val.ir_value(loc=loc, ip=ip)], + "cvt.rn.f16.f32 $0, $1;", + "=h,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + return Float32( + llvm.inline_asm( + T.f32(), + [f16_bits.ir_value(loc=loc, ip=ip)], + "cvt.f32.f16 $0, $1;", + "=f,h", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + return SimpleNamespace( + max_x2=max_x2, + abs_max_x2=abs_max_x2, + abs_max_scalar=abs_max_scalar, + bits_to_f32=bits_to_f32, + x2_lo_to_f32=x2_lo_to_f32, + x2_hi_to_f32=x2_hi_to_f32, + truncate_f32=truncate_f32, + ) + + +_BF16_KIT = _build_packed16_kit("bf16") +_F16_KIT = _build_packed16_kit("f16") + + +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 + + +def packed16_kit(dtype): + """Trace-time selector — pick a Packed16Kit for the input dtype.""" + if dtype is cutlass.Float16: + return _F16_KIT + return _BF16_KIT diff --git a/transformer_engine/common/CuTeDSL/utils_fp8.py b/transformer_engine/common/CuTeDSL/utils_fp8.py new file mode 100644 index 0000000000..f496ead553 --- /dev/null +++ b/transformer_engine/common/CuTeDSL/utils_fp8.py @@ -0,0 +1,380 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""FP8 conversion helpers (f32<->fp8 e4m3/e5m2/e8m0, fused mul+cvt PTX wrappers) for the CuTeDSL kernels.""" + +import logging + +import cutlass +from cutlass import Uint8, cute +from cutlass import Float32, Int64, Int32, Int16, Uint32 +from cutlass._mlir.dialects import arith as mlir_arith +from cutlass._mlir.dialects import llvm +from cutlass.cutlass_dsl import T, dsl_user_op + +from transformer_engine.common.CuTeDSL.utils import ( + FP32_MANTISSA_BITS, + _bitcast_f32_to_i32, + device_is_blackwell, +) + +logger = logging.getLogger("transformer_engine.cutedsl.utils_fp8") + + +def as_byte_tensor(t): + """View an FP8/E8M0 gmem tensor as raw bytes (the kernels build/store the bit patterns).""" + return cute.make_tensor(cute.recast_ptr(t.iterator, dtype=Uint8), t.layout) + + +@dsl_user_op +def cvt_f32_to_fp8e4m3(val: Float32, *, loc=None, ip=None) -> Int32: + """float32 -> fp8e4m3 conversion.""" + zero = Float32(0.0) + result_i16 = Int16( + llvm.inline_asm( + T.i16(), + [zero.ir_value(loc=loc, ip=ip), val.ir_value(loc=loc, ip=ip)], + "cvt.rn.satfinite.e4m3x2.f32 $0, $1, $2;", + "=h,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + result_i32 = Int32( + mlir_arith.extui(T.i32(), result_i16.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + ) + return result_i32 & Int32(0xFF) + + +@dsl_user_op +def cvt_f32_to_fp8e5m2(val: Float32, *, loc=None, ip=None) -> Int32: + """float32 -> fp8e5m2 conversion.""" + zero = Float32(0.0) + result_i16 = Int16( + llvm.inline_asm( + T.i16(), + [zero.ir_value(loc=loc, ip=ip), val.ir_value(loc=loc, ip=ip)], + "cvt.rn.satfinite.e5m2x2.f32 $0, $1, $2;", + "=h,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + result_i32 = Int32( + mlir_arith.extui(T.i32(), result_i16.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + ) + return result_i32 & Int32(0xFF) + + +@dsl_user_op +def cvt_f32_to_fp8e8m0_non_blackwell(val: Float32, *, loc=None, ip=None) -> Int32: + """float32 -> fp8e8m0 conversion (generic, pre-Blackwell). + + Software round-up of the biased exponent, mirroring ptx::float_to_e8m0's + non-Blackwell branch (transformer_engine/common/util/ptx.cuh).""" + val_i32 = _bitcast_f32_to_i32(val, loc=loc, ip=ip) + rounded = val_i32 + Int32(0x7FFFFF) + exponent = (rounded >> Int32(FP32_MANTISSA_BITS)) & Int32(0xFF) + return Int32( + mlir_arith.minsi( + exponent.ir_value(loc=loc, ip=ip), Int32(254).ir_value(loc=loc, ip=ip), loc=loc, ip=ip + ) + ) + + +@dsl_user_op +def cvt_f32_to_fp8e8m0_blackwell(val: Float32, *, loc=None, ip=None) -> Int32: + """float32 -> fp8e8m0 conversion (Blackwell, SM >= 100). + + Uses the hardware cvt.rp.satfinite.ue8m0x2.f32 instruction, mirroring + ptx::float_to_e8m0's Blackwell branch. The x2 form packs two e8m0 bytes; + we feed (0.0, val) so the low byte is e8m0(val) and mask it out.""" + zero = Float32(0.0) + result_i16 = Int16( + llvm.inline_asm( + T.i16(), + [zero.ir_value(loc=loc, ip=ip), val.ir_value(loc=loc, ip=ip)], + "cvt.rp.satfinite.ue8m0x2.f32 $0, $1, $2;", + "=h,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + result_i32 = Int32( + mlir_arith.extui(T.i32(), result_i16.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + ) + return result_i32 & Int32(0xFF) + + +def _build_mul_i64_cvt_f32x4_to_fp8x4(out_fmt: str, relu: bool = False): + """Build a fused 4-wide `f32x4 * f32x2 -> fp8x4` PTX wrapper. + + Multiplies four f32 inputs by a broadcast inverse scale (passed as an + f32x2 pack of (s, s)) and converts to FP8, packing the four bytes into one + uint32: byte i = fp8(v_i * s). Two `mul.f32x2` + two `cvt...x2.f32` — the + f32-input form of this op family (CUDA ptx::mul_cvt_4x). + """ + out_op = "e4m3x2" if out_fmt == "e4m3" else "e5m2x2" + asm = ( + "{\n" + ".reg.b64 vp0; .reg.b64 vp1; .reg.b64 vp2; .reg.b64 vp3;\n\t" + ".reg.b32 vs0; .reg.b32 vs1; .reg.b32 vs2; .reg.b32 vs3;\n\t" + ".reg.b16 vo0; .reg.b16 vo1;\n\t" + "mov.b64 vp0, {$1, $2};\n\t" + "mov.b64 vp2, {$3, $4};\n\t" + "mul.f32x2 vp1, vp0, $5;\n\t" + "mul.f32x2 vp3, vp2, $5;\n\t" + "mov.b64 {vs0, vs1}, vp1;\n\t" + "mov.b64 {vs2, vs3}, vp3;\n\t" + # cvt d, a, b => d[15:8]=fp8(a), d[7:0]=fp8(b); feed (hi, lo) so the low + # byte holds the earlier element. + f"cvt.rn.satfinite{".relu" if relu else ""}.{out_op}.f32 vo0, vs1, vs0;\n\t" + f"cvt.rn.satfinite{".relu" if relu else ""}.{out_op}.f32 vo1, vs3, vs2;\n\t" + "mov.b32 $0, {vo0, vo1};\n\t" + "}" + ) + + @dsl_user_op + def fn( + v0: Float32, v1: Float32, v2: Float32, v3: Float32, scale_2x: Int64, *, loc=None, ip=None + ) -> Uint32: + return Uint32( + llvm.inline_asm( + T.i32(), + [ + v0.ir_value(loc=loc, ip=ip), + v1.ir_value(loc=loc, ip=ip), + v2.ir_value(loc=loc, ip=ip), + v3.ir_value(loc=loc, ip=ip), + scale_2x.ir_value(loc=loc, ip=ip), + ], + asm, + "=r,f,f,f,f,l", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + return fn + + +def mul_i64_cvt_f32x4_to_fp8x4(fp8_dtype: str, relu: bool = False): + """Return the fused 4-wide f32->FP8 multiply+cast op for the given FP8 format. + + The op takes (v0, v1, v2, v3, scale_2x) and returns a uint32 of four packed + fp8 bytes, byte i = fp8(v_i * scale). `scale_2x` is pack_f32x2(s, s).""" + return _build_mul_i64_cvt_f32x4_to_fp8x4("e5m2" if fp8_dtype == "e5m2" else "e4m3", relu) + + +def _build_mul_f32x4_cvt_f32x4_to_fp8x4(out_fmt: str, relu: bool = False): + """Build a fused elementwise `f32x4 * f32x4 -> fp8x4` PTX wrapper. + + General elementwise multiply-and-convert: byte i = fp8(a_i * b_i). Two + fma.rn.f32x2 against packed zeros + two cvt...x2.f32 (same sequence as + CUDA's mul_cvt_4x(out, floatx4, floatx4) in ptx.cuh). fma (not mul) so a + -0 product flushes to +0.""" + out_op = "e4m3x2" if out_fmt == "e4m3" else "e5m2x2" + asm = ( + "{\n" + ".reg.b64 va0; .reg.b64 va1; .reg.b64 vb0; .reg.b64 vb1;\n\t" + ".reg.b64 vr0; .reg.b64 vr1; .reg.b64 zeros;\n\t" + ".reg.b32 vs0; .reg.b32 vs1; .reg.b32 vs2; .reg.b32 vs3;\n\t" + ".reg.b16 vo0; .reg.b16 vo1;\n\t" + "mov.b64 zeros, {0x0, 0x0};\n\t" + "mov.b64 va0, {$1, $2};\n\t" + "mov.b64 va1, {$3, $4};\n\t" + "mov.b64 vb0, {$5, $6};\n\t" + "mov.b64 vb1, {$7, $8};\n\t" + "fma.rn.f32x2 vr0, va0, vb0, zeros;\n\t" + "fma.rn.f32x2 vr1, va1, vb1, zeros;\n\t" + "mov.b64 {vs0, vs1}, vr0;\n\t" + "mov.b64 {vs2, vs3}, vr1;\n\t" + # cvt d, a, b => d[15:8]=fp8(a), d[7:0]=fp8(b); feed (hi, lo) so the low + # byte holds the earlier element. + f"cvt.rn.satfinite{".relu" if relu else ""}.{out_op}.f32 vo0, vs1, vs0;\n\t" + f"cvt.rn.satfinite{".relu" if relu else ""}.{out_op}.f32 vo1, vs3, vs2;\n\t" + "mov.b32 $0, {vo0, vo1};\n\t" + "}" + ) + + @dsl_user_op + def fn( + a0: Float32, + a1: Float32, + a2: Float32, + a3: Float32, + b0: Float32, + b1: Float32, + b2: Float32, + b3: Float32, + *, + loc=None, + ip=None, + ) -> Uint32: + return Uint32( + llvm.inline_asm( + T.i32(), + [ + a0.ir_value(loc=loc, ip=ip), + a1.ir_value(loc=loc, ip=ip), + a2.ir_value(loc=loc, ip=ip), + a3.ir_value(loc=loc, ip=ip), + b0.ir_value(loc=loc, ip=ip), + b1.ir_value(loc=loc, ip=ip), + b2.ir_value(loc=loc, ip=ip), + b3.ir_value(loc=loc, ip=ip), + ], + asm, + "=r,f,f,f,f,f,f,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + return fn + + +def mul_f32x4_cvt_f32x4_to_fp8x4(fp8_dtype: str, relu: bool = False): + """Return the fused elementwise f32x4*f32x4 -> FP8x4 op for the given FP8 format. + + The op takes (a0..a3, b0..b3) and returns a uint32 of four packed fp8 + bytes, byte i = fp8(a_i * b_i).""" + return _build_mul_f32x4_cvt_f32x4_to_fp8x4("e5m2" if fp8_dtype == "e5m2" else "e4m3", relu) + + +def _build_mul_i64_cvt_packed16x4_to_fp8x4(in_fmt: str, out_fmt: str, relu: bool = False): + """Build a fused `2x x2 * f32x2 -> fp8x4` PTX wrapper. + + 16-bit-input form of _build_mul_i64_cvt_f32x4_to_fp8x4: widens four packed + bf16/f16 elements to f32 inside the asm, multiplies by the broadcast + (s, s) pair, and converts: byte i = fp8(elt_i * s). The two u16 cvt + results combine into the u32 via a register-pair mov (free).""" + out_op = "e4m3x2" if out_fmt == "e4m3" else "e5m2x2" + asm = ( + "{\n" + ".reg.b64 vp0; .reg.b64 vp1; .reg.b64 vq0; .reg.b64 vq1;\n\t" + ".reg.b32 v1; .reg.b32 v2; .reg.b32 v3; .reg.b32 v4;\n\t" + ".reg.b16 vb1; .reg.b16 vb2; .reg.b16 vb3; .reg.b16 vb4;\n\t" + ".reg.b16 vo0; .reg.b16 vo1;\n\t" + "mov.b32 {vb1, vb2}, $1;\n\t" + "mov.b32 {vb3, vb4}, $2;\n\t" + f"cvt.f32.{in_fmt} v1, vb1;\n\t" + f"cvt.f32.{in_fmt} v2, vb2;\n\t" + f"cvt.f32.{in_fmt} v3, vb3;\n\t" + f"cvt.f32.{in_fmt} v4, vb4;\n\t" + "mov.b64 vp0, {v1, v2};\n\t" + "mov.b64 vq0, {v3, v4};\n\t" + "mul.f32x2 vp1, vp0, $3;\n\t" + "mul.f32x2 vq1, vq0, $3;\n\t" + "mov.b64 {v2, v1}, vp1;\n\t" + "mov.b64 {v4, v3}, vq1;\n\t" + # cvt d, a, b => d[15:8]=fp8(a), d[7:0]=fp8(b); feed (hi, lo) so the + # low byte holds the earlier element. + f"cvt.rn.satfinite{".relu" if relu else ""}.{out_op}.f32 vo0, v1, v2;\n\t" + f"cvt.rn.satfinite{".relu" if relu else ""}.{out_op}.f32 vo1, v3, v4;\n\t" + "mov.b32 $0, {vo0, vo1};\n\t" + "}" + ) + + @dsl_user_op + def fn(lo_2x: Int32, hi_2x: Int32, scale_2x: Int64, *, loc=None, ip=None) -> Uint32: + return Uint32( + llvm.inline_asm( + T.i32(), + [ + lo_2x.ir_value(loc=loc, ip=ip), + hi_2x.ir_value(loc=loc, ip=ip), + scale_2x.ir_value(loc=loc, ip=ip), + ], + asm, + "=r,r,r,l", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + return fn + + +def mul_i64_cvt_packed16x4_to_fp8x4(dtype, fp8_dtype: str, relu: bool = False): + """Return the fused packed16x4 * broadcast-scale -> FP8x4 op. + + The op takes (lo_2x, hi_2x, scale_2x): two i32s of packed bf16/f16 pairs + (elements 0-1, 2-3) and a pack_f32x2(s, s) pair; returns a uint32 of four + packed fp8 bytes, byte i = fp8(elt_i * s).""" + in_fmt = "f16" if dtype is cutlass.Float16 else "bf16" + return _build_mul_i64_cvt_packed16x4_to_fp8x4( + in_fmt, "e5m2" if fp8_dtype == "e5m2" else "e4m3", relu + ) + + +# Pick the appropriate float32 -> fp8e8m0 conversion function based on the target architecture. +# Blackwell (SM >= 100) has a hardware instruction for this, while older architectures require a software implementation. +cvt_f32_to_fp8e8m0 = ( + cvt_f32_to_fp8e8m0_blackwell if device_is_blackwell() else cvt_f32_to_fp8e8m0_non_blackwell +) + + +@dsl_user_op +def cvt_f32x2_to_fp8e4m3x2( + val_hi: Float32, val_lo: Float32, relu: bool = False, *, loc=None, ip=None +) -> Int32: + """Convert two float32 values to two packed fp8e4m3fn bytes in one instruction. + + Returns an int32 where bits [7:0] = fp8(val_lo), bits [15:8] = fp8(val_hi). + """ + result_i16 = Int16( + llvm.inline_asm( + T.i16(), + [val_hi.ir_value(loc=loc, ip=ip), val_lo.ir_value(loc=loc, ip=ip)], + f"cvt.rn.satfinite{".relu" if relu else ""}.e4m3x2.f32 $0, $1, $2;", + "=h,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + return Int32(mlir_arith.extui(T.i32(), result_i16.ir_value(loc=loc, ip=ip), loc=loc, ip=ip)) + + +@dsl_user_op +def cvt_f32x2_to_fp8e5m2x2( + val_hi: Float32, val_lo: Float32, relu: bool = False, *, loc=None, ip=None +) -> Int32: + """Convert two float32 values to two packed fp8e5m2 bytes in one instruction. + + Returns an int32 where bits [7:0] = fp8(val_lo), bits [15:8] = fp8(val_hi). + """ + result_i16 = Int16( + llvm.inline_asm( + T.i16(), + [val_hi.ir_value(loc=loc, ip=ip), val_lo.ir_value(loc=loc, ip=ip)], + f"cvt.rn.satfinite{".relu" if relu else ""}.e5m2x2.f32 $0, $1, $2;", + "=h,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + return Int32(mlir_arith.extui(T.i32(), result_i16.ir_value(loc=loc, ip=ip), loc=loc, ip=ip)) + + +def get_cvt_f32_to_fp8_func(fp8_dtype: str): + """Returns the float32 -> float8 conversion function for the given FP8 format.""" + if fp8_dtype == "e5m2": + return cvt_f32_to_fp8e5m2 + return cvt_f32_to_fp8e4m3 + + +def get_cvt_f32x2_to_fp8x2_func(fp8_dtype: str): + """Returns the float32x2 -> float8x2 conversion function for the given FP8 format.""" + if fp8_dtype == "e5m2": + return cvt_f32x2_to_fp8e5m2x2 + return cvt_f32x2_to_fp8e4m3x2 diff --git a/transformer_engine/common/__init__.py b/transformer_engine/common/__init__.py index 231680321e..e221039590 100644 --- a/transformer_engine/common/__init__.py +++ b/transformer_engine/common/__init__.py @@ -15,6 +15,7 @@ import subprocess import sys import sysconfig +import logging from typing import Optional, Tuple @@ -354,6 +355,52 @@ def _load_cuda_library(lib_name: str): raise RuntimeError(f"{lib_name} shared object not found.") +@functools.lru_cache(maxsize=None) +def _load_tvm_ffi_library() -> bool: + """ + Attempt to load the tvm-ffi shared library (libtvm_ffi.so) for the optional CuTeDSL backend. + Returns True if the library is successfully loaded. + Otherwise, return False. In this case C++ will fail to load libtvmffi.so via dlopen, + and skip dispatching to CuTeDSL kernels, falling back to the default TE CUDA C++ kernels. + """ + try: + import tvm_ffi.libinfo as li + + # No binding flag given, so dlopen defaults to RTLD_LAZY, and function symbols will not be resolved immediately here. + # RTLD_GLOBAL adds this lib's symbols to the global scope and TVMFFICentral in C++ will reuse this loaded library and bind to these symbols. + ctypes.CDLL(li.find_libtvm_ffi(), mode=ctypes.RTLD_GLOBAL) + return True + except Exception as e: # pylint: disable=broad-exception-caught + logging.warning( + "Failed to load tvm-ffi dynamic library for CuTeDSL backend: %s. Will fall back to" + " default TE CUDA C++ kernels", + e, + ) + return False + + +@functools.lru_cache(maxsize=None) +def _register_cutedsl_backends() -> bool: + """ + Attempt to register CuTeDSL backends for on-demand compilation via TVM-FFI. + Returns True if the CuTeDSL backends are successfully registered. + Otherwise, return False. In this case C++ will fail to retrieve CuTeDSL kernels via TVM-FFI, + and skip dispatching to CuTeDSL kernels, falling back to the default TE CUDA C++ kernels. + """ + try: + from transformer_engine.common import CuTeDSL + + CuTeDSL.register_cutedsl_backends() + return True + except Exception as e: # pylint: disable=broad-exception-caught + logging.warning( + "Failed to register CuTeDSL backends: %s. Will fall back to default TE CUDA C++" + " kernels", + e, + ) + return False + + @functools.lru_cache(maxsize=None) def _load_core_library(): """Load shared library with Transformer Engine C extensions""" @@ -379,6 +426,11 @@ def _load_core_library(): _CUDART_LIB_CTYPES = _load_cuda_library_from_python("cudart", strict=True) _CUDNN_ALL_LIB_CTYPES = _load_cuda_library_from_python("cudnn", strict=True) + # Prepare CuTeDSL backend for on-demand compilation via TVM-FFI if user enables it. + if os.environ.get("NVTE_ENABLE_CUTEDSL_QUANT_BACKEND", "0") != "0": + _load_tvm_ffi_library() + _register_cutedsl_backends() + _TE_LIB_CTYPES = _load_core_library() # Needed to find the correct headers for NVRTC kernels. diff --git a/transformer_engine/common/cast/dispatch/quantize.cuh b/transformer_engine/common/cast/dispatch/quantize.cuh index 031122d966..539f2c1b83 100644 --- a/transformer_engine/common/cast/dispatch/quantize.cuh +++ b/transformer_engine/common/cast/dispatch/quantize.cuh @@ -13,8 +13,12 @@ #include +#include +#include + #include "../../common.h" #include "../../transpose/cast_transpose.h" +#include "../../util/cuda_runtime.h" #include "../../util/vectorized_pointwise.h" #include "../core/common.cuh" #include "../fp8/group_quantize_fp8.cuh" @@ -22,6 +26,7 @@ #include "../fp8_blockwise/group_quantize_fp8_blockwise.cuh" #include "../mxfp8/group_quantize_mxfp8.cuh" #include "../mxfp8/quantize_mxfp8.cuh" +#include "../mxfp8/quantize_mxfp8_cutedsl.cuh" #include "../nvfp4/group_quantize_transpose_nvfp4.cuh" #include "../nvfp4/quantize_4over6_nvfp4.cuh" #include "../nvfp4/quantize_transpose_nvfp4.cuh" @@ -86,9 +91,16 @@ void quantize_fwd_helper(const NVTETensor input, NVTETensor output, const Tensor *dummy_input_tensor = nullptr; Tensor *dummy_dbias_tensor = nullptr; Tensor *dummy_workspace_tensor = nullptr; - mxfp8::quantize( - *input_tensor, dummy_input_tensor, noop_tensor, output_tensor, dummy_dbias_tensor, - dummy_workspace_tensor, stream); + bool quantized_with_cutedsl = + cutedsl_backend::mxfp8_quantize_cutedsl( + input_tensor, dummy_input_tensor, noop_tensor, output_tensor, dummy_dbias_tensor, + dummy_workspace_tensor, stream); + if (!quantized_with_cutedsl) { + mxfp8::quantize( + *input_tensor, dummy_input_tensor, noop_tensor, output_tensor, dummy_dbias_tensor, + dummy_workspace_tensor, stream); + } break; } case NVTE_NVFP4_1D_SCALING: { @@ -251,9 +263,15 @@ void quantize_bwd_helper(const NVTETensor grad, const NVTETensor input, NVTETens break; } case NVTE_MXFP8_1D_SCALING: { - mxfp8::quantize( - *grad_tensor, input_tensor, noop_tensor, output_tensor, dbias_tensor, workspace_tensor, - stream); + bool quantized_with_cutedsl = + cutedsl_backend::mxfp8_quantize_cutedsl( + grad_tensor, input_tensor, noop_tensor, output_tensor, dbias_tensor, workspace_tensor, + stream); + if (!quantized_with_cutedsl) { + mxfp8::quantize( + *grad_tensor, input_tensor, noop_tensor, output_tensor, dbias_tensor, workspace_tensor, + stream); + } break; } case NVTE_NVFP4_1D_SCALING: { diff --git a/transformer_engine/common/cast/mxfp8/quantize_mxfp8_cutedsl.cuh b/transformer_engine/common/cast/mxfp8/quantize_mxfp8_cutedsl.cuh new file mode 100644 index 0000000000..3b2a79864d --- /dev/null +++ b/transformer_engine/common/cast/mxfp8/quantize_mxfp8_cutedsl.cuh @@ -0,0 +1,277 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#ifndef TRANSFORMER_ENGINE_COMMON_CAST_MXFP8_QUANTIZE_MXFP8_CUTEDSL_CUH_ +#define TRANSFORMER_ENGINE_COMMON_CAST_MXFP8_QUANTIZE_MXFP8_CUTEDSL_CUH_ + +#include +#include + +#include +#include +#include + +#include "../../common.h" +#include "../../tvm_ffi_bridge.h" +#include "../../util/math.h" +#include "../core/common.cuh" // dispatch::common::reduce_dbias +#include "quantize_mxfp8.cuh" // dispatch::mxfp8::quantize_kernel::zero_scales_kernel + +namespace transformer_engine { +namespace cutedsl_backend { + +// Activation, te_dtype_to_str, activation_to_str, DLTensorWrapper, TVMFFICentral +// all live in transformer_engine::tvm_ffi_bridge (tvm_ffi_bridge.h). +using namespace tvm_ffi_bridge; + +struct MXFP8QuantConfig { + static constexpr const char *kEntrypointName = "get_mxfp8_quantization_function"; + + DType dtype; // The input format + DType fp8_dtype; // The fp8 output format + bool rowwise; // If quantize rowwisely + bool colwise; // If quantize columnwisely + bool swizzled; // If the scale output is used for cudnn's swizzled layout + bool with_amax; // If the kernel should return the amax + bool with_dbias = false; // If the dbias is computated (via the workspace tensor) + bool with_dact = false; // If an activation derivative operation is fused + bool with_act = false; // If an activation operation is fused + bool with_noop = false; // If a non-nullptr noop tensor is passed to the kernel + Activation activation = Activation::kNone; + + std::string to_key() const { + std::string key; + key.reserve(56); + key.append("cutedsl_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(swizzled ? "1" : "0") + .append("_") + .append(with_amax ? "1" : "0") + .append("_") + .append(with_dbias ? "1" : "0") + .append("_") + .append(with_dact ? "1" : "0") + .append("_") + .append(with_act ? "1" : "0") + .append("_") + .append(with_noop ? "1" : "0") + .append("_") + .append(activation_to_str(activation)); + return key; + } + + bool retrieve_func_from_python(const std::string &fn_name) const { + auto entrypoint = tvm::ffi::Function::GetGlobal(kEntrypointName); + if (!entrypoint.has_value()) { + return false; + } + tvm::ffi::Any result = + (*entrypoint)(tvm::ffi::String(fn_name), tvm::ffi::String(te_dtype_to_str(dtype)), + tvm::ffi::String(te_dtype_to_str(fp8_dtype)), rowwise, colwise, swizzled, + with_amax, with_dbias, with_dact, with_act, with_noop, + tvm::ffi::String(activation_to_str(activation))); + return result.try_cast().value_or(false); + } +}; + +template +struct MXFP8QuantFused { + static constexpr Activation activation = Activation::kNone; + // No fused activation / activation derivative op: plain quantize + static constexpr bool supported = (OP == nullptr) && !IS_DACT && !IS_ACT; +}; +template <> +struct MXFP8QuantFused> { + static constexpr Activation activation = Activation::kReLU; + static constexpr bool supported = true; +}; +template <> +struct MXFP8QuantFused> { + static constexpr Activation activation = Activation::kGeLU; + static constexpr bool supported = true; +}; +template <> +struct MXFP8QuantFused> { + static constexpr Activation activation = Activation::kSiLU; + static constexpr bool supported = true; +}; +template <> +struct MXFP8QuantFused> { + static constexpr Activation activation = Activation::kQGeLU; + static constexpr bool supported = true; +}; +template <> +struct MXFP8QuantFused> { + static constexpr Activation activation = Activation::kSReLU; + static constexpr bool supported = true; +}; +template +struct MXFP8QuantFused> { + static constexpr Activation activation = Activation::kDReLU; + static constexpr bool supported = true; +}; +template +struct MXFP8QuantFused> { + static constexpr Activation activation = Activation::kDGeLU; + static constexpr bool supported = true; +}; +template +struct MXFP8QuantFused> { + static constexpr Activation activation = Activation::kDSiLU; + static constexpr bool supported = true; +}; +template +struct MXFP8QuantFused> { + static constexpr Activation activation = Activation::kDQGeLU; + static constexpr bool supported = true; +}; +template +struct MXFP8QuantFused> { + static constexpr Activation activation = Activation::kDSReLU; + static constexpr bool supported = true; +}; + +// Signature mirrors mxfp8::quantize (input, act_input, noop, output, dbias, +// workspace, stream). Returns false to fall back to the CUDA kernel. +inline bool mxfp8_quantize_cutedsl(const MXFP8QuantConfig &config, const Tensor *input_tensor, + const Tensor *act_input_tensor, const Tensor *noop_tensor, + Tensor *output_tensor, Tensor *dbias_tensor, + Tensor *workspace_tensor, cudaStream_t stream) { + constexpr size_t kCuTeDSLMXFP8ShapeAlignment = 32; + const size_t flat_m = input_tensor->flat_first_dim(); + const size_t flat_n = input_tensor->flat_last_dim(); + if (flat_m % kCuTeDSLMXFP8ShapeAlignment != 0 || flat_n % kCuTeDSLMXFP8ShapeAlignment != 0) { + return false; + } + + std::optional mxfp8_quant_func_opt = + tvm_ffi_bridge::TVMFFICentral::getInstance().lazyload_function(config); + if (!mxfp8_quant_func_opt.has_value()) { + return false; + } + + // When only WITH_DBIAS is true, we use a larger tile size (align with CUDA C++ implementation) + const bool cast_dbias_only = config.with_dbias && !config.with_dact && !config.with_act; + const size_t chunk_rows = cast_dbias_only ? 128 : 64; // input rows reduced per CTA + // Each CTA writes one partial-dbias row, so the workspace (and the cross-CTA + // reduction below) has ceil(M / chunk_rows) rows. + const size_t workspace_rows = (flat_m + chunk_rows - 1) / chunk_rows; + + // dbias workspace-size query, mirroring mxfp8::quantize: the framework first + // calls with an unallocated workspace to learn its shape, allocates a buffer of + // that shape, then calls again to run. The kernel writes per-row-block partial + // dbias into this workspace; reducing it to the final dbias is a separate step. + if (config.with_dbias && workspace_tensor != nullptr && workspace_tensor->data.dptr == nullptr) { + workspace_tensor->data.shape = {workspace_rows, flat_n}; + workspace_tensor->data.dtype = DType::kFloat32; + return true; + } + + // Zero out swizzled scales if padding is needed. + // Mirrors the CUDA C++ implementation in quantize_mxfp8.cuh. Skip when noop flag is set. + if (config.swizzled && (flat_m % 128 != 0 || flat_n % 128 != 0)) { + const float *noop_ptr = (noop_tensor != nullptr && noop_tensor->data.dptr != nullptr) + ? reinterpret_cast(noop_tensor->data.dptr) + : nullptr; + constexpr size_t zero_threads = 256; + if (output_tensor->has_data()) { + const size_t size_bytes = output_tensor->scale_inv.buffer_size_bytes(); + if (size_bytes > 0) { + const size_t zero_blocks = (size_bytes + zero_threads - 1) / zero_threads; + dispatch::mxfp8::quantize_kernel:: + zero_scales_kernel<<>>( + reinterpret_cast(output_tensor->scale_inv.dptr), size_bytes, noop_ptr); + NVTE_CHECK_CUDA(cudaGetLastError()); + } + } + if (output_tensor->has_columnwise_data()) { + const size_t size_bytes = output_tensor->columnwise_scale_inv.buffer_size_bytes(); + if (size_bytes > 0) { + const size_t zero_blocks = (size_bytes + zero_threads - 1) / zero_threads; + dispatch::mxfp8::quantize_kernel:: + zero_scales_kernel<<>>( + reinterpret_cast(output_tensor->columnwise_scale_inv.dptr), size_bytes, + noop_ptr); + NVTE_CHECK_CUDA(cudaGetLastError()); + } + } + } + + // Data tensors auto-flatten to 2D (DLTensorWrapper's default), matching the + // kernel's flat (rows, cols) view; scale/amax/noop are rank <= 2 and pass through. + tvm_ffi_bridge::DLTensorWrapper mX(input_tensor->data); + tvm_ffi_bridge::DLTensorWrapper mO_row, mS_row, mO_col, mS_col, mAmax, mNoop, mActInput, + mWorkspace; + if (output_tensor->has_data()) { + mO_row = tvm_ffi_bridge::DLTensorWrapper(output_tensor->data); + mS_row = tvm_ffi_bridge::DLTensorWrapper(output_tensor->scale_inv); + } + if (output_tensor->has_columnwise_data()) { + mO_col = tvm_ffi_bridge::DLTensorWrapper(output_tensor->columnwise_data); + mS_col = tvm_ffi_bridge::DLTensorWrapper(output_tensor->columnwise_scale_inv); + } + if (output_tensor->amax.dptr != nullptr) + mAmax = tvm_ffi_bridge::DLTensorWrapper(output_tensor->amax); + if (noop_tensor != nullptr && noop_tensor->data.dptr != nullptr) + mNoop = tvm_ffi_bridge::DLTensorWrapper(noop_tensor->data); + if (act_input_tensor != nullptr && act_input_tensor->data.dptr != nullptr) + mActInput = tvm_ffi_bridge::DLTensorWrapper(act_input_tensor->data); + if (workspace_tensor != nullptr && workspace_tensor->data.dptr != nullptr) + mWorkspace = tvm_ffi_bridge::DLTensorWrapper(workspace_tensor->data); + // stream is a tvm-ffi opaque "handle"; pass the CUDA stream as void*. + (*mxfp8_quant_func_opt)(&mX, &mO_row, &mS_row, &mO_col, &mS_col, &mAmax, &mNoop, &mActInput, + &mWorkspace, static_cast(stream)); + + // If WITH_DBIAS, reduce the workspace partial dbias in CUDA C++ for now. + // This will not be affected by the noop flag, which is aligned with the CUDA C++ implementation + 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::reduce_dbias(workspace_ptr, dbias_tensor, workspace_rows, flat_n, + stream);) // NOLINT(*) + } + return true; +} + +template +bool mxfp8_quantize_cutedsl(const Tensor *input_tensor, const Tensor *act_input_tensor, + const Tensor *noop_tensor, Tensor *output_tensor, Tensor *dbias_tensor, + Tensor *workspace_tensor, cudaStream_t stream) { + using Fused = MXFP8QuantFused; + if constexpr (!Fused::supported) { + return false; + } else { + const bool with_noop = noop_tensor != nullptr && noop_tensor->data.dptr != nullptr; + const MXFP8QuantConfig config{/*dtype=*/input_tensor->dtype(), + /*fp8_dtype=*/output_tensor->dtype(), + /*rowwise=*/output_tensor->has_data(), + /*colwise=*/output_tensor->has_columnwise_data(), + /*swizzled=*/output_tensor->with_gemm_swizzled_scales, + /*with_amax=*/output_tensor->amax.dptr != nullptr, + /*with_dbias=*/IS_DBIAS, + /*with_dact=*/IS_DACT, + /*with_act=*/IS_ACT, + /*with_noop=*/with_noop, + /*activation=*/Fused::activation}; + return mxfp8_quantize_cutedsl(config, input_tensor, act_input_tensor, noop_tensor, + output_tensor, dbias_tensor, workspace_tensor, stream); + } +} + +} // namespace cutedsl_backend +} // namespace transformer_engine + +#endif // TRANSFORMER_ENGINE_COMMON_CAST_MXFP8_QUANTIZE_MXFP8_CUTEDSL_CUH_ diff --git a/transformer_engine/common/transformer_engine.cpp b/transformer_engine/common/transformer_engine.cpp index b3179d38fd..74279b99e8 100644 --- a/transformer_engine/common/transformer_engine.cpp +++ b/transformer_engine/common/transformer_engine.cpp @@ -21,6 +21,7 @@ #include "common.h" #include "common/util/cuda_runtime.h" #include "common/util/logging.h" +#include "tvm_ffi_bridge.h" namespace transformer_engine { @@ -1361,3 +1362,11 @@ NVTEShape nvte_get_grouped_tensor_logical_shape(const NVTEGroupedTensor tensor) const auto &t = *transformer_engine::convertNVTEGroupedTensorCheck(tensor); return t.logical_shape; } + +extern "C" __attribute__((visibility("default"))) void nvte_set_cutedsl_quant_backend(int enabled) { + // Runtime toggle of the CuTeDSL quantize backend, overriding the + // NVTE_ENABLE_CUTEDSL_QUANT_BACKEND env default. + // Used for tests to compare the result of CuTeDSL and the original CUDA implementation. + transformer_engine::tvm_ffi_bridge::TVMFFICentral::getInstance().set_cutedsl_backend_enabled( + enabled != 0); +} diff --git a/transformer_engine/common/tvm_ffi_bridge.h b/transformer_engine/common/tvm_ffi_bridge.h new file mode 100644 index 0000000000..b31db3c0fa --- /dev/null +++ b/transformer_engine/common/tvm_ffi_bridge.h @@ -0,0 +1,332 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#ifndef TRANSFORMER_ENGINE_COMMON_TVM_FFI_BRIDGE_H_ +#define TRANSFORMER_ENGINE_COMMON_TVM_FFI_BRIDGE_H_ + +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "transformer_engine/transformer_engine.h" +#include "util/cuda_runtime.h" +#include "util/logging.h" + +namespace transformer_engine { +namespace tvm_ffi_bridge { + +inline const char *te_dtype_to_str(DType dtype) { + switch (dtype) { + case DType::kFloat32: + return "fp32"; + case DType::kFloat16: + return "fp16"; + case DType::kBFloat16: + return "bf16"; + case DType::kFloat8E4M3: + return "e4m3"; + case DType::kFloat8E5M2: + return "e5m2"; + default: + return ""; + } +} + +// Fused activation token forwarded to Python. Encodes both the family and the +// forward-vs-derivative direction: "relu" is the forward activation, "drelu" its +// backward derivative (dact). This is why no separate is_act/is_dact flag is +// needed — the token carries it; only with_dbias (orthogonal) is a separate flag. +// The d-variants are slots for the not-yet-wired backward path; the forward +// tokens must match Python's SUPPORTED_ACTIVATIONS set. +enum class Activation { + kNone, + kReLU, + kGeLU, + kSiLU, + kQGeLU, + kSReLU, + kDReLU, + kDGeLU, + kDSiLU, + kDQGeLU, + kDSReLU +}; + +inline const char *activation_to_str(Activation act) { + switch (act) { + case Activation::kReLU: + return "relu"; + case Activation::kGeLU: + return "gelu"; + case Activation::kSiLU: + return "silu"; + case Activation::kQGeLU: + return "qgelu"; + case Activation::kSReLU: + return "srelu"; + case Activation::kDReLU: + return "drelu"; + case Activation::kDGeLU: + return "dgelu"; + case Activation::kDSiLU: + return "dsilu"; + case Activation::kDQGeLU: + return "dqgelu"; + case Activation::kDSReLU: + return "dsrelu"; + case Activation::kNone: + return "none"; + } + return "none"; +} + +inline DLDataType convert_to_dltype(NVTEDType type) { + switch (type) { + case kNVTEFloat32: + return DLDataType{kDLFloat, 32, 1}; + case kNVTEFloat16: + return DLDataType{kDLFloat, 16, 1}; + case kNVTEBFloat16: + return DLDataType{kDLBfloat, 16, 1}; + case kNVTEByte: + return DLDataType{kDLUInt, 8, 1}; + case kNVTEInt32: + return DLDataType{kDLInt, 32, 1}; + case kNVTEInt64: + return DLDataType{kDLInt, 64, 1}; + // Native DLPack (>= 1.1) FP8 codes. TE's E4M3 is CUDA's finite __nv_fp8_e4m3, + // i.e. the "fn" variant (kDLFloat8_e4m3 would be the IEEE-style type with + // infinities); E8M0 scales are the unsigned, finite, single-NaN MX format. + case kNVTEFloat8E4M3: + return DLDataType{kDLFloat8_e4m3fn, 8, 1}; + case kNVTEFloat8E5M2: + return DLDataType{kDLFloat8_e5m2, 8, 1}; + case kNVTEFloat8E8M0: + return DLDataType{kDLFloat8_e8m0fnu, 8, 1}; + default: + NVTE_ERROR("unsupported NVTEDType: ", static_cast(type)); + } +} + +class DLTensorWrapper : public DLTensor { + public: + // Null wrapper (data == nullptr): packs as TVM-FFI None, no allocation. + DLTensorWrapper() : DLTensor{} {} + + explicit DLTensorWrapper(const NVTEBasicTensor &tensor, bool flatten_2D = true) { + const int32_t device_index = transformer_engine::cuda::current_device(); + const int n = static_cast(tensor.shape.ndim); + if (flatten_2D && n > 2) { + int64_t flat_first = 1; + for (int i = 0; i + 1 < n; ++i) flat_first *= static_cast(tensor.shape.data[i]); + const int64_t flat_last = static_cast(tensor.shape.data[n - 1]); + shape_buf_ = std::make_unique(2); + strides_buf_ = std::make_unique(2); + shape_buf_[0] = flat_first; + shape_buf_[1] = flat_last; + strides_buf_[0] = flat_last; + strides_buf_[1] = 1; + this->ndim = 2; + } else { + shape_buf_ = std::make_unique(n); + strides_buf_ = std::make_unique(n); + int64_t stride = 1; + for (int i = n - 1; i >= 0; --i) { + shape_buf_[i] = static_cast(tensor.shape.data[i]); + strides_buf_[i] = stride; + stride *= shape_buf_[i]; + } + this->ndim = n; + } + this->data = tensor.data_ptr; + this->device = DLDevice{kDLCUDA, device_index}; + this->dtype = convert_to_dltype(tensor.dtype); + this->shape = shape_buf_.get(); + this->strides = strides_buf_.get(); + this->byte_offset = 0; + } + + ~DLTensorWrapper() = default; + DLTensorWrapper(const DLTensorWrapper &) = delete; + DLTensorWrapper &operator=(const DLTensorWrapper &) = delete; + DLTensorWrapper(DLTensorWrapper &&) = default; + DLTensorWrapper &operator=(DLTensorWrapper &&) = default; + + private: + std::unique_ptr shape_buf_; + std::unique_ptr strides_buf_; +}; + +} // namespace tvm_ffi_bridge +} // namespace transformer_engine + +namespace tvm { +namespace ffi { +// Make a (borrowed) DLTensorWrapper* a first-class TVM-FFI argument, so wrappers +// can be passed straight to Function::operator()(&w, ...). Like DLTensor* it is a +// non-owning DLTensorPtr view (the wrapper must outlive the call), but a null +// pointer OR a wrapper over an absent buffer (null data) packs as TVM-FFI None — +// so a kernel's optional args need no special handling at the call site. Only +// the pack-as-argument path (CopyToAnyView) is provided; reading back is unused. +// Declared after DLTensorWrapper: the specialization needs the complete type +// (it reads src->data and static_casts to its DLTensor base). +template <> +struct TypeTraits + : public TypeTraits { + TVM_FFI_INLINE static void CopyToAnyView(transformer_engine::tvm_ffi_bridge::DLTensorWrapper *src, + TVMFFIAny *result) { + if (src == nullptr || src->data == nullptr) { + TypeTraits::CopyToAnyView(nullptr, result); // -> TVM-FFI None + } else { + TypeTraits::CopyToAnyView(static_cast(src), result); + } + } +}; +} // namespace ffi +} // namespace tvm + +namespace transformer_engine { +namespace tvm_ffi_bridge { + +// Compile-time check that a config provides the lazy-loadable kernel API: +// - std::string to_key() const +// - bool retrieve_func_from_python(const std::string& key) const +// (compiles + globally registers the kernel under `key`; returns whether +// a kernel is now registered / the config is supported) +// Drives the static_assert in TVMFFICentral::lazyload_function so a config that +// is missing either method fails with a clear message instead of a deref-into- +// the-template error. +namespace detail { +template +struct is_lazyloadable_config : std::false_type {}; +template +struct is_lazyloadable_config< + T, std::void_t().to_key()), + decltype(std::declval().retrieve_func_from_python( + std::declval()))>> : std::true_type {}; +} // namespace detail + +class TVMFFICentral { + public: + static TVMFFICentral &getInstance() { + // Deliberately leaked (never deleted) because cache_ holds Python-backed + // tvm::ffi::Function handles whose decref must NOT run at static-destruction + // time -- Python / the tvm-ffi registry may already be finalized by then, + // which would be a use-after-free crash at process exit. + static TVMFFICentral *instance = new TVMFFICentral(); + return *instance; + } + + template + std::optional lazyload_function(const Config &cfg) { + static_assert(detail::is_lazyloadable_config::value, + "Config must define `std::string to_key() const` and " + "`bool retrieve_func_from_python(const std::string&) const`."); + // libtvm_ffi.so absent -> backend disabled, no tvm::ffi symbol is touched. + if (!tvm_ffi_available_) { + NVTE_WARN( + "Cannot dispatch to CuTeDSL kernels because libtvm_ffi.so is not successfully loaded." + " Will fall back to the default CUDA C++ kernels."); + return std::nullopt; + } + if (!cutedsl_backend_enabled_.load(std::memory_order_relaxed)) { + if (warn_cutedsl_backend_not_chosen_) { + NVTE_WARN("CuTeDSL kernel for config `", cfg.to_key(), + "` is not supported because the CuTeDSL backend is disabled. " + "Set NVTE_ENABLE_CUTEDSL_QUANT_BACKEND=1 to enable it."); + } + return std::nullopt; + } + const std::string key = cfg.to_key(); + { + std::shared_lock read_lock(mutex_); + auto it = cache_.find(key); + if (it != cache_.end()) { + // If the key is present, the value is either a valid tvm::ffi::Function or std::nullopt (indicating config not supported) + return it->second; + } + } + // First time we see this config since the key isn't present in the cache: ask Python to compile + register the kernel + // under `key`, then resolve it once and cache the Function (or nullopt if unsupported) + std::optional fn = + cfg.retrieve_func_from_python(key) ? tvm::ffi::Function::GetGlobal(key) : std::nullopt; + { + std::unique_lock write_lock(mutex_); + // emplace is a no-op if another thread populated this key meanwhile; the + // resolved value is identical, so either copy is fine. + cache_.emplace(key, fn); + } + if (!fn && warn_cutedsl_backend_not_chosen_) { + NVTE_WARN("TVM-FFI kernel for config `", key, "` is not supported."); + } + return fn; + } + + // Runtime override of NVTE_ENABLE_CUTEDSL_QUANT_BACKEND (exposed to Python as + // nvte_set_cutedsl_quant_backend; used by tests to compare both backends in + // one process). Safe to toggle at any time + void set_cutedsl_backend_enabled(bool enabled) { + cutedsl_backend_enabled_.store(enabled, std::memory_order_relaxed); + } + + private: + ~TVMFFICentral() = default; + TVMFFICentral() + : tvm_ffi_available_(load_tvm_ffi()), + cutedsl_backend_enabled_(is_cutedsl_backend_enabled()), + warn_cutedsl_backend_not_chosen_(warn_if_cutedsl_backend_not_chosen()) {} + + // Load all tvm-ffi symbols into the global namespace, which should be already loaded in common/__init__.py via ctypes.CDLL + // if user uses TE from a python environment. Otherwise, if user stays in C++ only without python, then CuTeDSL kernels + // will be unavailable either because we fail to load libtvm_ffi.so or CuTeDSL kernel entrypoints are not registered in Python. + // In either case, we will fall back to the default TE CUDA C++ kernels. + static bool load_tvm_ffi() { return dlopen("libtvm_ffi.so", RTLD_NOW | RTLD_GLOBAL) != nullptr; } + TVMFFICentral(const TVMFFICentral &) = delete; + TVMFFICentral &operator=(const TVMFFICentral &) = delete; + TVMFFICentral(TVMFFICentral &&) = delete; + TVMFFICentral &operator=(TVMFFICentral &&) = delete; + + static bool is_cutedsl_backend_enabled() { + // Off by default; set NVTE_ENABLE_CUTEDSL_QUANT_BACKEND=1 to enable. + const char *flag = std::getenv("NVTE_ENABLE_CUTEDSL_QUANT_BACKEND"); + return flag != nullptr && flag[0] != '0'; + } + + static bool warn_if_cutedsl_backend_not_chosen() { + const char *flag = std::getenv("NVTE_WARN_IF_CUTEDSL_BACKEND_NOT_CHOSEN"); + return flag != nullptr && flag[0] != '0'; + } + + const bool tvm_ffi_available_; // libtvm_ffi.so loaded; false disables the backend + std::atomic cutedsl_backend_enabled_; + const bool warn_cutedsl_backend_not_chosen_; + std::shared_mutex mutex_; + // Per-config resolved kernel: cfg.to_key() -> GetGlobal result (std::nullopt == + // unsupported). Holds Python-backed tvm::ffi::Function handles; safe ONLY because + // the singleton is deliberately leaked (see getInstance), so these are never + // decref'd at static teardown. + std::unordered_map> cache_; +}; + +} // namespace tvm_ffi_bridge +} // namespace transformer_engine + +#endif // TRANSFORMER_ENGINE_COMMON_TVM_FFI_BRIDGE_H_ diff --git a/transformer_engine/pytorch/__init__.py b/transformer_engine/pytorch/__init__.py index 06db28ee27..5847f5caa1 100644 --- a/transformer_engine/pytorch/__init__.py +++ b/transformer_engine/pytorch/__init__.py @@ -18,6 +18,7 @@ load_framework_extension("torch") from transformer_engine.pytorch import constants from transformer_engine.pytorch.constants import DType + from transformer_engine.pytorch.module import LayerNormLinear from transformer_engine.pytorch.module import Linear from transformer_engine.pytorch.module import LayerNormMLP diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp index 8d77a9e349..9b3987760a 100644 --- a/transformer_engine/pytorch/csrc/extensions/cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp @@ -51,6 +51,8 @@ void allreduce_nvfp4_amax_tensors(NVFP4Quantizer *nvfp4_quantizer_cpp, py::object quantize(const at::Tensor &tensor, py::handle quantizer, const py::object &output, std::optional noop_flag) { + init_extension(); + // Convert quantizer to C++ object auto quantizer_cpp = convert_quantizer(quantizer);