Guard fused QMoE CUDA routing

Keep fused QMoE activation inputs in the model dtype and fail fast for CUDA exports that require router_weights, since ORT CUDA QMoE currently ignores input 14. Add regression tests for activation dtype and the CUDA guard.\n\nCo-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Signed-off-by: justinchuby <justinchuby@users.noreply.github.com>

Author: justinchuby

src/mobius/components/_moe.py

@@ -14,6 +14,7 @@
 import torch
 from onnxscript import OpBuilder, nn
 
+from mobius._build_context import ep_capabilities
 from mobius._configs import ArchitectureConfig, QuantizationConfig
 from mobius._weight_utils import preprocess_awq_weights, preprocess_gptq_weights
 from mobius.components._mlp import MLP
@@ -279,6 +280,7 @@ def __init__(
         config: ArchitectureConfig,
         gate: nn.Module | None = None,
         linear_class: type | None = None,
+        enable_qmoe: bool = True,
     ):
         super().__init__()
         assert config.num_local_experts is not None
@@ -296,7 +298,11 @@ def __init__(
             if config.moe_intermediate_size is not None
             else config
         )
-        if self._qmoe_quantization is not None and hasattr(self.gate, "qmoe_routing"):
+        if (
+            enable_qmoe
+            and self._qmoe_quantization is not None
+            and hasattr(self.gate, "qmoe_routing")
+        ):
             self.experts = None
             self._init_qmoe_parameters(expert_config)
         else:
@@ -529,12 +535,20 @@ def __init__(
             )
 
     def forward(self, op: OpBuilder, hidden_states: ir.Value) -> ir.Value:
+        if ep_capabilities().name == "cuda":
+            raise ValueError(
+                "FusedQuantizedMoE is disabled for CUDA because ORT CUDA QMoE "
+                "currently ignores router_weights (input 14), which is required "
+                "for GLM/DeepSeek group-limited routing. Use the decomposed "
+                "MatMulNBits export path for CUDA."
+            )
         hidden = self._hidden
 
         # QMoE requires 2-D router_probs, so flatten [B, S, H] -> [rows, H].
         orig_shape = op.Shape(hidden_states)
         flat = op.Reshape(hidden_states, op.Constant(value_ints=[-1, hidden]))
-        # QMoE input/router_probs must be float32.
+        # Router math must be float32, but QMoE activation input stays in the
+        # model dtype: CUDA QMoE kernels are registered for fp16/bf16 inputs.
         flat_f32 = op.Cast(flat, to=1)
 
         scores_for_choice, routing_weights, selected_experts = self._route(op, flat_f32)
@@ -545,7 +559,7 @@ def forward(self, op: OpBuilder, hidden_states: ir.Value) -> ir.Value:
         aggregation = op.ScatterElements(zeros, selected_experts, routing_weights, axis=-1)
 
         moe_out = op.QMoE(
-            flat_f32,  # 0: input
+            flat,  # 0: input
             scores_for_choice,  # 1: router_probs (selection logits)
             self.fc1_experts_weights,  # 2
             op.Cast(self.fc1_scales, to=1),  # 3

src/mobius/components/_moe_test.py

@@ -186,6 +186,30 @@ def test_fused_qmoe_wires_asymmetric_zero_points():
     assert layer.fc2_experts_zero_points.shape == ir.Shape([4, 32, 1])
 
 
+def test_fused_qmoe_keeps_activation_input_model_dtype():
+    config = make_config(
+        hidden_size=32,
+        intermediate_size=16,
+        moe_intermediate_size=16,
+        num_local_experts=4,
+        num_experts_per_tok=2,
+        quantization=QuantizationConfig(
+            bits=4,
+            group_size=16,
+            quant_method="gptq",
+            sym=True,
+        ),
+    )
+    layer = FusedQuantizedMoE(config, DeepSeekMoEGate(config))
+    builder, op, graph = create_test_builder()
+    hidden = create_test_input(builder, "hidden", [1, 2, 32], dtype=ir.DataType.FLOAT16)
+    builder._adapt_outputs([layer(op, hidden)], "")
+
+    qmoe = next(node for node in graph if node.op_type == "QMoE")
+    assert qmoe.inputs[0].dtype == ir.DataType.FLOAT16
+    assert qmoe.inputs[1].dtype == ir.DataType.FLOAT
+
+
 def test_deepseek_fused_qmoe_graph_has_one_node_per_moe_layer():
     config = make_config(
         hidden_size=32,
@@ -230,6 +254,32 @@ def test_deepseek_fused_qmoe_graph_has_one_node_per_moe_layer():
     )
 
 
+def test_deepseek_fused_qmoe_rejects_cuda():
+    config = make_config(
+        hidden_size=32,
+        intermediate_size=32,
+        moe_intermediate_size=16,
+        num_hidden_layers=2,
+        first_k_dense_replace=1,
+        num_local_experts=4,
+        num_experts_per_tok=2,
+        n_shared_experts=1,
+        fused_quantized_moe=True,
+        quantization=QuantizationConfig(
+            bits=4,
+            group_size=16,
+            quant_method="gguf",
+            sym=True,
+        ),
+    )
+    with pytest.raises(ValueError, match="CUDA.*ignores router_weights"):
+        build_from_module(
+            DeepSeekV3CausalLMModel(config),
+            config,
+            execution_provider="cuda",
+        )
+
+
 def test_expert_major_packing_matches_static_64_expert_top6_reference():
     torch.manual_seed(0)
     experts, top_k = 64, 6