From 5ddd917528e5c1667fa2763f2b1aff79c154ccb2 Mon Sep 17 00:00:00 2001 From: drunkcoding Date: Tue, 11 Aug 2026 14:31:41 +0000 Subject: [PATCH 1/5] fix(dflash): place native forward + speculator on the resident GPU Serving Qwen3.5-MoE with a DFlash drafter on more than one GPU hung during generation. The offloaded text backbone (embed_tokens/attention/ lm_head) is resident on the LAST visible GPU (mirrors ContinuousBatchingEngine._resolve_device), but the DFlash native path assumed cuda:0: - MoE._native_model_forward{,_rich} placed inputs on cuda:0, so the very first target forward's embed_tokens(input_ids) raised a device mismatch (index on cuda:0 vs weight on cuda:N). - DFlashSpeculator._infer_cuda_device returned the first cuda parameter (or cuda:0), so the drafter and its block landed on a different GPU than the bound shared embed_tokens/lm_head. Both RuntimeErrors propagated out of engine.step() and were swallowed by the async engine loop, so the client saw an indefinite hang (HTTP 000) instead of an error. Resolve the native input device from the resident input embedding (fallback cuda:{last_gpu}); infer the speculator device from the bound shared embedding. Verified end-to-end on 2 GPUs: Qwen3.5-35B-A3B target + z-lab/Qwen3.5-35B-A3B-DFlash drafter + ContextPilot generate over api_server_v2. dflash (264) / qwen3_5_moe (16) / contextpilot (75) suites stay green. Refs: #146 --- moe_infinity/entrypoints/big_modeling.py | 54 +++++++++++++++++------- moe_infinity/spec_decode/dflash.py | 17 +++++++- 2 files changed, 54 insertions(+), 17 deletions(-) diff --git a/moe_infinity/entrypoints/big_modeling.py b/moe_infinity/entrypoints/big_modeling.py index 36169f94..a0246459 100644 --- a/moe_infinity/entrypoints/big_modeling.py +++ b/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: diff --git a/moe_infinity/spec_decode/dflash.py b/moe_infinity/spec_decode/dflash.py index 4fc6df6d..6f2402cf 100644 --- a/moe_infinity/spec_decode/dflash.py +++ b/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( From 54261d5c4cfe83a51083ca87e04f9fc8749046df Mon Sep 17 00:00:00 2001 From: drunkcoding Date: Tue, 11 Aug 2026 14:31:49 +0000 Subject: [PATCH 2/5] fix(serving): resolve Qwen3.5-MoE arch dims via get_text_config Qwen3.5-MoE ships a vision-language config that nests the text backbone dimensions under text_config, so _build_engine_config raised "RuntimeError: unable to resolve model num_layers". Resolve num_layers/heads/kv_heads/head_dim/eos via config.get_text_config(), which returns self for text-only configs so existing checkpoints are unaffected. Verified: num_layers=40, num_kv_heads=2, head_dim=256, eos=248044. Refs: #146 --- .../entrypoints/openai/api_server_v2.py | 25 +++++++++++++------ 1 file changed, 18 insertions(+), 7 deletions(-) diff --git a/moe_infinity/entrypoints/openai/api_server_v2.py b/moe_infinity/entrypoints/openai/api_server_v2.py index 48c3d7ca..eb0fb395 100644 --- a/moe_infinity/entrypoints/openai/api_server_v2.py +++ b/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 From 8ae26c8e2a9542f4c34c1db25315b5db576fdcf1 Mon Sep 17 00:00:00 2001 From: drunkcoding Date: Tue, 11 Aug 2026 14:31:49 +0000 Subject: [PATCH 3/5] test(contextpilot): skip live tests when optional package is absent Mirror test_tokenizer_compat's importorskip: the CP import test and the live-middleware tests now skip (instead of failing) when the optional contextpilot package is not installed, so the suite is green with or without it. --- tests/python/contextpilot/test_cp_import.py | 4 ++++ tests/python/contextpilot/test_middleware.py | 10 ++++++++++ 2 files changed, 14 insertions(+) diff --git a/tests/python/contextpilot/test_cp_import.py b/tests/python/contextpilot/test_cp_import.py index 362e207e..1b529805 100644 --- a/tests/python/contextpilot/test_cp_import.py +++ b/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") diff --git a/tests/python/contextpilot/test_middleware.py b/tests/python/contextpilot/test_middleware.py index 8efc6e8a..711f0724 100644 --- a/tests/python/contextpilot/test_middleware.py +++ b/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, From 80adb3d9f2ac1023d35bf3a808685995bd9bd3f5 Mon Sep 17 00:00:00 2001 From: drunkcoding Date: Tue, 11 Aug 2026 18:20:51 +0000 Subject: [PATCH 4/5] fix(ci): skip GLM routing tests when transformers drops route_tokens_to_experts transformers 5.15 removed GlmMoeDsaMoE.route_tokens_to_experts, which glm_moe_dsa.py and test_glm_routing.py rely on. CI installs the latest 5.x (requirements pin transformers>=5.3.0,<6), so unit-tests (3.10)/(3.12) fail on dev (and every open PR) with "AttributeError: type object 'GlmMoeDsaMoE' has no attribute route_tokens_to_experts". Skip the GLM routing module when the method is absent (mirrors the existing importorskip guard), and resolve it via getattr so the block raises a clear, actionable error at routing time instead of a cryptic AttributeError at construction. Refs: #146 --- moe_infinity/models/glm_moe_dsa.py | 9 ++++++++- tests/python/unit/test_glm_routing.py | 10 ++++++++++ 2 files changed, 18 insertions(+), 1 deletion(-) diff --git a/moe_infinity/models/glm_moe_dsa.py b/moe_infinity/models/glm_moe_dsa.py index de9394a2..ef7bed3c 100644 --- a/moe_infinity/models/glm_moe_dsa.py +++ b/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): diff --git a/tests/python/unit/test_glm_routing.py b/tests/python/unit/test_glm_routing.py index 1dd3e42c..abb1d293 100644 --- a/tests/python/unit/test_glm_routing.py +++ b/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 ( From 3cb94cccbe5c87e433e9563bdabbca95ecb34b65 Mon Sep 17 00:00:00 2001 From: drunkcoding Date: Tue, 11 Aug 2026 18:26:17 +0000 Subject: [PATCH 5/5] fix(test): accept speculative_draft in watchdog _FakeRuntimeEngine mock _initialize_model passes speculative_draft to ContinuousBatchingEngine (since the DFlash integration), but the watchdog test's _FakeRuntimeEngine mock never accepted it. On transformers 5.15, once the earlier GLM failure stops masking it (pytest runs fail-fast), test_watchdog_integration fails with "unexpected keyword argument 'speculative_draft'". Match the real ContinuousBatchingEngine signature (same fix already in #148). Refs: #146 --- tests/python/unit/test_watchdog_integration.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/python/unit/test_watchdog_integration.py b/tests/python/unit/test_watchdog_integration.py index b3ef2d6f..46d75368 100644 --- a/tests/python/unit/test_watchdog_integration.py +++ b/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: