1"""Distributed expert execution: a placement maps (layer, expert) -> worker name (or None = local). For each layer the2execution block splits the (row, k) pairs between local slots and remote workers, sends every worker its batch BEFORE waiting3for any reply (so a layer waits for its slowest worker, not for each in turn), then assembles y[row, k, :] in the original4order so the weighted sum is bit-identical to the single-node path. Workers that time out or fail raise; the caller may5retry locally (fallback) — the expert exists on local NVMe too.6"""7from __future__ import annotations8import time9from typing import Dict, List, Optional, Tuple10import numpy as np11import mlx.core as mx12from mlx_lm.models.activations import swiglu13from transport.client import WorkerClient141516class RemoteExecutor:17 def __init__(self, workers: Dict[str, WorkerClient], placement: Dict[Tuple[int, int], Optional[str]], fallback_local: bool = True):18 self.workers = workers19 self.placement = placement20 self.fallback_local = fallback_local21 self.remote_pairs = 0; self.local_pairs = 0; self.layers_remote = 0; self.remote_wait = 0.0; self.failures = 02223 def where(self, layer: int, expert: int) -> Optional[str]:24 return self.placement.get((layer, expert))2526 def split(self, layer: int, flat: List[int], k: int):27 """flat: expert ids in (row, k) order. Returns local pairs and per-worker pairs as (position, row, expert)."""28 local, remote = [], {}29 for pos, e in enumerate(flat):30 w = self.where(layer, e)31 if w is None or w not in self.workers:32 local.append((pos, pos // k, e))33 else:34 remote.setdefault(w, []).append((pos, pos // k, e))35 return local, remote3637 def run_layer(self, block, x_flat: mx.array, flat: List[int], shape, probs_np=None):38 """Execute one layer's (row, k) pairs locally and remotely; returns y of shape (*shape, D)."""39 k = shape[-1]40 local, remote = self.split(block.table.layer, flat, k)41 R, D = x_flat.shape42 outs = [None] * len(flat)43 # 1. dispatch every remote batch first (activations as float16 rows)44 if remote:45 x_np = np.array(x_flat.astype(mx.float16))46 for w, pairs in remote.items():47 self.workers[w].exec_async(block.table.layer, x_np, [(r, e) for _, r, e in pairs])48 self.layers_remote += 149 # 2. local experts through the slot table (same code path as single node)50 if local:51 lflat = [e for _, _, e in local]52 rows = [r for _, r, _ in local]53 uniq = list(dict.fromkeys(lflat))54 if block.prefetcher is not None:55 block.prefetcher.before_layer(block.table.layer, uniq); block.prefetcher.after_layer(block.table.layer, uniq, probs_np)56 slot_of = dict(zip(uniq, block.table.ensure(uniq)))57 sidx = mx.array([slot_of[e] for e in lflat], dtype=mx.uint32).reshape(-1, 1)58 xs = mx.expand_dims(x_flat[mx.array(rows, dtype=mx.uint32)], (-2, -3))59 up, gate = block._upgate(xs, sidx) if block.table.fused else (block._qmm(xs, "up", sidx, False), block._qmm(xs, "gate", sidx, False))60 yl = block._qmm(swiglu(gate, up), "down", sidx, False).squeeze(-2).squeeze(-2) # [P, D]61 for i, (pos, _, _) in enumerate(local):62 outs[pos] = yl[i]63 self.local_pairs += len(local)64 # 3. collect remote replies65 if remote:66 t = time.perf_counter()67 for w, pairs in remote.items():68 try:69 y, _ = self.workers[w].exec_wait()70 except Exception as exc:71 self.failures += 172 if not self.fallback_local:73 raise74 # local fallback: execute the failed pairs locally75 lflat = [e for _, _, e in pairs]; rows = [r for _, r, _ in pairs]76 uniq = list(dict.fromkeys(lflat)); slot_of = dict(zip(uniq, block.table.ensure(uniq)))77 sidx = mx.array([slot_of[e] for e in lflat], dtype=mx.uint32).reshape(-1, 1)78 xs = mx.expand_dims(x_flat[mx.array(rows, dtype=mx.uint32)], (-2, -3))79 up, gate = block._upgate(xs, sidx) if block.table.fused else (block._qmm(xs, "up", sidx, False), block._qmm(xs, "gate", sidx, False))80 yf = block._qmm(swiglu(gate, up), "down", sidx, False).squeeze(-2).squeeze(-2)81 for i, (pos, _, _) in enumerate(pairs):82 outs[pos] = yf[i]83 continue84 ym = mx.array(y)85 for i, (pos, _, _) in enumerate(pairs):86 outs[pos] = ym[i]87 self.remote_pairs += len(pairs)88 self.remote_wait += time.perf_counter() - t89 y = mx.stack(outs, 0).reshape(tuple(shape) + (D,))90 return y9192 def stats(self):93 return {"remote_pairs": self.remote_pairs, "local_pairs": self.local_pairs, "layers_remote": self.layers_remote,94 "remote_wait_s": self.remote_wait, "failures": self.failures, "workers": {n: w.stats() for n, w in self.workers.items()}}95