From 76524c2dd7b3c561d62e19ada21ce2b962d155c7 Mon Sep 17 00:00:00 2001
From: Justin Chu <justinchuby@users.noreply.github.com>
Date: Sat, 28 Mar 2026 07:39:48 -0700
Subject: [PATCH 1/5] Fix bfloat16/float type mismatch in OffsetRMSNorm and
 Gemma3nAltUp

Use CastLike to cast the float constant 1.0 to match the weight/coefs
dtype before Add. Without this, op.Add(bf16_tensor, 1.0) creates a
type mismatch since 1.0 becomes a float32 Constant.

Note: bf16 inference still fails on ORT CPU EP (no bf16 kernels) and
on CUDA EP due to ORT Attention op decomposition creating mixed-type
nodes internally. This is an ORT limitation, not a graph construction
issue. f16 works on both CPU and CUDA.

Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
---
 src/mobius/components/_rms_norm.py | 4 +++-
 src/mobius/models/gemma3n.py       | 2 +-
 2 files changed, 4 insertions(+), 2 deletions(-)

diff --git a/src/mobius/components/_rms_norm.py b/src/mobius/components/_rms_norm.py
index 3b6d1bb7..c31d42c3 100644
--- a/src/mobius/components/_rms_norm.py
+++ b/src/mobius/components/_rms_norm.py
@@ -37,7 +37,9 @@ def __init__(self, hidden_size: int, eps: float = 1e-6):
         self.variance_epsilon = eps
 
     def forward(self, op: builder.OpBuilder, hidden_states: ir.Value):
-        effective_weight = op.Add(self.weight, 1.0)
+        effective_weight = op.Add(
+            self.weight, op.CastLike(op.Constant(value_float=1.0), self.weight)
+        )
         return op.RMSNormalization(
             hidden_states,
             effective_weight,
diff --git a/src/mobius/models/gemma3n.py b/src/mobius/models/gemma3n.py
index f7bfa647..407b12f4 100644
--- a/src/mobius/models/gemma3n.py
+++ b/src/mobius/models/gemma3n.py
@@ -256,7 +256,7 @@ def correct(self, op: builder.OpBuilder, predictions: list, activated):
 
         # correction_coefs: [batch, seq, num_inputs] + 1
         all_coefs = self.correction_coefs(op, modalities)
-        all_coefs = op.Add(all_coefs, 1.0)
+        all_coefs = op.Add(all_coefs, op.CastLike(op.Constant(value_float=1.0), all_coefs))
 
         corrected = []
         for i in range(self.altup_num_inputs):

From 4d56ecad98cd860e8412572345ad1601d8d3b700 Mon Sep 17 00:00:00 2001
From: Justin Chu <justinchuby@users.noreply.github.com>
Date: Sat, 28 Mar 2026 07:50:20 -0700
Subject: [PATCH 2/5] Fix CastLike type mismatches: wrap float constants with
 op.CastLike

Ensure float32 literals are cast to match model-precision tensors before
use in arithmetic ops. Affected sites:

- _activations.py: quick_gelu scalar 1.702
- _diffusion.py: AdaLayerNormZero constant 1.0 added to scale
- _moe.py: threshold and neg_inf constants matched to scores dtype (x2)
- gemma3n.py: router_input_scale constant; sqrt(2) divisor in Laurel residual
- _ssm.py: epsilon constant in RMS variance computation
- _audio.py: 0.5 half-scaling constant in ConformerBlock

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
---
 src/mobius/components/_activations.py | 2 +-
 src/mobius/components/_audio.py       | 2 +-
 src/mobius/components/_diffusion.py   | 2 +-
 src/mobius/components/_moe.py         | 6 +++---
 src/mobius/models/gemma3n.py          | 6 ++++--
 5 files changed, 10 insertions(+), 8 deletions(-)

diff --git a/src/mobius/components/_activations.py b/src/mobius/components/_activations.py
index 2773bd0d..753f0e21 100644
--- a/src/mobius/components/_activations.py
+++ b/src/mobius/components/_activations.py
@@ -41,7 +41,7 @@ def relu2(op: builder.OpBuilder, x):
 
 def quick_gelu(op: builder.OpBuilder, x):
     """QuickGELU activation: x * sigmoid(1.702 * x)."""
-    return op.Mul(x, op.Sigmoid(op.Mul(1.702, x)))
+    return op.Mul(x, op.Sigmoid(op.Mul(op.CastLike(op.Constant(value_float=1.702), x), x)))
 
 
 def mish(op: builder.OpBuilder, x):
diff --git a/src/mobius/components/_audio.py b/src/mobius/components/_audio.py
index a509ba97..5871f1af 100644
--- a/src/mobius/components/_audio.py
+++ b/src/mobius/components/_audio.py
@@ -409,7 +409,7 @@ def __init__(
         self.layer_norm = LayerNorm(d_model)
 
     def forward(self, op: builder.OpBuilder, x: ir.Value, relative_attention_bias: ir.Value):
-        half = op.Constant(value_float=0.5)
+        half = op.CastLike(op.Constant(value_float=0.5), x)
 
         # Macaron feed-forward in
         x = op.Add(x, op.Mul(self.feed_forward_in(op, x), half))
diff --git a/src/mobius/components/_diffusion.py b/src/mobius/components/_diffusion.py
index 776fd818..738b5a35 100644
--- a/src/mobius/components/_diffusion.py
+++ b/src/mobius/components/_diffusion.py
@@ -83,7 +83,7 @@ def forward(self, op: builder.OpBuilder, hidden_states: ir.Value, timestep_emb:
         emb = self._silu(op, timestep_emb)
         emb = self.linear(op, emb)
         shift, scale = op.Split(emb, num_outputs=2, axis=-1, _outputs=2)
-        one = op.Constant(value_float=1.0)
+        one = op.CastLike(op.Constant(value_float=1.0), scale)
         hidden_states = self.norm(op, hidden_states)
         hidden_states = op.Mul(hidden_states, op.Add(one, op.Unsqueeze(scale, [1])))
         hidden_states = op.Add(hidden_states, op.Unsqueeze(shift, [1]))
diff --git a/src/mobius/components/_moe.py b/src/mobius/components/_moe.py
index 926e62aa..a1b05a0c 100644
--- a/src/mobius/components/_moe.py
+++ b/src/mobius/components/_moe.py
@@ -145,9 +145,9 @@ def _threshold_mask_and_select(self, op, scores, jitter_eps):
         factor = op.Max(abs_scores, max_score)
         diff = op.Sub(max_score, scores)
         ratio = op.Div(diff, factor)
-        threshold = op.Constant(value_float=2.0 * jitter_eps)
+        threshold = op.CastLike(op.Constant(value_float=2.0 * jitter_eps), scores)
         mask = op.Greater(ratio, threshold)
-        neg_inf = op.Constant(value_float=-1e30)
+        neg_inf = op.CastLike(op.Constant(value_float=-1e30), scores)
         masked_scores = op.Where(mask, neg_inf, scores)
         weights = op.Softmax(masked_scores, axis=-1)
         k_one = op.Constant(value_ints=[1])
@@ -169,7 +169,7 @@ def forward(self, op: builder.OpBuilder, hidden_states: ir.Value):
             )
             all_weights.append(weight_k)
             all_experts.append(expert_k)
-            neg_inf = op.Constant(value_float=-1e30)
+            neg_inf = op.CastLike(op.Constant(value_float=-1e30), current_scores)
             current_scores = op.ScatterElements(
                 current_scores,
                 expert_k,
diff --git a/src/mobius/models/gemma3n.py b/src/mobius/models/gemma3n.py
index 407b12f4..ac1c0973 100644
--- a/src/mobius/models/gemma3n.py
+++ b/src/mobius/models/gemma3n.py
@@ -196,7 +196,7 @@ def __init__(self, config: Gemma3nConfig):
     def _compute_router_modalities(self, op: builder.OpBuilder, x):
         """Compute router modalities: tanh(router(norm(x) * scale))."""
         router_input = self.router_norm(op, x)
-        scale = op.Constant(value_float=self.router_input_scale)
+        scale = op.CastLike(op.Constant(value_float=self.router_input_scale), router_input)
         router_input = op.Mul(router_input, scale)
         routed = self.modality_router(op, router_input)
         return op.Tanh(routed)
@@ -335,7 +335,9 @@ def forward(
         # Residual + Laurel (with sqrt(2) normalization)
         attn_gated = op.Add(active, attn_output)
         attn_laurel = op.Add(attn_gated, laurel_output)
-        attn_laurel = op.Div(attn_laurel, op.Constant(value_float=float(math.sqrt(2))))
+        attn_laurel = op.Div(
+            attn_laurel, op.CastLike(op.Constant(value_float=float(math.sqrt(2))), attn_laurel)
+        )
 
         # MLP
         mlp_input = self.pre_feedforward_layernorm(op, attn_laurel)

From 3a81f00096977b32f5372bb2f445c3740efec3b0 Mon Sep 17 00:00:00 2001
From: Justin Chu <justinchuby@users.noreply.github.com>
Date: Wed, 1 Apr 2026 09:46:47 -0700
Subject: [PATCH 3/5] Add tests verifying CastLike wrapping for non-float32
 dtypes

Tests build OffsetRMSNorm, quick_gelu, and AdaLayerNormOutput with
float16/bfloat16 inputs and assert CastLike ops are present. Verified
tests FAIL without the CastLike fix and PASS with it.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
---
 src/mobius/components/_castlike_dtype_test.py | 67 +++++++++++++++++++
 1 file changed, 67 insertions(+)
 create mode 100644 src/mobius/components/_castlike_dtype_test.py

diff --git a/src/mobius/components/_castlike_dtype_test.py b/src/mobius/components/_castlike_dtype_test.py
new file mode 100644
index 00000000..571cfe87
--- /dev/null
+++ b/src/mobius/components/_castlike_dtype_test.py
@@ -0,0 +1,67 @@
+# Copyright (c) ONNX Project Contributors
+# SPDX-License-Identifier: Apache-2.0
+
+"""Tests that float constants are CastLike-wrapped for non-float32 dtypes.
+
+PR #58 wraps bare float32 constants with op.CastLike() to prevent dtype
+mismatches when model tensors are float16 or bfloat16. These tests build
+component graphs with float16 inputs and verify CastLike ops are present.
+"""
+
+from __future__ import annotations
+
+import onnx_ir as ir
+import pytest
+
+from mobius._testing import count_op_type, create_test_builder, create_test_input
+from mobius.components._activations import quick_gelu
+from mobius.components._diffusion import AdaLayerNormOutput
+from mobius.components._rms_norm import OffsetRMSNorm
+
+
+class TestCastLikeDtypeSafety:
+    """Verify CastLike wrapping of float constants for non-float32 dtypes."""
+
+    @pytest.mark.parametrize("dtype", [ir.DataType.FLOAT16, ir.DataType.BFLOAT16])
+    def test_offset_rms_norm_casts_constant(self, dtype: ir.DataType):
+        """OffsetRMSNorm adds 1.0 to weight — must CastLike for non-f32."""
+        norm = OffsetRMSNorm(hidden_size=64, eps=1e-6)
+        builder_, op, graph = create_test_builder()
+        x = create_test_input(builder_, "x", [2, 3, 64], dtype)
+        result = norm(op, x)
+        graph.outputs.append(result)
+        assert count_op_type(graph, "CastLike") >= 1
+
+    @pytest.mark.parametrize("dtype", [ir.DataType.FLOAT16, ir.DataType.BFLOAT16])
+    def test_quick_gelu_casts_constant(self, dtype: ir.DataType):
+        """quick_gelu uses 1.702 constant — must CastLike for non-f32."""
+        builder_, op, graph = create_test_builder()
+        x = create_test_input(builder_, "x", [2, 3, 64], dtype)
+        result = quick_gelu(op, x)
+        graph.outputs.append(result)
+        assert count_op_type(graph, "CastLike") >= 1
+
+    @pytest.mark.parametrize("dtype", [ir.DataType.FLOAT16, ir.DataType.BFLOAT16])
+    def test_ada_layer_norm_output_casts_constant(self, dtype: ir.DataType):
+        """AdaLayerNormOutput adds 1.0 to scale — must CastLike for non-f32."""
+        mod = AdaLayerNormOutput(hidden_size=64, eps=1e-6)
+        builder_, op, graph = create_test_builder()
+        hidden = create_test_input(builder_, "hidden", [1, 4, 64], dtype)
+        temb = create_test_input(builder_, "temb", [1, 64], dtype)
+        result = mod(op, hidden, temb)
+        graph.outputs.append(result)
+        assert count_op_type(graph, "CastLike") >= 1
+
+    def test_no_castlike_needed_for_float32(self):
+        """With float32 inputs, CastLike is still present but benign.
+
+        This verifies the pattern is always applied regardless of dtype,
+        so there's no conditional logic that could miss a case.
+        """
+        norm = OffsetRMSNorm(hidden_size=64, eps=1e-6)
+        builder_, op, graph = create_test_builder()
+        x = create_test_input(builder_, "x", [2, 3, 64], ir.DataType.FLOAT)
+        result = norm(op, x)
+        graph.outputs.append(result)
+        # CastLike is always emitted (no-op for float32, but present)
+        assert count_op_type(graph, "CastLike") >= 1

From e65c4284896aa9817a42d466bfd2ae29d0f43ba9 Mon Sep 17 00:00:00 2001
From: Justin Chu <justinchuby@users.noreply.github.com>
Date: Wed, 1 Apr 2026 10:43:09 -0700
Subject: [PATCH 4/5] Replace op.Constant scalars with Python literals for
 auto-casting
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

onnxscript auto-casts Python scalars (int/float/bool) to match the dtype
of the other operand in binary ops. op.Constant(value_float=x) returns
ir.Value(FLOAT) which is NOT auto-cast, causing dtype mismatches in
BF16/FP16 builds.

Changes:
- _activations.py: 1.702 → Python literal (op.Mul auto-casts)
- _audio.py: 0.5 → inline literals (removes CastLike + intermediate var)
- _diffusion.py: 1.0 → Python literal (op.Add auto-casts)
- _moe.py: 1e-9, routed_scaling_factor, 2*jitter_eps → Python literals;
  -1e30 uses op.CastLike(-1e30, scores) with Python literal to avoid
  cache-key collision (Expand has no type-variable binding for its input)
- _rms_norm.py: 1.0 → Python literal (op.Add auto-casts)
- gemma3n.py: router_input_scale, 1.0, sqrt(2) → Python literals

Test update: TestCastLikeDtypeSafety → TestPythonLiteralAutocast.
New tests call _cast_module_dtype to simulate production dtype promotion,
then verify output dtype matches input dtype with no CastLike ops present.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
---
 src/mobius/components/_activations.py         |  2 +-
 src/mobius/components/_audio.py               |  6 +-
 src/mobius/components/_castlike_dtype_test.py | 61 ++++++++++++-------
 src/mobius/components/_diffusion.py           |  3 +-
 src/mobius/components/_moe.py                 | 17 +++---
 src/mobius/components/_rms_norm.py            |  4 +-
 src/mobius/models/gemma3n.py                  |  9 +--
 7 files changed, 54 insertions(+), 48 deletions(-)

diff --git a/src/mobius/components/_activations.py b/src/mobius/components/_activations.py
index 753f0e21..2773bd0d 100644
--- a/src/mobius/components/_activations.py
+++ b/src/mobius/components/_activations.py
@@ -41,7 +41,7 @@ def relu2(op: builder.OpBuilder, x):
 
 def quick_gelu(op: builder.OpBuilder, x):
     """QuickGELU activation: x * sigmoid(1.702 * x)."""
-    return op.Mul(x, op.Sigmoid(op.Mul(op.CastLike(op.Constant(value_float=1.702), x), x)))
+    return op.Mul(x, op.Sigmoid(op.Mul(1.702, x)))
 
 
 def mish(op: builder.OpBuilder, x):
diff --git a/src/mobius/components/_audio.py b/src/mobius/components/_audio.py
index 5871f1af..1e33ba0a 100644
--- a/src/mobius/components/_audio.py
+++ b/src/mobius/components/_audio.py
@@ -409,10 +409,8 @@ def __init__(
         self.layer_norm = LayerNorm(d_model)
 
     def forward(self, op: builder.OpBuilder, x: ir.Value, relative_attention_bias: ir.Value):
-        half = op.CastLike(op.Constant(value_float=0.5), x)
-
         # Macaron feed-forward in
-        x = op.Add(x, op.Mul(self.feed_forward_in(op, x), half))
+        x = op.Add(x, op.Mul(self.feed_forward_in(op, x), 0.5))
 
         # Multi-head attention with pre-norm
         norm_x = self.layer_norm_att(op, x)
@@ -422,7 +420,7 @@ def forward(self, op: builder.OpBuilder, x: ir.Value, relative_attention_bias: i
         x = op.Add(x, self.conv(op, x))
 
         # Macaron feed-forward out
-        x = op.Add(x, op.Mul(self.feed_forward_out(op, x), half))
+        x = op.Add(x, op.Mul(self.feed_forward_out(op, x), 0.5))
 
         return self.layer_norm(op, x)
 
diff --git a/src/mobius/components/_castlike_dtype_test.py b/src/mobius/components/_castlike_dtype_test.py
index 571cfe87..deffc0c4 100644
--- a/src/mobius/components/_castlike_dtype_test.py
+++ b/src/mobius/components/_castlike_dtype_test.py
@@ -1,11 +1,12 @@
 # Copyright (c) ONNX Project Contributors
 # SPDX-License-Identifier: Apache-2.0
 
-"""Tests that float constants are CastLike-wrapped for non-float32 dtypes.
+"""Tests that Python float literals auto-cast to match operand dtypes.
 
-PR #58 wraps bare float32 constants with op.CastLike() to prevent dtype
-mismatches when model tensors are float16 or bfloat16. These tests build
-component graphs with float16 inputs and verify CastLike ops are present.
+onnxscript auto-casts Python scalars (int/float/bool) to match the dtype of
+the other operand in a binary op. These tests verify that components using
+Python literals produce output in the correct dtype — no spurious FLOAT32
+constants widening BF16/FP16 computations.
 """
 
 from __future__ import annotations
@@ -13,55 +14,69 @@
 import onnx_ir as ir
 import pytest
 
+from mobius._builder import _cast_module_dtype
 from mobius._testing import count_op_type, create_test_builder, create_test_input
 from mobius.components._activations import quick_gelu
 from mobius.components._diffusion import AdaLayerNormOutput
 from mobius.components._rms_norm import OffsetRMSNorm
 
 
-class TestCastLikeDtypeSafety:
-    """Verify CastLike wrapping of float constants for non-float32 dtypes."""
+def _get_output_dtype(graph: ir.Graph) -> ir.DataType | None:
+    """Return the dtype of the first graph output, or None."""
+    if graph.outputs:
+        return graph.outputs[0].dtype
+    return None
+
+
+class TestPythonLiteralAutocast:
+    """Verify that Python float literals auto-cast to match tensor operand dtypes."""
 
     @pytest.mark.parametrize("dtype", [ir.DataType.FLOAT16, ir.DataType.BFLOAT16])
-    def test_offset_rms_norm_casts_constant(self, dtype: ir.DataType):
-        """OffsetRMSNorm adds 1.0 to weight — must CastLike for non-f32."""
+    def test_offset_rms_norm_constant_autocasts(self, dtype: ir.DataType):
+        """OffsetRMSNorm: `1.0` literal in op.Add auto-casts to weight dtype.
+
+        After _cast_module_dtype, the weight param is BF16/FP16. The Python
+        literal 1.0 in op.Add(self.weight, 1.0) must auto-cast to match.
+        """
         norm = OffsetRMSNorm(hidden_size=64, eps=1e-6)
+        _cast_module_dtype(norm, dtype)
         builder_, op, graph = create_test_builder()
         x = create_test_input(builder_, "x", [2, 3, 64], dtype)
         result = norm(op, x)
         graph.outputs.append(result)
-        assert count_op_type(graph, "CastLike") >= 1
+        # No CastLike — auto-cast handles it; output stays in the expected dtype
+        assert count_op_type(graph, "CastLike") == 0
+        assert _get_output_dtype(graph) == dtype
 
     @pytest.mark.parametrize("dtype", [ir.DataType.FLOAT16, ir.DataType.BFLOAT16])
-    def test_quick_gelu_casts_constant(self, dtype: ir.DataType):
-        """quick_gelu uses 1.702 constant — must CastLike for non-f32."""
+    def test_quick_gelu_constant_autocasts(self, dtype: ir.DataType):
+        """quick_gelu: `1.702` literal in op.Mul auto-casts to input dtype."""
         builder_, op, graph = create_test_builder()
         x = create_test_input(builder_, "x", [2, 3, 64], dtype)
         result = quick_gelu(op, x)
         graph.outputs.append(result)
-        assert count_op_type(graph, "CastLike") >= 1
+        assert count_op_type(graph, "CastLike") == 0
+        assert _get_output_dtype(graph) == dtype
 
     @pytest.mark.parametrize("dtype", [ir.DataType.FLOAT16, ir.DataType.BFLOAT16])
-    def test_ada_layer_norm_output_casts_constant(self, dtype: ir.DataType):
-        """AdaLayerNormOutput adds 1.0 to scale — must CastLike for non-f32."""
+    def test_ada_layer_norm_output_constant_autocasts(self, dtype: ir.DataType):
+        """AdaLayerNormOutput: `1.0` literal auto-casts to scale tensor dtype."""
         mod = AdaLayerNormOutput(hidden_size=64, eps=1e-6)
+        _cast_module_dtype(mod, dtype)
         builder_, op, graph = create_test_builder()
         hidden = create_test_input(builder_, "hidden", [1, 4, 64], dtype)
         temb = create_test_input(builder_, "temb", [1, 64], dtype)
         result = mod(op, hidden, temb)
         graph.outputs.append(result)
-        assert count_op_type(graph, "CastLike") >= 1
-
-    def test_no_castlike_needed_for_float32(self):
-        """With float32 inputs, CastLike is still present but benign.
+        assert count_op_type(graph, "CastLike") == 0
+        assert _get_output_dtype(graph) == dtype
 
-        This verifies the pattern is always applied regardless of dtype,
-        so there's no conditional logic that could miss a case.
-        """
+    def test_float32_inputs_produce_float32_output(self):
+        """Float32 inputs — no special casting needed, output stays float32."""
         norm = OffsetRMSNorm(hidden_size=64, eps=1e-6)
         builder_, op, graph = create_test_builder()
         x = create_test_input(builder_, "x", [2, 3, 64], ir.DataType.FLOAT)
         result = norm(op, x)
         graph.outputs.append(result)
-        # CastLike is always emitted (no-op for float32, but present)
-        assert count_op_type(graph, "CastLike") >= 1
+        assert count_op_type(graph, "CastLike") == 0
+        assert _get_output_dtype(graph) == ir.DataType.FLOAT
diff --git a/src/mobius/components/_diffusion.py b/src/mobius/components/_diffusion.py
index 738b5a35..62a64e4a 100644
--- a/src/mobius/components/_diffusion.py
+++ b/src/mobius/components/_diffusion.py
@@ -83,9 +83,8 @@ def forward(self, op: builder.OpBuilder, hidden_states: ir.Value, timestep_emb:
         emb = self._silu(op, timestep_emb)
         emb = self.linear(op, emb)
         shift, scale = op.Split(emb, num_outputs=2, axis=-1, _outputs=2)
-        one = op.CastLike(op.Constant(value_float=1.0), scale)
         hidden_states = self.norm(op, hidden_states)
-        hidden_states = op.Mul(hidden_states, op.Add(one, op.Unsqueeze(scale, [1])))
+        hidden_states = op.Mul(hidden_states, op.Add(1.0, op.Unsqueeze(scale, [1])))
         hidden_states = op.Add(hidden_states, op.Unsqueeze(shift, [1]))
         return hidden_states
 
diff --git a/src/mobius/components/_moe.py b/src/mobius/components/_moe.py
index a1b05a0c..5d541707 100644
--- a/src/mobius/components/_moe.py
+++ b/src/mobius/components/_moe.py
@@ -105,13 +105,9 @@ def forward(self, op: builder.OpBuilder, hidden_states: ir.Value):
         if self.norm_topk_prob:
             # Renormalize selected weights to sum to 1 (prevents vanishing gradients)
             weight_sum = op.ReduceSum(routing_weights, [-1], keepdims=True)
-            eps = op.CastLike(op.Constant(value_float=1e-9), routing_weights)
-            routing_weights = op.Div(routing_weights, op.Add(weight_sum, eps))
+            routing_weights = op.Div(routing_weights, op.Add(weight_sum, 1e-9))
         if self.routed_scaling_factor != 1.0:  # noqa: RUF069
-            scale = op.CastLike(
-                op.Constant(value_float=self.routed_scaling_factor), routing_weights
-            )
-            routing_weights = op.Mul(routing_weights, scale)
+            routing_weights = op.Mul(routing_weights, self.routed_scaling_factor)
         return routing_weights, selected_experts
 
 
@@ -145,9 +141,12 @@ def _threshold_mask_and_select(self, op, scores, jitter_eps):
         factor = op.Max(abs_scores, max_score)
         diff = op.Sub(max_score, scores)
         ratio = op.Div(diff, factor)
-        threshold = op.CastLike(op.Constant(value_float=2.0 * jitter_eps), scores)
+        threshold = 2.0 * jitter_eps
         mask = op.Greater(ratio, threshold)
-        neg_inf = op.CastLike(op.Constant(value_float=-1e30), scores)
+        # op.CastLike with Python literal: reuses a single constant, avoids cache-key
+        # collision that would occur if -1e30 were used as a plain literal in both
+        # op.Where (auto-cast to typed constant) and op.Expand (unbound → FLOAT).
+        neg_inf = op.CastLike(-1e30, scores)
         masked_scores = op.Where(mask, neg_inf, scores)
         weights = op.Softmax(masked_scores, axis=-1)
         k_one = op.Constant(value_ints=[1])
@@ -169,7 +168,7 @@ def forward(self, op: builder.OpBuilder, hidden_states: ir.Value):
             )
             all_weights.append(weight_k)
             all_experts.append(expert_k)
-            neg_inf = op.CastLike(op.Constant(value_float=-1e30), current_scores)
+            neg_inf = op.CastLike(-1e30, current_scores)
             current_scores = op.ScatterElements(
                 current_scores,
                 expert_k,
diff --git a/src/mobius/components/_rms_norm.py b/src/mobius/components/_rms_norm.py
index c31d42c3..3b6d1bb7 100644
--- a/src/mobius/components/_rms_norm.py
+++ b/src/mobius/components/_rms_norm.py
@@ -37,9 +37,7 @@ def __init__(self, hidden_size: int, eps: float = 1e-6):
         self.variance_epsilon = eps
 
     def forward(self, op: builder.OpBuilder, hidden_states: ir.Value):
-        effective_weight = op.Add(
-            self.weight, op.CastLike(op.Constant(value_float=1.0), self.weight)
-        )
+        effective_weight = op.Add(self.weight, 1.0)
         return op.RMSNormalization(
             hidden_states,
             effective_weight,
diff --git a/src/mobius/models/gemma3n.py b/src/mobius/models/gemma3n.py
index ac1c0973..33068d6a 100644
--- a/src/mobius/models/gemma3n.py
+++ b/src/mobius/models/gemma3n.py
@@ -196,8 +196,7 @@ def __init__(self, config: Gemma3nConfig):
     def _compute_router_modalities(self, op: builder.OpBuilder, x):
         """Compute router modalities: tanh(router(norm(x) * scale))."""
         router_input = self.router_norm(op, x)
-        scale = op.CastLike(op.Constant(value_float=self.router_input_scale), router_input)
-        router_input = op.Mul(router_input, scale)
+        router_input = op.Mul(router_input, self.router_input_scale)
         routed = self.modality_router(op, router_input)
         return op.Tanh(routed)
 
@@ -256,7 +255,7 @@ def correct(self, op: builder.OpBuilder, predictions: list, activated):
 
         # correction_coefs: [batch, seq, num_inputs] + 1
         all_coefs = self.correction_coefs(op, modalities)
-        all_coefs = op.Add(all_coefs, op.CastLike(op.Constant(value_float=1.0), all_coefs))
+        all_coefs = op.Add(all_coefs, 1.0)
 
         corrected = []
         for i in range(self.altup_num_inputs):
@@ -335,9 +334,7 @@ def forward(
         # Residual + Laurel (with sqrt(2) normalization)
         attn_gated = op.Add(active, attn_output)
         attn_laurel = op.Add(attn_gated, laurel_output)
-        attn_laurel = op.Div(
-            attn_laurel, op.CastLike(op.Constant(value_float=float(math.sqrt(2))), attn_laurel)
-        )
+        attn_laurel = op.Div(attn_laurel, float(math.sqrt(2)))
 
         # MLP
         mlp_input = self.pre_feedforward_layernorm(op, attn_laurel)

From 02e6fdb50aa457bdb8d6c8edaf4b40ce58c834dd Mon Sep 17 00:00:00 2001
From: Justin Chu <justinchuby@users.noreply.github.com>
Date: Wed, 1 Apr 2026 13:05:54 -0700
Subject: [PATCH 5/5] test: add BF16/FP16 tests for MoE gates and audio
 components

Add 10 new parametrized tests covering the 5 components identified by the
code reviewer as missing BF16/FP16 coverage:

SigmoidTopKGate (2 tests x 2 dtypes = 4):
- test_routing_weights_dtype: verifies 1e-9 epsilon in op.Add auto-casts
  so routing weights stay in FP16/BF16 (not widened to FP32)
- test_routed_scaling_factor_autocasts: verifies routed_scaling_factor
  Python float literal in op.Mul auto-casts to routing dtype

SparseMixerGate (1 test x 2 dtypes = 2):
- test_castlike_neg_inf_uses_input_dtype: verifies op.CastLike(-1e30, scores)
  correctly casts the -1e30 constant to the input dtype for use in
  op.Where and op.Expand, preventing FP32 type mismatches

ConformerEncoderLayer (1 test x 2 dtypes = 2):
- test_macaron_half_weight_autocasts: verifies 0.5 Macaron weight literals
  in op.Mul auto-cast to hidden state dtype (FP16/BF16)

Gemma3nAltUp (1 test x 2 dtypes = 2):
- test_router_input_scale_autocasts: verifies router_input_scale
  (hidden_size**-1.0 Python float) auto-casts in op.Mul during
  _compute_router_modalities

All 17 tests pass (10 new + 7 existing).

Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
---
 src/mobius/components/_castlike_dtype_test.py | 137 ++++++++++++++++++
 1 file changed, 137 insertions(+)

diff --git a/src/mobius/components/_castlike_dtype_test.py b/src/mobius/components/_castlike_dtype_test.py
index deffc0c4..cc97f3de 100644
--- a/src/mobius/components/_castlike_dtype_test.py
+++ b/src/mobius/components/_castlike_dtype_test.py
@@ -15,10 +15,14 @@
 import pytest
 
 from mobius._builder import _cast_module_dtype
+from mobius._configs import Gemma3nConfig
 from mobius._testing import count_op_type, create_test_builder, create_test_input
 from mobius.components._activations import quick_gelu
+from mobius.components._audio import ConformerEncoderLayer
 from mobius.components._diffusion import AdaLayerNormOutput
+from mobius.components._moe import SigmoidTopKGate, SparseMixerGate
 from mobius.components._rms_norm import OffsetRMSNorm
+from mobius.models.gemma3n import Gemma3nAltUp
 
 
 def _get_output_dtype(graph: ir.Graph) -> ir.DataType | None:
@@ -80,3 +84,136 @@ def test_float32_inputs_produce_float32_output(self):
         graph.outputs.append(result)
         assert count_op_type(graph, "CastLike") == 0
         assert _get_output_dtype(graph) == ir.DataType.FLOAT
+
+
+class TestSigmoidTopKGate:
+    """SigmoidTopKGate: verify 1e-9 epsilon and routed_scaling_factor auto-cast."""
+
+    @pytest.mark.parametrize("dtype", [ir.DataType.FLOAT16, ir.DataType.BFLOAT16])
+    def test_routing_weights_dtype(self, dtype: ir.DataType):
+        """Routing weights stay in input dtype — no FP32 widening from 1e-9 literal."""
+        gate = SigmoidTopKGate(
+            hidden_size=32,
+            num_experts=4,
+            top_k=2,
+            norm_topk_prob=True,
+        )
+        _cast_module_dtype(gate, dtype)
+        builder_, op, graph = create_test_builder()
+        x = create_test_input(builder_, "x", [1, 3, 32], dtype)
+        routing_weights, selected_experts = gate(op, x)
+        graph.outputs.extend([routing_weights, selected_experts])
+        # 1e-9 in op.Add(weight_sum, 1e-9) must auto-cast to routing dtype
+        assert count_op_type(graph, "CastLike") == 0
+        assert routing_weights.dtype == dtype
+
+    @pytest.mark.parametrize("dtype", [ir.DataType.FLOAT16, ir.DataType.BFLOAT16])
+    def test_routed_scaling_factor_autocasts(self, dtype: ir.DataType):
+        """routed_scaling_factor Python float literal auto-casts to routing dtype."""
+        gate = SigmoidTopKGate(
+            hidden_size=32,
+            num_experts=4,
+            top_k=2,
+            norm_topk_prob=False,
+            routed_scaling_factor=2.5,
+        )
+        _cast_module_dtype(gate, dtype)
+        builder_, op, graph = create_test_builder()
+        x = create_test_input(builder_, "x", [1, 3, 32], dtype)
+        routing_weights, _ = gate(op, x)
+        graph.outputs.append(routing_weights)
+        # routed_scaling_factor=2.5 in op.Mul must auto-cast, not widen to FP32
+        assert count_op_type(graph, "CastLike") == 0
+        assert routing_weights.dtype == dtype
+
+
+class TestSparseMixerGate:
+    """SparseMixerGate: verify CastLike(-1e30) preserves input dtype."""
+
+    @pytest.mark.parametrize("dtype", [ir.DataType.FLOAT16, ir.DataType.BFLOAT16])
+    def test_castlike_neg_inf_uses_input_dtype(self, dtype: ir.DataType):
+        """op.CastLike(-1e30, scores) must cast the constant to the scores dtype.
+
+        Without CastLike the -1e30 literal would be FLOAT32, causing type
+        mismatches in the op.Where and op.Expand downstream ops.
+        """
+        gate = SparseMixerGate(hidden_size=32, num_experts=4, top_k=2)
+        _cast_module_dtype(gate, dtype)
+        builder_, op, graph = create_test_builder()
+        x = create_test_input(builder_, "x", [1, 3, 32], dtype)
+        routing_weights, selected_experts = gate(op, x)
+        graph.outputs.extend([routing_weights, selected_experts])
+        # CastLike nodes should be present (the pattern is intentional for -1e30)
+        assert count_op_type(graph, "CastLike") > 0
+        # Final routing weights must remain in the input dtype
+        assert routing_weights.dtype == dtype
+
+
+class TestConformerEncoderLayer:
+    """ConformerEncoderLayer: verify 0.5 Macaron weight auto-casts."""
+
+    @pytest.mark.parametrize("dtype", [ir.DataType.FLOAT16, ir.DataType.BFLOAT16])
+    def test_macaron_half_weight_autocasts(self, dtype: ir.DataType):
+        """0.5 literal in op.Mul(feed_forward(x), 0.5) auto-casts to input dtype.
+
+        The Macaron structure applies half-weight feed-forward modules:
+        ``x += 0.5 * feed_forward_in(x)`` and ``x += 0.5 * feed_forward_out(x)``.
+        Both 0.5 literals must auto-cast to the hidden state dtype.
+        """
+        layer = ConformerEncoderLayer(d_model=32, num_heads=4, d_inner=64, kernel_size=3)
+        _cast_module_dtype(layer, dtype)
+        builder_, op, graph = create_test_builder()
+        x = create_test_input(builder_, "x", [1, 5, 32], dtype)
+        # relative_attention_bias: [num_heads, q_len, kv_len]
+        bias = create_test_input(builder_, "bias", [4, 5, 5], dtype)
+        result = layer(op, x, bias)
+        graph.outputs.append(result)
+        # No CastLike needed — Python float 0.5 auto-casts
+        assert count_op_type(graph, "CastLike") == 0
+        assert _get_output_dtype(graph) == dtype
+
+
+class TestGemma3nAltUp:
+    """Gemma3nAltUp: verify router_input_scale Python float auto-casts."""
+
+    def _make_config(self, hidden_size: int = 32) -> Gemma3nConfig:
+        from mobius._configs import ArchitectureConfig
+
+        base = ArchitectureConfig(
+            hidden_size=hidden_size,
+            intermediate_size=64,
+            num_attention_heads=4,
+            num_key_value_heads=2,
+            num_hidden_layers=1,
+            vocab_size=256,
+        )
+        return Gemma3nConfig(
+            **{k: getattr(base, k) for k in base.__dataclass_fields__ if hasattr(base, k)},
+            altup_num_inputs=2,
+            altup_active_idx=0,
+            altup_correct_scale=True,
+            laurel_rank=8,
+            hidden_size_per_layer_input=16,
+            vocab_size_per_layer_input=256,
+        )
+
+    @pytest.mark.parametrize("dtype", [ir.DataType.FLOAT16, ir.DataType.BFLOAT16])
+    def test_router_input_scale_autocasts(self, dtype: ir.DataType):
+        """router_input_scale = hidden_size**-1.0 auto-casts in op.Mul.
+
+        AltUp._compute_router_modalities multiplies a normalized hidden state
+        by self.router_input_scale (a Python float). This must not widen
+        BF16/FP16 computations to FP32.
+        """
+        config = self._make_config(hidden_size=32)
+        altup = Gemma3nAltUp(config)
+        _cast_module_dtype(altup, dtype)
+        builder_, op, graph = create_test_builder()
+        # altup_num_inputs=2 — provide two hidden state tensors
+        hs0 = create_test_input(builder_, "hs0", [1, 3, 32], dtype)
+        hs1 = create_test_input(builder_, "hs1", [1, 3, 32], dtype)
+        predicted = altup.predict(op, [hs0, hs1])
+        graph.outputs.extend(predicted)
+        # router_input_scale (float) in op.Mul must auto-cast
+        assert count_op_type(graph, "CastLike") == 0
+        assert predicted[0].dtype == dtype