diff --git a/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py b/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py index 63e69beafc..68199c7bfd 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py +++ b/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py @@ -240,8 +240,6 @@ def check_nvfp4_gemm_versus_reference( def check_nvfp4_row_scaled_grouped_gemm_matches_per_gemm( - x_dtype: torch.dtype, - w_dtype: torch.dtype, out_dtype: torch.dtype, m_splits: list[int], k: int, @@ -250,15 +248,15 @@ def check_nvfp4_row_scaled_grouped_gemm_matches_per_gemm( use_bias: bool, single_output: bool, use_4over6: bool = False, - nvfp4_4over6_err_mode: str = "MAE", + monkeypatch=None, + expected_error=None, ): te_dtype = te.DType.kFloat4E2M1 device = "cuda" torch.manual_seed(23) torch.cuda.manual_seed(23) - num_gemms = len(m_splits) - + dtype = torch.bfloat16 x_quantizer = NVFP4Quantizer( fp4_dtype=te_dtype, rowwise=True, @@ -269,7 +267,6 @@ def check_nvfp4_row_scaled_grouped_gemm_matches_per_gemm( with_post_rht_amax=False, row_scaled_nvfp4=True, nvfp4_use_4over6=use_4over6, - nvfp4_4over6_err_mode=nvfp4_4over6_err_mode, ) w_quantizer = NVFP4Quantizer( fp4_dtype=te_dtype, @@ -280,7 +277,6 @@ def check_nvfp4_row_scaled_grouped_gemm_matches_per_gemm( with_rht=False, with_post_rht_amax=False, nvfp4_use_4over6=use_4over6, - nvfp4_4over6_err_mode=nvfp4_4over6_err_mode, ) x_nvfp4 = [] @@ -288,53 +284,95 @@ def check_nvfp4_row_scaled_grouped_gemm_matches_per_gemm( bias = [] expected = [] for m in m_splits: - x = torch.randn((m, k), dtype=x_dtype, device=device) - w = torch.randn((n, k), dtype=w_dtype, device=device) + x = torch.randn((m, k), dtype=dtype, device=device) + w = torch.randn((n, k), dtype=dtype, device=device) x_nvfp4.append( x_quantizer.update_quantized( - x, x_quantizer.make_empty(x.shape, dtype=x_dtype, device=device) + x, x_quantizer.make_empty(x.shape, dtype=dtype, device=device) ) ) w_nvfp4.append( w_quantizer.update_quantized( - w, w_quantizer.make_empty(w.shape, dtype=w_dtype, device=device) + w, w_quantizer.make_empty(w.shape, dtype=dtype, device=device) ) ) - bias.append(torch.randn(n, dtype=torch.bfloat16, device=device) if use_bias else None) - expected.append( - general_gemm( - w_nvfp4[-1], - x_nvfp4[-1], - out_dtype=out_dtype, - layout="TN", - bias=bias[-1], - )[0] - ) + bias.append(torch.randn(n, dtype=dtype, device=device) if use_bias else None) + if expected_error is None: + expected.append( + general_gemm( + w_nvfp4[-1], + x_nvfp4[-1], + out_dtype=out_dtype, + layout="TN", + bias=bias[-1], + )[0] + ) if single_output: out = [torch.empty((sum(m_splits), n), dtype=out_dtype, device=device)] else: out = [torch.empty((m, n), dtype=out_dtype, device=device) for m in m_splits] - grouped_out, _, _ = general_grouped_gemm( - w_nvfp4, - x_nvfp4, - out, - quantization_params=[None] * num_gemms, - out_dtype=out_dtype, - layout="TN", - m_splits=m_splits, - bias=bias, - use_bias=use_bias, - single_output=single_output, + grouped_gemm_args = (w_nvfp4, x_nvfp4, out) + grouped_gemm_kwargs = { + "quantization_params": [None] * len(m_splits), + "out_dtype": out_dtype, + "layout": "TN", + "m_splits": m_splits, + "bias": bias, + "use_bias": use_bias, + "single_output": single_output, + } + if expected_error is not None: + with pytest.raises(NotImplementedError, match=expected_error): + general_grouped_gemm(*grouped_gemm_args, **grouped_gemm_kwargs) + return + + try: + import cudnn + except ImportError as exc: + pytest.skip(f"cudnn frontend unavailable: {exc}") + if not hasattr(cudnn, "grouped_gemm_quant_wrapper_sm100"): + pytest.skip("grouped_gemm_quant_wrapper_sm100 unavailable") + + calls = [] + original_wrapper = cudnn.grouped_gemm_quant_wrapper_sm100 + + def traced_wrapper(*args, **kwargs): + calls.append(kwargs) + return original_wrapper(*args, **kwargs) + + monkeypatch.setattr(cudnn, "grouped_gemm_quant_wrapper_sm100", traced_wrapper) + grouped_out, _, _ = general_grouped_gemm(*grouped_gemm_args, **grouped_gemm_kwargs) + + assert len(calls) == 1 + call = calls[0] + global_scale_denom = 448.0 * 6.0 + expected_alpha = ( + torch.cat([tensor._amax_rowwise.view(-1) for tensor in w_nvfp4]) / global_scale_denom + ) + expected_row_scale = ( + torch.cat([tensor._amax_rowwise.view(-1) for tensor in x_nvfp4]) / global_scale_denom + ) + torch.testing.assert_close(call["alpha_tensor"], expected_alpha, atol=0.0, rtol=0.0) + torch.testing.assert_close(call["row_scale_tensor"], expected_row_scale, atol=0.0, rtol=0.0) + torch.testing.assert_close( + call["padded_offsets"], + torch.tensor(m_splits, dtype=torch.int32, device=device).cumsum(0, dtype=torch.int32), + atol=0, + rtol=0, ) + assert call["sf_vec_size"] == 16 + assert call["b_major"] == "k" + assert torch.count_nonzero(call["prob_tensor"] != 1).item() == 0 + assert (call["bias_tensor"] is not None) == use_bias if single_output: grouped_slices = torch.split(grouped_out, m_splits, dim=0) else: grouped_slices = grouped_out for grouped, ref in zip(grouped_slices, expected): - torch.testing.assert_close(grouped, ref, atol=0.0, rtol=0.0) + torch.testing.assert_close(grouped, ref, atol=0.125, rtol=0.25) def check_nvfp4_row_scaled_gemm_matches_emulated( @@ -489,53 +527,54 @@ def test_nvfp4_gemm_versus_reference( ) +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +@pytest.mark.parametrize("use_bias", [False, True], ids=["no_bias", "bias"]) +@pytest.mark.parametrize("single_output", [False, True], ids=["list_output", "single_output"]) +def test_nvfp4_row_scaled_grouped_gemm_matches_per_gemm( + use_bias: bool, + single_output: bool, + monkeypatch, +): + if torch.cuda.get_device_capability() < (10, 0): + pytest.skip("Requires SM100+ for cuDNN grouped GEMM quant kernel.") + check_nvfp4_row_scaled_grouped_gemm_matches_per_gemm( + out_dtype=torch.bfloat16, + m_splits=[256, 256, 256, 256], + k=512, + n=512, + use_bias=use_bias, + single_output=single_output, + monkeypatch=monkeypatch, + ) + + @pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) @pytest.mark.parametrize( - "m_splits, k, n", + "m_splits, k, n, out_dtype, use_4over6, expected_error", [ - ([32, 48, 48], 128, 128), - ([64, 80, 112], 128, 256), - ([64, 80, 112], 256, 256), - ([64, 80, 112], 1024, 256), - ([256, 256, 512], 1024, 1024), - ([1024, 1536, 1536], 512, 3072), - ([16, 32, 64], 128, 96), - ([80, 96, 128], 640, 304), - ([320, 336, 352], 3072, 992), - ([64, 80, 112], 64, 256), - ([32, 48, 48], 128, 112), + ([128, 256], 128, 128, torch.bfloat16, False, "M multiples of 256"), + ([256, 256], 64, 128, torch.bfloat16, False, "K and N multiples of 128"), + ([256, 256], 128, 128, torch.float32, False, "BF16/FP16 outputs"), + ([256, 256], 128, 128, torch.bfloat16, True, "does not support 4over6"), ], ) -@pytest.mark.parametrize("x_dtype", [torch.float32, torch.bfloat16], ids=str) -@pytest.mark.parametrize("w_dtype", [torch.float32, torch.bfloat16], ids=str) -@pytest.mark.parametrize("out_dtype", [torch.float32, torch.bfloat16], ids=str) -@pytest.mark.parametrize("use_bias", [False, True], ids=["no_bias", "bias"]) -@pytest.mark.parametrize("single_output", [False, True], ids=["list_output", "single_output"]) -@pytest.mark.parametrize("use_4over6", [False, True], ids=["default", "4over6"]) -@pytest.mark.parametrize("nvfp4_4over6_err_mode", ["MAE", "MSE"], ids=["mae_err", "mse_err"]) -def test_nvfp4_row_scaled_grouped_gemm_matches_per_gemm( +def test_nvfp4_row_scaled_grouped_gemm_rejects_unsupported( m_splits: list[int], k: int, n: int, - x_dtype: torch.dtype, - w_dtype: torch.dtype, out_dtype: torch.dtype, - use_bias: bool, - single_output: bool, use_4over6: bool, - nvfp4_4over6_err_mode: str, + expected_error: str, ): check_nvfp4_row_scaled_grouped_gemm_matches_per_gemm( - x_dtype=x_dtype, - w_dtype=w_dtype, out_dtype=out_dtype, m_splits=m_splits, k=k, n=n, - use_bias=use_bias, - single_output=single_output, + use_bias=False, + single_output=True, use_4over6=use_4over6, - nvfp4_4over6_err_mode=nvfp4_4over6_err_mode, + expected_error=expected_error, ) diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index f3b066d50b..7fce054d1b 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -9,7 +9,7 @@ import functools import torch import transformer_engine_torch as tex -from ..constants import TE_DType, DType +from ..constants import NVFP4_BLOCK_SCALING_SIZE, TE_DType, DType from ..utils import get_sm_count, _empty_tensor from ..quantized_tensor import Quantizer @@ -102,6 +102,152 @@ def _nvfp4_row_scaled_gemm_inputs( ) +def _cudnn_row_scaled_nvfp4_grouped_gemm( + weights: List[NVFP4TensorStorage], + inputs: List[NVFP4TensorStorage], + outputs: List[torch.Tensor], + *, + m_splits: Optional[List[int]], + bias: Optional[List[torch.Tensor]], + single_output: bool, +) -> Union[torch.Tensor, List[torch.Tensor]]: + """Run tensor-scaled weights and row-scaled inputs with the cuDNN MoE kernel.""" + num_gemms = len(weights) + if num_gemms == 0 or len(inputs) != num_gemms: + raise ValueError("Grouped GEMM requires matching non-empty weight and input lists.") + if not all(isinstance(tensor, NVFP4TensorStorage) for tensor in weights + inputs): + raise TypeError("cuDNN row-scaled NVFP4 grouped GEMM requires NVFP4 inputs.") + if any(_is_nvfp4_row_scaled_tensor(tensor) for tensor in weights): + raise NotImplementedError( + "cuDNN row-scaled NVFP4 grouped GEMM requires tensor-scaled weights." + ) + if not all(_is_nvfp4_row_scaled_tensor(tensor) for tensor in inputs): + raise NotImplementedError("cuDNN row-scaled NVFP4 grouped GEMM requires row-scaled inputs.") + if any(tensor._nvfp4_use_4over6 for tensor in weights + inputs): + raise NotImplementedError("cuDNN row-scaled NVFP4 grouped GEMM does not support 4over6.") + + m_splits_list = ( + list(m_splits) if m_splits is not None else [int(tensor.size(0)) for tensor in inputs] + ) + if len(m_splits_list) != num_gemms: + raise ValueError("m_splits length must match the number of grouped GEMMs.") + if any(m % 256 != 0 for m in m_splits_list): + raise NotImplementedError( + "cuDNN row-scaled NVFP4 grouped GEMM requires M multiples of 256." + ) + + k = int(inputs[0].size(1)) + n = int(weights[0].size(0)) + if k % 128 != 0 or n % 128 != 0: + raise NotImplementedError( + "cuDNN row-scaled NVFP4 grouped GEMM requires K and N multiples of 128." + ) + if any(tuple(tensor.size()) != (n, k) for tensor in weights): + raise ValueError("All grouped GEMM weights must have the same (N, K) shape.") + if any( + int(tensor.size(0)) != m or int(tensor.size(1)) != k + for tensor, m in zip(inputs, m_splits_list) + ): + raise ValueError("Grouped GEMM input shapes must match m_splits and K.") + expected_output_rows = [sum(m_splits_list)] if single_output else m_splits_list + if len(outputs) != len(expected_output_rows) or any( + output.shape[-1] != n or output.numel() != m * n + for output, m in zip(outputs, expected_output_rows) + ): + raise ValueError("Grouped GEMM output shapes do not match m_splits and N.") + if outputs[0].dtype not in (torch.bfloat16, torch.float16) or any( + output.dtype != outputs[0].dtype for output in outputs + ): + raise NotImplementedError( + "cuDNN row-scaled NVFP4 grouped GEMM supports uniform BF16/FP16 outputs only." + ) + if bias is not None and ( + len(bias) != num_gemms + or any(tensor is None or tuple(tensor.size()) != (n,) for tensor in bias) + ): + raise ValueError("Grouped GEMM bias tensors must have shape (N,).") + + from cudnn import ( # pylint: disable=import-outside-toplevel,no-name-in-module + grouped_gemm_quant_wrapper_sm100, + ) + + device = inputs[0].device + total_m = sum(m_splits_list) + + a_data = torch.cat( + [tensor._rowwise_data.view(m, k // 2) for tensor, m in zip(inputs, m_splits_list)], + dim=0, + ) + a_tensor = a_data.view(torch.float4_e2m1fn_x2).unsqueeze(0).permute(1, 2, 0) + + sfa_tensor = ( + torch.cat([tensor._rowwise_scale_inv for tensor in inputs], dim=0) + .view(dtype=torch.float8_e4m3fn) + .view(1, total_m // 128, 4, 32, k // (4 * NVFP4_BLOCK_SCALING_SIZE), 4) + .permute(0, 1, 4, 3, 2, 5) + .contiguous() + .permute(3, 4, 1, 5, 2, 0) + ) + + b_ptrs, sfb_ptrs, _sfb_buffer = ( + tex.grouped_mlp_experimental.swizzle_scales_and_pack_ptrs_for_discrete_weights( + [tensor._rowwise_data for tensor in weights], + [tensor._rowwise_scale_inv for tensor in weights], + "nvfp4", + device, + ) + ) + + weight_amaxes = [tensor._amax_rowwise for tensor in weights] + input_amaxes = [tensor._amax_rowwise for tensor in inputs] + if any(amax is None or amax.numel() != 1 for amax in weight_amaxes): + raise ValueError("Row-scaled NVFP4 grouped GEMM requires tensor-scaled weights.") + if any(amax is None or amax.numel() != m for amax, m in zip(input_amaxes, m_splits_list)): + raise ValueError("Row-scaled NVFP4 grouped GEMM requires one input scale per row.") + global_scale_denom = 448.0 * 6.0 + alpha_tensor = torch.cat([amax.view(-1) for amax in weight_amaxes]) / global_scale_denom + row_scale_tensor = torch.cat([amax.view(-1) for amax in input_amaxes]) / global_scale_denom + + bias_tensor = None if bias is None else torch.stack(bias, dim=0).transpose(0, 1) + + padded_offsets = torch.tensor(m_splits_list, dtype=torch.int32, device=device).cumsum( + 0, dtype=torch.int32 + ) + prob_tensor = _get_fp32_ones_tensor(total_m, device).view(total_m, 1, 1) + + result = grouped_gemm_quant_wrapper_sm100( + a_tensor=a_tensor, + b_ptrs=b_ptrs, + sfa_tensor=sfa_tensor, + sfb_ptrs=sfb_ptrs, + padded_offsets=padded_offsets, + alpha_tensor=alpha_tensor, + bias_tensor=bias_tensor, + norm_const_tensor=None, + prob_tensor=prob_tensor, + row_scale_tensor=row_scale_tensor, + acc_dtype=torch.float32, + d_dtype=outputs[0].dtype, + cd_major="n", + sf_vec_size=NVFP4_BLOCK_SCALING_SIZE, + discrete_col_sfd=False, + b_dtype=torch.float4_e2m1fn_x2, + b_major="k", + n=n, + current_stream=torch.cuda.current_stream().cuda_stream, + use_dynamic_sched=True, + ) + d_tensor = result["d_tensor"].squeeze(-1) + + if single_output: + outputs[0].view(total_m, n).copy_(d_tensor) + return outputs[0] + + for output, output_data in zip(outputs, d_tensor.split(m_splits_list)): + output.view(-1, n).copy_(output_data) + return outputs + + def general_gemm( A: torch.Tensor, B: torch.Tensor, @@ -295,12 +441,45 @@ def general_grouped_gemm( """ num_gemms = len(A) - transa = layout[0] == "T" - transb = layout[1] == "T" - empty_tensor = _empty_tensor() empty_tensors = [empty_tensor] * num_gemms + 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): + if D_dtype is not None: + raise NotImplementedError( + "cuDNN row-scaled NVFP4 grouped GEMM does not support D_dtype." + ) + if layout != "TN": + raise NotImplementedError( + "cuDNN row-scaled NVFP4 grouped GEMM supports TN layout only." + ) + if grad or gelu or accumulate or use_split_accumulator: + raise NotImplementedError( + "cuDNN row-scaled NVFP4 grouped GEMM supports fprop without GELU, " + "accumulation, or split accumulator only." + ) + if any(quantizer is not None for quantizer in quantization_params): + raise NotImplementedError( + "cuDNN row-scaled NVFP4 grouped GEMM does not support output quantization." + ) + return ( + _cudnn_row_scaled_nvfp4_grouped_gemm( + A, + B, + out, + m_splits=m_splits, + bias=bias if use_bias else None, + single_output=single_output, + ), + empty_tensors, + empty_tensors, + ) + + transa = layout[0] == "T" + transb = layout[1] == "T" + # Use bfloat16 as default bias_dtype gelu_input = empty_tensors out_dtype = TE_DType[out[0].dtype] if D_dtype is None else D_dtype @@ -320,44 +499,6 @@ def general_grouped_gemm( else: bias_dtype = TE_DType[torch.bfloat16] - 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 - if isinstance(quantization_params[0], DebugQuantizer): assert not gelu, "GELU not supported in debug mode" if single_output: diff --git a/transformer_engine/pytorch/quantization.py b/transformer_engine/pytorch/quantization.py index f1bffb122f..2642ed92d9 100644 --- a/transformer_engine/pytorch/quantization.py +++ b/transformer_engine/pytorch/quantization.py @@ -264,6 +264,8 @@ def get_align_size_for_quantization(recipe: Recipe) -> int: if recipe.mxfp8(): return 32 if recipe.nvfp4(): + if recipe.row_scaled_activation: + return 256 return 128 return 16