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

2026-10-28 00:33:34 +01:00
parent 01dce2ffc6
commit 16c3bf81f7
4 changed files with 31 additions and 13 deletions
@@ -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),