From 90586f275c4587c317eba430e0f44055d675010f Mon Sep 17 00:00:00 2001 From: gspark Date: Wed, 21 Oct 2026 11:19:07 +0900 Subject: [PATCH 1/3] accel.moe: drop padding ids through the routing tables instead of per-call masks (#89) The owner/local tables get one extra entry whose owner is no tier, so padding ids (-1 and num_experts, as in transformers' expert parallelism) reach no tier: no per-call compare/clamp/mul kernels on the decode path. The routing EMA and tier counts no longer count padding as expert 0; the native dispatcher skips padding ids instead of indexing the owner table with -1. opt_grouped_mm starts from zeros (an inf in hidden_states no longer becomes NaN). tests/conftest.py skips tests that need a transformers 5 model family on transformers 4. Co-Authored-By: Claude Opus 5.5 --- opt_prime/opt_prime/accel/moe/batched.py | 5 ++- opt_prime/opt_prime/accel/moe/layer.py | 36 +++++++++++------- .../accel/platforms/cpu/csrc/dispatch.c | 6 ++- opt_prime/tests/conftest.py | 37 +++++++++++++++++++ opt_prime/tests/test_moe_offload.py | 28 ++++++++++++++ opt_prime/tests/test_moe_review.py | 35 ++++++++++++++++++ 6 files changed, 131 insertions(+), 16 deletions(-) create mode 100644 opt_prime/tests/conftest.py diff --git a/opt_prime/opt_prime/accel/moe/batched.py b/opt_prime/opt_prime/accel/moe/batched.py index a9ba3952..3e314946 100644 --- a/opt_prime/opt_prime/accel/moe/batched.py +++ b/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 diff --git a/opt_prime/opt_prime/accel/moe/layer.py b/opt_prime/opt_prime/accel/moe/layer.py index 4e1c5e39..8a41c0f5 100644 --- a/opt_prime/opt_prime/accel/moe/layer.py +++ b/opt_prime/opt_prime/accel/moe/layer.py @@ -61,9 +61,15 @@ 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.""" + 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 +79,12 @@ 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()) + 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._padded_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 +106,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 +117,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 +132,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 +165,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 +227,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: @@ -285,7 +292,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) diff --git a/opt_prime/opt_prime/accel/platforms/cpu/csrc/dispatch.c b/opt_prime/opt_prime/accel/platforms/cpu/csrc/dispatch.c index c0af6851..3eba61f8 100644 --- a/opt_prime/opt_prime/accel/platforms/cpu/csrc/dispatch.c +++ b/opt_prime/opt_prime/accel/platforms/cpu/csrc/dispatch.c @@ -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]; + int t = e < 0 ? ntiers : owner[e]; // owner has num_experts + 1 entries; padding -> ntiers + if (t >= ntiers) continue; 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 && 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; } diff --git a/opt_prime/tests/conftest.py b/opt_prime/tests/conftest.py new file mode 100644 index 00000000..53ae4d6a --- /dev/null +++ b/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 diff --git a/opt_prime/tests/test_moe_offload.py b/opt_prime/tests/test_moe_offload.py index eccb2700..df5d40a3 100644 --- a/opt_prime/tests/test_moe_offload.py +++ b/opt_prime/tests/test_moe_offload.py @@ -127,6 +127,34 @@ def test_native_host_kernel_matches_torch(rows, wdtype): assert (a != b).float().mean().item() < 0.05 +@pytest.mark.parametrize("pad", [-1, NUM_EXPERTS]) +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""" + 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")]) diff --git a/opt_prime/tests/test_moe_review.py b/opt_prime/tests/test_moe_review.py index ab0b44ab..76323501 100644 --- a/opt_prime/tests/test_moe_review.py +++ b/opt_prime/tests/test_moe_review.py @@ -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)) From 5877606ce282d0b383d4f3a6a98c7fc34a0e8f86 Mon Sep 17 00:00:00 2001 From: gspark Date: Fri, 23 Oct 2026 09:20:09 +0900 Subject: [PATCH 2/3] accel.moe: native dispatcher bounds expert ids; cache the padded host tables dispatch.c now takes num_experts and skips any id outside [0, num_experts) instead of reading owner[] out of bounds for ids above num_experts. _snapshot builds the padded owner/local tables once per placement instead of two torch.cat per layer call on the decode path. The routing-id contract ([-1, num_experts]) is documented on _padded_tables. Co-Authored-By: Claude Opus 5.5 --- opt_prime/opt_prime/accel/moe/layer.py | 11 +++++++---- .../opt_prime/accel/platforms/cpu/csrc/dispatch.c | 8 ++++---- opt_prime/opt_prime/accel/platforms/cpu/native.py | 2 +- opt_prime/tests/test_moe_offload.py | 5 +++-- 4 files changed, 15 insertions(+), 11 deletions(-) diff --git a/opt_prime/opt_prime/accel/moe/layer.py b/opt_prime/opt_prime/accel/moe/layer.py index 8a41c0f5..6681b503 100644 --- a/opt_prime/opt_prime/accel/moe/layer.py +++ b/opt_prime/opt_prime/accel/moe/layer.py @@ -63,7 +63,9 @@ def __init__(self, num_experts: int, tiers: List, handles: List, placement: List 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.""" + 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)])) @@ -79,11 +81,12 @@ 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))) - owner, local = self._padded_tables() + # 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 - owner, local = self._padded_tables() + owner, local = self._host_tables return owner, local, self._tables(device), self._has, self._blocking_any, self._lists def _pinned(self, name, like): @@ -276,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 diff --git a/opt_prime/opt_prime/accel/platforms/cpu/csrc/dispatch.c b/opt_prime/opt_prime/accel/platforms/cpu/csrc/dispatch.c index 3eba61f8..e223b5c0 100644 --- a/opt_prime/opt_prime/accel/platforms/cpu/csrc/dispatch.c +++ b/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}; @@ -113,8 +113,8 @@ int moe_disp_start(void *dp, int R, int k, int H, const int64_t *idx, const floa int cnt[MAX_TIERS] = {0}; for (int j = 0; j < k; j++) { int64_t e = idx[r * k + j]; - int t = e < 0 ? ntiers : owner[e]; // owner has num_experts + 1 entries; padding -> ntiers - if (t >= ntiers) continue; + 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]++; } @@ -146,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 (e >= 0 && 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; } diff --git a/opt_prime/opt_prime/accel/platforms/cpu/native.py b/opt_prime/opt_prime/accel/platforms/cpu/native.py index 69bfd7f9..e74259fb 100644 --- a/opt_prime/opt_prime/accel/platforms/cpu/native.py +++ b/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] diff --git a/opt_prime/tests/test_moe_offload.py b/opt_prime/tests/test_moe_offload.py index df5d40a3..5d0ff9f3 100644 --- a/opt_prime/tests/test_moe_offload.py +++ b/opt_prime/tests/test_moe_offload.py @@ -127,9 +127,10 @@ def test_native_host_kernel_matches_torch(rows, wdtype): assert (a != b).float().mean().item() < 0.05 -@pytest.mark.parametrize("pad", [-1, NUM_EXPERTS]) +@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""" + """#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 8d00d484574ba7268f87b6050b8e3a16e7c96c4e Mon Sep 17 00:00:00 2001 From: gspark Date: Fri, 30 Oct 2026 09:36:34 +0900 Subject: [PATCH 3/3] accel.moe: pass num_experts in the direct dispatcher test (#89) Co-Authored-By: Claude Opus 5.5 --- opt_prime/tests/test_moe_review.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/opt_prime/tests/test_moe_review.py b/opt_prime/tests/test_moe_review.py index 76323501..23408712 100644 --- a/opt_prime/tests/test_moe_review.py +++ b/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]))