feat(gpt-oss): expert offloading Phase 3 — wire dispatch + device_memory_ratio tunable

core/parallel/expert_module.cpp

@@ -196,6 +196,8 @@ torch::Tensor MoEMLP::forward(torch::Tensor hidden_states,
   DLOG_FATAL_IF(batch_size > kMaxTokens || batch_size <= 0,
                 "batch_size should be (0,", kMaxTokens, "] , but got",
                 batch_size);
+  TORCH_CHECK(hidden_states.scalar_type() == input_.scalar_type(),
+              "hidden_states dtype must match expert input dtype");
 
   // Use async copy with the provided execution stream
   cudaMemcpyAsync(input_.data_ptr(), hidden_states.data_ptr(),

moe_infinity/models/gpt_oss.py

@@ -36,6 +36,13 @@ def __init__(self, config):
 class SyncGptOssMLP(nn.Module):
     archer_config: Optional[ArcherConfig] = None
     layer_id: Optional[int] = None
+    expert_executor = None
+    expert_prefetcher = None
+    expert_tracer = None
+    expert_predictor = None
+    expert_tensor_map = None
+    archer_engine = None
+    lib = None
 
     def __init__(self, config):
         super().__init__()

moe_infinity/runtime/model_offload.py

@@ -174,6 +174,33 @@ def _out_block(experts_prefix: str) -> str:
     "down_proj_bias",
 )
 
+_GPT_OSS_RESIDENT_RATIO = 0.9
+
+
+def _gpt_oss_offload_enabled(model_type, archer_config):
+    if model_type != "gpt_oss":
+        return True
+    return archer_config.device_memory_ratio < _GPT_OSS_RESIDENT_RATIO
+
+
+def _wire_expert_collaborators(
+    module,
+    enabled,
+    *,
+    expert_executor,
+    expert_prefetcher,
+    expert_tracer,
+    expert_predictor,
+    expert_tensor_map,
+):
+    if not enabled:
+        return
+    module.expert_executor = expert_executor
+    module.expert_prefetcher = expert_prefetcher
+    module.expert_tracer = expert_tracer
+    module.expert_predictor = expert_predictor
+    module.expert_tensor_map = expert_tensor_map
+
 
 def _expand_gpt_oss_packed_experts(state_dict, config):
     if getattr(config, "model_type", "") != "gpt_oss":
@@ -542,6 +569,9 @@ def init(
         )
 
         self.archer_config = _archer_config
+        self.gpt_oss_offload_enabled = _gpt_oss_offload_enabled(
+            getattr(self.config, "model_type", ""), self.archer_config
+        )
         if _archer_config.trace_path is not None:
             self.expert_tracer.load_trace(_archer_config.trace_path)
 
@@ -1078,12 +1108,19 @@ def archer_from_pretrained(cls, *args, **kwargs):
                         module.is_gptq = is_gptq_quantized(self.config)
                         self.expert_modules.append(module)
 
-                        if not isinstance(module, SyncGptOssMLP):
-                            module.expert_executor = self.expert_executor
-                            module.expert_prefetcher = self.expert_prefetcher
-                            module.expert_tracer = self.expert_tracer
-                            module.expert_predictor = self.expert_predictor
-                            module.expert_tensor_map = self.expert_tensor_map
+                        use_dispatcher = (
+                            not isinstance(module, SyncGptOssMLP)
+                            or self.gpt_oss_offload_enabled
+                        )
+                        _wire_expert_collaborators(
+                            module,
+                            use_dispatcher,
+                            expert_executor=self.expert_executor,
+                            expert_prefetcher=self.expert_prefetcher,
+                            expert_tracer=self.expert_tracer,
+                            expert_predictor=self.expert_predictor,
+                            expert_tensor_map=self.expert_tensor_map,
+                        )
 
                         module.lib = self.prefetch_lib
 
@@ -1101,7 +1138,10 @@ def archer_from_pretrained(cls, *args, **kwargs):
                         if self._is_shared_expert_param(_name):
                             del self.name_id_map[_name]
 
-                if getattr(self.config, "model_type", "") == "gpt_oss":
+                if (
+                    getattr(self.config, "model_type", "") == "gpt_oss"
+                    and not self.gpt_oss_offload_enabled
+                ):
                     self._load_resident_gpt_oss(model)
 
                 self.setup_archer_hooks(model)

tests/python/unit/test_gpt_oss_mxfp4_dispatch.py

@@ -129,3 +129,137 @@ def test_dequantized_option_a_matches_resident_expert_forward():
     )
     assert ((option_a.float() - golden).abs() <= envelope).all()
     assert ((resident.float() - golden).abs() <= envelope).all()
+
+
+def _native_dispatch(tmp_path, hidden_dtype):
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA required")
+    try:
+        from moe_infinity.runtime.model_offload import _load_prefetch_lib
+
+        prefetch_lib = _load_prefetch_lib()
+    except Exception as exc:
+        pytest.skip(f"native Archer extension unavailable: {exc}")
+
+    from moe_infinity.models.gpt_oss import SyncGptOssMLP
+
+    class Config:
+        hidden_size = 64
+        intermediate_size = 32
+        num_local_experts = 2
+        num_experts_per_tok = 1
+
+    torch.manual_seed(137)
+    hidden = Config.hidden_size
+    intermediate = Config.intermediate_size
+    mlp = SyncGptOssMLP(Config())
+    tensors = [
+        torch.randint(
+            0, 256, (2 * intermediate, hidden // 2), dtype=torch.uint8
+        ),
+        torch.randint(
+            120, 135, (2 * intermediate, hidden // 32), dtype=torch.uint8
+        ),
+        torch.randn(2 * intermediate, dtype=torch.bfloat16),
+        torch.randint(0, 256, (hidden, intermediate // 2), dtype=torch.uint8),
+        torch.randint(
+            120, 135, (hidden, intermediate // 32), dtype=torch.uint8
+        ),
+        torch.randn(hidden, dtype=torch.bfloat16),
+    ]
+    mlp.experts.gate_up_proj.requires_grad_(False)
+    mlp.experts.down_proj.requires_grad_(False)
+    mlp.experts.gate_up_proj.data = tensors[0].unsqueeze(0).repeat(2, 1, 1)
+    mlp.experts.gate_up_proj_scales.data = (
+        tensors[1].unsqueeze(0).repeat(2, 1, 1)
+    )
+    mlp.experts.gate_up_proj_bias.data = tensors[2].unsqueeze(0).repeat(2, 1)
+    mlp.experts.down_proj.data = tensors[3].unsqueeze(0).repeat(2, 1, 1)
+    mlp.experts.down_proj_scales.data = tensors[4].unsqueeze(0).repeat(2, 1, 1)
+    mlp.experts.down_proj_bias.data = tensors[5].unsqueeze(0).repeat(2, 1)
+
+    hidden_states = torch.randn(3, hidden, dtype=hidden_dtype, device="cuda:0")
+    expected = mlp._expert_forward_mxfp4(hidden_states, 0)
+
+    engine = prefetch_lib.prefetch_handle(f"{tmp_path}/", 0.5)
+    expert_tensor_ids = [list(range(6)), list(range(6, 12))]
+    for tensor_id, tensor in zip(expert_tensor_ids[0], tensors):
+        engine.offload(tensor, tensor_id)
+    for tensor_id, tensor in zip(expert_tensor_ids[1], tensors):
+        engine.offload(tensor, tensor_id)
+    dense_tensor_ids = list(range(12, 12 + 2 * torch.cuda.device_count()))
+    for tensor_id in dense_tensor_ids:
+        engine.offload(torch.zeros(1, dtype=torch.bfloat16), tensor_id)
+    split = len(dense_tensor_ids) // 2
+    dense_before = [
+        (f"model.dense.{index}", [[tensor_id]])
+        for index, tensor_id in enumerate(dense_tensor_ids[:split])
+    ]
+    dense_after = [
+        (f"model.dense.{index + split}", [[tensor_id]])
+        for index, tensor_id in enumerate(dense_tensor_ids[split:])
+    ]
+    engine.set_topology(
+        dense_before
+        + [("model.layers.0.mlp.experts", expert_tensor_ids)]
+        + dense_after
+    )
+
+    dispatcher = prefetch_lib.expert_dispatcher(2, 1, 0, 6, 1)
+    dispatcher.register_expert(0, 0, expert_tensor_ids[0], "")
+    dispatcher.register_expert(0, 1, expert_tensor_ids[1], "")
+    torch.cuda.set_device(hidden_states.device)
+    router_mask = torch.zeros((3, 2), dtype=torch.bool, device="cuda:0")
+    router_mask[:, 0] = True
+    router_weights = torch.zeros((3, 2), dtype=torch.bfloat16, device="cuda:0")
+    router_weights[:, 0] = 1
+    dispatcher.set_inputs(
+        hidden_states,
+        router_mask,
+        router_weights,
+    )
+    dispatcher.set_expected_queue(1)
+    dispatcher.enqueue_expert(0, 0, 0, False)
+    dispatcher.notify_fetch_start()
+    actual = dispatcher.wait_expert()
+
+    return actual, expected, hidden_states, tensors
+
+
+@pytest.mark.gpu
+def test_native_dispatcher_matches_resident_expert_forward(tmp_path):
+    actual, resident, hidden_states, tensors = _native_dispatch(
+        tmp_path, torch.bfloat16
+    )
+    from moe_infinity._v4_fp4 import mxfp4_dequant
+
+    gate_weight = mxfp4_dequant(tensors[0].cuda(), tensors[1].cuda()).float()
+    down_weight = mxfp4_dequant(tensors[3].cuda(), tensors[4].cuda()).float()
+    golden_gate_up = (
+        hidden_states.float() @ gate_weight.t() + tensors[2].cuda().float()
+    )
+    golden_gate, golden_up = golden_gate_up[:, ::2], golden_gate_up[:, 1::2]
+    golden_activated = (golden_up.clamp(-7, 7) + 1) * (
+        golden_gate.clamp(max=7)
+        * torch.sigmoid(golden_gate.clamp(max=7) * 1.702)
+    )
+    golden = golden_activated @ down_weight.t() + tensors[5].cuda().float()
+    envelope = (
+        8
+        * (2**-8)
+        * (
+            golden_activated.abs() @ down_weight.abs().t()
+            + tensors[5].cuda().float().abs()
+        )
+        + 1e-2
+    )
+
+    assert ((actual.to(golden.device) - golden).abs() <= envelope).all()
+    assert ((resident.float() - golden).abs() <= envelope).all()
+
+
+@pytest.mark.gpu
+def test_native_dispatcher_rejects_mismatched_hidden_dtype(tmp_path):
+    actual, _, _, _ = _native_dispatch(tmp_path, torch.float32)
+
+    assert torch.count_nonzero(actual).item() == 0

tests/test_gpt_oss_offload_policy.py

@@ -0,0 +1,64 @@
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+from moe_infinity.runtime.model_offload import (
+    _gpt_oss_offload_enabled,
+    _wire_expert_collaborators,
+)
+
+
+def _engine_config(ratio):
+    return SimpleNamespace(device_memory_ratio=ratio)
+
+
+def test_gpt_oss_default_ratio_keeps_resident_fallback():
+    assert not _gpt_oss_offload_enabled("gpt_oss", _engine_config(0.9))
+
+
+def test_gpt_oss_low_ratio_enables_dispatcher():
+    assert _gpt_oss_offload_enabled("gpt_oss", _engine_config(0.5))
+
+
+def test_policy_does_not_disable_other_architectures():
+    assert _gpt_oss_offload_enabled("mixtral", _engine_config(0.9))
+
+
+def test_low_ratio_wires_all_gpt_oss_collaborators():
+    module = SimpleNamespace()
+    collaborators = {
+        "expert_executor": Mock(),
+        "expert_prefetcher": Mock(),
+        "expert_tracer": Mock(),
+        "expert_predictor": Mock(),
+        "expert_tensor_map": {(0, 0): 7},
+    }
+
+    _wire_expert_collaborators(module, True, **collaborators)
+
+    for name, value in collaborators.items():
+        assert getattr(module, name) is value
+
+
+def test_resident_mode_leaves_gpt_oss_collaborators_unset():
+    module = SimpleNamespace(
+        expert_executor=None,
+        expert_prefetcher=None,
+        expert_tracer=None,
+        expert_predictor=None,
+        expert_tensor_map=None,
+    )
+    _wire_expert_collaborators(
+        module,
+        False,
+        expert_executor=Mock(),
+        expert_prefetcher=Mock(),
+        expert_tracer=Mock(),
+        expert_predictor=Mock(),
+        expert_tensor_map={(0, 0): 7},
+    )
+
+    assert module.expert_executor is None
+    assert module.expert_prefetcher is None
+    assert module.expert_tracer is None
+    assert module.expert_predictor is None
+    assert module.expert_tensor_map is None

tests/test_gpt_oss_wrapper.py

@@ -1,7 +1,7 @@
 import importlib.machinery
 import sys
 import types
-from unittest.mock import MagicMock
+from unittest.mock import MagicMock, Mock
 
 import torch
 
@@ -34,6 +34,18 @@ def make_gpt_oss_config():
     return cfg
 
 
+def _deterministic_mlp():
+    from moe_infinity.models.gpt_oss import SyncGptOssMLP
+
+    torch.manual_seed(137)
+    mlp = SyncGptOssMLP(make_gpt_oss_config())
+    for parameter in mlp.parameters():
+        if parameter.numel():
+            torch.nn.init.normal_(parameter, std=0.01)
+    mlp.layer_id = 1
+    return mlp
+
+
 def test_sync_gpt_oss_mlp_instantiation():
     from moe_infinity.models.gpt_oss import SyncGptOssMLP
 
@@ -104,3 +116,59 @@ def test_sync_gpt_oss_mlp_router_has_bias():
     assert mlp.router.bias.shape == (
         config.num_local_experts,
     ), f"Router bias shape must be ({config.num_local_experts},)"
+
+
+def test_sync_gpt_oss_mlp_dispatches_when_executor_is_injected():
+    mlp = _deterministic_mlp()
+    hidden = torch.randn(1, 2, 64)
+    expected = torch.randn(2, 64)
+    executor = Mock()
+    executor.wait_dispatch_local.return_value = expected
+    mlp.expert_executor = executor
+
+    output, router_logits = mlp(hidden)
+
+    assert output.shape == hidden.shape
+    assert torch.equal(output.view(2, 64), expected)
+    executor.dispatch_local.assert_called_once()
+    args = executor.dispatch_local.call_args
+    assert args.args[0] == 1
+    assert args.args[1].shape == (2, 64)
+    assert args.args[2].shape == (2, 4)
+    assert args.args[3].shape == (2, 4)
+    assert torch.equal(args.kwargs["router_logits"], router_logits)
+    executor.wait_dispatch_local.assert_called_once_with()
+
+
+def test_sync_gpt_oss_mlp_resident_path_matches_existing_forward():
+    mlp = _deterministic_mlp()
+    hidden = torch.randn(1, 2, 64)
+
+    with torch.no_grad():
+        output, router_logits = mlp(hidden)
+        hidden_flat = hidden.view(-1, hidden.shape[-1])
+        routing_weights = torch.softmax(
+            router_logits, dim=-1, dtype=torch.float32
+        )
+        routing_weights, selected_experts = torch.topk(
+            routing_weights, mlp.top_k, dim=-1
+        )
+        routing_weights = routing_weights / routing_weights.sum(
+            dim=-1, keepdim=True
+        )
+        expected = torch.zeros_like(hidden_flat)
+        for expert_idx in range(mlp.num_experts):
+            token_positions, topk_positions = torch.where(
+                selected_experts == expert_idx
+            )
+            if token_positions.numel() == 0:
+                continue
+            expert_output = mlp._expert_forward(
+                hidden_flat[token_positions], expert_idx
+            )
+            weights = routing_weights[token_positions, topk_positions].to(
+                hidden.dtype
+            )
+            expected[token_positions] += weights.unsqueeze(-1) * expert_output
+
+    assert torch.equal(output, expected.view_as(hidden))