prain: README; example config with the measured eGPU defaults (miss auto <=2 tokens, 16 CPU threads, no pinned RAM tier)

README.md

@@ -0,0 +1,45 @@
+# prain
+
+Serve a mixture-of-experts model with vLLM when its experts do not fit in VRAM. A TOML file places
+each expert of each layer in VRAM (pinned or in a cache slot), in RAM, or in a pack file on disk.
+Missing experts are copied to the GPU or computed on the CPU (`[miss]`). Eviction, heat, prefetch
+and miss policies are configurable, and `prain tune` picks them from a routing trace.
+Target: Gemma 4 26B-A4B (30 layers x 128 experts, 42.5 GiB of experts) on an 8 GB RX 6600 XT.
+
+## Build (GPU host, user space)
+
+    scripts/build.sh wrappers sdk fetch torch torchvision vllm prain
+
+- `sdk`: venv with AMD's ROCm SDK wheels for `PRAIN_GPU_ARCH` (default gfx1032)
+- `torch`, `torchvision`: from source with clang and `-march=native` (hours)
+- `vllm`: v0.32.0 from source with gfx1032 added and `patches/vllm-*.patch` applied
+- `prain`: editable install and the test suite
+
+Check the stack: `scripts/smoke.sh MODEL_DIR --attention-backend TRITON_ATTN`
+
+## Serve
+
+    . scripts/env.sh
+    prain serve examples/gemma4-26b-a4b-rx6600xt.toml MODEL_DIR -- --served-model-name gemma4 \
+        --language-model-only --max-model-len 4096 --max-num-batched-tokens 1024 \
+        --enable-prefix-caching --attention-backend TRITON_ATTN
+
+The first start writes the pack (~3 min). Chat with `prain tui`, or `prain web --config FILE` for
+several users with drafts.
+
+## Tune for a machine
+
+    python scripts/measure.py PACK_DIR/layer029.w13 --threads 16 > cost.json   # server stopped
+    prain sim TRACE CONFIG MODEL_DIR/config.json --cost cost.json
+    prain tune TRACE MODEL_DIR/config.json --cost cost.json --expert-vram-gib 1.67 --ram-gib 0 \
+        --pack-dir PACK_DIR --cpu-threads 16 > tuned.toml
+    python scripts/bench.py http://127.0.0.1:8000
+
+The simulator treats every disk-tier expert as a disk read; the page cache makes real runs faster.
+
+## vLLM patches
+
+- `vllm-0.32-rocm-no-cuda-families`: gfx1032 reports capability 10.3, and vLLM's CUDA "family"
+  checks took it for Blackwell (B200 attention tiles: 128 KiB of LDS on a 64 KiB GPU).
+- `vllm-0.32-gemma4-fit-8gb`: Gemma 4 shares `embed_tokens` as `lm_head` instead of allocating a
+  second 1.4 GiB copy while loading, and sizes rope caches by `max_model_len` (saves ~0.4 GiB).

examples/gemma4-26b-a4b-rx6600xt.toml

@@ -1,6 +1,9 @@
 # prain placement: prithivMLmods/gemma-4-26B-A4B-it-Uncensored-MAX on an 8 GB RX 6600 XT
 # Model: 30 MoE layers x 128 experts, top-8 routing, 11.34 MiB per expert (BF16).
-# VRAM: 8 GiB - 4.45 GiB dense weights - 0.75 GiB KV cache - ~0.9 GiB runtime = ~1.7 GiB experts.
+# VRAM, measured: 6.22 GiB weights (4.46 dense + 5 slots x 30 layers) + 0.75 GiB KV cache: 7.9 of
+# 8 GiB in use. Serve with --attention-backend TRITON_ATTN (ROCM_ATTN needs more LDS than RDNA2 has).
+# RAM tier: pinning is very slow behind a strict IOMMU (Thunderbolt eGPU: minutes per GiB), so this
+# file keeps no RAM tier; the kernel page cache holds the hot part of the pack instead.
 #
 # Tiers:  vram = pinned in a VRAM slot for the whole run
 #         ram  = pinned host memory, copied into a cache slot when a token needs it
@@ -38,10 +41,11 @@ budget = 2            # experts per layer and forward
 min_prob = 0.05       # skip guesses below this router probability
 
 [miss]                # an expert that is not in VRAM
-mode = "copy"         # copy | cpu (always compute it on the CPU) | auto (CPU if it has
-cpu_max_tokens = 1    #   at most cpu_max_tokens tokens, else copy)
-cpu_threads = 4       # CPU threads of the vLLM worker (sets OMP_NUM_THREADS)
-admit = 1             # cpu/auto: copy this many CPU-run experts into VRAM afterwards
+mode = "auto"         # copy | cpu (always compute it on the CPU) | auto (CPU if it has
+cpu_max_tokens = 2    #   at most cpu_max_tokens tokens, else copy). Measured on the Thunderbolt
+                      #   eGPU host: a copy costs ~4.6 ms + disk read, the CPU ~2.5 ms per token
+cpu_threads = 16      # CPU threads of the vLLM worker (sets OMP_NUM_THREADS)
+admit = 0             # cpu/auto: copy this many CPU-run RAM-tier experts into VRAM afterwards
 
 [serve]               # parallel requests (`prain serve` passes these to vLLM)
 max_requests = 2      # requests vLLM runs at once (--max-num-seqs); warm-ups count too

patches/vllm-0.32-gemma4-fit-8gb.patch

@@ -0,0 +1,52 @@
+diff --git a/vllm/model_executor/models/gemma4.py b/vllm/model_executor/models/gemma4.py
+index d1498f1..745d1d3 100644
+--- a/vllm/model_executor/models/gemma4.py
++++ b/vllm/model_executor/models/gemma4.py
+@@ -26,7 +26,7 @@ import torch
+ from torch import nn
+ 
+ from vllm.compilation.decorators import support_torch_compile
+-from vllm.config import CacheConfig, VllmConfig
++from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config
+ from vllm.config.utils import getattr_iter
+ from vllm.distributed import (
+     get_pp_group,
+@@ -491,9 +491,12 @@ class Gemma4Attention(nn.Module):
+                     f"{kv_shared_layer_index}.self_attn.attn"
+                 )
+ 
++        # Gemma 4's rope types (default, proportional) do not depend on max_position, so
++        # size the cos/sin caches for max_model_len: 262144 positions take ~0.4 GiB.
++        max_model_len = get_current_vllm_config().model_config.max_model_len
+         self.rotary_emb = get_rope(
+             self.head_dim,
+-            max_position=max_position_embeddings,
++            max_position=min(max_position_embeddings, max_model_len),
+             rope_parameters=rope_parameters,
+             is_neox_style=True,
+         )
+@@ -1444,14 +1447,17 @@ class Gemma4ForCausalLM(
+             prefix=maybe_prefix(prefix, "model"),
+         )
+ 
+-        self.lm_head = ParallelLMHead(
+-            config.vocab_size,
+-            config.hidden_size,
+-            quant_config=quant_config,
+-            prefix=maybe_prefix(prefix, "lm_head"),
+-        )
+         if config.tie_word_embeddings:
+-            self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens)
++            # Share embed_tokens: a ParallelLMHead tied after construction would hold the
++            # vocab matrix (1.4 GiB) twice while loading, which overflows 8 GB GPUs.
++            self.lm_head = self.model.embed_tokens
++        else:
++            self.lm_head = ParallelLMHead(
++                config.vocab_size,
++                config.hidden_size,
++                quant_config=quant_config,
++                prefix=maybe_prefix(prefix, "lm_head"),
++            )
+ 
+         self.logits_processor = LogitsProcessor(
+             config.vocab_size,

prain/cli.py

@@ -43,6 +43,8 @@ def _variant_text(variant):
     policy, prefetch, miss = variant
     fetch = "none" if prefetch.mode == "none" else f"pregate d{prefetch.distance} b{prefetch.budget}"
     cpu = miss.mode if miss.mode == "copy" else f"{miss.mode} admit{miss.admit}"
+    if miss.mode == "auto":
+        cpu += f" <={miss.cpu_max_tokens}tok"
     return (f"{policy.eviction:5} {policy.heat:8} keep{policy.scan_keep} "
             f"promote{policy.promote}  {fetch:15} {cpu}")
 
@@ -64,13 +66,15 @@ def sim(trace, config_path, hf_config_path, cost_path, expert_mib):
 
 
 def tune_command(args):
+    from prain.config import Miss, Policy, Prefetch
     from prain.sim import describe, load_cost, read_forwards, tune
 
     _, num_experts, _, size = _model_dims(args.hf_config)
     cost = load_cost(args.cost, args.expert_mib or size / 2**20)
     forwards = read_forwards(args.trace)
     best, ram, per, final, bound, rows = tune(forwards, num_experts, args.expert_vram_gib,
-                                              args.ram_gib, cost)
+                                              args.ram_gib, cost,
+                                              base=(Policy(), Prefetch(), Miss(cpu_threads=args.cpu_threads)))
     print(" rank  total_s  decode_ms  hit%  eviction heat keep promote  prefetch  miss",
           file=sys.stderr)
     for rank, (variant, r) in enumerate(rows[: args.top], 1):
@@ -139,6 +143,8 @@ def main(argv=None):
             s.add_argument("--expert-vram-gib", type=float, required=True)
             s.add_argument("--ram-gib", type=float, required=True)
             s.add_argument("--pack-dir", required=True)
+            s.add_argument("--cpu-threads", type=int, default=4,
+                           help="[miss] cpu_threads to write: the --threads scripts/measure.py used")
             s.add_argument("--top", type=int, default=10, help="variants to list")
     y = sub.add_parser("synth", help="write a synthetic routing trace, to try sim and tune")
     y.add_argument("out")

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))
@@ -226,15 +231,19 @@ DIMENSIONS = (
     ("prefetch", "mode", ("none", "pregate")),
     ("prefetch", "distance", (1, 2)),
     ("prefetch", "budget", (1, 2, 4)),
-    ("miss", "mode", ("copy", "auto", "cpu")),
+    # One dimension: cpu_max_tokens only matters for auto, and coordinate descent would never
+    # reach auto with a high threshold once auto <= 1 lost to cpu.
+    ("miss", ("mode", "cpu_max_tokens"),
+     (("copy", 1), ("cpu", 1)) + tuple(("auto", n) for n in (1, 2, 4, 8, 16))),
     ("miss", "admit", (0, 1, 2)),
 )
 _PART = {"policy": 0, "prefetch": 1, "miss": 2}
 
 
 def _with(variant, table, name, value):
+    names, values = (name, value) if isinstance(name, tuple) else ((name,), (value,))
     parts = list(variant)
-    parts[_PART[table]] = replace(parts[_PART[table]], **{name: value})
+    parts[_PART[table]] = replace(parts[_PART[table]], **dict(zip(names, values)))
     return tuple(parts)
 
 

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))
@@ -81,3 +89,13 @@ class Tune(unittest.TestCase):
         self.assertLessEqual(final.total_ms, default.total_ms)
         self.assertEqual(rows[0][1].total_ms, min(r.total_ms for _, r in rows))
         self.assertGreaterEqual(bound.hit_rate, 0.0)
+
+    def test_tune_searches_the_auto_threshold_and_keeps_the_cpu_threads(self):
+        forwards = forwards_of(synthetic_trace(**SMALL))
+        cost = Cost(expert_mib=64.0, h2d_gbps=2.5, cpu_ms=0.5)  # eGPU: a copy is ~50 CPU routes
+        base = (Policy(), Prefetch(), Miss(cpu_threads=16))
+        best, ram, per, final, bound, rows = tune(forwards, 32, 1.0, 1.0, cost, base=base)
+        self.assertEqual(best[2].cpu_threads, 16)
+        self.assertNotEqual(best[2].mode, "copy")
+        tried = {(v[2].mode, v[2].cpu_max_tokens) for v, _ in rows}
+        self.assertLessEqual({("copy", 1), ("cpu", 1), ("auto", 1), ("auto", 8), ("auto", 16)}, tried)