From 708e152549949bec86183123c50017b32fb89a8e Mon Sep 17 00:00:00 2001 From: fabolous005 Date: Wed, 28 Oct 2026 22:07:25 +0100 Subject: [PATCH] prain: int8/int4 experts quantized once on the GPU, rounds through vLLM's expert_map, [miss] min_weight skips experts not in VRAM; scripts: tune_moe, quality (answer NLL on identical token ids), bench --context, measure --quant --- prain/cli.py | 20 +++-- prain/config.py | 26 ++++-- prain/experts.py | 171 ++++++++++++++++++++++++++----------- prain/quant.py | 89 +++++++++++++++++++ prain/tiered.py | 91 +++++++++++--------- scripts/bench.py | 17 +++- scripts/env.sh | 3 + scripts/measure.py | 21 ++++- scripts/quality.py | 89 +++++++++++++++++++ scripts/tune_moe.py | 111 ++++++++++++++++++++++++ tests/test_cli.py | 14 +++ tests/test_config.py | 17 ++++ tests/test_kernel_match.py | 61 +++++++++++++ tests/test_quant.py | 49 +++++++++++ tests/test_tiered.py | 44 ++++++++-- 15 files changed, 703 insertions(+), 120 deletions(-) create mode 100644 prain/quant.py create mode 100644 scripts/quality.py create mode 100644 scripts/tune_moe.py create mode 100644 tests/test_quant.py diff --git a/prain/cli.py b/prain/cli.py index a3743b4..7aed86c 100644 --- a/prain/cli.py +++ b/prain/cli.py @@ -7,19 +7,19 @@ import os import sys from dataclasses import asdict -from prain.config import (GIB, Miss, Policy, Prefetch, expert_bytes, format_ids, load_config, +from prain.config import (GIB, QUANTS, Miss, Policy, Prefetch, expert_bytes, format_ids, load_config, moe_dims, plan, render_toml) -def _model_dims(hf_config_path): +def _model_dims(hf_config_path, quant="bf16"): with open(hf_config_path) as f: num_layers, num_experts, top_k, hidden, inter = moe_dims(json.load(f)) - return num_layers, num_experts, top_k, expert_bytes(hidden, inter) + return num_layers, num_experts, top_k, expert_bytes(hidden, inter, quant) def check(config_path, hf_config_path): config = load_config(config_path) - num_layers, num_experts, top_k, size = _model_dims(hf_config_path) + num_layers, num_experts, top_k, size = _model_dims(hf_config_path, config.quant) print("layer vram ram disk cache_slots") vram_rows = ram_rows = few_slots = 0 for layer in range(num_layers): @@ -53,7 +53,7 @@ def sim(trace, config_path, hf_config_path, cost_path, expert_mib): from prain.sim import describe, load_cost, read_forwards, simulate config = load_config(config_path) - _, num_experts, _, size = _model_dims(hf_config_path) + _, num_experts, _, size = _model_dims(hf_config_path, config.quant) forwards = read_forwards(trace) cost = load_cost(cost_path, expert_mib or size / 2**20) layers = sorted({layer for fwd in forwards for layer in fwd}) @@ -69,7 +69,7 @@ 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) + _, num_experts, _, size = _model_dims(args.hf_config, args.quant) 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, @@ -88,7 +88,7 @@ def tune_command(args): for layer in per} print(render_toml( f"generated by `prain tune` from {os.path.basename(args.trace)}", - {"pack_dir": os.path.abspath(args.pack_dir), "default_tier": "disk", + {"pack_dir": os.path.abspath(args.pack_dir), "quant": None if args.quant == "bf16" else args.quant, "default_tier": "disk", "cache_slots": min(per.values()), "trace_path": os.path.abspath(args.trace)}, {"policy": asdict(policy), "prefetch": asdict(prefetch), "miss": asdict(miss)}, layers, @@ -130,6 +130,7 @@ def main(argv=None): p.add_argument("--ram-gib", type=float, required=True, help="pinned host memory for experts") p.add_argument("--cache-slots", type=int, default=16) p.add_argument("--pack-dir", required=True) + p.add_argument("--quant", choices=QUANTS, default="bf16", help="expert weight format") for name, text in (("sim", "simulate a placement file on a routing trace"), ("tune", "search policies and per-layer slots on a trace; prints a placement file")): s = sub.add_parser(name, help=text) @@ -143,6 +144,7 @@ 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("--quant", choices=QUANTS, default="bf16", help="expert weight format") 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") @@ -172,11 +174,11 @@ def main(argv=None): if args.command == "check": check(args.config, args.hf_config) elif args.command == "plan": - _, num_experts, _, size = _model_dims(args.hf_config) + _, num_experts, _, size = _model_dims(args.hf_config, args.quant) with open(args.stats) as f: stats = json.load(f) print(plan(stats, num_experts, size, args.expert_vram_gib, args.ram_gib, - args.cache_slots, os.path.abspath(args.pack_dir), os.path.abspath(args.stats)), end="") + args.cache_slots, os.path.abspath(args.pack_dir), os.path.abspath(args.stats), args.quant), end="") elif args.command == "sim": sim(args.trace, args.config, args.hf_config, args.cost, args.expert_mib) elif args.command == "tune": diff --git a/prain/config.py b/prain/config.py index 4150da3..1c9d5a1 100644 --- a/prain/config.py +++ b/prain/config.py @@ -8,8 +8,10 @@ import tomllib from dataclasses import dataclass, field, fields TIERS = ("vram", "ram", "disk") +QUANTS = ("bf16", "int8", "int4") # expert weight formats, see prain.quant +INT4_GROUP = 64 # int4 weights per scale GIB = 1024**3 -_TOP_KEYS = {"pack_dir", "default_tier", "cache_slots", "stats_path", "stats_every", +_TOP_KEYS = {"pack_dir", "quant", "default_tier", "cache_slots", "stats_path", "stats_every", "trace_path", "policy", "prefetch", "miss", "serve", "draft", "layers"} _LAYER_KEYS = {"cache_slots", *TIERS} @@ -104,6 +106,8 @@ class Miss: cpu_max_tokens: int = _opt(1, minimum=1) # auto: CPU runs experts with <= this many tokens cpu_threads: int = _opt(4, minimum=1) # OMP threads of the vLLM worker admit: int = _opt(1) # CPU-run experts copied into VRAM afterwards, hottest first + min_weight: float = _opt(1.0, maximum=1) # router weight every token must compute; below 1 + # it skips experts not in VRAM: faster, dumber @dataclass(frozen=True) @@ -143,6 +147,7 @@ def _section(cls, raw, where): @dataclass(frozen=True) class Config: pack_dir: str + quant: str = "bf16" default_tier: str = "ram" cache_slots: int = 16 stats_path: str | None = None @@ -221,6 +226,7 @@ def load_config(path): layers.append((selector, section)) config = Config( pack_dir=_path(path, raw["pack_dir"]), + quant=raw.get("quant", "bf16"), default_tier=raw.get("default_tier", "ram"), cache_slots=raw.get("cache_slots", 16), stats_path=_path(path, raw.get("stats_path")), @@ -236,6 +242,8 @@ def load_config(path): if config.draft.full_requests > config.serve.max_requests: raise ValueError(f"{path}: [draft] full_requests must be <= [serve] max_requests " f"({config.serve.max_requests})") + if config.quant not in QUANTS: + raise ValueError(f"{path}: quant must be one of {QUANTS}") if config.default_tier not in TIERS: raise ValueError(f"{path}: default_tier must be one of {TIERS}") _check_count(config.cache_slots, 0, f"{path}: cache_slots") @@ -262,9 +270,14 @@ def moe_dims(hf_config): ) -def expert_bytes(hidden, moe_intermediate, dtype_bytes=2): - """One gated expert: gate+up [2I, H] plus down [H, I].""" - return 3 * hidden * moe_intermediate * dtype_bytes +def expert_bytes(hidden, moe_intermediate, quant="bf16"): + """One gated expert, gate+up [2I, H] plus down [H, I], stored as `quant`.""" + weights = 3 * hidden * moe_intermediate + if quant == "bf16": + return 2 * weights + if quant == "int8": # plus one fp32 scale per output row + return weights + 4 * (2 * moe_intermediate + hidden) + return weights // 2 + 2 * weights // INT4_GROUP # plus one bf16 scale per group def render_toml(header, top, tables=None, layers=None): @@ -278,7 +291,8 @@ def render_toml(header, top, tables=None, layers=None): return "\n".join(lines) + "\n" -def plan(stats, num_experts, expert_size, expert_vram_gib, ram_gib, cache_slots, pack_dir, stats_path): +def plan(stats, num_experts, expert_size, expert_vram_gib, ram_gib, cache_slots, pack_dir, stats_path, + quant="bf16"): """Hottest experts (by steps that routed to them) pinned in VRAM, next in RAM, rest on disk.""" uses = {int(layer): s["uses"] for layer, s in stats["layers"].items()} vram_rows = int(expert_vram_gib * GIB // expert_size) - cache_slots * len(uses) @@ -298,7 +312,7 @@ def plan(stats, num_experts, expert_size, expert_vram_gib, ram_gib, cache_slots, } return render_toml( "generated by `prain plan`; experts not listed live on disk", - {"pack_dir": pack_dir, "default_tier": "disk", "cache_slots": cache_slots, + {"pack_dir": pack_dir, "quant": None if quant == "bf16" else quant, "default_tier": "disk", "cache_slots": cache_slots, "stats_path": stats_path}, layers=layers, ) diff --git a/prain/experts.py b/prain/experts.py index bc5eb04..102c242 100644 --- a/prain/experts.py +++ b/prain/experts.py @@ -2,14 +2,17 @@ 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: +checkpoint. With quant = "int8" or "int4" the first forward quantizes that bf16 pack once into +a second pack (prain.quant), and the VRAM slots, copies and CPU path use the quantized +experts. 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, + 3. copies the other missing experts into slots chosen by the [policy] config and runs + vLLM's Triton fused-MoE kernel on the slot tensors, in rounds if needed (each round + computes only its own routes), 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. @@ -26,11 +29,14 @@ 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.fused_moe import fused_experts 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.quant import (GROUP, cpu_matmul, layout, moe_quant_config, quantize, to_cpu_int4, + to_gpu_int4) from prain.sim import write_trace from prain.tiered import draft_rows, tiered_forward @@ -133,9 +139,16 @@ class TieredRoutedExperts(RoutedExperts): 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 + # Quantized slots are prain's own tensors; vLLM's bf16 ones shrink to one row. + quantized = config.quant != "bf16" + moe_config.num_experts = moe_config.num_local_experts = 1 if quantized else placement.num_slots super().__init__(layer_name, params_dtype, moe_config, quant_config, *args, **kwargs) self.local_num_experts = placement.num_slots + self.quant = config.quant + self._shapes = (tuple(self.w13_weight.shape[1:]), tuple(self.w2_weight.shape[1:])) + self._qslots = [[torch.zeros((placement.num_slots, *shape), dtype=dtype, device="cuda") + for shape, dtype in layout(config.quant, *dims)] + for dims in self._shapes] if quantized else None self.prain = config self.prain_layer = layer self.placement = placement @@ -155,27 +168,69 @@ class TieredRoutedExperts(RoutedExperts): 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", + placement.cache_slots, "" if self.pack_complete else + ", writing pack" if self._writing else f", quantizing the pack to {self.quant}", ) def _open_pack(self, pack_dir, layer, num_experts, model): + """The loader writes the bf16 pack; a quantized pack is made from it once. Once the + pack of config.quant is complete, the bf16 one is not read (it may be deleted).""" 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") + self._stem = os.path.join(pack_dir, f"layer{layer:03d}") + self._num_experts = num_experts + self._marker = {"model": model, "w13": [num_experts, *self._shapes[0]], + "w2": [num_experts, *self._shapes[1]], "dtype": str(self.w13_weight.dtype)} + self.pack_complete = self._complete(self.quant) + self._bf16 = None if self.pack_complete else self._pack_of("bf16") + self._writing = self._bf16 is not None and not self._complete("bf16") + + def _paths(self, quant): + """[[w13 files], [w2 files]] of a pack format: weights, then scales.""" + infix = "" if quant == "bf16" else f".{quant}" + return [[f"{self._stem}{infix}.{name}{'s' * i}" for i in range(len(layout(quant, *dims)))] + for name, dims in zip(("w13", "w2"), self._shapes)] + + def _pack_of(self, quant): + return [[_file_tensor(path, (self._num_experts, *shape), dtype) + for path, (shape, dtype) in zip(paths, layout(quant, *dims))] + for paths, dims in zip(self._paths(quant), self._shapes)] + + def _marker_of(self, quant): + if quant == "bf16": + return f"{self._stem}.done", self._marker + return f"{self._stem}.{quant}.done", {**self._marker, "quant": quant, "group": GROUP} + + def _complete(self, quant): + path, marker = self._marker_of(quant) + if not os.path.exists(path): + return False + with open(path) as f: + found = json.load(f) + if found != marker: + raise ValueError(f"prain: {path} is for {found}, not {marker}; delete the pack " + "or use another pack_dir") + return True + + def _seal(self, quant): + for path in sum(self._paths(quant), []): # mmap writes -> disk before we trust the pack + fd = os.open(path, os.O_RDONLY) + os.fsync(fd) + os.close(fd) + path, marker = self._marker_of(quant) + with open(path, "w") as f: + json.dump(marker, f) + + def _quantize_pack(self): + """bf16 pack -> config.quant pack, one expert at a time on the GPU (int4's clip search + takes ~20 min for the whole model on the CPU, under a minute on the GPU).""" + out = self._pack_of(self.quant) + for src, dst in zip(self._bf16, out): + for e in range(self._num_experts): + parts = quantize(src[0][e].cuda(), self.quant) + if self.quant == "int4": + parts[0] = to_cpu_int4(parts[0]) + for tensor, part in zip(dst, parts): + tensor[e] = part.cpu() def get_expert_mapping(self, ckpt_gate_proj_name=None, ckpt_down_proj_name=None, ckpt_up_proj_name=None, include_fused=False): @@ -198,8 +253,8 @@ class TieredRoutedExperts(RoutedExperts): 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 + if self._writing: + pack = self._bf16[0][0] if param is self.w13_weight else self._bf16[1][0] super().weight_loader(pack, loaded_weight, weight_name, shard_id, expert_id, return_success=True) return True if return_success else None @@ -208,15 +263,23 @@ class TieredRoutedExperts(RoutedExperts): 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._writing: + self._seal("bf16") + if self.quant != "bf16": + self._quantize_pack() + self._seal(self.quant) + self.pack_complete, self._writing, self._bf16 = True, False, None + self._pack = self._pack_of(self.quant) # [[w13 tensors], [w2 tensors]] per expert row + if self._qslots is None: + self._slots, self._qconfig = [[self.w13_weight.data], [self.w2_weight.data]], None + else: + self._slots = self._qslots + self._qconfig = moe_quant_config(self.quant, self._slots[0][1], self._slots[1][1]) + # vLLM's one-row bf16 w13/w2 are never read: free them (11 MiB per layer) + self.w13_weight.data = self.w13_weight.data.new_empty(0) + self.w2_weight.data = self.w2_weight.data.new_empty(0) 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 + row = sum(t[0].nbytes for mats in self._pack for t in mats) 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: @@ -225,30 +288,35 @@ class TieredRoutedExperts(RoutedExperts): "from ram to disk") ram = list(self.placement.ram) self._ram_row = {e: row for row, e in enumerate(ram)} - # 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] + # ponytail: pinning is slow on the eGPU host (minutes for a 10 GiB tier, even per + # expert); the page cache serves the pack about as fast, so keep the RAM tier small. + self._ram = [[[t[e].pin_memory() for t in mats] for mats in self._pack] 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): 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.""" + """[[w13 tensors], [w2 tensors]] of an expert on the CPU: 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] + return [[t[expert] for t in mats] for mats in self._pack] + return self._ram[row] + + def _copy(self, expert, slot): + for dsts, srcs in zip(self._slots, self._source(expert)): + for i, (dst, src) in enumerate(zip(dsts, srcs)): + if i == 0 and self.quant == "int4": # the pack keeps the CPU kernel's order + src = to_gpu_int4(src.to(dst.device, non_blocking=True)) + dst[slot].copy_(src, non_blocking=True) 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) + self._copy(expert, slot) def _load_async(self, expert, slot): """Copy on the side stream; only RAM-tier (pinned) experts are truly asynchronous.""" @@ -256,23 +324,23 @@ class TieredRoutedExperts(RoutedExperts): 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) + self._copy(expert, slot) event = torch.cuda.Event() event.record(stream) self._ready_events[slot] = event - def _run_slots(self, x, weights, slot_ids, used_slots): + def _run_slots(self, x, weights, ids, table, used_slots): + """vLLM's Triton fused-MoE kernel on the slot tensors; expert_map=table skips the + routes of experts outside this round.""" 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") + out = fused_experts(x, self._slots[0][0], self._slots[1][0], weights, ids, + activation=self.activation, global_num_experts=table.numel(), + expert_map=table, quant_config=self._qconfig) if self._async: event = torch.cuda.Event() event.record(current) @@ -296,7 +364,10 @@ class TieredRoutedExperts(RoutedExperts): self._predictions[d] = list(zip(e.tolist(), p.tolist())) # host sync def _cpu_weights(self, expert): - return self._source(expert) + mats = self._source(expert) + if self.quant == "bf16": + return mats[0][0], mats[1][0] + return tuple(cpu_matmul(self.quant, *parts) for parts in mats) def _after_enqueue(self, cpu, stats): miss, prefetch = self.prain.miss, self.prain.prefetch diff --git a/prain/quant.py b/prain/quant.py new file mode 100644 index 0000000..01016c4 --- /dev/null +++ b/prain/quant.py @@ -0,0 +1,89 @@ +"""Weight-only int8 / int4 quantization of expert weights. Needs torch. + +int8: q int8 [..., N, K] and one fp32 scale per output row: w = q * scale. vLLM's Triton + fused-MoE kernel and torch's _weight_int8pack_mm (CPU) both read this layout. +int4: one bf16 scale per GROUP inputs of a row: w = (nibble - 8) * scale, with the scale's + clip ratio chosen per group for the least squared error. Two layouts of the nibbles: + gpu: [..., N, K/2], two values per byte along K, even k in the low nibble (vLLM's + Triton GPTQ/AWQ MoE kernel); + cpu: [N/32, K, 16] bytes, rows d and d+16 of each 32-row block share a byte (torch's AVX2 + _weight_int4pack_mm_for_cpu); packs store it, slots convert with to_gpu_int4(). +""" + +import torch + +from prain.config import INT4_GROUP as GROUP # divides Gemma 4 26B-A4B's 2816 and 704 + +CLIP = (1.0, 0.95, 0.9, 0.85, 0.8, 0.75, 0.7) + + +def layout(quant, rows, cols): + """[(shape, dtype)] of the tensors that store one [rows, cols] weight matrix.""" + if quant == "bf16": + return [((rows, cols), torch.bfloat16)] + if quant == "int8": + return [((rows, cols), torch.int8), ((rows,), torch.float32)] + return [((rows, cols // 2), torch.uint8), ((rows, cols // GROUP), torch.bfloat16)] + + +def quantize(w, quant): + """Float weights [..., N, K] -> [q, scale], int4 in the gpu layout.""" + w = w.float() + if quant == "int8": + scale = w.abs().amax(-1).clamp_min(1e-12) / 127 + return [(w / scale[..., None]).round().to(torch.int8), scale] + g = w.unflatten(-1, (-1, GROUP)) + amax = g.abs().amax(-1, keepdim=True).clamp_min(1e-12) + best_err = best_scale = None + for ratio in CLIP: + scale = (amax * ratio / 7).bfloat16().float() # the kernels see the bf16 scale + err = ((g / scale).round().clamp(-8, 7) * scale - g).pow(2).sum(-1, keepdim=True) + if best_err is None: + best_err, best_scale = err, scale + else: + better = err < best_err + best_err, best_scale = torch.where(better, err, best_err), torch.where(better, scale, best_scale) + q = ((g / best_scale).round().clamp(-8, 7) + 8).flatten(-2).to(torch.uint8) + return [q[..., 0::2] | q[..., 1::2] << 4, best_scale.squeeze(-1).bfloat16()] + + +def dequantize(q, scale, quant, dtype=torch.bfloat16): + """The weights [..., N, K] of quantize()'s output (int4 in the gpu layout).""" + if quant == "int8": + return q.to(dtype) * scale.to(dtype)[..., None] + v = torch.stack((q & 15, q >> 4), -1).flatten(-2).to(dtype) - 8 + return (v.unflatten(-1, (-1, GROUP)) * scale.to(dtype)[..., None]).flatten(-2) + + +def to_cpu_int4(q): + """int4 nibbles [..., N, K/2]: gpu layout -> cpu layout (same shape).""" + *lead, n, half = q.shape + v = torch.stack((q & 15, q >> 4), -1).reshape(*lead, n // 32, 32, 2 * half).transpose(-1, -2) + return (v[..., :16] | v[..., 16:] << 4).reshape(q.shape) + + +def to_gpu_int4(c): + """int4 nibbles [..., N, K/2]: cpu layout -> gpu layout (same shape; works on the GPU).""" + *lead, n, half = c.shape + b = c.reshape(*lead, n // 32, 2 * half, 16) + v = torch.cat((b & 15, b >> 4), -1).transpose(-1, -2).reshape(*lead, n, 2 * half) + return v[..., 0::2] | v[..., 1::2] << 4 + + +def cpu_matmul(quant, q, scale): + """x -> x @ w.T on the CPU for stored int8 q / cpu-layout int4 q; x is bf16 [M, K].""" + if quant == "int8": + return lambda x: torch._weight_int8pack_mm(x, q, scale.to(x.dtype)) + zeros = torch.stack((scale.t(), torch.zeros_like(scale.t())), -1).contiguous() + return lambda x: torch._weight_int4pack_mm_for_cpu(x, q, GROUP, zeros) + + +def moe_quant_config(quant, w13_scale, w2_scale): + """vLLM's fused-MoE quant config for int8 / gpu-layout int4 slot tensors.""" + from vllm.model_executor.layers.fused_moe import config + + if quant == "int8": + return config.int8_w8a16_moe_quant_config(w1_scale=w13_scale, w2_scale=w2_scale, + w1_zp=None, w2_zp=None) + return config.int4_w4a16_moe_quant_config(w1_scale=w13_scale, w2_scale=w2_scale, w1_zp=None, + w2_zp=None, block_shape=[0, GROUP]) diff --git a/prain/tiered.py b/prain/tiered.py index 1925c8b..3386d40 100644 --- a/prain/tiered.py +++ b/prain/tiered.py @@ -49,7 +49,7 @@ def degrade(topk_weights, topk_ids, draft, available, min_weight): A draft token keeps its heaviest unavailable routes only while less than min_weight of its router weight would be computed; its kept routes are scaled up to its full router weight. - Dropped routes get id -1 and weight 0, which route_stats and remap ignore. + Dropped routes get id -1 and weight 0, which route_stats and slot_table ignore. draft: bool [tokens]; available: bool [num_experts]. Full tokens are left unchanged.""" ids = topk_ids.long() valid = ids >= 0 @@ -68,37 +68,36 @@ def degrade(topk_weights, topk_ids, draft, available, min_weight): return weights, torch.where(drop, torch.full_like(topk_ids, -1), topk_ids) -def remap(topk_weights, topk_ids, slot_of, num_experts): - """Point routes at VRAM slots for one round. - - Routes to experts outside this round (and padding routes with id -1) get slot 0 - and weight 0, so any fused-MoE kernel can run them unchanged.""" - table = torch.full((num_experts,), -1, dtype=torch.int64) +def slot_table(slot_of, num_experts, device): + """int32 [num_experts + 1]: VRAM slot of each expert in one round, -1 for the others. + vLLM's fused-MoE kernel takes it as expert_map and skips routes that map to -1; the extra + last row stands for the padding id (route id -1 arrives as num_experts).""" + table = torch.full((num_experts + 1,), -1, dtype=torch.int32) if slot_of: - table[torch.tensor(list(slot_of), dtype=torch.int64)] = torch.tensor( - list(slot_of.values()), dtype=torch.int64 - ) - table = table.to(topk_ids.device) - ids = topk_ids.long() - slots = torch.where(ids >= 0, table[ids.clamp_min(0)], -1) - weights = torch.where(slots >= 0, topk_weights, torch.zeros_like(topk_weights)) - return weights, slots.clamp_min(0).to(topk_ids.dtype) + table[torch.tensor(list(slot_of))] = torch.tensor(list(slot_of.values()), dtype=torch.int32) + return table.to(device) + + +def _matmul(w): + """x -> x @ w.T for a weight tensor; functions (quantized experts) pass through.""" + return w if callable(w) else (lambda v: v.to(w.dtype) @ w.t()) def cpu_experts(x, topk_weights, topk_ids, experts, weights_of): """The routes to `experts`, computed on the CPU like vLLM's fused kernel with activation "gelu_tanh": w13 = [gate; up], y = (gelu_tanh(x @ gate.T) * (x @ up.T)) @ w2.T times the - route weight. weights_of(e) -> (w13 [2I, H], w2 [H, I]). Returns [tokens, hidden] fp32.""" + route weight. weights_of(e) -> (w13 [2I, H], w2 [H, I]), or two functions v -> v @ w.T + (quantized experts). Returns [tokens, hidden] fp32.""" out = torch.zeros(x.shape, dtype=torch.float32) for e in sorted(experts): rows, cols = (topk_ids == e).nonzero(as_tuple=True) if rows.numel() == 0: continue - w13, w2 = weights_of(e) - h = x[rows].to(w13.dtype) @ w13.t() - half = w13.shape[0] // 2 + mm13, mm2 = map(_matmul, weights_of(e)) + h = mm13(x[rows]) + half = h.shape[1] // 2 act = F.gelu(h[:, :half].float(), approximate="tanh") * h[:, half:].float() - y = (act.to(w2.dtype) @ w2.t()).float() + y = mm2(act.to(h.dtype)).float() out.index_add_(0, rows, y * topk_weights[rows, cols].float()[:, None]) return out @@ -108,20 +107,32 @@ def tiered_forward(x, topk_weights, topk_ids, cache, load, kernel, heat="weight" after_enqueue=None, draft=None, min_weight=0.5): """One MoE forward over VRAM slots. - cache: SlotCache. load(expert, slot) copies an expert into a slot. kernel(x, weights, - slot_ids, used_slots) runs the stock fused-MoE kernel on the slot tensors. With a + cache: SlotCache. load(expert, slot) copies an expert into a slot. kernel(x, weights, ids, + table, used_slots) runs the fused-MoE kernel on the slot tensors for the routes whose + expert the slot_table() `table` maps to a slot. With a [miss] config other than "copy", the experts it picks run on the CPU (cpu_weights(e) gives their weights) while the GPU rounds run. on_stats(stats, tokens) sees the routing statistics; after_enqueue(cpu, stats) runs once all GPU work of this forward is queued. - draft: bool [tokens] for the rows of draft requests, or None; see degrade(). + draft: bool [tokens] for the rows of draft requests, or None; see degrade(). Draft rows keep + min_weight of their router weight, every row keeps [miss] min_weight (1 = everything). """ raw = None - if draft is not None: - if on_stats is not None: # the trace records the routing, not what drafts skipped + skip = miss is not None and miss.min_weight < 1 + if draft is not None or skip: + if on_stats is not None: # the trace records the routing, not what was skipped raw = route_stats(topk_weights, topk_ids, cache.num_experts, recency_half_life) - available = available_experts(topk_ids, draft, set(cache.pinned) | set(cache.slots), + rows = draft if draft is not None else torch.zeros(topk_ids.shape[0], dtype=torch.bool, + device=topk_ids.device) + need = min_weight + if skip: # every row may skip; draft rows by the smaller share + need = torch.where(rows, min(min_weight, miss.min_weight), miss.min_weight)[:, None] + rows = torch.ones_like(rows) + available = available_experts(topk_ids, rows, set(cache.pinned) | set(cache.slots), cache.num_experts) - topk_weights, topk_ids = degrade(topk_weights, topk_ids, draft, available, min_weight) + _, kept = degrade(topk_weights, topk_ids, rows, available, need) + # experts some row copies anyway serve every row for free + available |= torch.bincount(kept[kept >= 0].long(), minlength=cache.num_experts) > 0 + topk_weights, topk_ids = degrade(topk_weights, topk_ids, rows, available, need) stats = route_stats(topk_weights, topk_ids, cache.num_experts, recency_half_life) # host sync tokens = int(topk_ids.shape[0]) if on_stats is not None: @@ -131,22 +142,22 @@ def tiered_forward(x, topk_weights, topk_ids, cache, load, kernel, heat="weight" if cpu: # copy the inputs now: later the copy would wait for the queued GPU rounds host = x.cpu(), topk_weights.cpu(), topk_ids.cpu() index = HEAT_INDEX[heat] - outs = [] + ids = torch.where(topk_ids >= 0, topk_ids, cache.num_experts) # padding -> the table's last row + acc = torch.promote_types(x.dtype, torch.float32) + total = None # sum of the round (and CPU) outputs: one buffer, not one per round + + def add(out): + nonlocal total + total = out if total is None else total.to(acc).add_(out) # .to: a no-op once fp32 + for slot_of, loads in cache.plan({e: s[index] for e, s in stats.items()}, tokens, cpu): for expert, slot in loads: load(expert, slot) - if slot_of: - weights, slots = remap(topk_weights, topk_ids, slot_of, cache.num_experts) - outs.append(kernel(x, weights, slots, list(slot_of.values()))) + if slot_of: # each round computes only its own routes + table = slot_table(slot_of, cache.num_experts, x.device) + add(kernel(x, topk_weights, ids, table, list(slot_of.values()))) if after_enqueue is not None: after_enqueue(cpu, stats) if cpu: # runs while the GPU works through its rounds - outs.append(cpu_experts(*host, cpu, cpu_weights).to(x.device)) - if not outs: - return torch.zeros_like(x) - if len(outs) == 1: - return outs[0].to(x.dtype) - # ponytail: every extra round re-runs all tokens with weight 0 for experts it does not - # hold; only big batches need many rounds. Upgrade: compact tokens per round. - acc = torch.promote_types(x.dtype, torch.float32) - return torch.stack([o.to(acc) for o in outs]).sum(0).to(x.dtype) + add(cpu_experts(*host, cpu, cpu_weights).to(x.device)) + return torch.zeros_like(x) if total is None else total.to(x.dtype) diff --git a/scripts/bench.py b/scripts/bench.py index fc556f6..f60f247 100644 --- a/scripts/bench.py +++ b/scripts/bench.py @@ -3,7 +3,9 @@ usage (prain installed): python3 scripts/bench.py http://127.0.0.1:8000 [--tokens 128] [--runs 3] Each run sends a fresh prompt (a random prefix defeats the prefix cache) and decodes exactly --tokens tokens (ignore_eos), so numbers compare across runs and configs. The last line holds -the medians. +the medians. --context FILE... puts the files' text (up to --context-chars) before every +prompt and the random tag after it: run 1 prefills the whole context, later runs find it in +the prefix cache, like an agent's system prompt and tools. """ import argparse @@ -21,8 +23,8 @@ PROMPTS = [ ] -def run(base_url, model, prompt, tokens): - text = f"[{uuid.uuid4().hex[:8]}] {prompt}" +def run(base_url, model, prompt, tokens, context=""): + text = f"{context}[{uuid.uuid4().hex[:8]}] {prompt}" body = chat_body(model, [{"role": "user", "content": text}], max_tokens=tokens, temperature=0, ignore_eos=True) start = time.monotonic() @@ -39,11 +41,18 @@ def main(): parser.add_argument("base_url") parser.add_argument("--tokens", type=int, default=128) parser.add_argument("--runs", type=int, default=3) + parser.add_argument("--context", nargs="*", default=[], help="files to put before each prompt") + parser.add_argument("--context-chars", type=int, default=48000) args = parser.parse_args() model = first_model(args.base_url) + context = "" + for path in args.context: + with open(path) as f: + context += f.read() + context = context[: args.context_chars] + "\n\n" if context else "" results = [] for i in range(args.runs): - results.append(run(args.base_url, model, PROMPTS[i % len(PROMPTS)], args.tokens)) + results.append(run(args.base_url, model, PROMPTS[i % len(PROMPTS)], args.tokens, context)) print(json.dumps(results[-1]), flush=True) print(json.dumps({key: statistics.median(r[key] for r in results) for key in results[0]} | {"median_of": len(results), "tokens": args.tokens})) diff --git a/scripts/env.sh b/scripts/env.sh index 159d327..5f71096 100644 --- a/scripts/env.sh +++ b/scripts/env.sh @@ -23,4 +23,7 @@ export PATH # libc++/compiler-rt defaults: everything links libstdc++/libgcc like the ROCm libraries. export CC=$PRAIN_ROOT/build/bin/clang CXX=$PRAIN_ROOT/build/bin/clang++ export CFLAGS="-march=native" CXXFLAGS="-march=native" +# Fused-MoE kernel configs tuned for this GPU by scripts/tune_moe.py (vLLM's defaults when +# the files are missing). +export VLLM_TUNED_CONFIG_FOLDER=$PRAIN_ROOT/build/moe-configs unset _site _dev _core diff --git a/scripts/measure.py b/scripts/measure.py index 197bc39..ab5efbe 100644 --- a/scripts/measure.py +++ b/scripts/measure.py @@ -1,7 +1,7 @@ """Measure the prain cost model on this machine and print it as JSON for `prain sim --cost` and `prain tune --cost`. Needs the GPU, vLLM and prain. -usage: python3 scripts/measure.py PACK_FILE [--threads N] [--dense-gib 4.45] [--layers 30] > cost.json +usage: python3 scripts/measure.py PACK_FILE [--quant int4] [--threads N] [--dense-gib 4.45] [--layers 30] > cost.json PACK_FILE is any /layerNNN.w13 file; it is read cold (dropped from the page cache). """ @@ -15,6 +15,8 @@ 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.config import QUANTS, expert_bytes as size_of +from prain.quant import cpu_matmul, moe_quant_config, quantize, to_cpu_int4 from prain.tiered import cpu_experts HIDDEN, INTER, SLOTS = 2816, 704, 5 # Gemma 4 26B-A4B expert, 5 slots per layer @@ -37,10 +39,11 @@ def main(): parser.add_argument("--threads", type=int, default=4) parser.add_argument("--dense-gib", type=float, default=4.45, help="non-expert weights") parser.add_argument("--layers", type=int, default=30) + parser.add_argument("--quant", choices=QUANTS, default="bf16", help="expert weight format") args = parser.parse_args() torch.set_num_threads(args.threads) - expert_bytes = 3 * HIDDEN * INTER * 2 - host = torch.empty(expert_bytes // 2, dtype=torch.bfloat16).pin_memory() + expert_bytes = size_of(HIDDEN, INTER, args.quant) + host = torch.empty(expert_bytes, dtype=torch.uint8).pin_memory() device = torch.empty_like(host, device="cuda") h2d_gbps = expert_bytes / 1e6 / ms(lambda: device.copy_(host, non_blocking=True)) @@ -56,13 +59,17 @@ def main(): w1 = torch.randn(SLOTS, 2 * INTER, HIDDEN, dtype=torch.bfloat16, device="cuda") * 0.02 w2 = torch.randn(SLOTS, HIDDEN, INTER, dtype=torch.bfloat16, device="cuda") * 0.02 + quant_config = None + if args.quant != "bf16": + (w1, s1), (w2, s2) = quantize(w1, args.quant), quantize(w2, args.quant) + quant_config = moe_quant_config(args.quant, s1, s2) def moe(tokens): x = torch.randn(tokens, HIDDEN, dtype=torch.bfloat16, device="cuda") ids = torch.arange(SLOTS, device="cuda", dtype=torch.int32).repeat(tokens, 1) weights = torch.full((tokens, SLOTS), 1 / SLOTS, device="cuda") return lambda: fused_experts(x, w1, w2, weights, ids, activation=MoEActivation.GELU_TANH, - global_num_experts=SLOTS) + global_num_experts=SLOTS, quant_config=quant_config) with set_current_vllm_config(VllmConfig()): round_ms = ms(moe(1)) @@ -76,6 +83,12 @@ def main(): # 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)] + if args.quant != "bf16": # the CPU kernels on pack-layout weights, as the plugin runs them + def stored(w): + q, scale = quantize(w, args.quant) + return cpu_matmul(args.quant, to_cpu_int4(q) if args.quant == "int4" else q, scale) + + experts = [(stored(a), stored(b)) for a, b in experts] x = torch.randn(1, HIDDEN, dtype=torch.bfloat16) one = torch.ones(1, 1) ids = torch.zeros(1, 1, dtype=torch.int32) diff --git a/scripts/quality.py b/scripts/quality.py new file mode 100644 index 0000000..bec166a --- /dev/null +++ b/scripts/quality.py @@ -0,0 +1,89 @@ +"""Answer quality of a served config against reference answers from another config. + + python scripts/quality.py URL make REF.json greedy reference answers (run on the exact config) + python scripts/quality.py URL score REF.json how likely this server finds those answers + +Both work on token ids: the prompt is the chat template with its generation prompt, and `score` +re-scores the reference answer on exactly the ids it was generated from (prompt logprobs). A +finished assistant message renders differently from the generation prompt, so re-scoring +through the chat API skews the logprobs. `nll` is the mean negative log-likelihood per answer +token (the reference config scores about its own decode logprobs); `same_tokens` is how many +leading tokens this server's own greedy answer shares with the reference. +""" + +import argparse +import json +import urllib.request + +PROMPTS = [ + "Write a Python function that returns the n-th Fibonacci number iteratively.", + "Explain what `git rebase -i HEAD~3` does, in three sentences.", + "Write a bash one-liner that counts the lines of all .py files below the current directory.", + "On Gentoo, how do I rebuild every package that links against a library I just updated?", + "What is the difference between a process and a thread? Answer in two sentences.", +] +CODE_PROMPT = "Here is some Python code:\n\n{code}\n\nWhat does this code do? Answer in two sentences." + + +def post(url, path, body): + request = urllib.request.Request(url + path, json.dumps(body).encode(), + {"Content-Type": "application/json"}) + with urllib.request.urlopen(request, timeout=3600) as response: + return json.load(response) + + +def greedy(url, model, ids, tokens): + choice = post(url, "/v1/completions", { + "model": model, "prompt": ids, "max_tokens": tokens, "temperature": 0, "logprobs": 1, + "return_tokens_as_token_ids": True})["choices"][0] + return choice["text"], [int(t.split(":")[1]) for t in choice["logprobs"]["tokens"]] + + +def nll(url, model, ids, answer): + """Mean negative log-likelihood of the answer tokens after the prompt ids.""" + entries = post(url, "/v1/completions", { + "model": model, "prompt": ids + answer, "max_tokens": 1, "temperature": 0, + "prompt_logprobs": 0})["choices"][0]["prompt_logprobs"] + return -sum(entries[len(ids) + i][str(t)]["logprob"] for i, t in enumerate(answer)) / len(answer) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("url") + parser.add_argument("action", choices=("make", "score")) + parser.add_argument("ref") + parser.add_argument("--model", default="gemma4") + parser.add_argument("--tokens", type=int, default=48, help="answer tokens") + parser.add_argument("--code", default="prain/policy.py", help="file for the code question") + parser.add_argument("--code-chars", type=int, default=900) + args = parser.parse_args() + if args.action == "make": + with open(args.code) as f: + prompts = PROMPTS + [CODE_PROMPT.format(code=f.read()[: args.code_chars])] + refs = [] + for prompt in prompts: + ids = post(args.url, "/tokenize", {"model": args.model, "add_generation_prompt": True, + "messages": [{"role": "user", "content": prompt}]}) + ids = ids["tokens"] + text, answer = greedy(args.url, args.model, ids, args.tokens) + refs.append({"prompt": prompt, "ids": ids, "answer": answer, "text": text}) + print(json.dumps({"prompt": prompt[:60], "answer": text[:120]}), flush=True) + with open(args.ref, "w") as f: + json.dump(refs, f) + return + with open(args.ref) as f: + refs = json.load(f) + scores, same = [], [] + for ref in refs: + scores.append(nll(args.url, args.model, ref["ids"], ref["answer"])) + text, answer = greedy(args.url, args.model, ref["ids"], len(ref["answer"])) + same.append(next((i for i, (a, b) in enumerate(zip(answer, ref["answer"])) if a != b), + min(len(answer), len(ref["answer"])))) + print(json.dumps({"prompt": ref["prompt"][:40], "nll": round(scores[-1], 3), + "same_tokens": same[-1], "answer": text[:100]}), flush=True) + print(json.dumps({"mean_nll": round(sum(scores) / len(scores), 3), "same_tokens": same, + "identical": sum(s == len(r["answer"]) for s, r in zip(same, refs))})) + + +if __name__ == "__main__": + main() diff --git a/scripts/tune_moe.py b/scripts/tune_moe.py new file mode 100644 index 0000000..b3ebb65 --- /dev/null +++ b/scripts/tune_moe.py @@ -0,0 +1,111 @@ +"""Tune vLLM's Triton fused-MoE kernel for this GPU on prain's slot tensors. + +usage: python scripts/tune_moe.py OUT_DIR [FORMAT...] (scripts/env.sh points + VLLM_TUNED_CONFIG_FOLDER at build/moe-configs) +Times candidate configs on one prain round (16 slots of 128 routed experts, expert_map; +up to 32 tokens route into the slots like decode, more tokens spread like prefill), +Gemma 4 expert shapes, bf16/int8/int4 slots, and writes the fastest per batch size as vLLM +config files for every slot count E (the best config hardly depends on E). Every config in +the files compiled, ran and matched the default config's output on this GPU. +""" + +import itertools +import json +import os +import sys +import time + +import torch +from vllm.config import VllmConfig, set_current_vllm_config +from vllm.model_executor.layers.fused_moe import override_config +from vllm.model_executor.layers.fused_moe.activation import MoEActivation +from vllm.model_executor.layers.fused_moe.config import (int4_w4a16_moe_quant_config, + int8_w8a16_moe_quant_config) +from vllm.model_executor.layers.fused_moe.fused_moe import (fused_experts, get_config_file_name, + get_default_config) + +from prain.quant import GROUP, quantize + +H, I, TOPK, EXPERTS, SLOTS = 2816, 704, 8, 128, 16 +SIZES = (1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048) +FORMATS = {"bf16": (None, None), "int8": ("int8_w8a16", None), "int4": ("int4_w4a16", [0, GROUP])} + + +def candidates(tokens): + # ponytail: ~8 s per Triton compile on this host, so a small space (~70 compiles per + # format); widen it when a GPU compiles faster. + if tokens <= 32: + space = itertools.product((16,), (32, 64, 128), (64, 128, 256), (1,), (2, 4), (2,)) + else: + space = itertools.product((32, 64), (64, 128), (32, 64), (8,), (4, 8), (2,)) + for bm, bn, bk, group, warps, stages in space: + yield {"BLOCK_SIZE_M": bm, "BLOCK_SIZE_N": bn, "BLOCK_SIZE_K": bk, "GROUP_SIZE_M": group, + "SPLIT_K": 1, "num_warps": warps, "num_stages": stages} # int4 needs SPLIT_K + + +def timed(fn, repeat): + fn() + torch.cuda.synchronize() + start = time.perf_counter() + for _ in range(repeat): + fn() + torch.cuda.synchronize() + return (time.perf_counter() - start) * 1000 / repeat + + +def main(out_dir, formats): + os.makedirs(out_dir, exist_ok=True) + torch.manual_seed(0) + w13 = torch.randn(SLOTS, 2 * I, H, device="cuda").mul_(0.02).bfloat16() + w2 = torch.randn(SLOTS, H, I, device="cuda").mul_(0.02).bfloat16() + table = torch.full((EXPERTS + 1,), -1, dtype=torch.int32, device="cuda") + table[:SLOTS] = torch.arange(SLOTS, dtype=torch.int32, device="cuda") + for name in formats: + dtype, block_shape = FORMATS[name] + if name == "bf16": + a, b, quant_config = w13, w2, None + else: + (a, s13), (b, s2) = quantize(w13, name), quantize(w2, name) + make = int8_w8a16_moe_quant_config if name == "int8" else int4_w4a16_moe_quant_config + quant_config = make(w1_scale=s13, w2_scale=s2, w1_zp=None, w2_zp=None, + block_shape=block_shape) + best = {} + for tokens in SIZES: + x = torch.randn(tokens, H, device="cuda").mul_(0.5).bfloat16() + pool = SLOTS if tokens <= 32 else EXPERTS + ids = torch.stack([torch.randperm(pool, device="cuda")[:TOPK] + for _ in range(tokens)]).to(torch.int32) + weights = torch.softmax(torch.randn(tokens, TOPK, device="cuda"), -1) + + def run(): + return fused_experts(x, a, b, weights, ids, activation=MoEActivation.GELU_TANH, + global_num_experts=EXPERTS + 1, expert_map=table, + quant_config=quant_config) + + default = get_default_config(tokens, SLOTS, I, a.shape[2], TOPK, dtype, block_shape) + with override_config(default): # vLLM's own choice when no config file exists + default_ms = timed(run, 10 if tokens <= 32 else 3) + reference = run() + results = [(default_ms, json.dumps(default, sort_keys=True))] # never worse than it + for config in candidates(tokens): + # catch inside: vLLM's override_config skips its restore when the body raises + with override_config(config): + try: + ms = timed(run, 10 if tokens <= 32 else 3) + out = run() + except Exception: # noqa: BLE001 - configs that do not compile or fit + continue + if (out.float() - reference.float()).norm() <= 1e-2 * reference.float().norm() + 1e-6: + results.append((ms, json.dumps(config, sort_keys=True))) + ms, config = min(results) + best[str(tokens)] = json.loads(config) + print(f"{name} tokens={tokens:4d}: default {default_ms:8.3f} ms, best {ms:8.3f} ms " + f"({len(results)} configs) {config}", flush=True) + for e in range(1, EXPERTS + 1): # any slot count a placement can give a layer + with open(os.path.join(out_dir, get_config_file_name(e, I, dtype, block_shape)), "w") as f: + json.dump(best, f, indent=1) + + +if __name__ == "__main__": + with set_current_vllm_config(VllmConfig()): + main(sys.argv[1], sys.argv[2:] or list(FORMATS)) diff --git a/tests/test_cli.py b/tests/test_cli.py index 4acc410..aa2bec7 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -76,3 +76,17 @@ class Cli(TmpCase): self.assertEqual((tuned.pack_dir, tuned.trace_path), ("/packs", trace)) self.assertLessEqual(sum(tuned.placement(layer, 128).cache_slots for layer in range(3)), int(0.2 * 1024 / 11.34) // 3 * 3) + + def test_quantized_experts_check_and_tune(self): + hf = self.write("config.json", json.dumps(GEMMA4)) + out, _ = self.run_cli("check", self.write("q.toml", 'pack_dir = "/packs"\nquant = "int4"\n'), hf) + self.assertIn("expert size 3.01 MiB\n", out) + trace = os.path.join(self.dir, "trace.jsonl.gz") + self.run_cli("synth", trace, "--forwards", "40", "--layers", "3") + out, _ = self.run_cli("tune", trace, hf, "--quant", "int4", "--expert-vram-gib", "0.2", + "--ram-gib", "4", "--pack-dir", "/packs", "--top", "3") + tuned = load_config(self.write("tuned.toml", out)) + self.assertEqual(tuned.quant, "int4") + slots = sum(tuned.placement(layer, 128).cache_slots for layer in range(3)) + self.assertGreater(slots, int(0.2 * 1024 / 11.34)) # int4 experts are ~4x smaller + self.assertLessEqual(slots, int(0.2 * 1024 / 3.01)) diff --git a/tests/test_config.py b/tests/test_config.py index 0a23acc..c3546d6 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -107,6 +107,23 @@ mode = "auto" self.assertEqual(config.stats_path, os.path.expanduser("~/s.json")) self.assertEqual(config.trace_path, "/abs/t.gz") + def test_miss_min_weight(self): + config = mc.load_config(self.write("m.toml", 'pack_dir = "/p"\n[miss]\nmin_weight = 0.7\n')) + self.assertEqual(config.miss.min_weight, 0.7) + self.assertEqual(mc.Miss().min_weight, 1.0) # default: compute every routed expert + with self.assertRaisesRegex(ValueError, "min_weight must be a number in"): + mc.load_config(self.write("m2.toml", 'pack_dir = "/p"\n[miss]\nmin_weight = 1.5\n')) + + def test_quant_key_and_expert_sizes(self): + self.assertEqual(mc.load_config(self.write("q.toml", 'pack_dir = "/p"\nquant = "int4"\n')).quant, "int4") + self.assertEqual(mc.load_config(self.write("b.toml", 'pack_dir = "/p"\n')).quant, "bf16") + with self.assertRaisesRegex(ValueError, "quant must be one of"): + mc.load_config(self.write("bad.toml", 'pack_dir = "/p"\nquant = "int3"\n')) + h, i = 2816, 704 + self.assertEqual(mc.expert_bytes(h, i), 3 * h * i * 2) + self.assertEqual(mc.expert_bytes(h, i, "int8"), 3 * h * i + 4 * (2 * i + h)) # fp32 row scales + self.assertEqual(mc.expert_bytes(h, i, "int4"), 3 * h * i // 2 + 2 * 3 * h * i // 64) + def test_serve_and_draft_tables(self): config = mc.load_config(self.write("s.toml", 'pack_dir = "/p"\n[serve]\nmax_requests = 4\n' 'kv_cache_gib = 1.5\n[draft]\nfull_requests = 2\n' diff --git a/tests/test_kernel_match.py b/tests/test_kernel_match.py index e74099c..0e4fadb 100644 --- a/tests/test_kernel_match.py +++ b/tests/test_kernel_match.py @@ -31,3 +31,64 @@ class KernelMatch(unittest.TestCase): 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) + + def test_a_round_computes_only_the_experts_its_slot_table_maps(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, slot_table + + torch.manual_seed(0) + tokens, hidden, inter, experts = 64, 256, 128, 6 + 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_ids[0, 1] = experts # a padding route + topk_weights = torch.rand(tokens, 2) + table = slot_table({1: 0, 4: 1}, experts, "cuda") + with set_current_vllm_config(VllmConfig()): + gpu = fused_experts(x.cuda(), w13[[1, 4]].cuda(), w2[[1, 4]].cuda(), topk_weights.cuda(), + topk_ids.cuda(), activation=MoEActivation.GELU_TANH, + global_num_experts=table.numel(), expert_map=table).float().cpu() + cpu = cpu_experts(x, topk_weights, topk_ids, {1, 4}, lambda e: (w13[e], w2[e])) + torch.testing.assert_close(cpu, gpu, rtol=3e-2, atol=3e-2) + + def test_quantized_slots_compute_what_the_cpu_path_computes(self): + """Pack layout -> slot layout on the GPU -> vLLM's int8/int4 kernel for a round that + holds experts 1 and 3, against the CPU kernels on the pack layout and against the + dequantized weights.""" + 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.quant import (cpu_matmul, dequantize, moe_quant_config, quantize, to_cpu_int4, + to_gpu_int4) + from prain.tiered import cpu_experts, slot_table + + 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) * 0.1 + w2 = torch.randn(experts, hidden, inter) * 0.1 + topk_ids = torch.stack([torch.randperm(experts)[:2] for _ in range(tokens)]).int() + topk_weights = torch.rand(tokens, 2) + table = slot_table({1: 0, 3: 1}, experts, "cuda") + for quant in ("int8", "int4"): + parts = [quantize(w, quant) for w in (w13, w2)] # [q, scale], gpu layout + pack = [[to_cpu_int4(q) if quant == "int4" else q, s] for q, s in parts] + slots = [[to_gpu_int4(q[[1, 3]].cuda()) if quant == "int4" else q[[1, 3]].cuda(), + s[[1, 3]].cuda()] for q, s in pack] + with set_current_vllm_config(VllmConfig()): + gpu = fused_experts(x.cuda(), slots[0][0], slots[1][0], topk_weights.cuda(), + topk_ids.cuda(), activation=MoEActivation.GELU_TANH, + global_num_experts=table.numel(), expert_map=table, + quant_config=moe_quant_config(quant, slots[0][1], slots[1][1])) + cpu = cpu_experts(x, topk_weights, topk_ids, {1, 3}, lambda e: tuple( + cpu_matmul(quant, q[e], s[e]) for q, s in pack)) + exact = cpu_experts(x, topk_weights, topk_ids, {1, 3}, lambda e: tuple( + dequantize(q, s, quant)[e] for q, s in parts)) + self.assertGreater(exact.norm().item(), 1.0) + torch.testing.assert_close(cpu, exact, rtol=3e-2, atol=3e-2, msg=quant) + torch.testing.assert_close(gpu.float().cpu(), exact, rtol=3e-2, atol=3e-2, msg=quant) diff --git a/tests/test_quant.py b/tests/test_quant.py new file mode 100644 index 0000000..033d961 --- /dev/null +++ b/tests/test_quant.py @@ -0,0 +1,49 @@ +import unittest + +try: + import torch +except ImportError: # the planning host has no torch + torch = None + +if torch is not None: + from prain.quant import (GROUP, cpu_matmul, dequantize, layout, quantize, to_cpu_int4, + to_gpu_int4) + + +@unittest.skipIf(torch is None, "needs torch") +class Quant(unittest.TestCase): + def setUp(self): + torch.manual_seed(0) + self.w = torch.randn(64, 256) * 0.02 + + def test_layouts_match_what_quantize_returns(self): + for quant in ("int8", "int4"): + got = [(tuple(t.shape), t.dtype) for t in quantize(self.w, quant)] + self.assertEqual(got, layout(quant, 64, 256)) + + def test_round_trip_error(self): + for quant, bound in (("int8", 0.01), ("int4", 0.11)): + err = (dequantize(*quantize(self.w, quant), quant, torch.float32) - self.w).norm() + self.assertLess(err / self.w.norm(), bound) + + def test_int4_clipping_beats_plain_round_to_nearest(self): + g = self.w.unflatten(-1, (-1, GROUP)) + scale = (g.abs().amax(-1, keepdim=True) / 7).bfloat16().float() + plain = ((g / scale).round() * scale - g).norm() + clipped = (dequantize(*quantize(self.w, "int4"), "int4", torch.float32) - self.w).norm() + self.assertLess(clipped, plain) + + def test_cpu_layout_is_torch_avx2_int4pack_and_inverts(self): + q, _ = quantize(self.w, "int4") + nibbles = torch.stack((q & 15, q >> 4), -1).flatten(-2).to(torch.int32) + self.assertTrue(torch.equal(to_cpu_int4(q), torch._convert_weight_to_int4pack_for_cpu(nibbles, 1))) + self.assertTrue(torch.equal(to_gpu_int4(to_cpu_int4(q)), q)) + + def test_cpu_matmul_matches_dequantized_weights(self): + x = torch.randn(3, 256, dtype=torch.bfloat16) + for quant in ("int8", "int4"): + q, scale = quantize(self.w, quant) + ref = x.float() @ dequantize(q, scale, quant, torch.float32).t() + stored = to_cpu_int4(q) if quant == "int4" else q + got = cpu_matmul(quant, stored, scale)(x).float() + self.assertLess(((got - ref).norm() / ref.norm()).item(), 0.01, quant) diff --git a/tests/test_tiered.py b/tests/test_tiered.py index cafb16b..02c9081 100644 --- a/tests/test_tiered.py +++ b/tests/test_tiered.py @@ -9,7 +9,7 @@ from prain.config import Miss, Policy from prain.policy import SlotCache if torch is not None: - from prain.tiered import (available_experts, cpu_experts, degrade, remap, route_stats, + from prain.tiered import (available_experts, cpu_experts, degrade, route_stats, slot_table, tiered_forward) HIDDEN, INTER, EXPERTS = 6, 4, 10 @@ -52,8 +52,9 @@ class TieredForward(unittest.TestCase): loads.append(expert) slots13[slot], slots2[slot] = self.w13[expert], self.w2[expert] - def kernel(x, w, ids, used): - return naive_moe(x, slots13, slots2, w, ids) + def kernel(x, w, ids, table, used): + self.assertTrue(bool((ids >= 0).all())) # padding id -1 arrives as num_experts + return naive_moe(x, slots13, slots2, w, table[ids.long()]) out = tiered_forward(self.x, self.topk_weights, self.topk_ids, cache, load, kernel, miss=miss, cpu_weights=lambda e: (self.w13[e], self.w2[e]), @@ -80,6 +81,33 @@ class TieredForward(unittest.TestCase): self.assertEqual(cpu_experts(self.x, self.topk_weights, self.topk_ids, set(), None).abs().sum().item(), 0.0) + def test_miss_min_weight_skips_missing_experts_for_every_token(self): + """[miss] min_weight < 1: every token skips experts that are not in VRAM while it keeps + that share of its router weight; experts some token copies anyway serve all tokens.""" + everyone = torch.ones(6, dtype=torch.bool) + resident = torch.zeros(EXPERTS, dtype=torch.bool) + resident[[0, 1, 2]] = True + _, first = degrade(self.topk_weights, self.topk_ids, everyone, resident, 0.5) + free = resident.clone() + free[first[first >= 0].long()] = True + w, i = degrade(self.topk_weights, self.topk_ids, everyone, free, 0.5) + out, loads, _ = self.run_tiered(pinned=[0, 1, 2], cache_slots=3, miss=Miss(min_weight=0.5)) + torch.testing.assert_close(out, naive_moe(self.x, self.w13, self.w2, w, i)) + self.assertEqual(sorted(loads), sorted(set(i[i >= 0].tolist()) - {0, 1, 2})) + self.assertLess(len(loads), 7) # computing everything copies all 7 others + self.assertEqual(self.seen_experts, list(range(EXPERTS))) # the trace sees the routing + out, loads, _ = self.run_tiered(pinned=[0, 1, 2], cache_slots=3, miss=Miss(min_weight=1.0)) + torch.testing.assert_close(out, self.expected) + self.assertEqual(len(loads), 7) + + def test_cpu_experts_take_matmul_functions(self): + """Quantized experts hand the CPU path x -> x @ w.T functions instead of weights.""" + def matmuls(e): + return (lambda v: v @ self.w13[e].t()), (lambda v: v @ self.w2[e].t()) + + out = cpu_experts(self.x, self.topk_weights, self.topk_ids, set(range(EXPERTS)), matmuls) + torch.testing.assert_close(out, self.expected.float()) + def test_route_stats(self): stats = route_stats(torch.tensor([[0.5, 0.25], [1.0, 0.0]]), torch.tensor([[2, 1], [2, -1]]), 4, recency_half_life=1.0) @@ -89,10 +117,10 @@ class TieredForward(unittest.TestCase): self.assertAlmostEqual(stats[2][2], 0.5 * 0.5 + 1.0) # the first token is one half-life old self.assertAlmostEqual(stats[1][2], 0.125) - def test_remap_zeroes_routes_outside_the_round(self): - weights, slots = remap(torch.ones(1, 3), torch.tensor([[2, 7, -1]]), {7: 4}, 10) - self.assertEqual(weights.tolist(), [[0.0, 1.0, 0.0]]) - self.assertEqual(slots.tolist(), [[0, 4, 0]]) + def test_slot_table_maps_the_round_and_skips_the_rest(self): + table = slot_table({7: 4, 2: 0}, 10, "cpu") + self.assertEqual(table.tolist(), [-1, -1, 0, -1, -1, -1, -1, 4, -1, -1, -1]) + self.assertEqual(table.dtype, torch.int32) # one extra row: padding id 10 -> -1 def test_degrade_drops_unavailable_routes_of_draft_tokens_only(self): weights = torch.tensor([[0.5, 0.3, 0.2], [0.5, 0.3, 0.2]]) @@ -115,6 +143,8 @@ class TieredForward(unittest.TestCase): available[[0, 1, 2, 4, 5, 6, 7, 8]] = True self.assertTrue(torch.equal( available_experts(self.topk_ids, draft, {0}, EXPERTS), available)) + _, kept = degrade(self.topk_weights, self.topk_ids, draft, available, 0.5) + available[kept[kept >= 0].long()] = True # experts some row copies anyway serve all rows expected = naive_moe(self.x, self.w13, self.w2, *degrade(self.topk_weights, self.topk_ids, draft, available, 0.5)) torch.testing.assert_close(out, expected)