From 5bd8a54d9a0229ce8df480ff57bddfb904210677 Mon Sep 17 00:00:00 2001 From: Siddhartha Raman S Date: Tue, 27 Oct 2026 07:42:31 -0700 Subject: [PATCH] [PyTorch] Fuse NVFP4 global scales and GEMM alpha in the grouped MLP Compute both operands' per-group NVFP4 global scales and their product (or its square root for dGLU) in one Triton launch instead of several eager elementwise kernels. Bit-identical to the eager ops; falls back to them for unsupported inputs. Signed-off-by: Siddhartha Raman S Co-Authored-By: Claude Opus 5.5 --- .../pytorch/nvfp4/test_nvfp4_global_scale.py | 107 +++++++++++ .../common/triton/nvfp4_global_scale.py | 53 ++++++ .../pytorch/ops/fused/grouped_mlp.py | 169 ++++++++++-------- .../pytorch/triton/nvfp4_global_scale.py | 82 +++++++++ 4 files changed, 334 insertions(+), 77 deletions(-) create mode 100644 tests/pytorch/nvfp4/test_nvfp4_global_scale.py create mode 100644 transformer_engine/common/triton/nvfp4_global_scale.py create mode 100644 transformer_engine/pytorch/triton/nvfp4_global_scale.py diff --git a/tests/pytorch/nvfp4/test_nvfp4_global_scale.py b/tests/pytorch/nvfp4/test_nvfp4_global_scale.py new file mode 100644 index 00000000000..945e6ff7bc1 --- /dev/null +++ b/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)) diff --git a/transformer_engine/common/triton/nvfp4_global_scale.py b/transformer_engine/common/triton/nvfp4_global_scale.py new file mode 100644 index 00000000000..aee88be855e --- /dev/null +++ b/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) diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index d678d92749c..881803d14e8 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -37,6 +37,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, @@ -491,6 +492,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, @@ -813,17 +858,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, ) @@ -1681,22 +1725,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 @@ -2052,22 +2091,17 @@ def fuser_forward( sf_swizzled=grouped_fc2_x._with_gemm_swizzled_scales, ) - fc2_fwd_alpha_tensor = ( - _nvfp4_global_scale( - grouped_fc2_x, - fc2_input_quantizer, - columnwise=False, - num_groups=num_groups, - device=device, - ) - * _nvfp4_global_scale( - grouped_fc2_weight, - fc2_weight_quantizer, - columnwise=False, - num_groups=num_groups, - device=device, - ) - ).to(torch.float32) + fc2_fwd_alpha_tensor = _nvfp4_global_scale_pair( + grouped_fc2_x, + fc2_input_quantizer, + False, + grouped_fc2_weight, + fc2_weight_quantizer, + False, + num_groups=num_groups, + device=device, + combine="product", + ) fc2_scales_tensor = ( fc2_scales.detach().to(dtype=torch.float32).reshape(-1, 1, 1) @@ -2571,34 +2605,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: @@ -2998,22 +3018,17 @@ def fuser_backward( sf_swizzled=grouped_fc1_dy._with_gemm_swizzled_scales, ) - fc1_dgrad_alpha = ( - _nvfp4_global_scale( - grouped_fc1_dy, - fc1_grad_output_quantizer, - columnwise=False, - num_groups=num_groups, - device=device, - ) - * _nvfp4_global_scale( - grouped_fc1_weight, - fc1_weight_quantizer, - columnwise=True, - num_groups=num_groups, - device=device, - ) - ).to(torch.float32) + fc1_dgrad_alpha = _nvfp4_global_scale_pair( + grouped_fc1_dy, + fc1_grad_output_quantizer, + False, + grouped_fc1_weight, + fc1_weight_quantizer, + True, + num_groups=num_groups, + device=device, + combine="product", + ) fc1_dgrad_kwargs = { "a_tensor": fc1_dgrad_a_data, diff --git a/transformer_engine/pytorch/triton/nvfp4_global_scale.py b/transformer_engine/pytorch/triton/nvfp4_global_scale.py new file mode 100644 index 00000000000..4efa261ff9b --- /dev/null +++ b/transformer_engine/pytorch/triton/nvfp4_global_scale.py @@ -0,0 +1,82 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""PyTorch wrapper for the fused NVFP4 global-scale Triton kernel.""" + +from typing import Optional, Tuple, Union + +import torch +import triton + +from transformer_engine.common.triton.nvfp4_global_scale import ( + NVFP4_GLOBAL_SCALE_PAIR, + NVFP4_GLOBAL_SCALE_PRODUCT, + NVFP4_GLOBAL_SCALE_SQRT_PRODUCT, + _nvfp4_global_scale_pair_kernel, +) + +# One program covers every group; larger group counts take the eager path. +_MAX_GROUPS = 4096 + +_MODES = { + None: NVFP4_GLOBAL_SCALE_PAIR, + "product": NVFP4_GLOBAL_SCALE_PRODUCT, + "sqrt_product": NVFP4_GLOBAL_SCALE_SQRT_PRODUCT, +} + + +def nvfp4_global_scale_pair( + amax_a: torch.Tensor, + amax_b: torch.Tensor, + denom_a: float, + denom_b: float, + combine: Optional[str] = None, +) -> Optional[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]]: + """Per-group NVFP4 global scales ``amax / denom`` of two operands in one launch. + + Equivalent, bit for bit, to ``amax_a.float() / denom_a`` and ``amax_b.float() / denom_b`` + followed by the optional combination. ``amax_a`` and ``amax_b`` hold one FP32 amax per + group; with ``combine``, either may instead hold a single amax shared by all groups. + + Args: + amax_a, amax_b: per-group (or single) FP32 amaxes on the same CUDA device. + denom_a, denom_b: the operands' ``6 * scale_max`` denominators. + combine: ``None`` returns both scales, ``"product"`` their product and + ``"sqrt_product"`` the square root of the product. + + Returns: + The scale tensor(s), or ``None`` when the inputs are outside what the kernel + handles (non-FP32 amaxes, mismatched group counts, a broadcast amax without + ``combine``, more than 4096 groups); the caller then computes the scales eagerly. + """ + if amax_a.dtype != torch.float32 or amax_b.dtype != torch.float32: + return None + if amax_a.device != amax_b.device or not amax_a.is_cuda: + return None + num_a, num_b = amax_a.numel(), amax_b.numel() + n = max(num_a, num_b) + if n == 0 or n > _MAX_GROUPS or num_a not in (1, n) or num_b not in (1, n): + return None + mode = _MODES[combine] + if mode == NVFP4_GLOBAL_SCALE_PAIR and num_a != num_b: + # Keep each scale's shape: only the combined modes broadcast. + return None + out_a = torch.empty(n, dtype=torch.float32, device=amax_a.device) + out_b = torch.empty_like(out_a) if mode == NVFP4_GLOBAL_SCALE_PAIR else out_a + _nvfp4_global_scale_pair_kernel[(1,)]( + amax_a.contiguous(), + amax_b.contiguous(), + out_a, + out_b, + n, + num_a == 1 and n > 1, + num_b == 1 and n > 1, + float(denom_a), + float(denom_b), + MODE=mode, + BLOCK=max(16, triton.next_power_of_2(n)), + ) + if mode == NVFP4_GLOBAL_SCALE_PAIR: + return out_a, out_b + return out_a