[ops] fix: honor sigmoid activation in npu gated rmsnorm

.github/workflows/npu_unit_tests.yml

@@ -232,6 +232,9 @@ jobs:
           uv run --frozen pytest -v -s -x tests/ops/test_comp.py
           uv run --frozen pytest -v -s -x tests/ops/test_moe_hw_gate.py
           uv run --frozen pytest -v -s -x tests/ops/test_npu_kernels.py
+          # A missing native runtime/device must fail Ascend CI, not skip the parity gate.
+          uv run --frozen python -c "import torch; import torch_npu; assert torch.npu.is_available(), 'Ascend parity gate requires a live NPU'"
+          VEOMNI_REQUIRE_NPU=1 uv run --frozen pytest -v -s -x tests/ops/test_npu_rms_norm_gated.py
           # This numerical gate must execute on NPU; fail instead of silently skipping
           # when the vendor extension is missing from the CI image.
           uv run --frozen python -c "import cann_ops_transformer"

tests/ops/test_npu_kernels.py

@@ -328,8 +328,8 @@ def test_partial_matches_eager_bf16(self, B, H, S, D, rotary_dim):
 class TestNPURmsNormGated:
     """Tests for the ``rms_norm_gated`` NPU kernel (NPUFusedRMSNormGated)."""
 
-    @pytest.mark.parametrize("batch,seq,hidden,ffn_dim", [(2, 16, 128, 256), (1, 8, 64, 128)])
-    def test_matches_eager_bf16(self, batch, seq, hidden, ffn_dim):
+    @pytest.mark.parametrize("batch,seq,hidden", [(2, 16, 128), (1, 8, 64)])
+    def test_matches_eager_bf16(self, batch, seq, hidden):
         slot = OpSlot("rms_norm_gated", "standard")
         slot.bind("npu")
         # The bound kernel is the NPUFusedRMSNormGated class; instantiate it.
@@ -338,7 +338,7 @@ def test_matches_eager_bf16(self, batch, seq, hidden, ffn_dim):
         fused_module = fused_cls(hidden_size=hidden, eps=1e-6).to(device=DEVICE, dtype=torch.bfloat16)
 
         hidden_states = torch.randn(batch, seq, hidden, device=DEVICE, dtype=torch.bfloat16)
-        gate = torch.randn(batch, seq, ffn_dim, device=DEVICE, dtype=torch.bfloat16)
+        gate = torch.randn_like(hidden_states)
 
         out_fused = fused_module(hidden_states, gate=gate)
         out_eager = _eager_rms_norm_gated(hidden_states, fused_module.weight, fused_module.variance_epsilon, gate)

tests/ops/test_npu_kernels_extended.py

@@ -267,34 +267,27 @@ def test_partial_pass_through_preserved(self):
 class TestNPURmsNormGatedExtended:
     """Coverage for the NPU FusedRMSNormGated module beyond PR #818."""
 
-    @pytest.mark.parametrize("ffn_dim_factor", [2, 4])
-    def test_gating_output_shape_and_dtype(self, ffn_dim_factor):
-        """Output shape is (B, S, (ffn_dim + hidden) // 2).
-
-        NPUFusedRMSNormGated concatenates [gate, rms_norm(hidden)]
-        along the last dim (so the concat length is ffn_dim +
-        hidden), then npu_swiglu splits in half via SiLU(x[:h])
-        * x[h:].  The result is shape (B, S, (ffn_dim + hidden) / 2).
-        """
+    @pytest.mark.parametrize("hidden", [32, 128])
+    def test_gating_output_shape_and_dtype(self, hidden):
+        """Elementwise gating preserves the hidden-state shape and dtype."""
         from veomni.ops.kernels.gated_delta_rule.npu_rms_norm_gated import NPUFusedRMSNormGated
 
-        batch, seq, hidden, ffn_dim = 1, 4, 32, 32 * ffn_dim_factor
-        expected_dim = (ffn_dim + hidden) // 2
+        batch, seq = 1, 4
         fused_module = NPUFusedRMSNormGated(hidden_size=hidden, eps=1e-6).to(device=DEVICE, dtype=torch.bfloat16)
         hidden_states = torch.randn(batch, seq, hidden, device=DEVICE, dtype=torch.bfloat16)
-        gate = torch.randn(batch, seq, ffn_dim, device=DEVICE, dtype=torch.bfloat16)
+        gate = torch.randn(batch, seq, hidden, device=DEVICE, dtype=torch.bfloat16)
         out = fused_module(hidden_states, gate=gate)
-        assert out.shape == (batch, seq, expected_dim)
+        assert out.shape == (batch, seq, hidden)
         assert out.dtype == torch.bfloat16
 
     def test_gating_zero_gate_zeros_output(self):
         """When gate is zero, the output is 0 * normed = 0."""
         from veomni.ops.kernels.gated_delta_rule.npu_rms_norm_gated import NPUFusedRMSNormGated
 
-        batch, seq, hidden, ffn_dim = 1, 4, 32, 64
+        batch, seq, hidden = 1, 4, 32
         fused_module = NPUFusedRMSNormGated(hidden_size=hidden, eps=1e-6).to(device=DEVICE, dtype=torch.bfloat16)
         hidden_states = torch.randn(batch, seq, hidden, device=DEVICE, dtype=torch.bfloat16)
-        gate = torch.zeros(batch, seq, ffn_dim, device=DEVICE, dtype=torch.bfloat16)
+        gate = torch.zeros(batch, seq, hidden, device=DEVICE, dtype=torch.bfloat16)
         out = fused_module(hidden_states, gate=gate)
         # SiLU(0) * normed -> 0, so the output is bit-exact zero.
         assert out.abs().max().item() == 0.0, f"Output should be exactly zero, got max abs {out.abs().max().item()}"
@@ -303,10 +296,10 @@ def test_gating_zero_gate_zeros_output(self):
     def test_gating_eps_parameter_used(self, eps):
         from veomni.ops.kernels.gated_delta_rule.npu_rms_norm_gated import NPUFusedRMSNormGated
 
-        batch, seq, hidden, ffn_dim = 1, 4, 32, 64
+        batch, seq, hidden = 1, 4, 32
         fused_module = NPUFusedRMSNormGated(hidden_size=hidden, eps=eps).to(device=DEVICE, dtype=torch.bfloat16)
         hidden_states = torch.randn(batch, seq, hidden, device=DEVICE, dtype=torch.bfloat16)
-        gate = torch.randn(batch, seq, ffn_dim, device=DEVICE, dtype=torch.bfloat16)
+        gate = torch.randn(batch, seq, hidden, device=DEVICE, dtype=torch.bfloat16)
         out_fused = fused_module(hidden_states, gate=gate)
         out_eager = _eager_rms_norm_gated(hidden_states, fused_module.weight, eps, gate)
         # Compound op (RMSNorm + concat + SiLU gate) with bf16 rounding.

tests/ops/test_npu_rms_norm_gated.py

@@ -0,0 +1,133 @@
+"""CPU contract checks and native Ascend forward/backward parity."""
+
+import importlib.util
+import os
+import sys
+import types
+from pathlib import Path
+from unittest.mock import patch
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+
+def load_wrapper():
+    def rms(x, weight, eps):
+        normalized = x.float() * torch.rsqrt(x.float().square().mean(-1, keepdim=True) + eps)
+        return (normalized.to(x.dtype) * weight,)
+
+    def swiglu(x, dim):
+        gate, value = x.chunk(2, dim=dim)
+        return F.silu(gate) * value
+
+    fake = types.SimpleNamespace(npu_rms_norm=rms, npu_swiglu=swiglu)
+    path = Path(__file__).parents[2] / "veomni/ops/kernels/gated_delta_rule/npu_rms_norm_gated.py"
+    spec = importlib.util.spec_from_file_location("_gated_norm_test", path)
+    module = importlib.util.module_from_spec(spec)
+    with patch.dict(sys.modules, {"torch_npu": fake}):
+        spec.loader.exec_module(module)
+    return module.NPUFusedRMSNormGated
+
+
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32])
+def test_sigmoid_reference_and_gradients(dtype):
+    torch.manual_seed(73)
+    module = load_wrapper()(128, dtype=dtype, activation="sigmoid")
+    with torch.no_grad():
+        module.weight.copy_(torch.linspace(0.25, 1.75, 128).to(dtype))
+    x = torch.randn(7, 128, dtype=dtype).requires_grad_()
+    z = torch.linspace(-20, 20, 896).reshape(7, 128).to(dtype).requires_grad_()
+    r = x.float() * torch.rsqrt(x.float().square().mean(-1, keepdim=True) + module.variance_epsilon)
+    expected = (module.weight * r.to(dtype) * z.float().sigmoid()).to(dtype)
+    actual = module(x, z)
+    torch.testing.assert_close(actual, expected, atol=0, rtol=0)
+    dy = torch.randn_like(actual)
+    got = torch.autograd.grad(actual, (x, z, module.weight), dy)
+    ref = torch.autograd.grad(expected, (x, z, module.weight), dy)
+    for a, b in zip(got, ref, strict=True):
+        torch.testing.assert_close(a, b, atol=0, rtol=0)
+
+
+def test_zero_gate_is_half_normalized_and_silu_unchanged():
+    cls = load_wrapper()
+    x = torch.ones(2, 128)
+    z = torch.zeros_like(x)
+    torch.testing.assert_close(cls(128, activation="sigmoid")(x, z), x * 0.5 / (1 + 1e-6) ** 0.5)
+    torch.testing.assert_close(cls(128)(x, z), torch.zeros_like(x))
+    with pytest.raises(ValueError, match="activation"):
+        cls(128, activation="unsupported_gate")
+    with pytest.raises(ValueError, match="gate"):
+        cls(128, activation="sigmoid")(x)
+
+
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32])
+@pytest.mark.parametrize("activation", ["sigmoid", "silu"])
+@pytest.mark.parametrize("autocast_enabled", [False, True])
+def test_native_npu_reference_and_gradients(dtype, activation, autocast_enabled):
+    # Load the real wrapper without the CPU tests' torch_npu substitution.
+    if os.environ.get("VEOMNI_REQUIRE_NPU") == "1":
+        import torch_npu  # noqa: F401
+
+        assert torch.npu.is_available(), "Ascend parity gate requires a live NPU"
+    else:
+        pytest.importorskip("torch_npu")
+    if not torch.npu.is_available():
+        pytest.skip("Native Ascend device is required")
+    path = Path(__file__).parents[2] / "veomni/ops/kernels/gated_delta_rule/npu_rms_norm_gated.py"
+    spec = importlib.util.spec_from_file_location("_native_gated_norm_test", path)
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    torch.manual_seed(73)
+    norm = module.NPUFusedRMSNormGated(128, device="npu", dtype=dtype, activation=activation)
+    with torch.no_grad():
+        norm.weight.copy_(torch.linspace(0.25, 1.75, 128).to(device="npu", dtype=dtype))
+    x = torch.randn(2, 7, 128, device="npu", dtype=dtype).requires_grad_()
+    z = torch.linspace(-20, 20, x.numel()).reshape(x.shape).to(device="npu", dtype=dtype).requires_grad_()
+    rx, rz, rw = (tensor.detach().clone().requires_grad_() for tensor in (x, z, norm.weight))
+    # An independent torch reference preserves the architecture's cast boundary.
+    with torch.autocast("npu", enabled=False):
+        normalized = rx.float() * torch.rsqrt(rx.float().square().mean(-1, keepdim=True) + norm.variance_epsilon)
+        if activation == "sigmoid":
+            expected = (rw * normalized.to(dtype) * rz.float().sigmoid()).to(dtype)
+        else:
+            expected = F.silu(rz) * (normalized * rw.float()).to(dtype)
+    with torch.autocast("npu", dtype=torch.bfloat16, enabled=autocast_enabled):
+        actual = norm(x, z)
+    assert actual.dtype == dtype
+    dy = torch.randn_like(actual)
+    got = torch.autograd.grad(actual, (x, z, norm.weight), dy)
+    ref = torch.autograd.grad(expected, (rx, rz, rw), dy)
+    # Fused CANN reductions need tolerance; BF16 roundoff is not bitwise FP32.
+    rtol, atol = (3e-2, 2e-2) if dtype == torch.bfloat16 else (3e-5, 3e-5)
+    for a, b in zip((actual, *got), (expected, *ref), strict=True):
+        assert torch.isfinite(a).all()
+        torch.testing.assert_close(a, b, atol=atol, rtol=rtol)
+
+
+@pytest.mark.parametrize("activation", ["sigmoid", "silu"])
+def test_gate_must_match_input_shape(activation):
+    norm = load_wrapper()(32, activation=activation)
+    x = torch.ones(1, 4, 32)
+    for gate in [None, torch.zeros(1, 4, 64), torch.zeros(1, 1, 32)]:
+        with pytest.raises(ValueError, match="gate with the input shape"):
+            norm(x, gate)
+
+
+@pytest.mark.parametrize("activation", ["relu", "gelu", "tanh", "swish"])
+def test_upstream_activation_fallback_is_preserved(activation):
+    from transformers.activations import ACT2FN
+
+    torch.manual_seed(94)
+    norm = load_wrapper()(32, activation=activation)
+    x = torch.randn(2, 32, requires_grad=True)
+    z = torch.randn_like(x, requires_grad=True)
+    normalized = x * torch.rsqrt(x.square().mean(-1, keepdim=True) + norm.variance_epsilon)
+    expected = normalized * norm.weight * ACT2FN[activation](z)
+    actual = norm(x, z)
+    torch.testing.assert_close(actual, expected)
+    for got, ref in zip(
+        torch.autograd.grad(actual.sum(), (x, z, norm.weight)),
+        torch.autograd.grad(expected.sum(), (x, z, norm.weight)),
+    ):
+        torch.testing.assert_close(got, ref)

veomni/ops/kernels/gated_delta_rule/npu_rms_norm_gated.py

@@ -24,29 +24,30 @@
 from transformers.activations import ACT2FN
 
 
-# ``npu_swiglu`` hard-codes a SiLU gate, so it can only stand in for the gate
-# activation when the model asks for SiLU. Qwen3.5 does (``hidden_act``), but
-# Qwen4-Exp sets ``output_gate_type="sigmoid"`` — fusing there would silently
-# compute the wrong thing.
-_FUSED_GATE_ACTIVATIONS = frozenset({"silu", "swish"})
-
-
 class NPUFusedRMSNormGated(nn.Module):
-    """Ascend NPU RMSNorm with an activation-gated output projection."""
-
-    def __init__(self, hidden_size, eps=1e-6, activation: str = "silu", **kwargs):
-        """Initialize the norm weight and gate activation."""
+    def __init__(self, hidden_size, eps=1e-6, activation="silu", device=None, dtype=None, **kwargs):
         super().__init__()
-        self.weight = nn.Parameter(torch.ones(hidden_size))
+        if activation not in ACT2FN:
+            raise ValueError(f"Unsupported NPU gated RMSNorm activation: {activation!r}")
+        self.weight = nn.Parameter(torch.ones(hidden_size, device=device, dtype=dtype))
         self.variance_epsilon = eps
         self.activation = activation
 
     def forward(self, hidden_states, gate=None):
-        """Apply RMSNorm to ``hidden_states`` and multiply by the activated ``gate``."""
-        normed = torch_npu.npu_rms_norm(hidden_states, self.weight, self.variance_epsilon)[0]
-        if self.activation in _FUSED_GATE_ACTIVATIONS:
-            return torch_npu.npu_swiglu(torch.cat([gate, normed], dim=-1), dim=-1)
-
-        # Still fuses the norm; only the gate falls back to the eager formula
-        # (float32 activation, then cast back) that the reference modules use.
-        return (normed * ACT2FN[self.activation](gate.to(torch.float32))).to(hidden_states.dtype)
+        if gate is None or gate.shape != hidden_states.shape:
+            raise ValueError("NPU gated RMSNorm requires a gate with the input shape")
+        if self.activation == "sigmoid":
+            # Match Qwen4-Exp's rounding boundary: normalize, cast to the input
+            # dtype, then apply the learned weight. Fusing the learned weight
+            # into RMSNorm would move that boundary for BF16 compute weights.
+            unit_weight = torch.ones(hidden_states.shape[-1], device=hidden_states.device, dtype=hidden_states.dtype)
+            normalized = torch_npu.npu_rms_norm(hidden_states, unit_weight, self.variance_epsilon)[0]
+            weighted = self.weight * normalized.to(hidden_states.dtype)
+            return (weighted * gate.float().sigmoid()).to(hidden_states.dtype)
+        # Preserve the existing SiLU route for models such as Qwen3.6.
+        normalized = torch_npu.npu_rms_norm(hidden_states, self.weight, self.variance_epsilon)[0]
+        if self.activation in {"silu", "swish"}:
+            return torch_npu.npu_swiglu(torch.cat([gate, normalized], dim=-1), dim=-1)
+        # Preserve upstream's generic activation fallback without changing the
+        # Qwen4 sigmoid cast boundary or the established SiLU fused path.
+        return (normalized * ACT2FN[self.activation](gate.float())).to(hidden_states.dtype)