gt50 · Oct 4, 2026 · Oct 4, 2026
diff --git a/configs/strix_halo_large.yaml b/configs/strix_halo_large.yaml
@@ -96,8 +96,16 @@ pool:
   ram_cache: 0            # 0 = every expert stays in the tier, which here
                           # is device memory: disk is read once per expert
 
-  capacity_factor: 1.5    # most of a batch one expert may take, as a multiple
+  capacity_factor: 0      # most of a batch one expert may take, as a multiple
                           # of its fair share; the rest is dropped. 0 removes
+                          # the cap. The cap exists to bound memory on an 8 GB
+                          # card; with 128 experts on the card a fresh model's
+                          # router piles onto a few, and at 1.5 it dropped 59%
+                          # of routing in tools/bench_hw.py (76% in a real
+                          # run's first minutes) - most of what the experts
+                          # could learn. Uncapped, a step was 27% slower and
+                          # peaked at 22.8 GB instead of 17.8. Measured on the
+                          # 8060S: 3 dropped 41%, 6 dropped 25%
 
   # THE BALANCE TERM - how the whole pool gets used. Left alone the router
   # keeps choosing the same experts: reading held-out on its own it put 59 of

diff --git a/minagi/paged.py b/minagi/paged.py
@@ -287,6 +287,10 @@ def __init__(self, path, d_model, d_ff, n_experts, resident=16,
         self.pressure = 0.0
         self.want_k = 0.0
         self.swaps = 0
+        # slot order as a device tensor, rebuilt only when the slots change -
+        # see _slot_index
+        self._slot_key = None
+        self._slot_t = None
         self._opt = None
         self._sites = []          # the call sites routing into this pool
         # THE PRUNE CLOCK: the text count at which each expert was last
@@ -442,24 +446,39 @@ def router_rows(self):
         """One row per expert on disk; growth adds rows."""
         return self._n
 
+    def _slot_index(self, device):
+        """
+        The slots' experts as an index tensor on `device`, in slot order.
+
+        Every routing call asks for it several times, once per row of every
+        character, and the slots change only when an expert is loaded. It was
+        rebuilt from the Python list each time: a host-to-device copy per
+        call, which is most of what writing a character cost on an APU. Now
+        it is rebuilt when the list differs from the one it was built from.
+        Callers only read it.
+        """
+        key = tuple(self.slots)
+        if (self._slot_key != key or self._slot_t is None
+                or self._slot_t.device != torch.device(device)):
+            self._slot_t = torch.tensor([max(s, 0) for s in key],
+                                        device=device)
+            self._slot_key = key
+        return self._slot_t
+
     def resident_rows(self):
         """Which router rows the resident experts own, in slot order."""
-        return torch.tensor([max(s, 0) for s in self.slots],
-                            device=self.gate.d
8000
evice)
+        return self._slot_index(self.gate.device)
 
     def n_resident(self):
         return self.resident
 
     def routable_gate(self):
         """Gates of the resident experts, in slot order."""
-        idx = torch.tensor([max(s, 0) for s in self.slots],
-                           device=self.gate.device)
-        return self.gate[idx]
+        return self.gate[self._slot_index(self.gate.device)]
 
     def note_use(self, hit):
         """Routing counts arrive per SLOT; usage is kept per EXPERT."""
-        idx = torch.tensor([max(s, 0) for s in self.slots],
-                           device=self.use.device)
+        idx = self._slot_index(self.use.device)
         self.use.index_add_(0, idx, hit.to(self.use.dtype))
         self.age += 1
 
@@ -609,7 +628,11 @@ def admit(self, mass, draw=8):
         if free <= 0:
             return 0
         order = torch.argsort(m, descending=True).tolist()
-        new = [e for e in order if float(m[e]) > 0 and e not in adm][:free]
+        # as a list: indexing the tensor per expert was two torch ops each,
+        # for every expert in the pool at every row of every character -
+        # nearly half the operations writing a character took
+        ml = m.tolist()
+        new = [e for e in order if ml[e] > 0 and e not in adm][:free]
         if not new:
             return 0
         return self._place(new)
@@ -652,8 +675,10 @@ def _place(self, new):
         # the least recently used, which is the one least likely to be wanted
         # back
         victims = [s for s, e in enumerate(self.slots) if e < 0 or e not in adm]
+        # read once: indexing the device tensor per slot was a sync per slot
+        seen = self.last_seen.tolist()
         victims.sort(key=lambda s: (self.slots[s] >= 0,
-                                    float(self.last_seen[self.slots[s]])
+                                    seen[self.slots[s]]
                                     if self.slots[s] >= 0 else -1.0))
         plan = list(self.slots)
         for e in new:
@@ -666,8 +691,8 @@ def _place(self, new):
                     self.admits[e] += 1
             self._rearrange(plan, here)
             self.swaps += 1
-        for e in new:
-            self.ever[e] = True
+        # one indexed write, not one device write per expert
+        self.ever[torch.tensor(new, device=self.ever.device)] = True
         # THE PRUNE CLOCK: whatever admits an expert resets it
         self.last_seen[torch.tensor(new, device=self.last_seen.device)] = \
             float(self.segments)

diff --git a/minagi/pool.py b/minagi/pool.py
@@ -335,6 +335,10 @@ def __init__(self, pool, d_model, top_k=4, site=0, z_weight=1e-3,
         # ...and the denominator, so the counter reads as a share of the
         # work rather than as a number nobody can scale.
         self.routed = 0
+        # Run only the experts a call actually chose when they are at most
+        # half the card - see the dispatch. Off reproduces the old path, which
+        # the tests compare against.
+        self.subset_dispatch = True
 
     @property
     def pool(self):
@@ -444,7 +448,9 @@ def requested(scores):
         # measured BEFORE normalisation - afterwards it sums to 1 by
         # construction and carries no information
         with torch.no_grad():
-            kept = float(w.sum(-1).mean())
+            # kept on the device: these feed telemetry only, and reading them
+            # back here was a sync per row of every character
+            kept = w.sum(-1).mean()
             # How many experts the router actually wants: the smallest number
             # covering 90% of its probability mass. If that exceeds k, the
             # router is being forced to discard experts it would have used, and
@@ -494,8 +500,8 @@ def requested(scores):
             # wants to spread across more experts than k, the pool is too
             # small to express what it is trying to do - that is capacity
             # pressure, and it is visible without waiting for a plateau.
-            p.pressure = 0.9 * float(p.pressure) + 0.1 * (1.0 - kept)
-            p.want_k = 0.9 * float(p.want_k) + 0.1 * float(want)
+            p.pressure = 0.9 * p.pressure + 0.1 * (1.0 - kept)
+            p.want_k = 0.9 * p.want_k + 0.1 * want
         frac = F.one_hot(idx[:, 0], n).float().mean(0)
         self.aux = ((frac * probs.mean(0)).sum() * n
                     + self.z_weight * torch.logsumexp(logits, -1).pow(2).mean())
@@ -511,7 +517,13 @@ def requested(scores):
         order = torch.argsort(flat_e)
         e_sorted, t_sorted, w_sorted = flat_e[order], tok[order], flat_w[order]
         counts = torch.bincount(e_sorted, minlength=n)
-        cap = int(counts.max().item()) if n else 0
+        # One character's top_k are distinct experts, so each holds at most
+        # one of its assignments: no need to ask the device, which is a sync
+        # per row when a character is being written.
+        if N == 1:
+            cap = 1 if n else 0
+        else:
+            cap = int(counts.max().item()) if n else 0
         if cap == 0:
             return torch.zeros_like(x)
    
@@ -563,6 +575,37 @@ def requested(scores):
                 w_sorted, slot = w_sorted[keep], slot[keep]
                 cap = limit
 
+        # ONLY THE EXPERTS THAT WERE CHOSEN RUN. The buffer is [n, cap, D] and
+        # every matmul below reads every expert's weights - and casts them to
+        # the compute dtype - whether or not anything was routed to it. A
+        # training chunk of 2,048 characters reaches nearly every expert on
+        # the card, so that costs nothing there. A character being written is
+        # a forward of its own and reaches top_k of them: with 128 on the
+        # card, every row read and cast all 128 to run 8, and writing was
+        # bound by memory bandwidth - 2.9 characters a second on Strix Halo.
+        #
+        # So when at most half the experts are in use, the weights of just
+        # those are gathered first. Half, because the gather copies them in
+        # fp32 before the cast; below that the copy is no larger than the
+        # cast of the whole card it replaces, so peak memory never rises.
+        sub = None
+        if self.subset_dispatch and n > 1:
+            if N == 1:
+                # sorted and distinct already, and the count is known here
+                present = e_sorted
+            else:
+                present = torch.nonzero(counts > 0).squeeze(1)
+            m = int(present.numel())
+            if 0 < m * 2 <= n:
+                remap = torch.full((n,), -1, dtype=torch.long,
+                                   device=flat.device)
+                remap[present] = torch.arange(m, device=flat.device)
+                e_sorted = remap[e_sorted]
+                sub, n = present, m
+
+        def pick(W):
+            return W if sub is None else W.index_select(0, sub)
+
         def run(src, *ws):
             lv = [(ws[i], ws[i + 1], ws[i + 2]) for i in range(0, len(ws), 3)]
             # The slot buffer, and everything computed from it, in the COMPUTE
@@ -577,15 +620,15 @@ def run(src, *ws):
             buf = torch.zeros(n, cap, D, device=src.device, dtype=dt)
             buf[e_sorted, slot] = src[t_sorted].to(dt)
             if len(lv) == 1:
-                W1, W3, W2 = lv[0]
+                W1, W3, W2 = (pick(t) for t in lv[0])
                 h = F.silu(torch.bmm(buf, W1.transpose(1, 2).to(buf.dtype))) * \
                     torch.bmm(buf, W3.transpose(1, 2).to(buf.dtype))
                 y = torch.bmm(h, W2.transpose(1, 2).to(buf.dtype))  # [n,cap,D]
                 return y[e_sorted, slot].to(src.dtype)
             # deeper experts stack residually, so an expert of any depth can
             # still be added to the mixture without changing its scale
             x = buf
-            for W1, W3, W2 in lv:
+            for W1, W3, W2 in ((pick(a), pick(b), pick(c)) for a, b, c in lv):
                 h = F.silu(torch.bmm(x, W1.transpose(1, 2).to(x.dtype))) * \
                     torch.bmm(x, W3.transpose(1, 2).to(x.dtype))
                 x = x + torch.bmm(h, W2.transpose(1, 2).to(x.dtype))

diff --git a/tests/test_dispatch.py b/tests/test_dispatch.py
@@ -0,0 +1,48 @@
+"""
+Running only the chosen experts must compute what running all of them did:
+the same outputs, and the same gradients on every slot.
+"""
+
+import torch
+
+
+def _sites(model):
+    from minagi.pool import PooledMLP
+    return [m for m in model.modules() if isinstance(m, PooledMLP)]
+
+
+def _run(model, x, subset):
+    for s in _sites(model):
+        s.subset_dispatch = subset
+    model.zero_grad(set_to_none=True)
+    logits, loss = model(x[:, :-1], targets=x[:, 1:])
+    loss.backward()
+    grads = [p.grad.detach().clone() if p.grad is not None else None
+             for p in (model.pool.w1, model.pool.w3, model.pool.w2)]
+    return logits.detach(), grads
+
+
+def test_subset_dispatch_matches_full(tiny):
+    from minagi.build import build_paged
+    from minagi.create import create
+
+    create(tiny["weights"], verbose=False)
+    # 6 slots and top_k 2: a single character reaches 2 of them, under half,
+    # which is when the subset path runs
+    model, _, pool, _ = build_paged(tiny["weights"], "cpu", resident=6)
+    model.train()
+    for s in _sites(model):
+        s.grad_checkpoint = False
+
+    g = torch.Generator().manual_seed(3)
+    for n in (2, 3):                    # one and two characters per forward
+        x = torch.randint(32, 127, (1, n), generator=g)
+        torch.manual_seed(0)
+        a, ga = _run(model, x, subset=True)
+        torch.manual_seed(0)
+        b, gb = _run(model, x, subset=False)
+        assert torch.allclose(a, b, atol=1e-5), (a - b).abs().max()
+        for u, v in zip(ga, gb):
+            assert (u is None) == (v is None)
+            if u is not None:
+                assert torch.allclose(u, v, atol=1e-6), (u - v).abs().max()
diff --git a/tools/bench_hw.py b/tools/bench_hw.py
@@ -181,17 +181,66 @@ def step():
 
     if dev.type == "cuda":
         torch.cuda.reset_peak_memory_stats()
+    from minagi.pool import PooledMLP
+    sites = [m for m in model.modules() if isinstance(m, PooledMLP)]
+    for s in sites:
+        s.dropped = s.routed = 0
     ms = timed(step, dev, reps, warm=1)
+    routed = sum(s.routed for s in sites)
+    dropped = sum(s.dropped for s in sites) / max(routed, 1)
     peak = (torch.cuda.max_memory_allocated() / (1 << 30)
             if dev.type == "cuda" else None)
     rep = pool.report() if hasattr(pool, "report") else {}
     return {"chunk": chunk, "context": context, "step_ms": ms,
             "chars_per_s": chunk / (ms * 1e-3), "peak_gb": peak,
+            "capacity_factor": sites[0].capacity_factor if sites else None,
+            "dropped_share": dropped,
             "resident": pool.resident, "experts": pool.n_experts(),
             "loads": getattr(pool, "loads", None),
             "tier_hit_rate": rep.get("hit_rate")}
 
 
+def bench_decode(dev, wdir, n_chars, prompt_len=512):
+    """
+    Writing: one character per forward behind a cache, as serve.py and the
+    sample log write. Timed with the expert dispatch running every expert on
+    the card, and running only the ones each character chose.
+    """
+    from minagi.build import build_paged
+    from minagi.pool import PooledMLP
+    from minagi.precision import amp
+
+    model, cfg, pool, _ = build_paged(wdir, dev, read_only=True)
+    model.eval()
+    sites = [m for m in model.modules() if isinstance(m, PooledMLP)]
+    prompt_len = max(1, min(prompt_len, cfg.block - n_chars - 2))
+    g = torch.Generator().manual_seed(0)
+    prompt = torch.randint(32, 127, (1, prompt_len), generator=g).to(dev)
+    out = {"resident": pool.resident, "experts": pool.n_experts()}
+    for name, subset in (("all_experts", False), ("chosen_only", True)):
+        for s in sites:
+            s.subset_dispatch = subset
+        caches = model.empty_caches()
+        with torch.no_grad(), amp(dev):
+            logits = model(prompt, caches=caches)[0]
+            cur = logits[:, -1:].argmax(-1)
+            pos = prompt_len
+            model(cur, caches=caches, pos_offset=pos)       # warm
+            pos += 1
+            sync(dev)
+            t = time.perf_counter()
+            for _ in range(n_chars):
+                logits = model(cur, caches=caches, pos_offset=pos)[0]
+                cur = logits[:, -1:].argmax(-1)
+                pos += 1
+            sync(dev)
+        dt = time.perf_counter() - t
+        out[name] = {"chars_per_s": n_chars / dt,
+                     "ms_per_char": dt * 1e3 / n_chars,
+                     "rows_per_char": getattr(model, "last_steps", None)}
+    return out
+
+
 def bench_paging(dev, wdir, reps):
     from minagi.paged import Tiers
     import j
FA5F
son as _j
@@ -232,7 +281,10 @@ def main():
                     help="fewer repetitions and shorter windows")
     ap.add_argument("--out", default=None)
     ap.add_argument("--skip", default="",
-                    help="comma list of attention,experts,step,paging")
+                    help="comma list of attention,experts,step,paging,decode")
+    ap.add_argument("--weights", default=None,
+                    help="time writing on this weights directory (opened "
+                         "read-only) instead of a throwaway model")
     a = ap.parse_args()
 
     from minagi import device as hwdev
@@ -286,7 +338,20 @@ def main():
               f"({e['fwd_tflops']:.1f} TFLOP/s)  fwd+bwd "
               f"{e['fwd_bwd_ms']:.2f} ms ({e['fwd_bwd_tflops']:.1f} TFLOP/s)")
 
-    if not ({"step", "paging"} <= skip):
+    def show_decode(d):
+        for k in ("all_experts", "chosen_only"):
+            v = d[k]
+            print(f"  {k:12} {v['chars_per_s']:7.1f} char/s  "
+                  f"({v['ms_per_char']:.0f} ms a character)")
+
+    if "decode" not in skip and a.weights:
+        res["decode"] = bench_decode(dev, a.weights, 16 if a.quick else 64)
+        print(f"\nwriting, {a.weights} ({res['decode']['resident']} on the "
+              f"card of {res['decode']['experts']}):")
+        show_decode(res["decode"])
+
+    if not ({"step", "paging"} <= skip) or ("decode" not in skip
+                                             and not a.weights):
         from minagi.create import create
         tmp = tempfile.mkdtemp(prefix="minagi-bench-")
         wdir = os.path.join(tmp, "weights")
@@ -299,6 +364,10 @@ def main():
                     print(f"  tier in {k:6}: disk -> tier "
                           f"{v['disk_to_tier_ms']:.1f} ms, tier -> slot "
                           f"{v['tier_to_slot_ms_2_of_3_tensors']:.2f} ms")
+            if "decode" not in skip and not a.weights:
+                res["decode"] = bench_decode(dev, wdir, 16 if a.quick else 64)
+                print("\nwriting, throwaway model:")
+                show_decode(res["decode"])
             if "step" not in skip:
                 res["step"] = bench_step(dev, wdir, chunk, ctx,
                                          2 if a.quick else 5)
@@ -307,7 +376,9 @@ def main():
                       if s["peak_gb"] is not None else "")
                 print(f"\ntraining step: {s['chunk']} characters behind "
                       f"{s['context']}: {s['step_ms']:.0f} ms, "
-                      f"{s['chars_per_s']:,.0f} char/s{pk}")
+                      f"{s['chars_per_s']:,.0f} char/s{pk}, capacity "
+                      f"{s['capacity_factor']:g} dropped "
+                      f"{100 * s['dropped_share']:.1f}% of routing")
         finally:
             shutil.rmtree(tmp, ignore_errors=True)