From 7f23a626e9a6ea6991f4b9da276d622080fe835a Mon Sep 17 00:00:00 2001
From: janmenjayap <janmenjaya.panda.22@gmail.com>
Date: Sun, 11 Oct 2026 20:51:34 +0530
Subject: [PATCH 1/8] fix(moe): avoid affected Granite grouped experts on MPS

---
 .../test_granite_grouped_expert_dispatch.py   |  81 +++++++++
 .../test_moe_expert_dispatch.py               | 159 ++++++++++++++++++
 .../generalized_components/moe.py             |  57 +++++++
 3 files changed, 297 insertions(+)
 create mode 100644 tests/mps/test_granite_grouped_expert_dispatch.py
 create mode 100644 tests/unit/model_bridge/generalized_components/test_moe_expert_dispatch.py

diff --git a/tests/mps/test_granite_grouped_expert_dispatch.py b/tests/mps/test_granite_grouped_expert_dispatch.py
new file mode 100644
index 0000000000..870a287429
--- /dev/null
+++ b/tests/mps/test_granite_grouped_expert_dispatch.py
@@ -0,0 +1,81 @@
+"""Granite expert correctness on MPS with fixed routes and an independent oracle."""
+
+import warnings
+from types import SimpleNamespace
+
+import pytest
+import torch
+import transformers
+from packaging.version import Version
+from torch import nn
+from torch.nn import functional as F
+from transformers.models.granitemoehybrid.modeling_granitemoehybrid import (
+    GraniteMoeHybridExperts,
+)
+
+from transformer_lens.model_bridge.generalized_components.moe import MoEBridge
+
+pytestmark = pytest.mark.skipif(not torch.backends.mps.is_available(), reason="MPS not available")
+
+
+class _FixedRoutes(nn.Module):
+    def __init__(self):
+        super().__init__()
+        config = SimpleNamespace(
+            hidden_size=64, intermediate_size=32, num_local_experts=4, hidden_act="silu"
+        )
+        config._experts_implementation = "grouped_mm"
+        self.experts = GraniteMoeHybridExperts(config)
+        generator = torch.Generator().manual_seed(0)
+        with torch.no_grad():
+            for parameter in self.experts.parameters():
+                parameter.copy_(torch.randn(parameter.shape, generator=generator) * 0.1)
+        self.register_buffer(
+            "indices", torch.tensor([[0], [0], [1], [1], [1], [3], [3], [3], [3], [3]])
+        )
+        self.register_buffer("weights", torch.linspace(0.2, 1.0, 10).unsqueeze(1))
+
+    def forward(self, layer_input):
+        return self.experts(layer_input, self.indices, self.weights)
+
+
+def _oracle(block, inputs):
+    contributions = []
+    for token, indices, weights in zip(inputs, block.indices, block.weights):
+        routes = []
+        for index, weight in zip(indices, weights):
+            gate, up = F.linear(token, block.experts.gate_up_proj[index]).chunk(2, dim=-1)
+            routes.append(F.linear(F.silu(gate) * up, block.experts.down_proj[index]) * weight)
+        contributions.append(torch.stack(routes).sum(0))
+    return torch.stack(contributions)
+
+
+def test_fixed_route_experts_match_independent_oracle_and_repeat():
+    block = _FixedRoutes()
+    inputs = torch.randn(10, 64, generator=torch.Generator().manual_seed(1))
+    with torch.no_grad():
+        expected = _oracle(block, inputs)
+    block.to("mps")
+    bridge = MoEBridge("moe")
+    bridge.set_original_component(block)
+    observed = []
+    handle = block.experts.register_forward_hook(
+        lambda module, args, output: observed.append(output)
+    )
+    try:
+        with torch.no_grad(), warnings.catch_warnings(record=True) as caught:
+            warnings.simplefilter("always")
+            first = bridge(inputs.to("mps"))
+            second = bridge(layer_input=inputs.to("mps"))
+        assert first.device.type == "mps"
+        torch.testing.assert_close(first.cpu(), expected, rtol=1e-4, atol=1e-6)
+        torch.testing.assert_close(second, first, rtol=0, atol=0)
+        if (
+            Version(torch.__version__).public == "2.11.0"
+            and Version(transformers.__version__).public == "5.13.0"
+        ):
+            assert any("eager Granite" in str(warning.message) for warning in caught)
+        assert len(observed) == 2
+        assert block.experts.config._experts_implementation == "grouped_mm"
+    finally:
+        handle.remove()
diff --git a/tests/unit/model_bridge/generalized_components/test_moe_expert_dispatch.py b/tests/unit/model_bridge/generalized_components/test_moe_expert_dispatch.py
new file mode 100644
index 0000000000..a9ce9d6eda
--- /dev/null
+++ b/tests/unit/model_bridge/generalized_components/test_moe_expert_dispatch.py
@@ -0,0 +1,159 @@
+"""Scoped Granite expert dispatch without shared configuration mutation."""
+
+import copy
+import warnings
+from concurrent.futures import ThreadPoolExecutor
+from functools import wraps
+from types import SimpleNamespace
+
+import pytest
+import torch
+from torch import nn
+from transformers.models.granitemoehybrid.modeling_granitemoehybrid import (
+    GraniteMoeHybridExperts,
+)
+
+from transformer_lens.model_bridge.generalized_components.moe import (
+    _granite_mps_experts_forward,
+    _guard_granite_mps_experts,
+)
+
+
+class _MpsInput(torch.Tensor):
+    @property
+    def device(self):
+        return torch.device("mps")
+
+
+class _Experts(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.fail = False
+
+    def eager(self, hidden_states):
+        if self.fail:
+            raise ValueError("expert failure")
+        return "eager"
+
+    @wraps(eager)
+    def forward(self, hidden_states):
+        return self.config._experts_implementation
+
+
+def test_dispatch_does_not_mutate_shared_config():
+    config = SimpleNamespace(_experts_implementation="grouped_mm")
+    experts = [_Experts(config), _Experts(config)]
+    mps = torch.empty(1).as_subclass(_MpsInput)
+    cpu = torch.empty(1)
+    with pytest.warns(RuntimeWarning, match="eager Granite"):
+        with ThreadPoolExecutor(max_workers=2) as pool:
+            results = list(
+                pool.map(
+                    lambda pair: _granite_mps_experts_forward(*pair),
+                    [(experts[0], mps), (experts[1], cpu)],
+                )
+            )
+    assert results == ["eager", "grouped_mm"]
+    assert vars(config) == {"_experts_implementation": "grouped_mm"}
+
+
+@pytest.mark.parametrize("policy", ["eager", "batched_mm", None])
+def test_other_policies_are_unchanged(policy):
+    experts = _Experts(SimpleNamespace(_experts_implementation=policy))
+    assert _granite_mps_experts_forward(experts, torch.empty(1).as_subclass(_MpsInput)) == policy
+
+
+def test_exception_does_not_leave_dispatch_state():
+    experts = _Experts(SimpleNamespace(_experts_implementation="grouped_mm"))
+    experts.fail = True
+    with pytest.warns(RuntimeWarning), pytest.raises(ValueError, match="expert failure"):
+        _granite_mps_experts_forward(experts, hidden_states=torch.empty(1).as_subclass(_MpsInput))
+    assert experts.config._experts_implementation == "grouped_mm"
+    experts.fail = False
+    with pytest.warns(RuntimeWarning):
+        assert (
+            _granite_mps_experts_forward(experts, torch.empty(1).as_subclass(_MpsInput)) == "eager"
+        )
+
+
+def _block():
+    config = SimpleNamespace(
+        hidden_size=8, intermediate_size=4, num_local_experts=4, hidden_act="silu"
+    )
+    config._experts_implementation = "grouped_mm"
+    block = nn.Module()
+    block.experts = GraniteMoeHybridExperts(config)
+    return block
+
+
+@pytest.mark.parametrize(
+    "torch_version,hf_version,installed",
+    [
+        ("2.11.0", "5.13.0", True),
+        ("2.11.0+local", "5.13.0+local", True),
+        ("2.11.1", "5.13.0", False),
+        ("2.12.0", "5.13.0", False),
+        ("2.11.0", "5.14.0", False),
+        ("2.11.0rc1", "5.13.0", False),
+        ("2.11.0", "5.13.0.dev0", False),
+    ],
+)
+def test_installation_version_scope(monkeypatch, torch_version, hf_version, installed):
+    import transformers
+
+    monkeypatch.setattr(torch, "__version__", torch_version)
+    monkeypatch.setattr(transformers, "__version__", hf_version)
+    block = _block()
+    _guard_granite_mps_experts(block)
+    assert (block.experts.forward.__func__ is _granite_mps_experts_forward) == installed
+    assert block.experts.config._experts_implementation == "grouped_mm"
+
+
+def test_installation_is_idempotent_and_deepcopy_safe(monkeypatch):
+    import transformers
+
+    monkeypatch.setattr(torch, "__version__", "2.11.0")
+    monkeypatch.setattr(transformers, "__version__", "5.13.0")
+    block = _block()
+    _guard_granite_mps_experts(block)
+    forward = block.experts.forward
+    _guard_granite_mps_experts(block)
+    assert block.experts.forward is forward
+    cloned = copy.deepcopy(block)
+    assert cloned.experts.forward.__self__ is cloned.experts
+
+
+def test_guarded_cpu_experts_keep_grouped_execution_without_warning(monkeypatch):
+    import transformers
+
+    monkeypatch.setattr(torch, "__version__", "2.11.0")
+    monkeypatch.setattr(transformers, "__version__", "5.13.0")
+    block = _block()
+    with torch.no_grad():
+        for parameter in block.experts.parameters():
+            parameter.normal_(std=0.1)
+    inputs = torch.randn(5, 8)
+    indices = torch.tensor([[0], [0], [1], [3], [3]])
+    weights = torch.ones(5, 1)
+    with torch.no_grad():
+        expected = block.experts(inputs, indices, weights)
+    _guard_granite_mps_experts(block)
+    with warnings.catch_warnings(record=True) as caught, torch.no_grad():
+        warnings.simplefilter("always")
+        actual = block.experts(inputs, indices, weights)
+    assert not caught
+    torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+    assert block.experts.config._experts_implementation == "grouped_mm"
+
+
+def test_unrelated_experts_and_custom_forwards_are_unchanged():
+    block = nn.Module()
+    block.experts = _Experts(SimpleNamespace(_experts_implementation="grouped_mm"))
+    _guard_granite_mps_experts(block)
+    assert block.experts.forward.__func__ is _Experts.forward
+    block = _block()
+    custom = lambda hidden_states: hidden_states
+    block.experts.forward = custom
+    _guard_granite_mps_experts(block)
+    assert block.experts.forward is custom
diff --git a/transformer_lens/model_bridge/generalized_components/moe.py b/transformer_lens/model_bridge/generalized_components/moe.py
index ca96426d30..dc50c0441d 100644
--- a/transformer_lens/model_bridge/generalized_components/moe.py
+++ b/transformer_lens/model_bridge/generalized_components/moe.py
@@ -6,9 +6,12 @@
 from __future__ import annotations
 
 import logging
+import warnings
+from types import MethodType
 from typing import Any, Dict, List, Mapping, Optional, Tuple
 
 import torch
+from packaging.version import Version
 from torch import nn
 
 from transformer_lens.hook_points import HookPoint
@@ -30,6 +33,59 @@
 _BATCHED_PROJECTIONS = _BATCHED_INPUT_PROJECTIONS | _BATCHED_OUTPUT_PROJECTIONS
 
 
+def _granite_mps_experts_forward(experts: nn.Module, *args: Any, **kwargs: Any) -> Any:
+    """Dispatch locally, without changing the config shared by other expert layers."""
+    hidden_states = args[0] if args else kwargs.get("hidden_states")
+    forward = type(experts).forward
+    config: Any = getattr(experts, "config", None)
+    if (
+        isinstance(hidden_states, torch.Tensor)
+        and hidden_states.device.type == "mps"
+        and getattr(config, "_experts_implementation", None) == "grouped_mm"
+    ):
+        warnings.warn(
+            "Using eager Granite MoE Hybrid experts on MPS to avoid incorrect grouped "
+            "matrix products in torch 2.11.0; expert execution may be slower.",
+            RuntimeWarning,
+            stacklevel=2,
+        )
+        return getattr(forward, "__wrapped__")(experts, *args, **kwargs)
+    return forward(experts, *args, **kwargs)
+
+
+def _guard_granite_mps_experts(component: nn.Module) -> None:
+    """Install instance-local dispatch only for the verified HF expert implementation.
+
+    The bound method preserves nn.Module hooks and deepcopy semantics. Configs and
+    global dispatch tables remain untouched, including during concurrent forwards.
+    """
+    experts = getattr(component, "experts", None)
+    if not isinstance(experts, nn.Module):
+        return
+    experts = unwrap_bridge(experts)
+    if type(experts).__module__ != (
+        "transformers.models.granitemoehybrid.modeling_granitemoehybrid"
+    ):
+        return
+    import transformers
+    from transformers.models.granitemoehybrid.modeling_granitemoehybrid import (
+        GraniteMoeHybridExperts,
+    )
+
+    if (
+        type(experts) is not GraniteMoeHybridExperts
+        or Version(torch.__version__).public != "2.11.0"
+        or Version(transformers.__version__).public != "5.13.0"
+    ):
+        return
+    # Respect custom instance forwards rather than silently replacing them.
+    if getattr(experts.forward, "__func__", None) is not GraniteMoeHybridExperts.forward:
+        return
+    if not callable(getattr(GraniteMoeHybridExperts.forward, "__wrapped__", None)):
+        raise RuntimeError("Granite eager expert forward is unavailable for the MPS workaround")
+    setattr(experts, "forward", MethodType(_granite_mps_experts_forward, experts))
+
+
 class UnfoldableMoEParameter(Exception):
     """A parameter whose relationship to the MoE block's input cannot be established."""
 
@@ -431,6 +487,7 @@ def forward(self, *args: Any, **kwargs: Any) -> Any:
         elif "hidden_states" in kwargs:
             hooked = self.hook_in(kwargs["hidden_states"])
             kwargs = {**kwargs, "hidden_states": hooked}
+        _guard_granite_mps_experts(self.original_component)
         output = self.original_component(*args, **kwargs)
         if isinstance(output, tuple):
             if not output:

From 72a3833772c119527745efc1c3cf5a0f8f4028a0 Mon Sep 17 00:00:00 2001
From: janmenjayap <janmenjaya.panda.22@gmail.com>
Date: Sun, 11 Oct 2026 20:56:10 +0530
Subject: [PATCH 2/8] test(moe): cover Granite expert dispatch migration and
 state

---
 .../test_granite_grouped_expert_dispatch.py   | 154 +++++++++++++++++-
 .../test_moe_expert_dispatch.py               |  54 ++++++
 2 files changed, 204 insertions(+), 4 deletions(-)

diff --git a/tests/mps/test_granite_grouped_expert_dispatch.py b/tests/mps/test_granite_grouped_expert_dispatch.py
index 870a287429..6413c2e561 100644
--- a/tests/mps/test_granite_grouped_expert_dispatch.py
+++ b/tests/mps/test_granite_grouped_expert_dispatch.py
@@ -1,5 +1,6 @@
 """Granite expert correctness on MPS with fixed routes and an independent oracle."""
 
+import copy
 import warnings
 from types import SimpleNamespace
 
@@ -35,13 +36,15 @@ def __init__(self):
         )
         self.register_buffer("weights", torch.linspace(0.2, 1.0, 10).unsqueeze(1))
 
-    def forward(self, layer_input):
-        return self.experts(layer_input, self.indices, self.weights)
+    def forward(self, layer_input, top_k_weights=None):
+        weights = self.weights if top_k_weights is None else top_k_weights
+        return self.experts(layer_input, self.indices, weights)
 
 
-def _oracle(block, inputs):
+def _oracle(block, inputs, top_k_weights=None):
     contributions = []
-    for token, indices, weights in zip(inputs, block.indices, block.weights):
+    route_weights = block.weights if top_k_weights is None else top_k_weights
+    for token, indices, weights in zip(inputs, block.indices, route_weights):
         routes = []
         for index, weight in zip(indices, weights):
             gate, up = F.linear(token, block.experts.gate_up_proj[index]).chunk(2, dim=-1)
@@ -79,3 +82,146 @@ def test_fixed_route_experts_match_independent_oracle_and_repeat():
         assert block.experts.config._experts_implementation == "grouped_mm"
     finally:
         handle.remove()
+
+
+def test_dispatch_follows_cpu_mps_cpu_migration():
+    block = _FixedRoutes()
+    bridge = MoEBridge("moe")
+    bridge.set_original_component(block)
+    inputs = torch.randn(10, 64, generator=torch.Generator().manual_seed(1))
+    with torch.no_grad(), warnings.catch_warnings(record=True) as caught:
+        warnings.simplefilter("always")
+        cpu_before = bridge(inputs)
+    assert not caught
+    expected = _oracle(block, inputs)
+    torch.testing.assert_close(cpu_before, expected, rtol=1e-4, atol=1e-6)
+    bridge.to("mps")
+    with torch.no_grad():
+        metal = bridge(inputs.to("mps"))
+    assert metal.device.type == "mps"
+    torch.testing.assert_close(metal.cpu(), expected, rtol=1e-4, atol=1e-6)
+    bridge.to("cpu")
+    with torch.no_grad(), warnings.catch_warnings(record=True) as caught:
+        warnings.simplefilter("always")
+        cpu_after = bridge(inputs)
+    assert not caught
+    torch.testing.assert_close(cpu_after, cpu_before, rtol=0, atol=0)
+    assert block.experts.config._experts_implementation == "grouped_mm"
+
+
+def test_explicit_eager_policy_and_expert_hooks_are_preserved():
+    block = _FixedRoutes().to("mps")
+    block.experts.config._experts_implementation = "eager"
+    bridge = MoEBridge("moe")
+    bridge.set_original_component(block)
+    inputs = torch.randn(10, 64).to("mps")
+    indices, weights = block.indices.clone(), block.weights.clone()
+    calls = []
+
+    def observe(module, args, output):
+        calls.append(args)
+        return output * 2
+
+    handle = block.experts.register_forward_hook(observe)
+    try:
+        with torch.no_grad(), warnings.catch_warnings(record=True) as caught:
+            warnings.simplefilter("always")
+            expected = _oracle(block, inputs) * 2
+            actual = bridge(inputs)
+        assert not caught
+        assert len(calls) == 1
+        torch.testing.assert_close(calls[0][0], inputs, rtol=0, atol=0)
+        torch.testing.assert_close(calls[0][1], indices, rtol=0, atol=0)
+        torch.testing.assert_close(calls[0][2], weights, rtol=0, atol=0)
+        torch.testing.assert_close(actual, expected, rtol=1e-4, atol=1e-6)
+        assert actual.shape == inputs.shape
+        assert block.experts.config._experts_implementation == "eager"
+        torch.testing.assert_close(block.indices, indices, rtol=0, atol=0)
+        torch.testing.assert_close(block.weights, weights, rtol=0, atol=0)
+    finally:
+        handle.remove()
+
+
+def test_corrected_experts_preserve_input_route_and_parameter_gradients():
+    block = _FixedRoutes()
+    block.indices = torch.tensor(
+        [
+            [0, 1],
+            [0, 3],
+            [1, 0],
+            [1, 3],
+            [1, 0],
+            [3, 0],
+            [3, 1],
+            [3, 0],
+            [3, 1],
+            [3, 0],
+        ]
+    )
+    reference = copy.deepcopy(block)
+    block.to("mps")
+    bridge = MoEBridge("moe")
+    bridge.set_original_component(block)
+    data = torch.randn(10, 64, generator=torch.Generator().manual_seed(1))
+    cpu_input = data.clone().requires_grad_()
+    mps_input = data.to("mps").requires_grad_()
+    cpu_routes = torch.linspace(0.2, 1.0, 20).reshape(10, 2).requires_grad_()
+    mps_routes = torch.linspace(0.2, 1.0, 20, device="mps").reshape(10, 2).requires_grad_()
+    calls = []
+
+    def observe(module, args, output):
+        calls.append(args)
+        return output * 2
+
+    handle = block.experts.register_forward_hook(observe)
+    try:
+        expected = _oracle(reference, cpu_input, cpu_routes) * 2
+        actual = bridge(mps_input, top_k_weights=mps_routes)
+    finally:
+        handle.remove()
+    assert len(calls) == 1
+    assert calls[0][0] is mps_input
+    assert calls[0][1] is block.indices
+    assert calls[0][2] is mps_routes
+    assert actual.shape == mps_input.shape
+    torch.testing.assert_close(actual.cpu(), expected, rtol=1e-4, atol=1e-6)
+    cpu_targets = (cpu_input, cpu_routes, *reference.experts.parameters())
+    mps_targets = (mps_input, mps_routes, *block.experts.parameters())
+    cpu_grads = torch.autograd.grad(expected.square().sum(), cpu_targets)
+    mps_grads = torch.autograd.grad(actual.square().sum(), mps_targets)
+    for measured, oracle in zip(mps_grads, cpu_grads):
+        assert measured.device.type == "mps"
+        torch.testing.assert_close(measured.cpu(), oracle, rtol=1e-4, atol=1e-6)
+    assert cpu_grads[1].abs().max() > 0
+    assert all(gradient.abs().max() > 0 for gradient in cpu_grads[2:])
+    torch.testing.assert_close(block.indices.cpu(), reference.indices, rtol=0, atol=0)
+    for parameter, saved in zip(block.experts.parameters(), reference.experts.parameters()):
+        torch.testing.assert_close(parameter.cpu(), saved, rtol=0, atol=0)
+    assert block.experts.config._experts_implementation == "grouped_mm"
+
+
+def test_failed_expert_forward_leaves_policy_and_parameters_intact():
+    block = _FixedRoutes().to("mps")
+    bridge = MoEBridge("moe")
+    bridge.set_original_component(block)
+    inputs = torch.randn(10, 64).to("mps")
+    parameters = [parameter.clone() for parameter in block.experts.parameters()]
+    activation = block.experts.act_fn
+
+    class FailingActivation(nn.Module):
+        def forward(self, value):
+            raise ValueError("expert activation failed")
+
+    block.experts.act_fn = FailingActivation()
+    try:
+        with pytest.raises(ValueError, match="expert activation failed"):
+            bridge(inputs)
+    finally:
+        block.experts.act_fn = activation
+    assert block.experts.config._experts_implementation == "grouped_mm"
+    with torch.no_grad():
+        actual = bridge(inputs)
+        expected = _oracle(block, inputs)
+    torch.testing.assert_close(actual, expected, rtol=1e-4, atol=1e-6)
+    for parameter, saved in zip(block.experts.parameters(), parameters):
+        torch.testing.assert_close(parameter, saved, rtol=0, atol=0)
diff --git a/tests/unit/model_bridge/generalized_components/test_moe_expert_dispatch.py b/tests/unit/model_bridge/generalized_components/test_moe_expert_dispatch.py
index a9ce9d6eda..46cf128695 100644
--- a/tests/unit/model_bridge/generalized_components/test_moe_expert_dispatch.py
+++ b/tests/unit/model_bridge/generalized_components/test_moe_expert_dispatch.py
@@ -147,6 +147,60 @@ def test_guarded_cpu_experts_keep_grouped_execution_without_warning(monkeypatch)
     assert block.experts.config._experts_implementation == "grouped_mm"
 
 
+def test_dispatch_does_not_create_missing_policy():
+    class MissingPolicyExperts(_Experts):
+        def forward(self, hidden_states):
+            return hidden_states
+
+    experts = MissingPolicyExperts(SimpleNamespace())
+    inputs = torch.empty(1).as_subclass(_MpsInput)
+    with warnings.catch_warnings(record=True) as caught:
+        warnings.simplefilter("always")
+        assert _granite_mps_experts_forward(experts, inputs) is inputs
+    assert not caught
+    assert vars(experts.config) == {}
+
+
+def test_guarded_eager_cpu_forward_preserves_hooks_and_gradients(monkeypatch):
+    import transformers
+
+    monkeypatch.setattr(torch, "__version__", "2.11.0")
+    monkeypatch.setattr(transformers, "__version__", "5.13.0")
+    block = _block()
+    block.experts.config._experts_implementation = "eager"
+    with torch.no_grad():
+        for parameter in block.experts.parameters():
+            parameter.normal_(std=0.1)
+    reference = copy.deepcopy(block)
+    _guard_granite_mps_experts(block)
+    inputs = torch.randn(5, 8, requires_grad=True)
+    routes = torch.rand(5, 1, requires_grad=True)
+    indices = torch.tensor([[0], [0], [1], [3], [3]])
+    calls = []
+    handle = block.experts.register_forward_hook(lambda module, args, output: calls.append(args))
+    try:
+        with warnings.catch_warnings(record=True) as caught:
+            warnings.simplefilter("always")
+            actual = block.experts(inputs, indices, routes)
+        expected = reference.experts(inputs, indices, routes)
+        assert not caught
+        assert len(calls) == 1
+        assert calls[0][1] is indices
+        assert calls[0][2] is routes
+        torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+        actual_grads = torch.autograd.grad(
+            actual.sum(), (inputs, routes, *block.experts.parameters())
+        )
+        expected_grads = torch.autograd.grad(
+            expected.sum(), (inputs, routes, *reference.experts.parameters())
+        )
+        for measured, oracle in zip(actual_grads, expected_grads):
+            torch.testing.assert_close(measured, oracle, rtol=0, atol=0)
+        assert block.experts.config._experts_implementation == "eager"
+    finally:
+        handle.remove()
+
+
 def test_unrelated_experts_and_custom_forwards_are_unchanged():
     block = nn.Module()
     block.experts = _Experts(SimpleNamespace(_experts_implementation="grouped_mm"))

From caa8e21af505e9ccdb34a62dd039481999a82065 Mon Sep 17 00:00:00 2001
From: janmenjayap <janmenjaya.panda.22@gmail.com>
Date: Sun, 11 Oct 2026 21:52:22 +0530
Subject: [PATCH 3/8] docs(moe): document MPS grouped-expert correctness
 fallback

---
 docs/source/content/drivers.md | 90 ++++++++++++++++++++++++++++++++++
 1 file changed, 90 insertions(+)

diff --git a/docs/source/content/drivers.md b/docs/source/content/drivers.md
index 43918889c9..4652ef19ed 100644
--- a/docs/source/content/drivers.md
+++ b/docs/source/content/drivers.md
@@ -48,6 +48,96 @@ logits, cache = bridge.run_with_cache("Hello, world")
 
 There is also `TransformerBridge.boot_native(config)`, which builds a small randomly-initialized TL-native model on the same driver — no HuggingFace Hub call — useful for tests and toy-model experiments.
 
+#### Granite MoE Hybrid grouped experts on Apple MPS
+
+The local transformers backend applies a narrow correctness fallback for the exact
+HF `GraniteMoeHybridExperts` class when all of these conditions hold:
+
+- PyTorch release **2.11.0** and transformers release **5.13.0**, including local
+  version suffixes but excluding prerelease and development versions.
+- The live expert input is on **MPS**.
+- The expert config selects **`grouped_mm`**.
+
+In this verified environment, grouped expert matrix products can be incorrect or
+unstable. The HF grouped fallback also uses nonzero-offset `torch.mm(..., out=...)`
+slice writes that reproduce incorrect results, so disabling only native grouped
+GEMM is insufficient. The bridge selects HF's original ordinary eager expert
+forward instead, emitting a `RuntimeWarning` about the correctness fallback and
+its possible performance cost. Ordinary eager execution may be slower; no
+throughput guarantee is made.
+
+Dispatch is installed on the expert instance when its containing `MoEBridge`
+runs, and checks the live input device and expert policy on every expert call.
+CPU-to-MPS and MPS-to-CPU migration therefore follow the current device. Shared
+HF configs, global dispatch tables, weights, routing indices, and routing weights
+are not mutated. Expert module hooks still execute, including output edits, and
+input, routing-weight, and expert-parameter gradients remain connected.
+
+The fallback does not serialize forwards or mutate shared configuration while
+other layers run. This avoids a shared-policy race, but does not guarantee that
+arbitrary concurrent model calls, device migration, config edits, or hook edits
+are safe. Direct HF execution before the containing bridge's first call is not
+protected; after installation, calls to that same expert instance use the local
+wrapper. Custom instance forwards are left unchanged.
+
+CPU/CUDA inputs, explicit eager execution, other expert policies, and other
+expert classes retain their existing dispatch. The observed defect and regression
+validation cover only the versions above, Apple MPS, and float32. Older/newer
+releases are neither declared affected nor declared fixed, and CUDA behavior has
+not been device-validated for this fallback. Unverified versions receive no
+workaround. MPS remains opt-in via `TRANSFORMERLENS_ALLOW_MPS=1`; this fallback is
+not a general guarantee of MPS correctness.
+
+##### Contributor diagnostics
+
+From the repository root, the offline regression selection requires no Hub
+model download and checks independent expert-oracle agreement, repeatability,
+device migration, hooks, gradients, and failure recovery:
+
+```bash
+TRANSFORMERLENS_ALLOW_MPS=1 uv run pytest \
+  tests/mps/test_granite_grouped_expert_dispatch.py -q
+uv run pytest \
+  tests/unit/model_bridge/generalized_components/test_moe_expert_dispatch.py -q
+TRANSFORMERLENS_ALLOW_MPS=1 uv run pytest \
+  tests/integration/model_bridge/test_granite_moe_hybrid_adapter.py -q
+```
+
+The primitive diagnostic below is independent of TransformerLens and HF expert
+routing. It reports error rather than asserting that an upstream kernel must
+remain defective. Run it on an MPS-capable host and record the installed versions
+when evaluating an upstream fix:
+
+```python
+import platform
+
+import torch
+import transformers
+
+if not torch.backends.mps.is_available():
+    raise RuntimeError("This diagnostic requires MPS")
+
+print(platform.system(), platform.release(), platform.machine())
+print("torch:", torch.__version__, "transformers:", transformers.__version__)
+generator = torch.Generator().manual_seed(0)
+inputs = torch.randn(10, 64, generator=generator)
+weights = torch.randn(4, 64, 32, generator=generator)
+expected = inputs[5:] @ weights[3]
+x, w = inputs.to("mps"), weights.to("mps")
+with torch.no_grad():
+    output = torch.zeros(10, 32, device="mps")
+    torch.mm(x[5:], w[3], out=output[5:])
+    assigned = torch.zeros_like(output)
+    assigned[5:] = x[5:] @ w[3]
+    torch.mps.synchronize()
+    print("out-slice max error:", (output[5:].cpu() - expected).abs().max().item())
+    print("assignment max error:", (assigned[5:].cpu() - expected).abs().max().item())
+```
+
+No dependency pin change or confirmed upstream fixed version is implied by this
+workaround. Revalidate both the primitive and full expert regressions before
+changing its scope.
+
 ### vLLM — high-throughput capture + declarative interventions
 
 `RemoteBridge.boot_vllm` constructs a vLLM engine (PagedAttention, `torch.compile`, CUDA graphs) and installs capture hooks inside the worker *before* compilation. Activations come back over `collective_rpc` and replay through the bridge's HookPoint tree. This is the throughput path for SAE/probe data collection; unlike observation-only tools (vllm-lens), each hook also applies an affine transform `output = output * scale + bias`, so declarative interventions (`suppress` / `scale` / `add` / `set`) propagate to downstream layers.

From 7dcd402ac8147df09ab20d0d2be47274067fde89 Mon Sep 17 00:00:00 2001
From: janmenjayap <janmenjaya.panda.22@gmail.com>
Date: Tue, 13 Oct 2026 15:10:18 +0530
Subject: [PATCH 4/8] fix(moe): detect broken MPS output-slice writes at
 runtime

---
 .../test_granite_grouped_expert_dispatch.py   |  44 ++++-
 .../test_moe_expert_dispatch.py               | 167 ++++++++++++++++--
 .../generalized_components/moe.py             |  47 ++++-
 3 files changed, 231 insertions(+), 27 deletions(-)

diff --git a/tests/mps/test_granite_grouped_expert_dispatch.py b/tests/mps/test_granite_grouped_expert_dispatch.py
index 6413c2e561..014a0273bb 100644
--- a/tests/mps/test_granite_grouped_expert_dispatch.py
+++ b/tests/mps/test_granite_grouped_expert_dispatch.py
@@ -6,19 +6,50 @@
 
 import pytest
 import torch
-import transformers
-from packaging.version import Version
 from torch import nn
 from torch.nn import functional as F
 from transformers.models.granitemoehybrid.modeling_granitemoehybrid import (
     GraniteMoeHybridExperts,
 )
 
-from transformer_lens.model_bridge.generalized_components.moe import MoEBridge
+from transformer_lens.model_bridge.generalized_components import moe
+from transformer_lens.model_bridge.generalized_components.moe import (
+    MoEBridge,
+    _mps_output_slice_writes_are_broken,
+)
 
 pytestmark = pytest.mark.skipif(not torch.backends.mps.is_available(), reason="MPS not available")
 
 
+@pytest.fixture(autouse=True)
+def isolated_probe_cache():
+    moe._MPS_OUTPUT_SLICE_RESULTS.clear()
+    yield
+    moe._MPS_OUTPUT_SLICE_RESULTS.clear()
+
+
+@pytest.mark.parametrize("dtype", [torch.float32, torch.float16])
+def test_runtime_probe_matches_direct_slice_write_diagnostic(dtype):
+    generator = torch.Generator(device="cpu").manual_seed(0)
+    inputs = torch.randn(10, 64, generator=generator).to(device="mps", dtype=dtype)
+    weights = torch.randn(4, 64, 32, generator=generator).to(device="mps", dtype=dtype)
+    with torch.no_grad():
+        expected = inputs[5:] @ weights[3]
+        output = torch.zeros(10, 32, device="mps", dtype=dtype)
+        torch.mm(inputs[5:], weights[3], out=output[5:])
+    products_match = torch.allclose(
+        output[5:], expected, rtol=max(1e-5, 2 * torch.finfo(dtype).eps), atol=1e-5
+    )
+    prefix_intact = torch.count_nonzero(output[:5]).item() == 0
+    cpu_rng = torch.get_rng_state().clone()
+    mps_rng = torch.mps.get_rng_state().clone()
+    detected = _mps_output_slice_writes_are_broken(inputs.device, dtype)
+    assert detected == (not (products_match and prefix_intact))
+    assert _mps_output_slice_writes_are_broken(inputs.device, dtype) == detected
+    torch.testing.assert_close(torch.get_rng_state(), cpu_rng, rtol=0, atol=0)
+    torch.testing.assert_close(torch.mps.get_rng_state(), mps_rng, rtol=0, atol=0)
+
+
 class _FixedRoutes(nn.Module):
     def __init__(self):
         super().__init__()
@@ -73,11 +104,8 @@ def test_fixed_route_experts_match_independent_oracle_and_repeat():
         assert first.device.type == "mps"
         torch.testing.assert_close(first.cpu(), expected, rtol=1e-4, atol=1e-6)
         torch.testing.assert_close(second, first, rtol=0, atol=0)
-        if (
-            Version(torch.__version__).public == "2.11.0"
-            and Version(transformers.__version__).public == "5.13.0"
-        ):
-            assert any("eager Granite" in str(warning.message) for warning in caught)
+        warned = any("eager Granite" in str(warning.message) for warning in caught)
+        assert warned == _mps_output_slice_writes_are_broken(first.device, first.dtype)
         assert len(observed) == 2
         assert block.experts.config._experts_implementation == "grouped_mm"
     finally:
diff --git a/tests/unit/model_bridge/generalized_components/test_moe_expert_dispatch.py b/tests/unit/model_bridge/generalized_components/test_moe_expert_dispatch.py
index 46cf128695..0c12bb4a8f 100644
--- a/tests/unit/model_bridge/generalized_components/test_moe_expert_dispatch.py
+++ b/tests/unit/model_bridge/generalized_components/test_moe_expert_dispatch.py
@@ -13,12 +13,23 @@
     GraniteMoeHybridExperts,
 )
 
+from transformer_lens.model_bridge.generalized_components import moe
 from transformer_lens.model_bridge.generalized_components.moe import (
     _granite_mps_experts_forward,
     _guard_granite_mps_experts,
+    _mps_output_slice_writes_are_broken,
+    _probe_mps_output_slice_writes,
 )
 
 
+@pytest.fixture(autouse=True)
+def isolated_probe(monkeypatch):
+    moe._MPS_OUTPUT_SLICE_RESULTS.clear()
+    monkeypatch.setattr(moe, "_probe_mps_output_slice_writes", lambda device, dtype: True)
+    yield
+    moe._MPS_OUTPUT_SLICE_RESULTS.clear()
+
+
 class _MpsInput(torch.Tensor):
     @property
     def device(self):
@@ -88,26 +99,162 @@ def _block():
 
 
 @pytest.mark.parametrize(
-    "torch_version,hf_version,installed",
+    "torch_version,hf_version",
     [
-        ("2.11.0", "5.13.0", True),
-        ("2.11.0+local", "5.13.0+local", True),
-        ("2.11.1", "5.13.0", False),
-        ("2.12.0", "5.13.0", False),
-        ("2.11.0", "5.14.0", False),
-        ("2.11.0rc1", "5.13.0", False),
-        ("2.11.0", "5.13.0.dev0", False),
+        ("2.11.0", "5.13.0"),
+        ("2.11.0+local", "5.13.0+local"),
+        ("2.11.1", "5.13.0"),
+        ("2.12.0", "5.13.0"),
+        ("2.11.0", "5.14.0"),
+        ("2.11.0rc1", "5.13.0"),
+        ("2.11.0", "5.13.0.dev0"),
     ],
 )
-def test_installation_version_scope(monkeypatch, torch_version, hf_version, installed):
+@pytest.mark.parametrize("broken", [False, True])
+def test_dispatch_uses_probe_not_package_versions(monkeypatch, torch_version, hf_version, broken):
     import transformers
 
     monkeypatch.setattr(torch, "__version__", torch_version)
     monkeypatch.setattr(transformers, "__version__", hf_version)
+    probes = []
+
+    def probe(device, dtype):
+        probes.append((device, dtype))
+        return broken
+
+    monkeypatch.setattr(moe, "_probe_mps_output_slice_writes", probe)
     block = _block()
     _guard_granite_mps_experts(block)
-    assert (block.experts.forward.__func__ is _granite_mps_experts_forward) == installed
+    assert block.experts.forward.__func__ is _granite_mps_experts_forward
+    assert not probes
     assert block.experts.config._experts_implementation == "grouped_mm"
+    experts = _Experts(block.experts.config)
+    with warnings.catch_warnings(record=True) as caught:
+        warnings.simplefilter("always")
+        actual = _granite_mps_experts_forward(experts, torch.empty(1).as_subclass(_MpsInput))
+    assert actual == ("eager" if broken else "grouped_mm")
+    assert len(caught) == int(broken)
+    assert probes == [(torch.device("mps"), torch.float32)]
+
+
+@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
+@pytest.mark.parametrize("corruption", [None, "missing_write", "prefix_write"])
+def test_primitive_probe_checks_products_and_untouched_rows(monkeypatch, dtype, corruption):
+    original_mm = torch.mm
+
+    def mm(inputs, weights, *, out):
+        assert not torch.is_grad_enabled()
+        assert out.storage_offset() == 5 * 32
+        if corruption != "missing_write":
+            original_mm(inputs, weights, out=out)
+        if corruption == "prefix_write":
+            out._base[0, 0] = 1
+        return out
+
+    monkeypatch.setattr(torch, "mm", mm)
+    rng_before = torch.get_rng_state().clone()
+    assert _probe_mps_output_slice_writes(torch.device("cpu"), dtype) == (corruption is not None)
+    torch.testing.assert_close(torch.get_rng_state(), rng_before, rtol=0, atol=0)
+
+
+def test_probe_cache_is_per_device_and_dtype(monkeypatch):
+    calls = []
+
+    def probe(device, dtype):
+        calls.append((device, dtype))
+        return dtype == torch.float16
+
+    monkeypatch.setattr(moe, "_probe_mps_output_slice_writes", probe)
+    keys = [
+        (torch.device("mps:0"), torch.float32),
+        (torch.device("mps:0"), torch.float16),
+        (torch.device("mps:1"), torch.float32),
+    ]
+    for _ in range(2):
+        for device, dtype in keys:
+            assert _mps_output_slice_writes_are_broken(device, dtype) == (dtype == torch.float16)
+    assert calls == keys
+
+
+def test_concurrent_first_use_probes_once(monkeypatch):
+    calls = []
+
+    def probe(device, dtype):
+        calls.append((device, dtype))
+        return False
+
+    monkeypatch.setattr(moe, "_probe_mps_output_slice_writes", probe)
+    with ThreadPoolExecutor(max_workers=4) as pool:
+        results = list(
+            pool.map(
+                lambda _: _mps_output_slice_writes_are_broken(torch.device("mps"), torch.float32),
+                range(16),
+            )
+        )
+    assert results == [False] * 16
+    assert calls == [(torch.device("mps"), torch.float32)]
+
+
+@pytest.mark.parametrize("error_type", [RuntimeError, NotImplementedError, TypeError])
+def test_probe_failure_is_diagnostic_and_not_cached(monkeypatch, error_type):
+    error = error_type("unsupported diagnostic")
+
+    def probe(device, dtype):
+        raise error
+
+    monkeypatch.setattr(moe, "_probe_mps_output_slice_writes", probe)
+    experts = _Experts(SimpleNamespace(_experts_implementation="grouped_mm"))
+    with warnings.catch_warnings(record=True) as caught:
+        warnings.simplefilter("always")
+        with pytest.raises(RuntimeError, match="Select eager experts explicitly") as raised:
+            _granite_mps_experts_forward(experts, torch.empty(1).as_subclass(_MpsInput))
+    assert raised.value.__cause__ is error
+    assert not caught
+    assert not moe._MPS_OUTPUT_SLICE_RESULTS
+    assert experts.config._experts_implementation == "grouped_mm"
+    monkeypatch.setattr(moe, "_probe_mps_output_slice_writes", lambda device, dtype: False)
+    assert (
+        _granite_mps_experts_forward(experts, torch.empty(1).as_subclass(_MpsInput)) == "grouped_mm"
+    )
+
+
+def test_unexpected_probe_failure_propagates(monkeypatch):
+    def probe(device, dtype):
+        raise ValueError("unexpected diagnostic failure")
+
+    monkeypatch.setattr(moe, "_probe_mps_output_slice_writes", probe)
+    with pytest.raises(ValueError, match="unexpected diagnostic failure"):
+        _mps_output_slice_writes_are_broken(torch.device("mps"), torch.float32)
+    assert not moe._MPS_OUTPUT_SLICE_RESULTS
+
+
+@pytest.mark.parametrize(
+    "device,policy",
+    [
+        (device, policy)
+        for device in ["cpu", "cuda"]
+        for policy in ["grouped_mm", "eager", "batched_mm", None]
+    ]
+    + [("mps", policy) for policy in ["eager", "batched_mm", None]],
+)
+def test_unaffected_calls_do_not_probe(monkeypatch, policy, device):
+    class DeviceInput(torch.Tensor):
+        @property
+        def device(self):
+            return torch.device(device)
+
+    def probe(device, dtype):
+        raise AssertionError("unaffected execution must not probe")
+
+    monkeypatch.setattr(moe, "_probe_mps_output_slice_writes", probe)
+    experts = _Experts(SimpleNamespace(_experts_implementation=policy))
+    with warnings.catch_warnings(record=True) as caught:
+        warnings.simplefilter("always")
+        assert (
+            _granite_mps_experts_forward(experts, torch.empty(1).as_subclass(DeviceInput)) == policy
+        )
+    assert not caught
+    assert not moe._MPS_OUTPUT_SLICE_RESULTS
 
 
 def test_installation_is_idempotent_and_deepcopy_safe(monkeypatch):
diff --git a/transformer_lens/model_bridge/generalized_components/moe.py b/transformer_lens/model_bridge/generalized_components/moe.py
index dc50c0441d..fc805feea8 100644
--- a/transformer_lens/model_bridge/generalized_components/moe.py
+++ b/transformer_lens/model_bridge/generalized_components/moe.py
@@ -7,11 +7,11 @@
 
 import logging
 import warnings
+from threading import Lock
 from types import MethodType
 from typing import Any, Dict, List, Mapping, Optional, Tuple
 
 import torch
-from packaging.version import Version
 from torch import nn
 
 from transformer_lens.hook_points import HookPoint
@@ -32,6 +32,39 @@
 _BATCHED_OUTPUT_PROJECTIONS = frozenset({"down_proj"})
 _BATCHED_PROJECTIONS = _BATCHED_INPUT_PROJECTIONS | _BATCHED_OUTPUT_PROJECTIONS
 
+_MPS_OUTPUT_SLICE_RESULTS: Dict[Tuple[torch.device, torch.dtype], bool] = {}
+_MPS_OUTPUT_SLICE_LOCK = Lock()
+
+
+def _probe_mps_output_slice_writes(device: torch.device, dtype: torch.dtype) -> bool:
+    """Check for incorrect nonzero-offset matmul output writes using isolated tensors."""
+    with torch.no_grad():
+        generator = torch.Generator(device="cpu").manual_seed(0)
+        inputs = torch.randn(10, 64, generator=generator).to(device=device, dtype=dtype)
+        weights = torch.randn(4, 64, 32, generator=generator).to(device=device, dtype=dtype)
+        output = torch.zeros(10, 32, device=device, dtype=dtype)
+        expected = inputs[5:] @ weights[3]
+        torch.mm(inputs[5:], weights[3], out=output[5:])
+        rtol = max(1e-5, 2 * torch.finfo(dtype).eps)
+        prefix_intact = torch.count_nonzero(output[:5]).item() == 0
+        products_match = torch.allclose(output[5:], expected, rtol=rtol, atol=1e-5)
+        return not (prefix_intact and products_match)
+
+
+def _mps_output_slice_writes_are_broken(device: torch.device, dtype: torch.dtype) -> bool:
+    """Cache the primitive check once per device and dtype, including concurrent first use."""
+    key = (device, dtype)
+    with _MPS_OUTPUT_SLICE_LOCK:
+        if key not in _MPS_OUTPUT_SLICE_RESULTS:
+            try:
+                _MPS_OUTPUT_SLICE_RESULTS[key] = _probe_mps_output_slice_writes(device, dtype)
+            except (RuntimeError, TypeError) as error:
+                raise RuntimeError(
+                    f"Cannot verify MPS matmul output-slice writes on {device} with {dtype}. "
+                    "Select eager experts explicitly or use CPU execution."
+                ) from error
+        return _MPS_OUTPUT_SLICE_RESULTS[key]
+
 
 def _granite_mps_experts_forward(experts: nn.Module, *args: Any, **kwargs: Any) -> Any:
     """Dispatch locally, without changing the config shared by other expert layers."""
@@ -42,10 +75,11 @@ def _granite_mps_experts_forward(experts: nn.Module, *args: Any, **kwargs: Any)
         isinstance(hidden_states, torch.Tensor)
         and hidden_states.device.type == "mps"
         and getattr(config, "_experts_implementation", None) == "grouped_mm"
+        and _mps_output_slice_writes_are_broken(hidden_states.device, hidden_states.dtype)
     ):
         warnings.warn(
-            "Using eager Granite MoE Hybrid experts on MPS to avoid incorrect grouped "
-            "matrix products in torch 2.11.0; expert execution may be slower.",
+            "Using eager Granite MoE Hybrid experts on MPS after detecting incorrect "
+            "matmul output-slice writes; expert execution may be slower.",
             RuntimeWarning,
             stacklevel=2,
         )
@@ -67,16 +101,11 @@ def _guard_granite_mps_experts(component: nn.Module) -> None:
         "transformers.models.granitemoehybrid.modeling_granitemoehybrid"
     ):
         return
-    import transformers
     from transformers.models.granitemoehybrid.modeling_granitemoehybrid import (
         GraniteMoeHybridExperts,
     )
 
-    if (
-        type(experts) is not GraniteMoeHybridExperts
-        or Version(torch.__version__).public != "2.11.0"
-        or Version(transformers.__version__).public != "5.13.0"
-    ):
+    if type(experts) is not GraniteMoeHybridExperts:
         return
     # Respect custom instance forwards rather than silently replacing them.
     if getattr(experts.forward, "__func__", None) is not GraniteMoeHybridExperts.forward:

From a786fae35df9120032e3b90b97b0e3cee12f41fb Mon Sep 17 00:00:00 2001
From: janmenjayap <janmenjaya.panda.22@gmail.com>
Date: Tue, 13 Oct 2026 15:34:11 +0530
Subject: [PATCH 5/8] fix(moe): guard shared Hugging Face expert dispatch on
 MPS

---
 .../test_granite_grouped_expert_dispatch.py   |  78 +++++--
 .../test_moe_expert_dispatch.py               | 217 +++++++++++++++---
 .../generalized_components/moe.py             |  72 ++++--
 3 files changed, 295 insertions(+), 72 deletions(-)

diff --git a/tests/mps/test_granite_grouped_expert_dispatch.py b/tests/mps/test_granite_grouped_expert_dispatch.py
index 014a0273bb..c9d8c180f8 100644
--- a/tests/mps/test_granite_grouped_expert_dispatch.py
+++ b/tests/mps/test_granite_grouped_expert_dispatch.py
@@ -1,4 +1,4 @@
-"""Granite expert correctness on MPS with fixed routes and an independent oracle."""
+"""Shared HF expert correctness on MPS with fixed routes and an independent oracle."""
 
 import copy
 import warnings
@@ -8,9 +8,13 @@
 import torch
 from torch import nn
 from torch.nn import functional as F
+from transformers.integrations import moe as hf_moe
+from transformers.models.granitemoe.modeling_granitemoe import GraniteMoeExperts
 from transformers.models.granitemoehybrid.modeling_granitemoehybrid import (
     GraniteMoeHybridExperts,
 )
+from transformers.models.mixtral.modeling_mixtral import MixtralExperts
+from transformers.models.qwen3_moe.modeling_qwen3_moe import Qwen3MoeExperts
 
 from transformer_lens.model_bridge.generalized_components import moe
 from transformer_lens.model_bridge.generalized_components.moe import (
@@ -21,6 +25,11 @@
 pytestmark = pytest.mark.skipif(not torch.backends.mps.is_available(), reason="MPS not available")
 
 
+@pytest.fixture(params=[GraniteMoeHybridExperts, GraniteMoeExperts, MixtralExperts, Qwen3MoeExperts])
+def expert_class(request):
+    return request.param
+
+
 @pytest.fixture(autouse=True)
 def isolated_probe_cache():
     moe._MPS_OUTPUT_SLICE_RESULTS.clear()
@@ -51,13 +60,18 @@ def test_runtime_probe_matches_direct_slice_write_diagnostic(dtype):
 
 
 class _FixedRoutes(nn.Module):
-    def __init__(self):
+    def __init__(self, expert_class=GraniteMoeHybridExperts):
         super().__init__()
         config = SimpleNamespace(
-            hidden_size=64, intermediate_size=32, num_local_experts=4, hidden_act="silu"
+            hidden_size=64,
+            intermediate_size=32,
+            moe_intermediate_size=32,
+            num_local_experts=4,
+            num_experts=4,
+            hidden_act="silu",
         )
         config._experts_implementation = "grouped_mm"
-        self.experts = GraniteMoeHybridExperts(config)
+        self.experts = expert_class(config)
         generator = torch.Generator().manual_seed(0)
         with torch.no_grad():
             for parameter in self.experts.parameters():
@@ -84,8 +98,8 @@ def _oracle(block, inputs, top_k_weights=None):
     return torch.stack(contributions)
 
 
-def test_fixed_route_experts_match_independent_oracle_and_repeat():
-    block = _FixedRoutes()
+def test_fixed_route_experts_match_independent_oracle_and_repeat(expert_class):
+    block = _FixedRoutes(expert_class)
     inputs = torch.randn(10, 64, generator=torch.Generator().manual_seed(1))
     with torch.no_grad():
         expected = _oracle(block, inputs)
@@ -104,7 +118,7 @@ def test_fixed_route_experts_match_independent_oracle_and_repeat():
         assert first.device.type == "mps"
         torch.testing.assert_close(first.cpu(), expected, rtol=1e-4, atol=1e-6)
         torch.testing.assert_close(second, first, rtol=0, atol=0)
-        warned = any("eager Granite" in str(warning.message) for warning in caught)
+        warned = any("eager experts" in str(warning.message) for warning in caught)
         assert warned == _mps_output_slice_writes_are_broken(first.device, first.dtype)
         assert len(observed) == 2
         assert block.experts.config._experts_implementation == "grouped_mm"
@@ -112,8 +126,8 @@ def test_fixed_route_experts_match_independent_oracle_and_repeat():
         handle.remove()
 
 
-def test_dispatch_follows_cpu_mps_cpu_migration():
-    block = _FixedRoutes()
+def test_dispatch_follows_cpu_mps_cpu_migration(expert_class):
+    block = _FixedRoutes(expert_class)
     bridge = MoEBridge("moe")
     bridge.set_original_component(block)
     inputs = torch.randn(10, 64, generator=torch.Generator().manual_seed(1))
@@ -137,8 +151,8 @@ def test_dispatch_follows_cpu_mps_cpu_migration():
     assert block.experts.config._experts_implementation == "grouped_mm"
 
 
-def test_explicit_eager_policy_and_expert_hooks_are_preserved():
-    block = _FixedRoutes().to("mps")
+def test_explicit_eager_policy_and_expert_hooks_are_preserved(expert_class):
+    block = _FixedRoutes(expert_class).to("mps")
     block.experts.config._experts_implementation = "eager"
     bridge = MoEBridge("moe")
     bridge.set_original_component(block)
@@ -170,8 +184,8 @@ def observe(module, args, output):
         handle.remove()
 
 
-def test_corrected_experts_preserve_input_route_and_parameter_gradients():
-    block = _FixedRoutes()
+def test_corrected_experts_preserve_input_route_and_parameter_gradients(expert_class):
+    block = _FixedRoutes(expert_class)
     block.indices = torch.tensor(
         [
             [0, 1],
@@ -228,8 +242,8 @@ def observe(module, args, output):
     assert block.experts.config._experts_implementation == "grouped_mm"
 
 
-def test_failed_expert_forward_leaves_policy_and_parameters_intact():
-    block = _FixedRoutes().to("mps")
+def test_failed_expert_forward_leaves_policy_and_parameters_intact(expert_class):
+    block = _FixedRoutes(expert_class).to("mps")
     bridge = MoEBridge("moe")
     bridge.set_original_component(block)
     inputs = torch.randn(10, 64).to("mps")
@@ -253,3 +267,37 @@ def forward(self, value):
     torch.testing.assert_close(actual, expected, rtol=1e-4, atol=1e-6)
     for parameter, saved in zip(block.experts.parameters(), parameters):
         torch.testing.assert_close(parameter, saved, rtol=0, atol=0)
+
+
+@pytest.mark.parametrize("broken", [False, True])
+def test_real_expert_dispatch_selects_grouped_or_eager_from_probe(monkeypatch, expert_class, broken):
+    grouped_calls = []
+
+    def safe_grouped(inputs, weights, offs):
+        grouped_calls.append(offs.cpu().tolist())
+        products = []
+        start = 0
+        for index, end in enumerate(offs.cpu().tolist()):
+            products.append(inputs[start:end] @ weights[index])
+            start = end
+        return torch.cat(products)
+
+    monkeypatch.setattr(hf_moe, "_grouped_mm", safe_grouped)
+    monkeypatch.setattr(moe, "_probe_mps_output_slice_writes", lambda device, dtype: broken)
+    registry = dict(hf_moe.ALL_EXPERTS_FUNCTIONS)
+    block = _FixedRoutes(expert_class)
+    inputs = torch.randn(10, 64, generator=torch.Generator().manual_seed(1))
+    with torch.no_grad():
+        expected = _oracle(block, inputs)
+    block.to("mps")
+    bridge = MoEBridge("moe")
+    bridge.set_original_component(block)
+    with torch.no_grad(), warnings.catch_warnings(record=True) as caught:
+        warnings.simplefilter("always")
+        actual = bridge(layer_input=inputs.to("mps"))
+    assert block.experts.forward.__func__ is moe._mps_experts_forward
+    torch.testing.assert_close(actual.cpu(), expected, rtol=1e-4, atol=1e-6)
+    assert len(caught) == int(broken)
+    assert grouped_calls == ([] if broken else [[2, 5, 5, 10]] * 2)
+    assert block.experts.config._experts_implementation == "grouped_mm"
+    assert dict(hf_moe.ALL_EXPERTS_FUNCTIONS) == registry
diff --git a/tests/unit/model_bridge/generalized_components/test_moe_expert_dispatch.py b/tests/unit/model_bridge/generalized_components/test_moe_expert_dispatch.py
index 0c12bb4a8f..5b80c06131 100644
--- a/tests/unit/model_bridge/generalized_components/test_moe_expert_dispatch.py
+++ b/tests/unit/model_bridge/generalized_components/test_moe_expert_dispatch.py
@@ -1,4 +1,4 @@
-"""Scoped Granite expert dispatch without shared configuration mutation."""
+"""Shared HF expert dispatch without shared configuration mutation."""
 
 import copy
 import warnings
@@ -9,18 +9,26 @@
 import pytest
 import torch
 from torch import nn
+from transformers.integrations import moe as hf_moe
+from transformers.integrations.moe import ExpertsInterface, use_experts_implementation
+from transformers.models.granitemoe.modeling_granitemoe import GraniteMoeExperts
 from transformers.models.granitemoehybrid.modeling_granitemoehybrid import (
     GraniteMoeHybridExperts,
 )
+from transformers.models.mixtral.modeling_mixtral import MixtralExperts
+from transformers.models.qwen3_moe.modeling_qwen3_moe import Qwen3MoeExperts
 
 from transformer_lens.model_bridge.generalized_components import moe
 from transformer_lens.model_bridge.generalized_components.moe import (
-    _granite_mps_experts_forward,
-    _guard_granite_mps_experts,
+    _mps_experts_forward,
+    _guard_mps_experts,
+    _hf_eager_experts_forward,
     _mps_output_slice_writes_are_broken,
     _probe_mps_output_slice_writes,
 )
 
+_EXPERT_CLASSES = [GraniteMoeHybridExperts, GraniteMoeExperts, MixtralExperts, Qwen3MoeExperts]
+
 
 @pytest.fixture(autouse=True)
 def isolated_probe(monkeypatch):
@@ -52,16 +60,28 @@ def forward(self, hidden_states):
         return self.config._experts_implementation
 
 
-def test_dispatch_does_not_mutate_shared_config():
+@pytest.fixture
+def stub_dispatch(monkeypatch):
+    recognize = moe._hf_eager_experts_forward
+
+    def eager(forward):
+        if forward is _Experts.forward:
+            return forward.__wrapped__
+        return recognize(forward)
+
+    monkeypatch.setattr(moe, "_hf_eager_experts_forward", eager)
+
+
+def test_dispatch_does_not_mutate_shared_config(stub_dispatch):
     config = SimpleNamespace(_experts_implementation="grouped_mm")
     experts = [_Experts(config), _Experts(config)]
     mps = torch.empty(1).as_subclass(_MpsInput)
     cpu = torch.empty(1)
-    with pytest.warns(RuntimeWarning, match="eager Granite"):
+    with pytest.warns(RuntimeWarning, match="eager experts"):
         with ThreadPoolExecutor(max_workers=2) as pool:
             results = list(
                 pool.map(
-                    lambda pair: _granite_mps_experts_forward(*pair),
+                    lambda pair: _mps_experts_forward(*pair),
                     [(experts[0], mps), (experts[1], cpu)],
                 )
             )
@@ -72,29 +92,34 @@ def test_dispatch_does_not_mutate_shared_config():
 @pytest.mark.parametrize("policy", ["eager", "batched_mm", None])
 def test_other_policies_are_unchanged(policy):
     experts = _Experts(SimpleNamespace(_experts_implementation=policy))
-    assert _granite_mps_experts_forward(experts, torch.empty(1).as_subclass(_MpsInput)) == policy
+    assert _mps_experts_forward(experts, torch.empty(1).as_subclass(_MpsInput)) == policy
 
 
-def test_exception_does_not_leave_dispatch_state():
+def test_exception_does_not_leave_dispatch_state(stub_dispatch):
     experts = _Experts(SimpleNamespace(_experts_implementation="grouped_mm"))
     experts.fail = True
     with pytest.warns(RuntimeWarning), pytest.raises(ValueError, match="expert failure"):
-        _granite_mps_experts_forward(experts, hidden_states=torch.empty(1).as_subclass(_MpsInput))
+        _mps_experts_forward(experts, hidden_states=torch.empty(1).as_subclass(_MpsInput))
     assert experts.config._experts_implementation == "grouped_mm"
     experts.fail = False
     with pytest.warns(RuntimeWarning):
         assert (
-            _granite_mps_experts_forward(experts, torch.empty(1).as_subclass(_MpsInput)) == "eager"
+            _mps_experts_forward(experts, torch.empty(1).as_subclass(_MpsInput)) == "eager"
         )
 
 
-def _block():
+def _block(expert_class=GraniteMoeHybridExperts):
     config = SimpleNamespace(
-        hidden_size=8, intermediate_size=4, num_local_experts=4, hidden_act="silu"
+        hidden_size=8,
+        intermediate_size=4,
+        moe_intermediate_size=4,
+        num_local_experts=4,
+        num_experts=4,
+        hidden_act="silu",
     )
     config._experts_implementation = "grouped_mm"
     block = nn.Module()
-    block.experts = GraniteMoeHybridExperts(config)
+    block.experts = expert_class(config)
     return block
 
 
@@ -111,7 +136,9 @@ def _block():
     ],
 )
 @pytest.mark.parametrize("broken", [False, True])
-def test_dispatch_uses_probe_not_package_versions(monkeypatch, torch_version, hf_version, broken):
+def test_dispatch_uses_probe_not_package_versions(
+    monkeypatch, torch_version, hf_version, broken, stub_dispatch
+):
     import transformers
 
     monkeypatch.setattr(torch, "__version__", torch_version)
@@ -124,14 +151,14 @@ def probe(device, dtype):
 
     monkeypatch.setattr(moe, "_probe_mps_output_slice_writes", probe)
     block = _block()
-    _guard_granite_mps_experts(block)
-    assert block.experts.forward.__func__ is _granite_mps_experts_forward
+    _guard_mps_experts(block)
+    assert block.experts.forward.__func__ is _mps_experts_forward
     assert not probes
     assert block.experts.config._experts_implementation == "grouped_mm"
     experts = _Experts(block.experts.config)
     with warnings.catch_warnings(record=True) as caught:
         warnings.simplefilter("always")
-        actual = _granite_mps_experts_forward(experts, torch.empty(1).as_subclass(_MpsInput))
+        actual = _mps_experts_forward(experts, torch.empty(1).as_subclass(_MpsInput))
     assert actual == ("eager" if broken else "grouped_mm")
     assert len(caught) == int(broken)
     assert probes == [(torch.device("mps"), torch.float32)]
@@ -196,7 +223,7 @@ def probe(device, dtype):
 
 
 @pytest.mark.parametrize("error_type", [RuntimeError, NotImplementedError, TypeError])
-def test_probe_failure_is_diagnostic_and_not_cached(monkeypatch, error_type):
+def test_probe_failure_is_diagnostic_and_not_cached(monkeypatch, error_type, stub_dispatch):
     error = error_type("unsupported diagnostic")
 
     def probe(device, dtype):
@@ -207,14 +234,14 @@ def probe(device, dtype):
     with warnings.catch_warnings(record=True) as caught:
         warnings.simplefilter("always")
         with pytest.raises(RuntimeError, match="Select eager experts explicitly") as raised:
-            _granite_mps_experts_forward(experts, torch.empty(1).as_subclass(_MpsInput))
+            _mps_experts_forward(experts, torch.empty(1).as_subclass(_MpsInput))
     assert raised.value.__cause__ is error
     assert not caught
     assert not moe._MPS_OUTPUT_SLICE_RESULTS
     assert experts.config._experts_implementation == "grouped_mm"
     monkeypatch.setattr(moe, "_probe_mps_output_slice_writes", lambda device, dtype: False)
     assert (
-        _granite_mps_experts_forward(experts, torch.empty(1).as_subclass(_MpsInput)) == "grouped_mm"
+        _mps_experts_forward(experts, torch.empty(1).as_subclass(_MpsInput)) == "grouped_mm"
     )
 
 
@@ -251,32 +278,34 @@ def probe(device, dtype):
     with warnings.catch_warnings(record=True) as caught:
         warnings.simplefilter("always")
         assert (
-            _granite_mps_experts_forward(experts, torch.empty(1).as_subclass(DeviceInput)) == policy
+            _mps_experts_forward(experts, torch.empty(1).as_subclass(DeviceInput)) == policy
         )
     assert not caught
     assert not moe._MPS_OUTPUT_SLICE_RESULTS
 
 
-def test_installation_is_idempotent_and_deepcopy_safe(monkeypatch):
+@pytest.mark.parametrize("expert_class", _EXPERT_CLASSES)
+def test_installation_is_idempotent_and_deepcopy_safe(monkeypatch, expert_class):
     import transformers
 
     monkeypatch.setattr(torch, "__version__", "2.11.0")
     monkeypatch.setattr(transformers, "__version__", "5.13.0")
-    block = _block()
-    _guard_granite_mps_experts(block)
+    block = _block(expert_class)
+    _guard_mps_experts(block)
     forward = block.experts.forward
-    _guard_granite_mps_experts(block)
+    _guard_mps_experts(block)
     assert block.experts.forward is forward
     cloned = copy.deepcopy(block)
     assert cloned.experts.forward.__self__ is cloned.experts
 
 
-def test_guarded_cpu_experts_keep_grouped_execution_without_warning(monkeypatch):
+@pytest.mark.parametrize("expert_class", _EXPERT_CLASSES)
+def test_guarded_cpu_experts_keep_grouped_execution_without_warning(monkeypatch, expert_class):
     import transformers
 
     monkeypatch.setattr(torch, "__version__", "2.11.0")
     monkeypatch.setattr(transformers, "__version__", "5.13.0")
-    block = _block()
+    block = _block(expert_class)
     with torch.no_grad():
         for parameter in block.experts.parameters():
             parameter.normal_(std=0.1)
@@ -285,7 +314,7 @@ def test_guarded_cpu_experts_keep_grouped_execution_without_warning(monkeypatch)
     weights = torch.ones(5, 1)
     with torch.no_grad():
         expected = block.experts(inputs, indices, weights)
-    _guard_granite_mps_experts(block)
+    _guard_mps_experts(block)
     with warnings.catch_warnings(record=True) as caught, torch.no_grad():
         warnings.simplefilter("always")
         actual = block.experts(inputs, indices, weights)
@@ -303,23 +332,24 @@ def forward(self, hidden_states):
     inputs = torch.empty(1).as_subclass(_MpsInput)
     with warnings.catch_warnings(record=True) as caught:
         warnings.simplefilter("always")
-        assert _granite_mps_experts_forward(experts, inputs) is inputs
+        assert _mps_experts_forward(experts, inputs) is inputs
     assert not caught
     assert vars(experts.config) == {}
 
 
-def test_guarded_eager_cpu_forward_preserves_hooks_and_gradients(monkeypatch):
+@pytest.mark.parametrize("expert_class", _EXPERT_CLASSES)
+def test_guarded_eager_cpu_forward_preserves_hooks_and_gradients(monkeypatch, expert_class):
     import transformers
 
     monkeypatch.setattr(torch, "__version__", "2.11.0")
     monkeypatch.setattr(transformers, "__version__", "5.13.0")
-    block = _block()
+    block = _block(expert_class)
     block.experts.config._experts_implementation = "eager"
     with torch.no_grad():
         for parameter in block.experts.parameters():
             parameter.normal_(std=0.1)
     reference = copy.deepcopy(block)
-    _guard_granite_mps_experts(block)
+    _guard_mps_experts(block)
     inputs = torch.randn(5, 8, requires_grad=True)
     routes = torch.rand(5, 1, requires_grad=True)
     indices = torch.tensor([[0], [0], [1], [3], [3]])
@@ -351,10 +381,129 @@ def test_guarded_eager_cpu_forward_preserves_hooks_and_gradients(monkeypatch):
 def test_unrelated_experts_and_custom_forwards_are_unchanged():
     block = nn.Module()
     block.experts = _Experts(SimpleNamespace(_experts_implementation="grouped_mm"))
-    _guard_granite_mps_experts(block)
+    _guard_mps_experts(block)
     assert block.experts.forward.__func__ is _Experts.forward
     block = _block()
     custom = lambda hidden_states: hidden_states
     block.experts.forward = custom
-    _guard_granite_mps_experts(block)
+    _guard_mps_experts(block)
     assert block.experts.forward is custom
+
+
+def _decorated_block(interface, config=None):
+    @use_experts_implementation(experts_interface=interface)
+    class DecoratedExperts(nn.Module):
+        def __init__(self, config):
+            super().__init__()
+            self.received = None
+
+        def forward(self, hidden_states, token=None, *, payload=None):
+            self.received = (hidden_states, token, payload)
+            return "eager"
+
+    block = nn.Module()
+    block.experts = DecoratedExperts(
+        config if config is not None else SimpleNamespace(_experts_implementation="grouped_mm")
+    )
+    return block
+
+
+@pytest.mark.parametrize("broken", [False, True])
+@pytest.mark.parametrize("keyword_input", [False, True])
+def test_shared_decorator_dispatch_preserves_arguments_and_global_registry(
+    monkeypatch, broken, keyword_input
+):
+    def grouped(experts, hidden_states, token=None, *, payload=None):
+        experts.received = (hidden_states, token, payload)
+        return "grouped_mm"
+
+    registry = dict(hf_moe.ALL_EXPERTS_FUNCTIONS)
+    interface = ExpertsInterface()
+    interface["grouped_mm"] = grouped
+    monkeypatch.setattr(hf_moe, "grouped_mm_experts_forward", grouped)
+    monkeypatch.setattr(moe, "_probe_mps_output_slice_writes", lambda device, dtype: broken)
+    block = _decorated_block(interface)
+    _guard_mps_experts(block)
+    assert block.experts.forward.__func__ is _mps_experts_forward
+    inputs = torch.empty(1).as_subclass(_MpsInput)
+    token, payload = object(), object()
+    with warnings.catch_warnings(record=True) as caught:
+        warnings.simplefilter("always")
+        if keyword_input:
+            actual = block.experts(hidden_states=inputs, token=token, payload=payload)
+        else:
+            actual = block.experts(inputs, token, payload=payload)
+    assert actual == ("eager" if broken else "grouped_mm")
+    assert block.experts.received == (inputs, token, payload)
+    assert len(caught) == int(broken)
+    assert vars(block.experts.config) == {"_experts_implementation": "grouped_mm"}
+    assert dict(hf_moe.ALL_EXPERTS_FUNCTIONS) == registry
+
+
+def test_guarded_shared_config_supports_concurrent_cpu_and_mps_dispatch(monkeypatch):
+    def grouped(experts, hidden_states):
+        return "grouped_mm"
+
+    interface = ExpertsInterface()
+    interface["grouped_mm"] = grouped
+    monkeypatch.setattr(hf_moe, "grouped_mm_experts_forward", grouped)
+    config = SimpleNamespace(_experts_implementation="grouped_mm")
+    blocks = [_decorated_block(interface, config), _decorated_block(interface, config)]
+    for block in blocks:
+        _guard_mps_experts(block)
+    inputs = [torch.empty(1).as_subclass(_MpsInput), torch.empty(1)]
+    with pytest.warns(RuntimeWarning, match="eager experts"):
+        with ThreadPoolExecutor(max_workers=2) as pool:
+            actual = list(pool.map(lambda pair: pair[0].experts(pair[1]), zip(blocks, inputs)))
+    assert actual == ["eager", "grouped_mm"]
+    assert all(block.experts.config is config for block in blocks)
+    assert vars(config) == {"_experts_implementation": "grouped_mm"}
+
+
+@pytest.mark.parametrize("expert_class", _EXPERT_CLASSES)
+def test_recovers_real_expert_eager_forward(expert_class):
+    assert _hf_eager_experts_forward(expert_class.forward) is expert_class.forward.__wrapped__
+
+
+def test_custom_grouped_interface_is_not_overridden():
+    def custom(experts, hidden_states):
+        return "custom"
+
+    interface = ExpertsInterface()
+    interface["grouped_mm"] = custom
+    block = _decorated_block(interface)
+    forward = block.experts.forward
+    _guard_mps_experts(block)
+    assert block.experts.forward == forward
+    with warnings.catch_warnings(record=True) as caught:
+        assert block.experts(torch.empty(1).as_subclass(_MpsInput)) == "custom"
+    assert not caught
+    assert not moe._MPS_OUTPUT_SLICE_RESULTS
+
+
+def test_interface_changed_after_installation_is_respected():
+    interface = ExpertsInterface()
+    block = _decorated_block(interface)
+    _guard_mps_experts(block)
+    assert block.experts.forward.__func__ is _mps_experts_forward
+    interface["grouped_mm"] = lambda experts, hidden_states: "custom"
+    with warnings.catch_warnings(record=True) as caught:
+        assert block.experts(torch.empty(1).as_subclass(_MpsInput)) == "custom"
+    assert not caught
+    assert not moe._MPS_OUTPUT_SLICE_RESULTS
+
+
+def test_mismatched_wrapped_forward_is_not_recognized(monkeypatch):
+    block = _decorated_block(ExpertsInterface())
+    forward = type(block.experts).forward
+    monkeypatch.setattr(forward, "__wrapped__", lambda experts, inputs: inputs)
+    _guard_mps_experts(block)
+    assert block.experts.forward.__func__ is forward
+
+
+def test_ordinary_expert_lists_are_unchanged():
+    block = nn.Module()
+    block.experts = nn.ModuleList([nn.Linear(8, 4)])
+    forward = block.experts.forward
+    _guard_mps_experts(block)
+    assert block.experts.forward == forward
diff --git a/transformer_lens/model_bridge/generalized_components/moe.py b/transformer_lens/model_bridge/generalized_components/moe.py
index fc805feea8..fc4737500a 100644
--- a/transformer_lens/model_bridge/generalized_components/moe.py
+++ b/transformer_lens/model_bridge/generalized_components/moe.py
@@ -7,9 +7,10 @@
 
 import logging
 import warnings
+from inspect import getclosurevars
 from threading import Lock
-from types import MethodType
-from typing import Any, Dict, List, Mapping, Optional, Tuple
+from types import CodeType, FunctionType, MethodType
+from typing import Any, Callable, Dict, List, Mapping, Optional, Tuple
 
 import torch
 from torch import nn
@@ -66,29 +67,64 @@ def _mps_output_slice_writes_are_broken(device: torch.device, dtype: torch.dtype
         return _MPS_OUTPUT_SLICE_RESULTS[key]
 
 
-def _granite_mps_experts_forward(experts: nn.Module, *args: Any, **kwargs: Any) -> Any:
+def _mps_experts_forward(experts: nn.Module, *args: Any, **kwargs: Any) -> Any:
     """Dispatch locally, without changing the config shared by other expert layers."""
     hidden_states = args[0] if args else kwargs.get("hidden_states")
     forward = type(experts).forward
     config: Any = getattr(experts, "config", None)
+    eager = None
     if (
         isinstance(hidden_states, torch.Tensor)
         and hidden_states.device.type == "mps"
         and getattr(config, "_experts_implementation", None) == "grouped_mm"
-        and _mps_output_slice_writes_are_broken(hidden_states.device, hidden_states.dtype)
+    ):
+        eager = _hf_eager_experts_forward(forward)
+    if eager is not None and _mps_output_slice_writes_are_broken(
+        hidden_states.device, hidden_states.dtype
     ):
         warnings.warn(
-            "Using eager Granite MoE Hybrid experts on MPS after detecting incorrect "
+            "Using eager experts on MPS after detecting incorrect "
             "matmul output-slice writes; expert execution may be slower.",
             RuntimeWarning,
             stacklevel=2,
         )
-        return getattr(forward, "__wrapped__")(experts, *args, **kwargs)
+        return eager(experts, *args, **kwargs)
     return forward(experts, *args, **kwargs)
 
 
-def _guard_granite_mps_experts(component: nn.Module) -> None:
-    """Install instance-local dispatch only for the verified HF expert implementation.
+def _hf_eager_experts_forward(forward: Any) -> Optional[Callable[..., Any]]:
+    """Recognize HF's shared grouped dispatcher and recover its original eager forward."""
+    from transformers.integrations.moe import (
+        ExpertsInterface,
+        grouped_mm_experts_forward,
+        use_experts_implementation,
+    )
+
+    if not isinstance(forward, FunctionType):
+        return None
+    if not any(
+        forward.__code__ is candidate
+        for wrapper in use_experts_implementation.__code__.co_consts
+        if isinstance(wrapper, CodeType) and wrapper.co_name == "wrapper"
+        for candidate in wrapper.co_consts
+        if isinstance(candidate, CodeType) and candidate.co_name == "forward"
+    ):
+        return None
+    closure = getclosurevars(forward).nonlocals
+    original = closure.get("original_forward")
+    interface = closure.get("experts_interface")
+    if (
+        not callable(original)
+        or original is not getattr(forward, "__wrapped__", None)
+        or not isinstance(interface, ExpertsInterface)
+        or interface.get_interface("grouped_mm", original) is not grouped_mm_experts_forward
+    ):
+        return None
+    return original
+
+
+def _guard_mps_experts(component: nn.Module) -> None:
+    """Install instance-local dispatch only for HF's shared grouped expert implementation.
 
     The bound method preserves nn.Module hooks and deepcopy semantics. Configs and
     global dispatch tables remain untouched, including during concurrent forwards.
@@ -97,22 +133,12 @@ def _guard_granite_mps_experts(component: nn.Module) -> None:
     if not isinstance(experts, nn.Module):
         return
     experts = unwrap_bridge(experts)
-    if type(experts).__module__ != (
-        "transformers.models.granitemoehybrid.modeling_granitemoehybrid"
-    ):
-        return
-    from transformers.models.granitemoehybrid.modeling_granitemoehybrid import (
-        GraniteMoeHybridExperts,
-    )
-
-    if type(experts) is not GraniteMoeHybridExperts:
+    forward = type(experts).forward
+    if getattr(experts.forward, "__func__", None) is not forward:
         return
-    # Respect custom instance forwards rather than silently replacing them.
-    if getattr(experts.forward, "__func__", None) is not GraniteMoeHybridExperts.forward:
+    if _hf_eager_experts_forward(forward) is None:
         return
-    if not callable(getattr(GraniteMoeHybridExperts.forward, "__wrapped__", None)):
-        raise RuntimeError("Granite eager expert forward is unavailable for the MPS workaround")
-    setattr(experts, "forward", MethodType(_granite_mps_experts_forward, experts))
+    setattr(experts, "forward", MethodType(_mps_experts_forward, experts))
 
 
 class UnfoldableMoEParameter(Exception):
@@ -516,7 +542,7 @@ def forward(self, *args: Any, **kwargs: Any) -> Any:
         elif "hidden_states" in kwargs:
             hooked = self.hook_in(kwargs["hidden_states"])
             kwargs = {**kwargs, "hidden_states": hooked}
-        _guard_granite_mps_experts(self.original_component)
+        _guard_mps_experts(self.original_component)
         output = self.original_component(*args, **kwargs)
         if isinstance(output, tuple):
             if not output:

From c6ce797b27f335f916f9a2163f6c254ac5bc0ce7 Mon Sep 17 00:00:00 2001
From: janmenjayap <janmenjaya.panda.22@gmail.com>
Date: Tue, 13 Oct 2026 15:47:48 +0530
Subject: [PATCH 6/8] fix(moe): warn once for the MPS expert correctness
 fallback

---
 .../test_granite_grouped_expert_dispatch.py   |  10 +-
 .../test_moe_expert_dispatch.py               | 113 +++++++++++++++++-
 .../generalized_components/moe.py             |  40 +++++--
 3 files changed, 153 insertions(+), 10 deletions(-)

diff --git a/tests/mps/test_granite_grouped_expert_dispatch.py b/tests/mps/test_granite_grouped_expert_dispatch.py
index c9d8c180f8..193ad3aba8 100644
--- a/tests/mps/test_granite_grouped_expert_dispatch.py
+++ b/tests/mps/test_granite_grouped_expert_dispatch.py
@@ -33,8 +33,10 @@ def expert_class(request):
 @pytest.fixture(autouse=True)
 def isolated_probe_cache():
     moe._MPS_OUTPUT_SLICE_RESULTS.clear()
+    moe._MPS_EXPERTS_WARNED = False
     yield
     moe._MPS_OUTPUT_SLICE_RESULTS.clear()
+    moe._MPS_EXPERTS_WARNED = False
 
 
 @pytest.mark.parametrize("dtype", [torch.float32, torch.float16])
@@ -118,8 +120,12 @@ def test_fixed_route_experts_match_independent_oracle_and_repeat(expert_class):
         assert first.device.type == "mps"
         torch.testing.assert_close(first.cpu(), expected, rtol=1e-4, atol=1e-6)
         torch.testing.assert_close(second, first, rtol=0, atol=0)
-        warned = any("eager experts" in str(warning.message) for warning in caught)
-        assert warned == _mps_output_slice_writes_are_broken(first.device, first.dtype)
+        fallback_warnings = [
+            warning for warning in caught if "eager experts" in str(warning.message)
+        ]
+        assert len(fallback_warnings) == int(
+            _mps_output_slice_writes_are_broken(first.device, first.dtype)
+        )
         assert len(observed) == 2
         assert block.experts.config._experts_implementation == "grouped_mm"
     finally:
diff --git a/tests/unit/model_bridge/generalized_components/test_moe_expert_dispatch.py b/tests/unit/model_bridge/generalized_components/test_moe_expert_dispatch.py
index 5b80c06131..44e641d839 100644
--- a/tests/unit/model_bridge/generalized_components/test_moe_expert_dispatch.py
+++ b/tests/unit/model_bridge/generalized_components/test_moe_expert_dispatch.py
@@ -1,6 +1,7 @@
 """Shared HF expert dispatch without shared configuration mutation."""
 
 import copy
+import sys
 import warnings
 from concurrent.futures import ThreadPoolExecutor
 from functools import wraps
@@ -33,9 +34,11 @@
 @pytest.fixture(autouse=True)
 def isolated_probe(monkeypatch):
     moe._MPS_OUTPUT_SLICE_RESULTS.clear()
+    moe._MPS_EXPERTS_WARNED = False
     monkeypatch.setattr(moe, "_probe_mps_output_slice_writes", lambda device, dtype: True)
     yield
     moe._MPS_OUTPUT_SLICE_RESULTS.clear()
+    moe._MPS_EXPERTS_WARNED = False
 
 
 class _MpsInput(torch.Tensor):
@@ -102,10 +105,12 @@ def test_exception_does_not_leave_dispatch_state(stub_dispatch):
         _mps_experts_forward(experts, hidden_states=torch.empty(1).as_subclass(_MpsInput))
     assert experts.config._experts_implementation == "grouped_mm"
     experts.fail = False
-    with pytest.warns(RuntimeWarning):
+    with warnings.catch_warnings(record=True) as caught:
+        warnings.simplefilter("always")
         assert (
             _mps_experts_forward(experts, torch.empty(1).as_subclass(_MpsInput)) == "eager"
         )
+    assert not caught
 
 
 def _block(expert_class=GraniteMoeHybridExperts):
@@ -507,3 +512,109 @@ def test_ordinary_expert_lists_are_unchanged():
     forward = block.experts.forward
     _guard_mps_experts(block)
     assert block.experts.forward == forward
+
+
+@pytest.mark.parametrize("concurrent", [False, True])
+def test_warning_is_once_across_layers_models_and_forwards(concurrent):
+    blocks = [_decorated_block(ExpertsInterface()) for _ in range(4)]
+    for block in blocks:
+        _guard_mps_experts(block)
+    inputs = torch.empty(1).as_subclass(_MpsInput)
+    with warnings.catch_warnings(record=True) as caught:
+        warnings.simplefilter("always")
+        if concurrent:
+            with ThreadPoolExecutor(max_workers=8) as pool:
+                results = list(pool.map(lambda block: block.experts(inputs), blocks * 8))
+        else:
+            results = [block.experts(inputs) for block in blocks * 8]
+    assert results == ["eager"] * 32
+    assert len(caught) == 1
+    assert caught[0].category is RuntimeWarning
+    assert "eager experts" in str(caught[0].message)
+    assert moe._MPS_EXPERTS_WARNED
+
+
+def test_warning_identifies_external_caller_through_runtime_wrappers(monkeypatch):
+    modules = []
+
+    def currentframe():
+        frame = sys._getframe(1)
+        while frame is not None:
+            modules.append(frame.f_globals.get("__name__", ""))
+            frame = frame.f_back
+        return sys._getframe(1)
+
+    monkeypatch.setattr(moe, "currentframe", currentframe)
+    block = _decorated_block(ExpertsInterface())
+    _guard_mps_experts(block)
+    assert not moe._MPS_EXPERTS_WARNED
+    inputs = torch.empty(1).as_subclass(_MpsInput)
+    with warnings.catch_warnings(record=True) as caught:
+        warnings.simplefilter("always")
+        caller_line = sys._getframe().f_lineno + 1
+        result = block.experts(inputs)
+    assert result == "eager"
+    assert len(caught) == 1
+    assert caught[0].filename == __file__
+    assert caught[0].lineno == caller_line
+    assert any(module.startswith("jaxtyping.") for module in modules)
+    assert any(module.startswith("torch.") for module in modules)
+
+
+def test_warning_as_error_propagates_without_consuming_notification():
+    block = _decorated_block(ExpertsInterface())
+    _guard_mps_experts(block)
+    inputs = torch.empty(1).as_subclass(_MpsInput)
+    with warnings.catch_warnings():
+        warnings.simplefilter("error", RuntimeWarning)
+        for _ in range(2):
+            with pytest.raises(RuntimeWarning, match="eager experts"):
+                block.experts(inputs)
+            assert not moe._MPS_EXPERTS_WARNED
+            assert block.experts.received is None
+            assert block.experts.config._experts_implementation == "grouped_mm"
+    with warnings.catch_warnings(record=True) as caught:
+        warnings.simplefilter("always")
+        assert block.experts(inputs) == "eager"
+        assert block.experts(inputs) == "eager"
+    assert len(caught) == 1
+    assert moe._MPS_EXPERTS_WARNED
+
+
+def test_unaffected_calls_do_not_consume_fallback_notification(monkeypatch):
+    def grouped(experts, hidden_states):
+        return "grouped_mm"
+
+    interface = ExpertsInterface()
+    interface["grouped_mm"] = grouped
+    monkeypatch.setattr(hf_moe, "grouped_mm_experts_forward", grouped)
+    monkeypatch.setattr(moe, "_probe_mps_output_slice_writes", lambda device, dtype: False)
+    block = _decorated_block(interface)
+    _guard_mps_experts(block)
+    inputs = torch.empty(1).as_subclass(_MpsInput)
+    with warnings.catch_warnings(record=True) as caught:
+        warnings.simplefilter("always")
+        assert block.experts(torch.empty(1)) == "grouped_mm"
+        assert block.experts(inputs) == "grouped_mm"
+        block.experts.config._experts_implementation = "eager"
+        assert block.experts(inputs) == "eager"
+    assert not caught
+    assert not moe._MPS_EXPERTS_WARNED
+    block.experts.config._experts_implementation = "grouped_mm"
+    moe._MPS_OUTPUT_SLICE_RESULTS.clear()
+    monkeypatch.setattr(moe, "_probe_mps_output_slice_writes", lambda device, dtype: True)
+    with pytest.warns(RuntimeWarning, match="eager experts"):
+        assert block.experts(inputs) == "eager"
+
+
+def test_fallback_warning_does_not_suppress_unrelated_warnings():
+    block = _decorated_block(ExpertsInterface())
+    _guard_mps_experts(block)
+    inputs = torch.empty(1).as_subclass(_MpsInput)
+    with warnings.catch_warnings(record=True) as caught:
+        warnings.simplefilter("always")
+        assert block.experts(inputs) == "eager"
+        warnings.warn("Unrelated diagnostic", UserWarning)
+        assert block.experts(inputs) == "eager"
+        warnings.warn("Unrelated diagnostic", UserWarning)
+    assert [warning.category for warning in caught] == [RuntimeWarning, UserWarning, UserWarning]
diff --git a/transformer_lens/model_bridge/generalized_components/moe.py b/transformer_lens/model_bridge/generalized_components/moe.py
index fc4737500a..d4117db270 100644
--- a/transformer_lens/model_bridge/generalized_components/moe.py
+++ b/transformer_lens/model_bridge/generalized_components/moe.py
@@ -7,7 +7,7 @@
 
 import logging
 import warnings
-from inspect import getclosurevars
+from inspect import currentframe, getclosurevars
 from threading import Lock
 from types import CodeType, FunctionType, MethodType
 from typing import Any, Callable, Dict, List, Mapping, Optional, Tuple
@@ -35,6 +35,37 @@
 
 _MPS_OUTPUT_SLICE_RESULTS: Dict[Tuple[torch.device, torch.dtype], bool] = {}
 _MPS_OUTPUT_SLICE_LOCK = Lock()
+_MPS_EXPERTS_WARNED = False
+_MPS_EXPERTS_WARNING_LOCK = Lock()
+
+
+def _warn_mps_experts_once() -> None:
+    """Warn once at the first caller outside the model and runtime-checking libraries."""
+    global _MPS_EXPERTS_WARNED
+    with _MPS_EXPERTS_WARNING_LOCK:
+        if _MPS_EXPERTS_WARNED:
+            return
+        frame = currentframe()
+        stacklevel = 1
+        try:
+            while frame is not None:
+                module = frame.f_globals.get("__name__", "")
+                if not any(
+                    module == prefix or module.startswith(prefix + ".")
+                    for prefix in ("transformer_lens", "transformers", "torch", "jaxtyping", "beartype")
+                ):
+                    break
+                frame = frame.f_back
+                stacklevel += 1
+            warnings.warn(
+                "Using eager experts on MPS after detecting incorrect "
+                "matmul output-slice writes; expert execution may be slower.",
+                RuntimeWarning,
+                stacklevel=stacklevel,
+            )
+            _MPS_EXPERTS_WARNED = True
+        finally:
+            del frame
 
 
 def _probe_mps_output_slice_writes(device: torch.device, dtype: torch.dtype) -> bool:
@@ -82,12 +113,7 @@ def _mps_experts_forward(experts: nn.Module, *args: Any, **kwargs: Any) -> Any:
     if eager is not None and _mps_output_slice_writes_are_broken(
         hidden_states.device, hidden_states.dtype
     ):
-        warnings.warn(
-            "Using eager experts on MPS after detecting incorrect "
-            "matmul output-slice writes; expert execution may be slower.",
-            RuntimeWarning,
-            stacklevel=2,
-        )
+        _warn_mps_experts_once()
         return eager(experts, *args, **kwargs)
     return forward(experts, *args, **kwargs)
 

From 10224e37939913811fdacb9062203c71f0ca2cbf Mon Sep 17 00:00:00 2001
From: janmenjayap <janmenjaya.panda.22@gmail.com>
Date: Tue, 13 Oct 2026 15:52:17 +0530
Subject: [PATCH 7/8] docs: move grouped-expert fallback guidance to the MPS
 setup note

---
 docs/source/content/drivers.md         | 90 --------------------------
 docs/source/content/getting_started.md |  5 ++
 2 files changed, 5 insertions(+), 90 deletions(-)

diff --git a/docs/source/content/drivers.md b/docs/source/content/drivers.md
index 4652ef19ed..43918889c9 100644
--- a/docs/source/content/drivers.md
+++ b/docs/source/content/drivers.md
@@ -48,96 +48,6 @@ logits, cache = bridge.run_with_cache("Hello, world")
 
 There is also `TransformerBridge.boot_native(config)`, which builds a small randomly-initialized TL-native model on the same driver — no HuggingFace Hub call — useful for tests and toy-model experiments.
 
-#### Granite MoE Hybrid grouped experts on Apple MPS
-
-The local transformers backend applies a narrow correctness fallback for the exact
-HF `GraniteMoeHybridExperts` class when all of these conditions hold:
-
-- PyTorch release **2.11.0** and transformers release **5.13.0**, including local
-  version suffixes but excluding prerelease and development versions.
-- The live expert input is on **MPS**.
-- The expert config selects **`grouped_mm`**.
-
-In this verified environment, grouped expert matrix products can be incorrect or
-unstable. The HF grouped fallback also uses nonzero-offset `torch.mm(..., out=...)`
-slice writes that reproduce incorrect results, so disabling only native grouped
-GEMM is insufficient. The bridge selects HF's original ordinary eager expert
-forward instead, emitting a `RuntimeWarning` about the correctness fallback and
-its possible performance cost. Ordinary eager execution may be slower; no
-throughput guarantee is made.
-
-Dispatch is installed on the expert instance when its containing `MoEBridge`
-runs, and checks the live input device and expert policy on every expert call.
-CPU-to-MPS and MPS-to-CPU migration therefore follow the current device. Shared
-HF configs, global dispatch tables, weights, routing indices, and routing weights
-are not mutated. Expert module hooks still execute, including output edits, and
-input, routing-weight, and expert-parameter gradients remain connected.
-
-The fallback does not serialize forwards or mutate shared configuration while
-other layers run. This avoids a shared-policy race, but does not guarantee that
-arbitrary concurrent model calls, device migration, config edits, or hook edits
-are safe. Direct HF execution before the containing bridge's first call is not
-protected; after installation, calls to that same expert instance use the local
-wrapper. Custom instance forwards are left unchanged.
-
-CPU/CUDA inputs, explicit eager execution, other expert policies, and other
-expert classes retain their existing dispatch. The observed defect and regression
-validation cover only the versions above, Apple MPS, and float32. Older/newer
-releases are neither declared affected nor declared fixed, and CUDA behavior has
-not been device-validated for this fallback. Unverified versions receive no
-workaround. MPS remains opt-in via `TRANSFORMERLENS_ALLOW_MPS=1`; this fallback is
-not a general guarantee of MPS correctness.
-
-##### Contributor diagnostics
-
-From the repository root, the offline regression selection requires no Hub
-model download and checks independent expert-oracle agreement, repeatability,
-device migration, hooks, gradients, and failure recovery:
-
-```bash
-TRANSFORMERLENS_ALLOW_MPS=1 uv run pytest \
-  tests/mps/test_granite_grouped_expert_dispatch.py -q
-uv run pytest \
-  tests/unit/model_bridge/generalized_components/test_moe_expert_dispatch.py -q
-TRANSFORMERLENS_ALLOW_MPS=1 uv run pytest \
-  tests/integration/model_bridge/test_granite_moe_hybrid_adapter.py -q
-```
-
-The primitive diagnostic below is independent of TransformerLens and HF expert
-routing. It reports error rather than asserting that an upstream kernel must
-remain defective. Run it on an MPS-capable host and record the installed versions
-when evaluating an upstream fix:
-
-```python
-import platform
-
-import torch
-import transformers
-
-if not torch.backends.mps.is_available():
-    raise RuntimeError("This diagnostic requires MPS")
-
-print(platform.system(), platform.release(), platform.machine())
-print("torch:", torch.__version__, "transformers:", transformers.__version__)
-generator = torch.Generator().manual_seed(0)
-inputs = torch.randn(10, 64, generator=generator)
-weights = torch.randn(4, 64, 32, generator=generator)
-expected = inputs[5:] @ weights[3]
-x, w = inputs.to("mps"), weights.to("mps")
-with torch.no_grad():
-    output = torch.zeros(10, 32, device="mps")
-    torch.mm(x[5:], w[3], out=output[5:])
-    assigned = torch.zeros_like(output)
-    assigned[5:] = x[5:] @ w[3]
-    torch.mps.synchronize()
-    print("out-slice max error:", (output[5:].cpu() - expected).abs().max().item())
-    print("assignment max error:", (assigned[5:].cpu() - expected).abs().max().item())
-```
-
-No dependency pin change or confirmed upstream fixed version is implied by this
-workaround. Revalidate both the primitive and full expert regressions before
-changing its scope.
-
 ### vLLM — high-throughput capture + declarative interventions
 
 `RemoteBridge.boot_vllm` constructs a vLLM engine (PagedAttention, `torch.compile`, CUDA graphs) and installs capture hooks inside the worker *before* compilation. Activations come back over `collective_rpc` and replay through the bridge's HookPoint tree. This is the throughput path for SAE/probe data collection; unlike observation-only tools (vllm-lens), each hook also applies an affine transform `output = output * scale + bias`, so declarative interventions (`suppress` / `scale` / `add` / `set`) propagate to downstream layers.
diff --git a/docs/source/content/getting_started.md b/docs/source/content/getting_started.md
index 518daf4393..5160f622bb 100644
--- a/docs/source/content/getting_started.md
+++ b/docs/source/content/getting_started.md
@@ -65,6 +65,11 @@ Set to `"1"` to opt in to Apple Silicon (MPS) as a target device for model infer
 export TRANSFORMERLENS_ALLOW_MPS=1
 ```
 
+For Hugging Face MoE experts using the shared grouped dispatcher, `TransformerBridge`
+checks the MPS matmul output-slice primitive and uses eager experts only when incorrect
+writes are detected. This emits a once-per-process warning and may reduce performance.
+Setting `TRANSFORMERLENS_ALLOW_MPS=1` does not disable this correctness safeguard.
+
 ## Huggingface Gated Access
 
 For convenience, gated-model access depends only on `HF_TOKEN` above. Once you have set the token and accepted any model-specific agreements on the HuggingFace Hub, gated models load through TransformerLens with no additional configuration. The most popular gated families supported by TransformerLens are the Llama, Mistral/Mixtral, and Gemma models.

From 3eccd6b18beb36114f401afb264134e7904bfd00 Mon Sep 17 00:00:00 2001
From: janmenjayap <janmenjaya.panda.22@gmail.com>
Date: Sun, 18 Oct 2026 17:06:44 +0530
Subject: [PATCH 8/8] fix(moe): harden MPS expert dispatch

---
 .../test_granite_grouped_expert_dispatch.py   |  38 ++++-
 .../test_moe_expert_dispatch.py               | 140 ++++++++++++++++--
 .../generalized_components/moe.py             |  75 ++++++----
 3 files changed, 209 insertions(+), 44 deletions(-)

diff --git a/tests/mps/test_granite_grouped_expert_dispatch.py b/tests/mps/test_granite_grouped_expert_dispatch.py
index 193ad3aba8..dc0be79f5b 100644
--- a/tests/mps/test_granite_grouped_expert_dispatch.py
+++ b/tests/mps/test_granite_grouped_expert_dispatch.py
@@ -25,7 +25,9 @@
 pytestmark = pytest.mark.skipif(not torch.backends.mps.is_available(), reason="MPS not available")
 
 
-@pytest.fixture(params=[GraniteMoeHybridExperts, GraniteMoeExperts, MixtralExperts, Qwen3MoeExperts])
+@pytest.fixture(
+    params=[GraniteMoeHybridExperts, GraniteMoeExperts, MixtralExperts, Qwen3MoeExperts]
+)
 def expert_class(request):
     return request.param
 
@@ -276,7 +278,9 @@ def forward(self, value):
 
 
 @pytest.mark.parametrize("broken", [False, True])
-def test_real_expert_dispatch_selects_grouped_or_eager_from_probe(monkeypatch, expert_class, broken):
+def test_real_expert_dispatch_selects_grouped_or_eager_from_probe(
+    monkeypatch, expert_class, broken
+):
     grouped_calls = []
 
     def safe_grouped(inputs, weights, offs):
@@ -301,9 +305,37 @@ def safe_grouped(inputs, weights, offs):
     with torch.no_grad(), warnings.catch_warnings(record=True) as caught:
         warnings.simplefilter("always")
         actual = bridge(layer_input=inputs.to("mps"))
-    assert block.experts.forward.__func__ is moe._mps_experts_forward
+    assert block.experts.forward.__wrapped__ is type(block.experts).forward
     torch.testing.assert_close(actual.cpu(), expected, rtol=1e-4, atol=1e-6)
     assert len(caught) == int(broken)
     assert grouped_calls == ([] if broken else [[2, 5, 5, 10]] * 2)
     assert block.experts.config._experts_implementation == "grouped_mm"
     assert dict(hf_moe.ALL_EXPERTS_FUNCTIONS) == registry
+
+
+@pytest.mark.parametrize("weight_dtype", [torch.float32, torch.float16])
+def test_autocast_cold_and_warm_dispatch_matches_raw_hf_policy(expert_class, weight_dtype):
+    block = _FixedRoutes(expert_class).to("mps")
+    block.experts.to(dtype=weight_dtype)
+    reference = copy.deepcopy(block)
+    bridge = MoEBridge("moe")
+    bridge.set_original_component(block)
+    inputs = torch.randn(10, 64, generator=torch.Generator().manual_seed(1)).to("mps")
+    assert not moe._MPS_OUTPUT_SLICE_RESULTS
+    with torch.no_grad(), warnings.catch_warnings(record=True) as caught:
+        warnings.simplefilter("always")
+        with torch.autocast("mps", dtype=torch.float16):
+            first = bridge(inputs)
+            assert torch.is_autocast_enabled("mps")
+            key = (inputs.device, weight_dtype)
+            assert set(moe._MPS_OUTPUT_SLICE_RESULTS) == {key}
+            broken = moe._MPS_OUTPUT_SLICE_RESULTS[key]
+            reference.experts.config._experts_implementation = "eager" if broken else "grouped_mm"
+            expected = reference(inputs)
+            second = bridge(inputs)
+            assert torch.is_autocast_enabled("mps")
+    torch.testing.assert_close(first, expected, rtol=0, atol=0)
+    torch.testing.assert_close(second, first, rtol=0, atol=0)
+    assert len(caught) == int(broken)
+    assert set(moe._MPS_OUTPUT_SLICE_RESULTS) == {key}
+    assert block.experts.config._experts_implementation == "grouped_mm"
diff --git a/tests/unit/model_bridge/generalized_components/test_moe_expert_dispatch.py b/tests/unit/model_bridge/generalized_components/test_moe_expert_dispatch.py
index 44e641d839..3fd3b680e2 100644
--- a/tests/unit/model_bridge/generalized_components/test_moe_expert_dispatch.py
+++ b/tests/unit/model_bridge/generalized_components/test_moe_expert_dispatch.py
@@ -5,6 +5,8 @@
 import warnings
 from concurrent.futures import ThreadPoolExecutor
 from functools import wraps
+from inspect import signature
+from threading import Lock, Thread
 from types import SimpleNamespace
 
 import pytest
@@ -21,9 +23,9 @@
 
 from transformer_lens.model_bridge.generalized_components import moe
 from transformer_lens.model_bridge.generalized_components.moe import (
-    _mps_experts_forward,
     _guard_mps_experts,
     _hf_eager_experts_forward,
+    _mps_experts_forward,
     _mps_output_slice_writes_are_broken,
     _probe_mps_output_slice_writes,
 )
@@ -107,9 +109,7 @@ def test_exception_does_not_leave_dispatch_state(stub_dispatch):
     experts.fail = False
     with warnings.catch_warnings(record=True) as caught:
         warnings.simplefilter("always")
-        assert (
-            _mps_experts_forward(experts, torch.empty(1).as_subclass(_MpsInput)) == "eager"
-        )
+        assert _mps_experts_forward(experts, torch.empty(1).as_subclass(_MpsInput)) == "eager"
     assert not caught
 
 
@@ -157,7 +157,7 @@ def probe(device, dtype):
     monkeypatch.setattr(moe, "_probe_mps_output_slice_writes", probe)
     block = _block()
     _guard_mps_experts(block)
-    assert block.experts.forward.__func__ is _mps_experts_forward
+    assert block.experts.forward.__wrapped__ is type(block.experts).forward
     assert not probes
     assert block.experts.config._experts_implementation == "grouped_mm"
     experts = _Experts(block.experts.config)
@@ -245,9 +245,7 @@ def probe(device, dtype):
     assert not moe._MPS_OUTPUT_SLICE_RESULTS
     assert experts.config._experts_implementation == "grouped_mm"
     monkeypatch.setattr(moe, "_probe_mps_output_slice_writes", lambda device, dtype: False)
-    assert (
-        _mps_experts_forward(experts, torch.empty(1).as_subclass(_MpsInput)) == "grouped_mm"
-    )
+    assert _mps_experts_forward(experts, torch.empty(1).as_subclass(_MpsInput)) == "grouped_mm"
 
 
 def test_unexpected_probe_failure_propagates(monkeypatch):
@@ -282,9 +280,7 @@ def probe(device, dtype):
     experts = _Experts(SimpleNamespace(_experts_implementation=policy))
     with warnings.catch_warnings(record=True) as caught:
         warnings.simplefilter("always")
-        assert (
-            _mps_experts_forward(experts, torch.empty(1).as_subclass(DeviceInput)) == policy
-        )
+        assert _mps_experts_forward(experts, torch.empty(1).as_subclass(DeviceInput)) == policy
     assert not caught
     assert not moe._MPS_OUTPUT_SLICE_RESULTS
 
@@ -296,12 +292,15 @@ def test_installation_is_idempotent_and_deepcopy_safe(monkeypatch, expert_class)
     monkeypatch.setattr(torch, "__version__", "2.11.0")
     monkeypatch.setattr(transformers, "__version__", "5.13.0")
     block = _block(expert_class)
+    original_signature = signature(block.experts.forward)
     _guard_mps_experts(block)
     forward = block.experts.forward
+    assert signature(forward) == original_signature
     _guard_mps_experts(block)
     assert block.experts.forward is forward
     cloned = copy.deepcopy(block)
     assert cloned.experts.forward.__self__ is cloned.experts
+    assert signature(cloned.experts.forward) == original_signature
 
 
 @pytest.mark.parametrize("expert_class", _EXPERT_CLASSES)
@@ -429,7 +428,7 @@ def grouped(experts, hidden_states, token=None, *, payload=None):
     monkeypatch.setattr(moe, "_probe_mps_output_slice_writes", lambda device, dtype: broken)
     block = _decorated_block(interface)
     _guard_mps_experts(block)
-    assert block.experts.forward.__func__ is _mps_experts_forward
+    assert block.experts.forward.__wrapped__ is type(block.experts).forward
     inputs = torch.empty(1).as_subclass(_MpsInput)
     token, payload = object(), object()
     with warnings.catch_warnings(record=True) as caught:
@@ -490,7 +489,7 @@ def test_interface_changed_after_installation_is_respected():
     interface = ExpertsInterface()
     block = _decorated_block(interface)
     _guard_mps_experts(block)
-    assert block.experts.forward.__func__ is _mps_experts_forward
+    assert block.experts.forward.__wrapped__ is type(block.experts).forward
     interface["grouped_mm"] = lambda experts, hidden_states: "custom"
     with warnings.catch_warnings(record=True) as caught:
         assert block.experts(torch.empty(1).as_subclass(_MpsInput)) == "custom"
@@ -618,3 +617,118 @@ def test_fallback_warning_does_not_suppress_unrelated_warnings():
         assert block.experts(inputs) == "eager"
         warnings.warn("Unrelated diagnostic", UserWarning)
     assert [warning.category for warning in caught] == [RuntimeWarning, UserWarning, UserWarning]
+
+
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+@pytest.mark.parametrize("broken", [False, True])
+def test_primitive_probe_is_independent_of_ambient_autocast(monkeypatch, dtype, broken):
+    original_mm = torch.mm
+
+    def mm(inputs, weights, *, out):
+        assert not torch.is_autocast_enabled("cpu")
+        assert inputs.dtype == weights.dtype == out.dtype == dtype
+        if broken:
+            return out
+        return original_mm(inputs, weights, out=out)
+
+    monkeypatch.setattr(torch, "mm", mm)
+    with torch.autocast("cpu", dtype=torch.bfloat16):
+        assert _probe_mps_output_slice_writes(torch.device("cpu"), dtype) == broken
+        assert torch.is_autocast_enabled("cpu")
+
+
+def test_dispatch_probes_batched_weight_dtype_instead_of_input_dtype(monkeypatch):
+    block = _decorated_block(ExpertsInterface())
+    block.experts.register_parameter(
+        "gate_up_proj", nn.Parameter(torch.empty(4, 8, 8, dtype=torch.float16))
+    )
+    block.experts.register_parameter(
+        "down_proj", nn.Parameter(torch.empty(4, 8, 4, dtype=torch.float16))
+    )
+    probed = []
+
+    def probe(device, dtype):
+        probed.append(dtype)
+        return True
+
+    monkeypatch.setattr(moe, "_probe_mps_output_slice_writes", probe)
+    _guard_mps_experts(block)
+    with pytest.warns(RuntimeWarning, match="eager experts"):
+        assert block.experts(torch.empty(1).as_subclass(_MpsInput)) == "eager"
+    assert probed == [torch.float16]
+
+
+def test_warning_callback_can_reenter_guarded_experts(monkeypatch):
+    monkeypatch.setattr(moe, "_MPS_EXPERTS_WARNING_LOCK", Lock())
+    block = _decorated_block(ExpertsInterface())
+    _guard_mps_experts(block)
+    inputs = torch.empty(1).as_subclass(_MpsInput)
+    callbacks, results, errors = [], [], []
+
+    def showwarning(*args, **kwargs):
+        callbacks.append(block.experts(inputs))
+
+    def invoke():
+        try:
+            results.append(block.experts(inputs))
+        except Exception as error:
+            errors.append(error)
+
+    with warnings.catch_warnings():
+        warnings.simplefilter("always")
+        monkeypatch.setattr(warnings, "showwarning", showwarning)
+        worker = Thread(target=invoke, daemon=True)
+        worker.start()
+        worker.join(timeout=2)
+    assert not worker.is_alive(), "warning callback deadlocked on re-entry"
+    assert not errors
+    assert callbacks == ["eager"]
+    assert results == ["eager"]
+    assert moe._MPS_EXPERTS_WARNED
+
+
+def test_warning_callback_failure_releases_reservation(monkeypatch):
+    block = _decorated_block(ExpertsInterface())
+    _guard_mps_experts(block)
+    inputs = torch.empty(1).as_subclass(_MpsInput)
+
+    def showwarning(*args, **kwargs):
+        raise ValueError("warning callback failed")
+
+    with warnings.catch_warnings(), monkeypatch.context() as context:
+        warnings.simplefilter("always")
+        context.setattr(warnings, "showwarning", showwarning)
+        with pytest.raises(ValueError, match="warning callback failed"):
+            block.experts(inputs)
+    assert not moe._MPS_EXPERTS_WARNED
+    with pytest.warns(RuntimeWarning, match="eager experts"):
+        assert block.experts(inputs) == "eager"
+
+
+def test_dispatch_checks_both_projection_weight_dtypes(monkeypatch):
+    def grouped(experts, hidden_states):
+        return "grouped_mm"
+
+    interface = ExpertsInterface()
+    interface["grouped_mm"] = grouped
+    monkeypatch.setattr(hf_moe, "grouped_mm_experts_forward", grouped)
+    block = _decorated_block(interface)
+    block.experts.register_parameter(
+        "gate_up_proj", nn.Parameter(torch.empty(4, 8, 8, dtype=torch.float16))
+    )
+    block.experts.register_parameter(
+        "down_proj", nn.Parameter(torch.empty(4, 8, 4, dtype=torch.float32))
+    )
+    probed = []
+
+    def probe(device, dtype):
+        probed.append(dtype)
+        return False
+
+    monkeypatch.setattr(moe, "_probe_mps_output_slice_writes", probe)
+    _guard_mps_experts(block)
+    with warnings.catch_warnings(record=True) as caught:
+        warnings.simplefilter("always")
+        assert block.experts(torch.empty(1).as_subclass(_MpsInput)) == "grouped_mm"
+    assert set(probed) == {torch.float16, torch.float32}
+    assert not caught
diff --git a/transformer_lens/model_bridge/generalized_components/moe.py b/transformer_lens/model_bridge/generalized_components/moe.py
index d4117db270..be46ed1f55 100644
--- a/transformer_lens/model_bridge/generalized_components/moe.py
+++ b/transformer_lens/model_bridge/generalized_components/moe.py
@@ -7,6 +7,7 @@
 
 import logging
 import warnings
+from functools import wraps
 from inspect import currentframe, getclosurevars
 from threading import Lock
 from types import CodeType, FunctionType, MethodType
@@ -45,32 +46,38 @@ def _warn_mps_experts_once() -> None:
     with _MPS_EXPERTS_WARNING_LOCK:
         if _MPS_EXPERTS_WARNED:
             return
+        _MPS_EXPERTS_WARNED = True
+    frame = None
+    emitted = False
+    try:
         frame = currentframe()
         stacklevel = 1
-        try:
-            while frame is not None:
-                module = frame.f_globals.get("__name__", "")
-                if not any(
-                    module == prefix or module.startswith(prefix + ".")
-                    for prefix in ("transformer_lens", "transformers", "torch", "jaxtyping", "beartype")
-                ):
-                    break
-                frame = frame.f_back
-                stacklevel += 1
-            warnings.warn(
-                "Using eager experts on MPS after detecting incorrect "
-                "matmul output-slice writes; expert execution may be slower.",
-                RuntimeWarning,
-                stacklevel=stacklevel,
-            )
-            _MPS_EXPERTS_WARNED = True
-        finally:
-            del frame
+        while frame is not None:
+            module = frame.f_globals.get("__name__", "")
+            if not any(
+                module == prefix or module.startswith(prefix + ".")
+                for prefix in ("transformer_lens", "transformers", "torch", "jaxtyping", "beartype")
+            ):
+                break
+            frame = frame.f_back
+            stacklevel += 1
+        warnings.warn(
+            "Using eager experts on MPS after detecting incorrect "
+            "matmul output-slice writes; expert execution may be slower.",
+            RuntimeWarning,
+            stacklevel=stacklevel,
+        )
+        emitted = True
+    finally:
+        del frame
+        if not emitted:
+            with _MPS_EXPERTS_WARNING_LOCK:
+                _MPS_EXPERTS_WARNED = False
 
 
 def _probe_mps_output_slice_writes(device: torch.device, dtype: torch.dtype) -> bool:
     """Check for incorrect nonzero-offset matmul output writes using isolated tensors."""
-    with torch.no_grad():
+    with torch.no_grad(), torch.autocast(device_type=device.type, enabled=False):
         generator = torch.Generator(device="cpu").manual_seed(0)
         inputs = torch.randn(10, 64, generator=generator).to(device=device, dtype=dtype)
         weights = torch.randn(4, 64, 32, generator=generator).to(device=device, dtype=dtype)
@@ -110,11 +117,17 @@ def _mps_experts_forward(experts: nn.Module, *args: Any, **kwargs: Any) -> Any:
         and getattr(config, "_experts_implementation", None) == "grouped_mm"
     ):
         eager = _hf_eager_experts_forward(forward)
-    if eager is not None and _mps_output_slice_writes_are_broken(
-        hidden_states.device, hidden_states.dtype
-    ):
-        _warn_mps_experts_once()
-        return eager(experts, *args, **kwargs)
+        if eager is not None:
+            dtypes = {
+                parameter.dtype
+                for parameter in experts.parameters(recurse=False)
+                if parameter.ndim == 3 and parameter.is_floating_point()
+            } or {hidden_states.dtype}
+            if any(
+                _mps_output_slice_writes_are_broken(hidden_states.device, dtype) for dtype in dtypes
+            ):
+                _warn_mps_experts_once()
+                return eager(experts, *args, **kwargs)
     return forward(experts, *args, **kwargs)
 
 
@@ -152,8 +165,9 @@ def _hf_eager_experts_forward(forward: Any) -> Optional[Callable[..., Any]]:
 def _guard_mps_experts(component: nn.Module) -> None:
     """Install instance-local dispatch only for HF's shared grouped expert implementation.
 
-    The bound method preserves nn.Module hooks and deepcopy semantics. Configs and
-    global dispatch tables remain untouched, including during concurrent forwards.
+    The bound method preserves the forward signature, nn.Module hooks, and deepcopy
+    semantics. Configs and global dispatch tables remain untouched, including during
+    concurrent forwards.
     """
     experts = getattr(component, "experts", None)
     if not isinstance(experts, nn.Module):
@@ -164,7 +178,12 @@ def _guard_mps_experts(component: nn.Module) -> None:
         return
     if _hf_eager_experts_forward(forward) is None:
         return
-    setattr(experts, "forward", MethodType(_mps_experts_forward, experts))
+
+    @wraps(forward)
+    def guarded_forward(experts: nn.Module, *args: Any, **kwargs: Any) -> Any:
+        return _mps_experts_forward(experts, *args, **kwargs)
+
+    setattr(experts, "forward", MethodType(guarded_forward, experts))
 
 
 class UnfoldableMoEParameter(Exception):