From 9a21e3a3c1ed220768b28f82bbb4be53abedb38e Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Fri, 1 May 2026 21:23:27 +0000 Subject: [PATCH 1/5] Fix SmolLM3 inverted no_rope_layers logic and add L3 test SmolLM3's no_rope_layers config was interpreted backwards: the HF convention is no_rope_layers[i] == 1 means USE RoPE (despite the misleading name), but mobius was treating it as skip RoPE. This caused 27/36 layers to have wrong RoPE gating, resulting in max logit diffs of ~6.9 vs HF (vs 0.00005 after fix). Changes: - Fix inverted RoPE gating in SmolLM3TextModel.forward() - Remove L5 xfail now that generation matches HF exactly - Add L3 integration test entry for SmolLM3-3B - Mark smollm3 as representative in L1 test configs - Regenerate L4 golden data with updated HF reference Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu --- src/mobius/models/smollm.py | 17 ++++++++------ testdata/golden/causal-lm/smollm3-3b.json | 28 +++++++++++------------ tests/_test_configs.py | 2 +- tests/e2e_golden_test.py | 2 +- tests/integration_test.py | 2 ++ 5 files changed, 28 insertions(+), 23 deletions(-) diff --git a/src/mobius/models/smollm.py b/src/mobius/models/smollm.py index c5edf0a0b..14d1b48e5 100644 --- a/src/mobius/models/smollm.py +++ b/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, diff --git a/testdata/golden/causal-lm/smollm3-3b.json b/testdata/golden/causal-lm/smollm3-3b.json index d9c8e49c2..733c7b3aa 100644 --- a/testdata/golden/causal-lm/smollm3-3b.json +++ b/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, diff --git a/tests/_test_configs.py b/tests/_test_configs.py index f619787da..f26009964 100644 --- a/tests/_test_configs.py +++ b/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", {}, True), # per-layer RoPE gating via no_rope_layers # === Mixture of Experts === ( "phimoe", diff --git a/tests/e2e_golden_test.py b/tests/e2e_golden_test.py index ff4d0b15c..2ce719ac2 100644 --- a/tests/e2e_golden_test.py +++ b/tests/e2e_golden_test.py @@ -155,7 +155,7 @@ 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)", } diff --git a/tests/integration_test.py b/tests/integration_test.py index dd16a06dd..2227e4fd2 100644 --- a/tests/integration_test.py +++ b/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 From 2b01ac3fc8ff2da66cc494adb2fc8c10b45d2ca3 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Fri, 1 May 2026 21:27:12 +0000 Subject: [PATCH 2/5] Add no_rope_layers to SmolLM3 L1 test config Exercises per-layer RoPE gating in the unit test so the inverted polarity bug is directly covered at L1 level (layer 0 = use RoPE, layer 1 = skip RoPE). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu --- tests/_test_configs.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/_test_configs.py b/tests/_test_configs.py index f26009964..a40390c7e 100644 --- a/tests/_test_configs.py +++ b/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", {}, True), # per-layer RoPE gating via no_rope_layers + ("smollm3", {"no_rope_layers": [1, 0]}, True), # exercise per-layer RoPE gating # === Mixture of Experts === ( "phimoe", From 004ed070bd89fac4e6a4a69f9f7b964dbfaaa055 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Fri, 1 May 2026 21:49:26 +0000 Subject: [PATCH 3/5] Fix BF16 type mismatch in RoPE cos/sin embeddings The ONNX RotaryEmbedding op requires all tensor inputs (x, cos, sin) to share the same type T. The RoPE cos/sin cache is computed in FP32 (numpy limitation), so when the model uses BF16 or FP16, the RotaryEmbedding op received mixed types: BF16 query/key + FP32 cos/sin. This cascaded through Attention -> o_proj -> residual Add, causing ORT to throw: Type Error: Type parameter (T) of Optype (Add) bound to different types (tensor(bfloat16) and tensor(float)) Fix: Add a dtype parameter to BaseRope that all RoPE subclasses thread from config.dtype. BaseRope._cast_embeddings() inserts Cast ops after gathering cos/sin when the model dtype differs from FP32. For FP32 models, no Cast is added (zero overhead). Covers all RoPE variants: DefaultRope, ProportionalRope, LinearRope, DynamicNTKRope, Llama3Rope, LongRope, YarnRope, and MRope (chunked/interleaved). Also fixes YarnRope's attn_scale dtype derivation to use the stored model dtype instead of inspecting cos_cache.dtype (which stays FP32). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu --- src/mobius/components/_rotary_embedding.py | 84 ++++++++++++++++------ 1 file changed, 63 insertions(+), 21 deletions(-) diff --git a/src/mobius/components/_rotary_embedding.py b/src/mobius/components/_rotary_embedding.py index 8a831686a..559d7a7c7 100644 --- a/src/mobius/components/_rotary_embedding.py +++ b/src/mobius/components/_rotary_embedding.py @@ -96,10 +96,27 @@ 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): + 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 +128,35 @@ 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,7 +198,7 @@ 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): @@ -169,7 +206,7 @@ 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 +226,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 +255,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 +281,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 +291,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 +302,10 @@ 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,7 +382,7 @@ 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). @@ -350,9 +390,12 @@ def find_correction_dim(num_rotations): 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 +405,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 +432,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 +495,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): From 8537a2adebccf78bc5c20556de1d7bde98c416f6 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Fri, 1 May 2026 21:52:55 +0000 Subject: [PATCH 4/5] Add subclass guidance to BaseRope docstring Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu --- src/mobius/components/_rotary_embedding.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/src/mobius/components/_rotary_embedding.py b/src/mobius/components/_rotary_embedding.py index 559d7a7c7..24c40a01f 100644 --- a/src/mobius/components/_rotary_embedding.py +++ b/src/mobius/components/_rotary_embedding.py @@ -98,6 +98,10 @@ def apply_rotary_pos_emb( class BaseRope(nn.Module): """Base class for rotary position embeddings. + 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). From 1565b079698197bafe544876e406059644ffdc21 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Fri, 1 May 2026 23:32:12 +0000 Subject: [PATCH 5/5] Fix lint: formatting in RoPE and golden test Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu --- src/mobius/components/_rotary_embedding.py | 12 +++--------- tests/e2e_golden_test.py | 1 - 2 files changed, 3 insertions(+), 10 deletions(-) diff --git a/src/mobius/components/_rotary_embedding.py b/src/mobius/components/_rotary_embedding.py index 24c40a01f..542829772 100644 --- a/src/mobius/components/_rotary_embedding.py +++ b/src/mobius/components/_rotary_embedding.py @@ -150,9 +150,7 @@ def _cast_embeddings( return cos, sin 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) return self._cast_embeddings(op, cos, sin) @@ -306,9 +304,7 @@ 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) - 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) return self._cast_embeddings(op, cos, sin) @@ -394,9 +390,7 @@ def find_correction_dim(num_rotations): 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 diff --git a/tests/e2e_golden_test.py b/tests/e2e_golden_test.py index 2ce719ac2..120f78a13 100644 --- a/tests/e2e_golden_test.py +++ b/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", - # 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)", }