prain: pin the RAM tier per expert (big pinned blocks take minutes on eGPU hosts); sim charges disk reads on the CPU path; measure CPU cost on cold experts

prain/experts.py

@@ -225,12 +225,10 @@ class TieredRoutedExperts(RoutedExperts):
                                    "from ram to disk")
         ram = list(self.placement.ram)
         self._ram_row = {e: row for row, e in enumerate(ram)}
-        index = torch.tensor(ram, dtype=torch.int64, device="cpu")
-        self._ram_w13, self._ram_w2 = (
-            torch.index_select(pack, 0, index, out=torch.empty(
-                (len(ram), *pack.shape[1:]), dtype=pack.dtype, device="cpu", pin_memory=True))
-            for pack in (self._pack_w13, self._pack_w2)
-        )
+        # One pinned tensor per expert: pinning a single 340 MiB block took 6-23 s on an
+        # eGPU host (hipHostMalloc), small blocks pin at the normal rate.
+        self._ram_w13 = [self._pack_w13[e].pin_memory() for e in ram]
+        self._ram_w2 = [self._pack_w2[e].pin_memory() for e in ram]
         for slot, expert in enumerate(self.placement.vram):
             self._load(expert, slot)
         if self._router and (self.prain.prefetch.mode == "pregate" or self.prain.trace_path):

prain/sim.py

@@ -31,9 +31,13 @@ class Cost:
     sync_ms: float = 0.05  # host sync per layer
     cpu_ms: float = 0.5  # CPU work per (token, expert) route
 
+    def read_ms(self):
+        """Reading one expert from the pack (the CPU path pays this for disk experts too)."""
+        return self.expert_mib * 1.048576 / self.disk_gbps
+
     def copy_ms(self, disk):
         mb = self.expert_mib * 1.048576
-        return mb / self.h2d_gbps + (mb / self.disk_gbps if disk else 0.0)
+        return mb / self.h2d_gbps + (self.read_ms() if disk else 0.0)
 
 
 def load_cost(path, expert_mib):
@@ -175,7 +179,8 @@ def simulate(forwards, placements, policy, prefetch, miss, cost, belady=False):
                     gpu = link
                 land = max((ready.pop((layer, e), 0.0) for e in slot_of), default=0.0)
                 gpu = max(gpu, land) + cost.round_ms + cost.route_ms * sum(counts[e] for e in slot_of)
-            cpu_end = t + cost.sync_ms + cost.cpu_ms * sum(counts[e] for e in cpu)
+            cpu_end = (t + cost.sync_ms + cost.cpu_ms * sum(counts[e] for e in cpu)
+                       + cost.read_ms() * sum(e in disk[layer] for e in cpu))
             later = []
             if cpu and miss.admit:
                 hot = sorted((e for e in cpu if e not in disk[layer]), key=lambda e: (-heat[e], e))

scripts/measure.py

@@ -73,15 +73,22 @@ def main():
     other_ms = args.dense_gib * 2**30 / 1e6 / device_gbps / args.layers
     flag = torch.zeros(1, device="cuda")
     sync_ms = ms(lambda: flag.add_(1).item())
-    w13_cpu = torch.randn(2 * INTER, HIDDEN, dtype=torch.bfloat16).pin_memory()
-    w2_cpu = torch.randn(HIDDEN, INTER, dtype=torch.bfloat16).pin_memory()
+    # Serving computes a different expert each time, read from RAM (not from the CPU caches).
+    experts = [(torch.randn(2 * INTER, HIDDEN, dtype=torch.bfloat16),
+                torch.randn(HIDDEN, INTER, dtype=torch.bfloat16)) for _ in range(23)]
     x = torch.randn(1, HIDDEN, dtype=torch.bfloat16)
     one = torch.ones(1, 1)
     ids = torch.zeros(1, 1, dtype=torch.int32)
+
+    def cpu_route(weights):
+        cpu_experts(x, one, ids, {0}, lambda e: weights)
+
+    for weights in experts[:3]:  # warm up like ms(): thread pool and BLAS setup are one-time costs
+        cpu_route(weights)
     start = time.perf_counter()
-    for _ in range(20):
-        cpu_experts(x, one, ids, {0}, lambda e: (w13_cpu, w2_cpu))
-    cpu_ms = (time.perf_counter() - start) * 1000 / 20
+    for weights in experts[3:]:
+        cpu_route(weights)
+    cpu_ms = (time.perf_counter() - start) * 1000 / len(experts[3:])
     print(json.dumps({"h2d_gbps": round(h2d_gbps, 2), "disk_gbps": round(disk_gbps, 2),
                       "round_ms": round(round_ms, 4), "route_ms": round(route_ms, 6),
                       "other_ms": round(other_ms, 4), "sync_ms": round(sync_ms, 4),

tests/test_sim.py

@@ -57,6 +57,14 @@ class Simulate(unittest.TestCase):
         copy = self.run_sim().ms_per_decode
         self.assertLess(self.run_sim(miss=Miss(mode="auto")).ms_per_decode, copy)
 
+    def test_cpu_path_pays_the_disk_read_of_disk_experts(self):
+        on_disk = {layer: LayerPlacement((), (), tuple(range(32)), 3) for layer in range(4)}
+        cpu = Miss(mode="cpu", admit=0)
+        ram = simulate(self.forwards, all_ram(3), Policy(), Prefetch(), cpu, Cost())
+        disk = simulate(self.forwards, on_disk, Policy(), Prefetch(), cpu, Cost())
+        self.assertAlmostEqual(disk.total_ms - ram.total_ms, ram.counters["cpu"] * Cost().read_ms(),
+                               places=3)
+
     def test_prefetch_helps_when_predictions_are_right(self):
         none = self.run_sim().hit_rate
         good = self.run_sim(prefetch=Prefetch(mode="pregate", budget=4))