From 6611d356d9836fa26da21ba1f963e2befe8efe47 Mon Sep 17 00:00:00 2001 From: jiagaoxiang Date: Thu, 3 Sep 2026 16:43:53 +0000 Subject: [PATCH 1/3] test: cover fused clamped SwiGLU Triton autograd Closes PRPUNDIT-23. One GPU-gated suite for unweighted and weighted ClampedSwiGLUFunction paths; not the eager test_clamped_swiglu.py module. --- .../deepseek_v4/test_fused_clamped_swiglu.py | 121 ++++++++++++++++++ 1 file changed, 121 insertions(+) create mode 100644 tests/unit_tests/megatron/transformer/deepseek_v4/test_fused_clamped_swiglu.py diff --git a/tests/unit_tests/megatron/transformer/deepseek_v4/test_fused_clamped_swiglu.py b/tests/unit_tests/megatron/transformer/deepseek_v4/test_fused_clamped_swiglu.py new file mode 100644 index 000000000..2dc14147e --- /dev/null +++ b/tests/unit_tests/megatron/transformer/deepseek_v4/test_fused_clamped_swiglu.py @@ -0,0 +1,121 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Parameterized Triton clamped-SwiGLU autograd tests (PRPUNDIT-23). + +``test_clamped_swiglu.py`` covers the eager ``clamped_swiglu.py`` module. +This file covers the production Triton wrappers in ``fused_bias_swiglu.py``: +unweighted ``ClampedSwiGLUFunction`` (shared-expert MLP) and weighted +``ClampedWeightedSwiGLUFunction`` (grouped MLP). Linxwang consolidated both +into one ticket; the weighted path is dormant in shipped DeepSeek-V4 config +but stays in this shared suite. +""" + +from __future__ import annotations + +import pytest +import torch +import torch.nn.functional as F + +cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="clamped SwiGLU Triton kernels need CUDA/HIP") + + +def _eager_clamped_swiglu(y: torch.Tensor, alpha: float) -> torch.Tensor: + half = y.shape[-1] // 2 + gate, up = y[..., :half], y[..., half:] + return F.silu(torch.clamp(gate, max=alpha)) * torch.clamp(up, min=-alpha, max=alpha) + + +def _eager_clamped_weighted_swiglu(y: torch.Tensor, weights: torch.Tensor, alpha: float) -> torch.Tensor: + return _eager_clamped_swiglu(y, alpha) * weights.unsqueeze(-1) + + +@cuda +class TestClampedSwiGLUFunction: + def setup_method(self): + pytest.importorskip("triton") + from primus.backends.megatron.core.fusions.fused_bias_swiglu import ( + ClampedSwiGLUFunction, + swiglu_impl, + ) + + self.Fn = ClampedSwiGLUFunction + self.swiglu_impl = swiglu_impl + + @pytest.mark.parametrize("alpha", [7.0, 1.0, 0.5]) + @pytest.mark.parametrize("fp8_input_store", [False, True]) + def test_forward_matches_eager_reference(self, alpha: float, fp8_input_store: bool): + torch.manual_seed(2024) + y = (torch.randn(17, 64, device="cuda", dtype=torch.float32) * 5.0).requires_grad_(True) + out = self.Fn.apply(y, fp8_input_store, alpha) + torch.testing.assert_close(out, _eager_clamped_swiglu(y, alpha), atol=1e-5, rtol=1e-5) + + def test_backward_matches_eager_autograd(self): + torch.manual_seed(6) + alpha = 7.0 + y_fn = (torch.randn(5, 20, device="cuda", dtype=torch.float32) * 5.0).requires_grad_(True) + y_eager = y_fn.detach().clone().requires_grad_(True) + grad_out = torch.randn(5, 10, device="cuda", dtype=torch.float32) + + self.Fn.apply(y_fn, False, alpha).backward(grad_out) + _eager_clamped_swiglu(y_eager, alpha).backward(grad_out) + torch.testing.assert_close(y_fn.grad, y_eager.grad, atol=1e-5, rtol=1e-5) + + def test_swiglu_impl_routes_to_clamped_function(self): + torch.manual_seed(9) + x = torch.randn(3, 5, 16, device="cuda", dtype=torch.float32) + out = self.swiglu_impl(x, None, fp8_input_store=False, clamp_value=7.0) + ref = _eager_clamped_swiglu(x.view(-1, 16), 7.0).view(3, 5, 8) + torch.testing.assert_close(out, ref, atol=1e-5, rtol=1e-5) + + +@cuda +class TestClampedWeightedSwiGLUFunction: + def setup_method(self): + pytest.importorskip("triton") + from primus.backends.megatron.core.fusions.fused_bias_swiglu import ( + ClampedWeightedSwiGLUFunction, + weighted_bias_swiglu_impl, + ) + + self.Fn = ClampedWeightedSwiGLUFunction + self.impl = weighted_bias_swiglu_impl + + @pytest.mark.parametrize("alpha", [7.0, 1.0]) + @pytest.mark.parametrize("fp8_input_store", [False, True]) + def test_forward_matches_eager_reference(self, alpha: float, fp8_input_store: bool): + torch.manual_seed(2024) + M, half = 5, 16 + y = (torch.randn(M, 2 * half, dtype=torch.float32, device="cuda") * 5.0).requires_grad_() + weights = torch.rand(M, dtype=torch.float32, device="cuda") + 0.1 + out = self.Fn.apply(y, weights, fp8_input_store, alpha) + torch.testing.assert_close( + out, _eager_clamped_weighted_swiglu(y, weights, alpha), atol=1e-5, rtol=1e-5 + ) + saved_input, saved_weights = out.grad_fn.saved_tensors + assert saved_input.dtype == (torch.float8_e4m3fn if fp8_input_store else y.dtype) + assert saved_weights.dtype == weights.dtype + + def test_backward_matches_eager_autograd(self): + torch.manual_seed(8) + alpha = 7.0 + y_fn = (torch.randn(6, 16, device="cuda", dtype=torch.float32) * 4.0).requires_grad_(True) + w_fn = (torch.rand(6, device="cuda", dtype=torch.float32) + 0.1).requires_grad_(True) + y_e = y_fn.detach().clone().requires_grad_(True) + w_e = w_fn.detach().clone().requires_grad_(True) + grad_out = torch.randn(6, 8, device="cuda", dtype=torch.float32) + + self.Fn.apply(y_fn, w_fn, False, alpha).backward(grad_out) + _eager_clamped_weighted_swiglu(y_e, w_e, alpha).backward(grad_out) + torch.testing.assert_close(y_fn.grad, y_e.grad, atol=1e-5, rtol=1e-5) + torch.testing.assert_close(w_fn.grad, w_e.grad, atol=1e-5, rtol=1e-5) + + def test_impl_routes_to_clamped_weighted_function(self): + torch.manual_seed(3) + x = torch.randn(4, 16, device="cuda", dtype=torch.float32) + weights = torch.rand(4, device="cuda", dtype=torch.float32) + 0.1 + out = self.impl(x, None, weights, fp8_input_store=False, clamp_value=7.0) + torch.testing.assert_close(out, _eager_clamped_weighted_swiglu(x, weights, 7.0), atol=1e-5, rtol=1e-5) From 400980cf8af78d764088240917f5d808494b8878 Mon Sep 17 00:00:00 2001 From: Doug J Date: Thu, 3 Sep 2026 11:08:15 -0700 Subject: [PATCH 2/3] test: gate fp8_input_store=True on a working float8_e4m3fn cast Copilot review noted the fp8_input_store=True parametrized cases can hard-fail on GPUs/torch builds where float8 casting is unsupported, even when CUDA is available. Mirror the FP8 cast probe used in test_v4_fp8_indexer.py, but scope the skip to just the True case so fp8_input_store=False coverage still runs on non-FP8 devices. --- .../deepseek_v4/test_fused_clamped_swiglu.py | 28 +++++++++++++++++-- 1 file changed, 26 insertions(+), 2 deletions(-) diff --git a/tests/unit_tests/megatron/transformer/deepseek_v4/test_fused_clamped_swiglu.py b/tests/unit_tests/megatron/transformer/deepseek_v4/test_fused_clamped_swiglu.py index 2dc14147e..bdb53bf90 100644 --- a/tests/unit_tests/megatron/transformer/deepseek_v4/test_fused_clamped_swiglu.py +++ b/tests/unit_tests/megatron/transformer/deepseek_v4/test_fused_clamped_swiglu.py @@ -23,6 +23,30 @@ cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="clamped SwiGLU Triton kernels need CUDA/HIP") +def _fp8_cast_works() -> bool: + if not hasattr(torch, "float8_e4m3fn"): + return False + try: + torch.zeros(4).to(torch.float8_e4m3fn).to(torch.float32) + return True + except (RuntimeError, TypeError): + return False + + +# fp8_input_store=True needs a working float8_e4m3fn cast; gate just that +# parameter (mirrors test_v4_fp8_indexer.py) so fp8_input_store=False still +# runs on GPUs/torch builds without FP8 support instead of hard-failing. +_fp8_input_store_cases = [ + False, + pytest.param( + True, + marks=pytest.mark.skipif( + not _fp8_cast_works(), reason="torch.float8_e4m3fn cast unsupported on this build/device" + ), + ), +] + + def _eager_clamped_swiglu(y: torch.Tensor, alpha: float) -> torch.Tensor: half = y.shape[-1] // 2 gate, up = y[..., :half], y[..., half:] @@ -46,7 +70,7 @@ def setup_method(self): self.swiglu_impl = swiglu_impl @pytest.mark.parametrize("alpha", [7.0, 1.0, 0.5]) - @pytest.mark.parametrize("fp8_input_store", [False, True]) + @pytest.mark.parametrize("fp8_input_store", _fp8_input_store_cases) def test_forward_matches_eager_reference(self, alpha: float, fp8_input_store: bool): torch.manual_seed(2024) y = (torch.randn(17, 64, device="cuda", dtype=torch.float32) * 5.0).requires_grad_(True) @@ -85,7 +109,7 @@ def setup_method(self): self.impl = weighted_bias_swiglu_impl @pytest.mark.parametrize("alpha", [7.0, 1.0]) - @pytest.mark.parametrize("fp8_input_store", [False, True]) + @pytest.mark.parametrize("fp8_input_store", _fp8_input_store_cases) def test_forward_matches_eager_reference(self, alpha: float, fp8_input_store: bool): torch.manual_seed(2024) M, half = 5, 16 From e448f2e8c89f0d28dfe673f465e489a076823d1e Mon Sep 17 00:00:00 2001 From: Doug J Date: Thu, 3 Sep 2026 11:17:58 -0700 Subject: [PATCH 3/3] test: probe fp8 cast support on CUDA, not CPU Copilot flagged that _fp8_cast_works() probed float8_e4m3fn casting on a CPU tensor even though every parametrized case in this suite runs on device="cuda". On builds where FP8 casting is CUDA-only, the CPU probe would incorrectly report unsupported and skip valid fp8_input_store=True coverage. Probe on the CUDA device when available instead. --- .../transformer/deepseek_v4/test_fused_clamped_swiglu.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/tests/unit_tests/megatron/transformer/deepseek_v4/test_fused_clamped_swiglu.py b/tests/unit_tests/megatron/transformer/deepseek_v4/test_fused_clamped_swiglu.py index bdb53bf90..07575e63d 100644 --- a/tests/unit_tests/megatron/transformer/deepseek_v4/test_fused_clamped_swiglu.py +++ b/tests/unit_tests/megatron/transformer/deepseek_v4/test_fused_clamped_swiglu.py @@ -26,8 +26,14 @@ def _fp8_cast_works() -> bool: if not hasattr(torch, "float8_e4m3fn"): return False + # This suite only ever runs its parametrized cases on CUDA/HIP (the + # module is skipped entirely otherwise via the `cuda` mark below), so + # probe the cast on the device the kernels actually execute on. A + # CPU-only probe can incorrectly skip fp8_input_store=True on builds + # where FP8 casting is CUDA-only. + device = "cuda" if torch.cuda.is_available() else "cpu" try: - torch.zeros(4).to(torch.float8_e4m3fn).to(torch.float32) + torch.zeros(4, device=device).to(torch.float8_e4m3fn).to(torch.float32) return True except (RuntimeError, TypeError): return False