fix(dflash): unblock Qwen3.5-MoE + DFlash + ContextPilot serving hang (multi-GPU device placement)

moe_infinity/entrypoints/big_modeling.py

@@ -491,18 +491,47 @@ def _build_native_components(
             "generation_engine": generation_engine,
         }
 
-    def _native_model_forward(
-        self, token_ids: list[int], _attention_metadata: object
-    ) -> torch.Tensor:
-        input_tensor = torch.tensor([token_ids], dtype=torch.long)
-        if torch.cuda.is_available():
-            input_tensor = input_tensor.to("cuda:0")
-        else:
+    def _resolve_native_input_device(self) -> torch.device:
+        """Input device for the native forward (mirrors engine._resolve_device).
+
+        The OffloadEngine-managed backbone is resident on the LAST visible GPU,
+        so hard-coding ``cuda:0`` mismatches ``embed_tokens`` on multi-GPU runs.
+        """
+        if not torch.cuda.is_available():
             model_device = getattr(self.model, "device", None)
             if isinstance(
                 model_device, torch.device
             ) and model_device.type not in ("meta", "cpu"):
-                input_tensor = input_tensor.to(model_device)
+                return model_device
+            return torch.device("cpu")
+
+        get_embed = getattr(self.model, "get_input_embeddings", None)
+        if callable(get_embed):
+            try:
+                weight = getattr(get_embed(), "weight", None)
+                embed_device = getattr(weight, "device", None)
+            except Exception:
+                embed_device = None
+            if (
+                isinstance(embed_device, torch.device)
+                and embed_device.type == "cuda"
+            ):
+                return embed_device
+
+        model_device = getattr(self.model, "device", None)
+        if (
+            isinstance(model_device, torch.device)
+            and model_device.type == "cuda"
+        ):
+            return model_device
+
+        return torch.device(f"cuda:{torch.cuda.device_count() - 1}")
+
+    def _native_model_forward(
+        self, token_ids: list[int], _attention_metadata: object
+    ) -> torch.Tensor:
+        input_tensor = torch.tensor([token_ids], dtype=torch.long)
+        input_tensor = input_tensor.to(self._resolve_native_input_device())
 
         is_prefill = True
         if _attention_metadata is not None:
@@ -595,14 +624,7 @@ def _native_model_forward_rich(
         preserved — nothing here bypasses expert dispatch.
         """
         input_tensor = torch.tensor([token_ids], dtype=torch.long)
-        if torch.cuda.is_available():
-            input_tensor = input_tensor.to("cuda:0")
-        else:
-            model_device = getattr(self.model, "device", None)
-            if isinstance(
-                model_device, torch.device
-            ) and model_device.type not in ("meta", "cpu"):
-                input_tensor = input_tensor.to(model_device)
+        input_tensor = input_tensor.to(self._resolve_native_input_device())
 
         is_prefill = True
         if _attention_metadata is not None:

moe_infinity/entrypoints/openai/api_server_v2.py

@@ -1775,34 +1775,41 @@ def _build_engine_config(
         raise RuntimeError(
             "model config is required to initialize serving engine"
         )
+    # Multimodal MoE checkpoints (e.g. Qwen3.5-MoE VL) nest the text backbone
+    # dimensions under text_config; get_text_config() returns self otherwise.
+    text_config = (
+        model_config.get_text_config()
+        if hasattr(model_config, "get_text_config")
+        else model_config
+    )
 
     num_layers = _resolve_int_attr(
-        model_config,
+        text_config,
         "num_hidden_layers",
         "num_layers",
         "n_layer",
     )
     num_attention_heads = _resolve_int_attr(
-        model_config,
+        text_config,
         "num_attention_heads",
         "n_head",
     )
     num_kv_heads = _resolve_int_attr(
-        model_config,
+        text_config,
         "num_key_value_heads",
         "num_kv_heads",
         "n_head_kv",
     )
-    hidden_size = _resolve_int_attr(model_config, "hidden_size", "n_embd")
+    hidden_size = _resolve_int_attr(text_config, "hidden_size", "n_embd")
     max_seq_length = _resolve_int_attr(
-        model_config,
+        text_config,
         "max_position_embeddings",
         "max_seq_len",
         "max_sequence_length",
         "n_positions",
         "model_max_length",
     )
-    head_dim = _resolve_int_attr(model_config, "head_dim")
+    head_dim = _resolve_int_attr(text_config, "head_dim")
 
     if num_layers is None:
         raise RuntimeError("unable to resolve model num_layers")
@@ -1819,7 +1826,11 @@ def _build_engine_config(
         model_config, "eos_token_id"
     )
     if eos_token_id is None:
-        config_eos = getattr(model_config, "eos_token_id", None)
+        eos_token_id = _resolve_int_attr(text_config, "eos_token_id")
+    if eos_token_id is None:
+        config_eos = getattr(model_config, "eos_token_id", None) or getattr(
+            text_config, "eos_token_id", None
+        )
         if (
             isinstance(config_eos, list)
             and config_eos

moe_infinity/models/glm_moe_dsa.py

@@ -59,7 +59,9 @@ def __init__(self, config):
             intermediate_size=config.moe_intermediate_size
             * config.n_shared_experts,
         )
-        self._hf_route_tokens = GlmMoeDsaMoE.route_tokens_to_experts
+        self._hf_route_tokens = getattr(
+            GlmMoeDsaMoE, "route_tokens_to_experts", None
+        )
 
     def _route(self, hidden_flat: torch.Tensor):
         dev = hidden_flat.device
@@ -68,6 +70,11 @@ def _route(self, hidden_flat: torch.Tensor):
                 self.gate.e_score_correction_bias.to(dev)
             )
         router_logits = self.gate(hidden_flat)
+        if self._hf_route_tokens is None:
+            raise RuntimeError(
+                "GLM-MoE-DSA routing requires a transformers build providing "
+                "GlmMoeDsaMoE.route_tokens_to_experts (removed in 5.15+)"
+            )
         return self._hf_route_tokens(self, router_logits)
 
     def _local_experts(self, hidden_flat, router_mask, routing_weights_mask):

moe_infinity/spec_decode/dflash.py

@@ -205,13 +205,28 @@ def _infer_cuda_device(model: Any) -> str:
     dev = getattr(model, "device", None)
     if isinstance(dev, torch.device) and dev.type == "cuda":
         return str(dev)
+    # Match the bound shared embed_tokens/lm_head device: an offloaded backbone
+    # is resident on the LAST visible GPU, so first-cuda-param/cuda:0 would put
+    # the drafter's block on a different GPU than its shared weights.
+    try:
+        embed = _resolve_input_embeddings(model)
+        embed_device = getattr(getattr(embed, "weight", None), "device", None)
+        if (
+            isinstance(embed_device, torch.device)
+            and embed_device.type == "cuda"
+        ):
+            return str(embed_device)
+    except Exception:
+        pass
     try:
         for param in model.parameters():
             if param.device.type == "cuda":
                 return str(param.device)
     except Exception:
         pass
-    return "cuda:0" if torch.cuda.is_available() else "cpu"
+    if torch.cuda.is_available():
+        return f"cuda:{torch.cuda.device_count() - 1}"
+    return "cpu"
 
 
 def _resolve_stop_ids(

tests/python/contextpilot/test_cp_import.py

@@ -1,5 +1,9 @@
 import importlib
 
+import pytest
+
+pytest.importorskip("contextpilot", reason="contextpilot package not installed")
+
 
 def test_contextpilot_import_and_instantiation_with_moe_infinity():
     contextpilot = importlib.import_module("contextpilot")

tests/python/contextpilot/test_middleware.py

@@ -3,11 +3,18 @@
 import threading
 import time
 
+import pytest
 from _pytest.monkeypatch import MonkeyPatch
 
 import moe_infinity.serving.contextpilot_middleware as middleware_module
 from moe_infinity.serving.contextpilot_middleware import ContextPilotMiddleware
 
+# Skip (not fail) live-middleware tests that need the optional real package.
+requires_contextpilot = pytest.mark.skipif(
+    middleware_module.ContextPilot is None,
+    reason="contextpilot package not installed",
+)
+
 
 def test_process_chat_request_returns_messages(
     monkeypatch: MonkeyPatch,
@@ -185,6 +192,7 @@ def test_on_request_complete_doesnt_raise() -> None:
     middleware.on_request_complete("request-123")
 
 
+@requires_contextpilot
 def test_is_enabled_respects_flag() -> None:
     disabled = ContextPilotMiddleware(enabled=False)
     enabled = ContextPilotMiddleware(enabled=True)
@@ -246,6 +254,7 @@ def optimize(
     assert stats["total_tokens_saved"] > 0
 
 
+@requires_contextpilot
 def test_dedup_without_reorder() -> None:
     middleware = ContextPilotMiddleware(
         use_gpu=False,
@@ -270,6 +279,7 @@ def test_dedup_without_reorder() -> None:
     assert stats["total_tokens_saved"] > 0
 
 
+@requires_contextpilot
 def test_token_savings_tracked() -> None:
     middleware = ContextPilotMiddleware(
         use_gpu=False,

tests/python/unit/test_glm_routing.py

@@ -8,6 +8,16 @@
     reason="transformers >= 5.12 required",
 )
 
+from transformers.models.glm_moe_dsa.modeling_glm_moe_dsa import (  # noqa: E402
+    GlmMoeDsaMoE as _GlmMoeDsaMoE,
+)
+
+if not hasattr(_GlmMoeDsaMoE, "route_tokens_to_experts"):
+    pytest.skip(
+        "transformers dropped GlmMoeDsaMoE.route_tokens_to_experts (5.15+)",
+        allow_module_level=True,
+    )
+
 
 def _tiny_config():
     from transformers.models.glm_moe_dsa.modeling_glm_moe_dsa import (

tests/python/unit/test_watchdog_integration.py

@@ -43,11 +43,13 @@ def __init__(
         engine: object,
         config: dict[str, object],
         tokenizer: object,
+        speculative_draft: object = None,
     ) -> None:
         self.model = model
         self.engine = engine
         self.config = config
         self.tokenizer = tokenizer
+        self.speculative_draft = speculative_draft
 
 
 class _FakeMoE: