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..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") @@ -136,6 +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 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: @@ -3588,7 +3595,12 @@ 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, + ) 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: diff --git a/tests/pytorch/test_grouped_mlp.py b/tests/pytorch/test_grouped_mlp.py index d09c92ad49..e562dab8d0 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: + _grouped_mlp_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: @@ -118,16 +121,18 @@ 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 ( + 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 f3d97b7269..484f1c4503 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -4,13 +4,15 @@ """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 itertools +import math 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 @@ -188,6 +190,496 @@ 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.""" + 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: 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.""" + + # 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, in_features = out.size() + + # 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, + ) + + # 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 + + +@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 + (N, M) = (N, K) @ (K, M) + (N, M) + + B: + - "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, 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 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, \ + 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." + + # 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) + + # `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) + 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 + + # 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) + + # 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 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" + _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.view(N, M), + accumulate=accumulate, + 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" + + # cuDNN's grouped quant kernel requires M to be divisible by 256 so we need to pad it + 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(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) + 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). + 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=N_padded, + 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=M, + 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. + 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] == 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((M, 1), (1, M)) + + # Prepare for output + out = validate_or_alloc_output(out, out_shape, out_dtype, device) + 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(N, M).as_strided((N, M, 1), (M, 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([N_padded], 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) + + if N_padded != N: + # Drop the zero-padded rows. Safe to overwrite rather than accumulate: + # this path asserts accumulate is False above. + 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 + + def general_gemm( A: torch.Tensor, B: torch.Tensor, @@ -226,6 +718,36 @@ def general_gemm( 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, + out_dtype, + quantization_params, + gelu, + gelu_in, + alpha, + beta, + accumulate, + layout, + out, + bias, + use_split_accumulator, + grad, + ub, + ub_type, + extra_output, + bulk_overlap, + ) + 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" @@ -428,71 +950,49 @@ 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], + gelu=gelu, 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/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/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 9860d48237..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. @@ -501,10 +502,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 +515,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 fp8: + 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 activation_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( @@ -903,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 @@ -1003,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, @@ -2410,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, 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, 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"] 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, )