Fix SmolLM3 inverted no_rope_layers logic and add L3/L5 tests

src/mobius/components/_rotary_embedding.py

@@ -96,10 +96,31 @@ def apply_rotary_pos_emb(
 
 
 class BaseRope(nn.Module):
-    """Base class for rotary position embeddings."""
+    """Base class for rotary position embeddings.
 
-    def __init__(self, cos_cache_data: np.ndarray, sin_cache_data: np.ndarray):
+    Subclasses that accept an ``ArchitectureConfig`` should pass
+    ``dtype=config.dtype`` to ``super().__init__()`` so that cos/sin
+    embeddings are cast to the model compute dtype.
+
+    Args:
+        cos_cache_data: Precomputed cosine cache (numpy float32).
+        sin_cache_data: Precomputed sine cache (numpy float32).
+        dtype: Model compute dtype. When not FP32, a Cast is inserted
+            after gathering cos/sin embeddings so that the
+            ``RotaryEmbedding`` op receives inputs with matching types.
+            The cache itself stays FP32 for precision; only the gathered
+            per-position embeddings are cast.
+    """
+
+    def __init__(
+        self,
+        cos_cache_data: np.ndarray,
+        sin_cache_data: np.ndarray,
+        *,
+        dtype: ir.DataType = ir.DataType.FLOAT,
+    ):
         super().__init__()
+        self._dtype = dtype
         self.cos_cache = nn.Parameter(
             list(cos_cache_data.shape),
             name="cos_cache",
@@ -111,15 +132,33 @@ def __init__(self, cos_cache_data: np.ndarray, sin_cache_data: np.ndarray):
             data=ir.tensor(sin_cache_data),
         )
 
+    def _cast_embeddings(
+        self,
+        op: builder.OpBuilder,
+        cos: ir.Value,
+        sin: ir.Value,
+    ) -> tuple[ir.Value, ir.Value]:
+        """Cast cos/sin embeddings to model dtype if needed.
+
+        The RotaryEmbedding op requires all tensor inputs to share the
+        same type ``T``.  Since caches are stored in FP32, a Cast is
+        needed when the model uses FP16 or BF16.
+        """
+        if self._dtype != ir.DataType.FLOAT:
+            cos = op.Cast(cos, to=self._dtype)
+            sin = op.Cast(sin, to=self._dtype)
+        return cos, sin
+
     def forward(self, op: builder.OpBuilder, position_ids: ir.Value):
-        return get_rotary_pos_emb(op, position_ids, self.cos_cache, self.sin_cache)
+        cos, sin = get_rotary_pos_emb(op, position_ids, self.cos_cache, self.sin_cache)
+        return self._cast_embeddings(op, cos, sin)
 
 
 class DefaultRope(BaseRope):
     def __init__(self, config: ArchitectureConfig):
         inv_freq = _get_default_inv_freq(config)
         cos_cache, sin_cache = _get_cos_sin_cache(config.max_position_embeddings, inv_freq)
-        super().__init__(cos_cache, sin_cache)
+        super().__init__(cos_cache, sin_cache, dtype=config.dtype)
 
 
 class ProportionalRope(BaseRope):
@@ -161,15 +200,15 @@ def __init__(self, config: ArchitectureConfig):
 
         # cos/sin cache shape: [max_pos, head_dim] (emb = cat([freqs, freqs]))
         cos_cache, sin_cache = _get_cos_sin_cache(config.max_position_embeddings, inv_freq)
-        super().__init__(cos_cache, sin_cache)
+        super().__init__(cos_cache, sin_cache, dtype=config.dtype)
 
 
 class LinearRope(BaseRope):
     def __init__(self, config: ArchitectureConfig):
         inv_freq = _get_default_inv_freq(config)
         inv_freq = inv_freq / config.rope_scaling["factor"]
         cos_cache, sin_cache = _get_cos_sin_cache(config.max_position_embeddings, inv_freq)
-        super().__init__(cos_cache, sin_cache)
+        super().__init__(cos_cache, sin_cache, dtype=config.dtype)
 
 
 class DynamicNTKRope(BaseRope):
@@ -189,7 +228,7 @@ def __init__(self, config: ArchitectureConfig):
         new_theta = config.rope_theta * (factor ** (dim / (dim - 2)))
         inv_freq = 1.0 / (new_theta ** (np.arange(0, dim, 2, dtype=np.float32) / dim))
         cos_cache, sin_cache = _get_cos_sin_cache(config.max_position_embeddings, inv_freq)
-        super().__init__(cos_cache, sin_cache)
+        super().__init__(cos_cache, sin_cache, dtype=config.dtype)
 
 
 class Llama3Rope(BaseRope):
@@ -218,7 +257,7 @@ def __init__(self, config: ArchitectureConfig):
         cos_cache, sin_cache = _get_cos_sin_cache(
             config.max_position_embeddings, inv_freq_llama
         )
-        super().__init__(cos_cache, sin_cache)
+        super().__init__(cos_cache, sin_cache, dtype=config.dtype)
 
 
 class LongRope(BaseRope):
@@ -244,7 +283,7 @@ def __init__(self, config: ArchitectureConfig):
             original_max_pos, inv_freq / short_factor, attention_factor
         )
         if not self.has_long_cache:
-            super().__init__(short_cos, short_sin)
+            super().__init__(short_cos, short_sin, dtype=config.dtype)
             return
 
         long_cos, long_sin = _get_cos_sin_cache(
@@ -254,7 +293,7 @@ def __init__(self, config: ArchitectureConfig):
         )
         cos_cache = np.concatenate([short_cos, long_cos], axis=0)
         sin_cache = np.concatenate([short_sin, long_sin], axis=0)
-        super().__init__(cos_cache, sin_cache)
+        super().__init__(cos_cache, sin_cache, dtype=config.dtype)
 
     def forward(self, op: builder.OpBuilder, position_ids: ir.Value):
         if self.has_long_cache:
@@ -265,7 +304,8 @@ def forward(self, op: builder.OpBuilder, position_ids: ir.Value):
             )
             offset = op.Mul(use_long, self.original_max_position_embeddings)
             position_ids = op.Add(position_ids, offset)
-        return get_rotary_pos_emb(op, position_ids, self.cos_cache, self.sin_cache)
+        cos, sin = get_rotary_pos_emb(op, position_ids, self.cos_cache, self.sin_cache)
+        return self._cast_embeddings(op, cos, sin)
 
 
 class YarnRope(BaseRope):
@@ -342,17 +382,18 @@ def find_correction_dim(num_rotations):
         cos_cache, sin_cache = _get_cos_sin_cache(
             config.max_position_embeddings, inv_freq, attention_factor
         )
-        super().__init__(cos_cache, sin_cache)
+        super().__init__(cos_cache, sin_cache, dtype=config.dtype)
 
         # Store llama_4_scaling_beta for Ministral3 position-dependent query scaling.
         # When set, forward() returns (cos, sin, attn_scale) instead of (cos, sin).
         self._llama4_beta = rope_scaling.get("llama_4_scaling_beta")
         self._llama4_original_max_pos = float(original_max_pos)
 
     def forward(self, op: builder.OpBuilder, position_ids: ir.Value):
-        cos_sin = get_rotary_pos_emb(op, position_ids, self.cos_cache, self.sin_cache)
+        cos, sin = get_rotary_pos_emb(op, position_ids, self.cos_cache, self.sin_cache)
+        cos, sin = self._cast_embeddings(op, cos, sin)
         if self._llama4_beta is None:
-            return cos_sin
+            return cos, sin
 
         # Compute position-dependent attention scale for Ministral3/Mistral4:
         #   scale = 1 + beta * log(1 + floor(position_ids / original_max_pos))
@@ -362,13 +403,12 @@ def forward(self, op: builder.OpBuilder, position_ids: ir.Value):
         floored = op.Floor(op.Div(pos_float, float(self._llama4_original_max_pos)))
         log_term = op.Log(op.Add(floored, 1.0))
         attn_scale = op.Add(op.Mul(log_term, float(self._llama4_beta)), 1.0)
-        # Cast to match model dtype (e.g. FP16) — cos_cache has the target dtype
-        cos_dtype = self.cos_cache.dtype
-        if cos_dtype is not None and cos_dtype != ir.DataType.FLOAT:
-            attn_scale = op.Cast(attn_scale, to=cos_dtype)
+        # Cast to match model compute dtype
+        if self._dtype != ir.DataType.FLOAT:
+            attn_scale = op.Cast(attn_scale, to=self._dtype)
         # Unsqueeze to [batch, seq_len, 1] for broadcasting with 3D query states
         attn_scale = op.Unsqueeze(attn_scale, [-1])
-        return (cos_sin[0], cos_sin[1], attn_scale)
+        return (cos, sin, attn_scale)
 
 
 class _MRopeBase(BaseRope):
@@ -390,7 +430,7 @@ def __init__(
     ):
         inv_freq = _get_default_inv_freq(config)
         cos_cache, sin_cache = _get_cos_sin_cache(config.max_position_embeddings, inv_freq)
-        super().__init__(cos_cache, sin_cache)
+        super().__init__(cos_cache, sin_cache, dtype=config.dtype)
 
         rotary_dim = len(inv_freq)
         self.h_mask = nn.Parameter(
@@ -453,7 +493,7 @@ def forward(self, op: builder.OpBuilder, position_ids: ir.Value):
         sin = op.Where(self.h_mask, sin_h, sin_t)
         sin = op.Where(self.w_mask, sin_w, sin)
 
-        return cos, sin
+        return self._cast_embeddings(op, cos, sin)
 
 
 class ChunkedMRope(_MRopeBase):

src/mobius/models/smollm.py

@@ -78,14 +78,17 @@ def forward(
             attn_bias = (
                 sliding_attn_bias if layer_type == "sliding_attention" else full_attn_bias
             )
-            # SmolLM3 uses no_rope_layers to gate RoPE per layer:
-            # no_rope_layers[i] == 1 → skip RoPE, 0 → apply RoPE
-            skip_rope = (
-                self.no_rope_layers is not None
-                and i < len(self.no_rope_layers)
-                and self.no_rope_layers[i] == 1
+            # SmolLM3 uses no_rope_layers to gate RoPE per layer.
+            # Despite the name, the HF convention is:
+            #   no_rope_layers[i] == 1 → USE RoPE
+            #   no_rope_layers[i] == 0 → skip RoPE
+            # (HF assigns self.use_rope = config.no_rope_layers[layer_idx])
+            use_rope = (
+                self.no_rope_layers is None
+                or i >= len(self.no_rope_layers)
+                or self.no_rope_layers[i] == 1
             )
-            rope = None if skip_rope else position_embeddings
+            rope = position_embeddings if use_rope else None
 
             hidden_states, present_kv = layer(
                 op,

testdata/golden/causal-lm/smollm3-3b.json

@@ -14,22 +14,22 @@
     1131
   ],
   "top10_logits": [
-    "0x1.b324e60000000p+4",
-    "0x1.a30b7a0000000p+4",
-    "0x1.87b1fa0000000p+4",
-    "0x1.76de420000000p+4",
-    "0x1.6d041e0000000p+4",
-    "0x1.5c8ad40000000p+4",
-    "0x1.5144c40000000p+4",
-    "0x1.513a040000000p+4",
-    "0x1.5008a00000000p+4",
-    "0x1.4f57620000000p+4"
+    "0x1.b324e00000000p+4",
+    "0x1.a30b740000000p+4",
+    "0x1.87b1f00000000p+4",
+    "0x1.76de380000000p+4",
+    "0x1.6d04120000000p+4",
+    "0x1.5c8ad00000000p+4",
+    "0x1.5144c00000000p+4",
+    "0x1.5139fa0000000p+4",
+    "0x1.5008980000000p+4",
+    "0x1.4f575a0000000p+4"
   ],
   "logits_summary": [
-    "0x1.b324e60000000p+4",
-    "-0x1.d251fa0000000p+3",
-    "0x1.662fc1b5f12a7p+2",
-    "0x1.a02a75d17f04cp+1"
+    "0x1.b324e00000000p+4",
+    "-0x1.d2520c0000000p+3",
+    "0x1.662fad9be0416p+2",
+    "0x1.a02a76b63e452p+1"
   ],
   "input_ids": [
     12805,

tests/_test_configs.py

@@ -409,7 +409,7 @@ def _base_config(config_cls=None, **overrides) -> ArchitectureConfig:
         True,
     ),
     ("qwen3_vl_text", {"attn_qk_norm": True}, False),
-    ("smollm3", {}, False),
+    ("smollm3", {"no_rope_layers": [1, 0]}, True),  # exercise per-layer RoPE gating
     # === Mixture of Experts ===
     (
         "phimoe",

tests/e2e_golden_test.py

@@ -155,7 +155,6 @@ def _use_temp_hf_cache(tmp_path):
     "text-generation/helium-1-2b": "Helium decode loop diverges from HF after first token",
     "text-generation/nanochat-d20": "NanoChat decode loop diverges from HF after first token",
     "text-generation/ernie4_5-0_3b": "ERNIE 4.5 decode loop diverges from HF after first token",
-    "text-generation/smollm3-3b": "SmolLM3 3B decode loop diverges from HF (FP32 precision with 3B params)",
     # MLA compressed KV cache dimensions not yet handled by OnnxGenerator
     "text-generation/youtu-2b": "Youtu MLA KV cache dims differ from standard attention (v_head_dim != head_dim)",
 }

tests/integration_test.py

@@ -86,6 +86,8 @@ def _model_accessible(model_id: str) -> bool:
     # CausalLMModel (base: llama/mistral/qwen2 architecture)
     pytest.param("Qwen/Qwen2.5-0.5B", False, id="qwen2.5-0.5b"),
     pytest.param("HuggingFaceTB/SmolLM-135M", False, id="smollm-135m"),
+    # SmolLM3 (per-layer RoPE gating via no_rope_layers)
+    pytest.param("HuggingFaceTB/SmolLM3-3B", False, id="smollm3-3b"),
     # Gemma
     pytest.param("google/gemma-3-1b-pt", False, id="gemma3-1b"),
     # Granite