prain: tune searches the auto CPU threshold and writes --cpu-threads; vLLM patch: fit Gemma 4 in 8 GB (tied lm_head without a copy, rope caches sized by max_model_len)

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/sim.py

@@ -226,15 +226,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)
 
 

tests/test_sim.py

@@ -81,3 +81,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)