From 54fd20582c72088f32d1623b0cde64697252baac Mon Sep 17 00:00:00 2001 From: Jan Bielak Date: Mon, 2 Nov 2026 14:07:30 -0800 Subject: [PATCH 1/5] Port the CuTeDSL NVFP4 quantize backend to the current CuTeDSL infrastructure Bring the NVFP4 tuned-1D CuTeDSL backend from #3223 onto main, whose CuTeDSL plumbing has changed since that PR branched: - Build and dispatch the backend only with NVTE_WITH_CUTEDSL, and enable it with NVTE_ENABLE_CUTEDSL_BACKEND like MXFP8. - Give NVFP4QuantConfig a to_id() and SM architecture so it uses TVMFFIConfigCache, and pass the device index to DLTensorWrapper, whose FP4 packing main already provides. - Pass the transposed output as (cols, rows): for a rank > 2 input TE shapes it {cols, d0, d1, ...}, which the default 2D flattening of DLTensorWrapper collapses to the wrong extents. - Fall back to CUDA, with a reason under NVTE_WARN_IF_CUTEDSL_BACKEND_NOT_CHOSEN, for configurations the compiled ABI cannot take: non-BF16 input, GEMM-swizzled or non-E4M3 scales, missing amax tensors, and flat dims not divisible by 32. - Keep the NVFP4-only helpers in CuTeDSL/cast/nvfp4/utils.py and philox_rng.py instead of the shared utils.py, using main's inline_ptx style, and register the native TVM-FFI function as MXFP8 does. - Start the packed abs-max from a zero held in a register: as a literal it became an immediate operand, which ptxas rejects for max.xorsign.abs.bf16x2. - Switch the test to main's backend switches and config keys. Signed-off-by: Jan Bielak Signed-off-by: Teddy Do --- .../nvfp4/test_nvfp4_cutedsl_backend.py | 740 +++++++++++++++ transformer_engine/common/CuTeDSL/__init__.py | 6 + .../common/CuTeDSL/cast/__init__.py | 1 + .../common/CuTeDSL/cast/nvfp4/__init__.py | 7 + .../cast/nvfp4/quantize_transpose/__init__.py | 7 + .../cast/nvfp4/quantize_transpose/config.py | 33 + .../cast/nvfp4/quantize_transpose/kernel.py | 844 ++++++++++++++++++ .../cast/nvfp4/quantize_transpose/launcher.py | 195 ++++ .../common/CuTeDSL/cast/nvfp4/utils.py | 384 ++++++++ .../common/CuTeDSL/philox_rng.py | 96 ++ .../common/cast/dispatch/quantize.cuh | 49 +- .../quantize_transpose_nvfp4_cutedsl.cuh | 221 +++++ 12 files changed, 2575 insertions(+), 8 deletions(-) create mode 100644 tests/pytorch/nvfp4/test_nvfp4_cutedsl_backend.py create mode 100644 transformer_engine/common/CuTeDSL/cast/nvfp4/__init__.py create mode 100644 transformer_engine/common/CuTeDSL/cast/nvfp4/quantize_transpose/__init__.py create mode 100644 transformer_engine/common/CuTeDSL/cast/nvfp4/quantize_transpose/config.py create mode 100644 transformer_engine/common/CuTeDSL/cast/nvfp4/quantize_transpose/kernel.py create mode 100644 transformer_engine/common/CuTeDSL/cast/nvfp4/quantize_transpose/launcher.py create mode 100644 transformer_engine/common/CuTeDSL/cast/nvfp4/utils.py create mode 100644 transformer_engine/common/CuTeDSL/philox_rng.py create mode 100644 transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4_cutedsl.cuh 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..ab57771cb33 --- /dev/null +++ b/tests/pytorch/nvfp4/test_nvfp4_cutedsl_backend.py @@ -0,0 +1,740 @@ +# 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-transpose kernel. + +Every case quantizes the same input twice, once with the CuTeDSL backend disabled and once with +it enabled, and requires the two results to agree byte for byte. The CUDA C++ kernels are the +reference: `quantize_transpose_nvfp4_tuned_1D.cuh` is what the CuTeDSL kernel reimplements, and +its numerics are in turn pinned against a pure-PyTorch reference by test_nvfp4_quantize_exact.py, +so matching it bit for bit is the whole contract. + +Configurations the CuTeDSL kernel does not own are kept in the matrix, but it is worth being +precise about how little they check: both runs execute the same CUDA kernel, so the comparison +only says that turning the backend on is inert. That is not nothing, since the dispatcher walks a +wall of NVTE_CHECKs and builds a config key before it declines and the Python entrypoint compiles +without being allowed to let anything escape, but it says nothing about the CuTeDSL kernel itself. +What pins the kernel's coverage down is `test_kernel_coverage`, which asks the entrypoint +directly; a config moves between the two categories by editing `kernel_implements`. + +Which backend ran is therefore never inferred, it is observed: every kernel the CuTeDSL backend +can serve is wrapped in a counting proxy at import (see `_install_call_counters`), and every case +states how many times the dispatcher must call into one. A comparison that silently ran the CUDA +kernel twice would otherwise pass while testing nothing, which is the one failure this file cannot +afford -- and it is not hypothetical: a kernel that fails to compile, a dispatcher check that +rejects the config, and a backend switch that does not take effect all produce exactly that. +Counting catches all three, including the last, since the run that is meant to be the CUDA +reference asserts that no CuTeDSL kernel ran either. + +Stochastic rounding is deliberately not compared byte for byte, even though the CuTeDSL kernel +implements it (and, as it happens, currently matches the CUDA kernel bit for bit, having +replicated its Philox seeding, thread arrangement and draw order): which random bits an element +gets follows from the work decomposition, and the CUDA kernel seeds Philox from the chunk +coordinate a CTA starts on without reseeding when it steals a chunk through cluster launch +control, so two implementations, or even two launches, may legitimately disagree without either +being wrong. test_stochastic_rounding checks the property that defines the rounding instead, +and test_nvfp4_sr_quantize.py checks the rounding itself on the CUDA side. + +RHT is out of scope: it is a pre-transform with kernels and tests of its own, and it splits the +quantization into per-direction calls, so "which backend ran" stops being one fact per case. +""" + +# Optional CuTeDSL dependencies must be checked before importing Transformer Engine. +# pylint: disable=wrong-import-position + +import collections +import ctypes +import itertools +import math +import os +from contextlib import contextmanager +from dataclasses import dataclass, replace + +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 +backend_available = recipe_available and cutedsl_built and cutedsl_enabled +pytestmark = pytest.mark.skipif(not backend_available, reason=skip_reason) + +# Elements sharing one E4M3 block scale. +NVFP4_BLOCK_SIZE = 16 +# Flat-dim divisibility the tuned-1D path requires (16B alignment for its TMA descriptors). +TUNED_1D_ALIGNMENT = 32 +# Name the CuTeDSL NVFP4 compile-on-demand entrypoint is registered under (kEntrypointName in +# quantize_transpose_nvfp4_cutedsl.cuh). +ENTRYPOINT = "get_nvfp4_quantization_function" +# Filled into the output buffers before every quantization: anything a kernel fails to write then +# shows up as a difference instead of as whatever the caching allocator recycled into the block. +POISON_BYTE = 0xA5 +SEED = 1234 + + +@dataclass(frozen=True) +class Case: + """An NVFP4Quantizer configuration together with the input it is exercised with.""" + + id: str + shape: tuple = (256, 384) + dtype: torch.dtype = torch.bfloat16 + rowwise: bool = True + columnwise: bool = False + row_scaled: bool = False + two_d: bool = False + stochastic_rounding: bool = False + fast_math: bool = False + use_4over6: bool = False + + +# The configuration the CuTeDSL kernel currently covers; other cases are derived from it. +ROWWISE = Case(id="rowwise") + +get_case_id = lambda case: case.id +get_shape_id = lambda shape: "x".join(str(d) for d in shape) +get_dtype_id = lambda dtype: str(dtype).removeprefix("torch.") + + +def flat_dims(shape): + """The 2D dims the kernels see; a rank > 2 input is flattened over its leading dims.""" + return math.prod(shape[:-1]), shape[-1] + + +def ceil_div(numerator, denominator): + return -(-numerator // denominator) + + +def make_quantizer(case): + return NVFP4Quantizer( + rowwise=case.rowwise, + columnwise=case.columnwise, + with_2d_quantization=case.two_d, + stochastic_rounding=case.stochastic_rounding, + row_scaled_nvfp4=case.row_scaled, + nvfp4_use_4over6=case.use_4over6, + ) + + +def cutedsl_flags(case): + """The four booleans an NVFP4QuantConfig is made of, as the dispatcher reads them off the + quantization config and the output tensor: a transposed output is requested by allocating + columnwise data, and fast math comes from NVTE_USE_FAST_MATH.""" + return (case.stochastic_rounding, case.fast_math, case.row_scaled, case.columnwise) + + +def cutedsl_key(flags): + """Mirror of NVFP4QuantConfig::to_key: the name the compiled kernel for these flags is + registered under, and the key the C++ side caches its lookup result against.""" + major, minor = device_compute_capability() + return f"cutedsl_nvfp4_sm{10 * major + minor}_" + "_".join( + "1" if flag else "0" for flag in flags + ) + + +def kernel_implements(flags): + """Whether the CuTeDSL kernel compiles a kernel for these flags. Every feature of the CUDA + tuned-1D kernel is implemented; the one refusal left is the flag combination that is not a + real configuration (NVFP4QuantizeConfig raises for it): row-scaled quantization has no + transposed output.""" + _, _, row_scaled, return_transpose = flags + return not (row_scaled and return_transpose) + + +def dispatcher_offers(case): + """Whether the C++ dispatcher hands this config to the CuTeDSL backend at all. Mirrors the + NVTE_NVFP4_1D_SCALING branch of dispatch/quantize.cuh and the 2D early-out in + quantize_transpose_nvfp4_cutedsl.cuh: 4over6 has kernels of its own, the optimized (tuned-1D) + path is taken only for bf16 with 32-divisible flat dims and an allocated rowwise output, and + the CuTeDSL entry for 2D scaling declines before it even builds a config key.""" + rows, cols = flat_dims(case.shape) + return ( + not case.use_4over6 + and not case.two_d + and case.dtype is torch.bfloat16 + and rows % TUNED_1D_ALIGNMENT == 0 + and cols % TUNED_1D_ALIGNMENT == 0 + and case.rowwise + ) + + +def expect_cutedsl(case): + """Whether the CuTeDSL kernel is the one that produces this case's output.""" + return dispatcher_offers(case) and kernel_implements(cutedsl_flags(case)) + + +def set_cutedsl_backend(enabled): + CORE_LIB.nvte_set_cutedsl_backend(1 if enabled else 0) + + +# --- Observing which backend ran --- + +# How many times the C++ dispatcher has called into a CuTeDSL kernel, per config key. +_cutedsl_calls_by_key = collections.Counter() + + +def _install_call_counters(): + """Replace every kernel the CuTeDSL backend can serve with a proxy that counts its calls. + + A registered key says a kernel exists, not that the dispatcher reached it, and a key another + case registered answers for every case after it; counting invocations is what makes "the + CuTeDSL backend produced this output" a fact rather than an inference. The proxies have to go + in before the first quantization of the process: the C++ side resolves each key once through + `TVMFFIConfigCache::get_or_load` and caches the handle for good, so a proxy installed after + that is never seen. Compiling the kernels here rather than letting the dispatcher trigger it + is what makes that possible, and costs one compilation per supported config at import. + """ + from transformer_engine.common.CuTeDSL.cast.nvfp4.quantize_transpose import ( + get_nvfp4_quantization_function, + ) + + def counting_proxy(key, kernel): + def proxy(*args): + _cutedsl_calls_by_key[key] += 1 + return kernel(*args) + + return proxy + + for flags in itertools.product([False, True], repeat=4): + if not kernel_implements(flags): + continue + key = cutedsl_key(flags) + # A config kernel_implements() claims is served but that does not compile is left alone: + # test_kernel_coverage is where that contradiction is reported. + if not get_nvfp4_quantization_function(key, *flags): + continue + tvm_ffi.register_global_func( + key, counting_proxy(key, tvm_ffi.get_global_func(key)), override=True + ) + + +@contextmanager +def cutedsl_calls(expected, tag, key=None): + """Require exactly `expected` CuTeDSL kernel invocations inside the block, all of them of the + kernel compiled for `key` when one is named. + + Which key ran matters as much as whether any did: the flags are what a compiled kernel is, so + a case served by the kernel for another flag combination -- exact math answering for fast + math, say -- is testing something other than what it says. + """ + before = collections.Counter(_cutedsl_calls_by_key) + yield + made = _cutedsl_calls_by_key - before + actual = sum(made.values()) + if actual != expected: + if expected == 0: + raise AssertionError( + f"{tag}: the CuTeDSL backend ran {actual} time(s) ({dict(made)}) where it was " + "expected to stay out of the way, so this case did not exercise the backend it " + "names" + ) + raise AssertionError( + f"{tag}: the CuTeDSL backend ran {actual} time(s) instead of {expected}, so the CUDA " + "kernel served a config this case is about. Either the dispatcher declined it, or a " + "quantization earlier in this pytest session made the C++ side cache the kernel " + "handle before the counting proxies were installed" + ) + if key is not None and expected > 0 and set(made) != {key}: + raise AssertionError( + f"{tag}: the CuTeDSL kernels that ran were {dict(made)}, not {expected} call(s) of " + f"{key}, so this case was served by a kernel compiled for another config" + ) + + +if backend_available: + _install_call_counters() + + +@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")) + + +@contextmanager +def fast_math(enabled): + """NVTE_USE_FAST_MATH is what puts use_fast_math in the quantization config, and TE reads it + afresh on every quantize call, so setting it here is enough to switch the config.""" + name = "NVTE_USE_FAST_MATH" + previous = os.environ.get(name) + os.environ[name] = "1" if enabled else "0" + try: + yield + finally: + if previous is None: + os.environ.pop(name, None) + else: + os.environ[name] = previous + + +def make_input(case, seed=SEED): + """Uniform noise with a per-scaling-block power-of-two magnitude, so that neighbouring blocks + land on different E4M3 block scales and the elements at the top of a block saturate the E2M1 + range.""" + rows, cols = flat_dims(case.shape) + generator = torch.Generator(device="cuda").manual_seed(seed) + values = torch.empty((rows, cols), dtype=torch.float32, device="cuda").uniform_( + -6.0, 6.0, generator=generator + ) + if cols % NVFP4_BLOCK_SIZE == 0: + blocks = cols // NVFP4_BLOCK_SIZE + exponents = torch.randint( + -6, 7, (rows, blocks, 1), device="cuda", generator=generator + ).float() + values = (values.view(rows, blocks, NVFP4_BLOCK_SIZE) * torch.exp2(exponents)).view( + rows, cols + ) + return values.to(case.dtype).view(case.shape) + + +def make_input_with_zero_blocks(case, negative=False): + """Whole scaling blocks of zeros, which is where the block amax is zero and both kernels divide + by it and clamp the coefficient. Random data never lands there.""" + x = make_input(case) + x.view(-1, NVFP4_BLOCK_SIZE)[::3].fill_(-0.0 if negative else 0.0) + return x + + +def make_output(case): + """An output tensor of this config's shape, poisoned so that every byte the comparison looks + at has to have been written by the kernel under test.""" + out = make_quantizer(case).make_empty(case.shape, dtype=case.dtype, device="cuda") + for buffer in ( + out._rowwise_data, + out._columnwise_data, + out._rowwise_scale_inv, + out._columnwise_scale_inv, + ): + if buffer is not None: + buffer.view(torch.uint8).fill_(POISON_BYTE) + return out + + +def quantize_into(out, x, noop_flag=None): + """Quantize through the same C++ dispatch a plain quantizer call takes, into a preallocated + output so that its buffers can be poisoned and a cast-noop flag can be passed.""" + # The Philox state for stochastic rounding is drawn from the default CUDA generator when the + # quantize op runs, so seeding here is what makes the two backends see the same randomness. + torch.cuda.manual_seed(SEED) + return out.quantize_(x, noop_flag=noop_flag) + + +def quantize(case, x): + return quantize_into(make_output(case), x) + + +def quantized_parts(out, case, with_amax=True): + """The bytes of a quantized tensor that both backends must agree on. The scale buffers are + padded to (round_up(rows, 128), round_up(ceil(cols / 16), 4)) and no kernel writes the + padding, so only the valid region is taken.""" + rows, cols = flat_dims(case.shape) + parts = {} + if case.rowwise: + parts["rowwise data"] = out._rowwise_data.view(torch.uint8).clone() + parts["rowwise scales"] = out._rowwise_scale_inv.view(torch.uint8)[ + :rows, : ceil_div(cols, NVFP4_BLOCK_SIZE) + ].clone() + if case.columnwise: + parts["columnwise data"] = out._columnwise_data.view(torch.uint8).clone() + parts["columnwise scales"] = out._columnwise_scale_inv.view(torch.uint8)[ + :cols, : ceil_div(rows, NVFP4_BLOCK_SIZE) + ].clone() + if with_amax: + # The amax comes from a pass of its own rather than from the quantize kernel, but it is + # what fixes the global encode scale, so comparing it separates a wrong amax from wrong + # quantization when a case fails. + if out._amax_rowwise is not None: + parts["rowwise amax"] = out._amax_rowwise.clone() + if out._amax_columnwise is not None: + parts["columnwise amax"] = out._amax_columnwise.clone() + return parts + + +def assert_parts_equal(expected, actual, tag): + for name, expected_bytes in expected.items(): + assert torch.equal(actual[name], expected_bytes), f"{tag}: {name} does not match" + + +def relative_dequantize_error(out, x): + """How far the dequantized output sits from the input, relative to the input's norm.""" + dequantized = out.dequantize(dtype=torch.bfloat16).float() + reference = x.float() + error = torch.linalg.vector_norm(dequantized - reference) / torch.linalg.vector_norm(reference) + return error.item() + + +def run_case(case, x=None): + """Assert the CuTeDSL and CUDA backends produce bit-identical output for this config, and + that the case exercised the backend it was meant to.""" + if x is None: + x = make_input(case) + + with fast_math(case.fast_math): + set_cutedsl_backend(False) + # The reference run has to be CUDA, which is only true if the backend switch took effect. + with cutedsl_calls(0, f"{case.id} (CUDA reference run)"): + cuda_parts = quantized_parts(quantize(case, x), case) + + key = cutedsl_key(cutedsl_flags(case)) + set_cutedsl_backend(True) + try: + with cutedsl_calls(1 if expect_cutedsl(case) else 0, case.id, key): + cutedsl_parts = quantized_parts(quantize(case, x), case) + finally: + set_cutedsl_backend(False) + + registered = tvm_ffi.get_global_func(key, allow_missing=True) is not None + if expect_cutedsl(case): + # Redundant with the call count above, but it separates a kernel that failed to compile + # from a dispatcher that declined a kernel it has. + assert registered, ( + f"{case.id}: nothing is registered under {key}, so the CuTeDSL backend fell back to " + "CUDA and this case compared CUDA against itself" + ) + elif not kernel_implements(cutedsl_flags(case)): + assert not registered, ( + f"{case.id}: a kernel is registered under {key}, so the CuTeDSL backend now covers a " + "config this test expects it to decline; update kernel_implements()" + ) + # A config the dispatcher never offers (2D scaling, 4over6, non-bf16, flat dims not divisible + # by 32) shares its key with configs it does offer, so registration says nothing about it. + # What such a case checks is that turning the backend on leaves the CUDA result alone. + + assert_parts_equal(cuda_parts, cutedsl_parts, case.id) + + +# --- The kernel's own account of what it supports --- + + +@pytest.mark.parametrize( + "flags", + list(itertools.product([False, True], repeat=4)), + ids=lambda flags: cutedsl_key(flags).removeprefix("cutedsl_nvfp4_"), +) +def test_kernel_coverage(flags): + """Ask the compile-on-demand entrypoint directly which flag combinations it compiles a kernel + for. This is the one place the kernel's coverage is asserted, so a config that starts or stops + being supported is reported here instead of showing up as a puzzling fallback elsewhere. + + The probe registers under a name of its own so that it cannot make a later case's + registration check pass by proxy. + """ + # Imports cleanly only with the CuTeDSL stack present, which the module-level skip guarantees. + from transformer_engine.common.CuTeDSL.cast.nvfp4.quantize_transpose import ( + get_nvfp4_quantization_function, + ) + + assert tvm_ffi.get_global_func(ENTRYPOINT, allow_missing=True) is not None, ( + f"{ENTRYPOINT} is not in the tvm-ffi registry, so the C++ dispatcher has no way to ask for " + "a CuTeDSL NVFP4 kernel at all" + ) + + stochastic_rounding, use_fast_math, row_scaled, return_transpose = flags + supported = get_nvfp4_quantization_function( + "test_probe_" + cutedsl_key(flags), + stochastic_rounding, + use_fast_math, + row_scaled, + return_transpose, + ) + assert bool(supported) == kernel_implements(flags), ( + f"the CuTeDSL NVFP4 backend reports supported={bool(supported)} for {cutedsl_key(flags)}, " + "which contradicts kernel_implements()" + ) + + +# --- Shapes, over the config the kernel serves --- + +# Flat dims must be multiples of 32 for the tuned-1D path to be chosen at all; within that, the +# interesting axis is how the shape sits relative to the kernel's 64x64 chunk, since everything +# that does not tile exactly is covered by predication. +SHAPES = [ + (32, 32), # one partial chunk, in both dimensions + (32, 1024), # a single row of chunks, partial + (1024, 32), # a single column of chunks, partial + (96, 160), # partial chunks along both edges + (128, 128), + (256, 1024), + (512, 512), + (8, 32, 1024), # rank > 2, flattened to (256, 1024) + (8192, 7168), # a shape from a real model +] + + +@pytest.mark.parametrize("shape", SHAPES, ids=get_shape_id) +@pytest.mark.parametrize("use_fast_math", [False, True], ids=["exact", "fastmath"]) +def test_shapes(shape, use_fast_math): + variant = "fastmath" if use_fast_math else "exact" + case = replace( + ROWWISE, + id=f"{get_shape_id(shape)}-{variant}", + shape=shape, + fast_math=use_fast_math, + ) + assert expect_cutedsl(case) + run_case(case) + + +# Flat dims that are not multiples of 32 leave the optimized path entirely, so they exercise the +# dispatcher's other fallback rather than the CuTeDSL kernel. +@pytest.mark.parametrize("shape", [(48, 48), (16, 1024), (256, 400)], ids=get_shape_id) +def test_unaligned_shapes_fall_back(shape): + case = replace(ROWWISE, id=f"unaligned-{get_shape_id(shape)}", shape=shape) + assert not expect_cutedsl(case) + run_case(case) + + +# --- Input dtypes --- + + +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32, torch.float16], ids=get_dtype_id) +def test_input_dtypes(dtype): + """Only bf16 reaches the CuTeDSL kernel; fp32 and fp16 must take the CUDA fallback and be + left untouched by the backend choice.""" + case = replace(ROWWISE, id=f"dtype-{get_dtype_id(dtype)}", dtype=dtype) + assert expect_cutedsl(case) == (dtype is torch.bfloat16) + run_case(case) + + +# --- Quantizer variations --- + +# The NVFP4 quantize configurations reachable from the PyTorch API whose output is deterministic, +# minus RHT and amax reduction. The tuned-1D configs (plain, fast-math, columnwise, row-scaled) +# are served by the CuTeDSL kernel; the 2D and 4over6 ones check that enabling the backend is +# inert for a config it is never offered. Stochastic rounding is not here, since its bytes are +# not a cross-backend contract. +VARIANT_CASES = [ + ROWWISE, + replace(ROWWISE, id="rowwise-fastmath", fast_math=True), + replace(ROWWISE, id="rowwise-columnwise", columnwise=True), + replace(ROWWISE, id="rowwise-columnwise-fastmath", columnwise=True, fast_math=True), + replace(ROWWISE, id="columnwise-only", rowwise=False, columnwise=True), + replace(ROWWISE, id="row-scaled", row_scaled=True), + replace(ROWWISE, id="row-scaled-fastmath", row_scaled=True, fast_math=True), + replace(ROWWISE, id="2d-quantization", columnwise=True, two_d=True), + replace( + ROWWISE, + id="2d-quantization-columnwise-only", + rowwise=False, + columnwise=True, + two_d=True, + ), + replace(ROWWISE, id="4over6", use_4over6=True), + replace(ROWWISE, id="4over6-columnwise", use_4over6=True, columnwise=True), +] + + +@pytest.mark.parametrize("case", VARIANT_CASES, ids=get_case_id) +def test_config_variants(case): + run_case(case) + + +# Larger shapes for the variants that do produce a transposed output, since a transpose crosses +# chunk boundaries in a way a single chunk cannot show. +@pytest.mark.parametrize("shape", [(64, 64), (256, 1024), (512, 512)], ids=get_shape_id) +def test_transpose_shapes(shape): + case = replace(ROWWISE, id=f"transpose-{get_shape_id(shape)}", shape=shape, columnwise=True) + run_case(case) + + +# --- Stochastic rounding, which is not a byte-for-byte contract --- + + +def nibble_counts(data_bytes, probe_mask): + """Count E2M1 nibbles over the probe positions of a packed FP4 tensor.""" + lo = data_bytes & 0xF + hi = data_bytes >> 4 + nibbles = torch.stack([lo, hi], dim=-1).flatten() # element order: (byte, low-then-high nibble) + picked = nibbles[probe_mask.flatten()] + return torch.bincount(picked.long(), minlength=16), picked.numel() + + +@pytest.mark.parametrize("columnwise", [False, True], ids=["rowwise", "rowwise-columnwise"]) +def test_stochastic_rounding(columnwise): + """The CuTeDSL kernel serves stochastic rounding, but per the module docstring its bytes are + deliberately not compared against CUDA's, even though the current implementation happens to + match it byte for byte (same Philox seeding, thread arrangement and draw order): that + equality would not survive either kernel legitimately reorganizing its work. What is checked + instead is the property that defines stochastic rounding. The input is built so every scaling + block carries one exact 6.0 (making the encode coefficient exactly 1.0 in both directions) + and probe values of 2.75, which sit at 3/4 of the way from 2 to 3 on the E2M1 lattice, so a + correct rounder must send them to 3 with probability 0.75, and to 2 otherwise. With ~200k + probes the tolerance below is dozens of standard deviations, yet fails a round-to-nearest + (would give 1.0) or wrongly-seeded implementation outright. Determinism for a fixed Philox + seed is also asserted, since TE draws the seed from the framework RNG. + """ + case = replace( + ROWWISE, + id="stochastic-rounding" + ("-columnwise" if columnwise else ""), + shape=(256, 1024), + stochastic_rounding=True, + columnwise=columnwise, + ) + assert expect_cutedsl(case) + rows, cols = flat_dims(case.shape) + + # 6.0 on every 16th row and column anchors every scaling block of both directions; the rest + # are probes. The global amax 6.0 makes the global encode scale 448/6 . 6/448-exact, the + # anchored block amaxes make every block decode scale exactly 448 (an E4M3 lattice point), + # and the encode coefficient works out to exactly 1.0, so probes arrive at the converter + # as exactly 2.75. + x = torch.full((rows, cols), 2.75, dtype=torch.bfloat16, device="cuda") + x[::NVFP4_BLOCK_SIZE, :] = 6.0 + x[:, ::NVFP4_BLOCK_SIZE] = 6.0 + probe = torch.ones(rows, cols, dtype=torch.bool, device="cuda") + probe[::NVFP4_BLOCK_SIZE, :] = False + probe[:, ::NVFP4_BLOCK_SIZE] = False + + key = cutedsl_key(cutedsl_flags(case)) + set_cutedsl_backend(True) + try: + with cutedsl_calls(2, case.id, key): + out = quantize(case, x) + out_again = quantize(case, x) + finally: + set_cutedsl_backend(False) + + assert_parts_equal( + quantized_parts(out, case), + quantized_parts(out_again, case), + f"{case.id} (same seed must give the same bytes)", + ) + + checked = {"rowwise data": probe} + if columnwise: + checked["columnwise data"] = probe.T.contiguous() + for name, mask in checked.items(): + counts, total = nibble_counts(quantized_parts(out, case)[name].view(torch.uint8), mask) + # E2M1 codes: 4 -> 2.0, 5 -> 3.0. Everything else means the scales are wrong. + assert counts[4] + counts[5] == total, ( + f"{case.id}: {name} probes landed outside {{2.0, 3.0}}:" + f" {dict(enumerate(counts.tolist()))}" + ) + up_fraction = counts[5].item() / total + assert ( + abs(up_fraction - 0.75) < 0.02 + ), f"{case.id}: {name} rounded 2.75 up with frequency {up_fraction:.4f}, expected 0.75" + + error = relative_dequantize_error(out, x) + assert ( + error < 0.3 + ), f"{case.id}: relative error {error:.4f} is too large for an NVFP4 round trip" + + +# --- Zeros, where the block amax and the global encode scale degenerate --- + + +@pytest.mark.parametrize("use_fast_math", [False, True], ids=["exact", "fastmath"]) +@pytest.mark.parametrize("kind", ["zero-blocks", "negative-zero-blocks", "all-zero"]) +def test_zero_blocks(kind, use_fast_math): + """A block of zeros makes the block amax zero, which both kernels feed to a reciprocal and + clamp; an all-zero tensor does the same to the global encode scale. + + Negative zero is the input that decides how the scaling multiply is spelled. The CUDA kernels + scale with `mul.f32x2` for an f32 coefficient and with `fma.rn.f32.bf16 v, v_h, coeff, zero` + for a bf16 one, and IEEE addition of that +0 addend turns a -0 product into +0 where a plain + multiply keeps the sign. E2M1 has a signed zero, so the two encode a -0 element as 0x0 and + 0x8; the kernel uses the same instruction as CUDA on each path, which is what makes this case + pass for fast math as well. + """ + case = replace( + ROWWISE, + id=f"{kind}-{'fastmath' if use_fast_math else 'exact'}", + shape=(256, 1024), + fast_math=use_fast_math, + ) + if kind == "all-zero": + x = torch.zeros(case.shape, dtype=case.dtype, device="cuda") + else: + x = make_input_with_zero_blocks(case, negative=kind == "negative-zero-blocks") + run_case(case, x=x) + + +# --- The cast-noop flag --- + + +@pytest.mark.parametrize("flag_is_set", [True, False], ids=["noop-set", "noop-clear"]) +def test_cast_noop_flag(flag_is_set): + """The CuTeDSL kernel dereferences the cast-noop flag on device and must leave the output + exactly as it found it when the flag reads 1.0, and quantize as usual when it reads 0.0.""" + case = replace(ROWWISE, id=f"noop-{'set' if flag_is_set else 'clear'}", shape=(256, 1024)) + assert expect_cutedsl(case) + + x = make_input(case) + previous_x = make_input(case, seed=SEED + 1) + flag = torch.full((1,), 1.0 if flag_is_set else 0.0, dtype=torch.float32, device="cuda") + + set_cutedsl_backend(True) + try: + # Every one of the three quantizations has to be the CuTeDSL kernel's, the suppressed one + # included: a kernel that is not launched at all also leaves the output untouched. + with cutedsl_calls(3, case.id, cutedsl_key(cutedsl_flags(case))): + # Quantizing x with the flag clear is what the flag has to suppress, and the state the + # output is left in beforehand is the quantization of a different tensor. + quantized_x = quantized_parts(quantize(case, x), case, with_amax=False) + + out = make_output(case) + quantize_into(out, previous_x) + before = quantized_parts(out, case, with_amax=False) + quantize_into(out, x, noop_flag=flag) + after = quantized_parts(out, case, with_amax=False) + finally: + set_cutedsl_backend(False) + + if flag_is_set: + # Both halves matter: nothing changed, and something would have changed. + assert_parts_equal(before, after, case.id) + assert not torch.equal( + before["rowwise data"], quantized_x["rowwise data"] + ), f"{case.id}: the two inputs quantize to the same bytes, so the flag suppressed nothing" + else: + assert_parts_equal(quantized_x, after, case.id) + + +# --- A numerical floor, independent of the CUDA kernels --- + + +@pytest.mark.parametrize("use_fast_math", [False, True], ids=["exact", "fastmath"]) +def test_dequantized_output_tracks_input(use_fast_math): + """Dequantizing the CuTeDSL output has to land near the input. Bit-exactness against the CUDA + kernel is the real contract, but it is a comparison between two implementations of the same + idea; this is a floor that does not depend on either of them being right. + """ + case = replace(ROWWISE, id="dequantize", shape=(256, 1024), fast_math=use_fast_math) + assert expect_cutedsl(case) + + x = make_input(case) + with fast_math(use_fast_math): + set_cutedsl_backend(True) + try: + with cutedsl_calls(1, case.id, cutedsl_key(cutedsl_flags(case))): + out = quantize(case, x) + finally: + set_cutedsl_backend(False) + + # NVFP4 keeps about two mantissa bits per element, so a few percent of relative error is + # expected; an order of magnitude more than that means the output is not a quantization of x. + error = relative_dequantize_error(out, x) + assert error < 0.2, f"relative error {error:.4f} is too large for an NVFP4 round trip" diff --git a/transformer_engine/common/CuTeDSL/__init__.py b/transformer_engine/common/CuTeDSL/__init__.py index 50d536852c2..05286b2b908 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_transpose 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..902d2e9ed27 --- /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_transpose diff --git a/transformer_engine/common/CuTeDSL/cast/nvfp4/quantize_transpose/__init__.py b/transformer_engine/common/CuTeDSL/cast/nvfp4/quantize_transpose/__init__.py new file mode 100644 index 00000000000..e95698e425d --- /dev/null +++ b/transformer_engine/common/CuTeDSL/cast/nvfp4/quantize_transpose/__init__.py @@ -0,0 +1,7 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""CuTeDSL NVFP4 quantize transpose fused kernel.""" + +from .launcher import get_nvfp4_quantization_function diff --git a/transformer_engine/common/CuTeDSL/cast/nvfp4/quantize_transpose/config.py b/transformer_engine/common/CuTeDSL/cast/nvfp4/quantize_transpose/config.py new file mode 100644 index 00000000000..181c39220e8 --- /dev/null +++ b/transformer_engine/common/CuTeDSL/cast/nvfp4/quantize_transpose/config.py @@ -0,0 +1,33 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Configuration for NVFP4 quantize-transpose.""" + + +class NVFP4QuantizeConfig: + """Instantiation parameters of the CuTE DSL kernel""" + + 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__ diff --git a/transformer_engine/common/CuTeDSL/cast/nvfp4/quantize_transpose/kernel.py b/transformer_engine/common/CuTeDSL/cast/nvfp4/quantize_transpose/kernel.py new file mode 100644 index 00000000000..1ab8eabc255 --- /dev/null +++ b/transformer_engine/common/CuTeDSL/cast/nvfp4/quantize_transpose/kernel.py @@ -0,0 +1,844 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Device and host-launch logic for the tuned NVFP4 quantize-transpose kernel.""" + +from typing import Optional + +import cutlass +from cutlass import cute +from cutlass import ( + BFloat16, + Float4E2M1FN, + Float8E4M3FN, + Float32, + Int32, + Int64, + Uint8, + Uint32, + Uint64, +) +from cuda.bindings.driver import CUstream # pylint: disable=no-name-in-module + +from transformer_engine.common.CuTeDSL.cast.nvfp4.utils import ( + CUTEDSL_DEBUG_LOGGING, + NVFP4_BLOCK_SCALING_SIZE, + NVFP4_SCALE_PAD_INNER, + NVFP4_SCALE_PAD_OUTER, + compute_block_decode_sf, + compute_block_encode_sf, + compute_global_encode_sf, + cvt_f32x8_to_fp4x8, + cvt_f32x8_to_fp4x8_sr, + make_prmt_u32, + mul_cvt_bf16x8_to_fp4x8, + mul_cvt_bf16x8_to_fp4x8_sr, + noop_flag_is_set, + pack_u32x2, + validate_tensor, +) +from transformer_engine.common.CuTeDSL.philox_rng import PhiloxRng +from transformer_engine.common.CuTeDSL.utils import ( + abs_max_x2_bf16, + fabs_f32, + x2_hi_to_f32_bf16, + x2_lo_to_f32_bf16, +) + + +@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 NVFP4QuantizeTransposeTuned1DKernel: + """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 + + 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( + "[CuTeDSL] NVFP4QuantizeTransposeTuned1DKernel.__call__() with config:" + f" {self.cfg}\n" + ) + + ## Validation + + # Validate input and output tensor layouts + M, N = mX.shape + mX_layout = cute.make_ordered_layout((M, N), order=(1, 0)) + mO_row_layout = mX_layout + mO_col_layout = cute.make_ordered_layout((N, M), order=(1, 0)) + validate_tensor(mX, mX_layout, BFloat16) + validate_tensor(mO_row, mO_row_layout, Float4E2M1FN) + validate_tensor(mO_col, mO_col_layout, Float4E2M1FN) + + # Validate scaling factor tensor layouts + mS_row_layout = cute.make_ordered_layout( + ( + cute.round_up(M, NVFP4_SCALE_PAD_OUTER), + cute.round_up(cute.ceil_div(N, NVFP4_BLOCK_SCALING_SIZE), NVFP4_SCALE_PAD_INNER), + ), + order=(1, 0), + ) + mS_col_layout = cute.make_ordered_layout( + ( + cute.round_up(N, NVFP4_SCALE_PAD_OUTER), + cute.round_up(cute.ceil_div(M, NVFP4_BLOCK_SCALING_SIZE), NVFP4_SCALE_PAD_INNER), + ), + order=(1, 0), + ) + validate_tensor(mS_row, mS_row_layout, Float8E4M3FN) + validate_tensor(mS_col, mS_col_layout, Float8E4M3FN) + + # Validate amax tensor layouts + if cutlass.const_expr(self.cfg.ROW_SCALED_NVFP4): + mAmaxRow_layout = cute.make_layout((M,)) + validate_tensor(mAmaxRow, mAmaxRow_layout, Float32) + else: + mAmaxRow_layout = cute.make_layout((1,)) + validate_tensor(mAmaxRow, mAmaxRow_layout, Float32) + mAmaxCol_layout = cute.make_layout((1,)) + validate_tensor(mAmaxCol, mAmaxCol_layout, Float32) + + # Validate RNG state tensor layout + mRngState_layout = cute.make_layout((2,)) + validate_tensor(mRngState, mRngState_layout, Int64) + + ## 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: + 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: + 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 + buffered_tile = lambda inner_cols: 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() + out = cvt_f32x8_to_fp4x8_sr(*scaled, rbits03, rbits47) + else: + out = cvt_f32x8_to_fp4x8(*scaled) + 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) + ) + else: + outs.append(cvt_f32x8_to_fp4x8(*scaled[8 * e : 8 * e + 8])) + + 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_INNER = 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/cast/nvfp4/quantize_transpose/launcher.py b/transformer_engine/common/CuTeDSL/cast/nvfp4/quantize_transpose/launcher.py new file mode 100644 index 00000000000..41e7844f61b --- /dev/null +++ b/transformer_engine/common/CuTeDSL/cast/nvfp4/quantize_transpose/launcher.py @@ -0,0 +1,195 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Compilation and TVM-FFI registration for NVFP4 quantize-transpose.""" + +import logging + +from cutlass import cute +from cutlass import BFloat16, Float4E2M1FN, Float8E4M3FN, Float32, Int64 +import tvm_ffi + +from transformer_engine.common.CuTeDSL.cast.nvfp4.utils import ( + NVFP4_SCALE_PAD_INNER, + NVFP4_SCALE_PAD_OUTER, + NVFP4_SHAPE_ALIGNMENT, +) +from transformer_engine.common.CuTeDSL.utils import device_compute_capability + +from .config import NVFP4QuantizeConfig +from .kernel import NVFP4QuantizeTransposeTuned1DKernel + + +logger = logging.getLogger("transformer_engine.cutedsl.nvfp4") + + +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 = NVFP4QuantizeTransposeTuned1DKernel(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_OUTER), + cute.sym_int32(divisibility=NVFP4_SCALE_PAD_INNER), + ) + scale_col_shape = ( + cute.sym_int32(divisibility=NVFP4_SCALE_PAD_OUTER), + cute.sym_int32(divisibility=NVFP4_SCALE_PAD_INNER), + ) + + # 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/utils.py b/transformer_engine/common/CuTeDSL/cast/nvfp4/utils.py new file mode 100644 index 00000000000..0d1c7a84554 --- /dev/null +++ b/transformer_engine/common/CuTeDSL/cast/nvfp4/utils.py @@ -0,0 +1,384 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Constants, scaling math, and E2M1 conversion helpers for NVFP4 kernels.""" + +import os +from typing import Optional + +import cutlass +from cutlass import cute +from cutlass import BFloat16, Boolean, Float32, Float8E4M3FN, Int64, Uint32 +from cutlass._mlir.dialects import arith as mlir_arith +from cutlass._mlir.dialects import llvm +from cutlass.cutlass_dsl import T, dsl_user_op + + +CUTEDSL_DEBUG_LOGGING = os.environ.get("CUTEDSL_DEBUG_LOGGING", "0") == "1" + +FLOAT32_MAX = 3.4028234663852886e38 +BFLOAT16_MAX = 3.3895313892515355e38 +FLOAT8E4M3_MAX = 448.0 +FLOAT4E2M1_MAX = 6.0 + +# Number of elements per NVFP4 scale block (they share one E4M3 scale factor). +NVFP4_BLOCK_SCALING_SIZE = 16 +# Input row/col divisibility the CUDA tuned-1D kernel requires (16B TMA alignment). +NVFP4_SHAPE_ALIGNMENT = 32 +# Padding TE applies to the scale tensors' outer / inner dim (NVFP4Quantizer::get_scale_shape). +NVFP4_SCALE_PAD_OUTER = 128 +NVFP4_SCALE_PAD_INNER = 4 + + +# Runs if CUTE_DSL_ENABLE_ASSERTIONS=1 or --enable-assertions present in cute.compile +def validate_tensor(tensor: Optional[cute.Tensor], expected_layout: cute.Layout, expected_dtype): + """Assert at trace time that an optional tensor has the expected layout and dtype.""" + if tensor is None: + return + cute.testing.assert_(tensor.layout == expected_layout, "Tensor layout does not match") + cute.testing.assert_(tensor.element_type == expected_dtype, "Tensor dtype does not match") + + +@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 + + +@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 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 = FLOAT8E4M3_MAX * FLOAT4E2M1_MAX / 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 / FLOAT4E2M1_MAX)) + 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 + elif 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) + else: + 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, + ) + ) diff --git a/transformer_engine/common/CuTeDSL/philox_rng.py b/transformer_engine/common/CuTeDSL/philox_rng.py new file mode 100644 index 00000000000..8e8772a3297 --- /dev/null +++ b/transformer_engine/common/CuTeDSL/philox_rng.py @@ -0,0 +1,96 @@ +# 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 + +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) 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_ From 0c55146c6faa8e6ddd4d76eb1948ba61cad2a89b Mon Sep 17 00:00:00 2001 From: Teddy Do Date: Mon, 2 Nov 2026 14:11:09 -0800 Subject: [PATCH 2/5] Split the CuTeDSL NVFP4 backend into a dispatcher and per-kernel modules Lay out the NVFP4 backend the way the MXFP8 one is organized: quantize_nvfp4_common.py holds the config, a kernel base class with supports() / eligible_to_dispatch(), the NVFP4QuantizeEntry that picks a kernel at run time, and compilation and TVM-FFI registration. The TMA-staged tuned-1D kernel moves to quantize_nvfp4_g2s_general.py as NVFP4QuantizeGeneralG2SKernel, the fallback for every config, so that specialized kernels can be added beside it without touching it. No functional change. Signed-off-by: Teddy Do --- .../nvfp4/test_nvfp4_cutedsl_backend.py | 4 +- transformer_engine/common/CuTeDSL/__init__.py | 2 +- .../common/CuTeDSL/cast/nvfp4/__init__.py | 2 +- .../launcher.py => quantize_nvfp4_common.py} | 145 +++++++++++++++++- ...ernel.py => quantize_nvfp4_g2s_general.py} | 18 ++- .../cast/nvfp4/quantize_transpose/__init__.py | 7 - .../cast/nvfp4/quantize_transpose/config.py | 33 ---- .../common/CuTeDSL/cast/nvfp4/utils.py | 19 --- 8 files changed, 155 insertions(+), 75 deletions(-) rename transformer_engine/common/CuTeDSL/cast/nvfp4/{quantize_transpose/launcher.py => quantize_nvfp4_common.py} (58%) rename transformer_engine/common/CuTeDSL/cast/nvfp4/{quantize_transpose/kernel.py => quantize_nvfp4_g2s_general.py} (98%) delete mode 100644 transformer_engine/common/CuTeDSL/cast/nvfp4/quantize_transpose/__init__.py delete mode 100644 transformer_engine/common/CuTeDSL/cast/nvfp4/quantize_transpose/config.py diff --git a/tests/pytorch/nvfp4/test_nvfp4_cutedsl_backend.py b/tests/pytorch/nvfp4/test_nvfp4_cutedsl_backend.py index ab57771cb33..d7ed7f41c4a 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_cutedsl_backend.py +++ b/tests/pytorch/nvfp4/test_nvfp4_cutedsl_backend.py @@ -204,7 +204,7 @@ def _install_call_counters(): that is never seen. Compiling the kernels here rather than letting the dispatcher trigger it is what makes that possible, and costs one compilation per supported config at import. """ - from transformer_engine.common.CuTeDSL.cast.nvfp4.quantize_transpose import ( + from transformer_engine.common.CuTeDSL.cast.nvfp4.quantize_nvfp4_common import ( get_nvfp4_quantization_function, ) @@ -442,7 +442,7 @@ def test_kernel_coverage(flags): registration check pass by proxy. """ # Imports cleanly only with the CuTeDSL stack present, which the module-level skip guarantees. - from transformer_engine.common.CuTeDSL.cast.nvfp4.quantize_transpose import ( + from transformer_engine.common.CuTeDSL.cast.nvfp4.quantize_nvfp4_common import ( get_nvfp4_quantization_function, ) diff --git a/transformer_engine/common/CuTeDSL/__init__.py b/transformer_engine/common/CuTeDSL/__init__.py index 05286b2b908..09a11882ceb 100644 --- a/transformer_engine/common/CuTeDSL/__init__.py +++ b/transformer_engine/common/CuTeDSL/__init__.py @@ -13,7 +13,7 @@ from transformer_engine.common.CuTeDSL.cast.mxfp8.quantize_mxfp8 import ( get_mxfp8_quantization_function, ) -from transformer_engine.common.CuTeDSL.cast.nvfp4.quantize_transpose import ( +from transformer_engine.common.CuTeDSL.cast.nvfp4.quantize_nvfp4_common import ( get_nvfp4_quantization_function, ) diff --git a/transformer_engine/common/CuTeDSL/cast/nvfp4/__init__.py b/transformer_engine/common/CuTeDSL/cast/nvfp4/__init__.py index 902d2e9ed27..56b710bfa35 100644 --- a/transformer_engine/common/CuTeDSL/cast/nvfp4/__init__.py +++ b/transformer_engine/common/CuTeDSL/cast/nvfp4/__init__.py @@ -4,4 +4,4 @@ """CuTeDSL NVFP4 quantization kernels.""" -from . import quantize_transpose +from . import quantize_nvfp4_common diff --git a/transformer_engine/common/CuTeDSL/cast/nvfp4/quantize_transpose/launcher.py b/transformer_engine/common/CuTeDSL/cast/nvfp4/quantize_nvfp4_common.py similarity index 58% rename from transformer_engine/common/CuTeDSL/cast/nvfp4/quantize_transpose/launcher.py rename to transformer_engine/common/CuTeDSL/cast/nvfp4/quantize_nvfp4_common.py index 41e7844f61b..96dea33eda6 100644 --- a/transformer_engine/common/CuTeDSL/cast/nvfp4/quantize_transpose/launcher.py +++ b/transformer_engine/common/CuTeDSL/cast/nvfp4/quantize_nvfp4_common.py @@ -2,12 +2,16 @@ # # See LICENSE for license information. -"""Compilation and TVM-FFI registration for NVFP4 quantize-transpose.""" +"""Shared definitions, kernel dispatch and TVM-FFI registration for CuTeDSL NVFP4.""" +import abc import logging +import os +from typing import Optional from cutlass import cute -from cutlass import BFloat16, Float4E2M1FN, Float8E4M3FN, Float32, Int64 +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.cast.nvfp4.utils import ( @@ -17,13 +21,142 @@ ) from transformer_engine.common.CuTeDSL.utils import device_compute_capability -from .config import NVFP4QuantizeConfig -from .kernel import NVFP4QuantizeTransposeTuned1DKernel - +CUTEDSL_DEBUG_LOGGING = os.environ.get("CUTEDSL_DEBUG_LOGGING", "0") == "1" logger = logging.getLogger("transformer_engine.cutedsl.nvfp4") +@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_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 () 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. @@ -32,7 +165,7 @@ def compile_cutedsl_function_from_cfg(cfg): signature. """ - kernel_obj = NVFP4QuantizeTransposeTuned1DKernel(cfg) + kernel_obj = NVFP4QuantizeEntry(cfg) def _gmem(dtype, shape, stride_order, align): """Fake row-major GMEM tensor; `align` is the assumed base alignment in bytes.""" diff --git a/transformer_engine/common/CuTeDSL/cast/nvfp4/quantize_transpose/kernel.py b/transformer_engine/common/CuTeDSL/cast/nvfp4/quantize_nvfp4_g2s_general.py similarity index 98% rename from transformer_engine/common/CuTeDSL/cast/nvfp4/quantize_transpose/kernel.py rename to transformer_engine/common/CuTeDSL/cast/nvfp4/quantize_nvfp4_g2s_general.py index 1ab8eabc255..61a437535a4 100644 --- a/transformer_engine/common/CuTeDSL/cast/nvfp4/quantize_transpose/kernel.py +++ b/transformer_engine/common/CuTeDSL/cast/nvfp4/quantize_nvfp4_g2s_general.py @@ -2,7 +2,11 @@ # # See LICENSE for license information. -"""Device and host-launch logic for the tuned NVFP4 quantize-transpose kernel.""" +"""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 @@ -22,7 +26,6 @@ from cuda.bindings.driver import CUstream # pylint: disable=no-name-in-module from transformer_engine.common.CuTeDSL.cast.nvfp4.utils import ( - CUTEDSL_DEBUG_LOGGING, NVFP4_BLOCK_SCALING_SIZE, NVFP4_SCALE_PAD_INNER, NVFP4_SCALE_PAD_OUTER, @@ -34,11 +37,11 @@ make_prmt_u32, mul_cvt_bf16x8_to_fp4x8, mul_cvt_bf16x8_to_fp4x8_sr, - noop_flag_is_set, pack_u32x2, validate_tensor, ) from transformer_engine.common.CuTeDSL.philox_rng import PhiloxRng +from .quantize_nvfp4_common import CUTEDSL_DEBUG_LOGGING, NVFP4QuantizeKernelBase, noop_flag_is_set from transformer_engine.common.CuTeDSL.utils import ( abs_max_x2_bf16, fabs_f32, @@ -59,7 +62,7 @@ def zero_b32() -> Int32: prmt_hi_u32 = make_prmt_u32(0x7632) -class NVFP4QuantizeTransposeTuned1DKernel: +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 @@ -97,6 +100,10 @@ class NVFP4QuantizeTransposeTuned1DKernel: 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 @@ -127,8 +134,7 @@ def __call__( """ if cutlass.const_expr(CUTEDSL_DEBUG_LOGGING): cute.printf( - "[CuTeDSL] NVFP4QuantizeTransposeTuned1DKernel.__call__() with config:" - f" {self.cfg}\n" + f"[CuTeDSL] NVFP4QuantizeGeneralG2SKernel.__call__() with config: {self.cfg}\n" ) ## Validation diff --git a/transformer_engine/common/CuTeDSL/cast/nvfp4/quantize_transpose/__init__.py b/transformer_engine/common/CuTeDSL/cast/nvfp4/quantize_transpose/__init__.py deleted file mode 100644 index e95698e425d..00000000000 --- a/transformer_engine/common/CuTeDSL/cast/nvfp4/quantize_transpose/__init__.py +++ /dev/null @@ -1,7 +0,0 @@ -# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# -# See LICENSE for license information. - -"""CuTeDSL NVFP4 quantize transpose fused kernel.""" - -from .launcher import get_nvfp4_quantization_function diff --git a/transformer_engine/common/CuTeDSL/cast/nvfp4/quantize_transpose/config.py b/transformer_engine/common/CuTeDSL/cast/nvfp4/quantize_transpose/config.py deleted file mode 100644 index 181c39220e8..00000000000 --- a/transformer_engine/common/CuTeDSL/cast/nvfp4/quantize_transpose/config.py +++ /dev/null @@ -1,33 +0,0 @@ -# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# -# See LICENSE for license information. - -"""Configuration for NVFP4 quantize-transpose.""" - - -class NVFP4QuantizeConfig: - """Instantiation parameters of the CuTE DSL kernel""" - - 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__ diff --git a/transformer_engine/common/CuTeDSL/cast/nvfp4/utils.py b/transformer_engine/common/CuTeDSL/cast/nvfp4/utils.py index 0d1c7a84554..c57b5dd040e 100644 --- a/transformer_engine/common/CuTeDSL/cast/nvfp4/utils.py +++ b/transformer_engine/common/CuTeDSL/cast/nvfp4/utils.py @@ -4,7 +4,6 @@ """Constants, scaling math, and E2M1 conversion helpers for NVFP4 kernels.""" -import os from typing import Optional import cutlass @@ -14,9 +13,6 @@ from cutlass._mlir.dialects import llvm from cutlass.cutlass_dsl import T, dsl_user_op - -CUTEDSL_DEBUG_LOGGING = os.environ.get("CUTEDSL_DEBUG_LOGGING", "0") == "1" - FLOAT32_MAX = 3.4028234663852886e38 BFLOAT16_MAX = 3.3895313892515355e38 FLOAT8E4M3_MAX = 448.0 @@ -40,21 +36,6 @@ def validate_tensor(tensor: Optional[cute.Tensor], expected_layout: cute.Layout, cute.testing.assert_(tensor.element_type == expected_dtype, "Tensor dtype does not match") -@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 - - @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.""" From 7712feff12239108398d7b3a39d8d09895e52885 Mon Sep 17 00:00:00 2001 From: Teddy Do Date: Mon, 2 Nov 2026 15:27:38 -0800 Subject: [PATCH 3/5] Add register-resident CuTeDSL NVFP4 quantize kernels Add two specialized kernels to the CuTeDSL NVFP4 backend, the CuTeDSL counterparts of the CUDA register-resident 1D quantize kernel, and dispatch to them ahead of the general TMA-staged kernel for BF16 input without fast math or row-scaled NVFP4: - NVFP4QuantizeRowwiseG2RKernel (rowwise only) streams each 16-element block from one 32-byte load straight to its output, with the next wave of rows prefetched into L2. - NVFP4QuantizeBidimensionalG2RKernel (rowwise + columnwise) reads every row twice, once per direction, with the second read served by L2, and stages only the transposed output in shared memory. A cost model picks the number of 64-row slabs per CTA per launch. The per-block coefficient depends only on the block's scale byte, so its divide is computed once per CTA into a shared-memory table. Round-to- nearest output is byte-identical to the CUDA kernels. Stochastic rounding draws Philox once per output block, keyed by the block's coordinates and direction (philox_rng.py). Each draw is plain Philox4x32 of a known counter, as curanddx computes it, with the round count from NVTE_BUILD_NUM_PHILOX_ROUNDS; builds with fewer than 3 rounds keep the general kernel for stochastic rounding. The random bits for a given seed differ from the general kernel's, which follow its thread layout. The kernels launch with the driver's default shared-memory carveout: with min_blocks_per_mp set, CuTeDSL otherwise derives one that ignores the driver's per-CTA reservation and limits residency below the target. Test that the folded block draw equals plain Philox4x32, and that an input that is only 16-byte aligned is served by the general kernel. Signed-off-by: Teddy Do --- .../nvfp4/test_nvfp4_cutedsl_backend.py | 102 +++- .../cast/nvfp4/quantize_nvfp4_common.py | 16 +- .../cast/nvfp4/quantize_nvfp4_g2r_bidim.py | 531 ++++++++++++++++++ .../cast/nvfp4/quantize_nvfp4_g2r_rowwise.py | 266 +++++++++ .../common/CuTeDSL/cast/nvfp4/utils.py | 280 ++++++++- .../common/CuTeDSL/philox_rng.py | 129 ++++- 6 files changed, 1309 insertions(+), 15 deletions(-) create mode 100644 transformer_engine/common/CuTeDSL/cast/nvfp4/quantize_nvfp4_g2r_bidim.py create mode 100644 transformer_engine/common/CuTeDSL/cast/nvfp4/quantize_nvfp4_g2r_rowwise.py diff --git a/tests/pytorch/nvfp4/test_nvfp4_cutedsl_backend.py b/tests/pytorch/nvfp4/test_nvfp4_cutedsl_backend.py index d7ed7f41c4a..6f3b3f933e3 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_cutedsl_backend.py +++ b/tests/pytorch/nvfp4/test_nvfp4_cutedsl_backend.py @@ -2,13 +2,14 @@ # # See LICENSE for license information. -"""Cross-backend bit-exactness tests for the CuTeDSL NVFP4 quantize-transpose kernel. +"""Cross-backend bit-exactness tests for the CuTeDSL NVFP4 quantize-transpose kernels. Every case quantizes the same input twice, once with the CuTeDSL backend disabled and once with it enabled, and requires the two results to agree byte for byte. The CUDA C++ kernels are the -reference: `quantize_transpose_nvfp4_tuned_1D.cuh` is what the CuTeDSL kernel reimplements, and -its numerics are in turn pinned against a pure-PyTorch reference by test_nvfp4_quantize_exact.py, -so matching it bit for bit is the whole contract. +reference: `quantize_transpose_nvfp4_tuned_1D.cuh` is what the CuTeDSL kernels reproduce (the +general TMA-staged kernel, and for the configs they serve the register-resident rowwise and +rowwise + columnwise kernels), and its numerics are in turn pinned against a pure-PyTorch +reference by test_nvfp4_quantize_exact.py, so matching it bit for bit is the whole contract. Configurations the CuTeDSL kernel does not own are kept in the matrix, but it is worth being precise about how little they check: both runs execute the same CUDA kernel, so the comparison @@ -27,14 +28,15 @@ Counting catches all three, including the last, since the run that is meant to be the CUDA reference asserts that no CuTeDSL kernel ran either. -Stochastic rounding is deliberately not compared byte for byte, even though the CuTeDSL kernel -implements it (and, as it happens, currently matches the CUDA kernel bit for bit, having -replicated its Philox seeding, thread arrangement and draw order): which random bits an element -gets follows from the work decomposition, and the CUDA kernel seeds Philox from the chunk -coordinate a CTA starts on without reseeding when it steals a chunk through cluster launch -control, so two implementations, or even two launches, may legitimately disagree without either -being wrong. test_stochastic_rounding checks the property that defines the rounding instead, -and test_nvfp4_sr_quantize.py checks the rounding itself on the CUDA side. +Stochastic rounding is deliberately not compared byte for byte, even though the CuTeDSL kernels +implement it: which random bits an element gets follows from the work decomposition. The general +kernel replicates the CUDA kernel's Philox seeding, thread arrangement and draw order, but the +CUDA kernel does not reseed when it steals a chunk through cluster launch control, so two +launches may legitimately disagree; the register kernels draw once per output block, keyed by the +block's coordinates, so their stream differs from the CUDA kernel's by design. +test_stochastic_rounding checks the property that defines the rounding instead, +test_block_philox_matches_philox4x32 checks the register kernels' generator against plain +Philox4x32, and test_nvfp4_sr_quantize.py checks the rounding itself on the CUDA side. RHT is out of scope: it is a pre-transform with kernels and tests of its own, and it splits the quantization into per-direction calls, so "which backend ran" stops being one fact per case. @@ -738,3 +740,79 @@ def test_dequantized_output_tracks_input(use_fast_math): # expected; an order of magnitude more than that means the output is not a quantization of x. error = relative_dequantize_error(out, x) assert error < 0.2, f"relative error {error:.4f} is too large for an NVFP4 round trip" + + +@pytest.mark.parametrize("columnwise", [False, True], ids=["rowwise", "rowwise-columnwise"]) +def test_unaligned_input(columnwise): + """An input that is only 16-byte aligned is not eligible for the register kernels' 32-byte + loads, so the backend has to serve it with the general kernel, still bit for bit.""" + case = replace(ROWWISE, id="unaligned", columnwise=columnwise) + rows, cols = flat_dims(case.shape) + buffer = torch.empty(rows * cols + 8, dtype=case.dtype, device="cuda") + x = buffer[8:].view(case.shape) # 16 bytes past a 256-byte-aligned allocation + x.copy_(make_input(case)) + assert x.data_ptr() % 32 == 16 + run_case(case, x) + + +def _philox4x32(key, counter, rounds): + """Plain Philox4x32 (curanddx's philox4x32_native_state): `rounds` rounds of the counter.""" + 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 register kernels' folded block draw is plain Philox4x32 of the counter + {offset_lo, offset_hi, subsequence_lo, subsequence_hi} for key = seed, with + 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 folded 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/cast/nvfp4/quantize_nvfp4_common.py b/transformer_engine/common/CuTeDSL/cast/nvfp4/quantize_nvfp4_common.py index 96dea33eda6..34a326b9f5c 100644 --- a/transformer_engine/common/CuTeDSL/cast/nvfp4/quantize_nvfp4_common.py +++ b/transformer_engine/common/CuTeDSL/cast/nvfp4/quantize_nvfp4_common.py @@ -5,10 +5,12 @@ """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 @@ -26,6 +28,14 @@ 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. @@ -120,6 +130,8 @@ 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 @@ -129,7 +141,9 @@ def __init__(self, cfg: NVFP4QuantizeConfig): # order, take the calls they are eligible for. self.general_kernel = NVFP4QuantizeGeneralG2SKernel(cfg) self.specialized_kernels = [ - kernel_cls(cfg) for kernel_cls in () if kernel_cls.supports(cfg) + kernel_cls(cfg) + for kernel_cls in (NVFP4QuantizeRowwiseG2RKernel, NVFP4QuantizeBidimensionalG2RKernel) + if kernel_cls.supports(cfg) ] @cute.jit 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..c6531e24630 --- /dev/null +++ b/transformer_engine/common/CuTeDSL/cast/nvfp4/quantize_nvfp4_g2r_bidim.py @@ -0,0 +1,531 @@ +# 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.cast.nvfp4.utils import ( + NVFP4_BLOCK_SCALING_SIZE, + abs_max_tree_bf16x2, + block_scale_code, + coefficient_of, + compute_global_encode_sf, + e4m3_code, + ld_global_evict_last_v8, + ld_global_v2, + pack_column_block, + pack_row_block, + pack_u32x2, + prefetch_l2, + rcp_rn_f32, + row_offset, + splat_f32x2, + st_global_u16, + st_global_u32, + st_global_u8, + st_global_v2, + st_global_v4, + st_global_v8, + unpack_u64, +) +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, + NVFP4QuantizeConfig, + NVFP4QuantizeKernelBase, + device_sm_count, + noop_flag_is_set, +) + + +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_u64(staging[slot * (self._STAGING_WORDS // 2) + 4 * ll + t]) + d.append(lo) + d.append(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)] + coeff2_a = [splat_f32x2(table[Int32(codes_a[i])]) for i in range(4)] + 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] = e4m3_code((pair_amax << Uint32(16)).bitcast(Float32) * S_enc_col_div6) + codes_b[2 * p + 1] = e4m3_code( + (pair_amax & Uint32(0xFFFF0000)).bitcast(Float32) * S_enc_col_div6 + ) + coeff2_b = [splat_f32x2(table[THREADS + Int32(codes_b[k])]) for k in range(4)] + 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..eed9d38ee28 --- /dev/null +++ b/transformer_engine/common/CuTeDSL/cast/nvfp4/quantize_nvfp4_g2r_rowwise.py @@ -0,0 +1,266 @@ +# 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.cast.nvfp4.utils import ( + NVFP4_BLOCK_SCALING_SIZE, + block_scale_code, + coefficient_of, + compute_global_encode_sf, + ld_global_nc_evict_first_v8, + pack_row_block, + prefetch_l2, + rcp_rn_f32, + row_offset, + splat_f32x2, + st_global_u8, + st_global_v2, +) +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, + NVFP4QuantizeConfig, + NVFP4QuantizeKernelBase, + device_sm_count, + noop_flag_is_set, +) + + +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] = splat_f32x2(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): + coeff2 = splat_f32x2(table[Int32(codes[i])]) + 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/utils.py b/transformer_engine/common/CuTeDSL/cast/nvfp4/utils.py index c57b5dd040e..eaa306fa0ed 100644 --- a/transformer_engine/common/CuTeDSL/cast/nvfp4/utils.py +++ b/transformer_engine/common/CuTeDSL/cast/nvfp4/utils.py @@ -8,7 +8,7 @@ import cutlass from cutlass import cute -from cutlass import BFloat16, Boolean, Float32, Float8E4M3FN, Int64, Uint32 +from cutlass import BFloat16, Boolean, Float32, Float8E4M3FN, Int32, Int64, Uint32, Uint64 from cutlass._mlir.dialects import arith as mlir_arith from cutlass._mlir.dialects import llvm from cutlass.cutlass_dsl import T, dsl_user_op @@ -363,3 +363,281 @@ def mul_cvt_bf16x8_to_fp4x8_sr( asm_dialect=llvm.AsmDialect.AD_ATT, ) ) + + +# --- Register-resident kernels: global memory access by address --- +# +# The register 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 below 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 unpack_u64(v: Int64): + """The (low, high) words of a 64-bit value.""" + return cute.arch.inline_ptx( + "mov.b64 {{$w0}, {$w1}}, {$r0};", write_only_types=[Uint32, Uint32], read_only_args=[v] + ) + + +@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), + ) + + +# --- Register-resident 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 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] + ) + + +@cute.jit +def e4m3_code(x: Float32) -> Uint32: + """The E4M3 byte of x (round to nearest, saturating to 448), as static_cast does. + Its clamp to FLT_MAX in compute_decoding_scaling_factor cannot change the byte: the only + value it alters, +inf, saturates to the same 448.""" + return cute.arch.inline_ptx( + "{ .reg .b16 h; cvt.rn.satfinite.e4m3x2.f32 h, {$r0}, {$r0}; " + "cvt.u32.u16 {$w0}, h; and.b32 {$w0}, {$w0}, 255; }", + write_only_types=[Uint32], + read_only_args=[x], + ) + + +@cute.jit +def e4m3_value(code: Uint32) -> Float32: + """The value of the E4M3 byte `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], + ) + + +@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(e4m3_value(code) * S_dec), Float32(FLOAT32_MAX)) + + +@cute.jit +def splat_f32x2(x: Float32) -> Int64: + """x in both halves of a packed f32 pair, as mul.f32x2 consumes it.""" + return cute.arch.inline_ptx( + "mov.b64 {$w0}, {{$r0}, {$r0}};", write_only_types=[Int64], read_only_args=[x] + ) + + +@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], + ) + + +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).""" + # Imported here: utils.py of the package root is imported by every CuTeDSL module. + from transformer_engine.common.CuTeDSL.utils import ( # pylint: disable=import-outside-toplevel + abs_max_x2_bf16, + ) + + 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] + + +@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 e4m3_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/CuTeDSL/philox_rng.py b/transformer_engine/common/CuTeDSL/philox_rng.py index 8e8772a3297..afb757d6164 100644 --- a/transformer_engine/common/CuTeDSL/philox_rng.py +++ b/transformer_engine/common/CuTeDSL/philox_rng.py @@ -6,10 +6,10 @@ 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")) @@ -94,3 +94,130 @@ def get_rbits(self) -> Uint32: 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 From bcd15f129cc106463f58b9beccd92ae70df0a6f1 Mon Sep 17 00:00:00 2001 From: Teddy Do Date: Tue, 3 Nov 2026 16:52:37 -0800 Subject: [PATCH 4/5] Simplify the CuTeDSL NVFP4 backend tests to the MXFP8 test layout Follow test_mxfp8_cutedsl_backend.py: quantize with the CUDA and the CuTeDSL backend and compare bytes, guarding against a silent CUDA fallback by checking that the config's kernel is registered. Drop the call-counting proxies and the cases for configs the CuTeDSL backend does not serve, which only compared the CUDA kernel against itself. Add the g2r grid.y overflow cases that the g2s kernel must serve, and shorten the module docstring. Signed-off-by: Teddy Do --- .../nvfp4/test_nvfp4_cutedsl_backend.py | 846 ++++-------------- 1 file changed, 187 insertions(+), 659 deletions(-) diff --git a/tests/pytorch/nvfp4/test_nvfp4_cutedsl_backend.py b/tests/pytorch/nvfp4/test_nvfp4_cutedsl_backend.py index 6f3b3f933e3..a6bc5e03c39 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_cutedsl_backend.py +++ b/tests/pytorch/nvfp4/test_nvfp4_cutedsl_backend.py @@ -2,56 +2,21 @@ # # See LICENSE for license information. -"""Cross-backend bit-exactness tests for the CuTeDSL NVFP4 quantize-transpose kernels. - -Every case quantizes the same input twice, once with the CuTeDSL backend disabled and once with -it enabled, and requires the two results to agree byte for byte. The CUDA C++ kernels are the -reference: `quantize_transpose_nvfp4_tuned_1D.cuh` is what the CuTeDSL kernels reproduce (the -general TMA-staged kernel, and for the configs they serve the register-resident rowwise and -rowwise + columnwise kernels), and its numerics are in turn pinned against a pure-PyTorch -reference by test_nvfp4_quantize_exact.py, so matching it bit for bit is the whole contract. - -Configurations the CuTeDSL kernel does not own are kept in the matrix, but it is worth being -precise about how little they check: both runs execute the same CUDA kernel, so the comparison -only says that turning the backend on is inert. That is not nothing, since the dispatcher walks a -wall of NVTE_CHECKs and builds a config key before it declines and the Python entrypoint compiles -without being allowed to let anything escape, but it says nothing about the CuTeDSL kernel itself. -What pins the kernel's coverage down is `test_kernel_coverage`, which asks the entrypoint -directly; a config moves between the two categories by editing `kernel_implements`. - -Which backend ran is therefore never inferred, it is observed: every kernel the CuTeDSL backend -can serve is wrapped in a counting proxy at import (see `_install_call_counters`), and every case -states how many times the dispatcher must call into one. A comparison that silently ran the CUDA -kernel twice would otherwise pass while testing nothing, which is the one failure this file cannot -afford -- and it is not hypothetical: a kernel that fails to compile, a dispatcher check that -rejects the config, and a backend switch that does not take effect all produce exactly that. -Counting catches all three, including the last, since the run that is meant to be the CUDA -reference asserts that no CuTeDSL kernel ran either. - -Stochastic rounding is deliberately not compared byte for byte, even though the CuTeDSL kernels -implement it: which random bits an element gets follows from the work decomposition. The general -kernel replicates the CUDA kernel's Philox seeding, thread arrangement and draw order, but the -CUDA kernel does not reseed when it steals a chunk through cluster launch control, so two -launches may legitimately disagree; the register kernels draw once per output block, keyed by the -block's coordinates, so their stream differs from the CUDA kernel's by design. -test_stochastic_rounding checks the property that defines the rounding instead, -test_block_philox_matches_philox4x32 checks the register kernels' generator against plain -Philox4x32, and test_nvfp4_sr_quantize.py checks the rounding itself on the CUDA side. - -RHT is out of scope: it is a pre-transform with kernels and tests of its own, and it splits the -quantization into per-direction calls, so "which backend ran" stops being one fact per case. +"""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. """ -# Optional CuTeDSL dependencies must be checked before importing Transformer Engine. # pylint: disable=wrong-import-position -import collections import ctypes -import itertools -import math import os -from contextlib import contextmanager -from dataclasses import dataclass, replace import pytest import torch @@ -78,195 +43,38 @@ skip_reason = "NVTE_ENABLE_CUTEDSL_BACKEND is not set" else: skip_reason = reason_for_no_recipe -backend_available = recipe_available and cutedsl_built and cutedsl_enabled -pytestmark = pytest.mark.skipif(not backend_available, reason=skip_reason) +pytestmark = pytest.mark.skipif( + not (recipe_available and cutedsl_built and cutedsl_enabled), + reason=skip_reason, +) -# Elements sharing one E4M3 block scale. NVFP4_BLOCK_SIZE = 16 -# Flat-dim divisibility the tuned-1D path requires (16B alignment for its TMA descriptors). -TUNED_1D_ALIGNMENT = 32 -# Name the CuTeDSL NVFP4 compile-on-demand entrypoint is registered under (kEntrypointName in -# quantize_transpose_nvfp4_cutedsl.cuh). -ENTRYPOINT = "get_nvfp4_quantization_function" -# Filled into the output buffers before every quantization: anything a kernel fails to write then -# shows up as a difference instead of as whatever the caching allocator recycled into the block. -POISON_BYTE = 0xA5 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)] -@dataclass(frozen=True) -class Case: - """An NVFP4Quantizer configuration together with the input it is exercised with.""" - - id: str - shape: tuple = (256, 384) - dtype: torch.dtype = torch.bfloat16 - rowwise: bool = True - columnwise: bool = False - row_scaled: bool = False - two_d: bool = False - stochastic_rounding: bool = False - fast_math: bool = False - use_4over6: bool = False - - -# The configuration the CuTeDSL kernel currently covers; other cases are derived from it. -ROWWISE = Case(id="rowwise") - -get_case_id = lambda case: case.id -get_shape_id = lambda shape: "x".join(str(d) for d in shape) -get_dtype_id = lambda dtype: str(dtype).removeprefix("torch.") - - -def flat_dims(shape): - """The 2D dims the kernels see; a rank > 2 input is flattened over its leading dims.""" - return math.prod(shape[:-1]), shape[-1] - - -def ceil_div(numerator, denominator): - return -(-numerator // denominator) - - -def make_quantizer(case): - return NVFP4Quantizer( - rowwise=case.rowwise, - columnwise=case.columnwise, - with_2d_quantization=case.two_d, - stochastic_rounding=case.stochastic_rounding, - row_scaled_nvfp4=case.row_scaled, - nvfp4_use_4over6=case.use_4over6, - ) - - -def cutedsl_flags(case): - """The four booleans an NVFP4QuantConfig is made of, as the dispatcher reads them off the - quantization config and the output tensor: a transposed output is requested by allocating - columnwise data, and fast math comes from NVTE_USE_FAST_MATH.""" - return (case.stochastic_rounding, case.fast_math, case.row_scaled, case.columnwise) - - -def cutedsl_key(flags): - """Mirror of NVFP4QuantConfig::to_key: the name the compiled kernel for these flags is - registered under, and the key the C++ side caches its lookup result against.""" - major, minor = device_compute_capability() - return f"cutedsl_nvfp4_sm{10 * major + minor}_" + "_".join( - "1" if flag else "0" for flag in flags - ) - - -def kernel_implements(flags): - """Whether the CuTeDSL kernel compiles a kernel for these flags. Every feature of the CUDA - tuned-1D kernel is implemented; the one refusal left is the flag combination that is not a - real configuration (NVFP4QuantizeConfig raises for it): row-scaled quantization has no - transposed output.""" - _, _, row_scaled, return_transpose = flags - return not (row_scaled and return_transpose) - - -def dispatcher_offers(case): - """Whether the C++ dispatcher hands this config to the CuTeDSL backend at all. Mirrors the - NVTE_NVFP4_1D_SCALING branch of dispatch/quantize.cuh and the 2D early-out in - quantize_transpose_nvfp4_cutedsl.cuh: 4over6 has kernels of its own, the optimized (tuned-1D) - path is taken only for bf16 with 32-divisible flat dims and an allocated rowwise output, and - the CuTeDSL entry for 2D scaling declines before it even builds a config key.""" - rows, cols = flat_dims(case.shape) - return ( - not case.use_4over6 - and not case.two_d - and case.dtype is torch.bfloat16 - and rows % TUNED_1D_ALIGNMENT == 0 - and cols % TUNED_1D_ALIGNMENT == 0 - and case.rowwise - ) - - -def expect_cutedsl(case): - """Whether the CuTeDSL kernel is the one that produces this case's output.""" - return dispatcher_offers(case) and kernel_implements(cutedsl_flags(case)) +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) -# --- Observing which backend ran --- - -# How many times the C++ dispatcher has called into a CuTeDSL kernel, per config key. -_cutedsl_calls_by_key = collections.Counter() - - -def _install_call_counters(): - """Replace every kernel the CuTeDSL backend can serve with a proxy that counts its calls. - - A registered key says a kernel exists, not that the dispatcher reached it, and a key another - case registered answers for every case after it; counting invocations is what makes "the - CuTeDSL backend produced this output" a fact rather than an inference. The proxies have to go - in before the first quantization of the process: the C++ side resolves each key once through - `TVMFFIConfigCache::get_or_load` and caches the handle for good, so a proxy installed after - that is never seen. Compiling the kernels here rather than letting the dispatcher trigger it - is what makes that possible, and costs one compilation per supported config at import. - """ - from transformer_engine.common.CuTeDSL.cast.nvfp4.quantize_nvfp4_common import ( - get_nvfp4_quantization_function, - ) - - def counting_proxy(key, kernel): - def proxy(*args): - _cutedsl_calls_by_key[key] += 1 - return kernel(*args) - - return proxy - - for flags in itertools.product([False, True], repeat=4): - if not kernel_implements(flags): - continue - key = cutedsl_key(flags) - # A config kernel_implements() claims is served but that does not compile is left alone: - # test_kernel_coverage is where that contradiction is reported. - if not get_nvfp4_quantization_function(key, *flags): - continue - tvm_ffi.register_global_func( - key, counting_proxy(key, tvm_ffi.get_global_func(key)), override=True - ) - - -@contextmanager -def cutedsl_calls(expected, tag, key=None): - """Require exactly `expected` CuTeDSL kernel invocations inside the block, all of them of the - kernel compiled for `key` when one is named. - - Which key ran matters as much as whether any did: the flags are what a compiled kernel is, so - a case served by the kernel for another flag combination -- exact math answering for fast - math, say -- is testing something other than what it says. - """ - before = collections.Counter(_cutedsl_calls_by_key) - yield - made = _cutedsl_calls_by_key - before - actual = sum(made.values()) - if actual != expected: - if expected == 0: - raise AssertionError( - f"{tag}: the CuTeDSL backend ran {actual} time(s) ({dict(made)}) where it was " - "expected to stay out of the way, so this case did not exercise the backend it " - "names" - ) - raise AssertionError( - f"{tag}: the CuTeDSL backend ran {actual} time(s) instead of {expected}, so the CUDA " - "kernel served a config this case is about. Either the dispatcher declined it, or a " - "quantization earlier in this pytest session made the C++ side cache the kernel " - "handle before the counting proxies were installed" - ) - if key is not None and expected > 0 and set(made) != {key}: - raise AssertionError( - f"{tag}: the CuTeDSL kernels that ran were {dict(made)}, not {expected} call(s) of " - f"{key}, so this case was served by a kernel compiled for another config" - ) - - -if backend_available: - _install_call_counters() - - @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.""" @@ -275,488 +83,210 @@ def _restore_backend_choice_from_env(): set_cutedsl_backend(flag is not None and not flag.startswith("0")) -@contextmanager -def fast_math(enabled): - """NVTE_USE_FAST_MATH is what puts use_fast_math in the quantization config, and TE reads it - afresh on every quantize call, so setting it here is enough to switch the config.""" - name = "NVTE_USE_FAST_MATH" - previous = os.environ.get(name) - os.environ[name] = "1" if enabled else "0" - try: - yield - finally: - if previous is None: - os.environ.pop(name, None) - else: - os.environ[name] = previous - - -def make_input(case, seed=SEED): - """Uniform noise with a per-scaling-block power-of-two magnitude, so that neighbouring blocks - land on different E4M3 block scales and the elements at the top of a block saturate the E2M1 - range.""" - rows, cols = flat_dims(case.shape) - generator = torch.Generator(device="cuda").manual_seed(seed) - values = torch.empty((rows, cols), dtype=torch.float32, device="cuda").uniform_( - -6.0, 6.0, generator=generator +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 ) - if cols % NVFP4_BLOCK_SIZE == 0: - blocks = cols // NVFP4_BLOCK_SIZE - exponents = torch.randint( - -6, 7, (rows, blocks, 1), device="cuda", generator=generator - ).float() - values = (values.view(rows, blocks, NVFP4_BLOCK_SIZE) * torch.exp2(exponents)).view( - rows, cols - ) - return values.to(case.dtype).view(case.shape) - - -def make_input_with_zero_blocks(case, negative=False): - """Whole scaling blocks of zeros, which is where the block amax is zero and both kernels divide - by it and clamp the coefficient. Random data never lands there.""" - x = make_input(case) - x.view(-1, NVFP4_BLOCK_SIZE)[::3].fill_(-0.0 if negative else 0.0) - return x - - -def make_output(case): - """An output tensor of this config's shape, poisoned so that every byte the comparison looks - at has to have been written by the kernel under test.""" - out = make_quantizer(case).make_empty(case.shape, dtype=case.dtype, device="cuda") - for buffer in ( - out._rowwise_data, - out._columnwise_data, - out._rowwise_scale_inv, - out._columnwise_scale_inv, - ): - if buffer is not None: - buffer.view(torch.uint8).fill_(POISON_BYTE) - return out - - -def quantize_into(out, x, noop_flag=None): - """Quantize through the same C++ dispatch a plain quantizer call takes, into a preallocated - output so that its buffers can be poisoned and a cast-noop flag can be passed.""" - # The Philox state for stochastic rounding is drawn from the default CUDA generator when the - # quantize op runs, so seeding here is what makes the two backends see the same randomness. + 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 quantize(case, x): - return quantize_into(make_output(case), x) +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 quantized_parts(out, case, with_amax=True): - """The bytes of a quantized tensor that both backends must agree on. The scale buffers are - padded to (round_up(rows, 128), round_up(ceil(cols / 16), 4)) and no kernel writes the - padding, so only the valid region is taken.""" - rows, cols = flat_dims(case.shape) +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 case.rowwise: + if rowwise: parts["rowwise data"] = out._rowwise_data.view(torch.uint8).clone() parts["rowwise scales"] = out._rowwise_scale_inv.view(torch.uint8)[ - :rows, : ceil_div(cols, NVFP4_BLOCK_SIZE) + :rows, : cols // NVFP4_BLOCK_SIZE ].clone() - if case.columnwise: + 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, : ceil_div(rows, NVFP4_BLOCK_SIZE) + :cols, : rows // NVFP4_BLOCK_SIZE ].clone() - if with_amax: - # The amax comes from a pass of its own rather than from the quantize kernel, but it is - # what fixes the global encode scale, so comparing it separates a wrong amax from wrong - # quantization when a case fails. - if out._amax_rowwise is not None: - parts["rowwise amax"] = out._amax_rowwise.clone() - if out._amax_columnwise is not None: - parts["columnwise amax"] = out._amax_columnwise.clone() + parts["columnwise amax"] = out._amax_columnwise.clone() return parts -def assert_parts_equal(expected, actual, tag): - for name, expected_bytes in expected.items(): - assert torch.equal(actual[name], expected_bytes), f"{tag}: {name} does not match" - - -def relative_dequantize_error(out, x): - """How far the dequantized output sits from the input, relative to the input's norm.""" - dequantized = out.dequantize(dtype=torch.bfloat16).float() - reference = x.float() - error = torch.linalg.vector_norm(dequantized - reference) / torch.linalg.vector_norm(reference) - return error.item() - - -def run_case(case, x=None): - """Assert the CuTeDSL and CUDA backends produce bit-identical output for this config, and - that the case exercised the backend it was meant to.""" +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 = make_input(case) + x = generate_input(shape) + fast_math = os.environ.get("NVTE_USE_FAST_MATH", "0") == "1" + quantizer = make_quantizer(rowwise, columnwise, row_scaled) - with fast_math(case.fast_math): + 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) - # The reference run has to be CUDA, which is only true if the backend switch took effect. - with cutedsl_calls(0, f"{case.id} (CUDA reference run)"): - cuda_parts = quantized_parts(quantize(case, x), case) - - key = cutedsl_key(cutedsl_flags(case)) - set_cutedsl_backend(True) - try: - with cutedsl_calls(1 if expect_cutedsl(case) else 0, case.id, key): - cutedsl_parts = quantized_parts(quantize(case, x), case) - finally: - set_cutedsl_backend(False) - - registered = tvm_ffi.get_global_func(key, allow_missing=True) is not None - if expect_cutedsl(case): - # Redundant with the call count above, but it separates a kernel that failed to compile - # from a dispatcher that declined a kernel it has. - assert registered, ( - f"{case.id}: nothing is registered under {key}, so the CuTeDSL backend fell back to " - "CUDA and this case compared CUDA against itself" - ) - elif not kernel_implements(cutedsl_flags(case)): - assert not registered, ( - f"{case.id}: a kernel is registered under {key}, so the CuTeDSL backend now covers a " - "config this test expects it to decline; update kernel_implements()" - ) - # A config the dispatcher never offers (2D scaling, 4over6, non-bf16, flat dims not divisible - # by 32) shares its key with configs it does offer, so registration says nothing about it. - # What such a case checks is that turning the backend on leaves the CUDA result alone. - assert_parts_equal(cuda_parts, cutedsl_parts, case.id) + # 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" -# --- The kernel's own account of what it supports --- +@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( - "flags", - list(itertools.product([False, True], repeat=4)), - ids=lambda flags: cutedsl_key(flags).removeprefix("cutedsl_nvfp4_"), -) -def test_kernel_coverage(flags): - """Ask the compile-on-demand entrypoint directly which flag combinations it compiles a kernel - for. This is the one place the kernel's coverage is asserted, so a config that starts or stops - being supported is reported here instead of showing up as a puzzling fallback elsewhere. - - The probe registers under a name of its own so that it cannot make a later case's - registration check pass by proxy. - """ - # Imports cleanly only with the CuTeDSL stack present, which the module-level skip guarantees. - from transformer_engine.common.CuTeDSL.cast.nvfp4.quantize_nvfp4_common import ( - get_nvfp4_quantization_function, - ) - assert tvm_ffi.get_global_func(ENTRYPOINT, allow_missing=True) is not None, ( - f"{ENTRYPOINT} is not in the tvm-ffi registry, so the C++ dispatcher has no way to ask for " - "a CuTeDSL NVFP4 kernel at all" - ) +@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) - stochastic_rounding, use_fast_math, row_scaled, return_transpose = flags - supported = get_nvfp4_quantization_function( - "test_probe_" + cutedsl_key(flags), - stochastic_rounding, - use_fast_math, - row_scaled, - return_transpose, - ) - assert bool(supported) == kernel_implements(flags), ( - f"the CuTeDSL NVFP4 backend reports supported={bool(supported)} for {cutedsl_key(flags)}, " - "which contradicts kernel_implements()" - ) +@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) -# --- Shapes, over the config the kernel serves --- -# Flat dims must be multiples of 32 for the tuned-1D path to be chosen at all; within that, the -# interesting axis is how the shape sits relative to the kernel's 64x64 chunk, since everything -# that does not tile exactly is covered by predication. -SHAPES = [ - (32, 32), # one partial chunk, in both dimensions - (32, 1024), # a single row of chunks, partial - (1024, 32), # a single column of chunks, partial - (96, 160), # partial chunks along both edges - (128, 128), - (256, 1024), - (512, 512), - (8, 32, 1024), # rank > 2, flattened to (256, 1024) - (8192, 7168), # a shape from a real model -] +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) -@pytest.mark.parametrize("shape", SHAPES, ids=get_shape_id) -@pytest.mark.parametrize("use_fast_math", [False, True], ids=["exact", "fastmath"]) -def test_shapes(shape, use_fast_math): - variant = "fastmath" if use_fast_math else "exact" - case = replace( - ROWWISE, - id=f"{get_shape_id(shape)}-{variant}", - shape=shape, - fast_math=use_fast_math, - ) - assert expect_cutedsl(case) - run_case(case) - - -# Flat dims that are not multiples of 32 leave the optimized path entirely, so they exercise the -# dispatcher's other fallback rather than the CuTeDSL kernel. -@pytest.mark.parametrize("shape", [(48, 48), (16, 1024), (256, 400)], ids=get_shape_id) -def test_unaligned_shapes_fall_back(shape): - case = replace(ROWWISE, id=f"unaligned-{get_shape_id(shape)}", shape=shape) - assert not expect_cutedsl(case) - run_case(case) - - -# --- Input dtypes --- - - -@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32, torch.float16], ids=get_dtype_id) -def test_input_dtypes(dtype): - """Only bf16 reaches the CuTeDSL kernel; fp32 and fp16 must take the CUDA fallback and be - left untouched by the backend choice.""" - case = replace(ROWWISE, id=f"dtype-{get_dtype_id(dtype)}", dtype=dtype) - assert expect_cutedsl(case) == (dtype is torch.bfloat16) - run_case(case) - - -# --- Quantizer variations --- - -# The NVFP4 quantize configurations reachable from the PyTorch API whose output is deterministic, -# minus RHT and amax reduction. The tuned-1D configs (plain, fast-math, columnwise, row-scaled) -# are served by the CuTeDSL kernel; the 2D and 4over6 ones check that enabling the backend is -# inert for a config it is never offered. Stochastic rounding is not here, since its bytes are -# not a cross-backend contract. -VARIANT_CASES = [ - ROWWISE, - replace(ROWWISE, id="rowwise-fastmath", fast_math=True), - replace(ROWWISE, id="rowwise-columnwise", columnwise=True), - replace(ROWWISE, id="rowwise-columnwise-fastmath", columnwise=True, fast_math=True), - replace(ROWWISE, id="columnwise-only", rowwise=False, columnwise=True), - replace(ROWWISE, id="row-scaled", row_scaled=True), - replace(ROWWISE, id="row-scaled-fastmath", row_scaled=True, fast_math=True), - replace(ROWWISE, id="2d-quantization", columnwise=True, two_d=True), - replace( - ROWWISE, - id="2d-quantization-columnwise-only", - rowwise=False, - columnwise=True, - two_d=True, - ), - replace(ROWWISE, id="4over6", use_4over6=True), - replace(ROWWISE, id="4over6-columnwise", use_4over6=True, columnwise=True), -] +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("case", VARIANT_CASES, ids=get_case_id) -def test_config_variants(case): - run_case(case) - - -# Larger shapes for the variants that do produce a transposed output, since a transpose crosses -# chunk boundaries in a way a single chunk cannot show. -@pytest.mark.parametrize("shape", [(64, 64), (256, 1024), (512, 512)], ids=get_shape_id) -def test_transpose_shapes(shape): - case = replace(ROWWISE, id=f"transpose-{get_shape_id(shape)}", shape=shape, columnwise=True) - run_case(case) - - -# --- Stochastic rounding, which is not a byte-for-byte contract --- - - -def nibble_counts(data_bytes, probe_mask): - """Count E2M1 nibbles over the probe positions of a packed FP4 tensor.""" - lo = data_bytes & 0xF - hi = data_bytes >> 4 - nibbles = torch.stack([lo, hi], dim=-1).flatten() # element order: (byte, low-then-high nibble) - picked = nibbles[probe_mask.flatten()] - return torch.bincount(picked.long(), minlength=16), picked.numel() - - -@pytest.mark.parametrize("columnwise", [False, True], ids=["rowwise", "rowwise-columnwise"]) -def test_stochastic_rounding(columnwise): - """The CuTeDSL kernel serves stochastic rounding, but per the module docstring its bytes are - deliberately not compared against CUDA's, even though the current implementation happens to - match it byte for byte (same Philox seeding, thread arrangement and draw order): that - equality would not survive either kernel legitimately reorganizing its work. What is checked - instead is the property that defines stochastic rounding. The input is built so every scaling - block carries one exact 6.0 (making the encode coefficient exactly 1.0 in both directions) - and probe values of 2.75, which sit at 3/4 of the way from 2 to 3 on the E2M1 lattice, so a - correct rounder must send them to 3 with probability 0.75, and to 2 otherwise. With ~200k - probes the tolerance below is dozens of standard deviations, yet fails a round-to-nearest - (would give 1.0) or wrongly-seeded implementation outright. Determinism for a fixed Philox - seed is also asserted, since TE draws the seed from the framework RNG. - """ - case = replace( - ROWWISE, - id="stochastic-rounding" + ("-columnwise" if columnwise else ""), - shape=(256, 1024), - stochastic_rounding=True, - columnwise=columnwise, - ) - assert expect_cutedsl(case) - rows, cols = flat_dims(case.shape) - - # 6.0 on every 16th row and column anchors every scaling block of both directions; the rest - # are probes. The global amax 6.0 makes the global encode scale 448/6 . 6/448-exact, the - # anchored block amaxes make every block decode scale exactly 448 (an E4M3 lattice point), - # and the encode coefficient works out to exactly 1.0, so probes arrive at the converter - # as exactly 2.75. - x = torch.full((rows, cols), 2.75, dtype=torch.bfloat16, device="cuda") - x[::NVFP4_BLOCK_SIZE, :] = 6.0 - x[:, ::NVFP4_BLOCK_SIZE] = 6.0 - probe = torch.ones(rows, cols, dtype=torch.bool, device="cuda") - probe[::NVFP4_BLOCK_SIZE, :] = False - probe[:, ::NVFP4_BLOCK_SIZE] = False +@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") - key = cutedsl_key(cutedsl_flags(case)) set_cutedsl_backend(True) try: - with cutedsl_calls(2, case.id, key): - out = quantize(case, x) - out_again = quantize(case, x) + 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) - assert_parts_equal( - quantized_parts(out, case), - quantized_parts(out_again, case), - f"{case.id} (same seed must give the same bytes)", - ) - - checked = {"rowwise data": probe} - if columnwise: - checked["columnwise data"] = probe.T.contiguous() - for name, mask in checked.items(): - counts, total = nibble_counts(quantized_parts(out, case)[name].view(torch.uint8), mask) - # E2M1 codes: 4 -> 2.0, 5 -> 3.0. Everything else means the scales are wrong. - assert counts[4] + counts[5] == total, ( - f"{case.id}: {name} probes landed outside {{2.0, 3.0}}:" - f" {dict(enumerate(counts.tolist()))}" - ) - up_fraction = counts[5].item() / total - assert ( - abs(up_fraction - 0.75) < 0.02 - ), f"{case.id}: {name} rounded 2.75 up with frequency {up_fraction:.4f}, expected 0.75" - - error = relative_dequantize_error(out, x) - assert ( - error < 0.3 - ), f"{case.id}: relative error {error:.4f} is too large for an NVFP4 round trip" - - -# --- Zeros, where the block amax and the global encode scale degenerate --- - - -@pytest.mark.parametrize("use_fast_math", [False, True], ids=["exact", "fastmath"]) -@pytest.mark.parametrize("kind", ["zero-blocks", "negative-zero-blocks", "all-zero"]) -def test_zero_blocks(kind, use_fast_math): - """A block of zeros makes the block amax zero, which both kernels feed to a reciprocal and - clamp; an all-zero tensor does the same to the global encode scale. - - Negative zero is the input that decides how the scaling multiply is spelled. The CUDA kernels - scale with `mul.f32x2` for an f32 coefficient and with `fma.rn.f32.bf16 v, v_h, coeff, zero` - for a bf16 one, and IEEE addition of that +0 addend turns a -0 product into +0 where a plain - multiply keeps the sign. E2M1 has a signed zero, so the two encode a -0 element as 0x0 and - 0x8; the kernel uses the same instruction as CUDA on each path, which is what makes this case - pass for fast math as well. - """ - case = replace( - ROWWISE, - id=f"{kind}-{'fastmath' if use_fast_math else 'exact'}", - shape=(256, 1024), - fast_math=use_fast_math, - ) - if kind == "all-zero": - x = torch.zeros(case.shape, dtype=case.dtype, device="cuda") - else: - x = make_input_with_zero_blocks(case, negative=kind == "negative-zero-blocks") - run_case(case, x=x) + 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" -# --- The cast-noop flag --- - - -@pytest.mark.parametrize("flag_is_set", [True, False], ids=["noop-set", "noop-clear"]) -def test_cast_noop_flag(flag_is_set): - """The CuTeDSL kernel dereferences the cast-noop flag on device and must leave the output - exactly as it found it when the flag reads 1.0, and quantize as usual when it reads 0.0.""" - case = replace(ROWWISE, id=f"noop-{'set' if flag_is_set else 'clear'}", shape=(256, 1024)) - assert expect_cutedsl(case) - - x = make_input(case) - previous_x = make_input(case, seed=SEED + 1) - flag = torch.full((1,), 1.0 if flag_is_set else 0.0, dtype=torch.float32, device="cuda") +@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: - # Every one of the three quantizations has to be the CuTeDSL kernel's, the suppressed one - # included: a kernel that is not launched at all also leaves the output untouched. - with cutedsl_calls(3, case.id, cutedsl_key(cutedsl_flags(case))): - # Quantizing x with the flag clear is what the flag has to suppress, and the state the - # output is left in beforehand is the quantization of a different tensor. - quantized_x = quantized_parts(quantize(case, x), case, with_amax=False) - - out = make_output(case) - quantize_into(out, previous_x) - before = quantized_parts(out, case, with_amax=False) - quantize_into(out, x, noop_flag=flag) - after = quantized_parts(out, case, with_amax=False) + 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 - if flag_is_set: - # Both halves matter: nothing changed, and something would have changed. - assert_parts_equal(before, after, case.id) - assert not torch.equal( - before["rowwise data"], quantized_x["rowwise data"] - ), f"{case.id}: the two inputs quantize to the same bytes, so the flag suppressed nothing" - else: - assert_parts_equal(quantized_x, after, case.id) - - -# --- A numerical floor, independent of the CUDA kernels --- - - -@pytest.mark.parametrize("use_fast_math", [False, True], ids=["exact", "fastmath"]) -def test_dequantized_output_tracks_input(use_fast_math): - """Dequantizing the CuTeDSL output has to land near the input. Bit-exactness against the CUDA - kernel is the real contract, but it is a comparison between two implementations of the same - idea; this is a floor that does not depend on either of them being right. - """ - case = replace(ROWWISE, id="dequantize", shape=(256, 1024), fast_math=use_fast_math) - assert expect_cutedsl(case) - - x = make_input(case) - with fast_math(use_fast_math): - set_cutedsl_backend(True) - try: - with cutedsl_calls(1, case.id, cutedsl_key(cutedsl_flags(case))): - out = quantize(case, x) - finally: - set_cutedsl_backend(False) - - # NVFP4 keeps about two mantissa bits per element, so a few percent of relative error is - # expected; an order of magnitude more than that means the output is not a quantization of x. - error = relative_dequantize_error(out, x) - assert error < 0.2, f"relative error {error:.4f} is too large for an NVFP4 round trip" - - -@pytest.mark.parametrize("columnwise", [False, True], ids=["rowwise", "rowwise-columnwise"]) -def test_unaligned_input(columnwise): - """An input that is only 16-byte aligned is not eligible for the register kernels' 32-byte - loads, so the backend has to serve it with the general kernel, still bit for bit.""" - case = replace(ROWWISE, id="unaligned", columnwise=columnwise) - rows, cols = flat_dims(case.shape) - buffer = torch.empty(rows * cols + 8, dtype=case.dtype, device="cuda") - x = buffer[8:].view(case.shape) # 16 bytes past a 256-byte-aligned allocation - x.copy_(make_input(case)) - assert x.data_ptr() % 32 == 16 - run_case(case, x) + 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): `rounds` rounds of the counter.""" + """Plain Philox4x32 (curanddx's philox4x32_native_state).""" mask = 0xFFFFFFFF k0, k1 = key c = list(counter) @@ -769,8 +299,8 @@ def _philox4x32(key, counter, rounds): def test_block_philox_matches_philox4x32(): - """The register kernels' folded block draw is plain Philox4x32 of the counter - {offset_lo, offset_hi, subsequence_lo, subsequence_hi} for key = seed, with + """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 @@ -779,7 +309,7 @@ def test_block_philox_matches_philox4x32(): from transformer_engine.common.CuTeDSL import philox_rng if not philox_rng.BLOCK_PHILOX_FOLD_SUPPORTED: - pytest.skip("the folded block draw needs at least 3 Philox rounds") + pytest.skip("the per-block draw needs at least 3 Philox rounds") @cute.kernel def draw_kernel(mIn: cute.Tensor, mOut: cute.Tensor): @@ -811,8 +341,6 @@ def draw(mIn: cute.Tensor, mOut: cute.Tensor): 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, + (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}" From a18f3e90b32639473c8a09ac361babfca07d580a Mon Sep 17 00:00:00 2001 From: Teddy Do Date: Wed, 4 Nov 2026 14:38:48 -0800 Subject: [PATCH 5/5] Share CuTeDSL helpers between the NVFP4 and MXFP8 backends Lay the NVFP4 helpers out like the FP8 ones: cast/nvfp4/utils.py becomes utils_nvfp4.py next to utils_fp8.py and keeps only the NVFP4 scaling math and E2M1 conversions. Format-independent helpers (selects, packs, byte permutes, the correctly rounded reciprocal, the bf16x2 abs-max tree and the addressed global loads and stores) move to utils.py, and the E4M3 encode/decode helpers to utils_fp8.py, reusing the existing pack_f32x2, unpack_i64_to_i32x2 and cvt_f32x2_to_fp8e4m3x2 rather than NVFP4 copies of them. The block size, shape alignment and scale padding move to quantize_nvfp4_common.py, as the MXFP8 block size lives in its common module. Drop validate_tensor: it only ran under CuTeDSL assertions, the MXFP8 backend has no equivalent, and the C++ dispatcher already checks the tensors it passes. Fix the pylint findings in the general kernel. Signed-off-by: Teddy Do --- .../cast/nvfp4/quantize_nvfp4_common.py | 23 +- .../cast/nvfp4/quantize_nvfp4_g2r_bidim.py | 48 +-- .../cast/nvfp4/quantize_nvfp4_g2r_rowwise.py | 25 +- .../cast/nvfp4/quantize_nvfp4_g2s_general.py | 96 ++---- transformer_engine/common/CuTeDSL/utils.py | 185 +++++++++++- .../common/CuTeDSL/utils_fp8.py | 21 ++ .../{cast/nvfp4/utils.py => utils_nvfp4.py} | 279 ++---------------- 7 files changed, 317 insertions(+), 360 deletions(-) rename transformer_engine/common/CuTeDSL/{cast/nvfp4/utils.py => utils_nvfp4.py} (61%) diff --git a/transformer_engine/common/CuTeDSL/cast/nvfp4/quantize_nvfp4_common.py b/transformer_engine/common/CuTeDSL/cast/nvfp4/quantize_nvfp4_common.py index 34a326b9f5c..36fc26d6a86 100644 --- a/transformer_engine/common/CuTeDSL/cast/nvfp4/quantize_nvfp4_common.py +++ b/transformer_engine/common/CuTeDSL/cast/nvfp4/quantize_nvfp4_common.py @@ -16,15 +16,20 @@ from cuda.bindings.driver import CUstream # pylint: disable=no-name-in-module import tvm_ffi -from transformer_engine.common.CuTeDSL.cast.nvfp4.utils import ( - NVFP4_SCALE_PAD_INNER, - NVFP4_SCALE_PAD_OUTER, - NVFP4_SHAPE_ALIGNMENT, -) 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") @@ -202,12 +207,12 @@ def _gmem(dtype, shape, stride_order, align): # 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_OUTER), - cute.sym_int32(divisibility=NVFP4_SCALE_PAD_INNER), + 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_OUTER), - cute.sym_int32(divisibility=NVFP4_SCALE_PAD_INNER), + 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 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 index c6531e24630..c1f5b652d7e 100644 --- a/transformer_engine/common/CuTeDSL/cast/nvfp4/quantize_nvfp4_g2r_bidim.py +++ b/transformer_engine/common/CuTeDSL/cast/nvfp4/quantize_nvfp4_g2r_bidim.py @@ -13,29 +13,32 @@ 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.cast.nvfp4.utils import ( - NVFP4_BLOCK_SCALING_SIZE, +from transformer_engine.common.CuTeDSL.utils import ( abs_max_tree_bf16x2, - block_scale_code, - coefficient_of, - compute_global_encode_sf, - e4m3_code, ld_global_evict_last_v8, ld_global_v2, - pack_column_block, - pack_row_block, + pack_f32x2, pack_u32x2, prefetch_l2, rcp_rn_f32, row_offset, - splat_f32x2, st_global_u16, st_global_u32, st_global_u8, st_global_v2, st_global_v4, st_global_v8, - unpack_u64, + 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, @@ -48,10 +51,11 @@ ) from .quantize_nvfp4_common import ( CUTEDSL_DEBUG_LOGGING, - NVFP4QuantizeConfig, - NVFP4QuantizeKernelBase, + NVFP4_BLOCK_SCALING_SIZE, device_sm_count, noop_flag_is_set, + NVFP4QuantizeConfig, + NVFP4QuantizeKernelBase, ) @@ -379,9 +383,11 @@ def kernel( if cbase + q < K and nblk > 0: d = [] for t in cutlass.range_constexpr(4): - lo, hi = unpack_u64(staging[slot * (self._STAGING_WORDS // 2) + 4 * ll + t]) - d.append(lo) - d.append(hi) + 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 @@ -469,7 +475,8 @@ def _slab( 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)] - coeff2_a = [splat_f32x2(table[Int32(codes_a[i])]) 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): @@ -515,11 +522,14 @@ def _slab( 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] = e4m3_code((pair_amax << Uint32(16)).bitcast(Float32) * S_enc_col_div6) - codes_b[2 * p + 1] = e4m3_code( + 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 ) - coeff2_b = [splat_f32x2(table[THREADS + Int32(codes_b[k])]) for k in range(4)] + 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 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 index eed9d38ee28..50f50b7ef40 100644 --- a/transformer_engine/common/CuTeDSL/cast/nvfp4/quantize_nvfp4_g2r_rowwise.py +++ b/transformer_engine/common/CuTeDSL/cast/nvfp4/quantize_nvfp4_g2r_rowwise.py @@ -12,20 +12,21 @@ 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.cast.nvfp4.utils import ( - NVFP4_BLOCK_SCALING_SIZE, - block_scale_code, - coefficient_of, - compute_global_encode_sf, +from transformer_engine.common.CuTeDSL.utils import ( ld_global_nc_evict_first_v8, - pack_row_block, + pack_f32x2, prefetch_l2, rcp_rn_f32, row_offset, - splat_f32x2, 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, @@ -35,10 +36,11 @@ ) from .quantize_nvfp4_common import ( CUTEDSL_DEBUG_LOGGING, - NVFP4QuantizeConfig, - NVFP4QuantizeKernelBase, + NVFP4_BLOCK_SCALING_SIZE, device_sm_count, noop_flag_is_set, + NVFP4QuantizeConfig, + NVFP4QuantizeKernelBase, ) @@ -221,7 +223,7 @@ def kernel( if cutlass.const_expr(STOCHASTIC): table[tidx] = coeff else: - table[tidx] = splat_f32x2(coeff) + table[tidx] = pack_f32x2(coeff, coeff) active = j0 < blocks_per_row for w in cutlass.range_constexpr(0, ROWS, WAVE): @@ -256,7 +258,8 @@ def kernel( ) for i in cutlass.range_constexpr(g, g + 2): if cutlass.const_expr(STOCHASTIC): - coeff2 = splat_f32x2(table[Int32(codes[i])]) + 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) 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 index 61a437535a4..58fcdfd844a 100644 --- a/transformer_engine/common/CuTeDSL/cast/nvfp4/quantize_nvfp4_g2s_general.py +++ b/transformer_engine/common/CuTeDSL/cast/nvfp4/quantize_nvfp4_g2s_general.py @@ -14,7 +14,6 @@ from cutlass import cute from cutlass import ( BFloat16, - Float4E2M1FN, Float8E4M3FN, Float32, Int32, @@ -25,28 +24,29 @@ ) from cuda.bindings.driver import CUstream # pylint: disable=no-name-in-module -from transformer_engine.common.CuTeDSL.cast.nvfp4.utils import ( - NVFP4_BLOCK_SCALING_SIZE, - NVFP4_SCALE_PAD_INNER, - NVFP4_SCALE_PAD_OUTER, +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, - make_prmt_u32, mul_cvt_bf16x8_to_fp4x8, mul_cvt_bf16x8_to_fp4x8_sr, - pack_u32x2, - validate_tensor, ) from transformer_engine.common.CuTeDSL.philox_rng import PhiloxRng -from .quantize_nvfp4_common import CUTEDSL_DEBUG_LOGGING, NVFP4QuantizeKernelBase, noop_flag_is_set -from transformer_engine.common.CuTeDSL.utils import ( - abs_max_x2_bf16, - fabs_f32, - x2_hi_to_f32_bf16, - x2_lo_to_f32_bf16, +from .quantize_nvfp4_common import ( + CUTEDSL_DEBUG_LOGGING, + NVFP4_BLOCK_SCALING_SIZE, + noop_flag_is_set, + NVFP4QuantizeKernelBase, ) @@ -137,48 +137,7 @@ def __call__( f"[CuTeDSL] NVFP4QuantizeGeneralG2SKernel.__call__() with config: {self.cfg}\n" ) - ## Validation - - # Validate input and output tensor layouts M, N = mX.shape - mX_layout = cute.make_ordered_layout((M, N), order=(1, 0)) - mO_row_layout = mX_layout - mO_col_layout = cute.make_ordered_layout((N, M), order=(1, 0)) - validate_tensor(mX, mX_layout, BFloat16) - validate_tensor(mO_row, mO_row_layout, Float4E2M1FN) - validate_tensor(mO_col, mO_col_layout, Float4E2M1FN) - - # Validate scaling factor tensor layouts - mS_row_layout = cute.make_ordered_layout( - ( - cute.round_up(M, NVFP4_SCALE_PAD_OUTER), - cute.round_up(cute.ceil_div(N, NVFP4_BLOCK_SCALING_SIZE), NVFP4_SCALE_PAD_INNER), - ), - order=(1, 0), - ) - mS_col_layout = cute.make_ordered_layout( - ( - cute.round_up(N, NVFP4_SCALE_PAD_OUTER), - cute.round_up(cute.ceil_div(M, NVFP4_BLOCK_SCALING_SIZE), NVFP4_SCALE_PAD_INNER), - ), - order=(1, 0), - ) - validate_tensor(mS_row, mS_row_layout, Float8E4M3FN) - validate_tensor(mS_col, mS_col_layout, Float8E4M3FN) - - # Validate amax tensor layouts - if cutlass.const_expr(self.cfg.ROW_SCALED_NVFP4): - mAmaxRow_layout = cute.make_layout((M,)) - validate_tensor(mAmaxRow, mAmaxRow_layout, Float32) - else: - mAmaxRow_layout = cute.make_layout((1,)) - validate_tensor(mAmaxRow, mAmaxRow_layout, Float32) - mAmaxCol_layout = cute.make_layout((1,)) - validate_tensor(mAmaxCol, mAmaxCol_layout, Float32) - - # Validate RNG state tensor layout - mRngState_layout = cute.make_layout((2,)) - validate_tensor(mRngState, mRngState_layout, Int64) ## TMA descriptors # The FP4 outputs are moved as bytes: two elements to a byte, so the byte view of the @@ -300,7 +259,7 @@ def kernel_main( if cutlass.const_expr(cfg.RETURN_TRANSPOSE): @cute.struct - class SharedStorage: + 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 @@ -323,7 +282,7 @@ class SharedStorage: else: @cute.struct - class SharedStorage: + 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 @@ -341,10 +300,12 @@ class SharedStorage: storage = smem.allocate(SharedStorage) # Create views into the shared memory - buffered_tile = lambda inner_cols: cute.make_layout( - ((TILE, inner_cols), self.NUM_BUFFERS), - stride=((inner_cols, 1), TILE * inner_cols), - ) + 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( @@ -676,9 +637,10 @@ def rowwise_tile( 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) + out = cvt_f32x8_to_fp4x8(*scaled) # pylint: disable=no-value-for-parameter sO_u32[row * u32_per_row + group] = out @cute.jit @@ -795,10 +757,14 @@ def colwise_tile( rbits03 = rng.get_rbits() rbits47 = rng.get_rbits() outs.append( - cvt_f32x8_to_fp4x8_sr(*scaled[8 * e : 8 * e + 8], rbits03, rbits47) + 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])) + 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]) @@ -833,7 +799,7 @@ def flush_scales( 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_INNER = 4, so rows stay u32-aligned, and a chunk's first + # 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( 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/cast/nvfp4/utils.py b/transformer_engine/common/CuTeDSL/utils_nvfp4.py similarity index 61% rename from transformer_engine/common/CuTeDSL/cast/nvfp4/utils.py rename to transformer_engine/common/CuTeDSL/utils_nvfp4.py index eaa306fa0ed..c930a3f6598 100644 --- a/transformer_engine/common/CuTeDSL/cast/nvfp4/utils.py +++ b/transformer_engine/common/CuTeDSL/utils_nvfp4.py @@ -2,81 +2,30 @@ # # See LICENSE for license information. -"""Constants, scaling math, and E2M1 conversion helpers for NVFP4 kernels.""" - -from typing import Optional +"""NVFP4 scaling math and E2M1 conversion helpers for the CuTeDSL kernels.""" import cutlass from cutlass import cute -from cutlass import BFloat16, Boolean, Float32, Float8E4M3FN, Int32, Int64, Uint32, Uint64 -from cutlass._mlir.dialects import arith as mlir_arith +from cutlass import BFloat16, Float32, Float8E4M3FN, Int64, Uint32 from cutlass._mlir.dialects import llvm from cutlass.cutlass_dsl import T, dsl_user_op -FLOAT32_MAX = 3.4028234663852886e38 -BFLOAT16_MAX = 3.3895313892515355e38 -FLOAT8E4M3_MAX = 448.0 -FLOAT4E2M1_MAX = 6.0 - -# Number of elements per NVFP4 scale block (they share one E4M3 scale factor). -NVFP4_BLOCK_SCALING_SIZE = 16 -# Input row/col divisibility the CUDA tuned-1D kernel requires (16B TMA alignment). -NVFP4_SHAPE_ALIGNMENT = 32 -# Padding TE applies to the scale tensors' outer / inner dim (NVFP4Quantizer::get_scale_shape). -NVFP4_SCALE_PAD_OUTER = 128 -NVFP4_SCALE_PAD_INNER = 4 - - -# Runs if CUTE_DSL_ENABLE_ASSERTIONS=1 or --enable-assertions present in cute.compile -def validate_tensor(tensor: Optional[cute.Tensor], expected_layout: cute.Layout, expected_dtype): - """Assert at trace time that an optional tensor has the expected layout and dtype.""" - if tensor is None: - return - cute.testing.assert_(tensor.layout == expected_layout, "Tensor layout does not match") - cute.testing.assert_(tensor.element_type == expected_dtype, "Tensor dtype does not match") - - -@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], - ) +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, +) - return prmt_u32 +# Largest finite E2M1 value. +FP4E2M1_MAX_NORM = 6.0 @cute.jit @@ -87,7 +36,7 @@ def compute_global_encode_sf(global_amax: Float32) -> Float32: 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 = FLOAT8E4M3_MAX * FLOAT4E2M1_MAX / global_amax + 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) @@ -96,7 +45,7 @@ def compute_global_encode_sf(global_amax: Float32) -> Float32: @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 / FLOAT4E2M1_MAX)) + 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) @@ -113,12 +62,11 @@ def compute_block_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 - elif cutlass.const_expr(sf_type == BFloat16): + 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) - else: - raise ValueError("Unsupported scaling-factor type. Only FP32 and BF16 are supported.") + raise ValueError("Unsupported scaling-factor type. Only FP32 and BF16 are supported.") @dsl_user_op @@ -365,120 +313,7 @@ def mul_cvt_bf16x8_to_fp4x8_sr( ) -# --- Register-resident kernels: global memory access by address --- -# -# The register 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 below 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 unpack_u64(v: Int64): - """The (low, high) words of a 64-bit value.""" - return cute.arch.inline_ptx( - "mov.b64 {{$w0}, {$w1}}, {$r0};", write_only_types=[Uint32, Uint32], read_only_args=[v] - ) - - -@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), - ) - - -# --- Register-resident kernels: the numeric core --- +# --- 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 @@ -486,59 +321,11 @@ def st_global_v8(addr: Int64, words): # bit-identical to the CUDA kernels. -@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] - ) - - -@cute.jit -def e4m3_code(x: Float32) -> Uint32: - """The E4M3 byte of x (round to nearest, saturating to 448), as static_cast does. - Its clamp to FLT_MAX in compute_decoding_scaling_factor cannot change the byte: the only - value it alters, +inf, saturates to the same 448.""" - return cute.arch.inline_ptx( - "{ .reg .b16 h; cvt.rn.satfinite.e4m3x2.f32 h, {$r0}, {$r0}; " - "cvt.u32.u16 {$w0}, h; and.b32 {$w0}, {$w0}, 255; }", - write_only_types=[Uint32], - read_only_args=[x], - ) - - -@cute.jit -def e4m3_value(code: Uint32) -> Float32: - """The value of the E4M3 byte `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], - ) - - @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(e4m3_value(code) * S_dec), Float32(FLOAT32_MAX)) - - -@cute.jit -def splat_f32x2(x: Float32) -> Int64: - """x in both halves of a packed f32 pair, as mul.f32x2 consumes it.""" - return cute.arch.inline_ptx( - "mov.b64 {$w0}, {{$r0}, {$r0}};", write_only_types=[Int64], read_only_args=[x] - ) + return cute.arch.fmin(rcp_rn_f32(cvt_fp8e4m3_code_to_f32(code) * S_dec), Float32(FLOAT32_MAX)) @cute.jit @@ -570,24 +357,6 @@ def scale_column_bf16(a: Uint32, b: Uint32, coeff2: Int64, LOW: cutlass.Constexp ) -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).""" - # Imported here: utils.py of the package root is imported by every CuTeDSL module. - from transformer_engine.common.CuTeDSL.utils import ( # pylint: disable=import-outside-toplevel - abs_max_x2_bf16, - ) - - 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] - - @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 @@ -595,7 +364,7 @@ def block_scale_code(words, S_enc_div6: Float32) -> Uint32: 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 e4m3_code(folded.bitcast(Float32) * S_enc_div6) + return cvt_f32_to_fp8e4m3_code(folded.bitcast(Float32) * S_enc_div6) @cute.jit