Fix Gemma4 scale-free V norm FP16 overflow on CUDA (#253)

Gemma4's parameterless V normalization (`v / sqrt(mean(v²) + ε)`)
squared FP16 values directly. V projection outputs reach ~888, and 888²
overflows FP16 max (65504) → inf → 0. This caused all-zero V outputs on
CUDA (CPU uses FP32 internally).

**Fix**: Cast to FP32 before squaring, compute full RMSNorm in FP32,
CastLike back.

**Result**: F16 CUDA 151.5 tok/s on H200 (was NaN).

---------

Signed-off-by: Justin Chu <justinchu@microsoft.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Author: justinchuby

src/mobius/models/gemma4.py

@@ -100,27 +100,29 @@ class _Gemma4ScaleFreeRMSNorm(nn.Module):
     Used for V norms in the vision encoder, the vision projector pre-norm,
     and the audio pre-projection norm.
 
-    Implemented as manual ``x / sqrt(mean(x²) + ε)`` rather than
-    ``op.RMSNormalization`` to prevent ORT's graph optimizer from fusing
-    an upstream ``Add(bias)`` into ``SkipSimplifiedLayerNormalization``,
-    which requires the skip tensor to have the same shape as the input
-    (failing when the bias is 1D or has mismatched temporal dimension).
+    Uses ``op.RMSNormalization`` with ``stash_type=1`` (float32 accumulation)
+    to handle FP16 overflow: values > 256 squared exceed the FP16 max (65504).
     """
 
     def __init__(self, dim: int, eps: float = 1e-6):
         super().__init__()
         self.dim = dim
         self.eps = eps
+        # Constant all-ones scale (not a learnable parameter from HF).
+        self.weight = nn.Parameter([dim], data=ir.Tensor(np.ones(dim, dtype=np.float32)))
 
     def forward(self, op: builder.OpBuilder, hidden_states: ir.Value) -> ir.Value:
-        # Manual RMSNorm: x / sqrt(mean(x²) + ε), scale = 1.0 (scale-free).
-        # Using primitive ops avoids ORT's SkipLayerNorm fusion pattern which
-        # would corrupt the skip shape when an upstream Add uses a 1D bias.
-        square = op.Mul(hidden_states, hidden_states)
-        mean_sq = op.ReduceMean(square, op.Constant(value_ints=[-1]), keepdims=1)
-        eps = op.CastLike(op.Constant(value_float=self.eps), mean_sq)
-        rms = op.Sqrt(op.Add(mean_sq, eps))
-        return op.Div(hidden_states, rms)
+        # stash_type=1 means accumulate variance in float32, avoiding
+        # FP16 overflow when squaring large values.
+        # CastLike ensures the weight matches the input dtype.
+        scale = op.CastLike(self.weight, hidden_states)
+        return op.RMSNormalization(
+            hidden_states,
+            scale,
+            axis=-1,
+            epsilon=self.eps,
+            stash_type=1,
+        )
 
 
 # ---------------------------------------------------------------------------
@@ -773,16 +775,18 @@ def forward(
                 value_raw = key_raw
             else:
                 value_raw = self.v_proj(op, hidden_states)
-            # Parameterless per-head V normalisation
+            # Parameterless per-head V normalisation (FP32 accumulation to
+            # prevent FP16 overflow when squaring values > 256).
             value_states = op.Reshape(
                 value_raw,
                 op.Constant(value_ints=[0, 0, self.num_key_value_heads, self.head_dim]),
             )
-            sq = op.Mul(value_states, value_states)
+            v_f32 = op.Cast(value_states, to=ir.DataType.FLOAT)
+            sq = op.Mul(v_f32, v_f32)
             mean_sq = op.ReduceMean(sq, [-1], keepdims=1)
             eps = op.Constant(value_floats=[self._v_norm_eps])
-            rms = op.Sqrt(op.Add(mean_sq, op.CastLike(eps, mean_sq)))
-            value_states = op.Div(value_states, rms)
+            rms = op.Sqrt(op.Add(mean_sq, eps))
+            value_states = op.CastLike(op.Div(v_f32, rms), value_states)
             value_states = op.Reshape(value_states, [0, 0, -1])
 
             # Build GQA attributes
@@ -846,19 +850,21 @@ def forward(
                 value_raw = key_raw
             else:
                 value_raw = self.v_proj(op, hidden_states)
-            # Parameterless per-head V normalisation
+            # Parameterless per-head V normalisation (FP32 accumulation to
+            # prevent FP16 overflow when squaring values > 256).
             value_states = op.Reshape(
                 value_raw,
                 op.Constant(value_ints=[0, 0, self.num_key_value_heads, self.head_dim]),
             )
-            sq = op.Mul(value_states, value_states)
+            v_f32 = op.Cast(value_states, to=ir.DataType.FLOAT)
+            sq = op.Mul(v_f32, v_f32)
             mean_sq = op.ReduceMean(sq, [-1], keepdims=1)
             # Use op.Constant to create a 1D tensor node (not a scalar initializer).
             # Scalar Python floats use a type-keyed cache that can fail when upstream
             # type information is missing (e.g., after custom ops like com.microsoft.MoE).
             eps = op.Constant(value_floats=[self._v_norm_eps])
-            rms = op.Sqrt(op.Add(mean_sq, op.CastLike(eps, mean_sq)))
-            value_states = op.Div(value_states, rms)
+            rms = op.Sqrt(op.Add(mean_sq, eps))
+            value_states = op.CastLike(op.Div(v_f32, rms), value_states)
             value_states = op.Reshape(value_states, [0, 0, -1])
 
             attn_output, present_key, present_value = _apply_attention(

src/mobius/models/gemma4_test.py

@@ -5,6 +5,7 @@
 
 from __future__ import annotations
 
+import onnx_ir as ir
 import torch
 
 from mobius._configs import Gemma4Config
@@ -117,3 +118,47 @@ def test_per_expert_scale_not_folded(self):
         key = "decoder.model.layers.0.router.per_expert_scale"
         assert key in result
         assert torch.allclose(result[key], torch.ones(4))
+
+
+class TestScaleFreeRMSNormOverflow:
+    """V norm should handle FP16 overflow from squaring large values."""
+
+    def test_vnorm_fp16_no_nan(self):
+        """Values ~888 overflow FP16 when squared (888²=788K > 65504).
+
+        The scale-free RMSNorm must use stash_type=1 (float32 accumulation)
+        to avoid inf/NaN from the variance computation.
+        """
+        import numpy as np
+
+        from mobius._testing.ort_inference import OnnxModelSession
+        from mobius.models.gemma4 import _Gemma4ScaleFreeRMSNorm
+
+        dim = 64
+        norm = _Gemma4ScaleFreeRMSNorm(dim, eps=1e-6)
+
+        # Build a minimal ONNX graph for the norm
+        from mobius.tasks._base import _make_graph, _make_model
+
+        graph, builder = _make_graph()
+        op = builder.op
+        x = builder.input("x", dtype=ir.DataType.FLOAT16, shape=[1, 4, dim])
+        y = norm(op, x)
+        builder.add_output(y, "y")
+        model = _make_model(graph)
+
+        session = OnnxModelSession(model, device="cpu")
+
+        # Values that overflow FP16 when squared: 888² = 788,544 > 65504
+        test_input = np.full((1, 4, dim), 888.0, dtype=np.float16)
+        result = session.run({"x": test_input})
+        output = result["y"]
+
+        assert not np.any(np.isnan(output)), "V norm produced NaN for input 888"
+        assert not np.any(np.isinf(output)), "V norm produced Inf for input 888"
+        # RMSNorm of a constant vector: x/rms(x) = sign(x) ≈ 1.0
+        np.testing.assert_allclose(
+            output.astype(np.float32),
+            np.ones_like(output, dtype=np.float32),
+            atol=0.01,
+        )