Match Dots1 routing normalization floor

Preserve the pinned llama.cpp denominator clamp for normalized Dots1 expert weights while leaving native DeepSeek routing unchanged. Add a low-score ONNX Runtime parity regression for the boundary.

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

Author: justinchuby

src/mobius/_configs/_base.py

@@ -494,6 +494,7 @@ class ArchitectureConfig(BaseModelConfig):
     n_group: int = 1
     topk_group: int = 1
     routed_scaling_factor: float = 1.0
+    routing_weight_normalization_floor: float | None = None
     scoring_func: str = "softmax"
     topk_method: str = "greedy"
     first_k_dense_replace: int = 0

src/mobius/integrations/gguf/_config_mapping.py

@@ -1522,6 +1522,7 @@ def _conventional_shared_moe_postprocess(
         topk_method="greedy",
         n_group=1,
         topk_group=1,
+        routing_weight_normalization_floor=(6.103515625e-5 if arch == "dots1" else None),
         use_expert_bias=use_expert_bias,
         norm_topk_prob=norm_topk_prob,
         routed_scaling_factor=route_scale,

src/mobius/integrations/gguf/_config_mapping_test.py

@@ -1030,6 +1030,9 @@ def test_exact_shared_moe_contract(
         assert config.n_group == config.topk_group == 1
         assert config.norm_topk_prob is (architecture != "deepseek")
         assert config.routed_scaling_factor == pytest.approx(1.25)
+        assert config.routing_weight_normalization_floor == (
+            6.103515625e-5 if architecture == "dots1" else None
+        )
         assert config.attn_qk_norm is qk_norm
         assert config.tie_word_embeddings is False
 

src/mobius/models/deepseek.py

@@ -72,6 +72,7 @@ def __init__(self, config: ArchitectureConfig):
         self.topk_group = config.topk_group
         self.routed_scaling_factor = config.routed_scaling_factor
         self.norm_topk_prob = config.norm_topk_prob
+        self.routing_weight_normalization_floor = config.routing_weight_normalization_floor
         self.scoring_func = config.scoring_func
         self.topk_method = config.topk_method
         explicit_expert_bias = getattr(config, "use_expert_bias", None)
@@ -110,8 +111,13 @@ def forward(self, op: OpBuilder, hidden_states: ir.Value):
         # Normalize weights (V3 with norm_topk_prob=True)
         if self.norm_topk_prob:
             weight_sum = op.ReduceSum(routing_weights, [-1], keepdims=True)
-            eps = 1e-20
-            routing_weights = op.Div(routing_weights, op.Add(weight_sum, eps))
+            if self.routing_weight_normalization_floor is None:
+                denominator = op.Add(weight_sum, 1e-20)
+            else:
+                denominator = op.Max(
+                    weight_sum, float(self.routing_weight_normalization_floor)
+                )
+            routing_weights = op.Div(routing_weights, denominator)
 
         # Apply routing scale
         routing_weights = op.Mul(routing_weights, float(self.routed_scaling_factor))

src/mobius/models/deepseek_test.py

@@ -112,6 +112,48 @@ def test_sigmoid_gate_without_optional_correction_bias_has_no_bias_parameter():
     assert [name for name, _ in gate.named_parameters()] == ["weight"]
 
 
+def test_dots1_routing_normalization_uses_pinned_llama_floor():
+    config = make_config(
+        hidden_size=1,
+        num_local_experts=2,
+        num_experts_per_tok=2,
+        scoring_func="sigmoid",
+        use_expert_bias=False,
+        norm_topk_prob=True,
+        routing_weight_normalization_floor=6.103515625e-5,
+    )
+    gate = DeepSeekMoEGate(config)
+    gate.weight.const_value = ir.tensor(np.ones((2, 1), dtype=np.float32))
+
+    hidden = ir.Value(
+        name="hidden_states",
+        shape=ir.Shape([1, 1, 1]),
+        type=ir.TensorType(ir.DataType.FLOAT),
+    )
+    graph = ir.Graph(
+        inputs=[hidden],
+        outputs=[],
+        nodes=[],
+        name="dots1_routing_floor",
+        opset_imports={"": OPSET_VERSION},
+    )
+    builder = GraphBuilder(graph)
+    routing_weights, _ = gate(builder.op, hidden)
+    routing_weights.name = "routing_weights"
+    graph.outputs.append(routing_weights)
+
+    session = ort.InferenceSession(
+        ir.to_proto(ir.Model(graph, ir_version=11)).SerializeToString(),
+        providers=["CPUExecutionProvider"],
+    )
+    actual = session.run(None, {"hidden_states": np.array([[[-11.0]]], dtype=np.float32)})[0]
+    score = 1.0 / (1.0 + np.exp(11.0))
+    expected = np.full((1, 1, 2), score / 6.103515625e-5, dtype=np.float32)
+
+    np.testing.assert_allclose(actual, expected, rtol=2e-3, atol=1e-8)
+    assert actual.sum() < 1.0
+
+
 @pytest.mark.parametrize(
     ("scoring_func", "expects_bias"),
     [("softmax", False), ("sigmoid", True)],