[PyTorch] Fuse NVFP4 global scales and GEMM alpha in the grouped MLP

tests/pytorch/nvfp4/test_nvfp4_global_scale.py

@@ -0,0 +1,107 @@
+# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+"""Tests for the fused NVFP4 global-scale / GEMM-alpha Triton kernel.
+
+The kernel replaces the eager ``amax.float() / (6 * scale_max)`` divisions and the
+product / square root of the two operand scales used by the grouped MLP. It must be
+bit-identical to the eager ops, including with a broadcast (single-group) amax, and must
+be safe to capture in a CUDA graph."""
+
+import pytest
+import torch
+
+from transformer_engine.pytorch.triton.nvfp4_global_scale import nvfp4_global_scale_pair
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required")
+
+_DENOMS = [6.0 * 448.0, 6.0 * 256.0, 6.0 * 114688.0]
+
+
+def _random_amax(n, device):
+    # Wide dynamic range plus exact zeros (empty experts) and a few special magnitudes.
+    amax = torch.exp2(torch.empty(n, device=device).uniform_(-40.0, 40.0))
+    amax *= torch.rand(n, device=device) + 0.5
+    if n >= 4:
+        amax[0] = 0.0
+        amax[1] = 1.0
+        amax[2] = torch.finfo(torch.float32).tiny
+        amax[3] = 448.0 * 6.0
+    return amax
+
+
+def _reference(amax_a, amax_b, denom_a, denom_b, combine):
+    scale_a = amax_a.to(torch.float32) / denom_a
+    scale_b = amax_b.to(torch.float32) / denom_b
+    if combine is None:
+        return scale_a, scale_b
+    if combine == "sqrt_product":
+        return torch.sqrt(scale_a * scale_b)
+    return (scale_a * scale_b).to(torch.float32)
+
+
+def _assert_bitwise_equal(out, ref):
+    if isinstance(ref, tuple):
+        assert isinstance(out, tuple) and len(out) == len(ref)
+        for o, r in zip(out, ref):
+            _assert_bitwise_equal(o, r)
+        return
+    assert out.dtype == torch.float32
+    assert out.shape == ref.shape
+    assert torch.equal(out.view(torch.int32), ref.contiguous().view(torch.int32))
+
+
+@pytest.mark.parametrize("n", [1, 3, 8, 32, 257, 4096])
+@pytest.mark.parametrize("combine", [None, "product", "sqrt_product"])
+@pytest.mark.parametrize("denoms", [(_DENOMS[0], _DENOMS[0]), (_DENOMS[1], _DENOMS[2])])
+@pytest.mark.parametrize("broadcast", ["none", "a", "b"])
+def test_matches_eager(n, combine, denoms, broadcast):
+    torch.manual_seed(n)
+    device = torch.device("cuda")
+    amax_a = _random_amax(1 if broadcast == "a" else n, device)
+    amax_b = _random_amax(1 if broadcast == "b" else n, device)
+    out = nvfp4_global_scale_pair(amax_a, amax_b, denoms[0], denoms[1], combine=combine)
+    if combine is None and amax_a.numel() != amax_b.numel():
+        # Separate scales keep their own shapes, so the caller falls back to eager.
+        assert out is None
+        return
+    assert out is not None
+    _assert_bitwise_equal(out, _reference(amax_a, amax_b, denoms[0], denoms[1], combine))
+
+
+def test_falls_back_on_unsupported_inputs():
+    device = torch.device("cuda")
+    amax = torch.ones(8, device=device)
+    assert nvfp4_global_scale_pair(amax.half(), amax, 1.0, 1.0) is None
+    assert nvfp4_global_scale_pair(amax, amax[:3], 1.0, 1.0) is None
+    assert nvfp4_global_scale_pair(amax.cpu(), amax.cpu(), 1.0, 1.0) is None
+    big = torch.ones(4097, device=device)
+    assert nvfp4_global_scale_pair(big, big, 1.0, 1.0) is None
+
+
+@pytest.mark.parametrize("combine", [None, "product", "sqrt_product"])
+def test_cuda_graph_replay(combine):
+    device = torch.device("cuda")
+    n = 32
+    denom = _DENOMS[0]
+    amax_a = _random_amax(n, device)
+    amax_b = _random_amax(n, device)
+
+    # Warm up (compiles the kernel) on a side stream before capture.
+    stream = torch.cuda.Stream()
+    stream.wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(stream):
+        nvfp4_global_scale_pair(amax_a, amax_b, denom, denom, combine=combine)
+    torch.cuda.current_stream().wait_stream(stream)
+
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        out = nvfp4_global_scale_pair(amax_a, amax_b, denom, denom, combine=combine)
+
+    for seed in range(3):
+        torch.manual_seed(seed)
+        amax_a.copy_(_random_amax(n, device))
+        amax_b.copy_(_random_amax(n, device))
+        graph.replay()
+        torch.cuda.synchronize()
+        _assert_bitwise_equal(out, _reference(amax_a, amax_b, denom, denom, combine))

transformer_engine/common/triton/nvfp4_global_scale.py

@@ -0,0 +1,53 @@
+# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+
+"""Triton kernel for the per-group NVFP4 global scales of a GEMM operand pair."""
+
+import triton
+import triton.language as tl
+
+# Output modes of _nvfp4_global_scale_pair_kernel.
+NVFP4_GLOBAL_SCALE_PAIR = 0  # out_a = amax_a / denom_a, out_b = amax_b / denom_b
+NVFP4_GLOBAL_SCALE_PRODUCT = 1  # out_a = (amax_a / denom_a) * (amax_b / denom_b)
+NVFP4_GLOBAL_SCALE_SQRT_PRODUCT = 2  # out_a = sqrt of the product above
+
+
+@triton.jit
+def _nvfp4_global_scale_pair_kernel(
+    amax_a_ptr,
+    amax_b_ptr,
+    out_a_ptr,
+    out_b_ptr,
+    n,
+    a_broadcast,
+    b_broadcast,
+    denom_a,
+    denom_b,
+    MODE: tl.constexpr,
+    BLOCK: tl.constexpr,
+):
+    """Per-group NVFP4 global scales of two operands, and optionally their product.
+
+    One program handles all ``n`` groups. ``a_broadcast`` / ``b_broadcast`` mark an operand
+    with a single amax shared by every group. The arithmetic matches the eager PyTorch ops
+    bit for bit: dividing a CUDA tensor by a Python scalar multiplies by the FP32 reciprocal,
+    and the product and square root are IEEE round-to-nearest.
+    """
+    idx = tl.arange(0, BLOCK)
+    mask = idx < n
+    amax_a = tl.load(amax_a_ptr + tl.where(a_broadcast, 0, idx), mask=mask, other=0.0)
+    amax_b = tl.load(amax_b_ptr + tl.where(b_broadcast, 0, idx), mask=mask, other=0.0)
+    one = tl.full([BLOCK], 1.0, tl.float32)
+    rcp_a = tl.math.div_rn(one, tl.full([BLOCK], denom_a, tl.float32))
+    rcp_b = tl.math.div_rn(one, tl.full([BLOCK], denom_b, tl.float32))
+    scale_a = amax_a.to(tl.float32) * rcp_a
+    scale_b = amax_b.to(tl.float32) * rcp_b
+    if MODE == 0:
+        tl.store(out_a_ptr + idx, scale_a, mask=mask)
+        tl.store(out_b_ptr + idx, scale_b, mask=mask)
+    else:
+        product = scale_a * scale_b
+        if MODE == 2:
+            product = tl.math.sqrt_rn(product)
+        tl.store(out_a_ptr + idx, product, mask=mask)

transformer_engine/pytorch/ops/fused/grouped_mlp.py

@@ -36,6 +36,7 @@
 from ...tensor.mxfp8_tensor import MXFP8Quantizer, MXFP8Tensor
 from ...tensor.storage.grouped_tensor_storage import GroupedTensorStorage
 from ...triton.grouped_dbias_dscales import compute_grouped_dbias, compute_grouped_dbias_dscales
+from ...triton.nvfp4_global_scale import nvfp4_global_scale_pair
 from ...utils import (
     ceil_div,
     clear_tensor_data,
@@ -398,6 +399,50 @@ def _nvfp4_global_scale(
     return _nvfp4_amax(tensors, columnwise=columnwise).to(torch.float32) / denom
 
 
+def _nvfp4_global_scale_pair(
+    tensors_a: GroupedTensor | Iterable[NVFP4TensorStorage],
+    quantizer_a: Quantizer,
+    columnwise_a: bool,
+    tensors_b: GroupedTensor | Iterable[NVFP4TensorStorage],
+    quantizer_b: Quantizer,
+    columnwise_b: bool,
+    *,
+    num_groups: int,
+    device: torch.device,
+    combine: Optional[Literal["product", "sqrt_product"]] = None,
+) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+    """Per-group NVFP4 global scales of a GEMM operand pair.
+
+    Returns both scales, or with ``combine`` their product (the GEMM alpha) or its square
+    root. When both operands have a second-level scale, a single Triton kernel computes the
+    result instead of several small eager kernels; its output is bit-identical to
+    ``_nvfp4_global_scale`` and the eager combination.
+    """
+    if not getattr(quantizer_a, "disable_second_level_scale", False) and not getattr(
+        quantizer_b, "disable_second_level_scale", False
+    ):
+        fused = nvfp4_global_scale_pair(
+            _nvfp4_amax(tensors_a, columnwise=columnwise_a),
+            _nvfp4_amax(tensors_b, columnwise=columnwise_b),
+            6.0 * _nvfp4_scale_max(quantizer_a),
+            6.0 * _nvfp4_scale_max(quantizer_b),
+            combine=combine,
+        )
+        if fused is not None:
+            return fused
+    scale_a = _nvfp4_global_scale(
+        tensors_a, quantizer_a, columnwise=columnwise_a, num_groups=num_groups, device=device
+    )
+    scale_b = _nvfp4_global_scale(
+        tensors_b, quantizer_b, columnwise=columnwise_b, num_groups=num_groups, device=device
+    )
+    if combine is None:
+        return scale_a, scale_b
+    if combine == "sqrt_product":
+        return torch.sqrt(scale_a * scale_b)
+    return (scale_a * scale_b).to(torch.float32)
+
+
 def _single_quantized_tensor_from_grouped(
     grouped: GroupedTensor,
     quantizer: Optional[MXFP8Quantizer | NVFP4Quantizer] = None,
@@ -698,17 +743,16 @@ def _cudnn_compute_wgrad(
                 device=device,
             )
         else:
-            common_wgrad_kwargs["global_scale_a"] = _nvfp4_global_scale(
+            (
+                common_wgrad_kwargs["global_scale_a"],
+                common_wgrad_kwargs["global_scale_b"],
+            ) = _nvfp4_global_scale_pair(
                 grouped_dy,
                 dy_quantizer,
-                columnwise=True,
-                num_groups=num_groups,
-                device=device,
-            )
-            common_wgrad_kwargs["global_scale_b"] = _nvfp4_global_scale(
+                True,
                 grouped_x,
                 x_quantizer,
-                columnwise=True,
+                True,
                 num_groups=num_groups,
                 device=device,
             )
@@ -1598,22 +1642,17 @@ def fuser_forward(
             # cuDNN receives NVFP4 block-scaled inputs without TE's per-group
             # global scale factors, so alpha supplies the product of the two
             # operand global scales.
-            fc1_alpha_tensor = (
-                _nvfp4_global_scale(
-                    grouped_fc1_x,
-                    fc1_input_quantizer,
-                    columnwise=False,
-                    num_groups=num_groups,
-                    device=device,
-                )
-                * _nvfp4_global_scale(
-                    grouped_fc1_weight,
-                    fc1_weight_quantizer,
-                    columnwise=False,
-                    num_groups=num_groups,
-                    device=device,
-                )
-            ).to(torch.float32)
+            fc1_alpha_tensor = _nvfp4_global_scale_pair(
+                grouped_fc1_x,
+                fc1_input_quantizer,
+                False,
+                grouped_fc1_weight,
+                fc1_weight_quantizer,
+                False,
+                num_groups=num_groups,
+                device=device,
+                combine="product",
+            )
         else:
             fc1_alpha_tensor = alpha_tensor
 
@@ -2420,34 +2459,20 @@ def fuser_backward(
 
         fc2_d_dtype = torch.bfloat16 if use_nvfp4 else torch.float8_e4m3fn
         if use_nvfp4:
-            fc2_dy_global_scale = _nvfp4_global_scale(
+            # DSReLU applies alpha once, so it takes the full product of the two
+            # operand global scales. DGLU applies alpha to both gate branches, so
+            # the wrapper expects sqrt(product) to recover the same factor.
+            fc2_alpha_tensor = _nvfp4_global_scale_pair(
                 grouped_fc2_dy,
                 fc2_grad_output_quantizer,
-                columnwise=False,
-                num_groups=num_groups,
-                device=device,
-            )
-            fc2_weight_col_global_scale = _nvfp4_global_scale(
+                False,
                 grouped_fc2_weight,
                 fc2_weight_quantizer,
-                columnwise=True,
+                True,
                 num_groups=num_groups,
                 device=device,
-            )
-            if activation_is_srelu:
-                # DSReLU applies alpha once, so pass the full product of the
-                # two operand global scales.
-                fc2_alpha_tensor = (
-                    (fc2_dy_global_scale * fc2_weight_col_global_scale)
-                    .to(torch.float32)
-                    .expand(num_groups)
-                )
-            else:
-                # DGLU applies alpha to both gate branches, so the wrapper
-                # expects sqrt(product) to recover the same global-scale factor.
-                fc2_alpha_tensor = torch.sqrt(
-                    fc2_dy_global_scale * fc2_weight_col_global_scale
-                ).expand(num_groups)
+                combine="product" if activation_is_srelu else "sqrt_product",
+            ).expand(num_groups)
             fc2_beta_tensor = get_cached_ones_tensor(num_groups, torch.float32, device)
             fc2_norm_const_tensor = None
         else: