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 dt = x_flat.dtype46 x_np = np.array(x_flat.view(mx.uint16)) if dt == mx.bfloat16 else np.array(x_flat)47 for w, pairs in remote.items():48 self.workers[w].exec_async(block.table.layer, x_np, [(r, e) for _, r, e in pairs], dtype=str(dt).split(".")[-1])49 self.layers_remote += 150 # 2. local experts through the slot table (same code path as single node)51 if local:52 lflat = [e for _, _, e in local]53 rows = [r for _, r, _ in local]54 uniq = list(dict.fromkeys(lflat))55 if block.prefetcher is not None:56 block.prefetcher.before_layer(block.table.layer, uniq); block.prefetcher.after_layer(block.table.layer, uniq, probs_np)57 slot_of = dict(zip(uniq, block.table.ensure(uniq)))58 sidx = mx.array([slot_of[e] for e in lflat], dtype=mx.uint32).reshape(-1, 1)59 xs = mx.expand_dims(x_flat[mx.array(rows, dtype=mx.uint32)], (-2, -3))60 up, gate = block._upgate(xs, sidx) if block.table.fused else (block._qmm(xs, "up", sidx, False), block._qmm(xs, "gate", sidx, False))61 yl = block._qmm(swiglu(gate, up), "down", sidx, False).squeeze(-2).squeeze(-2) # [P, D]62 for i, (pos, _, _) in enumerate(local):63 outs[pos] = yl[i]64 self.local_pairs += len(local)65 # 3. collect remote replies66 if remote:67 t = time.perf_counter()68 for w, pairs in remote.items():69 try:70 y, _ = self.workers[w].exec_wait()71 except Exception as exc:72 self.failures += 173 if not self.fallback_local:74 raise75 # local fallback: execute the failed pairs locally76 lflat = [e for _, _, e in pairs]; rows = [r for _, r, _ in pairs]77 uniq = list(dict.fromkeys(lflat)); slot_of = dict(zip(uniq, block.table.ensure(uniq)))78 sidx = mx.array([slot_of[e] for e in lflat], dtype=mx.uint32).reshape(-1, 1)79 xs = mx.expand_dims(x_flat[mx.array(rows, dtype=mx.uint32)], (-2, -3))80 up, gate = block._upgate(xs, sidx) if block.table.fused else (block._qmm(xs, "up", sidx, False), block._qmm(xs, "gate", sidx, False))81 yf = block._qmm(swiglu(gate, up), "down", sidx, False).squeeze(-2).squeeze(-2)82 for i, (pos, _, _) in enumerate(pairs):83 outs[pos] = yf[i]84 continue85 ym = mx.array(y).view(mx.bfloat16) if x_flat.dtype == mx.bfloat16 else mx.array(y)86 for i, (pos, _, _) in enumerate(pairs):87 outs[pos] = ym[i]88 self.remote_pairs += len(pairs)89 self.remote_wait += time.perf_counter() - t90 y = mx.stack(outs, 0).reshape(tuple(shape) + (D,))91 return y9293 def stats(self):94 return {"remote_pairs": self.remote_pairs, "local_pairs": self.local_pairs, "layers_remote": self.layers_remote,95 "remote_wait_s": self.remote_wait, "failures": self.failures, "workers": {n: w.stats() for n, w in self.workers.items()}}96