fix(moe): avoid incorrect Granite grouped expert execution on MPS

[janmenjayap]

---

Description

Fixes #1886

Granite MoE Hybrid expert execution can produce incorrect or unstable outputs on Apple MPS when HuggingFace selects grouped_mm in the verified torch 2.11.0 / transformers 5.13.0 environment. This can make the existing SSM eager/fused parity regression fail even though the first substantial divergence occurs in the MoE expert output rather than the SSM scan.

A standalone fixed-data diagnostic isolates incorrect nonzero-offset torch.mm(..., out=output[5:]) writes. HF's grouped fallback also uses output-slice writes, so disabling native grouped execution alone is insufficient.

---

Changes

• Add an instance-local dispatch wrapper for the exact HF GraniteMoeHybridExperts class, installed through MoEBridge.forward.
• Select HF's original ordinary eager expert forward only for live MPS inputs with grouped_mm on torch release 2.11.0 and transformers release 5.13.0. Local version suffixes are accepted; prerelease/development versions are not.
• Leave shared HF configs and global dispatch tables untouched. No production activation transfers to CPU, routing reimplementation, dependency changes, or global kernel patches are introduced.
• Preserve CPU/CUDA dispatch, explicit eager execution, other policies/classes, and custom instance forwards.
• Add independent expert-oracle and repeatability tests, CPU-to-MPS-to-CPU migration coverage, input/routing-weight/expert-parameter gradient comparisons, observer/output-editing hook coverage, and failure-recovery/state-preservation tests.
• Document the scope, warning, performance cost, concurrency boundary, and primitive diagnostic in the execution-backends guide.

The existing Granite MoE Hybrid MPS scan-parity assertion passes through the corrected production path without a changed threshold, new skip/xfail, or forced eager policy in its fixture.

---

Scope and limitations

• Verified locally on Darwin 23.6.0 arm64, Python 3.12.14, torch 2.11.0, transformers 5.13.0, float32, Apple MPS.
• The fallback deliberately emits a RuntimeWarning; eager expert execution may be slower.
• Dispatch follows the live input device and policy. Avoiding shared-config mutation removes a policy race, but does not guarantee 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. Once installed, the same expert instance retains the wrapper.
• Older/newer dependency versions are neither declared affected nor declared fixed. CUDA and other dependency versions were not device-validated.
• This workaround does not establish the internal root cause of native grouped execution or imply a confirmed upstream fixed version.

---

Relationship to other work

This is separate from #1878, which fixes Granite embedding scaling in cache residual decomposition and component-level direct logit attribution without changing model forward execution. That analysis correction does not repair incorrect expert matrix products.

It is also independent of GPT-2 contextual decomposition.

A potentially related upstream report is pytorch/pytorch#193487, concerning allocator-dependent incorrect MPS matmul results. Its reproduction differs from the small nonzero-offset output-write reproduction here; a shared root cause has not been established. No upstream report or comment has been posted for this work.

---

Type of change

• Bug fix (non-breaking change which fixes an issue)
• This change requires a documentation update

---

Validation

Before and after

On baseline a8bc19e7, the CPU scan-parity case passed and the MPS case failed with relative difference 1.0326683320021979 (required < 1e-4). On this branch, the original test passes without weakened assertions.

The standalone primitive diagnostic measured:

MPS out-slice max error: 29.132984161376953
MPS slice-assignment max error: 4.76837158203125e-06

Focused regressions

TRANSFORMERLENS_ALLOW_MPS=1 uv run pytest \
  tests/mps/test_granite_grouped_expert_dispatch.py \
  tests/unit/model_bridge/generalized_components/test_moe_expert_dispatch.py \
  tests/integration/model_bridge/test_granite_moe_hybrid_adapter.py \
  -q --tb=short

Result: 49 passed, with the intended fallback warnings and other dependency/model warnings reported.

Standard checks

Full unit tier: 7322 passed, 58 skipped, 57 deselected, 3 xfailed
Docstring tier: 13 passed, 24 skipped
Acceptance tier: 163 passed, 69 skipped, 11 deselected
Mypy: no issues found in 391 source files
Docs: build succeeded, 19 warnings

make format, make check-format, and git diff --check passed. The skip/xfail counts above are suite outcomes, not new markers added to bypass this defect.

Full PR validation remains incomplete. make test-pr completed unit, docstring, and acceptance tiers, then hit the 30-minute command limit during integration testing. No integration failure had been reported before interruption, but the integration tier is not claimed complete or passing. The docs build is not claimed warning-free. CUDA was not tested.

---

Checklist

• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes
• I have not rewritten tests relating to key interfaces which would affect backward compatibility

The warning checkbox is intentionally unchecked because the affected MPS path emits an explicit correctness-fallback warning. Completing the full integration tier remains a validation follow-up before merge.

---

[jlarson4]

Hi @janmenjayap, the idea behind your fix is sound, but there are some fundamental flaws that need to be resolved before this can be merged. Please take a look at my comments below, and let me know if you have any questions

[jlarson4]

This gates on torch 2.11.0 and transformers 5.13.0, but the bad slice writes are not tied to those versions. On macOS 26.6 with exactly those packages, the primitive, native grouped_mm, and HF's fallback are all correct and the parity test passes without this change, yet the gate still fires, so every Granite Hybrid forward on MPS takes the eager loop and receive a false warning saying its products are wrong. Please key the fallback on a direct check of the slice-write primitive rather than on package versions.

[jlarson4]

Only GraniteMoeHybridExperts is recognized, but the primitive that fails is shared. Every @use_experts_implementation class routes through the same _grouped_mm, and HF defaults all of them to grouped_mm with no device check. On an affected host a GraniteMoe, Mixtral, or Qwen3-MoE bridge produces the same wrong products with no fallback. Please adjust to recognize the shared dispatch decorator.

[jlarson4]

The warning is raised inside every expert call, and stacklevel=2 lands in jaxtyping's wrapper, so users see a library frame rather than their own call and, under non-default filters, one line per layer per forward. devices.py already keeps its MPS warnings to once per process with a module flag; consider the same here.

[jlarson4]

This subsection sits under the transformers backend in the driver guide, but it is an Apple-specific workaround. Would you please move it to the existing TRANSFORMERLENS_ALLOW_MPS section in getting_started.md? Consider reducing to a short note there, with the diagnostic kept in the test file or the issue, instead of this subsection.