From 95503053dcf4c20237f90f98fff0e32c3d206db5 Mon Sep 17 00:00:00 2001 From: Siddhartha Raman S Date: Mon, 9 Nov 2026 11:34:36 -0800 Subject: [PATCH 1/9] [PyTorch] Grouped MLP: emit FP32 scale gradient in one kernel for FP32 scales The block-scaled fused grouped MLP backward rounds the FP32 scale (router probability) gradient to the compute dtype with `.to(dtype)`. When the scales input is FP32 (e.g. FP32 dispatched router probabilities), autograd then casts that BF16 gradient back to FP32: two cast kernels per layer and microbatch. Return float32(round(grad, compute dtype)) directly from one Triton kernel (`round_through_dtype`) when the scales are FP32. The values are bitwise identical to the previous `.to(bf16)` + autograd `.to(fp32)` pair; BF16/FP16 scales keep the old path. Signed-off-by: Siddhartha Raman S Co-Authored-By: Claude Opus 5.5 --- tests/pytorch/test_grouped_mlp.py | 208 ++++++++++++++++++ .../common/triton/round_through_dtype.py | 31 +++ .../pytorch/ops/fused/grouped_mlp.py | 31 ++- .../pytorch/triton/round_through_dtype.py | 39 ++++ 4 files changed, 308 insertions(+), 1 deletion(-) create mode 100644 transformer_engine/common/triton/round_through_dtype.py create mode 100644 transformer_engine/pytorch/triton/round_through_dtype.py diff --git a/tests/pytorch/test_grouped_mlp.py b/tests/pytorch/test_grouped_mlp.py index 18f9d3f772a..86f4c2a2a2d 100644 --- a/tests/pytorch/test_grouped_mlp.py +++ b/tests/pytorch/test_grouped_mlp.py @@ -48,6 +48,7 @@ is_bf16_available, ) from transformer_engine.pytorch.tensor.grouped_tensor import GroupedTensor +from transformer_engine.pytorch.triton.round_through_dtype import round_through_dtype from transformer_engine.pytorch.utils import get_device_compute_capability import transformer_engine_torch as tex @@ -5063,3 +5064,210 @@ def test_frozen_layer_gets_no_weight_grad(self, frozen: str) -> None: rtol=0, atol=0, ) + + +class TestScalesGradRounding: + """FP32 scale (router probability) gradient of the block-scaled fused grouped MLP. + + The backward rounds the FP32 scale gradient to the compute dtype. For FP32 + scales it returns the rounded values in FP32 from one kernel; this must be + bitwise equal to the previous ``.to(compute_dtype)`` followed by autograd's + cast back to FP32. + """ + + _special_values = ( + 0.0, + -0.0, + float("inf"), + -float("inf"), + 1e-45, + -1e-45, + 1e-40, + -3e-39, + 1.17549435e-38, + 3.3895314e38, + -3.3895314e38, + 3.4028235e38, + 1.0 + 2**-8, + 1.0 + 3 * 2**-9, + 1.0 + 2**-11, + 1.0 + 3 * 2**-12, + 65504.0, + 65520.0, + -65520.0, + 6e-8, + 3e-8, + ) + + @staticmethod + def _assert_bitwise_equal_or_both_nan(out: torch.Tensor, ref: torch.Tensor) -> None: + assert out.dtype == ref.dtype == torch.float32 + assert out.shape == ref.shape + nan = torch.isnan(ref) + assert torch.equal(torch.isnan(out), nan) + assert torch.equal(out[~nan].view(torch.int32), ref[~nan].view(torch.int32)) + + @pytest.mark.parametrize("dtype", (torch.bfloat16, torch.float16)) + @pytest.mark.parametrize("numel", (0, 1, 4093, 131072)) + @pytest.mark.parametrize("scale", (1e-38, 1e-3, 1.0, 1e30)) + def test_round_through_dtype_matches_two_casts( + self, dtype: torch.dtype, numel: int, scale: float + ) -> None: + """Kernel equals ``x.to(dtype).to(float32)`` bitwise, incl. special values.""" + x = torch.randn(numel, device="cuda", dtype=torch.float32) * scale + special = torch.tensor(self._special_values, device="cuda", dtype=torch.float32) + k = min(numel, special.numel()) + x[:k] = special[:k] + out = round_through_dtype(x, dtype) + self._assert_bitwise_equal_or_both_nan(out, x.to(dtype).to(torch.float32)) + + @pytest.mark.parametrize("dtype", (torch.bfloat16, torch.float16)) + def test_round_through_dtype_random_bit_patterns(self, dtype: torch.dtype) -> None: + """Random FP32 bit patterns (normals, denormals, inf, NaN payloads).""" + bits = torch.randint(-(2**31), 2**31 - 1, (1 << 22,), device="cuda", dtype=torch.int64) + bits = bits.to(torch.int32) + x = bits.view(torch.float32) + out = round_through_dtype(x, dtype) + self._assert_bitwise_equal_or_both_nan(out, x.to(dtype).to(torch.float32)) + + def test_round_through_dtype_non_contiguous_and_shape(self) -> None: + """Non-contiguous inputs are handled and the input shape is kept.""" + x = torch.randn(64, 96, device="cuda", dtype=torch.float32).t() + out = round_through_dtype(x, torch.bfloat16) + assert out.shape == x.shape + self._assert_bitwise_equal_or_both_nan(out, x.to(torch.bfloat16).to(torch.float32)) + + def test_round_through_dtype_cuda_graph(self) -> None: + """The kernel can be captured in a CUDA graph and replays on new data.""" + x = torch.randn(32768, device="cuda", dtype=torch.float32) + round_through_dtype(x, torch.bfloat16) # compile outside capture + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = round_through_dtype(x, torch.bfloat16) + for scale in (1.0, 1e-20, 1e20): + x.copy_(torch.randn_like(x) * scale) + graph.replay() + torch.cuda.synchronize() + self._assert_bitwise_equal_or_both_nan(out, x.to(torch.bfloat16).to(torch.float32)) + + @pytest.mark.parametrize("quantization", ("mxfp8", "nvfp4_rht")) + @pytest.mark.parametrize("scales_dtype", (torch.float32, torch.bfloat16)) + @pytest.mark.parametrize("bias", (False, True)) + def test_fused_grouped_mlp_scales_grad( + self, + monkeypatch, + quantization: str, + scales_dtype: torch.dtype, + bias: bool, + ) -> None: + """Scale gradient of the fused op equals the old cast path bitwise. + + The FP32 gradient produced by the fused kernels is recorded right before + it is rounded. For FP32 scales the gradient seen by the caller must be + ``fp32_grad.to(bf16).to(fp32)`` (what the old ``.to(bf16)`` plus autograd's + cast back produced) and no cast of the scale gradient may run in the + backward; for BF16 scales it must be ``fp32_grad.to(bf16)`` as before. + """ + if quantization == "mxfp8" and not mxfp8_available: + pytest.skip(reason_for_no_mxfp8) + if quantization == "nvfp4_rht" and not nvfp4_available: + pytest.skip(reason_for_no_nvfp4) + if quantization == "nvfp4_rht" and bias: + pytest.skip("NVFP4 RHT SwiGLU grouped MLP coverage is limited to no-bias") + if not ( + _cudnn_frontend_version_supported() + and (device_is_blackwell or device_is_rubin) + and GroupedMLP_CuTeGEMMGLU.is_supported() + ): + pytest.skip("Block-scaled fused grouped MLP is not supported on this system") + + dtype = torch.bfloat16 + group_size, hidden_size, split_alignment = 4, 128, 256 + split_sizes = [split_alignment * i for i in range(1, group_size + 1)] + random.shuffle(split_sizes) + split_sizes = torch.tensor(split_sizes, dtype=torch.int, device="cuda") + num_tokens = int(split_sizes.sum().item()) + + recipe = make_recipe(quantization) + with te.quantized_model_init(enabled=True, recipe=recipe): + fc1 = te.ops.GroupedLinear( + group_size, + hidden_size, + 2 * hidden_size, + bias=bias, + device="cuda", + dtype=dtype, + split_alignment=split_alignment, + ) + fc2 = te.ops.GroupedLinear( + group_size, + hidden_size, + hidden_size, + bias=bias, + device="cuda", + dtype=dtype, + scale_bias=bias, + split_alignment=split_alignment, + ) + module = te.ops.Sequential(fc1, te.ops.ScaledSwiGLU(glu_interleave_size=32), fc2) + + x = (torch.rand(num_tokens, hidden_size, device="cuda", dtype=dtype) - 0.5) / 2 + x.requires_grad_(True) + probs = torch.rand(num_tokens, device="cuda", dtype=scales_dtype).requires_grad_(True) + dy = (torch.rand(num_tokens, hidden_size, device="cuda", dtype=dtype) - 0.5) / 2 + + recorded = [] + cast_scales_grad = grouped_mlp_module._cast_scales_grad + + def _recording_cast_scales_grad(grad_scales, **kwargs): + out = cast_scales_grad(grad_scales, **kwargs) + recorded.append((grad_scales.clone(), out)) + return out + + monkeypatch.setattr(grouped_mlp_module, "_cast_scales_grad", _recording_cast_scales_grad) + + with te.autocast(enabled=True, recipe=recipe): + fc2_extra = (split_sizes, probs) if bias else (split_sizes,) + y = module(x, split_sizes, probs, *fc2_extra) + forward_ops = module._module_groups[0]._forward_ops + assert len(forward_ops) == 1 and isinstance(forward_ops[0][0], GroupedMLP_CuTeGEMMGLU) + + # Count dtype casts of the [num_tokens] scale gradient during the backward + from torch.utils._python_dispatch import TorchDispatchMode + + class _CountScaleGradCasts(TorchDispatchMode): + def __init__(self): + super().__init__() + self.count = 0 + + def __torch_dispatch__(self, func, types, args=(), kwargs=None): + out = func(*args, **(kwargs or {})) + if func in (torch.ops.aten._to_copy.default, torch.ops.aten.copy_.default): + src = args[1] if func is torch.ops.aten.copy_.default else args[0] + dst = args[0] if func is torch.ops.aten.copy_.default else out + if ( + isinstance(src, torch.Tensor) + and src.numel() == num_tokens + and src.dim() == 1 + and src.dtype != dst.dtype + ): + self.count += 1 + return out + + cast_counter = _CountScaleGradCasts() + with cast_counter: + y.backward(dy) + + assert len(recorded) == 1 + grad_fp32, returned = recorded[0] + assert grad_fp32.dtype == torch.float32 + assert probs.grad.dtype == scales_dtype + if scales_dtype == torch.float32: + expected = grad_fp32.to(dtype).to(torch.float32) + self._assert_bitwise_equal_or_both_nan(probs.grad, expected) + assert returned.dtype == torch.float32 + assert cast_counter.count == 0 + else: + expected = grad_fp32.to(dtype) + assert torch.equal(probs.grad.view(torch.int16), expected.view(torch.int16)) diff --git a/transformer_engine/common/triton/round_through_dtype.py b/transformer_engine/common/triton/round_through_dtype.py new file mode 100644 index 00000000000..9fe4df36963 --- /dev/null +++ b/transformer_engine/common/triton/round_through_dtype.py @@ -0,0 +1,31 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Triton kernel that rounds an FP32 tensor to a 16-bit float format, keeping FP32 storage.""" + +import triton +import triton.language as tl + + +@triton.jit +def _round_through_dtype_kernel( + x_ptr, + y_ptr, + n, + ROUND_TO_BF16: tl.constexpr, + BLOCK: tl.constexpr, +): + """y = float32(round_to_nearest_even(x, BF16 or FP16)) in one pass. + + Equivalent to ``x.to(torch.bfloat16).to(torch.float32)`` (or FP16), without + materializing the 16-bit intermediate. + """ + offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + mask = offs < n + x = tl.load(x_ptr + offs, mask=mask, other=0.0) + if ROUND_TO_BF16: + y = x.to(tl.bfloat16).to(tl.float32) + else: + y = x.to(tl.float16).to(tl.float32) + tl.store(y_ptr + offs, y, mask=mask) diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index d29ae1fd434..e8ba7174a99 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.round_through_dtype import round_through_dtype from ...utils import ( ceil_div, clear_tensor_data, @@ -492,6 +493,30 @@ def _nvfp4_global_scale( return _nvfp4_amax(tensors, columnwise=columnwise).to(torch.float32) / denom +def _cast_scales_grad( + grad_scales: torch.Tensor, + *, + compute_dtype: torch.dtype, + scales_dtype: Optional[torch.dtype], +) -> torch.Tensor: + """Round the FP32 scale (router probability) gradient to the compute dtype. + + The gradient is rounded to ``compute_dtype``. When the scales input is FP32 + while the compute dtype is BF16/FP16 (e.g. FP32 router probabilities from an + expert-parallel dispatcher), the rounded values are returned directly in FP32 + by one kernel. Casting to ``compute_dtype`` here would make autograd cast the + gradient back to the FP32 scales dtype, i.e. two cast kernels that produce + exactly the same values. + """ + if ( + scales_dtype == torch.float32 + and grad_scales.dtype == torch.float32 + and compute_dtype in (torch.bfloat16, torch.float16) + ): + return round_through_dtype(grad_scales, compute_dtype) + return grad_scales.to(dtype=compute_dtype) + + def _single_quantized_tensor_from_grouped( grouped: GroupedTensor, quantizer: Optional[MXFP8Quantizer | NVFP4Quantizer] = None, @@ -2889,7 +2914,11 @@ def fuser_backward( fc2_bias_grads = [fc2_dbias_packed[idx] for idx in range(num_groups)] if grad_scales is not None: - grad_scales = grad_scales.to(dtype=dtype) + grad_scales = _cast_scales_grad( + grad_scales, + compute_dtype=dtype, + scales_dtype=None if scales is None else scales.dtype, + ) fc1_bias_grads: Optional[list[Optional[torch.Tensor]]] = None fc1_bias_grad_packed: Optional[torch.Tensor] = None diff --git a/transformer_engine/pytorch/triton/round_through_dtype.py b/transformer_engine/pytorch/triton/round_through_dtype.py new file mode 100644 index 00000000000..9ab75f9683b --- /dev/null +++ b/transformer_engine/pytorch/triton/round_through_dtype.py @@ -0,0 +1,39 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""PyTorch wrapper for the round-through-dtype Triton kernel.""" + +import torch +import triton + +from transformer_engine.common.triton.round_through_dtype import _round_through_dtype_kernel + +_BLOCK = 1024 + + +def round_through_dtype(x: torch.Tensor, dtype: torch.dtype) -> torch.Tensor: + """Round an FP32 tensor to ``dtype`` (BF16 or FP16) and return it in FP32. + + Bitwise equal to ``x.to(dtype).to(torch.float32)`` for every non-NaN input + (NaN stays NaN), but runs as a single kernel and never materializes the + 16-bit intermediate. Returns a new contiguous tensor with the shape of ``x``. + """ + if x.dtype != torch.float32: + raise ValueError(f"round_through_dtype expects an FP32 tensor, got {x.dtype}") + if dtype not in (torch.bfloat16, torch.float16): + raise ValueError(f"round_through_dtype supports BF16 and FP16, got {dtype}") + if not x.is_cuda: + return x.to(dtype).to(torch.float32) + x = x.contiguous() + y = torch.empty_like(x) + n = x.numel() + if n > 0: + _round_through_dtype_kernel[(triton.cdiv(n, _BLOCK),)]( + x, + y, + n, + ROUND_TO_BF16=dtype == torch.bfloat16, + BLOCK=_BLOCK, + ) + return y From 423987ab0d8b26a8f5e7a0197f9fc8296e06d2c2 Mon Sep 17 00:00:00 2001 From: Siddhartha Raman S Date: Fri, 13 Nov 2026 12:15:15 -0800 Subject: [PATCH 2/9] [PyTorch] Grouped MLP: keep the FP32 scale gradient, without rounding it Review feedback: the block-scaled fused grouped MLP rounded the FP32 scale (router probability) gradient from the cuDNN kernels to the compute dtype with .to(dtype), a leftover from when the scales were always BF16. With FP32 scales, autograd then cast it back to FP32, so the gradient lost precision for no reason and cost two cast kernels. Return the gradient in the dtype of the scales input instead, as the BF16 fused grouped MLP path already does: FP32 scales get the kernels' FP32 values unchanged (no cast), BF16 scales get one cast as before. This drops the round_through_dtype Triton kernel. The test checks the FP32 gradient is returned bitwise unchanged with no cast in the backward, and the BF16 case. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Siddhartha Raman S --- tests/pytorch/test_grouped_mlp.py | 93 ++----------------- .../common/triton/round_through_dtype.py | 31 ------- .../pytorch/ops/fused/grouped_mlp.py | 32 ++----- .../pytorch/triton/round_through_dtype.py | 39 -------- 4 files changed, 19 insertions(+), 176 deletions(-) delete mode 100644 transformer_engine/common/triton/round_through_dtype.py delete mode 100644 transformer_engine/pytorch/triton/round_through_dtype.py diff --git a/tests/pytorch/test_grouped_mlp.py b/tests/pytorch/test_grouped_mlp.py index 86f4c2a2a2d..91adc939aa0 100644 --- a/tests/pytorch/test_grouped_mlp.py +++ b/tests/pytorch/test_grouped_mlp.py @@ -48,7 +48,6 @@ is_bf16_available, ) from transformer_engine.pytorch.tensor.grouped_tensor import GroupedTensor -from transformer_engine.pytorch.triton.round_through_dtype import round_through_dtype from transformer_engine.pytorch.utils import get_device_compute_capability import transformer_engine_torch as tex @@ -5066,39 +5065,14 @@ def test_frozen_layer_gets_no_weight_grad(self, frozen: str) -> None: ) -class TestScalesGradRounding: - """FP32 scale (router probability) gradient of the block-scaled fused grouped MLP. +class TestScalesGradDtype: + """Scale (router probability) gradient of the block-scaled fused grouped MLP. - The backward rounds the FP32 scale gradient to the compute dtype. For FP32 - scales it returns the rounded values in FP32 from one kernel; this must be - bitwise equal to the previous ``.to(compute_dtype)`` followed by autograd's - cast back to FP32. + The fused kernels compute it in FP32. It is returned in the dtype of the scales input: + unchanged for FP32 scales (no rounding through the compute dtype, no cast), cast once + for BF16 scales. """ - _special_values = ( - 0.0, - -0.0, - float("inf"), - -float("inf"), - 1e-45, - -1e-45, - 1e-40, - -3e-39, - 1.17549435e-38, - 3.3895314e38, - -3.3895314e38, - 3.4028235e38, - 1.0 + 2**-8, - 1.0 + 3 * 2**-9, - 1.0 + 2**-11, - 1.0 + 3 * 2**-12, - 65504.0, - 65520.0, - -65520.0, - 6e-8, - 3e-8, - ) - @staticmethod def _assert_bitwise_equal_or_both_nan(out: torch.Tensor, ref: torch.Tensor) -> None: assert out.dtype == ref.dtype == torch.float32 @@ -5107,50 +5081,6 @@ def _assert_bitwise_equal_or_both_nan(out: torch.Tensor, ref: torch.Tensor) -> N assert torch.equal(torch.isnan(out), nan) assert torch.equal(out[~nan].view(torch.int32), ref[~nan].view(torch.int32)) - @pytest.mark.parametrize("dtype", (torch.bfloat16, torch.float16)) - @pytest.mark.parametrize("numel", (0, 1, 4093, 131072)) - @pytest.mark.parametrize("scale", (1e-38, 1e-3, 1.0, 1e30)) - def test_round_through_dtype_matches_two_casts( - self, dtype: torch.dtype, numel: int, scale: float - ) -> None: - """Kernel equals ``x.to(dtype).to(float32)`` bitwise, incl. special values.""" - x = torch.randn(numel, device="cuda", dtype=torch.float32) * scale - special = torch.tensor(self._special_values, device="cuda", dtype=torch.float32) - k = min(numel, special.numel()) - x[:k] = special[:k] - out = round_through_dtype(x, dtype) - self._assert_bitwise_equal_or_both_nan(out, x.to(dtype).to(torch.float32)) - - @pytest.mark.parametrize("dtype", (torch.bfloat16, torch.float16)) - def test_round_through_dtype_random_bit_patterns(self, dtype: torch.dtype) -> None: - """Random FP32 bit patterns (normals, denormals, inf, NaN payloads).""" - bits = torch.randint(-(2**31), 2**31 - 1, (1 << 22,), device="cuda", dtype=torch.int64) - bits = bits.to(torch.int32) - x = bits.view(torch.float32) - out = round_through_dtype(x, dtype) - self._assert_bitwise_equal_or_both_nan(out, x.to(dtype).to(torch.float32)) - - def test_round_through_dtype_non_contiguous_and_shape(self) -> None: - """Non-contiguous inputs are handled and the input shape is kept.""" - x = torch.randn(64, 96, device="cuda", dtype=torch.float32).t() - out = round_through_dtype(x, torch.bfloat16) - assert out.shape == x.shape - self._assert_bitwise_equal_or_both_nan(out, x.to(torch.bfloat16).to(torch.float32)) - - def test_round_through_dtype_cuda_graph(self) -> None: - """The kernel can be captured in a CUDA graph and replays on new data.""" - x = torch.randn(32768, device="cuda", dtype=torch.float32) - round_through_dtype(x, torch.bfloat16) # compile outside capture - torch.cuda.synchronize() - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - out = round_through_dtype(x, torch.bfloat16) - for scale in (1.0, 1e-20, 1e20): - x.copy_(torch.randn_like(x) * scale) - graph.replay() - torch.cuda.synchronize() - self._assert_bitwise_equal_or_both_nan(out, x.to(torch.bfloat16).to(torch.float32)) - @pytest.mark.parametrize("quantization", ("mxfp8", "nvfp4_rht")) @pytest.mark.parametrize("scales_dtype", (torch.float32, torch.bfloat16)) @pytest.mark.parametrize("bias", (False, True)) @@ -5161,13 +5091,11 @@ def test_fused_grouped_mlp_scales_grad( scales_dtype: torch.dtype, bias: bool, ) -> None: - """Scale gradient of the fused op equals the old cast path bitwise. + """The scale gradient keeps the kernels' FP32 values. - The FP32 gradient produced by the fused kernels is recorded right before - it is rounded. For FP32 scales the gradient seen by the caller must be - ``fp32_grad.to(bf16).to(fp32)`` (what the old ``.to(bf16)`` plus autograd's - cast back produced) and no cast of the scale gradient may run in the - backward; for BF16 scales it must be ``fp32_grad.to(bf16)`` as before. + The FP32 gradient produced by the fused kernels is recorded before it is returned. + For FP32 scales the caller must get exactly those values and no cast of the scale + gradient may run in the backward; for BF16 scales it must be ``fp32_grad.to(bf16)``. """ if quantization == "mxfp8" and not mxfp8_available: pytest.skip(reason_for_no_mxfp8) @@ -5264,8 +5192,7 @@ def __torch_dispatch__(self, func, types, args=(), kwargs=None): assert grad_fp32.dtype == torch.float32 assert probs.grad.dtype == scales_dtype if scales_dtype == torch.float32: - expected = grad_fp32.to(dtype).to(torch.float32) - self._assert_bitwise_equal_or_both_nan(probs.grad, expected) + self._assert_bitwise_equal_or_both_nan(probs.grad, grad_fp32) assert returned.dtype == torch.float32 assert cast_counter.count == 0 else: diff --git a/transformer_engine/common/triton/round_through_dtype.py b/transformer_engine/common/triton/round_through_dtype.py deleted file mode 100644 index 9fe4df36963..00000000000 --- a/transformer_engine/common/triton/round_through_dtype.py +++ /dev/null @@ -1,31 +0,0 @@ -# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# -# See LICENSE for license information. - -"""Triton kernel that rounds an FP32 tensor to a 16-bit float format, keeping FP32 storage.""" - -import triton -import triton.language as tl - - -@triton.jit -def _round_through_dtype_kernel( - x_ptr, - y_ptr, - n, - ROUND_TO_BF16: tl.constexpr, - BLOCK: tl.constexpr, -): - """y = float32(round_to_nearest_even(x, BF16 or FP16)) in one pass. - - Equivalent to ``x.to(torch.bfloat16).to(torch.float32)`` (or FP16), without - materializing the 16-bit intermediate. - """ - offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) - mask = offs < n - x = tl.load(x_ptr + offs, mask=mask, other=0.0) - if ROUND_TO_BF16: - y = x.to(tl.bfloat16).to(tl.float32) - else: - y = x.to(tl.float16).to(tl.float32) - tl.store(y_ptr + offs, y, mask=mask) diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index e8ba7174a99..fa526750d71 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -37,7 +37,6 @@ 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.round_through_dtype import round_through_dtype from ...utils import ( ceil_div, clear_tensor_data, @@ -494,27 +493,16 @@ def _nvfp4_global_scale( def _cast_scales_grad( - grad_scales: torch.Tensor, - *, - compute_dtype: torch.dtype, - scales_dtype: Optional[torch.dtype], + grad_scales: torch.Tensor, *, scales_dtype: Optional[torch.dtype] ) -> torch.Tensor: - """Round the FP32 scale (router probability) gradient to the compute dtype. - - The gradient is rounded to ``compute_dtype``. When the scales input is FP32 - while the compute dtype is BF16/FP16 (e.g. FP32 router probabilities from an - expert-parallel dispatcher), the rounded values are returned directly in FP32 - by one kernel. Casting to ``compute_dtype`` here would make autograd cast the - gradient back to the FP32 scales dtype, i.e. two cast kernels that produce - exactly the same values. + """Return the scale (router probability) gradient in the dtype of the scales input. + + The fused kernels compute it in FP32. FP32 scales get it unchanged (no rounding, no cast); + lower-precision scales get one cast, as the BF16 fused grouped MLP path does. """ - if ( - scales_dtype == torch.float32 - and grad_scales.dtype == torch.float32 - and compute_dtype in (torch.bfloat16, torch.float16) - ): - return round_through_dtype(grad_scales, compute_dtype) - return grad_scales.to(dtype=compute_dtype) + if scales_dtype is None or grad_scales.dtype == scales_dtype: + return grad_scales + return grad_scales.to(dtype=scales_dtype) def _single_quantized_tensor_from_grouped( @@ -2915,9 +2903,7 @@ def fuser_backward( if grad_scales is not None: grad_scales = _cast_scales_grad( - grad_scales, - compute_dtype=dtype, - scales_dtype=None if scales is None else scales.dtype, + grad_scales, scales_dtype=None if scales is None else scales.dtype ) fc1_bias_grads: Optional[list[Optional[torch.Tensor]]] = None diff --git a/transformer_engine/pytorch/triton/round_through_dtype.py b/transformer_engine/pytorch/triton/round_through_dtype.py deleted file mode 100644 index 9ab75f9683b..00000000000 --- a/transformer_engine/pytorch/triton/round_through_dtype.py +++ /dev/null @@ -1,39 +0,0 @@ -# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# -# See LICENSE for license information. - -"""PyTorch wrapper for the round-through-dtype Triton kernel.""" - -import torch -import triton - -from transformer_engine.common.triton.round_through_dtype import _round_through_dtype_kernel - -_BLOCK = 1024 - - -def round_through_dtype(x: torch.Tensor, dtype: torch.dtype) -> torch.Tensor: - """Round an FP32 tensor to ``dtype`` (BF16 or FP16) and return it in FP32. - - Bitwise equal to ``x.to(dtype).to(torch.float32)`` for every non-NaN input - (NaN stays NaN), but runs as a single kernel and never materializes the - 16-bit intermediate. Returns a new contiguous tensor with the shape of ``x``. - """ - if x.dtype != torch.float32: - raise ValueError(f"round_through_dtype expects an FP32 tensor, got {x.dtype}") - if dtype not in (torch.bfloat16, torch.float16): - raise ValueError(f"round_through_dtype supports BF16 and FP16, got {dtype}") - if not x.is_cuda: - return x.to(dtype).to(torch.float32) - x = x.contiguous() - y = torch.empty_like(x) - n = x.numel() - if n > 0: - _round_through_dtype_kernel[(triton.cdiv(n, _BLOCK),)]( - x, - y, - n, - ROUND_TO_BF16=dtype == torch.bfloat16, - BLOCK=_BLOCK, - ) - return y From 6aa6f814559b689393e052ac749e8813f6f2d77b Mon Sep 17 00:00:00 2001 From: Siddhartha Raman S Date: Fri, 13 Nov 2026 16:29:12 -0800 Subject: [PATCH 3/9] [PyTorch] Grouped MLP: FP32 router-probability scales end to end, no casts Review feedback: the scale gradient from the cuDNN kernels is already FP32, so return it as is (no helper; autograd casts it for lower-precision scales), and pass the scales to the cuDNN kernels in FP32 in the forward for every recipe, as the NVFP4 path already did, instead of rounding them to the compute dtype for MXFP8. Drops the cast-counting test. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Siddhartha Raman S --- tests/pytorch/test_grouped_mlp.py | 135 ------------------ .../pytorch/ops/fused/grouped_mlp.py | 22 +-- 2 files changed, 1 insertion(+), 156 deletions(-) diff --git a/tests/pytorch/test_grouped_mlp.py b/tests/pytorch/test_grouped_mlp.py index 70d3f180cda..8c3bf5b2c0e 100644 --- a/tests/pytorch/test_grouped_mlp.py +++ b/tests/pytorch/test_grouped_mlp.py @@ -5004,138 +5004,3 @@ def test_frozen_layer_gets_no_weight_grad(self, frozen: str) -> None: rtol=0, atol=0, ) - - -class TestScalesGradDtype: - """Scale (router probability) gradient of the block-scaled fused grouped MLP. - - The fused kernels compute it in FP32. It is returned in the dtype of the scales input: - unchanged for FP32 scales (no rounding through the compute dtype, no cast), cast once - for BF16 scales. - """ - - @staticmethod - def _assert_bitwise_equal_or_both_nan(out: torch.Tensor, ref: torch.Tensor) -> None: - assert out.dtype == ref.dtype == torch.float32 - assert out.shape == ref.shape - nan = torch.isnan(ref) - assert torch.equal(torch.isnan(out), nan) - assert torch.equal(out[~nan].view(torch.int32), ref[~nan].view(torch.int32)) - - @pytest.mark.parametrize("quantization", ("mxfp8", "nvfp4_rht")) - @pytest.mark.parametrize("scales_dtype", (torch.float32, torch.bfloat16)) - @pytest.mark.parametrize("bias", (False, True)) - def test_fused_grouped_mlp_scales_grad( - self, - monkeypatch, - quantization: str, - scales_dtype: torch.dtype, - bias: bool, - ) -> None: - """The scale gradient keeps the kernels' FP32 values. - - The FP32 gradient produced by the fused kernels is recorded before it is returned. - For FP32 scales the caller must get exactly those values and no cast of the scale - gradient may run in the backward; for BF16 scales it must be ``fp32_grad.to(bf16)``. - """ - if quantization == "mxfp8" and not mxfp8_available: - pytest.skip(reason_for_no_mxfp8) - if quantization == "nvfp4_rht" and not nvfp4_available: - pytest.skip(reason_for_no_nvfp4) - if quantization == "nvfp4_rht" and bias: - pytest.skip("NVFP4 RHT SwiGLU grouped MLP coverage is limited to no-bias") - if not ( - _cudnn_frontend_version_supported() - and (device_is_blackwell or device_is_rubin) - and GroupedMLP_CuTeGEMMGLU.is_supported() - ): - pytest.skip("Block-scaled fused grouped MLP is not supported on this system") - - dtype = torch.bfloat16 - group_size, hidden_size, split_alignment = 4, 128, 256 - split_sizes = [split_alignment * i for i in range(1, group_size + 1)] - random.shuffle(split_sizes) - split_sizes = torch.tensor(split_sizes, dtype=torch.int, device="cuda") - num_tokens = int(split_sizes.sum().item()) - - recipe = make_recipe(quantization) - with te.quantized_model_init(enabled=True, recipe=recipe): - fc1 = te.ops.GroupedLinear( - group_size, - hidden_size, - 2 * hidden_size, - bias=bias, - device="cuda", - dtype=dtype, - split_alignment=split_alignment, - ) - fc2 = te.ops.GroupedLinear( - group_size, - hidden_size, - hidden_size, - bias=bias, - device="cuda", - dtype=dtype, - scale_bias=bias, - split_alignment=split_alignment, - ) - module = te.ops.Sequential(fc1, te.ops.ScaledSwiGLU(glu_interleave_size=32), fc2) - - x = (torch.rand(num_tokens, hidden_size, device="cuda", dtype=dtype) - 0.5) / 2 - x.requires_grad_(True) - probs = torch.rand(num_tokens, device="cuda", dtype=scales_dtype).requires_grad_(True) - dy = (torch.rand(num_tokens, hidden_size, device="cuda", dtype=dtype) - 0.5) / 2 - - recorded = [] - cast_scales_grad = grouped_mlp_module._cast_scales_grad - - def _recording_cast_scales_grad(grad_scales, **kwargs): - out = cast_scales_grad(grad_scales, **kwargs) - recorded.append((grad_scales.clone(), out)) - return out - - monkeypatch.setattr(grouped_mlp_module, "_cast_scales_grad", _recording_cast_scales_grad) - - with te.autocast(enabled=True, recipe=recipe): - fc2_extra = (split_sizes, probs) if bias else (split_sizes,) - y = module(x, split_sizes, probs, *fc2_extra) - forward_ops = module._module_groups[0]._forward_ops - assert len(forward_ops) == 1 and isinstance(forward_ops[0][0], GroupedMLP_CuTeGEMMGLU) - - # Count dtype casts of the [num_tokens] scale gradient during the backward - from torch.utils._python_dispatch import TorchDispatchMode - - class _CountScaleGradCasts(TorchDispatchMode): - def __init__(self): - super().__init__() - self.count = 0 - - def __torch_dispatch__(self, func, types, args=(), kwargs=None): - out = func(*args, **(kwargs or {})) - if func in (torch.ops.aten._to_copy.default, torch.ops.aten.copy_.default): - src = args[1] if func is torch.ops.aten.copy_.default else args[0] - dst = args[0] if func is torch.ops.aten.copy_.default else out - if ( - isinstance(src, torch.Tensor) - and src.numel() == num_tokens - and src.dim() == 1 - and src.dtype != dst.dtype - ): - self.count += 1 - return out - - cast_counter = _CountScaleGradCasts() - with cast_counter: - y.backward(dy) - - assert len(recorded) == 1 - grad_fp32, returned = recorded[0] - assert grad_fp32.dtype == torch.float32 - assert probs.grad.dtype == scales_dtype - if scales_dtype == torch.float32: - self._assert_bitwise_equal_or_both_nan(probs.grad, grad_fp32) - assert returned.dtype == torch.float32 - assert cast_counter.count == 0 - else: - expected = grad_fp32.to(dtype) - assert torch.equal(probs.grad.view(torch.int16), expected.view(torch.int16)) diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index 81d251771aa..20810307a4a 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -398,19 +398,6 @@ def _nvfp4_global_scale( return _nvfp4_amax(tensors, columnwise=columnwise).to(torch.float32) / denom -def _cast_scales_grad( - grad_scales: torch.Tensor, *, scales_dtype: Optional[torch.dtype] -) -> torch.Tensor: - """Return the scale (router probability) gradient in the dtype of the scales input. - - The fused kernels compute it in FP32. FP32 scales get it unchanged (no rounding, no cast); - lower-precision scales get one cast, as the BF16 fused grouped MLP path does. - """ - if scales_dtype is None or grad_scales.dtype == scales_dtype: - return grad_scales - return grad_scales.to(dtype=scales_dtype) - - def _single_quantized_tensor_from_grouped( grouped: GroupedTensor, quantizer: Optional[MXFP8Quantizer | NVFP4Quantizer] = None, @@ -1603,9 +1590,7 @@ def fuser_forward( fc1_prob_tensor = None if not unit_activation_scale: fc1_prob_tensor = ( - scales.detach() - .to(dtype=torch.float32 if use_nvfp4 else dtype) - .reshape((-1,) if canonical else (-1, 1, 1)) + scales.detach().to(dtype=torch.float32).reshape((-1,) if canonical else (-1, 1, 1)) ) if use_nvfp4: # cuDNN receives NVFP4 block-scaled inputs without TE's per-group @@ -2714,11 +2699,6 @@ def fuser_backward( else: fc2_bias_grads = [fc2_dbias_packed[idx] for idx in range(num_groups)] - if grad_scales is not None: - grad_scales = _cast_scales_grad( - grad_scales, scales_dtype=None if scales is None else scales.dtype - ) - fc1_bias_grads: Optional[list[Optional[torch.Tensor]]] = None fc1_bias_grad_packed: Optional[torch.Tensor] = None if fc1_op.has_bias: From 7dc8e0c176c5e129abffced717ace47a7da30a6c Mon Sep 17 00:00:00 2001 From: Siddhartha Raman S Date: Fri, 13 Nov 2026 17:36:15 -0800 Subject: [PATCH 4/9] [PyTorch] Test: the MXFP8 GLU forward now gets FP32 probabilities Co-Authored-By: Claude Opus 5.5 Signed-off-by: Siddhartha Raman S --- tests/pytorch/test_grouped_mlp.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/pytorch/test_grouped_mlp.py b/tests/pytorch/test_grouped_mlp.py index 8c3bf5b2c0e..b1e13fc5e9e 100644 --- a/tests/pytorch/test_grouped_mlp.py +++ b/tests/pytorch/test_grouped_mlp.py @@ -1804,7 +1804,7 @@ def test_grouped_mlp_glu_mxfp8_real_cudnn_fusion( "a_dtype": torch.float8_e4m3fn, "b_dtype": torch.float8_e4m3fn, "alpha_dtype": torch.bfloat16, - "prob_dtype": torch.bfloat16, + "prob_dtype": torch.float32, "c_dtype": torch.bfloat16, "d_dtype": torch.float8_e4m3fn, "act_func": act_func, From f871cab89000b92bab91e4cbcdb28ffd788b05df Mon Sep 17 00:00:00 2001 From: vthumbe1503 Date: Fri, 13 Nov 2026 20:47:55 -0800 Subject: [PATCH 5/9] Cast only for NVFP4 Signed-off-by: vthumbe1503 --- transformer_engine/pytorch/ops/fused/grouped_mlp.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index 20810307a4a..1fcf6c1c115 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -1588,10 +1588,14 @@ def fuser_forward( current_stream = torch.cuda.current_stream(device.index).cuda_stream fc1_prob_tensor = None - if not unit_activation_scale: + if not unit_activation_scale: fc1_prob_tensor = ( - scales.detach().to(dtype=torch.float32).reshape((-1,) if canonical else (-1, 1, 1)) + scales.detach().reshape((-1,) if canonical else (-1, 1, 1)) ) + if use_nvfp4: + # cuDNN NVFP4 kernel only accepts fp32 scales. + fc1_prob_tensor = fc1_prob_tensor.to(torch.float32) + if use_nvfp4: # cuDNN receives NVFP4 block-scaled inputs without TE's per-group # global scale factors, so alpha supplies the product of the two From 55152d32ab4e4f6ce9b258f1c6f8abe66a91e2cf Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 14 Nov 2026 04:49:19 +0000 Subject: [PATCH 6/9] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- transformer_engine/pytorch/ops/fused/grouped_mlp.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index 1fcf6c1c115..05ad6d178c9 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -1588,14 +1588,12 @@ def fuser_forward( current_stream = torch.cuda.current_stream(device.index).cuda_stream fc1_prob_tensor = None - if not unit_activation_scale: - fc1_prob_tensor = ( - scales.detach().reshape((-1,) if canonical else (-1, 1, 1)) - ) + if not unit_activation_scale: + fc1_prob_tensor = scales.detach().reshape((-1,) if canonical else (-1, 1, 1)) if use_nvfp4: # cuDNN NVFP4 kernel only accepts fp32 scales. fc1_prob_tensor = fc1_prob_tensor.to(torch.float32) - + if use_nvfp4: # cuDNN receives NVFP4 block-scaled inputs without TE's per-group # global scale factors, so alpha supplies the product of the two From 8f24294fd66f63baf8b957406c75fe6c69181a55 Mon Sep 17 00:00:00 2001 From: vthumbe1503 Date: Fri, 13 Nov 2026 20:50:40 -0800 Subject: [PATCH 7/9] Move fp32 scale conversion for NVFP4 condition Signed-off-by: vthumbe1503 --- transformer_engine/pytorch/ops/fused/grouped_mlp.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index 05ad6d178c9..b07c42562b3 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -1590,11 +1590,10 @@ def fuser_forward( fc1_prob_tensor = None if not unit_activation_scale: fc1_prob_tensor = scales.detach().reshape((-1,) if canonical else (-1, 1, 1)) - if use_nvfp4: - # cuDNN NVFP4 kernel only accepts fp32 scales. - fc1_prob_tensor = fc1_prob_tensor.to(torch.float32) if use_nvfp4: + # cuDNN NVFP4 kernel only accepts fp32 scales. + fc1_prob_tensor = fc1_prob_tensor.to(torch.float32) # 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. From 5c3553def3084462ff855ad5eaedd92ceb7a5b85 Mon Sep 17 00:00:00 2001 From: vthumbe1503 Date: Fri, 13 Nov 2026 21:14:29 -0800 Subject: [PATCH 8/9] revert Signed-off-by: vthumbe1503 --- transformer_engine/pytorch/ops/fused/grouped_mlp.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index b07c42562b3..1f31ca12440 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -1590,10 +1590,10 @@ def fuser_forward( fc1_prob_tensor = None if not unit_activation_scale: fc1_prob_tensor = scales.detach().reshape((-1,) if canonical else (-1, 1, 1)) - - if use_nvfp4: # cuDNN NVFP4 kernel only accepts fp32 scales. fc1_prob_tensor = fc1_prob_tensor.to(torch.float32) + + if use_nvfp4: # 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. From 4e11b5e2caf3f5cc7193c6cefeb860738202f2af Mon Sep 17 00:00:00 2001 From: vthumbe1503 Date: Fri, 13 Nov 2026 21:17:08 -0800 Subject: [PATCH 9/9] Revert properly Ensure fc1_prob_tensor is converted to fp32 only when using NVFP4. Signed-off-by: vthumbe1503 --- transformer_engine/pytorch/ops/fused/grouped_mlp.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index 1f31ca12440..05ad6d178c9 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -1590,8 +1590,9 @@ def fuser_forward( fc1_prob_tensor = None if not unit_activation_scale: fc1_prob_tensor = scales.detach().reshape((-1,) if canonical else (-1, 1, 1)) - # cuDNN NVFP4 kernel only accepts fp32 scales. - fc1_prob_tensor = fc1_prob_tensor.to(torch.float32) + if use_nvfp4: + # cuDNN NVFP4 kernel only accepts fp32 scales. + fc1_prob_tensor = fc1_prob_tensor.to(torch.float32) if use_nvfp4: # cuDNN receives NVFP4 block-scaled inputs without TE's per-group