Fix QMoE scale precision: scales silently forced entire MoE FFN onto CPU EP (#505)

## Bug

When exporting a hybrid MoE model (e.g. `Qwen/Qwen3.6-35B-A3B`) via the
native QMoE-emission path in `MoELayer` (`_moe.py`) at a non-float32
precision (e.g. `fp16` for CUDA), the `fc1_scales`/`fc2_scales`
parameters were pinned with `_keep_float32 = True`, which
`_cast_module_dtype()` (`_builder.py`) honors by skipping them during
the FLOAT->target-dtype cast pass. Every other parameter (weights,
activations) was correctly cast to FLOAT16, but the scales stayed
FLOAT32.

ONNX Runtime's `com.microsoft::QMoE` kernel (`quant_type="int"`)
requires its `T2` type constraint (used by `fc1_scales`/`fc2_scales`) to
be **exactly equal to `T`** (the activation dtype) — confirmed by
reading
`onnxruntime/contrib_ops/{cuda,cpu}/moe/moe_quantization*.{cc,h}`. There
is no registered kernel variant for `T2=FLOAT32` while `T=FLOAT16`.

Because of this mismatch, ONNX Runtime silently fails to find *any*
matching QMoE kernel — on **either** CPU or CUDA EP — and falls back to
placing every QMoE node on `CPUExecutionProvider`, inserting
`InsertedPrecisionFreeCast_*` bridging nodes. Since QMoE is the entire
MoE FFN compute (the dominant cost for a large MoE model), this silently
forced 100% of the MoE compute onto CPU even when the workflow
explicitly targeted CUDA — and also broke `enable_cuda_graph` in ONNX
Runtime GenAI, since not all decoder nodes were assigned to the same EP.

## Fix

- `_moe.py`: stop pinning `fc1_scales`/`fc2_scales` to FLOAT32; let them
be downcast to the model's target export precision (FP16/BF16) like
every other float parameter.
- `_qmoe_fusion.py` (the separate "dense-fallback -> QMoE" rewrite
rule): had the same class of bug — scales and router logits were
hardcoded to FLOAT/FLOAT32 regardless of the model's activation dtype.
Fixed to use the activation dtype instead.
- Updated `_moe_test.py` and `_qmoe_fusion_test.py` to reflect the
corrected (dtype-matching) behavior, and added explicit FP32/FP16/BF16
coverage in `_qmoe_fusion_test.py`.

## Validation

- `python -m pytest src/mobius/components/_moe_test.py
src/mobius/rewrite_rules/_qmoe_fusion_test.py -v` — 39 passed.
- `ruff` — passed.
- Real-model validation: re-ran the Olive `Rtn(moe=true)` +
`MobiusBuilder(precision=fp16, CUDA EP)` export against
`Qwen/Qwen3.6-35B-A3B`. All 430 QMoE scale initializers in the exported
graph are now FLOAT16 (previously FLOAT32).
- Loaded the re-exported model with `onnxruntime` (verbose session
logging): **zero** "CUDA kernel not found in registries" messages for
QMoE, and all 40 QMoE nodes are now correctly assigned to
`CUDAExecutionProvider` (previously all 40 fell back to CPU).
- `onnxruntime_genai.Model(...)` now loads successfully with
`enable_cuda_graph=1` in `genai_config.json`, with no manual workaround
needed (previously required manually stripping `enable_cuda_graph` from
the generated config).

Related: builds on top of the export fixes merged in #495.

Author: titaiwangms

src/mobius/components/_moe.py

@@ -371,15 +371,13 @@ def _init_qmoe_parameters(self, expert_config: ArchitectureConfig) -> None:
             dtype=ir.DataType.UINT8,
         )
         self.fc1_scales = nn.Parameter([self.num_experts, fc1_out, hidden_size // block_size])
-        self.fc1_scales._keep_float32 = True
         self.fc2_experts_weights = nn.Parameter(
             [self.num_experts, hidden_size, intermediate_size * bits // 8],
             dtype=ir.DataType.UINT8,
         )
         self.fc2_scales = nn.Parameter(
             [self.num_experts, hidden_size, intermediate_size // block_size]
         )
-        self.fc2_scales._keep_float32 = True
         if quantization.sym:
             self.fc1_experts_zero_points = None
             self.fc2_experts_zero_points = None

src/mobius/components/_moe_test.py

@@ -244,8 +244,8 @@ def test_int4_moe_emits_expert_major_qmoe(self):
 
         _cast_module_dtype(layer, ir.DataType.FLOAT16)
         assert layer.gate.weight.dtype == ir.DataType.FLOAT16
-        assert layer.fc1_scales.dtype == ir.DataType.FLOAT
-        assert layer.fc2_scales.dtype == ir.DataType.FLOAT
+        assert layer.fc1_scales.dtype == ir.DataType.FLOAT16
+        assert layer.fc2_scales.dtype == ir.DataType.FLOAT16
 
     def test_expert_major_packing_matches_static_64_expert_top6_reference(self):
         """Packed QMoE math matches the existing loop-over-experts semantics."""

src/mobius/rewrite_rules/_group_query_attention_test.py

@@ -327,12 +327,14 @@ def test_deepseek_v2_lite_int4_uses_one_qmoe_per_moe_layer(self):
             # computed in FLOAT32 for routing-score numerical stability.
             assert node.inputs[1] is not None
             assert node.inputs[1].dtype == ir.DataType.FLOAT16
-            # fc1_scales/fc2_scales (indices 3, 6) are quantization scale
-            # tensors, independent of the hidden_states/router dtype.
+            # fc1_scales/fc2_scales (indices 3, 6) must match the
+            # hidden_states/activation dtype (FLOAT16 here): QMoE's kernel
+            # registration requires T2 (scales) to exactly equal T
+            # (activation), so scales are no longer pinned to FLOAT32.
             assert node.inputs[3] is not None
-            assert node.inputs[3].dtype == ir.DataType.FLOAT
+            assert node.inputs[3].dtype == ir.DataType.FLOAT16
             assert node.inputs[6] is not None
-            assert node.inputs[6].dtype == ir.DataType.FLOAT
+            assert node.inputs[6].dtype == ir.DataType.FLOAT16
             assert node.inputs[14] is not None
             assert node.inputs[14].dtype == ir.DataType.FLOAT16
 

src/mobius/rewrite_rules/_qmoe_fusion.py

@@ -21,14 +21,15 @@
   (the packed ``uint8`` bytes are copied unchanged; ``swiglu_fusion=2`` tells the
   kernel the gate half precedes the up half).
 * ``fc2_experts_weights`` = per-expert ``flatten(down_w)`` (bytes unchanged).
-* scales are upcast ``float16 -> float32`` (value-exact, lossless) because the
-  QMoE kernel requires ``float32`` scales and ``router_probs``.
+* scales are cast to the activation dtype because QMoE's ``T2`` kernel type
+  constraint requires integer-quantization scales to match ``input``.
 * zero-points are copied unchanged (same low-nibble-first packing as
   ``MatMulNBits``).
 
-The router logits (a quantized ``MatMulNBits`` gate) are ``Cast`` to ``float32``
-and passed as ``router_probs`` with ``normalize_routing_weights=1`` -- matching
-:class:`mobius.components._moe.TopKGate` (``Softmax(TopK(logits))``).
+The router logits (a quantized ``MatMulNBits`` gate) are cast to the activation
+dtype and passed as ``router_probs`` with ``normalize_routing_weights=1`` --
+matching :class:`mobius.components._moe.TopKGate`
+(``Softmax(TopK(logits))``).
 
 Any surrounding **shared expert** (Qwen2-MoE style ``shared_expert`` +
 ``shared_expert_gate``) is left untouched: only the routed per-expert storm and
@@ -319,12 +320,14 @@ def _pack_projection(nodes: list[ir.Node], slot: int) -> np.ndarray:
     return np.stack(stacked, axis=0)
 
 
-def _pack_scales(nodes: list[ir.Node], slot: int, out_features: int) -> np.ndarray:
+def _pack_scales(
+    nodes: list[ir.Node], slot: int, out_features: int, dtype: ir.DataType
+) -> np.ndarray:
     stacked = []
     for node in nodes:
         arr = _array(node.inputs[slot]).reshape(out_features, -1)
-        stacked.append(arr.astype(np.float32))
-    return np.stack(stacked, axis=0)
+        stacked.append(arr)
+    return np.stack(stacked, axis=0).astype(dtype.numpy())
 
 
 def _pack_zero_points(nodes: list[ir.Node], slot: int, out_features: int) -> np.ndarray:
@@ -334,7 +337,9 @@ def _pack_zero_points(nodes: list[ir.Node], slot: int, out_features: int) -> np.
     return np.stack(stacked, axis=0)
 
 
-def _fuse_layer(graph: ir.Graph, layer: _DenseMoELayer, index: int) -> None:
+def _fuse_layer(
+    graph: ir.Graph, layer: _DenseMoELayer, index: int, activation_dtype: ir.DataType
+) -> None:
     ids = sorted(layer.experts)
     gate_nodes = [layer.experts[i].gate for i in ids]
     up_nodes = [layer.experts[i].up for i in ids]
@@ -350,10 +355,10 @@ def _fuse_layer(graph: ir.Graph, layer: _DenseMoELayer, index: int) -> None:
     fc1_w = np.concatenate([gate_w, up_w], axis=1)
     fc2_w = _pack_projection(down_nodes, 1)
 
-    gate_s = _pack_scales(gate_nodes, 2, inter)
-    up_s = _pack_scales(up_nodes, 2, inter)
+    gate_s = _pack_scales(gate_nodes, 2, inter, activation_dtype)
+    up_s = _pack_scales(up_nodes, 2, inter, activation_dtype)
     fc1_s = np.concatenate([gate_s, up_s], axis=1)
-    fc2_s = _pack_scales(down_nodes, 2, hidden_dim)
+    fc2_s = _pack_scales(down_nodes, 2, hidden_dim, activation_dtype)
 
     fc1_zp = fc2_zp = None
     if has_zp:
@@ -366,11 +371,11 @@ def _fuse_layer(graph: ir.Graph, layer: _DenseMoELayer, index: int) -> None:
     fc1_w_v = _make_initializer(
         graph, f"{prefix}.fc1_experts_weights", fc1_w, ir.DataType.UINT8
     )
-    fc1_s_v = _make_initializer(graph, f"{prefix}.fc1_scales", fc1_s, ir.DataType.FLOAT)
+    fc1_s_v = _make_initializer(graph, f"{prefix}.fc1_scales", fc1_s, activation_dtype)
     fc2_w_v = _make_initializer(
         graph, f"{prefix}.fc2_experts_weights", fc2_w, ir.DataType.UINT8
     )
-    fc2_s_v = _make_initializer(graph, f"{prefix}.fc2_scales", fc2_s, ir.DataType.FLOAT)
+    fc2_s_v = _make_initializer(graph, f"{prefix}.fc2_scales", fc2_s, activation_dtype)
     fc1_zp_v = (
         _make_initializer(graph, f"{prefix}.fc1_zero_points", fc1_zp, ir.DataType.UINT8)
         if fc1_zp is not None
@@ -385,7 +390,7 @@ def _fuse_layer(graph: ir.Graph, layer: _DenseMoELayer, index: int) -> None:
     cast = ir.node(
         "Cast",
         inputs=[layer.logits],
-        attributes={"to": ir.DataType.FLOAT.value},
+        attributes={"to": activation_dtype.value},
         num_outputs=1,
         name=f"{prefix}.router_probs_cast",
     )
@@ -460,9 +465,9 @@ def fuse_dense_moe_to_qmoe(model: ir.Model) -> int:
     """Fuse every dense-fallback MoE subgraph in ``model`` into ``QMoE`` nodes.
 
     Reuses the existing int4 ``MatMulNBits`` expert weights byte-for-byte (pure
-    expert-major concat/layout transform, no requantization). ``float16`` scales
-    are upcast to the ``float32`` the QMoE kernel requires (lossless). Any shared
-    expert is preserved.
+    expert-major concat/layout transform, no requantization). Scales and router
+    probabilities are cast to the activation dtype required by QMoE's kernel
+    type constraints. Any shared expert is preserved.
 
     Args:
         model: The IR model to rewrite in place.
@@ -492,7 +497,26 @@ def fuse_dense_moe_to_qmoe(model: ir.Model) -> int:
                 block_size,
             )
             continue
-        _fuse_layer(graph, layer, index)
+        activation_dtype = layer.hidden.dtype
+        if activation_dtype is None:
+            # `layer.hidden` is often a graph-internal, not-yet-type-inferred
+            # value in real models. Fall back to the dtype of an existing
+            # expert's MatMulNBits scales, which already matches the
+            # pre-fusion activation dtype (scales share QMoE's/MatMulNBits's
+            # "T" type constraint with the activation).
+            activation_dtype = down.inputs[2].dtype
+        if activation_dtype not in {
+            ir.DataType.FLOAT,
+            ir.DataType.FLOAT16,
+            ir.DataType.BFLOAT16,
+        }:
+            logger.warning(
+                "skipping MoE layer at %s: QMoE activation dtype is missing or unsupported (%s)",
+                layer.topk.name,
+                activation_dtype,
+            )
+            continue
+        _fuse_layer(graph, layer, index, activation_dtype)
         fused += 1
     if fused:
         _remove_dead_nodes(graph)

src/mobius/rewrite_rules/_qmoe_fusion_test.py

@@ -92,15 +92,24 @@ def _dequant(
 class _Quant:
     """A randomly-generated int4 ``MatMulNBits`` weight triple."""
 
-    def __init__(self, rng: np.random.Generator, n: int, kdim: int) -> None:
+    def __init__(
+        self,
+        rng: np.random.Generator,
+        n: int,
+        kdim: int,
+        scale_dtype: ir.DataType = ir.DataType.FLOAT16,
+    ) -> None:
         n_blocks = kdim // BLOCK
         codes = rng.integers(0, 16, size=(n, kdim), dtype=np.int32)
         zp_codes = rng.integers(0, 16, size=(n, n_blocks), dtype=np.int32)
         self.weight = _pack_weight(codes, BLOCK)
-        self.scales = rng.random((n, n_blocks), dtype=np.float32).astype(np.float16) * 0.1
+        self.scales = (rng.random((n, n_blocks), dtype=np.float32) * 0.1).astype(
+            scale_dtype.numpy()
+        )
         self.zero_points = _pack_zero_points(zp_codes)
         self.n = n
         self.kdim = kdim
+        self.scale_dtype = scale_dtype
 
     @property
     def dense(self) -> np.ndarray:
@@ -118,7 +127,7 @@ def _init(graph: ir.Graph, name: str, arr: np.ndarray, dtype: ir.DataType) -> ir
 
 def _matmulnbits(name: str, x: ir.Value, q: _Quant, graph: ir.Graph) -> ir.Value:
     w = _init(graph, f"{name}.weight", q.weight, ir.DataType.UINT8)
-    s = _init(graph, f"{name}.scales", q.scales, ir.DataType.FLOAT16)
+    s = _init(graph, f"{name}.scales", q.scales, q.scale_dtype)
     z = _init(graph, f"{name}.zero_points", q.zero_points, ir.DataType.UINT8)
     node = ir.node(
         "MatMulNBits",
@@ -145,21 +154,33 @@ def _constant_int(graph_nodes: list[ir.Node], name: str, value: int) -> ir.Value
     return node.outputs[0]
 
 
-def _build_dense_graph() -> tuple[ir.Model, dict[str, _Quant], np.ndarray]:
-    """Build a tiny dense-fallback Qwen35-MoE graph and return it with its weights."""
+def _build_dense_graph(
+    activation_dtype: ir.DataType = ir.DataType.FLOAT16,
+    *,
+    hidden_dtype: ir.DataType | None = None,
+) -> tuple[ir.Model, dict[str, _Quant], np.ndarray]:
+    """Build a tiny dense-fallback Qwen35-MoE graph and return it with its weights.
+
+    ``activation_dtype`` controls the dtype of MatMulNBits scales/router casts.
+    ``hidden_dtype`` (defaults to ``activation_dtype``) controls the declared
+    type of the ``hidden`` graph input; pass ``ir.DataType.UNDEFINED`` to
+    simulate a not-yet-type-inferred graph-internal value.
+    """
     rng = _rng()
     quants: dict[str, _Quant] = {}
     nodes: list[ir.Node] = []
+    if hidden_dtype is None:
+        hidden_dtype = activation_dtype
 
     hidden = ir.Value(
         name="hidden",
         shape=ir.Shape(["T", H]),
-        type=ir.TensorType(ir.DataType.FLOAT16),
+        type=ir.TensorType(hidden_dtype) if hidden_dtype != ir.DataType.UNDEFINED else None,
     )
     graph = ir.Graph([hidden], [], nodes=[], name="tiny_moe")
 
     def q(key: str, n: int, kdim: int) -> _Quant:
-        quants[key] = _Quant(rng, n, kdim)
+        quants[key] = _Quant(rng, n, kdim, scale_dtype=activation_dtype)
         return quants[key]
 
     # Router: quantized gate MatMulNBits -> TopK -> Softmax.
@@ -287,7 +308,7 @@ def q(key: str, n: int, kdim: int) -> _Quant:
     )
     final.outputs[0].name = "moe_out"
     final.outputs[0].shape = ir.Shape(["T", H])
-    final.outputs[0].type = ir.TensorType(ir.DataType.FLOAT16)
+    final.outputs[0].type = ir.TensorType(activation_dtype)
     nodes.append(final)
 
     for node in nodes:
@@ -425,22 +446,50 @@ def test_qmoe_weights_are_bit_identical_to_concatenated_experts() -> None:
         np.testing.assert_array_equal(fc2_z[e], quants[f"d{e}"].zero_points)
 
 
-def test_scales_are_lossless_float32_upcast() -> None:
-    model, quants, _ = _build_dense_graph()
+@pytest.mark.parametrize(
+    "activation_dtype",
+    [ir.DataType.FLOAT, ir.DataType.FLOAT16, ir.DataType.BFLOAT16],
+)
+def test_scales_match_activation_dtype(activation_dtype: ir.DataType) -> None:
+    model, quants, _ = _build_dense_graph(activation_dtype)
     fuse_dense_moe_to_qmoe(model)
     qmoe = next(n for n in model.graph if n.op_type == "QMoE")
 
     fc1_s = qmoe.inputs[3].const_value.numpy()
-    assert fc1_s.dtype == np.float32
+    fc2_s = qmoe.inputs[6].const_value.numpy()
+    assert qmoe.inputs[3].dtype == activation_dtype
+    assert qmoe.inputs[6].dtype == activation_dtype
+    assert fc1_s.dtype == activation_dtype.numpy()
+    assert fc2_s.dtype == activation_dtype.numpy()
     for e in range(E):
         expected = np.concatenate(
             [
                 quants[f"g{e}"].scales.reshape(INTER, -1),
                 quants[f"u{e}"].scales.reshape(INTER, -1),
             ],
             axis=0,
-        ).astype(np.float32)
+        ).astype(activation_dtype.numpy())
         np.testing.assert_array_equal(fc1_s[e], expected)
+        expected_fc2 = quants[f"d{e}"].scales.reshape(H, -1).astype(activation_dtype.numpy())
+        np.testing.assert_array_equal(fc2_s[e], expected_fc2)
+
+
+def test_fuses_when_hidden_dtype_is_untyped() -> None:
+    """Graph-internal `hidden` values are commonly untyped before shape inference.
+
+    The fusion must still succeed by falling back to the dtype of an existing
+    expert's MatMulNBits scales, rather than silently skipping the layer.
+    """
+    model, _, _ = _build_dense_graph(ir.DataType.FLOAT16, hidden_dtype=ir.DataType.UNDEFINED)
+    graph = model.graph
+    assert graph.inputs[0].dtype is None
+
+    fused = fuse_dense_moe_to_qmoe(model)
+
+    assert fused == 1
+    qmoe = next(n for n in graph if n.op_type == "QMoE")
+    assert qmoe.inputs[3].dtype == ir.DataType.FLOAT16
+    assert qmoe.inputs[6].dtype == ir.DataType.FLOAT16
 
 
 def test_rewritten_forward_matches_dense_forward() -> None:
@@ -451,15 +500,19 @@ def test_rewritten_forward_matches_dense_forward() -> None:
     np.testing.assert_allclose(got, expected, rtol=1e-5, atol=1e-5)
 
 
-def test_router_probs_cast_to_float32() -> None:
-    model, _, _ = _build_dense_graph()
+@pytest.mark.parametrize(
+    "activation_dtype",
+    [ir.DataType.FLOAT, ir.DataType.FLOAT16, ir.DataType.BFLOAT16],
+)
+def test_router_probs_cast_to_activation_dtype(activation_dtype: ir.DataType) -> None:
+    model, _, _ = _build_dense_graph(activation_dtype)
     fuse_dense_moe_to_qmoe(model)
     graph = model.graph
     qmoe = next(n for n in graph if n.op_type == "QMoE")
     router_probs = qmoe.inputs[1]
     cast = router_probs.producer()
     assert cast.op_type == "Cast"
-    assert cast.attributes["to"].value == ir.DataType.FLOAT.value
+    assert cast.attributes["to"].value == activation_dtype.value
 
 
 def test_attributes_match_qmoe_abi() -> None: