diff --git a/prain/experts.py b/prain/experts.py new file mode 100644 index 0000000..6fb7b3d --- /dev/null +++ b/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) diff --git a/prain/plugin.py b/prain/plugin.py new file mode 100644 index 0000000..e014b75 --- /dev/null +++ b/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") diff --git a/tests/test_kernel_match.py b/tests/test_kernel_match.py new file mode 100644 index 0000000..e74099c --- /dev/null +++ b/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) diff --git a/tests/test_plugin.py b/tests/test_plugin.py new file mode 100644 index 0000000..8bbe259 --- /dev/null +++ b/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__))