AutoEP: preserve native Qwen3.5 wrapper text routing

deepspeed/module_inject/auto_ep.py

@@ -126,11 +126,13 @@ def _has_3d_expert_params(module: nn.Module, preset: MoEModelPreset) -> bool:
 
 def _get_num_experts_from_config(model_config, preset: MoEModelPreset) -> int | None:
     """Extract num_experts from model.config using the preset's attribute name."""
+    model_config = getattr(model_config, "text_config", model_config)
     return getattr(model_config, preset.num_experts_attr, None)
 
 
 def _get_top_k_from_config(model_config, preset: MoEModelPreset) -> int | None:
     """Extract top_k from model.config using the preset's attribute name."""
+    model_config = getattr(model_config, "text_config", model_config)
     return getattr(model_config, preset.top_k_attr, None)
 
 
@@ -611,6 +613,8 @@ def ep_parser(self) -> list[MoELayerSpec]:
                     router_logits_capture_mode=forward_contract.router_logits_capture_mode,
                     moe_output_shape=forward_contract.moe_output_shape,
                     e_score_correction_bias_path=e_score_correction_bias_path,
+                    router_topk_sorted=preset.router_topk_sorted,
+                    router_scores_in_gate_dtype=preset.router_scores_in_gate_dtype,
                     expert_activation=expert_activation,
                     expert_activation_alpha=activation_alpha,
                     expert_activation_limit=activation_limit,

deepspeed/module_inject/auto_ep_layer.py

@@ -595,6 +595,8 @@ def __init__(
                 route_scale=spec.route_scale,
                 gate_bias=spec.gate_bias,
                 group_score_func=spec.group_score_func,
+                topk_sorted=spec.router_topk_sorted,
+                scores_in_gate_dtype=spec.router_scores_in_gate_dtype,
             )
             # Copy gate weights
             _copy_parameter_data(self.router.gate.weight, source_gate.weight)

deepspeed/module_inject/auto_ep_presets/base.py

@@ -58,6 +58,8 @@ class MoEModelPreset:
     min_transformers_version: str | None = None
     validated_transformers_versions: str = ""
     docs_support_notes: str = ""
+    router_topk_sorted: bool = False
+    router_scores_in_gate_dtype: bool = False
 
 
 @dataclass
@@ -97,6 +99,8 @@ class MoELayerSpec:
     router_logits_capture_mode: Literal["raw", "post_score"] = "post_score"
     moe_output_shape: Literal["batched", "flat"] = "batched"
     e_score_correction_bias_path: str | None = None
+    router_topk_sorted: bool = False
+    router_scores_in_gate_dtype: bool = False
     #: Expert MLP activation resolved for this layer (see MoEModelPreset.expert_activation).
     expert_activation: str = "swiglu"
     expert_activation_alpha: float = 1.702

deepspeed/module_inject/auto_ep_presets/qwen3_5_moe.py

@@ -20,7 +20,7 @@
 PRESET_NAME = "qwen3_5_moe"
 
 PRESET = MoEModelPreset(
-    moe_layer_pattern=r"model\.layers\.\d+\.mlp",
+    moe_layer_pattern=r"model\.(?:language_model\.)?layers\.\d+\.mlp",
     router_pattern="gate",
     experts_pattern="experts",
     expert_storage="fused_3d",
@@ -37,13 +37,11 @@
     shared_experts_pattern="shared_expert",
     shared_experts_gate_pattern="shared_expert_gate",
     preset_adapter="qwen3_5_moe",
-    hf_model_types=("qwen3_5_moe_text", ),
-    unsupported_hf_model_type_notes={
-        "qwen3_5_moe": ("AutoEP supports the Qwen3.5 text backbone preset path; pass the "
-                        "text-backbone model/config with model_type='qwen3_5_moe_text'.")
-    },
+    hf_model_types=("qwen3_5_moe_text", "qwen3_5_moe"),
     min_transformers_version="5.2.0",
-    docs_support_notes="Requires the Qwen3.5 text-backbone qwen3_5_moe_text model type.",
+    docs_support_notes="Text backbone and the language-model path of the Qwen3.5 MoE wrapper.",
+    router_topk_sorted=True,
+    router_scores_in_gate_dtype=True,
 )
 
 
@@ -61,6 +59,7 @@ def adjust_forward_contract(self, contract: ForwardContract) -> ForwardContract:
             return_router_logits=True,
             capture_target="router",
             capture_index=1,
+            router_logits_capture_mode="raw",
         )
 
     def retarget_transformers_output_recorders(

deepspeed/moe/ep_router.py

@@ -60,6 +60,8 @@ def __init__(
         route_scale: float,
         gate_bias: bool,
         group_score_func: str = "top2_sum",
+        topk_sorted: bool = False,
+        scores_in_gate_dtype: bool = False,
     ):
         super().__init__()
         self.gate = nn.Linear(dim, num_experts, bias=gate_bias)
@@ -71,6 +73,8 @@ def __init__(
         self.route_norm = route_norm
         self.route_scale = route_scale
         self.group_score_func = group_score_func
+        self.topk_sorted = topk_sorted
+        self.scores_in_gate_dtype = scores_in_gate_dtype
         # Trainable expert score correction bias (e.g. DeepSeek-V3/Moonlight noaux_tc).
         # Separate from the dynamic load-balancing expert_bias passed in forward().
         self.e_score_correction_bias = None
@@ -152,6 +156,7 @@ def forward(
         """
         # Gate projection -> (T, num_experts)
         scores = self.gate(x)
+        gate_dtype = scores.dtype
 
         # Scoring in float32 to avoid loss explosion
         if self.score_func == "sigmoid":
@@ -172,7 +177,7 @@ def forward(
             scores_for_choice = self._get_node_limited_routing_scores(scores_for_choice)
 
         # Select top-k experts per token
-        _, selected_experts_indices = torch.topk(scores_for_choice, k=self.top_k, dim=-1, sorted=False)
+        _, selected_experts_indices = torch.topk(scores_for_choice, k=self.top_k, dim=-1, sorted=self.topk_sorted)
 
         # Gather original (unbiased) scores for selected experts
         top_scores = scores.gather(dim=1, index=selected_experts_indices)
@@ -183,6 +188,8 @@ def forward(
             top_scores = top_scores / denominator
 
         top_scores = top_scores * self.route_scale
+        if self.scores_in_gate_dtype:
+            top_scores = top_scores.to(gate_dtype)
 
         num_tokens_per_expert = count_tokens_per_expert(selected_experts_indices, self.num_experts)
 

docs/code-docs/source/autoep.rst

@@ -33,7 +33,7 @@ Transformers build that exposes the matching config/model classes,
        but the tested ``4.x`` builds do not match the validated AutoEP layout.
    * - ``qwen3_5_moe``
      - ``5.2.0``
-     - Requires the Qwen3.5 text-backbone ``qwen3_5_moe_text`` model type;
+     - Supports text input through the text backbone or the MoE wrapper; vision input is untested;
        for performance on Qwen3.5's Gated DeltaNet layers, install optimized
        kernels. See the `Hugging Face Transformers kernel loading docs
        <https://huggingface.co/docs/transformers/kernel_doc/loading_kernels>`__

tests/unit/v1/moe/test_autoep_unit.py

@@ -1029,9 +1029,8 @@ def test_preset_registry_core_contracts(self):
         assert preset_name_for_hf_model_type("minimax_m3_vl_text") == "minimax_m3"
         assert preset_name_for_hf_model_type("llama4_text") is None
 
-        qwen35 = unsupported_preset_for_hf_model_type("qwen3_5_moe")
-        assert qwen35 is not None
-        assert "qwen3_5_moe_text" in qwen35[1].unsupported_hf_model_type_notes["qwen3_5_moe"]
+        assert preset_name_for_hf_model_type("qwen3_5_moe") == "qwen3_5_moe"
+        assert unsupported_preset_for_hf_model_type("qwen3_5_moe") is None
         minimax = unsupported_preset_for_hf_model_type("minimax_m3_vl")
         assert minimax is not None
         assert "minimax_m3_vl_text" in minimax[1].unsupported_hf_model_type_notes["minimax_m3_vl"]
@@ -2173,8 +2172,9 @@ def test_qwen_adapter_guards(self, monkeypatch):
         assert specs[0].model_family == "qwen3_moe"
 
         model.config.model_type = "qwen3_5_moe"
-        with pytest.raises(ValueError, match="qwen3_5_moe_text"):
-            AutoEP(model, _runtime_config(enabled=True, autoep_size=1))._resolve_presets()
+        monkeypatch.setattr(get_preset_adapter("qwen3_5_moe"), "_installed_transformers_version", lambda: "5.16.1")
+        specs = AutoEP(model, _runtime_config(enabled=True, autoep_size=1)).ep_parser()
+        assert specs[0].model_family == "qwen3_5_moe"
 
     def test_deepseek_v3_detection_and_score_correction_bias_copy(self, monkeypatch):
         FakeGatheredParameters.calls = []
@@ -2811,3 +2811,79 @@ def test_fused_gate_up_source_feeds_both_w1_and_w3(self):
         assert [id(p) for p in w1] == [id(experts_source.gate_up_proj)]
         assert [id(p) for p in w3] == [id(experts_source.gate_up_proj)]
         assert [id(p) for p in w2] == [id(experts_source.down_proj)]
+
+
+@pytest.mark.parametrize("wrapper", [False, True])
+@pytest.mark.parametrize("autocast", [False, True])
+def test_qwen35_autoep_preserves_native_router_and_text_outputs(wrapper, autocast):
+    # Catches wrong wrapper paths, top-k/dtype drift, and post-softmax logit capture.
+    transformers = pytest.importorskip("transformers")
+    required = ("Qwen3_5MoeTextConfig", "Qwen3_5MoeConfig", "Qwen3_5MoeForCausalLM",
+                "Qwen3_5MoeForConditionalGeneration")
+    if not all(hasattr(transformers, name) for name in required):
+        pytest.skip("Transformers does not provide the Qwen3.5 MoE model classes")
+    Qwen3_5MoeTextConfig = transformers.Qwen3_5MoeTextConfig
+    Qwen3_5MoeConfig = transformers.Qwen3_5MoeConfig
+    Qwen3_5MoeForCausalLM = transformers.Qwen3_5MoeForCausalLM
+    Qwen3_5MoeForConditionalGeneration = transformers.Qwen3_5MoeForConditionalGeneration
+
+    torch.manual_seed(42)
+    config = Qwen3_5MoeTextConfig(vocab_size=128,
+                                  hidden_size=128,
+                                  num_hidden_layers=1,
+                                  num_attention_heads=2,
+                                  num_key_value_heads=1,
+                                  head_dim=64,
+                                  moe_intermediate_size=128,
+                                  shared_expert_intermediate_size=128,
+                                  num_experts=4,
+                                  num_experts_per_tok=2,
+                                  layer_types=["full_attention"],
+                                  use_cache=False)
+    cls = Qwen3_5MoeForCausalLM
+    if wrapper:
+        config = Qwen3_5MoeConfig(text_config=config.to_dict(),
+                                  vision_config={
+                                      "depth": 1,
+                                      "hidden_size": 32,
+                                      "intermediate_size": 64,
+                                      "num_heads": 4,
+                                      "out_hidden_size": 128,
+                                      "num_position_embeddings": 16
+                                  })
+        cls = Qwen3_5MoeForConditionalGeneration
+    config._attn_implementation = "eager"
+    reference = cls(config).to(dtype=torch.bfloat16).eval()
+    candidate = copy.deepcopy(reference)
+    autoep = AutoEP(
+        candidate,
+        parse_autoep_config({
+            "enabled": True,
+            "autoep_size": 1,
+            "preset_model": "qwen3_5_moe",
+            "use_grouped_mm": False
+        }))
+    specs = autoep.ep_parser()
+    assert len(specs) == 1
+    path = "model.language_model.layers.0.mlp" if wrapper else "model.layers.0.mlp"
+    assert specs[0].moe_module_name == path
+    original = reference.get_submodule(path)
+    autoep.replace_moe_layer(specs[0], ep_size=1, ep_rank=0)
+    replaced = candidate.get_submodule(path)
+    x = torch.randn(13, 128, dtype=torch.float32 if autocast else torch.bfloat16)
+    with torch.autocast("cpu", dtype=torch.bfloat16, enabled=autocast):
+        logits, scores, indices = original.gate(x)
+        actual_scores, actual_indices, counts = replaced.router(x)
+    assert actual_scores.dtype == scores.dtype == torch.bfloat16
+    torch.testing.assert_close(actual_indices, indices, rtol=0, atol=0)
+    torch.testing.assert_close(actual_scores, scores, rtol=0, atol=0)
+    with torch.no_grad(), torch.autocast("cpu", dtype=torch.bfloat16, enabled=autocast):
+        _, returned_logits = replaced(x.unsqueeze(0))
+    torch.testing.assert_close(returned_logits, logits, rtol=0, atol=0)
+    assert counts.sum() == 26
+    ids = torch.tensor([[1, 5, 7, 9, 11]])
+    with torch.no_grad():
+        expected = reference(input_ids=ids, labels=ids)
+        actual = candidate(input_ids=ids, labels=ids)
+    torch.testing.assert_close(actual.loss, expected.loss, rtol=0, atol=0.005)
+    torch.testing.assert_close(actual.logits, expected.logits, rtol=0.02, atol=0.01)