Preserve native Bolt routing normalization parity

Select denominator routes with repeated ArgMax and lower-index tie breaking instead of CUDA TopK. Accumulate unbiased mixing weights in selected-rank FP32 order instead of a tree ReduceSum, preserving the validated V3 native routing arithmetic before scaling and the BF16 output store.

Add exact CPU/CUDA coefficient regressions for cutoff ties, non-associative sums, zero weights, top-k boundaries, and unnormalized scaling. The new sum-order cases fail against the previous emitter. Weights, recurrence, cache interfaces, and the ORT runtime are unchanged.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Signed-off-by: rui-ren <ruiren1225@gmail.com>

Author: rui-ren

src/mobius/models/_nemotron_h_omni_packed.py

@@ -334,9 +334,22 @@ def emit(op):
             # routing, prepare a row denominator from the same full selection
             # scores; MoEPacked still selects and gathers its own expert routes.
             if normalize:
-                _, chosen = op.TopK(scores, [topk], axis=-1, largest=1, sorted=1, _outputs=2)
-                chosen_weights = op.GatherElements(mixing, chosen, axis=-1)
-                total = op.ReduceSum(chosen_weights, [-1], keepdims=1)
+                # Match native lower-index tie breaking and rank-ordered FP32
+                # addition; CUDA TopK and tree reductions can change the denominator.
+                remaining = scores
+                total = None
+                mask = None
+                for rank in range(topk):
+                    chosen = op.ArgMax(remaining, axis=-1, keepdims=1, select_last_index=0)
+                    chosen_weight = op.GatherElements(mixing, chosen, axis=-1)
+                    total = chosen_weight if total is None else op.Add(total, chosen_weight)
+                    if rank + 1 < topk:
+                        if mask is None:
+                            mask = op.ConstantOfShape(
+                                op.Shape(chosen),
+                                value=ir.tensor([-float("inf")], dtype=ir.DataType.FLOAT),
+                            )
+                        remaining = op.ScatterElements(remaining, chosen, mask, axis=-1)
                 mixing = op.Div(mixing, op.Add(total, 1e-20))
             mixing = op.Mul(mixing, layer.routed_scaling_factor)
             normalize = 0

src/mobius/models/_nemotron_h_omni_packed_test.py

@@ -207,9 +207,7 @@ def test_compact_decoder_wires_full_fp32_routing_without_compact_ids(normalize,
             assert scaled.op_type == "Mul"
             assert scaled.inputs[1].const_value.numpy().item() == np.float32(scale)
             assert next(iter(consumer.outputs[0].uses())).node.op_type != "Cast"
-    assert sum(node.op_type == "TopK" for node in graph) == (
-        len(consumers) if normalize and scale != 1.0 else 0  # noqa: RUF069
-    )
+    assert not any(node.op_type == "TopK" for node in graph)
     assert not any(
         node.op_type in {"MoEPackedIQ2XXS", "MatMulPackedIQ2XXS", "MatMulPackedQ8"}
         for node in graph.all_nodes()
@@ -225,31 +223,55 @@ def test_compact_decoder_wires_full_fp32_routing_without_compact_ids(normalize,
     )
 
 
-@pytest.mark.parametrize("topk", [1, 2])
+@pytest.mark.parametrize("provider", ["CPUExecutionProvider", "CUDAExecutionProvider"])
+@pytest.mark.parametrize("topk", [1, 2, 8])
 @pytest.mark.parametrize("normalize", [False, True])
-def test_compact_decoder_prescales_selected_mixing_in_fp32(tmp_path, topk, normalize):
-    graph = optimize_nemotron_h_omni_cuda(_fixture(topk=topk, normalize=normalize, scale=5.0))[
-        "decoder"
-    ].graph
+def test_compact_decoder_prescales_selected_mixing_in_fp32(
+    tmp_path, topk, normalize, provider
+):
+    if provider not in ort.get_available_providers():
+        pytest.skip(f"{provider} is unavailable")
+    graph = optimize_nemotron_h_omni_cuda(
+        _fixture(topk=topk, experts=8, normalize=normalize, scale=5.0)
+    )["decoder"].graph
     router = next(node for node in graph if node.op_type == "MoEGroupRouter")
     consumer = next(node for node in graph if node.op_type == "MoEPacked")
     preparation = ir.convenience.extract(
         graph, inputs=list(router.outputs[1:]), outputs=[consumer.inputs[5]]
     )
     path = tmp_path / "mixing.onnx"
     ir.save(ir.Model(preparation, ir_version=10), path)
-    session = ort.InferenceSession(str(path), providers=["CPUExecutionProvider"])
-    scores = np.array([[0.7, 0.9, 0.9, 0.6], [0.2, 0.5, 0.1, 0.4]], np.float32)
-    mixing = np.array([[0.8, 0.2, 0.6, 0.9], [0, 0, 0, 0]], np.float32)
+    session = ort.InferenceSession(str(path), providers=[provider])
+    assert session.get_providers()[0] == provider
+    scores = np.array(
+        [
+            [30, 29.577760696411133, 29.577760696411133, 4, 3, 2, 1, 0],
+            [0, 1, 2, 3, 4, 5, 6, 7],
+            [7, 6, 5, 4, 3, 2, 1, 0],
+        ],
+        np.float32,
+    )
+    mixing = np.array(
+        [
+            [0.8, 0.2, 0.6, 0.9, 0.3, 0.4, 0.5, 0.1],
+            [0, 0, 0, 0, 0, 0, 0, 0],
+            [1, 2**-24, 2**-24, 2**-24, 2**-24, 2**-24, 2**-24, 2**-24],
+        ],
+        np.float32,
+    )
     actual = session.run(
         None, dict(zip([value.name for value in preparation.inputs], [scores, mixing]))
     )[0]
     chosen = np.argsort(-scores, axis=-1, kind="stable")[:, :topk]
-    expected = np.take_along_axis(mixing, chosen, axis=-1)
+    expected = mixing
     if normalize:
-        expected = expected / (expected.sum(axis=-1, keepdims=True) + np.float32(1e-20))
+        selected = np.take_along_axis(mixing, chosen, axis=-1)
+        total = selected[:, :1].copy()
+        for rank in range(1, topk):
+            total = total + selected[:, rank : rank + 1]
+        expected = expected / (total + np.float32(1e-20))
     expected = expected * np.float32(5.0)
-    np.testing.assert_array_equal(np.take_along_axis(actual, chosen, axis=-1), expected)
+    np.testing.assert_array_equal(actual, expected)
     assert actual.dtype == np.float32
     assert np.isfinite(actual).all()