From 01dce2ffc6c44026d8e5350ee8ec255ad43cd8fe Mon Sep 17 00:00:00 2001 From: fabolous005 Date: Tue, 27 Oct 2026 23:52:53 +0100 Subject: [PATCH] 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 | 52 ++++++++++++++++++++++++++ prain/cli.py | 8 +++- prain/sim.py | 8 +++- tests/test_sim.py | 10 +++++ 4 files changed, 75 insertions(+), 3 deletions(-) create mode 100644 patches/vllm-0.32-gemma4-fit-8gb.patch diff --git a/patches/vllm-0.32-gemma4-fit-8gb.patch b/patches/vllm-0.32-gemma4-fit-8gb.patch new file mode 100644 index 0000000..ef69a60 --- /dev/null +++ b/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, diff --git a/prain/cli.py b/prain/cli.py index 3bc586e..a3743b4 100644 --- a/prain/cli.py +++ b/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") diff --git a/prain/sim.py b/prain/sim.py index eb14cf8..939e5bb 100644 --- a/prain/sim.py +++ b/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) diff --git a/tests/test_sim.py b/tests/test_sim.py index 2874f1c..6a6adff 100644 --- a/tests/test_sim.py +++ b/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)