diff --git a/tests/pytorch/nvfp4/test_nvfp4_cutedsl_backend.py b/tests/pytorch/nvfp4/test_nvfp4_cutedsl_backend.py new file mode 100644 index 00000000000..a6bc5e03c39 --- /dev/null +++ b/tests/pytorch/nvfp4/test_nvfp4_cutedsl_backend.py @@ -0,0 +1,346 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Cross-backend bit-exactness tests for the CuTeDSL NVFP4 quantize kernels, driven from pytorch. + +Each case quantizes the same input with the CUDA backend and with the CuTeDSL backend and requires +identical bytes. The CuTeDSL backend picks one of three kernels per call: the register-resident +(g2r) rowwise or rowwise + columnwise kernel when the input is eligible, else the general +TMA-staged (g2s) kernel. Only configs the CuTeDSL backend serves are tested. + +Stochastic rounding is checked statistically rather than byte for byte, since the g2r kernels draw +Philox per output block and the CUDA kernel per thread. +""" + +# pylint: disable=wrong-import-position + +import ctypes +import os + +import pytest +import torch + +tvm_ffi = pytest.importorskip("tvm_ffi") +pytest.importorskip("cutlass") + +import transformer_engine.pytorch as te + +from transformer_engine.common import _get_shared_object_file +from transformer_engine.common.CuTeDSL.utils import device_compute_capability +from transformer_engine.pytorch import NVFP4Quantizer + +recipe_available, reason_for_no_recipe = te.is_nvfp4_available(return_reason=True) + +CORE_LIB = ctypes.CDLL(str(_get_shared_object_file("core"))) +cutedsl_built = hasattr(CORE_LIB, "nvte_is_cutedsl_backend_built") and bool( + CORE_LIB.nvte_is_cutedsl_backend_built() +) +cutedsl_enabled = os.environ.get("NVTE_ENABLE_CUTEDSL_BACKEND", "0") != "0" +if not cutedsl_built: + skip_reason = "Transformer Engine was built without NVTE_WITH_CUTEDSL=1" +elif not cutedsl_enabled: + skip_reason = "NVTE_ENABLE_CUTEDSL_BACKEND is not set" +else: + skip_reason = reason_for_no_recipe +pytestmark = pytest.mark.skipif( + not (recipe_available and cutedsl_built and cutedsl_enabled), + reason=skip_reason, +) + +NVFP4_BLOCK_SIZE = 16 +SEED = 1234 + +# The CuTeDSL backend takes BF16 input whose flattened dims are multiples of 32. +MATRIX_SIZES = [ + (32, 32), + (32, 1024), + (1024, 32), + (96, 160), + (128, 128), + (256, 1024), + (512, 512), + (8, 32, 1024), # rank > 2, flattened to (256, 1024) + (8192, 7168), +] +# (rowwise, columnwise). Row-scaled NVFP4 has no columnwise output. +DIRECTIONS = [(True, False), (True, True)] + +get_shape_id = lambda s: "x".join(str(d) for d in s) +get_direction_id = lambda d: "rowwise-columnwise" if d[1] else "rowwise" +get_fast_math_id = lambda f: "fastmath" if f else "exact" + + +def set_cutedsl_backend(enabled): + CORE_LIB.nvte_set_cutedsl_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_BACKEND") + set_cutedsl_backend(flag is not None and not flag.startswith("0")) + + +def flat_dims(shape): + return int(torch.tensor(shape[:-1]).prod()), shape[-1] + + +def generate_input(shape, seed=SEED): + """Uniform noise scaled by a random power of two per 16-element block, so neighbouring blocks + get different E4M3 scales.""" + rows, cols = flat_dims(shape) + g = torch.Generator(device="cuda").manual_seed(seed) + x = torch.empty(rows, cols, device="cuda").uniform_(-6.0, 6.0, generator=g) + exponents = torch.randint( + -6, 7, (rows, cols // NVFP4_BLOCK_SIZE, 1), device="cuda", generator=g + ) + x = (x.view(rows, -1, NVFP4_BLOCK_SIZE) * torch.exp2(exponents.float())).view(rows, cols) + return x.to(torch.bfloat16).view(shape) + + +def make_quantizer(rowwise, columnwise, row_scaled=False, stochastic_rounding=False): + return NVFP4Quantizer( + rowwise=rowwise, + columnwise=columnwise, + row_scaled_nvfp4=row_scaled, + stochastic_rounding=stochastic_rounding, + ) + + +def quantize(quantizer, x, noop_flag=None, out=None): + # Stochastic rounding draws its Philox seed from the default CUDA generator. + torch.cuda.manual_seed(SEED) + if out is None: + out = quantizer.make_empty(x.shape, dtype=x.dtype, device="cuda") + return out.quantize_(x, noop_flag=noop_flag) + + +def get_cfg_key(stochastic_rounding, fast_math, row_scaled, columnwise): + """Mirror of NVFP4QuantConfig::to_key (quantize_transpose_nvfp4_cutedsl.cuh).""" + major, minor = device_compute_capability() + flags = (stochastic_rounding, fast_math, row_scaled, columnwise) + return f"cutedsl_nvfp4_sm{major * 10 + minor}_" + "_".join("1" if f else "0" for f in flags) + + +def extract_quantized_output(out, shape, rowwise, columnwise): + """The bytes both backends must agree on. Scale padding is not written, so only the valid + region of the scales is taken.""" + rows, cols = flat_dims(shape) + parts = {} + if rowwise: + parts["rowwise data"] = out._rowwise_data.view(torch.uint8).clone() + parts["rowwise scales"] = out._rowwise_scale_inv.view(torch.uint8)[ + :rows, : cols // NVFP4_BLOCK_SIZE + ].clone() + parts["rowwise amax"] = out._amax_rowwise.clone() + if columnwise: + parts["columnwise data"] = out._columnwise_data.view(torch.uint8).clone() + parts["columnwise scales"] = out._columnwise_scale_inv.view(torch.uint8)[ + :cols, : rows // NVFP4_BLOCK_SIZE + ].clone() + parts["columnwise amax"] = out._amax_columnwise.clone() + return parts + + +def run_test_case(shape, rowwise, columnwise, row_scaled=False, x=None): + """Assert the CuTeDSL and CUDA backends produce bit-identical outputs for the same input.""" + if x is None: + x = generate_input(shape) + fast_math = os.environ.get("NVTE_USE_FAST_MATH", "0") == "1" + quantizer = make_quantizer(rowwise, columnwise, row_scaled) + + set_cutedsl_backend(False) + out_cuda = quantize(quantizer, x) + cuda_output = extract_quantized_output(out_cuda, shape, rowwise, columnwise) + + set_cutedsl_backend(True) + try: + out_cutedsl = quantize(quantizer, x) + cutedsl_output = extract_quantized_output(out_cutedsl, shape, rowwise, columnwise) + finally: + set_cutedsl_backend(False) + + # Guard against a silent CUDA fallback, which would compare CUDA against itself. + key = get_cfg_key(False, fast_math, row_scaled, columnwise) + assert ( + tvm_ffi.get_global_func(key, allow_missing=True) is not None + ), f"CuTeDSL kernel not registered for {key}; the backend fell back to CUDA" + + tag = f"{get_shape_id(shape)}/{get_direction_id((rowwise, columnwise))}" + for name, cuda_bytes in cuda_output.items(): + assert torch.equal(cutedsl_output[name], cuda_bytes), f"{tag}: {name} differ" + + +@pytest.mark.parametrize("shape", MATRIX_SIZES, ids=get_shape_id) +@pytest.mark.parametrize("direction", DIRECTIONS, ids=get_direction_id) +@pytest.mark.parametrize("fast_math", [False, True], ids=get_fast_math_id) +def test_cast(fast_math, direction, shape, monkeypatch): + monkeypatch.setenv("NVTE_USE_FAST_MATH", "1" if fast_math else "0") + run_test_case(shape, *direction) + + +@pytest.mark.parametrize("fast_math", [False, True], ids=get_fast_math_id) +def test_row_scaled(fast_math, monkeypatch): + monkeypatch.setenv("NVTE_USE_FAST_MATH", "1" if fast_math else "0") + run_test_case((256, 1024), True, False, row_scaled=True) + + +@pytest.mark.parametrize("direction", DIRECTIONS, ids=get_direction_id) +@pytest.mark.parametrize("kind", ["zero-blocks", "negative-zero-blocks", "all-zero"]) +def test_zero_blocks(kind, direction): + """Zero block amaxes (and, for all-zero, a zero global amax) hit the clamped reciprocal; + negative zero checks the sign of zero survives the scaling multiply as in CUDA.""" + shape = (256, 1024) + if kind == "all-zero": + x = torch.zeros(shape, dtype=torch.bfloat16, device="cuda") + else: + x = generate_input(shape) + x.view(-1, NVFP4_BLOCK_SIZE)[::3].fill_(-0.0 if kind == "negative-zero-blocks" else 0.0) + run_test_case(shape, *direction, x=x) + + +@pytest.mark.parametrize("direction", DIRECTIONS, ids=get_direction_id) +def test_unaligned_input_uses_g2s(direction): + """A 16-byte-aligned input is not eligible for the g2r kernels' 32-byte loads, so the g2s + kernel serves it.""" + shape = (256, 384) + rows, cols = shape + buffer = torch.empty(rows * cols + 8, dtype=torch.bfloat16, device="cuda") + x = buffer[8:].view(shape) + x.copy_(generate_input(shape)) + assert x.data_ptr() % 32 == 16 + run_test_case(shape, *direction, x=x) + + +def test_rowwise_g2r_grid_y_overflow_uses_g2s(): + """The rowwise g2r kernel covers 4 rows per CTA, so this shape needs grid.y = 65536 and the + g2s kernel must serve it.""" + run_test_case((4 * 65536, 32), True, False) + + +def test_bidimensional_g2r_grid_y_overflow_uses_g2s(): + """The rowwise + columnwise g2r kernel covers 64 rows per CTA, so this shape needs + grid.y = 65536 and the g2s kernel must serve it.""" + run_test_case((64 * 65536, 32), True, True) + + +@pytest.mark.parametrize("flag_is_set", [True, False], ids=["noop-set", "noop-clear"]) +def test_cast_noop_flag(flag_is_set): + """With the noop flag set the kernel must leave the output untouched.""" + shape = (256, 1024) + quantizer = make_quantizer(True, False) + x = generate_input(shape) + flag = torch.full((1,), 1.0 if flag_is_set else 0.0, dtype=torch.float32, device="cuda") + + set_cutedsl_backend(True) + try: + expected = extract_quantized_output(quantize(quantizer, x), shape, True, False) + out = quantize(quantizer, generate_input(shape, seed=SEED + 1)) + before = extract_quantized_output(out, shape, True, False) + quantize(quantizer, x, noop_flag=flag, out=out) + after = extract_quantized_output(out, shape, True, False) + finally: + set_cutedsl_backend(False) + + reference = before if flag_is_set else expected + for name in ("rowwise data", "rowwise scales"): + assert torch.equal(after[name], reference[name]), f"{name} differ" + + +@pytest.mark.parametrize("direction", DIRECTIONS, ids=get_direction_id) +def test_stochastic_rounding(direction): + """Every 16-element block holds one exact 6.0, which makes the encode coefficient exactly 1.0, + and the other elements are 2.75: 3/4 of the way from 2 to 3 on the E2M1 grid. Correct + stochastic rounding sends them to 3 with probability 0.75 and is deterministic for a fixed + seed.""" + rowwise, columnwise = direction + shape = (256, 1024) + x = torch.full(shape, 2.75, dtype=torch.bfloat16, device="cuda") + x[::NVFP4_BLOCK_SIZE, :] = 6.0 + x[:, ::NVFP4_BLOCK_SIZE] = 6.0 + probe = x == 2.75 + quantizer = make_quantizer(rowwise, columnwise, stochastic_rounding=True) + + set_cutedsl_backend(True) + try: + out = extract_quantized_output(quantize(quantizer, x), shape, rowwise, columnwise) + out_again = extract_quantized_output(quantize(quantizer, x), shape, rowwise, columnwise) + finally: + set_cutedsl_backend(False) + key = get_cfg_key(True, False, False, columnwise) + assert tvm_ffi.get_global_func(key, allow_missing=True) is not None, key + + masks = {"rowwise data": probe} + if columnwise: + masks["columnwise data"] = probe.T.contiguous() + for name, mask in masks.items(): + assert torch.equal(out[name], out_again[name]), f"{name}: not deterministic" + data = out[name] + nibbles = torch.stack([data & 0xF, data >> 4], dim=-1).flatten()[mask.flatten()] + # E2M1 codes 4 and 5 are 2.0 and 3.0. + assert bool(((nibbles == 4) | (nibbles == 5)).all()), f"{name}: probe outside {{2, 3}}" + up_fraction = (nibbles == 5).float().mean().item() + assert abs(up_fraction - 0.75) < 0.02, f"{name}: rounded up {up_fraction:.4f}, not 0.75" + + +def _philox4x32(key, counter, rounds): + """Plain Philox4x32 (curanddx's philox4x32_native_state).""" + mask = 0xFFFFFFFF + k0, k1 = key + c = list(counter) + for r in range(rounds): + if r > 0: + k0, k1 = (k0 + 0x9E3779B9) & mask, (k1 + 0xBB67AE85) & mask + p0, p1 = 0xD2511F53 * c[0], 0xCD9E8D57 * c[2] + c = [(p1 >> 32) ^ c[1] ^ k0, p1 & mask, (p0 >> 32) ^ c[3] ^ k1, p0 & mask] + return c + + +def test_block_philox_matches_philox4x32(): + """The g2r kernels' per-block draw is plain Philox4x32 of the counter + {offset_lo, offset_hi, subsequence_lo, subsequence_hi} with key = seed, where + subsequence_hi = logical_hi ^ hi(offset_lo * M0) ^ seed_hi (philox_rng.py).""" + # pylint: disable=import-outside-toplevel + import cutlass + from cutlass import cute, Int32, Uint32, Uint64 + from cutlass.cute.runtime import from_dlpack + from transformer_engine.common.CuTeDSL import philox_rng + + if not philox_rng.BLOCK_PHILOX_FOLD_SUPPORTED: + pytest.skip("the per-block draw needs at least 3 Philox rounds") + + @cute.kernel + def draw_kernel(mIn: cute.Tensor, mOut: cute.Tensor): + t, _, _ = cute.arch.thread_idx() + word = [] + for i in cutlass.range_constexpr(6): + word.append(Uint32(mIn[t, i])) + seed = Uint64(word[0]) | (Uint64(word[1]) << Uint64(32)) + offset = Uint64(word[2]) | (Uint64(word[3]) << Uint64(32)) + fold = philox_rng.make_block_philox_fold( + philox_rng.make_block_philox(seed, offset), word[4] + ) + draw = philox_rng.block_philox_draws([word[5]], fold)[0] + for i in cutlass.range_constexpr(4): + mOut[t, i] = Int32(draw[i]) + + @cute.jit + def draw(mIn: cute.Tensor, mOut: cute.Tensor): + draw_kernel(mIn, mOut).launch(grid=[1, 1, 1], block=[mIn.shape[0], 1, 1]) + + generator = torch.Generator().manual_seed(SEED) + words = torch.randint(0, 2**32, (128, 6), dtype=torch.int64, generator=generator) + words[0] = 0 + words[1] = 2**32 - 1 + draws = torch.zeros((128, 4), dtype=torch.int32, device="cuda") + draw(from_dlpack(words.to(torch.int32).cuda()), from_dlpack(draws)) + + for row, out in zip(words.tolist(), draws.cpu().tolist()): + seed_lo, seed_hi, offset_lo, offset_hi, sub_lo, logical_hi = row + sub_hi = logical_hi ^ ((offset_lo * 0xD2511F53) >> 32) ^ seed_hi + expected = _philox4x32( + (seed_lo, seed_hi), (offset_lo, offset_hi, sub_lo, sub_hi), philox_rng.NUM_PHILOX_ROUNDS + ) + assert [w & 0xFFFFFFFF for w in out] == expected, f"draw for {row}" diff --git a/transformer_engine/common/CuTeDSL/__init__.py b/transformer_engine/common/CuTeDSL/__init__.py index 50d536852c2..09a11882ceb 100644 --- a/transformer_engine/common/CuTeDSL/__init__.py +++ b/transformer_engine/common/CuTeDSL/__init__.py @@ -13,6 +13,9 @@ from transformer_engine.common.CuTeDSL.cast.mxfp8.quantize_mxfp8 import ( get_mxfp8_quantization_function, ) +from transformer_engine.common.CuTeDSL.cast.nvfp4.quantize_nvfp4_common import ( + get_nvfp4_quantization_function, +) def register_cutedsl_backends(): @@ -22,3 +25,6 @@ def register_cutedsl_backends(): tvm_ffi.register_global_func( "get_mxfp8_quantization_function", get_mxfp8_quantization_function, override=True ) + tvm_ffi.register_global_func( + "get_nvfp4_quantization_function", get_nvfp4_quantization_function, override=True + ) diff --git a/transformer_engine/common/CuTeDSL/cast/__init__.py b/transformer_engine/common/CuTeDSL/cast/__init__.py index aefd0b97715..12882ce7208 100644 --- a/transformer_engine/common/CuTeDSL/cast/__init__.py +++ b/transformer_engine/common/CuTeDSL/cast/__init__.py @@ -5,3 +5,4 @@ """CuTeDSL cast/quantize kernels.""" from . import mxfp8 +from . import nvfp4 diff --git a/transformer_engine/common/CuTeDSL/cast/nvfp4/__init__.py b/transformer_engine/common/CuTeDSL/cast/nvfp4/__init__.py new file mode 100644 index 00000000000..56b710bfa35 --- /dev/null +++ b/transformer_engine/common/CuTeDSL/cast/nvfp4/__init__.py @@ -0,0 +1,7 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""CuTeDSL NVFP4 quantization kernels.""" + +from . import quantize_nvfp4_common diff --git a/transformer_engine/common/CuTeDSL/cast/nvfp4/quantize_nvfp4_common.py b/transformer_engine/common/CuTeDSL/cast/nvfp4/quantize_nvfp4_common.py new file mode 100644 index 00000000000..36fc26d6a86 --- /dev/null +++ b/transformer_engine/common/CuTeDSL/cast/nvfp4/quantize_nvfp4_common.py @@ -0,0 +1,347 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Shared definitions, kernel dispatch and TVM-FFI registration for CuTeDSL NVFP4.""" + +import abc +import functools +import logging +import os +from typing import Optional + +import cutlass +from cutlass import cute +from cutlass import BFloat16, Boolean, Float4E2M1FN, Float8E4M3FN, Float32, Int64 +from cuda.bindings.driver import CUstream # pylint: disable=no-name-in-module +import tvm_ffi + +from transformer_engine.common.CuTeDSL.utils import device_compute_capability + +CUTEDSL_DEBUG_LOGGING = os.environ.get("CUTEDSL_DEBUG_LOGGING", "0") == "1" + +# Elements sharing one E4M3 block scale. +NVFP4_BLOCK_SCALING_SIZE = 16 +# The C++ dispatcher only offers flattened dims that are multiples of 32. +NVFP4_SHAPE_ALIGNMENT = 32 +# The scale tensors are allocated padded to (round_up(rows, 128), round_up(cols / 16, 4)), the +# 128x4 tile of the GEMM-swizzled layout (NVFP4Quantizer::get_scale_shape). Declaring that +# divisibility at compile time tells the kernels each scale row starts 4-byte aligned. +NVFP4_SCALE_PAD_ROWS = 128 +NVFP4_SCALE_PAD_COLS = 4 + +logger = logging.getLogger("transformer_engine.cutedsl.nvfp4") + + +@functools.lru_cache(maxsize=None) +def device_sm_count() -> int: + """SM count of the current device, which kernels size their grids against. Read when a + kernel is compiled, so a compiled kernel assumes the SM count of the device it was compiled + on; the C++ side keys kernels by SM architecture.""" + return cutlass.utils.HardwareInfo().get_device_multiprocessor_count() + + +@cute.jit +def noop_flag_is_set(mNoop: cute.Pointer) -> Boolean: + """Whether the cast_noop flag says this quantization is a no-op and must be skipped. + + mNoop is a pointer rather than a tensor so that one compiled kernel serves both a present and + an absent flag, hence the address is checked before it is dereferenced, exactly like the CUDA + C++ kernel's `noop != nullptr && noop[0] == 1.0f`. The two checks cannot be joined with `and`, + which the DSL lowers to a non-short-circuiting op that would load from the null pointer. + """ + flag_is_set = Boolean(False) + if mNoop.toint() != Int64(0): + flag_is_set = cute.make_tensor(mNoop, cute.make_layout((1,)))[0] == Float32(1.0) + return flag_is_set + + +class NVFP4QuantizeConfig: + """Instantiation parameters of the CuTeDSL kernel. These are fixed once the kernel is + compiled and behave as const expressions.""" + + def __init__( + self, + use_stochastic_rounding: bool, + use_fast_math: bool, + row_scaled_nvfp4: bool, + return_transpose: bool, + ): + self.USE_STOCHASTIC_ROUNDING = use_stochastic_rounding + self.USE_FAST_MATH = use_fast_math + if row_scaled_nvfp4 and return_transpose: + raise ValueError("row-scaled NVFP4 quantization does not produce a transposed output") + self.ROW_SCALED_NVFP4 = row_scaled_nvfp4 + self.RETURN_TRANSPOSE = return_transpose + + def __str__(self): + return ( + f"NVFP4QuantizeConfig(use_stochastic_rounding={self.USE_STOCHASTIC_ROUNDING}, " + f"use_fast_math={self.USE_FAST_MATH}, " + f"row_scaled_nvfp4={self.ROW_SCALED_NVFP4}, " + f"return_transpose={self.RETURN_TRANSPOSE})" + ) + + __repr__ = __str__ + + +class NVFP4QuantizeKernelBase(abc.ABC): + """Base class for NVFP4 quantize kernels.""" + + @staticmethod + def supports(cfg: NVFP4QuantizeConfig) -> bool: # pylint: disable=unused-argument + """Compile time check: whether this specialized kernel handles cfg at all.""" + return False + + @staticmethod + def eligible_to_dispatch( + cfg: NVFP4QuantizeConfig, + mX: cute.Tensor, + mO_row: cute.Tensor, + mS_row: cute.Tensor, + mO_col: Optional[cute.Tensor], + mS_col: Optional[cute.Tensor], + ) -> Boolean: + """Runtime check: whether we can dispatch to the instantiated kernel with this call's + shapes and alignments.""" + raise NotImplementedError + + @abc.abstractmethod + def __call__( + self, + mX: cute.Tensor, + mO_row: cute.Tensor, + mS_row: cute.Tensor, + mO_col: Optional[cute.Tensor], + mS_col: Optional[cute.Tensor], + mAmaxRow: cute.Tensor, + mAmaxCol: Optional[cute.Tensor], + mNoop: cute.Pointer, + mRngState: Optional[cute.Tensor], + stream: CUstream, + ): + """ + Compiled kernel entrypoint (decorate with @cute.jit). + All NVFP4 quantize kernels must implement this interface because it is the contract of + the C++ call site in `quantize_transpose_nvfp4_cutedsl.cuh`, which passes the arguments + in this exact order via tvm-ffi; a kernel must accept them even if it does not use them. + """ + + +class NVFP4QuantizeEntry(NVFP4QuantizeKernelBase): + """Select the NVFP4 quantization kernel from the configuration and the runtime shapes.""" + + def __init__(self, cfg: NVFP4QuantizeConfig): + # Kernels import shared definitions from this module. Resolve their + # implementations only when constructing the compile-time entry object. + # pylint: disable=import-outside-toplevel + from .quantize_nvfp4_g2r_bidim import NVFP4QuantizeBidimensionalG2RKernel + from .quantize_nvfp4_g2r_rowwise import NVFP4QuantizeRowwiseG2RKernel + from .quantize_nvfp4_g2s_general import NVFP4QuantizeGeneralG2SKernel + + # pylint: enable=import-outside-toplevel + + self.cfg = cfg + # The general kernel handles every config; the specialized ones, in dispatch priority + # order, take the calls they are eligible for. + self.general_kernel = NVFP4QuantizeGeneralG2SKernel(cfg) + self.specialized_kernels = [ + kernel_cls(cfg) + for kernel_cls in (NVFP4QuantizeRowwiseG2RKernel, NVFP4QuantizeBidimensionalG2RKernel) + if kernel_cls.supports(cfg) + ] + + @cute.jit + def __call__( + self, + mX: cute.Tensor, + mO_row: cute.Tensor, + mS_row: cute.Tensor, + mO_col: Optional[cute.Tensor], + mS_col: Optional[cute.Tensor], + mAmaxRow: cute.Tensor, + mAmaxCol: Optional[cute.Tensor], + mNoop: cute.Pointer, + mRngState: Optional[cute.Tensor], + stream: CUstream, + ): + args = (mX, mO_row, mS_row, mO_col, mS_col, mAmaxRow, mAmaxCol, mNoop, mRngState, stream) + dispatched_to_specialized = False + for kernel in self.specialized_kernels: + if not dispatched_to_specialized: + if kernel.eligible_to_dispatch(self.cfg, mX, mO_row, mS_row, mO_col, mS_col): + dispatched_to_specialized = True + kernel(*args) + if not dispatched_to_specialized: + self.general_kernel(*args) + + +def compile_cutedsl_function_from_cfg(cfg): + """ + Return the compiled CuTeDSL function object for the given NVFP4 quantization config. + + The fake tensors below are what fix the compiled ABI, so they must agree with `__call__`'s + signature. + """ + + kernel_obj = NVFP4QuantizeEntry(cfg) + + def _gmem(dtype, shape, stride_order, align): + """Fake row-major GMEM tensor; `align` is the assumed base alignment in bytes.""" + return cute.runtime.make_fake_compact_tensor( + dtype, + shape, + stride_order=stride_order, + memspace=cute.AddressSpace.gmem, + assumed_align=align, + ) + + # Flattened 2D input dims. Their divisibility is also what keeps every row stride a multiple + # of 16B as TMA requires: 2*N bytes for the bf16 input, N/2 and M/2 for the fp4 outputs. + sym_M = cute.sym_int32(divisibility=NVFP4_SHAPE_ALIGNMENT) + sym_N = cute.sym_int32(divisibility=NVFP4_SHAPE_ALIGNMENT) + + # Scale extents are TE's padded ones (NVFP4Quantizer::get_scale_shape): + # mS_row: (roundup(M, 128), roundup(ceil(N / 16), 4)) + # mS_col: (roundup(N, 128), roundup(ceil(M / 16), 4)) + # Neither is expressible in sym_M / sym_N (SymInt has no `//` or `+`), so the scales get + # fresh syms carrying only the divisibility the padding guarantees. + scale_row_shape = ( + cute.sym_int32(divisibility=NVFP4_SCALE_PAD_ROWS), + cute.sym_int32(divisibility=NVFP4_SCALE_PAD_COLS), + ) + scale_col_shape = ( + cute.sym_int32(divisibility=NVFP4_SCALE_PAD_ROWS), + cute.sym_int32(divisibility=NVFP4_SCALE_PAD_COLS), + ) + + # mX / mO_row / mS_row: (M, N) input, (M, N) rowwise fp4 output and its scales. The fp4 + # extents are logical element counts, so the row stride is N while only N/2 bytes are + # stored. That is the kernel's view of it; the compiled ABI is not, because the DSL rewrites + # every Float4E2M1FN argument into the packed float4_e2m1fnx2 form DLPack can express, with + # the contiguous extent and the strides above it halved. What the caller hands over is + # therefore the packed tensor, which is what tvm_ffi_bridge.h builds for an FP4 TE tensor. + in_fake = _gmem(BFloat16, (sym_M, sym_N), stride_order=(1, 0), align=16) + out_row_fake = _gmem(Float4E2M1FN, (sym_M, sym_N), stride_order=(1, 0), align=16) + scale_row_fake = _gmem(Float8E4M3FN, scale_row_shape, stride_order=(1, 0), align=4) + + # mO_col / mS_col: (N, M) transposed fp4 output and its scales, present iff a transpose is + # produced. + out_col_fake = ( + _gmem(Float4E2M1FN, (sym_N, sym_M), stride_order=(1, 0), align=16) + if cfg.RETURN_TRANSPOSE + else None + ) + scale_col_fake = ( + _gmem(Float8E4M3FN, scale_col_shape, stride_order=(1, 0), align=4) + if cfg.RETURN_TRANSPOSE + else None + ) + + # mAmaxRow: unpadded per-row amax (M,) when row-scaled, else a single global amax (1,). + amax_row_fake = ( + _gmem(Float32, (sym_M,), stride_order=(0,), align=4) + if cfg.ROW_SCALED_NVFP4 + else _gmem(Float32, (1,), stride_order=(0,), align=4) + ) + # mAmaxCol: (1,) global amax of the transposed output, present iff a transpose is produced. + amax_col_fake = ( + _gmem(Float32, (1,), stride_order=(0,), align=4) if cfg.RETURN_TRANSPOSE else None + ) + + # mNoop: always-present f32 pointer to the single-element cast-noop flag. A pointer rather + # than a tensor so the runtime value can be null without changing the compiled ABI -- the + # kernel does the null-check and the noop[0] == 1.0f test on device. The null address here + # is only a compile-time placeholder for the pointer the dispatcher passes at launch. + noop_fake = cute.runtime.nullptr(Float32, mem_space=cute.AddressSpace.gmem, assumed_align=4) + + # mRngState: Philox {seed, offset}; present (and consumed) only for stochastic rounding. + rng_state_fake = ( + _gmem(Int64, (2,), stride_order=(0,), align=8) if cfg.USE_STOCHASTIC_ROUNDING else None + ) + + compiled = cute.compile( + kernel_obj, + in_fake, # mX + out_row_fake, # mO_row + scale_row_fake, # mS_row + out_col_fake, # mO_col + scale_col_fake, # mS_col + amax_row_fake, # mAmaxRow + amax_col_fake, # mAmaxCol + noop_fake, # mNoop + rng_state_fake, # mRngState + cute.runtime.make_fake_stream(), # stream + options="--enable-tvm-ffi", + ) + return compiled + + +def get_nvfp4_quantization_function( + fn_name: str, + use_stochastic_rounding: bool, + use_fast_math: bool, + row_scaled_nvfp4: bool, + return_transpose: bool, +) -> bool: + """Compile the NVFP4 quantize kernel for this config and register it in the TVM-FFI global + registry under EXACTLY `fn_name` (the key the C++ dispatcher built). Returns True if a kernel + is registered under `fn_name`, False if the config is unsupported or compilation failed, in + which case 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 + + major, minor = device_compute_capability() + if major < 10: + logger.warning( + "CuTeDSL NVFP4 backend requires compute capability >= 10.0 (Blackwell), " + "but detected %d.%d; falling back to the CUDA C++ kernel.", + major, + minor, + ) + return False + + try: + cfg = NVFP4QuantizeConfig( + use_stochastic_rounding=use_stochastic_rounding, + use_fast_math=use_fast_math, + row_scaled_nvfp4=row_scaled_nvfp4, + return_transpose=return_transpose, + ) + except ValueError as e: + # The exception message states exactly why the config is unsupported. Surfacing it as a + # warning lets the C++ dispatcher's CUDA fallback be recognized as expected. + logger.warning( + "CuTeDSL NVFP4 backend does not support this config, " + "falling back to the CUDA C++ kernel: %s", + e, + ) + return False + + logger.debug("Compiling CuTeDSL NVFP4 quantization kernel for %s", cfg) + try: + compiled = compile_cutedsl_function_from_cfg(cfg) + except NotImplementedError as e: + # Raised while tracing a config this kernel does not cover yet, which is expected rather + # than a failure, so it is logged apart from a genuine compilation error. + logger.warning( + "CuTeDSL NVFP4 backend does not implement this config yet, " + "falling back to the CUDA C++ kernel: %s", + e, + ) + return False + except Exception as e: # pylint: disable=broad-exception-caught + logger.warning( + "CuTeDSL NVFP4 kernel compilation failed, falling back to the CUDA C++ kernel: %s", + e, + ) + return False + # Register the native function rather than the Python wrapper cute.compile returns: the C++ + # side already passes the arguments positionally, so the wrapper is pure launch overhead. + native = getattr(compiled, "__tvm_ffi_object__", lambda: None)() + tvm_ffi.register_global_func(fn_name, native if native is not None else compiled, override=True) + + return True diff --git a/transformer_engine/common/CuTeDSL/cast/nvfp4/quantize_nvfp4_g2r_bidim.py b/transformer_engine/common/CuTeDSL/cast/nvfp4/quantize_nvfp4_g2r_bidim.py new file mode 100644 index 00000000000..c1f5b652d7e --- /dev/null +++ b/transformer_engine/common/CuTeDSL/cast/nvfp4/quantize_nvfp4_g2r_bidim.py @@ -0,0 +1,541 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Register-resident rowwise + columnwise NVFP4 quantization: every row is read twice, once per +direction, with the second read served by L2, and only the transposed output is staged in shared +memory.""" + +from typing import Optional + +import cutlass +from cutlass import cute +from cutlass import Boolean, Float32, Int32, Int64, Uint32, Uint64 +from cuda.bindings.driver import CUstream # pylint: disable=no-name-in-module + +from transformer_engine.common.CuTeDSL.utils import ( + abs_max_tree_bf16x2, + ld_global_evict_last_v8, + ld_global_v2, + pack_f32x2, + pack_u32x2, + prefetch_l2, + rcp_rn_f32, + row_offset, + st_global_u16, + st_global_u32, + st_global_u8, + st_global_v2, + st_global_v4, + st_global_v8, + unpack_i64_to_i32x2, +) +from transformer_engine.common.CuTeDSL.utils_fp8 import ( + cvt_f32_to_fp8e4m3_code, +) +from transformer_engine.common.CuTeDSL.utils_nvfp4 import ( + block_scale_code, + coefficient_of, + compute_global_encode_sf, + pack_column_block, + pack_row_block, +) +from transformer_engine.common.CuTeDSL.philox_rng import ( + BLOCK_PHILOX_FOLD_SUPPORTED, + block_philox_draws, + columnwise_sub_hi, + columnwise_sub_lo, + make_block_philox, + make_block_philox_fold, + rowwise_sub_lo, +) +from .quantize_nvfp4_common import ( + CUTEDSL_DEBUG_LOGGING, + NVFP4_BLOCK_SCALING_SIZE, + device_sm_count, + noop_flag_is_set, + NVFP4QuantizeConfig, + NVFP4QuantizeKernelBase, +) + + +class NVFP4QuantizeBidimensionalG2RKernel(NVFP4QuantizeKernelBase): + """Rowwise + columnwise NVFP4 quantization from registers. + + A CTA owns _TILE_K columns and walks up to _MAX_SLABS slabs of _SLAB_M rows, keeping the + whole transposed FP4 result for its columns in shared memory: + A. rowwise: one block per thread per quarter slab, from a GROUP-CTA wide strip (adjacent + CTAs pool their rowwise work so a warp writes 256 contiguous payload bytes and whole + 32-byte scale sectors); + B. columnwise: four adjacent columns per thread over one 16-row block, re-reading the rows + phase A just brought into L2, into the transposed staging buffer; + C. drain: four lanes per column, each writing one 32-byte sector of its transposed run. + The slab count is chosen per launch by a cost model of streaming work against the per-CTA + drain. + """ + + _THREADS = 128 + _TILE_K = 128 + _SLAB_M = 64 + _MAX_SLABS = 4 + _GROUP_CTAS = 4 + # Cost model for the slab count: waves * (_SLAB_WORK * slabs + _DRAIN_WORK), in memory + # instructions per thread. + _SLAB_WORK = 36 + _DRAIN_WORK = 51 + # Minimum resident CTAs per SM (launch bounds): the occupancy each reaches anyway. + _MIN_CTAS_PER_SM_RNE = 7 + _MIN_CTAS_PER_SM_SR = 5 + _MAX_GRID_Y = 65535 + + _BLOCKS_PER_SLAB = _SLAB_M // NVFP4_BLOCK_SCALING_SIZE + # Transposed staging buffer: one row of _STAGING_WORDS words per column. A row stride of + # 2 (mod 32) words puts the sixteen lanes of an 8-byte shared access on sixteen disjoint bank + # pairs for both the staging store and the drain's gather, and an even stride keeps every + # slot 8-byte aligned. + _STAGING_WORDS = 34 + _STAGING_SCALE_BYTES = _BLOCKS_PER_SLAB * _MAX_SLABS + 4 + assert _STAGING_WORDS >= 2 * _BLOCKS_PER_SLAB * _MAX_SLABS + assert _STAGING_WORDS % 32 == 2 + assert _STAGING_SCALE_BYTES % 4 == 0 + + @staticmethod + def supports(cfg: NVFP4QuantizeConfig) -> bool: + return ( + cfg.RETURN_TRANSPOSE + and not cfg.ROW_SCALED_NVFP4 + and not cfg.USE_FAST_MATH + and (not cfg.USE_STOCHASTIC_ROUNDING or BLOCK_PHILOX_FOLD_SUPPORTED) + ) + + @staticmethod + @cute.jit + def eligible_to_dispatch( + cfg: cutlass.Constexpr[NVFP4QuantizeConfig], + mX: cute.Tensor, + mO_row: cute.Tensor, + mS_row: cute.Tensor, + mO_col: Optional[cute.Tensor], + mS_col: Optional[cute.Tensor], + ) -> Boolean: + M, K = mX.shape + # Row strides in elements stay 32-bit (16 rows of K per pass), and the 32-byte loads, + # 8-byte rowwise stores, 32-byte transposed stores and 4-byte transposed scale stores are + # aligned (rows and columns are multiples of 32, so only the bases matter). + return ( + K <= (2**31 - 1) // 16 + and cute.ceil_div(M, NVFP4QuantizeBidimensionalG2RKernel._SLAB_M) + <= NVFP4QuantizeBidimensionalG2RKernel._MAX_GRID_Y + and mX.iterator.toint() % 32 == 0 + and mO_row.iterator.toint() % 8 == 0 + and mO_col.iterator.toint() % 32 == 0 + and mS_col.iterator.toint() % 4 == 0 + ) + + def __init__(self, cfg: NVFP4QuantizeConfig): + self.cfg = cfg + + @cute.jit + def __call__( + self, + mX: cute.Tensor, + mO_row: cute.Tensor, + mS_row: cute.Tensor, + mO_col: Optional[cute.Tensor], + mS_col: Optional[cute.Tensor], + mAmaxRow: cute.Tensor, + mAmaxCol: Optional[cute.Tensor], + mNoop: cute.Pointer, + mRngState: Optional[cute.Tensor], + stream: CUstream, + ): + if cutlass.const_expr(CUTEDSL_DEBUG_LOGGING): + cute.printf( + "[CuTeDSL] NVFP4QuantizeBidimensionalG2RKernel.__call__() with config:" + f" {self.cfg}\n" + ) + M, K = mX.shape + args = (mX, mO_row, mS_row, mO_col, mS_col, mAmaxRow, mAmaxCol, mNoop, mRngState) + # The strip remap needs every column of a GROUP-CTA strip to exist. + if K % (self._GROUP_CTAS * self._TILE_K) == 0: + self._launch(args, M, K, stream, self._GROUP_CTAS) + else: + self._launch(args, M, K, stream, 1) + + @cute.jit + def _launch(self, args, M: Int32, K: Int32, stream: CUstream, GROUP: cutlass.Constexpr): + min_ctas = ( + self._MIN_CTAS_PER_SM_SR + if self.cfg.USE_STOCHASTIC_ROUNDING + else self._MIN_CTAS_PER_SM_RNE + ) + resident = Int64(min_ctas * device_sm_count()) + grid_x = cute.ceil_div(K, self._TILE_K) + # Taller CTAs make each transposed run longer (the drain costs one request per column + # whatever the run length) but shrink the grid, and a grid that no longer covers the + # device several times wastes its tail. Pick the slab count by + # waves * (slab streaming work + per-CTA drain). + num_slabs = Int32(1) + best = Int64(-1) + for n in cutlass.range_constexpr(3): + slabs = 1 << n + if cutlass.const_expr(slabs <= self._MAX_SLABS): + grid = Int64(grid_x) * cute.ceil_div(Int64(M), slabs * self._SLAB_M) + waves = (grid + resident - 1) // resident + filled = cutlass.min(grid, resident) + cost = waves * (self._SLAB_WORK * slabs + self._DRAIN_WORK) * resident // filled + if best < 0 or cost < best: + best = cost + num_slabs = Int32(slabs) + self.kernel(*args, num_slabs, GROUP).launch( + grid=[grid_x, cute.ceil_div(M, num_slabs * self._SLAB_M), 1], + block=[self._THREADS, 1, 1], + stream=stream, + min_blocks_per_mp=min_ctas, + # Leave the shared-memory carveout to the driver. Derived from min_blocks_per_mp, it + # ignores the driver's per-CTA reservation and caps residency below min_ctas. + preferred_smem_carveout=-1, + ) + + @cute.kernel + def kernel( + self, + mX: cute.Tensor, + mO_row: cute.Tensor, + mS_row: cute.Tensor, + mO_col: cute.Tensor, + mS_col: cute.Tensor, + mAmaxRow: cute.Tensor, + mAmaxCol: cute.Tensor, + mNoop: cute.Pointer, + mRngState: Optional[cute.Tensor], + num_slabs: Int32, + GROUP: cutlass.Constexpr, + ): + """Device entry for the rowwise + columnwise register kernel.""" + STOCHASTIC = self.cfg.USE_STOCHASTIC_ROUNDING + THREADS = self._THREADS + TILE_K = self._TILE_K + SLAB_M = self._SLAB_M + BLOCK = NVFP4_BLOCK_SCALING_SIZE + tidx, _, _ = cute.arch.thread_idx() + bidx, bidy, _ = cute.arch.block_idx() + M, K = mX.shape + blocks_per_row = K // BLOCK + s_row_pitch = cute.size(mS_row, mode=[1]) + s_col_pitch = cute.size(mS_col, mode=[1]) + x_base = mX.iterator.toint() + q_row_base = mO_row.iterator.toint() + s_row_base = mS_row.iterator.toint() + q_col_base = mO_col.iterator.toint() + s_col_base = mS_col.iterator.toint() + + smem = cutlass.memory.SmemAllocator() + # [0, 128) rowwise coefficients by scale byte, [128, 256) columnwise. + table = smem.allocate_tensor(Float32, cute.make_layout(2 * THREADS), byte_alignment=16) + # Transposed result, column-major: staging[column slot][word], as 8-byte pairs. + staging = smem.allocate_tensor( + Int64, cute.make_layout(TILE_K * self._STAGING_WORDS // 2), byte_alignment=16 + ) + staging_scales = smem.allocate_tensor( + Uint32, cute.make_layout(TILE_K * self._STAGING_SCALE_BYTES // 4), byte_alignment=16 + ) + staging_scales_u8 = cute.make_tensor( + cute.recast_ptr(staging_scales.iterator, dtype=cutlass.Uint8), + cute.make_layout(TILE_K * self._STAGING_SCALE_BYTES), + ) + + # Launch-wide scalars are read here but consumed further down: each heads a dependency + # chain off a global load, so reading them first overlaps those loads with the input's. + skip = noop_flag_is_set(mNoop) + amax_row = mAmaxRow[0] + amax_col = mAmaxCol[0] + rng_seed = Uint64(0) + rng_offset = Uint64(0) + if cutlass.const_expr(STOCHASTIC): + rng_seed = Uint64(mRngState[0]) + rng_offset = Uint64(mRngState[1]) + + c0 = bidx * TILE_K + r00 = bidy * (num_slabs * SLAB_M) + + # Phase A lanes: JA block columns per row x RG rows per pass, from the GROUP-wide strip + # this CTA belongs to (GROUP == 1 is the plain per-tile mapping). + JA = 8 * GROUP + RG = THREADS // JA + ROWS_A = SLAB_M // GROUP + qa = bidx % GROUP # position inside the strip + jbase = (bidx - qa) * (TILE_K // BLOCK) + jl = tidx % JA + rl = tidx // JA + j0 = jbase + jl + j = j0 + if j0 >= blocks_per_row: + j = jbase + arow = j0 < blocks_per_row + # A band of at most 32 rows that starts inside the matrix also ends inside it + # (M % 32 == 0), so clamping the band base is all the bounds work phase A's walk needs. + BAND_SAFE = ROWS_A <= 32 + x_pitch = K * 2 + q_pitch = K // 2 + x_col = x_base + Int64(j) * (BLOCK * 2) + q_col_off = q_row_base + Int64(j) * (BLOCK // 2) + s_col_off = s_row_base + Int64(j) + + # Phase B lanes: four adjacent columns per thread (a warp spans 128 columns, one + # 256-byte request per row) and one of the slab's four column blocks. + cp = tidx % 32 + mb = tidx // 32 + c = c0 + 4 * cp + cok = c < K + + # The first real load is inside the slab loop, so it cannot be hoisted above the + # encode-scale chain; a prefetch can, and warms the first slab meanwhile. + pb = r00 + ROWS_A * qa + rl + for i in cutlass.range_constexpr(4): + r = pb + RG * i + if cutlass.const_expr(BAND_SAFE): + if pb >= M: + r = r00 + RG * i + else: + if r >= M: + r = r00 + prefetch_l2(x_col + row_offset(r, x_pitch)) + + if not skip: # CTA-uniform, and no barrier has been reached yet + S_enc_row = compute_global_encode_sf(amax_row) + # The usual caller passes the same global amax in both directions; recognizing that + # from the raw bits skips a second divide chain in the CTA prologue. + S_enc_col = S_enc_row + if amax_col.bitcast(Uint32) != amax_row.bitcast(Uint32): + S_enc_col = compute_global_encode_sf(amax_col) + S_enc_row_div6 = S_enc_row * Float32(1.0 / 6.0) + S_enc_col_div6 = S_enc_col * Float32(1.0 / 6.0) + # Both halves fold Philox rounds 0..2 against a low word that is constant for the + # CTA: the rowwise half keys on its block column, the columnwise half on the base of + # its four columns, with the column's offset in the quad folded into the high word. + fold_row = None + fold_col = None + if cutlass.const_expr(STOCHASTIC): + state = make_block_philox(rng_seed, rng_offset) + fold_row = make_block_philox_fold(state, rowwise_sub_lo(Uint32(j))) + fold_col = make_block_philox_fold(state, columnwise_sub_lo(Uint32(c))) + + # Equal encode scales share one set of divides. + S_dec_row = rcp_rn_f32(S_enc_row) + row_coeff = coefficient_of(Uint32(tidx), S_dec_row) + col_coeff = row_coeff + if S_enc_col.bitcast(Uint32) != S_enc_row.bitcast(Uint32): + col_coeff = coefficient_of(Uint32(tidx), rcp_rn_f32(S_enc_col)) + table[tidx] = row_coeff + table[THREADS + tidx] = col_coeff + cute.arch.sync_threads() + + for sub in cutlass.range(num_slabs, unroll=1): + r0 = r00 + sub * SLAB_M + if r0 < M: + self._slab( + r0, + sub, + M, + x_base, + x_col, + q_col_off, + s_col_off, + x_pitch, + q_pitch, + s_row_pitch, + qa, + rl, + arow, + cp, + mb, + c, + cok, + num_slabs, + table, + staging, + staging_scales_u8, + S_enc_row_div6, + S_enc_col_div6, + fold_row, + fold_col, + GROUP, + BAND_SAFE, + STOCHASTIC, + ) + + cute.arch.sync_threads() + + # Phase C: four lanes cooperate on one column, each writing a whole 32-byte sector, + # so a column's transposed run leaves as one coalesced request. + valid_blocks = cutlass.min(num_slabs * self._BLOCKS_PER_SLAB, (M - r00) // BLOCK) + ll = tidx % 4 + cgrp = tidx // 4 # this lane group's column quad + cbase = c0 + 4 * cgrp + nblk = valid_blocks - 4 * ll # payload blocks this lane owns (always even) + q_col_pitch = M // 2 + wide = q_col_pitch % 32 == 0 # 32-byte alignment of a transposed row + q_ptr = q_col_base + row_offset(cbase, q_col_pitch) + Int64(r00 // 2 + 32 * ll) + s_ptr = s_col_base + row_offset(cbase, s_col_pitch) + Int64(r00 // BLOCK + 4 * ll) + for q in cutlass.range_constexpr(4): + slot = cgrp + 32 * q + if cbase + q < K and nblk > 0: + d = [] + for t in cutlass.range_constexpr(4): + lo, hi = unpack_i64_to_i32x2( + staging[slot * (self._STAGING_WORDS // 2) + 4 * ll + t] + ) + d.append(Uint32(lo)) + d.append(Uint32(hi)) + if nblk >= 4: + if wide: + st_global_v8(q_ptr, d) # one 32-byte sector + else: # rows are only 16-byte aligned when M % 64 != 0 + st_global_v4(q_ptr, d[0:4]) + st_global_v4(q_ptr + Int64(16), d[4:8]) + else: # tail slab: this lane owns exactly two blocks + st_global_v4(q_ptr, d[0:4]) + scale_word = staging_scales[slot * (self._STAGING_SCALE_BYTES // 4) + ll] + if nblk >= 4: + st_global_u32(s_ptr, scale_word) + else: + st_global_u16(s_ptr, scale_word) + q_ptr = q_ptr + Int64(q_col_pitch) + s_ptr = s_ptr + Int64(s_col_pitch) + + @cute.jit + def _slab( + self, + r0: Int32, + sub: Int32, + M: Int32, + x_base: Int64, + x_col: Int64, + q_col_off: Int64, + s_col_off: Int64, + x_pitch: Int32, + q_pitch: Int32, + s_row_pitch: Int32, + qa: Int32, + rl: Int32, + arow: Boolean, + cp: Int32, + mb: Int32, + c: Int32, + cok: Boolean, + num_slabs: Int32, + table: cute.Tensor, + staging: cute.Tensor, + staging_scales_u8: cute.Tensor, + S_enc_row_div6: Float32, + S_enc_col_div6: Float32, + fold_row, + fold_col, + GROUP: cutlass.Constexpr, + BAND_SAFE: cutlass.Constexpr, + STOCHASTIC: cutlass.Constexpr, + ): + """Phases A and B for the slab of rows [r0, r0 + _SLAB_M).""" + THREADS = self._THREADS + SLAB_M = self._SLAB_M + BLOCK = NVFP4_BLOCK_SCALING_SIZE + RG = THREADS // (8 * GROUP) + ROWS_A = SLAB_M // GROUP + + # Phase A: one rowwise block per thread per quarter slab. A full slab needs no per-row + # bounds test. + rbase = r0 + ROWS_A * qa + rl + afull = arow and r0 + SLAB_M <= M + rows = [rbase + RG * i for i in range(4)] + v = [] + for i in cutlass.range_constexpr(4): + r = rows[i] + if cutlass.const_expr(BAND_SAFE): + if rbase >= M: + r = r0 + RG * i + else: + if r >= M: + r = r0 + v.append(ld_global_evict_last_v8(x_col + row_offset(r, x_pitch))) + + # Warm L2 for the next slab under SR only: RNE streams near its DRAM ceiling and a + # lookahead competes with demand loads, while SR has the slack for it to land. + if cutlass.const_expr(STOCHASTIC): + if sub + 1 < num_slabs and r0 + 2 * SLAB_M <= M: + for i in cutlass.range_constexpr(4): + prefetch_l2(x_col + row_offset(rows[i] + SLAB_M, x_pitch)) + + # The Philox chain depends only on the row, never on the data, so it goes before the + # amax trees that stall on the loads. + rnd_row = [None] * 4 + if cutlass.const_expr(STOCHASTIC): + for g in cutlass.range_constexpr(0, 4, 2): + draws = block_philox_draws([Uint32(rows[g]), Uint32(rows[g + 1])], fold_row) + rnd_row[g] = draws[0] + rnd_row[g + 1] = draws[1] + codes_a = [block_scale_code(v[i], S_enc_row_div6) for i in range(4)] + coeffs_a = [table[Int32(codes_a[i])] for i in range(4)] + coeff2_a = [pack_f32x2(c, c) for c in coeffs_a] + for i in cutlass.range_constexpr(4): + w0, w1 = pack_row_block(v[i], coeff2_a[i], rnd_row[i], STOCHASTIC) + if afull or (rows[i] < M and arow): + st_global_v2(q_col_off + row_offset(rows[i], q_pitch), w0, w1) + st_global_u8(s_col_off + row_offset(rows[i], s_row_pitch), codes_a[i]) + + # Phase B: the rows phase A just read are still in L2, so this column-major pass is a + # cache re-read rather than a second DRAM stream. + rb = r0 + mb * BLOCK + bok = cok and rb + BLOCK <= M + cc = c + rr = rb + if not bok: + cc = Int32(0) + rr = r0 + x_ptr = x_base + row_offset(rr, x_pitch) + Int64(cc) * 2 + raw = [[], []] + for _ in cutlass.range_constexpr(BLOCK): + lo, hi = ld_global_v2(x_ptr) + raw[0].append(lo) + raw[1].append(hi) + x_ptr = x_ptr + Int64(x_pitch) + + word_offset = (sub * self._BLOCKS_PER_SLAB + mb) * 2 + byte_offset = sub * self._BLOCKS_PER_SLAB + mb + # High word = (row block, column inside the quad); the quad base is in the folded low + # word, so the four columns differ only in two bits. + rnd = [None] * 4 + if cutlass.const_expr(STOCHASTIC): + row_block = Uint32(rb) >> Uint32(4) + for g in cutlass.range_constexpr(0, 4, 2): + draws = block_philox_draws( + [ + columnwise_sub_hi(row_block, Uint32(g)), + columnwise_sub_hi(row_block, Uint32(g + 1)), + ], + fold_col, + ) + rnd[g] = draws[0] + rnd[g + 1] = draws[1] + # Both 16-row amax trees and all four coefficient lookups resolve before any column + # packs, so the four shared loads are in flight together. + codes_b = [None] * 4 + for p in cutlass.range_constexpr(2): # one packed word holds a pair of adjacent columns + pair_amax = abs_max_tree_bf16x2(raw[p]) & Uint32(0x7FFF7FFF) + codes_b[2 * p] = cvt_f32_to_fp8e4m3_code( + (pair_amax << Uint32(16)).bitcast(Float32) * S_enc_col_div6 + ) + codes_b[2 * p + 1] = cvt_f32_to_fp8e4m3_code( + (pair_amax & Uint32(0xFFFF0000)).bitcast(Float32) * S_enc_col_div6 + ) + coeffs_b = [table[THREADS + Int32(codes_b[k])] for k in range(4)] + coeff2_b = [pack_f32x2(c, c) for c in coeffs_b] + for p in cutlass.range_constexpr(2): + for q in cutlass.range_constexpr(2): + cs = 2 * p + q + w0, w1 = pack_column_block(raw[p], coeff2_b[cs], rnd[cs], q == 0, STOCHASTIC) + slot = cs * 32 + cp # one 8-byte staging store per column block + staging[slot * (self._STAGING_WORDS // 2) + word_offset // 2] = pack_u32x2(w0, w1) + staging_scales_u8[slot * self._STAGING_SCALE_BYTES + byte_offset] = cutlass.Uint8( + codes_b[cs] + ) diff --git a/transformer_engine/common/CuTeDSL/cast/nvfp4/quantize_nvfp4_g2r_rowwise.py b/transformer_engine/common/CuTeDSL/cast/nvfp4/quantize_nvfp4_g2r_rowwise.py new file mode 100644 index 00000000000..50f50b7ef40 --- /dev/null +++ b/transformer_engine/common/CuTeDSL/cast/nvfp4/quantize_nvfp4_g2r_rowwise.py @@ -0,0 +1,269 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Register-resident rowwise-only NVFP4 quantization: each 16-element block streams from one +32-byte load straight to its 8 output bytes, with no shared-memory staging of the data.""" + +from typing import Optional + +import cutlass +from cutlass import cute +from cutlass import Boolean, Float32, Int32, Int64, Uint32, Uint64 +from cuda.bindings.driver import CUstream # pylint: disable=no-name-in-module + +from transformer_engine.common.CuTeDSL.utils import ( + ld_global_nc_evict_first_v8, + pack_f32x2, + prefetch_l2, + rcp_rn_f32, + row_offset, + st_global_u8, + st_global_v2, +) +from transformer_engine.common.CuTeDSL.utils_nvfp4 import ( + block_scale_code, + coefficient_of, + compute_global_encode_sf, + pack_row_block, +) +from transformer_engine.common.CuTeDSL.philox_rng import ( + BLOCK_PHILOX_FOLD_SUPPORTED, + block_philox_draws, + make_block_philox, + make_block_philox_fold, + rowwise_sub_lo, +) +from .quantize_nvfp4_common import ( + CUTEDSL_DEBUG_LOGGING, + NVFP4_BLOCK_SCALING_SIZE, + device_sm_count, + noop_flag_is_set, + NVFP4QuantizeConfig, + NVFP4QuantizeKernelBase, +) + + +class NVFP4QuantizeRowwiseG2RKernel(NVFP4QuantizeKernelBase): + """Rowwise-only NVFP4 quantization from registers. + + A CTA owns _BLOCKS_PER_CTA consecutive blocks of ROWS rows, one block column per thread, and + walks its rows _WAVE at a time: each thread holds _WAVE 32-byte loads in flight while the + next wave is prefetched into L2. The per-block coefficient depends only on the block's scale + byte, so its IEEE divide is evaluated once per CTA into a shared-memory table indexed by that + byte. Stochastic rounding draws Philox once per block (philox_rng.make_block_philox_fold). + """ + + # A CTA covers 128 blocks = 2048 input bytes and 128 scale bytes per row. + _BLOCKS_PER_CTA = 128 + _THREADS = 128 + # Rows per CTA. Round-to-nearest is memory-bound, and 4 rows (4 loads in flight per thread) + # is where more CTAs stop buying latency hiding. Stochastic rounding is bound by the Philox + # chain, so taller CTAs amortize the per-CTA prologue; it doubles _ROWS_SR while the grid + # still fills the GPU. + _ROWS_RNE = 4 + _ROWS_SR = 8 + # Rows held in registers at once (one 32-byte load each). + _WAVE = 4 + # Minimum resident CTAs per SM (launch bounds): the occupancy each reaches anyway. Under SR, + # 64 registers instead of 72 buys a quarter more warps to hide Philox with. + _MIN_CTAS_PER_SM_RNE = 10 + _MIN_CTAS_PER_SM_SR = 8 + _MAX_GRID_Y = 65535 + + @staticmethod + def supports(cfg: NVFP4QuantizeConfig) -> bool: + return ( + not cfg.RETURN_TRANSPOSE + and not cfg.ROW_SCALED_NVFP4 + and not cfg.USE_FAST_MATH + and (not cfg.USE_STOCHASTIC_ROUNDING or BLOCK_PHILOX_FOLD_SUPPORTED) + ) + + @staticmethod + @cute.jit + def eligible_to_dispatch( + cfg: cutlass.Constexpr[NVFP4QuantizeConfig], + mX: cute.Tensor, + mO_row: cute.Tensor, + mS_row: cute.Tensor, + mO_col: Optional[cute.Tensor], + mS_col: Optional[cute.Tensor], + ) -> Boolean: + M = mX.shape[0] + # Rows fit grid.y at the shortest CTA, and the 32-byte loads and 8-byte stores are + # aligned (rows and columns are multiples of 32, so only the bases matter). + return ( + cute.ceil_div(M, NVFP4QuantizeRowwiseG2RKernel._ROWS_RNE) + <= NVFP4QuantizeRowwiseG2RKernel._MAX_GRID_Y + and mX.iterator.toint() % 32 == 0 + and mO_row.iterator.toint() % 8 == 0 + ) + + def __init__(self, cfg: NVFP4QuantizeConfig): + self.cfg = cfg + + @cute.jit + def __call__( + self, + mX: cute.Tensor, + mO_row: cute.Tensor, + mS_row: cute.Tensor, + mO_col: Optional[cute.Tensor], # Unused, kept for API compatibility + mS_col: Optional[cute.Tensor], # Unused, kept for API compatibility + mAmaxRow: cute.Tensor, + mAmaxCol: Optional[cute.Tensor], # Unused, kept for API compatibility + mNoop: cute.Pointer, + mRngState: Optional[cute.Tensor], + stream: CUstream, + ): + if cutlass.const_expr(CUTEDSL_DEBUG_LOGGING): + cute.printf( + f"[CuTeDSL] NVFP4QuantizeRowwiseG2RKernel.__call__() with config: {self.cfg}\n" + ) + M, N = mX.shape + grid_x = cute.ceil_div(N // NVFP4_BLOCK_SCALING_SIZE, self._BLOCKS_PER_CTA) + args = (mX, mO_row, mS_row, mAmaxRow, mNoop, mRngState) + if cutlass.const_expr(self.cfg.USE_STOCHASTIC_ROUNDING): + # Under SR the Philox chain, not DRAM, sets the pace, so the per-CTA prologue (amax + # load, divides, coefficient table, barrier) is a real fraction of each CTA. Doubling + # the stripe halves it, worth taking whenever the taller grid still fills the device. + resident = self._MIN_CTAS_PER_SM_SR * device_sm_count() + if Int64(grid_x) * cute.ceil_div(M, 2 * self._ROWS_SR) >= resident: + self._launch(args, grid_x, M, stream, 2 * self._ROWS_SR) + else: + self._launch(args, grid_x, M, stream, self._ROWS_SR) + else: + self._launch(args, grid_x, M, stream, self._ROWS_RNE) + + @cute.jit + def _launch(self, args, grid_x: Int32, M: Int32, stream: CUstream, ROWS: cutlass.Constexpr): + min_ctas = ( + self._MIN_CTAS_PER_SM_SR + if self.cfg.USE_STOCHASTIC_ROUNDING + else self._MIN_CTAS_PER_SM_RNE + ) + self.kernel(*args, ROWS).launch( + grid=[grid_x, cute.ceil_div(M, ROWS), 1], + block=[self._THREADS, 1, 1], + stream=stream, + min_blocks_per_mp=min_ctas, + # Leave the shared-memory carveout to the driver. Derived from min_blocks_per_mp, it + # ignores the driver's per-CTA reservation and caps residency below min_ctas. + preferred_smem_carveout=-1, + ) + + @cute.kernel + def kernel( + self, + mX: cute.Tensor, + mO_row: cute.Tensor, + mS_row: cute.Tensor, + mAmaxRow: cute.Tensor, + mNoop: cute.Pointer, + mRngState: Optional[cute.Tensor], + ROWS: cutlass.Constexpr, + ): + """Device entry for the rowwise-only register kernel.""" + STOCHASTIC = self.cfg.USE_STOCHASTIC_ROUNDING + WAVE = self._WAVE + tidx, _, _ = cute.arch.thread_idx() + bidx, bidy, _ = cute.arch.block_idx() + K = mX.shape[1] + blocks_per_row = K // NVFP4_BLOCK_SCALING_SIZE + s_pitch = cute.size(mS_row, mode=[1]) + + smem = cutlass.memory.SmemAllocator() + # Rowwise RNE consumes every coefficient as a packed pair, so it stores that form once; + # SR keeps the narrower float. + table = smem.allocate_tensor( + Float32 if STOCHASTIC else Int64, cute.make_layout(self._THREADS), byte_alignment=16 + ) + + # Launch-wide scalars are read here but consumed further down: each heads a dependency + # chain off a global load, so reading them first overlaps those loads with the input's. + skip = noop_flag_is_set(mNoop) + amax_row = mAmaxRow[0] + rng_seed = Uint64(0) + rng_offset = Uint64(0) + if cutlass.const_expr(STOCHASTIC): + rng_seed = Uint64(mRngState[0]) + rng_offset = Uint64(mRngState[1]) + + j0 = bidx * self._BLOCKS_PER_CTA + tidx + j = j0 + if j0 >= blocks_per_row: + j = Int32(0) # clamp so the loads stay in bounds; the stores are predicated + r0 = bidy * ROWS + + # The column term of every address is a thread constant, folded into the bases once. + x_j = mX.iterator.toint() + Int64(j) * (NVFP4_BLOCK_SCALING_SIZE * 2) + q_j = mO_row.iterator.toint() + Int64(j) * (NVFP4_BLOCK_SCALING_SIZE // 2) + s_j = mS_row.iterator.toint() + Int64(j) + x_pitch = K * 2 # bytes per input row + q_pitch = K // 2 # bytes per output row + + v = [ld_global_nc_evict_first_v8(x_j + row_offset(r0 + i, x_pitch)) for i in range(WAVE)] + + if not skip: # CTA-uniform, and no barrier has been reached yet + # Resolve the RNG first: its offset-dependent multiply has the longer wait, and the + # already-issued loads continue toward readiness behind it. The block column is a + # thread constant, so placing it in the subsequence's low word folds rounds 0..2. + fold = None + if cutlass.const_expr(STOCHASTIC): + fold = make_block_philox_fold( + make_block_philox(rng_seed, rng_offset), rowwise_sub_lo(Uint32(j)) + ) + + S_enc = compute_global_encode_sf(amax_row) + S_enc_div6 = S_enc * Float32(1.0 / 6.0) + # One table entry per thread; its barrier is deferred into the first wave, past that + # wave's amax trees, which depend only on the input already in flight. + coeff = coefficient_of(Uint32(tidx), rcp_rn_f32(S_enc)) + if cutlass.const_expr(STOCHASTIC): + table[tidx] = coeff + else: + table[tidx] = pack_f32x2(coeff, coeff) + + active = j0 < blocks_per_row + for w in cutlass.range_constexpr(0, ROWS, WAVE): + if cutlass.const_expr(w > 0): + v = [ + ld_global_nc_evict_first_v8(x_j + row_offset(r0 + w + i, x_pitch)) + for i in range(WAVE) + ] + # Register-free L2 warm-up for the next wave. Those rows are inside this CTA's + # stripe (ROWS divides 32 and M % 32 == 0), so no bounds test is needed. + if cutlass.const_expr(w + WAVE < ROWS): + for i in cutlass.range_constexpr(WAVE): + prefetch_l2(x_j + row_offset(r0 + w + WAVE + i, x_pitch)) + # The wave's first random pair needs none of the loaded data, so draw it before + # the amax trees that stall on the loads. + rnd0 = None + if cutlass.const_expr(STOCHASTIC): + rnd0 = block_philox_draws([Uint32(r0 + w), Uint32(r0 + w + 1)], fold) + # Every block's amax tree and scale byte resolve before any of them packs. + codes = [block_scale_code(v[i], S_enc_div6) for i in range(WAVE)] + if cutlass.const_expr(w == 0): + cute.arch.sync_threads() # the coefficient table is read only from here on + + for g in cutlass.range_constexpr(0, WAVE, 2): + rnd = [None, None] + if cutlass.const_expr(STOCHASTIC): + if cutlass.const_expr(g == 0): + rnd = rnd0 + else: + rnd = block_philox_draws( + [Uint32(r0 + w + g), Uint32(r0 + w + g + 1)], fold + ) + for i in cutlass.range_constexpr(g, g + 2): + if cutlass.const_expr(STOCHASTIC): + coeff = table[Int32(codes[i])] + coeff2 = pack_f32x2(coeff, coeff) + else: + coeff2 = table[Int32(codes[i])] + w0, w1 = pack_row_block(v[i], coeff2, rnd[i - g], STOCHASTIC) + row = r0 + w + i + if active: + st_global_v2(q_j + row_offset(row, q_pitch), w0, w1) + st_global_u8(s_j + row_offset(row, s_pitch), codes[i]) diff --git a/transformer_engine/common/CuTeDSL/cast/nvfp4/quantize_nvfp4_g2s_general.py b/transformer_engine/common/CuTeDSL/cast/nvfp4/quantize_nvfp4_g2s_general.py new file mode 100644 index 00000000000..58fcdfd844a --- /dev/null +++ b/transformer_engine/common/CuTeDSL/cast/nvfp4/quantize_nvfp4_g2s_general.py @@ -0,0 +1,816 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""General NVFP4 quantize-transpose kernel: TMA-staged tiles, the CuTeDSL counterpart of +quantize_transpose_nvfp4_tuned_1D.cuh. + +It handles every config. NVFP4QuantizeEntry in quantize_nvfp4_common.py dispatches the calls a +specialized kernel is eligible for to that kernel and falls back to this one otherwise.""" + +from typing import Optional + +import cutlass +from cutlass import cute +from cutlass import ( + BFloat16, + Float8E4M3FN, + Float32, + Int32, + Int64, + Uint8, + Uint32, + Uint64, +) +from cuda.bindings.driver import CUstream # pylint: disable=no-name-in-module + +from transformer_engine.common.CuTeDSL.utils import ( + abs_max_x2_bf16, + fabs_f32, + make_prmt_u32, + pack_u32x2, + x2_hi_to_f32_bf16, + x2_lo_to_f32_bf16, +) +from transformer_engine.common.CuTeDSL.utils_nvfp4 import ( + compute_block_decode_sf, + compute_block_encode_sf, + compute_global_encode_sf, + cvt_f32x8_to_fp4x8, + cvt_f32x8_to_fp4x8_sr, + mul_cvt_bf16x8_to_fp4x8, + mul_cvt_bf16x8_to_fp4x8_sr, +) +from transformer_engine.common.CuTeDSL.philox_rng import PhiloxRng +from .quantize_nvfp4_common import ( + CUTEDSL_DEBUG_LOGGING, + NVFP4_BLOCK_SCALING_SIZE, + noop_flag_is_set, + NVFP4QuantizeKernelBase, +) + + +@cute.jit +def zero_b32() -> Int32: + """A zero the packed abs-max can start from. As a literal it would be folded into an + immediate operand, which max.xorsign.abs.bf16x2 does not accept.""" + return cute.arch.inline_ptx("mov.b32 {$w0}, 0;", write_only_types=[Int32]) + + +# Interleave low/high bf16 halves of two pair-registers for the colwise pass. +prmt_lo_u32 = make_prmt_u32(0x5410) +prmt_hi_u32 = make_prmt_u32(0x7632) + + +class NVFP4QuantizeGeneralG2SKernel(NVFP4QuantizeKernelBase): + """Tuned kernel to cast to NVFP4 and transpose. + + Each thread block processes a CHUNK_DIM_Y x CHUNK_DIM_X _chunk_ of the input tensor as a grid + of TILE_DIM x TILE_DIM _tiles_, walked sequentially with PREFETCH_STAGES tiles in flight + beyond the current one: TMA loads a tile to SMEM, the CTA quantizes it into staged SMEM + output buffers, and TMA stores those while the next tile loads. + + The thread arrangements are the CUDA kernel's. Rowwise, a thread owns one whole scaling + block (16 row-adjacent elements) per iteration, making the block amax an intra-thread + reduction, with the 16-byte shared memory reads staggered by bank group. Columnwise, a + thread owns two adjacent columns of one 16-row scaling block, reading pairs, with the block + row staggered by warp so consecutive lanes read different SMEM rows. + """ + + CHUNK_DIM_Y = 128 + CHUNK_DIM_X = 128 + TILE_DIM = 64 + THREADS = 128 + PREFETCH_STAGES = 1 + + # Derived tiling constants (names follow the CUDA kernel). + NUM_BUFFERS = PREFETCH_STAGES + 1 + STAGES_Y = CHUNK_DIM_Y // TILE_DIM + STAGES_X = CHUNK_DIM_X // TILE_DIM + STAGES = STAGES_Y * STAGES_X + SCALES_PER_TILE = TILE_DIM // NVFP4_BLOCK_SCALING_SIZE # 4 + SCALES_PER_CHUNK_X = CHUNK_DIM_X // NVFP4_BLOCK_SCALING_SIZE # 8 + SCALES_PER_CHUNK_Y = CHUNK_DIM_Y // NVFP4_BLOCK_SCALING_SIZE # 8 + THREADS_X_ROWWISE = TILE_DIM // NVFP4_BLOCK_SCALING_SIZE # 4 + THREADS_Y_ROWWISE = THREADS // THREADS_X_ROWWISE # 32 + ITERATIONS_ROWWISE = TILE_DIM // THREADS_Y_ROWWISE # 2 + PACK_SIZE = 8 # elements per vectorized SMEM access + WAVES = NVFP4_BLOCK_SCALING_SIZE // PACK_SIZE # 2 + # Threads that span the 32 4-byte SMEM banks at 16 bf16 per thread (the rowwise swizzle). + THREADS_PER_BANK = (32 * 4 * 8) // 4 // NVFP4_BLOCK_SCALING_SIZE # 16 + assert NUM_BUFFERS <= STAGES # otherwise, prefetch loop would read OOB + + @staticmethod + def supports(cfg) -> bool: + return True + + def __init__(self, cfg): + self.cfg = cfg + + # Host-side kernel launch + @cute.jit + def __call__( + self, + mX: cute.Tensor, # (M, N) bf16 input + mO_row: cute.Tensor, # (M, N) fp4 rowwise output + mS_row: cute.Tensor, # (roundup(M, 128), roundup(ceil(N / 16), 4)) e4m3 scales + mO_col: Optional[cute.Tensor], # (N, M) fp4 transposed output + mS_col: Optional[cute.Tensor], # (roundup(N, 128), roundup(ceil(M / 16), 4)) e4m3 scales + mAmaxRow: cute.Tensor, # (M,) f32 per-row amax if ROW_SCALED_NVFP4, else (1,) global amax + mAmaxCol: Optional[cute.Tensor], # (1,) f32 global amax of the transposed output + mNoop: cute.Pointer, # f32 cast-noop flag; may be null, checked on device + mRngState: Optional[cute.Tensor], # (2,) i64 Philox {seed, offset} + stream: CUstream, + ): + """AOT-compiled host entrypoint. `quantize_transpose_nvfp4_cutedsl.cuh` passes these + arguments in this exact order via tvm-ffi, and the config fixes which of the optional + ones are present (see `compile_cutedsl_function_from_cfg`). + + M, N are the input's *flattened* 2D dims, both multiples of NVFP4_SHAPE_ALIGNMENT; a + rank > 2 input already arrives flattened. + All tensors are row-major (with rightmost stride 1). + FP4 extents are logical element counts, not the halved extents of the uint8 buffer TE + actually allocates. + """ + if cutlass.const_expr(CUTEDSL_DEBUG_LOGGING): + cute.printf( + f"[CuTeDSL] NVFP4QuantizeGeneralG2SKernel.__call__() with config: {self.cfg}\n" + ) + + M, N = mX.shape + + ## TMA descriptors + # The FP4 outputs are moved as bytes: two elements to a byte, so the byte view of the + # (M, N) rowwise output is (M, N/2) and a 64x64-element tile is a 64x32-byte box. The + # 32-divisibility of both dims is what keeps every row stride a multiple of 16B as TMA + # requires: 2*N bytes for the bf16 input, N/2 and M/2 for the fp4 outputs. + mO_row_bytes = cute.recast_tensor(mO_row, Uint8) + + tile_in_layout = cute.make_ordered_layout((self.TILE_DIM, self.TILE_DIM), order=(1, 0)) + tile_out_layout = cute.make_ordered_layout( + (self.TILE_DIM, self.TILE_DIM // 2), order=(1, 0) + ) + + op_load = cute.nvgpu.cpasync.CopyBulkTensorTileG2SOp() + op_store = cute.nvgpu.cpasync.CopyBulkTensorTileS2GOp() + + tma_atom_in, tma_view_in = cute.nvgpu.cpasync.make_tiled_tma_atom( + op_load, mX, tile_in_layout, (self.TILE_DIM, self.TILE_DIM), num_multicast=1 + ) + tma_atom_row, tma_view_row = cute.nvgpu.cpasync.make_tiled_tma_atom( + op_store, + mO_row_bytes, + tile_out_layout, + (self.TILE_DIM, self.TILE_DIM // 2), + num_multicast=1, + ) + if cutlass.const_expr(self.cfg.RETURN_TRANSPOSE): + mO_col_bytes = cute.recast_tensor(mO_col, Uint8) + tma_atom_col, tma_view_col = cute.nvgpu.cpasync.make_tiled_tma_atom( + op_store, + mO_col_bytes, + tile_out_layout, + (self.TILE_DIM, self.TILE_DIM // 2), + num_multicast=1, + ) + else: + tma_atom_col = None + tma_view_col = None + + ## Grid: one CTA per chunk; X indexes columns and Y rows, as in the CUDA kernel's + ## ctaid_X / ctaid_Y. + grid = [ + cute.ceil_div(N, self.CHUNK_DIM_X), + cute.ceil_div(M, self.CHUNK_DIM_Y), + 1, + ] + + self.kernel( + mX, + mS_row, + mS_col, + mAmaxRow, + mAmaxCol, + mNoop, + mRngState, + tma_atom_in, + tma_view_in, + tma_atom_row, + tma_view_row, + tma_atom_col, + tma_view_col, + ).launch(grid=grid, block=[self.THREADS, 1, 1], stream=stream) + + @cute.kernel + def kernel( + self, + mX: cute.Tensor, + mS_row: cute.Tensor, + mS_col: Optional[cute.Tensor], + mAmaxRow: cute.Tensor, + mAmaxCol: Optional[cute.Tensor], + mNoop: cute.Pointer, + mRngState: Optional[cute.Tensor], + tma_atom_in: cute.CopyAtom, + tma_view_in: cute.Tensor, + tma_atom_row: cute.CopyAtom, + tma_view_row: cute.Tensor, + tma_atom_col: Optional[cute.CopyAtom], + tma_view_col: Optional[cute.Tensor], + ): + """Device entry for the NVFP4 tuned-1D quantize-transpose kernel.""" + if not noop_flag_is_set(mNoop): + self.kernel_main( + mX, + mS_row, + mS_col, + mAmaxRow, + mAmaxCol, + mRngState, + tma_atom_in, + tma_view_in, + tma_atom_row, + tma_view_row, + tma_atom_col, + tma_view_col, + ) + + @cute.jit + def kernel_main( + self, + mX: cute.Tensor, + mS_row: cute.Tensor, + mS_col: Optional[cute.Tensor], + mAmaxRow: cute.Tensor, + mAmaxCol: Optional[cute.Tensor], + mRngState: Optional[cute.Tensor], + tma_atom_in: cute.CopyAtom, + tma_view_in: cute.Tensor, + tma_atom_row: cute.CopyAtom, + tma_view_row: cute.Tensor, + tma_atom_col: Optional[cute.CopyAtom], + tma_view_col: Optional[cute.Tensor], + ): + # -- Trace time -- + cfg = self.cfg + TILE = self.TILE_DIM + + # Shared memory layout + if cutlass.const_expr(cfg.RETURN_TRANSPOSE): + + @cute.struct + class SharedStorage: # pylint: disable=missing-class-docstring + mbar_storage: cute.struct.MemRange[cute.Int64, self.NUM_BUFFERS] + sX: cute.struct.Align[ + cute.struct.MemRange[BFloat16, TILE * TILE * self.NUM_BUFFERS], 128 + ] + sO_row: cute.struct.Align[ + cute.struct.MemRange[Uint8, TILE * (TILE // 2) * self.NUM_BUFFERS], 128 + ] + sO_col: cute.struct.Align[ + cute.struct.MemRange[Uint8, TILE * (TILE // 2) * self.NUM_BUFFERS], 128 + ] + sS_row: cute.struct.Align[ + cute.struct.MemRange[Float8E4M3FN, self.CHUNK_DIM_Y * self.SCALES_PER_CHUNK_X], + 16, + ] + sS_col: cute.struct.Align[ + cute.struct.MemRange[Float8E4M3FN, self.CHUNK_DIM_X * self.SCALES_PER_CHUNK_Y], + 16, + ] + + else: + + @cute.struct + class SharedStorage: # pylint: disable=missing-class-docstring + mbar_storage: cute.struct.MemRange[cute.Int64, self.NUM_BUFFERS] + sX: cute.struct.Align[ + cute.struct.MemRange[BFloat16, TILE * TILE * self.NUM_BUFFERS], 128 + ] + sO_row: cute.struct.Align[ + cute.struct.MemRange[Uint8, TILE * (TILE // 2) * self.NUM_BUFFERS], 128 + ] + sS_row: cute.struct.Align[ + cute.struct.MemRange[Float8E4M3FN, self.CHUNK_DIM_Y * self.SCALES_PER_CHUNK_X], + 16, + ] + + # "Allocate" shared memory + smem = cutlass.utils.SmemAllocator() + storage = smem.allocate(SharedStorage) + + # Create views into the shared memory + def buffered_tile(inner_cols): + return cute.make_layout( + ((TILE, inner_cols), self.NUM_BUFFERS), + stride=((inner_cols, 1), TILE * inner_cols), + ) + + sX = storage.sX.get_tensor(buffered_tile(TILE)) + sO_row = storage.sO_row.get_tensor(buffered_tile(TILE // 2)) + sS_row = storage.sS_row.get_tensor( + cute.make_layout( + (self.CHUNK_DIM_Y, self.SCALES_PER_CHUNK_X), stride=(self.SCALES_PER_CHUNK_X, 1) + ) + ) + if cutlass.const_expr(cfg.RETURN_TRANSPOSE): + sO_col = storage.sO_col.get_tensor(buffered_tile(TILE // 2)) + sS_col = storage.sS_col.get_tensor( + cute.make_layout( + (self.CHUNK_DIM_X, self.SCALES_PER_CHUNK_Y), + stride=(self.SCALES_PER_CHUNK_Y, 1), + ) + ) + else: + sO_col = None + sS_col = None + + # Bind GMEM and SMEM tensor views for TMA + gX_tiled = cute.zipped_divide(tma_view_in, (TILE, TILE)) + tXsX, tXgX = cute.nvgpu.cpasync.tma_partition( + tma_atom_in, 0, cute.make_layout(1), sX, gX_tiled + ) + gO_row_tiled = cute.zipped_divide(tma_view_row, (TILE, TILE // 2)) + tOsO_row, tOgO_row = cute.nvgpu.cpasync.tma_partition( + tma_atom_row, 0, cute.make_layout(1), sO_row, gO_row_tiled + ) + if cutlass.const_expr(cfg.RETURN_TRANSPOSE): + gO_col_tiled = cute.zipped_divide(tma_view_col, (TILE, TILE // 2)) + tOsO_col, tOgO_col = cute.nvgpu.cpasync.tma_partition( + tma_atom_col, 0, cute.make_layout(1), sO_col, gO_col_tiled + ) + + # -- Runtime -- + rows, cols = mX.shape + tidx, _, _ = cute.arch.thread_idx() + bidx, bidy, _ = cute.arch.block_idx() # (chunk_x, chunk_y) + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + + # Prefetch TMA descriptor of the input + if warp_idx == 0: + cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_in) + + # Initialize mbarriers for TMA G2S input tensor copy + mbar = storage.mbar_storage.data_ptr() + if warp_idx == 0: + with cute.arch.elect_one(): + for b in cutlass.range_constexpr(self.NUM_BUFFERS): + cute.arch.mbarrier_init(mbar + b, 1) + # release mbarrier_init from elected thread to all threads in cluster (here: CTA) + cute.arch.mbarrier_init_fence() + cute.arch.sync_threads() # acquire mbarrier_init by CTA (cluster) + + # Prefetch NUM_BUFFERS=PREFETCH_STAGES tiles + the first tile to process + tx_count = TILE * TILE * 2 # each load is a TILE x TILE x bf16 (2 bytes) tile + if warp_idx == 0: + for s in cutlass.range_constexpr(self.NUM_BUFFERS): + tile_coord = ( + bidy * self.STAGES_Y + s // self.STAGES_X, + bidx * self.STAGES_X + s % self.STAGES_X, + ) + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx(mbar + s, tx_count) + cute.copy( + tma_atom_in, + tXgX[(None, tile_coord)], + tXsX[(None, s)], + tma_bar_ptr=mbar + s, + ) + + # Compute global encode scales from supplied amax + if cutlass.const_expr(not cfg.ROW_SCALED_NVFP4): + S_enc_rowwise = compute_global_encode_sf(mAmaxRow[0]) + else: + # Per-row encode scales are drawn inside the rowwise pass instead. + S_enc_rowwise = Float32(1.0) + if cutlass.const_expr(cfg.RETURN_TRANSPOSE): + S_enc_colwise = compute_global_encode_sf(mAmaxCol[0]) + else: + S_enc_colwise = Float32(1.0) + + # Construct RNG state for SR + rng = None + if cutlass.const_expr(cfg.USE_STOCHASTIC_ROUNDING): + grid_dim_x, _, _ = cute.arch.grid_dim() + # Contrary to CUDA C++ version, calculate in Int64 for correctness + rng_sequence = ( + Int64(tidx) + + Int64(bidx) * self.THREADS + + Int64(bidy) * Int64(grid_dim_x) * self.THREADS + ) + rng = PhiloxRng( + seed=Uint64(mRngState[0].ir_value()), + subsequence=Uint64(rng_sequence.ir_value()), + offset=Uint64(mRngState[1].ir_value()), + ) + + # Main loop over tiles/stages in chunk + for stage in cutlass.range_constexpr(self.STAGES): + stage_y = stage // self.STAGES_X + stage_x = stage % self.STAGES_X + parity = (stage // self.NUM_BUFFERS) % 2 + buf = stage % self.NUM_BUFFERS + + # Wait for TMA G2S tile load to complete + cute.arch.mbarrier_wait(mbar + buf, parity) + + sX_tile = sX[(None, buf)] + sO_row_tile = sO_row[(None, buf)] + + self.rowwise_tile( + sX_tile, + sO_row_tile, + sS_row, + mAmaxRow, + S_enc_rowwise, + bidy * self.CHUNK_DIM_Y, + rows, + stage_y, + stage_x, + rng, + ) + if cutlass.const_expr(cfg.RETURN_TRANSPOSE): + self.colwise_tile( + sX_tile, + sO_col[(None, buf)], + sS_col, + S_enc_colwise, + stage_y, + stage_x, + rng, + ) + + # Make the SMEM output writes visible to the TMA async proxy. The dedicated + # fence.proxy.async.shared::cta, not the far costlier generic membar + # cute.arch.fence_proxy would emit. + cute.arch.fence_view_async_shared() + # Before the barrier: wait for every prior TMA store to have finished reading its + # staged output buffer. The barrier then broadcasts that to all threads, making the + # next stage's writes into that buffer safe, and orders this stage's SMEM output + # writes before the TMA store below. One barrier serves all three purposes; this is + # the only syncthreads in the loop. + if warp_idx == 0: + cute.arch.cp_async_bulk_wait_group(0, read=True) + cute.arch.sync_threads() + + if warp_idx == 0: + # Store this tile's outputs. The transposed tile of chunk-relative tile (y, x) + # lands at tile (x, y) of the transposed tensor. + cute.copy( + tma_atom_row, + tOsO_row[(None, buf)], + tOgO_row[ + ( + None, + (bidy * self.STAGES_Y + stage_y, bidx * self.STAGES_X + stage_x), + ) + ], + ) + if cutlass.const_expr(cfg.RETURN_TRANSPOSE): + cute.copy( + tma_atom_col, + tOsO_col[(None, buf)], + tOgO_col[ + ( + None, + (bidx * self.STAGES_X + stage_x, bidy * self.STAGES_Y + stage_y), + ) + ], + ) + cute.arch.cp_async_bulk_commit_group() + + # Refill the buffer this stage just consumed: every thread's reads of it + # finished before the syncthreads above. + if cutlass.const_expr(stage + self.NUM_BUFFERS < self.STAGES): + next_stage = stage + self.NUM_BUFFERS + tile_coord = ( + bidy * self.STAGES_Y + next_stage // self.STAGES_X, + bidx * self.STAGES_X + next_stage % self.STAGES_X, + ) + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx(mbar + buf, tx_count) + cute.copy( + tma_atom_in, + tXgX[(None, tile_coord)], + tXsX[(None, buf)], + tma_bar_ptr=mbar + buf, + ) + + ## Epilogue: flush the chunk's staged scales with predicated stores. The last stage's + ## syncthreads ordered every thread's scale writes before this point. This is the one + ## place anything is predicated: the data tensors are covered by TMA's out-of-bounds + ## handling, the scale tensors are not. + self.flush_scales( + sS_row, mS_row, bidy * self.CHUNK_DIM_Y, bidx * self.SCALES_PER_CHUNK_X, rows, cols + ) + if cutlass.const_expr(cfg.RETURN_TRANSPOSE): + self.flush_scales( + sS_col, mS_col, bidx * self.CHUNK_DIM_X, bidy * self.SCALES_PER_CHUNK_Y, cols, rows + ) + + # Wait for in-flight TMA stores before the kernel returns. + cute.arch.cp_async_bulk_wait_group(0, read=False) + + @cute.jit + def rowwise_tile( + self, + sX_tile: cute.Tensor, # (TILE, TILE) bf16 SMEM input tile + sO_row_tile: cute.Tensor, # (TILE, TILE/2) u8 SMEM staged rowwise output + sS_row: cute.Tensor, # (CHUNK_DIM_Y, SCALES_PER_CHUNK_X) e4m3 SMEM staged scales + mAmaxRow: cute.Tensor, # per-row amax (ROW_SCALED) or the global amax (unused here) + S_enc: Float32, # global rowwise encode scale (ignored when ROW_SCALED) + chunk_row0: Int32, # global row of the chunk's first row + rows: Int32, + stage_y: cutlass.Constexpr, + stage_x: cutlass.Constexpr, + rng, + ): + """Quantize one SMEM tile rowwise: a thread owns one 16-element scaling block per + iteration, so the block amax never leaves the thread. Mirrors rowwise_scaling in the + CUDA kernel, including the bank-group stagger of the two 16-byte reads per block.""" + cfg = self.cfg + tidx, _, _ = cute.arch.thread_idx() + lane = tidx % 32 + bank_group = lane // self.THREADS_PER_BANK # which 16-byte wave this thread starts on + tid_y = tidx // self.THREADS_X_ROWWISE + tid_x = tidx % self.THREADS_X_ROWWISE + + if cutlass.const_expr(cfg.USE_FAST_MATH): + sf_type = BFloat16 + else: + sf_type = Float32 + + # The tile as (row * 8-element group, element): a flat first mode so a slice needs only + # one dynamic index, each group one 16-byte vectorized SMEM access. A fresh view on the + # tile's pointer rather than cute.composition, whose right operand maps into the tile's + # colexicographic element order, not its memory order. + groups_per_row = self.TILE_DIM // self.PACK_SIZE + sX_groups = cute.make_tensor( + sX_tile.iterator, + cute.make_layout( + (self.TILE_DIM * groups_per_row, self.PACK_SIZE), stride=(self.PACK_SIZE, 1) + ), + ) + # The staged output as u32: one conversion of 8 elements fills one u32. + sO_u32 = cute.make_tensor( + cute.recast_ptr(sO_row_tile.iterator, dtype=Uint32), + cute.make_layout( + (self.TILE_DIM * self.TILE_DIM // 8,), + stride=(1,), + ), + ) + u32_per_row = self.TILE_DIM // 8 + + for it in cutlass.range_constexpr(self.ITERATIONS_ROWWISE): + row = tid_y + it * self.THREADS_Y_ROWWISE + + # Read this thread's scaling block, one staggered 8-element wave at a time, and + # accumulate the block amax as packed bf16 pairs (max.xorsign.abs.bf16x2). + frg = [] + frg_u32 = [] + for w in cutlass.range_constexpr(self.WAVES): + frg.append(cute.make_rmem_tensor(self.PACK_SIZE, BFloat16)) + frg_u32.append( + cute.make_tensor( + cute.recast_ptr(frg[w].iterator, dtype=Uint32), + cute.make_layout((self.PACK_SIZE // 2,), stride=(1,)), + ) + ) + amax_2x = zero_b32() + for w in cutlass.range_constexpr(self.WAVES): + group = tid_x * self.WAVES + ((w + bank_group) % self.WAVES) + cute.autovec_copy(sX_groups[row * groups_per_row + group, None], frg[w]) + for j in cutlass.range_constexpr(self.PACK_SIZE // 2): + amax_2x = abs_max_x2_bf16(amax_2x, Int32(frg_u32[w][j].ir_value())) + block_amax = cute.arch.fmax( + fabs_f32(x2_lo_to_f32_bf16(amax_2x)), + fabs_f32(x2_hi_to_f32_bf16(amax_2x)), + ) + + # The encode scale: global, or this row's own under row-scaled quantization. An + # out-of-bounds row (the CUDA kernel gives it encode scale 1.0) produces nothing + # observable -- its data rows are dropped by the output TMA and its scale rows are + # not flushed -- so the amax load clamps the index instead of branching. + if cutlass.const_expr(cfg.ROW_SCALED_NVFP4): + row_idx = chunk_row0 + stage_y * self.TILE_DIM + row + S_enc_block = compute_global_encode_sf(mAmaxRow[cutlass.min(row_idx, rows - 1)]) + else: + S_enc_block = S_enc + + block_decode_sf = compute_block_decode_sf(block_amax, S_enc_block) + sS_row[stage_y * self.TILE_DIM + row, stage_x * self.SCALES_PER_TILE + tid_x] = ( + block_decode_sf + ) + coeff = compute_block_encode_sf(block_decode_sf, S_enc_block, sf_type) + + # Scale and convert one wave (8 elements, one u32 of nibbles) at a time, storing to + # the same staggered group the wave was read from. Each coefficient type gets the + # multiply the CUDA kernel pairs it with: an fma against a zero addend for the bf16 + # one (which flushes a -0 product to +0, and E2M1 has a signed zero, so the + # instruction has to match) and a plain f32 multiply for the f32 one. + for w in cutlass.range_constexpr(self.WAVES): + group = tid_x * self.WAVES + ((w + bank_group) % self.WAVES) + if cutlass.const_expr(cfg.USE_FAST_MATH): + coeff_f32 = coeff.to(Float32) + if cutlass.const_expr(cfg.USE_STOCHASTIC_ROUNDING): + rbits03 = rng.get_rbits() + rbits47 = rng.get_rbits() + out = mul_cvt_bf16x8_to_fp4x8_sr( + frg_u32[w][0], + frg_u32[w][1], + frg_u32[w][2], + frg_u32[w][3], + coeff_f32, + rbits03, + rbits47, + ) + else: + out = mul_cvt_bf16x8_to_fp4x8( + frg_u32[w][0], frg_u32[w][1], frg_u32[w][2], frg_u32[w][3], coeff_f32 + ) + else: + scaled = [] + for j in cutlass.range_constexpr(self.PACK_SIZE // 2): + pair = Int32(frg_u32[w][j].ir_value()) + scaled.append(x2_lo_to_f32_bf16(pair) * coeff) + scaled.append(x2_hi_to_f32_bf16(pair) * coeff) + if cutlass.const_expr(cfg.USE_STOCHASTIC_ROUNDING): + rbits03 = rng.get_rbits() + rbits47 = rng.get_rbits() + # pylint: disable-next=no-value-for-parameter + out = cvt_f32x8_to_fp4x8_sr(*scaled, rbits03, rbits47) + else: + out = cvt_f32x8_to_fp4x8(*scaled) # pylint: disable=no-value-for-parameter + sO_u32[row * u32_per_row + group] = out + + @cute.jit + def colwise_tile( + self, + sX_tile: cute.Tensor, # (TILE, TILE) bf16 SMEM input tile + sO_col_tile: cute.Tensor, # (TILE, TILE/2) u8 SMEM staged transposed output + sS_col: cute.Tensor, # (CHUNK_DIM_X, SCALES_PER_CHUNK_Y) e4m3 SMEM staged scales + S_enc: Float32, + stage_y: cutlass.Constexpr, + stage_x: cutlass.Constexpr, + rng, + ): + """Quantize one SMEM tile columnwise into the transposed staged output: a thread owns + two adjacent columns of one 16-row scaling block, read as bf16 pairs, with the block row + staggered by warp so consecutive lanes hit different SMEM rows (conflict-free 4-byte + reads). Mirrors colwise_scaling in the CUDA kernel.""" + cfg = self.cfg + tidx, _, _ = cute.arch.thread_idx() + lane = tidx % 32 + warp = tidx // 32 + tid_y = (lane // 2 + warp) % self.SCALES_PER_TILE # which 16-row block + + if cutlass.const_expr(cfg.USE_FAST_MATH): + sf_type = BFloat16 + else: + sf_type = Float32 + + # The tile as (row, column pair): a pair-register read per row. + sX_pairs = cute.make_tensor( + cute.recast_ptr(sX_tile.iterator, dtype=Uint32), + cute.make_layout( + (self.TILE_DIM * (self.TILE_DIM // 2),), + stride=(1,), + ), + ) + pairs_per_row = self.TILE_DIM // 2 + # The transposed staged output as u64: one 16-element block of a transposed row is 8 + # bytes, one two-conversion store. + sO_u64 = cute.make_tensor( + cute.recast_ptr(sO_col_tile.iterator, dtype=Int64), + cute.make_layout( + (self.TILE_DIM * self.TILE_DIM // 16,), + stride=(1,), + ), + ) + u64_per_row = self.TILE_DIM // 16 + + # Read the two columns' 16-row block as pairs, accumulating both block amaxes at once + # in the packed halves. + row0 = tid_y * NVFP4_BLOCK_SCALING_SIZE + pairs = [] + amax_2x = zero_b32() + for i in cutlass.range_constexpr(NVFP4_BLOCK_SCALING_SIZE): + pair = sX_pairs[(row0 + i) * pairs_per_row + lane] + pairs.append(pair) + amax_2x = abs_max_x2_bf16(amax_2x, Int32(pair.ir_value())) + block_amax = [ + fabs_f32(x2_lo_to_f32_bf16(amax_2x)), + fabs_f32(x2_hi_to_f32_bf16(amax_2x)), + ] + + for w in cutlass.range_constexpr(2): + block_decode_sf = compute_block_decode_sf(block_amax[w], S_enc) + sS_col[ + stage_x * self.TILE_DIM + 2 * lane + w, + stage_y * self.SCALES_PER_TILE + tid_y, + ] = block_decode_sf + coeff = compute_block_encode_sf(block_decode_sf, S_enc, sf_type) + + outs = [] + if cutlass.const_expr(cfg.USE_FAST_MATH): + coeff_f32 = coeff.to(Float32) + # Repack this column's elements into adjacent bf16 pairs for the fma-based + # conversion: bytes of (row 2j, row 2j+1) picked by a byte permute. + prmt = prmt_lo_u32 if w == 0 else prmt_hi_u32 + packed = [ + prmt(pairs[2 * j], pairs[2 * j + 1]) + for j in range(NVFP4_BLOCK_SCALING_SIZE // 2) + ] + for e in cutlass.range_constexpr(NVFP4_BLOCK_SCALING_SIZE // 8): + if cutlass.const_expr(cfg.USE_STOCHASTIC_ROUNDING): + rbits03 = rng.get_rbits() + rbits47 = rng.get_rbits() + outs.append( + mul_cvt_bf16x8_to_fp4x8_sr( + packed[4 * e], + packed[4 * e + 1], + packed[4 * e + 2], + packed[4 * e + 3], + coeff_f32, + rbits03, + rbits47, + ) + ) + else: + outs.append( + mul_cvt_bf16x8_to_fp4x8( + packed[4 * e], + packed[4 * e + 1], + packed[4 * e + 2], + packed[4 * e + 3], + coeff_f32, + ) + ) + else: + widen = x2_lo_to_f32_bf16 if w == 0 else x2_hi_to_f32_bf16 + scaled = [ + widen(Int32(pairs[i].ir_value())) * coeff + for i in range(NVFP4_BLOCK_SCALING_SIZE) + ] + for e in cutlass.range_constexpr(NVFP4_BLOCK_SCALING_SIZE // 8): + if cutlass.const_expr(cfg.USE_STOCHASTIC_ROUNDING): + rbits03 = rng.get_rbits() + rbits47 = rng.get_rbits() + outs.append( + cvt_f32x8_to_fp4x8_sr( + *scaled[8 * e : 8 * e + 8], rbits03, rbits47 + ) # pylint: disable=no-value-for-parameter + ) + else: + outs.append( + cvt_f32x8_to_fp4x8(*scaled[8 * e : 8 * e + 8]) + ) # pylint: disable=no-value-for-parameter + + sO_u64[(2 * lane + w) * u64_per_row + tid_y] = pack_u32x2(outs[0], outs[1]) + + @cute.jit + def flush_scales( + self, + sS: cute.Tensor, # (128, 8) e4m3 SMEM staged scales for the whole chunk + mS: cute.Tensor, # padded gmem scale tensor + chunk_row0: Int32, # global row of the chunk's first scale row + chunk_sf_col0: Int32, # global scale column of the chunk's first scale + outer: Int32, # valid rows of the scale tensor (input rows / cols) + inner_elems: Int32, # input elements along the scaled direction (cols / rows) + # -> valid scale columns = inner_elems / 16 + ): + """Write one chunk's staged scale bytes out, one row per thread, predicated against the + input's real extents: the padding region of the scale tensors is left untouched and + rows/columns past the input's edge are skipped, like the CUDA kernel's Vec-based scale + store. An interior chunk takes the vectorized path -- one 8-byte SMEM read and two + 4-byte gmem stores (the padded scale stride only guarantees 4-byte alignment) -- since + the elementwise fallback's divergent byte stores are what an unoptimized epilogue + spends most of the kernel's stalls on.""" + tidx, _, _ = cute.arch.thread_idx() + row_global = chunk_row0 + tidx + count = cutlass.min( + Int32(self.SCALES_PER_CHUNK_X), + (inner_elems - chunk_sf_col0 * NVFP4_BLOCK_SCALING_SIZE) // NVFP4_BLOCK_SCALING_SIZE, + ) + frg = cute.make_rmem_tensor(self.SCALES_PER_CHUNK_X, Float8E4M3FN) + cute.autovec_copy(sS[tidx, None], frg) + frg_u32 = cute.make_tensor( + cute.recast_ptr(frg.iterator, dtype=Uint32), + cute.make_layout((self.SCALES_PER_CHUNK_X // 4,), stride=(1,)), + ) + # A u32 view of the scale tensor: its padded inner extent (= the row stride) is a + # multiple of NVFP4_SCALE_PAD_COLS = 4, so rows stay u32-aligned, and a chunk's first + # scale column is a multiple of SCALES_PER_CHUNK_X. + scale_stride = cute.size(mS.shape[1]) + mS_u32 = cute.make_tensor( + cute.recast_ptr(mS.iterator, dtype=Uint32), + cute.make_layout((mS.shape[0], scale_stride // 4), stride=(scale_stride // 4, 1)), + ) + if row_global < outer: + if count == self.SCALES_PER_CHUNK_X: + for c4 in cutlass.range_constexpr(self.SCALES_PER_CHUNK_X // 4): + mS_u32[row_global, chunk_sf_col0 // 4 + c4] = frg_u32[c4] + else: + for c in cutlass.range_constexpr(self.SCALES_PER_CHUNK_X): + if c < count: + mS[row_global, chunk_sf_col0 + c] = frg[c] diff --git a/transformer_engine/common/CuTeDSL/philox_rng.py b/transformer_engine/common/CuTeDSL/philox_rng.py new file mode 100644 index 00000000000..afb757d6164 --- /dev/null +++ b/transformer_engine/common/CuTeDSL/philox_rng.py @@ -0,0 +1,223 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Philox random-number generation for CuTeDSL kernels. Based on curanddx.hpp""" + +import os + +import cutlass +from cutlass import cute +from cutlass import Boolean, Uint32, Uint64 + +# per envvars.rst +NUM_PHILOX_ROUNDS = int(os.environ.get("NVTE_BUILD_NUM_PHILOX_ROUNDS", "10")) + +# per curanddx.hpp +PHILOX_W32_0 = 0x9E3779B9 +PHILOX_W32_1 = 0xBB67AE85 +PHILOX_M4x32_0 = 0xD2511F53 +PHILOX_M4x32_1 = 0xCD9E8D57 + + +@cute.jit +def umulhi_u32(a: Uint32, b: Uint32) -> Uint32: + """High 32 bits of the unsigned 32x32 product (`__umulhi`).""" + return cute.arch.inline_ptx( + "mul.hi.u32 {$w0}, {$r0}, {$r1};", + write_only_types=[Uint32], + read_only_args=[a, b], + ) + + +@cute.jit +def u64_lo32(v: Uint64) -> Uint32: + """Low 32 bits of a u64.""" + return Uint32(v) + + +@cute.jit +def u64_hi32(v: Uint64) -> Uint32: + """High 32 bits of a u64.""" + return Uint32(v >> Uint64(32)) + + +@cute.jit +def bool_to_u64(b: Boolean) -> Uint64: + """0 or 1 as a u64.""" + return Uint64(b) + + +class PhiloxRng: + """Trace-time replica of the curanddx philox4x32 generator the CUDA kernel rounds with. + + The counter is kept as two u64 halves: the low half is the offset the state was initialized + with and the high half the subsequence, exactly the state philox4x32_native_state::init + reaches (both increments start from a zero counter, so their carry chains vanish), and + incrementing the 128-bit counter is then one add-with-carry. Consumption is deterministic -- + every loop that draws random words is unrolled -- so the one-of-four word cycling of + core::get_rbits is a Python-side list rather than a device-side index. + """ + + def __init__(self, seed: Uint64, subsequence: Uint64, offset: Uint64): + self._key = (u64_lo32(seed), u64_hi32(seed)) + self._ctr_lo = offset + self._ctr_hi = subsequence + self._buf = [] + # Matches the CUDA kernel's `random_uint4 = rng.generate4()` at init. + self._generate4() + + def _generate4(self): + c = [ + u64_lo32(self._ctr_lo), + u64_hi32(self._ctr_lo), + u64_lo32(self._ctr_hi), + u64_hi32(self._ctr_hi), + ] + k0, k1 = self._key + for round_idx in range(NUM_PHILOX_ROUNDS): + if round_idx > 0: + k0 = k0 + Uint32(PHILOX_W32_0) + k1 = k1 + Uint32(PHILOX_W32_1) + lo0 = Uint32(PHILOX_M4x32_0) * c[0] + hi0 = umulhi_u32(Uint32(PHILOX_M4x32_0), c[0]) + lo1 = Uint32(PHILOX_M4x32_1) * c[2] + hi1 = umulhi_u32(Uint32(PHILOX_M4x32_1), c[2]) + c = [hi1 ^ c[1] ^ k0, lo1, hi0 ^ c[3] ^ k1, lo0] + self._buf = c + new_lo = self._ctr_lo + Uint64(1) + self._ctr_hi = self._ctr_hi + bool_to_u64(new_lo == Uint64(0)) + self._ctr_lo = new_lo + + def get_rbits(self) -> Uint32: + """The next 32 random bits, regenerating every four words like core::get_rbits.""" + if not self._buf: + self._generate4() + return self._buf.pop(0) + + +# --- Block draws: one Philox4x32 generate per output block, several streams at a time --- +# +# Each draw is exactly what curanddx's philox4x32_native_state returns from +# init(seed, subsequence, offset) followed by one generate4(): the counter is +# {offset_lo, offset_hi, subsequence_lo, subsequence_hi}. Only the cost differs from PhiloxRng: +# - every round-0 term that reads only the launch-wide offset is evaluated once +# (make_block_philox); +# - when a thread draws for many blocks that share one subsequence word (its block column, say), +# placing that word in the subsequence's LOW half makes both of round 0's products thread +# constants, and the constant parts of rounds 0..2 fold away (make_block_philox_fold). The +# other half carries the varying coordinate XORed with the launch constant `b`, a bijection that +# cancels round 0's own XOR. + +# make_block_philox_fold consumes rounds 0..2, so callers need a plain generator below 3 rounds. +BLOCK_PHILOX_FOLD_SUPPORTED = NUM_PHILOX_ROUNDS >= 3 + + +# Subsequence of each output block, as (low word, logical high word). The low word is the +# coordinate a thread keeps fixed while it walks the other, so make_block_philox_fold can fold +# it. Its parity tells the directions apart, so the rowwise and columnwise draws of one element +# never share a subsequence; within a direction the pair is the block's coordinates. +# rowwise block (row, block_col): (2 * block_col, row) +# columnwise block (row_block, col), col = 4q + i: (2 * (4q) + 1, 4 * row_block + i) +# The column's offset inside its group of four moves to the high word because a thread owns four +# adjacent columns and folds against their base. +@cute.jit +def rowwise_sub_lo(block_col: Uint32) -> Uint32: + """Subsequence low word of a rowwise block.""" + return block_col << Uint32(1) + + +@cute.jit +def columnwise_sub_lo(quad_base_col: Uint32) -> Uint32: + """Subsequence low word of the columnwise blocks of four adjacent columns.""" + return (quad_base_col << Uint32(1)) | Uint32(1) + + +@cute.jit +def columnwise_sub_hi(row_block: Uint32, col_in_quad: Uint32) -> Uint32: + """Logical subsequence high word of a columnwise block.""" + return (row_block << Uint32(2)) | col_in_quad + + +@cute.jit +def mulhilo_u32(a: Uint32, b: Uint32): + """Both halves of a 32 x 32 -> 64-bit product in one IMAD.WIDE.U32, as (lo, hi).""" + p = cute.arch.inline_ptx( + "mul.wide.u32 {$w0}, {$r0}, {$r1};", + write_only_types=[Uint64], + read_only_args=[a, b], + ) + return Uint32(p), Uint32(p >> Uint64(32)) + + +@cute.jit +def xor3_u32(a: Uint32, b: Uint32, c: Uint32) -> Uint32: + """Three-way XOR in one LOP3.""" + return cute.arch.inline_ptx( + "lop3.b32 {$w0}, {$r0}, {$r1}, {$r2}, 0x96;", + write_only_types=[Uint32], + read_only_args=[a, b, c], + ) + + +@cute.jit +def make_block_philox(seed: Uint64, offset: Uint64): + """Round-0 terms that read only the key and the launch-wide offset, as + (key0, key1, a, b, w): the key round 1 uses, a = offset_hi ^ k0 (round-0 c.x is + hi(sub_lo * M1) ^ a), b = hi(offset_lo * M0) ^ k1 (round-0 c.z is b ^ sub_hi) and + w = lo(offset_lo * M0) (round-0 c.w).""" + k0 = u64_lo32(seed) + k1 = u64_hi32(seed) + lo0, hi0 = mulhilo_u32(u64_lo32(offset), Uint32(PHILOX_M4x32_0)) + return ( + k0 + Uint32(PHILOX_W32_0), + k1 + Uint32(PHILOX_W32_1), + u64_hi32(offset) ^ k0, + hi0 ^ k1, + lo0, + ) + + +@cute.jit +def make_block_philox_fold(state, sub_lo: Uint32): + """Rounds 0..2 for a fixed subsequence low word, as (A, B, D, c1_2, k3x, k3y): + c0_1 = hi1 ^ A, c0_2 = B ^ c1_1, c2_2 = hi0 ^ D, c1_2 the round-2 c1 and k3 the round-3 key.""" + kx, ky, a, _, w = state + lo1, hi1 = mulhilo_u32(sub_lo, Uint32(PHILOX_M4x32_1)) # round 0: both products constant + c0_0 = hi1 ^ a + c1_0 = lo1 + lo0, hi0 = mulhilo_u32(c0_0, Uint32(PHILOX_M4x32_0)) # round 1: the .x product is constant + fold_a = c1_0 ^ kx + c2_1 = xor3_u32(hi0, w, ky) + c3_1 = lo0 + kx = kx + Uint32(PHILOX_W32_0) + ky = ky + Uint32(PHILOX_W32_1) # round-2 key + lo1, hi1 = mulhilo_u32(c2_1, Uint32(PHILOX_M4x32_1)) # round 2: the .z product is constant + fold_b = hi1 ^ kx + c1_2 = lo1 + fold_d = c3_1 ^ ky + kx = kx + Uint32(PHILOX_W32_0) + ky = ky + Uint32(PHILOX_W32_1) # round-3 key + return (fold_a, fold_b, fold_d, c1_2, kx, ky) + + +@cute.jit +def block_philox_draws(logical_hi, fold): + """One draw per entry of `logical_hi` (the varying coordinate of a block sharing the fold's + low word; the draw's subsequence high word is logical_hi ^ b), as a list of 4-word lists. + The draws advance together, so the key schedule is paid once per batch.""" + fold_a, fold_b, fold_d, c1_2, kx, ky = fold + c = [] + for hi in logical_hi: # the varying remnant of rounds 0..2 + lo1, hi1 = mulhilo_u32(hi, Uint32(PHILOX_M4x32_1)) + lo0, hi0 = mulhilo_u32(hi1 ^ fold_a, Uint32(PHILOX_M4x32_0)) + c.append([fold_b ^ lo1, c1_2, hi0 ^ fold_d, lo0]) + for r in cutlass.range_constexpr(3, NUM_PHILOX_ROUNDS): + for i in cutlass.range_constexpr(len(c)): + lo0, hi0 = mulhilo_u32(c[i][0], Uint32(PHILOX_M4x32_0)) + lo1, hi1 = mulhilo_u32(c[i][2], Uint32(PHILOX_M4x32_1)) + c[i] = [xor3_u32(hi1, c[i][1], kx), lo1, xor3_u32(hi0, c[i][3], ky), lo0] + if cutlass.const_expr(r < NUM_PHILOX_ROUNDS - 1): + kx = kx + Uint32(PHILOX_W32_0) + ky = ky + Uint32(PHILOX_W32_1) + return c diff --git a/transformer_engine/common/CuTeDSL/utils.py b/transformer_engine/common/CuTeDSL/utils.py index ed05aa1c3c6..12096d93a97 100644 --- a/transformer_engine/common/CuTeDSL/utils.py +++ b/transformer_engine/common/CuTeDSL/utils.py @@ -9,7 +9,9 @@ import cutlass from cutlass import cute -from cutlass import Float32, Int64, Int32, Uint16, Uint32 +from cutlass import Boolean, Float32, Int64, Int32, Uint16, Uint32, Uint64 +from cutlass._mlir.dialects import arith as mlir_arith +from cutlass.cutlass_dsl import dsl_user_op # This must match `inline std::string to_string(const DType type)`` in transformer_engine/common/common.h exactly # which is how we encode types to strings in C++ @@ -51,6 +53,10 @@ def cutlass_dtype_to_str(dtype): FP32_MANTISSA_BITS = 23 +# Largest finite values, which the scaling math clamps to as the CUDA kernels do with +# TypeExtrema::max. CuTeDSL's numeric types do not expose them. +FLOAT32_MAX = float.fromhex("0x1.fffffep127") +BFLOAT16_MAX = float.fromhex("0x1.fep127") @cute.jit @@ -258,3 +264,180 @@ def truncate_f32_f16(val: Float32) -> Float32: def is_packed16(dtype) -> bool: """True if `dtype` is one of the 16-bit packed input formats.""" return dtype is cutlass.BFloat16 or dtype is cutlass.Float16 + + +@dsl_user_op +def select_f32(cond: Boolean, if_true: Float32, if_false: Float32, *, loc=None, ip=None) -> Float32: + """`cond ? if_true : if_false` as one select, never a branch.""" + return Float32( + mlir_arith.select( + cond.ir_value(loc=loc, ip=ip), + if_true.ir_value(loc=loc, ip=ip), + if_false.ir_value(loc=loc, ip=ip), + loc=loc, + ip=ip, + ) + ) + + +@cute.jit +def pack_u32x2(lo: Uint32, hi: Uint32) -> Int64: + """Pack two u32 into one 64-bit register (register-pair move, no real instruction).""" + return cute.arch.inline_ptx( + "mov.b64 {$w0}, {{$r0}, {$r1}};", + write_only_types=[Int64], + read_only_args=[lo, hi], + ) + + +def make_prmt_u32(selector: int): + """A byte-permute op with the 16-bit selector baked in as an immediate. + + prmt.b32 indexes the eight source bytes {a0..a3, b0..b7} and each selector nibble picks the + byte for one destination position, low nibble first. 0x5410 interleaves the low halves of a + and b (a0 a1 b0 b1), 0x7632 the high halves. + """ + + @cute.jit + def prmt_u32(a: Uint32, b: Uint32) -> Uint32: + return cute.arch.inline_ptx( + f"prmt.b32 {{$w0}}, {{$r0}}, {{$r1}}, {selector:#x};", + write_only_types=[Uint32], + read_only_args=[a, b], + ) + + return prmt_u32 + + +@cute.jit +def swap_halves_b32(v: Uint32) -> Uint32: + """Exchange the two 16-bit halves of a word.""" + return cute.arch.inline_ptx( + "prmt.b32 {$w0}, {$r0}, {$r0}, 0x1032;", write_only_types=[Uint32], read_only_args=[v] + ) + + +@cute.jit +def rcp_rn_f32(x: Float32) -> Float32: + """Correctly rounded reciprocal, bit for bit the IEEE quotient 1 / x.""" + return cute.arch.inline_ptx( + "rcp.rn.f32 {$w0}, {$r0};", write_only_types=[Float32], read_only_args=[x] + ) + + +def abs_max_tree_bf16x2(words): + """Packed abs-max over a power-of-two count of words, as a balanced tree: the magnitudes of + both halves are the abs-maxes of the even and odd elements (signs are unreliable, + max.xorsign.abs leaves the XOR of the operands' signs).""" + level = list(words) + while len(level) > 1: + level = [ + Uint32(abs_max_x2_bf16(Int32(level[2 * i]), Int32(level[2 * i + 1]))) + for i in range(len(level) // 2) + ] + return level[0] + + +# --- Global memory access by address --- +# +# The register-resident (g2r) kernels address global memory with 64-bit byte addresses rather +# than through layouts: every access is one row offset from a thread-constant base, and the cache +# hints and vector widths are part of what makes them fast. + + +def _vec(prefix: str, count: int) -> str: + return "{" + ", ".join(f"{{${prefix}{i}}}" for i in range(count)) + "}" + + +@cute.jit +def row_offset(row: Int32, pitch: Int32) -> Int64: + """row * pitch for non-negative 32-bit operands. A product of two zero-extended words lowers + to one IMAD.WIDE.U32; plain arithmetic rather than inline PTX, so ptxas can recompute an + address instead of keeping it live.""" + return Int64(Uint64(Uint32(row)) * Uint64(Uint32(pitch))) + + +@cute.jit +def prefetch_l2(addr: Int64): + """Bring one global cache line into L2. Names no destination register, so it can be issued + ahead of a chain the thread is already stalled on.""" + cute.arch.inline_ptx("prefetch.global.L2 [{$r0}];", read_only_args=[addr]) + + +@cute.jit +def ld_global_nc_evict_first_v8(addr: Int64): + """32 bytes of input read once: read-only path, evicted first from L2. 32-byte aligned.""" + return list( + cute.arch.inline_ptx( + "ld.global.nc.L2::evict_first.v8.b32 " + _vec("w", 8) + ", [{$r0}];", + write_only_types=[Uint32] * 8, + read_only_args=[addr], + ) + ) + + +@cute.jit +def ld_global_evict_last_v8(addr: Int64): + """32 bytes of input that is read again soon: coherent path (the read-only path would not + keep the line for the re-read) and kept in L2 until then. 32-byte aligned.""" + return list( + cute.arch.inline_ptx( + "ld.global.L2::evict_last.v8.b32 " + _vec("w", 8) + ", [{$r0}];", + write_only_types=[Uint32] * 8, + read_only_args=[addr], + ) + ) + + +@cute.jit +def ld_global_v2(addr: Int64): + """8 bytes, 8-byte aligned.""" + return list( + cute.arch.inline_ptx( + "ld.global.v2.b32 " + _vec("w", 2) + ", [{$r0}];", + write_only_types=[Uint32] * 2, + read_only_args=[addr], + ) + ) + + +@cute.jit +def st_global_u8(addr: Int64, value: Uint32): + """Store the low byte of value.""" + cute.arch.inline_ptx("st.global.u8 [{$r0}], {$r1};", read_only_args=[addr, value]) + + +@cute.jit +def st_global_u16(addr: Int64, value: Uint32): + """Store the low two bytes of value, 2-byte aligned.""" + cute.arch.inline_ptx("st.global.u16 [{$r0}], {$r1};", read_only_args=[addr, value]) + + +@cute.jit +def st_global_u32(addr: Int64, value: Uint32): + """4-byte aligned.""" + cute.arch.inline_ptx("st.global.b32 [{$r0}], {$r1};", read_only_args=[addr, value]) + + +@cute.jit +def st_global_v2(addr: Int64, w0: Uint32, w1: Uint32): + """8 bytes, 8-byte aligned.""" + cute.arch.inline_ptx("st.global.v2.b32 [{$r0}], {{$r1}, {$r2}};", read_only_args=[addr, w0, w1]) + + +@cute.jit +def st_global_v4(addr: Int64, words): + """16 bytes, 16-byte aligned.""" + cute.arch.inline_ptx( + "st.global.v4.b32 [{$r0}], {{$r1}, {$r2}, {$r3}, {$r4}};", + read_only_args=[addr] + list(words), + ) + + +@cute.jit +def st_global_v8(addr: Int64, words): + """32 bytes (one sector), 32-byte aligned.""" + cute.arch.inline_ptx( + "st.global.v8.b32 [{$r0}], {" + ", ".join(f"{{$r{i}}}" for i in range(1, 9)) + "};", + read_only_args=[addr] + list(words), + ) diff --git a/transformer_engine/common/CuTeDSL/utils_fp8.py b/transformer_engine/common/CuTeDSL/utils_fp8.py index b3b73e7998d..c02d5ffece5 100644 --- a/transformer_engine/common/CuTeDSL/utils_fp8.py +++ b/transformer_engine/common/CuTeDSL/utils_fp8.py @@ -17,6 +17,10 @@ logger = logging.getLogger("transformer_engine.cutedsl.utils_fp8") +# Largest finite FP8 values. +FP8E4M3_MAX_NORM = 448.0 +FP8E5M2_MAX_NORM = 57344.0 + @cute.jit def cvt_f32_to_fp8e8m0fnu(val: Float32) -> Float8E8M0FNU: @@ -291,3 +295,20 @@ def get_cvt_f32x2_to_fp8x2_func(fp8_dtype, relu: bool = False) -> Callable[..., if relu: return functools.partial(fn, relu=True) return fn + + +@cute.jit +def cvt_f32_to_fp8e4m3_code(val: Float32) -> Uint32: + """The E4M3 byte of val (round to nearest, saturating to 448), as static_cast does.""" + return Uint32(cvt_f32x2_to_fp8e4m3x2(val, val)) & Uint32(0xFF) + + +@cute.jit +def cvt_fp8e4m3_code_to_f32(code: Uint32) -> Float32: + """The value of the E4M3 byte held in the low bits of `code`, exactly.""" + return cute.arch.inline_ptx( + "{ .reg .b16 c, lo, hi; .reg .b32 h; cvt.u16.u32 c, {$r0}; " + "cvt.rn.f16x2.e4m3x2 h, c; mov.b32 {lo, hi}, h; cvt.f32.f16 {$w0}, lo; }", + write_only_types=[Float32], + read_only_args=[code], + ) diff --git a/transformer_engine/common/CuTeDSL/utils_nvfp4.py b/transformer_engine/common/CuTeDSL/utils_nvfp4.py new file mode 100644 index 00000000000..c930a3f6598 --- /dev/null +++ b/transformer_engine/common/CuTeDSL/utils_nvfp4.py @@ -0,0 +1,412 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""NVFP4 scaling math and E2M1 conversion helpers for the CuTeDSL kernels.""" + +import cutlass +from cutlass import cute +from cutlass import BFloat16, Float32, Float8E4M3FN, Int64, Uint32 +from cutlass._mlir.dialects import llvm +from cutlass.cutlass_dsl import T, dsl_user_op + +from transformer_engine.common.CuTeDSL.utils import ( + BFLOAT16_MAX, + FLOAT32_MAX, + abs_max_tree_bf16x2, + rcp_rn_f32, + select_f32, + swap_halves_b32, +) +from transformer_engine.common.CuTeDSL.utils_fp8 import ( + FP8E4M3_MAX_NORM, + cvt_f32_to_fp8e4m3_code, + cvt_fp8e4m3_code_to_f32, +) + +# Largest finite E2M1 value. +FP4E2M1_MAX_NORM = 6.0 + + +@cute.jit +def compute_global_encode_sf(global_amax: Float32) -> Float32: + """like compute_global_encode_scaling_factor_FP4 in core_nvfp4.cuh + + Branch-free: the degenerate-amax fallback is a select rather than an `if`, because the + row-scaled configuration evaluates this per scaling block and the DSL lowers a dynamic `if` + to real divergent branches where the CUDA kernel's ternary gets predicated selects. + """ + global_encode_scale = FP8E4M3_MAX_NORM * FP4E2M1_MAX_NORM / global_amax + global_encode_scale = cute.arch.fmin(global_encode_scale, FLOAT32_MAX) + degenerate = (global_amax == 0.0) | (global_encode_scale == 0.0) + return select_f32(degenerate, Float32(1.0), global_encode_scale) + + +@cute.jit +def compute_block_decode_sf(block_amax: Float32, global_encode_sf: Float32) -> Float8E4M3FN: + "like quantization_and_transposition_SF::compute_decoding_scaling_factor in core_nvfp4.cuh" + block_decode_sf = block_amax * (global_encode_sf * (1.0 / FP4E2M1_MAX_NORM)) + block_decode_sf = cute.arch.fmin(block_decode_sf, FLOAT32_MAX) + return block_decode_sf.to(Float8E4M3FN) + + +@cute.jit +def compute_block_encode_sf( + block_decode_sf: Float8E4M3FN, + global_encode_sf: Float32, + sf_type: cutlass.Constexpr[type], +): + "like compute_nvfp4_scaling_coefficient in quantize_tranpose_nvfp4_tuned_1D.cuh" + if cutlass.const_expr(sf_type == Float32): + global_decode_sf = 1.0 / global_encode_sf + block_encode_sf = 1.0 / (block_decode_sf.to(Float32) * global_decode_sf) + block_encode_sf = cute.arch.fmin(block_encode_sf, FLOAT32_MAX) + return block_encode_sf + if cutlass.const_expr(sf_type == BFloat16): + block_encode_sf = global_encode_sf / block_decode_sf.to(Float32) + block_encode_sf = cute.arch.fmin(block_encode_sf, BFLOAT16_MAX) + return block_encode_sf.to(BFloat16) + raise ValueError("Unsupported scaling-factor type. Only FP32 and BF16 are supported.") + + +@dsl_user_op +def cvt_f32x8_to_fp4x8( + v0: Float32, + v1: Float32, + v2: Float32, + v3: Float32, + v4: Float32, + v5: Float32, + v6: Float32, + v7: Float32, + *, + loc=None, + ip=None, +) -> Uint32: + """Convert eight f32 values to eight E2M1 values packed into four bytes, v0 in the lowest. + + `cvt d, a, b` puts fp4(a) in d[7:4] and fp4(b) in d[3:0], so each pair is fed as (odd, even) + to keep the earlier element in the low nibble. This is the same instruction, and hence the + same rounding and saturation, as mul_cvt_bf16_to_fp4_8x_round_to_nearest in ptx.cuh. + The scaling multiply is left to the caller rather than fused, which is exact because the + instruction that one fuses is a plain `mul.f32x2`; the bf16 coefficient fuses an fma instead + and is not interchangeable that way, hence mul_cvt_bf16x8_to_fp4x8 below. + """ + asm = ( + "{\n" + ".reg.b8 f0; .reg.b8 f1; .reg.b8 f2; .reg.b8 f3;\n\t" + "cvt.rn.satfinite.e2m1x2.f32 f0, $2, $1;\n\t" + "cvt.rn.satfinite.e2m1x2.f32 f1, $4, $3;\n\t" + "cvt.rn.satfinite.e2m1x2.f32 f2, $6, $5;\n\t" + "cvt.rn.satfinite.e2m1x2.f32 f3, $8, $7;\n\t" + "mov.b32 $0, {f0, f1, f2, f3};\n\t" + "}" + ) + 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), + v4.ir_value(loc=loc, ip=ip), + v5.ir_value(loc=loc, ip=ip), + v6.ir_value(loc=loc, ip=ip), + v7.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, + ) + ) + + +@dsl_user_op +def cvt_f32x8_to_fp4x8_sr( + v0: Float32, + v1: Float32, + v2: Float32, + v3: Float32, + v4: Float32, + v5: Float32, + v6: Float32, + v7: Float32, + rbits03: Uint32, + rbits47: Uint32, + *, + loc=None, + ip=None, +) -> Uint32: + """Stochastic-rounding variant of cvt_f32x8_to_fp4x8: the same packing, but through + `cvt.rs.satfinite.e2m1x4.f32`, which rounds each conversion with random bits. One b16 result + holds four elements, `{v3, v2, v1, v0}` keeps v0 in the lowest nibble, and each four-element + group consumes one 32-bit random word, exactly like + mul_cvt_bf16_to_fp4_8x_stochastic_rounding in ptx.cuh (whose multiplies are plain + `mul.f32`, which is what the caller's scaling multiply lowers to). + """ + asm = ( + "{\n" + ".reg.b16 b03, b47;\n\t" + "cvt.rs.satfinite.e2m1x4.f32 b03, {$4, $3, $2, $1}, $9;\n\t" + "cvt.rs.satfinite.e2m1x4.f32 b47, {$8, $7, $6, $5}, $10;\n\t" + "mov.b32 $0, {b03, b47};\n\t" + "}" + ) + 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), + v4.ir_value(loc=loc, ip=ip), + v5.ir_value(loc=loc, ip=ip), + v6.ir_value(loc=loc, ip=ip), + v7.ir_value(loc=loc, ip=ip), + rbits03.ir_value(loc=loc, ip=ip), + rbits47.ir_value(loc=loc, ip=ip), + ], + asm, + "=r,f,f,f,f,f,f,f,f,r,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + +@dsl_user_op +def mul_cvt_bf16x8_to_fp4x8( + v01: Uint32, + v23: Uint32, + v45: Uint32, + v67: Uint32, + coeff: Float32, + *, + loc=None, + ip=None, +) -> Uint32: + """Scale eight bf16 values by a bf16 coefficient and pack them into eight E2M1 values. + + This is mul_cvt_bf16_to_fp4_8x_round_to_nearest from ptx.cuh instruction for + instruction, which matters for exactly one input: `fma` against a zero addend flushes a -0 + product to +0 where a plain multiply keeps the sign, and E2M1 has a signed zero, so scaling + any other way would encode a -0 element as 0x8 instead of 0x0. Everything else about it + agrees with widening to f32 and multiplying, since a bf16 product has at most 16 mantissa + bits. + + The values arrive as four u32s of bf16 pairs, which is how they already sit in the fragment. + The coefficient arrives widened to f32 and is rounded back to bf16 here, which is exact, and + keeps the operands to the f32 and u32 registers the other PTX wrappers here use. + """ + asm = ( + "{\n" + ".reg.f32 zero;\n\t" + "mov.b32 zero, 0;\n\t" + ".reg.b16 c;\n\t" + "cvt.rn.bf16.f32 c, $5;\n\t" + ".reg.b16 h0, h1, h2, h3, h4, h5, h6, h7;\n\t" + "mov.b32 {h0, h1}, $1;\n\t" + "mov.b32 {h2, h3}, $2;\n\t" + "mov.b32 {h4, h5}, $3;\n\t" + "mov.b32 {h6, h7}, $4;\n\t" + ".reg.f32 v0, v1, v2, v3, v4, v5, v6, v7;\n\t" + "fma.rn.f32.bf16 v0, h0, c, zero;\n\t" + "fma.rn.f32.bf16 v1, h1, c, zero;\n\t" + "fma.rn.f32.bf16 v2, h2, c, zero;\n\t" + "fma.rn.f32.bf16 v3, h3, c, zero;\n\t" + "fma.rn.f32.bf16 v4, h4, c, zero;\n\t" + "fma.rn.f32.bf16 v5, h5, c, zero;\n\t" + "fma.rn.f32.bf16 v6, h6, c, zero;\n\t" + "fma.rn.f32.bf16 v7, h7, c, zero;\n\t" + ".reg.b8 f0, f1, f2, f3;\n\t" + "cvt.rn.satfinite.e2m1x2.f32 f0, v1, v0;\n\t" + "cvt.rn.satfinite.e2m1x2.f32 f1, v3, v2;\n\t" + "cvt.rn.satfinite.e2m1x2.f32 f2, v5, v4;\n\t" + "cvt.rn.satfinite.e2m1x2.f32 f3, v7, v6;\n\t" + "mov.b32 $0, {f0, f1, f2, f3};\n\t" + "}" + ) + return Uint32( + llvm.inline_asm( + T.i32(), + [ + v01.ir_value(loc=loc, ip=ip), + v23.ir_value(loc=loc, ip=ip), + v45.ir_value(loc=loc, ip=ip), + v67.ir_value(loc=loc, ip=ip), + coeff.ir_value(loc=loc, ip=ip), + ], + asm, + "=r,r,r,r,r,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + +@dsl_user_op +def mul_cvt_bf16x8_to_fp4x8_sr( + v01: Uint32, + v23: Uint32, + v45: Uint32, + v67: Uint32, + coeff: Float32, + rbits03: Uint32, + rbits47: Uint32, + *, + loc=None, + ip=None, +) -> Uint32: + """Stochastic-rounding variant of mul_cvt_bf16x8_to_fp4x8: the same bf16-coefficient fma + scaling, but converted through `cvt.rs.satfinite.e2m1x4.f32` with random bits, mirroring + mul_cvt_bf16_to_fp4_8x_stochastic_rounding in ptx.cuh.""" + asm = ( + "{\n" + ".reg.f32 zero;\n\t" + "mov.b32 zero, 0;\n\t" + ".reg.b16 c;\n\t" + "cvt.rn.bf16.f32 c, $5;\n\t" + ".reg.b16 h0, h1, h2, h3, h4, h5, h6, h7;\n\t" + "mov.b32 {h0, h1}, $1;\n\t" + "mov.b32 {h2, h3}, $2;\n\t" + "mov.b32 {h4, h5}, $3;\n\t" + "mov.b32 {h6, h7}, $4;\n\t" + ".reg.f32 v0, v1, v2, v3, v4, v5, v6, v7;\n\t" + "fma.rn.f32.bf16 v0, h0, c, zero;\n\t" + "fma.rn.f32.bf16 v1, h1, c, zero;\n\t" + "fma.rn.f32.bf16 v2, h2, c, zero;\n\t" + "fma.rn.f32.bf16 v3, h3, c, zero;\n\t" + "fma.rn.f32.bf16 v4, h4, c, zero;\n\t" + "fma.rn.f32.bf16 v5, h5, c, zero;\n\t" + "fma.rn.f32.bf16 v6, h6, c, zero;\n\t" + "fma.rn.f32.bf16 v7, h7, c, zero;\n\t" + ".reg.b16 b03, b47;\n\t" + "cvt.rs.satfinite.e2m1x4.f32 b03, {v3, v2, v1, v0}, $6;\n\t" + "cvt.rs.satfinite.e2m1x4.f32 b47, {v7, v6, v5, v4}, $7;\n\t" + "mov.b32 $0, {b03, b47};\n\t" + "}" + ) + return Uint32( + llvm.inline_asm( + T.i32(), + [ + v01.ir_value(loc=loc, ip=ip), + v23.ir_value(loc=loc, ip=ip), + v45.ir_value(loc=loc, ip=ip), + v67.ir_value(loc=loc, ip=ip), + coeff.ir_value(loc=loc, ip=ip), + rbits03.ir_value(loc=loc, ip=ip), + rbits47.ir_value(loc=loc, ip=ip), + ], + asm, + "=r,r,r,r,r,f,r,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + +# --- Register-resident (g2r) kernels: the numeric core --- +# +# A 16-element block is eight packed bf16 words; word h holds element 2h in its low half and +# 2h+1 in its high half. Block scale -> coefficient -> E2M1 follows TE's formulas +# (compute_decoding_scaling_factor, compute_nvfp4_scaling_coefficient), so the output is +# bit-identical to the CUDA kernels. + + +@cute.jit +def coefficient_of(code: Uint32, S_dec: Float32) -> Float32: + """compute_nvfp4_scaling_coefficient for scale byte `code`: + min(1 / (scale * S_dec), FLT_MAX) with S_dec = 1 / S_enc computed once by the caller.""" + return cute.arch.fmin(rcp_rn_f32(cvt_fp8e4m3_code_to_f32(code) * S_dec), Float32(FLOAT32_MAX)) + + +@cute.jit +def scale_bf16x2(v: Uint32, coeff2: Int64): + """Both bf16 halves of v widened to f32 and multiplied by the packed coefficient, (lo, hi).""" + return cute.arch.inline_ptx( + "{ .reg .b32 lo, hi; .reg .b64 p, r; shl.b32 lo, {$r0}, 16; " + "and.b32 hi, {$r0}, 0xffff0000; mov.b64 p, {lo, hi}; mul.f32x2 r, p, {$r1}; " + "mov.b64 {{$w0}, {$w1}}, r; }", + write_only_types=[Float32, Float32], + read_only_args=[v, coeff2], + ) + + +@cute.jit +def scale_column_bf16(a: Uint32, b: Uint32, coeff2: Int64, LOW: cutlass.Constexpr[bool]): + """Two rows of one column of a packed column pair, scaled: the low (LOW) or high halves of + the row words a and b, as (row a, row b).""" + if cutlass.const_expr(LOW): + widen = "shl.b32 ea, {$r0}, 16; shl.b32 eb, {$r1}, 16; " + else: + widen = "and.b32 ea, {$r0}, 0xffff0000; and.b32 eb, {$r1}, 0xffff0000; " + return cute.arch.inline_ptx( + "{ .reg .b32 ea, eb; .reg .b64 p, r; " + + widen + + "mov.b64 p, {ea, eb}; mul.f32x2 r, p, {$r2}; mov.b64 {{$w0}, {$w1}}, r; }", + write_only_types=[Float32, Float32], + read_only_args=[a, b, coeff2], + ) + + +@cute.jit +def block_scale_code(words, S_enc_div6: Float32) -> Uint32: + """E4M3 scale byte of one rowwise block (eight packed words) for the encode scale + S_enc_div6 = S_enc * (1 / 6).""" + m = abs_max_tree_bf16x2(words) + # Fold the two halves with the same unit: swap them and take one more abs-max. + folded = abs_max_tree_bf16x2([m, swap_halves_b32(m)]) & Uint32(0x7FFF0000) + return cvt_f32_to_fp8e4m3_code(folded.bitcast(Float32) * S_enc_div6) + + +@cute.jit +def pack_half_block(floats, r0: Uint32, r1: Uint32, STOCHASTIC: cutlass.Constexpr[bool]) -> Uint32: + """Pack 8 scaled elements into one word of E2M1 nibbles, element 0 in the lowest; under + stochastic rounding r0 and r1 are the random words of elements 0..3 and 4..7.""" + if cutlass.const_expr(STOCHASTIC): + return cvt_f32x8_to_fp4x8_sr(*floats, r0, r1) + return cvt_f32x8_to_fp4x8(*floats) + + +@cute.jit +def pack_row_block(words, coeff2: Int64, rbits, STOCHASTIC: cutlass.Constexpr[bool]): + """Scale and pack one rowwise block (eight packed words) into two words; rbits is the block's + draw under stochastic rounding. Half a block at a time keeps 8 floats live, not 16.""" + halves = [] + for half in cutlass.range_constexpr(2): + floats = [] + for h in cutlass.range_constexpr(4 * half, 4 * half + 4): + floats.extend(scale_bf16x2(words[h], coeff2)) + r0 = rbits[2 * half] if cutlass.const_expr(STOCHASTIC) else None + r1 = rbits[2 * half + 1] if cutlass.const_expr(STOCHASTIC) else None + halves.append(pack_half_block(floats, r0, r1, STOCHASTIC)) + return halves[0], halves[1] + + +@cute.jit +def pack_column_block( + rows, + coeff2: Int64, + rbits, + LOW: cutlass.Constexpr[bool], + STOCHASTIC: cutlass.Constexpr[bool], +): + """Columnwise twin of pack_row_block: rows[t] is row t of a packed column pair, and LOW + selects its even column.""" + halves = [] + for half in cutlass.range_constexpr(2): + floats = [] + for h in cutlass.range_constexpr(4 * half, 4 * half + 4): + floats.extend(scale_column_bf16(rows[2 * h], rows[2 * h + 1], coeff2, LOW)) + r0 = rbits[2 * half] if cutlass.const_expr(STOCHASTIC) else None + r1 = rbits[2 * half + 1] if cutlass.const_expr(STOCHASTIC) else None + halves.append(pack_half_block(floats, r0, r1, STOCHASTIC)) + return halves[0], halves[1] diff --git a/transformer_engine/common/cast/dispatch/quantize.cuh b/transformer_engine/common/cast/dispatch/quantize.cuh index 13c3694d040..478d875adf8 100644 --- a/transformer_engine/common/cast/dispatch/quantize.cuh +++ b/transformer_engine/common/cast/dispatch/quantize.cuh @@ -33,6 +33,7 @@ #ifdef NVTE_WITH_CUTEDSL #include "../mxfp8/quantize_mxfp8_cutedsl.cuh" +#include "../nvfp4/quantize_transpose_nvfp4_cutedsl.cuh" #endif namespace transformer_engine { @@ -167,11 +168,27 @@ void quantize_fwd_helper(const NVTETensor input, NVTETensor output, } } else if (use_optimized_kernel) { if (quant_config_cpp.nvfp4_2d_quantization) { - nvfp4::quantize_transpose( - *input_tensor, noop_tensor, output_tensor, &quant_config_cpp, stream); + bool quantized_with_cutedsl = false; +#ifdef NVTE_WITH_CUTEDSL + quantized_with_cutedsl = + cutedsl_backend::nvfp4_quantize_transpose_cutedsl( + *input_tensor, noop_tensor, output_tensor, &quant_config_cpp, stream); +#endif + if (!quantized_with_cutedsl) { + nvfp4::quantize_transpose( + *input_tensor, noop_tensor, output_tensor, &quant_config_cpp, stream); + } } else { - nvfp4::quantize_transpose( - *input_tensor, noop_tensor, output_tensor, &quant_config_cpp, stream); + bool quantized_with_cutedsl = false; +#ifdef NVTE_WITH_CUTEDSL + quantized_with_cutedsl = + cutedsl_backend::nvfp4_quantize_transpose_cutedsl( + *input_tensor, noop_tensor, output_tensor, &quant_config_cpp, stream); +#endif + if (!quantized_with_cutedsl) { + nvfp4::quantize_transpose( + *input_tensor, noop_tensor, output_tensor, &quant_config_cpp, stream); + } } } else { auto &global_amax = (output_tensor->amax.dptr != nullptr) ? output_tensor->amax @@ -356,11 +373,27 @@ void quantize_bwd_helper(const NVTETensor grad, const NVTETensor input, NVTETens } } else if (use_optimized_kernel) { if (quant_config_cpp.nvfp4_2d_quantization) { - nvfp4::quantize_transpose( - *grad_tensor, noop_tensor, output_tensor, &quant_config_cpp, stream); + bool quantized_with_cutedsl = false; +#ifdef NVTE_WITH_CUTEDSL + quantized_with_cutedsl = + cutedsl_backend::nvfp4_quantize_transpose_cutedsl( + *grad_tensor, noop_tensor, output_tensor, &quant_config_cpp, stream); +#endif + if (!quantized_with_cutedsl) { + nvfp4::quantize_transpose( + *grad_tensor, noop_tensor, output_tensor, &quant_config_cpp, stream); + } } else { - nvfp4::quantize_transpose( - *grad_tensor, noop_tensor, output_tensor, &quant_config_cpp, stream); + bool quantized_with_cutedsl = false; +#ifdef NVTE_WITH_CUTEDSL + quantized_with_cutedsl = + cutedsl_backend::nvfp4_quantize_transpose_cutedsl( + *grad_tensor, noop_tensor, output_tensor, &quant_config_cpp, stream); +#endif + if (!quantized_with_cutedsl) { + nvfp4::quantize_transpose( + *grad_tensor, noop_tensor, output_tensor, &quant_config_cpp, stream); + } } } else { auto &global_amax = (output_tensor->amax.dptr != nullptr) ? output_tensor->amax diff --git a/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4_cutedsl.cuh b/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4_cutedsl.cuh new file mode 100644 index 00000000000..687ba3cb183 --- /dev/null +++ b/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4_cutedsl.cuh @@ -0,0 +1,221 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#ifndef TRANSFORMER_ENGINE_COMMON_CAST_NVFP4_QUANTIZE_TRANSPOSE_NVFP4_CUTEDSL_CUH_ +#define TRANSFORMER_ENGINE_COMMON_CAST_NVFP4_QUANTIZE_TRANSPOSE_NVFP4_CUTEDSL_CUH_ + +#include +#include + +#include +#include +#include +#include + +#include "../../common.h" +#include "../../tvm_ffi_bridge.h" + +namespace transformer_engine { +namespace cutedsl_backend { +#if FP4_TYPE_SUPPORTED + +// DLTensorWrapper and TVMFFICentral live in transformer_engine::tvm_ffi_bridge. +using namespace tvm_ffi_bridge; + +// Instantiation parameters of the CuTeDSL kernel +struct NVFP4QuantConfig { + static constexpr const char *kEntrypointName = "get_nvfp4_quantization_function"; + + bool use_stochastic_rounding; // If stochastic rounding is applied during quantization + bool use_fast_math; // If fast math approximations are used during quantization + bool row_scaled_nvfp4; // If scales are computed per row (row-scaled NVFP4) + bool return_transpose; // If columnwise (transposed) output is produced + uint32_t sm_arch = static_cast(cuda::sm_arch()); + + // Bit layout: flags [3:0], SM architecture [12:4] (9 used/reserved). + uint32_t to_id() const { + NVTE_CHECK(sm_arch < 512, "SM architecture no longer fits in the 9 bits to_id() gives it."); + return static_cast(use_stochastic_rounding) | + (static_cast(use_fast_math) << 1) | + (static_cast(row_scaled_nvfp4) << 2) | + (static_cast(return_transpose) << 3) | (sm_arch << 4); + } + + std::optional get_kernel() const { + static TVMFFIConfigCache &cache = TVMFFIConfigCache::create(); + return cache.get_or_load(*this); + } + + // Globally unique TVM-FFI registry key used when the CuTeDSL function is + // compiled and registered on a cache miss. + std::string to_key() const { + std::string key; + key.reserve(32); + key.append("cutedsl_nvfp4_sm") + .append(std::to_string(sm_arch)) + .append("_") + .append(use_stochastic_rounding ? "1" : "0") + .append("_") + .append(use_fast_math ? "1" : "0") + .append("_") + .append(row_scaled_nvfp4 ? "1" : "0") + .append("_") + .append(return_transpose ? "1" : "0"); + 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), use_stochastic_rounding, + use_fast_math, row_scaled_nvfp4, return_transpose); + return result.try_cast().value_or(false); + } +}; + +// Dispatch to Python side via TVM FFI +inline bool nvfp4_quantize_transpose_cutedsl(const NVFP4QuantConfig &config, const Tensor &input, + const Tensor *noop, Tensor *output, + const QuantizationConfig *quant_config, + cudaStream_t stream) { + std::optional nvfp4_quant_func_opt = config.get_kernel(); + if (!nvfp4_quant_func_opt.has_value()) { + return false; + } + + // Data tensors auto-flatten to 2D (DLTensorWrapper's default); the amaxes and the RNG state are + // rank-1 in the compiled ABI. + const int32_t device_index = transformer_engine::cuda::current_device(); + DLTensorWrapper mX(input.data, true, device_index); + DLTensorWrapper mO_row(output->data, true, device_index); + DLTensorWrapper mS_row(output->scale_inv, true, device_index); + DLTensorWrapper mAmaxRow(output->amax, false, device_index); + + DLTensorWrapper mO_col, mS_col, mAmaxCol; + if (config.return_transpose) { + // The transposed output of a rank > 2 input is shaped {cols, d0, d1, ...}, which the default + // flattening would collapse to (cols * d0, ...) rather than (cols, rows). + const auto [rows, cols] = input.flat_2d_dims(); + const SimpleTensor col_data(output->columnwise_data.dptr, {cols, rows}, + output->columnwise_data.dtype); + mO_col = DLTensorWrapper(col_data, false, device_index); + mS_col = DLTensorWrapper(output->columnwise_scale_inv, true, device_index); + mAmaxCol = DLTensorWrapper(output->columnwise_amax, false, device_index); + } + + // Optional Philox RNG state ({seed, offset}); required (and only used) for SR. + DLTensorWrapper mRngState; + if (config.use_stochastic_rounding) { + const NVTETensor rng_state_tensor = + (quant_config != nullptr) ? quant_config->rng_state : nullptr; + NVTE_CHECK(rng_state_tensor != nullptr, + "Stochastic rounding requires an RNG state tensor to be provided."); + Tensor &rng_state_te_tensor = *convertNVTETensor(rng_state_tensor); + mRngState = DLTensorWrapper(rng_state_te_tensor.data, false, device_index); + } + + // noop passed as ptr because the null-check must be done at runtime on device + void *noop_ptr = (noop != nullptr) ? noop->data.dptr : nullptr; + // noop and stream are tvm-ffi opaque "handles"; pass them as void*. + (*nvfp4_quant_func_opt)(&mX, &mO_row, &mS_row, &mO_col, &mS_col, &mAmaxRow, &mAmaxCol, noop_ptr, + &mRngState, static_cast(stream)); + return true; +} +#endif // FP4_TYPE_SUPPORTED + +// CuTeDSL counterpart of nvfp4::quantize_transpose. Returns false to fall back to the CUDA kernel. +template +bool nvfp4_quantize_transpose_cutedsl(const Tensor &input, const Tensor *noop, Tensor *output, + const QuantizationConfig *quant_config, cudaStream_t stream) { +#if FP4_TYPE_SUPPORTED + if (!TVMFFICentral::getInstance().get_cutedsl_backend_enabled()) { + maybe_warn_cutedsl_not_chosen("the CuTeDSL backend is disabled."); + return false; + } + // Currently, only the CuTeDSL counterpart of quantize_transpose_tuned_1D is supported + if constexpr (use_2d_quantization) { + maybe_warn_cutedsl_not_chosen("2D NVFP4 quantization is not supported."); + return false; + } else { + if (input.dtype() != DType::kBFloat16) { + maybe_warn_cutedsl_not_chosen("NVFP4 quantization supports only BF16 input."); + return false; + } + if (output->with_gemm_swizzled_scales) { + maybe_warn_cutedsl_not_chosen("NVFP4 quantization does not produce GEMM-swizzled scales."); + return false; + } + if (output->has_data() && output->scale_inv.dtype != DType::kFloat8E4M3) { + maybe_warn_cutedsl_not_chosen("NVFP4 quantization supports only E4M3 scales."); + return false; + } + const bool return_transpose = output->has_columnwise_data(); + // The compiled ABI takes the amaxes as tensors, not optionals. + if (output->amax.dptr == nullptr || + (return_transpose && output->columnwise_amax.dptr == nullptr)) { + maybe_warn_cutedsl_not_chosen("NVFP4 quantization requires the amax tensors."); + return false; + } + const auto [rows, cols] = input.flat_2d_dims(); + if (rows % 32 != 0 || cols % 32 != 0) { + maybe_warn_cutedsl_not_chosen("NVFP4 quantization needs rows and columns divisible by 32."); + return false; + } + + const bool use_stochastic_rounding = + quant_config != nullptr ? quant_config->stochastic_rounding : false; + const bool use_fast_math = quant_config != nullptr ? quant_config->use_fast_math : false; + const bool row_scaled_nvfp4 = output->row_scaled_nvfp4; + const NVFP4QuantConfig config{ + use_stochastic_rounding, + use_fast_math, + row_scaled_nvfp4, + return_transpose, + }; + + // Checks mirroring the CUDA version + checkCuDriverContext(stream); + CheckNoopTensor(*noop, "cast_noop"); + CheckInputTensor(input, "input"); + CheckOutputTensor(*output, "output", false); + + NVTE_CHECK(input.has_data(), "Cannot quantize tensor without rowwise data."); + NVTE_CHECK(output->has_data(), "NVFP4 output tensor must be allocated."); + NVTE_CHECK(is_fp4_dtype(output->data.dtype), "Output must have FP4 type."); + NVTE_CHECK(output->scale_inv.dptr != nullptr, "Scaling tensor must be allocated"); + NVTE_CHECK(!row_scaled_nvfp4 || !output->has_columnwise_data(), + "Row-scaled NVFP4 quantization does not produce columnwise output."); + + if (return_transpose) { + NVTE_CHECK(is_fp4_dtype(output->columnwise_data.dtype), + "Transposed output must have FP4 type."); + NVTE_CHECK(output->columnwise_scale_inv.dptr != nullptr, + "Transposed scaling tensor must be allocated"); + } + + const NVTETensor rng_state_tensor = + (quant_config != nullptr) ? quant_config->rng_state : nullptr; + if (rng_state_tensor != nullptr) { + Tensor &rng_state_te_tensor = *convertNVTETensor(rng_state_tensor); + NVTE_CHECK(rng_state_te_tensor.dtype() == DType::kInt64, + "RNG state should contain 2 64-bit values."); + NVTE_CHECK(rng_state_te_tensor.data.shape == Shape{2}, + "Shape of the RNG state should be [2], but got ", rng_state_te_tensor.data.shape); + } + + return nvfp4_quantize_transpose_cutedsl(config, input, noop, output, quant_config, stream); + } +#else + NVTE_ERROR("FP4 support requires CUDA 12.8+, but compile-time CUDA version is ", CUDA_VERSION); +#endif // FP4_TYPE_SUPPORTED +} + +} // namespace cutedsl_backend +} // namespace transformer_engine + +#endif // TRANSFORMER_ENGINE_COMMON_CAST_NVFP4_QUANTIZE_TRANSPOSE_NVFP4_CUTEDSL_CUH_