From d8f52acbd61ec05b78a819b310fbd7b508576f70 Mon Sep 17 00:00:00 2001 From: Ziang Li Date: Fri, 22 May 2026 21:35:06 -0700 Subject: [PATCH 1/7] Use cuDNN for row-scaled NVFP4 grouped GEMM Signed-off-by: Ziang Li --- tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py | 129 ++++++++++- .../pytorch/cpp_extensions/gemm.py | 202 +++++++++++++++++- 2 files changed, 329 insertions(+), 2 deletions(-) diff --git a/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py b/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py index eb480060e2..0d0dee4bce 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py +++ b/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py @@ -330,12 +330,139 @@ def check_nvfp4_row_scaled_grouped_gemm_matches_per_gemm( single_output=single_output, ) + if single_output: + grouped_slices = torch.split(grouped_out, m_splits, dim=0) + else: + grouped_slices = grouped_out + uses_cudnn_grouped_path = ( + out_dtype in (torch.bfloat16, torch.float16) + and not use_4over6 + and all(m % 256 == 0 for m in m_splits) + and k % 128 == 0 + and n % 128 == 0 + ) + atol = 0.5 if uses_cudnn_grouped_path else 0.0 + rtol = 0.25 if uses_cudnn_grouped_path else 0.0 + for grouped, ref in zip(grouped_slices, expected): + torch.testing.assert_close(grouped, ref, atol=atol, rtol=rtol) + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +@pytest.mark.parametrize( + "use_bias, single_output", + [(False, False), (True, True)], + ids=["no_bias_list_output", "bias_single_output"], +) +def test_nvfp4_row_scaled_grouped_gemm_uses_cudnn_quant_wrapper( + 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.") + + 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") + + te_dtype = tex.DType.kFloat4E2M1 + device = "cuda" + dtype = torch.bfloat16 + m_splits = [256, 512] + k = 128 + n = 128 + torch.manual_seed(29) + torch.cuda.manual_seed(29) + + x_quantizer = NVFP4Quantizer( + fp4_dtype=te_dtype, + rowwise=True, + columnwise=False, + with_amax_reduction=False, + amax_reduction_group=None, + with_rht=False, + with_post_rht_amax=False, + row_scaled_nvfp4=True, + ) + w_quantizer = NVFP4Quantizer( + fp4_dtype=te_dtype, + rowwise=True, + columnwise=True, + with_amax_reduction=False, + amax_reduction_group=None, + with_rht=False, + with_post_rht_amax=False, + ) + + x_nvfp4 = [] + w_nvfp4 = [] + bias = [] + expected = [] + for m in m_splits: + 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=dtype, device=device), + ) + ) + w_nvfp4.append( + w_quantizer.update_quantized( + 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=dtype, + layout="TN", + bias=bias[-1], + )[0] + ) + + 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) + if single_output: + out = [torch.empty((sum(m_splits), n), dtype=dtype, device=device)] + else: + out = [torch.empty((m, n), dtype=dtype, device=device) for m in m_splits] + grouped_out, _, _ = general_grouped_gemm( + w_nvfp4, + x_nvfp4, + out, + quantization_params=[None] * len(m_splits), + out_dtype=dtype, + layout="TN", + m_splits=m_splits, + bias=bias, + use_bias=use_bias, + single_output=single_output, + ) + + assert len(calls) == 1 + assert calls[0]["sf_vec_size"] == 16 + assert calls[0]["row_scale_tensor"].shape == (sum(m_splits),) + assert calls[0]["b_major"] == "k" + assert (calls[0]["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.5, rtol=0.25) def check_nvfp4_row_scaled_gemm_matches_emulated( diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index f3d97b7269..41bb8df0b5 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 QuantizedTensorStorage, Quantizer @@ -186,6 +186,188 @@ def _validate_native_gemm_output_quantizer(quantization_params): "Return a TE-native quantizer for output/grad_input roles or disable " "quantized GEMM output for this boundary." ) +def _ceil_div(a: int, b: int) -> int: + """Integer ceil division.""" + return (a + b - 1) // b + + +def _nvfp4_cudnn_scale_layout(scale: torch.Tensor, m: int, k: int) -> torch.Tensor: + """Pack compact NVFP4 scales into the cuDNN CuTe layout.""" + m_tiles = _ceil_div(m, 128) + sf_k = _ceil_div(k, NVFP4_BLOCK_SCALING_SIZE) + k_tiles = _ceil_div(sf_k, 4) + compact = scale.view(torch.float8_e4m3fn)[: m_tiles * 128, : k_tiles * 4].contiguous() + logical_layout = compact.view(1, m_tiles, 4, 32, k_tiles, 4).permute(3, 2, 1, 5, 4, 0) + base = torch.empty( + (1, m_tiles, k_tiles, 32, 4, 4), + dtype=torch.float8_e4m3fn, + device=scale.device, + ) + cudnn_layout = base.permute(3, 4, 1, 5, 2, 0) + cudnn_layout.copy_(logical_layout) + return cudnn_layout + + +def _nvfp4_rowwise_data_logical_view(tensor: NVFP4TensorStorage) -> torch.Tensor: + """Return a logical FP4 rowwise data view with the packed buffer's data pointer.""" + packed = tensor._rowwise_data + rows = int(tensor.size(0)) + cols = int(tensor.size(1)) + return torch.as_strided(packed, (rows, cols), (packed.stride(0), 0)) + + +def _try_cudnn_grouped_gemm_quant_for_row_scaled_nvfp4( + A: List[torch.Tensor], + B: List[torch.Tensor], + out: List[torch.Tensor], + *, + transa: bool, + transb: bool, + m_splits: Optional[List[int]], + bias: List[torch.Tensor], + use_bias: bool, + single_output: bool, + accumulate: bool, + gelu: bool, + grad: bool, + use_split_accumulator: bool, +) -> Optional[torch.Tensor]: + """Use cuDNN grouped GEMM quant for supported row-scaled NVFP4 grouped GEMMs. + + Returns ``None`` when the inputs are outside the currently supported cuDNN + path so callers can fall back to the existing per-GEMM implementation. + """ + if grad or gelu or accumulate or use_split_accumulator: + return None + if not transa or transb: + return None + if not out or out[0].dtype not in (torch.bfloat16, torch.float16): + return None + if not all(isinstance(tensor, NVFP4TensorStorage) for tensor in A + B): + return None + if any(_is_nvfp4_row_scaled_tensor(tensor) for tensor in A): + return None + if not all(_is_nvfp4_row_scaled_tensor(tensor) for tensor in B): + return None + if any(getattr(tensor, "_nvfp4_use_4over6", False) for tensor in A + B): + return None + + num_gemms = len(A) + m_splits_list = ( + list(m_splits) if m_splits is not None else [int(tensor.size(0)) for tensor in B] + ) + if len(m_splits_list) != num_gemms: + return None + if any(m % 256 != 0 for m in m_splits_list): + return None + + k = int(B[0].size(1)) + n = int(A[0].size(0)) + if k % 128 != 0 or n % 128 != 0: + return None + if any(tuple(tensor.size()) != (n, k) for tensor in A): + return None + if any( + int(tensor.size(0)) != m or int(tensor.size(1)) != k for tensor, m in zip(B, m_splits_list) + ): + return None + + try: + from cudnn import ( + grouped_gemm_quant_wrapper_sm100, + ) # pylint: disable=import-outside-toplevel + except ImportError: + return None + + device = B[0]._rowwise_data.device + total_m = sum(m_splits_list) + + a_data = torch.cat( + [tensor._rowwise_data.view(m, k // 2) for tensor, m in zip(B, m_splits_list)], + dim=0, + ) + a_tensor = a_data.view(torch.float4_e2m1fn_x2).unsqueeze(0).permute(1, 2, 0) + + sf_cols = _ceil_div(_ceil_div(k, NVFP4_BLOCK_SCALING_SIZE), 4) * 4 + sfa_compact = torch.cat( + [ + tensor._rowwise_scale_inv.view(torch.float8_e4m3fn)[:m, :sf_cols] + for tensor, m in zip(B, m_splits_list) + ], + dim=0, + ) + sfa_tensor = _nvfp4_cudnn_scale_layout(sfa_compact, total_m, k) + + sfb_tensors = [_nvfp4_cudnn_scale_layout(tensor._rowwise_scale_inv, n, k) for tensor in A] + b_ptrs, sfb_ptrs, _sfb_keepalive = tex.get_device_pointer_for_data_and_scales( + [_nvfp4_rowwise_data_logical_view(tensor) for tensor in A], + sfb_tensors, + False, + True, + A[0]._fp4_dtype, + ) + + row_scale = [] + for weight, activation in zip(A, B): + weight_amax = weight._amax_rowwise if transa else weight._amax_columnwise + if weight_amax is None or activation._amax_rowwise is None: + return None + if weight_amax.numel() != 1: + return None + activation_decode_scale = activation._amax_rowwise / ( + float(activation._nvfp4_e4m3_max) * 6.0 + ) + weight_decode_scale = weight_amax / (float(weight._nvfp4_e4m3_max) * 6.0) + row_scale.append((activation_decode_scale * weight_decode_scale).to(dtype=torch.float32)) + row_scale_tensor = torch.cat(row_scale).contiguous() + + bias_tensor = None + if use_bias: + if any(tensor.numel() == 0 for tensor in bias): + return None + bias_tensor = torch.stack(bias, dim=0).transpose(0, 1) + + padded_offsets = torch.tensor( + [sum(m_splits_list[: i + 1]) for i in range(num_gemms)], + dtype=torch.int32, + device=device, + ) + alpha_tensor = torch.ones(num_gemms, dtype=torch.float32, device=device) + prob_tensor = torch.ones(total_m, 1, 1, dtype=torch.float32, device=device) + + 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=out[0].dtype, + cd_major="n", + sf_vec_size=NVFP4_BLOCK_SCALING_SIZE, + discrete_col_sfd=True, + 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: + out[0].copy_(d_tensor) + return out[0] + + start = 0 + for output, m in zip(out, m_splits_list): + output.copy_(d_tensor[start : start + m]) + start += m + return out def general_gemm( @@ -434,6 +616,24 @@ def general_grouped_gemm( assert ( m_splits is not None ), "Row-scaled NVFP4 grouped GEMM requires m_splits with single output." + cudnn_out = _try_cudnn_grouped_gemm_quant_for_row_scaled_nvfp4( + A, + B, + out, + transa=transa, + transb=transb, + m_splits=m_splits, + bias=bias, + use_bias=use_bias, + single_output=single_output, + accumulate=accumulate, + gelu=gelu, + grad=grad, + use_split_accumulator=use_split_accumulator, + ) + if cudnn_out is not None: + return cudnn_out, grad_bias, gelu_input + out_init = out[0] if single_output else None if single_output: start_idx = 0 From 04bed003438ac0c62122d945cf40e850f6afbf1f Mon Sep 17 00:00:00 2001 From: Ziang Li Date: Mon, 25 May 2026 20:58:57 -0700 Subject: [PATCH 2/7] Require cuDNN for row-scaled NVFP4 grouped GEMM Signed-off-by: Ziang Li --- tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py | 68 ++++++---- .../pytorch/cpp_extensions/gemm.py | 124 +++++++----------- 2 files changed, 92 insertions(+), 100 deletions(-) diff --git a/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py b/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py index 0d0dee4bce..85309f10c3 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py +++ b/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py @@ -259,6 +259,13 @@ def check_nvfp4_row_scaled_grouped_gemm_matches_per_gemm( torch.cuda.manual_seed(23) num_gemms = len(m_splits) + uses_cudnn_grouped_path = ( + out_dtype in (torch.bfloat16, torch.float16) + and not use_4over6 + and all(m % 256 == 0 for m in m_splits) + and k % 128 == 0 + and n % 128 == 0 + ) x_quantizer = NVFP4Quantizer( fp4_dtype=te_dtype, @@ -302,49 +309,56 @@ def check_nvfp4_row_scaled_grouped_gemm_matches_per_gemm( ) ) 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] - ) + if uses_cudnn_grouped_path: + 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( + grouped_gemm_args = ( 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_kwargs = { + "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, + } + if not uses_cudnn_grouped_path: + with pytest.raises((NotImplementedError, ValueError)): + 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") + + grouped_out, _, _ = general_grouped_gemm(*grouped_gemm_args, **grouped_gemm_kwargs) if single_output: grouped_slices = torch.split(grouped_out, m_splits, dim=0) else: grouped_slices = grouped_out - uses_cudnn_grouped_path = ( - out_dtype in (torch.bfloat16, torch.float16) - and not use_4over6 - and all(m % 256 == 0 for m in m_splits) - and k % 128 == 0 - and n % 128 == 0 - ) - atol = 0.5 if uses_cudnn_grouped_path else 0.0 - rtol = 0.25 if uses_cudnn_grouped_path else 0.0 for grouped, ref in zip(grouped_slices, expected): - torch.testing.assert_close(grouped, ref, atol=atol, rtol=rtol) + torch.testing.assert_close(grouped, ref, atol=0.5, rtol=0.25) @pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index 41bb8df0b5..26383c2f6d 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -216,7 +216,7 @@ def _nvfp4_rowwise_data_logical_view(tensor: NVFP4TensorStorage) -> torch.Tensor return torch.as_strided(packed, (rows, cols), (packed.stride(0), 0)) -def _try_cudnn_grouped_gemm_quant_for_row_scaled_nvfp4( +def _cudnn_grouped_gemm_quant_for_row_scaled_nvfp4( A: List[torch.Tensor], B: List[torch.Tensor], out: List[torch.Tensor], @@ -231,53 +231,57 @@ def _try_cudnn_grouped_gemm_quant_for_row_scaled_nvfp4( gelu: bool, grad: bool, use_split_accumulator: bool, -) -> Optional[torch.Tensor]: - """Use cuDNN grouped GEMM quant for supported row-scaled NVFP4 grouped GEMMs. - - Returns ``None`` when the inputs are outside the currently supported cuDNN - path so callers can fall back to the existing per-GEMM implementation. - """ +) -> torch.Tensor: + """Use cuDNN grouped GEMM quant for row-scaled NVFP4 grouped GEMMs.""" if grad or gelu or accumulate or use_split_accumulator: - return None + raise NotImplementedError( + "cuDNN row-scaled NVFP4 grouped GEMM supports fprop without GELU, " + "accumulation, or split accumulator only." + ) if not transa or transb: - return None + raise NotImplementedError("cuDNN row-scaled NVFP4 grouped GEMM supports TN layout only.") if not out or out[0].dtype not in (torch.bfloat16, torch.float16): - return None + raise NotImplementedError( + "cuDNN row-scaled NVFP4 grouped GEMM supports BF16/FP16 outputs only." + ) if not all(isinstance(tensor, NVFP4TensorStorage) for tensor in A + B): - return None + raise TypeError("cuDNN row-scaled NVFP4 grouped GEMM requires NVFP4 inputs.") if any(_is_nvfp4_row_scaled_tensor(tensor) for tensor in A): - return None + raise NotImplementedError( + "cuDNN row-scaled NVFP4 grouped GEMM does not support row-scaled A." + ) if not all(_is_nvfp4_row_scaled_tensor(tensor) for tensor in B): - return None + raise NotImplementedError("cuDNN row-scaled NVFP4 grouped GEMM requires row-scaled B.") if any(getattr(tensor, "_nvfp4_use_4over6", False) for tensor in A + B): - return None + raise NotImplementedError("cuDNN row-scaled NVFP4 grouped GEMM does not support 4over6.") num_gemms = len(A) m_splits_list = ( list(m_splits) if m_splits is not None else [int(tensor.size(0)) for tensor in B] ) if len(m_splits_list) != num_gemms: - return None + raise ValueError("m_splits length must match the number of grouped GEMMs.") if any(m % 256 != 0 for m in m_splits_list): - return None + raise NotImplementedError( + "cuDNN row-scaled NVFP4 grouped GEMM requires M multiples of 256." + ) k = int(B[0].size(1)) n = int(A[0].size(0)) if k % 128 != 0 or n % 128 != 0: - return None + 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 A): - return None + raise ValueError("All grouped GEMM A tensors 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(B, m_splits_list) ): - return None + raise ValueError("Grouped GEMM B tensor shapes must match m_splits and K.") + + import cudnn # pylint: disable=import-outside-toplevel - try: - from cudnn import ( - grouped_gemm_quant_wrapper_sm100, - ) # pylint: disable=import-outside-toplevel - except ImportError: - return None + grouped_gemm_quant_wrapper_sm100 = cudnn.grouped_gemm_quant_wrapper_sm100 device = B[0]._rowwise_data.device total_m = sum(m_splits_list) @@ -311,9 +315,11 @@ def _try_cudnn_grouped_gemm_quant_for_row_scaled_nvfp4( for weight, activation in zip(A, B): weight_amax = weight._amax_rowwise if transa else weight._amax_columnwise if weight_amax is None or activation._amax_rowwise is None: - return None + raise ValueError( + "Row-scaled NVFP4 grouped GEMM requires activation and weight amax metadata." + ) if weight_amax.numel() != 1: - return None + raise ValueError("Row-scaled NVFP4 grouped GEMM requires tensor-scaled weights.") activation_decode_scale = activation._amax_rowwise / ( float(activation._nvfp4_e4m3_max) * 6.0 ) @@ -324,7 +330,7 @@ def _try_cudnn_grouped_gemm_quant_for_row_scaled_nvfp4( bias_tensor = None if use_bias: if any(tensor.numel() == 0 for tensor in bias): - return None + raise ValueError("Bias tensors must be non-empty when use_bias=True.") bias_tensor = torch.stack(bias, dim=0).transpose(0, 1) padded_offsets = torch.tensor( @@ -616,53 +622,25 @@ def general_grouped_gemm( assert ( m_splits is not None ), "Row-scaled NVFP4 grouped GEMM requires m_splits with single output." - cudnn_out = _try_cudnn_grouped_gemm_quant_for_row_scaled_nvfp4( - A, - B, - out, - transa=transa, - transb=transb, - m_splits=m_splits, - bias=bias, - use_bias=use_bias, - single_output=single_output, - accumulate=accumulate, - gelu=gelu, - grad=grad, - use_split_accumulator=use_split_accumulator, - ) - if cudnn_out is not None: - return cudnn_out, grad_bias, gelu_input - - 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, + return ( + _cudnn_grouped_gemm_quant_for_row_scaled_nvfp4( + A, + B, + out, + transa=transa, + transb=transb, + m_splits=m_splits, + bias=bias, + use_bias=use_bias, + single_output=single_output, accumulate=accumulate, - layout=layout, - bias=bias[i] if use_bias else None, - use_split_accumulator=use_split_accumulator, + gelu=gelu, grad=grad, - ) - if single_output: - out = out_init - return out, grad_bias, gelu_input + use_split_accumulator=use_split_accumulator, + ), + grad_bias, + gelu_input, + ) if isinstance(quantization_params[0], DebugQuantizer): assert not gelu, "GELU not supported in debug mode" From cc3fa3788ff9337359660f3a00f0e623cc6cfcc5 Mon Sep 17 00:00:00 2001 From: Ziang Li Date: Mon, 13 Jul 2026 14:50:49 -0700 Subject: [PATCH 3/7] Refactor row-scaled cuDNN grouped GEMM integration Signed-off-by: Ziang Li --- tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py | 218 +++++---------- .../pytorch/cpp_extensions/gemm.py | 253 ++++++++---------- transformer_engine/pytorch/quantization.py | 2 + 3 files changed, 186 insertions(+), 287 deletions(-) diff --git a/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py b/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py index 85309f10c3..f615d5ae29 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py +++ b/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py @@ -252,6 +252,8 @@ def check_nvfp4_row_scaled_grouped_gemm_matches_per_gemm( 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" @@ -259,14 +261,6 @@ def check_nvfp4_row_scaled_grouped_gemm_matches_per_gemm( torch.cuda.manual_seed(23) num_gemms = len(m_splits) - uses_cudnn_grouped_path = ( - out_dtype in (torch.bfloat16, torch.float16) - and not use_4over6 - and all(m % 256 == 0 for m in m_splits) - and k % 128 == 0 - and n % 128 == 0 - ) - x_quantizer = NVFP4Quantizer( fp4_dtype=te_dtype, rowwise=True, @@ -309,7 +303,7 @@ def check_nvfp4_row_scaled_grouped_gemm_matches_per_gemm( ) ) bias.append(torch.randn(n, dtype=torch.bfloat16, device=device) if use_bias else None) - if uses_cudnn_grouped_path: + if expected_error is None: expected.append( general_gemm( w_nvfp4[-1], @@ -339,8 +333,8 @@ def check_nvfp4_row_scaled_grouped_gemm_matches_per_gemm( "use_bias": use_bias, "single_output": single_output, } - if not uses_cudnn_grouped_path: - with pytest.raises((NotImplementedError, ValueError)): + if expected_error is not None: + with pytest.raises(NotImplementedError, match=expected_error): general_grouped_gemm(*grouped_gemm_args, **grouped_gemm_kwargs) return @@ -351,96 +345,6 @@ def check_nvfp4_row_scaled_grouped_gemm_matches_per_gemm( if not hasattr(cudnn, "grouped_gemm_quant_wrapper_sm100"): pytest.skip("grouped_gemm_quant_wrapper_sm100 unavailable") - grouped_out, _, _ = general_grouped_gemm(*grouped_gemm_args, **grouped_gemm_kwargs) - - 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.5, rtol=0.25) - - -@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) -@pytest.mark.parametrize( - "use_bias, single_output", - [(False, False), (True, True)], - ids=["no_bias_list_output", "bias_single_output"], -) -def test_nvfp4_row_scaled_grouped_gemm_uses_cudnn_quant_wrapper( - 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.") - - 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") - - te_dtype = tex.DType.kFloat4E2M1 - device = "cuda" - dtype = torch.bfloat16 - m_splits = [256, 512] - k = 128 - n = 128 - torch.manual_seed(29) - torch.cuda.manual_seed(29) - - x_quantizer = NVFP4Quantizer( - fp4_dtype=te_dtype, - rowwise=True, - columnwise=False, - with_amax_reduction=False, - amax_reduction_group=None, - with_rht=False, - with_post_rht_amax=False, - row_scaled_nvfp4=True, - ) - w_quantizer = NVFP4Quantizer( - fp4_dtype=te_dtype, - rowwise=True, - columnwise=True, - with_amax_reduction=False, - amax_reduction_group=None, - with_rht=False, - with_post_rht_amax=False, - ) - - x_nvfp4 = [] - w_nvfp4 = [] - bias = [] - expected = [] - for m in m_splits: - 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=dtype, device=device), - ) - ) - w_nvfp4.append( - w_quantizer.update_quantized( - 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=dtype, - layout="TN", - bias=bias[-1], - )[0] - ) - calls = [] original_wrapper = cudnn.grouped_gemm_quant_wrapper_sm100 @@ -449,28 +353,35 @@ def traced_wrapper(*args, **kwargs): return original_wrapper(*args, **kwargs) monkeypatch.setattr(cudnn, "grouped_gemm_quant_wrapper_sm100", traced_wrapper) - if single_output: - out = [torch.empty((sum(m_splits), n), dtype=dtype, device=device)] - else: - out = [torch.empty((m, n), dtype=dtype, device=device) for m in m_splits] - grouped_out, _, _ = general_grouped_gemm( - w_nvfp4, - x_nvfp4, - out, - quantization_params=[None] * len(m_splits), - out_dtype=dtype, - layout="TN", - m_splits=m_splits, - bias=bias, - use_bias=use_bias, - single_output=single_output, - ) + grouped_out, _, _ = general_grouped_gemm(*grouped_gemm_args, **grouped_gemm_kwargs) assert len(calls) == 1 - assert calls[0]["sf_vec_size"] == 16 - assert calls[0]["row_scale_tensor"].shape == (sum(m_splits),) - assert calls[0]["b_major"] == "k" - assert (calls[0]["bias_tensor"] is not None) == use_bias + call = calls[0] + expected_alpha = torch.cat( + [ + tensor._amax_rowwise.view(-1) / (float(tensor._nvfp4_e4m3_max) * 6.0) + for tensor in w_nvfp4 + ] + ).to(dtype=torch.float32) + expected_row_scale = torch.cat( + [ + tensor._amax_rowwise.view(-1) / (float(tensor._nvfp4_e4m3_max) * 6.0) + for tensor in x_nvfp4 + ] + ).to(dtype=torch.float32) + 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: @@ -690,53 +601,58 @@ 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( + x_dtype=torch.bfloat16, + w_dtype=torch.bfloat16, + 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, + x_dtype=torch.bfloat16, + w_dtype=torch.bfloat16, 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 26383c2f6d..663c527a13 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -186,78 +186,34 @@ def _validate_native_gemm_output_quantizer(quantization_params): "Return a TE-native quantizer for output/grad_input roles or disable " "quantized GEMM output for this boundary." ) -def _ceil_div(a: int, b: int) -> int: - """Integer ceil division.""" - return (a + b - 1) // b - - -def _nvfp4_cudnn_scale_layout(scale: torch.Tensor, m: int, k: int) -> torch.Tensor: - """Pack compact NVFP4 scales into the cuDNN CuTe layout.""" - m_tiles = _ceil_div(m, 128) - sf_k = _ceil_div(k, NVFP4_BLOCK_SCALING_SIZE) - k_tiles = _ceil_div(sf_k, 4) - compact = scale.view(torch.float8_e4m3fn)[: m_tiles * 128, : k_tiles * 4].contiguous() - logical_layout = compact.view(1, m_tiles, 4, 32, k_tiles, 4).permute(3, 2, 1, 5, 4, 0) - base = torch.empty( - (1, m_tiles, k_tiles, 32, 4, 4), - dtype=torch.float8_e4m3fn, - device=scale.device, - ) - cudnn_layout = base.permute(3, 4, 1, 5, 2, 0) - cudnn_layout.copy_(logical_layout) - return cudnn_layout - -def _nvfp4_rowwise_data_logical_view(tensor: NVFP4TensorStorage) -> torch.Tensor: - """Return a logical FP4 rowwise data view with the packed buffer's data pointer.""" - packed = tensor._rowwise_data - rows = int(tensor.size(0)) - cols = int(tensor.size(1)) - return torch.as_strided(packed, (rows, cols), (packed.stride(0), 0)) - -def _cudnn_grouped_gemm_quant_for_row_scaled_nvfp4( - A: List[torch.Tensor], - B: List[torch.Tensor], - out: List[torch.Tensor], +def _cudnn_row_scaled_nvfp4_grouped_gemm( + weights: List[NVFP4TensorStorage], + inputs: List[NVFP4TensorStorage], + outputs: List[torch.Tensor], *, - transa: bool, - transb: bool, m_splits: Optional[List[int]], - bias: List[torch.Tensor], - use_bias: bool, + bias: Optional[List[torch.Tensor]], single_output: bool, - accumulate: bool, - gelu: bool, - grad: bool, - use_split_accumulator: bool, -) -> torch.Tensor: - """Use cuDNN grouped GEMM quant for row-scaled NVFP4 grouped GEMMs.""" - 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 not transa or transb: - raise NotImplementedError("cuDNN row-scaled NVFP4 grouped GEMM supports TN layout only.") - if not out or out[0].dtype not in (torch.bfloat16, torch.float16): - raise NotImplementedError( - "cuDNN row-scaled NVFP4 grouped GEMM supports BF16/FP16 outputs only." - ) - if not all(isinstance(tensor, NVFP4TensorStorage) for tensor in A + B): +) -> 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 A): + if any(_is_nvfp4_row_scaled_tensor(tensor) for tensor in weights): raise NotImplementedError( - "cuDNN row-scaled NVFP4 grouped GEMM does not support row-scaled A." + "cuDNN row-scaled NVFP4 grouped GEMM requires tensor-scaled weights." ) - if not all(_is_nvfp4_row_scaled_tensor(tensor) for tensor in B): - raise NotImplementedError("cuDNN row-scaled NVFP4 grouped GEMM requires row-scaled B.") - if any(getattr(tensor, "_nvfp4_use_4over6", False) for tensor in A + B): + 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(getattr(tensor, "_nvfp4_use_4over6", False) for tensor in weights + inputs): raise NotImplementedError("cuDNN row-scaled NVFP4 grouped GEMM does not support 4over6.") - num_gemms = len(A) m_splits_list = ( - list(m_splits) if m_splits is not None else [int(tensor.size(0)) for tensor in B] + 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.") @@ -266,79 +222,101 @@ def _cudnn_grouped_gemm_quant_for_row_scaled_nvfp4( "cuDNN row-scaled NVFP4 grouped GEMM requires M multiples of 256." ) - k = int(B[0].size(1)) - n = int(A[0].size(0)) + 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 A): - raise ValueError("All grouped GEMM A tensors must have the same (N, K) shape.") + 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(B, m_splits_list) + int(tensor.size(0)) != m or int(tensor.size(1)) != k + for tensor, m in zip(inputs, m_splits_list) ): - raise ValueError("Grouped GEMM B tensor shapes must match m_splits and K.") - - import cudnn # pylint: disable=import-outside-toplevel + 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,).") - grouped_gemm_quant_wrapper_sm100 = cudnn.grouped_gemm_quant_wrapper_sm100 + from cudnn import ( # pylint: disable=import-outside-toplevel,no-name-in-module + grouped_gemm_quant_wrapper_sm100, + ) - device = B[0]._rowwise_data.device + device = inputs[0]._rowwise_data.device total_m = sum(m_splits_list) a_data = torch.cat( - [tensor._rowwise_data.view(m, k // 2) for tensor, m in zip(B, m_splits_list)], + [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) - sf_cols = _ceil_div(_ceil_div(k, NVFP4_BLOCK_SCALING_SIZE), 4) * 4 - sfa_compact = torch.cat( - [ - tensor._rowwise_scale_inv.view(torch.float8_e4m3fn)[:m, :sf_cols] - for tensor, m in zip(B, m_splits_list) - ], - dim=0, + sfa_compact = torch.cat([tensor._rowwise_scale_inv for tensor in inputs], dim=0) + sfa_logical = sfa_compact.view(dtype=torch.float8_e4m3fn).view( + 1, + total_m // 128, + 4, + 32, + k // (4 * NVFP4_BLOCK_SCALING_SIZE), + 4, ) - sfa_tensor = _nvfp4_cudnn_scale_layout(sfa_compact, total_m, k) - - sfb_tensors = [_nvfp4_cudnn_scale_layout(tensor._rowwise_scale_inv, n, k) for tensor in A] - b_ptrs, sfb_ptrs, _sfb_keepalive = tex.get_device_pointer_for_data_and_scales( - [_nvfp4_rowwise_data_logical_view(tensor) for tensor in A], - sfb_tensors, - False, - True, - A[0]._fp4_dtype, + sfa_logical = sfa_logical.permute(3, 2, 1, 5, 4, 0) + sfa_storage = torch.empty( + (1, total_m // 128, k // (4 * NVFP4_BLOCK_SCALING_SIZE), 32, 4, 4), + dtype=torch.float8_e4m3fn, + device=device, ) - - row_scale = [] - for weight, activation in zip(A, B): - weight_amax = weight._amax_rowwise if transa else weight._amax_columnwise - if weight_amax is None or activation._amax_rowwise is None: - raise ValueError( - "Row-scaled NVFP4 grouped GEMM requires activation and weight amax metadata." - ) - if weight_amax.numel() != 1: - raise ValueError("Row-scaled NVFP4 grouped GEMM requires tensor-scaled weights.") - activation_decode_scale = activation._amax_rowwise / ( - float(activation._nvfp4_e4m3_max) * 6.0 + sfa_tensor = sfa_storage.permute(3, 4, 1, 5, 2, 0) + sfa_tensor.copy_(sfa_logical) + + 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_decode_scale = weight_amax / (float(weight._nvfp4_e4m3_max) * 6.0) - row_scale.append((activation_decode_scale * weight_decode_scale).to(dtype=torch.float32)) - row_scale_tensor = torch.cat(row_scale).contiguous() + ) - bias_tensor = None - if use_bias: - if any(tensor.numel() == 0 for tensor in bias): - raise ValueError("Bias tensors must be non-empty when use_bias=True.") - bias_tensor = torch.stack(bias, dim=0).transpose(0, 1) + 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.") + alpha_tensor = torch.cat( + [ + amax.view(-1) / (float(tensor._nvfp4_e4m3_max) * 6.0) + for tensor, amax in zip(weights, weight_amaxes) + ] + ).to(dtype=torch.float32) + row_scale_tensor = torch.cat( + [ + amax.view(-1) / (float(tensor._nvfp4_e4m3_max) * 6.0) + for tensor, amax in zip(inputs, input_amaxes) + ] + ).to(dtype=torch.float32) - padded_offsets = torch.tensor( - [sum(m_splits_list[: i + 1]) for i in range(num_gemms)], - dtype=torch.int32, - device=device, + 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 ) - alpha_tensor = torch.ones(num_gemms, dtype=torch.float32, device=device) prob_tensor = torch.ones(total_m, 1, 1, dtype=torch.float32, device=device) result = grouped_gemm_quant_wrapper_sm100( @@ -353,10 +331,10 @@ def _cudnn_grouped_gemm_quant_for_row_scaled_nvfp4( prob_tensor=prob_tensor, row_scale_tensor=row_scale_tensor, acc_dtype=torch.float32, - d_dtype=out[0].dtype, + d_dtype=outputs[0].dtype, cd_major="n", sf_vec_size=NVFP4_BLOCK_SCALING_SIZE, - discrete_col_sfd=True, + discrete_col_sfd=False, b_dtype=torch.float4_e2m1fn_x2, b_major="k", n=n, @@ -366,14 +344,12 @@ def _cudnn_grouped_gemm_quant_for_row_scaled_nvfp4( d_tensor = result["d_tensor"].squeeze(-1) if single_output: - out[0].copy_(d_tensor) - return out[0] + outputs[0].view(total_m, n).copy_(d_tensor) + return outputs[0] - start = 0 - for output, m in zip(out, m_splits_list): - output.copy_(d_tensor[start : start + m]) - start += m - return out + 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( @@ -617,26 +593,31 @@ 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." + 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_grouped_gemm_quant_for_row_scaled_nvfp4( + _cudnn_row_scaled_nvfp4_grouped_gemm( A, B, out, - transa=transa, - transb=transb, m_splits=m_splits, - bias=bias, - use_bias=use_bias, + bias=bias if use_bias else None, single_output=single_output, - accumulate=accumulate, - gelu=gelu, - grad=grad, - use_split_accumulator=use_split_accumulator, ), grad_bias, gelu_input, diff --git a/transformer_engine/pytorch/quantization.py b/transformer_engine/pytorch/quantization.py index 07a3b80483..8263324083 100644 --- a/transformer_engine/pytorch/quantization.py +++ b/transformer_engine/pytorch/quantization.py @@ -274,6 +274,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 From 5f27dd46d922822a5aae9e7b45bf849b47dc0de2 Mon Sep 17 00:00:00 2001 From: Ziang Li Date: Mon, 13 Jul 2026 15:10:38 -0700 Subject: [PATCH 4/7] Tighten row-scaled cuDNN grouped GEMM integration Signed-off-by: Ziang Li --- tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py | 50 ++++------- .../pytorch/cpp_extensions/gemm.py | 86 ++++++++----------- 2 files changed, 50 insertions(+), 86 deletions(-) diff --git a/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py b/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py index f615d5ae29..605c6f84f3 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py +++ b/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py @@ -241,8 +241,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, @@ -251,7 +249,6 @@ 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, ): @@ -260,7 +257,7 @@ def check_nvfp4_row_scaled_grouped_gemm_matches_per_gemm( 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, @@ -271,7 +268,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, @@ -282,7 +278,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 = [] @@ -290,19 +285,19 @@ 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) + bias.append(torch.randn(n, dtype=dtype, device=device) if use_bias else None) if expected_error is None: expected.append( general_gemm( @@ -319,13 +314,9 @@ def check_nvfp4_row_scaled_grouped_gemm_matches_per_gemm( else: out = [torch.empty((m, n), dtype=out_dtype, device=device) for m in m_splits] - grouped_gemm_args = ( - w_nvfp4, - x_nvfp4, - out, - ) + grouped_gemm_args = (w_nvfp4, x_nvfp4, out) grouped_gemm_kwargs = { - "quantization_params": [None] * num_gemms, + "quantization_params": [None] * len(m_splits), "out_dtype": out_dtype, "layout": "TN", "m_splits": m_splits, @@ -357,18 +348,13 @@ def traced_wrapper(*args, **kwargs): assert len(calls) == 1 call = calls[0] - expected_alpha = torch.cat( - [ - tensor._amax_rowwise.view(-1) / (float(tensor._nvfp4_e4m3_max) * 6.0) - for tensor in w_nvfp4 - ] - ).to(dtype=torch.float32) - expected_row_scale = torch.cat( - [ - tensor._amax_rowwise.view(-1) / (float(tensor._nvfp4_e4m3_max) * 6.0) - for tensor in x_nvfp4 - ] - ).to(dtype=torch.float32) + 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( @@ -387,7 +373,7 @@ def traced_wrapper(*args, **kwargs): else: grouped_slices = grouped_out for grouped, ref in zip(grouped_slices, expected): - torch.testing.assert_close(grouped, ref, atol=0.5, rtol=0.25) + torch.testing.assert_close(grouped, ref, atol=0.125, rtol=0.25) def check_nvfp4_row_scaled_gemm_matches_emulated( @@ -612,8 +598,6 @@ def test_nvfp4_row_scaled_grouped_gemm_matches_per_gemm( 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( - x_dtype=torch.bfloat16, - w_dtype=torch.bfloat16, out_dtype=torch.bfloat16, m_splits=[256, 256, 256, 256], k=512, @@ -643,8 +627,6 @@ def test_nvfp4_row_scaled_grouped_gemm_rejects_unsupported( expected_error: str, ): check_nvfp4_row_scaled_grouped_gemm_matches_per_gemm( - x_dtype=torch.bfloat16, - w_dtype=torch.bfloat16, out_dtype=out_dtype, m_splits=m_splits, k=k, diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index 663c527a13..778a7e9ab6 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -209,7 +209,7 @@ def _cudnn_row_scaled_nvfp4_grouped_gemm( ) 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(getattr(tensor, "_nvfp4_use_4over6", False) for tensor in weights + 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 = ( @@ -257,7 +257,7 @@ def _cudnn_row_scaled_nvfp4_grouped_gemm( grouped_gemm_quant_wrapper_sm100, ) - device = inputs[0]._rowwise_data.device + device = inputs[0].device total_m = sum(m_splits_list) a_data = torch.cat( @@ -266,23 +266,14 @@ def _cudnn_row_scaled_nvfp4_grouped_gemm( ) a_tensor = a_data.view(torch.float4_e2m1fn_x2).unsqueeze(0).permute(1, 2, 0) - sfa_compact = torch.cat([tensor._rowwise_scale_inv for tensor in inputs], dim=0) - sfa_logical = sfa_compact.view(dtype=torch.float8_e4m3fn).view( - 1, - total_m // 128, - 4, - 32, - k // (4 * NVFP4_BLOCK_SCALING_SIZE), - 4, + 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) ) - sfa_logical = sfa_logical.permute(3, 2, 1, 5, 4, 0) - sfa_storage = torch.empty( - (1, total_m // 128, k // (4 * NVFP4_BLOCK_SCALING_SIZE), 32, 4, 4), - dtype=torch.float8_e4m3fn, - device=device, - ) - sfa_tensor = sfa_storage.permute(3, 4, 1, 5, 2, 0) - sfa_tensor.copy_(sfa_logical) b_ptrs, sfb_ptrs, _sfb_buffer = ( tex.grouped_mlp_experimental.swizzle_scales_and_pack_ptrs_for_discrete_weights( @@ -299,25 +290,16 @@ def _cudnn_row_scaled_nvfp4_grouped_gemm( 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.") - alpha_tensor = torch.cat( - [ - amax.view(-1) / (float(tensor._nvfp4_e4m3_max) * 6.0) - for tensor, amax in zip(weights, weight_amaxes) - ] - ).to(dtype=torch.float32) - row_scale_tensor = torch.cat( - [ - amax.view(-1) / (float(tensor._nvfp4_e4m3_max) * 6.0) - for tensor, amax in zip(inputs, input_amaxes) - ] - ).to(dtype=torch.float32) + 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 = torch.ones(total_m, 1, 1, dtype=torch.float32, device=device) + prob_tensor = _get_fp32_ones_tensor(total_m, device).view(total_m, 1, 1) result = grouped_gemm_quant_wrapper_sm100( a_tensor=a_tensor, @@ -571,25 +553,6 @@ def general_grouped_gemm( empty_tensor = _empty_tensor() empty_tensors = [empty_tensor] * num_gemms - # 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 - - sm_count = get_sm_count() - workspaces = get_cublas_workspace(A[0].device.index, False, True) - - if grad and use_bias: - grad_bias = [ - torch.empty(B[i].size(1), dtype=out[0].dtype, device="cuda") for i in range(num_gemms) - ] - else: - grad_bias = empty_tensors - bias = bias if use_bias else empty_tensors - if use_bias: - bias_dtype = TE_DType[grad_bias[0].dtype] if grad else TE_DType[bias[0].dtype] - 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): @@ -619,10 +582,29 @@ def general_grouped_gemm( bias=bias if use_bias else None, single_output=single_output, ), - grad_bias, - gelu_input, + empty_tensors, + empty_tensors, ) + # 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 + + sm_count = get_sm_count() + workspaces = get_cublas_workspace(A[0].device.index, False, True) + + if grad and use_bias: + grad_bias = [ + torch.empty(B[i].size(1), dtype=out[0].dtype, device="cuda") for i in range(num_gemms) + ] + else: + grad_bias = empty_tensors + bias = bias if use_bias else empty_tensors + if use_bias: + bias_dtype = TE_DType[grad_bias[0].dtype] if grad else TE_DType[bias[0].dtype] + else: + bias_dtype = TE_DType[torch.bfloat16] + if isinstance(quantization_params[0], DebugQuantizer): assert not gelu, "GELU not supported in debug mode" if single_output: From 6bb1abf472a1a01043d1b2e55c537cf64cd42a3e Mon Sep 17 00:00:00 2001 From: Ziang Li Date: Thu, 30 Jul 2026 21:17:03 -0700 Subject: [PATCH 5/7] Support small-K and FP32 cuDNN grouped GEMM Signed-off-by: Ziang Li --- tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py | 22 ++++++++++++++----- .../pytorch/cpp_extensions/gemm.py | 8 +++---- 2 files changed, 20 insertions(+), 10 deletions(-) diff --git a/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py b/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py index 605c6f84f3..f274eb7983 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py +++ b/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py @@ -588,9 +588,21 @@ def test_nvfp4_gemm_versus_reference( @pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +@pytest.mark.parametrize( + "m_splits, k, n, out_dtype", + [ + pytest.param([256, 256, 256, 256], 512, 512, torch.bfloat16, id="default_bf16"), + pytest.param([256, 256, 256, 256], 64, 256, torch.bfloat16, id="small_k_bf16"), + pytest.param([256, 256, 256, 256], 64, 256, torch.float32, id="small_k_fp32"), + ], +) @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( + m_splits: list[int], + k: int, + n: int, + out_dtype: torch.dtype, use_bias: bool, single_output: bool, monkeypatch, @@ -598,10 +610,10 @@ def test_nvfp4_row_scaled_grouped_gemm_matches_per_gemm( 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, + out_dtype=out_dtype, + m_splits=m_splits, + k=k, + n=n, use_bias=use_bias, single_output=single_output, monkeypatch=monkeypatch, @@ -613,8 +625,6 @@ def test_nvfp4_row_scaled_grouped_gemm_matches_per_gemm( "m_splits, k, n, out_dtype, use_4over6, expected_error", [ ([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"), ], ) diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index 778a7e9ab6..a72eb9ca22 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -224,9 +224,9 @@ def _cudnn_row_scaled_nvfp4_grouped_gemm( k = int(inputs[0].size(1)) n = int(weights[0].size(0)) - if k % 128 != 0 or n % 128 != 0: + if k % 64 != 0 or n % 128 != 0: raise NotImplementedError( - "cuDNN row-scaled NVFP4 grouped GEMM requires K and N multiples of 128." + "cuDNN row-scaled NVFP4 grouped GEMM requires K multiples of 64 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.") @@ -241,11 +241,11 @@ def _cudnn_row_scaled_nvfp4_grouped_gemm( 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( + if outputs[0].dtype not in (torch.float32, 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." + "cuDNN row-scaled NVFP4 grouped GEMM supports uniform FP32/BF16/FP16 outputs only." ) if bias is not None and ( len(bias) != num_gemms From b6e69356de05dd2560a0b0fede2d12e3cd16d138 Mon Sep 17 00:00:00 2001 From: Ziang Li Date: Thu, 6 Aug 2026 17:29:31 -0700 Subject: [PATCH 6/7] Tighten cuDNN row-scaled grouped GEMM Signed-off-by: Ziang Li --- tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py | 190 +++++----- .../pytorch/cpp_extensions/gemm.py | 336 ++++++++++++++---- transformer_engine/pytorch/quantization.py | 2 - 3 files changed, 360 insertions(+), 168 deletions(-) diff --git a/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py b/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py index f274eb7983..dfc19d0761 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py +++ b/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py @@ -241,6 +241,8 @@ 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, @@ -249,15 +251,17 @@ 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", + weight_scales_swizzled: bool = False, monkeypatch=None, - expected_error=None, ): te_dtype = te.DType.kFloat4E2M1 device = "cuda" torch.manual_seed(23) torch.cuda.manual_seed(23) - dtype = torch.bfloat16 + num_gemms = len(m_splits) + x_quantizer = NVFP4Quantizer( fp4_dtype=te_dtype, rowwise=True, @@ -268,6 +272,7 @@ 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, @@ -278,57 +283,43 @@ 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, ) + w_quantizer.optimize_for_gemm = weight_scales_swizzled x_nvfp4 = [] w_nvfp4 = [] bias = [] expected = [] for m in m_splits: - x = torch.randn((m, k), dtype=dtype, device=device) - w = torch.randn((n, k), dtype=dtype, device=device) + x = torch.randn((m, k), dtype=x_dtype, device=device) + w = torch.randn((n, k), dtype=w_dtype, device=device) x_nvfp4.append( x_quantizer.update_quantized( - x, x_quantizer.make_empty(x.shape, dtype=dtype, device=device) + x, x_quantizer.make_empty(x.shape, dtype=x_dtype, device=device) ) ) w_nvfp4.append( w_quantizer.update_quantized( - w, w_quantizer.make_empty(w.shape, dtype=dtype, device=device) + w, w_quantizer.make_empty(w.shape, dtype=w_dtype, device=device) ) ) - 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] - ) + 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] + ) 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_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: @@ -344,36 +335,61 @@ def traced_wrapper(*args, **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, + 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, ) - 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 torch.cuda.get_device_capability() == (10, 7): + assert not calls + else: + assert len(calls) == 1 + call = calls[0] + padded_m_splits = [((m + 255) // 256) * 256 if m > 0 else 0 for m in m_splits] + expected_alpha = torch.cat( + [tensor._amax_rowwise.view(-1) / (6.0 * tensor._nvfp4_e4m3_max) for tensor in w_nvfp4] + ) + expected_row_scale = torch.zeros(sum(padded_m_splits), dtype=torch.float32, device=device) + offset = 0 + for tensor, m, padded_m in zip(x_nvfp4, m_splits, padded_m_splits): + expected_row_scale[offset : offset + m].copy_( + tensor._amax_rowwise.view(-1) / (6.0 * tensor._nvfp4_e4m3_max) + ) + offset += padded_m + 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(padded_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 call["m_aligned"] == 256 + assert call["prob_tensor"] is None + assert (call["bias_tensor"] is not None) == use_bias + direct_output = single_output and padded_m_splits == m_splits + assert (call["d_tensor"] is not None) == direct_output + if direct_output: + assert call["d_tensor"].data_ptr() == out[0].data_ptr() 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.125, rtol=0.25) + torch.testing.assert_close(grouped, ref, atol=0.0, rtol=0.0) def check_nvfp4_row_scaled_gemm_matches_emulated( @@ -589,62 +605,62 @@ def test_nvfp4_gemm_versus_reference( @pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) @pytest.mark.parametrize( - "m_splits, k, n, out_dtype", + "m_splits, k, n", [ - pytest.param([256, 256, 256, 256], 512, 512, torch.bfloat16, id="default_bf16"), - pytest.param([256, 256, 256, 256], 64, 256, torch.bfloat16, id="small_k_bf16"), - pytest.param([256, 256, 256, 256], 64, 256, torch.float32, id="small_k_fp32"), + ([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), ], ) +@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"]) +@pytest.mark.parametrize( + "weight_scales_swizzled", + [False, True], + ids=["compact_weight_scales", "swizzled_weight_scales"], +) def test_nvfp4_row_scaled_grouped_gemm_matches_per_gemm( 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, + weight_scales_swizzled: 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( + 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, - monkeypatch=monkeypatch, - ) - - -@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) -@pytest.mark.parametrize( - "m_splits, k, n, out_dtype, use_4over6, expected_error", - [ - ([128, 256], 128, 128, torch.bfloat16, False, "M multiples of 256"), - ([256, 256], 128, 128, torch.bfloat16, True, "does not support 4over6"), - ], -) -def test_nvfp4_row_scaled_grouped_gemm_rejects_unsupported( - m_splits: list[int], - k: int, - n: int, - out_dtype: torch.dtype, - use_4over6: bool, - expected_error: str, -): - check_nvfp4_row_scaled_grouped_gemm_matches_per_gemm( - out_dtype=out_dtype, - m_splits=m_splits, - k=k, - n=n, - use_bias=False, - single_output=True, use_4over6=use_4over6, - expected_error=expected_error, + nvfp4_4over6_err_mode=nvfp4_4over6_err_mode, + weight_scales_swizzled=weight_scales_swizzled, + monkeypatch=monkeypatch, ) diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index a72eb9ca22..2f3650a111 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -10,7 +10,7 @@ import torch import transformer_engine_torch as tex from ..constants import NVFP4_BLOCK_SCALING_SIZE, TE_DType, DType -from ..utils import get_sm_count, _empty_tensor +from ..utils import ceil_div, get_sm_count, round_up_to_nearest_multiple, _empty_tensor from ..quantized_tensor import QuantizedTensorStorage, Quantizer from ..tensor.float8_blockwise_tensor import Float8BlockQuantizer @@ -198,6 +198,7 @@ def _cudnn_row_scaled_nvfp4_grouped_gemm( single_output: bool, ) -> Union[torch.Tensor, List[torch.Tensor]]: """Run tensor-scaled weights and row-scaled inputs with the cuDNN MoE kernel.""" + m_alignment = 256 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.") @@ -209,25 +210,21 @@ def _cudnn_row_scaled_nvfp4_grouped_gemm( ) 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] + [int(m) for m in 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." - ) + if any(m < 0 for m in m_splits_list): + raise ValueError("Grouped GEMM split sizes must be non-negative.") + if any(len(tensor.size()) != 2 for tensor in weights + inputs): + raise ValueError("cuDNN row-scaled NVFP4 grouped GEMM requires 2-D operands.") k = int(inputs[0].size(1)) n = int(weights[0].size(0)) - if k % 64 != 0 or n % 128 != 0: - raise NotImplementedError( - "cuDNN row-scaled NVFP4 grouped GEMM requires K multiples of 64 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( @@ -235,71 +232,168 @@ def _cudnn_row_scaled_nvfp4_grouped_gemm( for tensor, m in zip(inputs, m_splits_list) ): raise ValueError("Grouped GEMM input shapes must match m_splits and K.") + if k % (2 * NVFP4_BLOCK_SCALING_SIZE) != 0 or n % NVFP4_BLOCK_SCALING_SIZE != 0: + raise ValueError("NVFP4 grouped GEMM requires K multiples of 32 and N multiples of 16.") + + rowwise_tensors = weights + inputs + if any( + tensor._rowwise_data is None + or tensor._rowwise_scale_inv is None + or tensor._amax_rowwise is None + for tensor in rowwise_tensors + ): + raise ValueError("cuDNN row-scaled NVFP4 grouped GEMM requires rowwise NVFP4 data.") + if any(tensor._with_gemm_swizzled_scales for tensor in inputs): + raise ValueError("cuDNN row-scaled NVFP4 grouped GEMM requires compact input scales.") + + device = inputs[0].device + if device.type != "cuda" or any(tensor.device != device for tensor in rowwise_tensors): + raise ValueError("Grouped GEMM operands must be CUDA tensors on the same device.") + + expected_scale_cols = 4 * ceil_div(k, 4 * NVFP4_BLOCK_SCALING_SIZE) + for tensor, rows in zip(weights + inputs, [n] * num_gemms + m_splits_list): + expected_data_shape = (rows, k // 2) + expected_scale_shape = ( + round_up_to_nearest_multiple(rows, 128), + expected_scale_cols, + ) + if tuple(tensor._rowwise_data.shape) != expected_data_shape: + raise ValueError( + f"Expected rowwise NVFP4 data shape {expected_data_shape}, " + f"got {tuple(tensor._rowwise_data.shape)}." + ) + if tuple(tensor._rowwise_scale_inv.shape) != expected_scale_shape: + raise ValueError( + f"Expected rowwise NVFP4 scale shape {expected_scale_shape}, " + f"got {tuple(tensor._rowwise_scale_inv.shape)}." + ) + + weight_amaxes = [tensor._amax_rowwise for tensor in weights] + input_amaxes = [tensor._amax_rowwise for tensor in inputs] + if any(amax.dtype != torch.float32 or amax.numel() != 1 for amax in weight_amaxes): + raise ValueError("Row-scaled NVFP4 grouped GEMM requires FP32 tensor-scaled weights.") + if any( + amax.dtype != torch.float32 or amax.numel() != m + for amax, m in zip(input_amaxes, m_splits_list) + ): + raise ValueError("Row-scaled NVFP4 grouped GEMM requires one FP32 input scale per row.") + 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 + tuple(output.shape) != (m, n) or not output.is_contiguous() for output, m in zip(outputs, expected_output_rows) ): - raise ValueError("Grouped GEMM output shapes do not match m_splits and N.") + raise ValueError("Grouped GEMM outputs must be contiguous with shape (M, N).") if outputs[0].dtype not in (torch.float32, 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 FP32/BF16/FP16 outputs only." ) + if any(output.device != device for output in outputs): + raise ValueError("Grouped GEMM outputs must be on the operand device.") 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,).") + if bias is not None and ( + any(tensor.device != device for tensor in bias) + or any(tensor.dtype != bias[0].dtype for tensor in bias) + or bias[0].dtype not in (torch.bfloat16, torch.float16) + ): + raise ValueError("Grouped GEMM biases must have a uniform BF16/FP16 dtype and device.") 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, + padded_m_splits = [ + round_up_to_nearest_multiple(m, m_alignment) if m > 0 else 0 for m in m_splits_list + ] + total_padded_m = sum(padded_m_splits) + a_data = torch.zeros( + (total_padded_m, k // 2), + dtype=inputs[0]._rowwise_data.dtype, + device=device, + ) + a_scales = torch.zeros( + (total_padded_m, expected_scale_cols), + dtype=inputs[0]._rowwise_scale_inv.dtype, + device=device, ) + row_scale_tensor = torch.zeros(total_padded_m, dtype=torch.float32, device=device) + padded_offset = 0 + for tensor, amax, m, padded_m in zip( + inputs, + input_amaxes, + m_splits_list, + padded_m_splits, + ): + a_data[padded_offset : padded_offset + m].copy_(tensor._rowwise_data) + a_scales[padded_offset : padded_offset + m].copy_(tensor._rowwise_scale_inv[:m]) + row_scale_tensor[padded_offset : padded_offset + m].copy_( + amax.view(-1) / (6.0 * tensor._nvfp4_e4m3_max) + ) + padded_offset += padded_m + 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) + a_scales.view(dtype=torch.float8_e4m3fn) + .view( + 1, + total_padded_m // 128, + 4, + 32, + ceil_div(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", + weight_scales_swizzled = {bool(tensor._with_gemm_swizzled_scales) for tensor in weights} + if len(weight_scales_swizzled) != 1: + raise ValueError("Grouped GEMM weights must use a uniform scale layout.") + if weight_scales_swizzled.pop(): + packed_ptrs = tex.copy_data_ptrs_to_device( + [tensor._rowwise_data for tensor in weights] + + [tensor._rowwise_scale_inv for tensor in weights], device, ) - ) + b_ptrs = packed_ptrs[:num_gemms] + sfb_ptrs = packed_ptrs[num_gemms:] + else: + 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 + alpha_tensor = torch.cat( + [ + amax.view(-1) / (6.0 * tensor._nvfp4_e4m3_max) + for tensor, amax in zip(weights, weight_amaxes) + ] + ) 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( + padded_offsets = torch.tensor(padded_m_splits, dtype=torch.int32, device=device).cumsum( 0, dtype=torch.int32 ) - prob_tensor = _get_fp32_ones_tensor(total_m, device).view(total_m, 1, 1) + + cudnn_output = None + if single_output and padded_m_splits == m_splits_list: + cudnn_output = outputs[0].as_strided( + (total_padded_m, n, 1), + (n, 1, total_padded_m * n), + ) result = grouped_gemm_quant_wrapper_sm100( a_tensor=a_tensor, @@ -310,27 +404,35 @@ def _cudnn_row_scaled_nvfp4_grouped_gemm( alpha_tensor=alpha_tensor, bias_tensor=bias_tensor, norm_const_tensor=None, - prob_tensor=prob_tensor, + prob_tensor=None, row_scale_tensor=row_scale_tensor, acc_dtype=torch.float32, d_dtype=outputs[0].dtype, + d_tensor=cudnn_output, cd_major="n", sf_vec_size=NVFP4_BLOCK_SCALING_SIZE, + m_aligned=m_alignment, discrete_col_sfd=False, b_dtype=torch.float4_e2m1fn_x2, b_major="k", n=n, - current_stream=torch.cuda.current_stream().cuda_stream, + current_stream=torch.cuda.current_stream(device=device).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) + if cudnn_output is not None: return outputs[0] - for output, output_data in zip(outputs, d_tensor.split(m_splits_list)): - output.view(-1, n).copy_(output_data) + padded_offset = 0 + output_offset = 0 + for i, (m, padded_m) in enumerate(zip(m_splits_list, padded_m_splits)): + output = outputs[0][output_offset : output_offset + m] if single_output else outputs[i] + output.copy_(d_tensor[padded_offset : padded_offset + m]) + padded_offset += padded_m + output_offset += m + if single_output: + return outputs[0] return outputs @@ -463,6 +565,42 @@ def general_gemm( assert isinstance( B, NVFP4TensorStorage ), "Row-scaled NVFP4 GEMM currently requires NVFP4 B." + + cudnn_output_dtype = out.dtype if out is not None else out_dtype + if ( + layout == "TN" + and not _is_nvfp4_row_scaled_tensor(A) + and _is_nvfp4_row_scaled_tensor(B) + and len(A.size()) == 2 + and len(B.size()) == 2 + and B.size(1) % (2 * NVFP4_BLOCK_SCALING_SIZE) == 0 + and cudnn_output_dtype in (torch.float32, torch.bfloat16, torch.float16) + and (bias is None or bias.dtype in (torch.bfloat16, torch.float16)) + and alpha == 1.0 + and not accumulate + and not use_split_accumulator + and not grad + and torch.cuda.get_device_capability(A.device) != (10, 7) + ): + if out is None: + out = torch.empty( + (B.size(0), A.size(0)), + dtype=cudnn_output_dtype, + device=B.device, + ) + out = _cudnn_row_scaled_nvfp4_grouped_gemm( + [A], + [B], + [out], + m_splits=[B.size(0)], + bias=None if bias is None else [bias], + single_output=True, + ) + if debug_quantizer is not None: + out = debug_quantizer.process_gemm_output(out) + empty_tensor = _empty_tensor() + return out, empty_tensor, empty_tensor, None + # Reuse the per-tensor GEMM and apply selected row/column global scales # to the FP32 output. This extends #2931 without a dedicated GEMM kernel. gemm_A, gemm_B, output_row_scales, output_col_scales = _nvfp4_row_scaled_gemm_inputs( @@ -553,39 +691,6 @@ def general_grouped_gemm( 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, - ) - # 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 @@ -605,6 +710,79 @@ 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): + if torch.cuda.get_device_capability(A[0].device) != (10, 7): + if len(B) != num_gemms or len(quantization_params) != num_gemms: + raise ValueError( + "Grouped GEMM operands and output quantizers must have matching lengths." + ) + 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, + ), + grad_bias, + gelu_input, + ) + + 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 8263324083..07a3b80483 100644 --- a/transformer_engine/pytorch/quantization.py +++ b/transformer_engine/pytorch/quantization.py @@ -274,8 +274,6 @@ 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 From 7e5c95f3760b0908d5941cbe913529667431b2bb Mon Sep 17 00:00:00 2001 From: Ziang Li Date: Thu, 6 Aug 2026 18:00:02 -0700 Subject: [PATCH 7/7] Refactor cuDNN row-scaled grouped GEMM Signed-off-by: Ziang Li --- tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py | 42 +++--- .../pytorch/cpp_extensions/gemm.py | 141 ++++++++++++------ 2 files changed, 116 insertions(+), 67 deletions(-) diff --git a/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py b/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py index dfc19d0761..439c6daf61 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py +++ b/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py @@ -320,21 +320,29 @@ def check_nvfp4_row_scaled_grouped_gemm_matches_per_gemm( else: out = [torch.empty((m, n), dtype=out_dtype, device=device) for m in m_splits] - 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) + compute_capability = torch.cuda.get_device_capability() + use_cudnn = compute_capability[0] == 10 and compute_capability != (10, 7) + if use_cudnn: + 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") + + 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) + else: + monkeypatch.setattr( + "transformer_engine.pytorch.cpp_extensions.gemm._cudnn_row_scaled_nvfp4_grouped_gemm", + lambda *args, **kwargs: pytest.fail("cuDNN row-scaled GEMM dispatched on fallback"), + ) grouped_out, _, _ = general_grouped_gemm( w_nvfp4, x_nvfp4, @@ -348,9 +356,7 @@ def traced_wrapper(*args, **kwargs): single_output=single_output, ) - if torch.cuda.get_device_capability() == (10, 7): - assert not calls - else: + if use_cudnn: assert len(calls) == 1 call = calls[0] padded_m_splits = [((m + 255) // 256) * 256 if m > 0 else 0 for m in m_splits] @@ -647,7 +653,7 @@ def test_nvfp4_row_scaled_grouped_gemm_matches_per_gemm( monkeypatch, ): if torch.cuda.get_device_capability() < (10, 0): - pytest.skip("Requires SM100+ for cuDNN grouped GEMM quant kernel.") + pytest.skip("Requires SM100+ for NVFP4 grouped GEMM.") check_nvfp4_row_scaled_grouped_gemm_matches_per_gemm( x_dtype=x_dtype, w_dtype=w_dtype, diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index 2f3650a111..ad87f86c41 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -188,7 +188,16 @@ def _validate_native_gemm_output_quantizer(quantization_params): ) -def _cudnn_row_scaled_nvfp4_grouped_gemm( +_CUDNN_ROW_SCALED_NVFP4_M_ALIGNMENT = 256 + + +def _use_cudnn_row_scaled_nvfp4_gemm(device: torch.device) -> bool: + """Whether cuDNN supports row-scaled NVFP4 GEMM on this architecture.""" + compute_capability = torch.cuda.get_device_capability(device) + return compute_capability[0] == 10 and compute_capability != (10, 7) + + +def _validate_cudnn_row_scaled_nvfp4_grouped_gemm( weights: List[NVFP4TensorStorage], inputs: List[NVFP4TensorStorage], outputs: List[torch.Tensor], @@ -196,9 +205,8 @@ def _cudnn_row_scaled_nvfp4_grouped_gemm( 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.""" - m_alignment = 256 +) -> Tuple[List[int], int, int, torch.device]: + """Validate operands and return normalized split sizes, N, K, and device.""" 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.") @@ -304,41 +312,43 @@ def _cudnn_row_scaled_nvfp4_grouped_gemm( ): raise ValueError("Grouped GEMM biases must have a uniform BF16/FP16 dtype and device.") - from cudnn import ( # pylint: disable=import-outside-toplevel,no-name-in-module - grouped_gemm_quant_wrapper_sm100, - ) + return m_splits_list, n, k, device + +def _pack_cudnn_row_scaled_nvfp4_inputs( + inputs: List[NVFP4TensorStorage], + m_splits: List[int], + k: int, + device: torch.device, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, List[int], torch.Tensor]: + """Pack discrete row-scaled inputs into the padded cuDNN grouped-GEMM layout.""" padded_m_splits = [ - round_up_to_nearest_multiple(m, m_alignment) if m > 0 else 0 for m in m_splits_list + round_up_to_nearest_multiple(m, _CUDNN_ROW_SCALED_NVFP4_M_ALIGNMENT) if m > 0 else 0 + for m in m_splits ] total_padded_m = sum(padded_m_splits) + scale_cols = 4 * ceil_div(k, 4 * NVFP4_BLOCK_SCALING_SIZE) a_data = torch.zeros( (total_padded_m, k // 2), dtype=inputs[0]._rowwise_data.dtype, device=device, ) a_scales = torch.zeros( - (total_padded_m, expected_scale_cols), + (total_padded_m, scale_cols), dtype=inputs[0]._rowwise_scale_inv.dtype, device=device, ) - row_scale_tensor = torch.zeros(total_padded_m, dtype=torch.float32, device=device) - padded_offset = 0 - for tensor, amax, m, padded_m in zip( - inputs, - input_amaxes, - m_splits_list, - padded_m_splits, - ): - a_data[padded_offset : padded_offset + m].copy_(tensor._rowwise_data) - a_scales[padded_offset : padded_offset + m].copy_(tensor._rowwise_scale_inv[:m]) - row_scale_tensor[padded_offset : padded_offset + m].copy_( - amax.view(-1) / (6.0 * tensor._nvfp4_e4m3_max) + row_scales = torch.zeros(total_padded_m, dtype=torch.float32, device=device) + offset = 0 + for tensor, m, padded_m in zip(inputs, m_splits, padded_m_splits): + a_data[offset : offset + m].copy_(tensor._rowwise_data) + a_scales[offset : offset + m].copy_(tensor._rowwise_scale_inv[:m]) + row_scales[offset : offset + m].copy_( + tensor._amax_rowwise.view(-1) / (6.0 * tensor._nvfp4_e4m3_max) ) - padded_offset += padded_m + offset += padded_m a_tensor = a_data.view(torch.float4_e2m1fn_x2).unsqueeze(0).permute(1, 2, 0) - sfa_tensor = ( a_scales.view(dtype=torch.float8_e4m3fn) .view( @@ -353,11 +363,24 @@ def _cudnn_row_scaled_nvfp4_grouped_gemm( .contiguous() .permute(3, 4, 1, 5, 2, 0) ) + padded_offsets = torch.tensor(padded_m_splits, dtype=torch.int32, device=device).cumsum( + 0, dtype=torch.int32 + ) + return a_tensor, sfa_tensor, row_scales, padded_m_splits, padded_offsets - weight_scales_swizzled = {bool(tensor._with_gemm_swizzled_scales) for tensor in weights} - if len(weight_scales_swizzled) != 1: + +def _pack_cudnn_nvfp4_weights( + weights: List[NVFP4TensorStorage], + device: torch.device, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: + """Pack tensor-scaled weights and return their global scales.""" + num_gemms = len(weights) + scale_layouts = {bool(tensor._with_gemm_swizzled_scales) for tensor in weights} + if len(scale_layouts) != 1: raise ValueError("Grouped GEMM weights must use a uniform scale layout.") - if weight_scales_swizzled.pop(): + + scale_buffer = None + if scale_layouts.pop(): packed_ptrs = tex.copy_data_ptrs_to_device( [tensor._rowwise_data for tensor in weights] + [tensor._rowwise_scale_inv for tensor in weights], @@ -366,7 +389,7 @@ def _cudnn_row_scaled_nvfp4_grouped_gemm( b_ptrs = packed_ptrs[:num_gemms] sfb_ptrs = packed_ptrs[num_gemms:] else: - b_ptrs, sfb_ptrs, _sfb_buffer = ( + b_ptrs, sfb_ptrs, scale_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], @@ -376,20 +399,46 @@ def _cudnn_row_scaled_nvfp4_grouped_gemm( ) alpha_tensor = torch.cat( - [ - amax.view(-1) / (6.0 * tensor._nvfp4_e4m3_max) - for tensor, amax in zip(weights, weight_amaxes) - ] + [tensor._amax_rowwise.view(-1) / (6.0 * tensor._nvfp4_e4m3_max) for tensor in weights] ) + return b_ptrs, sfb_ptrs, alpha_tensor, scale_buffer - bias_tensor = None if bias is None else torch.stack(bias, dim=0).transpose(0, 1) - padded_offsets = torch.tensor(padded_m_splits, dtype=torch.int32, device=device).cumsum( - 0, dtype=torch.int32 +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.""" + m_splits, n, k, device = _validate_cudnn_row_scaled_nvfp4_grouped_gemm( + weights, + inputs, + outputs, + m_splits=m_splits, + bias=bias, + single_output=single_output, + ) + + from cudnn import ( # pylint: disable=import-outside-toplevel,no-name-in-module + grouped_gemm_quant_wrapper_sm100, + ) + + a_tensor, sfa_tensor, row_scale_tensor, padded_m_splits, padded_offsets = ( + _pack_cudnn_row_scaled_nvfp4_inputs(inputs, m_splits, k, device) + ) + b_ptrs, sfb_ptrs, alpha_tensor, _weight_scale_buffer = _pack_cudnn_nvfp4_weights( + weights, device ) + total_padded_m = sum(padded_m_splits) + bias_tensor = None if bias is None else torch.stack(bias, dim=0).transpose(0, 1) + cudnn_output = None - if single_output and padded_m_splits == m_splits_list: + if single_output and padded_m_splits == m_splits: cudnn_output = outputs[0].as_strided( (total_padded_m, n, 1), (n, 1, total_padded_m * n), @@ -411,7 +460,7 @@ def _cudnn_row_scaled_nvfp4_grouped_gemm( d_tensor=cudnn_output, cd_major="n", sf_vec_size=NVFP4_BLOCK_SCALING_SIZE, - m_aligned=m_alignment, + m_aligned=_CUDNN_ROW_SCALED_NVFP4_M_ALIGNMENT, discrete_col_sfd=False, b_dtype=torch.float4_e2m1fn_x2, b_major="k", @@ -424,16 +473,10 @@ def _cudnn_row_scaled_nvfp4_grouped_gemm( if cudnn_output is not None: return outputs[0] - padded_offset = 0 - output_offset = 0 - for i, (m, padded_m) in enumerate(zip(m_splits_list, padded_m_splits)): - output = outputs[0][output_offset : output_offset + m] if single_output else outputs[i] - output.copy_(d_tensor[padded_offset : padded_offset + m]) - padded_offset += padded_m - output_offset += m - if single_output: - return outputs[0] - return outputs + output_views = torch.split(outputs[0], m_splits) if single_output else outputs + for output, padded_output in zip(output_views, torch.split(d_tensor, padded_m_splits)): + output.copy_(padded_output[: output.size(0)]) + return outputs[0] if single_output else outputs def general_gemm( @@ -567,7 +610,7 @@ def general_gemm( ), "Row-scaled NVFP4 GEMM currently requires NVFP4 B." cudnn_output_dtype = out.dtype if out is not None else out_dtype - if ( + if ( # pylint: disable=too-many-boolean-expressions layout == "TN" and not _is_nvfp4_row_scaled_tensor(A) and _is_nvfp4_row_scaled_tensor(B) @@ -580,7 +623,7 @@ def general_gemm( and not accumulate and not use_split_accumulator and not grad - and torch.cuda.get_device_capability(A.device) != (10, 7) + and _use_cudnn_row_scaled_nvfp4_gemm(A.device) ): if out is None: out = torch.empty( @@ -713,7 +756,7 @@ 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): - if torch.cuda.get_device_capability(A[0].device) != (10, 7): + if _use_cudnn_row_scaled_nvfp4_gemm(A[0].device): if len(B) != num_gemms or len(quantization_params) != num_gemms: raise ValueError( "Grouped GEMM operands and output quantizers must have matching lengths."