diff --git a/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py b/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py index eb480060e2..439c6daf61 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", + weight_scales_swizzled: bool = False, + monkeypatch=None, ): te_dtype = te.DType.kFloat4E2M1 device = "cuda" @@ -283,6 +285,7 @@ def check_nvfp4_row_scaled_grouped_gemm_matches_per_gemm( 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 = [] @@ -317,6 +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] + calls = [] + 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, @@ -330,6 +356,40 @@ def check_nvfp4_row_scaled_grouped_gemm_matches_per_gemm( single_output=single_output, ) + 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] + 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: @@ -573,6 +633,11 @@ def test_nvfp4_gemm_versus_reference( @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, @@ -584,7 +649,11 @@ def test_nvfp4_row_scaled_grouped_gemm_matches_per_gemm( 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 NVFP4 grouped GEMM.") check_nvfp4_row_scaled_grouped_gemm_matches_per_gemm( x_dtype=x_dtype, w_dtype=w_dtype, @@ -596,6 +665,8 @@ def test_nvfp4_row_scaled_grouped_gemm_matches_per_gemm( single_output=single_output, use_4over6=use_4over6, 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 f3d97b7269..ad87f86c41 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -9,8 +9,8 @@ import functools import torch import transformer_engine_torch as tex -from ..constants import TE_DType, DType -from ..utils import get_sm_count, _empty_tensor +from ..constants import NVFP4_BLOCK_SCALING_SIZE, TE_DType, DType +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 @@ -188,6 +188,297 @@ def _validate_native_gemm_output_quantizer(quantization_params): ) +_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], + *, + m_splits: Optional[List[int]], + bias: Optional[List[torch.Tensor]], + single_output: bool, +) -> 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.") + if not all(isinstance(tensor, NVFP4TensorStorage) for tensor in weights + inputs): + raise TypeError("cuDNN row-scaled NVFP4 grouped GEMM requires NVFP4 inputs.") + if any(_is_nvfp4_row_scaled_tensor(tensor) for tensor in weights): + raise NotImplementedError( + "cuDNN row-scaled NVFP4 grouped GEMM requires tensor-scaled weights." + ) + if not all(_is_nvfp4_row_scaled_tensor(tensor) for tensor in inputs): + raise NotImplementedError("cuDNN row-scaled NVFP4 grouped GEMM requires row-scaled inputs.") + + m_splits_list = ( + [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 < 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 any(tuple(tensor.size()) != (n, k) for tensor in weights): + raise ValueError("All grouped GEMM weights must have the same (N, K) shape.") + if any( + int(tensor.size(0)) != m or int(tensor.size(1)) != k + for tensor, m in zip(inputs, m_splits_list) + ): + raise ValueError("Grouped GEMM input shapes must match m_splits and K.") + 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( + tuple(output.shape) != (m, n) or not output.is_contiguous() + for output, m in zip(outputs, expected_output_rows) + ): + 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.") + + 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, _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, scale_cols), + dtype=inputs[0]._rowwise_scale_inv.dtype, + device=device, + ) + 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) + ) + 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( + 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) + ) + 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 + + +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.") + + 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], + device, + ) + b_ptrs = packed_ptrs[:num_gemms] + sfb_ptrs = packed_ptrs[num_gemms:] + else: + 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], + "nvfp4", + device, + ) + ) + + alpha_tensor = torch.cat( + [tensor._amax_rowwise.view(-1) / (6.0 * tensor._nvfp4_e4m3_max) for tensor in weights] + ) + return b_ptrs, sfb_ptrs, alpha_tensor, scale_buffer + + +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: + 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, + 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=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=_CUDNN_ROW_SCALED_NVFP4_M_ALIGNMENT, + discrete_col_sfd=False, + b_dtype=torch.float4_e2m1fn_x2, + b_major="k", + n=n, + current_stream=torch.cuda.current_stream(device=device).cuda_stream, + use_dynamic_sched=True, + ) + d_tensor = result["d_tensor"].squeeze(-1) + + if cudnn_output is not None: + return outputs[0] + + 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( A: torch.Tensor, B: torch.Tensor, @@ -317,6 +608,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 ( # pylint: disable=too-many-boolean-expressions + 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 _use_cudnn_row_scaled_nvfp4_gemm(A.device) + ): + 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( @@ -429,6 +756,41 @@ 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 _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." + ) + 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 (