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/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 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/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( 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, 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 ( 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: