accel.moe: drop padding ids through the routing tables; padding not counted (#89)

opt_prime/opt_prime/accel/moe/batched.py

@@ -52,7 +52,10 @@ def opt_grouped_mm(self, hidden_states, top_k_index, top_k_weights, device=None)
     gated = getattr(self, "has_gate", hasattr(self, "gate_up_proj"))
     lin = lambda x, w, b: (x @ w if tr else x @ w.t()) + (b if b is not None else 0)
     get = lambda name, e: getattr(self, name)[e].to(dev) if getattr(self, name, None) is not None else None
-    out = hidden_states.float() * 0 + top_k_weights.float().sum() * 0
+    # zeros that stay connected to both inputs (all-padding calls still get zero gradients); empty slices, so an
+    # inf in hidden_states does not turn into NaN
+    out = (torch.zeros(T, H, dtype=torch.float32, device=hidden_states.device)
+           + hidden_states[:0].float().sum() + top_k_weights[:0].float().sum())
     for e in torch.unique(flat).tolist():
         if not 0 <= e < self.num_experts:
             continue

opt_prime/opt_prime/accel/moe/layer.py

@@ -61,9 +61,17 @@ def __init__(self, num_experts: int, tiers: List, handles: List, placement: List
         self.owner_c, self.local_c = owner, local
         self._dev_tables = {}
 
+    def _padded_tables(self):
+        """owner/local tables with one extra entry for padding ids (-1 and num_experts, as in transformers' expert
+        parallelism): its owner is len(tiers), which is no tier, so padded pairs reach no tier and add nothing.
+        Routing ids must lie in [-1, num_experts]: other negative ids would wrap to a real expert, larger ones raise
+        (the native dispatcher skips anything outside [0, num_experts))."""
+        return (torch.cat([self.owner_c, torch.tensor([len(self.tiers)])]),
+                torch.cat([self.local_c, torch.zeros(1, dtype=torch.long)]))
+
     def _tables(self, device):
         if device not in self._dev_tables:
-            self._dev_tables[device] = (self.owner_c.to(device), self.local_c.to(device))
+            self._dev_tables[device] = tuple(t.to(device) for t in self._padded_tables())
         return self._dev_tables[device]
 
     def _snapshot(self, device):
@@ -73,11 +81,13 @@ def _snapshot(self, device):
             if self._has is None:
                 self._has = [bool((self.owner_c == t).any()) for t in range(len(self.tiers))]
                 self._blocking_any = any(self._has[t] and self.tiers[t].blocking for t in range(len(self.tiers)))
-                self._lists = (self.owner_c.tolist(), self.local_c.tolist(),      # for the decode fast path
-                               self.owner_c.to(torch.int32).contiguous(), self.local_c.to(torch.int32).contiguous())
+                # new tensors, never written in place: _adapt edits owner_c/local_c and then resets _has
+                self._host_tables = owner, local = self._padded_tables()
+                self._lists = (owner.tolist(), local.tolist(),                    # for the decode fast path
+                               owner.to(torch.int32).contiguous(), local.to(torch.int32).contiguous())
                 self._native = None
-            return (self.owner_c.clone(), self.local_c.clone(), self._tables(device), self._has, self._blocking_any,
-                    self._lists)
+            owner, local = self._host_tables
+            return owner, local, self._tables(device), self._has, self._blocking_any, self._lists
 
     def _pinned(self, name, like):
         key = (name, tuple(like.shape), like.dtype)
@@ -99,7 +109,7 @@ def _host_buf(self, name, shape, pinned):
 
     def flush_stats(self, st):
         if self._dev_counts is not None:
-            for t, n in enumerate(self._dev_counts.tolist()):
+            for t, n in enumerate(self._dev_counts.tolist()[:len(self.tiers)]):     # the last bin is padding
                 if n:
                     st.tier_pairs[self.names[t]] = st.tier_pairs.get(self.names[t], 0) + int(n)
                     if self.tiers[t].blocking: st.host_pairs += int(n)
@@ -110,8 +120,7 @@ def forward(self, hidden, top_k_index, top_k_weights):
         with self._execution_lock:
             guard = _NATIVE_LOCK if self.cfg.native_dispatch else nullcontext()
             with guard:
-                valid = (top_k_index >= 0) & (top_k_index < self.E)
-                return self._forward(hidden, top_k_index.clamp(0, self.E - 1), top_k_weights * valid)
+                return self._forward(hidden, top_k_index, top_k_weights)     # padding ids: see _padded_tables
 
     def _forward(self, hidden, top_k_index, top_k_weights):
         cfg, st = self.cfg, self.cfg.stats
@@ -126,7 +135,8 @@ def _forward(self, hidden, top_k_index, top_k_weights):
         if cfg.adapt_every:                                 # routing EMA on the device (no host sync)
             if self._counts_d is None or self._counts_d.device != dev:
                 self._counts_d = torch.zeros(self.E, dtype=torch.float32, device=dev)
-            self._counts_d.mul_(cfg.adapt_decay).add_(torch.bincount(top_k_index.reshape(-1), minlength=self.E).float())
+            ids = top_k_index.reshape(-1).remainder(self.E + 1)            # padding (-1, E) -> bin E, not counted
+            self._counts_d.mul_(cfg.adapt_decay).add_(torch.bincount(ids, minlength=self.E + 1)[:self.E].float())
             self.calls += 1
 
         # 1) start copying the routing (and the hidden rows) to the host without waiting
@@ -158,9 +168,9 @@ def _forward(self, hidden, top_k_index, top_k_weights):
                 futs, rows_c = self._split(idx_p, w_p, h_p, owner_c, local_c)
         else:
             if self._dev_counts is None or self._dev_counts.device != dev:
-                self._dev_counts = torch.zeros(len(self.tiers), dtype=torch.long, device=dev)
+                self._dev_counts = torch.zeros(len(self.tiers) + 1, dtype=torch.long, device=dev)
                 st._sources.append(self)
-            self._dev_counts += torch.bincount(own_idx.reshape(-1), minlength=len(self.tiers))
+            self._dev_counts += torch.bincount(own_idx.reshape(-1), minlength=len(self.tiers) + 1)
 
         if prof is not None: tp = _ph(prof, "split", tp)
         # 3) launch every device tier that owns experts of this layer (needs nothing from the host)
@@ -220,7 +230,7 @@ def _forward(self, hidden, top_k_index, top_k_weights):
     def _split(self, idx_c, w_p, h_p, owner_c, local_c):
         """general path: tensors per blocking tier, all in workers."""
         own_c = owner_c[idx_c]
-        per_tier = torch.bincount(own_c.reshape(-1), minlength=len(self.tiers)).tolist()
+        per_tier = torch.bincount(own_c.reshape(-1), minlength=len(self.tiers) + 1).tolist()[:len(self.tiers)]
         self._count(per_tier)
         blocking = [t for t, n in enumerate(per_tier) if n and self.tiers[t].blocking]
         if not blocking:
@@ -269,7 +279,7 @@ def _native_start(self, nat, idx_p, w_p, h_p, lists):
         counts = torch.zeros(T, dtype=torch.int64)
         R, k = idx.shape
         _NATIVE["lib"].moe_disp_start(_NATIVE["disp"], R, k, h_p.shape[1], idx.data_ptr(), w.data_ptr(), h_p.data_ptr(),
-                                      lists[2].data_ptr(), lists[3].data_ptr(), T, blocking, fns, ctxs,
+                                      lists[2].data_ptr(), lists[3].data_ptr(), self.E, T, blocking, fns, ctxs,
                                       counts.data_ptr())
         self._count(counts.tolist())
         return True
@@ -285,7 +295,10 @@ def _split_small(self, idx_p, w_p, h_p, lists):
         counts = [0] * T
         for r, (ir, wr) in enumerate(zip(il, wl)):
             for e, a in zip(ir, wr):
-                t = owner_l[e]; counts[t] += 1
+                t = owner_l[e]
+                if t == T:                                # padding id
+                    continue
+                counts[t] += 1
                 if a != 0.0 and self.tiers[t].blocking:
                     per[t].setdefault(r, []).append((local_l[e], a))
         self._count(counts)

opt_prime/opt_prime/accel/platforms/cpu/csrc/dispatch.c

@@ -102,7 +102,7 @@ static void reserve(job_t *j, int R, int k, int H) {
 // Split the [R, k] routed pairs by owner[expert]; post one job per blocking tier with pairs. counts[t] += pairs routed
 // to tier t (all tiers). Returns the number of jobs posted (they run while the caller launches device work).
 int moe_disp_start(void *dp, int R, int k, int H, const int64_t *idx, const float *w, const uint16_t *x,
-                   const int32_t *owner, const int32_t *local, int ntiers, const int8_t *blocking, void *const *fns,
+                   const int32_t *owner, const int32_t *local, int nexp, int ntiers, const int8_t *blocking, void *const *fns,
                    void *const *ctxs, int64_t *counts) {
     disp_t *d = (disp_t *)dp;
     int njobs = 0, active[MAX_TIERS] = {0};
@@ -112,7 +112,9 @@ int moe_disp_start(void *dp, int R, int k, int H, const int64_t *idx, const floa
     for (int r = 0; r < R; r++) {
         int cnt[MAX_TIERS] = {0};
         for (int j = 0; j < k; j++) {
-            int t = owner[idx[r * k + j]];
+            int64_t e = idx[r * k + j];
+            if (e < 0 || e >= nexp) continue;           // padding ids (-1, num_experts) and anything out of range
+            int t = owner[e];
             counts[t]++;
             if (blocking[t] && w[r * k + j] != 0.0f) cnt[t]++;
         }
@@ -144,7 +146,7 @@ int moe_disp_start(void *dp, int R, int k, int H, const int64_t *idx, const floa
             float *wr = jb->w + (size_t)jb->R * jb->k;
             for (int j = 0; j < k; j++) {
                 int64_t e = idx[r * k + j];
-                if (owner[e] == t && w[r * k + j] != 0.0f) { ir[c] = local[e]; wr[c] = w[r * k + j]; c++; }
+                if (e >= 0 && e < nexp && owner[e] == t && w[r * k + j] != 0.0f) { ir[c] = local[e]; wr[c] = w[r * k + j]; c++; }
             }
             if (!c) continue;
             for (int j = c; j < jb->k; j++) { ir[j] = ir[0]; wr[j] = 0.0f; }

opt_prime/opt_prime/accel/platforms/cpu/native.py

@@ -47,7 +47,7 @@ def load_dispatch():
             lib = ctypes.CDLL(_build(os.path.join(_CSRC, "dispatch.c"), "moe_dispatch", ("-lpthread",)))
             P, I = ctypes.c_void_p, ctypes.c_int
             lib.moe_disp_new.restype = P
-            lib.moe_disp_start.argtypes = [P, I, I, I, P, P, P, P, P, I, P, P, P, P]
+            lib.moe_disp_start.argtypes = [P, I, I, I, P, P, P, P, P, I, I, P, P, P, P]
             lib.moe_disp_wait.argtypes = [P, I, I, P]
             lib.moe_cpu_ctx_new.restype = P
             lib.moe_cpu_ctx_new.argtypes = [P, P, P, I, I, I]

opt_prime/tests/conftest.py

@@ -0,0 +1,37 @@
+"""Shared pytest configuration.
+
+opt_prime.accel targets transformers 5.x. On an older transformers, tests that need a model family or config
+attribute it does not have are skipped instead of failing.
+"""
+import pytest
+
+try:
+    import transformers
+    _TF_MAJOR = int(transformers.__version__.split(".")[0])
+except Exception:                                   # transformers missing or unparsable version
+    _TF_MAJOR = None
+
+_TF4_GAPS = ("Unrecognized model identifier", "object has no attribute 'rope_parameters'",
+             "does not recognize this architecture")
+
+
+@pytest.hookimpl(wrapper=True)
+def pytest_runtest_setup(item):
+    try:
+        return (yield)
+    except (ValueError, AttributeError, KeyError) as exc:
+        _reraise_as_skip(exc)
+
+
+@pytest.hookimpl(wrapper=True)
+def pytest_runtest_call(item):
+    try:
+        return (yield)
+    except (ValueError, AttributeError, KeyError) as exc:
+        _reraise_as_skip(exc)
+
+
+def _reraise_as_skip(exc):
+    if _TF_MAJOR is not None and _TF_MAJOR < 5 and any(g in str(exc) for g in _TF4_GAPS):
+        pytest.skip(f"needs transformers>=5 (installed {transformers.__version__}): {str(exc)[:80]}")
+    raise exc

opt_prime/tests/test_moe_offload.py

@@ -127,6 +127,35 @@ def test_native_host_kernel_matches_torch(rows, wdtype):
     assert (a != b).float().mean().item() < 0.05
 
 
+@pytest.mark.parametrize("pad", [-1, NUM_EXPERTS, -5])
+def test_native_dispatch_drops_padding_ids(pad):
+    """#89: the C dispatcher must skip padding ids (-1 and num_experts) instead of indexing the owner table, and
+    skip other negative ids rather than read the table out of bounds"""
+    import opt_prime.accel.platforms.cpu.native as native
+    from opt_prime.accel.moe.backends import get_backend
+    from opt_prime.accel.moe.config import MoEOffloadConfig
+    from opt_prime.accel.moe.layer import OffloadedLayer
+    from opt_prime.accel.spec.adapters import adapter_for
+    if native.load() is None or native.load_dispatch() is None:
+        pytest.skip("no AVX-512 BF16 or no gcc")
+    m = build_tiny("qwen3_moe", torch.bfloat16)
+    ex = find_moe_experts(m)[0].module
+    ad = adapter_for(ex)
+    tier = get_backend("cpu", kernel="native")
+    handle = tier.load(ad.weights(ex, list(range(NUM_EXPERTS))))
+    layer = OffloadedLayer(NUM_EXPERTS, [tier], [handle], [list(range(NUM_EXPERTS))], ["host"],
+                           MoEOffloadConfig(native_dispatch=True), ad, ex)
+    h = torch.randn(3, ex.gate_up_proj.shape[2], dtype=torch.bfloat16)
+    idx = torch.tensor([[1, pad], [pad, pad], [2, 0]])
+    w = torch.rand(3, 2).to(torch.bfloat16)
+    assert layer._native_plan(layer._snapshot(h.device)[3], True) is not None
+    got = layer.forward(h, idx, w)
+    valid = (idx >= 0) & (idx < NUM_EXPERTS)
+    ref = tier.run(handle, h, idx.clamp(0, NUM_EXPERTS - 1), w * valid)
+    assert got[1].abs().max().item() == 0
+    assert (got.float() - ref.float()).abs().max().item() <= 2e-2 * ref.float().abs().max().item() + 1e-6
+
+
 # ---------- heterogeneous tiers, planner, backends ----------
 
 @pytest.mark.parametrize("fast_rows,kernel", [(0, "torch"), (64, "torch"), (64, "auto")])

opt_prime/tests/test_moe_review.py

@@ -229,7 +229,7 @@ def test_native_start_drains_jobs_of_a_call_that_was_not_waited_for():
     blocking, fns, ctxs = (ctypes.c_int8 * 1)(1), (ctypes.c_void_p * 1)(tier.native_tier(None)[0]), (ctypes.c_void_p * 1)()
     for _ in range(2):                          # the first call raised before moe_disp_wait
         lib.moe_disp_start(d, 2, 1, 4, idx.data_ptr(), w.data_ptr(), x.data_ptr(), owner.data_ptr(), local.data_ptr(),
-                           1, blocking, fns, ctxs, counts.data_ptr())
+                           1, 1, blocking, fns, ctxs, counts.data_ptr())   # num_experts, ntiers
     assert lib.moe_disp_wait(d, 2, 4, acc.data_ptr()) == 0
     assert tier.calls == 2
     torch.testing.assert_close(acc, torch.tensor([[1.] * 4, [.5] * 4]))
@@ -258,3 +258,38 @@ def test_offloaded_padding_routes_match_forward_and_backward():
     actual = layer.backward(x, ids, w, torch.ones_like(ref))
     for a, b in zip(actual, expected):
         torch.testing.assert_close(a, b)
+
+
+def _host_layer(native=False, adapt_every=0):
+    m, adapter = Experts(), GenericAdapter()
+    tier = CpuBackend(kernel='torch', threads=1)
+    handle = tier.load(adapter.weights(m, [0, 1]))
+    cfg = MoEOffloadConfig(native_dispatch=native, adapt_every=adapt_every)
+    return m, OffloadedLayer(2, [tier], [handle], [[0, 1]], ['host'], cfg, adapter, m)
+
+
+@pytest.mark.parametrize('pad', [-1, 2])
+def test_padding_ids_reach_no_tier_and_are_not_counted(pad):
+    # #89: padding is dropped through the routing tables (no per-call mask); it must not count as expert 0
+    m, layer = _host_layer(adapt_every=1000)
+    x, w = torch.ones(3, 2), torch.tensor([[0.4, 0.7], [0.5, 0.5], [1.0, 0.3]])
+    ids = torch.tensor([[0, pad], [pad, pad], [1, pad]])
+    ref = opt_batched_mm(m, x, ids.clamp(0, 1), w * ((ids >= 0) & (ids < 2)))
+    torch.testing.assert_close(layer.forward(x, ids, w), ref)
+    assert layer._counts_d.tolist() == [1.0, 1.0]
+
+
+def test_padding_with_large_rows_uses_general_split():
+    m, layer = _host_layer()
+    layer.cfg.fast_rows = 0                                  # force the general (tensor) host path
+    x = torch.ones(4, 2)
+    ids, w = torch.tensor([[0, -1], [2, 1], [-1, -1], [1, 0]]), torch.full((4, 2), 0.5)
+    ref = opt_batched_mm(m, x, ids.clamp(0, 1), w * ((ids >= 0) & (ids < 2)))
+    torch.testing.assert_close(layer.forward(x, ids, w), ref)
+
+
+def test_grouped_mm_all_padding_with_inf_hidden_is_zero_not_nan():
+    m = Experts()
+    x = torch.full((1, 2), float('inf'), requires_grad=True)
+    out = opt_grouped_mm(m, x, torch.tensor([[2]]), torch.ones(1, 1))
+    assert torch.equal(out, torch.zeros_like(out))