From abde002091453ca50ac2f0506db5741a4517e573 Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Fri, 7 Aug 2026 23:32:17 +0000 Subject: [PATCH 1/6] [PyTorch] Enable e5m3 fused GEMM kernels from cuDNN Signed-off-by: Kaining Zhong --- tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py | 109 +++++ tests/pytorch/test_fusible_ops.py | 85 +++- .../pytorch/cpp_extensions/gemm.py | 328 ++++++++++++++- .../pytorch/ops/fused/grouped_mlp.py | 380 ++++++++++++++++-- 4 files changed, 864 insertions(+), 38 deletions(-) diff --git a/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py b/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py index 639b2f752e..40c7c3bf82 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py +++ b/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py @@ -690,3 +690,112 @@ def test_nvfp4_row_scaled_gemm_matches_emulated( use_4over6=use_4over6, nvfp4_4over6_err_mode=nvfp4_4over6_err_mode, ) + + +def _check_ue5m3_gemm_versus_dequantized( + M, K, N, x_columnwise, w_columnwise, disable_second_level_scale +): + """Run an NVFP4/UE5M3 GEMM and compare against a dequantized FP32 reference.""" + if M % 256 != 0: + pytest.skip( + "cuDNN's grouped GEMM pads every group to 256 rows, so the UE5M3 path (which " + "routes there while cuBLAS lacks UE5M3 kernels) requires M % 256 == 0." + ) + torch.manual_seed(0) + device, dtype, out_dtype = "cuda", torch.bfloat16, torch.bfloat16 + x_shape = (K, M) if x_columnwise else (M, K) + w_shape = (K, N) if w_columnwise else (N, K) + x = torch.randn(x_shape, dtype=dtype, device=device) + w = torch.randn(w_shape, dtype=dtype, device=device) + + common = dict( + fp4_dtype=tex.DType.kFloat4E2M1, + scale_dtype=tex.DType.kFloat8UE5M3, + rowwise=True, + columnwise=True, + with_amax_reduction=False, + amax_reduction_group=None, + with_rht=False, + with_post_rht_amax=False, + ) + # disable_second_level_scale is given per operand, as (x, w). + xq = NVFP4Quantizer(**common, disable_second_level_scale=disable_second_level_scale[0]) + wq = NVFP4Quantizer(**common, disable_second_level_scale=disable_second_level_scale[1]) + x_q = xq.update_quantized(x, xq.make_empty(x_shape, dtype=dtype, device=device)) + w_q = wq.update_quantized(w, wq.make_empty(w_shape, dtype=dtype, device=device)) + + if disable_second_level_scale[0]: + assert x_q._amax_rowwise is None, "disable_second_level_scale should drop the amax" + if disable_second_level_scale[1]: + assert w_q._amax_rowwise is None, "disable_second_level_scale should drop the amax" + + # Reference: dequantize the orientation each operand is actually read in. + x_ref = _dequantize_nvfp4_usage(x_q, columnwise=x_columnwise) + w_ref = _dequantize_nvfp4_usage(w_q, columnwise=w_columnwise) + # _dequantize_nvfp4_usage returns each operand canonically as (rows, K), so + # the reference is the same expression for every layout. + ref = x_ref @ w_ref.t() + + if x_columnwise: + x_q.update_usage(rowwise_usage=False) + if w_columnwise: + w_q.update_usage(rowwise_usage=False) + transa, transb = not w_columnwise, x_columnwise + layout = ("T" if transa else "N") + ("T" if transb else "N") + y = general_gemm(w_q, x_q, out_dtype=out_dtype, layout=layout)[0] + + # Both sides see identically quantized operands, so quantization error cancels and + # only accumulation order and the bf16 output rounding differ. One bf16 ulp is + # already ~4e-3 relative, which no elementwise tolerance survives, so compare the + # whole result instead. + rel_err = (y.float() - ref).norm() / ref.norm() + assert rel_err < 5e-3, f"relative error {rel_err:.2e} is too large" + +ue5m3_available, reason_for_no_ue5m3 = te.is_fp8_ue5m3_available(return_reason=True) + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +@pytest.mark.skipif(not ue5m3_available, reason=reason_for_no_ue5m3) +@pytest.mark.parametrize( + "M, K, N", + [ + (256, 128, 256), + (256, 256, 256), + (256, 1024, 256), + (1024, 1024, 1024), + (4096, 512, 3072), + (112, 128, 96), + (304, 640, 304), + (1008, 3072, 992), + (256, 64, 256), + (128, 128, 112), + ], +) +@pytest.mark.parametrize( + "x_columnwise, w_columnwise", + [ + (False, False), # TN -- w rowwise, x rowwise (fprop) + (False, True), # NN -- w colwise, x rowwise (dgrad) + (True, True), # NT -- w colwise, x colwise (wgrad) + ], ids=["FF", "FT", "TT"] +) +@pytest.mark.parametrize( + "disable_second_level_scale", [ + (True, False), + ], ids=["TF"] +) +def test_nvfp4_ue5m3_gemm_versus_reference( + M: int, + K: int, + N: int, + x_columnwise: bool, + w_columnwise: bool, + disable_second_level_scale: bool, +): + """NVFP4 GEMM with UE5M3 block scales, with and without second-level scaling. + + UE5M3's wider range is what makes dropping the per-tensor global scale + viable, so both configurations must match the dequantized reference. + """ + _check_ue5m3_gemm_versus_dequantized( + M, K, N, x_columnwise, w_columnwise, disable_second_level_scale + ) diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py index cc8ded1cfb..7b8e77c02d 100644 --- a/tests/pytorch/test_fusible_ops.py +++ b/tests/pytorch/test_fusible_ops.py @@ -99,6 +99,7 @@ def maybe_skip_quantization( dims: Optional[Iterable[int] | int] = None, device: Optional[torch.device | str] = None, dtype: Optional[torch.dtype] = None, + fused_grouped_gemm: bool = False, ) -> None: """Skip test case if a quantization scheme is not supported""" @@ -136,6 +137,14 @@ def maybe_skip_quantization( elif quantization in nvfp4_variant_names: if math.prod(dims[:-1]) % 16 != 0 or dims[-1] % 16 != 0: pytest.skip("NVFP4 GEMMs require dims that are divisible by 16") + if quantization == "nvfp4_rht_ue5m3_scales" and not fused_grouped_gemm and ( + math.prod(dims[:-1]) % 256 != 0 or dims[-1] % 256 != 0 + ): + # UE5M3 has no cuBLAS kernels, so these GEMMs route to cuDNN's grouped + # kernel, which pads every group to 256 rows. Both dims are constrained, + # not just the leading one: wgrad contracts over the tokens, putting the + # feature dim in the kernel's ragged position. + pytest.skip("NVFP4 UE5M3 GEMMs require dims that are divisible by 256") # Check dtype if dtype is not None: @@ -143,6 +152,31 @@ def maybe_skip_quantization( pytest.skip("NVFP4 quantization is only supported with BF16 data") +def maybe_resize_for_ue5m3( + quantization: Optional[str], + weight_shape: tuple[int, int], + in_shape: Iterable[int], +) -> tuple[tuple[int, int], Iterable[int]]: + """Grow a test case's shapes to what the NVFP4 UE5M3 GEMM path can run. + + UE5M3 has no cuBLAS kernels, so its GEMMs route to cuDNN's grouped kernel, + which pads every group to 256 rows. Every dim is constrained, not just the + leading one: wgrad contracts over the tokens, so the feature dim also lands + in the kernel's ragged position. Tests share small default shapes to stay + fast, so rescale them for this recipe instead of losing the coverage to a + skip. Leading dims are collapsed to a single 256-row dim, keeping any extra + dims so the >2D cases still exercise the leading-dim handling. + """ + if quantization != "nvfp4_rht_ue5m3_scales": + return weight_shape, in_shape + leading = list(in_shape)[:-1] + if len(leading) == 0: + # A 1D input is a single row, which can never reach 256. + pytest.skip("NVFP4 UE5M3 GEMMs need at least 256 rows, so 1D inputs cannot run") + new_leading = [1] * (len(leading) - 1) + [256] + return (256, 256), (*new_leading, -1) + + @torch.no_grad() def make_reference_and_test_tensors( shape: int | Iterable[int], @@ -936,6 +970,8 @@ def _test_basic_linear( ) -> None: """Helper function for tests with GEMM""" + weight_shape, in_shape = maybe_resize_for_ue5m3(quantization, weight_shape, in_shape) + # Make input and weight shapes consistent out_features, in_features = weight_shape in_shape = list(in_shape)[:-1] + [in_features] @@ -944,6 +980,9 @@ def _test_basic_linear( # Skip invalid configurations maybe_skip_quantization(quantization, dims=in_shape, device=device, dtype=dtype) maybe_skip_quantization(quantization, dims=out_shape) + if quantization == "nvfp4_rht_ue5m3_scales" and accumulate_into_main_grad: + # The cuDNN kernel UE5M3 routes to writes its output, with no beta term. + pytest.skip("NVFP4 UE5M3 GEMMs cannot accumulate into an existing output") quantization_needed = any( ( quantized_compute, @@ -2145,6 +2184,11 @@ def test_grouped_linear( pytest.skip("Quantized group GEMM is only supported with BF16/FP16") if quantization == "nvfp4_4over6": pytest.skip("NVFP4 4over6 grouped quantization is not supported") + if quantization == "nvfp4_rht_ue5m3_scales": + # A standalone GroupedLinear issues grouped cuBLAS GEMMs, which have no + # UE5M3 kernels. UE5M3 grouped support is limited to the CuteDSL fused + # grouped MLP, which routes its GEMMs to cuDNN instead. + pytest.skip("UE5M3 scales are not supported by grouped cuBLAS GEMMs") # Random data x_ref, x_test = make_reference_and_test_tensors( @@ -2568,6 +2612,8 @@ def test_forward_linear_bias_activation( ) -> None: """Forward GEMM + bias + activation""" + weight_shape, in_shape = maybe_resize_for_ue5m3(quantization, weight_shape, in_shape) + # Make input and weight shapes consistent out_features, in_features = weight_shape in_shape = list(in_shape)[:-1] + [in_features] @@ -3557,6 +3603,7 @@ def test_layernorm_mlp( @pytest.mark.parametrize("bias", (False, True)) @pytest.mark.parametrize("dtype", _dtypes) @pytest.mark.parametrize("quantization", _quantization_list) + @pytest.mark.parametrize("glu_interleave_size", (None, 32)) def test_grouped_mlp( self, *, @@ -3568,6 +3615,7 @@ def test_grouped_mlp( device: torch.device = "cuda", split_alignment: int = 256, activation: str = "scaled_swiglu", + glu_interleave_size: Optional[int], ) -> None: """GroupedLinear + scaled activation + GroupedLinear""" @@ -3588,11 +3636,30 @@ def test_grouped_mlp( # Skip invalid configurations with_quantization = quantization is not None - maybe_skip_quantization(quantization, dims=in_shape, device=device, dtype=dtype) + maybe_skip_quantization( + quantization, + dims=in_shape, + device=device, + dtype=dtype, + fused_grouped_gemm=True, + ) if with_quantization and dtype not in (torch.bfloat16, torch.float16): pytest.skip("Quantized group GEMM is only supported with BF16/FP16") if activation == "scaled_srelu" and quantization == "nvfp4_rht" and bias: pytest.skip("NVFP4 RHT SReLU grouped MLP coverage is limited to no-bias") + if ( + quantization == "nvfp4_rht_ue5m3_scales" + and activation == "scaled_swiglu" + and glu_interleave_size is None + ): + # Without interleaving the GLU pattern never matches the CuteDSL fused + # grouped MLP, so the ops fall back to grouped cuBLAS, which has no UE5M3 + # kernels. The glu_interleave_size=32 variant covers this recipe. + pytest.skip("UE5M3 grouped MLP is only supported on the fused path") + if quantization == "nvfp4_rht_ue5m3_scales" and activation == "scaled_srelu": + # The fuser refuses UE5M3 with SReLU outright, since cuDNN has no UE5M3 + # SReLU kernel, so this would fall back to grouped cuBLAS. + pytest.skip("cuDNN has no UE5M3 SReLU kernel, so this cannot use the fused path") # Random data x_ref, x_test = make_reference_and_test_tensors( @@ -3678,7 +3745,18 @@ def test_grouped_mlp( x, fc1_ws_ref[group_idx], bias=fc1_bs_ref[group_idx] ) if activation == "scaled_swiglu": - act_in1, act_in2 = fc1_out.chunk(2, dim=-1) + if glu_interleave_size is not None: + # Undo the interleaving so the two GLU halves can be chunked. + act_in = fc1_out.reshape( + -1, + fc1_out_features // (2 * glu_interleave_size), + 2, + glu_interleave_size, + ) + act_in = act_in.transpose(1, 2).reshape(fc1_out.shape) + else: + act_in = fc1_out + act_in1, act_in2 = act_in.chunk(2, dim=-1) act_out = torch.nn.functional.silu(act_in1) * act_in2 elif activation == "scaled_srelu": act_out = torch.nn.functional.relu(fc1_out).square() @@ -3713,7 +3791,7 @@ def test_grouped_mlp( scale_bias=bias, ) if activation == "scaled_swiglu": - activation_op = te_ops.ScaledSwiGLU() + activation_op = te_ops.ScaledSwiGLU(glu_interleave_size=glu_interleave_size) elif activation == "scaled_srelu": activation_op = te_ops.ScaledSReLU() else: @@ -3763,6 +3841,7 @@ def test_grouped_mlp_nvfp4_rht_srelu( quantization="nvfp4_rht", device=device, activation="scaled_srelu", + glu_interleave_size=None, # SReLU is not a GLU, so interleaving does not apply ) diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index f3d97b7269..b6435efa29 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -4,13 +4,13 @@ """Python interface for GEMM extensions""" -from typing import Iterable, Literal, Optional, Tuple, Union, List +from typing import Callable, Iterable, Literal, Optional, Tuple, Union, List import os import functools import torch import transformer_engine_torch as tex -from ..constants import TE_DType, DType -from ..utils import get_sm_count, _empty_tensor +from ..constants import MXFP8_BLOCK_SCALING_SIZE, NVFP4_BLOCK_SCALING_SIZE, TE_DType, DType +from ..utils import ceil_div, get_cached_ones_tensor, get_sm_count, _empty_tensor from ..quantized_tensor import QuantizedTensorStorage, Quantizer from ..tensor.float8_blockwise_tensor import Float8BlockQuantizer @@ -25,6 +25,7 @@ from ..tensor.storage.nvfp4_tensor_storage import NVFP4TensorStorage from ..tensor.utils import is_custom from ..custom_recipes.gemm import custom_gemm +from ..ops._common import validate_or_alloc_output from ...debug.pytorch.debug_quantization import DebugQuantizedTensor, DebugQuantizer __all__ = [ @@ -188,6 +189,294 @@ def _validate_native_gemm_output_quantizer(quantization_params): ) +@functools.lru_cache(maxsize=None) +def grouped_gemm_quant_kernel() -> Callable: + """cuDNN CuTe DSL grouped GEMM kernel for block-scaled inputs.""" + from cudnn import grouped_gemm_quant_wrapper_sm100 # pylint: disable=no-name-in-module + + return grouped_gemm_quant_wrapper_sm100 + + +def convert_TE_MX_tensor_to_cuDNN_operand( + data: torch.Tensor, + scale_inv: torch.Tensor, + *, + data_dtype: torch.dtype, + scale_dtype: torch.dtype, + valid_M_or_N: int, + k_logical: int, + L: int = 1, + sf_swizzled: bool = False, + use_N_major_for_B: bool = False, +) -> tuple[torch.Tensor, torch.Tensor]: + """Reshape an plain buffer into the layout cuDNN's grouped GEMM expects. + + cuDNN requirements: + A: (valid_m, K, 1), K-major + B: (N, K, L), K-major (FP8 also supports N-major) + + SFA: (32, 4, ceil(valid_m/128), 4, ceil(ceil(K/sf_vec_size)/4), 1) + SFB: (32, 4, ceil(N/128), 4, ceil(ceil(K/sf_vec_size)/4), L) + + whereas TE stores flat buffers which can be intepreted as contiguous tensors + with the following layouts: + + Note: K_packed is K/2 for FP4 (two values per byte) and K for FP8 + + A (K-major): (1, valid_m, K_packed) + B (K-major): (L, N, K_packed) -- used for FP4 only now + B (N-major): (L, K, N) -- used for FP8 only now + + SFA (unswizzled): (1, ceil(valid_m/128), 4, 32, ceil(ceil(K/sf_vec_size)/4), 4) + SFB (unswizzled): (L, ceil(N/128), 4, 32, ceil(ceil(K/sf_vec_size)/4), 4) + SFA (swizzled): (1, ceil(valid_m/128), ceil(ceil(K/sf_vec_size)/4), 32, 4, 4) + SFB (swizzled): (L, ceil(N/128), ceil(ceil(K/sf_vec_size)/4), 32, 4, 4) + """ + + if use_N_major_for_B: + assert data_dtype in (torch.float8_e4m3fn, torch.float8_e5m2), \ + f"Using N-major layout for B is only supported for FP8, but got {data_dtype}." + + available_scalings = { + # NVFP4 recipe (UE5M3 rides as E4M3 since torch has no ue5m3 dtype) + (torch.float4_e2m1fn_x2, torch.float8_e4m3fn): NVFP4_BLOCK_SCALING_SIZE, + # MXFP8 recipe + (torch.float8_e4m3fn, torch.float8_e8m0fnu): MXFP8_BLOCK_SCALING_SIZE, + } + assert (data_dtype, scale_dtype) in available_scalings, ( + "Unsupported (data_dtype, scale_dtype) pair for a cuDNN block-scaled operand: " + f"({data_dtype}, {scale_dtype}). Expected NVFP4 (float4_e2m1fn_x2, " + "float8_e4m3fn) or MXFP8 (float8_e4m3fn, float8_e8m0fnu)." + ) + sf_vec_size = available_scalings[(data_dtype, scale_dtype)] + + k_sf_tiles = ceil_div(k_logical, 4 * sf_vec_size) + + if data_dtype == torch.float4_e2m1fn_x2: + k_packed = k_logical // 2 # fp4 packs two values per byte + else: + k_packed = k_logical # fp8 packs one value per byte + + data = data.view(dtype=data_dtype) + if use_N_major_for_B: + # B is stored untransposed, i.e. (L, K, N); permuting to (N, K, L) leaves + # stride 1 on N. Only FP8 accepts this, asserted above. + data = data.view(L, k_packed, valid_M_or_N) + data = data.permute(2, 1, 0) + else: + # (L, N, K) -> (N, K, L), stride 1 on K. + data = data.view(L, valid_M_or_N, k_packed) + data = data.permute(1, 2, 0) + + if sf_swizzled: + scale_inv = scale_inv.view(dtype=scale_dtype) + scale_inv = scale_inv.view( + L, + ceil_div(valid_M_or_N, 128), + k_sf_tiles, + 32, + 4, + 4, + ) + scale_inv = scale_inv.permute(3, 4, 1, 5, 2, 0) + return data, scale_inv + + scale_inv = scale_inv.view(dtype=scale_dtype) + scale_inv = scale_inv.view( + L, + ceil_div(valid_M_or_N, 128), + 4, + 32, + k_sf_tiles, + 4, + ) + scale_inv = scale_inv.permute(3, 2, 1, 5, 4, 0) + return data, scale_inv + + +def general_cuDNN_MX_gemm( + A: torch.Tensor, + B: torch.Tensor, + out_dtype: Optional[torch.dtype] = None, + quantization_params: Optional[Quantizer] = None, + gelu: bool = False, + gelu_in: torch.Tensor = None, + alpha: float = 1.0, + beta: Optional[float] = None, + accumulate: bool = False, + layout: str = "TN", + out: Optional[torch.Tensor] = None, + bias: Optional[torch.Tensor] = None, + use_split_accumulator: bool = False, + grad: bool = False, + ub: Union[tex.CommOverlap, tex.CommOverlapP2P] = None, + ub_type: tex.CommOverlapType = None, + extra_output: Optional[torch.Tensor] = None, + bulk_overlap: bool = False, +) -> Iterable[Optional[torch.Tensor]]: + """Perform GEMM via cuDNN kernels + + The parameters passed are in cuBLAS notation, where + D = alpha * op(B) @ op(A) + beta * C, where the shape is always + (M, N) = (M, K) @ (K, N) + (M, N) + + B: + - "N" is (M, K), which is always TE's rowwise data, and op(B) is B + - "T" is (K, M), which is always TE's colwise data, and op(B) is B.T + A + - "N" is (K, N), which is always TE's colwise data, and op(A) is A + - "T" is (N, K), which is always TE's rowwise data, and op(A) is A.T + + Note: layout string means layout of "A" and "B" respectively. + + We use cuDNN-frontend's grouped_gemm_quant_wrapper_sm100 API here which is supposed to be a grouped GEMM + but here we set groups = 1 so it is effectively a single GEMM. + """ + assert isinstance(A, NVFP4TensorStorage) and isinstance(B, NVFP4TensorStorage) and \ + A.get_metadata()["scale_dtype"] == DType.kFloat8UE5M3 and B.get_metadata()["scale_dtype"] == DType.kFloat8UE5M3, \ + f"cuDNN MX GEMM is only used for NVFP4 GEMM with e5m3 scale factors for now." + + assert quantization_params is None, "cuDNN GEMM currently does not support output quantization." + assert gelu is False and gelu_in is None, "cuDNN GEMM currently does not support fused GELU." + assert accumulate is False, "cuDNN GEMM currently does not support accumulation." + # `grad` only changes behaviour when a bias is supplied: it turns the bias slot + # into a bias-gradient output, which cuDNN has no epilogue for. Backward GEMMs + # that pass grad=True without a bias need nothing special. + assert not (grad and bias is not None), ( + "cuDNN GEMM currently does not support fused bias gradient." + ) + # use_split_accumulator is deliberately not checked: it is a cuBLAS knob for + # raising accumulator precision, and the cuDNN kernel always accumulates in + # FP32, so the request is already satisfied either way. + assert ub is None and ub_type is None, "cuDNN GEMM currently does not support CommOverlap." + assert extra_output is None, "cuDNN GEMM currently does not support extra output." + assert bulk_overlap is False, "cuDNN GEMM currently does not support bulk overlap." + + assert layout in ("TN", "NN", "NT"), f"GEMM layout {layout} not supported." + transa = layout[0] == "T" + transb = layout[1] == "T" + + assert out_dtype in (torch.float32, torch.float16, torch.bfloat16), \ + f"cuDNN MX GEMM currently only supports float32, float16, and bfloat16 outputs, but got {out_dtype}." + + device = A.device + + # cuDNN only accepts GEMM-swizzled scale factors -- an unswizzled buffer is + # rejected on its strides -- so swizzle first if the quantizer did not + # (optimize_for_gemm defaults to False). This mirrors what the cuBLAS path + # does in C++ via swizzle_scales_for_gemm. The call is in-place, swizzles + # both orientations, and no-ops when the tensor is already swizzled. + if not A._with_gemm_swizzled_scales: + tex.swizzle_scales_for_gemm_(A) + if not B._with_gemm_swizzled_scales: + tex.swizzle_scales_for_gemm_(B) + + # Pick the buffer whose block scales run along K. In every case the selected + # buffer is physically (rows, K_packed), so the reshape below is uniform. + # LHS is always (M, K) + if transb: + dataB, sfB, amaxB = B._columnwise_data, B._columnwise_scale_inv, B._amax_columnwise + else: + dataB, sfB, amaxB = B._rowwise_data, B._rowwise_scale_inv, B._amax_rowwise + # RHS is always (K, N) + if transa: + dataA, sfA, amaxA = A._rowwise_data, A._rowwise_scale_inv, A._amax_rowwise + else: + dataA, sfA, amaxA = A._columnwise_data, A._columnwise_scale_inv, A._amax_columnwise + + # Find the logical shapes (not the physical shapes where K is packed with 2 fp4 stored in 1 byte). + M, K = dataB.numel() // dataB.shape[-1], dataB.shape[-1] * 2 + N, k_from_a = dataA.numel() // dataA.shape[-1], dataA.shape[-1] * 2 + assert K == k_from_a, f"Contraction dims disagree: A implies {k_from_a}, B implies {K}." + + # The output keeps B's leading dims. They only exist when B is read row-wise: + # the column-wise NVFP4 buffer is physically transposed and always 2D, and the + # layouts that select it give B the logical shape (K, M). + out_shape = (*dataB.shape[:-1], N) if not transb else (M, N) + + # cuDNN's own operand names are the other way round: its "a" is the (M, K) + # activation-like operand (TE's B) and its "b" is the (N, K) weight-like one + # (TE's A). + cudnn_a, cudnn_sfa = convert_TE_MX_tensor_to_cuDNN_operand( + dataB, + sfB, + data_dtype=torch.float4_e2m1fn_x2, + scale_dtype=torch.float8_e4m3fn, # e5m3 rides as e4m3; torch has no ue5m3 + valid_M_or_N=M, + k_logical=K, + L=1, + sf_swizzled=True, # ensured above + ) + cudnn_b, cudnn_sfb = convert_TE_MX_tensor_to_cuDNN_operand( + dataA, + sfA, + data_dtype=torch.float4_e2m1fn_x2, + scale_dtype=torch.float8_e4m3fn, # e5m3 rides as e4m3; torch has no ue5m3 + valid_M_or_N=N, + k_logical=K, + L=1, + sf_swizzled=True, # ensured above + ) + + # Row-scaled NVFP4 stores one amax per row instead of one per tensor, which + # this path cannot express; general_gemm handles that mode separately. + for name, amax in (("A", amaxA), ("B", amaxB)): + assert amax is None or amax.numel() == 1, ( + f"cuDNN MX GEMM expects a per-tensor amax for {name}, but got {amax.numel()} " + "values. Row-scaled NVFP4 is not supported on this path." + ) + + # Prepare alpha. cuDNN applies the block scales but not TE's per-tensor global + # scale, so alpha carries the product of both operands'. A tensor quantized + # without second-level scaling has no amax and contributes a factor of one. + + # general_gemm normalizes alpha/beta only on its cuBLAS path, downstream of the + # dispatch here, so callers can still reach this with alpha=None meaning one. + alpha = validate_gemm_scale(alpha, True) + validate_gemm_scale(beta, accumulate) + nvfp4_global_scale = 6.0 * 114688.0 + ones = get_cached_ones_tensor(1, dtype=torch.float32, device=device) + scaleA = ones if amaxA is None else amaxA.to(torch.float32).reshape(1) / nvfp4_global_scale + scaleB = ones if amaxB is None else amaxB.to(torch.float32).reshape(1) / nvfp4_global_scale + alpha_tensor = (alpha * scaleA * scaleB).to(torch.float32) + + if bias is not None: + assert bias.dim() == 1 and bias.shape[0] == N, ( + f"cuDNN MX GEMM expects a ({N},) bias, but got {tuple(bias.shape)}." + ) + # cuDNN checks the stride literally, so (1, N) rather than reshape's (1, 1). + bias = bias.contiguous().as_strided((N, 1), (1, N)) + + # Prepare for output + out = validate_or_alloc_output(out, out_shape, out_dtype, device) + d_tensor = out.view(M, N).as_strided((M, N, 1), (N, 1, M * N)) + + gemm_kwargs = { + "a_tensor": cudnn_a, + "sfa_tensor": cudnn_sfa, + "b_tensor": cudnn_b, + "sfb_tensor": cudnn_sfb, + # One group, so the only padded end offset is the full row count. + "padded_offsets": torch.tensor([M], dtype=torch.int32, device=device), + "alpha_tensor": alpha_tensor, + "bias_tensor": bias, + "norm_const_tensor": None, # must be None for FP4 inputs + "acc_dtype": torch.float32, + "d_dtype": out_dtype, # high precision -> no output quantization + "d_tensor": d_tensor, + "cd_major": "n", # only "n" is supported by cuDNN + "sf_vec_size": NVFP4_BLOCK_SCALING_SIZE, # Hardcode to NVFP4 for now + "sf_fp8_dtype_override": "e5m3", # Hardcode for now + "current_stream": torch.cuda.current_stream().cuda_stream, + "discrete_col_sfd": False, + "use_dynamic_sched": True, + } + grouped_gemm_quant_kernel()(**gemm_kwargs) + + # Matches general_gemm's contract: (out, bias_grad, gelu_input, extra_output). + return out, None, None, None + + def general_gemm( A: torch.Tensor, B: torch.Tensor, @@ -210,6 +499,39 @@ def general_gemm( ) -> Iterable[Optional[torch.Tensor]]: """GEMM supporting fp8 inputs.""" + route_to_cuDNN = False + # Route NVFP4 GEMM with e5m3 scale factors to cuDNN since cuBLAS is not ready yet. + # Test against the storage class, not NVFP4Tensor: the ops and module paths hand + # this function bare NVFP4TensorStorage operands, and NVFP4Tensor subclasses it. + if isinstance(A, NVFP4TensorStorage) and isinstance(B, NVFP4TensorStorage): + if ( + A.get_metadata()["scale_dtype"] == DType.kFloat8UE5M3 + and B.get_metadata()["scale_dtype"] == DType.kFloat8UE5M3 + ): + route_to_cuDNN = True + + if route_to_cuDNN: + return general_cuDNN_MX_gemm( + A, + B, + out_dtype, + quantization_params, + gelu, + gelu_in, + alpha, + beta, + accumulate, + layout, + out, + bias, + use_split_accumulator, + grad, + ub, + ub_type, + extra_output, + bulk_overlap, + ) + assert layout in ("TN", "NN", "NT"), f"GEMM layout {layout} not supported." transa = layout[0] == "T" transb = layout[1] == "T" diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index a44bef0b2d..8b03a0587c 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -10,7 +10,7 @@ import functools import os from importlib.metadata import PackageNotFoundError, version as get_pkg_version -from typing import Any, Optional +from typing import Any, Literal, Optional import torch from packaging.version import Version as PkgVersion @@ -20,6 +20,7 @@ from ...constants import MXFP8_BLOCK_SCALING_SIZE, NVFP4_BLOCK_SCALING_SIZE, TE_DType from ...cpu_offload import is_cpu_offload_enabled, mark_activation_offload, start_offload from ...cpp_extensions import general_gemm, general_grouped_gemm_for_grouped_tensor +from ...cpp_extensions.gemm import convert_TE_MX_tensor_to_cuDNN_operand from ...distributed_weight import ( is_distributed_weight, materialize_weight_for_forward, @@ -251,6 +252,56 @@ def _nvfp4_amax( return torch.cat([amax.view(-1) for amax in amaxes], dim=0) +# TODO(kainingz): remove this temporary workaround after pytorch & tvm-ffi supports e5m3 GEMM +def _nvfp4_sf_dtype_override(quantizer: Optional[Quantizer]) -> Literal["e5m3"] | None: + """Returns a string to indicate the real scale factor dtype for cuDNN. + + Since pytorch doesn't have a native e5m3 dtype, we need let e5m3 pretend to be e4m3 and + use this string to indicate cuDNN to interpret the scale factors as e5m3 correctly when + it enters CuTeDSL region which has e5m3 support. + """ + if quantizer is None or not isinstance(quantizer, NVFP4Quantizer): + return None + if getattr(quantizer, "nvfp4_use_4over6", False): + # We don't use e5m3 for 4over6 + return None + scale_dtype = getattr(quantizer, "scale_dtype", None) + if scale_dtype is not None and scale_dtype == tex.DType.kFloat8UE5M3: + return "e5m3" + # If we don't use e5m3 we don't need to pass this string to override + return None + + +def _nvfp4_scale_max(quantizer: Quantizer) -> float: + """Return the maximum representable magnitude of an NVFP4 scale factor.""" + # 4over6 might override e4m3's max to 256 over default 448 + override_max = getattr(quantizer, "nvfp4_e4m3_max", None) + # NVFP4Quantizer's initialization sets nvfp4_e4m3_max to -1 if no override + if override_max is not None and override_max != -1: + return float(override_max) + scale_dtype = getattr(quantizer, "scale_dtype", None) + if scale_dtype is not None and scale_dtype == tex.DType.kFloat8UE5M3: + return 114688.0 + return 448.0 + + +def _nvfp4_global_scale( + tensors: GroupedTensor | Iterable[NVFP4TensorStorage], + quantizer: Quantizer, + *, + columnwise: bool, + num_groups: int, + device: torch.device, +) -> torch.Tensor: + """Return the per-group global scale factor for an NVFP4 operand.""" + if getattr(quantizer, "disable_second_level_scale", False): + # The second-level scale is disabled, so the global scale is always 1.0. + return get_cached_ones_tensor(num_groups, torch.float32, device) + # 6.0 is NVFP4_FP4_MAX + denom = 6.0 * _nvfp4_scale_max(quantizer) + return _nvfp4_amax(tensors, columnwise=columnwise).to(torch.float32) / denom + + def _single_quantized_tensor_from_grouped( grouped: GroupedTensor, quantizer: Optional[MXFP8Quantizer | NVFP4Quantizer] = None, @@ -303,6 +354,7 @@ def _single_quantized_tensor_from_grouped( with_gemm_swizzled_scales=grouped._with_gemm_swizzled_scales, ) + # TODO(kainingz): claude told me this doesn't pass the required param scale_dtype. Should check this later return NVFP4Tensor( shape=shape, dtype=grouped.get_dtype(), @@ -475,6 +527,10 @@ def _cudnn_compute_wgrad( out_features, in_features = weight_shape total_tokens = grouped_dy.logical_shape[0] + device = grouped_dy.columnwise_data.device + + dy_quantizer=getattr(grouped_dy, "quantizer", None) + x_quantizer=getattr(grouped_x, "quantizer", None) sfa_leading_dim = round_up_to_nearest_multiple(out_features, 128) sfb_leading_dim = round_up_to_nearest_multiple(in_features, 128) @@ -483,7 +539,6 @@ def _cudnn_compute_wgrad( # A workaround for the case with zero-token experts. # Even for this case, cuteDSL still requires the same # stride requirements for the input and scale tensors. - device = grouped_dy.columnwise_data.device a_tensor = torch.empty_strided( (out_features, 0), (16, 1), @@ -554,9 +609,9 @@ def _cudnn_compute_wgrad( "current_stream": current_stream, } if use_nvfp4: - global_scale_denom = 448.0 * 6.0 + num_groups = offsets.shape[0] if total_tokens == 0: - global_scale_shape = (offsets.shape[0],) + global_scale_shape = (num_groups,) common_wgrad_kwargs["global_scale_a"] = torch.zeros( global_scale_shape, dtype=torch.float32, @@ -568,13 +623,24 @@ def _cudnn_compute_wgrad( device=device, ) else: - common_wgrad_kwargs["global_scale_a"] = ( - _nvfp4_amax(grouped_dy, columnwise=True).to(torch.float32) / global_scale_denom + common_wgrad_kwargs["global_scale_a"] = _nvfp4_global_scale( + grouped_dy, + dy_quantizer, + columnwise=True, + num_groups=num_groups, + device=device, ) - common_wgrad_kwargs["global_scale_b"] = ( - _nvfp4_amax(grouped_x, columnwise=True).to(torch.float32) / global_scale_denom + common_wgrad_kwargs["global_scale_b"] = _nvfp4_global_scale( + grouped_x, + x_quantizer, + columnwise=True, + num_groups=num_groups, + device=device, ) common_wgrad_kwargs["input_order"] = "tensor_ragged" + wgrad_sf_dtype_override = _nvfp4_sf_dtype_override(dy_quantizer) + if wgrad_sf_dtype_override is not None: + common_wgrad_kwargs["sf_fp8_dtype_override"] = wgrad_sf_dtype_override # Prepare wgrad output if single_grouped_weight: @@ -820,20 +886,20 @@ def fuse_grouped_mlp_ops( elif not (recipe.mxfp8() or recipe.nvfp4()): return ops + if activation_op_types is None: + activation_op_types = (ScaledSwiGLU, ScaledClampedQGeGLU) + # Check for unsupported NVFP4 recipe configs if recipe.nvfp4(): if recipe.disable_rht: # Graph-safe grouped quantize is only supported with RHT return ops - if ( - recipe.row_scaled_activation - or recipe.nvfp4_4over6 - or recipe.fp8_format == RecipeFormat.UE5M3 - ): + if recipe.row_scaled_activation or recipe.nvfp4_4over6 != "none": + # 4over6 doesn't used fused kernels + return ops + if recipe.fp8_format == RecipeFormat.UE5M3 and ScaledSReLU in activation_op_types: + # cuDNN has no SReLU support for UE5M3 for now return ops - - if activation_op_types is None: - activation_op_types = (ScaledSwiGLU, ScaledClampedQGeGLU) # Scan ops through with sliding window out = [] @@ -1341,16 +1407,24 @@ def fuser_forward( ) fc1_norm_const_tensor = None if use_nvfp4 else norm_const_tensor if use_nvfp4: - nvfp4_fp4_max = 6.0 - nvfp4_fp8_max = 448.0 - nvfp4_global_scale_denom = nvfp4_fp4_max * nvfp4_fp8_max # cuDNN receives NVFP4 block-scaled inputs without TE's per-group # global scale factors, so alpha supplies the product of the two # operand global scales. fc1_alpha_tensor = ( - _nvfp4_amax(grouped_fc1_x, columnwise=False) - * _nvfp4_amax(grouped_fc1_weight, columnwise=False) - / (nvfp4_global_scale_denom**2) + _nvfp4_global_scale( + grouped_fc1_x, + fc1_input_quantizer, + columnwise=False, + num_groups=num_groups, + device=device, + ) + * _nvfp4_global_scale( + grouped_fc1_weight, + fc1_weight_quantizer, + columnwise=False, + num_groups=num_groups, + device=device, + ) ).to(torch.float32) else: fc1_alpha_tensor = alpha_tensor @@ -1363,6 +1437,8 @@ def fuser_forward( and isinstance(fc2_input_quantizer, NVFP4Quantizer) and fc2_input_quantizer.with_rht and fc2_input_quantizer.with_post_rht_amax + # If we don't have the second-level scaling we don't need the post-RHT amax in the kernel. + and not fc2_input_quantizer.disable_second_level_scale ) activation_is_srelu = isinstance(activation_op, ScaledSReLU) activation_supports_hadamard = self._cudnn_act_func == "swiglu" or ( @@ -1389,6 +1465,13 @@ def fuser_forward( "current_stream": current_stream, "use_dynamic_sched": True, } + fc1_sf_dtype_override = _nvfp4_sf_dtype_override(fc1_input_quantizer) + # Only override the dtype if we are using e5m3 and not using the Hadamard kernel, + # since the Hadamard fused GEEM kernel does not support e5m3. + # At the time of writing, the e5m3 recipe doesn't have the second level scaling enabled, + # which naturally leads to use_fc1_act_hadamard=False + if fc1_sf_dtype_override is not None and not use_fc1_act_hadamard: + fc1_activation_kwargs["sf_fp8_dtype_override"] = fc1_sf_dtype_override if use_fc1_act_hadamard_srelu: fc1_activation_kwargs["act_func"] = "srelu" elif self._cudnn_act_func is not None: @@ -1527,7 +1610,8 @@ def fuser_forward( fc2_out_shape = in_shape[:-1] + [fc2_weight_shape[0]] fc2_scales = basic_op_extra_inputs[2][1] if fc2_op._scale_bias else None - if use_nvfp4: + fc2_input_sf_override = _nvfp4_sf_dtype_override(fc2_input_quantizer) + if use_nvfp4 and fc2_input_sf_override is None: fc2_bias_for_gemm = None fc2_bias_scale = None if fc2_bias_packed is not None: @@ -1595,6 +1679,120 @@ def fuser_forward( bias_scale=fc2_bias_scale, ) fc2_out = fc2_out_buf + elif use_nvfp4 and fc2_input_sf_override is not None: # TODO(kainingz): remove this e5m3 workaround once cuBLAS is ready. + fc2_in = fc1_kernel_out["d_tensor"] + fc2_in = fc2_in.view(in_shape[0], fc2_weight_shape[1]).contiguous() + fc2_input_quantizer.set_usage(rowwise=True, columnwise=weight_requires_grad) + fc2_input_quantizer.optimize_for_gemm = True + + if use_fc1_act_hadamard: # Currently unreachable since e5m3 doesn't use second-level scaling + grouped_fc2_x = _group_quantize_with_amax_for_grouped_mlp( + fc2_in, + fc2_input_quantizer, + num_groups, + split_sizes, + fc1_kernel_out["amax_tensor"].view(-1), + fc1_kernel_out["post_rht_amax_tensor"].view(-1), + tensor_offsets=fc2_x_tensor_offsets, + ) + else: + grouped_fc2_x = _group_quantize_for_grouped_mlp( + fc2_in, + fc2_input_quantizer, + num_groups, + split_sizes, + tensor_offsets=fc2_x_tensor_offsets, + ) + + fc2_x_data, fc2_x_scales = convert_TE_MX_tensor_to_cuDNN_operand( + grouped_fc2_x.rowwise_data, + grouped_fc2_x.scale_inv, + data_dtype=data_dtype, + scale_dtype=scale_view_dtype, + valid_M_or_N=in_shape[0], + k_logical=fc2_weight_shape[1], + sf_swizzled=grouped_fc2_x._with_gemm_swizzled_scales, + ) + + fc2_fwd_alpha_tensor = ( + _nvfp4_global_scale( + grouped_fc2_x, + fc2_input_quantizer, + columnwise=False, + num_groups=num_groups, + device=device, + ) + * _nvfp4_global_scale( + grouped_fc2_weight, + fc2_weight_quantizer, + columnwise=False, + num_groups=num_groups, + device=device, + ) + ).to(torch.float32) + + fc2_scales_tensor = ( + fc2_scales.detach().to(dtype=torch.float32).reshape(-1, 1, 1) + if fc2_scales is not None + else torch.ones((in_shape[0], 1, 1), dtype=torch.float32, device=device) + ) + fc2_quant_kwargs = { + "a_tensor": fc2_x_data, + "sfa_tensor": fc2_x_scales, + "padded_offsets": split_points, + "alpha_tensor": fc2_fwd_alpha_tensor, + "bias_tensor": fc2_bias_packed, + "norm_const_tensor": None, + "prob_tensor": fc2_scales_tensor, + "acc_dtype": torch.float32, + "d_dtype": dtype, + "cd_major": "n", + "sf_vec_size": sf_vec_size, + "sf_fp8_dtype_override": fc2_input_sf_override, + "current_stream": current_stream, + "use_dynamic_sched": True, + } + + if fc2_op.single_grouped_weight: + # Clone and swizzle scales for GEMM (original stays unmodified + # for save_for_backward). + fc2_weight_for_gemm = grouped_fc2_weight.copy() + tex.grouped_swizzle_for_gemm(fc2_weight_for_gemm, rowwise=True, columnwise=False) + + fc2_w_data, fc2_w_scales = convert_TE_MX_tensor_to_cuDNN_operand( + fc2_weight_for_gemm.rowwise_data, + fc2_weight_for_gemm.scale_inv, + data_dtype=data_dtype, + scale_dtype=scale_view_dtype, + valid_M_or_N=fc2_weight_shape[0], + k_logical=fc2_weight_shape[1], + L=num_groups, + sf_swizzled=fc2_weight_for_gemm._with_gemm_swizzled_scales, + ) + fc2_quant_kwargs["b_tensor"] = fc2_w_data + fc2_quant_kwargs["sfb_tensor"] = fc2_w_scales + else: + fc2_b_ptrs, fc2_sfb_ptrs, _fc2_sfb_buffer = ( + tex.grouped_mlp_experimental.swizzle_scales_and_pack_ptrs_for_discrete_weights( + [w._rowwise_data for w in grouped_fc2_weight], + [w._rowwise_scale_inv for w in grouped_fc2_weight], + "nvfp4", + device, + ) + ) + fc2_quant_kwargs["b_ptrs"] = fc2_b_ptrs + fc2_quant_kwargs["sfb_ptrs"] = fc2_sfb_ptrs + fc2_quant_kwargs["n"] = fc2_weight_shape[0] + fc2_quant_kwargs["b_dtype"] = data_dtype + fc2_quant_kwargs["b_major"] = "k" + + output_buffer = validate_or_alloc_output(output_buffer, fc2_out_shape, dtype, device) + fc2_quant_kwargs["d_tensor"] = output_buffer.as_strided( + (in_shape[0], fc2_weight_shape[0], 1), + (fc2_weight_shape[0], 1, in_shape[0] * fc2_weight_shape[0]), + ) + self.grouped_gemm_quant_kernel()(**fc2_quant_kwargs) + fc2_out = output_buffer else: fc2_in_row_data = fc1_kernel_out["d_tensor"] fc2_in_row_data = fc2_in_row_data.view(in_shape[0], fc2_weight_shape[1]) @@ -1780,6 +1978,7 @@ def fuser_forward( ) fc1_ctx.input_quantizers = [fc1_input_quantizer] + fc1_ctx.weight_quantizers = [fc1_weight_quantizer] fc1_ctx.grad_output_quantizers = [fc1_grad_output_quantizer] fc1_ctx.dtype = dtype fc1_ctx.input_requires_grad = input_requires_grad @@ -1788,6 +1987,7 @@ def fuser_forward( fc2_ctx.input_quantizers = [fc2_input_quantizer] fc2_ctx.grad_output_quantizers = [fc2_grad_output_quantizer] + fc2_ctx.weight_quantizers = [fc2_weight_quantizer] fc2_ctx.dtype = dtype fc2_ctx.input_requires_grad = input_requires_grad fc2_ctx.weight_requires_grad = weight_requires_grad @@ -1888,6 +2088,7 @@ def fuser_backward( # Split grad output tensor and convert dtypes if needed fc2_grad_output_quantizer = fc2_ctx.grad_output_quantizers[0] + fc2_weight_quantizer = fc2_ctx.weight_quantizers[0] fc2_grad_output_quantizer.set_usage(rowwise=True, columnwise=fc2_ctx.weight_requires_grad) fc2_grad_output_quantizer.optimize_for_gemm = True output_fc2_dbias = fc2_op.has_bias @@ -2008,24 +2209,33 @@ def fuser_backward( fc2_d_dtype = torch.bfloat16 if use_nvfp4 else torch.float8_e4m3fn if use_nvfp4: - nvfp4_fp4_max = 6.0 - nvfp4_fp8_max = 448.0 - nvfp4_global_scale_denom = nvfp4_fp4_max * nvfp4_fp8_max - fc2_dy_amax = _nvfp4_amax(grouped_fc2_dy, columnwise=False) - fc2_weight_col_amax = _nvfp4_amax(grouped_fc2_weight, columnwise=True) + fc2_dy_global_scale = _nvfp4_global_scale( + grouped_fc2_dy, + fc2_grad_output_quantizer, + columnwise=False, + num_groups=num_groups, + device=device, + ) + fc2_weight_col_global_scale = _nvfp4_global_scale( + grouped_fc2_weight, + fc2_weight_quantizer, + columnwise=True, + num_groups=num_groups, + device=device, + ) if activation_is_srelu: # DSReLU applies alpha once, so pass the full product of the # two operand global scales. fc2_alpha_tensor = ( - (fc2_dy_amax * fc2_weight_col_amax / (nvfp4_global_scale_denom**2)) + (fc2_dy_global_scale * fc2_weight_col_global_scale) .to(torch.float32) .expand(num_groups) ) else: # DGLU applies alpha to both gate branches, so the wrapper # expects sqrt(product) to recover the same global-scale factor. - fc2_alpha_tensor = ( - torch.sqrt(fc2_dy_amax * fc2_weight_col_amax) / nvfp4_global_scale_denom + fc2_alpha_tensor = torch.sqrt( + fc2_dy_global_scale * fc2_weight_col_global_scale ).expand(num_groups) fc2_beta_tensor = get_cached_ones_tensor(num_groups, torch.float32, device) fc2_norm_const_tensor = None @@ -2054,6 +2264,9 @@ def fuser_backward( dactivation_kernel = self.grouped_gemm_dactivation_kernel() if _cudnn_frontend_supports_single_group_runtime_offsets(): fc2_dactivation_kwargs["use_single_group_runtime_offsets"] = num_groups == 1 + fc2_sf_dtype_override = _nvfp4_sf_dtype_override(fc2_grad_output_quantizer) + if fc2_sf_dtype_override is not None: + fc2_dactivation_kwargs["sf_fp8_dtype_override"] = fc2_sf_dtype_override if self._cudnn_dact_func is not None: fc2_dactivation_kwargs["beta_tensor"] = fc2_beta_tensor fc2_dactivation_kwargs["act_func"] = self._cudnn_dact_func @@ -2271,6 +2484,7 @@ def fuser_backward( # FC1 grad output for dgrad and wgrad GEMMs fc1_dy_tensor_offsets = fc1_out_tensor_offsets fc1_grad_output_quantizer = fc1_ctx.grad_output_quantizers[0] + fc1_weight_quantizer = fc1_ctx.weight_quantizers[0] if use_nvfp4: fc1_grad_output_quantizer.set_usage( rowwise=True, @@ -2344,6 +2558,8 @@ def fuser_backward( if is_distributed_weight(fc1_leader): grouped_fc1_weight = materialize_weight_for_backward(fc1_leader) + fc1_dgrad_sf_override = _nvfp4_sf_dtype_override(fc1_grad_output_quantizer) + use_single_group_dense_dgrad = num_groups == 1 if use_single_group_dense_dgrad: grad_input = validate_or_alloc_output(grad_input_buffer, in_shape, dtype, device) @@ -2354,7 +2570,7 @@ def fuser_backward( single_grouped_weight=fc1_op.single_grouped_weight, dtype=dtype, ) - elif use_nvfp4: + elif use_nvfp4 and fc1_dgrad_sf_override is None: grad_input = validate_or_alloc_output(grad_input_buffer, in_shape, dtype, device) grouped_grad_input = GroupedTensor( shape=(out_shape[0], fc1_weight_shape[1]), @@ -2371,6 +2587,106 @@ def fuser_backward( grouped_grad_input, layout="NN", ) + elif use_nvfp4: # TODO(kainingz): remove this e5m3 workaround once cuBLAS is ready + # This assertion should never fail because we set fc1_grad_output_quantizer.optimize_for_gemm = True + assert grouped_fc1_dy._with_gemm_swizzled_scales, ( + "cuDNN NVFP4 dgrad requires GEMM-swizzled grad-output scale factors." + ) + + grad_input_buffer = validate_or_alloc_output(grad_input_buffer, in_shape, dtype, device) + + dgrad_k = fc1_weight_shape[0] # contraction dim + dgrad_valid_m = out_shape[0] # batch dim + + # Create A and its sf tensor for cuDNN that satisfies its layout requirements + fc1_dgrad_a_data, fc1_dgrad_a_scales = convert_TE_MX_tensor_to_cuDNN_operand( + grouped_fc1_dy.rowwise_data, + grouped_fc1_dy.scale_inv, + data_dtype=data_dtype, + scale_dtype=scale_view_dtype, + valid_M_or_N=dgrad_valid_m, + k_logical=dgrad_k, + sf_swizzled=grouped_fc1_dy._with_gemm_swizzled_scales, + ) + + fc1_dgrad_alpha = ( + _nvfp4_global_scale( + grouped_fc1_dy, + fc1_grad_output_quantizer, + columnwise=False, + num_groups=num_groups, + device=device, + ) + * _nvfp4_global_scale( + grouped_fc1_weight, + fc1_weight_quantizer, + columnwise=True, + num_groups=num_groups, + device=device, + ) + ).to(torch.float32) + + fc1_dgrad_kwargs = { + "a_tensor": fc1_dgrad_a_data, + "sfa_tensor": fc1_dgrad_a_scales, + "padded_offsets": split_points, + "alpha_tensor": fc1_dgrad_alpha, + "norm_const_tensor": None, # must be None for FP4 inputs + "acc_dtype": torch.float32, + "d_dtype": dtype, # high precision -> no output quantization + "cd_major": "n", + "sf_vec_size": sf_vec_size, + "sf_fp8_dtype_override": fc1_dgrad_sf_override, + "current_stream": current_stream, + "discrete_col_sfd": False, + "use_dynamic_sched": True, + } + + if fc1_op.single_grouped_weight: + # Clone and swizzle scales for GEMM + fc1_weight_for_gemm = grouped_fc1_weight.copy() + tex.grouped_swizzle_for_gemm( + fc1_weight_for_gemm, rowwise=False, columnwise=True + ) + + # Create B and its sf tensor for cuDNN that satisfies its layout + # requirements. NVFP4 column-wise data is physically transposed, so + # it is already (in_features, out_features) and stays K-major. + fc1_w_data, fc1_w_scales = convert_TE_MX_tensor_to_cuDNN_operand( + fc1_weight_for_gemm.columnwise_data, + fc1_weight_for_gemm.columnwise_scale_inv, + data_dtype=data_dtype, + scale_dtype=scale_view_dtype, + valid_M_or_N=fc1_weight_shape[1], + k_logical=dgrad_k, + L=num_groups, + sf_swizzled=fc1_weight_for_gemm._with_gemm_swizzled_scales, + ) + fc1_dgrad_kwargs["b_tensor"] = fc1_w_data + fc1_dgrad_kwargs["sfb_tensor"] = fc1_w_scales + else: + fc1_b_ptrs, fc1_sfb_ptrs, _fc1_sfb_buffer = ( + tex.grouped_mlp_experimental.swizzle_scales_and_pack_ptrs_for_discrete_weights( + [w._columnwise_data for w in grouped_fc1_weight], + [w._columnwise_scale_inv for w in grouped_fc1_weight], + "nvfp4", + device, + ) + ) + fc1_dgrad_kwargs["b_ptrs"] = fc1_b_ptrs + fc1_dgrad_kwargs["sfb_ptrs"] = fc1_sfb_ptrs + fc1_dgrad_kwargs["n"] = fc1_weight_shape[1] + fc1_dgrad_kwargs["b_dtype"] = torch.float4_e2m1fn_x2 + # FP4 has no N-major operand support, and the column-wise buffer is + # already transposed, so it is K-major. + fc1_dgrad_kwargs["b_major"] = "k" + + fc1_dgrad_kwargs["d_tensor"] = grad_input_buffer.as_strided( + (out_shape[0], fc1_weight_shape[1], 1), + (fc1_weight_shape[1], 1, out_shape[0] * fc1_weight_shape[1]), + ) + self.grouped_gemm_quant_kernel()(**fc1_dgrad_kwargs) + grad_input = grad_input_buffer else: fc1_dgrad_a_data = fc2_dgrad_kernel_out["d_row_tensor"] fc1_dgrad_a_scales = fc2_dgrad_kernel_out["sfd_row_tensor"] From ccb900c3217db0b08f337ede2ff34a44c3a79008 Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Wed, 12 Aug 2026 00:49:19 +0000 Subject: [PATCH 2/6] have to pad to 256 to use cuDNN Signed-off-by: Kaining Zhong --- .../pytorch/cpp_extensions/gemm.py | 34 +++++++++++++++++-- 1 file changed, 31 insertions(+), 3 deletions(-) diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index b6435efa29..808d588715 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -394,6 +394,24 @@ def general_cuDNN_MX_gemm( # layouts that select it give B the logical shape (K, M). out_shape = (*dataB.shape[:-1], N) if not transb else (M, N) + # cuDNN's grouped quant kernel requires M to be divisible by 256 so we need to pad it + M_padded = ceil_div(M, 256) * 256 + if M_padded != M: + k_packed = dataB.shape[-1] + src = dataB.reshape(M, k_packed) + buf = src.new_zeros((M_padded, k_packed)) + buf[:M].copy_(src) + dataB = buf + + # Swizzled scales are blocked by 128 rows: + # (1, ceil(M/128), k_sf_tiles, 32, 4, 4) + per_block = ceil_div(K, 4 * NVFP4_BLOCK_SCALING_SIZE) * 32 * 4 * 4 + n_blk, n_blk_padded = ceil_div(M, 128), ceil_div(M_padded, 128) + src_sf = sfB.reshape(-1)[: n_blk * per_block].reshape(n_blk, per_block) + buf_sf = src_sf.new_zeros((n_blk_padded, per_block)) + buf_sf[:n_blk].copy_(src_sf) + sfB = buf_sf + # cuDNN's own operand names are the other way round: its "a" is the (M, K) # activation-like operand (TE's B) and its "b" is the (N, K) weight-like one # (TE's A). @@ -402,7 +420,7 @@ def general_cuDNN_MX_gemm( sfB, data_dtype=torch.float4_e2m1fn_x2, scale_dtype=torch.float8_e4m3fn, # e5m3 rides as e4m3; torch has no ue5m3 - valid_M_or_N=M, + valid_M_or_N=M_padded, k_logical=K, L=1, sf_swizzled=True, # ensured above @@ -449,7 +467,12 @@ def general_cuDNN_MX_gemm( # Prepare for output out = validate_or_alloc_output(out, out_shape, out_dtype, device) - d_tensor = out.view(M, N).as_strided((M, N, 1), (N, 1, M * N)) + if M_padded != M: + # The kernel writes M_padded rows, so it cannot target `out` directly. + d_buf = torch.empty((M_padded, N), dtype=out_dtype, device=device) + d_tensor = d_buf.as_strided((M_padded, N, 1), (N, 1, M_padded * N)) + else: + d_tensor = out.view(M, N).as_strided((M, N, 1), (N, 1, M * N)) gemm_kwargs = { "a_tensor": cudnn_a, @@ -457,7 +480,7 @@ def general_cuDNN_MX_gemm( "b_tensor": cudnn_b, "sfb_tensor": cudnn_sfb, # One group, so the only padded end offset is the full row count. - "padded_offsets": torch.tensor([M], dtype=torch.int32, device=device), + "padded_offsets": torch.tensor([M_padded], dtype=torch.int32, device=device), "alpha_tensor": alpha_tensor, "bias_tensor": bias, "norm_const_tensor": None, # must be None for FP4 inputs @@ -473,6 +496,11 @@ def general_cuDNN_MX_gemm( } grouped_gemm_quant_kernel()(**gemm_kwargs) + if M_padded != M: + # Drop the zero-padded rows. Safe to overwrite rather than accumulate: + # this path asserts accumulate is False above. + out.view(M, N).copy_(d_buf[:M]) + # Matches general_gemm's contract: (out, bias_grad, gelu_input, extra_output). return out, None, None, None From ee42aaea9622c3deb8c1568392d3c04d81ec3454 Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Wed, 12 Aug 2026 21:48:11 +0000 Subject: [PATCH 3/6] fix: need to pass scale_dtype Signed-off-by: Kaining Zhong --- transformer_engine/pytorch/csrc/extensions.h | 17 +++++++++----- .../csrc/extensions/nvfp4_2d_partial_cast.cpp | 11 +++++---- .../pytorch/csrc/extensions/pybind.cpp | 15 +++++++++--- .../pytorch/csrc/extensions/transpose.cpp | 23 ++++++++++++------- transformer_engine/pytorch/tensor/utils.py | 14 ++++++++++- 5 files changed, 58 insertions(+), 22 deletions(-) diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index 1a9fabb3bb..4f74a81b1a 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -217,19 +217,23 @@ void nvfp4_multi_tensor_compute_partial_amax( void nvfp4_expand_scale_to_fp8(at::Tensor input, at::Tensor output, int64_t tile_rows, int64_t tile_cols, int64_t rows_padded, int64_t block_len); -void nvfp4_compute_per_block_scale(at::Tensor block_amax, at::Tensor scale, at::Tensor global_amax); +void nvfp4_compute_per_block_scale(at::Tensor block_amax, at::Tensor scale, at::Tensor global_amax, + const DType scale_dtype = DType::kFloat8E4M3); void nvfp4_fused_scale(at::Tensor block_amax, at::Tensor global_amax, at::Tensor per_block_scale, at::Tensor target_scale, at::Tensor target_amax, int64_t tile_rows, - int64_t tile_cols, int64_t rows_padded, int64_t block_len); + int64_t tile_cols, int64_t rows_padded, int64_t block_len, + const DType scale_dtype = DType::kFloat8E4M3); void nvfp4_multi_tensor_fused_scale( std::vector block_amax_list, std::vector global_amax_list, std::vector per_block_scale_list, std::vector target_scale_list, std::vector target_amax_list, std::vector tile_rows_list, - std::vector tile_cols_list, std::vector rows_padded_list, int64_t block_len); + std::vector tile_cols_list, std::vector rows_padded_list, int64_t block_len, + const DType scale_dtype = DType::kFloat8E4M3); -void nvfp4_compute_global_scale(at::Tensor global_amax, at::Tensor global_scale); +void nvfp4_compute_global_scale(at::Tensor global_amax, at::Tensor global_scale, + const DType scale_dtype = DType::kFloat8E4M3); at::Tensor swap_first_dims(at::Tensor tensor, std::optional out = std::nullopt); @@ -489,14 +493,15 @@ void nvfp4_2d_compute_partial_amax(const at::Tensor &tensor, at::Tensor amax, si void nvfp4_2d_partial_cast(const at::Tensor &inp, py::handle out, const at::Tensor &scale, const at::Tensor &global_scale, size_t h, size_t w, size_t start_offset, - size_t block_len); + size_t block_len, const DType scale_dtype = DType::kFloat8E4M3); void nvfp4_multi_tensor_2d_partial_cast(std::vector inp_list, std::vector out_list, std::vector scale_list, std::vector global_scale_list, std::vector h_list, std::vector w_list, - std::vector start_offset_list, int64_t block_len); + std::vector start_offset_list, int64_t block_len, + const DType scale_dtype = DType::kFloat8E4M3); void mxfp8_scaling_compute_partial_amax(const at::Tensor &input, at::Tensor amax_rowwise, at::Tensor amax_colwise, int rows, int cols, size_t start_offset); diff --git a/transformer_engine/pytorch/csrc/extensions/nvfp4_2d_partial_cast.cpp b/transformer_engine/pytorch/csrc/extensions/nvfp4_2d_partial_cast.cpp index 685250d137..8b58299351 100644 --- a/transformer_engine/pytorch/csrc/extensions/nvfp4_2d_partial_cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/nvfp4_2d_partial_cast.cpp @@ -27,7 +27,7 @@ void nvfp4_2d_compute_partial_amax(const at::Tensor& tensor, at::Tensor amax, si void nvfp4_2d_partial_cast(const at::Tensor& inp, py::handle out, const at::Tensor& scale, const at::Tensor& global_scale, size_t h, size_t w, size_t start_offset, - size_t block_len) { + size_t block_len, const DType scale_dtype) { TORCH_CHECK(block_len == 16, "Currently only block_len = 16 is supported for NVFP4 2D"); TORCH_CHECK(scale.dim() == 2, "scale must be a 2D tensor"); TORCH_CHECK(scale.scalar_type() == at::ScalarType::Float, "scale must be a float tensor"); @@ -45,7 +45,8 @@ void nvfp4_2d_partial_cast(const at::Tensor& inp, py::handle out, const at::Tens nvte_nvfp4_2d_partial_cast(inp_cu.data(), out_cu.data(), scale_cu.data(), global_scale_cu.data(), h, w, scale.stride(0), scale.stride(1), start_offset, block_len, - at::cuda::getCurrentCUDAStream()); + at::cuda::getCurrentCUDAStream(), + static_cast(scale_dtype)); } void nvfp4_multi_tensor_2d_partial_cast(std::vector inp_list, @@ -53,7 +54,8 @@ void nvfp4_multi_tensor_2d_partial_cast(std::vector inp_list, std::vector scale_list, std::vector global_scale_list, std::vector h_list, std::vector w_list, - std::vector start_offset_list, int64_t block_len) { + std::vector start_offset_list, int64_t block_len, + const DType scale_dtype) { TORCH_CHECK(block_len == 16, "Currently only block_len = 16 is supported for NVFP4 2D"); const size_t num_tensors = inp_list.size(); @@ -95,7 +97,8 @@ void nvfp4_multi_tensor_2d_partial_cast(std::vector inp_list, nvte_nvfp4_2d_partial_cast(inp_cu.data(), out_cu.data(), scale_cu.data(), global_scale_cu.data(), h, w, scale.stride(0), scale.stride(1), - start_offset, static_cast(block_len), stream); + start_offset, static_cast(block_len), stream, + static_cast(scale_dtype)); } } diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp index 2173dff8b2..1cd9c11e8b 100644 --- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp +++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp @@ -412,15 +412,21 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("nvfp4_compute_per_block_scale", &transformer_engine::pytorch::nvfp4_compute_per_block_scale, "Compute per-block decode scale from block amax and global amax", py::arg("block_amax"), - py::arg("scale"), py::arg("global_amax"), py::call_guard()); + py::arg("scale"), py::arg("global_amax"), + py::arg("scale_dtype") = transformer_engine::DType::kFloat8E4M3, + py::call_guard()); m.def("nvfp4_compute_global_scale", &transformer_engine::pytorch::nvfp4_compute_global_scale, "Compute global encode scale from global amax", py::arg("global_amax"), - py::arg("global_scale"), py::call_guard()); + py::arg("global_scale"), + py::arg("scale_dtype") = transformer_engine::DType::kFloat8E4M3, + py::call_guard()); m.def("nvfp4_fused_scale", &transformer_engine::pytorch::nvfp4_fused_scale, "Fused kernel: compute per-block decode scale, copy global amax, expand to row-level FP8", py::arg("block_amax"), py::arg("global_amax"), py::arg("per_block_scale"), py::arg("target_scale"), py::arg("target_amax"), py::arg("tile_rows"), py::arg("tile_cols"), - py::arg("rows_padded"), py::arg("block_len"), py::call_guard()); + py::arg("rows_padded"), py::arg("block_len"), + py::arg("scale_dtype") = transformer_engine::DType::kFloat8E4M3, + py::call_guard()); m.def("nvfp4_multi_tensor_fused_scale", &transformer_engine::pytorch::nvfp4_multi_tensor_fused_scale, "Batched fused scale: compute per-block decode scale, copy global amax, expand to FP8 for " @@ -428,6 +434,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { py::arg("block_amax_list"), py::arg("global_amax_list"), py::arg("per_block_scale_list"), py::arg("target_scale_list"), py::arg("target_amax_list"), py::arg("tile_rows_list"), py::arg("tile_cols_list"), py::arg("rows_padded_list"), py::arg("block_len"), + py::arg("scale_dtype") = transformer_engine::DType::kFloat8E4M3, py::call_guard()); m.def("nvfp4_2d_multi_tensor_transpose", &transformer_engine::pytorch::nvfp4_2d_multi_tensor_transpose, @@ -474,12 +481,14 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { "Partial cast from master weights for NVFP4 2D", py::arg("inp"), py::arg("out"), py::arg("scale"), py::arg("global_scale"), py::arg("h"), py::arg("w"), py::arg("start_offset"), py::arg("block_len") = 16, + py::arg("scale_dtype") = transformer_engine::DType::kFloat8E4M3, py::call_guard()); m.def("nvfp4_multi_tensor_2d_partial_cast", &transformer_engine::pytorch::nvfp4_multi_tensor_2d_partial_cast, "Batched partial cast from master weights for NVFP4 2D", py::arg("inp_list"), py::arg("out_list"), py::arg("scale_list"), py::arg("global_scale_list"), py::arg("h_list"), py::arg("w_list"), py::arg("start_offset_list"), py::arg("block_len") = 16, + py::arg("scale_dtype") = transformer_engine::DType::kFloat8E4M3, py::call_guard()); m.def("mxfp8_scaling_compute_partial_amax", &transformer_engine::pytorch::mxfp8_scaling_compute_partial_amax, diff --git a/transformer_engine/pytorch/csrc/extensions/transpose.cpp b/transformer_engine/pytorch/csrc/extensions/transpose.cpp index 0318978195..4b887c3749 100644 --- a/transformer_engine/pytorch/csrc/extensions/transpose.cpp +++ b/transformer_engine/pytorch/csrc/extensions/transpose.cpp @@ -145,7 +145,7 @@ void nvfp4_expand_scale_to_fp8(at::Tensor input, at::Tensor output, int64_t tile } void nvfp4_compute_per_block_scale(at::Tensor block_amax, at::Tensor scale, - at::Tensor global_amax) { + at::Tensor global_amax, const DType scale_dtype) { init_extension(); // block_amax and scale: [tile_rows, tile_cols], float32 @@ -160,12 +160,14 @@ void nvfp4_compute_per_block_scale(at::Tensor block_amax, at::Tensor scale, auto global_amax_cu = makeTransformerEngineTensor(global_amax); nvte_nvfp4_compute_per_block_scale(block_amax_cu.data(), scale_cu.data(), global_amax_cu.data(), - at::cuda::getCurrentCUDAStream()); + at::cuda::getCurrentCUDAStream(), + static_cast(scale_dtype)); } void nvfp4_fused_scale(at::Tensor block_amax, at::Tensor global_amax, at::Tensor per_block_scale, at::Tensor target_scale, at::Tensor target_amax, int64_t tile_rows, - int64_t tile_cols, int64_t rows_padded, int64_t block_len) { + int64_t tile_cols, int64_t rows_padded, int64_t block_len, + const DType scale_dtype) { init_extension(); // block_amax: [tile_rows, tile_cols], float32 @@ -191,14 +193,16 @@ void nvfp4_fused_scale(at::Tensor block_amax, at::Tensor global_amax, at::Tensor target_scale_cu.data(), target_amax_cu.data(), static_cast(tile_rows), static_cast(tile_cols), static_cast(rows_padded), static_cast(block_len), - at::cuda::getCurrentCUDAStream()); + at::cuda::getCurrentCUDAStream(), + static_cast(scale_dtype)); } void nvfp4_multi_tensor_fused_scale( std::vector block_amax_list, std::vector global_amax_list, std::vector per_block_scale_list, std::vector target_scale_list, std::vector target_amax_list, std::vector tile_rows_list, - std::vector tile_cols_list, std::vector rows_padded_list, int64_t block_len) { + std::vector tile_cols_list, std::vector rows_padded_list, int64_t block_len, + const DType scale_dtype) { init_extension(); const size_t num_tensors = block_amax_list.size(); @@ -242,11 +246,13 @@ void nvfp4_multi_tensor_fused_scale( nvte_nvfp4_fused_scale(block_amax_cu.data(), global_amax_cu.data(), per_block_scale_cu.data(), target_scale_cu.data(), target_amax_cu.data(), tile_rows, tile_cols, - rows_padded, static_cast(block_len), stream); + rows_padded, static_cast(block_len), stream, + static_cast(scale_dtype)); } } -void nvfp4_compute_global_scale(at::Tensor global_amax, at::Tensor global_scale) { +void nvfp4_compute_global_scale(at::Tensor global_amax, at::Tensor global_scale, + const DType scale_dtype) { init_extension(); // global_amax and global_scale: [num_params], float32 @@ -257,7 +263,8 @@ void nvfp4_compute_global_scale(at::Tensor global_amax, at::Tensor global_scale) auto global_scale_cu = makeTransformerEngineTensor(global_scale); nvte_nvfp4_compute_global_scale(global_amax_cu.data(), global_scale_cu.data(), - at::cuda::getCurrentCUDAStream()); + at::cuda::getCurrentCUDAStream(), + static_cast(scale_dtype)); } at::Tensor swap_first_dims(at::Tensor tensor, std::optional out) { diff --git a/transformer_engine/pytorch/tensor/utils.py b/transformer_engine/pytorch/tensor/utils.py index e35d57b363..fa77fdd1e3 100644 --- a/transformer_engine/pytorch/tensor/utils.py +++ b/transformer_engine/pytorch/tensor/utils.py @@ -851,7 +851,17 @@ def _cast_master_weights_to_nvfp4_2d( # This replaces multiple Python tensor operations with a single kernel global_scale_tensor = torch.empty_like(global_amaxes) - tex.nvfp4_compute_global_scale(global_amaxes, global_scale_tensor) + # There should only be one scale dtype for all quantizers in the params list if using the same NVFP4 recipe. + scale_dtypes = {p[0]._get_quantizer().scale_dtype for p in params} + if len(scale_dtypes) != 1: + raise ValueError( + "quantize_master_weights requires a single NVFP4 scale dtype per call, " + f"but got {scale_dtypes}." + ) + # NVFP4Quantizer.scale_dtype is the pure-python constants.DType; the pybind + # entry points want tex.DType. The enum values are shared, so map by value. + scale_dtype = tex.DType(int(scale_dtypes.pop())) + tex.nvfp4_compute_global_scale(global_amaxes, global_scale_tensor, scale_dtype=scale_dtype) global_scale_views = [global_scale_tensor[i : i + 1] for i in range(len(params))] # Collect tensors for batched fused scale kernel @@ -949,6 +959,7 @@ def _cast_master_weights_to_nvfp4_2d( fused_scale_tile_cols_list, fused_scale_rows_padded_list, block_len, + scale_dtype=scale_dtype, ) # Batched multi-tensor call for partial cast @@ -962,6 +973,7 @@ def _cast_master_weights_to_nvfp4_2d( partial_cast_w_list, partial_cast_start_offset_list, block_len, + scale_dtype=scale_dtype, ) From 419791084f7b14a50efa4699d41da23f62ba9772 Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Thu, 13 Aug 2026 02:46:17 +0000 Subject: [PATCH 4/6] route wgrad to cuDNN's wgrad API Signed-off-by: Kaining Zhong --- .../pytorch/cpp_extensions/gemm.py | 136 ++++++++++++++++-- 1 file changed, 122 insertions(+), 14 deletions(-) diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index 808d588715..7a464f87f3 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -189,6 +189,83 @@ def _validate_native_gemm_output_quantizer(quantization_params): ) +@functools.lru_cache(maxsize=None) +def grouped_gemm_wgrad_kernel() -> Callable: + """cuDNN CuTe DSL grouped wgrad kernel for block-scaled inputs.""" + from cudnn import grouped_gemm_wgrad_wrapper_sm100 # pylint: disable=no-name-in-module + + return grouped_gemm_wgrad_wrapper_sm100 + + +def _cuDNN_wgrad_gemm( + a_tensor: torch.Tensor, + b_tensor: torch.Tensor, + sfa: torch.Tensor, + sfb: torch.Tensor, + amax_a: Optional[torch.Tensor], + amax_b: Optional[torch.Tensor], + out_dtype: torch.dtype, + out: Optional[torch.Tensor], + accumulate: bool, + alpha: Optional[float] = None, +) -> Iterable[Optional[torch.Tensor]]: + """Compute dw = dy^T @ x with cuDNN's purpose-built grouped wgrad kernel.""" + + # Column-wise NVFP4 buffers are physically (features, tokens), FP4-packed + # two values per byte along the token dim. + tokens_packed = a_tensor.shape[-1] + tokens = tokens_packed * 2 + out_features = a_tensor.numel() // tokens_packed + in_features = b_tensor.numel() // tokens_packed + + # grouped_gemm_wgrad_wrapper_sm100 wants: + # a_tensor (feature_out, tokens) K-major, FP4-packed + # b_tensor (tokens, feature_in) + # sfa (round_up(feature_out, 128), scale_cols) + # sfb (round_up(feature_in, 128), scale_cols) + fp4 = torch.float4_e2m1fn_x2 + a_tensor = a_tensor.view(dtype=fp4).view(out_features, tokens_packed) + b_tensor = b_tensor.view(dtype=fp4).view(in_features, tokens_packed).T + + # Create the scale factor tensors with the logical layout cuDNN expects + # In general_cuDNN_MX_gemm we've already ensured they are swizzled physically + def _sf(scale_inv, features): + leading = ceil_div(features, 128) * 128 + return scale_inv.view(leading, -1).view(dtype=torch.float8_e4m3fn) + + # grouped_gemm_wgrad_wrapper_sm100 expects two separate global_scale + ones = get_cached_ones_tensor(1, dtype=torch.float32, device=a_tensor.device) + denom = 6.0 * 114688.0 # fp4_max * fp8_max(UE5M3) + global_scale_a = ones if amax_a is None else amax_a.to(torch.float32).reshape(1) / denom + global_scale_b = ones if amax_b is None else amax_b.to(torch.float32).reshape(1) / denom + # Fold alpha into one of them if it's given + if alpha is not None and alpha != 1.0: + global_scale_a = global_scale_a * alpha + + out = validate_or_alloc_output(out, (out_features, in_features), out_dtype, a_tensor.device) + grouped_gemm_wgrad_kernel()( + a_tensor=a_tensor, + b_tensor=b_tensor, + sfa_tensor=_sf(sfa, out_features), + sfb_tensor=_sf(sfb, in_features), + offsets_tensor=torch.tensor([tokens], dtype=torch.int32, device=a_tensor.device), + global_scale_a=global_scale_a, + global_scale_b=global_scale_b, + acc_dtype=torch.float32, + wgrad_dtype=out.dtype, + output_mode="dense", + wgrad_tensor=out.view(1, out_features, in_features), + sf_vec_size=NVFP4_BLOCK_SCALING_SIZE, + sf_fp8_dtype_override="e5m3", + input_order="tensor_ragged", + accumulate_on_output=accumulate, + current_stream=torch.cuda.current_stream().cuda_stream, + ) + + # Matches general_gemm's contract: (out, bias_grad, gelu_input, extra_output). + return out, None, None, None + + @functools.lru_cache(maxsize=None) def grouped_gemm_quant_kernel() -> Callable: """cuDNN CuTe DSL grouped GEMM kernel for block-scaled inputs.""" @@ -329,8 +406,17 @@ def general_cuDNN_MX_gemm( Note: layout string means layout of "A" and "B" respectively. - We use cuDNN-frontend's grouped_gemm_quant_wrapper_sm100 API here which is supposed to be a grouped GEMM - but here we set groups = 1 so it is effectively a single GEMM. + TE stores x (token, feature_in), w (feature_out, feature_in) and dy (token, feature_out) in physical rowwise direction. + For cuBLAS: + fprop = x @ wT: token is M, feature_in is K, feature_out is N, so it's TN (x as B, w transposed to wT as A) + dgrad = dy @ w: token is M, feature_out is K, feature_in is N, so it's NN (dy as B, w as A) + wgrad = dyT @ x: feature_out is M, token is K, feature_in is N, so it's NT (dy transposed to dyT as B, x as A) + + We use cuDNN-frontend's APIs here which are supposed to be used for grouped GEMM but we set groups = 1 + so it is effectively a single GEMM. + + Naming convention: uppercase letters (A, B) are used for cuBLAS notation, lowercase letters (a, b) are used for cuDNN notation. + where cuBLAS's B is cuDNN's a, and cuBLAS's A is cuDNN's b (their notation is inverted). """ assert isinstance(A, NVFP4TensorStorage) and isinstance(B, NVFP4TensorStorage) and \ A.get_metadata()["scale_dtype"] == DType.kFloat8UE5M3 and B.get_metadata()["scale_dtype"] == DType.kFloat8UE5M3, \ @@ -338,13 +424,7 @@ def general_cuDNN_MX_gemm( assert quantization_params is None, "cuDNN GEMM currently does not support output quantization." assert gelu is False and gelu_in is None, "cuDNN GEMM currently does not support fused GELU." - assert accumulate is False, "cuDNN GEMM currently does not support accumulation." - # `grad` only changes behaviour when a bias is supplied: it turns the bias slot - # into a bias-gradient output, which cuDNN has no epilogue for. Backward GEMMs - # that pass grad=True without a bias need nothing special. - assert not (grad and bias is not None), ( - "cuDNN GEMM currently does not support fused bias gradient." - ) + # use_split_accumulator is deliberately not checked: it is a cuBLAS knob for # raising accumulator precision, and the cuDNN kernel always accumulates in # FP32, so the request is already satisfied either way. @@ -371,6 +451,13 @@ def general_cuDNN_MX_gemm( if not B._with_gemm_swizzled_scales: tex.swizzle_scales_for_gemm_(B) + # `grad` only changes behaviour when a bias is supplied: it turns the bias slot + # into a bias-gradient output, which cuDNN has no epilogue for. Backward GEMMs + # that pass grad=True without a bias need nothing special. + assert not (grad and bias is not None), ( + "cuDNN GEMM currently does not support fused bias gradient." + ) + # Pick the buffer whose block scales run along K. In every case the selected # buffer is physically (rows, K_packed), so the reshape below is uniform. # LHS is always (M, K) @@ -389,6 +476,32 @@ def general_cuDNN_MX_gemm( N, k_from_a = dataA.numel() // dataA.shape[-1], dataA.shape[-1] * 2 assert K == k_from_a, f"Contraction dims disagree: A implies {k_from_a}, B implies {K}." + # Route NT (implying wgrad) to cuDNN-FE's wgrad API instead of the general GEMM one + if layout == "NT" and bias is None: # The wgrad kernel has no bias epilogue + alpha = alpha if alpha is not None else 1.0 + # This path uses cuDNN's wgrad (grouped_gemm_wgrad_wrapper_sm100) which supports grad accumulation + if accumulate: # Accumulate GEMM's result to the out tensorn + assert beta in (1.0, None), "beta must be one or None if accumulate is True" + else: # Overwrite GEMM's result to the out tensor + assert beta in (0.0, None), "beta must be zero or None if not accumulate" + return _cuDNN_wgrad_gemm( + a_tensor=dataB, + b_tensor=dataA, + sfa=sfB, + sfb=sfA, + amax_a=amaxB, + amax_b=amaxA, + out_dtype=out_dtype, + out=out, + accumulate=accumulate, + alpha=alpha + ) + + alpha = alpha if alpha is not None else 1.0 + # cuDNN's general GEMM path (grouped_gemm_quant_wrapper_sm100) doesn't support accumulation + assert accumulate is False, "cuDNN GEMM currently does not support accumulation for this operation." + assert beta in (0.0, None), "beta must be zero or None if not accumulate" + # The output keeps B's leading dims. They only exist when B is read row-wise: # the column-wise NVFP4 buffer is physically transposed and always 2D, and the # layouts that select it give B the logical shape (K, M). @@ -447,11 +560,6 @@ def general_cuDNN_MX_gemm( # Prepare alpha. cuDNN applies the block scales but not TE's per-tensor global # scale, so alpha carries the product of both operands'. A tensor quantized # without second-level scaling has no amax and contributes a factor of one. - - # general_gemm normalizes alpha/beta only on its cuBLAS path, downstream of the - # dispatch here, so callers can still reach this with alpha=None meaning one. - alpha = validate_gemm_scale(alpha, True) - validate_gemm_scale(beta, accumulate) nvfp4_global_scale = 6.0 * 114688.0 ones = get_cached_ones_tensor(1, dtype=torch.float32, device=device) scaleA = ones if amaxA is None else amaxA.to(torch.float32).reshape(1) / nvfp4_global_scale From 403d218b34d4a4a8856084d800544a5d4af0d17e Mon Sep 17 00:00:00 2001 From: Tim Moon Date: Thu, 13 Aug 2026 10:58:49 +0000 Subject: [PATCH 5/6] Support grouped linear with NVFP4-UE5M3 NVFP4-UE5M3 grouped GEMM falls back to dense GEMMs. Generalize usage of wgrad kernel and use when tensors sizes are not 256-aligned. Fix inconsistent m,n,k GEMM notation. Remove ue5m3 hacks in op fuser tests. Add ue5m3 to grouped MLP tests. Signed-off-by: Tim Moon --- tests/pytorch/test_fusible_ops.py | 83 +---- tests/pytorch/test_grouped_mlp.py | 37 ++- .../pytorch/cpp_extensions/gemm.py | 289 ++++++++++-------- .../pytorch/module/grouped_linear.py | 93 +++--- .../pytorch/ops/basic/grouped_linear.py | 75 +++-- 5 files changed, 299 insertions(+), 278 deletions(-) diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py index 7b8e77c02d..1ae33c0403 100644 --- a/tests/pytorch/test_fusible_ops.py +++ b/tests/pytorch/test_fusible_ops.py @@ -82,6 +82,8 @@ if nvfp4_available: _quantization_list.append("nvfp4") _quantization_list.append("nvfp4_4over6") + if fp8_ue5m3_available: + _quantization_list.append("nvfp4_rht_ue5m3") if fp8_block_scaling_available: _quantization_list.append("fp8_block_scaling") @@ -99,7 +101,6 @@ def maybe_skip_quantization( dims: Optional[Iterable[int] | int] = None, device: Optional[torch.device | str] = None, dtype: Optional[torch.dtype] = None, - fused_grouped_gemm: bool = False, ) -> None: """Skip test case if a quantization scheme is not supported""" @@ -137,14 +138,11 @@ def maybe_skip_quantization( elif quantization in nvfp4_variant_names: if math.prod(dims[:-1]) % 16 != 0 or dims[-1] % 16 != 0: pytest.skip("NVFP4 GEMMs require dims that are divisible by 16") - if quantization == "nvfp4_rht_ue5m3_scales" and not fused_grouped_gemm and ( - math.prod(dims[:-1]) % 256 != 0 or dims[-1] % 256 != 0 + if ( + quantization in ("nvfp4_ue5m3", "nvfp4_rht_ue5m3") + and (math.prod(dims[:-1]) % 64 != 0 or dims[-1] % 64 != 0) ): - # UE5M3 has no cuBLAS kernels, so these GEMMs route to cuDNN's grouped - # kernel, which pads every group to 256 rows. Both dims are constrained, - # not just the leading one: wgrad contracts over the tokens, putting the - # feature dim in the kernel's ragged position. - pytest.skip("NVFP4 UE5M3 GEMMs require dims that are divisible by 256") + pytest.skip("cuDNN FE NVFP4-UE5M3 GEMMs produce incorrect values with 32x32 tensors") # Check dtype if dtype is not None: @@ -152,31 +150,6 @@ def maybe_skip_quantization( pytest.skip("NVFP4 quantization is only supported with BF16 data") -def maybe_resize_for_ue5m3( - quantization: Optional[str], - weight_shape: tuple[int, int], - in_shape: Iterable[int], -) -> tuple[tuple[int, int], Iterable[int]]: - """Grow a test case's shapes to what the NVFP4 UE5M3 GEMM path can run. - - UE5M3 has no cuBLAS kernels, so its GEMMs route to cuDNN's grouped kernel, - which pads every group to 256 rows. Every dim is constrained, not just the - leading one: wgrad contracts over the tokens, so the feature dim also lands - in the kernel's ragged position. Tests share small default shapes to stay - fast, so rescale them for this recipe instead of losing the coverage to a - skip. Leading dims are collapsed to a single 256-row dim, keeping any extra - dims so the >2D cases still exercise the leading-dim handling. - """ - if quantization != "nvfp4_rht_ue5m3_scales": - return weight_shape, in_shape - leading = list(in_shape)[:-1] - if len(leading) == 0: - # A 1D input is a single row, which can never reach 256. - pytest.skip("NVFP4 UE5M3 GEMMs need at least 256 rows, so 1D inputs cannot run") - new_leading = [1] * (len(leading) - 1) + [256] - return (256, 256), (*new_leading, -1) - - @torch.no_grad() def make_reference_and_test_tensors( shape: int | Iterable[int], @@ -970,8 +943,6 @@ def _test_basic_linear( ) -> None: """Helper function for tests with GEMM""" - weight_shape, in_shape = maybe_resize_for_ue5m3(quantization, weight_shape, in_shape) - # Make input and weight shapes consistent out_features, in_features = weight_shape in_shape = list(in_shape)[:-1] + [in_features] @@ -980,9 +951,6 @@ def _test_basic_linear( # Skip invalid configurations maybe_skip_quantization(quantization, dims=in_shape, device=device, dtype=dtype) maybe_skip_quantization(quantization, dims=out_shape) - if quantization == "nvfp4_rht_ue5m3_scales" and accumulate_into_main_grad: - # The cuDNN kernel UE5M3 routes to writes its output, with no beta term. - pytest.skip("NVFP4 UE5M3 GEMMs cannot accumulate into an existing output") quantization_needed = any( ( quantized_compute, @@ -2184,11 +2152,6 @@ def test_grouped_linear( pytest.skip("Quantized group GEMM is only supported with BF16/FP16") if quantization == "nvfp4_4over6": pytest.skip("NVFP4 4over6 grouped quantization is not supported") - if quantization == "nvfp4_rht_ue5m3_scales": - # A standalone GroupedLinear issues grouped cuBLAS GEMMs, which have no - # UE5M3 kernels. UE5M3 grouped support is limited to the CuteDSL fused - # grouped MLP, which routes its GEMMs to cuDNN instead. - pytest.skip("UE5M3 scales are not supported by grouped cuBLAS GEMMs") # Random data x_ref, x_test = make_reference_and_test_tensors( @@ -2612,8 +2575,6 @@ def test_forward_linear_bias_activation( ) -> None: """Forward GEMM + bias + activation""" - weight_shape, in_shape = maybe_resize_for_ue5m3(quantization, weight_shape, in_shape) - # Make input and weight shapes consistent out_features, in_features = weight_shape in_shape = list(in_shape)[:-1] + [in_features] @@ -3603,7 +3564,6 @@ def test_layernorm_mlp( @pytest.mark.parametrize("bias", (False, True)) @pytest.mark.parametrize("dtype", _dtypes) @pytest.mark.parametrize("quantization", _quantization_list) - @pytest.mark.parametrize("glu_interleave_size", (None, 32)) def test_grouped_mlp( self, *, @@ -3615,7 +3575,6 @@ def test_grouped_mlp( device: torch.device = "cuda", split_alignment: int = 256, activation: str = "scaled_swiglu", - glu_interleave_size: Optional[int], ) -> None: """GroupedLinear + scaled activation + GroupedLinear""" @@ -3641,25 +3600,11 @@ def test_grouped_mlp( dims=in_shape, device=device, dtype=dtype, - fused_grouped_gemm=True, ) if with_quantization and dtype not in (torch.bfloat16, torch.float16): pytest.skip("Quantized group GEMM is only supported with BF16/FP16") if activation == "scaled_srelu" and quantization == "nvfp4_rht" and bias: pytest.skip("NVFP4 RHT SReLU grouped MLP coverage is limited to no-bias") - if ( - quantization == "nvfp4_rht_ue5m3_scales" - and activation == "scaled_swiglu" - and glu_interleave_size is None - ): - # Without interleaving the GLU pattern never matches the CuteDSL fused - # grouped MLP, so the ops fall back to grouped cuBLAS, which has no UE5M3 - # kernels. The glu_interleave_size=32 variant covers this recipe. - pytest.skip("UE5M3 grouped MLP is only supported on the fused path") - if quantization == "nvfp4_rht_ue5m3_scales" and activation == "scaled_srelu": - # The fuser refuses UE5M3 with SReLU outright, since cuDNN has no UE5M3 - # SReLU kernel, so this would fall back to grouped cuBLAS. - pytest.skip("cuDNN has no UE5M3 SReLU kernel, so this cannot use the fused path") # Random data x_ref, x_test = make_reference_and_test_tensors( @@ -3745,18 +3690,7 @@ def test_grouped_mlp( x, fc1_ws_ref[group_idx], bias=fc1_bs_ref[group_idx] ) if activation == "scaled_swiglu": - if glu_interleave_size is not None: - # Undo the interleaving so the two GLU halves can be chunked. - act_in = fc1_out.reshape( - -1, - fc1_out_features // (2 * glu_interleave_size), - 2, - glu_interleave_size, - ) - act_in = act_in.transpose(1, 2).reshape(fc1_out.shape) - else: - act_in = fc1_out - act_in1, act_in2 = act_in.chunk(2, dim=-1) + act_in1, act_in2 = fc1_out.chunk(2, dim=-1) act_out = torch.nn.functional.silu(act_in1) * act_in2 elif activation == "scaled_srelu": act_out = torch.nn.functional.relu(fc1_out).square() @@ -3791,7 +3725,7 @@ def test_grouped_mlp( scale_bias=bias, ) if activation == "scaled_swiglu": - activation_op = te_ops.ScaledSwiGLU(glu_interleave_size=glu_interleave_size) + activation_op = te_ops.ScaledSwiGLU() elif activation == "scaled_srelu": activation_op = te_ops.ScaledSReLU() else: @@ -3841,7 +3775,6 @@ def test_grouped_mlp_nvfp4_rht_srelu( quantization="nvfp4_rht", device=device, activation="scaled_srelu", - glu_interleave_size=None, # SReLU is not a GLU, so interleaving does not apply ) diff --git a/tests/pytorch/test_grouped_mlp.py b/tests/pytorch/test_grouped_mlp.py index d09c92ad49..90f4263100 100644 --- a/tests/pytorch/test_grouped_mlp.py +++ b/tests/pytorch/test_grouped_mlp.py @@ -43,6 +43,7 @@ dtype_tols, make_recipe, MegatronTrainingHelper, + nvfp4_variant_names, quantization_tols, reset_rng_states, ) @@ -51,6 +52,7 @@ fp8_available, reason_for_no_fp8 = te.is_fp8_available(return_reason=True) mxfp8_available, reason_for_no_mxfp8 = te.is_mxfp8_available(return_reason=True) nvfp4_available, reason_for_no_nvfp4 = te.is_nvfp4_available(return_reason=True) +fp8_ue5m3_available, reason_for_no_fp8_ue5m3 = te.is_fp8_ue5m3_available(return_reason=True) # Supported data types _dtypes: list[torch.dtype] = [torch.float32, torch.float16] @@ -73,6 +75,8 @@ _grouped_mlp_quantization_list.append("mxfp8") if nvfp4_available: _grouped_mlp_quantization_list.append("nvfp4_rht") + if fp8_ue5m3_available: + _quantization_list.append("nvfp4_rht_ue5m3") @pytest.fixture(autouse=True, scope="function") @@ -102,11 +106,10 @@ def maybe_skip_quantization( pytest.skip(reason_for_no_fp8) if quantization == "mxfp8" and not mxfp8_available: pytest.skip(reason_for_no_mxfp8) - if ( - quantization in ("nvfp4", "nvfp4_row_scaled", "nvfp4_4over6", "nvfp4_rht") - and not nvfp4_available - ): + if quantization in nvfp4_variant_names and not nvfp4_available: pytest.skip(reason_for_no_nvfp4) + if quantization in ("nvfp4_ue5m3", "nvfp4_rht_ue5m3") and not fp8_ue5m3_available: + pytest.skip(reason_for_no_fp8_ue5m3) # Check dims if dims is not None: @@ -121,13 +124,15 @@ def maybe_skip_quantization( elif quantization in ("nvfp4", "nvfp4_row_scaled", "nvfp4_4over6", "nvfp4_rht"): if math.prod(dims[:-1]) % 16 != 0 or dims[-1] % 16 != 0: pytest.skip("NVFP4 GEMMs require dims that are divisible by 16") + if ( + quantization in ("nvfp4_ue5m3", "nvfp4_rht_ue5m3") + and (math.prod(dims[:-1]) % 64 != 0 or dims[-1] % 64 != 0) + ): + pytest.skip("cuDNN FE NVFP4-UE5M3 GEMMs produce incorrect values with 32x32 tensors") # Check dtype if dtype is not None: - if ( - quantization in ("nvfp4", "nvfp4_row_scaled", "nvfp4_4over6", "nvfp4_rht") - and dtype != torch.bfloat16 - ): + if quantization in nvfp4_variant_names and dtype != torch.bfloat16: pytest.skip("NVFP4 quantization is only supported with BF16 data") @@ -183,17 +188,27 @@ def make_reference_and_test_tensors( test = quantizer(test) elif quantization == "mxfp8": test = MXFP8Quantizer(fp8_dtype=te.DType.kFloat8E4M3)(test) - elif quantization in ("nvfp4", "nvfp4_row_scaled", "nvfp4_rht"): + elif quantization in ( + "nvfp4", + "nvfp4_row_scaled", + "nvfp4_rht", + "nvfp4_ue5m3", + "nvfp4_rht_ue5m3", + ): tensor_type = "input" if quantizer_role is not None: tensor_type = quantizer_role.tensor_type - with_rht = quantization == "nvfp4_rht" and tensor_type != "weight" + with_rht = quantization in ("nvfp4_rht", "nvfp4_rht_ue5m3") and tensor_type != "weight" + scale_dtype = ( + te.DType.kFloat8UE5M3 if quantization == "nvfp4_rht_ue5m3" else te.DType.kFloat8E4M3 + ) test = NVFP4Quantizer( + scale_dtype=scale_dtype, with_rht=with_rht, with_post_rht_amax=with_rht, with_2d_quantization=False, stochastic_rounding=False, - with_random_sign_mask=False, + with_random_sign_mask=with_rht, )(test) elif quantization == "nvfp4_4over6": tensor_type = "input" diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index 7a464f87f3..7c7aa567fb 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -5,6 +5,8 @@ """Python interface for GEMM extensions""" from typing import Callable, Iterable, Literal, Optional, Tuple, Union, List +import itertools +import math import os import functools import torch @@ -25,7 +27,6 @@ from ..tensor.storage.nvfp4_tensor_storage import NVFP4TensorStorage from ..tensor.utils import is_custom from ..custom_recipes.gemm import custom_gemm -from ..ops._common import validate_or_alloc_output from ...debug.pytorch.debug_quantization import DebugQuantizedTensor, DebugQuantizer __all__ = [ @@ -189,6 +190,32 @@ def _validate_native_gemm_output_quantizer(quantization_params): ) +def validate_or_alloc_output( + buffer: Optional[torch.Tensor], + shape: tuple[int, ...] | list[int], + dtype: torch.dtype, + device: torch.device, +) -> torch.Tensor: + """Return the caller's output buffer, or allocate one if it is None. + + The buffer must be a contiguous tensor matching the required + shape, dtype, and device. + + """ + shape = tuple(shape) + if buffer is None: + return torch.empty(shape, dtype=dtype, device=device) + if tuple(buffer.shape) != shape: + raise ValueError(f"Output buffer shape {tuple(buffer.shape)} does not match {shape}.") + if buffer.dtype != dtype: + raise ValueError(f"Output buffer dtype {buffer.dtype} does not match {dtype}.") + if buffer.device != device: + raise ValueError(f"Output buffer device {buffer.device} does not match {device}.") + if not buffer.is_contiguous(): + raise ValueError("Output buffer must be contiguous.") + return buffer + + @functools.lru_cache(maxsize=None) def grouped_gemm_wgrad_kernel() -> Callable: """cuDNN CuTe DSL grouped wgrad kernel for block-scaled inputs.""" @@ -205,9 +232,11 @@ def _cuDNN_wgrad_gemm( amax_a: Optional[torch.Tensor], amax_b: Optional[torch.Tensor], out_dtype: torch.dtype, - out: Optional[torch.Tensor], + out: torch.Tensor, accumulate: bool, alpha: Optional[float] = None, + beta: Optional[float] = None, + bias: Optional[torch.Tensor] = None, ) -> Iterable[Optional[torch.Tensor]]: """Compute dw = dy^T @ x with cuDNN's purpose-built grouped wgrad kernel.""" @@ -215,8 +244,7 @@ def _cuDNN_wgrad_gemm( # two values per byte along the token dim. tokens_packed = a_tensor.shape[-1] tokens = tokens_packed * 2 - out_features = a_tensor.numel() // tokens_packed - in_features = b_tensor.numel() // tokens_packed + out_features, in_features = out.size() # grouped_gemm_wgrad_wrapper_sm100 wants: # a_tensor (feature_out, tokens) K-major, FP4-packed @@ -262,6 +290,10 @@ def _sf(scale_inv, features): current_stream=torch.cuda.current_stream().cuda_stream, ) + # Apply bias + if bias is not None: + out += bias.view(1, in_features) + # Matches general_gemm's contract: (out, bias_grad, gelu_input, extra_output). return out, None, None, None @@ -395,28 +427,34 @@ def general_cuDNN_MX_gemm( The parameters passed are in cuBLAS notation, where D = alpha * op(B) @ op(A) + beta * C, where the shape is always - (M, N) = (M, K) @ (K, N) + (M, N) + (N, M) = (N, K) @ (K, M) + (N, M) B: - - "N" is (M, K), which is always TE's rowwise data, and op(B) is B - - "T" is (K, M), which is always TE's colwise data, and op(B) is B.T + - "N" is (N, K), which is always TE's rowwise data, and op(B) is B + - "T" is (K, N), which is always TE's colwise data, and op(B) is B.T A - - "N" is (K, N), which is always TE's colwise data, and op(A) is A - - "T" is (N, K), which is always TE's rowwise data, and op(A) is A.T + - "N" is (K, M), which is always TE's colwise data, and op(A) is A + - "T" is (M, K), which is always TE's rowwise data, and op(A) is A.T Note: layout string means layout of "A" and "B" respectively. TE stores x (token, feature_in), w (feature_out, feature_in) and dy (token, feature_out) in physical rowwise direction. For cuBLAS: - fprop = x @ wT: token is M, feature_in is K, feature_out is N, so it's TN (x as B, w transposed to wT as A) - dgrad = dy @ w: token is M, feature_out is K, feature_in is N, so it's NN (dy as B, w as A) - wgrad = dyT @ x: feature_out is M, token is K, feature_in is N, so it's NT (dy transposed to dyT as B, x as A) + fprop = x @ wT: token is N, feature_in is K, feature_out is M, so it's TN (x as B, w transposed to wT as A) + dgrad = dy @ w: token is N, feature_out is K, feature_in is M, so it's NN (dy as B, w as A) + wgrad = dyT @ x: feature_out is N, token is K, feature_in is M, so it's NT (dy transposed to dyT as B, x as A) We use cuDNN-frontend's APIs here which are supposed to be used for grouped GEMM but we set groups = 1 so it is effectively a single GEMM. Naming convention: uppercase letters (A, B) are used for cuBLAS notation, lowercase letters (a, b) are used for cuDNN notation. where cuBLAS's B is cuDNN's a, and cuBLAS's A is cuDNN's b (their notation is inverted). + + This function is a temporary hack until TE supports NVFP4-UE5M3 + GEMMs natively. This should not be used externally and once native + GEMM support is added then this function (and related helper + functions) should be removed entirely. + """ assert isinstance(A, NVFP4TensorStorage) and isinstance(B, NVFP4TensorStorage) and \ A.get_metadata()["scale_dtype"] == DType.kFloat8UE5M3 and B.get_metadata()["scale_dtype"] == DType.kFloat8UE5M3, \ @@ -471,55 +509,84 @@ def general_cuDNN_MX_gemm( else: dataA, sfA, amaxA = A._columnwise_data, A._columnwise_scale_inv, A._amax_columnwise - # Find the logical shapes (not the physical shapes where K is packed with 2 fp4 stored in 1 byte). - M, K = dataB.numel() // dataB.shape[-1], dataB.shape[-1] * 2 - N, k_from_a = dataA.numel() // dataA.shape[-1], dataA.shape[-1] * 2 - assert K == k_from_a, f"Contraction dims disagree: A implies {k_from_a}, B implies {K}." + # Input tensor dims + A_shape = list(dataA.size()) + A_shape[-1] *= 2 + B_shape = list(dataB.size()) + B_shape[-1] *= 2 + + # GEMM dimensions + M_full = A_shape[:-1] if transa else [A_shape[0]] + N_full = [B_shape[0]] if transb else B_shape[:-1] + K_full = [A_shape[-1]] if transa else A_shape[1:] + K_full_b = B_shape[1:] if transb else [B_shape[-1]] + assert K_full == K_full_b, f"Contraction dims disagree: A implies {K_full}, B implies {K_full_b}." + M = math.prod(M_full) + N = math.prod(N_full) + K = math.prod(K_full) + + # Allocate output tensor if needed + out_shape = N_full + M_full + out = validate_or_alloc_output(out, out_shape, out_dtype, device) - # Route NT (implying wgrad) to cuDNN-FE's wgrad API instead of the general GEMM one - if layout == "NT" and bias is None: # The wgrad kernel has no bias epilogue + # Trivial cases + if K == 0: + if bias is not None: + out_2d = out.view(N, M) + bias_2d = bias.view(1, M) + if accumulate: + out_2d += bias_2d + else: + out_2d.copy_(bias_2d) + elif not accumulate: + out.zero_() + return out, None, None, None + if M == 0 or N == 0: + return out, None, None, None + + # Route to cuDNN-FE's wgrad API for cases not supported by the + # grouped GEMM (accumulation to output tensor, insufficient + # alignment). The wgrad kernel has no bias epilogue, so any bias + # has to be applied after the GEMM. + if accumulate or N % 256 != 0: alpha = alpha if alpha is not None else 1.0 # This path uses cuDNN's wgrad (grouped_gemm_wgrad_wrapper_sm100) which supports grad accumulation - if accumulate: # Accumulate GEMM's result to the out tensorn + if accumulate: # Accumulate GEMM's result to the out tensor assert beta in (1.0, None), "beta must be one or None if accumulate is True" else: # Overwrite GEMM's result to the out tensor assert beta in (0.0, None), "beta must be zero or None if not accumulate" - return _cuDNN_wgrad_gemm( - a_tensor=dataB, - b_tensor=dataA, + _cuDNN_wgrad_gemm( + a_tensor=dataB.view(N, K // 2), + b_tensor=dataA.view(M, K // 2), sfa=sfB, sfb=sfA, amax_a=amaxB, amax_b=amaxA, out_dtype=out_dtype, - out=out, + out=out.view(N, M), accumulate=accumulate, - alpha=alpha + alpha=alpha, + bias=bias, ) + return out, None, None, None alpha = alpha if alpha is not None else 1.0 # cuDNN's general GEMM path (grouped_gemm_quant_wrapper_sm100) doesn't support accumulation assert accumulate is False, "cuDNN GEMM currently does not support accumulation for this operation." assert beta in (0.0, None), "beta must be zero or None if not accumulate" - # The output keeps B's leading dims. They only exist when B is read row-wise: - # the column-wise NVFP4 buffer is physically transposed and always 2D, and the - # layouts that select it give B the logical shape (K, M). - out_shape = (*dataB.shape[:-1], N) if not transb else (M, N) - # cuDNN's grouped quant kernel requires M to be divisible by 256 so we need to pad it - M_padded = ceil_div(M, 256) * 256 - if M_padded != M: - k_packed = dataB.shape[-1] - src = dataB.reshape(M, k_packed) - buf = src.new_zeros((M_padded, k_packed)) - buf[:M].copy_(src) + N_padded = ceil_div(N, 256) * 256 + if N_padded != N: + src = dataB.reshape(N, K // 2) + buf = src.new_zeros((N_padded, K // 2)) + buf[:N].copy_(src) dataB = buf # Swizzled scales are blocked by 128 rows: # (1, ceil(M/128), k_sf_tiles, 32, 4, 4) per_block = ceil_div(K, 4 * NVFP4_BLOCK_SCALING_SIZE) * 32 * 4 * 4 - n_blk, n_blk_padded = ceil_div(M, 128), ceil_div(M_padded, 128) + n_blk, n_blk_padded = ceil_div(N, 128), ceil_div(N_padded, 128) src_sf = sfB.reshape(-1)[: n_blk * per_block].reshape(n_blk, per_block) buf_sf = src_sf.new_zeros((n_blk_padded, per_block)) buf_sf[:n_blk].copy_(src_sf) @@ -533,7 +600,7 @@ def general_cuDNN_MX_gemm( sfB, data_dtype=torch.float4_e2m1fn_x2, scale_dtype=torch.float8_e4m3fn, # e5m3 rides as e4m3; torch has no ue5m3 - valid_M_or_N=M_padded, + valid_M_or_N=N_padded, k_logical=K, L=1, sf_swizzled=True, # ensured above @@ -543,7 +610,7 @@ def general_cuDNN_MX_gemm( sfA, data_dtype=torch.float4_e2m1fn_x2, scale_dtype=torch.float8_e4m3fn, # e5m3 rides as e4m3; torch has no ue5m3 - valid_M_or_N=N, + valid_M_or_N=M, k_logical=K, L=1, sf_swizzled=True, # ensured above @@ -567,20 +634,20 @@ def general_cuDNN_MX_gemm( alpha_tensor = (alpha * scaleA * scaleB).to(torch.float32) if bias is not None: - assert bias.dim() == 1 and bias.shape[0] == N, ( - f"cuDNN MX GEMM expects a ({N},) bias, but got {tuple(bias.shape)}." + assert bias.dim() == 1 and bias.shape[0] == M, ( + f"cuDNN MX GEMM expects a ({M},) bias, but got {tuple(bias.shape)}." ) # cuDNN checks the stride literally, so (1, N) rather than reshape's (1, 1). - bias = bias.contiguous().as_strided((N, 1), (1, N)) + bias = bias.contiguous().as_strided((M, 1), (1, M)) # Prepare for output out = validate_or_alloc_output(out, out_shape, out_dtype, device) - if M_padded != M: - # The kernel writes M_padded rows, so it cannot target `out` directly. - d_buf = torch.empty((M_padded, N), dtype=out_dtype, device=device) - d_tensor = d_buf.as_strided((M_padded, N, 1), (N, 1, M_padded * N)) + if N_padded != N: + # The kernel writes N_padded rows, so it cannot target `out` directly. + d_buf = torch.empty((N_padded, M), dtype=out_dtype, device=device) + d_tensor = d_buf.as_strided((N_padded, M, 1), (M, 1, N_padded * M)) else: - d_tensor = out.view(M, N).as_strided((M, N, 1), (N, 1, M * N)) + d_tensor = out.view(N, M).as_strided((N, M, 1), (M, 1, M * N)) gemm_kwargs = { "a_tensor": cudnn_a, @@ -588,7 +655,7 @@ def general_cuDNN_MX_gemm( "b_tensor": cudnn_b, "sfb_tensor": cudnn_sfb, # One group, so the only padded end offset is the full row count. - "padded_offsets": torch.tensor([M_padded], dtype=torch.int32, device=device), + "padded_offsets": torch.tensor([N_padded], dtype=torch.int32, device=device), "alpha_tensor": alpha_tensor, "bias_tensor": bias, "norm_const_tensor": None, # must be None for FP4 inputs @@ -604,10 +671,10 @@ def general_cuDNN_MX_gemm( } grouped_gemm_quant_kernel()(**gemm_kwargs) - if M_padded != M: + if N_padded != N: # Drop the zero-padded rows. Safe to overwrite rather than accumulate: # this path asserts accumulate is False above. - out.view(M, N).copy_(d_buf[:M]) + out.view(N, M).copy_(d_buf[:N]) # Matches general_gemm's contract: (out, bias_grad, gelu_input, extra_output). return out, None, None, None @@ -635,18 +702,31 @@ def general_gemm( ) -> Iterable[Optional[torch.Tensor]]: """GEMM supporting fp8 inputs.""" - route_to_cuDNN = False - # Route NVFP4 GEMM with e5m3 scale factors to cuDNN since cuBLAS is not ready yet. - # Test against the storage class, not NVFP4Tensor: the ops and module paths hand - # this function bare NVFP4TensorStorage operands, and NVFP4Tensor subclasses it. - if isinstance(A, NVFP4TensorStorage) and isinstance(B, NVFP4TensorStorage): - if ( - A.get_metadata()["scale_dtype"] == DType.kFloat8UE5M3 - and B.get_metadata()["scale_dtype"] == DType.kFloat8UE5M3 - ): - route_to_cuDNN = True - - if route_to_cuDNN: + assert layout in ("TN", "NN", "NT"), f"GEMM layout {layout} not supported." + transa = layout[0] == "T" + transb = layout[1] == "T" + + debug_quantizer = None + if isinstance(quantization_params, DebugQuantizer): + debug_quantizer = quantization_params + quantization_params = quantization_params.parent_quantizer + + A = _unwrap_tensor(A, "rowwise" if transa else "columnwise") + B = _unwrap_tensor(B, "columnwise" if transb else "rowwise") + + alpha = validate_gemm_scale(alpha, True) + beta = validate_gemm_scale(beta, accumulate) + workspace = get_cublas_workspace(A.device.index, ub is not None, False) + + # Temporary hack to route NVFP4 GEMM with UE5M3 scale factors to + # cuDNN Frontend kernels. UE5M3-specific logic should be removed + # in its entirety once TE supports NVFP4-UE5M3 GEMMs natively. + if ( + isinstance(A, NVFP4TensorStorage) + and isinstance(B, NVFP4TensorStorage) + and A._scale_dtype == DType.kFloat8UE5M3 + and B._scale_dtype == DType.kFloat8UE5M3 + ): return general_cuDNN_MX_gemm( A, B, @@ -668,22 +748,6 @@ def general_gemm( bulk_overlap, ) - assert layout in ("TN", "NN", "NT"), f"GEMM layout {layout} not supported." - transa = layout[0] == "T" - transb = layout[1] == "T" - - debug_quantizer = None - if isinstance(quantization_params, DebugQuantizer): - debug_quantizer = quantization_params - quantization_params = quantization_params.parent_quantizer - - A = _unwrap_tensor(A, "rowwise" if transa else "columnwise") - B = _unwrap_tensor(B, "columnwise" if transb else "rowwise") - - alpha = validate_gemm_scale(alpha, True) - beta = validate_gemm_scale(beta, accumulate) - workspace = get_cublas_workspace(A.device.index, ub is not None, False) - if ub_type is not None: assert ub is not None, ( f"{'AG+GEMM' if ub_type == tex.CommOverlapType.AG else 'GEMM+RS'} overlap requires" @@ -798,7 +862,7 @@ def general_gemm( gemm_kwargs = dict(kwargs) gemm_kwargs["beta"] = 0.0 out, bias_grad, gelu_input, extra_output = tex.generic_gemm(*gemm_args, **gemm_kwargs) - out_2d = out.reshape(-1, out.shape[-1]) + out_2d = out.view(N, M) assert output_row_scales.numel() in (1, out_2d.shape[0]) assert output_col_scales.numel() in (1, out_2d.shape[1]) @@ -886,71 +950,48 @@ def general_grouped_gemm( if any(_is_nvfp4_row_scaled_tensor(tensor) for tensor in A): raise NotImplementedError("Row-scaled NVFP4 grouped GEMM does not support row-scaled A.") - if any(_is_nvfp4_row_scaled_tensor(tensor) for tensor in B): - assert D_dtype is None, "Row-scaled NVFP4 grouped GEMM currently does not support D_dtype." - if single_output: - assert ( - m_splits is not None - ), "Row-scaled NVFP4 grouped GEMM requires m_splits with single output." - out_init = out[0] if single_output else None - if single_output: - start_idx = 0 - out_views = [] - for i in range(num_gemms): - size = m_splits[i] - out_views.append(out_init[start_idx : start_idx + size]) - start_idx += size - else: - out_views = out - for i in range(num_gemms): - if out_views[i].numel() == 0: - continue - general_gemm( - A[i], - B[i], - quantization_params=quantization_params[i], - out_dtype=out_views[i].dtype, - out=out_views[i], - gelu=gelu, - accumulate=accumulate, - layout=layout, - bias=bias[i] if use_bias else None, - use_split_accumulator=use_split_accumulator, - grad=grad, - ) - if single_output: - out = out_init - return out, grad_bias, gelu_input + # Determine whether to repeatedly call general_gemm + use_general_gemm_impl = False if isinstance(quantization_params[0], DebugQuantizer): - assert not gelu, "GELU not supported in debug mode" + use_general_gemm_impl = True + elif any(_is_nvfp4_row_scaled_tensor(tensor) for tensor in B): + use_general_gemm_impl = True + elif any( + isinstance(t, NVFP4TensorStorage) and t._scale_dtype == DType.kFloat8UE5M3 + for t in itertools.chain(A, B) + ): + use_general_gemm_impl = True + + # Repeatedly call general_gemm if needed + if use_general_gemm_impl: + out_views = out if single_output: - out_init = out[0] start_idx = 0 - out = [None] * num_gemms + out = out[0] + out_views = [None] * num_gemms for i in range(num_gemms): size = m_splits[i] - out[i] = out_init[start_idx : start_idx + size] + out_views[i] = out[start_idx : start_idx + size] start_idx += size for i in range(num_gemms): - _, bias_or_grad, _, _ = general_gemm( + _, bias_or_grad, gelu_input_i, _ = general_gemm( A[i], B[i], quantization_params=quantization_params[i], - out_dtype=out[0].dtype, + out_dtype=out_views[i].dtype, layout=layout, accumulate=accumulate, - out=out[i], + out=out_views[i], bias=bias[i] if use_bias else None, use_split_accumulator=use_split_accumulator, grad=grad, ) if grad and use_bias: grad_bias[i] = bias_or_grad - if single_output: - out = out_init - - return out, grad_bias if grad else bias, None + if gelu: + gelu_input[i] = gelu_input_i + return out, grad_bias if grad else bias, gelu_input if gelu: gelu_input = [ diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 9860d48237..3da7056764 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -501,10 +501,11 @@ def _is_grouped_tensor_path_supported( Input/weight/grad_output quantizers are assumed to be of the same type, otherwise it would trigger a fatal error in the cuBLASLt grouped GEMM check. """ - # 1. Filter by environment variable + # Filter by environment variable if not bool(int(os.getenv("NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM", "0"))): return False - # 2. Filter out advanced features + + # Filter out advanced features if ( debug or cpu_offloading @@ -513,47 +514,59 @@ def _is_grouped_tensor_path_supported( or save_original_input ): return False - # 3. Filter by compute capability and cuBLAS version - device_capability = get_device_compute_capability() - if not (9, 0) <= device_capability <= (11, 0): - return False - cublaslt_version = tex.get_cublasLt_version() - if cublaslt_version < 130300: - return False - if device_capability < (10, 0) and cublaslt_version < 130400: - return False - # 4. Output quantization is not supported. + + # Output quantization is not supported. if any(q is not None for q in output_quantizers): return False - # 5. Filter by quantization recipes. - if fp8: - if all(isinstance(q, Float8CurrentScalingQuantizer) for q in input_quantizers): - # FP8 per-tensor scaling grouped GEMM on Hopper requires cuBLAS 13.5+. - if device_capability < (10, 0) and cublaslt_version < 130500: - return False - return True - if all(isinstance(q, Float8BlockQuantizer) for q in input_quantizers): - # Grouped FP8 block-scaling quantize kernels and cuBLASLt grouped GEMM - # scale modes are Hopper-only, and the fused path has no MXFP8-broadcast - # emulation. On Blackwell (SM100/SM110, the only other arch that reaches - # this branch) fail loudly rather than silently falling back to the - # unfused path the user explicitly opted out of. - if get_device_compute_capability() >= (10, 0): - raise RuntimeError( - "NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM=1 does not support the" - " FP8 block-scaling recipe on Blackwell GPUs: the fused grouped" - " FP8 block-scaling path is Hopper-only. Unset" - " NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM to use the unfused" - " path (emulated via MXFP8 GEMM on Blackwell)." - ) - return True - # MXFP8 and NVFP4 require Blackwell+. - if not (10, 0) <= device_capability <= (11, 0): + + device_arch = get_device_compute_capability() + + # Unquantized compute + if not with_quantized_compute: + if not (9, 0) <= device_arch <= (11, 0): + # cuBLAS supports grouped GEMM on Hopper+ return False - return all(isinstance(q, MXFP8Quantizer) for q in input_quantizers) or all( - isinstance(q, NVFP4Quantizer) and q.with_rht for q in input_quantizers - ) - return activation_dtype in (torch.bfloat16, torch.float16) + return dtype in (torch.bfloat16, torch.float16) + + # FP8 current scaling + if all(isinstance(q, Float8CurrentScalingQuantizer) for q in input_quantizers): + if not (9, 0) <= device_arch <= (11, 0): + # cuBLAS supports grouped GEMM on Hopper+ + return False + if device_arch[0] == 9 and tex.get_cublasLt_version() < 130500: + # Hopper support for grouped GEMM requires cuBLAS 13.5+ + return False + return True + + # FP8 block scaling + if all(isinstance(q, Float8BlockQuantizer) for q in input_quantizers): + # Grouped GEMM requires Hopper and cuBLAS 13.4+ + return device_arch[0] == 9 and tex.get_cublasLt_version() >= 130400 + + # MXFP8 + if all(isinstance(q, MXFP8Quantizer) for q in input_quantizers): + # MXFP8 grouped quantization requires Blackwell + return (10, 0) <= device_arch <= (11, 0) + + # NVFP4 + if all(isinstance(q, NVFP4Quantizer) for q in input_quantizers): + if not (10, 0) <= device_arch <= (11, 0): + # NVFP4 grouped quantization requires Blackwell + return False + if single_grouped_weight: + # NVFP4 graph-safe grouped quantization only supports discrete weights + return False + for q in input_quantizers: + if not q.with_rht: + # NVFP4 graph-safe grouped quantization requires RHT + return False + if q.scale_dtype != DType.kFloat8E4M3: + # NVFP4 grouped GEMM is only supported with E4M3 scales + return False + return True + + # Fall back to non-graph-safe implementation + return False @staticmethod def _make_grouped_tensor( diff --git a/transformer_engine/pytorch/ops/basic/grouped_linear.py b/transformer_engine/pytorch/ops/basic/grouped_linear.py index e1980d2943..cbff2abe81 100644 --- a/transformer_engine/pytorch/ops/basic/grouped_linear.py +++ b/transformer_engine/pytorch/ops/basic/grouped_linear.py @@ -822,36 +822,55 @@ def _is_graph_safe_path_supported( * Input/weight/grad_output quantizers are assumed to be of the same type, otherwise it would trigger a fatal error in the cuBLASLt grouped GEMM check. """ - if not (9, 0) <= get_device_compute_capability() <= (11, 0): - return False - if with_quantized_compute: - # FP8 per-tensor current scaling runs on the Hopper and Blackwell grouped GEMM - # path; the compute-capability range was already checked above. On Hopper it - # requires cuBLAS 13.5+; fall back to the legacy flow on older cuBLAS. - if all(isinstance(q, Float8CurrentScalingQuantizer) for q in input_quantizers): - if ( - get_device_compute_capability() < (10, 0) - and tex.get_cublasLt_version() < 130500 - ): + + device_arch = get_device_compute_capability() + + # Unquantized compute + if not with_quantized_compute: + if not (9, 0) <= device_arch <= (11, 0): + # cuBLAS supports grouped GEMM on Hopper+ + return False + return dtype in (torch.bfloat16, torch.float16) + + # FP8 current scaling + if all(isinstance(q, Float8CurrentScalingQuantizer) for q in input_quantizers): + if not (9, 0) <= device_arch <= (11, 0): + # cuBLAS supports grouped GEMM on Hopper+ + return False + if device_arch[0] == 9 and tex.get_cublasLt_version() < 130500: + # Hopper support for grouped GEMM requires cuBLAS 13.5+ + return False + return True + + # FP8 block scaling + if all(isinstance(q, Float8BlockQuantizer) for q in input_quantizers): + # Grouped GEMM requires Hopper and cuBLAS 13.4+ + return device_arch[0] == 9 and tex.get_cublasLt_version() >= 130400 + + # MXFP8 + if all(isinstance(q, MXFP8Quantizer) for q in input_quantizers): + # MXFP8 grouped quantization requires Blackwell + return (10, 0) <= device_arch <= (11, 0) + + # NVFP4 + if all(isinstance(q, NVFP4Quantizer) for q in input_quantizers): + if not (10, 0) <= device_arch <= (11, 0): + # NVFP4 grouped quantization requires Blackwell + return False + if single_grouped_weight: + # NVFP4 graph-safe grouped quantization only supports discrete weights + return False + for q in input_quantizers: + if not q.with_rht: + # NVFP4 graph-safe grouped quantization requires RHT return False - return True - if all(isinstance(q, Float8BlockQuantizer) for q in input_quantizers): - # Grouped FP8 block scaling is Hopper-only and needs cuBLAS 13.4+; elsewhere - # fall back to the split-quantize (MXFP8-emulated) flow. - if get_device_compute_capability() >= (10, 0): + if q.scale_dtype != DType.kFloat8E4M3: + # NVFP4 grouped GEMM is only supported with E4M3 scales return False - return tex.get_cublasLt_version() >= 130400 - # MXFP8 and NVFP4 grouped quantization kernels require Blackwell. - if not (10, 0) <= get_device_compute_capability() <= (11, 0): - return False - if all(isinstance(q, MXFP8Quantizer) for q in input_quantizers): - return True - # NVFP4 graph-safe grouped quantization requires RHT and only supports - # discrete weights; otherwise fall back to the split-quantize flow. - if all(isinstance(q, NVFP4Quantizer) and q.with_rht for q in input_quantizers): - return not single_grouped_weight - return False - return dtype in (torch.bfloat16, torch.float16) + return True + + # Fall back to non-graph-safe implementation + return False def _get_grouped_weight_for_gemm( self, From 67903fae4c85f75e67ff8c36b9d06dc3d6ae215f Mon Sep 17 00:00:00 2001 From: Tim Moon Date: Thu, 13 Aug 2026 11:43:47 +0000 Subject: [PATCH 6/6] Fix typos Co-authored-by: Codex Signed-off-by: Tim Moon --- tests/pytorch/test_grouped_mlp.py | 4 ++-- transformer_engine/pytorch/cpp_extensions/gemm.py | 3 ++- transformer_engine/pytorch/module/grouped_linear.py | 10 +++++++--- 3 files changed, 11 insertions(+), 6 deletions(-) diff --git a/tests/pytorch/test_grouped_mlp.py b/tests/pytorch/test_grouped_mlp.py index 90f4263100..e562dab8d0 100644 --- a/tests/pytorch/test_grouped_mlp.py +++ b/tests/pytorch/test_grouped_mlp.py @@ -76,7 +76,7 @@ if nvfp4_available: _grouped_mlp_quantization_list.append("nvfp4_rht") if fp8_ue5m3_available: - _quantization_list.append("nvfp4_rht_ue5m3") + _grouped_mlp_quantization_list.append("nvfp4_rht_ue5m3") @pytest.fixture(autouse=True, scope="function") @@ -121,7 +121,7 @@ def maybe_skip_quantization( elif quantization == "mxfp8": if math.prod(dims[:-1]) % 32 != 0 or dims[-1] % 32 != 0: pytest.skip("MXFP8 GEMMs require dims that are divisible by 32") - elif quantization in ("nvfp4", "nvfp4_row_scaled", "nvfp4_4over6", "nvfp4_rht"): + elif quantization in nvfp4_variant_names: if math.prod(dims[:-1]) % 16 != 0 or dims[-1] % 16 != 0: pytest.skip("NVFP4 GEMMs require dims that are divisible by 16") if ( diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index 7c7aa567fb..484f1c4503 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -862,7 +862,7 @@ def general_gemm( gemm_kwargs = dict(kwargs) gemm_kwargs["beta"] = 0.0 out, bias_grad, gelu_input, extra_output = tex.generic_gemm(*gemm_args, **gemm_kwargs) - out_2d = out.view(N, M) + out_2d = out.reshape(-1, out.shape[-1]) assert output_row_scales.numel() in (1, out_2d.shape[0]) assert output_col_scales.numel() in (1, out_2d.shape[1]) @@ -983,6 +983,7 @@ def general_grouped_gemm( layout=layout, accumulate=accumulate, out=out_views[i], + gelu=gelu, bias=bias[i] if use_bias else None, use_split_accumulator=use_split_accumulator, grad=grad, diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 3da7056764..fe871e0182 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -59,7 +59,7 @@ general_grouped_gemm, general_grouped_gemm_for_grouped_tensor, ) -from ..constants import GemmParallelModes, dist_group_type +from ..constants import DType, GemmParallelModes, dist_group_type from ..jit import no_torch_dynamo from ..cpu_offload import is_cpu_offload_enabled, mark_not_offload, start_offload from ..triton.grouped_dbias_dscales import compute_grouped_dbias @@ -476,6 +476,7 @@ def _is_grouped_tensor_path_supported( activation_dtype: torch.dtype, input_quantizers: List[Optional[Quantizer]], output_quantizers: List[Optional[Quantizer]], + single_grouped_weight: bool, ) -> bool: """Whether to use cuBLASLt grouped GEMM through GroupedTensor metadata. @@ -522,11 +523,11 @@ def _is_grouped_tensor_path_supported( device_arch = get_device_compute_capability() # Unquantized compute - if not with_quantized_compute: + if not fp8: if not (9, 0) <= device_arch <= (11, 0): # cuBLAS supports grouped GEMM on Hopper+ return False - return dtype in (torch.bfloat16, torch.float16) + return activation_dtype in (torch.bfloat16, torch.float16) # FP8 current scaling if all(isinstance(q, Float8CurrentScalingQuantizer) for q in input_quantizers): @@ -916,6 +917,7 @@ def forward( delayed_scaling_input_quantizer, unsafe_requantization_input_quantizer, debug, + single_grouped_weight, ) = non_tensor_args if fp8: backward_override = FP8GlobalStateManager.get_fp8_recipe().backward_override @@ -1016,6 +1018,7 @@ def forward( activation_dtype=activation_dtype, input_quantizers=input_quantizers, output_quantizers=output_quantizers, + single_grouped_weight=single_grouped_weight, ): return _GroupedLinear._forward_grouped_tensor( ctx, @@ -2423,6 +2426,7 @@ def forward( self._delayed_scaling_input_quantizer, self._unsafe_requantization_input_quantizer, debug, + self.single_grouped_weight, ) out, new_workspaces = linear_fn( *autograd_ctx,