Fix audio encoder SkipSimplifiedLayerNormalization CUDA crash

ORT fuses Add(output_proj.bias) + RMSNormalization into
SkipSimplifiedLayerNormalization, placing the 1D bias as the 'skip'
input. ORT's CUDA kernel rejects 1D skip (requires 2D+), while the
CPU kernel accepts it.

This fusion was enabled by PR #253 which changed _Gemma4ScaleFreeRMSNorm
from manual primitive ops to op.RMSNormalization(stash_type=1). The
RMSNormalization op is recognized by ORT's SkipLayerNorm fusion pattern.

Fix: inline manual RMSNorm ops in _Gemma4AudioEncoderModel.forward()
for the pre_projection_norm, preventing ORT from recognizing the fusion
pattern. This preserves the FP32 accumulation for numerical stability.

Tested: audio encoder runs on CUDA with correct output (standalone ORT
and GenAI at 117 tok/s).

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

Author: justinchuby

src/mobius/components/_gemma4_audio.py

@@ -784,10 +784,12 @@ def __init__(
             ]
         )
         # HF uses nn.Linear(..., bias=True) for the output projection.
-        # The bias would normally cause ORT to fuse Add(1D bias) + LayerNorm into
-        # SkipSimplifiedLayerNorm (with 1D skip, which ORT rejects).  This is avoided
-        # because _Gemma4ScaleFreeRMSNorm uses manual primitive ops rather than
-        # op.RMSNormalization, preventing ORT from recognizing the fusion pattern.
+        # ORT fuses Add(1D bias) + RMSNormalization into
+        # SkipSimplifiedLayerNormalization, placing the 1D bias in the
+        # "skip" input position. The CUDA kernel rejects 1D skip.
+        # Keep bias=True here; the caller's pre_projection_norm must use
+        # manual primitive ops (not op.RMSNormalization) to prevent this
+        # fusion pattern.
         self.output_proj = Linear(hidden_size, output_proj_dims, bias=True)
 
     def forward(

src/mobius/models/gemma4.py

@@ -1970,7 +1970,10 @@ def __init__(self, config: Gemma4Config):
         )
         # Scale-free RMSNorm applied before the projection (HF embed_audio.embedding_pre_projection_norm).
         # with_scale=False in HF → no learnable weight → no checkpoint key, no ONNX initializer.
-        self.pre_projection_norm = _Gemma4ScaleFreeRMSNorm(output_proj_dims, eps=rms_norm_eps)
+        # NOTE: We inline the RMSNorm in forward() using manual ops to prevent
+        # ORT from fusing Add(output_proj.bias) + RMSNormalization into
+        # SkipSimplifiedLayerNormalization (CUDA rejects 1D skip).
+        self._rms_norm_eps = rms_norm_eps
         # Learned projection from encoder output space → text hidden size.
         # Corresponds to HF's embed_audio.embedding_projection (no bias).
         self.projector = Linear(output_proj_dims, config.hidden_size, bias=False)
@@ -1985,8 +1988,17 @@ def forward(
         audio_features, downsampled_mask = self.encoder(
             op, input_features, input_features_mask=input_features_mask
         )
-        # Scale-free RMSNorm before projection (HF embed_audio.embedding_pre_projection_norm)
-        audio_features = self.pre_projection_norm(op, audio_features)
+        # Scale-free RMSNorm before projection (HF embed_audio.embedding_pre_projection_norm).
+        # Use manual primitive ops instead of op.RMSNormalization to prevent
+        # ORT from fusing Add(output_proj.bias) + RMSNorm into
+        # SkipSimplifiedLayerNormalization with a 1D bias as skip input
+        # (CUDA kernel rejects 1D skip, CPU kernel accepts it).
+        x_f32 = op.Cast(audio_features, to=ir.DataType.FLOAT)
+        sq = op.Mul(x_f32, x_f32)
+        mean_sq = op.ReduceMean(sq, op.Constant(value_ints=[-1]), keepdims=1)
+        eps = op.Constant(value_float=self._rms_norm_eps)
+        rms = op.Sqrt(op.Add(mean_sq, eps))
+        audio_features = op.CastLike(op.Div(x_f32, rms), audio_features)
         # → projector → [B, T//4, text_hidden_size]
         return self.projector(op, audio_features), downsampled_mask