prain: vLLM tiered experts layer and plugin

prain/experts.py

@@ -0,0 +1,361 @@
+"""TieredRoutedExperts: vLLM's RoutedExperts with only some experts in VRAM.
+
+vLLM builds this class instead of RoutedExperts once prain.plugin.register() ran.
+Every expert is written once to a per-layer pack file on disk while vLLM loads the
+checkpoint. On the first forward the RAM tier is copied from the pack into pinned host
+memory and the pinned VRAM slots are filled. After that each forward:
+  0. drops the routes of draft tokens to experts that are not available ([draft] min_weight;
+     prain.drafts knows which rows are drafts),
+  1. counts routes per expert (one host sync),
+  2. lets the [miss] config send some missing experts to the CPU,
+  3. copies the other missing experts into slots chosen by the [policy] config and runs the
+     stock fused-MoE kernel on the slot tensors, in rounds if needed,
+  4. queues async copies on a side stream: [prefetch] guesses for a later layer, and the
+     hottest CPU-run experts (admit-behind),
+  5. adds the CPU result.
+Slot safety with async copies: a copy into a slot waits for the last kernel that read it
+(`_last_use`), and a kernel or demand copy waits for a pending async copy into its slot
+(`_ready_events`).
+"""
+
+import atexit
+import json
+import math
+import os
+
+import torch
+from vllm.config import get_current_vllm_config
+from vllm.logger import init_logger
+from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts
+
+from prain import drafts
+from prain.config import layer_index, load_config
+from prain.policy import SlotCache, write_stats
+from prain.sim import write_trace
+from prain.tiered import draft_rows, tiered_forward
+
+logger = init_logger("vllm.prain")  # under "vllm." so vLLM's log handler prints it
+LAYERS = {}  # layer index -> TieredRoutedExperts
+_COPY_STREAM = []  # one side stream shared by all layers: async copies stay in issue order
+_TRACE = []  # trace records not written yet
+_DRAFT_MASK = [None, None]  # (spans, their bool row mask on the GPU) of the current forward
+_DRAFT_TOKENS = [0]  # draft token rows seen, for the stats log
+
+
+def _install_draft_hook():
+    """Wrap the V2 model runner's prepare_inputs so prain.drafts sees each step's batch
+    layout: which request ids sit at which batch positions, and their token rows."""
+    from vllm.v1.worker.gpu.model_runner import GPUModelRunner
+
+    original = GPUModelRunner.prepare_inputs
+    if getattr(original, "prain", False):
+        return
+
+    def prepare_inputs(self, scheduler_output, *args, **kwargs):
+        batch = original(self, scheduler_output, *args, **kwargs)
+        if hasattr(scheduler_output, "scheduled_new_reqs"):
+            drafts.note_step(scheduler_output, batch.req_ids, batch.query_start_loc_np)
+        return batch
+
+    prepare_inputs.prain = True
+    GPUModelRunner.prepare_inputs = prepare_inputs
+
+
+def _draft_mask(x):
+    """Bool row mask of draft tokens for this forward, built once and shared by all layers."""
+    spans = drafts.spans()
+    if spans is None:
+        return None
+    if max(end for _, end in spans) > x.shape[0]:  # not this forward's layout: run full
+        return None
+    if _DRAFT_MASK[0] is not spans:
+        _DRAFT_MASK[:] = [spans, draft_rows(spans, x.shape[0], x.device)]
+        _DRAFT_TOKENS[0] += sum(end - first for first, end in spans)
+    return _DRAFT_MASK[1]
+
+
+def _copy_stream():
+    if not _COPY_STREAM:
+        _COPY_STREAM.append(torch.cuda.Stream())
+    return _COPY_STREAM[0]
+
+
+def _file_tensor(path, shape, dtype):
+    """A CPU tensor backed by `path` (created if missing); writes land in the file."""
+    return torch.from_file(
+        path, shared=True, size=math.prod(shape), dtype=dtype, device="cpu"
+    ).view(shape)
+
+
+def _mem_available():
+    with open("/proc/meminfo") as f:
+        for line in f:
+            if line.startswith("MemAvailable:"):
+                return int(line.split()[1]) * 1024
+    return None
+
+
+def _flush_trace():
+    if _TRACE and LAYERS:
+        write_trace(next(iter(LAYERS.values())).prain.trace_path, _TRACE, mode="at")
+        _TRACE.clear()
+
+
+atexit.register(_flush_trace)
+
+
+class TieredRoutedExperts(RoutedExperts):
+    def __init__(self, layer_name, params_dtype, moe_config, quant_config, *args, **kwargs):
+        vllm_config = get_current_vllm_config()
+        parallel = moe_config.moe_parallel_config
+        layer = layer_index(layer_name)
+        if layer is None:
+            raise ValueError(f"prain: cannot find the layer number in {layer_name!r}")
+        if quant_config is not None:
+            raise NotImplementedError("prain: quantized MoE weights are not supported")
+        if not vllm_config.model_config.enforce_eager:
+            raise ValueError("prain: pass --enforce-eager (expert copies wait on the host)")
+        if vllm_config.kernel_config.moe_backend != "triton":
+            raise ValueError("prain: pass --moe-backend triton")
+        if parallel.tp_size > 1 or parallel.dp_size > 1 or parallel.use_ep:
+            raise NotImplementedError("prain: one GPU only (no TP, DP or EP)")
+        if vllm_config.use_v2_model_runner:
+            _install_draft_hook()
+        elif not LAYERS:
+            logger.warning("prain: drafts need vLLM's V2 model runner; with "
+                           "VLLM_USE_V2_MODEL_RUNNER=0 every request runs with all its experts")
+        config = load_config(os.environ["PRAIN_CONFIG"])
+        num_experts = moe_config.num_experts
+        placement = config.placement(layer, num_experts)
+        # Gemma 4 passes its router module's routing_function: keep the module for pregate.
+        router = getattr(kwargs.get("custom_routing_function"), "__self__", None)
+        if config.prefetch.mode == "pregate" and not isinstance(router, torch.nn.Module):
+            raise NotImplementedError("prain: prefetch.mode = pregate needs a router module "
+                                      "passed as custom_routing_function (Gemma 4 does)")
+        # vLLM sizes w13/w2 and the kernel by num_experts: make that the VRAM slot count.
+        moe_config.num_experts = moe_config.num_local_experts = placement.num_slots
+        super().__init__(layer_name, params_dtype, moe_config, quant_config, *args, **kwargs)
+        self.local_num_experts = placement.num_slots
+        self.prain = config
+        self.prain_layer = layer
+        self.placement = placement
+        disk = set(placement.disk)
+        cost = [config.policy.disk_cost if e in disk else 1.0 for e in range(num_experts)]
+        self.cache = SlotCache(placement.vram, num_experts, placement.cache_slots, config.policy, cost)
+        self._router = [router] if isinstance(router, torch.nn.Module) else []  # not a submodule
+        self._async = config.prefetch.mode != "none" or (config.miss.mode != "copy" and config.miss.admit)
+        self._ready_events = {}  # slot -> event: an async copy into the slot lands
+        self._last_use = {}  # slot -> event: the last kernel reading the slot finished
+        self._predictions = {}  # distance -> [(expert, probability)] from this layer's router
+        self._ram_row = None
+        self._open_pack(config.pack_dir, layer, num_experts, vllm_config.model_config.model)
+        self._ready = False
+        self._steps = 0
+        LAYERS[layer] = self
+        logger.info(
+            "prain layer %d: %d vram, %d ram, %d disk, %d cache slots%s",
+            layer, len(placement.vram), len(placement.ram), len(placement.disk),
+            placement.cache_slots, "" if self.pack_complete else ", writing pack",
+        )
+
+    def _open_pack(self, pack_dir, layer, num_experts, model):
+        os.makedirs(pack_dir, exist_ok=True)
+        stem = os.path.join(pack_dir, f"layer{layer:03d}")
+        w13_shape = (num_experts, *self.w13_weight.shape[1:])
+        w2_shape = (num_experts, *self.w2_weight.shape[1:])
+        self._pack_files = (f"{stem}.w13", f"{stem}.w2")
+        self._pack_w13 = _file_tensor(self._pack_files[0], w13_shape, self.w13_weight.dtype)
+        self._pack_w2 = _file_tensor(self._pack_files[1], w2_shape, self.w2_weight.dtype)
+        self._marker_path = f"{stem}.done"
+        self._marker = {"model": model, "w13": list(w13_shape), "w2": list(w2_shape),
+                        "dtype": str(self.w13_weight.dtype)}
+        self.pack_complete = os.path.exists(self._marker_path)
+        if self.pack_complete:
+            with open(self._marker_path) as f:
+                found = json.load(f)
+            if found != self._marker:
+                raise ValueError(f"prain: {self._marker_path} is for {found}, not "
+                                 f"{self._marker}; delete {pack_dir} or use another pack_dir")
+
+    def get_expert_mapping(self, ckpt_gate_proj_name=None, ckpt_down_proj_name=None,
+                           ckpt_up_proj_name=None, include_fused=False):
+        # The base class derives "redundant experts" as num_experts - num_logical_experts,
+        # which is negative here because num_experts counts VRAM slots.
+        return self.build_expert_params_mapping(
+            ckpt_gate_proj_name or self.ckpt_gate_proj_name,
+            ckpt_down_proj_name or self.ckpt_down_proj_name,
+            ckpt_up_proj_name or self.ckpt_up_proj_name,
+            num_experts=self.moe_config.num_logical_experts,
+            routed_experts_prefix="",
+            lora_base_layer_prefix=self.lora_base_layer_prefix,
+            include_fused=include_fused,
+            is_gated=self.moe_config.is_act_and_mul,
+        )
+
+    def _map_global_expert_id_to_local_expert_id(self, expert_id):
+        return expert_id  # the loader only ever writes into the [num_experts, ...] pack
+
+    def weight_loader(self, param, loaded_weight, weight_name, shard_id, expert_id,
+                      return_success=False):
+        # Every expert goes to the pack; the VRAM and RAM tiers fill from it later.
+        if not self.pack_complete:
+            pack = self._pack_w13 if param is self.w13_weight else self._pack_w2
+            super().weight_loader(pack, loaded_weight, weight_name, shard_id, expert_id,
+                                  return_success=True)
+        return True if return_success else None
+
+    weight_loader.supports_moe_loading = True
+
+    def _materialize(self):
+        if not self.pack_complete:
+            for path in self._pack_files:  # mmap writes -> disk before we trust the pack
+                fd = os.open(path, os.O_RDONLY)
+                os.fsync(fd)
+                os.close(fd)
+            with open(self._marker_path, "w") as f:
+                json.dump(self._marker, f)
+            self.pack_complete = True
+        if self.prain_layer == min(LAYERS):  # check the whole RAM tier once, before pinning
+            row = self._pack_w13[0].nbytes + self._pack_w2[0].nbytes
+            need = row * sum(len(experts.placement.ram) for experts in LAYERS.values())
+            free = _mem_available()
+            if free is not None and need > 0.9 * free:
+                raise RuntimeError(f"prain: the RAM tier needs {need / 2**30:.1f} GiB pinned, "
+                                   f"but only {free / 2**30:.1f} GiB are available; move experts "
+                                   "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)
+        )
+        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):
+            self._router[0].register_forward_pre_hook(self._predict)
+
+    def _source(self, expert):
+        """CPU tensors of an expert: pinned rows for the RAM tier, else the mmap'd pack."""
+        row = self._ram_row.get(expert)
+        if row is None:
+            return self._pack_w13[expert], self._pack_w2[expert]
+        return self._ram_w13[row], self._ram_w2[row]
+
+    def _load(self, expert, slot):
+        """Demand copy on the compute stream."""
+        pending = self._ready_events.pop(slot, None)
+        if pending is not None:  # an older async copy into this slot must land first
+            torch.cuda.current_stream().wait_event(pending)
+        w13, w2 = self._source(expert)
+        self.w13_weight.data[slot].copy_(w13, non_blocking=True)
+        self.w2_weight.data[slot].copy_(w2, non_blocking=True)
+
+    def _load_async(self, expert, slot):
+        """Copy on the side stream; only RAM-tier (pinned) experts are truly asynchronous."""
+        stream = _copy_stream()
+        last = self._last_use.get(slot)
+        if last is not None:  # a queued kernel may still read the slot's old expert
+            stream.wait_event(last)
+        w13, w2 = self._source(expert)
+        with torch.cuda.stream(stream):
+            self.w13_weight.data[slot].copy_(w13, non_blocking=True)
+            self.w2_weight.data[slot].copy_(w2, non_blocking=True)
+        event = torch.cuda.Event()
+        event.record(stream)
+        self._ready_events[slot] = event
+
+    def _run_slots(self, x, weights, slot_ids, used_slots):
+        current = torch.cuda.current_stream()
+        for slot in used_slots:
+            pending = self._ready_events.pop(slot, None)
+            if pending is not None:
+                current.wait_event(pending)
+        out = super().forward_modular(x, weights, slot_ids)
+        if not isinstance(out, torch.Tensor):
+            raise NotImplementedError("prain: deferred MoE finalize is not supported")
+        if self._async:
+            event = torch.cuda.Event()
+            event.record(current)
+            for slot in used_slots:
+                self._last_use[slot] = event
+        return out
+
+    def _predict(self, module, args):
+        """Pre-hook of this layer's router: guess the experts of later layers from the same
+        residual stream (their routers normalize their input, so the guess is cheap and
+        usually close)."""
+        distances = (1, 2) if self.prain.trace_path else (self.prain.prefetch.distance,)
+        self._predictions = {}
+        for d in distances:
+            target = LAYERS.get(self.prain_layer + d)
+            if target is None or not target._router:
+                continue
+            logits = target._router[0].forward(args[0])  # .forward: skip the target's own hook
+            probs = torch.softmax(logits.float(), dim=-1).mean(0)
+            p, e = probs.topk(min(2 * self.top_k, probs.numel()))
+            self._predictions[d] = list(zip(e.tolist(), p.tolist()))  # host sync
+
+    def _cpu_weights(self, expert):
+        return self._source(expert)
+
+    def _after_enqueue(self, cpu, stats):
+        miss, prefetch = self.prain.miss, self.prain.prefetch
+        if cpu and miss.admit:  # admit-behind: the hottest CPU-run RAM experts go to VRAM
+            hot = sorted((e for e in cpu if e in self._ram_row), key=lambda e: (-stats[e][1], e))
+            for e, slot in self.cache.prefetch([(e, 1.0) for e in hot], miss.admit):
+                self._load_async(e, slot)
+        target = LAYERS.get(self.prain_layer + prefetch.distance)
+        if prefetch.mode == "pregate" and target is not None and target._ready:
+            guesses = [(e, p) for e, p in self._predictions.get(prefetch.distance, ())
+                       if p >= prefetch.min_prob and e in target._ram_row]
+            for e, slot in target.cache.prefetch(guesses, prefetch.budget):
+                target._load_async(e, slot)
+
+    def _record(self, stats, tokens):
+        ids = sorted(stats)
+        record = {"l": self.prain_layer, "t": tokens, "e": ids,
+                  "c": [stats[e][0] for e in ids], "w": [round(stats[e][1], 5) for e in ids],
+                  "r": [round(stats[e][2], 5) for e in ids]}
+        if self._predictions:
+            record["p"] = {str(d): [[e, round(p, 5)] for e, p in guesses]
+                           for d, guesses in self._predictions.items()}
+        _TRACE.append(record)
+
+    def forward_modular(self, x, topk_weights, topk_ids, shared_experts=None,
+                        shared_experts_input=None):
+        if shared_experts is not None:
+            raise NotImplementedError("prain: fused shared experts are not supported")
+        if not self._ready:
+            self._materialize()
+            self._ready = True
+        policy = self.prain.policy
+        out = tiered_forward(
+            x, topk_weights, topk_ids, self.cache, self._load, self._run_slots,
+            policy.heat, policy.recency_half_life, self.prain.miss, self._cpu_weights,
+            on_stats=self._record if self.prain.trace_path else None,
+            after_enqueue=self._after_enqueue,
+            draft=_draft_mask(x), min_weight=self.prain.draft.min_weight,
+        )
+        if self.prain_layer == max(LAYERS):
+            drafts.clear()  # the next forward may be a warm-up run without a batch layout
+        self._maybe_write_stats()
+        return out
+
+    def _maybe_write_stats(self):
+        self._steps += 1
+        if self._steps % self.prain.stats_every or self.prain_layer != min(LAYERS):
+            return
+        if self.prain.trace_path:
+            _flush_trace()
+        if not self.prain.stats_path:
+            return
+        write_stats(self.prain.stats_path, {layer: e.cache for layer, e in LAYERS.items()})
+        c = {k: sum(getattr(e.cache, k) for e in LAYERS.values())
+             for k in ("hits", "misses", "cpu", "prefetches", "prefetch_hits")}
+        logger.info("prain: VRAM hit rate %.1f%% (%d hits, %d copies, %d on CPU, "
+                    "%d prefetched, %d of them used, %d draft tokens), stats in %s",
+                    100 * c["hits"] / max(1, c["hits"] + c["misses"] + c["cpu"]), c["hits"],
+                    c["misses"], c["cpu"], c["prefetches"], c["prefetch_hits"],
+                    _DRAFT_TOKENS[0], self.prain.stats_path)

prain/plugin.py

@@ -0,0 +1,16 @@
+"""vLLM general plugin entry point (see pyproject.toml). Inactive unless PRAIN_CONFIG is set."""
+
+import os
+
+
+def register():
+    if not os.environ.get("PRAIN_CONFIG"):
+        return
+    from vllm.model_executor.custom_op import PluggableLayer, op_registry_oot
+
+    if "RoutedExperts" in op_registry_oot:  # vLLM may call register() more than once
+        return
+    from prain.experts import TieredRoutedExperts
+
+    # vLLM's documented out-of-tree hook: every RoutedExperts(...) now builds ours.
+    PluggableLayer.register_oot(TieredRoutedExperts, name="RoutedExperts")

tests/test_kernel_match.py

@@ -0,0 +1,33 @@
+import importlib.util
+import unittest
+
+try:
+    import torch
+except ImportError:
+    torch = None
+
+GPU = torch is not None and torch.cuda.is_available() and importlib.util.find_spec("vllm") is not None
+
+
+@unittest.skipUnless(GPU, "needs vLLM and a GPU")
+class KernelMatch(unittest.TestCase):
+    def test_cpu_experts_compute_what_the_vllm_kernel_computes(self):
+        from vllm.config import VllmConfig, set_current_vllm_config
+        from vllm.model_executor.layers.fused_moe.activation import MoEActivation
+        from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts
+
+        from prain.tiered import cpu_experts
+
+        torch.manual_seed(0)
+        tokens, hidden, inter, experts = 8, 256, 128, 4
+        x = torch.randn(tokens, hidden, dtype=torch.bfloat16)
+        w13 = torch.randn(experts, 2 * inter, hidden, dtype=torch.bfloat16) * 0.1
+        w2 = torch.randn(experts, hidden, inter, dtype=torch.bfloat16) * 0.1
+        topk_ids = torch.stack([torch.randperm(experts)[:2] for _ in range(tokens)]).int()
+        topk_weights = torch.rand(tokens, 2)
+        with set_current_vllm_config(VllmConfig()):
+            gpu = fused_experts(x.cuda(), w13.cuda(), w2.cuda(), topk_weights.cuda(),
+                                topk_ids.cuda(), activation=MoEActivation.GELU_TANH,
+                                global_num_experts=experts).float().cpu()
+        cpu = cpu_experts(x, topk_weights, topk_ids, set(range(experts)), lambda e: (w13[e], w2[e]))
+        torch.testing.assert_close(cpu, gpu, rtol=3e-2, atol=3e-2)

tests/test_plugin.py

@@ -0,0 +1,40 @@
+import importlib.metadata
+import importlib.util
+import os
+import unittest
+from unittest import mock
+
+from tests.helpers import TmpCase
+
+
+@unittest.skipIf(importlib.util.find_spec("vllm") is None, "needs vllm")
+class Plugin(TmpCase):
+    def test_entry_point_is_installed(self):
+        plugins = importlib.metadata.entry_points(group="vllm.general_plugins")
+        self.assertIn("prain", [p.name for p in plugins])
+
+    def test_register_swaps_in_tiered_experts(self):
+        config = self.write("p.toml", f'pack_dir = "{self.dir}/pack"\n')
+        with mock.patch.dict(os.environ, {"PRAIN_CONFIG": config}):
+            from prain.plugin import register
+
+            register()
+            register()  # vLLM may call it again; must not raise
+        from vllm.model_executor.custom_op import op_registry_oot
+
+        from prain.experts import TieredRoutedExperts
+
+        self.assertIs(op_registry_oot["RoutedExperts"], TieredRoutedExperts)
+
+    def test_draft_hook_wraps_the_v2_model_runner_once(self):
+        from vllm.v1.worker.gpu.input_batch import InputBatch
+        from vllm.v1.worker.gpu.model_runner import GPUModelRunner
+
+        from prain.experts import _install_draft_hook
+
+        _install_draft_hook()
+        wrapped = GPUModelRunner.prepare_inputs
+        _install_draft_hook()
+        self.assertIs(GPUModelRunner.prepare_inputs, wrapped)
+        self.assertTrue(wrapped.prain)
+        self.assertLessEqual({"req_ids", "query_start_loc_np"}, set(InputBatch.__dataclass_fields__))