[MoE] Count tokens per expert with scatter_add on MPS (torch.histc is ~45 ms per call there)

[nassersala]

What

grouped_mm_experts_forward (integrations/moe.py) computes tokens_per_expert with torch.histc. On MPS, histc
costs about 0.17 ms per bin regardless of input size: 11 ms at 64 experts, 22 ms at 128, 45 ms at 256, once per MoE
layer, so once per layer per generated token. grouped_mm is the default experts implementation, so this hits every
MoE model run on Apple Silicon.

This PR counts with scatter_add_ on MPS instead. Ids at or above num_experts (EP sentinels) go to an extra bin that is
dropped, which is what histc(max=num_experts - 1) does with them. CPU and CUDA paths are unchanged, and there is no
data-dependent shape, so nothing changes for graph capture.

-    histc_input = expert_ids_g.float() if device.type in ("cpu", "mps") else expert_ids_g.int()
-    tokens_per_expert = torch.histc(histc_input, bins=self.num_experts, min=0, max=self.num_experts - 1)
+    if device.type == "mps":
+        # torch.histc costs ~0.17 ms per bin on MPS; count with scatter_add. Out-of-range ids (EP sentinels) land in
+        # an extra bin that is dropped, as histc(max=num_experts - 1) drops them.
+        tokens_per_expert = torch.zeros(self.num_experts + 1, device=device, dtype=torch.int32).scatter_add_(
+            0, expert_ids_g.clamp(max=self.num_experts), torch.ones_like(expert_ids_g, dtype=torch.int32)
+        )[:-1]
+    else:
+        histc_input = expert_ids_g.float() if device.type == "cpu" else expert_ids_g.int()
+        tokens_per_expert = torch.histc(histc_input, bins=self.num_experts, min=0, max=self.num_experts - 1)

Measurements (Apple M5 Max, macOS 26.5.2, torch 2.14.0, transformers 5.17.0)

histc vs scatter_add alone (bench_mps.py below):

ids	experts	histc	scatter_add
8	64	11.5 ms	0.024 ms
1024	128	21.8 ms	0.023 ms
944	256	43.9 ms	0.024 ms
65536	256	45.1 ms	0.059 ms

Greedy decoding, random-weight Qwen3-MoE with Qwen3-30B-A3B's expert layout (48 layers, 128 experts, top 8), bf16
(bench_moe_decode.py below):

experts implementation	ms / token
grouped_mm (current default)	48.7
grouped_mm, this PR	14.4
eager	451.3

Same tokens in all three. On a real model (Mapika/decider-35b-a3b, Qwen3.5-35B-A3B architecture, 256 experts), prefill of
a 96-token input goes from 2.8 s to 1.0 s with this change alone, with bit-identical output probabilities.

Checked on this branch (current main): the same model decodes at 46.8 -> 13.0 ms/token against 5.17.0, with bit-identical
logits and tokens.

Correctness: the cumulative offsets match histc exactly on 50 random id sets of up to 5,000 ids that include sentinels.

Related PyTorch issue: pytorch/pytorch#198304.

Follow-up (separate PR or issue, happy to open it)

The gated-delta-rule reference path (torch_chunk_gated_delta_rule, used by qwen3_5, qwen3_5_moe, qwen3_next, olmo_hybrid,
qwen4_exp) calls torch.linalg.solve_triangular twice per layer. On MPS that is 18 ms for 32 heads x 10 chunks and
118 ms for 32 x 64, against 1.4 ms and 8.8 ms on CPU. Two MPS options, both checked against the MPS solver:

• the existing is_torchdynamo_exporting() forward-substitution loop: 3.3 ms / 11.8 ms, no new code;
• a block-doubling inverse, [[A11,0],[A21,A22]]^-1 = [[X11,0],[-X22 A21 X11, X22]], all batched matmuls: 0.40 ms /
  1.4 ms, max abs difference 6e-7 on realistic systems.

On decider-35b-a3b both changes together take a 600-token prefill from 3.9 s to 0.47 s. Output probabilities move by at
most 0.033, less than switching to the exact CPU solver moves them (0.042); no argmax changed on six items, and the
model's JevBench public numbers are reproduced exactly (48/48, 70/72, 75/111).

(A Neumann-series inverse (I - N)(I + N^2)... looks tempting but gives NaNs on real inputs, whose entries are close to 1.)

bench_mps.py

"""MPS slow paths hit by transformers MoE / gated-delta-rule models. Run: python bench_mps.py"""
import platform, subprocess, time
import torch


def ms(f, n=20):
    f(); torch.mps.synchronize(); t = time.perf_counter()
    for _ in range(n): f()
    torch.mps.synchronize(); return (time.perf_counter() - t) / n * 1000


print(f"torch {torch.__version__}, macOS {platform.mac_ver()[0]}, "
      f"{subprocess.run(['sysctl', '-n', 'machdep.cpu.brand_string'], capture_output=True, text=True).stdout.strip()}")

print("\n1. histc of expert ids (float, bins = num_experts, min 0, max num_experts - 1)")
print(f"{'ids':>8} {'experts':>8} {'histc':>10} {'scatter_add':>12} {'bincount':>10}   equal")
for n_ids, E in ((8, 64), (1024, 128), (944, 256), (8192, 256), (65536, 256)):
    ids = torch.randint(0, E, (n_ids,), device="mps")
    f = ids.float()
    h = lambda: torch.histc(f, bins=E, min=0, max=E - 1)
    s = lambda: torch.zeros(E, device="mps").scatter_add_(0, ids, torch.ones(n_ids, device="mps"))
    b = lambda: torch.bincount(ids, minlength=E)
    print(f"{n_ids:>8} {E:>8} {ms(h):>8.2f}ms {ms(s):>10.3f}ms {ms(b):>8.3f}ms   {torch.equal(h(), s())}")

print("\n2. unit lower-triangular solve, gated delta rule shapes (fp32, chunk 64, v_dim 128)")


def block_inverse(A):  # [[A11,0],[A21,A22]]^-1 = [[X11,0],[-X22 A21 X11, X22]], block size 1 -> n
    n = A.shape[-1]; inv = torch.ones(*A.shape[:-2], n, 1, 1, device=A.device, dtype=A.dtype); b = 1
    while b < n:
        nb = n // (2 * b)
        blk = A.reshape(*A.shape[:-2], nb, 2 * b, nb, 2 * b).diagonal(dim1=-4, dim2=-2).movedim(-1, -3)
        x11, x22 = inv[..., 0::2, :, :], inv[..., 1::2, :, :]
        new = torch.zeros(*A.shape[:-2], nb, 2 * b, 2 * b, device=A.device, dtype=A.dtype)
        new[..., :b, :b], new[..., b:, b:], new[..., b:, :b] = x11, x22, -(x22 @ blk[..., b:, :b] @ x11)
        inv, b = new, 2 * b
    return inv[..., 0, :, :]


def export_loop(A):  # transformers' existing is_torchdynamo_exporting() path
    u = -A.tril(-1)
    for i in range(1, A.shape[-1]):
        row, sub = u[..., i, :i].clone(), u[..., :i, :i].clone()
        u[..., i, :i] = row + (row.unsqueeze(-1) * sub).sum(-2)
    return u + torch.eye(A.shape[-1], device=A.device)


print(f"{'heads x chunks':>15} {'MPS solve':>10} {'CPU solve':>10} {'export loop':>12} {'block inv':>10}   max |diff| block vs MPS solve")
for heads, chunks in ((32, 2), (32, 10), (32, 64)):
    g = torch.Generator().manual_seed(0)
    k = torch.nn.functional.normalize(torch.randn(1, heads, chunks, 64, 128, generator=g), dim=-1)
    beta, decay = torch.rand(1, heads, chunks, 64, 1, generator=g), torch.rand(1, heads, chunks, 64, generator=g) * -0.05
    cd = decay.cumsum(-1); pw = (cd.unsqueeze(-1) - cd.unsqueeze(-2)).masked_fill(torch.ones(64, 64, dtype=torch.bool).triu(1), float("-inf")).exp()
    A = ((k * beta) @ k.transpose(-1, -2) * pw).tril(-1) + torch.eye(64)  # realistic: entries up to ~1
    B = torch.randn(1, heads, chunks, 64, 128, generator=g)
    Am, Bm = A.to("mps"), B.to("mps")
    ref = torch.linalg.solve_triangular(Am, Bm, upper=False, unitriangular=True)
    t_mps = ms(lambda: torch.linalg.solve_triangular(Am, Bm, upper=False, unitriangular=True))
    t0 = time.perf_counter(); [torch.linalg.solve_triangular(A, B, upper=False, unitriangular=True) for _ in range(20)]; t_cpu = (time.perf_counter() - t0) / 20 * 1000
    t_loop = ms(lambda: export_loop(Am.clone()) @ Bm, n=3)
    t_blk = ms(lambda: block_inverse(Am) @ Bm)
    diff = (block_inverse(Am) @ Bm - ref).abs().max().item()
    print(f"{heads:>7} x {chunks:<5} {t_mps:>8.2f}ms {t_cpu:>8.2f}ms {t_loop:>10.2f}ms {t_blk:>8.2f}ms   {diff:.1e} (|x| max {ref.abs().max().item():.1f})")

bench_moe_decode.py

"""Decode speed of a small random-weight Qwen3-MoE (Qwen3-30B-A3B's expert layout: 48 layers, 128 experts, top 8) on MPS,
with the default experts implementation, with transformers' histc replaced by scatter_add, and with eager experts."""
import time
import torch
from transformers import Qwen3MoeConfig, Qwen3MoeForCausalLM

cfg = Qwen3MoeConfig(vocab_size=1000, hidden_size=256, intermediate_size=512, moe_intermediate_size=128, num_hidden_layers=48,
                     num_attention_heads=4, num_key_value_heads=2, head_dim=64, num_experts=128, num_experts_per_tok=8)
torch.manual_seed(0)
_histc = torch.histc


def histc(x, bins=100, min=0, max=0, **kw):
    if x.device.type == "mps" and x.dim() == 1 and min == 0 and bins == max + 1 and not kw:
        # histc drops out-of-range ids (EP sentinels): send them to an extra bin and cut it
        return torch.zeros(bins + 1, device=x.device, dtype=x.dtype).scatter_add_(0, x.long().clamp(max=bins), torch.ones_like(x))[:bins]
    return _histc(x, bins=bins, min=min, max=max, **kw)


ids = torch.randint(0, 1000, (1, 32), device="mps")
out = {}
for name, impl, patch in (("grouped_mm (default)", None, False), ("grouped_mm, histc -> scatter_add", None, True), ("eager experts", "eager", False)):
    torch.histc = histc if patch else _histc
    m = Qwen3MoeForCausalLM(cfg) if impl is None else Qwen3MoeForCausalLM._from_config(cfg, experts_implementation=impl)
    m = m.to("mps", torch.bfloat16).eval()
    torch.manual_seed(0); m.load_state_dict(ref_sd) if out else None
    ref_sd = m.state_dict() if not out else ref_sd
    with torch.no_grad():
        m.generate(ids, max_new_tokens=4, do_sample=False)
        torch.mps.synchronize(); t = time.perf_counter()
        g = m.generate(ids, max_new_tokens=32, min_new_tokens=32, do_sample=False)
        torch.mps.synchronize(); dt = time.perf_counter() - t
    out[name] = g
    print(f"{name:<34} {dt / 32 * 1000:7.1f} ms/token   {m.config._experts_implementation if hasattr(m.config, '_experts_implementation') else ''}")
torch.histc = _histc
first = next(iter(out.values()))
print("same tokens as default:", {k: torch.equal(v, first) for k, v in out.items()})

🤖 Generated with Claude Code

[github-actions]

---

CI recap

Dashboard: View test results in Grafana
Latest run: 35806753970:2
Result: success | Jobs: 16 | Tests: 189,590 | Failures: 0 | Duration: 16h 43m

[Isalia20]

Please hold off from merging before i check the numbers here:
pytorch/pytorch#198304

[Isalia20]

Let's try to make histc better in pytorch rather than add a workaround, I'll take a look