diff --git a/benchmarks/README.md b/benchmarks/README.md index 6218903f23..817f99da18 100644 --- a/benchmarks/README.md +++ b/benchmarks/README.md @@ -27,5 +27,67 @@ batch size x miss rate. python benchmarks/bench_offload_cache_copy.py ``` +**`bench_scored_ensure.py`** — synthetic (no checkpoint): us per decode `ensure` call at a +fixed miss count, flashlib `lru_ensure` vs the `--moe-cache-policy` kernels (lru, kd, kdfb, +rule), from a CUDA graph of two decode steps on a full cache. `--stats` times it with +`--moe-collect-stats` accumulation on. + +```bash +python benchmarks/bench_scored_ensure.py --slots 1650 --k 10,20 +``` + +**`bench_moe_copy_overlap.py`** — synthetic (no checkpoint): `FREETOKEN_MOE_COPY_OVERLAP` +off vs on for Qwen4ExpMoE decode layers at the Qwen3.8-Flash-Next expert geometry, as CUDA +graphs. Checks the outputs are bitwise identical, times replays ABAB, and profiles the copy's +start after the ensure and how much of the side-stream shared expert runs inside it. + +```bash +python benchmarks/bench_moe_copy_overlap.py --bs 1 --slots 160 +``` + +**`bench_moe_prefetch.py`** — synthetic (no checkpoint): `FREETOKEN_MOE_PREFETCH=measure` +cost. The router-lookahead predictor alone per layer (GEMV `[bs, 2560] x [2560, 512]`, top-K +select, and the count kernel), then prefetch off vs measure on a real-geometry Qwen4ExpMoE decode +stack (512 experts, rule eviction, copy overlap on) with forced routing that misses exactly m +experts per layer, as CUDA graphs. `--on` times whole 48-layer decode steps with prefetch off vs +on instead: a GEMV per layer stands in for attention (`--filler-us`), each layer prefetches 3 or +4 candidates of which `--useful` are its misses, and `--delay-us` holds every prefetch copy back +to price a late one. It reports ms/step, demand misses/step and the late fraction. `--steady +--trace DIR` runs steady-state decode on a recorded routing trace instead (experts.npy, +logit_idx.npy, logit_val.npy, req_index.npy): no cache restore, so wrong prefetches and evictions +carry over, and a noisy copy of each layer's traced logits (`--sigma`) stands in for the +lookahead. Arms such as `off,measure,on,on:4/4` (budgets before GDN / QSA layers) replay step by +step in turn, each on its own copy of the cache state; it adds pollution (non-resident demands +per layer against off's misses) and the late fraction per layer kind. Slots and sigma were +calibrated to production's measure line (misses/layer ~4.3, precision ~0.56) at 24 layers. + +```bash +python benchmarks/bench_moe_prefetch.py --bs 1,2 --misses 0,4 +python benchmarks/bench_moe_prefetch.py --on --bs 1,2 --misses 4 --delay-us 0,50,150,300 +python benchmarks/bench_moe_prefetch.py --steady --trace DIR --bs 1,2 --steady-slots 740 --sigma 0.55 \ + --arms off,measure,on,on:empty,on:3/3,on:4/4,on:4/5 +``` + +**`bench_expert_copy_corun.py`** — synthetic (no checkpoint): the fused expert copy kernel +alone and beside compute, on registered host banks with the qwen4_exp NVFP4 row sizes. Sweeps +`fast_index_copy_multi` (threads x blocks/bank) and `fast_index_copy_multi_slim` +(blocks x threads x unroll) over 1..10-expert plans (GB/s, "hot" back-to-back and "cold" +spacer-separated copies), then reports the stretch sigma of bf16/fp8 GEMVs, a small matmul and +the NVFP4 routed decode GEMV while a copy loop runs on another stream. `--json` keeps the raw +per-trial times. + +```bash +python benchmarks/bench_expert_copy_corun.py --trials 5 --json copy_corun.json +``` + For host RAM vs PCIe bandwidth and the offload/hybrid backend pick, use `ft bench bw` instead — it writes the JSON profile the engine reads. + +**`bench_kv_quant.py`** compares BF16, FP8 and NVFP4 KV storage bytes, one-step +scatter latency and paged decode latency on synthetic inputs. No checkpoint is +required. Keep the GPU idle and use identical arguments for A/B comparisons; +this does not measure model quality or end-to-end serving throughput. + +```bash +PYTHONPATH=python:. uv run python benchmarks/bench_kv_quant.py --lengths 1024,8192,32768 +``` diff --git a/benchmarks/bench_expert_copy_corun.py b/benchmarks/bench_expert_copy_corun.py new file mode 100644 index 0000000000..b76b82ee6f --- /dev/null +++ b/benchmarks/bench_expert_copy_corun.py @@ -0,0 +1,518 @@ +"""Expert-copy kernel sweep and co-run stretch on real-size NVFP4 expert banks. + +Two measurements, no checkpoint needed: + +* sweep: GB/s of one fused expert copy (all 6 banks, 1..10 experts per plan) for each + ``fast_index_copy_multi`` (num_threads, blocks_per_bank) config, and for each + ``fast_index_copy_multi_slim`` (blocks, threads, unroll) config when that kernel exists; +* co-run: the stretch ``sigma = t_corun / t_alone - 1`` of decode compute kernels (bf16 and + fp8 GEMV [1,H]x[H,N], a small bf16 matmul, the NVFP4 routed decode GEMV) while a copy loop + runs on another stream, and the copy's GB/s while each compute loop runs. + +The host banks are registered anonymous mmaps (``HostBank``, like the serving banks) with +the qwen4_exp triton-NVFP4 row sizes. Source rows are not reused within a loop, the GEMV +weights are cycled over copies that together exceed L2, and every loop is a CUDA graph. +Alone and co-run trials are interleaved; each cell reports min and median over trials. + +Copy loops run in two modes. "hot": copy nodes back to back. "cold": a one-element spacer +kernel after every copy, as the ensure kernel precedes every decode copy. On the RTX 5080, in +graph replay, a copy node that follows another kernel's node runs ~8 us longer than one that +follows a copy node, for every variant (CUPTI puts it in the copy's own duration; eager +launches do not show it; cause not identified). Cold GB/s subtracts the spacer's alone time. + +Run: CUDA_VISIBLE_DEVICES=0 PYTHONPATH=python python benchmarks/bench_expert_copy_corun.py +""" + +from __future__ import annotations + +import argparse +import json +import math +import os +import statistics +import tempfile +from dataclasses import dataclass +from typing import Callable + +import torch +import torch.nn.functional as F + +from freetoken.gpu_select import assign_gpu, bind_assigned_gpu, single_gpu_arg +from freetoken.kernel import fast_index_copy as fic +from freetoken.kernel.pinned import device_ptr +from freetoken.moe.host_banks import HostBank +from freetoken.moe.offload_cache import _BANK_BYTES_PER_EXPERT + +L2_BYTES = 64 << 20 +SLEEP_CYCLES = 2_000_000 # ~1 ms head start so the host enqueues every graph before the GPU starts +LMAX = 16 # plan buffer length; the kernels read the live count from num_indices + + +def nvfp4_bank_bytes(h: int, i: int) -> list[int]: + # triton NVFP4 layout() order: gate_up, gate_up_scale, gate_up_global, down, down_scale, down_global + return [2 * i * (h // 2), 2 * i * (h // 16), 2 * i * 2, h * (i // 2), h * (i // 16), h * 2] + + +@dataclass(frozen=True) +class CopyVariant: + name: str + launch: Callable[..., None] + + +def multi_variant(threads: int, blocks_per_bank: int) -> CopyVariant: + def launch(dst, src, feat, di, si, num): + fic.fast_index_copy_multi_jit( + dst, src, feat, di, si, num, num_threads=threads, blocks_per_bank=blocks_per_bank + ) + + return CopyVariant(f"multi {threads}x{blocks_per_bank}", launch) + + +def slim_variant(blocks: int, threads: int, unroll: int) -> CopyVariant: + def launch(dst, src, feat, di, si, num): + fic.fast_index_copy_multi_slim_jit( + dst, src, feat, di, si, num, num_blocks=blocks, num_threads=threads, unroll=unroll + ) + + return CopyVariant(f"slim {blocks}x{threads}u{unroll}", launch) + + +class CopyRig: + """Registered host banks, GPU slot caches and the fused-copy descriptor tensors.""" + + def __init__(self, feats: list[int], n_src: int, n_slots: int, device: torch.device, seed: int): + self.feats = feats + self.expert_bytes = sum(feats) + self.n_src = n_src + self.n_slots = n_slots + self.device = device + gen = torch.Generator().manual_seed(seed) + self.host = [] + for feat in feats: + bank = HostBank((n_src, feat), torch.uint8) + bank.tensor.copy_(torch.randint(0, 256, (n_src, feat), dtype=torch.uint8, generator=gen)) + bank.pin() + self.host.append(bank) + self.slots = [torch.empty((n_slots, feat), dtype=torch.uint8, device=device) for feat in feats] + self.dst_ptrs = torch.tensor([c.data_ptr() for c in self.slots], dtype=torch.int64, device=device) + self.src_ptrs = torch.tensor([device_ptr(b.tensor) for b in self.host], dtype=torch.int64, device=device) + self.feat_bytes = torch.tensor(feats, dtype=torch.int64, device=device) + self.spacer_buf = torch.zeros(1, dtype=torch.int32, device=device) + self.gen = gen + self._cursor = 0 + self._perm = torch.randperm(n_src, generator=gen) + + def plans(self, n_launch: int, experts: int) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """[n_launch, LMAX] dst/src plans: source rows walk a permutation of the host rows + (no reuse until all are used), destination slots are distinct within a plan.""" + src = torch.zeros((n_launch, LMAX), dtype=torch.int32) + dst = torch.zeros((n_launch, LMAX), dtype=torch.int32) + for r in range(n_launch): + for j in range(experts): + if self._cursor == self.n_src: + self._perm = torch.randperm(self.n_src, generator=self.gen) + self._cursor = 0 + src[r, j] = self._perm[self._cursor] + self._cursor += 1 + dst[r, :experts] = torch.randperm(self.n_slots, generator=self.gen)[:experts] + num = torch.full((1,), experts, dtype=torch.int64, device=self.device) + return dst.to(self.device), src.to(self.device), num + + def copy_loop(self, variant: CopyVariant, n_launch: int, experts: int, cold: bool) -> Callable[[], None]: + dst, src, num = self.plans(n_launch, experts) + + def run(): + for r in range(n_launch): + variant.launch(self.dst_ptrs, self.src_ptrs, self.feat_bytes, dst[r], src[r], num) + if cold: + self.spacer_buf.add_(1) + + return run + + def spacer_us(self, stream: torch.cuda.Stream, trials: int) -> float: + n = 200 + + def run(): + for _ in range(n): + self.spacer_buf.add_(1) + + g = capture(run, stream) + return statistics.median(time_alone(g, n, stream) for _ in range(trials)) + + +@dataclass +class ComputeOp: + name: str + make: Callable[[int], Callable[[], None]] # n_launch -> loop body + single: Callable[[], None] + + +def _weight_copies(nbytes: int) -> int: + return max(2, math.ceil(2 * L2_BYTES / nbytes)) + + +def gemv_bf16(h: int, n: int, device) -> ComputeOp: + x = torch.randn(1, h, dtype=torch.bfloat16, device=device) + ws = [torch.randn(n, h, dtype=torch.bfloat16, device=device) for _ in range(_weight_copies(n * h * 2))] + + def make(n_launch): + def run(): + for r in range(n_launch): + F.linear(x, ws[r % len(ws)]) + + return run + + return ComputeOp(f"gemv_bf16 1x{h}x{n}", make, lambda: F.linear(x, ws[0])) + + +def gemv_fp8(h: int, n: int, device) -> ComputeOp: + fp8 = torch.float8_e4m3fn + x = torch.randn(1, h, device=device).to(fp8) + ws = [torch.randn(n, h, device=device).to(fp8) for _ in range(_weight_copies(n * h))] + one = torch.ones((), dtype=torch.float32, device=device) + + def mm(w): + torch._scaled_mm(x, w.t(), scale_a=one, scale_b=one, out_dtype=torch.bfloat16) + + def make(n_launch): + def run(): + for r in range(n_launch): + mm(ws[r % len(ws)]) + + return run + + return ComputeOp(f"gemv_fp8 1x{h}x{n}", make, lambda: mm(ws[0])) + + +def matmul_bf16(m: int, h: int, n: int, device) -> ComputeOp: + a = torch.randn(m, h, dtype=torch.bfloat16, device=device) + ws = [torch.randn(n, h, dtype=torch.bfloat16, device=device) for _ in range(_weight_copies(n * h * 2))] + + def make(n_launch): + def run(): + for r in range(n_launch): + F.linear(a, ws[r % len(ws)]) + + return run + + return ComputeOp(f"mm_bf16 {m}x{h}x{n}", make, lambda: F.linear(a, ws[0])) + + +def nvfp4_routed(h: int, i: int, top_k: int, n_slots: int, device) -> ComputeOp: + """The decode routed-expert GEMV pair (gate_up -> act -> down) over its own slot banks.""" + from freetoken.moe.fused_nvfp4 import fused_experts_decode_nvfp4_marlin + + fp8 = torch.float8_e4m3fn + gu = torch.randint(0, 256, (n_slots, 2 * i, h // 2), dtype=torch.uint8, device=device) + gus = torch.full((n_slots, 2 * i, h // 16), 1.0, device=device).to(fp8) + gug = torch.full((n_slots, 2 * i), 1.0, dtype=torch.float16, device=device) + dn = torch.randint(0, 256, (n_slots, h, i // 2), dtype=torch.uint8, device=device) + dns = torch.full((n_slots, h, i // 16), 1.0, device=device).to(fp8) + dng = torch.full((n_slots, h), 1.0, dtype=torch.float16, device=device) + x = torch.randn(1, h, dtype=torch.bfloat16, device=device) * 0.01 + w = torch.full((1, top_k), 1.0 / top_k, dtype=torch.float32, device=device) + groups = max(1, n_slots // top_k) + ids = [ + (torch.arange(top_k, dtype=torch.int32, device=device) + g * top_k).view(1, top_k) % n_slots + for g in range(groups) + ] + + def call(t): + fused_experts_decode_nvfp4_marlin(x, gu, gus, gug, dn, dns, dng, w, t, "silu", False) + + def make(n_launch): + def run(): + for r in range(n_launch): + call(ids[r % groups]) + + return run + + return ComputeOp(f"nvfp4_routed k={top_k}", make, lambda: call(ids[0])) + + +def capture(fn: Callable[[], None], stream: torch.cuda.Stream) -> torch.cuda.CUDAGraph: + with torch.cuda.stream(stream): + fn() # warm-up (JIT / Triton compile, cuBLAS heuristics) outside capture + stream.synchronize() + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g, stream=stream): + fn() + torch.cuda.synchronize() + g.keepalive = fn # the closure owns the plan tensors the graph reads + return g + + +def _event() -> torch.cuda.Event: + return torch.cuda.Event(enable_timing=True) + + +def time_alone(g: torch.cuda.CUDAGraph, n: int, stream: torch.cuda.Stream) -> float: + """us per launch of a graph holding ``n`` launches, replayed alone.""" + start, end = _event(), _event() + with torch.cuda.stream(stream): + torch.cuda._sleep(SLEEP_CYCLES) + start.record(stream) + g.replay() + end.record(stream) + end.synchronize() + return start.elapsed_time(end) * 1e3 / n + + +def time_corun( + g_fg: torch.cuda.CUDAGraph, n_fg: int, s_fg: torch.cuda.Stream, + g_bg: torch.cuda.CUDAGraph, s_bg: torch.cuda.Stream, +) -> tuple[float, bool]: + """us per launch of the foreground graph while the background graph runs on another + stream, and whether the background outlasted the foreground (a valid co-run).""" + ev0, bg_end, fg_start, fg_end = _event(), _event(), _event(), _event() + s0 = torch.cuda.current_stream() + torch.cuda._sleep(SLEEP_CYCLES) + ev0.record(s0) + s_bg.wait_event(ev0) + s_fg.wait_event(ev0) + with torch.cuda.stream(s_bg): + g_bg.replay() + bg_end.record(s_bg) + with torch.cuda.stream(s_fg): + torch.cuda._sleep(SLEEP_CYCLES // 20) # let the background reach steady state + fg_start.record(s_fg) + g_fg.replay() + fg_end.record(s_fg) + torch.cuda.synchronize() + covered = ev0.elapsed_time(bg_end) >= ev0.elapsed_time(fg_end) + return fg_start.elapsed_time(fg_end) * 1e3 / n_fg, covered + + +def kernel_info(fn: Callable[[], None]) -> list[dict]: + """(name, grid, block, regs/thread, smem) of every kernel one call launches.""" + from torch.profiler import ProfilerActivity, profile + + fn() + torch.cuda.synchronize() + with profile(activities=[ProfilerActivity.CUDA]) as prof: + fn() + torch.cuda.synchronize() + with tempfile.TemporaryDirectory() as d: + path = os.path.join(d, "trace.json") + prof.export_chrome_trace(path) + with open(path) as f: + events = json.load(f)["traceEvents"] + out = [] + for e in events: + if e.get("cat") != "kernel": + continue + a = e.get("args", {}) + out.append({ + "name": e["name"][:80], + "grid": a.get("grid"), + "block": a.get("block"), + "regs": a.get("registers per thread"), + "smem": a.get("shared memory"), + }) + return out + + +def _fmt_info(info: list[dict]) -> str: + return "; ".join( + f"{k['name'][:48]} grid={k['grid']} block={k['block']} regs={k['regs']} smem={k['smem']}" for k in info + ) + + +def _stats(xs: list[float]) -> tuple[float, float]: + return min(xs), statistics.median(xs) + + +def parse_pairs(text: str, n: int) -> list[tuple[int, ...]]: + out = [] + for item in text.split(","): + item = item.strip() + if not item: + continue + parts = tuple(int(p) for p in item.lower().replace("u", "x").split("x")) + if len(parts) != n: + raise argparse.ArgumentTypeError(f"{item!r}: expected {n} 'x'-separated ints") + out.append(parts) + return out + + +def parse_args() -> argparse.Namespace: + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("--gpu", type=single_gpu_arg, default=None, + help="GPU UUID or nvidia-smi index (default: the first visible GPU)") + p.add_argument("--hidden", type=int, default=2560) + p.add_argument("--inter", type=int, default=640) + p.add_argument("--src-experts", type=int, default=256, help="host rows per bank (pinned)") + p.add_argument("--slots", type=int, default=64, help="GPU slot-cache rows per bank") + p.add_argument("--multi", default="1024x8,1024x4,512x8,512x4,256x16,256x8,128x32", + help="fast_index_copy_multi configs, THREADSxBLOCKS_PER_BANK") + p.add_argument("--slim", default="8x256x4,16x256x2,12x256x4,16x256x4,8x256x2,4x256x4", + help="fast_index_copy_multi_slim configs, BLOCKSxTHREADSxUNROLL " + "(skipped if the kernel is absent)") + p.add_argument("--experts", default="1,2,3,4,5,6,7,8,9,10", help="plan sizes for the sweep") + p.add_argument("--corun-experts", default="4", help="plan sizes for the co-run copy loop") + p.add_argument("--link", choices=["hot", "cold", "both"], default="both", help="sweep link mode(s)") + p.add_argument("--corun-link", choices=["hot", "cold"], default="cold", help="co-run copy loop link mode") + p.add_argument("--trials", type=int, default=5) + p.add_argument("--copy-ms", type=float, default=6.0, help="target length of one timed copy loop") + p.add_argument("--compute-ms", type=float, default=3.0, help="target length of one timed compute loop") + p.add_argument("--cover", type=float, default=3.0, + help="background loop length as a multiple of the foreground's alone time") + p.add_argument("--skip-sweep", action="store_true") + p.add_argument("--skip-corun", action="store_true") + p.add_argument("--ops", default="all", + help="comma list of co-run ops by prefix: gemv_bf16,gemv_fp8,mm_bf16,nvfp4_routed") + p.add_argument("--json", default=None, help="write all results to this JSON path") + p.add_argument("--seed", type=int, default=0) + return p.parse_args() + + +def build_variants(args) -> list[CopyVariant]: + variants = [multi_variant(t, b) for t, b in parse_pairs(args.multi, 2)] + if hasattr(fic, "fast_index_copy_multi_slim_jit"): + variants += [slim_variant(g, t, u) for g, t, u in parse_pairs(args.slim, 3)] + else: + print("fast_index_copy_multi_slim_jit not available: slim configs skipped", flush=True) + return variants + + +def run_sweep(rig: CopyRig, variants, experts_list, args, stream, spacer_us, results) -> None: + modes = ["hot", "cold"] if args.link == "both" else [args.link] + for mode in modes: + cold = mode == "cold" + sub = spacer_us if cold else 0.0 + print(f"\n== sweep ({mode} link{f', spacer {sub:.2f} us subtracted' if cold else ''}): copy alone, " + f"expert={rig.expert_bytes} B, trials={args.trials}; cells are GB/s at min-time / " + f"median-time (us/copy median)", flush=True) + print("variant".ljust(18) + "".join(f"{e:>20d}" for e in experts_list), flush=True) + for v in variants: + row = v.name.ljust(18) + for e in experts_list: + n_launch = max(8, math.ceil(args.copy_ms * 1e3 / (e * 55.0))) + g = capture(rig.copy_loop(v, n_launch, e, cold), stream) + ts = [time_alone(g, n_launch, stream) - sub for _ in range(args.trials)] + t_min, t_med = _stats(ts) + nbytes = e * rig.expert_bytes + results["sweep"].append({ + "link": mode, "variant": v.name, "experts": e, "us_min": t_min, "us_median": t_med, + "gbps_best": nbytes / (t_min * 1e3), "gbps_median": nbytes / (t_med * 1e3), "us_all": ts, + }) + row += f"{nbytes / (t_min * 1e3):>7.1f}/{nbytes / (t_med * 1e3):5.1f} ({t_med:5.0f})" + del g + print(row, flush=True) + + +def run_corun(rig: CopyRig, variants, ops, corun_experts, args, s_fg, s_bg, spacer_us, results) -> None: + cold = args.corun_link == "cold" + sub = spacer_us if cold else 0.0 + print(f"\n== co-run ({args.corun_link} link copy loop): compute on one stream, copy loop on another; " + f"trials={args.trials} (interleaved alone/co-run); sigma = t_corun / t_alone - 1", flush=True) + for op in ops: + n_fg_probe = 20 + g_probe = capture(op.make(n_fg_probe), s_fg) + t_probe = statistics.median(time_alone(g_probe, n_fg_probe, s_fg) for _ in range(3)) + del g_probe + n_op = max(20, math.ceil(args.compute_ms * 1e3 / t_probe)) + g_op = capture(op.make(n_op), s_fg) + print(f"\n-- {op.name}: ~{t_probe:.1f} us/launch alone, loop of {n_op}", flush=True) + results["kernels"][op.name] = kernel_info(op.single) + print(f" kernels: {_fmt_info(results['kernels'][op.name])}", flush=True) + print(" " + "copy variant".ljust(18) + "E".rjust(3) + + " op_alone_us(min/med) op_corun_us(min/med) sigma(min/med)" + + " copy_GBps alone(med) -> under op(med) valid (! = alone median >10% over the probe:" + + " other GPU load, trust the min column)", flush=True) + for v in variants: + for e in corun_experts: + copy_launch_us = e * 55.0 + n_copy_long = max(4, math.ceil(args.cover * n_op * t_probe / copy_launch_us)) + g_copy_long = capture(rig.copy_loop(v, n_copy_long, e, cold), s_bg) + n_copy = max(8, math.ceil(args.copy_ms * 1e3 / copy_launch_us)) + g_copy = capture(rig.copy_loop(v, n_copy, e, cold), s_bg) + n_op_long = max(n_op, math.ceil(args.cover * n_copy * copy_launch_us / t_probe)) + g_op_long = capture(op.make(n_op_long), s_fg) if n_op_long > n_op else g_op + op_alone, op_co, cp_alone, cp_co = [], [], [], [] + valid_op = valid_cp = 0 + for _ in range(args.trials): + op_alone.append(time_alone(g_op, n_op, s_fg)) + t, ok = time_corun(g_op, n_op, s_fg, g_copy_long, s_bg) + op_co.append(t) + valid_op += ok + cp_alone.append(time_alone(g_copy, n_copy, s_bg) - sub) + t, ok = time_corun(g_copy, n_copy, s_bg, g_op_long, s_fg) + cp_co.append(t - sub) + valid_cp += ok + a_min, a_med = _stats(op_alone) + c_min, c_med = _stats(op_co) + sig_min, sig_med = c_min / a_min - 1, c_med / a_med - 1 + nbytes = e * rig.expert_bytes + gb_alone = nbytes / (statistics.median(cp_alone) * 1e3) + gb_co = nbytes / (statistics.median(cp_co) * 1e3) + suspect = a_med > 1.1 * t_probe + results["corun"].append({ + "link": args.corun_link, "op": op.name, "variant": v.name, "experts": e, "suspect": suspect, + "op_alone_us": op_alone, "op_corun_us": op_co, + "sigma_min": sig_min, "sigma_median": sig_med, + "copy_alone_us": cp_alone, "copy_corun_us": cp_co, + "copy_gbps_alone_median": gb_alone, "copy_gbps_corun_median": gb_co, + "valid_op_trials": valid_op, "valid_copy_trials": valid_cp, + }) + print(f" {v.name.ljust(18)}{e:>3d} {a_min:8.2f}/{a_med:8.2f} {c_min:8.2f}/{c_med:8.2f}" + f" {sig_min:+6.3f}/{sig_med:+6.3f} {gb_alone:6.1f} -> {gb_co:6.1f}" + f" {valid_op}/{valid_cp}{' !' if suspect else ''}", flush=True) + del g_copy_long, g_copy, g_op_long + del g_op + + +def main() -> None: + args = parse_args() + assert torch.cuda.is_available(), "CUDA is required" + try: + assign_gpu(args.gpu) + device = bind_assigned_gpu() + except (ValueError, RuntimeError) as e: + raise SystemExit(f"error: {e}") from e + + feats = nvfp4_bank_bytes(args.hidden, args.inter) + assert sum(feats) == _BANK_BYTES_PER_EXPERT["nvfp4"](args.hidden, args.inter) + if (args.hidden, args.inter) == (2560, 640): + assert sum(feats) == 2_772_480 + print("gpu", torch.cuda.get_device_name(device), "| banks", feats, "| expert", sum(feats), "B", flush=True) + rig = CopyRig(feats, args.src_experts, args.slots, device, args.seed) + variants = build_variants(args) + results = {"feats": feats, "sweep": [], "corun": [], "kernels": {}} + + s_fg = torch.cuda.Stream(device=device) + s_bg = torch.cuda.Stream(device=device) + spacer_us = rig.spacer_us(s_bg, args.trials) + results["spacer_us"] = spacer_us + print(f"spacer kernel alone: {spacer_us:.2f} us", flush=True) + for v in variants: + dst, src, num = rig.plans(1, 4) + info = kernel_info(lambda: v.launch(rig.dst_ptrs, rig.src_ptrs, rig.feat_bytes, dst[0], src[0], num)) + results["kernels"][v.name] = info + print(f"{v.name.ljust(18)} {_fmt_info(info)}", flush=True) + + if not args.skip_sweep: + experts_list = [int(x) for x in args.experts.split(",") if x] + run_sweep(rig, variants, experts_list, args, s_bg, spacer_us, results) + + if not args.skip_corun: + h, i = args.hidden, args.inter + all_ops = [ + gemv_bf16(h, 5120, device), gemv_bf16(h, 12288, device), + gemv_fp8(h, 5120, device), gemv_fp8(h, 12288, device), + matmul_bf16(256, h, h, device), + nvfp4_routed(h, i, 10, args.slots, device), + ] + wanted = [w.strip() for w in args.ops.split(",") if w.strip()] + ops = all_ops if wanted == ["all"] else [o for o in all_ops if any(o.name.startswith(w) for w in wanted)] + corun_experts = [int(x) for x in args.corun_experts.split(",") if x] + run_corun(rig, variants, ops, corun_experts, args, s_fg, s_bg, spacer_us, results) + + if args.json: + with open(args.json, "w") as f: + json.dump(results, f, indent=1) + print(f"\nwrote {args.json}", flush=True) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/bench_kv_quant.py b/benchmarks/bench_kv_quant.py new file mode 100644 index 0000000000..18707eb074 --- /dev/null +++ b/benchmarks/bench_kv_quant.py @@ -0,0 +1,77 @@ +"""Paged KV storage/scatter/decode microbenchmark, independent of model weights. + +Run with PYTHONPATH=python:. uv run python benchmarks/bench_kv_quant.py. +Compare identical arguments on the baseline and candidate; this does not measure +end-to-end model quality, TTFT, or serving throughput. +""" + +import argparse +import json + +import torch +import triton.testing + +from freetoken.distributed import set_tp_info +from freetoken.kernel.triton.attention import decode_paged_attention +from freetoken.kvcache.mha_pool import MHAKVCache + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--formats", default="none,fp8,nvfp4") + parser.add_argument("--lengths", default="1024,8192,32768") + parser.add_argument("--batch", type=int, default=1) + parser.add_argument("--heads", type=int, default=4) + parser.add_argument("--group", type=int, default=4) + parser.add_argument("--dim", type=int, default=128) + args = parser.parse_args() + set_tp_info(rank=0, size=1) + torch.manual_seed(42) + batch, heads, dim = args.batch, args.heads, args.dim + qheads = heads * args.group + device = torch.device("cuda") + results = [] + for length in map(int, args.lengths.split(",")): + slots = batch * length + k, v = [torch.randn(slots, heads * dim, device=device, dtype=torch.bfloat16) + for _ in range(2)] + loc = torch.arange(slots, device=device, dtype=torch.int32) + q = torch.randn(batch, qheads, dim, device=device, dtype=torch.bfloat16) + indptr = torch.arange(batch + 1, device=device, dtype=torch.int32) * length + pos = torch.full((batch,), length - 1, device=device, dtype=torch.int32) + scratch = torch.empty(batch, qheads, 8, dim, device=device) + lse = torch.empty(batch, qheads, 8, device=device) + splits = torch.full((batch,), 8, device=device, dtype=torch.int32) + out = torch.empty_like(q) + for quant in args.formats.split(","): + pool = MHAKVCache(heads, 1, dim, slots, 1, q.dtype, device, kv_quant=quant) + pool.store_kv(k, v, loc, 0) + extra = {} + if quant == "nvfp4": + extra = dict(kv_quant=quant, k_block_scale=pool.k_block_scale(0), + v_block_scale=pool.v_block_scale(0)) + kc, vc = [getattr(pool, name)(0).flatten(0, 1) for name in ("k_cache", "v_cache")] + + def decode(): + return decode_paged_attention(q, kc, vc, indptr, loc, pos, + scratch, lse, splits, 8, dim ** -.5, out=out, + k_scale=pool.k_scale(0), v_scale=pool.v_scale(0), **extra) + + decode() + decode_ms = triton.testing.do_bench(decode, warmup=100, rep=300) + store_ms = triton.testing.do_bench( + lambda: pool.store_kv(k[-batch:], v[-batch:], loc[-batch:], 0), + warmup=100, rep=300) + record = dict(format=quant, length=length, batch=batch, + kv_bytes=pool.unit_bytes()[0] * slots, + decode_ms=decode_ms, store_ms=store_ms) + results.append(record) + print(json.dumps(record), flush=True) + del pool, kc, vc + del k, v + print(json.dumps(dict(gpu=torch.cuda.get_device_name(), torch=torch.__version__, + heads=heads, group=args.group, dim=dim, results=results))) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/bench_moe_copy_overlap.py b/benchmarks/bench_moe_copy_overlap.py new file mode 100644 index 0000000000..535f7b1bd2 --- /dev/null +++ b/benchmarks/bench_moe_copy_overlap.py @@ -0,0 +1,233 @@ +"""Decode MoE copy/compute overlap bench (FREETOKEN_MOE_COPY_OVERLAP). + +Qwen4ExpMoE layers at the Qwen3.8-Flash-Next expert geometry (hidden 2560, moe and shared +intermediate 640, top-10, NVFP4 triton banks: 2,772,480 B per expert) over one offload cache, +with fewer experts per layer than the model so host and GPU memory stay small. Each arm +(overlap off / on) captures one decode step through every layer as a CUDA graph, then: + 1. bitwise check of the two arms over the same fresh inputs, + 2. replay time per step with CUDA events, arms interleaved ABAB, fresh inputs every round, + 3. a torch.profiler (CUPTI) trace of graph replays: per layer, the copy kernel's start slip + after the ensure, the shared expert's span (side stream) and whether it ends inside the copy, the join + gap, and the copy's duration against the off arm. +No attention or hyper-connection work runs, so the percentages overstate the model's. + +Run: CUDA_VISIBLE_DEVICES=0 PYTHONPATH=python python benchmarks/bench_moe_copy_overlap.py +""" + +from __future__ import annotations + +import argparse +import statistics +from types import SimpleNamespace + +import torch + +H, I, K = 2560, 640, 10 + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--layers", type=int, default=4) + parser.add_argument("--experts", type=int, default=64, help="experts per layer") + parser.add_argument("--slots", type=int, default=160, help="cache slots (160 of 4 x 64 gives ~4.4 misses/layer)") + parser.add_argument("--bs", type=int, default=1) + parser.add_argument("--steps", type=int, default=40, help="decode steps per timing round") + parser.add_argument("--rounds", type=int, default=6) + return parser.parse_args() + + +def nvfp4_quant(): + from freetoken.layers.quantization import QuantBackend, QuantConfig, set_quant_backend + + set_quant_backend(QuantBackend.parse("moe.nvfp4=triton")) + return QuantConfig.from_hf( + {"quantization_config": {"quant_method": "modelopt", "quant_algo": "NVFP4", "ignore": ["lm_head"]}} + ) + + +def build_arm(args, overlap: bool, banks) -> SimpleNamespace: + import freetoken.core as core + from freetoken.core import Context, set_global_ctx + from freetoken.layers.moe import OffloadMoELayer + from freetoken.models.qwen4_exp.moe import Qwen4ExpMoE + from freetoken.moe.offload_cache import OffloadMoeCache + from freetoken.utils.torch_utils import torch_dtype + + E, L = args.experts, args.layers + config = SimpleNamespace( + hidden_size=H, num_experts=E, num_experts_per_tok=K, moe_intermediate_size=I, + shared_expert_intermediate_size=I, norm_topk_prob=True, quant=None, + moe_strategy="offload", decode_target="gpu", + ) + device = torch.device("cuda") + with torch.device(device), torch_dtype(torch.bfloat16): + moes = [Qwen4ExpMoE(config, l, prefix=f"model.layers.{l}.mlp") for l in range(L)] + gen = torch.Generator(device=device).manual_seed(11) + for moe in moes: + for t in moe.state_dict().values(): + if t.is_floating_point(): + t.normal_(0.0, 0.02, generator=gen) + quant = nvfp4_quant() + experts = [ + OffloadMoELayer(l, E, K, H, I, quant_config=quant, prefix=f"model.layers.{l}.mlp.experts") + for l in range(L) + ] + cache = OffloadMoeCache( + num_layers=L, num_experts=E, cache_size=args.slots, device=device, + quant_format=banks.quant_format, layout=banks.layout, + max_slots=experts[0].quant_method.slot_limit(), decode_copy_overlap=overlap, + ) + cache.set_bank_sources(banks.sources) + cache.set_alphas(banks.gate_up_alpha, banks.down_alpha) + cache.collect_stats = True + for moe, ex in zip(moes, experts): + ex.offload_cache = cache + moe.experts = ex + cache.reset() + core._GLOBAL_CTX = None + ctx = Context(page_size=1) + set_global_ctx(ctx) + ctx._batch = SimpleNamespace(is_prefill=False) + + static = torch.randn(args.bs, H, device=device, dtype=torch.bfloat16) + out = torch.empty(L, args.bs, H, device=device, dtype=torch.bfloat16) + + def step(): + for l, moe in enumerate(moes): + out[l].copy_(moe.forward(static.clone())) + + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + step() # the eager warm-up the graph runner also does + torch.cuda.current_stream().wait_stream(side) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + step() + cache.reset() + cache.reset_stats() + # the graph holds raw pointers only: the modules must outlive it or replays read freed weights + return SimpleNamespace(graph=graph, static=static, out=out, cache=cache, moes=moes) + + +def fresh_inputs(args, n: int, seed: int) -> list[torch.Tensor]: + gen = torch.Generator(device="cuda").manual_seed(seed) + return [torch.randn(args.bs, H, device="cuda", dtype=torch.bfloat16, generator=gen) for _ in range(n)] + + +def replay(arm, inputs, keep: bool = False) -> list[torch.Tensor]: + outs = [] + for x in inputs: + arm.static.copy_(x) + arm.graph.replay() + if keep: + outs.append(arm.out.clone()) + return outs + + +def time_arm(arm, inputs) -> float: + start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True) + start.record() + replay(arm, inputs) + end.record() + torch.cuda.synchronize() + return start.elapsed_time(end) / len(inputs) + + +def trace(arm, inputs) -> list[tuple[float, float, str]]: + from torch.profiler import ProfilerActivity, profile + + with profile(activities=[ProfilerActivity.CUDA]) as prof: + replay(arm, inputs) + torch.cuda.synchronize() + return sorted( + (e.time_range.start, e.time_range.end, e.name) for e in prof.events() if e.device_type.name == "CUDA" + ) + + +def layer_calls(ev, overlap: bool) -> list[dict]: + """Per MoE layer call, located by kernel name (graph replay spreads branches over streams).""" + rows = [] + for i, (_, ens_end, name) in enumerate(ev): + if "_lru_ensure_kernel" not in name: + continue + copy = next((k for k in ev[i + 1:] if "fast_index_copy" in k[2]), None) + gemm = next((k for k in ev[i + 1:] if "_decode_nvfp4" in k[2]), None) + if copy is None or gemm is None: + continue + # the shared expert chain starts at the gate kernel before the ensure; with the overlap it runs on + # the side stream beside the routed path, without it serially ahead of that path's topk + j = max(j for j in range(i) if "_gate_sigmoid_kernel" in ev[j][2]) + routed_path = ("router", "topk", "ensure", "fast_index_copy") + chain = [k for k in (ev[j:] if overlap else ev[j:i]) if k[1] <= gemm[0] + and not any(m in k[2] for m in routed_path)] + rows.append({ + "copy_us": copy[1] - copy[0], + "slip_us": copy[0] - ens_end, + "join_gap_us": gemm[0] - copy[1], + "shared_span_us": max(k[1] for k in chain) - chain[0][0], + "shared_ends_in_copy": max(k[1] for k in chain) <= copy[1], + }) + return rows + + +def pct(xs, q: float) -> float: + xs = sorted(xs) + return xs[min(len(xs) - 1, int(q * len(xs)))] + + +def main() -> None: + args = parse_args() + from freetoken.distributed import set_tp_info + from freetoken.layers.moe import OffloadMoELayer + from freetoken.moe.expert_banks import build_expert_banks + + set_tp_info(rank=0, size=1) + torch.manual_seed(0) + method = OffloadMoELayer(0, args.experts, K, H, I, quant_config=nvfp4_quant(), prefix="model.layers.0.mlp.experts").quant_method + banks = build_expert_banks(method, args.layers, None, device=torch.device("cuda"), dummy=True) + expert_bytes = sum(layers[0][0].numel() * layers[0].element_size() for layers in banks.sources.values()) + print(f"expert {expert_bytes} B; {args.layers} layers x {args.experts} experts, {args.slots} slots, bs {args.bs}") + arms = {False: build_arm(args, False, banks), True: build_arm(args, True, banks)} + + seed = 100 + inputs = fresh_inputs(args, args.steps, seed) + outs = {k: replay(arm, inputs, keep=True) for k, arm in arms.items()} + same = all(torch.equal(a, b) for a, b in zip(outs[True], outs[False])) + active, missed, calls = arms[False].cache.lru_stats.sum(0).tolist() + print(f"bitwise identical on/off over {args.steps} steps: {same}; misses per layer call {missed / calls:.2f}") + + times = {False: [], True: []} + for r in range(args.rounds): + seed += 1 + inputs = fresh_inputs(args, args.steps, seed) + for k in (False, True) if r % 2 == 0 else (True, False): + times[k].append(time_arm(arms[k], inputs)) + for k in (False, True): + t = times[k] + print(f"overlap {'on ' if k else 'off'}: {statistics.median(t):.4f} ms/step median (min {min(t):.4f}, max {max(t):.4f}, n={len(t)})") + per_layer = [1000 * (a - b) / args.layers for a, b in zip(times[False], times[True])] + print(f"saving per layer, paired rounds (us): {[round(x, 1) for x in per_layer]}, median {statistics.median(per_layer):.1f}") + + inputs = fresh_inputs(args, 12, seed + 1) + rows = {k: layer_calls(trace(arms[k], inputs), k) for k in (False, True)} + for k in (False, True): + r = [x for x in rows[k] if x["copy_us"] > 5.0] # layer calls with at least one miss + if not r: + continue + slips = [x["slip_us"] for x in r] + print( + f"overlap {'on ' if k else 'off'}: {len(r)} layer calls with misses; copy {statistics.median(x['copy_us'] for x in r):.1f} us; " + f"copy start - ensure end median {statistics.median(slips):.2f} us (p90 {pct(slips, 0.9):.2f}, <=5 us in {100 * sum(s <= 5 for s in slips) / len(slips):.0f}%); " + f"GEMM start - copy end median {statistics.median(x['join_gap_us'] for x in r):.2f} us; " + f"shared expert span median {statistics.median(x['shared_span_us'] for x in r):.1f} us" + + (f", ends inside the copy in {100 * sum(x['shared_ends_in_copy'] for x in r) / len(r):.0f}%" if k else "") + ) + pairs = [(a["copy_us"], b["copy_us"]) for a, b in zip(rows[False], rows[True]) if a["copy_us"] > 5.0] + if pairs: + ratio = [b / a for a, b in pairs] + print(f"copy duration on/off: median {statistics.median(ratio):.3f} (p10 {pct(ratio, 0.1):.3f}, p90 {pct(ratio, 0.9):.3f})") + + +if __name__ == "__main__": + main() diff --git a/benchmarks/bench_moe_prefetch.py b/benchmarks/bench_moe_prefetch.py new file mode 100644 index 0000000000..be5d38ebd7 --- /dev/null +++ b/benchmarks/bench_moe_prefetch.py @@ -0,0 +1,580 @@ +"""FREETOKEN_MOE_PREFETCH cost and gain: the router-lookahead predictor alone, its effect on the +decode compute stream in measure mode, and (--on) whole decode steps with prefetch off vs on. + +1. Predictor alone, per layer: the lookahead GEMV ([bs, 2560] x [2560, 512] bf16, a distinct router + weight per layer as in the model), the top-K + select kernel, and both; one stream, CUDA graph. +2. The count kernel alone (it runs on the compute stream in measure mode). +3. A real-geometry decode stack (Qwen3.8-Flash-Next: hidden 2560, 512 experts, top-10, moe and + shared intermediate 640, NVFP4 triton banks) of N Qwen4ExpMoE blocks: router GEMV, shared + expert (on the FREETOKEN_MOE_COPY_OVERLAP side stream by default), top-k, rule ensure, miss + copy of exactly m experts from pinned host, grouped GEMM. The routing is forced (10 - m hot + experts + m cold ones per row) and the cache restored before every replay, so both arms copy + the same bytes. One CUDA graph per step, prefetch off vs measure. No attention runs, so the + percentages overstate the model's. +4. --on: the same stack at the model's depth (48 layers), with a bf16 GEMV "filler" before every + block standing in for attention and the rest of the layer (--filler-us, ~225 us between one + demand copy and the next in production). Prefetch on installs a fixed candidate list per layer: + 3 before GDN / 4 before attention layers, of which --useful are the layer's cold experts and the + rest never-routed ids (production: ~3.24 issued, ~1.75 useful, ~4.2 misses per layer at bs 1). + Reports ms/step, demand misses/step and the late fraction; --delay-us adds a sleep ahead of every + prefetch copy to price the compute stream's wait on a late copy. +5. --steady --trace DIR: steady-state decode on a recorded routing trace (no restore, so wrong + prefetches and evictions carry over), --steady-layers layers of it with --steady-slots slots, + a GEMV filler of --filler-gdn-us / --filler-qsa-us before GDN / QSA layers, and a lookahead + that is each layer's traced top-32 logits plus N(0, --sigma) noise. Every arm (off, measure, + on[:gdn/qsa budgets | :empty]) keeps its own cache state and all arms replay step by step in + turn, so drift hits them alike. Reports ms/step (mean, p50, p90), demand misses per layer, + issued/useful/precision, non-resident demands per layer (misses + prefetch hits: above off's + misses is pollution) and the late fraction before GDN and QSA layers. + +Run: CUDA_VISIBLE_DEVICES=0 PYTHONPATH=python python benchmarks/bench_moe_prefetch.py + ... benchmarks/bench_moe_prefetch.py --on --bs 1,2 --misses 4 + ... benchmarks/bench_moe_prefetch.py --steady --trace DIR --steady-slots 740 --sigma 0.55 +""" + +from __future__ import annotations + +import argparse +import gc +import statistics + +import torch +import torch.nn.functional as F + +H, E, TOP_K, INTER = 2560, 512, 10, 640 + + +def time_graph(run, reps: int, before=None) -> float: + """Median ms of one replay of ``run`` captured as a CUDA graph (after an eager warm-up).""" + if before: + before() + run() + torch.cuda.synchronize() + if before: + before() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run() + start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True) + times = [] + for rep in range(reps + 3): + if before: + before() + torch.cuda.synchronize() + start.record() + graph.replay() + end.record() + end.synchronize() + if rep >= 3: + times.append(start.elapsed_time(end)) + return statistics.median(times) + + +def bench_predictor(bs: int, layers: int, k: int, budget: int, reps: int) -> dict: + from freetoken.moe.prefetch import lookahead_select, prefetch_count + + dev = torch.device("cuda") + gen = torch.Generator(device=dev).manual_seed(bs) + weights = [torch.randn(E, H, device=dev, dtype=torch.bfloat16, generator=gen) * 0.02 for _ in range(layers)] + x = torch.randn(bs, H, device=dev, dtype=torch.bfloat16, generator=gen) + resident = torch.where(torch.rand(E, device=dev, generator=gen) < 0.4, 1, -1).to(torch.int32) + sel = torch.full((layers, 2 * k), -1, dtype=torch.int32, device=dev) + res = sel.clone() + logits = [F.linear(x, w) for w in weights] + + def gemv(): + for w in weights: + F.linear(x, w) + + def select(): + for layer in range(layers): + lookahead_select(logits[layer], resident, sel[layer], res[layer], k=k, budget=budget) + + def both(): + for layer, w in enumerate(weights): + lookahead_select(F.linear(x, w), resident, sel[layer], res[layer], k=k, budget=budget) + + slots = torch.randint(0, 1024, (bs * TOP_K,), device=dev, dtype=torch.int32, generator=gen) + id_of_slot = torch.randint(0, E, (1024,), device=dev, dtype=torch.int32, generator=gen) + misses = torch.ones(1, dtype=torch.int64, device=dev) + counters = torch.zeros((layers, 6), dtype=torch.int64, device=dev) + + def count(): + for layer in range(layers): + prefetch_count(slots, id_of_slot, misses, sel[layer], res[layer], counters[layer], id_base=0, rows=bs) + + per = 1000.0 / layers + return {name: time_graph(fn, reps) * per for name, fn in + (("gemv", gemv), ("select", select), ("predictor", both), ("count", count))} + + +class ForcedGate: + """A layer's real router GEMV whose logits are then replaced, so every replay routes (and + misses) the same experts; the next layer's lookahead runs through it too.""" + + def __init__(self, gate, logits: torch.Tensor): + self.gate, self.logits = gate, logits + + def forward(self, x: torch.Tensor) -> torch.Tensor: + self.gate.forward(x) + return self.logits + + +class DecodeStack: + """N Qwen4ExpMoE blocks (router, shared expert, nvfp4 offload experts) over one shared cache, + wired like qwen4_exp for the lookahead.""" + + def __init__(self, layers: int, slots: int, bs: int, misses: int, mode: str, overlap: bool, filler=None, + routing=None, budgets=(3, 4)): + import freetoken.core as core + from types import SimpleNamespace + + from freetoken.core import Context, set_global_ctx + from freetoken.distributed import set_tp_info, try_get_tp_info + from freetoken.layers.moe import OffloadMoELayer + from freetoken.layers.quantization import QuantBackend, QuantConfig, set_quant_backend + from freetoken.models.qwen4_exp.moe import Qwen4ExpMoE + from freetoken.moe.expert_banks import build_expert_banks + from freetoken.moe.offload_cache import OffloadMoeCache + from freetoken.utils.torch_utils import torch_dtype + + if try_get_tp_info() is None: + set_tp_info(rank=0, size=1) + set_quant_backend(QuantBackend.parse("moe.nvfp4=triton")) + quant = QuantConfig.from_hf({"quantization_config": {"quant_method": "modelopt", "quant_algo": "NVFP4", "ignore": ["lm_head"]}}) + dev = torch.device("cuda") + config = SimpleNamespace( + hidden_size=H, num_experts=E, num_experts_per_tok=TOP_K, moe_intermediate_size=INTER, + shared_expert_intermediate_size=INTER, norm_topk_prob=True, quant=None, + moe_strategy="offload", decode_target="gpu", + ) + with torch.device(dev), torch_dtype(torch.bfloat16): + self.moes = [Qwen4ExpMoE(config, i, prefix=f"model.layers.{i}.mlp") for i in range(layers)] + gen = torch.Generator(device=dev).manual_seed(0) + for moe in self.moes: + for t in moe.state_dict().values(): + if t.is_floating_point(): + t.normal_(0.0, 0.02, generator=gen) + experts = [ + OffloadMoELayer(i, E, TOP_K, H, INTER, quant_config=quant, prefix=f"model.layers.{i}.mlp.experts") + for i in range(layers) + ] + method = experts[0].quant_method + if not hasattr(DecodeStack, "_banks"): + DecodeStack._banks = build_expert_banks(method, 1, None, device=dev, dummy=True) # one layer, shared + banks = DecodeStack._banks + self.cache = cache = OffloadMoeCache( + num_layers=layers, num_experts=E, cache_size=slots, device=dev, cache_policy="rule", + quant_format=banks.quant_format, layout=banks.layout, max_slots=method.slot_limit(), + decode_copy_overlap=overlap, prefetch_mode=mode, + ) + cache.set_bank_sources({role: per_layer * layers for role, per_layer in banks.sources.items()}) + cache.set_alphas(banks.gate_up_alpha, banks.down_alpha) + core._GLOBAL_CTX = None + ctx = Context(page_size=1) + set_global_ctx(ctx) + ctx._batch = SimpleNamespace(is_prefill=False) + self.x = [torch.randn(bs, H, device=dev, dtype=torch.bfloat16, generator=gen) * 0.5 for _ in range(layers)] + self.filler = filler + if routing is not None: + # --steady: every layer routes this step's row of route_buf and layer L's lookahead (from + # layer L-1) returns look_buf[L]; both run the real router GEMV first + route_buf, look_buf = routing + gates = [moe.gate for moe in self.moes] + for i, (moe, ex) in enumerate(zip(self.moes, experts)): + ex.offload_cache = cache + moe.experts = ex + moe.gate = ForcedGate(gates[i], route_buf[i]) + for i in range(1, layers): + experts[i - 1].set_lookahead(ForcedGate(gates[i], look_buf[i]), i, budgets[(i + 1) % 4 == 0]) + self.bs = bs + cache.reset() + torch.cuda.synchronize() + return + # forced routing: row r routes hot experts [r*hot, (r+1)*hot) and m fresh ones past all rows' hot ones + hot = TOP_K - misses + for moe, ex in zip(self.moes, experts): + ex.offload_cache = cache + moe.experts = ex + lg = torch.zeros(bs, E, device=dev, dtype=torch.bfloat16) + for r in range(bs): + lg[r, list(range(r * hot, (r + 1) * hot))] = 10.0 + lg[r, list(range(bs * hot + r * misses, bs * hot + (r + 1) * misses))] = 10.0 + moe.gate = ForcedGate(moe.gate, lg) + for i in range(1, layers): + experts[i - 1].set_lookahead(self.moes[i].gate, i, budgets[(i + 1) % 4 == 0]) + self.bs, self.hot, self.misses = bs, hot, misses + # warm the hot experts in, then drop the fresh ones: every replay then misses exactly those + cache.reset() + self.step() + cache.slot_for_id[:, bs * hot:].fill_(-1) + fresh = (cache.id_of_slot >= 0) & (cache.id_of_slot % E >= bs * hot) + cache.id_of_slot[fresh] = -1 + cache.usage[fresh] = 0 + torch.cuda.synchronize() + self.state = [cache.slot_for_id, cache.id_of_slot, cache.usage, cache.step] + for name in ("evict_tok", "evict_last_tok", "evict_lc", "evict_ct", "evict_slot_owner", + "evict_slot_last_tok", "evict_slot_lc", "evict_slot_ct"): + if getattr(cache, name) is not None: + self.state.append(getattr(cache, name)) + self.snap = [t.clone() for t in self.state] + + def restore(self): + for dst, src in zip(self.state, self.snap): + dst.copy_(src) + + def step(self): + for i, (moe, x) in enumerate(zip(self.moes, self.x)): + filler = self.filler[i] if isinstance(self.filler, list) else self.filler + if filler is not None: + F.linear(x, filler) + moe.forward(x.clone()) + + def cold(self, layer: int) -> list[int]: + """The experts ``layer`` misses every replay, merged rank by rank across rows like the select.""" + per_row = [list(range(self.bs * self.hot + r * self.misses, self.bs * self.hot + (r + 1) * self.misses)) + for r in range(self.bs)] + return [row[i] for i in range(self.misses) for row in per_row] + + +def bench_stack(layers, slots, bs, misses, reps, overlap) -> dict: + from flashlib.kernels.slot_cache import Stat + + out = {} + for mode in ("off", "measure"): + stack = DecodeStack(layers, slots, bs, misses, mode, overlap) + stack.cache.collect_stats = True + stack.restore() + stack.cache.lru_stats.zero_() + stack.step() + got = stack.cache.lru_stats[:, Stat.MISS].tolist() + assert got == [bs * misses] * layers, f"expected {bs * misses} misses per layer, got {got}" + stack.cache.collect_stats = False + out[mode] = time_graph(stack.step, reps, before=stack.restore) * 1000.0 / layers + if stack.cache.prefetch is not None: + stack.cache.prefetch.counters.zero_() + stack.restore() + stack.step() + c = stack.cache.prefetch.counters.sum(0).tolist() + out["issued/layer"] = c[0] / max(c[3], 1) + del stack + torch.cuda.empty_cache() + return out + + +def budget_of(layer: int) -> int: + return 3 if (layer + 1) % 4 else 4 + + +def bench_on(layers, slots, bs, misses, useful, reps, filler_us, delays_us) -> list[dict]: + """ms per decode step, prefetch off vs on (and on with every prefetch copy delayed).""" + from flashlib.kernels.slot_cache import Stat + + from freetoken.moe.prefetch import COPIED, ISSUED, LATE, USEFUL + + dev = torch.device("cuda") + x = torch.randn(1, H, device=dev, dtype=torch.bfloat16) + probe = torch.randn(32768, H, device=dev, dtype=torch.bfloat16) + per_row_us = time_graph(lambda: F.linear(x, probe), reps) * 1000.0 / 32768 + del probe + filler = torch.randn(max(int(filler_us / per_row_us), 1), H, device=dev, dtype=torch.bfloat16) * 0.01 + cycles_per_us = 1_000_000 / (time_graph(lambda: torch.cuda._sleep(1_000_000), 5) * 1000.0) + out = [] + for arm in ["off", "on", *[f"on+{d}us" for d in delays_us if d]]: + stack = DecodeStack(layers, slots, bs, misses, "off" if arm == "off" else "on", True, filler=filler) + cache = stack.cache + want = [bs * misses] * layers + if arm != "off": + pf = cache.prefetch + pf.sel_override = torch.full_like(pf.sel, -1) + for layer in range(1, layers): + u = min(useful[layer % len(useful)], bs * misses) + wrong = list(range(E - 1, E - 1 - (budget_of(layer) - u), -1)) + ids = stack.cold(layer)[:u] + wrong + pf.sel_override[layer, : len(ids)] = torch.tensor(ids, dtype=torch.int32) + want[layer] -= u + if arm.startswith("on+"): + pf.delay_copy_cycles = int(float(arm[3:-2]) * cycles_per_us) + cache.collect_stats = True + stack.restore() + cache.lru_stats.zero_() + stack.step() + got = cache.lru_stats[:, Stat.MISS].tolist() + assert got == want, f"{arm}: expected demand misses {want}, got {got}" + cache.collect_stats = False + stack.restore() + stack.step() + torch.cuda.synchronize() + stack.restore() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + stack.step() + if cache.prefetch is not None: + cache.prefetch.stats.zero_() + start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True) + times = [] + for rep in range(reps + 3): + stack.restore() + torch.cuda.synchronize() + start.record() + graph.replay() + end.record() + end.synchronize() + if rep >= 3: + times.append(start.elapsed_time(end)) + row = {"arm": arm, "ms": statistics.median(times), "misses": sum(want)} + if cache.prefetch is not None: + st = cache.prefetch.stats.sum(0).tolist() + calls = max(reps + 3, 1) * (layers - 1) + row.update(issued=st[ISSUED] / calls, useful=st[USEFUL] / calls, late=st[LATE] / max(st[COPIED], 1)) + out.append(row) + del graph, stack, cache + gc.collect() # the module graph holds cycles: free this arm's slot banks before the next one + torch.cuda.empty_cache() + return out + + +def load_trace(path: str, layers: int, offset: int, bs: int, steps: int, skip: int, sigma: float, seed: int): + """Per-step routing and lookahead logits ``[steps, layers, bs, E]`` from a routing trace. + + ``path`` holds experts.npy ``[N, 48, 10]`` (routed ids), logit_idx.npy / logit_val.npy + ``[N, 48, 32]`` (each layer's own top-32 router logits) and req_index.npy ``[N]``. Row r of the + batch follows the r-th longest request from its ``skip``-th token on. The routing logits pick + exactly the traced top-10; the lookahead of layer L is layer L's traced top-32 logits plus + N(0, sigma) noise (everything else far below), a stand-in for W_gate[L] . x_{L-1}.""" + import os + from collections import Counter + + import numpy as np + + ex = np.load(os.path.join(path, "experts.npy"), mmap_mode="r") + li = np.load(os.path.join(path, "logit_idx.npy"), mmap_mode="r") + lv = np.load(os.path.join(path, "logit_val.npy"), mmap_mode="r") + req = np.load(os.path.join(path, "req_index.npy")) + reqs = [r for r, n in Counter(req.tolist()).most_common(bs)] + rng = np.random.default_rng(seed) + route = np.full((steps, layers, bs, E), -10.0, np.float32) + look = np.full((steps, layers, bs, E), -30.0, np.float32) + lsl = slice(offset, offset + layers) + for r, rid in enumerate(reqs): + rows = np.flatnonzero(req == rid)[skip: skip + steps] + assert len(rows) == steps, f"request {rid} has {len(rows) + skip} tokens, fewer than skip + steps" + routed = np.asarray(ex[rows, lsl]).astype(np.int64) # [steps, layers, 10] + np.put_along_axis(route[:, :, r], routed, 10.0 - 0.1 * np.arange(TOP_K), -1) + idx = np.asarray(li[rows, lsl]).astype(np.int64) + val = np.asarray(lv[rows, lsl]).astype(np.float32) + rng.normal(0.0, sigma, idx.shape).astype(np.float32) + np.put_along_axis(look[:, :, r], idx, val, -1) + dev = torch.device("cuda") + return (torch.from_numpy(route).to(dev, torch.bfloat16), torch.from_numpy(look).to(dev, torch.bfloat16), reqs) + + +def bench_steady(args, bs: int, arms: list[str]) -> list[dict]: + """ms per decode step in steady state on traced routing: no restore, the cache evolves. + + One stack, one graph per arm, each arm with its own copy of the cache's slot maps, eviction + state and counters: every step replays all arms back to back (rotating their order), so + clock or link drift hits every arm alike. Slot bytes are shared (timing does not read them).""" + from flashlib.kernels.slot_cache import Stat + + from freetoken.moe.prefetch import COPIED, ISSUED, LATE, MISSES, USEFUL + + dev = torch.device("cuda") + steps = args.warm + args.timed + route_all, look_all, reqs = load_trace(args.trace, args.steady_layers, args.layer_offset, bs, steps, + args.skip, args.sigma, args.seed) + layers = args.steady_layers + x = torch.randn(1, H, device=dev, dtype=torch.bfloat16) + probe = torch.randn(32768, H, device=dev, dtype=torch.bfloat16) + per_row_us = time_graph(lambda: F.linear(x, probe), args.reps) * 1000.0 / 32768 + del probe + fill = {us: torch.randn(max(int(us / per_row_us), 1), H, device=dev, dtype=torch.bfloat16) * 0.01 + for us in (args.filler_gdn_us, args.filler_qsa_us)} + # the filler of layer i stands in for its attention: QSA layers are 3, 7, ... (i % 4 == 3) + fillers = [fill[args.filler_qsa_us if i % 4 == 3 else args.filler_gdn_us] for i in range(layers)] + qsa = torch.tensor([i % 4 == 3 for i in range(layers)]) + route_buf = torch.empty_like(route_all[0]) + look_buf = torch.empty_like(look_all[0]) + stack = DecodeStack(layers, args.steady_slots, bs, 0, "on", True, filler=fillers, routing=(route_buf, look_buf)) + cache, pf = stack.cache, stack.cache.prefetch + cache.collect_stats = True + experts = [moe.experts for moe in stack.moes] + state = [cache.slot_for_id, cache.id_of_slot, cache.usage, cache.step, cache.lru_stats, pf.stats] + for name in ("evict_tok", "evict_last_tok", "evict_lc", "evict_ct", "evict_slot_owner", + "evict_slot_last_tok", "evict_slot_lc", "evict_slot_ct"): + if getattr(cache, name) is not None: + state.append(getattr(cache, name)) + init = [t.clone() for t in state] + + def load(snap): + for dst, src in zip(state, snap): + dst.copy_(src) + + def save(snap): + for dst, src in zip(snap, state): + dst.copy_(src) + + route_buf.copy_(route_all[0]) + look_buf.copy_(look_all[0]) + graphs, snaps, keep = {}, {}, [] + for arm in dict.fromkeys(arms): + mode, _, spec = arm.partition(":") + budgets = tuple(int(b) for b in spec.split("/")) if spec and spec != "empty" else (3, 4) + for i in range(1, layers): + experts[i - 1].set_lookahead(experts[i - 1]._lookahead[0], i, budgets[(i + 1) % 4 == 0]) + # every mode switch is read at enqueue time, so each arm's graph bakes its own path in + pf.mode = "on" if mode == "off" else mode + pf.sel_override = torch.full_like(pf.sel, -1) if spec == "empty" else None + keep.append(pf.sel_override) + cache.prefetch = None if mode == "off" else pf + load(init) + stack.step() # compile outside capture + torch.cuda.synchronize() + load(init) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + stack.step() + graphs[arm], snaps[arm] = graph, [t.clone() for t in init] + cache.prefetch, pf.mode, pf.sel_override = pf, "on", None + names = list(graphs) + events = {arm: [] for arm in names} + for t in range(steps): + if t == args.warm: + for arm in names: + snaps[arm][4].zero_() + snaps[arm][5].zero_() + route_buf.copy_(route_all[t]) + look_buf.copy_(look_all[t]) + for j in range(len(names)): + arm = names[(t + j) % len(names)] + load(snaps[arm]) + ev = (torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)) + ev[0].record() + graphs[arm].replay() + ev[1].record() + save(snaps[arm]) + if t >= args.warm: + events[arm].append(ev) + if t % 50 == 49: + torch.cuda.synchronize() # bound the event backlog + torch.cuda.synchronize() + out = [] + n = args.timed + for arm in names: + mode = arm.partition(":")[0] + times = sorted(a.elapsed_time(b) for a, b in events[arm]) + miss = snaps[arm][4][:, Stat.MISS].double().cpu() + row = {"arm": arm, "mean": statistics.fmean(times), "p50": times[len(times) // 2], + "p90": times[int(len(times) * 0.9)], "miss": miss.sum().item() / (n * layers), + "miss1": miss[1:].sum().item() / (n * (layers - 1))} + if mode != "off": + st = snaps[arm][5].double().cpu() + calls = n * (layers - 1) + row.update(issued=st[:, ISSUED].sum().item() / calls, useful=st[:, USEFUL].sum().item() / calls) + row["precision"] = row["useful"] / row["issued"] if row["issued"] else 0.0 + if mode == "measure": + row["coverage"] = st[:, USEFUL].sum().item() / max(st[:, MISSES].sum().item(), 1) + if mode == "on": + # every routed id that was not a normal resident: residual misses plus prefetch hits + row["nonres1"] = row["miss1"] + row["useful"] + for name, mask in (("late_gdn", ~qsa), ("late_qsa", qsa)): + m = mask.clone() + m[0] = False + row[name] = st[m, LATE].sum().item() / max(st[m, COPIED].sum().item(), 1) + row["late"] = st[:, LATE].sum().item() / max(st[:, COPIED].sum().item(), 1) + out.append(row) + print(f"# peak allocated {torch.cuda.max_memory_allocated(dev) / 2**30:.2f} GiB, requests {reqs}", flush=True) + del graphs, stack, cache, pf, events + gc.collect() + torch.cuda.empty_cache() + return out + + +def main() -> None: + ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument("--bs", default="1,2") + ap.add_argument("--k", type=int, default=16) + ap.add_argument("--budget", type=int, default=4) + ap.add_argument("--layers", type=int, default=8, help="MoE layers in the decode stack") + ap.add_argument("--slots", type=int, default=512, help="expert slots (2.77 MB each)") + ap.add_argument("--misses", default="0,4", help="demand misses per layer per row") + ap.add_argument("--reps", type=int, default=30) + ap.add_argument("--overlap", type=int, default=1, help="FREETOKEN_MOE_COPY_OVERLAP in the stack (production: 1)") + ap.add_argument("--on", action="store_true", help="time whole decode steps with prefetch off vs on instead") + ap.add_argument("--on-layers", type=int, default=48) + ap.add_argument("--on-slots", type=int, default=0, help="0 = every layer's routed experts + 64 (the warm-up must not evict)") + ap.add_argument("--filler-us", type=float, default=160.0, help="per-layer compute standing in for attention") + ap.add_argument("--useful", default="2,1,2,2", help="useful prefetches per layer, cycled over layers") + ap.add_argument("--delay-us", default="0,150,300", help="sleep ahead of every prefetch copy (0 = none)") + ap.add_argument("--steady", action="store_true", help="steady-state decode on a routing trace instead") + ap.add_argument("--trace", default=None, help="--steady: dir with experts/logit_idx/logit_val/req_index .npy") + ap.add_argument("--steady-layers", type=int, default=24, help="--steady: layers (VRAM: slots * 2.77 MB)") + ap.add_argument("--layer-offset", type=int, default=0, help="--steady: first traced layer") + ap.add_argument("--steady-slots", type=int, default=825, help="--steady: slots (1650 * 24 / 48 by default)") + ap.add_argument("--warm", type=int, default=300, help="--steady: untimed steps before timing") + ap.add_argument("--timed", type=int, default=500) + ap.add_argument("--skip", type=int, default=200, help="--steady: tokens skipped at each request's start") + ap.add_argument("--sigma", type=float, default=0.3, help="--steady: lookahead noise on the traced logits") + ap.add_argument("--seed", type=int, default=0) + ap.add_argument("--filler-gdn-us", type=float, default=171.0, help="--steady: compute before a GDN layer's MoE") + ap.add_argument("--filler-qsa-us", type=float, default=212.0, help="--steady: compute before a QSA layer's MoE") + ap.add_argument("--arms", default="off,measure,on,on:empty,on:2/3,on:3/3,on:4/4", + help="--steady: mode[:gdn/qsa budgets | :empty (on, nothing selected)]") + args = ap.parse_args() + from freetoken.env import ENV + + ENV.MOE_PREFETCH_K.value, ENV.MOE_PREFETCH_BUDGET.value = args.k, 0 if args.steady else args.budget + dev = torch.device("cuda") + if args.steady: + print(f"# {torch.cuda.get_device_name(dev)}, steady state on {args.trace}: {args.steady_layers} layers from " + f"{args.layer_offset}, {args.steady_slots} slots, rule, copy overlap on, filler {args.filler_gdn_us:.0f}/" + f"{args.filler_qsa_us:.0f} us (GDN/QSA), sigma {args.sigma}, {args.warm} warm + {args.timed} timed steps") + print(f"{'bs':>3} {'arm':>10} {'mean ms':>8} {'vs off':>7} {'p50':>7} {'p90':>7} {'miss/L':>7} {'miss/L1+':>8} " + f"{'iss/L':>6} {'use/L':>6} {'prec':>6} {'nonres/L1+':>10} {'late':>6} {'lateGDN':>7} {'lateQSA':>7}") + for bs in (int(b) for b in args.bs.split(",")): + rows = bench_steady(args, bs, args.arms.split(",")) + base = rows[0]["mean"] + for r in rows: + g = r.get + print(f"{bs:>3} {r['arm']:>10} {r['mean']:>8.3f} {r['mean'] / base - 1:>+7.1%} {r['p50']:>7.3f} " + f"{r['p90']:>7.3f} {r['miss']:>7.3f} {r['miss1']:>8.3f} {g('issued', 0):>6.2f} {g('useful', 0):>6.2f} " + f"{g('precision', 0):>6.3f} {g('nonres1', 0):>10.3f} {g('late', 0):>6.3f} {g('late_gdn', 0):>7.3f} " + f"{g('late_qsa', 0):>7.3f}", flush=True) + return + if args.on: + useful = [int(u) for u in args.useful.split(",")] + delays = [float(d) for d in args.delay_us.split(",")] + print(f"# {torch.cuda.get_device_name(dev)}, {args.on_layers} Qwen4ExpMoE nvfp4 layers, " + f"{args.on_slots or 'layers * bs * top-k + 64'} slots, " + f"rule, copy overlap on, filler {args.filler_us:.0f} us/layer, useful {useful}, median of {args.reps} replays") + print(f"{'bs':>3} {'m':>3} {'arm':>10} {'ms/step':>9} {'vs off':>8} {'misses/step':>12} " + f"{'issued/layer':>13} {'useful/layer':>13} {'late':>6}") + for bs in (int(b) for b in args.bs.split(",")): + for m in (int(x) for x in args.misses.split(",")): + slots = args.on_slots or args.on_layers * bs * TOP_K + 64 + rows = bench_on(args.on_layers, slots, bs, m, useful, args.reps, args.filler_us, delays) + base = rows[0]["ms"] + for r in rows: + extra = (f"{r['issued']:>13.2f} {r['useful']:>13.2f} {r['late']:>6.3f}" if "late" in r + else f"{'':>13} {'':>13} {'':>6}") + print(f"{bs:>3} {m:>3} {r['arm']:>10} {r['ms']:>9.3f} {r['ms'] / base - 1:>+8.1%} " + f"{r['misses']:>12d} {extra}", flush=True) + print(f"# peak allocated {torch.cuda.max_memory_allocated(dev) / 2**30:.2f} GiB") + return + print(f"# {torch.cuda.get_device_name(dev)}, K={args.k}, budget={args.budget}, median of {args.reps} graph replays") + print("# predictor alone, us per layer (47 layers, one stream)") + print(f"{'bs':>3} {'gemv':>8} {'select':>8} {'predictor':>10} {'count':>8}") + for bs in (int(b) for b in args.bs.split(",")): + r = bench_predictor(bs, 47, args.k, args.budget, args.reps) + print(f"{bs:>3} {r['gemv']:>8.2f} {r['select']:>8.2f} {r['predictor']:>10.2f} {r['count']:>8.2f}", flush=True) + print(f"# decode stack ({args.layers} Qwen4ExpMoE nvfp4 layers, {args.slots} slots, rule, overlap={args.overlap}), " + "us per layer, prefetch off vs measure") + print(f"{'bs':>3} {'m':>3} {'off':>9} {'measure':>9} {'delta':>8} {'issued/layer':>13}") + for bs in (int(b) for b in args.bs.split(",")): + for m in (int(x) for x in args.misses.split(",")): + r = bench_stack(args.layers, args.slots, bs, m, args.reps, bool(args.overlap)) + print(f"{bs:>3} {m:>3} {r['off']:>9.2f} {r['measure']:>9.2f} {r['measure'] - r['off']:>+8.2f} " + f"{r.get('issued/layer', 0.0):>13.2f}", flush=True) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/bench_scored_ensure.py b/benchmarks/bench_scored_ensure.py new file mode 100644 index 0000000000..5aa852716e --- /dev/null +++ b/benchmarks/bench_scored_ensure.py @@ -0,0 +1,158 @@ +"""Decode ensure cost per call: flashlib lru_ensure vs the scored --moe-cache-policy kernels. + +One ensure call per MoE layer per decode step, on a full cache (S slots over L x E experts). +Each timed call routes K ids of which exactly m miss: the hits are hot, recently used experts +and the victims are cold filler, under every policy. A CUDA graph of 2 decode steps (2 * L +calls; 1 step when that many fresh ids would not fit) is replayed from a restored snapshot, +so the number is the graph-node cost the decode graph pays. The miss count is verified +through the kernels' stats before timing. + +Run: CUDA_VISIBLE_DEVICES=0 PYTHONPATH=python python benchmarks/bench_scored_ensure.py +""" + +from __future__ import annotations + +import argparse +import statistics + +import numpy as np +import torch +from flashlib.kernels.slot_cache import N_STATS, Stat, lru_ensure + +from freetoken.moe import scored_ensure as se + +L, E, TOP_K = 48, 512, 10 +VARIANTS = ("flashlib-lru", "lru", "kd", "kdfb", "rule") + + +def build(variant: str, slots: int, k: int, m: int, dev: torch.device): + """Full-cache state plus 1-2 steps of L queries of k ids, the first k - m hot, the rest fresh misses.""" + policy = 0 if variant.endswith("lru") else se.POLICY_IDS[variant] + hot = k - m + calls = 2 * L if hot * L + 2 * L * m <= slots else L + step0, tok0 = 1_000_000, 1000 + slot_of_id = np.full(L * E, -1, np.int32) + id_of_slot = np.full(slots, -1, np.int32) + usage = np.zeros(slots, np.int64) + last_tok = np.full(L * E, se.LAST_TOK_NEVER, np.int64) + lc = np.full(L * E, se.LC_NEVER, np.int32) + ct = np.zeros(L * E, np.int64) + ids = [layer * E + e for layer in range(L) for e in range(hot)] + assert len(ids) + calls * m <= slots, "fresh ids would evict hot ones" + filler = (layer * E + e for e in range(k + m, E) for layer in range(L)) + ids += [next(filler) for _ in range(slots - len(ids))] + for s, key in enumerate(ids): + id_of_slot[s], slot_of_id[key] = key, s + is_hot = key % E < hot + usage[s] = step0 - 1 if is_hot else 1 + s + last_tok[key] = tok0 if is_hot else tok0 - 500 - s % 7 + if is_hot: + lc[key], ct[key] = 5 << se.Q, tok0 * L + key // E + queries = np.empty((calls, k), np.int32) + for i in range(calls): + visit = i // L + queries[i] = np.r_[np.arange(hot), hot + visit * m + np.arange(m)] + rng = np.random.default_rng(0) + logits = rng.normal(size=(calls, k // TOP_K, E)).astype(np.float32) + rows = queries.reshape(calls, k // TOP_K, -1).astype(np.int64) + np.put_along_axis(logits, rows, np.take_along_axis(logits, rows, -1) + 3.0, -1) + held = id_of_slot.astype(np.int64) + t = lambda a, dtype=None: torch.from_numpy(a).to(dev, dtype) # noqa: E731 + state = dict( + owner=t(id_of_slot.copy()), slot_last=t(last_tok[held]), slot_lc=t(lc[held]), slot_ct=t(ct[held]), + slot_of_id=t(slot_of_id), id_of_slot=t(id_of_slot), usage=t(usage), + step=torch.tensor(step0, dtype=torch.int64, device=dev), + tok=torch.tensor(tok0, dtype=torch.int64, device=dev), + last_tok=t(last_tok), lc=t(lc), ct=t(ct), queries=t(queries), + src=torch.zeros(slots, dtype=torch.int32, device=dev), + dst=torch.zeros(slots, dtype=torch.int32, device=dev), + num=torch.zeros((), dtype=torch.int64, device=dev), + stats=torch.zeros(L, N_STATS, dtype=torch.int64, device=dev), + ) + return policy, state, t(logits, torch.bfloat16), se.softplus2_table(dev) + + +def step_fn(variant, policy, st, logits, g_table, stats: bool): + def run(): + for i in range(st["queries"].shape[0]): + layer = i % L + q = st["queries"][i] + s = st["stats"][layer] if stats else None + if variant == "flashlib-lru": + lru_ensure(q, st["slot_of_id"], st["id_of_slot"], st["usage"], st["step"], q, + st["src"], st["dst"], st["num"], stats=s, id_base=layer * E) + continue + se.scored_ensure( + q, st["slot_of_id"], st["id_of_slot"], st["usage"], st["step"], q, + st["src"], st["dst"], st["num"], policy=policy, num_layers=L, num_experts=E, + tok=st["tok"], last_tok=st["last_tok"], lc=st["lc"], ct=st["ct"], g_table=g_table, + slot_owner=st["owner"], slot_last_tok=st["slot_last"], slot_lc=st["slot_lc"], slot_ct=st["slot_ct"], + router_logits=logits[i] if policy == 3 else None, near_miss_thr=0.25, + stats=s, id_base=layer * E, bump_tok=layer == 0, + ) + return run + + +def bench(variant, slots, k, m, dev, reps, stats): + policy, st, logits, g_table = build(variant, slots, k, m, dev) + snap = {name: v.clone() for name, v in st.items()} + + def restore(): + for name, v in snap.items(): + st[name].copy_(v) + + # verify the miss count (a stats build of the same calls), then time the requested build + restore() + step_fn(variant, policy, st, logits, g_table, True)() + got = int(st["stats"][:, Stat.MISS].sum()) + calls = st["queries"].shape[0] + assert got == calls * m, f"{variant} K={k} m={m}: {got} misses, expected {calls * m}" + restore() + run = step_fn(variant, policy, st, logits, g_table, stats) + run() # compile outside capture + restore() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run() + times = [] + start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True) + for rep in range(reps + 3): + restore() + torch.cuda.synchronize() + start.record() + graph.replay() + end.record() + end.synchronize() + if rep >= 3: + times.append(start.elapsed_time(end) * 1000.0 / calls) + return statistics.median(times) + + +def main() -> None: + ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument("--slots", type=int, default=1650) + ap.add_argument("--k", default="10,20", help="routed ids per call (10 = bs 1, 20 = bs 2)") + ap.add_argument("--misses", default="0-10", help="list or range; capped at K per call") + ap.add_argument("--variants", default=",".join(VARIANTS)) + ap.add_argument("--reps", type=int, default=30) + ap.add_argument("--stats", action="store_true", help="time with --moe-collect-stats accumulation on") + args = ap.parse_args() + dev = torch.device("cuda") + variants = args.variants.split(",") + misses = [] + for part in args.misses.split(","): + lo, _, hi = part.partition("-") + misses += range(int(lo), int(hi or lo) + 1) + print(f"# {torch.cuda.get_device_name(dev)}, S={args.slots}, us per ensure call " + f"(median of {args.reps} graph replays), stats={'on' if args.stats else 'off'}") + print(f"{'K':>3} {'m':>3} " + " ".join(f"{v:>13}" for v in variants) + " max delta vs flashlib-lru") + for k in (int(x) for x in args.k.split(",")): + for m in (x for x in misses if x <= k): + us = {v: bench(v, args.slots, k, m, dev, args.reps, args.stats) for v in variants} + base = us.get("flashlib-lru") + delta = "" if base is None else f" {max(us[v] - base for v in variants):+.2f}" + print(f"{k:>3} {m:>3} " + " ".join(f"{us[v]:>13.2f}" for v in variants) + delta, flush=True) + + +if __name__ == "__main__": + main() diff --git a/docs/cli.md b/docs/cli.md index b4c318e424..0a301a7777 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -42,6 +42,7 @@ parsers all resolve automatically from the checkpoint and the GPU. | `--host` | 127.0.0.1 | Bind address | | `--port` | 1919 | Bind port | | `--api-key` | `$FREETOKEN_API_KEY`, else off | Require the key on every route except `/health`, as `Authorization: Bearer ` or `x-api-key: `; other requests get a 401 | +| `--dist-port` | `--port` + 1 | Internal TP rendezvous port (loopback-only regardless of `--host`) | | `--gpu` | GPU 0 | GPU to run on: a UUID from `nvidia-smi -L` or an `nvidia-smi` index; see [below](#choosing-a-gpu) | | `--max-running-requests` | 4 | Max concurrently running requests | | `--max-output-tokens` | 32768 | Default output budget for requests that omit one | @@ -73,8 +74,32 @@ ft serve --model ... --gpu GPU-9e8d7c6b # the same card by UUID (a unique prefi | `--num-pages` / `--num-tokens` | auto | KV capacity override in pages / tokens (mutually exclusive; auto sizes from VRAM left after weights and MoE cache) | | `--page-size` | 1 | KV page size; DSV4 forces 128, the TRTLLM backend needs 16/32/64, SWA models require 1 | | `--cache-type` | radix | `radix` (prefix reuse; SWA/GDN-aware variants picked automatically) or `naive` | +| `--kv-cache-dtype` | bf16 | `bf16`, `fp8`, or `nvfp4` (see [NVFP4 KV cache](#nvfp4-kv-cache)): FP8 stores the KV cache as e4m3 codes plus one fp32 scale per (token, kv head), roughly doubling the tokens that fit in the same VRAM; see [FP8 KV cache](#fp8-kv-cache) | | `--attention-backend`, `--attn` | auto | `trtllm`/`fi`/`fa`/`triton`/`dsv4_sparse`/`dsa`; `prefill,decode` pair allowed; auto picks per model + GPU | +### FP8 KV cache + +`ft serve --kv-cache-dtype fp8` halves the bytes per cached token (8-bit codes instead +of 16), so a card that held N tokens holds close to 2N. Each `(token, kv head)` row +keeps its own fp32 scale, which costs ~3% back at `head_dim=128`. Requirements and +trade-offs: + +- Needs the **triton** attention backend; `--attn auto` selects it (and refuses an + explicit `fi`/`fa`/`trtllm`, which cannot be shown to apply these scales). +- Works on the plain paged, hybrid-SWA and QSA sparse (Qwen3.8-Flash-Next) KV pools. + On QSA the block-selection index keys stay 16-bit; only the selected K/V rows are + read back as codes. MLA/DSA latent KV, DeepSeek-V4's tiered pool and the block-sparse + MiniMax-M3 pool stay 16-bit; asking for fp8 there fails at startup rather than + silently ignoring the flag. +- The same bytes on every GPU FreeToken targets: the codes sit in a plain byte buffer + and are decoded in software, so the cache holds identical data and produces identical + numbers on any card (the fp8 type is deliberately kept out of the kernels, which is + also what makes the feature work on the RTX 30 series). +- Accuracy is checkpoint-dependent. Expect it to matter most on long contexts and on + models with outlier key channels; keep `bf16` when a run must be bit-reproducible. +- `ft ctl stats` / `/v1/cache/status` report the smaller `kv_bytes_per_token`, and + `ft ctl cache --kv N` moves the same (now cheaper) pool. + ### MoE offload See [models.md](models.md#moe-strategies) for what each strategy does. @@ -84,7 +109,9 @@ See [models.md](models.md#moe-strategies) for what each strategy does. | `--moe-strategy` | auto | `fused`/`offload`/`cpu`/`hybrid`; auto → offload, or hybrid with a `ft bench bw` profile; fused on unified-memory GPUs (GB10), see [models.md](models.md#unified-memory-gpus-gb10--dgx-spark). `--moe-backend` is the deprecated old spelling | | `--quant-backend` | auto | Kernel per quantized layer type, `layer[.kind]=name` entries: `linear=marlin,moe=b12x` or `moe.nvfp4=triton`. A layer-level entry applies to every kind whose table lists the name | | `--nvfp4-backend` | — | Deprecated: stands in for `--quant-backend moe.nvfp4=` (`flashinfer` means b12x); cannot be combined with `--quant-backend` | -| `--moe-cache-size` / `--moe-cache-rate` / `--moe-cache-auto` | auto | GPU expert-cache size as slots / fraction of all experts / sized from free VRAM (mutually exclusive; auto is enabled by default for offload-family strategies) | +| `--moe-cache-size` / `--moe-cache-rate` / `--moe-cache-mem` / `--moe-cache-auto` | auto | GPU expert-cache size as slots / fraction of all experts / memory such as `12GiB` / sized from free VRAM (mutually exclusive; auto is enabled by default for offload-family strategies) | +| `--moe-cache-scope` | unified | `unified`: one LRU pool of slots per expert quant (quant type and bank layout), so all MoE layers share one pool when they share a quant format; `per-layer`: each layer has its own pool. `per-layer` does not combine with `--moe-strategy cpu` or `--moe-cpu-layers` | +| `--moe-cache-split` | equal | How the pools share the slots: `equal` (in proportion to their layer counts), or `profile:`: every pool keeps `top_k` slots and the rest go to the most used experts in the profile's per-layer expert counts. With `offload`, a pool smaller than the expert count lowers `--max-running-requests` to its slots / top_k | | `--kv-reserve-tokens` | 8192 | KV token floor reserved before `--moe-cache-auto` fills experts | | `--moe-cpu-threads` | physical cores | CPU worker threads for the cpu/hybrid executor | | `--moe-cpu-layers` | all on GPU | With `offload`: which MoE layers decode on CPU (`3,7,11`, a count, a fraction, or `auto`). `auto` is for Windows/WSL only, where CUDA pinned memory is capped; every value needs an expert format the CPU executor serves (bf16, fp8_block, nvfp4, mxfp4) | @@ -102,6 +129,62 @@ See [models.md](models.md#moe-strategies) for what each strategy does. | `--reasoning-parser` | auto | Splits chain-of-thought into `reasoning_content`; auto-inferred; `off` disables | | `--enable-cache-report` | off | Report prefix-cache hits in each response's usage block | | `--anthropic-inline-system` | auto | Placement of late Anthropic system instructions: `auto`, `preserve`, or `fold` | +| `--enable-metrics-report` | off | Serve a per-request `metrics` object next to usage ([below](#per-request-performance-metrics)) | + +### Per-request performance metrics + +`--enable-metrics-report` adds a `metrics` object to each response on `/v1/chat/completions`, +`/v1/messages` and `/v1/responses`. It is not part of any of those protocols, which is why it is +off by default; clients that do not read it are unaffected. + +```json +{ + "usage": { "prompt_tokens": 8192, "completion_tokens": 256, "total_tokens": 8448 }, + "metrics": { + "ttft_ms": 320.5, + "prefill_tokens": 2048, + "cached_prompt_tokens": 6144, + "prefill_time_ms": 410.0, + "prefill_tokens_per_second": 4995.12, + "decode_tokens": 256, + "decode_time_ms": 5120.8, + "decode_tokens_per_second": 49.8, + "total_time_ms": 5480.2 + } +} +``` + +| Field | Meaning | +| --- | --- | +| `ttft_ms` | Time to the first generated token, from the point the request entered the generation path | +| `prefill_tokens` | Prompt tokens actually forwarded: `prompt_tokens - cached_prompt_tokens` | +| `cached_prompt_tokens` | Prompt tokens served from the prefix cache instead of recomputed | +| `prefill_time_ms` | The scheduler's own prefill span: admission to the token sampled off the last prefill chunk | +| `prefill_tokens_per_second` | `prefill_tokens / prefill_time_ms` | +| `decode_tokens` | Generated tokens (same as `usage.completion_tokens`) | +| `decode_time_ms` | First generated token to the last | +| `decode_tokens_per_second` | `(decode_tokens - 1) / decode_time_ms`; the first token comes out of prefill, so it spans no decode interval | +| `total_time_ms` | Request latency up to the terminal engine reply, excluding HTTP framing | + +Reading the numbers: + +- `prefill_time_ms` is measured inside the scheduler, which is the point of the flag: a client + timing SSE frames cannot see when prefill starts or ends. The others are measured at the API + layer and carry the same IPC hops a client's own timestamps would. +- `ttft_ms` is larger than `prefill_time_ms` by the time the request spent queued before + admission. Under load that gap is the queue, not the model. +- Every span is **this request's share of shared work**. Its prefill chunk is co-scheduled with + other prompts and its decode steps are batched with other requests, so the throughputs + describe this request under that load, not isolated engine benchmarks. For those, + send one request at a time. +- `cached_prompt_tokens` here always reports the real prefix-cache hit, whether or not + `--enable-cache-report` is set; that flag governs the billing fields in `usage`. Without it + `prefill_tokens_per_second` would be computed over tokens that were never forwarded. + +Streaming responses carry `metrics` on the same final message as usage, so the request has to ask +for usage too: `stream_options: {"include_usage": true}` on `/v1/chat/completions`. On +`/v1/messages` it rides `message_delta`; on `/v1/responses`, `response.completed`. +`/v1/completions` does not serve `metrics`: it does not go through the shared generation core. ### Image input @@ -122,7 +205,7 @@ so a client can gate its attachment controls without reading the checkpoint conf | `--image-min-tokens`, `--image-max-tokens` | processor defaults | Per-image token budget: the image processor resizes every image to take between these many tokens, converted to the family's own units by its processor. A family with fixed budgets honors the maximum only and refuses one below its smallest budget at start-up | | `--mm-processor-kwargs` | none | JSON object of extra keyword arguments for the checkpoint's image processor call, for knobs the token budget does not cover; applied after the budget, so an explicit key wins | | `--mm-embed-cache-device` | cpu | Where encoded image embeddings live between prefill chunks. `cpu` keeps them out of the VRAM budget; `cuda` skips the copy back | -| `--allowed-media-domains` | any | Comma-separated hostname allowlist for image URLs; requests for other domains are rejected with a 400. Empty allows any domain | +| `--allowed-media-domains` | any | Comma-separated hostname allowlist for image URLs; requests for other domains are rejected with a 400. Empty allows any domain. Redirects are re-checked against the list and may not land on a private or loopback address | | `--allowed-local-media-path` | off | Directory `file://` image refs may be read from; unset rejects local files | ## ft shell @@ -213,3 +296,35 @@ profile that `ft serve --moe-strategy auto` and `--moe-hybrid-max-fetch -1` then On NVLink-C2C hosts, see [models.md](models.md#nvlink-c2c-hosts-gh200--gb200) for `FREETOKEN_H2D_BLOCKS_PER_BANK` if `PCIe-gather` sits well below the linear H2D ceiling. + + + +### NVFP4 KV cache + +`ft serve --model --kv-cache-dtype nvfp4 --attention-backend triton` +opts into packed E2M1 KV storage. The initial implementation supports plain paged +FULL attention (MHA/GQA), hybrid-SWA, and the full-attention portion of hybrid-linear +models, QSA, and MLA/DSA (including GLM-5.3-Flash). Head dimensions must be +divisible by 16. DSV4 and BSA pools are rejected at startup. `auto` selects +Triton, QSA sparse, or DSA attention for supported models. For MLA/DSA use +`--attention-backend auto` or `--attention-backend dsa`. Only the latent slab is +quantized; indexer keys, kpool tails/gates, and recurrent states retain their +existing precision. A 512-element latent row occupies 292 bytes instead of 1024 +bytes in BF16, excluding those other tiers. + +Each K or V row stores `head_dim / 2` packed bytes, `head_dim / 16` E4M3 block-scale +bytes, and one FP32 row scale. At head_dim 128 this is 76 bytes, versus 256 for +BF16 and 132 for the existing FP8 format. Pool management, recurrent states, +attention workspace and model weights consume additional memory. + +The second-level scale is dynamic per token/head, so appending a token never +rescales an existing prefix. This is a FreeToken KV layout, not an external +NVFP4 checkpoint or attention-library ABI. K/V are restored inside attention; +Q and attention arithmetic retain their compute precision. The MoE weight option +`--nvfp4-backend` is independent. Paged MHA prefill uses fresh compute-dtype K/V +while cached prefixes are restored, as in the FP8 path. MLA/DSA stores fresh +latent rows first and reads the quantized cache in both prefill and decode. + +NVFP4 is opt-in: assess quality on your checkpoint and workload before using it +for long-context inference. Capacity savings do not guarantee faster decode; +packing, reconstruction, and the selected attention backend affect throughput. diff --git a/docs/install_amd.md b/docs/install_amd.md index 81d1ad5009..ed934ad657 100644 --- a/docs/install_amd.md +++ b/docs/install_amd.md @@ -45,3 +45,11 @@ On ROCm, these paths use the built-in fallbacks; do not install the `[accel]`, `[fi]`, or `[sgl]` extras. Forcing `--attention-backend fi`, `fa`, or `trtllm`, or NVFP4 Marlin/b12x (e.g. `--quant-backend moe.nvfp4=marlin` or `--quant-backend moe.nvfp4=b12x`), fails with a ROCm-specific error. + +## CPU and hybrid MoE + +`--moe-strategy cpu`, `hybrid` and `--moe-cpu-layers` sync the GPU with the CPU +expert pool through HIP stream memory operations. ROCm 10 records them into CUDA +graphs itself; on ROCm 7.14, which does not, FreeToken adds the graph nodes +explicitly. When the startup probe of that path fails, CUDA graph capture stops +with an error; run those strategies with `--cuda-graph-max-bs 0`. diff --git a/docs/models.md b/docs/models.md index cdadc2dcb1..d82f8f02f7 100644 --- a/docs/models.md +++ b/docs/models.md @@ -120,3 +120,9 @@ unset on PCIe GPUs: wider grids add no bandwidth there. "reasoning_effort": 37}`, and the checkpoint's encoder validates it. Runtime window-cache controls use 128-token pages. - Qwen3.8-Flash-Next keeps a 47.7 GiB PLE n-gram table pinned in host RAM. +- `--kv-cache-dtype fp8` (see [cli.md](cli.md#fp8-kv-cache)) covers the plain paged, + hybrid-SWA and QSA sparse KV pools — gpt-oss, Qwen3/3.5/3.6, GLM-4.x, Gemma-4, + MiniMax-M2.5, Muse-Glimmer, Llama/Qwen2/Mistral, Qwen3.8-Flash-Next (on QSA only the + selected K/V rows are read back as codes; block selection keeps 16-bit index keys). + MLA/DSA (GLM-5.2), DeepSeek-V4's tiered pool and MiniMax-M3's block-sparse pool stay + 16-bit and reject it. diff --git a/docs/openai_api.md b/docs/openai_api.md new file mode 100644 index 0000000000..0b1cc9d93c --- /dev/null +++ b/docs/openai_api.md @@ -0,0 +1,101 @@ +# OpenAI-compatible API: endpoints and parameters + +What `ft serve` accepts on its OpenAI-shaped routes, what it ignores, and what it refuses, +next to the same surface in vLLM, SGLang and llama.cpp's server. "honoured" means the value +changes the generation; "accepted" means the request is not rejected but the value has no +effect; "400" means the server refuses the request with a clear message. + +## Endpoints + +| route | FreeToken | vLLM | SGLang | llama.cpp | +|---|---|---|---|---| +| `POST /v1/chat/completions` | yes | yes | yes | yes | +| `POST /v1/completions` | yes | yes | yes | yes | +| `GET /v1/models`, `GET /v1/models/{id}` | yes | yes | yes | yes | +| `POST /v1/responses` (+ get, cancel) | yes | yes | yes | yes | +| `POST /v1/messages` (Anthropic), `count_tokens` | yes | yes | no | yes | +| `POST /tokenize`, `POST /detokenize` (also under `/v1/`) | yes | yes | yes | yes (`/tokenize`, `/detokenize`) | +| `GET /metrics` (Prometheus) | yes | yes | yes | yes | +| `GET /version` | yes | yes | yes | no (`/props`) | +| `GET /health` | yes | yes | yes | yes | +| `GET /v1/stats`, `/v1/cache/*`, `/v1/requests`, `/v1/admin/*` | yes (FreeToken control plane) | `/metrics` only | `/get_server_info`, `/flush_cache` | `/slots`, `/props`, `/metrics` | +| `POST /v1/embeddings`, `/rerank`, `/score`, `/pooling` | no (generative models only) | embedding / reranker models | embedding / reranker models | yes, any model with pooling | +| `POST /v1/audio/*`, `/v1/files`, `/v1/batches` | no | transcriptions, batch runner (offline) | transcriptions, files, batches | no | +| `POST /infill`, `/apply-template`, `/completion` (non-OpenAI) | `suffix` on `/v1/completions`; `/tokenize` with `messages` renders the template | no | no | yes | + +## Chat completions: request + +| parameter | FreeToken | vLLM | SGLang | llama.cpp | +|---|---|---|---|---| +| `model`, `messages` (text + `image_url` parts) | honoured | honoured | honoured | honoured | +| `max_tokens`, `max_completion_tokens` | honoured (default 32k) | honoured | honoured | `max_tokens` | +| `temperature`, `top_p`, `top_k` | honoured | honoured | honoured | honoured | +| `min_p` | honoured (prob. below `min_p` x top prob dropped) | honoured | honoured | honoured | +| `presence_penalty`, `frequency_penalty` | honoured, over generated tokens | honoured | honoured | honoured | +| `repetition_penalty` | honoured, over prompt + generated tokens (HF semantics) | honoured | honoured | `repeat_penalty` | +| `logit_bias` | honoured (token id -> bias, clamped to [-100, 100]) | honoured | honoured | honoured | +| `min_tokens` | honoured: no EOS / stop token before this many generated tokens | honoured | `min_new_tokens` | no | +| `stop` (strings), `stop_token_ids` | honoured | honoured | honoured | `stop` only | +| `include_stop_str_in_output` | honoured | honoured | `no_stop_trim` | no | +| `skip_special_tokens` | honoured; default off here, the reasoning and tool parsers read the tags | honoured (default on) | honoured (default on) | n/a | +| `n` | honoured, 1..16, one generation per choice (the prefix cache shares the prompt) | honoured | honoured | `n_cmpl` | +| `stream`, `stream_options.include_usage` | honoured | honoured | honoured | honoured | +| `logprobs`, `top_logprobs` | 400 on this route (see #224); `/v1/completions` `logprobs` 400 | honoured | honoured | `n_probs` | +| `tools`, `tool_choice` (`auto` / `none` / `required` / named), `parallel_tool_calls` | honoured | honoured | honoured | honoured | +| `reasoning_effort`, `thinking`, `chat_template_kwargs` | honoured | `chat_template_kwargs` | honoured | honoured | +| `continue_final_message` | honoured (the last assistant message is continued, no generation prompt) | honoured | honoured | no | +| `response_format` `json_object` / `json_schema`, guided decoding | 400 (no constrained decoding) | honoured (xgrammar / guidance) | honoured (xgrammar / outlines / llguidance) | `grammar`, `json_schema` | +| `seed` | accepted, not honoured (one batched sampling kernel, no per-request generator) | honoured | `sampling_seed` | honoured | +| `user`, `metadata`, `store`, `service_tier` | accepted, ignored | accepted, ignored | accepted, ignored | accepted, ignored | +| `request_id` | honoured: becomes the response id | honoured | `rid` | `id_slot` (different meaning) | +| `function_call` (legacy) | 400 | 400 | 400 | no | +| `best_of`, `use_beam_search`, `length_penalty`, `prompt_logprobs`, `echo` (chat), `truncate_prompt_tokens`, `priority`, `cache_salt`, `lora` | no | honoured | partly | partly | +| `ignore_eos` | honoured (benchmarking) | honoured | honoured | honoured | + +## Legacy completions: request + +`prompt` as a string or a list of strings (token-id prompts are refused); the same sampling +parameters as above; plus: + +| parameter | FreeToken | vLLM | SGLang | llama.cpp | +|---|---|---|---|---| +| `echo` | honoured (the prompt leads the text; not with `logprobs`) | honoured | honoured | no | +| `suffix` | honoured on models with fill-in-the-middle tokens (`<|fim_prefix|>` ... `<|fim_middle|>`): the prompt becomes prefix-suffix-middle | 400 | 400 | `/infill` | +| `n` | honoured, prompt-major choice order (`prompt i`, sample `k` -> `index i*n+k`) | honoured | honoured | no | +| `logprobs` | 400 (see #224 for the sampled-token logprobs) | honoured | honoured | `n_probs` | + +## Responses + +| field | FreeToken | notes | +|---|---|---| +| `id` | `chatcmpl-` / `cmpl-`, or `-` when the client sent one | | +| `created`, `model`, `object` | yes | | +| `system_fingerprint` | `null` | vLLM and SGLang also send `null` | +| `choices[].message.reasoning_content` | yes, when the model reasoned | vLLM `reasoning_content`, llama.cpp `reasoning_content` | +| `choices[].message.tool_calls` | yes | | +| `choices[].finish_reason` | `stop`, `length`, `tool_calls` | | +| `usage.prompt_tokens`, `completion_tokens`, `total_tokens` | yes | for `n > 1` the prompt is counted once, completions summed | +| `usage.prompt_tokens_details.cached_tokens` | with `--enable-cache-report` | vLLM, SGLang: always | +| `usage.completion_tokens_details.reasoning_tokens` | yes, when the model emitted a reasoning end tag | vLLM, SGLang | +| streaming: one `[DONE]`, an optional final usage chunk | yes | | +| llama.cpp `timings` | no; `/v1/stats` and `/metrics` carry the server-side rates | | + +## Tokenize / detokenize + +``` +POST /tokenize {"prompt": "text", "add_special_tokens": true, "return_token_strs": false} +POST /tokenize {"messages": [...], "add_generation_prompt": true, "chat_template_kwargs": {}} + -> {"count": N, "max_model_len": L, "tokens": [...], "token_strs": [...]?} +POST /detokenize {"tokens": [...], "skip_special_tokens": false} -> {"prompt": "text"} +``` + +`messages` are rendered through the model's chat template exactly as a generation would +render them, so `count` equals the `prompt_tokens` that generation reports. + +## Not implemented, on purpose + +- Constrained decoding (`response_format` json / json_schema, regex, grammar): needs a + grammar engine on the sampling path. +- Prompt logprobs and `echo` with `logprobs`: need the prefill logits. +- Per-request `seed`: the batched sampling kernel draws from one generator. +- Embeddings, reranking, scoring, audio: not generation. diff --git a/python/freetoken/attention/__init__.py b/python/freetoken/attention/__init__.py index eda3b61f55..c696e7d6a3 100644 --- a/python/freetoken/attention/__init__.py +++ b/python/freetoken/attention/__init__.py @@ -26,6 +26,7 @@ class BackendInfo: supported_types: frozenset[AttnType] requires_flashinfer: bool = False requires_sgl_kernel: bool = False + requires_sm90: bool = False requires_sm100: bool = False # Allowed page sizes (None -> any). Config-time resolution coerces to the last # entry when the resolved page_size is not in the list. @@ -33,6 +34,12 @@ class BackendInfo: # Whether forward() honors a per-call AttentionSpec (window/sm_scale/sinks). # Non-consumers raise on a non-None spec instead of silently dropping it. consumes_attn_spec: bool = False + # Whether forward() reads an fp8 KV pool (codes + per-token/per-head scales). + # Backends that hand the cache to an external kernel must opt out until that + # kernel is proven to apply our scale layout; the engine then refuses (or auto- + # avoids) them for --kv-cache-dtype fp8. + supports_fp8_kv: bool = False + supports_nvfp4_kv: bool = False SUPPORTED_ATTENTION_BACKENDS = Registry[BackendCreator]("Attention Backend") @@ -71,6 +78,7 @@ def create_fi_backend(config: ModelConfig): BackendInfo( supported_types=frozenset({AttnType.FULL}), requires_sgl_kernel=True, + requires_sm90=True, ), ) def create_fa_backend(config: ModelConfig): @@ -84,6 +92,8 @@ def create_fa_backend(config: ModelConfig): BackendInfo( supported_types=frozenset({AttnType.FULL, AttnType.SWA}), consumes_attn_spec=True, + supports_fp8_kv=True, + supports_nvfp4_kv=True, ), ) def create_triton_backend(config: ModelConfig): @@ -114,7 +124,11 @@ def create_dsv41_sparse_backend(config: ModelConfig): @SUPPORTED_ATTENTION_BACKENDS.register( "dsa", - BackendInfo(supported_types=frozenset({AttnType.MLA, AttnType.DSA})), + BackendInfo( + supported_types=frozenset({AttnType.MLA, AttnType.DSA}), + supports_fp8_kv=True, + supports_nvfp4_kv=True, + ), ) def create_dsa_backend(config: ModelConfig): # MLA with a grouped index (index_ratio > 1) is the kpool indexer layout. @@ -147,6 +161,10 @@ def create_m3_sparse_backend(config: ModelConfig): "qsa_sparse", BackendInfo( supported_types=frozenset({AttnType.QSA}), + # The attend kernel dequantizes on load (kernel/triton/qsa/attend.py); the + # compressed index keys it scores against are a separate, always-16-bit tier. + supports_fp8_kv=True, + supports_nvfp4_kv=True, # 64-token pages: a 4-token compress group never straddles a page, so the # compressed row of a group is page_base // 4 + block-in-page. page_sizes=(64,), diff --git a/python/freetoken/attention/dsa.py b/python/freetoken/attention/dsa.py index 398d170022..2d6781da51 100644 --- a/python/freetoken/attention/dsa.py +++ b/python/freetoken/attention/dsa.py @@ -201,6 +201,9 @@ def _attend( return glm_dsa_sparse_attn( q_cat, self.kvcache.latent_rows(layer_id), sel, self.sm_scale, counts=cnt, d_v=self.kv_lora_rank, + pool_scale=self.kvcache.latent_scale(layer_id), + kv_quant=self.kvcache.kv_quant, + pool_block_scale=self.kvcache.latent_block_scale(layer_id), ) def mla_forward( diff --git a/python/freetoken/attention/linear.py b/python/freetoken/attention/linear.py index f3717e58e7..c174c976ab 100644 --- a/python/freetoken/attention/linear.py +++ b/python/freetoken/attention/linear.py @@ -26,12 +26,16 @@ class FLAMetadata: fresh_state_indices prefill only: the state-pool slots whose sequence is fresh (cached_len == 0) and must be zeroed before the chunk kernel reads them in place. None if there are none / for decode. + max_seq_len prefill only: the longest extend_len, as a host int. It sizes the + varlen conv's launch grid; carrying it here keeps that launch free + of the D2H sync a device-side max would need (illegal under capture). """ cu_seqlens: torch.Tensor cache_indices: torch.Tensor has_initial_state: torch.Tensor | None = None fresh_state_indices: torch.Tensor | None = None + max_seq_len: int | None = None # --- hybrid-radix track-checkpoint (extra_buffer) fields; all None when not caching --- # For each request crossing a chunk-aligned (×CHUNK) boundary this forward, snapshot its @@ -87,6 +91,7 @@ def gdn_slot(r): fresh_state_indices=( fresh_host.to(device, non_blocking=True) if fresh_host is not None else None ), + max_seq_len=max(lens), **track, ) diff --git a/python/freetoken/attention/qsa_sparse.py b/python/freetoken/attention/qsa_sparse.py index 8ce63c4242..7d20092efc 100644 --- a/python/freetoken/attention/qsa_sparse.py +++ b/python/freetoken/attention/qsa_sparse.py @@ -92,6 +92,10 @@ class QSASparseMetadata(BaseAttnMetadata): rope_rows: torch.Tensor | None = None # [T] int32 arange q_rope_cache: torch.Tensor | None = None # [T, rotary_dim] float32 k_rope_cache: torch.Tensor | None = None # [T, rotary_dim] float32 + # KV host offload (--kv-host-pages): this forward's write pages (re-included in every + # fetch query so no mid-forward eviction can take them) and the physical out_loc. + write_pages: torch.Tensor | None = None + out_loc_gpu: torch.Tensor | None = None # fmt: on def get_last_indices(self, bs: int) -> torch.Tensor: @@ -112,7 +116,17 @@ def __init__(self, config: ModelConfig) -> None: f"qsa_sparse backend needs a QSA pool, got {type(self.kvcache).__name__}" ) self.device = self.kvcache.device + # The pool's COMPUTE dtype, never its store dtype (the contract lives in + # kvcache/base.py). These buffers feed the indexer -- qsa_index_norm_rope and + # qsa_mqa_paged -- whose tl.dot has no fp8 path, so an e4m3 q_index does not + # fail here, it fails at CUDA-graph capture with "Unsupported rhs dtype + # fp8e4nv". --kv-cache-dtype fp8 quantizes only the KV tiers; the index tiers + # stay 16-bit by design (kvcache/qsa_pool.py). self.dtype = self.kvcache.dtype + assert self.dtype.itemsize == 2, ( + f"QSA block selection needs a 16-bit compute dtype, got {self.dtype} -- " + "the KV pool must report its compute dtype, not e4m3 codes" + ) self.index_head_dim = self.kvcache.index_head_dim self.ratio = self.kvcache.index_ratio self.ring_capacity = self.kvcache.ring_capacity @@ -142,6 +156,12 @@ def __init__(self, config: ModelConfig) -> None: ).to(self.device) self._block_topk_kernel = _resolve_block_topk() + # KV host offload (None unless --kv-host-pages > 0): the GPU K/V pool is an LRU cache + # over the pinned host mirror; reads/writes translate logical -> physical pages. + self.kv_offloader = getattr(get_global_ctx(), "kv_offloader", None) + if self.kv_offloader is not None: + assert self.kv_offloader.pool is self.kvcache + self.kv_offloader.set_select_width(self.select_width, self.block_topk) # decode staging (static buffers under CUDA graphs; eager decode snapshots per step) self._graph: dict[str, torch.Tensor] = {} self.capture_bs: List[int] = [] @@ -287,7 +307,14 @@ def qsa_forward( md = batch.attn_metadata assert isinstance(md, QSASparseMetadata) slot = self._idx_slot[layer_id] - self.kvcache.store_kv(k, v, batch.out_loc, layer_id) + off = self.kv_offloader + if off is not None and slot == 0: + # Bind this forward's write pages before the first store; they ride every fetch + # query below, so no mid-forward eviction can take them (and the write-through + # mirror makes any other eviction lossless). + md.write_pages = off.ensure_write_pages(batch.out_loc, md.is_decode) + md.out_loc_gpu = off.translate_slots(batch.out_loc, md.is_decode) + self.kvcache.store_kv(k, v, batch.out_loc if off is None else md.out_loc_gpu, layer_id) if md.block_table is None: self._snapshot_decode(md, batch) if slot == 0 or md.cmp_rows is None: @@ -297,7 +324,14 @@ def qsa_forward( self._plan_index_writes(md, batch) self._update_index_cache(index, md, slot) + if off is not None: + # Write-through: the pinned host mirror (the backing store) sees every token. + off.mirror_store(k, v, batch.out_loc, layer_id, getattr(md, "out_loc_gpu", None)) indices = self._select(index, md, slot) + # K/V scale tensors are independent of the BF16 index tier, so selection is + # quantization-agnostic; only sparse K/V attention reconstructs the codes. + if off is not None: + return self._attend_offloaded(qsa_sparse_paged_attention, q, indices, md, layer_id) return qsa_sparse_paged_attention( q, self.kvcache.k_cache(layer_id), @@ -306,8 +340,50 @@ def qsa_forward( md.block_table, md.token_to_req, torch.empty_like(q), + k_scale=self.kvcache.k_scale(layer_id), + v_scale=self.kvcache.v_scale(layer_id), + kv_quant=self.kvcache.kv_quant, + k_block_scale=self.kvcache.k_block_scale(layer_id), + v_block_scale=self.kvcache.v_block_scale(layer_id), ) + def _attend_offloaded(self, attend, q, indices, md, layer_id) -> torch.Tensor: + """Attend through the host-mirrored page space: fetch the selected pages into + physical slots, translate the block table, call the attend kernel unmodified. + Prefill attends one row at a time (a row's scattered selection nearly fills the + physical pool) over a once-per-layer compacted page list -- no host syncs.""" + off = self.kv_offloader + k_cache = self.kvcache.k_cache(layer_id) + v_cache = self.kvcache.v_cache(layer_id) + # fp8 pool (PR #354): the attend kernel dequantizes with these; None on a bf16 pool. + ks = self.kvcache.k_scale(layer_id) + vs = self.kvcache.v_scale(layer_id) + if md.is_decode: + eff = off.fetch_for_attend(indices, md) + return attend(q, k_cache, v_cache, indices, eff, md.token_to_req, + torch.empty_like(q), k_scale=ks, v_scale=vs) + rows = q.shape[0] + if rows == 0: + return attend(q, k_cache, v_cache, indices, md.block_table, md.token_to_req, + torch.empty_like(q), k_scale=ks, v_scale=vs) + # Row-group loop: one compact for the whole layer (grid = rows, no sync), then the + # offloader packs rows into residency-feasible groups CPU-side with REAL page unions + # (one D2H per layer) -- consecutive rows select mostly the same pages, so groups + # pack several rows. Per group: one tight ensure + one translate + one attend. + sel_all = off.compact_all(indices, md.token_to_req, md.block_table) + out = torch.empty_like(q) + for (start, end), eff in off.iter_prefill_groups(sel_all, md.write_pages, md.block_table): + attend(q[start:end], k_cache, v_cache, indices[start:end], eff, + md.token_to_req[start:end], out[start:end], k_scale=ks, v_scale=vs) + dropped = off.trunc_count() # one sync per layer + if dropped: + logger.warning( + f"KV host offload: {dropped} selected pages dropped this layer " + f"(selection spans > {off.max_sel_pages} pages); the affected rows attended " + "over fallback slots" + ) + return out + def _plan_index_writes(self, md: QSASparseMetadata, batch: Batch) -> None: """Per-token slab row and ring row for this forward; the other QSA layers reuse it (it is layer-invariant). Pure device arithmetic: no host sync, graph-capturable.""" @@ -517,6 +593,8 @@ def empty(*shape: int, dtype: torch.dtype) -> torch.Tensor: } if topk_scratch: self._graph["topk_scratch"] = empty(chunk, topk_scratch, dtype=torch.int32) + if self.kv_offloader is not None: + self.kv_offloader.init_graph_buffers(max_bs, pages, self.select_width) def prepare_for_capture(self, batch: Batch) -> None: self.prepare_metadata(batch) diff --git a/python/freetoken/attention/triton.py b/python/freetoken/attention/triton.py index 8bf9792d55..2b0137cb4b 100644 --- a/python/freetoken/attention/triton.py +++ b/python/freetoken/attention/triton.py @@ -151,10 +151,19 @@ def forward( k_raw = self.kvcache.k_cache(layer_id) v_raw = self.kvcache.v_cache(layer_id) - kv_heads, head_dim = k_raw.shape[-2], k_raw.shape[-1] - assert head_dim == q.shape[-1] - k_cache = k_raw.view(-1, kv_heads, head_dim) - v_cache = v_raw.view(-1, kv_heads, head_dim) + kv_heads, stored_dim = k_raw.shape[-2], k_raw.shape[-1] + head_dim = q.shape[-1] + kv_quant = getattr(self.kvcache, "kv_quant", "none") + assert stored_dim == (head_dim // 2 if kv_quant == "nvfp4" else head_dim) + k_cache = k_raw.view(-1, kv_heads, stored_dim) + v_cache = v_raw.view(-1, kv_heads, stored_dim) + k_block_scale = self.kvcache.k_block_scale(layer_id) if kv_quant == "nvfp4" else None + v_block_scale = self.kvcache.v_block_scale(layer_id) if kv_quant == "nvfp4" else None + # An fp8 KV pool hands us its per-(token, head) scales; a 16-bit pool returns + # None and every kernel below keeps its original (scale-free) code path. + k_scale = self.kvcache.k_scale(layer_id) + v_scale = self.kvcache.v_scale(layer_id) + assert (k_scale is None) == (v_scale is None), "K and V scales come as a pair" spec = attn_spec or AttentionSpec() block_ends = batch.mm_block_ends if spec.bidirectional_mm_blocks else None @@ -182,6 +191,11 @@ def forward( sm_scale=scale, sliding_window=spec.sliding_window, sinks=spec.sinks, + k_scale=k_scale, + v_scale=v_scale, + kv_quant=kv_quant, + k_block_scale=k_block_scale, + v_block_scale=v_block_scale, ) if ( (not metadata.is_decode) @@ -202,6 +216,11 @@ def forward( sinks=spec.sinks, k_extend=k.view(q.shape[0], kv_heads, head_dim), v_extend=v.view(q.shape[0], kv_heads, head_dim), + k_scale=k_scale, + v_scale=v_scale, + kv_quant=kv_quant, + k_block_scale=k_block_scale, + v_block_scale=v_block_scale, block_ends=block_ends, ) if block_ends is not None: @@ -217,6 +236,11 @@ def forward( sm_scale=scale, sliding_window=spec.sliding_window, sinks=spec.sinks, + k_scale=k_scale, + v_scale=v_scale, + kv_quant=kv_quant, + k_block_scale=k_block_scale, + v_block_scale=v_block_scale, ) def prepare_metadata(self, batch: Batch) -> None: diff --git a/python/freetoken/checkpoint/__init__.py b/python/freetoken/checkpoint/__init__.py index f3ce70fcdc..406fdf997a 100644 --- a/python/freetoken/checkpoint/__init__.py +++ b/python/freetoken/checkpoint/__init__.py @@ -5,6 +5,7 @@ """ from .ftw import ( + FTWFormatError, FTWReader, FTWWriter, is_ftw_checkpoint, @@ -14,6 +15,6 @@ from .convert import convert_checkpoint __all__ = [ - "FTWReader", "FTWWriter", "is_ftw_checkpoint", + "FTWFormatError", "FTWReader", "FTWWriter", "is_ftw_checkpoint", "iter_ftw_weights", "load_ftw_banks", "convert_checkpoint", ] diff --git a/python/freetoken/checkpoint/convert.py b/python/freetoken/checkpoint/convert.py index 152f3871c0..191f294509 100644 --- a/python/freetoken/checkpoint/convert.py +++ b/python/freetoken/checkpoint/convert.py @@ -103,6 +103,25 @@ def _copy_metadata(model_path: str, out_dir: str) -> list[str]: return copied +def _refuse_occupied_output(out_dir: str) -> None: + """Refuse an output dir that already holds FTW output: the index of a finished conversion, + or shards left by an interrupted one. There is no resume; the writer would open shard + 00000 with "wb" and a stale index could then describe the wrong bytes.""" + from .ftw import INDEX_NAME, is_ftw_checkpoint + + if is_ftw_checkpoint(out_dir): + found = INDEX_NAME + else: + shards = glob.glob(os.path.join(glob.escape(out_dir), "freetoken-*.ftw")) + if not shards: + return + found = os.path.basename(min(shards)) + raise SystemExit( + f"{out_dir} already holds an FTW checkpoint ({found}); " + f"convert into an empty directory or remove the old one first" + ) + + class _ConvertSink: """Layer-completion sink for ``load_expert_banks(layer_sink=...)``: writes each completed layer's banks as their own FTW entries immediately (name @@ -184,6 +203,7 @@ def convert_checkpoint( if is_ftw_checkpoint(model_path): raise SystemExit(f"{model_path} is already an FTW checkpoint") + _refuse_occupied_output(out_dir) tp = try_get_tp_info() if tp is None: set_tp_info(rank=0, size=1) diff --git a/python/freetoken/checkpoint/ftw.py b/python/freetoken/checkpoint/ftw.py index e31e9d8278..31080d7384 100644 --- a/python/freetoken/checkpoint/ftw.py +++ b/python/freetoken/checkpoint/ftw.py @@ -37,6 +37,7 @@ import mmap import os import re +import stat import threading from collections.abc import Callable from concurrent.futures import ThreadPoolExecutor @@ -59,6 +60,14 @@ # Per-layer expert-bank entry name (converter streaming path, see checkpoint/convert.py): # each layer of a bank is its own FTW tensor instead of one flat [num_layers*E, ...] region. _LAYER_ENTRY_RE = re.compile(r"^(?P.+)#L(?P\d{5})$") +_MAX_INDEX_INT = (1 << 63) - 1 +_INDEX_FIELDS = { + "format", "version", "align", "shard_limit", "total_bytes", "tensors", "shards", +} + + +class FTWFormatError(ValueError): + """The FTW v1 index or one of its declared shards is structurally invalid.""" def layer_bank_entry_name(bank_name: str, layer_id: int) -> str: @@ -102,6 +111,321 @@ def _elsize(dt: torch.dtype) -> int: return torch.empty((), dtype=dt).element_size() +def _index_int(value, field: str, *, minimum: int = 0) -> int: + """Return a JSON integer that is safe to pass to file/tensor APIs. + + ``bool`` is deliberately excluded even though it subclasses ``int`` in Python. + File offsets, byte counts, and dimensions ultimately cross signed 64-bit APIs, so + accepting larger arbitrary-precision JSON numbers only postpones a less useful + overflow failure until allocation or I/O. + """ + if type(value) is not int or not minimum <= value <= _MAX_INDEX_INT: + raise FTWFormatError( + f"{field} must be an integer in [{minimum}, {_MAX_INDEX_INT}], got {value!r}" + ) + return value + + +def _required(obj: dict, key: str, owner: str): + if key not in obj: + raise FTWFormatError(f"{owner} is missing required field {key!r}") + return obj[key] + + +def _json_object(pairs: list[tuple[str, object]]) -> dict: + """Build one JSON object while rejecting ambiguous duplicate keys.""" + obj = {} + for key, value in pairs: + if key in obj: + raise FTWFormatError(f"FTW index contains duplicate JSON key {key!r}") + obj[key] = value + return obj + + +def _invalid_json_constant(value: str): + raise FTWFormatError(f"FTW index contains non-standard JSON constant {value}") + + +def _json_float(value: str) -> float: + parsed = float(value) + if not math.isfinite(parsed): + raise FTWFormatError(f"FTW index floating-point value is out of range: {value}") + return parsed + + +def _validate_metadata(value, field: str = "metadata") -> None: + """Reject metadata that cannot round-trip through strict JSON unambiguously.""" + if isinstance(value, dict): + for key, item in value.items(): + if not isinstance(key, str): + raise ValueError(f"{field} keys must be strings, got {key!r}") + _validate_metadata(item, f"{field}.{key}") + elif isinstance(value, (list, tuple)): + for i, item in enumerate(value): + _validate_metadata(item, f"{field}[{i}]") + elif isinstance(value, float): + if not math.isfinite(value): + raise ValueError(f"{field} must not contain NaN or infinity") + elif value is not None and not isinstance(value, (str, int, bool)): + raise ValueError( + f"{field} must contain only JSON-compatible values, got {type(value).__name__}" + ) + + +def _safe_shard_basename(value, field: str) -> str: + if (not isinstance(value, str) or not value or value in {".", ".."} + or os.path.basename(value) != value or "/" in value or "\\" in value + or "\x00" in value): + raise FTWFormatError(f"{field} must be a safe shard basename, got {value!r}") + return value + + +def _validate_ftw_v1_index(path: str, index) -> tuple[list[dict], list[dict]]: + """Eagerly validate all v1 structure used by :class:`FTWReader`. + + This is intentionally a structural check, not a content-authenticity check. It + prevents malformed JSON from turning into path traversal, integer overflow, + overlapping tensor views, or delayed mmap/O_DIRECT failures. A future format can + add hashes without weakening the valid v1 contract checked here. + """ + if not isinstance(index, dict): + raise FTWFormatError(f"FTW index root must be an object, got {type(index).__name__}") + + fmt = _required(index, "format", "FTW index") + if fmt != FORMAT_TAG: + raise FTWFormatError(f"not a {FORMAT_TAG} checkpoint: {path} (format={fmt!r})") + + version = _index_int(_required(index, "version", "FTW index"), "version") + if version != FORMAT_VERSION: + raise FTWFormatError( + f"unsupported FTW version {version}; this reader supports {FORMAT_VERSION}" + ) + + align = _index_int(_required(index, "align", "FTW index"), "align", minimum=1) + if align != ALIGN: + raise FTWFormatError(f"unsupported FTW alignment {align}; expected {ALIGN}") + + shard_limit = _index_int( + _required(index, "shard_limit", "FTW index"), "shard_limit", minimum=ALIGN + ) + if shard_limit % ALIGN: + raise FTWFormatError(f"shard_limit must be a multiple of {ALIGN}, got {shard_limit}") + + total_bytes = _index_int( + _required(index, "total_bytes", "FTW index"), "total_bytes" + ) + if total_bytes % ALIGN: + raise FTWFormatError(f"total_bytes must be {ALIGN}-aligned, got {total_bytes}") + + raw_shards = _required(index, "shards", "FTW index") + if not isinstance(raw_shards, list): + raise FTWFormatError("shards must be an array") + + shard_names: set[str] = set() + shards: list[dict] = [] + for i, shard in enumerate(raw_shards): + owner = f"shards[{i}]" + if not isinstance(shard, dict): + raise FTWFormatError(f"{owner} must be an object") + file = _safe_shard_basename(_required(shard, "file", owner), f"{owner}.file") + if file in shard_names: + raise FTWFormatError(f"duplicate shard file {file!r}") + shard_names.add(file) + + global_off = _index_int( + _required(shard, "global_off", owner), f"{owner}.global_off" + ) + nbytes = _index_int(_required(shard, "nbytes", owner), f"{owner}.nbytes") + if global_off % ALIGN: + raise FTWFormatError( + f"{owner}.global_off must be {ALIGN}-aligned, got {global_off}" + ) + if nbytes % ALIGN: + raise FTWFormatError(f"{owner}.nbytes must be {ALIGN}-aligned, got {nbytes}") + if nbytes > shard_limit: + raise FTWFormatError( + f"{owner}.nbytes {nbytes} exceeds shard_limit {shard_limit}" + ) + if global_off > _MAX_INDEX_INT - nbytes: + raise FTWFormatError(f"{owner} byte range exceeds signed 64-bit range") + + shard_path = os.path.join(path, file) + try: + # Follow symlinks: Hugging Face hub snapshots are symlink farms, so a shard + # entry routinely resolves into blobs/. The regular-file and size checks + # below therefore apply to the link target. + shard_stat = os.stat(shard_path) + except (OSError, UnicodeError) as exc: + # UnicodeError: a JSON-escaped lone surrogate is a valid str but not a + # filesystem name; keep it inside the format-error contract. + raise FTWFormatError(f"cannot stat FTW shard {file!r}: {exc}") from exc + if not stat.S_ISREG(shard_stat.st_mode): + raise FTWFormatError(f"FTW shard {file!r} is not a regular file") + if shard_stat.st_size != nbytes: + raise FTWFormatError( + f"FTW shard {file!r} size mismatch: index declares {nbytes}, " + f"file has {shard_stat.st_size}" + ) + shards.append(shard) + + shards.sort(key=lambda shard: shard["global_off"]) + if any(shard["nbytes"] == 0 for shard in shards): + # FTWWriter can produce one empty shard for a checkpoint containing only + # zero-sized tensors. It cannot produce empty shards in a non-empty stream. + if total_bytes != 0 or len(shards) != 1: + raise FTWFormatError("zero-length shards are only valid as one empty FTW shard") + + cursor = 0 + for shard in shards: + if shard["global_off"] != cursor: + relation = "an overlap" if shard["global_off"] < cursor else "a gap" + raise FTWFormatError( + f"FTW shard coverage has {relation}: expected offset {cursor}, " + f"got {shard['global_off']} for {shard['file']!r}" + ) + cursor += shard["nbytes"] + if cursor != total_bytes: + raise FTWFormatError( + f"shards cover {cursor} logical bytes but total_bytes is {total_bytes}" + ) + + raw_tensors = _required(index, "tensors", "FTW index") + if not isinstance(raw_tensors, list): + raise FTWFormatError("tensors must be an array") + + tensor_names: set[str] = set() + tensors: list[dict] = [] + allocations: list[tuple[int, int, str]] = [] + zero_allocations: list[tuple[int, str]] = [] + dtype_sizes: dict[torch.dtype, int] = {} + for i, tensor in enumerate(raw_tensors): + owner = f"tensors[{i}]" + if not isinstance(tensor, dict): + raise FTWFormatError(f"{owner} must be an object") + + name = _required(tensor, "name", owner) + if not isinstance(name, str): + raise FTWFormatError(f"{owner}.name must be a string") + if name in tensor_names: + raise FTWFormatError(f"duplicate tensor name {name!r}") + tensor_names.add(name) + + kind = _required(tensor, "kind", owner) + if not isinstance(kind, str): + raise FTWFormatError(f"{owner}.kind must be a string") + + dtype_name = _required(tensor, "dtype", owner) + if not isinstance(dtype_name, str): + raise FTWFormatError(f"{owner}.dtype must be a string") + dtype = getattr(torch, dtype_name, None) + if not isinstance(dtype, torch.dtype): + raise FTWFormatError(f"{owner}.dtype is not a known torch dtype: {dtype_name!r}") + element_size = dtype_sizes.get(dtype) + if element_size is None: + try: + element_size = _elsize(dtype) + # Match the actual reconstruction primitive in iter_ftw_weights rather + # than accepting every object torch happens to classify as a dtype. + torch.frombuffer(bytearray(element_size), dtype=dtype, count=1) + except (RuntimeError, TypeError, ValueError) as exc: + raise FTWFormatError( + f"{owner}.dtype is not readable FTW storage: {dtype_name!r}" + ) from exc + dtype_sizes[dtype] = element_size + + shape = _required(tensor, "shape", owner) + if not isinstance(shape, list): + raise FTWFormatError(f"{owner}.shape must be an array") + dims = [_index_int(dim, f"{owner}.shape[{j}]") for j, dim in enumerate(shape)] + numel = 1 + for dim in dims: + if dim == 0: + numel = 0 + elif numel and numel > _MAX_INDEX_INT // dim: + raise FTWFormatError(f"{owner}.shape element count exceeds signed 64-bit range") + else: + numel *= dim + if numel == 0: + try: + # Zero-numel shapes allocate no storage, so checking PyTorch's stride + # arithmetic here is cheap and catches dimensions that cannot be rebuilt. + torch.empty(tuple(dims), dtype=dtype) + except (RuntimeError, TypeError, ValueError) as exc: + raise FTWFormatError( + f"{owner}.shape cannot be reconstructed by torch: {shape!r}" + ) from exc + if numel and numel > _MAX_INDEX_INT // element_size: + raise FTWFormatError(f"{owner} byte size exceeds signed 64-bit range") + expected_nbytes = numel * element_size + + global_off = _index_int( + _required(tensor, "global_off", owner), f"{owner}.global_off" + ) + nbytes = _index_int(_required(tensor, "nbytes", owner), f"{owner}.nbytes") + if nbytes != expected_nbytes: + raise FTWFormatError( + f"{owner}.nbytes is {nbytes}, but shape {shape!r} and dtype " + f"{dtype_name!r} require {expected_nbytes}" + ) + if global_off % ALIGN: + raise FTWFormatError( + f"{owner}.global_off must be {ALIGN}-aligned, got {global_off}" + ) + if global_off > _MAX_INDEX_INT - nbytes: + raise FTWFormatError(f"{owner} byte range exceeds signed 64-bit range") + end = global_off + nbytes + padded_end = _align_up(end) + if end > total_bytes or padded_end > total_bytes: + raise FTWFormatError( + f"{owner} range [{global_off}, {end}) (padded to {padded_end}) " + f"exceeds total_bytes {total_bytes}" + ) + if nbytes: + allocations.append((global_off, padded_end, name)) + else: + zero_allocations.append((global_off, name)) + tensors.append(tensor) + + allocations.sort(key=lambda item: (item[0], item[1], item[2])) + allocation_boundaries = {0} + cursor = 0 + previous_name: str | None = None + for start, end, name in allocations: + if start < cursor: + raise FTWFormatError( + f"tensor {name!r} overlaps padded allocation for {previous_name!r}" + ) + if start > cursor: + raise FTWFormatError( + f"FTW tensor coverage has a gap: expected offset {cursor}, " + f"got {start} for {name!r}" + ) + allocation_boundaries.add(start) + allocation_boundaries.add(end) + cursor = end + previous_name = name + if cursor != total_bytes: + raise FTWFormatError( + f"tensors cover {cursor} padded bytes but total_bytes is {total_bytes}" + ) + + for offset, name in zero_allocations: + if offset not in allocation_boundaries: + raise FTWFormatError( + f"zero-sized tensor {name!r} at offset {offset} is not on a tensor boundary" + ) + + # Preserve both zero-byte forms emitted by FTWWriter: finalize-without-add_tensor + # has no shards, while one or more zero-sized tensors produce one empty shard. + if total_bytes == 0: + if tensors and not shards: + raise FTWFormatError("a zero-sized tensor checkpoint must contain one empty shard") + if not tensors and shards: + raise FTWFormatError("an empty checkpoint must not contain shard entries") + + return shards, tensors + + def is_ftw_checkpoint(path: str) -> bool: """True if ``path`` is a directory holding a FreeToken Weight (FTW) index.""" return os.path.isfile(os.path.join(path, INDEX_NAME)) @@ -133,11 +457,17 @@ class FTWWriter: """ def __init__(self, out_dir: str, *, shard_limit: int = DEFAULT_SHARD_LIMIT): - assert shard_limit % ALIGN == 0, "shard_limit must be a multiple of ALIGN" + if (type(shard_limit) is not int or shard_limit < ALIGN + or shard_limit > _MAX_INDEX_INT or shard_limit % ALIGN): + raise ValueError( + f"shard_limit must be an integer multiple of {ALIGN} in " + f"[{ALIGN}, {_MAX_INDEX_INT}], got {shard_limit!r}" + ) os.makedirs(out_dir, exist_ok=True) self.out_dir = out_dir self.shard_limit = shard_limit self._tensors: list[dict] = [] + self._tensor_names: set[str] = set() self._shards: list[dict] = [] self._global = 0 # running FTW offset (incl. padding) self._f = None # current shard file handle @@ -171,6 +501,10 @@ def _write_raw(self, data: memoryview) -> None: self._global += take def add_tensor(self, name: str, tensor: torch.Tensor, kind: str = "weight") -> None: + if not isinstance(name, str) or not isinstance(kind, str): + raise ValueError("FTW tensor name and kind must be strings") + if name in self._tensor_names: + raise ValueError(f"duplicate FTW tensor name {name!r}") t = tensor.detach().cpu().contiguous() raw = t.reshape(-1).view(torch.uint8) nbytes = int(raw.numel()) @@ -179,16 +513,24 @@ def add_tensor(self, name: str, tensor: torch.Tensor, kind: str = "weight") -> N and self._cur + nbytes > self.shard_limit): self._roll() global_off = self._global - assert global_off % ALIGN == 0, "tensor start must be aligned (invariant)" + if global_off % ALIGN: + raise RuntimeError("FTW writer invariant failed: tensor start is not aligned") self._write_raw(memoryview(raw.numpy())) self._tensors.append({"name": name, "kind": kind, "dtype": _dtype_str(t.dtype), "shape": list(t.shape), "global_off": global_off, "nbytes": nbytes}) + self._tensor_names.add(name) # pad to ALIGN so the next tensor starts aligned pad = _align_up(self._global) - self._global if pad: self._write_raw(memoryview(bytes(pad))) def finalize(self, meta: dict) -> dict: + if not isinstance(meta, dict): + raise ValueError(f"FTW metadata must be an object, got {type(meta).__name__}") + collisions = _INDEX_FIELDS & meta.keys() + if collisions: + raise ValueError(f"FTW metadata cannot override index fields: {sorted(collisions)}") + _validate_metadata(meta) if self._f is not None: self._shards.append({"file": _SHARD_FMT.format(self._shard_idx), "global_off": self._shard_start, "nbytes": self._cur}) @@ -199,7 +541,7 @@ def finalize(self, meta: dict) -> dict: "tensors": self._tensors, "shards": self._shards, **meta} tmp = os.path.join(self.out_dir, INDEX_NAME + ".tmp") with open(tmp, "w") as f: - json.dump(index, f) + json.dump(index, f, allow_nan=False) os.replace(tmp, os.path.join(self.out_dir, INDEX_NAME)) return index @@ -215,12 +557,22 @@ class FTWReader: (the rounding reads into the region's padding, which is discarded by the tensor view).""" def __init__(self, path: str): - with open(os.path.join(path, INDEX_NAME)) as f: - self.index = json.load(f) - assert self.index.get("format") == FORMAT_TAG, f"not a {FORMAT_TAG}: {path}" + index_path = os.path.join(path, INDEX_NAME) + try: + with open(index_path, encoding="utf-8") as f: + self.index = json.load( + f, + object_pairs_hook=_json_object, + parse_constant=_invalid_json_constant, + parse_float=_json_float, + ) + except FTWFormatError: + raise + except (ValueError, RecursionError) as exc: + raise FTWFormatError(f"malformed FTW index {index_path!r}: {exc}") from exc self.dir = path - self.shards = sorted(self.index["shards"], key=lambda s: s["global_off"]) - self.tensors = {t["name"]: t for t in self.index["tensors"]} + self.shards, tensors = _validate_ftw_v1_index(path, self.index) + self.tensors = {tensor["name"]: tensor for tensor in tensors} self._fds: dict[str, int] = {} self._maps: dict[str, tuple[mmap.mmap, memoryview]] = {} # O_DIRECT (DMA straight from disk, bypassing the page cache) is the fast path but a @@ -359,7 +711,9 @@ def rd(job): def _transient_buffer(nbytes: int) -> mmap.mmap: - return mmap.mmap(-1, _align_up(nbytes)) + # mmap rejects a zero-length mapping, but FTWWriter legitimately emits zero-sized + # tensors. Give those tensors one aligned backing page while exposing count=0 below. + return mmap.mmap(-1, _align_up(nbytes) or ALIGN) def iter_ftw_weights(path: str, *, kinds=("weight",), keep: Callable[[str], bool] | None = None, @@ -402,9 +756,14 @@ def _producer(): buf = _transient_buffer(e["nbytes"]) reader.read_into(memoryview(buf), e, workers=workers, chunk=chunk) dt = _dtype_of(e["dtype"]) - t = torch.frombuffer(buf, dtype=dt, count=e["nbytes"] // _elsize(dt)) - # a 0-d entry (a per-tensor scale) comes back 0-d, not [1] - if not _put((e["name"], t.view(*e["shape"]) if e["shape"] else t.view(()), buf, e["nbytes"])): + if e["nbytes"]: + t = torch.frombuffer(buf, dtype=dt, count=e["nbytes"] // _elsize(dt)) + # a 0-d entry (a per-tensor scale) comes back 0-d, not [1] + t = t.view(tuple(e["shape"])) if e["shape"] else t.view(()) + else: + # torch.frombuffer rejects count=0 even with a non-empty backing map. + t = torch.empty(tuple(e["shape"]), dtype=dt) + if not _put((e["name"], t, buf, e["nbytes"])): return except BaseException as ex: # surface to consumer err.append(ex) @@ -460,12 +819,14 @@ def _split_bank_entries(reader: FTWReader, path: str, num_layers: int): per_layer_groups.setdefault(m.group("base"), {})[int(m.group("layer"))] = e mixed = {e["name"] for e in flat_entries} & per_layer_groups.keys() - assert not mixed, f"FTW bank(s) mix flat and per-layer row layouts: {sorted(mixed)}" + if mixed: + raise FTWFormatError(f"FTW bank(s) mix flat and per-layer row layouts: {sorted(mixed)}") for base, by_layer in per_layer_groups.items(): - assert sorted(by_layer) == list(range(num_layers)), ( - f"FTW bank {base!r} has per-layer entries for layers {sorted(by_layer)}, " - f"expected exactly range({num_layers})" - ) + if sorted(by_layer) != list(range(num_layers)): + raise FTWFormatError( + f"FTW bank {base!r} has per-layer entries for layers {sorted(by_layer)}, " + f"expected exactly range({num_layers})" + ) return bank_entries, alpha_entries, flat_entries, per_layer_groups @@ -502,13 +863,20 @@ def load_ftw_banks( vectors, unaffected by the row split (fixed GPU residency; see ``cache_budget.expert_bytes_per_slot``). """ + if type(num_layers) is not int or num_layers < 1: + raise ValueError(f"num_layers must be a positive integer, got {num_layers!r}") + if layer_residency is not None and len(layer_residency) != num_layers: + raise ValueError( + f"layer_residency has {len(layer_residency)} entries; expected {num_layers}" + ) + from freetoken.moe.host_banks import ( HostBank, HostResidency, PinPipeline, alloc_banks, born_pinned_default, ) from freetoken.utils.progress import byte_bar - residency = layer_residency or [HostResidency.PINNED.value] * num_layers - assert len(residency) == num_layers, (len(residency), num_layers) + residency = ([HostResidency.PINNED.value] * num_layers + if layer_residency is None else layer_residency) # PINNED layers are born-pinned (cudaHostAlloc) where that wins (see born_pinned_default); LOCKED/PAGEABLE layers stay lazy mmaps born = born_pinned_default() @@ -543,13 +911,25 @@ def _backing(layer_id: int) -> str: for e in flat_entries: name = e["name"] + if not e["shape"]: + reader.close() + raise FTWFormatError(f"FTW bank {name!r} must have at least one dimension") total, *row_shape = e["shape"] - assert total % num_layers == 0, (name, total, num_layers) + if total % num_layers: + reader.close() + raise FTWFormatError( + f"FTW bank {name!r} has {total} rows, not divisible by " + f"num_layers={num_layers}" + ) num_experts = total // num_layers dtype = _dtype_of(e["dtype"]) row_bytes = (math.prod(row_shape) if row_shape else 1) * _elsize(dtype) layer_bytes = num_experts * row_bytes - assert layer_bytes * num_layers == e["nbytes"], (name, layer_bytes, num_layers, e["nbytes"]) + if layer_bytes * num_layers != e["nbytes"]: + reader.close() + raise FTWFormatError( + f"FTW bank {name!r} byte geometry does not match its layer layout" + ) row_hb[name] = [] row_view_args[name] = [] for layer_id in range(num_layers): @@ -567,7 +947,11 @@ def _backing(layer_id: int) -> str: row_view_args[base] = [] for layer_id in range(num_layers): e = by_layer[layer_id] - assert e["global_off"] % ALIGN == 0, (base, layer_id, e["global_off"]) # writer invariant + if e["global_off"] % ALIGN: + reader.close() + raise FTWFormatError( + f"FTW bank {base!r} layer {layer_id} is not {ALIGN}-aligned" + ) bank = HostBank(tuple(e["shape"]), _dtype_of(e["dtype"]), backing=_backing(layer_id)) row_hb[base].append(bank) row_view_args[base].append(None) @@ -737,6 +1121,6 @@ def _read(job) -> torch.Tensor: __all__ = [ "INDEX_NAME", "FORMAT_TAG", "FORMAT_VERSION", "ALIGN", "DEFAULT_SHARD_LIMIT", - "is_ftw_checkpoint", "ftw_tensor_names", "FTWWriter", "FTWReader", + "FTWFormatError", "is_ftw_checkpoint", "ftw_tensor_names", "FTWWriter", "FTWReader", "iter_ftw_weights", "load_ftw_banks", "load_ftw_banks_to_device", "layer_bank_entry_name", ] diff --git a/python/freetoken/core.py b/python/freetoken/core.py index 812f9a4349..a7811ed712 100644 --- a/python/freetoken/core.py +++ b/python/freetoken/core.py @@ -10,8 +10,10 @@ from freetoken.attention import BaseAttnBackend, BaseAttnMetadata from freetoken.attention.linear import FLAMetadata from freetoken.kvcache import BaseCacheHandle, BaseKVCachePool + from freetoken.kvcache.kv_host_offload import KVHostOffloader from freetoken.kvcache.linear_state_pool import LinearStatePool from freetoken.moe.offload_cache import OffloadMoeCache + from freetoken.tokenizer.detokenize import DecodeStatus @dataclass @@ -24,13 +26,49 @@ class SamplingParams: # Stop strings (OpenAI `stop` / Anthropic `stop_sequences`). Generation finishes when one # appears in the decoded output; the matched substring (and anything after) is trimmed. stop_strs: list[str] = field(default_factory=list) + # ---- logits processors (engine/sample.py), all off by default ---- + # min_p: drop tokens whose probability is below min_p x the top probability. + min_p: float = 0.0 + # OpenAI presence/frequency penalties over the generated tokens; repetition_penalty is the + # HF/vLLM multiplicative penalty over prompt + generated tokens (1.0 = off). The sampler + # keeps their state on the device (Req.output_token_counts, Req.prompt_token_mask). presence_penalty: float = 0.0 frequency_penalty: float = 0.0 + repetition_penalty: float = 1.0 + # OpenAI logit_bias as [[token id, bias], ...] (clamped to [-100, 100] by the API). A list, + # not a dict: the messages cross process boundaries as msgpack, which refuses int map keys. + logit_bias: list[list[float]] | None = None + # No EOS / stop token before this many generated tokens (vLLM min_tokens, SGLang + # min_new_tokens). The scheduler fills min_tokens_stop_ids with the ids to mask. + min_tokens: int = 0 + min_tokens_stop_ids: list[int] = field(default_factory=list) + # Extra token ids that end generation with finish_reason "stop" (vLLM stop_token_ids). + stop_token_ids: list[int] = field(default_factory=list) + # Keep the matched stop string in the output instead of trimming it (vLLM + # include_stop_str_in_output, SGLang no_stop_trim). + include_stop_str_in_output: bool = False + # Decode with skip_special_tokens. Off by default here: the reasoning and tool parsers + # consume the /tool tags from the decoded text. + skip_special_tokens: bool = False + # Sampled-token logprobs (OpenAI `logprobs`/`top_logprobs`): when on, the sampler + # reports the chosen token's raw (pre-temperature) logprob and top-k alternatives. + logprobs: bool = False + top_logprobs: int = 0 @property def is_greedy(self) -> bool: return self.temperature <= 0.0 or self.top_k == 1 + @property + def needs_logits_processing(self) -> bool: + # the penalties are not LogitsPlan work: the sampler applies them from its own + # device-side state (Sampler.prepare's penalties). + return bool( + self.min_p > 0.0 + or self.logit_bias + or self.min_tokens > 0 + ) + @dataclass(eq=False) class Req: @@ -69,6 +107,9 @@ class Req: # _process_last_data frees the request when the batch drains (after copy_done.synchronize). aborted: bool = False output_token_counts: torch.Tensor | None = field(default=None, init=False, repr=False) + stop_decode_status: DecodeStatus | None = None + # [vocab] bool on the device: the prompt's ids, for repetition_penalty. Built once. + prompt_token_mask: torch.Tensor | None = field(default=None, init=False, repr=False) def __post_init__(self) -> None: assert self.input_ids.is_cpu @@ -189,6 +230,9 @@ class Context: attn_backend: BaseAttnBackend = field(init=False) moe_offload_cache: OffloadMoeCache | None = None kv_cache: BaseKVCachePool = field(init=False) + # KV host offload (--kv-host-pages): the QSA backend's logical->physical page bridge. + # Set by the engine before create_attention_backend; None when disabled. + kv_offloader: KVHostOffloader | None = None # Per-request recurrent state for GatedDeltaNet layers; set by the engine for # hybrid linear-attention models, otherwise None. linear_state_pool: LinearStatePool | None = None diff --git a/python/freetoken/engine/cache_budget.py b/python/freetoken/engine/cache_budget.py index b5f07ee8ab..9e824f972b 100644 --- a/python/freetoken/engine/cache_budget.py +++ b/python/freetoken/engine/cache_budget.py @@ -14,8 +14,8 @@ import torch -def expert_bytes_per_slot(sources: dict[str, "list[torch.Tensor]"]) -> int: - """Bytes one expert slot occupies on GPU: summed row bytes over all banks. +def layer_slot_bytes(sources: dict[str, "list[torch.Tensor]"]) -> list[int]: + """Bytes one expert slot of each layer occupies on GPU: its row bytes summed over all banks. Each bank source is per-layer ``[num_experts, *row_shape]`` tensors and is already TP-sharded upstream, so the per-row byte count is the per-rank slot @@ -25,7 +25,14 @@ def expert_bytes_per_slot(sources: dict[str, "list[torch.Tensor]"]) -> int: # with cache_size), so they are intentionally excluded from the per-slot growth term. # tensor[0].numel() is the per-row element count (one expert slot); see the matching # slot-byte idiom in kvcache/linear_state_pool.py and kvcache/dsv4/v4_pool.py. - return sum(t[0][0].numel() * t[0].element_size() for t in sources.values()) + num_layers = len(next(iter(sources.values()))) + return [sum(t[l][0].numel() * t[l].element_size() for t in sources.values()) for l in range(num_layers)] + + +def expert_bytes_per_slot(sources: dict[str, "list[torch.Tensor]"]) -> int: + """Bytes one expert slot occupies on GPU, rounded up to the mean over layers when their bank layouts differ.""" + per_layer = layer_slot_bytes(sources) + return -(-sum(per_layer) // len(per_layer)) def net_cache_budget_bytes( @@ -54,6 +61,7 @@ def plan_cache_budget( prefill_overlap: bool, kv_reserve_pages: int, max_slots: int, + page_size: int = 1, ) -> tuple[int, int, bool]: """Split ``budget_bytes`` MoE-first into (moe_cache_size, num_pages, prefill_overlap). @@ -61,7 +69,8 @@ def plan_cache_budget( weights + fixed_cache_size; the (1-memory_ratio) remainder is the graph headroom). Experts greedily fill the budget after reserving ``kv_reserve_pages`` for KV, clamped to ``[floor, min(total_experts, max_slots)]`` (floor is ``2*num_experts`` when prefill - overlap is feasible else ``num_experts``); KV pages take whatever remains. + overlap is feasible, in slots and in bytes, else ``num_experts``); KV pages take whatever + remains. ``page_size`` (tokens per page) only feeds the error message. """ assert per_expert_bytes > 0, "per_expert_bytes must be positive" assert cache_per_page > 0, "cache_per_page must be positive (owned-KV models unsupported here)" @@ -74,6 +83,10 @@ def plan_cache_budget( assert hi >= lo, f"slot cap {hi} below the minimum {lo} slots" kv_reserve_bytes = kv_reserve_pages * cache_per_page + # Overlap is a prefill speedup, not a requirement: when its 2*num_experts floor leaves no + # room for the KV reserve (and the 2 pages asserted below), plan without it. + if overlap and 2 * num_experts * per_expert_bytes + max(kv_reserve_pages, 2) * cache_per_page > budget_bytes: + overlap, lo = False, num_experts # MoE-priority: reserve KV first, then experts greedily take the remaining budget. raw = (budget_bytes - kv_reserve_bytes) // per_expert_bytes moe_cache_size = max(lo, min(raw, hi)) @@ -86,10 +99,13 @@ def plan_cache_budget( # the reserve (or negative), yielding a plan that exceeds budget_bytes. Reject here so # --moe-cache-auto fails in arithmetic instead of OOMing in a later CUDA allocation. total = moe_cache_size * per_expert_bytes + num_pages * cache_per_page + fit_pages = max(0, (budget_bytes - moe_cache_size * per_expert_bytes) // cache_per_page) assert total <= budget_bytes, ( f"cache budget too small: minimum plan (moe={moe_cache_size} slots, " - f"kv={num_pages} pages) needs {total} B > budget {budget_bytes} B " - "(raise memory_ratio, lower kv_reserve_tokens, or free GPU memory)" + f"kv={num_pages} pages) needs {total / 2**30:.2f} GiB > budget {budget_bytes / 2**30:.2f} GiB; " + f"beside {moe_cache_size} slots at most {fit_pages} KV pages ({fit_pages * page_size} tokens) fit. " + "Lower --kv-reserve-tokens, use a smaller --kv-cache-dtype, raise --memory-ratio, " + "or set --moe-cache-size and --max-seq-len-override explicitly" ) assert num_pages > 1, "not enough memory for KV cache after MoE allocation" return moe_cache_size, num_pages, overlap @@ -120,7 +136,9 @@ def resolve_moe_cache_auto( """ budget_bytes = net_cache_budget_bytes(memory_ratio, baseline_free, weights_bytes, fixed_cache_size) max_slots = total_experts if max_slots is None else min(max_slots, total_experts) - kv_reserve_pages = div_ceil(kv_reserve_tokens, page_size) + # Every pool keeps page 0 as an unreachable dummy/sentinel. The CLI floor is expressed in + # usable tokens, so reserve that internal page in addition to the user-visible capacity. + kv_reserve_pages = div_ceil(kv_reserve_tokens, page_size) + 1 return plan_cache_budget( budget_bytes=budget_bytes, per_expert_bytes=per_expert_bytes, @@ -130,4 +148,5 @@ def resolve_moe_cache_auto( prefill_overlap=prefill_overlap, kv_reserve_pages=kv_reserve_pages, max_slots=max_slots, + page_size=page_size, ) diff --git a/python/freetoken/engine/config.py b/python/freetoken/engine/config.py index 4b30fe7fda..19c7436ab7 100644 --- a/python/freetoken/engine/config.py +++ b/python/freetoken/engine/config.py @@ -1,6 +1,7 @@ from __future__ import annotations import copy +import os from dataclasses import dataclass, field, replace from functools import cached_property from typing import TYPE_CHECKING, Any, List, Mapping @@ -37,6 +38,8 @@ class EngineConfig: # DeepSeek-V4.1 Decoder SWA Bounded Replay: "bounded" (default) runs the decoder layers on each # prompt's last window; "exact" runs them on every token. swa_decoder_replay: str = "bounded" + # --embed-weights: "host" keeps the token embedding in pinned host RAM, gathered over PCIe. + embed_weights: str = "gpu" # Expert-bank host load (--expert-load): auto|serial|parallel. "auto" reads scattered # experts in parallel but falls back to serial when free RAM can't cover the banks + the # parallel reader's extra (non-reclaimable) whole-shard buffer; "serial" forces the @@ -45,6 +48,12 @@ class EngineConfig: moe_cache_size: int = 0 moe_cache_rate: float | None = None moe_cache_auto: bool = False + # --moe-cache-mem: slot-cache budget in bytes, turned into a slot count once the expert banks are loaded + moe_cache_mem: int | None = None + # --moe-cache-scope: unified (one LRU pool per expert quant) or per-layer (one pool per layer) + moe_cache_scope: str = "unified" + # --moe-cache-split: how pools share the slots; equal or profile: (by expert use from a profile file) + moe_cache_split: str = "equal" kv_reserve_tokens: int = 8192 # KV floor for --moe-cache-auto; small by design (MoE-priority) moe_cache_policy: str = "lru" moe_prefill_demand_tokens: int = 0 @@ -57,6 +66,8 @@ class EngineConfig: # CPU MoE backend (--moe-strategy cpu): number of CPU worker threads computing # the decode experts. 0 = auto (physical cores). Ignored by other backends. moe_cpu_threads: int = 0 + # skip the startup prefill warmup forwards; first requests then pay the kernel module loads (and the full JIT compile on a cold start) + prefill_warmup: bool = True # Hybrid CPU/GPU decode (--moe-strategy offload only): which MoE layers decode on # the CPU executor instead of the GPU offload/PCIe path. Spec is an explicit id # list ("3,7,11"), a count ("8" -> 8 layers evenly strided across depth), or a @@ -72,22 +83,34 @@ class EngineConfig: cuda_graph_bs: List[int] | None = None cuda_graph_max_bs: int | None = None page_size: int = 1 + # KV-cache storage quantization: "none" stores the compute dtype, "fp8" stores e4m3 + # codes plus one fp32 scale per (token, slab, layer, kv head) -- about 2x the tokens + # per GiB, at a small accuracy cost. --kv-cache-dtype; resolved from "auto" by + # _adjust_config, which also refuses it on a pool family or attention backend that + # cannot read the scales. + kv_quant: str = "none" memory_ratio: float = 0.9 # Hybrid GDN models default to the HybridRadixCache (cross-request GDN-state prefix reuse); # `--cache-type naive` opts out. linear_state_cache_ratio sizes the GDN snapshot cache as # ceil(ratio * max_running_req) extra slots. - linear_state_cache_ratio: float = 2.0 + linear_state_cache_ratio: float = float(os.environ.get("FT_LINEAR_STATE_CACHE_RATIO", "2.0")) # Window/full ratio for the SWA radix cache (`--cache-type radix` on SWA models) and the DSV4 # window tier: the DEFAULT window-pool size = max(working-set floor, ratio x full-pool tokens). # < 1.0 trades retained window-prefix capacity for memory savings; must be in (0, 1]. It is the # DSV4 window/full ratio directly. Used only when swa_num_pages_override is None (a runtime # rebuild can pin an absolute window instead). swa_full_tokens_ratio: float = 0.2 + # Force one decode step after this many consecutive prefill steps, so a long chunked + # prefill cannot starve in-flight decodes. None/0 keeps the historical prefill-first + # order (the scheduler's own TODO names this: "support other policies: e.g. DECODE + # first"). At 8 the cost is ~1% of prefill wall time. + decode_interleave_every: int | None = None # Absolute window-pool size in the pool's own pages (usable, dummy excluded); None -> use the # ratio default above. A runtime cache rebuild sets this (num_swa_pages) to pin the window # regardless of the full anchor; the ratio is the startup default and the fallback. swa_num_pages_override: int | None = None - distributed_timeout: float = 60.0 + distributed_timeout: float = 1800.0 # ranks reach the first collective minutes apart on a 100+ GiB offload load + distributed_port: int = 2333 use_dummy_weight: bool = False use_pynccl: bool = True max_seq_len_override: int | None = None @@ -95,6 +118,11 @@ class EngineConfig: # KV capacity in tokens; resolved into num_page_override by _adjust_config once page_size # is final. Mutually exclusive with num_page_override. num_token_override: int | None = None + # Extra KV pages mirrored to pinned host RAM (QSA paged pools, e.g. Qwen3.8-Flash-Next): + # the GPU pool becomes an LRU cache over the logical page space (scheduler, page table and + # radix tree all see num_pages + kv_host_pages pages), extending context past VRAM + # capacity. 0 (default) = off. + kv_host_pages: int = 0 # Runtime knobs of the multimodal path; the architecture side (vision_config, mrope) lives in ModelConfig. mm: MultimodalConfig = field(default_factory=MultimodalConfig) @@ -142,8 +170,18 @@ def model_config(self) -> ModelConfig: quant = checkpoint_quant_config(self.model_path, hf_config, spec) set_quant_config(quant) model_config = _load_attr(spec.module, spec.parse_config)(hf_config) + self._check_embed_weights(model_config) return replace(model_config, quant=quant) + def _check_embed_weights(self, model_config) -> None: + # the parsed flag, not the hf key: gemma4 ties by default and gguf infers it from the tensors + if self.embed_weights != "host": + return + if model_config.tie_word_embeddings: + raise ValueError("--embed-weights host needs an untied embedding (a tied lm_head reads the table on the GPU)") + if self.tp_info.size > 1: + raise ValueError("--embed-weights host serves a single GPU only (--tensor-parallel-size 1)") + @property def max_seq_len(self) -> int: if self.max_seq_len_override is not None: @@ -156,4 +194,4 @@ def max_forward_len(self) -> int: @property def distributed_addr(self) -> str: - return "tcp://127.0.0.1:2333" + return f"tcp://127.0.0.1:{self.distributed_port}" diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index 5e02fb2ed8..d26141a1eb 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -5,6 +5,7 @@ import gc import math import os +import socket from datetime import timedelta from typing import Any, Dict, Iterable, NamedTuple, Tuple @@ -12,6 +13,7 @@ from freetoken.attention import AttnType, attention_backend_info, create_attention_backend from freetoken.core import Batch, Context, Req, set_global_ctx from freetoken.distributed import destroy_distributed, enable_pynccl_distributed, set_tp_info +from freetoken.env import ENV from freetoken.gpu_select import gpu_identity from freetoken.checkpoint.ftw import ftw_tensor_names, is_ftw_checkpoint from freetoken.layers import OffloadMoELayer, iter_moe_layers, set_rope_device @@ -23,7 +25,8 @@ from freetoken.moe import is_offload_moe_strategy from freetoken.moe.expert_banks import attach_resident_banks, load_expert_banks from freetoken.moe.host_banks import PinFailed -from freetoken.moe.offload_cache import OffloadMoeCache, attach_offload_moe_cache +from freetoken.moe.offload_cache import OffloadMoeCache, attach_offload_moe_cache, bank_layouts +from freetoken.moe.pool_policy import group_layers, make_pools, pool_of_layer, resolve_profile, split_slots from freetoken.utils import align_ceil, init_logger, is_sm90_family, is_sm100_family, mem_GB, torch_dtype from .config import EngineConfig @@ -33,11 +36,35 @@ from freetoken.kvcache.base import CacheRebuildRejected from freetoken.kvcache.cache_status import _supports_swa_ratio from freetoken.kvcache.linear_state_pool import ( - _linear_pool_min_slots, _linear_pool_num_slots, state_pool_bytes, + _linear_pool_min_slots, _linear_pool_num_slots, gdn_prefill_workspace_bytes, + state_pool_bytes, ) logger = init_logger(__name__) +def _clamp_decode_batch_to_pools(config: EngineConfig, pool_sizes: list[int]) -> None: + """Clamp the decode batch so every GPU-decoded layer's routed experts fit in its pool.""" + num_experts = config.model_config.num_experts + small = [s for s in pool_sizes if s < num_experts] + if not small: + return + limit = min(small) // config.model_config.num_experts_per_tok + if config.max_running_req > limit: + logger.warning_rank0( + f"the smallest MoE slot pool has {min(small)} slots; lowering max_running_req " + f"{config.max_running_req} -> {limit} so a decode step fits in it" + ) + object.__setattr__(config, "max_running_req", limit) + if config.cuda_graph_max_bs is None or config.cuda_graph_max_bs > limit: + object.__setattr__(config, "cuda_graph_max_bs", limit) + if config.cuda_graph_bs is not None: + object.__setattr__(config, "cuda_graph_bs", [bs for bs in config.cuda_graph_bs if bs <= limit]) +# Decode steps between --moe-collect-stats reports. Reading the counters costs a host sync, +# so the window is coarse; it is short enough that a normal-length reply still produces a +# few reports rather than one at the very end. +MOE_STATS_INTERVAL = 256 + + def _require_offload_cache_size(cache_size: int, num_experts: int) -> None: """The offload MoE cache needs at least one slot per expert per layer. A too-small size (e.g. a bare offload run with moe_cache_size unset and auto disabled) must fail loudly.""" @@ -50,6 +77,12 @@ def _require_offload_cache_size(cache_size: int, num_experts: int) -> None: f"cache-sizing flags)." ) +# Headroom the KV host offload keeps out of the GPU pool solve: the bigger logical index slab +# is priced exactly, but prefill transients (GDN chunk workspaces, QSA score logits) also grow +# with the logical width and need this slack (a 0.8 GiB post-capture margin OOMed a 4096-token +# prefill; the pre-offload runs sat at ~1.05 GiB). +_KV_OFFLOAD_SERVING_MARGIN = 512 << 20 + def _flashinfer_available() -> bool: from freetoken.kernel.backend import is_flashinfer_installed @@ -74,6 +107,42 @@ def _sgl_flash_attn_available() -> bool: return True +def _place_embeddings_on_host(model) -> int: + """--embed-weights host: every input embedding table moves to pinned host RAM, except one an lm_head is tied to. + + Returns the pinned bytes.""" + from freetoken.layers.base import BaseOP + from freetoken.layers.embedding import ParallelLMHead, VocabParallelEmbedding + + found: list[VocabParallelEmbedding] = [] + tied: set[int] = set() + stack, seen = [model], set() + while stack: + op = stack.pop() + if id(op) in seen: + continue + seen.add(id(op)) + if isinstance(op, ParallelLMHead): + if op.tied_embedding is not None: + tied.add(id(op.tied_embedding)) + continue + if isinstance(op, VocabParallelEmbedding): + found.append(op) + continue + for value in vars(op).values(): + items = value if isinstance(value, (list, tuple)) else (value,) + stack.extend(v for v in items if isinstance(v, BaseOP)) + moved = [e for e in found if id(e) not in tied] + if not moved: + raise ValueError("--embed-weights host: the model has no untied input embedding to move") + for emb in moved: + emb.place_on_host() + torch.cuda.empty_cache() + size = sum(e.weight.numel() * e.weight.element_size() for e in moved) + logger.info_rank0(f"Token embedding in pinned host RAM ({mem_GB(size)}), gathered over PCIe") + return size + + def _startup_kv_budget(memory_ratio: float, init_free_memory: int, new_free_memory: int) -> int: """Bytes available to the KV pool at startup: ratio-scaled pre-load free memory minus what the resident model consumed. Kept as a pure function so the composition with the @@ -120,15 +189,49 @@ def _backend_requirements_met(name: str) -> bool: return False if any(i.requires_sgl_kernel for i in infos) and not _sgl_flash_attn_available(): return False + # requires_sm90 also admits sm_100 (FA4). Hopper-only would reject B200. + if any(i.requires_sm90 for i in infos) and not (is_sm90_family() or is_sm100_family()): + return False if any(i.requires_sm100 for i in infos) and not is_sm100_family(): return False return True -def _resolve_auto_attention_backend(required: frozenset[AttnType]) -> str: +# --kv-cache-dtype spellings -> the stored EngineConfig.kv_quant value. +KV_QUANT_ALIASES = {"auto": "none", "bf16": "none", "none": "none", "fp8": "fp8", "nvfp4": "nvfp4"} + + +def _resolve_kv_quant(value: str | None) -> str: + """Normalize a --kv-cache-dtype spelling to EngineConfig.kv_quant.""" + key = (value or "auto").strip().lower() + if key not in KV_QUANT_ALIASES: + raise ValueError( + f"unknown --kv-cache-dtype {value!r}; expected one of " + f"{', '.join(sorted(KV_QUANT_ALIASES))}" + ) + return KV_QUANT_ALIASES[key] + + +def _backend_supports_kv_quant(name: str, kv_quant: str) -> bool: + """Whether every comma part of an attention-backend string can read a quantized + KV pool (an unquantized pool needs nothing from the backend).""" + if kv_quant == "none": + return True + if kv_quant not in ("fp8", "nvfp4"): + return False + return all( + getattr(attention_backend_info(part.strip()), f"supports_{kv_quant}_kv") + for part in name.split(",") + ) + + +def _resolve_auto_attention_backend( + required: frozenset[AttnType], *, kv_quant: str = "none" +) -> str: """First candidate (in per-type priority order) whose arch condition holds, - whose packages are installed, and whose every comma part serves ALL required - types. Reproduces the historical hardware tree for FULL-only models: + whose packages are installed, whose every comma part serves ALL required + types, and which can decode a quantized KV cache when one is configured. + Reproduces the historical hardware tree for FULL-only models: sm_100 -> trtllm, sm_90+sgl_kernel -> "fa,fi", flashinfer -> fi, else triton.""" candidates: list[tuple[str, bool]] = [] if AttnType.DSV4 in required: @@ -157,10 +260,18 @@ def _resolve_auto_attention_backend(required: frozenset[AttnType]) -> str: continue if not _backend_requirements_met(name): continue + if not _backend_supports_kv_quant(name, kv_quant): + continue return name raise RuntimeError( "No attention backend can serve attention types " - f"{sorted(t.value for t in required)} on this machine." + f"{sorted(t.value for t in required)} on this machine" + + ( + f" with a {kv_quant} KV cache" + if kv_quant != "none" + else "" + ) + + "." ) @@ -206,6 +317,23 @@ def _validate_attention_backend_choice(config, override, required: frozenset[Att f"SWA models require, got {config.attention_backend!r}." ) + # A quantized KV pool is only readable by a backend that applies its per-(token, + # head) scales; one that hands the cache to an external kernel would silently + # attend to raw e4m3 codes. Rejected here, before any weight is resident. + kv_quant = getattr(config, "kv_quant", "none") + if not _backend_supports_kv_quant(config.attention_backend, kv_quant): + fp8_backends = [ + name + for name in ("trtllm", "fi", "fa", "triton") + if required <= attention_backend_info(name).supported_types + and _backend_supports_kv_quant(name, kv_quant) + ] + raise ValueError( + f"--kv-cache-dtype {kv_quant} needs an attention backend that decodes the KV " + f"scales; {config.attention_backend!r} does not. Valid for this model: " + f"{', '.join(fp8_backends) or 'none'} (or use --kv-cache-dtype bf16)." + ) + # An explicitly-selected backend may require a package that isn't installed. Auto # never resolves to one of these when its package is missing, so this only fires for # explicit --attention-backend choices. @@ -233,6 +361,13 @@ def _validate_attention_backend_choice(config, override, required: frozenset[Att "not installed. Install it with `pip install 'freetoken[sgl]'` (or " "'freetoken[accel]'), or use --attention-backend triton." ) + if info.requires_sm90 and not (is_sm90_family() or is_sm100_family()): + raise RuntimeError( + f"Attention backend {config.attention_backend!r} requires a compute capability " + "9.x GPU (FA3, sm_90 cubins) or 10.x GPU (FA4). Ada (sm_89) and consumer " + "Blackwell (sm_12x) are not in the published sgl_kernel flash_ops set. " + "Use --attention-backend fi (or triton) instead." + ) if info.requires_sm100 and not is_sm100_family(): raise RuntimeError( f"Attention backend {config.attention_backend!r} requires a compute capability " @@ -340,6 +475,11 @@ class ForwardOutput(NamedTuple): next_tokens_gpu: torch.Tensor next_tokens_cpu: torch.Tensor copy_done_event: torch.cuda.Event + # Sampled-token logprobs (None unless some request in the batch asked): CPU + # copies covered by copy_done_event, padded to the batch max top_logprobs. + chosen_logprobs_cpu: torch.Tensor | None = None + top_ids_cpu: torch.Tensor | None = None + top_logprobs_cpu: torch.Tensor | None = None class Engine: @@ -372,6 +512,9 @@ def __init__(self, config: EngineConfig): logger.info_rank0(f"Free memory before loading model: {mem_GB(init_free_memory)}") # ======================= Model initialization ======================== + from freetoken.kernel.triton.autotune_cache import bound_autotune_flush_buffer + + bound_autotune_flush_buffer(self.device.index) set_rope_device(self.device) with torch.device("meta"), torch_dtype(config.dtype): self.model = create_model(config.model_config) @@ -386,6 +529,7 @@ def __init__(self, config: EngineConfig): ) # before the residency snapshot, so streamed blocks are not charged as resident weights self.model.place_encoder_weights(config.mm.encoder_weights) + embed_host_bytes = _place_embeddings_on_host(self.model) if config.embed_weights == "host" else 0 post_weights_free = self._sync_get_memory()[0] self._weights_bytes = self._baseline_free - post_weights_free # Pool-budget baseline for the desktop cache sliders: free VRAM after the weights are @@ -395,14 +539,15 @@ def __init__(self, config: EngineConfig): # graphs, or other processes. Cross-rank MIN, deterministic across ranks. self._post_weights_free = post_weights_free self.moe_offload_cache = None + self._moe_stats_step = 0 self.cpu_moe_executor = None # Host-side auxiliary stores (qwen4_exp's pinned PLE table): after the weights so a # load failure is not masked, before the MoE offload cache so the bank residency - # planning sees the pin quota the table already spent. - self._host_tables_bytes = 0 + # planning sees the pin quota the table (and a host embedding) already spent. + self._host_tables_bytes = embed_host_bytes if hasattr(self.model, "load_host_tables"): with _weight_load_context(): - self._host_tables_bytes = int(self.model.load_host_tables(config) or 0) + self._host_tables_bytes += int(self.model.load_host_tables(config) or 0) if is_offload_moe_strategy(config.moe_strategy): self._init_offload_moe_cache(config) if hasattr(self.model, "prepare_for_runtime"): @@ -435,10 +580,38 @@ def __init__(self, config: EngineConfig): # off it; the KV pool family owns every geometry-specific formula behind the rest. available_memory = _startup_kv_budget(config.memory_ratio, init_free_memory, new_free) available_memory -= state_pool_bytes(config) - self.num_pages = self._pool_cls.solve_num_pages(config, available_memory) + # Reserve part of the GDN prefill transient (chunk states + v_new); the rest of the + # prefill activations still live in the (1 - memory_ratio) headroom. + available_memory -= gdn_prefill_workspace_bytes(config) + num_gpu_pages = self._pool_cls.solve_num_pages(config, available_memory) + if config.kv_host_pages > 0: + # KV host offload: the scheduler, page table and radix tree run in a LOGICAL page + # space of num_gpu_pages + kv_host_pages; the GPU K/V buffer stays a + # num_gpu_pages LRU cache over a pinned host mirror (kvcache/kv_host_offload.py). + # The QSA index slab is per LOGICAL page, so reserve its growth before solving + # the GPU page count. + from freetoken.kvcache.qsa_pool import QSAKVCache + + if not issubclass(self._pool_cls, QSAKVCache): + raise ValueError( + f"--kv-host-pages requires the QSA paged KV pool, got {self._pool_cls.__name__}" + ) + spec = next(s for s in config.model_config.kv_cache_group_specs() if s.num_layers > 0) + slab_per_page = ( + (config.page_size // spec.index_ratio) + * spec.num_index_layers * spec.index_head_dim * 2 + ) + # The two-pass solve prices the logical-space slab growth AND hands back the + # serving-margin the growth would otherwise eat (prefill transients need it). + reserve = config.kv_host_pages * slab_per_page + _KV_OFFLOAD_SERVING_MARGIN + num_gpu_pages = self._pool_cls.solve_num_pages(config, available_memory - reserve) + self.num_pages = num_gpu_pages + config.kv_host_pages + else: + self.num_pages = num_gpu_pages num_tokens = self.num_pages * config.page_size self.ctx.kv_cache = self.kv_cache = create_kv_pool( - config, self.num_pages, device=self.device, dtype=self.dtype + config, num_gpu_pages, device=self.device, dtype=self.dtype, + num_index_pages=self.num_pages if config.kv_host_pages > 0 else None, ) # ======================= Linear (GatedDeltaNet) state initialization ======================== @@ -472,6 +645,21 @@ def __init__(self, config: EngineConfig): # through them belongs to the attention backend, built later in init_capture_graph. self.kv_cache.attach_page_table(self.page_table) + # ======================= KV host offload ======================== + # Attached before create_attention_backend: the QSA backend reads ctx.kv_offloader + # at construction. + self.kv_offloader = None + if config.kv_host_pages > 0: + from freetoken.kvcache.kv_host_offload import KVHostOffloader + + self.kv_offloader = KVHostOffloader(self.kv_cache, self.num_pages, self.device) + self.ctx.kv_offloader = self.kv_offloader + logger.info_rank0( + f"Allocating {num_tokens} tokens for KV cache (logical): " + f"{num_gpu_pages} GPU-resident pages + {config.kv_host_pages} host pages, " + f"pinned mirror {mem_GB(self.kv_offloader.mirror_bytes)}" + ) + # ======================= Attention backend initialization ======================== self.ctx.attn_backend = self.attn_backend = create_attention_backend( config.attention_backend, config.model_config @@ -511,19 +699,43 @@ def __init__(self, config: EngineConfig): moe_offload_cache=self.moe_offload_cache, mrope=config.model_config.model_is_mrope, ) - if config.attention_backend.split(",")[0] == "triton": - # Prefill runs on the first comma part; warm its autotune cache. + if config.prefill_warmup: self._warmup_prefill() - def _init_communication(self, config: EngineConfig) -> torch.distributed.ProcessGroup: - if config.tp_info.size == 1 or config.use_pynccl: - torch.distributed.init_process_group( - backend="gloo", - rank=config.tp_info.rank, - world_size=config.tp_info.size, - timeout=timedelta(seconds=config.distributed_timeout), - init_method=config.distributed_addr, + def _make_distributed_store(self, config: EngineConfig) -> torch.distributed.Store: + """The rendezvous store's C10d server ignores the host it's given and always listens + on every interface, so an unauthenticated TCPStore ends up reachable off-box even + when distributed_addr says 127.0.0.1. Rank 0 pre-binds the listening socket to + loopback itself and hands the fd to TCPStore (master_listen_fd) to force that.""" + timeout = timedelta(seconds=config.distributed_timeout) + if not config.tp_info.is_primary(): + return torch.distributed.TCPStore( + "127.0.0.1", config.distributed_port, config.tp_info.size, + is_master=False, timeout=timeout, multi_tenant=True, ) + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + sock.bind(("127.0.0.1", config.distributed_port)) + sock.listen(1024) + self._distributed_listen_addr = sock.getsockname() + # TCPStore takes ownership of the fd and closes it with the store; detach so the + # Python socket never closes that number again after it has been reused. + return torch.distributed.TCPStore( + "127.0.0.1", config.distributed_port, config.tp_info.size, + is_master=True, timeout=timeout, multi_tenant=True, master_listen_fd=sock.detach(), + ) + + def _init_communication(self, config: EngineConfig) -> torch.distributed.ProcessGroup: + store = self._make_distributed_store(config) + use_gloo = config.tp_info.size == 1 or config.use_pynccl + torch.distributed.init_process_group( + backend="gloo" if use_gloo else "nccl", + rank=config.tp_info.rank, + world_size=config.tp_info.size, + timeout=timedelta(seconds=config.distributed_timeout), + store=store, + ) + if use_gloo: tp_cpu_group = torch.distributed.group.WORLD assert tp_cpu_group is not None max_bytes = ( @@ -531,13 +743,6 @@ def _init_communication(self, config: EngineConfig) -> torch.distributed.Process ) enable_pynccl_distributed(config.tp_info, tp_cpu_group, max_bytes) else: - torch.distributed.init_process_group( - backend="nccl", - rank=config.tp_info.rank, - world_size=config.tp_info.size, - timeout=timedelta(seconds=config.distributed_timeout), - init_method=config.distributed_addr, - ) tp_cpu_group = torch.distributed.new_group(backend="gloo") assert tp_cpu_group is not None return tp_cpu_group @@ -623,7 +828,7 @@ def _run_mm_encoder(self, batch: Batch) -> None: if parts: batch.mm_embeds = torch.cat([p.to(self.dtype) for p in parts], dim=0) - def _resolve_auto_moe_cache_size(self, config: EngineConfig, banks, method=None) -> tuple[int, int, bool]: + def _resolve_auto_moe_cache_size(self, config: EngineConfig, banks, method=None, max_slots=None) -> tuple[int, int, bool]: """Resolve --moe-cache-auto into (moe_cache_size, num_pages, prefill_overlap). Pure glue over the Phase-1 budget policy; isolated here so it is unit-testable @@ -633,6 +838,7 @@ def _resolve_auto_moe_cache_size(self, config: EngineConfig, banks, method=None) cache_per_page, fixed_cache_size, page_tokens, min_reserve = self._pool_cls.kv_cost(config) fixed_cache_size += state_pool_bytes(config) # sibling GDN state pool, engine-summed + fixed_cache_size += gdn_prefill_workspace_bytes(config) # GDN prefill transient num_experts = config.model_config.num_experts total_experts = config.model_config.num_moe_layers * num_experts return resolve_moe_cache_auto( @@ -645,9 +851,13 @@ def _resolve_auto_moe_cache_size(self, config: EngineConfig, banks, method=None) num_experts=num_experts, total_experts=total_experts, prefill_overlap=config.moe_prefill_overlap, - kv_reserve_tokens=max(config.kv_reserve_tokens, min_reserve), + kv_reserve_tokens=max( + config.kv_reserve_tokens, + min_reserve, + (config.num_page_override or 0) * page_tokens, + ), page_size=page_tokens, - max_slots=method.slot_limit() if method is not None else None, + max_slots=max_slots if max_slots is not None or method is None else method.slot_limit(), ) def _init_offload_moe_cache(self, config: EngineConfig) -> OffloadMoeCache: @@ -703,6 +913,11 @@ def _init_offload_moe_cache(self, config: EngineConfig) -> OffloadMoeCache: else HostResidency.PINNED.value for i in range(config.model_config.num_moe_layers) ] + # read before the expert load, so a bad profile fails in seconds + expert_counts = resolve_profile( + config.moe_cache_split, config.model_path, + config.model_config.num_moe_layers, config.model_config.num_experts, + ) try: with _weight_load_context(): banks = load_expert_banks( @@ -718,8 +933,32 @@ def _init_offload_moe_cache(self, config: EngineConfig) -> OffloadMoeCache: ) except PinFailed as exc: raise RuntimeError(f"{exc}; {_pin_hint(self._host_tables_bytes)}") from exc + from freetoken.engine.cache_budget import expert_bytes_per_slot, layer_slot_bytes + + self._moe_layer_slot_bytes = layer_slot_bytes(banks.sources) + self._moe_pool_groups = group_layers(config.moe_cache_scope, bank_layouts(banks.sources, banks.layer_tags), self._moe_layer_slot_bytes) + if len(self._moe_pool_groups) > 1 and cpu_layer_ids: + raise ValueError( + "several MoE slot pools (--moe-cache-scope per-layer, or experts in several quant formats) need every MoE " + "layer to decode from the GPU slot cache; they do not combine with --moe-strategy cpu or --moe-cpu-layers" + ) + self._moe_expert_counts = expert_counts if len(self._moe_pool_groups) > 1 else None + if len(self._moe_pool_groups) == 1 and config.moe_cache_split != "equal": + logger.warning_rank0( + f"--moe-cache-split {config.moe_cache_split} is ignored: all MoE layers share one slot pool " + f"(--moe-cache-scope per-layer gives each layer its own)" + ) + self._moe_slot_limit = method.slot_limit() if method is not None else None + # with GPU decode, a pool below num_experts slots bounds the batch; the startup split keeps room for one token per step + pools = self._moe_pools(1) + size_cap = sum(p.max_num_slots for p in pools) if config.moe_cache_auto: - size, pages, overlap = self._resolve_auto_moe_cache_size(config, banks, method) + size, pages, overlap = self._resolve_auto_moe_cache_size(config, banks, method, max_slots=size_cap) + if config.moe_prefill_overlap and not overlap: + logger.info_rank0( + f"--moe-cache-auto: prefill overlap disabled, its {2 * config.model_config.num_experts} " + "slot floor does not fit beside the KV reserve" + ) object.__setattr__(config, "moe_cache_size", size) object.__setattr__(config, "moe_prefill_overlap", overlap) if config.num_page_override is None: @@ -735,7 +974,28 @@ def _init_offload_moe_cache(self, config: EngineConfig) -> OffloadMoeCache: f"--moe-cache-auto resolved moe_cache_size={size} " f"num_pages={pages} (prefill_overlap={overlap})" ) + elif config.moe_cache_mem is not None: + per_slot = expert_bytes_per_slot(banks.sources) + size = min(config.moe_cache_mem // per_slot, size_cap) + # per_slot is a mean over layers, so with pools of different slot sizes the split can exceed the budget: shrink until it fits + while size > 0 and self._moe_pool_bytes(split_slots(size, pools, self._moe_expert_counts)) > config.moe_cache_mem: + size -= 1 + overlap = config.moe_prefill_overlap and size >= 2 * config.model_config.num_experts + object.__setattr__(config, "moe_cache_size", size) + object.__setattr__(config, "moe_prefill_overlap", overlap) + logger.info_rank0( + f"--moe-cache-mem {mem_GB(config.moe_cache_mem)} resolved moe_cache_size={size} " + f"({per_slot / 2**20:.2f} MiB per slot, prefill_overlap={overlap})" + ) _require_offload_cache_size(config.moe_cache_size, config.model_config.num_experts) + pool_sizes = split_slots(config.moe_cache_size, pools, self._moe_expert_counts) + if len(pools) > 1: + logger.info_rank0( + f"MoE slot pools: {config.moe_cache_scope}, {len(pools)} pools of {min(pool_sizes)}-{max(pool_sizes)} " + f"slots ({config.moe_cache_split})" + ) + if decode_target == "gpu": + _clamp_decode_batch_to_pools(config, pool_sizes) layout = max_slots = None if method is not None: if banks.kind is not None and (banks.kind, banks.kernel) != (method.kind, method.kernel.name): @@ -761,24 +1021,44 @@ def _init_offload_moe_cache(self, config: EngineConfig) -> OffloadMoeCache: hybrid_max_fetch=config.moe_hybrid_max_fetch, layout=layout, max_slots=max_slots, + pool_of_layer=pool_of_layer(pools, config.model_config.num_moe_layers), + pool_sizes=pool_sizes, ) # before set_bank_sources: the residency validation and the copy plan's skip of non-pinned layers key on the CPU-layer set cache.cpu_layer_ids = cpu_layer_ids - cache.set_bank_sources(banks.sources, layer_residency=banks.layer_residency) + cache.set_bank_sources(banks.sources, layer_residency=banks.layer_residency, layer_tags=banks.layer_tags) cache.set_alphas(banks.gate_up_alpha, banks.down_alpha) if decode_target == "hybrid": self._resolve_hybrid_fetch(config, cache) # Must be set before CUDA graph capture so the (device-side) accumulation ops are # captured and re-run on every decode replay. cache.collect_stats = config.moe_collect_stats + # The routing histogram rides the same switch: on its own the miss rate says how + # often we fetch, but not whether a smarter policy could have avoided the fetch. + # decode_routing_stats turns it into routing-skew numbers, including the hit rate of + # the best fixed expert set per layer (not a ceiling: LRU's temporal locality can beat it). + cache.collect_decode_freq = config.moe_collect_stats + # attach_offload_moe_cache walks for OffloadMoELayers, or defers to a model's + # _iter_offload_moe_layers() hook when its MoE blocks are bespoke nn.Modules (DSV4). layers = attach_offload_moe_cache(self.model, cache) assert len(layers) == config.model_config.num_moe_layers + cache.zero_inactive_route_slots = cache.decode_target == "hybrid" and not all(layer.skips_inactive_slots for layer in layers) if cache.decode_target in ("cpu", "hybrid"): self._init_cpu_moe_executor(config, cache, layers) self.ctx.moe_offload_cache = cache self.moe_offload_cache = cache return cache + def _moe_pool_bytes(self, pool_sizes: list[int]) -> int: + return sum(s * self._moe_layer_slot_bytes[g[0]] for g, s in zip(self._moe_pool_groups, pool_sizes)) + + def _moe_pools(self, batch: int) -> list: + """The slot pools, each with room for the routed experts of ``batch`` decode tokens.""" + model_config = self.config.model_config + num_experts = model_config.num_experts + min_num_slots = min(num_experts, model_config.num_experts_per_tok * batch) + return make_pools(self._moe_pool_groups, num_experts, min_num_slots, self._moe_slot_limit) + def _resolve_hybrid_fetch(self, config: EngineConfig, cache) -> None: """Resolve --moe-hybrid-max-fetch -1 (auto) into a bandwidth-matched fetch fraction. @@ -814,9 +1094,8 @@ def _init_cpu_moe_executor(self, config: EngineConfig, cache, layers) -> None: """Build the persistent CPU MoE executor (decode-time expert compute). Must run before CUDA graph capture: the worker pool has to be live for the - eager warmup forward, and the pinned IO buffers / host-func task pointers - must be stable for the captured nodes. Buffers/tasks themselves are - allocated lazily on the first (eager) forward at each batch size. + eager warmup forward, and the pinned IO block must be stable for the captured + nodes. """ from freetoken.moe.cpu_executor import CpuMoeExecutor @@ -829,7 +1108,7 @@ def _init_cpu_moe_executor(self, config: EngineConfig, cache, layers) -> None: ) # Decode batches never exceed max_running_req, but CUDA-graph padding can # round a batch up to the largest captured size; cover both. - max_tokens = max(config.max_running_req, config.cuda_graph_max_bs or 0, 1) + max_tokens = max(config.max_running_req, config.cuda_graph_max_bs or 0, *(config.cuda_graph_bs or []), 1) method = sample.quant_method executor = CpuMoeExecutor( cache, @@ -869,19 +1148,18 @@ def _sync_get_memory(self) -> Tuple[int, int]: return min_free_memory, max_free_memory - def _target_moe_and_expert_bytes(self, moe_cache_size: int | None) -> tuple[int, int]: - from freetoken.engine.cache_budget import expert_bytes_per_slot - - target_moe = ( - moe_cache_size - if moe_cache_size is not None - else (self.moe_offload_cache.cache_size if self.moe_offload_cache else 0) - ) - per_expert_bytes = ( - expert_bytes_per_slot(self.moe_offload_cache.bank_sources) - if self.moe_offload_cache is not None else 0 - ) - return target_moe, per_expert_bytes + def _target_moe_and_expert_bytes(self, moe_cache_size: int | None, pool_sizes: list[int] | None) -> tuple[int, int]: + cache = self.moe_offload_cache + if cache is None: + return moe_cache_size or 0, 0 + if moe_cache_size is None: + target_moe = cache.cache_size + nbytes = sum(c.numel() * c.element_size() for _, c in cache.banks) + else: + target_moe = moe_cache_size + nbytes = self._moe_pool_bytes(pool_sizes) + # pools of several bank layouts have different slot sizes; round up so target * bytes covers them + return target_moe, -(-nbytes // target_moe) def _resize_kv_pool(self, config, num_pages: int, num_swa_pages: int | None) -> None: # IN-PLACE, identity-preserving: the CacheManager's swa_pool reference, ctx.kv_cache and @@ -940,11 +1218,22 @@ def rebuild_runtime_cache( "moe_cache_size requested but this model has no MoE offload cache" ) try: - self.moe_offload_cache.validate_rebuild(moe_cache_size) + # a rebuild does not lower max_running_req again, so with GPU decode every pool must hold a full batch; hybrid caps its fetches instead + batch = config.max_running_req if self.moe_offload_cache.decode_target == "gpu" else 1 + moe_pool_sizes = split_slots(moe_cache_size, self._moe_pools(batch), self._moe_expert_counts) + self.moe_offload_cache.validate_rebuild(moe_cache_size, moe_pool_sizes) except ValueError as e: raise CacheRebuildRejected(str(e)) from e if num_pages is not None and num_pages <= 0: raise CacheRebuildRejected(f"num_pages must be positive, got {num_pages}") + if num_pages is not None and self.kv_offloader is not None: + # The logical/physical split (host mirror, LRU maps, QSA index slab sized to the + # logical space) is built once at startup; a runtime resize would have to rebuild + # all of it. Reject recoverably instead. + raise CacheRebuildRejected( + "num_pages rebuild is not supported with --kv-host-pages (restart with a " + "different --kv-host-pages / --memory-ratio instead)" + ) if num_mamba_slots is not None: if self.linear_state_pool is None: raise CacheRebuildRejected( @@ -977,7 +1266,9 @@ def rebuild_runtime_cache( # must reject (recoverable) so the old caches stay intact and serving continues, # rather than freeing and then OOMing into permanent failure. The engine supplies # the memory account; the pool answers whether its target geometry fits. - target_moe, per_expert_bytes = self._target_moe_and_expert_bytes(moe_cache_size) + target_moe, per_expert_bytes = self._target_moe_and_expert_bytes( + moe_cache_size, moe_pool_sizes if moe_cache_size is not None else None + ) # Price the sibling GDN state pool at ITS target (physical slots = usable + padding # sink) and hand the bytes in -- the KV pool only budgets its own tiers. target_mamba = ( @@ -991,7 +1282,8 @@ def rebuild_runtime_cache( per_expert_bytes=per_expert_bytes, baseline_free=self._baseline_free, weights_bytes=self._weights_bytes, current_num_pages=self.num_pages, extra_fixed_bytes=( - state_pool_bytes(config, target_mamba) if target_mamba is not None else 0 + (state_pool_bytes(config, target_mamba) if target_mamba is not None else 0) + + gdn_prefill_workspace_bytes(config) ), extra_note=( f", mamba={target_mamba - 1} slots" if target_mamba is not None else "" @@ -1012,29 +1304,54 @@ def rebuild_runtime_cache( # 1. Tear down CUDA graphs + backend capture scratch (free-before-alloc). self.attn_backend.reset_capture() self.graph_runner.destroy_cuda_graphs() - # 2. Resize caches in place (each frees its old GPU tensors before allocating). + # 2. Resize caches in place (each frees its old GPU tensors before allocating), + # shrinking resizes before growing ones: a pool's grow must not run while a + # still-large sibling is resident, or a target that passed the fit check above can + # OOM mid-rebuild (#643). With shrinks first the transient peak never exceeds + # max(current, target) totals. Footprints are compared BEFORE the window-override + # write below (the old side must price the pre-override config). + kv_old_bytes = type(self.kv_cache).rebuild_footprint(config, self.num_pages, None) + kv_new_bytes = type(self.kv_cache).rebuild_footprint( + config, num_pages if num_pages is not None else self.num_pages, num_swa_pages + ) + + def _resize_moe() -> None: + assert self.moe_offload_cache is not None, "no MoE offload cache to resize" + self.moe_offload_cache.rebuild(moe_cache_size, moe_pool_sizes) + + def _resize_kv() -> None: + # sets self.num_pages (rebuilds KV + window); a window-only change passes the + # CURRENT page count so only the window pool is re-derived + self._resize_kv_pool( + config, num_pages if num_pages is not None else self.num_pages, num_swa_pages + ) + + def _resize_mamba() -> None: + # Like every resize here it must sit between graph teardown and re-capture so + # the recaptured graphs bind the new state tensors. +1 for the reserved + # padding sink: num_mamba_slots is the usable count. + self.linear_state_pool.rebuild(num_mamba_slots + 1) + + resize_ops = [ + (moe_cache_size is not None, + moe_cache_size is None or moe_cache_size <= self.moe_offload_cache.cache_size, + _resize_moe), + (num_pages is not None or num_swa_pages is not None, + kv_new_bytes <= kv_old_bytes, _resize_kv), + (num_mamba_slots is not None, + num_mamba_slots is None or num_mamba_slots + 1 <= self.linear_state_pool.num_slots, + _resize_mamba), + ] # Pin the new window first (validated above) so any KV-pool rebuild below sizes the window # to it (_dsv4_pool_sizes / _swa_paged_num_tokens read config.swa_num_pages_override). # frozen EngineConfig — mutate in place like the moe_cache_size path; `config.x = y` raises # FrozenInstanceError, which here aborts the rebuild after the CUDA graphs are gone (→ 503). if num_swa_pages is not None: object.__setattr__(config, "swa_num_pages_override", num_swa_pages) - if moe_cache_size is not None: - assert self.moe_offload_cache is not None, "no MoE offload cache to resize" - self.moe_offload_cache.rebuild(moe_cache_size) - if num_pages is not None: - # sets self.num_pages (rebuilds KV + window) - self._resize_kv_pool(config, num_pages, num_swa_pages) - elif num_swa_pages is not None: - # Window-only change: no page-count change, but re-derive the window pool at the new - # pin against the CURRENT page count. This re-allocs the same-size full pool and - # the resized window, both inside the pool's own rebuild_from_config. - self._resize_kv_pool(config, self.num_pages, num_swa_pages) - if num_mamba_slots is not None: - # Reallocate the GDN state pool (frees old tensors first). Must sit between graph - # teardown and re-capture so the recaptured graphs bind the new state tensors. - # +1 for the reserved padding sink: num_mamba_slots is the usable count. - self.linear_state_pool.rebuild(num_mamba_slots + 1) + for shrinking_phase in (True, False): + for requested, shrinks, op in resize_ops: + if requested and shrinks == shrinking_phase: + op() # 3. Refresh max_seq_len (+ generic page table) for the new token budget. self._refresh_seq_state(config) aligned_max_seq_len = _page_table_width(self.max_seq_len, config.page_size) @@ -1064,6 +1381,9 @@ def forward_batch(self, batch: Batch, args: BatchSamplingArgs) -> ForwardOutput: use_graph = self.graph_runner.can_use_cuda_graph(batch) with self.ctx.forward_batch(batch), self.model.forward_host_ctx(batch, use_graph): logits = self.graph_runner.replay(batch) if use_graph else self.model.forward() + audit = self.moe_offload_cache.audit if self.moe_offload_cache is not None else None + if audit is not None and batch.is_decode: + audit.end_decode_step(self.moe_offload_cache) # every side stream has joined the compute stream if self.cpu_moe_executor is not None: # One pinned read: surfaces a fired flag-handshake watchdog (dead coordinator # -> stale expert outputs) as a loud error instead of silent corruption. @@ -1075,9 +1395,147 @@ def forward_batch(self, batch: Batch, args: BatchSamplingArgs) -> ForwardOutput: batch_logits = logits[: batch.size] next_tokens_gpu = self.sampler.sample(batch_logits, args).to(torch.int32) next_tokens_cpu = next_tokens_gpu.to("cpu", non_blocking=True) + logprobs_out = self.sampler.compute_logprobs(batch_logits, next_tokens_gpu, args) copy_done_event = torch.cuda.Event() copy_done_event.record(self.stream) - return ForwardOutput(next_tokens_gpu, next_tokens_cpu, copy_done_event) + cache = self.moe_offload_cache + if cache is not None and batch.is_decode and ( + cache.collect_stats or cache.prefetch is not None or cache.verify is not None or cache.audit is not None + ): + self._moe_stats_step += 1 + if self._moe_stats_step >= MOE_STATS_INTERVAL: + self._moe_stats_step = 0 + if cache.prefetch is not None: + self._emit_prefetch_stats() + if cache.verify is not None: + self._emit_verify_stats() + if cache.audit is not None: + self._emit_audit_stats() + if cache.collect_stats: + self._emit_moe_stats() + if logprobs_out is None: + return ForwardOutput(next_tokens_gpu, next_tokens_cpu, copy_done_event) + chosen_logprobs, top_ids, top_logprobs = logprobs_out + return ForwardOutput( + next_tokens_gpu, next_tokens_cpu, copy_done_event, + chosen_logprobs_cpu=chosen_logprobs, top_ids_cpu=top_ids, + top_logprobs_cpu=top_logprobs, + ) + + def _warmup_prefill_lens(self) -> list[int]: + """Prefill lengths that cross every size bucket the in-repo Triton prefill kernels specialize on.""" + cap = min(self.max_seq_len, self.config.max_forward_len) + if ENV.WARMUP_MAX_LEN.value > 0: + cap = min(cap, ENV.WARMUP_MAX_LEN.value) + mc = self.config.model_config + # 16/32/64 BLOCK_M ladders (nvfp4_linear, fused MoE) plus an odd twin each: triton re-specializes int args on % 16 == 0 + lens = {9, 16, 17, 32, 33, 48, 65, 80} + if mc.is_moe: + top_k = max(1, mc.num_experts_per_tok) + # moe_align leaves its single-CTA path at numel = tokens * top_k > 1024 + lens.add(1024 // top_k + 1) + import triton + + from freetoken.kernel.backend import is_sgl_kernel_installed + + if not is_sgl_kernel_installed(): + # the single-CTA path keys on (next_pow2(numel), num_warps); walk each cell once + seen = set() + for tokens in range(2, 1024 // top_k + 1): + numel = tokens * top_k + cell = (triton.next_power_of_2(numel), triton.next_power_of_2(min(16, max(2, numel // 32)))) + if cell not in seen: + seen.add(cell) + lens.add(tokens) + if cap >= 4096: + # act_and_mul flips BLOCK_D at M >= 4096; also GDN chunk_delta_h NT bucket 1 (NT > 32) + lens.add(4096) + elif mc.has_linear_attention and cap >= 2049: + lens.add(2049) + if mc.has_linear_attention and cap >= 8193: + lens.add(8193) # chunk_delta_h NT bucket 2 (NT > 128) + return sorted(n for n in lens if 2 <= n <= cap) + + def _emit_prefetch_stats(self) -> None: + """Report one window of the FREETOKEN_MOE_PREFETCH counters (one host sync), then restart them.""" + from freetoken.moe.prefetch import format_per_layer, format_summary, summarize + + prefetch = self.moe_offload_cache.prefetch + window = prefetch.take_window() + stats = summarize(window, prefetch.mode) + if not stats["layer_calls"]: + return + logger.info_rank0(f"MoE prefetch {prefetch.mode} ({MOE_STATS_INTERVAL} decode steps): {format_summary(stats)}") + if ENV.MOE_PREFETCH_DEBUG: + logger.info_rank0(f"MoE prefetch per layer (useful/issued/misses per call): {format_per_layer(window)}") + + def _emit_verify_stats(self) -> None: + """Report one window of the FREETOKEN_MOE_PREFETCH_VERIFY counts and the new bad records (host syncs).""" + _log_lines(self.moe_offload_cache.verify.report_window(MOE_STATS_INTERVAL)) + + def _emit_audit_stats(self) -> None: + """Report one window of FREETOKEN_MOE_SLOT_AUDIT totals and the new bad slots (host syncs).""" + _log_lines(self.moe_offload_cache.audit.report_window(MOE_STATS_INTERVAL)) + + def _log_prefetch_totals(self) -> None: + from freetoken.moe.prefetch import format_per_layer, format_summary, summarize + + prefetch = self.moe_offload_cache.prefetch if self.moe_offload_cache is not None else None + if prefetch is None: + return + totals = prefetch.totals + prefetch.stats.cpu() + stats = summarize(totals, prefetch.mode) + if stats["layer_calls"]: + logger.info_rank0(f"MoE prefetch {prefetch.mode} (session, {stats['tokens']} tokens): {format_summary(stats)}") + logger.info_rank0(f"MoE prefetch per layer (useful/issued/misses per call): {format_per_layer(totals)}") + + def _emit_moe_stats(self) -> None: + """Report one window of expert-cache behaviour, then reset the miss counters. + + Accumulation is device-side and captured into the decode graph, so it is free to + leave running; reading it is not (the counters have to come back to the host), which + is why this only fires every MOE_STATS_INTERVAL decode steps. The routing histogram + is deliberately *not* reset -- it describes the whole run's distribution, not one + window's, so the static top-k figure flattens as a session covers more topics. + """ + cache = self.moe_offload_cache + agg = cache.decode_miss_stats() + if not agg["layer_calls"]: + return + parts = [ + f"miss_rate={agg['miss_rate']:.3f}", + f"active/layer={agg['active_per_layer']:.1f}", + f"missing/layer={agg['missing_per_layer']:.1f}", + ] + if cache.decode_target == "hybrid": + # How the misses split: PCIe-fetched to the GPU vs handed to the CPU kernels. + parts.append(f"fetch_rate={agg['fetch_rate']:.3f}") + parts.append(f"cpu/layer={agg['cpu_per_layer']:.1f}") + logger.info_rank0(f"MoE cache ({MOE_STATS_INTERVAL} decode steps): " + ", ".join(parts)) + + per_layer = [L for L in cache.decode_miss_stats_per_layer()["per_layer"] if L["steps"]] + worst = sorted(per_layer, key=lambda L: -L["miss_rate"])[:5] + if worst: + logger.info_rank0( + "MoE cache worst layers: " + + ", ".join(f"L{L['layer']}={L['miss_rate']:.3f}" for L in worst) + ) + routing = cache.decode_routing_stats() + if routing: + # static_topk_hit is what pinning each layer's most frequent experts would score over + # the whole run. A realized hit rate above it means the cache is winning on temporal + # locality; it says nothing about how close to optimal eviction is. + logger.info_rank0( + "MoE routing: " + f"static_topk_hit={routing['static_topk_hit_at_slots']:.3f} " + f"(realized {1.0 - agg['miss_rate']:.3f}), " + f"slots/layer={routing['slots_per_layer']:.1f}, " + f"working_set={routing['working_set_mean']:.1f}" + f"/{routing['working_set_max']}, " + f"experts_for_90pct={routing['experts_for_90pct']:.1f}, " + f"norm_entropy={routing['norm_entropy']:.3f}" + ) + cache.reset_stats() @torch.inference_mode() def _warmup_prefill(self) -> None: @@ -1088,18 +1546,18 @@ def _warmup_prefill(self) -> None: restore it afterwards so padded decode graph replay keeps using the dedicated dummy KV slot. """ - if self.max_seq_len < 2: + if AttnType.DSV4 in _required_attn_types(self.config.model_config): + # DSV4 prefill reads pool.full_loc_map, which only the scheduler's paged allocation fills + logger.info_rank0("Prefill warmup skipped for DSV4.") return - - warmup_lens = [min(80, self.max_seq_len)] - if self.max_seq_len >= 128: - warmup_lens.append(128) - warmup_lens = sorted({length for length in warmup_lens if length >= 2}) + warmup_lens = self._warmup_prefill_lens() if not warmup_lens: return + logger.info_rank0(f"Prefill warmup over {len(warmup_lens)} lengths: {warmup_lens}") dummy_row = self.page_table[self.dummy_req.table_idx] dummy_slot = int(dummy_row[0].item()) + done: list[int] = [] started = torch.cuda.Event(enable_timing=True) ended = torch.cuda.Event(enable_timing=True) started.record(self.stream) @@ -1117,6 +1575,9 @@ def _warmup_prefill(self) -> None: sampling_params=None, # type: ignore[arg-type] cache_handle=None, # type: ignore[arg-type] ) + if self.linear_state_pool is not None: + # like dummy_req: GDN state writes land in the padding sink, not a real slot + warm_req.linear_slot_idx = self.linear_state_pool.padding_slot batch = Batch(reqs=[warm_req], phase="prefill") batch.padded_reqs = batch.reqs batch.input_ids = torch.zeros(length, dtype=torch.int32, device=self.device) @@ -1129,23 +1590,53 @@ def _warmup_prefill(self) -> None: self.attn_backend.prepare_metadata(batch) with self.ctx.forward_batch(batch): self.model.forward() + # surface async faults here, at the failing length, not past the except + torch.cuda.synchronize(self.device) + done.append(length) + except Exception as exc: + if self.config.tp_info.size > 1: + raise + logger.warning_rank0(f"Prefill warmup stopped after {done}: {type(exc).__name__}: {exc}") finally: dummy_row.fill_(dummy_slot) if self.moe_offload_cache is not None: self.moe_offload_cache.reset() + if len(done) != len(warmup_lens): + return ended.record(self.stream) torch.cuda.synchronize(self.device) logger.info_rank0( - f"Prefill warmup complete for lengths {warmup_lens} " + f"Prefill warmup complete for lengths {done} " f"in {started.elapsed_time(ended) / 1000.0:.3f} s" ) def shutdown(self) -> None: + try: + self._log_prefetch_totals() + except Exception as exc: # noqa: BLE001 -- a diagnostic must never block shutdown + logger.warning(f"MoE prefetch totals unavailable at shutdown: {exc}") + try: + verify = self.moe_offload_cache.verify if self.moe_offload_cache is not None else None + if verify is not None: + _log_lines(verify.report_session()) + except Exception as exc: # noqa: BLE001 + logger.warning(f"MoE verify totals unavailable at shutdown: {exc}") + try: + audit = self.moe_offload_cache.audit if self.moe_offload_cache is not None else None + if audit is not None: + _log_lines(audit.report_session()) + except Exception as exc: # noqa: BLE001 + logger.warning(f"MoE slot audit totals unavailable at shutdown: {exc}") self.graph_runner.destroy_cuda_graphs() torch.distributed.destroy_process_group() destroy_distributed() +def _log_lines(lines: list[tuple[str, str]]) -> None: + for level, line in lines: + (logger.warning_rank0 if level == "warning" else logger.info_rank0)(line) + + def _profile_gpu(index: "int | None" = None) -> Tuple[str | None, str | None]: """(name, uuid) of visible device ``index`` (default: the current, i.e. bound, device); (None, None) without CUDA.""" if not torch.cuda.is_available(): @@ -1366,10 +1857,10 @@ def _cpu_moe_executor_viable(model_config) -> bool: """Whether an automatic CPU-decode decision may target the CPU MoE executor. A default boot must degrade to GPU offload instead of crashing in CpuMoeExecutor after the whole load; explicit cpu/hybrid/--moe-cpu-layers picks still fail loudly.""" - from freetoken.moe.cpu_executor import _WFMT_IDS, compiled_extension_supports + from freetoken.moe.cpu_executor import compiled_extension_supports try: - from freetoken.kernel import _cpu_moe # noqa: F401 + from freetoken.kernel import _cpu_moe except ImportError: return False act = getattr(model_config, "hidden_act", "silu") @@ -1380,7 +1871,7 @@ def _cpu_moe_executor_viable(model_config) -> bool: return False expert_quant = getattr(model_config, "expert_quant", "none") fmt = expert_quant if expert_quant != "none" else (moe_wfmt or "bf16") - return fmt == "mxfp4" or fmt in _WFMT_IDS + return fmt == "mxfp4" or fmt in _cpu_moe.weight_formats() def _pin_budget_bytes(reserved: int = 0) -> int | None: @@ -1457,6 +1948,9 @@ def _auto_cpu_layers(config: EngineConfig, num_moe_layers: int, *, reserved: int "moe_cache_size": 0, "moe_cache_rate": None, "moe_cache_auto": False, + "moe_cache_mem": None, + "moe_cache_scope": "unified", + "moe_cache_split": "equal", "moe_cpu_layers": None, "moe_cpu_threads": 0, "moe_hybrid_max_fetch": -1, @@ -1604,6 +2098,34 @@ def override(attr: str, value: Any): # this is dangerous, use with caution # lists, then validate whatever is now selected (explicit or auto) -- every # comma part must serve every required type, with packages/arch available. required_attn_types = _required_attn_types(model_config) + # Resolve KV quantization BEFORE the backend tree: a quantized pool narrows both + # which pool families are usable and which backend auto may pick. + kv_quant = _resolve_kv_quant(getattr(config, "kv_quant", "none")) + override("kv_quant", kv_quant) + if kv_quant == "nvfp4": + if required_attn_types - {AttnType.FULL, AttnType.SWA, AttnType.QSA, AttnType.MLA, AttnType.DSA}: + raise ValueError( + "--kv-cache-dtype nvfp4 requires a paged FULL, hybrid-SWA, QSA, or MLA/DSA KV pool" + ) + for spec in model_config.kv_cache_group_specs(): + if spec.head_dim % 16: + raise ValueError("--kv-cache-dtype nvfp4 requires head_dim divisible by 16") + if kv_quant != "none": + # Quantized codes are wired through the pools that hand their rows to a Triton + # kernel: plain paged and hybrid-SWA, QSA sparse, and DSA/MLA. QSA's index + # tier and DSA's index-key/tail tiers stay bf16; the DSA kernel dequantizes + # selected latent rows with their per-token scale. Other sparse families have + # no scale-read path and remain rejected before weights load. + quant_unsupported = required_attn_types - { + AttnType.FULL, AttnType.SWA, AttnType.QSA, AttnType.MLA, AttnType.DSA, + } + if quant_unsupported: + raise ValueError( + f"--kv-cache-dtype {kv_quant} is implemented for the plain paged, " + "hybrid-SWA, QSA sparse and DSA/MLA KV pools; this model also needs " + f"{', '.join(sorted(t.value for t in quant_unsupported))} attention " + "(use --kv-cache-dtype bf16)." + ) _dtype = getattr(config, "dtype", None) # duck-typed test configs omit it if ( required_attn_types & {AttnType.BSA, AttnType.QSA} @@ -1631,7 +2153,7 @@ def override(attr: str, value: Any): # this is dangerous, use with caution if config.attention_backend == "auto": override( "attention_backend", - _resolve_auto_attention_backend(required_attn_types), + _resolve_auto_attention_backend(required_attn_types, kv_quant=kv_quant), ) logger.info_rank0(f"Auto-selected attention backend: {config.attention_backend}") _validate_attention_backend_choice(config, override, required_attn_types) @@ -1753,6 +2275,7 @@ def override(attr: str, value: Any): # this is dangerous, use with caution is_offload_moe_strategy(config.moe_strategy) and config.moe_cache_size <= 0 and config.moe_cache_rate is None + and config.moe_cache_mem is None and not getattr(config, "moe_cache_auto", False) ): # args.py's "no sizing flag -> default --moe-cache-auto" only fires when the @@ -1786,11 +2309,13 @@ def override(attr: str, value: Any): # this is dangerous, use with caution # An explicit 'fused' keeps the experts resident, so there is no slot cache to size. The # sizing flags no longer redirect the backend, so ignore them here and say so -- the # geometry the user asked for is what runs. Report the flag actually passed: --moe-cache- - # rate was already folded into moe_cache_size above, and the three are mutually exclusive. + # rate was already folded into moe_cache_size above, and the sizing flags are mutually exclusive. if config.moe_cache_rate: inert = f"--moe-cache-rate={config.moe_cache_rate}" elif config.moe_cache_size: inert = f"--moe-cache-size={config.moe_cache_size}" + elif config.moe_cache_mem is not None: + inert = f"--moe-cache-mem={config.moe_cache_mem}" elif getattr(config, "moe_cache_auto", False): inert = "--moe-cache-auto" else: @@ -1802,6 +2327,7 @@ def override(attr: str, value: Any): # this is dangerous, use with caution ) override("moe_cache_size", 0) override("moe_cache_rate", None) + override("moe_cache_mem", None) override("moe_cache_auto", False) if is_moe and config.moe_strategy == "cpu": @@ -1812,6 +2338,7 @@ def override(attr: str, value: Any): # this is dangerous, use with caution num_experts = config.model_config.num_experts if getattr(config, "moe_cache_auto", False): override("moe_cache_auto", False) + override("moe_cache_mem", None) override("moe_cache_size", 2 * num_experts) override("moe_prefill_overlap", True) logger.info_rank0( diff --git a/python/freetoken/engine/sample.py b/python/freetoken/engine/sample.py index c6fb2f40ee..9d12b5f571 100644 --- a/python/freetoken/engine/sample.py +++ b/python/freetoken/engine/sample.py @@ -1,7 +1,7 @@ from __future__ import annotations from dataclasses import dataclass, field -from typing import TYPE_CHECKING, List +from typing import TYPE_CHECKING, List, NamedTuple import torch from freetoken.utils import is_sm90_supported, nvtx_annotate @@ -10,19 +10,116 @@ from freetoken.core import Batch +@dataclass +class LogitsPlan: + """Inputs of the logits processors for the rows of one batch that asked for any + (logit_bias, min_tokens, min_p). The penalties are not here: the sampler applies + them from per-request device state (``RowPenalty``). + + Built on the host in ``Sampler.prepare`` from the requests' SamplingParams, applied + on the device by ``apply_logits_processors`` before the sampling kernel. ``None`` + when no request in the batch uses a processor, so the default path stays exactly + as it was. Row indices inside the plan are LOCAL (0..r-1, in the order of ``rows``); + ``rows`` maps them back to the batch. + """ + + rows: torch.Tensor # [r] int64 batch rows + temps: torch.Tensor # [r, 1] float32 temperatures (min_p works on probabilities) + bias_rows: torch.Tensor | None = None # [b] int64 local rows + bias_ids: torch.Tensor | None = None # [b] int64 + bias_vals: torch.Tensor | None = None # [b] float32 + # Rows still under min_tokens and the ids (EOS + stop_token_ids, padded with + # vocab_size, a scratch column that is dropped) they must not sample yet. + min_rows: torch.Tensor | None = None # [m] int64 local rows + min_ids: torch.Tensor | None = None # [m, S] int64 + min_p: torch.Tensor | None = None # [r, 1] float32 (0 = off for that row) + + +class RowPenalty(NamedTuple): + """One batch row's presence / frequency / repetition penalty. ``counts`` are the + request's generated-token counts ([V] int32), updated on the device after each draw so + overlap scheduling never sees them late; ``prompt_mask`` ([V] bool) marks the prompt's + ids for the repetition penalty and is None when that is off.""" + + row: int + counts: torch.Tensor + presence: float + frequency: float + repetition: float = 1.0 + prompt_mask: torch.Tensor | None = None + + @dataclass class BatchSamplingArgs: temperatures: torch.Tensor | None top_k: torch.Tensor | None = None top_p: torch.Tensor | None = None greedy_mask: torch.Tensor | None = None - penalties: list[tuple[int, torch.Tensor, float, float]] = field(default_factory=list) + penalties: list[RowPenalty] = field(default_factory=list) + plan: LogitsPlan | None = None + logprob_rows: torch.Tensor | None = None + max_top_logprobs: int = 0 def make_device_tensor(data: List, dtype: torch.dtype, device: torch.device) -> torch.Tensor: + if device.type != "cuda": # the unit tests build plans on the CPU + return torch.tensor(data, dtype=dtype, device=device) return torch.tensor(data, dtype=dtype, pin_memory=True).to(device, non_blocking=True) +def _padded_rows(rows: list[list[int]], pad: int, device: torch.device) -> torch.Tensor: + """[len(rows), max_len] int64 with ``pad`` on the right.""" + width = max(1, max((len(r) for r in rows), default=1)) + return make_device_tensor([r + [pad] * (width - len(r)) for r in rows], torch.int64, device) + + +def apply_penalties(logits: torch.Tensor, penalties: list[RowPenalty]) -> torch.Tensor: + """A float32 copy of ``logits`` with each row's penalties applied in vLLM's order: + repetition (over prompt + generated ids), then frequency and presence (generated).""" + logits = logits.float().clone() + for row, counts, presence, frequency, repetition, prompt_mask in penalties: + if repetition != 1.0: + x = logits[row] + seen = counts > 0 + if prompt_mask is not None: + seen = seen | prompt_mask + # HF / vLLM semantics: divide positive logits, multiply negative ones. + penalized = torch.where(x > 0, x / repetition, x * repetition) + logits[row] = torch.where(seen, penalized, x) + if presence or frequency: + logits[row] -= frequency * counts + presence * (counts > 0) + return logits + + +def apply_logits_processors( + logits: torch.Tensor, plan: LogitsPlan, vocab_size: int +) -> torch.Tensor: + """A float32 copy of ``logits`` with the plan applied to its rows: logit_bias, the + min_tokens mask, then min_p. Rows outside the plan are copied unchanged.""" + out = logits.to(torch.float32, copy=True) + sub = out[plan.rows] # [r, V] (advanced indexing copies) + device = sub.device + neg_inf = float("-inf") + + if plan.bias_ids is not None: + sub.index_put_((plan.bias_rows, plan.bias_ids), plan.bias_vals, accumulate=True) + + if plan.min_ids is not None: + block = torch.zeros((plan.min_ids.shape[0], vocab_size + 1), dtype=torch.bool, device=device) + block.scatter_(1, plan.min_ids, torch.ones_like(plan.min_ids, dtype=torch.bool)) + rows = sub[plan.min_rows] + sub[plan.min_rows] = torch.where(block[:, :vocab_size], torch.full_like(rows, neg_inf), rows) + + if plan.min_p is not None: + probs = torch.softmax(sub / plan.temps, dim=-1) + keep = probs >= plan.min_p * probs.amax(dim=-1, keepdim=True) + keep = keep | (plan.min_p <= 0.0) # rows without min_p stay untouched + sub = torch.where(keep, sub, torch.full_like(sub, neg_inf)) + + out[plan.rows] = sub + return out + + def sample_impl( logits: torch.Tensor, temperatures: torch.Tensor, @@ -57,23 +154,112 @@ class Sampler: device: torch.device vocab_size: int + def _plan(self, batch: Batch, params) -> LogitsPlan | None: + """The LogitsPlan for this batch, or None when no request asks for a processor.""" + MIN_T = 1e-6 + rows = [i for i, p in enumerate(params) if p.needs_logits_processing] + if not rows: + return None + picked = [(batch.reqs[i], params[i]) for i in rows] + temps = [[max(0.0 if p.is_greedy else p.temperature, MIN_T)] for _, p in picked] + plan = LogitsPlan( + rows=make_device_tensor(rows, torch.int64, self.device), + temps=make_device_tensor(temps, torch.float32, self.device), + ) + pad = self.vocab_size + + bias_rows: list[int] = [] + bias_ids: list[int] = [] + bias_vals: list[float] = [] + for local, (_, p) in enumerate(picked): + for tid, val in p.logit_bias or (): + if 0 <= int(tid) < self.vocab_size: + bias_rows.append(local) + bias_ids.append(int(tid)) + bias_vals.append(float(val)) + if bias_ids: + plan.bias_rows = make_device_tensor(bias_rows, torch.int64, self.device) + plan.bias_ids = make_device_tensor(bias_ids, torch.int64, self.device) + plan.bias_vals = make_device_tensor(bias_vals, torch.float32, self.device) + + min_rows: list[int] = [] + min_ids: list[list[int]] = [] + for local, (req, p) in enumerate(picked): + if p.min_tokens > 0 and p.min_tokens_stop_ids: + # device_len, not the host ids: under overlap scheduling the previous token + # reaches req.input_ids only after this batch is prepared + generated = req.device_len - (req.max_device_len - req.output_len) + if generated < p.min_tokens: + min_rows.append(local) + min_ids.append([t for t in p.min_tokens_stop_ids if 0 <= t < self.vocab_size]) + if min_rows: + plan.min_rows = make_device_tensor(min_rows, torch.int64, self.device) + plan.min_ids = _padded_rows(min_ids, pad, self.device) + + if any(p.min_p > 0.0 for _, p in picked): + plan.min_p = make_device_tensor([[p.min_p] for _, p in picked], torch.float32, self.device) + return plan + + def _prompt_mask(self, req) -> torch.Tensor: + """[V] bool marking the request's prompt ids, uploaded once per request.""" + prompt_len = req.max_device_len - req.output_len + # multimodal placeholders are pseudo-ids >= vocab_size (MM_PAD_SHIFT_VALUE); + # send them to the scratch column, not past the scatter target + ids = req.input_ids[:prompt_len].to(torch.int64).clamp(max=self.vocab_size) + if self.device.type == "cuda": + ids = ids.pin_memory().to(self.device, non_blocking=True) + mask = torch.zeros(self.vocab_size + 1, dtype=torch.bool, device=self.device) + mask[ids] = True + return mask[: self.vocab_size] + def prepare(self, batch: Batch) -> BatchSamplingArgs: params = [r.sampling_params for r in batch.reqs] is_greedy = [p.is_greedy for p in params] penalties = [] for row, req in enumerate(batch.reqs): p = req.sampling_params - if not (p.presence_penalty or p.frequency_penalty) or not req.can_decode: + repetition = p.repetition_penalty + if not (p.presence_penalty or p.frequency_penalty or repetition != 1.0): + continue + if not req.can_decode: continue if req.output_token_counts is None: req.output_token_counts = torch.zeros( self.vocab_size, dtype=torch.int32, device=self.device ) + prompt_mask = None + if repetition != 1.0: + if req.prompt_token_mask is None: + req.prompt_token_mask = self._prompt_mask(req) + prompt_mask = req.prompt_token_mask penalties.append( - (row, req.output_token_counts, p.presence_penalty, p.frequency_penalty) + RowPenalty( + row, + req.output_token_counts, + p.presence_penalty, + p.frequency_penalty, + repetition, + prompt_mask, + ) ) + plan = self._plan(batch, params) + want_logprobs = [p.logprobs for p in params] + logprob_rows = ( + make_device_tensor(want_logprobs, torch.bool, self.device) + if any(want_logprobs) + else None + ) + max_top_logprobs = min( + max((p.top_logprobs for p in params if p.logprobs), default=0), self.vocab_size + ) if all(is_greedy): - return BatchSamplingArgs(temperatures=None, penalties=penalties) + return BatchSamplingArgs( + temperatures=None, + penalties=penalties, + plan=plan, + logprob_rows=logprob_rows, + max_top_logprobs=max_top_logprobs, + ) MIN_P = MIN_T = 1e-6 # Greedy outputs are selected explicitly in sample(); use neutral sampling @@ -97,16 +283,24 @@ def prepare(self, batch: Batch) -> BatchSamplingArgs: make_device_tensor(is_greedy, torch.bool, self.device) if any(is_greedy) else None ) return BatchSamplingArgs( - temperatures, top_k=top_k, top_p=top_p, greedy_mask=greedy_mask, penalties=penalties + temperatures, + top_k=top_k, + top_p=top_p, + greedy_mask=greedy_mask, + penalties=penalties, + plan=plan, + logprob_rows=logprob_rows, + max_top_logprobs=max_top_logprobs, ) @nvtx_annotate("Sampler") def sample(self, logits: torch.Tensor, args: BatchSamplingArgs) -> torch.Tensor: with torch.cuda.nvtx.range("Sampler"): if args.penalties: - logits = logits.float().clone() - for row, counts, presence, frequency in args.penalties: - logits[row] -= frequency * counts + presence * (counts > 0) + logits = apply_penalties(logits, args.penalties) + # after the penalties: min_p must see the final distribution + if args.plan is not None: + logits = apply_logits_processors(logits, args.plan, self.vocab_size) if args.temperatures is None: # greedy sampling tokens = torch.argmax(logits, dim=-1) else: @@ -118,8 +312,64 @@ def sample(self, logits: torch.Tensor, args: BatchSamplingArgs) -> torch.Tensor: tokens = torch.where(args.greedy_mask, greedy_tokens, tokens) # Update on the sampling stream: overlapped scheduling can prepare the next # batch before the previous token reaches Req.input_ids on the CPU. - for row, counts, _, _ in args.penalties: + for row, counts, *_ in args.penalties: counts.scatter_add_( 0, tokens[row : row + 1].long(), counts.new_ones(1) ) return tokens + + def compute_logprobs( + self, + logits: torch.Tensor, + sampled_tokens: torch.Tensor, + args: BatchSamplingArgs, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None: + if args.logprob_rows is None: + return None + + requested_rows = torch.nonzero(args.logprob_rows, as_tuple=False).flatten() + if requested_rows.numel() == 0: + return None + + request_logits = logits.index_select(0, requested_rows).float() + # Reported values are raw model logprobs (pre-temperature log_softmax over logits). + request_logprobs = torch.log_softmax(request_logits, dim=-1) + + request_tokens = sampled_tokens.to(dtype=torch.long, device=logits.device).index_select( + 0, requested_rows + ) + request_row_idx = torch.arange(requested_rows.numel(), device=logits.device) + request_chosen_logprobs = request_logprobs[request_row_idx, request_tokens] + + chosen_logprobs = torch.full( + (logits.shape[0],), float("nan"), dtype=torch.float32, device=logits.device + ) + chosen_logprobs.index_copy_(0, requested_rows, request_chosen_logprobs) + + if args.max_top_logprobs > 0: + request_top_logprobs, request_top_ids = torch.topk( + request_logprobs, k=args.max_top_logprobs, dim=-1 + ) + top_ids = torch.full( + (logits.shape[0], args.max_top_logprobs), + -1, + dtype=torch.int32, + device=logits.device, + ) + top_logprobs = torch.full( + (logits.shape[0], args.max_top_logprobs), + float("-inf"), + dtype=torch.float32, + device=logits.device, + ) + top_ids[requested_rows] = request_top_ids.to(torch.int32) + top_logprobs[requested_rows] = request_top_logprobs + else: + top_ids = torch.empty((logits.shape[0], 0), dtype=torch.int32, device=logits.device) + top_logprobs = torch.empty((logits.shape[0], 0), dtype=torch.float32, device=logits.device) + + return ( + chosen_logprobs.to("cpu", non_blocking=True), + top_ids.to("cpu", non_blocking=True), + top_logprobs.to("cpu", non_blocking=True), + ) diff --git a/python/freetoken/env.py b/python/freetoken/env.py index 6878dc4a19..9a06b1c029 100644 --- a/python/freetoken/env.py +++ b/python/freetoken/env.py @@ -75,6 +75,24 @@ class EnvClassSingleton: # fp32 matches the Qwen3.x configs (mamba_ssm_dtype); fp16/bf16 halves the GDN state # pool at some precision cost on the long recurrence (mirrors SGLang's mamba_ssm_dtype). MAMBA_SSM_DTYPE = EnvStr("float32") + # cap the longest prefill warmup forward; 0 = the engine's own chunk cap + WARMUP_MAX_LEN = EnvInt(0) + # GPU-decode MoE: run a layer's expert miss copy on a side stream under the shared expert + MOE_COPY_OVERLAP = EnvBool(False) + # GPU-decode MoE cross-layer prefetch: off | measure (predict and count, no copies) | on (predict and copy) + MOE_PREFETCH = EnvStr("off") + # router-lookahead candidates per layer, and the non-resident ones kept per target layer + # (0 = 3 before GDN layers, 4 before full-attention layers) + MOE_PREFETCH_K = EnvInt(16) + MOE_PREFETCH_BUDGET = EnvInt(0) + # measure mode: also log the per-layer breakdown with every summary + MOE_PREFETCH_DEBUG = EnvBool(False) + # debug: check every GPU decode layer call's slot maps (meta) and slot bytes before/after the GEMM (1) + MOE_PREFETCH_VERIFY = EnvStr("0") + # routed experts a call byte-checks (scratch rows, ~2.8 MB each on Qwen3.8-Flash) + MOE_PREFETCH_VERIFY_ROWS = EnvInt(20) + # debug: record every expert slot writer and audit every held slot each N decode steps (0 = off) + MOE_SLOT_AUDIT = EnvInt(0) def __new__(cls): # single instance diff --git a/python/freetoken/kernel/aot.py b/python/freetoken/kernel/aot.py index bbc2d2c669..54539d952c 100644 --- a/python/freetoken/kernel/aot.py +++ b/python/freetoken/kernel/aot.py @@ -124,6 +124,23 @@ def build(build_directory: pathlib.Path) -> object: return KernelSpec(name=_make_name("fast_index_copy_multi", *args), build=build) +def _fast_index_copy_multi_slim_spec() -> KernelSpec: + from .fast_index_copy import SLIM_COPY_BLOCKS, SLIM_COPY_THREADS, SLIM_COPY_UNROLL + + args = make_cpp_args(SLIM_COPY_BLOCKS, SLIM_COPY_THREADS, SLIM_COPY_UNROLL) + + def build(build_directory: pathlib.Path) -> object: + return load_jit( + "fast_index_copy_multi_slim", + *args, + cuda_files=["fast_index_copy.cuh"], + cuda_wrappers=[("launch", f"&MultiIndexCopySlimKernel<{args}>::run")], + build_directory=str(build_directory), + ) + + return KernelSpec(name=_make_name("fast_index_copy_multi_slim", *args), build=build) + + def _batch_memcpy_spec() -> KernelSpec: def build(build_directory: pathlib.Path) -> object: return load_jit( @@ -156,6 +173,8 @@ def default_kernel_specs() -> tuple[KernelSpec, ...]: specs.append(_fast_index_copy_multi_spec(num_threads=1024, blocks_per_bank=32)) # prefill hit-D2D gather (HBM-bound: wide grid) + its miss-side batch H2D binding. specs.append(_fast_index_copy_multi_spec(num_threads=1024, blocks_per_bank=64)) + # opt-in decode copy (FREETOKEN_MOE_SLIM_COPY=1) + specs.append(_fast_index_copy_multi_slim_spec()) specs.append(_batch_memcpy_spec()) specs.append(_radix_spec()) diff --git a/python/freetoken/kernel/aot_models.py b/python/freetoken/kernel/aot_models.py index ecd1f8502b..8e0df9f59b 100644 --- a/python/freetoken/kernel/aot_models.py +++ b/python/freetoken/kernel/aot_models.py @@ -11,9 +11,11 @@ (a drifted derivation misses the prebuilt cache by spec name and falls back to JIT, which needs nvcc): -- store: ``element_size = num_kv_heads * head_dim * 2`` (bf16 KV row), one per - paged-KV attention group (kvcache/mha_pool.py, kvcache/hybrid_swa_pool.py). - DSV4 writes its MLA latent via torch scatter and contributes nothing. +- store: ``element_size = num_kv_heads * head_dim * dtype_bytes``, one per + paged-KV attention group (kvcache/mha_pool.py, kvcache/hybrid_swa_pool.py) and + per KV width: 2 for the 16-bit cache, 1 for an fp8 one (``--kv-cache-dtype + fp8``, kvcache/mha_pool.py). DSV4 writes its MLA latent via torch scatter and + contributes nothing. - index: ``element_size = hidden_size * 2`` (bf16 embedding row) paired with the runtime ``num_splits_for`` rule (layers/embedding.py -> kernel/index.py). DSV4 (plain nn.Embedding) and GGUF embeddings (GGUFEmbedding) bypass it. @@ -33,7 +35,8 @@ from .index import num_splits_for -KV_CACHE_DTYPE_BYTES = 2 # every current model allocates bf16 paged KV +KV_CACHE_DTYPE_BYTES = 2 # the default paged KV is the 16-bit compute dtype +FP8_KV_CACHE_DTYPE_BYTES = 1 # --kv-cache-dtype fp8 stores one e4m3 code per element EMBED_DTYPE_BYTES = 2 # embedding weights stay bf16 on the indexing() path @@ -406,8 +409,9 @@ def expert_bank_row_bytes(fmt: str, hidden_size: int, moe_intermediate_size: int ) -def store_element_sizes(model: AotModel) -> set[int]: - return {kv * hd * KV_CACHE_DTYPE_BYTES for kv, hd in model.kv_groups} +def store_element_sizes(model: AotModel, dtype_bytes: int = KV_CACHE_DTYPE_BYTES) -> set[int]: + """Store-kernel row sizes for one model's paged-KV groups at a given bytes/elem.""" + return {kv * hd * dtype_bytes for kv, hd in model.kv_groups} def index_variants(model: AotModel) -> set[tuple[int, int]]: @@ -429,9 +433,17 @@ def fast_index_copy_feature_sizes(model: AotModel) -> set[int]: def aggregate_store_element_sizes() -> tuple[int, ...]: + """Every store row size the runtime can ask for. + + Both KV widths ship: the 16-bit default and the fp8 (``--kv-cache-dtype fp8``) + code buffer, whose rows are exactly half as wide. A missing size is not a + correctness bug -- it is a kernel-cache miss that falls back to JIT and fails + the ``FREETOKEN_DISABLE_JIT=1`` release gate. + """ sizes: set[int] = set() for model in SUPPORTED_MODELS: - sizes.update(store_element_sizes(model)) + for dtype_bytes in (KV_CACHE_DTYPE_BYTES, FP8_KV_CACHE_DTYPE_BYTES): + sizes.update(store_element_sizes(model, dtype_bytes)) return tuple(sorted(sizes)) diff --git a/python/freetoken/kernel/batch_memcpy.py b/python/freetoken/kernel/batch_memcpy.py index b39e5cde24..67e2bf9870 100644 --- a/python/freetoken/kernel/batch_memcpy.py +++ b/python/freetoken/kernel/batch_memcpy.py @@ -28,6 +28,12 @@ def _probe(fn) -> None: src = torch.arange(16, dtype=torch.uint8).pin_memory() dst = torch.zeros(16, dtype=torch.uint8, device="cuda") stream = torch.cuda.Stream() + # torch.zeros filled dst on the CURRENT stream; the copy below runs on this one. + # Without the join the fill can land after the copy and clobber it, leaving the + # probe to read back zeros -- a false "copied wrong bytes" on hardware that + # supports the API. A cold process hides this (the first torch.zeros and + # torch.cuda.Stream() each synchronize the device); a warmed-up server does not. + stream.wait_stream(torch.cuda.current_stream()) fn( torch.tensor([dst.data_ptr()]), torch.tensor([src.data_ptr()]), diff --git a/python/freetoken/kernel/causal_conv1d.py b/python/freetoken/kernel/causal_conv1d.py index 81f1fab7b4..c196658648 100644 --- a/python/freetoken/kernel/causal_conv1d.py +++ b/python/freetoken/kernel/causal_conv1d.py @@ -22,9 +22,19 @@ def causal_conv1d_varlen( cu_seqlens: torch.Tensor, # [batch+1] int32 prefix sums of per-request lengths cache_indices: torch.Tensor, # [batch] int32 slot id per request has_initial_state: torch.Tensor, # [batch] bool (carry conv state across chunks) + max_seq_len: int | None = None, # host-known longest extend_len; None derives it on device ) -> torch.Tensor: - """Varlen (prefill) depthwise causal conv with silu; writes silu(conv) into ``x`` - in place and refreshes ``conv_states[cache_indices]`` with each request's tail.""" + """Varlen (prefill) depthwise causal conv with silu; returns silu(conv) and refreshes + ``conv_states[cache_indices]`` with each request's tail. + + Whether ``x`` survives depends on the backend: sgl_kernel convolves ``x`` in place and + returns it, while the triton fallback leaves ``x`` untouched and returns a new tensor. + Callers must use the returned tensor and must not rely on ``x`` afterwards (every + current call site builds ``x`` fresh and never reads it again). + + ``max_seq_len`` only sizes the triton launch grid. Deriving it from ``cu_seqlens`` costs a + D2H sync per prefill, which is illegal under CUDA graph capture, so the scheduler passes the + host value it already has.""" from freetoken.kernel.backend import is_sgl_kernel_installed if not is_sgl_kernel_installed(): @@ -33,7 +43,8 @@ def causal_conv1d_varlen( ) return triton_causal_conv1d_varlen( - x, weight, conv_states, cu_seqlens, cache_indices, has_initial_state + x, weight, conv_states, cu_seqlens, cache_indices, has_initial_state, + max_seq_len=max_seq_len, ) from sgl_kernel import causal_conv1d_fwd diff --git a/python/freetoken/kernel/csrc/cpu_moe/common.h b/python/freetoken/kernel/csrc/cpu_moe/common.h new file mode 100644 index 0000000000..adf7e14584 --- /dev/null +++ b/python/freetoken/kernel/csrc/cpu_moe/common.h @@ -0,0 +1,235 @@ +// Numeric helpers and ISA detection shared by the CPU MoE formats. +#pragma once + +#include +#include +#include +#include +#include +#include +#include + +#if defined(__x86_64__) || defined(__i386__) +#include +#define CPU_MOE_X86 1 +#else +#define CPU_MOE_X86 0 +#endif + +// Each SIMD kernel is compiled for its own ISA and picked at runtime (pick_isa), so the build sets no global -march. +#define CPU_MOE_TARGET(...) __attribute__((target(__VA_ARGS__))) + +#if CPU_MOE_X86 && ((defined(__GNUC__) && __GNUC__ >= 10) || defined(__clang__)) +#define CPU_MOE_HAS_AVX512BF16 1 +#define CPU_MOE_HAS_AVX512VNNI 1 +#endif + +namespace freetoken::cpu_moe { + +// Scratch rows are read with 64-byte SIMD loads, so their base is cache-line aligned whatever the heap state; std::vector alone guarantees 16 bytes. +template +struct CacheLineAllocator { + using value_type = T; + CacheLineAllocator() = default; + template + CacheLineAllocator(const CacheLineAllocator&) {} + T* allocate(size_t n) { return static_cast(::operator new(n * sizeof(T), std::align_val_t(64))); } + void deallocate(T* p, size_t) { ::operator delete(p, std::align_val_t(64)); } + template + bool operator==(const CacheLineAllocator&) const { return true; } +}; + +template +using AlignedVector = std::vector>; + +using bf16_t = uint16_t; + +inline float bf16_to_f32(bf16_t v) { + uint32_t u = static_cast(v) << 16; + float f; + std::memcpy(&f, &u, sizeof(f)); + return f; +} + +inline bf16_t f32_to_bf16(float f) { + uint32_t u; + std::memcpy(&u, &f, sizeof(u)); + // round-to-nearest-even + const uint32_t lsb = (u >> 16) & 1u; + u += 0x7fffu + lsb; + return static_cast(u >> 16); +} + +// ACT_SWIGLUOAI is the clamped (up + 1) swiglu (gpt-oss "swigluoai" / +// MiniMax-M3): gate/up are combined jointly with the runtime alpha/limit +// scalars, so it is handled in the gated_rows epilogue (act_apply never sees it; +// the mxfp4 kernel additionally fuses its own copy of the same math). +// ACT_SWIGLU_CLAMP (GLM-5.3 "swiglu_limit") is the same clamped form WITHOUT +// the (up + 1) bias: clamp(gate, max=lim) * sigmoid(alpha*gate) * clamp(up, +-lim). +enum ActKind { + ACT_SILU = 0, + ACT_GELU = 1, + ACT_GELU_TANH = 2, + ACT_SWIGLUOAI = 3, + ACT_SWIGLU_CLAMP = 4, +}; + +inline float act_apply(int act, float x) { + if (act == ACT_SILU) return x / (1.0f + std::exp(-x)); + if (act == ACT_GELU) + return 0.5f * x * (1.0f + std::erf(x * 0.70710678118654752440f)); + // gelu_tanh + const float k0 = 0.79788456080286535588f; // sqrt(2/pi) + const float inner = k0 * (x + 0.044715f * x * x * x); + return 0.5f * x * (1.0f + std::tanh(inner)); +} + +// Software prefetch distance (bytes) ahead of the current weight row stream. The +// weight stream is the bandwidth bottleneck (read once, never reused); nudging the +// HW prefetcher with a few cache lines of lookahead raises sustained throughput. +constexpr int PF_AHEAD = 512; + +const float kE2M1[16] = {0.0f, 0.5f, 1.0f, 1.5f, 2.0f, 3.0f, 4.0f, 6.0f, + -0.0f, -0.5f, -1.0f, -1.5f, -2.0f, -3.0f, -4.0f, -6.0f}; + +inline float fp16_to_f32(uint16_t h) { + const uint32_t sign = (uint32_t)(h & 0x8000u) << 16; + uint32_t exp = (h >> 10) & 0x1Fu; + uint32_t man = h & 0x3FFu; + uint32_t f; + if (exp == 0) { + if (man == 0) { + f = sign; + } else { + exp = 127 - 15 + 1; + while ((man & 0x400u) == 0) { + man <<= 1; + --exp; + } + man &= 0x3FFu; + f = sign | (exp << 23) | (man << 13); + } + } else if (exp == 0x1Fu) { + f = sign | 0x7F800000u | (man << 13); + } else { + f = sign | ((exp + (127 - 15)) << 23) | (man << 13); + } + float out; + std::memcpy(&out, &f, sizeof(out)); + return out; +} + +// e4m3 (OCP "fn": finite, max-normal 448, exp bias 7). Decoded into a 256-entry LUT. +inline float e4m3_decode(uint8_t v) { + const float sign = (v & 0x80u) ? -1.0f : 1.0f; + const uint32_t exp = (v >> 3) & 0xFu; + const uint32_t man = v & 0x7u; + if (exp == 0) return sign * (man / 8.0f) * 0.015625f; // 2^(1-7) = 2^-6 + return sign * (1.0f + man / 8.0f) * std::ldexp(1.0f, (int)exp - 7); +} + +#if CPU_MOE_X86 +CPU_MOE_TARGET("avx2,fma") +inline float hsum256(__m256 v) { + __m128 lo = _mm256_castps256_ps128(v); + lo = _mm_add_ps(lo, _mm256_extractf128_ps(v, 1)); + lo = _mm_add_ps(lo, _mm_movehl_ps(lo, lo)); + lo = _mm_add_ss(lo, _mm_shuffle_ps(lo, lo, 0x55)); + return _mm_cvtss_f32(lo); +} + +// AVX2 e2m1 nibble decode: codes (0..15) in 8 int32 lanes -> fp32. AVX2 vpermps is +// only 8-wide, so instead of a 16-entry LUT we use the e2m1 sign/magnitude symmetry: +// value = (code&8 ? - : +) * mag8[code&7], mag8 = e2m1[0..7]. The sign is bit 3 of +// the code shifted into the fp32 sign bit (bit 31). Bit-identical to the e2m1 LUT. +CPU_MOE_TARGET("avx2,fma") +inline __m256 e2m1_decode8(__m256i codes, __m256 mag8) { + __m256 mag = _mm256_permutevar8x32_ps(mag8, _mm256_and_si256(codes, _mm256_set1_epi32(7))); + __m256i sgn = _mm256_slli_epi32(_mm256_and_si256(codes, _mm256_set1_epi32(8)), 28); + return _mm256_xor_ps(mag, _mm256_castsi256_ps(sgn)); +} +#endif // CPU_MOE_X86 + +// SIMD tiers, ascending. Each format picks the highest tier <= the one chosen by +// pick_isa() that it implements (fp4 and fp8_block formats have no bf16-specific tier, so the +// avx512bf16 tier maps to their avx512 kernel). +enum IsaTier { ISA_SCALAR = 0, ISA_AVX2 = 1, ISA_AVX512 = 2, ISA_AVX512BF16 = 3 }; + +// Best tier the CPU+build supports, optionally capped DOWN by +// FREETOKEN_CPU_MOE_ISA={scalar,avx2,avx512,avx512bf16} (A/B testing on a machine +// that supports more). FREETOKEN_CPU_MOE_SCALAR=1 forces scalar (legacy alias). +inline IsaTier pick_isa() { +#if CPU_MOE_X86 + if (getenv("FREETOKEN_CPU_MOE_SCALAR")) return ISA_SCALAR; + IsaTier best = ISA_SCALAR; + if (__builtin_cpu_supports("avx2") && __builtin_cpu_supports("fma")) best = ISA_AVX2; + if (best >= ISA_AVX2 && __builtin_cpu_supports("avx512f")) best = ISA_AVX512; +#ifdef CPU_MOE_HAS_AVX512BF16 + if (best >= ISA_AVX512 && __builtin_cpu_supports("avx512bf16")) best = ISA_AVX512BF16; +#endif + if (const char* f = getenv("FREETOKEN_CPU_MOE_ISA")) { + IsaTier want = best; + if (!std::strcmp(f, "scalar")) want = ISA_SCALAR; + else if (!std::strcmp(f, "avx2")) want = ISA_AVX2; + else if (!std::strcmp(f, "avx512")) want = ISA_AVX512; + else if (!std::strcmp(f, "avx512bf16")) want = ISA_AVX512BF16; + if (want < best) best = want; // cap downward; never force above hw/build support + } + return best; +#else + return ISA_SCALAR; +#endif +} + +inline const char* tier_name(IsaTier t) { +#if CPU_MOE_X86 +#ifdef CPU_MOE_HAS_AVX512BF16 + if (t >= ISA_AVX512BF16) return "avx512bf16"; +#endif + if (t >= ISA_AVX512) return "avx512f"; + if (t >= ISA_AVX2) return "avx2"; +#endif + (void)t; + return "scalar"; +} + +// AVX-VNNI (VEX-256 VPDPBUSD) availability: Alder/Raptor Lake, Sapphire Rapids+, Zen5. +// Distinct from AVX-512 VNNI. Opt out with FREETOKEN_CPU_MOE_NO_VNNI=1 (A/B the W4A8 path). +inline bool cpu_has_avxvnni() { +#if CPU_MOE_X86 + const char* no = getenv("FREETOKEN_CPU_MOE_NO_VNNI"); + if (no && no[0] && no[0] != '0') return false; // ignore unset/empty/"0" + return __builtin_cpu_supports("avx2") && __builtin_cpu_supports("avxvnni"); +#else + return false; +#endif +} + +// AVX-512 VNNI (512-bit VPDPBUSD): Cascade Lake+, Ice Lake, Sapphire/Emerald Rapids, +// Zen4+. 2x the 256-bit AVX-VNNI width. FREETOKEN_CPU_MOE_NO_AVX512VNNI=1 forces the +// 256-bit path (A/B the two W4A8 kernels on the same box); FREETOKEN_CPU_MOE_NO_VNNI=1 +// still disables the whole W4A8 family (back to the faithful fp32 nvdot), so it is +// honored here too. +inline bool cpu_has_avx512vnni() { +#if CPU_MOE_X86 && defined(CPU_MOE_HAS_AVX512VNNI) + const char* no = getenv("FREETOKEN_CPU_MOE_NO_AVX512VNNI"); + if (no && no[0] && no[0] != '0') return false; + const char* no_vnni = getenv("FREETOKEN_CPU_MOE_NO_VNNI"); + if (no_vnni && no_vnni[0] && no_vnni[0] != '0') return false; + return __builtin_cpu_supports("avx512vnni"); +#else + return false; +#endif +} + +// Split an interleaved bf16 row into fp32 even/odd halves (even[m]=src[2m]). +// bf16->fp32 is exact, so this only reorders -- done once per token/route and +// reused across every output row of the GEMV. +inline void deinterleave_bf16_f32(const bf16_t* src, float* even, float* odd, int K) { + for (int m = 0; m < K / 2; ++m) { + even[m] = bf16_to_f32(src[2 * m]); + odd[m] = bf16_to_f32(src[2 * m + 1]); + } +} + +} // namespace freetoken::cpu_moe diff --git a/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp b/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp index 6f85971e9c..c8b4d1b465 100644 --- a/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp +++ b/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp @@ -1,32 +1,27 @@ -// CPU-compute MoE executor for the "cpu" offload backend. +// CPU-compute MoE executor for the cpu and hybrid MoE strategies, and its Python binding. // // Decode ships activations to the CPU, computes the routed experts here (reading -// the pinned host expert banks at full RAM bandwidth), and ships the results -// back. To keep the whole decode path inside a single CUDA graph we expose -// submit/sync as host nodes via cudaLaunchHostFunc -- the callbacks only touch a -// CPU worker pool + pinned host buffers and never call any CUDA API. +// the pinned host expert banks at full RAM bandwidth), and ships the results back. +// One task (one MoE layer) is in flight at a time: submit() wakes the persistent worker +// pool, the pool runs pass 1 (gate_up) and pass 2 (down), and sync() waits for it. +// The GPU rings a doorbell to submit and waits on it to sync (GpuSync, gpu_sync.h), so the +// whole decode step stays in one CUDA graph; cudaLaunchHostFunc nodes are the fallback. // -// One task is in flight at a time (per MoE layer): submit() wakes the pool, -// sync() blocks the host-func thread until the pool drains. The heavy GEMV runs -// on the persistent worker threads, not the host-func thread. -// -// Weight formats: bf16, NVFP4, MXFP4, ds_fp4, Q4_0 and block-FP8 expert banks (see WFmt and -// the per-format bank schemas). Compute is FP32-accumulate; the intermediate is -// stored bf16 to match the GPU decode path. ISA is chosen once at construction -// (AVX-512-BF16 dpbf16 -> AVX-512F widening -> AVX2+FMA -> scalar). +// Weight formats live in kernels/, one file each behind the MoeFormat interface (format.h). +// Compute is FP32-accumulate; the intermediate is stored bf16 to match the GPU decode path. #include #include #include -#include #include #include -#include #include +#include #include #include #include #include +#include #include #include @@ -40,1379 +35,99 @@ #define CPU_MOE_HAS_AFFINITY 0 #endif -#if defined(__x86_64__) || defined(__i386__) -#include -#define CPU_MOE_X86 1 -#else -#define CPU_MOE_X86 0 -#endif - -namespace { - -inline void cpu_relax() { -#if CPU_MOE_X86 - _mm_pause(); -#elif defined(__aarch64__) - asm volatile("yield" ::: "memory"); -#endif -} - -using bf16_t = uint16_t; - -inline float bf16_to_f32(bf16_t v) { - uint32_t u = static_cast(v) << 16; - float f; - std::memcpy(&f, &u, sizeof(f)); - return f; -} - -inline bf16_t f32_to_bf16(float f) { - uint32_t u; - std::memcpy(&u, &f, sizeof(u)); - // round-to-nearest-even - const uint32_t lsb = (u >> 16) & 1u; - u += 0x7fffu + lsb; - return static_cast(u >> 16); -} - -// ACT_SWIGLUOAI is the clamped (up + 1) swiglu (gpt-oss "swigluoai" / -// MiniMax-M3): gate/up are combined jointly with the runtime alpha/limit -// scalars, so it is handled in the do_pass1 epilogue (act_apply never sees it; -// the mxfp4 kernel additionally fuses its own copy of the same math). -// ACT_SWIGLU_CLAMP (GLM-5.3 "swiglu_limit") is the same clamped form WITHOUT -// the (up + 1) bias: clamp(gate, max=lim) * sigmoid(alpha*gate) * clamp(up, +-lim). -enum ActKind { - ACT_SILU = 0, - ACT_GELU = 1, - ACT_GELU_TANH = 2, - ACT_SWIGLUOAI = 3, - ACT_SWIGLU_CLAMP = 4, -}; - -inline float act_apply(int act, float x) { - if (act == ACT_SILU) return x / (1.0f + std::exp(-x)); - if (act == ACT_GELU) - return 0.5f * x * (1.0f + std::erf(x * 0.70710678118654752440f)); - // gelu_tanh - const float k0 = 0.79788456080286535588f; // sqrt(2/pi) - const float inner = k0 * (x + 0.044715f * x * x * x); - return 0.5f * x * (1.0f + std::tanh(inner)); -} - -// ------------------------------- dot products ------------------------------- -// dot(weight[bf16], act[bf16], n) -> fp32. The selected impl is a function -// pointer chosen at runtime; the per-row call overhead is negligible vs n. - -using dot_fn = float (*)(const bf16_t*, const bf16_t*, int); - -float dot_scalar(const bf16_t* w, const bf16_t* x, int n) { - float acc = 0.0f; - for (int i = 0; i < n; ++i) acc += bf16_to_f32(w[i]) * bf16_to_f32(x[i]); - return acc; -} - -// Software prefetch distance (bytes) ahead of the current weight row stream. The -// weight stream is the bandwidth bottleneck (read once, never reused); nudging the -// HW prefetcher with a few cache lines of lookahead raises sustained throughput. -constexpr int PF_AHEAD = 512; - -#if CPU_MOE_X86 -__attribute__((target("avx512f"))) -float dot_avx512f(const bf16_t* w, const bf16_t* x, int n) { - // 4 independent accumulators -> more in-flight loads (memory-level parallelism), - // which is what lifts a bandwidth-bound GEMV toward peak. - __m512 a0 = _mm512_setzero_ps(), a1 = _mm512_setzero_ps(); - __m512 a2 = _mm512_setzero_ps(), a3 = _mm512_setzero_ps(); - int i = 0; - for (; i + 64 <= n; i += 64) { - _mm_prefetch(reinterpret_cast(w + i) + PF_AHEAD, _MM_HINT_T0); - for (int j = 0; j < 64; j += 16) { - __m256i wi = _mm256_loadu_si256(reinterpret_cast(w + i + j)); - __m256i xi = _mm256_loadu_si256(reinterpret_cast(x + i + j)); - __m512 wf = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(wi), 16)); - __m512 xf = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(xi), 16)); - __m512& acc = (j == 0) ? a0 : (j == 16) ? a1 : (j == 32) ? a2 : a3; - acc = _mm512_fmadd_ps(wf, xf, acc); - } - } - for (; i + 16 <= n; i += 16) { - __m256i wi = _mm256_loadu_si256(reinterpret_cast(w + i)); - __m256i xi = _mm256_loadu_si256(reinterpret_cast(x + i)); - __m512 wf = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(wi), 16)); - __m512 xf = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(xi), 16)); - a0 = _mm512_fmadd_ps(wf, xf, a0); - } - float s = _mm512_reduce_add_ps(_mm512_add_ps(_mm512_add_ps(a0, a1), _mm512_add_ps(a2, a3))); - for (; i < n; ++i) s += bf16_to_f32(w[i]) * bf16_to_f32(x[i]); - return s; -} - -#if (defined(__GNUC__) && __GNUC__ >= 10) || defined(__clang__) -#define CPU_MOE_HAS_AVX512BF16 1 -__attribute__((target("avx512bf16,avx512f"))) -static inline __m512bh load_bh(const bf16_t* p) { - __m512i raw = _mm512_loadu_si512(reinterpret_cast(p)); - __m512bh out; - std::memcpy(&out, &raw, sizeof(out)); - return out; -} - -__attribute__((target("avx512bf16,avx512f"))) -float dot_avx512bf16(const bf16_t* w, const bf16_t* x, int n) { - // 4 accumulators (128 bf16/iter) for memory-level parallelism + a prefetch nudge. - __m512 a0 = _mm512_setzero_ps(), a1 = _mm512_setzero_ps(); - __m512 a2 = _mm512_setzero_ps(), a3 = _mm512_setzero_ps(); - int i = 0; - for (; i + 128 <= n; i += 128) { - _mm_prefetch(reinterpret_cast(w + i) + PF_AHEAD, _MM_HINT_T0); - a0 = _mm512_dpbf16_ps(a0, load_bh(w + i), load_bh(x + i)); - a1 = _mm512_dpbf16_ps(a1, load_bh(w + i + 32), load_bh(x + i + 32)); - a2 = _mm512_dpbf16_ps(a2, load_bh(w + i + 64), load_bh(x + i + 64)); - a3 = _mm512_dpbf16_ps(a3, load_bh(w + i + 96), load_bh(x + i + 96)); - } - for (; i + 32 <= n; i += 32) { - a0 = _mm512_dpbf16_ps(a0, load_bh(w + i), load_bh(x + i)); - } - float s = _mm512_reduce_add_ps(_mm512_add_ps(_mm512_add_ps(a0, a1), _mm512_add_ps(a2, a3))); - for (; i < n; ++i) s += bf16_to_f32(w[i]) * bf16_to_f32(x[i]); - return s; -} -#endif // avx512bf16 available - -// --------- AVX2 + FMA (256-bit fallback for CPUs without AVX-512) ---------- -// Covers Intel 12-14th gen / Arrow Lake (AVX-512 fused off) and AMD Zen<4. The -// bf16->fp32 widen is a zero-extend + <<16; with 4 independent accumulators the -// GEMV is memory-bandwidth bound, same as the AVX-512 path (just half the width). -__attribute__((target("avx2,fma"))) -inline float hsum256(__m256 v) { - __m128 lo = _mm256_castps256_ps128(v); - lo = _mm_add_ps(lo, _mm256_extractf128_ps(v, 1)); - lo = _mm_add_ps(lo, _mm_movehl_ps(lo, lo)); - lo = _mm_add_ss(lo, _mm_shuffle_ps(lo, lo, 0x55)); - return _mm_cvtss_f32(lo); -} - -__attribute__((target("avx2,fma"))) -float dot_avx2(const bf16_t* w, const bf16_t* x, int n) { - __m256 a0 = _mm256_setzero_ps(), a1 = _mm256_setzero_ps(); - __m256 a2 = _mm256_setzero_ps(), a3 = _mm256_setzero_ps(); - int i = 0; - for (; i + 32 <= n; i += 32) { - _mm_prefetch(reinterpret_cast(w + i) + PF_AHEAD, _MM_HINT_T0); - for (int j = 0; j < 32; j += 8) { - __m128i wi = _mm_loadu_si128(reinterpret_cast(w + i + j)); - __m128i xi = _mm_loadu_si128(reinterpret_cast(x + i + j)); - __m256 wf = _mm256_castsi256_ps(_mm256_slli_epi32(_mm256_cvtepu16_epi32(wi), 16)); - __m256 xf = _mm256_castsi256_ps(_mm256_slli_epi32(_mm256_cvtepu16_epi32(xi), 16)); - __m256& acc = (j == 0) ? a0 : (j == 8) ? a1 : (j == 16) ? a2 : a3; - acc = _mm256_fmadd_ps(wf, xf, acc); - } - } - for (; i + 8 <= n; i += 8) { - __m128i wi = _mm_loadu_si128(reinterpret_cast(w + i)); - __m128i xi = _mm_loadu_si128(reinterpret_cast(x + i)); - __m256 wf = _mm256_castsi256_ps(_mm256_slli_epi32(_mm256_cvtepu16_epi32(wi), 16)); - __m256 xf = _mm256_castsi256_ps(_mm256_slli_epi32(_mm256_cvtepu16_epi32(xi), 16)); - a0 = _mm256_fmadd_ps(wf, xf, a0); - } - float s = hsum256(_mm256_add_ps(_mm256_add_ps(a0, a1), _mm256_add_ps(a2, a3))); - for (; i < n; ++i) s += bf16_to_f32(w[i]) * bf16_to_f32(x[i]); - return s; -} -#endif // CPU_MOE_X86 +#include "format.h" +#include "gpu_sync.h" +#include "gpu_sync_io.h" -// --------------------------- NVFP4 (W4A16) dequant --------------------------- -// Weights: e2m1 4-bit codes (2/byte, low nibble first), per-16 block scale in -// fp8-e4m3, per-output-row global scale in fp16. Dequant matches the GPU kernels -// (freetoken/kernel/triton/nvfp4_dequant.py): w = E2M1[code] * e4m3(scale) * global. -// Activations stay bf16 (W4A16); the GEMV dequantizes weights inside the K-loop. +namespace freetoken::cpu_moe { -const float kE2M1[16] = {0.0f, 0.5f, 1.0f, 1.5f, 2.0f, 3.0f, 4.0f, 6.0f, - -0.0f, -0.5f, -1.0f, -1.5f, -2.0f, -3.0f, -4.0f, -6.0f}; +std::unique_ptr make_bf16_format(const MoeParams& p, const FormatArgs& a); +std::unique_ptr make_nvfp4_format(const MoeParams& p, const FormatArgs& a); +std::unique_ptr make_mxfp4_format(const MoeParams& p, const FormatArgs& a); +std::unique_ptr make_ds_fp4_format(const MoeParams& p, const FormatArgs& a); +std::unique_ptr make_q4_0_format(const MoeParams& p, const FormatArgs& a); +std::unique_ptr make_fp8_block_format(const MoeParams& p, const FormatArgs& a); -// e2m1 * 2 as exact int8 (all e2m1 values are multiples of 0.5). Used by the AVX-VNNI -// W4A8 path: nibble -> int8 weight via PSHUFB LUT, then VPDPBUSD against int8 activations; -// the *2 is undone by a 0.5 folded into the final scale. Mirrors ggml's kvalues_mxfp4. -alignas(16) const int8_t kE2M1x2[16] = {0, 1, 2, 3, 4, 6, 8, 12, 0, -1, -2, -3, -4, -6, -8, -12}; - -inline float fp16_to_f32(uint16_t h) { - const uint32_t sign = (uint32_t)(h & 0x8000u) << 16; - uint32_t exp = (h >> 10) & 0x1Fu; - uint32_t man = h & 0x3FFu; - uint32_t f; - if (exp == 0) { - if (man == 0) { - f = sign; - } else { - exp = 127 - 15 + 1; - while ((man & 0x400u) == 0) { - man <<= 1; - --exp; - } - man &= 0x3FFu; - f = sign | (exp << 23) | (man << 13); - } - } else if (exp == 0x1Fu) { - f = sign | 0x7F800000u | (man << 13); - } else { - f = sign | ((exp + (127 - 15)) << 23) | (man << 13); - } - float out; - std::memcpy(&out, &f, sizeof(out)); - return out; -} - -// e4m3 (OCP "fn": finite, max-normal 448, exp bias 7). Decoded into a 256-entry LUT. -inline float e4m3_decode(uint8_t v) { - const float sign = (v & 0x80u) ? -1.0f : 1.0f; - const uint32_t exp = (v >> 3) & 0xFu; - const uint32_t man = v & 0x7u; - if (exp == 0) return sign * (man / 8.0f) * 0.015625f; // 2^(1-7) = 2^-6 - return sign * (1.0f + man / 8.0f) * std::ldexp(1.0f, (int)exp - 7); -} - -// ------------------------ block-FP8 (W8A16) dequant ------------------------ -// fp8-e4m3 weights row-major [rows, K] plus one bf16 scale per 128x128 block. The scale tensor is named weight_scale_inv but multiplies, as in kernel/triton/fp8_blockscale_moe.py. -// Whole 128-blocks read the activations as fp32 even/odd halves (xe[m] = x[2m], xo[m] = x[2m+1]); a ragged tail reads the bf16 row x. -// 0x7F/0xFF (NaN in torch float8_e4m3fn) decode to +-480 in every kernel, as e4m3_decode does. -using fp8_block_dot_fn = float (*)(const uint8_t*, const float*, const float*, const bf16_t*, - const bf16_t*, int, const float*); - -constexpr int FP8_BLK = 128; // K elements covered by one weight scale - -inline float fp8_block_tail(const uint8_t* w, const bf16_t* x, int k0, int K, const float* e4m3) { - float blk = 0.0f; - for (int k = k0; k < K; ++k) blk += e4m3[w[k]] * bf16_to_f32(x[k]); - return blk; -} - -float dot_fp8_block_scalar(const uint8_t* w, const float* xe, const float* xo, const bf16_t* x, - const bf16_t* s, int K, const float* e4m3) { - float acc = 0.0f; - int b = 0, k0 = 0; - for (; k0 + FP8_BLK <= K; k0 += FP8_BLK, ++b) { - // 8 independent sums: a single running sum is bound by the add latency - float a[8] = {}; - const uint8_t* wb = w + k0; - const float* pe = xe + k0 / 2; - const float* po = xo + k0 / 2; - for (int j = 0; j < FP8_BLK / 2; j += 4) - for (int l = 0; l < 4; ++l) { - a[2 * l] += e4m3[wb[2 * (j + l)]] * pe[j + l]; - a[2 * l + 1] += e4m3[wb[2 * (j + l) + 1]] * po[j + l]; - } - acc += (((a[0] + a[1]) + (a[2] + a[3])) + ((a[4] + a[5]) + (a[6] + a[7]))) * - bf16_to_f32(s[b]); - } - if (k0 < K) acc += fp8_block_tail(w, x, k0, K, e4m3) * bf16_to_f32(s[b]); - return acc; -} - -#if CPU_MOE_X86 -// A byte b read as fp16 bits (b & 0x80) << 8 | (b & 0x7F) << 7 is its e4m3 value times 2^-8, subnormals included, so the block scale carries the 256. -// In a 16-bit lane holding bytes (2m, 2m+1), srai 1 then & 0xBF80 builds that for byte 2m+1; a left shift by 8 first does byte 2m. Every AVX2 CPU has F16C. -__attribute__((target("avx2,fma,f16c"))) -float dot_fp8_block_avx2(const uint8_t* w, const float* xe, const float* xo, const bf16_t* x, - const bf16_t* s, int K, const float* e4m3) { - const __m256i keep = _mm256_set1_epi16((short)0xBF80); - __m256 total = _mm256_setzero_ps(); - int b = 0, k0 = 0; - for (; k0 + FP8_BLK <= K; k0 += FP8_BLK, ++b) { - _mm_prefetch(reinterpret_cast(w + k0) + PF_AHEAD, _MM_HINT_T0); - __m256 a0 = _mm256_setzero_ps(), a1 = _mm256_setzero_ps(); - __m256 a2 = _mm256_setzero_ps(), a3 = _mm256_setzero_ps(); - for (int j = 0; j < FP8_BLK; j += 32) { - const __m256i v = _mm256_loadu_si256(reinterpret_cast(w + k0 + j)); - const __m256i odd = _mm256_and_si256(_mm256_srai_epi16(v, 1), keep); - const __m256i even = _mm256_and_si256(_mm256_srai_epi16(_mm256_slli_epi16(v, 8), 1), keep); - const float* pe = xe + (k0 + j) / 2; - const float* po = xo + (k0 + j) / 2; - a0 = _mm256_fmadd_ps(_mm256_cvtph_ps(_mm256_castsi256_si128(even)), _mm256_loadu_ps(pe), a0); - a1 = _mm256_fmadd_ps(_mm256_cvtph_ps(_mm256_extracti128_si256(even, 1)), _mm256_loadu_ps(pe + 8), a1); - a2 = _mm256_fmadd_ps(_mm256_cvtph_ps(_mm256_castsi256_si128(odd)), _mm256_loadu_ps(po), a2); - a3 = _mm256_fmadd_ps(_mm256_cvtph_ps(_mm256_extracti128_si256(odd, 1)), _mm256_loadu_ps(po + 8), a3); - } - const __m256 blk = _mm256_add_ps(_mm256_add_ps(a0, a1), _mm256_add_ps(a2, a3)); - total = _mm256_fmadd_ps(blk, _mm256_set1_ps(bf16_to_f32(s[b]) * 256.0f), total); - } - float acc = hsum256(total); - if (k0 < K) acc += fp8_block_tail(w, x, k0, K, e4m3) * bf16_to_f32(s[b]); - return acc; -} - -__attribute__((target("avx512f,avx2,fma,f16c"))) -float dot_fp8_block_avx512f(const uint8_t* w, const float* xe, const float* xo, const bf16_t* x, - const bf16_t* s, int K, const float* e4m3) { - const __m256i keep = _mm256_set1_epi16((short)0xBF80); - __m512 total = _mm512_setzero_ps(); - int b = 0, k0 = 0; - for (; k0 + FP8_BLK <= K; k0 += FP8_BLK, ++b) { - _mm_prefetch(reinterpret_cast(w + k0) + PF_AHEAD, _MM_HINT_T0); - __m512 a0 = _mm512_setzero_ps(), a1 = _mm512_setzero_ps(); - __m512 a2 = _mm512_setzero_ps(), a3 = _mm512_setzero_ps(); - for (int j = 0; j < FP8_BLK; j += 64) { - const __m256i v0 = _mm256_loadu_si256(reinterpret_cast(w + k0 + j)); - const __m256i v1 = _mm256_loadu_si256(reinterpret_cast(w + k0 + j + 32)); - const __m256i o0 = _mm256_and_si256(_mm256_srai_epi16(v0, 1), keep); - const __m256i e0 = _mm256_and_si256(_mm256_srai_epi16(_mm256_slli_epi16(v0, 8), 1), keep); - const __m256i o1 = _mm256_and_si256(_mm256_srai_epi16(v1, 1), keep); - const __m256i e1 = _mm256_and_si256(_mm256_srai_epi16(_mm256_slli_epi16(v1, 8), 1), keep); - const float* pe = xe + (k0 + j) / 2; - const float* po = xo + (k0 + j) / 2; - a0 = _mm512_fmadd_ps(_mm512_cvtph_ps(e0), _mm512_loadu_ps(pe), a0); - a1 = _mm512_fmadd_ps(_mm512_cvtph_ps(o0), _mm512_loadu_ps(po), a1); - a2 = _mm512_fmadd_ps(_mm512_cvtph_ps(e1), _mm512_loadu_ps(pe + 16), a2); - a3 = _mm512_fmadd_ps(_mm512_cvtph_ps(o1), _mm512_loadu_ps(po + 16), a3); - } - const __m512 blk = _mm512_add_ps(_mm512_add_ps(a0, a1), _mm512_add_ps(a2, a3)); - total = _mm512_fmadd_ps(blk, _mm512_set1_ps(bf16_to_f32(s[b]) * 256.0f), total); - } - float acc = _mm512_reduce_add_ps(total); - if (k0 < K) acc += fp8_block_tail(w, x, k0, K, e4m3) * bf16_to_f32(s[b]); - return acc; -} -#endif // CPU_MOE_X86 - -// Activations pre-deinterleaved to fp32 (xe[m]=x[2m], xo[m]=x[2m+1]); see the -// ds_fp4 dot below for why (drops the hot loop to ~1.5 shuffle ops / 16 weights). -using nvdot_fn = float (*)(const uint8_t*, const uint8_t*, float, const float*, const float*, - int, const float*, const float*); - -float dot_nvfp4_scalar(const uint8_t* packed, const uint8_t* scale, float global, - const float* xe, const float* xo, int K, const float* e2m1, - const float* e4m3) { - float acc = 0.0f; - const int nb = K / 16; - for (int b = 0; b < nb; ++b) { - const float bs = e4m3[scale[b]]; - const uint8_t* pk = packed + (size_t)b * 8; - const float* xeb = xe + (size_t)b * 8; // 16 K -> 8 even + 8 odd - const float* xob = xo + (size_t)b * 8; - float bsum = 0.0f; - for (int j = 0; j < 8; ++j) { - const uint8_t byte = pk[j]; - bsum += e2m1[byte & 0xF] * xeb[j]; - bsum += e2m1[byte >> 4] * xob[j]; - } - acc += bs * bsum; - } - return acc * global; -} - -// ---- NVFP4 W4A8 (int8 activations) dot: nibble->int8 via LUT, per-16 act scale ---- -// asi8: int8 activations laid out per-16 block as [even(8), odd(8)]; asb[b] = per-block -// activation scale (absmax/127). Result folds the e2m1*2 -> *0.5 into the scale. -using nvi8dot_fn = float (*)(const uint8_t*, const uint8_t*, float, const int8_t*, int, - const float*, const float*); - -[[maybe_unused]] float dot_nvfp4_i8_scalar(const uint8_t* packed, const uint8_t* scale, - float global, const int8_t* asi8, int K, const float* e4m3, - const float* asb) { - float acc = 0.0f; - const int nb = K / 16; - for (int b = 0; b < nb; ++b) { - const uint8_t* pk = packed + (size_t)b * 8; - const int8_t* ae = asi8 + (size_t)b * 16; // even(8) - const int8_t* ao = ae + 8; // odd(8) - int isum = 0; - for (int j = 0; j < 8; ++j) { - isum += (int)kE2M1x2[pk[j] & 0xF] * (int)ae[j]; - isum += (int)kE2M1x2[pk[j] >> 4] * (int)ao[j]; - } - acc += (e4m3[scale[b]] * asb[b]) * (float)isum; - } - return acc * (0.5f * global); -} - -#if CPU_MOE_X86 -// AVX2 e2m1 nibble decode: codes (0..15) in 8 int32 lanes -> fp32. AVX2 vpermps is -// only 8-wide, so instead of a 16-entry LUT we use the e2m1 sign/magnitude symmetry: -// value = (code&8 ? - : +) * mag8[code&7], mag8 = e2m1[0..7]. The sign is bit 3 of -// the code shifted into the fp32 sign bit (bit 31). Bit-identical to the e2m1 LUT. -__attribute__((target("avx2,fma"))) -inline __m256 e2m1_decode8(__m256i codes, __m256 mag8) { - __m256 mag = _mm256_permutevar8x32_ps(mag8, _mm256_and_si256(codes, _mm256_set1_epi32(7))); - __m256i sgn = _mm256_slli_epi32(_mm256_and_si256(codes, _mm256_set1_epi32(8)), 28); - return _mm256_xor_ps(mag, _mm256_castsi256_ps(sgn)); -} - -// Two 16-K blocks (16 packed bytes) per iter: lo nibbles -> even-K, hi -> odd-K, -// gathered via two vpermps. The per-16 e4m3 scale differs across the two blocks, so -// it is applied per lane (low 8 lanes = block b, high 8 = block b+1). -__attribute__((target("avx512f"))) -inline __m512 nvfp4_blk2(const uint8_t* pk, const float* xeb, const float* xob, __m512 lut, - __m512i loma, float s0, float s1) { - __m512i wi = _mm512_cvtepu8_epi32(_mm_loadu_si128(reinterpret_cast(pk))); - __m512 vlo = _mm512_permutexvar_ps(_mm512_and_si512(wi, loma), lut); - __m512 vhi = _mm512_permutexvar_ps(_mm512_and_si512(_mm512_srli_epi32(wi, 4), loma), lut); - __m512 prod = _mm512_fmadd_ps(vlo, _mm512_loadu_ps(xeb), _mm512_mul_ps(vhi, _mm512_loadu_ps(xob))); - // lanes 0-7 (block b) -> s0, lanes 8-15 (block b+1) -> s1 (pure AVX512F mask move) - __m512 scv = _mm512_mask_mov_ps(_mm512_set1_ps(s0), 0xFF00, _mm512_set1_ps(s1)); - return _mm512_mul_ps(prod, scv); -} - -__attribute__((target("avx512f"))) -float dot_nvfp4_avx512(const uint8_t* packed, const uint8_t* scale, float global, - const float* xe, const float* xo, int K, const float* e2m1, - const float* e4m3) { - const __m512 lut = _mm512_loadu_ps(e2m1); // 16 e2m1 values - const __m512i loma = _mm512_set1_epi32(0xF); - __m512 acc0 = _mm512_setzero_ps(), acc1 = _mm512_setzero_ps(); - const int nb = K / 16; // 8 packed bytes + one e4m3 scale per block - int b = 0; - for (; b + 4 <= nb; b += 4) { // two blk2 calls -> 4 blocks, 2 accumulators - acc0 = _mm512_add_ps(acc0, nvfp4_blk2(packed + (size_t)b * 8, xe + (size_t)b * 8, - xo + (size_t)b * 8, lut, loma, e4m3[scale[b]], e4m3[scale[b + 1]])); - acc1 = _mm512_add_ps(acc1, nvfp4_blk2(packed + (size_t)(b + 2) * 8, xe + (size_t)(b + 2) * 8, - xo + (size_t)(b + 2) * 8, lut, loma, e4m3[scale[b + 2]], e4m3[scale[b + 3]])); - } - for (; b + 2 <= nb; b += 2) - acc0 = _mm512_add_ps(acc0, nvfp4_blk2(packed + (size_t)b * 8, xe + (size_t)b * 8, - xo + (size_t)b * 8, lut, loma, e4m3[scale[b]], e4m3[scale[b + 1]])); - float s = _mm512_reduce_add_ps(_mm512_add_ps(acc0, acc1)); - for (; b < nb; ++b) { // odd final 16-K block - const uint8_t* pk = packed + (size_t)b * 8; - const float* xeb = xe + (size_t)b * 8; - const float* xob = xo + (size_t)b * 8; - float bsum = 0.0f; - for (int j = 0; j < 8; ++j) { - bsum += e2m1[pk[j] & 0xF] * xeb[j] + e2m1[pk[j] >> 4] * xob[j]; - } - s += e4m3[scale[b]] * bsum; - } - return s * global; -} - -// AVX2: one 16-K block (8 packed bytes) per call, 8 even + 8 odd lanes. -__attribute__((target("avx2,fma"))) -inline __m256 nvfp4_blk_avx2(const uint8_t* pk, const float* xeb, const float* xob, - __m256 mag8, float sc) { - __m256i wi = _mm256_cvtepu8_epi32(_mm_loadl_epi64(reinterpret_cast(pk))); - __m256 vlo = e2m1_decode8(_mm256_and_si256(wi, _mm256_set1_epi32(0xF)), mag8); - __m256 vhi = e2m1_decode8(_mm256_srli_epi32(wi, 4), mag8); - __m256 prod = _mm256_fmadd_ps(vlo, _mm256_loadu_ps(xeb), _mm256_mul_ps(vhi, _mm256_loadu_ps(xob))); - return _mm256_mul_ps(prod, _mm256_set1_ps(sc)); -} - -__attribute__((target("avx2,fma"))) -float dot_nvfp4_avx2(const uint8_t* packed, const uint8_t* scale, float global, - const float* xe, const float* xo, int K, const float* e2m1, - const float* e4m3) { - const __m256 mag8 = _mm256_loadu_ps(e2m1); // e2m1[0..7] magnitudes - __m256 acc0 = _mm256_setzero_ps(), acc1 = _mm256_setzero_ps(); - const int nb = K / 16; - int b = 0; - for (; b + 2 <= nb; b += 2) { - acc0 = _mm256_add_ps(acc0, nvfp4_blk_avx2(packed + (size_t)b * 8, xe + (size_t)b * 8, - xo + (size_t)b * 8, mag8, e4m3[scale[b]])); - acc1 = _mm256_add_ps(acc1, nvfp4_blk_avx2(packed + (size_t)(b + 1) * 8, xe + (size_t)(b + 1) * 8, - xo + (size_t)(b + 1) * 8, mag8, e4m3[scale[b + 1]])); - } - for (; b < nb; ++b) - acc0 = _mm256_add_ps(acc0, nvfp4_blk_avx2(packed + (size_t)b * 8, xe + (size_t)b * 8, - xo + (size_t)b * 8, mag8, e4m3[scale[b]])); - return hsum256(_mm256_add_ps(acc0, acc1)) * global; -} - -// AVX-VNNI W4A8: decode 8 packed bytes (16 nibbles) of one 16-block to int8 [lo(8),hi(8)] -// via PSHUFB against the e2m1*2 LUT (replaces the 2 vpermps fp32 expands -- ~4x less -// port-5 traffic). lo=even-K weights, hi=odd-K, matching the [even(8),odd(8)] act layout. -__attribute__((target("avx2,avxvnni,fma"))) -inline __m128i nvfp4_decode_block_i8(const uint8_t* pk, __m128i lut) { - __m128i b = _mm_loadl_epi64(reinterpret_cast(pk)); // 8 bytes - __m128i lo = _mm_and_si128(b, _mm_set1_epi8(0x0F)); - __m128i hi = _mm_and_si128(_mm_srli_epi16(b, 4), _mm_set1_epi8(0x0F)); - return _mm_shuffle_epi8(lut, _mm_unpacklo_epi64(lo, hi)); // [lo(8),hi(8)] -> int8 -} - -// Two 16-blocks per VPDPBUSD (32 int8). Sign trick (ggml mul_add_epi8): |w|*(sign(w)*a)=w*a, -// so u8*s8 needs no offset/correction term. Per-block scale (e4m3 * act-scale) folded in fp32 -// (lanes 0-3 -> block b, 4-7 -> block b+1). Bit-faithful weight; only the int8 activation quant -// (W4A8) differs from the bf16 reference. -__attribute__((target("avx2,avxvnni,fma"))) -float dot_nvfp4_i8_vnni(const uint8_t* packed, const uint8_t* scale, float global, - const int8_t* asi8, int K, const float* e4m3, const float* asb) { - const __m128i lut = _mm_loadu_si128(reinterpret_cast(kE2M1x2)); - __m256 accF = _mm256_setzero_ps(); - const int nb = K / 16; - int b = 0; - for (; b + 2 <= nb; b += 2) { - __m128i wb = nvfp4_decode_block_i8(packed + (size_t)b * 8, lut); - __m128i wb1 = nvfp4_decode_block_i8(packed + (size_t)(b + 1) * 8, lut); - __m256i w = _mm256_set_m128i(wb1, wb); // [blk b | blk b+1] - __m256i a = _mm256_loadu_si256(reinterpret_cast(asi8 + (size_t)b * 16)); - __m256i aw = _mm256_sign_epi8(w, w); // |w| (u8 operand) - __m256i sa = _mm256_sign_epi8(a, w); // sign(w)*a (s8 operand) - __m256i di = _mm256_dpbusd_avx_epi32(_mm256_setzero_si256(), aw, sa); - __m256 scv = _mm256_blend_ps(_mm256_set1_ps(e4m3[scale[b]] * asb[b]), - _mm256_set1_ps(e4m3[scale[b + 1]] * asb[b + 1]), 0xF0); - accF = _mm256_fmadd_ps(_mm256_cvtepi32_ps(di), scv, accF); - } - float s = hsum256(accF); - for (; b < nb; ++b) { // tail (odd block count) - const uint8_t* pk = packed + (size_t)b * 8; - const int8_t* ae = asi8 + (size_t)b * 16; const int8_t* ao = ae + 8; - int isum = 0; - for (int j = 0; j < 8; ++j) - isum += (int)kE2M1x2[pk[j] & 0xF] * (int)ae[j] + (int)kE2M1x2[pk[j] >> 4] * (int)ao[j]; - s += (e4m3[scale[b]] * asb[b]) * (float)isum; - } - return s * (0.5f * global); -} - -#if (defined(__GNUC__) && __GNUC__ >= 10) || defined(__clang__) -#define CPU_MOE_HAS_AVX512VNNI 1 - -// Software-prefetch distance for the W4A8 weight stream, in 16-K blocks (8 packed -// bytes each). Returns -1 when FREETOKEN_CPU_MOE_PF_BLOCKS is unset: the kernel then -// uses the built-in default min(512 blocks = 4 KB, 2 rows) -- 4 KB is the empirical -// optimum on large-row machines (Emerald Rapids sweep), while the 2-row cap keeps a small-row -// model's overshoot bounded (the executor works in 32-row tiles, so a fixed byte -// distance otherwise prefetches another worker's tile: duplicated DRAM traffic that -// regresses at the bandwidth ceiling). An EXPLICIT env value is honored verbatim -// (no clamp; 0 disables): the per-machine optimum can sit past the safe default (+20% -// at 4 KB on a 24-thread Ice Lake with 256B rows), so the escape hatch must reach it. -// Prefetch never faults, so overshooting a row/bank tail is safe. -static int nvfp4_pf_blocks() { - static const int v = [] { - const char* s = getenv("FREETOKEN_CPU_MOE_PF_BLOCKS"); - return (s && s[0]) ? atoi(s) : -1; - }(); - return v; -} - -// AVX-512 VNNI W4A8: FOUR 16-K blocks per VPDPBUSD (64 int8) -- 2x the AVX-VNNI -// (256-bit) path. Decode 32 packed bytes -> 64 int8 with a single _mm512_shuffle_epi8 -// (e2m1*2 LUT replicated to all 4 128-bit lanes). AVX-512 has no _mm512_sign_epi8, so -// the u8*s8 sign trick (|w| as u8, sign(w)*a as s8) uses abs_epi8 + a masked negate. -// Bit-faithful weight; only the int8 activation quant (W4A8) differs from bf16. -// One 4-block group (64 int8) -> scaled fp32 partial (16 lanes). Isolated as a helper so -// the caller can run several independent chains into separate accumulators (the decode -> -// dpbusd -> scale chain is long, so a single accumulator leaves the core latency-bound). -__attribute__((target("avx512f,avx512bw,avx512vnni,avx2"))) -static inline __m512 nvfp4_i8_grp4(const uint8_t* packed, const uint8_t* scale, - const int8_t* asi8, const float* e4m3, const float* asb, - int b, __m512i lut, __m512i idx, __m512i mask0F, - __m512i idxsc) { - const __mmask64 hi_half = 0xFF00FF00FF00FF00ULL; // bytes 8-15 of each 128b lane -> hi nibbles - __m256i raw = _mm256_loadu_si256(reinterpret_cast(packed + (size_t)b * 8)); - __m512i src = _mm512_permutexvar_epi64(idx, _mm512_castsi256_si512(raw)); - __m512i lo = _mm512_and_si512(src, mask0F); - __m512i hi = _mm512_and_si512(_mm512_srli_epi16(src, 4), mask0F); - __m512i comb = _mm512_mask_blend_epi8(hi_half, lo, hi); - __m512i w = _mm512_shuffle_epi8(lut, comb); // 64 int8 weights (e2m1*2) - // u8*s8 sign trick without _mm512_sign_epi8: aw=|w|, sa = (w<0 ? -a : a) (a moot at w==0) - __m512i a = _mm512_loadu_si512(reinterpret_cast(asi8 + (size_t)b * 16)); - __m512i aw = _mm512_abs_epi8(w); - __mmask64 neg = _mm512_movepi8_mask(w); - __m512i sa = _mm512_mask_sub_epi8(a, neg, _mm512_setzero_si512(), a); - __m512i di = _mm512_dpbusd_epi32(_mm512_setzero_si512(), aw, sa); // 16 int32 (groups of 4) - // int32[0..3]->blk b, [4..7]->b+1, [8..11]->b+2, [12..15]->b+3. The 4 block scales - // (e4m3 LUT x per-block act scale) are computed vectorized: a 4-byte load + epu8->epi32 - // widen + one 4-lane LUT gather + one mul replaces 8 scalar loads + 4 scalar muls + - // a set_ps assembly, which otherwise dominates the per-group op count. - int sc_raw; - memcpy(&sc_raw, scale + b, 4); - __m128i sc4 = _mm_cvtepu8_epi32(_mm_cvtsi32_si128(sc_raw)); - __m128 s4 = _mm_mul_ps(_mm_i32gather_ps(e4m3, sc4, 4), _mm_loadu_ps(asb + b)); - __m512 scv = _mm512_permutexvar_ps(idxsc, _mm512_castps128_ps512(s4)); - return _mm512_mul_ps(_mm512_cvtepi32_ps(di), scv); -} - -__attribute__((target("avx512f,avx512bw,avx512vnni,avx2"))) -float dot_nvfp4_i8_avx512vnni(const uint8_t* packed, const uint8_t* scale, float global, - const int8_t* asi8, int K, const float* e4m3, const float* asb) { - const __m512i lut = _mm512_broadcast_i32x4( - _mm_loadu_si128(reinterpret_cast(kE2M1x2))); - const __m512i idx = _mm512_set_epi64(3, 3, 2, 2, 1, 1, 0, 0); // block i -> 128b lane i - const __m512i mask0F = _mm512_set1_epi8(0x0F); - const __m512i idxsc = _mm512_set_epi32(3, 3, 3, 3, 2, 2, 2, 2, 1, 1, 1, 1, 0, 0, 0, 0); - __m512 acc0 = _mm512_setzero_ps(), acc1 = _mm512_setzero_ps(); - __m512 acc2 = _mm512_setzero_ps(), acc3 = _mm512_setzero_ps(); - const int nb = K / 16; - // Built-in default: min(4 KB, 2 rows) -- see nvfp4_pf_blocks(). An explicitly set - // env value is used verbatim so a per-machine sweep can reach operating points past - // the conservative default. - const int pfb = nvfp4_pf_blocks(); - const int pf = (pfb < 0) ? std::min(512, 2 * nb) : pfb; - int b = 0; - // Four independent 4-block groups per iter -> four accumulator chains hide the ~10-op - // decode->dpbusd->scale latency (1-2 chains leave the loop latency-bound: the per-core - // rate sat at ~60% of the core's achievable DRAM stream rate). - for (; b + 16 <= nb; b += 16) { - // Prefetch the weight stream ahead: the interleaved decode lowers the L1-miss - // concurrency the HW prefetcher sustains on its own (134 -> 167 GB/s on Emerald Rapids). - if (pf > 0) { - _mm_prefetch(reinterpret_cast(packed + ((size_t)b + (size_t)pf) * 8), - _MM_HINT_T0); - _mm_prefetch(reinterpret_cast(packed + ((size_t)b + (size_t)pf) * 8 + 64), - _MM_HINT_T0); - } - acc0 = _mm512_add_ps(acc0, nvfp4_i8_grp4(packed, scale, asi8, e4m3, asb, b, lut, idx, - mask0F, idxsc)); - acc1 = _mm512_add_ps(acc1, nvfp4_i8_grp4(packed, scale, asi8, e4m3, asb, b + 4, lut, idx, - mask0F, idxsc)); - acc2 = _mm512_add_ps(acc2, nvfp4_i8_grp4(packed, scale, asi8, e4m3, asb, b + 8, lut, idx, - mask0F, idxsc)); - acc3 = _mm512_add_ps(acc3, nvfp4_i8_grp4(packed, scale, asi8, e4m3, asb, b + 12, lut, idx, - mask0F, idxsc)); - } - for (; b + 4 <= nb; b += 4) - acc0 = _mm512_add_ps(acc0, nvfp4_i8_grp4(packed, scale, asi8, e4m3, asb, b, lut, idx, - mask0F, idxsc)); - float s = _mm512_reduce_add_ps( - _mm512_add_ps(_mm512_add_ps(acc0, acc1), _mm512_add_ps(acc2, acc3))); - for (; b < nb; ++b) { // tail (<4 remaining 16-K blocks) - const uint8_t* pk = packed + (size_t)b * 8; - const int8_t* ae = asi8 + (size_t)b * 16; const int8_t* ao = ae + 8; - int isum = 0; - for (int j = 0; j < 8; ++j) - isum += (int)kE2M1x2[pk[j] & 0xF] * (int)ae[j] + (int)kE2M1x2[pk[j] >> 4] * (int)ao[j]; - s += (e4m3[scale[b]] * asb[b]) * (float)isum; - } - return s * (0.5f * global); -} -#endif // avx512vnni available -#endif - -// ===================================================================================== -// CUDA stream memory operations (driver API, resolved via dlopen -- no link-time or -// toolchain dependence). The GPU side of the flag handshake: submit = WRITE_VALUE -// (done[slot]=0 then ready[slot]=1), sync = WAIT_VALUE(done[slot] >= 1). The wait is -// executed by the GPU front-end (no SM-resident kernel), so GPU "utilization" stays -// truthful during CPU compute windows -- a resident spin kernel pinned it at 99%, -// which laptop CPU/GPU dynamic power schedulers answered by clamping the CPU's max -// frequency (GEMV workers -1.5x: the reported edge regression). Availability is -// probed functionally at startup (memops_probe); anything unsupported (Windows WDDM, -// vGPU, old drivers) falls back to the cudaLaunchHostFunc path. -#if defined(_WIN32) -#include -static void* cumemop_dlopen() { return (void*)::LoadLibraryA("nvcuda.dll"); } -static void* cumemop_dlsym(void* h, const char* n) { - return (void*)::GetProcAddress((HMODULE)h, n); -} -#else -#include -static void* cumemop_dlopen() { - void* h = dlopen("libcuda.so.1", RTLD_LAZY | RTLD_LOCAL); - if (h == nullptr) h = dlopen("libcuda.so", RTLD_LAZY | RTLD_LOCAL); - return h; -} -static void* cumemop_dlsym(void* h, const char* n) { return dlsym(h, n); } -#endif - -using cuMemOp64_fn = int (*)(void* stream, unsigned long long addr, unsigned long long value, - unsigned int flags); -static cuMemOp64_fn g_cu_write64 = nullptr; -static cuMemOp64_fn g_cu_wait64 = nullptr; -static constexpr unsigned int kCuWaitValueGeq = 0x0; // CU_STREAM_WAIT_VALUE_GEQ -static constexpr unsigned int kCuWriteDefault = 0x0; // CU_STREAM_WRITE_VALUE_DEFAULT - -static bool cumemop_resolve() { - static bool resolved = [] { - void* h = cumemop_dlopen(); - if (h == nullptr) return false; - // 11.7+ made the v2 entry points the default; older drivers export only the v1 - // names with the same signature. - g_cu_write64 = reinterpret_cast(cumemop_dlsym(h, "cuStreamWriteValue64_v2")); - if (g_cu_write64 == nullptr) - g_cu_write64 = reinterpret_cast(cumemop_dlsym(h, "cuStreamWriteValue64")); - g_cu_wait64 = reinterpret_cast(cumemop_dlsym(h, "cuStreamWaitValue64_v2")); - if (g_cu_wait64 == nullptr) - g_cu_wait64 = reinterpret_cast(cumemop_dlsym(h, "cuStreamWaitValue64")); - return g_cu_write64 != nullptr && g_cu_wait64 != nullptr; - }(); - return resolved; -} - -// Functional probe on a scratch pinned int64: enqueue WRITE(7) + WAIT(>=7) + sync. -// Returns true only if the whole memop path works on THIS stream/device/driver. -static bool cumemops_probe(uintptr_t stream, uintptr_t scratch_addr) { - if (!cumemop_resolve()) return false; - auto* s = reinterpret_cast(stream); - if (g_cu_write64(s, (unsigned long long)scratch_addr, 7ULL, kCuWriteDefault) != 0) return false; - if (g_cu_wait64(s, (unsigned long long)scratch_addr, 7ULL, kCuWaitValueGeq) != 0) return false; - return cudaStreamSynchronize(reinterpret_cast(stream)) == cudaSuccess; -} - -// GPU side of the flag handshake (see the block comment above): enqueued on the -// caller's (possibly capturing) stream; the WAIT immediate is the constant 1, -// replay-safe under CUDA graphs. -// The startup probe validates EAGER memops; a driver could still reject them at graph -// capture time. Those enqueue errors would otherwise be swallowed here and surface only -// as a later EndCapture failure -- log the first CUresult so triage is one step. -static void cumemop_check(int rc, const char* what) { - static std::atomic warned{false}; - if (rc != 0 && !warned.exchange(true)) { - std::fprintf(stderr, - "[freetoken/cpu_moe] %s failed with CUresult=%d (first occurrence; " - "subsequent errors are not repeated). If this happened during CUDA " - "graph capture, the driver lacks capture support for stream memops -- " - "set FREETOKEN_CPU_MOE_FLAG_SYNC=0.\n", - what, rc); - } -} - -static void cumemop_submit(uintptr_t stream, uintptr_t done_addr, uintptr_t ready_addr, - int64_t slot) { - auto* s = reinterpret_cast(stream); - // Order matters and is preserved by the front end: reset done BEFORE raising ready, - // so the coordinator's completion write for THIS step can never be wiped. - cumemop_check(g_cu_write64(s, (unsigned long long)(done_addr + (size_t)slot * 8), 0ULL, - kCuWriteDefault), - "cuStreamWriteValue64(done)"); - cumemop_check(g_cu_write64(s, (unsigned long long)(ready_addr + (size_t)slot * 8), 1ULL, - kCuWriteDefault), - "cuStreamWriteValue64(ready)"); -} - -static void cumemop_sync(uintptr_t stream, uintptr_t done_addr, int64_t slot) { - cumemop_check(g_cu_wait64(reinterpret_cast(stream), - (unsigned long long)(done_addr + (size_t)slot * 8), 1ULL, - kCuWaitValueGeq), - "cuStreamWaitValue64(done)"); -} - -struct DotChoice { - dot_fn fn; - const char* name; -}; - -// SIMD tiers, ascending. Each format picks the highest tier <= the one chosen by -// pick_isa() that it implements (fp4 and fp8_block formats have no bf16-specific tier, so the -// avx512bf16 tier maps to their avx512 kernel). -enum IsaTier { ISA_SCALAR = 0, ISA_AVX2 = 1, ISA_AVX512 = 2, ISA_AVX512BF16 = 3 }; - -// Best tier the CPU+build supports, optionally capped DOWN by -// FREETOKEN_CPU_MOE_ISA={scalar,avx2,avx512,avx512bf16} (A/B testing on a machine -// that supports more). FREETOKEN_CPU_MOE_SCALAR=1 forces scalar (legacy alias). -inline IsaTier pick_isa() { -#if CPU_MOE_X86 - if (getenv("FREETOKEN_CPU_MOE_SCALAR")) return ISA_SCALAR; - IsaTier best = ISA_SCALAR; - if (__builtin_cpu_supports("avx2") && __builtin_cpu_supports("fma")) best = ISA_AVX2; - if (best >= ISA_AVX2 && __builtin_cpu_supports("avx512f")) best = ISA_AVX512; -#ifdef CPU_MOE_HAS_AVX512BF16 - if (best >= ISA_AVX512 && __builtin_cpu_supports("avx512bf16")) best = ISA_AVX512BF16; -#endif - if (const char* f = getenv("FREETOKEN_CPU_MOE_ISA")) { - IsaTier want = best; - if (!std::strcmp(f, "scalar")) want = ISA_SCALAR; - else if (!std::strcmp(f, "avx2")) want = ISA_AVX2; - else if (!std::strcmp(f, "avx512")) want = ISA_AVX512; - else if (!std::strcmp(f, "avx512bf16")) want = ISA_AVX512BF16; - if (want < best) best = want; // cap downward; never force above hw/build support - } - return best; -#else - return ISA_SCALAR; -#endif -} - -DotChoice select_dot() { - const IsaTier t = pick_isa(); -#if CPU_MOE_X86 -#ifdef CPU_MOE_HAS_AVX512BF16 - if (t >= ISA_AVX512BF16) return {dot_avx512bf16, "avx512bf16"}; -#endif - if (t >= ISA_AVX512) return {dot_avx512f, "avx512f"}; - if (t >= ISA_AVX2) return {dot_avx2, "avx2"}; -#endif - (void)t; - return {dot_scalar, "scalar"}; -} - -nvdot_fn select_nvdot() { - const IsaTier t = pick_isa(); -#if CPU_MOE_X86 - if (t >= ISA_AVX512) return dot_nvfp4_avx512; - if (t >= ISA_AVX2) return dot_nvfp4_avx2; -#endif - (void)t; - return dot_nvfp4_scalar; -} - -// AVX-VNNI (VEX-256 VPDPBUSD) availability: Alder/Raptor Lake, Sapphire Rapids+, Zen5. -// Distinct from AVX-512 VNNI. Opt out with FREETOKEN_CPU_MOE_NO_VNNI=1 (A/B the W4A8 path). -inline bool cpu_has_avxvnni() { -#if CPU_MOE_X86 - const char* no = getenv("FREETOKEN_CPU_MOE_NO_VNNI"); - if (no && no[0] && no[0] != '0') return false; // ignore unset/empty/"0" - return __builtin_cpu_supports("avx2") && __builtin_cpu_supports("avxvnni"); -#else - return false; -#endif -} +namespace { -// AVX-512 VNNI (512-bit VPDPBUSD): Cascade Lake+, Ice Lake, Sapphire/Emerald Rapids, -// Zen4+. 2x the 256-bit AVX-VNNI width. FREETOKEN_CPU_MOE_NO_AVX512VNNI=1 forces the -// 256-bit path (A/B the two W4A8 kernels on the same box); FREETOKEN_CPU_MOE_NO_VNNI=1 -// still disables the whole W4A8 family (back to the faithful fp32 nvdot), so it is -// honored here too. -inline bool cpu_has_avx512vnni() { -#if CPU_MOE_X86 && defined(CPU_MOE_HAS_AVX512VNNI) - const char* no = getenv("FREETOKEN_CPU_MOE_NO_AVX512VNNI"); - if (no && no[0] && no[0] != '0') return false; - const char* no_vnni = getenv("FREETOKEN_CPU_MOE_NO_VNNI"); - if (no_vnni && no_vnni[0] && no_vnni[0] != '0') return false; - return __builtin_cpu_supports("avx512vnni"); +inline void pin_this_thread(int cpu) { +#if CPU_MOE_HAS_AFFINITY + cpu_set_t set; + CPU_ZERO(&set); + CPU_SET(cpu, &set); + pthread_setaffinity_np(pthread_self(), sizeof(set), &set); #else - return false; -#endif -} - -// Best W4A8 (int8-activation) nvfp4 dot, or nullptr if no SIMD VNNI (caller keeps the -// faithful fp32 nvdot path). The scalar i8 dot exists only as a correctness reference. -nvi8dot_fn select_nvi8dot() { -#if CPU_MOE_X86 -#if defined(CPU_MOE_HAS_AVX512VNNI) - if (cpu_has_avx512vnni()) return dot_nvfp4_i8_avx512vnni; -#endif - if (cpu_has_avxvnni()) return dot_nvfp4_i8_vnni; -#endif - return nullptr; -} - -// ----------------------- DeepSeek-V4 ds_fp4 (W4A8) --------------------------- -// Row-major e2m1 (2/byte, low nibble first) + e8m0 per-32 block scale, no global -// (w = E2M1[code] * 2^(e8m0-127)); activations are FP8-e4m3 round-tripped (per-128 -// block, ue8m0 scale) before each GEMM. Matches kernel/triton/dsv4 (fused_moe + -// fp8_linear): silu(clamp(gate,max=lim)) * clamp(up,-lim,lim), router weight on the -// down output. - -// Activations are pre-deinterleaved to fp32 (xe[m]=x[2m], xo[m]=x[2m+1]) once per -// token/route and reused across all output rows. This drops the hot dot to one -// vpmovzxbd + two vpermps per 32 weights (1.5 shuffle ops / 16 vs 4 for a bf16, -// dup-permute, per-element gather), so the row-major fp4 GEMV stops being port-5 -// bound and approaches the bf16 memory-bandwidth ceiling. -using dsdot_fn = float (*)(const uint8_t*, const uint8_t*, const float*, const float*, int, - const float*, const float*); - -float dot_dsfp4_scalar(const uint8_t* packed, const uint8_t* scale, const float* xe, - const float* xo, int K, const float* e2m1, const float* e8m0) { - float acc = 0.0f; - const int nb = K / 32; - for (int b = 0; b < nb; ++b) { - const float sc = e8m0[scale[b]]; - const uint8_t* pk = packed + (size_t)b * 16; // 16 bytes = 32 codes - const float* xeb = xe + (size_t)b * 16; - const float* xob = xo + (size_t)b * 16; - float bsum = 0.0f; - for (int j = 0; j < 16; ++j) { - const uint8_t byte = pk[j]; - bsum += e2m1[byte & 0xF] * xeb[j]; // low nibble -> even-K activation - bsum += e2m1[byte >> 4] * xob[j]; // high nibble -> odd-K activation - } - acc += sc * bsum; - } - return acc; -} - -#if CPU_MOE_X86 -// One 32-block: 16 bytes -> 16 low + 16 high nibble values via two vpermps, times -// the pre-split even/odd fp32 activations, folded by the per-32 e8m0 scale. -__attribute__((target("avx512f"))) -inline __m512 dsfp4_blk(const uint8_t* pk, const float* xeb, const float* xob, __m512 lut, - __m512i loma, float sc) { - __m512i wi = _mm512_cvtepu8_epi32(_mm_loadu_si128(reinterpret_cast(pk))); - __m512 vlo = _mm512_permutexvar_ps(_mm512_and_si512(wi, loma), lut); - __m512 vhi = _mm512_permutexvar_ps(_mm512_and_si512(_mm512_srli_epi32(wi, 4), loma), lut); - __m512 prod = _mm512_fmadd_ps(vlo, _mm512_loadu_ps(xeb), _mm512_mul_ps(vhi, _mm512_loadu_ps(xob))); - return _mm512_mul_ps(prod, _mm512_set1_ps(sc)); -} - -__attribute__((target("avx512f"))) -float dot_dsfp4_avx512(const uint8_t* packed, const uint8_t* scale, const float* xe, - const float* xo, int K, const float* e2m1, const float* e8m0) { - const __m512 lut = _mm512_loadu_ps(e2m1); - const __m512i loma = _mm512_set1_epi32(0xF); - __m512 acc0 = _mm512_setzero_ps(), acc1 = _mm512_setzero_ps(); - const int nb = K / 32; // 16 packed bytes + one e8m0 scale per block - int b = 0; - for (; b + 2 <= nb; b += 2) { // two independent accumulators hide FMA latency - acc0 = _mm512_add_ps(acc0, dsfp4_blk(packed + (size_t)b * 16, xe + (size_t)b * 16, - xo + (size_t)b * 16, lut, loma, e8m0[scale[b]])); - acc1 = _mm512_add_ps(acc1, dsfp4_blk(packed + (size_t)(b + 1) * 16, xe + (size_t)(b + 1) * 16, - xo + (size_t)(b + 1) * 16, lut, loma, e8m0[scale[b + 1]])); - } - for (; b < nb; ++b) - acc0 = _mm512_add_ps(acc0, dsfp4_blk(packed + (size_t)b * 16, xe + (size_t)b * 16, - xo + (size_t)b * 16, lut, loma, e8m0[scale[b]])); - return _mm512_reduce_add_ps(_mm512_add_ps(acc0, acc1)); -} - -// AVX2: a 32-K block is 16 bytes -> two 8-lane halves (8 even + 8 odd each). -__attribute__((target("avx2,fma"))) -inline __m256 dsfp4_half_avx2(const uint8_t* pk, const float* xeb, const float* xob, __m256 mag8) { - __m256i wi = _mm256_cvtepu8_epi32(_mm_loadl_epi64(reinterpret_cast(pk))); - __m256 vlo = e2m1_decode8(_mm256_and_si256(wi, _mm256_set1_epi32(0xF)), mag8); - __m256 vhi = e2m1_decode8(_mm256_srli_epi32(wi, 4), mag8); - return _mm256_fmadd_ps(vlo, _mm256_loadu_ps(xeb), _mm256_mul_ps(vhi, _mm256_loadu_ps(xob))); -} - -__attribute__((target("avx2,fma"))) -float dot_dsfp4_avx2(const uint8_t* packed, const uint8_t* scale, const float* xe, - const float* xo, int K, const float* e2m1, const float* e8m0) { - const __m256 mag8 = _mm256_loadu_ps(e2m1); - __m256 acc0 = _mm256_setzero_ps(), acc1 = _mm256_setzero_ps(); - const int nb = K / 32; - for (int b = 0; b < nb; ++b) { - const uint8_t* pk = packed + (size_t)b * 16; - const float* xeb = xe + (size_t)b * 16; - const float* xob = xo + (size_t)b * 16; - const __m256 sc = _mm256_set1_ps(e8m0[scale[b]]); - acc0 = _mm256_fmadd_ps(dsfp4_half_avx2(pk, xeb, xob, mag8), sc, acc0); - acc1 = _mm256_fmadd_ps(dsfp4_half_avx2(pk + 8, xeb + 8, xob + 8, mag8), sc, acc1); - } - return hsum256(_mm256_add_ps(acc0, acc1)); -} -#endif - -dsdot_fn select_dsdot() { - const IsaTier t = pick_isa(); -#if CPU_MOE_X86 - if (t >= ISA_AVX512) return dot_dsfp4_avx512; - if (t >= ISA_AVX2) return dot_dsfp4_avx2; -#endif - (void)t; - return dot_dsfp4_scalar; -} - -// ------------------------- mxfp4 (gpt-oss) GEMV ----------------------------- -// Transposed split-K layout: blk[Kpairs, N2] (N innermost), scl[Kpairs/16, N2] -// e8m0 per 32-K. Computes out[c] = sum_kb (E2M1[lo]*x[2kb] + E2M1[hi]*x[2kb+1]) -// * 2^(e8m0-127) for a contiguous column tile (blk/scl already offset to col 0 of -// the tile). Vectorized over N (16 columns / __m512), K stays the outer (cache- -// sequential) loop. Used by both gate_up (K=H) and down (K=I). -using mxgemv_fn = void (*)(float*, const uint8_t*, const uint8_t*, const bf16_t*, int, int, - int, const float*, const float*); - -void mxfp4_gemv_scalar(float* out, const uint8_t* blk, const uint8_t* scl, const bf16_t* x, - int Kpairs, int N2, int ncol, const float* e2m1, const float* e8m0) { - for (int c = 0; c < ncol; ++c) out[c] = 0.0f; - for (int kb = 0; kb < Kpairs; ++kb) { - const uint8_t* w = blk + (size_t)kb * N2; - const uint8_t* s = scl + (size_t)(kb >> 4) * N2; - const float xl = bf16_to_f32(x[2 * kb]); - const float xh = bf16_to_f32(x[2 * kb + 1]); - for (int c = 0; c < ncol; ++c) { - const uint8_t byte = w[c]; - out[c] += (e2m1[byte & 0xF] * xl + e2m1[byte >> 4] * xh) * e8m0[s[c]]; - } - } -} - -#if CPU_MOE_X86 -__attribute__((target("avx512f"))) -void mxfp4_gemv_avx512(float* out, const uint8_t* blk, const uint8_t* scl, const bf16_t* x, - int Kpairs, int N2, int ncol, const float* e2m1, const float* e8m0) { - (void)e8m0; // e8m0[c]=2^(c-127) computed via bit construction (no gather) - const __m512 lut = _mm512_loadu_ps(e2m1); - const __m512i loma = _mm512_set1_epi32(0xF); - // K-outer / N-inner: each kb cache line is read once and all live column chunks - // (up to 4 -> 64 cols) accumulate from registers, so DRAM/L2 stream the tile once. - int c0 = 0; - for (; c0 + 16 <= ncol; c0 += 64) { - const int nchunk = std::min(4, (ncol - c0) / 16); - __m512 acc[4]; - for (int ci = 0; ci < nchunk; ++ci) acc[ci] = _mm512_setzero_ps(); - for (int kblk = 0; kblk < Kpairs; kblk += 16) { // 16 K-pairs = 32 K = one scale row - __m512 sc[4]; - for (int ci = 0; ci < nchunk; ++ci) { - __m128i sraw = _mm_loadu_si128(reinterpret_cast( - scl + (size_t)(kblk >> 4) * N2 + c0 + ci * 16)); - sc[ci] = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu8_epi32(sraw), 23)); - } - __m512 blk_acc[4]; - for (int ci = 0; ci < nchunk; ++ci) blk_acc[ci] = _mm512_setzero_ps(); - for (int kk = 0; kk < 16; ++kk) { - const int kb = kblk + kk; - const uint8_t* wbase = blk + (size_t)kb * N2 + c0; - // The transposed layout strides K by N2 bytes; prefetch ahead so the strided - // reads are not exposed to DRAM latency (the HW streamer misses big strides). - constexpr int PFD = 8; - if (kb + PFD < Kpairs) - _mm_prefetch(reinterpret_cast(blk + (size_t)(kb + PFD) * N2 + c0), - _MM_HINT_T0); - const __m512 xl = _mm512_set1_ps(bf16_to_f32(x[2 * kb])); - const __m512 xh = _mm512_set1_ps(bf16_to_f32(x[2 * kb + 1])); - for (int ci = 0; ci < nchunk; ++ci) { - __m512i wi = _mm512_cvtepu8_epi32( - _mm_loadu_si128(reinterpret_cast(wbase + ci * 16))); - __m512 vlo = _mm512_permutexvar_ps(_mm512_and_si512(wi, loma), lut); - __m512 vhi = _mm512_permutexvar_ps(_mm512_and_si512(_mm512_srli_epi32(wi, 4), loma), lut); - blk_acc[ci] = _mm512_fmadd_ps(vlo, xl, blk_acc[ci]); - blk_acc[ci] = _mm512_fmadd_ps(vhi, xh, blk_acc[ci]); - } - } - for (int ci = 0; ci < nchunk; ++ci) acc[ci] = _mm512_fmadd_ps(blk_acc[ci], sc[ci], acc[ci]); - } - for (int ci = 0; ci < nchunk; ++ci) _mm512_storeu_ps(out + c0 + ci * 16, acc[ci]); - } - for (int c = c0; c < ncol; ++c) { // tail columns (< 16) - float o = 0.0f; - for (int kb = 0; kb < Kpairs; ++kb) { - const uint8_t byte = blk[(size_t)kb * N2 + c]; - uint32_t bits = (uint32_t)scl[(size_t)(kb >> 4) * N2 + c] << 23; - float sc; - std::memcpy(&sc, &bits, 4); - o += (e2m1[byte & 0xF] * bf16_to_f32(x[2 * kb]) + - e2m1[byte >> 4] * bf16_to_f32(x[2 * kb + 1])) * sc; - } - out[c] = o; - } -} - -__attribute__((target("avx2,fma"))) -void mxfp4_gemv_avx2(float* out, const uint8_t* blk, const uint8_t* scl, const bf16_t* x, - int Kpairs, int N2, int ncol, const float* e2m1, const float* e8m0) { - (void)e8m0; // e8m0[s]=2^(s-127) built via s<<23 (no gather) - const __m256 mag8 = _mm256_loadu_ps(e2m1); - int c0 = 0; - for (; c0 + 8 <= ncol; c0 += 32) { // up to 4 chunks of 8 = 32 cols - const int nchunk = std::min(4, (ncol - c0) / 8); - __m256 acc[4]; - for (int ci = 0; ci < nchunk; ++ci) acc[ci] = _mm256_setzero_ps(); - for (int kblk = 0; kblk < Kpairs; kblk += 16) { // 16 K-pairs = one scale row - __m256 sc[4]; - for (int ci = 0; ci < nchunk; ++ci) { - __m128i sraw = _mm_loadl_epi64(reinterpret_cast( - scl + (size_t)(kblk >> 4) * N2 + c0 + ci * 8)); - sc[ci] = _mm256_castsi256_ps(_mm256_slli_epi32(_mm256_cvtepu8_epi32(sraw), 23)); - } - __m256 blk_acc[4]; - for (int ci = 0; ci < nchunk; ++ci) blk_acc[ci] = _mm256_setzero_ps(); - for (int kk = 0; kk < 16; ++kk) { - const int kb = kblk + kk; - const uint8_t* wbase = blk + (size_t)kb * N2 + c0; - constexpr int PFD = 8; - if (kb + PFD < Kpairs) - _mm_prefetch(reinterpret_cast(blk + (size_t)(kb + PFD) * N2 + c0), - _MM_HINT_T0); - const __m256 xl = _mm256_set1_ps(bf16_to_f32(x[2 * kb])); - const __m256 xh = _mm256_set1_ps(bf16_to_f32(x[2 * kb + 1])); - for (int ci = 0; ci < nchunk; ++ci) { - __m256i wi = _mm256_cvtepu8_epi32( - _mm_loadl_epi64(reinterpret_cast(wbase + ci * 8))); - __m256 vlo = e2m1_decode8(_mm256_and_si256(wi, _mm256_set1_epi32(0xF)), mag8); - __m256 vhi = e2m1_decode8(_mm256_srli_epi32(wi, 4), mag8); - blk_acc[ci] = _mm256_fmadd_ps(vlo, xl, blk_acc[ci]); - blk_acc[ci] = _mm256_fmadd_ps(vhi, xh, blk_acc[ci]); - } - } - for (int ci = 0; ci < nchunk; ++ci) acc[ci] = _mm256_fmadd_ps(blk_acc[ci], sc[ci], acc[ci]); - } - for (int ci = 0; ci < nchunk; ++ci) _mm256_storeu_ps(out + c0 + ci * 8, acc[ci]); - } - for (int c = c0; c < ncol; ++c) { // tail columns (< 8); none when ncol%8==0 - float o = 0.0f; - for (int kb = 0; kb < Kpairs; ++kb) { - const uint8_t byte = blk[(size_t)kb * N2 + c]; - uint32_t bits = (uint32_t)scl[(size_t)(kb >> 4) * N2 + c] << 23; - float sc; - std::memcpy(&sc, &bits, 4); - o += (e2m1[byte & 0xF] * bf16_to_f32(x[2 * kb]) + - e2m1[byte >> 4] * bf16_to_f32(x[2 * kb + 1])) * sc; - } - out[c] = o; - } -} -#endif - -fp8_block_dot_fn select_fp8_block_dot() { - const IsaTier t = pick_isa(); -#if CPU_MOE_X86 - // the avx512bf16 tier keeps the fp32 kernel: on Sapphire Rapids it is faster than widening to bf16 for dpbf16 - if (t >= ISA_AVX512) return dot_fp8_block_avx512f; - if (t >= ISA_AVX2) return dot_fp8_block_avx2; + (void)cpu; #endif - (void)t; - return dot_fp8_block_scalar; } -mxgemv_fn select_mxgemv() { - const IsaTier t = pick_isa(); +inline void cpu_relax() { #if CPU_MOE_X86 - if (t >= ISA_AVX512) return mxfp4_gemv_avx512; - if (t >= ISA_AVX2) return mxfp4_gemv_avx2; + _mm_pause(); +#elif defined(__aarch64__) + asm volatile("yield" ::: "memory"); #endif - (void)t; - return mxfp4_gemv_scalar; -} - -// Round a clamped |x|<=448 to nearest float8-e4m3 (RNE), back to fp32. Matches -// torch.float8_e4m3fn / triton .to(float8e4nv). -inline float e4m3_round(float x) { - const float sign = x < 0.0f ? -1.0f : 1.0f; - const float a = std::fabs(x); - if (a == 0.0f) return 0.0f; - if (a >= 448.0f) return sign * 448.0f; - int e; - std::frexp(a, &e); // a in [2^(e-1), 2^e) - float step = std::ldexp(1.0f, e - 4); - const float min_step = std::ldexp(1.0f, -9); // e4m3 subnormal step (2^-9) - if (step < min_step) step = min_step; - float r = std::nearbyint(a / step) * step; - if (r > 448.0f) r = 448.0f; - return sign * r; -} - -// IEEE ceil(log2(v)) for v>0 (matches dsv4 _log2_ceil / fast_round_scale). -inline int ceil_log2_pos(float v) { - uint32_t bits; - std::memcpy(&bits, &v, sizeof(bits)); - const int exp = (int)((bits >> 23) & 0xFF); - const int man = (int)(bits & 0x7FFFFF); - return exp - 127 + (man != 0 ? 1 : 0); -} - -// Split an interleaved bf16 row into fp32 even/odd halves (even[m]=src[2m]). -// bf16->fp32 is exact, so this only reorders -- done once per token/route and -// reused across every output row of the GEMV. -inline void deinterleave_bf16_f32(const bf16_t* src, float* even, float* odd, int K) { - for (int m = 0; m < K / 2; ++m) { - even[m] = bf16_to_f32(src[2 * m]); - odd[m] = bf16_to_f32(src[2 * m + 1]); - } -} - -// DeepSeek-V4 / V4.1 activation FP8 round-trip (bf16 in/out): per `block` (the checkpoint's -// fp8 block: 128 on V4, 32 on V4.1), s = 2^ceil(log2(max(|x|,1e-4)/448)); -// y = round_e4m3(clamp(x/s,+-448)) * s. -void fp8_roundtrip_bf16(const bf16_t* src, bf16_t* dst, int K, int block) { - for (int b0 = 0; b0 < K; b0 += block) { - const int b1 = std::min(K, b0 + block); - float amax = 1e-4f; - for (int i = b0; i < b1; ++i) amax = std::max(amax, std::fabs(bf16_to_f32(src[i]))); - const float s = std::ldexp(1.0f, ceil_log2_pos(amax * (1.0f / 448.0f))); - const float inv_s = 1.0f / s; - for (int i = b0; i < b1; ++i) { - float q = bf16_to_f32(src[i]) * inv_s; - q = std::min(448.0f, std::max(-448.0f, q)); - dst[i] = f32_to_bf16(e4m3_round(q) * s); - } - } } -// --------------------------------- executor --------------------------------- - -struct CpuMoeExecutor; +constexpr int kMaxTopK = 64; +// One MoE layer on the first num_tokens rows of the pinned IO. +// It is also the doorbell word the GPU writes, a constant baked into each captured graph; num_tokens >= 1 keeps it nonzero, since 0 means no doorbell. struct MoeTask { - CpuMoeExecutor* exec; - int layer_id; - int num_tokens; - const bf16_t* x; // [num_tokens, H] - const int32_t* ids; // [num_tokens, top_k] (raw expert ids; <0 = skip) - const float* w; // [num_tokens, top_k] - bf16_t* y; // [num_tokens, H] + uint64_t layer : 32; + uint64_t num_tokens : 32; }; +static_assert(sizeof(MoeTask) == sizeof(uint64_t), "MoeTask must pack into one 64-bit word"); -// Output-row tiling. Small enough to give every worker independent work even at -// batch size 1; large enough to amortize the atomic work-grab. -// -// Bandwidth notes (Sapphire Rapids 8480+, 13 cores): the two passes already read -// every expert weight byte exactly once per token (each output row block is owned -// by one worker), and x stays hot in L1 across a (token,expert)'s rows -- so the -// kernel is single-read bandwidth-optimal at bs=1 (~205 GB/s vs ~55 GB/s PCIe). -// One worker per *physical* core, pinned, is the sweet spot; SMT oversubscription -// thrashes the spin-barrier. Deferred (not worth it here / for this workload): -// - AMX-bf16: a GEMM tile engine; decode is M=1 GEMV so tiles sit idle. It would -// only pay off in a grouped/batched (dedup) path. -// - expert dedup for bs>1: read each distinct expert once and GEMM its tokens. -// Helps locality+bytes when bs is large; decode batches here are tiny (<=4). -// - NUMA: a single node is assumed. Multi-socket machines would split each -// expert's K dimension per node (banks are already per-row contiguous). -constexpr int IBLK = 32; -constexpr int HBLK = 32; - -// -------------------------------- Q4_0 (W4A8) -------------------------------- -// Native GGUF Q4_0 experts (gemma4 GGUF): per-32 block = fp16 scale d + 16 packed -// bytes; byte j holds element j in its low nibble and j+16 in its high nibble, so a -// block's storage order is [lo0..lo15, hi0..hi15] and w = (nibble - 8) * d. Matches -// the reference dequant (models/gguf/dequant.py) and the packed banks the GPU offload -// path streams. -// -// llama.cpp ggml_vec_dot_q4_0_q8_0: W4A8. The activation is pre-quantized to Q8_0 -// (per-32-block int8 ``aq`` + fp32 scale ``asb``); each block unpacks its 16 bytes to -// 32 int8 weights in [-8,7] (bytes_from_nibbles_32: low nibbles -> elems 0..15, high -// -> 16..31) and runs an integer block dot -- VPDPBUSD (AVX-VNNI) or VPMADDUBSW+VPMADDWD -// (AVX2) with the ggml sign trick |w|*(sign(w)*a)=w*a, or a scalar int loop -- then -// scales the block sum by wd*xd in fp32. No fp weight dequant / shuffle chain. The GPU -// offload path (ggml_moe_a8_vec / MMVQ) is also W4A8, so cpu and hybrid stay close. -using q4dot_fn = float (*)(const uint8_t*, const int8_t*, const float*, int); - -float q4_0_dot_i8_scalar(const uint8_t* w, const int8_t* aq, const float* asb, int K) { - float acc = 0.0f; - const int nb = K / 32; - for (int b = 0; b < nb; ++b) { - const uint8_t* blk = w + (size_t)b * 18; - uint16_t dh; - std::memcpy(&dh, blk, sizeof(dh)); - const uint8_t* q = blk + 2; // 16 nibble bytes - const int8_t* a = aq + (size_t)b * 32; - int isum = 0; - for (int j = 0; j < 16; ++j) { - isum += ((int)(q[j] & 0x0F) - 8) * (int)a[j]; // elem j - isum += ((int)(q[j] >> 4) - 8) * (int)a[16 + j]; // elem 16+j - } - acc += fp16_to_f32(dh) * asb[b] * (float)isum; - } - return acc; -} - -#if CPU_MOE_X86 -// fp16 block scale -> fp32 via HW F16C (single value in lane 0). -__attribute__((target("f16c"))) -static inline float q4_scale(uint16_t h) { - return _mm_cvtss_f32(_mm_cvtph_ps(_mm_cvtsi32_si128((int)h))); +inline uint64_t to_word(MoeTask t) { + uint64_t w; + std::memcpy(&w, &t, sizeof(w)); + return w; } -// Unpack one Q4_0 block's 16 bytes -> 32 int8 weights in [-8,7] (elems 0..15 = low -// nibbles, 16..31 = high nibbles). ``eight`` = _mm256_set1_epi8(8). -__attribute__((target("avx2"))) -static inline __m256i q4_unpack32(const uint8_t* blk, __m128i mask, __m256i eight) { - const __m128i qb = _mm_loadu_si128(reinterpret_cast(blk + 2)); - const __m128i lo = _mm_and_si128(qb, mask); - const __m128i hi = _mm_and_si128(_mm_srli_epi16(qb, 4), mask); - return _mm256_sub_epi8(_mm256_set_m128i(hi, lo), eight); +inline MoeTask from_word(uint64_t w) { + MoeTask t; + std::memcpy(&t, &w, sizeof(t)); + return t; } -// AVX2 W4A8 (llama.cpp non-VNNI mul_sum_i8_pairs): integer block dot via VPMADDUBSW + -// VPMADDWD (sign trick), scaled by wd*xd. |aw*sa| pair sums <= 8*127*2 < 32767 -> no -// int16 saturation. This is the fast path on AVX2 CPUs without AVX-VNNI (and the -// avx512-tier fallback, since the block dot is 256-bit either way). -__attribute__((target("avx2,fma,f16c"))) -float q4_0_dot_i8_avx2(const uint8_t* w, const int8_t* aq, const float* asb, int K) { - const __m128i mask = _mm_set1_epi8(0x0F); - const __m256i eight = _mm256_set1_epi8(8); - const __m256i ones16 = _mm256_set1_epi16(1); - __m256 accF = _mm256_setzero_ps(); - const int nb = K / 32; - for (int b = 0; b < nb; ++b) { - const uint8_t* blk = w + (size_t)b * 18; - _mm_prefetch(reinterpret_cast(blk) + 512, _MM_HINT_T0); - uint16_t dh; - std::memcpy(&dh, blk, sizeof(dh)); - __m256i wq = q4_unpack32(blk, mask, eight); - __m256i a = _mm256_loadu_si256(reinterpret_cast(aq + (size_t)b * 32)); - __m256i aw = _mm256_sign_epi8(wq, wq); // |wq| (unsigned operand) - __m256i sa = _mm256_sign_epi8(a, wq); // sign(wq) * a (signed operand) - __m256i d32 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, sa), ones16); // 8 int32 - accF = _mm256_fmadd_ps(_mm256_cvtepi32_ps(d32), _mm256_set1_ps(q4_scale(dh) * asb[b]), accF); - } - return hsum256(accF); -} - -// AVX-VNNI W4A8: one VPDPBUSD per block (the fast path on modern CPUs). -__attribute__((target("avx2,avxvnni,fma,f16c"))) -float q4_0_dot_i8_vnni(const uint8_t* w, const int8_t* aq, const float* asb, int K) { - const __m128i mask = _mm_set1_epi8(0x0F); - const __m256i eight = _mm256_set1_epi8(8); - __m256 accF = _mm256_setzero_ps(); - const int nb = K / 32; - for (int b = 0; b < nb; ++b) { - const uint8_t* blk = w + (size_t)b * 18; - _mm_prefetch(reinterpret_cast(blk) + 512, _MM_HINT_T0); - uint16_t dh; - std::memcpy(&dh, blk, sizeof(dh)); - __m256i wq = q4_unpack32(blk, mask, eight); - __m256i a = _mm256_loadu_si256(reinterpret_cast(aq + (size_t)b * 32)); - __m256i aw = _mm256_sign_epi8(wq, wq); // |wq| (unsigned operand) - __m256i sa = _mm256_sign_epi8(a, wq); // sign(wq) * a (signed operand) - __m256i di = _mm256_dpbusd_avx_epi32(_mm256_setzero_si256(), aw, sa); - // All 32 elems of the block share wd*xd; distribute over di's 8 partial sums and - // reduce at the end (equivalent to scale * block_total). - accF = _mm256_fmadd_ps(_mm256_cvtepi32_ps(di), _mm256_set1_ps(q4_scale(dh) * asb[b]), accF); - } - return hsum256(accF); -} -#endif // CPU_MOE_X86 - -// All tiers are W4A8 (int8 activations pre-quantized to Q8_0). AVX-VNNI is orthogonal to -// the ISA tier (gated by cpu_has_avxvnni() / FREETOKEN_CPU_MOE_NO_VNNI), so it wins when -// present; otherwise the 256-bit VPMADDUBSW kernel covers both the avx2 and avx512 tiers. -q4dot_fn select_q4dot() { - const IsaTier t = pick_isa(); -#if CPU_MOE_X86 - if (cpu_has_avxvnni()) return q4_0_dot_i8_vnni; - if (t >= ISA_AVX2) return q4_0_dot_i8_avx2; -#endif - (void)t; - return q4_0_dot_i8_scalar; -} - -enum WFmt { WF_BF16 = 0, WF_NVFP4 = 1, WF_MXFP4 = 2, WF_DSFP4 = 3, WF_Q4_0 = 4, - WF_FP8_BLOCK = 5 }; +// Every weight format this module serves, keyed by the name cpu_executor.py passes; weight_formats() exports the names. +const std::pair (*)(const MoeParams&, const FormatArgs&)> kFormats[] = { + {"bf16", make_bf16_format}, {"nvfp4", make_nvfp4_format}, {"mxfp4_triton", make_mxfp4_format}, + {"ds_fp4", make_ds_fp4_format}, {"q4_0", make_q4_0_format}, {"fp8_block", make_fp8_block_format}, +}; -// Each ctor pointer arg is the address of a CPU int64 array of length -// num_layers (one base address per layer, built by cpu_executor.py's -// _make_table), not a single flat bank. tbl_at resolves -// tbl[layer_id] once per task/pass; a null table (bank unused by this fmt, ptr -// arg 0) resolves to nullptr without dereferencing. -inline const void* tbl_at(const uint64_t* tbl, int layer_id) { - return tbl ? reinterpret_cast(tbl[layer_id]) : nullptr; +std::unique_ptr make_format(const std::string& name, const MoeParams& p, const FormatArgs& a) { + for (const auto& [n, make] : kFormats) + if (name == n) return make(p, a); + throw std::invalid_argument("unknown CPU MoE weight format " + name); } struct CpuMoeExecutor { + // A host-function callback gets one pointer, so it carries the executor with the task. + struct HostTask { + CpuMoeExecutor* exec; + MoeTask task; + }; + + const int num_layers, num_experts; + const bool apply_on_input; + MoeParams params; // the format keeps a reference; declared before fmt + std::unique_ptr fmt; + bool preps_inter = false; + bool rounds_routes = false; + int n_iblk = 0, n_hblk = 0; // kRowBlock blocks per intermediate row and per output row + + // Pinned IO shared by every layer and batch size (one task in flight): [max_tokens, H] x and y, [max_tokens, top_k] ids and w. + const bf16_t* io_x; + const int32_t* io_ids; + const float* io_w; + bf16_t* io_y; + AlignedVector g_scratch; // [max_tokens * top_k * I] intermediate + + // Worker pool. int num_threads; - int num_layers, num_experts, top_k; - int H, I; - int act, apply_on_input; - int fmt; // WFmt - bool needs_di = false; // pre-deinterleave activations to fp32 (nvfp4/ds_fp4/fp8_block) - // Per-layer pointer tables (one base address per layer, see tbl_at). gate_up_tbl - // doubles as the bf16 gate_up table and the nvfp4/mxfp4/q4_0/ds_fp4 packed-gate_up - // table (down_tbl likewise for down); which reinterpretation applies is picked by - // fmt at each resolve site (see gemm1_dot/gemm2_dot/do_pass1_mxfp4/do_pass1_dsfp4). - const uint64_t* gate_up_tbl; // bf16: [E,2I,H] rows; else: packed e2m1/mxfp4-blocks - const uint64_t* down_tbl; // bf16: [E,H,I] rows; else: packed e2m1/mxfp4-blocks - const uint64_t* gu_scale_tbl; // nvfp4/mxfp4/ds_fp4: [E,2I,*] block scales - const uint64_t* gu_global_tbl; // nvfp4: [E,2I] fp16 row globals - const uint64_t* dn_scale_tbl; // nvfp4/mxfp4/ds_fp4: [E,H,*] block scales - const uint64_t* dn_global_tbl; // nvfp4: [E,H] fp16 row globals - const uint64_t* gu_bias_tbl; // mxfp4: [E,2I] bf16 biases - const uint64_t* dn_bias_tbl; // mxfp4: [E,H] bf16 biases - float swiglu_alpha; - float swiglu_limit; // +inf == no clamp - dot_fn dot; - nvdot_fn nvdot; - nvi8dot_fn nvi8dot = nullptr; // AVX-VNNI W4A8 nvfp4 dot (nullptr -> use fp32 nvdot) - bool use_vnni = false; // nvfp4 + AVX-VNNI: decode via int8 VPDPBUSD (W4A8) - bool use_q4a8 = false; // q4_0: always W4A8 (llama.cpp Q4_0 x Q8_0); int8 pre-quant - dsdot_fn dsdot; - mxgemv_fn mxgemv; - q4dot_fn q4dot; - fp8_block_dot_fn fp8_block_dot; - // fp8_block: bf16 scale per 128x128 weight block. These are the scale tensor's - // trailing strides, including any padding used by the host bank layout. - int fp8_block_gu_scale_row = 0, fp8_block_dn_scale_row = 0; - // ds_fp4: the caller already FP8-round-tripped the input activations on the GPU - // (same reference grid), so submit() must not repeat it on the host-callback - // thread. That scalar per-element pass is single-threaded ON THE DECODE CRITICAL - // PATH (~0.3ms/layer at H=4096, every worker and the GPU waiting on it); moving - // it to a captured GPU elementwise kernel removes it while keeping the official - // W4A8 numerics bit-exact. Set via set_input_prequant (see cpu_executor.py). - bool input_prequant = false; - // ds_fp4 activation round-trip block (the checkpoint's fp8 block); set via set_act_block. - int act_block = 128; - // Q4_0 packed-row byte strides (H/32*18 for gate_up over K=H, I/32*18 for down over K=I). - int q4_gu_row_bytes = 0, q4_dn_row_bytes = 0; - float e2m1_lut[16]; - float e4m3_lut[256]; - float e8m0_lut[256]; // mxfp4 block scale: 2^(s-127), s clamped to [0,254] - const char* isa; - - std::vector g_scratch; // [max_tokens * top_k * I] intermediate - std::vector xq_scratch; // [max_tokens * H] ds_fp4 fp8-roundtripped input - // nvfp4/ds_fp4/fp8_block activations pre-deinterleaved to fp32 (even/odd K) for the row-major dot. - std::vector xe_scratch, xo_scratch; // [max_tokens * H/2] (input) - std::vector ge_scratch, go_scratch; // [max_tokens*top_k*I/2] (intermediate) - // AVX-VNNI W4A8: per-16-block int8 activations [even(8),odd(8)] + per-block scale. - std::vector xi8_scratch, gi8_scratch; // [max_tokens*H], [max_tokens*top_k*I] - std::vector xas_scratch, gas_scratch; // [max_tokens*H/16], [..*top_k*I/16] - std::string isa_str; - + std::vector core_ids; // worker tid -> logical CPU to pin to (may be empty) std::vector workers; std::mutex task_mtx; std::condition_variable task_cv; @@ -1421,7 +136,7 @@ struct CpuMoeExecutor { std::atomic stop{false}; // spinning workers read it without the lock uint64_t cur_gen = 0; - MoeTask* cur_task = nullptr; + MoeTask cur_task{}; // Separate cache lines: idle workers poll `submitted` while busy workers hit p1_next/p2_next. alignas(64) std::atomic submitted{0}; alignas(64) std::atomic completed{0}; @@ -1431,264 +146,56 @@ struct CpuMoeExecutor { std::atomic spin_wait{false}; bool coord_pinned = false; + // Work split of the current task. alignas(64) std::atomic p1_next{0}; std::atomic p2_next{0}; - std::atomic prt_next{0}; // intermediate-row prep phase (prep_g_row) + std::atomic prt_next{0}; // intermediate-row prep phase (prep_inter) int64_t p1_total = 0, p2_total = 0, prt_total = 0; - int n_iblk = 0, n_hblk = 0; std::atomic done_count{0}; std::atomic bar_count{0}; std::atomic bar_sense{0}; - std::vector owned_tasks; // persistent task descriptors (graph-stable) - std::vector core_ids; // worker tid -> logical CPU to pin to (may be empty) - // ---- Flag-based GPU<->CPU handshake (replaces the per-layer cudaLaunchHostFunc pair) ---- - // A tiny GPU kernel bumps ready_flags[slot] at submit; this coordinator thread busy-polls - // it, runs the slot's task on the worker pool, and sets done_flags[slot], which a GPU - // spin-wait kernel polls at sync. This removes the ~2x30-50us host-func dispatch round - // trips per MoE layer per decode step that otherwise idle the GPU (~6 ms/step on a - // 75-layer model). One slot per (layer, decode batch size) pair -- the Python side - // allocates slots as tasks are created. Flags live in mapped-pinned host memory (UVA: - // the same pointers are used by the GPU kernels and by this thread). + // The GPU rings the doorbell with a task word at submit; this coordinator thread busy-polls it, runs the task on the worker pool, and completes it, which the GPU waits on at sync. + // This removes the ~2x30-50us host-func dispatch round trips per MoE layer per decode step that otherwise idle the GPU (~6 ms/step on a 75-layer model). + std::unique_ptr gpu_sync; + std::unique_ptr kernel_io; // replaces the copies around the doorbell once enable_kernel_io() runs std::thread coord_thread; std::atomic coord_stop{false}; - volatile int64_t* ready_flags = nullptr; // GPU increments, this thread polls - volatile int64_t* done_flags = nullptr; // this thread sets, GPU spin-waits - int coord_num_slots = 0; - std::vector flag_task; // slot -> task (registered lazily) - std::vector flag_served; // slot -> completed dispatch count (tests/debug) - std::mutex flag_task_mtx; - - // Portable ordering for the flag handshake: "ready observed => the DMA'd inputs that - // preceded the bump are visible" and "y stores are visible before done". Plain - // volatile loads lean on x86 TSO; acquire/release makes it hold on aarch64 too - // (GH200/Jetson) at zero x86 cost. (MSVC branch is x86-only today: compiler barrier - // + TSO.) - static int64_t flag_load_acquire(const volatile int64_t* p) { -#if defined(_MSC_VER) - const int64_t v = *p; - _ReadWriteBarrier(); - return v; -#else - return __atomic_load_n(const_cast(p), __ATOMIC_ACQUIRE); -#endif - } - - static void flag_store_release(volatile int64_t* p, int64_t v) { -#if defined(_MSC_VER) - _ReadWriteBarrier(); - *p = v; -#else - __atomic_store_n(const_cast(p), v, __ATOMIC_RELEASE); -#endif - } + std::atomic flag_served{0}; // completed doorbells (tests, watchdog) + std::map host_tasks; // host-function userData, keyed by doorbell word (graph-stable) CpuMoeExecutor(int num_threads_, int num_layers_, int num_experts_, int top_k_, - int hidden_size, int inter_size, int max_tokens, int activation_id, - int apply_router_weight_on_input, int weight_format, - uintptr_t gate_up_ptr, uintptr_t down_ptr, uintptr_t gate_up_scale_ptr, - uintptr_t gate_up_global_ptr, uintptr_t down_scale_ptr, - uintptr_t down_global_ptr, uintptr_t gate_up_bias_ptr, - uintptr_t down_bias_ptr, double swiglu_alpha_, double swiglu_limit_, - std::vector core_ids_, int fp8_block_gu_scale_stride_, - int fp8_block_dn_scale_stride_) - : num_threads(num_threads_ > 0 ? num_threads_ : 1), - num_layers(num_layers_), + int hidden_size, int inter_size, int max_tokens_, int activation_id, + int apply_router_weight_on_input, const std::string& weight_format, + std::map banks, std::map options, + double swiglu_alpha_, double swiglu_limit_, + std::vector core_ids_, uintptr_t x_ptr, uintptr_t ids_ptr, uintptr_t w_ptr, + uintptr_t y_ptr) + : num_layers(num_layers_), num_experts(num_experts_), - top_k(top_k_), - H(hidden_size), - I(inter_size), - act(activation_id), - apply_on_input(apply_router_weight_on_input), - fmt(weight_format), - gate_up_tbl(reinterpret_cast(gate_up_ptr)), - down_tbl(reinterpret_cast(down_ptr)), - gu_scale_tbl(reinterpret_cast(gate_up_scale_ptr)), - gu_global_tbl(reinterpret_cast(gate_up_global_ptr)), - dn_scale_tbl(reinterpret_cast(down_scale_ptr)), - dn_global_tbl(reinterpret_cast(down_global_ptr)), - gu_bias_tbl(reinterpret_cast(gate_up_bias_ptr)), - dn_bias_tbl(reinterpret_cast(down_bias_ptr)), - swiglu_alpha(static_cast(swiglu_alpha_)), - swiglu_limit(static_cast(swiglu_limit_)), - fp8_block_gu_scale_row(fp8_block_gu_scale_stride_), - fp8_block_dn_scale_row(fp8_block_dn_scale_stride_), + apply_on_input(apply_router_weight_on_input != 0), + params{top_k_, max_tokens_, hidden_size, inter_size, activation_id, static_cast(swiglu_alpha_), + static_cast(swiglu_limit_)}, + io_x(reinterpret_cast(x_ptr)), + io_ids(reinterpret_cast(ids_ptr)), + io_w(reinterpret_cast(w_ptr)), + io_y(reinterpret_cast(y_ptr)), + num_threads(num_threads_ > 0 ? num_threads_ : 1), core_ids(std::move(core_ids_)) { - DotChoice c = select_dot(); - dot = c.fn; - nvdot = select_nvdot(); - dsdot = select_dsdot(); - mxgemv = select_mxgemv(); - q4dot = select_q4dot(); - fp8_block_dot = select_fp8_block_dot(); - if (weight_format == WF_Q4_0) { - if (H % 32 != 0 || I % 32 != 0) - throw std::runtime_error("Q4_0 CPU MoE requires H and I to be multiples of 32"); - q4_gu_row_bytes = (H / 32) * 18; // K = H (gate_up rows) - q4_dn_row_bytes = (I / 32) * 18; // K = I (down rows) - } - isa = c.name; - // nvfp4 (AVX-VNNI only): W4A8 int8 decode when the CPU supports it. q4_0 is always - // W4A8 (activations pre-quantized to Q8_0); select_q4dot picks VPDPBUSD / VPMADDUBSW - // / scalar for the tier, so the tag reflects which of those q4dot resolved to. - nvi8dot = select_nvi8dot(); - use_vnni = (weight_format == WF_NVFP4) && (nvi8dot != nullptr); - use_q4a8 = (weight_format == WF_Q4_0); - const char* q4tag = use_q4a8 ? (cpu_has_avxvnni() ? "+vnni(q4_0-w4a8)" : "+q4_0-w4a8") : ""; - const char* vnni_tag = - cpu_has_avx512vnni() ? "+avx512vnni(nvfp4-w4a8)" : "+vnni(nvfp4-w4a8)"; - isa_str = std::string(c.name) + (use_vnni ? vnni_tag : "") + q4tag; - isa = isa_str.c_str(); - for (int i = 0; i < 16; ++i) e2m1_lut[i] = kE2M1[i]; - for (int i = 0; i < 256; ++i) e4m3_lut[i] = e4m3_decode((uint8_t)i); - // e8m0 (mxfp4 block scale) = 2^(s-127); the GPU GEMV clamps s to [0,254]. - for (int i = 0; i < 256; ++i) e8m0_lut[i] = std::ldexp(1.0f, std::min(i, 254) - 127); - g_scratch.assign(static_cast(max_tokens) * top_k * I, 0); - // Row-major fp4 (nvfp4/ds_fp4) and fp8_block pre-deinterleave activations to fp32 even/odd. - needs_di = (fmt == WF_NVFP4 || fmt == WF_DSFP4 || fmt == WF_FP8_BLOCK); - if (needs_di) { - if (fmt == WF_DSFP4) xq_scratch.assign(static_cast(max_tokens) * H, 0); - xe_scratch.assign(static_cast(max_tokens) * (H / 2), 0); - xo_scratch.assign(static_cast(max_tokens) * (H / 2), 0); - ge_scratch.assign(static_cast(max_tokens) * top_k * (I / 2), 0); - go_scratch.assign(static_cast(max_tokens) * top_k * (I / 2), 0); - if (use_vnni) { - xi8_scratch.assign(static_cast(max_tokens) * H, 0); - xas_scratch.assign(static_cast(max_tokens) * (H / 16), 0); - gi8_scratch.assign(static_cast(max_tokens) * top_k * I, 0); - gas_scratch.assign(static_cast(max_tokens) * top_k * (I / 16), 0); - } - } - // q4_0 W4A8: per-32-block Q8_0 activations (int8 + fp32 scale) for input + intermediate. - if (use_q4a8) { - xi8_scratch.assign(static_cast(max_tokens) * H, 0); - xas_scratch.assign(static_cast(max_tokens) * (H / 32), 0); - gi8_scratch.assign(static_cast(max_tokens) * top_k * I, 0); - gas_scratch.assign(static_cast(max_tokens) * top_k * (I / 32), 0); - } + if (params.top_k > kMaxTopK) throw std::invalid_argument("CPU MoE supports top_k <= " + std::to_string(kMaxTopK)); + fmt = make_format(weight_format, params, FormatArgs{weight_format, std::move(banks), std::move(options)}); + preps_inter = fmt->preps_inter(); + rounds_routes = fmt->rounds_routes(); + n_iblk = (params.I + kRowBlock - 1) / kRowBlock; + n_hblk = (params.H + kRowBlock - 1) / kRowBlock; + g_scratch.assign(static_cast(params.max_tokens) * params.top_k * params.I, 0); for (int t = 0; t < num_threads; ++t) workers.emplace_back([this, t] { worker_loop(t); }); } - // Quantize a bf16 activation row to Q8_0 (llama.cpp): per-32-block symmetric int8 in - // natural order + one fp32 scale (amax/127) per block. Done once per token/route, - // amortized over every output row that the q4_0 W4A8 GEMV reads. K % 32 == 0. - void quant_q8_0(const bf16_t* x, int K, int8_t* aq, float* asb) { - const int nb = K / 32; - for (int b = 0; b < nb; ++b) { - const bf16_t* xb = x + (size_t)b * 32; - float xf[32], amax = 0.0f; - for (int j = 0; j < 32; ++j) { - xf[j] = bf16_to_f32(xb[j]); - amax = std::max(amax, std::fabs(xf[j])); - } - const float d = amax > 0.0f ? amax / 127.0f : 1.0f; - asb[b] = d; - const float inv = amax > 0.0f ? 1.0f / d : 0.0f; - int8_t* o = aq + (size_t)b * 32; - for (int j = 0; j < 32; ++j) - o[j] = (int8_t)std::max(-127, std::min(127, (int)std::lround(xf[j] * inv))); - } - } - - // Quantize the pre-deinterleaved fp32 even/odd activations to per-16-block int8 in the - // [even(8),odd(8)] layout the VNNI dot expects. Done once per token/route (amortized over - // every output row), so a scalar pass is fine relative to the GEMV. - void quant_i8_pg16(const float* xe, const float* xo, int K, int8_t* asi8, float* asb) { - const int nb = K / 16; - for (int b = 0; b < nb; ++b) { - const float* xeb = xe + (size_t)b * 8; - const float* xob = xo + (size_t)b * 8; - float amax = 0.0f; - for (int j = 0; j < 8; ++j) - amax = std::max(amax, std::max(std::fabs(xeb[j]), std::fabs(xob[j]))); - const float s = amax > 0.0f ? amax / 127.0f : 1.0f; - asb[b] = s; - const float inv = 1.0f / s; - int8_t* ae = asi8 + (size_t)b * 16; - for (int j = 0; j < 8; ++j) { - int qe = (int)std::lround(xeb[j] * inv), qo = (int)std::lround(xob[j] * inv); - ae[j] = (int8_t)std::max(-127, std::min(127, qe)); - ae[8 + j] = (int8_t)std::max(-127, std::min(127, qo)); - } - } - } - - // gate_up output row `row` (in [0, 2I)) dotted with activation over K = H. ``e`` is - // the layer-local expert row (0..num_experts); the layer bases (already resolved - // once per task/pass by the caller via tbl_at) pick the layer's own tensors. - // bf16 uses the interleaved bf16 row; nvfp4 uses the pre-split fp32 even/odd halves - // (or, with AVX-VNNI, the per-16-block int8 activations). - // fp8_block uses the halves for whole 128-blocks and the bf16 row for a ragged tail. - inline float gemm1_dot(const bf16_t* gate_up_l, const uint8_t* gu_packed_l, - const uint8_t* gu_scale_l, const uint16_t* gu_global_l, int e, int row, - const bf16_t* x, const float* xe, const float* xo, const int8_t* xi8, - const float* xas) { - if (fmt == WF_BF16) { - const bf16_t* w = gate_up_l + ((size_t)e * (2 * I) + row) * H; - return dot(w, x, H); - } - if (fmt == WF_Q4_0) { - const uint8_t* w = - gu_packed_l + ((size_t)e * (2 * I) + row) * (size_t)q4_gu_row_bytes; - return q4dot(w, xi8, xas, H); // W4A8: int8 activations (Q8_0), scale in xas - } - if (fmt == WF_FP8_BLOCK) { - const uint8_t* w = gu_packed_l + ((size_t)e * (2 * I) + row) * (size_t)H; - const bf16_t* sc = reinterpret_cast(gu_scale_l) + - ((size_t)e * ((2 * I + FP8_BLK - 1) / FP8_BLK) + row / FP8_BLK) * - (size_t)fp8_block_gu_scale_row; - return fp8_block_dot(w, xe, xo, x, sc, H, e4m3_lut); - } - const size_t r = (size_t)e * (2 * I) + row; - if (use_vnni) - return nvi8dot(gu_packed_l + r * (size_t)(H / 2), gu_scale_l + r * (size_t)(H / 16), - fp16_to_f32(gu_global_l[r]), xi8, H, e4m3_lut, xas); - return nvdot(gu_packed_l + r * (size_t)(H / 2), gu_scale_l + r * (size_t)(H / 16), - fp16_to_f32(gu_global_l[r]), xe, xo, H, e2m1_lut, e4m3_lut); - } - - // down output row `row` (in [0, H)) dotted with the intermediate over K = I. Same - // layer-local-base convention as gemm1_dot. - inline float gemm2_dot(const bf16_t* down_l, const uint8_t* dn_packed_l, - const uint8_t* dn_scale_l, const uint16_t* dn_global_l, int e, int row, - const bf16_t* g, const float* ge, const float* go, const int8_t* gi8, - const float* gas) { - if (fmt == WF_BF16) { - const bf16_t* w = down_l + ((size_t)e * H + row) * I; - return dot(w, g, I); - } - if (fmt == WF_Q4_0) { - const uint8_t* w = dn_packed_l + ((size_t)e * H + row) * (size_t)q4_dn_row_bytes; - return q4dot(w, gi8, gas, I); // W4A8: int8 activations (Q8_0), scale in gas - } - if (fmt == WF_FP8_BLOCK) { - const uint8_t* w = dn_packed_l + ((size_t)e * H + row) * (size_t)I; - const bf16_t* sc = reinterpret_cast(dn_scale_l) + - ((size_t)e * ((H + FP8_BLK - 1) / FP8_BLK) + row / FP8_BLK) * - (size_t)fp8_block_dn_scale_row; - return fp8_block_dot(w, ge, go, g, sc, I, e4m3_lut); - } - const size_t r = (size_t)e * H + row; - if (use_vnni) - return nvi8dot(dn_packed_l + r * (size_t)(I / 2), dn_scale_l + r * (size_t)(I / 16), - fp16_to_f32(dn_global_l[r]), gi8, I, e4m3_lut, gas); - return nvdot(dn_packed_l + r * (size_t)(I / 2), dn_scale_l + r * (size_t)(I / 16), - fp16_to_f32(dn_global_l[r]), ge, go, I, e2m1_lut, e4m3_lut); - } - void pin_self(int tid) { -#if CPU_MOE_HAS_AFFINITY - if (core_ids.empty()) return; - const int cpu = core_ids[tid % static_cast(core_ids.size())]; - cpu_set_t set; - CPU_ZERO(&set); - CPU_SET(cpu, &set); - pthread_setaffinity_np(pthread_self(), sizeof(set), &set); -#else - (void)tid; -#endif + if (!core_ids.empty()) pin_this_thread(core_ids[tid % static_cast(core_ids.size())]); } ~CpuMoeExecutor() { @@ -1701,30 +208,23 @@ struct CpuMoeExecutor { task_cv.notify_all(); for (auto& th : workers) if (th.joinable()) th.join(); - for (MoeTask* t : owned_tasks) delete t; } - uintptr_t create_task(int layer_id, int num_tokens, uintptr_t x_ptr, - uintptr_t ids_ptr, uintptr_t w_ptr, uintptr_t y_ptr) { - MoeTask* t = new MoeTask{this, - layer_id, - num_tokens, - reinterpret_cast(x_ptr), - reinterpret_cast(ids_ptr), - reinterpret_cast(w_ptr), - reinterpret_cast(y_ptr)}; - owned_tasks.push_back(t); - return reinterpret_cast(t); + MoeTask task(int layer_id, int num_tokens) const { + if (layer_id < 0 || layer_id >= num_layers || num_tokens < 1 || num_tokens > params.max_tokens) + throw std::out_of_range("CPU MoE task (layer " + std::to_string(layer_id) + ", " + + std::to_string(num_tokens) + " tokens) is outside the executor's layers or IO"); + return MoeTask{static_cast(layer_id), static_cast(num_tokens)}; } - const char* isa_name() const { return isa; } + std::string isa_name() const { return fmt->isa_name(); } - const char* fp8_block_isa_name() const { -#if CPU_MOE_X86 - if (fp8_block_dot == dot_fp8_block_avx512f) return "avx512f"; - if (fp8_block_dot == dot_fp8_block_avx2) return "avx2"; -#endif - return "scalar"; + // Whether any route of the task names an expert; do_pass1 and do_pass2 skip the others. + bool has_routes(const MoeTask& t) const { + const int64_t n = static_cast(t.num_tokens) * params.top_k; + for (int64_t i = 0; i < n; ++i) + if (io_ids[i] >= 0 && io_ids[i] < num_experts) return true; + return false; } void barrier(int& local_sense) { @@ -1733,275 +233,58 @@ struct CpuMoeExecutor { bar_count.store(0); bar_sense.store(local_sense); } else { - while (bar_sense.load() != local_sense) { -#if CPU_MOE_X86 - _mm_pause(); -#endif - } + while (bar_sense.load() != local_sense) cpu_relax(); } } - void do_pass1(const MoeTask* t, int64_t p) { - if (fmt == WF_MXFP4) { - do_pass1_mxfp4(t, p); - return; - } - if (fmt == WF_DSFP4) { - do_pass1_dsfp4(t, p); - return; - } + // Pass 1 item p = one kRowBlock block of one (token, route)'s intermediate row. + void do_pass1(const MoeTask& t, int64_t p) { + const int top_k = params.top_k, H = params.H, I = params.I; const int64_t ib = p % n_iblk; const int64_t tk = p / n_iblk; const int k = static_cast(tk % top_k); const int tok = static_cast(tk / top_k); - const int e = t->ids[static_cast(tok) * top_k + k]; + const int e = io_ids[static_cast(tok) * top_k + k]; if (e < 0 || e >= num_experts) return; - const float w_in = apply_on_input ? t->w[static_cast(tok) * top_k + k] : 1.0f; - // Resolve this task's layer base once; row indexing below is layer-local (e). - const bf16_t* gate_up_l = reinterpret_cast(tbl_at(gate_up_tbl, t->layer_id)); - const uint8_t* gu_packed_l = reinterpret_cast(tbl_at(gate_up_tbl, t->layer_id)); - const uint8_t* gu_scale_l = reinterpret_cast(tbl_at(gu_scale_tbl, t->layer_id)); - const uint16_t* gu_global_l = - reinterpret_cast(tbl_at(gu_global_tbl, t->layer_id)); - const bf16_t* x_row = t->x + (size_t)tok * H; - const float* xe = needs_di ? xe_scratch.data() + (size_t)tok * (H / 2) : nullptr; - const float* xo = needs_di ? xo_scratch.data() + (size_t)tok * (H / 2) : nullptr; - const int8_t* xi8 = - (use_vnni || use_q4a8) ? xi8_scratch.data() + (size_t)tok * H : nullptr; - const float* xas = use_vnni ? xas_scratch.data() + (size_t)tok * (H / 16) - : use_q4a8 ? xas_scratch.data() + (size_t)tok * (H / 32) - : nullptr; - bf16_t* g_row = g_scratch.data() + ((size_t)tok * top_k + k) * I; - const int i0 = static_cast(ib) * IBLK; - const int i1 = std::min(I, i0 + IBLK); - const bool clamped = act == ACT_SWIGLUOAI || act == ACT_SWIGLU_CLAMP; - const float up_bias = act == ACT_SWIGLUOAI ? 1.0f : 0.0f; - const float lim = swiglu_limit, alpha = swiglu_alpha; - for (int i = i0; i < i1; ++i) { - // gate = row i, up = row I+i - float gate = - gemm1_dot(gate_up_l, gu_packed_l, gu_scale_l, gu_global_l, e, i, x_row, xe, xo, xi8, xas) * w_in; - float up = gemm1_dot(gate_up_l, gu_packed_l, gu_scale_l, gu_global_l, e, I + i, x_row, - xe, xo, xi8, xas) * w_in; - if (clamped) { - // clamp(gate, max=lim) * sigmoid(alpha * gate) * (clamp(up, +-lim) + up_bias) - // -- swigluoai carries the +1 up bias (gpt-oss/MiniMax); swiglu_clamp - // (GLM-5.3) does not. lim == +inf: no clamp. - if (gate > lim) gate = lim; - if (up > lim) up = lim; - else if (up < -lim) up = -lim; - const float glu = gate / (1.0f + std::exp(-gate * alpha)); - g_row[i] = f32_to_bf16(glu * (up + up_bias)); - } else { - g_row[i] = f32_to_bf16(act_apply(act, gate) * up); - } - } + const float w_in = apply_on_input ? io_w[static_cast(tok) * top_k + k] : 1.0f; + const int i0 = static_cast(ib) * kRowBlock; + const int i1 = std::min(I, i0 + kRowBlock); + fmt->gate_up(static_cast(t.layer), e, tok, i0, i1, io_x + (size_t)tok * H, w_in, + g_scratch.data() + ((size_t)tok * top_k + k) * I); } - void do_pass2(const MoeTask* t, int64_t p) { - if (fmt == WF_MXFP4) { - do_pass2_mxfp4(t, p); - return; - } - if (fmt == WF_DSFP4) { - do_pass2_dsfp4(t, p); - return; - } - const int64_t hb = p % n_hblk; - const int tok = static_cast(p / n_hblk); - const int h0 = static_cast(hb) * HBLK; - const int h1 = std::min(H, h0 + HBLK); - // Resolve this task's layer base once; row indexing below is layer-local (e). - const bf16_t* down_l = reinterpret_cast(tbl_at(down_tbl, t->layer_id)); - const uint8_t* dn_packed_l = reinterpret_cast(tbl_at(down_tbl, t->layer_id)); - const uint8_t* dn_scale_l = reinterpret_cast(tbl_at(dn_scale_tbl, t->layer_id)); - const uint16_t* dn_global_l = - reinterpret_cast(tbl_at(dn_global_tbl, t->layer_id)); - bf16_t* y_row = t->y + (size_t)tok * H; - for (int h = h0; h < h1; ++h) { - float acc = 0.0f; - for (int k = 0; k < top_k; ++k) { - const int e = t->ids[static_cast(tok) * top_k + k]; - if (e < 0 || e >= num_experts) continue; - const float w_out = apply_on_input ? 1.0f : t->w[static_cast(tok) * top_k + k]; - const size_t gr = (size_t)tok * top_k + k; - const bf16_t* g_row = g_scratch.data() + gr * I; - const float* ge = needs_di ? ge_scratch.data() + gr * (I / 2) : nullptr; - const float* go = needs_di ? go_scratch.data() + gr * (I / 2) : nullptr; - const int8_t* gi8 = (use_vnni || use_q4a8) ? gi8_scratch.data() + gr * I : nullptr; - const float* gas = use_vnni ? gas_scratch.data() + gr * (I / 16) - : use_q4a8 ? gas_scratch.data() + gr * (I / 32) - : nullptr; - acc += gemm2_dot(down_l, dn_packed_l, dn_scale_l, dn_global_l, e, h, g_row, ge, go, gi8, - gas) * w_out; - } - y_row[h] = f32_to_bf16(acc); - } - } - - // ----------------------------- mxfp4 (gpt-oss) ----------------------------- - // Transposed split-K layout (N innermost), so the GEMV streams K and accumulates - // a contiguous block of N output columns -> cache-line-efficient, no repack, no - // extra host memory. Pass 1 fuses gate_up + clamped-swiglu(+bias); pass 2 fuses - // down(+bias) * router-weight, summed over the token's routes. - // - // Dequant: w = E2M1[code] * 2^(e8m0_scale - 127); two codes per byte (low nibble - // first), one e8m0 scale per 32 contiguous K. Matches kernel/triton/mxfp4_moe.py. - - void do_pass1_mxfp4(const MoeTask* t, int64_t p) { - const int64_t ib = p % n_iblk; - const int64_t tk = p / n_iblk; - const int k = static_cast(tk % top_k); - const int tok = static_cast(tk / top_k); - const int e = t->ids[static_cast(tok) * top_k + k]; - if (e < 0 || e >= num_experts) return; - // Resolve this task's layer base once; row indexing below is layer-local (e). - const uint8_t* gu_packed_l = reinterpret_cast(tbl_at(gate_up_tbl, t->layer_id)); - const uint8_t* gu_scale_l = reinterpret_cast(tbl_at(gu_scale_tbl, t->layer_id)); - const bf16_t* gu_bias_l = reinterpret_cast(tbl_at(gu_bias_tbl, t->layer_id)); - const int N2 = 2 * I; // gate_up output width (gate/up interleaved) - const int Hh = H / 2; // packed-K rows - const int i0 = static_cast(ib) * IBLK; - const int i1 = std::min(I, i0 + IBLK); - const int nunit = i1 - i0; // intermediate units owned by this tile - const int col0 = 2 * i0; // first gate_up column - const int ncol = 2 * nunit; // gate_up columns owned by this tile - const bf16_t* x_row = t->x + (size_t)tok * H; - const uint8_t* blk_e = gu_packed_l + (size_t)e * Hh * N2; - const uint8_t* scl_e = gu_scale_l + (size_t)e * (size_t)(H / 32) * N2; - float gu[2 * IBLK]; - mxgemv(gu, blk_e + col0, scl_e + col0, x_row, Hh, N2, ncol, e2m1_lut, e8m0_lut); - const bf16_t* bias_e = gu_bias_l + (size_t)e * N2 + col0; - bf16_t* g_row = g_scratch.data() + ((size_t)tok * top_k + k) * I; - const float lim = swiglu_limit, alpha = swiglu_alpha; - for (int j = 0; j < nunit; ++j) { - float gate = gu[2 * j] + bf16_to_f32(bias_e[2 * j]); - float up = gu[2 * j + 1] + bf16_to_f32(bias_e[2 * j + 1]); - if (gate > lim) gate = lim; - if (up > lim) up = lim; - else if (up < -lim) up = -lim; - const float glu = gate / (1.0f + std::exp(-gate * alpha)); // gate * sigmoid(alpha*gate) - g_row[i0 + j] = f32_to_bf16(glu * (up + 1.0f)); - } - } - - void do_pass2_mxfp4(const MoeTask* t, int64_t p) { + // Pass 2 item p = one kRowBlock block of one token's output row, summed over its routes. + void do_pass2(const MoeTask& t, int64_t p) { + const int top_k = params.top_k, H = params.H, I = params.I; const int64_t hb = p % n_hblk; const int tok = static_cast(p / n_hblk); - const int Ih = I / 2; // packed-K rows - const int h0 = static_cast(hb) * HBLK; - const int h1 = std::min(H, h0 + HBLK); - const int nh = h1 - h0; - // Resolve this task's layer base once; row indexing below is layer-local (e). - const uint8_t* dn_packed_l = reinterpret_cast(tbl_at(down_tbl, t->layer_id)); - const uint8_t* dn_scale_l = reinterpret_cast(tbl_at(dn_scale_tbl, t->layer_id)); - const bf16_t* dn_bias_l = reinterpret_cast(tbl_at(dn_bias_tbl, t->layer_id)); - float acc[HBLK]; - for (int c = 0; c < nh; ++c) acc[c] = 0.0f; + const int h0 = static_cast(hb) * kRowBlock; + const int h1 = std::min(H, h0 + kRowBlock); + DownRoute routes[kMaxTopK]; + float w_out[kMaxTopK]; + int n = 0; for (int k = 0; k < top_k; ++k) { - const int e = t->ids[static_cast(tok) * top_k + k]; + const int e = io_ids[static_cast(tok) * top_k + k]; if (e < 0 || e >= num_experts) continue; - const float wt = t->w[static_cast(tok) * top_k + k]; - const bf16_t* g_row = g_scratch.data() + ((size_t)tok * top_k + k) * I; - const uint8_t* blk_e = dn_packed_l + (size_t)e * Ih * H; - const uint8_t* scl_e = dn_scale_l + (size_t)e * (size_t)(I / 32) * H; - float part[HBLK]; - mxgemv(part, blk_e + h0, scl_e + h0, g_row, Ih, H, nh, e2m1_lut, e8m0_lut); - const bf16_t* bias_e = dn_bias_l + (size_t)e * H + h0; - for (int c = 0; c < nh; ++c) acc[c] += (part[c] + bf16_to_f32(bias_e[c])) * wt; - } - bf16_t* y_row = t->y + (size_t)tok * H; - for (int c = 0; c < nh; ++c) y_row[h0 + c] = f32_to_bf16(acc[c]); - } - - // ----------------------------- ds_fp4 (DSV4) ------------------------------- - // Row-major e2m1 + e8m0/32 (no global); silu-swiglu with clamp; FP8-roundtripped - // activations (x once in submit -> xq_scratch; the intermediate g in a dedicated - // round-trip phase between the two passes). Router weight applies on the down output. - - void do_pass1_dsfp4(const MoeTask* t, int64_t p) { - const int64_t ib = p % n_iblk; - const int64_t tk = p / n_iblk; - const int k = static_cast(tk % top_k); - const int tok = static_cast(tk / top_k); - const int e = t->ids[static_cast(tok) * top_k + k]; - if (e < 0 || e >= num_experts) return; - // Resolve this task's layer base once; row indexing below is layer-local (e). - const uint8_t* gu_packed_l = reinterpret_cast(tbl_at(gate_up_tbl, t->layer_id)); - const uint8_t* gu_scale_l = reinterpret_cast(tbl_at(gu_scale_tbl, t->layer_id)); - const int N2 = 2 * I, Hh = H / 2, Hs = H / 32; - // fp8-roundtripped input, pre-deinterleaved to fp32 even/odd halves. - const float* xe = xe_scratch.data() + (size_t)tok * (H / 2); - const float* xo = xo_scratch.data() + (size_t)tok * (H / 2); - const uint8_t* gp = gu_packed_l + (size_t)e * N2 * Hh; - const uint8_t* gs = gu_scale_l + (size_t)e * N2 * Hs; - bf16_t* g_row = g_scratch.data() + ((size_t)tok * top_k + k) * I; - const int i0 = static_cast(ib) * IBLK; - const int i1 = std::min(I, i0 + IBLK); - const float lim = swiglu_limit; - for (int i = i0; i < i1; ++i) { - // gate_up is stored bf16 by the reference GEMV before swiglu; round to match. - float gate = bf16_to_f32(f32_to_bf16( - dsdot(gp + (size_t)i * Hh, gs + (size_t)i * Hs, xe, xo, H, e2m1_lut, e8m0_lut))); - float up = bf16_to_f32(f32_to_bf16(dsdot( - gp + (size_t)(I + i) * Hh, gs + (size_t)(I + i) * Hs, xe, xo, H, e2m1_lut, e8m0_lut))); - if (lim > 0.0f) { - if (gate > lim) gate = lim; - if (up > lim) up = lim; - else if (up < -lim) up = -lim; - } - const float glu = gate / (1.0f + std::exp(-gate)); // silu(gate) - g_row[i] = f32_to_bf16(glu * up); - } - } - - // Prepare one intermediate row (token,route) for the down GEMV, reused across every down output row. - // ds_fp4 first FP8 round-trips it (DSV4 act_quant); nvfp4/ds_fp4/fp8_block deinterleave it to fp32 even/odd; q4_0 Q8_0-quantizes it. - void prep_g_row(int64_t r) { - bf16_t* g = g_scratch.data() + (size_t)r * I; - if (use_q4a8) { // q4_0 W4A8: Q8_0-quantize the intermediate row for the down GEMV. - quant_q8_0(g, I, gi8_scratch.data() + (size_t)r * I, - gas_scratch.data() + (size_t)r * (I / 32)); - return; - } - if (fmt == WF_DSFP4) fp8_roundtrip_bf16(g, g, I, act_block); - float* ge = ge_scratch.data() + (size_t)r * (I / 2); - float* go = go_scratch.data() + (size_t)r * (I / 2); - deinterleave_bf16_f32(g, ge, go, I); - if (use_vnni) - quant_i8_pg16(ge, go, I, gi8_scratch.data() + (size_t)r * I, - gas_scratch.data() + (size_t)r * (I / 16)); - } - - void do_pass2_dsfp4(const MoeTask* t, int64_t p) { - const int64_t hb = p % n_hblk; - const int tok = static_cast(p / n_hblk); - const int Ih = I / 2, Is = I / 32; - const int h0 = static_cast(hb) * HBLK; - const int h1 = std::min(H, h0 + HBLK); - // Resolve this task's layer base once; row indexing below is layer-local (e). - const uint8_t* dn_packed_l = reinterpret_cast(tbl_at(down_tbl, t->layer_id)); - const uint8_t* dn_scale_l = reinterpret_cast(tbl_at(dn_scale_tbl, t->layer_id)); - bf16_t* y_row = t->y + (size_t)tok * H; + const size_t r = (size_t)tok * top_k + k; + routes[n] = {e, static_cast(r), g_scratch.data() + r * I}; + w_out[n] = apply_on_input ? 1.0f : io_w[r]; + ++n; + } + float part[kMaxTopK * kRowBlock]; + if (n > 0) fmt->down(static_cast(t.layer), h0, h1, routes, n, part); + bf16_t* y_row = io_y + (size_t)tok * H; for (int h = h0; h < h1; ++h) { float acc = 0.0f; - for (int k = 0; k < top_k; ++k) { - const int e = t->ids[static_cast(tok) * top_k + k]; - if (e < 0 || e >= num_experts) continue; - const float wt = t->w[static_cast(tok) * top_k + k]; - const float* ge = ge_scratch.data() + ((size_t)tok * top_k + k) * (I / 2); - const float* go = go_scratch.data() + ((size_t)tok * top_k + k) * (I / 2); - const uint8_t* dp = dn_packed_l + (size_t)e * (size_t)H * Ih + (size_t)h * Ih; - const uint8_t* ds = dn_scale_l + (size_t)e * (size_t)H * Is + (size_t)h * Is; - // The reference rounds each route's weighted down output to bf16 before the - // fp32 sum over routes (down [T, top_k, H] bf16 -> .sum(dim=1)). - acc += bf16_to_f32(f32_to_bf16(dsdot(dp, ds, ge, go, I, e2m1_lut, e8m0_lut) * wt)); + for (int j = 0; j < n; ++j) { + const float v = part[j * kRowBlock + h - h0] * w_out[j]; + acc += rounds_routes ? bf16_to_f32(f32_to_bf16(v)) : v; } y_row[h] = f32_to_bf16(acc); } } - void run_task_body(const MoeTask* t) { + void run_task_body(const MoeTask& t) { int local_sense = 0; for (;;) { int64_t p = p1_next.fetch_add(1, std::memory_order_relaxed); @@ -2009,13 +292,13 @@ struct CpuMoeExecutor { do_pass1(t, p); } barrier(local_sense); - // Prepare the intermediate rows (per token,route) for the down GEMV, see prep_g_row. + // Prepare the intermediate rows (per token,route) for the down GEMV, see prep_inter. // Needs all of pass1 done (a full row spans every iblk). - if (needs_di || use_q4a8) { + if (preps_inter) { for (;;) { int64_t r = prt_next.fetch_add(1, std::memory_order_relaxed); if (r >= prt_total) break; - prep_g_row(r); + fmt->prep_inter(r, g_scratch.data() + (size_t)r * params.I); } barrier(local_sense); } @@ -2051,8 +334,7 @@ struct CpuMoeExecutor { } // submit() writes cur_task before the release store to `submitted`. my_gen = gen; - MoeTask* t = cur_task; - run_task_body(t); + run_task_body(cur_task); last_active = wclk::now(); if (done_count.fetch_add(1) + 1 == num_threads) { completed.store(my_gen, std::memory_order_release); @@ -2064,69 +346,20 @@ struct CpuMoeExecutor { } } - void submit(MoeTask* t) { - n_iblk = (I + IBLK - 1) / IBLK; - n_hblk = (H + HBLK - 1) / HBLK; - // Grow the per-token intermediate scratch if a larger batch shows up than the - // construction-time hint (CUDA-graph capture warms the largest bs first, so - // this happens at most once, before any capture, while the pool is idle). - const size_t need = static_cast(t->num_tokens) * top_k * I; - if (need > g_scratch.size()) g_scratch.resize(need); - p1_total = static_cast(t->num_tokens) * top_k * n_iblk; - p2_total = static_cast(t->num_tokens) * n_hblk; - prt_total = (needs_di || use_q4a8) ? static_cast(t->num_tokens) * top_k : 0; + void submit(MoeTask t) { + const int num_tokens = static_cast(t.num_tokens); + p1_total = static_cast(num_tokens) * params.top_k * n_iblk; + p2_total = static_cast(num_tokens) * n_hblk; + prt_total = preps_inter ? static_cast(num_tokens) * params.top_k : 0; p1_next.store(0, std::memory_order_relaxed); p2_next.store(0, std::memory_order_relaxed); prt_next.store(0, std::memory_order_relaxed); done_count.store(0, std::memory_order_relaxed); bar_count.store(0, std::memory_order_relaxed); bar_sense.store(0, std::memory_order_relaxed); - // Deinterleave the per-token input once, up front (ds_fp4 FP8-round-trips it first). + // Prepare the per-token input once, up front. // Single-threaded: tiny for decode, and done before the workers are woken below. - if (needs_di) { - const size_t xn = static_cast(t->num_tokens) * H; - if (xn / 2 > xe_scratch.size()) { - xe_scratch.resize(xn / 2); - xo_scratch.resize(xn / 2); - ge_scratch.resize(static_cast(t->num_tokens) * top_k * (I / 2)); - go_scratch.resize(static_cast(t->num_tokens) * top_k * (I / 2)); - if (use_vnni) { - xi8_scratch.resize(xn); - xas_scratch.resize(xn / 16); - gi8_scratch.resize(static_cast(t->num_tokens) * top_k * I); - gas_scratch.resize(static_cast(t->num_tokens) * top_k * (I / 16)); - } - } - const bool ds = (fmt == WF_DSFP4) && !input_prequant; - if (ds && xn > xq_scratch.size()) xq_scratch.resize(xn); - for (int tok = 0; tok < t->num_tokens; ++tok) { - const bf16_t* src = t->x + (size_t)tok * H; - if (ds) { // DSV4 FP8-round-trips the input before the gate_up GEMV - bf16_t* xq = xq_scratch.data() + (size_t)tok * H; - fp8_roundtrip_bf16(src, xq, H, act_block); - src = xq; - } - float* xe = xe_scratch.data() + (size_t)tok * (H / 2); - float* xo = xo_scratch.data() + (size_t)tok * (H / 2); - deinterleave_bf16_f32(src, xe, xo, H); - if (use_vnni) - quant_i8_pg16(xe, xo, H, xi8_scratch.data() + (size_t)tok * H, - xas_scratch.data() + (size_t)tok * (H / 16)); - } - } - // q4_0 W4A8: Q8_0-quantize the per-token input once (single-threaded, tiny for decode). - if (use_q4a8) { - const size_t xn = static_cast(t->num_tokens) * H; - if (xn > xi8_scratch.size()) { - xi8_scratch.resize(xn); - xas_scratch.resize(xn / 32); - gi8_scratch.resize(static_cast(t->num_tokens) * top_k * I); - gas_scratch.resize(static_cast(t->num_tokens) * top_k * (I / 32)); - } - for (int tok = 0; tok < t->num_tokens; ++tok) - quant_q8_0(t->x + (size_t)tok * H, H, xi8_scratch.data() + (size_t)tok * H, - xas_scratch.data() + (size_t)tok * (H / 32)); - } + for (int tok = 0; tok < num_tokens; ++tok) fmt->prep_input(tok, io_x + (size_t)tok * params.H); { std::lock_guard lk(task_mtx); cur_task = t; @@ -2153,62 +386,56 @@ struct CpuMoeExecutor { sync_cv.wait(lk, [&] { return completed.load(std::memory_order_acquire) >= target; }); } - void submit_with_cuda_stream(uintptr_t stream, uintptr_t task) { - cudaLaunchHostFunc(reinterpret_cast(stream), &CpuMoeExecutor::submit_cb, - reinterpret_cast(task)); + void submit_with_cuda_stream(uintptr_t stream, int layer_id, int num_tokens) { + const MoeTask t = task(layer_id, num_tokens); + HostTask& h = host_tasks.try_emplace(to_word(t), HostTask{this, t}).first->second; + cudaLaunchHostFunc(reinterpret_cast(stream), &CpuMoeExecutor::submit_cb, &h); } - void sync_with_cuda_stream(uintptr_t stream, uintptr_t task) { - cudaLaunchHostFunc(reinterpret_cast(stream), &CpuMoeExecutor::sync_cb, - reinterpret_cast(task)); + void sync_with_cuda_stream(uintptr_t stream) { + cudaLaunchHostFunc(reinterpret_cast(stream), &CpuMoeExecutor::sync_cb, this); } - // Register a (layer, batch-size) slot's task so the coordinator can dispatch it on a - // flag bump. - void register_flag_task(int slot, uintptr_t task) { - std::lock_guard lk(flag_task_mtx); - if (static_cast(flag_task.size()) <= slot) flag_task.resize(slot + 1, nullptr); - flag_task[slot] = reinterpret_cast(task); - } - - int64_t flag_served_count(int slot) const { - return (slot >= 0 && slot < static_cast(flag_served.size())) ? flag_served[slot] : 0; - } - - // Start the busy-poll coordinator over the mapped-pinned flag arrays. ``pin_core`` >= 0 + // Start the busy-poll coordinator over the doorbell. ``pin_core`` >= 0 // pins the coordinator to that logical CPU (the worker auto-sizing reserves it), so its // polling never migrates onto / contends with a GEMV worker's core. - void start_flag_coordinator(uintptr_t ready_ptr, uintptr_t done_ptr, int num_slots, - int pin_core) { - ready_flags = reinterpret_cast(ready_ptr); - done_flags = reinterpret_cast(done_ptr); - coord_num_slots = num_slots; - { - std::lock_guard lk(flag_task_mtx); - if (static_cast(flag_task.size()) < num_slots) flag_task.resize(num_slots, nullptr); - } - flag_served.assign(num_slots, 0); + void start_flag_coordinator(int pin_core) { + gpu_sync = std::make_unique(); coord_stop.store(false); coord_pinned = CPU_MOE_HAS_AFFINITY && pin_core >= 0; coord_thread = std::thread([this, pin_core] { -#if CPU_MOE_HAS_AFFINITY - if (pin_core >= 0) { - cpu_set_t set; - CPU_ZERO(&set); - CPU_SET(pin_core, &set); - pthread_setaffinity_np(pthread_self(), sizeof(set), &set); - } -#endif + if (pin_core >= 0) pin_this_thread(pin_core); coordinator_loop(); }); } + void flag_signal(uintptr_t stream, int layer_id, int num_tokens) { + gpu_sync->gpu_submit(stream, to_word(task(layer_id, num_tokens))); + } + void flag_wait(uintptr_t stream) { gpu_sync->gpu_wait(stream); } + void enable_kernel_io() { + kernel_io = std::make_unique(gpu_sync->ready_word(), gpu_sync->done_word(), io_x, io_ids, io_w, io_y, params.H, + params.top_k); + } + void flag_publish(uintptr_t stream, uintptr_t x, uintptr_t ids, uintptr_t w, int layer_id, int num_tokens) { + kernel_io->publish(stream, reinterpret_cast(x), reinterpret_cast(ids), + reinterpret_cast(w), num_tokens, to_word(task(layer_id, num_tokens))); + } + void flag_combine(uintptr_t stream, uintptr_t out, int num_tokens, bool accumulate) { + kernel_io->combine(stream, reinterpret_cast(out), num_tokens, accumulate); + } + bool flag_pending() const { return gpu_sync && gpu_sync->pending(); } + // Complete a doorbell the coordinator never answered so the stream can make progress. + void poison_flag() { + if (gpu_sync) gpu_sync->cpu_done(); + } + void coordinator_loop() { // Idle backoff (edge-device friendliness): spin hot only while decode traffic is - // flowing, with a TIME-based hot window (an iteration count varies 2-6x with slot - // count and pause cost, so some configs dozed every token). 50 ms since the last - // flag comfortably covers intra- and inter-token gaps -- including host-func-only - // stretches (prefill bursts, --moe-cpu-layers subsets) -- so steady decode never + // flowing, with a TIME-based hot window (an iteration count varies 2-6x with pause + // cost, so some configs dozed every token). 50 ms since the last + // flag comfortably covers intra- and inter-token gaps -- including prefill bursts and + // --moe-cpu-layers gaps -- so steady decode never // sleeps and keeps the ~us-level wakeup. Past it, i.e. the engine is actually idle, // escalate to timed sleeps capped at 2 ms: a dozing coordinator costs <0.1% of a // core instead of 100%, and the only price is a <=2 ms discovery delay on the FIRST @@ -2222,30 +449,19 @@ struct CpuMoeExecutor { int64_t sleep_us = 100; bool dozing = false; while (!coord_stop.load(std::memory_order_relaxed)) { - bool any = false; - for (int L = 0; L < coord_num_slots; ++L) { - // Binary handshake (memop-compatible: the GPU-side WAIT compares against an - // immediate baked at graph capture, so the protocol resets per step instead of - // counting). Acquire: everything the GPU made visible before setting ready -- - // the D2H input copies -- is visible to the worker pool after this read. - if (flag_load_acquire(&ready_flags[L]) != 0) { - flag_store_release(&ready_flags[L], 0); // consume this step's doorbell - MoeTask* t; - { - std::lock_guard lk(flag_task_mtx); - t = (L < static_cast(flag_task.size())) ? flag_task[L] : nullptr; - } - if (t != nullptr) { - submit(t); - sync(coord_pinned && spin_wait.load(std::memory_order_relaxed)); - } - // Release: the workers' y stores are visible before the GPU sees done. - flag_store_release(&done_flags[L], 1); - if (L < static_cast(flag_served.size())) ++flag_served[L]; - any = true; + const uint64_t word = gpu_sync->cpu_poll(); + if (word != 0) { + const MoeTask t = from_word(word); + if (has_routes(t)) { + submit(t); + sync(coord_pinned && spin_wait.load(std::memory_order_relaxed)); + } else { + // What the pool would write for a task without routes, minus its fork and join. + std::memset(io_y, 0, sizeof(bf16_t) * t.num_tokens * params.H); } - } - if (any) { + // Release: the workers' y stores are visible before the GPU sees done. + gpu_sync->cpu_done(); + flag_served.fetch_add(1, std::memory_order_relaxed); last_active = coord_clock::now(); empty_polls = 0; sleep_us = 100; @@ -2254,15 +470,13 @@ struct CpuMoeExecutor { } // Hot phase: pause-spin, consulting the clock only every 1024 empty polls (the // amortization must gate the CLOCK, not the sleep -- gating the sleep decision - // let 1023/1024 idle iterations hot-scan all slots, ~20% of a core at 992 slots). + // let 1023/1024 idle iterations hot-poll). // Doze phase: sleep on EVERY iteration until traffic returns; one wake-up scan // (~0.5 us) per 2 ms sleep is ~0.03% duty. if (!dozing) { if ((++empty_polls & 1023u) != 0 || coord_clock::now() - last_active < kHotWindow) { -#if CPU_MOE_X86 - _mm_pause(); -#endif + cpu_relax(); continue; } dozing = true; @@ -2270,92 +484,79 @@ struct CpuMoeExecutor { std::this_thread::sleep_for(std::chrono::microseconds(sleep_us)); sleep_us = std::min(sleep_us * 2, kSleepCapUs); } - // Teardown: release any in-flight (or future) spin-wait immediately so a replay - // caught mid-shutdown exits its sync kernel now instead of owning the watchdog - // stall. Runs before the destructor's join() returns, while the flag arrays are - // still alive on the Python side. - for (int L = 0; L < coord_num_slots; ++L) { - flag_store_release(&done_flags[L], INT64_MAX); - } + // Teardown: complete any in-flight wait so a replay caught mid-shutdown does not stall + // until the watchdog fires. Runs before the destructor's join() returns, while the doorbell is still alive. + gpu_sync->cpu_done(); } // Eager (non-graph) path: run one task to completion on the pool. - void run_task(uintptr_t task) { - MoeTask* t = reinterpret_cast(task); - submit(t); + void run(int layer_id, int num_tokens) { + submit(task(layer_id, num_tokens)); sync(); } static void CUDART_CB submit_cb(void* ud) { - MoeTask* t = reinterpret_cast(ud); - t->exec->submit(t); - } - static void CUDART_CB sync_cb(void* ud) { - MoeTask* t = reinterpret_cast(ud); - t->exec->sync(); + HostTask* h = reinterpret_cast(ud); + h->exec->submit(h->task); } + static void CUDART_CB sync_cb(void* ud) { reinterpret_cast(ud)->sync(); } }; } // namespace +} // namespace freetoken::cpu_moe PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { namespace py = pybind11; + using freetoken::cpu_moe::CpuMoeExecutor; + using freetoken::cpu_moe::GpuSync; py::class_(m, "CpuMoeExecutor") - .def(py::init, int, int>(), + .def(py::init, std::map, + double, double, std::vector, uintptr_t, uintptr_t, uintptr_t, uintptr_t>(), py::arg("num_threads"), py::arg("num_layers"), py::arg("num_experts"), py::arg("top_k"), py::arg("hidden_size"), py::arg("inter_size"), py::arg("max_tokens"), py::arg("activation_id"), py::arg("apply_router_weight_on_input"), py::arg("weight_format"), - py::arg("gate_up_ptr"), py::arg("down_ptr"), py::arg("gate_up_scale_ptr"), - py::arg("gate_up_global_ptr"), py::arg("down_scale_ptr"), - py::arg("down_global_ptr"), py::arg("gate_up_bias_ptr"), - py::arg("down_bias_ptr"), py::arg("swiglu_alpha"), py::arg("swiglu_limit"), - py::arg("core_ids"), py::arg("fp8_block_gu_scale_stride") = 0, - py::arg("fp8_block_dn_scale_stride") = 0) - .def("create_task", &CpuMoeExecutor::create_task, py::arg("layer_id"), - py::arg("num_tokens"), py::arg("x_ptr"), py::arg("ids_ptr"), py::arg("w_ptr"), + py::arg("banks"), py::arg("options"), py::arg("swiglu_alpha"), py::arg("swiglu_limit"), + py::arg("core_ids"), py::arg("x_ptr"), py::arg("ids_ptr"), py::arg("w_ptr"), py::arg("y_ptr")) .def("submit_with_cuda_stream", &CpuMoeExecutor::submit_with_cuda_stream, - py::arg("stream"), py::arg("task"), py::call_guard()) + py::arg("stream"), py::arg("layer_id"), py::arg("num_tokens"), + py::call_guard()) .def("sync_with_cuda_stream", &CpuMoeExecutor::sync_with_cuda_stream, - py::arg("stream"), py::arg("task"), py::call_guard()) - .def("run_task", &CpuMoeExecutor::run_task, py::arg("task"), + py::arg("stream"), py::call_guard()) + .def("run", &CpuMoeExecutor::run, py::arg("layer_id"), py::arg("num_tokens"), py::call_guard()) - .def("register_flag_task", &CpuMoeExecutor::register_flag_task, - py::arg("slot"), py::arg("task")) - .def("flag_served_count", &CpuMoeExecutor::flag_served_count, py::arg("slot")) - .def("start_flag_coordinator", &CpuMoeExecutor::start_flag_coordinator, - py::arg("ready_ptr"), py::arg("done_ptr"), py::arg("num_slots"), - py::arg("pin_core")) - .def("set_input_prequant", - [](CpuMoeExecutor& e, bool v) { e.input_prequant = v; }, - py::arg("value")) - .def("set_act_block", - [](CpuMoeExecutor& e, int v) { e.act_block = v; }, - py::arg("value")) + .def("flag_served", [](const CpuMoeExecutor& e) { return e.flag_served.load(); }) + .def("start_flag_coordinator", &CpuMoeExecutor::start_flag_coordinator, py::arg("pin_core")) + .def("flag_signal", &CpuMoeExecutor::flag_signal, py::arg("stream"), py::arg("layer_id"), + py::arg("num_tokens")) + .def("flag_wait", &CpuMoeExecutor::flag_wait, py::arg("stream")) + .def("flag_pending", &CpuMoeExecutor::flag_pending) + .def("enable_kernel_io", &CpuMoeExecutor::enable_kernel_io) + .def("flag_publish", &CpuMoeExecutor::flag_publish, py::arg("stream"), py::arg("x"), py::arg("ids"), py::arg("w"), + py::arg("layer_id"), py::arg("num_tokens")) + .def("flag_combine", &CpuMoeExecutor::flag_combine, py::arg("stream"), py::arg("out"), py::arg("num_tokens"), + py::arg("accumulate")) + .def("poison_flag", &CpuMoeExecutor::poison_flag) .def("isa_name", &CpuMoeExecutor::isa_name) - .def("fp8_block_isa_name", &CpuMoeExecutor::fp8_block_isa_name) .def("set_spin_wait", [](CpuMoeExecutor& e, bool v) { e.spin_wait.store(v, std::memory_order_relaxed); }, py::arg("enabled")) .def("get_spin_wait", [](const CpuMoeExecutor& e) { return e.spin_wait.load(); }); - m.def("memops_probe", &cumemops_probe, py::arg("stream"), py::arg("scratch_addr")); - m.def("memop_submit", &cumemop_submit, py::arg("stream"), py::arg("done_addr"), - py::arg("ready_addr"), py::arg("slot")); - m.def("memop_sync", &cumemop_sync, py::arg("stream"), py::arg("done_addr"), - py::arg("slot")); + m.def("memops_probe", &GpuSync::probe, py::arg("stream")); + m.def("kernel_io_probe", &freetoken::cpu_moe::KernelIo::probe, py::arg("stream")); + m.def("host_func_capture_safe", &GpuSync::host_func_capture_safe); // ABI capability marker: the highest ActKind this build implements in the // GENERIC epilogue. CpuMoeExecutor.__init__ probes it before requesting an act // id the epilogue must handle -- a prebuilt .so from before ACT_SWIGLUOAI // accepts id 3 without error and silently computes the wrong activation // (act_apply falls through to gelu_tanh); the probe turns a stale extension // into a loud rebuild instruction instead of wrong model outputs. - m.def("max_generic_act_id", []() { return static_cast(ACT_SWIGLU_CLAMP); }); - // Weight-format ABI probe. A stale prebuilt .so accepts a newer WFmt id without - // complaint and then falls through to the wrong dequant branch -- silently wrong - // numbers, not a crash. The Python executor refuses to build against an extension - // that predates the format it was asked for. - m.def("max_weight_format_id", []() { return static_cast(WF_FP8_BLOCK); }); + m.def("max_generic_act_id", []() { return static_cast(freetoken::cpu_moe::ACT_SWIGLU_CLAMP); }); + m.def("weight_formats", []() { + py::list names; + for (const auto& f : freetoken::cpu_moe::kFormats) names.append(f.first); + return names; + }); } diff --git a/python/freetoken/kernel/csrc/cpu_moe/format.h b/python/freetoken/kernel/csrc/cpu_moe/format.h new file mode 100644 index 0000000000..286ae110f7 --- /dev/null +++ b/python/freetoken/kernel/csrc/cpu_moe/format.h @@ -0,0 +1,131 @@ +// The interface between the CPU MoE executor and one expert weight format. +#pragma once + +#include +#include +#include +#include + +#include "common.h" + +namespace freetoken::cpu_moe { + +// Output-row tiling. Small enough to give every worker independent work even at +// batch size 1; large enough to amortize the atomic work-grab. +// +// Bandwidth notes (Sapphire Rapids 8480+, 13 cores): the two passes already read +// every expert weight byte exactly once per token (each output row block is owned +// by one worker), and x stays hot in L1 across a (token,expert)'s rows -- so the +// kernel is single-read bandwidth-optimal at bs=1 (~205 GB/s vs ~55 GB/s PCIe). +// One worker per *physical* core, pinned, is the sweet spot; SMT oversubscription +// thrashes the spin-barrier. Deferred (not worth it here / for this workload): +// - AMX-bf16: a GEMM tile engine; decode is M=1 GEMV so tiles sit idle. It would +// only pay off in a grouped/batched (dedup) path. +// - expert dedup for bs>1: read each distinct expert once and GEMM its tokens. +// Helps locality+bytes when bs is large; decode batches here are tiny (<=4). +// - NUMA: a single node is assumed. Multi-socket machines would split each +// expert's K dimension per node (banks are already per-row contiguous). +constexpr int kRowBlock = 32; + +struct MoeParams { + int top_k = 0, max_tokens = 0; + int H = 0, I = 0; + int act = ACT_SILU; + float swiglu_alpha = 1.0f; + float swiglu_limit = INFINITY; // +inf == no clamp +}; + +// What cpu_executor.py passes for one format, by name: banks map canonical roles (gate_up, gate_up_scale, ...) to per-layer tables, options hold small integers. +// A format reads what it needs once, in its constructor. +struct FormatArgs { + std::string format; + std::map banks; + std::map options; + + const uint64_t* bank(const std::string& role) const { + const auto it = banks.find(role); + if (it == banks.end() || it->second == 0) + throw std::invalid_argument("CPU MoE " + format + " experts need the " + role + " bank"); + return reinterpret_cast(it->second); + } + int64_t option(const std::string& name) const { + const auto it = options.find(name); + if (it == options.end()) throw std::invalid_argument("CPU MoE " + format + " experts need the " + name + " option"); + return it->second; + } + int64_t option(const std::string& name, int64_t fallback) const { + const auto it = options.find(name); + return it == options.end() ? fallback : it->second; + } +}; + +// A bank table is the address of a CPU int64 array of length +// num_layers (one base address per layer, built by cpu_executor.py's +// _make_table), not a single flat bank. tbl_at resolves +// tbl[layer_id] once per task/pass. +inline const void* tbl_at(const uint64_t* tbl, int layer_id) { + return reinterpret_cast(tbl[layer_id]); +} + +// One route of a pass 2 item: expert e on intermediate row r = tok * top_k + k, whose bf16 values are g. +struct DownRoute { + int e; + int64_t r; + const bf16_t* g; +}; + +// g[i] for i in [i0, i1) from the gate row i and the up row I + i; dot(row) is the row's GEMV. +// ACT_SWIGLUOAI and ACT_SWIGLU_CLAMP clamp both inputs; the other kinds go through act_apply. +template +inline void gated_rows(const MoeParams& p, int i0, int i1, float w_in, bf16_t* g, RowDot dot) { + const bool clamped = p.act == ACT_SWIGLUOAI || p.act == ACT_SWIGLU_CLAMP; + const float up_bias = p.act == ACT_SWIGLUOAI ? 1.0f : 0.0f; + const float lim = p.swiglu_limit, alpha = p.swiglu_alpha; + for (int i = i0; i < i1; ++i) { + // gate = row i, up = row I+i + float gate = dot(i) * w_in; + float up = dot(p.I + i) * w_in; + if (clamped) { + // clamp(gate, max=lim) * sigmoid(alpha * gate) * (clamp(up, +-lim) + up_bias) + // -- swigluoai carries the +1 up bias (gpt-oss/MiniMax); swiglu_clamp + // (GLM-5.3) does not. lim == +inf: no clamp. + if (gate > lim) gate = lim; + if (up > lim) up = lim; + else if (up < -lim) up = -lim; + const float glu = gate / (1.0f + std::exp(-gate * alpha)); + g[i] = f32_to_bf16(glu * (up + up_bias)); + } else { + g[i] = f32_to_bf16(act_apply(p.act, gate) * up); + } + } +} + +// One expert weight format: its kernels, its activation layout and the scratch for it. +// The executor owns the threads and the work split. +// Scratch is sized once from max_tokens; gate_up, down and prep_inter run on the worker threads, prep_input before they wake. +class MoeFormat { + public: + explicit MoeFormat(const MoeParams& p) : p(p) {} + virtual ~MoeFormat() = default; + + virtual std::string isa_name() const = 0; + // Once per input row x[tok] before pass 1. + virtual void prep_input(int /*tok*/, const bf16_t* /*x*/) {} + // Whether pass 2 needs prep_inter on every intermediate row. + virtual bool preps_inter() const { return false; } + virtual void prep_inter(int64_t /*r*/, bf16_t* /*g*/) {} + // Pass 1: g[i] for i in [i0, i1) (at most kRowBlock rows), expert e of `layer` on token tok. + // w_in scales the gate and up dots: the router weight when it applies on the input, else 1. + virtual void gate_up(int layer, int e, int tok, int i0, int i1, const bf16_t* x, float w_in, + bf16_t* g) const = 0; + // Pass 2: out[j * kRowBlock + h - h0] for h in [h0, h1) (at most kRowBlock rows) and each of the n routes, before the router weight. + // Visiting the routes' rows h together keeps n weight streams in flight, which the down GEMV needs to reach DRAM bandwidth. + virtual void down(int layer, int h0, int h1, const DownRoute* routes, int n, float* out) const = 0; + // Round each router-weighted down output to bf16 before the sum over routes. + virtual bool rounds_routes() const { return false; } + + protected: + const MoeParams& p; +}; + +} // namespace freetoken::cpu_moe diff --git a/python/freetoken/kernel/csrc/cpu_moe/gpu_sync.h b/python/freetoken/kernel/csrc/cpu_moe/gpu_sync.h new file mode 100644 index 0000000000..829a8f34c3 --- /dev/null +++ b/python/freetoken/kernel/csrc/cpu_moe/gpu_sync.h @@ -0,0 +1,78 @@ +// GPU <-> CPU handshake for the CPU MoE coordinator, one implementation per GPU runtime (gpu_sync_cuda.cpp, gpu_sync_hip.cpp). +// +// A backend keeps two contracts, in eager mode and in graph replay: +// 1. submit, GPU -> CPU: after gpu_submit(stream, task), cpu_poll() returns task once, and every earlier write of the stream is visible to the CPU. +// 2. wait, CPU -> GPU: gpu_wait(stream) holds later work of the stream until cpu_done(), and that work sees every CPU write made before cpu_done(). +// task is a nonzero word, opaque here; the executor encodes the layer and token count in it, as a constant baked into each captured graph. +// One task is in flight at a time: decode runs on one stream and every layer waits before the next one signals. +// KernelIo (gpu_sync_io.h) may replace the copies around submit and wait: its kernel copies the inputs into the IO rows and does gpu_submit's writes itself. +// It skips a step without a CPU route: done = 1 and no ready, so cpu_poll() never sees the step and gpu_wait() passes at once. +#pragma once + +#include +#include + +#if defined(_MSC_VER) +#include +#endif + +namespace freetoken::cpu_moe { + +// Portable ordering for the flag handshake: "ready observed => the DMA'd inputs that +// preceded the bump are visible" and "y stores are visible before done". Plain +// volatile loads lean on x86 TSO; acquire/release makes it hold on aarch64 too +// (GH200/Jetson) at zero x86 cost. (MSVC branch is x86-only today: compiler barrier +// + TSO.) +inline int64_t flag_load_acquire(const volatile int64_t* p) { +#if defined(_MSC_VER) + const int64_t v = *p; + _ReadWriteBarrier(); + return v; +#else + return __atomic_load_n(const_cast(p), __ATOMIC_ACQUIRE); +#endif +} + +inline void flag_store_release(volatile int64_t* p, int64_t v) { +#if defined(_MSC_VER) + _ReadWriteBarrier(); + *p = v; +#else + __atomic_store_n(const_cast(p), v, __ATOMIC_RELEASE); +#endif +} + +class GpuSync { + public: + // Allocates the handshake memory; throws if this runtime cannot. + GpuSync(); + ~GpuSync(); + GpuSync(const GpuSync&) = delete; + GpuSync& operator=(const GpuSync&) = delete; + + // Whether both contracts hold on this device and stream, including graph replay where the runtime needs a real check. + static bool probe(uintptr_t stream); + // Whether the cudaLaunchHostFunc fallback is replayed correctly from a captured graph. + static bool host_func_capture_safe(); + + // 1. submit. gpu_submit() is enqueued on `stream`, which may be capturing; cpu_poll() returns 0 when nothing is submitted. + void gpu_submit(uintptr_t stream, uint64_t task); + uint64_t cpu_poll(); + + // 2. wait. cpu_done() also unblocks the stream on teardown and from the watchdog. + void gpu_wait(uintptr_t stream); + void cpu_done(); + + // Submitted but not polled yet (watchdog only). + bool pending() const; + + // Device addresses of the two words, for KernelIo. + uintptr_t ready_word() const; + uintptr_t done_word() const; + + private: + struct Impl; + std::unique_ptr impl_; +}; + +} // namespace freetoken::cpu_moe diff --git a/python/freetoken/kernel/csrc/cpu_moe/gpu_sync_cuda.cpp b/python/freetoken/kernel/csrc/cpu_moe/gpu_sync_cuda.cpp new file mode 100644 index 0000000000..a2a089bf67 --- /dev/null +++ b/python/freetoken/kernel/csrc/cpu_moe/gpu_sync_cuda.cpp @@ -0,0 +1,166 @@ +#include +#include +#include +#include +#include +#include + +#if defined(_WIN32) +#include +#else +#include +#endif + +#include + +#include "gpu_sync.h" + +namespace freetoken::cpu_moe { +namespace { + +// ===================================================================================== +// CUDA stream memory operations (driver API, resolved via dlopen -- no link-time or +// toolchain dependence). The GPU side of the flag handshake: submit = WRITE_VALUE +// (done=0 then ready=task word), sync = WAIT_VALUE(done >= 1). The wait is +// executed by the GPU front-end (no SM-resident kernel), so GPU "utilization" stays +// truthful during CPU compute windows -- a resident spin kernel pinned it at 99%, +// which laptop CPU/GPU dynamic power schedulers answered by clamping the CPU's max +// frequency (GEMV workers -1.5x: the reported edge regression). Availability is +// probed functionally at startup (memops_probe); anything unsupported (some driver models, +// vGPU, old drivers) falls back to the cudaLaunchHostFunc path. +#if defined(_WIN32) +static void* cumemop_dlopen() { return (void*)::LoadLibraryA("nvcuda.dll"); } +static void* cumemop_dlsym(void* h, const char* n) { + return (void*)::GetProcAddress((HMODULE)h, n); +} +#else +static void* cumemop_dlopen() { + void* h = dlopen("libcuda.so.1", RTLD_LAZY | RTLD_LOCAL); + if (h == nullptr) h = dlopen("libcuda.so", RTLD_LAZY | RTLD_LOCAL); + return h; +} +static void* cumemop_dlsym(void* h, const char* n) { return dlsym(h, n); } +#endif + +using cuMemOp64_fn = int (*)(void* stream, unsigned long long addr, unsigned long long value, + unsigned int flags); +static cuMemOp64_fn g_cu_write64 = nullptr; +static cuMemOp64_fn g_cu_wait64 = nullptr; +static constexpr unsigned int kCuWaitValueGeq = 0x0; // CU_STREAM_WAIT_VALUE_GEQ +static constexpr unsigned int kCuWriteDefault = 0x0; // CU_STREAM_WRITE_VALUE_DEFAULT + +static bool cumemop_resolve() { + static bool resolved = [] { + void* h = cumemop_dlopen(); + if (h == nullptr) return false; + // 11.7+ made the v2 entry points the default; older drivers export only the v1 + // names with the same signature. + g_cu_write64 = reinterpret_cast(cumemop_dlsym(h, "cuStreamWriteValue64_v2")); + if (g_cu_write64 == nullptr) + g_cu_write64 = reinterpret_cast(cumemop_dlsym(h, "cuStreamWriteValue64")); + g_cu_wait64 = reinterpret_cast(cumemop_dlsym(h, "cuStreamWaitValue64_v2")); + if (g_cu_wait64 == nullptr) + g_cu_wait64 = reinterpret_cast(cumemop_dlsym(h, "cuStreamWaitValue64")); + return g_cu_write64 != nullptr && g_cu_wait64 != nullptr; + }(); + return resolved; +} + +// GPU side of the flag handshake (see the block comment above): enqueued on the +// caller's (possibly capturing) stream; the WAIT immediate is the constant 1, +// replay-safe under CUDA graphs. +// The startup probe validates EAGER memops; a driver could still reject them at graph +// capture time. Those enqueue errors would otherwise be swallowed here and surface only +// as a later EndCapture failure -- log the first CUresult so triage is one step. +static void cumemop_check(int rc, const char* what) { + static std::atomic warned{false}; + if (rc != 0 && !warned.exchange(true)) { + std::fprintf(stderr, + "[freetoken/cpu_moe] %s failed with CUresult=%d (first occurrence; " + "subsequent errors are not repeated). If this happened during CUDA " + "graph capture, the driver lacks capture support for stream memops -- " + "set FREETOKEN_CPU_MOE_FLAG_SYNC=0.\n", + what, rc); + } +} + +} // namespace + +// ready and done live on separate cache lines of one mapped-pinned block; with UVA the GPU memops and the coordinator use the same addresses. +struct GpuSync::Impl { + void* block = nullptr; + void* device_block = nullptr; + volatile int64_t* ready = nullptr; + volatile int64_t* done = nullptr; +}; + +GpuSync::GpuSync() : impl_(std::make_unique()) { + if (!cumemop_resolve()) throw std::runtime_error("CPU MoE flag sync: CUDA stream memops are unavailable"); + const cudaError_t rc = cudaHostAlloc(&impl_->block, 128, cudaHostAllocPortable | cudaHostAllocMapped); + if (rc != cudaSuccess) + throw std::runtime_error(std::string("CPU MoE flag allocation failed: ") + cudaGetErrorString(rc)); + std::memset(impl_->block, 0, 128); + impl_->ready = static_cast(impl_->block); + impl_->done = impl_->ready + 8; + const cudaError_t dev_rc = cudaHostGetDevicePointer(&impl_->device_block, impl_->block, 0); + if (dev_rc != cudaSuccess) { + cudaFreeHost(impl_->block); + throw std::runtime_error(std::string("CPU MoE flag device address: ") + cudaGetErrorString(dev_rc)); + } +} + +GpuSync::~GpuSync() { cudaFreeHost(impl_->block); } + +// Functional probe on a scratch pinned int64: enqueue WRITE(7) + WAIT(>=7) + sync. +// Returns true only if the whole memop path works on THIS stream/device/driver. +bool GpuSync::probe(uintptr_t stream) { + if (!cumemop_resolve()) return false; + void* scratch = nullptr; + if (cudaHostAlloc(&scratch, sizeof(int64_t), cudaHostAllocPortable | cudaHostAllocMapped) != cudaSuccess) + return false; + *static_cast(scratch) = 0; + auto* s = reinterpret_cast(stream); + const auto addr = reinterpret_cast(scratch); + const bool ok = g_cu_write64(s, addr, 7ULL, kCuWriteDefault) == 0 && g_cu_wait64(s, addr, 7ULL, kCuWaitValueGeq) == 0 && + cudaStreamSynchronize(reinterpret_cast(stream)) == cudaSuccess; + // a failed enqueue can leave the WRITE pending on the scratch word + cudaStreamSynchronize(reinterpret_cast(stream)); + cudaFreeHost(scratch); + return ok; +} + +bool GpuSync::host_func_capture_safe() { return true; } + +void GpuSync::gpu_submit(uintptr_t stream, uint64_t task) { + auto* s = reinterpret_cast(stream); + // Order matters and is preserved by the front end: reset done BEFORE raising ready, + // so the coordinator's completion write for THIS step can never be wiped. + cumemop_check(g_cu_write64(s, reinterpret_cast(impl_->done), 0ULL, kCuWriteDefault), + "cuStreamWriteValue64(done)"); + cumemop_check(g_cu_write64(s, reinterpret_cast(impl_->ready), task, kCuWriteDefault), + "cuStreamWriteValue64(ready)"); +} + +void GpuSync::gpu_wait(uintptr_t stream) { + cumemop_check(g_cu_wait64(reinterpret_cast(stream), reinterpret_cast(impl_->done), 1ULL, + kCuWaitValueGeq), + "cuStreamWaitValue64(done)"); +} + +uint64_t GpuSync::cpu_poll() { + // Acquire: everything the GPU made visible before writing ready -- the D2H input copies -- is visible to the worker pool after this read. + const int64_t task = flag_load_acquire(impl_->ready); + if (task == 0) return 0; + flag_store_release(impl_->ready, 0); // consume this step's doorbell + return static_cast(task); +} + +void GpuSync::cpu_done() { flag_store_release(impl_->done, 1); } + +bool GpuSync::pending() const { return flag_load_acquire(impl_->ready) != 0 && flag_load_acquire(impl_->done) == 0; } + +uintptr_t GpuSync::ready_word() const { return reinterpret_cast(impl_->device_block); } + +uintptr_t GpuSync::done_word() const { return reinterpret_cast(impl_->device_block) + 8 * sizeof(int64_t); } + +} // namespace freetoken::cpu_moe diff --git a/python/freetoken/kernel/csrc/cpu_moe/gpu_sync_hip.cpp b/python/freetoken/kernel/csrc/cpu_moe/gpu_sync_hip.cpp new file mode 100644 index 0000000000..31ff9be539 --- /dev/null +++ b/python/freetoken/kernel/csrc/cpu_moe/gpu_sync_hip.cpp @@ -0,0 +1,275 @@ +// HIP stream memory operations for the CPU MoE flag handshake. +// +// HIP WAIT addresses must come from hipMallocSignalMemory, one 64-bit word per allocation. +// +// ROCm 10 records hipStreamWrite/WaitValue64 into the graph under stream capture, one node per call. +// ROCm 7.14 runs them while the stream is being captured but records nothing, so there we add explicit batch-memop nodes and splice them into the stream's dependency set. +// Which of the two applies is measured once at runtime (records_native_memops), not read from the version number. +// The probe then does a real capture + instantiate + replay through the chosen path. +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +#include "gpu_sync.h" + +namespace freetoken::cpu_moe { + +namespace { + +void hip_check(hipError_t rc, const char* what) { + if (rc != hipSuccess) + throw std::runtime_error(std::string("CPU MoE HIP flag sync: ") + what + " failed: " + hipGetErrorString(rc)); +} + +// Capability failures must not poison eager fallback or hide a caller's earlier error. +struct ProbeErrorScope { + ProbeErrorScope() { hip_check(hipPeekAtLastError(), "pending error before probe"); } + ~ProbeErrorScope() { (void)hipGetLastError(); } +}; + +hipError_t alloc_signal(uint64_t** out) { + void* storage = nullptr; + const hipError_t rc = hipExtMallocWithFlags(&storage, sizeof(uint64_t), hipMallocSignalMemory); + if (rc == hipSuccess) { + *out = static_cast(storage); + __atomic_store_n(*out, 0ULL, __ATOMIC_RELEASE); + } + return rc; +} + +void free_signal(uint64_t* signal) { + if (signal != nullptr) (void)hipFree(signal); +} + +volatile int64_t* as_flag(uint64_t* signal) { return reinterpret_cast(signal); } + +bool capturing(hipStream_t stream) { + hipStreamCaptureStatus status = hipStreamCaptureStatusNone; + hip_check(hipStreamIsCapturing(stream, &status), "capture status query"); + return status == hipStreamCaptureStatusActive; +} + +// Whether this runtime records a native hipStreamWriteValue64 issued under capture as a graph node. +bool records_native_memops() { + static const bool native = [] { + ProbeErrorScope errors; + uint64_t* word = nullptr; + hipStream_t stream = nullptr; + hipGraph_t graph = nullptr; + hipError_t rc = alloc_signal(&word); + if (rc == hipSuccess) rc = hipStreamCreateWithFlags(&stream, hipStreamNonBlocking); + if (rc == hipSuccess) rc = hipStreamBeginCapture(stream, hipStreamCaptureModeRelaxed); + const bool capture_started = rc == hipSuccess; + if (rc == hipSuccess) rc = hipStreamWriteValue64(stream, word, 1, hipStreamWriteValueDefault); + if (capture_started) { + const hipError_t end_rc = hipStreamEndCapture(stream, &graph); + if (rc == hipSuccess) rc = end_rc; + } + size_t nodes = 0; + if (rc == hipSuccess) rc = hipGraphGetNodes(graph, nullptr, &nodes); + if (graph != nullptr) (void)hipGraphDestroy(graph); + if (stream != nullptr) { + (void)hipStreamSynchronize(stream); + (void)hipStreamDestroy(stream); + } + free_signal(word); + return rc == hipSuccess && nodes == 1; + }(); + return native; +} + +// Appends one batch-memop node to the capture of `stream`; the node keeps a pointer to `ops`. +hipError_t add_capture_memop_node(hipStream_t stream, hipStreamBatchMemOpParams* ops, unsigned int count) { + hipStreamCaptureStatus status = hipStreamCaptureStatusNone; + hipGraph_t graph = nullptr; + const hipGraphNode_t* dependencies = nullptr; + size_t dependency_count = 0; + hipError_t rc = hipStreamGetCaptureInfo_v2(stream, &status, nullptr, &graph, &dependencies, &dependency_count); + if (rc != hipSuccess) return rc; + if (status != hipStreamCaptureStatusActive || graph == nullptr) return hipErrorIllegalState; + + hipCtx_t context = nullptr; +#if defined(__GNUC__) +#pragma GCC diagnostic push +#pragma GCC diagnostic ignored "-Wdeprecated-declarations" +#endif + rc = hipCtxGetCurrent(&context); +#if defined(__GNUC__) +#pragma GCC diagnostic pop +#endif + if (rc != hipSuccess) return rc; + + hipBatchMemOpNodeParams params{}; + params.ctx = context; + params.count = count; + params.paramArray = ops; + hipGraphNode_t node = nullptr; + rc = hipGraphAddBatchMemOpNode(&node, graph, dependencies, dependency_count, ¶ms); + if (rc != hipSuccess) return rc; + rc = hipStreamUpdateCaptureDependencies(stream, &node, 1, hipStreamSetCaptureDependencies); + // EndCapture rejects a node that was added but never joined to the stream. + if (rc != hipSuccess && node != nullptr) (void)hipGraphDestroyNode(node); + return rc; +} + +hipStreamBatchMemOpParams write_op(uint64_t* address, uint64_t value) { + hipStreamBatchMemOpParams op{}; + op.writeValue.operation = hipStreamMemOpWriteValue64; + op.writeValue.address = reinterpret_cast(address); + op.writeValue.value64 = value; + return op; +} + +hipStreamBatchMemOpParams wait_op(uint64_t* address) { + hipStreamBatchMemOpParams op{}; + op.waitValue.operation = hipStreamMemOpWaitValue64; + op.waitValue.address = reinterpret_cast(address); + op.waitValue.value64 = 1; + op.waitValue.flags = hipStreamWaitValueGte; + return op; +} + +bool needs_explicit_nodes(hipStream_t stream) { return capturing(stream) && !records_native_memops(); } + +// node_op is null for a native call, else the parameter slot of an explicit node, which the node keeps a pointer to. +// The ops of one batch node run in parallel, so each write gets its own node and consecutive writes stay in order. +hipError_t write_value(hipStream_t stream, uint64_t* address, uint64_t value, hipStreamBatchMemOpParams* node_op) { + if (node_op == nullptr) return hipStreamWriteValue64(stream, address, value, hipStreamWriteValueDefault); + *node_op = write_op(address, value); + return add_capture_memop_node(stream, node_op, 1); +} + +hipError_t wait_done(hipStream_t stream, uint64_t* done, hipStreamBatchMemOpParams* node_op) { + if (node_op == nullptr) return hipStreamWaitValue64(stream, done, 1, hipStreamWaitValueGte, UINT64_MAX); + *node_op = wait_op(done); + return add_capture_memop_node(stream, node_op, 1); +} + +} // namespace + +struct GpuSync::Impl { + uint64_t* ready = nullptr; // signal memory + uint64_t* done = nullptr; + // Explicit nodes keep these addresses, so one parameter pair per distinct task, never moved (std::map nodes are stable). + std::map> submit_ops; // done = 0, ready = task + hipStreamBatchMemOpParams wait_op{}; +}; + +GpuSync::GpuSync() : impl_(std::make_unique()) { + hipError_t rc = alloc_signal(&impl_->ready); + if (rc == hipSuccess) rc = alloc_signal(&impl_->done); + if (rc != hipSuccess) { + free_signal(impl_->ready); + free_signal(impl_->done); + hip_check(rc, "signal allocation"); + } + // Measure now: the measurement captures and allocates, which a graph capture in progress would reject. + (void)records_native_memops(); +} + +GpuSync::~GpuSync() { + free_signal(impl_->ready); + free_signal(impl_->done); +} + +bool GpuSync::probe(uintptr_t /*stream*/) { + ProbeErrorScope errors; + int device = 0; + int can_wait = 0; + if (hipGetDevice(&device) != hipSuccess || + hipDeviceGetAttribute(&can_wait, hipDeviceAttributeCanUseStreamWaitValue, device) != hipSuccess || can_wait == 0) + return false; + const bool explicit_nodes = !records_native_memops(); + + uint64_t* ready = nullptr; + uint64_t* done = nullptr; + hipStream_t stream = nullptr; + hipGraph_t graph = nullptr; + hipGraphExec_t graph_exec = nullptr; + std::array node_ops{}; + hipError_t rc = alloc_signal(&ready); + if (rc == hipSuccess) rc = alloc_signal(&done); + if (rc == hipSuccess) rc = hipStreamCreateWithFlags(&stream, hipStreamNonBlocking); + // A runtime that runs the WAIT during capture instead of recording it must not block here. + if (rc == hipSuccess) __atomic_store_n(done, 1ULL, __ATOMIC_RELEASE); + if (rc == hipSuccess) rc = hipStreamBeginCapture(stream, hipStreamCaptureModeRelaxed); + const bool capture_started = rc == hipSuccess; + if (rc == hipSuccess) rc = write_value(stream, ready, 1, explicit_nodes ? &node_ops[0] : nullptr); + if (rc == hipSuccess) rc = wait_done(stream, done, explicit_nodes ? &node_ops[1] : nullptr); + // End the capture even on failure so its stream and graph are not left active during cleanup. + if (capture_started) { + const hipError_t end_rc = hipStreamEndCapture(stream, &graph); + if (rc == hipSuccess) rc = end_rc; + } + if (rc == hipSuccess) rc = hipStreamSynchronize(stream); + size_t nodes = 0; + if (rc == hipSuccess) rc = hipGraphGetNodes(graph, nullptr, &nodes); + if (rc == hipSuccess && nodes != 2) rc = hipErrorNotSupported; + if (rc == hipSuccess) rc = hipGraphInstantiate(&graph_exec, graph, nullptr, nullptr, 0); + + bool served = false; + if (rc == hipSuccess) { + __atomic_store_n(ready, 0ULL, __ATOMIC_RELEASE); + __atomic_store_n(done, 0ULL, __ATOMIC_RELEASE); + std::thread responder([&] { + const auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(2); + while (__atomic_load_n(ready, __ATOMIC_ACQUIRE) == 0 && std::chrono::steady_clock::now() < deadline) + std::this_thread::yield(); + served = __atomic_load_n(ready, __ATOMIC_ACQUIRE) == 1; + // The probe graph never resets done, so this also releases a launch that arrives after the deadline. + __atomic_store_n(done, 1ULL, __ATOMIC_RELEASE); + }); + rc = hipGraphLaunch(graph_exec, stream); + if (rc == hipSuccess) rc = hipStreamSynchronize(stream); + responder.join(); + } + const bool ok = rc == hipSuccess && served; + + if (graph_exec != nullptr) (void)hipGraphExecDestroy(graph_exec); + if (graph != nullptr) (void)hipGraphDestroy(graph); + if (stream != nullptr) (void)hipStreamDestroy(stream); + free_signal(ready); + free_signal(done); + return ok; +} + +void GpuSync::gpu_submit(uintptr_t stream, uint64_t task) { + auto s = reinterpret_cast(stream); + hipStreamBatchMemOpParams* ops = needs_explicit_nodes(s) ? impl_->submit_ops[task].data() : nullptr; + // Reset done before raising ready, so the coordinator's completion for this step can never be wiped. + hip_check(write_value(s, impl_->done, 0, ops), "write done"); + hip_check(write_value(s, impl_->ready, task, ops ? ops + 1 : nullptr), "write ready"); +} + +void GpuSync::gpu_wait(uintptr_t stream) { + auto s = reinterpret_cast(stream); + hip_check(wait_done(s, impl_->done, needs_explicit_nodes(s) ? &impl_->wait_op : nullptr), "wait done"); +} + +bool GpuSync::host_func_capture_safe() { return false; } + +uint64_t GpuSync::cpu_poll() { + const int64_t task = flag_load_acquire(as_flag(impl_->ready)); + if (task == 0) return 0; + flag_store_release(as_flag(impl_->ready), 0); + return static_cast(task); +} + +void GpuSync::cpu_done() { flag_store_release(as_flag(impl_->done), 1); } + +bool GpuSync::pending() const { + return flag_load_acquire(as_flag(impl_->ready)) != 0 && flag_load_acquire(as_flag(impl_->done)) == 0; +} + +uintptr_t GpuSync::ready_word() const { return reinterpret_cast(impl_->ready); } + +uintptr_t GpuSync::done_word() const { return reinterpret_cast(impl_->done); } + +} // namespace freetoken::cpu_moe diff --git a/python/freetoken/kernel/csrc/cpu_moe/gpu_sync_io.cu b/python/freetoken/kernel/csrc/cpu_moe/gpu_sync_io.cu new file mode 100644 index 0000000000..b41961f0bf --- /dev/null +++ b/python/freetoken/kernel/csrc/cpu_moe/gpu_sync_io.cu @@ -0,0 +1,199 @@ +#include +#include +#include + +#include + +#include "gpu_sync_io.h" + +namespace freetoken::cpu_moe { +namespace { + +__device__ inline float bf16_bits_to_float(uint16_t b) { return __uint_as_float(static_cast(b) << 16); } + +// Round to nearest even, the same rounding as a torch bf16 add. +__device__ inline uint16_t float_to_bf16_bits(float f) { + const uint32_t u = __float_as_uint(f); + if ((u & 0x7fffffffu) > 0x7f800000u) return 0x7fc0; + return static_cast((u + 0x7fffu + ((u >> 16) & 1u)) >> 16); +} + +// One block. Copies x/ids/w into the IO rows, then done = 0 and ready = task. +// A step without a CPU route writes done = 1 and rings nothing: the coordinator stays asleep and the GPU wait passes at once. +__global__ void publish_kernel(const uint16_t* __restrict__ x, const int32_t* __restrict__ ids, const float* __restrict__ w, + int64_t n_x, int64_t n_ids, uint16_t* io_x, int32_t* io_ids, float* io_w, + volatile int64_t* ready, volatile int64_t* done, int64_t task, int32_t* has_cpu) { + __shared__ int any; + if (threadIdx.x == 0) any = 0; + __syncthreads(); + for (int64_t i = threadIdx.x; i < n_ids; i += blockDim.x) + if (ids[i] >= 0) any = 1; + __syncthreads(); + if (!any) { + if (threadIdx.x == 0) { + *has_cpu = 0; + *done = 1; + __threadfence_system(); + } + return; + } + const bool vec = (reinterpret_cast(x) | reinterpret_cast(io_x)) % 16 == 0 && n_x % 8 == 0; + if (vec) { + for (int64_t i = threadIdx.x; i < n_x / 8; i += blockDim.x) + reinterpret_cast(io_x)[i] = reinterpret_cast(x)[i]; + } else { + for (int64_t i = threadIdx.x; i < n_x; i += blockDim.x) io_x[i] = x[i]; + } + for (int64_t i = threadIdx.x; i < n_ids; i += blockDim.x) { + io_ids[i] = ids[i]; + io_w[i] = w[i]; + } + __threadfence_system(); + __syncthreads(); + if (threadIdx.x == 0) { + *has_cpu = 1; + *done = 0; + __threadfence_system(); + *ready = task; + __threadfence_system(); + } +} + +// out = CPU rows, or out += CPU rows with accumulate; a step publish skipped contributes zeros. +// io_y is volatile: the same host rows hold a new step on every call. +__global__ void combine_kernel(uint16_t* out, const volatile uint16_t* io_y, const int32_t* has_cpu, int64_t n, bool accumulate) { + const bool on = *has_cpu != 0; + if (accumulate && !on) return; + for (int64_t i = blockIdx.x * static_cast(blockDim.x) + threadIdx.x; i < n; i += static_cast(gridDim.x) * blockDim.x) { + float v = accumulate ? bf16_bits_to_float(out[i]) : 0.0f; + if (on) v += bf16_bits_to_float(io_y[i]); + out[i] = float_to_bf16_bits(v); + } +} + +cudaError_t launch_publish(cudaStream_t stream, const uint16_t* x, const int32_t* ids, const float* w, int64_t n_x, int64_t n_ids, + uint16_t* io_x, int32_t* io_ids, float* io_w, volatile int64_t* ready, volatile int64_t* done, int64_t task, + int32_t* has_cpu) { + publish_kernel<<<1, 256, 0, stream>>>(x, ids, w, n_x, n_ids, io_x, io_ids, io_w, ready, done, task, has_cpu); + return cudaGetLastError(); +} + +cudaError_t launch_combine(cudaStream_t stream, uint16_t* out, const uint16_t* io_y, const int32_t* has_cpu, int64_t n, bool accumulate) { + if (n == 0) return cudaSuccess; + const auto blocks = static_cast(std::min((n + 255) / 256, 1024)); + combine_kernel<<>>(out, io_y, has_cpu, n, accumulate); + return cudaGetLastError(); +} + +void check(cudaError_t rc, const char* what) { + if (rc != cudaSuccess) throw std::runtime_error(std::string("CPU MoE kernel IO: ") + what + ": " + cudaGetErrorString(rc)); +} + +template +T* device_address(const void* host) { + void* dev = nullptr; + check(cudaHostGetDevicePointer(&dev, const_cast(host), 0), "device address of the IO rows"); + return static_cast(dev); +} + +} // namespace + +KernelIo::KernelIo(uintptr_t ready, uintptr_t done, const void* x, const int32_t* ids, const float* w, const void* y, int hidden, + int top_k) + : ready_(reinterpret_cast(ready)), + done_(reinterpret_cast(done)), + x_(device_address(x)), + ids_(device_address(ids)), + w_(device_address(w)), + y_(device_address(y)), + hidden_(hidden), + top_k_(top_k) { + check(cudaMalloc(reinterpret_cast(&has_cpu_), sizeof(int32_t)), "has_cpu allocation"); + const cudaError_t rc = cudaMemset(has_cpu_, 0, sizeof(int32_t)); + if (rc != cudaSuccess) { + (void)cudaFree(has_cpu_); + check(rc, "has_cpu reset"); + } +} + +KernelIo::~KernelIo() { (void)cudaFree(has_cpu_); } + +void KernelIo::publish(uintptr_t stream, const void* x, const int32_t* ids, const float* w, int num_tokens, uint64_t task) { + check(launch_publish(reinterpret_cast(stream), static_cast(x), ids, w, + static_cast(num_tokens) * hidden_, static_cast(num_tokens) * top_k_, x_, ids_, w_, ready_, done_, + static_cast(task), has_cpu_), + "publish launch"); +} + +void KernelIo::combine(uintptr_t stream, void* out, int num_tokens, bool accumulate) { + check(launch_combine(reinterpret_cast(stream), static_cast(out), y_, has_cpu_, + static_cast(num_tokens) * hidden_, accumulate), + "combine launch"); +} + +// A ringing step, its combine and a skipped step on scratch mapped memory, checking every word they write. +std::string KernelIo::probe(uintptr_t stream_handle) { + struct alignas(16) Area { + int64_t ready, done; + int32_t has_cpu, pad[3]; + uint16_t x[8], io_x[8], io_y[8], out[8]; + int32_t ids[2], io_ids[2]; + float w[2], io_w[2]; + }; + const auto stream = reinterpret_cast(stream_handle); + void* host = nullptr; + void* dev = nullptr; + if (cudaHostAlloc(&host, sizeof(Area), cudaHostAllocPortable | cudaHostAllocMapped) != cudaSuccess) + return "pinned scratch allocation failed"; + std::string reason; + auto ok = [&](cudaError_t rc, const char* what) { + if (rc != cudaSuccess && reason.empty()) reason = std::string(what) + ": " + cudaGetErrorString(rc); + return rc == cudaSuccess; + }; + auto fail = [&](const char* why) { + if (reason.empty()) reason = why; + }; + if (ok(cudaHostGetDevicePointer(&dev, host, 0), "device address of the scratch")) { + auto* h = static_cast(host); + auto* d = static_cast(dev); + std::memset(h, 0, sizeof(Area)); + for (int i = 0; i < 8; ++i) { + h->x[i] = static_cast(0x3f80 + i); + h->io_y[i] = 0x3f80; // 1.0 + h->out[i] = 0x4000; // 2.0 + } + h->ids[0] = 3; + h->ids[1] = -1; + h->w[0] = 0.25f; + h->w[1] = 0.75f; + auto publish = [&] { + return launch_publish(stream, d->x, d->ids, d->w, 8, 2, d->io_x, d->io_ids, d->io_w, &d->ready, &d->done, 5, &d->has_cpu); + }; + + if (ok(publish(), "publish launch") && ok(cudaStreamSynchronize(stream), "publish")) { + if (h->ready != 5 || h->done != 0 || h->has_cpu != 1) fail("publish wrote wrong handshake words"); + if (std::memcmp(h->io_x, h->x, sizeof(h->x)) || std::memcmp(h->io_ids, h->ids, sizeof(h->ids)) || + std::memcmp(h->io_w, h->w, sizeof(h->w))) + fail("publish wrote wrong IO rows"); + } + if (reason.empty() && ok(launch_combine(stream, d->out, d->io_y, &d->has_cpu, 8, true), "combine launch") && + ok(cudaStreamSynchronize(stream), "combine")) { + for (int i = 0; i < 8; ++i) + if (h->out[i] != 0x4040) fail("combine returned a wrong sum"); // 3.0 + } + if (reason.empty()) { + h->ready = 0; + h->ids[0] = -1; + if (ok(publish(), "skip launch") && ok(launch_combine(stream, d->out, d->io_y, &d->has_cpu, 8, false), "skip combine launch") && + ok(cudaStreamSynchronize(stream), "skipped step")) { + bool zeros = true; + for (int i = 0; i < 8; ++i) zeros = zeros && h->out[i] == 0; + if (h->ready != 0 || h->done != 1 || h->has_cpu != 0 || !zeros) fail("a skipped step wrote wrong words"); + } + } + } + (void)cudaFreeHost(host); + return reason; +} + +} // namespace freetoken::cpu_moe diff --git a/python/freetoken/kernel/csrc/cpu_moe/gpu_sync_io.h b/python/freetoken/kernel/csrc/cpu_moe/gpu_sync_io.h new file mode 100644 index 0000000000..6eacb3936b --- /dev/null +++ b/python/freetoken/kernel/csrc/cpu_moe/gpu_sync_io.h @@ -0,0 +1,39 @@ +// Kernel IO for the GPU <-> CPU handshake (gpu_sync.h): kernels on the executor's mapped pinned IO rows in place of the copies around submit and wait. +// One source for both runtimes: nvcc builds gpu_sync_io.cu on CUDA, hipcc on ROCm through freetoken/hip_compat.h. +#pragma once + +#include +#include + +namespace freetoken::cpu_moe { + +class KernelIo { + public: + // ready, done: device addresses of GpuSync's words. x, ids, w, y: host addresses of the pinned IO rows. Throws if setup fails. + KernelIo(uintptr_t ready, uintptr_t done, const void* x, const int32_t* ids, const float* w, const void* y, int hidden, int top_k); + ~KernelIo(); + KernelIo(const KernelIo&) = delete; + KernelIo& operator=(const KernelIo&) = delete; + + // Copies the step's device inputs into the IO rows and does gpu_submit's writes, or skips the step when no route goes to the CPU. + // x [n, hidden] bf16, ids [n, top_k] int32 (< 0 = not a CPU route), w [n, top_k] fp32. + void publish(uintptr_t stream, const void* x, const int32_t* ids, const float* w, int num_tokens, uint64_t task); + // Follows gpu_wait(): writes the CPU rows to out [n, hidden] bf16, or adds them to it; a skipped step gives zeros. + void combine(uintptr_t stream, void* out, int num_tokens, bool accumulate); + + // Whether both kernels run and write the right words on this device: "" or the reason. + static std::string probe(uintptr_t stream); + + private: + volatile int64_t* ready_; + volatile int64_t* done_; + uint16_t* x_; + int32_t* ids_; + float* w_; + const uint16_t* y_; + int hidden_; + int top_k_; + int32_t* has_cpu_ = nullptr; // device word: whether the last published step has CPU routes +}; + +} // namespace freetoken::cpu_moe diff --git a/python/freetoken/kernel/csrc/cpu_moe/kernels/bf16.cpp b/python/freetoken/kernel/csrc/cpu_moe/kernels/bf16.cpp new file mode 100644 index 0000000000..273c394609 --- /dev/null +++ b/python/freetoken/kernel/csrc/cpu_moe/kernels/bf16.cpp @@ -0,0 +1,159 @@ +#include "../format.h" + +namespace freetoken::cpu_moe { +namespace { + +// ------------------------------- dot products ------------------------------- +// dot(weight[bf16], act[bf16], n) -> fp32. The selected impl is a function +// pointer chosen at runtime; the per-row call overhead is negligible vs n. + +using dot_fn = float (*)(const bf16_t*, const bf16_t*, int); + +float dot_scalar(const bf16_t* w, const bf16_t* x, int n) { + float acc = 0.0f; + for (int i = 0; i < n; ++i) acc += bf16_to_f32(w[i]) * bf16_to_f32(x[i]); + return acc; +} + +#if CPU_MOE_X86 +CPU_MOE_TARGET("avx512f") +float dot_avx512f(const bf16_t* w, const bf16_t* x, int n) { + // 4 independent accumulators -> more in-flight loads (memory-level parallelism), + // which is what lifts a bandwidth-bound GEMV toward peak. + __m512 a0 = _mm512_setzero_ps(), a1 = _mm512_setzero_ps(); + __m512 a2 = _mm512_setzero_ps(), a3 = _mm512_setzero_ps(); + int i = 0; + for (; i + 64 <= n; i += 64) { + _mm_prefetch(reinterpret_cast(w + i) + PF_AHEAD, _MM_HINT_T0); + for (int j = 0; j < 64; j += 16) { + __m256i wi = _mm256_loadu_si256(reinterpret_cast(w + i + j)); + __m256i xi = _mm256_loadu_si256(reinterpret_cast(x + i + j)); + __m512 wf = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(wi), 16)); + __m512 xf = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(xi), 16)); + __m512& acc = (j == 0) ? a0 : (j == 16) ? a1 : (j == 32) ? a2 : a3; + acc = _mm512_fmadd_ps(wf, xf, acc); + } + } + for (; i + 16 <= n; i += 16) { + __m256i wi = _mm256_loadu_si256(reinterpret_cast(w + i)); + __m256i xi = _mm256_loadu_si256(reinterpret_cast(x + i)); + __m512 wf = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(wi), 16)); + __m512 xf = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(xi), 16)); + a0 = _mm512_fmadd_ps(wf, xf, a0); + } + float s = _mm512_reduce_add_ps(_mm512_add_ps(_mm512_add_ps(a0, a1), _mm512_add_ps(a2, a3))); + for (; i < n; ++i) s += bf16_to_f32(w[i]) * bf16_to_f32(x[i]); + return s; +} + +#ifdef CPU_MOE_HAS_AVX512BF16 +CPU_MOE_TARGET("avx512bf16,avx512f") +static inline __m512bh load_bh(const bf16_t* p) { + __m512i raw = _mm512_loadu_si512(reinterpret_cast(p)); + __m512bh out; + std::memcpy(&out, &raw, sizeof(out)); + return out; +} + +CPU_MOE_TARGET("avx512bf16,avx512f") +float dot_avx512bf16(const bf16_t* w, const bf16_t* x, int n) { + // 4 accumulators (128 bf16/iter) for memory-level parallelism + a prefetch nudge. + __m512 a0 = _mm512_setzero_ps(), a1 = _mm512_setzero_ps(); + __m512 a2 = _mm512_setzero_ps(), a3 = _mm512_setzero_ps(); + int i = 0; + for (; i + 128 <= n; i += 128) { + _mm_prefetch(reinterpret_cast(w + i) + PF_AHEAD, _MM_HINT_T0); + a0 = _mm512_dpbf16_ps(a0, load_bh(w + i), load_bh(x + i)); + a1 = _mm512_dpbf16_ps(a1, load_bh(w + i + 32), load_bh(x + i + 32)); + a2 = _mm512_dpbf16_ps(a2, load_bh(w + i + 64), load_bh(x + i + 64)); + a3 = _mm512_dpbf16_ps(a3, load_bh(w + i + 96), load_bh(x + i + 96)); + } + for (; i + 32 <= n; i += 32) { + a0 = _mm512_dpbf16_ps(a0, load_bh(w + i), load_bh(x + i)); + } + float s = _mm512_reduce_add_ps(_mm512_add_ps(_mm512_add_ps(a0, a1), _mm512_add_ps(a2, a3))); + for (; i < n; ++i) s += bf16_to_f32(w[i]) * bf16_to_f32(x[i]); + return s; +} +#endif // avx512bf16 available + +// --------- AVX2 + FMA (256-bit fallback for CPUs without AVX-512) ---------- +// Covers Intel 12-14th gen / Arrow Lake (AVX-512 fused off) and AMD Zen<4. The +// bf16->fp32 widen is a zero-extend + <<16; with 4 independent accumulators the +// GEMV is memory-bandwidth bound, same as the AVX-512 path (just half the width). +CPU_MOE_TARGET("avx2,fma") +float dot_avx2(const bf16_t* w, const bf16_t* x, int n) { + __m256 a0 = _mm256_setzero_ps(), a1 = _mm256_setzero_ps(); + __m256 a2 = _mm256_setzero_ps(), a3 = _mm256_setzero_ps(); + int i = 0; + for (; i + 32 <= n; i += 32) { + _mm_prefetch(reinterpret_cast(w + i) + PF_AHEAD, _MM_HINT_T0); + for (int j = 0; j < 32; j += 8) { + __m128i wi = _mm_loadu_si128(reinterpret_cast(w + i + j)); + __m128i xi = _mm_loadu_si128(reinterpret_cast(x + i + j)); + __m256 wf = _mm256_castsi256_ps(_mm256_slli_epi32(_mm256_cvtepu16_epi32(wi), 16)); + __m256 xf = _mm256_castsi256_ps(_mm256_slli_epi32(_mm256_cvtepu16_epi32(xi), 16)); + __m256& acc = (j == 0) ? a0 : (j == 8) ? a1 : (j == 16) ? a2 : a3; + acc = _mm256_fmadd_ps(wf, xf, acc); + } + } + for (; i + 8 <= n; i += 8) { + __m128i wi = _mm_loadu_si128(reinterpret_cast(w + i)); + __m128i xi = _mm_loadu_si128(reinterpret_cast(x + i)); + __m256 wf = _mm256_castsi256_ps(_mm256_slli_epi32(_mm256_cvtepu16_epi32(wi), 16)); + __m256 xf = _mm256_castsi256_ps(_mm256_slli_epi32(_mm256_cvtepu16_epi32(xi), 16)); + a0 = _mm256_fmadd_ps(wf, xf, a0); + } + float s = hsum256(_mm256_add_ps(_mm256_add_ps(a0, a1), _mm256_add_ps(a2, a3))); + for (; i < n; ++i) s += bf16_to_f32(w[i]) * bf16_to_f32(x[i]); + return s; +} +#endif // CPU_MOE_X86 + +dot_fn select_dot() { + const IsaTier t = pick_isa(); +#if CPU_MOE_X86 +#ifdef CPU_MOE_HAS_AVX512BF16 + if (t >= ISA_AVX512BF16) return dot_avx512bf16; +#endif + if (t >= ISA_AVX512) return dot_avx512f; + if (t >= ISA_AVX2) return dot_avx2; +#endif + (void)t; + return dot_scalar; +} + +class Bf16Format final : public MoeFormat { + public: + Bf16Format(const MoeParams& p, const FormatArgs& a) + : MoeFormat(p), dot_(select_dot()), gate_up_(a.bank("gate_up")), down_(a.bank("down")) {} + + std::string isa_name() const override { return tier_name(pick_isa()); } + + void gate_up(int layer, int e, int /*tok*/, int i0, int i1, const bf16_t* x, float w_in, + bf16_t* g) const override { + const bf16_t* w = static_cast(tbl_at(gate_up_, layer)) + (size_t)e * (2 * p.I) * p.H; + gated_rows(p, i0, i1, w_in, g, [&](int row) { return dot_(w + (size_t)row * p.H, x, p.H); }); + } + + void down(int layer, int h0, int h1, const DownRoute* routes, int n, float* out) const override { + const bf16_t* w = static_cast(tbl_at(down_, layer)); + for (int h = h0; h < h1; ++h) + for (int j = 0; j < n; ++j) + out[j * kRowBlock + h - h0] = dot_(w + ((size_t)routes[j].e * p.H + h) * p.I, routes[j].g, p.I); + } + + private: + dot_fn dot_; + // bf16 rows: gate_up [E, 2I, H], down [E, H, I]. + const uint64_t* gate_up_; + const uint64_t* down_; +}; + +} // namespace + +std::unique_ptr make_bf16_format(const MoeParams& p, const FormatArgs& a) { + return std::make_unique(p, a); +} + +} // namespace freetoken::cpu_moe diff --git a/python/freetoken/kernel/csrc/cpu_moe/kernels/ds_fp4.cpp b/python/freetoken/kernel/csrc/cpu_moe/kernels/ds_fp4.cpp new file mode 100644 index 0000000000..a7afe9804d --- /dev/null +++ b/python/freetoken/kernel/csrc/cpu_moe/kernels/ds_fp4.cpp @@ -0,0 +1,269 @@ +#include + +#include "../format.h" + +namespace freetoken::cpu_moe { +namespace { + +// ----------------------- DeepSeek-V4 ds_fp4 (W4A8) --------------------------- +// Row-major e2m1 (2/byte, low nibble first) + e8m0 per-32 block scale, no global +// (w = E2M1[code] * 2^(e8m0-127)); activations are FP8-e4m3 round-tripped (per-128 +// block, ue8m0 scale) before each GEMM. Matches kernel/triton/dsv4 (fused_moe + +// fp8_linear): silu(clamp(gate,max=lim)) * clamp(up,-lim,lim), router weight on the +// down output. + +// Activations are pre-deinterleaved to fp32 (xe[m]=x[2m], xo[m]=x[2m+1]) once per +// token/route and reused across all output rows. This drops the hot dot to one +// vpmovzxbd + two vpermps per 32 weights (1.5 shuffle ops / 16 vs 4 for a bf16, +// dup-permute, per-element gather), so the row-major fp4 GEMV stops being port-5 +// bound and approaches the bf16 memory-bandwidth ceiling. +using dsdot_fn = float (*)(const uint8_t*, const uint8_t*, const float*, const float*, int, + const float*, const float*); + +float dot_dsfp4_scalar(const uint8_t* packed, const uint8_t* scale, const float* xe, + const float* xo, int K, const float* e2m1, const float* e8m0) { + float acc = 0.0f; + const int nb = K / 32; + for (int b = 0; b < nb; ++b) { + const float sc = e8m0[scale[b]]; + const uint8_t* pk = packed + (size_t)b * 16; // 16 bytes = 32 codes + const float* xeb = xe + (size_t)b * 16; + const float* xob = xo + (size_t)b * 16; + float bsum = 0.0f; + for (int j = 0; j < 16; ++j) { + const uint8_t byte = pk[j]; + bsum += e2m1[byte & 0xF] * xeb[j]; // low nibble -> even-K activation + bsum += e2m1[byte >> 4] * xob[j]; // high nibble -> odd-K activation + } + acc += sc * bsum; + } + return acc; +} + +#if CPU_MOE_X86 +// One 32-block: 16 bytes -> 16 low + 16 high nibble values via two vpermps, times +// the pre-split even/odd fp32 activations, folded by the per-32 e8m0 scale. +CPU_MOE_TARGET("avx512f") +inline __m512 dsfp4_blk(const uint8_t* pk, const float* xeb, const float* xob, __m512 lut, + __m512i loma, float sc) { + __m512i wi = _mm512_cvtepu8_epi32(_mm_loadu_si128(reinterpret_cast(pk))); + __m512 vlo = _mm512_permutexvar_ps(_mm512_and_si512(wi, loma), lut); + __m512 vhi = _mm512_permutexvar_ps(_mm512_and_si512(_mm512_srli_epi32(wi, 4), loma), lut); + __m512 prod = _mm512_fmadd_ps(vlo, _mm512_loadu_ps(xeb), _mm512_mul_ps(vhi, _mm512_loadu_ps(xob))); + return _mm512_mul_ps(prod, _mm512_set1_ps(sc)); +} + +CPU_MOE_TARGET("avx512f") +float dot_dsfp4_avx512(const uint8_t* packed, const uint8_t* scale, const float* xe, + const float* xo, int K, const float* e2m1, const float* e8m0) { + const __m512 lut = _mm512_loadu_ps(e2m1); + const __m512i loma = _mm512_set1_epi32(0xF); + __m512 acc0 = _mm512_setzero_ps(), acc1 = _mm512_setzero_ps(); + const int nb = K / 32; // 16 packed bytes + one e8m0 scale per block + int b = 0; + for (; b + 2 <= nb; b += 2) { // two independent accumulators hide FMA latency + acc0 = _mm512_add_ps(acc0, dsfp4_blk(packed + (size_t)b * 16, xe + (size_t)b * 16, + xo + (size_t)b * 16, lut, loma, e8m0[scale[b]])); + acc1 = _mm512_add_ps(acc1, dsfp4_blk(packed + (size_t)(b + 1) * 16, xe + (size_t)(b + 1) * 16, + xo + (size_t)(b + 1) * 16, lut, loma, e8m0[scale[b + 1]])); + } + for (; b < nb; ++b) + acc0 = _mm512_add_ps(acc0, dsfp4_blk(packed + (size_t)b * 16, xe + (size_t)b * 16, + xo + (size_t)b * 16, lut, loma, e8m0[scale[b]])); + return _mm512_reduce_add_ps(_mm512_add_ps(acc0, acc1)); +} + +// AVX2: a 32-K block is 16 bytes -> two 8-lane halves (8 even + 8 odd each). +CPU_MOE_TARGET("avx2,fma") +inline __m256 dsfp4_half_avx2(const uint8_t* pk, const float* xeb, const float* xob, __m256 mag8) { + __m256i wi = _mm256_cvtepu8_epi32(_mm_loadl_epi64(reinterpret_cast(pk))); + __m256 vlo = e2m1_decode8(_mm256_and_si256(wi, _mm256_set1_epi32(0xF)), mag8); + __m256 vhi = e2m1_decode8(_mm256_srli_epi32(wi, 4), mag8); + return _mm256_fmadd_ps(vlo, _mm256_loadu_ps(xeb), _mm256_mul_ps(vhi, _mm256_loadu_ps(xob))); +} + +CPU_MOE_TARGET("avx2,fma") +float dot_dsfp4_avx2(const uint8_t* packed, const uint8_t* scale, const float* xe, + const float* xo, int K, const float* e2m1, const float* e8m0) { + const __m256 mag8 = _mm256_loadu_ps(e2m1); + __m256 acc0 = _mm256_setzero_ps(), acc1 = _mm256_setzero_ps(); + const int nb = K / 32; + for (int b = 0; b < nb; ++b) { + const uint8_t* pk = packed + (size_t)b * 16; + const float* xeb = xe + (size_t)b * 16; + const float* xob = xo + (size_t)b * 16; + const __m256 sc = _mm256_set1_ps(e8m0[scale[b]]); + acc0 = _mm256_fmadd_ps(dsfp4_half_avx2(pk, xeb, xob, mag8), sc, acc0); + acc1 = _mm256_fmadd_ps(dsfp4_half_avx2(pk + 8, xeb + 8, xob + 8, mag8), sc, acc1); + } + return hsum256(_mm256_add_ps(acc0, acc1)); +} +#endif + +dsdot_fn select_dsdot() { + const IsaTier t = pick_isa(); +#if CPU_MOE_X86 + if (t >= ISA_AVX512) return dot_dsfp4_avx512; + if (t >= ISA_AVX2) return dot_dsfp4_avx2; +#endif + (void)t; + return dot_dsfp4_scalar; +} + +// Round a clamped |x|<=448 to nearest float8-e4m3 (RNE), back to fp32. Matches +// torch.float8_e4m3fn / triton .to(float8e4nv). +inline float e4m3_round(float x) { + const float sign = x < 0.0f ? -1.0f : 1.0f; + const float a = std::fabs(x); + if (a == 0.0f) return 0.0f; + if (a >= 448.0f) return sign * 448.0f; + int e; + std::frexp(a, &e); // a in [2^(e-1), 2^e) + float step = std::ldexp(1.0f, e - 4); + const float min_step = std::ldexp(1.0f, -9); // e4m3 subnormal step (2^-9) + if (step < min_step) step = min_step; + float r = std::nearbyint(a / step) * step; + if (r > 448.0f) r = 448.0f; + return sign * r; +} + +// IEEE ceil(log2(v)) for v>0 (matches dsv4 _log2_ceil / fast_round_scale). +inline int ceil_log2_pos(float v) { + uint32_t bits; + std::memcpy(&bits, &v, sizeof(bits)); + const int exp = (int)((bits >> 23) & 0xFF); + const int man = (int)(bits & 0x7FFFFF); + return exp - 127 + (man != 0 ? 1 : 0); +} + +// DeepSeek-V4 / V4.1 activation FP8 round-trip (bf16 in/out): per `block` (the checkpoint's +// fp8 block: 128 on V4, 32 on V4.1), s = 2^ceil(log2(max(|x|,1e-4)/448)); +// y = round_e4m3(clamp(x/s,+-448)) * s. +void fp8_roundtrip_bf16(const bf16_t* src, bf16_t* dst, int K, int block) { + for (int b0 = 0; b0 < K; b0 += block) { + const int b1 = std::min(K, b0 + block); + float amax = 1e-4f; + for (int i = b0; i < b1; ++i) amax = std::max(amax, std::fabs(bf16_to_f32(src[i]))); + const float s = std::ldexp(1.0f, ceil_log2_pos(amax * (1.0f / 448.0f))); + const float inv_s = 1.0f / s; + for (int i = b0; i < b1; ++i) { + float q = bf16_to_f32(src[i]) * inv_s; + q = std::min(448.0f, std::max(-448.0f, q)); + dst[i] = f32_to_bf16(e4m3_round(q) * s); + } + } +} + +class DsFp4Format final : public MoeFormat { + public: + DsFp4Format(const MoeParams& p, const FormatArgs& a) : MoeFormat(p), dot_(select_dsdot()) { + gate_up_ = a.bank("gate_up"); + gu_scale_ = a.bank("gate_up_scale"); + down_ = a.bank("down"); + dn_scale_ = a.bank("down_scale"); + act_block_ = static_cast(a.option("act_block")); + input_prequant_ = a.option("input_prequant", 0) != 0; + const size_t tokens = p.max_tokens, routes = (size_t)p.max_tokens * p.top_k; + xq_.resize(tokens * p.H); + xe_.resize(tokens * (p.H / 2)); + xo_.resize(tokens * (p.H / 2)); + ge_.resize(routes * (p.I / 2)); + go_.resize(routes * (p.I / 2)); + for (int i = 0; i < 16; ++i) e2m1_lut_[i] = kE2M1[i]; + for (int i = 0; i < 256; ++i) e8m0_lut_[i] = std::ldexp(1.0f, std::min(i, 254) - 127); + } + + std::string isa_name() const override { return tier_name(pick_isa()); } + + + // DSV4 FP8-round-trips the input before the gate_up GEMV, unless the GPU already did. + void prep_input(int tok, const bf16_t* x) override { + const bf16_t* src = x; + if (!input_prequant_) { + bf16_t* xq = xq_.data() + (size_t)tok * p.H; + fp8_roundtrip_bf16(x, xq, p.H, act_block_); + src = xq; + } + deinterleave_bf16_f32(src, xe_.data() + (size_t)tok * (p.H / 2), xo_.data() + (size_t)tok * (p.H / 2), p.H); + } + + bool preps_inter() const override { return true; } + + void prep_inter(int64_t r, bf16_t* g) override { + fp8_roundtrip_bf16(g, g, p.I, act_block_); + deinterleave_bf16_f32(g, ge_.data() + (size_t)r * (p.I / 2), go_.data() + (size_t)r * (p.I / 2), p.I); + } + + void gate_up(int layer, int e, int tok, int i0, int i1, const bf16_t* /*x*/, float /*w_in*/, + bf16_t* g) const override { + const int H = p.H, I = p.I; + const uint8_t* gu_packed_l = static_cast(tbl_at(gate_up_, layer)); + const uint8_t* gu_scale_l = static_cast(tbl_at(gu_scale_, layer)); + const int N2 = 2 * I, Hh = H / 2, Hs = H / 32; + const float* xe = xe_.data() + (size_t)tok * (H / 2); + const float* xo = xo_.data() + (size_t)tok * (H / 2); + const uint8_t* gp = gu_packed_l + (size_t)e * N2 * Hh; + const uint8_t* gs = gu_scale_l + (size_t)e * N2 * Hs; + const float lim = p.swiglu_limit; + for (int i = i0; i < i1; ++i) { + // gate_up is stored bf16 by the reference GEMV before swiglu; round to match. + float gate = bf16_to_f32(f32_to_bf16( + dot_(gp + (size_t)i * Hh, gs + (size_t)i * Hs, xe, xo, H, e2m1_lut_, e8m0_lut_))); + float up = bf16_to_f32(f32_to_bf16(dot_( + gp + (size_t)(I + i) * Hh, gs + (size_t)(I + i) * Hs, xe, xo, H, e2m1_lut_, e8m0_lut_))); + if (lim > 0.0f) { + if (gate > lim) gate = lim; + if (up > lim) up = lim; + else if (up < -lim) up = -lim; + } + const float glu = gate / (1.0f + std::exp(-gate)); // silu(gate) + g[i] = f32_to_bf16(glu * up); + } + } + + void down(int layer, int h0, int h1, const DownRoute* routes, int n, float* out) const override { + const int H = p.H, I = p.I, Ih = I / 2, Is = I / 32; + const uint8_t* dn_packed_l = static_cast(tbl_at(down_, layer)); + const uint8_t* dn_scale_l = static_cast(tbl_at(dn_scale_, layer)); + for (int h = h0; h < h1; ++h) { + for (int j = 0; j < n; ++j) { + const size_t e = (size_t)routes[j].e, r = (size_t)routes[j].r; + const uint8_t* dp = dn_packed_l + e * (size_t)H * Ih + (size_t)h * Ih; + const uint8_t* ds = dn_scale_l + e * (size_t)H * Is + (size_t)h * Is; + out[j * kRowBlock + h - h0] = dot_(dp, ds, ge_.data() + r * Ih, go_.data() + r * Ih, I, e2m1_lut_, e8m0_lut_); + } + } + } + + // The reference rounds each route's weighted down output to bf16 before the fp32 sum over routes. + bool rounds_routes() const override { return true; } + + private: + dsdot_fn dot_; + // e2m1 rows [E,2I,H/2] with e8m0 scales per 32 along K [E,2I,H/32]; down likewise with K = I. + const uint64_t* gate_up_ = nullptr; + const uint64_t* gu_scale_ = nullptr; + const uint64_t* down_ = nullptr; + const uint64_t* dn_scale_ = nullptr; + // Activation FP8 round-trip block, the checkpoint's fp8 block. + int act_block_ = 128; + // The caller already FP8-round-tripped the input activations on the GPU + // (same reference grid), so prep_input must not repeat it on the host-callback + // thread. That scalar per-element pass is single-threaded ON THE DECODE CRITICAL + // PATH (~0.3ms/layer at H=4096, every worker and the GPU waiting on it); moving + // it to a captured GPU elementwise kernel removes it while keeping the official + // W4A8 numerics bit-exact (see cpu_executor.py). + bool input_prequant_ = false; + float e2m1_lut_[16]; + float e8m0_lut_[256]; + AlignedVector xq_; + AlignedVector xe_, xo_, ge_, go_; +}; + +} // namespace + +std::unique_ptr make_ds_fp4_format(const MoeParams& p, const FormatArgs& a) { + return std::make_unique(p, a); +} + +} // namespace freetoken::cpu_moe diff --git a/python/freetoken/kernel/csrc/cpu_moe/kernels/fp8_block.cpp b/python/freetoken/kernel/csrc/cpu_moe/kernels/fp8_block.cpp new file mode 100644 index 0000000000..d27fb1868e --- /dev/null +++ b/python/freetoken/kernel/csrc/cpu_moe/kernels/fp8_block.cpp @@ -0,0 +1,203 @@ +#include +#include + +#include "../format.h" + +namespace freetoken::cpu_moe { +namespace { + +// ------------------------ block-FP8 (W8A16) dequant ------------------------ +// fp8-e4m3 weights row-major [rows, K] plus one bf16 scale per 128x128 block. The scale tensor is named weight_scale_inv but multiplies, as in kernel/triton/fp8_blockscale_moe.py. +// Whole 128-blocks read the activations as fp32 even/odd halves (xe[m] = x[2m], xo[m] = x[2m+1]); a ragged tail reads the bf16 row x. +// 0x7F/0xFF (NaN in torch float8_e4m3fn) decode to +-480 in every kernel, as e4m3_decode does. +using fp8_block_dot_fn = float (*)(const uint8_t*, const float*, const float*, const bf16_t*, + const bf16_t*, int, const float*); + +constexpr int FP8_BLK = 128; // K elements covered by one weight scale + +inline float fp8_block_tail(const uint8_t* w, const bf16_t* x, int k0, int K, const float* e4m3) { + float blk = 0.0f; + for (int k = k0; k < K; ++k) blk += e4m3[w[k]] * bf16_to_f32(x[k]); + return blk; +} + +float dot_fp8_block_scalar(const uint8_t* w, const float* xe, const float* xo, const bf16_t* x, + const bf16_t* s, int K, const float* e4m3) { + float acc = 0.0f; + int b = 0, k0 = 0; + for (; k0 + FP8_BLK <= K; k0 += FP8_BLK, ++b) { + // 8 independent sums: a single running sum is bound by the add latency + float a[8] = {}; + const uint8_t* wb = w + k0; + const float* pe = xe + k0 / 2; + const float* po = xo + k0 / 2; + for (int j = 0; j < FP8_BLK / 2; j += 4) + for (int l = 0; l < 4; ++l) { + a[2 * l] += e4m3[wb[2 * (j + l)]] * pe[j + l]; + a[2 * l + 1] += e4m3[wb[2 * (j + l) + 1]] * po[j + l]; + } + acc += (((a[0] + a[1]) + (a[2] + a[3])) + ((a[4] + a[5]) + (a[6] + a[7]))) * + bf16_to_f32(s[b]); + } + if (k0 < K) acc += fp8_block_tail(w, x, k0, K, e4m3) * bf16_to_f32(s[b]); + return acc; +} + +#if CPU_MOE_X86 +// A byte b read as fp16 bits (b & 0x80) << 8 | (b & 0x7F) << 7 is its e4m3 value times 2^-8, subnormals included, so the block scale carries the 256. +// In a 16-bit lane holding bytes (2m, 2m+1), srai 1 then & 0xBF80 builds that for byte 2m+1; a left shift by 8 first does byte 2m. Every AVX2 CPU has F16C. +CPU_MOE_TARGET("avx2,fma,f16c") +float dot_fp8_block_avx2(const uint8_t* w, const float* xe, const float* xo, const bf16_t* x, + const bf16_t* s, int K, const float* e4m3) { + const __m256i keep = _mm256_set1_epi16((short)0xBF80); + __m256 total = _mm256_setzero_ps(); + int b = 0, k0 = 0; + for (; k0 + FP8_BLK <= K; k0 += FP8_BLK, ++b) { + _mm_prefetch(reinterpret_cast(w + k0) + PF_AHEAD, _MM_HINT_T0); + __m256 a0 = _mm256_setzero_ps(), a1 = _mm256_setzero_ps(); + __m256 a2 = _mm256_setzero_ps(), a3 = _mm256_setzero_ps(); + for (int j = 0; j < FP8_BLK; j += 32) { + const __m256i v = _mm256_loadu_si256(reinterpret_cast(w + k0 + j)); + const __m256i odd = _mm256_and_si256(_mm256_srai_epi16(v, 1), keep); + const __m256i even = _mm256_and_si256(_mm256_srai_epi16(_mm256_slli_epi16(v, 8), 1), keep); + const float* pe = xe + (k0 + j) / 2; + const float* po = xo + (k0 + j) / 2; + a0 = _mm256_fmadd_ps(_mm256_cvtph_ps(_mm256_castsi256_si128(even)), _mm256_loadu_ps(pe), a0); + a1 = _mm256_fmadd_ps(_mm256_cvtph_ps(_mm256_extracti128_si256(even, 1)), _mm256_loadu_ps(pe + 8), a1); + a2 = _mm256_fmadd_ps(_mm256_cvtph_ps(_mm256_castsi256_si128(odd)), _mm256_loadu_ps(po), a2); + a3 = _mm256_fmadd_ps(_mm256_cvtph_ps(_mm256_extracti128_si256(odd, 1)), _mm256_loadu_ps(po + 8), a3); + } + const __m256 blk = _mm256_add_ps(_mm256_add_ps(a0, a1), _mm256_add_ps(a2, a3)); + total = _mm256_fmadd_ps(blk, _mm256_set1_ps(bf16_to_f32(s[b]) * 256.0f), total); + } + float acc = hsum256(total); + if (k0 < K) acc += fp8_block_tail(w, x, k0, K, e4m3) * bf16_to_f32(s[b]); + return acc; +} + +CPU_MOE_TARGET("avx512f,avx2,fma,f16c") +float dot_fp8_block_avx512f(const uint8_t* w, const float* xe, const float* xo, const bf16_t* x, + const bf16_t* s, int K, const float* e4m3) { + const __m256i keep = _mm256_set1_epi16((short)0xBF80); + __m512 total = _mm512_setzero_ps(); + int b = 0, k0 = 0; + for (; k0 + FP8_BLK <= K; k0 += FP8_BLK, ++b) { + _mm_prefetch(reinterpret_cast(w + k0) + PF_AHEAD, _MM_HINT_T0); + __m512 a0 = _mm512_setzero_ps(), a1 = _mm512_setzero_ps(); + __m512 a2 = _mm512_setzero_ps(), a3 = _mm512_setzero_ps(); + for (int j = 0; j < FP8_BLK; j += 64) { + const __m256i v0 = _mm256_loadu_si256(reinterpret_cast(w + k0 + j)); + const __m256i v1 = _mm256_loadu_si256(reinterpret_cast(w + k0 + j + 32)); + const __m256i o0 = _mm256_and_si256(_mm256_srai_epi16(v0, 1), keep); + const __m256i e0 = _mm256_and_si256(_mm256_srai_epi16(_mm256_slli_epi16(v0, 8), 1), keep); + const __m256i o1 = _mm256_and_si256(_mm256_srai_epi16(v1, 1), keep); + const __m256i e1 = _mm256_and_si256(_mm256_srai_epi16(_mm256_slli_epi16(v1, 8), 1), keep); + const float* pe = xe + (k0 + j) / 2; + const float* po = xo + (k0 + j) / 2; + a0 = _mm512_fmadd_ps(_mm512_cvtph_ps(e0), _mm512_loadu_ps(pe), a0); + a1 = _mm512_fmadd_ps(_mm512_cvtph_ps(o0), _mm512_loadu_ps(po), a1); + a2 = _mm512_fmadd_ps(_mm512_cvtph_ps(e1), _mm512_loadu_ps(pe + 16), a2); + a3 = _mm512_fmadd_ps(_mm512_cvtph_ps(o1), _mm512_loadu_ps(po + 16), a3); + } + const __m512 blk = _mm512_add_ps(_mm512_add_ps(a0, a1), _mm512_add_ps(a2, a3)); + total = _mm512_fmadd_ps(blk, _mm512_set1_ps(bf16_to_f32(s[b]) * 256.0f), total); + } + float acc = _mm512_reduce_add_ps(total); + if (k0 < K) acc += fp8_block_tail(w, x, k0, K, e4m3) * bf16_to_f32(s[b]); + return acc; +} +#endif // CPU_MOE_X86 + +fp8_block_dot_fn select_fp8_block_dot() { + const IsaTier t = pick_isa(); +#if CPU_MOE_X86 + // the avx512bf16 tier keeps the fp32 kernel: on Sapphire Rapids it is faster than widening to bf16 for dpbf16 + if (t >= ISA_AVX512) return dot_fp8_block_avx512f; + if (t >= ISA_AVX2) return dot_fp8_block_avx2; +#endif + (void)t; + return dot_fp8_block_scalar; +} + +class Fp8BlockFormat final : public MoeFormat { + public: + Fp8BlockFormat(const MoeParams& p, const FormatArgs& a) : MoeFormat(p), dot_(select_fp8_block_dot()) { + gate_up_ = a.bank("gate_up"); + gu_scale_ = a.bank("gate_up_scale"); + down_ = a.bank("down"); + dn_scale_ = a.bank("down_scale"); + gu_scale_stride_ = a.option("gu_scale_stride"); + dn_scale_stride_ = a.option("dn_scale_stride"); + const size_t tokens = p.max_tokens, routes = (size_t)p.max_tokens * p.top_k; + xe_.resize(tokens * (p.H / 2)); + xo_.resize(tokens * (p.H / 2)); + ge_.resize(routes * (p.I / 2)); + go_.resize(routes * (p.I / 2)); + for (int i = 0; i < 256; ++i) e4m3_lut_[i] = e4m3_decode((uint8_t)i); + } + + std::string isa_name() const override { return tier_name(pick_isa()); } + + + void prep_input(int tok, const bf16_t* x) override { + deinterleave_bf16_f32(x, xe_.data() + (size_t)tok * (p.H / 2), xo_.data() + (size_t)tok * (p.H / 2), p.H); + } + + bool preps_inter() const override { return true; } + + void prep_inter(int64_t r, bf16_t* g) override { + deinterleave_bf16_f32(g, ge_.data() + (size_t)r * (p.I / 2), go_.data() + (size_t)r * (p.I / 2), p.I); + } + + // The halves cover whole 128-blocks; the bf16 row covers a ragged tail. + void gate_up(int layer, int e, int tok, int i0, int i1, const bf16_t* x, float w_in, + bf16_t* g) const override { + const int H = p.H, I = p.I; + const uint8_t* packed = static_cast(tbl_at(gate_up_, layer)); + const bf16_t* scale = static_cast(tbl_at(gu_scale_, layer)); + const float* xe = xe_.data() + (size_t)tok * (H / 2); + const float* xo = xo_.data() + (size_t)tok * (H / 2); + gated_rows(p, i0, i1, w_in, g, [&](int row) { + const uint8_t* w = packed + ((size_t)e * (2 * I) + row) * (size_t)H; + const bf16_t* sc = scale + ((size_t)e * ((2 * I + FP8_BLK - 1) / FP8_BLK) + row / FP8_BLK) * + (size_t)gu_scale_stride_; + return dot_(w, xe, xo, x, sc, H, e4m3_lut_); + }); + } + + void down(int layer, int h0, int h1, const DownRoute* routes, int n, float* out) const override { + const int H = p.H, I = p.I; + const uint8_t* packed = static_cast(tbl_at(down_, layer)); + const bf16_t* scale = static_cast(tbl_at(dn_scale_, layer)); + for (int h = h0; h < h1; ++h) { + for (int j = 0; j < n; ++j) { + const size_t e = (size_t)routes[j].e, r = (size_t)routes[j].r; + const uint8_t* w = packed + (e * H + h) * (size_t)I; + const bf16_t* sc = scale + (e * ((H + FP8_BLK - 1) / FP8_BLK) + h / FP8_BLK) * + (size_t)dn_scale_stride_; + out[j * kRowBlock + h - h0] = + dot_(w, ge_.data() + r * (I / 2), go_.data() + r * (I / 2), routes[j].g, sc, I, e4m3_lut_); + } + } + } + + private: + fp8_block_dot_fn dot_; + // e4m3 rows [E,2I,H] and [E,H,I] with one bf16 scale per 128x128 block. + const uint64_t* gate_up_ = nullptr; + const uint64_t* gu_scale_ = nullptr; + const uint64_t* down_ = nullptr; + const uint64_t* dn_scale_ = nullptr; + // Row strides of the scale tensors, including any padding of the host bank layout. + int64_t gu_scale_stride_ = 0, dn_scale_stride_ = 0; + float e4m3_lut_[256]; + AlignedVector xe_, xo_, ge_, go_; +}; + +} // namespace + +std::unique_ptr make_fp8_block_format(const MoeParams& p, const FormatArgs& a) { + return std::make_unique(p, a); +} + +} // namespace freetoken::cpu_moe diff --git a/python/freetoken/kernel/csrc/cpu_moe/kernels/mxfp4.cpp b/python/freetoken/kernel/csrc/cpu_moe/kernels/mxfp4.cpp new file mode 100644 index 0000000000..484d6dc001 --- /dev/null +++ b/python/freetoken/kernel/csrc/cpu_moe/kernels/mxfp4.cpp @@ -0,0 +1,247 @@ +#include "../format.h" + +namespace freetoken::cpu_moe { +namespace { + +// ------------------------- mxfp4 (gpt-oss) GEMV ----------------------------- +// Transposed split-K layout: blk[Kpairs, N2] (N innermost), scl[Kpairs/16, N2] +// e8m0 per 32-K. Computes out[c] = sum_kb (E2M1[lo]*x[2kb] + E2M1[hi]*x[2kb+1]) +// * 2^(e8m0-127) for a contiguous column tile (blk/scl already offset to col 0 of +// the tile). Vectorized over N (16 columns / __m512), K stays the outer (cache- +// sequential) loop. Used by both gate_up (K=H) and down (K=I). +using mxgemv_fn = void (*)(float*, const uint8_t*, const uint8_t*, const bf16_t*, int, int, + int, const float*, const float*); + +void mxfp4_gemv_scalar(float* out, const uint8_t* blk, const uint8_t* scl, const bf16_t* x, + int Kpairs, int N2, int ncol, const float* e2m1, const float* e8m0) { + for (int c = 0; c < ncol; ++c) out[c] = 0.0f; + for (int kb = 0; kb < Kpairs; ++kb) { + const uint8_t* w = blk + (size_t)kb * N2; + const uint8_t* s = scl + (size_t)(kb >> 4) * N2; + const float xl = bf16_to_f32(x[2 * kb]); + const float xh = bf16_to_f32(x[2 * kb + 1]); + for (int c = 0; c < ncol; ++c) { + const uint8_t byte = w[c]; + out[c] += (e2m1[byte & 0xF] * xl + e2m1[byte >> 4] * xh) * e8m0[s[c]]; + } + } +} + +#if CPU_MOE_X86 +CPU_MOE_TARGET("avx512f") +void mxfp4_gemv_avx512(float* out, const uint8_t* blk, const uint8_t* scl, const bf16_t* x, + int Kpairs, int N2, int ncol, const float* e2m1, const float* e8m0) { + (void)e8m0; // e8m0[c]=2^(c-127) computed via bit construction (no gather) + const __m512 lut = _mm512_loadu_ps(e2m1); + const __m512i loma = _mm512_set1_epi32(0xF); + // K-outer / N-inner: each kb cache line is read once and all live column chunks + // (up to 4 -> 64 cols) accumulate from registers, so DRAM/L2 stream the tile once. + int c0 = 0; + for (; c0 + 16 <= ncol; c0 += 64) { + const int nchunk = std::min(4, (ncol - c0) / 16); + __m512 acc[4]; + for (int ci = 0; ci < nchunk; ++ci) acc[ci] = _mm512_setzero_ps(); + for (int kblk = 0; kblk < Kpairs; kblk += 16) { // 16 K-pairs = 32 K = one scale row + __m512 sc[4]; + for (int ci = 0; ci < nchunk; ++ci) { + __m128i sraw = _mm_loadu_si128(reinterpret_cast( + scl + (size_t)(kblk >> 4) * N2 + c0 + ci * 16)); + sc[ci] = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu8_epi32(sraw), 23)); + } + __m512 blk_acc[4]; + for (int ci = 0; ci < nchunk; ++ci) blk_acc[ci] = _mm512_setzero_ps(); + for (int kk = 0; kk < 16; ++kk) { + const int kb = kblk + kk; + const uint8_t* wbase = blk + (size_t)kb * N2 + c0; + // The transposed layout strides K by N2 bytes; prefetch ahead so the strided + // reads are not exposed to DRAM latency (the HW streamer misses big strides). + constexpr int PFD = 8; + if (kb + PFD < Kpairs) + _mm_prefetch(reinterpret_cast(blk + (size_t)(kb + PFD) * N2 + c0), + _MM_HINT_T0); + const __m512 xl = _mm512_set1_ps(bf16_to_f32(x[2 * kb])); + const __m512 xh = _mm512_set1_ps(bf16_to_f32(x[2 * kb + 1])); + for (int ci = 0; ci < nchunk; ++ci) { + __m512i wi = _mm512_cvtepu8_epi32( + _mm_loadu_si128(reinterpret_cast(wbase + ci * 16))); + __m512 vlo = _mm512_permutexvar_ps(_mm512_and_si512(wi, loma), lut); + __m512 vhi = _mm512_permutexvar_ps(_mm512_and_si512(_mm512_srli_epi32(wi, 4), loma), lut); + blk_acc[ci] = _mm512_fmadd_ps(vlo, xl, blk_acc[ci]); + blk_acc[ci] = _mm512_fmadd_ps(vhi, xh, blk_acc[ci]); + } + } + for (int ci = 0; ci < nchunk; ++ci) acc[ci] = _mm512_fmadd_ps(blk_acc[ci], sc[ci], acc[ci]); + } + for (int ci = 0; ci < nchunk; ++ci) _mm512_storeu_ps(out + c0 + ci * 16, acc[ci]); + } + for (int c = c0; c < ncol; ++c) { // tail columns (< 16) + float o = 0.0f; + for (int kb = 0; kb < Kpairs; ++kb) { + const uint8_t byte = blk[(size_t)kb * N2 + c]; + uint32_t bits = (uint32_t)scl[(size_t)(kb >> 4) * N2 + c] << 23; + float sc; + std::memcpy(&sc, &bits, 4); + o += (e2m1[byte & 0xF] * bf16_to_f32(x[2 * kb]) + + e2m1[byte >> 4] * bf16_to_f32(x[2 * kb + 1])) * sc; + } + out[c] = o; + } +} + +CPU_MOE_TARGET("avx2,fma") +void mxfp4_gemv_avx2(float* out, const uint8_t* blk, const uint8_t* scl, const bf16_t* x, + int Kpairs, int N2, int ncol, const float* e2m1, const float* e8m0) { + (void)e8m0; // e8m0[s]=2^(s-127) built via s<<23 (no gather) + const __m256 mag8 = _mm256_loadu_ps(e2m1); + int c0 = 0; + for (; c0 + 8 <= ncol; c0 += 32) { // up to 4 chunks of 8 = 32 cols + const int nchunk = std::min(4, (ncol - c0) / 8); + __m256 acc[4]; + for (int ci = 0; ci < nchunk; ++ci) acc[ci] = _mm256_setzero_ps(); + for (int kblk = 0; kblk < Kpairs; kblk += 16) { // 16 K-pairs = one scale row + __m256 sc[4]; + for (int ci = 0; ci < nchunk; ++ci) { + __m128i sraw = _mm_loadl_epi64(reinterpret_cast( + scl + (size_t)(kblk >> 4) * N2 + c0 + ci * 8)); + sc[ci] = _mm256_castsi256_ps(_mm256_slli_epi32(_mm256_cvtepu8_epi32(sraw), 23)); + } + __m256 blk_acc[4]; + for (int ci = 0; ci < nchunk; ++ci) blk_acc[ci] = _mm256_setzero_ps(); + for (int kk = 0; kk < 16; ++kk) { + const int kb = kblk + kk; + const uint8_t* wbase = blk + (size_t)kb * N2 + c0; + constexpr int PFD = 8; + if (kb + PFD < Kpairs) + _mm_prefetch(reinterpret_cast(blk + (size_t)(kb + PFD) * N2 + c0), + _MM_HINT_T0); + const __m256 xl = _mm256_set1_ps(bf16_to_f32(x[2 * kb])); + const __m256 xh = _mm256_set1_ps(bf16_to_f32(x[2 * kb + 1])); + for (int ci = 0; ci < nchunk; ++ci) { + __m256i wi = _mm256_cvtepu8_epi32( + _mm_loadl_epi64(reinterpret_cast(wbase + ci * 8))); + __m256 vlo = e2m1_decode8(_mm256_and_si256(wi, _mm256_set1_epi32(0xF)), mag8); + __m256 vhi = e2m1_decode8(_mm256_srli_epi32(wi, 4), mag8); + blk_acc[ci] = _mm256_fmadd_ps(vlo, xl, blk_acc[ci]); + blk_acc[ci] = _mm256_fmadd_ps(vhi, xh, blk_acc[ci]); + } + } + for (int ci = 0; ci < nchunk; ++ci) acc[ci] = _mm256_fmadd_ps(blk_acc[ci], sc[ci], acc[ci]); + } + for (int ci = 0; ci < nchunk; ++ci) _mm256_storeu_ps(out + c0 + ci * 8, acc[ci]); + } + for (int c = c0; c < ncol; ++c) { // tail columns (< 8); none when ncol%8==0 + float o = 0.0f; + for (int kb = 0; kb < Kpairs; ++kb) { + const uint8_t byte = blk[(size_t)kb * N2 + c]; + uint32_t bits = (uint32_t)scl[(size_t)(kb >> 4) * N2 + c] << 23; + float sc; + std::memcpy(&sc, &bits, 4); + o += (e2m1[byte & 0xF] * bf16_to_f32(x[2 * kb]) + + e2m1[byte >> 4] * bf16_to_f32(x[2 * kb + 1])) * sc; + } + out[c] = o; + } +} +#endif + +mxgemv_fn select_mxgemv() { + const IsaTier t = pick_isa(); +#if CPU_MOE_X86 + if (t >= ISA_AVX512) return mxfp4_gemv_avx512; + if (t >= ISA_AVX2) return mxfp4_gemv_avx2; +#endif + (void)t; + return mxfp4_gemv_scalar; +} + +// ----------------------------- mxfp4 (gpt-oss) ----------------------------- +// Transposed split-K layout (N innermost), so the GEMV streams K and accumulates +// a contiguous block of N output columns -> cache-line-efficient, no repack, no +// extra host memory. Pass 1 fuses gate_up + clamped-swiglu(+bias); pass 2 adds the +// down bias, and the executor applies the router weight and sums the routes. +// +// Dequant: w = E2M1[code] * 2^(e8m0_scale - 127); two codes per byte (low nibble +// first), one e8m0 scale per 32 contiguous K. Matches kernel/triton/mxfp4_moe.py. +class Mxfp4Format final : public MoeFormat { + public: + Mxfp4Format(const MoeParams& p, const FormatArgs& a) : MoeFormat(p), gemv_(select_mxgemv()) { + gate_up_ = a.bank("gate_up"); + gu_scale_ = a.bank("gate_up_scale"); + gu_bias_ = a.bank("gate_up_bias"); + down_ = a.bank("down"); + dn_scale_ = a.bank("down_scale"); + dn_bias_ = a.bank("down_bias"); + for (int i = 0; i < 16; ++i) e2m1_lut_[i] = kE2M1[i]; + // e8m0 (mxfp4 block scale) = 2^(s-127); the GPU GEMV clamps s to [0,254]. + for (int i = 0; i < 256; ++i) e8m0_lut_[i] = std::ldexp(1.0f, std::min(i, 254) - 127); + } + + std::string isa_name() const override { return tier_name(pick_isa()); } + + void gate_up(int layer, int e, int /*tok*/, int i0, int i1, const bf16_t* x, float /*w_in*/, + bf16_t* g) const override { + const int H = p.H, I = p.I; + const uint8_t* gu_packed_l = static_cast(tbl_at(gate_up_, layer)); + const uint8_t* gu_scale_l = static_cast(tbl_at(gu_scale_, layer)); + const bf16_t* gu_bias_l = static_cast(tbl_at(gu_bias_, layer)); + const int N2 = 2 * I; // gate_up output width (gate/up interleaved) + const int Hh = H / 2; // packed-K rows + const int nunit = i1 - i0; // intermediate units owned by this tile + const int col0 = 2 * i0; // first gate_up column + const int ncol = 2 * nunit; // gate_up columns owned by this tile + const uint8_t* blk_e = gu_packed_l + (size_t)e * Hh * N2; + const uint8_t* scl_e = gu_scale_l + (size_t)e * (size_t)(H / 32) * N2; + float gu[2 * kRowBlock]; + gemv_(gu, blk_e + col0, scl_e + col0, x, Hh, N2, ncol, e2m1_lut_, e8m0_lut_); + const bf16_t* bias_e = gu_bias_l + (size_t)e * N2 + col0; + const float lim = p.swiglu_limit, alpha = p.swiglu_alpha; + for (int j = 0; j < nunit; ++j) { + float gate = gu[2 * j] + bf16_to_f32(bias_e[2 * j]); + float up = gu[2 * j + 1] + bf16_to_f32(bias_e[2 * j + 1]); + if (gate > lim) gate = lim; + if (up > lim) up = lim; + else if (up < -lim) up = -lim; + const float glu = gate / (1.0f + std::exp(-gate * alpha)); // gate * sigmoid(alpha*gate) + g[i0 + j] = f32_to_bf16(glu * (up + 1.0f)); + } + } + + // The transposed layout already streams a whole column block per call, so the routes go one after another. + void down(int layer, int h0, int h1, const DownRoute* routes, int n, float* out) const override { + const int H = p.H, I = p.I; + const int Ih = I / 2; // packed-K rows + const int nh = h1 - h0; + const uint8_t* dn_packed_l = static_cast(tbl_at(down_, layer)); + const uint8_t* dn_scale_l = static_cast(tbl_at(dn_scale_, layer)); + const bf16_t* dn_bias_l = static_cast(tbl_at(dn_bias_, layer)); + for (int j = 0; j < n; ++j) { + const size_t e = (size_t)routes[j].e; + const uint8_t* blk_e = dn_packed_l + e * Ih * H; + const uint8_t* scl_e = dn_scale_l + e * (size_t)(I / 32) * H; + float* part = out + j * kRowBlock; + gemv_(part, blk_e + h0, scl_e + h0, routes[j].g, Ih, H, nh, e2m1_lut_, e8m0_lut_); + const bf16_t* bias_e = dn_bias_l + e * H + h0; + for (int c = 0; c < nh; ++c) part[c] += bf16_to_f32(bias_e[c]); + } + } + + private: + mxgemv_fn gemv_; + // gpt-oss split-K layout, N innermost: blocks [E,H/2,2I], e8m0 scales [E,H/32,2I], bf16 bias [E,2I]; down likewise with K = I, N = H. + const uint64_t* gate_up_ = nullptr; + const uint64_t* gu_scale_ = nullptr; + const uint64_t* gu_bias_ = nullptr; + const uint64_t* down_ = nullptr; + const uint64_t* dn_scale_ = nullptr; + const uint64_t* dn_bias_ = nullptr; + float e2m1_lut_[16]; + float e8m0_lut_[256]; +}; + +} // namespace + +std::unique_ptr make_mxfp4_format(const MoeParams& p, const FormatArgs& a) { + return std::make_unique(p, a); +} + +} // namespace freetoken::cpu_moe diff --git a/python/freetoken/kernel/csrc/cpu_moe/kernels/nvfp4.cpp b/python/freetoken/kernel/csrc/cpu_moe/kernels/nvfp4.cpp new file mode 100644 index 0000000000..429955faf3 --- /dev/null +++ b/python/freetoken/kernel/csrc/cpu_moe/kernels/nvfp4.cpp @@ -0,0 +1,466 @@ +#include + +#include "../format.h" + +namespace freetoken::cpu_moe { +namespace { + +// --------------------------- NVFP4 (W4A16) dequant --------------------------- +// Weights: e2m1 4-bit codes (2/byte, low nibble first), per-16 block scale in +// fp8-e4m3, per-output-row global scale in fp16. Dequant matches the GPU kernels +// (freetoken/kernel/triton/nvfp4_dequant.py): w = E2M1[code] * e4m3(scale) * global. +// Activations stay bf16 (W4A16); the GEMV dequantizes weights inside the K-loop. + + +// e2m1 * 2 as exact int8 (all e2m1 values are multiples of 0.5). Used by the AVX-VNNI +// W4A8 path: nibble -> int8 weight via PSHUFB LUT, then VPDPBUSD against int8 activations; +// the *2 is undone by a 0.5 folded into the final scale. Mirrors ggml's kvalues_mxfp4. +alignas(16) const int8_t kE2M1x2[16] = {0, 1, 2, 3, 4, 6, 8, 12, 0, -1, -2, -3, -4, -6, -8, -12}; + +// Activations pre-deinterleaved to fp32 (xe[m]=x[2m], xo[m]=x[2m+1]); see the +// ds_fp4 dot in ds_fp4.cpp for why (drops the hot loop to ~1.5 shuffle ops / 16 weights). +using nvdot_fn = float (*)(const uint8_t*, const uint8_t*, float, const float*, const float*, + int, const float*, const float*); + +float dot_nvfp4_scalar(const uint8_t* packed, const uint8_t* scale, float global, + const float* xe, const float* xo, int K, const float* e2m1, + const float* e4m3) { + float acc = 0.0f; + const int nb = K / 16; + for (int b = 0; b < nb; ++b) { + const float bs = e4m3[scale[b]]; + const uint8_t* pk = packed + (size_t)b * 8; + const float* xeb = xe + (size_t)b * 8; // 16 K -> 8 even + 8 odd + const float* xob = xo + (size_t)b * 8; + float bsum = 0.0f; + for (int j = 0; j < 8; ++j) { + const uint8_t byte = pk[j]; + bsum += e2m1[byte & 0xF] * xeb[j]; + bsum += e2m1[byte >> 4] * xob[j]; + } + acc += bs * bsum; + } + return acc * global; +} + +// ---- NVFP4 W4A8 (int8 activations) dot: nibble->int8 via LUT, per-16 act scale ---- +// asi8: int8 activations laid out per-16 block as [even(8), odd(8)]; asb[b] = per-block +// activation scale (absmax/127). Result folds the e2m1*2 -> *0.5 into the scale. +using nvi8dot_fn = float (*)(const uint8_t*, const uint8_t*, float, const int8_t*, int, + const float*, const float*); + +[[maybe_unused]] float dot_nvfp4_i8_scalar(const uint8_t* packed, const uint8_t* scale, + float global, const int8_t* asi8, int K, const float* e4m3, + const float* asb) { + float acc = 0.0f; + const int nb = K / 16; + for (int b = 0; b < nb; ++b) { + const uint8_t* pk = packed + (size_t)b * 8; + const int8_t* ae = asi8 + (size_t)b * 16; // even(8) + const int8_t* ao = ae + 8; // odd(8) + int isum = 0; + for (int j = 0; j < 8; ++j) { + isum += (int)kE2M1x2[pk[j] & 0xF] * (int)ae[j]; + isum += (int)kE2M1x2[pk[j] >> 4] * (int)ao[j]; + } + acc += (e4m3[scale[b]] * asb[b]) * (float)isum; + } + return acc * (0.5f * global); +} + +#if CPU_MOE_X86 +// Two 16-K blocks (16 packed bytes) per iter: lo nibbles -> even-K, hi -> odd-K, +// gathered via two vpermps. The per-16 e4m3 scale differs across the two blocks, so +// it is applied per lane (low 8 lanes = block b, high 8 = block b+1). +CPU_MOE_TARGET("avx512f") +inline __m512 nvfp4_blk2(const uint8_t* pk, const float* xeb, const float* xob, __m512 lut, + __m512i loma, float s0, float s1) { + __m512i wi = _mm512_cvtepu8_epi32(_mm_loadu_si128(reinterpret_cast(pk))); + __m512 vlo = _mm512_permutexvar_ps(_mm512_and_si512(wi, loma), lut); + __m512 vhi = _mm512_permutexvar_ps(_mm512_and_si512(_mm512_srli_epi32(wi, 4), loma), lut); + __m512 prod = _mm512_fmadd_ps(vlo, _mm512_loadu_ps(xeb), _mm512_mul_ps(vhi, _mm512_loadu_ps(xob))); + // lanes 0-7 (block b) -> s0, lanes 8-15 (block b+1) -> s1 (pure AVX512F mask move) + __m512 scv = _mm512_mask_mov_ps(_mm512_set1_ps(s0), 0xFF00, _mm512_set1_ps(s1)); + return _mm512_mul_ps(prod, scv); +} + +CPU_MOE_TARGET("avx512f") +float dot_nvfp4_avx512(const uint8_t* packed, const uint8_t* scale, float global, + const float* xe, const float* xo, int K, const float* e2m1, + const float* e4m3) { + const __m512 lut = _mm512_loadu_ps(e2m1); // 16 e2m1 values + const __m512i loma = _mm512_set1_epi32(0xF); + __m512 acc0 = _mm512_setzero_ps(), acc1 = _mm512_setzero_ps(); + const int nb = K / 16; // 8 packed bytes + one e4m3 scale per block + int b = 0; + for (; b + 4 <= nb; b += 4) { // two blk2 calls -> 4 blocks, 2 accumulators + acc0 = _mm512_add_ps(acc0, nvfp4_blk2(packed + (size_t)b * 8, xe + (size_t)b * 8, + xo + (size_t)b * 8, lut, loma, e4m3[scale[b]], e4m3[scale[b + 1]])); + acc1 = _mm512_add_ps(acc1, nvfp4_blk2(packed + (size_t)(b + 2) * 8, xe + (size_t)(b + 2) * 8, + xo + (size_t)(b + 2) * 8, lut, loma, e4m3[scale[b + 2]], e4m3[scale[b + 3]])); + } + for (; b + 2 <= nb; b += 2) + acc0 = _mm512_add_ps(acc0, nvfp4_blk2(packed + (size_t)b * 8, xe + (size_t)b * 8, + xo + (size_t)b * 8, lut, loma, e4m3[scale[b]], e4m3[scale[b + 1]])); + float s = _mm512_reduce_add_ps(_mm512_add_ps(acc0, acc1)); + for (; b < nb; ++b) { // odd final 16-K block + const uint8_t* pk = packed + (size_t)b * 8; + const float* xeb = xe + (size_t)b * 8; + const float* xob = xo + (size_t)b * 8; + float bsum = 0.0f; + for (int j = 0; j < 8; ++j) { + bsum += e2m1[pk[j] & 0xF] * xeb[j] + e2m1[pk[j] >> 4] * xob[j]; + } + s += e4m3[scale[b]] * bsum; + } + return s * global; +} + +// AVX2: one 16-K block (8 packed bytes) per call, 8 even + 8 odd lanes. +CPU_MOE_TARGET("avx2,fma") +inline __m256 nvfp4_blk_avx2(const uint8_t* pk, const float* xeb, const float* xob, + __m256 mag8, float sc) { + __m256i wi = _mm256_cvtepu8_epi32(_mm_loadl_epi64(reinterpret_cast(pk))); + __m256 vlo = e2m1_decode8(_mm256_and_si256(wi, _mm256_set1_epi32(0xF)), mag8); + __m256 vhi = e2m1_decode8(_mm256_srli_epi32(wi, 4), mag8); + __m256 prod = _mm256_fmadd_ps(vlo, _mm256_loadu_ps(xeb), _mm256_mul_ps(vhi, _mm256_loadu_ps(xob))); + return _mm256_mul_ps(prod, _mm256_set1_ps(sc)); +} + +CPU_MOE_TARGET("avx2,fma") +float dot_nvfp4_avx2(const uint8_t* packed, const uint8_t* scale, float global, + const float* xe, const float* xo, int K, const float* e2m1, + const float* e4m3) { + const __m256 mag8 = _mm256_loadu_ps(e2m1); // e2m1[0..7] magnitudes + __m256 acc0 = _mm256_setzero_ps(), acc1 = _mm256_setzero_ps(); + const int nb = K / 16; + int b = 0; + for (; b + 2 <= nb; b += 2) { + acc0 = _mm256_add_ps(acc0, nvfp4_blk_avx2(packed + (size_t)b * 8, xe + (size_t)b * 8, + xo + (size_t)b * 8, mag8, e4m3[scale[b]])); + acc1 = _mm256_add_ps(acc1, nvfp4_blk_avx2(packed + (size_t)(b + 1) * 8, xe + (size_t)(b + 1) * 8, + xo + (size_t)(b + 1) * 8, mag8, e4m3[scale[b + 1]])); + } + for (; b < nb; ++b) + acc0 = _mm256_add_ps(acc0, nvfp4_blk_avx2(packed + (size_t)b * 8, xe + (size_t)b * 8, + xo + (size_t)b * 8, mag8, e4m3[scale[b]])); + return hsum256(_mm256_add_ps(acc0, acc1)) * global; +} + +// AVX-VNNI W4A8: decode 8 packed bytes (16 nibbles) of one 16-block to int8 [lo(8),hi(8)] +// via PSHUFB against the e2m1*2 LUT (replaces the 2 vpermps fp32 expands -- ~4x less +// port-5 traffic). lo=even-K weights, hi=odd-K, matching the [even(8),odd(8)] act layout. +CPU_MOE_TARGET("avx2,avxvnni,fma") +inline __m128i nvfp4_decode_block_i8(const uint8_t* pk, __m128i lut) { + __m128i b = _mm_loadl_epi64(reinterpret_cast(pk)); // 8 bytes + __m128i lo = _mm_and_si128(b, _mm_set1_epi8(0x0F)); + __m128i hi = _mm_and_si128(_mm_srli_epi16(b, 4), _mm_set1_epi8(0x0F)); + return _mm_shuffle_epi8(lut, _mm_unpacklo_epi64(lo, hi)); // [lo(8),hi(8)] -> int8 +} + +// Two 16-blocks per VPDPBUSD (32 int8). Sign trick (ggml mul_add_epi8): |w|*(sign(w)*a)=w*a, +// so u8*s8 needs no offset/correction term. Per-block scale (e4m3 * act-scale) folded in fp32 +// (lanes 0-3 -> block b, 4-7 -> block b+1). Bit-faithful weight; only the int8 activation quant +// (W4A8) differs from the bf16 reference. +CPU_MOE_TARGET("avx2,avxvnni,fma") +float dot_nvfp4_i8_vnni(const uint8_t* packed, const uint8_t* scale, float global, + const int8_t* asi8, int K, const float* e4m3, const float* asb) { + const __m128i lut = _mm_loadu_si128(reinterpret_cast(kE2M1x2)); + __m256 accF = _mm256_setzero_ps(); + const int nb = K / 16; + int b = 0; + for (; b + 2 <= nb; b += 2) { + __m128i wb = nvfp4_decode_block_i8(packed + (size_t)b * 8, lut); + __m128i wb1 = nvfp4_decode_block_i8(packed + (size_t)(b + 1) * 8, lut); + __m256i w = _mm256_set_m128i(wb1, wb); // [blk b | blk b+1] + __m256i a = _mm256_loadu_si256(reinterpret_cast(asi8 + (size_t)b * 16)); + __m256i aw = _mm256_sign_epi8(w, w); // |w| (u8 operand) + __m256i sa = _mm256_sign_epi8(a, w); // sign(w)*a (s8 operand) + __m256i di = _mm256_dpbusd_avx_epi32(_mm256_setzero_si256(), aw, sa); + __m256 scv = _mm256_blend_ps(_mm256_set1_ps(e4m3[scale[b]] * asb[b]), + _mm256_set1_ps(e4m3[scale[b + 1]] * asb[b + 1]), 0xF0); + accF = _mm256_fmadd_ps(_mm256_cvtepi32_ps(di), scv, accF); + } + float s = hsum256(accF); + for (; b < nb; ++b) { // tail (odd block count) + const uint8_t* pk = packed + (size_t)b * 8; + const int8_t* ae = asi8 + (size_t)b * 16; const int8_t* ao = ae + 8; + int isum = 0; + for (int j = 0; j < 8; ++j) + isum += (int)kE2M1x2[pk[j] & 0xF] * (int)ae[j] + (int)kE2M1x2[pk[j] >> 4] * (int)ao[j]; + s += (e4m3[scale[b]] * asb[b]) * (float)isum; + } + return s * (0.5f * global); +} + +#ifdef CPU_MOE_HAS_AVX512VNNI + +// Software-prefetch distance for the W4A8 weight stream, in 16-K blocks (8 packed +// bytes each). Returns -1 when FREETOKEN_CPU_MOE_PF_BLOCKS is unset: the kernel then +// uses the built-in default min(512 blocks = 4 KB, 2 rows) -- 4 KB is the empirical +// optimum on large-row machines (Emerald Rapids sweep), while the 2-row cap keeps a small-row +// model's overshoot bounded (the executor works in 32-row tiles, so a fixed byte +// distance otherwise prefetches another worker's tile: duplicated DRAM traffic that +// regresses at the bandwidth ceiling). An EXPLICIT env value is honored verbatim +// (no clamp; 0 disables): the per-machine optimum can sit past the safe default (+20% +// at 4 KB on a 24-thread Ice Lake with 256B rows), so the escape hatch must reach it. +// Prefetch never faults, so overshooting a row/bank tail is safe. +static int nvfp4_pf_blocks() { + static const int v = [] { + const char* s = getenv("FREETOKEN_CPU_MOE_PF_BLOCKS"); + return (s && s[0]) ? atoi(s) : -1; + }(); + return v; +} + +// AVX-512 VNNI W4A8: FOUR 16-K blocks per VPDPBUSD (64 int8) -- 2x the AVX-VNNI +// (256-bit) path. Decode 32 packed bytes -> 64 int8 with a single _mm512_shuffle_epi8 +// (e2m1*2 LUT replicated to all 4 128-bit lanes). AVX-512 has no _mm512_sign_epi8, so +// the u8*s8 sign trick (|w| as u8, sign(w)*a as s8) uses abs_epi8 + a masked negate. +// Bit-faithful weight; only the int8 activation quant (W4A8) differs from bf16. +// One 4-block group (64 int8) -> scaled fp32 partial (16 lanes). Isolated as a helper so +// the caller can run several independent chains into separate accumulators (the decode -> +// dpbusd -> scale chain is long, so a single accumulator leaves the core latency-bound). +CPU_MOE_TARGET("avx512f,avx512bw,avx512vnni,avx2") +static inline __m512 nvfp4_i8_grp4(const uint8_t* packed, const uint8_t* scale, + const int8_t* asi8, const float* e4m3, const float* asb, + int b, __m512i lut, __m512i idx, __m512i mask0F, + __m512i idxsc) { + const __mmask64 hi_half = 0xFF00FF00FF00FF00ULL; // bytes 8-15 of each 128b lane -> hi nibbles + __m256i raw = _mm256_loadu_si256(reinterpret_cast(packed + (size_t)b * 8)); + __m512i src = _mm512_permutexvar_epi64(idx, _mm512_castsi256_si512(raw)); + __m512i lo = _mm512_and_si512(src, mask0F); + __m512i hi = _mm512_and_si512(_mm512_srli_epi16(src, 4), mask0F); + __m512i comb = _mm512_mask_blend_epi8(hi_half, lo, hi); + __m512i w = _mm512_shuffle_epi8(lut, comb); // 64 int8 weights (e2m1*2) + // u8*s8 sign trick without _mm512_sign_epi8: aw=|w|, sa = (w<0 ? -a : a) (a moot at w==0) + __m512i a = _mm512_loadu_si512(reinterpret_cast(asi8 + (size_t)b * 16)); + __m512i aw = _mm512_abs_epi8(w); + __mmask64 neg = _mm512_movepi8_mask(w); + __m512i sa = _mm512_mask_sub_epi8(a, neg, _mm512_setzero_si512(), a); + __m512i di = _mm512_dpbusd_epi32(_mm512_setzero_si512(), aw, sa); // 16 int32 (groups of 4) + // int32[0..3]->blk b, [4..7]->b+1, [8..11]->b+2, [12..15]->b+3. The 4 block scales + // (e4m3 LUT x per-block act scale) are computed vectorized: a 4-byte load + epu8->epi32 + // widen + one 4-lane LUT gather + one mul replaces 8 scalar loads + 4 scalar muls + + // a set_ps assembly, which otherwise dominates the per-group op count. + int sc_raw; + memcpy(&sc_raw, scale + b, 4); + __m128i sc4 = _mm_cvtepu8_epi32(_mm_cvtsi32_si128(sc_raw)); + __m128 s4 = _mm_mul_ps(_mm_i32gather_ps(e4m3, sc4, 4), _mm_loadu_ps(asb + b)); + __m512 scv = _mm512_permutexvar_ps(idxsc, _mm512_castps128_ps512(s4)); + return _mm512_mul_ps(_mm512_cvtepi32_ps(di), scv); +} + +CPU_MOE_TARGET("avx512f,avx512bw,avx512vnni,avx2") +float dot_nvfp4_i8_avx512vnni(const uint8_t* packed, const uint8_t* scale, float global, + const int8_t* asi8, int K, const float* e4m3, const float* asb) { + const __m512i lut = _mm512_broadcast_i32x4( + _mm_loadu_si128(reinterpret_cast(kE2M1x2))); + const __m512i idx = _mm512_set_epi64(3, 3, 2, 2, 1, 1, 0, 0); // block i -> 128b lane i + const __m512i mask0F = _mm512_set1_epi8(0x0F); + const __m512i idxsc = _mm512_set_epi32(3, 3, 3, 3, 2, 2, 2, 2, 1, 1, 1, 1, 0, 0, 0, 0); + __m512 acc0 = _mm512_setzero_ps(), acc1 = _mm512_setzero_ps(); + __m512 acc2 = _mm512_setzero_ps(), acc3 = _mm512_setzero_ps(); + const int nb = K / 16; + // Built-in default: min(4 KB, 2 rows) -- see nvfp4_pf_blocks(). An explicitly set + // env value is used verbatim so a per-machine sweep can reach operating points past + // the conservative default. + const int pfb = nvfp4_pf_blocks(); + const int pf = (pfb < 0) ? std::min(512, 2 * nb) : pfb; + int b = 0; + // Four independent 4-block groups per iter -> four accumulator chains hide the ~10-op + // decode->dpbusd->scale latency (1-2 chains leave the loop latency-bound: the per-core + // rate sat at ~60% of the core's achievable DRAM stream rate). + for (; b + 16 <= nb; b += 16) { + // Prefetch the weight stream ahead: the interleaved decode lowers the L1-miss + // concurrency the HW prefetcher sustains on its own (134 -> 167 GB/s on Emerald Rapids). + if (pf > 0) { + _mm_prefetch(reinterpret_cast(packed + ((size_t)b + (size_t)pf) * 8), + _MM_HINT_T0); + _mm_prefetch(reinterpret_cast(packed + ((size_t)b + (size_t)pf) * 8 + 64), + _MM_HINT_T0); + } + acc0 = _mm512_add_ps(acc0, nvfp4_i8_grp4(packed, scale, asi8, e4m3, asb, b, lut, idx, + mask0F, idxsc)); + acc1 = _mm512_add_ps(acc1, nvfp4_i8_grp4(packed, scale, asi8, e4m3, asb, b + 4, lut, idx, + mask0F, idxsc)); + acc2 = _mm512_add_ps(acc2, nvfp4_i8_grp4(packed, scale, asi8, e4m3, asb, b + 8, lut, idx, + mask0F, idxsc)); + acc3 = _mm512_add_ps(acc3, nvfp4_i8_grp4(packed, scale, asi8, e4m3, asb, b + 12, lut, idx, + mask0F, idxsc)); + } + for (; b + 4 <= nb; b += 4) + acc0 = _mm512_add_ps(acc0, nvfp4_i8_grp4(packed, scale, asi8, e4m3, asb, b, lut, idx, + mask0F, idxsc)); + float s = _mm512_reduce_add_ps( + _mm512_add_ps(_mm512_add_ps(acc0, acc1), _mm512_add_ps(acc2, acc3))); + for (; b < nb; ++b) { // tail (<4 remaining 16-K blocks) + const uint8_t* pk = packed + (size_t)b * 8; + const int8_t* ae = asi8 + (size_t)b * 16; const int8_t* ao = ae + 8; + int isum = 0; + for (int j = 0; j < 8; ++j) + isum += (int)kE2M1x2[pk[j] & 0xF] * (int)ae[j] + (int)kE2M1x2[pk[j] >> 4] * (int)ao[j]; + s += (e4m3[scale[b]] * asb[b]) * (float)isum; + } + return s * (0.5f * global); +} +#endif // avx512vnni available +#endif + +nvdot_fn select_nvdot() { + const IsaTier t = pick_isa(); +#if CPU_MOE_X86 + if (t >= ISA_AVX512) return dot_nvfp4_avx512; + if (t >= ISA_AVX2) return dot_nvfp4_avx2; +#endif + (void)t; + return dot_nvfp4_scalar; +} + +// Best W4A8 (int8-activation) nvfp4 dot, or nullptr if no SIMD VNNI (caller keeps the +// faithful fp32 nvdot path). The scalar i8 dot exists only as a correctness reference. +nvi8dot_fn select_nvi8dot() { +#if CPU_MOE_X86 +#if defined(CPU_MOE_HAS_AVX512VNNI) + if (cpu_has_avx512vnni()) return dot_nvfp4_i8_avx512vnni; +#endif + if (cpu_has_avxvnni()) return dot_nvfp4_i8_vnni; +#endif + return nullptr; +} + +// Quantize the pre-deinterleaved fp32 even/odd activations to per-16-block int8 in the +// [even(8),odd(8)] layout the VNNI dot expects. Done once per token/route (amortized over +// every output row), so a scalar pass is fine relative to the GEMV. +void quant_i8_pg16(const float* xe, const float* xo, int K, int8_t* asi8, float* asb) { + const int nb = K / 16; + for (int b = 0; b < nb; ++b) { + const float* xeb = xe + (size_t)b * 8; + const float* xob = xo + (size_t)b * 8; + float amax = 0.0f; + for (int j = 0; j < 8; ++j) + amax = std::max(amax, std::max(std::fabs(xeb[j]), std::fabs(xob[j]))); + const float s = amax > 0.0f ? amax / 127.0f : 1.0f; + asb[b] = s; + const float inv = 1.0f / s; + int8_t* ae = asi8 + (size_t)b * 16; + for (int j = 0; j < 8; ++j) { + int qe = (int)std::lround(xeb[j] * inv), qo = (int)std::lround(xob[j] * inv); + ae[j] = (int8_t)std::max(-127, std::min(127, qe)); + ae[8 + j] = (int8_t)std::max(-127, std::min(127, qo)); + } + } +} + +class Nvfp4Format final : public MoeFormat { + public: + Nvfp4Format(const MoeParams& p, const FormatArgs& a) + : MoeFormat(p), nvdot_(select_nvdot()), nvi8dot_(select_nvi8dot()) { + gate_up_ = a.bank("gate_up"); + gu_scale_ = a.bank("gate_up_scale"); + gu_global_ = a.bank("gate_up_global"); + down_ = a.bank("down"); + dn_scale_ = a.bank("down_scale"); + dn_global_ = a.bank("down_global"); + const size_t tokens = p.max_tokens, routes = (size_t)p.max_tokens * p.top_k; + xe_.resize(tokens * (p.H / 2)); + xo_.resize(tokens * (p.H / 2)); + ge_.resize(routes * (p.I / 2)); + go_.resize(routes * (p.I / 2)); + if (nvi8dot_ != nullptr) { + xi8_.resize(tokens * p.H); + xas_.resize(tokens * (p.H / 16)); + gi8_.resize(routes * p.I); + gas_.resize(routes * (p.I / 16)); + } + isa_ = tier_name(pick_isa()); + if (nvi8dot_ != nullptr) isa_ += cpu_has_avx512vnni() ? "+avx512vnni(nvfp4-w4a8)" : "+vnni(nvfp4-w4a8)"; + for (int i = 0; i < 16; ++i) e2m1_lut_[i] = kE2M1[i]; + for (int i = 0; i < 256; ++i) e4m3_lut_[i] = e4m3_decode((uint8_t)i); + } + + std::string isa_name() const override { return isa_; } + + + void prep_input(int tok, const bf16_t* x) override { + float* xe = xe_.data() + (size_t)tok * (p.H / 2); + float* xo = xo_.data() + (size_t)tok * (p.H / 2); + deinterleave_bf16_f32(x, xe, xo, p.H); + if (nvi8dot_ != nullptr) + quant_i8_pg16(xe, xo, p.H, xi8_.data() + (size_t)tok * p.H, xas_.data() + (size_t)tok * (p.H / 16)); + } + + bool preps_inter() const override { return true; } + + void prep_inter(int64_t r, bf16_t* g) override { + float* ge = ge_.data() + (size_t)r * (p.I / 2); + float* go = go_.data() + (size_t)r * (p.I / 2); + deinterleave_bf16_f32(g, ge, go, p.I); + if (nvi8dot_ != nullptr) + quant_i8_pg16(ge, go, p.I, gi8_.data() + (size_t)r * p.I, gas_.data() + (size_t)r * (p.I / 16)); + } + + void gate_up(int layer, int e, int tok, int i0, int i1, const bf16_t* /*x*/, float w_in, + bf16_t* g) const override { + const int H = p.H; + const uint8_t* packed = static_cast(tbl_at(gate_up_, layer)); + const uint8_t* scale = static_cast(tbl_at(gu_scale_, layer)); + const uint16_t* global = static_cast(tbl_at(gu_global_, layer)); + const float* xe = xe_.data() + (size_t)tok * (H / 2); + const float* xo = xo_.data() + (size_t)tok * (H / 2); + const int8_t* xi8 = nvi8dot_ != nullptr ? xi8_.data() + (size_t)tok * H : nullptr; + const float* xas = nvi8dot_ != nullptr ? xas_.data() + (size_t)tok * (H / 16) : nullptr; + gated_rows(p, i0, i1, w_in, g, [&](int row) { + const size_t r = (size_t)e * (2 * p.I) + row; + if (nvi8dot_ != nullptr) + return nvi8dot_(packed + r * (size_t)(H / 2), scale + r * (size_t)(H / 16), + fp16_to_f32(global[r]), xi8, H, e4m3_lut_, xas); + return nvdot_(packed + r * (size_t)(H / 2), scale + r * (size_t)(H / 16), + fp16_to_f32(global[r]), xe, xo, H, e2m1_lut_, e4m3_lut_); + }); + } + + void down(int layer, int h0, int h1, const DownRoute* routes, int n, float* out) const override { + const int I = p.I; + const uint8_t* packed = static_cast(tbl_at(down_, layer)); + const uint8_t* scale = static_cast(tbl_at(dn_scale_, layer)); + const uint16_t* global = static_cast(tbl_at(dn_global_, layer)); + for (int h = h0; h < h1; ++h) { + for (int j = 0; j < n; ++j) { + const size_t r = (size_t)routes[j].r; + const size_t row = (size_t)routes[j].e * p.H + h; + out[j * kRowBlock + h - h0] = nvi8dot_ != nullptr + ? nvi8dot_(packed + row * (size_t)(I / 2), scale + row * (size_t)(I / 16), fp16_to_f32(global[row]), + gi8_.data() + r * I, I, e4m3_lut_, gas_.data() + r * (I / 16)) + : nvdot_(packed + row * (size_t)(I / 2), scale + row * (size_t)(I / 16), fp16_to_f32(global[row]), + ge_.data() + r * (I / 2), go_.data() + r * (I / 2), I, e2m1_lut_, e4m3_lut_); + } + } + } + + private: + nvdot_fn nvdot_; + nvi8dot_fn nvi8dot_; // AVX-VNNI W4A8 nvfp4 dot (nullptr -> use fp32 nvdot) + // ModelOpt rows: packed e2m1 [E,2I,H/2], e4m3 scales per 16 along K [E,2I,H/16], fp16 row globals [E,2I]; down likewise with K = I. + const uint64_t* gate_up_ = nullptr; + const uint64_t* gu_scale_ = nullptr; + const uint64_t* gu_global_ = nullptr; + const uint64_t* down_ = nullptr; + const uint64_t* dn_scale_ = nullptr; + const uint64_t* dn_global_ = nullptr; + std::string isa_; + float e2m1_lut_[16]; + float e4m3_lut_[256]; + // Activations pre-deinterleaved to fp32 even/odd for the row-major dot; with VNNI also per-16-block int8 + scale. + AlignedVector xe_, xo_, ge_, go_; + AlignedVector xi8_, gi8_; + AlignedVector xas_, gas_; +}; + +} // namespace + +std::unique_ptr make_nvfp4_format(const MoeParams& p, const FormatArgs& a) { + return std::make_unique(p, a); +} + +} // namespace freetoken::cpu_moe diff --git a/python/freetoken/kernel/csrc/cpu_moe/kernels/q4_0.cpp b/python/freetoken/kernel/csrc/cpu_moe/kernels/q4_0.cpp new file mode 100644 index 0000000000..600a194ebc --- /dev/null +++ b/python/freetoken/kernel/csrc/cpu_moe/kernels/q4_0.cpp @@ -0,0 +1,216 @@ +#include +#include + +#include "../format.h" + +namespace freetoken::cpu_moe { +namespace { + +// -------------------------------- Q4_0 (W4A8) -------------------------------- +// Native GGUF Q4_0 experts (gemma4 GGUF): per-32 block = fp16 scale d + 16 packed +// bytes; byte j holds element j in its low nibble and j+16 in its high nibble, so a +// block's storage order is [lo0..lo15, hi0..hi15] and w = (nibble - 8) * d. Matches +// the reference dequant (models/gguf/dequant.py) and the packed banks the GPU offload +// path streams. +// +// llama.cpp ggml_vec_dot_q4_0_q8_0: W4A8. The activation is pre-quantized to Q8_0 +// (per-32-block int8 ``aq`` + fp32 scale ``asb``); each block unpacks its 16 bytes to +// 32 int8 weights in [-8,7] (bytes_from_nibbles_32: low nibbles -> elems 0..15, high +// -> 16..31) and runs an integer block dot -- VPDPBUSD (AVX-VNNI) or VPMADDUBSW+VPMADDWD +// (AVX2) with the ggml sign trick |w|*(sign(w)*a)=w*a, or a scalar int loop -- then +// scales the block sum by wd*xd in fp32. No fp weight dequant / shuffle chain. The GPU +// offload path (ggml_moe_a8_vec / MMVQ) is also W4A8, so cpu and hybrid stay close. +using q4dot_fn = float (*)(const uint8_t*, const int8_t*, const float*, int); + +float q4_0_dot_i8_scalar(const uint8_t* w, const int8_t* aq, const float* asb, int K) { + float acc = 0.0f; + const int nb = K / 32; + for (int b = 0; b < nb; ++b) { + const uint8_t* blk = w + (size_t)b * 18; + uint16_t dh; + std::memcpy(&dh, blk, sizeof(dh)); + const uint8_t* q = blk + 2; // 16 nibble bytes + const int8_t* a = aq + (size_t)b * 32; + int isum = 0; + for (int j = 0; j < 16; ++j) { + isum += ((int)(q[j] & 0x0F) - 8) * (int)a[j]; // elem j + isum += ((int)(q[j] >> 4) - 8) * (int)a[16 + j]; // elem 16+j + } + acc += fp16_to_f32(dh) * asb[b] * (float)isum; + } + return acc; +} + +#if CPU_MOE_X86 +// fp16 block scale -> fp32 via HW F16C (single value in lane 0). +CPU_MOE_TARGET("f16c") +static inline float q4_scale(uint16_t h) { + return _mm_cvtss_f32(_mm_cvtph_ps(_mm_cvtsi32_si128((int)h))); +} + +// Unpack one Q4_0 block's 16 bytes -> 32 int8 weights in [-8,7] (elems 0..15 = low +// nibbles, 16..31 = high nibbles). ``eight`` = _mm256_set1_epi8(8). +CPU_MOE_TARGET("avx2") +static inline __m256i q4_unpack32(const uint8_t* blk, __m128i mask, __m256i eight) { + const __m128i qb = _mm_loadu_si128(reinterpret_cast(blk + 2)); + const __m128i lo = _mm_and_si128(qb, mask); + const __m128i hi = _mm_and_si128(_mm_srli_epi16(qb, 4), mask); + return _mm256_sub_epi8(_mm256_set_m128i(hi, lo), eight); +} + +// AVX2 W4A8 (llama.cpp non-VNNI mul_sum_i8_pairs): integer block dot via VPMADDUBSW + +// VPMADDWD (sign trick), scaled by wd*xd. |aw*sa| pair sums <= 8*127*2 < 32767 -> no +// int16 saturation. This is the fast path on AVX2 CPUs without AVX-VNNI (and the +// avx512-tier fallback, since the block dot is 256-bit either way). +CPU_MOE_TARGET("avx2,fma,f16c") +float q4_0_dot_i8_avx2(const uint8_t* w, const int8_t* aq, const float* asb, int K) { + const __m128i mask = _mm_set1_epi8(0x0F); + const __m256i eight = _mm256_set1_epi8(8); + const __m256i ones16 = _mm256_set1_epi16(1); + __m256 accF = _mm256_setzero_ps(); + const int nb = K / 32; + for (int b = 0; b < nb; ++b) { + const uint8_t* blk = w + (size_t)b * 18; + _mm_prefetch(reinterpret_cast(blk) + 512, _MM_HINT_T0); + uint16_t dh; + std::memcpy(&dh, blk, sizeof(dh)); + __m256i wq = q4_unpack32(blk, mask, eight); + __m256i a = _mm256_loadu_si256(reinterpret_cast(aq + (size_t)b * 32)); + __m256i aw = _mm256_sign_epi8(wq, wq); // |wq| (unsigned operand) + __m256i sa = _mm256_sign_epi8(a, wq); // sign(wq) * a (signed operand) + __m256i d32 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, sa), ones16); // 8 int32 + accF = _mm256_fmadd_ps(_mm256_cvtepi32_ps(d32), _mm256_set1_ps(q4_scale(dh) * asb[b]), accF); + } + return hsum256(accF); +} + +// AVX-VNNI W4A8: one VPDPBUSD per block (the fast path on modern CPUs). +CPU_MOE_TARGET("avx2,avxvnni,fma,f16c") +float q4_0_dot_i8_vnni(const uint8_t* w, const int8_t* aq, const float* asb, int K) { + const __m128i mask = _mm_set1_epi8(0x0F); + const __m256i eight = _mm256_set1_epi8(8); + __m256 accF = _mm256_setzero_ps(); + const int nb = K / 32; + for (int b = 0; b < nb; ++b) { + const uint8_t* blk = w + (size_t)b * 18; + _mm_prefetch(reinterpret_cast(blk) + 512, _MM_HINT_T0); + uint16_t dh; + std::memcpy(&dh, blk, sizeof(dh)); + __m256i wq = q4_unpack32(blk, mask, eight); + __m256i a = _mm256_loadu_si256(reinterpret_cast(aq + (size_t)b * 32)); + __m256i aw = _mm256_sign_epi8(wq, wq); // |wq| (unsigned operand) + __m256i sa = _mm256_sign_epi8(a, wq); // sign(wq) * a (signed operand) + __m256i di = _mm256_dpbusd_avx_epi32(_mm256_setzero_si256(), aw, sa); + // All 32 elems of the block share wd*xd; distribute over di's 8 partial sums and + // reduce at the end (equivalent to scale * block_total). + accF = _mm256_fmadd_ps(_mm256_cvtepi32_ps(di), _mm256_set1_ps(q4_scale(dh) * asb[b]), accF); + } + return hsum256(accF); +} +#endif // CPU_MOE_X86 + +// All tiers are W4A8 (int8 activations pre-quantized to Q8_0). AVX-VNNI is orthogonal to +// the ISA tier (gated by cpu_has_avxvnni() / FREETOKEN_CPU_MOE_NO_VNNI), so it wins when +// present; otherwise the 256-bit VPMADDUBSW kernel covers both the avx2 and avx512 tiers. +q4dot_fn select_q4dot() { + const IsaTier t = pick_isa(); +#if CPU_MOE_X86 + if (cpu_has_avxvnni()) return q4_0_dot_i8_vnni; + if (t >= ISA_AVX2) return q4_0_dot_i8_avx2; +#endif + (void)t; + return q4_0_dot_i8_scalar; +} + +// Quantize a bf16 activation row to Q8_0 (llama.cpp): per-32-block symmetric int8 in +// natural order + one fp32 scale (amax/127) per block. Done once per token/route, +// amortized over every output row that the q4_0 W4A8 GEMV reads. K % 32 == 0. +void quant_q8_0(const bf16_t* x, int K, int8_t* aq, float* asb) { + const int nb = K / 32; + for (int b = 0; b < nb; ++b) { + const bf16_t* xb = x + (size_t)b * 32; + float xf[32], amax = 0.0f; + for (int j = 0; j < 32; ++j) { + xf[j] = bf16_to_f32(xb[j]); + amax = std::max(amax, std::fabs(xf[j])); + } + const float d = amax > 0.0f ? amax / 127.0f : 1.0f; + asb[b] = d; + const float inv = amax > 0.0f ? 1.0f / d : 0.0f; + int8_t* o = aq + (size_t)b * 32; + for (int j = 0; j < 32; ++j) + o[j] = (int8_t)std::max(-127, std::min(127, (int)std::lround(xf[j] * inv))); + } +} + +class Q4_0Format final : public MoeFormat { + public: + Q4_0Format(const MoeParams& p, const FormatArgs& a) : MoeFormat(p), dot_(select_q4dot()) { + gate_up_ = a.bank("gate_up"); + down_ = a.bank("down"); + const size_t tokens = p.max_tokens, routes = (size_t)p.max_tokens * p.top_k; + xi8_.resize(tokens * p.H); + xas_.resize(tokens * (p.H / 32)); + gi8_.resize(routes * p.I); + gas_.resize(routes * (p.I / 32)); + if (p.H % 32 != 0 || p.I % 32 != 0) + throw std::runtime_error("Q4_0 CPU MoE requires H and I to be multiples of 32"); + gu_row_bytes_ = (p.H / 32) * 18; // K = H (gate_up rows) + dn_row_bytes_ = (p.I / 32) * 18; // K = I (down rows) + // select_q4dot picks VPDPBUSD / VPMADDUBSW / scalar for the tier, so the tag reflects which of those it resolved to. + isa_ = std::string(tier_name(pick_isa())) + (cpu_has_avxvnni() ? "+vnni(q4_0-w4a8)" : "+q4_0-w4a8"); + } + + std::string isa_name() const override { return isa_; } + + + void prep_input(int tok, const bf16_t* x) override { + quant_q8_0(x, p.H, xi8_.data() + (size_t)tok * p.H, xas_.data() + (size_t)tok * (p.H / 32)); + } + + bool preps_inter() const override { return true; } + + void prep_inter(int64_t r, bf16_t* g) override { + quant_q8_0(g, p.I, gi8_.data() + (size_t)r * p.I, gas_.data() + (size_t)r * (p.I / 32)); + } + + void gate_up(int layer, int e, int tok, int i0, int i1, const bf16_t* /*x*/, float w_in, + bf16_t* g) const override { + const uint8_t* packed = static_cast(tbl_at(gate_up_, layer)); + const int8_t* xi8 = xi8_.data() + (size_t)tok * p.H; + const float* xas = xas_.data() + (size_t)tok * (p.H / 32); + gated_rows(p, i0, i1, w_in, g, [&](int row) { + return dot_(packed + ((size_t)e * (2 * p.I) + row) * (size_t)gu_row_bytes_, xi8, xas, p.H); + }); + } + + void down(int layer, int h0, int h1, const DownRoute* routes, int n, float* out) const override { + const uint8_t* packed = static_cast(tbl_at(down_, layer)); + for (int h = h0; h < h1; ++h) { + for (int j = 0; j < n; ++j) { + const size_t r = (size_t)routes[j].r; + out[j * kRowBlock + h - h0] = dot_(packed + ((size_t)routes[j].e * p.H + h) * (size_t)dn_row_bytes_, + gi8_.data() + r * p.I, gas_.data() + r * (p.I / 32), p.I); + } + } + } + + private: + q4dot_fn dot_; + // GGUF Q4_0 rows of 32-element blocks (fp16 scale + 16 nibble bytes): gate_up [E,2I,H/32*18], down [E,H,I/32*18]. + const uint64_t* gate_up_ = nullptr; + const uint64_t* down_ = nullptr; + int gu_row_bytes_ = 0, dn_row_bytes_ = 0; + std::string isa_; + // Q8_0 activations: per-32-block int8 + fp32 scale, for the input and the intermediate. + AlignedVector xi8_, gi8_; + AlignedVector xas_, gas_; +}; + +} // namespace + +std::unique_ptr make_q4_0_format(const MoeParams& p, const FormatArgs& a) { + return std::make_unique(p, a); +} + +} // namespace freetoken::cpu_moe diff --git a/python/freetoken/kernel/csrc/include/freetoken/hip_compat.h b/python/freetoken/kernel/csrc/include/freetoken/hip_compat.h index 1898b41189..8abcd1a098 100644 --- a/python/freetoken/kernel/csrc/include/freetoken/hip_compat.h +++ b/python/freetoken/kernel/csrc/include/freetoken/hip_compat.h @@ -44,6 +44,18 @@ #define cudaGetLastError hipGetLastError #endif +#ifndef cudaMalloc +#define cudaMalloc hipMalloc +#endif + +#ifndef cudaFree +#define cudaFree hipFree +#endif + +#ifndef cudaMemset +#define cudaMemset hipMemset +#endif + #ifndef cudaMallocHost #define cudaMallocHost hipHostMalloc #endif diff --git a/python/freetoken/kernel/csrc/jit/fast_index_copy.cuh b/python/freetoken/kernel/csrc/jit/fast_index_copy.cuh index 43f3fc1ba5..e37e3d48ad 100644 --- a/python/freetoken/kernel/csrc/jit/fast_index_copy.cuh +++ b/python/freetoken/kernel/csrc/jit/fast_index_copy.cuh @@ -595,3 +595,112 @@ struct MultiIndexCopyKernel { device.unwrap())(kernel, params); } }; + + +// --------------------------------------------------------------------------- +// Slim multi-bank copy: the fused copy from a small grid, so it can share the GPU with compute. The +// grid strides each (row, bank) segment in turn, carrying the phase across, so no bank runs alone at the tail. +template +__always_inline __device__ void slim_copy_segment( + uint4* __restrict__ dst, const uint4* __restrict__ src, int32_t count, int32_t i, int32_t stride +) { + constexpr int32_t kU = static_cast(kUnroll); + // kUnroll independent loads in flight per thread: host-link latency, not issue rate, bounds a small grid + for (; i + (kU - 1) * stride < count; i += kU * stride) { + uint4 v[kUnroll]; +#pragma unroll + for (int32_t k = 0; k < kU; ++k) { + v[k] = details::load_nc(src + i + k * stride); + } +#pragma unroll + for (int32_t k = 0; k < kU; ++k) { + dst[i + k * stride] = v[k]; + } + } + uint4 v[kUnroll]; +#pragma unroll + for (int32_t k = 0; k < kU; ++k) { + if (i + k * stride < count) { + v[k] = details::load_nc(src + i + k * stride); + } + } +#pragma unroll + for (int32_t k = 0; k < kU; ++k) { + if (i + k * stride < count) { + dst[i + k * stride] = v[k]; + } + } +} + +template +__global__ __launch_bounds__(kNumThreads) void fast_index_copy_multi_slim( + const __grid_constant__ MultiIndexCopyParams p +) { + const int64_t n = p.valid_length ? p.valid_length[0] : p.length; + const auto stride = static_cast(gridDim.x * kNumThreads); + const auto tid = static_cast(blockIdx.x * kNumThreads + threadIdx.x); + const auto* di = static_cast(p.dst_indices); + const auto* si = static_cast(p.src_indices); + int32_t phase = 0; // grid thread that takes the next segment's first chunk + for (int64_t row = 0; row < n; ++row) { + const auto ps = static_cast(si[row]); + const auto pd = static_cast(di[row]); + for (int b = 0; b < p.num_banks; ++b) { + const int64_t feat = p.feat_bytes[b]; + const auto units = static_cast(feat >> 4); // feat % 16 == 0; int32 chunk index: rows < 32 GiB + const auto* src = reinterpret_cast(reinterpret_cast(p.src_ptrs[b]) + ps * feat); + auto* dst = reinterpret_cast(reinterpret_cast(p.dst_ptrs[b]) + pd * feat); + const int32_t first = tid >= phase ? tid - phase : tid - phase + stride; + slim_copy_segment(dst, src, units, first, stride); + phase = static_cast((static_cast(phase) + units) % stride); + } + } +} + +template +struct MultiIndexCopySlimKernel { + static void run( + tvm::ffi::TensorView dst_ptrs, + tvm::ffi::TensorView src_ptrs, + tvm::ffi::TensorView feat_bytes, + tvm::ffi::TensorView dst_indices, + tvm::ffi::TensorView src_indices, + tvm::ffi::Optional num_indices + ) { + using namespace host; + auto device = SymbolicDevice{}; + auto B = SymbolicSize{"num_banks"}; + auto L = SymbolicSize{"indices length"}; + auto ptr_dtype = SymbolicDType{}; + auto indices_dtype = SymbolicDType{}; + auto num_indices_dtype = SymbolicDType{}; + + TensorMatcher({B}).with_dtype(ptr_dtype).with_device(device) + .verify(dst_ptrs).verify(src_ptrs).verify(feat_bytes); + TensorMatcher({L}).with_dtype(indices_dtype).with_device(device) + .verify(dst_indices).verify(src_indices); + + const int64_t* valid_length = nullptr; + if (num_indices.has_value()) { + TensorMatcher({1}).with_dtype(num_indices_dtype).with_device(device) + .verify(num_indices.value()); + valid_length = static_cast(num_indices.value().data_ptr()); + } + + const auto params = MultiIndexCopyParams{ + static_cast(dst_ptrs.data_ptr()), + static_cast(src_ptrs.data_ptr()), + static_cast(feat_bytes.data_ptr()), + dst_indices.data_ptr(), + src_indices.data_ptr(), + valid_length, + static_cast(L.unwrap()), + static_cast(B.unwrap()), + }; + const auto use_int32 = indices_dtype.unwrap().bits == 32; + const auto kernel = use_int32 + ? fast_index_copy_multi_slim + : fast_index_copy_multi_slim; + LaunchKernel(kNumBlocks, kNumThreads, device.unwrap())(kernel, params); + } +}; diff --git a/python/freetoken/kernel/fast_index_copy.py b/python/freetoken/kernel/fast_index_copy.py index b2030620ee..3047947c78 100644 --- a/python/freetoken/kernel/fast_index_copy.py +++ b/python/freetoken/kernel/fast_index_copy.py @@ -14,6 +14,11 @@ DEFAULT_NUM_BLOCKS = 4 +# fast_index_copy_multi_slim grid (8K loads in flight, the least that holds ~51 GB/s on an +# RTX 5080 at 4 experts); the AOT spec (kernel/aot.py) builds exactly this config +SLIM_COPY_BLOCKS = 8 +SLIM_COPY_THREADS = 256 +SLIM_COPY_UNROLL = 4 SKIP_FAST_INDEX_COPY_ENV = "FREETOKEN_SKIP_FAST_INDEX_COPY" _TRUE_VALUES = {"1", "true", "yes", "on"} @@ -202,6 +207,43 @@ def fast_index_copy_multi_jit( module.launch(dst_ptrs, src_ptrs, feat_bytes, dst_indices, src_indices, num_indices) +@lru_cache(maxsize=None) +def _jit_fast_index_copy_multi_slim_module(*, num_blocks: int, num_threads: int, unroll: int) -> Module: + args = make_cpp_args(num_blocks, num_threads, unroll) + return load_jit( + "fast_index_copy_multi_slim", + *args, + cuda_files=["fast_index_copy.cuh"], + cuda_wrappers=[("launch", f"&MultiIndexCopySlimKernel<{args}>::run")], + ) + + +def fast_index_copy_multi_slim_jit( + dst_ptrs: torch.Tensor, + src_ptrs: torch.Tensor, + feat_bytes: torch.Tensor, + dst_indices: torch.Tensor, + src_indices: torch.Tensor, + num_indices: torch.Tensor | None = None, + *, + num_blocks: int = SLIM_COPY_BLOCKS, + num_threads: int = SLIM_COPY_THREADS, + unroll: int = SLIM_COPY_UNROLL, +) -> None: + """:func:`fast_index_copy_multi_jit` from a small fixed grid, for copies that run beside compute. + + ``num_blocks`` CTAs split the concatenated (row, bank) 16-byte chunk range evenly and each + thread keeps ``unroll`` loads in flight, so a few SMs still saturate the host link while + the rest stay free for concurrent kernels. Same arguments and alignment contract. + """ + if _skip_fast_index_copy_enabled(): + return + module = _jit_fast_index_copy_multi_slim_module( + num_blocks=num_blocks, num_threads=num_threads, unroll=unroll + ) + module.launch(dst_ptrs, src_ptrs, feat_bytes, dst_indices, src_indices, num_indices) + + def update_copy_flag_jit(sync_flag: torch.Tensor, delta: int) -> None: assert sync_flag.is_cuda assert sync_flag.numel() == 1 diff --git a/python/freetoken/kernel/triton/attention.py b/python/freetoken/kernel/triton/attention.py index c9087e03c2..b90048f079 100644 --- a/python/freetoken/kernel/triton/attention.py +++ b/python/freetoken/kernel/triton/attention.py @@ -6,6 +6,33 @@ import triton import triton.language as tl +from freetoken.kernel.triton.e4m3_compat import KV_TILE_SCALE +from freetoken.kernel.triton.e4m3_compat import kv_load_e4m3_tile_f32 as _kv_load_f32 +from freetoken.kernel.triton.e4m3_compat import ( + kv_load_e4m3_tile_scaled16 as _kv_load_s16, +) +from freetoken.kernel.triton.kv_nvfp4 import load_nvfp4 + + +@triton.jit +def _kv_dequant_scale(scale_ptr, slots, stride, kv_head, mask): + """Per-(token, kv_head) dequant scale for an fp8 KV tile, pre-multiplied by the + 2**8 that :func:`kv_load_e4m3_tile_scaled16` leaves on the tile it returns. + + Applied to a dot's OUTPUT rather than to K or V. The scale is constant down the + reduction dim, so ``scores[m,n] = (sum_d q[m,d]*k[d,n]) * s_k[n]`` and + ``p @ (diag(s_v) @ v) = (p * s_v[None,:]) @ v`` both hold: scaling the + ``BLOCK_M x BLOCK_N`` result costs less than scaling the ``BLOCK_D x BLOCK_N`` K + tile or the ``BLOCK_N x BLOCK_DV`` V tile, by a factor of head_dim / BLOCK_M. + + It is also the more accurate order, which is why the 16-bit tile is safe: every + value the loader returns carries at most the code's own 3 mantissa bits and lies + in +-1.75, so narrowing it to the compute dtype is lossless, whereas multiplying + by a general scale first and narrowing after rounds the product. + """ + s = tl.load(scale_ptr + slots * stride + kv_head, mask=mask, other=0.0) + return s * KV_TILE_SCALE + _MAX_KV_SPLITS = 8 _MIN_BLOCK_KV = 32 @@ -23,25 +50,34 @@ def _num_sms(device_index: int) -> int: return int(torch.cuda.get_device_properties(device_index).multi_processor_count) -def _select_extend_tile(head_dim: int, block_d: int, smem_optin: int) -> tuple[int, int]: +def _select_extend_tile( + head_dim: int, block_d: int, smem_optin: int, kv_bytes: int = 2 +) -> tuple[int, int]: """Pick ``(BLOCK_M, BLOCK_N)`` for the extend/prefill kernel, shared-memory aware. - Larger tiles run materially faster (~2x for head_dim 512 on H100) but their bf16 - q/k/v tiles need about ``(BLOCK_M + 2 * BLOCK_N) * BLOCK_D * 2`` bytes of shared - memory, which overflows consumer GPUs (sm_89 ~99KB opt-in) once head_dim >= 256. - Keep the fast tiles where the device's opt-in shared memory fits them (datacenter - A100/H100); shrink only where it does not. ``smem_optin == 0`` (unknown) conservatively - selects the small tiles, i.e. the prior consumer-safe behavior. + Larger tiles run materially faster (~2x for head_dim 512 on H100) but their q/k/v + tiles need shared memory, which overflows consumer GPUs (sm_89 ~99KB opt-in) once + head_dim >= 256. Keep the fast tiles where the device's opt-in shared memory fits + them (datacenter A100/H100); shrink only where it does not. ``smem_optin == 0`` + (unknown) conservatively selects the small tiles, i.e. the prior consumer-safe + behavior. + + ``kv_bytes`` is the KV cache's element size: q is always 2 bytes/element but K and + V follow the cache, so a 1-byte fp8 cache fits a tile a 16-bit one cannot. Passing + 2 reproduces the previous budget exactly, so the 16-bit ladder is unchanged. """ budget = smem_optin * 0.8 # headroom for scores/acc/alignment/triton scratch def fits(block_m: int, block_n: int) -> bool: - return (block_m + 2 * block_n) * block_d * 2 <= budget + return (block_m * 2 + 2 * block_n * kv_bytes) * block_d <= budget if head_dim <= 128: return 128, 64 if head_dim <= 256: - return (128, 64) if fits(128, 64) else (64, 32) + if fits(128, 64): + return 128, 64 + # Reachable by an fp8 cache where a 16-bit one falls through to 64x32. + return (64, 64) if fits(64, 64) else (64, 32) if head_dim <= 384: return (32, 64) if fits(32, 64) else (32, 32) return (32, 64) if fits(32, 64) else (16, 16) @@ -52,6 +88,10 @@ def _paged_attention_kernel( q_ptr, k_ptr, v_ptr, + k_scale_ptr, + v_scale_ptr, + k_block_ptr, + v_block_ptr, o_ptr, indptr_ptr, indices_ptr, @@ -65,6 +105,8 @@ def _paged_attention_kernel( stride_kh, stride_vs, stride_vh, + stride_kss, + stride_vss, stride_ot, stride_oh, GROUP: tl.constexpr, @@ -73,6 +115,8 @@ def _paged_attention_kernel( BLOCK_N: tl.constexpr, SLIDING_WINDOW: tl.constexpr, HAS_SINKS: tl.constexpr, + HAS_KV_SCALE: tl.constexpr, + KV_NVFP4: tl.constexpr, ): q_tok = tl.program_id(0) q_head = tl.program_id(1) @@ -112,14 +156,36 @@ def _paged_attention_kernel( skip_tile = tl.max(mask_n.to(tl.int32), axis=0) == 0 if not skip_tile: slots = tl.load(indices_ptr + kv_start + offs_n, mask=offs_n < kv_len, other=0) - k = tl.load( - k_ptr - + slots[:, None] * stride_ks - + kv_head * stride_kh - + offs_d[None, :], - mask=(offs_n[:, None] < kv_len) & mask_d[None, :], - other=0.0, - ).to(tl.float32) + if KV_NVFP4: + k = load_nvfp4( + k_ptr, k_block_ptr, k_scale_ptr, slots[:, None], kv_head, + offs_d[None, :], + offs_n[:, None] < kv_len, stride_ks, stride_kh, stride_kss, D, + ) + elif HAS_KV_SCALE: + # fp8 KV: the codes carry magnitude, the per-(token, head) fp32 scale + # restores it. Index math stays int32 like the 16-bit path below. + s_k = tl.load( + k_scale_ptr + slots * stride_kss + kv_head, + mask=offs_n < kv_len, + other=0.0, + ) + k = _kv_load_f32( + k_ptr + + slots[:, None] * stride_ks + + kv_head * stride_kh + + offs_d[None, :], + (offs_n[:, None] < kv_len) & mask_d[None, :], + ) * s_k[:, None] + else: + k = tl.load( + k_ptr + + slots[:, None] * stride_ks + + kv_head * stride_kh + + offs_d[None, :], + mask=(offs_n[:, None] < kv_len) & mask_d[None, :], + other=0.0, + ).to(tl.float32) scores = tl.sum(q[None, :] * k, axis=1) * sm_scale scores = tl.where(mask_n, scores, -float("inf")) @@ -129,14 +195,34 @@ def _paged_attention_kernel( alpha = tl.exp(m_i - m_new) p = tl.exp(scores - m_new) - v = tl.load( - v_ptr - + slots[:, None] * stride_vs - + kv_head * stride_vh - + offs_d[None, :], - mask=(offs_n[:, None] < kv_len) & mask_d[None, :], - other=0.0, - ).to(tl.float32) + if KV_NVFP4: + v = load_nvfp4( + v_ptr, v_block_ptr, v_scale_ptr, slots[:, None], kv_head, + offs_d[None, :], + offs_n[:, None] < kv_len, stride_vs, stride_vh, stride_vss, D, + ) + elif HAS_KV_SCALE: + s_v = tl.load( + v_scale_ptr + slots * stride_vss + kv_head, + mask=offs_n < kv_len, + other=0.0, + ) + v = _kv_load_f32( + v_ptr + + slots[:, None] * stride_vs + + kv_head * stride_vh + + offs_d[None, :], + (offs_n[:, None] < kv_len) & mask_d[None, :], + ) * s_v[:, None] + else: + v = tl.load( + v_ptr + + slots[:, None] * stride_vs + + kv_head * stride_vh + + offs_d[None, :], + mask=(offs_n[:, None] < kv_len) & mask_d[None, :], + other=0.0, + ).to(tl.float32) acc = acc * alpha + tl.sum(p[:, None] * v, axis=0) l_i = l_i * alpha + tl.sum(p, axis=0) m_i = m_new @@ -154,6 +240,10 @@ def _decode_grouped_stage1_kernel( q_ptr, k_ptr, v_ptr, + k_scale_ptr, + v_scale_ptr, + k_block_ptr, + v_block_ptr, sm_scale, indptr_ptr, indices_ptr, @@ -167,6 +257,8 @@ def _decode_grouped_stage1_kernel( stride_kh, stride_vs, stride_vh, + stride_kss, + stride_vss, stride_mid_ob, stride_mid_oh, stride_mid_os, @@ -184,6 +276,8 @@ def _decode_grouped_stage1_kernel( D: tl.constexpr, DV: tl.constexpr, SLIDING_WINDOW: tl.constexpr, + HAS_KV_SCALE: tl.constexpr, + KV_NVFP4: tl.constexpr, ): batch_id = tl.program_id(0) head_block_id = tl.program_id(1) @@ -229,7 +323,10 @@ def _decode_grouped_stage1_kernel( if split_end > split_start: q = tl.load(q_ptr + q_offsets, mask=mask_h[:, None] & mask_d[None, :], other=0.0) - q = q.to(k_ptr.dtype.element_ty) + if not HAS_KV_SCALE: + # A 16-bit cache feeds tl.dot as-is; an fp8 cache is decoded up to q's own + # compute dtype below, so q must NOT be narrowed to the (1-byte) cache type. + q = q.to(k_ptr.dtype.element_ty) for rel_start in tl.range(split_start, split_end, BLOCK_N): rel_offs = rel_start + tl.arange(0, BLOCK_N) @@ -237,24 +334,56 @@ def _decode_grouped_stage1_kernel( logical_offs = effective_start + rel_offs slots = tl.load(indices_ptr + kv_start + logical_offs, mask=mask_n, other=0) - k = tl.load( - k_ptr + slots[None, :] * stride_ks + k_base_offsets, - mask=mask_n[None, :] & mask_d[:, None], - other=0.0, - ) + if KV_NVFP4: + k = load_nvfp4( + k_ptr, k_block_ptr, k_scale_ptr, slots[None, :], kv_head, + offs_d[:, None], + mask_n[None, :], stride_ks, stride_kh, stride_kss, D, + ).to(q.dtype) + elif HAS_KV_SCALE: + s_k = _kv_dequant_scale(k_scale_ptr, slots, stride_kss, kv_head, mask_n) + k = _kv_load_s16( + k_ptr + slots[None, :] * stride_ks + k_base_offsets, + mask_n[None, :] & mask_d[:, None], + ).to(q.dtype) + else: + k = tl.load( + k_ptr + slots[None, :] * stride_ks + k_base_offsets, + mask=mask_n[None, :] & mask_d[:, None], + other=0.0, + ) scores = tl.dot(q, k) * sm_scale + # NVFP4's loader already applies both row and block scales. + if HAS_KV_SCALE and not KV_NVFP4: + scores = scores * s_k[None, :] scores = tl.where(mask_h[:, None] & mask_n[None, :], scores, -float("inf")) - v = tl.load( - v_ptr + slots[:, None] * stride_vs + v_base_offsets, - mask=mask_n[:, None] & mask_dv[None, :], - other=0.0, - ) + if KV_NVFP4: + v = load_nvfp4( + v_ptr, v_block_ptr, v_scale_ptr, slots[:, None], kv_head, + offs_dv[None, :], + mask_n[:, None], stride_vs, stride_vh, stride_vss, D, + ).to(q.dtype) + elif HAS_KV_SCALE: + s_v = _kv_dequant_scale(v_scale_ptr, slots, stride_vss, kv_head, mask_n) + v = _kv_load_s16( + v_ptr + slots[:, None] * stride_vs + v_base_offsets, + mask_n[:, None] & mask_dv[None, :], + ).to(q.dtype) + else: + v = tl.load( + v_ptr + slots[:, None] * stride_vs + v_base_offsets, + mask=mask_n[:, None] & mask_dv[None, :], + other=0.0, + ) m_new = tl.maximum(tl.max(scores, axis=1), m_i) alpha = tl.exp(m_i - m_new) p = tl.exp(scores - m_new[:, None]) - acc = acc * alpha[:, None] + tl.dot(p.to(v.dtype), v) + # p itself stays unscaled: l_i is the softmax denominator and knows + # nothing about V's quantization. + pv = (p * s_v[None, :]) if HAS_KV_SCALE and not KV_NVFP4 else p + acc = acc * alpha[:, None] + tl.dot(pv.to(v.dtype), v) l_i = l_i * alpha + tl.sum(p, axis=1) m_i = m_new @@ -354,6 +483,27 @@ def _decode_stage2_kernel( ) +def _validate_kv_format(q, k, v, kr, vr, quant, kb, vb): + quant = quant if quant is not None else ("fp8" if kr is not None else "none") + if quant not in ("none", "fp8", "nvfp4"): + raise ValueError(f"unknown kv_quant {quant!r}") + d = q.shape[-1] + assert k.shape == v.shape + assert k.stride(-1) == v.stride(-1) == 1 + assert (kr is not None) == (vr is not None) == (quant != "none") + if quant == "nvfp4": + assert d % 16 == 0 and k.shape[-1] == d // 2 + for codes, row, block in ((k, kr, kb), (v, vr, vb)): + assert codes.dtype == torch.uint8 + assert row.dtype == torch.float32 and row.is_contiguous() + assert block is not None and block.dtype == torch.uint8 + assert block.shape == (*codes.shape[:2], d // 16) and block.is_contiguous() + assert codes.device == row.device == block.device == q.device + else: + assert k.shape[-1] == d and kb is None and vb is None + return quant == "nvfp4" + + def decode_paged_attention( q: torch.Tensor, k_cache: torch.Tensor, @@ -369,16 +519,31 @@ def decode_paged_attention( sliding_window: int | None = None, sinks: torch.Tensor | None = None, out: torch.Tensor | None = None, + k_scale: torch.Tensor | None = None, + v_scale: torch.Tensor | None = None, + kv_quant: str | None = None, + k_block_scale: torch.Tensor | None = None, + v_block_scale: torch.Tensor | None = None, ) -> torch.Tensor: - """SGLang-style split-k grouped decode attention for one query per request.""" + """SGLang-style split-k grouped decode attention for one query per request. + + ``kv_quant="nvfp4"`` requires packed half-width codes, FP32 row scales and + uint8 E4M3 block scales. Omitted ``kv_quant`` retains FP8 scale inference. + + ``k_scale`` / ``v_scale`` (``[num_slots, num_kv_heads]`` fp32) turn the cache into + an fp8 KV cache: every code row is multiplied by its own token/head scale. Both + must be given together; ``None`` keeps the 16-bit path byte-identical. + """ assert q.is_cuda and k_cache.is_cuda and v_cache.is_cuda + assert (k_scale is None) == (v_scale is None), "K and V scales come as a pair" assert q.dim() == 3 and k_cache.dim() == 3 and v_cache.dim() == 3 batch, num_q_heads, head_dim = q.shape num_kv_heads = k_cache.shape[1] assert batch == indptr.numel() - 1 assert v_cache.shape[1] == num_kv_heads - assert k_cache.shape[-1] == head_dim and v_cache.shape[-1] == head_dim + nvfp4 = _validate_kv_format(q, k_cache, v_cache, k_scale, v_scale, + kv_quant, k_block_scale, v_block_scale) assert num_q_heads % num_kv_heads == 0 assert attn_logits.shape[0] >= batch assert attn_logits.shape[1] >= num_q_heads @@ -396,6 +561,17 @@ def decode_paged_attention( o = out if out is not None else torch.empty_like(q) sinks_arg = sinks if sinks is not None else q group = num_q_heads // num_kv_heads + # Unused pointer args still need a real tensor (same convention as sinks_arg). + has_kv_scale = k_scale is not None + k_scale_arg = k_scale if has_kv_scale else k_cache + v_scale_arg = v_scale if has_kv_scale else v_cache + k_block_arg = k_block_scale if nvfp4 else k_cache + v_block_arg = v_block_scale if nvfp4 else v_cache + if has_kv_scale: + assert k_scale.shape == v_scale.shape == (k_cache.shape[0], num_kv_heads), ( + tuple(k_scale.shape), + tuple(k_cache.shape), + ) # valid_block_h = heads computed per program (drives the grid + head indexing); block_h = # power-of-two tile size for tl.arange. They differ only for non-power-of-two GQA groups # (e.g. 6), where block_h rounds up and the kernel masks the extra lanes. @@ -415,6 +591,10 @@ def decode_paged_attention( q, k_cache, v_cache, + k_scale_arg, + v_scale_arg, + k_block_arg, + v_block_arg, sm_scale, indptr, indices, @@ -428,6 +608,8 @@ def decode_paged_attention( k_cache.stride(1), v_cache.stride(0), v_cache.stride(1), + k_scale_arg.stride(0), + v_scale_arg.stride(0), attn_logits.stride(0), attn_logits.stride(1), attn_logits.stride(2), @@ -445,6 +627,8 @@ def decode_paged_attention( D=head_dim, DV=head_dim, SLIDING_WINDOW=sliding_window or 0, + HAS_KV_SCALE=has_kv_scale, + KV_NVFP4=nvfp4, num_warps=4, num_stages=2, ) @@ -481,6 +665,10 @@ def _extend_attention_kernel( q_ptr, k_ptr, v_ptr, + k_scale_ptr, + v_scale_ptr, + k_block_ptr, + v_block_ptr, o_ptr, qo_indptr_ptr, kv_indptr_ptr, @@ -495,6 +683,8 @@ def _extend_attention_kernel( stride_kh, stride_vs, stride_vh, + stride_kss, + stride_vss, stride_ot, stride_oh, GROUP: tl.constexpr, @@ -505,6 +695,8 @@ def _extend_attention_kernel( BLOCK_N: tl.constexpr, SLIDING_WINDOW: tl.constexpr, HAS_SINKS: tl.constexpr, + HAS_KV_SCALE: tl.constexpr, + KV_NVFP4: tl.constexpr, HAS_BLOCKS: tl.constexpr, ): seq_id = tl.program_id(0) @@ -565,15 +757,33 @@ def _extend_attention_kernel( skip_tile = tl.max(tl.max(final_mask.to(tl.int32), axis=1), axis=0) == 0 if not skip_tile: slots = tl.load(kv_indices_ptr + kv_start + kv_offsets, mask=mask_n, other=0) - k = tl.load( - k_ptr - + slots[None, :] * stride_ks - + kv_head * stride_kh - + offs_d[:, None], - mask=mask_n[None, :] & mask_d[:, None], - other=0.0, - ) + if KV_NVFP4: + k = load_nvfp4( + k_ptr, k_block_ptr, k_scale_ptr, slots[None, :], kv_head, + offs_d[:, None], + mask_n[None, :], stride_ks, stride_kh, stride_kss, D, + ).to(q.dtype) + elif HAS_KV_SCALE: + s_k = _kv_dequant_scale(k_scale_ptr, slots, stride_kss, kv_head, mask_n) + k = _kv_load_s16( + k_ptr + + slots[None, :] * stride_ks + + kv_head * stride_kh + + offs_d[:, None], + mask_n[None, :] & mask_d[:, None], + ).to(q.dtype) + else: + k = tl.load( + k_ptr + + slots[None, :] * stride_ks + + kv_head * stride_kh + + offs_d[:, None], + mask=mask_n[None, :] & mask_d[:, None], + other=0.0, + ) scores = tl.dot(q.to(k.dtype), k) * sm_scale + if HAS_KV_SCALE and not KV_NVFP4: + scores = scores * s_k[None, :] scores = tl.where(final_mask, scores, -float("inf")) row_max = tl.max(scores, axis=1) @@ -582,15 +792,33 @@ def _extend_attention_kernel( alpha = tl.exp(m_i - m_new) p = tl.exp(scores - m_new[:, None]) - v = tl.load( - v_ptr - + slots[:, None] * stride_vs - + kv_head * stride_vh - + offs_dv[None, :], - mask=mask_n[:, None] & mask_dv[None, :], - other=0.0, - ) - acc = acc * alpha[:, None] + tl.dot(p.to(v.dtype), v) + if KV_NVFP4: + v = load_nvfp4( + v_ptr, v_block_ptr, v_scale_ptr, slots[:, None], kv_head, + offs_dv[None, :], + mask_n[:, None], stride_vs, stride_vh, stride_vss, D, + ).to(q.dtype) + elif HAS_KV_SCALE: + s_v = _kv_dequant_scale(v_scale_ptr, slots, stride_vss, kv_head, mask_n) + v = _kv_load_s16( + v_ptr + + slots[:, None] * stride_vs + + kv_head * stride_vh + + offs_dv[None, :], + mask_n[:, None] & mask_dv[None, :], + ).to(q.dtype) + else: + v = tl.load( + v_ptr + + slots[:, None] * stride_vs + + kv_head * stride_vh + + offs_dv[None, :], + mask=mask_n[:, None] & mask_dv[None, :], + other=0.0, + ) + # p stays unscaled for the l_i denominator below. + pv = (p * s_v[None, :]) if HAS_KV_SCALE and not KV_NVFP4 else p + acc = acc * alpha[:, None] + tl.dot(pv.to(v.dtype), v) l_i = l_i * alpha + tl.sum(p, axis=1) m_i = m_new @@ -631,6 +859,56 @@ def _extend_attend_tile( return m_new, l_i, acc +@triton.jit +def _extend_attend_cache_tile( + q, m_i, l_i, acc, qk_scale, k_cache_ptr, v_cache_ptr, k_scale_ptr, v_scale_ptr, k_block_ptr, v_block_ptr, + slots, kv_head, stride_kcs, stride_kch, stride_vcs, stride_vch, stride_kss, stride_vss, + offs_d, offs_dv, mask_d, mask_dv, mask_n, mask, D: tl.constexpr, MASKED: tl.constexpr, KV_NVFP4: tl.constexpr, +): + """:func:`_extend_attend_tile` over an fp8 or NVFP4 cache tile: K/V are decoded to q's dtype, and an fp8 + tile's per-(token, kv_head) scales go on the dot outputs (see :func:`_kv_dequant_scale`). ``mask_n`` is + required even when ``MASKED`` is false, since the scale loads take it.""" + if KV_NVFP4: + k = load_nvfp4( + k_cache_ptr, k_block_ptr, k_scale_ptr, slots[None, :], kv_head, offs_d[:, None], + mask_n[None, :], stride_kcs, stride_kch, stride_kss, D, + ).to(q.dtype) + else: + s_k = _kv_dequant_scale(k_scale_ptr, slots, stride_kss, kv_head, mask_n) + k = _kv_load_s16( + k_cache_ptr + slots[None, :] * stride_kcs + kv_head * stride_kch + offs_d[:, None], + mask_n[None, :] & mask_d[:, None], + ).to(q.dtype) + scores = tl.dot(q.to(k.dtype), k) + # NVFP4's loader already applies both row and block scales. + if not KV_NVFP4: + scores = scores * s_k[None, :] + if MASKED: + scores = tl.where(mask, scores, -float("inf")) + m_new = tl.maximum(m_i, tl.max(scores, axis=1) * qk_scale) + if MASKED: + m_new = tl.where(m_new == -float("inf"), -1e20, m_new) + alpha = tl.math.exp2(m_i - m_new) + p = tl.math.exp2(scores * qk_scale - m_new[:, None]) + if KV_NVFP4: + v = load_nvfp4( + v_cache_ptr, v_block_ptr, v_scale_ptr, slots[:, None], kv_head, offs_dv[None, :], + mask_n[:, None], stride_vcs, stride_vch, stride_vss, D, + ).to(q.dtype) + pv = p + else: + s_v = _kv_dequant_scale(v_scale_ptr, slots, stride_vss, kv_head, mask_n) + v = _kv_load_s16( + v_cache_ptr + slots[:, None] * stride_vcs + kv_head * stride_vch + offs_dv[None, :], + mask_n[:, None] & mask_dv[None, :], + ).to(q.dtype) + # p stays unscaled for the l_i denominator below. + pv = p * s_v[None, :] + acc = acc * alpha[:, None] + tl.dot(pv.to(v.dtype), v) + l_i = l_i * alpha + tl.sum(p, axis=1) + return m_new, l_i, acc + + @triton.jit def _extend_tile_ranges( first_row, diag, end, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, SLIDING_WINDOW: tl.constexpr @@ -654,6 +932,10 @@ def _extend_attention_split_kernel( v_extend_ptr, k_cache_ptr, v_cache_ptr, + k_scale_ptr, + v_scale_ptr, + k_block_ptr, + v_block_ptr, o_ptr, qo_indptr_ptr, kv_indptr_ptr, @@ -672,6 +954,8 @@ def _extend_attention_split_kernel( stride_kch, stride_vcs, stride_vch, + stride_kss, + stride_vss, stride_ot, stride_oh, GROUP: tl.constexpr, @@ -682,6 +966,8 @@ def _extend_attention_split_kernel( BLOCK_N: tl.constexpr, SLIDING_WINDOW: tl.constexpr, HAS_SINKS: tl.constexpr, + HAS_KV_SCALE: tl.constexpr, + KV_NVFP4: tl.constexpr, HAS_BLOCKS: tl.constexpr, ): seq_id = tl.program_id(0) @@ -732,10 +1018,17 @@ def _extend_attention_split_kernel( ) for start_n in tl.range(u0, u1, BLOCK_N): slots = tl.load(kv_indices_ptr + kv_start + start_n + offs_n) - m_i, l_i, acc = _extend_attend_tile( - q, m_i, l_i, acc, qk_scale, kc_base, vc_base, slots * stride_kcs, slots * stride_vcs, - offs_d, offs_dv, mask_d, mask_dv, None, None, False, - ) + if HAS_KV_SCALE: + m_i, l_i, acc = _extend_attend_cache_tile( + q, m_i, l_i, acc, qk_scale, k_cache_ptr, v_cache_ptr, k_scale_ptr, v_scale_ptr, k_block_ptr, + v_block_ptr, slots, kv_head, stride_kcs, stride_kch, stride_vcs, stride_vch, stride_kss, stride_vss, + offs_d, offs_dv, mask_d, mask_dv, offs_n < BLOCK_N, None, D, False, KV_NVFP4, + ) + else: + m_i, l_i, acc = _extend_attend_tile( + q, m_i, l_i, acc, qk_scale, kc_base, vc_base, slots * stride_kcs, slots * stride_vcs, + offs_d, offs_dv, mask_d, mask_dv, None, None, False, + ) n_left = (u0 - lo) // BLOCK_N for i in tl.range(0, n_left + tl.cdiv(prefix_end - u1, BLOCK_N)): kv_offsets = tl.where(i < n_left, lo + i * BLOCK_N, u1 + (i - n_left) * BLOCK_N) + offs_n @@ -744,10 +1037,17 @@ def _extend_attention_split_kernel( if SLIDING_WINDOW > 0: final_mask = final_mask & ((kv_offsets[None, :] + SLIDING_WINDOW) > q_abs_pos[:, None]) slots = tl.load(kv_indices_ptr + kv_start + kv_offsets, mask=mask_n, other=0) - m_i, l_i, acc = _extend_attend_tile( - q, m_i, l_i, acc, qk_scale, kc_base, vc_base, slots * stride_kcs, slots * stride_vcs, - offs_d, offs_dv, mask_d, mask_dv, mask_n, final_mask, True, - ) + if HAS_KV_SCALE: + m_i, l_i, acc = _extend_attend_cache_tile( + q, m_i, l_i, acc, qk_scale, k_cache_ptr, v_cache_ptr, k_scale_ptr, v_scale_ptr, k_block_ptr, + v_block_ptr, slots, kv_head, stride_kcs, stride_kch, stride_vcs, stride_vch, stride_kss, stride_vss, + offs_d, offs_dv, mask_d, mask_dv, mask_n, final_mask, D, True, KV_NVFP4, + ) + else: + m_i, l_i, acc = _extend_attend_tile( + q, m_i, l_i, acc, qk_scale, kc_base, vc_base, slots * stride_kcs, slots * stride_vcs, + offs_d, offs_dv, mask_d, mask_dv, mask_n, final_mask, True, + ) ke_base = k_extend_ptr + kv_head * stride_keh ve_base = v_extend_ptr + kv_head * stride_veh @@ -814,18 +1114,32 @@ def extend_paged_attention( out: torch.Tensor | None = None, k_extend: torch.Tensor | None = None, v_extend: torch.Tensor | None = None, + k_scale: torch.Tensor | None = None, + v_scale: torch.Tensor | None = None, + kv_quant: str | None = None, + k_block_scale: torch.Tensor | None = None, + v_block_scale: torch.Tensor | None = None, block_ends: torch.Tensor | None = None, ) -> torch.Tensor: - """Block-tiled causal prefill/extend attention over paged KV cache; block_ends holds per query token the end of the multimodal span it sits in (0 for none), whose later keys the row also attends.""" + """Block-tiled causal prefill/extend attention over paged KV cache. + + ``k_scale`` / ``v_scale`` mark an fp8 KV cache (see ``decode_paged_attention``). + The ``k_extend`` / ``v_extend`` rows are the current request's own K/V and stay in + the compute dtype either way, so only the cached prefix is decoded. + ``block_ends`` holds per query token the end of the multimodal span it sits in (0 for + none), whose later keys the row also attends. + """ assert q.is_cuda and k_cache.is_cuda and v_cache.is_cuda + assert (k_scale is None) == (v_scale is None), "K and V scales come as a pair" assert q.dim() == 3 and k_cache.dim() == 3 and v_cache.dim() == 3 num_q_tokens, num_q_heads, head_dim = q.shape num_kv_heads = k_cache.shape[1] assert qo_indptr.numel() == kv_indptr.numel() assert prefix_lens.numel() == qo_indptr.numel() - 1 assert v_cache.shape[1] == num_kv_heads - assert k_cache.shape[-1] == head_dim and v_cache.shape[-1] == head_dim + nvfp4 = _validate_kv_format(q, k_cache, v_cache, k_scale, v_scale, + kv_quant, k_block_scale, v_block_scale) assert num_q_heads % num_kv_heads == 0 if sinks is not None: assert sinks.is_cuda @@ -837,14 +1151,27 @@ def extend_paged_attention( assert block_ends.is_cuda and block_ends.dtype == torch.int32 and block_ends.numel() == num_q_tokens o = out if out is not None else torch.empty_like(q) sinks_arg = sinks if sinks is not None else q + has_kv_scale = k_scale is not None + # Unused pointer args still need a real tensor (same convention as sinks_arg). + k_scale_arg = k_scale if has_kv_scale else k_cache + v_scale_arg = v_scale if has_kv_scale else v_cache + k_block_arg = k_block_scale if nvfp4 else k_cache + v_block_arg = v_block_scale if nvfp4 else v_cache + if has_kv_scale: + assert k_scale.shape == v_scale.shape == (k_cache.shape[0], num_kv_heads), ( + tuple(k_scale.shape), + tuple(k_cache.shape), + ) block_ends_arg = block_ends if block_ends is not None else qo_indptr block_d = triton.next_power_of_2(head_dim) block_dv = triton.next_power_of_2(head_dim) # Tile size is shared-memory bound: keep the fast (large) tiles on GPUs whose opt-in # shared memory fits them, shrink on consumer GPUs (sm_89 ~99KB) where the default # 128x64 overflows once head_dim >= 256 (e.g. gemma4: SWA 256, full-attention 512). + # NVFP4 restores full-width tiles with block scales; keep its 16-bit budget. + kv_tile_bytes = 2 if nvfp4 else k_cache.element_size() block_m, block_n = _select_extend_tile( - head_dim, block_d, _optin_smem_bytes(q.device.index) + head_dim, block_d, _optin_smem_bytes(q.device.index), kv_tile_bytes ) grid = (qo_indptr.numel() - 1, num_q_heads, triton.cdiv(max_q_len, block_m)) if k_extend is not None or v_extend is not None: @@ -860,6 +1187,10 @@ def extend_paged_attention( v_extend, k_cache, v_cache, + k_scale_arg, + v_scale_arg, + k_block_arg, + v_block_arg, o, qo_indptr, kv_indptr, @@ -878,6 +1209,8 @@ def extend_paged_attention( k_cache.stride(1), v_cache.stride(0), v_cache.stride(1), + k_scale_arg.stride(0), + v_scale_arg.stride(0), o.stride(0), o.stride(1), GROUP=num_q_heads // num_kv_heads, @@ -888,6 +1221,8 @@ def extend_paged_attention( BLOCK_N=block_n, SLIDING_WINDOW=sliding_window or 0, HAS_SINKS=sinks is not None, + HAS_KV_SCALE=has_kv_scale, + KV_NVFP4=nvfp4, HAS_BLOCKS=block_ends is not None, num_warps=8, num_stages=1, @@ -898,6 +1233,10 @@ def extend_paged_attention( q, k_cache, v_cache, + k_scale_arg, + v_scale_arg, + k_block_arg, + v_block_arg, o, qo_indptr, kv_indptr, @@ -912,6 +1251,8 @@ def extend_paged_attention( k_cache.stride(1), v_cache.stride(0), v_cache.stride(1), + k_scale_arg.stride(0), + v_scale_arg.stride(0), o.stride(0), o.stride(1), GROUP=num_q_heads // num_kv_heads, @@ -922,6 +1263,8 @@ def extend_paged_attention( BLOCK_N=block_n, SLIDING_WINDOW=sliding_window or 0, HAS_SINKS=sinks is not None, + HAS_KV_SCALE=has_kv_scale, + KV_NVFP4=nvfp4, HAS_BLOCKS=block_ends is not None, num_warps=8, num_stages=1, @@ -942,20 +1285,28 @@ def paged_attention( sinks: torch.Tensor | None = None, out: torch.Tensor | None = None, block_n: int = 32, + k_scale: torch.Tensor | None = None, + v_scale: torch.Tensor | None = None, + kv_quant: str | None = None, + k_block_scale: torch.Tensor | None = None, + v_block_scale: torch.Tensor | None = None, ) -> torch.Tensor: """Paged causal attention for one layer. ``q`` is ``[num_query_tokens, num_q_heads, head_dim]``. KV cache tensors are flattened to ``[num_slots, num_kv_heads, head_dim]``. ``indptr`` and - ``indices`` describe each request's logical KV slots in order. + ``indices`` describe each request's logical KV slots in order. ``k_scale`` / + ``v_scale`` (``[num_slots, num_kv_heads]`` fp32) mark an fp8 codes cache. """ assert q.is_cuda and k_cache.is_cuda and v_cache.is_cuda + assert (k_scale is None) == (v_scale is None), "K and V scales come as a pair" assert q.dim() == 3 and k_cache.dim() == 3 and v_cache.dim() == 3 num_tokens, num_q_heads, head_dim = q.shape num_kv_heads = k_cache.shape[1] assert v_cache.shape[1] == num_kv_heads - assert k_cache.shape[-1] == head_dim and v_cache.shape[-1] == head_dim + nvfp4 = _validate_kv_format(q, k_cache, v_cache, k_scale, v_scale, + kv_quant, k_block_scale, v_block_scale) assert num_q_heads % num_kv_heads == 0 if sinks is not None: assert sinks.is_cuda @@ -965,12 +1316,26 @@ def paged_attention( o = out if out is not None else torch.empty_like(q) sinks_arg = sinks if sinks is not None else q + has_kv_scale = k_scale is not None + k_scale_arg = k_scale if has_kv_scale else k_cache + v_scale_arg = v_scale if has_kv_scale else v_cache + k_block_arg = k_block_scale if nvfp4 else k_cache + v_block_arg = v_block_scale if nvfp4 else v_cache + if has_kv_scale: + assert k_scale.shape == v_scale.shape == (k_cache.shape[0], num_kv_heads), ( + tuple(k_scale.shape), + tuple(k_cache.shape), + ) block_d = triton.next_power_of_2(head_dim) grid = (num_tokens, num_q_heads) _paged_attention_kernel[grid]( q, k_cache, v_cache, + k_scale_arg, + v_scale_arg, + k_block_arg, + v_block_arg, o, indptr, indices, @@ -984,6 +1349,8 @@ def paged_attention( k_cache.stride(1), v_cache.stride(0), v_cache.stride(1), + k_scale_arg.stride(0), + v_scale_arg.stride(0), o.stride(0), o.stride(1), GROUP=num_q_heads // num_kv_heads, @@ -992,6 +1359,8 @@ def paged_attention( BLOCK_N=block_n, SLIDING_WINDOW=sliding_window or 0, HAS_SINKS=sinks is not None, + HAS_KV_SCALE=has_kv_scale, + KV_NVFP4=nvfp4, num_warps=8 if head_dim >= 256 else 4, num_stages=2, ) diff --git a/python/freetoken/kernel/triton/autotune_cache.py b/python/freetoken/kernel/triton/autotune_cache.py index ea95768368..2e57fadde9 100644 --- a/python/freetoken/kernel/triton/autotune_cache.py +++ b/python/freetoken/kernel/triton/autotune_cache.py @@ -34,4 +34,28 @@ {"cache_results": _CACHE_RESULTS} if _SUPPORTS_AUTOTUNE_CACHE else {} ) -__all__ = ["autotune_cache_kwargs"] + + +def bound_autotune_flush_buffer(device_index: int) -> None: + """Size triton's L2-flush buffer for autotune benchmarks to the device instead of a flat 256 MB. + + A kernel autotunes on its first launch of each key, which can be the first request, after the + cache planner has handed out nearly all VRAM; a 256 MB transient there is an OOM on a small + card. Four times the L2 still evicts everything between timed runs. + """ + import torch + + try: + from triton.backends.nvidia.driver import CudaDriver + except ImportError: + return + l2 = int(getattr(torch.cuda.get_device_properties(device_index), "L2_cache_size", 0) or 0) + nbytes = min(256 << 20, max(8 << 20, 4 * l2)) + + def get_empty_cache_for_benchmark(self): + return torch.empty(nbytes // 4, dtype=torch.int, device="cuda") + + CudaDriver.get_empty_cache_for_benchmark = get_empty_cache_for_benchmark + + +__all__ = ["autotune_cache_kwargs", "bound_autotune_flush_buffer"] diff --git a/python/freetoken/kernel/triton/e4m3_compat.py b/python/freetoken/kernel/triton/e4m3_compat.py index 14a7acd835..f0a4464a5a 100644 --- a/python/freetoken/kernel/triton/e4m3_compat.py +++ b/python/freetoken/kernel/triton/e4m3_compat.py @@ -5,8 +5,18 @@ Affected kernels branch on :func:`e4m3_native_cx` (a compile-time constexpr): the native branch stays byte-identical on sm_89+, the emulated branch is dead-code eliminated there. When the emulated branch is active, wrappers must pass e4m3 -tensors as ``.view(torch.uint8)`` and allocate act-quant outputs as bf16 -- use -the host-side twin :func:`e4m3_native` for those decisions. +tensors as ``.view(torch.uint8)`` and allocate act-quant outputs as bf16. + +TWO independent probes answer "is fp8e4nv native here": :func:`e4m3_native` (host, +torch's device capability -- decides what a buffer is ALLOCATED as) and +:func:`e4m3_native_cx` (compile-time, triton's target -- decides which arm a kernel +compiles). They cannot be merged, because a constexpr function that referenced the +host one would not survive triton's cache-key AST walk, so on a box where the probes +disagree the host holds real fp8 tensors while kernels take the emulated arm. +:func:`warn_if_probes_disagree` says so once at startup. Code handed a buffer whose +dtype the host already settled -- the KV pool -- must not re-ask at all and picks its +decode from the pointer: :func:`kv_load_e4m3_tile_f32`. Trusting the probe over the +tensor is what produced "cannot cast int32 to fp8e4nv" at CUDA graph capture. ``FREETOKEN_FORCE_E4M3_EMU=1`` (or true/yes/on) forces the emulated path on any GPU (for A/B validation against the native fp8 unit). The flag is read ONCE at @@ -47,8 +57,9 @@ def _env_force() -> bool: def e4m3_native() -> bool: - """Host-side twin of :func:`e4m3_native_cx`: True when kernels take fp8e4nv - tensors directly. False: pass ``.view(torch.uint8)`` and bf16 act buffers.""" + """Host-side probe: does THIS device take fp8e4nv tensors directly? True: kernels + get fp8 tensors, False: pass ``.view(torch.uint8)`` and bf16 act buffers. NOT + necessarily the answer :func:`e4m3_native_cx` gives -- see this module's header.""" global _native if _env_force() != FORCE_EMU: raise RuntimeError( @@ -68,9 +79,47 @@ def e4m3_native() -> bool: # one process runs on one GPU, so its convention is that GPU's; None (-> the current device) only before the process binds _native = torch.cuda.get_device_capability(assigned_visible_gpu()) >= (8, 9) + warn_if_probes_disagree() return _native +_warned_disagree = False + + +def warn_if_probes_disagree() -> None: + """Log once when the two native-fp8e4nv probes answer differently. + + :func:`e4m3_native` decides what buffers the host ALLOCATES while + :func:`e4m3_native_cx` decides which arm a kernel compiles, and the two cannot be + unified (that function's docstring explains triton's cache-key walk). A box where + triton's probe under-reports therefore runs every e4m3 kernel through the software + decode -- bit-exact per this module's header, but slower -- and any kernel that + trusts the probe over the tensor it was handed stops compiling outright. Reads the + latched ``_native`` rather than calling back into :func:`e4m3_native`.""" + global _warned_disagree + if _warned_disagree: + return + _warned_disagree = True + if FORCE_EMU: # emulating by request is not a disagreement + return + try: + triton_native = target_info.cuda_capability_geq(8, 9) + if triton_native == bool(_native): + return + major, minor = torch.cuda.get_device_capability() + except Exception: # noqa: BLE001 -- no driver/no target yet: nothing to compare + return + from freetoken.utils import init_logger + + init_logger(__name__).warning( + "native fp8e4nv disagreement: torch reports sm_%d%d for this device but " + "triton's target probe says %s, so e4m3 kernels compile the software-decode " + "branch (bit-exact, slower). The fp8 KV cache is unaffected -- it follows the " + "buffer it was given.", + major, minor, "supported" if triton_native else "unsupported", + ) + + def e4m3_kernel_view(t: torch.Tensor) -> torch.Tensor: """An e4m3 tensor as the branched kernels expect it: unchanged when native, the uint8 view otherwise (the fp8 pointer type is illegal pre-sm_89).""" @@ -88,7 +137,17 @@ def e4m3_native_cx(): """Compile-time: does the compilation target have native fp8e4nv (sm_89+)? Delegates to ``target_info`` (reads the active driver's target, so cross-compilation tests that patch ``driver.active.get_current_target`` - resolve consistently).""" + resolve consistently). + + It CANNOT defer to :func:`e4m3_native`, however much one verdict per process is + what we want: triton hashes a constexpr function by walking its AST + (runtime/jit.py: cache_key -> record_reference), and a bare reference to a plain + python function raises "Unsupported function referenced: " + -- trying that once disabled every e4m3 kernel at once, PLE gather included. + Module attributes (``target_info.whatever``) survive the walk, plain functions do + not. The probes therefore stay separate, :func:`warn_if_probes_disagree` reports + when they disagree, and code handed a buffer the host already typed -- the KV + pool -- ignores this function and follows the pointer: kv_load_e4m3_tile_f32.""" return not FORCE_EMU and target_info.cuda_capability_geq(8, 9) @@ -123,6 +182,16 @@ def e4m3_f32_to_u8(x): return sign | tl.where(tl.abs(x) >= 0.015625, normal, sub) +@jit +def e4m3_to_f16_x128(raw): + """An e4m3 value as a branched kernel loads it (fp8 on sm_89+, its uint8 bits + before) -> fp16 pre-scaled by 128, the same value on both.""" + if e4m3_native_cx(): + return raw.to(tl.float16) * 128.0 + else: + return e4m3_u8_to_f16_x128(raw) + + @jit def round_e4m3(x): """Round fp32 onto the e4m3 value grid (RNE), fp32 -> fp32, in a SINGLE @@ -140,3 +209,57 @@ def round_e4m3(x): y_norm = ((b + 524287 + lsb) & 0xFFF00000).to(tl.float32, bitcast=True) y_sub = (x + 24576.0) - 24576.0 return tl.where(tl.abs(x) >= 0.015625, y_norm, y_sub) + + +@jit +def kv_load_e4m3_tile_f32(ptrs, mask): + """Load a tile of KV e4m3 codes and widen it to fp32. + + Straight-line on purpose: no probe, no dtype test, so there is no arm left to + prune. The pools keep their codes in a plain byte buffer on EVERY architecture + (kv_quant.kv_codes_dtype), so the fp8e4nv type never reaches Triton through here. + Both ways of choosing an arm were tried on real hardware and each broke the run: + the compile-time fp8-native answer is a second, independent verdict that can + disagree with the host that allocated the buffer, and a comparison against the + pointer's element type is NOT statically pruned -- Triton type-checks the arm that + should have been dead, and an int mask fill against an fp8 pointer is rejected + ("cannot cast int32 to fp8e4nv", raised at CUDA graph capture on sm_100). + + What remains is the decode that already runs wherever the fp8 type is unavailable, + bit-exact per this module's header: the same load-with-int-fill and software + widening as kernel/triton/ple.py and nvfp4_linear.py. Callers use this only in + their quantized branch -- the 16-bit path keeps its own tl.load, so bf16 attention + is untouched instruction for instruction -- and the dense paged kernels and the + QSA sparse one read the same pool, hence this helper lives here. + """ + return e4m3_u8_to_f32(tl.load(ptrs, mask=mask, other=0)) + + +# What kv_load_e4m3_tile_scaled16 leaves on its tile for the caller to repay. +# tl.constexpr, not a plain float: a @jit function that references a module-level +# python float fails triton's cache-key AST walk with a CompilationError, the same +# way this module's header describes for @constexpr_function and host functions. +# Host-side callers (tests) want ``KV_TILE_SCALE.value``. +KV_TILE_SCALE = tl.constexpr(256.0) + + +@jit +def kv_load_e4m3_tile_scaled16(ptrs, mask): + """Load a tile of KV e4m3 codes as fp16 holding the value times 1/KV_TILE_SCALE. + + :func:`kv_load_e4m3_tile_f32` without its last two steps: the widen to fp32 and + the ``* 256.0`` that puts the tile back on the true e4m3 scale exist only so the + result is directly usable, and they are what force the tile to 32 bits. A caller + that must apply a per-(token, kv_head) dequant scale anyway can fold ``2**8`` into + that scale instead -- exactly, since it is a power of two -- and keep the tile + 16-bit. Callers owe that fold; see :data:`KV_TILE_SCALE`. + + Same straight-line load-with-int-fill as that function, for the same reason (see + its docstring), and the same bit placement, folded: for ``v = 128s + r`` it builds + ``32768s + 128r`` with two masks, two widens, two shifts and an or, while + ``(v + (v & 0x80)) << 7 = (256s + r) << 7`` is the same number in one widen, one + mask, one add and one shift. Verified identical on all 256 codes, NaN patterns + (0x7F/0xFF -> +-480 after the caller's fold) included. + """ + w = tl.load(ptrs, mask=mask, other=0).to(tl.uint16) + return ((w + (w & 0x80)) << 7).to(tl.float16, bitcast=True) diff --git a/python/freetoken/kernel/triton/glm_dsa_sparse.py b/python/freetoken/kernel/triton/glm_dsa_sparse.py index 8aaeeb9ba9..bac2cd7eed 100644 --- a/python/freetoken/kernel/triton/glm_dsa_sparse.py +++ b/python/freetoken/kernel/triton/glm_dsa_sparse.py @@ -28,6 +28,9 @@ import triton import triton.language as tl +from freetoken.kernel.triton.e4m3_compat import kv_load_e4m3_tile_f32 +from freetoken.kernel.triton.kv_nvfp4 import load_nvfp4 + BLOCK_H = 16 BLOCK_T = 32 MAX_SPLITS = 32 @@ -36,11 +39,12 @@ @triton.jit def _glm_dsa_sparse_kernel( - q_ptr, pool_ptr, o_ptr, idx_ptr, cnt_ptr, + q_ptr, pool_ptr, pool_scale_ptr, pool_block_ptr, o_ptr, idx_ptr, cnt_ptr, scale, H, TOPK, stride_qb, stride_qm, stride_qh, stride_qd, stride_pn, stride_pd, + stride_ps, stride_ob, stride_om, stride_oh, stride_od, stride_ib, stride_im, stride_it, stride_nb, stride_nm, @@ -50,6 +54,8 @@ def _glm_dsa_sparse_kernel( BLOCK_T: tl.constexpr, HAS_COUNTS: tl.constexpr, HAS_ROPE: tl.constexpr, + HAS_FP8: tl.constexpr, + HAS_NVFP4: tl.constexpr, ): pid_m = tl.program_id(0) pid_b = tl.program_id(1) @@ -61,11 +67,15 @@ def _glm_dsa_sparse_kernel( q_base = q_ptr + pid_b * stride_qb + pid_m * stride_qm + offs_h[:, None] * stride_qh q_v = tl.load(q_base + offs_v[None, :] * stride_qd, mask=h_mask[:, None], other=0.0).to(tl.float32) + if HAS_NVFP4: + q_v = q_v.to(q_ptr.dtype.element_ty) if HAS_ROPE: # NoPE checkpoints (glm5_next) have D_R == 0: tl.arange needs a non-empty # span, so the whole rope half is compiled out on the constexpr. offs_r = tl.arange(0, D_R) q_r = tl.load(q_base + (D_V + offs_r[None, :]) * stride_qd, mask=h_mask[:, None], other=0.0).to(tl.float32) + if HAS_NVFP4: + q_r = q_r.to(q_ptr.dtype.element_ty) m_i = tl.full((BLOCK_H,), -float("inf"), dtype=tl.float32) l_i = tl.zeros((BLOCK_H,), dtype=tl.float32) @@ -82,11 +92,34 @@ def _glm_dsa_sparse_kernel( idxs = tl.load(idx_base + offs_t * stride_it, mask=t_mask, other=-1) valid = idxs >= 0 kv_base = pool_ptr + idxs[:, None] * stride_pn - kv_v = tl.load(kv_base + offs_v[None, :] * stride_pd, mask=valid[:, None], other=0.0).to(tl.float32) + if HAS_NVFP4: + kv_v = load_nvfp4( + pool_ptr, pool_block_ptr, pool_scale_ptr, + idxs[:, None], 0, offs_v[None, :], valid[:, None], + stride_pn, 0, stride_ps, D_V + D_R, + ).to(q_ptr.dtype.element_ty) + elif HAS_FP8: + row_scale = tl.load(pool_scale_ptr + idxs * stride_ps, mask=valid, other=0.0) + kv_v = kv_load_e4m3_tile_f32( + kv_base + offs_v[None, :] * stride_pd, valid[:, None] + ) * row_scale[:, None] + else: + kv_v = tl.load(kv_base + offs_v[None, :] * stride_pd, mask=valid[:, None], other=0.0).to(tl.float32) scores = tl.dot(q_v, tl.trans(kv_v)) if HAS_ROPE: - kv_r = tl.load(kv_base + (D_V + offs_r[None, :]) * stride_pd, mask=valid[:, None], other=0.0).to(tl.float32) + if HAS_NVFP4: + kv_r = load_nvfp4( + pool_ptr, pool_block_ptr, pool_scale_ptr, + idxs[:, None], 0, offs_r[None, :], valid[:, None], + stride_pn, 0, stride_ps, D_V + D_R, DIM_OFFSET=D_V, + ).to(q_ptr.dtype.element_ty) + elif HAS_FP8: + kv_r = kv_load_e4m3_tile_f32( + kv_base + (D_V + offs_r[None, :]) * stride_pd, valid[:, None] + ) * row_scale[:, None] + else: + kv_r = tl.load(kv_base + (D_V + offs_r[None, :]) * stride_pd, mask=valid[:, None], other=0.0).to(tl.float32) scores += tl.dot(q_r, tl.trans(kv_r)) scores = scores * scale scores = tl.where(valid[None, :], scores, -float("inf")) @@ -99,7 +132,7 @@ def _glm_dsa_sparse_kernel( acc = acc * alpha[:, None] + tl.dot(p.to(kv_v.dtype), kv_v) m_i = m_new - o = acc / l_i[:, None] + o = tl.where(l_i[:, None] > 0, acc / l_i[:, None], 0.0) o_ptrs = o_ptr + pid_b * stride_ob + pid_m * stride_om + offs_h[:, None] * stride_oh + offs_v[None, :] * stride_od tl.store(o_ptrs, o.to(o_ptr.dtype.element_ty), mask=h_mask[:, None]) @@ -200,11 +233,12 @@ def glm_dsa_decode_logits( @triton.jit def _glm_dsa_splitk_kernel( - q_ptr, pool_ptr, mid_o_ptr, mid_lse_ptr, idx_ptr, cnt_ptr, + q_ptr, pool_ptr, pool_scale_ptr, pool_block_ptr, mid_o_ptr, mid_lse_ptr, idx_ptr, cnt_ptr, scale, H, TOPK, stride_qb, stride_qm, stride_qh, stride_qd, stride_pn, stride_pd, + stride_ps, stride_mb, stride_mm, stride_mh, stride_ms, stride_md, stride_lb, stride_lm, stride_lh, stride_ls, stride_ib, stride_im, stride_it, @@ -215,6 +249,8 @@ def _glm_dsa_splitk_kernel( BLOCK_T: tl.constexpr, HAS_COUNTS: tl.constexpr, HAS_ROPE: tl.constexpr, + HAS_FP8: tl.constexpr, + HAS_NVFP4: tl.constexpr, NUM_SPLITS: tl.constexpr, ): """Stage 1 (decode flash-decoding): each program reduces one BLOCK_T-aligned slice of @@ -245,9 +281,13 @@ def _glm_dsa_splitk_kernel( if split_end > split_start: q_base = q_ptr + pid_b * stride_qb + pid_m * stride_qm + offs_h[:, None] * stride_qh q_v = tl.load(q_base + offs_v[None, :] * stride_qd, mask=h_mask[:, None], other=0.0).to(tl.float32) + if HAS_NVFP4: + q_v = q_v.to(q_ptr.dtype.element_ty) if HAS_ROPE: offs_r = tl.arange(0, D_R) q_r = tl.load(q_base + (D_V + offs_r[None, :]) * stride_qd, mask=h_mask[:, None], other=0.0).to(tl.float32) + if HAS_NVFP4: + q_r = q_r.to(q_ptr.dtype.element_ty) idx_base = idx_ptr + pid_b * stride_ib + pid_m * stride_im for start in range(split_start, split_end, BLOCK_T): @@ -256,11 +296,34 @@ def _glm_dsa_splitk_kernel( idxs = tl.load(idx_base + offs_t * stride_it, mask=t_mask, other=-1) valid = idxs >= 0 kv_base = pool_ptr + idxs[:, None] * stride_pn - kv_v = tl.load(kv_base + offs_v[None, :] * stride_pd, mask=valid[:, None], other=0.0).to(tl.float32) + if HAS_NVFP4: + kv_v = load_nvfp4( + pool_ptr, pool_block_ptr, pool_scale_ptr, + idxs[:, None], 0, offs_v[None, :], valid[:, None], + stride_pn, 0, stride_ps, D_V + D_R, + ).to(q_ptr.dtype.element_ty) + elif HAS_FP8: + row_scale = tl.load(pool_scale_ptr + idxs * stride_ps, mask=valid, other=0.0) + kv_v = kv_load_e4m3_tile_f32( + kv_base + offs_v[None, :] * stride_pd, valid[:, None] + ) * row_scale[:, None] + else: + kv_v = tl.load(kv_base + offs_v[None, :] * stride_pd, mask=valid[:, None], other=0.0).to(tl.float32) scores = tl.dot(q_v, tl.trans(kv_v)) if HAS_ROPE: - kv_r = tl.load(kv_base + (D_V + offs_r[None, :]) * stride_pd, mask=valid[:, None], other=0.0).to(tl.float32) + if HAS_NVFP4: + kv_r = load_nvfp4( + pool_ptr, pool_block_ptr, pool_scale_ptr, + idxs[:, None], 0, offs_r[None, :], valid[:, None], + stride_pn, 0, stride_ps, D_V + D_R, DIM_OFFSET=D_V, + ).to(q_ptr.dtype.element_ty) + elif HAS_FP8: + kv_r = kv_load_e4m3_tile_f32( + kv_base + (D_V + offs_r[None, :]) * stride_pd, valid[:, None] + ) * row_scale[:, None] + else: + kv_r = tl.load(kv_base + (D_V + offs_r[None, :]) * stride_pd, mask=valid[:, None], other=0.0).to(tl.float32) scores += tl.dot(q_r, tl.trans(kv_r)) scores = scores * scale scores = tl.where(valid[None, :], scores, -float("inf")) @@ -323,7 +386,7 @@ def _glm_dsa_merge_kernel( l_i = l_i * alpha + beta m_i = m_new - o = acc / l_i + o = tl.where(l_i > 0, acc / l_i, 0.0) o_ptrs = ( o_ptr + pid_b * stride_ob + pid_m * stride_om + pid_h * stride_oh + offs_v * stride_od @@ -351,7 +414,10 @@ def glm_dsa_sparse_attn( softmax_scale: float, counts: torch.Tensor | None = None, # [b, m] int32 live columns per query (device-read) d_v: int = 512, + pool_scale: torch.Tensor | None = None, # [rows] fp32, one scale per quantized latent row force_splits: int | None = None, # tests only: 0 = single-program, N = split-k N + kv_quant: str | None = None, + pool_block_scale: torch.Tensor | None = None, ) -> torch.Tensor: """Sparse MLA attention over gathered latent rows; returns ``[b, m, h, d_v]``. @@ -359,13 +425,45 @@ def glm_dsa_sparse_attn( list across all queries with STRIDE 0 -- the identity-selection dense path: every query reads the same position-ordered row list, causally bounded by its own ``counts[q] = position + 1``, with zero per-query index materialization. + + NVFP4 restores tiles to Q's compute dtype before the dot products, with FP32 + accumulators. FP32 dot operands exceed consumer GPU shared memory at width 576. """ b, m, h, d = q.shape d_r = d - d_v topk = topk_idxs.shape[-1] - assert pool.shape[-1] == d, (pool.shape, d) + if kv_quant is None: + kv_quant = "fp8" if pool_scale is not None else "none" + assert kv_quant in ("none", "fp8", "nvfp4"), kv_quant + has_fp8 = kv_quant == "fp8" + has_nvfp4 = kv_quant == "nvfp4" + stored_dim = d // 2 if has_nvfp4 else d + assert pool.ndim == 2 and pool.shape[-1] == stored_dim, (pool.shape, d) + if has_fp8 or has_nvfp4: + assert pool_scale is not None + assert pool.dtype == torch.uint8, pool.dtype + assert pool_scale.shape == (pool.shape[0],), (pool_scale.shape, pool.shape) + assert pool_scale.dtype is torch.float32, pool_scale.dtype + scale_pool = pool_scale.contiguous() + else: + assert pool_scale is None + assert pool.is_floating_point() + scale_pool = pool + if has_nvfp4: + assert d % 16 == 0 + assert pool_block_scale is not None + assert pool_block_scale.shape == (pool.shape[0], d // 16) + assert pool_block_scale.dtype == torch.uint8 + assert pool_block_scale.device == pool.device + block_pool = pool_block_scale.contiguous() + else: + assert pool_block_scale is None + block_pool = pool + assert pool.device == q.device + if has_fp8 or has_nvfp4: + assert pool_scale.device == pool.device q = q.contiguous() - pool_2d = pool.reshape(-1, d) + pool_2d = pool.reshape(-1, stored_dim) assert pool_2d.stride(-1) == 1 idx = topk_idxs.contiguous().to(torch.int32) broadcast_m = idx.shape[1] == 1 and m > 1 @@ -379,25 +477,30 @@ def glm_dsa_sparse_attn( else: cnt, stride_nb, stride_nm = idx, 0, 0 + # Packed gathers need additional layout conversions at the 512-wide latent size. + block_t = 16 if has_nvfp4 else BLOCK_T n_splits = _split_count(b, m, h, topk, q.device) if force_splits is None else force_splits if n_splits: mid_o = q.new_empty(b, m, h, n_splits, d_v, dtype=torch.float32) mid_lse = q.new_empty(b, m, h, n_splits, dtype=torch.float32) grid1 = (m * n_splits, b, triton.cdiv(h, BLOCK_H)) _glm_dsa_splitk_kernel[grid1]( - q, pool_2d, mid_o, mid_lse, idx, cnt, + q, pool_2d, scale_pool, block_pool, mid_o, mid_lse, idx, cnt, float(softmax_scale), h, topk, q.stride(0), q.stride(1), q.stride(2), q.stride(3), pool_2d.stride(0), pool_2d.stride(1), + scale_pool.stride(0) if has_fp8 or has_nvfp4 else 0, mid_o.stride(0), mid_o.stride(1), mid_o.stride(2), mid_o.stride(3), mid_o.stride(4), mid_lse.stride(0), mid_lse.stride(1), mid_lse.stride(2), mid_lse.stride(3), idx.stride(0), 0 if broadcast_m else idx.stride(1), idx.stride(2), stride_nb, stride_nm, D_V=d_v, D_R=d_r, - BLOCK_H=BLOCK_H, BLOCK_T=BLOCK_T, - HAS_COUNTS=has_counts, HAS_ROPE=d_r > 0, NUM_SPLITS=n_splits, - num_warps=4, num_stages=2, + BLOCK_H=BLOCK_H, BLOCK_T=block_t, + HAS_COUNTS=has_counts, HAS_ROPE=d_r > 0, HAS_FP8=has_fp8, HAS_NVFP4=has_nvfp4, NUM_SPLITS=n_splits, + # Consumer Blackwell caps kernel shared memory at 99 KiB. Single-stage, the float path + # needs 100 KiB; two stages (its setting before quantized KV) fit. + num_warps=4, num_stages=1 if has_fp8 or has_nvfp4 else 2, ) grid2 = (m, b, h) _glm_dsa_merge_kernel[grid2]( @@ -412,18 +515,20 @@ def glm_dsa_sparse_attn( grid = (m, b, triton.cdiv(h, BLOCK_H)) _glm_dsa_sparse_kernel[grid]( - q, pool_2d, o, idx, cnt, + q, pool_2d, scale_pool, block_pool, o, idx, cnt, float(softmax_scale), h, topk, q.stride(0), q.stride(1), q.stride(2), q.stride(3), pool_2d.stride(0), pool_2d.stride(1), + scale_pool.stride(0) if has_fp8 or has_nvfp4 else 0, o.stride(0), o.stride(1), o.stride(2), o.stride(3), idx.stride(0), 0 if broadcast_m else idx.stride(1), idx.stride(2), stride_nb, stride_nm, D_V=d_v, D_R=d_r, - BLOCK_H=BLOCK_H, BLOCK_T=BLOCK_T, - HAS_COUNTS=has_counts, HAS_ROPE=d_r > 0, - num_warps=4, num_stages=2, + BLOCK_H=BLOCK_H, BLOCK_T=block_t, + HAS_COUNTS=has_counts, HAS_ROPE=d_r > 0, HAS_FP8=has_fp8, HAS_NVFP4=has_nvfp4, + # same shared-memory split as the split-k launch above + num_warps=4, num_stages=1 if has_fp8 or has_nvfp4 else 2, ) return o diff --git a/python/freetoken/kernel/triton/hc.py b/python/freetoken/kernel/triton/hc.py index 45d71bd731..acae534c45 100644 --- a/python/freetoken/kernel/triton/hc.py +++ b/python/freetoken/kernel/triton/hc.py @@ -326,9 +326,11 @@ def _hc_combine_norm_kernel( tl.store(out_ptr + row * stride_out + offs, out, mask=mask_inner) out = out.to(tl.float32) - # Keep the two-axis reduction: flattening the padded tile is ~40% slower - # at decode sizes. - sum_sq = tl.sum(tl.sum(out * out, axis=1), axis=0) + # The standalone grouped RMSNorm reduces one padded [next_pow2(HC_DIM)] + # vector; reshape keeps the identical tree (same element + # order, same masked zeros), so the fused statistic is bit-exact vs the split. + flat = tl.reshape(out * out, [NUM_TILES_PAD * BLOCK_SIZE]) + sum_sq = tl.sum(flat) rrms = tl.rsqrt(sum_sq / HC_DIM + EPS) if launch_pdl: diff --git a/python/freetoken/kernel/triton/kv_nvfp4.py b/python/freetoken/kernel/triton/kv_nvfp4.py new file mode 100644 index 0000000000..04727deec9 --- /dev/null +++ b/python/freetoken/kernel/triton/kv_nvfp4.py @@ -0,0 +1,154 @@ +"""Row-scaled NVFP4 KV storage, with low nibble first and 16-wide blocks. + +Each (token, head) row has an FP32 scale; each block has an E4M3 scale +stored as uint8. Reconstruction is E2M1 * block_scale * row_scale. +Row-local second-level scales keep appends independent of the cached prefix. +""" + +import torch +import triton +import triton.language as tl + +from .e4m3_compat import e4m3_f32_to_u8, e4m3_u8_to_f32, round_e4m3 + + +@triton.jit +def _encode_e2m1(x): + a = tl.abs(x) + code = tl.full(x.shape, 0, tl.int32) + code = tl.where(a > 0.25, 1, code) + code = tl.where(a >= 0.75, 2, code) + code = tl.where(a > 1.25, 3, code) + code = tl.where(a >= 1.75, 4, code) + code = tl.where(a > 2.5, 5, code) + code = tl.where(a >= 3.5, 6, code) + code = tl.where(a > 5.0, 7, code) + return code | tl.where(x < 0, 8, 0) + + +@triton.jit +def _decode_e2m1(code): + # Place E2M1 bits in FP16, then compensate the exponent-bias difference (15 - 1). + bits = ((code & 8).to(tl.uint16) << 12) | ((code & 7).to(tl.uint16) << 9) + return bits.to(tl.float16, bitcast=True).to(tl.float32) * 16384.0 + + +@triton.jit +def load_nvfp4(ptr, block_ptr, row_ptr, slots, head, dims, slot_mask, + stride_slot, stride_head, stride_row, D: tl.constexpr, + DIM_OFFSET: tl.constexpr = 0): + TRANSPOSE: tl.constexpr = dims.shape[0] != 1 + WIDTH: tl.constexpr = dims.shape[0] if TRANSPOSE else dims.shape[1] + TOKENS: tl.constexpr = slots.shape[0] * slots.shape[1] + slot = slots.reshape(TOKENS).to(tl.int64) + valid = slot_mask.reshape(TOKENS) + dim = DIM_OFFSET + tl.arange(0, WIDTH) + packed = tl.load( + ptr + slot[:, None] * stride_slot + head * stride_head + (dim[None, :] // 2), + valid[:, None] & (dim[None, :] < D), other=0, + ).to(tl.int32) + codes = tl.where((dim[None, :] & 1) == 0, packed & 15, packed >> 4) + block_dim = dim // 16 + block = tl.load( + block_ptr + (slot[:, None] * stride_row + head) * (D // 16) + block_dim[None, :], + valid[:, None] & (dim[None, :] < D), other=0, + ) + row = tl.load(row_ptr + slot * stride_row + head, valid, other=0) + value = _decode_e2m1(codes) * e4m3_u8_to_f32(block) * row[:, None] + if TRANSPOSE: + return tl.trans(value) + else: + return value + + +@triton.jit +def _quantize_row(src, dst, block_ptr, row_ptr, t, h, slot, stride_src, + HEADS: tl.constexpr, D: tl.constexpr, BLOCKS: tl.constexpr): + blocks = tl.arange(0, BLOCKS) + dims = blocks[:, None] * 16 + tl.arange(0, 16)[None, :] + x = tl.load(src + t * stride_src + h * D + dims, dims < D, other=0).to(tl.float32) + amax = tl.max(tl.abs(x), 1) + row_scale = tl.maximum(tl.max(amax, 0), 1e-10) / (6.0 * 448.0) + block_scale = round_e4m3(tl.minimum(tl.div_rn(amax, 6.0 * row_scale), 448.0)) + # Quantize against the scale actually stored, including E4M3 rounding/underflow. + denom = tl.where(block_scale > 0, block_scale * row_scale, 1.0) + normalized = tl.where(block_scale[:, None] > 0, tl.div_rn(x, denom[:, None]), 0.0) + codes = _encode_e2m1(normalized).reshape(BLOCKS, 8, 2) + lo, hi = tl.split(codes) + packed = lo | (hi << 4) + byte_dims = blocks[:, None] * 8 + tl.arange(0, 8)[None, :] + tl.store(dst + (slot * HEADS + h) * (D // 2) + byte_dims, packed, byte_dims < D // 2) + tl.store(block_ptr + (slot * HEADS + h) * (D // 16) + blocks, + e4m3_f32_to_u8(block_scale), blocks < D // 16) + tl.store(row_ptr + slot * HEADS + h, row_scale) + + +@triton.jit +def _scatter_rows(src, dst, block, row, indices, stride_src, + D: tl.constexpr, BLOCKS: tl.constexpr): + t = tl.program_id(0) + slot = tl.load(indices + t).to(tl.int64) + _quantize_row(src, dst, block, row, t, 0, slot, stride_src, 1, D, BLOCKS) + + +def quantize_nvfp4_rows_to_cache(rows, out_loc, cache, scales, block_scales) -> None: + """Quantize a single MLA latent slab, with one second-level scale per token.""" + tokens, dim = rows.shape + slots = cache.shape[0] + assert dim % 16 == 0 and rows.stride(1) == 1 + assert rows.dtype in (torch.float16, torch.bfloat16, torch.float32) + assert cache.shape == (slots, dim // 2) and cache.dtype == torch.uint8 + assert scales.shape == (slots,) and scales.dtype == torch.float32 + assert block_scales.shape == (slots, dim // 16) and block_scales.dtype == torch.uint8 + assert out_loc.shape == (tokens,) and out_loc.dtype in (torch.int32, torch.int64) + for tensor in (cache, scales, block_scales, out_loc): + assert tensor.is_contiguous() and tensor.device == rows.device + assert rows.is_cuda + if tokens: + _scatter_rows[(tokens,)]( + rows, cache, block_scales, scales, out_loc, rows.stride(0), + D=dim, BLOCKS=triton.next_power_of_2(dim // 16), + num_warps=4, enable_fp_fusion=False, + ) + + +@triton.jit +def _scatter(k, v, kc, vc, kb, vb, kr, vr, indices, stride_k, stride_v, + HEADS: tl.constexpr, D: tl.constexpr, BLOCKS: tl.constexpr): + t, h = tl.program_id(0), tl.program_id(1) + slot = tl.load(indices + t).to(tl.int64) + _quantize_row(k, kc, kb, kr, t, h, slot, stride_k, HEADS, D, BLOCKS) + _quantize_row(v, vc, vb, vr, t, h, slot, stride_v, HEADS, D, BLOCKS) + + +def quantize_nvfp4_to_cache( + k: torch.Tensor, + v: torch.Tensor, + out_loc: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + k_scale: torch.Tensor, + v_scale: torch.Tensor, + k_block_scale: torch.Tensor, + v_block_scale: torch.Tensor, +) -> None: + tokens, width = k.shape + slots, heads, packed_dim = k_cache.shape + dim = packed_dim * 2 + assert dim % 16 == 0 and width == heads * dim + assert v.shape == k.shape and k.stride(1) == v.stride(1) == 1 + assert k.dtype in (torch.float16, torch.bfloat16, torch.float32) and v.dtype == k.dtype + assert out_loc.shape == (tokens,) and out_loc.dtype in (torch.int32, torch.int64) + assert out_loc.is_contiguous() + for codes, row, block in ((k_cache, k_scale, k_block_scale), (v_cache, v_scale, v_block_scale)): + assert codes.shape == (slots, heads, packed_dim) and codes.dtype == torch.uint8 + assert row.shape == (slots, heads) and row.dtype == torch.float32 + assert block.shape == (slots, heads, dim // 16) and block.dtype == torch.uint8 + assert codes.is_contiguous() and row.is_contiguous() and block.is_contiguous() + assert codes.device == row.device == block.device == k.device + assert k.is_cuda and v.device == out_loc.device == k.device + if tokens: + _scatter[(tokens, heads)](k, v, k_cache, v_cache, k_block_scale, v_block_scale, + k_scale, v_scale, out_loc, k.stride(0), v.stride(0), + HEADS=heads, D=dim, BLOCKS=triton.next_power_of_2(dim // 16), + num_warps=4, enable_fp_fusion=False) diff --git a/python/freetoken/kernel/triton/kv_quant.py b/python/freetoken/kernel/triton/kv_quant.py new file mode 100644 index 0000000000..c10ad9c4b6 --- /dev/null +++ b/python/freetoken/kernel/triton/kv_quant.py @@ -0,0 +1,249 @@ +"""FP8 (e4m3) KV-cache storage: per-token, per-head symmetric quantization. + +One KV row is one ``(token, kv_head)`` slice of ``head_dim`` elements. It is stored +as ``head_dim`` e4m3 bytes plus ONE fp32 scale shared by the whole row: + + scale = max(amax(row) / 448, eps) # 448 == e4m3 finite max + code = round_e4m3(clamp(row / scale)) # RNE onto the e4m3 grid + read = code.to(f32) * scale # in the attention kernels + +Granularity rationale: an fp32 scale per (token, head) costs ``4 / head_dim`` bytes +per element (3% at head_dim 128, 6% at 64) while tracking each key's own magnitude, +which is what keeps a quantized KV from collapsing on outlier heads. A coarser +per-tensor scale needs no storage at all but has no headroom for them; a finer +per-element "scale" is the format itself. + +Architectures below sm_89 have no fp8e4nv type in Triton (see +:mod:`freetoken.kernel.triton.e4m3_compat`), so the codes live in a plain ``uint8`` +buffer on EVERY architecture and are decoded by :func:`e4m3_u8_to_f32`. That keeps one +set of bytes and one set of numbers across GPUs, and it is why :func:`kv_codes_dtype` +is a constant rather than a question: the fp8 type never appears in a kernel +signature, so nothing here can disagree with the host that allocated the buffer. +Choosing the encode/decode per target -- by an arch probe, or by testing the pointer's +element type -- is what broke this feature twice on real hardware (see +:func:`kv_load_e4m3_tile_f32`). + +The write path replaces ``kernel.store_cache`` for a quantized pool: the plain store +kernel is a raw byte copy that requires the source and the cache to share a dtype, +and quantization is exactly the step where the two diverge. Folding the scatter into +the quantization kernel keeps that to a single launch (and a single HBM round trip) +under CUDA-graph capture, where the slot ids arrive as a device tensor. +""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + +from freetoken.kernel.triton.e4m3_compat import ( + e4m3_f32_to_u8, + round_e4m3, +) + +FP8 = torch.float8_e4m3fn +KV_SCALE_DTYPE = torch.float32 +KV_QUANT_FP8 = "fp8" + + +def kv_codes_dtype() -> torch.dtype: + """Storage dtype of one quantized KV element: e4m3 bytes in a uint8 buffer. + + A constant, deliberately. Every attempt to answer this per target -- triton's + compile-time probe, then testing the pointer's element type at the load -- ended up + picking an arm that did not match the buffer the host had just allocated (see the + module header). torch still reads these bytes as fp8 whenever real numbers are + wanted: :func:`codes_to_f32`. + """ + return torch.uint8 + + +def alloc_codes(shape: tuple[int, ...], device: torch.device) -> torch.Tensor: + """A zero-filled code buffer of :func:`kv_codes_dtype` -- bytes, on every arch. + + Zero-filling matters because the pools read slots that were never written: a stale + byte decodes to a real number (0x7F/0xFF even to NaN once reinterpreted as fp8), + while 0x00 is exactly 0.0 (same reasoning as kvcache/bsa_pool.py). + """ + return torch.zeros(shape, dtype=kv_codes_dtype(), device=device) + + +def codes_to_f32(codes: torch.Tensor) -> torch.Tensor: + """Decode a code buffer to fp32 ON THE HOST (torch's own e4m3 cast). + + Works on either storage dtype -- uint8 bytes are reinterpreted as fp8 first -- so + a test or a debugging tool reads the same numbers on sm_86 as on sm_90, and does + it through torch rather than through the software decoder it is checking. + """ + if codes.dtype is not FP8: + codes = codes.view(FP8) + return codes.to(torch.float32) + + +@triton.jit +def _kv_quant_scatter_kernel( + k_src, + v_src, + k_dst, + v_dst, + k_scale, + v_scale, + idx_ptr, + stride_xs, # K source row pitch, in elements (the qkv slice is wider than one row) + stride_vx, # V source row pitch. K and V need not share one: a .clamp() on one side + # leaves it densely packed, so reusing K's pitch reads V off its rows. + stride_kd, # K cache row pitch, in elements (== HEADS * D) + stride_vd, + stride_ks, # scale row pitch, in elements (== HEADS) + stride_vs, + D: tl.constexpr, + BLOCK_D: tl.constexpr, +): + """One program per (token, kv_head): quantize the row and write it to slot + ``idx_ptr[token]`` (the same ``out_loc`` the bf16 store scatters through).""" + t = tl.program_id(0) + h = tl.program_id(1) + pos = tl.load(idx_ptr + t).to(tl.int64) # int64: slots * row can pass 2**31 + d = tl.arange(0, BLOCK_D) + mask = d < D + + off = h * D + d + xk = tl.load(k_src + t * stride_xs + off, mask=mask, other=0.0).to(tl.float32) + xv = tl.load(v_src + t * stride_vx + off, mask=mask, other=0.0).to(tl.float32) + + # 448 == e4m3 finite max; 1e-10 is the amax floor of the activation quant in + # kernel/triton/fp8_block_linear.py (literals keep the kernel self-contained). + sk = tl.maximum(tl.max(tl.abs(xk), axis=0), 1e-10) / 448.0 + sv = tl.maximum(tl.max(tl.abs(xv), axis=0), 1e-10) / 448.0 + qk = tl.clamp(xk / sk, -448.0, 448.0) + qv = tl.clamp(xv / sv, -448.0, 448.0) + + # Straight-line, like the reader: round onto the e4m3 grid in ONE step (RNE) and + # pack the bits into the byte buffer. No fp8 type on either side -- the reason this + # feature stopped compiling twice is explained in kv_codes_dtype's docstring. + out_k = e4m3_f32_to_u8(round_e4m3(qk)) + out_v = e4m3_f32_to_u8(round_e4m3(qv)) + + tl.store(k_dst + pos * stride_kd + h * D + d, out_k, mask=mask) + tl.store(v_dst + pos * stride_vd + h * D + d, out_v, mask=mask) + tl.store(k_scale + pos * stride_ks + h, sk) + tl.store(v_scale + pos * stride_vs + h, sv) + + +def quantize_kv_to_cache( + k: torch.Tensor, + v: torch.Tensor, + out_loc: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + k_scale: torch.Tensor, + v_scale: torch.Tensor, +) -> None: + """Quantize fresh K/V rows into an fp8 KV pool. + + ``k``/``v`` : ``[T, num_kv_heads * head_dim]`` compute-dtype rows -- exactly what + the attention backends hand to ``store_kv`` (a slice of the qkv projection, + so the row pitch may be wider than the row itself). + ``out_loc`` : ``[T]`` device slot index per row (int32 or int64). + ``*_cache`` : ``[num_slots, num_kv_heads, head_dim]`` of :func:`kv_codes_dtype`. + ``*_scale`` : ``[num_slots, num_kv_heads]`` fp32, indexed by the SAME slot. + """ + tokens = k.shape[0] + if tokens == 0: + return + assert k.dim() == 2 and v.shape == k.shape, (k.shape, v.shape) + assert k.stride(1) == 1 and v.stride(1) == 1, "K/V rows must be contiguous" + heads, dim = k_cache.shape[1], k_cache.shape[2] + assert k.shape[1] == heads * dim, (tuple(k.shape), tuple(k_cache.shape)) + assert k_cache.shape == v_cache.shape, (tuple(k_cache.shape), tuple(v_cache.shape)) + assert k_scale.shape == v_scale.shape == (k_cache.shape[0], heads), ( + tuple(k_scale.shape), + tuple(k_cache.shape), + ) + assert k_cache.dtype == v_cache.dtype == kv_codes_dtype(), ( + k_cache.dtype, + kv_codes_dtype(), + ) + assert k_scale.dtype == KV_SCALE_DTYPE, k_scale.dtype + _kv_quant_scatter_kernel[(tokens, heads)]( + k, + v, + k_cache, + v_cache, + k_scale, + v_scale, + out_loc, + k.stride(0), + v.stride(0), + k_cache.stride(0), + v_cache.stride(0), + k_scale.stride(0), + v_scale.stride(0), + D=dim, + BLOCK_D=triton.next_power_of_2(dim), + num_warps=1, + ) + + +@triton.jit +def _kv_quant_rows_scatter_kernel( + src, dst, scale_ptr, idx_ptr, + stride_src, stride_dst, stride_scale, + D: tl.constexpr, + BLOCK_D: tl.constexpr, +): + """Quantize one whole latent row per program and scatter it by ``out_loc``.""" + t = tl.program_id(0) + pos = tl.load(idx_ptr + t).to(tl.int64) + d = tl.arange(0, BLOCK_D) + mask = d < D + x = tl.load(src + t * stride_src + d, mask=mask, other=0.0).to(tl.float32) + s = tl.maximum(tl.max(tl.abs(x), axis=0), 1e-10) / 448.0 + q = tl.clamp(x / s, -448.0, 448.0) + tl.store(dst + pos * stride_dst + d, e4m3_f32_to_u8(round_e4m3(q)), mask=mask) + tl.store(scale_ptr + pos * stride_scale, s) + + +def quantize_rows_to_cache( + rows: torch.Tensor, + out_loc: torch.Tensor, + cache: torch.Tensor, + scales: torch.Tensor, +) -> None: + """Quantize and scatter one scale-bearing FP8 row per token. + + MLA stores its absorbed K/V state as one latent row, rather than separate K and + V heads. Its scale granularity is therefore one FP32 value per ``(token, + layer)`` row, which is the single-head MLA case already priced by + ``kv_scale_bytes_per_token``. + """ + tokens = rows.shape[0] + if tokens == 0: + return + assert rows.dim() == 2 and rows.stride(1) == 1, tuple(rows.shape) + assert cache.dim() == 2 and cache.shape[1] == rows.shape[1], ( + tuple(cache.shape), tuple(rows.shape), + ) + assert scales.shape == (cache.shape[0],), (tuple(scales.shape), tuple(cache.shape)) + assert cache.dtype == kv_codes_dtype(), cache.dtype + assert scales.dtype == KV_SCALE_DTYPE, scales.dtype + dim = rows.shape[1] + _kv_quant_rows_scatter_kernel[(tokens,)]( + rows, cache, scales, out_loc, + rows.stride(0), cache.stride(0), scales.stride(0), + D=dim, BLOCK_D=triton.next_power_of_2(dim), + num_warps=4 if dim > 256 else 1, + ) + + +__all__ = [ + "FP8", + "KV_QUANT_FP8", + "KV_SCALE_DTYPE", + "alloc_codes", + "codes_to_f32", + "kv_codes_dtype", + "quantize_kv_to_cache", + "quantize_rows_to_cache", +] + diff --git a/python/freetoken/kernel/triton/moe/__init__.py b/python/freetoken/kernel/triton/moe/__init__.py new file mode 100644 index 0000000000..3f620fec6e --- /dev/null +++ b/python/freetoken/kernel/triton/moe/__init__.py @@ -0,0 +1,5 @@ +"""Triton kernels for the MoE offload cache (``moe/offload_cache.py``).""" + +from .invalidate import invalidate_prefill_slots + +__all__ = ["invalidate_prefill_slots"] diff --git a/python/freetoken/kernel/triton/moe/invalidate.py b/python/freetoken/kernel/triton/moe/invalidate.py new file mode 100644 index 0000000000..24c5fbf3f9 --- /dev/null +++ b/python/freetoken/kernel/triton/moe/invalidate.py @@ -0,0 +1,75 @@ +"""Single-launch invalidation of the MoE prefill overlap buffers. + +``OffloadMoeCache._invalidate_prefill_buffer`` runs twice per prefill chunk (once per +double-buffer slot) and used to clear the buffer's slot map with a boolean-mask index: + + slot_for_id.view(-1)[old_ids[old_ids >= 0].long()] = -1 + +A boolean index produces a data-dependent shape, so every call hid a device-to-host +synchronization. With two buffer reuses per chunk x 48 layers, the host waited -- per +layer -- for ALL enqueued GPU work (the cached-context attention included) before +prefetching the next layer, which serialized the prefill-overlap pipeline. On long +cached contexts that turned a ~3 s chunk into a 20-150 s turn (py-spy: 92% of the +scheduler inside this invalidation + ``synchronize``, GPU 0-3% utilized). + +This kernel is fixed-shape: one launch, no synchronization, identical result. Zeroing +``usage`` here (instead of a separate ``zero_``) keeps the buffer's slots as the +argmin(usage) eviction victims, exactly like the code it replaces. +""" + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _invalidate_slots_kernel( + id_of_slot_ptr, # (num_slots,) int32 -- slot -> expert id, -1 if empty + slot_for_id_ptr, # (num_layers * num_experts) int32 -- expert id -> slot, -1 if not resident + usage_ptr, # (num_slots,) int64 -- last-used step; zeroed so these slots are + # the argmin(usage) victims in ensure_experts' eviction + slot_start, + n, + BLOCK: tl.constexpr, +): + i = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + m = i < n + e = tl.load(id_of_slot_ptr + slot_start + i, mask=m, other=-1) + # clear the expert's residency flag; empty slots (-1) are masked out + tl.store(slot_for_id_ptr + e, -1, mask=m & (e >= 0)) + tl.store(id_of_slot_ptr + slot_start + i, -1, mask=m) + tl.store(usage_ptr + slot_start + i, 0, mask=m) + + +def invalidate_prefill_slots( + id_of_slot: torch.Tensor, + slot_for_id: torch.Tensor, + usage: torch.Tensor, + slot_start: int, + num_experts: int, +) -> None: + """Free the ``num_experts`` slots starting at ``slot_start`` (fixed shape, no sync). + + ``slot_for_id`` may be the 2-D ``(num_layers, num_experts)`` map: ids are flat + (``layer * num_experts + expert``) and the tensor is contiguous, so the kernel's + flat addressing matches the ``view(-1)`` the eager path used. + """ + # The raw-pointer kernel cannot clamp like the eager slices did; refuse a buffer + # that does not fit instead of reading past the cache. + if slot_start + num_experts > id_of_slot.numel(): + raise ValueError( + f"prefill buffer [{slot_start}, {slot_start + num_experts}) exceeds " + f"id_of_slot size {id_of_slot.numel()}" + ) + # The kernel needs a CUDA context; CPU test caches keep the old eager path, where + # the hidden sync is harmless (no enqueued GPU work to drain). + if id_of_slot.device.type != "cuda": + old_ids = id_of_slot[slot_start : slot_start + num_experts] + slot_for_id.view(-1)[old_ids[old_ids >= 0].long()] = -1 + old_ids.fill_(-1) + usage[slot_start : slot_start + num_experts].zero_() + return + grid = ((num_experts + 255) // 256,) + _invalidate_slots_kernel[grid]( + id_of_slot, slot_for_id, usage, slot_start, num_experts, BLOCK=256 + ) diff --git a/python/freetoken/kernel/triton/nvfp4_fused_moe.py b/python/freetoken/kernel/triton/nvfp4_fused_moe.py index b1de443787..f3625f065c 100644 --- a/python/freetoken/kernel/triton/nvfp4_fused_moe.py +++ b/python/freetoken/kernel/triton/nvfp4_fused_moe.py @@ -19,6 +19,10 @@ MiniMax-M2 decode toward the RTX 5090 read-bandwidth ceiling (~87%); the residual gap is FP4 dequant ALU, which a swizzled-layout tensor-core path (marlin / flashinfer b12x) closes but those need sm_80-99 / CUDA>=13 respectively. + +Prefill (:func:`_prefill_nvfp4_moe_kernel`) is tensor-core bound, so it dequantizes with +``nvfp4_linear.py``'s fp16 bit placement rather than the table, and loads its activations as +one contiguous tile, split into the even and odd k the two nibbles of a byte multiply. """ from __future__ import annotations @@ -29,8 +33,8 @@ import triton import triton.language as tl -from freetoken.kernel.triton.e4m3_compat import e4m3_native_cx, e4m3_u8_to_f32 -from freetoken.kernel.triton.nvfp4_linear import _nvfp4_pair_f32 +from freetoken.kernel.triton.e4m3_compat import e4m3_native_cx, e4m3_to_f16_x128, e4m3_u8_to_f32 +from freetoken.kernel.triton.nvfp4_linear import _nvfp4_pair_f16, _nvfp4_pair_f32 _E2M1_VALUES = [ 0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, @@ -45,6 +49,13 @@ def _e2m1_lut(device_index: int) -> torch.Tensor: ) +@triton.jit +def _e2m1_byte_f16_x2pow_neg14(packed): + """The two e2m1 codes of each byte (int32, low nibble first) -> (lo, hi) fp16 = value * 2^-14: + ``_nvfp4_pair_f16`` once the high nibble also sits at bits [19:16].""" + return _nvfp4_pair_f16(packed | (packed << 12)) + + @triton.jit def _decode_nvfp4_moe_kernel( a_ptr, # [M, K] activations (compute dtype) @@ -303,7 +314,6 @@ def _prefill_nvfp4_moe_kernel( sorted_token_ids_ptr, expert_ids_ptr, # cache slot per M-block num_tokens_post_padded_ptr, - lut_ptr, N, K, EM, @@ -321,6 +331,7 @@ def _prefill_nvfp4_moe_kernel( MUL_ROUTED_WEIGHT: tl.constexpr, top_k: tl.constexpr, compute_type: tl.constexpr, + EVEN_K: tl.constexpr, ): pid = tl.program_id(axis=0) num_pid_m = tl.cdiv(EM, BLOCK_SIZE_M) @@ -342,8 +353,8 @@ def _prefill_nvfp4_moe_kernel( offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N offs_kb = tl.arange(0, BLOCK_SIZE_KB) - a_ptrs_lo = a_ptr + (offs_token[:, None] // top_k * stride_am + (2 * offs_kb)[None, :] * stride_ak) - a_ptrs_hi = a_ptr + (offs_token[:, None] // top_k * stride_am + (2 * offs_kb + 1)[None, :] * stride_ak) + offs_k = tl.arange(0, 2 * BLOCK_SIZE_KB) + a_ptrs = a_ptr + (offs_token[:, None] // top_k * stride_am + offs_k[None, :] * stride_ak) slot = tl.load(expert_ids_ptr + pid_m).to(tl.int64) packed_base = packed_ptr + slot * stride_pe + offs_bn[None, :] * stride_pn @@ -353,30 +364,32 @@ def _prefill_nvfp4_moe_kernel( K_BYTES = K // 2 for kb in range(0, tl.cdiv(K_BYTES, BLOCK_SIZE_KB)): byte_idx = kb * BLOCK_SIZE_KB + offs_kb - byte_mask = byte_idx < K_BYTES - p_ptrs = packed_base + byte_idx[:, None] * stride_pkb - bytes_ = tl.load(p_ptrs, mask=byte_mask[:, None], other=0).to(tl.int32) - lo = bytes_ & 0xF - hi = (bytes_ >> 4) & 0xF - sblk = byte_idx // 8 - s_ptrs = scale_base + sblk[:, None] * stride_sblk - if e4m3_native_cx(): - scale = tl.load(s_ptrs, mask=byte_mask[:, None], other=0.0).to(tl.float32) + s_ptrs = scale_base + (byte_idx // 8)[:, None] * stride_sblk + if EVEN_K: + bytes_ = tl.load(p_ptrs).to(tl.int32) + raw_scale = tl.load(s_ptrs) else: - scale = e4m3_u8_to_f32(tl.load(s_ptrs, mask=byte_mask[:, None], other=0)) - b_lo = tl.load(lut_ptr + lo) * scale # [BLOCK_KB, BLOCK_N] - b_hi = tl.load(lut_ptr + hi) * scale - - a_lo = tl.load(a_ptrs_lo, mask=token_mask[:, None] & byte_mask[None, :], other=0.0) - a_hi = tl.load(a_ptrs_hi, mask=token_mask[:, None] & byte_mask[None, :], other=0.0) - accumulator += tl.dot(a_lo, b_lo.to(a_lo.dtype)) - accumulator += tl.dot(a_hi, b_hi.to(a_hi.dtype)) - - a_ptrs_lo += BLOCK_SIZE_KB * 2 * stride_ak - a_ptrs_hi += BLOCK_SIZE_KB * 2 * stride_ak - - g = tl.load(global_ptr + slot * stride_ge + offs_bn * stride_gn).to(tl.float32) + byte_mask = byte_idx < K_BYTES + bytes_ = tl.load(p_ptrs, mask=byte_mask[:, None], other=0).to(tl.int32) + raw_scale = tl.load(s_ptrs, mask=byte_mask[:, None], other=0.0) + # [BLOCK_KB, BLOCK_N]; every e2m1 * e4m3 product is exact in fp16 and in the dot dtype + scale = e4m3_to_f16_x128(raw_scale) + lo, hi = _e2m1_byte_f16_x2pow_neg14(bytes_) + b_lo = (lo * scale).to(a_ptr.dtype.element_ty) + b_hi = (hi * scale).to(a_ptr.dtype.element_ty) + if EVEN_K: + a = tl.load(a_ptrs, mask=token_mask[:, None], other=0.0) + else: + k_mask = kb * 2 * BLOCK_SIZE_KB + offs_k < K + a = tl.load(a_ptrs, mask=token_mask[:, None] & k_mask[None, :], other=0.0) + a_lo, a_hi = tl.split(tl.reshape(a, (BLOCK_SIZE_M, BLOCK_SIZE_KB, 2))) + accumulator += tl.dot(a_lo, b_lo) + accumulator += tl.dot(a_hi, b_hi) + a_ptrs += BLOCK_SIZE_KB * 2 * stride_ak + + # the placement's 2^-14 and the scale's 2^7 leave the products 2^-7 short + g = tl.load(global_ptr + slot * stride_ge + offs_bn * stride_gn).to(tl.float32) * 128.0 accumulator = accumulator * g[None, :] if MUL_ROUTED_WEIGHT: diff --git a/python/freetoken/kernel/triton/nvfp4_linear.py b/python/freetoken/kernel/triton/nvfp4_linear.py index 53abdedc71..943a6e50e8 100644 --- a/python/freetoken/kernel/triton/nvfp4_linear.py +++ b/python/freetoken/kernel/triton/nvfp4_linear.py @@ -43,7 +43,7 @@ from freetoken.kernel.triton.e4m3_compat import ( e4m3_kernel_view, e4m3_native_cx, - e4m3_u8_to_f16_x128, + e4m3_to_f16_x128, e4m3_u8_to_f32, ) @@ -426,19 +426,13 @@ def _nvfp4_gemm_kernel( for t in range(tiles_per): if EVEN_K: word = tl.load(b_ptrs, mask=n_mask[:, None], other=0) - if e4m3_native_cx(): - s128 = tl.load(s_ptrs, mask=n_mask[:, None], other=0.0).to(tl.float16) * 128.0 - else: - s128 = e4m3_u8_to_f16_x128(tl.load(s_ptrs, mask=n_mask[:, None], other=0)) + s128 = e4m3_to_f16_x128(tl.load(s_ptrs, mask=n_mask[:, None], other=0.0)) a_tile = tl.load(a_ptrs, mask=m_mask[:, None], other=0.0) else: kw_ok = (tile0 + t) * BLOCK_KW + offs_kw < K_WORDS word = tl.load(b_ptrs, mask=n_mask[:, None] & kw_ok[None, :], other=0) s_mask = n_mask[:, None] & kw_ok[None, :] - if e4m3_native_cx(): - s128 = tl.load(s_ptrs, mask=s_mask, other=0.0).to(tl.float16) * 128.0 - else: - s128 = e4m3_u8_to_f16_x128(tl.load(s_ptrs, mask=s_mask, other=0)) + s128 = e4m3_to_f16_x128(tl.load(s_ptrs, mask=s_mask, other=0.0)) ak_ok = (tile0 + t) * BLOCK_KW * 8 + offs_ak < K a_tile = tl.load(a_ptrs, mask=m_mask[:, None] & ak_ok[None, :], other=0.0) diff --git a/python/freetoken/kernel/triton/ple.py b/python/freetoken/kernel/triton/ple.py index 2314e94e72..466fce11a0 100644 --- a/python/freetoken/kernel/triton/ple.py +++ b/python/freetoken/kernel/triton/ple.py @@ -21,6 +21,7 @@ # Latency-bound over PCIe, so keep the block small and let many of them be in flight. _NUM_WARPS = 1 +_TL_TABLE_DTYPES = {torch.bfloat16: tl.bfloat16, torch.float16: tl.float16, torch.float32: tl.float32} @triton.jit @@ -33,8 +34,10 @@ def _ple_gather_kernel( EMB_DIM: tl.constexpr, IS_FP8: tl.constexpr, BLOCK_D: tl.constexpr, + TABLE_DTYPE: tl.constexpr, ): - row = tl.program_id(0) + # the output offset row * EMB_DIM can pass 2^31 for a long prefill at a wide hidden size + row = tl.program_id(0).to(tl.int64) idx = tl.load(ids_ptr + row).to(tl.int64) in_range = (idx >= 0) & (idx < num_rows) idx = tl.where(in_range, idx, 0) @@ -50,7 +53,7 @@ def _ple_gather_kernel( base = table_ptr.to(tl.int64).to(tl.pointer_type(tl.uint8)) values = e4m3_u8_to_f32(tl.load(base + idx * EMB_DIM + offsets, mask=mask, other=0)) else: - base = table_ptr.to(tl.int64).to(tl.pointer_type(tl.bfloat16)) + base = table_ptr.to(tl.int64).to(tl.pointer_type(TABLE_DTYPE)) values = tl.load(base + idx * EMB_DIM + offsets, mask=mask, other=0.0).to(tl.float32) values = tl.where(in_range, values * scale, 0.0) tl.store( @@ -68,11 +71,12 @@ def ple_gather_rows( out: torch.Tensor, scale: float = 1.0, is_fp8: bool = True, + table_dtype: torch.dtype = torch.bfloat16, ) -> torch.Tensor: """Gather ``row_ids`` from the host-resident table at ``table_ptr`` into ``out``. ``row_ids`` is a flat device int tensor; ``out`` is ``[row_ids.numel(), embed_dim]`` - bf16 on the same device. ``table_ptr`` is the address the GPU must dereference + on the same device. A non-fp8 table holds ``table_dtype`` elements. ``table_ptr`` is the address the GPU must dereference (``kernel/pinned.device_ptr``), not necessarily the host ``data_ptr``. """ n = row_ids.numel() @@ -87,6 +91,7 @@ def ple_gather_rows( EMB_DIM=embed_dim, IS_FP8=is_fp8, BLOCK_D=triton.next_power_of_2(embed_dim), + TABLE_DTYPE=_TL_TABLE_DTYPES[table_dtype], num_warps=_NUM_WARPS, ) return out diff --git a/python/freetoken/kernel/triton/ple_hash.py b/python/freetoken/kernel/triton/ple_hash.py new file mode 100644 index 0000000000..c8c5e6302d --- /dev/null +++ b/python/freetoken/kernel/triton/ple_hash.py @@ -0,0 +1,145 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Fused n-gram hash -> PLE table row ids. + +The eager form of this hash (``NGramEmbedding._window`` + ``_shift_ignore_eos`` + the per-ngram +XOR/multiply/remainder/offset loop) is 39 tiny CUDA kernels per PLE layer per step -- ~400 us of +launch wall for a few us of GPU work. This is the same arithmetic as one kernel, one program per +token; ``NGramEmbedding.row_ids_reference`` keeps the torch-op form as the oracle and CPU path. + +The key observation that collapses the ``cummax``-over-the-whole-window boundary scan into a +``ngram_size-1`` step walk: ``_shift_ignore_eos`` marks shift ``s`` valid at position ``p`` iff +``p-s >= 0`` and no boundary token sits anywhere in ``[p-s, p-1]``. Only shifts ``< ngram_size`` +are ever used, so the scan never needs to look further back than that, and the predicate is +built incrementally as the walk goes. + +Window addressing avoids materializing the ``[B, ctx+max_len]`` packed window entirely. Token +``t`` of the forward belongs to request ``req[t]`` at intra-request offset ``local[t]``, so the +token ``s`` places to its left is ``input_ids[t - s]`` when ``local[t] >= s`` and +``ngram_context[req[t], ctx_len + local[t] - s]`` otherwise -- out of range on the left is the +boundary token, exactly as the eos-filled packed window was. + +Capture-safe: fixed shapes, every input on device, no host reads. +""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _ple_row_ids_kernel( + ids_ptr, # [T] int64 -- this forward's tokens, ragged, in request order + ctx_ptr, # [B, CTX_LEN] int64 -- the tokens immediately before each request's first + req_ptr, # [T] int32 -- request index of each token + local_ptr, # [T] int32 -- intra-request offset of each token + mult_ptr, # [NGRAM] int64 + vocab_ptr, # [NUM_HEADS] int64 + off_ptr, # [NUM_HEADS] int64 + out_ptr, # [T, NUM_HEADS] int64 + EOS: tl.constexpr, + CTX_LEN: tl.constexpr, + NGRAM: tl.constexpr, + HEADS_PER: tl.constexpr, + NUM_HEADS: tl.constexpr, + BLOCK_H: tl.constexpr, +): + token = tl.program_id(0).to(tl.int64) + req = tl.load(req_ptr + token).to(tl.int64) + local = tl.load(local_ptr + token).to(tl.int64) + + head = tl.arange(0, BLOCK_H) + head_ok = head < NUM_HEADS + + # shift 0 is always the token itself; the eos-crossing rule never masks it + mixed = tl.load(ids_ptr + token).to(tl.int64) * tl.load(mult_ptr).to(tl.int64) + acc = tl.zeros([BLOCK_H], dtype=tl.int64) + + valid = 1 + for shift in tl.static_range(1, NGRAM): + column = CTX_LEN + local - shift # column in the (virtual) packed window + from_ids = column >= CTX_LEN + from_ctx = (column >= 0) & (column < CTX_LEN) + token_ids = tl.load(ids_ptr + (token - shift), mask=from_ids, other=0) + token_ctx = tl.load(ctx_ptr + req * CTX_LEN + column, mask=from_ctx, other=EOS) + raw = tl.where(from_ids, token_ids, token_ctx).to(tl.int64) + # the window may not cross a boundary token, and a boundary token is itself the wall + valid = valid * tl.where((column >= 0) & (raw != EOS), 1, 0) + mixed = mixed ^ ( + tl.where(valid == 1, raw, EOS) * tl.load(mult_ptr + shift).to(tl.int64) + ) + # after ``shift`` taps the (shift+1)-gram mix is complete; it owns one head block + ngram = shift + 1 + block = (head >= (ngram - 2) * HEADS_PER) & (head < (ngram - 1) * HEADS_PER) + acc = tl.where(block, mixed, acc) + + vocab = tl.load(vocab_ptr + head, mask=head_ok, other=1).to(tl.int64) + offset = tl.load(off_ptr + head, mask=head_ok, other=0).to(tl.int64) + # torch.remainder is floored, triton's % is truncated; the divisor is always positive + rem = acc % vocab + rem = tl.where(rem < 0, rem + vocab, rem) + tl.store(out_ptr + token * NUM_HEADS + head, rem + offset, mask=head_ok) + + +def ple_row_ids( + input_ids: torch.Tensor, + ngram_context: torch.Tensor, + req_index: torch.Tensor, + local_index: torch.Tensor, + multipliers: torch.Tensor, + vocab_sizes: torch.Tensor, + offsets: torch.Tensor, + *, + eos_token_id: int, + heads_per_ngram: int, + out: torch.Tensor | None = None, +) -> torch.Tensor: + """``[T, num_heads]`` int64 global table rows for this forward's tokens. + + ``input_ids`` [T] and ``ngram_context`` [B, ngram_size-1] are int64 device tensors; + ``req_index`` / ``local_index`` are [T] int32 device tensors naming each token's request and + its offset within that request. The three hash constant tensors are int64 and on the same + device. ``out``, when given, is the destination (a CUDA graph replays into a fixed buffer). + """ + tokens = input_ids.numel() + ngram_size = int(multipliers.numel()) + num_heads = int(vocab_sizes.numel()) + ctx_len = int(ngram_context.shape[-1]) + # Checked as raises, not asserts: the kernel addresses the context row and the head + # blocks by these, and ``python -O`` must not turn a geometry mismatch into an OOB read. + if ctx_len != ngram_size - 1: + raise ValueError( + f"PLE hash: ngram_context has {ctx_len} context ids but ngram_size {ngram_size} " + f"needs {ngram_size - 1}" + ) + if num_heads != heads_per_ngram * (ngram_size - 1): + raise ValueError( + f"PLE hash: {num_heads} heads is not heads_per_ngram {heads_per_ngram} x " + f"{ngram_size - 1} n-gram orders" + ) + if out is None: + out = torch.empty((tokens, num_heads), dtype=torch.int64, device=input_ids.device) + if tokens == 0: + return out + _ple_row_ids_kernel[(tokens,)]( + input_ids, + ngram_context, + req_index, + local_index, + multipliers, + vocab_sizes, + offsets, + out, + EOS=int(eos_token_id), + CTX_LEN=ctx_len, + NGRAM=ngram_size, + HEADS_PER=int(heads_per_ngram), + NUM_HEADS=num_heads, + BLOCK_H=triton.next_power_of_2(num_heads), + num_warps=1, + ) + return out + + +__all__ = ["ple_row_ids"] diff --git a/python/freetoken/kernel/triton/qsa/attend.py b/python/freetoken/kernel/triton/qsa/attend.py index 541e27c680..9414ca01c3 100644 --- a/python/freetoken/kernel/triton/qsa/attend.py +++ b/python/freetoken/kernel/triton/qsa/attend.py @@ -9,12 +9,19 @@ import triton import triton.language as tl +from freetoken.kernel.triton.e4m3_compat import kv_load_e4m3_tile_f32 +from freetoken.kernel.triton.kv_nvfp4 import load_nvfp4 + @triton.jit def _qsa_sparse_paged_gqa_splitk_kernel( q_ptr, k_cache_ptr, v_cache_ptr, + k_scale_ptr, + v_scale_ptr, + k_block_scale_ptr, + v_block_scale_ptr, indices_ptr, block_table_ptr, token_to_req_ptr, @@ -29,6 +36,8 @@ def _qsa_sparse_paged_gqa_splitk_kernel( stride_v_block, stride_v_token, stride_v_head, + stride_kss, + stride_vss, stride_indices_row, stride_table_req, stride_output_row, @@ -46,6 +55,11 @@ def _qsa_sparse_paged_gqa_splitk_kernel( NUM_TILES: tl.constexpr, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, + # e4m3 KV pool (kvcache/mha_pool.py): read codes + per-token row scales instead of + # 16-bit values. The bf16 branch below stays exactly as it was, instruction for + # instruction, for the unquantized default. + HAS_KV_SCALE: tl.constexpr, + KV_NVFP4: tl.constexpr, ) -> None: # row * stride can overflow int32 for large row counts. row = tl.program_id(0).to(tl.int64) @@ -101,24 +115,94 @@ def _qsa_sparse_paged_gqa_splitk_kernel( valid &= (physical_page >= 0) & (physical_page < num_cache_blocks) # physical_page * block stride can overflow int32 for large caches. safe_page = tl.maximum(physical_page, 0).to(tl.int64) - keys = tl.load( - k_cache_ptr - + safe_page[None, :] * stride_k_block - + page_offset[None, :] * stride_k_token - + kv_head * stride_k_head - + dim_offsets[:, None], - mask=valid[None, :], - other=0.0, - ) - values = tl.load( - v_cache_ptr - + safe_page[:, None] * stride_v_block - + page_offset[:, None] * stride_v_token - + kv_head * stride_v_head - + dim_offsets[None, :], - mask=valid[:, None], - other=0.0, - ) + if KV_NVFP4: + scale_slot = safe_page * PAGE_SIZE + page_offset + keys = load_nvfp4( + k_cache_ptr, + k_block_scale_ptr, + k_scale_ptr, + scale_slot[None, :], + kv_head, + dim_offsets[:, None], + valid[None, :], + stride_k_token, + stride_k_head, + stride_kss, + HEAD_DIM, + ).to(query.dtype) + values = load_nvfp4( + v_cache_ptr, + v_block_scale_ptr, + v_scale_ptr, + scale_slot[:, None], + kv_head, + dim_offsets[None, :], + valid[:, None], + stride_v_token, + stride_v_head, + stride_vss, + HEAD_DIM, + ).to(query.dtype) + elif HAS_KV_SCALE: + # The scale row is the slot the code lives in: QSA pins page_size to this + # kernel's PAGE_SIZE (attention/__init__.py registers page_sizes=(64,)), so + # slot = page * PAGE_SIZE + offset addresses k_scale/v_scale exactly. + scale_slot = safe_page * PAGE_SIZE + page_offset # safe_page is int64 + # Invalid columns mask the codes to 0.0 and the scale to 1.0, and slots that + # were never written read back 0.0 * 0.0 -- both buffers are zero-filled. + # Either way the operand stays finite, so the -inf row mask below is what + # decides such a column's fate rather than a NaN poisoning the row. + s_k = tl.load( + k_scale_ptr + scale_slot[None, :] * stride_kss + kv_head, + mask=valid[None, :], + other=1.0, + ) + s_v = tl.load( + v_scale_ptr + scale_slot[:, None] * stride_vss + kv_head, + mask=valid[:, None], + other=1.0, + ) + keys = ( + kv_load_e4m3_tile_f32( + k_cache_ptr + + safe_page[None, :] * stride_k_block + + page_offset[None, :] * stride_k_token + + kv_head * stride_k_head + + dim_offsets[:, None], + valid[None, :], + ) + * s_k + ).to(query.dtype) + values = ( + kv_load_e4m3_tile_f32( + v_cache_ptr + + safe_page[:, None] * stride_v_block + + page_offset[:, None] * stride_v_token + + kv_head * stride_v_head + + dim_offsets[None, :], + valid[:, None], + ) + * s_v + ).to(query.dtype) + else: + keys = tl.load( + k_cache_ptr + + safe_page[None, :] * stride_k_block + + page_offset[None, :] * stride_k_token + + kv_head * stride_k_head + + dim_offsets[:, None], + mask=valid[None, :], + other=0.0, + ) + values = tl.load( + v_cache_ptr + + safe_page[:, None] * stride_v_block + + page_offset[:, None] * stride_v_token + + kv_head * stride_v_head + + dim_offsets[None, :], + mask=valid[:, None], + other=0.0, + ) scores = tl.dot(query, keys) # Scaling scores avoids re-quantizing a scaled query to BF16. scores *= softmax_scale_log2 @@ -232,8 +316,13 @@ def qsa_sparse_paged_attention( block_table: torch.Tensor, token_to_req: torch.Tensor, out: torch.Tensor | None = None, + k_scale: torch.Tensor | None = None, + v_scale: torch.Tensor | None = None, + kv_quant: str | None = None, + k_block_scale: torch.Tensor | None = None, + v_block_scale: torch.Tensor | None = None, ) -> torch.Tensor: - """Run sparse GQA directly over paged BF16 K/V caches.""" + """Run sparse GQA over bf16, FP8, or packed NVFP4 K/V caches.""" if q.ndim != 3 or k_cache.ndim != 4 or v_cache.shape != k_cache.shape: raise ValueError("QSA sparse attention received invalid Q/K/V shapes") @@ -243,11 +332,52 @@ def qsa_sparse_paged_attention( raise ValueError("QSA sparse attention metadata has invalid shapes") if logical_indices.shape[1] <= 0: raise ValueError("QSA sparse attention requires a positive selection width") - if q.shape[2] != k_cache.shape[3] or q.shape[1] % k_cache.shape[2]: + kv_quant = kv_quant if kv_quant is not None else ("fp8" if k_scale is not None else "none") + if kv_quant not in ("none", "fp8", "nvfp4"): + raise ValueError(f"unknown QSA kv_quant {kv_quant!r}") + stored_dim = q.shape[2] // 2 if kv_quant == "nvfp4" else q.shape[2] + if ( + (kv_quant == "nvfp4" and q.shape[2] % 16) + or k_cache.shape[3] != stored_dim + or q.shape[1] % k_cache.shape[2] + ): raise ValueError("QSA sparse attention requires valid grouped-query heads") head_dim = q.shape[2] assert head_dim >= 16 and (head_dim & (head_dim - 1)) == 0 - assert q.dtype == k_cache.dtype == v_cache.dtype + if (k_scale is None) != (v_scale is None) or (k_block_scale is None) != (v_block_scale is None): + raise ValueError("QSA sparse attention requires both KV scale tensors") + nvfp4 = kv_quant == "nvfp4" + if nvfp4 and (k_scale is None or k_block_scale is None): + raise ValueError("QSA NVFP4 requires row and block scale tensors") + if not nvfp4 and (k_block_scale is not None or v_block_scale is not None): + raise ValueError("QSA block scales require kv_quant='nvfp4'") + if k_scale is not None: + # The pool hands out 1-byte e4m3 codes plus one fp32 row scale per + # (slot, kv_head); the kernel rebuilds that slot as + # page * PAGE_SIZE + page_offset, which is exact only because QSA pins + # page_size to PAGE_SIZE (page_sizes=(64,) in attention/__init__.py). + if k_cache.element_size() != 1 or v_cache.element_size() != 1: + raise ValueError("QSA KV scales require 1-byte e4m3 code caches") + if k_scale.dtype is not torch.float32 or v_scale.dtype is not torch.float32: + raise ValueError("QSA KV scale tensors must be float32") + want = (k_cache.shape[0] * k_cache.shape[1], k_cache.shape[2]) + if k_scale.shape != want or v_scale.shape != want: + raise ValueError(f"QSA KV scale tensors must have shape {want}") + assert k_scale.stride(1) == v_scale.stride(1) == 1 + if nvfp4: + block_want = (*want, head_dim // 16) + if ( + k_cache.dtype is not torch.uint8 + or v_cache.dtype is not torch.uint8 + or k_block_scale.dtype is not torch.uint8 + or v_block_scale.dtype is not torch.uint8 + or k_block_scale.shape != block_want + or v_block_scale.shape != block_want + ): + raise ValueError(f"QSA NVFP4 block scales must have shape {block_want}") + assert k_block_scale.stride(2) == v_block_scale.stride(2) == 1 + else: + assert q.dtype == k_cache.dtype == v_cache.dtype assert logical_indices.dtype == block_table.dtype == torch.int32 assert token_to_req.dtype == torch.int32 assert q.stride(2) == k_cache.stride(3) == v_cache.stride(3) == 1 @@ -277,6 +407,12 @@ def qsa_sparse_paged_attention( else: block_n, target_splits, partial_warps = 64, 1, 2 + if nvfp4 and block_n == 64 and head_dim >= 256: + # The wide bf16 profile has load_nvfp4 decode a BLOCK_N x HEAD_DIM fp32 tile for K and V + # across 64 threads (256 values per thread at HEAD_DIM 256): Triton 3.8 compiles that to + # 40 registers with ~1.4 KB of spills, ~10x slower and a ~550 MiB local-memory reservation + # outside the torch allocator. Narrow tiles at 4 warps do not spill and are faster on 3.6 too. + block_n, partial_warps = 16, 4 num_tiles = triton.cdiv(logical_indices.shape[1], block_n) # Avoid empty splits when the selection width is smaller than the profile. max_useful_splits = 1 << (num_tiles.bit_length() - 1) @@ -303,6 +439,12 @@ def qsa_sparse_paged_attention( q, k_cache, v_cache, + # Never dereferenced while HAS_KV_SCALE is False -- pass the caches so the + # launch stays type-valid without a second None-handling path. + k_cache if k_scale is None else k_scale, + v_cache if v_scale is None else v_scale, + k_cache if k_block_scale is None else k_block_scale, + v_cache if v_block_scale is None else v_block_scale, logical_indices, block_table, token_to_req, @@ -317,6 +459,8 @@ def qsa_sparse_paged_attention( v_cache.stride(0), v_cache.stride(1), v_cache.stride(2), + 0 if k_scale is None else k_scale.stride(0), + 0 if v_scale is None else v_scale.stride(0), logical_indices.stride(0), block_table.stride(0), out.stride(0), @@ -334,6 +478,8 @@ def qsa_sparse_paged_attention( NUM_TILES=num_tiles, BLOCK_M=block_m, BLOCK_N=block_n, + HAS_KV_SCALE=k_scale is not None, + KV_NVFP4=nvfp4, num_warps=partial_warps, num_stages=2, ) diff --git a/python/freetoken/kernel/triton/qsa/offload.py b/python/freetoken/kernel/triton/qsa/offload.py new file mode 100644 index 0000000000..12223f48b8 --- /dev/null +++ b/python/freetoken/kernel/triton/qsa/offload.py @@ -0,0 +1,244 @@ +"""Triton kernels for the QSA KV host offload (``kvcache/kv_host_offload.py``). + +The page table, radix tree and compressed-index slab all run in LOGICAL page space; the GPU +KV buffer is an LRU cache of physical slots over the pinned host mirror. These kernels are +the fixed-shape, CUDA-graph-safe glue between the two spaces: + +* :func:`write_pages` -- logical page id of every written token slot (``out_loc // page_size``). +* :func:`compact_selected_pages` -- per query row, the distinct logical pages the sparse + selection touches (``block_table[req, indices//page_size]``), first-element padded. +* :func:`translate_table` -- logical block table -> physical slot table (post-ensure). +* :func:`mirror_store` -- write-through of this forward's K/V rows into the pinned host + mirror (UVA stores), so a physical slot is always drop-clean. +""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _write_pages_kernel(out_loc, write_pages, page_size: tl.constexpr): + i = tl.program_id(0) + tl.store(write_pages + i, tl.load(out_loc + i) // page_size) + + +def write_pages(out_loc: torch.Tensor, out: torch.Tensor, page_size: int) -> None: + """``out[i] = out_loc[i] // page_size`` for i < out.numel() (out_loc may be wider).""" + n = out.numel() + if n: + _write_pages_kernel[(n,)](out_loc, out, page_size) + + +@triton.jit +def _compact_sel_kernel( + indices, # [rows, SEL] int32: selected LOGICAL token ids, -1 padded + token_to_req, # [rows] int32 + block_table, # [reqs, W] int32: logical page ids + out, # [rows, MAXP] int32: distinct logical pages per row + trunc, # [1] int32: rows whose selection overflowed MAXP (0 = none dropped) + counts, # [rows] int32: distinct pages stored per row (<= MAXP) + sel_width: tl.constexpr, + width, # W: block_table row stride + page_size: tl.constexpr, + maxp: tl.constexpr, # rows of `out` per row written (<= MAXV) + maxv: tl.constexpr, # next_pow2(maxp) + dummy_page, # logical id used when a row selects nothing (never happens in practice) +): + r = tl.program_id(0) + req = tl.load(token_to_req + r).to(tl.int64) + lanes = tl.arange(0, maxv) + seen = tl.full([maxv], -1, tl.int32) + first = tl.full((), dummy_page, tl.int32) + count = tl.zeros((), tl.int32) + ndrop = tl.zeros((), tl.int32) + for j in range(sel_width): + tok = tl.load(indices + r * sel_width + j) + col = tok // page_size + pg = tl.load(block_table + req * width + col, mask=tok >= 0, other=-1) + new = (tok >= 0) & (tl.sum((seen == pg).to(tl.int32)) == 0) + take = new & (count < maxp) + ndrop += (new & (count >= maxp)).to(tl.int32) + seen = tl.where((lanes == count) & take, pg, seen) + first = tl.where((count == 0) & take, pg, first) + count += take.to(tl.int32) + row = tl.where(lanes < count, seen, first) + tl.store(out + r * maxp + lanes, row, mask=lanes < maxp) + tl.store(counts + r, count) + tl.atomic_max(trunc, ndrop) + + +def compact_selected_pages( + indices: torch.Tensor, + token_to_req: torch.Tensor, + block_table: torch.Tensor, + out: torch.Tensor, + page_size: int, + dummy_page: int, + trunc: torch.Tensor | None = None, + counts: torch.Tensor | None = None, +) -> None: + """Per query row, the distinct logical pages its selected tokens live on. + + ``out`` is ``[rows, maxp]`` int32; rows with fewer hits are padded with the row's first + page (duplicates collapse in lru_ensure), or ``dummy_page`` when the row selected nothing. + ``trunc`` (optional [1] int32, device) accumulates how many pages were dropped per row + because they exceeded ``maxp`` -- 0 means no truncation happened. ``counts`` (optional + [rows] int32, device) receives each row's stored distinct-page count. + """ + rows, sel = indices.shape + maxp = out.shape[1] + if trunc is None: + trunc = getattr(compact_selected_pages, "_dummy", None) + if trunc is None or trunc.device != out.device: + trunc = compact_selected_pages._dummy = torch.zeros(1, dtype=torch.int32, device=out.device) + if counts is None: + c = getattr(compact_selected_pages, "_dummy_counts", None) + if c is None or c.numel() < rows or c.device != out.device: + c = compact_selected_pages._dummy_counts = torch.empty( + max(rows, 4096), dtype=torch.int32, device=out.device + ) + counts = c + if rows: + _compact_sel_kernel[(rows,)]( + indices, token_to_req, block_table, out, trunc, counts, + sel, block_table.stride(0), page_size, maxp, triton.next_power_of_2(maxp), + dummy_page, + ) + + +@triton.jit +def _translate_kernel(table, phys_of, out, total, BLOCK: tl.constexpr): + i = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + mask = i < total + pg = tl.load(table + i, mask=mask, other=0) + ph = tl.load(phys_of + pg, mask=mask, other=0) + # A -1 entry means the page was not ensured (only possible for a truncated selection, + # reported via the trunc counter): point it at slot 0 rather than indexing garbage. + tl.store(out + i, tl.where(ph < 0, 0, ph), mask=mask) + + +def translate_table(table: torch.Tensor, phys_of: torch.Tensor, out: torch.Tensor) -> None: + """``out = phys_of[table]`` elementwise over ``out.numel()`` entries (all int32).""" + n = out.numel() + if n: + _translate_kernel[(triton.cdiv(n, 1024),)](table, phys_of, out, n, 1024) + + +@triton.jit +def _translate_slots_kernel(out_loc, phys_of, out, page_size: tl.constexpr): + i = tl.program_id(0) + slot = tl.load(out_loc + i) + tl.store(out + i, tl.load(phys_of + slot // page_size) * page_size + slot % page_size) + + +def translate_slots(out_loc: torch.Tensor, phys_of: torch.Tensor, out: torch.Tensor, page_size: int) -> None: + """Logical token slot -> physical token slot (page translate, offset preserved).""" + n = out.numel() + if n: + _translate_slots_kernel[(n,)](out_loc, phys_of, out, page_size) + + +@triton.jit +def _mirror_store_kernel( + k_bytes, # [T, ROW] uint8 view of this layer's K rows + v_bytes, # [T, ROW] uint8 view of this layer's V rows + out_loc, # [T] int32 logical token slots + host_ptrs, # [2*L] int64: per dense layer the K/V host bank base addresses + layer, + row_bytes: tl.constexpr, + block: tl.constexpr, +): + t = tl.program_id(0) + slot = tl.load(out_loc + t).to(tl.int64) + offs = tl.arange(0, block) + mask = offs < row_bytes + kb = tl.load(host_ptrs + 2 * layer).to(tl.pointer_type(tl.uint8)) + vb = tl.load(host_ptrs + 2 * layer + 1).to(tl.pointer_type(tl.uint8)) + tl.store(kb + slot * row_bytes + offs, tl.load(k_bytes + t * row_bytes + offs, mask=mask), mask=mask) + tl.store(vb + slot * row_bytes + offs, tl.load(v_bytes + t * row_bytes + offs, mask=mask), mask=mask) + + +def mirror_store( + k: torch.Tensor, + v: torch.Tensor, + out_loc: torch.Tensor, + host_ptrs: torch.Tensor, + dense_layer: int, +) -> None: + """Write-through: this layer's K/V rows of every token in ``out_loc`` into the host mirror.""" + k = k.reshape(k.shape[0], -1) + v = v.reshape(v.shape[0], -1) + if not k.is_contiguous(): + k = k.contiguous() + if not v.is_contiguous(): + v = v.contiguous() + assert k.shape == v.shape and k.dtype == v.dtype + t = k.shape[0] + row_bytes = k.shape[1] * k.element_size() + if t: + _mirror_store_kernel[(t,)]( + k.view(torch.uint8), v.view(torch.uint8), out_loc, host_ptrs, dense_layer, + row_bytes, triton.next_power_of_2(row_bytes), + ) + + +@triton.jit +def _mirror_store_quant_kernel( + src_ptrs, # [4L] int64 GPU: por camada, bases pool (kcode, kscale, vcode, vscale) + host_ptrs, # [4L] int64 GPU: por camada, bases host UVA (kcode, kscale, vcode, vscale) + out_phys, # [T] int32: slots físicos (pool) de cada token + out_log, # [T] int32: slots lógicos (espelho) de cada token + layer, + code_bytes: tl.constexpr, scale_bytes: tl.constexpr, + BLOCK_C: tl.constexpr, BLOCK_S: tl.constexpr, +): + """fp8 write-through: lê códigos+scales do pool (pós-quantização) e espelha no host.""" + t = tl.program_id(0) + phys = tl.load(out_phys + t).to(tl.int64) + logi = tl.load(out_log + t).to(tl.int64) + oc = tl.arange(0, BLOCK_C) + mc = oc < code_bytes + os_ = tl.arange(0, BLOCK_S) + ms = os_ < scale_bytes + kc_h = tl.load(host_ptrs + 4 * layer).to(tl.pointer_type(tl.uint8)) + ks_h = tl.load(host_ptrs + 4 * layer + 1).to(tl.pointer_type(tl.uint8)) + vc_h = tl.load(host_ptrs + 4 * layer + 2).to(tl.pointer_type(tl.uint8)) + vs_h = tl.load(host_ptrs + 4 * layer + 3).to(tl.pointer_type(tl.uint8)) + kc_s = tl.load(src_ptrs + 4 * layer).to(tl.pointer_type(tl.uint8)) + ks_s = tl.load(src_ptrs + 4 * layer + 1).to(tl.pointer_type(tl.uint8)) + vc_s = tl.load(src_ptrs + 4 * layer + 2).to(tl.pointer_type(tl.uint8)) + vs_s = tl.load(src_ptrs + 4 * layer + 3).to(tl.pointer_type(tl.uint8)) + tl.store(kc_h + logi * code_bytes + oc, tl.load(kc_s + phys * code_bytes + oc, mask=mc), mask=mc) + tl.store(vc_h + logi * code_bytes + oc, tl.load(vc_s + phys * code_bytes + oc, mask=mc), mask=mc) + tl.store(ks_h + logi * scale_bytes + os_, tl.load(ks_s + phys * scale_bytes + os_, mask=ms), mask=ms) + tl.store(vs_h + logi * scale_bytes + os_, tl.load(vs_s + phys * scale_bytes + os_, mask=ms), mask=ms) + + +def mirror_store_quant( + kc_pool: torch.Tensor, # [P*page, H*D] e4m3 (view) + ks_pool: torch.Tensor, # [P*page, H] fp32 + vc_pool: torch.Tensor, + vs_pool: torch.Tensor, + out_phys: torch.Tensor, # [T] int32 slots físicos + out_log: torch.Tensor, # [T] int32 slots lógicos + src_ptrs: torch.Tensor, # [4L] int64 GPU + host_ptrs: torch.Tensor, # [4L] int64 GPU + dense_layer: int, +) -> None: + """Espelha códigos+scales de uma camada fp8 pro host (graph-safe: UVA stores).""" + t = out_phys.numel() + if not t: + return + code_bytes = kc_pool.shape[1] * kc_pool.element_size() + scale_bytes = ks_pool.shape[1] * ks_pool.element_size() + _mirror_store_quant_kernel[(t,)]( + src_ptrs, host_ptrs, out_phys, out_log, dense_layer, + code_bytes, scale_bytes, + triton.next_power_of_2(code_bytes), triton.next_power_of_2(scale_bytes), + ) + + +__all__ = ["write_pages", "compact_selected_pages", "translate_table", "translate_slots", "mirror_store", "mirror_store_quant"] diff --git a/python/freetoken/kernel/triton/qsa/score.py b/python/freetoken/kernel/triton/qsa/score.py index 49d7020828..7b65da8d6f 100644 --- a/python/freetoken/kernel/triton/qsa/score.py +++ b/python/freetoken/kernel/triton/qsa/score.py @@ -139,6 +139,14 @@ def qsa_mqa_paged( raise ValueError("QSA request mapping and positions must match query rows") if sequence_lengths.shape != (page_table.shape[0],): raise ValueError("QSA sequence lengths must match page-table requests") + if q.dtype not in (torch.bfloat16, torch.float16) or k_cache.dtype is not q.dtype: + # The dot below is a plain 16-bit matmul: an fp8 operand is not a slow path, it + # is a triton compile error that surfaces mid-CUDA-graph-capture. The KV cache + # may well be e4m3 codes (--kv-cache-dtype fp8) -- what must never reach here + # are those codes; the compressed index keys are their own 16-bit tier. + raise ValueError( + f"QSA scoring is 16-bit only, got query={q.dtype} keys={k_cache.dtype}" + ) score_divisor = math.sqrt(q.shape[2]) if score_scale is None else score_scale columns = logits.shape[1] if not q.shape[0] or not columns: diff --git a/python/freetoken/kvcache/__init__.py b/python/freetoken/kvcache/__init__.py index e0cf606db1..32f7195dbb 100644 --- a/python/freetoken/kvcache/__init__.py +++ b/python/freetoken/kvcache/__init__.py @@ -80,19 +80,36 @@ def resolve_pool_class(model_config: ModelConfig) -> type[BaseKVCachePool]: return MHAKVCache -def create_kv_pool(config, num_pages: int, device: torch.device, dtype: torch.dtype): +def _reject_unsupported_quant(pool: str, kv_quant: str) -> None: + """A pool family that has no fp8 store/scale-read path must say so at startup, + not silently serve a 16-bit cache the budget priced for an fp8 one.""" + if kv_quant != "none": + raise ValueError( + f"--kv-cache-dtype {kv_quant} is not implemented for the {pool} KV pool " + "(only the plain paged / hybrid-SWA pools, served by the triton attention " + "backend); use --kv-cache-dtype bf16." + ) + + +def create_kv_pool(config, num_pages: int, device: torch.device, dtype: torch.dtype, + num_index_pages: int | None = None): """Build the engine's KV pool for ``num_pages`` USABLE pages (the dummy page and every secondary tier -- window pool, index slab, state rings -- are derived here or inside - the pool). Single factory entry for all pool families, DSV4 included.""" + the pool). Single factory entry for all pool families, DSV4 included. + + ``num_index_pages`` (KV host offload only) sizes the QSA compressed-index slab in the + LOGICAL page space while the K/V slabs stay at ``num_pages`` physical pages.""" from .dsv4.v4_cost_model import _dsv4_pool_sizes from .hybrid_swa_pool import _naive_swa_num_tokens, _swa_paged_num_tokens from .dsv4.v4_pool import DSV4PagedKVCache model_config = config.model_config + kv_quant = getattr(config, "kv_quant", "none") if resolve_pool_class(model_config) is DSV4PagedKVCache: # DSV4 is driven by the generic CacheManager over the shared page table; the pool is # the only DSV4-specific piece (the swa_pool plug-in: window tier + cmp/idx/state # shadows). Sizing reads dsv4_args, never the group spec. + _reject_unsupported_quant("DSV4 paged", kv_quant) pool = DSV4PagedKVCache( sizes=_dsv4_pool_sizes(config, num_pages + 1), # +1 for dummy page args=model_config.dsv4_args, @@ -136,6 +153,8 @@ def create_kv_pool(config, num_pages: int, device: torch.device, dtype: torch.dt device=device, dtype=dtype, num_req_slots=config.max_running_req + 1, # + 1 for the dummy request row + kv_quant=kv_quant, + num_index_pages=None if num_index_pages is None else num_index_pages + 1, ) @@ -147,7 +166,21 @@ def create_kvcache_pool( device: torch.device, num_swa_tokens: int | None = None, num_req_slots: int | None = None, + kv_quant: str = "none", + num_index_pages: int | None = None, ) -> BaseKVCachePool: + if kv_quant == "nvfp4": + from freetoken.attention import AttnType + + if any( + spec.attn_type not in (AttnType.FULL, AttnType.SWA, AttnType.QSA, AttnType.MLA, AttnType.DSA) + or spec.head_dim % 16 + for spec in model_config.kv_cache_group_specs() + ): + raise ValueError( + "--kv-cache-dtype nvfp4 requires paged FULL, hybrid-SWA, QSA, or MLA/DSA groups " + "with head_dim divisible by 16" + ) if model_config.has_swa_attention: from .hybrid_swa_pool import HybridSWAKVCache @@ -159,6 +192,7 @@ def create_kvcache_pool( num_swa_tokens=num_swa_tokens, device=device, dtype=dtype, + kv_quant=kv_quant, ) from .mha_pool import MHAKVCache @@ -186,6 +220,7 @@ def create_kvcache_pool( if len(kv_specs) == 1 and kv_specs[0].attn_type == _AttnType.BSA: from .bsa_pool import BSAKVCache + _reject_unsupported_quant("block-sparse (BSA)", kv_quant) spec = kv_specs[0] assert layer_ids is None, "hybrid-linear x BSA has no pool support yet" return BSAKVCache( @@ -224,6 +259,10 @@ def create_kvcache_pool( num_req_slots=num_req_slots, layer_ids=spec.layer_ids, mrope=model_config.model_is_mrope, + # Quantizes the KV tiers only -- the compressed index slab the score kernel + # reads stays the engine dtype (kvcache/qsa_pool.py). + kv_quant=kv_quant, + num_index_pages=num_index_pages, ) if len(kv_specs) == 1 and kv_specs[0].mla: @@ -244,6 +283,7 @@ def create_kvcache_pool( index_head_dim=spec.index_head_dim, num_index_layers=spec.num_index_layers, layer_ids=layer_ids, + kv_quant=kv_quant, ) if spec.index_ratio > 1: # kpool tail rings are keyed by Req.table_idx; + 1 covers the dummy request row. @@ -263,6 +303,7 @@ def create_kvcache_pool( dtype=dtype, device=device, layer_ids=layer_ids, + kv_quant=kv_quant, ) spec = kv_specs[0] if len(kv_specs) == 1 else None @@ -275,6 +316,7 @@ def create_kvcache_pool( device=device, dtype=dtype, layer_ids=layer_ids, + kv_quant=kv_quant, ) diff --git a/python/freetoken/kvcache/base.py b/python/freetoken/kvcache/base.py index 569da1cb6c..7da2011bb3 100644 --- a/python/freetoken/kvcache/base.py +++ b/python/freetoken/kvcache/base.py @@ -9,6 +9,9 @@ logger = init_logger(__name__) +# One fp32 scale per (token, slab, layer, kv head) rides alongside fp8 KV codes. +FP8_KV_SCALE_BYTES = 4 + class CacheRebuildRejected(Exception): """A runtime cache rebuild was rejected BEFORE any destructive free (e.g. the @@ -16,23 +19,63 @@ class CacheRebuildRejected(Exception): this is recoverable, unlike a failure after the free.""" +def kv_storage_bytes_per_elem(config) -> int: + """Storage bytes of ONE cached KV element under the configured quantization. + + ``kv_quant == "fp8"`` stores e4m3 codes (1 byte) instead of the 16-bit compute + dtype; anything else is the compute dtype itself. Single source for the pool's + allocation, the budget math below, and the AOT kernel-shape table. + """ + quant = getattr(config, "kv_quant", "none") + if quant == "fp8": + return 1 + if quant != "none": + raise ValueError(f"unknown kv_quant {quant!r}") + return config.dtype.itemsize + + +def kv_scale_bytes_per_token(spec, config) -> int: + """Sidecar bytes per token: FP32 row scales plus NVFP4 E4M3 block scales. + The unquantized pool has no scales. + + Priced here rather than inside the pool so ``kv_cost`` and the pool's own + allocation can never disagree -- the same rule the 16-bit path follows.""" + if getattr(config, "kv_quant", "none") not in ("fp8", "nvfp4"): + return 0 + heads = div_even(spec.num_kv_heads, config.tp_info.size, allow_replicate=True) + scale_bytes = FP8_KV_SCALE_BYTES + if getattr(config, "kv_quant", "none") == "nvfp4": + scale_bytes += spec.head_dim // 16 + return (1 if spec.mla else 2) * spec.num_layers * heads * scale_bytes + + def spec_kv_bytes_per_token(spec, config) -> int: """One paged-KV group's bytes per token: (1|2 slabs) x head_dim x local kv heads x dtype - x layers, plus the bf16 DSA index-key slab when the spec carries indexer dims. Pure - per-spec arithmetic -- pool families compose it over THEIR OWN groups; no family - branching here. (2 bytes/elem == the torch.bfloat16 dsa_pool.DSAKVCache._alloc - hardcodes; keep the two in lockstep if the slab dtype ever changes.) + x layers, plus the fp8 scale sidecar when the cache is quantized, plus the bf16 DSA + index-key slab when the spec carries indexer dims. Pure per-spec arithmetic -- pool + families compose it over THEIR OWN groups; no family branching here. (2 bytes/elem == + the torch.bfloat16 dsa_pool.DSAKVCache._alloc hardcodes; keep the two in lockstep if the + slab dtype ever changes.) ``index_ratio`` > 1 (QSA) stores one index key per token group, not per token; that slab's ring and scratch rows are fixed-size and priced in QSAKVCache.kv_cost instead.""" + if getattr(config, "kv_quant", "none") == "nvfp4": + if spec.head_dim % 16: + raise ValueError("NVFP4 KV requires head_dim divisible by 16") + row_bytes = spec.head_dim // 2 + else: + row_bytes = spec.head_dim * kv_storage_bytes_per_elem(config) per_token = ( (1 if spec.mla else 2) # MLA latent groups store one slab (V aliases K) - * spec.head_dim + * row_bytes * div_even(spec.num_kv_heads, config.tp_info.size, allow_replicate=True) - * config.dtype.itemsize * spec.num_layers ) - return per_token + spec.index_head_dim * spec.num_index_layers * 2 // spec.index_ratio + return ( + per_token + + kv_scale_bytes_per_token(spec, config) + + spec.index_head_dim * spec.num_index_layers * 2 // spec.index_ratio + ) @dataclass(frozen=True) @@ -62,6 +105,12 @@ def window_pages(self) -> int: """Current usable window pages in the units declared by window_spec.""" return 0 + # KV storage quantization: "none" keeps the compute dtype, "fp8" keeps e4m3 codes plus + # one fp32 scale per (token, slab, layer, kv head). A pool that does not implement the + # quantized allocation/store/scale-view trio stays at "none"; create_kv_pool rejects a + # quantization the family does not implement, so nothing here is ever silently ignored. + kv_quant: str = "none" + # ---- sizing/cost classmethods: run BEFORE the pool exists (startup budget solve, # --moe-cache-auto). The engine measures memory and passes bytes in; each pool family # implements kv_cost for ITS OWN buffers only (the engine sums families, e.g. adds @@ -136,6 +185,24 @@ def validate_rebuild( f"needs {mem_GB(need)} > budget {mem_GB(budget)}; old cache kept, still serving" ) + @classmethod + def rebuild_footprint( + cls, config, num_pages: int, num_swa_pages: int | None + ) -> int: + """Whole-family GPU bytes at ``(num_pages, num_swa_pages)`` -- the pool's own + tiers priced together, same kv_cost plumbing validate_rebuild uses. The engine + compares old-vs-new footprints to order a live rebuild's resizes + shrinking-first (#643); ``num_swa_pages=None`` prices the config's current + window pin.""" + import inspect + + cost_params = inspect.signature(cls.kv_cost).parameters + cost_kwargs = ( + {"num_swa_pages": num_swa_pages} if "num_swa_pages" in cost_params else {} + ) + cache_per_page, fixed_cache_size, _, _ = cls.kv_cost(config, **cost_kwargs) + return num_pages * cache_per_page + fixed_cache_size + @abstractmethod def rebuild_from_config( self, config, num_pages: int, *, num_swa_pages: int | None = None @@ -171,13 +238,44 @@ def store_kv( layer_id: int, ) -> None: ... + def k_scale(self, index: int) -> torch.Tensor | None: + """fp32 ``[num_slots, local_kv_heads]`` scale of ``k_cache(index)``, indexed by the + same slot; None when the pool is not quantized.""" + return None + + def v_scale(self, index: int) -> torch.Tensor | None: + """fp32 ``[num_slots, local_kv_heads]`` scale of ``v_cache(index)``; see k_scale.""" + return None + + def k_block_scale(self, index: int) -> torch.Tensor | None: + """NVFP4 E4M3 scales for ``k_cache(index)``, or ``None`` for other layouts.""" + return None + + def v_block_scale(self, index: int) -> torch.Tensor | None: + """NVFP4 E4M3 scales for ``v_cache(index)``, or ``None`` for other layouts.""" + return None + @property @abstractmethod def device(self) -> torch.device: ... @property @abstractmethod - def dtype(self) -> torch.dtype: ... + def dtype(self) -> torch.dtype: + """The pool's COMPUTE dtype: the dtype of the K/V rows a backend hands to + ``store_kv``, and what backends size their scratch with. A quantized (e4m3) + pool still answers 16-bit here -- code bytes handed to a scratch buffer end up + as the rhs of a ``tl.dot`` that has no fp8 path (QSA's indexer died this way at + graph capture). See :attr:`store_dtype` for what the buffer holds.""" + ... + + @property + def store_dtype(self) -> torch.dtype: + """Element type of the KV buffer: e4m3/uint8 codes on a quantized pool, else + :attr:`dtype`. Only code that touches the buffer itself needs this; attention + backends that cannot apply the row scales are refused a quantized pool up + front (``BackendInfo.supports_fp8_kv``).""" + return self.dtype @property @abstractmethod diff --git a/python/freetoken/kvcache/cache_status.py b/python/freetoken/kvcache/cache_status.py index 5b531ce265..8c29dccab5 100644 --- a/python/freetoken/kvcache/cache_status.py +++ b/python/freetoken/kvcache/cache_status.py @@ -40,6 +40,9 @@ def _moe() -> int: banks = getattr(engine.moe_offload_cache, "bank_caches", None) if not banks: return 0 + if getattr(engine.moe_offload_cache, "_raw_banks", False): + # pools of several bank layouts share byte segments: report the mean slot size + return int(sum(t.numel() for t in banks.values()) // engine.moe_offload_cache.cache_size) # Each bank cache is (cache_size, *row_shape); one slot's bytes = row bytes summed over # the format's banks (== cache_budget.expert_bytes_per_slot on the source rows). return int(sum(t[0].numel() * t.element_size() for t in banks.values())) diff --git a/python/freetoken/kvcache/dsa_pool.py b/python/freetoken/kvcache/dsa_pool.py index d7c36fd734..f693227401 100644 --- a/python/freetoken/kvcache/dsa_pool.py +++ b/python/freetoken/kvcache/dsa_pool.py @@ -43,8 +43,12 @@ def __init__( dtype: torch.dtype, device: torch.device, layer_ids: "tuple[int, ...] | None" = None, + kv_quant: str = "none", ) -> None: self._latent_dim = latent_dim + if kv_quant == "nvfp4" and latent_dim % 16: + raise ValueError("NVFP4 KV requires latent_dim divisible by 16") + self._stored_dim = latent_dim // 2 if kv_quant == "nvfp4" else latent_dim if layer_ids is None: self._num_layers = num_layers self._layer_index: dict[int, int] | None = None @@ -54,6 +58,7 @@ def __init__( self._page_size = page_size self._dtype = dtype self._device = device + self.kv_quant = kv_quant self._alloc(num_pages) def _local_layer(self, layer_id: int) -> int: @@ -61,15 +66,29 @@ def _local_layer(self, layer_id: int) -> int: def _alloc(self, num_pages: int) -> None: self._num_pages = num_pages - self._kv_buffer = torch.empty( - (1, self._num_layers, num_pages, self._page_size, 1, self._latent_dim), - device=self._device, - dtype=self._dtype, - ) + shape = (1, self._num_layers, num_pages, self._page_size, 1, self._stored_dim) + self._block_scale_buffer = None + if self.kv_quant in ("fp8", "nvfp4"): + self._kv_buffer = torch.zeros(shape, device=self._device, dtype=torch.uint8) + self._scale_buffer = torch.zeros( + (self._num_layers, num_pages * self._page_size), + device=self._device, + dtype=torch.float32, + ) + if self.kv_quant == "nvfp4": + self._block_scale_buffer = torch.zeros( + (self._num_layers, num_pages * self._page_size, self._latent_dim // 16), + device=self._device, dtype=torch.uint8, + ) + elif self.kv_quant == "none": + self._kv_buffer = torch.empty(shape, device=self._device, dtype=self._dtype) + self._scale_buffer = None + else: + raise ValueError(f"unknown kv_quant {self.kv_quant!r}") # -- views (addressed by GLOBAL layer id; remapped when layer_ids was given) -- def k_cache(self, layer_id: int) -> torch.Tensor: - """Paged latent view ``[num_pages, page_size, latent_dim]``.""" + """Paged latent view; NVFP4 stores ``latent_dim // 2`` bytes per row.""" return self._kv_buffer[0, self._local_layer(layer_id)].view( self._num_pages, self._page_size, -1 ) @@ -79,8 +98,20 @@ def v_cache(self, layer_id: int) -> torch.Tensor: return self.k_cache(layer_id) def latent_rows(self, layer_id: int) -> torch.Tensor: - """Row-flat latent view ``[num_pages * page_size, latent_dim]``.""" - return self._kv_buffer[0, self._local_layer(layer_id)].view(-1, self._latent_dim) + """Row-flat latent view ``[num_pages * page_size, stored_dim]``.""" + return self._kv_buffer[0, self._local_layer(layer_id)].view(-1, self._stored_dim) + + def latent_scale(self, layer_id: int) -> torch.Tensor | None: + """One FP32 scale per latent row, or ``None`` for the compute-dtype pool.""" + if self._scale_buffer is None: + return None + return self._scale_buffer[self._local_layer(layer_id)] + + def latent_block_scale(self, layer_id: int) -> torch.Tensor | None: + """E4M3 bytes per 16 latent elements, present only for NVFP4.""" + if self._block_scale_buffer is None: + return None + return self._block_scale_buffer[self._local_layer(layer_id)] # -- writes ----------------------------------------------------------------- def store_kv( @@ -97,7 +128,24 @@ def store_kv( store.cu (two-width store). """ rows = self.latent_rows(layer_id) - split = rows.shape[1] - k_rope.shape[-1] + split = self._latent_dim - k_rope.shape[-1] + assert c_kv.shape == (out_loc.numel(), split) + assert k_rope.shape[0] == c_kv.shape[0] + if self.kv_quant == "nvfp4": + from freetoken.kernel.triton.kv_nvfp4 import quantize_nvfp4_rows_to_cache + + latent = torch.cat((c_kv, k_rope), dim=-1) if k_rope.shape[-1] else c_kv + quantize_nvfp4_rows_to_cache( + latent, out_loc, rows, self.latent_scale(layer_id), + self.latent_block_scale(layer_id), + ) + return + if self.kv_quant == "fp8": + from freetoken.kernel.triton.kv_quant import quantize_rows_to_cache + + latent = torch.cat((c_kv, k_rope), dim=-1).contiguous() + quantize_rows_to_cache(latent, out_loc, rows, self.latent_scale(layer_id)) + return rows[out_loc, :split] = c_kv rows[out_loc, split:] = k_rope @@ -105,6 +153,8 @@ def rebuild(self, num_pages: int) -> None: """In-place resize (frees the old slab first; object identity preserved -- callers re-derive views per forward, same contract as MHAKVCache.rebuild).""" self._kv_buffer = None + self._scale_buffer = None + self._block_scale_buffer = None if self._device.type == "cuda": torch.cuda.synchronize(self._device) torch.cuda.empty_cache() @@ -127,7 +177,13 @@ def rebuild_from_config( def unit_bytes(self) -> tuple[int, int]: buf = self._kv_buffer - return int(buf.numel() * buf.element_size()) // (self._num_pages * self._page_size), 0 + tokens = self._num_pages * self._page_size + kv = int(buf.numel() * buf.element_size()) // tokens + if self._scale_buffer is not None: + kv += int(self._scale_buffer.numel() * self._scale_buffer.element_size()) // tokens + if self._block_scale_buffer is not None: + kv += self._block_scale_buffer.numel() // tokens + return kv, 0 # -- pool properties ---------------------------------------------------------- @property @@ -138,6 +194,10 @@ def device(self) -> torch.device: def dtype(self) -> torch.dtype: return self._dtype + @property + def store_dtype(self) -> torch.dtype: + return self._kv_buffer.dtype + @property def num_layers(self) -> int: return self._num_layers @@ -158,12 +218,13 @@ def __init__( index_head_dim: int, num_index_layers: int, layer_ids: "tuple[int, ...] | None" = None, + kv_quant: str = "none", ) -> None: self._index_head_dim = index_head_dim self._num_index_layers = num_index_layers super().__init__( latent_dim, num_layers, num_pages, page_size, dtype, device, - layer_ids=layer_ids, + layer_ids=layer_ids, kv_quant=kv_quant, ) def _index_rows(self, num_pages: int) -> int: @@ -231,6 +292,21 @@ def _index_rows(self, num_pages: int) -> int: # 1/ratio shadow of every token slot + one scratch row per request slot. return num_pages * self._page_size // self._index_ratio + self._num_req_slots + @classmethod + def kv_cost(cls, config) -> tuple[int, int, int, int]: + per_page, fixed, page_size, reserve = super().kv_cost(config) + for spec in config.model_config.kv_cache_group_specs(): + if spec.mla and spec.index_ratio > 1: + row_bytes = spec.index_head_dim * spec.num_index_layers * 2 + fixed += (config.max_running_req + 1) * row_bytes * (2 * spec.index_ratio + 1) + return per_page, fixed, page_size, reserve + + def unit_bytes(self) -> tuple[int, int]: + # Scratch rows and the two tail rings are fixed costs, not token capacity. + kv, swa = MLAKVCache.unit_bytes(self) + index_bytes = self._num_index_layers * self._index_head_dim * 2 // self._index_ratio + return kv + index_bytes, swa + @property def cmp_scratch_base(self) -> int: """First scratch row (== shadow row count); request ``table_idx`` offsets it.""" diff --git a/python/freetoken/kvcache/dsv4/v4_pool.py b/python/freetoken/kvcache/dsv4/v4_pool.py index d5eae6d080..484dcbbb29 100644 --- a/python/freetoken/kvcache/dsv4/v4_pool.py +++ b/python/freetoken/kvcache/dsv4/v4_pool.py @@ -279,6 +279,16 @@ def validate_rebuild( f"{mem_GB(need)} > budget {mem_GB(budget)}; old cache kept, still serving" ) + @classmethod + def rebuild_footprint(cls, config, num_pages: int, num_swa_pages: int | None) -> int: + # The DSV4 kv_cost does not take num_swa_pages, so price the family with the same + # byte model validate_rebuild uses (num_pages is USABLE; +1 for the dummy page). + from .v4_cost_model import _dsv4_pool_sizes, dsv4_pool_bytes + + dsv4_args = config.model_config.dsv4_args + sizes = _dsv4_pool_sizes(config, num_pages + 1, num_swa_pages=num_swa_pages) + return dsv4_pool_bytes(sizes, dsv4_args, config.max_running_req + 1) + def rebuild_from_config( self, config, num_pages: int, *, num_swa_pages: int | None = None ) -> None: diff --git a/python/freetoken/kvcache/host_prefix_tier.py b/python/freetoken/kvcache/host_prefix_tier.py new file mode 100644 index 0000000000..4295aa8e4b --- /dev/null +++ b/python/freetoken/kvcache/host_prefix_tier.py @@ -0,0 +1,342 @@ +"""Host-side cold-prefix tier (Nível 3 — Degrau 1: snapshots GDN; Degrau 3: KV de prefixo). + +Dois artefatos em RAM pinned: +1. um banco de snapshots GDN (write na evicção do tree, read no match/admissão); +2. um índice de prefixos frios: token ids + slots KV host + slot GDN host. + +Degrau 3: o KV vive num **slot space host separado** (banco pinned [2, L, S, page, H, D] de +códigos + scales, mesma geometria do pool GPU). Ids lógicos de página são reutilizados pelo +scheduler, então o conteúdo é COPIADO (D2H) na hora do evict do tree — nunca referenciado por +id lógico depois. A reidratação aloca páginas novas no pool e copia H2D de volta. + +Ver IMPLEMENTACAO-NIVEL3.md para o design completo. +""" + +from __future__ import annotations + +from collections import OrderedDict +from dataclasses import dataclass + +import torch + +from freetoken.moe.host_banks import HostBank + + +class HostMatch: + __slots__ = ("host_len", "kv_pages", "gdn_slot") + + def __init__(self, host_len: int, kv_pages: torch.Tensor | None, gdn_slot: int | None) -> None: + self.host_len = host_len # tokens cobertos (page-aligned) + self.kv_pages = kv_pages # int32 CPU: slots KV do tier (não ids lógicos!) + self.gdn_slot = gdn_slot # slot host do snapshot GDN (None = sem snapshot) + + +@dataclass +class _Entry: + key_tokens: torch.Tensor # CPU, o prefixo (page-aligned) + kv_slots: torch.Tensor | None = None # int32 CPU: um slot host por página + gdn_slot: int | None = None # snapshot GDN host do boundary + + +class HostPrefixTier: + """Índice + storage RAM pra prefixos evictados do radix tree GPU. + + O índice é LRU por entrada (hit re-marca); quando a RAM do tier enche (KV) ou os slots + GDN acabam, drop-oldest libera a entrada mais fria inteira (KV slots + slot GDN). + As páginas KV são copiadas para slots host próprios na evicção — o conteúdo sobrevive à + reutilização dos ids lógicos pelo scheduler. + + Geometria GDN: os tensores do LinearStatePool são [layers, slots, ...]; cada banco host + é [host_slots, layers, ...] — a cópia de um slot é [layers, ...] <-> [layers, ...]. + Geometria KV: pool._kv_buffer [2, L, P, page, H, D] (códigos e4m3 no pool fp8) + + pool._scale_buffer [2, L, P*page, H] fp32; os banks host espelham isso com S host slots. + """ + + def __init__(self, state_pool, kv_pool=None, *, gdn_slots_host: int = 64, + kv_budget_bytes: int = 0, device: torch.device | str = "cuda") -> None: + self.device = torch.device(device) + self.pool = state_pool + tensors = {"conv": state_pool.conv_states, "rec": state_pool.recurrent_states} + tensors.update(state_pool.slot_states) + self._order = list(tensors) + # banks only: the pool tensors are re-read on every copy, since + # LinearStatePool.rebuild replaces them (a runtime cache rebuild) + self._banks: dict[str, HostBank] = {} + for name, t in tensors.items(): + # pool: [layers, slots, *rest] -> bank: [host_slots, layers, *rest] + bank = HostBank((gdn_slots_host, t.shape[0], *t.shape[2:]), t.dtype) + if self.device.type == "cuda" and torch.cuda.is_available(): + bank.pin() + self._banks[name] = bank + self.num_slots = gdn_slots_host + self._free = list(range(gdn_slots_host)) + self._in_use: set[int] = set() + self._index: OrderedDict[bytes, _Entry] = OrderedDict() + # key: bytes dos token ids page-aligned -> _Entry + + # ---- KV page store (Degrau 3) ---- + self.kv_pool = kv_pool + self.kv_page_bytes = 0 + self._kv_codes: HostBank | None = None + self._kv_scales: HostBank | None = None + self._kv_free: list[int] = [] + if kv_pool is not None and kv_budget_bytes > 0: + from freetoken.kvcache.kv_host_offload import _reject_block_scaled_pool + + _reject_block_scaled_pool(kv_pool, "FT_PREFIX_HOST") + if getattr(kv_pool, "_cmp_k_buffer", None) is not None: + # a QSA pool's indexer slab is keyed by the same pages and is not mirrored here + raise ValueError( + "FT_PREFIX_HOST cannot restore a sparse-attention pool's compressed index " + "slab; keep FT_GDN_HOST_TIER alone for this model" + ) + codes = kv_pool._kv_buffer # [2, L, P, page, H, D] + scales = getattr(kv_pool, "_scale_buffer", None) # [2, L, P*page, H] | None (fp8) + _, layers, _, page_size, heads, dim = codes.shape + page_bytes = 2 * layers * page_size * heads * dim * codes.element_size() + if scales is not None: + page_bytes += 2 * layers * page_size * heads * scales.element_size() + n = max(1, kv_budget_bytes // page_bytes) + self._kv_codes = HostBank((2, layers, n, page_size, heads, dim), codes.dtype) + if scales is not None: + self._kv_scales = HostBank((2, layers, n, page_size, heads), scales.dtype) + if self.device.type == "cuda" and torch.cuda.is_available(): + self._kv_codes.pin() + if self._kv_scales is not None: + self._kv_scales.pin() + self._kv_free = list(range(n)) + self.kv_page_bytes = page_bytes + self.copy_stream = ( + torch.cuda.Stream(self.device) if self.device.type == "cuda" else None) + + # ---------------- GDN snapshot storage ---------------- + def _pairs(self): + """(live pool tensor [layers, slots, ...], host bank) per state tensor.""" + pool = self.pool + live = {"conv": pool.conv_states, "rec": pool.recurrent_states, **pool.slot_states} + return [(live[name], self._banks[name]) for name in self._order] + + def _copy(self, copies) -> None: + if self.copy_stream is None: + for dst, src in copies: + dst.copy_(src) + return + # Order the copies after the caller's stream: the scheduler stream (itself made to wait + # on the engine stream) or the engine stream. Overlap launches the next batch before + # the previous one is drained, so a just-freed slot may still be written by it. + self.copy_stream.wait_stream(torch.cuda.current_stream(self.device)) + with torch.cuda.stream(self.copy_stream): + for dst, src in copies: + dst.copy_(src, non_blocking=True) + self.copy_stream.synchronize() + + def write_gdn(self, gpu_slot: int) -> int: + """D2H do snapshot do gpu_slot pra um slot host livre. Retorna o host_slot. + Deve ser enfileirado ANTES de state_pool.free(gpu_slot).""" + if not self._free: + raise RuntimeError("HostPrefixTier: sem slots host GDN livres (chame evict_host)") + hs = self._free.pop() + self._in_use.add(hs) + self._copy([(bank.tensor[hs], t[:, gpu_slot]) for t, bank in self._pairs()]) + return hs + + def read_gdn(self, host_slot: int, dst_gpu_slot: int) -> None: + """H2D host_slot -> dst_gpu_slot do pool (já alocado pelo caller).""" + if host_slot not in self._in_use: + raise KeyError(f"HostPrefixTier: slot host {host_slot} não está em uso") + self._copy([(t[:, dst_gpu_slot], bank.tensor[host_slot]) for t, bank in self._pairs()]) + + def free_gdn(self, host_slot: int) -> None: + self._in_use.discard(host_slot) + self._free.append(host_slot) + + # ---------------- KV page storage (Degrau 3) ---------------- + @property + def has_kv(self) -> bool: + return self._kv_codes is not None + + def write_kv(self, kv_indices: torch.Tensor) -> torch.Tensor | None: + """Copia as páginas (token slots do pool, page-aligned) para slots host livres. + Retorna os slots host (int32 CPU), ou None se o tier não tem KV store. + Deve rodar DURANTE o evict do tree, com os bytes ainda válidos no buffer GPU.""" + if not self.has_kv: + return None + page_size = self._kv_codes.tensor.shape[3] + pages = (kv_indices[::page_size] // page_size).long().to(self.device) + n = pages.numel() + if n == 0: + return None + codes, scales = self.kv_pool._kv_buffer, self.kv_pool._scale_buffer + # Defesa: um slot fora do pool indicaria um bug de lifetime no tree. Pular a escrita + # (o prefixo simplesmente não vai pro tier) em vez de derrubar o scheduler. + max_page = int(pages.max().item()) + if max_page >= codes.shape[2]: + from freetoken.utils import init_logger + init_logger(__name__).warning( + f"HostPrefixTier.write_kv: page {max_page} fora do pool " + f"({codes.shape[2]} páginas): n={n}, kv_slots min/max = " + f"{int(kv_indices.min())}/{int(kv_indices.max())}, prefixo de {kv_indices.numel()} tokens -- " + "escrita ignorada" + ) + return None + while len(self._kv_free) < n: + self._drop_oldest() # libera os slots da entrada mais fria (inclui GDN) + slots = torch.tensor([self._kv_free.pop() for _ in range(n)], dtype=torch.int64) + # NB: sem staging do prefixo inteiro — cada fatia calcula seus token-ids. + + def _go(): + # Chunked: NUNCA estagiar o prefixo inteiro (o evict acontece sob pressão de + # memória — um staging GB-sized explode a VRAM e derruba o scheduler). + CH = 128 # páginas por fatia (~100 MB de temp por fatia em fp8) + for i in range(0, n, CH): + p = pages[i : i + CH] + s = slots[i : i + CH] + staged = codes[:, :, p].contiguous() + self._kv_codes.tensor[:, :, s] = staged.cpu() if staged.is_cuda else staged + del staged + if scales is not None: + for i in range(0, n, CH): + p = pages[i : i + CH] + s = slots[i : i + CH] + tk = (p.unsqueeze(1) * page_size + + torch.arange(page_size, device=self.device)).flatten() + staged_s = scales[:, :, tk].view( + 2, scales.shape[1], p.numel(), page_size, scales.shape[3]).contiguous() + self._kv_scales.tensor[:, :, s] = ( + staged_s.cpu() if staged_s.is_cuda else staged_s) + del staged_s + + if self.copy_stream is not None: + self.copy_stream.wait_stream(torch.cuda.current_stream(self.device)) + with torch.cuda.stream(self.copy_stream): + _go() + self.copy_stream.synchronize() + else: + _go() + return slots.to(torch.int32).cpu() + + def read_kv(self, host_slots: torch.Tensor, dst_pages: torch.Tensor) -> None: + """H2D dos slots host para as páginas dst (ids de página do pool, device).""" + assert self.has_kv + slots = host_slots.long() + pages = dst_pages.long().to(self.device) + codes, scales = self.kv_pool._kv_buffer, self.kv_pool._scale_buffer + page_size = self._kv_codes.tensor.shape[3] + n = slots.numel() + + def _go(): + CH = 128 + for i in range(0, n, CH): + s = slots[i : i + CH] + p = pages[i : i + CH] + staged = self._kv_codes.tensor[:, :, s].to(self.device, non_blocking=True) + codes[:, :, p] = staged + del staged + if scales is not None: + for i in range(0, n, CH): + s = slots[i : i + CH] + p = pages[i : i + CH] + src = self._kv_scales.tensor[:, :, s].reshape( + 2, scales.shape[1], p.numel() * page_size, scales.shape[3]) + tk = (p.unsqueeze(1) * page_size + + torch.arange(page_size, device=self.device)).flatten() + staged_s = src.to(self.device, non_blocking=True) + scales[:, :, tk] = staged_s + del staged_s + + if self.copy_stream is not None: + self.copy_stream.wait_stream(torch.cuda.current_stream(self.device)) + with torch.cuda.stream(self.copy_stream): + _go() + self.copy_stream.synchronize() + else: + _go() + + def read_kv_to_mirror(self, host_slots: torch.Tensor, mirror_codes, + mirror_scales, logical_pages: torch.Tensor) -> None: + """H2H (pinned memcpy): slots KV do tier -> espelho do offloader nas páginas + lógicas dadas. Usado quando o offloader está ativo: reidratar NÃO sobe bytes pra + GPU — o fetch lazy do offloader cuida da residência na atenção.""" + assert self.has_kv + s = host_slots.long() + lp = logical_pages.long().cpu() + page = self._kv_codes.tensor.shape[3] + mirror_codes[:, :, lp] = self._kv_codes.tensor[:, :, s] + if self._kv_scales is not None and mirror_scales is not None: + t_s = self._kv_scales.tensor.view(2, -1, self._kv_codes.tensor.shape[2], page, + self._kv_scales.tensor.shape[-1]) + m_s = mirror_scales.view(2, -1, mirror_codes.shape[2], page, + self._kv_scales.tensor.shape[-1]) + m_s[:, :, lp] = t_s[:, :, s] + + def _drop_oldest(self) -> None: + if not self._index: + raise RuntimeError("HostPrefixTier: KV host store cheio e índice vazio") + _, ent = self._index.popitem(last=False) # o mais frio + if ent.kv_slots is not None: + self._kv_free.extend(ent.kv_slots.tolist()) + ent.kv_slots = None + if ent.gdn_slot is not None: + self.free_gdn(ent.gdn_slot) + ent.gdn_slot = None + + # ---------------- índice de prefixos ---------------- + @staticmethod + def _key_of(tokens: torch.Tensor) -> bytes: + return tokens.to(torch.int64).cpu().numpy().tobytes() + + def put_prefix(self, key_tokens: torch.Tensor, kv_pages: torch.Tensor | None = None, + host_gdn_slot: int | None = None) -> None: + """Registra (ou faz merge numa entrada existente) um prefixo frio. key_tokens deve + ser page-aligned (o tree já alinha). kv_pages = slots host do KV store.""" + key = self._key_of(key_tokens) + ent = self._index.get(key) + if ent is None: + ent = _Entry(key_tokens.cpu()) + self._index[key] = ent + else: + self._index.move_to_end(key) # refresh LRU + if kv_pages is not None: + if ent.kv_slots is not None and ent.kv_slots is not kv_pages: + self._kv_free.extend(ent.kv_slots.tolist()) # substituição: libera os velhos + ent.kv_slots = kv_pages.to(torch.int32).cpu() + if host_gdn_slot is not None: + if ent.gdn_slot is not None and ent.gdn_slot != host_gdn_slot: + self.free_gdn(ent.gdn_slot) # a cópia nova substitui a anterior + ent.gdn_slot = host_gdn_slot + + def lookup(self, input_ids: torch.Tensor, page_size: int = 64) -> HostMatch | None: + """Maior prefixo page-aligned presente no índice (python/CPU, fora do hot path).""" + toks = input_ids.to(torch.int64).cpu() + n = (toks.numel() // page_size) * page_size + while n > 0: + key = self._key_of(toks[:n]) + hit = self._index.get(key) + if hit is not None: + self._index.move_to_end(key) + return HostMatch(n, hit.kv_slots, hit.gdn_slot) + n -= page_size + return None + + def lookup_exact(self, key_tokens: torch.Tensor) -> HostMatch | None: + """The entry for exactly this prefix: one hash, where lookup re-hashes every shorter + page-aligned prefix on a miss.""" + key = self._key_of(key_tokens) + hit = self._index.get(key) + if hit is None: + return None + self._index.move_to_end(key) + return HostMatch(key_tokens.numel(), hit.kv_slots, hit.gdn_slot) + + def evict_host(self, key_tokens: torch.Tensor) -> None: + """Remove uma entrada do índice e libera os slots host (KV + GDN).""" + ent = self._index.pop(self._key_of(key_tokens), None) + if ent is None: + return + if ent.kv_slots is not None: + self._kv_free.extend(ent.kv_slots.tolist()) + if ent.gdn_slot is not None: + self.free_gdn(ent.gdn_slot) + + +__all__ = ["HostPrefixTier", "HostMatch"] diff --git a/python/freetoken/kvcache/hybrid_radix_cache.py b/python/freetoken/kvcache/hybrid_radix_cache.py index 0b66253b17..c09800d956 100644 --- a/python/freetoken/kvcache/hybrid_radix_cache.py +++ b/python/freetoken/kvcache/hybrid_radix_cache.py @@ -71,6 +71,16 @@ def __init__(self, device: torch.device, page_size: int) -> None: self.full_protected = 0 self.mamba_evictable = 0 # number of live, unlocked snapshots self.mamba_protected = 0 + # Host tier (FT_GDN_HOST_TIER): wired by CacheManager. ``on_evict_node`` fires with + # the node IMMEDIATELY before its snapshot is freed (the slot is still valid -- + # enqueue the D2H before returning). ``on_match_dead_snapshot`` may rehydrate a + # tombstone during match: return a fresh GPU slot to attach, None to keep walking. + self.on_evict_node = None + self.on_match_dead_snapshot = None + # KV host tier (FT_PREFIX_HOST): fires for every node whose KV pages LEAVE the tree + # (evict_full's leaf + cascade-exposed tombstone parents) -- BEFORE the free, while + # the bytes are still valid in the GPU buffer (the D2H must happen inside the call). + self.on_evict_kv = None # ---------------------------------------------------------------- match / insert def match_prefix(self, input_ids: torch.Tensor) -> HybridMatch: @@ -83,6 +93,14 @@ def match_prefix(self, input_ids: torch.Tensor) -> HybridMatch: while not cur.is_root(): if cur.mamba_value is not None: return HybridMatch(self._collect_kv(cur), end_len, cur.mamba_value, cur) + if self.on_match_dead_snapshot is not None: + slot = self.on_match_dead_snapshot(cur) + if slot is not None: + # revive the tombstone: the rehydrated snapshot rides the tree again + cur.mamba_value = slot + if cur.mamba_ref_count == 0: + self.mamba_evictable += 1 + return HybridMatch(self._collect_kv(cur), end_len, slot, cur) end_len -= cur.length cur = cur.parent return HybridMatch(self.empty, 0, None, self.root) @@ -159,8 +177,12 @@ def evict_full(self, num_tokens: int) -> EvictResult: if node.ref_count != 0 or not node.is_leaf() or node.is_root(): continue freed += node.length + if self.on_evict_kv is not None: + self.on_evict_kv(node) kv.append(node.value) self.full_evictable -= node.length + if node.mamba_value is not None and self.on_evict_node is not None: + self.on_evict_node(node) self._free_node_mamba(node, mamba) parent, casc = self._cascade_tombstone_leaves(self._unlink(node), kv) freed += casc @@ -180,14 +202,19 @@ def evict_mamba(self, num: int) -> EvictResult: node = heapq.heappop(cands) if node.mamba_value is None or node.mamba_ref_count != 0 or node.is_root(): continue - if node.is_leaf() and node.ref_count == 0: + if self.on_evict_node is not None: + self.on_evict_node(node) + if node.is_leaf() and node.ref_count == 0 and self.on_evict_node is None: kv.append(node.value) self.full_evictable -= node.length self._free_node_mamba(node, mamba) freed += 1 self._cascade_tombstone_leaves(self._unlink(node), kv) else: - self._free_node_mamba(node, mamba) # tombstone internal (or locked-KV) node + # internal (or locked-KV) node -> tombstone; with a host tier wired, LEAF too: + # the snapshot is safely on the host, and keeping the leaf's KV in the tree + # is what lets a later match rehydrate the boundary instead of re-prefilling. + self._free_node_mamba(node, mamba) freed += 1 return EvictResult(torch.cat(kv) if kv else self.empty, mamba) @@ -233,6 +260,8 @@ def _cascade_tombstone_leaves(self, parent: RadixTreeNode, kv_out: List[torch.Te freed = 0 while (parent.mamba_value is None and parent.is_leaf() and parent.ref_count == 0 and not parent.is_root()): + if self.on_evict_kv is not None: + self.on_evict_kv(parent) kv_out.append(parent.value) self.full_evictable -= parent.length freed += parent.length @@ -255,6 +284,16 @@ def _collect_kv(self, node: RadixTreeNode) -> torch.Tensor: vals.reverse() return torch.cat(vals) if vals else self.empty + def _collect_key(self, node: RadixTreeNode) -> torch.Tensor: + """The full token prefix ending at ``node`` (segment keys concatenated root->node).""" + keys: List[torch.Tensor] = [] + n = node + while not n.is_root(): + keys.append(n._key) + n = n.parent + keys.reverse() + return torch.cat(keys) if keys else torch.empty(0, dtype=torch.int32) + def _leaves(self) -> List[RadixTreeNode]: out, stack = [], [self.root] while stack: diff --git a/python/freetoken/kvcache/hybrid_swa_pool.py b/python/freetoken/kvcache/hybrid_swa_pool.py index 022eb49918..57a5734aa4 100644 --- a/python/freetoken/kvcache/hybrid_swa_pool.py +++ b/python/freetoken/kvcache/hybrid_swa_pool.py @@ -23,6 +23,62 @@ class _KVGroupStorage: k_buffer: torch.Tensor v_buffer: torch.Tensor storage_shape: tuple[int, int, int] + logical_head_dim: int + # (2, num_layers, num_slots, local_kv_heads) fp32, or None for an unquantized group. + scale_buffer: torch.Tensor | None = None + # (2, num_layers, num_slots, local_kv_heads, head_dim // 16) uint8, or None. + block_scale_buffer: torch.Tensor | None = None + + +def _alloc_group_storage( + *, + num_layers: int, + local_kv_heads: int, + head_dim: int, + device: torch.device, + store_dtype: torch.dtype, + outer_size: int, + inner_size: int, + kv_quant: str, +) -> _KVGroupStorage: + """One group's code buffer (+ fp8 scale buffer), shared by the initial allocation + and the in-place rebuild so the two can never drift. + + A quantized buffer is zero-filled: a stale e4m3 code decodes to a real number, so + an unwritten slot (the dummy page, a padded request row) would poison attention, + while code 0x00 is exactly 0.0. One memset per allocation, same as + kvcache/bsa_pool.py. The 16-bit buffer keeps torch.empty. + """ + stored_dim = head_dim // 2 if kv_quant == "nvfp4" else head_dim + shape = (2, num_layers, outer_size, inner_size, local_kv_heads, stored_dim) + if kv_quant != "none": + from freetoken.kernel.triton.kv_quant import alloc_codes + + buffer = alloc_codes(shape, device) + scale = torch.zeros( + (2, num_layers, outer_size * inner_size, local_kv_heads), + device=device, + dtype=torch.float32, + ) + else: + buffer = torch.empty(shape, device=device, dtype=store_dtype) + scale = None + block_scale = None + if kv_quant == "nvfp4": + block_scale = torch.zeros( + (2, num_layers, outer_size * inner_size, local_kv_heads, head_dim // 16), + device=device, + dtype=torch.uint8, + ) + return _KVGroupStorage( + buffer=buffer, + k_buffer=buffer[0], + v_buffer=buffer[1], + storage_shape=(outer_size * inner_size, local_kv_heads, stored_dim), + logical_head_dim=head_dim, + scale_buffer=scale, + block_scale_buffer=block_scale, + ) class HybridSWAKVCache(BaseKVCachePool): @@ -37,14 +93,26 @@ def __init__( dtype: torch.dtype, device: torch.device, num_swa_tokens: int | None = None, + kv_quant: str = "none", ) -> None: + if kv_quant not in ("none", "fp8", "nvfp4"): + raise ValueError(f"unsupported hybrid-SWA kv_quant {kv_quant!r}") specs = {group.name: group for group in groups if group.num_layers > 0} if set(specs) != {"full", "swa"}: raise ValueError(f"HybridSWAKVCache requires full and swa groups, got {sorted(specs)}") + if kv_quant == "nvfp4" and any(spec.head_dim % 16 for spec in specs.values()): + raise ValueError("NVFP4 KV requires head_dim divisible by 16") + + from .mha_pool import _kv_store_dtype self._num_layers = num_layers self._device = device - self._dtype = dtype + self.kv_quant = kv_quant + self._compute_dtype = dtype + # What the BUFFER holds -- fp8/uint8 codes when quantized -- reported as + # store_dtype. The dtype property keeps answering the compute dtype, same + # contract as MHAKVCache (kvcache/base.py): backends size their scratch with it. + self._store_dtype = _kv_store_dtype(dtype, kv_quant) self._full_num_tokens = num_full_pages * page_size self._swa_num_tokens = num_swa_tokens if num_swa_tokens is not None else self._full_num_tokens self._page_size = page_size @@ -63,6 +131,7 @@ def __init__( inner_size=page_size, dtype=dtype, device=device, + kv_quant=kv_quant, ) self.swa_kv_pool = self._allocate_group( specs["swa"], @@ -71,6 +140,7 @@ def __init__( inner_size=1, dtype=dtype, device=device, + kv_quant=kv_quant, ) self._storages = { "full": self.full_kv_pool, @@ -88,18 +158,20 @@ def _allocate_group( inner_size: int, dtype: torch.dtype, device: torch.device, + kv_quant: str = "none", ) -> _KVGroupStorage: + from .mha_pool import _kv_store_dtype + local_kv_heads = div_even(spec.num_kv_heads, tp_size, allow_replicate=True) - buffer = torch.empty( - (2, spec.num_layers, outer_size, inner_size, local_kv_heads, spec.head_dim), + return _alloc_group_storage( + num_layers=spec.num_layers, + local_kv_heads=local_kv_heads, + head_dim=spec.head_dim, device=device, - dtype=dtype, - ) - return _KVGroupStorage( - buffer=buffer, - k_buffer=buffer[0], - v_buffer=buffer[1], - storage_shape=(outer_size * inner_size, local_kv_heads, spec.head_dim), + store_dtype=_kv_store_dtype(dtype, kv_quant), + outer_size=outer_size, + inner_size=inner_size, + kv_quant=kv_quant, ) @staticmethod @@ -200,6 +272,26 @@ def v_cache(self, index: int) -> torch.Tensor: ref = self.layers_mapping[index] return self._storages[ref.group].v_buffer[ref.index] + def k_scale(self, index: int) -> torch.Tensor | None: + ref = self.layers_mapping[index] + scale = self._storages[ref.group].scale_buffer + return None if scale is None else scale[0][ref.index] + + def v_scale(self, index: int) -> torch.Tensor | None: + ref = self.layers_mapping[index] + scale = self._storages[ref.group].scale_buffer + return None if scale is None else scale[1][ref.index] + + def k_block_scale(self, index: int) -> torch.Tensor | None: + ref = self.layers_mapping[index] + scale = self._storages[ref.group].block_scale_buffer + return None if scale is None else scale[0][ref.index] + + def v_block_scale(self, index: int) -> torch.Tensor | None: + ref = self.layers_mapping[index] + scale = self._storages[ref.group].block_scale_buffer + return None if scale is None else scale[1][ref.index] + def store_kv( self, k: torch.Tensor, @@ -207,13 +299,41 @@ def store_kv( out_loc: torch.Tensor, layer_id: int, ) -> None: - from freetoken.kernel import store_cache - ref = self.layers_mapping[layer_id] storage = self._storages[ref.group] indices = out_loc if ref.group == "swa": indices = self.translate_loc_from_full_to_swa(out_loc) + if self.kv_quant == "nvfp4": + from freetoken.kernel.triton.kv_nvfp4 import quantize_nvfp4_to_cache + + quantize_nvfp4_to_cache( + k=k, + v=v, + out_loc=indices, + k_cache=storage.k_buffer[ref.index].view(storage.storage_shape), + v_cache=storage.v_buffer[ref.index].view(storage.storage_shape), + k_scale=storage.scale_buffer[0][ref.index], + v_scale=storage.scale_buffer[1][ref.index], + k_block_scale=storage.block_scale_buffer[0][ref.index], + v_block_scale=storage.block_scale_buffer[1][ref.index], + ) + return + if self.kv_quant == "fp8": + from freetoken.kernel.triton.kv_quant import quantize_kv_to_cache + + quantize_kv_to_cache( + k=k, + v=v, + out_loc=indices, + k_cache=storage.k_buffer[ref.index].view(storage.storage_shape), + v_cache=storage.v_buffer[ref.index].view(storage.storage_shape), + k_scale=storage.scale_buffer[0][ref.index], + v_scale=storage.scale_buffer[1][ref.index], + ) + return + from freetoken.kernel import store_cache + store_cache( k_cache=storage.k_buffer[ref.index].view(storage.storage_shape), v_cache=storage.v_buffer[ref.index].view(storage.storage_shape), @@ -236,7 +356,12 @@ def device(self) -> torch.device: @property def dtype(self) -> torch.dtype: - return self._dtype + return self._compute_dtype + + @property + def store_dtype(self) -> torch.dtype: + """Element type of the KV buffers: e4m3/uint8 codes when quantized.""" + return self._store_dtype @property def num_layers(self) -> int: @@ -245,24 +370,31 @@ def num_layers(self) -> int: @staticmethod def _group_geometry(group: _KVGroupStorage) -> tuple: # Everything the realloc needs that does NOT pin the old buffer alive: layer count, - # kv heads, head_dim, device, dtype. (Plain ints + device/dtype handles, no tensor.) - _, num_layers, _old_outer, _old_inner, local_kv_heads, head_dim = group.buffer.shape - return (num_layers, local_kv_heads, head_dim, group.buffer.device, group.buffer.dtype) + # kv heads, logical head_dim, device, storage dtype, and quantized sidecars. + # (Plain ints + device/dtype handles, no tensor.) + _, num_layers, _old_outer, _old_inner, local_kv_heads, _stored_dim = group.buffer.shape + return ( + num_layers, + local_kv_heads, + group.logical_head_dim, + group.buffer.device, + group.buffer.dtype, + "nvfp4" if group.block_scale_buffer is not None else ("fp8" if group.scale_buffer is not None else "none"), + ) @staticmethod def _alloc_group(geom: tuple, outer_size: int, inner_size: int) -> _KVGroupStorage: # Only the outer (page/token) dimension changes; the rest comes from ``geom``. - num_layers, local_kv_heads, head_dim, device, dtype = geom - buffer = torch.empty( - (2, num_layers, outer_size, inner_size, local_kv_heads, head_dim), + num_layers, local_kv_heads, head_dim, device, store_dtype, kv_quant = geom + return _alloc_group_storage( + num_layers=num_layers, + local_kv_heads=local_kv_heads, + head_dim=head_dim, device=device, - dtype=dtype, - ) - return _KVGroupStorage( - buffer=buffer, - k_buffer=buffer[0], - v_buffer=buffer[1], - storage_shape=(outer_size * inner_size, local_kv_heads, head_dim), + store_dtype=store_dtype, + outer_size=outer_size, + inner_size=inner_size, + kv_quant=kv_quant, ) def rebuild(self, num_full_pages: int, num_swa_tokens: int | None = None) -> None: @@ -356,10 +488,22 @@ def unit_bytes(self) -> tuple[int, int]: full = self.full_kv_pool.buffer swa = self.swa_kv_pool.buffer full_tokens = int(full.shape[2]) * int(full.shape[3]) - return ( - int(full.numel() * full.element_size()) // full_tokens, - int(swa.numel() * swa.element_size()) // self._swa_num_tokens, - ) + kv = int(full.numel() * full.element_size()) // full_tokens + swa_b = int(swa.numel() * swa.element_size()) // self._swa_num_tokens + # fp8 codes are priced with their scale sidecar, matching kv_cost exactly. + if self.full_kv_pool.scale_buffer is not None: + fs = self.full_kv_pool.scale_buffer + kv += int(fs.numel() * fs.element_size()) // full_tokens + if self.swa_kv_pool.scale_buffer is not None: + ss = self.swa_kv_pool.scale_buffer + swa_b += int(ss.numel() * ss.element_size()) // self._swa_num_tokens + if self.full_kv_pool.block_scale_buffer is not None: + fs = self.full_kv_pool.block_scale_buffer + kv += int(fs.numel() * fs.element_size()) // full_tokens + if self.swa_kv_pool.block_scale_buffer is not None: + ss = self.swa_kv_pool.block_scale_buffer + swa_b += int(ss.numel() * ss.element_size()) // self._swa_num_tokens + return kv, swa_b # ---- SWA pool sizing (pure arithmetic; the pool family's geometry formulas) ---- diff --git a/python/freetoken/kvcache/kv_host_offload.py b/python/freetoken/kvcache/kv_host_offload.py new file mode 100644 index 0000000000..1f8c358426 --- /dev/null +++ b/python/freetoken/kvcache/kv_host_offload.py @@ -0,0 +1,449 @@ +"""KV host offload for the QSA paged pool (Qwen3.8-Flash-Next, ``--kv-host-pages``). + +The engine keeps TWO page counts: ``num_pages`` logical pages (what the scheduler, the page +table and the radix tree see) and a smaller GPU-resident physical pool. This module is the +bridge: + +* the pinned host **mirror** holds one K/V page per logical page (the backing store); +* the GPU K/V buffer is an **LRU cache** of physical slots over the mirror, driven by + flashlib's ``lru_ensure`` (device-side, fixed-shape, CUDA-graph safe -- the same machinery + the MoE expert offload uses); +* every ``store_kv`` is **write-through**: a Triton kernel mirrors the token's K/V rows into + the host bank over UVA, so a physical slot is always clean and eviction is a pure + rebinding (no write-back, no dirty tracking); +* reads translate: the QSA backend selects pages in logical space (the compressed index slab + is per logical page and always resident), ensures the selected pages are resident + (``lru_ensure`` + one fused H2D multi-bank copy), rewrites the block table to physical + slots, and calls the attend kernel unmodified. + +The dummy page is an ordinary cacheable id (``num_logical_pages``); its mirror page stays +zeroed. Stale bindings of freed/reallocated pages are harmless: a new tenant's pages are +fully written (or position-masked) before they are read, and the mirror catches up through +the same write-through. +""" + +from __future__ import annotations + +import torch +from flashlib.kernels.slot_cache import lru_ensure + +from freetoken.moe.host_banks import HostBank +from freetoken.utils import init_logger + +logger = init_logger(__name__) + + +def _reject_block_scaled_pool(pool, knob: str) -> None: + """The host mirrors carry codes plus the fp8 pool's per-token scales; an nvfp4 pool + (--kv-cache-dtype nvfp4) keeps per-block scales in _block_scale_buffer, which no mirror + copies, so a rehydrated page would be codes without scales. Refuse instead of decoding + garbage.""" + if getattr(pool, "_block_scale_buffer", None) is not None: + raise ValueError( + f"{knob} needs a KV pool with per-token scales (bf16 or fp8); the nvfp4 pool keeps " + "block scales the host mirror does not carry -- leave the tier off or serve with " + "--kv-cache-dtype fp8" + ) + + +class KVHostOffloader: + """GPU-slot LRU cache over the pinned host mirror of the logical KV page space.""" + + def __init__(self, pool, num_logical_pages: int, device: torch.device) -> None: + _reject_block_scaled_pool(pool, "--kv-host-pages") + buf = pool._kv_buffer # [2, L, P_gpu + 1, page_size, kv_heads, head_dim] + assert buf.shape[0] == 2 and buf.is_contiguous() + _, num_layers, num_slots, page_size, kv_heads, head_dim = buf.shape + self.pool = pool + self.page_size = page_size + self.num_layers = num_layers + self.num_slots = num_slots # buffer pages, dummy GPU page included + self.num_logical = num_logical_pages # scheduler-visible pages; dummy id == this + self.dummy_page = num_logical_pages + self.device = device + # The selection touches whole pages: select_width tokens can span at most + # select_width // page_size + 2 distinct pages (boundary straddle + open group). + self.max_sel_pages = 0 # set by set_select_width (backend init) + # Static decode buffers, created by init_graph_buffers (capture) or lazily. + self._g: dict[str, torch.Tensor] = {} + + self.mirror = HostBank( + (2, num_layers, num_logical_pages + 1, page_size, kv_heads, head_dim), buf.dtype + ) + self.mirror.tensor.zero_() # fresh/dummy pages read back as zeros, never stale bytes + self.mirror.pin() + + # fp8 pool (kv_quant="fp8", PR#354): the pool stores e4m3 CODES + fp32 per-token + # scales. The mirror must carry BOTH or rehydrated pages decode garbage (the read + # path below restores codes+scales into the pool; a bf16-era mirror of raw k/v + # would silently corrupt). The raw-k/v Triton mirror_store only runs on bf16 pools. + self._scale_buf = getattr(pool, "_scale_buffer", None) # [2, L, P*page, H] | None + self.mirror_scales: HostBank | None = None + if self._scale_buf is not None: + scale_heads = self._scale_buf.shape[-1] + self.mirror_scales = HostBank( + (2, num_layers, (num_logical_pages + 1) * page_size, scale_heads), + self._scale_buf.dtype, + ) + self.mirror_scales.tensor.zero_() + self.mirror_scales.pin() + + # LRU slot maps (lru_ensure's id space is logical pages, slot space is buffer pages). + self.phys_of = torch.full((num_logical_pages + 1,), -1, dtype=torch.int32, device=device) + self.logical_of = torch.full((num_slots,), -1, dtype=torch.int32, device=device) + self.usage = torch.zeros((num_slots,), dtype=torch.int64, device=device) + self.step = torch.zeros((), dtype=torch.int64, device=device) + self.src_ids = torch.empty((num_slots,), dtype=torch.int32, device=device) + self.dst_slots = torch.empty((num_slots,), dtype=torch.int32, device=device) + self.num_copy = torch.zeros((1,), dtype=torch.int64, device=device) + self._out_scratch = torch.empty((num_slots,), dtype=torch.int32, device=device) + # Eager-prefill scratch (never captured): per-layer compacted selections (+ their + # pinned host copy for the CPU group packer), the per-group tight query (pinned + # staging + device buffer) and the translated block table. Lazy, grow-only. + self._sel_all: torch.Tensor | None = None + self._counts_dev: torch.Tensor | None = None + self._sel_host = None + self._counts_host = None + self._query_host = None + self._query_buf: torch.Tensor | None = None + self._eff_buf: torch.Tensor | None = None + self._trunc_eager = torch.zeros(1, dtype=torch.int32, device=device) + + from freetoken.kernel.pinned import device_ptr + + # Fused-copy descriptor: 2 * num_layers banks (K/V of each layer), one page row each. + row_bytes = page_size * kv_heads * head_dim * buf.dtype.itemsize + assert row_bytes % 16 == 0 + mirror = self.mirror.tensor + dst_ptrs, src_ptrs, host_ptrs = [], [], [] + for kv in range(2): + for layer in range(num_layers): + assert buf[kv, layer].is_contiguous() and mirror[kv, layer].is_contiguous() + dst_ptrs.append(buf[kv, layer].data_ptr()) + src_ptrs.append(device_ptr(mirror[kv, layer])) + for layer in range(num_layers): # K/V interleaved, the mirror-store kernel's layout + host_ptrs.append(device_ptr(mirror[0, layer])) + host_ptrs.append(device_ptr(mirror[1, layer])) + self._dst_ptrs = torch.tensor(dst_ptrs, dtype=torch.int64, device=device) + self._src_ptrs = torch.tensor(src_ptrs, dtype=torch.int64, device=device) + self._feats = torch.full((2 * num_layers,), row_bytes, dtype=torch.int64, device=device) + self._host_ptrs = torch.tensor(host_ptrs, dtype=torch.int64, device=device) + # fp8: estende o descritor com os bancos de scales (mesma lista de páginas — + # o fused copy roda uma linha por banco por página; feats heterogêneos ok). + # E prepara os ponteiros do mirror_store quantizado (graph-safe, UVA). + self._quant_src_ptrs = None + self._quant_host_ptrs = None + if self.mirror_scales is not None: + sc = self._scale_buf # [2, L, P*page, H] -> por página: [page, H] + scale_row = page_size * sc.shape[-1] * sc.element_size() + assert scale_row % 16 == 0 + qsrc, qhost = [], [] + for kv in range(2): + for layer in range(num_layers): + pool_v = sc[kv, layer].view(-1, page_size, sc.shape[-1]) + mir_v = self.mirror_scales.tensor[kv, layer].view( + -1, page_size, sc.shape[-1]) + assert pool_v.is_contiguous() and mir_v.is_contiguous() + dst_ptrs.append(pool_v.data_ptr()) + src_ptrs.append(device_ptr(mir_v)) + for layer in range(num_layers): + # ordem por camada: (kcode, kscale, vcode, vscale) — casa com o kernel + qsrc += [buf[0, layer].data_ptr(), sc[0, layer].data_ptr(), + buf[1, layer].data_ptr(), sc[1, layer].data_ptr()] + qhost += [device_ptr(mirror[0, layer]), + device_ptr(self.mirror_scales.tensor[0, layer]), + device_ptr(mirror[1, layer]), + device_ptr(self.mirror_scales.tensor[1, layer])] + self._quant_src_ptrs = torch.tensor(qsrc, dtype=torch.int64, device=device) + self._quant_host_ptrs = torch.tensor(qhost, dtype=torch.int64, device=device) + self._dst_ptrs = torch.tensor(dst_ptrs, dtype=torch.int64, device=device) + self._src_ptrs = torch.tensor(src_ptrs, dtype=torch.int64, device=device) + self._feats = torch.cat([ + self._feats, + torch.full((2 * num_layers,), scale_row, dtype=torch.int64, device=device), + ]) + self.mirror_bytes = self.mirror.nbytes + ( + self.mirror_scales.nbytes if self.mirror_scales is not None else 0) + + # ----- geometry --------------------------------------------------------- + def set_select_width(self, select_width: int, block_topk: int = 0, + max_bs: int = 1) -> None: + # The sparse selection is block_topk short runs of consecutive tokens SCATTERED + # across the sequence -- each run can touch its own page (two if it straddles a + # boundary). Contiguity (select_width // page_size) badly underestimates this. + # Hard cap: the WHOLE decode batch's query (max_bs x (1 + sel)) must fit the pool, + # or mid-ensure evictions corrupt the attend. Divide by the batch, not the row. + worst = 2 * (block_topk or (select_width // 4)) + 8 + per_batch = max(16, (self.num_slots - 73) // max(1, max_bs)) + self.max_sel_pages = min(worst, per_batch) + + def prefill_fetch_rows(self, n_write_pages: int) -> int: + """Query rows per prefill attend sub-chunk: the write pages plus this many rows' worth + of selected pages must always fit the physical pool (lru_ensure's distinct bound).""" + margin = 8 + budget = self.num_slots - margin - n_write_pages + if budget < self.max_sel_pages: + raise RuntimeError( + f"KV host offload: {self.num_slots} physical pages cannot hold {n_write_pages} " + f"write pages + one row's selection ({self.max_sel_pages}); the GPU pool is too small" + ) + return max(1, budget // self.max_sel_pages) + + # ----- graph buffers ------------------------------------------------------ + def init_graph_buffers(self, max_bs: int, width_pages: int, select_width: int | None = None) -> None: + """Static decode buffers (capture-safe addresses); sliced to the live bs per step.""" + if select_width is not None: + self.set_select_width(select_width, max_bs=max_bs) + assert self.max_sel_pages > 0, "set_select_width must run before init_graph_buffers" + d = self.device + need = max(self.num_slots, max_bs * (1 + self.max_sel_pages)) + if self._out_scratch.numel() < need: + self._out_scratch = torch.empty((need,), dtype=torch.int32, device=d) + self._g = { + "write_pages": torch.zeros(max_bs, dtype=torch.int32, device=d), + "out_loc_gpu": torch.zeros(max_bs, dtype=torch.int32, device=d), + "sel_pages": torch.zeros((max_bs, self.max_sel_pages), dtype=torch.int32, device=d), + "query": torch.zeros(max_bs * (1 + self.max_sel_pages), dtype=torch.int32, device=d), + "eff_table": torch.zeros((max_bs, width_pages), dtype=torch.int32, device=d), + "trunc": torch.zeros(1, dtype=torch.int32, device=d), + } + + def invalidate_pages(self, logical_pages: torch.Tensor) -> None: + """Marca páginas lógicas como frias (phys_of=-1), desfazendo bindings defasados. + Necessário quando o espelho foi reescrito por fora (reidratação host->espelho): + sem isso, um binding antigo apontaria pra um slot físico com conteúdo de outro + inquilino e o fetch seria pulado. Só desfaz se o slot ainda aponta pra esta página.""" + lp = logical_pages.long().to(self.device) + old = self.phys_of[lp] + mask = old >= 0 + if not bool(mask.any()): + return + old_slots = old[mask].long() + cur = self.logical_of[old_slots] + still = cur == lp[mask].long() + self.logical_of[old_slots[still]] = -1 + self.phys_of[lp[mask]] = -1 + + def trunc_count(self) -> int: + """Dropped selection pages accumulated since the last read (0 = no truncation).""" + n = int(self._trunc_eager.item()) + if n: + self._trunc_eager.zero_() + t = self._g.get("trunc") + if t is not None: + m = int(t.item()) + if m: + t.zero_() + n += m + return n + + def _graph_buffers(self, bs: int, width_pages: int) -> None: + """Lazy (eager-only) static-buffer init; capture paths get them from init_graph_buffers.""" + g = self._g + if not g or g["write_pages"].shape[0] < bs or g["eff_table"].shape[1] != width_pages: + self.init_graph_buffers(max(bs, g["write_pages"].shape[0] if g else 0), width_pages) + + # ----- residency engine --------------------------------------------------- + # lru_ensure's phase-1 register tile is O(K^2): the kernel is tuned for K<=512 + # (MoE's whole-layer materialize). Chunk bigger queries into calls inside that envelope. + _ENSURE_CHUNK = 512 + + def _ensure(self, query: torch.Tensor) -> None: + """Bind every logical page in ``query`` to a physical slot and fill the misses.""" + from freetoken.kernel.fast_index_copy import fast_index_copy_multi_jit + + n = query.numel() + # The chunks are disjoint slices, so distinct(query) <= n; keeping n <= num_slots + # makes the kernel's |distinct| <= num_cached requirement hold by construction. + assert n <= self.num_slots, ( + f"ensure query of {n} pages exceeds the {self.num_slots} physical slots" + ) + if n > self._out_scratch.numel(): + # eager-only path (prefill); capture-time queries always fit the pre-sized buffer + self._out_scratch = torch.empty((n,), dtype=torch.int32, device=self.device) + for off in range(0, n, self._ENSURE_CHUNK): + q = query[off : off + self._ENSURE_CHUNK] + k = q.numel() + lru_ensure( + q, + self.phys_of, + self.logical_of, + self.usage, + self.step, + self._out_scratch[:k], + self.src_ids, + self.dst_slots, + self.num_copy, + ) + fast_index_copy_multi_jit( + self._dst_ptrs, self._src_ptrs, self._feats, + self.dst_slots, self.src_ids, self.num_copy, + ) + + # ----- write path --------------------------------------------------------- + def ensure_write_pages(self, out_loc: torch.Tensor, is_decode: bool) -> torch.Tensor: + """Make every page ``out_loc`` writes to resident; returns the write-page query + (the fetch path re-includes it so a write page can never be evicted mid-forward).""" + from freetoken.kernel.triton.qsa.offload import write_pages + + if is_decode: + n = out_loc.numel() + wp = self._g["write_pages"][:n] + write_pages(out_loc, wp, self.page_size) + self._ensure(wp) + return wp + # The chunk's tokens are position-contiguous, so every 64th slot names a written + # page exactly once -- no torch.unique (a GPU sync) needed. + pages = (out_loc[:: self.page_size] // self.page_size).contiguous() + self._ensure(pages) + return pages + + def translate_slots(self, out_loc: torch.Tensor, is_decode: bool) -> torch.Tensor: + """Logical token slots -> physical (post-ensure; the write pages are pinned-hot).""" + from freetoken.kernel.triton.qsa.offload import translate_slots + + if is_decode: + n = out_loc.numel() + out = self._g["out_loc_gpu"][:n] + else: + out = torch.empty_like(out_loc) + translate_slots(out_loc, self.phys_of, out, self.page_size) + return out + + def mirror_store(self, k: torch.Tensor, v: torch.Tensor, out_loc: torch.Tensor, + layer_id: int, out_loc_gpu: torch.Tensor | None = None) -> None: + if self.mirror_scales is None: + from freetoken.kernel.triton.qsa.offload import mirror_store + + mirror_store(k, v, out_loc, self._host_ptrs, self.pool._dense(layer_id)) + return + # fp8 pool: o pool JÁ quantizou (store_kv rodou antes). Espelhamos codes+scales + # LIDOS DO POOL (a única fonte da verdade pós-quantização) — nunca re-quantizar. + assert out_loc_gpu is not None, "fp8 mirror_store precisa dos slots físicos" + self._mirror_store_fp8(out_loc, out_loc_gpu, self.pool._dense(layer_id)) + + def _mirror_store_fp8(self, out_loc_logical: torch.Tensor, out_loc_phys: torch.Tensor, + dense_layer: int) -> None: + """codes+scales do pool (slots físicos, pós-quantização) -> espelho (slots lógicos). + Graph-safe: kernel Triton com stores UVA, nada de .cpu() no caminho.""" + from freetoken.kernel.triton.qsa.offload import mirror_store_quant + + H, D = self.mirror.tensor.shape[4], self.mirror.tensor.shape[5] + scH = self._scale_buf.shape[-1] + mirror_store_quant( + self.pool._kv_buffer[0, dense_layer].view(-1, H * D), + self._scale_buf[0, dense_layer], + self.pool._kv_buffer[1, dense_layer].view(-1, H * D), + self._scale_buf[1, dense_layer], + out_loc_phys, out_loc_logical, + self._quant_src_ptrs, self._quant_host_ptrs, dense_layer, + ) + + # ----- read path ------------------------------------------------------------ + def compact_all(self, indices: torch.Tensor, token_to_req: torch.Tensor, + block_table: torch.Tensor) -> torch.Tensor: + """One launch for the whole layer: per query row, the distinct logical pages its + selection touches ([rows, max_sel_pages], first-page padded) plus per-row distinct + counts. No host sync.""" + from freetoken.kernel.triton.qsa.offload import compact_selected_pages + + rows = indices.shape[0] + buf = self._sel_all + if buf is None or buf.shape[0] < rows: + n = max(rows, 4096) + buf = self._sel_all = torch.empty((n, self.max_sel_pages), dtype=torch.int32, device=self.device) + self._counts_dev = torch.empty((n,), dtype=torch.int32, device=self.device) + sel = buf[:rows] + compact_selected_pages( + indices, token_to_req, block_table, sel, self.page_size, self.dummy_page, + self._trunc_eager, counts=self._counts_dev[:rows], + ) + return sel + + def iter_prefill_groups(self, sel_all: torch.Tensor, write_pages: torch.Tensor, + block_table: torch.Tensor): + """Pack the layer's query rows into residency-feasible groups, then per group upload + the tight query (write set + the group's REAL union), ensure, translate; yields + ``((row_start, row_end), eff_table)`` and the caller attends between iterations. + + One D2H per layer (the compacted page lists); the CPU packer then works with exact + sets, so the per-group distinct bound (|write| + |union| + slack <= num_slots) is + VERIFIED, not estimated -- lru_ensure's capacity requirement holds by construction. + Consecutive rows select mostly the same pages, so groups pack several rows each. + Buffers are reused across groups; single-stream ordering (attend of group i is + enqueued before group i+1's H2D/ensure/translate) makes that safe. + """ + import numpy as np + + rows = sel_all.shape[0] + host = self._sel_host + if host is None or host.shape[0] < rows: + n = max(rows, 4096) + host = self._sel_host = torch.empty((n, self.max_sel_pages), dtype=torch.int32, pin_memory=True) + self._counts_host = torch.empty((n,), dtype=torch.int32, pin_memory=True) + sel_h, cnt_h = host[:rows], self._counts_host[:rows] + sel_h.copy_(sel_all, non_blocking=True) + cnt_h.copy_(self._counts_dev[:rows], non_blocking=True) + torch.cuda.synchronize(self.device) # the one sync per layer + sel_np, cnt_np = sel_h.numpy(), cnt_h.numpy() + write_np = write_pages.cpu().numpy().astype(np.int64) + # |write| + |union of the group's pages| + slack <= physical slots + cap = self.num_slots - write_np.size - 9 + n_ids = self.num_logical + 1 + + gmask = np.zeros(n_ids, dtype=bool) + gmask[write_np] = True + gstart, gcost = 0, 0 + groups: list[tuple[int, int, np.ndarray]] = [] + for r in range(rows): + pages = sel_np[r, : cnt_np[r]].astype(np.int64) + extra = int((~gmask[pages]).sum()) + if gcost + extra > cap and r > gstart: + groups.append((gstart, r, np.flatnonzero(gmask))) + gmask[:] = False + gmask[write_np] = True + gstart, gcost = r, 0 + extra = int((~gmask[pages]).sum()) + gmask[pages] = True + gcost += extra + groups.append((gstart, rows, np.flatnonzero(gmask))) + + from freetoken.kernel.triton.qsa.offload import translate_table + + for gstart, gend, pages_np in groups: + n = pages_np.size + qh = self._query_host + if qh is None or qh.numel() < n: + qh = self._query_host = torch.empty((max(n, 2048),), dtype=torch.int32, pin_memory=True) + self._query_buf = torch.empty_like(qh, device=self.device) + qh[:n] = torch.from_numpy(pages_np.astype(np.int32)) + qd = self._query_buf[:n] + qd.copy_(qh[:n], non_blocking=True) + self._ensure(qd) + eff = self._eff_buf + if eff is None or eff.shape != block_table.shape: + eff = self._eff_buf = torch.empty_like(block_table) + translate_table(block_table, self.phys_of, eff) + yield (gstart, gend), eff + + def fetch_for_attend(self, indices: torch.Tensor, md) -> torch.Tensor: + """Decode (graph-safe): ensure the selected pages resident, return the PHYSICAL + block table for the attend kernel. Fixed shapes throughout.""" + from freetoken.kernel.triton.qsa.offload import compact_selected_pages, translate_table + + rows = indices.shape[0] + self._graph_buffers(rows, md.block_table.shape[1]) + g = self._g + sel = g["sel_pages"][:rows] + compact_selected_pages(indices, md.token_to_req, md.block_table, sel, self.page_size, self.dummy_page, g["trunc"]) + per_row = 1 + self.max_sel_pages + query = g["query"][: rows * per_row].view(rows, per_row) + query[:, 0].copy_(md.write_pages) + query[:, 1:].copy_(sel) + self._ensure(query.view(-1)) + eff = g["eff_table"][:rows] + translate_table(md.block_table, self.phys_of, eff) + return eff + + +__all__ = ["KVHostOffloader"] diff --git a/python/freetoken/kvcache/linear_state_pool.py b/python/freetoken/kvcache/linear_state_pool.py index ff7539410c..df5a3396da 100644 --- a/python/freetoken/kvcache/linear_state_pool.py +++ b/python/freetoken/kvcache/linear_state_pool.py @@ -1,6 +1,7 @@ from __future__ import annotations import math +import os import torch from freetoken.distributed import get_tp_info @@ -260,6 +261,34 @@ def linear_state_bytes_per_req( __all__ = ["LinearStatePool", "linear_state_bytes_per_req"] +_GDN_CHUNK_SIZE = 64 # keep in lockstep with kernel/fla/chunk_delta_h.py CHUNK_SIZE + + +def gdn_prefill_workspace_bytes(config) -> int: + """Bytes of the GDN chunked-prefill buffers that scale with the extend length. + + ``chunk_gated_delta_rule_fwd_h`` materializes the per-chunk recurrent states for + the WHOLE extend at once -- ``[NT, H, V, K]`` with ``NT = ceil(extend / 64)`` -- + plus the ``v_new`` value buffer (288 MiB at an 8192-token extend for qwen4_exp). + This is a lower bound on the GDN prefill peak, not the peak: w/u/A/o and the + in_proj/conv activations still come out of the ``(1 - memory_ratio)`` headroom. + 0 for non-GDN models.""" + linear_group = config.model_config.linear_attention_group() + if linear_group is None: + return 0 + # EngineConfig has no max_extend_tokens (the served config is a SchedulerConfig subclass) + extent = min(getattr(config, "max_extend_tokens", 8192), config.max_seq_len) + nt = -(-extent // _GDN_CHUNK_SIZE) + states = ( + nt + * linear_group.num_value_heads + * linear_group.value_head_dim + * linear_group.key_head_dim + ) + v_new = extent * linear_group.num_value_heads * linear_group.value_head_dim + return (states + v_new) * config.dtype.itemsize + + def state_pool_bytes(config, num_slots: int | None = None) -> int: """Total GDN state-pool bytes at ``num_slots`` PHYSICAL slots (default: the startup slot count). The engine adds this to the KV family's fixed cost when budgeting -- @@ -285,7 +314,10 @@ def _linear_pool_num_slots(config) -> int: if config.cache_type != "hybrid_radix": return mr + 1 # live + dummy/padding ratio = config.linear_state_cache_ratio - n_cache = max(4, int(ratio * mr)) + # With the host snapshot tier (FT_GDN_HOST_TIER=1) evicted snapshots spill to RAM and come + # back on a hit, so the GPU cache need not keep four of them resident. + floor = 1 if os.environ.get("FT_GDN_HOST_TIER", "0") == "1" else 4 + n_cache = max(floor, int(ratio * mr)) return 4 * mr + n_cache + 1 # live + 2 ping-pong + locked committed snapshot + cache + padding diff --git a/python/freetoken/kvcache/mha_pool.py b/python/freetoken/kvcache/mha_pool.py index 8ed280b966..f198887a7f 100644 --- a/python/freetoken/kvcache/mha_pool.py +++ b/python/freetoken/kvcache/mha_pool.py @@ -9,6 +9,19 @@ from .base import BaseKVCachePool +def _kv_store_dtype(dtype: torch.dtype, kv_quant: str) -> torch.dtype: + """Storage dtype of the KV buffer for a quantization mode.""" + if kv_quant == "none": + return dtype + if kv_quant == "nvfp4": + return torch.uint8 + if kv_quant == "fp8": + from freetoken.kernel.triton.kv_quant import kv_codes_dtype + + return kv_codes_dtype() + raise ValueError(f"unknown kv_quant {kv_quant!r}") + + class MHAKVCache(BaseKVCachePool): """ Base class for key-value caches. @@ -20,6 +33,16 @@ class MHAKVCache(BaseKVCachePool): that hold no paged KV; passing the full-attention layer ids here allocates one storage slab per KV layer (not per model layer) and remaps the global id to its dense slot, avoiding a multiple-x over-allocation of unused slabs. + + ``kv_quant="fp8"`` halves the cache: rows become e4m3 codes and every + ``(token, slab, layer, kv head)`` row carries one fp32 scale (see + :mod:`freetoken.kernel.triton.kv_quant`). The codes buffer keeps the exact same + shape as the 16-bit one, so ``k_cache``/``v_cache`` and every index into them are + unchanged -- only the element type, and ``store_kv``'s write path, differ. + + ``kv_quant="nvfp4"`` packs two E2M1 values per byte and adds one E4M3 scale + per 16 values, alongside the FP32 row scale. Logical head_dim is retained + separately so rebuild never mistakes the packed width for the model width. """ def __init__( @@ -32,10 +55,17 @@ def __init__( dtype: torch.dtype, device: torch.device, layer_ids: Sequence[int] | None = None, + kv_quant: str = "none", ) -> None: tp_info = get_tp_info() local_kv_heads = div_even(num_kv_heads, tp_info.size, allow_replicate=True) + _kv_store_dtype(dtype, kv_quant) + if kv_quant == "nvfp4" and head_dim % 16: + raise ValueError("NVFP4 KV requires head_dim divisible by 16") + self._head_dim = head_dim self._num_layers = num_layers + self.kv_quant = kv_quant + self._compute_dtype = dtype if layer_ids is None: num_storage_layers = num_layers self._layer_map: list[int] | None = None @@ -47,15 +77,49 @@ def __init__( raise ValueError(f"KV layer id {global_id} outside [0, {num_layers})") layer_map[global_id] = dense self._layer_map = layer_map - self._kv_buffer = torch.empty( - (2, num_storage_layers, num_pages, page_size, local_kv_heads, head_dim), - device=device, - dtype=dtype, - ) + self._device = device + self._alloc(num_pages, page_size, num_storage_layers, local_kv_heads, head_dim) + + def _alloc( + self, + num_pages: int, + page_size: int, + num_storage_layers: int, + local_kv_heads: int, + head_dim: int, + ) -> None: + """Allocate the code buffer (and, when quantized, the scale buffer). + + A quantized buffer is zero-filled -- e4m3 has NaN bit patterns, so an + unwritten slot (the dummy page, a padded request's row) must not read back as + one. The 16-bit buffer keeps ``torch.empty``: it is bytes-sized, never + interpreted, and the memset would cost real startup time on a large cache. + """ + stored_dim = head_dim // 2 if self.kv_quant == "nvfp4" else head_dim + shape = (2, num_storage_layers, num_pages, page_size, local_kv_heads, stored_dim) + self._block_scale_buffer = None + if self.kv_quant in ("fp8", "nvfp4"): + from freetoken.kernel.triton.kv_quant import alloc_codes + + self._kv_buffer = alloc_codes(shape, self._device) + self._scale_buffer = torch.zeros( + (2, num_storage_layers, num_pages * page_size, local_kv_heads), + device=self._device, + dtype=torch.float32, + ) + else: + self._kv_buffer = torch.empty( + shape, device=self._device, dtype=self._compute_dtype + ) + self._scale_buffer = None + if self.kv_quant == "nvfp4": + self._block_scale_buffer = torch.zeros( + (2, num_storage_layers, num_pages * page_size, local_kv_heads, head_dim // 16), + device=self._device, dtype=torch.uint8, + ) self._k_buffer = self._kv_buffer[0] self._v_buffer = self._kv_buffer[1] - self._device = device - self._storage_shape = (num_pages * page_size, local_kv_heads, head_dim) + self._storage_shape = (num_pages * page_size, local_kv_heads, stored_dim) def rebuild(self, num_pages: int) -> None: """Reallocate the KV buffer for ``num_pages`` pages IN PLACE. @@ -64,23 +128,18 @@ def rebuild(self, num_pages: int) -> None: existing buffer; only the page count changes. Views and ``_storage_shape`` are refreshed. Object identity is preserved so cached backend references stay valid. """ - _, num_storage_layers, _old_pages, page_size, local_kv_heads, head_dim = self._kv_buffer.shape - dtype = self._kv_buffer.dtype + _, num_storage_layers, _old_pages, page_size, local_kv_heads, _stored_dim = self._kv_buffer.shape + head_dim = self._head_dim device = self._device + self._block_scale_buffer = None self._k_buffer = None self._v_buffer = None self._kv_buffer = None + self._scale_buffer = None if device.type == "cuda": torch.cuda.synchronize(device) torch.cuda.empty_cache() - self._kv_buffer = torch.empty( - (2, num_storage_layers, num_pages, page_size, local_kv_heads, head_dim), - device=device, - dtype=dtype, - ) - self._k_buffer = self._kv_buffer[0] - self._v_buffer = self._kv_buffer[1] - self._storage_shape = (num_pages * page_size, local_kv_heads, head_dim) + self._alloc(num_pages, page_size, num_storage_layers, local_kv_heads, head_dim) @classmethod def kv_cost(cls, config) -> tuple[int, int, int, int]: @@ -101,7 +160,13 @@ def rebuild_from_config( def unit_bytes(self) -> tuple[int, int]: buf = self._kv_buffer tokens = int(buf.shape[2]) * int(buf.shape[3]) - return int(buf.numel() * buf.element_size()) // tokens, 0 + kv = int(buf.numel() * buf.element_size()) // tokens + if self._scale_buffer is not None: + sc = self._scale_buffer + kv += int(sc.numel() * sc.element_size()) // tokens + if self._block_scale_buffer is not None: + kv += self._block_scale_buffer.numel() // tokens + return kv, 0 def _dense(self, layer_id: int) -> int: if self._layer_map is None: @@ -117,6 +182,26 @@ def k_cache(self, index: int) -> torch.Tensor: def v_cache(self, index: int) -> torch.Tensor: return self._v_buffer[self._dense(index)] + def k_scale(self, index: int) -> torch.Tensor | None: + if self._scale_buffer is None: + return None + return self._scale_buffer[0][self._dense(index)] + + def v_scale(self, index: int) -> torch.Tensor | None: + if self._scale_buffer is None: + return None + return self._scale_buffer[1][self._dense(index)] + + def k_block_scale(self, index: int) -> torch.Tensor | None: + if self._block_scale_buffer is None: + return None + return self._block_scale_buffer[0][self._dense(index)] + + def v_block_scale(self, index: int) -> torch.Tensor | None: + if self._block_scale_buffer is None: + return None + return self._block_scale_buffer[1][self._dense(index)] + def store_kv( self, k: torch.Tensor, @@ -124,9 +209,33 @@ def store_kv( out_loc: torch.Tensor, layer_id: int, ) -> None: + dense = self._dense(layer_id) + if self.kv_quant == "nvfp4": + from freetoken.kernel.triton.kv_nvfp4 import quantize_nvfp4_to_cache + + quantize_nvfp4_to_cache( + k, v, out_loc, + self._k_buffer[dense].view(self._storage_shape), + self._v_buffer[dense].view(self._storage_shape), + self.k_scale(layer_id), self.v_scale(layer_id), + self.k_block_scale(layer_id), self.v_block_scale(layer_id), + ) + return + if self.kv_quant == "fp8": + from freetoken.kernel.triton.kv_quant import quantize_kv_to_cache + + quantize_kv_to_cache( + k=k, + v=v, + out_loc=out_loc, + k_cache=self._k_buffer[dense].view(self._storage_shape), + v_cache=self._v_buffer[dense].view(self._storage_shape), + k_scale=self._scale_buffer[0][dense], + v_scale=self._scale_buffer[1][dense], + ) + return from freetoken.kernel import store_cache - dense = self._dense(layer_id) store_cache( k_cache=self._k_buffer[dense].view(self._storage_shape), v_cache=self._v_buffer[dense].view(self._storage_shape), @@ -141,6 +250,15 @@ def device(self) -> torch.device: @property def dtype(self) -> torch.dtype: + """The COMPUTE dtype (kvcache/base.py): what ``store_kv`` receives and what + backends size their scratch with -- still 16-bit on an fp8 pool.""" + return self._compute_dtype + + @property + def store_dtype(self) -> torch.dtype: + """Element type of ``_kv_buffer``: fp8/uint8 codes when quantized, else the + compute dtype. Reading the buffer itself needs THIS, and a backend that cannot + apply the row scales never gets a quantized pool (supports_fp8_kv gate).""" return self._kv_buffer.dtype @property diff --git a/python/freetoken/kvcache/qsa_pool.py b/python/freetoken/kvcache/qsa_pool.py index 7de70977d5..adcb4f5ae3 100644 --- a/python/freetoken/kvcache/qsa_pool.py +++ b/python/freetoken/kvcache/qsa_pool.py @@ -66,7 +66,11 @@ def __init__( ring_capacity: int | None = None, layer_ids: Sequence[int] | None = None, mrope: bool = False, + kv_quant: str = "none", + num_index_pages: int | None = None, ) -> None: + if kv_quant not in ("none", "fp8", "nvfp4"): + raise ValueError(f"unsupported QSA kv_quant {kv_quant!r}") if index_ratio < 1 or page_size % index_ratio != 0: # slot // index_ratio only names one group when a group never straddles a page. raise ValueError( @@ -93,6 +97,10 @@ def __init__( self._index_dtype = dtype self._page_size = page_size self._mrope = mrope + # KV host offload: the K/V slabs stay at num_pages PHYSICAL pages while the + # compressed index slab covers the (larger) logical page space, so selection keeps + # working for host-resident pages. None -> the slab follows the K/V pages. + self._num_index_pages = num_pages if num_index_pages is None else num_index_pages super().__init__( num_kv_heads=num_kv_heads, num_layers=num_layers, @@ -102,15 +110,26 @@ def __init__( dtype=dtype, device=device, layer_ids=layer_ids, + # "fp8" quantizes ONLY the KV tiers: codes + one scale per (slot, kv_head) + # replace the bf16 K/V, and store_kv's fused writer replaces the + # separate qsa_store_rows. The three index tiers below stay bf16 no matter + # what is asked for -- block selection reads a different tensor + # (models/qwen3_8_flash_next.py builds index_k in the engine dtype), so + # --kv-cache-dtype fp8 leaves retrieval quality, and the score kernel's + # dtype asserts, untouched. + kv_quant=kv_quant, ) self._zero_kv_slabs() - self._alloc_index_tiers(num_pages) + self._alloc_index_tiers(self._num_index_pages) def _zero_kv_slabs(self) -> None: # Defense-in-depth: the attend kernels pos-mask every K/V load (the real fix for # torch.empty's recycled NaN/Inf bit patterns), but a zeroed slab keeps any future - # unmasked read finite instead of model-poisoning. One memset per (re)allocation. - self._kv_buffer.zero_() + # unmasked read finite instead of model-poisoning. One memset per (re)allocation, + # through the byte view so it works whether the slab holds bf16 values or e4m3 + # codes -- a memset is the one op both representations support, and the codes + # live in bytes on every architecture (kv_quant.kv_codes_dtype). + self._kv_buffer.view(torch.uint8).zero_() def _alloc_index_tiers(self, num_pages: int) -> None: # ZERO-initialized: the score kernel reads whole rows of blocks unmasked and relies on @@ -156,6 +175,10 @@ def rebuild(self, num_pages: int) -> None: self._kv_buffer = None self._k_buffer = None self._v_buffer = None + # Same reason as above on an fp8 pool: a grown K/V slab whose scales are gone + # would serve quantized rows at the wrong scale rather than fail. + self._scale_buffer = None + self._block_scale_buffer = None raise @classmethod diff --git a/python/freetoken/layers/embedding.py b/python/freetoken/layers/embedding.py index a22216548a..860c220010 100644 --- a/python/freetoken/layers/embedding.py +++ b/python/freetoken/layers/embedding.py @@ -36,16 +36,40 @@ def __init__( self._embed_scale = embed_scale self._embed_scale_t: torch.Tensor | None = None self._comm = DistributedCommunicator() + self._host_ptr: int | None = None + + def place_on_host(self) -> None: + """Move the table into pinned host RAM; lookups then read their rows over PCIe (UVA).""" + from freetoken.kernel.pinned import alloc_pinned_tensor, device_ptr + from freetoken.kernel.triton.ple import _TL_TABLE_DTYPES + + if self.tp_size > 1: + raise ValueError("a host-resident embedding serves a single GPU only") + if self.weight.dtype not in _TL_TABLE_DTYPES: + raise ValueError(f"a host-resident embedding cannot gather a {self.weight.dtype} table") + host = alloc_pinned_tensor(*self.weight.shape, dtype=self.weight.dtype) + host.copy_(self.weight) + self.weight = host + self._host_ptr = device_ptr(host) @nvtx_annotate("Embedding") def forward(self, x: torch.Tensor) -> torch.Tensor: from freetoken.kernel import indexing - y = indexing( - weights=self.weight, - indices=x, - vocab_range=self.vocab_range if self.tp_size > 1 else None, - ) + if self._host_ptr is not None: + from freetoken.kernel.triton.ple import ple_gather_rows + + rows, dim = self.weight.shape + y = torch.empty(x.numel(), dim, dtype=self.weight.dtype, device=x.device) + ple_gather_rows(self._host_ptr, rows, dim, x.reshape(-1), y, is_fp8=False, + table_dtype=self.weight.dtype) + y = y.view(*x.shape, dim) + else: + y = indexing( + weights=self.weight, + indices=x, + vocab_range=self.vocab_range if self.tp_size > 1 else None, + ) if self.tp_size > 1: y = self._comm.all_reduce(y) @@ -121,6 +145,17 @@ def state_dict( return super().state_dict(prefix=prefix, result=result) return {} if result is None else result + def _logits(self, x: torch.Tensor) -> torch.Tensor: + """The local vocab-shard GEMM; the seam a quantized head overrides.""" + module = self.tied_embedding or self + return F.linear(x, module.weight, self.bias) + + def _logits(self, x: torch.Tensor) -> torch.Tensor: + """The local vocab-shard GEMM; the seam a quantized head overrides.""" + if self.tied_embedding is not None: + return F.linear(x, self.tied_embedding.weight, self.bias) + return self.quant_method.apply(self, x) + @nvtx_annotate("LMHead") def forward(self, x: torch.Tensor) -> torch.Tensor: ctx = get_global_ctx() @@ -131,10 +166,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: x = x[indices].contiguous() del indices - if self.tied_embedding is not None: - logits = F.linear(x, self.tied_embedding.weight, self.bias) - else: - logits = self.quant_method.apply(self, x) + logits = self._logits(x) if self.tp_size == 1: return logits input_shape = logits.shape diff --git a/python/freetoken/layers/fp8_dynamic.py b/python/freetoken/layers/fp8_dynamic.py new file mode 100644 index 0000000000..caeebb3221 --- /dev/null +++ b/python/freetoken/layers/fp8_dynamic.py @@ -0,0 +1,211 @@ +"""Load-time per-tensor FP8 dense linear (W8A8 through cuBLASLt ``torch._scaled_mm``). + +The weight is e4m3 ``[out_local, in_local]`` with one fp32 ``weight_scale`` (shape ``()``), +produced by the model's weight reader from the bf16 checkpoint tensor (after TP sharding). +The activation is quantized per call with a dynamic per-tensor scale: one Triton launch (a +single program walking the tensor) for tiny inputs, two above ``_SPLIT_MIN_ELEMENTS`` (a +block-parallel amax, then a reduce+cast). Both compute the same arithmetic, so the split is +a pure speed choice. No host sync anywhere (the scales stay on the device), so the decode +path is CUDA-graph safe; the branch between the two paths is on the tensor *shape*, never +on its values. + +Measured on an RTX 6000 Ada (sm_89, torch 2.11.0+cu130) at the qwen4_exp TP=2 shapes, per +decode step per rank over 48 layers: bf16 cuBLAS 3.2-3.4 ms, this path 1.9 ms at M=1/8/16. +Requires sm_89+ (``_scaled_mm``'s floor; Ampere has no FP8 tensor cores). +""" + +from __future__ import annotations + +from typing import List + +import torch +import triton +import triton.language as tl +from freetoken.distributed import DistributedCommunicator, get_tp_info +from freetoken.utils import div_even + +from .base import BaseOP +from .embedding import ParallelLMHead + +FP8 = torch.float8_e4m3fn +E4M3_MAX = 448.0 +_BLOCK = 4096 +# One program per block above this, two launches (partial amax, then reduce+cast); below it +# a single program walks the whole tensor. The one-program kernel is serial over the tensor, +# so it must not be given the [T, hc_count*hidden] hyper-connection activations. +_SPLIT_MIN_ELEMENTS = 16384 +# Partial-amax programs, and so the constexpr width of the fold in pass 2. Capping it (rather +# than letting it follow the input) keeps BOTH split kernels to ONE compiled Triton variant: +# the partial count is a runtime argument, so a server that sees a new prompt length does not +# compile a new kernel mid-generation. +_MAX_PARTS = 512 + + +@triton.jit +def _quant_fused_kernel(x_ptr, out_ptr, scale_ptr, n, BLOCK: tl.constexpr): + """One program: max|x| over the tensor, then the cast. Tiny inputs only (see SPLIT_MIN).""" + acc = tl.zeros([BLOCK], dtype=tl.float32) + for start in range(0, n, BLOCK): + offs = start + tl.arange(0, BLOCK) + v = tl.load(x_ptr + offs, mask=offs < n, other=0.0).to(tl.float32) + acc = tl.maximum(acc, tl.abs(v)) + amax = tl.maximum(tl.max(acc, axis=0), 1e-12) + tl.store(scale_ptr, amax / 448.0) + inv = 448.0 / amax + for start in range(0, n, BLOCK): + offs = start + tl.arange(0, BLOCK) + mask = offs < n + v = tl.load(x_ptr + offs, mask=mask, other=0.0).to(tl.float32) * inv + v = tl.minimum(tl.maximum(v, -448.0), 448.0) + tl.store(out_ptr + offs, v.to(tl.float8e4nv), mask=mask) + + +@triton.jit +def _amax_partial_kernel(x_ptr, part_ptr, n, nprog, BLOCK: tl.constexpr): + """max|x| -> one fp32 partial per program (pass 1). Grid-strided, so ``nprog`` bounds the + partial count however large the tensor is.""" + pid = tl.program_id(0) + acc = tl.zeros([BLOCK], dtype=tl.float32) + for start in range(pid * BLOCK, n, nprog * BLOCK): + offs = start + tl.arange(0, BLOCK) + v = tl.load(x_ptr + offs, mask=offs < n, other=0.0).to(tl.float32) + acc = tl.maximum(acc, tl.abs(v)) + tl.store(part_ptr + pid, tl.max(acc, axis=0)) + + +@triton.jit +def _reduce_cast_kernel( + x_ptr, out_ptr, part_ptr, scale_ptr, n, nprog, + MAX_PARTS: tl.constexpr, BLOCK: tl.constexpr, +): + """Pass 2: fold the partials to the tensor amax, then cast this program's block. + + Every program repeats the (nprog-element, L2-resident) fold instead of paying a third + launch for it; program 0 also publishes the scale. The arithmetic is the one + ``_quant_fused_kernel`` uses, so both paths quantize a given tensor identically. + """ + poffs = tl.arange(0, MAX_PARTS) + amax = tl.max(tl.load(part_ptr + poffs, mask=poffs < nprog, other=0.0), axis=0) + amax = tl.maximum(amax, 1e-12) + pid = tl.program_id(0) + if pid == 0: + tl.store(scale_ptr, amax / 448.0) + offs = pid * BLOCK + tl.arange(0, BLOCK) + mask = offs < n + v = tl.load(x_ptr + offs, mask=mask, other=0.0).to(tl.float32) * (448.0 / amax) + v = tl.minimum(tl.maximum(v, -448.0), 448.0) + tl.store(out_ptr + offs, v.to(tl.float8e4nv), mask=mask) + + +def quant_per_tensor(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """``(x_fp8, scale)`` with ``x ~= x_fp8 * scale``; ``x`` contiguous, ``scale`` fp32 ``()``.""" + n = x.numel() + out = torch.empty_like(x, dtype=FP8) + scale = torch.empty((), dtype=torch.float32, device=x.device) + if n <= _SPLIT_MIN_ELEMENTS: + # below this, the second launch costs more than the parallelism buys + _quant_fused_kernel[(1,)](x, out, scale, n, BLOCK=_BLOCK, num_warps=8) + return out, scale + nblock = triton.cdiv(n, _BLOCK) + nprog = min(nblock, _MAX_PARTS) + part = torch.empty(nprog, dtype=torch.float32, device=x.device) + _amax_partial_kernel[(nprog,)](x, part, n, nprog, BLOCK=_BLOCK, num_warps=8) + _reduce_cast_kernel[(nblock,)]( + x, out, part, scale, n, nprog, MAX_PARTS=_MAX_PARTS, BLOCK=_BLOCK, num_warps=8, + ) + return out, scale + + +def fp8_dynamic_linear( + x: torch.Tensor, weight: torch.Tensor, weight_scale: torch.Tensor +) -> torch.Tensor: + """``x @ (weight * weight_scale)^T`` in W8A8; ``weight`` [N, K] e4m3 row-major, whose ``.t()`` + is the column-major operand cuBLASLt wants (a stride change, never a copy).""" + *lead, k = x.shape + x2 = x.reshape(-1, k).contiguous() + x8, scale = quant_per_tensor(x2) + y = torch._scaled_mm( + x8, weight.t(), scale_a=scale, scale_b=weight_scale, out_dtype=x.dtype + ) + return y.reshape(*lead, weight.shape[0]) + + +class Fp8DynamicLinear(BaseOP): + """Per-tensor FP8 linear over the local shard; ``all_reduce`` adds the TP sum (row-parallel).""" + + def __init__(self, local_isize: int, local_osize: int, *, all_reduce: bool = False): + assert local_isize % 16 == 0 and local_osize % 16 == 0, ( + local_isize, + local_osize, + ) + self.local_input_size = local_isize + self.local_output_size = local_osize + self.weight = torch.empty(local_osize, local_isize, dtype=FP8) + self.weight_scale = torch.empty((), dtype=torch.float32) + self._comm = ( + DistributedCommunicator() if all_reduce and get_tp_info().size > 1 else None + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + y = fp8_dynamic_linear(x, self.weight, self.weight_scale) + if self._comm is not None: + y = self._comm.all_reduce(y) + return y + + +class Fp8DynamicColMerged(Fp8DynamicLinear): + """Drop-in for ``LinearColParallelMerged``: one weight concatenating several projections + along the output dim; the caller splits the output by the local sizes as before.""" + + def __init__( + self, + input_size: int, + output_sizes: List[int], + local_output_sizes: List[int] | None = None, + ): + tp = get_tp_info() + if local_output_sizes is None: + local_output_sizes = [div_even(size, tp.size) for size in output_sizes] + self.output_sizes = list(output_sizes) + self.local_output_sizes = list(local_output_sizes) + super().__init__(input_size, sum(local_output_sizes)) + + +class Fp8DynamicRowParallel(Fp8DynamicLinear): + """Drop-in for ``LinearOProj`` / ``LinearRowParallel``: the input dim is sharded, the + all-reduce runs after the local GEMM (each rank scales its own shard).""" + + def __init__(self, input_size: int, output_size: int): + super().__init__( + div_even(input_size, get_tp_info().size), output_size, all_reduce=True + ) + + +class Fp8ParallelLMHead(ParallelLMHead): + """``ParallelLMHead`` whose vocab shard is a per-tensor e4m3 weight (FREETOKEN_FP8_LMHEAD=1). + + Only the GEMM changes; the vocab-parallel all_gather of the logits above it is untouched. + Untied embeddings only -- a tied head shares the bf16 embedding table, which the lookup + side still reads as bf16. + """ + + def __init__(self, num_embeddings: int, embedding_dim: int): + super().__init__(num_embeddings, embedding_dim) + self.weight = torch.empty(self.num_embeddings_tp, embedding_dim, dtype=FP8) + self.weight_scale = torch.empty((), dtype=torch.float32) + + def _logits(self, x: torch.Tensor) -> torch.Tensor: + y = fp8_dynamic_linear(x, self.weight, self.weight_scale) + return y if self.bias is None else y + self.bias + + +__all__ = [ + "FP8", + "E4M3_MAX", + "Fp8ParallelLMHead", + "Fp8DynamicColMerged", + "Fp8DynamicLinear", + "Fp8DynamicRowParallel", + "fp8_dynamic_linear", + "quant_per_tensor", +] diff --git a/python/freetoken/layers/linear.py b/python/freetoken/layers/linear.py index d8a3e8619f..40ebe934c0 100644 --- a/python/freetoken/layers/linear.py +++ b/python/freetoken/layers/linear.py @@ -83,13 +83,17 @@ def __init__( input_size: int, output_sizes: List[int], has_bias: bool, + local_output_sizes: List[int] | None = None, *, quant_config: QuantConfig | None = None, prefix: str = "", ): - # check that all output sizes are divisible by tp_size + # check that all output sizes are divisible by tp_size (a caller that replicates + # GQA kv heads across ranks passes the per-rank sizes explicitly) tp_info = get_tp_info() - tp_output_sizes = [div_even(size, tp_info.size) for size in output_sizes] + if local_output_sizes is None: + local_output_sizes = [div_even(size, tp_info.size) for size in output_sizes] + tp_output_sizes = local_output_sizes output_size = sum(output_sizes) tp_output_size = sum(tp_output_sizes) super().__init__( diff --git a/python/freetoken/layers/moe.py b/python/freetoken/layers/moe.py index 710a44c79a..29dbddfdff 100644 --- a/python/freetoken/layers/moe.py +++ b/python/freetoken/layers/moe.py @@ -24,6 +24,12 @@ # default. Set FREETOKEN_HYBRID_OVERLAP=0 to force the serial path (CPU sync before the # GPU work) -- a measurement-only escape hatch to A/B the overlap benefit. _HYBRID_OVERLAP = os.getenv("FREETOKEN_HYBRID_OVERLAP", "1") != "0" +# Prefills of at most this many tokens load only their routed experts (the decode path's +# on-demand LRU) instead of streaming every expert of every layer; 0 = off. +_SMALL_PREFILL_TOKENS = int(os.getenv("FREETOKEN_MOE_SMALL_PREFILL_TOKENS", "0")) +# ids per LRU ensure in that path: the kernel's block is next_pow2(ids) x next_pow2(slots), and a +# block much wider than decode's spills to local memory whose launch reservation can OOM +_ENSURE_CHUNK_IDS = 32 class MoELayer(BaseOP): @@ -188,6 +194,17 @@ def __init__( prefix=prefix, ) self.offload_cache: OffloadMoeCache | None = None + # FREETOKEN_MOE_PREFETCH: (next MoE layer's router op, its layer id, its budget); the + # underscore keeps the borrowed router out of this layer's state dict + self._lookahead: tuple[BaseOP, int, int] | None = None + + def set_lookahead(self, gate: BaseOP, target_layer: int, budget: int) -> None: + """Let this layer's GPU decode predict ``target_layer``'s experts with that layer's router.""" + self._lookahead = (gate, target_layer, budget) + + @property + def skips_inactive_slots(self) -> bool: + return self.quant_method is not None and self.quant_method.kernel.supports_inactive_slots def forward( self, @@ -235,7 +252,19 @@ def decode_forward( topk=self.top_k, renormalize=self.renormalize, ) - return self._decode_routed(hidden_states, topk_weights, topk_ids) + return self._decode_routed(hidden_states, topk_weights, topk_ids, router_logits=router_logits) + + def decode_side_applies(self) -> bool: + """Whether FREETOKEN_MOE_COPY_OVERLAP applies to this layer's next forward: a GPU decode on a + cache with a side stream (not prefill, CPU/hybrid decode or TP > 1); see decode_side.""" + cache = self.offload_cache + return not ( + cache is None + or cache.decode_copy_stream is None + or self.tp_size > 1 + or get_global_ctx().batch.is_prefill + or cache.is_cpu_layer(self.layer_id) + ) def prefill_forward( self, @@ -264,6 +293,7 @@ def _decode_routed( hidden_states: torch.Tensor, topk_weights: torch.Tensor, topk_ids: torch.Tensor, + router_logits: torch.Tensor | None = None, ) -> torch.Tensor: """On-demand load: ``ensure_experts`` rewrites ``topk_ids`` into cache slot ids in place (loading missing experts), then the GEMM reads the full slot @@ -283,18 +313,57 @@ def _decode_routed( return executor.decode(self.layer_id, hidden_states, topk_weights, topk_ids) if cache.decode_target == "hybrid": return self._decode_hybrid(cache, hidden_states, topk_weights, topk_ids) - cache.ensure_experts(self.layer_id, topk_ids) + prefetch, verify = cache.prefetch, cache.verify + if verify is not None: + verify.begin(cache, self.layer_id, topk_ids, prefetch and prefetch.plan_of(self.layer_id)) + cache.ensure_experts(self.layer_id, topk_ids, router_logits=router_logits) + target = None + if prefetch is not None: + # on: install the next layer's prediction (forked at this block's start) on this stream, the slot + # maps' only writer; before this layer's copy join, so a late prefetch copy never delays it + target = prefetch.install(self._lookahead and self._lookahead[1]) + prefetch.join_copy(self.layer_id) cache.copy_missing() - return self._expert_gemm( + if prefetch is not None: + prefetch.join_and_count(self.layer_id, topk_ids, cache.id_of_slot, cache.num_indices) + if target is not None: + prefetch.issue_copy(target) + if verify is not None: + # after issue_copy: the next layer's prefetch copy starts when it would without the checks + verify.before_gemm(cache, topk_ids, None if target is None else prefetch.plan_of(target)) + out = self._expert_gemm( cache, hidden_states, topk_weights, topk_ids, - views=cache.bank_views(), + views=cache.pool_bank_views(self.layer_id), n=None, alphas=cache.alphas_for_slots(self.layer_id), is_prefill=False, ) + if verify is not None: + verify.after_gemm(cache) + return out + + def fork_lookahead(self, hidden_states: torch.Tensor) -> None: + """FREETOKEN_MOE_PREFETCH: at the start of this layer's MoE block, predict the next layer's + experts from this block's router input on the predictor stream, beside this layer's router + and ensure. Only on the GPU decode path that ``_decode_routed`` takes and not under TP, and + never toward a CPU-decoded layer, whose ensure would not use it.""" + cache = self.offload_cache + if ( + self._lookahead is None + or cache is None + or cache.prefetch is None + or self.tp_size > 1 + or get_global_ctx().batch.is_prefill + or cache.is_cpu_layer(self.layer_id) + or cache.decode_target == "hybrid" + ): + return + gate, target, budget = self._lookahead + if not cache.is_cpu_layer(target): + cache.prefetch.fork(target, hidden_states, gate, cache.slot_for_id[target], budget) def _decode_hybrid( self, @@ -308,7 +377,7 @@ def _decode_hybrid( The CPU pool is kicked off (``decode_submit``) before the GPU PCIe fetch + GEMM so the CPU overflow GEMV runs concurrently with the GPU work. Capture-safe: the - routing split is device-side elementwise and the CPU submit/sync are host nodes. + routing split runs inside the hybrid ensure kernel and the CPU submit/sync are host nodes. Each route is computed exactly once -- CPU-assigned routes reach the GPU as slot -1 with weight 0 (a kernel with ``supports_inactive_slots`` skips them without a read; any other reads slot 0, whose storage is a payload or the bank's zero fill) and the CPU ids are -1 for @@ -316,14 +385,11 @@ def _decode_hybrid( """ executor = cache.cpu_executor assert executor is not None, "CPU MoE executor was not initialized" - raw_ids = topk_ids.clone() - cache.ensure_experts_hybrid(self.layer_id, topk_ids) - on_gpu = topk_ids >= 0 - cpu_ids = torch.where(on_gpu, -1, raw_ids).contiguous() - gpu_w = torch.where(on_gpu, topk_weights, 0.0).contiguous() - gpu_slots = topk_ids - if cache.collect_stats: - cache.record_decode_stats_hybrid(self.layer_id) + # the kernel reads every route's slot: give the inactive ones slot 0, which holds a + # complete payload or the bank's zero fill (never uninitialized storage), times weight 0 + # with several quants this holds only because the cache zeroes staged slot 0 rows (zero_inactive_route_slots) + clamp = not self.skips_inactive_slots + gpu_slots, gpu_w, cpu_ids = cache.ensure_experts_hybrid(self.layer_id, topk_ids, topk_weights, clamp=clamp) pending = executor.decode_submit(self.layer_id, hidden_states, topk_weights, cpu_ids) # Measurement knob: FREETOKEN_HYBRID_OVERLAP=0 syncs the CPU pool *before* the @@ -333,22 +399,19 @@ def _decode_hybrid( ) cache.copy_missing() - if not (self.quant_method is not None and self.quant_method.kernel.supports_inactive_slots): - # the kernel reads every route's slot: give the inactive ones slot 0, which holds a - # complete payload or the bank's zero fill (never uninitialized storage), times weight 0 - gpu_slots = gpu_slots.clamp_min(0) gpu_routed = self._expert_gemm( cache, hidden_states, gpu_w, gpu_slots, - views=cache.bank_views(), + views=cache.pool_bank_views(self.layer_id), n=None, alphas=cache.alphas_for_slots(self.layer_id), is_prefill=False, ) - cpu_routed = cpu_routed_early if not _HYBRID_OVERLAP else executor.decode_sync(pending) - return gpu_routed + cpu_routed + if not _HYBRID_OVERLAP: + return gpu_routed + cpu_routed_early + return executor.decode_sync(pending, add_to=gpu_routed) def _prefill_routed( self, @@ -363,14 +426,22 @@ def _prefill_routed( position == expert id, so the routing ids pass through unmapped.""" cache = self.offload_cache assert cache is not None - if cache.loads_prefill_on_demand(hidden_states.shape[0]): - cache.ensure_experts(self.layer_id, topk_ids, counts_as_decode=False) - cache.copy_missing() + if ( + 0 < hidden_states.shape[0] <= max(cache.prefill_demand_tokens, _SMALL_PREFILL_TOKENS) + and cache.decode_target in ("gpu", "hybrid") + and not cache.is_unpinned_layer(self.layer_id) + ): + # A short extension (an agent turn over a cached prefix) touches a fraction of each + # layer's experts; fetching those beats streaming all of them (~68 GB here). It runs + # the prefill GEMM over the slot cache (#676), so it matches the streamed layer bit for + # bit. An unpinned layer's copy_missing presumes position == expert id, so it streams + # below; unpinned layers only exist with the prefill overlap off, so no layer can enter + # the overlap double buffer after an earlier layer skipped begin_prefill. return self._expert_gemm( cache, hidden_states, topk_weights, - topk_ids, + self._ensure_unique(cache, topk_ids), views=cache.bank_views(), n=cache.cache_size, alphas=cache.alphas_for_slots(self.layer_id), @@ -392,16 +463,39 @@ def _prefill_routed( return out cache.materialize_layer(self.layer_id) cache.copy_missing() - return self._expert_gemm( + out = self._expert_gemm( cache, hidden_states, topk_weights, topk_ids, - views=cache.bank_views(self.num_experts), + views=cache.stage_views(self.layer_id, 0), n=self.num_experts, alphas=cache.alphas_for_layer(self.layer_id), is_prefill=True, ) + cache.release_prefill_layer(self.layer_id) + return out + + def _ensure_unique(self, cache: OffloadMoeCache, topk_ids: torch.Tensor) -> torch.Tensor: + """Make this layer's routed experts resident and return ``topk_ids`` mapped to slots. + + The LRU kernel is sized for decode: its block is next_pow2(ids) x next_pow2(slots), so a + prefill's tokens x top_k ids overflow Triton's element limit, and even a block under it + can fail to launch on a full GPU. It gets the unique ids (at most num_experts) in + decode-width chunks; each chunk's misses are copied before the next ensure reuses the plan.""" + uniq, inverse = torch.unique(topk_ids.reshape(-1), return_inverse=True) + n = uniq.numel() + # Each ensure stamps a newer LRU step, so a later chunk evicts the earlier chunks' slots + # only once every other slot is gone; the engine's num_experts slot floor rules that out. + assert n <= cache.cache_size, f"{n} routed experts exceed the {cache.cache_size}-slot cache" + uniq = uniq.to(torch.int32) + # scored eviction ranks an earlier chunk's fresh installs coldest, so pin them explicitly + pin_since = cache.step.clone() if cache.cache_policy_id else None + for start in range(0, n, _ENSURE_CHUNK_IDS): + part = uniq[start : start + _ENSURE_CHUNK_IDS] # a contiguous view: ensure rewrites it in place + cache.ensure_experts(self.layer_id, part, update_state=False, pin_since=pin_since, counts_as_decode=False) # not a decode step + cache.copy_missing() + return uniq[inverse].view_as(topk_ids) def _wait_prefill_overlap(self, cache: OffloadMoeCache) -> tuple[torch.Tensor, ...]: """Double-buffer choreography for this layer's overlap prefill: kick off the diff --git a/python/freetoken/layers/quantization/moe/nvfp4.py b/python/freetoken/layers/quantization/moe/nvfp4.py index 1cd1a1ac73..27236d2612 100644 --- a/python/freetoken/layers/quantization/moe/nvfp4.py +++ b/python/freetoken/layers/quantization/moe/nvfp4.py @@ -41,14 +41,20 @@ class TritonNvfp4MoEKernel(MoEKernel): cpu_format = "nvfp4" def unusable_reason(self, cfg: MoEConfig) -> str | None: - reason = self._common_reject(cfg, tp_ok=False, cpu_ok=True, plain_silu_only=False) + # tp_ok: the source stream is sliced along the intermediate axis per rank + # (nvfp4_banks._tp_shard) and this kernel sizes its banks from + # cfg.local_intermediate, so a rank holds and reads only its own half. The routed + # output is then a partial sum, which the MoE layer reduces (_maybe_all_reduce, or + # one combined all-reduce in qwen4_exp's block). marlin and b12x stay off: their + # pack() repacks the rows and neither has been checked against a sharded bank. + reason = self._common_reject(cfg, tp_ok=True, cpu_ok=True, plain_silu_only=False) if reason: return reason reason = gated_epilogue_reason(cfg) return f"triton nvfp4 MoE kernel: {reason}" if reason else None def layout(self, cfg: MoEConfig) -> dict[str, BankSpec]: - i, h = cfg.intermediate, cfg.hidden + i, h = cfg.local_intermediate, cfg.hidden return { "gate_up": BankSpec((2 * i, h // 2), torch.uint8), "gate_up_scale": BankSpec((2 * i, h // GROUP), FP8), @@ -61,7 +67,7 @@ def layout(self, cfg: MoEConfig) -> dict[str, BankSpec]: def pack(self, pieces, cfg: MoEConfig, out): out["gate_up"].copy_(fused_piece(pieces, "gate_up")) out["gate_up_scale"].copy_(fused_piece(pieces, "gate_up_scale")) - out["gate_up_global"].copy_(fused_global(pieces, cfg.intermediate)) + out["gate_up_global"].copy_(fused_global(pieces, cfg.local_intermediate)) out["down"].copy_(pieces["down"]) out["down_scale"].copy_(pieces["down_scale"]) out["down_global"].copy_(global_rows(pieces["down_global"], cfg.hidden)) @@ -243,7 +249,7 @@ def worth_it(self, cfg: MoEConfig) -> bool: return (8, 0) <= backend.device_capability() < (10, 0) def layout(self, cfg: MoEConfig) -> dict[str, BankSpec]: - i, h = cfg.intermediate, cfg.hidden + i, h = cfg.local_intermediate, cfg.hidden return { "gate_up": BankSpec((h // GROUP, 4 * i), torch.int32), "gate_up_scale": BankSpec((h // GROUP, 2 * i), FP8), @@ -254,7 +260,7 @@ def layout(self, cfg: MoEConfig) -> dict[str, BankSpec]: } def pack(self, pieces, cfg: MoEConfig, out): - i, h = cfg.intermediate, cfg.hidden + i, h = cfg.local_intermediate, cfg.hidden device = torch.device("cuda") gu, gus, gug = fused_piece(pieces, "gate_up"), fused_piece(pieces, "gate_up_scale"), fused_global(pieces, i) dn, dns, dng = pieces["down"], pieces["down_scale"], global_rows(pieces["down_global"], h) @@ -507,7 +513,7 @@ def unusable_reason(self, cfg: MoEConfig) -> str | None: def worth_it(self, cfg: MoEConfig) -> bool: # NOTE: never auto-selected. flashinfer's cute launcher indexes each bank with int32 element offsets, so the GPU slot cache is capped at (2^31 - 1) / elements-per-slot (about 1000 slots for GLM-5.3-Flash's 12960 experts); until flashinfer lifts that, b12x stays behind triton in the table; the selection does not weigh the cap, cache-auto clamps to slot_limit() and OffloadMoeCache refuses a larger cache. - return cfg.intermediate >= B12X_MIN_INTERMEDIATE + return cfg.local_intermediate >= B12X_MIN_INTERMEDIATE def slot_limit(self, cfg: MoEConfig) -> int | None: # the cute launcher indexes each bank with int32 element offsets @@ -516,7 +522,7 @@ def slot_limit(self, cfg: MoEConfig) -> int | None: def layout(self, cfg: MoEConfig) -> dict[str, BankSpec]: # flashinfer's prepared tiles: the Marlin shapes, byte-identical to the native rows - i, h = cfg.intermediate, cfg.hidden + i, h = cfg.local_intermediate, cfg.hidden return { "gate_up": BankSpec((h // GROUP, 4 * i), torch.int32), "gate_up_scale": BankSpec((h // GROUP, 2 * i), FP8), @@ -529,7 +535,7 @@ def layout(self, cfg: MoEConfig) -> dict[str, BankSpec]: def pack(self, pieces, cfg: MoEConfig, out): from flashinfer.fused_moe.cute_dsl.blackwell_sm12x.moe_w4a16_prepare import prepare_w4a16_packed_weights - i, h = cfg.intermediate, cfg.hidden + i, h = cfg.local_intermediate, cfg.hidden device = torch.device("cuda") gu = fused_piece(pieces, "gate_up").to(device) gug = fused_global(pieces, i).to(device).float() diff --git a/python/freetoken/message/frontend.py b/python/freetoken/message/frontend.py index 24725567b9..d70453e910 100644 --- a/python/freetoken/message/frontend.py +++ b/python/freetoken/message/frontend.py @@ -45,6 +45,9 @@ class UserReply(BaseFrontendMsg): swa_total_tokens: int = 0 # Bytes the engine process holds on the GPU (torch reserved pool). 0 when not reported. gpu_mem_bytes: int = 0 + # Scheduler-measured prefill span (see DetokenizeMsg.prefill_ms). Arrives once, on the + # reply carrying the request's first generated token; 0.0 on every other reply. + prefill_ms: float = 0.0 # Set (with finished=True) when a request failed before producing output — e.g. a chat # template that the tokenizer cannot render, or a prompt that exceeds the KV budget the # scheduler can serve. Carries a human-readable reason. Without this, such a request would @@ -56,6 +59,12 @@ class UserReply(BaseFrontendMsg): finish_reason: str | None = None # The stop string that ended generation (Anthropic reports it as stop_reason='stop_sequence'). matched_stop: str | None = None + # On the finished reply: tokens generated up to and including the reasoning end tag + # (0 when the model emitted none) -> usage.completion_tokens_details.reasoning_tokens. + reasoning_tokens: int = 0 + # Neutral sampled-token logprobs entry for this token (see + # tokenizer.detokenize.build_logprobs_entry); None when the request did not ask. + logprobs: dict | None = None @dataclass diff --git a/python/freetoken/message/tokenizer.py b/python/freetoken/message/tokenizer.py index d52d658e63..1b3b49eeeb 100644 --- a/python/freetoken/message/tokenizer.py +++ b/python/freetoken/message/tokenizer.py @@ -36,6 +36,15 @@ class DetokenizeMsg(BaseTokenizerMsg): # The request's stop strings (None when it has none), so the detokenizer can hold back # a trailing partial-stop prefix instead of streaming it and then needing to retract. stop_strs: list[str] | None = None + # Keep the matched stop string in the output (SamplingParams.include_stop_str_in_output). + keep_stop_str: bool = False + # Decode this request with skip_special_tokens (SamplingParams.skip_special_tokens). + skip_special_tokens: bool = False + # Sampled-token logprobs (None unless the request asked): the chosen token's raw + # logprob and the top alternatives, already cut to this request's top_logprobs. + chosen_logprob: float | None = None + top_ids: list[int] | None = None + top_logprobs: list[float] | None = None # KV page-pool usage snapshot at this step (not-evictable used/total), passed # through to the frontend for the shell status bar. 0/0 for owned-KV models. kv_used_pages: int = 0 @@ -48,6 +57,11 @@ class DetokenizeMsg(BaseTokenizerMsg): swa_total_tokens: int = 0 # Bytes this engine process holds on the GPU (torch reserved pool). 0 on CPU. gpu_mem_bytes: int = 0 + # Scheduler-measured prefill span for this request: admission (the first prefill batch + # was prepared) to the token sampled off its last prefill chunk. Set on that first token + # only, 0.0 on every later one -- it is the one span a client cannot derive from HTTP + # timestamps, which is why it rides the wire instead of being re-estimated frontend-side. + prefill_ms: float = 0.0 @dataclass diff --git a/python/freetoken/mm/media.py b/python/freetoken/mm/media.py index 5c3c4ef789..a51e5bc98b 100644 --- a/python/freetoken/mm/media.py +++ b/python/freetoken/mm/media.py @@ -1,6 +1,8 @@ """Client-supplied image handling: ref collection, the accept gate, and byte fetching. Raises plain ValueError on bad input; the server layer maps it to its wire error type. +The message names the class of failure only: no local path, allowlist root, errno or +library text reaches the client (that detail goes to the server log). """ from __future__ import annotations @@ -8,10 +10,15 @@ import base64 from typing import TYPE_CHECKING, Any +from freetoken.utils import init_logger + if TYPE_CHECKING: from freetoken.server.args import ServerArgs +logger = init_logger(__name__) + _MAX_IMAGE_BYTES = 32 << 20 +_MAX_REDIRECTS = 5 def collect_image_refs(messages: list[dict[str, Any]]) -> list[dict[str, Any]]: @@ -56,6 +63,57 @@ def _check_media_domain(url: str, config: ServerArgs) -> None: ) +def _check_redirect_target(url: str, config: ServerArgs) -> None: + """A redirect hop is the remote server's choice, not the client's URL: it must pass the + allowlist like the original, and may never land on a loopback / private / link-local + address (an allowed public host 302-ing into the operator's network is the SSRF shape).""" + import ipaddress + from urllib.parse import urlparse + + if not url.startswith(("http://", "https://")): + raise ValueError("the URL redirects to a non-http(s) location") + try: + _check_media_domain(url, config) + except ValueError: + raise ValueError("the URL redirects outside the allowed domains") from None + host = (urlparse(url).hostname or "").rstrip(".").lower() + if host == "localhost" or host.endswith(".localhost"): + raise ValueError("the URL redirects to a private or loopback address") + try: + ip = ipaddress.ip_address(host) + except ValueError: + return # a name; the allowlist above is the gate for names (no DNS resolution here) + if isinstance(ip, ipaddress.IPv6Address) and ip.ipv4_mapped is not None: + ip = ip.ipv4_mapped + if not ip.is_global or ip.is_multicast: # covers loopback, private, link-local, CGNAT, unspecified + raise ValueError("the URL redirects to a private or loopback address") + + +async def _fetch_remote_media(url: str, config: ServerArgs) -> bytes: + """GET an http(s) ref, following redirects one hop at a time so every target is re-checked.""" + import httpx + + _check_media_domain(url, config) + async with httpx.AsyncClient(timeout=30.0, follow_redirects=False) as client: + for _ in range(_MAX_REDIRECTS + 1): + async with client.stream("GET", url) as resp: + if resp.next_request is not None: # a 3xx with a Location: httpx built the hop + url = str(resp.next_request.url) + _check_redirect_target(url, config) + continue + resp.raise_for_status() + declared = resp.headers.get("content-length", "") + if declared.isdigit() and int(declared) > _MAX_IMAGE_BYTES: + raise ValueError(f"image exceeds {_MAX_IMAGE_BYTES} bytes") + buf = bytearray() + async for chunk in resp.aiter_bytes(): + buf += chunk + if len(buf) > _MAX_IMAGE_BYTES: + raise ValueError(f"image exceeds {_MAX_IMAGE_BYTES} bytes") + return bytes(buf) + raise ValueError(f"the URL redirects more than {_MAX_REDIRECTS} times") + + def _load_local_media(url: str, config: ServerArgs) -> bytes: """Read a file:// ref; the resolved path must be a strict subpath of --allowed-local-media-path.""" from pathlib import Path @@ -70,41 +128,53 @@ def _load_local_media(url: str, config: ServerArgs) -> bytes: # resolve() follows symlinks, so a link inside the root escaping it is rejected too resolved = filepath.resolve() if Path(root).resolve() not in resolved.parents: - raise ValueError( - f"the file path {filepath} must be a subpath of " - f"--allowed-local-media-path {root}" - ) + raise ValueError("the file path must be a subpath of --allowed-local-media-path") return resolved.read_bytes() +async def _fetch_one(ref: dict[str, Any], config: ServerArgs) -> bytes: + """One ref to bytes. Our own ValueErrors are already client-safe and pass through; anything + a library raises is reduced to its failure class and logged with the detail.""" + data = ref.get("data") or "" + try: + if ref.get("kind") == "b64": + return base64.b64decode(data) + if data.startswith("data:"): + return base64.b64decode(data.split(",", 1)[1]) + except Exception as exc: # noqa: BLE001 -- binascii.Error / IndexError, input-driven + raise ValueError("image data is not valid base64") from exc + if data.startswith(("http://", "https://")): + import httpx + + try: + return await _fetch_remote_media(data, config) + except ValueError: + raise + except httpx.HTTPStatusError as exc: + raise ValueError(f"the image URL returned HTTP {exc.response.status_code}") from exc + except httpx.TimeoutException as exc: + raise ValueError("timed out fetching the image URL") from exc + except Exception as exc: # noqa: BLE001 -- transport / protocol / bad URL + logger.warning("image URL fetch failed: %r", exc) + raise ValueError("the image URL is unreachable") from exc + if data.startswith("file://"): + try: + return _load_local_media(data, config) + except ValueError: + raise + except OSError as exc: + logger.warning("local image read failed: %r", exc) + raise ValueError("the local file is missing or unreadable") from exc + raise ValueError("unsupported image source (expect http(s)/file url or base64)") + + async def fetch_image_bytes(refs: list[dict[str, Any]], config: ServerArgs) -> list[bytes]: """Resolve collected image refs (URLs / base64) to raw bytes, in order.""" out: list[bytes] = [] for ref in refs: - data = ref.get("data") or "" try: - if ref.get("kind") == "b64": - out.append(base64.b64decode(data)) - elif data.startswith("data:"): - out.append(base64.b64decode(data.split(",", 1)[1])) - elif data.startswith(("http://", "https://")): - import httpx - - _check_media_domain(data, config) - # redirect targets are not re-checked against the allowlist - async with httpx.AsyncClient(timeout=30.0, follow_redirects=True) as client: - resp = await client.get(data) - resp.raise_for_status() - if len(resp.content) > _MAX_IMAGE_BYTES: - raise ValueError(f"image exceeds {_MAX_IMAGE_BYTES} bytes") - out.append(resp.content) - elif data.startswith("file://"): - out.append(_load_local_media(data, config)) - else: - raise ValueError( - "unsupported image source (expect http(s)/file url or base64)" - ) - except Exception as exc: # noqa: BLE001 -- input-driven, client-classifiable + out.append(await _fetch_one(ref, config)) + except ValueError as exc: raise ValueError(f"could not load image: {exc}") from exc return out diff --git a/python/freetoken/models/config.py b/python/freetoken/models/config.py index 25bc73f5e4..0d4faae420 100644 --- a/python/freetoken/models/config.py +++ b/python/freetoken/models/config.py @@ -1,4 +1,6 @@ from __future__ import annotations + +import os from dataclasses import dataclass from typing import Any, ClassVar, Dict, List, Literal, Tuple, TypeAlias @@ -7,6 +9,23 @@ # State-dict key prefixes of the vision stack; load_weight drops them when the engine serves text-only. VISION_KEY_PREFIXES = ("vision_tower.", "embed_vision.", "vision_embedder.", "visual.") +_ENV_TRUE = {"1", "true", "yes", "on"} + + +def fp8_dense_enabled() -> bool: + """Load-time FP8 for a model's bf16 attention / GDN projections (opt-in, default OFF): + per-tensor e4m3 weights, per-tensor dynamic activation scale, cuBLASLt W8A8 GEMMs + (``torch._scaled_mm``, sm_89+). ``FREETOKEN_FP8_DENSE=1``.""" + return os.getenv("FREETOKEN_FP8_DENSE", "0").strip().lower() in _ENV_TRUE + + +def fp8_lmhead_enabled() -> bool: + """Load-time FP8 for the lm_head as well (opt-in, default OFF, needs FREETOKEN_FP8_DENSE=1 + for the rest). Separate from :func:`fp8_dense_enabled` because this one moves the logits: + a per-tensor e4m3 vocab matrix changes every sampled token's score, so it carries its own + quality gate. ``FREETOKEN_FP8_LMHEAD=1``.""" + return os.getenv("FREETOKEN_FP8_LMHEAD", "0").strip().lower() in _ENV_TRUE + def detect_expert_quant(hf_config: Any) -> str: """Routed-expert quantization from a checkpoint's ``quantization_config``: ``"nvfp4"`` for diff --git a/python/freetoken/models/deepseek_v4/moe.py b/python/freetoken/models/deepseek_v4/moe.py index 35a316f103..65de7f85fe 100644 --- a/python/freetoken/models/deepseek_v4/moe.py +++ b/python/freetoken/models/deepseek_v4/moe.py @@ -120,12 +120,13 @@ def _prefill_routed( cache = self.offload_cache assert cache is not None # unpinned (LOCKED) layers must take the base materialize path: their copy_missing is the whole-layer pageable branch with position == expert id, which ensure_experts's LRU slot remap would contradict (the GEMM would gather other experts' weights) + # one threshold for all layers (the smallest pool): with overlap prefill, the layers of a chunk must all stream or all take the slot path if ( - hidden_states.shape[0] * self.top_k >= self.num_experts + hidden_states.shape[0] * self.top_k >= min(self.num_experts, *cache.pool_sizes) or cache.is_unpinned_layer(self.layer_id) ): return super()._prefill_routed(hidden_states, topk_weights, topk_ids) - cache.ensure_experts(self.layer_id, topk_ids) # in-place expert-id -> slot + cache.ensure_experts(self.layer_id, topk_ids, update_state=False) # in-place expert-id -> slot; not a decode step cache.copy_missing() if cache.collect_stats: cache.record_decode_stats(self.layer_id) @@ -134,7 +135,7 @@ def _prefill_routed( hidden_states, topk_weights, topk_ids, - views=cache.bank_views(), + views=cache.pool_bank_views(self.layer_id), n=None, alphas=cache.alphas_for_slots(self.layer_id), is_prefill=True, diff --git a/python/freetoken/models/glm5_next/kda.py b/python/freetoken/models/glm5_next/kda.py index ba722ae287..4d86892034 100644 --- a/python/freetoken/models/glm5_next/kda.py +++ b/python/freetoken/models/glm5_next/kda.py @@ -142,6 +142,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: mixed = causal_conv1d_varlen( x, self._conv_weight(), pool.conv_states[li], fla.cu_seqlens, fla.cache_indices, fla.has_initial_state, + max_seq_len=fla.max_seq_len, # host-known: no D2H sync sizing the launch ).transpose(0, 1) q, k, v = ( t.reshape(1, total, h, d).to(dtype) diff --git a/python/freetoken/models/nvfp4_banks.py b/python/freetoken/models/nvfp4_banks.py index 2dfdaaa35f..2565dce58f 100644 --- a/python/freetoken/models/nvfp4_banks.py +++ b/python/freetoken/models/nvfp4_banks.py @@ -59,6 +59,36 @@ def _bank_layer(spec: Nvfp4ExpertSourceSpec, layer: int, config) -> int | None: return bank_layer +def _tp_slice() -> tuple[int, int] | None: + """``(i_local, i_lo)``: this rank's slice of the expert intermediate axis, or None at TP=1. + + TP shards every routed expert along I (gate/up output rows, down input columns), so each + rank caches only its own half and the MoE layer all-reduces the partial sum. The expert + kernel sizes its banks from ``MoEConfig.local_intermediate``, which is the same split. + """ + from freetoken.distributed import try_get_tp_info + + tp = try_get_tp_info() + if tp is None or tp.size == 1: + return None + return tp.size, tp.rank + + +def _tp_shard(role: str, tensor: torch.Tensor, inter: int, tp_size: int, tp_rank: int): + """This rank's I-slice of one checkpoint expert tensor. ``_global`` is a per-tensor scalar + and has no I axis; gate/up carry I on the row axis, down on the column axis (halved for the + packed FP4 codes, sixteenthed for the fp8 block scales).""" + if role.endswith("_global"): + return tensor + n = inter // tp_size + assert n % 16 == 0, f"NVFP4 TP shard {n} must cover whole 16-wide scale blocks" + lo = tp_rank * n + if role.startswith("down"): + d = 2 if role == "down" else 16 + return tensor[:, lo // d : (lo + n) // d] + return tensor[lo : lo + n] + + def _kind_suffix(kind: str) -> str: return {"weight": "", "weight_scale": "_scale", "weight_scale_2": "_global"}[kind] @@ -132,7 +162,19 @@ def _parallel(): tensor = _ingest_global(spec, tensor) yield name, tensor - return per_expert_pieces(_parallel() if parallel else _serial(), wanted.get, tensors_per_expert=9) + stream = _parallel() if parallel else _serial() + shard = _tp_slice() + if shard is not None: + tp_size, tp_rank = shard + inter = config.moe_intermediate_size + base = stream + + def _sharded(): + for name, tensor in base: + yield name, _tp_shard(wanted[name][2], tensor, inter, tp_size, tp_rank) + + stream = _sharded() + return per_expert_pieces(stream, wanted.get, tensors_per_expert=9) __all__ = ["Nvfp4ExpertSourceSpec", "iter_nvfp4_expert_pieces"] diff --git a/python/freetoken/models/qwen3_5_moe/gdn.py b/python/freetoken/models/qwen3_5_moe/gdn.py index d5fe2aa8fe..af3fbfbcc4 100644 --- a/python/freetoken/models/qwen3_5_moe/gdn.py +++ b/python/freetoken/models/qwen3_5_moe/gdn.py @@ -90,14 +90,17 @@ def _gate_params(self, a: torch.Tensor, b: torch.Tensor): def _conv_weight(self) -> torch.Tensor: return self.conv1d.weight.squeeze(1) # [conv_dim, kernel] for the fused kernel - def _conv_prefill(self, conv_in, pool, cu_seqlens, cache_indices, has_initial_state) -> torch.Tensor: + def _conv_prefill(self, conv_in, pool, cu_seqlens, cache_indices, has_initial_state, + max_seq_len=None) -> torch.Tensor: """Varlen causal conv (fused sgl_kernel) with silu; reads/updates each request's conv state in place by ``cache_indices`` slot. ``conv_in`` [total, conv_dim]. - ``cu_seqlens`` / ``cache_indices`` / ``has_initial_state`` come from FLAMetadata.""" + ``cu_seqlens`` / ``cache_indices`` / ``has_initial_state`` / ``max_seq_len`` come from + FLAMetadata; the last is host-known so the launch needs no D2H sync.""" li = pool.local_index(self.layer_id) x = conv_in.transpose(0, 1).contiguous() # [conv_dim, total] out = causal_conv1d_varlen(x, self._conv_weight(), pool.conv_states[li], - cu_seqlens, cache_indices, has_initial_state) + cu_seqlens, cache_indices, has_initial_state, + max_seq_len=max_seq_len) return out.transpose(0, 1) # [total, conv_dim] def _conv_decode(self, conv_in: torch.Tensor, table_idx: torch.Tensor, pool) -> torch.Tensor: @@ -166,7 +169,8 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: ) else: mixed = self._conv_prefill( - conv_in, pool, fla.cu_seqlens, fla.cache_indices, fla.has_initial_state) + conv_in, pool, fla.cu_seqlens, fla.cache_indices, fla.has_initial_state, + fla.max_seq_len) # fla chunk handles GQA in-kernel: q/k stay at num_k_heads, v at num_v_heads. qf, kf, vf = torch.split(mixed, [self.key_dim, self.key_dim, self.value_dim], dim=-1) q = qf.reshape(1, total, self.num_k_heads, self.head_k_dim).to(dtype) diff --git a/python/freetoken/models/qwen4_exp/__init__.py b/python/freetoken/models/qwen4_exp/__init__.py index 0fe4a4ffa8..84154b26f4 100644 --- a/python/freetoken/models/qwen4_exp/__init__.py +++ b/python/freetoken/models/qwen4_exp/__init__.py @@ -12,6 +12,7 @@ from .config import parse_config from .model import Qwen4ExpForCausalLM, Qwen4ExpForConditionalGeneration from .weight import ( + check_ftw_weights, ftw_side_files, nvfp4_expert_spec, iter_vision_weights, @@ -24,6 +25,7 @@ from freetoken.models.qwen3_5_moe.weight import iter_expert_pieces __all__ = [ + "check_ftw_weights", "ftw_side_files", "nvfp4_expert_spec", "Qwen4ExpForCausalLM", diff --git a/python/freetoken/models/qwen4_exp/attention.py b/python/freetoken/models/qwen4_exp/attention.py index 8f6b938870..11292e347e 100644 --- a/python/freetoken/models/qwen4_exp/attention.py +++ b/python/freetoken/models/qwen4_exp/attention.py @@ -19,9 +19,17 @@ import torch from freetoken.core import get_global_ctx -from freetoken.layers import BaseOP, GemmaPlusOneRMSNorm, LinearColParallelMerged, LinearReplicated +from freetoken.distributed import get_tp_info +from freetoken.layers import ( + BaseOP, + GemmaPlusOneRMSNorm, + LinearColParallelMerged, + LinearOProj, + LinearReplicated, +) +from freetoken.layers.fp8_dynamic import Fp8DynamicColMerged, Fp8DynamicRowParallel from freetoken.layers.rotary import get_rope -from freetoken.utils import nvtx_annotate +from freetoken.utils import div_even, nvtx_annotate if TYPE_CHECKING: from freetoken.core import Batch @@ -123,15 +131,38 @@ def __init__(self, config: ModelConfig, layer_id: int, *, prefix: str = "") -> N self.head_dim = config.head_dim self.qo_attn_dim = self.num_q * self.head_dim self.kv_attn_dim = self.num_kv * self.head_dim - self._qkv_split = [self.qo_attn_dim * 2, self.kv_attn_dim, self.kv_attn_dim] - self.qkv_proj = LinearColParallelMerged( - config.hidden_size, self._qkv_split, has_bias=False, - quant_config=config.quant, prefix=f"{prefix}.qkv_proj", - ) - self.o_proj = LinearReplicated( - self.qo_attn_dim, config.hidden_size, has_bias=False, - quant_config=config.quant, prefix=f"{prefix}.o_proj", + # TP: q heads split across ranks, kv heads split or (num_kv < tp) replicated; the + # indexer stays replicated so every rank selects the same blocks. + tp = get_tp_info() + self._local_num_q = div_even(self.num_q, tp.size) + self._local_num_kv = div_even(self.num_kv, tp.size, allow_replicate=True) + self._local_qo_dim = self._local_num_q * self.head_dim + self._local_kv_dim = self._local_num_kv * self.head_dim + self._qkv_split = [self._local_qo_dim * 2, self._local_kv_dim, self._local_kv_dim] + qkv_sizes = [self.qo_attn_dim * 2, self.kv_attn_dim, self.kv_attn_dim] + # Load-time per-tensor FP8 applies only where the checkpoint declares no scheme of + # its own; a quantized checkpoint always wins, and its modules go through QuantConfig. + fp8_dyn = config.attn_quant == "fp8_dynamic" and ( + config.quant is None or config.quant.scheme_for(f"{prefix}.qkv_proj") is None ) + if fp8_dyn: # W8A8 through _scaled_mm, weights quantized by the loader + self.qkv_proj = Fp8DynamicColMerged( + config.hidden_size, qkv_sizes, local_output_sizes=self._qkv_split + ) + self.o_proj = Fp8DynamicRowParallel(self.qo_attn_dim, config.hidden_size) + else: + self.qkv_proj = LinearColParallelMerged( + config.hidden_size, qkv_sizes, has_bias=False, + local_output_sizes=self._qkv_split, + quant_config=config.quant, prefix=f"{prefix}.qkv_proj", + ) + # row-parallel, NOT LinearReplicated: under TP>1 qkv_proj is column-parallel, so + # this rank's attention output is its own head slice and o_proj must consume the + # sharded input dim and all-reduce. At TP=1 the two are equivalent. + self.o_proj = LinearOProj( + self.qo_attn_dim, config.hidden_size, has_bias=False, + quant_config=config.quant, prefix=f"{prefix}.o_proj", + ) self.q_norm = GemmaPlusOneRMSNorm(self.head_dim, eps=config.rms_norm_eps) self.k_norm = GemmaPlusOneRMSNorm(self.head_dim, eps=config.rms_norm_eps) rotary = config.rotary_config @@ -149,21 +180,21 @@ def __init__(self, config: ModelConfig, layer_id: int, *, prefix: str = "") -> N @nvtx_annotate("QSA") def forward(self, x: torch.Tensor, batch: Batch) -> torch.Tensor: qg, k, v = self.qkv_proj.forward(x).split(self._qkv_split, dim=-1) - qg = qg.view(-1, self.num_q, self.head_dim * 2) + qg = qg.view(-1, self._local_num_q, self.head_dim * 2) q = qg[..., : self.head_dim].contiguous() - gate = qg[..., self.head_dim :].reshape(-1, self.qo_attn_dim) - k = k.contiguous().view(-1, self.num_kv, self.head_dim) + gate = qg[..., self.head_dim :].reshape(-1, self._local_qo_dim) + k = k.contiguous().view(-1, self._local_num_kv, self.head_dim) v = v.contiguous() self.q_norm.forward_inplace(q) self.k_norm.forward_inplace(k) q, k = self.rotary.forward( - batch.get_attn_positions(), q.view(-1, self.qo_attn_dim), k.view(-1, self.kv_attn_dim) + batch.get_attn_positions(), q.view(-1, self._local_qo_dim), k.view(-1, self._local_kv_dim) ) index = self.indexer.forward(x) o = get_global_ctx().attn_backend.qsa_forward( - q.view(-1, self.num_q, self.head_dim), k, v, index, self.layer_id, batch + q.view(-1, self._local_num_q, self.head_dim), k, v, index, self.layer_id, batch ) - gated = o.reshape(-1, self.qo_attn_dim) * torch.sigmoid(gate) + gated = o.reshape(-1, self._local_qo_dim) * torch.sigmoid(gate) return self.o_proj.forward(gated) diff --git a/python/freetoken/models/qwen4_exp/config.py b/python/freetoken/models/qwen4_exp/config.py index e413786266..61d4c0da3a 100644 --- a/python/freetoken/models/qwen4_exp/config.py +++ b/python/freetoken/models/qwen4_exp/config.py @@ -7,6 +7,8 @@ from freetoken.layers.quantization import QuantConfig from freetoken.models.config import ( + fp8_dense_enabled, + fp8_lmhead_enabled, mrope_layout_from_rope_params, FullAttentionGroupConfig, LinearGatedDeltaGroupConfig, @@ -112,6 +114,23 @@ def _layer_types(text: Any) -> list[str]: ] +def use_fp8_lmhead(config: ModelConfig) -> bool: + """Whether to build/load the lm_head as load-time per-tensor FP8. + + The MODEL BUILDER and the WEIGHT READER must agree exactly: if one says yes and the + other no, the state dict gains or loses ``lm_head.weight_scale`` and load fails. They + used to test this separately and drifted. Conditions: the operator asked + (FREETOKEN_FP8_LMHEAD=1), the head is untied (a tied head shares the bf16 embedding + table), and the checkpoint declares no scheme OF ITS OWN for lm_head -- a quantized head + goes through QuantConfig instead. Note it is the lm_head scheme that matters, not whether + the checkpoint has a QuantConfig at all: this model ships NVFP4 experts with lm_head in + the modelopt ignore list.""" + if not fp8_lmhead_enabled() or config.tie_word_embeddings: + return False + quant = getattr(config, "quant", None) + return quant is None or quant.scheme_for("lm_head") is None + + def parse_config(hf_config: Any) -> ModelConfig: text = getattr(hf_config, "text_config", hf_config) @@ -141,6 +160,12 @@ def parse_config(hf_config: Any) -> ModelConfig: else {k: v for k, v in rope_params.items() if not isinstance(v, (list, dict))} ) + # FREETOKEN_FP8_DENSE=1: the bf16 attention / GDN projections are quantized at LOAD to + # per-tensor e4m3 and served W8A8 through cuBLASLt (layers/fp8_dynamic.py). This is a + # synthetic scheme for a checkpoint that ships those modules unquantized -- distinct from + # the QuantConfig path, which serves what the checkpoint declares. The module builders + # apply it only where the checkpoint has no scheme of its own, so the two never collide. + attn_quant = "fp8_dynamic" if fp8_dense_enabled() else "none" layer_types = _layer_types(text) full_ids = tuple(i for i, t in enumerate(layer_types) if t == "full_attention") linear_ids = tuple(i for i, t in enumerate(layer_types) if t == "linear_attention") @@ -256,6 +281,7 @@ def parse_config(hf_config: Any) -> ModelConfig: image_token_id=getattr(hf_config, "image_token_id", None), attention_groups=groups, expert_quant=expert_quant, + attn_quant=attn_quant, qwen4_args=qwen4_args, slot_states=ple_slot_states(qwen4_args), ) diff --git a/python/freetoken/models/qwen4_exp/gdn.py b/python/freetoken/models/qwen4_exp/gdn.py index ed49dd42f0..3b39abab91 100644 --- a/python/freetoken/models/qwen4_exp/gdn.py +++ b/python/freetoken/models/qwen4_exp/gdn.py @@ -3,9 +3,12 @@ import torch import torch.nn.functional as F from freetoken.core import get_global_ctx +from freetoken.distributed import get_tp_info from freetoken.kernel.causal_conv1d import causal_conv1d_decode, causal_conv1d_varlen -from freetoken.layers import BaseOP, GatedRMSNorm, LinearColParallelMerged, LinearReplicated +from freetoken.layers import BaseOP, GatedRMSNorm, LinearColParallelMerged, LinearOProj from freetoken.layers.quantization import QuantConfig +from freetoken.utils import div_even +from freetoken.layers.fp8_dynamic import Fp8DynamicColMerged, Fp8DynamicRowParallel from freetoken.models.qwen3_5_moe.gdn_kernels import gdn_decode_fla, gdn_prefill_chunk_fla @@ -33,6 +36,7 @@ def __init__( self, hidden_size, num_k_heads, num_v_heads, head_k_dim, head_v_dim, conv_kernel_size, rms_norm_eps, layer_id, output_gate: str = "sigmoid", *, quant_config: QuantConfig | None = None, prefix: str = "", + attn_quant: str = "none", ): self.layer_id = layer_id # The fla chunk/decode kernels read+write the recurrent state and the per-chunk h as @@ -50,13 +54,42 @@ def __init__( self.value_dim = num_v_heads * head_v_dim self.conv_dim = 2 * self.key_dim + self.value_dim self.conv_kernel_size = conv_kernel_size + # TP-local head counts: k heads and their v-head groups split evenly across ranks + # (the fla kernels take the GQA ratio from the shapes); the state pool is sharded the + # same way (kvcache.linear_state_pool._linear_local_dims). + tp = get_tp_info() + self._local_num_k_heads = div_even(num_k_heads, tp.size, allow_replicate=True) + self._local_num_v_heads = div_even(num_v_heads, tp.size, allow_replicate=True) + self._local_key_dim = self._local_num_k_heads * head_k_dim + self._local_value_dim = self._local_num_v_heads * head_v_dim + self._local_conv_dim = 2 * self._local_key_dim + self._local_value_dim # quantized checkpoints quantize qkv|z but not b|a, so the fusion splits into a qkvz GEMM and a ba GEMM with their own schemes (matches sglang / vLLM) self._split_in_proj = ( quant_config is not None and quant_config.scheme_for(f"{prefix}.in_proj_qkvz") is not None ) - - self._in_proj_split = [self.conv_dim, self.value_dim, num_v_heads, num_v_heads] - if self._split_in_proj: + # Load-time per-tensor FP8 splits the fusion the same way, but it is synthetic: it + # applies only where the checkpoint declares no scheme of its own. + self._dynamic_fp8 = attn_quant == "fp8_dynamic" and not self._split_in_proj + + self._in_proj_split = [ + self._local_conv_dim, self._local_value_dim, + self._local_num_v_heads, self._local_num_v_heads, + ] + # A checkpoint-declared scheme has no TP-aware variant here: those two linears are + # built from the checkpoint's own scheme and are not sharded. Load-time fp8 is. + assert not (tp.size > 1 and self._split_in_proj), ( + "qwen4_exp TP shards bf16 or load-time fp8 GDN projections only" + ) + if self._dynamic_fp8: + self.in_proj_qkvz = Fp8DynamicColMerged( + hidden_size, [self.conv_dim, self.value_dim], + local_output_sizes=[self._local_conv_dim, self._local_value_dim], + ) + self.in_proj_ba = LinearColParallelMerged( + hidden_size, [num_v_heads, num_v_heads], has_bias=False, + local_output_sizes=[self._local_num_v_heads, self._local_num_v_heads], + ) + elif self._split_in_proj: self.in_proj_qkvz = LinearColParallelMerged( hidden_size, [self.conv_dim, self.value_dim], has_bias=False, quant_config=quant_config, prefix=f"{prefix}.in_proj_qkvz", @@ -68,21 +101,29 @@ def __init__( else: # Fused input projection (one GEMM instead of four): qkv | z | b | a. self.in_proj = LinearColParallelMerged( - hidden_size, self._in_proj_split, has_bias=False, + hidden_size, [self.conv_dim, self.value_dim, num_v_heads, num_v_heads], + has_bias=False, local_output_sizes=self._in_proj_split, quant_config=quant_config, prefix=f"{prefix}.in_proj", ) - self.conv1d = _DepthwiseConv1d(self.conv_dim, conv_kernel_size) + self.conv1d = _DepthwiseConv1d(self._local_conv_dim, conv_kernel_size) # Recurrence-gating params kept in fp32 (exp/softplus is precision-sensitive, # and the fla kernel reads them as fp32) -- matches HF/sglang, and avoids a # per-call .float() upcast in the decode wrapper. The weight loader exempts # *.A_log / *.dt_bias from the model-dtype downcast. - self.dt_bias = torch.empty(num_v_heads, dtype=torch.float32) - self.A_log = torch.empty(num_v_heads, dtype=torch.float32) + self.dt_bias = torch.empty(self._local_num_v_heads, dtype=torch.float32) + self.A_log = torch.empty(self._local_num_v_heads, dtype=torch.float32) self.norm = GatedRMSNorm(head_v_dim, eps=rms_norm_eps, activation=output_gate) - self.out_proj = LinearReplicated( - self.value_dim, hidden_size, has_bias=False, - quant_config=quant_config, prefix=f"{prefix}.out_proj", - ) + # Row-parallel over the local v heads, all-reduce inside. At TP=1 LinearOProj is + # exactly LinearReplicated (div_even(x, 1) == x and the reduction is skipped), so one + # class covers both -- and under TP>1 a replicated o_proj would silently take the + # wrong input width and skip the reduction. + if self._dynamic_fp8: + self.out_proj = Fp8DynamicRowParallel(self.value_dim, hidden_size) + else: + self.out_proj = LinearOProj( + self.value_dim, hidden_size, has_bias=False, + quant_config=quant_config, prefix=f"{prefix}.out_proj", + ) def _gate_params(self, a: torch.Tensor, b: torch.Tensor): beta = b.sigmoid() @@ -92,14 +133,17 @@ def _gate_params(self, a: torch.Tensor, b: torch.Tensor): def _conv_weight(self) -> torch.Tensor: return self.conv1d.weight.squeeze(1) # [conv_dim, kernel] for the fused kernel - def _conv_prefill(self, conv_in, pool, cu_seqlens, cache_indices, has_initial_state) -> torch.Tensor: + def _conv_prefill(self, conv_in, pool, cu_seqlens, cache_indices, has_initial_state, + max_seq_len=None) -> torch.Tensor: """Varlen causal conv (fused sgl_kernel) with silu; reads/updates each request's conv state in place by ``cache_indices`` slot. ``conv_in`` [total, conv_dim]. - ``cu_seqlens`` / ``cache_indices`` / ``has_initial_state`` come from FLAMetadata.""" + ``cu_seqlens`` / ``cache_indices`` / ``has_initial_state`` / ``max_seq_len`` come from + FLAMetadata; the last is host-known so the launch needs no D2H sync (graph capture).""" li = pool.local_index(self.layer_id) x = conv_in.transpose(0, 1).contiguous() # [conv_dim, total] out = causal_conv1d_varlen(x, self._conv_weight(), pool.conv_states[li], - cu_seqlens, cache_indices, has_initial_state) + cu_seqlens, cache_indices, has_initial_state, + max_seq_len=max_seq_len) return out.transpose(0, 1) # [total, conv_dim] def _conv_decode(self, conv_in: torch.Tensor, table_idx: torch.Tensor, pool) -> torch.Tensor: @@ -140,15 +184,17 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: fla = build_fla_metadata(batch, hidden_states.device) batch.fla_metadata = fla - if self._split_in_proj: + nk, nv = self._local_num_k_heads, self._local_num_v_heads + if self._split_in_proj or self._dynamic_fp8: qkvz = self.in_proj_qkvz.forward(hidden_states) - conv_in, z = torch.split(qkvz, [self.conv_dim, self.value_dim], dim=-1) + conv_in, z = torch.split(qkvz, self._in_proj_split[:2], dim=-1) ba = self.in_proj_ba.forward(hidden_states) - b, a = torch.split(ba, [self.num_v_heads, self.num_v_heads], dim=-1) + b, a = torch.split(ba, [nv, nv], dim=-1) else: proj = self.in_proj.forward(hidden_states) conv_in, z, b, a = torch.split(proj, self._in_proj_split, dim=-1) - z = z.reshape(total, self.num_v_heads, self.head_v_dim) + kd, vd = self._local_key_dim, self._local_value_dim + z = z.reshape(total, nv, self.head_v_dim) li = pool.local_index(self.layer_id) if batch.is_decode: @@ -157,10 +203,10 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: # no clone, no external l2norm). q/k stay at num_k_heads (kernel handles GQA). mixed = self._conv_decode(conv_in, fla.cache_indices, pool) # [B, conv_dim] B = mixed.shape[0] - qf, kf, vf = torch.split(mixed, [self.key_dim, self.key_dim, self.value_dim], dim=-1) - q = qf.reshape(1, B, self.num_k_heads, self.head_k_dim).to(dtype) - k = kf.reshape(1, B, self.num_k_heads, self.head_k_dim).to(dtype) - v = vf.reshape(1, B, self.num_v_heads, self.head_v_dim).to(dtype) + qf, kf, vf = torch.split(mixed, [kd, kd, vd], dim=-1) + q = qf.reshape(1, B, nk, self.head_k_dim).to(dtype) + k = kf.reshape(1, B, nk, self.head_k_dim).to(dtype) + v = vf.reshape(1, B, nv, self.head_v_dim).to(dtype) core_out = gdn_decode_fla( q, k, v, a, b, A_log=self.A_log, dt_bias=self.dt_bias, state_source=pool.recurrent_states[li], indices=fla.cache_indices, @@ -168,15 +214,16 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: ) else: mixed = self._conv_prefill( - conv_in, pool, fla.cu_seqlens, fla.cache_indices, fla.has_initial_state) + conv_in, pool, fla.cu_seqlens, fla.cache_indices, fla.has_initial_state, + fla.max_seq_len) # fla chunk handles GQA in-kernel: q/k stay at num_k_heads, v at num_v_heads. - qf, kf, vf = torch.split(mixed, [self.key_dim, self.key_dim, self.value_dim], dim=-1) - q = qf.reshape(1, total, self.num_k_heads, self.head_k_dim).to(dtype) - k = kf.reshape(1, total, self.num_k_heads, self.head_k_dim).to(dtype) - v = vf.reshape(1, total, self.num_v_heads, self.head_v_dim).to(dtype) + qf, kf, vf = torch.split(mixed, [kd, kd, vd], dim=-1) + q = qf.reshape(1, total, nk, self.head_k_dim).to(dtype) + k = kf.reshape(1, total, nk, self.head_k_dim).to(dtype) + v = vf.reshape(1, total, nv, self.head_v_dim).to(dtype) g, beta = self._gate_params(a, b) - g = g.reshape(1, total, self.num_v_heads) - beta = beta.float().reshape(1, total, self.num_v_heads) + g = g.reshape(1, total, nv) + beta = beta.float().reshape(1, total, nv) # The chunk kernel reads + writes back initial_state[cache_indices] in place; # fresh sequences (cached_len==0) must start from a zeroed slot. if fla.fresh_state_indices is not None: diff --git a/python/freetoken/models/qwen4_exp/hc.py b/python/freetoken/models/qwen4_exp/hc.py index f13808dccd..1574cd473f 100644 --- a/python/freetoken/models/qwen4_exp/hc.py +++ b/python/freetoken/models/qwen4_exp/hc.py @@ -17,6 +17,7 @@ from freetoken.kernel.triton.hc import ( grouped_gemma_rmsnorm, hc_combine, + hc_combine_norm, hc_gate_mix, hc_silu, ) @@ -119,14 +120,16 @@ def _down(self, rn: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor | None]: # both slices keep unit inner stride, so the kernels read them without a copy return down[:, : self.lowrank], down[:, self.lowrank : self.lowrank + self.hc_count] - def _mix_kernel(self, R: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor | None]: - rn = grouped_gemma_rmsnorm(R, self.hc_norm.weight, self.hc_norm.eps, self.hc_count) + def _mix_kernel(self, R: torch.Tensor, rn: torch.Tensor | None = None) -> Tuple[torch.Tensor, torch.Tensor | None]: + if rn is None: + rn = grouped_gemma_rmsnorm(R, self.hc_norm.weight, self.hc_norm.eps, self.hc_count) lora, s = self._down(rn) gate = self.input_mix_weight_up.forward(hc_silu(lora, self.hc_count)) return hc_gate_mix(rn, gate, self.hc_count), s - def _mix_torch(self, R: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor | None]: - rn = grouped_plus_one_rms_norm(R, self.hc_norm.weight, self.hc_norm.eps, self.hc_count) + def _mix_torch(self, R: torch.Tensor, rn: torch.Tensor | None = None) -> Tuple[torch.Tensor, torch.Tensor | None]: + if rn is None: + rn = grouped_plus_one_rms_norm(R, self.hc_norm.weight, self.hc_norm.eps, self.hc_count) lora, s = self._down(rn) lora = F.silu(lora.float() / self.hc_count) gate = self.input_mix_weight_up.forward(lora.to(R.dtype)) @@ -140,9 +143,12 @@ def _combine_torch(self, R: torch.Tensor, y: torch.Tensor, s: torch.Tensor) -> t out = out + y.float().unsqueeze(-2) * inject.unsqueeze(-1) return out.flatten(-2).to(R.dtype) - def mix(self, R: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor | None]: - """Return the block input ``x [T, hidden]`` and the inject logits ``s [T, hc_count]`` (None if no combine).""" - return self._mix_kernel(R) if R.is_cuda else self._mix_torch(R) + def mix(self, R: torch.Tensor, rn: torch.Tensor | None = None) -> Tuple[torch.Tensor, torch.Tensor | None]: + """Return the block input ``x [T, hidden]`` and the inject logits ``s [T, hc_count]`` (None if no combine). + + ``rn`` is the pre-normalized R from a preceding :meth:`combine_norm` (same norm weight as + this block's ``hc_norm``), letting the caller fuse the two kernels; None recomputes it.""" + return self._mix_kernel(R, rn) if R.is_cuda else self._mix_torch(R, rn) def combine(self, R: torch.Tensor, y: torch.Tensor, s: torch.Tensor) -> torch.Tensor: """Inject the block output ``y [T, hidden]`` back into every stream of ``R``.""" @@ -150,5 +156,15 @@ def combine(self, R: torch.Tensor, y: torch.Tensor, s: torch.Tensor) -> torch.Te return hc_combine(R, y, s, self.hc_count) return self._combine_torch(R, y, s) + def combine_norm( + self, R: torch.Tensor, y: torch.Tensor, s: torch.Tensor, next_block: "GatedResidual" + ) -> Tuple[torch.Tensor, torch.Tensor | None]: + """combine() fused with ``next_block``'s hc_norm; returns (R', Rn'). CPU keeps the split path (Rn' = None).""" + if not R.is_cuda: + return self.combine(R, y, s), None + return hc_combine_norm( + R, y, s, next_block.hc_norm.weight, next_block.hc_norm.eps, self.hc_count + ) + __all__ = ["GatedResidual", "GroupedPlusOneRMSNorm", "grouped_plus_one_rms_norm"] diff --git a/python/freetoken/models/qwen4_exp/model.py b/python/freetoken/models/qwen4_exp/model.py index 258badac42..6ce2da985b 100644 --- a/python/freetoken/models/qwen4_exp/model.py +++ b/python/freetoken/models/qwen4_exp/model.py @@ -21,11 +21,12 @@ from freetoken.core import get_global_ctx from freetoken.layers import BaseOP, OPList, ParallelLMHead, VocabParallelEmbedding from freetoken.models.blocks import BaseLLMModel +from .config import use_fp8_lmhead from freetoken.utils import nvtx_annotate from .attention import Qwen4ExpAttention from .hc import GatedResidual -from .moe import Qwen4ExpMoE +from .moe import Qwen4ExpMoE, wire_router_lookahead from .ple import PLELayer from freetoken.models.blocks import embed_input_ids from freetoken.models.qwen3_vl.vision import Qwen3VLVisionModel, QwenVLVisionMixin @@ -52,6 +53,7 @@ def build_linear_mixer(config: ModelConfig, layer_id: int, prefix: str) -> BaseO output_gate=g.output_gate, quant_config=config.quant, prefix=prefix, + attn_quant=config.attn_quant, ) @@ -73,17 +75,19 @@ def __init__(self, config: ModelConfig, layer_id: int, *, prefix: str = "") -> N ) @nvtx_annotate("Layer_{}", layer_id_field="_layer_id") - def forward(self, hidden: torch.Tensor, batch: Batch) -> torch.Tensor: + def forward(self, hidden: torch.Tensor, batch: Batch, rn: torch.Tensor | None = None): + """Returns (R, mlp_out, inject) un-combined: the caller fuses the mlp combine with the + next consumer's norm (next layer's attn mix, or the final mixer).""" if self.ple is not None: hidden = hidden + self.ple.forward(hidden, batch) - block_input, inject = self.attn_hyper_connection.mix(hidden) + block_input, inject = self.attn_hyper_connection.mix(hidden, rn=rn) if self._is_linear: block_output = self.linear_attn.forward(block_input) else: block_output = self.self_attn.forward(block_input, batch) - hidden = self.attn_hyper_connection.combine(hidden, block_output, inject) - block_input, inject = self.mlp_hyper_connection.mix(hidden) - return self.mlp_hyper_connection.combine(hidden, self.mlp.forward(block_input), inject) + hidden, rn_mlp = self.attn_hyper_connection.combine_norm(hidden, block_output, inject, self.mlp_hyper_connection) + block_input, inject = self.mlp_hyper_connection.mix(hidden, rn=rn_mlp) + return hidden, self.mlp.forward(block_input), inject class Qwen4ExpModel(BaseOP): @@ -99,6 +103,7 @@ def __init__(self, config: ModelConfig, *, prefix: str = "model") -> None: for layer_id in range(config.num_layers) ] ) + wire_router_lookahead([layer.mlp for layer in self.layers.op_list], config) self.hyper_connection_mixer = GatedResidual(config, use_combine=False, prefix=f"{prefix}.hyper_connection_mixer") # plain tuple (not an OP child), so it never shows up in the state dict self._ple = tuple(layer.ple for layer in self.layers.op_list if layer.ple is not None) @@ -118,27 +123,48 @@ def forward(self, input_ids: torch.Tensor, batch: Batch) -> torch.Tensor: meta = build_ple_metadata(batch, self._ple[0].args, input_ids.device) for ple in self._ple: # gather the pinned-host PLE rows while the early layers run ple.start_prefetch(batch, meta) - for layer in self.layers.op_list: - hidden = layer.forward(hidden, batch) + # The last block of each layer combines through the layer boundary, fusing the + # combine with the next consumer's hc_norm -- except before a PLE layer, whose + # `hidden + ple(hidden)` happens after the boundary and changes what the norm sees. + layers = self.layers.op_list + rn = None + for i, layer in enumerate(layers): + hidden, mlp_out, inject = layer.forward(hidden, batch, rn=rn) + nxt = layers[i + 1] if i + 1 < len(layers) else None + if nxt is not None and nxt.ple is not None: + hidden = layer.mlp_hyper_connection.combine(hidden, mlp_out, inject) + rn = None + else: + consumer = nxt.attn_hyper_connection if nxt is not None else self.hyper_connection_mixer + hidden, rn = layer.mlp_hyper_connection.combine_norm(hidden, mlp_out, inject, consumer) if meta is not None: # single writer: the layers only read the context, so a second PLE layer's # prefetch sees the un-rolled window commit_ngram_context(meta, getattr(batch, "fla_metadata", None)) - return self.hyper_connection_mixer.mix(hidden)[0] + return self.hyper_connection_mixer.mix(hidden, rn=rn)[0] class Qwen4ExpForCausalLM(BaseLLMModel): def __init__(self, config: ModelConfig) -> None: self._config = config self.model = Qwen4ExpModel(config) - self.lm_head = ParallelLMHead( - num_embeddings=config.vocab_size, - embedding_dim=config.hidden_size, - tie_word_embeddings=config.tie_word_embeddings, - tied_embedding=self.model.embed_tokens if config.tie_word_embeddings else None, - quant_config=config.quant, - prefix="lm_head", - ) + if use_fp8_lmhead(config): + # Load-time per-tensor FP8 head, for a checkpoint that ships lm_head unquantized. + # A checkpoint that declares its own scheme goes through QuantConfig instead. + from freetoken.layers.fp8_dynamic import Fp8ParallelLMHead + + self.lm_head = Fp8ParallelLMHead( + num_embeddings=config.vocab_size, embedding_dim=config.hidden_size + ) + else: + self.lm_head = ParallelLMHead( + num_embeddings=config.vocab_size, + embedding_dim=config.hidden_size, + tie_word_embeddings=config.tie_word_embeddings, + tied_embedding=self.model.embed_tokens if config.tie_word_embeddings else None, + quant_config=config.quant, + prefix="lm_head", + ) super().__init__() def load_host_tables(self, engine_config) -> int: diff --git a/python/freetoken/models/qwen4_exp/moe.py b/python/freetoken/models/qwen4_exp/moe.py index 3687f28035..4a1907d2ce 100644 --- a/python/freetoken/models/qwen4_exp/moe.py +++ b/python/freetoken/models/qwen4_exp/moe.py @@ -3,7 +3,12 @@ from typing import TYPE_CHECKING import torch +import torch.nn.functional as F +from freetoken.core import get_global_ctx +from freetoken.distributed import DistributedCommunicator, get_tp_info from freetoken.kernel.triton.moe_shared_gate import shared_gate_mul_add, shared_gate_sigmoid +from freetoken.layers import LinearRowParallel, silu_and_mul +from freetoken.layers.moe import OffloadMoELayer from freetoken.models.qwen3_5_moe.moe import Qwen3_5MoE if TYPE_CHECKING: @@ -14,16 +19,65 @@ class Qwen4ExpMoE(Qwen3_5MoE): """Qwen3_5MoE with the shared-expert gate on triton instead of gemv + sigmoid + mul + add. Same weights, same state dict. The gate reduction stays ahead of the routed experts, which may write into ``hidden_states`` in place. + + TP: the offload experts are sharded along the intermediate axis and the bf16 shared expert is + row-parallel, so both produce partial sums; ``routed + gate * shared`` is linear in them and + is reduced once (one all-reduce per MoE layer instead of two). """ + def __init__( + self, config: ModelConfig, layer_id: int | None = None, *, prefix: str = "" + ) -> None: + super().__init__(config, layer_id=layer_id, prefix=prefix) + self._comm = DistributedCommunicator() + self._tp_size = get_tp_info().size + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: num_tokens, hidden_dim = hidden_states.shape hidden_states = hidden_states.view(-1, hidden_dim) + se, ex = self.shared_expert, self.experts + if isinstance(ex, OffloadMoELayer): + # FREETOKEN_MOE_PREFETCH: the next layer's prediction needs only this block's input, so it + # starts here and runs beside this layer's router and ensure + ex.fork_lookahead(hidden_states) + if isinstance(ex, OffloadMoELayer) and ex.decode_side_applies(): + # FREETOKEN_MOE_COPY_OVERLAP: the shared expert runs on a side stream while the router, + # ensure and miss copy stay back to back on this one; joined before the combine + cache = ex.offload_cache + with cache.decode_side(ex.layer_id): + gate = shared_gate_sigmoid(hidden_states, self.shared_expert_gate.weight.view(-1)) + shared = se.forward(hidden_states) + routed = ex.forward(hidden_states=hidden_states, router_logits=self.gate.forward(hidden_states)) + cache.join_decode_side(ex.layer_id) + return shared_gate_mul_add(routed, shared, gate).view(num_tokens, hidden_dim) router_logits = self.gate.forward(hidden_states) - shared = self.shared_expert.forward(hidden_states) gate = shared_gate_sigmoid(hidden_states, self.shared_expert_gate.weight.view(-1)) - routed = self.experts.forward(hidden_states=hidden_states, router_logits=router_logits) + if ( + self._tp_size > 1 + and isinstance(se.down_proj, LinearRowParallel) + and isinstance(ex, OffloadMoELayer) + ): + shared = F.linear(silu_and_mul(se.gate_up_proj.forward(hidden_states)), se.down_proj.weight) + if get_global_ctx().batch.is_prefill: + routed = ex.prefill_forward(hidden_states, router_logits) + else: + routed = ex.decode_forward(hidden_states, router_logits) + out = self._comm.all_reduce(shared_gate_mul_add(routed, shared, gate)) + return out.view(num_tokens, hidden_dim) + shared = se.forward(hidden_states) + routed = ex.forward(hidden_states=hidden_states, router_logits=router_logits) return shared_gate_mul_add(routed, shared, gate).view(num_tokens, hidden_dim) -__all__ = ["Qwen4ExpMoE"] +def wire_router_lookahead(moes: list[Qwen4ExpMoE], config: ModelConfig) -> None: + """Give each offloaded MoE block the next layer's router, for FREETOKEN_MOE_PREFETCH's + lookahead (inert when it is off). ``moes[i]`` is decoder layer ``i``'s block.""" + from freetoken.moe.prefetch import default_budget + + for target in range(1, len(moes)): + experts = moes[target - 1].experts + if isinstance(experts, OffloadMoELayer): + experts.set_lookahead(moes[target].gate, target, default_budget(config.is_linear_layer(target))) + + +__all__ = ["Qwen4ExpMoE", "wire_router_lookahead"] diff --git a/python/freetoken/models/qwen4_exp/ple.py b/python/freetoken/models/qwen4_exp/ple.py index c944c4946f..22bda9e6c8 100644 --- a/python/freetoken/models/qwen4_exp/ple.py +++ b/python/freetoken/models/qwen4_exp/ple.py @@ -19,6 +19,7 @@ from __future__ import annotations import math +import os from dataclasses import dataclass from typing import TYPE_CHECKING, List, Protocol, Sequence, Tuple @@ -43,6 +44,18 @@ _SPLITMIX_M1 = 0xBF58476D1CE4E5B9 _SPLITMIX_M2 = 0x94D049BB133111EB _PLE_LAYER_PRIME = 10007 +# Distinct (is_decode, shape, device) keys the n-gram row-id token index is memoized for. +_TOKEN_INDEX_CACHE_SIZE = 64 +_FUSED_HASH_ENV = "FREETOKEN_PLE_FUSED_HASH" + + +def _capturing(device: torch.device) -> bool: + return device.type == "cuda" and torch.cuda.is_current_stream_capturing() + + +def _fused_row_ids_enabled() -> bool: + """The fused hash kernel, on unless ``FREETOKEN_PLE_FUSED_HASH=0`` takes it back to torch ops.""" + return (os.getenv(_FUSED_HASH_ENV) or "1").strip() not in ("0", "false", "False") class PLETableBackend(Protocol): @@ -421,6 +434,7 @@ def __init__(self, args: Qwen4ExpArgs, table: PLETableBackend | None = None) -> self.ngram_heads_vocab_sizes = torch.empty(self.num_heads, dtype=torch.int64) self.ngram_heads_offsets = torch.empty(self.num_heads, dtype=torch.int64) self._table = table + self._token_index_cache: dict[tuple, Tuple[torch.Tensor, torch.Tensor]] = {} def attach_table(self, table: PLETableBackend) -> None: self._table = table @@ -467,8 +481,93 @@ def _shift_ignore_eos(self, packed: torch.Tensor) -> List[torch.Tensor]: shifted.append(torch.where(valid, gathered, packed.new_full((), self.eos_token_id))) return shifted - def row_ids(self, meta: PLEMetadata) -> torch.Tensor: - """Global table row per (token, hash head): ``[T, num_ngram_heads]`` int64.""" + def _token_index(self, meta: PLEMetadata) -> Tuple[torch.Tensor, torch.Tensor]: + """``(req[T], local[T])`` int32: each token's request, and its offset inside it. + + The fused kernel addresses the hash window through these instead of materializing the + ``[B, ctx+max_len]`` packed window. Memoized on the shape (which is all they depend on) + for eager forwards only: a capture always builds its own, which the graph recomputes + into its pool on every replay. A cached pair (the eager warmup forward right before + capture leaves one) would be baked in by address, then evicted and reused. + """ + device = meta.input_ids.device + num_tokens = meta.input_ids.numel() + capturing = _capturing(device) + # is_decode is part of the key, not just the shape: a decode of B requests and a + # prefill of ONE B-token request are the same [T] and mean opposite things (one token + # per request at offset 0 vs B offsets inside one request). + key = ( + meta.is_decode, + (num_tokens,) if meta.is_decode else tuple(meta.seq_lens), + str(device), + ) + cached = None if capturing else self._token_index_cache.get(key) + if cached is not None: + return cached + if meta.is_decode: # one token per request, each at offset 0 + index = ( + torch.arange(num_tokens, dtype=torch.int32, device=device), + torch.zeros(num_tokens, dtype=torch.int32, device=device), + ) + else: + cu = meta.cu_seqlens.long() + flat_pos = torch.arange(num_tokens, device=device) + req = (torch.searchsorted(cu, flat_pos, right=True) - 1).clamp_( + max=len(meta.seq_lens) - 1 + ) + index = ((req).to(torch.int32), (flat_pos - cu[req]).to(torch.int32)) + if not capturing: + if len(self._token_index_cache) >= _TOKEN_INDEX_CACHE_SIZE: + self._token_index_cache.pop(next(iter(self._token_index_cache))) + self._token_index_cache[key] = index + return index + + def _use_fused_row_ids(self, meta: PLEMetadata) -> bool: + device = meta.input_ids.device + if device.type != "cuda": + return False + if not _fused_row_ids_enabled(): + return False + return all( + t.device == device + for t in ( + meta.ngram_context, + self.layer_multipliers, + self.ngram_heads_vocab_sizes, + self.ngram_heads_offsets, + ) + ) + + def row_ids(self, meta: PLEMetadata, out: torch.Tensor | None = None) -> torch.Tensor: + """Global table row per (token, hash head): ``[T, num_ngram_heads]`` int64. + + One Triton program per token on CUDA; ``row_ids_reference`` is the same arithmetic in + torch ops and stays the oracle (and the CPU path). + """ + if self._use_fused_row_ids(meta): + from freetoken.kernel.triton.ple_hash import ple_row_ids + + req, local = self._token_index(meta) + return ple_row_ids( + meta.input_ids.long(), + meta.ngram_context, + req, + local, + self.layer_multipliers, + self.ngram_heads_vocab_sizes, + self.ngram_heads_offsets, + eos_token_id=self.eos_token_id, + heads_per_ngram=self.heads_per_ngram, + out=out, + ) + rows = self.row_ids_reference(meta) + if out is None: + return rows + out.copy_(rows) + return out + + def row_ids_reference(self, meta: PLEMetadata) -> torch.Tensor: + """Torch-op transcription of the hash; the oracle the fused kernel is diffed against.""" packed, select = self._window(meta) tokens = [select(s) for s in self._shift_ignore_eos(packed)] blocks = [] @@ -482,8 +581,19 @@ def row_ids(self, meta: PLEMetadata) -> torch.Tensor: blocks.append(head_ids + self.ngram_heads_offsets[start:end]) return torch.cat(blocks, dim=-1) + def lookup_row_ids(self, meta: PLEMetadata) -> torch.Tensor: + """``row_ids`` as the attached table's ``prefetch``/``lookup`` take them. A table that + hashes on the host (``reads_row_ids = False``) reads only their shape: skip the hash.""" + if not getattr(self.table, "reads_row_ids", True): + return torch.empty( + (meta.input_ids.numel(), self.num_heads), + dtype=torch.int64, + device=meta.input_ids.device, + ) + return self.row_ids(meta) + def forward(self, meta: PLEMetadata, out: torch.Tensor | None = None) -> torch.Tensor: - return self.table.lookup(self.row_ids(meta), out) + return self.table.lookup(self.lookup_row_ids(meta), out) class _DepthwiseConv1d(BaseOP): @@ -581,7 +691,7 @@ def start_prefetch(self, batch: Batch, meta: PLEMetadata | None = None) -> None: """Hash this forward's n-grams and start the table gather on the side stream.""" if meta is None: meta = build_ple_metadata(batch, self.args, batch.input_ids.device) - row_ids = self.ple_embedding.row_ids(meta) + row_ids = self.ple_embedding.lookup_row_ids(meta) self._pending = (meta, row_ids) self.ple_embedding.table.prefetch(row_ids) @@ -602,7 +712,7 @@ def forward( elif pending is not None and pending[0] is meta: row_ids = pending[1] if row_ids is None: - row_ids = self.ple_embedding.row_ids(meta) + row_ids = self.ple_embedding.lookup_row_ids(meta) embeddings = self.ple_embedding.table.lookup(row_ids).to(R.dtype) key = self.norm_key.forward(self.key_proj.forward(embeddings)) diff --git a/python/freetoken/models/qwen4_exp/ple_disk.py b/python/freetoken/models/qwen4_exp/ple_disk.py index 69ee9fa93a..a0c71f7046 100644 --- a/python/freetoken/models/qwen4_exp/ple_disk.py +++ b/python/freetoken/models/qwen4_exp/ple_disk.py @@ -6,6 +6,10 @@ from __future__ import annotations import os +import queue +import threading +from collections import deque +from concurrent.futures import Future from contextlib import contextmanager from dataclasses import dataclass from typing import Sequence @@ -100,9 +104,62 @@ def resolve_row_source(folder: str) -> PleRowSource: return source_from_safetensors(folder) +class _Filler: + """The one thread that drives the PleStore (engine-thread-only C++, no locks), FIFO. + + Graph fills run here, submitted before the launch: the fill releases the captured decode's + WAIT, and cuGraphLaunch may block until the GPU drains, so it must not wait for the launch + to return (a large graph deadlocked that way; fork issue #65).""" + + def __init__(self) -> None: + self._jobs: queue.SimpleQueue = queue.SimpleQueue() + self._thread = threading.Thread(target=self._run, name="ple-filler", daemon=True) + self._thread.start() + + def submit(self, fn) -> Future: + fut: Future = Future() + self._jobs.put((fn, fut)) + return fut + + def _run(self) -> None: + while (item := self._jobs.get()) is not None: + fn, fut = item + try: + fut.set_result(fn()) + except BaseException as exc: # noqa: BLE001 -- surfaced to the engine thread + fut.set_exception(exc) + + def close(self) -> None: + self._jobs.put(None) + self._thread.join(timeout=5) + + +class _PendingFill: + """A graph fill submitted before its launch; ``cancel`` after a failed launch.""" + + def __init__(self, table: "DiskRowTable", future: Future, cancelled: threading.Event) -> None: + self._table, self._future, self._cancelled = table, future, cancelled + + def cancel(self) -> None: + # no WAIT consumed the signal: clear it, or the next step's WAIT passes on stale rows + self._cancelled.set() + try: + self._future.result() + except BaseException: # noqa: BLE001 -- the launch error is the one to surface + pass + pending = self._table._pending_fills() + if self._future in pending: + pending.remove(self._future) # its outcome was consumed here + self._table._retire(self._future) + self._table._flag.zero_() + + class DiskRowTable: """``PLETableBackend`` whose rows are read from disk per fill (--ple-backend disk).""" + # the fill hashes on the host; lookup/prefetch read only the row_ids shape + reads_row_ids = False + def __init__( self, source: PleRowSource, @@ -150,9 +207,6 @@ def __init__( self._graph_dev = torch.empty( max_graph_rows * self._token_bytes, dtype=torch.uint8, device=self._device ) - # completes once the last graph launch that read ``_graph_pinned`` has run (recorded by the - # dispatch context after the step's fill: a record inside the capture would not be observable) - self._graph_consumed = torch.cuda.Event() eager_bytes = max_extend_tokens * self._token_bytes # the overlap scheduler fills batch k+1 while batch k's lookup copy may still be queued: the # eager staging alternates two pinned buffers, and each is rewritten only once the H2D copy @@ -167,20 +221,62 @@ def __init__( from freetoken.kernel.row_store import probe_wait_sync self._wait_sync = probe_wait_sync(os.getenv(_SYNC_ENV, "auto"), self._device) - # one flag for all graphs: the readback event orders a fill after the previous graph, so signals never overlap + # one flag for all graphs: each fill's readback event orders it after the previous graph, so signals never overlap self._flag = alloc_pinned_tensor(1, dtype=torch.int64) self._flag.zero_() self._token_readback = alloc_pinned_tensor(max_graph_rows, dtype=torch.int32) - self._readback_event = torch.cuda.Event() sync = "wait-sync" if self._wait_sync else "launch-gating" logger.info_rank0(f"PLE disk backend: {self._store.io_backend()}, {sync}") # ---------------- host side (engine thread, before the forward launches) ---------------- + def _filler(self) -> _Filler: + filler = self.__dict__.get("_filler_thread") + if filler is None: + filler = self._filler_thread = _Filler() + return filler + + def close(self) -> None: + filler = self.__dict__.pop("_filler_thread", None) + if filler is not None: + filler.close() + + def _current_stream(self): + return torch.cuda.current_stream(self._device) + + def _pending_fills(self) -> deque: + return self.__dict__.setdefault("_pending_fill_queue", deque()) + + def _readback_pool(self) -> list: + return self.__dict__.setdefault("_readback_events", []) + + def _retire(self, future: Future) -> None: + """Return a finished fill's readback event for reuse (engine thread only).""" + event = getattr(future, "readback", None) + if event is not None: + self._readback_pool().append(event) + future.readback = None + + def _raise_failed_fill(self) -> None: + """Surface a finished graph fill's error on the engine thread, oldest first; a failed + fill signalled its WAIT, so that graph ran on stale rows.""" + pending = self._pending_fills() + while pending and pending[0].done(): + future = pending.popleft() + self._retire(future) + exc = future.exception() + if exc is not None: + raise exc + + def _fill_now(self, runs: Sequence[torch.Tensor], *, graph: bool) -> None: + """A fill that completes before its launch, on the store's thread (after any queued graph fill).""" + self._filler().submit(lambda: self.fill(runs, graph=graph)).result() + def fill(self, runs: Sequence[torch.Tensor], *, graph: bool) -> None: """Stage per-request token runs (two context ids, then the new tokens) in batch order.""" if graph: - self._graph_consumed.synchronize() + # no launch fence: the job already waited on the readback recorded after the last graph, + # and a fence recorded after this step's launch would wait on the WAIT this fill releases pinned = self._graph_pinned else: self._eager_slot ^= 1 @@ -202,22 +298,33 @@ def _ple_context(self, ids: torch.Tensor, position: int) -> list[int]: return _context(ids, position, self.eos_token_id) return [self.image_token_id if t >= MM_PAD_SHIFT_VALUE else t for t in _context(ids, position, self.eos_token_id)] - def host_fill_batch(self, batch: Batch, use_graph: bool): - """Stage this batch's rows; returns the post-dispatch fill callable under flag-sync, else None.""" - eos = self.eos_token_id + def host_fill_batch(self, batch: Batch, use_graph: bool) -> _PendingFill | None: + """Stage this batch's rows before the dispatch. Under flag-sync a graph fill is only submitted: + it completes on the filler thread, independently of the launch; returns its handle, else None.""" + self._raise_failed_fill() if batch.is_decode: reqs = list(batch.reqs) + # the context comes from the request as it is NOW: the engine mutates it once the launch returns + contexts = [self._ple_context(r.input_ids, r.device_len - 1) for r in reqs] if use_graph and self._wait_sync: bs = batch.padded_size self._token_readback[:bs].copy_(batch.input_ids, non_blocking=True) - self._readback_event.record(torch.cuda.current_stream(self._device)) - - def _complete() -> None: + # one event per fill: a shared one re-recorded by the next step before this job + # reached synchronize() would wait on a graph parked on this job's own WAIT + # recycled, never destroyed while serving: cuEventDestroy from the filler blocks on + # the driver lock a blocked cuGraphLaunch holds, and the fill it waits for never runs + pool = self._readback_pool() + readback = pool.pop() if pool else torch.cuda.Event() + readback.record(self._current_stream()) + cancelled = threading.Event() + + def _job() -> None: try: - self._readback_event.synchronize() + readback.synchronize() tokens = self._token_readback[:bs].to(torch.int64).tolist() - runs = [torch.tensor([*self._ple_context(r.input_ids, r.device_len - 1), t], dtype=torch.int64) - for r, t in zip(reqs, tokens)] + if cancelled.is_set(): + return + runs = [torch.tensor([*ctx, t], dtype=torch.int64) for ctx, t in zip(contexts, tokens)] self.fill(runs, graph=True) except BaseException: from freetoken.kernel.row_store import signal @@ -226,12 +333,14 @@ def _complete() -> None: signal(self._flag) raise - return _complete + future = self._filler().submit(_job) + future.readback = readback + self._pending_fills().append(future) + return _PendingFill(self, future, cancelled) # launch-gating: this D2H is the step's readback and orders the fill after sampling tokens = batch.input_ids.to("cpu").to(torch.int64).tolist() - runs = [torch.tensor([*self._ple_context(r.input_ids, r.device_len - 1), t], dtype=torch.int64) - for r, t in zip(reqs, tokens)] - self.fill(runs, graph=use_graph) + runs = [torch.tensor([*ctx, t], dtype=torch.int64) for ctx, t in zip(contexts, tokens)] + self._fill_now(runs, graph=use_graph) return None runs = [ torch.cat(( @@ -240,20 +349,19 @@ def _complete() -> None: )) for req in batch.padded_reqs ] - self.fill(runs, graph=False) + self._fill_now(runs, graph=False) return None @contextmanager def forward_host_ctx(self, batch: Batch, use_graph: bool): - """Around one dispatch: stage on enter, run the deferred fill+signal on exit.""" - deferred = self.host_fill_batch(batch, use_graph) - yield - # no try/finally: a failed launch leaves no WAIT pending, so the fill must not run - if deferred is not None: - deferred() - if use_graph: - # only after this step's fill: the next fill waits on this launch, not the fill on it - self._graph_consumed.record(torch.cuda.current_stream(self._device)) + """Around one dispatch: stage (or submit the graph fill) on enter; cancel a submitted fill if the launch fails.""" + pending = self.host_fill_batch(batch, use_graph) + try: + yield + except BaseException: + if pending is not None: + pending.cancel() + raise # ---------------- device side (PLETableBackend protocol) ---------------- diff --git a/python/freetoken/models/qwen4_exp/weight.py b/python/freetoken/models/qwen4_exp/weight.py index 691651d4cc..26f47af508 100644 --- a/python/freetoken/models/qwen4_exp/weight.py +++ b/python/freetoken/models/qwen4_exp/weight.py @@ -24,31 +24,29 @@ from freetoken.models.qwen3_vl.weight import rename_vl_prefix from freetoken.models.config import VISION_KEY_PREFIXES -from freetoken.models.loader import drop_page_cache, iter_weight_files +from freetoken.models.loader import drop_page_cache, iter_weight_files, shard_tensor from freetoken.models.nvfp4_banks import ( Nvfp4ExpertSourceSpec, ) -from freetoken.layers.quantization import get_quant_config +from freetoken.layers.quantization import QuantKind, get_quant_config from freetoken.models.register import get_model_spec from freetoken.moe.host_banks import HostBank, read_range_into -from freetoken.utils import cached_load_hf_config, download_hf_weight +from freetoken.utils import cached_load_hf_config, div_even, download_hf_weight, init_logger from freetoken.utils.progress import byte_bar from tqdm import tqdm -# Routed NVFP4 experts (nvidia modelopt layout): per-expert, un-fused. Matched against the RAW -# weight_map key in nvfp4_banks. The ``model.language_model.`` anchor excludes the MTP head's -# stacked ``mtp.layers.N.mlp.experts.*`` tensors. -_EXPERT_KEY_RE = re.compile( +logger = init_logger(__name__) + +# Routed NVFP4 experts: per-expert, un-fused, under the checkpoint's dialect names (nvidia modelopt +# or llm-compressor). Matched against the RAW weight_map key in nvfp4_banks. The +# ``model.language_model.`` anchor excludes the MTP head's stacked ``mtp.layers.N.mlp.experts.*``. +_EXPERT_KEY_RE = ( r"^model\.language_model\.layers\.(?P\d+)\.mlp\.experts\.(?P\d+)\." - r"(?Pgate_proj|up_proj|down_proj)\.(?Pweight|weight_scale|weight_scale_2)$" + r"(?Pgate_proj|up_proj|down_proj)\.(?P{kinds})$" ) +# role -> the expert bank reader's canonical (modelopt) tensor kind +_BANK_KINDS = {"weight": "weight", "weight_scale": "weight_scale", "weight_global": "weight_scale_2"} _EXPERT_RE = re.compile(r"\.mlp\.experts\.\d+\.") -_NVFP4_SOURCE_SPEC = Nvfp4ExpertSourceSpec( - key_pattern=_EXPERT_KEY_RE, - proj_to_role={"gate_proj": "gate", "up_proj": "up", "down_proj": "down"}, - layer_to_bank=lambda layer, config: layer, # every layer is MoE - desc="Qwen3.8-Flash-Next NVFP4 experts", -) # Per-tensor modelopt quant scales; consumed with their ``.weight`` (experts) or unused. _SCALE_SUFFIXES = (".weight_scale", ".weight_scale_2", ".input_scale") @@ -188,6 +186,107 @@ def fuse(self, name: str, tensor: torch.Tensor) -> list[tuple[str, torch.Tensor] return [(fused + kind, torch.cat(rows, dim=0))] +def _shard_rows( + t: torch.Tensor, parts: list[tuple[int, int]], rank: int, world: int +) -> torch.Tensor: + """Column-parallel slice (dim 0) of a ``[part0 | part1 | ...]`` fusion; ``parts`` gives each + part as ``(heads, rows_per_head)``. Heads split evenly across ranks; a part with fewer heads + than ranks (GQA kv) replicates head ``rank * heads // world``, the ``div_even(..., + allow_replicate=True)`` convention of the TP-aware layers.""" + out, off = [], 0 + for heads, rows in parts: + local = div_even(heads, world, allow_replicate=True) + first = rank * heads // world + out.append(t[off + first * rows : off + (first + local) * rows]) + off += heads * rows + assert off == t.shape[0], f"fusion parts {parts} cover {off} rows, tensor has {t.shape[0]}" + return torch.cat(out, dim=0) + + +def _shard(name: str, t: torch.Tensor, config, rank: int, world: int) -> torch.Tensor: + """TP shard of one state-dict tensor (fused projections included); identity at TP=1. + + Column-parallel (dim 0, by head): attention ``qkv_proj`` [q|gate per head | k | v], GDN + ``in_proj`` [q | k | v | z | b | a] and the matching ``conv1d`` channels, ``A_log`` / + ``dt_bias``, shared-expert ``gate_up_proj``. Row-parallel (dim 1): ``o_proj``, + ``out_proj``, shared-expert ``down_proj``. Vocab rows: ``embed_tokens`` / ``lm_head``. + Everything else (router, indexer, norms, HC, PLE, shared-expert gate) is replicated. + """ + if world == 1: + return t + if name.endswith(".self_attn.qkv_proj.weight"): + q = (config.num_qo_heads, 2 * config.head_dim) + kv = (config.num_kv_heads, config.head_dim) + return _shard_rows(t, [q, kv, kv], rank, world) + if ".linear_attn." in name: + g = config.linear_attention_group() + k = (g.num_key_heads, g.key_head_dim) + v = (g.num_value_heads, g.value_head_dim) + if name.endswith(".in_proj.weight"): + return _shard_rows(t, [k, k, v, v, (v[0], 1), (v[0], 1)], rank, world) + if name.endswith(".conv1d.weight"): + return _shard_rows(t, [k, k, v], rank, world) + if name.endswith((".A_log", ".dt_bias")): + return _shard_rows(t, [(v[0], 1)], rank, world) + if name.endswith(".out_proj.weight"): + return t.chunk(world, dim=1)[rank].clone() + return t + if name.endswith(".shared_expert.gate_up_proj.weight"): + half = t.shape[0] // 2 + return _shard_rows(t, [(half, 1), (half, 1)], rank, world) + # o_proj / down_proj: dim 1; embed_tokens / lm_head: vocab rows; others unchanged. + return shard_tensor(name, t, rank=rank, world_size=world, num_kv_heads=None) + + +# Load-time per-tensor FP8 (attn_quant == "fp8_dynamic", layers/fp8_dynamic.py): these keep +# their name and gain a sibling ``weight_scale``; GDN ``in_proj`` splits into the fp8 +# ``in_proj_qkvz`` and the bf16 ``in_proj_ba`` (the gate projections stay bf16, as in the +# block-fp8 checkpoints and in sglang / vLLM). +_FP8_DENSE_SUFFIXES = ( + ".self_attn.qkv_proj.weight", + ".self_attn.o_proj.weight", + ".linear_attn.out_proj.weight", +) +# Behind its own flag: this one moves the logits (see models.config.fp8_lmhead_enabled). +_FP8_LMHEAD_SUFFIX = "lm_head.weight" +_E4M3_MAX = 448.0 + + +def _quantize_per_tensor(w: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """``(e4m3 weight, fp32 scale ())`` with ``w ~= weight * scale``.""" + w = w.float() + scale = (w.abs().amax() / _E4M3_MAX).clamp_min(1e-12) + return (w / scale).clamp_(-_E4M3_MAX, _E4M3_MAX).to(torch.float8_e4m3fn), scale.reshape(()) + + +def _fp8_dense( + name: str, t: torch.Tensor, config, world: int, *, + dense: bool = True, lm_head: bool = False, +) -> Iterator[tuple[str, torch.Tensor]]: + """The (already sharded) dense tensor as the fp8_dynamic model expects it.""" + if dense and name.endswith(".linear_attn.in_proj.weight"): + g = config.linear_attention_group() + nk = div_even(g.num_key_heads, world, allow_replicate=True) + nv = div_even(g.num_value_heads, world, allow_replicate=True) + qkvz = 2 * nk * g.key_head_dim + 2 * nv * g.value_head_dim # [q | k | v | z] local rows + assert t.shape[0] == qkvz + 2 * nv, (name, t.shape, qkvz, nv) + base = name[: -len("in_proj.weight")] + w8, scale = _quantize_per_tensor(t[:qkvz]) + yield base + "in_proj_qkvz.weight", w8 + yield base + "in_proj_qkvz.weight_scale", scale + # clone, not contiguous(): a contiguous row slice IS contiguous, so .contiguous() would + # hand back a view that keeps the whole bf16 in_proj (36 x 42 MB per rank) resident + yield base + "in_proj_ba.weight", t[qkvz:].clone() + elif (dense and name.endswith(_FP8_DENSE_SUFFIXES)) or ( + lm_head and name.endswith(_FP8_LMHEAD_SUFFIX) + ): + w8, scale = _quantize_per_tensor(t) + yield name, w8 + yield name[: -len("weight")] + "weight_scale", scale + else: + yield name, t + + def iter_weights( model_path: str, device: torch.device, @@ -203,18 +302,44 @@ def iter_weights( Fusions, per kind: attention q|k|v -> ``qkv_proj``; GDN ``in_proj_{qkv,z,b,a}`` -> ``in_proj``, or ``in_proj_qkvz`` + bf16 ``in_proj_ba`` when qkv|z is quantized; shared-expert gate|up -> ``gate_up_proj``; each per-layer HC's ``input_mix_weight_down`` | ``block_inject_weight`` -> a zero-padded ``input_mix_weight_down_block_inject``. ``include_moe_experts`` is accepted for the loader contract but never yields anything: the routed experts are NVFP4 and always come from the offload cache's expert reader. """ - if get_tp_info().size > 1: - raise NotImplementedError("qwen4_exp weight loading supports TP=1 only") if not include_non_moe: return + from freetoken.models.config import fp8_dense_enabled, fp8_lmhead_enabled + + from .config import parse_config, use_fp8_lmhead + + tp = get_tp_info() hf_config = cached_load_hf_config(model_path) spec = get_model_spec(hf_config.architectures[0]) fuser = _DenseFuser(get_quant_config(), spec.packed_modules_mapping) + # the sharding geometry and the fp8 split need the parsed config; TP=1 bf16 never does + config = ( + parse_config(hf_config) + if tp.size > 1 or fp8_dense_enabled() or fp8_lmhead_enabled() + else None + ) + # the SAME predicate the model builder uses, so the emitted keys always match its slots + lmhead = config is not None and use_fp8_lmhead(config) + fp8 = config is not None and config.attn_quant == "fp8_dynamic" + if fp8: + logger.info( + "qwen4_exp dense projections: load-time per-tensor FP8 (W8A8 via _scaled_mm), " + "FREETOKEN_FP8_DENSE=1" + ) + if lmhead: + logger.info("qwen4_exp lm_head: load-time per-tensor FP8, FREETOKEN_FP8_LMHEAD=1") + + def emit(name: str, tensor: torch.Tensor): + tensor = _shard(name, tensor, config, tp.rank, tp.size) + if fp8 or lmhead: + yield from _fp8_dense(name, tensor, config, tp.size, dense=fp8, lm_head=lmhead) + else: + yield name, tensor for file in tqdm( iter_weight_files(model_path), desc="Loading weights", - disable=not get_tp_info().is_primary(), + disable=not tp.is_primary(), ): with safetensors.safe_open(file, framework="pt", device=str(device)) as f: for raw_name in f.keys(): @@ -227,11 +352,17 @@ def iter_weights( fused = fuser.fuse(name, tensor) if fused is None: fuser.check_unfused(name, tensor) - yield name, tensor + yield from emit(name, tensor) else: - yield from fused + for fused_name, fused_tensor in fused: + yield from emit(fused_name, fused_tensor) assert not fuser.buf, f"Incomplete projection fusions: {sorted(k[0] + k[1] for k in fuser.buf)}" + if (fp8 or lmhead) and device.type == "cuda": + # The bf16 originals and fp32 temporaries of the quantization sit in the caching + # allocator; hand them back so the expert-cache planner (free VRAM after load) sees + # the halved dense footprint instead of the slack. + torch.cuda.empty_cache() def iter_vision_weights(model_path: str, device: torch.device) -> Iterator[tuple[str, torch.Tensor]]: @@ -282,6 +413,32 @@ def _ple_table_files(folder: str) -> list[str]: return sorted(os.path.join(folder, shard) for shard in files) +def check_ftw_weights(config, model_path: str, names: set[str]) -> None: + """Refuse an FTW checkpoint whose stored FP8 split disagrees with this process's FP8 flags. + + The FTW holds the tensors ``iter_weights`` emitted at conversion, so a mismatch otherwise + surfaces only after the dense read, as a KeyError on the first fused GDN projection.""" + from .config import use_fp8_lmhead + + stored = { + "FREETOKEN_FP8_DENSE": any(n.endswith(".in_proj_qkvz.weight_scale") for n in names), + "FREETOKEN_FP8_LMHEAD": "lm_head.weight_scale" in names, + } + wanted = { + "FREETOKEN_FP8_DENSE": config.attn_quant == "fp8_dynamic", + "FREETOKEN_FP8_LMHEAD": use_fp8_lmhead(config), + } + wrong = [var for var in stored if stored[var] != wanted[var]] + if wrong: + was = ", ".join(f"{var}={int(stored[var])}" for var in wrong) + now = ", ".join(f"{var}={int(wanted[var])}" for var in wrong) + raise ValueError( + f"FTW checkpoint {model_path} stores its dense weights as converted with {was}, " + f"but this process has {now}; serve it with the conversion-time setting, or " + "reconvert with ft checkpoint under the setting you want" + ) + + def ftw_side_files(model_path: str, out_dir: str) -> list[str]: """Write the PLE n-gram table tensors, and only those, into ``ple-table-*.safetensors`` next to an FTW checkpoint. @@ -382,8 +539,19 @@ def load_ple_table(model_path: str, qwen4_args, *, pin: bool = True, # ====================================================================================== -def nvfp4_expert_spec(model_path: str, config): - return _NVFP4_SOURCE_SPEC +def nvfp4_expert_spec(model_path: str, config) -> Nvfp4ExpertSourceSpec: + """The per-expert NVFP4 layout under the checkpoint's dialect names (modelopt or llm-compressor).""" + quant = get_quant_config() + stored = quant.stored_tensors(QuantKind.NVFP4) + kind_map = {stored[role].name: kind for role, kind in _BANK_KINDS.items()} + return Nvfp4ExpertSourceSpec( + key_pattern=re.compile(_EXPERT_KEY_RE.format(kinds="|".join(map(re.escape, kind_map)))), + proj_to_role={"gate_proj": "gate", "up_proj": "up", "down_proj": "down"}, + layer_to_bank=lambda layer, config: layer, # every layer is MoE + desc=f"Qwen3.8-Flash-Next NVFP4 experts ({quant.dialect})", + kind_map=kind_map, + global_reciprocal=stored["weight_global"].reciprocal, + ) __all__ = [ diff --git a/python/freetoken/models/weight.py b/python/freetoken/models/weight.py index e3382b1b59..29cce4eefa 100644 --- a/python/freetoken/models/weight.py +++ b/python/freetoken/models/weight.py @@ -218,12 +218,24 @@ def load_weight( # decided at conversion (offload -> experts live in banks, not here); a backend mismatch # fails loudly in load_state_dict (strict missing/unexpected expert keys), so the reader # just yields the stored weight tensors regardless of the include_moe_experts flag. - from freetoken.checkpoint.ftw import is_ftw_checkpoint, iter_ftw_weights + from freetoken.checkpoint.ftw import FTWReader, is_ftw_checkpoint, iter_ftw_weights from freetoken.models.config import VISION_KEY_PREFIXES # a text-only engine never built the tower, so its tensors are not even read keep = None if include_vision else (lambda name: not name.startswith(VISION_KEY_PREFIXES)) if is_ftw_checkpoint(model_path): + # a family whose stored tensors depend on load-time settings checks them against the index before the + # read; a bare FTW without a readable HF config (synthetic test checkpoints) has no family to ask + try: + spec = get_model_spec(cached_load_hf_config(model_path).architectures[0]) + except (OSError, ValueError): + spec = None + check = _model_override(spec, "check_ftw_weights") if spec is not None else None + if check is not None: + reader = FTWReader(model_path) + names = {e["name"] for e in reader.entries("weight")} + reader.close() + check(_spec_for_model_path(model_path)[0], model_path, names) weights = iter_ftw_weights(model_path, keep=keep) else: _config, spec = _spec_for_model_path(model_path) diff --git a/python/freetoken/moe/benchbw.py b/python/freetoken/moe/benchbw.py index e70cad31d3..cd260cb2f6 100644 --- a/python/freetoken/moe/benchbw.py +++ b/python/freetoken/moe/benchbw.py @@ -413,14 +413,11 @@ def _build_gather_rig(fmt: str, wl: Workload, device: torch.device): specs = _offload_bank_specs(fmt, H, I) cache = OffloadMoeCache(num_layers=1, num_experts=E, cache_size=E, device=device, quant_format=fmt) total_bytes = 0 + sources = {} for name, (elems, dtype) in specs.items(): - src = alloc_pinned_tensor(E, elems, dtype=dtype) # cudaHostAlloc (mapped) like the loaders - dst = torch.empty(E, elems, dtype=dtype, device=device) - cache.bank_sources[name] = [src] - cache.bank_caches[name] = dst - cache.banks.append(([src], dst)) + sources[name] = [alloc_pinned_tensor(E, elems, dtype=dtype)] # cudaHostAlloc (mapped) like the loaders total_bytes += E * elems * dtype.itemsize - cache._build_copy_plan() # enable the fused path (production default) when aligned + cache.set_bank_sources(sources) # builds the fused copy plan (production default) when aligned cache.evict_slots[:E].copy_(torch.arange(E, dtype=torch.int32, device=device)) cache.src_indices[:E].copy_(torch.randperm(E, device=device).to(torch.int32)) @@ -490,17 +487,16 @@ def _cpu_moe_step_fns(ex, wl: Workload, E: int): disjoint experts (working set >> LLC) -- the DRAM-bound regime hybrid decode actually runs in; ``run()`` executes one step (releases the GIL inside the extension).""" bs, top_k = 1, wl.top_k - io = ex._io_for(bs) + io = ex._io io["x"].copy_(torch.randn(bs, wl.hidden, dtype=torch.bfloat16) * 0.1) io["w"].copy_(torch.rand(bs, top_k, dtype=torch.float32)) - task = ex._task_for(0, bs) base = torch.arange(top_k, dtype=torch.int32) def set_ids(i: int) -> None: io["ids"].copy_((((i * top_k) + base) % E).view(bs, top_k)) def run() -> None: - ex._ext.run_task(task) + ex._ext.run(0, bs) return set_ids, run @@ -566,7 +562,7 @@ def measure_overlap_bw(fmt: str, wl: Workload, device: torch.device, backend its bandwidth-matched fetch split: fetched : cpu-computed = pcie_ov : cpu_ov. Both sides hammer their own pinned banks flat out for ``seconds`` from a shared - barrier: a worker thread loops bs=1 CPU decode steps (``run_task`` releases the GIL) + barrier: a worker thread loops bs=1 CPU decode steps (``run`` releases the GIL) while the main thread loops full-layer ``copy_missing`` gathers (synchronized per copy so the DMA is really in flight, not just enqueued). Each side reports bytes / its own elapsed; the two windows differ by at most one CPU step + one gather. @@ -663,7 +659,7 @@ def _bench_format(fmt: str, wl: Workload, device: torch.device, threshold: float entry["isa_sweep"] = c["isa_sweep"] # The tier label comes from the bf16 dot chooser, so it can over-report tiers # the actual per-format kernel doesn't distinguish. Disclose the real story. - if fmt == "nvfp4": # WF_NVFP4 is the only use_vnni format + if fmt == "nvfp4": # nvfp4 is the only use_vnni format _note(entry, "isa tier labels are nominal: this W4A8 kernel rides AVX-VNNI " "regardless of the forced tier") elif fmt in ("mxfp4_triton", "ds_fp4", "fp8_block"): diff --git a/python/freetoken/moe/cpu_executor.py b/python/freetoken/moe/cpu_executor.py index f519687ed6..bb9ee56385 100644 --- a/python/freetoken/moe/cpu_executor.py +++ b/python/freetoken/moe/cpu_executor.py @@ -1,18 +1,15 @@ """Python wrapper around the ``_cpu_moe`` C++ executor (--moe-strategy cpu). -Owns the persistent CPU worker pool, the per-batch-size pinned IO buffers, and -the per-(layer, batch-size) host-func task descriptors. ``decode`` issues the -whole CUDA-graph-capturable sequence on the current stream: +Owns the persistent CPU worker pool and one pinned IO block of ``max_tokens`` rows that +every layer and batch size shares (one task is in flight at a time). ``decode`` issues +the whole CUDA-graph-capturable sequence on the current stream: - D2H (hidden, topk_ids, topk_weights -> pinned) - -> submit host node (cudaLaunchHostFunc: enqueue MoE task to the pool) - -> sync host node (cudaLaunchHostFunc: spin until the pool drains) + D2H (hidden, topk_ids, topk_weights -> pinned IO rows [:bs]) + -> submit (doorbell naming (layer, bs), or a cudaLaunchHostFunc host node) + -> sync (wait for the CPU pool to finish) -> H2D (pinned expert output -> GPU) -Buffers and tasks are allocated lazily per batch size. GraphRunner runs an eager -``model.forward()`` at each batch size immediately before capturing it, so the -first (eager) call materializes the stable pinned buffers + task pointers that -the subsequent capture embeds in its host/memcpy nodes. +With the doorbell on CUDA, a publish kernel does the D2H and the doorbell, and skips the step when no route goes to the CPU; a combine kernel does the H2D (csrc/cpu_moe/gpu_sync_io.h). """ from __future__ import annotations @@ -35,28 +32,27 @@ # on a 75-layer model). Instead the GPU raises a mapped-pinned "ready" flag at submit; a # persistent CPU coordinator (in _cpu_moe) polls it, runs the layer, and sets a "done" # flag the GPU waits on at sync -- no host-func round-trip. Both GPU-side operations are -# STREAM MEMORY OPERATIONS (cuStreamWriteValue64 / cuStreamWaitValue64, resolved from the -# driver at runtime): they execute on the GPU front end with no SM-resident kernel, so +# STREAM MEMORY OPERATIONS (cuStreamWriteValue64 / cuStreamWaitValue64 on CUDA, see +# gpu_sync.h): they execute on the GPU front end with no SM-resident kernel, so # GPU utilization stays truthful during the CPU compute window. (The first cut used a # spin-wait kernel; that pinned reported utilization at 99% and laptop CPU/GPU dynamic # power schedulers responded by clamping the CPU frequency -- a net decode regression on -# power-coupled edge devices.) Each (layer, decode batch size) pair gets its own flag -# slot, so every captured decode graph rides the handshake. Where memops are unavailable -# (Windows WDDM, vGPU, old drivers -- functionally probed at startup) or the slot -# capacity is exceeded, decode keeps the host-func path (functional, slower). A Python -# watchdog thread turns a wedged coordinator into a loud RuntimeError (via err[] + -# raise_if_unhealthy) instead of an indefinite stream stall. +# power-coupled edge devices.) The doorbell word names the (layer, decode batch size) +# task, so one ready/done pair serves every layer and every captured graph (gpu_sync.h). +# Where memops are unavailable (some driver models, vGPU, old drivers -- functionally probed at +# startup), decode keeps the host-func path (functional, slower). ROCm 7.14 cannot replay +# that path from a graph, so capture fails loudly there instead. A Python watchdog thread +# turns a wedged coordinator into a loud RuntimeError (via raise_if_unhealthy) instead of +# an indefinite stream stall. # Caveat: the coordinator busy-polls one core while decode traffic flows (idle backoff # otherwise); FREETOKEN_CPU_MOE_FLAG_SYNC=0 opts out entirely. _FLAG_SYNC = os.getenv("FREETOKEN_CPU_MOE_FLAG_SYNC", "1") != "0" -# Flag slots per MoE layer: covers this many distinct decode batch sizes (captured graph -# sizes plus any eager padded sizes); more than that is unheard of, and the overflow -# just keeps the host-func path for the extra combos. -_FLAG_SLOTS_PER_LAYER = 16 # Pool threads spin up to 50 ms before they park. FREETOKEN_CPU_MOE_SPIN=0 opts out. _SPIN_WAIT = os.getenv("FREETOKEN_CPU_MOE_SPIN", "1") != "0" +# ds_fp4 runs its input FP8 round-trip on the GPU when it can; tests turn this off to compare against the CPU round-trip. +_DSFP4_GPU_PREQUANT = True -# Activation ids must match ActKind in csrc/cpu_moe/cpu_moe_ext.cpp. Id 3 is the +# Activation ids must match ActKind in csrc/cpu_moe/common.h. Id 3 is the # clamped (up + 1) swiglu: "swigluoai" runs it in the generic GEMV epilogue, # "gpt_oss_swiglu" is the same math fused inside the mxfp4 kernel. _ACT_IDS = { @@ -70,10 +66,6 @@ "swiglu_clamp": 4, } -# Weight-format ids must match WFmt in csrc/cpu_moe/cpu_moe_ext.cpp. -_WFMT_IDS = {"bf16": 0, "nvfp4": 1, "mxfp4_triton": 2, "ds_fp4": 3, "q4_0": 4, - "fp8_block": 5} - def compiled_extension_supports(activation: str) -> bool: """Whether the compiled ``_cpu_moe`` extension can serve ``activation`` @@ -157,7 +149,7 @@ def resolve_threads_and_affinity(requested: int) -> tuple[int, list[int]]: class CpuMoeExecutor: """Decode-time CPU expert compute over an ``OffloadMoeCache``'s host banks - (bf16, nvfp4, mxfp4_triton, ds_fp4, q4_0 or fp8_block — see ``_WFMT_IDS`` / ``_resolve_banks``).""" + (the formats ``_cpu_moe.weight_formats()`` lists — see ``_resolve_banks``).""" def __init__( self, @@ -178,26 +170,20 @@ def __init__( from freetoken.moe.legacy_format import canonical_role fmt = fmt or cache.quant_format - if fmt not in _WFMT_IDS: + if fmt not in _cpu_moe.weight_formats(): raise NotImplementedError( f"--moe-strategy cpu/hybrid computes experts on the CPU and supports " - f"{sorted(_WFMT_IDS)} formats, but this checkpoint's experts are " + f"{sorted(_cpu_moe.weight_formats())} formats, but this checkpoint's experts are " f"{fmt!r}; use --moe-strategy offload (GPU-side dequant) instead." ) if activation not in _ACT_IDS: raise NotImplementedError(f"CPU MoE backend: unsupported activation {activation!r}") + if apply_router_weight_on_input and fmt in ("mxfp4_triton", "ds_fp4"): + raise NotImplementedError(f"CPU MoE backend: {fmt} experts apply the router weight on the output only") # ABI probe: a stale prebuilt _cpu_moe.so accepts newer act ids without # error and silently computes the wrong activation in the generic # epilogue -- fail loudly with the rebuild instruction instead. (mxfp4 # handles its act inside the kernel and predates the marker.) - # A prebuilt .so that predates this format would accept the id and then take - # the wrong dequant branch -- silently wrong output rather than a crash. - if _WFMT_IDS[fmt] > getattr(_cpu_moe, "max_weight_format_id", lambda: 4)(): - raise RuntimeError( - f"the compiled _cpu_moe extension predates the {fmt!r} expert format; " - "rebuild it with `python setup.py build_ext --inplace` (or reinstall " - "the wheel) before serving this checkpoint on the cpu/hybrid backend." - ) if _ACT_IDS[activation] >= 3 and fmt != "mxfp4_triton": supported = getattr(_cpu_moe, "max_generic_act_id", lambda: 2)() if _ACT_IDS[activation] > supported: @@ -218,28 +204,42 @@ def __init__( # The per-layer tensors and their pointer tables must outlive the executor # (C++ holds raw addresses into both). self._banks: list[torch.Tensor] = [] - ptrs, (self.H, self.I) = self._resolve_banks( + banks, options, (self.H, self.I) = self._resolve_banks( {canonical_role(name): per_layer for name, per_layer in cache.bank_sources.items()}, fmt ) + # ds_fp4: the reference FP8 activation round-trip is a scalar per-element chain + # that the C++ side runs single-threaded on the CUDA host-callback thread -- + # straight on the decode critical path (~0.3ms/layer at H=4096, every worker and + # the GPU waiting on it). When a GPU is present we run the numerically identical + # round-trip as a captured GPU kernel BEFORE the D2H (see decode_submit) and tell + # the C++ side to skip its own. Measured on DeepSeek-V4-Flash bs=1 decode: + # 12.85 -> 15.65 tok/s, output bit-identical (tests/moe/test_dsfp4_prequant.py). + self._gpu_prequant = fmt == "ds_fp4" and device.type == "cuda" and _DSFP4_GPU_PREQUANT + # the W4A8 activation round-trip block follows the checkpoint's fp8 block (128 on V4, 32 on V4.1) + self._act_block = act_block + if fmt == "ds_fp4": + assert act_block is not None, "ds_fp4 experts need act_block, the checkpoint's fp8 activation quant block" + options.update(act_block=int(act_block), input_prequant=int(self._gpu_prequant)) + if self._gpu_prequant: + logger.info_rank0( + "ds_fp4: input FP8 round-trip moved to the GPU " + "(bit-identical grid; the CPU-side scalar round-trip is skipped)" + ) + # Decide the flag handshake up front (env + device + a functional stream-memop # probe): its coordinator needs a core of its own, which the auto thread sizing # below reserves (a coordinator time-slicing against the GEMV workers measurably # destabilizes throughput on fully-subscribed boxes). self._flag_sync = _FLAG_SYNC and device.type == "cuda" - self._cpu_moe = _cpu_moe # module ref for the decode-path memop calls - if self._flag_sync: - probe_scratch = alloc_pinned_tensor(1, dtype=torch.int64) - probe_scratch.zero_() - if not _cpu_moe.memops_probe( - torch.cuda.current_stream().cuda_stream, probe_scratch.data_ptr() - ): - logger.info_rank0( - "cpu-moe flag handshake unavailable: CUDA stream memory operations " - "are not supported here (Windows WDDM / vGPU / old driver); using " - "the cudaLaunchHostFunc sync" - ) - self._flag_sync = False + if self._flag_sync and not _cpu_moe.memops_probe(torch.cuda.current_stream().cuda_stream): + logger.info_rank0( + "cpu-moe flag handshake unavailable: stream memory operations failed the " + "startup probe; using the cudaLaunchHostFunc sync" + ) + self._flag_sync = False + # ROCm 7.14 returns wrong results when it replays host-function nodes, so there only the doorbell is graph-safe. + self._graph_capture_safe = self._flag_sync or _cpu_moe.host_func_capture_safe() nthreads, core_ids = resolve_threads_and_affinity(num_threads) coord_core = -1 @@ -250,6 +250,17 @@ def __init__( nthreads -= 1 core_ids = core_ids[:-1] self._coord_core = coord_core + alloc_io = ( + (lambda *shape, dtype: torch.empty(*shape, dtype=dtype)) + if device.type == "cpu" + else alloc_pinned_tensor + ) + self._io = { + "x": alloc_io(self.max_tokens, self.H, dtype=torch.bfloat16), + "ids": alloc_io(self.max_tokens, self.top_k, dtype=torch.int32), + "w": alloc_io(self.max_tokens, self.top_k, dtype=torch.float32), + "y": alloc_io(self.max_tokens, self.H, dtype=torch.bfloat16), + } self._ext = _cpu_moe.CpuMoeExecutor( num_threads=nthreads, num_layers=self.num_layers, @@ -260,11 +271,16 @@ def __init__( max_tokens=self.max_tokens, activation_id=_ACT_IDS[activation], apply_router_weight_on_input=1 if apply_router_weight_on_input else 0, - weight_format=_WFMT_IDS[fmt], + weight_format=fmt, + banks=banks, + options=options, swiglu_alpha=float(swiglu_alpha), swiglu_limit=float(swiglu_limit) if swiglu_limit is not None else float("inf"), core_ids=core_ids, - **ptrs, + x_ptr=self._io["x"].data_ptr(), + ids_ptr=self._io["ids"].data_ptr(), + w_ptr=self._io["w"].data_ptr(), + y_ptr=self._io["y"].data_ptr(), ) self.num_threads = nthreads self.core_ids = core_ids @@ -281,34 +297,24 @@ def __init__( ) torch.set_num_threads(clamp) - self._io: dict[int, dict[str, torch.Tensor]] = {} - self._tasks: dict[tuple[int, int], int] = {} - - # Flag-based handshake: mapped-pinned ready/done/err int64 arrays (one slot per - # (MoE layer, decode batch size) pair, allocated as tasks are created) + a - # persistent CPU coordinator that polls ready[], runs the slot's task on the - # pool, and sets done[]. Binary per-step protocol (GPU memops: done=0, ready=1; - # coordinator: consume ready, run, done=1; GPU waits done>=1 -- the WAIT - # immediate is constant, so CUDA-graph replays are safe). err[] is raised by the - # WATCHDOG thread when a ready flag stays unanswered (dead coordinator): it - # poisons done to unblock the stream and raise_if_unhealthy() turns the step - # into a loud error instead of silent stale output. Buffers are kept alive on - # self so the coordinator's pinned pointers stay valid for the executor's - # lifetime (flag_sync itself was decided above, before thread sizing). - self._ready = self._done = self._err = None - self._flag_slots: dict[tuple[int, int], int] = {} # (layer_id, bs) -> slot - self._flag_capacity = self.num_layers * _FLAG_SLOTS_PER_LAYER + # Flag-based handshake: the extension owns the ready/done words (gpu_sync.h) and a + # persistent CPU coordinator that polls the doorbell, runs the task it names on the + # pool, and sets done for the GPU. The WATCHDOG thread marks the executor unhealthy + # when a doorbell stays unanswered (dead coordinator): it poisons done to unblock the + # stream and raise_if_unhealthy() turns the step into a loud error instead of silent + # stale output (flag_sync itself was decided above, before thread sizing). + self._unhealthy = False + self._kernel_io = False if self._flag_sync: - self._ready = alloc_pinned_tensor(self._flag_capacity, dtype=torch.int64) - self._done = alloc_pinned_tensor(self._flag_capacity, dtype=torch.int64) - self._err = alloc_pinned_tensor(self._flag_capacity, dtype=torch.int64) - self._ready.zero_() - self._done.zero_() - self._err.zero_() - self._ext.start_flag_coordinator( - self._ready.data_ptr(), self._done.data_ptr(), self._flag_capacity, - self._coord_core, - ) + self._ext.start_flag_coordinator(self._coord_core) + # ROCm keeps the copies until the kernel IO has been checked on AMD GPUs. + if torch.version.hip is None: + reason = _cpu_moe.kernel_io_probe(torch.cuda.current_stream().cuda_stream) + if reason: + logger.warning_rank0(f"cpu-moe kernel IO unavailable ({reason}); using D2H/H2D copies around the doorbell") + else: + self._ext.enable_kernel_io() + self._kernel_io = True self._watchdog_stop = False # The thread target holds a WEAKREF and re-derefs it each tick: a bound # method would strong-reference the executor forever (the loop never ends @@ -325,32 +331,14 @@ def __init__( ) self._watchdog.start() - # ds_fp4: the reference FP8 activation round-trip is a scalar per-element chain - # that the C++ side runs single-threaded on the CUDA host-callback thread -- - # straight on the decode critical path (~0.3ms/layer at H=4096, every worker and - # the GPU waiting on it). When a GPU is present we run the numerically identical - # round-trip as a captured GPU kernel BEFORE the D2H (see decode_submit) and tell - # the C++ side to skip its own. Measured on DeepSeek-V4-Flash bs=1 decode: - # 12.85 -> 15.65 tok/s, output bit-identical (tests/moe/test_dsfp4_prequant.py). - self._gpu_prequant = fmt == "ds_fp4" and device.type == "cuda" - # the W4A8 activation round-trip block follows the checkpoint's fp8 block (128 on V4, 32 on V4.1) - self._act_block = act_block - if fmt == "ds_fp4": - assert act_block is not None, "ds_fp4 experts need act_block, the checkpoint's fp8 activation quant block" - self._ext.set_act_block(act_block) - if self._gpu_prequant: - self._ext.set_input_prequant(True) - logger.info_rank0( - "ds_fp4: input FP8 round-trip moved to the GPU " - "(bit-identical grid; the CPU-side scalar round-trip is skipped)" - ) logger.info_rank0( f"CPU MoE executor ready: threads={nthreads} (pinned to cores " f"{core_ids[0]}..{core_ids[-1]}) isa={self.isa} fmt={fmt} " f"H={self.H} I={self.I} experts={self.num_experts} layers={self.num_layers} " f"top_k={self.top_k} act={activation} max_tokens={self.max_tokens} " - f"spin_wait={self.spin_wait}" + f"spin_wait={self.spin_wait} " + f"handshake={'kernel' if self._kernel_io else 'memop' if self._flag_sync else 'host-func'}" ) def _make_table(self, layers: list[torch.Tensor]) -> torch.Tensor: @@ -368,7 +356,11 @@ def _make_table(self, layers: list[torch.Tensor]) -> torch.Tensor: self._banks.extend(layers) return table - def _resolve_fp8_block_banks(self, banks: dict) -> tuple[dict, tuple[int, int]]: + def _tables(self, **roles: list[torch.Tensor]) -> dict[str, int]: + """Per-layer table addresses keyed by canonical bank role.""" + return {role: self._make_table(layers).data_ptr() for role, layers in roles.items()} + + def _resolve_fp8_block_banks(self, banks: dict) -> tuple[dict, dict, tuple[int, int]]: """Block-FP8 schema: row-major e4m3 + one bf16 scale per 128x128 block; the scale banks pad their last dim, so C++ gets the real row stride.""" gu, gus = banks["gate_up"], banks["gate_up_scale"] dn, dns = banks["down"], banks["down_scale"] @@ -380,28 +372,19 @@ def _resolve_fp8_block_banks(self, banks: dict) -> tuple[dict, tuple[int, int]]: assert tuple(dn[0].shape[1:]) == (H, I), (dn[0].shape, H, I) assert gus[0].shape[1] == (2 * I + 127) // 128 and gus[0].shape[2] >= (H + 127) // 128, (gus[0].shape, H, I) assert dns[0].shape[1] == (H + 127) // 128 and dns[0].shape[2] >= (I + 127) // 128, (dns[0].shape, H, I) - ptrs = dict( - gate_up_ptr=self._make_table(gu).data_ptr(), - down_ptr=self._make_table(dn).data_ptr(), - gate_up_scale_ptr=self._make_table(gus).data_ptr(), - gate_up_global_ptr=0, - down_scale_ptr=self._make_table(dns).data_ptr(), - down_global_ptr=0, - gate_up_bias_ptr=0, - down_bias_ptr=0, - fp8_block_gu_scale_stride=int(gus[0].stride(1)), - fp8_block_dn_scale_stride=int(dns[0].stride(1)), - ) - return ptrs, (H, I) + banks = self._tables(gate_up=gu, down=dn, gate_up_scale=gus, down_scale=dns) + options = dict(gu_scale_stride=int(gus[0].stride(1)), dn_scale_stride=int(dns[0].stride(1))) + return banks, options, (H, I) - def _resolve_banks(self, banks: dict, fmt: str) -> tuple[dict, tuple[int, int]]: - """Return (pointer kwargs for the C++ ctor, (H, I)) for the given format. + def _resolve_banks(self, banks: dict, fmt: str) -> tuple[dict, dict, tuple[int, int]]: + """Return (banks, options, (H, I)) for the C++ ctor: banks maps each canonical role + the format reads to its per-layer table's address, options holds its integer settings. ``banks[name]`` is a list of ``num_layers`` ``[num_experts, ...]`` tensors (the per-layer host bank contract); shapes are read from the first layer so - per-partition (TP) sizes are exact. Unused pointers are 0. Every pointer kwarg - is actually a per-layer table's address (see ``_make_table``), not a single - bank's -- the C++ ctor resolves ``tbl[layer_id]`` per task. + per-partition (TP) sizes are exact. Every bank address is a per-layer table + (see ``_make_table``), not a single bank's -- the C++ side resolves + ``tbl[layer_id]`` per task. """ if fmt == "bf16": gate_up = banks["gate_up"] @@ -414,17 +397,8 @@ def _resolve_banks(self, banks: dict, fmt: str) -> tuple[dict, tuple[int, int]]: I = int(gate_up[0].shape[1] // 2) assert gate_up[0].shape[1] == 2 * I assert tuple(down[0].shape[1:]) == (H, I), (down[0].shape, H, I) - ptrs = dict( - gate_up_ptr=self._make_table(gate_up).data_ptr(), - down_ptr=self._make_table(down).data_ptr(), - gate_up_scale_ptr=0, - gate_up_global_ptr=0, - down_scale_ptr=0, - down_global_ptr=0, - gate_up_bias_ptr=0, - down_bias_ptr=0, - ) - return ptrs, (H, I) + banks = self._tables(gate_up=gate_up, down=down) + return banks, {}, (H, I) if fmt == "q4_0": return self._resolve_q4_0_banks(banks) @@ -452,19 +426,10 @@ def _resolve_banks(self, banks: dict, fmt: str) -> tuple[dict, tuple[int, int]]: assert tuple(gus[0].shape[1:]) == (2 * I, H // 16), (gus[0].shape, I, H) assert tuple(dns[0].shape[1:]) == (H, I // 16), (dns[0].shape, H, I) assert tuple(gug[0].shape[1:]) == (2 * I,) and tuple(dng[0].shape[1:]) == (H,) - ptrs = dict( - gate_up_ptr=self._make_table(gup).data_ptr(), - down_ptr=self._make_table(dnp).data_ptr(), - gate_up_scale_ptr=self._make_table(gus).data_ptr(), - gate_up_global_ptr=self._make_table(gug).data_ptr(), - down_scale_ptr=self._make_table(dns).data_ptr(), - down_global_ptr=self._make_table(dng).data_ptr(), - gate_up_bias_ptr=0, - down_bias_ptr=0, - ) - return ptrs, (H, I) + banks = self._tables(gate_up=gup, down=dnp, gate_up_scale=gus, gate_up_global=gug, down_scale=dns, down_global=dng) + return banks, {}, (H, I) - def _resolve_q4_0_banks(self, banks: dict) -> tuple[dict, tuple[int, int]]: + def _resolve_q4_0_banks(self, banks: dict) -> tuple[dict, dict, tuple[int, int]]: """Native GGUF Q4_0 schema (gemma4 GGUF): per-32 blocks (fp16 scale + 16 nibble bytes), row-major over K -- the *same* packed banks the GPU offload path streams. gate_up is [S, 2I, H//32*18], down is [S, H, I//32*18]; the C++ W4A16 GEMV reads a @@ -479,19 +444,10 @@ def _resolve_q4_0_banks(self, banks: dict) -> tuple[dict, tuple[int, int]]: assert H % 32 == 0 and I % 32 == 0, (H, I) assert int(gate_up[0].shape[2]) == (H // 32) * 18, (gate_up[0].shape, H) assert int(down[0].shape[2]) == (I // 32) * 18, (down[0].shape, I) - ptrs = dict( - gate_up_ptr=self._make_table(gate_up).data_ptr(), - down_ptr=self._make_table(down).data_ptr(), - gate_up_scale_ptr=0, - gate_up_global_ptr=0, - down_scale_ptr=0, - down_global_ptr=0, - gate_up_bias_ptr=0, - down_bias_ptr=0, - ) - return ptrs, (H, I) + banks = self._tables(gate_up=gate_up, down=down) + return banks, {}, (H, I) - def _resolve_mxfp4_banks(self, banks: dict) -> tuple[dict, tuple[int, int]]: + def _resolve_mxfp4_banks(self, banks: dict) -> tuple[dict, dict, tuple[int, int]]: """gpt-oss mxfp4 ``mxfp4_triton`` schema: transposed split-K blocks/scales (N innermost) + per-output-row biases. The C++ kernel streams K and accumulates a contiguous N-block, so the GPU-tiled layout is read in place @@ -510,19 +466,10 @@ def _resolve_mxfp4_banks(self, banks: dict) -> tuple[dict, tuple[int, int]]: assert tuple(gus[0].shape[1:]) == (H // 32, 2 * I), (gus[0].shape, H, I) assert tuple(dns[0].shape[1:]) == (I // 32, H), (dns[0].shape, H, I) assert tuple(gob[0].shape[1:]) == (2 * I,) and tuple(dob[0].shape[1:]) == (H,) - ptrs = dict( - gate_up_ptr=self._make_table(gub).data_ptr(), - down_ptr=self._make_table(dnb).data_ptr(), - gate_up_scale_ptr=self._make_table(gus).data_ptr(), - gate_up_global_ptr=0, - down_scale_ptr=self._make_table(dns).data_ptr(), - down_global_ptr=0, - gate_up_bias_ptr=self._make_table(gob).data_ptr(), - down_bias_ptr=self._make_table(dob).data_ptr(), - ) - return ptrs, (H, I) + banks = self._tables(gate_up=gub, down=dnb, gate_up_scale=gus, down_scale=dns, gate_up_bias=gob, down_bias=dob) + return banks, {}, (H, I) - def _resolve_dsfp4_banks(self, banks: dict) -> tuple[dict, tuple[int, int]]: + def _resolve_dsfp4_banks(self, banks: dict) -> tuple[dict, dict, tuple[int, int]]: """DeepSeek-V4 ``ds_fp4`` schema: row-major e2m1 (2/byte) + e8m0 per-32 block scales, no global, no bias. Layout matches nvfp4 (K contiguous per output row), so the C++ GEMV reads it in place. The kernel additionally FP8-round-trips the @@ -538,57 +485,8 @@ def _resolve_dsfp4_banks(self, banks: dict) -> tuple[dict, tuple[int, int]]: assert tuple(dnp[0].shape[1:]) == (H, I // 2), (dnp[0].shape, H, I) assert tuple(gus[0].shape[1:]) == (2 * I, H // 32), (gus[0].shape, I, H) assert tuple(dns[0].shape[1:]) == (H, I // 32), (dns[0].shape, H, I) - ptrs = dict( - gate_up_ptr=self._make_table(gup).data_ptr(), - down_ptr=self._make_table(dnp).data_ptr(), - gate_up_scale_ptr=self._make_table(gus).data_ptr(), - gate_up_global_ptr=0, - down_scale_ptr=self._make_table(dns).data_ptr(), - down_global_ptr=0, - gate_up_bias_ptr=0, - down_bias_ptr=0, - ) - return ptrs, (H, I) - - def _io_for(self, bs: int) -> dict[str, torch.Tensor]: - io = self._io.get(bs) - if io is None: - alloc_io = ( - (lambda *shape, dtype: torch.empty(*shape, dtype=dtype)) - if self.device.type == "cpu" - else alloc_pinned_tensor - ) - io = { - "x": alloc_io(bs, self.H, dtype=torch.bfloat16), - "ids": alloc_io(bs, self.top_k, dtype=torch.int32), - "w": alloc_io(bs, self.top_k, dtype=torch.float32), - "y": alloc_io(bs, self.H, dtype=torch.bfloat16), - } - self._io[bs] = io - return io - - def _task_for(self, layer_id: int, bs: int) -> int: - key = (layer_id, bs) - task = self._tasks.get(key) - if task is None: - io = self._io_for(bs) - task = self._ext.create_task( - layer_id, - bs, - io["x"].data_ptr(), - io["ids"].data_ptr(), - io["w"].data_ptr(), - io["y"].data_ptr(), - ) - self._tasks[key] = task - # Allocate this (layer, bs) combo a flag slot and register its task with the - # coordinator. Combos past the slot capacity keep the host-func path. - if self._flag_sync and key not in self._flag_slots: - slot = len(self._flag_slots) - if slot < self._flag_capacity: - self._flag_slots[key] = slot - self._ext.register_flag_task(slot, task) - return task + banks = self._tables(gate_up=gup, down=dnp, gate_up_scale=gus, down_scale=dns) + return banks, {}, (H, I) def decode( self, @@ -612,7 +510,7 @@ def decode_submit( topk_weights: torch.Tensor, topk_ids: torch.Tensor, ) -> tuple: - """Issue the D2H copies + the CPU-pool submit host node, then return without + """Issue the D2H copies + the CPU-pool submit host node (or the publish kernel), then return without waiting. Lets a caller (the hybrid backend) enqueue GPU work between this and :meth:`decode_sync` so the CPU compute overlaps the GPU GEMM / PCIe fetch. @@ -621,7 +519,14 @@ def decode_submit( :meth:`decode_sync`. The output tensor is allocated here so it stays live (and distinct from the interleaved GPU work) across the overlap window.""" bs = hidden_states.shape[0] - io = self._io_for(bs) + if bs > self.max_tokens: + raise ValueError(f"CPU MoE decode batch {bs} exceeds the executor's max_tokens {self.max_tokens}") + if not self._graph_capture_safe and torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "CPU MoE decode cannot be captured into a CUDA graph here: the stream-memop " + "doorbell is off or failed its startup probe, and the host-function sync is not " + "replay-safe on this runtime; run with --cuda-graph-max-bs 0" + ) if self._gpu_prequant: # DSV4: apply the reference FP8 round-trip on the GPU (the same kernel the @@ -631,91 +536,86 @@ def decode_submit( hidden_states = act_quant_fp8_roundtrip(hidden_states, block=self._act_block) + stream = torch.cuda.current_stream().cuda_stream + if self._kernel_io and hidden_states.dtype == torch.bfloat16: + x = hidden_states.contiguous() + ids = topk_ids.to(torch.int32).contiguous() + w = topk_weights.to(torch.float32).contiguous() + assert x.shape == (bs, self.H) and ids.shape == w.shape == (bs, self.top_k), (x.shape, ids.shape, w.shape) + self._ext.flag_publish(stream, x.data_ptr(), ids.data_ptr(), w.data_ptr(), layer_id, bs) + return (bs, torch.empty_like(x), True) + # D2H: ship this step's activations + routing to pinned host memory. - io["x"].copy_(hidden_states, non_blocking=True) - io["ids"].copy_(topk_ids.to(torch.int32), non_blocking=True) - io["w"].copy_(topk_weights.to(torch.float32), non_blocking=True) + io = self._io + io["x"][:bs].copy_(hidden_states, non_blocking=True) + io["ids"][:bs].copy_(topk_ids.to(torch.int32), non_blocking=True) + io["w"][:bs].copy_(topk_weights.to(torch.float32), non_blocking=True) - task = self._task_for(layer_id, bs) out = torch.empty_like(hidden_states) - slot = self._flag_slots.get((layer_id, bs)) if self._flag_sync else None - if slot is not None: - # Front-end memops: done[slot]=0 then ready[slot]=1 (the coordinator's - # doorbell). No kernel launched; no host-func round trip. - self._cpu_moe.memop_submit( - torch.cuda.current_stream().cuda_stream, - self._done.data_ptr(), self._ready.data_ptr(), slot, - ) + if self._flag_sync: + # Front-end memops: reset done, then ring the doorbell with (layer, bs). + # No kernel launched; no host-func round trip. + self._ext.flag_signal(stream, layer_id, bs) else: - stream = torch.cuda.current_stream().cuda_stream - self._ext.submit_with_cuda_stream(stream, task) - return (bs, task, out, slot) + self._ext.submit_with_cuda_stream(stream, layer_id, bs) + return (bs, out, False) - def decode_sync(self, pending: tuple) -> torch.Tensor: + def decode_sync(self, pending: tuple, add_to: torch.Tensor | None = None) -> torch.Tensor: """Issue the CPU-pool sync + the H2D result copy for a prior :meth:`decode_submit`, and return the GPU output tensor. With flag-sync the wait is a front-end stream - memop on done[slot] (set by the CPU coordinator); otherwise a cudaLaunchHostFunc.""" - bs, task, out, slot = pending - if slot is not None: - # Front-end WAIT(done[slot] >= 1): blocks this stream's later nodes without + memop on done (set by the CPU coordinator); otherwise a cudaLaunchHostFunc. + With ``add_to``, return ``add_to`` + the CPU output; the kernel path adds in place into ``add_to``.""" + bs, out, kernel_io = pending + stream = torch.cuda.current_stream().cuda_stream + if self._flag_sync: + # Front-end WAIT(done >= 1): blocks this stream's later nodes without # occupying an SM, so GPU utilization stays truthful during the CPU window. - self._cpu_moe.memop_sync( - torch.cuda.current_stream().cuda_stream, self._done.data_ptr(), slot, - ) + self._ext.flag_wait(stream) else: - stream = torch.cuda.current_stream().cuda_stream - self._ext.sync_with_cuda_stream(stream, task) - io = self._io[bs] - out.copy_(io["y"], non_blocking=True) - return out - - def _watchdog_tick(self, suspects: dict) -> None: + self._ext.sync_with_cuda_stream(stream) + if kernel_io: + if add_to is not None and add_to.dtype == torch.bfloat16 and add_to.is_contiguous() and add_to.shape == out.shape: + self._ext.flag_combine(stream, add_to.data_ptr(), bs, True) + return add_to + self._ext.flag_combine(stream, out.data_ptr(), bs, False) + return out if add_to is None else add_to + out + out.copy_(self._io["y"][:bs], non_blocking=True) + return out if add_to is None else add_to + out + + def _watchdog_tick(self, suspect: dict) -> None: """One watchdog sampling round (called every 2 s by ``_watchdog_main``). - A slot is only declared dead when THREE things hold across >=10 s: its doorbell - is still pending (ready==1 && done==0), it was already pending when first - suspected, and the coordinator has served NOTHING on it since (flag_served_count - unchanged). The served-count criterion kills the false-positive window: two - point samples can land on the same slot's (independent, us-scale) pending - windows under heavy external load, but a coordinator that made progress in - between is alive by definition. ``suspects`` maps slot -> (first_seen, - served_at_first_sight) and persists across ticks.""" - stuck = (self._ready == 1) & (self._done == 0) - if not bool(stuck.any()): - suspects.clear() + The doorbell is only declared dead when THREE things hold across >=10 s: it is + still pending (ready set, done==0), it was already pending when first suspected, and + the coordinator has served NOTHING since (flag_served unchanged). The served-count + criterion kills the false-positive window: two point samples can land on + (independent, us-scale) pending windows under heavy external load, but a + coordinator that made progress in between is alive by definition. ``suspect`` + holds (first_seen, served_at_first_sight) and persists across ticks.""" + if not self._ext.flag_pending(): + suspect.clear() return now = time.monotonic() - pending = set(stuck.nonzero().flatten().tolist()) - for slot in list(suspects): - if slot not in pending: - del suspects[slot] - dead = [] - for slot in pending: - served = self._ext.flag_served_count(slot) - first_seen, served_then = suspects.get(slot, (None, None)) - if first_seen is None or served != served_then: - suspects[slot] = (now, served) # new suspect, or alive-but-loaded: rearm - continue - if now - first_seen >= 10.0: - dead.append(slot) - if not dead: + served = self._ext.flag_served() + if not suspect or suspect["served"] != served: + suspect.update(first_seen=now, served=served) # new suspect, or alive-but-loaded: rearm + return + if now - suspect["first_seen"] < 10.0: return logger.error( - f"cpu-moe flag watchdog: slots {dead} unanswered for >10s with no coordinator " - "progress (wedged/dead); poisoning done[] and failing the next step" + "cpu-moe flag watchdog: doorbell unanswered for >10s with no coordinator " + "progress (wedged/dead); poisoning done and failing the next step" ) - for i in dead: - self._err[i] = 1 - for i in dead: - self._done[i] = 1 # after err: unblock the stream into a checked failure - suspects.pop(i, None) + self._unhealthy = True + self._ext.poison_flag() # after the mark: unblock the stream into a checked failure + suspect.clear() def raise_if_unhealthy(self) -> None: """Raise if the flag watchdog fired (a doorbell stayed unanswered because the coordinator never responded). Called by the engine once per forward -- a single - pinned read -- so a dead coordinator surfaces as a loud error on the next step + flag read -- so a dead coordinator surfaces as a loud error on the next step instead of silently shipping stale expert outputs.""" - if self._err is not None and bool((self._err != 0).any()): + if self._unhealthy: raise RuntimeError( "CPU MoE flag-handshake watchdog fired: a decode step's doorbell was " "never answered by the coordinator thread (its outputs cannot be " @@ -728,13 +628,13 @@ def raise_if_unhealthy(self) -> None: def _watchdog_main(executor_ref) -> None: """Watchdog daemon body: weakref-deref per tick so the thread never keeps a dead executor alive (see the start site in ``CpuMoeExecutor.__init__``).""" - suspects: dict = {} + suspect: dict = {} while True: time.sleep(2.0) executor = executor_ref() - if executor is None or executor._watchdog_stop or executor._ready is None: + if executor is None or executor._watchdog_stop: return try: - executor._watchdog_tick(suspects) + executor._watchdog_tick(suspect) finally: del executor # drop the strong ref before the next sleep diff --git a/python/freetoken/moe/expert_banks.py b/python/freetoken/moe/expert_banks.py index 7500174ae4..114884c489 100644 --- a/python/freetoken/moe/expert_banks.py +++ b/python/freetoken/moe/expert_banks.py @@ -11,6 +11,7 @@ import glob import math import os +from contextlib import contextmanager, nullcontext from dataclasses import dataclass, field import torch @@ -27,6 +28,19 @@ _PARALLEL_READER_SUPPORTED = hasattr(os, "O_DIRECT") and hasattr(os, "preadv") +@contextmanager +def _single_threaded_host_copies(): + """Keep small NVFP4 bank placements off PyTorch's intra-op worker pool.""" + previous = torch.get_num_threads() + if previous > 1: + torch.set_num_threads(1) + try: + yield + finally: + if previous > 1: + torch.set_num_threads(previous) + + @dataclass(frozen=True) class ExpertBanks: """Loaded expert banks, normalized for ``OffloadMoeCache`` wiring.""" @@ -48,6 +62,8 @@ class ExpertBanks: kind: QuantKind | None = None kernel: str | None = None layout: dict | None = None + # per-layer quant tag (e.g. the gate_up and down GGUF types), for layers whose banks match in dtype and row shape but not in quant; None = one quant + layer_tags: list | None = None def _dummy_fill(role: str, tensor: torch.Tensor) -> None: @@ -141,13 +157,22 @@ def _fill(sink) -> None: if missing: raise ValueError(f"expert banks were not filled: {len(missing)} (layer, expert) rows missing (first {missing[:4]})") - if layer_sink is not None: - _fill(layer_sink) - elif torch.cuda.is_available() and not resident: - with PinPipeline() as pins: - _fill(pins) - else: - _fill(None) + # Per-expert NVFP4 pieces produce thousands of small CPU copies. On high-core + # hosts the intra-op pool costs more than the copies and contends with the + # checkpoint reader. Restore its setting even if packing or pinning fails. + copy_threads = ( + _single_threaded_host_copies() + if method.kind is QuantKind.NVFP4 and not resident + else nullcontext() + ) + with copy_threads: + if layer_sink is not None: + _fill(layer_sink) + elif torch.cuda.is_available() and not resident: + with PinPipeline() as pins: + _fill(pins) + else: + _fill(None) return ExpertBanks( legacy_format_for(method.kind, kernel.name), banks, diff --git a/python/freetoken/moe/fused.py b/python/freetoken/moe/fused.py index a90aa4c792..2f8af47c89 100644 --- a/python/freetoken/moe/fused.py +++ b/python/freetoken/moe/fused.py @@ -44,6 +44,9 @@ def fused_topk( return fused_topk_softmax(gating_output, topk, renormalize, num_token_non_padded) +_SGL_ALIGN_MAX_EXPERTS = 1024 + + def moe_align_block_size( topk_ids: torch.Tensor, block_size: int, num_experts: int ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: @@ -84,9 +87,23 @@ def moe_align_block_size( - The padding ensures that the total number of tokens is now divisible by block_size for proper block matrix operations. """ - from freetoken.kernel.backend import is_sgl_kernel_installed + from freetoken.kernel.backend import is_rocm, is_sgl_kernel_installed + + if is_rocm(): + import triton + from packaging.version import Version + + # FIXME: remove once Triton 3.8 is unsupported. Its AMD backend folds compares on tl.histogram results to false. + # https://github.com/triton-lang/triton/pull/11246 + if Version(triton.__version__).release[:2] == (3, 8): + from freetoken.kernel import moe_align_block_size_triton + + return moe_align_block_size_triton(topk_ids, block_size, num_experts) - if not is_sgl_kernel_installed(): + # sgl_kernel's align kernel takes num_experts + 1 counters in one block and reads out of + # bounds past 1024 of them (0.4.5 and 0.4.9). A prefill over an offload slot cache passes the + # slot count here (#676), which is thousands; the Triton kernel has no such limit. + if not is_sgl_kernel_installed() or num_experts + 1 > _SGL_ALIGN_MAX_EXPERTS: from freetoken.kernel.triton.moe_align import ( moe_align_block_size as triton_moe_align_block_size, ) diff --git a/python/freetoken/moe/fused_nvfp4.py b/python/freetoken/moe/fused_nvfp4.py index 375d2a547a..e8ad477021 100644 --- a/python/freetoken/moe/fused_nvfp4.py +++ b/python/freetoken/moe/fused_nvfp4.py @@ -270,7 +270,6 @@ def _prefill_gemm( _prefill_nvfp4_moe_kernel[grid]( a, packed, scale, glob, c, topk_weights_flat, sorted_ids, expert_ids, num_tokens_post_padded, - _e2m1_lut(a.device.index), N, K, EM, num_valid_tokens, a.stride(0), a.stride(1), packed.stride(0), packed.stride(1), packed.stride(2), @@ -281,6 +280,7 @@ def _prefill_gemm( MUL_ROUTED_WEIGHT=mul_routed_weight, top_k=kernel_top_k, compute_type=_tl_dtype(c.dtype), + EVEN_K=(K // 2) % cfg["BLOCK_SIZE_KB"] == 0, **cfg, ) diff --git a/python/freetoken/moe/offload_cache.py b/python/freetoken/moe/offload_cache.py index 2afbb1a79f..f7eb6c659e 100644 --- a/python/freetoken/moe/offload_cache.py +++ b/python/freetoken/moe/offload_cache.py @@ -2,8 +2,10 @@ import math import os +import weakref +from contextlib import contextmanager from dataclasses import dataclass -from typing import Iterator +from typing import TYPE_CHECKING, Iterator import torch from flashlib.kernels.slot_cache import N_STATS, Stat @@ -13,6 +15,9 @@ # (kept for A/B profiling). Falls back to per-bank automatically if a bank's row bytes or # base address are not 16-byte aligned. _FUSED_COPY = os.getenv("FREETOKEN_FUSED_COPY", "1").strip().lower() not in {"0", "false", "no", "off"} +# Opt-in: run the fused copy_missing on the slim kernel (a few small CTAs, leaves most SMs to +# concurrent compute) instead of the 8-CTAs-per-bank one. Off by default. +_SLIM_COPY = os.getenv("FREETOKEN_MOE_SLIM_COPY", "0").strip().lower() in {"1", "true", "yes", "on"} # cudaMemcpyBatchAsync silently degrades to a SYNCHRONOUS copy when a batch mixes # large entries with sub-~256KB entries on registered host memory (H100 + CUDA 13.0, @@ -26,8 +31,12 @@ _SMALL_BANK_FEAT_BYTES = 256 * 1024 from freetoken.kernel.fast_index_copy import default_blocks_per_bank +from freetoken.env import ENV from freetoken.utils import init_logger +if TYPE_CHECKING: + from freetoken.moe.offload_kernels import HybridRoutes + logger = init_logger(__name__) # quant_format -> bank names, in registration order: the single place a format's bank @@ -74,7 +83,7 @@ "down_bias", ), # DeepSeek-V4 FP4: packed e2m1 codes + e8m0 per-32 block scales, no global scale - # (4 banks). Read by DeepSeek-V4's own DS-FP4 grouped GEMV kernels via bank_views(). + # (4 banks). Read by DeepSeek-V4's own DS-FP4 grouped GEMV kernels via pool_bank_views(). "ds_fp4": ("gate_up_packed", "gate_up_scale", "down_packed", "down_scale"), } @@ -101,6 +110,14 @@ MARLIN_MAX_CACHE_SIZE = 992 +def bank_layouts(sources: dict[str, list[torch.Tensor]], layer_tags: list | None = None) -> list[tuple]: + """Each layer's layout key: its quant tag and every bank's (dtype, row shape). A slot pool holds layers of one key.""" + num_layers = len(next(iter(sources.values()))) + tags = layer_tags if layer_tags is not None else [None] * num_layers + assert len(tags) == num_layers, (len(tags), num_layers) + return [(tags[l], tuple((t[l].dtype, tuple(t[l].shape[1:])) for t in sources.values())) for l in range(num_layers)] + + @dataclass class OffloadMoeCache: num_layers: int @@ -148,17 +165,44 @@ class OffloadMoeCache: # bank layout from the expert kernel (a BankSpec per role); when given it replaces the _BANK_SCHEMAS lookup and the slot cap comes from max_slots layout: dict | None = None max_slots: int | None = None + # Each layer belongs to one pool: a contiguous slot range with its own LRU scope; pool p owns slots [sum(pool_sizes[:p]), +pool_sizes[p]). Slot ids in slot_for_id are local to the pool. None = one pool for all layers. + pool_of_layer: list[int] | None = None + pool_sizes: list[int] | None = None + # GPU decode only: the caller's work that is independent of the routed experts (qwen4_exp: the + # shared expert) runs on decode_copy_stream beside the miss copy. None = FREETOKEN_MOE_COPY_OVERLAP. + decode_copy_overlap: bool | None = None + # GPU decode only: cross-layer expert prefetch (moe/prefetch.py). None = FREETOKEN_MOE_PREFETCH. + prefetch_mode: str | None = None + # GPU decode only: debug slot checks (moe/verify.py). None = FREETOKEN_MOE_PREFETCH_VERIFY. + verify_mode: str | None = None + # GPU decode only: debug slot writer history + audit every N decode steps (moe/slot_audit.py). + # None = FREETOKEN_MOE_SLOT_AUDIT, 0 = off. + slot_audit: int | None = None def __post_init__(self) -> None: - policy_ids = {"lru": 0} - assert self.cache_policy in policy_ids + from freetoken.moe.scored_ensure import POLICY_IDS + + assert self.cache_policy in POLICY_IDS, self.cache_policy assert self.decode_target in ("gpu", "cpu", "hybrid"), self.decode_target + if self.cache_policy != "lru" and self.decode_target != "gpu": + # the hybrid and CPU decode paths run their own kernels and would silently stay LRU + raise ValueError( + f"--moe-cache-policy {self.cache_policy} needs GPU decode, not {self.decode_target!r}" + ) if self.prefill_demand_tokens and self.decode_target != "gpu": raise ValueError( f"--moe-prefill-demand-tokens needs GPU decode; this cache decodes on {self.decode_target!r}" ) if self.layout is None: assert self.quant_format in _BANK_SCHEMAS, f"unknown quant_format {self.quant_format!r}" + if self.pool_sizes is None: + self.pool_of_layer = [0] * self.num_layers + self.pool_sizes = [self.cache_size] + self.pool_of_layer = list(self.pool_of_layer) + self.pool_sizes = list(self.pool_sizes) + assert len(self.pool_of_layer) == self.num_layers, (len(self.pool_of_layer), self.num_layers) + assert sorted(set(self.pool_of_layer)) == list(range(len(self.pool_sizes))), self.pool_of_layer + assert sum(self.pool_sizes) == self.cache_size and min(self.pool_sizes) > 0, (self.pool_sizes, self.cache_size) # Attached by the engine for decode_target == "cpu" (CpuMoeExecutor); None # for the GPU decode path. self.cpu_executor = None @@ -166,14 +210,26 @@ def __post_init__(self) -> None: # offload/PCIe path. Set by the engine after construction (empty = all-GPU, # all layers = the plain --moe-strategy cpu case). self.cpu_layer_ids: frozenset = frozenset() + # Bank row layouts, set by set_bank_sources: layouts[g][k] is (dtype, row shape) of bank k (layouts of two quant tags may be equal), layer_layout[l] is layer l's g. + # With one (dtype, row shape) per bank a bank cache is typed [cache_size, *row]; with several it is a byte segment each pool views with its own layout. + self.layouts: list[tuple[tuple[torch.dtype, tuple[int, ...]], ...]] = [] + self.layer_layout: list[int] = [0] * self.num_layers + self._raw_banks = False + # Set by the engine when hybrid decode hands CPU routes to a kernel that reads slot 0 with weight 0 (no supports_inactive_slots). + # With several quants a staging buffer can leave another quant's bytes in a pool's slot 0, which may decode to NaN, so the last prefill layer zeroes those rows. + self.zero_inactive_route_slots = False # num_experts floor + nvfp4_marlin slot cap, shared with the runtime-rebuild path. - self.validate_rebuild(self.cache_size) + self.validate_rebuild(self.cache_size, self.pool_sizes) assert not self.prefill_overlap or self.cache_size >= 2 * self.num_experts, ( "Prefill overlap borrows two full expert-layer buffers from the unified MoE " "cache, so cache_size must be at least 2 * num_experts " "(raise moe_cache_size or disable moe_prefill_overlap)" ) - self.cache_policy_id = policy_ids[self.cache_policy] + self.cache_policy_id = POLICY_IDS[self.cache_policy] + # rule: an unrouted expert within this many logits of the top-k-th is a near miss + self.near_miss_thr = ( + float(os.getenv("FREETOKEN_MOE_NEAR_MISS_THR", "0.25")) if self.cache_policy == "rule" else None + ) self.slot_for_id = torch.full( (self.num_layers, self.num_experts), -1, @@ -207,6 +263,7 @@ def __post_init__(self) -> None: self.expert_recency = torch.full( (self.num_layers, self.num_experts), -1, dtype=torch.int64, device=self.device ) + self._init_evict_state() # Host source banks (one [num_experts, ...] tensor per layer, so layers can # carry independent host attributes -- see layer_residency) and their GPU # slot caches, keyed by the format's bank schema (attached by @@ -231,24 +288,23 @@ def __post_init__(self) -> None: # is a one-off warm-up increment at capture time (<0.1% over a session). self.collect_stats = False # [num_layers, N_STATS] -- ensure_experts passes lru_stats[layer_id] straight to - # the kernel, which accumulates in the same launch. The stat_* tensors below stay - # for the hybrid path, whose kernel is still ours. + # the kernel, which accumulates in the same launch. The stat_* tensors below are + # the hybrid path's: views of one buffer its kernel accumulates into the same way. self.lru_stats = torch.zeros( (self.num_layers, N_STATS), dtype=torch.int64, device=self.device ) - self.stat_missing = torch.zeros((), dtype=torch.int64, device=self.device) - self.stat_active = torch.zeros((), dtype=torch.int64, device=self.device) - self.stat_calls = torch.zeros((), dtype=torch.int64, device=self.device) - # hybrid only: experts actually fetched over PCIe (<= stat_missing). The CPU - # computes stat_missing - stat_fetched of them. - self.stat_fetched = torch.zeros((), dtype=torch.int64, device=self.device) + # (missing, fetched, active, calls), then the same per layer + stat_buf = torch.zeros(4 * (self.num_layers + 1), dtype=torch.int64, device=self.device) + self.hybrid_stats = (stat_buf[:4], stat_buf[4:]) + # hybrid only: stat_fetched counts the experts actually fetched over PCIe (<= stat_missing). + # The CPU computes stat_missing - stat_fetched of them. + self.stat_missing, self.stat_fetched, self.stat_active, self.stat_calls = stat_buf[:4].unbind() # Per-layer counterparts of the scalars above (indexed by MoE-layer id). Same # device-side accumulation (graph-safe: layer_id is a static index per graph node), # so one req's per-layer miss rate is readable via decode_miss_stats_per_layer(). - self.stat_missing_layer = torch.zeros(self.num_layers, dtype=torch.int64, device=self.device) - self.stat_active_layer = torch.zeros(self.num_layers, dtype=torch.int64, device=self.device) - self.stat_fetched_layer = torch.zeros(self.num_layers, dtype=torch.int64, device=self.device) - self.stat_steps_layer = torch.zeros(self.num_layers, dtype=torch.int64, device=self.device) + per_layer = stat_buf[4:].view(self.num_layers, 4) + self.stat_missing_layer, self.stat_fetched_layer = per_layer[:, 0], per_layer[:, 1] + self.stat_active_layer, self.stat_steps_layer = per_layer[:, 2], per_layer[:, 3] # Opt-in decode routing histogram (per layer, per expert) for cache-skew # analysis. Accumulated in ``ensure_experts`` from the raw expert ids before the # kernel rewrites them to slots. Only accurate with CUDA graphs disabled (the @@ -259,25 +315,31 @@ def __post_init__(self) -> None: ) # (per-layer sources, cache) per bank, in schema order. Every piece of cache # machinery that moves bank bytes (copy_missing, the prefill double buffers, - # bank_views) iterates this list, so the slot cache is bank-count agnostic. + # the pool and staging views) iterates this list, so the slot cache is bank-count agnostic. self.banks: list[tuple[list[torch.Tensor], torch.Tensor]] = [] # Fused multi-bank copy descriptor (built by set_bank_sources/_build_copy_plan). # Source pointers are per layer (_copy_src_ptrs[layer_id] -> [num_banks] device - # tensor); dst/feat are layer-invariant. + # tensor); dst pointers are per pool and row bytes per layout. self._copy_fused_ok = False self._copy_blocks_per_bank = 8 self._copy_dst_ptrs: torch.Tensor | None = None + self._copy_dst_ptrs_pool: list[torch.Tensor] | None = None self._copy_src_ptrs: list[torch.Tensor] | None = None - self._copy_feat_bytes: torch.Tensor | None = None + self._copy_feat_bytes: list[torch.Tensor] | None = None # The layer whose misses ensure_experts/materialize_layer staged last; consumed # by copy_missing to pick the per-layer source (part of the same pending-copy # state as evict_slots/src_indices/num_indices). # _pending_whole_layer records WHICH staged it: the pageable branch is only sound after materialize_layer self._pending_src_layer: int | None = None self._pending_whole_layer = False - # Per-bank [2, num_experts, ...] double-buffer views over the slot cache's - # first 2 * num_experts slots (set up when prefill_overlap is enabled). - self.prefill_bank_buffers: list[torch.Tensor] = [] + # The two prefill staging buffers are the first 2 * _stage_bytes[k] bytes of every bank, on top of the first pools. + # _stage_views[g][b] are buffer b's per-bank [num_experts, ...] views for layout g; _inval[b] are the global slots it overwrites. + self._stage_bytes: list[int] = [] + self._stage_views: list[list[tuple[torch.Tensor, ...] | None]] = [] + self._inval: list[slice | torch.Tensor] = [] + # the slot audit's label for the staged copy, and the most entries its plan can have + self._pending_kind = 0 + self._pending_width: int | None = None self.prefill_copy_stream: torch.cuda.Stream | None = None self.prefill_begin_event: torch.cuda.Event | None = None self.prefill_ready_events: list[torch.cuda.Event] = [] @@ -287,7 +349,7 @@ def __post_init__(self) -> None: self._prefill_buffer_has_release_event: list[bool] = [False, False] # hit-D2D split state: pinned begin-of-chunk snapshot of slot_for_id (the # classification input; frozen for the chunk -- no decode runs inside one, - # and buffer invalidation only clears slot < 2E entries, which classify as + # and buffer invalidation only clears slots under the staging buffers, which classify as # miss regardless), the lazily resolved batch-memcpy entry point (False = # unavailable), and row counters for cache reports. self._prefill_slot_snapshot: torch.Tensor | None = None @@ -297,11 +359,294 @@ def __post_init__(self) -> None: self._batch_memcpy = None self.prefill_hit_rows = 0 self.prefill_total_rows = 0 + self._slice_pools() + # Decode copy overlap: one side stream plus a (forked, done) event pair per layer. + # Neither is cache_size-shaped, so rebuild() keeps them. + if self.decode_copy_overlap is None: + self.decode_copy_overlap = bool(ENV.MOE_COPY_OVERLAP) + self.decode_copy_stream: torch.cuda.Stream | None = None + self._decode_copy_events: list[tuple[torch.cuda.Event, torch.cuda.Event]] = [] + if self.decode_copy_overlap and self.device.type == "cuda" and self.decode_target == "gpu": + self._init_decode_copy_overlap() + self.prefetch = self._init_prefetch() + self.verify = self._init_verify() + self.audit = self._init_audit() + self._refuse_single_pool_features() + + def _refuse_single_pool_features(self) -> None: + """The scored policies, prefetch, its debug checks and the on-demand prefill GEMM address + the cache by global slot id over one typed bank, so with several slot pools (#689) they + would install and read pool-local slots in the wrong rows: refuse that combination.""" + if len(self.pool_sizes) == 1: + return + used = [ + name + for name, on in ( + (f"--moe-cache-policy {self.cache_policy}", self.cache_policy != "lru"), + ("FREETOKEN_MOE_PREFETCH", self.prefetch_mode not in (None, "off")), + ("FREETOKEN_MOE_PREFETCH_VERIFY", self.verify_mode not in (None, "off")), + ("FREETOKEN_MOE_SLOT_AUDIT", bool(self.slot_audit)), + ("--moe-prefill-demand-tokens", bool(self.prefill_demand_tokens)), + ) + if on + ] + if used: + raise ValueError( + f"{', '.join(used)} need one MoE slot pool, but this cache has {len(self.pool_sizes)} " + "(--moe-cache-scope per-layer, or experts in several quant formats)" + ) + + def _init_audit(self): + from freetoken.kernel.fast_index_copy import _skip_fast_index_copy_enabled + from freetoken.moe.slot_audit import SlotAuditor, resolve_interval + + from_env = self.slot_audit is None + self.slot_audit = resolve_interval(self.slot_audit) + if not self.slot_audit or self.device.type != "cuda": + return None + if self.decode_target != "gpu": + if not from_env: + raise ValueError(f"FREETOKEN_MOE_SLOT_AUDIT={self.slot_audit} needs GPU decode, not {self.decode_target!r}") + logger.warning(f"FREETOKEN_MOE_SLOT_AUDIT={self.slot_audit} ignored: decode target is {self.decode_target!r}") + self.slot_audit = 0 + return None + if _skip_fast_index_copy_enabled(): + logger.warning("FREETOKEN_MOE_SLOT_AUDIT ignored: FREETOKEN_SKIP_FAST_INDEX_COPY makes every copy a no-op") + self.slot_audit = 0 + return None + return SlotAuditor(self.num_layers, self.num_experts, self.device, interval=self.slot_audit) + + def _init_verify(self): + from freetoken.moe.verify import ExpertVerifier, resolve_mode + + from_env = self.verify_mode is None + self.verify_mode = resolve_mode(self.verify_mode) + if self.verify_mode == "off" or self.device.type != "cuda": + return None + if self.decode_target != "gpu": + if not from_env: + raise ValueError(f"FREETOKEN_MOE_PREFETCH_VERIFY={self.verify_mode} needs GPU decode, not {self.decode_target!r}") + logger.warning(f"FREETOKEN_MOE_PREFETCH_VERIFY={self.verify_mode} ignored: decode target is {self.decode_target!r}") + self.verify_mode = "off" + return None + return ExpertVerifier(self.num_layers, self.num_experts, self.device, mode=self.verify_mode) + + def _init_prefetch(self): + from freetoken.moe.prefetch import ExpertPrefetcher, resolve_mode + + from_env = self.prefetch_mode is None + self.prefetch_mode = resolve_mode(self.prefetch_mode) + if self.prefetch_mode == "off": + return None + if self.decode_target != "gpu": + if not from_env: + raise ValueError(f"FREETOKEN_MOE_PREFETCH={self.prefetch_mode} needs GPU decode, not {self.decode_target!r}") + # the process-wide flag skips hybrid/CPU-decoded caches, like FREETOKEN_MOE_COPY_OVERLAP + logger.warning(f"FREETOKEN_MOE_PREFETCH={self.prefetch_mode} ignored: decode target is {self.decode_target!r}") + self.prefetch_mode = "off" + return None + if self.device.type != "cuda": + return None + prefetch = ExpertPrefetcher(self.num_layers, self.num_experts, self.device, mode=self.prefetch_mode) + # a proxy, not a cycle: dropping the cache must free its slot banks without waiting for gc + prefetch.cache = weakref.proxy(self) + return prefetch + + @property + def prefetch_on(self) -> bool: + """FREETOKEN_MOE_PREFETCH=on: every ensure keys prefetched slots low priority.""" + return self.prefetch is not None and self.prefetch.mode == "on" + + def _pool_layout(self, pool: int) -> int: + return self.layer_layout[self.pool_of_layer.index(pool)] + + def _bank_rows(self, k: int, layout: int, offset: int, nrows: int) -> torch.Tensor: + """``nrows`` rows of bank ``k`` in ``layout``, starting ``offset`` bytes into the bank cache.""" + cache = self.banks[k][1] + row = self._row_bytes[layout][k] + if not self._raw_banks: + first = offset // row + return cache if first == 0 and nrows == cache.shape[0] else cache[first : first + nrows] + dtype, shape = self.layouts[layout][k] + return cache[offset : offset + nrows * row].view(dtype).view(nrows, *shape) + + def _stage_fits(self, buffers: int) -> bool: + return all(c.numel() * c.element_size() >= buffers * s for (_, c), s in zip(self.banks, self._stage_bytes)) + + def _slice_pools(self) -> None: + E = self.num_experts + bases = [0] + for size in self.pool_sizes[:-1]: + bases.append(bases[-1] + size) + self.pool_bases = bases + spans = list(zip(bases, self.pool_sizes)) + + def cut(t: torch.Tensor, b: int, s: int) -> torch.Tensor: + return t if s == t.shape[0] else t[b : b + s] + + self._pool_id_of_slot = [cut(self.id_of_slot, b, s) for b, s in spans] + self._pool_usage = [cut(self.usage, b, s) for b, s in spans] + # rows[k][p]: bytes of one slot of pool p in bank k; before the banks exist, one unit per slot with a stage of E + if self.banks: + rows = [[self._row_bytes[self._pool_layout(p)][k] for p in range(len(spans))] for k in range(len(self.banks))] + stage = self._stage_bytes + else: + rows, stage = [[1] * len(spans)], [E] + offsets = [] + for row in rows: + offsets.append([0]) + for s, r in zip(self.pool_sizes[:-1], row): + offsets[-1].append(offsets[-1][-1] + s * r) + self._pool_offsets = offsets + # local slots below this index overlap the prefill staging area in some bank + self._pool_stable_from = [ + max(min(max(-(-(2 * stage[k] - offsets[k][p]) // rows[k][p]), 0), s) for k in range(len(rows))) + for p, s in enumerate(self.pool_sizes) + ] + self._inval = [] + masks = [] + for b in range(2): + mask = torch.zeros(self.cache_size, dtype=torch.bool) + masks.append(mask) + for k in range(len(rows)): + lo, hi = b * stage[k], (b + 1) * stage[k] + for p, (base, s) in enumerate(spans): + first = max((lo - offsets[k][p]) // rows[k][p], 0) + end = min(-(-(hi - offsets[k][p]) // rows[k][p]), s) + mask[base + first : base + max(first, end)] = True + idx = mask.nonzero().flatten() + if idx.numel() == 0 or int(idx[-1] - idx[0]) + 1 == idx.numel(): + self._inval.append(slice(int(idx[0]), int(idx[-1]) + 1) if idx.numel() else slice(0, 0)) + else: + self._inval.append(idx.to(self.device)) + self._stage_dirty = [False, False] + self._stage_slot0_rows = [[], []] + if not self.banks: + self._pool_bank_views, self._stage_views = [], [] + return + self._pool_bank_views = [ + tuple(self._bank_rows(k, self._pool_layout(p), offsets[k][p], s) for k in range(len(self.banks))) + for p, s in enumerate(self.pool_sizes) + ] + self._stage_views = [ + [ + tuple(self._bank_rows(k, g, b * stage[k], E) for k in range(len(self.banks))) if self._stage_fits(b + 1) else None + for b in range(2) + ] + for g in range(len(self.layouts)) + ] + if len(self.layouts) > 1: + self._stage_slot0_rows = [ + [self._bank_rows(k, self._pool_layout(p), offsets[k][p], 1) for p, base in enumerate(bases) if masks[b][base] for k in range(len(self.banks))] + for b in range(2) + ] + + def stage_views(self, layer_id: int, buffer_id: int) -> tuple[torch.Tensor, ...]: + """Per-bank ``[num_experts, ...]`` views of prefill staging buffer ``buffer_id`` in ``layer_id``'s layout; row == expert id.""" + views = self._stage_views[self.layer_layout[layer_id]][buffer_id] + assert views is not None, f"the slot cache is too small for staging buffer {buffer_id}" + return views + + def pool_state(self, layer_id: int) -> tuple[torch.Tensor, torch.Tensor]: + """``(id_of_slot, usage)`` views of the pool that serves ``layer_id``.""" + p = self.pool_of_layer[layer_id] + return self._pool_id_of_slot[p], self._pool_usage[p] + + def pool_slots(self, layer_id: int) -> int: + return self.pool_sizes[self.pool_of_layer[layer_id]] + + def bank_views(self, n: int | None = None) -> tuple[torch.Tensor, ...]: + """Per-bank views of the whole ``[S]`` slot cache, or of its first ``n`` slots, for a cache + with one pool and one bank layout, where a slot id is its bank row (the on-demand prefill + GEMM over the slot cache and the fork's slot checks).""" + assert self.banks, "set_bank_sources must register the banks first" + assert len(self.pool_sizes) == 1 and not self._raw_banks, "bank_views needs one slot pool and one bank layout" + if n is None: + return tuple(cache for _, cache in self.banks) + return tuple(cache[:n] for _, cache in self.banks) + + def pool_bank_views(self, layer_id: int) -> tuple[torch.Tensor, ...]: + """Per-bank slot views of the pool that serves ``layer_id``, indexed by its local slot ids.""" + assert self.banks, "set_bank_sources must register the banks first" + return self._pool_bank_views[self.pool_of_layer[layer_id]] + def _init_decode_copy_overlap(self) -> None: + # a dedicated stream, not one of torch's 32 pooled ones: a pooled stream can alias the engine + # or capture stream, and the side work would then run serially (correct, but no overlap) + from cuda.bindings import runtime as cudart + + with torch.cuda.device(self.device): + err, handle = cudart.cudaStreamCreateWithFlags(cudart.cudaStreamNonBlocking) + if err != cudart.cudaError_t.cudaSuccess: + raise RuntimeError(f"cudaStreamCreateWithFlags failed: {err}") + self.decode_copy_stream = torch.cuda.ExternalStream(int(handle), device=self.device) + self._decode_copy_events = [ + (torch.cuda.Event(), torch.cuda.Event()) for _ in range(self.num_layers) + ] + # torch creates the CUDA event on its first record: do that now, never inside a capture + for pair in self._decode_copy_events: + for event in pair: + event.record(self.decode_copy_stream) + logger.info("MoE decode copy overlap on: expert misses copy on a side stream") + + def _init_evict_state(self) -> None: + """Scored policies' per-(layer, expert) state (moe/scored_ensure.py); None under LRU. + + It is not slot-shaped: it survives eviction and prefill invalidation, and a rebuild + only resets it. ``evict_tok`` is the decode-step clock the first ensure of each step bumps. + """ + from freetoken.moe import scored_ensure as se + + self.evict_tok = self.evict_last_tok = self.evict_lc = self.evict_ct = None + self.evict_g_table = None + n, dev = self.num_layers * self.num_experts, self.device + if self.cache_policy_id >= 1: + self.evict_tok = torch.zeros((), dtype=torch.int64, device=dev) + self.evict_last_tok = torch.full((n,), se.LAST_TOK_NEVER, dtype=torch.int64, device=dev) + if self.cache_policy_id >= 2: + self.evict_lc = torch.full((n,), se.LC_NEVER, dtype=torch.int32, device=dev) + self.evict_ct = torch.zeros((n,), dtype=torch.int64, device=dev) + self.evict_g_table = se.softplus2_table(dev) + self._alloc_evict_mirrors() + + def _alloc_evict_mirrors(self) -> None: + """Slot-shaped copies of each slot owner's state, read coalesced by the victim scan. + + A mirror is valid while ``evict_slot_owner`` equals ``id_of_slot``; the kernel keeps it + so, and anything else that installs an id into a slot must set the owner to -1. + """ + self.evict_slot_owner = self.evict_slot_last_tok = self.evict_slot_lc = self.evict_slot_ct = None + if self.evict_tok is None: + return + size, dev = self.cache_size, self.device + self.evict_slot_owner = torch.full((size,), -1, dtype=torch.int32, device=dev) + self.evict_slot_last_tok = torch.zeros((size,), dtype=torch.int64, device=dev) + if self.evict_lc is not None: + self.evict_slot_lc = torch.zeros((size,), dtype=torch.int32, device=dev) + self.evict_slot_ct = torch.zeros((size,), dtype=torch.int64, device=dev) + + def reset_evict_state(self) -> None: + if self.evict_tok is None: + return + from freetoken.moe import scored_ensure as se + + self.evict_tok.zero_() + self.evict_last_tok.fill_(se.LAST_TOK_NEVER) + self.evict_slot_owner.fill_(-1) + if self.evict_lc is not None: + self.evict_lc.fill_(se.LC_NEVER) + self.evict_ct.zero_() + + def evict_clock_layer(self) -> int: + """The first layer that ensures on the GPU each decode step; its call bumps ``evict_tok``.""" + if not self.cpu_layer_ids: + return 0 + return min(set(range(self.num_layers)) - self.cpu_layer_ids, default=0) def set_bank_sources( self, sources: dict[str, list[torch.Tensor]], layer_residency: list[str] | None = None, + layer_tags: list | None = None, ) -> None: """Attach the host (CPU pinned) expert source banks and allocate a GPU slot cache per bank, following the format's bank schema. @@ -315,6 +660,8 @@ def set_bank_sources( ``layer_residency`` labels each layer with a ``HostResidency`` value (default: all pinned). Non-pinned (LOCKED/PAGEABLE) layers have no device address: they must already be routed to the CPU executor (``cpu_layer_ids``, set BEFORE this call), the copy plan skips their rows, and their only movement is ``copy_missing``'s whole-layer pageable prefill branch -- which is why prefill overlap is incompatible with them. + + ``layer_tags`` names each layer's quant (see :func:`bank_layouts`); layers of different tags never share a pool. """ from freetoken.moe.legacy_format import canonical_role from freetoken.moe.host_banks import HostResidency @@ -348,43 +695,66 @@ def set_bank_sources( for name in self.bank_schema: per_layer = sources[name] assert len(per_layer) == self.num_layers, (name, len(per_layer)) - head = per_layer[0] - if self.layout is not None: - spec = self.layout[name] - if tuple(head.shape[1:]) != tuple(spec.shape) or head.dtype != spec.dtype: - raise ValueError( - f"bank {name!r} rows are {tuple(head.shape[1:])} {head.dtype} but the expert kernel's layout " - f"wants {tuple(spec.shape)} {spec.dtype}; the banks were packed for another kernel" - ) for layer_id, source in enumerate(per_layer): assert source.is_contiguous(), f"bank {name!r} layer {layer_id} must be contiguous" assert source.size(0) == self.num_experts, (name, layer_id, source.shape) - assert source.shape == head.shape and source.dtype == head.dtype, ( - name, layer_id, source.shape, source.dtype, - ) + if self.layout is not None: + spec = self.layout[name] + if tuple(source.shape[1:]) != tuple(spec.shape) or source.dtype != spec.dtype: + raise ValueError( + f"bank {name!r} rows are {tuple(source.shape[1:])} {source.dtype} but the expert kernel's layout " + f"wants {tuple(spec.shape)} {spec.dtype}; the banks were packed for another kernel" + ) self.bank_sources[name] = list(per_layer) - self.bank_caches[name] = torch.zeros( - (self.cache_size, *head.shape[1:]), - dtype=head.dtype, - device=self.device, - ) - self.banks = [(self.bank_sources[n], self.bank_caches[n]) for n in self.bank_schema] + keys = bank_layouts(sources, layer_tags) + unique = list(dict.fromkeys(keys)) + self.layouts = [banks for _, banks in unique] + self.layer_layout = [unique.index(key) for key in keys] + for p in range(len(self.pool_sizes)): + mixed = {self.layer_layout[l] for l, q in enumerate(self.pool_of_layer) if q == p} + assert len(mixed) == 1, f"slot pool {p} holds layers of {len(mixed)} layouts; a pool needs one" + self._raw_banks = len(set(self.layouts)) > 1 + self._row_bytes = [[math.prod(shape) * dtype.itemsize for dtype, shape in layout] for layout in self.layouts] + if self._raw_banks: + # every pool view and staging buffer starts on a row boundary of its own layout + assert all(r % 16 == 0 for rows in self._row_bytes for r in rows), self._row_bytes + self._stage_bytes = [self.num_experts * max(rows[k] for rows in self._row_bytes) for k in range(len(self.bank_schema))] + self._alloc_banks() + if not self._stage_fits(1): + raise ValueError(f"MoE cache of {self.cache_size} slots cannot stage one whole layer of experts") + if self.prefill_overlap and not self._stage_fits(2): + logger.warning("Disabling MoE prefill overlap: the slot cache cannot hold two whole layers of experts") + self.prefill_overlap = False + self._slice_pools() self._build_copy_plan() + if self.verify is not None: + self.verify.bind(self) + if self.audit is not None: + self.audit.bind(self) if self.prefill_overlap: self._init_prefill_overlap_buffers() + def _alloc_banks(self) -> None: + for k, name in enumerate(self.bank_schema): + if self._raw_banks: + nbytes = sum(s * self._row_bytes[self._pool_layout(p)][k] for p, s in enumerate(self.pool_sizes)) + self.bank_caches[name] = torch.zeros(nbytes, dtype=torch.uint8, device=self.device) + else: + dtype, shape = self.layouts[0][k] + self.bank_caches[name] = torch.zeros((self.cache_size, *shape), dtype=dtype, device=self.device) + self.banks = [(self.bank_sources[n], self.bank_caches[n]) for n in self.bank_schema] + def _build_copy_plan(self) -> None: self._build_fused_copy_plan() if self._copy_fused_ok or self.device.type != "cuda" or not self.banks: return - for name in self.bank_schema: - cache = self.bank_caches[name] - feat = math.prod(cache.shape[1:]) * cache.element_size() - if feat % 128: - raise RuntimeError( - f"MoE bank {name!r} rows are {feat} bytes (not a multiple of 128): " - f"only the fused multi-bank copy can move them, but it is disabled" - ) + for k, name in enumerate(self.bank_schema): + for rows in self._row_bytes: + if rows[k] % 128: + raise RuntimeError( + f"MoE bank {name!r} rows are {rows[k]} bytes (not a multiple of 128): " + f"only the fused multi-bank copy can move them, but it is disabled" + ) def _build_fused_copy_plan(self) -> None: """Precompute the fused multi-bank copy descriptor (base addrs + per-row bytes). @@ -396,24 +766,25 @@ def _build_fused_copy_plan(self) -> None: """ self._copy_fused_ok = False self._copy_dst_ptrs = None + self._copy_dst_ptrs_pool = None self._copy_src_ptrs = None self._copy_feat_bytes = None self._copy_dst_ptrs_host: list[int] = [] self._copy_src_ptrs_host: list[list[int]] = [] - self._copy_feat_bytes_host: list[int] = [] - self._gather_bank_ids: list[int] = [] - self._gather_dst_ptrs: torch.Tensor | None = None - self._gather_feat_bytes: torch.Tensor | None = None + self._gather_bank_ids: list[list[int]] = [] + self._gather_feat_bytes: list[torch.Tensor | None] = [] + self._gather_stage_ptrs: list[list[torch.Tensor] | None] = [] + self._gather_pool_ptrs: list[torch.Tensor | None] = [] if not _FUSED_COPY or self.device.type != "cuda" or not self.banks: return from freetoken.kernel.pinned import device_ptr - dst_ptrs, feats = [], [] + aligned = [self._stage_bytes] + self._row_bytes + self._pool_offsets + if any(n % 16 for ns in aligned for n in ns) or any(c.data_ptr() % 16 for _, c in self.banks): + return # leave fused disabled; copy_missing uses the per-bank path + dst_ptrs = [cache.data_ptr() for _, cache in self.banks] layer_src_ptrs = [[] for _ in range(self.num_layers)] - for per_layer, cache in self.banks: - feat = math.prod(per_layer[0].shape[1:]) * per_layer[0].element_size() - if feat % 16 != 0 or cache.data_ptr() % 16 != 0: - return # leave fused disabled; copy_missing uses the per-bank path + for per_layer, _ in self.banks: for layer_id, source in enumerate(per_layer): if layer_id in self._unpinned_layers: # unregistered layer: no device alias exists, and the row is never consumed (CPU decode; pageable prefill) @@ -427,52 +798,65 @@ def _build_fused_copy_plan(self) -> None: if src_dev % 16 != 0: return layer_src_ptrs[layer_id].append(src_dev) - dst_ptrs.append(cache.data_ptr()) - feats.append(feat) - self._copy_dst_ptrs = torch.tensor(dst_ptrs, dtype=torch.int64, device=self.device) - self._copy_src_ptrs = [ - torch.tensor(ptrs, dtype=torch.int64, device=self.device) - for ptrs in layer_src_ptrs + + def ptrs(values: list[int]) -> torch.Tensor: + return torch.tensor(values, dtype=torch.int64, device=self.device) + + self._copy_dst_ptrs = ptrs(dst_ptrs) + self._copy_dst_ptrs_pool = [ + self._copy_dst_ptrs if all(o[p] == 0 for o in self._pool_offsets) + else ptrs([d + o[p] for d, o in zip(dst_ptrs, self._pool_offsets)]) + for p in range(len(self.pool_sizes)) ] - self._copy_feat_bytes = torch.tensor(feats, dtype=torch.int64, device=self.device) + self._copy_src_ptrs = [ptrs(values) for values in layer_src_ptrs] + self._copy_feat_bytes = [ptrs(rows) for rows in self._row_bytes] self._copy_dst_ptrs_host = dst_ptrs self._copy_src_ptrs_host = layer_src_ptrs - self._copy_feat_bytes_host = feats # hit-D2D gather serves only the big banks; small banks are whole-layer # H2D entries (see _SMALL_BANK_FEAT_BYTES), so their rows never need D2D. - self._gather_bank_ids = [i for i, f in enumerate(feats) if f >= _SMALL_BANK_FEAT_BYTES] - if len(self._gather_bank_ids) == len(feats): - self._gather_dst_ptrs = self._copy_dst_ptrs - self._gather_feat_bytes = self._copy_feat_bytes - elif self._gather_bank_ids: - self._gather_dst_ptrs = self._copy_dst_ptrs[self._gather_bank_ids].contiguous() - self._gather_feat_bytes = self._copy_feat_bytes[self._gather_bank_ids].contiguous() + # The gather copies cache rows of the layer's pool into its staging buffer, row == expert id. + self._gather_bank_ids = [[k for k, r in enumerate(rows) if r >= _SMALL_BANK_FEAT_BYTES] for rows in self._row_bytes] + for g, ids in enumerate(self._gather_bank_ids): + self._gather_feat_bytes.append(ptrs([self._row_bytes[g][k] for k in ids]) if ids else None) + self._gather_stage_ptrs.append( + [ptrs([dst_ptrs[k] + b * self._stage_bytes[k] for k in ids]) for b in range(2)] if ids else None + ) + for p in range(len(self.pool_sizes)): + ids = self._gather_bank_ids[self._pool_layout(p)] + self._gather_pool_ptrs.append(ptrs([dst_ptrs[k] + self._pool_offsets[k][p] for k in ids]) if ids else None) self._copy_blocks_per_bank = default_blocks_per_bank(self.device) self._copy_fused_ok = True - def validate_rebuild(self, cache_size: int) -> None: + def validate_rebuild(self, cache_size: int, pool_sizes: list[int] | None = None) -> None: """Pure geometry validation of a rebuild target (no GPU side effects). Raises ``ValueError`` if ``cache_size`` is below the ``num_experts`` floor or - above the marlin slot cap. Called by :meth:`rebuild` and by the engine's + a pool is above the expert kernel's slot cap, or the bank byte segments cannot stage one layer. Called by :meth:`rebuild` and by the engine's pre-teardown check, so an invalid target rejects with the old cache intact (no destructive free first). """ if cache_size < self.num_experts: raise ValueError(f"cache_size {cache_size} < num_experts {self.num_experts}") - if self.max_slots is not None and cache_size > self.max_slots: + # decode kernels see one pool's slots, so the cap applies per pool + largest = cache_size if pool_sizes is None else max(pool_sizes) + what = f"moe_cache_size={cache_size}" if largest == cache_size else f"a MoE slot pool of {largest} slots" + if self.max_slots is not None and largest > self.max_slots: raise ValueError( - f"moe_cache_size={cache_size} exceeds the expert kernel's slot limit of {self.max_slots}; " + f"{what} exceeds the expert kernel's slot limit of {self.max_slots}; " f"pass --moe-cache-size {self.max_slots} or less, or let the default kernel serve the experts" ) - if self.layout is None and self.quant_format == "nvfp4_marlin" and cache_size > MARLIN_MAX_CACHE_SIZE: + if self.layout is None and self.quant_format == "nvfp4_marlin" and largest > MARLIN_MAX_CACHE_SIZE: raise ValueError( - f"moe_cache_size={cache_size} exceeds the marlin backend's slot limit of " + f"{what} exceeds the marlin backend's slot limit of " f"{MARLIN_MAX_CACHE_SIZE} (vLLM moe_align_block_size caps padded experts at " "1024); reduce moe_cache_size or force --quant-backend moe.nvfp4=triton" ) + if self._raw_banks and pool_sizes is not None: + for k, stage in enumerate(self._stage_bytes): + if sum(s * self._row_bytes[self._pool_layout(p)][k] for p, s in enumerate(pool_sizes)) < stage: + raise ValueError(f"moe_cache_size={cache_size} cannot stage one whole layer of experts") - def rebuild(self, cache_size: int) -> None: + def rebuild(self, cache_size: int, pool_sizes: list[int] | None = None) -> None: """Resize the GPU slot cache + bookkeeping to ``cache_size`` IN PLACE. Keeps the CPU/pinned ``bank_sources`` and the GPU-resident alphas; never @@ -482,9 +866,10 @@ def rebuild(self, cache_size: int) -> None: ``ctx.moe_offload_cache`` stay valid. """ assert self.bank_sources, "set_bank_sources must run before rebuild" - self.validate_rebuild(cache_size) + sizes = [cache_size] if pool_sizes is None else list(pool_sizes) + assert len(sizes) == len(self.pool_sizes) and sum(sizes) == cache_size and min(sizes) > 0, (sizes, cache_size) + self.validate_rebuild(cache_size, sizes) # 1. Tear down prefill-overlap (its buffer views alias the old bank_caches). - self.prefill_bank_buffers = [] self.prefill_copy_stream = None self.prefill_begin_event = None self.prefill_ready_events = [] @@ -495,22 +880,24 @@ def rebuild(self, cache_size: int) -> None: # 2. Drop old GPU tensors (free-before-alloc). self.banks = [] self.bank_caches = {} + self._pool_bank_views = [] + self._stage_views = [] + self._stage_slot0_rows = [[], []] self.cache_size = cache_size + self.pool_sizes = sizes if self.device.type == "cuda": torch.cuda.synchronize(self.device) torch.cuda.empty_cache() # 3. Reallocate the slot cache from the retained host sources. - for name in self.bank_schema: - head = self.bank_sources[name][0] - self.bank_caches[name] = torch.zeros( - (cache_size, *head.shape[1:]), dtype=head.dtype, device=self.device - ) - self.banks = [(self.bank_sources[n], self.bank_caches[n]) for n in self.bank_schema] - self._build_copy_plan() # slot caches were reallocated -> refresh fused-copy addrs + self._alloc_banks() # 4. Reallocate cache_size-shaped bookkeeping; reset the slot map (cold start). self.slot_for_id.fill_(-1) self.id_of_slot = torch.full((cache_size,), -1, dtype=torch.int32, device=self.device) self.usage = torch.zeros((cache_size,), dtype=torch.int64, device=self.device) + self._slice_pools() + self._build_copy_plan() # slot caches were reallocated -> refresh fused-copy addrs + if self.verify is not None: + self.verify.bind(self) plan_slots = max(self.num_experts, cache_size) self.evict_slots = torch.empty((plan_slots,), dtype=torch.int32, device=self.device) self.src_indices = torch.empty((plan_slots,), dtype=torch.int32, device=self.device) @@ -519,6 +906,8 @@ def rebuild(self, cache_size: int) -> None: self.num_indices.zero_() self.num_missing_full.zero_() self.expert_recency.fill_(-1) + self._alloc_evict_mirrors() + self.reset_evict_state() self.stat_missing.zero_() self.stat_active.zero_() self.stat_calls.zero_() @@ -530,14 +919,18 @@ def rebuild(self, cache_size: int) -> None: self.stat_fetched_layer.zero_() self.stat_steps_layer.zero_() self.decode_freq.zero_() + if self.prefetch is not None: + self.prefetch.reset() + if self.audit is not None: + self.audit.bind(self) # new slot arrays: a new history self.prefill_hit_rows = 0 self.prefill_total_rows = 0 self._hit_d2d_fallback_logged = False # geometry changed; re-log if still unusable # 5. Re-evaluate prefill overlap against the new size. - if self.prefill_overlap and cache_size < 2 * self.num_experts: + if self.prefill_overlap and not self._stage_fits(2): logger.warning( f"Disabling MoE prefill overlap on rebuild: cache_size {cache_size} " - f"< 2*num_experts {2 * self.num_experts}." + f"cannot hold two whole layers of experts (2*num_experts {2 * self.num_experts} slots of one layout)." ) self.prefill_overlap = False if self.prefill_overlap: @@ -590,8 +983,9 @@ def alphas_for_slots(self, layer_id: int) -> tuple[torch.Tensor, torch.Tensor] | ``layer_id`` -- are ever read by the grouped GEMM.""" if self.gate_up_alpha is None: return None + id_of_slot, _ = self.pool_state(layer_id) idx = layer_id * self.num_experts + ( - self.id_of_slot.clamp(min=0).long() % self.num_experts + id_of_slot.clamp(min=0).long() % self.num_experts ) return self.gate_up_alpha[idx], self.down_alpha[idx] @@ -605,25 +999,12 @@ def alphas_for_layer(self, layer_id: int) -> tuple[torch.Tensor, torch.Tensor] | hi = lo + self.num_experts return self.gate_up_alpha[lo:hi], self.down_alpha[lo:hi] - def bank_views(self, n: int | None = None) -> tuple[torch.Tensor, ...]: - """Per-bank cache views in registration order: the full ``[S]`` slot cache - (decode), or its first ``n`` slots (materialized layer).""" - assert self.banks, "set_bank_sources must register the banks first" - if n is None: - return tuple(cache for _, cache in self.banks) - return tuple(cache[:n] for _, cache in self.banks) - def _init_prefill_overlap_buffers(self) -> None: assert self.banks, "set_bank_sources must register the banks first" self._prefill_buffer_layer = [None, None] self._prefill_buffer_released = [True, True] self._prefill_buffer_has_release_event = [False, False] - # The double buffers borrow the slot cache's first 2 * num_experts slots - # (one full expert layer per buffer), one view per registered bank. - self.prefill_bank_buffers = [ - cache[: 2 * self.num_experts].view(2, self.num_experts, *cache.shape[1:]) - for _, cache in self.banks - ] + # the double buffers are the staging views (_slice_pools), on top of the first pools' slots if self.device.type == "cuda": self.prefill_copy_stream = torch.cuda.Stream(device=self.device) self.prefill_ready_events = [torch.cuda.Event() for _ in range(2)] @@ -642,15 +1023,32 @@ def _init_prefill_overlap_buffers(self) -> None: ) self._prefill_hit_num = torch.zeros((1,), dtype=torch.int64, device=self.device) - def _invalidate_prefill_buffer(self, buffer_id: int) -> None: - slot_start = buffer_id * self.num_experts - slot_end = slot_start + self.num_experts - old_ids = self.id_of_slot[slot_start:slot_end] - self.slot_for_id.view(-1)[old_ids[old_ids >= 0].long()] = -1 - old_ids.fill_(-1) - # usage=0 makes these slots the oldest, so the argmin(usage) victim selection in - # ensure_experts evicts them first. - self.usage[slot_start:slot_end].zero_() + def _invalidate_prefill_buffer(self, buffer_id: int, layer_id: int = -1) -> None: + slots = self._inval[buffer_id] + if isinstance(slots, slice): + if self.audit is not None: + from freetoken.moe.slot_audit import INVALIDATE + + self.audit.record_clear(self, INVALIDATE, layer_id, slots.start, slots.stop - slots.start) + # One fixed-shape launch. The boolean-mask index + # (slot_for_id[old_ids[old_ids >= 0]] = -1) has a data-dependent shape, so every + # call hid a device-to-host sync -- and with two buffer reuses per chunk x 48 + # layers the host waited, per layer, for ALL enqueued GPU work (the cached-context + # attention included), serializing the prefill-overlap pipeline. + from freetoken.kernel.triton.moe import invalidate_prefill_slots + + if slots.stop > slots.start: + invalidate_prefill_slots( + self.id_of_slot, self.slot_for_id, self.usage, slots.start, slots.stop - slots.start + ) + else: + old_ids = self.id_of_slot[slots] + self.slot_for_id.view(-1)[old_ids[old_ids >= 0].long()] = -1 + self.id_of_slot[slots] = -1 + # usage=0 makes these slots the oldest, so the argmin(usage) victim selection in + # ensure_experts evicts them first. + self.usage[slots] = 0 + self._stage_dirty[buffer_id] = True def begin_prefill(self) -> None: if not self.prefill_overlap: @@ -681,7 +1079,7 @@ def prefetch_prefill_layer(self, layer_id: int) -> None: if layer_id < 0: raise ValueError(f"Invalid prefill layer id: {layer_id}") - assert self.banks and self.prefill_bank_buffers + assert self.banks and self._stage_views buffer_id = layer_id % 2 if self._prefill_buffer_layer[buffer_id] == layer_id: @@ -692,9 +1090,13 @@ def prefetch_prefill_layer(self, layer_id: int) -> None: ) def copy() -> None: - self._invalidate_prefill_buffer(buffer_id) - for (per_layer, _), buffer in zip(self.banks, self.prefill_bank_buffers): - buffer[buffer_id].copy_(per_layer[layer_id], non_blocking=True) + self._invalidate_prefill_buffer(buffer_id, layer_id) + for (per_layer, _), view in zip(self.banks, self.stage_views(layer_id, buffer_id)): + view.copy_(per_layer[layer_id], non_blocking=True) + if self.audit is not None: + from freetoken.moe.slot_audit import PREFILL_BUFFER + + self.audit.record_range(self, PREFILL_BUFFER, layer_id, buffer_id * self.num_experts, self.num_experts, meta=False) if self._prefill_hit_d2d_active: self._prefetch_split(layer_id, buffer_id) @@ -725,10 +1127,10 @@ def _hit_d2d_usable(self) -> bool: reason = "FREETOKEN_SKIP_FAST_INDEX_COPY is set (the hit gather would be a no-op)" elif not self._copy_fused_ok: reason = "the fused copy plan is unavailable (bank alignment or FREETOKEN_FUSED_COPY=0)" - elif self.cache_size <= 2 * self.num_experts: + elif all(stable >= s for stable, s in zip(self._pool_stable_from, self.pool_sizes)): reason = ( f"cache_size {self.cache_size} leaves no hit region " - f"(needs > {2 * self.num_experts} slots)" + f"(every slot lies under the prefill staging buffers of {2 * self.num_experts} experts)" ) elif not self._resolve_batch_memcpy(): reason = "cudaMemcpyBatchAsync is unavailable" # resolve logged the specifics @@ -766,7 +1168,7 @@ def _prefetch_split(self, layer_id: int, buffer_id: int) -> None: run list comes from the begin-of-chunk snapshot because the batch API takes HOST pointer arrays. Live-vs-snapshot cannot disagree: the only chunk-internal writer (buffer invalidation) rewrites slots already below - the 2E threshold, and slots < 2E (including -1) are misses on both sides + the pool's staging threshold, and those slots (and -1) are misses on both sides -- the buffers own those slots, so their bytes are volatile within the chunk. Hit and miss row sets are disjoint, so the streams need no ordering against each other. @@ -777,43 +1179,51 @@ def _prefetch_split(self, layer_id: int, buffer_id: int) -> None: from freetoken.moe.offload_kernels import prefill_hit_compact E = self.num_experts + layout = self.layer_layout[layer_id] + pool = self.pool_of_layer[layer_id] snap = self._prefill_snapshot_np[layer_id] - hit_mask = snap >= 2 * E + hit_mask = snap >= self._pool_stable_from[pool] self.prefill_hit_rows += int(hit_mask.sum()) self.prefill_total_rows += E - if self._gather_dst_ptrs is not None: - prefill_hit_compact(self, layer_id, buffer_id) + if self._gather_feat_bytes[layout] is not None: + prefill_hit_compact(self, layer_id) # blocks_per_bank=64 vs the PCIe-tuned default of 8: HBM D2D needs the # wider grid (~22 GB/s per 1024-thread block on H100). fast_index_copy_multi_jit( - self._gather_dst_ptrs, - self._gather_dst_ptrs, - self._gather_feat_bytes, + self._gather_stage_ptrs[layout][buffer_id], + self._gather_pool_ptrs[pool], + self._gather_feat_bytes[layout], self._prefill_hit_dst, self._prefill_hit_src, self._prefill_hit_num, blocks_per_bank=64, ) + if self.audit is not None: + from freetoken.moe.slot_audit import PREFILL_HIT_D2D + + self.audit.record_plan(self, PREFILL_HIT_D2D, layer_id, self._prefill_hit_dst, None, self._prefill_hit_num, + meta=False, aux=self._prefill_hit_src, row_base=buffer_id * E) miss = np.nonzero(~hit_mask)[0] with torch.cuda.stream(self.prefill_copy_stream): if self._prefill_buffer_has_release_event[buffer_id]: self.prefill_copy_stream.wait_event(self.prefill_release_events[buffer_id]) - self._invalidate_prefill_buffer(buffer_id) + self._invalidate_prefill_buffer(buffer_id, layer_id) if miss.size: run_starts = np.concatenate(([0], np.nonzero(np.diff(miss) != 1)[0] + 1)) starts = miss[run_starts] lengths = np.diff(np.concatenate((run_starts, [miss.size]))) dst, src, nbytes = [], [], [] - for b, feat in enumerate(self._copy_feat_bytes_host): + for b, feat in enumerate(self._row_bytes[layout]): + stage = self._copy_dst_ptrs_host[b] + buffer_id * self._stage_bytes[b] if feat < _SMALL_BANK_FEAT_BYTES: # Whole layer as one entry, EVEN with zero misses: it keeps every # batch entry above the driver's async floor and covers the hit # rows the gather skips for these banks. - dst.append(self._copy_dst_ptrs_host[b] + buffer_id * E * feat) + dst.append(stage) src.append(self._copy_src_ptrs_host[layer_id][b]) nbytes.append(E * feat) elif miss.size: - dst.extend(self._copy_dst_ptrs_host[b] + (buffer_id * E + starts) * feat) + dst.extend(stage + starts * feat) src.extend(self._copy_src_ptrs_host[layer_id][b] + starts * feat) nbytes.extend(lengths * feat) if dst: @@ -823,6 +1233,11 @@ def _prefetch_split(self, layer_id: int, buffer_id: int) -> None: torch.tensor(nbytes, dtype=torch.int64), torch.cuda.current_stream(self.device).cuda_stream, ) + if self.audit is not None: + from freetoken.moe.slot_audit import PREFILL_SPLIT_H2D + + # every row gets at least its small banks from the host, the misses all of them + self.audit.record_range(self, PREFILL_SPLIT_H2D, layer_id, buffer_id * E, E, meta=False) self.prefill_ready_events[buffer_id].record(self.prefill_copy_stream) def wait_prefill_layer(self, layer_id: int) -> tuple[torch.Tensor, ...]: @@ -831,15 +1246,22 @@ def wait_prefill_layer(self, layer_id: int) -> tuple[torch.Tensor, ...]: marlin/b12x ``(gate_up_packed, gate_up_scale, down_packed, down_scale)``; nvfp4 native adds the two global banks after each scale bank.""" assert self.prefill_overlap - assert self.prefill_bank_buffers + assert self._stage_views self.prefetch_prefill_layer(layer_id) buffer_id = layer_id % 2 assert self._prefill_buffer_layer[buffer_id] == layer_id if self.prefill_ready_events: torch.cuda.current_stream(self.device).wait_event(self.prefill_ready_events[buffer_id]) - return tuple(buffer[buffer_id] for buffer in self.prefill_bank_buffers) + return self.stage_views(layer_id, buffer_id) def release_prefill_layer(self, layer_id: int) -> None: + if self.zero_inactive_route_slots and layer_id == self.num_layers - 1: + # the staged slots were invalidated and nothing registers them within the forward, so their bytes are free to clear + for b in range(2): + if self._stage_dirty[b]: + for row in self._stage_slot0_rows[b]: + row.zero_() + self._stage_dirty[b] = False if not self.prefill_overlap: return buffer_id = layer_id % 2 @@ -851,21 +1273,79 @@ def release_prefill_layer(self, layer_id: int) -> None: self._prefill_buffer_released[buffer_id] = True def loads_prefill_on_demand(self, num_tokens: int) -> bool: - return num_tokens <= self.prefill_demand_tokens + return 0 < num_tokens <= self.prefill_demand_tokens - def ensure_experts(self, layer_id: int, expert_ids: torch.Tensor, *, counts_as_decode: bool = True) -> None: - from freetoken.moe.offload_kernels import ensure_experts - - if self.collect_decode_freq and counts_as_decode: + def ensure_experts( + self, + layer_id: int, + expert_ids: torch.Tensor, + *, + update_state: bool = True, + router_logits: torch.Tensor | None = None, + pin_since: torch.Tensor | None = None, + counts_as_decode: bool = True, + ) -> None: + """``update_state=False`` (small-prefill ensures) keeps a scored policy's state and + decode-step clock still; ``router_logits`` (``[rows, num_experts]``) feed the ``rule`` + policy's near misses; ``pin_since`` (a ``step`` value) keeps every slot touched after it + (a small prefill's earlier chunks). All are ignored under LRU, whose recency order + already keeps those slots.""" + from freetoken.moe.offload_kernels import ensure_experts, ensure_experts_scored + + if self.collect_decode_freq and counts_as_decode and update_state: # ``expert_ids`` still holds raw expert ids here (the kernel rewrites them to # slot ids in place), so snapshot the routing histogram before that happens. ids = expert_ids.reshape(-1).long() self.decode_freq[layer_id].scatter_add_(0, ids, torch.ones_like(ids)) self._pending_src_layer = layer_id self._pending_whole_layer = False - ensure_experts(self, layer_id, expert_ids, record_stats=counts_as_decode) + audit = self.audit + if audit is not None and update_state and layer_id == self.evict_clock_layer(): + audit.begin_step() + lowpri = self.prefetch_on + width = expert_ids.numel() + if self.cache_policy_id == 0 and not lowpri: + ensure_experts(self, layer_id, expert_ids, record_stats=counts_as_decode and update_state) + else: + # prefetch on: LRU runs the vendored kernel (flashlib's cannot key low-priority slots) + ensure_experts_scored( + self, + layer_id, + expert_ids, + bump_tok=update_state and layer_id == self.evict_clock_layer(), + update_state=update_state, + router_logits=router_logits, + pin_since=pin_since, + lowpri=lowpri, + pf_count=self.prefetch.count_args(layer_id) if lowpri and update_state else None, + ) + if audit is not None: + from freetoken.moe import slot_audit as sa + + kind = sa.DEMAND_INSTALL if update_state else sa.PREFILL_INSTALL + self._pending_kind = sa.DEMAND_COPY if update_state else sa.PREFILL_COPY + self._pending_width = width + audit.record_plan(self, kind, layer_id, self.evict_slots, self.src_indices, self.num_indices, meta=True, width=width) + + def prefetch_ensure( + self, layer_id: int, query: torch.Tensor, dst_slots: torch.Tensor, src_rows: torch.Tensor, + num: torch.Tensor, stats_row: torch.Tensor, ready: torch.Tensor, budget: int | None = None, + ) -> None: + """Install the first ``budget`` of ``query``'s non-resident experts of ``layer_id`` as + low-priority slots and plan their copy into ``dst_slots``/``src_rows``/``num`` (never the + demand plan, which copy(L-1) still reads). Runs on the compute stream like every other slot + map writer.""" + from freetoken.moe.offload_kernels import prefetch_ensure_experts + + prefetch_ensure_experts(self, layer_id, query, dst_slots, src_rows, num, stats_row, ready, budget=budget) + if self.audit is not None: + from freetoken.moe.slot_audit import PREFETCH_INSTALL - def ensure_experts_hybrid(self, layer_id: int, expert_ids: torch.Tensor) -> None: + self.audit.record_plan(self, PREFETCH_INSTALL, layer_id, dst_slots, src_rows, num, meta=True) + + def ensure_experts_hybrid( + self, layer_id: int, expert_ids: torch.Tensor, weights: torch.Tensor, clamp: bool = False + ) -> HybridRoutes: """Capped-fetch LRU for the hybrid backend. Like :meth:`ensure_experts` but assigns slots to (and schedules copies for) at @@ -874,7 +1354,8 @@ def ensure_experts_hybrid(self, layer_id: int, expert_ids: torch.Tensor) -> None left non-resident and ``expert_ids`` is rewritten to their cache slot (hit or freshly fetched) or ``-1`` (overflow -> compute on the CPU). ``num_indices`` holds the capped fetch count (for ``copy_missing``); ``num_missing_full`` the pre-cap - miss count (for stats). All device-side / fixed-shape, so it is CUDA-graph safe.""" + miss count (for stats). All device-side / fixed-shape, so it is CUDA-graph safe. + Returns the routes split between GPU and CPU (:class:`freetoken.moe.offload_kernels.HybridRoutes`).""" from freetoken.moe.offload_kernels import ensure_experts_hybrid if self.collect_decode_freq: @@ -882,8 +1363,8 @@ def ensure_experts_hybrid(self, layer_id: int, expert_ids: torch.Tensor) -> None self.decode_freq[layer_id].scatter_add_(0, ids, torch.ones_like(ids)) self._pending_src_layer = layer_id self._pending_whole_layer = False - ensure_experts_hybrid( - self, layer_id, expert_ids, self.hybrid_max_fetch, self.hybrid_fetch_fraction + return ensure_experts_hybrid( + self, layer_id, expert_ids, weights, self.hybrid_max_fetch, self.hybrid_fetch_fraction, clamp ) def materialize_layer(self, layer_id: int) -> None: @@ -891,15 +1372,36 @@ def materialize_layer(self, layer_id: int) -> None: self._pending_src_layer = layer_id self._pending_whole_layer = True + audit = self.audit + if audit is not None: + from freetoken.moe import slot_audit as sa + + # the kernel also frees this layer's slots above E; record them while they still name it + audit.record_clear(self, sa.MATERIALIZE_CLEAR, layer_id, 0, self.cache_size, layer_above_e=True) materialize_layer(self, layer_id) + if self.evict_slot_owner is not None: + self.evict_slot_owner[: self.num_experts].fill_(-1) # it installs the layer into these slots + if audit is not None: + self._pending_kind = sa.MATERIALIZE_COPY + self._pending_width = self.num_experts + audit.record_plan(self, sa.MATERIALIZE_INSTALL, layer_id, self.evict_slots, self.src_indices, self.num_indices, + meta=True, width=self.num_experts) def reset(self) -> None: from freetoken.moe.offload_kernels import reset_cache + if self.audit is not None: + from freetoken.moe.slot_audit import RESET + + self.audit.record_clear(self, RESET, -1, 0, self.cache_size) reset_cache(self) # Per-expert recency is not cache_size-shaped, so reset_cache leaves it alone; wipe # it here so a new sequence starts with cold hybrid fetch priorities. self.expert_recency.fill_(-1) + self.reset_evict_state() + if self.prefetch is not None: + # the graph runner resets after each capture: drop the capture-time counts too + self.prefetch.reset_counters() def reset_stats(self) -> None: self.prefill_hit_rows = 0 @@ -921,23 +1423,6 @@ def record_decode_stats(self, layer_id: int) -> None: was eight torch ops per layer per step, all captured into the decode graph. """ - def record_decode_stats_hybrid(self, layer_id: int) -> None: - """Hybrid stats: full miss count (pre-cap), the PCIe-fetched count (capped), and - the active count. The CPU computes (missing - fetched) experts. Device-side; - accumulates both the scalar totals and the per-layer breakdown.""" - assert 0 <= layer_id < self.num_layers, f"layer_id {layer_id} out of range [0, {self.num_layers})" - missing = self.num_missing_full.sum() - fetched = self.num_indices.sum() - active = self.active_mask.sum() - self.stat_missing += missing - self.stat_fetched += fetched - self.stat_active += active - self.stat_calls += 1 - self.stat_missing_layer[layer_id] += missing - self.stat_fetched_layer[layer_id] += fetched - self.stat_active_layer[layer_id] += active - self.stat_steps_layer[layer_id] += 1 - def decode_miss_stats(self) -> dict: if self.decode_target == "hybrid": active = int(self.stat_active.item()) @@ -981,6 +1466,8 @@ def decode_miss_stats_per_layer(self) -> dict: s, m, a, f = steps[L], missing[L], active[L], fetched[L] per_layer.append({ "layer": L, + "pool": self.pool_of_layer[L], + "pool_slots": self.pool_slots(L), "steps": s, "active_per_step": (a / s) if s else 0.0, "missing_per_step": (m / s) if s else 0.0, @@ -992,10 +1479,10 @@ def decode_miss_stats_per_layer(self) -> dict: def decode_routing_stats(self) -> dict: """Per-layer decode routing concentration, for cache-skew analysis. - Uses the histogram from ``collect_decode_freq``. The ``oracle_hit`` is the best a - per-layer LRU holding ``cache_size/num_layers`` slots could achieve on the observed - (stationary) routing distribution -- i.e. an upper bound on hit rate that depends - purely on how skewed routing is, independent of any LRU/LFU dynamics. + Uses the histogram from ``collect_decode_freq``. ``static_topk_hit_at_slots`` is the + hit rate of the best *fixed* expert set: each layer's ``cache_size/num_layers`` (rounded) + most frequent experts over the whole histogram. It measures routing skew only. It is + not an upper bound on a dynamic cache: LRU exploits temporal locality and can beat it. """ freq = self.decode_freq.float() total = freq.sum(dim=1) @@ -1005,7 +1492,7 @@ def decode_routing_stats(self) -> dict: slots_per_layer = self.cache_size / self.num_layers C = max(1, int(round(slots_per_layer))) sorted_f, _ = torch.sort(freq, dim=1, descending=True) - oracle_hit = (sorted_f[:, :C].sum(dim=1)[valid] / total[valid]).mean().item() + static_hit = (sorted_f[:, :C].sum(dim=1)[valid] / total[valid]).mean().item() ws = (freq > 0).sum(dim=1).float() cdf = torch.cumsum(sorted_f, dim=1) / total.clamp(min=1).unsqueeze(1) cover90 = ((cdf < 0.9).sum(dim=1).float() + 1)[valid] @@ -1017,13 +1504,16 @@ def decode_routing_stats(self) -> dict: "working_set_mean": ws[valid].mean().item(), "working_set_max": int(ws[valid].max().item()), "experts_for_90pct": cover90.mean().item(), - "oracle_hit_at_slots": oracle_hit, + "static_topk_hit_at_slots": static_hit, "norm_entropy": norm_ent, } - def copy_missing(self) -> None: + def copy_missing(self, layer_id: int | None = None) -> None: + """Copy the misses the last ensure_experts/materialize_layer staged, from ``layer_id``'s + host banks (default: the layer that staged them).""" assert self.banks, "set_bank_sources must register the banks first" - layer_id = self._pending_src_layer + if layer_id is None: + layer_id = self._pending_src_layer assert layer_id is not None, "no staged misses (ensure_experts/materialize_layer first)" if layer_id in self._unpinned_layers: if not self._pending_whole_layer: @@ -1032,39 +1522,78 @@ def copy_missing(self) -> None: f"pageable materialize (position == expert id); ensure_experts's " f"LRU slot remap cannot be honored without a device alias" ) - # the only copy a non-pinned layer ever needs is the non-overlap prefill materialize, which schedules the whole layer into slots [0, num_experts) with position == expert id -- a plain synchronous pageable H2D copy + # the only copy a non-pinned layer ever needs is the non-overlap prefill materialize, which schedules the whole layer into staging buffer 0 with position == expert id -- a plain synchronous pageable H2D copy # never CUDA-graph captured: prefill is not captured, and decode never reaches this branch (it routes to the CPU executor) - for per_layer, cache in self.banks: - cache[: self.num_experts].copy_(per_layer[layer_id]) + for (per_layer, _), view in zip(self.banks, self.stage_views(layer_id, 0)): + view.copy_(per_layer[layer_id]) + if self.audit is not None: + from freetoken.moe.slot_audit import PAGEABLE_COPY + + self.audit.record_range(self, PAGEABLE_COPY, layer_id, 0, self.num_experts, meta=False) return + # materialize stages rows of staging buffer 0; ensure stages slots local to the layer's pool + self.copy_rows(layer_id, self.evict_slots, self.src_indices, self.num_indices, slim=_SLIM_COPY, + kind=self._pending_kind, width=self._pending_width, whole=self._pending_whole_layer) + + def copy_rows( + self, layer_id: int, dst_slots: torch.Tensor, src_rows: torch.Tensor, num: torch.Tensor, *, slim: bool, + kind: int = 0, width: int | None = None, whole: bool = False, + ) -> None: + """Copy host rows ``src_rows[:num]`` of ``layer_id``'s banks into slots ``dst_slots[:num]`` + on the current stream (the demand plan, or a prefetch plan on the prefetch copy stream); + ``kind`` labels the copy for the slot audit and ``width`` bounds ``num`` for it. ``dst_slots`` + are local to the layer's pool, or rows of staging buffer 0 when ``whole`` (materialize).""" + self._copy_rows(layer_id, dst_slots, src_rows, num, slim=slim, whole=whole) + if self.audit is not None: + from freetoken.moe.slot_audit import OTHER_COPY + + self.audit.record_plan(self, kind or OTHER_COPY, layer_id, dst_slots, src_rows, num, meta=False, width=width) + + def _copy_rows( + self, layer_id: int, dst_slots: torch.Tensor, src_rows: torch.Tensor, num: torch.Tensor, *, slim: bool, + whole: bool = False, + ) -> None: if self._copy_fused_ok: - from freetoken.kernel.fast_index_copy import fast_index_copy_multi_jit + from freetoken.kernel.fast_index_copy import fast_index_copy_multi_jit, fast_index_copy_multi_slim_jit # One launch copies the missing rows for every bank (instead of one launch per - # bank). evict_slots/src_indices/num_indices are shared across banks; - # src_indices holds layer-local expert rows, resolved against this layer's - # source pointers (layer_id is a static int per captured graph node). - fast_index_copy_multi_jit( - self._copy_dst_ptrs, - self._copy_src_ptrs[layer_id], - self._copy_feat_bytes, - self.evict_slots, - self.src_indices, - self.num_indices, - blocks_per_bank=self._copy_blocks_per_bank, - ) + # bank). The slot/row plan is shared across banks; src_rows holds layer-local + # expert rows, resolved against this layer's source pointers (layer_id is a + # static int per captured graph node). + dst_ptrs = self._copy_dst_ptrs if whole else self._copy_dst_ptrs_pool[self.pool_of_layer[layer_id]] + feat_bytes = self._copy_feat_bytes[self.layer_layout[layer_id]] + if slim: + fast_index_copy_multi_slim_jit(dst_ptrs, self._copy_src_ptrs[layer_id], feat_bytes, dst_slots, src_rows, num) + else: + fast_index_copy_multi_jit( + dst_ptrs, self._copy_src_ptrs[layer_id], feat_bytes, dst_slots, src_rows, num, + blocks_per_bank=self._copy_blocks_per_bank, + ) return from freetoken.kernel import fast_index_copy_jit - for per_layer, cache in self.banks: - fast_index_copy_jit( - cache, - self.evict_slots, - per_layer[layer_id], - self.src_indices, - self.num_indices, - ) + views = self.stage_views(layer_id, 0) if whole else self.pool_bank_views(layer_id) + for (per_layer, _), cache in zip(self.banks, views): + fast_index_copy_jit(cache, dst_slots, per_layer[layer_id], src_rows, num) + + @contextmanager + def decode_side(self, layer_id: int) -> Iterator[None]: + """Run the body on ``decode_copy_stream``, ordered after the work already on the current + stream; ``join_decode_side`` must follow before its results are read. + + The copy stays on the current stream right after its ensure, so it never waits on a + cross-stream edge; the side work (slack-rich) takes that latency and any SM contention. + Under capture the event pair becomes a fork and a join edge of the graph.""" + forked, done = self._decode_copy_events[layer_id] + forked.record(torch.cuda.current_stream(self.device)) + self.decode_copy_stream.wait_event(forked) + with torch.cuda.stream(self.decode_copy_stream): + yield + done.record(self.decode_copy_stream) + + def join_decode_side(self, layer_id: int) -> None: + torch.cuda.current_stream(self.device).wait_event(self._decode_copy_events[layer_id][1]) def iter_offload_moe_layers(model) -> Iterator: diff --git a/python/freetoken/moe/offload_kernels.py b/python/freetoken/moe/offload_kernels.py index 5df2ffec68..59a8f9637b 100644 --- a/python/freetoken/moe/offload_kernels.py +++ b/python/freetoken/moe/offload_kernels.py @@ -1,6 +1,7 @@ from __future__ import annotations import os +from typing import NamedTuple import torch import triton @@ -25,7 +26,49 @@ def ensure_experts(cache, layer_id: int, expert_ids: torch.Tensor, *, record_sta tensor. ``out_indices`` aliases the input, preserving the in-place rewrite every downstream GEMM depends on. """ + id_of_slot, usage = cache.pool_state(layer_id) lru_ensure( + expert_ids, + cache.slot_for_id.view(-1), + id_of_slot, + usage, + cache.step, + expert_ids, + cache.src_indices, + cache.evict_slots, + cache.num_indices, + stats=cache.lru_stats[layer_id] if cache.collect_stats and record_stats else None, + id_base=layer_id * cache.num_experts, + ) + + +class HybridRoutes(NamedTuple): + """One layer's routes split between the GPU and the CPU, each [bs, top_k].""" + + gpu_slots: torch.Tensor # cache slot of each route; a CPU route is -1, or 0 with clamp + gpu_w: torch.Tensor # router weights, 0 for a CPU route + cpu_ids: torch.Tensor # expert ids, -1 for a GPU route + + +def ensure_experts_scored( + cache, + layer_id: int, + expert_ids: torch.Tensor, + *, + bump_tok: bool, + update_state: bool, + router_logits: torch.Tensor | None = None, + pin_since: torch.Tensor | None = None, + lowpri: bool = False, + pf_count: tuple | None = None, +) -> None: + """``ensure_experts`` for the scored ``--moe-cache-policy`` kinds (kd, kdfb, rule), and for + every policy under FREETOKEN_MOE_PREFETCH=on (``lowpri``; ``pf_count`` is the prefetcher's + ``(stats_row, ready, num, rows)`` when this layer was prefetched this step).""" + from freetoken.moe.scored_ensure import scored_ensure + + stats_row, ready, num, rows = pf_count if pf_count is not None else (None, None, None, 0) + scored_ensure( expert_ids, cache.slot_for_id.view(-1), cache.id_of_slot, @@ -35,14 +78,62 @@ def ensure_experts(cache, layer_id: int, expert_ids: torch.Tensor, *, record_sta cache.src_indices, cache.evict_slots, cache.num_indices, - stats=cache.lru_stats[layer_id] if cache.collect_stats and record_stats else None, + router_logits=router_logits if cache.near_miss_thr is not None else None, + near_miss_thr=cache.near_miss_thr or 0.0, + stats=cache.lru_stats[layer_id] if cache.collect_stats else None, + bump_tok=bump_tok, + update_state=update_state, + pin_since=pin_since, + lowpri=lowpri, + pf_stats=stats_row, + pf_ready=ready, + pf_num=num, + pf_rows=rows, + **_evict_state(cache, layer_id), + ) + + +def _evict_state(cache, layer_id: int) -> dict: + return dict( + policy=cache.cache_policy_id, + num_layers=cache.num_layers, + num_experts=cache.num_experts, + tok=cache.evict_tok, + last_tok=cache.evict_last_tok, + lc=cache.evict_lc, + ct=cache.evict_ct, + g_table=cache.evict_g_table, + slot_owner=cache.evict_slot_owner, + slot_last_tok=cache.evict_slot_last_tok, + slot_lc=cache.evict_slot_lc, + slot_ct=cache.evict_slot_ct, id_base=layer_id * cache.num_experts, ) -def ensure_experts_hybrid( - cache, layer_id: int, expert_ids: torch.Tensor, max_fetch: int, fetch_fraction: float = 0.0 +def prefetch_ensure_experts( + cache, layer_id: int, query: torch.Tensor, dst_slots: torch.Tensor, src_rows: torch.Tensor, + num: torch.Tensor, stats_row: torch.Tensor, ready: torch.Tensor, budget: int | None = None, ) -> None: + """FREETOKEN_MOE_PREFETCH=on: install the first ``budget`` of ``query``'s non-resident experts + of ``layer_id`` as low-priority slots and plan their copy (scored_ensure.prefetch_ensure).""" + from freetoken.moe.scored_ensure import prefetch_ensure + + prefetch_ensure( + query, cache.slot_for_id.view(-1), cache.id_of_slot, cache.usage, cache.step, + dst_slots, src_rows, num, stats_row, ready, budget=budget, **_evict_state(cache, layer_id), + ) + + +def ensure_experts_hybrid( + cache, + layer_id: int, + expert_ids: torch.Tensor, + weights: torch.Tensor, + max_fetch: int, + fetch_fraction: float = 0.0, + clamp: bool = False, +) -> HybridRoutes: """Capped-fetch variant of ``ensure_experts`` (hybrid backend). Identical LRU bookkeeping, but only the first ``max_fetch`` of this step's missing @@ -52,21 +143,26 @@ def ensure_experts_hybrid( split (fraction = pcie_bw / cpu_bw): fetch ~fraction of the step's misses, rounded to the integer that makes the PCIe fetch and the CPU overflow compute finish closest to together. ``num_indices`` = capped fetch count (copy_missing); ``num_missing_full`` = - pre-cap miss count (stats).""" - # Q16 fixed point so the GPU kernel and the CPU reference cap identically (no float). + pre-cap miss count (stats). + + The same launch splits the routes: ``expert_ids`` becomes the GPU slots in place, and + ``clamp`` gives CPU routes slot 0 instead of -1 for GEMM kernels that read every slot.""" + assert expert_ids.is_cuda, "the hybrid ensure kernel runs on the GPU" + # Q16 fixed point: the kernel caps in integers. frac_q16 = min(1 << 16, max(0, round(fetch_fraction * (1 << 16)))) - if not expert_ids.is_cuda: - return _ensure_experts_hybrid_cpu(cache, layer_id, expert_ids, max_fetch, frac_q16) - _ensure_experts_hybrid_gpu(cache, layer_id, expert_ids, max_fetch, frac_q16) + weights = weights.contiguous() + routes = HybridRoutes(expert_ids, torch.empty_like(weights), torch.empty_like(expert_ids)) + _ensure_experts_hybrid_gpu(cache, layer_id, expert_ids, weights, routes, clamp, max_fetch, frac_q16) + return routes -def prefill_hit_compact(cache, layer_id: int, buffer_id: int) -> None: +def prefill_hit_compact(cache, layer_id: int) -> None: """Compact this layer's cache-resident experts into gather indices, device-side. - hit = slot_for_id[layer_id][e] >= 2 * num_experts (the double buffer owns the + hit = slot_for_id[layer_id][e] >= the pool's staging threshold (the double buffer owns the slots below, so those bytes are volatile within a prefill chunk and classify as miss). Writes fixed-shape ``_prefill_hit_dst``/``_prefill_hit_src`` (buffer - row / cache slot) and the count into ``_prefill_hit_num``; one launch on the + row / pool-local slot) and the count into ``_prefill_hit_num``; one launch on the current stream, no host sync. Safe against the concurrent buffer invalidation on the copy stream: that only rewrites entries already below the threshold.""" num_experts = cache.num_experts @@ -75,15 +171,23 @@ def prefill_hit_compact(cache, layer_id: int, buffer_id: int) -> None: cache._prefill_hit_dst, cache._prefill_hit_src, cache._prefill_hit_num, - buffer_id * num_experts, - 2 * num_experts, + cache._pool_stable_from[cache.pool_of_layer[layer_id]], num_experts, BLOCK=triton.next_power_of_2(num_experts), ) def materialize_layer(cache, layer_id: int) -> None: - _materialize_layer_gpu(cache, layer_id) + if len(cache.pool_sizes) == 1: + _materialize_layer_gpu(cache, layer_id) + return + # staging buffer 0 may overlap several pools, so it only holds the layer for this forward and registers nothing + E = cache.num_experts + cache._invalidate_prefill_buffer(0) + rows = torch.arange(E, dtype=torch.int32, device=cache.device) + cache.evict_slots[:E].copy_(rows) + cache.src_indices[:E].copy_(rows) + cache.num_indices.fill_(E) def reset_cache(cache) -> None: @@ -95,16 +199,25 @@ def reset_cache(cache) -> None: def _ensure_experts_hybrid_gpu( - cache, layer_id: int, expert_ids: torch.Tensor, max_fetch: int, frac_q16: int + cache, + layer_id: int, + expert_ids: torch.Tensor, + weights: torch.Tensor, + routes: HybridRoutes, + clamp: bool, + max_fetch: int, + frac_q16: int, ) -> None: + id_of_slot, usage = cache.pool_state(layer_id) + num_slots = id_of_slot.numel() block_e = triton.next_power_of_2(cache.num_experts) - block_c = triton.next_power_of_2(cache.cache_size) + block_c = triton.next_power_of_2(num_slots) num_warps = 8 if block_c >= 2048 else 4 _ensure_experts_hybrid_kernel[(1,)]( expert_ids, cache.slot_for_id, - cache.id_of_slot, - cache.usage, + id_of_slot, + usage, cache.step, cache.active_mask, cache.evict_slots, @@ -112,82 +225,25 @@ def _ensure_experts_hybrid_gpu( cache.num_indices, cache.num_missing_full, cache.expert_recency, + weights, + routes.cpu_ids, + routes.gpu_w, + *(cache.hybrid_stats if cache.collect_stats else (expert_ids,) * 2), layer_id, expert_ids.numel(), int(max_fetch), int(frac_q16), cache.num_experts, - cache.cache_size, + num_slots, BLOCK_E=block_e, BLOCK_C=block_c, BY_RECENCY=_HYBRID_FETCH_BY_RECENCY, + CLAMP=clamp, + STATS=cache.collect_stats, num_warps=num_warps, ) -def _ensure_experts_hybrid_cpu( - cache, layer_id: int, expert_ids: torch.Tensor, max_fetch: int, frac_q16: int -) -> None: - """CPU reference mirror of the hybrid kernel (eviction/fetch decisions bit-identical to - the GPU path; see tests/test_offload_lru_kernels.py). Fetches at most ``max_fetch`` (or - the bandwidth-matched ``~frac_q16/2^16 * misses`` when ``frac_q16`` > 0) of the missing - experts; overflow misses are rewritten to -1. With ``BY_RECENCY`` the fetch set is the - most-recently-active misses (ties -> lower id); else the lowest ids.""" - seen = [] - for expert in expert_ids.view(-1).tolist(): - if expert not in seen: - seen.append(expert) - - cache.active_mask.zero_() - step = int(cache.step.item()) + 1 - cache.step.fill_(step) - for expert in seen: - cache.active_mask[expert] = 1 - - for expert in seen: - slot = int(cache.slot_for_id[layer_id, expert].item()) - if slot != -1: - cache.usage[slot] = step - - missing = [e for e in seen if int(cache.slot_for_id[layer_id, e].item()) == -1] - if _HYBRID_FETCH_BY_RECENCY: - rec = cache.expert_recency[layer_id].tolist() - missing.sort(key=lambda e: (-rec[e], e)) - else: - missing.sort() - if frac_q16 > 0: - m, q = len(missing), 1 << 16 - lo = (m * frac_q16) >> 16 - cost = lambda f: max(f * (q - frac_q16), (m - f) * frac_q16) # noqa: E731 - max_fetch = lo if cost(lo) <= cost(lo + 1) else lo + 1 - num_fetch = min(len(missing), int(max_fetch)) - cache.num_missing_full.fill_(len(missing)) - cache.num_indices.fill_(num_fetch) - - usage = cache.usage.tolist() - for idx in range(num_fetch): - expert = missing[idx] - victim = min(range(cache.cache_size), key=lambda s: (usage[s], s)) - old_id = int(cache.id_of_slot[victim].item()) - if old_id >= 0: - cache.slot_for_id.view(-1)[old_id] = -1 - cache.id_of_slot[victim] = layer_id * cache.num_experts + expert - cache.slot_for_id[layer_id, expert] = victim - cache.usage[victim] = step - usage[victim] = step - cache.evict_slots[idx] = victim - cache.src_indices[idx] = expert # layer-local row - - if _HYBRID_FETCH_BY_RECENCY: - for expert in seen: - cache.expert_recency[layer_id, expert] = step - - # Overflow misses keep slot_for_id == -1, so the rewrite below yields -1 for them. - flat = expert_ids.view(-1) - for i in range(flat.numel()): - flat[i] = int(cache.slot_for_id[layer_id, int(flat[i].item())].item()) - - def _materialize_layer_gpu(cache, layer_id: int) -> None: block = triton.next_power_of_2(max(cache.num_experts, cache.cache_size)) _materialize_layer_kernel[(1,)]( @@ -287,7 +343,7 @@ def _materialize_layer_kernel( -@triton.jit(do_not_specialize=["layer_id", "num_active", "max_fetch", "fetch_frac_q16"]) +@triton.jit(do_not_specialize=["layer_id", "num_active", "max_fetch", "fetch_frac_q16", "cache_size", "id_of_slot_ptr", "usage_ptr"]) def _ensure_experts_hybrid_kernel( expert_ids_ptr, slot_for_id_ptr, @@ -300,15 +356,22 @@ def _ensure_experts_hybrid_kernel( num_indices_ptr, num_missing_full_ptr, expert_recency_ptr, + weights_ptr, + cpu_ids_ptr, + gpu_w_ptr, + stats_ptr, + stats_layer_ptr, layer_id, num_active, max_fetch, fetch_frac_q16, num_experts: tl.constexpr, - cache_size: tl.constexpr, + cache_size, BLOCK_E: tl.constexpr, BLOCK_C: tl.constexpr, BY_RECENCY: tl.constexpr, + CLAMP: tl.constexpr, + STATS: tl.constexpr, ): """Capped-fetch timestamp-LRU (hybrid backend). @@ -352,10 +415,11 @@ def _ensure_experts_hybrid_kernel( (lo + 1) * ((1 << 16) - fetch_frac_q16), (num_missing - lo - 1) * fetch_frac_q16 ) max_fetch = tl.where(cost_lo <= cost_hi, lo, lo + 1) - num_fetch = tl.minimum(num_missing, max_fetch) + is_hit = is_active & (slot >= 0) + # a pool smaller than num_experts may hold fewer free slots than misses + num_fetch = tl.minimum(tl.minimum(num_missing, max_fetch), cache_size - tl.sum(is_hit.to(tl.int32))) tl.store(num_missing_full_ptr, num_missing.to(tl.int64)) tl.store(num_indices_ptr, num_fetch.to(tl.int64)) - is_hit = is_active & (slot >= 0) tl.store(usage_ptr + slot, step, mask=is_hit) # Fetch-selection priority: encode (recency desc, id asc) into one strictly-ordered @@ -401,8 +465,26 @@ def _ensure_experts_hybrid_kernel( for i in tl.range(num_active): e = tl.load(expert_ids_ptr + i) s = tl.load(slot_for_id_ptr + base + e) + on_gpu = s >= 0 + tl.store(cpu_ids_ptr + i, tl.where(on_gpu, -1, e)) + tl.store(gpu_w_ptr + i, tl.where(on_gpu, tl.load(weights_ptr + i), 0.0)) + if CLAMP: + s = tl.maximum(s, 0) tl.store(expert_ids_ptr + i, s) + # stats_ptr: (missing, fetched, active, calls); stats_layer_ptr: the same per layer, [num_layers, 4] + if STATS: + n_active = tl.sum((is_active & e_mask).to(tl.int64)) + tl.store(stats_ptr + 0, tl.load(stats_ptr + 0) + num_missing) + tl.store(stats_ptr + 1, tl.load(stats_ptr + 1) + num_fetch) + tl.store(stats_ptr + 2, tl.load(stats_ptr + 2) + n_active) + tl.store(stats_ptr + 3, tl.load(stats_ptr + 3) + 1) + row = stats_layer_ptr + layer_id * 4 + tl.store(row + 0, tl.load(row + 0) + num_missing) + tl.store(row + 1, tl.load(row + 1) + num_fetch) + tl.store(row + 2, tl.load(row + 2) + n_active) + tl.store(row + 3, tl.load(row + 3) + 1) + # Bump every active expert's recency to this step (LRU on the expert): an overflow miss # computed on the CPU now ranks high if it recurs, so it gets fetched next time. if BY_RECENCY: @@ -410,14 +492,13 @@ def _ensure_experts_hybrid_kernel( tl.store(expert_recency_ptr + base + off_e, step_vec, mask=is_active & e_mask) -@triton.jit(do_not_specialize=["buffer_base"]) +@triton.jit(do_not_specialize=["threshold"]) def _prefill_hit_compact_kernel( slot_ptr, # [num_experts] int32: this layer's slot_for_id row - dst_ptr, # [num_experts] int32 out: buffer rows, compacted - src_ptr, # [num_experts] int32 out: cache slots, compacted + dst_ptr, # [num_experts] int32 out: buffer rows (expert ids), compacted + src_ptr, # [num_experts] int32 out: pool-local slots, compacted num_ptr, # [1] int64 out: hit count - buffer_base, # buffer_id * num_experts - threshold, # 2 * num_experts + threshold, # first pool-local slot outside the staging buffers num_experts, BLOCK: tl.constexpr, ): @@ -426,6 +507,6 @@ def _prefill_hit_compact_kernel( slots = tl.load(slot_ptr + offs, mask=lane, other=-1) is_hit = lane & (slots >= threshold) pos = tl.cumsum(is_hit.to(tl.int32)) - 1 - tl.store(dst_ptr + pos, (buffer_base + offs).to(tl.int32), mask=is_hit) + tl.store(dst_ptr + pos, offs.to(tl.int32), mask=is_hit) tl.store(src_ptr + pos, slots, mask=is_hit) tl.store(num_ptr, tl.sum(is_hit.to(tl.int64))) diff --git a/python/freetoken/moe/pool_policy.py b/python/freetoken/moe/pool_policy.py new file mode 100644 index 0000000000..990c18ade2 --- /dev/null +++ b/python/freetoken/moe/pool_policy.py @@ -0,0 +1,149 @@ +"""Slot pool policy for the offload MoE cache: which layers share a pool, and how many slots each pool gets. + +Pure functions. The engine supplies the constraints and hands the result to ``OffloadMoeCache``: + + groups = group_layers(scope, layouts, slot_bytes) # --moe-cache-scope, once per load + pools = make_pools(groups, num_experts, min_num_slots, slot_limit) # the minimum follows the decode batch + sizes = split_slots(total, pools, counts) # --moe-cache-split, counts from resolve_profile + OffloadMoeCache(pool_of_layer=pool_of_layer(pools, num_layers), pool_sizes=sizes, ...) +""" + +from __future__ import annotations + +import json +from dataclasses import dataclass +from pathlib import Path + +import torch + +from freetoken.utils import init_logger + +logger = init_logger(__name__) + +SCOPES = ("unified", "per-layer") + + +@dataclass(frozen=True) +class PoolSpec: + layers: tuple[int, ...] + min_num_slots: int + max_num_slots: int + + +def group_layers(scope: str, layouts: list, slot_bytes: list[int]) -> list[list[int]]: + """The layers of each pool, in placement order. + + ``unified``: one pool per layout key (see ``bank_layouts``); ``per-layer``: one pool per layer. + Pools with wider slots come first: the prefill staging buffers sit at the start of every bank, so they then cover fewer slots. + """ + if scope == "unified": + groups: dict = {} + for layer, layout in enumerate(layouts): + groups.setdefault(layout, []).append(layer) + pools = list(groups.values()) + elif scope == "per-layer": + pools = [[layer] for layer in range(len(layouts))] + else: + raise ValueError(f"unknown --moe-cache-scope {scope!r}; want one of {SCOPES}") + return sorted(pools, key=lambda layers: (-slot_bytes[layers[0]], layers[0])) + + +def make_pools(groups: list[list[int]], num_experts: int, min_num_slots: int, slot_limit: int | None) -> list[PoolSpec]: + """Each group's slot range: at least ``min_num_slots``, at most one slot per expert of its layers and the kernel's ``slot_limit``.""" + pools = [] + for layers in groups: + max_num_slots = len(layers) * num_experts + if slot_limit is not None: + max_num_slots = min(max_num_slots, slot_limit) + pools.append(PoolSpec(tuple(layers), min(min_num_slots, max_num_slots), max_num_slots)) + return pools + + +def pool_of_layer(pools: list[PoolSpec], num_layers: int) -> list[int]: + """Each layer's pool index, the form ``OffloadMoeCache`` takes.""" + owner = [-1] * num_layers + for p, pool in enumerate(pools): + for layer in pool.layers: + owner[layer] = p + assert min(owner) >= 0, owner + return owner + + +def split_slots(total: int, pools: list[PoolSpec], counts: torch.Tensor | None = None) -> list[int]: + """Slots per pool, each in [min_num_slots, max_num_slots], summing to ``total``. + + With ``counts`` (per-layer expert use from ``resolve_profile``) the slots follow expert use, else the pools' layer counts. + """ + if len(pools) == 1: + return [total] + _check_budget(total, pools) + return _split_equal(total, pools) if counts is None else _split_hot(total, pools, counts) + + +def _check_budget(total: int, pools: list[PoolSpec]) -> None: + need = sum(p.min_num_slots for p in pools) + room = sum(p.max_num_slots for p in pools) + if total < need: + raise ValueError(f"MoE cache of {total} slots is below what the pools need ({need} slots)") + if total > room: + raise ValueError(f"MoE cache of {total} slots exceeds what the pools can hold ({room} slots)") + + +def _split_equal(total: int, pools: list[PoolSpec]) -> list[int]: + sizes = [p.min_num_slots for p in pools] + weights = [len(p.layers) for p in pools] + left = total - sum(sizes) + while left > 0: + open_ = [i for i, p in enumerate(pools) if sizes[i] < p.max_num_slots] + # give the next slot to the pool with the fewest slots per layer + i = min(open_, key=lambda j: (sizes[j] / weights[j], j)) + sizes[i] += 1 + left -= 1 + return sizes + + +def _split_hot(total: int, pools: list[PoolSpec], counts: torch.Tensor) -> list[int]: + sizes = [p.min_num_slots for p in pools] + # a pool's minimum covers its hottest experts; the rest, up to its maximum, competes across pools + cand = [] + for p, pool in enumerate(pools): + hot = counts[list(pool.layers)].flatten().sort(descending=True).values + cand.extend((float(c), p) for c in hot[pool.min_num_slots : pool.max_num_slots]) + cand.sort(key=lambda c: -c[0]) + for _, p in cand[: total - sum(sizes)]: + sizes[p] += 1 + return sizes + + +def resolve_profile(spec: str, model_path: str, num_layers: int, num_experts: int) -> torch.Tensor | None: + """Expert use for ``--moe-cache-split``: None for ``equal``, else the profile at ``profile:``.""" + if spec == "equal": + return None + path = Path(spec.removeprefix("profile:")) + logger.info_rank0(f"MoE pool profile: {path}") + return load_profile(path, num_layers, num_experts, model_path) + + +def load_profile(path: str | Path, num_layers: int, num_experts: int, model_path: str = "") -> torch.Tensor: + """Per-layer expert use ``[num_layers, num_experts]`` from a profile, each row scaled to sum to 1.""" + data = json.loads(Path(path).read_text()) + model = str(data.get("model", "")) + if model and model_path and not any(n.lower().startswith(model.lower()) for n in _model_names(model_path)): + logger.warning_rank0(f"MoE pool profile {path} was taken on {model!r}, not on {model_path}") + counts = torch.tensor(data["expert_counts"], dtype=torch.float64) + if tuple(counts.shape) != (num_layers, num_experts): + raise ValueError( + f"profile {path} has {tuple(counts.shape)} expert counts; this model needs ({num_layers}, {num_experts})" + ) + # every layer decodes the same tokens, so equal row sums keep the layers comparable + return counts / counts.sum(dim=1, keepdim=True).clamp(min=1e-12) + + +def _model_names(model_path: str) -> list[str]: + # every path component, so an HF cache snapshot (models--Org--Name/snapshots/) still finds its name + names = [] + for part in reversed(Path(model_path).parts): + name = part.split("--")[-1].removesuffix(".gguf") + if name: + names.append(name) + return names diff --git a/python/freetoken/moe/prefetch.py b/python/freetoken/moe/prefetch.py new file mode 100644 index 0000000000..01129d3997 --- /dev/null +++ b/python/freetoken/moe/prefetch.py @@ -0,0 +1,438 @@ +"""Cross-layer expert prefetch for GPU decode (``FREETOKEN_MOE_PREFETCH``). + +A training-free router lookahead predicts layer L's experts from the tensor layer L-1's router +consumed (``W_gate[L] . x_{L-1}``, top-K) and keeps the first BUDGET candidates that are not +resident. At bs 2 both rows' top-K lists merge rank by rank (row 0's first, then row 1's +first, ...), duplicates dropped, under one budget. + +``measure``: after layer L's ensure, count how many of them layer L actually routed. Nothing is +copied and the cache is never touched. The predictor reads ``slot_for_id[L]`` without ordering +against the ensures; a late read only skews the counters, never the model. + +The prediction forks at the start of layer L-1's MoE block, where x_{L-1} first exists, so it runs +beside layer L-1's own router and ensure: + + compute: -fork(L)-> router_{L-1} ensure(L-1) copy(L-1) GEMM(L-1) ... ensure(L) copy(L) -join-> count(L) GEMM(L) + predict: \\-> router_L(x_{L-1}) -> top-K -> select(L) -rec sel(L) -/ + +``on``: the predictor only merges the candidates (it reads no cache state); the compute stream +installs the first BUDGET of them that are not resident as low-priority slots (prefetch_ensure: +held, usage 0, keyed between the empty slots and every resident) and a third stream copies them +from layer L's host banks with the slim copy kernel: + + compute: -fork(L)-> ... ensure(L-1) [wait sel(L)] prefetch_ensure(L) [wait pfcopy(L-1)] copy(L-1) -rec dcopy(L)-> GEMM(L-1) ... ensure(L) [wait pfcopy(L)] copy(L) GEMM(L) + predict: \\-> router_L -> top-K -> merge(L) -rec sel(L) + copy: [wait dcopy(L)] slim copy(L) -> ready(L) -rec pfcopy(L) + +Every writer of the slot maps and their mirrors runs on the compute stream, so the maps have one +writer in stream order and no cross-stream edge guards them, and the installs follow layer L's +residency right after ensure(L-1) however late the predictor runs. dcopy(L) follows +prefetch_ensure(L) and the demand copy(L-1) on the compute stream, so the copy stream reads a +finished plan and owns the host link second. prefetch_ensure(L) keeps usage == step, i.e. layer +L-1's routed slots, pinned while GEMM(L-1) may read them; the compute stream joins the prefetch copy +before its own copy(L), so ensure(L) may reuse a prefetched slot as a demand victim and GEMM(L) only +ever reads landed bytes. A demand hit on a prefetched slot sets its usage (a normal resident) and +counts as useful. Layer 0 is never predicted. +""" +from __future__ import annotations + +import torch +import triton +import triton.language as tl + +from freetoken.env import ENV +from freetoken.moe.slot_audit import PREFETCH_COPY +from freetoken.utils import init_logger + +logger = init_logger(__name__) + +MODES = ("off", "measure", "on") +# non-resident candidates per target layer: GDN layers leave a shorter copy window than QSA ones +GDN_BUDGET = 3 +ATTN_BUDGET = 4 +# counters[layer, col]; on mode adds LATE and COPIED (calls whose plan had a copy) in a wider stats row +ISSUED, USEFUL, RESIDENT_HITS, CALLS, ROWS, MISSES, LATE, COPIED = range(8) +NUM_COLS = 6 +NUM_STAT_COLS = 8 +# the select kernel's dedup tile is (rows * k)^2; larger (eager-only) decode batches skip the prediction +MAX_ROWS = 8 + + +def resolve_mode(mode: str | None = None) -> str: + mode = (ENV.MOE_PREFETCH.value if mode is None else mode).strip().lower() + if mode not in MODES: + raise ValueError(f"FREETOKEN_MOE_PREFETCH={mode!r}: expected one of {', '.join(MODES)}") + return mode + + +def default_budget(before_linear: bool) -> int: + """Non-resident candidates kept for a target layer, by the target's attention kind.""" + return GDN_BUDGET if before_linear else ATTN_BUDGET + + +def dedicated_stream(device: torch.device, *, highest_priority: bool = False) -> torch.cuda.ExternalStream: + # not one of torch's 32 pooled streams: a pooled one can alias the engine or capture stream + from cuda.bindings import runtime as cudart + + with torch.cuda.device(device): + if highest_priority: + err, _, greatest = cudart.cudaDeviceGetStreamPriorityRange() + if err == cudart.cudaError_t.cudaSuccess: + err, handle = cudart.cudaStreamCreateWithPriority(cudart.cudaStreamNonBlocking, greatest) + else: + err, handle = cudart.cudaStreamCreateWithFlags(cudart.cudaStreamNonBlocking) + if err != cudart.cudaError_t.cudaSuccess: + raise RuntimeError(f"creating a dedicated CUDA stream failed: {err}") + return torch.cuda.ExternalStream(int(handle), device=device) + + +@triton.jit(do_not_specialize=["stride", "budget"]) +def _lookahead_select_kernel( + logits_ptr, stride, resident_ptr, sel_ptr, res_ptr, budget, + ROWS: tl.constexpr, ROWS_P: tl.constexpr, E: tl.constexpr, E_P: tl.constexpr, + K: tl.constexpr, N_P: tl.constexpr, W: tl.constexpr, W_P: tl.constexpr, +): + b = tl.arange(0, ROWS_P) + e = tl.arange(0, E_P) + live = (b < ROWS)[:, None] & (e < E)[None, :] + v = tl.load(logits_ptr + b[:, None] * stride + e[None, :], mask=live, other=float("-inf")).to(tl.float32) + # candidates in merged order: position r * ROWS + b holds row b's rank-r id + j = tl.arange(0, N_P) + cand = tl.full([N_P], -1, tl.int32) + for r in tl.static_range(K): + top = tl.argmax(v, axis=1).to(tl.int32) # ties go to the lowest id + at = (j[:, None] == r * ROWS + b[None, :]) & (b < ROWS)[None, :] + cand = tl.where(tl.sum(at.to(tl.int32), axis=1) > 0, tl.sum(tl.where(at, top[None, :], 0), axis=1), cand) + v = tl.where(e[None, :] == top[:, None], float("-inf"), v) + valid = j < ROWS * K + earlier = (cand[:, None] == cand[None, :]) & (j[None, :] < j[:, None]) & valid[None, :] + first = valid & (tl.sum(earlier.to(tl.int32), axis=1) == 0) + held = tl.load(resident_ptr + cand, mask=first, other=-1) >= 0 + pick = first & (~held) + rank = tl.cumsum(pick.to(tl.int32), axis=0) - 1 + kept = first & held & (tl.cumsum(first.to(tl.int32), axis=0) <= budget) + krank = tl.cumsum(kept.to(tl.int32), axis=0) - 1 + p = tl.arange(0, W_P) + m_sel = pick[None, :] & (rank[None, :] == p[:, None]) & (p < budget)[:, None] + m_res = kept[None, :] & (krank[None, :] == p[:, None]) + # at most one candidate matches each output position + sel = tl.sum(tl.where(m_sel, cand[None, :] + 1, 0), axis=1) - 1 + res = tl.sum(tl.where(m_res, cand[None, :] + 1, 0), axis=1) - 1 + tl.store(sel_ptr + p, sel, mask=p < W) + tl.store(res_ptr + p, res, mask=p < W) + + +@triton.jit(do_not_specialize=["id_base"]) +def _prefetch_count_kernel( + slots_ptr, id_of_slot_ptr, misses_ptr, sel_ptr, res_ptr, counters_ptr, id_base, + N: tl.constexpr, N_P: tl.constexpr, W: tl.constexpr, W_P: tl.constexpr, ROWS: tl.constexpr, +): + i = tl.arange(0, N_P) + slot = tl.load(slots_ptr + i, mask=i < N, other=-1) + # ensure rewrote the routed ids to slots, all pinned to this layer's ids until the next ensure + routed = tl.where(slot >= 0, tl.load(id_of_slot_ptr + slot, mask=slot >= 0, other=0) - id_base, -1) + p = tl.arange(0, W_P) + sel = tl.load(sel_ptr + p, mask=p < W, other=-1) + res = tl.load(res_ptr + p, mask=p < W, other=-1) + sel_hit = tl.sum(((sel[:, None] == routed[None, :]) & (i < N)[None, :]).to(tl.int32), axis=1) > 0 + res_hit = tl.sum(((res[:, None] == routed[None, :]) & (i < N)[None, :]).to(tl.int32), axis=1) > 0 + issued = tl.sum((sel >= 0).to(tl.int32)) + useful = tl.sum((sel_hit & (sel >= 0)).to(tl.int32)) + res_hits = tl.sum((res_hit & (res >= 0)).to(tl.int32)) + misses = tl.load(misses_ptr) + c = tl.arange(0, 8) + add = tl.where(c == 0, issued, tl.where(c == 1, useful, tl.where(c == 2, res_hits, tl.where(c == 3, 1, ROWS)))) + add = tl.where(c == 5, misses, add.to(tl.int64)) + # one writer per layer row, stream-ordered: no atomics needed + old = tl.load(counters_ptr + c, mask=c < 6, other=0) + tl.store(counters_ptr + c, old + add, mask=c < 6) + + +def lookahead_select( + logits: torch.Tensor, resident: torch.Tensor, sel: torch.Tensor, res: torch.Tensor, *, k: int, budget: int +) -> None: + """Top-``k`` of each ``logits`` row, merged rank by rank; ``sel`` gets the first ``budget`` + distinct ids with ``resident[id] < 0`` and ``res`` the resident ids among the first ``budget`` + distinct ones, both in merge order and padded with -1. One CTA, no host sync.""" + rows, num_experts = logits.shape + width = sel.numel() + assert logits.stride(1) == 1 and resident.numel() == num_experts and resident.dtype == torch.int32 + assert sel.dtype == res.dtype == torch.int32 and res.numel() == width and 1 <= k <= num_experts + _lookahead_select_kernel[(1,)]( + logits, logits.stride(0), resident, sel, res, min(budget, width), + ROWS=rows, ROWS_P=triton.next_power_of_2(rows), + E=num_experts, E_P=triton.next_power_of_2(num_experts), + K=k, N_P=triton.next_power_of_2(rows * k), + W=width, W_P=triton.next_power_of_2(width), + num_warps=4, + ) + + +def prefetch_count( + slots: torch.Tensor, id_of_slot: torch.Tensor, misses: torch.Tensor, + sel: torch.Tensor, res: torch.Tensor, counters: torch.Tensor, *, id_base: int, rows: int, +) -> None: + """Add one call's issued / useful / resident-hit / call / row / miss counts to ``counters``.""" + assert slots.dtype == torch.int32 and slots.is_contiguous() and counters.dtype == torch.int64 + n, width = slots.numel(), sel.numel() + _prefetch_count_kernel[(1,)]( + slots, id_of_slot, misses, sel, res, counters, id_base, + N=n, N_P=triton.next_power_of_2(n), W=width, W_P=triton.next_power_of_2(width), ROWS=rows, + num_warps=1, + ) + + +@triton.jit +def _mark_ready_kernel(ready_ptr): + tl.store(ready_ptr, 1) + + +class ExpertPrefetcher: + """Predictor stream, per-layer buffers and device counters for the router lookahead. + + ``fork`` runs layer L's prediction on the dedicated stream at the start of layer L-1's MoE + block. ``measure``: ``join_and_count`` waits for it after layer L's copy is enqueued and counts + on the compute stream. ``on``: ``install`` (right after layer L-1's ensure) waits for it and + runs prefetch_ensure(L) on the compute stream; ``issue_copy`` puts its copy on a second + dedicated stream after layer L-1's demand copy; the compute stream joins the copy before its + copy(L) (``join_copy``) and ensure(L) counts the call (``count_args``). The events are created + up front, the shapes are fixed and nothing syncs the host, so every side captures into the + decode graph. Nothing here is cache_size-shaped: a rebuild keeps it and only resets the counters. + """ + + def __init__( + self, num_layers: int, num_experts: int, device: torch.device, *, + mode: str = "measure", k: int | None = None, budget: int | None = None, + ) -> None: + self.mode = mode + self.num_layers = num_layers + self.num_experts = num_experts + self.device = device + k = ENV.MOE_PREFETCH_K.value if k is None else int(k) + if k < 1: + raise ValueError(f"FREETOKEN_MOE_PREFETCH_K={k} must be >= 1") + self.k = min(k, num_experts) + # 0 = the per-layer default the model wires in + self.budget_override = ENV.MOE_PREFETCH_BUDGET.value if budget is None else int(budget) + if self.budget_override < 0: + raise ValueError(f"FREETOKEN_MOE_PREFETCH_BUDGET={self.budget_override} must be >= 0") + self.stream = dedicated_stream(device) + # (forked, selected): the predictor starts after forked and records selected after the select + self._events = [(torch.cuda.Event(), torch.cuda.Event()) for _ in range(num_layers)] + # torch creates the CUDA event on its first record: do that now, never inside a capture + for pair in self._events: + for event in pair: + event.record(self.stream) + # on mode's select sees nothing resident: prefetch_ensure applies residency and the budget + self._none_resident = torch.full((num_experts,), -1, dtype=torch.int32, device=device) + self._budget = [0] * num_layers + width = 2 * self.k # room for the whole bs-2 union + self.sel = torch.full((num_layers, width), -1, dtype=torch.int32, device=device) + self.res = torch.full((num_layers, width), -1, dtype=torch.int32, device=device) + self.stats = torch.zeros((num_layers, NUM_STAT_COLS), dtype=torch.int64, device=device) + # measure mode's six columns, a view: its count kernel and readers keep the stage-A layout + self.counters = self.stats[:, :NUM_COLS] + self.totals = torch.zeros((num_layers, NUM_STAT_COLS), dtype=torch.int64) + # the router input of each forked prediction, alive until the compute stream has waited for + # the select so the allocator cannot hand its block to later work the predictor still reads + self._inflight: list[torch.Tensor | None] = [None] * num_layers + # on mode: whether install ran prefetch_ensure for the forked prediction (its copy is then due) + self._installed = [False] * num_layers + # on mode: the cache whose slot maps prefetch_ensure writes and whose banks the copy reads + self.cache = None + self.copy_stream = None + if mode == "on": + # the slim copy is an SM kernel: at default priority a late copy queues for SMs behind the next + # compute grid while the compute stream waits on it (bench --on: a 50 us slip cost 14% per step) + self.copy_stream = dedicated_stream(device, highest_priority=True) + # (dcopy: after prefetch_ensure(L) and layer L-1's demand copy on the compute stream, + # pfcopy: after layer L's prefetch copy) + self._copy_events = [(torch.cuda.Event(), torch.cuda.Event()) for _ in range(num_layers)] + for pair in self._copy_events: + for event in pair: + event.record(self.copy_stream) + # one plan per target layer: prefetch_ensure(L+1) must not overwrite the plan copy(L) reads + self.pf_slots = torch.full((num_layers, width), -1, dtype=torch.int32, device=device) + self.pf_src = torch.full((num_layers, width), -1, dtype=torch.int32, device=device) + self.pf_num = torch.zeros((num_layers, 1), dtype=torch.int64, device=device) + self.pf_ready = torch.zeros((num_layers, 1), dtype=torch.int32, device=device) + # test hooks, read at enqueue (capture) time: a [num_layers, width] device buffer replacing + # sel (on mode still installs at most the budget), and sleeps put on the predictor / copy + # stream ahead of their work + self.sel_override: torch.Tensor | None = None + self.delay_predict_cycles = 0 + self.delay_copy_cycles = 0 + logger.info( + f"MoE prefetch {mode}: router lookahead top-{self.k}, budget " + + (str(self.budget_override) if self.budget_override else f"{GDN_BUDGET} before GDN / {ATTN_BUDGET} before attention layers") + + (", slim copies on a dedicated high-priority stream" if mode == "on" else "") + ) + + def budget_for(self, wired: int) -> int: + return min(self.budget_override or wired, self.sel.shape[1]) + + def fork(self, target: int, x: torch.Tensor, gate, resident: torch.Tensor, budget: int) -> bool: + """Predict ``target``'s experts from ``x`` (the previous layer's router input) on the + predictor stream, ordered after the work already on the current stream. Measure mode skips + the ids ``resident`` holds (a racy read of the cache); on mode reads no cache state and + leaves residency and the budget to ``install``. Returns whether it forked.""" + if x.shape[0] > MAX_ROWS: + return False + budget = self.budget_for(budget) + if self.mode == "on": + self._budget[target] = budget + resident, budget = self._none_resident, self.sel.shape[1] + forked, selected = self._events[target] + forked.record(torch.cuda.current_stream(self.device)) + self.stream.wait_event(forked) + with torch.cuda.stream(self.stream): + if self.delay_predict_cycles: + torch.cuda._sleep(self.delay_predict_cycles) + logits = gate.forward(x) + lookahead_select(logits, resident, self.sel[target], self.res[target], k=self.k, budget=budget) + if self.sel_override is not None: + self.sel[target].copy_(self.sel_override[target]) + selected.record(self.stream) + self._inflight[target] = x + self._installed[target] = False + return True + + def install(self, target: int | None) -> int | None: + """On mode, right after the previous layer's ensure: wait for ``target``'s prediction and + install its first budget non-resident candidates with prefetch_ensure on the current + (compute) stream. Returns ``target`` when its copy is due (``issue_copy``), else None.""" + if target is None or self.mode != "on" or self._inflight[target] is None: + return None + torch.cuda.current_stream(self.device).wait_event(self._events[target][1]) + self.cache.prefetch_ensure( + target, self.sel[target], self.pf_slots[target], self.pf_src[target], + self.pf_num[target], self.stats[target], self.pf_ready[target], budget=self._budget[target], + ) + self._installed[target] = True + return target + + def _pending(self, layer_id: int) -> bool: + """On mode: ``layer_id``'s prediction was installed and its copy is not joined yet.""" + return self.mode == "on" and self._installed[layer_id] and self._inflight[layer_id] is not None + + def count_args(self, layer_id: int) -> tuple | None: + """On mode: what ``layer_id``'s demand ensure counts into, or None when it was not prefetched.""" + if not self._pending(layer_id): + return None + return self.stats[layer_id], self.pf_ready[layer_id], self.pf_num[layer_id], self._inflight[layer_id].shape[0] + + def plan_of(self, layer_id: int) -> tuple[torch.Tensor, torch.Tensor] | None: + """On mode: ``layer_id``'s prefetch copy plan (slots, count) while its prefetch is in flight.""" + if not self._pending(layer_id): + return None + return self.pf_slots[layer_id], self.pf_num[layer_id] + + def join_copy(self, layer_id: int) -> None: + """On mode: the compute stream waits for ``layer_id``'s prefetch copy before its own copy, so + a prefetched slot ensure(L) picked as a demand victim is never written by two copies.""" + if self.mode != "on" or self._inflight[layer_id] is None: + return + # a prediction that was never installed still has to rejoin (a capture needs every fork joined) + done = self._copy_events[layer_id][1] if self._pending(layer_id) else self._events[layer_id][1] + torch.cuda.current_stream(self.device).wait_event(done) + self._inflight[layer_id] = None + self._installed[layer_id] = False + + def issue_copy(self, target: int) -> None: + """On mode: copy ``target``'s prefetch plan on the copy stream after the work already on the + current stream: its prefetch_ensure and the demand copy (which owns the host link first).""" + if not self._pending(target): + return + dcopy, pfcopy = self._copy_events[target] + dcopy.record(torch.cuda.current_stream(self.device)) + self.copy_stream.wait_event(dcopy) + with torch.cuda.stream(self.copy_stream): + if self.delay_copy_cycles: + torch.cuda._sleep(self.delay_copy_cycles) + self.cache.copy_rows(target, self.pf_slots[target], self.pf_src[target], self.pf_num[target], slim=True, + kind=PREFETCH_COPY) + _mark_ready_kernel[(1,)](self.pf_ready[target]) + pfcopy.record(self.copy_stream) + + def join_and_count(self, layer_id: int, slots: torch.Tensor, id_of_slot: torch.Tensor, misses: torch.Tensor) -> None: + """Measure mode: join ``layer_id``'s prediction (if one was forked this step) and count it + against the routed slots ``slots`` of this layer's ensure.""" + x = self._inflight[layer_id] if self.mode == "measure" else None + if x is None: + return + torch.cuda.current_stream(self.device).wait_event(self._events[layer_id][1]) + prefetch_count( + slots.view(-1), id_of_slot, misses, self.sel[layer_id], self.res[layer_id], self.counters[layer_id], + id_base=layer_id * self.num_experts, rows=x.shape[0], + ) + self._inflight[layer_id] = None + + def reset_counters(self) -> None: + self.stats.zero_() + + def reset(self) -> None: + """Counters and in-flight state; the streams, events and buffers stay.""" + self.stats.zero_() + self.totals.zero_() + self._inflight = [None] * self.num_layers + self._installed = [False] * self.num_layers + + def take_window(self) -> torch.Tensor: + """This window's ``[num_layers, NUM_STAT_COLS]`` stats on the host (one sync), added to + ``totals``; the device ones restart.""" + window = self.stats.cpu() + self.stats.zero_() + self.totals += window + return window + + +def summarize(counters: torch.Tensor, mode: str = "measure") -> dict: + """Aggregate ``[num_layers, NUM_COLS or NUM_STAT_COLS]`` counters into per-layer-call and + per-token rates. In on mode the misses are what was left after the prefetch, so coverage is + useful / (useful + misses), and late is the share of copied calls whose copy had not landed + when the compute stream reached it.""" + total = counters.sum(0).tolist() + calls, issued, useful = total[CALLS], total[ISSUED], total[USEFUL] + tokens = int(counters[:, ROWS].max()) if counters.numel() else 0 + wanted = useful + total[MISSES] if mode == "on" else total[MISSES] + stats = { + "layer_calls": calls, + "tokens": tokens, + "issued_per_layer": issued / calls if calls else 0.0, + "useful_per_layer": useful / calls if calls else 0.0, + "resident_hits_per_layer": total[RESIDENT_HITS] / calls if calls else 0.0, + "misses_per_layer": total[MISSES] / calls if calls else 0.0, + "precision": useful / issued if issued else 0.0, + "coverage": useful / wanted if wanted else 0.0, + "useful_per_token": useful / tokens if tokens else 0.0, + } + if mode == "on": + copied = total[COPIED] if len(total) > COPIED else 0 + stats["copied_calls"] = copied + stats["late_fraction"] = total[LATE] / copied if copied else 0.0 + return stats + + +def format_summary(stats: dict) -> str: + if "late_fraction" in stats: + return ( + f"issued/layer={stats['issued_per_layer']:.2f}, useful/layer={stats['useful_per_layer']:.2f}, " + f"precision={stats['precision']:.3f}, useful/token={stats['useful_per_token']:.1f}, " + f"misses/layer={stats['misses_per_layer']:.2f}, coverage={stats['coverage']:.3f}, " + f"late={stats['late_fraction']:.3f} of {stats['copied_calls']} copies" + ) + return ( + f"issued/layer={stats['issued_per_layer']:.2f}, useful/layer={stats['useful_per_layer']:.2f}, " + f"precision={stats['precision']:.3f}, useful/token={stats['useful_per_token']:.1f}, " + f"misses/layer={stats['misses_per_layer']:.2f}, coverage={stats['coverage']:.3f}, " + f"resident_hits/layer={stats['resident_hits_per_layer']:.2f}" + ) + + +def format_per_layer(counters: torch.Tensor) -> str: + parts = [] + for layer, row in enumerate(counters.tolist()): + if row[CALLS]: + c = row[CALLS] + parts.append(f"L{layer}={row[USEFUL] / c:.2f}/{row[ISSUED] / c:.2f}/{row[MISSES] / c:.2f}") + return ", ".join(parts) diff --git a/python/freetoken/moe/scored_ensure.py b/python/freetoken/moe/scored_ensure.py new file mode 100644 index 0000000000..ba570b5169 --- /dev/null +++ b/python/freetoken/moe/scored_ensure.py @@ -0,0 +1,533 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright 2026 The FlashLib Authors +# Adapted from flashlib 0.3.0 (kernels/slot_cache/triton/lru_ensure.py: _phase1, _stats, +# _lru_ensure_kernel, _num_warps_for), https://github.com/FlashML-org/flashlib +"""Scored-victim slot-cache admission for the MoE expert cache (``--moe-cache-policy``). + +flashlib's sequential-argmin ``lru_ensure`` with the victim key made pluggable. POLICY 0 is +LRU and matches flashlib bit for bit; the others evict the lowest packed score key: + +- 1 ``kd``: ``-(k + beta * d / L)``. ``k`` is the tokens the expert has been passed over + since its last use, ``d`` the layer-steps until its layer comes up again. +- 2 ``kdfb``: kd plus ``w * log2(decayed use count)`` with a ``halflife``-token decay. +- 3 ``rule``: kdfb plus the near-miss refresh: this layer's experts whose router logit is + within a margin of the top-k-th count as used one token ago (``last_tok``), unpinned. + +The per-(layer, expert) state survives eviction; slot-indexed mirrors of it (tagged with the +owning id) let the victim scan read it coalesced instead of gathering by id. Scores and +counts are int32 Q16 fixed point and the count update ``log2(2^x + 1)`` reads a host-built +table, so a CPU reference can reproduce every victim exactly (``tl.exp2``/``tl.log2`` are +approximate). +""" +from __future__ import annotations + +import math + +import torch +import triton +import triton.language as tl + +POLICY_IDS = {"lru": 0, "kd": 1, "kdfb": 2, "rule": 3} + +# Replay-tuned on sep5_s01 at 1650 slots; the plateau is flat for beta 1-8, w 2-4, halflife 32-128. +BETA = 4.0 +W = 3.0 +HALFLIFE = 64 + +Q = 16 # fraction bits of scores and counts +LAST_TOK_NEVER = -1_000_000 +LC_NEVER = -(1 << 30) +LC_FLOOR = -(30 << Q) +K_MAX = 1 << 14 # tokens; with the caps on beta and w below, a score fits int32 +MAX_BETA, MAX_W = 64.0, 16.0 +DECAY_SHIFT = 8 +SCORE_BIAS = 1 << 31 +MAX_SLOT_BITS = 30 +# g(u) = log2(1 + 2^-u) sampled every 2^-G_STEP_BITS for u in [0, G_RANGE]; zero beyond +G_STEP_BITS = 6 +G_RANGE = 24 +G_LAST = G_RANGE << G_STEP_BITS +# jit code may only read globals that are constexpr +_Q = tl.constexpr(Q) +_LC_NEVER = tl.constexpr(LC_NEVER) +_LC_FLOOR = tl.constexpr(LC_FLOOR) +_K_MAX = tl.constexpr(K_MAX) +_DECAY_SHIFT = tl.constexpr(DECAY_SHIFT) +_SCORE_BIAS = tl.constexpr(SCORE_BIAS) +_G_SHIFT = tl.constexpr(Q - G_STEP_BITS) +_G_LAST = tl.constexpr(G_LAST) + + +def softplus2_table(device=None) -> torch.Tensor: + """Q16 ``log2(1 + 2^-u)`` at ``u = i / 64``, two entries past ``G_LAST`` for the interpolation.""" + step = 1 << G_STEP_BITS + vals = [round(math.log2(1.0 + 2.0 ** (-i / step)) * (1 << Q)) for i in range(G_LAST + 2)] + return torch.tensor(vals, dtype=torch.int32, device=device) + + +def score_params(beta: float, w: float, halflife: int, num_layers: int) -> tuple[int, int, int, int]: + """``(BETA_STEP, W_Q4, DECAY_MUL, DT_MAX)`` shared by the kernel and the CPU reference. + + ``beta * d / L`` is ``d * BETA_STEP`` in Q16; the decay over ``dt`` layer-steps is + ``(dt * DECAY_MUL) >> DECAY_SHIFT``, with ``dt`` capped at 64 half-lives (past the -30 floor). + """ + assert 0 <= beta <= MAX_BETA and 0 <= w <= MAX_W and halflife >= 1 + span = num_layers * int(halflife) + decay_mul = round((1 << (Q + DECAY_SHIFT)) / span) + dt_max = 64 * span + assert dt_max * decay_mul < (1 << 31) + return round(beta * (1 << Q) / num_layers), round(w * 16), decay_mul, dt_max + + +@triton.jit +def _phase1(query_ptr, slot_of_id_ptr, lru_usage_ptr, num_copy_ptr, step, K, + BLOCK_K: tl.constexpr, id_base): + """Vendored unchanged: dedup the query, split hit/miss, rank the misses, bump the hits.""" + k = tl.arange(0, BLOCK_K) + kmask = k < K + q = tl.load(query_ptr + k, mask=kmask, other=-1) + id_base + s = tl.load(slot_of_id_ptr + q, mask=kmask, other=-1) + hit = kmask & (s >= 0) + miss = kmask & (s == -1) + same = ( + (q[:, None] == q[None, :]) + & (k[:, None] > k[None, :]) + & kmask[:, None] + & kmask[None, :] + ) + first = kmask & (tl.sum(same.to(tl.int32), axis=1) == 0) + first_miss = miss & first + smaller = (q[None, :] < q[:, None]) & first_miss[None, :] + rank = tl.sum(smaller.to(tl.int32), axis=1) + num_missing = tl.sum(first_miss.to(tl.int32)) + tl.store(num_copy_ptr, num_missing.to(tl.int64)) + # Duplicated hits write the same value to the same slot -- idempotent. + tl.store(lru_usage_ptr + s, step, mask=hit) + return q, kmask, miss, first, first_miss, rank, num_missing, tl.where(hit, s, -1) + + +@triton.jit +def _prefetch_phase1(query_ptr, slot_of_id_ptr, K, BLOCK_K: tl.constexpr, id_base): + """_phase1 for a prefetch: -1 entries are padding, hits are left alone (a prefetch is not an + access), and the misses rank in query order so the select's best candidates install first.""" + k = tl.arange(0, BLOCK_K) + raw = tl.load(query_ptr + k, mask=k < K, other=-1) + kmask = raw >= 0 + q = tl.where(kmask, raw + id_base, 0) + s = tl.load(slot_of_id_ptr + q, mask=kmask, other=0) + miss = kmask & (s == -1) + same = (q[:, None] == q[None, :]) & (k[:, None] > k[None, :]) & kmask[:, None] & kmask[None, :] + first_miss = miss & (tl.sum(same.to(tl.int32), axis=1) == 0) + rank = tl.cumsum(first_miss.to(tl.int32), axis=0) - 1 + return q, first_miss, rank, tl.sum(first_miss.to(tl.int32)) + + +@triton.jit +def _stats(stats_ptr, first, num_missing): + # One vectorized atomic over 3 lanes, not three scalar ones: a scalar atomic serializes the CTA. + si = tl.arange(0, 4) + v = tl.where(si == 0, tl.sum(first.to(tl.int32)), tl.where(si == 1, num_missing, 1)) + # relaxed: only later kernels read the counters; flashlib's acq_rel waits on this call's stores (~0.3 us) + tl.atomic_add(stats_ptr + si, v.to(tl.int64), mask=si < 3, sem="relaxed") + + +@triton.jit +def _decayed(lc, ct, now, DECAY_MUL: tl.constexpr, DT_MAX: tl.constexpr): + """int32 Q16 log2 count decayed from ``ct`` to ``now`` (int64 layer-steps).""" + dt = tl.minimum(tl.maximum(now - ct, 0), DT_MAX).to(tl.int32) + return lc - ((dt * DECAY_MUL) >> _DECAY_SHIFT) + + +@triton.jit +def _softplus2(x, g_ptr, mask): + """Q16 ``log2(2^x + 1)`` = ``max(x, 0) + g(|x|)``, g linearly interpolated from the table.""" + u = tl.abs(x) + i = tl.minimum(u >> _G_SHIFT, _G_LAST) + frac = u & ((1 << _G_SHIFT) - 1) + g0 = tl.load(g_ptr + i, mask=mask, other=0) + g1 = tl.load(g_ptr + i + 1, mask=mask, other=0) + g = tl.where(i < _G_LAST, g0 + (((g1 - g0) * frac) >> _G_SHIFT), 0) + return tl.maximum(x, 0) + g + + +@triton.jit +def _owner_state(mirror_ptr, state_ptr, c, oid, mirrored, stale, other): + """The slot owner's state: its coalesced mirror, or a gather by id where the mirror is stale.""" + return tl.where( + mirrored, + tl.load(mirror_ptr + c, mask=mirrored, other=other), + tl.load(state_ptr + oid, mask=stale, other=other), + ) + + +@triton.jit(do_not_specialize=["K", "num_cached", "id_base", "nm_topk", "nm_stride", "pf_rows", "pf_budget"]) +def _scored_ensure_kernel( + query_ptr, slot_of_id_ptr, id_of_slot_ptr, lru_usage_ptr, lru_step_ptr, + out_ptr, src_ptr, dst_ptr, num_copy_ptr, stats_ptr, + tok_ptr, last_tok_ptr, lc_ptr, ct_ptr, g_ptr, logits_ptr, + owner_ptr, slot_last_ptr, slot_lc_ptr, slot_ct_ptr, pin_ptr, + pf_stats_ptr, pf_ready_ptr, pf_num_ptr, + K, num_cached, id_base, nm_topk, nm_stride, nm_thr, pf_rows, pf_budget, + BLOCK_K: tl.constexpr, BLOCK_C: tl.constexpr, BLOCK_E: tl.constexpr, BLOCK_TOPK: tl.constexpr, + NM_ROWS: tl.constexpr, + USAGE_MAX: tl.constexpr, COLLECT_STATS: tl.constexpr, + POLICY: tl.constexpr, BUMP_TOK: tl.constexpr, UPDATE_STATE: tl.constexpr, PIN_SINCE: tl.constexpr, + NUM_LAYERS: tl.constexpr, NUM_EXPERTS: tl.constexpr, + BETA_STEP: tl.constexpr, W_Q4: tl.constexpr, DECAY_MUL: tl.constexpr, DT_MAX: tl.constexpr, + SLOT_BITS: tl.constexpr, + LOWPRI: tl.constexpr, PF_COUNT: tl.constexpr, PREFETCH: tl.constexpr, +): + if PREFETCH: + # not an access: the clock stays, so usage == step still pins the last demand call's slots + # (at step 0, right after a reset, every slot reads pinned: nothing installs before a demand call) + step = tl.load(lru_step_ptr) + tl.store(num_copy_ptr, tl.zeros([], tl.int64)) + tl.store(pf_ready_ptr, 0) + else: + step = tl.load(lru_step_ptr) + 1 + tl.store(lru_step_ptr, step) + if POLICY != 0: + tok = tl.load(tok_ptr) + if BUMP_TOK: + tok = tok + 1 + tl.store(tok_ptr, tok) + layer = id_base // NUM_EXPERTS + if PREFETCH: + # score from where the pass is, right after layer - 1's demand call: the target layer's + # residents are next up, not a whole pass away (else it evicts what the layer is about to route) + layer = layer - 1 + now = tok * NUM_LAYERS + layer + if POLICY == 3: + if UPDATE_STATE: + # Routed ids qualify too; they are pinned now and set to tok below, so that is harmless. + ex = tl.arange(0, BLOCK_E) + j = tl.arange(0, BLOCK_TOPK) + prev = tl.load(last_tok_ptr + id_base + ex, mask=ex < NUM_EXPERTS, other=0) + held_at = tl.load(slot_of_id_ptr + id_base + ex, mask=ex < NUM_EXPERTS, other=-1) + near = ex < 0 + for r in tl.static_range(NM_ROWS): + routed = tl.load(query_ptr + r * nm_topk + j, mask=j < nm_topk, other=-1) + row_ptr = logits_ptr + r * nm_stride + kth = tl.min(tl.load(row_ptr + routed, mask=routed >= 0, other=float("inf")).to(tl.float32), axis=0) + row = tl.load(row_ptr + ex, mask=ex < NUM_EXPERTS, other=float("-inf")).to(tl.float32) + near = near | (row >= kth - nm_thr) + refreshed = tl.maximum(prev, tok - 1) + tl.store(last_tok_ptr + id_base + ex, refreshed, mask=near) + tl.store(slot_last_ptr + held_at, refreshed, mask=near & (held_at >= 0)) + # the victim scan and the routed update below read and rewrite what this wrote + tl.debug_barrier() + if PREFETCH: + q, first_miss, rank, num_missing = _prefetch_phase1(query_ptr, slot_of_id_ptr, K, BLOCK_K, id_base) + else: + if PF_COUNT: + kk = tl.arange(0, BLOCK_K) + sk = tl.load(slot_of_id_ptr + tl.load(query_ptr + kk, mask=kk < K, other=0) + id_base, mask=kk < K, other=-1) + # a hit on a held slot with usage 0 consumes a prefetch; read it before _phase1 bumps it + was_pf = (sk >= 0) & (tl.load(lru_usage_ptr + sk, mask=sk >= 0, other=1) == 0) + tl.debug_barrier() + q, kmask, miss, first, first_miss, rank, num_missing, out = _phase1( + query_ptr, slot_of_id_ptr, lru_usage_ptr, num_copy_ptr, step, K, BLOCK_K, + id_base) + + if num_missing > 0: + # REQUIRED (flashlib): the reload below must see _phase1's hit bump, or a hit is evicted. + tl.debug_barrier() + c = tl.arange(0, BLOCK_C) + cmask = c < num_cached + u = tl.load(lru_usage_ptr + c, mask=cmask, other=USAGE_MAX) + if POLICY == 0: + umax = tl.full([BLOCK_C], USAGE_MAX, u.dtype) + if LOWPRI: + held0 = cmask & (tl.load(id_of_slot_ptr + c, mask=cmask, other=-1) >= 0) + # empty slots, then unconsumed prefetches (held, usage 0), then residents by recency + u = tl.where((u == step) | (~cmask), umax, tl.where(held0, u + 1, 0)) + else: + u = tl.where((u == step) | (~cmask), umax, u) + n_free = tl.sum((u != umax).to(tl.int32)) + else: + oid = tl.load(id_of_slot_ptr + c, mask=cmask, other=-1) + held = cmask & (oid >= 0) + owner = tl.load(owner_ptr + c, mask=cmask, other=-1) + mirrored = held & (owner == oid) + stale = held & (owner != oid) + lk = oid // NUM_EXPERTS + ahead = lk > layer + d = tl.where(ahead, lk - layer, NUM_LAYERS - layer + lk) + last = _owner_state(slot_last_ptr, last_tok_ptr, c, oid, mirrored, stale, 0) + k = tl.minimum(tl.maximum(tl.where(ahead, tok - 1, tok) - last, 0), _K_MAX).to(tl.int32) + # int32 throughout: |score| < 2^31 by the caps on k, beta and w + score = -((k << _Q) + d * BETA_STEP) + if POLICY >= 2: + lc = _owner_state(slot_lc_ptr, lc_ptr, c, oid, mirrored, stale, _LC_NEVER) + ct = _owner_state(slot_ct_ptr, ct_ptr, c, oid, mirrored, stale, 0) + lcv = _decayed(lc, ct, now, DECAY_MUL, DT_MAX) + lcv = tl.where(lc == _LC_NEVER, _LC_FLOOR, tl.maximum(lcv, _LC_FLOOR)) + score += (W_Q4 * lcv) >> 4 + # Every key is distinct and ties go to the lowest slot; an empty slot keys below any held one. + key = tl.where(held, ((score.to(tl.int64) + _SCORE_BIAS) << SLOT_BITS) | c, c.to(tl.int64)) + if LOWPRI: + # an unconsumed prefetch (held, usage 0) keys above every empty slot and below every + # resident: the caps keep score + 2^31 > 2^30 - 2^25 - 2^22, far above this band's 1 + key = tl.where(held & (u == 0), (1 << SLOT_BITS) | c.to(tl.int64), key) + evictable = cmask & (u != step) + if PIN_SINCE: + # the score ignores recency, so slots an earlier call touched since *pin_ptr stay pinned too + evictable = evictable & (u <= tl.load(pin_ptr)) + # the pinned key's slot bits are 0, so an all-pinned scan falls back to slot 0 like flashlib's argmin + pinned_key = 0x7FFFFFFFFFFFFFFF ^ ((1 << SLOT_BITS) - 1) + key = tl.where(evictable, key, pinned_key) + n_free = tl.sum(evictable.to(tl.int32)) + if PREFETCH: + # never the all-pinned fallback: that slot may be under the running GEMM; misses rank in + # query order, so the budget keeps the first non-resident candidates + n_iter = tl.minimum(tl.minimum(num_missing, n_free), pf_budget) + tl.store(num_copy_ptr, n_iter.to(tl.int64)) + tl.store(pf_stats_ptr, tl.load(pf_stats_ptr) + n_iter) + else: + n_iter = num_missing + for i in tl.range(n_iter): + if POLICY == 0: + victim = tl.argmin(u, axis=0).to(tl.int32) + else: + victim = (tl.min(key, axis=0) & ((1 << SLOT_BITS) - 1)).to(tl.int32) + # Scalar load: victims are distinct, so no earlier iteration wrote this slot. + old = tl.load(id_of_slot_ptr + victim) + if old >= 0: + tl.store(slot_of_id_ptr + old, -1) + e = tl.sum(tl.where((rank == i) & first_miss, q, 0)) + tl.store(id_of_slot_ptr + victim, e) + tl.store(slot_of_id_ptr + e, victim) + if PREFETCH: + # usage 0 on a held slot marks it low priority until a demand call touches it + tl.store(lru_usage_ptr + victim, 0) + if POLICY != 0: + tl.store(owner_ptr + victim, -1) + else: + tl.store(lru_usage_ptr + victim, step) + tl.store(dst_ptr + i, victim) + tl.store(src_ptr + i, e - id_base) # back to the caller's id space + if not PREFETCH: + out = tl.where((rank == i) & miss, victim, out) + if POLICY == 0: + u = tl.where(c == victim, umax, u) # claim in-register + else: + key = tl.where(c == victim, pinned_key, key) + + if not PREFETCH: + # Written from registers, never re-read from slot_of_id, so out_ptr may alias query_ptr. + tl.store(out_ptr + tl.arange(0, BLOCK_K), out, mask=kmask) + if PF_COUNT: + # one writer per column: pf_ensure adds ISSUED (column 0), earlier on this stream + pn = tl.load(pf_num_ptr) + late = (pn > 0) & (tl.load(pf_ready_ptr, volatile=True) == 0) + col = tl.arange(0, 8) + add = tl.where(col == 1, tl.sum((first & was_pf).to(tl.int32)).to(tl.int64), 0) + add = tl.where(col == 3, 1, add) + add = tl.where(col == 4, pf_rows, add) + add = tl.where(col == 5, num_missing.to(tl.int64), add) + add = tl.where(col == 6, late.to(tl.int64), add) + add = tl.where(col == 7, (pn > 0).to(tl.int64), add) + cols = (col == 1) | (col >= 3) + tl.store(pf_stats_ptr + col, tl.load(pf_stats_ptr + col, mask=cols, other=0) + add, mask=cols) + if POLICY != 0: + if UPDATE_STATE: + # distinct ids only, so a duplicate at bs > 1 counts once per call; out is each id's final slot + tl.store(last_tok_ptr + q, tok, mask=first) + tl.store(owner_ptr + out, q, mask=first) + tl.store(slot_last_ptr + out, tok, mask=first) + if POLICY >= 2: + lc = tl.load(lc_ptr + q, mask=first, other=_LC_NEVER) + x = _decayed(lc, tl.load(ct_ptr + q, mask=first, other=0), now, DECAY_MUL, DT_MAX) + lc_new = tl.where(lc == _LC_NEVER, 0, _softplus2(x, g_ptr, first)) + tl.store(lc_ptr + q, lc_new, mask=first) + tl.store(ct_ptr + q, now, mask=first) + tl.store(slot_lc_ptr + out, lc_new, mask=first) + tl.store(slot_ct_ptr + out, now, mask=first) + elif not PREFETCH: + # a slot filled without a state update reads the per-id state until a decode hit re-mirrors it + tl.store(owner_ptr + out, -1, mask=first_miss) + if COLLECT_STATS: + _stats(stats_ptr, first, num_missing) + + +def _num_warps_for(block_c: int) -> int: + # flashlib's rule: 8 warps from 2048 slots, more only past ~32 elements per thread. + warps = 8 if block_c >= 2048 else 4 + return max(warps, min(triton.next_power_of_2(max(block_c // 1024, 1)), 32)) + + +def scored_ensure( + query: torch.Tensor, + slot_of_id: torch.Tensor, + id_of_slot: torch.Tensor, + lru_usage: torch.Tensor, + lru_step: torch.Tensor, + out_indices: torch.Tensor, + src_indices: torch.Tensor, + dst_indices: torch.Tensor, + num_copy: torch.Tensor, + *, + policy: int, + num_layers: int, + num_experts: int, + tok: torch.Tensor | None = None, + last_tok: torch.Tensor | None = None, + lc: torch.Tensor | None = None, + ct: torch.Tensor | None = None, + g_table: torch.Tensor | None = None, + slot_owner: torch.Tensor | None = None, + slot_last_tok: torch.Tensor | None = None, + slot_lc: torch.Tensor | None = None, + slot_ct: torch.Tensor | None = None, + router_logits: torch.Tensor | None = None, + near_miss_thr: float = 0.0, + stats: torch.Tensor | None = None, + id_base: int = 0, + bump_tok: bool = False, + update_state: bool = True, + pin_since: torch.Tensor | None = None, + beta: float = BETA, + w: float = W, + halflife: int = HALFLIFE, + lowpri: bool = False, + pf_stats: torch.Tensor | None = None, + pf_ready: torch.Tensor | None = None, + pf_num: torch.Tensor | None = None, + pf_rows: int = 0, + pf_budget: int | None = None, + prefetch: bool = False, +) -> None: + """``flashlib.lru_ensure`` (sequential strategy) with the victim picked by ``policy``. + + Same contract as ``lru_ensure``; ``policy`` > 0 also reads and updates the per-(layer, + expert) state ``last_tok``/``lc``/``ct`` (flat ``layer * num_experts + expert``) against + the device token counter ``tok``, and keeps its ``[num_cached]`` mirrors ``slot_last_tok``/ + ``slot_lc``/``slot_ct`` valid where ``slot_owner`` equals ``id_of_slot`` (-1 marks a stale + mirror). ``bump_tok`` advances the counter first (once per decode step); + ``update_state=False`` leaves the state untouched (small-prefill ensures). + ``pin_since`` (a ``lru_step`` value) also pins every slot touched after it, so a small + prefill's chunked calls cannot evict each other's experts; policy > 0 only. + Policy 3 refreshes the experts whose ``router_logits`` (``[rows, num_experts]``, rows of + ``query``) are within ``near_miss_thr`` of their row's lowest routed logit. + + FREETOKEN_MOE_PREFETCH=on: ``lowpri`` keys held slots with usage 0 (unconsumed prefetches) + between the empty slots and every resident. ``pf_stats`` (a ``[8]`` int64 prefetch stats row) + also counts this demand call: hits on such slots, the call, ``pf_rows``, the misses, and whether + the prefetch plan ``pf_num`` had rows whose copy had not set ``pf_ready`` yet. + ``prefetch=True`` is :func:`prefetch_ensure`'s install mode instead, installing at most + ``pf_budget`` ids (default: all). + """ + if prefetch: + assert lowpri and pf_stats is not None and pf_ready is not None and stats is None + assert not update_state and not bump_tok and pin_since is None and router_logits is None + if pf_stats is not None: + assert lowpri and pf_stats.dtype == torch.int64 and pf_stats.numel() == 8 and pf_stats.is_contiguous() + assert pf_ready is not None and pf_ready.dtype == torch.int32 and pf_ready.numel() == 1 + assert prefetch or (pf_num is not None and pf_num.dtype == torch.int64 and pf_num.numel() == 1) + k = query.numel() + num_cached = id_of_slot.numel() + assert query.dtype == torch.int32 and query.is_contiguous() + assert slot_of_id.dtype == torch.int32 and id_of_slot.dtype == torch.int32 + assert out_indices.dtype == torch.int32 and out_indices.numel() == k + assert lru_usage.dtype == lru_step.dtype and lru_usage.numel() == num_cached + plan = min(k, num_cached) + assert src_indices.numel() >= plan and dst_indices.numel() >= plan + block_c = triton.next_power_of_2(num_cached) + slot_bits = max(block_c.bit_length() - 1, 1) + beta_step, w_q4, decay_mul, dt_max = score_params(beta, w, halflife, num_layers) + if policy != 0: + assert slot_bits <= MAX_SLOT_BITS, f"{num_cached} slots overflow the packed score key" + assert tok is not None and last_tok is not None + assert last_tok.dtype == torch.int64 and last_tok.numel() == num_layers * num_experts + assert slot_owner is not None and slot_last_tok is not None + assert slot_owner.dtype == torch.int32 and slot_owner.numel() == num_cached + assert slot_last_tok.dtype == torch.int64 and slot_last_tok.numel() == num_cached + if pin_since is not None: + assert policy != 0 and pin_since.dtype == lru_step.dtype and pin_since.numel() == 1 + if policy >= 2: + assert lc is not None and ct is not None and g_table is not None + assert lc.dtype == torch.int32 and ct.dtype == torch.int64 + assert slot_lc is not None and slot_lc.dtype == torch.int32 and slot_lc.numel() == num_cached + assert slot_ct is not None and slot_ct.dtype == torch.int64 and slot_ct.numel() == num_cached + nm_rows = nm_topk = nm_stride = 0 + if policy == 3 and router_logits is not None: + assert router_logits.ndim == 2 and router_logits.shape[1] == num_experts + assert router_logits.stride(1) == 1 and k % router_logits.shape[0] == 0 + nm_rows, nm_stride = router_logits.shape[0], router_logits.stride(0) + nm_topk = k // nm_rows + # unused pointers still need a tensor argument + dummy = lru_step + _scored_ensure_kernel[(1,)]( + query, slot_of_id, id_of_slot, lru_usage, lru_step, + out_indices, src_indices, dst_indices, num_copy, stats, + dummy if tok is None else tok, + dummy if last_tok is None else last_tok, + dummy if lc is None else lc, + dummy if ct is None else ct, + dummy if g_table is None else g_table, + dummy if nm_rows == 0 else router_logits, + dummy if slot_owner is None else slot_owner, + dummy if slot_last_tok is None else slot_last_tok, + dummy if slot_lc is None else slot_lc, + dummy if slot_ct is None else slot_ct, + dummy if pin_since is None else pin_since, + dummy if pf_stats is None else pf_stats, + dummy if pf_ready is None else pf_ready, + dummy if pf_num is None else pf_num, + k, num_cached, id_base, nm_topk, nm_stride, float(near_miss_thr), int(pf_rows), + k if pf_budget is None else int(pf_budget), + BLOCK_K=triton.next_power_of_2(k), + BLOCK_C=block_c, + BLOCK_E=triton.next_power_of_2(num_experts), + BLOCK_TOPK=triton.next_power_of_2(max(nm_topk, 1)), + NM_ROWS=nm_rows, + USAGE_MAX=torch.iinfo(lru_usage.dtype).max, + COLLECT_STATS=stats is not None, + POLICY=policy, + BUMP_TOK=bool(bump_tok) and policy != 0, + UPDATE_STATE=bool(update_state) and policy != 0, + PIN_SINCE=pin_since is not None, + NUM_LAYERS=num_layers, + NUM_EXPERTS=num_experts, + BETA_STEP=beta_step, + W_Q4=w_q4, + DECAY_MUL=decay_mul, + DT_MAX=dt_max, + SLOT_BITS=slot_bits, + LOWPRI=bool(lowpri), + PF_COUNT=pf_stats is not None and not prefetch, + PREFETCH=bool(prefetch), + num_warps=_num_warps_for(block_c), + ) + + +def prefetch_ensure( + query: torch.Tensor, + slot_of_id: torch.Tensor, + id_of_slot: torch.Tensor, + lru_usage: torch.Tensor, + lru_step: torch.Tensor, + dst_slots: torch.Tensor, + src_rows: torch.Tensor, + num_copy: torch.Tensor, + pf_stats: torch.Tensor, + pf_ready: torch.Tensor, + budget: int | None = None, + **state, +) -> None: + """Install the ids of ``query`` (-1 = padding) that are not resident as low-priority slots + (held, usage 0), at most the first ``budget`` of them in query order (default: all), and + write their copy plan to ``dst_slots``/``src_rows``/``num_copy``. + + Unlike a demand call it bumps no clock, touches no hit and no per-id state, and so keeps + every slot of the last demand call (usage == step) pinned; victims follow the demand key + (empty, then earlier unconsumed prefetches, then the coldest residents), scored as of the + previous layer's call, the position the prefetch runs at. With fewer + evictable slots than misses it installs only the first ones, in query order. Adds the + installs to ``pf_stats[0]`` and clears ``pf_ready``. + """ + scored_ensure( + query, slot_of_id, id_of_slot, lru_usage, lru_step, query, src_rows, dst_slots, num_copy, + update_state=False, lowpri=True, pf_stats=pf_stats, pf_ready=pf_ready, pf_budget=budget, prefetch=True, + **state, + ) diff --git a/python/freetoken/moe/slot_audit.py b/python/freetoken/moe/slot_audit.py new file mode 100644 index 0000000000..12d9f40ad9 --- /dev/null +++ b/python/freetoken/moe/slot_audit.py @@ -0,0 +1,688 @@ +"""Expert slot audit (``FREETOKEN_MOE_SLOT_AUDIT=N``): a debug detector for PERSISTENT slot corruption. + +Two parts, both on the device and off by default: + +- writer history: every writer of the slot maps (demand and prefetch ensures, prefill ensures, the + whole-layer materialize, prefill buffer invalidation, reset) and every writer of the slot banks + (demand, prefetch and prefill copies, the prefill double buffers) is followed on its writer's + stream by a tiny kernel keyed on the same plan (or slot range) that appends ``(seq, kind, decode + step, lru step, layer, expert row, ...)`` to a per-slot ring of the last ``DEPTH`` metadata events + and one of the last ``DEPTH`` byte events. The map writers all run on the compute stream except the + prefill buffer invalidation (on the prefill copy stream, fenced behind it); the byte writers also + run on the prefetch and prefill copy streams. ``seq`` is one global atomic counter, so the two rings + merge into one cross-stream order. Clears are recorded just before their writer, so they carry the + old owner. + Fixed shapes and no host sync: the recorders capture into the decode graph. +- audit: every N decode steps, after the forward on the compute stream, every held slot is checked: + its bytes in every bank against the host row of the (layer, expert) its ``id_of_slot`` names + (host rows gathered with the fused copy into a scratch of fewer than 64 MiB, one layer chunk at a + time), ``slot_for_id[id_of_slot[s]] == s``, and, over every id, ``id_of_slot[slot_for_id[id]] == id`` + (a stale entry routes that id to another expert's bytes). The first ``RING`` bad (slot, owner) + pairs are kept with both history rings; the engine logs them and the totals every + MOE_STATS_INTERVAL decode steps and at shutdown, with a diagnosis read from the history. +""" +from __future__ import annotations + +import math + +import torch +import triton +import triton.language as tl + +from freetoken.env import ENV +from freetoken.utils import init_logger + +logger = init_logger(__name__) + +# events kept per slot and ring (metadata / bytes) +DEPTH = 4 +# event fields +(E_SEQ, E_KIND, E_DSTEP, E_LSTEP, E_LAYER, E_ROW, E_AUX, E_OWNER, E_USAGE, E_POS) = range(10) +EVENT_FIELDS = 10 +# metadata kinds (installs name the new owner; clears the old one, recorded before the writer) +DEMAND_INSTALL, PREFETCH_INSTALL, PREFILL_INSTALL, MATERIALIZE_INSTALL = 1, 2, 3, 4 +MATERIALIZE_CLEAR, INVALIDATE, RESET = 5, 6, 7 +# byte kinds +DEMAND_COPY, PREFETCH_COPY, PREFILL_COPY, MATERIALIZE_COPY, PAGEABLE_COPY = 11, 12, 13, 14, 15 +PREFILL_BUFFER, PREFILL_SPLIT_H2D, PREFILL_HIT_D2D, OTHER_COPY = 16, 17, 18, 19 +KIND_NAMES = { + DEMAND_INSTALL: "demand install", PREFETCH_INSTALL: "prefetch install", PREFILL_INSTALL: "prefill ensure install", + MATERIALIZE_INSTALL: "materialize install", MATERIALIZE_CLEAR: "materialize clear", + INVALIDATE: "prefill buffer invalidate", RESET: "reset", + DEMAND_COPY: "demand copy", PREFETCH_COPY: "prefetch copy", PREFILL_COPY: "prefill ensure copy", + MATERIALIZE_COPY: "materialize copy", PAGEABLE_COPY: "pageable materialize copy", + PREFILL_BUFFER: "prefill buffer copy", PREFILL_SPLIT_H2D: "prefill buffer H2D (misses and small banks)", + PREFILL_HIT_D2D: "prefill buffer D2D", OTHER_COPY: "copy", +} +INSTALLS = (DEMAND_INSTALL, PREFETCH_INSTALL, PREFILL_INSTALL, MATERIALIZE_INSTALL) +# the copy each install expects next +COPY_FOR = { + DEMAND_INSTALL: DEMAND_COPY, PREFETCH_INSTALL: PREFETCH_COPY, PREFILL_INSTALL: PREFILL_COPY, + MATERIALIZE_INSTALL: MATERIALIZE_COPY, +} +# recorder row sources: the plan's src rows, slot - row_base, the range index, the old owner +ROW_PLAN, ROW_FROM_SLOT, ROW_INDEX, ROW_OLD_OWNER = range(4) +# recorder predicates: every slot, held slots, slots >= E held by the layer +PRED_ALL, PRED_HELD, PRED_LAYER_ABOVE_E = range(3) +# totals: audits, held slots seen, slots byte-checked, held slots not byte-checked, bad slots found +# (summed over audits), of them with bad bytes / a forward map miss / named by a stale slot_for_id / +# an owner out of range, slot_for_id entries past the cache, new records, repeat findings, events +# recorded, plan entries naming no slot +(T_AUDITS, T_HELD, T_BYTES_CHECKED, T_UNCHECKED, T_BAD, T_BYTES, T_FWD, T_STALE, T_OWNER_RANGE, T_MAP_RANGE, + T_NEW, T_STILL, T_EVENTS, T_BAD_PLAN) = range(14) +NUM_TOTALS = 16 +# per-slot audit flags +F_BYTES, F_FWD, F_STALE, F_OWNER_RANGE = 1, 2, 4, 8 +# record header, then the metadata ring and the byte ring as stored ([DEPTH, EVENT_FIELDS] each) +(H_AUDIT, H_DSTEP, H_LSTEP, H_SLOT, H_OWNER, H_MAPPED, H_USAGE, H_FLAGS, H_BANKS, H_FIRST_BANK, H_FIRST_BYTE, + H_BAD_WORDS, H_STALE_ID, H_STALE_N, H_META_N, H_BYTES_N) = range(16) +HEADER = 16 +RING_EVENTS = DEPTH * EVENT_FIELDS +RECORD_FIELDS = HEADER + 2 * RING_EVENTS +RING = 32 +SCRATCH_BYTES = 64 << 20 +MAX_CHUNK_ROWS = 64 +COMPARE_WORDS = 4096 +_NO_BAD = tl.constexpr(0x7FFFFFFF) +_D = tl.constexpr(DEPTH) +_F = tl.constexpr(EVENT_FIELDS) + + +def resolve_interval(interval: int | None = None) -> int: + value = ENV.MOE_SLOT_AUDIT.value if interval is None else interval + value = int(value or 0) + if value < 0: + raise ValueError(f"FREETOKEN_MOE_SLOT_AUDIT={value}: expected 0 (off) or the decode steps between audits") + return value + + +@triton.jit +def _put(vals, f, idx: tl.constexpr, v): + return tl.where(f[None, :] == idx, v.to(tl.int64)[:, None], vals) + + +@triton.jit +def _bump_step_kernel(dstep_ptr): + tl.store(dstep_ptr, tl.load(dstep_ptr) + 1) + + +@triton.jit(do_not_specialize=["kind", "layer", "start", "row_base", "width", "num_slots"]) +def _record_kernel( + hist_ptr, cur_ptr, seq_ptr, dstep_ptr, lstep_ptr, id_of_slot_ptr, usage_ptr, totals_ptr, + dst_ptr, src_ptr, aux_ptr, num_ptr, + kind, layer, start, row_base, width, num_slots, + PLAN: tl.constexpr, ROW: tl.constexpr, HAS_AUX: tl.constexpr, HAS_NUM: tl.constexpr, PRED: tl.constexpr, + E: tl.constexpr, BLOCK: tl.constexpr, +): + """Append one event per written slot to that slot's ring: plan entries ``i < min(*num, width)`` + (``PLAN``) or the slots ``start + i`` for ``i < width``.""" + i = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + n = width + if HAS_NUM: + n = tl.minimum(tl.load(num_ptr).to(tl.int32), width) + valid = i < n + if PLAN: + slot = tl.load(dst_ptr + i, mask=valid, other=-1) + else: + slot = start + i + ok = valid & (slot >= 0) & (slot < num_slots) + bad_plan = valid & ~ok + owner = tl.load(id_of_slot_ptr + slot, mask=ok, other=-3) + usage = tl.load(usage_ptr + slot, mask=ok, other=-3) + if PRED == 1: + ok = ok & (owner >= 0) + if PRED == 2: + ok = ok & (owner >= layer * E) & (owner < layer * E + E) & (slot >= E) + lay = tl.zeros([BLOCK], tl.int32) + layer + aux = tl.zeros([BLOCK], tl.int32) - 1 + if ROW == 0: + row = tl.load(src_ptr + i, mask=ok, other=-1) + elif ROW == 1: + row = slot - row_base + elif ROW == 2: + row = i + else: + row = tl.where(owner >= 0, owner % E, -1) + lay = tl.where(owner >= 0, owner // E, -1) + aux = aux + 1 + layer + if HAS_AUX: + aux = tl.load(aux_ptr + i, mask=ok, other=-1) + cnt = tl.sum(ok.to(tl.int64), axis=0) + seq = tl.atomic_add(seq_ptr, cnt) + tl.cumsum(ok.to(tl.int64), axis=0) - 1 + pos = tl.atomic_add(cur_ptr + slot, 1, mask=ok) + f = tl.arange(0, 16) + z = tl.zeros([BLOCK], tl.int64) + vals = tl.zeros([BLOCK, 16], tl.int64) - 1 + vals = _put(vals, f, 0, seq) + vals = _put(vals, f, 1, z + kind) + vals = _put(vals, f, 2, z + tl.load(dstep_ptr)) + vals = _put(vals, f, 3, z + tl.load(lstep_ptr)) + vals = _put(vals, f, 4, lay) + vals = _put(vals, f, 5, row) + vals = _put(vals, f, 6, aux) + vals = _put(vals, f, 7, owner) + vals = _put(vals, f, 8, usage) + vals = _put(vals, f, 9, i) + at = (slot.to(tl.int64) * _D + (pos % _D)) * _F + tl.store(hist_ptr + at[:, None] + f[None, :], vals, mask=ok[:, None] & (f[None, :] < _F)) + t = tl.arange(0, 2) + add = tl.where(t == 0, cnt, tl.sum(bad_plan.to(tl.int64), axis=0)) + tl.atomic_add(totals_ptr + 12 + t, add, sem="relaxed") + + +@triton.jit(do_not_specialize=["num_slots", "num_ids"]) +def _audit_meta_kernel( + id_of_slot_ptr, slot_for_id_ptr, layer_ok_ptr, count_ptr, rows_ptr, slots_ptr, + flags_ptr, stale_id_ptr, stale_n_ptr, totals_ptr, num_slots, num_ids, + E: tl.constexpr, CAP: tl.constexpr, BLOCK: tl.constexpr, +): + """Both map checks, and each byte-checkable held slot appended to its layer's list.""" + i = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + ms = i < num_slots + o = tl.load(id_of_slot_ptr + i, mask=ms, other=-1) + held = ms & (o >= 0) + in_range = held & (o < num_ids) + mapped = tl.load(slot_for_id_ptr + o, mask=in_range, other=-2) + fwd = in_range & (mapped != i) + out_of_range = held & ~in_range + fl = tl.where(fwd, 2, 0) | tl.where(out_of_range, 8, 0) + tl.atomic_or(flags_ptr + i, fl, mask=ms & (fl != 0)) + layer = tl.where(in_range, o // E, 0) + ok = in_range & (tl.load(layer_ok_ptr + layer, mask=in_range, other=0) != 0) + p = tl.atomic_add(count_ptr + layer, 1, mask=ok) + take = ok & (p < CAP) + tl.store(rows_ptr + layer * CAP + p, o - layer * E, mask=take) + tl.store(slots_ptr + layer * CAP + p, i, mask=take) + mi = i < num_ids + s = tl.load(slot_for_id_ptr + i, mask=mi, other=-1) + live = mi & (s >= 0) + s_ok = live & (s < num_slots) + owner = tl.load(id_of_slot_ptr + s, mask=s_ok, other=-1) + stale = s_ok & (owner != i) + tl.atomic_or(flags_ptr + s, 4, mask=stale) + tl.store(stale_id_ptr + s, i, mask=stale) + tl.atomic_add(stale_n_ptr + s, 1, mask=stale) + t = tl.arange(0, 16) + add = tl.where(t == 1, tl.sum(held.to(tl.int64), axis=0), 0) + add = tl.where(t == 2, tl.sum(take.to(tl.int64), axis=0), add) + add = tl.where(t == 3, tl.sum((held & ~take).to(tl.int64), axis=0), add) + add = tl.where(t == 6, tl.sum(fwd.to(tl.int64), axis=0), add) + add = tl.where(t == 7, tl.sum(stale.to(tl.int64), axis=0), add) + add = tl.where(t == 8, tl.sum(out_of_range.to(tl.int64), axis=0), add) + add = tl.where(t == 9, tl.sum((live & ~s_ok).to(tl.int64), axis=0), add) + add = tl.where((t == 0) & (tl.program_id(0) == 0), 1, add) + tl.atomic_add(totals_ptr + t, add, mask=add != 0, sem="relaxed") + + +@triton.jit +def _audit_nums_kernel(count_ptr, nums_ptr, L: tl.constexpr, C: tl.constexpr, R: tl.constexpr, BLOCK: tl.constexpr): + j = tl.arange(0, BLOCK) + m = j < L * C + cnt = tl.load(count_ptr + j // C, mask=m, other=0) + tl.store(nums_ptr + j, tl.minimum(tl.maximum(cnt - (j % C) * R, 0), R).to(tl.int64), mask=m) + + +@triton.jit +def _audit_compare_kernel(cache_ptrs, ref_ptrs, words_ptr, slots_ptr, num_ptr, bad_cnt_ptr, bad_first_ptr, + NB: tl.constexpr, BLOCK: tl.constexpr): + """Program (row, bank, chunk): one chunk of a listed slot's bank row against its gathered host row.""" + r = tl.program_id(0) + if r < tl.load(num_ptr): + b = tl.program_id(1) + words = tl.load(words_ptr + b) + off = tl.program_id(2).to(tl.int64) * BLOCK + tl.arange(0, BLOCK) + m = off < words + slot = tl.load(slots_ptr + r).to(tl.int64) + got = tl.load(tl.load(cache_ptrs + b).to(tl.pointer_type(tl.int32)) + slot * words + off, mask=m, other=0) + want = tl.load(tl.load(ref_ptrs + b).to(tl.pointer_type(tl.int32)) + r * words + off, mask=m, other=0) + bad = m & (got != want) + nbad = tl.sum(bad.to(tl.int32), axis=0) + if nbad > 0: + tl.atomic_add(bad_cnt_ptr + slot * NB + b, nbad) + tl.atomic_min(bad_first_ptr + slot * NB + b, tl.min(tl.where(bad, off, _NO_BAD), axis=0).to(tl.int32)) + + +@triton.jit(do_not_specialize=["num_slots", "num_ids"]) +def _audit_report_kernel( + id_of_slot_ptr, slot_for_id_ptr, usage_ptr, lstep_ptr, dstep_ptr, + flags_ptr, stale_id_ptr, stale_n_ptr, bad_cnt_ptr, bad_first_ptr, reported_ptr, + meta_hist_ptr, meta_cur_ptr, bytes_hist_ptr, bytes_cur_ptr, + ring_ptr, cursor_ptr, totals_ptr, count_ptr, + num_slots, num_ids, + NB: tl.constexpr, NB_P: tl.constexpr, L: tl.constexpr, L_P: tl.constexpr, RING_N: tl.constexpr, + BLOCK: tl.constexpr, +): + """Keep each new bad (slot, owner) with its history, count the findings and clear the work arrays.""" + s = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + ms = s < num_slots + b = tl.arange(0, NB_P) + cell = ms[:, None] & (b < NB)[None, :] + cnt = tl.load(bad_cnt_ptr + s[:, None] * NB + b[None, :], mask=cell, other=0) + first = tl.load(bad_first_ptr + s[:, None] * NB + b[None, :], mask=cell, other=_NO_BAD) + bank_bad = cnt > 0 + bytes_bad = tl.sum(bank_bad.to(tl.int32), axis=1) > 0 + bank_mask = tl.sum(tl.where(bank_bad, 1 << b[None, :], 0), axis=1) + first_bank = tl.min(tl.where(bank_bad, b[None, :], NB_P), axis=1) + first_word = tl.sum(tl.where(b[None, :] == first_bank[:, None], first, 0), axis=1) + fl = tl.load(flags_ptr + s, mask=ms, other=0) | tl.where(bytes_bad, 1, 0) + bad = ms & (fl != 0) + owner = tl.load(id_of_slot_ptr + s, mask=ms, other=-1) + sid = tl.load(stale_id_ptr + s, mask=ms, other=-1) + key = (owner.to(tl.int64) + 1) * (num_ids + 2) + sid + 1 + new = bad & (tl.load(reported_ptr + s, mask=ms, other=-1) != key) + tl.store(reported_ptr + s, tl.where(bad, key, -1), mask=ms) + pos = tl.atomic_add(cursor_ptr + s * 0, 1, mask=new) + keep = new & (pos < RING_N) + in_range = ms & (owner >= 0) & (owner < num_ids) + f = tl.arange(0, 16) + z = tl.zeros([BLOCK], tl.int64) + vals = tl.zeros([BLOCK, 16], tl.int64) - 1 + vals = _put(vals, f, 0, z + tl.load(totals_ptr) - 1) + vals = _put(vals, f, 1, z + tl.load(dstep_ptr)) + vals = _put(vals, f, 2, z + tl.load(lstep_ptr)) + vals = _put(vals, f, 3, s) + vals = _put(vals, f, 4, owner) + vals = _put(vals, f, 5, tl.load(slot_for_id_ptr + owner, mask=in_range, other=-2)) + vals = _put(vals, f, 6, tl.load(usage_ptr + s, mask=ms, other=-1)) + vals = _put(vals, f, 7, fl) + vals = _put(vals, f, 8, bank_mask) + vals = _put(vals, f, 9, tl.where(bytes_bad, first_bank, -1)) + vals = _put(vals, f, 10, tl.where(bytes_bad, first_word.to(tl.int64) * 4, -1)) + vals = _put(vals, f, 11, tl.sum(cnt, axis=1)) + vals = _put(vals, f, 12, sid) + vals = _put(vals, f, 13, tl.load(stale_n_ptr + s, mask=ms, other=0)) + vals = _put(vals, f, 14, tl.load(meta_cur_ptr + s, mask=ms, other=0)) + vals = _put(vals, f, 15, tl.load(bytes_cur_ptr + s, mask=ms, other=0)) + rec = ring_ptr + pos.to(tl.int64) * (16 + 2 * _D * _F) + tl.store(rec[:, None] + f[None, :], vals, mask=keep[:, None]) + j = tl.arange(0, 64) + jm = keep[:, None] & (j < _D * _F)[None, :] + src = s.to(tl.int64)[:, None] * (_D * _F) + j[None, :] + tl.store(rec[:, None] + 16 + j[None, :], tl.load(meta_hist_ptr + src, mask=jm, other=-1), mask=jm) + tl.store(rec[:, None] + 16 + _D * _F + j[None, :], tl.load(bytes_hist_ptr + src, mask=jm, other=-1), mask=jm) + t = tl.arange(0, 16) + add = tl.where(t == 4, tl.sum(bad.to(tl.int64), axis=0), 0) + add = tl.where(t == 5, tl.sum((ms & bytes_bad).to(tl.int64), axis=0), add) + add = tl.where(t == 10, tl.sum(new.to(tl.int64), axis=0), add) + add = tl.where(t == 11, tl.sum((bad & ~new).to(tl.int64), axis=0), add) + tl.atomic_add(totals_ptr + t, add, mask=add != 0, sem="relaxed") + tl.store(flags_ptr + s, tl.zeros([BLOCK], tl.int32), mask=ms) + tl.store(stale_id_ptr + s, tl.zeros([BLOCK], tl.int32) - 1, mask=ms) + tl.store(stale_n_ptr + s, tl.zeros([BLOCK], tl.int32), mask=ms) + tl.store(bad_cnt_ptr + s[:, None] * NB + b[None, :], tl.zeros([BLOCK, NB_P], tl.int32), mask=cell) + tl.store(bad_first_ptr + s[:, None] * NB + b[None, :], tl.zeros([BLOCK, NB_P], tl.int32) + _NO_BAD, mask=cell) + if tl.program_id(0) == 0: + ll = tl.arange(0, L_P) + tl.store(count_ptr + ll, tl.zeros([L_P], tl.int32), mask=ll < L) + + +class SlotAuditor: + """Per-slot writer history (recorded by ``record_*`` on the writer's stream) and the periodic + audit (``end_decode_step`` -> ``scan`` on the compute stream). ``bind`` sizes everything for a + cache's slots and banks; a rebuild rebinds and starts a new history, the totals and records stay. + Holds no reference to its cache: every call takes it.""" + + def __init__(self, num_layers: int, num_experts: int, device: torch.device, *, interval: int, + ring: int | None = None, scratch_bytes: int = SCRATCH_BYTES) -> None: + self.num_layers = num_layers + self.num_experts = num_experts + self.device = device + self.interval = interval + self.ring_n = RING if ring is None else ring + self.scratch_bytes = scratch_bytes + ring = self.ring_n + self.dstep = torch.zeros((1,), dtype=torch.int64, device=device) + self.seq = torch.zeros((1,), dtype=torch.int64, device=device) + self.totals = torch.zeros((NUM_TOTALS,), dtype=torch.int64, device=device) + self.ring = torch.full((ring, RECORD_FIELDS), -1, dtype=torch.int64, device=device) + self.cursor = torch.zeros((1,), dtype=torch.int64, device=device) + self.count = torch.zeros((num_layers,), dtype=torch.int32, device=device) + self.num_slots = 0 + self.bank_names: list[str] = [] + self._feats: list[int] = [] + self._steps = 0 + self._printed = 0 + self._last = torch.zeros((NUM_TOTALS,), dtype=torch.int64) + self._dummy = torch.zeros((1,), dtype=torch.int32, device=device) + logger.info(f"MoE slot audit: recording every slot writer, auditing every held slot each {interval} decode steps") + + # ------------------------------------------------------------------ setup + + def bind(self, cache) -> None: + """Size the history and audit arrays for ``cache``'s slots and banks (a rebuild starts a new + history) and point the byte checks at its banks.""" + caches = [bank for _, bank in cache.banks] + feats = [math.prod(bank.shape[1:]) * bank.element_size() for bank in caches] + assert all(f % 4 == 0 for f in feats), feats + if self.bank_names: + assert list(cache.bank_schema) == self.bank_names and feats == self._feats, "the bank layout changed" + else: + self.bank_names = list(cache.bank_schema) + self._feats = feats + self._alloc_scratch(cache, feats) + self.cache_ptrs = torch.tensor([bank.data_ptr() for bank in caches], dtype=torch.int64, device=self.device) + S, dev = cache.cache_size, self.device + self.num_slots = S + self.meta_hist = torch.full((S, DEPTH, EVENT_FIELDS), -1, dtype=torch.int64, device=dev) + self.bytes_hist = torch.full((S, DEPTH, EVENT_FIELDS), -1, dtype=torch.int64, device=dev) + self.meta_cur = torch.zeros((S,), dtype=torch.int32, device=dev) + self.bytes_cur = torch.zeros((S,), dtype=torch.int32, device=dev) + self.flags = torch.zeros((S,), dtype=torch.int32, device=dev) + self.stale_id = torch.full((S,), -1, dtype=torch.int32, device=dev) + self.stale_n = torch.zeros((S,), dtype=torch.int32, device=dev) + self.reported = torch.full((S,), -1, dtype=torch.int64, device=dev) + nb = len(feats) + self.bad_cnt = torch.zeros((S * nb,), dtype=torch.int32, device=dev) + self.bad_first = torch.full((S * nb,), 0x7FFFFFFF, dtype=torch.int32, device=dev) + # a correct cache holds at most num_experts slots of a layer; more are counted unchecked + per_layer = min(self.num_experts, S) + self.chunks = triton.cdiv(per_layer, self.rows) + self.cap = self.chunks * self.rows + self.list_rows = torch.zeros((self.num_layers, self.cap), dtype=torch.int32, device=dev) + self.list_slots = torch.zeros((self.num_layers, self.cap), dtype=torch.int32, device=dev) + self.nums = torch.zeros((self.num_layers * self.chunks,), dtype=torch.int64, device=dev) + # layers without a device alias (LOCKED/PAGEABLE host banks) cannot be gathered + self.layer_ok_host = [layer not in cache._unpinned_layers for layer in range(self.num_layers)] + self.layer_ok = torch.tensor(self.layer_ok_host, dtype=torch.int32, device=dev) + self.count.zero_() + + def _alloc_scratch(self, cache, feats: list[int]) -> None: + row = sum(feats) + self.rows = max(1, min(MAX_CHUNK_ROWS, (self.scratch_bytes - 1) // row)) + rows = self.rows + self.scratch = torch.empty((rows * row,), dtype=torch.uint8, device=self.device) + offsets = [rows * sum(feats[:b]) for b in range(len(feats))] + self.scratch_views = [ + self.scratch[o : o + rows * f].view(bank.dtype).view(rows, *bank.shape[1:]) + for o, f, (_, bank) in zip(offsets, feats, cache.banks) + ] + self.scratch_ptrs = torch.tensor([self.scratch.data_ptr() + o for o in offsets], dtype=torch.int64, device=self.device) + self.words = torch.tensor([f // 4 for f in feats], dtype=torch.int64, device=self.device) + self.iota = torch.arange(rows, dtype=torch.int32, device=self.device) + self._word_chunks = triton.cdiv(max(feats) // 4, COMPARE_WORDS) + logger.info(f"MoE slot audit scratch: {self.scratch.numel() / 2**20:.1f} MiB, {rows} experts a chunk") + + # ------------------------------------------------------------------ recording (graph-capturable) + + def begin_step(self) -> None: + """Advance the decode-step clock (the first GPU layer's decode ensure calls this).""" + _bump_step_kernel[(1,)](self.dstep) + + def _launch(self, cache, kind: int, layer: int, *, meta: bool, width: int, dst=None, src=None, aux=None, + num=None, start: int = 0, row: int = ROW_PLAN, row_base: int = 0, pred: int = PRED_ALL) -> None: + if width <= 0 or not self.num_slots: + return + block = min(triton.next_power_of_2(width), 128) + d = self._dummy + _record_kernel[(triton.cdiv(width, block),)]( + self.meta_hist if meta else self.bytes_hist, self.meta_cur if meta else self.bytes_cur, + self.seq, self.dstep, cache.step, cache.id_of_slot, cache.usage, self.totals, + d if dst is None else dst, d if src is None else src, d if aux is None else aux, d if num is None else num, + kind, layer, start, row_base, width, self.num_slots, + PLAN=dst is not None, ROW=row, HAS_AUX=aux is not None, HAS_NUM=num is not None, PRED=pred, + E=self.num_experts, BLOCK=block, num_warps=4, + ) + + def record_plan(self, cache, kind: int, layer: int, dst: torch.Tensor, src: torch.Tensor | None, num: torch.Tensor, + *, meta: bool, width: int | None = None, aux: torch.Tensor | None = None, row_base: int = 0) -> None: + """The writer just wrote slots ``dst[:num]`` (rows ``src[:num]``, else ``dst - row_base``).""" + width = dst.numel() if width is None else min(width, dst.numel()) + self._launch(cache, kind, layer, meta=meta, width=width, dst=dst, src=src, aux=aux, num=num, + row=ROW_PLAN if src is not None else ROW_FROM_SLOT, row_base=row_base) + + def record_range(self, cache, kind: int, layer: int, start: int, n: int, *, meta: bool) -> None: + """The writer just wrote expert ``i`` of ``layer`` into slot ``start + i`` for ``i < n``.""" + self._launch(cache, kind, layer, meta=meta, width=n, start=start, row=ROW_INDEX) + + def record_clear(self, cache, kind: int, context_layer: int, start: int, n: int, *, layer_above_e: bool = False) -> None: + """Before a writer clears held slots in ``[start, start + n)`` (with ``layer_above_e``: only + those >= E held by ``context_layer``); the event names the old owner.""" + self._launch(cache, kind, context_layer, meta=True, width=n, start=start, row=ROW_OLD_OWNER, + pred=PRED_LAYER_ABOVE_E if layer_above_e else PRED_HELD) + + # ------------------------------------------------------------------ audit + + def end_decode_step(self, cache) -> bool: + """After a decode forward, on the compute stream: audit every ``interval`` steps.""" + self._steps += 1 + if self._steps % self.interval: + return False + self.scan(cache) + return True + + def scan(self, cache) -> None: + """Audit every held slot now (enqueued on the current stream; no host sync).""" + S, num_ids = self.num_slots, self.num_layers * self.num_experts + assert S == cache.cache_size and cache.id_of_slot.numel() == S, "bind after a rebuild" + block = 1024 + _audit_meta_kernel[(triton.cdiv(max(S, num_ids), block),)]( + cache.id_of_slot, cache.slot_for_id.view(-1), self.layer_ok, self.count, self.list_rows, self.list_slots, + self.flags, self.stale_id, self.stale_n, self.totals, S, num_ids, + E=self.num_experts, CAP=self.cap, BLOCK=block, num_warps=4, + ) + n_lists = self.num_layers * self.chunks + _audit_nums_kernel[(1,)](self.count, self.nums, L=self.num_layers, C=self.chunks, R=self.rows, + BLOCK=triton.next_power_of_2(n_lists)) + nb = len(self.bank_names) + R = self.rows + for layer in range(self.num_layers): + if not self.layer_ok_host[layer]: + continue + for c in range(self.chunks): + rows = self.list_rows[layer, c * R : (c + 1) * R] + slots = self.list_slots[layer, c * R : (c + 1) * R] + num = self.nums[layer * self.chunks + c : layer * self.chunks + c + 1] + self._gather(cache, layer, rows, num) + _audit_compare_kernel[(R, nb, self._word_chunks)]( + self.cache_ptrs, self.scratch_ptrs, self.words, slots, num, self.bad_cnt, self.bad_first, + NB=nb, BLOCK=COMPARE_WORDS, num_warps=8, + ) + rblock = 256 + _audit_report_kernel[(triton.cdiv(S, rblock),)]( + cache.id_of_slot, cache.slot_for_id.view(-1), cache.usage, cache.step, self.dstep, + self.flags, self.stale_id, self.stale_n, self.bad_cnt, self.bad_first, self.reported, + self.meta_hist, self.meta_cur, self.bytes_hist, self.bytes_cur, + self.ring, self.cursor, self.totals, self.count, S, num_ids, + NB=nb, NB_P=triton.next_power_of_2(nb), L=self.num_layers, L_P=triton.next_power_of_2(self.num_layers), + RING_N=self.ring_n, BLOCK=rblock, num_warps=4, + ) + + def _gather(self, cache, layer: int, rows: torch.Tensor, num: torch.Tensor) -> None: + # the plain fused copy, not the slim one prefetch copies use: a slim-kernel bug cannot hide itself + if cache._copy_fused_ok: + from freetoken.kernel.fast_index_copy import fast_index_copy_multi_jit + + fast_index_copy_multi_jit( + self.scratch_ptrs, cache._copy_src_ptrs[layer], cache._copy_feat_bytes[cache.layer_layout[layer]], self.iota, rows, num, + ) + return + from freetoken.kernel import fast_index_copy_jit + + for (per_layer, _), view in zip(cache.banks, self.scratch_views): + fast_index_copy_jit(view, self.iota, per_layer[layer], rows, num) + + # ------------------------------------------------------------------ reporting (host syncs) + + def take_window(self) -> tuple[list[int], list[int], list[list[int]], int]: + """(totals now, totals of this window, records kept since the last call, records in all).""" + totals = self.totals.cpu() + window = (totals - self._last).tolist() + self._last = totals + cursor = int(self.cursor.item()) + upto = min(cursor, self.ring_n) + records = self.ring[self._printed : upto].cpu().tolist() if upto > self._printed else [] + self._printed = max(self._printed, upto) + return totals.tolist(), window, records, cursor + + def report_window(self, steps: int) -> list[tuple[str, str]]: + """(level, line) pairs for one MOE_STATS_INTERVAL window; nothing when no audit ran.""" + totals, window, records, cursor = self.take_window() + if not window[T_AUDITS] and not records: + return [] + bad = bool(window[T_BAD] or window[T_MAP_RANGE] or window[T_BAD_PLAN] or records) + lines = [("warning" if bad else "info", + f"MoE slot audit ({steps} decode steps, every {self.interval}): {self.format_counts(window)}; " + f"session: {totals[T_AUDITS]} audits, {totals[T_NEW]} bad (slot, owner) pairs")] + lines += [("warning", "MoE slot audit record: " + self.format_record(r)) for r in records] + if records and cursor > self.ring_n: + lines.append(("warning", f"MoE slot audit: {cursor} bad (slot, owner) pairs in all, the first {self.ring_n} are kept")) + return lines + + def report_session(self) -> list[tuple[str, str]]: + """Session totals and every kept record (shutdown).""" + totals = self.totals.cpu().tolist() + cursor = int(self.cursor.item()) + if not totals[T_AUDITS]: + return [] + bad = bool(totals[T_BAD] or totals[T_MAP_RANGE] or totals[T_BAD_PLAN]) + lines = [("warning" if bad else "info", f"MoE slot audit (session, every {self.interval} decode steps): " + f"{self.format_counts(totals)}")] + records = self.ring[: min(cursor, self.ring_n)].cpu().tolist() + lines += [("warning", "MoE slot audit record: " + self.format_record(r)) for r in records] + if cursor > self.ring_n: + lines.append(("warning", f"MoE slot audit: {cursor} bad (slot, owner) pairs in all, the first {self.ring_n} are kept")) + return lines + + @staticmethod + def format_counts(t: list[int]) -> str: + audits = t[T_AUDITS] + per = (lambda v: f"{v / audits:.0f}") if audits else str + text = ( + f"{audits} audits, {per(t[T_HELD])} held slots an audit ({per(t[T_BYTES_CHECKED])} byte-checked), " + f"bad slots={t[T_BAD]} (bytes={t[T_BYTES]}, map={t[T_FWD]}, stale slot_for_id={t[T_STALE]}, " + f"owner out of range={t[T_OWNER_RANGE]}), new={t[T_NEW]}, repeat={t[T_STILL]}, events={t[T_EVENTS]}" + ) + if t[T_UNCHECKED]: + text += f", held slots not byte-checked={t[T_UNCHECKED]}" + if t[T_MAP_RANGE]: + text += f", slot_for_id entries past the cache={t[T_MAP_RANGE]}" + if t[T_BAD_PLAN]: + text += f", plan entries naming no slot={t[T_BAD_PLAN]}" + return text + + # ------------------------------------------------------------------ record decoding + + def _name(self, flat: int) -> str: + e = self.num_experts + return f"L{flat // e}/e{flat % e}" if flat >= 0 else str(flat) + + @staticmethod + def events(rec: list[int]) -> list[dict]: + """Both rings of a record as event dicts in ``seq`` order (``ring`` 'meta' or 'bytes').""" + out = [] + for ring, base in (("meta", HEADER), ("bytes", HEADER + RING_EVENTS)): + for d in range(DEPTH): + ev = rec[base + d * EVENT_FIELDS : base + (d + 1) * EVENT_FIELDS] + if ev[E_SEQ] >= 0: + out.append({ + "ring": ring, "seq": ev[E_SEQ], "kind": ev[E_KIND], "dstep": ev[E_DSTEP], "lstep": ev[E_LSTEP], + "layer": ev[E_LAYER], "row": ev[E_ROW], "aux": ev[E_AUX], "owner": ev[E_OWNER], + "usage": ev[E_USAGE], "pos": ev[E_POS], + }) + return sorted(out, key=lambda ev: ev["seq"]) + + def _flat(self, ev: dict) -> int: + return ev["layer"] * self.num_experts + ev["row"] if ev["layer"] >= 0 and ev["row"] >= 0 else -1 + + def format_event(self, ev: dict) -> str: + e = self.num_experts + text = f"#{ev['seq']} {KIND_NAMES.get(ev['kind'], ev['kind'])} L{ev['layer']}/e{ev['row']}" + if ev["kind"] == PREFILL_HIT_D2D: + text += f" from slot {ev['aux']}" + elif ev["kind"] in (MATERIALIZE_CLEAR, INVALIDATE) and ev["aux"] >= 0: + text += f" (for layer {ev['aux']})" + seen = "slot held" if ev["ring"] == "bytes" else "then held" + if ev["kind"] in (MATERIALIZE_CLEAR, INVALIDATE, RESET): + seen = "held" + owner = f"L{ev['owner'] // e}/e{ev['owner'] % e}" if ev["owner"] >= 0 else "nothing" + return f"{text} @step {ev['dstep']}/lru {ev['lstep']} ({seen} {owner}, usage {ev['usage']})" + + def _kind(self, ev: dict) -> str: + return KIND_NAMES.get(ev["kind"], str(ev["kind"])) + + def _pending_install(self, copy: dict, meta: list[dict], byte: list[dict], before: dict) -> dict | None: + """The earlier install (before ``before``) that ``copy`` completes: an install of the id it + copied with no copy of that id between them.""" + flat = self._flat(copy) + for inst in reversed([ev for ev in meta if ev["kind"] in INSTALLS and ev["seq"] < before["seq"]]): + if self._flat(inst) == flat: + done = any(self._flat(ev) == flat and inst["seq"] < ev["seq"] < copy["seq"] for ev in byte) + return None if done else inst + return None + + def _bytes_story(self, owner: int, meta: list[dict], byte: list[dict]) -> str: + insts = [ev for ev in meta if ev["kind"] in INSTALLS and self._flat(ev) == owner] + if not insts: + return (f"no recorded install of {self._name(owner)} in the kept history " + f"(an unrecorded map writer, or it is older than the last {DEPTH} events)") + inst = insts[-1] + name = self._kind(inst) + after = [ev for ev in byte if ev["seq"] > inst["seq"]] + own = [ev for ev in after if self._flat(ev) == owner] + if own: + last = after[-1] + if self._flat(last) == owner: + return (f"{name} and {self._kind(last)} of {self._name(owner)} (#{last['seq']}), then the bytes changed " + f"with no recorded write: overwritten after the last recorded write, or that copy did not land") + text = (f"{name} and {self._kind(own[-1])} of {self._name(owner)} (#{own[-1]['seq']}), then " + f"{self._kind(last)} of {self._name(self._flat(last))} (#{last['seq']}) overwrote it") + prior = self._pending_install(last, meta, byte, inst) + if prior is not None: + text += f": the late copy of the earlier {self._kind(prior)} of {self._name(self._flat(prior))} (#{prior['seq']})" + return text + text = f"{name}, no {KIND_NAMES[COPY_FOR[inst['kind']]]}" + if not after: + return text + last = after[-1] + prior = self._pending_install(last, meta, byte, inst) + if prior is not None: + return (f"{text}; the last write, {self._kind(last)} of {self._name(self._flat(last))} (#{last['seq']}), " + f"is the copy of the earlier {self._kind(prior)} (#{prior['seq']}) and landed after this install") + where = (f"from layer {last['layer']}'s source" if last["row"] == owner % self.num_experts + else f"of {self._name(self._flat(last))}") + return f"{name} of {self._name(owner)}, then {self._kind(last)} {where} ({self._name(self._flat(last))}, #{last['seq']})" + + def diagnose(self, rec: list[int]) -> str: + """What the history says about a bad slot, in words.""" + owner, fl = rec[H_OWNER], rec[H_FLAGS] + evs = self.events(rec) + meta = [ev for ev in evs if ev["ring"] == "meta"] + byte = [ev for ev in evs if ev["ring"] == "bytes"] + parts = [] + if fl & F_OWNER_RANGE: + parts.append(f"id_of_slot={owner} names no expert") + if fl & F_BYTES and owner >= 0: + parts.append(self._bytes_story(owner, meta, byte)) + if fl & F_FWD: + parts.append(f"slot_for_id[{self._name(owner)}]={rec[H_MAPPED]}, not this slot") + if fl & F_STALE: + sid = rec[H_STALE_ID] + text = (f"stale slot_for_id: {self._name(sid)} -> this slot, which holds {self._name(owner)}" + + (f" ({rec[H_STALE_N]} such ids)" if rec[H_STALE_N] > 1 else "")) + gone = [ev for ev in meta if ev["kind"] in INSTALLS and self._flat(ev) == sid] + if gone: + later = [ev for ev in meta if ev["seq"] > gone[-1]["seq"]] + text += (f"; {self._name(sid)} was installed here by {self._kind(gone[-1])} #{gone[-1]['seq']}" + + (f" and replaced by {self._kind(later[0])} #{later[0]['seq']} without clearing " + f"its map entry" if later else "")) + parts.append(text) + # an install's recorder reads the slot right after the install on the same stream + for ev in meta: + if ev["kind"] in INSTALLS and ev["owner"] != self._flat(ev): + parts.append(f"right after the {self._kind(ev)} of {self._name(self._flat(ev))} (#{ev['seq']}) the slot " + f"named {self._name(ev['owner'])}: another map writer ran concurrently") + return "; ".join(parts) or "clean" + + def format_record(self, rec: list[int]) -> str: + fl = rec[H_FLAGS] + text = (f"audit={rec[H_AUDIT]} step={rec[H_DSTEP]} lru={rec[H_LSTEP]} slot={rec[H_SLOT]} " + f"holds {self._name(rec[H_OWNER])} (slot_for_id={rec[H_MAPPED]}, usage={rec[H_USAGE]})") + if fl & F_BYTES: + banks = [self.bank_names[b] if b < len(self.bank_names) else str(b) + for b in range(max(len(self.bank_names), 1)) if rec[H_BANKS] >> b & 1] + first = self.bank_names[rec[H_FIRST_BANK]] if 0 <= rec[H_FIRST_BANK] < len(self.bank_names) else rec[H_FIRST_BANK] + text += (f" | bytes differ in {', '.join(banks)} (first at {first}+{rec[H_FIRST_BYTE]}, " + f"{rec[H_BAD_WORDS]} bad words)") + text += f" | diagnosis: {self.diagnose(rec)}" + history = "; ".join(self.format_event(ev) for ev in self.events(rec)) + return text + f" | history ({rec[H_META_N]} map / {rec[H_BYTES_N]} byte events in all): {history or 'none'}" diff --git a/python/freetoken/moe/verify.py b/python/freetoken/moe/verify.py new file mode 100644 index 0000000000..cc71e19d9d --- /dev/null +++ b/python/freetoken/moe/verify.py @@ -0,0 +1,516 @@ +"""Decode expert-cache verifier (``FREETOKEN_MOE_PREFETCH_VERIFY``): a debug detector for slot races. + +For every GPU decode layer call it checks, on the compute stream and in its order, what the +grouped GEMM is about to read and has just read: + +- meta (before and after the GEMM): each routed entry's slot maps to the expert the router chose, + in both slot maps (``id_of_slot[s] == layer * E + id`` and ``slot_for_id[layer, id] == s``); the + router's ids are copied before ensure rewrites them in place; +- bytes before the GEMM (``full``): every routed slot's row in every bank equals the host bank row + of that expert for that layer, gathered into a scratch buffer with the fused copy; +- bytes after the GEMM (``full``): the same slots against the same scratch, so a slot overwritten + while the GEMM ran shows up even when a later copy restores it. + +Nothing here records or waits on an event: the prefetch streams keep exactly the edges they have +without it. The checks do delay the compute stream (``full`` gathers every routed expert over +PCIe), which moves where the other streams land relative to it. ``meta`` skips the byte checks and +costs a few one-CTA launches per layer. Counts accumulate per layer on the device and the first +``RING`` offending records are kept with the slot state seen at check time; the engine reports +them every MOE_STATS_INTERVAL decode steps and at shutdown. +""" +from __future__ import annotations + +import math + +import torch +import triton +import triton.language as tl + +from freetoken.env import ENV +from freetoken.utils import init_logger + +logger = init_logger(__name__) + +MODES = ("off", "meta", "full") +# counters[layer, col]: checked calls, offending records per check, routed experts left unchecked +CHECKS, META_BAD, PRE_BAD, POST_BAD, POST_META_BAD, UNCHECKED = range(6) +NUM_COLS = 6 +COL_NAMES = ("checks", "meta_bad", "pre_bad", "post_bad", "post_meta_bad", "unchecked") +META_PRE, BYTES_PRE, META_POST, BYTES_POST = range(4) +KINDS = ("meta_pre", "bytes_pre", "meta_post", "bytes_post") +# record fields; owner / mapped / usage / step are read when the check runs +(R_CALL, R_LAYER, R_ROW, R_EXPERT, R_SLOT, R_BANK, R_KIND, R_OWNER, R_MAPPED, R_USAGE, R_STEP, + R_SOURCE, R_FIRST_BYTE, R_BAD_WORDS) = range(14) +NUM_FIELDS = 14 +# R_SOURCE bits: the slot is in this call's demand copy plan, this layer's prefetch copy plan, or +# the next layer's prefetch plan (installed earlier on the compute stream) +SRC_DEMAND, SRC_PREFETCH, SRC_NEXT_PREFETCH = 1, 2, 4 +RING = 64 +# routed entries a call may have (bs * top_k); wider eager batches are counted unchecked +MAX_ENTRIES = 64 +CHUNK_WORDS = 4096 +_NF = tl.constexpr(NUM_FIELDS) +_NC = tl.constexpr(NUM_COLS) +_NO_BAD = tl.constexpr(0x7FFFFFFF) + + +def resolve_mode(mode: str | None = None) -> str: + raw = (ENV.MOE_PREFETCH_VERIFY.value if mode is None else mode).strip().lower() + if raw in ("", "0", "off", "false", "no"): + return "off" + if raw in ("1", "on", "true", "yes", "full"): + return "full" + if raw == "meta": + return "meta" + raise ValueError(f"FREETOKEN_MOE_PREFETCH_VERIFY={raw!r}: expected 0, 1 (full) or meta") + + +@triton.jit +def _put(vals, f, idx: tl.constexpr, v): + return tl.where(f[None, :] == idx, v.to(tl.int64)[:, None], vals) + + +@triton.jit +def _in_plan(s, plan_ptr, num_ptr, W: tl.constexpr, W_P: tl.constexpr): + j = tl.arange(0, W_P) + plan = tl.load(plan_ptr + j, mask=j < tl.minimum(tl.load(num_ptr), W), other=-3) + return tl.sum((s[:, None] == plan[None, :]).to(tl.int32), axis=1) > 0 + + +@triton.jit +def _source(s, dem_ptr, dem_num_ptr, pf_ptr, pf_num_ptr, nx_ptr, nx_num_ptr, + DEM_W: tl.constexpr, DEM_P: tl.constexpr, PF_W: tl.constexpr, PF_P: tl.constexpr, + HAS_PF: tl.constexpr, HAS_NX: tl.constexpr): + src = tl.where(_in_plan(s, dem_ptr, dem_num_ptr, DEM_W, DEM_P), 1, 0) + if HAS_PF: + src += tl.where(_in_plan(s, pf_ptr, pf_num_ptr, PF_W, PF_P), 2, 0) + if HAS_NX: + src += tl.where(_in_plan(s, nx_ptr, nx_num_ptr, PF_W, PF_P), 4, 0) + return src + + +@triton.jit +def _write(ring_ptr, cursor, bad, vals, RING_N: tl.constexpr): + """Append the ``bad`` rows of ``vals`` after ``cursor`` while the ring has room; the new cursor.""" + f = tl.arange(0, 16) + pos = cursor + tl.cumsum(bad.to(tl.int64), axis=0) - 1 + keep = bad & (pos < RING_N) + tl.store(ring_ptr + pos[:, None] * _NF + f[None, :], vals, mask=keep[:, None] & (f[None, :] < _NF)) + return cursor + tl.sum(bad.to(tl.int64), axis=0) + + +@triton.jit +def _meta(ids_ptr, slots_ptr, id_of_slot_ptr, slot_for_id_ptr, usage_ptr, step_ptr, + dem_ptr, dem_num_ptr, pf_ptr, pf_num_ptr, nx_ptr, nx_num_ptr, + ring_ptr, cursor, counters_ptr, col: tl.constexpr, kind: tl.constexpr, call, layer, num_slots, + N: tl.constexpr, N_P: tl.constexpr, E: tl.constexpr, + DEM_W: tl.constexpr, DEM_P: tl.constexpr, PF_W: tl.constexpr, PF_P: tl.constexpr, + HAS_PF: tl.constexpr, HAS_NX: tl.constexpr, RING_N: tl.constexpr): + """Check every routed entry's slot against both maps; returns the new ring cursor.""" + i = tl.arange(0, N_P) + m = i < N + ids = tl.load(ids_ptr + i, mask=m, other=0) + s = tl.load(slots_ptr + i, mask=m, other=-1) + ok_s = m & (s >= 0) & (s < num_slots) + ok_id = m & (ids >= 0) & (ids < E) + fid = layer * E + ids + owner = tl.load(id_of_slot_ptr + s, mask=ok_s, other=-2) + mapped = tl.load(slot_for_id_ptr + fid, mask=ok_id, other=-2) + bad = m & ~(ok_s & ok_id & (owner == fid) & (mapped == s)) + usage = tl.load(usage_ptr + s, mask=ok_s, other=-1) + src = _source(s, dem_ptr, dem_num_ptr, pf_ptr, pf_num_ptr, nx_ptr, nx_num_ptr, DEM_W, DEM_P, PF_W, PF_P, HAS_PF, HAS_NX) + f = tl.arange(0, 16) + z = tl.zeros([N_P], tl.int64) + vals = tl.zeros([N_P, 16], tl.int64) - 1 + vals = _put(vals, f, 0, z + call) + vals = _put(vals, f, 1, z + layer) + vals = _put(vals, f, 2, i) + vals = _put(vals, f, 3, ids) + vals = _put(vals, f, 4, s) + vals = _put(vals, f, 6, z + kind) + vals = _put(vals, f, 7, owner) + vals = _put(vals, f, 8, mapped) + vals = _put(vals, f, 9, usage) + vals = _put(vals, f, 10, z + tl.load(step_ptr)) + vals = _put(vals, f, 11, src) + nbad = tl.sum(bad.to(tl.int64), axis=0) + tl.store(counters_ptr + col, tl.load(counters_ptr + col) + nbad) + return _write(ring_ptr, cursor, bad, vals, RING_N) + + +@triton.jit(do_not_specialize=["layer", "num_slots"]) +def _plan_kernel( + ids_ptr, slots_ptr, id_of_slot_ptr, slot_for_id_ptr, usage_ptr, step_ptr, + dem_ptr, dem_num_ptr, pf_ptr, pf_num_ptr, nx_ptr, nx_num_ptr, + ring_ptr, cursor_ptr, counters_ptr, calls_ptr, + u_ids_ptr, u_slots_ptr, u_rows_ptr, u_num_ptr, slots_copy_ptr, + layer, num_slots, + N: tl.constexpr, N_P: tl.constexpr, E: tl.constexpr, CAP: tl.constexpr, + DEM_W: tl.constexpr, DEM_P: tl.constexpr, PF_W: tl.constexpr, PF_P: tl.constexpr, + HAS_PF: tl.constexpr, HAS_NX: tl.constexpr, BYTES: tl.constexpr, RING_N: tl.constexpr, +): + """Before the GEMM: count the call, check the maps, keep the slots for the after-GEMM checks and + list the distinct (expert, slot) pairs to byte-check (first CAP of them).""" + call = tl.load(calls_ptr + layer) + tl.store(calls_ptr + layer, call + 1) + row = counters_ptr + layer * _NC + tl.store(row + 0, tl.load(row + 0) + 1) + cursor = tl.load(cursor_ptr) + cursor = _meta(ids_ptr, slots_ptr, id_of_slot_ptr, slot_for_id_ptr, usage_ptr, step_ptr, + dem_ptr, dem_num_ptr, pf_ptr, pf_num_ptr, nx_ptr, nx_num_ptr, + ring_ptr, cursor, row, 1, 0, call, layer, num_slots, + N, N_P, E, DEM_W, DEM_P, PF_W, PF_P, HAS_PF, HAS_NX, RING_N) + tl.store(cursor_ptr, cursor) + i = tl.arange(0, N_P) + m = i < N + ids = tl.load(ids_ptr + i, mask=m, other=-1) + s = tl.load(slots_ptr + i, mask=m, other=-1) + tl.store(slots_copy_ptr + i, s, mask=m) + if BYTES: + # an unreadable slot or id is already a meta record; a duplicate pair shares its first row + ok = m & (s >= 0) & (s < num_slots) & (ids >= 0) & (ids < E) + same = (ids[:, None] == ids[None, :]) & (s[:, None] == s[None, :]) & (i[None, :] < i[:, None]) & ok[None, :] + first = ok & (tl.sum(same.to(tl.int32), axis=1) == 0) + rank = tl.cumsum(first.to(tl.int32), axis=0) - 1 + take = first & (rank < CAP) + tl.store(u_ids_ptr + rank, ids, mask=take) + tl.store(u_slots_ptr + rank, s, mask=take) + tl.store(u_rows_ptr + rank, i, mask=take) + total = tl.sum(first.to(tl.int64), axis=0) + kept = tl.minimum(total, CAP) + tl.store(u_num_ptr, kept) + tl.store(row + 5, tl.load(row + 5) + total - kept) + + +@triton.jit +def _compare_kernel(cache_ptrs, ref_ptrs, words_ptr, u_slots_ptr, u_num_ptr, bad_cnt_ptr, bad_first_ptr, + NB: tl.constexpr, BLOCK: tl.constexpr): + """Program (row, bank, chunk): compare one chunk of a routed slot's row with its host copy.""" + r = tl.program_id(0) + b = tl.program_id(1) + words = tl.load(words_ptr + b) + off = tl.program_id(2).to(tl.int64) * BLOCK + tl.arange(0, BLOCK) + m = (r < tl.load(u_num_ptr)) & (off < words) + slot = tl.load(u_slots_ptr + r).to(tl.int64) + got = tl.load(tl.load(cache_ptrs + b).to(tl.pointer_type(tl.int32)) + slot * words + off, mask=m, other=0) + want = tl.load(tl.load(ref_ptrs + b).to(tl.pointer_type(tl.int32)) + r * words + off, mask=m, other=0) + bad = m & (got != want) + nbad = tl.sum(bad.to(tl.int32), axis=0) + if nbad > 0: + tl.atomic_add(bad_cnt_ptr + r * NB + b, nbad) + tl.atomic_min(bad_first_ptr + r * NB + b, tl.min(tl.where(bad, off, _NO_BAD), axis=0).to(tl.int32)) + + +@triton.jit(do_not_specialize=["layer", "num_slots"]) +def _tally_kernel( + ids_ptr, slots_copy_ptr, id_of_slot_ptr, slot_for_id_ptr, usage_ptr, step_ptr, + dem_ptr, dem_num_ptr, pf_ptr, pf_num_ptr, nx_ptr, nx_num_ptr, + ring_ptr, cursor_ptr, counters_ptr, calls_ptr, + u_ids_ptr, u_slots_ptr, u_rows_ptr, u_num_ptr, bad_cnt_ptr, bad_first_ptr, + layer, num_slots, + N: tl.constexpr, N_P: tl.constexpr, E: tl.constexpr, CAP: tl.constexpr, CAP_P: tl.constexpr, + NB: tl.constexpr, NB_P: tl.constexpr, + DEM_W: tl.constexpr, DEM_P: tl.constexpr, PF_W: tl.constexpr, PF_P: tl.constexpr, + HAS_PF: tl.constexpr, HAS_NX: tl.constexpr, BYTES: tl.constexpr, POST: tl.constexpr, + BYTES_KIND: tl.constexpr, BYTES_COL: tl.constexpr, RING_N: tl.constexpr, +): + """Turn the compare flags into counts and records (and clear them); after the GEMM also recheck + the maps of the slots it read.""" + call = tl.load(calls_ptr + layer) - 1 + row = counters_ptr + layer * _NC + cursor = tl.load(cursor_ptr) + if POST: + cursor = _meta(ids_ptr, slots_copy_ptr, id_of_slot_ptr, slot_for_id_ptr, usage_ptr, step_ptr, + dem_ptr, dem_num_ptr, pf_ptr, pf_num_ptr, nx_ptr, nx_num_ptr, + ring_ptr, cursor, row, 4, 2, call, layer, num_slots, + N, N_P, E, DEM_W, DEM_P, PF_W, PF_P, HAS_PF, HAS_NX, RING_N) + if BYTES: + j = tl.arange(0, CAP_P * NB_P) + r = j // NB_P + b = j % NB_P + cell = (r < CAP) & (b < NB) + live = cell & (r < tl.load(u_num_ptr)) + cnt = tl.load(bad_cnt_ptr + r * NB + b, mask=live, other=0) + first = tl.load(bad_first_ptr + r * NB + b, mask=live, other=_NO_BAD) + tl.store(bad_cnt_ptr + r * NB + b, tl.zeros([CAP_P * NB_P], tl.int32), mask=cell) + tl.store(bad_first_ptr + r * NB + b, tl.zeros([CAP_P * NB_P], tl.int32) + _NO_BAD, mask=cell) + bad = live & (cnt > 0) + ids = tl.load(u_ids_ptr + r, mask=live, other=0) + s = tl.load(u_slots_ptr + r, mask=live, other=-1) + ok_s = live & (s >= 0) & (s < num_slots) + src = _source(s, dem_ptr, dem_num_ptr, pf_ptr, pf_num_ptr, nx_ptr, nx_num_ptr, DEM_W, DEM_P, PF_W, PF_P, HAS_PF, HAS_NX) + f = tl.arange(0, 16) + z = tl.zeros([CAP_P * NB_P], tl.int64) + vals = tl.zeros([CAP_P * NB_P, 16], tl.int64) - 1 + vals = _put(vals, f, 0, z + call) + vals = _put(vals, f, 1, z + layer) + vals = _put(vals, f, 2, tl.load(u_rows_ptr + r, mask=live, other=-1)) + vals = _put(vals, f, 3, ids) + vals = _put(vals, f, 4, s) + vals = _put(vals, f, 5, b) + vals = _put(vals, f, 6, z + BYTES_KIND) + vals = _put(vals, f, 7, tl.load(id_of_slot_ptr + s, mask=ok_s, other=-2)) + vals = _put(vals, f, 8, tl.load(slot_for_id_ptr + layer * E + ids, mask=live, other=-2)) + vals = _put(vals, f, 9, tl.load(usage_ptr + s, mask=ok_s, other=-1)) + vals = _put(vals, f, 10, z + tl.load(step_ptr)) + vals = _put(vals, f, 11, src) + vals = _put(vals, f, 12, first.to(tl.int64) * 4) + vals = _put(vals, f, 13, cnt) + tl.store(row + BYTES_COL, tl.load(row + BYTES_COL) + tl.sum(bad.to(tl.int64), axis=0)) + cursor = _write(ring_ptr, cursor, bad, vals, RING_N) + tl.store(cursor_ptr, cursor) + + +class ExpertVerifier: + """Device counters, the record ring and (``full``) the scratch the byte checks compare against. + + ``begin`` copies the router's ids before ensure rewrites them, ``before_gemm`` runs after the + joins and the demand copy, ``after_gemm`` right after the GEMM; all three enqueue on the current + (compute) stream with fixed shapes and no host sync, so they capture into the decode graph. + ``bind`` sizes the scratch from the banks once and refreshes the slot bank addresses in place + (a rebuild reallocates the banks); counters survive rebuilds and resets. + """ + + def __init__(self, num_layers: int, num_experts: int, device: torch.device, *, mode: str = "full", + rows: int | None = None) -> None: + from freetoken.kernel.fast_index_copy import _skip_fast_index_copy_enabled + + if mode == "full" and _skip_fast_index_copy_enabled(): + logger.warning( + "FREETOKEN_MOE_PREFETCH_VERIFY=1 with FREETOKEN_SKIP_FAST_INDEX_COPY: the host rows cannot be " + "gathered, so only the slot maps are checked (meta)" + ) + mode = "meta" + self.mode = mode + self.num_layers = num_layers + self.num_experts = num_experts + self.device = device + self.rows = ENV.MOE_PREFETCH_VERIFY_ROWS.value if rows is None else int(rows) + if self.rows < 1: + raise ValueError(f"FREETOKEN_MOE_PREFETCH_VERIFY_ROWS={self.rows} must be >= 1") + self.counters = torch.zeros((num_layers, NUM_COLS), dtype=torch.int64, device=device) + self.totals = torch.zeros((num_layers, NUM_COLS), dtype=torch.int64) + self.calls = torch.zeros((num_layers,), dtype=torch.int64, device=device) + self.ring = torch.full((RING, NUM_FIELDS), -1, dtype=torch.int64, device=device) + self.cursor = torch.zeros((1,), dtype=torch.int64, device=device) + self.ids = torch.full((MAX_ENTRIES,), -1, dtype=torch.int32, device=device) + self.slots = torch.full((MAX_ENTRIES,), -1, dtype=torch.int32, device=device) + self.u_ids = torch.zeros((self.rows,), dtype=torch.int32, device=device) + self.u_slots = torch.zeros((self.rows,), dtype=torch.int32, device=device) + self.u_rows = torch.zeros((self.rows,), dtype=torch.int32, device=device) + self.u_num = torch.zeros((1,), dtype=torch.int64, device=device) + self.iota = torch.arange(self.rows, dtype=torch.int32, device=device) + self.bank_names: list[str] = [] + self.scratch: torch.Tensor | None = None + self._printed = 0 + self._top_k = 1 + self._call = None # (layer, entries, this layer's prefetch plan or None) between begin and after_gemm + self._next = None + logger.info( + f"MoE verify {self.mode}: checking every GPU decode layer call's slot maps" + + (f" and slot bytes before and after the GEMM (up to {self.rows} experts a call)" if self.mode == "full" else "") + ) + + def bind(self, cache) -> None: + """Size the scratch for ``cache``'s banks (first call) or point the checks at its new banks.""" + caches = [bank for _, bank in cache.banks] + feats = [math.prod(bank.shape[1:]) * bank.element_size() for bank in caches] + if self.bank_names: + assert list(cache.bank_schema) == self.bank_names and feats == self._feats, "the bank layout changed" + else: + assert all(f % 4 == 0 for f in feats), feats + self.bank_names = list(cache.bank_schema) + self._feats = feats + self.cache_ptrs = torch.zeros((len(feats),), dtype=torch.int64, device=self.device) + if self.mode == "full": + self._alloc_scratch(cache, feats) + self.cache_ptrs.copy_(torch.tensor([bank.data_ptr() for bank in caches], dtype=torch.int64)) + + def _alloc_scratch(self, cache, feats: list[int]) -> None: + rows = self.rows + self.scratch = torch.empty((rows * sum(feats),), dtype=torch.uint8, device=self.device) + offsets = [rows * sum(feats[:b]) for b in range(len(feats))] + self.scratch_views = [ + self.scratch[o : o + rows * f].view(bank.dtype).view(rows, *bank.shape[1:]) + for o, f, (_, bank) in zip(offsets, feats, cache.banks) + ] + self.scratch_ptrs = torch.tensor([self.scratch.data_ptr() + o for o in offsets], dtype=torch.int64, device=self.device) + self.words = torch.tensor([f // 4 for f in feats], dtype=torch.int64, device=self.device) + self.bad_cnt = torch.zeros((rows * len(feats),), dtype=torch.int32, device=self.device) + self.bad_first = torch.full((rows * len(feats),), 0x7FFFFFFF, dtype=torch.int32, device=self.device) + self._chunks = triton.cdiv(max(feats) // 4, CHUNK_WORDS) + logger.info(f"MoE verify scratch: {self.scratch.numel() / 2**20:.1f} MiB for {rows} experts") + + def begin(self, cache, layer_id: int, topk_ids: torch.Tensor, plan: tuple | None) -> None: + """Copy the router's ids of ``layer_id`` (``plan``: this layer's prefetch copy plan, if forked).""" + assert topk_ids.dtype == torch.int32 and topk_ids.is_contiguous() + n = topk_ids.numel() + self._top_k = topk_ids.shape[-1] + if n > MAX_ENTRIES: + self.counters[layer_id, UNCHECKED] += n + self._call = None + return + self.ids[:n].copy_(topk_ids.view(-1)) + self._call = (layer_id, n, plan) + + def _plans(self, cache) -> tuple: + _, n, plan = self._call + pf, nx = plan, self._next + dummy = (self.u_slots, self.u_num) + width = (pf or nx or dummy)[0].numel() + return ( + cache.evict_slots, cache.num_indices, *(pf or dummy), *(nx or dummy), + dict(DEM_W=n, DEM_P=triton.next_power_of_2(n), PF_W=width, PF_P=triton.next_power_of_2(width), + HAS_PF=pf is not None, HAS_NX=nx is not None), + ) + + def before_gemm(self, cache, slots: torch.Tensor, next_plan: tuple | None) -> None: + """Check the maps and (``full``) the bytes of the slots ``slots`` the GEMM is about to read + (``next_plan``: the next layer's prefetch plan, if one was issued).""" + if self._call is None: + return + layer, n, _ = self._call + self._next = next_plan + *plans, plan_kw = self._plans(cache) + full = self.mode == "full" + _plan_kernel[(1,)]( + self.ids, slots, cache.id_of_slot, cache.slot_for_id.view(-1), cache.usage, cache.step, *plans, + self.ring, self.cursor, self.counters, self.calls, + self.u_ids, self.u_slots, self.u_rows, self.u_num, self.slots, + layer, cache.cache_size, + N=n, N_P=triton.next_power_of_2(n), E=self.num_experts, CAP=self.rows, + BYTES=full, RING_N=RING, num_warps=4, **plan_kw, + ) + if not full: + return + self._gather(cache, layer) + self._compare_and_tally(cache, post=False) + + def after_gemm(self, cache) -> None: + """Recheck the maps and (``full``) the bytes of the slots the GEMM just read.""" + if self._call is None: + return + if self.mode == "full": + self._compare_and_tally(cache, post=True) + else: + self._tally(cache, post=True, bytes_=False) + self._call = self._next = None + + def _gather(self, cache, layer: int) -> None: + if cache._copy_fused_ok: + from freetoken.kernel.fast_index_copy import fast_index_copy_multi_jit + + fast_index_copy_multi_jit( + self.scratch_ptrs, cache._copy_src_ptrs[layer], cache._copy_feat_bytes[cache.layer_layout[layer]], self.iota, self.u_ids, self.u_num, + ) + return + from freetoken.kernel import fast_index_copy_jit + + for (per_layer, _), view in zip(cache.banks, self.scratch_views): + fast_index_copy_jit(view, self.iota, per_layer[layer], self.u_ids, self.u_num) + + def _compare_and_tally(self, cache, *, post: bool) -> None: + nb = len(self.bank_names) + _compare_kernel[(self.rows, nb, self._chunks)]( + self.cache_ptrs, self.scratch_ptrs, self.words, self.u_slots, self.u_num, self.bad_cnt, self.bad_first, + NB=nb, BLOCK=CHUNK_WORDS, num_warps=8, + ) + self._tally(cache, post=post, bytes_=True) + + def _tally(self, cache, *, post: bool, bytes_: bool) -> None: + layer, n, _ = self._call + *plans, plan_kw = self._plans(cache) + nb = max(len(self.bank_names), 1) + _tally_kernel[(1,)]( + self.ids, self.slots, cache.id_of_slot, cache.slot_for_id.view(-1), cache.usage, cache.step, *plans, + self.ring, self.cursor, self.counters, self.calls, + self.u_ids, self.u_slots, self.u_rows, self.u_num, + self.bad_cnt if bytes_ else self.u_rows, self.bad_first if bytes_ else self.u_rows, + layer, cache.cache_size, + N=n, N_P=triton.next_power_of_2(n), E=self.num_experts, CAP=self.rows, + CAP_P=triton.next_power_of_2(self.rows), NB=nb, NB_P=triton.next_power_of_2(nb), + BYTES=bytes_, POST=post, BYTES_KIND=BYTES_POST if post else BYTES_PRE, BYTES_COL=POST_BAD if post else PRE_BAD, + RING_N=RING, num_warps=4, **plan_kw, + ) + + # ------------------------------------------------------------------ reporting (host syncs) + + def take_window(self) -> tuple[torch.Tensor, list[list[int]], int]: + """This window's ``[num_layers, NUM_COLS]`` counts (added to ``totals``; the device ones + restart), the records kept since the last call, and how many records there were in all.""" + window = self.counters.cpu() + self.counters.zero_() + self.totals += window + cursor = int(self.cursor.item()) + upto = min(cursor, RING) + records = self.ring[self._printed : upto].cpu().tolist() if upto > self._printed else [] + self._printed = max(self._printed, upto) + return window, records, cursor + + def report_window(self, steps: int) -> list[tuple[str, str]]: + """(level, line) pairs for one MOE_STATS_INTERVAL window; nothing when no call was checked.""" + window, records, cursor = self.take_window() + if not int(window[:, CHECKS].sum()) and not records: + return [] + bad = self._bad(window) or bool(records) + session = self._bad(self.totals) + lines = [("warning" if bad else "info", + f"MoE verify {self.mode} ({steps} decode steps): {self.format_counts(window)}; session bad={session}")] + lines += [("warning", "MoE verify record: " + self.format_record(r)) for r in records] + if bad and cursor > RING: + lines.append(("warning", f"MoE verify: {cursor} bad records in all, the first {RING} are kept")) + return lines + + def report_session(self) -> list[tuple[str, str]]: + """Session counts and every kept record (shutdown).""" + totals = self.totals + self.counters.cpu() + cursor = int(self.cursor.item()) + if not int(totals[:, CHECKS].sum()): + return [] + bad = self._bad(totals) + lines = [("warning" if bad else "info", f"MoE verify {self.mode} (session): {self.format_counts(totals)}")] + records = self.ring[: min(cursor, RING)].cpu().tolist() + lines += [("warning", "MoE verify record: " + self.format_record(r)) for r in records] + if cursor > RING: + lines.append(("warning", f"MoE verify: {cursor} bad records in all, the first {RING} are kept")) + return lines + + @staticmethod + def _bad(counts: torch.Tensor) -> int: + return int(counts[:, [META_BAD, PRE_BAD, POST_BAD, POST_META_BAD]].sum()) + + def format_counts(self, counts: torch.Tensor) -> str: + total = counts.sum(0).tolist() + text = f"{total[CHECKS]} layer calls checked, " + ", ".join( + f"{COL_NAMES[c]}={total[c]}" for c in (META_BAD, PRE_BAD, POST_BAD, POST_META_BAD) + ) + if total[UNCHECKED]: + text += f", unchecked experts={total[UNCHECKED]}" + per_layer = [ + f"L{layer}=" + "/".join(str(row[c]) for c in (META_BAD, PRE_BAD, POST_BAD, POST_META_BAD)) + for layer, row in enumerate(counts.tolist()) + if any(row[c] for c in (META_BAD, PRE_BAD, POST_BAD, POST_META_BAD)) + ] + if per_layer: + text += " (meta/pre/post/post_meta by layer: " + ", ".join(per_layer) + ")" + return text + + def format_record(self, rec: list[int]) -> str: + e = self.num_experts + kind = KINDS[rec[R_KIND]] if 0 <= rec[R_KIND] < len(KINDS) else str(rec[R_KIND]) + owner = rec[R_OWNER] + owner_text = f"L{owner // e}/e{owner % e}" if owner >= 0 else str(owner) + src = [name for bit, name in ((SRC_DEMAND, "demand copy"), (SRC_PREFETCH, "prefetch copy"), + (SRC_NEXT_PREFETCH, "next layer's prefetch plan")) if rec[R_SOURCE] & bit] + row = rec[R_ROW] + text = ( + f"{kind} call={rec[R_CALL]} layer={rec[R_LAYER]} row={row} (token {row // self._top_k}, rank {row % self._top_k}) " + f"expert={rec[R_EXPERT]} slot={rec[R_SLOT]}" + ) + if rec[R_BANK] >= 0: + bank = self.bank_names[rec[R_BANK]] if rec[R_BANK] < len(self.bank_names) else rec[R_BANK] + text += f" bank={bank} first_bad_byte={rec[R_FIRST_BYTE]} bad_words={rec[R_BAD_WORDS]}" + return text + ( + f" | slot now holds {owner_text}, slot_for_id={rec[R_MAPPED]}, usage={rec[R_USAGE]}, lru_step={rec[R_STEP]}, " + f"slot in: {', '.join(src) or 'no copy plan (a resident hit)'}" + ) diff --git a/python/freetoken/scheduler/cache.py b/python/freetoken/scheduler/cache.py index 772297e069..ff4467328e 100644 --- a/python/freetoken/scheduler/cache.py +++ b/python/freetoken/scheduler/cache.py @@ -56,6 +56,31 @@ def __init__(self, num_pages: int, page_size: int, page_table: torch.Tensor, typ if swa_pool is not None: self.prefill_chunk_budget = getattr(swa_pool, "prefill_chunk_budget", None) self.prefix_cache = self._make_prefix_cache(device, page_size, type) + # Host prefix tier: FT_GDN_HOST_TIER=1 guarda snapshots GDN na RAM (Degrau 2); + # FT_PREFIX_HOST=1 + FT_PREFIX_HOST_GB adiciona o KV de prefixo num slot space host + # separado (Degrau 3 — o radix nunca perde o mapeamento). Hybrid path only. + self.host_tier = None + gdn_on = os.environ.get("FT_GDN_HOST_TIER", "0") == "1" + kv_gb = (float(os.environ.get("FT_PREFIX_HOST_GB", "16")) + if os.environ.get("FT_PREFIX_HOST", "0") == "1" else 0) + if self.is_hybrid and (gdn_on or kv_gb > 0): + from freetoken.core import get_global_ctx + from freetoken.kvcache.host_prefix_tier import HostPrefixTier + from freetoken.utils import init_logger + + slots = int(os.environ.get("FT_GDN_HOST_SLOTS", "32")) + kv_pool = get_global_ctx().kv_cache if kv_gb > 0 else None + self.host_tier = HostPrefixTier( + linear_state_pool, kv_pool, gdn_slots_host=slots, + kv_budget_bytes=int(kv_gb * (1 << 30)), device=device) + self.prefix_cache.on_evict_node = self._on_evict_node_to_host + self.prefix_cache.on_match_dead_snapshot = self._rehydrate_snapshot + if kv_pool is not None: + self.prefix_cache.on_evict_kv = self._on_evict_kv_to_host + init_logger(__name__).info( + f"host prefix tier ON: {slots} GDN slots" + + (f" + {kv_gb:.0f} GiB KV store" if kv_pool is not None else "") + ) self.device = device self.num_pages = num_pages self.page_table = page_table @@ -106,6 +131,16 @@ def match_req(self, req: PendingReq, max_len: int | None = None) -> MatchResult: if self.is_hybrid: from freetoken.kvcache.hybrid_radix_cache import HybridCacheHandle m = self.prefix_cache.match_prefix(ids) + # Host prefix tier (FT_PREFIX_HOST): a prefix evicted from the tree lives in RAM + # (KV bytes + GDN snapshot). If the tree match fell short, rehydrate the deeper + # host entry into fresh pages, re-insert it into the tree and re-match. + tier = self.host_tier + if tier is not None and tier.has_kv: + hit = tier.lookup(ids, self.page_size) + if (hit is not None and hit.kv_pages is not None and hit.gdn_slot is not None + and hit.host_len > m.cached_len): + self._rehydrate_host_prefix(ids, hit) + m = self.prefix_cache.match_prefix(ids) return MatchResult( HybridCacheHandle(m.cached_len, m.node, m.kv_indices), mamba_value=m.mamba_value) return self.prefix_cache.match_prefix(ids) @@ -553,6 +588,105 @@ def check_integrity(self) -> None: if self.page_size > 1: assert torch.all(self.free_slots % self.page_size == 0) + # ----- host GDN-snapshot tier (FT_GDN_HOST_TIER) ------------------------------- + def _on_evict_node_to_host(self, node) -> None: + """Tree hook: the node's GDN snapshot is about to be freed. Copy it to the host tier + and index the node's prefix (the GPU slot is freed only after we return).""" + tier = self.host_tier + prefix = self.prefix_cache._collect_key(node) + try: + hs = tier.write_gdn(node.mamba_value) + except RuntimeError: + # host tier full: drop the OLDEST indexed prefix and retry once + if not tier._index: + return + oldest = next(iter(tier._index)) + tier.evict_host(tier._index[oldest].key_tokens) + try: + hs = tier.write_gdn(node.mamba_value) + except RuntimeError: + return + tier.put_prefix(prefix, None, hs) + + def _on_evict_kv_to_host(self, node) -> None: + """Tree hook (FT_PREFIX_HOST): the node's KV pages are about to be freed. Copy the + WHOLE prefix's pages to the host store while the bytes are still valid (eviction is + bottom-up, so the ancestors' pages are still live), and index the prefix.""" + tier = self.host_tier + prefix = self.prefix_cache._collect_key(node) + kv = self.prefix_cache._collect_kv(node) + slots = tier.write_kv(kv) + if slots is not None: + tier.put_prefix(prefix, slots, None) + + def _rehydrate_snapshot(self, node) -> int | None: + """Tree hook (match path): the node is a snapshot tombstone. If the host tier has its + prefix, copy the snapshot back into a fresh pool slot and return the slot; None = + keep walking (the match truncates as before).""" + tier = self.host_tier + key = self.prefix_cache._collect_key(node) + # exact only: a shorter entry's state would resume this node's full depth with the + # tokens in between missing from the recurrence + hit = tier.lookup_exact(key) + if hit is None or hit.gdn_slot is None: + return None + pool = self.linear_state_pool + if pool.num_free_slots < 1: + # make room: evicting another snapshot writes IT to the host tier first, which may + # drop this entry and hand its host slot to the victim, so probe again + self.ensure_mamba_slots(1) + hit = tier.lookup_exact(key) + if pool.num_free_slots < 1 or hit is None or hit.gdn_slot is None: + return None + slot = pool.alloc(1)[0] + tier.read_gdn(hit.gdn_slot, slot) + return slot + + def _rehydrate_host_prefix(self, ids: torch.Tensor, hit) -> None: + """FT_PREFIX_HOST: sobe um prefixo frio do tier pro tree+GPU -- aloca páginas novas + (pode evictar: a vítima é escrita no host primeiro), copia os bytes KV H2D, restaura + o snapshot GDN num slot novo e re-insere o prefixo no tree (o match seguinte acha). + """ + tier = self.host_tier + n_pages = hit.host_len // self.page_size + kv_new = self._page_to_token(self._allocate(n_pages)) + pool = self.linear_state_pool + if hit.gdn_slot is not None: + self.ensure_mamba_slots(1) + # _allocate and ensure_mamba_slots write their victims to the tier, which may drop + # this entry and hand its host slots to them: use what the key maps to now + hit = tier.lookup_exact(ids[: hit.host_len]) + if (pool.num_free_slots < 1 or hit is None or hit.kv_pages is None + or hit.gdn_slot is None): + self._free(kv_new) # sem slot GDN não dá pra reidratar; devolve as páginas + return + slot = pool.alloc(1)[0] + tier.read_gdn(hit.gdn_slot, slot) + else: + self._free(kv_new) + return # o modelo híbrido não resume sem snapshot GDN + # Com o KVHostOffloader ativo, as páginas recém-alocadas são LÓGICAS: subir bytes + # pra GPU em bulk estoura o pool físico quando o prefixo > pool (o fetch lazy já + # resolve residência na atenção). Copiamos host->host pro espelho e invalidamos + # bindings defasados, em vez de H2D. + from freetoken.core import get_global_ctx + off = getattr(get_global_ctx(), "kv_offloader", None) + new_pages = kv_new[:: self.page_size] // self.page_size + if off is not None: + tier.read_kv_to_mirror(hit.kv_pages, off.mirror.tensor, + off.mirror_scales.tensor if off.mirror_scales is not None else None, + new_pages) + off.invalidate_pages(new_pages) + else: + tier.read_kv(hit.kv_pages, new_pages) + prefix_len, mamba_exist = self.prefix_cache.insert(ids[: hit.host_len], kv_new, slot) + if prefix_len > 0: + # o tree já tinha [0, prefix_len) com as páginas antigas (dedup do insert) -- + # as páginas frescas que aloquei pra essa parte voltam pro free list + self._free(kv_new[:prefix_len]) + if mamba_exist: + pool.free([slot]) # o boundary já tinha snapshot vivo (dedup do insert) + def rebuild(self, num_pages: int, page_table: torch.Tensor) -> None: """Re-point the page table and reset page accounting + prefix cache IN PLACE. @@ -566,6 +700,12 @@ def rebuild(self, num_pages: int, page_table: torch.Tensor) -> None: self.page_table = page_table self.free_slots = torch.arange(num_pages, dtype=torch.int32, device=device) * self.page_size self.prefix_cache = self._make_prefix_cache(device, self.page_size, self.cache_type) + if self.host_tier is not None: + # the rebuilt tree gets the same hooks (the host index survives a rebuild) + self.prefix_cache.on_evict_node = self._on_evict_node_to_host + self.prefix_cache.on_match_dead_snapshot = self._rehydrate_snapshot + if self.host_tier.has_kv: + self.prefix_cache.on_evict_kv = self._on_evict_kv_to_host # The discarded hybrid tree owned donated GDN-snapshot slots; rebuild is idle-only, so # reclaim the whole LinearStatePool free-list (else those slots leak -> admission hangs). if self.is_hybrid: diff --git a/python/freetoken/scheduler/interleave.py b/python/freetoken/scheduler/interleave.py new file mode 100644 index 0000000000..76f293c8f5 --- /dev/null +++ b/python/freetoken/scheduler/interleave.py @@ -0,0 +1,73 @@ +"""Decode-interleave scheduling policy. + +Background (measured on production, 2026-09-15) +----------------------------------------------- +``Scheduler._schedule_next_batch`` picks:: + + batch = (prefill_manager.schedule_next_batch(budget) + or decode_manager.schedule_next_batch()) + +Prefill wins unconditionally. A long prompt is chunked at the window pool's +budget (384 tokens in production), so a 300k-token context is 797 consecutive +prefill steps -- and while they run, *decoding requests already in flight are +never scheduled*. Measured on the live service: + + four long cold prefills: 680/797/506/250 chunks, 267/317/190/93 s + prefill-burst share of a 15-minute window: 70% + +so an in-flight request waits up to ~5 minutes for its next token. The +scheduler already carries the marker for this:: + + # TODO: support other policies: e.g. DECODE first + +Reversing the order outright is wrong -- it starves prefill, and a request that +never prefills never starts. What is needed is *interleaving*: let prefill keep +priority, but force a decode step every N prefill steps so in-flight requests +make progress. The cost is small by construction: measured chunk cost is +0.388 s, and one decode step for a handful of requests is a fraction of that, +so every 8th step costs ~1% of prefill wall time while bounding the stall at +8 chunks (~3 s) instead of the whole burst (~300 s). + +This module holds the *policy* only -- a counter and a decision -- so it can be +unit-tested without a Scheduler, an engine, or a GPU. +""" + +from __future__ import annotations + + +class DecodeInterleavePolicy: + """Decide whether to force a decode step between prefill steps. + + ``every`` is the number of consecutive prefill steps allowed before a decode + step is forced. ``None`` or <= 0 disables the policy, reproducing the + historical prefill-first behaviour bit-for-bit. + + The counter only advances on prefill steps that were actually scheduled, so + a scheduler with nothing to prefill never accumulates credit, and the first + decode of a burst is not delayed. + """ + + def __init__(self, every: int | None = None) -> None: + self.every = int(every) if every else None + self._prefill_streak = 0 + + @property + def enabled(self) -> bool: + return self.every is not None and self.every > 0 + + def note_prefill(self) -> None: + """Record that a prefill step was scheduled in this slot.""" + self._prefill_streak += 1 + + def note_decode(self) -> None: + """Record that a decode step was scheduled in this slot (streak resets).""" + self._prefill_streak = 0 + + def wants_decode(self) -> bool: + """True when the streak has reached the threshold and a decode is due.""" + if not self.enabled: + return False + return self._prefill_streak >= self.every + + def reset(self) -> None: + self._prefill_streak = 0 diff --git a/python/freetoken/scheduler/prefill.py b/python/freetoken/scheduler/prefill.py index de4bddfb4a..9019604707 100644 --- a/python/freetoken/scheduler/prefill.py +++ b/python/freetoken/scheduler/prefill.py @@ -143,6 +143,7 @@ def _add_one_req( next_track_idx: int = 0, restore_src: int | None = None, swa_evicted_seqlen: int = 0, + last_track_seqlen: int | None = None, ) -> Req | None: remain_len = pending_req.input_len - cached_len chunk_size = min(self.token_budget, remain_len) @@ -225,6 +226,9 @@ def _add_one_req( req.linear_slot_idx = linear_slot_idx req.mamba_ping_pong = ping_pong req.mamba_next_track_idx = next_track_idx + # A short continuation cannot make a new snapshot; retain the pending one + # together with its ping-pong cursor until the final prefill commit. + req.mamba_last_track_seqlen = last_track_seqlen req.mamba_restore_src = restore_src req.swa_evicted_seqlen = swa_evicted_seqlen # carry the extend-free watermark across chunks return req @@ -242,6 +246,7 @@ def try_add_one(self, pending_req: PendingReq) -> Req | None: linear_slot_idx=chunked_req.linear_slot_idx, ping_pong=chunked_req.mamba_ping_pong, next_track_idx=chunked_req.mamba_next_track_idx, + last_track_seqlen=chunked_req.mamba_last_track_seqlen, restore_src=None, # continuation chunk already has live state swa_evicted_seqlen=chunked_req.swa_evicted_seqlen, # extend-free watermark so far ) diff --git a/python/freetoken/scheduler/scheduler.py b/python/freetoken/scheduler/scheduler.py index d057d3b908..cb6085d344 100644 --- a/python/freetoken/scheduler/scheduler.py +++ b/python/freetoken/scheduler/scheduler.py @@ -1,6 +1,7 @@ from __future__ import annotations import os +import time from typing import TYPE_CHECKING, List, NamedTuple, NoReturn, Set, Tuple, TypeAlias @@ -31,6 +32,7 @@ from .cache import CacheManager from .config import SchedulerConfig from .decode import DecodeManager +from .interleave import DecodeInterleavePolicy from .io import SchedulerIOMixin from .mm import cut_image_spans, plan_mm_batch from .prefill import ChunkedReq, PrefillManager @@ -91,6 +93,14 @@ def __init__(self, config: SchedulerConfig): ) or getattr(self.engine.kv_cache, "sliding_window_size", None), ) self.decode_manager = DecodeManager(config.page_size) + # Bound how long a run of prefill steps may starve in-flight decodes. A long + # prompt chunked at the window budget is hundreds of consecutive prefill steps; + # measured on an 8xRTX4090 DSV4 deployment, that left an already-decoding + # request unscheduled for 267-317 s. Unset reproduces the historical + # prefill-first order exactly. + self._interleave = DecodeInterleavePolicy( + getattr(config, "decode_interleave_every", None) + ) self._bidirectional_mm = any(getattr(g, "bidirectional_mm_blocks", False) for g in config.model_config.attention_groups) self.prefill_manager = PrefillManager( self.cache_manager, @@ -111,6 +121,11 @@ def __init__(self, config: SchedulerConfig): # tombstone so an abort-before-admission request can never be resurrected after its # terminal accounting acknowledgement has already been published. self._abort_tombstones: dict[int, None] = {} + # uid -> monotonic clock at admission, for the per-request prefill span shipped on the + # first sampled token. An entry existing IS the "this is the first token" test, so it is + # popped there; _free_req_resources pops too, so a request aborted mid-prefill (which + # never samples) cannot leave one behind. + self._prefill_start: dict[int, float] = {} self._forward_iter = 0 # global forward counter; drives the SWA proactive-eviction cadence # The launched-but-not-yet-drained batch (overlap): set at the top of each overlap_loop # iteration so the abort handler can tell whether a request's forward is still in flight @@ -313,8 +328,9 @@ def _process_last_data(self, last_data: ForwardData | None) -> None: if last_data is None: return - batch, (_, next_tokens_cpu, copy_done) = last_data[0].batch, last_data[1] - copy_done.synchronize() + batch, outputs = last_data[0].batch, last_data[1] + next_tokens_cpu = outputs.next_tokens_cpu + outputs.copy_done_event.synchronize() reply: List[DetokenizeMsg] = [] new_finished_reqs: Set[Req] = set() with self.cache_manager.lazy_free_region(): @@ -347,13 +363,26 @@ def _process_last_data(self, last_data: ForwardData | None) -> None: next_token = next_tokens_cpu[i] req.append_host(next_token.unsqueeze(0)) next_token = int(next_token.item()) + + row_chosen_logprob: float | None = None + row_top_ids: list[int] | None = None + row_top_logprobs: list[float] | None = None + if req.sampling_params.logprobs and outputs.chosen_logprobs_cpu is not None: + row_chosen_logprob = float(outputs.chosen_logprobs_cpu[i].item()) + requested_top = req.sampling_params.top_logprobs + if requested_top > 0 and outputs.top_ids_cpu is not None: + row_top_ids = [int(t) for t in outputs.top_ids_cpu[i, :requested_top].tolist()] + row_top_logprobs = outputs.top_logprobs_cpu[i, :requested_top].tolist() + else: + row_top_ids = [] + row_top_logprobs = [] # EOS / stop-string -> "stop", output budget exhausted -> "length"; # EOS and stop strings win over length. # Overlap can advance device_len one token ahead of this host reply. hit_length = req.input_ids.numel() >= req.max_device_len hit_eos = ( not req.sampling_params.ignore_eos and next_token in self.eos_token_ids - ) + ) or next_token in req.sampling_params.stop_token_ids matched_stop = ( self._match_stop_str(req) if not hit_eos and req.sampling_params.stop_strs @@ -371,6 +400,10 @@ def _process_last_data(self, last_data: ForwardData | None) -> None: and not finished ): req.toolcall_anchor_len = req.input_ids.numel() + # Present only until this request's first token is sampled, i.e. off its last + # prefill chunk -- so popping it both ends the prefill span and marks this as + # the one reply that carries it. + started_at = self._prefill_start.pop(req.uid, None) reply.append( DetokenizeMsg( uid=req.uid, @@ -379,6 +412,15 @@ def _process_last_data(self, last_data: ForwardData | None) -> None: finish_reason=finish_reason, matched_stop=matched_stop, stop_strs=req.sampling_params.stop_strs or None, + keep_stop_str=req.sampling_params.include_stop_str_in_output, + skip_special_tokens=req.sampling_params.skip_special_tokens, + prefill_ms=( + 0.0 if started_at is None + else (time.monotonic() - started_at) * 1000.0 + ), + chosen_logprob=row_chosen_logprob, + top_ids=row_top_ids, + top_logprobs=row_top_logprobs, ) ) @@ -429,16 +471,38 @@ def _process_last_data(self, last_data: ForwardData | None) -> None: self.send_result(reply) def _match_stop_str(self, req: Req) -> str | None: - """First stop string present in this request's generated tail, else None. Decodes - only a short suffix (bounded by the longest stop string's char length, so a stop of - N chars spans at most N tokens) to keep the per-step cost small.""" + """Match stops against the same incrementally decoded text as the frontend.""" + from freetoken.tokenizer.detokenize import DecodeStatus + stop_strs = req.sampling_params.stop_strs prompt_len = req.max_device_len - req.output_len - if len(req.input_ids) <= prompt_len: + end = len(req.input_ids) + # The frontend omits a terminal EOS, including at the output limit with ignore_eos. + # The limit is the delivered length, as in hit_length: under overlap device_len (and + # so can_decode) runs a step ahead of the host. + if ( + end > prompt_len + and end >= req.max_device_len + and int(req.input_ids[-1]) in self.eos_token_ids + ): + end -= 1 + if end <= prompt_len: return None + if req.stop_decode_status is None: + req.stop_decode_status = DecodeStatus( + decoded_ids=[], + decoded_str="", + read_offset=0, + surr_offset=0, + sent_offset=0, + skip_special=req.sampling_params.skip_special_tokens, + ) + state = req.stop_decode_status + state.decoded_ids.extend(req.input_ids[prompt_len + len(state.decoded_ids) : end].tolist()) + tail = state.decode(self.tokenizer) max_chars = max(len(s) for s in stop_strs) - tail_start = max(prompt_len, len(req.input_ids) - (max_chars + 1)) - tail = self.tokenizer.decode(req.input_ids[tail_start:].tolist()) + # Bound the text used for matching, retaining token context for incomplete characters. + state.decoded_str = state.decoded_str[-max_chars:] for s in stop_strs: if s in tail: return s @@ -512,10 +576,17 @@ def _process_one_msg(self, msg: BaseBackendMsg) -> None: ) return input_len, max_seq_len = len(msg.input_ids), self.engine.max_seq_len - max_output_len = max_seq_len - input_len + # max_seq_len is the model's advertised context, which can far exceed the + # KV pool actually allocated (see issue #111): a prompt that passes this + # check but can never be granted enough pages is queued forever with no + # error and no log line. Clamp admission to the real pool so oversized + # prompts fail loudly with the same error clients already understand. + pool_tokens = self.engine.num_pages * self.config.page_size + effective_max = min(max_seq_len, pool_tokens) + max_output_len = effective_max - input_len if max_output_len <= 0: logger.warning_rank0( - f"Input sequence length {input_len} exceeds {max_seq_len}, " + f"Input sequence length {input_len} exceeds {effective_max}, " f"request {msg.uid} is dropped." ) # Tell the client instead of dropping silently — otherwise its wait_for_ack @@ -527,7 +598,7 @@ def _process_one_msg(self, msg: BaseBackendMsg) -> None: # "prompt is too long: N tokens > M" is the phrasing Claude Code and # OpenClaw match on; the Anthropic wire has no error code to read. error=( - f"prompt is too long: {input_len} tokens > {max_seq_len} maximum " + f"prompt is too long: {input_len} tokens > {effective_max} maximum " f"(prompt + generation); shorten the prompt or increase the KV " f"cache budget" ), @@ -542,6 +613,11 @@ def _process_one_msg(self, msg: BaseBackendMsg) -> None: logger.warning_rank0( f"Adjust max_tokens to {max_output_len} for request {msg.uid}." ) + sp = msg.sampling_params + if sp.min_tokens > 0: + # min_tokens: the sampler masks these ids while the output is shorter. + sp.min_tokens = min(sp.min_tokens, sp.max_tokens) + sp.min_tokens_stop_ids = sorted(set(self.eos_token_ids) | set(sp.stop_token_ids)) self.prefill_manager.add_one_req(msg) elif isinstance(msg, AbortBackendMsg): logger.debug_rank0("Aborting request %d", msg.uid) @@ -632,6 +708,9 @@ def _free_req_resources(self, req: Req) -> None: # slots to two later requests. table_idx == -1 marks an already-freed request. if req.table_idx == -1: return + # A request freed without ever sampling (aborted mid-prefill) would otherwise keep its + # admission timestamp forever; a normally-finished one popped it at its first token. + self._prefill_start.pop(req.uid, None) # Polymorphic free: the DSV4 manager returns the request's window pages + cmp/idx blocks # to their tier free-lists; the generic manager frees its KV pages (it reads # page_table[req.table_idx], so free the table entry after). @@ -861,11 +940,28 @@ def _gather_multimodal(self, batch: Batch) -> None: ) def _schedule_next_batch(self) -> ForwardInput | None: - # TODO: support other policies: e.g. DECODE first - batch = ( - self.prefill_manager.schedule_next_batch(self.prefill_budget) - or self.decode_manager.schedule_next_batch() - ) + # Prefill keeps priority -- a request that never prefills never starts -- but a + # long run of prefill chunks must not starve in-flight decodes. When the policy + # says a decode is due, take one if a decode is actually runnable; otherwise stay + # with prefill (never spend a slot idle just to keep a promise). + # getattr, not attribute access: the accounting tests drive this method on a + # partially built Scheduler, and interleaving must be an opt-in that a stripped + # object simply does not have rather than something that breaks it. A missing + # policy is the historical prefill-first order. + policy = getattr(self, "_interleave", None) + batch = None + if policy is not None and policy.wants_decode(): + batch = self.decode_manager.schedule_next_batch() + if batch is not None: + policy.note_decode() + if batch is None: + batch = self.prefill_manager.schedule_next_batch(self.prefill_budget) + if batch is not None and policy is not None: + policy.note_prefill() + if batch is None: + batch = self.decode_manager.schedule_next_batch() + if batch is not None and policy is not None: + policy.note_decode() if batch is None: return None forward_input = self._prepare_batch(batch) @@ -880,6 +976,11 @@ def _report_prompt_admissions(self, batch: Batch) -> None: """ if not batch.is_prefill or not batch.prompt_admissions: return + # Before the forward, which is what makes this the prefill START: _schedule_next_batch + # has prepared the batch but not run it. + now = time.monotonic() + for uid, _, _ in batch.prompt_admissions: + self._prefill_start[uid] = now self.send_result( [ PromptAdmittedMsg(uid=uid, prompt_tokens=prompt_tokens, cached_tokens=cached_tokens) diff --git a/python/freetoken/server/anthropic_api.py b/python/freetoken/server/anthropic_api.py index d9a6972f09..35d902b9d6 100644 --- a/python/freetoken/server/anthropic_api.py +++ b/python/freetoken/server/anthropic_api.py @@ -44,15 +44,18 @@ ToolCallArgsDelta, ToolCallsDelta, ToolCallStart, + build_metrics, count_prompt_tokens, generate_events, generate_full, + metrics_enabled, render_messages, resolve_sampling, split_tool_lists, submit_generation, with_keepalive, ) +from .openai_api import _await_watching_disconnect from .request_logger import log_request # Emit a protocol-native `ping` event after this many seconds of stream silence, @@ -121,26 +124,35 @@ async def handle_anthropic_messages( getattr(state.config, "max_output_tokens", None) or DEFAULT_MAX_OUTPUT_TOKENS ), inline_system_policy=getattr(state.config, "anthropic_inline_system", "auto"), + default_thinking_mode=getattr(state.config, "default_thinking_mode", "auto"), ) uid = await submit_generation(spec, state) except ValueError as exc: return _anthropic_error_response(400, "invalid_request_error", str(exc)) cache_report = getattr(state.config, "enable_cache_report", False) + metrics = metrics_enabled(state) if req.stream: events = anthropic_event_stream( generate_events(uid, spec, state, source="/v1/messages"), - req.model, uid, cache_report=cache_report, + req.model, uid, cache_report=cache_report, metrics=metrics, ) if request is not None: events = state.stream_with_cancellation(events, request, uid) return StreamingResponse(events, media_type="text/event-stream") try: - result = await generate_full(uid, spec, state, source="/v1/messages") + # an abandoned request must not keep decoding to max_tokens (#222's watcher) + result = await _await_watching_disconnect( + generate_full(uid, spec, state, source="/v1/messages"), request, state, [uid] + ) except GenerationError as exc: return _anthropic_error_response(400, "invalid_request_error", str(exc)) - response = anthropic_full_response(result, req.model, uid, cache_report=cache_report) + if result is None: + return _anthropic_error_response(499, "api_error", "client disconnected before the response was ready") + response = anthropic_full_response( + result, req.model, uid, cache_report=cache_report, metrics=metrics + ) return JSONResponse(content=response.model_dump(exclude_none=True)) @@ -154,7 +166,9 @@ async def handle_anthropic_count_tokens(req: AnthropicCountTokensRequest, state: # so it must not fall into the convert/empty-prompt ValueError branch. try: messages, template_tools, _, ctk = convert_anthropic_prompt( - req, reasoning_parser=getattr(state.config, "reasoning_parser", None) + req, + reasoning_parser=getattr(state.config, "reasoning_parser", None), + default_thinking_mode=getattr(state.config, "default_thinking_mode", "auto"), ) except ValueError as exc: return _anthropic_error_response(400, "invalid_request_error", str(exc)) @@ -186,6 +200,7 @@ async def handle_anthropic_count_tokens(req: AnthropicCountTokensRequest, state: def convert_anthropic_prompt( req: AnthropicMessagesRequest | AnthropicCountTokensRequest, reasoning_parser: str | None = None, + default_thinking_mode: str | None = None, ) -> tuple[list[dict[str, Any]], list[dict[str, Any]] | None, list[dict[str, Any]] | None, dict[str, Any]]: """(messages, template_tools, parser_tools, chat_template_kwargs) — the prompt side of the conversion, shared by /v1/messages and /v1/messages/count_tokens so @@ -298,7 +313,7 @@ def convert_anthropic_prompt( # Native extended-thinking toggle -> template kwargs, broadcast in every # spelling the ecosystem's templates read (a bare enable_thinking bool is # inert for templates that read a different knob, e.g. M3's thinking_mode). - from .model_meta import thinking_toggle_kwargs + from .model_meta import apply_default_thinking_mode, thinking_toggle_kwargs ctk: dict[str, Any] = {} if req.thinking: @@ -306,8 +321,13 @@ def convert_anthropic_prompt( ctk = thinking_toggle_kwargs(True) elif req.thinking.get("type") == "disabled": ctk = thinking_toggle_kwargs(False) + else: + # Any thinking block (e.g. "adaptive") is the client's explicit choice; the server + # default only covers requests that leave thinking out. + ctk = apply_default_thinking_mode(ctk, default_thinking_mode) - return render_messages(messages), template_tools, parser_tools, ctk + # late system reminders stay in place: the tokenizer folds them instead (prefix reuse) + return render_messages(messages, hoist_system=False), template_tools, parser_tools, ctk def convert_anthropic_to_genspec( @@ -316,9 +336,10 @@ def convert_anthropic_to_genspec( reasoning_parser: str | None = None, default_max_tokens: int = DEFAULT_MAX_OUTPUT_TOKENS, inline_system_policy: str = "auto", + default_thinking_mode: str | None = None, ) -> GenSpec: messages, template_tools, parser_tools, ctk = convert_anthropic_prompt( - req, reasoning_parser=reasoning_parser + req, reasoning_parser=reasoning_parser, default_thinking_mode=default_thinking_mode ) return GenSpec( messages=messages, @@ -384,7 +405,7 @@ def _tool_result_parts(content) -> tuple[str, list[dict[str, Any]]]: # Output formatting: GenResult / GenEvent -> Anthropic response / events # --------------------------------------------------------------------------- # def anthropic_full_response( - result: GenResult, model: str, uid: int, cache_report: bool = False + result: GenResult, model: str, uid: int, cache_report: bool = False, metrics: bool = False ) -> AnthropicMessagesResponse: content: list[AnthropicContentBlock] = [] if result.reasoning: @@ -411,6 +432,12 @@ def anthropic_full_response( usage=_anthropic_usage( result.prompt_tokens, result.completion_tokens, result.cached_tokens, cache_report ), + metrics=build_metrics( + prompt_tokens=result.prompt_tokens, + completion_tokens=result.completion_tokens, + cached_tokens=result.cached_tokens, + timings=result.timings, + ) if metrics else None, ) @@ -429,7 +456,8 @@ def _anthropic_usage( async def anthropic_event_stream( - events: AsyncIterator[Any], model: str, uid: int, cache_report: bool = False + events: AsyncIterator[Any], model: str, uid: int, cache_report: bool = False, + metrics: bool = False, ) -> AsyncIterator[str]: """Format the protocol-neutral GenEvent stream into Anthropic SSE events. @@ -453,7 +481,9 @@ def _open_thinking() -> str: block_open = "thinking" return _event(AnthropicStreamEvent( type="content_block_start", index=block_index, - content_block=AnthropicContentBlock(type="thinking", thinking=""), + # The SDK types signature as a required string on the opening block, before any + # signature_delta. Local reasoning is unsigned, like the full-response path. + content_block=AnthropicContentBlock(type="thinking", thinking="", signature=""), )) def _stop_block() -> list[str]: @@ -585,6 +615,12 @@ def _tool_args_delta(fragment: str) -> str: usage=_anthropic_usage( ev.prompt_tokens, ev.completion_tokens, ev.cached_tokens, cache_report ), + metrics=build_metrics( + prompt_tokens=ev.prompt_tokens, + completion_tokens=ev.completion_tokens, + cached_tokens=ev.cached_tokens, + timings=ev.timings, + ) if metrics else None, )) yield _event(AnthropicStreamEvent(type="message_stop")) # Anthropic streams terminate on message_stop — no OpenAI-style diff --git a/python/freetoken/server/anthropic_models.py b/python/freetoken/server/anthropic_models.py index 42ab545bec..2a6232ee6f 100644 --- a/python/freetoken/server/anthropic_models.py +++ b/python/freetoken/server/anthropic_models.py @@ -153,6 +153,10 @@ class AnthropicMessagesResponse(BaseModel): ) = None stop_sequence: str | None = None usage: AnthropicUsage | None = None + # FreeToken extension, served only under --enable-metrics-report: per-request inference + # timings. Anthropic's wire has no such field, so it stays absent by default rather than + # being emitted as null -- both response paths dump with exclude_none. + metrics: dict[str, Any] | None = None def model_post_init(self, __context): if not self.id: @@ -176,3 +180,5 @@ class AnthropicStreamEvent(BaseModel): index: int | None = None error: AnthropicError | None = None usage: AnthropicUsage | None = None + # Rides message_delta next to usage, under --enable-metrics-report only. + metrics: dict[str, Any] | None = None diff --git a/python/freetoken/server/api_models.py b/python/freetoken/server/api_models.py index 6c3ade0924..21b1d1f860 100644 --- a/python/freetoken/server/api_models.py +++ b/python/freetoken/server/api_models.py @@ -76,6 +76,12 @@ class ChatCompletionRequest(BaseModel): stop: str | list[str] | None = None presence_penalty: float = Field(default=0.0, ge=-2.0, le=2.0, allow_inf_nan=False) frequency_penalty: float = Field(default=0.0, ge=-2.0, le=2.0, allow_inf_nan=False) + # Sampled-token logprobs (OpenAI chat semantics): top_logprobs (0..20) requires + # logprobs=true; entries cover generated tokens only (no prompt logprobs). + # Fail-closed (400) when a reasoning parser or tool parsing is active: entries + # must align 1:1 with visible content tokens (see _chat_logprobs_conflict). + logprobs: bool = False + top_logprobs: int | None = None chat_template_kwargs: dict[str, Any] = Field(default_factory=dict) reasoning_effort: str | None = None # DeepSeek-wire thinking toggle ({"type": "enabled"|"disabled"}). Any so a @@ -90,11 +96,30 @@ class ChatCompletionRequest(BaseModel): function_call: Any | None = None logit_bias: dict[str, float] | None = None response_format: dict[str, Any] | None = None + # The last message is an assistant prefix the model continues (no generation prompt). + continue_final_message: bool = False + # ---- sampling extras shared with vLLM / SGLang / llama.cpp (all optional) ---- + min_p: float | None = None + repetition_penalty: float | None = None + min_tokens: int = 0 + stop_token_ids: list[int] | None = None + include_stop_str_in_output: bool = False + skip_special_tokens: bool | None = None + # Accepted for OpenAI compatibility, not honoured: the batched sampling kernel has no + # per-request generator (seed); user / metadata / store are opaque to the server. + seed: int | None = None + user: str | None = None + # A client-chosen id echoed as the response id (vLLM request_id / SGLang rid). + request_id: str | None = None @model_validator(mode="after") def _sync_max_completion_tokens(self) -> "ChatCompletionRequest": if self.max_completion_tokens is not None: self.max_tokens = self.max_completion_tokens + if self.continue_final_message and ( + not self.messages or self.messages[-1].role != "assistant" + ): + raise ValueError("continue_final_message needs a final assistant message") return self @@ -120,6 +145,19 @@ class CompletionRequest(BaseModel): suffix: str | None = None logit_bias: dict[str, float] | None = None response_format: dict[str, Any] | None = None + # ---- sampling extras shared with vLLM / SGLang / llama.cpp (all optional) ---- + min_p: float | None = None + repetition_penalty: float | None = None + min_tokens: int = 0 + stop_token_ids: list[int] | None = None + include_stop_str_in_output: bool = False + skip_special_tokens: bool | None = None + # Accepted for OpenAI compatibility, not honoured: the batched sampling kernel has no + # per-request generator (seed); user / metadata / store are opaque to the server. + seed: int | None = None + user: str | None = None + # A client-chosen id echoed as the response id (vLLM request_id / SGLang rid). + request_id: str | None = None @model_validator(mode="after") def _sync_max_completion_tokens(self) -> "CompletionRequest": @@ -128,6 +166,32 @@ def _sync_max_completion_tokens(self) -> "CompletionRequest": return self +MAX_N = 16 # choices per request; each is a full generation, so keep the fan-out bounded + + +class TokenizeRequest(BaseModel): + """``POST /tokenize`` (the vLLM / SGLang shape): a raw prompt, or chat messages that + are rendered through the model's chat template first.""" + + model_config = ConfigDict(extra="allow") + + model: str | None = None + prompt: str | None = None + messages: list[Message] | None = None + add_special_tokens: bool = True + add_generation_prompt: bool = True + return_token_strs: bool = False + chat_template_kwargs: dict[str, Any] = Field(default_factory=dict) + + +class DetokenizeRequest(BaseModel): + model_config = ConfigDict(extra="allow") + + model: str | None = None + tokens: list[int] + skip_special_tokens: bool = False + + class ModelCard(BaseModel): id: str object: str = "model" diff --git a/python/freetoken/server/api_server.py b/python/freetoken/server/api_server.py index a9b16adfd7..29ba88e1c7 100644 --- a/python/freetoken/server/api_server.py +++ b/python/freetoken/server/api_server.py @@ -35,7 +35,7 @@ init_logger, load_generation_sampling, ) -from pydantic import BaseModel +from pydantic import BaseModel, Field from starlette.datastructures import Headers from .args import ServerArgs @@ -60,6 +60,10 @@ # shutdown is treated as expected — no ERROR log, no "failed" latch. See run_backend_supervisor. _SHUTTING_DOWN = threading.Event() BACKEND_DEATH_EXIT_GRACE_S = 10.0 +# The longest a stream stays silent before it sends an SSE comment. A long prompt prefills for +# minutes before its first token, and clients time an idle body out (Node's fetch after 300 s). +STREAM_KEEPALIVE_SECONDS = 15.0 +_STREAM_KEEPALIVE = b": keep-alive\n\n" def get_global_state() -> FrontendManager: @@ -140,6 +144,13 @@ class FrontendManager: initialized: bool = False ack_map: Dict[int, List[UserReply]] = field(default_factory=dict) event_map: Dict[int, asyncio.Event] = field(default_factory=dict) + # Abort claims live HERE, not in the maps above: a disconnect cancels the response + # task at its innermost await (inside wait_for_ack, whose finally empties both maps + # on the way out), so by the time any abort path runs, a map-based claim would no-op + # exactly when the AbortMsg matters most. Insertion-ordered, FIFO-bounded like the + # scheduler's abort tombstones; a duplicate AbortMsg for an evicted or finished uid + # is safe (the scheduler acks aborts for uids it no longer has). + aborted_uids: Dict[int, None] = field(default_factory=dict) # Stable identity for this serve process. Generated before the backend is ready so every # /health state (loading/ok/error) and /v1/stats can identify the same engine generation. instance_id: str = field(default_factory=lambda: str(uuid.uuid4())) @@ -361,24 +372,56 @@ async def stream_generate(self, uid: int): yield b"data: [DONE]\n\n" logger.debug("Finished streaming response for user %s", uid) - async def stream_with_cancellation(self, generator, request: Request, uid: int): + async def stream_with_cancellation(self, generator, request: Request, uid): + """``uid`` is one request id or a sequence of them (an ``n > 1`` fan-out streams + several generations into one response): a disconnect aborts every one. While the + generator is silent (a long prefill), send an SSE comment every + ``STREAM_KEEPALIVE_SECONDS`` and check whether the client has gone.""" + uids = list(uid) if isinstance(uid, (list, tuple)) else [uid] + finished = False + chunks = generator.__aiter__() + pending: asyncio.Future | None = None try: - async for chunk in generator: - # detect if the client has disconnected + while True: + if pending is None: + pending = asyncio.ensure_future(chunks.__anext__()) + done, _ = await asyncio.wait({pending}, timeout=STREAM_KEEPALIVE_SECONDS) if await request.is_disconnected(): - logger.info("Client disconnected for user %s", uid) + logger.info("Client disconnected for user %s", uids) raise asyncio.CancelledError + if not done: + yield _STREAM_KEEPALIVE + continue + try: + chunk = pending.result() + except StopAsyncIteration: + finished = True + return + pending = None yield chunk - except asyncio.CancelledError: - asyncio.create_task(self.abort_user(uid)) - raise + finally: + if pending is not None and not pending.done(): + pending.cancel() + # finally, not `except CancelledError`: the abort must go out however the + # stream dies -- the server cancelling the response task, a late close + # delivering GeneratorExit, or an exception out of the generator. Only a + # stream that ran to completion leaves the engine with nothing to stop. + if not finished: + for one in uids: + try: + await asyncio.shield(self.abort_user(one)) + except Exception: # noqa: BLE001 + logger.exception("Failed to deliver abort for user %s", one) async def abort_user(self, uid: int): + if uid in self.aborted_uids: + return + self.aborted_uids[uid] = None + while len(self.aborted_uids) > 65_536: + self.aborted_uids.pop(next(iter(self.aborted_uids))) + self.event_map.pop(uid, None) + self.ack_map.pop(uid, None) await asyncio.sleep(0.1) - if uid in self.ack_map: - del self.ack_map[uid] - if uid in self.event_map: - del self.event_map[uid] self.stats.on_abort(uid) logger.warning("Aborting request for user %s", uid) await self.send_one(AbortMsg(uid=uid)) @@ -482,38 +525,57 @@ def _served_model_name() -> str | None: return getattr(cfg, "served_model_name", None) -@app.middleware("http") -async def _record_request_middleware(request: Request, call_next): +class _RecordRequestMiddleware: """Time every generation request into the ring for /v1/requests + /v1/stats p95. Single- model server, so model = served_model_name; stream is inferred from the response media - type. Token counts are P3 (SSE usage arrives after the handler returns) — kept as None.""" - path = request.url.path - if path.startswith(_UNTRACKED_REQUEST_PREFIXES) or not path.startswith( - _TRACKED_REQUEST_PREFIXES - ): - return await call_next(request) - import time as _time - - start = _time.monotonic() - response = await call_next(request) - duration_ms = int((_time.monotonic() - start) * 1000) - ctype = response.headers.get("content-type", "") - request_ring.record_request( - request_ring.RequestRecord( - ts=_time.strftime("%Y-%m-%dT%H:%M:%SZ", _time.gmtime()), - method=request.method, - path=path, - status=response.status_code, - model=_served_model_name(), - duration_ms=duration_ms, - ttft_ms=None, - prompt_tokens=None, - completion_tokens=None, - stream=ctype.startswith("text/event-stream"), - error=None, - ) - ) - return response + type. Token counts are P3 (SSE usage arrives after the handler returns) — kept as None. + + Pure ASGI on purpose: Starlette's ``@app.middleware("http")`` (BaseHTTPMiddleware) hands the + handler a wrapped receive channel that never yields the client's ``http.disconnect`` while a + non-streaming request is still computing, so ``request.is_disconnected()`` stays False and an + abandoned request runs to max_tokens. The recorder is the only such middleware here.""" + + def __init__(self, app): + self.app = app + + async def __call__(self, scope, receive, send): + path = scope.get("path", "") if scope["type"] == "http" else "" + if not path.startswith(_TRACKED_REQUEST_PREFIXES) or path.startswith(_UNTRACKED_REQUEST_PREFIXES): + await self.app(scope, receive, send) + return + import time as _time + + start = _time.monotonic() + + async def send_recording(message): + # the base-class version recorded once the handler returned its response object, + # i.e. at response start, before a streaming body -- keep that timing + if message["type"] == "http.response.start": + ctype = "" + for name, value in message.get("headers", ()): + if name.lower() == b"content-type": + ctype = value.decode("latin-1") + request_ring.record_request( + request_ring.RequestRecord( + ts=_time.strftime("%Y-%m-%dT%H:%M:%SZ", _time.gmtime()), + method=scope["method"], + path=path, + status=message["status"], + model=_served_model_name(), + duration_ms=int((_time.monotonic() - start) * 1000), + ttft_ms=None, + prompt_tokens=None, + completion_tokens=None, + stream=ctype.startswith("text/event-stream"), + error=None, + ) + ) + await send(message) + + await self.app(scope, receive, send_recording) + + +app.add_middleware(_RecordRequestMiddleware) class CacheRebuildRequest(BaseModel): @@ -533,7 +595,7 @@ class CacheRebuildRequest(BaseModel): # is deferred (needs the drain-gate machinery); constraining the Literal makes an # unsupported value fail fast with a 422 at the API layer instead of a generic 503. mode: Literal["if_idle"] = "if_idle" - timeout: float = 300.0 + timeout: float = Field(default=300.0, gt=0, le=3600) async def dispatch_rebuild( @@ -754,6 +816,30 @@ def _reasoning_geometry(state: Any) -> dict | None: return {"gears": list(gears), "default": default, "kwargs": kwargs} +def kv_pool_geometry(state: Any) -> tuple[int, int]: + """The allocated KV pool as ``(num_pages, page_size)``, most-recent-truth first: the last + rebuild's result, else the running UserReply snapshot, else the load-time ("meta", …) ack. + ``num_pages`` is 0 until one of those lands (still loading / older engine build). + + Split out of cache_geometry so /v1/models can clamp the advertised context to the pool + without recomputing the whole panel payload -- both callers must agree on the resolution + order or they would report different capacities for the same server. + """ + stats = getattr(state, "stats", None) + last = getattr(state, "last_rebuild", None) or {} + pools = getattr(state, "cache_pools", None) or {} + num_pages = int( + last.get("num_pages") + or getattr(stats, "kv_total_pages", 0) + or pools.get("num_pages", 0) + or 0 + ) + page_size = int( + pools.get("page_size", 0) or getattr(getattr(state, "config", None), "page_size", 1) or 1 + ) + return num_pages, page_size + + def cache_geometry(state: Any) -> dict: """Current cache geometry for the desktop cache panel. Each pool size resolves most-recent-truth first: the last rebuild's result, else the running UserReply snapshot @@ -769,11 +855,10 @@ def cache_geometry(state: Any) -> dict: config = state.config last = getattr(state, "last_rebuild", None) or {} pools = getattr(state, "cache_pools", None) or {} - num_pages = int(last.get("num_pages") or tr.kv_total_pages or pools.get("num_pages", 0) or 0) + num_pages, page_size = kv_pool_geometry(state) num_mamba_slots = int( last.get("mamba_slots") or tr.mamba_total_slots or pools.get("num_mamba_slots", 0) or 0 ) - page_size = int(pools.get("page_size", 0) or getattr(config, "page_size", 1) or 1) moe_cache_size = last.get("moe_cache_size") if moe_cache_size is None: moe_cache_size = int(pools.get("moe_cache_size", 0) or 0) or configured_moe_cache_size( diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py index 519447081b..41b8d05092 100644 --- a/python/freetoken/server/args.py +++ b/python/freetoken/server/args.py @@ -51,6 +51,12 @@ class ServerArgs(SchedulerConfig): # Reasoning parser that splits reasoning from content for OpenAI # responses. None disables it (default for models without a reasoning protocol). reasoning_parser: str | None = None + # Server-wide default thinking mode for reasoning-capable models: "auto" keeps the + # current per-request behavior; "chat" turns thinking off for every request that does + # not choose itself (template kwargs or a protocol-level effort/thinking field), for + # OpenAI-compatible clients that never send template kwargs, like Vercel AI SDK or + # llama-swap; "thinking" turns it on the same way. + default_thinking_mode: str = "auto" # "model": fill unspecified request sampling params from generation_config.json # (temperature/top_k/top_p), like sglang. "none": use framework defaults only. sampling_defaults: str = "model" @@ -62,6 +68,10 @@ class ServerArgs(SchedulerConfig): # input_tokens_details.cached_tokens). Mirrors sglang's --enable-cache-report. enable_cache_report: bool = False anthropic_inline_system: str = "auto" + # Serve a per-request `metrics` object (TTFT, prefill/decode times and throughputs, + # prefix-cache hit) alongside usage. Off by default: it is a non-standard field on + # every protocol we speak. + enable_metrics_report: bool = False # Comma-separated hostname allowlist for client-supplied image URLs; empty admits any domain. allowed_media_domains: str = "" # Directory file:// image refs may be read from; empty rejects local files. @@ -102,10 +112,6 @@ def backend_create_detokenizer_link(self) -> bool: def frontend_create_tokenizer_link(self) -> bool: return not self.share_tokenizer - @property - def distributed_addr(self) -> str: - return f"tcp://127.0.0.1:{self.server_port + 1}" - def _json_object(text: str) -> dict: try: @@ -134,6 +140,7 @@ def parse_args( from freetoken.attention import validate_attn_backend from freetoken.kvcache import SUPPORTED_CACHE_MANAGER from freetoken.moe import MOE_STRATEGIES + from freetoken.moe.pool_policy import SCOPES def _parse_quant_backend(value: str) -> str: from freetoken.layers.quantization import QuantBackend @@ -153,6 +160,27 @@ def _parse_moe_cache_rate(value: str) -> float: raise argparse.ArgumentTypeError("must be in [0, 1]") return rate + def _parse_mem_size(value: str) -> int: + units = { + "b": 1, "k": 1 << 10, "kb": 1 << 10, "kib": 1 << 10, "m": 1 << 20, "mb": 1 << 20, "mib": 1 << 20, + "g": 1 << 30, "gb": 1 << 30, "gib": 1 << 30, + } + text = value.strip().lower() + number = text.rstrip("abgikm") + unit = text[len(number):] + try: + size = int(float(number) * units[unit]) + except (KeyError, ValueError, OverflowError) as exc: + raise argparse.ArgumentTypeError("must be a size with a unit, e.g. 12GiB or 800MiB") from exc + if size <= 0: + raise argparse.ArgumentTypeError("must be positive") + return size + + def _parse_cache_split(value: str) -> str: + if value == "equal" or (value.startswith("profile:") and len(value) > len("profile:")): + return value + raise argparse.ArgumentTypeError("must be equal or profile:") + def _positive_int(value: str) -> int: try: n = int(value) @@ -162,6 +190,15 @@ def _positive_int(value: str) -> int: raise argparse.ArgumentTypeError("must be >= 1") return n + def _valid_port(value: str) -> int: + try: + n = int(value) + except ValueError as exc: + raise argparse.ArgumentTypeError("must be an integer") from exc + if not 0 <= n <= 65535: + raise argparse.ArgumentTypeError("must be between 0 and 65535") + return n + def _lazy_gpu_arg(value: str) -> tuple[str, ...]: from freetoken.gpu_select import gpu_arg @@ -379,6 +416,18 @@ def _infer_reasoning_parser(model_path: str) -> str | None: "unset serves without authentication.", ) + parser.add_argument( + "--dist-port", + type=_valid_port, + dest="distributed_port", + default=None, + help=( + "Port for the internal TP rendezvous store, loopback-only regardless of --host. " + "Defaults to --port + 1; override when that collides with another instance or " + "service on the same host." + ), + ) + parser.add_argument( "--cuda-graph-max-bs", "--graph", @@ -432,6 +481,20 @@ def _infer_reasoning_parser(model_path: str) -> str | None: ), ) + parser.add_argument( + "--decode-interleave-every", + type=int, + default=ServerArgs.decode_interleave_every, + help=( + "Force one decode step after this many consecutive prefill steps. A long " + "prompt is prefilled in chunks, so it occupies that many consecutive " + "scheduler steps; measured on an 8xRTX4090 DSV4 deployment, a 300k-token " + "context left an already-decoding request unscheduled for 267-317 s. " + "Unset keeps the historical prefill-first order (the scheduler's own TODO " + "names this policy)." + ), + ) + parser.add_argument( "--page-size", type=int, @@ -439,6 +502,23 @@ def _infer_reasoning_parser(model_path: str) -> str | None: help="Set the page size for system management.", ) + parser.add_argument( + "--kv-cache-dtype", + dest="kv_quant", + type=str, + default=ServerArgs.kv_quant, + choices=["auto", "bf16", "fp8", "nvfp4"], + help=( + "KV-cache storage format. 'bf16' (default) stores the compute dtype; 'fp8'" + " stores e4m3 codes plus one fp32 scale per (token, kv head), roughly " + "doubling the tokens that fit in the same VRAM. Requires the triton" + " attention backend and a plain paged, hybrid-SWA or QSA sparse KV pool" + " (not MLA/DSA, DSV4, or MiniMax-M3 block-sparse models)." + " 'nvfp4' stores packed E2M1 with block/row scales; supports only" + " paged FULL, hybrid-SWA, or QSA sparse attention with head_dim divisible by 16." + ), + ) + parser.add_argument( "--attention-backend", "--attn", @@ -557,6 +637,22 @@ def _infer_reasoning_parser(model_path: str) -> str | None: "or fold them into nearby user/tool content without hoisting the prompt prefix.", ) + parser.add_argument( + "--enable-metrics-report", + action="store_true", + default=ServerArgs.enable_metrics_report, + help=( + "Serve a per-request `metrics` object next to usage on /v1/chat/completions, " + "/v1/messages and /v1/responses: ttft_ms, prefill_time_ms and " + "prefill_tokens_per_second (the prefill span measured by the scheduler itself), " + "decode_time_ms and decode_tokens_per_second, cached_prompt_tokens and " + "total_time_ms. Streaming responses carry it on the same final chunk as usage, so " + "the request must also ask for usage (OpenAI stream_options.include_usage). " + "Non-standard on every protocol, hence opt-in. Under concurrency the spans are " + "this request's share of shared batches, not isolated engine throughput." + ), + ) + parser.add_argument( "--sampling-defaults", type=str, @@ -617,6 +713,21 @@ def _infer_reasoning_parser(model_path: str) -> str | None: ), ) + parser.add_argument( + "--default-thinking-mode", + type=str, + default=ServerArgs.default_thinking_mode, + choices=["auto", "chat", "thinking"], + help=( + "Server-wide default thinking mode for reasoning-capable models. 'auto' keeps " + "the current per-request behavior; 'chat' turns thinking off for every request " + "that does not choose itself (chat_template_kwargs, reasoning_effort, " + "reasoning.effort, thinking) -- for OpenAI-compatible clients that never send " + "template kwargs, like Vercel AI SDK or llama-swap; 'thinking' turns it on the " + "same way." + ), + ) + parser.add_argument( "--moe-strategy", default=ServerArgs.moe_strategy, @@ -670,6 +781,16 @@ def _infer_reasoning_parser(model_path: str) -> str | None: ), ) + parser.add_argument( + "--embed-weights", + default=ServerArgs.embed_weights, + choices=["gpu", "host"], + help=( + "Where the token embedding table lives. 'host' keeps it in pinned host RAM and gathers " + "the rows each step needs over PCIe, freeing its VRAM (single GPU only)." + ), + ) + parser.add_argument( "--nvfp4-backend", action=_DeprecatedAlias, @@ -705,6 +826,12 @@ def _infer_reasoning_parser(model_path: str) -> str | None: default=ServerArgs.moe_cache_rate, help="The fraction of all MoE experts to keep in GPU cache.", ) + moe_cache_group.add_argument( + "--moe-cache-mem", + type=_parse_mem_size, + default=ServerArgs.moe_cache_mem, + help="GPU memory for the MoE expert slot cache, e.g. 12GiB or 800MiB (G, GB and GiB all mean GiB); the slot count is this over the bytes per slot.", + ) moe_cache_group.add_argument( "--moe-cache-auto", action="store_true", @@ -719,14 +846,67 @@ def _infer_reasoning_parser(model_path: str) -> str | None: "--kv-reserve-tokens", type=int, default=ServerArgs.kv_reserve_tokens, - help="KV-cache token floor reserved before --moe-cache-auto fills experts.", + help=( + "Usable KV-cache token floor reserved before --moe-cache-auto fills experts " + "(the internal dummy page is additional)." + ), + ) + + parser.add_argument( + "--moe-cache-scope", + choices=SCOPES, + default=ServerArgs.moe_cache_scope, + help=( + "How MoE layers share the GPU expert slot cache: unified = one LRU pool per expert quant " + "(all layers when they share one quant format); per-layer = each layer has its own pool. " + "Pools are sized by --moe-cache-split." + ), + ) + parser.add_argument( + "--moe-cache-split", + type=_parse_cache_split, + default=ServerArgs.moe_cache_split, + help=( + "How pools share the slots: equal (by layer count), or profile: = by expert use from a " + "profile file. With --moe-strategy offload, a pool smaller than the expert count limits the " + "decode batch to its slots / top_k." + ), + ) + + parser.add_argument( + "--kv-host-pages", + type=int, + default=ServerArgs.kv_host_pages, + help=( + "Extra KV pages mirrored to pinned host RAM (QSA models): the GPU pool becomes " + "an LRU cache over the logical page space, extending context past VRAM capacity. " + "0 = off." + ), ) parser.add_argument( "--moe-cache-policy", default=ServerArgs.moe_cache_policy, - choices=["lru"], - help="The unified MoE cache eviction policy.", + choices=["lru", "kd", "kdfb", "rule"], + help=( + "The unified MoE cache eviction policy. lru (default); kd: evict the expert passed " + "over for the most tokens, then the one whose layer comes up last; kdfb: kd plus a " + "decayed per-expert use count; rule: kdfb plus router near misses counted as recent " + "(margin FREETOKEN_MOE_NEAR_MISS_THR logits). GPU decode only." + ), + ) + + parser.add_argument( + "--moe-collect-stats", + action="store_true", + default=ServerArgs.moe_collect_stats, + help=( + "Log MoE expert-cache miss rate and routing skew during decode. The counters " + "are captured into the decode CUDA graph, so this can only be chosen at startup. " + "Measured cost is below noise (46.1 vs 45.8 tok/s median on Ornith-35B-A3B " + "IQ3_S), but it stays off by default since it is a diagnostic and the readout " + "costs a host sync." + ), ) parser.add_argument( @@ -769,6 +949,17 @@ def _infer_reasoning_parser(model_path: str) -> str | None: ), ) + parser.add_argument( + "--skip-prefill-warmup", + action="store_false", + dest="prefill_warmup", + default=ServerArgs.prefill_warmup, + help=( + "Skip the startup prefill warmup forwards. The first request at each " + "new size then pays the Triton compile/load cost mid-request." + ), + ) + parser.add_argument( "--disable-moe-prefill-overlap", action="store_false", @@ -871,6 +1062,9 @@ def _infer_reasoning_parser(model_path: str) -> str | None: kwargs["api_key"] = resolve_api_key(kwargs["api_key"]) + if kwargs["distributed_port"] is None: + kwargs["distributed_port"] = kwargs["server_port"] + 1 + if kwargs["model_path"].startswith("~"): kwargs["model_path"] = os.path.expanduser(kwargs["model_path"]) @@ -903,21 +1097,32 @@ def _infer_reasoning_parser(model_path: str) -> str | None: _no_cache_flag = ( kwargs["moe_cache_size"] == 0 and not kwargs["moe_cache_auto"] + and kwargs["moe_cache_mem"] is None and (kwargs["moe_cache_rate"] is None or kwargs["moe_cache_rate"] == 0) ) if is_offload_moe_strategy(kwargs["moe_strategy"]) and _no_cache_flag: kwargs["moe_cache_auto"] = True - if kwargs["model_source"] == "modelscope": - model_path = kwargs["model_path"] - if not os.path.isdir(model_path): + # Resolve a hub repo id ("org/model") to a local checkpoint directory. This has to + # happen here, after served_model_name and the parser cascade have read the repo id + # (they want "DeepSeek-V4-Flash-0731", not a snapshot hash) and before anything opens + # the checkpoint: config parsing reads model-specific files straight off disk, so a + # bare repo id reaches `open()` as a relative path and dies with a FileNotFoundError. + model_path = kwargs["model_path"] + if not os.path.isdir(model_path): + if kwargs["model_source"] == "modelscope": from modelscope import snapshot_download ignore_patterns = [] if kwargs["use_dummy_weight"]: ignore_patterns = ["*.bin", "*.safetensors", "*.pt", "*.ckpt"] - model_path = snapshot_download(model_path, ignore_patterns=ignore_patterns) - kwargs["model_path"] = model_path + kwargs["model_path"] = snapshot_download(model_path, ignore_patterns=ignore_patterns) + else: + from freetoken.utils import download_hf_checkpoint + + kwargs["model_path"] = download_hf_checkpoint( + model_path, dummy_weight=kwargs["use_dummy_weight"] + ) del kwargs["model_source"] # "auto" (or an unspecified dtype) resolves to the checkpoint's dtype. Multimodal / diff --git a/python/freetoken/server/control_api.py b/python/freetoken/server/control_api.py index 7158e4e5fa..cef632306e 100644 --- a/python/freetoken/server/control_api.py +++ b/python/freetoken/server/control_api.py @@ -68,6 +68,24 @@ async def list_requests(since: int = 0, limit: int = 100): from .stats import build_stats + @app.get("/metrics") + async def metrics(): + """Prometheus text exposition of /v1/stats (the llama.cpp / vLLM convention).""" + from fastapi.responses import PlainTextResponse + + from .stats import prometheus_text + + doc = build_stats( + get_state(), request_ring.requests_p95_ms(), request_ring.requests_ttft_mean_ms() + ) + return PlainTextResponse(prometheus_text(doc), media_type="text/plain; version=0.0.4") + + @app.get("/version") + async def version(): + from freetoken import __version__ + + return {"version": __version__} + @app.get("/v1/stats") async def stats(): doc = build_stats( diff --git a/python/freetoken/server/function_call_parser.py b/python/freetoken/server/function_call_parser.py index 296cd560dc..88cc65b367 100644 --- a/python/freetoken/server/function_call_parser.py +++ b/python/freetoken/server/function_call_parser.py @@ -349,17 +349,136 @@ def _ends_with_partial_token(self, buffer: str, bot_token: str) -> int: return i return 0 - def _get_param_config(self, func_name: str, tools: List[Tool]) -> Dict: - """Extract the parameter properties (JSON schema) for one tool.""" + def _resolve_local_ref(self, schema: Dict, root_schema: Dict, _seen: Optional[set[str]] = None,) -> Dict: + """ + Resolve local JSON Schema references. + + Unresolvable or cyclic references are returned unresolved. Their values + therefore follow the parser's existing loose/untyped conversion behavior. + """ + if not isinstance(schema, dict): + return schema + ref = schema.get("$ref") + if not isinstance(ref, str) or not ref.startswith("#/"): + return schema + if _seen is None: + _seen = set() + if ref in _seen: + # Cyclic ref: stop resolution and preserve the unresolved schema. + return schema + _seen.add(ref) + node = root_schema + try: + for part in ref[2:].split("/"): + part = part.replace("~1", "/").replace("~0", "~") + node = node[part] + except (KeyError, TypeError): + return schema + if not isinstance(node, dict): + return schema + merged = dict(node) + merged.update({k: v for k, v in schema.items() if k != "$ref"}) + if "$ref" in merged: + return self._resolve_local_ref(merged, root_schema, _seen) + return merged + + @staticmethod + def _literal_schema_type(schema: Dict) -> Optional[str]: + """Type of an untyped ``enum``/``const`` schema, inferred from its literals.""" + if "const" in schema: + values = [schema["const"]] + elif isinstance(schema.get("enum"), list): + values = schema["enum"] + else: + return None + values = [v for v in values if v is not None] + if not values: + return None + if all(isinstance(v, str) for v in values): + return "string" + if all(isinstance(v, bool) for v in values): + return "boolean" + if all(isinstance(v, int) and not isinstance(v, bool) for v in values): + return "integer" + return "loose" + + def _normalize_param_schema(self, schema: Dict, root_schema: Dict, _depth: int = 0) -> Dict: + """Resolve the schema parts needed for tool argument typing.""" + if not isinstance(schema, dict): + return {"type": "string"} + schema = self._resolve_local_ref(schema, root_schema) + param_type = schema.get("type") + if isinstance(param_type, list): + non_null = [t for t in param_type if t != "null"] + result = dict(schema) + result["type"] = non_null[0] if len(non_null) == 1 else "loose" + return result + if isinstance(param_type, str): + return schema + all_of = schema.get("allOf") + if ( + isinstance(all_of, list) + and len(all_of) == 1 + and isinstance(all_of[0], dict) + and _depth < 16 # bounds a self-referential allOf chain + ): + # Pydantic v1 / OpenAPI 3.0 wrap a $ref in a one-element allOf to attach + # siblings such as default/description; type it like the bare $ref. + merged = dict(all_of[0]) + merged.update({k: v for k, v in schema.items() if k != "allOf"}) + return self._normalize_param_schema(merged, root_schema, _depth + 1) + for keyword in ("oneOf", "anyOf"): + branches = schema.get(keyword) + if not isinstance(branches, list): + continue + types = [] + for branch in branches: + if not isinstance(branch, dict): + continue + branch = self._resolve_local_ref(branch, root_schema) + branch_type = branch.get("type") + if not isinstance(branch_type, str): + branch_type = self._literal_schema_type(branch) + if isinstance(branch_type, str) and branch_type != "null": + types.append(branch_type) + types = list(dict.fromkeys(types)) + result = dict(schema) + result["type"] = types[0] if len(types) == 1 else "loose" + return result + if isinstance(schema.get("properties"), dict): + result = dict(schema) + result["type"] = "object" + return result + if isinstance(schema.get("items"), (dict, list)): + result = dict(schema) + result["type"] = "array" + return result + result = dict(schema) + result["type"] = self._literal_schema_type(schema) or "loose" + return result + + def _get_tool_schema(self, func_name: str, tools: List[Tool]) -> Dict: + """Return the root JSON schema for one tool.""" for tool in tools: - if tool.function.name == func_name and tool.function.parameters: - params = tool.function.parameters - if isinstance(params, dict) and "properties" in params: - return params["properties"] - elif isinstance(params, dict): - return params + if tool.function.name != func_name or not tool.function.parameters: + continue + params = tool.function.parameters + return params if isinstance(params, dict) else {} return {} + def _get_param_config(self, func_name: str, tools: List[Tool]) -> Dict: + """Extract and normalize parameter properties for one tool.""" + params = self._get_tool_schema(func_name, tools) + if not params: + return {} + properties = params.get("properties") + if not isinstance(properties, dict): + return params + return { + name: self._normalize_param_schema(prop, params) + for name, prop in properties.items() + } + def _convert_param_value(self, value: str, param_name: str, param_config: Dict, func_name: str) -> Any: """Convert parameter value based on schema type. Safe alternative to eval().""" if value.lower() == "null": @@ -394,6 +513,8 @@ def _convert_param_value(self, value: str, param_name: str, param_config: Dict, return ast.literal_eval(value) except (ValueError, SyntaxError, TypeError): return value + elif param_type == "loose": + return _parse_loose_json_value(value) return value def _schema_param_type(self, param_name: str, param_config: Dict, missing: str = "string") -> str: @@ -2601,18 +2722,23 @@ def _union_leaf(self, raw: str, subs: list) -> Any: return v return raw - def _nested_value(self, raw: str, schema: Any = None) -> Any: + def _nested_value(self, raw: str, schema: Any = None, root_schema: Any = None) -> Any: + if isinstance(schema, dict) and isinstance(root_schema, dict): + schema = self._normalize_param_schema(schema, root_schema) items, stray = self._scan_elements(raw) if not items: return self._typed_leaf(raw, schema) - return self._structure(items, schema, stray=stray) + return self._structure(items, schema, root_schema=root_schema, stray=stray) - def _structure(self, items: List[tuple], schema: Any = None, stray: str = "") -> Any: + def _structure(self, items: List[tuple], schema: Any = None, root_schema: Any = None, stray: str = "") -> Any: """Composite value from scanned elements, threading the JSON schema down (``properties`` / ``items`` / ``additionalProperties``). Only ```` children render as a bare array -- the template's array convention; repeated siblings under any other tag stay an object with an array-valued key. Mixed text+children keeps the text under ``"$text"``.""" + if isinstance(schema, dict) and isinstance(root_schema, dict): + schema = self._normalize_param_schema(schema, root_schema) + props = schema.get("properties") if isinstance(schema, dict) else None item_schema = schema.get("items") if isinstance(schema, dict) else None @@ -2622,11 +2748,13 @@ def _sub_schema(key: str) -> Any: ap = schema.get("additionalProperties") if isinstance(ap, dict): sub = ap + if isinstance(sub, dict) and isinstance(root_schema, dict): + sub = self._normalize_param_schema(sub, root_schema) return sub names = {k for k, _ in items} if names == {"item"}: - return [self._nested_value(raw, item_schema) for _, raw in items] + return [self._nested_value(raw, item_schema, root_schema=root_schema) for _, raw in items] counts: Dict[str, int] = {} for k, _ in items: counts[k] = counts.get(k, 0) + 1 @@ -2642,7 +2770,7 @@ def _sub_schema(key: str) -> Any: and str(sub.get("type", "")).lower() == "array" ): sub = sub.get("items") - value = self._nested_value(raw, sub) + value = self._nested_value(raw, sub, root_schema=root_schema) if key in out: prev = out[key] out[key] = (prev if isinstance(prev, list) else [prev]) + [value] @@ -2653,16 +2781,19 @@ def _sub_schema(key: str) -> Any: return out def _args_from_items( - self, func_name: str, items: List[tuple], tools: List[Tool] + self, func_name: str, items: List[tuple], tools: List[Tool], ) -> Dict: + root_schema = self._get_tool_schema(func_name, tools) config = self._get_param_config(func_name, tools) args: Dict[str, Any] = {} for key, raw in items: prop = config.get(key) if isinstance(config, dict) and key in config else None nested, stray = self._scan_elements(raw) if nested: - value: Any = self._structure(nested, prop, stray=stray) + value: Any = self._structure(nested, prop, root_schema=root_schema, stray=stray) else: + if isinstance(prop, dict) and isinstance(root_schema, dict): + prop = self._normalize_param_schema(prop, root_schema) value = self._typed_leaf(raw, prop) if key in args: prev = args[key] diff --git a/python/freetoken/server/generation.py b/python/freetoken/server/generation.py index 2e342ff899..00d94abebf 100644 --- a/python/freetoken/server/generation.py +++ b/python/freetoken/server/generation.py @@ -15,6 +15,7 @@ import asyncio import json +import math import time from collections.abc import AsyncIterator from dataclasses import dataclass, field @@ -24,7 +25,7 @@ from freetoken.core import SamplingParams from freetoken.message import TokenizeMsg from freetoken.mm.media import collect_image_refs, fetch_image_bytes, image_reject_reason -from freetoken.tokenizer.tokenize import resolve_thinking_mode +from freetoken.tokenizer.tokenize import join_system_contents, resolve_thinking_mode from freetoken.tokenizer.inline_system import InlineSystemError try: @@ -44,6 +45,250 @@ ) +# Qwen's decoder intentionally keeps protocol markers visible so reasoning and +# tool-call parsers can consume their (non-special-token) tags. At the semantic +# API waist, after those parsers, only transport/internal multimodal markers are +# suppressed. Grounding markers (object_ref/box/quad), think tags, tool tags, +# and FIM/repository markers are deliberately absent from this list. +_QWEN_SEMANTIC_HIDDEN_TOKENS = ( + "<|endoftext|>", + "<|im_start|>", + "<|im_end|>", + "<|vision_start|>", + "<|vision_end|>", + "<|vision_pad|>", + "<|image_pad|>", + "<|video_pad|>", + "<|audio_start|>", + "<|audio_end|>", + "<|audio_pad|>", + "", + "", + "", + "", +) + + +def _uses_qwen_semantic_protocol(state: Any) -> bool: + """Whether this server is configured for a Qwen semantic output dialect. + + Parser selection is the stable frontend-visible family signal available in + ``ServerArgs``. Either half is enough because users may explicitly disable + reasoning or omit tools for an individual request. + """ + config = state.config + return ( + getattr(config, "reasoning_parser", None) == "qwen3" + or getattr(config, "tool_call_parser", None) in {"qwen", "qwen25", "qwen3_coder"} + ) + + +class _SemanticSpecialTokenFilter: + """Request-local, chunk-safe suppression of semantic protocol markers. + + A proper marker prefix is retained until the next detokenizer delta, so a + marker split at any byte boundary cannot leak. A marker inside a quoted/code + literal is retained only once the matching closer arrives; an unclosed quote + is sanitized at end-of-stream. Structured events can be held inside either + ambiguity so their ordering relative to surrounding content stays exact. + """ + + def __init__(self, tokens: tuple[str, ...]) -> None: + self._tokens = tokens + self._quote: str | None = None + self._escaped = False + self._last_char = "" + self._candidate_text = "" + self._candidate_items: list[Any] = [] + self._candidate_quoted = False + self._literal_items: list[Any] | None = None + + @staticmethod + def _append_item(items: list[Any], item: Any) -> None: + if item == "": + return + if isinstance(item, str) and items and isinstance(items[-1], str): + items[-1] += item + else: + items.append(item) + + @classmethod + def _strip_items(cls, items: list[Any], tokens: tuple[str, ...]) -> list[Any]: + """Strip markers across text segments while preserving event positions.""" + out: list[Any] = [] + candidate_text = "" + candidate_items: list[Any] = [] + + def flush(*, keep_text: bool) -> None: + nonlocal candidate_text, candidate_items + for pending in candidate_items: + if keep_text or not isinstance(pending, str): + cls._append_item(out, pending) + candidate_text = "" + candidate_items = [] + + for item in items: + if not isinstance(item, str): + if candidate_text: + cls._append_item(candidate_items, item) + else: + cls._append_item(out, item) + continue + index = 0 + while index < len(item): + ch = item[index] + if not candidate_text: + if any(token.startswith(ch) for token in tokens): + candidate_text = ch + cls._append_item(candidate_items, ch) + else: + cls._append_item(out, ch) + index += 1 + continue + extended = candidate_text + ch + if extended in tokens: + candidate_text = extended + cls._append_item(candidate_items, ch) + flush(keep_text=False) + index += 1 + elif any(token.startswith(extended) for token in tokens): + candidate_text = extended + cls._append_item(candidate_items, ch) + index += 1 + else: + flush(keep_text=True) + flush(keep_text=True) + return out + + def _start_candidate(self, ch: str, *, quoted: bool) -> None: + self._candidate_text = ch + self._candidate_items = [ch] + self._candidate_quoted = quoted + + def _extend_candidate(self, ch: str) -> tuple[list[Any], bool]: + extended = self._candidate_text + ch + if extended in self._tokens: + self._candidate_text = extended + self._append_item(self._candidate_items, ch) + if self._candidate_quoted: + self._literal_items = self._candidate_items + out: list[Any] = [] + else: + out = [item for item in self._candidate_items if not isinstance(item, str)] + self._candidate_text = "" + self._candidate_items = [] + self._candidate_quoted = False + return out, True + if any(token.startswith(extended) for token in self._tokens): + self._candidate_text = extended + self._append_item(self._candidate_items, ch) + return [], True + out = self._candidate_items + self._last_char = self._candidate_text[-1] + self._candidate_text = "" + self._candidate_items = [] + self._candidate_quoted = False + return out, False + + def feed_items(self, text: str, *, final: bool = False) -> list[Any]: + if not self._tokens: + return [text] if text else [] + out: list[Any] = [] + index = 0 + while index < len(text): + ch = text[index] + if self._literal_items is not None: + self._append_item(self._literal_items, ch) + if self._escaped: + self._escaped = False + elif ch == "\\": + self._escaped = True + elif ch == self._quote: + out.extend(self._literal_items) + self._literal_items = None + self._quote = None + self._last_char = ch + index += 1 + continue + + if self._candidate_text: + emitted, consumed = self._extend_candidate(ch) + out.extend(emitted) + if consumed: + index += 1 + continue + + if self._quote is not None: + if any(token.startswith(ch) for token in self._tokens): + self._escaped = False + self._start_candidate(ch, quoted=True) + index += 1 + continue + if self._escaped: + self._escaped = False + elif ch == "\\": + self._escaped = True + elif ch == self._quote: + self._quote = None + self._append_item(out, ch) + self._last_char = ch + index += 1 + continue + + opens_quote = ch in {'"', "`", "“", "‘"} or ( + ch == "'" and not self._last_char.isalnum() + ) + if opens_quote: + self._quote = {"“": "”", "‘": "’"}.get(ch, ch) + self._append_item(out, ch) + elif any(token.startswith(ch) for token in self._tokens): + self._start_candidate(ch, quoted=False) + index += 1 + continue + else: + self._append_item(out, ch) + self._last_char = ch + index += 1 + + if final: + if self._literal_items is not None: + out.extend(self._strip_items(self._literal_items, self._tokens)) + self._literal_items = None + elif self._candidate_text: + out.extend(self._candidate_items) + self._candidate_text = "" + self._candidate_items = [] + self._candidate_quoted = False + self._quote = None + self._escaped = False + return out + + def feed(self, text: str, *, final: bool = False) -> str: + items = self.feed_items(text, final=final) + assert all(isinstance(item, str) for item in items) + return "".join(items) + + def boundary(self, item: Any) -> list[Any]: + if self._literal_items is not None: + self._append_item(self._literal_items, item) + return [] + if self._candidate_text: + self._append_item(self._candidate_items, item) + return [] + return [item] + + def finish(self) -> str: + return self.feed("", final=True) + + def finish_items(self) -> list[Any]: + return self.feed_items("", final=True) + + +def _make_semantic_special_token_filter(state: Any) -> _SemanticSpecialTokenFilter: + tokens = _QWEN_SEMANTIC_HIDDEN_TOKENS if _uses_qwen_semantic_protocol(state) else () + return _SemanticSpecialTokenFilter(tokens) + + class GenerationError(ValueError): """A request failed before producing output (surfaced via ``UserReply.error`` — e.g. a chat template the tokenizer cannot render, or a prompt that exceeds the KV budget). Each @@ -74,6 +319,10 @@ class ReasoningDelta: @dataclass class ContentDelta: text: str + # Neutral sampled-token logprob entries riding this delta (see UserReply.logprobs); + # None when the request did not ask. Parser buffering can attach several entries + # to one delta. + logprobs: list[dict] | None = None @dataclass @@ -108,6 +357,26 @@ class ToolCallsDelta: calls: list[ToolCallItem] +@dataclass +class GenTimings: + """Per-request inference timings, in milliseconds. + + ``prefill_ms`` is the scheduler's own span (admission to the token sampled off the + request's last prefill chunk), shipped over the wire because it is the one boundary a + client cannot see; the rest are measured here, around the engine ack loop, so they carry + the same IPC hops the client's own HTTP timestamps would. + + Under concurrency every span is this request's SHARE of shared work -- its prefill chunk + is co-scheduled with other prompts, its decode steps batched with other requests -- so + the derived rates describe this request under that load, not isolated engine throughput. + """ + + ttft_ms: float = 0.0 + prefill_ms: float = 0.0 + decode_ms: float = 0.0 + total_ms: float = 0.0 + + @dataclass class GenDone: finish_reason: str @@ -115,6 +384,8 @@ class GenDone: completion_tokens: int matched_stop: str | None = None cached_tokens: int = 0 + reasoning_tokens: int = 0 + timings: GenTimings = field(default_factory=GenTimings) GenEvent = ReasoningDelta | ContentDelta | ToolCallStart | ToolCallArgsDelta | ToolCallsDelta | GenDone @@ -130,6 +401,11 @@ class GenResult: completion_tokens: int matched_stop: str | None = None cached_tokens: int = 0 + reasoning_tokens: int = 0 + timings: GenTimings = field(default_factory=GenTimings) + # Neutral sampled-token logprob entries, one per sampled token (empty when the + # request did not ask). + logprobs: list[dict] = field(default_factory=list) @dataclass @@ -169,11 +445,19 @@ def resolve_sampling( model_sampling: dict[str, Any], stop: str | list[str] | None = None, default_max_tokens: int = DEFAULT_MAX_OUTPUT_TOKENS, + min_p: float | None = None, presence_penalty: float = 0.0, frequency_penalty: float = 0.0, + repetition_penalty: float | None = None, + logit_bias: dict[str, float] | dict[int, float] | None = None, + min_tokens: int = 0, + stop_token_ids: list[int] | None = None, + include_stop_str_in_output: bool = False, + skip_special_tokens: bool | None = None, ) -> SamplingParams: """Map a protocol's sampling fields onto the engine's neutral SamplingParams, - filling unspecified fields from the checkpoint's recommended defaults.""" + filling unspecified fields from the checkpoint's recommended defaults. Range errors + raise ValueError (the adapters answer 400).""" def pick(value, key, framework): return value if value is not None else model_sampling.get(key, framework) @@ -185,23 +469,91 @@ def pick(value, key, framework): # non-positive value is a client error. if max_tokens is not None and max_tokens < 1: raise ValueError(f"max_tokens must be at least 1, got {max_tokens}") - return SamplingParams( + resolved_temperature = pick(temperature, "temperature", 0.0) + if not math.isfinite(resolved_temperature) or resolved_temperature < 0: + raise ValueError(f"temperature must be a finite number >= 0, got {resolved_temperature}") + resolved_top_p = pick(top_p, "top_p", 1.0) + if not math.isfinite(resolved_top_p) or not 0 < resolved_top_p <= 1: + raise ValueError(f"top_p must be in (0, 1], got {resolved_top_p}") + resolved_top_k = pick(top_k, "top_k", -1) + # Any top_k <= 0 means "no top-k filter", as in vLLM and the other OpenAI-compatible + # servers: Ollama-style clients (Open WebUI) and some SDKs send top_k=0 for "off", and a + # checkpoint's generation_config.json can carry it too. Map it to the engine's -1. + if resolved_top_k <= 0: + resolved_top_k = -1 + resolved_min_p = float(pick(min_p, "min_p", 0.0)) + resolved_rep = float(pick(repetition_penalty, "repetition_penalty", 1.0)) + if not 0.0 <= resolved_min_p <= 1.0: + raise ValueError(f"min_p must be in [0, 1], got {resolved_min_p}") + for name, value in (("presence_penalty", presence_penalty), ("frequency_penalty", frequency_penalty)): + if not -2.0 <= float(value) <= 2.0: + raise ValueError(f"{name} must be in [-2, 2], got {value}") + # not `<= 0`: NaN compares False and would poison every seen token's logit + if not (math.isfinite(resolved_rep) and resolved_rep > 0.0): + raise ValueError(f"repetition_penalty must be a positive finite number, got {resolved_rep}") + if min_tokens < 0: + raise ValueError(f"min_tokens must be >= 0, got {min_tokens}") + bias: list[list[float]] | None = None + if logit_bias: + bias = [] + for key, value in logit_bias.items(): + try: + tid = int(key) + except (TypeError, ValueError): + raise ValueError(f"logit_bias keys must be token ids, got {key!r}") from None + if tid < 0: + raise ValueError(f"logit_bias token id must be >= 0, got {tid}") + bias.append([tid, max(-100.0, min(100.0, float(value)))]) # the OpenAI range + ids = [int(t) for t in (stop_token_ids or [])] + if any(t < 0 for t in ids): + raise ValueError("stop_token_ids must be non-negative token ids") + params = SamplingParams( ignore_eos=ignore_eos, max_tokens=default_max_tokens if max_tokens is None else max_tokens, - temperature=pick(temperature, "temperature", 0.0), - top_k=pick(top_k, "top_k", -1), - top_p=pick(top_p, "top_p", 1.0), + temperature=resolved_temperature, + top_k=resolved_top_k, + top_p=resolved_top_p, stop_strs=[s for s in stop_list if s], # drop empty strings (would match everything) - presence_penalty=presence_penalty, - frequency_penalty=frequency_penalty, + min_p=resolved_min_p, + presence_penalty=float(presence_penalty), + frequency_penalty=float(frequency_penalty), + repetition_penalty=resolved_rep, + logit_bias=bias, + min_tokens=int(min_tokens), + stop_token_ids=ids, + include_stop_str_in_output=bool(include_stop_str_in_output), ) + if skip_special_tokens is not None: + params.skip_special_tokens = bool(skip_special_tokens) + return params -def render_messages(messages: list[dict[str, Any]]) -> list[dict[str, Any]]: +def render_messages(messages: list[dict[str, Any]], *, hoist_system: bool = True) -> list[dict[str, Any]]: """Normalize OpenAI-shaped message dicts for the chat template: flatten text content parts to a string and decode tool-call arguments from JSON. Raises - ValueError on a non-text content part (text-only server). Shared by all adapters.""" - return [_render_message(m) for m in messages] + ValueError on a non-text content part (text-only server). Shared by all adapters. + + Some chat templates (e.g. Qwen3.6) require the system message at index 0; + hoist system messages to the front and merge multiples into one to satisfy + that constraint. ``hoist_system=False`` leaves late system messages in place + for the tokenizer's inline-system policy (the Anthropic path).""" + rendered = [_render_message(m) for m in messages] + if not hoist_system or any(m.get("role") == "developer" for m in rendered): + # The tokenizer maps developer per template and then merges in order + # (_map_developer_role); hoisting system alone here would put it ahead of + # a developer message that came first. + return rendered + system_msgs = [m for m in rendered if m.get("role") == "system"] + if system_msgs and rendered[0].get("role") != "system": + non_system = [m for m in rendered if m.get("role") != "system"] + rendered = system_msgs + non_system + # Merge multiple system messages into one (Qwen3.6 template rejects + # system messages that are not at loop.first, i.e. after the first). + system_msgs = [m for m in rendered if m.get("role") == "system"] + if len(system_msgs) > 1: + non_system = [m for m in rendered if m.get("role") != "system"] + rendered = [{"role": "system", "content": join_system_contents(system_msgs)}] + non_system + return rendered def _render_message(message: dict[str, Any]) -> dict[str, Any]: @@ -395,7 +747,12 @@ def _make_reasoning_parser(spec: GenSpec, state: Any) -> ReasoningParser | None: # The qwen3 chat template opens an implicit (thinking on) unless # enable_thinking is explicitly false, so the model emits only the closing # . Mirror that default here, else the chain-of-thought leaks into content. - force_reasoning = (spec.chat_template_kwargs or {}).get("enable_thinking") is not False + # continue_final_message is the exception: the template renders the final assistant + # turn as ... + content, so the continuation starts in content. + ctk = spec.chat_template_kwargs or {} + force_reasoning = ctk.get("enable_thinking") is not False and not ctk.get( + "continue_final_message" + ) elif parser_name == "glm": # GLM's template honors enable_thinking (default on) even with tools; the # generic fallback would force thinking and mislabel disabled output as reasoning. @@ -504,6 +861,85 @@ async def with_keepalive(events: AsyncIterator[GenEvent], interval: float): task.cancel() +class _Timer: + """Times one generation off the raw ack stream. + + Deliberately fed from acks rather than from emitted GenEvents: a reasoning or tool-call + parser holds text back, sometimes to the end of the stream, so the first EVENT can lag the + first generated token by a lot. Acks are also the one path both the streaming and the + buffered generator share, so TTFT means the same thing on either. + """ + + def __init__(self) -> None: + self.start = time.monotonic() + self.first_token_at: float | None = None + self.end: float | None = None + self.prefill_ms = 0.0 + + def observe(self, ack: Any) -> None: + # an older peer (or a reply built before #504) carries no prefill span + prefill_ms = getattr(ack, "prefill_ms", 0.0) + if prefill_ms: + self.prefill_ms = prefill_ms + # completion_tokens_delta, not incremental_output: the detokenizer holds back a + # trailing partial-stop prefix, so a real generated token can arrive with empty text. + if self.first_token_at is None and ack.completion_tokens_delta: + self.first_token_at = time.monotonic() + if ack.finished: + # Stamped here, not in finish(): the callers run the reasoning split and the + # tool-call drain between the terminal ack and building their result, and that + # post-processing is not inference time. + self.end = time.monotonic() + + def finish(self) -> GenTimings: + # A stream cut short (client disconnect) never saw a terminal ack; time it to here. + end = self.end if self.end is not None else time.monotonic() + first = self.first_token_at + return GenTimings( + ttft_ms=0.0 if first is None else (first - self.start) * 1000.0, + prefill_ms=self.prefill_ms, + decode_ms=0.0 if first is None else (end - first) * 1000.0, + total_ms=(end - self.start) * 1000.0, + ) + + +def build_metrics( + *, prompt_tokens: int, completion_tokens: int, cached_tokens: int, timings: GenTimings +) -> dict[str, Any]: + """The `metrics` object served under --enable-metrics-report, identical on every protocol. + + ``cached_prompt_tokens`` is the real prefix-cache hit whatever --enable-cache-report says: + that flag governs BILLING fields in `usage`, and gating the hit here instead would leave + prefill_tokens_per_second silently computed over tokens that were never forwarded. + + Decode throughput divides by completion_tokens - 1, not completion_tokens: the first token + falls out of prefill, so decode_time_ms spans only the intervals after it (the convention + vLLM's and sglang's serving benchmarks use). + """ + prefill_tokens = max(prompt_tokens - cached_tokens, 0) + prefill_s = timings.prefill_ms / 1000.0 + decode_s = timings.decode_ms / 1000.0 + return { + "ttft_ms": round(timings.ttft_ms, 3), + "prefill_tokens": prefill_tokens, + "cached_prompt_tokens": cached_tokens, + "prefill_time_ms": round(timings.prefill_ms, 3), + "prefill_tokens_per_second": round(prefill_tokens / prefill_s, 2) if prefill_s > 0 else 0.0, + "decode_tokens": completion_tokens, + "decode_time_ms": round(timings.decode_ms, 3), + "decode_tokens_per_second": ( + round((completion_tokens - 1) / decode_s, 2) + if decode_s > 0 and completion_tokens > 1 + else 0.0 + ), + "total_time_ms": round(timings.total_ms, 3), + } + + +def metrics_enabled(state: Any) -> bool: + return bool(getattr(state.config, "enable_metrics_report", False)) + + def _record_generation( *, source: str | None, @@ -545,27 +981,24 @@ async def generate_events( """Wraps `_generate_events_impl` to log the request with its totals, read off the terminal `GenDone`. The `finally` still records the row on a mid-stream disconnect — but with 0 tokens if the drop lands before `GenDone`, the only event carrying the totals.""" - start = time.monotonic() + timer = _Timer() prompt_tokens = 0 completion_tokens = 0 - first_token_at: float | None = None error: str | None = None try: - async for ev in _generate_events_impl(uid, spec, state): + async for ev in _generate_events_impl(uid, spec, state, timer): if isinstance(ev, GenDone): prompt_tokens = ev.prompt_tokens completion_tokens = ev.completion_tokens - elif first_token_at is None: - first_token_at = time.monotonic() yield ev except GenerationError as exc: error = str(exc) raise finally: _record_generation( - source=source, stream=True, start=start, + source=source, stream=True, start=timer.start, prompt_tokens=prompt_tokens, completion_tokens=completion_tokens, error=error, - first_token_at=first_token_at, + first_token_at=timer.first_token_at, ) @@ -574,25 +1007,29 @@ async def generate_full( ) -> GenResult: """Wraps `_generate_full_impl` to log the request with its totals; the `finally` also records a `GenerationError` as a failed row.""" - start = time.monotonic() + timer = _Timer() result: GenResult | None = None error: str | None = None try: - result = await _generate_full_impl(uid, spec, state) + result = await _generate_full_impl(uid, spec, state, timer) return result except GenerationError as exc: error = str(exc) raise finally: _record_generation( - source=source, stream=False, start=start, + source=source, stream=False, start=timer.start, prompt_tokens=result.prompt_tokens if result else 0, completion_tokens=result.completion_tokens if result else 0, error=error, + # No first_token_at: non-streaming rows deliberately carry no TTFT, because + # requests_ttft_mean_ms averages only the rows that have one. `metrics` still does. ) -async def _generate_events_impl(uid: int, spec: GenSpec, state: Any) -> AsyncIterator[GenEvent]: +async def _generate_events_impl( + uid: int, spec: GenSpec, state: Any, timer: _Timer +) -> AsyncIterator[GenEvent]: """Protocol-neutral streaming generation. Yields semantic events (reasoning / content / tool-call deltas) terminated by exactly one GenDone. Produces no wire format — the OpenAI/Anthropic/Responses streamers format these into their own. @@ -607,10 +1044,26 @@ async def _generate_events_impl(uid: int, spec: GenSpec, state: Any) -> AsyncIte prompt_tokens = 0 completion_tokens = 0 cached_tokens = 0 + reasoning_tokens = 0 pending = "" + pending_logprobs: list[dict] = [] parse_tools = spec.parse_tools reasoning_parser = _make_reasoning_parser(spec, state) specials = _leaked_special_tokens(state) + reasoning_special_filter = _make_semantic_special_token_filter(state) + content_special_filter = _make_semantic_special_token_filter(state) + + def _filter_reasoning(piece: str) -> str: + return reasoning_special_filter.feed(strip_special_tokens(piece, specials)) + + def _filter_content(piece: str) -> str: + return content_special_filter.feed(strip_special_tokens(piece, specials)) + + def _content_delta(text: str) -> ContentDelta: + nonlocal pending_logprobs + logprobs = pending_logprobs or None + pending_logprobs = [] + return ContentDelta(text, logprobs=logprobs) tool_parser: FunctionCallParser | None = None if parse_tools: @@ -622,6 +1075,14 @@ async def _generate_events_impl(uid: int, spec: GenSpec, state: Any) -> AsyncIte tool_parser = candidate frag_stable = tool_parser.args_fragments_prefix_stable() if tool_parser else True + # Logprob entries are collected only on the passthrough path: with a reasoning + # parser or tool parsing active, text can be hidden, held back, or reclassified, + # so an entry could come to describe a token that never reaches visible content. + # The API layer fail-closes those request combinations; this guard is the + # structural half of the contract — a hidden-token entry must never ride a + # later visible delta, whatever the caller. + collect_logprobs = reasoning_parser is None and tool_parser is None and not parse_tools + # Streaming tool-call assembly: detectors emit fragments (name first, then # argument diffs); they accumulate here and the call is emitted complete when # the next call starts, trailing text arrives, or the stream ends. @@ -632,6 +1093,28 @@ async def _generate_events_impl(uid: int, spec: GenSpec, state: Any) -> AsyncIte # chunk is a legitimate word gap. suppress_ws = False + def _semantic_content_events(items: list[Any]) -> list[GenEvent]: + nonlocal suppress_ws + out: list[GenEvent] = [] + for item in items: + if isinstance(item, str): + if item and not (item.strip() == "" and suppress_ws): + out.append(_content_delta(item)) + if item.strip(): + suppress_ws = False + else: + out.append(item) + if isinstance(item, ToolCallsDelta): + suppress_ws = True + return out + + def _filter_tool_text(piece: str) -> list[GenEvent]: + text = strip_special_tokens(piece, specials) + return _semantic_content_events(content_special_filter.feed_items(text)) + + def _filter_tool_boundary(event: GenEvent) -> list[GenEvent]: + return _semantic_content_events(content_special_filter.boundary(event)) + def _close_open_call() -> ToolCallsDelta | None: nonlocal open_call, calls_emitted if open_call is None: @@ -647,8 +1130,6 @@ def _close_open_call() -> ToolCallsDelta | None: ) open_call = None calls_emitted += 1 - nonlocal suppress_ws - suppress_ws = True return ToolCallsDelta([call]) def _route_tool_text(piece: str) -> list[GenEvent]: @@ -662,12 +1143,8 @@ def _route_tool_text(piece: str) -> list[GenEvent]: # close it first so wire order matches generation order. done = _close_open_call() if done is not None: - out.append(done) - stripped = strip_special_tokens(payload, specials) - if stripped and not (stripped.strip() == "" and suppress_ws): - out.append(ContentDelta(stripped)) - if stripped.strip(): - suppress_ws = False + out.extend(_filter_tool_boundary(done)) + out.extend(_filter_tool_text(payload)) continue for frag in payload: starts_new = frag.name is not None and ( @@ -676,7 +1153,7 @@ def _route_tool_text(piece: str) -> list[GenEvent]: if starts_new: done = _close_open_call() if done is not None: - out.append(done) + out.extend(_filter_tool_boundary(done)) if open_call is None or starts_new: open_call = { "detector_index": frag.tool_index, @@ -684,14 +1161,16 @@ def _route_tool_text(piece: str) -> list[GenEvent]: "params": "", "ordinal": calls_emitted, } - out.append(ToolCallStart( - tool_index=calls_emitted, name=frag.name, args_prefix_stable=frag_stable, - )) + out.extend(_filter_tool_boundary(ToolCallStart( + tool_index=calls_emitted, + name=frag.name, + args_prefix_stable=frag_stable, + ))) if frag.parameters: open_call["params"] += frag.parameters - out.append( + out.extend(_filter_tool_boundary( ToolCallArgsDelta(tool_index=open_call["ordinal"], fragment=frag.parameters) - ) + )) return out engine_finish_reason: str | None = None @@ -699,14 +1178,18 @@ def _route_tool_text(piece: str) -> list[GenEvent]: async for ack in state.wait_for_ack(uid): if getattr(ack, "error", None): raise GenerationError(ack.error, getattr(ack, "error_code", None)) + timer.observe(ack) prompt_tokens += ack.prompt_tokens_delta completion_tokens += ack.completion_tokens_delta cached_tokens += ack.cached_tokens + reasoning_tokens = max(reasoning_tokens, getattr(ack, "reasoning_tokens", 0)) + if collect_logprobs and getattr(ack, "logprobs", None) is not None: + pending_logprobs.append(ack.logprobs) content_delta = ack.incremental_output if reasoning_parser is not None and content_delta: reasoning_delta, content_delta = reasoning_parser.parse_stream_chunk(content_delta) if reasoning_delta: - stripped_reasoning = strip_special_tokens(reasoning_delta, specials) + stripped_reasoning = _filter_reasoning(reasoning_delta) if stripped_reasoning: # a bare special token must not open a thinking block yield ReasoningDelta(stripped_reasoning) if content_delta: @@ -716,7 +1199,9 @@ def _route_tool_text(piece: str) -> list[GenEvent]: elif parse_tools: pending += content_delta else: - yield ContentDelta(strip_special_tokens(content_delta, specials)) + stripped_content = _filter_content(content_delta) + if stripped_content: + yield _content_delta(stripped_content) if ack.finished: engine_finish_reason = getattr(ack, "finish_reason", None) engine_matched_stop = getattr(ack, "matched_stop", None) @@ -727,7 +1212,7 @@ def _route_tool_text(piece: str) -> list[GenEvent]: if reasoning_parser is not None: flush_reasoning, flush_content = reasoning_parser.flush() if flush_reasoning: - stripped_reasoning = strip_special_tokens(flush_reasoning, specials) + stripped_reasoning = _filter_reasoning(flush_reasoning) if stripped_reasoning: yield ReasoningDelta(stripped_reasoning) if flush_content: @@ -737,7 +1222,9 @@ def _route_tool_text(piece: str) -> list[GenEvent]: elif parse_tools: pending += flush_content else: - yield ContentDelta(strip_special_tokens(flush_content, specials)) + stripped_content = _filter_content(flush_content) + if stripped_content: + yield _content_delta(stripped_content) # Engine reason ("stop"/"length"); a tool call overrides it, but a truncation (length) wins. finish_reason = engine_finish_reason or "stop" @@ -751,58 +1238,92 @@ def _route_tool_text(piece: str) -> list[GenEvent]: for frag in tool_parser.finalize_stream(): if open_call is not None and frag.parameters: open_call["params"] += frag.parameters - yield ToolCallArgsDelta(tool_index=open_call["ordinal"], fragment=frag.parameters) + for ev in _filter_tool_boundary( + ToolCallArgsDelta(tool_index=open_call["ordinal"], fragment=frag.parameters) + ): + yield ev done = _close_open_call() if done is not None: - yield done + for ev in _filter_tool_boundary(done): + yield ev for item in tool_parser.recover_truncated_call(): - yield ToolCallStart(tool_index=calls_emitted, name=item.name) - yield ToolCallsDelta( - [ToolCallItem(tool_index=calls_emitted, name=item.name, parameters=item.parameters)] - ) + recovered = [ + ToolCallStart(tool_index=calls_emitted, name=item.name), + ToolCallsDelta([ + ToolCallItem( + tool_index=calls_emitted, + name=item.name, + parameters=item.parameters, + ) + ]), + ] + for event in recovered: + for ev in _filter_tool_boundary(event): + yield ev calls_emitted += 1 residual = tool_parser.finish_stream() if residual: - stripped = strip_special_tokens(residual, specials) - if stripped and not (stripped.strip() == "" and suppress_ws): - yield ContentDelta(stripped) + for ev in _filter_tool_text(residual): + yield ev if calls_emitted and finish_reason != "length": finish_reason = "tool_calls" else: parsed = _parse_tool_response(pending, spec, state) if parse_tools else None if parsed is not None: normal_text, tool_calls = parsed - normal_text = strip_special_tokens(normal_text, specials) + normal_text = _filter_content(normal_text) if normal_text: - yield ContentDelta(normal_text) + yield _content_delta(normal_text) yield ToolCallsDelta(tool_calls) if finish_reason != "length": finish_reason = "tool_calls" elif parse_tools and pending: - yield ContentDelta(strip_special_tokens(pending, specials)) + stripped_pending = _filter_content(pending) + if stripped_pending: + yield _content_delta(stripped_pending) + + trailing_reasoning = reasoning_special_filter.finish() + if trailing_reasoning: + yield ReasoningDelta(trailing_reasoning) + if tool_parser is not None: + for ev in _semantic_content_events(content_special_filter.finish_items()): + yield ev + else: + trailing_content = content_special_filter.finish() + if trailing_content: + yield _content_delta(trailing_content) + # Entries without a content delta are intentionally dropped in streaming mode. yield GenDone( finish_reason, prompt_tokens, completion_tokens, matched_stop=engine_matched_stop, cached_tokens=cached_tokens, + reasoning_tokens=reasoning_tokens, + timings=timer.finish(), ) -async def _generate_full_impl(uid: int, spec: GenSpec, state: Any) -> GenResult: +async def _generate_full_impl(uid: int, spec: GenSpec, state: Any, timer: _Timer) -> GenResult: """Protocol-neutral non-streaming generation: accumulate, split reasoning, parse tool calls, strip special tokens. The adapters format the GenResult into their wire.""" full_content = "" + logprob_entries: list[dict] = [] prompt_tokens = 0 completion_tokens = 0 cached_tokens = 0 + reasoning_tokens = 0 engine_finish_reason: str | None = None engine_matched_stop: str | None = None async for ack in state.wait_for_ack(uid): if getattr(ack, "error", None): raise GenerationError(ack.error, getattr(ack, "error_code", None)) + timer.observe(ack) prompt_tokens += ack.prompt_tokens_delta completion_tokens += ack.completion_tokens_delta cached_tokens += ack.cached_tokens + reasoning_tokens = max(reasoning_tokens, getattr(ack, "reasoning_tokens", 0)) full_content += ack.incremental_output + if getattr(ack, "logprobs", None) is not None: + logprob_entries.append(ack.logprobs) if ack.finished: engine_finish_reason = getattr(ack, "finish_reason", None) engine_matched_stop = getattr(ack, "matched_stop", None) @@ -819,13 +1340,22 @@ async def _generate_full_impl(uid: int, spec: GenSpec, state: Any) -> GenResult: finish_reason = "tool_calls" specials = _leaked_special_tokens(state) + reasoning_special_filter = _make_semantic_special_token_filter(state) + content_special_filter = _make_semantic_special_token_filter(state) return GenResult( - reasoning=strip_special_tokens(reasoning_text, specials).strip(), - content=strip_special_tokens(content_text, specials), + reasoning=reasoning_special_filter.feed( + strip_special_tokens(reasoning_text, specials), final=True + ).strip(), + content=content_special_filter.feed( + strip_special_tokens(content_text, specials), final=True + ), tool_calls=tool_calls, finish_reason=finish_reason, prompt_tokens=prompt_tokens, completion_tokens=completion_tokens, matched_stop=engine_matched_stop, cached_tokens=cached_tokens, + reasoning_tokens=reasoning_tokens, + timings=timer.finish(), + logprobs=logprob_entries, ) diff --git a/python/freetoken/server/launch.py b/python/freetoken/server/launch.py index bb0c256f46..b02876c332 100644 --- a/python/freetoken/server/launch.py +++ b/python/freetoken/server/launch.py @@ -52,13 +52,32 @@ def _run_tokenize_worker(detach: bool, **kwargs) -> None: _detach_process_group() from freetoken.tokenizer import tokenize_worker - tokenize_worker(**kwargs) + ack_queue = kwargs.get("ack_queue") + try: + tokenize_worker(**kwargs) + except Exception as exc: # noqa: BLE001 -- surface the reason, then let it propagate + # Same contract as _run_scheduler: without this, a tokenizer/detokenizer that dies + # before its ready ack reaches the parent only as a dead process, so the supervisor + # reports the generic "backend worker freetoken-detokenizer-0 exited during load" and + # the real cause (an unreadable tokenizer.json, a missing checkpoint file, a ZMQ bind + # clash) is lost. The traceback still prints and the process still exits non-zero. + if ack_queue is not None: + _report_startup_error(ack_queue, exc) + raise def _run_scheduler(args: ServerArgs, ack_queue: mp.Queue[str]) -> None: if args.shell_mode: _detach_process_group() + import threading + + from tqdm import tqdm + + # Only rank 0 draws a bar, so no bar lock spans processes; tqdm's default one is a multiprocessing + # semaphore, which this worker leaves to resource_tracker as leaked when a signal ends it. + tqdm.set_lock(threading.RLock()) + # published (not bound) here: the engine binds it after the allocator setup from freetoken.gpu_select import set_assigned_gpu diff --git a/python/freetoken/server/logprobs.py b/python/freetoken/server/logprobs.py new file mode 100644 index 0000000000..8618569544 --- /dev/null +++ b/python/freetoken/server/logprobs.py @@ -0,0 +1,60 @@ +"""OpenAI logprobs formatting: neutral engine entries -> wire shapes.""" + +from __future__ import annotations + +from typing import Any + + +def chat_logprobs_error(req: Any) -> str | None: + if req.top_logprobs is not None and not req.logprobs: + return "top_logprobs requires logprobs=true" + if req.top_logprobs is not None and not 0 <= req.top_logprobs <= 20: + return "top_logprobs must be between 0 and 20" + return None + + +def completion_logprobs_error(req: Any) -> str | None: + if req.logprobs is not None and not 0 <= req.logprobs <= 5: + return "logprobs must be between 0 and 5" + if req.echo and req.logprobs is not None: + return "echo with logprobs is not supported" + return None + + +def chat_content_entry(entry: dict) -> dict: + return { + "token": entry["token"], + "logprob": entry["logprob"], + "bytes": entry["bytes"], + "top_logprobs": [ + { + "token": candidate["token"], + "logprob": candidate["logprob"], + "bytes": candidate["bytes"], + } + for candidate in entry["top"] + ], + } + + +def completions_logprobs(entries: list[dict], start_offset: int = 0) -> dict: + tokens: list[str] = [] + token_logprobs: list[float] = [] + top_logprobs: list[dict[str, float]] = [] + text_offset: list[int] = [] + offset = start_offset + + for entry in entries: + token = entry["token"] + tokens.append(token) + token_logprobs.append(entry["logprob"]) + top_logprobs.append({candidate["token"]: candidate["logprob"] for candidate in entry["top"]}) + text_offset.append(offset) + offset += len(token) + + return { + "tokens": tokens, + "token_logprobs": token_logprobs, + "top_logprobs": top_logprobs, + "text_offset": text_offset, + } diff --git a/python/freetoken/server/model_meta.py b/python/freetoken/server/model_meta.py index 8ce8009a05..c9187ad629 100644 --- a/python/freetoken/server/model_meta.py +++ b/python/freetoken/server/model_meta.py @@ -119,6 +119,25 @@ def effort_toggle_kwargs( return mapped +def apply_default_thinking_mode( + ctk: dict[str, Any] | None, + default_mode: str | None, +) -> dict[str, Any] | None: + """Merge the server's --default-thinking-mode into a request's template kwargs. + + An explicit per-request value always wins; the server default only fills + in what the request left unset. "auto" (the default) is a no-op. Adapters + apply it after folding their protocol-level toggles (effort_toggle_kwargs), + so those count as explicit too. + """ + if default_mode not in ("chat", "thinking"): + return ctk + ctk = dict(ctk or {}) + if not any(key in ctk for key in _THINKING_KWARG_KEYS): + ctk.update(thinking_toggle_kwargs(default_mode == "thinking")) + return ctk + + def moe_total_experts(config: Any) -> int: """Total routed-expert slots the model has: experts per layer x MoE layers. Matches the engine's own basis (``Engine._resolve_auto_moe_cache_size``), so a residency rate derived diff --git a/python/freetoken/server/openai_api.py b/python/freetoken/server/openai_api.py index 99a4283549..0fc8113bc8 100644 --- a/python/freetoken/server/openai_api.py +++ b/python/freetoken/server/openai_api.py @@ -2,9 +2,10 @@ import asyncio import json +import logging import time import uuid -from collections.abc import AsyncIterator, Callable +from collections.abc import AsyncIterator, Callable, Sequence from typing import Any from fastapi import FastAPI, Request @@ -14,6 +15,9 @@ from freetoken.tokenizer.effort import EFFORT_SCALE, KNOWN_REASONING_EFFORTS from .api_models import ( + MAX_N, + DetokenizeRequest, + TokenizeRequest, ChatCompletionRequest, CompletionRequest, ModelCard, @@ -22,24 +26,35 @@ ) from .function_call_parser import ToolCallItem from .request_logger import log_request +from .logprobs import ( + chat_content_entry, + chat_logprobs_error, + completion_logprobs_error, + completions_logprobs, +) from .generation import ( DEFAULT_MAX_OUTPUT_TOKENS, ContentDelta, GenDone, GenerationError, GenSpec, + GenTimings, ReasoningDelta, ToolCallArgsDelta, ToolCallsDelta, ToolCallStart, + build_metrics, generate_events, generate_full, + metrics_enabled, prerender_error, render_messages, resolve_sampling, submit_generation, ) +logger = logging.getLogger(__name__) + #: The wire superset plus "off", DeepSeek's disable synonym that #: effort_toggle_kwargs has always honored. _ACCEPTED_EFFORTS = (*KNOWN_REASONING_EFFORTS, "off") @@ -53,35 +68,28 @@ def _thinking_type(req: Any) -> str | None: if value in ("enabled", "disabled"): return value return None - - - def chat_request_to_genspec( req: ChatCompletionRequest, model_sampling: dict[str, Any], default_max_tokens: int = DEFAULT_MAX_OUTPUT_TOKENS, + default_thinking_mode: str | None = None, ) -> GenSpec: """OpenAI ChatCompletionRequest -> GenSpec (the OpenAI 'to_sampling_params').""" - from .model_meta import effort_toggle_kwargs + from .model_meta import apply_default_thinking_mode, effort_toggle_kwargs ctk = req.chat_template_kwargs thinking_type = _thinking_type(req) if req.reasoning_effort or thinking_type: ctk = effort_toggle_kwargs(req.reasoning_effort, ctk, thinking_type=thinking_type) + ctk = apply_default_thinking_mode(ctk, default_thinking_mode) + if req.continue_final_message: + ctk = {**ctk, "continue_final_message": True} + sampling_params = _resolve_sampling(req, model_sampling, default_max_tokens=default_max_tokens) + sampling_params.logprobs = bool(req.logprobs) + sampling_params.top_logprobs = req.top_logprobs or 0 return GenSpec( messages=render_messages([m.model_dump(exclude_none=True) for m in req.messages]), - sampling_params=resolve_sampling( - temperature=req.temperature, - top_k=req.top_k, - top_p=req.top_p, - max_tokens=req.max_tokens, - ignore_eos=req.ignore_eos, - model_sampling=model_sampling, - stop=req.stop, - default_max_tokens=default_max_tokens, - presence_penalty=req.presence_penalty, - frequency_penalty=req.frequency_penalty, - ), + sampling_params=sampling_params, chat_template_kwargs=ctk, template_tools=_tools_for_template(req), parser_tools=(_all_tool_dicts(req.tools) if _should_parse_tools(req) else None), @@ -133,6 +141,16 @@ async def v1_completions(req: CompletionRequest, request: Request): return gate return await handle_completion(req, request, state, get_model_sampling()) + @app.post("/tokenize") + @app.post("/v1/tokenize") + async def tokenize(req: TokenizeRequest): + return await handle_tokenize(req, get_state()) + + @app.post("/detokenize") + @app.post("/v1/detokenize") + async def detokenize(req: DetokenizeRequest): + return await handle_detokenize(req, get_state()) + @app.get("/v1/models") async def v1_models(): state = get_state() @@ -148,6 +166,126 @@ async def v1_models(): default_reasoning_effort=default_effort, )]) + @app.get("/v1/models/{model_id:path}") + async def v1_model(model_id: str): + """Retrieve metadata for the currently served model. + + OpenAI-compatible clients may use this endpoint to validate a model + name before sending completion requests. + """ + if not model_id: + # The :path converter also matches "/v1/models/", which used to redirect to the list. + return await v1_models() + state = get_state() + served_model_id = _served_model_name(state) + + if model_id != served_model_id: + return create_error_response( + f"The model '{model_id}' does not exist", + status_code=404, + err_type="invalid_request_error", + param="model", + code="model_not_found", + ) + + ctx = _model_context_length(state) + efforts, default_effort = await _effort_fields(state) + return ModelCard( + id=served_model_id, + root=state.config.model_path, + max_model_len=ctx, + context_length=ctx, + supported_reasoning_efforts=efforts, + default_reasoning_effort=default_effort, + ) + +# How often a non-streaming handler looks at the transport while the engine +# generates. The streaming path checks per chunk; one second bounds an abandoned +# request's extra decode work without measurable polling overhead. +_DISCONNECT_POLL_SECONDS = 1.0 + + +async def _await_watching_disconnect( + awaitable, request: Request | None, state: Any, uids: Sequence[int] +): + """Run an in-flight generation awaitable while watching the transport — the + non-streaming twin of stream_with_cancellation. Returns the awaitable's result, + or None when the client disconnected first: the same shielded abort_user is + delivered for every uid (an ``n > 1`` fan-out) so an abandoned request stops + burning decode slots instead of running to max_tokens (with + --max-running-requests 1 that is a full outage for its remaining budget).""" + gen = asyncio.ensure_future(awaitable) + if request is None: + return await gen + + async def abort_all(): + for uid in uids: + try: + await asyncio.shield(state.abort_user(uid)) + except Exception: # noqa: BLE001 + logger.exception("Failed to deliver abort for user %s", uid) + + try: + while True: + done, _ = await asyncio.wait({gen}, timeout=_DISCONNECT_POLL_SECONDS) + if done: + return gen.result() + if await request.is_disconnected(): + break + except asyncio.CancelledError: + # The handler itself was cancelled (server shutdown): deliver the abort, + # then let the cancellation propagate — mirrors stream_with_cancellation. + await abort_all() + gen.cancel() + raise + # Client gone. Deliver the AbortMsg first (the decode stops sooner), then + # wind down the drain task; its result is undeliverable either way. + await abort_all() + gen.cancel() + try: + await gen + except asyncio.CancelledError: + pass + except Exception: # noqa: BLE001 — the response is undeliverable, only cleanup matters + pass + return None + + +def _client_disconnected_response() -> JSONResponse: + # 499 — nginx's "client closed request". The client is gone; this status is + # for the access log, not the wire. + return create_error_response( + "client disconnected before the response was ready", + status_code=499, + err_type="client_disconnected", + ) + + +def _chat_logprobs_conflict(req: ChatCompletionRequest, state: Any) -> str | None: + """Fail-close chat logprobs when a semantic parsing layer can hide or reclassify + generated tokens. A reasoning parser routes tokens out of visible content (and its + stream buffering re-chunks text), and tool parsing consumes tokens into tool_calls, + so logprob entries could not be aligned 1:1 with visible content tokens — and hidden + reasoning-token strings must never surface through logprobs entries. The raw + /v1/completions path has no semantic layer and keeps full logprobs support.""" + if not req.logprobs: + return None + if getattr(state.config, "reasoning_parser", None): + return ( + "logprobs on /v1/chat/completions are not supported when the server runs a " + "reasoning parser: reasoning tokens are hidden from message content, so " + "logprob entries cannot be aligned with it. Use /v1/completions for raw " + "token logprobs." + ) + if _should_parse_tools(req): + return ( + "logprobs with tool parsing on /v1/chat/completions are not supported: " + "tool-call tokens are consumed into tool_calls, so logprob entries cannot " + "be aligned with message content. Send tool_choice='none' or use " + "/v1/completions for raw token logprobs." + ) + return None + async def handle_chat_completion( req: ChatCompletionRequest, @@ -157,15 +295,21 @@ async def handle_chat_completion( ): if req.function_call is not None: return create_error_response("function_call is not supported; use tools/tool_choice instead") - if req.logit_bias is not None: - return create_error_response("logit_bias is not supported") + logprobs_error = chat_logprobs_error(req) + if logprobs_error is not None: + return create_error_response(logprobs_error, param="top_logprobs") + logprobs_conflict = _chat_logprobs_conflict(req, state) + if logprobs_conflict is not None: + return create_error_response(logprobs_conflict, param="logprobs") if _response_format_unsupported(req.response_format): + rtype = (req.response_format or {}).get("type") return create_error_response( - "response_format json_object/json_schema is not supported (no constrained decoding)", + f"response_format.type {rtype!r} is not supported " + "(supported: text, json_object, json_schema)", param="response_format", ) - if req.n != 1: - return create_error_response("Only n=1 is supported", param="n") + if not 1 <= req.n <= MAX_N: + return create_error_response(f"n must be between 1 and {MAX_N}", param="n") # Case/whitespace and the "off" disable synonym stay accepted here because # effort_toggle_kwargs normalizes and honors them downstream. effort = req.reasoning_effort.strip().lower() if isinstance(req.reasoning_effort, str) else None @@ -187,7 +331,12 @@ async def handle_chat_completion( default_max_tokens = ( getattr(state.config, "max_output_tokens", None) or DEFAULT_MAX_OUTPUT_TOKENS ) - spec = chat_request_to_genspec(req, model_sampling, default_max_tokens=default_max_tokens) + spec = chat_request_to_genspec( + req, + model_sampling, + default_max_tokens=default_max_tokens, + default_thinking_mode=getattr(state.config, "default_thinking_mode", "auto"), + ) except ValueError as exc: return create_error_response(str(exc)) @@ -198,45 +347,86 @@ async def handle_chat_completion( if err is not None: return create_error_response(str(err), code=err.code) + # n > 1: one generation per choice, submitted together (the prefix cache serves the + # shared prompt after the first prefill). + uids: list[int] = [] try: - uid = await submit_generation(spec, state) + for _ in range(req.n): + uids.append(await submit_generation(spec, state)) except GenerationError as exc: + for u in uids: # choices already queued before a later one failed + await state.abort_user(u) return create_error_response(str(exc), code=exc.code) + uid = uids[0] if req.stream: - chunks = stream_chat_completion_chunks(uid, req, state, spec) + if req.n == 1: + chunks = stream_chat_completion_chunks(uid, req, state, spec) + else: + chunks = _merge_streams( + [ + stream_chat_completion_chunks(u, req, state, spec, index=i, terminal=False) + for i, u in enumerate(uids) + ], + lambda: _chat_chunk(req, uid, []), + include_usage=bool(req.stream_options and req.stream_options.include_usage), + state=state, + ) if request is not None: - chunks = state.stream_with_cancellation(chunks, request, uid) + chunks = state.stream_with_cancellation(chunks, request, uids) return StreamingResponse(chunks, media_type="text/event-stream") + tasks = [ + asyncio.ensure_future(generate_full(u, spec, state, source="/v1/chat/completions")) + for u in uids + ] try: - result = await generate_full(uid, spec, state, source="/v1/chat/completions") + results = await _await_watching_disconnect(asyncio.gather(*tasks), request, state, uids) except GenerationError as exc: + # gather does not cancel the other choices; the aborts below drop the events they + # wait on, so without the cancel they would stay pending forever + for task in tasks: + task.cancel() + for u in uids: + await state.abort_user(u) return create_error_response(str(exc), code=exc.code) - message: dict[str, Any] = {"role": "assistant", "content": result.content} - if result.reasoning: - message["reasoning_content"] = result.reasoning - if result.tool_calls: - message["tool_calls"] = _tool_calls_to_openai(result.tool_calls) - - return { - "id": f"chatcmpl-{uid}", + if results is None: + return _client_disconnected_response() + choices: list[dict[str, Any]] = [] + for index, result in enumerate(results): + message: dict[str, Any] = {"role": "assistant", "content": result.content} + if result.reasoning: + message["reasoning_content"] = result.reasoning + if result.tool_calls: + message["tool_calls"] = _tool_calls_to_openai(result.tool_calls) + choice: dict[str, Any] = {"index": index, "message": message, "finish_reason": result.finish_reason} + if req.logprobs: + choice["logprobs"] = {"content": [chat_content_entry(e) for e in result.logprobs]} + choices.append(choice) + + first = results[0] + response: dict[str, Any] = { + "id": _response_id("chatcmpl", req, uid), "object": "chat.completion", "created": int(time.time()), "model": req.model, - "choices": [ - { - "index": 0, - "message": message, - "finish_reason": result.finish_reason, - } - ], + "system_fingerprint": None, + "choices": choices, "usage": _usage( - result.prompt_tokens, - result.completion_tokens, - _reported_cached(state, result.cached_tokens), + first.prompt_tokens, + sum(r.completion_tokens for r in results), + _reported_cached(state, max(r.cached_tokens for r in results)), + reasoning_tokens=sum(r.reasoning_tokens for r in results), ), } + if metrics_enabled(state): + response["metrics"] = build_metrics( + prompt_tokens=first.prompt_tokens, + completion_tokens=first.completion_tokens, + cached_tokens=first.cached_tokens, + timings=first.timings, + ) + return response async def stream_chat_completion_chunks( @@ -244,21 +434,33 @@ async def stream_chat_completion_chunks( req: ChatCompletionRequest, state: Any, spec: GenSpec | None = None, + index: int = 0, + terminal: bool = True, ) -> AsyncIterator[bytes]: - """Format generate_events() into the OpenAI chat.completion.chunk SSE stream.""" + """Format generate_events() into the OpenAI chat.completion.chunk SSE stream. + + ``index`` is the choice index (``n > 1`` runs one of these per choice); with + ``terminal=False`` the usage chunk and ``[DONE]`` are left to the merger, which reads + this stream's usage from the ``_StreamUsage`` yielded last.""" if spec is None: - spec = chat_request_to_genspec(req, {}) + spec = chat_request_to_genspec( + req, + {}, + default_thinking_mode=getattr(state.config, "default_thinking_mode", "auto"), + ) yield _sse( _chat_chunk( req, uid, - [{"delta": {"role": "assistant", "content": ""}, "index": 0, "finish_reason": None}], + [{"delta": {"role": "assistant", "content": ""}, "index": index, "finish_reason": None}], ) ) prompt_tokens = 0 completion_tokens = 0 cached_tokens = 0 + reasoning_tokens = 0 + timings = GenTimings() tool_calls_sent = 0 open_tool: dict[str, Any] | None = None events = generate_events(uid, spec, state, source="/v1/chat/completions") @@ -279,17 +481,14 @@ async def stream_chat_completion_chunks( _chat_chunk( req, uid, - [{"delta": {"reasoning_content": ev.text}, "index": 0, "finish_reason": None}], + [{"delta": {"reasoning_content": ev.text}, "index": index, "finish_reason": None}], ) ) elif isinstance(ev, ContentDelta): - yield _sse( - _chat_chunk( - req, - uid, - [{"delta": {"content": ev.text}, "index": 0, "finish_reason": None}], - ) - ) + choice: dict[str, Any] = {"delta": {"content": ev.text}, "index": index, "finish_reason": None} + if ev.logprobs: + choice["logprobs"] = {"content": [chat_content_entry(e) for e in ev.logprobs]} + yield _sse(_chat_chunk(req, uid, [choice])) elif isinstance(ev, ToolCallStart): open_tool = { "index": tool_calls_sent, @@ -307,7 +506,7 @@ async def stream_chat_completion_chunks( "type": "function", "function": {"name": ev.name, "arguments": ""}, }]}, - "index": 0, "finish_reason": None, + "index": index, "finish_reason": None, }], ) ) @@ -324,7 +523,7 @@ async def stream_chat_completion_chunks( "index": open_tool["index"], "function": {"arguments": ev.fragment}, }]}, - "index": 0, "finish_reason": None, + "index": index, "finish_reason": None, }], ) ) @@ -348,7 +547,7 @@ async def stream_chat_completion_chunks( "index": open_tool["index"], "function": {"arguments": remainder}, }]}, - "index": 0, "finish_reason": None, + "index": index, "finish_reason": None, }], ) ) @@ -360,7 +559,7 @@ async def stream_chat_completion_chunks( yield _sse( _chat_chunk( req, uid, - [{"delta": {"tool_calls": [delta]}, "index": 0, "finish_reason": None}], + [{"delta": {"tool_calls": [delta]}, "index": index, "finish_reason": None}], ) ) tool_calls_sent += 1 @@ -368,21 +567,27 @@ async def stream_chat_completion_chunks( prompt_tokens = ev.prompt_tokens completion_tokens = ev.completion_tokens cached_tokens = ev.cached_tokens - yield _sse(_chat_chunk(req, uid, [{"delta": {}, "index": 0, "finish_reason": ev.finish_reason}])) + reasoning_tokens = ev.reasoning_tokens + yield _sse(_chat_chunk(req, uid, [{"delta": {}, "index": index, "finish_reason": ev.finish_reason}])) + timings = ev.timings + usage = _usage( + prompt_tokens, completion_tokens, _reported_cached(state, cached_tokens), + reasoning_tokens=reasoning_tokens, + ) + if not terminal: + yield _StreamUsage(usage) + return if req.stream_options and req.stream_options.include_usage: - yield _sse( - { - "id": f"chatcmpl-{uid}", - "object": "chat.completion.chunk", - "created": int(time.time()), - "model": req.model, - "choices": [], - "usage": _usage( - prompt_tokens, completion_tokens, _reported_cached(state, cached_tokens) - ), - } - ) + final: dict[str, Any] = {**_chat_chunk(req, uid, []), "usage": usage} + if metrics_enabled(state): + final["metrics"] = build_metrics( + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + cached_tokens=cached_tokens, + timings=timings, + ) + yield _sse(final) yield b"data: [DONE]\n\n" @@ -396,47 +601,64 @@ async def handle_completion( unsupported = _completion_unsupported_reason(req) if unsupported is not None: return create_error_response(unsupported) - try: # surfaces an out-of-range max_tokens as a 400 rather than a 500 from the worker + if not 1 <= req.n <= MAX_N: + return create_error_response(f"n must be between 1 and {MAX_N}", param="n") + try: # surfaces an out-of-range value as a 400 rather than a 500 from the worker default_max_tokens = ( getattr(state.config, "max_output_tokens", None) or DEFAULT_MAX_OUTPUT_TOKENS ) - _resolve_sampling(req, model_sampling, default_max_tokens=default_max_tokens) + sampling = _resolve_sampling(req, model_sampling, default_max_tokens=default_max_tokens) except ValueError as exc: - return create_error_response(str(exc), param="max_tokens") + return create_error_response(str(exc)) prompts = [req.prompt] if isinstance(req.prompt, str) else req.prompt assert isinstance(prompts, list) + if req.suffix is not None: + markers = await _fim_markers(state) + if markers is None: + return create_error_response( + "suffix needs a model with fill-in-the-middle tokens " + "(<|fim_prefix|> / <|fim_suffix|> / <|fim_middle|>)", + param="suffix", + ) + prompts = [_fim_prompt(p, req.suffix, markers) for p in prompts] + echo_texts = list(prompts) if req.echo else [""] * len(prompts) + if req.suffix is not None and req.echo: + # echo returns what the client sent, not the FIM-wrapped prompt + echo_texts = [req.prompt] if isinstance(req.prompt, str) else list(req.prompt) + + # choices are ordered prompt-major: prompt i, sample k -> index i * n + k (OpenAI) + uids: list[int] = [] + for prompt in prompts: + for _ in range(req.n): + uid = state.new_user() + await state.send_one(TokenizeMsg(uid=uid, text=prompt, sampling_params=sampling)) + uids.append(uid) + if req.stream: - if len(prompts) != 1: - return create_error_response("Streaming completions only support a single text prompt") - uid = state.new_user() - await state.send_one( - TokenizeMsg(uid=uid, text=prompts[0], sampling_params=_resolve_sampling( - req, model_sampling, default_max_tokens=default_max_tokens - )) - ) - chunks = stream_completion_chunks(uid, req, state) + if len(uids) == 1: + chunks = stream_completion_chunks(uids[0], req, state, echo_text=echo_texts[0]) + else: + chunks = _merge_streams( + [ + stream_completion_chunks( + u, req, state, index=i, terminal=False, echo_text=echo_texts[i // req.n] + ) + for i, u in enumerate(uids) + ], + lambda: _completion_chunk(req, uids[0], []), + include_usage=bool(req.stream_options and req.stream_options.include_usage), + state=state, + ) if request is not None: - chunks = state.stream_with_cancellation(chunks, request, uid) + chunks = state.stream_with_cancellation(chunks, request, uids) return StreamingResponse(chunks, media_type="text/event-stream") - choices: list[dict[str, Any]] = [] - prompt_tokens = 0 - completion_tokens = 0 - cached_tokens = 0 - for index, prompt in enumerate(prompts): - uid = state.new_user() - await state.send_one( - TokenizeMsg( - uid=uid, - text=prompt, - sampling_params=_resolve_sampling( - req, model_sampling, default_max_tokens=default_max_tokens - ), - ) - ) + async def _drain_one(uid: int, index: int) -> dict[str, Any] | JSONResponse: text = "" + entries: list[dict] = [] finish_reason = "stop" + prompt_tokens = completion_tokens = cached_tokens = 0 async for ack in state.wait_for_ack(uid): if getattr(ack, "error", None): return create_error_response(ack.error) @@ -444,80 +666,245 @@ async def handle_completion( completion_tokens += ack.completion_tokens_delta cached_tokens += ack.cached_tokens text += ack.incremental_output + if req.logprobs is not None and ack.logprobs is not None: + entries.append(ack.logprobs) if ack.finished: finish_reason = getattr(ack, "finish_reason", None) or "stop" break - choices.append({"index": index, "text": text, "finish_reason": finish_reason, "logprobs": None}) + return { + "index": index, + "text": echo_texts[index // req.n] + text, + "finish_reason": finish_reason, + "logprobs": completions_logprobs(entries) if req.logprobs is not None else None, + "_usage": (prompt_tokens, completion_tokens, cached_tokens), + } + + drained = await _await_watching_disconnect( + asyncio.gather(*(_drain_one(u, i) for i, u in enumerate(uids))), request, state, uids + ) + if drained is None: + return _client_disconnected_response() + choices: list[dict[str, Any]] = [] + prompt_tokens = completion_tokens = cached_tokens = 0 + for choice in drained: + if isinstance(choice, JSONResponse): + return choice + p, c, cached = choice.pop("_usage") + # the prompt is counted once per prompt, not once per sample of it + if choice["index"] % req.n == 0: + prompt_tokens += p + completion_tokens += c + cached_tokens = max(cached_tokens, cached) + choices.append(choice) return { - "id": f"cmpl-{uuid.uuid4().hex}", + "id": _response_id("cmpl", req, uuid.uuid4().hex), "object": "text_completion", "created": int(time.time()), "model": req.model, + "system_fingerprint": None, "choices": choices, "usage": _usage(prompt_tokens, completion_tokens, _reported_cached(state, cached_tokens)), } -async def stream_completion_chunks(uid: int, req: CompletionRequest, state: Any) -> AsyncIterator[bytes]: +def _completion_chunk(req: CompletionRequest, uid: Any, choices: list[dict[str, Any]]) -> dict[str, Any]: + return { + "id": _response_id("cmpl", req, uid), + "object": "text_completion.chunk", + "created": int(time.time()), + "model": req.model, + "choices": choices, + } + + +async def stream_completion_chunks( + uid: int, + req: CompletionRequest, + state: Any, + index: int = 0, + terminal: bool = True, + echo_text: str = "", +) -> AsyncIterator[bytes]: prompt_tokens = 0 + text_offset = 0 completion_tokens = 0 cached_tokens = 0 finish_reason = "stop" + if echo_text: + # echo: the prompt leads the stream (vLLM sends it in the first chunk) + yield _sse(_completion_chunk(req, uid, [{"text": echo_text, "index": index, "finish_reason": None, "logprobs": None}])) async for ack in state.wait_for_ack(uid): if getattr(ack, "error", None): yield _sse({"error": {"message": ack.error, "type": "invalid_request_error", "code": None}}) - yield b"data: [DONE]\n\n" + if terminal: + yield b"data: [DONE]\n\n" return prompt_tokens += ack.prompt_tokens_delta completion_tokens += ack.completion_tokens_delta cached_tokens += ack.cached_tokens if ack.incremental_output: yield _sse( - { - "id": f"cmpl-{uid}", - "object": "text_completion.chunk", - "created": int(time.time()), - "model": req.model, - "choices": [ - { - "text": ack.incremental_output, - "index": 0, - "finish_reason": None, - "logprobs": None, - } - ], - } + _completion_chunk( + req, uid, + [{ + "text": ack.incremental_output, + "index": index, + "finish_reason": None, + "logprobs": ( + completions_logprobs([ack.logprobs], text_offset) + if req.logprobs is not None and ack.logprobs is not None + else None + ), + }], + ) ) + text_offset += len(ack.incremental_output) if ack.finished: finish_reason = getattr(ack, "finish_reason", None) or "stop" break - yield _sse( - { - "id": f"cmpl-{uid}", - "object": "text_completion.chunk", - "created": int(time.time()), - "model": req.model, - "choices": [{"text": "", "index": 0, "finish_reason": finish_reason, "logprobs": None}], - } - ) + yield _sse(_completion_chunk(req, uid, [{"text": "", "index": index, "finish_reason": finish_reason, "logprobs": None}])) + usage = _usage(prompt_tokens, completion_tokens, _reported_cached(state, cached_tokens)) + if not terminal: + yield _StreamUsage(usage) + return if req.stream_options and req.stream_options.include_usage: - yield _sse( - { - "id": f"cmpl-{uid}", - "object": "text_completion.chunk", - "created": int(time.time()), - "model": req.model, - "choices": [], - "usage": _usage( - prompt_tokens, completion_tokens, _reported_cached(state, cached_tokens) - ), - } + yield _sse({**_completion_chunk(req, uid, []), "usage": usage}) + yield b"data: [DONE]\n\n" + + +class _StreamUsage: + """Yielded last by a non-terminal sub-stream: its usage dict for the merger.""" + + __slots__ = ("usage",) + + def __init__(self, usage: dict[str, Any]) -> None: + self.usage = usage + + +async def _merge_streams( + streams: list[AsyncIterator[Any]], + envelope: Callable[[], dict[str, Any]], + *, + include_usage: bool, + state: Any, +) -> AsyncIterator[bytes]: + """Interleave the SSE chunks of several choice streams into one response (``n > 1``), + then one usage chunk (prompt counted once, completions summed) and one ``[DONE]``.""" + queue: asyncio.Queue[tuple[int, Any]] = asyncio.Queue() + usages: dict[int, dict[str, Any]] = {} + + async def pump(slot: int, stream: AsyncIterator[Any]) -> None: + try: + async for item in stream: + if isinstance(item, _StreamUsage): + usages[slot] = item.usage + else: + await queue.put((slot, item)) + finally: + await queue.put((slot, None)) + + tasks = [asyncio.create_task(pump(i, st)) for i, st in enumerate(streams)] + try: + open_slots = len(streams) + while open_slots: + slot, item = await queue.get() + if item is None: + open_slots -= 1 + continue + yield item + finally: + for task in tasks: + task.cancel() + if include_usage and usages: + first = usages[min(usages)] + merged = _usage( + first["prompt_tokens"], + sum(u["completion_tokens"] for u in usages.values()), + max(u.get("prompt_tokens_details", {}).get("cached_tokens", 0) for u in usages.values()), + reasoning_tokens=sum( + u.get("completion_tokens_details", {}).get("reasoning_tokens", 0) for u in usages.values() + ), ) + yield _sse({**envelope(), "usage": merged}) yield b"data: [DONE]\n\n" +_FIM_TOKENS = ("<|fim_prefix|>", "<|fim_suffix|>", "<|fim_middle|>") + + +async def _fim_markers(state: Any) -> tuple[str, str, str] | None: + """The model's fill-in-the-middle marker strings when its vocabulary has them + (Qwen / DeepSeek-Coder style), else None.""" + build = getattr(state, "frontend_tokenizer", None) + if build is None: + return None + try: + manager = await asyncio.to_thread(build) + tok = manager.tokenizer + unk = getattr(tok, "unk_token_id", None) + for name in _FIM_TOKENS: + tid = tok.convert_tokens_to_ids(name) + if not isinstance(tid, int) or tid < 0 or tid == unk: + return None + except Exception: # noqa: BLE001 -- no tokenizer here means no FIM + return None + return _FIM_TOKENS + + +def _fim_prompt(prompt: str, suffix: str, markers: tuple[str, str, str]) -> str: + """The prefix-suffix-middle prompt of OpenAI's ``suffix`` (llama.cpp's /infill).""" + prefix_tok, suffix_tok, middle_tok = markers + return f"{prefix_tok}{prompt}{suffix_tok}{suffix}{middle_tok}" + + +async def handle_tokenize(req: TokenizeRequest, state: Any): + """``POST /tokenize``: the vLLM / SGLang shape. A raw ``prompt`` is encoded as is; chat + ``messages`` are rendered through the model's chat template first (the exact prompt a + generation would tokenize).""" + if (req.prompt is None) == (req.messages is None): + return create_error_response("pass exactly one of prompt or messages") + build = getattr(state, "frontend_tokenizer", None) + if build is None: + return create_error_response("no tokenizer on this server", status_code=503, err_type="api_error") + try: + manager = await asyncio.to_thread(build) + if req.prompt is not None: + ids = await asyncio.to_thread( + manager.tokenizer.encode, req.prompt, add_special_tokens=req.add_special_tokens + ) + else: + messages = render_messages([m.model_dump(exclude_none=True) for m in req.messages]) + ctk = dict(req.chat_template_kwargs) + if not req.add_generation_prompt: + ctk["continue_final_message"] = True + msg = TokenizeMsg(uid=-1, text=messages, sampling_params=SamplingParams(), chat_template_kwargs=ctk) + text = await asyncio.to_thread(manager.render_prompt, msg) + ids = await asyncio.to_thread(manager.tokenizer.encode, text, add_special_tokens=False) + except ValueError as exc: + return create_error_response(str(exc)) + body: dict[str, Any] = { + "count": len(ids), + "max_model_len": _model_context_length(state), + "tokens": [int(t) for t in ids], + } + if req.return_token_strs: + body["token_strs"] = manager.tokenizer.convert_ids_to_tokens(ids) + return body + + +async def handle_detokenize(req: DetokenizeRequest, state: Any): + build = getattr(state, "frontend_tokenizer", None) + if build is None: + return create_error_response("no tokenizer on this server", status_code=503, err_type="api_error") + manager = await asyncio.to_thread(build) + text = await asyncio.to_thread( + manager.tokenizer.decode, req.tokens, skip_special_tokens=req.skip_special_tokens + ) + return {"prompt": text} + + def create_error_response( message: str, status_code: int = 400, @@ -543,7 +930,7 @@ def _resolve_sampling( model_sampling: dict[str, Any], default_max_tokens: int = DEFAULT_MAX_OUTPUT_TOKENS, ) -> SamplingParams: - return resolve_sampling( + sampling_params = resolve_sampling( temperature=req.temperature, top_k=req.top_k, top_p=req.top_p, @@ -552,9 +939,27 @@ def _resolve_sampling( model_sampling=model_sampling, stop=req.stop, default_max_tokens=default_max_tokens, + min_p=req.min_p, presence_penalty=req.presence_penalty, frequency_penalty=req.frequency_penalty, + repetition_penalty=req.repetition_penalty, + logit_bias=req.logit_bias, + min_tokens=req.min_tokens, + stop_token_ids=req.stop_token_ids, + include_stop_str_in_output=req.include_stop_str_in_output, + skip_special_tokens=req.skip_special_tokens, ) + if isinstance(req, CompletionRequest) and req.logprobs is not None: + sampling_params.logprobs = True + sampling_params.top_logprobs = req.logprobs + return sampling_params + + +def _response_id(prefix: str, req: Any, fallback: Any) -> str: + """``request_id`` from the client when given (vLLM request_id / SGLang rid), else the + server's own id.""" + rid = getattr(req, "request_id", None) + return f"{prefix}-{rid}" if rid else f"{prefix}-{fallback}" def _tools_for_template(req: ChatCompletionRequest) -> list[dict[str, Any]] | None: @@ -617,10 +1022,11 @@ def _tool_call_id(name: str | None, index: int) -> str: def _chat_chunk(req: ChatCompletionRequest, uid: int, choices: list[dict[str, Any]]) -> dict[str, Any]: return { - "id": f"chatcmpl-{uid}", + "id": _response_id("chatcmpl", req, uid), "object": "chat.completion.chunk", "created": int(time.time()), "model": req.model, + "system_fingerprint": None, "choices": choices, } @@ -634,7 +1040,9 @@ def _reported_cached(state: Any, cached_tokens: int) -> int: return cached_tokens if getattr(state.config, "enable_cache_report", False) else 0 -def _usage(prompt_tokens: int, completion_tokens: int, cached_tokens: int = 0) -> dict[str, Any]: +def _usage( + prompt_tokens: int, completion_tokens: int, cached_tokens: int = 0, reasoning_tokens: int = 0 +) -> dict[str, Any]: usage: dict[str, Any] = { "prompt_tokens": prompt_tokens, "completion_tokens": completion_tokens, @@ -644,27 +1052,34 @@ def _usage(prompt_tokens: int, completion_tokens: int, cached_tokens: int = 0) - # disabled report and a 0-token hit serialize identically. if cached_tokens > 0: usage["prompt_tokens_details"] = {"cached_tokens": cached_tokens} + # OpenAI's completion_tokens_details.reasoning_tokens: tokens up to and including the + # reasoning end tag, only when the model emitted one. + if reasoning_tokens > 0: + usage["completion_tokens_details"] = {"reasoning_tokens": reasoning_tokens} return usage def _response_format_unsupported(response_format: dict[str, Any] | None) -> bool: - # We have no constrained/guided decoding; only plain text ('text' or unset) is honored. - return response_format is not None and response_format.get("type") not in (None, "text") + # Soft-accept OpenAI json_object/json_schema. FreeToken has no constrained + # decoding; clients that already ask for JSON in the prompt still need the + # wire field accepted (HTTP 200) instead of a hard 400. Unknown types 400. + if response_format is None: + return False + return response_format.get("type") not in (None, "text", "json_object", "json_schema") def _completion_unsupported_reason(req: CompletionRequest) -> str | None: if _is_token_prompt(req.prompt): return "OpenAI token-id prompt inputs are not supported; pass text prompt strings instead" - if req.logprobs is not None: - return "logprobs is not supported" - if req.echo: - return "echo is not supported" - if req.suffix is not None: - return "suffix is not supported" - if req.logit_bias is not None: - return "logit_bias is not supported" + logprobs_error = completion_logprobs_error(req) + if logprobs_error is not None: + return logprobs_error if _response_format_unsupported(req.response_format): - return "response_format json_object/json_schema is not supported (no constrained decoding)" + rtype = (req.response_format or {}).get("type") + return ( + f"response_format.type {rtype!r} is not supported " + "(supported: text, json_object, json_schema)" + ) return None @@ -705,10 +1120,31 @@ def _served_model_name(state: Any) -> str: def _model_context_length(state: Any) -> int | None: - """The model ceiling, not `min(ceiling, KV budget)`: a rebuild moves the latter, and agents - read this once at startup.""" + """The usable context: ``min(model ceiling, allocated KV tokens)``. + + This is the clamp the engine already applies to its own ``max_seq_len`` (Engine.__init__, + and again in _refresh_seq_state after a rebuild), and therefore the limit the scheduler + admits against. The frontend process holds an unclamped ServerArgs copy, so publishing + ``config.max_seq_len`` here advertised the model's positional ceiling instead -- e.g. 262144 + against a 178176-token pool, so the engine and this route disagreed about one quantity. + + `ft launch` reads this to size each agent's context window (opencode's ``limit.context``, + codex's ``context_window``, ``CLAUDE_CODE_MAX_CONTEXT_TOKENS``), so overstating it stops + those agents compacting before the KV pool runs out. + """ try: # never 500 a metadata route: max_seq_len walks into the HF config on some builds value = int(state.config.max_seq_len) except Exception: # noqa: BLE001 return None - return value if value > 0 else None + if value <= 0: + return None + try: + # Local import: api_server imports this module, so a module-level one would cycle. + from .api_server import kv_pool_geometry + + num_pages, page_size = kv_pool_geometry(state) + kv_tokens = num_pages * page_size + except Exception: # noqa: BLE001 + kv_tokens = 0 + # 0 before the ("meta", …) ack lands; the ceiling is the best available answer until then. + return min(value, kv_tokens) if kv_tokens > 0 else value diff --git a/python/freetoken/server/responses_api.py b/python/freetoken/server/responses_api.py index 71d6eed873..a84c68cb6d 100644 --- a/python/freetoken/server/responses_api.py +++ b/python/freetoken/server/responses_api.py @@ -69,14 +69,17 @@ ToolCallArgsDelta, ToolCallsDelta, ToolCallStart, + build_metrics, generate_events, generate_full, + metrics_enabled, render_messages, resolve_sampling, split_tool_lists, submit_generation, with_keepalive, ) +from .openai_api import _await_watching_disconnect from .request_logger import log_request # Seconds of event silence before a keep-alive frame is emitted on the stream. @@ -155,6 +158,7 @@ async def handle_responses( spec = convert_responses_to_genspec( req, model_sampling, default_max_tokens=default_max, reasoning_parser=getattr(state.config, "reasoning_parser", None), + default_thinking_mode=getattr(state.config, "default_thinking_mode", "auto"), ) uid = await submit_generation(spec, state) except GenerationError as exc: @@ -163,20 +167,28 @@ async def handle_responses( return _error_response(400, str(exc)) cache_report = getattr(state.config, "enable_cache_report", False) + metrics = metrics_enabled(state) if req.stream: events = responses_stream_generator( generate_events(uid, spec, state, source="/v1/responses"), req, response_id, created, - cache_report=cache_report, + cache_report=cache_report, metrics=metrics, ) if request is not None: events = state.stream_with_cancellation(events, request, uid) return StreamingResponse(events, media_type="text/event-stream") try: - result = await generate_full(uid, spec, state, source="/v1/responses") + # an abandoned request must not keep decoding to max_tokens (#222's watcher) + result = await _await_watching_disconnect( + generate_full(uid, spec, state, source="/v1/responses"), request, state, [uid] + ) except GenerationError as exc: return _error_response(400, str(exc), exc.code) - response = build_responses_response(result, req, response_id, created, cache_report=cache_report) + if result is None: + return _error_response(499, "client disconnected before the response was ready") + response = build_responses_response( + result, req, response_id, created, cache_report=cache_report, metrics=metrics + ) return JSONResponse(content=response.model_dump(mode="json")) @@ -190,6 +202,7 @@ def convert_responses_to_genspec( model_sampling: dict[str, Any], default_max_tokens: int = DEFAULT_MAX_OUTPUT_TOKENS, reasoning_parser: str | None = None, + default_thinking_mode: str | None = None, ) -> GenSpec: # Collect every system/developer text — the top-level `instructions` PLUS any # system/developer-role input items (codex sends both: a system prompt as `instructions` @@ -232,11 +245,12 @@ def convert_responses_to_genspec( else: template_tools, parser_tools = split_tool_lists(raw_tools, selected) - from .model_meta import effort_toggle_kwargs + from .model_meta import apply_default_thinking_mode, effort_toggle_kwargs ctk = dict(getattr(req, "chat_template_kwargs", None) or {}) if req.reasoning: ctk = effort_toggle_kwargs(req.reasoning.get("effort"), ctk) + ctk = apply_default_thinking_mode(ctk, default_thinking_mode) return GenSpec( messages=render_messages(messages), @@ -410,6 +424,7 @@ def build_responses_response( response_id: str, created: int, cache_report: bool = False, + metrics: bool = False, ) -> Response: truncated = result.finish_reason == "length" item_status = "incomplete" if truncated else "completed" @@ -454,15 +469,21 @@ def build_responses_response( result.cached_tokens if cache_report else 0, ), incomplete_reason="max_output_tokens" if truncated else None, + metrics=build_metrics( + prompt_tokens=result.prompt_tokens, + completion_tokens=result.completion_tokens, + cached_tokens=result.cached_tokens, + timings=result.timings, + ) if metrics else None, ) def _response_obj( response_id: str, created: int, model: str, output: list[Any], *, status: str, usage: ResponseUsage | None, error: ResponseError | None = None, - incomplete_reason: str | None = None, + incomplete_reason: str | None = None, metrics: dict[str, Any] | None = None, ) -> Response: - return Response( + response = Response( id=response_id, created_at=created, model=model, @@ -476,6 +497,11 @@ def _response_obj( tool_choice="auto", tools=[], ) + if metrics is not None: + # The SDK model allows extras, so the FreeToken-only `metrics` object rides the real + # Response type rather than forcing a parallel dict-shaped response path. + response.metrics = metrics + return response def _usage(prompt_tokens: int, completion_tokens: int, cached_tokens: int = 0) -> ResponseUsage: @@ -499,13 +525,15 @@ async def responses_stream_generator( response_id: str, created: int, cache_report: bool = False, + metrics: bool = False, ) -> AsyncIterator[str]: seq = _Seq() - def snapshot(status, output, usage=None, incomplete_reason=None): + def snapshot(status, output, usage=None, incomplete_reason=None, request_metrics=None): return _response_obj( response_id, created, req.model, output, status=status, usage=usage, incomplete_reason=incomplete_reason, + metrics=request_metrics, ) yield _sse(ResponseCreatedEvent( @@ -707,6 +735,12 @@ def args_delta_frame(fragment: str) -> str: finish_reason = ev.finish_reason usage_pt, usage_ct = ev.prompt_tokens, ev.completion_tokens usage_cached = ev.cached_tokens if cache_report else 0 + request_metrics = build_metrics( + prompt_tokens=ev.prompt_tokens, + completion_tokens=ev.completion_tokens, + cached_tokens=ev.cached_tokens, + timings=ev.timings, + ) if metrics else None for f in close_current(): yield f if finish_reason == "length": @@ -716,6 +750,7 @@ def args_delta_frame(fragment: str) -> str: "incomplete", output_items, usage=_usage(usage_pt, usage_ct, usage_cached), incomplete_reason="max_output_tokens", + request_metrics=request_metrics, ), )) else: @@ -724,6 +759,7 @@ def args_delta_frame(fragment: str) -> str: response=snapshot( "completed", output_items, usage=_usage(usage_pt, usage_ct, usage_cached), + request_metrics=request_metrics, ), )) except GenerationError as exc: diff --git a/python/freetoken/server/stats.py b/python/freetoken/server/stats.py index 3d48ee96a0..c9d0212c32 100644 --- a/python/freetoken/server/stats.py +++ b/python/freetoken/server/stats.py @@ -174,3 +174,44 @@ def build_stats(state: Any, p95_ms: int, ttft_mean_ms: int) -> dict: "completion_tokens_total": tr.completion_tokens_total, }, } + + + +def prometheus_text(doc: dict) -> str: + """The /v1/stats document as Prometheus text exposition (gauges and counters).""" + lines: list[str] = [] + + def gauge(name: str, value, help_text: str, kind: str = "gauge") -> None: + if value is None: + return + lines.append(f"# HELP {name} {help_text}") + lines.append(f"# TYPE {name} {kind}") + lines.append(f"{name} {value}") + + model = doc.get("model") or {} + model_id = str(model.get("id") or model.get("name") or "").replace('"', "'") + lines.append("# HELP freetoken_info Server identity.") + lines.append("# TYPE freetoken_info gauge") + lines.append(f'freetoken_info{{instance_id="{doc.get("instance_id")}",model="{model_id}"}} 1') + gauge("freetoken_uptime_seconds", doc.get("uptime_s"), "Seconds since the server became ready.") + req = doc.get("requests") or {} + gauge("freetoken_requests_active", req.get("active"), "Requests in flight.") + gauge("freetoken_requests_completed_total", req.get("completed"), "Requests completed.", "counter") + gauge("freetoken_request_latency_p95_milliseconds", req.get("p95_ms"), "p95 request latency over the recent window.") + gauge("freetoken_ttft_mean_milliseconds", req.get("ttft_mean_ms"), "Mean time to first token over the recent window.") + gauge("freetoken_prompt_tokens_total", req.get("prompt_tokens_total"), "Prompt tokens processed.", "counter") + gauge("freetoken_completion_tokens_total", req.get("completion_tokens_total"), "Tokens generated.", "counter") + tp = doc.get("throughput") or {} + gauge("freetoken_decode_tokens_per_second", tp.get("decode_tps"), "Decode throughput, sliding window.") + gauge("freetoken_prefill_tokens_per_second", tp.get("prefill_tps"), "Prefill throughput, sliding window.") + for pool in ("kv", "swa"): + p = doc.get(pool) + if p: + gauge(f"freetoken_{pool}_pages_used", p.get("used_pages"), f"{pool} pages in use.") + gauge(f"freetoken_{pool}_pages_total", p.get("total_pages"), f"{pool} pages in the pool.") + mamba = doc.get("mamba") + if mamba: + gauge("freetoken_mamba_slots_used", mamba.get("used_slots"), "GDN state slots in use.") + gauge("freetoken_mamba_slots_total", mamba.get("total_slots"), "GDN state slots in the pool.") + gauge("freetoken_vram_bytes", doc.get("vram_bytes"), "GPU memory in use by the engine.") + return "\n".join(lines) + "\n" diff --git a/python/freetoken/server/supervisor.py b/python/freetoken/server/supervisor.py index 322419a65f..20deb7d132 100644 --- a/python/freetoken/server/supervisor.py +++ b/python/freetoken/server/supervisor.py @@ -55,6 +55,12 @@ class BackendHandle: processes: List[Any] = field(default_factory=list) expected_acks: int = 0 + def release_acks(self) -> None: + """Drop the ack queue once readiness has read it for the last time. A multiprocessing queue's + semaphores are unlinked only when it is collected, and a main process that a signal ends never + collects it, so resource_tracker reports them as leaked.""" + self.ack_queue = None + class WorkerDied(Exception): def __init__(self, message: str) -> None: @@ -172,6 +178,7 @@ def _shutting_down() -> bool: on_failure(exc.message) return + handle.release_acks() on_ready() while True: diff --git a/python/freetoken/tokenizer/detokenize.py b/python/freetoken/tokenizer/detokenize.py index d138ac1a68..2eb3029df0 100644 --- a/python/freetoken/tokenizer/detokenize.py +++ b/python/freetoken/tokenizer/detokenize.py @@ -71,11 +71,34 @@ class DecodeStatus: read_offset: int # length of read ids surr_offset: int # length of surr ids sent_offset: int # length of sent out string + skip_special: bool = False # decode this request with skip_special_tokens + think_end_at: int | None = None # generated tokens up to and incl. the reasoning end tag + + def update(self, read_str: str, surr_str: str) -> str: + new_text = read_str[len(surr_str) :] + if len(new_text) > 0 and not new_text.endswith("�"): + self.decoded_str += new_text + self.surr_offset = self.read_offset + self.read_offset = len(self.decoded_ids) + return self.decoded_str + return self.decoded_str + find_printable_text(new_text) + + def decode(self, tokenizer: PreTrainedTokenizerBase) -> str: + skip = self.skip_special + return self.update( + tokenizer.decode(self.decoded_ids[self.surr_offset :], skip_special_tokens=skip), + tokenizer.decode( + self.decoded_ids[self.surr_offset : self.read_offset], skip_special_tokens=skip + ), + ) class DetokenizeManager: def __init__( - self, tokenizer: PreTrainedTokenizerBase, eos_token_ids: FrozenSet[int] | None = None + self, + tokenizer: PreTrainedTokenizerBase, + eos_token_ids: FrozenSet[int] | None = None, + think_end_id: int | None = None, ) -> None: # uid -> DecodeStatus self.decode_map: Dict[int, DecodeStatus] = {} @@ -83,6 +106,9 @@ def __init__( self.eos_token_ids = ( eos_token_ids if eos_token_ids is not None else frozenset({tokenizer.eos_token_id}) ) + # The reasoning end tag's id (e.g. ````): tokens generated up to and + # including it are reported as usage.completion_tokens_details.reasoning_tokens. + self.think_end_id = think_end_id if think_end_id is None else int(think_end_id) def discard(self, uid: int) -> None: """Drop a uid's decode state without a finished DetokenizeMsg. An aborted or errored @@ -91,8 +117,30 @@ def discard(self, uid: int) -> None: self.decode_map.pop(uid, None) def detokenize(self, msgs: List[DetokenizeMsg]) -> List[str]: + return self.detokenize_with_meta(msgs)[0] + + def _batch_decode(self, ids: List[List[int]], skip: List[bool]) -> List[str]: + """batch_decode honouring a per-request skip_special_tokens flag (one call per + distinct flag value, results back in order).""" + out: List[str | None] = [None] * len(ids) + for flag in (False, True): + idx = [i for i, f in enumerate(skip) if f == flag] + if idx: + texts = self.tokenizer.batch_decode([ids[i] for i in idx], skip_special_tokens=flag) + for i, t in zip(idx, texts, strict=True): + out[i] = t + return [t if t is not None else "" for t in out] + + def detokenize_with_meta(self, msgs: List[DetokenizeMsg]) -> tuple[List[str], List[int]]: + """Incremental text per message plus, for a finished message, its reasoning token + count (0 otherwise or when the model emitted no reasoning end tag).""" + # Each message must advance its request's decode state before the next one. + if len({msg.uid for msg in msgs}) != len(msgs): + parts = [self.detokenize_with_meta([msg]) for msg in msgs] + return [p[0][0] for p in parts], [p[1][0] for p in parts] read_ids: List[List[int]] = [] surr_ids: List[List[int]] = [] + skip: List[bool] = [] for msg in msgs: if msg.uid not in self.decode_map: self.decode_map[msg.uid] = DecodeStatus( @@ -101,36 +149,44 @@ def detokenize(self, msgs: List[DetokenizeMsg]) -> List[str]: read_offset=0, surr_offset=0, sent_offset=0, + skip_special=bool(getattr(msg, "skip_special_tokens", False)), ) s = self.decode_map[msg.uid] - if not (msg.finished and msg.next_token in self.eos_token_ids): + if not ( + msg.finished and not msg.matched_stop and msg.next_token in self.eos_token_ids + ): s.decoded_ids.append(msg.next_token) + if ( + self.think_end_id is not None + and s.think_end_at is None + and msg.next_token == self.think_end_id + ): + s.think_end_at = len(s.decoded_ids) read_ids.append(s.decoded_ids[s.surr_offset :]) surr_ids.append(s.decoded_ids[s.surr_offset : s.read_offset]) + skip.append(s.skip_special) - read_texts = self.tokenizer.batch_decode(read_ids) - surr_texts = self.tokenizer.batch_decode(surr_ids) + read_texts = self._batch_decode(read_ids, skip) + surr_texts = self._batch_decode(surr_ids, skip) incremental_strs: List[str] = [] + reasoning_tokens: List[int] = [] for msg, read_str, surr_str in zip(msgs, read_texts, surr_texts, strict=True): s = self.decode_map[msg.uid] - new_text = read_str[len(surr_str) :] - # Streaming chunk: update the decode status - if len(new_text) > 0 and not new_text.endswith("�"): - output_str = s.decoded_str + new_text - s.decoded_str = output_str - s.surr_offset = s.read_offset - s.read_offset = len(s.decoded_ids) - else: - new_text = find_printable_text(new_text) - output_str = s.decoded_str + new_text + output_str = s.update(read_str, surr_str) prev_sent = s.sent_offset if msg.finished: - # Generation is over: flush everything, trimming at the matched stop string. + # Generation is over: flush everything, trimming at the matched stop string + # (or keeping it, for include_stop_str_in_output). if msg.matched_stop: cut = output_str.find(msg.matched_stop) - emit_end = cut if cut >= 0 else len(output_str) + if cut < 0: + emit_end = len(output_str) + elif getattr(msg, "keep_stop_str", False): + emit_end = cut + len(msg.matched_stop) + else: + emit_end = cut else: emit_end = len(output_str) elif msg.stop_strs: @@ -141,7 +197,37 @@ def detokenize(self, msgs: List[DetokenizeMsg]) -> List[str]: incremental_output = output_str[prev_sent:emit_end] if emit_end > prev_sent else "" s.sent_offset = max(prev_sent, emit_end) incremental_strs.append(incremental_output) + reasoning_tokens.append((s.think_end_at or 0) if msg.finished else 0) if msg.finished: del self.decode_map[msg.uid] - return incremental_strs + return incremental_strs, reasoning_tokens + + +def build_logprobs_entry( + tokenizer: PreTrainedTokenizerBase, + token_id: int, + chosen_logprob: float, + top_ids: list[int] | None, + top_logprobs: list[float] | None, +) -> dict: + """Neutral sampled-token logprobs entry (UserReply.logprobs). Values are the raw + pre-temperature logprobs the sampler computed; token text comes from a single-id + decode, with UTF-8 bytes alongside so clients can reassemble partial-UTF-8 pieces.""" + + def _fields(tid: int) -> tuple[str, list[int]]: + text = tokenizer.decode([tid]) + return text, list(text.encode("utf-8")) + + token, token_bytes = _fields(token_id) + top = [] + for tid, logprob in zip(top_ids or [], top_logprobs or []): + text, data = _fields(int(tid)) + top.append({"token_id": int(tid), "token": text, "bytes": data, "logprob": float(logprob)}) + return { + "token_id": token_id, + "token": token, + "bytes": token_bytes, + "logprob": float(chosen_logprob), + "top": top, + } diff --git a/python/freetoken/tokenizer/server.py b/python/freetoken/tokenizer/server.py index d3365bfecd..4c650042b4 100644 --- a/python/freetoken/tokenizer/server.py +++ b/python/freetoken/tokenizer/server.py @@ -81,6 +81,19 @@ def _send_generation_replies( _put_user_replies(send_frontend, terminal_errors) +def _reasoning_end_id(tokenizer) -> int | None: + """The id of the reasoning end tag when the vocabulary has one (```` for the + Qwen / DeepSeek / GLM families), else None. Only used to count reasoning tokens.""" + for tag in ("", "<|end_of_thinking|>"): + try: + tid = tokenizer.convert_tokens_to_ids(tag) + except Exception: # noqa: BLE001 -- a tokenizer without the tag answers unk or raises + continue + if isinstance(tid, int) and tid >= 0 and tid != getattr(tokenizer, "unk_token_id", -1): + return tid + return None + + def _tokenize_requests( tokenize_manager: Any, messages: List[TokenizeMsg], @@ -146,12 +159,14 @@ def tokenize_worker( from freetoken.mm.processor import get_mm_processor - from .detokenize import DetokenizeManager + from .detokenize import DetokenizeManager, build_logprobs_entry from .tokenize import TokenizeManager tokenize_manager = TokenizeManager(tokenizer, get_mm_processor(tokenizer_path, mm)) detokenize_manager = DetokenizeManager( - tokenizer, load_eos_token_ids(tokenizer_path, tokenizer) + tokenizer, + load_eos_token_ids(tokenizer_path, tokenizer), + think_end_id=_reasoning_end_id(tokenizer), ) if ack_queue is not None: @@ -209,7 +224,7 @@ def tokenize_worker( ) sampled_replies: List[UserReply] = [] if len(detokenize_msg) > 0: - replies = detokenize_manager.detokenize(detokenize_msg) + replies, reasoning_counts = detokenize_manager.detokenize_with_meta(detokenize_msg) sampled_replies = [ UserReply( uid=msg.uid, @@ -217,6 +232,7 @@ def tokenize_worker( finished=msg.finished, finish_reason=msg.finish_reason, matched_stop=msg.matched_stop, + reasoning_tokens=reasoning_tokens, completion_tokens_delta=1, kv_used_pages=msg.kv_used_pages, kv_total_pages=msg.kv_total_pages, @@ -225,8 +241,24 @@ def tokenize_worker( swa_used_tokens=msg.swa_used_tokens, swa_total_tokens=msg.swa_total_tokens, gpu_mem_bytes=msg.gpu_mem_bytes, + prefill_ms=msg.prefill_ms, + # Stop-string trimming can hide final visible text; keep one logprob + # entry per sampled token. + logprobs=( + build_logprobs_entry( + detokenize_manager.tokenizer, + msg.next_token, + msg.chosen_logprob, + msg.top_ids, + msg.top_logprobs, + ) + if msg.chosen_logprob is not None + else None + ), + ) + for msg, reply, reasoning_tokens in zip( + detokenize_msg, replies, reasoning_counts, strict=True ) - for msg, reply in zip(detokenize_msg, replies, strict=True) ] # An error reply and a client abort are both terminal for their uid, and neither diff --git a/python/freetoken/tokenizer/tokenize.py b/python/freetoken/tokenizer/tokenize.py index 772472573a..31b74bdac2 100644 --- a/python/freetoken/tokenizer/tokenize.py +++ b/python/freetoken/tokenizer/tokenize.py @@ -38,8 +38,14 @@ def resolve_thinking_mode(chat_template_kwargs: dict[str, Any] | None, tools: An one implementation prevents the two sides from disagreeing. Thinking is on when tools are offered (dsv4 only emits well-formed tool calls in thinking mode) or when the caller requests it via ``chat_template_kwargs``. + An explicit ``enable_thinking=False`` (or ``thinking=False``) from the + caller always wins, even when tools are present — the caller knows their + intent better than the heuristic. """ ctk = chat_template_kwargs or {} + # Explicit disable takes precedence over the tools-presumed-thinking heuristic. + if ctk.get("enable_thinking") is False or ctk.get("thinking") is False: + return "chat" mode = str(ctk.get("thinking_mode") or "chat") if tools or ctk.get("enable_thinking") or ctk.get("thinking"): mode = "thinking" @@ -51,6 +57,36 @@ def resolve_thinking_mode(chat_template_kwargs: dict[str, Any] | None, tools: An _EFFORT_PROBE_MESSAGES = [{"role": "user", "content": "ping"}] +def join_system_contents(messages: list[dict[str, Any]]) -> str: + """The contents of several system/developer messages as one system prompt, in order. + Image content (a part list left by text flattening) cannot be joined into text: raise + the error templates give for an image in the system turn.""" + contents = [m.get("content") for m in messages if m.get("content")] + if any(not isinstance(c, str) for c in contents): + raise ValueError("System message cannot contain images") + return "\n\n".join(contents) + + +def _map_developer_role( + messages: list[dict[str, Any]], chat_template: str | None +) -> list[dict[str, Any]]: + """OpenAI's ``developer`` role is the current spelling of ``system`` (clients such as pi + send the system prompt that way by default). Most chat templates do not know it and raise + ("Unexpected message role"), so map it to ``system`` unless the template handles it.""" + if "developer" in (chat_template or "") or not any( + m.get("role") == "developer" for m in messages + ): + return messages + mapped = [{**m, "role": "system"} if m.get("role") == "developer" else m for m in messages] + # The frontend hoisted and merged system messages before this mapping ran (templates such + # as Qwen's raise on a system message that is not first), so do it again for the new ones. + system = [m for m in mapped if m.get("role") == "system"] + if len(system) == 1 and mapped[0] is system[0]: + return mapped + rest = [m for m in mapped if m.get("role") != "system"] + return [{"role": "system", "content": join_system_contents(system)}] + rest + + class TokenizeManager: def __init__(self, tokenizer: PreTrainedTokenizerBase, mm_processor: MMProcessor | None = None) -> None: self.tokenizer = tokenizer @@ -156,10 +192,17 @@ def _render( ) if tools is not None: chat_template_kwargs = {**chat_template_kwargs, "tools": tools} + messages = _map_developer_role(messages, getattr(self.tokenizer, "chat_template", None)) + # continue_final_message (an OpenAI-compatible extra, as in vLLM / SGLang): the last + # message is an assistant prefix the model must continue, so no generation prompt. + chat_template_kwargs = dict(chat_template_kwargs) + continue_final = bool(chat_template_kwargs.pop("continue_final_message", False)) + if continue_final: + chat_template_kwargs["continue_final_message"] = True prompt = self.tokenizer.apply_chat_template( messages, tokenize=False, - add_generation_prompt=True, + add_generation_prompt=not continue_final, **chat_template_kwargs, ) assert isinstance(prompt, str) diff --git a/python/freetoken/utils/__init__.py b/python/freetoken/utils/__init__.py index af68bf5816..a1dd10b7c3 100644 --- a/python/freetoken/utils/__init__.py +++ b/python/freetoken/utils/__init__.py @@ -9,6 +9,7 @@ ) from .hf import ( cached_load_hf_config, + download_hf_checkpoint, download_hf_weight, load_eos_token_ids, load_generation_sampling, @@ -30,6 +31,7 @@ __all__ = [ "cached_load_hf_config", + "download_hf_checkpoint", "download_hf_weight", "load_eos_token_ids", "load_generation_sampling", diff --git a/python/freetoken/utils/hf.py b/python/freetoken/utils/hf.py index 489585e8c2..6c6a5f7a00 100644 --- a/python/freetoken/utils/hf.py +++ b/python/freetoken/utils/hf.py @@ -2,6 +2,7 @@ import functools import json import os +import re from typing import Any, FrozenSet, Mapping from huggingface_hub import hf_hub_download, snapshot_download @@ -296,4 +297,46 @@ def download_hf_weight(model_path: str) -> str: except Exception as e: raise ValueError( f"Model path '{model_path}' is neither a local directory nor a valid model ID: {e}" - ) \ No newline at end of file + ) + + +# Weight formats FreeToken never loads -- it reads safetensors (plus GGUF for Gemma-4). +# Plenty of repos still ship a duplicate .bin/.pth copy of the same tensors, so skipping +# them can halve a first-time pull. +_UNLOADABLE_WEIGHTS = ("*.bin", "*.pt", "*.pth", "*.h5", "*.msgpack", "*.onnx") + + +# A hub repo id is "org/name" and nothing else. Anything with a leading separator, a +# "~", or more than one "/" is a filesystem path the user got wrong, and turning that +# into a hub lookup replaces "no such directory" with a confusing download error. +_HUB_REPO_ID = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]*/[A-Za-z0-9][A-Za-z0-9._-]*$") + + +def looks_like_hub_repo_id(model_path: str) -> bool: + """Whether ``model_path`` should be resolved against the Hub rather than the disk.""" + return bool(_HUB_REPO_ID.match(model_path)) + + +def download_hf_checkpoint(model_path: str, *, dummy_weight: bool = False) -> str: + """Resolve a Hugging Face repo id to a complete local checkpoint directory. + + ``download_hf_weight`` fetches only ``*.safetensors``, which is all the weight + loaders need once the engine is running. Config parsing happens earlier and wants + more than that: the tokenizer, and the model-specific subdirs some architectures + read directly off disk -- DeepSeek-V4 fails without ``inference/config.json`` + (see ``models/deepseek_v4/args.py``). So this pulls the whole snapshot. + + A path that is already a local directory is returned untouched. + """ + if os.path.isdir(model_path) or not looks_like_hub_repo_id(model_path): + return model_path + ignore_patterns = list(_UNLOADABLE_WEIGHTS) + if dummy_weight: + # Config and tokenizer are still needed; the tensors are not. + ignore_patterns += ["*.safetensors", "*.gguf"] + try: + return snapshot_download(model_path, ignore_patterns=ignore_patterns) + except Exception as e: + raise ValueError( + f"Model path '{model_path}' is neither a local directory nor a valid model ID: {e}" + ) diff --git a/python/freetoken/utils/mp.py b/python/freetoken/utils/mp.py index 181baf1a07..ee35c057c1 100644 --- a/python/freetoken/utils/mp.py +++ b/python/freetoken/utils/mp.py @@ -1,6 +1,6 @@ from __future__ import annotations -from typing import Callable, Dict, Generic, TypeVar +from typing import Any, Callable, Dict, Generic, TypeVar import msgpack import zmq @@ -9,6 +9,38 @@ T = TypeVar("T") +class _CoalescedUnpacker: + """ + Buffered msgpack decoder tolerant of coalesced frames. + + A single ZMQ frame should carry exactly one msgpack object, but on some platforms + (issue #452: Windows, offload MoE) a read can surface a frame containing two packed + objects back to back. ``msgpack.unpackb`` raises ``ExtraData`` in that case and kills + the worker process. Feeding every received frame through one ``Unpacker`` instead + decodes the first object and buffers the remainder, so the next ``get()`` returns it + in order instead of crashing. + """ + + def __init__(self) -> None: + # 0 = effectively uncapped (2**31-1), as unpackb had; the Unpacker default of 100 MiB would + # kill the worker on a large message (e.g. an image's pixel values) + self._unpacker = msgpack.Unpacker(raw=False, max_buffer_size=0) + self._pending: list[Any] = [] + + def feed(self, frame: bytes) -> None: + """Queue up every msgpack object contained in ``frame`` (usually exactly one).""" + self._unpacker.feed(frame) + for obj in self._unpacker: + self._pending.append(obj) + + def take(self) -> Any: + """Pop the oldest undelivered object; raises StopIteration when the buffer is empty.""" + return self._pending.pop(0) + + def __len__(self) -> int: + return len(self._pending) + + class ZmqPushQueue(Generic[T]): def __init__( self, @@ -62,19 +94,30 @@ def __init__( self.socket = self.context.socket(zmq.PULL) self.socket.bind(addr) if create else self.socket.connect(addr) self.decoder = decoder + self._unpacker = _CoalescedUnpacker() def get(self) -> T: - event = self.socket.recv() - return self.decoder(msgpack.unpackb(event, raw=False)) + if not len(self._unpacker): + event = self.socket.recv() + self._unpacker.feed(event) + obj = self._unpacker.take() + return self.decoder(obj) def get_raw(self) -> bytes: + # Raw path (multi-rank broadcast) must stay unbuffered: callers pair empty()/get_raw() + # with a rank-wide count, so a buffered remainder would desynchronize the loop. + if len(self._unpacker): + raise RuntimeError( + "get_raw() cannot be mixed with buffered get() on the same queue: " + "the unpacker holds undelivered objects that would be skipped." + ) return self.socket.recv() def decode(self, raw: bytes) -> T: return self.decoder(msgpack.unpackb(raw, raw=False)) def empty(self) -> bool: - return self.socket.poll(timeout=0) == 0 + return len(self._unpacker) == 0 and self.socket.poll(timeout=0) == 0 def stop(self): self.socket.close() @@ -92,10 +135,14 @@ def __init__( self.socket = self.context.socket(zmq.PULL) self.socket.bind(addr) if create else self.socket.connect(addr) self.decoder = decoder + self._unpacker = _CoalescedUnpacker() async def get(self) -> T: - event = await self.socket.recv() - return self.decoder(msgpack.unpackb(event, raw=False)) + if not len(self._unpacker): + event = await self.socket.recv() + self._unpacker.feed(event) + obj = self._unpacker.take() + return self.decoder(obj) def stop(self): self.socket.close() @@ -138,14 +185,18 @@ def __init__( self.socket.bind(addr) if create else self.socket.connect(addr) self.socket.setsockopt_string(zmq.SUBSCRIBE, "") self.decoder = decoder + self._unpacker = _CoalescedUnpacker() def get(self) -> T: - event = self.socket.recv() - return self.decoder(msgpack.unpackb(event, raw=False)) + if not len(self._unpacker): + event = self.socket.recv() + self._unpacker.feed(event) + obj = self._unpacker.take() + return self.decoder(obj) def empty(self) -> bool: - return self.socket.poll(timeout=0) == 0 + return len(self._unpacker) == 0 and self.socket.poll(timeout=0) == 0 def stop(self): self.socket.close() - self.context.term() + self.context.term() \ No newline at end of file diff --git a/setup.py b/setup.py index b8e4693c89..b968373151 100644 --- a/setup.py +++ b/setup.py @@ -77,15 +77,22 @@ def _cuda_runtime_paths() -> tuple[list[str], list[str]]: if IS_ROCM: runtime_include_dirs, runtime_library_dirs, runtime_lib = _rocm_paths() runtime_link_args = [f"-Wl,-rpath,{runtime_library_dirs[0]}"] - # These extensions contain host code only. BuildExtension supplies the ROCm + # Host sources. BuildExtension supplies the ROCm # platform defines to the C++ compiler; offload architecture flags belong on # HIP device sources and would be rejected by the host compiler here. extra_compile = ["-O3", "-std=c++20"] + # Device sources (.cu, built by hipcc); BuildExtension adds the offload architectures. + device_compile = ["-O3", "-std=c++20"] else: runtime_include_dirs, runtime_library_dirs = _cuda_runtime_paths() runtime_lib = "cudart" runtime_link_args = [] extra_compile = ["-O3", "-std=c++20"] + # Device sources (.cu, built by nvcc): SASS for each major arch, which also runs on the newer minors of that major, + # plus PTX for the oldest arch nvcc 13 supports, which the driver builds for any other GPU at load. + # Naming archs also stops torch from picking them, which fails on a build machine without a GPU. + device_compile = ["-O3", "-std=c++20"] + [f"-gencode=arch=compute_{a},code=sm_{a}" for a in (75, 80, 90, 100, 120)] + device_compile.append("-gencode=arch=compute_75,code=compute_75") _check_toolchain() @@ -112,11 +119,17 @@ def _cuda_runtime_paths() -> tuple[list[str], list[str]]: name="freetoken.kernel._cpu_moe", sources=[ "python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp", + f"python/freetoken/kernel/csrc/cpu_moe/gpu_sync_{'hip' if IS_ROCM else 'cuda'}.cpp", + "python/freetoken/kernel/csrc/cpu_moe/gpu_sync_io.cu", + *sorted( + str(p.relative_to(ROOT)) + for p in (ROOT / "python/freetoken/kernel/csrc/cpu_moe/kernels").glob("*.cpp") + ), ], include_dirs=[KERNEL_INCLUDE, *runtime_include_dirs], library_dirs=runtime_library_dirs, libraries=[runtime_lib], - extra_compile_args=extra_compile + ["-pthread"], + extra_compile_args={"cxx": extra_compile + ["-pthread"], "nvcc": device_compile}, extra_link_args=runtime_link_args, ), # disk-backed row store (PLE / Engram tables); Linux-only until the TableFile/BatchReader seams grow Windows bodies diff --git a/tests/README.md b/tests/README.md index adc87035d9..d892c36dfa 100644 --- a/tests/README.md +++ b/tests/README.md @@ -65,6 +65,7 @@ checkpoint is set: | `FREETOKEN_TEST_MODELS_ROOT` | `models/test_quant_config.py` — directory whose checkpoints the stored-tensor scheme scan checks (default `/mnt/nvme/models`) | | `FREETOKEN_DSV4_VL_MODEL` | `tokenizer/test_mm_tokenize.py` - local DeepSeek-V4 vision checkpoint | | `FREETOKEN_DSV4_VISION_CKPT` | `models/test_dsv4_vision.py` - the same checkpoint, for the parity tests against its `inference/` reference code | +| `FREETOKEN_MOE_ROUTING_TRACE` | `moe/test_scored_ensure.py` — a compact decode routing trace directory (`experts.npy`, `logit_idx.npy`, `logit_val.npy`); its rows are replayed through the scored eviction kernel and the CPU reference | `test_aime.py` takes its sampling protocol from the checkpoint's own `generation_config.json` (pass@N at the recommended temperature, or a single greedy diff --git a/tests/attention/test_dsa_kpool.py b/tests/attention/test_dsa_kpool.py index fb856e1757..bed1bbdaef 100644 --- a/tests/attention/test_dsa_kpool.py +++ b/tests/attention/test_dsa_kpool.py @@ -49,8 +49,8 @@ def _args(num_layers=1): ) -@pytest.fixture() -def harness(monkeypatch): +@pytest.fixture(params=["none", "nvfp4"]) +def harness(monkeypatch, request): from freetoken.attention.dsa_indexer_kpool import Glm5NextDSABackend from freetoken.kvcache.dsa_pool import KpoolDSAKVCache @@ -58,7 +58,7 @@ def harness(monkeypatch): latent_dim=LATENT, num_layers=1, num_pages=8, page_size=64, dtype=torch.bfloat16, device=torch.device(DEV), index_head_dim=DI, num_index_layers=1, - index_ratio=KPOOL, num_req_slots=4, + index_ratio=KPOOL, num_req_slots=4, kv_quant=request.param, ) page_table = torch.full((4, 512), -1, dtype=torch.int32, device=DEV) page_table[0, :512] = torch.arange(512, dtype=torch.int32, device=DEV) @@ -76,6 +76,23 @@ def harness(monkeypatch): return backend, pool, ape +@pytest.mark.parametrize("kv_quant", ["fp8", "nvfp4"]) +def test_quantized_latent_cache_keeps_kpool_index_tiers_bf16(kv_quant): + from freetoken.kvcache.dsa_pool import KpoolDSAKVCache + + pool = KpoolDSAKVCache( + latent_dim=LATENT, num_layers=1, num_pages=8, page_size=64, + dtype=torch.bfloat16, device=torch.device(DEV), + index_head_dim=DI, num_index_layers=1, + index_ratio=KPOOL, num_req_slots=4, kv_quant=kv_quant, + ) + assert pool.latent_rows(0).dtype is torch.uint8 + assert pool.latent_scale(0).dtype is torch.float32 + assert pool.index_k_cache(0).dtype is torch.bfloat16 + assert pool.tail_k(0).dtype is torch.bfloat16 + assert pool.tail_gate(0).dtype is torch.bfloat16 + + def _req(device_len, cached_len=0): return SimpleNamespace( table_idx=0, device_len=device_len, extend_len=device_len - cached_len, @@ -113,10 +130,15 @@ def _rand_seq(total, seed=1): def _run(backend, batch, d, sl, ape): + d["kv_quant"] = backend.kvcache.kv_quant + backend.prepare_metadata(batch) + return _forward(backend, batch, d, sl, ape) + + +def _forward(backend, batch, d, sl, ape): from freetoken.attention.dsa import DSAIndexerInputs t = batch.positions.shape[0] - backend.prepare_metadata(batch) return backend.mla_forward( d["q_nope"][sl], d["q_nope"].new_empty(t, H, 0), d["c_kv"][sl], d["c_kv"].new_empty(t, 0), @@ -149,6 +171,14 @@ def _ref_scores(d, ape, q_idx_t, w_t, n_pools): def _ref_attend(d, q_t, positions): """Full softmax MLA over latent rows at ``positions`` for one query [H, LATENT].""" lat = d["c_kv"][positions].float() # [n, LATENT] + if d.get("kv_quant") == "nvfp4": + from tests.kernels.test_kv_nvfp4 import _reference + + packed, block, row = _reference(lat) + codes = torch.stack((packed & 15, packed >> 4), -1).flatten(-2).long() + grid = lat.new_tensor([0, .5, 1, 1.5, 2, 3, 4, 6, + 0, -.5, -1, -1.5, -2, -3, -4, -6]) + lat = grid[codes] * block.view(torch.float8_e4m3fn).float().repeat_interleave(16, -1) * row[:, None] logits = q_t.float() @ lat.T * SM_SCALE # [H, n] p = torch.softmax(logits, dim=-1) return (p @ lat).to(torch.bfloat16) @@ -367,3 +397,52 @@ def test_padding_and_empty_batch_leave_shadow_rows_clean(): ratio=KPOOL, ) assert torch.equal(slab, before) + + +def test_decode_graph_replay_tracks_slots_and_lengths(harness): + backend, pool, ape = harness + total, extra = 60, 6 + d = _rand_seq(total + extra, seed=81) + _run(backend, _prefill_batch(0, total), d, slice(0, total), ape) + static = {k: v[total:total + 1].clone() for k, v in d.items() if isinstance(v, torch.Tensor)} + batch = _decode_batch(total) + batch.size = batch.padded_size = 1 + backend.init_capture_graph(512, [1]) + backend.prepare_for_capture(batch) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(3): + _forward(backend, batch, static, slice(None), ape) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = _forward(backend, batch, static, slice(None), ape) + base = 0 + for pos in range(total, total + extra): + if pos == total + 2: + from freetoken.attention.dsa import get_global_ctx + + base = 256 + pool.latent_rows(0)[base:base + 128].copy_(pool.latent_rows(0)[:128]) + if pool.kv_quant == "nvfp4": + pool.latent_scale(0)[base:base + 128].copy_(pool.latent_scale(0)[:128]) + pool.latent_block_scale(0)[base:base + 128].copy_(pool.latent_block_scale(0)[:128]) + pool.index_k_cache(0)[64:96].copy_(pool.index_k_cache(0)[:32]) + pool.tail_k(0)[1].copy_(pool.tail_k(0)[0]) + pool.tail_gate(0)[1].copy_(pool.tail_gate(0)[0]) + get_global_ctx().page_table[1, :128] = torch.arange(base, base + 128, device=DEV) + batch.active_table_idx.fill_(1) + batch.padded_reqs[0].table_idx = 1 + for key, tensor in static.items(): + tensor.copy_(d[key][pos:pos + 1]) + batch.positions.fill_(pos) + batch.out_loc.fill_(base + pos) + batch.padded_reqs[0].device_len = pos + 1 + backend.prepare_metadata(batch) + backend.prepare_for_replay(batch) + graph.replay() + selected = _ref_selected_positions(d, ape, d["qi"][pos], d["wi"][pos], pos) + ref = _ref_attend(d, d["q_nope"][pos], selected) + torch.testing.assert_close(out[0], ref, atol=3e-2, rtol=1e-2) + backend.reset_capture() diff --git a/tests/checkpoint/__init__.py b/tests/checkpoint/__init__.py new file mode 100644 index 0000000000..5b8db6dafd --- /dev/null +++ b/tests/checkpoint/__init__.py @@ -0,0 +1 @@ +"""Checkpoint-format tests.""" diff --git a/tests/checkpoint/test_convert_refusal.py b/tests/checkpoint/test_convert_refusal.py new file mode 100644 index 0000000000..6c09f79ad8 --- /dev/null +++ b/tests/checkpoint/test_convert_refusal.py @@ -0,0 +1,60 @@ +"""convert_checkpoint refuses an output directory that already holds FTW output, before it writes anything. + +No GPU: the refusal sits ahead of the TP/CUDA setup, so a CPU box reaches it with an empty source dir.""" + +import os + +import pytest +import torch + +from freetoken.checkpoint.convert import _refuse_occupied_output, convert_checkpoint +from freetoken.checkpoint.ftw import INDEX_NAME, FTWWriter + + +def _listing(d): + return sorted((name, os.stat(os.path.join(d, name)).st_size, os.stat(os.path.join(d, name)).st_mtime_ns) + for name in os.listdir(d)) + + +def _refused(src, out): + before = _listing(out) + with pytest.raises(SystemExit, match="already holds an FTW checkpoint") as ei: + convert_checkpoint(str(src), str(out)) + assert _listing(out) == before # nothing written, nothing touched + return str(ei.value) + + +def test_finished_ftw_in_the_output_dir_is_refused_by_name(tmp_path): + out = tmp_path / "out" + writer = FTWWriter(str(out)) + writer.add_tensor("model.a.weight", torch.zeros(4, dtype=torch.bfloat16)) + writer.finalize({}) + src = tmp_path / "src" + src.mkdir() + msg = _refused(src, out) + assert str(out) in msg and INDEX_NAME in msg + + +def test_shards_from_an_interrupted_conversion_are_refused_by_name(tmp_path): + out = tmp_path / "out" + out.mkdir() + (out / "freetoken-00001.ftw").write_bytes(b"\0" * 64) + (out / "freetoken-00000.ftw").write_bytes(b"\0" * 64) + src = tmp_path / "src" + src.mkdir() + msg = _refused(src, out) + assert str(out) in msg and "freetoken-00000.ftw" in msg and INDEX_NAME not in msg + + +def test_empty_absent_or_unrelated_output_dirs_pass(tmp_path): + empty = tmp_path / "empty" + empty.mkdir() + assert _refuse_occupied_output(str(empty)) is None + absent = tmp_path / "absent" + assert _refuse_occupied_output(str(absent)) is None + assert not absent.exists() # the check creates nothing + other = tmp_path / "other" + other.mkdir() + (other / "config.json").write_text("{}") # copied metadata alone is not FTW output + (other / "notes.ftw.bak").write_bytes(b"") + assert _refuse_occupied_output(str(other)) is None diff --git a/tests/checkpoint/test_ftw_validation.py b/tests/checkpoint/test_ftw_validation.py new file mode 100644 index 0000000000..0f9acb4463 --- /dev/null +++ b/tests/checkpoint/test_ftw_validation.py @@ -0,0 +1,640 @@ +from __future__ import annotations + +import copy +import json +import mmap +import os +from pathlib import Path +import subprocess +import sys +import types + +import pytest +import torch + +from freetoken.checkpoint.ftw import ( + ALIGN, + FORMAT_TAG, + FORMAT_VERSION, + INDEX_NAME, + FTWFormatError, + FTWReader, + FTWWriter, + iter_ftw_weights, +) + +MAX_INDEX_INT = (1 << 63) - 1 + + +def _write_index(path, index) -> None: + with open(path / INDEX_NAME, "w") as f: + json.dump(index, f) + + +def _write_raw_index(path, data: bytes) -> None: + with open(path / INDEX_NAME, "wb") as f: + f.write(data) + + +def _two_tensor_checkpoint(tmp_path): + path = tmp_path / "checkpoint" + writer = FTWWriter(str(path), shard_limit=ALIGN) + writer.add_tensor("first", torch.arange(1000, dtype=torch.float32)) # 4000 + padding + writer.add_tensor("second", torch.arange(4, dtype=torch.int16)) # 8 + padding + return path, writer.finalize({"model_type": "test"}) + + +def _three_tensor_checkpoint(tmp_path): + path = tmp_path / "checkpoint" + writer = FTWWriter(str(path), shard_limit=ALIGN) + for name in ("first", "middle", "last"): + writer.add_tensor(name, torch.arange(8, dtype=torch.uint8)) + return path, writer.finalize({}) + + +def _expect_invalid(path, index, match: str) -> None: + _write_index(path, index) + with pytest.raises(FTWFormatError, match=match): + FTWReader(str(path)) + + +def test_valid_writer_output_reads_tensor_spanning_shards(tmp_path): + path = tmp_path / "checkpoint" + expected = (torch.arange(5000, dtype=torch.int64) % 251).to(torch.uint8) + writer = FTWWriter(str(path), shard_limit=ALIGN) + writer.add_tensor("spanning", expected) + index = writer.finalize({}) + + assert len(index["shards"]) == 2 + reader = FTWReader(str(path)) + reader._direct = 0 # make this portable across tmpfs/overlayfs and CI kernels + entry = reader.tensors["spanning"] + pieces = list(reader._pieces(entry["global_off"], entry["nbytes"])) + assert [piece[3] for piece in pieces] == [ALIGN, 5000 - ALIGN] + + buf = mmap.mmap(-1, 2 * ALIGN) + dest = memoryview(buf) + try: + reader.read_into(dest, entry, workers=1, chunk=ALIGN) + actual = torch.frombuffer(buf, dtype=torch.uint8, count=5000).clone() + assert torch.equal(actual, expected) + finally: + dest.release() + buf.close() + reader.close() + + +def test_valid_empty_and_zero_tensor_writer_outputs(tmp_path): + empty_path = tmp_path / "empty" + empty_index = FTWWriter(str(empty_path), shard_limit=ALIGN).finalize({}) + assert empty_index["shards"] == [] + empty_reader = FTWReader(str(empty_path)) + assert empty_reader.shards == [] + empty_reader.close() + + zero_path = tmp_path / "zero" + zero_writer = FTWWriter(str(zero_path), shard_limit=ALIGN) + zero_writer.add_tensor("zero", torch.empty(0, dtype=torch.float32)) + zero_index = zero_writer.finalize({}) + assert zero_index["shards"][0]["nbytes"] == 0 + zero_reader = FTWReader(str(zero_path)) + assert zero_reader.tensors["zero"]["nbytes"] == 0 + zero_reader.close() + + +def test_valid_zero_tensors_at_writer_allocation_boundaries(tmp_path): + path = tmp_path / "mixed-zero" + writer = FTWWriter(str(path), shard_limit=ALIGN) + writer.add_tensor("zero-before", torch.empty(0, dtype=torch.float32)) + writer.add_tensor("first", torch.arange(1000, dtype=torch.float32)) + writer.add_tensor("zero-between", torch.empty((2, 0, 3), dtype=torch.int16)) + writer.add_tensor("second", torch.arange(4, dtype=torch.int16)) + writer.add_tensor("zero-after", torch.empty(0, dtype=torch.uint8)) + index = writer.finalize({}) + + reader = FTWReader(str(path)) + assert [reader.tensors[name]["global_off"] for name in ( + "zero-before", "zero-between", "zero-after", + )] == [0, ALIGN, 2 * ALIGN] + assert index["total_bytes"] == 2 * ALIGN + reader.close() + + +def test_zero_tensor_can_be_iterated(tmp_path, monkeypatch): + path = tmp_path / "zero-iteration" + writer = FTWWriter(str(path), shard_limit=ALIGN) + writer.add_tensor("zero", torch.empty((2, 0, 3), dtype=torch.float32)) + writer.finalize({}) + + class _Progress: + def update(self, _nbytes): + pass + + def close(self): + pass + + progress = types.ModuleType("freetoken.utils.progress") + progress.byte_bar = lambda *_args, **_kwargs: _Progress() + monkeypatch.setitem(sys.modules, "freetoken.utils.progress", progress) + + loaded = list(iter_ftw_weights(str(path), workers=1, chunk=ALIGN, prefetch=1)) + assert len(loaded) == 1 + assert loaded[0][0] == "zero" + assert loaded[0][1].shape == (2, 0, 3) + + +def test_scalar_tensor_round_trips_with_scalar_shape(tmp_path, monkeypatch): + path = tmp_path / "scalar-iteration" + writer = FTWWriter(str(path), shard_limit=ALIGN) + writer.add_tensor("scalar", torch.tensor(17, dtype=torch.int32)) + writer.finalize({}) + + class _Progress: + def update(self, _nbytes): + pass + + def close(self): + pass + + progress = types.ModuleType("freetoken.utils.progress") + progress.byte_bar = lambda *_args, **_kwargs: _Progress() + monkeypatch.setitem(sys.modules, "freetoken.utils.progress", progress) + + loaded = list(iter_ftw_weights(str(path), workers=1, chunk=ALIGN, prefetch=1)) + assert loaded[0][0] == "scalar" + assert loaded[0][1].shape == () + assert loaded[0][1].item() == 17 + + +def test_shards_may_be_declared_out_of_order(tmp_path): + path, index = _two_tensor_checkpoint(tmp_path) + index["shards"].reverse() + _write_index(path, index) + + reader = FTWReader(str(path)) + assert [shard["global_off"] for shard in reader.shards] == [0, ALIGN] + reader.close() + + +@pytest.mark.parametrize( + ("field", "value", "match"), + [ + ("format", "other", "not a freetoken_weight"), + ("version", FORMAT_VERSION + 1, "unsupported FTW version"), + ("version", True, "version must be an integer"), + ("align", ALIGN * 2, "unsupported FTW alignment"), + ("shard_limit", ALIGN + 1, "shard_limit must be a multiple"), + ("shard_limit", 0, "shard_limit must be an integer"), + ("total_bytes", ALIGN + 1, "total_bytes must be 4096-aligned"), + ("total_bytes", -1, "total_bytes must be an integer"), + ("total_bytes", MAX_INDEX_INT + 1, "total_bytes must be an integer"), + ], +) +def test_rejects_invalid_header_fields(tmp_path, field, value, match): + path, index = _two_tensor_checkpoint(tmp_path) + index[field] = value + _expect_invalid(path, index, match) + + +def test_rejects_non_object_root_and_missing_fields(tmp_path): + path, index = _two_tensor_checkpoint(tmp_path) + _expect_invalid(path, [], "index root must be an object") + + index.pop("tensors") + _expect_invalid(path, index, "missing required field 'tensors'") + + +@pytest.mark.parametrize("raw", [b"{", b"\xff"]) +def test_rejects_malformed_json_and_encoding(tmp_path, raw): + path, _index = _two_tensor_checkpoint(tmp_path) + _write_raw_index(path, raw) + with pytest.raises(FTWFormatError, match="malformed FTW index"): + FTWReader(str(path)) + + +def test_rejects_duplicate_json_object_keys(tmp_path): + path, _index = _two_tensor_checkpoint(tmp_path) + _write_raw_index( + path, + b'{"format":"freetoken_weight","version":1,"version":1}', + ) + with pytest.raises(FTWFormatError, match="duplicate JSON key 'version'"): + FTWReader(str(path)) + + +@pytest.mark.parametrize("constant", ["NaN", "Infinity", "-Infinity"]) +def test_rejects_non_standard_json_constants_even_in_metadata(tmp_path, constant): + path, index = _two_tensor_checkpoint(tmp_path) + raw = json.dumps(index).removesuffix("}") + f', "metadata_value": {constant}}}' + _write_raw_index(path, raw.encode()) + with pytest.raises(FTWFormatError, match="non-standard JSON constant"): + FTWReader(str(path)) + + +def test_rejects_json_float_that_overflows_python_float(tmp_path): + path, index = _two_tensor_checkpoint(tmp_path) + raw = json.dumps(index).removesuffix("}") + ', "metadata_value": 1e999}' + _write_raw_index(path, raw.encode()) + with pytest.raises(FTWFormatError, match="floating-point value is out of range"): + FTWReader(str(path)) + + +@pytest.mark.parametrize( + ("field", "value", "match"), + [ + ("shards", {}, "shards must be an array"), + ("tensors", {}, "tensors must be an array"), + ], +) +def test_rejects_non_array_index_collections(tmp_path, field, value, match): + path, index = _two_tensor_checkpoint(tmp_path) + index[field] = value + _expect_invalid(path, index, match) + + +@pytest.mark.parametrize( + "file", + ["../outside.ftw", "/tmp/outside.ftw", "sub/file.ftw", "sub\\file.ftw", "..", "bad\x00.ftw"], +) +def test_rejects_unsafe_shard_names_before_filesystem_access(tmp_path, file): + path, index = _two_tensor_checkpoint(tmp_path) + index["shards"][0]["file"] = file + _expect_invalid(path, index, "safe shard basename") + + +def test_rejects_duplicate_shard_names(tmp_path): + path, index = _two_tensor_checkpoint(tmp_path) + index["shards"][1]["file"] = index["shards"][0]["file"] + _expect_invalid(path, index, "duplicate shard file") + + +@pytest.mark.parametrize( + ("entry", "match"), + [ + (1, r"shards\[0\] must be an object"), + ({"global_off": 0, "nbytes": ALIGN}, "missing required field 'file'"), + ({"file": 1, "global_off": 0, "nbytes": ALIGN}, "safe shard basename"), + ], +) +def test_rejects_malformed_shard_entries(tmp_path, entry, match): + path, index = _two_tensor_checkpoint(tmp_path) + index["shards"][0] = entry + _expect_invalid(path, index, match) + + +def test_rejects_missing_truncated_extended_and_non_regular_shards_eagerly(tmp_path): + path, index = _two_tensor_checkpoint(tmp_path) + first = path / index["shards"][0]["file"] + first.unlink() + _expect_invalid(path, index, "cannot stat FTW shard") + + path, index = _two_tensor_checkpoint(tmp_path / "truncated") + first = path / index["shards"][0]["file"] + with open(first, "r+b") as f: + f.truncate(ALIGN - 1) + _expect_invalid(path, index, "size mismatch") + + path, index = _two_tensor_checkpoint(tmp_path / "extended") + first = path / index["shards"][0]["file"] + with open(first, "ab") as f: + f.write(b"\x00") + _expect_invalid(path, index, "size mismatch") + + path, index = _two_tensor_checkpoint(tmp_path / "directory") + first = path / index["shards"][0]["file"] + first.unlink() + first.mkdir() + _expect_invalid(path, index, "not a regular file") + + +def test_accepts_symlinked_shard_whose_target_matches(tmp_path): + # Hugging Face hub snapshots are symlink farms (every file links into blobs/), so a + # shard entry that resolves through a symlink to a matching regular file must load. + path, index = _two_tensor_checkpoint(tmp_path) + blobs = tmp_path / "blobs" + blobs.mkdir() + first = path / index["shards"][0]["file"] + target = blobs / "first-blob" + target.write_bytes(first.read_bytes()) + first.unlink() + first.symlink_to(target) + + reader = FTWReader(str(path)) + reader._direct = 0 + buf = mmap.mmap(-1, ALIGN) + dest = memoryview(buf) + try: + reader.read_into(dest, reader.tensors["first"], workers=1, chunk=ALIGN) + actual = torch.frombuffer(buf, dtype=torch.float32, count=1000).clone() + assert torch.equal(actual, torch.arange(1000, dtype=torch.float32)) + finally: + dest.release() + buf.close() + reader.close() + + +def test_symlinked_shard_is_checked_against_its_target(tmp_path): + path, index = _two_tensor_checkpoint(tmp_path) + external = tmp_path / "outside.ftw" + external.write_bytes(b"\x00" * (ALIGN + 1)) + first = path / index["shards"][0]["file"] + first.unlink() + first.symlink_to(external) + + _expect_invalid(path, index, "size mismatch") + + +def test_rejects_shard_name_that_is_not_a_filesystem_name(tmp_path): + # A JSON-escaped lone surrogate decodes to a valid str that os.stat cannot encode; + # it must surface as FTWFormatError, not UnicodeEncodeError. + path, index = _two_tensor_checkpoint(tmp_path) + index["shards"][0]["file"] = "\ud800.ftw" + + _expect_invalid(path, index, "cannot stat FTW shard") + + +@pytest.mark.parametrize( + ("mutation", "match"), + [ + (lambda index: index["shards"][1].update(global_off=3 * ALIGN), "coverage has a gap"), + (lambda index: index["shards"][1].update(global_off=0), "coverage has an overlap"), + (lambda index: index.update(total_bytes=3 * ALIGN), "shards cover .* total_bytes"), + (lambda index: index["shards"][0].update(global_off=1), "global_off must be 4096-aligned"), + (lambda index: index["shards"][0].update(global_off=MAX_INDEX_INT + 1), "global_off must be an integer"), + (lambda index: index["shards"][0].update( + global_off=(MAX_INDEX_INT // ALIGN) * ALIGN, + ), "byte range exceeds signed 64-bit"), + (lambda index: index["shards"][0].update(nbytes=2 * ALIGN), "exceeds shard_limit"), + (lambda index: index["shards"][0].update(nbytes=ALIGN - 1), "nbytes must be 4096-aligned"), + (lambda index: index["shards"][0].update(nbytes=True), "nbytes must be an integer"), + ], +) +def test_rejects_invalid_shard_geometry(tmp_path, mutation, match): + path, index = _two_tensor_checkpoint(tmp_path) + mutation(index) + _expect_invalid(path, index, match) + + +def test_rejects_zero_length_shard_in_nonempty_stream(tmp_path): + path, index = _two_tensor_checkpoint(tmp_path) + empty_name = "empty.ftw" + (path / empty_name).touch() + index["shards"].insert(0, {"file": empty_name, "global_off": 0, "nbytes": 0}) + _expect_invalid(path, index, "zero-length shards") + + +def test_rejects_non_writer_zero_byte_shard_forms(tmp_path): + empty_path = tmp_path / "empty" + empty_index = FTWWriter(str(empty_path), shard_limit=ALIGN).finalize({}) + (empty_path / "empty.ftw").touch() + empty_index["shards"] = [{"file": "empty.ftw", "global_off": 0, "nbytes": 0}] + _expect_invalid(empty_path, empty_index, "empty checkpoint must not contain shard") + + zero_path = tmp_path / "zero" + zero_writer = FTWWriter(str(zero_path), shard_limit=ALIGN) + zero_writer.add_tensor("zero", torch.empty(0, dtype=torch.float32)) + zero_index = zero_writer.finalize({}) + zero_index["shards"] = [] + _expect_invalid(zero_path, zero_index, "zero-sized tensor checkpoint must contain") + + +@pytest.mark.parametrize( + ("mutation", "match"), + [ + (lambda index: index["tensors"][1].update(name="first"), "duplicate tensor name"), + (lambda index: index["tensors"][0].update(name=1), "name must be a string"), + (lambda index: index["tensors"][0].update(kind=1), "kind must be a string"), + (lambda index: index["tensors"][0].update(dtype="not_a_dtype"), "not a known torch dtype"), + (lambda index: index["tensors"][0].update(dtype="strided"), "not a known torch dtype"), + (lambda index: index["tensors"][0].update(dtype=1), "dtype must be a string"), + (lambda index: index["tensors"][0].update(shape="1000"), "shape must be an array"), + (lambda index: index["tensors"][0].update(shape=[True]), r"shape\[0\] must be an integer"), + (lambda index: index["tensors"][0].update(shape=[1.0]), r"shape\[0\] must be an integer"), + (lambda index: index["tensors"][0].update(shape=[-1]), r"shape\[0\] must be an integer"), + (lambda index: index["tensors"][0].update(shape=[(1 << 63) - 1, 2]), "element count exceeds"), + (lambda index: index["tensors"][0].update( + shape=[0, MAX_INDEX_INT, MAX_INDEX_INT], nbytes=0, + ), "shape cannot be reconstructed by torch"), + (lambda index: index["tensors"][0].update(shape=[MAX_INDEX_INT], dtype="float32"), "byte size exceeds"), + (lambda index: index["tensors"][0].update(nbytes=4), "shape .* require 4000"), + (lambda index: index["tensors"][0].update(global_off=1), "global_off must be 4096-aligned"), + (lambda index: index["tensors"][0].update( + dtype="uint8", shape=[ALIGN], nbytes=ALIGN, + global_off=(MAX_INDEX_INT // ALIGN) * ALIGN, + ), "byte range exceeds signed 64-bit"), + (lambda index: index["tensors"][1].update(global_off=2 * ALIGN), "exceeds total_bytes"), + (lambda index: index["tensors"][0].update(nbytes=True), "nbytes must be an integer"), + ], +) +def test_rejects_invalid_tensor_entries(tmp_path, mutation, match): + path, index = _two_tensor_checkpoint(tmp_path) + mutation(index) + _expect_invalid(path, index, match) + + +@pytest.mark.parametrize( + ("entry", "match"), + [ + (1, r"tensors\[0\] must be an object"), + ({"kind": "weight", "dtype": "float32", "shape": [1], + "global_off": 0, "nbytes": 4}, "missing required field 'name'"), + ], +) +def test_rejects_malformed_tensor_entries(tmp_path, entry, match): + path, index = _two_tensor_checkpoint(tmp_path) + index["tensors"][0] = entry + _expect_invalid(path, index, match) + + +def test_rejects_dtype_that_torch_cannot_reconstruct(tmp_path, monkeypatch): + path, index = _two_tensor_checkpoint(tmp_path) + original = torch.frombuffer + + def reject_float32(buffer, *, dtype, count): + if dtype == torch.float32: + raise RuntimeError("unsupported storage dtype") + return original(buffer, dtype=dtype, count=count) + + monkeypatch.setattr(torch, "frombuffer", reject_float32) + _expect_invalid(path, index, "dtype is not readable FTW storage") + + +def test_rejects_tensor_overlap_including_padded_allocation(tmp_path): + path, index = _two_tensor_checkpoint(tmp_path) + index["tensors"][1]["global_off"] = 0 + _expect_invalid(path, index, "overlaps padded allocation") + + +@pytest.mark.parametrize( + ("remove", "match"), + [ + (0, "tensor coverage has a gap"), + (1, "tensor coverage has a gap"), + (2, "tensors cover .* total_bytes"), + ], +) +def test_rejects_leading_internal_and_trailing_tensor_gaps(tmp_path, remove, match): + path, index = _three_tensor_checkpoint(tmp_path) + index["tensors"].pop(remove) + _expect_invalid(path, index, match) + + +def test_rejects_zero_sized_tensor_inside_nonzero_allocation(tmp_path): + path = tmp_path / "checkpoint" + writer = FTWWriter(str(path), shard_limit=ALIGN) + writer.add_tensor("spanning", torch.arange(5000, dtype=torch.uint8)) + index = writer.finalize({}) + index["tensors"].append({ + "name": "zero-inside", + "kind": "weight", + "dtype": "float32", + "shape": [0], + "global_off": ALIGN, + "nbytes": 0, + }) + _expect_invalid(path, index, "zero-sized tensor .* is not on a tensor boundary") + + +def test_index_validation_does_not_mutate_tensor_order(tmp_path): + path, index = _two_tensor_checkpoint(tmp_path) + index["tensors"].reverse() + expected = copy.deepcopy(index["tensors"]) + _write_index(path, index) + + reader = FTWReader(str(path)) + assert reader.index["tensors"] == expected + assert [entry["name"] for entry in reader.entries()] == ["second", "first"] + reader.close() + + +def test_writer_emits_declared_v1_header(tmp_path): + path, index = _two_tensor_checkpoint(tmp_path) + assert index["format"] == FORMAT_TAG + assert index["version"] == FORMAT_VERSION + assert index["align"] == ALIGN + assert sum(os.path.getsize(path / shard["file"]) for shard in index["shards"]) == index["total_bytes"] + + +@pytest.mark.parametrize("shard_limit", [0, ALIGN - 1, ALIGN + 1, True, MAX_INDEX_INT + 1]) +def test_writer_rejects_invalid_shard_limit_without_assertions(tmp_path, shard_limit): + with pytest.raises(ValueError, match="shard_limit must be an integer multiple"): + FTWWriter(str(tmp_path / "checkpoint"), shard_limit=shard_limit) + + +def test_writer_does_not_emit_indexes_its_reader_would_reject(tmp_path): + path = tmp_path / "checkpoint" + writer = FTWWriter(str(path), shard_limit=ALIGN) + writer.add_tensor("tensor", torch.arange(4, dtype=torch.float32)) + + with pytest.raises(ValueError, match="duplicate FTW tensor name"): + writer.add_tensor("tensor", torch.arange(4, dtype=torch.float32)) + with pytest.raises(ValueError, match="cannot override index fields"): + writer.finalize({"version": FORMAT_VERSION}) + with pytest.raises(ValueError, match="must not contain NaN or infinity"): + writer.finalize({"metadata_value": float("nan")}) + with pytest.raises(ValueError, match="keys must be strings"): + writer.finalize({"nested": {1: "ambiguous with a JSON string key"}}) + with pytest.raises(ValueError, match="JSON-compatible values"): + writer.finalize({"metadata_value": object()}) + + writer.finalize({"metadata_value": 1.5}) + FTWReader(str(path)).close() + + +def test_validation_remains_active_under_optimized_python(tmp_path): + repo_root = Path(__file__).resolve().parents[2] + script = r''' +import json +import logging +from pathlib import Path +import sys +import types + +import torch + +utils = types.ModuleType("freetoken.utils") +utils.init_logger = logging.getLogger +sys.modules["freetoken.utils"] = utils + +from freetoken.checkpoint.ftw import FTWFormatError, FTWReader, FTWWriter, INDEX_NAME, load_ftw_banks + +moe = types.ModuleType("freetoken.moe") +moe.__path__ = [] +host_banks = types.ModuleType("freetoken.moe.host_banks") +host_banks.HostBank = object +host_banks.HostResidency = types.SimpleNamespace(PINNED=types.SimpleNamespace(value="pinned")) +host_banks.PinPipeline = object +host_banks.alloc_banks = lambda _specs: {} +host_banks.born_pinned_default = lambda: False +progress = types.ModuleType("freetoken.utils.progress") +progress.byte_bar = lambda *_args, **_kwargs: None +sys.modules["freetoken.moe"] = moe +sys.modules["freetoken.moe.host_banks"] = host_banks +sys.modules["freetoken.utils.progress"] = progress + +path = Path(sys.argv[1]) / "optimized" +path.mkdir() +(path / INDEX_NAME).write_text(json.dumps({"format": "wrong"})) + +try: + FTWReader(str(path)) +except FTWFormatError: + pass +else: + raise SystemExit("malformed index was accepted under -O") + +for action in ( + lambda: FTWWriter(str(path / "writer"), shard_limit=0), + lambda: load_ftw_banks(str(path), num_layers=0), + lambda: load_ftw_banks(str(path), num_layers=1, layer_residency=[]), +): + try: + action() + except ValueError: + pass + else: + raise SystemExit("invalid writer/reader argument was accepted under -O") + +schema_cases = [] + +mixed = path / "mixed" +writer = FTWWriter(str(mixed), shard_limit=4096) +writer.add_tensor("bank", torch.ones(1), kind="experts_bank") +writer.add_tensor("bank#L00000", torch.ones(1), kind="experts_bank") +writer.finalize({}) +schema_cases.append((mixed, 1)) + +nondivisible = path / "nondivisible" +writer = FTWWriter(str(nondivisible), shard_limit=4096) +writer.add_tensor("bank", torch.ones(3), kind="experts_bank") +writer.finalize({}) +schema_cases.append((nondivisible, 2)) + +missing_layer = path / "missing-layer" +writer = FTWWriter(str(missing_layer), shard_limit=4096) +writer.add_tensor("bank#L00000", torch.ones(1), kind="experts_bank") +writer.finalize({}) +schema_cases.append((missing_layer, 2)) + +for checkpoint, layers in schema_cases: + try: + load_ftw_banks(str(checkpoint), num_layers=layers) + except FTWFormatError: + pass + else: + raise SystemExit("malformed expert-bank schema was accepted under -O") +''' + env = os.environ.copy() + env["PYTHONPATH"] = os.pathsep.join(filter(None, [ + str(repo_root / "python"), env.get("PYTHONPATH", ""), + ])) + result = subprocess.run( + [sys.executable, "-O", "-c", script, str(tmp_path)], + cwd=repo_root, + env=env, + capture_output=True, + text=True, + timeout=30, + check=False, + ) + assert result.returncode == 0, result.stdout + result.stderr diff --git a/tests/engine/test_attention_backend_arch.py b/tests/engine/test_attention_backend_arch.py index 25eb93547c..db62d907eb 100644 --- a/tests/engine/test_attention_backend_arch.py +++ b/tests/engine/test_attention_backend_arch.py @@ -110,6 +110,42 @@ def test_explicit_trtllm_allowed_on_sm100_family(monkeypatch): assert config.page_size in (16, 32, 64) +def test_explicit_fa_rejected_outside_sm90_family(monkeypatch): + from freetoken.engine.engine import _adjust_config + + _patch_env(monkeypatch, major=8) + config = _engine_config(attention_backend="fa") + with pytest.raises(RuntimeError, match="9.x"): + _adjust_config(config) + + +def test_explicit_fa_allowed_on_sm90_family(monkeypatch): + from freetoken.engine.engine import _adjust_config + + _patch_env(monkeypatch, major=9) + config = _engine_config(attention_backend="fa") + _adjust_config(config) + assert config.attention_backend == "fa" + + +def test_explicit_fa_allowed_on_sm100_family(monkeypatch): + from freetoken.engine.engine import _adjust_config + + _patch_env(monkeypatch, major=10) + config = _engine_config(attention_backend="fa") + _adjust_config(config) + assert config.attention_backend == "fa" + + +def test_explicit_fa_rejected_on_consumer_blackwell(monkeypatch): + from freetoken.engine.engine import _adjust_config + + _patch_env(monkeypatch, major=12) + config = _engine_config(attention_backend="fa") + with pytest.raises(RuntimeError, match="9.x"): + _adjust_config(config) + + def test_fa_backend_rejects_cc12(monkeypatch): from freetoken.attention.fa import FlashAttentionBackend diff --git a/tests/engine/test_cache_budget.py b/tests/engine/test_cache_budget.py index a77409f013..e997d2ba90 100644 --- a/tests/engine/test_cache_budget.py +++ b/tests/engine/test_cache_budget.py @@ -80,6 +80,43 @@ def test_budget_too_small_for_min_moe_plus_reserve_raises(): ) # min moe = 4 slots (400 B) + reserve (10 pages = 100 B) = 500 B > 300 B budget +def test_overlap_floor_that_does_not_fit_falls_back_to_num_experts(): + # #4: 2*num_experts slots (800 B) + the KV reserve (100 B) exceed the 700 B budget, but + # num_experts slots do fit. Plan without overlap instead of refusing to start. + size, pages, overlap = plan_cache_budget( + budget_bytes=700, per_expert_bytes=100, cache_per_page=10, + num_experts=4, total_experts=50, prefill_overlap=True, + kv_reserve_pages=10, max_slots=50, + ) + assert overlap is False + assert size == 6 # (700 - 100) // 100, above the num_experts floor + assert pages == 10 + assert size * 100 + pages * 10 <= 700 + + +def test_overlap_kept_when_its_floor_fits_the_budget(): + size, pages, overlap = plan_cache_budget( + budget_bytes=900, per_expert_bytes=100, cache_per_page=10, + num_experts=4, total_experts=50, prefill_overlap=True, + kv_reserve_pages=10, max_slots=50, + ) + assert (size, pages, overlap) == (8, 10, True) + + +def test_budget_too_small_message_names_what_fits_and_the_flags(): + # num_experts slots (400 B) + 10 reserved pages (100 B) > 450 B: 5 pages fit beside them. + with pytest.raises(AssertionError) as err: + plan_cache_budget( + budget_bytes=450, per_expert_bytes=100, cache_per_page=10, + num_experts=4, total_experts=50, prefill_overlap=True, + kv_reserve_pages=10, max_slots=50, page_size=64, + ) + msg = str(err.value) + assert "beside 4 slots at most 5 KV pages (320 tokens) fit" in msg + for flag in ("--kv-reserve-tokens", "--kv-cache-dtype", "--memory-ratio", "--moe-cache-size"): + assert flag in msg + + def test_prefill_overlap_false_is_honored(): # Even when the cache could fit 2*num_experts, an explicit False stays False. size, pages, overlap = plan_cache_budget( @@ -112,6 +149,25 @@ def test_resolve_auto_applies_ratio_once(): assert size == 8 and pages == 40 and overlap is True +def test_resolve_auto_reserves_usable_tokens_beyond_the_dummy_page(): + size, pages, overlap = resolve_moe_cache_auto( + baseline_free=940, + weights_bytes=0, + memory_ratio=1.0, + cache_per_page=10, + fixed_cache_size=0, + per_expert_bytes=100, + num_experts=2, + total_experts=50, + prefill_overlap=False, + kv_reserve_tokens=256, + page_size=64, + ) + assert overlap is False + assert (pages - 1) * 64 >= 256 + assert size == 8 and pages == 14 + + def test_resolve_auto_caps_slots_at_the_kernel_limit(): size, _, _ = resolve_moe_cache_auto( baseline_free=10_000_000, weights_bytes=0, memory_ratio=1.0, @@ -266,7 +322,7 @@ class tp_info: assert fixed == 0 -def test_engine_resolve_auto_moe_cache_size_maps_kwargs(): +def test_engine_resolve_auto_moe_cache_size_maps_kwargs(monkeypatch): import torch from freetoken.engine.engine import Engine @@ -293,6 +349,7 @@ class StubConfig: memory_ratio = 0.9 moe_prefill_overlap = True kv_reserve_tokens = 0 + num_page_override = 64 swa_full_tokens_ratio = 0.2 swa_num_pages_override = None model_config = StubModelConfig() @@ -314,29 +371,70 @@ class StubBanks: engine._weights_bytes = 1_000_000 engine._pool_cls = MHAKVCache # __init__ skipped -> install the generic pool family - size, pages, overlap = engine._resolve_auto_moe_cache_size(StubConfig(), StubBanks()) + captured = {} - # cross-check against the same pure functions, proving the kwarg mapping is faithful - from freetoken.engine.cache_budget import expert_bytes_per_slot, resolve_moe_cache_auto - from freetoken.kvcache.mha_pool import MHAKVCache + def fake_resolve(**kwargs): + captured.update(kwargs) + return 8, 64, True - cache_per_page, fixed, _, _ = MHAKVCache.kv_cost(StubConfig()) - expected = resolve_moe_cache_auto( - baseline_free=10_000_000, weights_bytes=1_000_000, memory_ratio=0.9, - cache_per_page=cache_per_page, fixed_cache_size=fixed, - per_expert_bytes=expert_bytes_per_slot(StubBanks.sources), - num_experts=4, total_experts=8, prefill_overlap=True, - kv_reserve_tokens=0, page_size=16, + monkeypatch.setattr( + "freetoken.engine.cache_budget.resolve_moe_cache_auto", fake_resolve ) - assert (size, pages, overlap) == expected + got = engine._resolve_auto_moe_cache_size(StubConfig(), StubBanks()) + + assert got == (8, 64, True) + assert captured["kv_reserve_tokens"] == 64 * 16 + assert captured["page_size"] == 16 + assert captured["num_experts"] == 4 + assert captured["total_experts"] == 8 + assert captured["per_expert_bytes"] == 512 + 256 class StubMethod: def slot_limit(self): return 5 - size, _, _ = engine._resolve_auto_moe_cache_size(StubConfig(), StubBanks(), StubMethod()) - assert size == 5 + # the kernel's slot limit rides the same mapping (fork #20) + engine._resolve_auto_moe_cache_size(StubConfig(), StubBanks(), StubMethod()) + assert captured["max_slots"] == 5 + + + +def test_rebuild_fit_check_reserves_the_gdn_prefill_workspace(): + """The runtime-rebuild fit-check prices the GDN prefill workspace like the startup sizing, + so a rebuild cannot grow the pools back into the space reserved for it.""" + from freetoken.engine.engine import Engine + from freetoken.kvcache.base import CacheRebuildRejected + from freetoken.kvcache.linear_state_pool import gdn_prefill_workspace_bytes, state_pool_bytes + from freetoken.models.config import LinearGatedDeltaGroupConfig + + group = LinearGatedDeltaGroupConfig( + name="linear", layer_ids=(0, 1), num_key_heads=2, num_value_heads=4, + key_head_dim=16, value_head_dim=16, conv_kernel_dim=4, output_gate="silu", + ) + config = SimpleNamespace( + model_config=SimpleNamespace(linear_attention_group=lambda: group), + dtype=torch.bfloat16, tp_info=SimpleNamespace(size=1), + max_extend_tokens=1024, max_seq_len=4096, + ) + captured = {} + def validate_rebuild(config, **kwargs): + captured.update(kwargs) + raise CacheRebuildRejected("stop before teardown") + + engine = SimpleNamespace( + config=config, moe_offload_cache=None, kv_offloader=None, + linear_state_pool=SimpleNamespace(num_slots=5), + kv_cache=SimpleNamespace(validate_rebuild=validate_rebuild), + _baseline_free=1 << 30, _weights_bytes=0, num_pages=64, + ) + engine._target_moe_and_expert_bytes = lambda size, sizes=None: (0, 0) + with pytest.raises(CacheRebuildRejected): + Engine.rebuild_runtime_cache(engine, num_pages=32) + + workspace = gdn_prefill_workspace_bytes(config) + assert workspace > 0 + assert captured["extra_fixed_bytes"] == state_pool_bytes(config, 5) + workspace # --------------------------------------------------------------------------- # offload-cache sizing guard + auto-resolution (_require_offload_cache_size / _adjust_config), @@ -523,3 +621,71 @@ def test_uncapped_platform_stays_uncapped(monkeypatch): if hasattr(os, "uname") and "microsoft" in os.uname().release.lower(): pytest.skip("WSL caps pinning") assert _pin_budget_bytes(reserved=2**30) is None + + +def test_moe_slot_pool_split(tmp_path): + import json + + from freetoken.moe.pool_policy import group_layers, load_profile, make_pools, split_slots + + pools = make_pools(group_layers("per-layer", [0] * 4, [1] * 4), num_experts=16, min_num_slots=8, slot_limit=None) + assert split_slots(41, pools) == [11, 10, 10, 10] + with pytest.raises(ValueError, match="need"): + split_slots(31, pools) + with pytest.raises(ValueError, match="hold"): + split_slots(65, pools) + # one pool takes the whole cache, as before pools existed + assert split_slots(100, make_pools(group_layers("unified", [0] * 4, [1] * 4), 16, 8, None)) == [100] + # unified pools one bank layout each, widest slots first + assert group_layers("unified", ["a", "b", "a", "b"], [1, 3, 1, 3]) == [[1, 3], [0, 2]] + + # layer 0 uses two experts, layer 1 has a hot half and a cold half, layers 2 and 3 use all 16 evenly + counts = [[5] * 2 + [0] * 14, [3] * 8 + [1] * 8, [1] * 16, [1] * 16] + path = tmp_path / "profile.json" + path.write_text(json.dumps({"model": "toy", "expert_counts": counts})) + assert split_slots(48, pools, load_profile(path, 4, 16)) == [8, 8, 16, 16] + with pytest.raises(ValueError, match="needs"): + load_profile(path, 4, 32) + + + +def test_moe_pool_flags_and_decode_batch_clamp(): + from unittest.mock import patch + + from freetoken.engine.engine import _clamp_decode_batch_to_pools + from freetoken.server.args import parse_args + + hf = SimpleNamespace(to_dict=lambda: {"architectures": ["Qwen3MoeForCausalLM"], "torch_dtype": "bfloat16"}) + with patch("freetoken.utils.cached_load_hf_config", lambda _path: hf): + args, _ = parse_args([ + "--model", "/models/anon", "--moe-cache-mem", "1.5GiB", + "--moe-cache-scope", "per-layer", "--moe-cache-split", "profile:/p.json", + ]) + assert (args.moe_cache_mem, args.moe_cache_scope, args.moe_cache_split) == (3 << 29, "per-layer", "profile:/p.json") + for bad in (["--moe-cache-mem", "12"], ["--moe-cache-mem", "1GiB", "--moe-cache-size", "512"]): + with pytest.raises(SystemExit): + parse_args(["--model", "/models/anon", *bad]) + + # GPU decode needs batch * top_k distinct experts to fit the smallest pool below num_experts + config = SimpleNamespace( + model_config=SimpleNamespace(num_experts=16, num_experts_per_tok=4), + max_running_req=4, cuda_graph_max_bs=None, cuda_graph_bs=[1, 2, 4], + ) + _clamp_decode_batch_to_pools(config, [12, 16, 9]) + assert (config.max_running_req, config.cuda_graph_max_bs, config.cuda_graph_bs) == (2, 2, [1, 2]) + + +def test_host_embedding_placement_reports_the_bytes_it_pinned(monkeypatch): + # the engine charges these bytes to the pin budget before it plans the expert banks + from freetoken.distributed import set_tp_info, try_get_tp_info + from freetoken.engine.engine import _place_embeddings_on_host + from freetoken.layers.embedding import ParallelLMHead, VocabParallelEmbedding + + if try_get_tp_info() is None: + set_tp_info(rank=0, size=1) + monkeypatch.setattr(VocabParallelEmbedding, "place_on_host", lambda self: None) + embed, tied = VocabParallelEmbedding(100, 64), VocabParallelEmbedding(50, 64) + embed.weight = torch.empty(100, 64, dtype=torch.bfloat16) + head = ParallelLMHead(50, 64, tie_word_embeddings=True, tied_embedding=tied) + model = SimpleNamespace(embed_tokens=embed, per_layer_embed=tied, lm_head=head) + assert _place_embeddings_on_host(model) == 100 * 64 * 2 diff --git a/tests/engine/test_distributed_store.py b/tests/engine/test_distributed_store.py new file mode 100644 index 0000000000..4fc2473762 --- /dev/null +++ b/tests/engine/test_distributed_store.py @@ -0,0 +1,113 @@ +"""Regression test for #301: the TP rendezvous store must not fall back to a wildcard +bind. torch's tcp:// rendezvous handler ignores the host it's given when binding the +master's listening socket and always listens on every interface, so +init_method="tcp://127.0.0.1:PORT" looks loopback-only but isn't -- see +Engine._make_distributed_store. This drives the real store-construction path across two +spawned local ranks (CPU/gloo) and checks the actual bind address, not just the URL. +""" + +from __future__ import annotations + +import gc +import multiprocessing as mp +import os +from datetime import timedelta +from types import SimpleNamespace + +import torch +import torch.distributed as dist + +from freetoken.distributed import DistributedInfo +from freetoken.engine.engine import Engine + +WORLD_SIZE = 2 +TIMEOUT = timedelta(seconds=15) + + +def _config(rank: int, port: int) -> SimpleNamespace: + return SimpleNamespace( + tp_info=DistributedInfo(rank, WORLD_SIZE), + distributed_timeout=TIMEOUT.total_seconds(), + distributed_port=port, + ) + + +def _run_rank(rank: int, port: int, result_q: mp.Queue) -> None: + engine = Engine.__new__(Engine) # only _make_distributed_store is under test + store = engine._make_distributed_store(_config(rank, port)) + dist.init_process_group( + backend="gloo", rank=rank, world_size=WORLD_SIZE, timeout=TIMEOUT, store=store + ) + total = torch.tensor([float(rank)]) + dist.all_reduce(total) + bind_host = engine._distributed_listen_addr[0] if rank == 0 else None + result_q.put((rank, total.item(), bind_host)) + dist.destroy_process_group() + + +def test_two_ranks_rendezvous_with_loopback_only_store(): + ctx = mp.get_context("spawn") + port = 29511 + result_q = ctx.Queue() + procs = [ctx.Process(target=_run_rank, args=(rank, port, result_q)) for rank in range(WORLD_SIZE)] + for p in procs: + p.start() + + results = {} + for _ in procs: + rank, total, bind_host = result_q.get(timeout=30) + results[rank] = (total, bind_host) + for p in procs: + p.join(timeout=30) + assert p.exitcode == 0 + + # both ranks completed the same all_reduce over the store -- rendezvous worked + assert results[0][0] == results[1][0] == 1.0 # sum(0, 1) + # the master's listening socket was pre-bound to loopback, not the wildcard + # address the C10d TCPStore server binds to on its own + assert results[0][1] == "127.0.0.1" + + +def _teardown_rank(port: int, result_q: mp.Queue) -> None: + listen_fds = [] + real_store = dist.TCPStore + + def recording_store(*args, **kwargs): + listen_fds.append(kwargs.get("master_listen_fd")) + return real_store(*args, **kwargs) + + dist.TCPStore = recording_store + engine = Engine.__new__(Engine) + config = SimpleNamespace( + tp_info=DistributedInfo(0, 1), distributed_timeout=TIMEOUT.total_seconds(), + distributed_port=port, + ) + store = engine._make_distributed_store(config) + dist.init_process_group(backend="gloo", rank=0, world_size=1, timeout=TIMEOUT, store=store) + dist.destroy_process_group() + del store + gc.collect() + # the listening fd number is reused by an unrelated file, as a later open() would + fd = listen_fds[0] + r, w = os.pipe() + os.dup2(r, fd) + del engine + gc.collect() + try: + os.fstat(fd) + result_q.put("open") + except OSError as exc: + result_q.put(f"closed: {exc}") + + +def test_the_engine_does_not_close_the_listen_fd_the_store_owns(): + # TCPStore takes ownership of master_listen_fd and closes it itself; an engine that also + # kept the Python socket closed whatever file reused that number when it was freed + ctx = mp.get_context("spawn") + result_q = ctx.Queue() + proc = ctx.Process(target=_teardown_rank, args=(29513, result_q)) + proc.start() + outcome = result_q.get(timeout=30) + proc.join(timeout=30) + assert proc.exitcode == 0 + assert outcome == "open" diff --git a/tests/engine/test_embed_weights.py b/tests/engine/test_embed_weights.py new file mode 100644 index 0000000000..d3aa102f94 --- /dev/null +++ b/tests/engine/test_embed_weights.py @@ -0,0 +1,39 @@ +"""--embed-weights host: a setup the host gather cannot serve is refused while the config is still a dataclass, not after the weight load.""" + +from types import SimpleNamespace + +import pytest +import torch + +from freetoken.distributed import DistributedInfo +from freetoken.engine.config import EngineConfig + + +def _config(monkeypatch, *, tie: bool, tp: int = 1, embed_weights: str = "host"): + import freetoken.engine.config as engine_config + + hf = SimpleNamespace( + architectures=["Qwen3ForCausalLM"], model_type="qwen3", num_hidden_layers=2, num_attention_heads=4, + num_key_value_heads=2, hidden_size=64, vocab_size=128, intermediate_size=128, rms_norm_eps=1e-6, + max_position_embeddings=1024, hidden_act="silu", tie_word_embeddings=tie, + ) + monkeypatch.setattr(engine_config, "cached_load_hf_config", lambda path, *_: hf) + monkeypatch.setattr(engine_config, "checkpoint_quant_config", lambda *args: None) + return EngineConfig( + model_path="/fake", tp_info=DistributedInfo(rank=0, size=tp), dtype=torch.bfloat16, embed_weights=embed_weights + ) + + +def test_a_tied_embedding_is_refused(monkeypatch): + with pytest.raises(ValueError, match="needs an untied embedding"): + _config(monkeypatch, tie=True).model_config + + +def test_tensor_parallel_is_refused(monkeypatch): + with pytest.raises(ValueError, match="single GPU only"): + _config(monkeypatch, tie=False, tp=2).model_config + + +def test_an_untied_single_gpu_model_and_the_gpu_default_pass(monkeypatch): + assert not _config(monkeypatch, tie=False).model_config.tie_word_embeddings + assert _config(monkeypatch, tie=True, tp=2, embed_weights="gpu").model_config.tie_word_embeddings diff --git a/tests/engine/test_kv_quant_config.py b/tests/engine/test_kv_quant_config.py new file mode 100644 index 0000000000..1f460166d1 --- /dev/null +++ b/tests/engine/test_kv_quant_config.py @@ -0,0 +1,270 @@ +"""Config-time gates for ``--kv-cache-dtype`` (EngineConfig.kv_quant). + +fp8 KV is only half a feature: the pool has to store it AND the attention backend has +to read the scales. Everything here must fail while the config is still a dataclass -- +after weights are resident, a wrong combination has already cost a load and (worse) +fi/fa/trtllm would happily attend over raw e4m3 codes and produce plausible garbage. +""" + +from types import SimpleNamespace + +import pytest +import torch + +from freetoken.attention import AttnType +from freetoken.models.config import KVCacheGroupSpec + + +def _spec(name, attn_type, *, mla=False, sliding_window=None, index_head_dim=0, index_ratio=1): + return KVCacheGroupSpec( + name=name, + layer_ids=(0, 1), + num_kv_heads=1, + head_dim=64, + sliding_window=sliding_window, + mla=mla, + index_head_dim=index_head_dim, + num_index_layers=2 if index_head_dim else 0, + index_ratio=index_ratio, + attn_type=attn_type, + ) + + +def _model_config(kind): + mc = SimpleNamespace( + model_type=kind, + single_stream_only=False, + is_moe=False, + expert_quant="none", + has_swa_attention=False, + has_linear_attention=False, + num_layers=4, + rotary_config=SimpleNamespace(max_position=1024, table_positions=1024), + ) + specs = { + "full": (_spec("full", AttnType.FULL),), + "swa": ( + _spec("full", AttnType.FULL), + _spec("swa", AttnType.SWA, sliding_window=128), + ), + "mla": (_spec("full", AttnType.MLA, mla=True),), + "dsa": (_spec("full", AttnType.DSA, mla=True, index_head_dim=128),), + "kpool": (_spec("full", AttnType.DSA, mla=True, index_head_dim=128, index_ratio=4),), + "dsv4": (_spec("dsv4", AttnType.DSV4, sliding_window=128),), + "bsa": (_spec("full", AttnType.BSA, index_head_dim=128),), + "qsa": (_spec("full", AttnType.QSA, index_head_dim=128, index_ratio=4),), + }[kind] + if kind == "swa": + mc.has_swa_attention = True + if kind == "dsv4": + mc.dsv4_args = SimpleNamespace(window_size=128) + if kind in ("qsa", "kpool"): + mc.has_linear_attention = True + mc.kv_cache_group_specs = lambda: specs + return mc + + +def _config(kind, **overrides): + from freetoken.distributed import DistributedInfo + from freetoken.engine.config import EngineConfig + + config = EngineConfig( + model_path="/tmp/freetoken-test-model", + tp_info=DistributedInfo(rank=0, size=1), + dtype=torch.bfloat16, + **overrides, + ) + object.__setattr__(config, "model_config", _model_config(kind)) + return config + + +def _patch_fast_machine(monkeypatch): + """A machine where every fast external backend is available, so an fp8 result can + only come from the gate and not from a missing package.""" + from freetoken.engine import engine + + monkeypatch.setattr(engine, "is_sm100_family", lambda: False) + monkeypatch.setattr(engine, "is_sm90_family", lambda: True) + monkeypatch.setattr(engine, "_flashinfer_available", lambda: True) + monkeypatch.setattr(engine, "_sgl_flash_attn_available", lambda: True) + + +def test_kv_quant_spellings(): + from freetoken.engine.engine import _resolve_kv_quant + + assert _resolve_kv_quant("auto") == "none" + assert _resolve_kv_quant("bf16") == "none" + assert _resolve_kv_quant("FP8") == "fp8" + assert _resolve_kv_quant("NVFP4") == "nvfp4" + assert _resolve_kv_quant(None) == "none" + with pytest.raises(ValueError, match="kv-cache-dtype"): + _resolve_kv_quant("q8") + + +def test_only_the_backends_that_read_scales_declare_fp8_support(): + from freetoken.attention import SUPPORTED_ATTENTION_BACKENDS, attention_backend_info + + fp8 = { + name + for name in SUPPORTED_ATTENTION_BACKENDS.supported_names() + if attention_backend_info(name).supports_fp8_kv + } + # triton serves plain paged / hybrid-SWA pools, qsa_sparse dequantizes selected + # rows, and dsa dequantizes MLA latent rows. External backends and sparse + # families without a scale path must stay out of this set. + assert fp8 == {"triton", "qsa_sparse", "dsa"} + + nvfp4 = { + name + for name in SUPPORTED_ATTENTION_BACKENDS.supported_names() + if attention_backend_info(name).supports_nvfp4_kv + } + assert nvfp4 == {"triton", "qsa_sparse", "dsa"} + + +def test_auto_avoids_the_fast_backends_for_fp8(monkeypatch): + from freetoken.engine.engine import _adjust_config + + _patch_fast_machine(monkeypatch) + # Same machine where a 16-bit cache auto-selects the sm90 "fa,fi" tree... + plain = _config("full", attention_backend="auto") + _adjust_config(plain) + assert plain.kv_quant == "none" + assert plain.attention_backend == "fa,fi" + + quantized = _config("full", attention_backend="auto", kv_quant="fp8") + _adjust_config(quantized) + assert quantized.attention_backend == "triton" + + +@pytest.mark.parametrize("backend", ["fi", "fa", "trtllm", "fi,triton", "triton,fi"]) +def test_explicit_unsupported_backend_is_rejected(monkeypatch, backend): + from freetoken.engine.engine import _adjust_config + + _patch_fast_machine(monkeypatch) + monkeypatch.setattr( + "freetoken.engine.engine.is_sm100_family", lambda: True + ) # let trtllm clear its own arch gate first + config = _config("full", attention_backend=backend, kv_quant="fp8", page_size=1) + with pytest.raises(ValueError, match="kv-cache-dtype fp8"): + _adjust_config(config) + + +def test_explicit_triton_is_accepted(monkeypatch): + from freetoken.engine.engine import _adjust_config + + _patch_fast_machine(monkeypatch) + config = _config("full", attention_backend="triton", kv_quant="fp8") + _adjust_config(config) + assert config.kv_quant == "fp8" + + +def test_nvfp4_auto_selects_triton(monkeypatch): + from freetoken.engine.engine import _adjust_config + + _patch_fast_machine(monkeypatch) + config = _config("full", attention_backend="auto", kv_quant="nvfp4") + _adjust_config(config) + assert config.attention_backend == "triton" + + +@pytest.mark.parametrize("kind", ["dsv4", "bsa"]) +def test_nvfp4_rejects_unsupported_pools_before_allocation(monkeypatch, kind): + from freetoken.engine.engine import _adjust_config + from freetoken.kvcache import create_kvcache_pool + + _patch_fast_machine(monkeypatch) + config = _config(kind, attention_backend="auto", kv_quant="nvfp4") + with pytest.raises(ValueError, match="nvfp4"): + _adjust_config(config) + with pytest.raises(ValueError, match="nvfp4"): + create_kvcache_pool(config.model_config, num_pages=4, page_size=1, + dtype=torch.bfloat16, device=torch.device("cpu"), kv_quant="nvfp4") + + +@pytest.mark.parametrize("backend", ["fi", "fa", "trtllm", "fi,triton", "triton,fi"]) +def test_nvfp4_rejects_backends_without_its_layout(monkeypatch, backend): + from freetoken.engine.engine import _adjust_config + + _patch_fast_machine(monkeypatch) + monkeypatch.setattr("freetoken.engine.engine.is_sm100_family", lambda: True) + config = _config("full", attention_backend=backend, kv_quant="nvfp4") + with pytest.raises(ValueError, match="nvfp4"): + _adjust_config(config) + + +def test_nvfp4_rejects_partial_blocks(monkeypatch): + from dataclasses import replace + from freetoken.engine.engine import _adjust_config + + _patch_fast_machine(monkeypatch) + config = _config("full", attention_backend="auto", kv_quant="nvfp4") + spec = replace(config.model_config.kv_cache_group_specs()[0], head_dim=72) + config.model_config.kv_cache_group_specs = lambda: (spec,) + with pytest.raises(ValueError, match="divisible by 16"): + _adjust_config(config) + + +def test_nvfp4_accepts_hybrid_swa(monkeypatch): + from freetoken.engine.engine import _adjust_config + + _patch_fast_machine(monkeypatch) + config = _config("swa", attention_backend="auto", kv_quant="nvfp4") + _adjust_config(config) + assert config.kv_quant == "nvfp4" + assert config.attention_backend == "triton" + + +def test_nvfp4_accepts_qsa(monkeypatch): + from freetoken.engine.engine import _adjust_config + + _patch_fast_machine(monkeypatch) + config = _config("qsa", attention_backend="auto", kv_quant="nvfp4") + _adjust_config(config) + assert config.kv_quant == "nvfp4" + assert config.attention_backend == "qsa_sparse" + + +@pytest.mark.parametrize("kind", ["mla", "dsa", "kpool"]) +@pytest.mark.parametrize("backend", ["auto", "dsa"]) +@pytest.mark.parametrize("kv_quant", ["fp8", "nvfp4"]) +def test_mla_and_dsa_select_the_scale_reading_backend(monkeypatch, kind, backend, kv_quant): + from freetoken.engine.engine import _adjust_config + + _patch_fast_machine(monkeypatch) + config = _config(kind, attention_backend=backend, kv_quant=kv_quant) + _adjust_config(config) + assert config.attention_backend == "dsa" + assert config.page_size == (64 if kind == "kpool" else 1) + + +@pytest.mark.parametrize("kind", ["dsv4", "bsa"]) +def test_pool_families_without_a_scale_read_path_are_rejected(monkeypatch, kind): + from freetoken.engine.engine import _adjust_config + + _patch_fast_machine(monkeypatch) # the rejection must not depend on this box's wheels + config = _config(kind, attention_backend="auto", kv_quant="fp8") + with pytest.raises(ValueError, match="kv-cache-dtype fp8"): + _adjust_config(config) + + +def test_qsa_keeps_fp8_available(monkeypatch): + """The QSA pool is the block-sparse family whose K/V rows do reach a kernel that can + dequantize them; the compressed index keys selection scores against are a separate, + always-16-bit tier, so the gate has nothing left to refuse here.""" + from freetoken.engine.engine import _adjust_config + + _patch_fast_machine(monkeypatch) + config = _config("qsa", attention_backend="auto", kv_quant="fp8") + _adjust_config(config) + assert config.kv_quant == "fp8" + assert "qsa_sparse" in config.attention_backend + + +def test_swa_pool_keeps_fp8_available(monkeypatch): + from freetoken.engine.engine import _adjust_config + + _patch_fast_machine(monkeypatch) + config = _config("swa", attention_backend="auto", kv_quant="fp8") + _adjust_config(config) + assert config.kv_quant == "fp8" and config.attention_backend == "triton" diff --git a/tests/engine/test_logits_processors.py b/tests/engine/test_logits_processors.py new file mode 100644 index 0000000000..1829f35f32 --- /dev/null +++ b/tests/engine/test_logits_processors.py @@ -0,0 +1,234 @@ +"""The logits processors behind min_p, the penalties, logit_bias and min_tokens: pure +torch on the CPU, so the math is checked without a GPU or the sampling kernels.""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch +from freetoken.core import SamplingParams +from freetoken.engine.sample import ( + LogitsPlan, + RowPenalty, + Sampler, + apply_logits_processors, + apply_penalties, +) + +V = 8 +CPU = torch.device("cpu") + + +def _plan(rows, **fields) -> LogitsPlan: + plan = LogitsPlan( + rows=torch.tensor(rows, dtype=torch.int64), + temps=torch.ones((len(rows), 1), dtype=torch.float32), + ) + for k, v in fields.items(): + setattr(plan, k, v) + return plan + + +def _ids(rows: list[list[int]]) -> torch.Tensor: + width = max(len(r) for r in rows) + out = torch.full((len(rows), width), V, dtype=torch.int64) + for i, r in enumerate(rows): + out[i, : len(r)] = torch.tensor(r) + return out + + +def test_repetition_penalty_covers_prompt_and_output_with_hf_semantics(): + logits = torch.tensor([[2.0, -2.0, 2.0, -2.0, 1.0, 1.0, 1.0, 1.0]]) + counts = torch.zeros(V, dtype=torch.int32) + counts[[0, 1]] = 1 # generated: 0 (positive), 1 (negative) + prompt = torch.zeros(V, dtype=torch.bool) + prompt[[2, 3]] = True # prompt: 2 (positive), 3 (negative) + out = apply_penalties(logits, [RowPenalty(0, counts, 0.0, 0.0, 2.0, prompt)]) + assert out[0, 0].item() == 1.0 and out[0, 2].item() == 1.0 # divided + assert out[0, 1].item() == -4.0 and out[0, 3].item() == -4.0 # multiplied + assert out[0, 4].item() == 1.0 # unseen + assert logits[0, 0].item() == 2.0 # the caller's logits stay untouched + + +def test_logit_bias_adds_per_row(): + logits = torch.zeros((2, V)) + plan = _plan( + [0, 1], + bias_rows=torch.tensor([0, 1, 1]), + bias_ids=torch.tensor([2, 2, 6]), + bias_vals=torch.tensor([5.0, -100.0, 1.5]), + ) + out = apply_logits_processors(logits, plan, V) + assert ( + out[0, 2].item() == 5.0 + and out[1, 2].item() == -100.0 + and out[1, 6].item() == 1.5 + ) + + +def test_min_tokens_masks_the_stop_ids_of_the_row_only(): + logits = torch.zeros((2, V)) + plan = _plan([0, 1], min_rows=torch.tensor([1]), min_ids=_ids([[0, 7]])) + out = apply_logits_processors(logits, plan, V) + assert out[1, 0].item() == float("-inf") and out[1, 7].item() == float("-inf") + assert out[1, 3].item() == 0.0 + assert torch.isfinite(out[0]).all() + + +@pytest.mark.parametrize("overlap", [True, False]) +def test_min_tokens_releases_the_stop_ids_at_output_index_min_tokens(overlap): + """Overlap scheduling prepares batch N before batch N-1's token reaches req.input_ids on + the host; the count must not lag by that token and hold EOS one step too long.""" + from freetoken.core import Req + + sp = SamplingParams(min_tokens=3, max_tokens=10) + sp.min_tokens_stop_ids = [0] + req = Req(input_ids=torch.tensor([5, 6], dtype=torch.int32), table_idx=0, cached_len=0, + output_len=10, uid=1, sampling_params=sp, cache_handle=None) + sampler = Sampler(CPU, V) + masked, pending = [], None + for _ in range(6): + plan = sampler.prepare(SimpleNamespace(reqs=[req])).plan + masked.append(plan is not None and plan.min_rows is not None) + req.complete_one() # the forward launch advances the device length + if overlap: + if pending is not None: # batch N-1 drains after batch N launched + req.append_host(pending) + pending = torch.tensor([7], dtype=torch.int32) + else: + req.append_host(torch.tensor([7], dtype=torch.int32)) + assert [i for i, m in enumerate(masked) if m] == [0, 1, 2] + + +def test_min_p_drops_tokens_below_the_fraction_of_the_top_probability(): + logits = torch.log(torch.tensor([[0.5, 0.3, 0.1, 0.05, 0.03, 0.01, 0.005, 0.005]])) + plan = _plan( + [0], min_p=torch.tensor([[0.15]]) + ) # threshold 0.075, between 0.1 and 0.05 + out = apply_logits_processors(logits, plan, V) + kept = torch.isfinite(out[0]) + assert kept.tolist() == [True, True, True, False, False, False, False, False] + # a row with min_p 0 is left alone + plan0 = _plan([0], min_p=torch.tensor([[0.0]])) + assert torch.isfinite(apply_logits_processors(logits, plan0, V)[0]).all() + + +def _req( + prompt: list[int], generated: list[int], max_tokens: int, **sp +) -> SimpleNamespace: + ids = torch.tensor(prompt + generated, dtype=torch.int32) + return SimpleNamespace( + input_ids=ids, + device_len=len(ids), + output_len=max_tokens, + max_device_len=len(prompt) + max_tokens, + sampling_params=SamplingParams(**sp), + can_decode=True, + output_token_counts=None, + prompt_token_mask=None, + ) + + +def test_prepare_builds_no_plan_without_processors_and_a_plan_with_them(): + sampler = Sampler(CPU, V) + plain = SimpleNamespace(reqs=[_req([1, 2], [3], 4)]) + assert sampler.prepare(plain).plan is None + + # all greedy: the argmax path runs on the CPU, the sampling kernels need a GPU + batch = SimpleNamespace( + reqs=[ + _req([1, 2], [3], 4), + _req( + [1, 2], + [3, 3], + 4, + presence_penalty=1.0, + repetition_penalty=1.5, + logit_bias=[[5, 2.0]], + ), + _req([4], [], 4, min_tokens=2, min_tokens_stop_ids=[0, 7]), + _req( + [4], [6, 6], 4, min_tokens=2, min_tokens_stop_ids=[0] + ), # already past min_tokens + ] + ) + args = sampler.prepare(batch) + plan = args.plan + assert plan is not None + assert plan.rows.tolist() == [1, 2, 3] + assert [pen.row for pen in args.penalties] == [1] + assert batch.reqs[1].prompt_token_mask.nonzero().flatten().tolist() == [1, 2] + assert plan.bias_rows.tolist() == [0] and plan.bias_ids.tolist() == [5] + assert plan.min_rows.tolist() == [1] # local row of the third request only + assert plan.min_ids[0].tolist() == [0, 7] + + logits = torch.zeros((4, V)) + out = apply_logits_processors(logits, plan, V) + assert out[1, 3].item() == 0.0 # the penalties are not plan work + assert out[1, 5].item() == 2.0 + assert out[2, 0].item() == float("-inf") and out[2, 7].item() == float("-inf") + assert torch.isfinite(out[3]).all() + # greedy path: argmax over the processed logits + assert sampler.sample(logits, args)[1].item() == 5 + + +def test_penalties_ride_the_device_state_not_the_plan(): + # the plan must not apply them a second time on top of the sampler's own counts + sampler = Sampler(CPU, V) + only = SimpleNamespace(reqs=[_req([1], [], 4, presence_penalty=1.0, frequency_penalty=0.5)]) + args = sampler.prepare(only) + assert args.plan is None and [row for row, *_ in args.penalties] == [0] + + req = _req([1], [], 4, presence_penalty=1.0, logit_bias=[[5, 0.25]]) + args = sampler.prepare(SimpleNamespace(reqs=[req])) + assert args.plan is not None and [row for row, *_ in args.penalties] == [0] + req.output_token_counts[3] = 1 + logits = torch.zeros((1, V)) + logits[0, 3] = 1.5 + # once: 3 -> 0.5 beats 5 -> 0.25; twice would drop 3 to -0.5 and pick 5 + assert sampler.sample(logits, args).item() == 3 + assert req.output_token_counts[3].item() == 2 + + +def test_repetition_penalty_ignores_multimodal_placeholder_ids_in_the_prompt(): + # image tokens carry pseudo-ids >= MM_PAD_SHIFT_VALUE; scattering them into the + # [rows, V+1] prompt mask indexed out of bounds (a device-side assert on CUDA) + sampler = Sampler(CPU, V) + batch = SimpleNamespace(reqs=[_req([1, 1_000_000, 2], [3], 4, repetition_penalty=2.0)]) + args = sampler.prepare(batch) + out = apply_penalties(torch.ones(1, V), args.penalties) + assert out[0, 1] == 0.5 and out[0, 2] == 0.5 and out[0, 4] == 1.0 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +def test_min_p_with_temperature_draws_only_the_kept_tokens(): + # min_p leaves -inf over nearly the whole row; the sampling softmax must not turn that + # into NaN (which the Pascal fallback then samples as a uniform row) + V = 32768 + sampler = Sampler(torch.device("cuda"), V) + req = _req([1, 2], [], 4, temperature=0.8, min_p=0.1) + args = sampler.prepare(SimpleNamespace(reqs=[req])) + logits = torch.full((1, V), -10.0, device="cuda") + kept = [100, 20000, 31000] + logits[0, kept] = torch.tensor([5.0, 4.5, 4.0], device="cuda") + for _ in range(32): + assert sampler.sample(logits, args).item() in kept + + +def test_repetition_penalty_reads_the_prompt_once_and_counts_output_on_the_device(): + # the prompt (up to the whole context) is uploaded once per request, not every step, + # and the generated side comes from the counts sample() keeps, not the host history + sampler = Sampler(CPU, V) + req = _req([1, 2], [], 4, repetition_penalty=2.0) + batch = SimpleNamespace(reqs=[req]) + args = sampler.prepare(batch) + mask = req.prompt_token_mask + logits = torch.tensor([[0.0, 1.0, 1.0, 4.0, 0.5, 0.0, 0.0, 0.0]]) + assert sampler.sample(logits, args).item() == 3 + assert req.output_token_counts.tolist() == [0, 0, 0, 1, 0, 0, 0, 0] + req.input_ids = torch.tensor([9, 9, 3], dtype=torch.int32) # a stale host view is not read + args = sampler.prepare(batch) + assert req.prompt_token_mask is mask + out = apply_penalties(logits, args.penalties) + assert out[0].tolist() == [0.0, 0.5, 0.5, 2.0, 0.5, 0.0, 0.0, 0.0] diff --git a/tests/engine/test_rebuild_runtime_cache_ordering.py b/tests/engine/test_rebuild_runtime_cache_ordering.py new file mode 100644 index 0000000000..ab01f8a18f --- /dev/null +++ b/tests/engine/test_rebuild_runtime_cache_ordering.py @@ -0,0 +1,209 @@ +"""Engine.rebuild_runtime_cache resize ordering (#643): shrinking resizes must run +before growing ones. Each pool frees only its own old tensors before reallocating, so a +grow that runs while a still-large sibling pool is resident can OOM even though the +TARGET geometry passed validate_rebuild. Drives the real Engine method unbound on a +SimpleNamespace engine shell (the idiom of test_kv_cache_rebuild.py's page-table test), +with fake pools mirroring each real pool's free-own-then-alloc sequence under a shared +accounting ledger, and the real BaseKVCachePool.validate_rebuild as the fit check.""" + +from __future__ import annotations + +import types + +import pytest +import torch + +from freetoken.engine.engine import Engine +from freetoken.kvcache.base import BaseKVCachePool + + +class _Ledger: + """Shared GPU-memory stand-in: cap in units (1 unit = 1 MoE slot / KV page).""" + + def __init__(self, cap: int): + self.cap = cap + self.used = 0 + self.peak = 0 + self.events: list[tuple] = [] + + def alloc(self, tag: str, n: int) -> None: + if self.used + n > self.cap: + self.events.append(("OOM", tag, n, self.used)) + raise torch.OutOfMemoryError( + f"alloc {tag}={n} would peak at {self.used + n} > cap {self.cap}" + ) + self.used += n + self.peak = max(self.peak, self.used) + self.events.append(("alloc", tag, n, self.used)) + + def free(self, tag: str, n: int) -> None: + self.used -= n + self.events.append(("free", tag, n, self.used)) + + +class _FakeMoeCache: + """Mirrors OffloadMoeCache.rebuild: validate -> free own tensors -> cache_size + mutated (offload_cache.py:491, before the alloc loop) -> reallocate.""" + + def __init__(self, ledger: _Ledger, size: int): + self.ledger = ledger + self.num_experts = 8 + self.max_slots = None + self.layout = None + self.quant_format = "none" + # one uint8 per expert row -> expert_bytes_per_slot == 1 unit + self.bank_sources = {"w13": [torch.zeros((self.num_experts, 1), dtype=torch.uint8)]} + self.cache_size = size + self.decode_target = "gpu" + self._held = size + ledger.alloc("moe", size) + + def validate_rebuild(self, cache_size: int, pool_sizes: list[int] | None = None) -> None: + if cache_size < self.num_experts: + raise ValueError(f"cache_size {cache_size} < num_experts {self.num_experts}") + + def rebuild(self, cache_size: int, pool_sizes: list[int] | None = None) -> None: + self.validate_rebuild(cache_size) + self.ledger.free("moe", self._held) + self._held = 0 + self.cache_size = cache_size + self.ledger.alloc("moe", cache_size) + self._held = cache_size + + +class _FakeKvPool(BaseKVCachePool): + """Mirrors MHAKVCache.rebuild (free the one buffer, then realloc); reuses the REAL + BaseKVCachePool.validate_rebuild with 1 unit per page and no fixed term.""" + + needs_rebind_on_rebuild = False + + def __init__(self, ledger: _Ledger, num_pages: int): + self.ledger = ledger + self.pages = num_pages + ledger.alloc("kv", num_pages) + + @classmethod + def kv_cost(cls, config): + return 1, 0, 16, 0 + + def rebuild_from_config(self, config, num_pages, *, num_swa_pages=None): + self.ledger.free("kv", self.pages) + self.ledger.alloc("kv", num_pages) + self.pages = num_pages + + def attach_page_table(self, page_table): + pass + + def unit_bytes(self): + return 1, 0 + + def k_cache(self, index): ... + def v_cache(self, index): ... + def store_kv(self, k, v, out_loc, layer_id): ... + + @property + def device(self): + return torch.device("cpu") + + @property + def dtype(self): + return torch.float16 + + @property + def num_layers(self): + return 1 + + +class _FakeGraphRunner: + def __init__(self, **kwargs): + self.graph_bs_list = [1, 2, 4] + self.destroy_cuda_graphs = lambda: None + + +def _build_engine(monkeypatch, ledger: _Ledger, *, moe: int, kv: int, budget: int): + monkeypatch.setattr(torch.cuda, "synchronize", lambda *a, **k: None) + monkeypatch.setattr("freetoken.engine.engine.GraphRunner", _FakeGraphRunner) + config = types.SimpleNamespace( + page_size=16, max_seq_len=4096, max_running_req=4, memory_ratio=1.0, + cuda_graph_max_bs=4, swa_num_pages_override=None, + model_config=types.SimpleNamespace( + vocab_size=128, model_is_mrope=False, linear_attention_group=lambda: None + ), + ) + page_table = torch.zeros((config.max_running_req + 1, 256), dtype=torch.int32) + eng = types.SimpleNamespace( + config=config, moe_offload_cache=_FakeMoeCache(ledger, moe), kv_cache=_FakeKvPool(ledger, kv), + linear_state_pool=None, kv_offloader=None, num_pages=kv, device=torch.device("cpu"), stream=None, model=None, + _baseline_free=budget, _weights_bytes=0, page_table=page_table, max_seq_len=4096, + ctx=types.SimpleNamespace(page_table=page_table), + dummy_req=types.SimpleNamespace(table_idx=0), + attn_backend=types.SimpleNamespace(reset_capture=lambda: None), + graph_runner=_FakeGraphRunner(), + _sync_get_memory=lambda: (0, 0), + ) + # one MoE slot pool (#689): the rebuild splits the target into its pool sizes first + from freetoken.moe.pool_policy import make_pools + + eng._moe_pools = lambda batch: make_pools([[0]], 8, 1, None) + eng._moe_expert_counts = None + eng._moe_pool_groups, eng._moe_layer_slot_bytes = [[0]], [1] # one uint8 per expert row + eng._moe_pool_bytes = Engine._moe_pool_bytes.__get__(eng) + eng._target_moe_and_expert_bytes = Engine._target_moe_and_expert_bytes.__get__(eng) + eng._resize_kv_pool = Engine._resize_kv_pool.__get__(eng) + eng._refresh_seq_state = Engine._refresh_seq_state.__get__(eng) + return eng + + +def _alloc_index(ledger: _Ledger, tag: str, n: int) -> int: + return next( + i for i, ev in enumerate(ledger.events) if ev[0] == "alloc" and ev[1] == tag and ev[2] == n + ) + + +def test_rebuild_runtime_cache_applies_shrinks_before_grows(monkeypatch): + # moe 40->60 grows, kv 50->30 shrinks, target total 90 <= cap 100: the fit check + # passes, so the rebuild must complete live instead of OOMing on the MoE grow. + ledger = _Ledger(cap=100) + eng = _build_engine(monkeypatch, ledger, moe=40, kv=50, budget=100) + + Engine.rebuild_runtime_cache(eng, moe_cache_size=60, num_pages=30) + + assert eng.moe_offload_cache.cache_size == 60 + assert eng.kv_cache.pages == 30 and eng.num_pages == 30 + assert ledger.used == 90 # target total resident + assert ledger.peak == 90 # never above max(current, target) totals + assert _alloc_index(ledger, "kv", 30) < _alloc_index(ledger, "moe", 60) # shrink ran first + + +def test_rebuild_runtime_cache_reverse_direction_still_fits(monkeypatch): + # The mirror request (moe shrinks, kv grows) worked before the fix because the old + # fixed order happened to shrink MoE first; a naive static reorder would break it. + ledger = _Ledger(cap=100) + eng = _build_engine(monkeypatch, ledger, moe=60, kv=30, budget=100) + + Engine.rebuild_runtime_cache(eng, moe_cache_size=40, num_pages=50) + + assert eng.moe_offload_cache.cache_size == 40 + assert eng.kv_cache.pages == 50 and eng.num_pages == 50 + assert ledger.used == 90 and ledger.peak == 90 + assert _alloc_index(ledger, "moe", 40) < _alloc_index(ledger, "kv", 50) # shrink ran first + + +def test_rebuild_runtime_cache_grow_failure_rolls_back(monkeypatch): + # Budget 110 prices the 105-unit target as fitting but the real cap is 100 (fragmentation + # / unaccounted allocations), so the grow phase OOMs after the shrink phase. The + # scheduler's rollback (rebuild_cache(**prior) -> another rebuild_runtime_cache) must + # restore the snapshot geometry. An OOM inside a KV/mamba realloc still wedges the + # rollback itself; that is #526, out of scope here. + ledger = _Ledger(cap=100) + eng = _build_engine(monkeypatch, ledger, moe=40, kv=50, budget=110) + + with pytest.raises(torch.OutOfMemoryError): + Engine.rebuild_runtime_cache(eng, moe_cache_size=70, num_pages=35) + assert eng.rebuild_teardown_started # the scheduler's mid-teardown signal + + Engine.rebuild_runtime_cache(eng, moe_cache_size=40, num_pages=50) # rollback to prior + + assert eng.moe_offload_cache.cache_size == 40 + assert eng.kv_cache.pages == 50 and eng.num_pages == 50 + assert ledger.used == 90 and ledger.peak <= 100 # serving geometry restored under the cap diff --git a/tests/engine/test_sample.py b/tests/engine/test_sample.py new file mode 100644 index 0000000000..28e7195b58 --- /dev/null +++ b/tests/engine/test_sample.py @@ -0,0 +1,69 @@ +"""Presence/frequency penalties in engine.sample.Sampler (greedy and sampled rows).""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch + +from freetoken.core import SamplingParams +from freetoken.engine import sample as S + + +def _batch(*params: SamplingParams, can_decode: bool = True): + reqs = [ + SimpleNamespace(sampling_params=p, can_decode=can_decode, output_token_counts=None) + for p in params + ] + return SimpleNamespace(reqs=reqs) + + +def test_no_penalty_allocates_nothing(): + sampler = S.Sampler(torch.device("cpu"), 5) + batch = _batch(SamplingParams(temperature=0.0)) + args = sampler.prepare(batch) + assert args.penalties == [] and batch.reqs[0].output_token_counts is None + + +def test_frequency_penalty_counts_generated_tokens_on_the_greedy_path(): + sampler = S.Sampler(torch.device("cpu"), 5) + batch = _batch(SamplingParams(temperature=0.0, frequency_penalty=10.0)) + logits = torch.tensor([[0.0, 5.0, 1.0, 0.0, 0.0]]) + + assert sampler.sample(logits, sampler.prepare(batch)).tolist() == [1] + # Token 1 now carries -10, so the runner-up wins; the caller's logits stay untouched. + assert sampler.sample(logits, sampler.prepare(batch)).tolist() == [2] + assert batch.reqs[0].output_token_counts.tolist() == [0, 1, 1, 0, 0] + assert logits[0, 1] == 5.0 + + +def test_intermediate_prefill_chunks_are_not_counted(): + sampler = S.Sampler(torch.device("cpu"), 5) + batch = _batch(SamplingParams(temperature=0.0, presence_penalty=1.0), can_decode=False) + assert sampler.prepare(batch).penalties == [] + + +def test_penalized_logits_reach_the_probability_sampler(monkeypatch): + """Sampled rows (the torch top-k/top-p fallback on Pascal) see the penalized logits.""" + seen = {} + + def fake_sample_impl(logits, temperatures, top_k, top_p): + seen["logits"] = logits.clone() + return torch.argmax(logits, dim=-1) + + monkeypatch.setattr(S, "sample_impl", fake_sample_impl) + monkeypatch.setattr(S, "make_device_tensor", lambda data, dtype, device: torch.tensor(data, dtype=dtype)) + sampler = S.Sampler(torch.device("cpu"), 4) + batch = _batch( + SamplingParams(temperature=1.0, top_k=2, presence_penalty=3.0), + SamplingParams(temperature=1.0, top_k=2), + ) + batch.reqs[0].output_token_counts = torch.tensor([0, 2, 0, 0], dtype=torch.int32) + logits = torch.tensor([[0.0, 4.0, 2.0, 0.0], [0.0, 4.0, 2.0, 0.0]]) + + tokens = sampler.sample(logits, sampler.prepare(batch)) + + assert seen["logits"][0].tolist() == [0.0, 1.0, 2.0, 0.0] # presence: -3 once, not per count + assert seen["logits"][1].tolist() == [0.0, 4.0, 2.0, 0.0] + assert tokens.tolist() == [2, 1] diff --git a/tests/engine/test_sample_logprobs.py b/tests/engine/test_sample_logprobs.py new file mode 100644 index 0000000000..88fffcc2ed --- /dev/null +++ b/tests/engine/test_sample_logprobs.py @@ -0,0 +1,168 @@ +from types import SimpleNamespace + +import torch + +from freetoken.core import SamplingParams +from freetoken.engine.sample import BatchSamplingArgs, Sampler + + +def test_compute_logprobs_matches_raw_log_softmax_and_sorted_top() -> None: + sampler = Sampler(torch.device("cpu"), vocab_size=4) + + logits = torch.tensor( + [ + [2.0, 0.0, 1.0, -1.0], + [0.0, -1.0, 1.0, 3.0], + [1.0, 2.0, 3.0, 4.0], + ], + dtype=torch.float32, + ) + sampled_tokens = torch.tensor([2, 3, 0], dtype=torch.long) + args = BatchSamplingArgs( + temperatures=torch.tensor([1.0, 1.0, 1.0], dtype=torch.float32), + logprob_rows=torch.tensor([True, False, True], dtype=torch.bool), + max_top_logprobs=3, + ) + + result = sampler.compute_logprobs(logits, sampled_tokens, args) + + assert result is not None + chosen_logprobs, top_ids, top_logprobs = result + expected = torch.log_softmax(logits.float(), dim=-1) + + assert torch.isclose(chosen_logprobs[0], expected[0, sampled_tokens[0]]) + assert torch.isnan(chosen_logprobs[1]) + assert torch.isclose(chosen_logprobs[2], expected[2, sampled_tokens[2]]) + + expected_top0 = torch.topk(expected[0], k=3) + expected_top2 = torch.topk(expected[2], k=3) + assert torch.equal(top_ids[0], expected_top0.indices) + assert torch.allclose(top_logprobs[0], expected_top0.values) + assert torch.equal(top_ids[2], expected_top2.indices) + assert torch.allclose(top_logprobs[2], expected_top2.values) + + assert torch.equal(top_ids[1], torch.full((3,), -1, dtype=torch.int32)) + assert torch.isneginf(top_logprobs[1]).all() + assert top_ids.shape == (3, 3) + assert top_logprobs.shape == (3, 3) + + +def test_compute_logprobs_ignores_temperatures() -> None: + sampler = Sampler(torch.device("cpu"), vocab_size=4) + + logits = torch.tensor( + [ + [0.5, 1.0, 2.0, 3.0], + [4.0, 3.0, 2.0, 1.0], + ], + dtype=torch.float32, + ) + sampled_tokens = torch.tensor([1, 2], dtype=torch.long) + logprob_rows = torch.tensor([True, True], dtype=torch.bool) + + cold = BatchSamplingArgs( + temperatures=torch.full((2,), 0.3, dtype=torch.float32), + logprob_rows=logprob_rows, + max_top_logprobs=2, + ) + hot = BatchSamplingArgs( + temperatures=torch.full((2,), 2.5, dtype=torch.float32), + logprob_rows=logprob_rows, + max_top_logprobs=2, + ) + + cold_out = sampler.compute_logprobs(logits, sampled_tokens, cold) + hot_out = sampler.compute_logprobs(logits, sampled_tokens, hot) + + assert cold_out is not None and hot_out is not None + cold_chosen, cold_top_ids, cold_top_logprobs = cold_out + hot_chosen, hot_top_ids, hot_top_logprobs = hot_out + + assert torch.allclose(cold_chosen, hot_chosen) + assert torch.equal(cold_top_ids, hot_top_ids) + assert torch.allclose(cold_top_logprobs, hot_top_logprobs) + + +def test_compute_logprobs_returns_none_without_requested_rows() -> None: + sampler = Sampler(torch.device("cpu"), vocab_size=4) + logits = torch.zeros((2, 4), dtype=torch.float32) + sampled_tokens = torch.tensor([0, 1], dtype=torch.long) + args = BatchSamplingArgs( + temperatures=torch.ones(2), + logprob_rows=torch.zeros(2, dtype=torch.bool), + max_top_logprobs=2, + ) + + assert sampler.compute_logprobs(logits, sampled_tokens, args) is None + + +def test_compute_logprobs_handles_zero_top_logprobs() -> None: + sampler = Sampler(torch.device("cpu"), vocab_size=4) + logits = torch.tensor( + [ + [1.0, 2.0, 3.0, 4.0], + [4.0, 3.0, 2.0, 1.0], + [0.1, 0.2, 0.3, 0.4], + ], + dtype=torch.float32, + ) + sampled_tokens = torch.tensor([3, 0, 1], dtype=torch.long) + args = BatchSamplingArgs( + temperatures=torch.full((3,), 1.0), + logprob_rows=torch.tensor([True, False, True], dtype=torch.bool), + max_top_logprobs=0, + ) + + result = sampler.compute_logprobs(logits, sampled_tokens, args) + + assert result is not None + chosen_logprobs, top_ids, top_logprobs = result + expected = torch.log_softmax(logits, dim=-1) + assert torch.isclose(chosen_logprobs[0], expected[0, 3]) + assert torch.isnan(chosen_logprobs[1]) + assert torch.isclose(chosen_logprobs[2], expected[2, 1]) + assert top_ids.shape == (3, 0) + assert top_logprobs.shape == (3, 0) + + +def test_compute_logprobs_with_prepare_keeps_max_top_logprobs_clamped(monkeypatch) -> None: + # prepare() pins host tensors for the H2D copy; pinning needs a CUDA context, + # so stub the transfer helper to keep this test host-agnostic. + import freetoken.engine.sample as sample_mod + + monkeypatch.setattr( + sample_mod, "make_device_tensor", + lambda data, dtype, device: torch.tensor(data, dtype=dtype), + ) + sampler = Sampler(torch.device("cpu"), vocab_size=5) + + batch = SimpleNamespace( + reqs=[ + SimpleNamespace(sampling_params=SamplingParams(logprobs=True, top_logprobs=17)), + SimpleNamespace(sampling_params=SamplingParams(logprobs=False, top_logprobs=0)), + ] + ) + args = sampler.prepare(batch) + assert args.max_top_logprobs == sampler.vocab_size + + logits = torch.tensor( + [ + [1.0, 0.0, -1.0, 0.5, 2.0], + [2.0, 1.0, 0.0, -1.0, -2.0], + ], + dtype=torch.float32, + ) + sampled_tokens = torch.tensor([4, 0], dtype=torch.long) + + result = sampler.compute_logprobs(logits, sampled_tokens, args) + assert result is not None + chosen_logprobs, top_ids, top_logprobs = result + + assert chosen_logprobs.shape == (2,) + assert torch.isclose(chosen_logprobs[0], torch.log_softmax(logits[0], dim=-1)[4]) + assert torch.isnan(chosen_logprobs[1]) + + assert top_ids.shape == (2, sampler.vocab_size) + assert top_logprobs.shape == (2, sampler.vocab_size) + assert torch.equal(top_ids[1], torch.full((sampler.vocab_size,), -1, dtype=torch.int32)) + assert torch.isneginf(top_logprobs[1]).all() diff --git a/tests/kernels/test_autotune_cache.py b/tests/kernels/test_autotune_cache.py new file mode 100644 index 0000000000..3416117f79 --- /dev/null +++ b/tests/kernels/test_autotune_cache.py @@ -0,0 +1,29 @@ +"""The autotune L2-flush buffer is sized to the device instead of triton's flat 256 MB.""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch + +MB = 1 << 20 + + +@pytest.mark.parametrize("l2_bytes, flush_bytes", [ + (2816 * 1024, 11 * MB), # GTX 1080 Ti: 4x its 2.75 MB L2 + (0, 8 * MB), # no L2 size reported: the floor + (128 * MB, 256 * MB), # a large L2: triton's own 256 MB is the cap +]) +def test_flush_buffer_is_four_times_the_l2_within_bounds(monkeypatch, l2_bytes, flush_bytes): + from triton.backends.nvidia.driver import CudaDriver + + from freetoken.kernel.triton.autotune_cache import bound_autotune_flush_buffer + + monkeypatch.setattr(CudaDriver, "get_empty_cache_for_benchmark", CudaDriver.get_empty_cache_for_benchmark) + monkeypatch.setattr(torch.cuda, "get_device_properties", lambda index: SimpleNamespace(L2_cache_size=l2_bytes)) + bound_autotune_flush_buffer(0) + bound_autotune_flush_buffer(0) # the engine may call it again; the patch must not stack + monkeypatch.setattr(torch, "empty", lambda n, dtype, device: SimpleNamespace(nbytes=n * 4, dtype=dtype, device=device)) + buf = CudaDriver.get_empty_cache_for_benchmark(None) + assert (buf.nbytes, buf.dtype, buf.device) == (flush_bytes, torch.int, "cuda") diff --git a/tests/kernels/test_batch_memcpy.py b/tests/kernels/test_batch_memcpy.py new file mode 100644 index 0000000000..bc20c13c3c --- /dev/null +++ b/tests/kernels/test_batch_memcpy.py @@ -0,0 +1,71 @@ +"""The cudaMemcpyBatchAsync load probe must be ordered against the current stream. + +_probe allocates its destination with torch.zeros (a fill enqueued on the CURRENT +stream) but enqueues the copy on a fresh probe stream. With nothing joining the two, +the fill can land *after* the copy and clobber it, so the probe reads back zeros and +load_batch_memcpy raises "probe copied wrong bytes" on a GPU that supports the API +perfectly well. It only reproduces once the caching allocator and the stream pool are +warm -- in a cold process torch.zeros and torch.cuda.Stream() each synchronize the +device and hide the race, which is why the probe passes standalone but fails inside a +server that is mid-warmup. +""" + +from __future__ import annotations + +import os + +import pytest +import torch + +CUDA = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +JIT = pytest.mark.skipif( + os.getenv("FREETOKEN_DISABLE_JIT", "").strip().lower() in {"1", "true", "yes", "on"}, + reason="batch_memcpy has no AOT prebuild; needs runtime JIT", +) + + +def _cuda_at_least(major: int, minor: int) -> bool: + cuda = torch.version.cuda + if cuda is None: + return False + return tuple(int(x) for x in cuda.split(".")[:2]) >= (major, minor) + + +BATCH_API = pytest.mark.skipif( + not _cuda_at_least(13, 0), reason="the cudaMemcpyBatchAsync binding needs CUDA >= 13.0" +) + +# ~1s of spin on the current stream. The probe's own copy is microseconds, so any +# ordering bug shows up as the whole backlog landing on top of the copied bytes. +_BACKLOG_CYCLES = 1_500_000_000 + + +def _warm_device_state() -> None: + """Reach the state a server is in at prefill warmup: the caching allocator holds a + free block of the probe's 16-byte size and the per-device stream pool is populated, + so neither torch.zeros nor torch.cuda.Stream() synchronizes on the probe's behalf.""" + for _ in range(4): + block = torch.zeros(16, dtype=torch.uint8, device="cuda") + del block + pinned = torch.arange(16, dtype=torch.uint8).pin_memory() + del pinned + pool = [torch.cuda.Stream() for _ in range(40)] + torch.cuda.synchronize() + del pool + + +@CUDA +@JIT +@BATCH_API +def test_probe_survives_a_busy_current_stream(): + from freetoken.kernel.batch_memcpy import load_batch_memcpy + + _warm_device_state() + try: + torch.cuda._sleep(_BACKLOG_CYCLES) + # Guard against a vacuous test: if something drained the queue the race is + # not being exercised at all and a pass would mean nothing. + assert not torch.cuda.current_stream().query(), "backlog drained before the probe ran" + load_batch_memcpy() + finally: + torch.cuda.synchronize() diff --git a/tests/kernels/test_causal_conv1d_capture.py b/tests/kernels/test_causal_conv1d_capture.py new file mode 100644 index 0000000000..7489333377 --- /dev/null +++ b/tests/kernels/test_causal_conv1d_capture.py @@ -0,0 +1,163 @@ +"""The varlen GDN conv must be CUDA-graph capturable: its ``max_seq_len`` is host-known +metadata everywhere the scheduler calls it, and deriving it on device costs a D2H sync that +is illegal inside a capture region.""" + +from __future__ import annotations + +import pytest +import torch + +from freetoken.kernel.causal_conv1d import causal_conv1d_varlen + +requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") + + +def _conv_inputs(device, *, lens=(4, 2), conv_dim=8, kernel=4): + total = sum(lens) + cu = torch.tensor([0, *lens], dtype=torch.int32).cumsum(0).to(torch.int32).to(device) + return dict( + x=torch.randn(conv_dim, total, device=device, dtype=torch.bfloat16), + weight=torch.randn(conv_dim, kernel, device=device, dtype=torch.bfloat16), + conv_states=torch.randn( + len(lens) + 1, conv_dim, kernel - 1, device=device, dtype=torch.bfloat16 + ), + cu_seqlens=cu, + cache_indices=torch.arange(1, len(lens) + 1, dtype=torch.int32, device=device), + has_initial_state=torch.ones(len(lens), dtype=torch.bool, device=device), + ) + + +def _call(inputs, **extra): + return causal_conv1d_varlen( + inputs["x"], + inputs["weight"], + inputs["conv_states"], + inputs["cu_seqlens"], + inputs["cache_indices"], + inputs["has_initial_state"], + **extra, + ) + + +@requires_cuda +def test_varlen_conv_skips_the_device_to_host_sync_when_max_seq_len_is_given(monkeypatch): + device = torch.device("cuda") + inputs = _conv_inputs(device) + original_item = torch.Tensor.item + calls = [] + + def counted_item(self): + calls.append(tuple(self.shape)) + return original_item(self) + + monkeypatch.setattr(torch.Tensor, "item", counted_item) + + _call(inputs, max_seq_len=4) + torch.cuda.synchronize() + + assert calls == [] + + +@requires_cuda +def test_varlen_conv_still_derives_max_seq_len_on_device_by_default(monkeypatch): + import freetoken.kernel.backend as _backend + + device = torch.device("cuda") + inputs = _conv_inputs(device) + original_item = torch.Tensor.item + calls = [] + + def counted_item(self): + calls.append(tuple(self.shape)) + return original_item(self) + + # The sync this asserts on lives in the triton fallback; on an install with sgl_kernel + # the native kernel needs no max_seq_len and calls no .item(). Force the fallback so the + # test proves the claim on any install rather than only where the fallback is selected. + monkeypatch.setattr(_backend, "is_sgl_kernel_installed", lambda: False) + monkeypatch.setattr(torch.Tensor, "item", counted_item) + + _call(inputs) + torch.cuda.synchronize() + + assert calls, "the default path must still work without host-side metadata" + + +@requires_cuda +def test_varlen_conv_with_host_metadata_matches_the_device_derived_result(): + device = torch.device("cuda") + inputs = _conv_inputs(device) + baseline_states = inputs["conv_states"].clone() + # sgl_kernel convolves ``x`` in place (the triton fallback returns a fresh tensor), so + # both inputs have to be reset between the two calls being compared. + baseline_x = inputs["x"].clone() + + device_derived = _call(inputs).clone() + device_states = inputs["conv_states"].clone() + + inputs["conv_states"].copy_(baseline_states) + inputs["x"].copy_(baseline_x) + host_known = _call(inputs, max_seq_len=4).clone() + + assert torch.equal(host_known, device_derived) + assert torch.equal(inputs["conv_states"], device_states) + + +@requires_cuda +def test_varlen_conv_replays_inside_a_cuda_graph(): + device = torch.device("cuda") + inputs = _conv_inputs(device) + baseline_states = inputs["conv_states"].clone() + baseline_x = inputs["x"].clone() + + def reset(): + inputs["conv_states"].copy_(baseline_states) + inputs["x"].copy_(baseline_x) # sgl_kernel mutates x in place; the fallback does not + + expected = _call(inputs, max_seq_len=4).clone() + expected_states = inputs["conv_states"].clone() + + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(3): + reset() + _call(inputs, max_seq_len=4) + torch.cuda.current_stream().wait_stream(stream) + torch.cuda.synchronize() + + graph = torch.cuda.CUDAGraph() + reset() + with torch.cuda.graph(graph, stream=stream, capture_error_mode="thread_local"): + captured = _call(inputs, max_seq_len=4) + reset() + graph.replay() + torch.cuda.synchronize() + + assert torch.equal(captured, expected) + assert torch.equal(inputs["conv_states"], expected_states) + + +def test_prefill_fla_metadata_carries_the_host_known_longest_extend_len(): + from freetoken.attention.linear import build_fla_metadata + from freetoken.core import Batch, Req, SamplingParams + + reqs = [] + for index, length in enumerate((5, 2, 9)): + req = Req( + input_ids=torch.arange(length, dtype=torch.int32), + table_idx=index, + cached_len=0, + output_len=1, + uid=index, + sampling_params=SamplingParams(), + cache_handle=None, + ) + req.linear_slot_idx = index + 1 + reqs.append(req) + batch = Batch(reqs=reqs, phase="prefill") + batch.padded_reqs = batch.reqs + + fla = build_fla_metadata(batch, torch.device("cpu")) + + assert fla.max_seq_len == 9 diff --git a/tests/kernels/test_e4m3_compat.py b/tests/kernels/test_e4m3_compat.py index fa8a1f1d32..3ee6ee7fa8 100644 --- a/tests/kernels/test_e4m3_compat.py +++ b/tests/kernels/test_e4m3_compat.py @@ -16,6 +16,9 @@ capability patched to sm_80/86/89/120 and every triton launch forced into warmup (compile-only), the full wrapper->kernel paths -- uint8 views, PDL gates, e4m3 branches -- must compile for the foreign arch. +4. A cache-key guard: the @constexpr_function probes must not reference plain + python functions. triton's AST walk rejects that, and it disables every kernel + that branches on e4m3 native-ness at once, so the failure is not local. """ from __future__ import annotations @@ -64,6 +67,32 @@ def gate_kernel(output): assert output.item() == int(e4m3_native()) +def test_constexpr_probe_never_references_a_host_function(): + """triton hashes an @constexpr_function by walking its AST (runtime/jit.py: + cache_key -> record_reference), and a bare reference to a plain python function + raises "Unsupported function referenced: ". Making + e4m3_native_cx defer to the host probe did exactly that, and it took out EVERY + kernel that branches on e4m3 native-ness at once (here: the PLE gather, inside + CUDA graph capture). Modules (``target_info.cuda_capability_geq``) survive the + walk, functions do not -- so unifying the two probes has to go the other way: the + code that owns a buffer follows the buffer (tests/kernels/test_kv_fp8.py).""" + import inspect + + from freetoken.kernel.triton import e4m3_compat + + lines = inspect.getsource(e4m3_compat).splitlines() + start = next(i for i, ln in enumerate(lines) if ln.startswith("def e4m3_native_cx(")) + body = [] + for line in lines[start + 1:]: + if line and not line.startswith((" ", "\t")): + break + body.append(line) + assert "e4m3_native(" not in "\n".join(body), ( + "a constexpr_function may not call a host function: triton's cache-key walk " + "rejects it and every e4m3 kernel stops compiling" + ) + + # ====================================================================================== # 1. Primitives vs the native fp8 unit (needs sm_89+ hardware for the reference). # ====================================================================================== @@ -136,6 +165,49 @@ def k(x_ptr, y_ptr, N, BLOCK: tl.constexpr): assert int(((y != ref) & ~(y.isnan() & ref.isnan())).sum()) == 0 +def test_kv_tile_scaled16_agrees_with_the_f32_loader(): + """kv_load_e4m3_tile_scaled16 is kv_load_e4m3_tile_f32 with its last two steps -- + the widen to fp32 and the ``* 256.0`` -- left for the caller to fold into the + per-(token, kv_head) dequant scale it has to apply anyway. That fold is only legal + if the two loaders agree on every code, so pin it: all 256, NaN patterns included. + + Second assertion pins the property that lets the 16-bit tile be narrowed to a bf16 + compute dtype for free: every value carries at most the code's own 3 mantissa bits + and lies in +-1.75, so bf16 holds it exactly. That is what makes scaling AFTER the + dot strictly more accurate than upstream's scale-then-narrow. + + Third pins masked lanes at zero in both, so a masked tile contributes nothing once + the scale is applied to the dot output instead of to the tile.""" + import triton + import triton.language as tl + + from freetoken.kernel.triton.e4m3_compat import ( + KV_TILE_SCALE, + kv_load_e4m3_tile_f32, + kv_load_e4m3_tile_scaled16, + ) + + @triton.jit + def k(v_ptr, wide_ptr, narrow_ptr, N, BLOCK: tl.constexpr): + offs = tl.arange(0, BLOCK) + mask = offs < N + tl.store(wide_ptr + offs, kv_load_e4m3_tile_f32(v_ptr + offs, mask)) + tl.store(narrow_ptr + offs, kv_load_e4m3_tile_scaled16(v_ptr + offs, mask)) + + n = 256 + v = torch.zeros(2 * n, dtype=torch.uint8, device="cuda") + v[:n] = torch.arange(n, dtype=torch.uint8, device="cuda") + wide = torch.empty(2 * n, dtype=torch.float32, device="cuda") + narrow = torch.empty(2 * n, dtype=torch.float16, device="cuda") + k[(1,)](v, wide, narrow, n, BLOCK=2 * n) + + assert torch.equal(narrow.to(torch.float32) * KV_TILE_SCALE.value, wide) + assert torch.equal( + narrow.to(torch.bfloat16).to(torch.float32), narrow.to(torch.float32) + ) + assert not wide[n:].any() and not narrow[n:].any() + + # ====================================================================================== # Shared emit path: every affected wrapper, deterministic inputs. # ====================================================================================== @@ -251,6 +323,78 @@ def _emit_all(path: str) -> None: out["nv_moe_prefill"] = f32(fused_experts_nvfp4( hid4, gup, gus4, gug, dnp, dns4, dng, tw, tids, S)) + # fp8 KV cache: the fused quantize+scatter writer the pools call in store_kv, and the + # QSA sparse reader that turns codes + row scales back into operands. Codes are plain + # bytes on every arch (kv_quant.kv_codes_dtype), so comparing them as BYTES is what + # pins a native run and a forced-EMU run to the very same encoding, and the two + # QSA runs -- one over codes, one over the same real numbers pre-rounded into bf16 -- + # must produce identical bits: the reader casts back to the query dtype before the + # dot, so any decode difference lands on the comparison instead of hiding in a + # tolerance. + from freetoken.kernel.triton.kv_quant import ( + alloc_codes, codes_to_f32, quantize_kv_to_cache, + ) + from freetoken.kernel.triton.qsa import qsa_sparse_paged_attention + + slots, kvh, hd, page = 128, 2, 64, 64 + kv_in = (torch.randn(slots, kvh * hd, device=dev, dtype=torch.float32) * 4.0).to( + torch.bfloat16 + ) + k_codes = alloc_codes((slots, kvh, hd), dev) + v_codes = alloc_codes((slots, kvh, hd), dev) + k_sc = torch.zeros((slots, kvh), dtype=torch.float32, device=dev) + v_sc = torch.zeros_like(k_sc) + quantize_kv_to_cache( + k=kv_in, + v=kv_in.flip(-1).contiguous(), + out_loc=torch.arange(slots, dtype=torch.int32, device=dev), + k_cache=k_codes, + v_cache=v_codes, + k_scale=k_sc, + v_scale=v_sc, + ) + out["kvfp8_codes"] = k_codes.view(torch.uint8).to(torch.int16).cpu() + out["kvfp8_scale"] = k_sc.cpu() + + pages = slots // page + kc4, vc4 = k_codes.view(pages, page, kvh, hd), v_codes.view(pages, page, kvh, hd) + q = torch.randn(2, 2 * kvh, hd, device=dev, dtype=torch.bfloat16) + sel = ( + torch.arange(2 * page, dtype=torch.int32, device=dev)[None, :] + .repeat(2, 1) + .contiguous() + ) + table = torch.arange(pages, dtype=torch.int32, device=dev)[None, :].contiguous() + t2r = torch.zeros(2, dtype=torch.int32, device=dev) + out["kvfp8_qsa"] = f32(qsa_sparse_paged_attention( + q, kc4, vc4, sel, table, t2r, k_scale=k_sc, v_scale=v_sc)) + # The very same numbers, pre-rounded into a bf16 cache. The reader casts its + # dequantized operands to the query dtype before tl.dot, so the two runs must end up + # bit-identical (asserted where launches really execute: tests/kernels/test_qsa_fp8.py + # -- the compile gate below runs warmup-only, where outputs are never written). + out["kvfp8_qsa_bf16"] = f32(qsa_sparse_paged_attention( + q, + (codes_to_f32(kc4) * k_sc.view(pages, page, kvh, 1)).to(torch.bfloat16), + (codes_to_f32(vc4) * v_sc.view(pages, page, kvh, 1)).to(torch.bfloat16), + sel, table, t2r)) + + + # a K no tile divides takes the masked loads, whose fill value must suit an fp8 pointer + Kr, H5, I5 = 272, 272, 80 + pkr = torch.randint(0, 256, (N, Kr // 2), dtype=torch.uint8, device=dev) + scr = (torch.rand(N, Kr // 16, device=dev) * 2 + 0.1).to(FP8) + out["nv_gemm_inkernel_ragged_k"] = f32(nvfp4_dense_linear( + torch.randn(8, Kr, dtype=torch.bfloat16, device=dev), pkr, scr, g)) + gup5 = torch.randint(0, 256, (S, 2 * I5, H5 // 2), dtype=torch.uint8, device=dev) + gus5 = (torch.rand(S, 2 * I5, H5 // 16, device=dev) + 0.1).to(FP8) + gug5 = (torch.rand(S, 2 * I5, device=dev) * 0.05 + 0.01).to(torch.float16) + dnp5 = torch.randint(0, 256, (S, H5, I5 // 2), dtype=torch.uint8, device=dev) + dns5 = (torch.rand(S, H5, I5 // 16, device=dev) + 0.1).to(FP8) + dng5 = (torch.rand(S, H5, device=dev) * 0.05 + 0.01).to(torch.float16) + out["nv_moe_prefill_ragged_k"] = f32(fused_experts_nvfp4( + torch.randn(2, H5, dtype=torch.bfloat16, device=dev), + gup5, gus5, gug5, dnp5, dns5, dng5, tw, tids, S)) + torch.save(out, path) diff --git a/tests/kernels/test_hc_combine_norm.py b/tests/kernels/test_hc_combine_norm.py new file mode 100644 index 0000000000..7f46755479 --- /dev/null +++ b/tests/kernels/test_hc_combine_norm.py @@ -0,0 +1,48 @@ +"""Pins for the fused hc combine + grouped RMSNorm path (hc_combine_norm). + +hc_combine_norm replaces the standalone pair (hc_combine -> grouped_gemma_rmsnorm) +in the Qwen4-exp decoder loop. The swap is only sound if it is BITWISE identical: +the hyper-connection streams feed every later layer, and a bf16 flip here compounds +into different greedy tokens by the end of the stack. We assert torch.equal, not +assert_close: an approx-tolerance test would not have caught the regression this +path shipped with at first attempt -- the fused kernel reduced squares over the +[8, block] tile view (a two-axis tree) while the standalone kernel reduces the +flat row tree; fp32 ulps that differ between trees flip bf16 rounding. Both paths +on the same data must return the same bytes: fails before the reduction-tree fix +in _hc_combine_norm_kernel, passes after. +""" + +from __future__ import annotations + +import pytest +import torch + +from freetoken.kernel.triton.hc import hc_combine, grouped_gemma_rmsnorm, hc_combine_norm + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") + +HC, H = 4, 2560 # model geometry: the reduction tree specializes on the block layout +DIM = HC * H +EPS = 1e-6 + + +def _inputs(seed: int, T: int): + g = torch.Generator(device="cuda").manual_seed(seed) + bf = lambda shape: torch.randn(shape, generator=g, device="cuda", dtype=torch.float32).to(torch.bfloat16) + return ( + bf((T, DIM)), # residual streams R + bf((T, H)), # block output y + bf((T, HC)), # inject scores s + (torch.randn(DIM, generator=g, device="cuda", dtype=torch.float32) * 0.1).to(torch.bfloat16), # norm weight + ) + + +@pytest.mark.parametrize("T", [1, 8]) +@pytest.mark.parametrize("seed", range(6)) +def test_combine_norm_matches_split_bitwise(seed: int, T: int) -> None: + R, y, s, w = _inputs(seed, T) + r2_split = hc_combine(R, y, s, HC) + rn_split = grouped_gemma_rmsnorm(r2_split, w, EPS, HC) + r2_fused, rn_fused = hc_combine_norm(R, y, s, w, EPS, HC) + assert torch.equal(r2_fused, r2_split) + assert torch.equal(rn_fused, rn_split) diff --git a/tests/kernels/test_host_embedding.py b/tests/kernels/test_host_embedding.py new file mode 100644 index 0000000000..866687a9f1 --- /dev/null +++ b/tests/kernels/test_host_embedding.py @@ -0,0 +1,64 @@ +"""--embed-weights host: the UVA row gather from pinned host RAM must equal the GPU table lookup.""" + +import pytest +import torch + +needs_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") + + +@pytest.fixture(autouse=True) +def _tp(): + from freetoken.distributed import set_tp_info, try_get_tp_info + + if try_get_tp_info() is None: + set_tp_info(rank=0, size=1) + + +def _embedding(dtype: torch.dtype, rows: int = 1000, dim: int = 2560, device: str = "cuda"): + from freetoken.layers.embedding import VocabParallelEmbedding + + torch.manual_seed(0) + emb = VocabParallelEmbedding(rows, dim) + emb.weight = torch.randn(rows, dim, device=device, dtype=dtype) + return emb + + +@needs_cuda +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32]) +def test_host_rows_equal_the_gpu_lookup(dtype): + emb = _embedding(dtype) + ids = torch.randint(0, 1000, (37,), device="cuda", dtype=torch.int32) + want = emb.forward(ids) + emb.place_on_host() + assert not emb.weight.is_cuda and emb.weight.is_pinned() + got = emb.forward(ids) + assert got.is_cuda and got.dtype is dtype + assert torch.equal(got, want) + + +def test_a_table_the_gather_cannot_read_is_refused_before_pinning(): + emb = _embedding(torch.float64, rows=4, dim=8, device="cpu") + with pytest.raises(ValueError, match="cannot gather a torch.float64 table"): + emb.place_on_host() + assert emb._host_ptr is None + + +@needs_cuda +def test_host_lookup_replays_in_a_cuda_graph_with_new_ids(): + emb = _embedding(torch.bfloat16) + table = emb.weight.clone() + emb.place_on_host() + ids = torch.zeros(8, device="cuda", dtype=torch.int64) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + emb.forward(ids) # compile outside the capture + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = emb.forward(ids) + for seed in range(3): + ids.copy_(torch.randint(0, 1000, (8,), generator=torch.Generator().manual_seed(seed))) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(out, table[ids]) diff --git a/tests/kernels/test_kv_fp8.py b/tests/kernels/test_kv_fp8.py new file mode 100644 index 0000000000..8176468868 --- /dev/null +++ b/tests/kernels/test_kv_fp8.py @@ -0,0 +1,263 @@ +"""FP8 (e4m3) KV quantization: the store kernel against an independent oracle. + +Expectations come from a brute-force nearest-code search over an e4m3 table decoded +from first principles (sign / exponent / mantissa), NOT from the kernels' own +rounding helpers -- so a drift in ``round_e4m3`` or in the new ``e4m3_f32_to_u8`` +encoder fails here instead of being blessed by itself. +""" + +from __future__ import annotations + +import pytest +import torch + +if not torch.cuda.is_available(): # pragma: no cover + pytest.skip("CUDA required", allow_module_level=True) + +from freetoken.distributed import set_tp_info, try_get_tp_info +from freetoken.kernel.triton.kv_quant import ( + KV_SCALE_DTYPE, + alloc_codes, + codes_to_f32, + kv_codes_dtype, + quantize_kv_to_cache, +) + +DEV = torch.device("cuda") +FP8_MAX = 448.0 +CODE_448 = 0x7E + + +def _init_tp() -> None: + if try_get_tp_info() is None: + set_tp_info(rank=0, size=1) + + +@pytest.fixture(autouse=True) +def _tp(): + _init_tp() + + +def _e4m3_table() -> dict[int, float]: + """Every finite e4m3fn value, decoded by hand from its bit fields. + + code = S EEEE MMM, exponent bias 7: normal (E>0) -> (-1)^S * 2^(E-7) * (1+M/8); + subnormal (E==0) -> (-1)^S * 2^-6 * (M/8). E==15 & M==7 (0x7F/0xFF) is the NaN + pattern and is dropped, which caps the format at +-448. + """ + values: dict[int, float] = {} + for code in range(256): + sign = -1.0 if (code >> 7) & 1 else 1.0 + exp, mant = (code >> 3) & 0x0F, code & 0x07 + if exp == 15 and mant == 7: + continue + values[code] = sign * ( + (2.0**-6) * (mant / 8.0) if exp == 0 else (2.0 ** (exp - 7)) * (1.0 + mant / 8.0) + ) + return values + + +E4M3_VALUES = _e4m3_table() +CODES = torch.tensor(sorted(E4M3_VALUES), dtype=torch.int32) +GRID = torch.tensor([E4M3_VALUES[c] for c in CODES.tolist()], dtype=torch.float64) +NEGATIVE_ZERO = 0x80 + + +def _as_bytes(t: torch.Tensor) -> torch.Tensor: + """A code buffer as raw bytes (the fp8 view on sm_89+, uint8 below it).""" + return t if t.dtype == torch.uint8 else t.view(torch.uint8) + + +def _canonical_zero(codes: torch.Tensor) -> torch.Tensor: + """Fold -0.0 (0x80) onto +0.0 (0x00). + + The two store paths legitimately disagree on zero's sign bit: sm_89+ converts with + ``x.to(fp8e4nv)`` (keeps it), the emulated path rounds first and + ``round_e4m3(-0.0)`` is documented to return +0.0. They decode to the same number, + so a test that compares bytes must not care which one it got. + """ + return torch.where(codes == NEGATIVE_ZERO, torch.zeros_like(codes), codes) + + +def _ref_codes(x: torch.Tensor) -> torch.Tensor: + """Independent encoder: nearest grid value by brute force, ties resolved to the + EVEN code (RNE). ``x`` is any shape; returns int32 codes.""" + grid = GRID.to(x.device) + codes = CODES.to(x.device) + dist = (x.to(torch.float64).reshape(-1, 1) - grid.unsqueeze(0)).abs() + near = dist == dist.min(dim=-1, keepdim=True).values + big = 1 << 30 + codes = codes.unsqueeze(0).expand_as(dist) + any_code = torch.where(near, codes, torch.full_like(codes, big)) + even = near & (codes % 2 == 0) + even_code = torch.where(even, codes, torch.full_like(codes, big)) + has_even = (even_code < big).any(dim=-1) + return torch.where(has_even, even_code.min(dim=-1).values, any_code.min(dim=-1).values) + + +def _store(rows_k: torch.Tensor, rows_v: torch.Tensor): + """Quantize ``[T, heads, dim]`` rows into a fresh code buffer, returning + ``(k_codes, v_codes, k_scales, v_scales)``.""" + _init_tp() + tokens, heads, dim = rows_k.shape + k_cache = alloc_codes((tokens, heads, dim), DEV) + v_cache = alloc_codes((tokens, heads, dim), DEV) + k_scale = torch.zeros((tokens, heads), dtype=KV_SCALE_DTYPE, device=DEV) + v_scale = torch.zeros((tokens, heads), dtype=KV_SCALE_DTYPE, device=DEV) + quantize_kv_to_cache( + k=rows_k.reshape(tokens, -1), + v=rows_v.reshape(tokens, -1), + out_loc=torch.arange(tokens, dtype=torch.int32, device=DEV), + k_cache=k_cache, + v_cache=v_cache, + k_scale=k_scale, + v_scale=v_scale, + ) + torch.cuda.synchronize() + return k_cache, v_cache, k_scale, v_scale + + +def test_scale_is_amax_over_e4m3_max(): + torch.manual_seed(0) + k = (torch.randn(4, 2, 64, device=DEV, dtype=torch.bfloat16) * 3.0).to(torch.bfloat16) + v = torch.randn_like(k) + _, _, k_scale, v_scale = _store(k, v) + # The kernel widens to fp32 before the amax/divide, so the reference must too: + # a bf16 intermediate would round the expected scale and hide a precision bug. + torch.testing.assert_close( + k_scale, k.abs().to(torch.float32).amax(dim=-1) / FP8_MAX, rtol=1e-6, atol=0 + ) + torch.testing.assert_close( + v_scale, v.abs().to(torch.float32).amax(dim=-1) / FP8_MAX, rtol=1e-6, atol=0 + ) + + +def test_codes_match_the_reference_quantizer_and_reconstruction_is_close(): + torch.manual_seed(1) + tokens, heads, dim = 8, 3, 128 + # Feed the rows as a qkv slice, the way the attention backends really do: the + # row pitch is then wider than the row, which the store kernel must honour. + qkv = torch.randn(tokens, heads * dim * 3, device=DEV, dtype=torch.bfloat16) + qkv[:, heads * dim : 2 * heads * dim] *= 5.0 # K: large magnitude + qkv[:, 2 * heads * dim :] *= 0.01 # V: subnormal end of the e4m3 grid + _, k_rows, v_rows = qkv.split(heads * dim, dim=-1) + k = k_rows.view(tokens, heads, dim) + v = v_rows.view(tokens, heads, dim).clamp(-FP8_MAX, FP8_MAX) + k_cache, v_cache, k_scale, v_scale = _store(k, v) + + for rows, cache, scale in ((k, k_cache, k_scale), (v, v_cache, v_scale)): + f32 = rows.to(torch.float32) + ref_scale = f32.abs().amax(dim=-1, keepdim=True) / FP8_MAX + expected = _ref_codes((f32 / ref_scale).clamp(-FP8_MAX, FP8_MAX)) + got = _canonical_zero(_as_bytes(cache).reshape(-1).to(torch.int32)) + expected = _canonical_zero(expected) + assert torch.equal(got, expected), ( + f"{int((got != expected).sum())} code mismatches of {got.numel()}" + ) + deq = codes_to_f32(cache) * scale.unsqueeze(-1) + err = (deq - f32).abs().max(dim=-1).values + assert torch.all(err <= 0.08 * f32.abs().amax(dim=-1)), float(err.max()) + + +def test_encoder_inverts_the_grid_through_the_scale_one_path(): + """Pack every e4m3 grid value into a row that also holds 448.0: the row scale is + then exactly 1.0, so each stored byte IS the encoder's answer for that value.""" + dim, per_row = 256, 255 + pairs = sorted(E4M3_VALUES.items()) + tokens = -(-len(pairs) // per_row) + rows = torch.zeros(tokens, 1, dim, dtype=torch.float32) + expected = torch.zeros(tokens, dim, dtype=torch.uint8) + for t in range(tokens): + rows[t, 0, 0] = FP8_MAX # the amax anchor + expected[t, 0] = CODE_448 + for j in range(per_row): + i = t * per_row + j + if i >= len(pairs): + break + code, value = pairs[i] + rows[t, 0, j + 1] = value + expected[t, j + 1] = code + + k_cache, _, k_scale, _ = _store( + rows.to(DEV, dtype=torch.bfloat16), + torch.zeros(tokens, 1, dim, dtype=torch.bfloat16, device=DEV), + ) + assert torch.equal(k_scale[:, 0], torch.ones_like(k_scale[:, 0])) + got = _canonical_zero(_as_bytes(k_cache)[:, 0, :]) + want = _canonical_zero(expected.to(DEV)) + bad = got != want + assert not bad.any().item(), ( + f"{int(bad.sum())} of {want.numel()} grid values round-tripped wrong; " + f"first at {bad.nonzero()[0].tolist()}: expected " + f"{want[bad][0].item():#x} got {got[bad][0].item():#x}" + ) + + +def test_zero_row_stays_finite_and_exact(): + k = torch.zeros(2, 2, 32, device=DEV, dtype=torch.bfloat16) + k_cache, _, k_scale, _ = _store(k, k.clone()) + assert torch.isfinite(k_scale).all() + assert (k_scale > 0).all(), "an all-zero row must still store a usable scale" + assert (codes_to_f32(k_cache) == 0).all() + + +def test_codes_are_plain_bytes_and_the_kernel_decode_matches_torch(): + """The KV codec never puts an fp8 type in front of Triton. + + Codes live in a uint8 buffer on EVERY architecture and the kernel widens them with + the software decoder, while the expectation below is torch's OWN e4m3 cast of those + very bytes. That pins the one claim the design rests on: byte for byte, the + software decode reads what a native fp8 unit would -- which is what lets the + quantized cache behave identically on GPUs where the fp8 type is illegal. + """ + import triton + import triton.language as tl + + from freetoken.kernel.triton.e4m3_compat import kv_load_e4m3_tile_f32 + + rows = torch.randn(5, 2, 64, device=DEV, dtype=torch.bfloat16) * 3.0 + k_cache, _, _, _ = _store(rows, rows.clone()) + assert kv_codes_dtype() is torch.uint8, "keep the fp8 type out of kernel signatures" + assert k_cache.dtype is torch.uint8 and k_cache.element_size() == 1 + want = codes_to_f32(k_cache) # torch reinterprets these bytes as e4m3 and casts + + @triton.jit + def read_out(codes_ptr, out_ptr, n, BLOCK: tl.constexpr): + offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + got = kv_load_e4m3_tile_f32(codes_ptr + offs, offs < n) + tl.store(out_ptr + offs, got, mask=offs < n) + + n = k_cache.numel() + out = torch.zeros(n, dtype=torch.float32, device=DEV) + read_out[(triton.cdiv(n, 256),)](k_cache, out, n, BLOCK=256) + flat = want.reshape(-1) + assert torch.equal(out, flat), ( + f"{int((out != flat).sum())} of {n} codes decode differently from torch's cast" + ) + + +def test_kv_codec_has_no_arch_or_dtype_branch(): + """The two rejected designs had one thing in common: they chose an arm. + + The compile-time fp8-native probe answers a question the allocator already + answered -- and on one box answered wrongly -- while a test against the pointer's + element type is NOT pruned by triton, so the dead arm still gets type-checked. + That is how an int mask fill ended up in front of an fp8 pointer, twice. Both + codecs are straight-line now; pin that, plus the identifiers of the two rejected + designs, so neither creeps back in as a "fast path". + """ + import inspect + + from freetoken.kernel.triton import kv_quant + from freetoken.kernel.triton.e4m3_compat import kv_load_e4m3_tile_f32 + + for obj in (kv_load_e4m3_tile_f32, kv_quant._kv_quant_scatter_kernel): + src = inspect.getsource(getattr(obj, "fn", obj)) + for banned in ("e4m3_native", "dtype.element_ty"): + assert banned not in src, f"{banned} is back in {obj.__name__}" + body = src.split('"""')[-1].splitlines() + arms = [ + line.strip() for line in body + if line.strip().startswith(("if ", "elif ", "else")) + ] + assert not arms, f"{obj.__name__} must not branch: {arms}" diff --git a/tests/kernels/test_kv_nvfp4.py b/tests/kernels/test_kv_nvfp4.py new file mode 100644 index 0000000000..662b44bf26 --- /dev/null +++ b/tests/kernels/test_kv_nvfp4.py @@ -0,0 +1,349 @@ +"""NVFP4 KV against independent torch rounding and dense attention references.""" + +from types import SimpleNamespace + +import pytest +import torch + +from freetoken.distributed import set_tp_info, try_get_tp_info +from freetoken.kvcache.mha_pool import MHAKVCache + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +def _pool(dim=128, heads=2, slots=96, layer_ids=None): + if try_get_tp_info() is None: + set_tp_info(rank=0, size=1) + return MHAKVCache(heads, 4, dim, slots // 4, 4, torch.bfloat16, + torch.device("cuda"), layer_ids=layer_ids, kv_quant="nvfp4") + + +def _reference(x): + shape = x.shape + x = x.float().reshape(*shape[:-1], -1, 16) + row = x.abs().flatten(-2).amax(-1).clamp_min(1e-10) / 2688.0 + block = (x.abs().amax(-1) / (6 * row[..., None])).clamp_max(448).to(torch.float8_e4m3fn) + denom = block.float() * row[..., None] + normalized = torch.where(denom[..., None] > 0, x / denom[..., None], 0) + grid = torch.tensor([0, .5, 1, 1.5, 2, 3, 4, 6], device=x.device) + distance = (normalized.abs()[..., None] - grid).abs() + nearest = distance == distance.amin(-1, keepdim=True) + codes = torch.arange(8, device=x.device).expand_as(distance) + # Prefer even codes on ties, independently of the kernel's threshold encoding. + rank = torch.where(nearest, codes % 2 * 8 + codes, 32) + code = rank.argmin(-1) | ((normalized < 0).long() * 8) + code = code.reshape(shape) + packed = (code[..., ::2] | (code[..., 1::2] << 4)).to(torch.uint8) + return packed, block.view(torch.uint8), row + + +def _decode(pool, which, layer=1): + codes = getattr(pool, f"{which}_cache")(layer).flatten(0, 1) + block = getattr(pool, f"{which}_block_scale")(layer).view(torch.float8_e4m3fn).float() + row = getattr(pool, f"{which}_scale")(layer) + code = torch.stack((codes & 15, codes >> 4), -1).flatten(-2).long() + grid = torch.tensor([0, .5, 1, 1.5, 2, 3, 4, 6, + 0, -.5, -1, -1.5, -2, -3, -4, -6], device=codes.device) + return grid[code] * block.repeat_interleave(16, -1) * row[..., None] + + +def _decode_latent(pool, layer=0): + codes = pool.latent_rows(layer) + code = torch.stack((codes & 15, codes >> 4), -1).flatten(-2).long() + grid = torch.tensor([0, .5, 1, 1.5, 2, 3, 4, 6, + 0, -.5, -1, -1.5, -2, -3, -4, -6], device=codes.device) + block = pool.latent_block_scale(layer).view(torch.float8_e4m3fn).float() + return grid[code] * block.repeat_interleave(16, -1) * pool.latent_scale(layer)[:, None] + + +@pytest.mark.parametrize("rope", [0, 64]) +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +def test_latent_scatter_matches_reference_and_preserves_prefix(rope, dtype): + from freetoken.kvcache.dsa_pool import MLAKVCache + + torch.manual_seed(73) + dim = 512 + rope + pool = MLAKVCache(dim, 4, 2, 64, dtype, torch.device("cuda"), + layer_ids=(1, 3), kv_quant="nvfp4") + x = torch.randn(7, dim + 32, device="cuda", dtype=dtype)[:, :dim] + x[0].zero_() + x[1, :16] *= 100 + x[2] *= 1e-5 + loc = torch.tensor([64, 3, 127, 14, 6, 90, 31], device="cuda") + pool.store_kv(x[:4, :512], x[:4, 512:], loc[:4], 3) + before = [v.clone() for v in (pool.latent_rows(3), pool.latent_scale(3), + pool.latent_block_scale(3))] + pool.store_kv(x[4:, :512], x[4:, 512:], loc[4:], 3) + pool.store_kv(x[:0, :512], x[:0, 512:], loc[:0], 3) + for got, saved in zip((pool.latent_rows(3), pool.latent_scale(3), + pool.latent_block_scale(3)), before): + torch.testing.assert_close(got[loc[:4]], saved[loc[:4]], rtol=0, atol=0) + packed, block, row = _reference(x) + torch.testing.assert_close(pool.latent_rows(3)[loc], packed) + torch.testing.assert_close(pool.latent_block_scale(3)[loc], block) + torch.testing.assert_close(pool.latent_scale(3)[loc], row) + assert torch.count_nonzero(pool.latent_rows(1)) == 0 + assert torch.isfinite(_decode_latent(pool, 3)).all() + assert pool.k_cache(3).data_ptr() == pool.v_cache(3).data_ptr() + # Reused physical slots replace codes and both scales together. + x[4:].mul_(0.125) + pool.store_kv(x[4:, :512], x[4:, 512:], loc[4:], 3) + packed, block, row = _reference(x[4:]) + torch.testing.assert_close(pool.latent_rows(3)[loc[4:]], packed) + torch.testing.assert_close(pool.latent_block_scale(3)[loc[4:]], block) + torch.testing.assert_close(pool.latent_scale(3)[loc[4:]], row) + + +@pytest.mark.parametrize("rope", [0, 64]) +@pytest.mark.parametrize("splits", [0, 4]) +@pytest.mark.parametrize("queries,broadcast", [(1, False), (5, False), (5, True)]) +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +def test_sparse_mla_nvfp4_matches_restored_reference(rope, splits, queries, broadcast, dtype): + from freetoken.kernel.triton.glm_dsa_sparse import glm_dsa_sparse_attn + from freetoken.kvcache.dsa_pool import MLAKVCache + + torch.manual_seed(74) + dim, n = 512 + rope, 67 + pool = MLAKVCache(dim, 1, 2, 64, torch.bfloat16, torch.device("cuda"), kv_quant="nvfp4") + x = torch.randn(n, dim, device="cuda", dtype=torch.bfloat16) + x *= torch.linspace(.2, 2, n, device="cuda")[:, None] + if rope: + x[:, 512:] *= 3 + loc = torch.randperm(128, device="cuda")[:n] + pool.store_kv(x[:, :512], x[:, 512:], loc, 0) + q = torch.randn(2, queries, 19, dim, device="cuda", dtype=dtype) + sel = loc.repeat(2, 1 if broadcast else queries, 1).to(torch.int32) + sel[1] = sel[1].flip(-1) + sel[..., 5] = -1 + cnt = torch.full((2, queries), n, device="cuda", dtype=torch.int32) + cnt[0, 0] = 0 + cnt[1, 0] = 13 + out = glm_dsa_sparse_attn( + q, pool.latent_rows(0), sel, .04, counts=cnt, d_v=512, + pool_scale=pool.latent_scale(0), pool_block_scale=pool.latent_block_scale(0), + kv_quant="nvfp4", force_splits=splits, + ) + decoded = _decode_latent(pool) + ref = torch.zeros_like(out, dtype=torch.float32) + for b in range(2): + for m in range(queries): + rows = sel[b, 0 if broadcast else m, :int(cnt[b, m])] + rows = rows[rows >= 0].long() + if rows.numel(): + kv = decoded[rows] + ref[b, m] = (q[b, m].float() @ kv.T * .04).softmax(-1) @ kv[:, :512] + torch.testing.assert_close(out.float(), ref, atol=2e-2, rtol=1e-2) + + +@pytest.mark.parametrize("dim", [16, 48, 64, 128, 256, 512]) +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +def test_scatter_matches_independent_reference(dim, dtype): + torch.manual_seed(41) + pool = _pool(dim) + # K is a projection slice; V is independently contiguous. + k = torch.randn(5, 4 * dim, device="cuda", dtype=dtype)[:, :2 * dim] + v = torch.randn(5, 2 * dim, device="cuda", dtype=dtype) + k[0].zero_() + k[1, :16] *= 100 + v[2] *= 1e-5 + loc = torch.tensor([7, 1, 95, 31, 12], device="cuda", dtype=torch.int64) + pool.store_kv(k, v, loc, 1) + for which, source in (("k", k), ("v", v)): + packed, block, row = _reference(source.reshape(5, 2, dim)) + torch.testing.assert_close(getattr(pool, f"{which}_cache")(1).flatten(0, 1)[loc], packed) + torch.testing.assert_close(getattr(pool, f"{which}_block_scale")(1)[loc], block) + torch.testing.assert_close(getattr(pool, f"{which}_scale")(1)[loc], row) + assert torch.isfinite(_decode(pool, which)).all() + assert torch.count_nonzero(_decode(pool, which)[0]) == 0 + + +def test_e2m1_grid_and_round_to_even_boundaries(): + positive = torch.tensor([0, .25, .5, .75, 1, 1.25, 1.5, 1.75, + 2, 2.5, 3, 3.5, 4, 5, 6], device="cuda") + values = torch.cat((positive, -positive)) + values = torch.cat((values, values.nextafter(torch.full_like(values, float("inf"))), + values.nextafter(torch.full_like(values, -float("inf"))))) + pool = _pool(dim=32, slots=192) + rows = torch.zeros(values.numel(), 2, 32, device="cuda") + rows[:, :, 0] = values[:, None] + rows[:, :, 15] = 6 + rows[:, :, 31] = 2688 # Forces row_scale=1 and the first block_scale=1. + loc = torch.arange(values.numel(), device="cuda", dtype=torch.int32) + pool.store_kv(rows.flatten(1), rows.flatten(1), loc, 1) + packed, block, row = _reference(rows) + torch.testing.assert_close(pool.k_cache(1).flatten(0, 1)[loc], packed) + torch.testing.assert_close(pool.k_block_scale(1)[loc], block) + torch.testing.assert_close(pool.k_scale(1)[loc], row) + + +@pytest.mark.parametrize("dim", [64, 128, 256, 512]) +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +@pytest.mark.parametrize("mode", ["paged", "decode", "extend", "split"]) +def test_attention_reads_packed_cache(dim, dtype, mode): + from freetoken.kernel.triton.attention import ( + decode_paged_attention, extend_paged_attention, paged_attention, + ) + + torch.manual_seed(42) + pool = _pool(dim) + n, prefix = 67, 35 + k, v = [torch.randn(n, 2 * dim, device="cuda", dtype=dtype) for _ in range(2)] + # Distinct token/head and block scales expose accidental FP8 rescaling of NVFP4. + row_gain = torch.linspace(0.25, 2.0, n * 2, device="cuda").view(n, 2, 1) + block_gain = torch.linspace(0.5, 1.5, dim // 16, device="cuda").repeat_interleave(16) + k.view(n, 2, dim).mul_(row_gain * block_gain) + v.view(n, 2, dim).mul_(row_gain.flip(0) * block_gain.flip(0)) + loc = torch.randperm(95, device="cuda")[:n] + 1 + pool.store_kv(k, v, loc, 1) + tokens = 1 if mode in ("paged", "decode") else n - prefix + q = torch.randn(tokens, 6, dim, device="cuda", dtype=dtype) + indptr = torch.tensor([0, n], device="cuda", dtype=torch.int32) + pos = torch.arange(n - tokens, n, device="cuda") + args = dict(q=q, k_cache=pool.k_cache(1).flatten(0, 1), + v_cache=pool.v_cache(1).flatten(0, 1), k_scale=pool.k_scale(1), + v_scale=pool.v_scale(1), k_block_scale=pool.k_block_scale(1), + v_block_scale=pool.v_block_scale(1), kv_quant="nvfp4", sm_scale=dim ** -.5) + kd, vd = _decode(pool, "k")[loc], _decode(pool, "v")[loc] + if mode == "split": + kd[prefix:] = k[prefix:].view(-1, 2, dim).float() + vd[prefix:] = v[prefix:].view(-1, 2, dim).float() + # Tensor-core paths round restored K/V to the query dtype before dot products. + if mode != "paged": + kd, vd = kd.to(q.dtype).float(), vd.to(q.dtype).float() + kd, vd = [x.repeat_interleave(3, 1).transpose(0, 1) for x in (kd, vd)] + score = torch.einsum("thd,hnd->thn", q.float(), kd) * dim ** -.5 + score.masked_fill_(torch.arange(n, device="cuda")[None, None, :] > pos[:, None, None], -float("inf")) + ref = torch.einsum("thn,hnd->thd", score.softmax(-1), vd).to(q.dtype) + if mode == "paged": + actual = paged_attention(**args, indptr=indptr, indices=loc, + q_to_req=torch.zeros(tokens, device="cuda", dtype=torch.int32), q_positions=pos) + elif mode == "decode": + actual = decode_paged_attention(**args, indptr=indptr, indices=loc, q_positions=pos, + attn_logits=torch.empty(1, 6, 8, dim, device="cuda"), + attn_lse=torch.empty(1, 6, 8, device="cuda"), + num_kv_splits=torch.tensor([8], device="cuda", dtype=torch.int32), max_kv_splits=8) + else: + extra = {} if mode == "extend" else dict(k_extend=k[prefix:].view(-1, 2, dim), v_extend=v[prefix:].view(-1, 2, dim)) + actual = extend_paged_attention(**args, **extra, + qo_indptr=torch.tensor([0, tokens], device="cuda", dtype=torch.int32), + kv_indptr=indptr, kv_indices=loc, + prefix_lens=torch.tensor([prefix], device="cuda", dtype=torch.int32), max_q_len=tokens) + torch.testing.assert_close(actual, ref, atol=0.008, rtol=0.025) + + +def test_pool_budget_rebuild_and_layer_mapping(): + from freetoken.kvcache.base import spec_kv_bytes_per_token + from freetoken.models.config import KVCacheGroupSpec + + pool = _pool(layer_ids=(1, 3)) + spec = KVCacheGroupSpec(name="full", layer_ids=(1, 3), num_kv_heads=2, head_dim=128, sliding_window=None) + cfg = SimpleNamespace(kv_quant="nvfp4", dtype=torch.bfloat16, tp_info=SimpleNamespace(size=1)) + assert spec_kv_bytes_per_token(spec, cfg) == pool.unit_bytes()[0] == 2 * 2 * 2 * 76 + with pytest.raises(KeyError): + pool.k_block_scale(0) + for pages in (32, 8): + pool.rebuild(pages) + assert pool.k_cache(3).shape == (pages, 4, 2, 64) + assert pool.k_block_scale(3).shape == (pages * 4, 2, 8) + assert pool.unit_bytes()[0] == 608 + assert torch.count_nonzero(_decode(pool, "k", 3)) == 0 + + +def test_store_cuda_graph_replay_changes_slots(): + pool = _pool() + k = torch.randn(2, 256, device="cuda", dtype=torch.bfloat16) + loc = torch.tensor([1, 2], device="cuda", dtype=torch.int32) + pool.store_kv(k, k, loc, 1) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + pool.store_kv(k, k, loc, 1) + k.mul_(2) + loc.copy_(torch.tensor([4, 9], device="cuda", dtype=torch.int32)) + graph.replay() + expected, _, _ = _reference(k.view(2, 2, 128)) + torch.testing.assert_close(pool.k_cache(1).flatten(0, 1)[loc], expected) + + +def test_hybrid_swa_pool_packs_both_groups_and_rebuilds(monkeypatch): + from freetoken.distributed.info import DistributedInfo + from freetoken.kvcache.hybrid_swa_pool import HybridSWAKVCache + from freetoken.models.config import KVCacheGroupSpec + + monkeypatch.setattr( + "freetoken.kvcache.hybrid_swa_pool.get_tp_info", + lambda: DistributedInfo(rank=0, size=1), + ) + groups = ( + KVCacheGroupSpec("full", (1,), 2, 64, None), + KVCacheGroupSpec("swa", (0,), 2, 128, 32), + ) + pool = HybridSWAKVCache( + groups=groups, + num_layers=2, + num_full_pages=12, + page_size=1, + num_swa_tokens=18, + dtype=torch.bfloat16, + device=torch.device("cuda"), + kv_quant="nvfp4", + ) + assert pool.k_cache(1).shape == (12, 1, 2, 32) + assert pool.k_cache(0).shape == (18, 1, 2, 64) + assert pool.k_block_scale(1).shape == (12, 2, 4) + assert pool.k_block_scale(0).shape == (18, 2, 8) + assert pool.unit_bytes() == (160, 304) + + full_loc = torch.tensor([4, 9], device="cuda", dtype=torch.int32) + swa_loc = torch.tensor([2, 7], device="cuda", dtype=torch.int32) + pool.alloc_swa(swa_loc) + for layer, dim, loc, cache_loc in ((1, 64, full_loc, full_loc), (0, 128, swa_loc, None)): + rows = torch.randn(2, 2 * dim, device="cuda", dtype=torch.bfloat16) + pool.store_kv(rows, rows, loc, layer) + if cache_loc is None: + cache_loc = pool.translate_loc_from_full_to_swa(loc) + expected, block, row = _reference(rows.view(2, 2, dim)) + codes = pool.k_cache(layer).flatten(0, 1)[cache_loc] + torch.testing.assert_close(codes, expected) + torch.testing.assert_close(pool.k_block_scale(layer)[cache_loc], block) + torch.testing.assert_close(pool.k_scale(layer)[cache_loc], row) + + pool.rebuild(num_full_pages=6, num_swa_tokens=10) + assert pool.k_cache(1).shape == (6, 1, 2, 32) + assert pool.k_cache(0).shape == (10, 1, 2, 64) + assert pool.unit_bytes() == (160, 304) + assert pool.swa_available_size() == 9 + + +def test_backend_decode_graph_replays_new_kv_and_page_tables(monkeypatch): + from freetoken.attention.triton import TritonAttentionBackend, TritonMetadata + from freetoken.kernel.triton.attention import paged_attention + + pool = _pool() + monkeypatch.setattr("freetoken.attention.triton.get_global_ctx", + lambda: SimpleNamespace(kv_cache=pool)) + backend = TritonAttentionBackend(SimpleNamespace(num_qo_heads=6, head_dim=128)) + q = torch.randn(2, 6, 128, device="cuda", dtype=torch.bfloat16) + k, v = [torch.randn(2, 256, device="cuda", dtype=q.dtype) for _ in range(2)] + loc = torch.tensor([7, 11], device="cuda", dtype=torch.int32) + indptr = torch.tensor([0, 1, 2], device="cuda", dtype=torch.int32) + positions = torch.zeros(2, device="cuda", dtype=torch.int32) + indices = loc.clone() + req = torch.arange(2, device="cuda", dtype=torch.int32) + metadata = TritonMetadata(cu_seqlens_q_gpu=indptr, indptr=indptr, indices=indices, + q_to_req=req, q_positions=positions, is_decode=True, prefix_lens=positions, + max_q_len=1) + batch = SimpleNamespace(out_loc=loc, attn_metadata=metadata) + backend.forward(q, k, v, 1, batch) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + actual = backend.forward(q, k, v, 1, batch) + for new_slots in ([31, 44], [7, 11]): + k.normal_() + v.normal_() + q.normal_() + loc.copy_(torch.tensor(new_slots, device="cuda", dtype=torch.int32)) + indices.copy_(loc) + graph.replay() + ref = paged_attention(q, _decode(pool, "k").to(q.dtype), + _decode(pool, "v").to(q.dtype), indptr, indices, req, positions, 128 ** -.5) + torch.testing.assert_close(actual, ref, atol=0.008, rtol=0.025) diff --git a/tests/kernels/test_qsa_fp8.py b/tests/kernels/test_qsa_fp8.py new file mode 100644 index 0000000000..1fa6669f86 --- /dev/null +++ b/tests/kernels/test_qsa_fp8.py @@ -0,0 +1,278 @@ +"""QSA sparse attention reading an fp8 KV pool (kernel/triton/qsa/attend.py). + +The oracle is the SAME data in a bf16 cache, not a tolerance. The kernel widens the e4m3 +codes to fp32, multiplies by the row scale, and casts back to the query dtype before +``tl.dot`` -- so a cache holding ``c * s`` and a cache holding ``c`` with scale ``s`` feed +the matmul bit-identical operands, and the two runs must agree bit for bit. Anything wrong +in scale indexing (the page/offset -> slot arithmetic), in the K-vs-V broadcast direction, +or in the masked-fill values shows up as a mismatch instead of slop inside an epsilon. + +Two quantizer variants are covered: + * hand-made codes with power-of-two scales, which also keeps the real values on the e4m3 + grid, so the fp8 buffer, the bf16 buffer and the test agree on the numbers; + * the fused store kernel the pools actually call (``kv_quant.quantize_kv_to_cache``, + amax/448 scales) over ordinary gaussian rows -- whose quantization error against the + ORIGINAL rows is bounded separately, because that part is quality, not exactness. + +Each parametrization also covers a different launcher profile: a small +``rows * kv_heads`` pushes it into split-K (partials + merge kernel), a large one takes +the direct-write path. +""" + +from __future__ import annotations + +import pytest +import torch + +from freetoken.kernel.triton.kv_quant import ( + FP8, + alloc_codes, + codes_to_f32, + kv_codes_dtype, + quantize_kv_to_cache, +) +from freetoken.kernel.triton.qsa import qsa_sparse_paged_attention + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available(), reason="Triton attention needs CUDA" +) + +PAGE = 64 # the page size the qsa_sparse backend registers +HEAD_DIM = 64 + + +def _on_grid(shape, device, generator) -> torch.Tensor: + """Values of the form +-m * 2^e with m in [8, 16): four significant bits, i.e. every + one of them is representable in e4m3 AND in bf16, so encoding is lossless and both + caches hold the very same real numbers.""" + mantissa = torch.randint( + 8, 16, shape, device=device, generator=generator, dtype=torch.int32 + ) + exponent = torch.randint( + -6, 5, shape, device=device, generator=generator, dtype=torch.int32 + ) + sign = torch.where( + torch.randint(0, 2, shape, device=device, generator=generator).bool(), 1.0, -1.0 + ) + return sign * mantissa.to(torch.float32) * torch.pow(2.0, exponent.to(torch.float32)) + + +def _code_buffer(values_f32: torch.Tensor) -> torch.Tensor: + """Encode e4m3-exact fp32 values into a buffer of the pool's code dtype.""" + codes = values_f32.to(FP8) + if kv_codes_dtype() is torch.uint8: + codes = codes.view(torch.uint8) + assert codes.dtype == kv_codes_dtype(), codes.dtype + return codes.contiguous() + + +def _layout(rows: int, topk: int, num_req: int): + """(indices, block_table, token_to_req) over three 64-token pages. + + Tokens 32..31+topk straddle pages 0 and 1, and the two requests map their logical + pages to physical ones in OPPOSITE order -- a page-table or page/offset slip cannot + hide behind a symmetric fixture. + """ + device = "cuda" + block_table = torch.tensor( + [[0, 1, 2], [2, 1, 0]], dtype=torch.int32, device=device + )[:num_req].contiguous() + indices = ( + torch.arange(topk, dtype=torch.int32, device=device)[None, :] + 32 + ).repeat(rows, 1).contiguous() + token_to_req = ( + torch.arange(rows, dtype=torch.int32, device=device) % num_req + ).contiguous() + return indices, block_table, token_to_req + + +def _run(codes, scales, q, indices, block_table, token_to_req): + return qsa_sparse_paged_attention( + q, + codes[0], + codes[1], + indices, + block_table, + token_to_req, + k_scale=scales[0], + v_scale=scales[1], + ) + + + +# rows 1 x 1 kv head -> base_programs 1 -> BLOCK_N 16, 4 tiles -> NUM_SPLITS 4 (split-K). +# rows 16 x 2 -> base_programs 32 -> BLOCK_N 64, 1 tile -> NUM_SPLITS 1 (direct). +@pytest.mark.parametrize(("rows", "kv_heads", "topk"), [(1, 1, 64), (16, 2, 64)]) +def test_fp8_codes_match_the_bf16_cache_bit_for_bit(rows, kv_heads, topk): + torch.manual_seed(3) + device = torch.device("cuda") + num_pages, num_query_heads = 3, 2 * kv_heads + slots = num_pages * PAGE + + k_values = _on_grid((num_pages, PAGE, kv_heads, HEAD_DIM), device, None) + v_values = _on_grid((num_pages, PAGE, kv_heads, HEAD_DIM), device, None) + # Alternate the scale exponent by slot -- a scale broadcast along the wrong axis then + # changes the answer instead of cancelling out -- and give K and V opposite parities. + parity = torch.arange(slots, device=device, dtype=torch.float32) % 2 + k_scale = ( + torch.pow(2.0, parity * 5 - 3).unsqueeze(-1).expand(slots, kv_heads).contiguous() + ) + v_scale = ( + torch.pow(2.0, (1 - parity) * 4 - 2).unsqueeze(-1).expand(slots, kv_heads).contiguous() + ) + k_ref = (k_values * k_scale.view(num_pages, PAGE, kv_heads, 1)).to(torch.bfloat16) + v_ref = (v_values * v_scale.view(num_pages, PAGE, kv_heads, 1)).to(torch.bfloat16) + + q = torch.randn(rows, num_query_heads, HEAD_DIM, device=device, dtype=torch.bfloat16) + indices, block_table, token_to_req = _layout(rows, topk, num_req=2) + + got = _run( + (_code_buffer(k_values), _code_buffer(v_values)), + (k_scale, v_scale), + q, + indices, + block_table, + token_to_req, + ) + want = qsa_sparse_paged_attention( + q, k_ref, v_ref, indices, block_table, token_to_req + ) + assert torch.equal(got, want), ( + "fp8 QSA attend diverged from the same data in a bf16 cache (max diff " + f"{(got.float() - want.float()).abs().max().item():.3e})" + ) + + +@pytest.mark.parametrize(("rows", "kv_heads", "topk"), [(1, 1, 64), (16, 2, 64)]) +def test_pool_writer_codes_match_the_bf16_cache_bit_for_bit(rows, kv_heads, topk): + """Ordinary rows through the fused quantize+scatter writer store_kv calls: the + codes/scales it produces, decoded on the host with torch's own e4m3 cast (never the + decoder under test), must feed the dot identical operands.""" + torch.manual_seed(5) + device = torch.device("cuda") + num_pages, num_query_heads = 3, 2 * kv_heads + slots = num_pages * PAGE + + k_rows = torch.randn(slots, kv_heads * HEAD_DIM, device=device, dtype=torch.float32) + v_rows = torch.randn(slots, kv_heads * HEAD_DIM, device=device, dtype=torch.float32) + # Wide amplitude spread: every 8th row carries a 64x outlier. That is what a per-row + # amax scale exists to absorb, and what a scale read from the wrong slot blows up on. + outlier = (torch.arange(slots, device=device) % 8 == 0).unsqueeze(-1) + k_rows = k_rows * torch.where(outlier, 64.0, 1.0) + v_rows = v_rows * torch.where(outlier.flip(0), 32.0, 0.5) + + k_flat = alloc_codes((slots, kv_heads, HEAD_DIM), device) + v_flat = alloc_codes((slots, kv_heads, HEAD_DIM), device) + k_scale = torch.zeros((slots, kv_heads), dtype=torch.float32, device=device) + v_scale = torch.zeros_like(k_scale) + quantize_kv_to_cache( + k=k_rows, + v=v_rows, + out_loc=torch.arange(slots, dtype=torch.int32, device=device), + k_cache=k_flat, + v_cache=v_flat, + k_scale=k_scale, + v_scale=v_scale, + ) + torch.cuda.synchronize() + + k_ref = (codes_to_f32(k_flat) * k_scale.unsqueeze(-1)).to(torch.bfloat16) + v_ref = (codes_to_f32(v_flat) * v_scale.unsqueeze(-1)).to(torch.bfloat16) + q = torch.randn(rows, num_query_heads, HEAD_DIM, device=device, dtype=torch.bfloat16) + indices, block_table, token_to_req = _layout(rows, topk, num_req=2) + + shape = (num_pages, PAGE, kv_heads, HEAD_DIM) + got = _run( + (k_flat.view(shape), v_flat.view(shape)), + (k_scale, v_scale), + q, + indices, + block_table, + token_to_req, + ) + want = qsa_sparse_paged_attention( + q, k_ref.view(shape), v_ref.view(shape), indices, block_table, token_to_req + ) + assert torch.equal(got, want), ( + "fused-writer codes diverged (max diff " + f"{(got.float() - want.float()).abs().max().item():.3e})" + ) + + # Quantization QUALITY, bounded per row: e4m3's half-ulp is 2^-4 of the binade a + # value lands in, and the writer scales each row so its amax sits at 448. + for source, codes, scales in ((k_rows, k_flat, k_scale), (v_rows, v_flat, v_scale)): + assert torch.isfinite(scales).all() and (scales > 0).all() + decoded = codes_to_f32(codes) * scales.unsqueeze(-1) + original = source.view(slots, kv_heads, HEAD_DIM) + amax = original.abs().amax(dim=-1, keepdim=True) + rel = ((decoded - original).abs() / amax.clamp_min(1e-9)).max() + assert rel.item() <= 0.07, f"e4m3 grid error {rel.item():.4f} above its half-ulp" + + +@pytest.mark.parametrize("which", ["k_only", "v_only", "bf16_cache", "wrong_shape"]) +def test_scale_arguments_are_validated(which): + """A half-supplied or mismatched scale pair must fail loudly: attending over raw + codes while believing they are bf16 is the exact failure mode this feature cannot be + allowed to have, and it produces plausible garbage rather than an error.""" + device = "cuda" + kv_heads, rows = 1, 1 + k = torch.zeros(2, PAGE, kv_heads, HEAD_DIM, device=device, dtype=kv_codes_dtype()) + v = torch.zeros_like(k) + scale = torch.ones(2 * PAGE, kv_heads, dtype=torch.float32, device=device) + q = torch.zeros(rows, 2 * kv_heads, HEAD_DIM, device=device, dtype=torch.bfloat16) + indices = torch.zeros(rows, 8, dtype=torch.int32, device=device) + block_table = torch.zeros(1, 2, dtype=torch.int32, device=device) + token_to_req = torch.zeros(rows, dtype=torch.int32, device=device) + + kwargs = {"k_scale": scale, "v_scale": scale} + if which == "k_only": + kwargs.pop("v_scale") + elif which == "v_only": + kwargs.pop("k_scale") + elif which == "bf16_cache": + k, v = k.to(torch.bfloat16), v.to(torch.bfloat16) + else: + kwargs["k_scale"] = scale[:-1] + + with pytest.raises(ValueError, match="QSA"): + qsa_sparse_paged_attention( + q, k, v, indices, block_table, token_to_req, **kwargs + ) + + +def test_qsa_scoring_refuses_fp8_operands(): + """The indexer's scoring dot is 16-bit only, and the wrapper has to say so. + + In the field an e4m3 ``q_index`` -- produced by a backend that took its scratch dtype + from a pool reporting its STORE dtype -- surfaced as ``CompilationError: Unsupported + rhs dtype fp8e4nv`` inside CUDA-graph capture: a dead scheduler and a stopped API + server, forty seconds after the pool allocated fine. Same mistake, now stopped at + the call with a name on it. + """ + from freetoken.kernel.triton.qsa import qsa_mqa_paged + + device = "cuda" + rows, heads, dim, pages, cmp_page = 2, 2, 64, 2, 16 + good = { + "q": torch.zeros(rows, heads, dim, device=device, dtype=torch.bfloat16), + "k_cache": torch.zeros( + pages, cmp_page, 1, dim, device=device, dtype=torch.bfloat16 + ), + "page_table": torch.zeros(1, pages, dtype=torch.int32, device=device), + "token_to_req": torch.zeros(rows, dtype=torch.int32, device=device), + "query_positions": torch.arange(rows, dtype=torch.int32, device=device), + "sequence_lengths": torch.zeros(1, dtype=torch.int32, device=device), + "compress_ratio": cmp_page, + # Zero columns -> the wrapper returns before launching. This control proves the + # new check still lets the dtype it exists to allow through. + "logits": torch.zeros(rows, 0, dtype=torch.float32, device=device), + "visible_blocks": torch.zeros(rows, dtype=torch.int32, device=device), + } + qsa_mqa_paged(**good) + + for name in ("q", "k_cache"): + bad = dict(good) + bad[name] = alloc_codes(tuple(good[name].shape), device) + with pytest.raises(ValueError, match="16-bit only"): + qsa_mqa_paged(**bad) + diff --git a/tests/kernels/test_qsa_nvfp4.py b/tests/kernels/test_qsa_nvfp4.py new file mode 100644 index 0000000000..ec96e34563 --- /dev/null +++ b/tests/kernels/test_qsa_nvfp4.py @@ -0,0 +1,196 @@ +"""QSA sparse attention over the packed NVFP4 KV tier.""" + +from __future__ import annotations + +import pytest +import torch + +from freetoken.kernel.triton.kv_nvfp4 import quantize_nvfp4_to_cache +from freetoken.kernel.triton.qsa import qsa_sparse_paged_attention + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available(), reason="Triton attention needs CUDA" +) + +PAGE = 64 +HEAD_DIM = 64 + + +def _layout(rows: int, topk: int): + block_table = torch.tensor( + [[0, 1, 2], [2, 1, 0]], dtype=torch.int32, device="cuda" + ).contiguous() + indices = ( + torch.arange(topk, dtype=torch.int32, device="cuda")[None, :] + 32 + ).repeat(rows, 1).contiguous() + token_to_req = torch.arange(rows, dtype=torch.int32, device="cuda") % 2 + return indices, block_table, token_to_req.contiguous() + + +def _decode(codes: torch.Tensor, row_scale: torch.Tensor, block_scale: torch.Tensor): + """Decode with torch, independently of the Triton read path under test.""" + lo = (codes & 15).to(torch.long) + hi = (codes >> 4).to(torch.long) + code = torch.stack((lo, hi), dim=-1).reshape(*codes.shape[:-1], HEAD_DIM) + magnitude = torch.tensor( + [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0], device=codes.device + )[code & 7] + value = torch.where(code < 8, magnitude, -magnitude) + block = block_scale.view(torch.float8_e4m3fn).to(torch.float32) + return (value * block.repeat_interleave(16, dim=-1) * row_scale.unsqueeze(-1)).to( + torch.bfloat16 + ) + + +@pytest.mark.parametrize(("rows", "kv_heads"), [(1, 1), (16, 2)]) +def test_qsa_nvfp4_matches_its_bf16_decode(rows, kv_heads): + """Page-table indirection must address the same packed slot and both scale tiers.""" + torch.manual_seed(19) + pages, query_heads, slots = 3, 2 * kv_heads, 3 * PAGE + k = torch.randn(slots, kv_heads * HEAD_DIM, device="cuda", dtype=torch.bfloat16) + v = torch.randn_like(k) * 0.25 + out_loc = torch.arange(slots, dtype=torch.int32, device="cuda") + code_shape = (slots, kv_heads, HEAD_DIM // 2) + block_shape = (slots, kv_heads, HEAD_DIM // 16) + k_codes = torch.zeros(code_shape, dtype=torch.uint8, device="cuda") + v_codes = torch.zeros_like(k_codes) + k_scale = torch.zeros((slots, kv_heads), dtype=torch.float32, device="cuda") + v_scale = torch.zeros_like(k_scale) + k_block = torch.zeros(block_shape, dtype=torch.uint8, device="cuda") + v_block = torch.zeros_like(k_block) + quantize_nvfp4_to_cache( + k, v, out_loc, k_codes, v_codes, k_scale, v_scale, k_block, v_block + ) + torch.cuda.synchronize() + + q = torch.randn(rows, query_heads, HEAD_DIM, device="cuda", dtype=torch.bfloat16) + indices, block_table, token_to_req = _layout(rows, topk=64) + packed_shape = (pages, PAGE, kv_heads, HEAD_DIM // 2) + dense_shape = (pages, PAGE, kv_heads, HEAD_DIM) + got = qsa_sparse_paged_attention( + q, + k_codes.view(packed_shape), + v_codes.view(packed_shape), + indices, + block_table, + token_to_req, + k_scale=k_scale, + v_scale=v_scale, + kv_quant="nvfp4", + k_block_scale=k_block, + v_block_scale=v_block, + ) + want = qsa_sparse_paged_attention( + q, + _decode(k_codes, k_scale, k_block).view(dense_shape), + _decode(v_codes, v_scale, v_block).view(dense_shape), + indices, + block_table, + token_to_req, + ) + assert torch.equal(got, want), ( + "NVFP4 QSA attend diverged from its bf16 decode (max diff " + f"{(got.float() - want.float()).abs().max().item():.3e})" + ) + + +def test_qsa_nvfp4_requires_both_block_scale_tensors(): + q = torch.zeros(1, 2, HEAD_DIM, device="cuda", dtype=torch.bfloat16) + codes = torch.zeros(2, PAGE, 1, HEAD_DIM // 2, device="cuda", dtype=torch.uint8) + scales = torch.ones(2 * PAGE, 1, device="cuda", dtype=torch.float32) + indices = torch.zeros(1, 8, device="cuda", dtype=torch.int32) + block_table = torch.zeros(1, 2, device="cuda", dtype=torch.int32) + token_to_req = torch.zeros(1, device="cuda", dtype=torch.int32) + + with pytest.raises(ValueError, match="NVFP4"): + qsa_sparse_paged_attention( + q, + codes, + codes, + indices, + block_table, + token_to_req, + k_scale=scales, + v_scale=scales, + kv_quant="nvfp4", + ) + + +def test_qsa_nvfp4_splitk_reads_value_rows(): + """A zero query makes the split-K result the selected V-row average.""" + torch.manual_seed(23) + slots = 3 * PAGE + v = torch.randn(slots, HEAD_DIM, device="cuda", dtype=torch.bfloat16) + codes = torch.zeros((slots, 1, HEAD_DIM // 2), dtype=torch.uint8, device="cuda") + row = torch.zeros((slots, 1), dtype=torch.float32, device="cuda") + block = torch.zeros((slots, 1, HEAD_DIM // 16), dtype=torch.uint8, device="cuda") + quantize_nvfp4_to_cache( + v, v, torch.arange(slots, dtype=torch.int32, device="cuda"), codes, codes, + row, row, block, block, + ) + q = torch.zeros(1, 2, HEAD_DIM, device="cuda", dtype=torch.bfloat16) + indices, block_table, token_to_req = _layout(1, topk=64) + args = (indices, block_table, token_to_req) + got = qsa_sparse_paged_attention( + q, codes.view(3, PAGE, 1, HEAD_DIM // 2), codes.view(3, PAGE, 1, HEAD_DIM // 2), + *args, k_scale=row, v_scale=row, kv_quant="nvfp4", k_block_scale=block, + v_block_scale=block, + ) + ref = _decode(codes, row, block).view(3, PAGE, 1, HEAD_DIM) + want = qsa_sparse_paged_attention(q, ref, ref, *args) + assert torch.equal(got, want) + + +def test_qsa_nvfp4_wide_head_prefill_uses_narrow_tiles(monkeypatch): + """At HEAD_DIM 256 the prefill profile must not decode 64-token fp32 tiles on 2 warps: + Triton 3.8 spills that variant (~10x slower, ~550 MiB of driver local memory).""" + from freetoken.kernel.triton.qsa import attend + + head_dim, kv_heads, query_heads, rows, pages = 256, 2, 24, 32, 3 + seen = [] + kernel = attend._qsa_sparse_paged_gqa_splitk_kernel + + class Spy: + def __getitem__(self, grid): + launch = kernel[grid] + + def run(*args, **kwargs): + seen.append((kwargs["BLOCK_N"], kwargs["num_warps"])) + return launch(*args, **kwargs) + return run + + torch.manual_seed(29) + slots = pages * PAGE + k = torch.randn(slots, kv_heads * head_dim, device="cuda", dtype=torch.bfloat16) + v = torch.randn_like(k) * 0.25 + codes = [torch.zeros((slots, kv_heads, head_dim // 2), dtype=torch.uint8, device="cuda") for _ in range(2)] + rows_s = [torch.zeros((slots, kv_heads), dtype=torch.float32, device="cuda") for _ in range(2)] + blocks = [torch.zeros((slots, kv_heads, head_dim // 16), dtype=torch.uint8, device="cuda") for _ in range(2)] + quantize_nvfp4_to_cache( + k, v, torch.arange(slots, dtype=torch.int32, device="cuda"), codes[0], codes[1], + rows_s[0], rows_s[1], blocks[0], blocks[1], + ) + q = torch.randn(rows, query_heads, head_dim, device="cuda", dtype=torch.bfloat16) + indices, block_table, token_to_req = _layout(rows, topk=128) + packed = (pages, PAGE, kv_heads, head_dim // 2) + monkeypatch.setattr(attend, "_qsa_sparse_paged_gqa_splitk_kernel", Spy()) + got = qsa_sparse_paged_attention( + q, codes[0].view(packed), codes[1].view(packed), indices, block_table, token_to_req, + k_scale=rows_s[0], v_scale=rows_s[1], kv_quant="nvfp4", + k_block_scale=blocks[0], v_block_scale=blocks[1], + ) + assert seen == [(16, 4)] + + def decode(c, r, b): + lo, hi = (c & 15).to(torch.long), (c >> 4).to(torch.long) + code = torch.stack((lo, hi), dim=-1).reshape(*c.shape[:-1], head_dim) + mag = torch.tensor([0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0], device=c.device)[code & 7] + val = torch.where(code < 8, mag, -mag) + blk = b.view(torch.float8_e4m3fn).to(torch.float32).repeat_interleave(16, dim=-1) + return (val * blk * r.unsqueeze(-1)).to(torch.bfloat16).view(pages, PAGE, kv_heads, head_dim) + + want = qsa_sparse_paged_attention( + q, decode(codes[0], rows_s[0], blocks[0]), decode(codes[1], rows_s[1], blocks[1]), + indices, block_table, token_to_req, + ) + torch.testing.assert_close(got.float(), want.float(), atol=2e-3, rtol=2e-2) diff --git a/tests/kernels/test_triton_attention.py b/tests/kernels/test_triton_attention.py index a29f621ddc..3176ad71ce 100644 --- a/tests/kernels/test_triton_attention.py +++ b/tests/kernels/test_triton_attention.py @@ -70,6 +70,14 @@ def k_cache(self, layer_id): def v_cache(self, layer_id): return self.v + # A 16-bit pool answers None here. The backend reads it on every path since the + # fp8 store landed, so the double answers it too. + def k_scale(self, layer_id): + return None + + def v_scale(self, layer_id): + return None + kv_cache = FakeKVCache() monkeypatch.setattr( "freetoken.attention.triton.get_global_ctx", @@ -666,6 +674,36 @@ def test_select_extend_tile_is_shared_memory_aware(head_dim, smem_optin, expecte assert _select_extend_tile(head_dim, block_d, smem_optin) == expected +@pytest.mark.parametrize( + ("head_dim", "smem_optin", "expected_16bit", "expected_fp8"), + [ + # consumer opt-in smem (sm_86/sm_89 ~99KB): a 1-byte cache buys BLOCK_N 32 -> 64 + (256, 101376, (64, 32), (64, 64)), + # where the fast tile already fits, the cache dtype changes nothing + (256, 232448, (128, 64), (128, 64)), + # unknown budget stays conservative for both + (256, 0, (64, 32), (64, 32)), + ], +) +def test_select_extend_tile_uses_kv_cache_element_size( + head_dim, smem_optin, expected_16bit, expected_fp8 +): + """The q tile is always 2 bytes/element but K and V follow the cache, so charging + K/V at 2 bytes regardless makes an fp8 cache run a smaller tile than it has shared + memory for. On an RTX 3070 (sm_86, 99KB opt-in, head_dim 256) that was BLOCK_N 32 + where 64 fits, worth ~9% of prefill time at 99k context. + + The 16-bit column is the no-regression half: ``kv_bytes=2`` reproduces the previous + budget identically, since (M + 2N) * D * 2 == (M*2 + 2N*2) * D.""" + import triton + + from freetoken.kernel.triton.attention import _select_extend_tile + + block_d = triton.next_power_of_2(head_dim) + assert _select_extend_tile(head_dim, block_d, smem_optin, 2) == expected_16bit + assert _select_extend_tile(head_dim, block_d, smem_optin, 1) == expected_fp8 + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="Triton attention needs CUDA") def test_triton_backend_stores_kv_and_matches_reference(monkeypatch): from freetoken.attention import AttentionSpec @@ -688,6 +726,14 @@ def k_cache(self, layer_id): def v_cache(self, layer_id): return self.v + # A 16-bit pool answers None here. The backend reads it on every path since the + # fp8 store landed, so the double answers it too. + def k_scale(self, layer_id): + return None + + def v_scale(self, layer_id): + return None + device = torch.device("cuda") head_dim = 256 page_table = torch.tensor([[0, 1], [2, 3]], dtype=torch.int32, device=device) @@ -758,6 +804,12 @@ def k_cache(self, layer_id): def v_cache(self, layer_id): return self.v + def k_scale(self, layer_id): # 16-bit pool: no fp8 scales + return None + + def v_scale(self, layer_id): + return None + device = torch.device("cuda") head_dim = 256 kv_cache = FakeKVCache(device, head_dim) @@ -884,3 +936,271 @@ def test_triton_metadata_keeps_full_indices_and_optional_swa_indices(monkeypatch assert metadata.indices.tolist() == [10, 11, 20, 21, 22] assert metadata.swa_indices is not None assert metadata.swa_indices.tolist() == [110, 111, 120, 121, 122] + + +def _fp8_cache(k_rows: torch.Tensor, v_rows: torch.Tensor): + """Quantize ``[slots, heads, dim]`` KV into codes + scales, in the layout the + attention kernels expect.""" + from freetoken.kernel.triton.kv_quant import alloc_codes, quantize_kv_to_cache + + slots, heads, dim = k_rows.shape + k_cache = alloc_codes((slots, heads, dim), k_rows.device) + v_cache = alloc_codes((slots, heads, dim), k_rows.device) + k_scale = torch.zeros((slots, heads), dtype=torch.float32, device=k_rows.device) + v_scale = torch.zeros((slots, heads), dtype=torch.float32, device=k_rows.device) + quantize_kv_to_cache( + k=k_rows.reshape(slots, heads * dim), + v=v_rows.reshape(slots, heads * dim), + out_loc=torch.arange(slots, dtype=torch.int32, device=k_rows.device), + k_cache=k_cache, + v_cache=v_cache, + k_scale=k_scale, + v_scale=v_scale, + ) + torch.cuda.synchronize() + return k_cache, v_cache, k_scale, v_scale + + +def _dequantized(codes: torch.Tensor, scale: torch.Tensor) -> torch.Tensor: + """Host-side decode (torch's e4m3 cast), so the reference never reuses the + software decoder the kernel is being tested against.""" + from freetoken.kernel.triton.kv_quant import codes_to_f32 + + return codes_to_f32(codes) * scale.unsqueeze(-1) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Triton attention needs CUDA") +@pytest.mark.parametrize(("head_dim", "num_kv_heads"), [(64, 4), (128, 2)]) +def test_decode_paged_attention_decodes_fp8_scales(head_dim: int, num_kv_heads: int): + """fp8 decode must equal the SAME data dequantized by hand. + + Not compared against the bf16 cache: the fp8-vs-bf16 gap is quantization error, + already bounded in tests/kernels/test_kv_fp8.py. What this pins is that the + kernel applies the right scale to the right row -- a dropped or mis-indexed + scale is off by a factor of amax/448, which no tolerance hides. + """ + from freetoken.kernel.triton.attention import decode_paged_attention + + torch.manual_seed(11) + device = torch.device("cuda") + batch, num_q_heads, max_kv_splits = 2, 8, 8 + seq_lens = [6, 9] + total_kv = sum(seq_lens) + q = torch.randn(batch, num_q_heads, head_dim, device=device, dtype=torch.bfloat16) + # Row magnitudes far from 1, so a missing scale cannot pass by luck. + k_rows = (torch.randn(total_kv, num_kv_heads, head_dim, device=device) * 8.0).to( + torch.bfloat16 + ) + v_rows = (torch.randn(total_kv, num_kv_heads, head_dim, device=device) * 0.05).to( + torch.bfloat16 + ) + indptr = torch.tensor([0, seq_lens[0], total_kv], dtype=torch.int32, device=device) + indices = torch.arange(total_kv, dtype=torch.int32, device=device) + q_positions = torch.tensor( + [seq_lens[0] - 1, seq_lens[1] - 1], dtype=torch.int64, device=device + ) + q_to_req = torch.arange(batch, dtype=torch.int32, device=device) + attn_logits = torch.empty( + batch, num_q_heads, max_kv_splits, head_dim, dtype=torch.float32, device=device + ) + attn_lse = torch.empty(batch, num_q_heads, max_kv_splits, dtype=torch.float32, device=device) + num_kv_splits = torch.full((batch,), max_kv_splits, dtype=torch.int32, device=device) + sm_scale = head_dim**-0.5 + + k_codes, v_codes, k_scale, v_scale = _fp8_cache(k_rows, v_rows) + actual = decode_paged_attention( + q, + k_codes, + v_codes, + indptr, + indices, + q_positions, + attn_logits, + attn_lse, + num_kv_splits, + max_kv_splits, + sm_scale, + k_scale=k_scale, + v_scale=v_scale, + ) + expected = _reference_paged_attention( + q, + _dequantized(k_codes, k_scale), + _dequantized(v_codes, v_scale), + indptr, + indices, + q_to_req, + q_positions, + sm_scale, + None, + ) + # The kernel rounds the decoded operands to the compute dtype before tl.dot; the + # reference stays in fp32, hence the same tolerance the bf16 tests already use. + torch.testing.assert_close(actual.float(), expected.float(), atol=2e-2, rtol=2e-2) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Triton attention needs CUDA") +def test_paged_attention_decodes_fp8_scales(): + """The non-tl.dot fallback kernel takes the same scales (it serves head_dim > 256 + with a short prefill, and every decode batch when the grouped path is not used).""" + from freetoken.kernel.triton.attention import paged_attention + + torch.manual_seed(13) + device = torch.device("cuda") + head_dim, num_kv_heads, num_q_heads = 64, 2, 4 + seq_lens = [5, 3] + total_kv = sum(seq_lens) + q = torch.randn(total_kv, num_q_heads, head_dim, device=device, dtype=torch.bfloat16) + k_rows = (torch.randn(total_kv, num_kv_heads, head_dim, device=device) * 16.0).to( + torch.bfloat16 + ) + v_rows = (torch.randn(total_kv, num_kv_heads, head_dim, device=device) * 0.02).to( + torch.bfloat16 + ) + indptr = torch.tensor([0, seq_lens[0], total_kv], dtype=torch.int32, device=device) + indices = torch.arange(total_kv, dtype=torch.int32, device=device) + q_to_req = torch.tensor( + [0] * seq_lens[0] + [1] * seq_lens[1], dtype=torch.int32, device=device + ) + q_positions = torch.cat( + [ + torch.arange(seq_lens[0], dtype=torch.int64, device=device), + torch.arange(seq_lens[1], dtype=torch.int64, device=device), + ] + ) + k_codes, v_codes, k_scale, v_scale = _fp8_cache(k_rows, v_rows) + actual = paged_attention( + q, + k_codes, + v_codes, + indptr, + indices, + q_to_req, + q_positions, + head_dim**-0.5, + k_scale=k_scale, + v_scale=v_scale, + ) + expected = _reference_paged_attention( + q, + _dequantized(k_codes, k_scale), + _dequantized(v_codes, v_scale), + indptr, + indices, + q_to_req, + q_positions, + head_dim**-0.5, + None, + ) + torch.testing.assert_close(actual.float(), expected.float(), atol=2e-2, rtol=2e-2) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Triton attention needs CUDA") +@pytest.mark.parametrize("use_split_inputs", [False, True]) +def test_extend_paged_attention_decodes_fp8_scales(use_split_inputs: bool): + """Prefill over an fp8 cache. + + With ``use_split_inputs`` the kernel reads a request's own new tokens from + ``k_extend`` (compute dtype, never quantized) and only the cached prefix from the + fp8 codes; without it every row is served from the codes. The reference mirrors + that split, so a scale leaking onto the extend path -- or failing to apply to the + cache path -- cannot pass. + """ + from freetoken.kernel.triton.attention import extend_paged_attention + + torch.manual_seed(14) + device = torch.device("cuda") + head_dim, num_kv_heads, num_q_heads = 64, 2, 8 + cached_lens, extend_lens = [4, 2], [3, 2] + seq_lens = [c + e for c, e in zip(cached_lens, extend_lens)] + total_q, total_kv = sum(extend_lens), sum(seq_lens) + q = torch.randn(total_q, num_q_heads, head_dim, device=device, dtype=torch.bfloat16) + k_extend = (torch.randn(total_q, num_kv_heads, head_dim, device=device) * 3.0).to( + torch.bfloat16 + ) + v_extend = (torch.randn(total_q, num_kv_heads, head_dim, device=device) * 3.0).to( + torch.bfloat16 + ) + # Magnitudes far from 1: a dropped scale is then off by orders of magnitude. + k_cache = (torch.randn(total_kv, num_kv_heads, head_dim, device=device) * 0.3).to( + torch.bfloat16 + ) + v_cache = (torch.randn(total_kv, num_kv_heads, head_dim, device=device) * 3.0).to( + torch.bfloat16 + ) + qo_indptr = torch.tensor([0] + extend_lens, dtype=torch.int32, device=device).cumsum_(0) + kv_indptr = torch.tensor([0] + seq_lens, dtype=torch.int32, device=device).cumsum_(0) + # KV positions are logical, but FP8 codes and their scales are addressed by the + # physical slots from the page table. A contiguous table masks a regression that + # looks scales up with the logical position instead of the slot. + indices = torch.tensor( + [10, 3, 8, 1, 9, 0, 7, 2, 6, 4, 5], dtype=torch.int32, device=device + ) + prefix_lens = torch.tensor(cached_lens, dtype=torch.int32, device=device) + q_to_req = torch.empty(total_q, dtype=torch.int32, device=device) + q_positions = torch.empty(total_q, dtype=torch.int64, device=device) + q_off = kv_off = 0 + for req_idx, (cached_len, extend_len) in enumerate(zip(cached_lens, extend_lens)): + q_to_req[q_off : q_off + extend_len].fill_(req_idx) + q_positions[q_off : q_off + extend_len] = torch.arange( + cached_len, cached_len + extend_len, dtype=torch.int64, device=device + ) + # The step's own tokens are what the engine stores at the tail of the span. + k_cache[kv_off + cached_len : kv_off + cached_len + extend_len] = k_extend[ + q_off : q_off + extend_len + ] + v_cache[kv_off + cached_len : kv_off + cached_len + extend_len] = v_extend[ + q_off : q_off + extend_len + ] + q_off += extend_len + kv_off += cached_len + extend_len + sm_scale = head_dim**-0.5 + + num_slots = total_kv + 1 + k_slots = torch.zeros( + num_slots, num_kv_heads, head_dim, dtype=k_cache.dtype, device=device + ) + v_slots = torch.zeros_like(k_slots) + k_slots[indices.to(torch.long)] = k_cache + v_slots[indices.to(torch.long)] = v_cache + k_codes, v_codes, k_scale, v_scale = _fp8_cache(k_slots, v_slots) + k_ref = _dequantized(k_codes, k_scale).clone() + v_ref = _dequantized(v_codes, v_scale).clone() + if use_split_inputs: + q_off = kv_off = 0 + for cached_len, extend_len in zip(cached_lens, extend_lens): + current_slots = indices[ + kv_off + cached_len : kv_off + cached_len + extend_len + ].to(torch.long) + k_ref[current_slots] = k_extend[q_off : q_off + extend_len].float() + v_ref[current_slots] = v_extend[q_off : q_off + extend_len].float() + q_off += extend_len + kv_off += cached_len + extend_len + + actual = extend_paged_attention( + q=q, + k_cache=k_codes, + v_cache=v_codes, + qo_indptr=qo_indptr, + kv_indptr=kv_indptr, + kv_indices=indices, + prefix_lens=prefix_lens, + max_q_len=max(extend_lens), + sm_scale=sm_scale, + k_extend=k_extend if use_split_inputs else None, + v_extend=v_extend if use_split_inputs else None, + k_scale=k_scale, + v_scale=v_scale, + ) + expected = _reference_paged_attention( + q, + k_ref, + v_ref, + kv_indptr, + indices, + q_to_req, + q_positions, + sm_scale, + None, + ) + torch.testing.assert_close(actual.float(), expected.float(), atol=3e-2, rtol=3e-2) diff --git a/tests/kvcache/test_dsa_pool.py b/tests/kvcache/test_dsa_pool.py index 7f4b1c7091..8bcdb8d5e9 100644 --- a/tests/kvcache/test_dsa_pool.py +++ b/tests/kvcache/test_dsa_pool.py @@ -16,13 +16,14 @@ LATENT, IDX_DIM = 80, 32 # latent 64+16: kernel spans need pow-2 dims (512+64 in the real model) -def _pool(num_pages: int): +def _pool(num_pages: int, kv_quant: str = "none"): from freetoken.kvcache.dsa_pool import DSAKVCache return DSAKVCache( latent_dim=LATENT, num_layers=2, num_pages=num_pages, page_size=1, dtype=torch.bfloat16, device=torch.device("cuda"), index_head_dim=IDX_DIM, num_index_layers=1, + kv_quant=kv_quant, ) @@ -91,6 +92,39 @@ def test_sparse_decode_reads_grown_pool_through_backend_kernels(): assert (o[0, 0].float() - ref_o).abs().max().item() < 2e-2 +def test_fp8_latent_rows_decode_with_their_scales(): + """The DSA kernel must read codes and scales from matching physical rows.""" + from freetoken.kernel.triton.glm_dsa_sparse import glm_dsa_sparse_attn + from freetoken.kernel.triton.kv_quant import codes_to_f32 + + torch.manual_seed(7) + pool = _pool(128, kv_quant="fp8") + rows = torch.randperm(128, device="cuda")[:96].to(torch.int32) + c_kv = torch.randn(96, LATENT - 16, device="cuda", dtype=torch.bfloat16) + k_rope = torch.randn(96, 16, device="cuda", dtype=torch.bfloat16) + pool.store_kv(c_kv, k_rope, rows, layer_id=0) + assert pool.latent_rows(0).dtype is torch.uint8 + assert pool.latent_scale(0).dtype is torch.float32 + + q = torch.randn(1, 1, 4, LATENT, device="cuda", dtype=torch.bfloat16) + sel = rows[:64].view(1, 1, -1) + cnt = torch.tensor([[64]], device="cuda", dtype=torch.int32) + out = glm_dsa_sparse_attn( + q, pool.latent_rows(0), sel, 0.1, counts=cnt, d_v=LATENT - 16, + pool_scale=pool.latent_scale(0), + ) + out_split = glm_dsa_sparse_attn( + q, pool.latent_rows(0), sel, 0.1, counts=cnt, d_v=LATENT - 16, + pool_scale=pool.latent_scale(0), force_splits=4, + ) + decoded = codes_to_f32(pool.latent_rows(0)) * pool.latent_scale(0).unsqueeze(-1) + picked = decoded[sel.view(-1).long()] + logits = (q[0, 0].float() @ picked.T) * 0.1 + ref = logits.softmax(-1) @ picked[:, : LATENT - 16] + assert (out[0, 0].float() - ref).abs().max().item() < 2e-2 + assert (out_split.float() - ref).abs().max().item() < 2e-2 + + def test_mla_pool_selected_by_group_spec(): """The factory keys MLA/DSA pools off the attention-group spec, never the model payload -- and zeroed index dims (the dense ablation) fall back to MLAKVCache.""" @@ -120,6 +154,10 @@ def cfg(index_dim, n_idx): dsa = create_kvcache_pool(model_config=cfg(IDX_DIM, 1), num_pages=8, page_size=1, device=torch.device("cuda"), dtype=torch.bfloat16) assert isinstance(dsa, DSAKVCache) + fp8_dsa = create_kvcache_pool(model_config=cfg(IDX_DIM, 1), num_pages=8, page_size=1, + device=torch.device("cuda"), dtype=torch.bfloat16, + kv_quant="fp8") + assert isinstance(fp8_dsa, DSAKVCache) and fp8_dsa.store_dtype is torch.uint8 mla = create_kvcache_pool(model_config=cfg(0, 0), num_pages=8, page_size=1, device=torch.device("cuda"), dtype=torch.bfloat16) assert isinstance(mla, MLAKVCache) and not isinstance(mla, DSAKVCache) @@ -136,3 +174,57 @@ def test_rebuild_shrink_and_engine_wiring(): pool.rebuild_from_config(config=None, num_pages=63) assert pool.latent_rows(0).shape[0] == 64 # 63 + 1 dummy page assert pool.index_k_cache(0).shape[0] == 64 + + +@pytest.mark.parametrize("kv_quant", ["none", "fp8", "nvfp4"]) +@pytest.mark.parametrize("ratio", [1, 4]) +def test_latent_budget_matches_allocations_and_rebuild(kv_quant, ratio): + from types import SimpleNamespace + from freetoken.attention import AttnType + from freetoken.kvcache.dsa_pool import DSAKVCache, KpoolDSAKVCache + from freetoken.models.config import KVCacheGroupSpec + + cls = KpoolDSAKVCache if ratio > 1 else DSAKVCache + spec = KVCacheGroupSpec( + name="full", layer_ids=(3, 7), num_kv_heads=1, head_dim=512, sliding_window=None, + mla=True, index_head_dim=128, num_index_layers=2, index_ratio=ratio, + attn_type=AttnType.DSA, + ) + cfg = SimpleNamespace( + kv_quant=kv_quant, dtype=torch.bfloat16, tp_info=SimpleNamespace(size=1), + max_running_req=3, page_size=64, + model_config=SimpleNamespace(kv_cache_group_specs=lambda: (spec,)), + ) + extra = dict(index_ratio=ratio, num_req_slots=4) if ratio > 1 else {} + pool = cls(512, 8, 2, 64, torch.bfloat16, torch.device("cuda"), + index_head_dim=128, num_index_layers=2, layer_ids=(3, 7), + kv_quant=kv_quant, **extra) + page_bytes, fixed, page_size, _ = cls.kv_cost(cfg) + for pages in (2, 5, 1): + pool.rebuild_from_config(cfg, pages) + allocated_pages = pages + 1 + buffers = (pool._kv_buffer, pool._scale_buffer, pool._block_scale_buffer, + pool._index_k_buffer) + if ratio > 1: + buffers += (pool._tail_k, pool._tail_gate) + assert pool.cmp_scratch_base == allocated_pages * 64 // ratio + assert pool.tail_k(0).dtype == torch.bfloat16 + actual = sum(b.numel() * b.element_size() for b in buffers if b is not None) + assert actual == allocated_pages * page_bytes + fixed + assert pool.unit_bytes() == (page_bytes // page_size, 0) + assert pool.latent_rows(7).shape == (allocated_pages * 64, 256 if kv_quant == "nvfp4" else 512) + loc = torch.tensor([allocated_pages * 64 - 1], device="cuda") + x = torch.ones(1, 512, device="cuda", dtype=torch.bfloat16) + pool.store_kv(x, x[:, :0], loc, 7) + if kv_quant == "nvfp4": + from tests.kernels.test_kv_nvfp4 import _decode_latent + + torch.testing.assert_close(_decode_latent(pool, 7)[loc], x.float()) + assert pool.k_cache(7).data_ptr() == pool.v_cache(7).data_ptr() + + +def test_nvfp4_latent_rejects_partial_blocks(): + from freetoken.kvcache.dsa_pool import MLAKVCache + + with pytest.raises(ValueError, match="divisible by 16"): + MLAKVCache(72, 1, 1, 1, torch.bfloat16, torch.device("cpu"), kv_quant="nvfp4") diff --git a/tests/kvcache/test_host_prefix_tier.py b/tests/kvcache/test_host_prefix_tier.py new file mode 100644 index 0000000000..fdb552e8e0 --- /dev/null +++ b/tests/kvcache/test_host_prefix_tier.py @@ -0,0 +1,60 @@ +"""Degrau 1 — teste standalone do HostPrefixTier (roda sem servidor; GPU tiny).""" +import torch +from types import SimpleNamespace + +from freetoken.kvcache.host_prefix_tier import HostPrefixTier + +def main(): + import os + dev = "cuda" if (os.environ.get("TIER_TEST_CPU") is None and torch.cuda.is_available()) else "cpu" + # pool fake: mesmas geometrias, mini + L, S = 4, 3 # 4 "layers", 3 slots GPU + pool = SimpleNamespace( + conv_states=torch.randn(L, S, 8, 3, dtype=torch.bfloat16, device=dev), + recurrent_states=torch.randn(L, S, 2, 4, 4, dtype=torch.float32, device=dev), + slot_states={"extra": torch.randn(L, S, 5, dtype=torch.bfloat16, device=dev)}, + ) + tier = HostPrefixTier(pool, gdn_slots_host=4, device=dev) + + # 1) round-trip D2H->H2D preserva bytes + original = [t.clone() for t in (pool.conv_states, pool.recurrent_states, pool.slot_states["extra"])] + hs = tier.write_gdn(1) # snapshot do slot 1 -> host + # suja o slot GPU pra provar que o restore vem do host + for t in (pool.conv_states, pool.recurrent_states, pool.slot_states["extra"]): + t[:, 1].zero_() + tier.read_gdn(hs, 2) # host -> slot 2 + ok = all(torch.equal(t[:, 1].cpu(), u[:, 2].cpu()) + for t, u in zip(original, (pool.conv_states, pool.recurrent_states, pool.slot_states["extra"]))) + assert ok, "snapshot divergiu no round-trip!" + print("PASS 1: snapshot GDN round-trip D2H->H2D bit-exato") + + # 2) índice: put + lookup longest-prefix page-aligned + toks = torch.arange(64 * 5, dtype=torch.int64) # 5 páginas + pages = torch.tensor([10, 11, 12, 13, 14], dtype=torch.int32) + tier.put_prefix(toks[: 64 * 4], pages[:4], hs) # registro 4 páginas + m = tier.lookup(torch.arange(64 * 5, dtype=torch.int64)) # query 5 páginas + assert m is not None and m.host_len == 64 * 4 and m.gdn_slot == hs + assert torch.equal(m.kv_pages, pages[:4]) + print("PASS 2: lookup acha o maior prefixo page-aligned") + + # 3) miss quando nada bate + assert tier.lookup(torch.full((64 * 3,), 999999, dtype=torch.int64)) is None + print("PASS 3: lookup retorna None em miss total") + + # 4) evict_host libera o slot GDN host + tier.evict_host(toks[: 64 * 4]) + assert hs in tier._free + print("PASS 4: evict_host libera slot host") + + # 5) proteção: tier novo com 2 slots, 3ª escrita levanta erro claro + tier2 = HostPrefixTier(pool, gdn_slots_host=2, device=dev) + tier2.write_gdn(0); tier2.write_gdn(0) + try: + tier2.write_gdn(0) + raise SystemExit("ERA PRA TER FALHADO") + except RuntimeError: + print("PASS 5: capacidade cheia levanta RuntimeError claro") + + print("\nTODOS OS TESTES PASSARAM ✓") + +main() diff --git a/tests/kvcache/test_host_tier_live_pool.py b/tests/kvcache/test_host_tier_live_pool.py new file mode 100644 index 0000000000..0eea1990e3 --- /dev/null +++ b/tests/kvcache/test_host_tier_live_pool.py @@ -0,0 +1,72 @@ +"""The GDN snapshot tier (upstream #525) must copy against the pool's CURRENT tensors, which +LinearStatePool.rebuild replaces, and must order its copies after the engine stream.""" + +import contextlib +from types import SimpleNamespace + +import torch + +import freetoken.kvcache.host_prefix_tier as tier_module +from freetoken.kvcache.host_prefix_tier import HostPrefixTier + + +def _pool(slots: int): + return SimpleNamespace( + conv_states=torch.zeros(2, slots, 4, 3), + recurrent_states=torch.zeros(2, slots, 2, 4), + slot_states={}, + ) + + +def test_read_after_a_pool_rebuild_lands_in_the_new_tensors(): + pool = _pool(3) + tier = HostPrefixTier(pool, gdn_slots_host=2, device="cpu") + pool.conv_states[:, 1] = 5.0 + pool.recurrent_states[:, 1] = 7.0 + host = tier.write_gdn(1) + # a runtime rebuild swaps in fresh (larger) state tensors + pool.conv_states = torch.zeros(2, 6, 4, 3) + pool.recurrent_states = torch.zeros(2, 6, 2, 4) + tier.read_gdn(host, 5) + assert torch.all(pool.conv_states[:, 5] == 5.0) + assert torch.all(pool.recurrent_states[:, 5] == 7.0) + + +def _fake_copy_stream(monkeypatch): + events = [] + engine = object() + + class FakeStream: + def wait_stream(self, other): + events.append(("wait", other)) + + def synchronize(self): + events.append(("sync",)) + + monkeypatch.setattr(tier_module.torch.cuda, "current_stream", lambda device=None: engine) + monkeypatch.setattr(tier_module.torch.cuda, "stream", lambda s: contextlib.nullcontext()) + return FakeStream(), events, engine + + +def test_copies_wait_for_the_engine_stream(monkeypatch): + stream, events, engine = _fake_copy_stream(monkeypatch) + pool = _pool(3) + tier = HostPrefixTier(pool, gdn_slots_host=2, device="cpu") + tier.copy_stream = stream + host = tier.write_gdn(1) + tier.read_gdn(host, 2) + assert events == [("wait", engine), ("sync",), ("wait", engine), ("sync",)] + + +def test_kv_page_copies_wait_for_the_engine_stream(monkeypatch): + stream, events, engine = _fake_copy_stream(monkeypatch) + kv_pool = SimpleNamespace( + _kv_buffer=torch.zeros(2, 1, 4, 4, 1, 2), # [2, L, P, page, H, D] + _scale_buffer=torch.zeros(2, 1, 16, 1), # [2, L, P*page, H] + _block_scale_buffer=None, + ) + tier = HostPrefixTier(_pool(3), kv_pool, gdn_slots_host=2, kv_budget_bytes=1 << 12, device="cpu") + tier.copy_stream = stream + slots = tier.write_kv(torch.arange(4, 8, dtype=torch.int32)) + tier.read_kv(slots, torch.tensor([2])) + assert events == [("wait", engine), ("sync",), ("wait", engine), ("sync",)] diff --git a/tests/kvcache/test_host_tier_nvfp4_guard.py b/tests/kvcache/test_host_tier_nvfp4_guard.py new file mode 100644 index 0000000000..50227267a6 --- /dev/null +++ b/tests/kvcache/test_host_tier_nvfp4_guard.py @@ -0,0 +1,51 @@ +"""The host KV tiers (upstream #499 / #525) mirror codes plus the fp8 pool's per-token scales; +an nvfp4 pool keeps block scales they do not copy, so both must refuse it up front.""" + +from types import SimpleNamespace + +import pytest +import torch + +from freetoken.kvcache.host_prefix_tier import HostPrefixTier +from freetoken.kvcache.kv_host_offload import KVHostOffloader + + +def _kv_pool(block_scaled: bool): + return SimpleNamespace( + _kv_buffer=torch.zeros(2, 1, 3, 4, 1, 8, dtype=torch.uint8), # [2, L, P+1, page, H, D] + _scale_buffer=None, + _block_scale_buffer=torch.zeros(1) if block_scaled else None, + ) + + +def _state_pool(): + return SimpleNamespace( + conv_states=torch.zeros(2, 2, 4, 3), + recurrent_states=torch.zeros(2, 2, 2, 4), + slot_states={}, + ) + + +def test_kv_host_offloader_refuses_an_nvfp4_pool(): + with pytest.raises(ValueError, match="--kv-host-pages .* nvfp4"): + KVHostOffloader(_kv_pool(block_scaled=True), num_logical_pages=4, device=torch.device("cpu")) + + +def test_prefix_tier_refuses_kv_pages_on_an_nvfp4_pool(): + with pytest.raises(ValueError, match="FT_PREFIX_HOST .* nvfp4"): + HostPrefixTier(_state_pool(), kv_pool=_kv_pool(block_scaled=True), gdn_slots_host=2, + kv_budget_bytes=1 << 16, device="cpu") + + +def test_prefix_tier_gdn_snapshots_do_not_touch_the_kv_pool(): + # the snapshot-only tier (no KV budget) never reads the pool, nvfp4 or not + tier = HostPrefixTier(_state_pool(), kv_pool=_kv_pool(block_scaled=True), gdn_slots_host=2, + kv_budget_bytes=0, device="cpu") + assert tier.num_slots == 2 and not tier.has_kv + + +def test_prefix_tier_refuses_kv_pages_on_a_qsa_pool(): + # the KV-page store mirrors codes and scales, not the QSA indexer's compressed-key slab + pool = SimpleNamespace(**vars(_kv_pool(block_scaled=False)), _cmp_k_buffer=torch.zeros(1)) + with pytest.raises(ValueError, match="FT_PREFIX_HOST .* index slab"): + HostPrefixTier(_state_pool(), kv_pool=pool, gdn_slots_host=2, kv_budget_bytes=1 << 16, device="cpu") diff --git a/tests/kvcache/test_host_tier_rehydrate_exact.py b/tests/kvcache/test_host_tier_rehydrate_exact.py new file mode 100644 index 0000000000..77dae0fd08 --- /dev/null +++ b/tests/kvcache/test_host_tier_rehydrate_exact.py @@ -0,0 +1,125 @@ +"""The GDN snapshot tier (upstream #525) may only revive a tombstone from a snapshot of +exactly that node's prefix: a shorter snapshot would resume the recurrence with the tokens +in between silently missing. Making room for the revival evicts other prefixes into the +tier, which may recycle the entry's own host slots, so it must not read stale ones.""" + +from types import SimpleNamespace + +import torch + +import freetoken.core as core +from freetoken.kvcache.host_prefix_tier import HostPrefixTier +from freetoken.kvcache.linear_state_pool import LinearStatePool +from freetoken.models.config import LinearGatedDeltaGroupConfig +from freetoken.scheduler.cache import CacheManager + + +def _manager(key_len: int = 128): + state = SimpleNamespace(conv_states=torch.zeros(1, 4, 1, 1), + recurrent_states=torch.zeros(1, 4, 1, 1), slot_states={}) + tier = HostPrefixTier(state, gdn_slots_host=4, device="cpu") + pool = SimpleNamespace(num_free_slots=4, alloc=lambda n: [3]) + manager = SimpleNamespace( + host_tier=tier, + prefix_cache=SimpleNamespace(_collect_key=lambda node: torch.arange(key_len)), + page_size=64, + linear_state_pool=pool, + ensure_mamba_slots=lambda n: None, + ) + return manager, tier, state + + +def _snapshot(tier, state, key, value): + state.conv_states[:, 1] = value + tier.put_prefix(key, None, tier.write_gdn(1)) + + +def test_a_shorter_host_snapshot_neither_revives_nor_refreshes_a_deeper_node(): + manager, tier, state = _manager(key_len=128) + _snapshot(tier, state, torch.arange(64), 1.0) + _snapshot(tier, state, torch.arange(1000, 1064), 2.0) + assert CacheManager._rehydrate_snapshot(manager, node=object()) is None + # the rejected entry keeps its LRU age + assert [len(e.key_tokens) for e in tier._index.values()] == [64, 64] + assert next(iter(tier._index.values())).key_tokens[0] == 0 + assert torch.all(state.conv_states[:, 3] == 0) + + +def test_an_exact_host_snapshot_revives_the_node(): + manager, tier, state = _manager(key_len=128) + _snapshot(tier, state, torch.arange(128), 5.0) + assert CacheManager._rehydrate_snapshot(manager, node=object()) == 3 + assert torch.all(state.conv_states[:, 3] == 5.0) + + +def _cache_manager(monkeypatch, num_pages, page_size, state_slots, **env): + for k, v in env.items(): + monkeypatch.setenv(k, v) + g = LinearGatedDeltaGroupConfig( + name="linear", layer_ids=(0,), num_key_heads=2, num_value_heads=4, + key_head_dim=16, value_head_dim=16, conv_kernel_dim=4, output_gate="silu", + ) + pool = LinearStatePool(group=g, num_slots=state_slots, dtype=torch.bfloat16, + device=torch.device("cpu"), tp_size=1) + page_table = torch.zeros(2, 64, dtype=torch.int32) + cm = CacheManager(num_pages, page_size, page_table, "hybrid_radix", linear_state_pool=pool) + return cm, pool + + +def _put(cm, pool, ids, state, stamp, kv_pool=None, kv_value=None): + kv = cm._page_to_token(cm._allocate(len(ids) // cm.page_size)) + if kv_pool is not None: + kv_pool._kv_buffer[:, :, kv[:: cm.page_size] // cm.page_size] = kv_value + slot = pool.alloc(1)[0] + pool.conv_states[:, slot] = state + cm.prefix_cache.insert(ids, kv, slot) + cm.prefix_cache.match_prefix(ids).node.timestamp = stamp + + +def _match(cm, ids): + tail = torch.tensor([7, 9], dtype=torch.int32) + return cm.match_req(SimpleNamespace(input_ids=torch.cat([ids, tail]), input_len=len(ids) + 2)) + + +def test_making_room_to_revive_does_not_restore_the_victim_into_the_host_slot(monkeypatch): + # one host slot: evicting B for the revived slot drops A's entry and writes B there + cm, pool = _cache_manager(monkeypatch, 8, 64, 3, FT_GDN_HOST_TIER="1", FT_GDN_HOST_SLOTS="1") + a = torch.arange(128, dtype=torch.int32) + b = torch.arange(1000, 1128, dtype=torch.int32) + _put(cm, pool, a, 1.0, stamp=1) + _put(cm, pool, b, 2.0, stamp=2) + cm.ensure_mamba_slots(1) # A -> host tier, A's node becomes a tombstone + pool.conv_states[:, pool.alloc(1)[0]] = 9.0 # someone else takes the freed slot + + mr = _match(cm, a) + restored = None if mr.mamba_value is None else pool.conv_states[:, mr.mamba_value].unique().tolist() + assert restored in (None, [1.0]), (mr.cuda_handle.cached_len, restored) + assert cm.host_tier.lookup_exact(b) is not None + + +def test_making_room_to_rehydrate_a_host_prefix_does_not_restore_the_victim(monkeypatch): + # a 4-page KV store: writing the 4-page victim V drops the 2-page hit H and reuses its slots + page, n_pages = 4, 6 + kv_pool = SimpleNamespace(_kv_buffer=torch.zeros(2, 1, n_pages, page, 1, 2), + _scale_buffer=None, _block_scale_buffer=None) + monkeypatch.setattr(core, "_GLOBAL_CTX", SimpleNamespace(kv_cache=kv_pool, kv_offloader=None)) + page_bytes = kv_pool._kv_buffer[:, :, 0].numel() * kv_pool._kv_buffer.element_size() + cm, pool = _cache_manager(monkeypatch, n_pages, page, 4, FT_PREFIX_HOST="1", + FT_PREFIX_HOST_GB=repr(4 * page_bytes / (1 << 30)), + FT_GDN_HOST_SLOTS="4") + h = torch.arange(2 * page, dtype=torch.int32) + v = torch.arange(100, 100 + 4 * page, dtype=torch.int32) + _put(cm, pool, h, 1.0, stamp=1, kv_pool=kv_pool, kv_value=1.0) + er = cm.prefix_cache.evict_full(len(h)) # H -> host tier (KV + snapshot), out of the tree + pool.free(er.mamba_slots) + cm._free(er.kv_indices) + _put(cm, pool, v, 2.0, stamp=2, kv_pool=kv_pool, kv_value=2.0) + cm._allocate(len(cm.free_slots)) # a running request holds the rest of the GPU pages + + mr = _match(cm, h) + handle = mr.cuda_handle + if handle.cached_len: + pages = handle.kv_indices[::page] // page + assert pool.conv_states[:, mr.mamba_value].unique().tolist() == [1.0] + assert kv_pool._kv_buffer[:, :, pages].unique().tolist() == [1.0] + assert cm.host_tier.lookup_exact(v) is not None diff --git a/tests/kvcache/test_kv_cache_rebuild.py b/tests/kvcache/test_kv_cache_rebuild.py index 0e3b7050e8..d80782fef4 100644 --- a/tests/kvcache/test_kv_cache_rebuild.py +++ b/tests/kvcache/test_kv_cache_rebuild.py @@ -64,6 +64,13 @@ def test_mla_and_dsa_rebuild_from_config_and_unit_bytes(): # the index slab's per-token bytes ride on top of the latent slab's, each floored on its own assert dsa.unit_bytes() == (layers * latent * 2 + n_idx * idx_dim * 2, 0) + fp8 = MLAKVCache(latent_dim=latent, num_layers=layers, num_pages=8, page_size=1, + dtype=torch.bfloat16, device=torch.device("cpu"), kv_quant="fp8") + fp8.rebuild_from_config(config=None, num_pages=20) + assert fp8.latent_rows(0).dtype is torch.uint8 + # One byte per latent element plus one FP32 scale for each latent layer. + assert fp8.unit_bytes() == (layers * latent + layers * 4, 0) + def _hybrid_groups(): from freetoken.models.config import KVCacheGroupSpec @@ -269,9 +276,15 @@ def test_every_kv_pool_answers_the_sizing_surface(): from freetoken.kvcache.mha_pool import MHAKVCache for cls in (MHAKVCache, MLAKVCache, DSAKVCache, HybridSWAKVCache, DSV4PagedKVCache): - for hook in ("kv_cost", "solve_num_pages", "min_kv_tokens", "validate_rebuild"): + for hook in ( + "kv_cost", "solve_num_pages", "min_kv_tokens", "validate_rebuild", + "rebuild_footprint", + ): assert callable(getattr(cls, hook, None)), f"{cls.__name__} is missing {hook}" - for hook in ("kv_cost", "solve_num_pages", "min_kv_tokens", "validate_rebuild"): + for hook in ( + "kv_cost", "solve_num_pages", "min_kv_tokens", "validate_rebuild", + "rebuild_footprint", + ): assert hook in DSV4PagedKVCache.__dict__, f"DSV4 lost its {hook} override" diff --git a/tests/kvcache/test_linear_state_pool_alloc.py b/tests/kvcache/test_linear_state_pool_alloc.py index e9fe405426..e0edd0438e 100644 --- a/tests/kvcache/test_linear_state_pool_alloc.py +++ b/tests/kvcache/test_linear_state_pool_alloc.py @@ -9,6 +9,7 @@ from freetoken.kvcache.linear_state_pool import ( LinearStatePool, + gdn_prefill_workspace_bytes, linear_state_bytes_per_req, state_pool_bytes, ) @@ -189,3 +190,44 @@ def test_slot_state_bytes_for_the_real_geometry(): delta = linear_state_bytes_per_req(group, 1, torch.bfloat16, (spec,)) - \ linear_state_bytes_per_req(group, 1, torch.bfloat16) assert delta == 4 * 2560 * 9 * 2 == 180 * 1024 + + +def _gdn_config(**fields): + # qwen4_exp GDN geometry: 48 value heads of 128x128 + group = LinearGatedDeltaGroupConfig( + name="linear", layer_ids=(0, 1), + num_key_heads=16, num_value_heads=48, + key_head_dim=128, value_head_dim=128, conv_kernel_dim=4, output_gate="silu", + ) + mc = SimpleNamespace(linear_attention_group=lambda: group) + return SimpleNamespace(model_config=mc, dtype=torch.bfloat16, max_seq_len=262144, **fields) + + +def test_gdn_prefill_workspace_for_the_real_geometry(): + # h [8192/64 chunks, 48, 128, 128] bf16 = 192 MiB, plus v_new [8192, 48, 128] = 96 MiB + assert gdn_prefill_workspace_bytes(_gdn_config(max_extend_tokens=8192)) == 288 * 2**20 + assert gdn_prefill_workspace_bytes(_gdn_config(max_extend_tokens=2048)) == 72 * 2**20 + + +def test_gdn_prefill_workspace_without_a_scheduler_extend_budget(): + # a bare EngineConfig has no max_extend_tokens: price the scheduler's 8192 default + assert gdn_prefill_workspace_bytes(_gdn_config()) == 288 * 2**20 + + +def test_gdn_prefill_workspace_is_capped_by_the_sequence_length(): + config = _gdn_config(max_extend_tokens=8192) + config.max_seq_len = 100 # 2 chunks, 100 value rows + assert gdn_prefill_workspace_bytes(config) == (2 * 48 * 128 * 128 + 100 * 48 * 128) * 2 + + +def test_gdn_prefill_workspace_is_zero_without_a_linear_group(): + config = _gdn_config(max_extend_tokens=8192) + config.model_config.linear_attention_group = lambda: None + assert gdn_prefill_workspace_bytes(config) == 0 + + +def test_gdn_prefill_workspace_chunk_matches_the_kernel(): + from freetoken.kernel.fla.chunk_delta_h import CHUNK_SIZE + from freetoken.kvcache.linear_state_pool import _GDN_CHUNK_SIZE + + assert _GDN_CHUNK_SIZE == CHUNK_SIZE diff --git a/tests/kvcache/test_mha_pool_fp8.py b/tests/kvcache/test_mha_pool_fp8.py new file mode 100644 index 0000000000..8d74b0edd3 --- /dev/null +++ b/tests/kvcache/test_mha_pool_fp8.py @@ -0,0 +1,270 @@ +"""The fp8 KV pool: code buffer + per-(token, head) scales, sized and rebuilt together. + +The pool-side half of --kv-cache-dtype fp8. The interesting failures are the silent +ones: a scale buffer that misses the layer_ids remap, a rebuild that resizes the codes +but not the scales, or a ``unit_bytes`` that drifts from ``kv_cost`` (which is what the +VRAM budget, the cache sliders and the runtime rebuild all divide by). +""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch + +from freetoken.distributed import set_tp_info, try_get_tp_info +from freetoken.models.config import KVCacheGroupSpec + +if not torch.cuda.is_available(): # pragma: no cover + pytest.skip("CUDA required", allow_module_level=True) + +from freetoken.kernel.triton.kv_quant import kv_codes_dtype + +DEV = torch.device("cuda") +HEADS, DIM, LAYERS, PAGES, PAGE_SIZE = 4, 64, 3, 6, 8 + + +def _init_tp() -> None: + if try_get_tp_info() is None: + set_tp_info(rank=0, size=1) + + +@pytest.fixture(autouse=True) +def _tp(): + _init_tp() + + +def _pool(kv_quant="fp8", num_pages=PAGES, layer_ids=None, num_layers=LAYERS): + from freetoken.kvcache.mha_pool import MHAKVCache + + _init_tp() + return MHAKVCache( + num_kv_heads=HEADS, + num_layers=num_layers, + head_dim=DIM, + num_pages=num_pages, + page_size=PAGE_SIZE, + dtype=torch.bfloat16, + device=DEV, + layer_ids=layer_ids, + kv_quant=kv_quant, + ) + + +def test_kv_store_dtype_selection(): + from freetoken.kvcache.mha_pool import _kv_store_dtype + + assert _kv_store_dtype(torch.bfloat16, "none") is torch.bfloat16 + assert _kv_store_dtype(torch.bfloat16, "fp8") is kv_codes_dtype() + with pytest.raises(ValueError, match="kv_quant"): + _kv_store_dtype(torch.bfloat16, "q6") + + +def test_fp8_pool_keeps_geometry_and_adds_scale_views(): + pool = _pool() + assert pool.kv_quant == "fp8" + # Same shape as the 16-bit pool -- only the element type changed. + assert pool._kv_buffer.shape == (2, LAYERS, PAGES, PAGE_SIZE, HEADS, DIM) + assert pool._kv_buffer.dtype == kv_codes_dtype() + assert pool.store_dtype == kv_codes_dtype() + # ``dtype`` stays the COMPUTE dtype (kvcache/base.py). Backends size their scratch + # with it -- reporting codes here hands e4m3 to a 16-bit tl.dot, which does not fail + # until CUDA-graph capture (exactly how QSA's indexer died in the field). + assert pool.dtype == torch.bfloat16 + slots = PAGES * PAGE_SIZE + for layer in range(LAYERS): + assert pool.k_scale(layer).shape == (slots, HEADS) + assert pool.v_scale(layer).shape == (slots, HEADS) + assert pool.k_scale(layer).dtype == torch.float32 + # A 16-bit pool exposes no scales at all. + assert _pool(kv_quant="none").k_scale(0) is None + + +def test_layer_ids_remap_applies_to_scales_too(): + # Hybrid GDN models back only their full-attention layers; a scale view that + # forgot the remap would hand layer 7's rows to layer 2's attention. + # (1, 3) names a subset of a model, so layer id 3 has to be inside it: LAYERS is 3, which + # makes 3 one past the end, so this pool is told the depth the ids imply. + layer_ids = (1, 3) + pool = _pool(layer_ids=layer_ids, num_layers=4) + assert pool._kv_buffer.shape[1] == 2 + with pytest.raises(KeyError): + pool.k_scale(0) +def test_store_kv_scatters_codes_and_scales(): + from freetoken.kernel.triton.kv_quant import codes_to_f32 + + tokens = 5 + torch.manual_seed(7) + rows = torch.randn(tokens, HEADS * DIM, device=DEV, dtype=torch.bfloat16) * 2.0 + out_loc = torch.tensor( + [3, 0, PAGES * PAGE_SIZE - 1, 40, 17], device=DEV, dtype=torch.int32 + ) + pool = _pool() + pool.store_kv(rows, rows.clone(), out_loc, layer_id=0) + torch.cuda.synchronize() + + codes = codes_to_f32(pool.k_cache(0).view(-1, HEADS, DIM)) + scale = pool.k_scale(0) + f32 = rows.view(tokens, HEADS, DIM).to(torch.float32) + deq = codes[out_loc.long()] * scale[out_loc.long()].unsqueeze(-1) + amax = f32.abs().amax(dim=-1, keepdim=True) + # The scales are per (token, head), so the bound is relative to each row's max. + err = ((deq - f32).abs() / amax.clamp_min(1e-6)).max() + assert float(err) < 0.08, float(err) + # Rows nobody wrote must stay exactly zero -- the scatter touched only out_loc. + untouched = torch.ones(PAGES * PAGE_SIZE, dtype=torch.bool, device=DEV) + untouched[out_loc.long()] = False + assert (codes[untouched] == 0).all() + assert (scale[untouched] == 0).all() + + +def test_rebuild_resizes_codes_and_scales_together(): + pool = _pool() + before = id(pool) + pool.rebuild(11) + assert id(pool) == before # identity preserved (backends cache the object) + assert pool._kv_buffer.shape == (2, LAYERS, 11, PAGE_SIZE, HEADS, DIM) + slots = 11 * PAGE_SIZE + assert pool.k_scale(0).shape == (slots, HEADS) + assert pool.k_cache(0).shape[0] == 11 + # Per-token cost is page-count invariant, codes and scales alike. + assert pool.unit_bytes() == _pool().unit_bytes() + + +def _sizing_config(kv_quant): + from freetoken.attention import AttnType + + spec = KVCacheGroupSpec( + name="full", + layer_ids=tuple(range(LAYERS)), + num_kv_heads=HEADS, + head_dim=DIM, + sliding_window=None, + attn_type=AttnType.FULL, + ) + mc = SimpleNamespace( + has_swa_attention=False, + has_linear_attention=False, + num_layers=LAYERS, + num_kv_heads=HEADS, + head_dim=DIM, + kv_cache_group_specs=lambda: (spec,), + ) + return SimpleNamespace( + model_config=mc, + page_size=PAGE_SIZE, + dtype=torch.bfloat16, + tp_info=SimpleNamespace(size=1), + kv_quant=kv_quant, + ) + + +@pytest.mark.parametrize("kv_quant", ["none", "fp8"]) +def test_unit_bytes_matches_the_cost_model_that_sized_the_pool(kv_quant): + """The budget solve and the live pool must agree byte for byte.""" + from freetoken.kvcache.base import spec_kv_bytes_per_token + from freetoken.kvcache.mha_pool import MHAKVCache + + config = _sizing_config(kv_quant) + (spec,) = config.model_config.kv_cache_group_specs() + per_token = spec_kv_bytes_per_token(spec, config) + assert MHAKVCache.kv_cost(config)[0] == per_token * PAGE_SIZE + + pool = _pool(kv_quant=kv_quant) + assert pool.unit_bytes() == (per_token, 0) + + +def test_fp8_lands_just_above_half_the_bytes(): + """Half the code bytes, plus a scale sidecar of 4 B per (token, slab, layer, head).""" + plain, quantized = _pool("none").unit_bytes()[0], _pool("fp8").unit_bytes()[0] + scales = 2 * LAYERS * HEADS * 4 + assert quantized == plain // 2 + scales, (plain, quantized, scales) + + +def test_latent_kv_pool_accepts_fp8(): + """MLA carries one latent row per token, so it uses one FP32 scale per row.""" + from freetoken.attention import AttnType + from freetoken.kvcache import create_kvcache_pool + + spec = KVCacheGroupSpec( + name="full", + layer_ids=(0, 1), + num_kv_heads=HEADS, + head_dim=DIM, + sliding_window=None, + mla=True, + attn_type=AttnType.MLA, + ) + mc = SimpleNamespace( + has_swa_attention=False, + has_linear_attention=False, + num_layers=2, + num_kv_heads=HEADS, + head_dim=DIM, + kv_cache_group_specs=lambda: (spec,), + ) + fp8 = create_kvcache_pool( + model_config=mc, + num_pages=4, + page_size=1, + dtype=torch.bfloat16, + device=DEV, + kv_quant="fp8", + ) + assert fp8.kv_quant == "fp8" and fp8.store_dtype is torch.uint8 + pool = create_kvcache_pool( + model_config=mc, + num_pages=4, + page_size=1, + dtype=torch.bfloat16, + device=DEV, + kv_quant="none", + ) + assert pool.kv_quant == "none" + + +def test_hybrid_swa_pool_also_separates_compute_and_store_dtype(monkeypatch): + """The hybrid-SWA pool is the other family that stores codes, so it must report the + same pair. Any backend sizing scratch off ``dtype`` would take e4m3 home with it + here too -- what that looked like in practice is QSA's indexer (kvcache/base.py).""" + from freetoken.distributed.info import DistributedInfo + from freetoken.kvcache.hybrid_swa_pool import HybridSWAKVCache + from freetoken.models.config import KVCacheGroupSpec + + monkeypatch.setattr( + "freetoken.kvcache.hybrid_swa_pool.get_tp_info", + lambda: DistributedInfo(rank=0, size=1), + ) + groups = ( + KVCacheGroupSpec( + name="full", layer_ids=(2, 5), num_kv_heads=2, head_dim=DIM, + sliding_window=None, + ), + KVCacheGroupSpec( + name="swa", layer_ids=(0, 1, 3, 4), num_kv_heads=2, head_dim=DIM, + sliding_window=32, + ), + ) + + def build(kv_quant: str): + return HybridSWAKVCache( + groups=groups, + num_layers=6, + num_full_pages=PAGES, + page_size=PAGE_SIZE, + num_swa_tokens=PAGES * PAGE_SIZE, + dtype=torch.bfloat16, + device=DEV, + kv_quant=kv_quant, + ) + + quantized, plain = build("fp8"), build("none") + assert quantized.dtype is torch.bfloat16 + assert plain.dtype is torch.bfloat16 + assert quantized.store_dtype == kv_codes_dtype() + assert plain.store_dtype is torch.bfloat16 + # ...while the buffers really did shrink: the two properties must not be aliases. + assert quantized.k_cache(2).element_size() == 1 + assert plain.k_cache(2).element_size() == 2 + assert quantized.k_scale(2) is not None and plain.k_scale(2) is None diff --git a/tests/kvcache/test_qsa_pool_fp8.py b/tests/kvcache/test_qsa_pool_fp8.py new file mode 100644 index 0000000000..02b525f679 --- /dev/null +++ b/tests/kvcache/test_qsa_pool_fp8.py @@ -0,0 +1,221 @@ +"""The QSA pool under ``--kv-cache-dtype fp8``: quantized K/V, 16-bit index tiers. + +Only the paged K/V slabs change. The compressed index slab, the per-request ring and +their scratch rows must stay 16-bit whatever the KV store does -- block selection scores +against them and the score kernel asserts their dtype -- while the byte account the +startup budget, the cache sliders and the runtime rebuild all divide by has to keep +telling the two halves apart. A drift here does not crash: it silently buys the wrong +number of pages. +""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch + +from freetoken.attention import AttnType +from freetoken.kvcache.base import spec_kv_bytes_per_token +from freetoken.models.config import KVCacheGroupSpec + +if not torch.cuda.is_available(): # pragma: no cover + pytest.skip("CUDA required", allow_module_level=True) + +from freetoken.kernel.triton.kv_quant import codes_to_f32, kv_codes_dtype +from freetoken.kvcache.qsa_pool import QSAKVCache + +DEV = torch.device("cuda") +PAGE_SIZE = 64 # the page size the qsa_sparse backend registers +LAYER_IDS = (1, 3, 5, 7) +HEADS, DIM, INDEX_DIM, INDEX_LAYERS, RATIO = 2, 64, 32, 4, 4 + + +@pytest.fixture(autouse=True) +def _tp(monkeypatch): + from freetoken.distributed.info import DistributedInfo + + monkeypatch.setattr( + "freetoken.kvcache.mha_pool.get_tp_info", + lambda: DistributedInfo(rank=0, size=1), + ) + + +def _pool(kv_quant="fp8", num_pages=4, num_req_slots=4): + return QSAKVCache( + num_kv_heads=HEADS, + num_layers=8, + head_dim=DIM, + num_pages=num_pages, + page_size=PAGE_SIZE, + dtype=torch.bfloat16, + device=DEV, + index_head_dim=INDEX_DIM, + num_index_layers=INDEX_LAYERS, + index_ratio=RATIO, + num_req_slots=num_req_slots, + layer_ids=LAYER_IDS, + kv_quant=kv_quant, + ) + + +def _spec(): + return KVCacheGroupSpec( + name="full", + layer_ids=LAYER_IDS, + num_kv_heads=HEADS, + head_dim=DIM, + sliding_window=None, + index_head_dim=INDEX_DIM, + num_index_layers=INDEX_LAYERS, + index_ratio=RATIO, + attn_type=AttnType.QSA, + ) + + +def _config(kv_quant, *, page_size=PAGE_SIZE, max_running_req=3): + mc = SimpleNamespace(num_layers=8, has_swa_attention=False, has_linear_attention=True, model_is_mrope=False) + mc.kv_cache_group_specs = lambda: (_spec(),) + return SimpleNamespace( + model_config=mc, + page_size=page_size, + dtype=torch.bfloat16, + tp_info=SimpleNamespace(size=1), + max_running_req=max_running_req, + kv_quant=kv_quant, + ) + + +def test_fp8_replaces_the_kv_slab_and_adds_scale_views(): + pool = _pool() + bf16 = _pool(kv_quant="none") + assert pool.kv_quant == "fp8" + # Same geometry as the 16-bit pool -- only the element type changed, because the + # attend kernels index codes exactly the way they index values. + assert pool._kv_buffer.shape == bf16._kv_buffer.shape + assert pool._kv_buffer.dtype == kv_codes_dtype() + # dtype = compute dtype (what the backend sizes its INDEXER scratch with), + # store_dtype = what the buffer holds. Swapping them is the bug that compiled an + # e4m3 operand into QSA's scoring dot and died at graph capture. + assert pool.dtype is torch.bfloat16 + assert pool.store_dtype == kv_codes_dtype() + assert bf16.dtype is torch.bfloat16 and bf16.store_dtype is torch.bfloat16 + assert pool.k_cache(3).shape == (4, PAGE_SIZE, HEADS, DIM) + assert pool.k_cache(3).element_size() == 1 + slots = 4 * PAGE_SIZE + assert pool.k_scale(3).shape == (slots, HEADS) + assert pool.k_scale(3).dtype is torch.float32 + assert pool.v_scale(3).shape == (slots, HEADS) + # Zero-filled: e4m3 has NaN bit patterns, and the dummy page / unwritten tail rows + # must never read back as one. + assert pool.k_scale(3).abs().sum().item() == 0.0 + assert codes_to_f32(pool.k_cache(3)).abs().sum().item() == 0.0 + # A 16-bit pool keeps answering None, so the backends' k_scale(...) pass-through is + # the only branch that ever differs between the two. + assert bf16.k_scale(3) is None and bf16.v_scale(3) is None + + +@pytest.mark.parametrize("kv_quant", ["none", "fp8", "nvfp4"]) +def test_index_tiers_stay_16_bit_whatever_the_kv_store_does(kv_quant): + """Block selection is quantization-agnostic by construction: it reads the compressed + index keys, not the KV rows, so fp8 must not touch these three buffers.""" + pool = _pool(kv_quant=kv_quant) + assert pool.cmp_k_cache(0).dtype is torch.bfloat16 + assert pool.pending_ring(0).dtype is torch.bfloat16 + assert pool.cmp_k_cache(0).shape == (4 * PAGE_SIZE // RATIO + 4, INDEX_DIM) + # ...and the byte account says so too: only the 1-byte KV codes got cheaper. + spec = _spec() + cost = spec_kv_bytes_per_token(spec, _config(kv_quant)) + plain = spec_kv_bytes_per_token(spec, _config("none")) + kv_layers = len(LAYER_IDS) + kv_16bit = 2 * DIM * HEADS * 2 * kv_layers # two slabs, 16-bit codes + scale_term = 2 * kv_layers * HEADS * 4 # one fp32 scale per (slot, head) + index_term = INDEX_DIM * INDEX_LAYERS * 2 // RATIO # the untouched 16-bit slab + assert plain == kv_16bit + index_term + if kv_quant == "fp8": + assert cost == kv_16bit // 2 + scale_term + index_term + elif kv_quant == "nvfp4": + block_term = 2 * kv_layers * HEADS * (DIM // 16) + assert cost == kv_16bit // 4 + scale_term + block_term + index_term + else: + assert cost == plain + + +def test_unit_bytes_and_kv_cost_still_agree_when_quantized(): + """The pool's own allocation and the budget model must divide the same way -- the + scale sidecar is priced in base.spec_kv_bytes_per_token, not here.""" + spec, config = _spec(), _config("fp8") + pool = _pool() + kv_bytes, swa_bytes = pool.unit_bytes() + assert swa_bytes == 0 + assert kv_bytes == spec_kv_bytes_per_token(spec, config) + assert kv_bytes * PAGE_SIZE == QSAKVCache.kv_cost(config)[0] + # The 16-bit pool's per-token figure is the reference: codes halve the KV term, the + # scale sidecar and the untouched index slab keep the total above a clean half. + plain_kv = spec_kv_bytes_per_token(spec, _config("none")) + assert plain_kv // 2 < kv_bytes < plain_kv + + +def test_rebuild_resizes_codes_and_scales_together(): + pool = _pool(num_pages=4) + ident = id(pool) + pool.rebuild(16) + assert id(pool) == ident + assert pool.k_cache(1).shape == (16, PAGE_SIZE, HEADS, DIM) + assert pool.k_scale(1).shape == (16 * PAGE_SIZE, HEADS) + assert pool.k_scale(1).abs().sum().item() == 0.0 + assert pool.cmp_k_cache(0).shape == (16 * PAGE_SIZE // RATIO + 4, INDEX_DIM) + + +def test_store_kv_writes_the_slot_the_attend_kernel_will_read(): + """out_loc numbering (page * page_size + offset) is the contract between the fused + writer and the attend kernel's scale slot arithmetic -- this is that round trip.""" + torch.manual_seed(0) + pool = _pool(num_pages=5) + slots = 5 * PAGE_SIZE + rows = (0, 1, 63, 64, 255, 256) # page boundaries included: 63/64 and 255/256 + k = torch.randn(len(rows), HEADS * DIM, device=DEV, dtype=torch.bfloat16) * 3.0 + v = torch.randn(len(rows), HEADS * DIM, device=DEV, dtype=torch.bfloat16) * 0.25 + out_loc = torch.tensor(rows, dtype=torch.int32, device=DEV) + + pool.store_kv(k, v, out_loc, layer_id=3) + torch.cuda.synchronize() + + codes = codes_to_f32(pool.k_cache(3).view(slots, HEADS, DIM)) + decoded_k = codes[out_loc] * pool.k_scale(3)[out_loc].unsqueeze(-1) + decoded_v = codes_to_f32(pool.v_cache(3).view(slots, HEADS, DIM))[out_loc] * pool.v_scale(3)[ + out_loc + ].unsqueeze(-1) + for want, got in ((k, decoded_k), (v, decoded_v)): + w = want.view(len(rows), HEADS, DIM).to(torch.float32) + amax = w.abs().amax(dim=-1, keepdim=True) + assert (((got - w).abs() / amax).max().item()) < 0.07 + # Rows nobody wrote stay zero rather than NaN -- the dummy page depends on it. + untouched = torch.tensor([r for r in range(64) if r not in rows], device=DEV) + assert pool.k_scale(3)[untouched].abs().sum().item() == 0.0 + + +def test_factory_threads_kv_quant_into_the_qsa_pool(): + from freetoken.kvcache import create_kvcache_pool + + mc = SimpleNamespace( + num_layers=8, has_swa_attention=False, has_linear_attention=True, + num_kv_heads=HEADS, head_dim=DIM, dsv4_args=None, model_is_mrope=False, + ) + mc.kv_cache_group_specs = lambda: (_spec(),) + pool = create_kvcache_pool( + mc, num_pages=4, page_size=PAGE_SIZE, dtype=torch.bfloat16, device=DEV, + num_req_slots=4, kv_quant="fp8", + ) + assert isinstance(pool, QSAKVCache) and pool.kv_quant == "fp8" + assert pool.k_cache(1).element_size() == 1 and pool.k_scale(1) is not None + + +def test_nvfp4_replaces_only_qsa_kv_tiers(): + pool = _pool(kv_quant="nvfp4") + slots = 4 * PAGE_SIZE + assert pool.k_cache(3).shape == (4, PAGE_SIZE, HEADS, DIM // 2) + assert pool.k_block_scale(3).shape == (slots, HEADS, DIM // 16) + assert pool.v_block_scale(3).dtype is torch.uint8 + assert pool.cmp_k_cache(0).dtype is torch.bfloat16 + diff --git a/tests/kvcache/test_tier3.py b/tests/kvcache/test_tier3.py new file mode 100644 index 0000000000..335a0c7094 --- /dev/null +++ b/tests/kvcache/test_tier3.py @@ -0,0 +1,85 @@ +"""Degrau 3 — teste standalone do KV store host do HostPrefixTier (CPU). + +Cobre: round-trip D2H->H2D de páginas KV (códigos + scales), a serpente dos ids lógicos +(a página lógica reescrita depois do evict NÃO contamina a cópia host), merge de entrada +(gdn depois kv), e drop-oldest quando o budget enche. +""" +import torch +from types import SimpleNamespace + +from freetoken.kvcache.host_prefix_tier import HostPrefixTier + + +def main(): + # pool fake: [2, L=2, P=8, page=4, H=2, D=3] códigos bf16 + scales [2, L, P*4, 2] fp32 + L, P, PAGE, H, D = 2, 8, 4, 2, 3 + kv_pool = SimpleNamespace( + _kv_buffer=torch.randn(2, L, P, PAGE, H, D, dtype=torch.bfloat16), + _scale_buffer=torch.randn(2, L, P * PAGE, H, dtype=torch.float32), + ) + state_pool = SimpleNamespace( + conv_states=torch.randn(L, 3, 8, 3, dtype=torch.bfloat16), + recurrent_states=torch.randn(L, 3, 2, 4, 4, dtype=torch.float32), + slot_states={}, + ) + page_bytes = 2 * L * PAGE * H * D * 2 + 2 * L * PAGE * H * 4 + # budget pra 3 páginas (força drop-oldest cedo) + tier = HostPrefixTier(state_pool, kv_pool, gdn_slots_host=2, + kv_budget_bytes=3 * page_bytes, device="cpu") + assert tier.has_kv and len(tier._kv_free) == 3 + + # 1) write/read round-trip: prefixo de 2 páginas (páginas lógicas 3 e 5) + kv_idx = torch.cat([torch.arange(3 * PAGE, 4 * PAGE), torch.arange(5 * PAGE, 6 * PAGE)]).to(torch.int32) + orig_codes = kv_pool._kv_buffer[:, :, 3].clone(), kv_pool._kv_buffer[:, :, 5].clone() + orig_scales = (kv_pool._scale_buffer[:, :, 3 * PAGE:4 * PAGE].clone(), + kv_pool._scale_buffer[:, :, 5 * PAGE:6 * PAGE].clone()) + hs = tier.write_kv(kv_idx) + assert hs is not None and hs.numel() == 2 + + # 2) A SERPENTE: sobrescreve as páginas lógicas 3 e 5 (novo inquilino) + kv_pool._kv_buffer[:, :, 3].fill_(99) + kv_pool._kv_buffer[:, :, 5].fill_(99) + kv_pool._scale_buffer[:, :, 3 * PAGE:6 * PAGE].fill_(99) + # lê de volta pra páginas 0 e 7 + tier.read_kv(hs, torch.tensor([0, 7], dtype=torch.int32)) + assert torch.equal(kv_pool._kv_buffer[:, :, 0], orig_codes[0]) + assert torch.equal(kv_pool._kv_buffer[:, :, 7], orig_codes[1]) + assert torch.equal(kv_pool._scale_buffer[:, :, 0 * PAGE:1 * PAGE], orig_scales[0]) + assert torch.equal(kv_pool._scale_buffer[:, :, 7 * PAGE:8 * PAGE], orig_scales[1]) + print("PASS 1: round-trip KV host + ids lógicos reescritos não contaminam") + + # 3) merge: entrada com gdn primeiro, kv depois + toks = torch.arange(64 * 2, dtype=torch.int64) + g = tier.write_gdn(1) + tier.put_prefix(toks, None, g) + tier.put_prefix(toks, hs, None) # merge: adota kv, mantém gdn + m = tier.lookup(toks) + assert m is not None and m.gdn_slot == g and torch.equal(m.kv_pages, hs) + print("PASS 2: merge gdn+kv na mesma entrada") + + # 4) drop-oldest quando o budget enche: escreve 2 páginas (resta 1), depois 2 de novo + # -> a entrada mais velha (toks) é dropada INTEIRA (KV slots + slot GDN) + kv2 = torch.cat([torch.arange(0 * PAGE, 1 * PAGE), torch.arange(2 * PAGE, 3 * PAGE)]).to(torch.int32) + toks2 = torch.arange(64 * 2, 64 * 3, dtype=torch.int64) + s2 = tier.write_kv(kv2) # precisa de 2 slots, tem 1 -> drop de toks (inteiro) + tier.put_prefix(toks2, s2, None) + assert tier.lookup(toks) is None, "o mais velho devia ter sido dropado inteiro" + assert g in tier._free, "o slot GDN do dropado devia voltar pro free-list" + kv3 = torch.arange(1 * PAGE, 3 * PAGE).to(torch.int32) + toks3 = torch.arange(64 * 3, 64 * 4, dtype=torch.int64) + s3 = tier.write_kv(kv3) # precisa de 2, tem 1 -> drop de toks2 + tier.put_prefix(toks3, s3, None) + assert tier.lookup(toks2) is None + m3 = tier.lookup(toks3) + assert m3 is not None and torch.equal(m3.kv_pages, s3) + print("PASS 3: drop-oldest remove a entrada mais fria inteira (KV + GDN)") + + # 5) evict_host libera tudo + tier.evict_host(toks3) + assert tier.lookup(toks3) is None and len(tier._kv_free) == 3 + print("PASS 4: evict_host remove e libera") + + print("\nTODOS OS TESTES PASSARAM ✓") + + +main() diff --git a/tests/kvcache/test_tree_hooks.py b/tests/kvcache/test_tree_hooks.py new file mode 100644 index 0000000000..840629310f --- /dev/null +++ b/tests/kvcache/test_tree_hooks.py @@ -0,0 +1,77 @@ +"""Degrau 2 — teste standalone dos hooks do HybridRadixCache (CPU, sem servidor).""" +import torch +from freetoken.kvcache.hybrid_radix_cache import HybridRadixCache + + +def main(): + tree = HybridRadixCache(torch.device("cpu"), 64) + evicted = [] # (prefix_len, mamba_slot) capturados pelo hook de evicção + rehydrated = [] # prefix_len reidratado no match + host_index = set() + + tree.on_evict_node = lambda node: evicted.append( + (len(tree._collect_key(node)), node.mamba_value)) + next_slot = [100] + + def fake_rehydrate(node): + n = len(tree._collect_key(node)) + if n in host_index: + rehydrated.append(n) + s = next_slot[0] + next_slot[0] += 1 + return s + return None + + tree.on_match_dead_snapshot = fake_rehydrate + + # A (4096 toks, snapshot 7) + filho B (mais 2048, snapshot 8) -> A é nó INTERNO + idsA = torch.arange(4096, dtype=torch.int32) + kvA = torch.arange(4096, dtype=torch.int32) * 10 + assert tree.insert(idsA, kvA, mamba_value=7) == (0, False) + idsB = torch.cat([idsA, torch.full((2048,), 555, dtype=torch.int32)]) + kvB = torch.arange(6144, dtype=torch.int32) * 10 + assert tree.insert(idsB, kvB, mamba_value=8) == (4096, False) + + # match enquanto tudo vivo: ganha o snapshot mais profundo (B, 6144) + m = tree.match_prefix(idsB) + assert m.cached_len == 6144 and m.mamba_value == 8 + + # evicta os 2 snapshots: A (interno -> tombstone) e B (folha -> com tier, tombstone + # in-place: o KV da folha FICA no tree, não é liberado) + er = tree.evict_mamba(2) + assert sorted(er.mamba_slots) == [7, 8], er + assert sorted(evicted) == [(4096, 7), (6144, 8)], evicted + assert er.kv_indices.numel() == 0, "com o tier, nenhuma página KV é liberada" + print("PASS 1: on_evict_node dispara; folha vira tombstone e mantém o KV:", evicted) + + # sem cópia host: match trunca pra 0 (A e B viraram tombstones) + m = tree.match_prefix(idsA) + assert m.cached_len == 0 and m.mamba_value is None + print("PASS 2: match trunca sem tier host") + + # host tem A: match(idsA) reidrata A + host_index.add(4096) + m = tree.match_prefix(idsA) + assert m.cached_len == 4096 and m.mamba_value == 100, (m.cached_len, m.mamba_value) + assert rehydrated == [4096] + print("PASS 3: match reidrata tombstone interno via tier host (slot 100)") + + # host tem B também: match(idsB) reidrata a FOLHA tombstone (mais profunda) + host_index.add(6144) + m = tree.match_prefix(idsB) + assert m.cached_len == 6144 and m.mamba_value == 101, (m.cached_len, m.mamba_value) + print("PASS 4: match reidrata folha-tombstone (KV preservado) -> 6144") + + # nós reidratados ficam vivos: segundo match não chama o hook + m2 = tree.match_prefix(idsB) + assert m2.mamba_value == 101 and rehydrated == [4096, 6144] + print("PASS 5: nós reidratados ficam vivos (sem H2D repetido)") + + # integridade de contagem: os dois snapshots revividos voltam a contar como evictable + assert tree.mamba_evictable == 2, tree.mamba_evictable + print("PASS 6: contagem mamba_evictable consistente") + + print("\nTODOS OS TESTES PASSARAM ✓") + + +main() diff --git a/tests/models/qwen4_exp/common.py b/tests/models/qwen4_exp/common.py index 0a52fb0e01..f3bddd7a93 100644 --- a/tests/models/qwen4_exp/common.py +++ b/tests/models/qwen4_exp/common.py @@ -154,6 +154,7 @@ def __init__( device: str = "cuda", dtype: torch.dtype = torch.bfloat16, page_size: int = 64, + kv_quant: str = "none", ) -> None: from freetoken.attention.qsa_sparse import QSASparseAttnBackend from freetoken.kvcache import create_kvcache_pool @@ -170,6 +171,7 @@ def __init__( dtype=dtype, device=self.device, num_req_slots=self.num_req_slots, + kv_quant=kv_quant, ) self.page_table = torch.zeros( (self.num_req_slots, num_pages * page_size), dtype=torch.int32, device=self.device diff --git a/tests/models/qwen4_exp/conftest.py b/tests/models/qwen4_exp/conftest.py index 2b9e80214f..4e2bb67bb8 100644 --- a/tests/models/qwen4_exp/conftest.py +++ b/tests/models/qwen4_exp/conftest.py @@ -13,3 +13,16 @@ def _runtime(): core._GLOBAL_CTX = None yield core._GLOBAL_CTX = None + import torch + + if torch.cuda.is_available() and torch.cuda.is_initialized(): + import gc + + # torch keeps a cuBLAS workspace per stream for good: each offload cache's dedicated + # streams would otherwise pin ~16 MiB for the rest of the session + gc.collect() + torch.cuda.synchronize() + clear = getattr(torch._C, "_cuda_clearCublasWorkspaces", None) # private torch API + if clear is not None: + clear() + torch.cuda.empty_cache() diff --git a/tests/models/qwen4_exp/test_build_cpu.py b/tests/models/qwen4_exp/test_build_cpu.py new file mode 100644 index 0000000000..85366638a5 --- /dev/null +++ b/tests/models/qwen4_exp/test_build_cpu.py @@ -0,0 +1,78 @@ +"""The whole model must CONSTRUCT without a GPU. + +Every other model test that builds a decoder layer is behind requires_cuda, so a constructor +signature mismatch between Qwen4ExpDecoderLayer and the ops it builds is invisible to a +CPU-only run -- it only shows up when a server boots. That happened: Qwen4ExpMoE did not +accept the `prefix` kwarg model.py passes, and a full CPU suite still went green. + +Building on the meta device costs no memory and no GPU, so there is no reason not to. +""" + +import torch +from freetoken.distributed import set_tp_info, try_get_tp_info +from freetoken.layers import rotary, set_rope_device +from freetoken.models.qwen4_exp.config import parse_config +from freetoken.models.qwen4_exp.model import Qwen4ExpForCausalLM + +from .common import toy_hf_config + + +def _build(): + if try_get_tp_info() is None: + set_tp_info(rank=0, size=1) + previous = rotary._ROPE_DEVICE + set_rope_device(torch.device("cpu")) + try: + with torch.device("meta"): + return Qwen4ExpForCausalLM(parse_config(toy_hf_config())) + finally: + # get_rope is cached per geometry: a CPU cos/sin table left behind feeds later CUDA tests + rotary.get_rope.cache_clear() + set_rope_device(previous) + + +def test_model_constructs_on_meta_without_cuda(): + model = _build() + sd = model.state_dict() + assert sd, "state dict is empty" + # both layer families and the head must be present + assert any(".self_attn." in k for k in sd), "no full-attention layer built" + assert any(".linear_attn." in k for k in sd), "no GDN layer built" + assert any(k.startswith("lm_head") for k in sd), "no lm_head built" + + +def test_every_layer_gets_its_own_prefixed_weights(): + # a prefix that is dropped or shared silently collapses layers onto one another + sd = _build().state_dict() + mlp_keys = {k for k in sd if ".mlp." in k} + layers = {k.split(".layers.")[1].split(".")[0] for k in mlp_keys if ".layers." in k} + assert len(layers) > 1, f"MoE weights landed on a single layer: {sorted(layers)}" + + +def test_offloaded_layers_borrow_the_next_router_for_the_lookahead(): + # FREETOKEN_MOE_PREFETCH: layer L-1's experts hold layer L's router, outside the state dict + # and without the engine's offload-layer walk finding extra layers through it + import dataclasses + + from freetoken.moe.offload_cache import iter_offload_moe_layers + + if try_get_tp_info() is None: + set_tp_info(rank=0, size=1) + previous = rotary._ROPE_DEVICE + set_rope_device(torch.device("cpu")) + try: + config = dataclasses.replace(parse_config(toy_hf_config()), moe_strategy="offload") + with torch.device("meta"): + model = Qwen4ExpForCausalLM(config) + finally: + rotary.get_rope.cache_clear() + set_rope_device(previous) + mlps = [layer.mlp for layer in model.model.layers.op_list] + assert len(list(iter_offload_moe_layers(model))) == len(mlps) == 4 + wired = [mlp.experts._lookahead for mlp in mlps] + assert wired[-1] is None + assert [(gate is mlps[t].gate, t, budget) for gate, t, budget in wired[:-1]] == [ + (True, 1, 3), (True, 2, 3), (True, 3, 4), # layer 3 is the full-attention one + ] + assert not any("lookahead" in k for k in model.state_dict()) + assert sum(k.endswith(".mlp.gate.weight") for k in model.state_dict()) == 4 diff --git a/tests/models/qwen4_exp/test_fp8_dense.py b/tests/models/qwen4_exp/test_fp8_dense.py new file mode 100644 index 0000000000..74da973341 --- /dev/null +++ b/tests/models/qwen4_exp/test_fp8_dense.py @@ -0,0 +1,222 @@ +"""Load-time per-tensor FP8 for the qwen4_exp dense projections (FREETOKEN_FP8_DENSE=1).""" + +from types import SimpleNamespace + +import pytest +import torch +from freetoken.models.qwen4_exp.weight import _fp8_dense, _quantize_per_tensor + +from .common import requires_cuda + +E4M3_MAX = 448.0 + + +def _cfg(): + g = SimpleNamespace( + num_key_heads=4, key_head_dim=8, num_value_heads=12, value_head_dim=8 + ) + return SimpleNamespace(linear_attention_group=lambda: g) + + +def _assert_e4m3_close( + deq: torch.Tensor, ref: torch.Tensor, scale: torch.Tensor +) -> None: + # e4m3 keeps 3 mantissa bits (rel 2^-4); below scale * 2^-6 it is subnormal (abs 2^-9 steps) + tol = ref.abs() * 2**-4 + float(scale) * 2**-9 + 1e-7 + assert ((deq - ref).abs() <= tol).all() + + +def test_quantize_per_tensor_round_trips_within_e4m3(): + w = torch.randn(64, 32) * 0.02 + w8, scale = _quantize_per_tensor(w) + assert ( + w8.dtype == torch.float8_e4m3fn + and scale.dtype == torch.float32 + and scale.shape == () + ) + assert torch.isclose(scale * E4M3_MAX, w.abs().max()) + _assert_e4m3_close(w8.float() * scale, w, scale) + + +def test_in_proj_splits_into_fp8_qkvz_and_bf16_ba_per_rank(): + cfg, world = ( + _cfg(), + 2, + ) # local: 2 k heads, 6 v heads -> qkvz = 2*2*8 + 2*6*8 = 128, ba = 12 + t = torch.randn(128 + 12, 16, dtype=torch.bfloat16) + out = dict(_fp8_dense("model.layers.3.linear_attn.in_proj.weight", t, cfg, world)) + assert sorted(out) == [ + "model.layers.3.linear_attn.in_proj_ba.weight", + "model.layers.3.linear_attn.in_proj_qkvz.weight", + "model.layers.3.linear_attn.in_proj_qkvz.weight_scale", + ] + w8 = out["model.layers.3.linear_attn.in_proj_qkvz.weight"] + scale = out["model.layers.3.linear_attn.in_proj_qkvz.weight_scale"] + assert w8.shape == (128, 16) and w8.dtype == torch.float8_e4m3fn + _assert_e4m3_close(w8.float() * scale, t[:128].float(), scale) + ba = out["model.layers.3.linear_attn.in_proj_ba.weight"] + assert ba.dtype == torch.bfloat16 and torch.equal(ba, t[128:]) + # its own storage: a view would keep the whole bf16 in_proj alive next to the fp8 copy + assert ba.untyped_storage().data_ptr() != t.untyped_storage().data_ptr() + + +def test_other_projections_gain_a_scale_and_the_rest_pass_through(): + cfg = _cfg() + t = torch.randn(32, 16, dtype=torch.bfloat16) + for name in ( + "x.self_attn.qkv_proj.weight", + "x.self_attn.o_proj.weight", + "x.linear_attn.out_proj.weight", + ): + out = dict(_fp8_dense(name, t, cfg, 1)) + assert sorted(out) == [name, name[: -len("weight")] + "weight_scale"] + assert out[name].dtype == torch.float8_e4m3fn + out = dict(_fp8_dense("x.mlp.shared_expert.gate_up_proj.weight", t, cfg, 1)) + assert ( + list(out) == ["x.mlp.shared_expert.gate_up_proj.weight"] + and out.popitem()[1] is t + ) + + +def test_ops_declare_fp8_weight_and_scalar_scale(): + from freetoken.distributed import set_tp_info, try_get_tp_info + from freetoken.layers.fp8_dynamic import Fp8DynamicColMerged, Fp8DynamicRowParallel + + if try_get_tp_info() is None: + set_tp_info(rank=0, size=1) + col = Fp8DynamicColMerged(32, [64, 16, 16]) + row = Fp8DynamicRowParallel(64, 32) + for op in (col, row): + sd = op.state_dict() + assert set(sd) == {"weight", "weight_scale"} + assert ( + sd["weight"].dtype == torch.float8_e4m3fn and sd["weight_scale"].shape == () + ) + assert col.weight.shape == (96, 32) and row.weight.shape == (32, 64) + + +@requires_cuda +@pytest.mark.parametrize("rows", [1, 16, 300], ids=["decode-1", "decode-16", "prefill"]) +def test_fp8_linear_matches_bf16_on_the_dequantized_weight(rows: int): + from freetoken.layers.fp8_dynamic import fp8_dynamic_linear + + torch.manual_seed(0) + w = (torch.randn(256, 128, device="cuda") * 0.02).to(torch.bfloat16) + w8, scale = _quantize_per_tensor(w) + x = torch.randn(rows, 128, device="cuda", dtype=torch.bfloat16) + ref = torch.nn.functional.linear(x, (w8.float() * scale).to(torch.bfloat16)) + got = fp8_dynamic_linear(x, w8, scale.to("cuda")) + assert got.shape == ref.shape and got.dtype == torch.bfloat16 + # the per-tensor activation cast is the only extra rounding: 2^-4 relative on the inputs + torch.testing.assert_close( + got.float(), ref.float(), rtol=0.1, atol=0.08 * ref.abs().max().item() + ) + + +@requires_cuda +def test_quant_per_tensor_zero_input_is_finite(): + from freetoken.layers.fp8_dynamic import quant_per_tensor + + x8, scale = quant_per_tensor( + torch.zeros(16, 128, device="cuda", dtype=torch.bfloat16) + ) + assert torch.isfinite(scale).item() and (x8.float() == 0).all() + + +@requires_cuda +@pytest.mark.parametrize("n", [4096, 16384, 20480, 81920, 200000]) +def test_split_quant_path_agrees_with_the_single_program_one(n: int): + """The two-launch path above _SPLIT_MIN_ELEMENTS must quantize bit for bit like the + one-program kernel: same amax (max is exact under any grouping), same arithmetic.""" + import freetoken.layers.fp8_dynamic as m + + torch.manual_seed(0) + x = (torch.randn(n, device="cuda") * 3.0).to(torch.bfloat16) + got8, got_scale = m.quant_per_tensor(x) + ref8 = torch.empty_like(x, dtype=m.FP8) + ref_scale = torch.empty((), dtype=torch.float32, device="cuda") + m._quant_fused_kernel[(1,)](x, ref8, ref_scale, n, BLOCK=m._BLOCK, num_warps=8) + assert torch.equal(got_scale, ref_scale) + assert torch.equal(got8.view(torch.uint8), ref8.view(torch.uint8)) + + +def test_reader_gates_lm_head_behind_its_own_flag(): + cfg = _cfg() + t = torch.randn(64, 32, dtype=torch.bfloat16) + assert list(dict(_fp8_dense("lm_head.weight", t, cfg, 1))) == ["lm_head.weight"] + out = dict(_fp8_dense("lm_head.weight", t, cfg, 1, dense=False, lm_head=True)) + assert sorted(out) == ["lm_head.weight", "lm_head.weight_scale"] + assert out["lm_head.weight"].dtype == torch.float8_e4m3fn + # the lm_head flag alone must not pull in the dense rewrites + passthrough = dict(_fp8_dense("x.self_attn.qkv_proj.weight", t, cfg, 1, + dense=False, lm_head=True)) + assert list(passthrough) == ["x.self_attn.qkv_proj.weight"] + + +def _ftw_with(tmp_path, names): + from freetoken.checkpoint.ftw import FTWWriter + + writer = FTWWriter(str(tmp_path)) + for name in names: + writer.add_tensor(name, torch.zeros(4, dtype=torch.bfloat16)) + writer.finalize({}) + return str(tmp_path) + + +_BF16_NAMES = ["model.layers.0.linear_attn.in_proj.weight", "lm_head.weight"] +_FP8_NAMES = [ + "model.layers.0.linear_attn.in_proj_qkvz.weight", + "model.layers.0.linear_attn.in_proj_qkvz.weight_scale", + "model.layers.0.linear_attn.in_proj_ba.weight", + "lm_head.weight", + "lm_head.weight_scale", +] + + +@pytest.mark.parametrize( + "names, env, bad", + [ + (_BF16_NAMES, {"FREETOKEN_FP8_DENSE": "1", "FREETOKEN_FP8_LMHEAD": "1"}, + "FREETOKEN_FP8_DENSE=0, FREETOKEN_FP8_LMHEAD=0"), + (_FP8_NAMES, {}, "FREETOKEN_FP8_DENSE=1, FREETOKEN_FP8_LMHEAD=1"), + (_FP8_NAMES, {"FREETOKEN_FP8_DENSE": "1"}, "converted with FREETOKEN_FP8_LMHEAD=1, but"), + ], +) +def test_ftw_load_refuses_fp8_flags_that_differ_from_conversion(tmp_path, monkeypatch, names, env, bad): + """#8: the FTW replays conversion-time tensors, so the FP8 flags must match before any byte is read.""" + from freetoken.models import weight as model_weight + + for var in ("FREETOKEN_FP8_DENSE", "FREETOKEN_FP8_LMHEAD"): + monkeypatch.delenv(var, raising=False) + for var, value in env.items(): + monkeypatch.setenv(var, value) + path = _ftw_with(tmp_path, names) + spec = SimpleNamespace(module="freetoken.models.qwen4_exp") + config = SimpleNamespace( + attn_quant="fp8_dynamic" if env.get("FREETOKEN_FP8_DENSE") else "none", + tie_word_embeddings=False, + quant=None, + ) + monkeypatch.setattr(model_weight, "cached_load_hf_config", lambda _: SimpleNamespace(architectures=["x"])) + monkeypatch.setattr(model_weight, "get_model_spec", lambda _: spec) + monkeypatch.setattr(model_weight, "_spec_for_model_path", lambda _: (config, spec)) + with pytest.raises(ValueError, match=bad): + next(model_weight.load_weight(path, torch.device("cpu"))) + + +@pytest.mark.parametrize("names, fp8", [(_BF16_NAMES, False), (_FP8_NAMES, True)]) +def test_ftw_load_accepts_matching_fp8_flags(tmp_path, monkeypatch, names, fp8): + from freetoken.models import weight as model_weight + + for var in ("FREETOKEN_FP8_DENSE", "FREETOKEN_FP8_LMHEAD"): + if fp8: + monkeypatch.setenv(var, "1") + else: + monkeypatch.delenv(var, raising=False) + path = _ftw_with(tmp_path, names) + spec = SimpleNamespace(module="freetoken.models.qwen4_exp") + config = SimpleNamespace(attn_quant="fp8_dynamic" if fp8 else "none", tie_word_embeddings=False, quant=None) + monkeypatch.setattr(model_weight, "cached_load_hf_config", lambda _: SimpleNamespace(architectures=["x"])) + monkeypatch.setattr(model_weight, "get_model_spec", lambda _: spec) + monkeypatch.setattr(model_weight, "_spec_for_model_path", lambda _: (config, spec)) + assert [name for name, _ in model_weight.load_weight(path, torch.device("cpu"))] == names diff --git a/tests/models/qwen4_exp/test_ple.py b/tests/models/qwen4_exp/test_ple.py index 0d2461427e..954d898945 100644 --- a/tests/models/qwen4_exp/test_ple.py +++ b/tests/models/qwen4_exp/test_ple.py @@ -17,6 +17,7 @@ import pytest import torch +import freetoken.models.qwen4_exp.ple as ple_module from freetoken.models.config import ModelConfig from freetoken.models.qwen4_exp.config import parse_config from freetoken.models.qwen4_exp.ple import ( @@ -173,6 +174,211 @@ def test_decode_hash_matches_prefill_hash(): assert torch.equal(got[i], prefill[i * len(sequences[0]) + step]) +_FUSED_HASH_SEQS = [[3, 4, EOS, 5, 6, 8], [2, EOS, 11, 12, 13, 14], [7]] +_FUSED_HASH_CTX = [[EOS, EOS], [21, 22], [31, EOS]] + + +def test_token_index_addresses_the_same_window_as_the_packed_build(): + """``_token_index`` names the (request, column) the packed window would have used.""" + layer = _make_layer(_config()) + embedding = layer.ple_embedding + meta = _meta(_FUSED_HASH_SEQS, _FUSED_HASH_CTX) + packed, select = embedding._window(meta) + req, local = embedding._token_index(meta) + ctx_len = embedding.ngram_size - 1 + assert torch.equal(packed[req.long(), local.long() + ctx_len], meta.input_ids) + assert torch.equal(select(packed), meta.input_ids) + + +def test_token_index_cache_is_keyed_by_shape(): + """A repeat of the same shape reuses the buffers; a different one does not.""" + layer = _make_layer(_config()) + embedding = layer.ple_embedding + meta = _meta(_FUSED_HASH_SEQS, _FUSED_HASH_CTX) + first = embedding._token_index(meta) + assert embedding._token_index(meta)[0] is first[0] + other = _meta([[3, 4]], [[EOS, EOS]]) + assert embedding._token_index(other)[0] is not first[0] + + +def test_token_index_capture_never_reuses_an_evictable_cache_entry(monkeypatch): + """The eager warmup forward caches the decode pair right before capture; a capture that + reused it would bake an address the FIFO later evicts while the graph still reads it.""" + embedding = _make_layer(_config()).ple_embedding + decode = _meta([[5]], [[EOS, EOS]], decode=True) + eager = embedding._token_index(decode) # the warmup forward + monkeypatch.setattr(ple_module, "_capturing", lambda device: True) + captured = embedding._token_index(decode) + assert captured[0] is not eager[0] and captured[1] is not eager[1] + assert all(v[0] is not captured[0] for v in embedding._token_index_cache.values()) + assert captured[0].tolist() == [0] and captured[1].tolist() == [0] + + +def test_token_index_does_not_confuse_a_decode_with_a_one_request_prefill(): + """Same [T], opposite meaning: B decode rows at offset 0 vs B offsets in one request.""" + layer = _make_layer(_config()) + embedding = layer.ple_embedding + tokens = [5, 6, 7] + context = [[EOS, EOS], [21, 22], [31, EOS]] + decode_req, decode_local = embedding._token_index( + _meta([[t] for t in tokens], context, decode=True) + ) + prefill_req, prefill_local = embedding._token_index(_meta([tokens], context[:1])) + assert decode_req.tolist() == [0, 1, 2] and decode_local.tolist() == [0, 0, 0] + assert prefill_req.tolist() == [0, 0, 0] and prefill_local.tolist() == [0, 1, 2] + + +def test_token_index_cache_is_bounded(): + """The memo is an LRU-ish ring, not an unbounded map keyed by every prefill shape.""" + embedding = _make_layer(_config()).ple_embedding + for length in range(4 * ple_module._TOKEN_INDEX_CACHE_SIZE): + embedding._token_index(_meta([list(range(length + 1))], [[EOS, EOS]])) + assert len(embedding._token_index_cache) == ple_module._TOKEN_INDEX_CACHE_SIZE + + +def test_fused_hash_is_off_on_cpu(): + """The Triton path needs CUDA; a CPU meta stays on the torch reference.""" + layer = _make_layer(_config()) + meta = _meta(_FUSED_HASH_SEQS, _FUSED_HASH_CTX) + assert not layer.ple_embedding._use_fused_row_ids(meta) + assert torch.equal( + layer.ple_embedding.row_ids(meta), layer.ple_embedding.row_ids_reference(meta) + ) + + +@requires_cuda +@pytest.mark.parametrize("decode", [False, True]) +def test_fused_hash_matches_the_torch_reference(decode): + """The fused kernel is byte-identical to the torch-op hash, prefill and decode.""" + layer = _make_layer(_config(), device="cuda") + embedding = layer.ple_embedding + if decode: + meta = _meta( + [[s[0]] for s in _FUSED_HASH_SEQS], _FUSED_HASH_CTX, device="cuda", decode=True + ) + else: + meta = _meta(_FUSED_HASH_SEQS, _FUSED_HASH_CTX, device="cuda") + assert embedding._use_fused_row_ids(meta) + assert torch.equal(embedding.row_ids(meta), embedding.row_ids_reference(meta)) + + +@requires_cuda +def test_fused_hash_captures_and_replays_in_a_cuda_graph(): + """Fixed shapes, no host reads: the hash is part of the captured decode step.""" + layer = _make_layer(_config(), device="cuda") + embedding = layer.ple_embedding + ids = torch.randint(0, VOCAB, (4,), dtype=torch.int64, device="cuda") + context = torch.randint(0, VOCAB, (4, embedding.ngram_size - 1), dtype=torch.int64, + device="cuda") + meta = PLEMetadata( + input_ids=ids, + cu_seqlens=torch.arange(5, dtype=torch.int32, device="cuda"), + seq_lens=(1,) * 4, + ngram_context=context, + state_slots=torch.arange(4, dtype=torch.int64, device="cuda"), + fresh_slots=None, + is_decode=True, + ) + out = torch.zeros(4, embedding.num_heads, dtype=torch.int64, device="cuda") + embedding.row_ids(meta, out) # prime the index cache and the JIT before capture + + graph = torch.cuda.CUDAGraph() + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + embedding.row_ids(meta, out) + torch.cuda.current_stream().wait_stream(side) + with torch.cuda.graph(graph): + embedding.row_ids(meta, out) + + for seed in range(3): + torch.manual_seed(seed) + ids.copy_(torch.randint(0, VOCAB, (4,), dtype=torch.int64, device="cuda")) + context.copy_( + torch.randint(0, VOCAB, context.shape, dtype=torch.int64, device="cuda") + ) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(out, embedding.row_ids_reference(meta)), seed + + +@requires_cuda +def test_fused_hash_env_switch_restores_the_torch_path(monkeypatch): + """FREETOKEN_PLE_FUSED_HASH=0 takes the hash back to the reference ops.""" + layer = _make_layer(_config(), device="cuda") + meta = _meta(_FUSED_HASH_SEQS, _FUSED_HASH_CTX, device="cuda") + monkeypatch.setenv("FREETOKEN_PLE_FUSED_HASH", "0") + assert not layer.ple_embedding._use_fused_row_ids(meta) + assert torch.equal( + layer.ple_embedding.row_ids(meta), layer.ple_embedding.row_ids_reference(meta) + ) + + +@pytest.mark.parametrize( + "ctx_len, heads_per_ngram, message", + [(1, 8, "context"), (2, 5, "heads")], +) +def test_the_fused_hash_refuses_a_geometry_it_cannot_address(ctx_len, heads_per_ngram, message): + """The kernel's block layout (``heads_per_ngram`` heads per n-gram order, ``ngram_size-1`` + context ids) is checked BEFORE anything is launched, as a sentence -- not as an ``assert`` + that ``python -O`` strips, leaving the kernel to read past the context row.""" + from freetoken.kernel.triton.ple_hash import ple_row_ids + + tokens = 3 + with pytest.raises(ValueError, match=message): + ple_row_ids( + torch.arange(tokens, dtype=torch.int64), + torch.zeros(1, ctx_len, dtype=torch.int64), + torch.zeros(tokens, dtype=torch.int32), + torch.arange(tokens, dtype=torch.int32), + torch.tensor([3, 5, 7], dtype=torch.int64), # ngram_size 3 + torch.full((16,), 11, dtype=torch.int64), # 16 heads + torch.arange(16, dtype=torch.int64) * 11, + eos_token_id=EOS, + heads_per_ngram=heads_per_ngram, + ) + + + +def test_a_host_hashing_table_skips_the_device_hash(monkeypatch): + """The disk table hashes on the host and reads only the row_ids shape, so neither the + prefetch nor the lookup path runs the device hash for it.""" + from freetoken.models.qwen4_exp.ple import ZeroTable + from freetoken.models.qwen4_exp.ple_disk import DiskRowTable + + assert DiskRowTable.reads_row_ids is False + config = _config() + args = config.qwen4_args + seen = [] + + class HostHashedTable(ZeroTable): + reads_row_ids = False + + def lookup(self, row_ids, out=None): + seen.append(("lookup", tuple(row_ids.shape), row_ids.dtype)) + return super().lookup(row_ids, out) + + def prefetch(self, row_ids): + seen.append(("prefetch", tuple(row_ids.shape), row_ids.dtype)) + + layer = _make_layer(config, table=HostHashedTable(_padded_vocab(args), args.ngram_head_dim)) + embedding = layer.ple_embedding + + def no_hash(*a, **k): + raise AssertionError("the device hash ran for a table that never reads it") + + monkeypatch.setattr(embedding, "row_ids", no_hash) + sequences, contexts = [[3, 4, EOS, 5], [2, 7]], [[EOS, EOS], [21, 22]] + meta = _meta(sequences, contexts) + total = sum(len(s) for s in sequences) + states = torch.zeros(len(sequences), args.ple_state_width, args.ple_conv_state_len) + + layer.start_prefetch(None, meta) + _forward(layer, torch.randn(total, layer.hc_count * layer.hidden_size), meta, states) + embedding.forward(meta) + rows = (total, embedding.num_heads) + assert seen == [("prefetch", rows, torch.int64), ("lookup", rows, torch.int64), ("lookup", rows, torch.int64)] + # -------------------------------------------------------------------------------------- # table backends # -------------------------------------------------------------------------------------- diff --git a/tests/models/qwen4_exp/test_ple_disk_fill_order.py b/tests/models/qwen4_exp/test_ple_disk_fill_order.py new file mode 100644 index 0000000000..9d979b64eb --- /dev/null +++ b/tests/models/qwen4_exp/test_ple_disk_fill_order.py @@ -0,0 +1,254 @@ +"""The PLE disk backend's wait-sync fill must not depend on the graph launch returning. + +The captured decode graph WAITs on ``_flag``; ``cuGraphLaunch`` may block until the GPU drains +(a large graph fills the launch queue), and the GPU is parked on that WAIT. So the fill that sets +the flag has to run on its own thread, started before the launch (fork issue #65). +These tests stand in for the GPU with host fakes: the "launch" blocks until the flag is set. +""" + +from __future__ import annotations + +import queue +import threading +import time +from types import SimpleNamespace + +import pytest +import torch + +import freetoken.models.qwen4_exp.ple_disk as ple_disk +from freetoken.models.qwen4_exp.ple_disk import DiskRowTable + +EOS = 9 + + +class _Store: + """Stands in for the C++ PleStore: records calls and the thread that made them.""" + + def __init__(self, flag: torch.Tensor): + self.flag = flag + self.calls: list[tuple[str, list[int] | None, str]] = [] + + def stage(self, tokens_addr, n, staging_addr): + tokens = ctypes_read(tokens_addr, n + 2) + self.calls.append(("stage", tokens, threading.current_thread().name)) + + def flush(self, signal_addr): + self.calls.append(("flush", None, threading.current_thread().name)) + if signal_addr: + self.flag[0] = 1 + + +def ctypes_read(addr: int, n: int) -> list[int]: + import ctypes + + return list((ctypes.c_int64 * n).from_address(addr)) + + +class _Event: + def record(self, stream=None): + pass + + def synchronize(self): + pass + + +def _table(monkeypatch) -> DiskRowTable: + # host-only: no CUDA stream behind the (faked) readback event + monkeypatch.setattr(ple_disk.torch.cuda, "current_stream", lambda *a, **k: None) + monkeypatch.setattr(ple_disk.torch.cuda, "Event", _Event) + t = DiskRowTable.__new__(DiskRowTable) + t._device = None + t._wait_sync = True + t._flag = torch.zeros(1, dtype=torch.int64) + t._store = _Store(t._flag) + t._token_readback = torch.zeros(8, dtype=torch.int32) + t._graph_pinned = torch.zeros(64, dtype=torch.uint8) + t._eager_pinned = [torch.zeros(64, dtype=torch.uint8), torch.zeros(64, dtype=torch.uint8)] + t._eager_read = [_Event(), _Event()] + t._eager_slot = 0 + t._token_bytes = 1 + t.eos_token_id = EOS + t.image_token_id = None + return t + + +def _close(t: DiskRowTable) -> None: + close = getattr(t, "close", None) + if close is not None: + close() + + +def _decode_batch(prev_tokens: list[int], new_token: int): + req = SimpleNamespace(input_ids=torch.tensor(prev_tokens + [new_token]), device_len=len(prev_tokens) + 1) + return SimpleNamespace(is_decode=True, reqs=[req], padded_size=1, + input_ids=torch.tensor([new_token], dtype=torch.int32)) + + +def _blocking_launch(flag: torch.Tensor, timeout: float = 5.0) -> bool: + """A cuGraphLaunch that cannot return until the GPU passes the graph's WAIT(flag).""" + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if int(flag[0]) == 1: + flag[0] = 0 # the graph's WAIT-then-RESET + return True + time.sleep(0.001) + return False + + +def test_graph_fill_does_not_wait_for_the_launch_to_return(monkeypatch): + t = _table(monkeypatch) + with t.forward_host_ctx(_decode_batch([1, 2, 3], 7), use_graph=True): + assert _blocking_launch(t._flag), "deadlock: the fill only runs after the launch returns" + _close(t) + + +def test_fills_run_in_step_order_on_one_thread_with_contexts_taken_before_the_launch(monkeypatch): + t = _table(monkeypatch) + for prev, new in (([1, 2, 3], 7), ([1, 2, 3, 7], 8), ([1, 2, 3, 7, 8], 4)): + batch = _decode_batch(prev, new) + with t.forward_host_ctx(batch, use_graph=True): + # the engine mutates the request as soon as the launch returns (overlap scheduling) + batch.reqs[0].input_ids = torch.tensor([0, 0, 0, 0, 0, 0]) + batch.reqs[0].device_len += 1 + assert _blocking_launch(t._flag) + _close(t) + staged = [c for c in t._store.calls if c[0] == "stage"] + assert [c[1] for c in staged] == [[2, 3, 7], [3, 7, 8], [7, 8, 4]] + assert len({c[2] for c in t._store.calls}) == 1 # the store stays single-threaded + + +def test_eager_fill_completes_before_the_launch(monkeypatch): + t = _table(monkeypatch) + req = SimpleNamespace(input_ids=torch.tensor([1, 2, 3, 4]), cached_len=1, device_len=4) + batch = SimpleNamespace(is_decode=False, padded_reqs=[req]) + with t.forward_host_ctx(batch, use_graph=False): + assert [c[0] for c in t._store.calls] == ["stage", "flush"] + _close(t) + + +def test_a_failed_launch_leaves_no_stale_signal(monkeypatch): + t = _table(monkeypatch) + with pytest.raises(RuntimeError, match="launch failed"): + with t.forward_host_ctx(_decode_batch([1, 2, 3], 7), use_graph=True): + raise RuntimeError("launch failed") + # no WAIT consumed it, so a leftover 1 would let the next step read stale rows + assert int(t._flag[0]) == 0 + _close(t) + + +class _Gpu: + """One in-order stream: each launched graph parks on WAIT(flag), then RESETs it.""" + + def __init__(self, flag: torch.Tensor): + self.flag = flag + self.enqueued = self.completed = 0 + self.stalled = False + self.cv = threading.Condition() + self._launches: queue.SimpleQueue = queue.SimpleQueue() + threading.Thread(target=self._run, daemon=True).start() + + def launch(self) -> None: + self.enqueued += 1 + self._launches.put(True) + + def close(self) -> None: + self._launches.put(False) + + def _run(self) -> None: + while self._launches.get(): + self.stalled |= not _blocking_launch(self.flag) + with self.cv: + self.completed += 1 + self.cv.notify_all() + + def wait_completed(self, n: int, timeout: float = 5.0) -> bool: + with self.cv: + return self.cv.wait_for(lambda: self.completed >= n, timeout) + + +class _StreamEvent: + """cudaEvent semantics: synchronize() waits for the work enqueued before the LAST record.""" + + def __init__(self, gpu: _Gpu): + self.gpu, self.target = gpu, 0 + + def record(self, stream=None): + self.target = self.gpu.enqueued + + def synchronize(self): + if not self.gpu.wait_completed(self.target, timeout=1.0): + raise TimeoutError("the readback waits on a graph parked on this fill's own WAIT") + + +def _two_steps_behind_a_slow_filler(t: DiskRowTable, gpu: _Gpu) -> None: + """Steps N and N+1 both launch before the filler thread reaches fill N.""" + gate = threading.Event() + t._filler().submit(gate.wait) + for prev, new in (([1, 2, 3], 7), ([1, 2, 3, 7], 8)): + with t.forward_host_ctx(_decode_batch(prev, new), use_graph=True): + gpu.launch() + gate.set() + assert gpu.wait_completed(2), "decode deadlocked" + t._filler().submit(lambda: None).result() # every submitted fill has finished + gpu.close() + + +def test_each_fill_waits_on_its_own_steps_readback(monkeypatch): + """A readback event shared across steps is re-recorded by step N+1 behind graph N, so + fill N would wait on the graph that is parked on fill N's own WAIT.""" + t = _table(monkeypatch) + gpu = _Gpu(t._flag) + monkeypatch.setattr(ple_disk.torch.cuda, "Event", lambda: _StreamEvent(gpu)) + _two_steps_behind_a_slow_filler(t, gpu) + t._raise_failed_fill() + _close(t) + assert not gpu.stalled + # the fake D2H readback is not stream-ordered, so compare only the contexts + assert [c[1][:2] for c in t._store.calls if c[0] == "stage"] == [[2, 3], [3, 7]] + + +def test_a_failed_fill_is_raised_after_the_next_step_was_submitted(monkeypatch): + """Fill N fails (its graph runs on stale rows) after step N+1 already submitted: the error + must still reach the engine thread at the next launch.""" + t = _table(monkeypatch) + gpu = _Gpu(t._flag) + monkeypatch.setattr(ple_disk.torch.cuda, "Event", lambda: _StreamEvent(gpu)) + flush = t._store.flush + + def fail_step_n(signal_addr): + if t._store.calls[-1][1][:2] == [2, 3]: # the run staged just before is step N's + raise OSError("PLE row read failed") + flush(signal_addr) + + t._store.flush = fail_step_n + _two_steps_behind_a_slow_filler(t, gpu) + req = SimpleNamespace(input_ids=torch.tensor([1, 2, 3, 4]), cached_len=1, device_len=4) + with pytest.raises(OSError, match="PLE row read failed"): + t.host_fill_batch(SimpleNamespace(is_decode=False, padded_reqs=[req]), use_graph=False) + _close(t) + + +def test_the_readback_event_is_never_freed_on_the_filler_thread(monkeypatch): + """cuEventDestroy on the filler blocks on the driver lock a blocked cuGraphLaunch holds (host + nodes in the graph make the launch wait for the GPU, which is parked on the next fill's WAIT).""" + import gc + + freed_on: list[str] = [] + + class _TrackedEvent(_Event): + def __del__(self): + freed_on.append(threading.current_thread().name) + + t = _table(monkeypatch) + monkeypatch.setattr(ple_disk.torch.cuda, "Event", _TrackedEvent) + for prev, new in (([1, 2, 3], 7), ([1, 2, 3, 7], 8)): + with t.forward_host_ctx(_decode_batch(prev, new), use_graph=True): + assert _blocking_launch(t._flag) + time.sleep(0.05) # let the filler drop its work items + gc.collect() + t._raise_failed_fill() + gc.collect() + assert "ple-filler" not in freed_on, freed_on + assert len(t._readback_pool()) >= 1 # finished fills hand their events back for reuse + _close(t) diff --git a/tests/models/qwen4_exp/test_ple_disk_no_graph_fence.py b/tests/models/qwen4_exp/test_ple_disk_no_graph_fence.py new file mode 100644 index 0000000000..cb9bc32356 --- /dev/null +++ b/tests/models/qwen4_exp/test_ple_disk_no_graph_fence.py @@ -0,0 +1,136 @@ +"""The wait-sync graph fill must never wait on an event recorded after its own launch. + +Upstream #588 fences ``_graph_pinned`` with ``_graph_consumed``, recorded once the launch returns and +synchronized by the next graph fill. Upstream runs the fill after the launch returns; here (fork +#65/#67) the fill runs on the filler thread BEFORE its own graph can pass its WAIT, so that fence +makes fill k wait on graph k, which waits on fill k: decode hangs on the second overlapped step. +The previous graph's read of ``_graph_pinned`` is already ordered before the next fill by the +per-fill readback event. This test drives the overlapped decode loop against an asynchronous, +in-order fake stream (launch returns at once; events record a stream position) and fails if any +step stalls, including when the fence is reintroduced. +""" + +from __future__ import annotations + +import threading +import time +from types import SimpleNamespace + +import torch + +import freetoken.models.qwen4_exp.ple_disk as ple_disk +from freetoken.models.qwen4_exp.ple_disk import DiskRowTable + +EOS = 9 + + +class _Stream: + """One in-order stream: a graph item WAITs on the flag, then RESETs it; launch is async.""" + + def __init__(self, flag: torch.Tensor): + self.flag = flag + self.cv = threading.Condition() + self.queue: list[str] = [] + self.enqueued = self.done = 0 + self.stop = False + threading.Thread(target=self._run, daemon=True).start() + + def enqueue(self, kind: str) -> None: + with self.cv: + self.queue.append(kind) + self.enqueued += 1 + self.cv.notify_all() + + def _run(self) -> None: + while True: + with self.cv: + self.cv.wait_for(lambda: self.queue or self.stop) + if self.stop: + return + kind = self.queue.pop(0) + if kind == "graph": + deadline = time.monotonic() + 10.0 + while int(self.flag[0]) != 1 and time.monotonic() < deadline: + time.sleep(0.0005) + self.flag[0] = 0 + with self.cv: + self.done += 1 + self.cv.notify_all() + + def wait_done(self, n: int, timeout: float) -> bool: + with self.cv: + return self.cv.wait_for(lambda: self.done >= n, timeout) + + def close(self) -> None: + with self.cv: + self.stop = True + self.cv.notify_all() + + +class _Event: + """cudaEvent semantics: synchronize() waits for the work enqueued before the LAST record().""" + + def __init__(self, stream: _Stream): + self.stream, self.pos = stream, 0 + + def record(self, _stream=None) -> None: + self.pos = self.stream.enqueued + + def synchronize(self) -> None: + if not self.stream.wait_done(self.pos, 2.0): + raise TimeoutError(f"a fill waited on stream position {self.pos}, which never completed") + + +class _Store: + def __init__(self, flag: torch.Tensor): + self.flag = flag + + def stage(self, tokens_addr, n, staging_addr): + pass + + def flush(self, signal_addr): + if signal_addr: + self.flag[0] = 1 + + +def test_overlapped_graph_fills_never_wait_on_their_own_launch(monkeypatch): + flag = torch.zeros(1, dtype=torch.int64) + stream = _Stream(flag) + monkeypatch.setattr(ple_disk.torch.cuda, "current_stream", lambda *a, **k: None) + monkeypatch.setattr(ple_disk.torch.cuda, "Event", lambda *a, **k: _Event(stream)) + t = DiskRowTable.__new__(DiskRowTable) + t._device = None + t._wait_sync = True + t._flag = flag + t._store = _Store(flag) + t._token_readback = torch.zeros(8, dtype=torch.int32) + t._graph_pinned = torch.zeros(64, dtype=torch.uint8) + t._eager_pinned = [torch.zeros(64, dtype=torch.uint8), torch.zeros(64, dtype=torch.uint8)] + t._eager_read = [_Event(stream), _Event(stream)] + t._eager_slot = 0 + # if upstream's post-launch fence is ever reintroduced, model it with real stream semantics + # so this test fails by hanging the way production would, not on a missing attribute + t._graph_consumed = _Event(stream) + t._token_bytes = 1 + t.eos_token_id = EOS + t.image_token_id = None + + tokens = [1, 2, 3] + steps = 12 + try: + for step in range(steps): + new = 10 + step + req = SimpleNamespace(input_ids=torch.tensor(tokens + [new]), device_len=len(tokens) + 1) + batch = SimpleNamespace(is_decode=True, reqs=[req], padded_size=1, + input_ids=torch.tensor([new], dtype=torch.int32)) + with t.forward_host_ctx(batch, use_graph=True): + stream.enqueue("graph") # cuGraphLaunch returns before the GPU passes the WAIT + tokens.append(new) + # the overlap loop keeps the engine at most one launch ahead of the GPU + assert stream.wait_done(step, 4.0), f"decode hung: graph {step - 1} never completed" + t._raise_failed_fill() + assert stream.wait_done(steps, 4.0), "decode hung on the last step" + t._raise_failed_fill() + finally: + t.close() + stream.close() diff --git a/tests/models/qwen4_exp/test_qsa_backend.py b/tests/models/qwen4_exp/test_qsa_backend.py index 236610791d..8e836f5632 100644 --- a/tests/models/qwen4_exp/test_qsa_backend.py +++ b/tests/models/qwen4_exp/test_qsa_backend.py @@ -4,8 +4,12 @@ tokens every complete block is selected, so QSA IS dense attention: the selection must be exactly the causal prefix and the layer output must match ``TorchDenseQSAReference`` (fp32) and a flashinfer dense run over the same pool; -(b) chunked prefill at unaligned cut points equals one-shot prefill (the dual-source compress); -(c) a captured decode replay equals the eager decode step. +(b) chunked prefill at unaligned cut points matches one-shot prefill to bf16 tolerance (the + dual-source compress); +(c) a captured decode replay equals the eager decode step; +(d) an fp8 KV pool (``--kv-cache-dtype fp8``) keeps block selection bit-identical to the + 16-bit run -- only the selected K/V rows are read back as e4m3 codes -- and the layer + output stays within quantization error of it. """ from __future__ import annotations @@ -133,7 +137,11 @@ def test_flashinfer_dense_matches_the_sparse_path(): @requires_cuda @pytest.mark.parametrize("cut", [1001, 4096, 4097], ids=["unaligned", "page-boundary", "boundary+1"]) def test_chunked_prefill_matches_one_shot(cut: int): - """Cut points that are not multiples of index_ratio exercise the dual-source compress.""" + """Cut points that are not multiples of index_ratio exercise the dual-source compress. + + Not bit for bit: the projections see a different M (5000 vs 5000 - cut rows) and cuBLAS + may pick a different kernel per shape, so the indexer's q and k round differently and a + few rows flip one block at the top-k margin. Same tolerance as the oracle tests above.""" config = parsed_config() fixture = Fixture(config, num_pages=512) attn = fixture.layer(QSA_LAYER) @@ -145,7 +153,19 @@ def test_chunked_prefill_matches_one_shot(cut: int): attn.forward(x[:cut], fixture.batch([head], "prefill")) tail = fixture.req(1, cut, length) got = attn.forward(x[cut:], fixture.batch([tail], "prefill")) - assert torch.equal(got, one_shot[cut:]) + torch.testing.assert_close(got.float(), one_shot[cut:].float(), rtol=2e-2, atol=2e-2) + + # A wrong key for the group straddling the cut only swaps one block in a few rows, which + # the output tolerance hides; the unit-scale slab rows show it as an O(1) error. + ratio = config.qwen4_args.index_ratio + slab = fixture.pool.cmp_k_cache(0) + group_starts = slice(0, length - length % ratio, ratio) + one_shot_rows, chunked_rows = ( + (fixture.page_table[table_idx, group_starts] // ratio).long() for table_idx in (0, 1) + ) + torch.testing.assert_close( + slab[chunked_rows].float(), slab[one_shot_rows].float(), rtol=2e-2, atol=2e-2 + ) @requires_cuda @@ -249,3 +269,64 @@ def test_two_qsa_layers_keep_separate_slab_slots(monkeypatch): slab = fixture.pool.cmp_k_cache assert not torch.equal(slab(0), slab(1)) + + +def _prefill_under_kv(monkeypatch, config, kv_quant: str, lengths): + """One prefill of the QSA layer under a given KV store. + + Each call builds its own Fixture on purpose: a Fixture owns the global ctx (pool, + page table, backend), so two KV stores cannot share one scenario. The weight seed + (``Fixture.layer``) and the input seed (``_inputs``) are fixed, so the two runs differ + ONLY in how the K/V rows are stored. + """ + fixture = Fixture(config, num_pages=128, kv_quant=kv_quant) + attn = fixture.layer(QSA_LAYER) + seen = selection_spy(monkeypatch, fixture.backend) + inputs = _inputs(fixture, lengths) + x = torch.cat([row[:n] for row, n in zip(inputs, lengths)]) + reqs = [fixture.req(i, 0, n) for i, n in enumerate(lengths)] + batch = fixture.batch(reqs, "prefill") + out = attn.forward(x, batch) + # the selection lives in a scratch buffer the next forward overwrites + return fixture, out.clone(), seen["indices"].clone(), batch.positions.clone() + + +@requires_cuda +def test_fp8_kv_pool_keeps_selection_and_output(monkeypatch): + """--kv-cache-dtype fp8 through the real layer: e4m3 codes + per-row scales in, same + answer out to within quantization error -- and, because block selection scores 16-bit + compressed index keys that fp8 never touches, the SAME selection bit for bit.""" + config = parsed_config() + lengths = [2051, 1000, 137] # every complete block is selected here + + plain, plain_out, plain_idx, _ = _prefill_under_kv(monkeypatch, config, "none", lengths) + quant, quant_out, quant_idx, positions = _prefill_under_kv( + monkeypatch, config, "fp8", lengths + ) + + # The tripwire for the field failure: the backend sizes its indexer scratch with + # pool.dtype, which must stay the COMPUTE dtype even when store_dtype is e4m3. An + # fp8 q_index compiles into qsa_mqa_paged's dot and dies at graph capture. + assert quant.backend.dtype is torch.bfloat16 + assert plain.backend.dtype is torch.bfloat16 + assert quant.pool.store_dtype != torch.bfloat16 + assert quant.pool.kv_quant == "fp8" and plain.pool.kv_quant == "none" + assert quant.pool.k_cache(QSA_LAYER).element_size() == 1 + assert quant.pool.v_cache(QSA_LAYER).element_size() == 1 + assert plain.pool.k_scale(QSA_LAYER) is None and plain.pool.v_scale(QSA_LAYER) is None + pages, page_size, kv_heads = quant.pool.k_cache(QSA_LAYER).shape[:3] + assert quant.pool.k_scale(QSA_LAYER).shape == (pages * page_size, kv_heads) + assert quant.pool.k_scale(QSA_LAYER).dtype is torch.float32 + + for pool in (plain.pool, quant.pool): + assert pool.cmp_k_cache(0).dtype is torch.bfloat16 + assert torch.equal(quant_idx, plain_idx), ( + "quantizing the KV rows changed which blocks the indexer selected -- the index " + "tier is supposed to be 16-bit in both runs" + ) + _assert_selection_is_causal_prefix(quant_idx, positions) + + # Looser than the 2e-2 the 16-bit run needs against the same reference: e4m3 carries + # four significant bits, so ~1e-2 relative per stored element is the floor here. + torch.testing.assert_close(quant_out.float(), plain_out.float(), rtol=4e-2, atol=4e-2) + diff --git a/tests/models/qwen4_exp/test_skeleton.py b/tests/models/qwen4_exp/test_skeleton.py index 8fb9112df2..0803cbca28 100644 --- a/tests/models/qwen4_exp/test_skeleton.py +++ b/tests/models/qwen4_exp/test_skeleton.py @@ -432,6 +432,2041 @@ def test_shared_gate_kernels_match_torch(num_tokens, hidden, dtype): assert (fused.float() - ref).abs().max() <= (eager.float() - ref).abs().max() + 1e-6 +def _offload_experts(kind: str, layer_id: int, num_experts: int, top_k: int, hidden: int, inter: int): + from freetoken.layers.moe import OffloadMoELayer + from freetoken.layers.quantization import NoQuantConfig, QuantBackend, QuantConfig, set_quant_backend + + if kind == "bf16": + quant = NoQuantConfig() + else: + set_quant_backend(QuantBackend.parse("moe.nvfp4=triton")) # the sm_120 decode kernel + quant = QuantConfig.from_hf({"quantization_config": {"quant_method": "modelopt", "quant_algo": "NVFP4", "ignore": ["lm_head"]}}) + return OffloadMoELayer( + layer_id, num_experts, top_k, hidden, inter, quant_config=quant, + prefix=f"model.layers.{layer_id}.mlp.experts", + ) + + +def _copy_overlap_run(kind, banks, overlap, bs, graph, inputs, monkeypatch): + """Two Qwen4ExpMoE layers sharing one offload cache as small as a layer, so every step evicts + and copies: warm up eagerly (and capture), reset the cache, then decode ``inputs``.""" + import freetoken.kernel.fast_index_copy as fic + from flashlib.kernels.slot_cache import Stat + from freetoken.models.qwen4_exp.moe import Qwen4ExpMoE + from freetoken.moe.offload_cache import OffloadMoeCache + from freetoken.utils.torch_utils import torch_dtype + + num_experts, top_k, hidden, inter = 8, 2, 256, 128 + config = parse_config(toy_hf_config( + hidden_size=hidden, num_experts=num_experts, num_experts_per_tok=top_k, + moe_intermediate_size=inter, shared_expert_intermediate_size=inter, + )) + device = torch.device("cuda") + with torch.device(device), torch_dtype(torch.bfloat16): + moes = [Qwen4ExpMoE(config, layer_id) for layer_id in range(2)] + gen = torch.Generator(device=device).manual_seed(41) + for moe in moes: + _fill(moe, gen, scale=0.2) + experts = [_offload_experts(kind, i, num_experts, top_k, hidden, inter) for i in range(2)] + cache = OffloadMoeCache( + num_layers=2, num_experts=num_experts, cache_size=num_experts, device=device, + quant_format=banks.quant_format, layout=banks.layout, + max_slots=experts[0].quant_method.slot_limit(), decode_copy_overlap=overlap, + ) + cache.set_bank_sources(banks.sources) + cache.set_alphas(banks.gate_up_alpha, banks.down_alpha) + cache.collect_stats = True + for moe, ex in zip(moes, experts): + ex.offload_cache = cache + moe.experts = ex + cache.reset() + + import freetoken.models.qwen4_exp.moe as qmoe + + copy_streams, gate_streams = [], [] + real_gate = qmoe.shared_gate_sigmoid + for name in ("fast_index_copy_multi_jit", "fast_index_copy_multi_slim_jit"): # FREETOKEN_MOE_SLIM_COPY + real_copy = getattr(fic, name) + monkeypatch.setattr( + fic, name, + lambda *a, _real=real_copy, **k: (copy_streams.append(torch.cuda.current_stream().cuda_stream), _real(*a, **k))[1], + ) + monkeypatch.setattr( + qmoe, "shared_gate_sigmoid", + lambda *a, **k: (gate_streams.append(torch.cuda.current_stream().cuda_stream), real_gate(*a, **k))[1], + ) + + def step(x0, x1): + return moes[0].forward(x0.clone()), moes[1].forward(x1.clone()) + + _fresh_ctx(_batch=SimpleNamespace(is_prefill=False)) + static = [torch.randn(bs, hidden, device=device, dtype=torch.bfloat16) for _ in range(2)] + step(*static) # the eager warm-up the graph runner also does + if graph: + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + outs = step(*static) + cache.reset() + cache.reset_stats() + results = [] + for x0, x1 in inputs: + if graph: + static[0].copy_(x0) + static[1].copy_(x1) + g.replay() + else: + outs = step(x0, x1) + results.append(tuple(o.clone() for o in outs)) + torch.cuda.synchronize() + state = [t.clone() for t in (cache.slot_for_id, cache.id_of_slot, cache.usage, *cache.bank_views())] + misses = int(cache.lru_stats[:, Stat.MISS].sum()) + side = getattr(cache.decode_copy_stream, "cuda_stream", None) + return results, state, misses, [s == side for s in copy_streams], [s == side for s in gate_streams] + + +@requires_cuda +@pytest.mark.parametrize("kind", ["bf16", "nvfp4"]) +@pytest.mark.parametrize("graph", [False, True], ids=["eager", "graph"]) +@pytest.mark.parametrize("bs", [1, 2]) +def test_moe_copy_overlap_is_bitwise_identical(kind, graph, bs, monkeypatch): + """FREETOKEN_MOE_COPY_OVERLAP moves the shared expert onto a side stream beside the decode miss + copy; the MoE outputs, the slot map and the slot contents must not change by a single bit.""" + from freetoken.moe.expert_banks import build_expert_banks + + torch.manual_seed(7) + method = _offload_experts(kind, 0, 8, 2, 256, 128).quant_method + banks = build_expert_banks(method, 2, None, device=torch.device("cuda"), dummy=True) + gen = torch.Generator(device="cuda").manual_seed(bs) + inputs = [ + tuple(torch.randn(bs, 256, device="cuda", dtype=torch.bfloat16, generator=gen) * 0.5 for _ in range(2)) + for _ in range(6) + ] + off = _copy_overlap_run(kind, banks, False, bs, graph, inputs, monkeypatch) + on = _copy_overlap_run(kind, banks, True, bs, graph, inputs, monkeypatch) + + assert off[2] > 0 and on[2] == off[2], "the sequence must miss, and copy, in both runs" + assert off[3] and not any(off[3]) and not any(off[4]), "flag off: everything stays on the compute stream" + assert on[3] and not any(on[3]), "flag on: the miss copy stays right behind its ensure" + assert on[4] and all(on[4]), "flag on: the shared expert runs on the cache's side stream" + for got, want in zip(on[0], off[0]): + for a, b in zip(got, want): + assert torch.isfinite(b.float()).all() + assert torch.equal(a, b) + for a, b in zip(on[1], off[1]): + assert torch.equal(a, b) + + +# FREETOKEN_MOE_PREFETCH stack: 4 layers (full attention on layer 3), 16 experts, top-2, and one +# shared cache as small as a layer so every step evicts +_PF_LAYERS, _PF_EXPERTS, _PF_TOPK, _PF_HIDDEN, _PF_INTER = 4, 16, 2, 256, 128 +# ~50-100 us between layers: the predictor must finish before the next ensure even at toy size +_PF_SLEEP_CYCLES = 200_000 + + +def _prefetch_banks(kind): + from freetoken.moe.expert_banks import build_expert_banks + + torch.manual_seed(7) + method = _offload_experts(kind, 0, _PF_EXPERTS, _PF_TOPK, _PF_HIDDEN, _PF_INTER).quant_method + return build_expert_banks(method, _PF_LAYERS, None, device=torch.device("cuda"), dummy=True) + + +def _prefetch_stack(kind, banks, mode, monkeypatch, *, k=6, budget=0, wire=True, overlap=False, policy="rule", + cache_size=_PF_EXPERTS, **cache_kw): + from freetoken.env import ENV + from freetoken.models.qwen4_exp.moe import Qwen4ExpMoE, wire_router_lookahead + from freetoken.moe.offload_cache import OffloadMoeCache + from freetoken.utils.torch_utils import torch_dtype + + monkeypatch.setattr(ENV.MOE_PREFETCH_K, "value", k) + monkeypatch.setattr(ENV.MOE_PREFETCH_BUDGET, "value", budget) + config = parse_config(toy_hf_config( + _PF_LAYERS, hidden_size=_PF_HIDDEN, num_experts=_PF_EXPERTS, num_experts_per_tok=_PF_TOPK, + moe_intermediate_size=_PF_INTER, shared_expert_intermediate_size=_PF_INTER, + )) + device = torch.device("cuda") + with torch.device(device), torch_dtype(torch.bfloat16): + moes = [Qwen4ExpMoE(config, layer_id) for layer_id in range(_PF_LAYERS)] + gen = torch.Generator(device=device).manual_seed(43) + for moe in moes: + _fill(moe, gen, scale=0.2) + experts = [_offload_experts(kind, i, _PF_EXPERTS, _PF_TOPK, _PF_HIDDEN, _PF_INTER) for i in range(_PF_LAYERS)] + cache = OffloadMoeCache( + num_layers=_PF_LAYERS, num_experts=_PF_EXPERTS, cache_size=cache_size, device=device, + cache_policy=policy, quant_format=banks.quant_format, layout=banks.layout, + max_slots=experts[0].quant_method.slot_limit(), prefetch_mode=mode, decode_copy_overlap=overlap, + **cache_kw, + ) + cache.set_bank_sources(banks.sources) + cache.set_alphas(banks.gate_up_alpha, banks.down_alpha) + cache.collect_stats = True + for moe, ex in zip(moes, experts): + ex.offload_cache = cache + moe.experts = ex + if wire: + wire_router_lookahead(moes, config) + cache.reset() + return moes, cache + + +def _prefetch_run(moes, cache, bs, graph, inputs, before_decode=None, sleep=_PF_SLEEP_CYCLES, before_step=None): + """Warm up eagerly (and capture), reset the cache, then decode ``inputs`` (per step, one + ``[bs, hidden]`` input per layer); ``sleep`` cycles between layers, ``before_step(i)`` before step i.""" + from flashlib.kernels.slot_cache import Stat + + _fresh_ctx(_batch=SimpleNamespace(is_prefill=False)) + + def step(xs): + outs = [] + for layer, (moe, x) in enumerate(zip(moes, xs)): + if layer and sleep: + torch.cuda._sleep(sleep) + outs.append(moe.forward(x.clone())) + return outs + + static = [torch.randn(bs, _PF_HIDDEN, device="cuda", dtype=torch.bfloat16) for _ in moes] + step(static) # the eager warm-up the graph runner also does + if graph: + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + outs = step(static) + cache.reset() + cache.reset_stats() + if before_decode is not None: + before_decode() + results = [] + for i, xs in enumerate(inputs): + if before_step is not None: + before_step(i) + if graph: + for s, x in zip(static, xs): + s.copy_(x) + g.replay() + else: + outs = step(xs) + if cache.audit is not None: + cache.audit.end_decode_step(cache) # where the engine calls it: after the forward + results.append([o.clone() for o in outs]) + torch.cuda.synchronize() + state = [t.clone() for t in (cache.slot_for_id, cache.id_of_slot, cache.usage, *cache.bank_views())] + counters = None if cache.prefetch is None else cache.prefetch.counters.cpu() + return results, state, counters, cache.lru_stats[:, Stat.MISS].cpu() + + +def _prefetch_inputs(bs, steps=8, chain=False, seed=0): + gen = torch.Generator(device="cuda").manual_seed(seed * 10 + bs) + + def x(): + return torch.randn(bs, _PF_HIDDEN, device="cuda", dtype=torch.bfloat16, generator=gen) * 0.5 + + if chain: # every layer sees the same input, so the lookahead equals the next router + return [[x()] * _PF_LAYERS for _ in range(steps)] + return [[x() for _ in range(_PF_LAYERS)] for _ in range(steps)] + + +@requires_cuda +@pytest.mark.parametrize("kind", ["bf16", "nvfp4"]) +@pytest.mark.parametrize("graph", [False, True], ids=["eager", "graph"]) +@pytest.mark.parametrize("bs", [1, 2]) +@pytest.mark.parametrize("overlap,policy", [(False, "rule"), (True, "rule"), (True, "lru")], ids=["serial", "overlap", "overlap-lru"]) +def test_moe_prefetch_measure_is_bitwise_identical(kind, graph, bs, overlap, policy, monkeypatch): + """FREETOKEN_MOE_PREFETCH=measure predicts and counts beside the decode (with and without + FREETOKEN_MOE_COPY_OVERLAP's side stream); the MoE outputs, the slot map, the usage clock and + the slot contents must not change by a single bit.""" + banks = _prefetch_banks(kind) + inputs = _prefetch_inputs(bs) + off = _prefetch_run(*_prefetch_stack(kind, banks, "off", monkeypatch, overlap=overlap, policy=policy), bs, graph, inputs) + on = _prefetch_run(*_prefetch_stack(kind, banks, "measure", monkeypatch, overlap=overlap, policy=policy), bs, graph, inputs) + + assert off[2] is None and int(off[3].sum()) > 0 and torch.equal(on[3], off[3]) + for got, want in zip(on[0], off[0]): + for a, b in zip(got, want): + assert torch.isfinite(b.float()).all() + assert torch.equal(a, b) + for a, b in zip(on[1], off[1]): + assert torch.equal(a, b) + from freetoken.moe.prefetch import CALLS, ISSUED, ROWS + + counters = on[2] + assert counters[0].tolist() == [0] * counters.shape[1], "layer 0 has no predecessor" + assert counters[1:, CALLS].tolist() == [len(inputs)] * (_PF_LAYERS - 1) + assert counters[1:, ROWS].tolist() == [len(inputs) * bs] * (_PF_LAYERS - 1) + assert int(counters[1:, ISSUED].sum()) > 0 + + +def _reference_counters(selects, ensures, budget_override, k): + """CPU counters from the spied predictor logits and the residency its select read, and each + target layer's routing and residency at its ensure.""" + from freetoken.moe.prefetch import NUM_COLS, default_budget + from tests.moe.ref_prefetch import ref_count, ref_select + + config = parse_config(toy_hf_config(_PF_LAYERS)) + counters = torch.zeros((_PF_LAYERS, NUM_COLS), dtype=torch.int64) + for layer in range(1, _PF_LAYERS): + assert len(selects[layer]) == len(ensures[layer]) > 0 + want_budget = budget_override or default_budget(config.is_linear_layer(layer)) + for (logits, budget, seen), (raw, resident) in zip(selects[layer], ensures[layer]): + assert budget == want_budget + sel, res = ref_select(logits.tolist(), seen.tolist(), k, budget) + routed = raw.view(-1).tolist() + misses = len({e for e in routed if resident[e] < 0}) + counters[layer] += torch.tensor(ref_count(sel, res, routed, misses, logits.shape[0])) + return counters + + +@requires_cuda +@pytest.mark.parametrize("kind", ["bf16", "nvfp4"]) +@pytest.mark.parametrize("bs", [1, 2]) +@pytest.mark.parametrize("budget", [0, 5]) +def test_moe_prefetch_counters_match_reference(kind, bs, budget, monkeypatch): + """Eager counters equal a CPU reference built from the predictor's own logits and the residency + its select read, and each layer's routing and pre-ensure residency; the captured graph counts + the same. The predictor runs on its dedicated stream, the count on the compute stream.""" + import freetoken.moe.prefetch as pf_mod + + banks = _prefetch_banks(kind) + inputs = _prefetch_inputs(bs, seed=1) + k = 6 + moes, cache = _prefetch_stack(kind, banks, "measure", monkeypatch, k=k, budget=budget) + selects = {layer: [] for layer in range(_PF_LAYERS)} + ensures = {layer: [] for layer in range(_PF_LAYERS)} + streams = {"select": set(), "count": set()} + real_select, real_count, real_ensure = pf_mod.lookahead_select, pf_mod.prefetch_count, cache.ensure_experts + + def spy_select(logits, resident, sel, res, *, k, budget): + # the select reads residency beside layer L-1's ensure: syncing on both sides pins what it saw + torch.cuda.synchronize() + layer = (sel.data_ptr() - cache.prefetch.sel.data_ptr()) // cache.prefetch.sel[0].nbytes + selects[layer].append((logits.float().cpu(), budget, resident.cpu())) + streams["select"].add(torch.cuda.current_stream().cuda_stream) + real_select(logits, resident, sel, res, k=k, budget=budget) + torch.cuda.synchronize() + + def spy_count(*args, **kwargs): + streams["count"].add(torch.cuda.current_stream().cuda_stream) + return real_count(*args, **kwargs) + + def spy_ensure(layer_id, expert_ids, **kwargs): + torch.cuda.synchronize() # the prediction forked after the previous ensure has landed + ensures[layer_id].append((expert_ids.cpu(), cache.slot_for_id[layer_id].cpu())) + return real_ensure(layer_id, expert_ids, **kwargs) + + def install(): + monkeypatch.setattr(pf_mod, "lookahead_select", spy_select) + monkeypatch.setattr(pf_mod, "prefetch_count", spy_count) + monkeypatch.setattr(cache, "ensure_experts", spy_ensure) + + eager = _prefetch_run(moes, cache, bs, False, inputs, before_decode=install) + want = _reference_counters(selects, ensures, budget, k) + assert torch.equal(eager[2], want), (eager[2], want) + assert int(want[:, 1].sum()) > 0, "the sequence must have useful predictions" + assert streams["select"] == {cache.prefetch.stream.cuda_stream} + assert streams["count"] == {torch.cuda.current_stream().cuda_stream} + + monkeypatch.undo() + moes, cache = _prefetch_stack(kind, banks, "measure", monkeypatch, k=k, budget=budget) + graph = _prefetch_run(moes, cache, bs, True, inputs) + assert torch.equal(graph[2], eager[2]) + for got, want_out in zip(graph[0], eager[0]): + for a, b in zip(got, want_out): + assert torch.equal(a, b) + + +@requires_cuda +@pytest.mark.parametrize("graph", [False, True], ids=["eager", "graph"]) +@pytest.mark.parametrize("bs", [1, 2]) +def test_moe_prefetch_perfect_lookahead_counts_every_miss(graph, bs, monkeypatch): + """Known answer: when layer L-1's router input equals layer L's, the lookahead is layer L's own + router, so with the budget covering every candidate each demand miss is a useful prediction.""" + from freetoken.moe.prefetch import ISSUED, MISSES, RESIDENT_HITS, USEFUL + + banks = _prefetch_banks("nvfp4") + moes, cache = _prefetch_stack("nvfp4", banks, "measure", monkeypatch, k=4, budget=8) + _, _, counters, lru_misses = _prefetch_run(moes, cache, bs, graph, _prefetch_inputs(bs, chain=True, seed=2)) + assert int(counters[1:, MISSES].sum()) > 0 + assert counters[1:, USEFUL].tolist() == counters[1:, MISSES].tolist() == lru_misses[1:].tolist() + assert (counters[1:, ISSUED] >= counters[1:, USEFUL]).all() + assert int(counters[1:, RESIDENT_HITS].sum()) > 0, "routed hits among the top candidates" + + +@requires_cuda +def test_moe_prefetch_off_launches_nothing_new(monkeypatch): + """With the flag off the wired lookahead is inert: no prefetcher, and the decode launches + exactly the kernels of an unwired stack (measure adds some, so the check can see them).""" + from torch.profiler import ProfilerActivity, profile + + banks = _prefetch_banks("bf16") + inputs = _prefetch_inputs(1, steps=2) + + def kernels(mode, wire): + moes, cache = _prefetch_stack("bf16", banks, mode, monkeypatch, wire=wire) + _prefetch_run(moes, cache, 1, False, inputs[:1]) # warm up and compile outside the trace + with profile(activities=[ProfilerActivity.CUDA]) as prof: + _prefetch_run(moes, cache, 1, False, inputs) + events = [e for e in prof.events() if e.device_type == torch.autograd.DeviceType.CUDA] + return cache, [e.name for e in sorted(events, key=lambda e: e.time_range.start)] + + cache, wired_off = kernels("off", True) + assert cache.prefetch is None + _, unwired_off = kernels("off", False) + _, measured = kernels("measure", True) + assert wired_off == unwired_off + assert len(measured) > len(wired_off) + + +@requires_cuda +def test_moe_prefetch_survives_a_rebuild(monkeypatch): + """A runtime rebuild resets the counters and keeps the predictor stream; a recaptured graph + counts like an eager run on the rebuilt cache.""" + banks = _prefetch_banks("bf16") + inputs = _prefetch_inputs(1, seed=3) + moes, cache = _prefetch_stack("bf16", banks, "measure", monkeypatch) + before = _prefetch_run(moes, cache, 1, True, inputs)[2] + stream = cache.prefetch.stream + assert int(before.sum()) > 0 + cache.rebuild(_PF_EXPERTS + 4) + assert int(cache.prefetch.counters.abs().sum()) == 0 and cache.prefetch.stream is stream + graph = _prefetch_run(moes, cache, 1, True, inputs) + eager = _prefetch_run(moes, cache, 1, False, inputs) + assert int(graph[2].sum()) > 0 and torch.equal(graph[2], eager[2]) + for a, b in zip(graph[1], eager[1]): + assert torch.equal(a, b) + + +@requires_cuda +@pytest.mark.parametrize("kind", ["bf16", "nvfp4"]) +@pytest.mark.parametrize("bs", [1, 2]) +def test_moe_prefetch_predicts_layer_l_from_layer_l_minus_1_input(kind, bs, monkeypatch): + """With a distinct input per layer, the prediction for layer L reads exactly layer L-1's router + input through layer L's router, and the select sees those logits and layer L's residency.""" + import freetoken.moe.prefetch as pf_mod + + banks = _prefetch_banks(kind) + inputs = _prefetch_inputs(bs, steps=3, seed=4) + moes, cache = _prefetch_stack(kind, banks, "measure", monkeypatch) + forks, selects = [], [] + real_fork, real_select = cache.prefetch.fork, pf_mod.lookahead_select + + def spy_fork(target, x, gate, resident, budget): + forks.append((target, x.clone(), gate, resident.data_ptr())) + return real_fork(target, x, gate, resident, budget) + + def spy_select(logits, resident, sel, res, *, k, budget): + selects.append(logits.clone()) + return real_select(logits, resident, sel, res, k=k, budget=budget) + + def install(): + monkeypatch.setattr(cache.prefetch, "fork", spy_fork) + monkeypatch.setattr(pf_mod, "lookahead_select", spy_select) + + _prefetch_run(moes, cache, bs, False, inputs, before_decode=install) + assert [f[0] for f in forks] == list(range(1, _PF_LAYERS)) * len(inputs) and len(selects) == len(forks) + for call, ((target, x, gate, resident), logits) in enumerate(zip(forks, selects)): + assert torch.equal(x, inputs[call // (_PF_LAYERS - 1)][target - 1]) + assert gate is moes[target].gate and resident == cache.slot_for_id[target].data_ptr() + assert torch.equal(logits, F.linear(x, moes[target].gate.weight)) + + +# FREETOKEN_MOE_PREFETCH=on: the same stack, with copies. No sleep between layers, so the prefetch +# streams race the compute stream at toy size and only the joins keep the GEMMs on landed bytes. +_PF_ON_STEPS = 200 + + +def _assert_slots_hold_their_experts(cache): + """Lossless: every held slot's bytes equal the host rows of the expert its id names.""" + held = (cache.id_of_slot >= 0).nonzero().view(-1).cpu() + ids = cache.id_of_slot[held.cuda()].long().cpu() + assert held.numel() > 0 + for (per_layer, slots), view in zip(cache.banks, cache.bank_views()): + host = torch.stack(per_layer).flatten(0, 1) + want = host[ids].contiguous().view(torch.uint8) + got = view[held.cuda()].cpu().contiguous().view(torch.uint8) + assert torch.equal(got, want) + + +def _routed_ids(moes, inputs): + """[step][layer] -> the experts each layer routes (the router alone decides; no cache input).""" + from freetoken.moe.fused import fused_topk + + out = [] + for xs in inputs: + row = [] + for moe, x in zip(moes, xs): + _, ids = fused_topk(hidden_states=x, gating_output=moe.gate.forward(x), topk=_PF_TOPK, renormalize=True) + row.append(ids.view(-1).tolist()) + out.append(row) + return out + + +def _garbage(routed, width, kind, seed): + """Per step a [layers, width] candidate override: 'random' ids (and -1), 'adversarial' ones + that mix layer L's routed ids with ids it does not route, duplicates and padding, or 'flood': + every expert of layer L, so a prefetch wants more slots than the unpinned ones.""" + gen = torch.Generator().manual_seed(seed) + steps = [] + for row in routed: + sel = torch.full((_PF_LAYERS, width), -1, dtype=torch.int32) + for layer in range(1, _PF_LAYERS): + if kind == "random": + ids = torch.randint(-1, _PF_EXPERTS, (width,), generator=gen) + elif kind == "flood": + ids = torch.randperm(_PF_EXPERTS, generator=gen)[:width] + else: + hit = list(dict.fromkeys(row[layer])) + miss = [e for e in torch.randperm(_PF_EXPERTS, generator=gen).tolist() if e not in hit][:3] + ids = torch.tensor((hit[:3] + miss + hit[:1] + [-1] + miss[:1]) * 2)[:width] + ids = ids[torch.randperm(ids.numel(), generator=gen)] + sel[layer, : ids.numel()] = ids.to(torch.int32) + steps.append(sel) + return steps + + +def _on_stack(kind, banks, monkeypatch, *, policy, overlap, k=6, budget=0): + return _prefetch_stack(kind, banks, "on", monkeypatch, k=k, budget=budget, overlap=overlap, policy=policy) + + +def _with_override(cache, steps, **delays): + """before_step hook that loads step i's candidate override into the prefetcher's static buffer.""" + pf = cache.prefetch + pf.sel_override = torch.full_like(pf.sel, -1) + for name, cycles in delays.items(): + setattr(pf, name, cycles) + + def before_step(i): + pf.sel_override.copy_(steps[i][:, : pf.sel.shape[1]]) + + return before_step + + +def _assert_same_outputs(got, want): + for a_step, b_step in zip(got, want): + for a, b in zip(a_step, b_step): + assert torch.isfinite(b.float()).all() + assert torch.equal(a, b) + + +@requires_cuda +@pytest.mark.parametrize("kind", ["bf16", "nvfp4"]) +@pytest.mark.parametrize("graph", [False, True], ids=["eager", "graph"]) +@pytest.mark.parametrize("bs", [1, 2]) +@pytest.mark.parametrize("policy", ["rule", "lru"]) +@pytest.mark.parametrize("overlap", [False, True], ids=["serial", "overlap"]) +def test_moe_prefetch_on_is_bitwise_identical(kind, graph, bs, policy, overlap, monkeypatch): + """FREETOKEN_MOE_PREFETCH=on copies predicted experts into low-priority slots beside the decode: + over 200 steps of a cache that evicts every step the MoE outputs are those of prefetch off to + the bit, every held slot holds its expert's bytes, and the prefetches were used.""" + from freetoken.moe.prefetch import CALLS, COPIED, ISSUED, USEFUL + + banks = _prefetch_banks(kind) + inputs = _prefetch_inputs(bs, steps=_PF_ON_STEPS, seed=5) + moes, cache = _prefetch_stack(kind, banks, "off", monkeypatch, overlap=overlap, policy=policy) + off = _prefetch_run(moes, cache, bs, graph, inputs, sleep=0) + _assert_slots_hold_their_experts(cache) + moes, cache = _on_stack(kind, banks, monkeypatch, policy=policy, overlap=overlap) + on = _prefetch_run(moes, cache, bs, graph, inputs, sleep=0) + _assert_same_outputs(on[0], off[0]) + _assert_slots_hold_their_experts(cache) + stats = cache.prefetch.stats.cpu() + assert stats[0].tolist() == [0] * stats.shape[1], "layer 0 is never prefetched" + assert stats[1:, CALLS].tolist() == [_PF_ON_STEPS] * (_PF_LAYERS - 1) + assert int(stats[:, ISSUED].sum()) > 0 and int(stats[:, USEFUL].sum()) > 0 and int(stats[:, COPIED].sum()) > 0 + assert int(on[3].sum()) < int(off[3].sum()), "useful prefetches remove demand misses" + + +@requires_cuda +@pytest.mark.parametrize("graph", [False, True], ids=["eager", "graph"]) +@pytest.mark.parametrize("bs", [1, 2]) +@pytest.mark.parametrize("policy", ["rule", "lru"]) +@pytest.mark.parametrize("garbage", ["random", "adversarial"]) +def test_moe_prefetch_on_survives_a_garbage_predictor(graph, bs, policy, garbage, monkeypatch): + """Whatever the candidates (random ids and padding, or the next layer's routed ids mixed with ids + it will not route, duplicates and -1), prefetch on changes no output bit and no slot's bytes.""" + from freetoken.moe.prefetch import ISSUED, USEFUL + + banks = _prefetch_banks("nvfp4") + inputs = _prefetch_inputs(bs, steps=_PF_ON_STEPS, seed=6) + moes, cache = _prefetch_stack("nvfp4", banks, "off", monkeypatch, overlap=True, policy=policy) + off = _prefetch_run(moes, cache, bs, graph, inputs, sleep=0) + moes, cache = _on_stack("nvfp4", banks, monkeypatch, policy=policy, overlap=True, budget=8) + steps = _garbage(_routed_ids(moes, inputs), 12, garbage, seed=bs) + on = _prefetch_run(moes, cache, bs, graph, inputs, sleep=0, before_step=_with_override(cache, steps)) + _assert_same_outputs(on[0], off[0]) + _assert_slots_hold_their_experts(cache) + stats = cache.prefetch.stats.cpu() + assert int(stats[:, ISSUED].sum()) > 0 and int(stats[:, USEFUL].sum()) > 0 + assert int(stats[:, USEFUL].sum()) < int(stats[:, ISSUED].sum()), "some prefetches must be wrong" + + +# ~0.5 ms at the RTX 5080's clock: far longer than a toy layer, so every delayed stage is late +_PF_DELAY_CYCLES = 1_000_000 + + +@requires_cuda +@pytest.mark.parametrize("graph", [False, True], ids=["eager", "graph"]) +@pytest.mark.parametrize("bs", [1, 2]) +@pytest.mark.parametrize("policy", ["rule", "lru"]) +@pytest.mark.parametrize("delay", ["copy", "predict", "both"]) +def test_moe_prefetch_on_joins_hold_when_the_prefetch_streams_run_late(graph, bs, policy, delay, monkeypatch): + """Race test: a sleep on the prefetch copy stream (every copy lands after its layer's ensure) or on + the predictor stream (prefetch_ensure runs long after the fork) must not change a bit. The + candidates include each layer's routed ids, so GEMMs do read prefetched slots. The joins also + make the slot maps timing-free: the delayed run ends with the undelayed run's maps and counts.""" + from freetoken.moe.prefetch import COPIED, LATE, NUM_STAT_COLS, USEFUL + + banks = _prefetch_banks("nvfp4") + inputs = _prefetch_inputs(bs, steps=_PF_ON_STEPS, seed=7) + moes, cache = _prefetch_stack("nvfp4", banks, "off", monkeypatch, overlap=True, policy=policy) + off = _prefetch_run(moes, cache, bs, graph, inputs, sleep=0) + runs = [] + for cycles in (0, _PF_DELAY_CYCLES): + moes, cache = _on_stack("nvfp4", banks, monkeypatch, policy=policy, overlap=True, budget=8) + steps = _garbage(_routed_ids(moes, inputs), 12, "adversarial", seed=bs + 10) + delays = {"delay_copy_cycles": cycles * (delay in ("copy", "both")), + "delay_predict_cycles": cycles * (delay in ("predict", "both"))} + on = _prefetch_run(moes, cache, bs, graph, inputs, sleep=0, before_step=_with_override(cache, steps, **delays)) + _assert_same_outputs(on[0], off[0]) + _assert_slots_hold_their_experts(cache) + runs.append((on[1][:3], cache.prefetch.stats.cpu())) + (maps, stats), (delayed_maps, delayed_stats) = runs + assert int(stats[:, USEFUL].sum()) > 0 + for a, b in zip(maps, delayed_maps): + assert torch.equal(a, b) + cols = [c for c in range(NUM_STAT_COLS) if c != LATE] + assert torch.equal(stats[:, cols], delayed_stats[:, cols]) + if delay == "copy": + # the late counter sees what the join waited for: the copies were still sleeping at ensure(L) + # (all of them under replay; in eager a loaded host can take longer than the sleep per layer) + late, copied = int(delayed_stats[:, LATE].sum()), int(delayed_stats[:, COPIED].sum()) + assert copied > 0 and (late == copied if graph else late > 0) + + +@requires_cuda +def test_moe_prefetch_race_test_catches_a_missing_copy_join(monkeypatch): + """Negative control for the race test: without the compute stream's wait on the prefetch copy, + GEMMs read prefetched slots before their bytes land and the outputs change.""" + from freetoken.moe.prefetch import ExpertPrefetcher + + banks = _prefetch_banks("nvfp4") + inputs = _prefetch_inputs(1, steps=40, seed=7) + moes, cache = _prefetch_stack("nvfp4", banks, "off", monkeypatch, overlap=True) + off = _prefetch_run(moes, cache, 1, False, inputs, sleep=0) + moes, cache = _on_stack("nvfp4", banks, monkeypatch, policy="rule", overlap=True, budget=8) + + def no_wait(self, layer_id): + self._inflight[layer_id] = None + + monkeypatch.setattr(ExpertPrefetcher, "join_copy", no_wait) + steps = _garbage(_routed_ids(moes, inputs), 12, "adversarial", seed=11) + on = _prefetch_run(moes, cache, 1, False, inputs, sleep=0, + before_step=_with_override(cache, steps, delay_copy_cycles=_PF_DELAY_CYCLES)) + assert any(not torch.equal(a, b) for got, want in zip(on[0], off[0]) for a, b in zip(got, want)) + + +def _replay_on_reference(policy, calls, rows): + """Replay the spied prefetch_ensure / ensure sequence on the CPU reference; per-layer stats.""" + from freetoken.moe.prefetch import CALLS, COPIED, ISSUED, MISSES, ROWS, USEFUL + from freetoken.moe.scored_ensure import POLICY_IDS + from tests.moe.ref_scored_cache import RefScoredCache + + ref = RefScoredCache(_PF_LAYERS, _PF_EXPERTS, _PF_EXPERTS, POLICY_IDS[policy]) + stats = torch.zeros((_PF_LAYERS, 8), dtype=torch.int64) + planned = {} + for call in calls: + if call[0] == "prefetch": + _, layer, sel, budget = call + src, _ = ref.prefetch(layer, sel.tolist(), budget) + stats[layer, ISSUED] += src.size + planned[layer] = src.size + else: + _, layer, ids, logits = call + _, src, _ = ref.ensure(layer, ids.view(-1).numpy(), bump_tok=layer == 0, lowpri=True, + logits=None if logits is None else logits.float().numpy(), thr=0.25) + if layer in planned: + stats[layer, [USEFUL, CALLS, ROWS, MISSES, COPIED]] += torch.tensor( + [ref.useful, 1, rows, src.size, int(planned.pop(layer) > 0)]) + return ref, stats + + +@requires_cuda +@pytest.mark.parametrize("bs", [1, 2]) +@pytest.mark.parametrize("policy", ["rule", "lru"]) +def test_moe_prefetch_on_counters_match_reference(bs, policy, monkeypatch): + """issued / useful / calls / rows / misses / copied equal a CPU replay of the spied candidates + and routings, the slot maps end identical, late is 0 when each ensure waits for the copies + (known answer), and with the candidates fixed a captured graph counts the same as eager. + prefetch_ensure and the demand copy run on the compute stream, the slim prefetch copy on the + copy stream.""" + import freetoken.kernel.fast_index_copy as fic + from freetoken.moe.prefetch import LATE, NUM_STAT_COLS, RESIDENT_HITS + + banks = _prefetch_banks("nvfp4") + inputs = _prefetch_inputs(bs, steps=60, seed=8) + moes, cache = _on_stack("nvfp4", banks, monkeypatch, policy=policy, overlap=True) + calls, streams = [], {"pf_ensure": set(), "slim": set(), "demand": set()} + real_pf, real_ensure = cache.prefetch_ensure, cache.ensure_experts + real_slim, real_multi = fic.fast_index_copy_multi_slim_jit, fic.fast_index_copy_multi_jit + + def spy_pf(layer, query, *args, **kw): + streams["pf_ensure"].add(torch.cuda.current_stream().cuda_stream) + calls.append(("prefetch", layer, query.clone(), kw.get("budget"))) # on the compute stream after its wait on the select + return real_pf(layer, query, *args, **kw) + + def spy_ensure(layer, ids, **kw): + torch.cuda.synchronize() # the copies forked before this ensure have landed: late must be 0 + logits = kw.get("router_logits") + calls.append(("ensure", layer, ids.clone(), None if logits is None or policy != "rule" else logits.clone())) + return real_ensure(layer, ids, **kw) + + def spy_slim(*args, **kw): + streams["slim"].add(torch.cuda.current_stream().cuda_stream) + return real_slim(*args, **kw) + + def spy_multi(*args, **kw): + streams["demand"].add(torch.cuda.current_stream().cuda_stream) + return real_multi(*args, **kw) + + def install(): + monkeypatch.setattr(cache, "prefetch_ensure", spy_pf) + monkeypatch.setattr(cache, "ensure_experts", spy_ensure) + monkeypatch.setattr(fic, "fast_index_copy_multi_slim_jit", spy_slim) + monkeypatch.setattr(fic, "fast_index_copy_multi_jit", spy_multi) + + eager = _prefetch_run(moes, cache, bs, False, inputs, before_decode=install, sleep=0) + torch.cuda.synchronize() + calls = [c if c[0] == "prefetch" else (c[0], c[1], c[2].cpu(), None if c[3] is None else c[3].cpu()) for c in calls] + calls = [(c[0], c[1], c[2].cpu(), c[3]) if c[0] == "prefetch" else c for c in calls] + ref, want = _replay_on_reference(policy, calls, bs) + got = cache.prefetch.stats.cpu() + assert got.shape == (_PF_LAYERS, NUM_STAT_COLS) + assert int(got[:, LATE].sum()) == 0 and int(got[:, RESIDENT_HITS].sum()) == 0 + assert torch.equal(got, want), (got, want) + assert int(want[:, 1].sum()) > 0, "the sequence must have useful prefetches" + assert torch.equal(cache.slot_for_id.view(-1).cpu().long(), torch.from_numpy(ref.slot_of_id)) + assert torch.equal(cache.id_of_slot.cpu().long(), torch.from_numpy(ref.id_of_slot)) + assert torch.equal(cache.usage.cpu(), torch.from_numpy(ref.usage)) + assert streams["pf_ensure"] == {torch.cuda.current_stream().cuda_stream} + assert streams["slim"] == {cache.prefetch.copy_stream.cuda_stream} + assert streams["demand"] == {torch.cuda.current_stream().cuda_stream} + + # the select reads residency beside layer L-1's ensure, so two runs may pick different candidates: + # with the candidates fixed, a captured graph must install, count and end like eager + monkeypatch.undo() + runs = [] + for graph in (False, True): + moes, cache = _on_stack("nvfp4", banks, monkeypatch, policy=policy, overlap=True) + steps = _garbage(_routed_ids(moes, inputs), 12, "adversarial", seed=bs + 8) + run = _prefetch_run(moes, cache, bs, graph, inputs, sleep=_PF_SLEEP_CYCLES, before_step=_with_override(cache, steps)) + _assert_same_outputs(run[0], eager[0]) + runs.append((run[1][:3], cache.prefetch.stats.cpu())) + cols = [c for c in range(NUM_STAT_COLS) if c != LATE] + (maps, stats), (graph_maps, graph_stats) = runs + for a, b in zip(maps, graph_maps): + assert torch.equal(a, b) + assert torch.equal(graph_stats[:, cols], stats[:, cols]) and int(stats[:, 1].sum()) > 0 + + +@requires_cuda +def test_moe_prefetch_on_rebuild_and_reset_clear_the_prefetch_state(monkeypatch): + """A rebuild keeps both prefetch streams, the events and plan buffers, and drops the counters, + the in-flight state and every low-priority slot; a recaptured graph then matches eager and off. + A reset (as after a capture) also drops the counters and low-priority slots.""" + banks = _prefetch_banks("bf16") + inputs = _prefetch_inputs(1, steps=40, seed=9) + moes, cache = _on_stack("bf16", banks, monkeypatch, policy="rule", overlap=True) + _prefetch_run(moes, cache, 1, True, inputs, sleep=0) + pf = cache.prefetch + streams, plan = (pf.stream, pf.copy_stream), (pf.pf_slots, pf.pf_src, pf.pf_num, pf.pf_ready) + assert int(pf.stats.sum()) > 0 and bool(((cache.id_of_slot >= 0) & (cache.usage == 0)).any()) + pf._inflight[2] = torch.zeros(1) + cache.rebuild(_PF_EXPERTS + 4) + assert (pf.stream, pf.copy_stream) == streams and (pf.pf_slots, pf.pf_src, pf.pf_num, pf.pf_ready) == plan + assert int(pf.stats.abs().sum()) == 0 and int(pf.totals.abs().sum()) == 0 and pf._inflight == [None] * _PF_LAYERS + assert not bool(((cache.id_of_slot >= 0) & (cache.usage == 0)).any()) + graph = _prefetch_run(moes, cache, 1, True, inputs, sleep=0) + eager = _prefetch_run(moes, cache, 1, False, inputs, sleep=0) + _assert_same_outputs(graph[0], eager[0]) + _assert_slots_hold_their_experts(cache) + assert int(pf.stats.sum()) > 0 + cache.reset() + assert int(pf.stats.abs().sum()) == 0 and not bool(((cache.id_of_slot >= 0) & (cache.usage == 0)).any()) + moes_off, cache_off = _prefetch_stack("bf16", banks, "off", monkeypatch, overlap=True) + cache_off.rebuild(_PF_EXPERTS + 4) + off = _prefetch_run(moes_off, cache_off, 1, False, inputs, sleep=0) + _assert_same_outputs(eager[0], off[0]) + + +@requires_cuda +def test_moe_prefetch_on_skips_batches_past_the_select_tile(monkeypatch): + """An eager decode wider than MAX_ROWS forks nothing: no join, no count, the same outputs.""" + from freetoken.moe.prefetch import MAX_ROWS + + banks = _prefetch_banks("bf16") + bs = MAX_ROWS + 1 + inputs = _prefetch_inputs(bs, steps=6, seed=10) + off = _prefetch_run(*_prefetch_stack("bf16", banks, "off", monkeypatch), bs, False, inputs, sleep=0) + moes, cache = _on_stack("bf16", banks, monkeypatch, policy="rule", overlap=True) + on = _prefetch_run(moes, cache, bs, False, inputs, sleep=0) + _assert_same_outputs(on[0], off[0]) + assert int(cache.prefetch.stats.abs().sum()) == 0 and cache.prefetch._inflight == [None] * _PF_LAYERS + + +def _stall_hooks(monkeypatch): + """Sleeps (cycles, read at enqueue time) on the compute stream right before each decode GEMM, i.e. + after issue_copy(L+1), and on the stream that runs the shared expert gate (the side stream).""" + import freetoken.models.qwen4_exp.moe as q4moe + from freetoken.layers.moe import OffloadMoELayer + + stall = {"gemm": 0, "shared": 0} + real_gemm, real_gate = OffloadMoELayer._expert_gemm, q4moe.shared_gate_sigmoid + + def gemm(self, cache, *args, is_prefill, **kw): + if stall["gemm"] and not is_prefill: + torch.cuda._sleep(stall["gemm"]) + return real_gemm(self, cache, *args, is_prefill=is_prefill, **kw) + + def gate(*args, **kw): + if stall["shared"]: + torch.cuda._sleep(stall["shared"]) + return real_gate(*args, **kw) + + monkeypatch.setattr(OffloadMoELayer, "_expert_gemm", gemm) + monkeypatch.setattr(q4moe, "shared_gate_sigmoid", gate) + return stall + + +@requires_cuda +@pytest.mark.parametrize("graph", [False, True], ids=["eager", "graph"]) +@pytest.mark.parametrize("bs", [1, 2]) +@pytest.mark.parametrize("policy", ["rule", "lru"]) +@pytest.mark.parametrize("where", ["gemm", "shared"]) +@pytest.mark.parametrize("cands", ["adversarial", "flood"]) +def test_moe_prefetch_on_holds_when_the_compute_side_stalls(graph, bs, policy, where, cands, monkeypatch): + """Race test from the other side: a sleep on the compute stream between issue_copy(L+1) and + GEMM(L) lets copy(L+1) land before GEMM(L) reads, so a prefetch into any of layer L's routed + slots changes outputs deterministically; a sleep on the shared expert's side stream must not + matter either. 'flood' asks for all 16 experts a layer, more than the unpinned slots, so even + LRU's recency order would reach the pinned ones. Bit-identical to off, and the same final maps + as the unstalled run.""" + k, width = (8, 16) if cands == "flood" else (6, 12) + banks = _prefetch_banks("nvfp4") + inputs = _prefetch_inputs(bs, steps=80, seed=21) + moes, cache = _prefetch_stack("nvfp4", banks, "off", monkeypatch, overlap=True, policy=policy) + off = _prefetch_run(moes, cache, bs, graph, inputs, sleep=0) + stall = _stall_hooks(monkeypatch) + runs = [] + for cycles in (0, _PF_DELAY_CYCLES): + stall[where] = cycles + moes, cache = _on_stack("nvfp4", banks, monkeypatch, policy=policy, overlap=True, k=k, budget=width) + steps = _garbage(_routed_ids(moes, inputs), width, cands, seed=bs + 20) + on = _prefetch_run(moes, cache, bs, graph, inputs, sleep=0, before_step=_with_override(cache, steps)) + _assert_same_outputs(on[0], off[0]) + _assert_slots_hold_their_experts(cache) + runs.append(on[1][:3]) + for a, b in zip(*runs): + assert torch.equal(a, b) + + +def _assert_held_slots_hold_their_experts(cache): + """_assert_slots_hold_their_experts, but a prefill may leave no slot held.""" + if bool((cache.id_of_slot >= 0).any()): + _assert_slots_hold_their_experts(cache) + + +def _mixed_run(moes, cache, bs, graph, schedule, before_step=None, late=None, settle=True): + """Warm up (and capture) the decode like _prefetch_run, reset, then run ``schedule``: ("decode", xs) + is one decode step (a replay under ``graph``), ("prefill", xs) an eager prefill of xs's tokens. + ``late=(cycles, steps)`` sleeps that long on the predictor stream in those decode steps (a second + graph captured with the sleep, in the same pool, replays them); ``settle`` syncs and checks the + slots after each prefill, else the next decode is enqueued right behind it.""" + ctx = _fresh_ctx(_batch=SimpleNamespace(is_prefill=False)) + + def step(xs): + return [moe.forward(x.clone()) for moe, x in zip(moes, xs)] + + static = [torch.randn(bs, _PF_HIDDEN, device="cuda", dtype=torch.bfloat16) for _ in moes] + step(static) + if graph: + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + graphs = {False: (g, step(static))} + if late is not None: + cache.prefetch.delay_predict_cycles = late[0] + g_late = torch.cuda.CUDAGraph() + with torch.cuda.graph(g_late, pool=g.pool()): + graphs[True] = (g_late, step(static)) + cache.prefetch.delay_predict_cycles = 0 + cache.reset() + results = [] + for i, (phase, xs) in enumerate(schedule): + if before_step is not None: + before_step(i) + ctx._batch.is_prefill = phase == "prefill" + slow = late is not None and i in late[1] + if phase == "decode" and graph: + for s, x in zip(static, xs): + s.copy_(x) + g, outs = graphs[slow] + g.replay() + else: + if late is not None: + cache.prefetch.delay_predict_cycles = late[0] if slow else 0 + outs = step(xs) + if phase == "decode" and cache.audit is not None: + cache.audit.end_decode_step(cache) + results.append([o.clone() for o in outs]) + if phase == "prefill" and settle: + torch.cuda.synchronize() + _assert_held_slots_hold_their_experts(cache) + torch.cuda.synchronize() + _assert_held_slots_hold_their_experts(cache) + return results + + +_PREFILL_PATHS = { + # (cache_size, prefill tokens, cache kwargs, FREETOKEN_MOE_SMALL_PREFILL_TOKENS) + "materialize": (_PF_EXPERTS + 4, 24, {}, 0), + "overlap": (2 * _PF_EXPERTS + 4, 24, {"prefill_overlap": True}, 0), + "overlap_d2d": (2 * _PF_EXPERTS + 4, 24, {"prefill_overlap": True, "prefill_hit_d2d": True}, 0), + "small": (_PF_EXPERTS, 3, {}, 8), +} + + +@requires_cuda +@pytest.mark.parametrize("graph", [False, True], ids=["eager", "graph"]) +@pytest.mark.parametrize("bs", [1, 2]) +@pytest.mark.parametrize("policy", ["rule", "lru"]) +@pytest.mark.parametrize("path", list(_PREFILL_PATHS)) +def test_moe_prefetch_on_survives_prefills_between_decodes(graph, bs, policy, path, monkeypatch): + """Prefills between decode steps rewrite slots behind the prefetch's back: the whole-layer + materialize installs slots [0, E), the overlap double buffers invalidate and stream [0, 2E) + (with the hit-D2D gather reading resident slots), and small prefills ensure their experts + under the low-priority keys. Decode, prefill, decode, prefill, decode: every output of prefetch + on equals off to the bit, and after each prefill every held slot holds its expert.""" + import freetoken.layers.moe as layers_moe + + cache_size, tokens, cache_kw, small = _PREFILL_PATHS[path] + monkeypatch.setattr(layers_moe, "_SMALL_PREFILL_TOKENS", small) + banks = _prefetch_banks("nvfp4") + decode = _prefetch_inputs(bs, steps=18, seed=31) + gen = torch.Generator(device="cuda").manual_seed(32) + prefills = [[torch.randn(tokens, _PF_HIDDEN, device="cuda", dtype=torch.bfloat16, generator=gen) * 0.5 + for _ in range(_PF_LAYERS)] for _ in range(2)] + schedule = ([("decode", xs) for xs in decode[:6]] + [("prefill", prefills[0])] + + [("decode", xs) for xs in decode[6:12]] + [("prefill", prefills[1])] + + [("decode", xs) for xs in decode[12:]]) + kw = dict(overlap=True, policy=policy, cache_size=cache_size, **cache_kw) + moes, cache = _prefetch_stack("nvfp4", banks, "off", monkeypatch, **kw) + off = _mixed_run(moes, cache, bs, graph, schedule) + moes, cache = _prefetch_stack("nvfp4", banks, "on", monkeypatch, budget=8, **kw) + steps = _garbage(_routed_ids(moes, [xs for _, xs in schedule]), 12, "adversarial", seed=bs + 30) + on = _mixed_run(moes, cache, bs, graph, schedule, before_step=_with_override(cache, steps)) + _assert_same_outputs(on, off) + if path == "overlap_d2d": + assert cache.prefill_hit_rows > 0, "the hit-D2D gather must have read resident slots" + assert int(cache.prefetch.stats[:, 0].sum()) > 0 + + +def _record_compute_streams(monkeypatch, moes): + """The streams the MoE blocks are entered on (the default stream, and the capture stream under a + graph): the compute streams of the decode.""" + compute = set() + for moe in moes: + real = moe.forward + + def forward(x, real=real): + compute.add(torch.cuda.current_stream().cuda_stream) + return real(x) + + monkeypatch.setattr(moe, "forward", forward) + return compute + + +def _spy_map_writers(monkeypatch): + """(writer, stream) of every launch that writes the slot maps or their mirrors, at enqueue time.""" + import freetoken.kernel.triton.moe as triton_moe + import freetoken.moe.offload_kernels as ok + import freetoken.moe.scored_ensure as se + + seen = [] + + def spy(name, real): + def launch(*args, **kwargs): + seen.append((name(kwargs) if callable(name) else name, torch.cuda.current_stream().cuda_stream)) + return real(*args, **kwargs) + + return launch + + # scored_ensure is the kernel of both the demand ensures and prefetch_ensure (prefetch=True) + monkeypatch.setattr(se, "scored_ensure", spy(lambda kw: "prefetch_ensure" if kw.get("prefetch") else "ensure", + se.scored_ensure)) + monkeypatch.setattr(ok, "lru_ensure", spy("ensure", ok.lru_ensure)) + monkeypatch.setattr(ok, "materialize_layer", spy("materialize", ok.materialize_layer)) + monkeypatch.setattr(ok, "reset_cache", spy("reset", ok.reset_cache)) + monkeypatch.setattr(triton_moe, "invalidate_prefill_slots", spy("invalidate", triton_moe.invalidate_prefill_slots)) + return seen + + +@requires_cuda +@pytest.mark.parametrize("graph", [False, True], ids=["eager", "graph"]) +@pytest.mark.parametrize("bs", [1, 2]) +@pytest.mark.parametrize("path", list(_PREFILL_PATHS)) +def test_moe_prefetch_on_writes_the_slot_maps_only_on_the_compute_stream(graph, bs, path, monkeypatch): + """Single writer: every launch that writes the slot maps (the demand and prefetch ensures, prefill + ensures, the whole-layer materialize, reset) runs on the stream the MoE blocks run on (the capture + stream under a graph), never on the predictor, prefetch copy or shared-expert side stream. The only + other writer is the overlap prefill's buffer invalidation, on the prefill copy stream it fences. + The predictor's select reads no cache state.""" + import freetoken.layers.moe as layers_moe + import freetoken.moe.prefetch as pf_mod + from freetoken.moe.prefetch import ISSUED + + cache_size, tokens, cache_kw, small = _PREFILL_PATHS[path] + monkeypatch.setattr(layers_moe, "_SMALL_PREFILL_TOKENS", small) + banks = _prefetch_banks("nvfp4") + decode = _prefetch_inputs(bs, steps=6, seed=95) + gen = torch.Generator(device="cuda").manual_seed(96) + prefill = [torch.randn(tokens, _PF_HIDDEN, device="cuda", dtype=torch.bfloat16, generator=gen) * 0.5 + for _ in range(_PF_LAYERS)] + schedule = [("decode", xs) for xs in decode[:3]] + [("prefill", prefill)] + [("decode", xs) for xs in decode[3:]] + moes, cache = _prefetch_stack("nvfp4", banks, "on", monkeypatch, budget=8, overlap=True, cache_size=cache_size, + **cache_kw) + steps = _garbage(_routed_ids(moes, [xs for _, xs in schedule]), 12, "adversarial", seed=bs + 97) + seen = _spy_map_writers(monkeypatch) + compute = _record_compute_streams(monkeypatch, moes) + selects, real_select = [], pf_mod.lookahead_select + + def spy_select(logits, resident, *args, **kwargs): + selects.append((resident.data_ptr(), torch.cuda.current_stream().cuda_stream)) + return real_select(logits, resident, *args, **kwargs) + + monkeypatch.setattr(pf_mod, "lookahead_select", spy_select) + _mixed_run(moes, cache, bs, graph, schedule, before_step=_with_override(cache, steps)) + pf = cache.prefetch + assert not compute & {pf.stream.cuda_stream, pf.copy_stream.cuda_stream, cache.decode_copy_stream.cuda_stream} + assert len(compute) == 1 + graph + writers = {name for name, _ in seen} + want = {"ensure", "prefetch_ensure", "reset"} | {"materialize": {"materialize"}, "overlap": {"invalidate"}, + "overlap_d2d": {"invalidate"}, "small": set()}[path] + assert writers == want, writers + off_compute = [(name, stream) for name, stream in seen if stream not in compute] + fence = {cache.prefill_copy_stream.cuda_stream} if cache.prefill_copy_stream is not None else set() + assert all(name == "invalidate" and stream in fence for name, stream in off_compute), off_compute + assert {stream for name, stream in seen if name == "prefetch_ensure"} == compute # eager, and the capture + assert int(pf.stats[:, ISSUED].sum()) > 0 + # after the writer checks, which are what a predictor-side install fails + assert selects and set(selects) == {(pf._none_resident.data_ptr(), pf.stream.cuda_stream)} + + +# ~4 ms at the RTX 5080's clock: a select far later than its whole layer +_PF_LATE_CYCLES = 8 * _PF_DELAY_CYCLES + + +def _step_starts(times): + """before_step hook recording a timed event on the compute stream at the start of each step.""" + def before_step(i): + times.append(torch.cuda.Event(enable_timing=True)) + times[-1].record() + + return before_step + + +@requires_cuda +@pytest.mark.parametrize("graph", [False, True], ids=["eager", "graph"]) +@pytest.mark.parametrize("bs", [1, 2]) +@pytest.mark.parametrize("policy", ["rule", "lru"]) +def test_moe_prefetch_on_waits_for_a_late_select(graph, bs, policy, monkeypatch): + """The compute stream waits for the select before it installs, however late the predictor runs: + with every select ~4 ms late (far past its whole layer) a step lasts at least its three late + selects, and the outputs equal off's and the maps and counts (all but LATE) those of the run + without the delay.""" + from freetoken.moe.prefetch import LATE, NUM_STAT_COLS, USEFUL + + banks = _prefetch_banks("nvfp4") + inputs = _prefetch_inputs(bs, steps=24, seed=100) + moes, cache = _prefetch_stack("nvfp4", banks, "off", monkeypatch, overlap=True, policy=policy) + off = _prefetch_run(moes, cache, bs, graph, inputs, sleep=0) + start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True) + start.record() + torch.cuda._sleep(_PF_LATE_CYCLES) + end.record() + end.synchronize() + late_ms = start.elapsed_time(end) + runs = [] + for cycles in (0, _PF_LATE_CYCLES): + moes, cache = _on_stack("nvfp4", banks, monkeypatch, policy=policy, overlap=True, budget=8) + steps = _garbage(_routed_ids(moes, inputs), 12, "adversarial", seed=bs + 101) + load, times = _with_override(cache, steps, delay_predict_cycles=cycles), [] + mark = _step_starts(times) + + def before_step(i, load=load, mark=mark): + mark(i) + load(i) + + on = _prefetch_run(moes, cache, bs, graph, inputs, sleep=0, before_step=before_step) + _assert_same_outputs(on[0], off[0]) + _assert_slots_hold_their_experts(cache) + step_ms = sorted(a.elapsed_time(b) for a, b in zip(times, times[1:])) + runs.append((on[1][:3], cache.prefetch.stats.cpu(), step_ms[len(step_ms) // 2])) + (maps, stats, fast_ms), (late_maps, late_stats, slow_ms) = runs + assert int(stats[:, USEFUL].sum()) > 0 + for a, b in zip(maps, late_maps): + assert torch.equal(a, b) + cols = [c for c in range(NUM_STAT_COLS) if c != LATE] + assert torch.equal(stats[:, cols], late_stats[:, cols]) + # layers 1-3 are predicted; each prediction forks after the previous one was waited for + assert slow_ms >= 0.9 * (_PF_LAYERS - 1) * late_ms > 4 * fast_ms, (slow_ms, late_ms, fast_ms) + + +@requires_cuda +def test_moe_prefetch_on_late_select_without_the_wait_changes_the_maps_but_no_output(monkeypatch): + """Negative control for the test above (eager: a capture needs the predictor joined): without the + compute stream's wait, prefetch_ensure installs stale candidates, so the maps differ from the + undelayed run's, yet every output still equals off's: the slot maps have one writer, so a wrong + prediction only costs a copy.""" + from freetoken.moe.prefetch import ExpertPrefetcher + + banks = _prefetch_banks("nvfp4") + inputs = _prefetch_inputs(1, steps=24, seed=100) + moes, cache = _prefetch_stack("nvfp4", banks, "off", monkeypatch, overlap=True) + off = _prefetch_run(moes, cache, 1, False, inputs, sleep=0) + + def no_wait(self, target): + """install without the compute stream's wait on the select.""" + if target is None or self._inflight[target] is None: + return None + self.cache.prefetch_ensure( + target, self.sel[target], self.pf_slots[target], self.pf_src[target], + self.pf_num[target], self.stats[target], self.pf_ready[target], budget=self._budget[target], + ) + self._installed[target] = True + return target + + runs = [] + for skip in (False, True): + moes, cache = _on_stack("nvfp4", banks, monkeypatch, policy="rule", overlap=True, budget=8) + if skip: + monkeypatch.setattr(ExpertPrefetcher, "install", no_wait) + steps = _garbage(_routed_ids(moes, inputs), 12, "adversarial", seed=102) + on = _prefetch_run(moes, cache, 1, False, inputs, sleep=0, + before_step=_with_override(cache, steps, delay_predict_cycles=_PF_LATE_CYCLES * skip)) + _assert_same_outputs(on[0], off[0]) + _assert_slots_hold_their_experts(cache) + runs.append(on[1][:3]) + assert any(not torch.equal(a, b) for a, b in zip(*runs)) + + +@requires_cuda +@pytest.mark.parametrize("graph", [False, True], ids=["eager", "graph"]) +@pytest.mark.parametrize("bs", [1, 2]) +@pytest.mark.parametrize("path", list(_PREFILL_PATHS)) +def test_moe_prefetch_on_late_select_right_after_a_prefill(graph, bs, path, monkeypatch): + """The suspected production trigger: a decode step enqueued right behind a prefill (which + rewrote slots behind the prefetch's back, the overlap buffers' [0, 2E) included) with the + predictor ~4 ms late. Every output equals off's to the bit, the maps equal those of the same run + with no late step, every held slot holds its expert and the decode verifier finds nothing.""" + import freetoken.layers.moe as layers_moe + from freetoken.moe.verify import CHECKS + + cache_size, tokens, cache_kw, small = _PREFILL_PATHS[path] + monkeypatch.setattr(layers_moe, "_SMALL_PREFILL_TOKENS", small) + banks = _prefetch_banks("nvfp4") + decode = _prefetch_inputs(bs, steps=15, seed=105) + gen = torch.Generator(device="cuda").manual_seed(106) + prefills = [[torch.randn(tokens, _PF_HIDDEN, device="cuda", dtype=torch.bfloat16, generator=gen) * 0.5 + for _ in range(_PF_LAYERS)] for _ in range(2)] + schedule = ([("decode", xs) for xs in decode[:5]] + [("prefill", prefills[0])] + + [("decode", xs) for xs in decode[5:10]] + [("prefill", prefills[1])] + + [("decode", xs) for xs in decode[10:]]) + after_prefill = {i + 1 for i, (phase, _) in enumerate(schedule) if phase == "prefill"} + kw = dict(overlap=True, cache_size=cache_size, **cache_kw) + moes, cache = _prefetch_stack("nvfp4", banks, "off", monkeypatch, **kw) + off = _mixed_run(moes, cache, bs, graph, schedule) + runs = [] + for late in (None, (_PF_LATE_CYCLES, after_prefill)): + moes, cache = _prefetch_stack("nvfp4", banks, "on", monkeypatch, budget=8, verify_mode="meta", **kw) + steps = _garbage(_routed_ids(moes, [xs for _, xs in schedule]), 12, "adversarial", seed=bs + 107) + on = _mixed_run(moes, cache, bs, graph, schedule, before_step=_with_override(cache, steps), late=late, + settle=False) + _assert_same_outputs(on, off) + counts, records = _verify_state(cache) + assert counts[CHECKS] > 0 and counts[CHECKS + 1 :] == [0] * (len(counts) - 1) and records == [], records + runs.append([t.clone() for t in (cache.slot_for_id, cache.id_of_slot, cache.usage)]) + for a, b in zip(*runs): + assert torch.equal(a, b) + if path == "overlap_d2d": + assert cache.prefill_hit_rows > 0, "the hit-D2D gather must have read resident slots" + + +@requires_cuda +@pytest.mark.parametrize("graph", [False, True], ids=["eager", "graph"]) +@pytest.mark.parametrize("bs", [1, 2]) +@pytest.mark.parametrize("policy", ["rule", "lru"]) +def test_moe_prefetch_on_install_may_evict_a_slot_whose_prefetch_copy_is_in_flight(graph, bs, policy, monkeypatch): + """prefetch_ensure(L) runs on the compute stream before it joins copy(L-1), so it may evict a slot + that plan L-1 installed and whose slim copy is still writing it. The next writer of that slot's + bytes (copy(L) on the same copy stream, or a demand copy after join_copy) is ordered after copy(L-1): + with every copy late and flood candidates the eviction happens (counted on the device, so it holds + under replay too), and still no output bit and no slot's bytes change.""" + from freetoken.moe.prefetch import ExpertPrefetcher + + banks = _prefetch_banks("nvfp4") + inputs = _prefetch_inputs(bs, steps=40, seed=110) + moes, cache = _prefetch_stack("nvfp4", banks, "off", monkeypatch, overlap=True, policy=policy) + off = _prefetch_run(moes, cache, bs, graph, inputs, sleep=0) + moes, cache = _on_stack("nvfp4", banks, monkeypatch, policy=policy, overlap=True, budget=8) + pf = cache.prefetch + evicted_in_flight = torch.zeros((), dtype=torch.int64, device="cuda") + col = torch.arange(pf.pf_slots.shape[1], device="cuda") + real_install = ExpertPrefetcher.install + + def install(self, target): + got = real_install(self, target) + if got is not None and got >= 2: + # on the compute stream right after prefetch_ensure(got), before join_copy(got - 1) + new = torch.where(col < self.pf_num[got], self.pf_slots[got], -2) + old = torch.where(col < self.pf_num[got - 1], self.pf_slots[got - 1], -3) + shared = (new[:, None] == old[None, :]).any() + evicted_in_flight.add_((shared & (self.pf_ready[got - 1][0] == 0)).long()) + return got + + monkeypatch.setattr(ExpertPrefetcher, "install", install) + steps = _garbage(_routed_ids(moes, inputs), 12, "flood", seed=bs + 111) + on = _prefetch_run(moes, cache, bs, graph, inputs, sleep=0, + before_step=_with_override(cache, steps, delay_copy_cycles=_PF_DELAY_CYCLES)) + _assert_same_outputs(on[0], off[0]) + _assert_slots_hold_their_experts(cache) + assert int(evicted_in_flight) > 0, "no install evicted a slot of the previous layer's in-flight copy" + + +@requires_cuda +def test_moe_prefetch_in_flight_eviction_test_catches_an_unordered_next_copy(monkeypatch): + """Negative control for the test above: order each prefetch copy only after its own install (not + after copy(L-1)'s join) and put odd and even layers' copies on two streams, only the even ones + late. A late copy then lands on a slot the next install took over, and the outputs or the held + slots' bytes go wrong.""" + from freetoken.moe.prefetch import ExpertPrefetcher, _mark_ready_kernel, dedicated_stream + from freetoken.moe.slot_audit import PREFETCH_COPY + + banks = _prefetch_banks("nvfp4") + inputs = _prefetch_inputs(1, steps=40, seed=110) + moes, cache = _prefetch_stack("nvfp4", banks, "off", monkeypatch, overlap=True) + off = _prefetch_run(moes, cache, 1, False, inputs, sleep=0) + moes, cache = _on_stack("nvfp4", banks, monkeypatch, policy="rule", overlap=True, budget=8) + pf = cache.prefetch + streams = [pf.copy_stream, dedicated_stream(torch.device("cuda"), highest_priority=True)] + installed = [torch.cuda.Event() for _ in range(_PF_LAYERS)] + real_install = ExpertPrefetcher.install + + def install(self, target): + got = real_install(self, target) + if got is not None: + installed[got].record(torch.cuda.current_stream()) + return got + + def issue_copy(self, target): + if not self._pending(target): + return + stream = streams[target % 2] + stream.wait_event(installed[target]) + with torch.cuda.stream(stream): + if target % 2 == 0: + torch.cuda._sleep(_PF_DELAY_CYCLES) + self.cache.copy_rows(target, self.pf_slots[target], self.pf_src[target], self.pf_num[target], slim=True, + kind=PREFETCH_COPY) + _mark_ready_kernel[(1,)](self.pf_ready[target]) + self._copy_events[target][1].record(stream) + + monkeypatch.setattr(ExpertPrefetcher, "install", install) + monkeypatch.setattr(ExpertPrefetcher, "issue_copy", issue_copy) + steps = _garbage(_routed_ids(moes, inputs), 12, "flood", seed=112) + on = _prefetch_run(moes, cache, 1, False, inputs, sleep=0, before_step=_with_override(cache, steps)) + try: + _assert_same_outputs(on[0], off[0]) + _assert_slots_hold_their_experts(cache) + except AssertionError: + return + pytest.fail("a prefetch copy landing on a slot the next install took over went unnoticed") + + +@requires_cuda +@pytest.mark.parametrize("kind", ["bf16", "nvfp4"]) +@pytest.mark.parametrize("graph", [False, True], ids=["eager", "graph"]) +@pytest.mark.parametrize("bs", [1, 2]) +def test_moe_prefetch_on_per_bank_copy_fallback_is_bitwise_identical(kind, graph, bs, monkeypatch): + """Without the fused copy plan (FREETOKEN_FUSED_COPY=0 or misaligned banks) the prefetch copy + takes copy_rows' per-bank launches on the copy stream; late copies still change no bit.""" + import freetoken.moe.offload_cache as offload_cache + + monkeypatch.setattr(offload_cache, "_FUSED_COPY", False) + banks = _prefetch_banks(kind) + inputs = _prefetch_inputs(bs, steps=60, seed=5) + moes, cache = _prefetch_stack(kind, banks, "off", monkeypatch, overlap=True) + off = _prefetch_run(moes, cache, bs, graph, inputs, sleep=0) + moes, cache = _on_stack(kind, banks, monkeypatch, policy="rule", overlap=True, budget=8) + assert not cache._copy_fused_ok + steps = _garbage(_routed_ids(moes, inputs), 12, "adversarial", seed=bs) + on = _prefetch_run(moes, cache, bs, graph, inputs, sleep=0, + before_step=_with_override(cache, steps, delay_copy_cycles=_PF_DELAY_CYCLES // 2)) + _assert_same_outputs(on[0], off[0]) + _assert_slots_hold_their_experts(cache) + assert int(cache.prefetch.stats[:, 0].sum()) > 0 + + +def _multi_graph_run(moes, cache, schedule, before_step=None): + """The graph runner's sequence: per graph size (largest first) an eager warm-up, a capture into + one shared pool and a reset; then ``schedule`` of (bs, xs, replay) steps, replays and eager + steps interleaved.""" + _fresh_ctx(_batch=SimpleNamespace(is_prefill=False)) + + def step(xs): + return [moe.forward(x.clone()) for moe, x in zip(moes, xs)] + + graphs, pool = {}, None + for bs in (2, 1): + static = [torch.randn(bs, _PF_HIDDEN, device="cuda", dtype=torch.bfloat16) for _ in moes] + step(static) + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g, pool=pool): + outs = step(static) + cache.reset() + pool = g.pool() if pool is None else pool + graphs[bs] = (g, static, outs) + results = [] + for i, (bs, xs, replay) in enumerate(schedule): + if before_step is not None: + before_step(i) + if replay: + g, static, outs = graphs[bs] + for s, x in zip(static, xs): + s.copy_(x) + g.replay() + else: + outs = step(xs) + if cache.audit is not None: + cache.audit.end_decode_step(cache) + results.append([o.clone() for o in outs]) + torch.cuda.synchronize() + return results + + +@requires_cuda +@pytest.mark.parametrize("policy", ["rule", "lru"]) +def test_moe_prefetch_on_interleaves_graph_sizes_and_eager_steps(policy, monkeypatch): + """bs 1 and bs 2 graphs captured into one pool (as the graph runner does), replayed in an + arbitrary order with eager steps between them, including an eager bs 9 step that forks nothing: + every output equals off's to the bit and every held slot holds its expert.""" + from freetoken.moe.prefetch import MAX_ROWS + + banks = _prefetch_banks("nvfp4") + gen = torch.Generator().manual_seed(41) + order = [(int(b), bool(r)) for b, r in zip(torch.randint(0, 4, (60,), generator=gen), torch.randint(0, 2, (60,), generator=gen))] + sizes = {0: 1, 1: 2, 2: 3, 3: MAX_ROWS + 1} + xs_by_bs = {bs: iter(_prefetch_inputs(bs, steps=60, seed=42 + bs)) for bs in sizes.values()} + schedule = [] + for b, r in order: + bs = sizes[b] + schedule.append((bs, next(xs_by_bs[bs]), r and bs <= 2)) + moes, cache = _prefetch_stack("nvfp4", banks, "off", monkeypatch, overlap=True, policy=policy) + off = _multi_graph_run(moes, cache, schedule) + moes, cache = _on_stack("nvfp4", banks, monkeypatch, policy=policy, overlap=True, budget=8) + steps = _garbage(_routed_ids(moes, [xs for _, xs, _ in schedule]), 12, "adversarial", seed=43) + on = _multi_graph_run(moes, cache, schedule, before_step=_with_override(cache, steps)) + _assert_same_outputs(on, off) + _assert_slots_hold_their_experts(cache) + assert cache.prefetch._inflight == [None] * _PF_LAYERS + assert int(cache.prefetch.stats[:, 0].sum()) > 0 and int(cache.prefetch.stats[:, 1].sum()) > 0 + + +@requires_cuda +@pytest.mark.parametrize("seed", range(4)) +@pytest.mark.parametrize("rows", ["bs1", "bs2", "bs2_same"]) +@pytest.mark.parametrize("policy", ["rule", "lru"]) +def test_moe_prefetch_on_fuzzes_per_layer_stream_delays(seed, rows, policy, monkeypatch): + """Eager fuzz: before every layer a random choice of sleeps on the predictor stream, the prefetch + copy stream, the compute stream before the GEMM and the shared expert's side stream, with up to + 12 candidates a layer (most of the 16 slots re-installed every call) and, for bs2_same, two + identical rows. Outputs equal off's, held slots hold their experts, and the final maps and + stats (all but LATE) equal those of the same run without delays.""" + from freetoken.moe.prefetch import LATE, NUM_STAT_COLS + + bs = 1 if rows == "bs1" else 2 + inputs = _prefetch_inputs(bs, steps=60, seed=50 + seed) + if rows == "bs2_same": + inputs = [[x[:1].expand(2, -1).contiguous() for x in xs] for xs in inputs] + banks = _prefetch_banks("nvfp4") + moes, cache = _prefetch_stack("nvfp4", banks, "off", monkeypatch, overlap=True, policy=policy) + off = _prefetch_run(moes, cache, bs, False, inputs, sleep=0) + stall = _stall_hooks(monkeypatch) + gen = torch.Generator().manual_seed(seed) + choices = torch.randint(0, 16, (len(inputs), _PF_LAYERS), generator=gen).tolist() + runs = [] + for delayed in (False, True): + moes, cache = _on_stack("nvfp4", banks, monkeypatch, policy=policy, overlap=True, budget=12) + steps = _garbage(_routed_ids(moes, inputs), 12, "adversarial", seed=60 + seed) + pf = cache.prefetch + load = _with_override(cache, steps) + at = {"step": -1} + + def before_step(i, load=load): + at["step"] = i + load(i) + + for layer, moe in enumerate(moes): + real = moe.forward + + def forward(x, layer=layer, real=real, pf=pf): + bits = choices[at["step"]][layer] if delayed and at["step"] >= 0 else 0 + # layer L's fork runs layer L+1's prediction and copy: these delays hit them + pf.delay_predict_cycles = _PF_DELAY_CYCLES // 2 if bits & 1 else 0 + pf.delay_copy_cycles = _PF_DELAY_CYCLES // 2 if bits & 2 else 0 + stall["gemm"] = _PF_DELAY_CYCLES // 2 if bits & 4 else 0 + stall["shared"] = _PF_DELAY_CYCLES // 2 if bits & 8 else 0 + return real(x) + + monkeypatch.setattr(moe, "forward", forward) + on = _prefetch_run(moes, cache, bs, False, inputs, sleep=0, before_step=before_step) + _assert_same_outputs(on[0], off[0]) + _assert_slots_hold_their_experts(cache) + runs.append((on[1][:3], pf.stats.cpu())) + (maps, stats), (delayed_maps, delayed_stats) = runs + for a, b in zip(maps, delayed_maps): + assert torch.equal(a, b) + cols = [c for c in range(NUM_STAT_COLS) if c != LATE] + assert torch.equal(stats[:, cols], delayed_stats[:, cols]) + + +# FREETOKEN_MOE_PREFETCH_VERIFY: the debug slot checks on the same stacks, quiet on a correct decode +# and loud on each kind of fault they exist for +_VERIFY_STEPS = 60 + + +def _verify_state(cache): + """(counters summed over layers, kept records) of the stack's verifier.""" + from freetoken.moe.verify import RING + + v = cache.verify + torch.cuda.synchronize() + return v.counters.cpu().sum(0).tolist(), v.ring[: min(int(v.cursor), RING)].cpu().tolist() + + +@requires_cuda +@pytest.mark.parametrize("graph", [False, True], ids=["eager", "graph"]) +@pytest.mark.parametrize("bs", [1, 2]) +@pytest.mark.parametrize("prefetch", ["off", "on"]) +@pytest.mark.parametrize("verify", ["full", "meta"]) +def test_moe_prefetch_verify_is_quiet_on_a_correct_decode(graph, bs, prefetch, verify, monkeypatch): + """FREETOKEN_MOE_PREFETCH_VERIFY checks every decode layer call and finds nothing on a correct + decode: prefetch off, and prefetch on with candidates the GEMMs read, late copies and a late + predictor. The checks only read, so the outputs equal those of verify off to the bit.""" + from freetoken.moe.prefetch import USEFUL + from freetoken.moe.verify import CHECKS + + banks = _prefetch_banks("nvfp4") + inputs = _prefetch_inputs(bs, steps=_VERIFY_STEPS, seed=70) + runs = [] + for mode in ("off", verify): + moes, cache = _prefetch_stack("nvfp4", banks, prefetch, monkeypatch, overlap=True, budget=8, verify_mode=mode) + before_step = None + if prefetch == "on": + steps = _garbage(_routed_ids(moes, inputs), 12, "adversarial", seed=bs + 70) + before_step = _with_override(cache, steps, delay_copy_cycles=_PF_DELAY_CYCLES // 2, + delay_predict_cycles=_PF_DELAY_CYCLES // 4) + runs.append(_prefetch_run(moes, cache, bs, graph, inputs, sleep=0, before_step=before_step)) + _assert_same_outputs(runs[1][0], runs[0][0]) + counts, records = _verify_state(cache) + assert counts[CHECKS] == (_VERIFY_STEPS + 1) * _PF_LAYERS # the eager warm-up is checked too + assert counts[CHECKS + 1 :] == [0] * (len(counts) - 1) and records == [], records + assert (cache.verify.scratch is None) == (verify == "meta") + if prefetch == "on": + assert int(cache.prefetch.stats[:, USEFUL].sum()) > 0, "the GEMMs must have read prefetched slots" + + +@requires_cuda +@pytest.mark.parametrize("graph", [False, True], ids=["eager", "graph"]) +@pytest.mark.parametrize("prefetch", ["off", "on"]) +@pytest.mark.parametrize("path", list(_PREFILL_PATHS)) +def test_moe_prefetch_verify_is_quiet_across_prefills(graph, prefetch, path, monkeypatch): + """Every request prefills before it decodes: the whole-layer materialize, the overlap double + buffers (with the hit-D2D gather) and small prefills rewrite slots between decode steps, and the + decode checks after each of them still find nothing.""" + import freetoken.layers.moe as layers_moe + from freetoken.moe.verify import CHECKS + + cache_size, tokens, cache_kw, small = _PREFILL_PATHS[path] + monkeypatch.setattr(layers_moe, "_SMALL_PREFILL_TOKENS", small) + banks = _prefetch_banks("nvfp4") + decode = _prefetch_inputs(2, steps=18, seed=75) + gen = torch.Generator(device="cuda").manual_seed(76) + prefills = [[torch.randn(tokens, _PF_HIDDEN, device="cuda", dtype=torch.bfloat16, generator=gen) * 0.5 + for _ in range(_PF_LAYERS)] for _ in range(2)] + schedule = ([("decode", xs) for xs in decode[:6]] + [("prefill", prefills[0])] + + [("decode", xs) for xs in decode[6:12]] + [("prefill", prefills[1])] + + [("decode", xs) for xs in decode[12:]]) + moes, cache = _prefetch_stack("nvfp4", banks, prefetch, monkeypatch, budget=8, overlap=True, cache_size=cache_size, + verify_mode="full", **cache_kw) + before_step = None + if prefetch == "on": + steps = _garbage(_routed_ids(moes, [xs for _, xs in schedule]), 12, "adversarial", seed=77) + before_step = _with_override(cache, steps) + _mixed_run(moes, cache, 2, graph, schedule, before_step=before_step) + counts, records = _verify_state(cache) + if path == "overlap_d2d": + assert cache.prefill_hit_rows > 0, "the hit-D2D gather must have read resident slots" + assert counts[CHECKS] == (len(decode) + 1) * _PF_LAYERS + assert counts[CHECKS + 1 :] == [0] * (len(counts) - 1) and records == [], records + + +@requires_cuda +@pytest.mark.parametrize("graph", [False, True], ids=["eager", "graph"]) +def test_moe_prefetch_verify_catches_a_missing_copy_join(graph, monkeypatch): + """Negative control: without the compute stream's wait on the prefetch copy (and with that copy + late), GEMMs read prefetched slots before their bytes land. The byte check before the GEMM flags + them, and the records place the slot in the layer's prefetch copy plan.""" + from freetoken.moe.prefetch import ExpertPrefetcher + from freetoken.moe.verify import BYTES_PRE, META_BAD, PRE_BAD, R_KIND, R_SOURCE, SRC_PREFETCH + + banks = _prefetch_banks("nvfp4") + inputs = _prefetch_inputs(1, steps=40, seed=7) + moes, cache = _prefetch_stack("nvfp4", banks, "on", monkeypatch, overlap=True, budget=8, verify_mode="full") + + def no_wait(self, layer_id): + self._inflight[layer_id] = None + + monkeypatch.setattr(ExpertPrefetcher, "join_copy", no_wait) + if graph: + # a capture needs the copy stream joined back: after the last layer, past every GEMM it races + last, copy_stream = moes[-1].forward, cache.prefetch.copy_stream + + def forward(x): + out = last(x) + torch.cuda.current_stream().wait_stream(copy_stream) + return out + + monkeypatch.setattr(moes[-1], "forward", forward) + steps = _garbage(_routed_ids(moes, inputs), 12, "adversarial", seed=11) + # ~2.5 ms a copy: the copy stream falls ever further behind even an eager host's slow enqueue + _prefetch_run(moes, cache, 1, graph, inputs, sleep=0, + before_step=_with_override(cache, steps, delay_copy_cycles=5 * _PF_DELAY_CYCLES)) + counts, records = _verify_state(cache) + assert counts[PRE_BAD] > 0 and counts[META_BAD] == 0 + assert any(r[R_KIND] == BYTES_PRE and r[R_SOURCE] & SRC_PREFETCH for r in records) + + +def _on_call(layer, call, action): + """Wrap ExpertVerifier.before_gemm so ``action(verifier, cache, slots, real)`` replaces it on the + ``call``-th decode call of ``layer`` (call 0 is the eager warm-up).""" + from freetoken.moe.verify import ExpertVerifier + + real, seen = ExpertVerifier.before_gemm, {} + + def before_gemm(self, cache, slots, next_plan): + at = self._call[0] + seen[at] = seen.get(at, -1) + 1 + if (at, seen[at]) == (layer, call): + return action(self, cache, slots, lambda: real(self, cache, slots, next_plan)) + return real(self, cache, slots, next_plan) + + return before_gemm + + +@requires_cuda +@pytest.mark.parametrize("prefetch", ["off", "on"]) +def test_moe_prefetch_verify_catches_a_corrupt_slot(prefetch, monkeypatch): + """Negative control: bytes flipped in one bank of one routed slot right before the checks are + flagged before the GEMM, with the call, layer, slot, expert, bank and byte offset of the flip.""" + from freetoken.moe.verify import (BYTES_PRE, META_BAD, PRE_BAD, R_BANK, R_CALL, R_EXPERT, R_FIRST_BYTE, + R_KIND, R_LAYER, R_SLOT, ExpertVerifier) + + banks = _prefetch_banks("nvfp4") + moes, cache = _prefetch_stack("nvfp4", banks, prefetch, monkeypatch, overlap=True, verify_mode="full") + hit = {} + + def corrupt(verifier, cache, slots, real): + torch.cuda.synchronize() + hit.update(slot=int(slots.view(-1)[0]), expert=int(verifier.ids[0])) + cache.bank_views()[3][hit["slot"]].view(torch.uint8).view(-1)[100:108].bitwise_not_() + real() + + monkeypatch.setattr(ExpertVerifier, "before_gemm", _on_call(2, 10, corrupt)) + _prefetch_run(moes, cache, 1, False, _prefetch_inputs(1, steps=20, seed=71), sleep=0) + counts, records = _verify_state(cache) + assert counts[PRE_BAD] >= 1 and counts[META_BAD] == 0 + first = records[0] + assert [first[f] for f in (R_KIND, R_CALL, R_LAYER, R_SLOT, R_EXPERT, R_BANK, R_FIRST_BYTE)] == [ + BYTES_PRE, 10, 2, hit["slot"], hit["expert"], 3, 100] + + +@requires_cuda +@pytest.mark.parametrize("verify", ["full", "meta"]) +def test_moe_prefetch_verify_catches_a_wrong_slot_owner(verify, monkeypatch): + """Negative control: a routed slot's id_of_slot names another layer's expert while the checks + before the GEMM run (restored right after them): exactly that entry is flagged, with the owner it + saw, and the checks after the GEMM find the maps whole again.""" + from freetoken.moe.verify import (CHECKS, META_BAD, META_PRE, R_CALL, R_EXPERT, R_KIND, R_LAYER, R_MAPPED, + R_OWNER, R_ROW, R_SLOT, ExpertVerifier) + + banks = _prefetch_banks("nvfp4") + moes, cache = _prefetch_stack("nvfp4", banks, "off", monkeypatch, overlap=True, verify_mode=verify) + hit = {} + + def corrupt(verifier, cache, slots, real): + torch.cuda.synchronize() + s, e = int(slots.view(-1)[1]), int(verifier.ids[1]) + hit.update(slot=s, expert=e, saved=int(cache.id_of_slot[s])) + cache.id_of_slot[s] = 3 * _PF_EXPERTS + e # in range: an ensure reading it stays in bounds + real() + cache.id_of_slot[s] = hit["saved"] + + monkeypatch.setattr(ExpertVerifier, "before_gemm", _on_call(1, 7, corrupt)) + _prefetch_run(moes, cache, 1, False, _prefetch_inputs(1, steps=12, seed=72), sleep=0) + counts, records = _verify_state(cache) + assert hit["saved"] == _PF_EXPERTS + hit["expert"] + assert counts[META_BAD] == 1 and sum(counts[CHECKS + 1 :]) == 1 and len(records) == 1 + rec = records[0] + assert [rec[f] for f in (R_KIND, R_CALL, R_LAYER, R_ROW, R_SLOT, R_EXPERT, R_OWNER, R_MAPPED)] == [ + META_PRE, 7, 1, 1, hit["slot"], hit["expert"], 3 * _PF_EXPERTS + hit["expert"], hit["slot"]] + + +@requires_cuda +@pytest.mark.parametrize("graph", [False, True], ids=["eager", "graph"]) +@pytest.mark.parametrize("prefetch", ["off", "on"]) +def test_moe_prefetch_verify_catches_a_write_during_the_gemm(graph, prefetch, monkeypatch): + """Negative control: another stream rewrites a routed slot while the GEMM runs, after the checks + before it (a sleep on the compute stream holds the GEMM until the write lands). The byte check + after the GEMM flags it; the check before it saw nothing.""" + from freetoken.layers.moe import OffloadMoELayer + from freetoken.moe.verify import (BYTES_POST, BYTES_PRE, POST_BAD, R_BANK, R_CALL, R_FIRST_BYTE, R_KIND, R_LAYER, + R_SLOT, ExpertVerifier) + + banks = _prefetch_banks("nvfp4") + moes, cache = _prefetch_stack("nvfp4", banks, prefetch, monkeypatch, overlap=True, verify_mode="full") + side, real, armed = torch.cuda.Stream(), OffloadMoELayer._expert_gemm, {"on": False} + slot = torch.zeros((1,), dtype=torch.int64, device="cuda") + offset = torch.arange(16, device="cuda") + + def gemm(self, cache, hidden, weights, topk_ids, *, is_prefill, **kw): + if armed["on"] and self.layer_id == 3 and not is_prefill: + slot.copy_(topk_ids.view(-1)[:1]) # device side, so a captured graph rewrites this step's slot + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + torch.cuda._sleep(_PF_DELAY_CYCLES // 20) + row = cache.bank_views()[0].view(cache.cache_size, -1) + row.index_put_((slot, offset), row[slot, offset].bitwise_not()) + torch.cuda._sleep(_PF_DELAY_CYCLES) + out = real(self, cache, hidden, weights, topk_ids, is_prefill=is_prefill, **kw) + torch.cuda.current_stream().wait_stream(side) # a capture must join the side stream back + return out + return real(self, cache, hidden, weights, topk_ids, is_prefill=is_prefill, **kw) + + monkeypatch.setattr(OffloadMoELayer, "_expert_gemm", gemm) + real_run = ExpertVerifier.before_gemm + + def arm_at_call_5(self, cache, slots, next_plan): + # graph: the arm must be set while capturing, so the rewrite replays every step; eager: once at call 5 + armed["on"] = self._call[0] == 3 and (graph or int(self.calls[3]) == 5) + return real_run(self, cache, slots, next_plan) + + monkeypatch.setattr(ExpertVerifier, "before_gemm", arm_at_call_5) + _prefetch_run(moes, cache, 1, graph, _prefetch_inputs(1, steps=8, seed=73), sleep=0) + torch.cuda.synchronize() + counts, records = _verify_state(cache) + assert counts[POST_BAD] >= 1 + first = records[0] + assert [first[f] for f in (R_KIND, R_LAYER, R_BANK, R_FIRST_BYTE)] == [BYTES_POST, 3, 0, 0] + assert not any(r[R_KIND] == BYTES_PRE and (r[R_CALL], r[R_SLOT]) == (first[R_CALL], first[R_SLOT]) for r in records) + + +@requires_cuda +def test_moe_prefetch_verify_off_builds_and_launches_nothing(monkeypatch): + """With the flag off there is no verifier and a decode launches only what it launches without + one; meta and full only add launches (their kernels, the id copy and the host-row gather).""" + from collections import Counter + + from torch.profiler import ProfilerActivity, profile + + from freetoken.env import ENV + + monkeypatch.setattr(ENV.MOE_PREFETCH_VERIFY, "value", "0") + banks = _prefetch_banks("bf16") + inputs = _prefetch_inputs(1, steps=2) + + def kernels(verify): + moes, cache = _prefetch_stack("bf16", banks, "off", monkeypatch, verify_mode=verify) + _prefetch_run(moes, cache, 1, False, inputs[:1]) # warm up and compile outside the trace + with profile(activities=[ProfilerActivity.CUDA]) as prof: + _prefetch_run(moes, cache, 1, False, inputs) + return cache, Counter(e.name for e in prof.events() if e.device_type == torch.autograd.DeviceType.CUDA) + + cache, off = kernels(None) + assert cache.verify is None and cache.verify_mode == "off" + ours = ("_plan_kernel", "_compare_kernel", "_tally_kernel") + assert not any(k in name for name in off for k in ours) + for mode in ("meta", "full"): + _, on = kernels(mode) + assert not off - on, f"{mode} dropped a launch" + added = on - off + assert all(any(k in name for k in (*ours, "Memcpy", "fast_index_copy_multi")) for name in added), added + assert any("_compare_kernel" in name for name in added) == (mode == "full") + + +@requires_cuda +def test_moe_prefetch_verify_follows_a_rebuild(monkeypatch): + """A rebuild reallocates the slot banks (production rebuilds to its final slot count after start): + the verifier keeps its scratch and counts, points its byte checks at the new banks, and a + recaptured graph checks clean.""" + from freetoken.moe.verify import CHECKS + + banks = _prefetch_banks("nvfp4") + inputs = _prefetch_inputs(2, steps=30, seed=74) + moes, cache = _prefetch_stack("nvfp4", banks, "on", monkeypatch, overlap=True, verify_mode="full") + v = cache.verify + scratch = v.scratch + _prefetch_run(moes, cache, 2, True, inputs, sleep=0) + cache.rebuild(_PF_EXPERTS + 4) + assert cache.verify is v and v.scratch is scratch + assert v.cache_ptrs.tolist() == [bank.data_ptr() for bank in cache.bank_views()] + _prefetch_run(moes, cache, 2, True, inputs, sleep=0) + counts, records = _verify_state(cache) + assert counts[CHECKS] == 2 * (30 + 1) * _PF_LAYERS + assert counts[CHECKS + 1 :] == [0] * (len(counts) - 1) and records == [] + + +# FREETOKEN_MOE_SLOT_AUDIT (moe/slot_audit.py): quiet and output-neutral on a correct decode with +# prefills between the steps, and each injected persistent fault found with the history that explains it + + +def _audit_records(cache): + """(totals, kept records) of the stack's slot audit.""" + a = cache.audit + torch.cuda.synchronize() + return a.totals.cpu().tolist(), a.ring[: min(int(a.cursor), a.ring_n)].cpu().tolist() + + +def _bad_slots(cache): + """Ground truth: the (slot, owner) pairs whose bytes differ from the owner's host rows.""" + torch.cuda.synchronize() + held = (cache.id_of_slot >= 0).nonzero().view(-1).cpu() + ids = cache.id_of_slot[held.cuda()].long().cpu() + bad = torch.zeros(held.numel(), dtype=torch.bool) + for (per_layer, _), view in zip(cache.banks, cache.bank_views()): + want = torch.stack(per_layer).flatten(0, 1)[ids].contiguous().view(torch.uint8).view(held.numel(), -1) + got = view[held.cuda()].cpu().contiguous().view(torch.uint8).view(held.numel(), -1) + bad |= (got != want).any(dim=1) + return {(int(s), int(o)) for s, o, b in zip(held, ids, bad) if b} + + +def _spy_audit_kinds(monkeypatch): + """The event kinds the recorders were asked for, at enqueue time, each with the streams its + recorders ran on (a recorder follows its writer on the writer's stream).""" + from freetoken.moe.slot_audit import SlotAuditor + + kinds, real = {}, SlotAuditor._launch + + def launch(self, cache, kind, *args, **kwargs): + kinds.setdefault(kind, set()).add(torch.cuda.current_stream().cuda_stream) + return real(self, cache, kind, *args, **kwargs) + + monkeypatch.setattr(SlotAuditor, "_launch", launch) + return kinds + + +@requires_cuda +@pytest.mark.parametrize("graph", [False, True], ids=["eager", "graph"]) +@pytest.mark.parametrize("bs", [1, 2]) +@pytest.mark.parametrize("prefetch", ["off", "on"]) +@pytest.mark.parametrize("path", list(_PREFILL_PATHS)) +def test_moe_slot_audit_is_quiet_and_changes_no_bit(graph, bs, prefetch, path, monkeypatch): + """FREETOKEN_MOE_SLOT_AUDIT=1 records every slot writer (the decode and prefetch ensures and copies, + and between the decode steps the whole-layer materialize, the overlap double buffers with the + hit-D2D gather, or small prefills) and audits every held slot after each decode step: it finds + nothing, and the outputs, slot maps and slot bytes equal those of the audit off to the bit. Every + map writer it records ran on the compute stream but the prefill buffer invalidation (on the + prefill copy stream), and the prefetch copy on the prefetch copy stream.""" + import freetoken.layers.moe as layers_moe + import freetoken.moe.offload_cache as offload_cache + from freetoken.moe import slot_audit as sa + + cache_size, tokens, cache_kw, small = _PREFILL_PATHS[path] + monkeypatch.setattr(layers_moe, "_SMALL_PREFILL_TOKENS", small) + if path == "overlap_d2d": + # every toy bank is under the 256 KiB small-bank floor, which would leave the hit gather idle + monkeypatch.setattr(offload_cache, "_SMALL_BANK_FEAT_BYTES", 8192) + banks = _prefetch_banks("nvfp4") + decode = _prefetch_inputs(bs, steps=18, seed=81) + gen = torch.Generator(device="cuda").manual_seed(82) + prefills = [[torch.randn(tokens, _PF_HIDDEN, device="cuda", dtype=torch.bfloat16, generator=gen) * 0.5 + for _ in range(_PF_LAYERS)] for _ in range(2)] + schedule = ([("decode", xs) for xs in decode[:6]] + [("prefill", prefills[0])] + + [("decode", xs) for xs in decode[6:12]] + [("prefill", prefills[1])] + + [("decode", xs) for xs in decode[12:]]) + kinds = _spy_audit_kinds(monkeypatch) + runs = [] + for interval in (0, 1): + moes, cache = _prefetch_stack("nvfp4", banks, prefetch, monkeypatch, budget=8, overlap=True, + cache_size=cache_size, slot_audit=interval, **cache_kw) + assert (cache.audit is None) == (interval == 0) + before_step = None + if prefetch == "on": + steps = _garbage(_routed_ids(moes, [xs for _, xs in schedule]), 12, "adversarial", seed=bs + 83) + before_step = _with_override(cache, steps) + compute = _record_compute_streams(monkeypatch, moes) + out = _mixed_run(moes, cache, bs, graph, schedule, before_step=before_step) + runs.append((out, [t.clone() for t in (cache.slot_for_id, cache.id_of_slot, cache.usage, *cache.bank_views())])) + _assert_same_outputs(runs[1][0], runs[0][0]) + for a, b in zip(runs[1][1], runs[0][1]): + assert torch.equal(a, b) + totals, records = _audit_records(cache) + assert records == [] and totals[sa.T_BAD] == 0, [cache.audit.format_record(r) for r in records] + assert totals[sa.T_MAP_RANGE] == totals[sa.T_BAD_PLAN] == totals[sa.T_UNCHECKED] == 0 + assert totals[sa.T_AUDITS] == len(decode) + assert totals[sa.T_HELD] > 0 and totals[sa.T_BYTES_CHECKED] == totals[sa.T_HELD] and totals[sa.T_EVENTS] > 0 + expected = {sa.DEMAND_INSTALL, sa.DEMAND_COPY, sa.RESET} + if prefetch == "on": + expected |= {sa.PREFETCH_INSTALL, sa.PREFETCH_COPY} + expected |= { + "materialize": {sa.MATERIALIZE_CLEAR, sa.MATERIALIZE_INSTALL, sa.MATERIALIZE_COPY}, + "overlap": {sa.INVALIDATE, sa.PREFILL_BUFFER}, + "overlap_d2d": {sa.INVALIDATE, sa.PREFILL_SPLIT_H2D, sa.PREFILL_HIT_D2D}, + "small": {sa.PREFILL_INSTALL, sa.PREFILL_COPY}, + }[path] + assert set(kinds) == expected, (set(kinds) - expected, expected - set(kinds)) + fence = set() if cache.prefill_copy_stream is None else {cache.prefill_copy_stream.cuda_stream} + for kind in (sa.DEMAND_INSTALL, sa.PREFETCH_INSTALL, sa.PREFILL_INSTALL, sa.MATERIALIZE_INSTALL, + sa.MATERIALIZE_CLEAR, sa.RESET, sa.INVALIDATE): + assert kinds.get(kind, set()) <= (fence if kind == sa.INVALIDATE else compute), kind + if prefetch == "on": + assert kinds[sa.PREFETCH_COPY] == {cache.prefetch.copy_stream.cuda_stream} + + +@requires_cuda +@pytest.mark.parametrize("policy", ["rule", "lru"]) +def test_moe_slot_audit_is_quiet_across_graph_sizes_and_eager_steps(policy, monkeypatch): + """The graph runner's shape: bs 1 and bs 2 graphs in one pool, replayed in any order with eager + steps between them (a bs 9 one forks no prefetch); audited after every step, with the decode + verifier on beside it, neither finds anything.""" + from freetoken.moe import slot_audit as sa + from freetoken.moe.prefetch import MAX_ROWS + from freetoken.moe.verify import CHECKS + + banks = _prefetch_banks("nvfp4") + gen = torch.Generator().manual_seed(90) + order = [(int(b), bool(r)) for b, r in zip(torch.randint(0, 4, (40,), generator=gen), torch.randint(0, 2, (40,), generator=gen))] + sizes = {0: 1, 1: 2, 2: 3, 3: MAX_ROWS + 1} + xs_by_bs = {bs: iter(_prefetch_inputs(bs, steps=40, seed=91 + bs)) for bs in sizes.values()} + schedule = [(sizes[b], next(xs_by_bs[sizes[b]]), r and sizes[b] <= 2) for b, r in order] + moes, cache = _prefetch_stack("nvfp4", banks, "on", monkeypatch, policy=policy, overlap=True, budget=8, slot_audit=1, + verify_mode="meta") + steps = _garbage(_routed_ids(moes, [xs for _, xs, _ in schedule]), 12, "adversarial", seed=92) + _multi_graph_run(moes, cache, schedule, before_step=_with_override(cache, steps)) + totals, records = _audit_records(cache) + assert records == [] and totals[sa.T_BAD] == 0 and totals[sa.T_AUDITS] == len(schedule) + assert totals[sa.T_BYTES_CHECKED] == totals[sa.T_HELD] > 0 + counts, verify_records = _verify_state(cache) + assert counts[CHECKS] > 0 and counts[CHECKS + 1 :] == [0] * (len(counts) - 1) and verify_records == [] + + +@requires_cuda +@pytest.mark.parametrize("graph", [False, True], ids=["eager", "graph"]) +@pytest.mark.parametrize("policy", ["rule", "lru"]) +def test_moe_slot_audit_explains_a_skipped_prefetch_copy(graph, policy, monkeypatch): + """Fault injection: layer 2's prefetch copies copy nothing. Every slot the audit finds bad holds a + layer-2 expert whose history reads 'prefetch install, no prefetch copy', and every bad slot still + held at the end was found.""" + from freetoken.moe import slot_audit as sa + + monkeypatch.setattr(sa, "RING", 1024) + banks = _prefetch_banks("nvfp4") + inputs = _prefetch_inputs(1, steps=40, seed=84) + moes, cache = _prefetch_stack("nvfp4", banks, "on", monkeypatch, policy=policy, overlap=True, budget=8, slot_audit=1) + real = cache.copy_rows + + def copy_rows(layer_id, dst, src, num, *, slim, **kw): + if kw.get("kind") == sa.PREFETCH_COPY and layer_id == 2: + return None # the fault: no bytes move (and so nothing is recorded) + return real(layer_id, dst, src, num, slim=slim, **kw) + + monkeypatch.setattr(cache, "copy_rows", copy_rows) + steps = _garbage(_routed_ids(moes, inputs), 12, "adversarial", seed=85) + _prefetch_run(moes, cache, 1, graph, inputs, sleep=0, before_step=_with_override(cache, steps)) + truth = _bad_slots(cache) + totals, records = _audit_records(cache) + assert records and totals[sa.T_BAD] > 0 and totals[sa.T_NEW] == len(records) + a = cache.audit + for rec in records: + assert rec[sa.H_FLAGS] == sa.F_BYTES and rec[sa.H_OWNER] // _PF_EXPERTS == 2, a.format_record(rec) + # a layer-1 prefetch copy still in flight may land after the install and is named as such + assert a.diagnose(rec).startswith("prefetch install, no prefetch copy"), a.format_record(rec) + assert truth <= {(r[sa.H_SLOT], r[sa.H_OWNER]) for r in records} + + +@requires_cuda +@pytest.mark.parametrize("graph", [False, True], ids=["eager", "graph"]) +@pytest.mark.parametrize("prefetch", ["off", "on"]) +def test_moe_slot_audit_explains_a_demand_copy_from_the_wrong_layer(graph, prefetch, monkeypatch): + """Fault injection: layer 2's demand copies read layer 1's host banks (the right rows). The bad + slots are layer-2 demand installs whose last copy came from layer 1's source.""" + from freetoken.moe import slot_audit as sa + from freetoken.moe.offload_cache import OffloadMoeCache + + monkeypatch.setattr(sa, "RING", 1024) + real = OffloadMoeCache.copy_missing + + def copy_missing(self, layer_id=None): + if self._pending_src_layer == 2 and not self._pending_whole_layer: + return self.copy_rows(1, self.evict_slots, self.src_indices, self.num_indices, slim=False, kind=self._pending_kind) + return real(self, layer_id) + + monkeypatch.setattr(OffloadMoeCache, "copy_missing", copy_missing) + banks = _prefetch_banks("nvfp4") + inputs = _prefetch_inputs(1, steps=30, seed=86) + moes, cache = _prefetch_stack("nvfp4", banks, prefetch, monkeypatch, overlap=True, budget=8, slot_audit=1) + _prefetch_run(moes, cache, 1, graph, inputs, sleep=0) + truth = _bad_slots(cache) + totals, records = _audit_records(cache) + assert records and truth <= {(r[sa.H_SLOT], r[sa.H_OWNER]) for r in records} + for rec in records: + owner = rec[sa.H_OWNER] + assert owner // _PF_EXPERTS == 2 and rec[sa.H_FLAGS] == sa.F_BYTES + e = owner % _PF_EXPERTS + assert cache.audit.diagnose(rec).startswith( + f"demand install of L2/e{e}, then demand copy from layer 1's source (L1/e{e}, #"), cache.audit.format_record(rec) + + +@requires_cuda +@pytest.mark.parametrize("prefetch", ["off", "on"]) +def test_moe_slot_audit_explains_bytes_overwritten_after_install(prefetch, monkeypatch): + """Fault injection: bytes of a held slot flipped behind every recorder's back. The next audit keeps + one record (bank, byte offset, bad words) whose history ends in the owner's install and copy, a + later audit counts it again without a second record, and a clean slot is forgotten.""" + from freetoken.moe import slot_audit as sa + + banks = _prefetch_banks("nvfp4") + moes, cache = _prefetch_stack("nvfp4", banks, prefetch, monkeypatch, overlap=True, budget=8, slot_audit=1000) + _prefetch_run(moes, cache, 1, False, _prefetch_inputs(1, steps=12, seed=87), sleep=0) + a = cache.audit + a.scan(cache) + assert _audit_records(cache)[1] == [] + slot = int((cache.id_of_slot >= 0).nonzero()[0]) + owner = int(cache.id_of_slot[slot]) + row = cache.bank_views()[3][slot].view(torch.uint8).view(-1) + row[100:108].bitwise_not_() + a.scan(cache) + a.scan(cache) + totals, records = _audit_records(cache) + assert len(records) == 1 and [totals[t] for t in (sa.T_AUDITS, sa.T_BAD, sa.T_NEW, sa.T_STILL)] == [3, 2, 1, 1] + rec = records[0] + assert [rec[f] for f in (sa.H_AUDIT, sa.H_SLOT, sa.H_OWNER, sa.H_FLAGS, sa.H_BANKS, sa.H_FIRST_BANK, + sa.H_FIRST_BYTE, sa.H_BAD_WORDS)] == [1, slot, owner, sa.F_BYTES, 1 << 3, 3, 100, 2] + events = a.events(rec) + install = [ev for ev in events if ev["kind"] in sa.INSTALLS][-1] + copy = [ev for ev in events if ev["ring"] == "bytes"][-1] + assert a._flat(install) == owner == a._flat(copy) and copy["seq"] > install["seq"] + assert "then the bytes changed with no recorded write" in a.diagnose(rec) + row[100:108].bitwise_not_() + a.scan(cache) + row[100:108].bitwise_not_() + a.scan(cache) + totals, records = _audit_records(cache) + assert len(records) == 2 and records[1][sa.H_AUDIT] == 4 and totals[sa.T_BAD] == 3 + + +@requires_cuda +def test_moe_slot_audit_explains_stale_and_missing_map_entries(monkeypatch): + """Fault injection on the maps: slot_for_id still names a slot for its previous owner (an eviction + that forgot to clear it: routing that expert would read another's bytes), and a held slot's owner + has lost its entry. Both are found and the stale one is traced to the install it outlived.""" + from freetoken.moe import slot_audit as sa + + banks = _prefetch_banks("nvfp4") + moes, cache = _prefetch_stack("nvfp4", banks, "on", monkeypatch, overlap=True, budget=8, slot_audit=1000) + _prefetch_run(moes, cache, 1, False, _prefetch_inputs(1, steps=12, seed=88), sleep=0) + a = cache.audit + flat = cache.slot_for_id.view(-1) + held = (cache.id_of_slot >= 0).nonzero().view(-1).tolist() + found = None + for s in held: + for ev in sorted(a.meta_hist[s].tolist()): + fid = ev[sa.E_LAYER] * _PF_EXPERTS + ev[sa.E_ROW] + if ev[sa.E_KIND] in sa.INSTALLS and int(flat[fid]) == -1: + found = (s, fid) + if found: + break + assert found, "some slot must have outlived an earlier owner" + stale_slot, stale_id = found + fwd_slot = next(t for t in held if t != stale_slot) + fwd_owner = int(cache.id_of_slot[fwd_slot]) + flat[stale_id] = stale_slot + flat[fwd_owner] = -1 + a.scan(cache) + totals, records = _audit_records(cache) + assert [totals[t] for t in (sa.T_BAD, sa.T_STALE, sa.T_FWD, sa.T_BYTES)] == [2, 1, 1, 0] and len(records) == 2 + by_slot = {r[sa.H_SLOT]: r for r in records} + rec = by_slot[stale_slot] + assert rec[sa.H_FLAGS] == sa.F_STALE and rec[sa.H_STALE_ID] == stale_id + holder = a._name(int(cache.id_of_slot[stale_slot])) + text = a.diagnose(rec) + assert text.startswith(f"stale slot_for_id: {a._name(stale_id)} -> this slot, which holds {holder}; " + f"{a._name(stale_id)} was installed here by "), text + assert text.endswith("without clearing its map entry"), text + rec = by_slot[fwd_slot] + assert rec[sa.H_FLAGS] == sa.F_FWD and rec[sa.H_MAPPED] == -1 + assert a.diagnose(rec) == f"slot_for_id[{a._name(fwd_owner)}]=-1, not this slot" + + +@requires_cuda +def test_moe_slot_audit_off_builds_and_launches_nothing(monkeypatch): + """With the flag off there is no auditor and a decode launches only what it launches without one; + on, a decode only adds the recorders (and the audit its own kernels and host-row gathers).""" + from collections import Counter + + from torch.profiler import ProfilerActivity, profile + + from freetoken.env import ENV + + monkeypatch.setattr(ENV.MOE_SLOT_AUDIT, "value", 0) + banks = _prefetch_banks("bf16") + inputs = _prefetch_inputs(1, steps=2) + + def kernels(interval): + moes, cache = _prefetch_stack("bf16", banks, "on", monkeypatch, slot_audit=interval) + _prefetch_run(moes, cache, 1, False, inputs[:1]) # warm up and compile outside the trace + with profile(activities=[ProfilerActivity.CUDA]) as prof: + _prefetch_run(moes, cache, 1, False, inputs) + return cache, Counter(e.name for e in prof.events() if e.device_type == torch.autograd.DeviceType.CUDA) + + cache, off = kernels(None) + assert cache.audit is None and cache.slot_audit == 0 + ours = ("_record_kernel", "_bump_step_kernel", "_audit_") + assert not any(k in name for name in off for k in ours) + _, on = kernels(1) + assert not off - on, "the audit dropped a launch" + added = on - off + assert all(any(k in name for k in (*ours, "fast_index_copy_multi")) for name in added), added + assert any("_audit_report_kernel" in name for name in added) + + +@requires_cuda +def test_moe_slot_audit_follows_a_rebuild(monkeypatch): + """A rebuild reallocates the slot banks and maps: the auditor keeps its scratch and totals, starts a + new history sized for the new slot count, points its byte checks at the new banks, and a recaptured + graph audits clean.""" + from freetoken.moe import slot_audit as sa + + banks = _prefetch_banks("nvfp4") + inputs = _prefetch_inputs(2, steps=20, seed=89) + moes, cache = _prefetch_stack("nvfp4", banks, "on", monkeypatch, overlap=True, slot_audit=1) + a = cache.audit + scratch = a.scratch + _prefetch_run(moes, cache, 2, True, inputs, sleep=0) + cache.rebuild(_PF_EXPERTS + 4) + assert cache.audit is a and a.scratch is scratch and a.num_slots == _PF_EXPERTS + 4 + assert a.meta_hist.shape[0] == a.bytes_cur.shape[0] == _PF_EXPERTS + 4 and int(a.meta_cur.sum()) == 0 + assert a.cache_ptrs.tolist() == [bank.data_ptr() for bank in cache.bank_views()] + _prefetch_run(moes, cache, 2, True, inputs, sleep=0) + totals, records = _audit_records(cache) + assert records == [] and totals[sa.T_BAD] == 0 and totals[sa.T_AUDITS] == 2 * len(inputs) + + @requires_cuda def test_decoder_stack_prefill_and_decode(monkeypatch): """Ragged bs=3 prefill then a bs=3 decode step through the whole model with dummy weights.""" diff --git a/tests/models/qwen4_exp/test_tp_shard.py b/tests/models/qwen4_exp/test_tp_shard.py new file mode 100644 index 0000000000..87a305ef0c --- /dev/null +++ b/tests/models/qwen4_exp/test_tp_shard.py @@ -0,0 +1,92 @@ +"""TP sharding of the qwen4_exp dense weights (pure tensor math, no TP runtime needed).""" + +from types import SimpleNamespace + +import torch +from freetoken.models.qwen4_exp.weight import _shard, _shard_rows + + +def _cfg(): + g = SimpleNamespace( + num_key_heads=4, key_head_dim=8, num_value_heads=12, value_head_dim=8 + ) + return SimpleNamespace( + num_qo_heads=6, num_kv_heads=2, head_dim=8, linear_attention_group=lambda: g + ) + + +def _gather(name, t, cfg, world, dim=0): + return torch.cat([_shard(name, t, cfg, r, world) for r in range(world)], dim=dim) + + +def test_shard_rows_splits_each_part_by_head(): + t = torch.arange(4 * 4 + 2 * 4).reshape( + -1, 1 + ) # part A: 4 heads x 4 rows, part B: 2 heads x 4 + s = [_shard_rows(t, [(4, 4), (2, 4)], r, 4) for r in range(4)] + assert all(x.shape[0] == 8 for x in s) + assert s[0][:4].flatten().tolist() == [0, 1, 2, 3] + assert s[3][:4].flatten().tolist() == [12, 13, 14, 15] + # 2 kv heads over 4 ranks replicate: rank * heads // world -> heads 0, 0, 1, 1 + assert s[0][4:].equal(s[1][4:]) and s[2][4:].flatten().tolist() == [20, 21, 22, 23] + + +def test_shard_round_trips_the_fused_projections(): + cfg, world = _cfg(), 2 + qkv = torch.randn(6 * 16 + 2 * 8 + 2 * 8, 5) + assert ( + _gather("model.layers.0.self_attn.qkv_proj.weight", qkv, cfg, world).shape + == qkv.shape + ) + q0 = _shard("model.layers.0.self_attn.qkv_proj.weight", qkv, cfg, 0, world) + assert q0.shape[0] == 3 * 16 + 8 + 8 + torch.testing.assert_close(q0[:48], qkv[:48]) # q heads 0-2 (16 rows each: q|gate) + torch.testing.assert_close(q0[48:56], qkv[96:104]) # kv head 0 of k + torch.testing.assert_close(q0[56:64], qkv[112:120]) # kv head 0 of v + kd, vd, nv = 4 * 8, 12 * 8, 12 + in_proj = torch.randn(kd + kd + vd + vd + nv + nv, 5) + s = _shard("model.layers.1.linear_attn.in_proj.weight", in_proj, cfg, 1, world) + assert s.shape[0] == (kd + kd + vd + vd + nv + nv) // 2 + torch.testing.assert_close(s[:16], in_proj[16:32]) # q: k heads 2,3 + torch.testing.assert_close(s[-6:], in_proj[-6:]) # a: v heads 6-11 + conv = torch.randn(kd + kd + vd, 1, 4) + assert _shard( + "model.layers.1.linear_attn.conv1d.weight", conv, cfg, 0, world + ).shape == ((kd + kd + vd) // 2, 1, 4) + a_log = torch.randn(nv) + torch.testing.assert_close( + _shard("model.layers.1.linear_attn.A_log", a_log, cfg, 1, world), a_log[6:] + ) + + +def test_shard_row_parallel_and_vocab_and_replicated(): + cfg, world = _cfg(), 2 + for name in ( + "model.layers.0.self_attn.o_proj.weight", + "model.layers.1.linear_attn.out_proj.weight", + "model.layers.0.mlp.shared_expert.down_proj.weight", + ): + t = torch.randn(3, 8) + torch.testing.assert_close(_gather(name, t, cfg, world, dim=1), t) + gate_up = torch.randn(2 * 6, 3) + s1 = _shard( + "model.layers.0.mlp.shared_expert.gate_up_proj.weight", gate_up, cfg, 1, world + ) + torch.testing.assert_close(s1, torch.cat([gate_up[3:6], gate_up[9:12]])) + emb = torch.randn(10, 3) + torch.testing.assert_close( + _gather("model.embed_tokens.weight", emb, cfg, world), emb + ) + torch.testing.assert_close(_gather("lm_head.weight", emb, cfg, world), emb) + for name in ( + "model.layers.0.mlp.gate.weight", + "model.layers.0.self_attn.indexer.index_qk_proj.weight", + "model.layers.0.attn_hyper_connection.input_mix_weight_down_block_inject.weight", + "model.layers.1.ple.value_proj.weight", + "model.layers.0.mlp.shared_expert_gate.weight", + ): + t = torch.randn(4, 6) + assert _shard(name, t, cfg, 1, world).equal(t) + assert _shard("model.layers.0.self_attn.qkv_proj.weight", gate_up, cfg, 0, 1).equal( + gate_up + ) diff --git a/tests/models/qwen4_exp/test_weight.py b/tests/models/qwen4_exp/test_weight.py index 1a9a100edb..1167eb1b4b 100644 --- a/tests/models/qwen4_exp/test_weight.py +++ b/tests/models/qwen4_exp/test_weight.py @@ -518,3 +518,62 @@ def test_checkpoint_disagreeing_with_its_quant_config_is_rejected(tmp_path, quan (tmp_path / "config.json").write_text(json.dumps(_config_json(quantization_config))) with pytest.raises(ValueError, match=match): _load(str(tmp_path)) + + +_CT_NVFP4_EXPERTS = { + "quant_method": "compressed-tensors", "format": "nvfp4-pack-quantized", "quantization_status": "compressed", + "config_groups": {"group_0": { + "format": "nvfp4-pack-quantized", "targets": [r"re:.*mlp\.experts\..*(gate|up|down)_proj$"], + "weights": {"num_bits": 4, "type": "float", "symmetric": True, "group_size": 16, + "strategy": "tensor_group", "dynamic": False, "scale_dtype": "torch.float8_e4m3fn"}, + "input_activations": {"num_bits": 4, "type": "float", "symmetric": True, "group_size": 16, + "strategy": "tensor_group", "dynamic": "local", "scale_dtype": "torch.float8_e4m3fn"}, + }}, + "ignore": ["lm_head"], +} + + +def test_compressed_tensors_experts_load_like_their_modelopt_twin(tmp_path, monkeypatch): + """An llm-compressor export (RedHatAI) names the routed experts weight_packed / weight_global_scale, + stores the quant-side global (the reciprocal of modelopt's weight_scale_2) and adds an unused + input_global_scale: its pieces must equal the modelopt twin's bit for bit.""" + from freetoken.layers.quantization import QuantConfig, set_quant_config + from freetoken.layers.quantization import configs as qconfigs + from freetoken.models.nvfp4_banks import iter_nvfp4_expert_pieces + from freetoken.models.qwen4_exp.weight import nvfp4_expert_spec + + monkeypatch.setattr(qconfigs, "_QUANT_CONFIG", qconfigs._QUANT_CONFIG) # restored at teardown + layers, experts, hid, inter = 2, 3, 32, 16 + g = torch.Generator().manual_seed(0) + ct, mo = {}, {} + for layer in range(layers): + for e in range(experts): + for proj, (rows, cols) in {"gate_proj": (inter, hid), "up_proj": (inter, hid), "down_proj": (hid, inter)}.items(): + base = f"{LM}.layers.{layer}.mlp.experts.{e}.{proj}" + packed = torch.randint(0, 256, (rows, cols // 2), dtype=torch.uint8, generator=g) + scale = (torch.rand(rows, cols // 16, generator=g) * 4 + 0.1).to(torch.float8_e4m3fn) + glob = torch.rand(1, generator=g) * 1000 + 1 + ct |= {f"{base}.weight_packed": packed, f"{base}.weight_scale": scale, + f"{base}.weight_global_scale": glob, f"{base}.input_global_scale": torch.rand(1, generator=g)} + mo |= {f"{base}.weight": packed, f"{base}.weight_scale": scale, + f"{base}.weight_scale_2": 1.0 / glob, f"{base}.input_scale": torch.rand(1, generator=g)} + config = SimpleNamespace(num_experts=experts, num_moe_layers=layers, moe_intermediate_size=inter) + + def load(name, tensors, quant): + folder = tmp_path / name + folder.mkdir() + save_file(tensors, str(folder / "model.safetensors")) + set_quant_config(QuantConfig.from_hf({"quantization_config": quant})) + spec = nvfp4_expert_spec(str(folder), config) + pieces = iter_nvfp4_expert_pieces(str(folder), config, spec, drop_page_cache=lambda path: None, primary=False) + return spec, {(lay, e): piece for lay, e, _, piece in pieces} + + ct_spec, ct_pieces = load("ct", ct, _CT_NVFP4_EXPERTS) + mo_spec, mo_pieces = load("mo", mo, {"quant_method": "modelopt", "quant_algo": "NVFP4", "ignore": ["lm_head"]}) + assert ct_spec.global_reciprocal and not mo_spec.global_reciprocal + assert sorted(ct_pieces) == sorted(mo_pieces) == [(lay, e) for lay in range(layers) for e in range(experts)] + for key, piece in ct_pieces.items(): + assert sorted(piece) == sorted(mo_pieces[key]) + for role, tensor in piece.items(): + twin = mo_pieces[key][role] + assert tensor.dtype == twin.dtype and torch.equal(tensor.view(torch.uint8), twin.view(torch.uint8)), (key, role) diff --git a/tests/models/qwen4_exp/test_weight_ckpt.py b/tests/models/qwen4_exp/test_weight_ckpt.py index 605c40f18b..38d32bd1af 100644 --- a/tests/models/qwen4_exp/test_weight_ckpt.py +++ b/tests/models/qwen4_exp/test_weight_ckpt.py @@ -22,10 +22,10 @@ from freetoken.models.nvfp4_banks import iter_nvfp4_expert_pieces from freetoken.models.qwen4_exp.config import parse_config from freetoken.models.qwen4_exp.weight import ( - _NVFP4_SOURCE_SPEC, _ZERO_CENTERED_NORM_SUFFIXES, iter_weights, load_ple_table, + nvfp4_expert_spec, ) from freetoken.moe.expert_banks import build_expert_banks from freetoken.moe.host_banks import HostResidency @@ -271,8 +271,9 @@ def layer0_expert_banks(): """The real NVFP4 source-bank loader, restricted to layer 0 (1.32 GiB instead of 63.5).""" if not torch.cuda.is_available(): pytest.skip("expert bank pinning needs CUDA") + install_quant_config(MODEL_PATH) spec = dataclasses.replace( - _NVFP4_SOURCE_SPEC, layer_to_bank=lambda layer, config: 0 if layer == 0 else None + nvfp4_expert_spec(MODEL_PATH, None), layer_to_bank=lambda layer, config: 0 if layer == 0 else None ) config = SimpleNamespace(num_experts=E, hidden_size=H, moe_intermediate_size=I, num_moe_layers=1) diff --git a/tests/models/test_glm5_next_config.py b/tests/models/test_glm5_next_config.py index ad56f7e148..6a4ac84020 100644 --- a/tests/models/test_glm5_next_config.py +++ b/tests/models/test_glm5_next_config.py @@ -23,6 +23,28 @@ _KDA_IDS = tuple(i for i in range(_NUM_LAYERS) if i not in _DSA_IDS) +def test_nvfp4_pool_factory_preserves_glm5_hybrid_geometry(): + import torch + from freetoken.kvcache import create_kvcache_pool + from freetoken.kvcache.dsa_pool import KpoolDSAKVCache + + cfg = parse_config(_hf_config()) + pool = create_kvcache_pool( + cfg, num_pages=2, page_size=64, dtype=torch.bfloat16, + device=torch.device("cpu"), num_req_slots=3, kv_quant="nvfp4", + ) + assert isinstance(pool, KpoolDSAKVCache) + spec, = cfg.kv_cache_group_specs() + assert pool.num_layers == len(_DSA_IDS) + for layer in _DSA_IDS: + assert pool.latent_rows(layer).shape == (128, spec.head_dim // 2) + assert pool.latent_block_scale(layer).shape == (128, spec.head_dim // 16) + with pytest.raises(KeyError): + pool.latent_rows(_KDA_IDS[0]) + assert pool.index_k_cache(0).dtype == torch.bfloat16 + assert pool.tail_gate(0).dtype == torch.bfloat16 + + def _layer_types() -> list[str]: return [ "deepseek_sparse_attention" if i in _DSA_IDS else "linear_attention" diff --git a/tests/models/test_glm_dsa.py b/tests/models/test_glm_dsa.py index 0dceabadd0..f09a738225 100644 --- a/tests/models/test_glm_dsa.py +++ b/tests/models/test_glm_dsa.py @@ -225,7 +225,7 @@ def test_splitk_matches_single_program(): assert (o_single.float() - o_split.float()).abs().max().item() < 1e-2 -def _make_backend(dsa: bool, latent=80, dv=64, idx_dim=32, idx_heads=16, topk=64, pages=400): +def _make_backend(dsa: bool, latent=80, dv=64, idx_dim=32, idx_heads=16, topk=64, pages=400, kv_quant="none"): """Minimal ctx + pool + DSAAttnBackend (no engine).""" from types import SimpleNamespace @@ -239,9 +239,9 @@ def _make_backend(dsa: bool, latent=80, dv=64, idx_dim=32, idx_heads=16, topk=64 ctx.page_table = torch.zeros(4, pages, dtype=torch.int32, device="cuda") if dsa: ctx.kv_cache = DSAKVCache(latent, 2, pages, 1, torch.bfloat16, torch.device("cuda"), - index_head_dim=idx_dim, num_index_layers=1) + index_head_dim=idx_dim, num_index_layers=1, kv_quant=kv_quant) else: - ctx.kv_cache = MLAKVCache(latent, 2, pages, 1, torch.bfloat16, torch.device("cuda")) + ctx.kv_cache = MLAKVCache(latent, 2, pages, 1, torch.bfloat16, torch.device("cuda"), kv_quant=kv_quant) set_global_ctx(ctx) args = SimpleNamespace( kv_lora_rank=dv, qk_rope_head_dim=latent - dv, qk_head_dim=latent, @@ -258,7 +258,8 @@ def _ref_attend(q_cat, pool_rows, live_rows, scale, dv): return s.softmax(-1) @ k[:, :dv] -def test_backend_ragged_prefill_identity_and_selection(): +@pytest.mark.parametrize("kv_quant", ["none", "nvfp4"]) +def test_backend_ragged_prefill_identity_and_selection(kv_quant): """Two-request ragged prefill through the BACKEND (page-table slicing, counts = positions + 1, per-request segmentation, leader/follower reuse): request A stays under index_topk (selection == identity == dense), request B @@ -267,7 +268,7 @@ def test_backend_ragged_prefill_identity_and_selection(): torch.manual_seed(5) dv, dr, h, idx_h, idx_d, topk = 64, 16, 8, 16, 32, 64 - backend, ctx = _make_backend(dsa=True, topk=topk) + backend, ctx = _make_backend(dsa=True, topk=topk, kv_quant=kv_quant) pool = ctx.kv_cache scale = backend.sm_scale @@ -311,6 +312,10 @@ def test_backend_ragged_prefill_identity_and_selection(): # request A (kv <= topk): selection covers all live -> equals dense reference q_cat = torch.cat([q_nope, q_pe], -1) slab = pool.latent_rows(0) + if kv_quant == "nvfp4": + from tests.kernels.test_kv_nvfp4 import _decode_latent + + slab = _decode_latent(pool) for j in range(8): # A's queries, positions 32..39 live = ctx.page_table[0, : 33 + j] ref = _ref_attend(q_cat[j], slab, live, scale, dv) @@ -335,15 +340,17 @@ def test_backend_ragged_prefill_identity_and_selection(): assert (o0.float() - o1.float()).abs().max().item() < 3e-2 # identity wiring (dense ablation): same batch through an MLAKVCache backend - backend_d, ctx_d = _make_backend(dsa=False) + backend_d, ctx_d = _make_backend(dsa=False, kv_quant=kv_quant) ctx_d.page_table.copy_(ctx.page_table) - for lid in (0, 1): - ctx_d.kv_cache._kv_buffer.copy_(pool._kv_buffer) + ctx_d.kv_cache._kv_buffer.copy_(pool._kv_buffer) + if kv_quant == "nvfp4": + ctx_d.kv_cache._scale_buffer.copy_(pool._scale_buffer) + ctx_d.kv_cache._block_scale_buffer.copy_(pool._block_scale_buffer) batch_d = SimpleNamespace(reqs=reqs, positions=positions, out_loc=out_loc, active_table_idx=None, attn_metadata=None) backend_d.prepare_metadata(batch_d) od = backend_d.mla_forward(q_nope, q_pe, c_kv, k_rope, 0, batch_d, indexer_inputs=None) - slab_d = ctx_d.kv_cache.latent_rows(0) + slab_d = _decode_latent(ctx_d.kv_cache) if kv_quant == "nvfp4" else ctx_d.kv_cache.latent_rows(0) for j in range(8): live = ctx_d.page_table[0, : 33 + j] ref = _ref_attend(q_cat[j], slab_d, live, scale, dv) diff --git a/tests/models/test_nvfp4_banks_tp.py b/tests/models/test_nvfp4_banks_tp.py new file mode 100644 index 0000000000..b943cff1c0 --- /dev/null +++ b/tests/models/test_nvfp4_banks_tp.py @@ -0,0 +1,64 @@ +"""TP sharding of the NVFP4 routed-expert stream. + +Every rank reads the same checkpoint tensors and keeps only its own slice of the expert +intermediate axis, so the ranks' slices must partition the original exactly -- the MoE layer +all-reduces the partial sums, which is only correct if nothing is dropped or double-counted. +The expert kernel sizes its banks from ``MoEConfig.local_intermediate``, the same split. +""" + +import pytest +import torch +from freetoken.layers.quantization.moe.base import MoEConfig +from freetoken.models.nvfp4_banks import _tp_shard + +INTER, HIDDEN = 640, 2560 + + +def _roles(): + """One expert's checkpoint tensors, in the modelopt layout.""" + return { + "gate": torch.arange(INTER * (HIDDEN // 2), dtype=torch.uint8).reshape(INTER, HIDDEN // 2), + "gate_scale": torch.arange(INTER * (HIDDEN // 16), dtype=torch.uint8).reshape(INTER, HIDDEN // 16), + "gate_global": torch.tensor([2.5]), + "up": torch.arange(INTER * (HIDDEN // 2), dtype=torch.uint8).reshape(INTER, HIDDEN // 2), + "up_scale": torch.arange(INTER * (HIDDEN // 16), dtype=torch.uint8).reshape(INTER, HIDDEN // 16), + "up_global": torch.tensor([2.5]), + "down": torch.arange(HIDDEN * (INTER // 2), dtype=torch.uint8).reshape(HIDDEN, INTER // 2), + "down_scale": torch.arange(HIDDEN * (INTER // 16), dtype=torch.uint8).reshape(HIDDEN, INTER // 16), + "down_global": torch.tensor([3.5]), + } + + +@pytest.mark.parametrize("tp_size", [1, 2, 4]) +def test_rank_slices_partition_the_original_tensor(tp_size): + full = _roles() + for role, tensor in full.items(): + parts = [_tp_shard(role, tensor, INTER, tp_size, r) for r in range(tp_size)] + if role.endswith("_global"): + # per-tensor scalar: no I axis, every rank keeps it whole + assert all(torch.equal(p, tensor) for p in parts) + continue + axis = 1 if role.startswith("down") else 0 + assert torch.equal(torch.cat(parts, dim=axis), tensor), role + assert all(p.shape[axis] == tensor.shape[axis] // tp_size for p in parts), role + + +def test_tp1_is_the_identity(): + for role, tensor in _roles().items(): + assert _tp_shard(role, tensor, INTER, 1, 0) is tensor or torch.equal( + _tp_shard(role, tensor, INTER, 1, 0), tensor + ) + + +def test_a_shard_that_splits_a_scale_block_is_rejected(): + # scales cover 16 values; a rank slice that is not a whole number of blocks cannot be + # sliced consistently across the weight and its scale + with pytest.raises(AssertionError, match="16-wide scale blocks"): + _tp_shard("gate", torch.zeros(24, 8, dtype=torch.uint8), 24, 3, 0) + + +def test_moe_config_local_intermediate_matches_the_shard(): + for tp_size in (1, 2, 4): + cfg = MoEConfig(num_experts=8, hidden=HIDDEN, intermediate=INTER, top_k=2, tp_size=tp_size) + assert cfg.local_intermediate == INTER // tp_size + assert _tp_shard("gate", _roles()["gate"], INTER, tp_size, 0).shape[0] == cfg.local_intermediate diff --git a/tests/moe/hip_probe_faults.cpp b/tests/moe/hip_probe_faults.cpp new file mode 100644 index 0000000000..4426bf0d14 --- /dev/null +++ b/tests/moe/hip_probe_faults.cpp @@ -0,0 +1,75 @@ +#include +#include + +#include +#include +#include + +namespace { +std::atomic begins{0}, ends{0}, injected{0}, removed_nodes{0}, last_end_rc{0}; + +bool mode(const char* value) { + const char* selected = std::getenv("FREETOKEN_TEST_HIP_PROBE_FAULT"); + return selected != nullptr && std::strcmp(selected, value) == 0; +} + +template +Function real(const char* name) { + return reinterpret_cast(dlsym(RTLD_NEXT, name)); +} +} // namespace + +extern "C" hipError_t hipStreamBeginCapture(hipStream_t stream, hipStreamCaptureMode capture_mode) { + const hipError_t rc = real("hipStreamBeginCapture")(stream, capture_mode); + if (rc == hipSuccess) ++begins; + return rc; +} + +extern "C" hipError_t hipStreamWriteValue64(hipStream_t stream, void* address, uint64_t value, unsigned flags) { + // Force the compatibility path on runtimes that normally record native writes. + if (mode("dependency") && begins == 1) { + ++injected; + return hipSuccess; + } + return real("hipStreamWriteValue64")(stream, address, value, flags); +} + +extern "C" hipError_t hipStreamUpdateCaptureDependencies( + hipStream_t stream, hipGraphNode_t* dependencies, size_t count, unsigned flags) { + if (mode("dependency")) { + ++injected; + return hipErrorNotSupported; + } + return real("hipStreamUpdateCaptureDependencies")( + stream, dependencies, count, flags); +} + +extern "C" hipError_t hipGraphDestroyNode(hipGraphNode_t node) { + const hipError_t rc = real("hipGraphDestroyNode")(node); + if (rc == hipSuccess) ++removed_nodes; + return rc; +} + +extern "C" hipError_t hipStreamEndCapture(hipStream_t stream, hipGraph_t* graph) { + ++ends; + hipError_t rc = real("hipStreamEndCapture")(stream, graph); + if (mode("end-capture") && ends == 2 && rc == hipSuccess) { + // A real invalid API call also sets HIP's thread-local last-error state. + rc = real("hipGraphGetNodes")(*graph, nullptr, nullptr); + ++injected; + } + last_end_rc = static_cast(rc); + return rc; +} + +extern "C" int probe_test_injected() { return injected; } +extern "C" int probe_test_begins() { return begins; } +extern "C" int probe_test_ends() { return ends; } +extern "C" int probe_test_removed_nodes() { return removed_nodes; } +extern "C" int probe_test_last_end_rc() { return last_end_rc; } +extern "C" int probe_test_peek_error() { + return static_cast(real("hipPeekAtLastError")()); +} +extern "C" int probe_test_make_error() { + return static_cast(real("hipGraphGetNodes")(nullptr, nullptr, nullptr)); +} diff --git a/tests/moe/ref_prefetch.py b/tests/moe/ref_prefetch.py new file mode 100644 index 0000000000..c40d9351a8 --- /dev/null +++ b/tests/moe/ref_prefetch.py @@ -0,0 +1,22 @@ +"""CPU reference for the router-lookahead select and count (freetoken.moe.prefetch).""" +from __future__ import annotations + + +def ref_select(logits, resident, k: int, budget: int): + """Each row's top-k (ties to the lowest id), merged rank by rank, duplicates dropped; the first + ``budget`` non-resident ids, and the resident ids among the first ``budget`` distinct ones.""" + rows = [sorted(range(len(row)), key=lambda e: (-row[e], e))[:k] for row in logits] + merged = [] + for r in range(k): + for row in rows: + if row[r] not in merged: + merged.append(row[r]) + sel = [e for e in merged if resident[e] < 0][:budget] + res = [e for e in merged[:budget] if resident[e] >= 0] + return sel, res + + +def ref_count(sel, res, routed, misses: int, rows: int): + """``[issued, useful, resident_hits, calls, rows, misses]`` for one layer call.""" + routed = set(routed) + return [len(sel), sum(e in routed for e in sel), sum(e in routed for e in res), 1, rows, misses] diff --git a/tests/moe/ref_scored_cache.py b/tests/moe/ref_scored_cache.py new file mode 100644 index 0000000000..15153c86d5 --- /dev/null +++ b/tests/moe/ref_scored_cache.py @@ -0,0 +1,191 @@ +"""CPU reference for the scored slot-cache ensure (freetoken.moe.scored_ensure). + +A port of the adversarial-synthesis replay (global slot cache with lru_ensure semantics: the +call's ids pinned, empty slots first, ties to the lowest slot, missing ids ascending paired with +ascending victims) onto the kernel's Q16 integer arithmetic, so victims match the GPU exactly. +""" +from __future__ import annotations + +import numpy as np + +from freetoken.moe import scored_ensure as se + +_KEY_MAX = np.iinfo(np.int64).max + + +class RefScoredCache: + def __init__(self, num_layers, num_experts, cache_size, policy, + beta=se.BETA, w=se.W, halflife=se.HALFLIFE): + self.L, self.E, self.S, self.policy = num_layers, num_experts, cache_size, policy + self.beta_step, self.w_q4, self.decay_mul, self.dt_max = se.score_params(beta, w, halflife, num_layers) + block_c = 1 << max(cache_size - 1, 1).bit_length() + self.slot_bits = max(block_c.bit_length() - 1, 1) + self.g = se.softplus2_table().numpy().astype(np.int64) + self.reset() + + def reset(self): + n = self.L * self.E + self.slot_of_id = np.full(n, -1, np.int64) + self.id_of_slot = np.full(self.S, -1, np.int64) + self.usage = np.zeros(self.S, np.int64) + self.step = 0 + self.tok = 0 + self.last_tok = np.full(n, se.LAST_TOK_NEVER, np.int64) + self.lc = np.full(n, se.LC_NEVER, np.int64) + self.ct = np.zeros(n, np.int64) + + def _decayed(self, ids, now): + dt = np.minimum(np.maximum(now - self.ct[ids], 0), self.dt_max) + return self.lc[ids] - ((dt * self.decay_mul) >> se.DECAY_SHIFT) + + def _softplus2(self, x): + shift = se.Q - se.G_STEP_BITS + u = np.abs(x) + i = np.minimum(u >> shift, se.G_LAST) + frac = u & ((1 << shift) - 1) + g0, g1 = self.g[i], self.g[i + 1] + g = np.where(i < se.G_LAST, g0 + (((g1 - g0) * frac) >> shift), 0) + return np.maximum(x, 0) + g + + def _keys(self, layer, tok, now, pin_since=None, lowpri=False): + c = np.arange(self.S, dtype=np.int64) + oid = self.id_of_slot + held = oid >= 0 + if self.policy == 0: + # lowpri: empty (0) < unconsumed prefetch (held, usage 0 -> 1) < resident (usage + 1) + key = (np.where(held, self.usage + 1, 0) << self.slot_bits) | c if lowpri else (self.usage << self.slot_bits) | c + else: + ids = np.where(held, oid, 0) + lk = ids // self.E + ahead = lk > layer + d = np.where(ahead, lk - layer, self.L - layer + lk) + k = np.where(ahead, tok - 1, tok) - self.last_tok[ids] + k = np.minimum(np.maximum(k, 0), se.K_MAX) + score = -((k << se.Q) + d * self.beta_step) + if self.policy >= 2: + lc = self.lc[ids] + lcv = self._decayed(ids, now) + lcv = np.where(lc == se.LC_NEVER, se.LC_FLOOR, np.maximum(lcv, se.LC_FLOOR)) + score = score + ((self.w_q4 * lcv) >> 4) + assert np.abs(score).max() < 2**31 # the kernel computes it in int32 + key = np.where(held, ((score + se.SCORE_BIAS) << self.slot_bits) | c, c) + if lowpri: + key = np.where(held & (self.usage == 0), (1 << self.slot_bits) | c, key) + evictable = self.usage != self.step + if pin_since is not None: + evictable &= self.usage <= pin_since + return np.where(evictable, key, _KEY_MAX) + + def ensure(self, layer, ids, *, bump_tok=False, update_state=True, logits=None, thr=0.0, near_miss=None, + pin_since=None, lowpri=False): + """One ensure call; returns ``(out_slots, src_ids, dst_slots)`` like the kernel's plan. + + ``pin_since`` (a ``step`` value) also pins every slot touched after it. ``lowpri`` + (FREETOKEN_MOE_PREFETCH=on) keys unconsumed prefetches between the empty slots and the + residents; ``self.useful`` is then the call's distinct hits on such slots. + + Policy 3 refreshes the ids whose ``logits`` (``[rows, E]``) are within ``thr`` of their + row's lowest routed logit, in fp32 as the kernel does, or the ids of a given ``near_miss`` mask. + """ + self.step += 1 + base = layer * self.E + tok = now = 0 + if self.policy != 0: + if bump_tok: + self.tok += 1 + tok = self.tok + now = tok * self.L + layer + if self.policy == 3 and update_state and logits is not None: + lg = np.asarray(logits, dtype=np.float32).reshape(-1, self.E) + routed = np.asarray(ids, dtype=np.int64).reshape(lg.shape[0], -1) + kth = lg[np.arange(lg.shape[0])[:, None], routed].min(axis=1) + near_miss = lg >= (kth - np.float32(thr))[:, None] + if self.policy == 3 and update_state and near_miss is not None: + near = base + np.flatnonzero(np.asarray(near_miss).reshape(-1, self.E).any(0)) + self.last_tok[near] = np.maximum(self.last_tok[near], tok - 1) + q = np.asarray(ids, dtype=np.int64).reshape(-1) + base + slots = self.slot_of_id[q] + uniq = np.unique(q) + held_at = self.slot_of_id[uniq] + self.useful = int(((held_at >= 0) & (self.usage[np.maximum(held_at, 0)] == 0)).sum()) + self.usage[slots[slots >= 0]] = self.step + missing = uniq[self.slot_of_id[uniq] < 0] # ascending, as the kernel ranks them + src = missing - base + dst = np.empty(missing.size, np.int64) + if missing.size: + keys = self._keys(layer, tok, now, pin_since, lowpri) + part = np.argpartition(keys, missing.size - 1)[: missing.size] + victims = part[np.argsort(keys[part])] + for i, (e, v) in enumerate(zip(missing, victims)): + old = self.id_of_slot[v] + if old >= 0: + self.slot_of_id[old] = -1 + self.id_of_slot[v] = e + self.slot_of_id[e] = v + self.usage[v] = self.step + dst[i] = v + out = self.slot_of_id[q] + if self.policy != 0 and update_state: + self.last_tok[uniq] = tok + if self.policy >= 2: + lc = self.lc[uniq] + x = self._decayed(uniq, now) + self.lc[uniq] = np.where(lc == se.LC_NEVER, 0, self._softplus2(x)) + self.ct[uniq] = now + return out, src, dst + + def prefetch(self, layer, sel, budget=None): + """``prefetch_ensure``: install ``sel``'s non-resident ids (-1 = padding, query order), at most + the first ``budget``, as held slots with usage 0, without bumping a clock or touching a hit; + the last demand call's slots (usage == step) stay pinned. Returns ``(src_ids, dst_slots)``, + the copy plan.""" + base = layer * self.E + tok = self.tok if self.policy else 0 + pos = layer - 1 # the kernel scores from the previous layer's call, where the prefetch runs + now = tok * self.L + pos + order = [] + for e in sel: + if e >= 0 and e + base not in order: + order.append(int(e) + base) + missing = [i for i in order if self.slot_of_id[i] < 0] + keys = self._keys(pos, tok, now, lowpri=True) + n = min(len(missing), int((keys != _KEY_MAX).sum()), len(missing) if budget is None else budget) + victims = np.argsort(keys, kind="stable")[:n] + for e, v in zip(missing[:n], victims): + old = self.id_of_slot[v] + if old >= 0: + self.slot_of_id[old] = -1 + self.id_of_slot[v] = e + self.slot_of_id[e] = v + self.usage[v] = 0 + return np.asarray(missing[:n], np.int64) - base, victims.astype(np.int64) + + def materialize(self, layer): + """The prefill materialize kernel: the whole layer into slots [0, E), position == expert id.""" + base = layer * self.E + same = (self.id_of_slot >= base) & (self.id_of_slot < base + self.E) + self.id_of_slot[same] = -1 + self.usage[same] = 0 + old = self.id_of_slot[: self.E].copy() + self.slot_of_id[old[old >= 0]] = -1 + self.step += 1 + self.id_of_slot[: self.E] = base + np.arange(self.E) + self.slot_of_id[base : base + self.E] = np.arange(self.E) + self.usage[: self.E] = self.step + + +def replay(cache: RefScoredCache, rows, logits=None, thr=0.0, near_miss=None): + """Decode ``rows`` ``[N, L, K]`` (one step per row, all layers); per-row miss counts. + + ``rows`` may be ``[N, L, B, K]`` for a batch of B rows per step, one K*B call per layer. + Policy 3 takes the matching ``[N, L, B, E]`` router ``logits`` (or a ``near_miss`` mask). + """ + rows = np.asarray(rows) + misses = np.zeros(rows.shape[0], np.int64) + for r in range(rows.shape[0]): + for layer in range(rows.shape[1]): + lg = None if logits is None else logits[r, layer] + nm = None if near_miss is None else near_miss[r, layer] + _, src, _ = cache.ensure(layer, rows[r, layer], bump_tok=layer == 0, logits=lg, thr=thr, near_miss=nm) + misses[r] += src.size + return misses diff --git a/tests/moe/test_cpu_moe.py b/tests/moe/test_cpu_moe.py index 547c94c60a..38e79b2743 100644 --- a/tests/moe/test_cpu_moe.py +++ b/tests/moe/test_cpu_moe.py @@ -440,6 +440,111 @@ def _reference(cache, layer, hidden, ids, w, act="silu", apply_in=False): return out +@pytest.fixture(scope="module") +def hip_probe_fault_library(tmp_path_factory): + import importlib.util + import os + from pathlib import Path + import shutil + import subprocess + + if not torch.version.hip: + pytest.skip("requires the HIP sync backend") + compiler = shutil.which("c++") + if compiler is None: + pytest.skip("requires a C++ compiler for HIP fault injection") + candidates = [Path(os.getenv("ROCM_HOME", "/opt/rocm")), Path("/opt/rocm")] + spec = importlib.util.find_spec("_rocm_sdk_core") + if spec and spec.submodule_search_locations: + candidates.append(Path(next(iter(spec.submodule_search_locations)))) + include = next((root / "include" for root in candidates if (root / "include/hip/hip_runtime.h").exists()), None) + if include is None: + pytest.skip("requires HIP SDK headers for fault injection") + library = tmp_path_factory.mktemp("hip_probe") / "faults.so" + command = [ + compiler, "-std=c++20", "-shared", "-fPIC", "-D__HIP_PLATFORM_AMD__=1", + f"-I{include}", str(Path(__file__).with_name("hip_probe_faults.cpp")), + "-ldl", "-o", str(library), + ] + result = subprocess.run(command, capture_output=True, text=True, timeout=60) + assert result.returncode == 0, result.stdout + result.stderr + return library + + +@pytest.mark.parametrize("fault", ["dependency", "end-capture", "pending"]) +def test_cpu_moe_hip_probe_error_recovery(hip_probe_fault_library, fault): + """Rejected probes leave eager usable and preserve errors predating the probe.""" + import os + from pathlib import Path + import subprocess + import sys + import textwrap + + script = textwrap.dedent(""" + import ctypes + import os + import runpy + import sys + + import pytest + import torch + from freetoken.kernel import _cpu_moe + from freetoken.moe.cpu_executor import CpuMoeExecutor + + shim = ctypes.CDLL(os.environ["LD_PRELOAD"]) + stream_ptr = torch.cuda.current_stream().cuda_stream + fault = os.environ["FREETOKEN_TEST_HIP_PROBE_FAULT"] + if fault == "pending": + assert shim.probe_test_make_error() != 0 + pending = shim.probe_test_peek_error() + with pytest.raises(RuntimeError, match="pending error before probe"): + _cpu_moe.memops_probe(stream_ptr) + assert shim.probe_test_peek_error() == pending + assert shim.probe_test_begins() == 0 + sys.exit(0) + + helpers = runpy.run_path(sys.argv[1]) + torch.manual_seed(17) + cache = helpers["_make_cache"](2, 8, 256, 128, scale=0.03) + ex = CpuMoeExecutor( + cache, top_k=4, activation="silu", apply_router_weight_on_input=False, + num_threads=4, max_tokens=3, device=torch.device("cuda"), + ) + assert not ex._flag_sync and not ex._graph_capture_safe + assert shim.probe_test_injected() > 0 + assert shim.probe_test_begins() == shim.probe_test_ends() + assert shim.probe_test_peek_error() == 0 + if fault == "dependency": + assert shim.probe_test_removed_nodes() == 1 + assert shim.probe_test_last_end_rc() == 0 + else: + assert shim.probe_test_last_end_rc() != 0 + + for index in range(2): + torch.manual_seed(23 + index) + hidden = (torch.randn(3, 256) * 0.3).bfloat16().cuda() + ids = torch.randint(0, 8, (3, 4), dtype=torch.int32).cuda() + weights = torch.rand(3, 4).cuda() + actual = ex.decode(1, hidden, weights, ids).float().cpu() + expected = helpers["_reference"](cache, 1, hidden, ids.cpu(), weights.cpu()) + relative = (actual - expected).abs().max() / (expected.abs().max() + 1e-6) + assert relative < 2e-2, relative + + stream = torch.cuda.Stream() + with pytest.raises(RuntimeError, match="--cuda-graph-max-bs 0"): + with torch.cuda.graph(torch.cuda.CUDAGraph(), stream=stream): + ex.decode(1, hidden, weights, ids) + assert torch.isfinite(ex.decode(1, hidden, weights, ids).cpu()).all() + """) + env = dict(os.environ, LD_PRELOAD=str(hip_probe_fault_library), + FREETOKEN_TEST_HIP_PROBE_FAULT=fault, FREETOKEN_CPU_MOE_FLAG_SYNC="1") + result = subprocess.run( + [sys.executable, "-c", script, str(Path(__file__).resolve())], + env=env, capture_output=True, text=True, timeout=30, + ) + assert result.returncode == 0, result.stdout + result.stderr + + def test_cpu_moe_decode_cuda_graph_replay(): from freetoken.moe.cpu_executor import CpuMoeExecutor @@ -467,7 +572,7 @@ def test_cpu_moe_decode_cuda_graph_replay(): ids = torch.randint(0, E, (bs, top_k), device=dev, dtype=torch.int32) w = torch.rand(bs, top_k, device=dev, dtype=torch.float32) - # Eager forward: also materializes the per-bs pinned buffers + per-(layer,bs) task. + # Eager forward. out_eager = ex.decode(layer, hidden, w, ids) torch.cuda.synchronize() ref_eager = _reference(cache, layer, hidden, ids.cpu(), w.cpu()) @@ -500,11 +605,9 @@ def test_cpu_moe_decode_cuda_graph_replay(): # must actually have exercised it -- guard against a silent fallback to the # host-func path making this test vacuous for the capture-embedded memops. if ex._flag_sync: - assert (layer, bs) in ex._flag_slots, "flag slot expected for the decode task" - slot = ex._flag_slots[(layer, bs)] - assert ex._ext.flag_served_count(slot) >= 4, "1 eager + 3 replay dispatches expected" - assert int(ex._done[slot]) == 1 and int(ex._ready[slot]) == 0, "handshake at rest" - assert int(ex._err.sum()) == 0, "watchdog must not fire in normal operation" + assert ex._ext.flag_served() >= 4, "1 eager + 3 replay dispatches expected" + assert not ex._ext.flag_pending(), "handshake at rest" + assert not ex._unhealthy, "watchdog must not fire in normal operation" ex.raise_if_unhealthy() print("cpu moe cuda graph replay OK") diff --git a/tests/moe/test_cpu_moe_fp8_block.py b/tests/moe/test_cpu_moe_fp8_block.py index 06f954d8d9..baecb86f9a 100644 --- a/tests/moe/test_cpu_moe_fp8_block.py +++ b/tests/moe/test_cpu_moe_fp8_block.py @@ -74,11 +74,11 @@ def test_cpu_fp8_block_accepts_qwen38_padded_scale_bank(): device=torch.device("cpu"), ) hidden = torch.randn(1, H, dtype=torch.bfloat16) - io = ex._io_for(1) + io = ex._io io["x"].copy_(hidden) io["ids"].zero_() io["w"].fill_(1.0) - ex._ext.run_task(ex._task_for(0, 1)) + ex._ext.run(0, 1) cpu_out = io["y"][0].float().to(torch.float64) banks = cache.bank_sources gate_up = _reference_fp8_block_gemv(banks["gate_up"][0], banks["gate_up_scale"][0], hidden[0]) @@ -140,11 +140,8 @@ def _reference_fp8_block_gemv(weights, scales, x): @pytest.mark.parametrize("codes", ["finite", "subnormal"]) -@pytest.mark.parametrize( - "isa, expected_isa", - [("scalar", "scalar"), ("avx2", "avx2"), ("avx512", "avx512f"), ("avx512bf16", "avx512f")], -) -def test_cpu_fp8_block_matches_float64_reference(monkeypatch, isa, expected_isa, codes): +@pytest.mark.parametrize("isa", ["scalar", "avx2", "avx512"]) +def test_cpu_fp8_block_matches_float64_reference(monkeypatch, isa, codes): from freetoken.moe.cpu_executor import CpuMoeExecutor H, I = 259, 257 # ragged K for both GEMVs; 257 down rows exercise many reductions @@ -159,17 +156,16 @@ def test_cpu_fp8_block_matches_float64_reference(monkeypatch, isa, expected_isa, max_tokens=1, device=torch.device("cpu"), ) - tiers = {"scalar": 0, "avx2": 1, "avx512": 2, "avx512bf16": 3} + tiers = {"scalar": 0, "avx2": 1, "avx512": 2} if tiers[isa] > ("scalar", "avx2", "avx512f", "avx512bf16").index(ex.isa): pytest.skip(f"requested {isa}, host/build selected {ex.isa}") - assert ex._ext.fp8_block_isa_name() == expected_isa hidden = torch.randn(1, H, dtype=torch.bfloat16) - io = ex._io_for(1) + io = ex._io io["x"].copy_(hidden) io["ids"].zero_() io["w"].fill_(1.0) - ex._ext.run_task(ex._task_for(0, 1)) + ex._ext.run(0, 1) cpu_out = io["y"][0].float().to(torch.float64) banks = cache.bank_sources @@ -233,9 +229,8 @@ def test_cpu_fp8_block_matches_gpu_decode_kernel(bs, isa, monkeypatch): max_tokens=bs, device=dev, ) - got = ex._ext.fp8_block_isa_name() - if got != _ISA_TIERS[isa]: - pytest.skip(f"asked for {isa}, kernel selected {got}: not supported here") + if ex.isa != _ISA_TIERS[isa]: + pytest.skip(f"asked for {isa}, host/build selected {ex.isa}: not supported here") hidden = torch.randn(bs, H, device=dev, dtype=torch.bfloat16) ids = torch.stack([torch.randperm(E, device=dev)[:top_k] for _ in range(bs)]).to(torch.int32) @@ -254,7 +249,7 @@ def test_cpu_fp8_block_matches_gpu_decode_kernel(bs, isa, monkeypatch): ).float() rel = (cpu_out - gpu_out).abs().max() / (gpu_out.abs().max() + 1e-6) - assert rel < 2e-2, f"bs={bs} isa={got} rel err {rel.item()}" + assert rel < 2e-2, f"bs={bs} isa={ex.isa} rel err {rel.item()}" def test_fp16_bits_carry_every_e4m3_byte(): diff --git a/tests/moe/test_cpu_moe_q4_0.py b/tests/moe/test_cpu_moe_q4_0.py index b5ec2f09b4..7babb5d4ce 100644 --- a/tests/moe/test_cpu_moe_q4_0.py +++ b/tests/moe/test_cpu_moe_q4_0.py @@ -1,7 +1,7 @@ """CPU MoE executor -- native GGUF Q4_0 experts (``--moe-backend cpu``/``hybrid`` for the gemma4 GGUF checkpoint). -The CPU W4A16 GEMV (``q4_0_dot`` in csrc/cpu_moe/cpu_moe_ext.cpp) reads the *same* +The CPU W4A16 GEMV (``q4_0_dot`` in csrc/cpu_moe/kernels/q4_0.cpp) reads the *same* packed Q4_0 banks the GPU offload path streams and dequantizes weights inside the K-loop. We check it against the reference dequant (models/gguf/dequant.py) + the production bf16 GPU decode kernel on byte-identical banks: both are W4A16, so the diff --git a/tests/moe/test_dsfp4_prequant.py b/tests/moe/test_dsfp4_prequant.py index 11519502e2..8187402d3f 100644 --- a/tests/moe/test_dsfp4_prequant.py +++ b/tests/moe/test_dsfp4_prequant.py @@ -17,8 +17,9 @@ @pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") -def test_dsfp4_gpu_prequant_bit_parity(): +def test_dsfp4_gpu_prequant_bit_parity(monkeypatch): from freetoken.kernel.pinned import alloc_pinned_tensor + from freetoken.moe import cpu_executor from freetoken.moe.cpu_executor import CpuMoeExecutor torch.manual_seed(0) @@ -44,14 +45,13 @@ def test_dsfp4_gpu_prequant_bit_parity(): wts = torch.rand(2, TOPK, device=dev) def run(cpu_roundtrip: bool) -> torch.Tensor: + # control: the original CPU-side scalar round-trip + monkeypatch.setattr(cpu_executor, "_DSFP4_GPU_PREQUANT", not cpu_roundtrip) ex = CpuMoeExecutor( cache, top_k=TOPK, activation="silu", apply_router_weight_on_input=False, num_threads=4, max_tokens=2, device=dev, swiglu_limit=10.0, act_block=128, ) - assert ex._gpu_prequant, "ds_fp4 on CUDA should enable GPU prequant" - if cpu_roundtrip: # control: the original CPU-side scalar round-trip - ex._gpu_prequant = False - ex._ext.set_input_prequant(False) + assert ex._gpu_prequant != cpu_roundtrip, "ds_fp4 on CUDA should enable GPU prequant" out = ex.decode(0, x, wts, ids.clone()).clone() torch.cuda.synchronize() del ex diff --git a/tests/moe/test_dsfp4_reference_semantics.py b/tests/moe/test_dsfp4_reference_semantics.py index 2f89cf334d..8aef4e661e 100644 --- a/tests/moe/test_dsfp4_reference_semantics.py +++ b/tests/moe/test_dsfp4_reference_semantics.py @@ -50,13 +50,16 @@ def _fp8_roundtrip(x: torch.Tensor, block: int) -> torch.Tensor: def reference(x, slots, weights, banks, block): """The MoE over dequantized banks: per-route bf16 expert outputs accumulated into an fp32 ``y``, - returned in fp32 (the kernels round that sum to bf16).""" + returned in fp32 (the kernels round that sum to bf16). ``tie`` is per element the largest bf16 + ulp among routes sitting on a rounding tie, where a different fp32 accumulation order can flip + the rounded result.""" gu_p, gu_s, dn_p, dn_s = banks W13 = _dequant(gu_p, gu_s) # [E, 2I, H] W2 = _dequant(dn_p, dn_s) # [E, H, I] xq = _fp8_roundtrip(x, block).float() T = x.shape[0] y = torch.zeros(T, H, dtype=torch.float32, device=x.device) + tie = torch.zeros_like(y) for t in range(T): for r in range(TOP_K): e = int(slots[t, r]) @@ -64,8 +67,25 @@ def reference(x, slots, weights, banks, block): gate, up = gu[:I].clamp(max=LIMIT), gu[I:].clamp(-LIMIT, LIMIT) h = (torch.nn.functional.silu(gate) * up).to(torch.bfloat16) hq = _fp8_roundtrip(h.view(1, -1), block).float().view(-1) - y[t] += (weights[t, r].float() * (hq @ W2[e].T)).to(torch.bfloat16).float() - return y + raw = weights[t, r].float() * (hq @ W2[e].T) + out = raw.to(torch.bfloat16).float() + y[t] += out + ulp = torch.exp2(torch.floor(torch.log2(out.abs())) - 7.0) + # a differing fp32 sum order (~1e-5 relative) only flips a value within noise of a + # rounding tie: distance to the rounded value at ulp/2; 10% of an ulp is generous + flippable = 0.5 * ulp - (raw - out).abs() < 0.1 * ulp + tie[t] = torch.maximum(tie[t], ulp * flippable) + return y, tie + + +def _assert_close(got, want, tie, atol=2e-2, rtol=2e-2): + # torch.testing.assert_close(atol, rtol) widened per element by ``tie``: a different fp32 + # accumulation order may flip one per-route bf16 rounding tie, moving y by one ulp + bad = (got - want).abs() > atol + rtol * want.abs() + tie + assert not bad.any(), ( + f"Mismatched elements: {int(bad.sum())} / {want.numel()}, " + f"greatest absolute difference: {(got - want).abs().max():.6f}" + ) def _inputs(T, device, seed=1): @@ -84,8 +104,8 @@ def test_gemv_path_matches_the_transcription(block): banks = _banks("cuda") x, slots, w = _inputs(6, "cuda") got = routed_experts_fp4(x, slots, w, *banks, LIMIT, act_block=block).float() - want = reference(x, slots, w, banks, block) - torch.testing.assert_close(got, want, atol=2e-2, rtol=2e-2) + want, tie = reference(x, slots, w, banks, block) + _assert_close(got, want, tie) @pytest.mark.parametrize("block", [32, 128]) @@ -109,8 +129,8 @@ def test_cpu_executor_matches_the_transcription(block): got = ex.decode(0, x, w, slots.clone()).clone().float() torch.cuda.synchronize() del ex - want = reference(x, slots, w, (gu_p, gu_s, dn_p, dn_s), block) - torch.testing.assert_close(got, want, atol=2e-2, rtol=2e-2) + want, tie = reference(x, slots, w, (gu_p, gu_s, dn_p, dn_s), block) + _assert_close(got, want, tie) def test_grouped_prefill_path_matches_the_transcription(): @@ -119,8 +139,8 @@ def test_grouped_prefill_path_matches_the_transcription(): banks = _banks("cuda") x, slots, w = _inputs(64, "cuda") got = fused_ds_fp4.routed_experts_fp4_prefill(x, slots, w, *banks, LIMIT, E, act_block=32).float() - want = reference(x, slots, w, banks, 32) - torch.testing.assert_close(got, want, atol=2e-2, rtol=2e-2) + want, tie = reference(x, slots, w, banks, 32) + _assert_close(got, want, tie) @pytest.mark.parametrize("block", [32, 128]) @@ -154,7 +174,8 @@ def test_inactive_routes_contribute_zero_without_reading_their_slots(block): # and each share matches the reference over its own routes (a zero weight drops a route) got_clean = routed_experts_fp4(x, gpu_slots, gpu_w, *clean, LIMIT, act_block=block) assert torch.equal(got, got_clean) - torch.testing.assert_close(got.float(), reference(x, slots, gpu_w, clean, block), atol=2e-2, rtol=2e-2) + want, tie = reference(x, slots, gpu_w, clean, block) + _assert_close(got.float(), want, tie) # the CPU executor computes the complementary split from the same raw ids (ids < 0 skipped) pinned = {} for name, t in (("gate_up_packed", clean[0]), ("gate_up_scale", clean[1]), ("down_packed", clean[2]), ("down_scale", clean[3])): @@ -167,4 +188,5 @@ def test_inactive_routes_contribute_zero_without_reading_their_slots(block): cpu_out = ex.decode(0, x, w, cpu_ids.clone()).clone() torch.cuda.synchronize() del ex - torch.testing.assert_close(cpu_out.float(), reference(x, slots, torch.where(on_gpu, 0.0, w), clean, block), atol=2e-2, rtol=2e-2) + want, tie = reference(x, slots, torch.where(on_gpu, 0.0, w), clean, block) + _assert_close(cpu_out.float(), want, tie) diff --git a/tests/moe/test_expert_bank_threading.py b/tests/moe/test_expert_bank_threading.py new file mode 100644 index 0000000000..da74719448 --- /dev/null +++ b/tests/moe/test_expert_bank_threading.py @@ -0,0 +1,57 @@ +from types import SimpleNamespace + +import pytest +import torch +from freetoken.layers.quantization import QuantKind +from freetoken.moe.expert_banks import build_expert_banks + + +@pytest.mark.parametrize("fail", [False, True]) +def test_nvfp4_host_bank_copy_threads_are_restored(monkeypatch, fail): + threads = 24 + changes = [] + observed = [] + + def set_threads(value): + nonlocal threads + threads = value + changes.append(value) + + monkeypatch.setattr(torch, "get_num_threads", lambda: threads) + monkeypatch.setattr(torch, "set_num_threads", set_threads) + monkeypatch.setattr(torch.cuda, "is_available", lambda: False) + + class Method: + kind = QuantKind.NVFP4 + kernel = SimpleNamespace(name="triton") + cfg = SimpleNamespace(num_experts=2) + + @staticmethod + def layout(): + spec = SimpleNamespace(shape=(2,), dtype=torch.uint8, resident=False) + return {"gate_up": spec} + + @staticmethod + def pack(piece, out): + observed.append(threads) + if fail and piece["gate_up"][0, 0] == 2: + raise RuntimeError("failed during expert placement") + out["gate_up"].copy_(piece["gate_up"]) + return {} + + pieces = ( + (0, expert, expert + 1, { + "gate_up": torch.full((1, 2), expert + 1, dtype=torch.uint8), + }) + for expert in range(2) + ) + if fail: + with pytest.raises(RuntimeError, match="failed during expert placement"): + build_expert_banks(Method(), 1, pieces, device=torch.device("cpu")) + else: + banks = build_expert_banks(Method(), 1, pieces, device=torch.device("cpu")) + assert banks.sources["gate_up"][0].tolist() == [[1, 1], [2, 2]] + + assert observed == [1, 1] + assert changes == [1, 24] + assert threads == 24 diff --git a/tests/moe/test_fused_copy.py b/tests/moe/test_fused_copy.py index 7a0c086444..c5336ea117 100644 --- a/tests/moe/test_fused_copy.py +++ b/tests/moe/test_fused_copy.py @@ -1,6 +1,8 @@ """The fused multi-bank ``copy_missing`` path must move exactly the same bytes as the legacy per-bank ``fast_index_copy_jit`` loop, for every miss count (including the -zero-copy case), across banks of differing per-row sizes. +zero-copy case), across banks of differing per-row sizes. The opt-in slim kernel +(``FREETOKEN_MOE_SLIM_COPY``) must move exactly the same bytes as the fused one, also when +replayed from a CUDA graph with a device-side plan. """ from __future__ import annotations @@ -8,6 +10,7 @@ import pytest import torch +from freetoken.moe import offload_cache from freetoken.moe.offload_cache import _BANK_SCHEMAS, OffloadMoeCache CUDA = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") @@ -17,22 +20,30 @@ # kernel/aot_models.py must be listed in its TEST_FEATURE_SIZES so the per-bank kernels # stay prebuilt under FREETOKEN_DISABLE_JIT=1. FEATS = [8192, 512, 256, 4096, 512, 256] +# qwen4_exp (Qwen3.8-Flash-Next) triton-NVFP4 rows in _BANK_SCHEMAS["nvfp4"] order: gate_up +# packed/scale/global, down packed/scale/global -- 2,772,480 B per expert. +NVFP4_FEATS = [1638400, 204800, 2560, 819200, 102400, 5120] +# 16 * odd, several rows shorter than one grid pass, so segments end mid-round and the +# stride phase wraps across banks and rows. +ODD_FEATS = [16 * 1031, 16 * 7, 16 * 250, 16, 16 * 4099, 16 * 33] -def _build_cache(num_layers, num_experts, cache_size): +def _build_cache(num_layers, num_experts, cache_size, quant_format="mxfp4_triton", feats=FEATS, pinned=False): dev = torch.device("cuda") cache = OffloadMoeCache( num_layers=num_layers, num_experts=num_experts, cache_size=cache_size, - device=dev, cache_policy="lru", prefill_overlap=False, quant_format="mxfp4_triton", + device=dev, cache_policy="lru", prefill_overlap=False, quant_format=quant_format, ) - schema = _BANK_SCHEMAS["mxfp4_triton"] + schema = _BANK_SCHEMAS[quant_format] # Views into one flat tensor: only per-layer addressing matters here, not # independent allocations. - sources = { - name: list(torch.randint(0, 256, (num_layers * num_experts, feat), dtype=torch.uint8, device=dev) - .split(num_experts)) - for name, feat in zip(schema, FEATS) - } + sources = {} + for name, feat in zip(schema, feats): + if pinned: + flat = torch.randint(0, 256, (num_layers * num_experts, feat), dtype=torch.uint8).pin_memory() + else: + flat = torch.randint(0, 256, (num_layers * num_experts, feat), dtype=torch.uint8, device=dev) + sources[name] = list(flat.split(num_experts)) cache.set_bank_sources(sources) # also builds the fused-copy descriptor return cache @@ -74,3 +85,147 @@ def test_fused_copy_matches_per_bank(num_indices): for b, (r, (_, c)) in enumerate(zip(ref, cache.banks)): assert torch.equal(r, c), f"bank {b} (feat={FEATS[b]}) fused != per-bank at num_indices={num_indices}" + + +def _host_banks(feats, n_src, seed): + from freetoken.kernel.pinned import device_ptr + + gen = torch.Generator().manual_seed(seed) + src = [torch.randint(0, 256, (n_src, f), dtype=torch.uint8, generator=gen).pin_memory() for f in feats] + dev = torch.device("cuda") + src_ptrs = torch.tensor([device_ptr(t) for t in src], dtype=torch.int64, device=dev) + feat_bytes = torch.tensor(feats, dtype=torch.int64, device=dev) + return src, src_ptrs, feat_bytes + + +def _slot_banks(feats, n_slots): + dst = [torch.full((n_slots, f), 0xAB, dtype=torch.uint8, device="cuda") for f in feats] + return dst, torch.tensor([t.data_ptr() for t in dst], dtype=torch.int64, device="cuda") + + +def _expected(src, feats, n_slots, di, si, num): + # untouched slots keep the 0xAB fill; plan rows land at their dst slots + out = [] + for b, f in enumerate(feats): + e = torch.full((n_slots, f), 0xAB, dtype=torch.uint8) + e[di[:num].long().cpu()] = src[b][si[:num].long().cpu()] + out.append(e.cuda()) + return out + + +def _random_plan(gen, n_src, n_slots, plan, dtype): + di = torch.randperm(n_slots, generator=gen)[:plan].to(dtype).cuda() # distinct: no write races + si = torch.randint(0, n_src, (plan,), generator=gen).to(dtype).cuda() + return di, si + + +def _check_slim_matches_multi(feats, idx_dtype, **grid): + from freetoken.kernel.fast_index_copy import fast_index_copy_multi_jit, fast_index_copy_multi_slim_jit + + n_src, n_slots, plan = 24, 16, 16 + src, src_ptrs, feat_bytes = _host_banks(feats, n_src, seed=0) + gen = torch.Generator().manual_seed(1) + for num in [0, 1, 3, 10, plan, None]: # None: no num_indices, the whole index length + di, si = _random_plan(gen, n_src, n_slots, plan, idx_dtype) + count = None if num is None else torch.tensor([num], dtype=torch.int64, device="cuda") + ref, ref_ptrs = _slot_banks(feats, n_slots) + fast_index_copy_multi_jit(ref_ptrs, src_ptrs, feat_bytes, di, si, count) + out, out_ptrs = _slot_banks(feats, n_slots) + fast_index_copy_multi_slim_jit(out_ptrs, src_ptrs, feat_bytes, di, si, count, **grid) + torch.cuda.synchronize() + exp = _expected(src, feats, n_slots, di, si, plan if num is None else num) + for b, f in enumerate(feats): + assert torch.equal(ref[b], exp[b]), f"multi reference wrong: bank {b} (feat={f}) num={num}" + assert torch.equal(out[b], ref[b]), f"slim != multi: bank {b} (feat={f}) num={num} grid={grid}" + + +@CUDA +@pytest.mark.parametrize("idx_dtype", [torch.int32, torch.int64]) +@pytest.mark.parametrize("feats", [NVFP4_FEATS, ODD_FEATS], ids=["nvfp4", "odd"]) +def test_slim_copy_matches_multi(feats, idx_dtype): + _check_slim_matches_multi(feats, idx_dtype) + + +@CUDA +@pytest.mark.slow +@pytest.mark.parametrize("grid", [(3, 128, 2), (16, 256, 2), (5, 64, 8)]) +@pytest.mark.parametrize("feats", [NVFP4_FEATS, ODD_FEATS], ids=["nvfp4", "odd"]) +def test_slim_copy_matches_multi_other_grids(feats, grid): + blocks, threads, unroll = grid + _check_slim_matches_multi(feats, torch.int32, num_blocks=blocks, num_threads=threads, unroll=unroll) + + +@CUDA +def test_slim_copy_graph_replay_follows_device_plan(): + from freetoken.kernel.fast_index_copy import fast_index_copy_multi_slim_jit + + feats, n_src, n_slots, plan = NVFP4_FEATS, 24, 16, 16 + src, src_ptrs, feat_bytes = _host_banks(feats, n_src, seed=2) + dst, dst_ptrs = _slot_banks(feats, n_slots) + di = torch.zeros(plan, dtype=torch.int32, device="cuda") + si = torch.zeros(plan, dtype=torch.int32, device="cuda") + count = torch.zeros(1, dtype=torch.int64, device="cuda") + fast_index_copy_multi_slim_jit(dst_ptrs, src_ptrs, feat_bytes, di, si, count) # JIT outside capture + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + fast_index_copy_multi_slim_jit(dst_ptrs, src_ptrs, feat_bytes, di, si, count) + + gen = torch.Generator().manual_seed(3) + for num in [4, 0, plan, 1, 10]: + new_di, new_si = _random_plan(gen, n_src, n_slots, plan, torch.int32) + di.copy_(new_di) + si.copy_(new_si) + count.fill_(num) + for d in dst: + d.fill_(0xAB) + graph.replay() + torch.cuda.synchronize() + exp = _expected(src, feats, n_slots, di, si, num) + for b, f in enumerate(feats): + assert torch.equal(dst[b], exp[b]), f"graph replay: bank {b} (feat={f}) num={num}" + + +def test_slim_copy_flag_defaults_off(monkeypatch): + import importlib.util + + monkeypatch.delenv("FREETOKEN_MOE_SLIM_COPY", raising=False) + spec = importlib.util.find_spec("freetoken.moe.offload_cache") + fresh = importlib.util.module_from_spec(spec) + spec.loader.exec_module(fresh) # a private copy: the imported module and its classes stay untouched + assert fresh._SLIM_COPY is False + monkeypatch.setenv("FREETOKEN_MOE_SLIM_COPY", "1") + spec.loader.exec_module(fresh) + assert fresh._SLIM_COPY is True + + +@CUDA +@pytest.mark.parametrize("slim", [False, True], ids=["multi", "slim"]) +def test_copy_missing_slim_switch(monkeypatch, slim): + from freetoken.kernel import fast_index_copy as fic + + calls = [] + for name in ("fast_index_copy_multi_jit", "fast_index_copy_multi_slim_jit"): + real = getattr(fic, name) + monkeypatch.setattr(fic, name, lambda *a, _n=name, _r=real, **k: (calls.append(_n), _r(*a, **k))) + monkeypatch.setattr(offload_cache, "_SLIM_COPY", slim) + + num_layers, num_experts, cache_size, layer_id, num = 2, 8, 16, 1, 5 + cache = _build_cache(num_layers, num_experts, cache_size, "nvfp4", NVFP4_FEATS, pinned=True) + assert cache._copy_fused_ok + cache._pending_src_layer = layer_id + dev = torch.device("cuda") + cache.num_indices.fill_(num) + cache.evict_slots[:num] = torch.tensor([9, 2, 15, 0, 7], dtype=torch.int32, device=dev) + cache.src_indices[:num] = torch.tensor([3, 3, 0, 7, 5], dtype=torch.int32, device=dev) + for _, c in cache.banks: + c.fill_(0xAB) + cache.copy_missing() + torch.cuda.synchronize() + + assert calls == ["fast_index_copy_multi_slim_jit" if slim else "fast_index_copy_multi_jit"] + slots, rows = cache.evict_slots[:num].long().cpu(), cache.src_indices[:num].long().cpu() + for b, (per_layer, c) in enumerate(cache.banks): + exp = torch.full(c.shape, 0xAB, dtype=torch.uint8) + exp[slots] = per_layer[layer_id][rows] + assert torch.equal(c.cpu(), exp), f"bank {b} (feat={NVFP4_FEATS[b]})" diff --git a/tests/moe/test_hybrid_fetch.py b/tests/moe/test_hybrid_fetch.py index a26630b894..fad4fa7843 100644 --- a/tests/moe/test_hybrid_fetch.py +++ b/tests/moe/test_hybrid_fetch.py @@ -1,43 +1,15 @@ """Hybrid decode's bandwidth-matched fetch split. -Covers the two halves of --moe-hybrid-max-fetch auto: the profile reader that turns -`ft bench bw` kernel bandwidths into a fetch fraction, and the ensure kernel's -per-step integer split (GPU kernel vs CPU reference mirror, and the balance rule). +Covers the profile reader of --moe-hybrid-max-fetch auto, which turns `ft bench bw` +kernel bandwidths into a fetch fraction. """ import json import os import pytest -import torch from freetoken.moe.bench_profile import PROFILE_VERSION, default_profile_path, load_backend_recommendation, load_hybrid_fetch_fraction -from freetoken.moe.offload_cache import OffloadMoeCache - -Q = 1 << 16 - - -def _balanced_fetch(num_missing: int, frac_q16: int) -> int: - """Reference split: F ~ frac * misses, rounded to whichever integer neighbor - minimizes the slower overlapped side (fetch ~ F*(1-frac), CPU ~ (M-F)*frac).""" - lo = (num_missing * frac_q16) >> 16 - cost = lambda f: max(f * (Q - frac_q16), (num_missing - f) * frac_q16) # noqa: E731 - return min(num_missing, lo if cost(lo) <= cost(lo + 1) else lo + 1) - - -def test_balanced_fetch_tracks_fraction(): - # The split follows fetched : cpu = pcie : (cpu - pcie) up to integer rounding, and - # never over/under-shoots by more than one expert. - for frac in (0.1, 0.415, 0.454, 0.7, 1.0): - q = round(frac * Q) - for m in range(0, 65): - f = _balanced_fetch(m, q) - assert 0 <= f <= m - assert abs(f - frac * m) <= 1.0 - # ceil would over-fetch here (the regression this rule fixed): 41.5% of 3 misses is - # 1.24 -> fetching 2 makes the PCIe side ~1.6x slower than balance; keep it at 1. - assert _balanced_fetch(3, round(0.415 * Q)) == 1 - assert _balanced_fetch(4, round(0.415 * Q)) == 2 def test_load_hybrid_fetch_fraction(tmp_path): @@ -87,46 +59,3 @@ def write(path, name, verdict): # this card's own file wins over the legacy one write(default_profile_path(uuid), "FAKE GPU", "offload") assert load_backend_recommendation("bf16", gpu_name="FAKE GPU", gpu_uuid=uuid) == "offload" - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") -def test_hybrid_fraction_gpu_matches_cpu_reference(): - torch.manual_seed(0) - num_experts, cache_size, top_k, frac = 32, 40, 8, 0.415 - - def make(): - return OffloadMoeCache( - num_layers=2, num_experts=num_experts, cache_size=cache_size, - device=torch.device("cuda"), quant_format="bf16", decode_target="hybrid", - hybrid_max_fetch=num_experts, hybrid_fetch_fraction=frac, - ) - - gpu, ref = make(), make() - frac_q16 = round(frac * Q) - for step in range(64): - ids = torch.randperm(num_experts)[:top_k].to(torch.int32) - g, c = ids.clone().cuda(), ids.clone() # a CPU ids tensor drives the reference path - gpu.ensure_experts_hybrid(0, g) - ref.ensure_experts_hybrid(0, c) - missing = int(gpu.num_missing_full.item()) - fetched = int(gpu.num_indices.item()) - assert missing == int(ref.num_missing_full.item()) - assert fetched == int(ref.num_indices.item()) == _balanced_fetch(missing, frac_q16) - # slot rewrites (hit/fetched -> slot, overflow -> -1) and LRU state stay identical - assert torch.equal(g.cpu(), c) - assert torch.equal(gpu.slot_for_id.cpu(), ref.slot_for_id.cpu()) - assert torch.equal(gpu.id_of_slot.cpu(), ref.id_of_slot.cpu()) - assert (g >= 0).sum().item() == len(set(ids.tolist())) - (missing - fetched) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") -def test_hybrid_fixed_cap_unchanged(): - # fraction 0 (no profile / explicit --moe-hybrid-max-fetch) keeps the fixed cap. - cache = OffloadMoeCache( - num_layers=1, num_experts=32, cache_size=40, device=torch.device("cuda"), - quant_format="bf16", decode_target="hybrid", hybrid_max_fetch=1, - ) - ids = torch.arange(8, dtype=torch.int32).cuda() - cache.ensure_experts_hybrid(0, ids) - assert int(cache.num_missing_full.item()) == 8 - assert int(cache.num_indices.item()) == 1 diff --git a/tests/moe/test_moe_collect_stats.py b/tests/moe/test_moe_collect_stats.py new file mode 100644 index 0000000000..bd1c4e9b3f --- /dev/null +++ b/tests/moe/test_moe_collect_stats.py @@ -0,0 +1,131 @@ +"""--moe-collect-stats: the flag, and the report it produces. + +The counters themselves are accumulated device-side inside ``ensure_experts`` and were +already covered; what was missing until this flag existed was any way to turn them on from +the command line or read them back. These tests cover that wiring -- the flag reaching +``ServerArgs``, and the emit formatting the numbers and resetting the window afterwards. +""" + +import contextlib +import io +from types import SimpleNamespace + +import torch + +from freetoken.engine.engine import MOE_STATS_INTERVAL, Engine +from freetoken.moe.offload_cache import OffloadMoeCache +from freetoken.server.args import ServerArgs, parse_args + + +def test_flag_is_registered_and_defaults_off(): + """``--help`` short-circuits before the model is resolved, so this needs no checkpoint.""" + buf = io.StringIO() + with contextlib.redirect_stdout(buf), contextlib.suppress(SystemExit): + parse_args(["--help"]) + assert "--moe-collect-stats" in buf.getvalue() + # Off unless asked for: the counters ride in the decode CUDA graph and cost throughput. + assert ServerArgs.moe_collect_stats is False + + +class _StubCache: + """Just enough cache to exercise the emit: the four readers plus the window reset.""" + + def __init__(self, decode_target="gpu", layer_calls=512): + self.collect_stats = True + self.decode_target = decode_target + self._layer_calls = layer_calls + self.reset_calls = 0 + + def decode_miss_stats(self): + return { + "layer_calls": self._layer_calls, + "active_per_layer": 8.0, + "missing_per_layer": 2.0, + "miss_rate": 0.25, + "fetched_per_layer": 1.5, + "cpu_per_layer": 0.5, + "fetch_rate": 0.75, + "prefill_hit_rows": 0, + "prefill_rows": 0, + } + + def decode_miss_stats_per_layer(self): + return { + "per_layer": [ + {"layer": 0, "steps": 4, "miss_rate": 0.5}, + {"layer": 1, "steps": 4, "miss_rate": 0.1}, + # steps == 0 means the layer never ran in this window; it must not be + # ranked as a 0.0-miss-rate "best" layer. + {"layer": 2, "steps": 0, "miss_rate": 0.0}, + ] + } + + def decode_routing_stats(self): + return { + "slots_per_layer": 56.7, + "working_set_mean": 173.1, + "working_set_max": 243, + "experts_for_90pct": 92.3, + "static_topk_hit_at_slots": 0.764, + "norm_entropy": 0.813, + } + + def reset_stats(self): + self.reset_calls += 1 + + +def _emit(cache, caplog): + engine = SimpleNamespace(moe_offload_cache=cache, _emit_moe_stats=None) + with caplog.at_level("INFO"): + Engine._emit_moe_stats(engine) + return "\n".join(r.getMessage() for r in caplog.records) + + +def test_emit_reports_and_resets_the_window(caplog): + cache = _StubCache() + out = _emit(cache, caplog) + assert "miss_rate=0.250" in out + # Labelled as the fixed-set figure it is; "oracle" implied an upper bound it is not. + assert "static_topk_hit=0.764" in out + assert "oracle" not in out + assert "(realized 0.750)" in out + # Ranked worst-first, and the layer that never ran is left out entirely. + assert "L0=0.500, L1=0.100" in out + assert "L2" not in out + assert cache.reset_calls == 1 + + +def test_hybrid_split_only_reported_for_hybrid(caplog): + assert "fetch_rate" not in _emit(_StubCache(decode_target="gpu"), caplog) + caplog.clear() + assert "fetch_rate=0.750" in _emit(_StubCache(decode_target="hybrid"), caplog) + + +def test_idle_window_emits_nothing_and_keeps_counters(caplog): + """No decode ran, so there is nothing to report -- and nothing to reset either.""" + cache = _StubCache(layer_calls=0) + assert _emit(cache, caplog) == "" + assert cache.reset_calls == 0 + + +def test_interval_is_a_sane_window(): + assert MOE_STATS_INTERVAL >= 1 + + +def test_static_topk_hit_is_not_an_upper_bound_on_lru(): + """One slot, four experts each routed in a run of four tokens: A A A A B B B B ... + + The best fixed expert catches 4 of 16 tokens; an LRU holding one slot misses only on + each run's first token. The figure is routing skew, not a ceiling on a dynamic cache. + """ + trace = [e for e in range(4) for _ in range(4)] + cached, lru_hits = None, 0 + for e in trace: + lru_hits += e == cached + cached = e + freq = torch.bincount(torch.tensor(trace), minlength=4).unsqueeze(0) + cache = SimpleNamespace(decode_freq=freq, cache_size=1, num_layers=1, num_experts=4) + stats = OffloadMoeCache.decode_routing_stats(cache) + assert stats["static_topk_hit_at_slots"] == 0.25 + assert "oracle_hit_at_slots" not in stats + assert lru_hits / len(trace) == 0.75 diff --git a/tests/moe/test_nvfp4_backends.py b/tests/moe/test_nvfp4_backends.py index 91abeee1a8..8dd08cd71e 100644 --- a/tests/moe/test_nvfp4_backends.py +++ b/tests/moe/test_nvfp4_backends.py @@ -46,7 +46,9 @@ def _dequant_ref(packed: torch.Tensor, scale: torch.Tensor, row_global: torch.Te return w * s * row_global.float().unsqueeze(1) -def _make_native_sources(device: torch.device, seed: int = 0) -> dict[str, list[torch.Tensor]]: +def _make_native_sources( + device: torch.device, seed: int = 0, hidden: int = H, inter: int = I +) -> dict[str, list[torch.Tensor]]: """Random ModelOpt-style banks, CPU pinned, with one expert whose w1/w3 globals differ. One flat ``[L*E, ...]`` RNG draw (so seeding is unaffected) split into L @@ -61,15 +63,15 @@ def rand_u8(*shape): def rand_scale(*shape): return (torch.rand(*shape, generator=g) * 1.5 + 0.25).to(torch.float8_e4m3fn) - gate_up_global = torch.full((total, 2 * I), 1.0, dtype=torch.float16) - gate_up_global[:, I:] = 0.5 # w3 global != w1 global: exercises the alpha fold - down_global = torch.full((total, H), 0.75, dtype=torch.float16) + gate_up_global = torch.full((total, 2 * inter), 1.0, dtype=torch.float16) + gate_up_global[:, inter:] = 0.5 # w3 global != w1 global: exercises the alpha fold + down_global = torch.full((total, hidden), 0.75, dtype=torch.float16) flat = { - "gate_up_packed": rand_u8(total, 2 * I, H // 2), - "gate_up_scale": rand_scale(total, 2 * I, H // 16), + "gate_up_packed": rand_u8(total, 2 * inter, hidden // 2), + "gate_up_scale": rand_scale(total, 2 * inter, hidden // 16), "gate_up_global": gate_up_global, - "down_packed": rand_u8(total, H, I // 2), - "down_scale": rand_scale(total, H, I // 16), + "down_packed": rand_u8(total, hidden, inter // 2), + "down_scale": rand_scale(total, hidden, inter // 16), "down_global": down_global, } return {name: list(t.pin_memory().split(E)) for name, t in flat.items()} @@ -84,8 +86,9 @@ def _assert_close(out: torch.Tensor, ref: torch.Tensor) -> None: def _swigluoai_ref(h: torch.Tensor, alpha: float = 1.702, limit: float = 7.0) -> torch.Tensor: """MiniMax-M3 / gpt-oss clamped swiglu over UNINTERLEAVED [gate; up] halves.""" - gate = h[:I].clamp(max=limit) - up = h[I:].clamp(-limit, limit) + gate, up = h.chunk(2) + gate = gate.clamp(max=limit) + up = up.clamp(-limit, limit) return gate * torch.sigmoid(gate * alpha) * (up + 1.0) @@ -110,7 +113,8 @@ def _ref_moe(sources, layer_id, hidden, topk_weights, topk_ids, activation="silu if activation == "swigluoai": act = _swigluoai_ref(h) else: - act = torch.nn.functional.silu(h[:I]) * h[I:] + gate, up = h.chunk(2) + act = torch.nn.functional.silu(gate) * up out[t] += float(topk_weights[t, j]) * (dn @ act) return out @@ -187,7 +191,7 @@ def test_marlin_prefill_matches_dequant_reference(): cache.materialize_layer(0) cache.copy_missing() g1, g2 = cache.alphas_for_layer(0) - gu_p, gu_s, dn_p, dn_s = cache.bank_views(E) + gu_p, gu_s, dn_p, dn_s = cache.stage_views(0, 0) out = marlin_fused_experts( hidden, gu_p, gu_s, g1, dn_p, dn_s, g2, topk_weights, topk_ids, "silu", False, @@ -214,7 +218,7 @@ def decode(layer_id, experts): cache.ensure_experts(layer_id, ids) # rewrites ids -> slots in place cache.copy_missing() g1, g2 = cache.alphas_for_slots(layer_id) - gu_p, gu_s, dn_p, dn_s = cache.bank_views() + gu_p, gu_s, dn_p, dn_s = cache.pool_bank_views(layer_id) out = marlin_fused_experts( hidden, gu_p, gu_s, g1, dn_p, dn_s, g2, topk_weights, ids, "silu", False, @@ -278,7 +282,7 @@ def test_marlin_overlap_prefill_matches_dequant_reference(): assert int(cache.num_indices.item()) == 2, "stale slot mappings survived the prefetch" cache.copy_missing() g1, g2 = cache.alphas_for_slots(0) - gu_p, gu_s, dn_p, dn_s = cache.bank_views() + gu_p, gu_s, dn_p, dn_s = cache.pool_bank_views(0) out = marlin_fused_experts( dec_hidden, gu_p, gu_s, g1, dn_p, dn_s, g2, dec_weights, ids, "silu", False, @@ -325,6 +329,74 @@ def test_triton_overlap_prefill_matches_dequant_reference(): cache.release_prefill_layer(layer_id) +@cuda +def test_triton_prefill_dequant_reads_every_code_and_scale_exactly(): + """The prefill GEMM decodes e2m1 by placing its bits in fp16 rather than through a table: + every code must read 2^-14 of its value, and every product with an e4m3 block scale must be + exact in bf16 -- the property that keeps the GEMM bit-identical to the table decode.""" + import triton + import triton.language as tl + + from freetoken.kernel.triton.e4m3_compat import e4m3_kernel_view, e4m3_to_f16_x128 + from freetoken.kernel.triton.nvfp4_fused_moe import _e2m1_byte_f16_x2pow_neg14 + + @triton.jit + def dequant(bytes_ptr, scales_ptr, placed_ptr, weights_ptr, BLOCK: tl.constexpr): + offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + lo, hi = _e2m1_byte_f16_x2pow_neg14(tl.load(bytes_ptr + offs)) + scale = e4m3_to_f16_x128(tl.load(scales_ptr + offs)) + tl.store(placed_ptr + 2 * offs, lo) + tl.store(placed_ptr + 2 * offs + 1, hi) + tl.store(weights_ptr + 2 * offs, (lo * scale).to(tl.bfloat16)) + tl.store(weights_ptr + 2 * offs + 1, (hi * scale).to(tl.bfloat16)) + + device = torch.device("cuda") + # every byte against every scale: each (low code, high code) pair meets all 256 scale codes + packed = torch.arange(256, dtype=torch.int32, device=device).repeat_interleave(256) + scales = torch.arange(256, dtype=torch.uint8, device=device).repeat(256).view(torch.float8_e4m3fn) + placed = torch.empty(2 * packed.numel(), dtype=torch.float16, device=device) + weights = torch.empty(2 * packed.numel(), dtype=torch.bfloat16, device=device) + dequant[(packed.numel() // 1024,)](packed, e4m3_kernel_view(scales), placed, weights, BLOCK=1024) + + codes = torch.stack([packed & 0xF, packed >> 4], dim=1).reshape(-1) + value = _E2M1.double().to(device)[codes.long()] + assert torch.equal(placed.double(), value * 2.0**-14) + assert torch.equal(placed.signbit(), value.signbit()) + scale = scales.double().repeat_interleave(2) + finite = ~scale.isnan() + assert int(finite.sum()) == 2 * 256 * 254 + assert torch.equal(weights.double()[finite], (value * scale * 2.0**-7)[finite]) + + +@cuda +@pytest.mark.parametrize( + ("hidden", "inter", "tokens"), + [(H, I, 96), (272, 80, 96)], + ids=["tiled-k", "k-not-a-tile-multiple"], +) +def test_triton_prefill_large_batch_matches_dequant_reference(hidden, inter, tokens): + """A batch past the small-batch tile (> 64 tokens) through the Triton prefill grouped GEMM, + at a K the byte tile divides and at one it does not (272 and 80 leave 8 bytes over).""" + from freetoken.moe.fused_nvfp4 import fused_experts_nvfp4 + + device = torch.device("cuda") + sources = _make_native_sources(device, seed=13, hidden=hidden, inter=inter) + torch.manual_seed(14) + x = torch.randn(tokens, hidden, dtype=torch.bfloat16, device=device) / 4 + topk_ids = torch.randint(0, E, (tokens, TOPK), dtype=torch.int32, device=device) + topk_weights = torch.rand(tokens, TOPK, dtype=torch.float32, device=device) + banks = [ + sources[name][0].to(device) + for name in ( + "gate_up_packed", "gate_up_scale", "gate_up_global", + "down_packed", "down_scale", "down_global", + ) + ] + ref = _ref_moe(sources, 0, x, topk_weights, topk_ids) + out = fused_experts_nvfp4(x, *banks, topk_weights, topk_ids, E, "silu", False) + _assert_close(out, ref) + + @cuda def test_triton_demand_prefill_matches_the_streamed_layer_bit_for_bit(): """A short prefill chunk loads only its routed experts, through the slot cache, and runs the @@ -438,7 +510,7 @@ def decode(layer_id, experts): ref = _ref_moe(ref_sources, layer_id, hidden, topk_weights, ids) cache.ensure_experts(layer_id, ids) # rewrites ids -> slots in place cache.copy_missing() - gu_p, gu_s, gu_g, dn_p, dn_s, dn_g = cache.bank_views() + gu_p, gu_s, gu_g, dn_p, dn_s, dn_g = cache.pool_bank_views(layer_id) out = fused_experts_decode_nvfp4_marlin( hidden, gu_p, gu_s, gu_g, dn_p, dn_s, dn_g, topk_weights, ids, "silu", False ) @@ -468,7 +540,7 @@ def test_triton_decode_marlin_matches_baseline_kernel(): ids = torch.tensor([[1, 6]], dtype=torch.int32, device=device) cache.ensure_experts(0, ids) cache.copy_missing() - banks = cache.bank_views() + banks = cache.pool_bank_views(0) marlin = fused_experts_decode_nvfp4_marlin(hidden, *banks, topk_weights, ids, "silu", False) base = fused_experts_decode_nvfp4_serial(hidden, *banks, topk_weights, ids, "silu", False) torch.testing.assert_close(marlin.float(), base.float(), rtol=2e-3, atol=2e-3) @@ -496,7 +568,7 @@ def test_b12x_decode_matches_dequant_reference(): cache.ensure_experts(0, ids) cache.copy_missing() g1, g2 = cache.alphas_for_slots(0) - gu_p, gu_s, dn_p, dn_s = cache.bank_views() + gu_p, gu_s, dn_p, dn_s = cache.pool_bank_views(0) out = b12x_fused_experts( hidden, gu_p, gu_s, g1, dn_p, dn_s, g2, topk_weights, ids, "silu", False ) @@ -567,3 +639,22 @@ def test_b12x_pack_keeps_per_layer_banks_and_flat_alphas(): assert len(cache.bank_sources["gate_up"]) == L assert sum(t.shape[0] for t in cache.bank_sources["gate_up"]) == total assert cache.gate_up_alpha.shape == (total,) + + +@cuda +@pytest.mark.parametrize("num_experts", [1023, 2100]) +def test_moe_align_handles_a_slot_cache_wider_than_1024(num_experts): + """A prefill over the offload slot cache aligns against its slot count (#676): sgl_kernel's + align reads out of bounds past 1024 counters, so a 2100-slot cache must still align.""" + from freetoken.moe.fused import moe_align_block_size + + torch.manual_seed(0) + ids = torch.randint(0, num_experts, (300, 10), device="cuda", dtype=torch.int32) + sorted_ids, expert_ids, padded = moe_align_block_size(ids, 16, num_experts) + n = int(padded) + flat = ids.flatten() + real = sorted_ids[:n][sorted_ids[:n] < flat.numel()] + assert sorted(real.tolist()) == list(range(flat.numel())) + block_of = torch.arange(n, device="cuda") // 16 + rows = sorted_ids[:n] < flat.numel() + assert torch.equal(flat[sorted_ids[:n][rows].long()], expert_ids[block_of[rows]]) diff --git a/tests/moe/test_offload.py b/tests/moe/test_offload.py index 2cbd29eb29..818dd37096 100644 --- a/tests/moe/test_offload.py +++ b/tests/moe/test_offload.py @@ -138,9 +138,10 @@ def test_offload_moe_layer_short_prefill_loads_its_experts_through_the_slot_cach ) monkeypatch.setattr(cache, "materialize_layer", lambda layer_id: calls.setdefault("materialized", layer_id)) - def ensure(layer_id, expert_ids, *, counts_as_decode=True): + def ensure(layer_id, expert_ids, *, counts_as_decode=True, **kwargs): + # the layer hands the unique routed ids (decode-width chunks); rewrite each to a slot calls["ensured"] = (layer_id, expert_ids.tolist(), counts_as_decode) - expert_ids.copy_(torch.tensor([[5, 3]], dtype=torch.int32)) + expert_ids.add_(3) monkeypatch.setattr(cache, "ensure_experts", ensure) monkeypatch.setattr(cache, "copy_missing", lambda: calls.setdefault("copied", True)) @@ -157,11 +158,11 @@ def fake_fused(hidden_states, w1, w2, got_topk_weights, got_topk_ids, activation assert out is hidden_states assert "materialized" not in calls - assert calls["ensured"] == (0, [[2, 1]], False) + assert calls["ensured"] == (0, [1, 2], False) assert calls["copied"] is True # the grouped GEMM reads the whole slot cache at the slots the admission wrote assert calls["w1"].shape[0] == cache.cache_size - assert calls["topk_ids"].tolist() == [[5, 3]] + assert calls["topk_ids"].tolist() == [[5, 4]] calls.clear() layer.prefill_forward(torch.randn(2, 8), torch.randn(2, 4)) @@ -252,7 +253,7 @@ def fake_fused( assert fused_calls[0]["w2_ptr"] == fused_calls[2]["w2_ptr"] assert fused_calls[0]["w1_ptr"] != fused_calls[1]["w1_ptr"] assert fused_calls[0]["w2_ptr"] != fused_calls[1]["w2_ptr"] - prefill_gate_up_buffer, prefill_down_buffer = cache.prefill_bank_buffers + prefill_gate_up_buffer, prefill_down_buffer = cache.stage_views(0, 0) assert prefill_gate_up_buffer.data_ptr() == cache.bank_caches["gate_up"].data_ptr() assert prefill_down_buffer.data_ptr() == cache.bank_caches["down"].data_ptr() @@ -336,6 +337,50 @@ def test_prefill_overlap_prefetch_invalidates_borrowed_unified_cache_slots(): assert torch.equal(cache.bank_caches["down"][:num_experts], down_source[0]) +def test_invalidate_prefill_slots_kernel_matches_eager_reference(): + from freetoken.kernel.triton.moe import invalidate_prefill_slots + + if not torch.cuda.is_available(): + pytest.skip("CUDA is required for the invalidation kernel") + + for num_layers, num_experts, cache_size, slot_start in [ + (3, 4, 8, 0), + (3, 4, 8, 4), + (2, 256, 1024, 512), + ]: + for trial in range(20): + torch.manual_seed(1234 + trial) + ids = torch.full((cache_size,), -1, dtype=torch.int32) + k = torch.randint(0, num_experts + 1, (1,)).item() + slots = torch.randperm(cache_size)[:k] + ids[slots] = torch.randint(0, num_layers * num_experts, (k,), dtype=torch.int32) + slot_for_id = torch.full((num_layers, num_experts), -1, dtype=torch.int32) + for slot, expert in enumerate(ids.tolist()): + if expert >= 0: + slot_for_id.view(-1)[expert] = slot + usage = torch.randint(0, 999, (cache_size,), dtype=torch.int64) + + ids_c, sfi_c, usage_c = ids.cuda(), slot_for_id.cuda(), usage.cuda() + + ref_sfi = slot_for_id.clone() + ref_usage = usage.clone() + ref_ids = ids.clone() + buf = ref_ids[slot_start : slot_start + num_experts] + ref_sfi.view(-1)[buf[buf >= 0].long()] = -1 + buf.fill_(-1) + ref_usage[slot_start : slot_start + num_experts].zero_() + + invalidate_prefill_slots(ids_c, sfi_c, usage_c, slot_start, num_experts) + torch.cuda.synchronize() + + assert torch.equal(ref_sfi, sfi_c.cpu()) + assert torch.equal(ref_ids, ids_c.cpu()) + assert torch.equal(ref_usage, usage_c.cpu()) + + with pytest.raises(ValueError): + invalidate_prefill_slots(ids.new_full((1,), -1), slot_for_id, usage.new_zeros(1), 0, 2) + + def test_prefill_overlap_waits_for_previous_prefill_release_after_begin(monkeypatch): from freetoken.moe.offload_cache import OffloadMoeCache @@ -402,7 +447,7 @@ def test_offload_moe_layer_decode_forward_uses_remapped_slot_ids(monkeypatch): lambda *, hidden_states, gating_output, topk, renormalize: (topk_weights, topk_ids), ) - def fake_ensure(layer_id, expert_ids): + def fake_ensure(layer_id, expert_ids, **kw): calls["ensure_layer_id"] = layer_id calls["ensure_expert_ids"] = expert_ids.clone() expert_ids.copy_(torch.tensor([[5, 0]], dtype=torch.int32)) @@ -442,6 +487,121 @@ def fake_fused_decode( assert calls["topk_ids"].tolist() == [[5, 0]] +def test_decode_copy_overlap_follows_the_flag_and_the_decode_target(monkeypatch): + from freetoken.env import ENV + from freetoken.moe.offload_cache import OffloadMoeCache + + def cache(device, **kw): + return OffloadMoeCache(num_layers=3, num_experts=4, cache_size=8, device=torch.device(device), **kw) + + monkeypatch.setattr(ENV.MOE_COPY_OVERLAP, "value", True) + assert cache("cpu").decode_copy_overlap is True # unset -> FREETOKEN_MOE_COPY_OVERLAP + assert cache("cpu").decode_copy_stream is None + monkeypatch.setattr(ENV.MOE_COPY_OVERLAP, "value", False) + assert cache("cpu").decode_copy_overlap is False + if not torch.cuda.is_available(): + return + assert cache("cuda").decode_copy_stream is None + for target in ("cpu", "hybrid"): + assert cache("cuda", decode_copy_overlap=True, decode_target=target).decode_copy_stream is None + on = cache("cuda", decode_copy_overlap=True) + assert on.decode_copy_stream is not None and len(on._decode_copy_events) == 3 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +def test_copy_missing_takes_an_explicit_layer_id(): + # a second copy path (the overlap fork, later prefetch) names its source layer instead of trusting whichever call staged last + from freetoken.moe.offload_cache import OffloadMoeCache + + num_layers, num_experts = 2, 4 + fingerprint = torch.arange(num_layers * num_experts, dtype=torch.float32).view(num_layers, num_experts, 1, 1) + sources = { + "gate_up": [fingerprint[l].expand(num_experts, 32, 8).contiguous().pin_memory() for l in range(num_layers)], + "down": [fingerprint[l].expand(num_experts, 8, 16).contiguous().pin_memory() for l in range(num_layers)], + } + cache = OffloadMoeCache( + num_layers=num_layers, num_experts=num_experts, cache_size=num_experts, device=torch.device("cuda") + ) + cache.set_bank_sources(sources) + cache.reset() + ids = torch.tensor([2, 3], dtype=torch.int32, device="cuda") + cache.ensure_experts(1, ids) + cache._pending_src_layer = 0 + cache.copy_missing(1) + torch.cuda.synchronize() + for bank in cache.bank_views(): + assert [int(bank[s].flatten()[0]) for s in ids.tolist()] == [num_experts + 2, num_experts + 3] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +def test_decode_side_applies_only_to_gpu_decode(monkeypatch): + # prefill, CPU-routed layers, TP > 1 and the flag off keep the shared expert on the compute stream + from types import SimpleNamespace + + import freetoken.core as core + from freetoken.core import Context + from freetoken.moe.offload_cache import OffloadMoeCache + + _init_tp() + layer = _bf16_offload_layer(0, 4, 2, 8, 16) + ctx = Context(page_size=1) + monkeypatch.setattr(core, "_GLOBAL_CTX", ctx) + + def applies(is_prefill=False, cpu_layers=(), **cache_kw): + layer.offload_cache = OffloadMoeCache( + num_layers=1, num_experts=4, cache_size=4, device=torch.device("cuda"), **cache_kw + ) + layer.offload_cache.cpu_layer_ids = frozenset(cpu_layers) + ctx._batch = SimpleNamespace(is_prefill=is_prefill) + return layer.decode_side_applies() + + assert not applies(decode_copy_overlap=False) + assert not applies(decode_copy_overlap=True, decode_target="hybrid") + assert not applies(is_prefill=True, decode_copy_overlap=True) + assert not applies(cpu_layers=(0,), decode_copy_overlap=True) + monkeypatch.setattr(layer, "tp_size", 2) + assert not applies(decode_copy_overlap=True) + monkeypatch.setattr(layer, "tp_size", 1) + assert applies(decode_copy_overlap=True) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +def test_decode_side_stream_never_aliases_a_pooled_stream(): + # torch.cuda.Stream() hands out 32 pooled streams per priority round robin; a pooled side stream + # can come back as the engine or capture stream, and then the side work silently runs serially + from freetoken.moe.offload_cache import OffloadMoeCache + + cache = OffloadMoeCache( + num_layers=1, num_experts=4, cache_size=4, device=torch.device("cuda"), decode_copy_overlap=True + ) + side = cache.decode_copy_stream.cuda_stream + pooled = {torch.cuda.Stream().cuda_stream for _ in range(64)} + assert side not in pooled and side != torch.cuda.current_stream().cuda_stream + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +def test_decode_side_survives_a_rebuild(): + # the side stream and events are not cache_size-shaped, so a runtime rebuild keeps them working + from freetoken.moe.offload_cache import OffloadMoeCache + + cache = OffloadMoeCache( + num_layers=2, num_experts=4, cache_size=4, device=torch.device("cuda"), decode_copy_overlap=True + ) + cache.set_bank_sources({ + "gate_up": [torch.zeros(4, 32, 8, dtype=torch.bfloat16).pin_memory() for _ in range(2)], + "down": [torch.zeros(4, 8, 16, dtype=torch.bfloat16).pin_memory() for _ in range(2)], + }) + stream = cache.decode_copy_stream + cache.rebuild(6) + assert cache.decode_copy_stream is stream is not None + x = torch.arange(1024, dtype=torch.float32, device="cuda") + for layer_id in range(2): + with cache.decode_side(layer_id): + assert torch.cuda.current_stream() == stream + y = x * 2 + cache.join_decode_side(layer_id) + assert torch.equal(y + 0, x * 2) + def test_lru_gpu_cache_assigns_unique_slots_for_large_miss_batch(): import pytest @@ -729,6 +889,136 @@ def test_offload_cache_rebuild_resizes_and_preserves_sources(): assert torch.all(cache.id_of_slot == -1) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +def test_slot_pools_of_two_bank_layouts_keep_their_rows_through_prefill(): + from freetoken.moe.offload_cache import OffloadMoeCache + + _init_tp() + torch.manual_seed(0) + L, E, top_k = 4, 8, 4 + # layers 1 and 3 keep wider down rows (uint8, like Q6_K next to Q4_K); their pools come first + wide = (1, 3) + sources = { + "gate_up": [torch.randn(E, 1024, 256, dtype=torch.bfloat16).pin_memory() for _ in range(L)], + "down": [ + torch.randint(0, 255, (E, 512, 768), dtype=torch.uint8).pin_memory() if layer in wide + else torch.randn(E, 512, 256, dtype=torch.bfloat16).pin_memory() + for layer in range(L) + ], + } + cache = OffloadMoeCache( + num_layers=L, num_experts=E, cache_size=33, device=torch.device("cuda"), + prefill_overlap=True, prefill_hit_d2d=True, pool_of_layer=[2, 0, 3, 1], pool_sizes=[6, 5, 12, 10], + ) + cache.set_bank_sources(sources) + assert [v[1].dtype for v in map(cache.pool_bank_views, range(L))] == [torch.bfloat16, torch.uint8] * 2 + # staging buffers take 2 x 8 rows of every bank: all of the wide pools, the first 8 slots of layer 0's pool + assert cache._pool_stable_from == [6, 5, 8, 0] + + def check_resident() -> None: + torch.cuda.synchronize() + for layer in range(L): + for e, s in enumerate(cache.slot_for_id[layer].tolist()): + if s >= 0: + for (src, _), view in zip(cache.banks, cache.pool_bank_views(layer)): + assert torch.equal(view[s].cpu(), src[layer][e]), (layer, e, s) + + def decode(steps: int) -> None: + for _ in range(steps): + for layer in range(L): + cache.ensure_experts(layer, torch.randperm(E)[:top_k].to(torch.int32).view(1, top_k).cuda()) + cache.copy_missing() + check_resident() + + decode(20) + cache.zero_inactive_route_slots = True + cache.begin_prefill() + for layer in range(L): + cache.prefetch_prefill_layer(layer) + cache.prefetch_prefill_layer(layer + 1) + views = cache.wait_prefill_layer(layer) + torch.cuda.synchronize() + for (src, _), view in zip(cache.banks, views): + assert torch.equal(view.cpu(), src[layer]) + cache.release_prefill_layer(layer) + assert cache.prefill_hit_rows > 0 + for p, stable in enumerate(cache._pool_stable_from): + if stable: + assert not any(view[0].any() for view in cache.pool_bank_views(cache.pool_of_layer.index(p))) + check_resident() + decode(20) + + cache.materialize_layer(3) + cache.copy_missing() + for (src, _), view in zip(cache.banks, cache.stage_views(3, 0)): + assert torch.equal(view.cpu(), src[3]) + check_resident() + + cache.rebuild(32, [5, 5, 10, 12]) + decode(20) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +def test_slot_pools_split_layers_of_one_bank_layout_by_quant_tag(): + from freetoken.moe.offload_cache import OffloadMoeCache, bank_layouts + from freetoken.moe.pool_policy import group_layers + + _init_tp() + L, E = 4, 8 + # uint8 rows of one width, as GGUF Q4_K and IQ4_NL rows are + sources = { + "gate_up": [torch.randint(1, 255, (E, 64, 144), dtype=torch.uint8).pin_memory() for _ in range(L)], + "down": [torch.randint(1, 255, (E, 32, 144), dtype=torch.uint8).pin_memory() for _ in range(L)], + } + tags = [("ggml", "Q4_K", "Q4_K"), ("ggml", "IQ4_NL", "IQ4_NL")] * 2 + assert group_layers("unified", bank_layouts(sources, tags), [1] * L) == [[0, 2], [1, 3]] + + def make(pool_of_layer, pool_sizes): + return OffloadMoeCache( + num_layers=L, num_experts=E, cache_size=sum(pool_sizes), device=torch.device("cuda"), + prefill_overlap=True, pool_of_layer=pool_of_layer, pool_sizes=pool_sizes, + ) + + with pytest.raises(AssertionError, match="layouts"): + make([0] * L, [20]).set_bank_sources(sources, layer_tags=tags) + cache = make([0, 1, 0, 1], [10, 10]) + cache.set_bank_sources(sources, layer_tags=tags) + assert cache.bank_caches["down"].shape == (20, 32, 144) + # the staging buffers cover pool 0 and slots [0, 6) of pool 1, so both slot 0 rows end up with another quant's bytes + cache.zero_inactive_route_slots = True + cache.begin_prefill() + for layer in range(L): + cache.prefetch_prefill_layer(layer) + cache.prefetch_prefill_layer(layer + 1) + cache.wait_prefill_layer(layer) + cache.release_prefill_layer(layer) + torch.cuda.synchronize() + for layer in (0, 1): + for view in cache.pool_bank_views(layer): + assert not view[0].any() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +def test_offload_cache_rebuild_frees_the_old_slots_first(): + from freetoken.moe.offload_cache import OffloadMoeCache + + _init_tp() + cache = OffloadMoeCache( + num_layers=2, num_experts=4, cache_size=32, device=torch.device("cuda"), + pool_of_layer=[0, 1], pool_sizes=[16, 16], + ) + cache.set_bank_sources({ + "gate_up": [torch.zeros(4, 1024, 1024, dtype=torch.bfloat16).pin_memory() for _ in range(2)], + "down": [torch.zeros(4, 512, 1024, dtype=torch.bfloat16).pin_memory() for _ in range(2)], + }) + torch.cuda.synchronize() + before = torch.cuda.memory_allocated() + torch.cuda.reset_peak_memory_stats() + cache.rebuild(32, [8, 24]) + # 96 MiB of slots: holding the old ones while allocating the new ones would double it + assert torch.cuda.max_memory_allocated() - before < 8 << 20 + + def test_offload_cache_rebuild_disables_prefill_overlap_when_too_small(): from freetoken.moe.offload_cache import OffloadMoeCache @@ -744,7 +1034,7 @@ def test_offload_cache_rebuild_disables_prefill_overlap_when_too_small(): assert cache.cache_size == 5 assert cache.prefill_overlap is False - assert cache.prefill_bank_buffers == [] + assert cache._stage_views[0][1] is None # no room for the second staging buffer def test_offload_cache_rebuild_keeps_overlap_at_boundary(): @@ -780,6 +1070,58 @@ def test_offload_cache_validate_rebuild_enforces_marlin_cap_and_floor(): bf16.validate_rebuild(3) # below the num_experts floor +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +@pytest.mark.parametrize("target", ["gpu", "hybrid"]) +def test_slot_pools_keep_each_layer_in_its_own_slots(target): + from freetoken.moe.offload_cache import OffloadMoeCache + + _init_tp() + torch.manual_seed(0) + L, E, sizes = 3, 16, [5, 7, 20] + # hybrid routes more experts than the smallest pool holds; GPU decode must fit them + top_k = 2 if target == "gpu" else 8 + sources = { + "gate_up": [torch.randn(E, 64, 32, dtype=torch.bfloat16).pin_memory() for _ in range(L)], + "down": [torch.randn(E, 32, 32, dtype=torch.bfloat16).pin_memory() for _ in range(L)], + } + + cache = OffloadMoeCache( + num_layers=L, num_experts=E, cache_size=sum(sizes), device=torch.device("cuda"), + decode_target=target, hybrid_max_fetch=E, pool_of_layer=[0, 1, 2], pool_sizes=sizes, + ) + cache.set_bank_sources(sources) + for _ in range(50): + for layer in range(L): + experts = torch.randperm(E)[:top_k].to(torch.int32) + slots = experts.cuda() + if target == "gpu": + cache.ensure_experts(layer, slots) + else: + cache.ensure_experts_hybrid(layer, slots, torch.rand(top_k, device="cuda")) + cache.copy_missing() + resident = [(e, s) for e, s in zip(experts.tolist(), slots.tolist()) if s >= 0] + assert len({s for _, s in resident}) == len(resident) + for e, s in resident: + assert s < sizes[layer] + for (src, _), view in zip(cache.banks, cache.pool_bank_views(layer)): + assert torch.equal(view[s].cpu(), src[layer][e]) + + # with several pools a whole-layer prefill stages slots [0, E) without registering them in any pool + cache.materialize_layer(1) + cache.copy_missing() + for (src, _), view in zip(cache.banks, cache.stage_views(1, 0)): + assert torch.equal(view.cpu(), src[1]) + slot_for_id, id_of_slot = cache.slot_for_id.cpu(), cache.id_of_slot.cpu() + for layer in range(L): + for e, s in enumerate(slot_for_id[layer].tolist()): + if s >= 0: + assert cache.pool_bases[layer] + s >= E + assert id_of_slot[cache.pool_bases[layer] + s] == layer * E + e + + cache.rebuild(40, [8, 8, 24]) + assert [v[0].shape[0] for v in map(cache.pool_bank_views, range(L))] == [8, 8, 24] + + def _make_split_cache(num_layers=2, locked=(1,), prefill_overlap=False, device="cpu"): """A [gate_up, down] bf16 cache with the given layers LOCKED (rest pinned).""" from freetoken.moe.host_banks import HostResidency @@ -971,3 +1313,26 @@ def boom(addr, nbytes): with hb.PinPipeline() as pins: pins(1, {"gate_up": hb.HostBank((4,), torch.uint8)}) assert plan2.actual == {1: hb.HostResidency.PAGEABLE.value} + + +@pytest.mark.parametrize( + "kwargs, flag", + [ + ({"cache_policy": "rule"}, "--moe-cache-policy rule"), + ({"prefetch_mode": "on"}, "FREETOKEN_MOE_PREFETCH"), + ({"verify_mode": "meta"}, "FREETOKEN_MOE_PREFETCH_VERIFY"), + ({"slot_audit": 8}, "FREETOKEN_MOE_SLOT_AUDIT"), + ({"prefill_demand_tokens": 64}, "--moe-prefill-demand-tokens"), + ], +) +def test_offload_cache_refuses_global_slot_features_with_several_pools(kwargs, flag): + # these address slots globally over one typed bank: with pool-local slot ids they would + # install and read experts in another pool's rows + from freetoken.moe.offload_cache import OffloadMoeCache + + _init_tp() + common = dict(num_layers=2, num_experts=4, cache_size=16, device=torch.device("cpu"), **kwargs) + with pytest.raises(ValueError, match=f"{flag}.*one MoE slot pool"): + OffloadMoeCache(pool_of_layer=[0, 1], pool_sizes=[8, 8], **common) + OffloadMoeCache(**common) + OffloadMoeCache(num_layers=2, num_experts=4, cache_size=16, device=torch.device("cpu"), pool_of_layer=[0, 1], pool_sizes=[8, 8]) diff --git a/tests/moe/test_prefetch.py b/tests/moe/test_prefetch.py new file mode 100644 index 0000000000..47c26066a4 --- /dev/null +++ b/tests/moe/test_prefetch.py @@ -0,0 +1,424 @@ +"""FREETOKEN_MOE_PREFETCH stage A: the router-lookahead select / count kernels against a Python +reference, the mode switch, and the predictor stream's lifetime across resets and rebuilds. +The qwen4_exp stack tests (bit-identical outputs, counters over a real decode) live in +tests/models/qwen4_exp/test_skeleton.py.""" + +from __future__ import annotations + +import weakref + +import pytest +import torch + +from .ref_prefetch import ref_count, ref_select + +requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") + + +def _padded(ids, width): + return ids + [-1] * (width - len(ids)) + + +def _select(logits, resident, k, budget, width=16): + from freetoken.moe.prefetch import lookahead_select + + sel = torch.full((width,), -7, dtype=torch.int32, device="cuda") + res = torch.full((width,), -7, dtype=torch.int32, device="cuda") + lookahead_select( + torch.tensor(logits, dtype=torch.bfloat16, device="cuda"), + torch.tensor(resident, dtype=torch.int32, device="cuda"), + sel, res, k=k, budget=budget, + ) + return sel.tolist(), res.tolist() + + +def test_mode_switch(): + from freetoken.moe.prefetch import default_budget, resolve_mode + + assert resolve_mode("off") == "off" and resolve_mode(" Measure ") == "measure" and resolve_mode("ON") == "on" + with pytest.raises(ValueError): + resolve_mode("lookahead") + assert (default_budget(True), default_budget(False)) == (3, 4) + + +@requires_cuda +def test_cache_builds_the_prefetcher_only_when_asked(): + from freetoken.moe.offload_cache import OffloadMoeCache + + def cache(**kw): + return OffloadMoeCache(num_layers=2, num_experts=4, cache_size=4, device=torch.device("cuda"), **kw) + + assert cache(prefetch_mode="off").prefetch is None + assert cache(prefetch_mode="measure").prefetch.copy_stream is None + on = cache(prefetch_mode="on") + assert on.prefetch.mode == "on" and on.prefetch.cache.prefetch is on.prefetch and on.prefetch_on + # no reference cycle: dropping the cache frees its tensors at once + ref = weakref.ref(on) + del on + assert ref() is None + for mode in ("measure", "on"): + with pytest.raises(ValueError): + cache(prefetch_mode=mode, decode_target="hybrid") + + +@requires_cuda +def test_the_prefetch_env_flag_skips_caches_that_do_not_decode_on_the_gpu(monkeypatch): + # a process-wide flag must not break hybrid or CPU-decoded caches, as FREETOKEN_MOE_COPY_OVERLAP does not + from freetoken.env import ENV + from freetoken.moe.offload_cache import OffloadMoeCache + + monkeypatch.setattr(ENV.MOE_PREFETCH, "value", "on") + for target in ("hybrid", "cpu"): + c = OffloadMoeCache(num_layers=2, num_experts=4, cache_size=4, device=torch.device("cuda"), decode_target=target) + assert c.prefetch is None and not c.prefetch_on + assert OffloadMoeCache(num_layers=2, num_experts=4, cache_size=4, device=torch.device("cuda")).prefetch_on + + +@requires_cuda +def test_select_known_answers(): + # order by logit: 1, 4, 6, 2, 3, 7, 0, 5; experts 4 and 2 are resident + logits = [[0.1, 0.9, 0.5, 0.3, 0.8, 0.0, 0.7, 0.2]] + resident = [-1, -1, 7, -1, 3, -1, -1, -1] + assert _select(logits, resident, 5, 2) == (_padded([1, 6], 16), _padded([4], 16)) + assert _select(logits, resident, 5, 4) == (_padded([1, 6, 3], 16), _padded([4, 2], 16)) + # the top-k cap: 3 and 7 are out of a top-4 list + assert _select(logits, resident, 4, 4) == (_padded([1, 6], 16), _padded([4, 2], 16)) + # all-equal logits rank by id + assert _select([[0.5] * 8], [-1] * 8, 8, 3) == (_padded([0, 1, 2], 16), _padded([], 16)) + # bs 2: row 0 ranks 1, 4, 6 and row 1 ranks 4, 0, 5 -> merged 1, 4, 0, 6, 5; expert 0 resident + rows = [[0.0, 0.9, 0.1, 0.1, 0.8, 0.1, 0.7, 0.1], [0.8, 0.1, 0.1, 0.1, 0.9, 0.7, 0.1, 0.1]] + resident = [2, -1, -1, -1, -1, -1, -1, -1] + assert _select(rows, resident, 3, 3) == (_padded([1, 4, 6], 16), _padded([0], 16)) + assert _select(rows, resident, 3, 16) == (_padded([1, 4, 6, 5], 16), _padded([0], 16)) + + +@requires_cuda +@pytest.mark.parametrize("rows", [1, 2]) +@pytest.mark.parametrize("k,budget", [(3, 3), (16, 4), (16, 3), (8, 16)]) +def test_select_matches_reference(rows, k, budget): + gen = torch.Generator().manual_seed(rows * 100 + k + budget) + for trial in range(20): + # coarse values make bf16 ties common, so the tie rule is exercised + logits = (torch.randint(-24, 24, (rows, 512), generator=gen).float() / 8).tolist() + resident = torch.where(torch.rand(512, generator=gen) < 0.4, 1, -1).tolist() + want = ref_select(logits, resident, k, budget) + width = 2 * k + got = _select(logits, resident, k, budget, width=width) + assert got == (_padded(want[0], width), _padded(want[1], width)), trial + + +@requires_cuda +def test_count_known_answers(): + from freetoken.moe.prefetch import prefetch_count + + num_experts, layer = 8, 2 + base = layer * num_experts + # slots 0, 1, 3 hold this layer's experts 3, 1, 6; slot 4 another layer's expert + id_of_slot = torch.tensor([base + 3, base + 1, -1, base + 6, 5, -1], dtype=torch.int32, device="cuda") + slots = torch.tensor([0, 1, 3], dtype=torch.int32, device="cuda") + sel = torch.tensor(_padded([3, 5], 6), dtype=torch.int32, device="cuda") + res = torch.tensor(_padded([1, 4], 6), dtype=torch.int32, device="cuda") + misses = torch.tensor([2], dtype=torch.int64, device="cuda") + counters = torch.zeros(6, dtype=torch.int64, device="cuda") + for _ in range(2): + prefetch_count(slots, id_of_slot, misses, sel, res, counters, id_base=base, rows=1) + want = ref_count([3, 5], [1, 4], [3, 1, 6], 2, 1) + assert counters.tolist() == [2 * v for v in want] == [4, 2, 2, 2, 2, 4] + # bs 2 with a duplicate routed expert counts it once + counters.zero_() + slots2 = torch.tensor([0, 1, 0, 3], dtype=torch.int32, device="cuda") + prefetch_count(slots2, id_of_slot, misses, sel, res, counters, id_base=base, rows=2) + assert counters.tolist() == ref_count([3, 5], [1, 4], [3, 1, 3, 6], 2, 2) + + +@requires_cuda +def test_prefetch_stream_never_aliases_a_pooled_stream(): + from freetoken.moe.offload_cache import OffloadMoeCache + + cache = OffloadMoeCache( + num_layers=2, num_experts=4, cache_size=4, device=torch.device("cuda"), + decode_copy_overlap=True, prefetch_mode="measure", + ) + pred = cache.prefetch.stream.cuda_stream + pooled = {torch.cuda.Stream().cuda_stream for _ in range(64)} + assert pred not in pooled and pred != torch.cuda.current_stream().cuda_stream + assert pred != cache.decode_copy_stream.cuda_stream + + +@requires_cuda +def test_prefetch_copy_stream_is_a_third_dedicated_stream(): + from freetoken.moe.offload_cache import OffloadMoeCache + + cache = OffloadMoeCache( + num_layers=2, num_experts=4, cache_size=4, device=torch.device("cuda"), + decode_copy_overlap=True, prefetch_mode="on", + ) + pf = cache.prefetch + copy = pf.copy_stream.cuda_stream + pooled = {torch.cuda.Stream().cuda_stream for _ in range(64)} + assert copy not in pooled and pf.stream.cuda_stream not in pooled + assert len({copy, pf.stream.cuda_stream, cache.decode_copy_stream.cuda_stream, torch.cuda.current_stream().cuda_stream}) == 4 + # its CTAs go ahead of pending compute ones, so a late copy does not also wait for SMs + from cuda.bindings import runtime as cudart + + _, _, greatest = cudart.cudaDeviceGetStreamPriorityRange() + assert cudart.cudaStreamGetPriority(copy)[1] == greatest < 0 + assert cudart.cudaStreamGetPriority(pf.stream.cuda_stream)[1] == 0 + + +@requires_cuda +def test_rebuild_and_reset_clear_the_counters(): + from freetoken.moe.offload_cache import OffloadMoeCache + + cache = OffloadMoeCache( + num_layers=2, num_experts=4, cache_size=4, device=torch.device("cuda"), prefetch_mode="measure" + ) + cache.set_bank_sources({ + "gate_up": [torch.zeros(4, 32, 8, dtype=torch.bfloat16).pin_memory() for _ in range(2)], + "down": [torch.zeros(4, 8, 16, dtype=torch.bfloat16).pin_memory() for _ in range(2)], + }) + pf = cache.prefetch + stream, sel = pf.stream, pf.sel + pf.counters.fill_(3) + cache.reset() + assert int(pf.counters.abs().sum()) == 0 + pf.counters.fill_(5) + pf.take_window() + assert int(pf.totals.sum()) > 0 + pf.counters.fill_(2) + pf._inflight[1] = torch.zeros(1) + cache.rebuild(6) + assert cache.prefetch is pf and pf.stream is stream and pf.sel is sel + assert int(pf.counters.abs().sum()) == 0 and int(pf.totals.abs().sum()) == 0 + assert pf._inflight == [None, None] + + +def _emit_prefetch(window, monkeypatch, *, debug=False, mode="measure"): + from types import SimpleNamespace + + from freetoken.engine import engine as engine_mod + from freetoken.env import ENV + + taken, lines = [], [] + prefetch = SimpleNamespace(mode=mode, take_window=lambda: taken.append(1) or window) + engine = SimpleNamespace(moe_offload_cache=SimpleNamespace(prefetch=prefetch)) + monkeypatch.setattr(ENV.MOE_PREFETCH_DEBUG, "value", debug) + # the freetoken loggers do not propagate, so caplog only sees them in some import orders + monkeypatch.setattr(engine_mod.logger, "info_rank0", lambda msg, *a, **k: lines.append(msg)) + engine_mod.Engine._emit_prefetch_stats(engine) + return "\n".join(lines), taken + + +def test_engine_reports_one_prefetch_window(monkeypatch): + from freetoken.engine.engine import MOE_STATS_INTERVAL + + window = torch.tensor([[0] * 6, [30, 12, 4, 10, 10, 40], [40, 20, 6, 10, 10, 30]]) + out, taken = _emit_prefetch(window, monkeypatch) + assert taken == [1] + assert f"MoE prefetch measure ({MOE_STATS_INTERVAL} decode steps): issued/layer=3.50, useful/layer=1.60" in out + assert "per layer" not in out + out, _ = _emit_prefetch(window, monkeypatch, debug=True) + assert "L1=1.20/3.00/4.00, L2=2.00/4.00/3.00" in out and "L0" not in out + assert _emit_prefetch(torch.zeros((3, 6), dtype=torch.int64), monkeypatch)[0] == "" + + +def test_engine_reports_the_on_mode_window(monkeypatch): + from freetoken.engine.engine import MOE_STATS_INTERVAL + + # issued, useful, resident_hits, calls, rows, misses, late, copied + window = torch.tensor([[0] * 8, [30, 12, 0, 10, 10, 18, 2, 10], [40, 20, 0, 10, 10, 20, 0, 6]]) + out, _ = _emit_prefetch(window, monkeypatch, mode="on") + assert f"MoE prefetch on ({MOE_STATS_INTERVAL} decode steps): issued/layer=3.50, useful/layer=1.60" in out + assert "misses/layer=1.90, coverage=0.457, late=0.125 of 16 copies" in out and "resident_hits" not in out + + +def test_summary_rates(): + from freetoken.moe.prefetch import format_per_layer, format_summary, summarize + + # layer 0 never predicts; layers 1 and 2 over 10 bs-1 steps + counters = torch.tensor([[0] * 6, [30, 12, 4, 10, 10, 40], [40, 20, 6, 10, 10, 30]]) + stats = summarize(counters) + assert stats["layer_calls"] == 20 and stats["tokens"] == 10 + assert stats["issued_per_layer"] == 3.5 and stats["useful_per_layer"] == 1.6 + assert stats["precision"] == 32 / 70 and stats["coverage"] == 32 / 70 + assert stats["useful_per_token"] == 3.2 and stats["resident_hits_per_layer"] == 0.5 + assert "useful/token=3.2" in format_summary(stats) + assert format_per_layer(counters) == "L1=1.20/3.00/4.00, L2=2.00/4.00/3.00" + + +# FREETOKEN_MOE_PREFETCH_VERIFY (moe/verify.py): the mode switch, where the verifier is built, and +# the engine's report; the checks themselves run on the qwen4_exp stacks in test_skeleton.py + + +def test_verify_mode_switch(): + from freetoken.moe.verify import resolve_mode + + assert [resolve_mode(m) for m in ("0", "", "off", "1", " True ", "full", "META")] == [ + "off", "off", "off", "full", "full", "full", "meta"] + with pytest.raises(ValueError): + resolve_mode("bytes") + + +@requires_cuda +def test_cache_builds_the_verifier_only_for_gpu_decode(monkeypatch): + from freetoken.env import ENV + from freetoken.moe.offload_cache import OffloadMoeCache + + def cache(**kw): + return OffloadMoeCache(num_layers=2, num_experts=4, cache_size=4, device=torch.device("cuda"), **kw) + + monkeypatch.setattr(ENV.MOE_PREFETCH_VERIFY, "value", "0") + assert cache().verify is None and cache(verify_mode="off").verify is None + full = cache(verify_mode="full", prefetch_mode="on") + assert full.verify.mode == "full" and cache(verify_mode="meta").verify.mode == "meta" + # the verifier holds no reference to its cache + ref = weakref.ref(full) + del full + assert ref() is None + with pytest.raises(ValueError): + cache(verify_mode="full", decode_target="hybrid") + monkeypatch.setattr(ENV.MOE_PREFETCH_VERIFY, "value", "1") + for target in ("hybrid", "cpu"): + assert cache(decode_target=target).verify is None + assert cache().verify.mode == "full" + + +@requires_cuda +def test_engine_reports_the_verify_window_and_session(monkeypatch): + from types import SimpleNamespace + + from freetoken.engine import engine as engine_mod + from freetoken.engine.engine import MOE_STATS_INTERVAL + from freetoken.moe.offload_cache import OffloadMoeCache + from freetoken.moe.verify import BYTES_PRE, CHECKS, NUM_FIELDS, PRE_BAD + + cache = OffloadMoeCache(num_layers=3, num_experts=4, cache_size=4, device=torch.device("cuda"), verify_mode="full") + cache.set_bank_sources({ + "gate_up": [torch.randn(4, 32, 8, dtype=torch.bfloat16).pin_memory() for _ in range(3)], + "down": [torch.randn(4, 8, 16, dtype=torch.bfloat16).pin_memory() for _ in range(3)], + }) + v = cache.verify + lines = [] + for level in ("info_rank0", "warning_rank0"): + monkeypatch.setattr(engine_mod.logger, level, lambda msg, *a, level=level, **k: lines.append((level, msg))) + engine = SimpleNamespace(moe_offload_cache=cache) + + v.counters[:, CHECKS] = 256 + engine_mod.Engine._emit_verify_stats(engine) + assert lines == [("info_rank0", f"MoE verify full ({MOE_STATS_INTERVAL} decode steps): 768 layer calls checked, " + "meta_bad=0, pre_bad=0, post_bad=0, post_meta_bad=0; session bad=0")] + lines.clear() + v.counters[:, CHECKS] = 256 + v.counters[1, PRE_BAD] = 1 + # call 7, layer 1, row 3 of a top-2 call, expert 2, slot 1, bank down, before the GEMM; slot 1 holds + # layer 1's expert 2, usage 9 at lru step 9, demand-copied; bytes 64.. differ in 4 words + record = [7, 1, 3, 2, 1, 1, BYTES_PRE, 6, 1, 9, 9, 1, 64, 4] + assert len(record) == NUM_FIELDS + v.ring[0] = torch.tensor(record) + v.cursor.fill_(1) + v._top_k = 2 + engine_mod.Engine._emit_verify_stats(engine) + assert [level for level, _ in lines] == ["warning_rank0", "warning_rank0"] + assert "pre_bad=1" in lines[0][1] and "(meta/pre/post/post_meta by layer: L1=0/1/0/0); session bad=1" in lines[0][1] + assert lines[1][1] == ( + "MoE verify record: bytes_pre call=7 layer=1 row=3 (token 1, rank 1) expert=2 slot=1 bank=down " + "first_bad_byte=64 bad_words=4 | slot now holds L1/e2, slot_for_id=1, usage=9, lru_step=9, slot in: demand copy") + lines.clear() + engine_mod.Engine._emit_verify_stats(engine) # nothing checked and no new record: silent + assert lines == [] + session = v.report_session() + assert session[0] == ("warning", "MoE verify full (session): 1536 layer calls checked, meta_bad=0, pre_bad=1, " + "post_bad=0, post_meta_bad=0 (meta/pre/post/post_meta by layer: L1=0/1/0/0)") + assert session[1:] == [("warning", "MoE verify record: " + v.format_record(record))] + + +# FREETOKEN_MOE_SLOT_AUDIT (moe/slot_audit.py): the switch, where the auditor is built, and the +# engine's report; the recorders and the audit run on the qwen4_exp stacks in test_skeleton.py + + +def test_slot_audit_interval_switch(monkeypatch): + from freetoken.env import ENV + from freetoken.moe.slot_audit import resolve_interval + + monkeypatch.setattr(ENV.MOE_SLOT_AUDIT, "value", 64) + assert [resolve_interval(v) for v in (None, 0, 1, 8)] == [64, 0, 1, 8] + with pytest.raises(ValueError): + resolve_interval(-1) + + +@requires_cuda +def test_cache_builds_the_auditor_only_for_gpu_decode(monkeypatch): + from freetoken.env import ENV + from freetoken.moe.offload_cache import OffloadMoeCache + + def cache(**kw): + return OffloadMoeCache(num_layers=2, num_experts=4, cache_size=4, device=torch.device("cuda"), **kw) + + monkeypatch.setattr(ENV.MOE_SLOT_AUDIT, "value", 0) + assert cache().audit is None and cache(slot_audit=0).audit is None + on = cache(slot_audit=8, prefetch_mode="on") + assert on.audit.interval == 8 and on.slot_audit == 8 + ref = weakref.ref(on) # the auditor holds no reference to its cache + del on + assert ref() is None + with pytest.raises(ValueError): + cache(slot_audit=8, decode_target="hybrid") + monkeypatch.setattr(ENV.MOE_SLOT_AUDIT, "value", 16) + for target in ("hybrid", "cpu"): + assert cache(decode_target=target).audit is None + assert cache().audit.interval == 16 + monkeypatch.setenv("FREETOKEN_SKIP_FAST_INDEX_COPY", "1") + assert cache().audit is None, "copies that move nothing cannot be audited" + + +@requires_cuda +def test_engine_reports_the_audit_window_and_session(monkeypatch): + from types import SimpleNamespace + + from freetoken.engine import engine as engine_mod + from freetoken.engine.engine import MOE_STATS_INTERVAL + from freetoken.moe.offload_cache import OffloadMoeCache + + cache = OffloadMoeCache(num_layers=3, num_experts=4, cache_size=6, device=torch.device("cuda"), slot_audit=2) + cache.set_bank_sources({ + "gate_up": [torch.randn(4, 32, 8, dtype=torch.bfloat16).pin_memory() for _ in range(3)], + "down": [torch.randn(4, 8, 16, dtype=torch.bfloat16).pin_memory() for _ in range(3)], + }) + cache.reset() + lines = [] + for level in ("info_rank0", "warning_rank0"): + monkeypatch.setattr(engine_mod.logger, level, lambda msg, *a, level=level, **k: lines.append((level, msg))) + engine = SimpleNamespace(moe_offload_cache=cache) + engine_mod.Engine._emit_audit_stats(engine) + assert lines == [], "no audit ran: nothing to say" + + # one decode step of layer 1 routing experts 2 and 0: two demand installs and their copy + ids = torch.tensor([[2, 0]], dtype=torch.int32, device="cuda") + cache.ensure_experts(1, ids) + cache.copy_missing() + audit = cache.audit + assert not audit.end_decode_step(cache) and audit.end_decode_step(cache), "every 2 decode steps" + engine_mod.Engine._emit_audit_stats(engine) + assert lines == [("info_rank0", f"MoE slot audit ({MOE_STATS_INTERVAL} decode steps, every 2): 1 audits, 2 held slots " + "an audit (2 byte-checked), bad slots=0 (bytes=0, map=0, stale slot_for_id=0, owner out " + "of range=0), new=0, repeat=0, events=4; session: 1 audits, 0 bad (slot, owner) pairs")] + lines.clear() + slot = int(ids[0, 0]) + cache.bank_views()[1][slot].view(torch.uint8).view(-1)[8:12].bitwise_not_() + audit.scan(cache) + engine_mod.Engine._emit_audit_stats(engine) + assert [level for level, _ in lines] == ["warning_rank0", "warning_rank0"] + assert "1 audits, 2 held slots an audit (2 byte-checked), bad slots=1 (bytes=1," in lines[0][1] + assert lines[1][1].startswith( + f"MoE slot audit record: audit=1 step=0 lru=1 slot={slot} holds L1/e2 (slot_for_id={slot}, usage=1) | bytes differ in " + "down (first at down+8, 1 bad words) | diagnosis: demand install and demand copy of L1/e2 (#"), lines[1][1] + assert "then the bytes changed with no recorded write" in lines[1][1] + # the kernel ranks misses by id: e0 installs first (#0), e2 second (#1); the copy records after both + assert lines[1][1].endswith("history (1 map / 1 byte events in all): #1 demand install L1/e2 @step 0/lru 1 (then held " + "L1/e2, usage 1); #3 demand copy L1/e2 @step 0/lru 1 (slot held L1/e2, usage 1)"), lines[1][1] + lines.clear() + engine_mod.Engine._emit_audit_stats(engine) + assert lines == [] + session = audit.report_session() + assert session[0][0] == "warning" and session[0][1].startswith("MoE slot audit (session, every 2 decode steps): 2 audits") + assert session[1:] == [("warning", "MoE slot audit record: " + audit.format_record(audit.ring[0].tolist()))] diff --git a/tests/moe/test_scored_ensure.py b/tests/moe/test_scored_ensure.py new file mode 100644 index 0000000000..21318c257b --- /dev/null +++ b/tests/moe/test_scored_ensure.py @@ -0,0 +1,941 @@ +"""--moe-cache-policy kd / kdfb / rule: the vendored scored-victim ensure kernel. + +The kernel must reproduce flashlib's LRU bit for bit under POLICY=lru, and match the CPU +reference (ref_scored_cache.py, a port of the eviction replay) victim for victim under the +scored policies. Set FREETOKEN_MOE_ROUTING_TRACE to a compact routing trace directory +(experts.npy, logit_idx.npy, logit_val.npy) to also replay real decode rows. +""" +import os +from types import SimpleNamespace +from pathlib import Path + +import numpy as np +import pytest +import torch + +L, E, S, TOP_K = 48, 512, 1650, 10 +NEAR_MISS_THR = 0.25 +_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +_TRACE = os.environ.get("FREETOKEN_MOE_ROUTING_TRACE", "") + + +def _trace(): + if not _TRACE or not (Path(_TRACE) / "experts.npy").exists(): + pytest.skip("FREETOKEN_MOE_ROUTING_TRACE is not set to a routing trace directory") + load = lambda name: np.load(Path(_TRACE) / name, mmap_mode="r") # noqa: E731 + return load("experts.npy"), load("logit_idx.npy"), load("logit_val.npy") + + +def _trace_logits(li, lv): + """Router logits rebuilt from the trace's recorded top-32 (the rest far below any margin).""" + li, lv = np.asarray(li).astype(np.int64), np.asarray(lv).astype(np.float32) + logits = np.full(li.shape[:-1] + (E,), -1e4, np.float32) + np.put_along_axis(logits, li, lv, axis=-1) + return logits + + +def _zipf_rows(n, rows_per_step=1, seed=0, num_layers=L, num_experts=E, top_k=TOP_K, skew=0.9): + """Skewed decode routing ``[n, L, rows_per_step, top_k]``: distinct ids within a row, duplicates across rows.""" + rng = np.random.default_rng(seed) + p = 1.0 / np.arange(1, num_experts + 1) ** skew + perm = [rng.permutation(num_experts) for _ in range(num_layers)] + out = np.empty((n, num_layers, rows_per_step, top_k), np.int64) + for r in range(n): + for layer in range(num_layers): + for b in range(rows_per_step): + pick = rng.choice(num_experts, top_k, replace=False, p=p / p.sum()) + out[r, layer, b] = perm[layer][pick] + return out + + +def _random_logits(rows, seed=1, num_experts=E): + """bf16-exact router logits shaped like ``rows`` with the routed ids on top (~10% near misses at 0.25).""" + rng = np.random.default_rng(seed) + logits = rng.normal(size=rows.shape[:-1] + (num_experts,)).astype(np.float32) + np.put_along_axis(logits, rows, np.take_along_axis(logits, rows, axis=-1) + 3.0, axis=-1) + return torch.from_numpy(logits).bfloat16().float().numpy() + + +def _cache(policy, cache_size=S, num_layers=L, num_experts=E, device="cuda", **kw): + from freetoken.moe.offload_cache import OffloadMoeCache + + return OffloadMoeCache( + num_layers=num_layers, num_experts=num_experts, cache_size=cache_size, + device=torch.device(device), cache_policy=policy, **kw, + ) + + +def _ref(policy, cache_size=S, num_layers=L, num_experts=E): + from freetoken.moe.scored_ensure import POLICY_IDS + + from .ref_scored_cache import RefScoredCache + + return RefScoredCache(num_layers, num_experts, cache_size, POLICY_IDS[policy]) + + +def _run_gpu(cache, rows, logits=None, dtype=torch.bfloat16): + """Decode ``rows`` ``[N, L, B, K]`` through ``cache.ensure_experts``; per-row miss counts.""" + dev = cache.device + ids = torch.from_numpy(np.ascontiguousarray(rows, dtype=np.int32)).to(dev) + lg = None if logits is None else torch.from_numpy(np.ascontiguousarray(logits)).to(dev, dtype) + cache.collect_stats = True + cache.lru_stats.zero_() + marks = torch.zeros(rows.shape[0], dtype=torch.int64, device=dev) + for r in range(rows.shape[0]): + for layer in range(rows.shape[1]): + cache.ensure_experts(layer, ids[r, layer], router_logits=None if lg is None else lg[r, layer]) + marks[r] = cache.lru_stats[:, 1].sum() + return np.diff(marks.cpu().numpy(), prepend=0) + + +def _run_ref(ref, rows, logits=None): + from .ref_scored_cache import replay + + return replay(ref, rows, logits, NEAR_MISS_THR) + + +def _assert_same_tables(cache, ref): + got = lambda t: t.detach().cpu().numpy().astype(np.int64) # noqa: E731 + np.testing.assert_array_equal(got(cache.slot_for_id.view(-1)), ref.slot_of_id) + np.testing.assert_array_equal(got(cache.id_of_slot), ref.id_of_slot) + np.testing.assert_array_equal(got(cache.usage), ref.usage) + assert int(cache.step) == ref.step + if ref.policy: + assert int(cache.evict_tok) == ref.tok + np.testing.assert_array_equal(got(cache.evict_last_tok), ref.last_tok) + if ref.policy >= 2: + np.testing.assert_array_equal(got(cache.evict_lc), ref.lc) + np.testing.assert_array_equal(got(cache.evict_ct), ref.ct) + if ref.policy: + _assert_mirrors_hold(cache) + + +def _assert_mirrors_hold(cache): + """Every slot whose owner tag matches its id mirrors that id's state.""" + got = lambda t: t.detach().cpu().numpy().astype(np.int64) # noqa: E731 + ids, owner = got(cache.id_of_slot), got(cache.evict_slot_owner) + tagged = (ids >= 0) & (owner == ids) + np.testing.assert_array_equal(got(cache.evict_slot_last_tok)[tagged], got(cache.evict_last_tok)[ids[tagged]]) + if cache.evict_lc is not None: + np.testing.assert_array_equal(got(cache.evict_slot_lc)[tagged], got(cache.evict_lc)[ids[tagged]]) + np.testing.assert_array_equal(got(cache.evict_slot_ct)[tagged], got(cache.evict_ct)[ids[tagged]]) + return tagged + + +def _check_against_ref(policy, rows, logits=None, dtype=torch.bfloat16): + cache, ref = _cache(policy), _ref(policy) + assert policy != "rule" or cache.near_miss_thr == NEAR_MISS_THR + got = _run_gpu(cache, rows, logits, dtype) + want = _run_ref(ref, rows, logits) + np.testing.assert_array_equal(got, want) + _assert_same_tables(cache, ref) + return got + + +# (a) POLICY=lru is flashlib's kernel ------------------------------------------------------- + + +@_cuda +@pytest.mark.parametrize("rows_per_step", [1, 2]) +def test_vendored_lru_is_bit_identical_to_flashlib(rows_per_step): + from flashlib.kernels.slot_cache import N_STATS, lru_ensure + + from freetoken.moe.scored_ensure import scored_ensure + + dev = torch.device("cuda") + k = rows_per_step * TOP_K + # skewed enough that some calls hit every id + stream = _zipf_rows(63, rows_per_step, seed=rows_per_step, skew=1.6).reshape(-1, k)[:3000] + layers = np.arange(stream.shape[0]) % L + + def state(): + return dict( + slot_of_id=torch.full((L * E,), -1, dtype=torch.int32, device=dev), + id_of_slot=torch.full((S,), -1, dtype=torch.int32, device=dev), + usage=torch.zeros(S, dtype=torch.int64, device=dev), + step=torch.zeros((), dtype=torch.int64, device=dev), + src=torch.zeros(S, dtype=torch.int32, device=dev), + dst=torch.zeros(S, dtype=torch.int32, device=dev), + num=torch.zeros((), dtype=torch.int64, device=dev), + stats=torch.zeros(N_STATS, dtype=torch.int64, device=dev), + ) + + a, b = state(), state() + queries = torch.from_numpy(stream.astype(np.int32)).to(dev) + qa, qb = queries.clone(), queries.clone() + plans = ([], []) + for i in range(stream.shape[0]): + base = int(layers[i]) * E + lru_ensure(qa[i], a["slot_of_id"], a["id_of_slot"], a["usage"], a["step"], qa[i], + a["src"], a["dst"], a["num"], stats=a["stats"], id_base=base) + scored_ensure(qb[i], b["slot_of_id"], b["id_of_slot"], b["usage"], b["step"], qb[i], + b["src"], b["dst"], b["num"], policy=0, num_layers=L, num_experts=E, + stats=b["stats"], id_base=base) + for st, plan in zip((a, b), plans): + plan.append(torch.cat([st["num"].view(1), st["src"][:k].long(), st["dst"][:k].long()])) + assert torch.equal(qa, qb), "out slots differ" + pa, pb = (torch.stack(plan) for plan in plans) + n = pa[:, 0] + assert int(n.sum()) > 1000 and int((n == 0).sum()) > 0 # the stream exercises misses and all-hit calls + for i in range(pa.shape[0]): + m = int(n[i]) + assert torch.equal(pa[i, : 1 + m], pb[i, : 1 + m]) and torch.equal(pa[i, 1 + k : 1 + k + m], pb[i, 1 + k : 1 + k + m]) + for key in ("slot_of_id", "id_of_slot", "usage", "step", "stats"): + assert torch.equal(a[key], b[key]), key + + +def test_lru_policy_stays_on_flashlib(monkeypatch): + import freetoken.moe.offload_kernels as ok + + cache = _cache("lru", device="cpu") + calls = [] + monkeypatch.setattr(ok, "lru_ensure", lambda *a, **kw: calls.append("flashlib")) + monkeypatch.setattr(ok, "ensure_experts_scored", lambda *a, **kw: calls.append("scored")) + cache.ensure_experts(3, torch.zeros(TOP_K, dtype=torch.int32)) + cache.ensure_experts(3, torch.zeros(TOP_K, dtype=torch.int32), update_state=False) + assert calls == ["flashlib", "flashlib"] + assert cache.evict_tok is None and cache.evict_last_tok is None and cache.near_miss_thr is None + + +# (b) scored policies == CPU reference ------------------------------------------------------ + + +@_cuda +@pytest.mark.parametrize("policy", ["lru", "kd", "kdfb", "rule"]) +def test_scored_kernel_matches_cpu_reference_on_synthetic_routing(policy): + rows = _zipf_rows(40, seed=3) + got = _check_against_ref(policy, rows, _random_logits(rows) if policy == "rule" else None) + assert got.sum() > 0 + + +@_cuda +@pytest.mark.slow +@pytest.mark.parametrize("policy", ["kd", "kdfb", "rule"]) +def test_scored_kernel_matches_cpu_reference_on_trace_rows(policy): + experts, li, lv = _trace() + rows = np.asarray(experts[:2000]).astype(np.int64)[:, :, None, :] + logits = _trace_logits(li[:2000], lv[:2000])[:, :, None, :] if policy == "rule" else None + got = _check_against_ref(policy, rows, logits, torch.float32) + lru = _run_ref(_ref("lru"), rows) + assert got.sum() < lru.sum() # the point of the policy + + +# (c) victims: never the call's own ids, empty slots first ---------------------------------- + + +@_cuda +@pytest.mark.parametrize("policy", ["kd", "kdfb", "rule"]) +@pytest.mark.parametrize("rows_per_step", [1, 2]) +def test_victims_skip_the_call_and_take_empty_slots_first(policy, rows_per_step): + from freetoken.moe import scored_ensure as se + + num_layers, num_experts, cache_size = 4, 64, 100 + cache = _cache(policy, cache_size, num_layers, num_experts) + ref = _ref(policy, cache_size, num_layers, num_experts) + warm = _zipf_rows(6, rows_per_step, seed=5, num_layers=num_layers, num_experts=num_experts) + _run_gpu(cache, warm) + _run_ref(ref, warm) + assert int((cache.id_of_slot >= 0).sum()) == cache_size + + layer = 2 + held = cache.id_of_slot.cpu().numpy() + resident = sorted(int(i) - layer * num_experts for i in held if layer * num_experts <= i < (layer + 1) * num_experts) + absent = sorted(set(range(num_experts)) - set(resident)) + k = rows_per_step * TOP_K + hits = resident[: k // 2] + ids = np.array((hits + absent)[:k], np.int64) + if rows_per_step == 2: + ids[-1] = ids[0] # a duplicate hit + # make this call's hits the stalest ids in the cache, so only the pin keeps them + hit_keys = [layer * num_experts + e for e in hits] + for key in hit_keys: + cache.evict_last_tok[key] = se.LAST_TOK_NEVER + ref.last_tok[key] = se.LAST_TOK_NEVER + if cache.evict_lc is not None: + cache.evict_lc[key] = se.LC_NEVER + ref.lc[key] = se.LC_NEVER + hit_slots = {int(cache.slot_for_id[layer, e]) for e in hits} + # invalidate three slots (as the prefill overlap does): they must be the first victims + empty = sorted([s for s in (7, 40, 91, 13, 55, 77) if s not in hit_slots][:3]) + for slot in empty: + for t, v in ((cache.id_of_slot, held[slot]), (ref.id_of_slot, held[slot])): + t[slot] = -1 + cache.slot_for_id.view(-1)[int(held[slot])] = -1 + ref.slot_of_id[int(held[slot])] = -1 + cache.usage[slot] = 0 + ref.usage[slot] = 0 + + q = torch.tensor(ids, dtype=torch.int32, device="cuda") + lg = _random_logits(ids.reshape(rows_per_step, TOP_K), seed=9, num_experts=num_experts) + cache.ensure_experts(layer, q, router_logits=torch.from_numpy(lg).cuda()) + out, src, dst = ref.ensure(layer, ids, logits=lg, thr=NEAR_MISS_THR) + + n = int(cache.num_indices) + got_dst = cache.evict_slots[:n].tolist() + assert n == len(set(ids.tolist()) - set(hits)) == src.size + assert got_dst[: len(empty)] == empty + assert not hit_slots & set(got_dst) + assert q.tolist() == out.tolist() + assert all(int(cache.slot_for_id[layer, e]) >= 0 for e in ids.tolist()) + assert got_dst == dst.tolist() + _assert_same_tables(cache, ref) + + +# (d) the decode-step clock ----------------------------------------------------------------- + + +@_cuda +@pytest.mark.parametrize("rows_per_step", [1, 2]) +def test_token_clock_moves_once_per_decode_step(rows_per_step): + cache = _cache("kdfb") + rows = _zipf_rows(5, rows_per_step, seed=7) + ids = torch.from_numpy(rows.astype(np.int32)).cuda() + for r in range(rows.shape[0]): + for layer in range(L): + cache.ensure_experts(layer, ids[r, layer]) + assert int(cache.evict_tok) == r + 1 + # a CPU-decoded first layer moves the clock to the first GPU layer + cache.reset() + cache.cpu_layer_ids = frozenset({0, 1}) + ids = torch.from_numpy(rows.astype(np.int32)).cuda() + for layer in range(2, L): + cache.ensure_experts(layer, ids[0, layer]) + assert int(cache.evict_tok) == 1 + + +@_cuda +def test_small_prefill_ensures_leave_the_state_alone(monkeypatch): + import freetoken.layers.moe as moe_mod + from freetoken.distributed import set_tp_info, try_get_tp_info + from freetoken.layers.moe import OffloadMoELayer + from freetoken.layers.quantization import NoQuantConfig + + if try_get_tp_info() is None: + set_tp_info(rank=0, size=1) + dev = torch.device("cuda") + experts, top_k, hidden, inter = 64, 4, 16, 32 + layer = OffloadMoELayer(0, experts, top_k, hidden, inter, quant_config=NoQuantConfig(), + prefix="model.layers.0.mlp.experts") + cache = _cache("kdfb", 96, 1, experts) + g = torch.Generator().manual_seed(0) + cache.set_bank_sources({ + "gate_up": [torch.randn(experts, 2 * inter, hidden, generator=g, dtype=torch.bfloat16) * 0.1], + "down": [torch.randn(experts, hidden, inter, generator=g, dtype=torch.bfloat16) * 0.1], + }) + layer.offload_cache = cache + decode = torch.tensor([[1, 2, 3, 4]], dtype=torch.int32, device=dev) + cache.ensure_experts(0, decode) + before = [t.clone() for t in (cache.evict_tok, cache.evict_last_tok, cache.evict_lc, cache.evict_ct)] + + monkeypatch.setattr(moe_mod, "_SMALL_PREFILL_TOKENS", 64) + tokens = 40 + ids = torch.randint(0, experts // 2, (tokens, top_k), dtype=torch.int32, device=dev) + w = torch.full((tokens, top_k), 1.0 / top_k, device=dev) + layer._prefill_routed(torch.randn(tokens, hidden, dtype=torch.bfloat16, device=dev), w, ids.clone()) + assert bool((cache.slot_for_id[0, ids.long().unique()] >= 0).all()) # it did make them resident + assert int((cache.slot_for_id[0] >= 0).sum()) < experts # through the on-demand path, not the whole layer + after = (cache.evict_tok, cache.evict_last_tok, cache.evict_lc, cache.evict_ct) + for b, a in zip(before, after): + assert torch.equal(b, a) + + +@_cuda +def test_dsv4_on_demand_prefill_leaves_the_state_alone(): + from types import SimpleNamespace + + from freetoken.models.deepseek_v4.moe import DSV4OffloadMoELayer + + experts, top_k = 64, 4 + cache = _cache("kdfb", 96, 1, experts) + cache.set_bank_sources({"gate_up": [torch.zeros(experts, 2, 4)], "down": [torch.zeros(experts, 4, 2)]}) + cache.ensure_experts(0, torch.tensor([1, 2, 3, 4], dtype=torch.int32, device="cuda")) + before = [t.clone() for t in (cache.evict_tok, cache.evict_last_tok, cache.evict_lc, cache.evict_ct)] + # few enough routes (3 x 4 < 64) for the slot path rather than whole-layer streaming + layer = SimpleNamespace(offload_cache=cache, num_experts=experts, top_k=top_k, layer_id=0, + _expert_gemm=lambda *a, **kw: None) + ids = torch.tensor([[5, 6, 7, 8], [9, 10, 11, 12], [5, 13, 14, 15]], dtype=torch.int32, device="cuda") + DSV4OffloadMoELayer._prefill_routed(layer, torch.zeros(3, 8, device="cuda"), torch.full((3, top_k), 0.25), ids) + assert bool((cache.slot_for_id[0, 5:16] >= 0).all()) + for b, a in zip(before, (cache.evict_tok, cache.evict_last_tok, cache.evict_lc, cache.evict_ct)): + assert torch.equal(b, a) + + +# (e) reset() and rebuild() ----------------------------------------------------------------- + + +@_cuda +@pytest.mark.parametrize("policy", ["kd", "kdfb", "rule"]) +def test_reset_and_rebuild_clear_the_state(policy): + from freetoken.moe import scored_ensure as se + + cache = _cache(policy, 1480) + cache.set_bank_sources({"gate_up": [torch.zeros(E, 2, 4)] * L, "down": [torch.zeros(E, 4, 2)] * L}) + rows = _zipf_rows(3, seed=11) + fresh = _ref(policy) + + def assert_cold(): + assert int(cache.evict_tok) == 0 + assert bool((cache.evict_last_tok == se.LAST_TOK_NEVER).all()) + assert bool((cache.evict_slot_owner == -1).all()) and cache.evict_slot_owner.numel() == cache.cache_size + if cache.evict_lc is not None: + assert bool((cache.evict_lc == se.LC_NEVER).all()) and bool((cache.evict_ct == 0).all()) + + for clear in (cache.reset, lambda: cache.rebuild(S)): + _run_gpu(cache, rows) + assert int(cache.evict_tok) == rows.shape[0] + clear() + assert_cold() + assert cache.cache_size == S + # a rebuilt cache replays exactly like a fresh one + np.testing.assert_array_equal(_run_gpu(cache, rows), _run_ref(fresh, rows)) + _assert_same_tables(cache, fresh) + + +# (f) CUDA graph replay == eager ------------------------------------------------------------ + + +@_cuda +@pytest.mark.parametrize("policy", ["kdfb", "rule"]) +@pytest.mark.parametrize("rows_per_step", [1, 2]) +def test_graph_replay_matches_eager(policy, rows_per_step): + steps = 50 + rows = _zipf_rows(steps + 1, rows_per_step, seed=13) + ids = torch.from_numpy(rows.astype(np.int32)).cuda() + logits = torch.from_numpy(_random_logits(rows, seed=14)).cuda().bfloat16() + + eager = _cache(policy) + for r in range(1, steps + 1): + for layer in range(L): + eager.ensure_experts(layer, ids[r, layer].clone(), router_logits=logits[r, layer]) + + graphed = _cache(policy) + static_ids = ids[0].clone() + static_logits = logits[0].clone() + + def step(): + for layer in range(L): + graphed.ensure_experts(layer, static_ids[layer], router_logits=static_logits[layer]) + + step() # warm-up compiles the kernel variants outside capture + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + step() + graphed.reset() # as graph.py does after capture + for r in range(1, steps + 1): + static_ids.copy_(ids[r]) + static_logits.copy_(logits[r]) + graph.replay() + torch.cuda.synchronize() + + assert int(graphed.evict_tok) == steps + for name in ("slot_for_id", "id_of_slot", "usage", "step", "evict_tok", "evict_last_tok", "evict_lc", "evict_ct", + "evict_slot_owner", "evict_slot_last_tok", "evict_slot_lc", "evict_slot_ct"): + assert torch.equal(getattr(eager, name), getattr(graphed, name)), name + + +# (g) bs=2: two rows per call --------------------------------------------------------------- + + +@_cuda +@pytest.mark.parametrize("policy", ["kd", "kdfb", "rule"]) +def test_two_rows_per_call_match_cpu_reference(policy): + rows = _zipf_rows(30, 2, seed=17) + assert any(len(set(rows[r, layer].ravel())) < 2 * TOP_K for r in range(30) for layer in range(L)) + _check_against_ref(policy, rows, _random_logits(rows, seed=18) if policy == "rule" else None) + + +@_cuda +@pytest.mark.slow +@pytest.mark.parametrize("policy", ["kdfb", "rule"]) +def test_two_trace_streams_per_call_match_cpu_reference(policy): + experts, li, lv = _trace() + a, b = slice(0, 600), slice(45000, 45600) # two different requests decoding together + rows = np.stack([np.asarray(experts[a]), np.asarray(experts[b])], axis=2).astype(np.int64) + logits = None + if policy == "rule": + logits = np.stack([_trace_logits(li[a], lv[a]), _trace_logits(li[b], lv[b])], axis=2) + _check_against_ref(policy, rows, logits, torch.float32) + + +# the slot mirrors: foreign installs and state-free ensures fall back to the per-id state ----- + + +@_cuda +@pytest.mark.parametrize("policy", ["kd", "kdfb", "rule"]) +def test_materialized_prefill_between_decodes_matches_cpu_reference(policy): + cache, ref = _cache(policy), _ref(policy) + rows = _zipf_rows(24, seed=21) + logits = _random_logits(rows, seed=22) if policy == "rule" else None + # layer 5 comes back to slots it lost to layer 9: its tags there would otherwise look current + for part, layer in ((slice(0, 6), 5), (slice(6, 12), 9), (slice(12, 18), 5), (slice(18, 24), None)): + lg = None if logits is None else logits[part] + np.testing.assert_array_equal(_run_gpu(cache, rows[part], lg), _run_ref(ref, rows[part], lg)) + _assert_same_tables(cache, ref) + if layer is not None: + # the whole-layer prefill path re-seats this layer's experts at slots [0, E) + cache.materialize_layer(layer) + ref.materialize(layer) + assert not bool((cache.evict_slot_owner[:E] >= 0).any()) + _assert_same_tables(cache, ref) + + +@_cuda +@pytest.mark.parametrize("policy", ["kd", "kdfb"]) +def test_materialize_untags_the_slots_it_fills(policy): + # expert 3 of layer 0 leaves slot 3, is used elsewhere, then materialize puts it back at slot 3 + cache = _cache(policy, 24, 2, 8) + dev = torch.device("cuda") + ids = lambda *e: torch.tensor(e, dtype=torch.int32, device=dev) # noqa: E731 + cache.materialize_layer(0) + cache.ensure_experts(0, ids(3, 4)) # a hit at slot 3 tags it with expert 3 + assert int(cache.evict_slot_owner[3]) == 3 + cache.materialize_layer(1) # slot 3 now holds layer 1's expert 3 + cache.ensure_experts(1, ids(0, 1)) + cache.ensure_experts(0, ids(3, 5)) # expert 3 comes back elsewhere, with new state + assert int(cache.slot_for_id[0, 3]) != 3 + cache.materialize_layer(0) + assert int(cache.slot_for_id[0, 3]) == 3 + _assert_mirrors_hold(cache) + + +@_cuda +@pytest.mark.parametrize("policy", ["kd", "kdfb", "rule"]) +@pytest.mark.parametrize("chunks", [1, 2]) +def test_state_free_ensures_between_decodes_match_cpu_reference(policy, chunks): + cache, ref = _cache(policy), _ref(policy) + rows = _zipf_rows(12, seed=23) + extra = _zipf_rows(4, 3, seed=24, skew=0.3) + logits = _random_logits(rows, seed=25) if policy == "rule" else None + for r in range(rows.shape[0]): + lg = None if logits is None else logits[r : r + 1] + np.testing.assert_array_equal(_run_gpu(cache, rows[r : r + 1], lg), _run_ref(ref, rows[r : r + 1], lg)) + if r % 3 == 2: + for layer in range(0, L, 7): # a small prefill's chunked ensures + ids = np.unique(extra[r // 3, layer]) + # two chunks pin each other's slots as _ensure_unique does; one chunk keeps the plain pin + pin = cache.step.clone() if chunks > 1 else None + ref_pin = ref.step if chunks > 1 else None + for part in np.array_split(ids, chunks): + q = torch.from_numpy(part.astype(np.int32)).cuda() + cache.ensure_experts(layer, q, update_state=False, pin_since=pin) + out, _, dst = ref.ensure(layer, part, update_state=False, pin_since=ref_pin) + assert q.tolist() == out.tolist() + assert bool((cache.evict_slot_owner[torch.from_numpy(dst).cuda()] == -1).all()) + assert bool((cache.slot_for_id[layer, torch.from_numpy(ids).cuda()] >= 0).all()) + _assert_same_tables(cache, ref) + assert not _assert_mirrors_hold(cache).all() # some slots still read the per-id state + + +# rule: the near misses come from the router logits ---------------------------------------- + + +@_cuda +@pytest.mark.parametrize("rows_per_step", [1, 2]) +def test_rule_refreshes_the_router_near_misses(rows_per_step): + from freetoken.moe import scored_ensure as se + + cache = _cache("rule") + layer = 5 + g = torch.Generator(device="cuda").manual_seed(rows_per_step) + logits = (torch.randn(rows_per_step, E, device="cuda", generator=g) * 1.5).bfloat16() + ids = torch.topk(logits.float(), TOP_K, dim=-1).indices.to(torch.int32) + cache.ensure_experts(0, torch.arange(TOP_K, dtype=torch.int32, device="cuda")) # step 1 starts + cache.ensure_experts(layer, ids.clone(), router_logits=logits) + kth = logits.float().gather(1, ids.long()).min(dim=1, keepdim=True).values + near = (logits.float() >= kth - NEAR_MISS_THR).any(0) + routed = torch.zeros(E, dtype=torch.bool, device="cuda") + routed[ids.long().view(-1)] = True + last = cache.evict_last_tok[layer * E : (layer + 1) * E] + assert bool((last[routed] == 1).all()) + assert bool((last[near & ~routed] == 0).all()) + assert bool((last[~near & ~routed] == se.LAST_TOK_NEVER).all()) + assert 0 < int((near & ~routed).sum()) < E - TOP_K * rows_per_step + + +@pytest.mark.parametrize("policy", ["lru", "kdfb", "rule"]) +def test_decode_forward_hands_the_router_logits_to_ensure(monkeypatch, policy): + from freetoken.distributed import set_tp_info, try_get_tp_info + from freetoken.layers.moe import OffloadMoELayer + from freetoken.layers.quantization import NoQuantConfig + + if try_get_tp_info() is None: + set_tp_info(rank=0, size=1) + layer = OffloadMoELayer(0, 8, 2, 8, 16, quant_config=NoQuantConfig(), prefix="model.layers.0.mlp.experts") + layer.offload_cache = cache = _cache(policy, 12, 1, 8, device="cpu") + seen = {} + monkeypatch.setattr( + "freetoken.layers.moe.fused_topk", + lambda **kw: (torch.full((2, 2), 0.5), torch.zeros((2, 2), dtype=torch.int32)), + ) + monkeypatch.setattr(cache, "ensure_experts", lambda lid, ids, **kw: seen.update(kw)) + monkeypatch.setattr(cache, "copy_missing", lambda: None) + monkeypatch.setattr(cache, "pool_bank_views", lambda layer_id: ()) + monkeypatch.setattr(layer, "_expert_gemm", lambda *a, **kw: None) + router_logits = torch.randn(2, 8) + layer.decode_forward(torch.randn(2, 8), router_logits) + assert seen["router_logits"] is router_logits + assert cache.near_miss_thr == (NEAR_MISS_THR if policy == "rule" else None) + + +def test_policy_flag_and_gpu_only(): + import contextlib + import io + + from freetoken.server.args import parse_args + + buf = io.StringIO() + with contextlib.redirect_stdout(buf), contextlib.suppress(SystemExit): + parse_args(["--help"]) + assert "{lru,kd,kdfb,rule}" in buf.getvalue() + for policy in ("kd", "kdfb", "rule"): + _cache(policy, 12, 1, 8, device="cpu") + with pytest.raises(ValueError, match="GPU decode"): + _cache(policy, 12, 1, 8, device="cpu", decode_target="hybrid") + + +@_cuda +@pytest.mark.parametrize("policy", ["kd", "kdfb", "rule"]) +def test_an_all_pinned_call_never_stores_past_the_slot_arrays(policy): + # 5 slots pad to an 8-lane block; with every slot pinned the sentinel key must not decode to lane 7. + # The cache enforces cache_size >= num_experts, so the kernel gets 5-slot views of a 16-slot cache. + from freetoken.moe.offload_kernels import ensure_experts_scored + + s, pad = 5, 3 + cache = _cache(policy, cache_size=16, num_layers=1, num_experts=16) + guarded = {} + for name in ("id_of_slot", "usage", "evict_slot_owner", "evict_slot_last_tok", "evict_slot_lc", "evict_slot_ct"): + t = getattr(cache, name) + if t is None: # kd keeps no lc/ct mirrors + continue + buf = torch.full((s + pad,), 7, dtype=t.dtype, device=t.device) + buf[:s] = t[:s] + setattr(cache, name, buf[:s]) + guarded[name] = buf + dev = cache.id_of_slot.device + ensure_experts_scored(cache, 0, torch.arange(s, dtype=torch.int32, device=dev), bump_tok=True, update_state=True) + # 6 distinct ids against 5 slots breaks the caller contract; the kernel must still stay in bounds + q = torch.arange(s + 1, dtype=torch.int32, device=dev) + ensure_experts_scored(cache, 0, q, bump_tok=True, update_state=True) + torch.cuda.synchronize() + for name, buf in guarded.items(): + assert torch.all(buf[s:] == 7), f"{name} written past slot {s - 1}" + assert int(cache.num_indices) == 1 and 0 <= int(cache.evict_slots[0]) < s + + +# (h) FREETOKEN_MOE_PREFETCH=on: prefetch installs and low-priority keys ---------------------- + +_PF_L, _PF_E, _PF_S, _PF_K, _PF_W = 4, 64, 100, 6, 16 + + +def _pf_cache(policy, cache_size=_PF_S, num_layers=_PF_L, num_experts=_PF_E): + return _cache(policy, cache_size, num_layers, num_experts, prefetch_mode="on") + + +def _gpu_prefetch(cache, layer, sel, stats_row, width=_PF_W, budget=None): + """prefetch_ensure on the current stream; returns (src_ids, dst_slots), the plan it wrote.""" + from freetoken.moe.offload_kernels import prefetch_ensure_experts + + dev = cache.id_of_slot.device + query = torch.tensor(list(sel) + [-1] * (width - len(sel)), dtype=torch.int32, device=dev) + dst = torch.full((width,), -7, dtype=torch.int32, device=dev) + src = torch.full((width,), -7, dtype=torch.int32, device=dev) + num = torch.full((1,), 99, dtype=torch.int64, device=dev) + ready = torch.ones(1, dtype=torch.int32, device=dev) + prefetch_ensure_experts(cache, layer, query, dst, src, num, stats_row, ready, budget=budget) + n = int(num) + assert int(ready) == 0, "prefetch_ensure clears the copy's ready flag" + assert dst[n:].tolist() == [-7] * (width - n) + return src[:n].cpu().numpy().astype(np.int64), dst[:n].cpu().numpy().astype(np.int64), num + + +def _lowpri_slots(cache): + return set(((cache.id_of_slot >= 0) & (cache.usage == 0)).nonzero().view(-1).tolist()) + + +@_cuda +@pytest.mark.parametrize("policy", ["lru", "kd", "kdfb", "rule"]) +@pytest.mark.parametrize("rows_per_step", [1, 2]) +def test_prefetch_mode_matches_cpu_reference(policy, rows_per_step): + """prefetch_ensure and the low-priority demand ensure against the CPU reference victim for victim, + with adversarial candidates (next routed ids, resident ids, duplicates, padding, junk) and a + materialized prefill in between; the per-layer prefetch stats match the reference's counts.""" + from freetoken.moe.offload_kernels import ensure_experts_scored + from freetoken.moe.prefetch import CALLS, COPIED, ISSUED, LATE, MISSES, RESIDENT_HITS, ROWS, USEFUL + + rng = np.random.default_rng(31 + rows_per_step) + rows = _zipf_rows(40, rows_per_step, seed=33, num_layers=_PF_L, num_experts=_PF_E, top_k=_PF_K) + logits = _random_logits(rows, seed=34, num_experts=_PF_E) if policy == "rule" else None + cache, ref = _pf_cache(policy), _ref(policy, _PF_S, _PF_L, _PF_E) + dev = cache.id_of_slot.device + stats = torch.zeros((_PF_L, 8), dtype=torch.int64, device=dev) + want = np.zeros((_PF_L, 8), np.int64) + ready = torch.zeros(1, dtype=torch.int32, device=dev) + for r in range(rows.shape[0]): + if r == 20: + cache.materialize_layer(2) + ref.materialize(2) + assert not {s for s in _lowpri_slots(cache) if s < _PF_E}, "materialize installs at usage step" + for layer in range(_PF_L): + ids = rows[r, layer] + plan = None + if layer and rng.random() < 0.9: + routed = list(dict.fromkeys(ids.ravel().tolist())) + held = [e for e in range(_PF_E) if ref.slot_of_id[layer * _PF_E + e] >= 0] + sel = rng.choice(routed, size=min(3, len(routed)), replace=False).tolist() + sel += rng.choice(_PF_E, size=4).tolist() + held[:2] + sel[:1] + [-1] + rng.shuffle(sel) + pinned = set(np.flatnonzero((ref.usage == ref.step) & (ref.id_of_slot >= 0)).tolist()) + budget = [None, 0, 1, 2, 4][int(rng.integers(5))] + src, dst, num = _gpu_prefetch(cache, layer, sel, stats[layer], budget=budget) + want_src, want_dst = ref.prefetch(layer, sel, budget) + np.testing.assert_array_equal(src, want_src) + np.testing.assert_array_equal(dst, want_dst) + assert not pinned & set(dst.tolist()), "a prefetch evicted a slot of the last demand call" + assert set(dst.tolist()) <= _lowpri_slots(cache) + want[layer, ISSUED] += len(want_src) + _assert_same_tables(cache, ref) + plan = num + ready.fill_(int(rng.random() < 0.5)) # a copy that has (1) or has not (0) landed + q = torch.from_numpy(np.ascontiguousarray(ids, dtype=np.int32)).to(dev).view(-1) + lg = None if logits is None else torch.from_numpy(np.ascontiguousarray(logits[r, layer])).to(dev, torch.bfloat16) + lowpri_before = _lowpri_slots(cache) + ensure_experts_scored( + cache, layer, q, bump_tok=layer == 0, update_state=True, router_logits=lg, lowpri=True, + pf_count=None if plan is None else (stats[layer], ready, plan, rows_per_step), + ) + out, src, _ = ref.ensure(layer, ids, bump_tok=layer == 0, logits=None if lg is None else lg.float().cpu().numpy(), + thr=NEAR_MISS_THR, lowpri=True) + np.testing.assert_array_equal(q.cpu().numpy(), out) + _assert_same_tables(cache, ref) + # a demand hit makes a prefetched slot an ordinary resident, tagged with its owner's state + converted = lowpri_before & set(out.tolist()) + assert all(int(cache.usage[s]) == int(cache.step) for s in converted) + if policy != "lru": + assert all(int(cache.evict_slot_owner[s]) == int(cache.id_of_slot[s]) for s in converted) + if plan is not None: + n = int(plan) + want[layer, [USEFUL, CALLS, ROWS, MISSES]] += [ref.useful, 1, rows_per_step, src.size] + want[layer, LATE] += int(n > 0 and int(ready) == 0) + want[layer, COPIED] += int(n > 0) + got = stats.cpu().numpy() + np.testing.assert_array_equal(got, want) + assert got[:, USEFUL].sum() > 0 and got[:, LATE].sum() > 0 and got[:, RESIDENT_HITS].sum() == 0 + assert got[:, COPIED].sum() > got[:, LATE].sum() + + +@_cuda +@pytest.mark.parametrize("policy", ["lru", "rule"]) +@pytest.mark.parametrize("rows_per_step", [1, 2, 8]) +def test_prefetch_mode_matches_cpu_reference_at_production_geometry(policy, rows_per_step): + """The same victim-for-victim check at the deployment's shape: 48 layers x 512 experts, 1480 + slots (a 2048-wide, 8-warp victim scan), top-10, up to 8 rows (every graph size that forks) + and 32-wide prefetch queries, with small-prefill ensures (pinned since their first chunk) and a + materialized layer in between. No prefetch ever evicts a slot of the last demand call.""" + from freetoken.moe.offload_kernels import ensure_experts_scored + + num_layers, num_experts, size, width = L, E, 1480, 32 + rng = np.random.default_rng(71 + rows_per_step) + rows = _zipf_rows(4, rows_per_step, seed=72, num_layers=num_layers, num_experts=num_experts, top_k=TOP_K) + logits = _random_logits(rows, seed=73, num_experts=num_experts) if policy == "rule" else None + cache, ref = _pf_cache(policy, size, num_layers, num_experts), _ref(policy, size, num_layers, num_experts) + dev = cache.id_of_slot.device + stats = torch.zeros((num_layers, 8), dtype=torch.int64, device=dev) + ready = torch.zeros(1, dtype=torch.int32, device=dev) + installs = 0 + for r in range(rows.shape[0]): + if r == 2: + cache.materialize_layer(5) + ref.materialize(5) + for layer in range(num_layers): + ids = rows[r, layer] + plan = None + if layer: + routed = list(dict.fromkeys(ids.ravel().tolist())) + held = [e for e in range(num_experts) if ref.slot_of_id[layer * num_experts + e] >= 0] + sel = rng.choice(routed, size=min(6, len(routed)), replace=False).tolist() + sel += rng.choice(num_experts, size=16).tolist() + held[:4] + sel[:2] + [-1, -1] + sel = sel[:width] + rng.shuffle(sel) + pinned = set(np.flatnonzero((ref.usage == ref.step) & (ref.id_of_slot >= 0)).tolist()) + src, dst, plan = _gpu_prefetch(cache, layer, sel, stats[layer], width=width) + want_src, want_dst = ref.prefetch(layer, sel) + np.testing.assert_array_equal(src, want_src) + np.testing.assert_array_equal(dst, want_dst) + assert not pinned & set(dst.tolist()), "a prefetch evicted a slot of the last demand call" + installs += len(dst) + _assert_same_tables(cache, ref) + q = torch.from_numpy(np.ascontiguousarray(ids, dtype=np.int32)).to(dev).view(-1) + lg = None if logits is None else torch.from_numpy(np.ascontiguousarray(logits[r, layer])).to(dev, torch.bfloat16) + ensure_experts_scored( + cache, layer, q, bump_tok=layer == 0, update_state=True, router_logits=lg, lowpri=True, + pf_count=None if plan is None else (stats[layer], ready, plan, rows_per_step), + ) + out, _, _ = ref.ensure(layer, ids, bump_tok=layer == 0, logits=None if lg is None else lg.float().cpu().numpy(), + thr=NEAR_MISS_THR, lowpri=True) + np.testing.assert_array_equal(q.cpu().numpy(), out) + _assert_same_tables(cache, ref) + # a short prefill of one layer between decode steps: chunked, pinned since its first chunk + layer = int(rng.integers(num_layers)) + uniq = np.unique(rng.choice(num_experts, size=60, replace=False)) + pin = cache.step.clone() if policy != "lru" else None + ref_pin = ref.step if policy != "lru" else None + for start in range(0, uniq.size, 32): + part = torch.from_numpy(uniq[start : start + 32].astype(np.int32)).to(dev) + cache.ensure_experts(layer, part, update_state=False, pin_since=pin) + out, _, _ = ref.ensure(layer, uniq[start : start + 32], update_state=False, pin_since=ref_pin, lowpri=True) + np.testing.assert_array_equal(part.cpu().numpy(), out) + _assert_same_tables(cache, ref) + assert installs > 0 and int((ref.id_of_slot >= 0).sum()) == size, "the cache must be full and evicting" + + +def _plant(cache, ref, slot, flat_id, usage): + for t_id, t_slot in ((cache.id_of_slot, cache.slot_for_id.view(-1)), (ref.id_of_slot, ref.slot_of_id)): + t_id[slot] = flat_id + t_slot[flat_id] = slot + cache.usage[slot] = usage + ref.usage[slot] = usage + + +@_cuda +@pytest.mark.parametrize("policy", ["lru", "kd", "kdfb", "rule"]) +def test_lowpri_band_sits_between_empty_slots_and_every_resident(policy): + """An unconsumed prefetch is evicted after every empty slot and before any resident, even the + coldest resident the score caps allow, whatever the slot order.""" + from freetoken.moe import scored_ensure as se + from freetoken.moe.offload_kernels import ensure_experts_scored + + num_layers, num_experts, size = 2, 16, 16 + cache, ref = _pf_cache(policy, size, num_layers, num_experts), _ref(policy, size, num_layers, num_experts) + dev = cache.id_of_slot.device + cache.step.fill_(5) + ref.step = 5 + # slots 0-12: layer-0 residents at the oldest usage 1 with never-seen state (the lowest possible + # score); slot 13 empty; slots 14, 15 unconsumed prefetches + for slot in range(13): + _plant(cache, ref, slot, slot, 1) + _plant(cache, ref, 14, 13, 0) + _plant(cache, ref, 15, 14, 0) + if policy != "lru": + cache.evict_slot_owner.fill_(-1) + victims = [] + for e in (15, 16 + 0, 16 + 1, 16 + 2): # four misses, one per call + layer, expert = divmod(e, num_experts) + q = torch.tensor([expert], dtype=torch.int32, device=dev) + ensure_experts_scored(cache, layer, q, bump_tok=False, update_state=True, lowpri=True) + ref.ensure(layer, [expert], lowpri=True) + victims.append(int(cache.evict_slots[0])) + _assert_same_tables(cache, ref) + assert victims[:3] == [13, 14, 15] and victims[3] < 13 + # the band's keys are below 2 << SLOT_BITS; the caps' worst resident score keeps score + 2^31 >= 2 + beta_step, w_q4, _, _ = se.score_params(se.MAX_BETA, se.MAX_W, 1, 1) + worst = -((se.K_MAX << se.Q) + 1 * beta_step) + ((w_q4 * se.LC_FLOOR) >> 4) + assert worst + se.SCORE_BIAS >= 2 + + +@_cuda +@pytest.mark.parametrize("policy", ["lru", "rule"]) +def test_prefetch_budget_keeps_the_first_non_resident_candidates(policy): + """The install budget counts only non-resident candidates, in query order: resident ids, padding + and duplicates ahead of them take none of it, and ISSUED counts the installs.""" + num_experts = 16 + cache, ref = _pf_cache(policy, num_experts, 2, num_experts), _ref(policy, num_experts, 2, num_experts) + stats = torch.zeros(8, dtype=torch.int64, device="cuda") + cache.ensure_experts(0, torch.tensor([0, 1], dtype=torch.int32, device="cuda")) # step 0 pins every slot + ref.ensure(0, np.arange(2), bump_tok=True, lowpri=True) + src, _, _ = _gpu_prefetch(cache, 1, [4, 6], stats) + ref.prefetch(1, [4, 6]) + assert src.tolist() == [4, 6] + sel = [4, -1, 6, 9, 9, 3, 7, 1] + src, dst, _ = _gpu_prefetch(cache, 1, sel, stats, budget=2) + assert src.tolist() == [9, 3] and int(stats[0]) == 4 + np.testing.assert_array_equal(np.stack(ref.prefetch(1, sel, 2)), np.stack([src, dst])) + _assert_same_tables(cache, ref) + src, _, _ = _gpu_prefetch(cache, 1, [5, 8], stats, budget=0) + assert src.size == 0 and int(stats[0]) == 4 + + +@_cuda +@pytest.mark.parametrize("policy", ["lru", "rule"]) +def test_prefetch_installs_only_into_evictable_slots(policy): + """With fewer evictable slots than candidates a prefetch installs the first candidates only, and + with every slot pinned by the last demand call it installs nothing (no all-pinned fallback).""" + num_experts = 16 + cache, ref = _pf_cache(policy, num_experts, 2, num_experts), _ref(policy, num_experts, 2, num_experts) + stats = torch.zeros(8, dtype=torch.int64, device="cuda") + ids = np.arange(12) + cache.ensure_experts(0, torch.from_numpy(ids.astype(np.int32)).cuda()) + ref.ensure(0, ids, bump_tok=True, lowpri=True) + pinned = cache.id_of_slot.clone() + src, dst, _ = _gpu_prefetch(cache, 1, [9, 3, 7, 1, 5, 11, 2], stats) + assert src.tolist() == [9, 3, 7, 1] and sorted(dst.tolist()) == [12, 13, 14, 15] + np.testing.assert_array_equal(np.stack(ref.prefetch(1, [9, 3, 7, 1, 5, 11, 2])), np.stack([src, dst])) + assert torch.equal(cache.id_of_slot[:12], pinned[:12]) and int(stats[0]) == 4 + all16 = np.arange(16) + cache.ensure_experts(0, torch.from_numpy(all16.astype(np.int32)).cuda()) + before = [t.clone() for t in (cache.id_of_slot, cache.slot_for_id, cache.usage)] + src, dst, num = _gpu_prefetch(cache, 1, [0, 1, 2], stats) + assert int(num) == 0 and src.size == 0 and int(stats[0]) == 4 + for a, b in zip(before, (cache.id_of_slot, cache.slot_for_id, cache.usage)): + assert torch.equal(a, b) + + +@_cuda +@pytest.mark.parametrize("policy", ["kd", "kdfb", "rule"]) +def test_prefetch_scores_victims_from_the_previous_layers_position(policy): + """prefetch_ensure(L) runs right after layer L-1's demand call and scores victims from there: + layer L's residents are next up (d = 1), so one used last token outlives an older resident of + a later layer instead of looking a whole pass away.""" + num_layers, num_experts, size = 8, 16, 16 + cache, ref = _pf_cache(policy, size, num_layers, num_experts), _ref(policy, size, num_layers, num_experts) + cache.step.fill_(7) + ref.step = 7 + cache.evict_tok.fill_(10) + ref.tok = 10 + cache.evict_slot_owner.fill_(-1) + # slots 0-13: layer 1's routed experts (pinned), slot 14: a layer-2 expert used last token, + # slot 15: a layer-5 expert last used three tokens ago; all with the same (never-counted) use count + planted = [(slot, 16 + slot, 7, 10) for slot in range(14)] + [(14, 2 * 16 + 3, 6, 9), (15, 5 * 16 + 4, 5, 7)] + for slot, flat, usage, last in planted: + _plant(cache, ref, slot, flat, usage) + cache.evict_last_tok[flat] = last + ref.last_tok[flat] = last + stats = torch.zeros(8, dtype=torch.int64, device="cuda") + src, dst, _ = _gpu_prefetch(cache, 2, [9], stats) + assert src.tolist() == [9] and dst.tolist() == [15], "the prefetch evicted the target layer's fresh resident" + np.testing.assert_array_equal(np.stack(ref.prefetch(2, [9])), np.stack([src, dst])) + _assert_same_tables(cache, ref) + + +@_cuda +@pytest.mark.parametrize("rows_per_step", [1, 2]) +def test_lowpri_lru_without_prefetches_is_flashlib(rows_per_step): + """Prefetch on routes LRU through the vendored kernel with the low-priority key; while no slot is + a prefetch it evicts exactly what flashlib's lru_ensure does.""" + num_layers, num_experts, size = 8, 64, 200 + on = _pf_cache("lru", size, num_layers, num_experts) + off = _cache("lru", size, num_layers, num_experts) + rows = _zipf_rows(30, rows_per_step, seed=41, num_layers=num_layers, num_experts=num_experts) + ids = torch.from_numpy(rows.astype(np.int32)).cuda() + for r in range(rows.shape[0]): + for layer in range(num_layers): + a, b = ids[r, layer].clone(), ids[r, layer].clone() + on.ensure_experts(layer, a) + off.ensure_experts(layer, b) + assert torch.equal(a, b) + assert torch.equal(on.evict_slots[: int(on.num_indices)], off.evict_slots[: int(off.num_indices)]) + for name in ("slot_for_id", "id_of_slot", "usage", "step"): + assert torch.equal(getattr(on, name), getattr(off, name)), name + + +def test_lru_policy_leaves_flashlib_only_when_prefetch_is_on(monkeypatch): + import freetoken.moe.offload_kernels as ok + + calls = [] + monkeypatch.setattr(ok, "lru_ensure", lambda *a, **kw: calls.append("flashlib")) + monkeypatch.setattr(ok, "ensure_experts_scored", lambda *a, **kw: calls.append(("scored", kw["lowpri"]))) + for mode in ("off", "measure", "on"): + cache = _cache("lru", 12, 2, 8, device="cpu", prefetch_mode=mode) + cache.prefetch = SimpleNamespace(mode=mode, count_args=lambda layer: None) if mode != "off" else None + cache.ensure_experts(1, torch.zeros(2, dtype=torch.int32)) + assert calls == ["flashlib", "flashlib", ("scored", True)] diff --git a/tests/moe/test_small_prefill_on_demand.py b/tests/moe/test_small_prefill_on_demand.py new file mode 100644 index 0000000000..f7971c67d5 --- /dev/null +++ b/tests/moe/test_small_prefill_on_demand.py @@ -0,0 +1,187 @@ +"""Short prefills load only their routed experts (FREETOKEN_MOE_SMALL_PREFILL_TOKENS) instead of +streaming whole layers; above the threshold, or with it off, the whole-layer path is unchanged.""" + +import pytest +import torch + +import freetoken.layers.moe as moe_mod +from freetoken.distributed import set_tp_info, try_get_tp_info + + +def _layer_and_cache(device, num_experts=8, top_k=2, hidden=16, inter=32, cache_size=None, policy="lru", num_layers=1): + from freetoken.layers.moe import OffloadMoELayer + from freetoken.layers.quantization import NoQuantConfig + from freetoken.moe.offload_cache import OffloadMoeCache + + if try_get_tp_info() is None: + set_tp_info(rank=0, size=1) + layer = OffloadMoELayer(0, num_experts, top_k, hidden, inter, quant_config=NoQuantConfig(), + prefix="model.layers.0.mlp.experts") + cache = OffloadMoeCache(num_layers=num_layers, num_experts=num_experts, cache_size=cache_size or num_experts, + device=device, cache_policy=policy) + g = torch.Generator().manual_seed(0) + cache.set_bank_sources({ + "gate_up": [torch.randn(num_experts, 2 * inter, hidden, generator=g, dtype=torch.bfloat16) * 0.1 + for _ in range(num_layers)], + "down": [torch.randn(num_experts, hidden, inter, generator=g, dtype=torch.bfloat16) * 0.1 + for _ in range(num_layers)], + }) + layer.offload_cache = cache + return layer, cache + + +@pytest.mark.parametrize("tokens, threshold, on_demand", [ + (3, 8, True), (8, 8, True), (9, 8, False), (3, 0, False), +]) +def test_small_prefill_dispatch(monkeypatch, tokens, threshold, on_demand): + layer, cache = _layer_and_cache(torch.device("cpu")) + monkeypatch.setattr(moe_mod, "_SMALL_PREFILL_TOKENS", threshold) + calls = [] + monkeypatch.setattr(cache, "ensure_experts", lambda layer_id, ids, **kw: calls.append("ensure")) + monkeypatch.setattr(cache, "materialize_layer", lambda layer_id: calls.append("materialize")) + monkeypatch.setattr(cache, "copy_missing", lambda: calls.append("copy")) + monkeypatch.setattr(cache, "begin_prefill", lambda: calls.append("begin_prefill"), raising=False) + got = {} + + def fake_gemm(cache_, hs, w, ids, *, views, n, alphas, is_prefill): + got.update(n=n, is_prefill=is_prefill) + return hs + + monkeypatch.setattr(layer, "_expert_gemm", fake_gemm) + hs = torch.randn(tokens, 16) + layer._prefill_routed(hs, torch.full((tokens, 2), 0.5), torch.zeros(tokens, 2, dtype=torch.int32)) + if on_demand: + assert calls == ["ensure", "copy"] and got == {"n": cache.cache_size, "is_prefill": True} + else: + assert "ensure" not in calls and "begin_prefill" not in calls[1:] and got["n"] == layer.num_experts + + +def _record_paths(monkeypatch, layer, cache): + """Stub the cache movement and the GEMM; log each call and what the GEMM was handed.""" + calls, got = [], {} + + def ensure(layer_id, ids, **_): + calls.append(("ensure", ids.tolist())) + ids.add_(100) # the LRU kernel rewrites its query to slot ids in place + + monkeypatch.setattr(cache, "ensure_experts", ensure) + monkeypatch.setattr(cache, "materialize_layer", lambda layer_id: calls.append(("materialize",))) + monkeypatch.setattr(cache, "copy_missing", lambda: calls.append(("copy",))) + + def fake_gemm(cache_, hs, w, ids, *, views, n, alphas, is_prefill): + got.update(ids=ids.tolist(), n=n, is_prefill=is_prefill) + return hs + + monkeypatch.setattr(layer, "_expert_gemm", fake_gemm) + return calls, got + + +def test_small_prefill_leaves_an_unpinned_layer_on_the_whole_layer_path(monkeypatch): + # an unpinned (locked) layer's copy_missing is the whole-layer pageable copy, which presumes + # position == expert id and cannot serve an LRU slot remap + layer, cache = _layer_and_cache(torch.device("cpu")) + monkeypatch.setattr(moe_mod, "_SMALL_PREFILL_TOKENS", 8) + cache._unpinned_layers = frozenset({layer.layer_id}) + calls, got = _record_paths(monkeypatch, layer, cache) + layer._prefill_routed(torch.randn(3, 16), torch.full((3, 2), 0.5), torch.zeros(3, 2, dtype=torch.int32)) + assert calls == [("materialize",), ("copy",)] and got["n"] == layer.num_experts + + +def test_small_prefill_path_does_not_depend_on_how_many_experts_a_layer_routes_to(monkeypatch): + # 6 distinct experts in an 8-slot cache: every layer of a short prefill still takes the + # on-demand path, so no layer can fall back into an overlap buffer that was never begun + layer, cache = _layer_and_cache(torch.device("cpu")) + monkeypatch.setattr(moe_mod, "_SMALL_PREFILL_TOKENS", 8) + monkeypatch.setattr(moe_mod, "_ENSURE_CHUNK_IDS", 4) + calls, got = _record_paths(monkeypatch, layer, cache) + ids = torch.tensor([[5, 0], [1, 7], [2, 5], [4, 0]], dtype=torch.int32) + layer._prefill_routed(torch.randn(4, 16), torch.full((4, 2), 0.5), ids.clone()) + assert calls == [("ensure", [0, 1, 2, 4]), ("copy",), ("ensure", [5, 7]), ("copy",)] + assert got == {"ids": (ids + 100).tolist(), "n": cache.cache_size, "is_prefill": True} + + +def test_small_prefill_refuses_more_experts_than_the_cache_holds(monkeypatch): + layer, cache = _layer_and_cache(torch.device("cpu")) + monkeypatch.setattr(moe_mod, "_SMALL_PREFILL_TOKENS", 8) + _record_paths(monkeypatch, layer, cache) + cache.cache_size = 4 # below the engine's num_experts floor + ids = torch.tensor([[5, 0], [1, 7], [2, 5], [4, 0]], dtype=torch.int32) + with pytest.raises(AssertionError, match="exceed"): + layer._prefill_routed(torch.randn(4, 16), torch.full((4, 2), 0.5), ids) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +def test_small_prefill_matches_the_whole_layer_path(monkeypatch): + from freetoken.core import get_global_ctx # noqa: F401 (layer reads the ctx only in forward()) + + dev = torch.device("cuda") + hs = torch.randn(5, 16, dtype=torch.bfloat16, device=dev) * 0.5 + ids = torch.tensor([[0, 3], [1, 3], [7, 2], [3, 4], [5, 0]], dtype=torch.int32, device=dev) + w = torch.tensor([[0.6, 0.4], [0.5, 0.5], [0.9, 0.1], [0.3, 0.7], [0.2, 0.8]], device=dev) + outs = {} + for threshold in (0, 8): + monkeypatch.setattr(moe_mod, "_SMALL_PREFILL_TOKENS", threshold) + layer, cache = _layer_and_cache(dev, cache_size=32) + outs[threshold] = _run_and_check_path(layer, hs, w, ids, on_demand=threshold > 0) + torch.testing.assert_close(outs[8], outs[0], rtol=2e-2, atol=2e-2) + + +def _run_and_check_path(layer, hs, w, ids, *, on_demand): + taken = [] + real = getattr(layer, "_ensure_unique", None) + if real is not None: + layer._ensure_unique = lambda cache, t: taken.append(r := real(cache, t)) or r + out = layer._prefill_routed(hs.clone(), w.clone(), ids.clone()).float() + assert (bool(taken) and taken[0] is not None) == on_demand + return out + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +def test_a_real_sized_small_prefill_fits_the_lru_kernel(monkeypatch): + # 200 tokens x top-10 = 2000 routed ids against a cache that rounds up to 2048 slots: the + # LRU kernel's block is next_pow2(ids) x next_pow2(slots), so the raw ids overflow Triton's + # 1M-element limit, and a 512-id block failed to launch on a full GPU (CUDA OOM). The path + # must hand it the unique ids in decode-width chunks. + dev = torch.device("cuda") + tokens, experts, top_k = 200, 512, 10 + g = torch.Generator(device=dev).manual_seed(1) + hs = torch.randn(tokens, 16, dtype=torch.bfloat16, device=dev, generator=g) * 0.5 + ids = torch.randint(0, experts, (tokens, top_k), dtype=torch.int32, device=dev, generator=g) + w = torch.softmax(torch.randn(tokens, top_k, device=dev, generator=g), dim=-1) + outs = {} + for threshold in (0, 1024): + monkeypatch.setattr(moe_mod, "_SMALL_PREFILL_TOKENS", threshold) + layer, cache = _layer_and_cache(dev, num_experts=experts, top_k=top_k, cache_size=1100) + widths, ensure = [], cache.ensure_experts + monkeypatch.setattr(cache, "ensure_experts", lambda lid, q, **kw: (widths.append(q.numel()), ensure(lid, q, **kw))[1]) + outs[threshold] = _run_and_check_path(layer, hs, w, ids, on_demand=threshold > 0) + assert widths and max(widths) <= 32, f"LRU ensure widths {sorted(set(widths))}: wider than decode's" + torch.testing.assert_close(outs[1024], outs[0], rtol=2e-2, atol=2e-2) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +@pytest.mark.parametrize("policy", ["lru", "kd", "kdfb", "rule"]) +def test_a_later_chunk_keeps_the_experts_an_earlier_chunk_loaded(monkeypatch, policy): + # Decode fills every slot, then 40 never-decoded experts take two ensure chunks. A scored + # policy ranks the first chunk's fresh installs as its coldest slots, so they need a pin. + dev = torch.device("cuda") + experts, top_k, slots, tokens = 64, 4, 80, 10 + g = torch.Generator(device=dev).manual_seed(2) + hs = torch.randn(tokens, 16, dtype=torch.bfloat16, device=dev, generator=g) * 4 # a wrong expert is ~1.0 off + ids = torch.arange(16, 16 + tokens * top_k, dtype=torch.int32, device=dev).view(tokens, top_k) + w = torch.softmax(torch.randn(tokens, top_k, device=dev, generator=g), dim=-1) + outs = {} + for threshold in (0, 64): + monkeypatch.setattr(moe_mod, "_SMALL_PREFILL_TOKENS", threshold) + layer, cache = _layer_and_cache(dev, num_experts=experts, top_k=top_k, cache_size=slots, policy=policy, + num_layers=2) + for layer_id, first, last in ((1, 0, experts), (0, 0, 16)): + for start in range(first, last, 8): + cache.ensure_experts(layer_id, torch.arange(start, start + 8, dtype=torch.int32, device=dev)) + cache.copy_missing() + assert bool((cache.id_of_slot >= 0).all()) + outs[threshold] = _run_and_check_path(layer, hs, w, ids, on_demand=threshold > 0) + if threshold: + slot_of = cache.slot_for_id[0, ids.long()] + assert bool((slot_of >= 0).all()), f"{int((slot_of < 0).sum())} routed experts were evicted" + torch.testing.assert_close(outs[64], outs[0], rtol=2e-2, atol=0.1) diff --git a/tests/scheduler/test_abort_inflight_prefill.py b/tests/scheduler/test_abort_inflight_prefill.py index 85c217859f..11b2038a76 100644 --- a/tests/scheduler/test_abort_inflight_prefill.py +++ b/tests/scheduler/test_abort_inflight_prefill.py @@ -63,6 +63,7 @@ def _setup(): config=SimpleNamespace(page_size=1), status_reporter=SimpleNamespace(report_batch=lambda *_, **__: None), send_result=sent.extend, + _prefill_start={}, _kv_usage_pages=cm.page_usage, _mamba_slot_usage=lambda: None, _swa_token_usage=lambda: None, @@ -94,11 +95,19 @@ def _launch_req(pool, cm, tm, prompt, *, cls=Req, track_seqlen=None): def _as_last_data(batch): - return ( - SimpleNamespace(batch=batch), - (None, torch.tensor([42], dtype=torch.int32), - SimpleNamespace(synchronize=lambda: None)), + # Mirrors engine.ForwardOutput's drain-facing fields. _process_last_data reads + # attributes (next_tokens_cpu, copy_done_event, the *_cpu logprob columns), so a + # bare tuple no longer works; a namespace keeps this test off the engine import + # chain, whose kernel packages a scheduler unit test should not require. + outputs = SimpleNamespace( + next_tokens_gpu=None, + next_tokens_cpu=torch.tensor([42], dtype=torch.int32), + copy_done_event=SimpleNamespace(synchronize=lambda: None), + chosen_logprobs_cpu=None, + top_ids_cpu=None, + top_logprobs_cpu=None, ) + return (SimpleNamespace(batch=batch), outputs) def test_abort_inflight_final_chunk_marks_then_drains(): diff --git a/tests/scheduler/test_cost_accounting_core.py b/tests/scheduler/test_cost_accounting_core.py index 045010c6b4..512bfa4049 100644 --- a/tests/scheduler/test_cost_accounting_core.py +++ b/tests/scheduler/test_cost_accounting_core.py @@ -94,6 +94,7 @@ def test_schedule_reports_admission_only_after_prepare_succeeds(): scheduler.prefill_budget = 99 scheduler.prefill_manager = SimpleNamespace(schedule_next_batch=lambda budget: batch) scheduler.decode_manager = SimpleNamespace(schedule_next_batch=lambda: None) + scheduler._prefill_start = {} events = [] def prepare(value): @@ -110,6 +111,8 @@ def send(messages): assert events[0] == ("prepared", batch) sent = events[1][1] assert [(m.uid, m.prompt_tokens, m.cached_tokens) for m in sent] == [(1, 12, 4), (2, 34, 0)] + # Admission also opens each request's prefill span, closed at its first sampled token. + assert sorted(scheduler._prefill_start) == [1, 2] def test_prepare_failure_emits_no_prompt_admission(): @@ -130,9 +133,24 @@ def fail_prepare(_batch): assert sent == [] +def test_freeing_a_request_that_never_sampled_closes_its_prefill_span(): + """The first sampled token is what normally pops the span. A request aborted mid-prefill + never reaches that point, so the free path has to pop it -- otherwise a long-lived server + leaks one entry per aborted request.""" + scheduler = Scheduler.__new__(Scheduler) + scheduler._prefill_start = {7: 100.0, 9: 200.0} + scheduler.cache_manager = SimpleNamespace(cache_req=lambda req, finished: None) + scheduler.table_manager = SimpleNamespace(free=lambda table_idx: None) + + Scheduler._free_req_resources(scheduler, SimpleNamespace(uid=7, table_idx=3)) + + assert scheduler._prefill_start == {9: 200.0} # only the freed request's span is dropped + + def test_scheduler_rejection_emits_error_but_no_admission(): scheduler = Scheduler.__new__(Scheduler) - scheduler.engine = SimpleNamespace(max_seq_len=4) + scheduler.engine = SimpleNamespace(max_seq_len=4, num_pages=4) + scheduler.config = SimpleNamespace(page_size=1) added = [] scheduler.prefill_manager = SimpleNamespace(add_one_req=added.append) sent = [] diff --git a/tests/scheduler/test_hybrid_cache_manager.py b/tests/scheduler/test_hybrid_cache_manager.py index 098f27660c..165689ab3b 100644 --- a/tests/scheduler/test_hybrid_cache_manager.py +++ b/tests/scheduler/test_hybrid_cache_manager.py @@ -6,6 +6,7 @@ from types import SimpleNamespace import torch +import pytest from freetoken.core import Req, SamplingParams from freetoken.kvcache.linear_state_pool import LinearStatePool @@ -153,6 +154,83 @@ def test_naive_cache_does_not_align_prefill_chunks(): assert adder.try_add_one(pending).extend_len == 100 +@pytest.mark.parametrize("full_chunks", [2, 3]) +@pytest.mark.parametrize("tail", [1, 59, 64, 65]) +@pytest.mark.parametrize("finish_early", [False, True]) +def test_chunked_prefill_retains_resumable_snapshot(monkeypatch, full_chunks, tail, finish_early): + """A short final extend must retain the prior state; a longer one must replace it.""" + import freetoken.core as core + from freetoken.attention.linear import build_fla_metadata + from freetoken.scheduler.decode import DecodeManager + from freetoken.scheduler.prefill import ChunkedReq, PrefillManager + from freetoken.scheduler.scheduler import Scheduler + from freetoken.scheduler.table import TableManager + from freetoken.scheduler.utils import PendingReq + + pool = _pool() + pt = torch.zeros(2, 1024, dtype=torch.int32) + cm = CacheManager(32, 64, pt, "hybrid_radix", linear_state_pool=pool) + tm = TableManager(max_running_reqs=1, page_table=pt) + pm = PrefillManager(cm, tm, DecodeManager(page_size=64)) + monkeypatch.setattr(core, "_GLOBAL_CTX", core.Context(page_size=64, linear_state_pool=pool)) + prompt = torch.arange(full_chunks * 128 + tail, dtype=torch.int32) + pm.pending_list = [PendingReq(0, prompt, SamplingParams(max_tokens=8))] + last_batch = None + states = {} + final = None + + while pm.runnable or last_batch is not None: + # Match overlap ordering: prepare the continuation before draining the prior chunk. + batch = pm.schedule_next_batch(128) + if batch is not None: + batch.padded_reqs = batch.reqs + cm.allocate_paged(batch.reqs) + req = batch.reqs[0] + meta = build_fla_metadata(batch, torch.device("cpu")) + if meta.track_dst is not None: + slot = meta.track_dst.item() + value = len(states) + 1 + states[req.mamba_last_track_seqlen] = value + pool.recurrent_states[:, slot].fill_(value) + pool.conv_states[:, slot].fill_(value) + pool.recurrent_states[:, req.linear_slot_idx].fill_(-9) + pool.conv_states[:, req.linear_slot_idx].fill_(-9) + req.complete_one() + if last_batch is not None: + previous = last_batch.reqs[0] + if not isinstance(previous, ChunkedReq): + final = previous + if not finish_early: + cm.cache_req(final, finished=False) + last_batch = batch + + # Before the finish path can donate an aligned live state, require the frozen checkpoint. + expected = full_chunks * 128 - 64 if tail <= 64 else full_chunks * 128 + 64 + if not finish_early: + match = cm.match_req(_pend(prompt.tolist() + [999])) + assert match.cuda_handle.cached_len == expected + assert torch.all(pool.recurrent_states[:, match.mamba_value] == states[expected]) + assert torch.all(pool.conv_states[:, match.mamba_value] == states[expected]) + + # Finish/abort may occur before the prefill commit; use the real idempotent cleanup path. + stub = SimpleNamespace(cache_manager=cm, table_manager=tm, _prefill_start={}) + Scheduler._free_req_resources(stub, final) + Scheduler._free_req_resources(stub, final) + match = cm.match_req(_pend(prompt.tolist() + [999])) + expected_finish = len(prompt) if tail == 64 else expected + assert match.cuda_handle.cached_len == expected_finish + value = -9 if tail == 64 else states[expected] + assert torch.all(pool.recurrent_states[:, match.mamba_value] == value) + assert torch.all(pool.conv_states[:, match.mamba_value] == value) + assert tm.available_size == 1 + cm.check_integrity() + assert pool.num_free_slots + cm.prefix_cache.mamba_evictable_size == pool.num_slots - 1 + cm.ensure_mamba_slots(pool.num_slots - 1) + cm.check_integrity() + assert pool.num_free_slots == pool.num_slots - 1 + assert len(cm.free_slots) == cm.num_pages + + def test_pool_sizing_covers_4mr_floor(): """C6: pool must reserve the 4-slot-per-request non-evictable floor even at a tiny ratio.""" from types import SimpleNamespace @@ -163,8 +241,16 @@ def test_pool_sizing_covers_4mr_floor(): assert _linear_pool_num_slots(c) >= 4 * mr + 1, (mr, _linear_pool_num_slots(c)) +def test_host_snapshot_tier_lowers_the_gpu_snapshot_floor(monkeypatch): + """With FT_GDN_HOST_TIER evicted snapshots live in RAM, so one GPU snapshot slot is enough.""" + from types import SimpleNamespace + from freetoken.kvcache.linear_state_pool import _linear_pool_num_slots + c = SimpleNamespace(max_running_req=1, cache_type="hybrid_radix", linear_state_cache_ratio=0.0) + monkeypatch.delenv("FT_GDN_HOST_TIER", raising=False) + assert _linear_pool_num_slots(c) == 4 + 4 + 1 + monkeypatch.setenv("FT_GDN_HOST_TIER", "1") + assert _linear_pool_num_slots(c) == 4 + 1 + 1 + + if __name__ == "__main__": - for name, fn in list(globals().items()): - if name.startswith("test_") and callable(fn): - fn() - print(f"{name}: PASS") + raise SystemExit(pytest.main([__file__])) diff --git a/tests/scheduler/test_interleave.py b/tests/scheduler/test_interleave.py new file mode 100644 index 0000000000..42a3f7221e --- /dev/null +++ b/tests/scheduler/test_interleave.py @@ -0,0 +1,346 @@ +"""Tests for DecodeInterleavePolicy and its wiring into Scheduler._schedule_next_batch. + +The production symptom this exists for (measured 2026-09-15): a 300k-token prompt +is chunked at the window pool budget (384 tok), so it becomes 797 consecutive +prefill steps, and an already-decoding request is not scheduled once during the +267-317 s that takes. Prefill must keep priority (a request that never prefills +never starts) but must yield to decode periodically. +""" + +from __future__ import annotations + +import pytest + +from freetoken.scheduler.interleave import DecodeInterleavePolicy + + +# --------------------------------------------------------------------------- # +# policy, in isolation +# --------------------------------------------------------------------------- # +def test_disabled_by_default_and_never_forces_decode(): + """Unset reproduces historical prefill-first behaviour exactly.""" + p = DecodeInterleavePolicy() + assert not p.enabled + for _ in range(10_000): + p.note_prefill() + assert not p.wants_decode() + + +@pytest.mark.parametrize("every", [1, 2, 8, 16]) +def test_forces_decode_after_exactly_every_prefills(every): + p = DecodeInterleavePolicy(every) + assert p.enabled + for i in range(1, every): + p.note_prefill() + assert not p.wants_decode(), f"fired early at {i} < {every}" + p.note_prefill() + assert p.wants_decode(), f"did not fire at {every}" + + +def test_decode_resets_the_streak(): + p = DecodeInterleavePolicy(4) + for _ in range(4): + p.note_prefill() + assert p.wants_decode() + p.note_decode() + assert not p.wants_decode(), "streak must reset after a decode step" + for _ in range(3): + p.note_prefill() + assert not p.wants_decode() + p.note_prefill() + assert p.wants_decode() + + +def test_every_one_interleaves_strictly(): + """every=1 is the extreme: decode after every single prefill step.""" + p = DecodeInterleavePolicy(1) + p.note_prefill() + assert p.wants_decode() + + +@pytest.mark.parametrize("bad", [0, -1, -100]) +def test_non_positive_disables(bad): + p = DecodeInterleavePolicy(bad) + assert not p.enabled + for _ in range(50): + p.note_prefill() + assert not p.wants_decode() + + +def test_reset_clears_credit(): + p = DecodeInterleavePolicy(8) + for _ in range(7): + p.note_prefill() + p.reset() + assert not p.wants_decode() + p.note_prefill() + assert not p.wants_decode() + + +# --------------------------------------------------------------------------- # +# the decision the scheduler actually makes +# --------------------------------------------------------------------------- # +def _decide(policy, has_prefill, has_decode, every=8): + """Mirror of the scheduler's new decision, as a pure function. + + prefill wins unless the policy says decode is due AND a decode is available. + """ + if has_prefill and not (policy.wants_decode() and has_decode): + policy.note_prefill() + return "prefill" + if has_decode: + policy.note_decode() + return "decode" + if has_prefill: + policy.note_prefill() + return "prefill" + return None + + +def test_burst_is_punctuated_by_decode_every_n(): + """797 prefill steps in a row -- the production shape -- must contain decode steps.""" + policy = DecodeInterleavePolicy(8) + seq = [] + for _ in range(797): + seq.append(_decide(policy, has_prefill=True, has_decode=True)) + n_dec = seq.count("decode") + # A decode step resets the streak, so a cycle is every+1 slots (8 prefill + 1 decode): + # 797 slots hold 88 full cycles, i.e. 88 decode steps. + assert n_dec == 797 // 9, f"expected {797 // 9} decode steps, got {n_dec}" + # and no run of prefills exceeds the threshold + run = 0 + worst = 0 + for s in seq: + run = run + 1 if s == "prefill" else 0 + worst = max(worst, run) + assert worst == 8, f"longest prefill run {worst} exceeds the threshold" + + +def test_disabled_policy_yields_the_historical_sequence(): + policy = DecodeInterleavePolicy() + seq = [_decide(policy, True, True) for _ in range(100)] + assert seq.count("decode") == 0, "disabled must never take a decode slot" + + +def test_decode_is_not_forced_when_none_is_runnable(): + """Nothing to decode -> prefill keeps the slot (no idle step).""" + policy = DecodeInterleavePolicy(2) + seq = [_decide(policy, has_prefill=True, has_decode=False) for _ in range(10)] + assert seq == ["prefill"] * 10 + + +def test_prefill_keeps_priority_below_the_threshold(): + """The first N-1 slots of a burst must still go to prefill.""" + policy = DecodeInterleavePolicy(8) + seq = [_decide(policy, True, True) for _ in range(7)] + assert seq.count("decode") == 0 + + +def test_pure_decode_phase_unaffected(): + """With nothing to prefill, every slot is decode (no policy interference).""" + policy = DecodeInterleavePolicy(8) + seq = [_decide(policy, has_prefill=False, has_decode=True) for _ in range(50)] + assert seq == ["decode"] * 50 + + +def test_scheduler_without_policy_keeps_prefill_first(): + """A Scheduler built without __init__ (as the accounting tests do) must not break. + + _schedule_next_batch reads the policy through getattr because the accounting tests + build a Scheduler this way; a stripped object keeps the historical prefill-first + order instead of raising. + """ + from types import SimpleNamespace + + from freetoken.scheduler.scheduler import Scheduler + + prefill_batch = SimpleNamespace(is_prefill=True, prompt_admissions=[]) + decode_batch = SimpleNamespace(is_prefill=False, prompt_admissions=[]) + s = Scheduler.__new__(Scheduler) + s.prefill_budget = 384 + s.prefill_manager = SimpleNamespace(schedule_next_batch=lambda budget: prefill_batch) + s.decode_manager = SimpleNamespace(schedule_next_batch=lambda: decode_batch) + s._prepare_batch = lambda value: value + s.send_result = lambda messages: None + # no _interleave attribute at all + assert not hasattr(s, "_interleave") + assert Scheduler._schedule_next_batch(s) is prefill_batch, "must stay prefill-first" + + +def test_scheduler_with_policy_and_no_decode_stays_prefill(): + """When no decode is runnable the policy must not steal a slot.""" + from types import SimpleNamespace + + from freetoken.scheduler.scheduler import Scheduler + from freetoken.scheduler.interleave import DecodeInterleavePolicy + + prefill_batch = SimpleNamespace(is_prefill=True, prompt_admissions=[]) + s = Scheduler.__new__(Scheduler) + s.prefill_budget = 384 + s.prefill_manager = SimpleNamespace(schedule_next_batch=lambda budget: prefill_batch) + s.decode_manager = SimpleNamespace(schedule_next_batch=lambda: None) + s._prepare_batch = lambda value: value + s.send_result = lambda messages: None + s._interleave = DecodeInterleavePolicy(1) # fires every step + for _ in range(5): + assert Scheduler._schedule_next_batch(s) is prefill_batch + + +def test_scheduler_with_policy_takes_decode_when_due(): + """Once the streak reaches the threshold, a runnable decode gets the slot.""" + from types import SimpleNamespace + + from freetoken.scheduler.scheduler import Scheduler + from freetoken.scheduler.interleave import DecodeInterleavePolicy + + prefill_batch = SimpleNamespace(is_prefill=True, prompt_admissions=[]) + decode_batch = SimpleNamespace(is_prefill=False, prompt_admissions=[]) + s = Scheduler.__new__(Scheduler) + s.prefill_budget = 384 + s.prefill_manager = SimpleNamespace(schedule_next_batch=lambda budget: prefill_batch) + s.decode_manager = SimpleNamespace(schedule_next_batch=lambda: decode_batch) + s._prepare_batch = lambda value: value + s.send_result = lambda messages: None + s._interleave = DecodeInterleavePolicy(3) + got = [Scheduler._schedule_next_batch(s) for _ in range(4)] + assert got[0] is prefill_batch and got[1] is prefill_batch and got[2] is prefill_batch + assert got[3] is decode_batch, "4th slot must be the forced decode" + + +# --------------------------------------------------------------------------- # +# Truth table against the REAL scheduler (codex review: _decide is a model of +# the decision, not the decision -- it can drift from the code it describes). +# --------------------------------------------------------------------------- # +def _real_scheduler(every=None, prefill=True, decode=True): + """Build a Scheduler whose two managers are stubs with call counters.""" + from types import SimpleNamespace + + from freetoken.scheduler.scheduler import Scheduler + from freetoken.scheduler.interleave import DecodeInterleavePolicy + + calls = {"prefill": 0, "decode": 0} + pf = SimpleNamespace(is_prefill=True, prompt_admissions=[]) + dc = SimpleNamespace(is_prefill=False, prompt_admissions=[]) + + def _pf(budget): + calls["prefill"] += 1 + return pf if prefill else None + + def _dc(): + calls["decode"] += 1 + return dc if decode else None + + s = Scheduler.__new__(Scheduler) + s.prefill_budget = 384 + s.prefill_manager = SimpleNamespace(schedule_next_batch=_pf) + s.decode_manager = SimpleNamespace(schedule_next_batch=_dc) + s._prepare_batch = lambda value: value + s.send_result = lambda messages: None + if every is not None: + s._interleave = DecodeInterleavePolicy(every) + return s, calls, pf, dc + + +@pytest.mark.parametrize("every,expected", [ + # (policy N, expected sequence of slot winners over 9 slots) + (None, ["p"] * 9), # disabled -> always prefill + # streak starts at 0, so slot 0 is prefill; then strict alternation + (1, ["p", "d", "p", "d", "p", "d", "p", "d", "p"]), + (3, ["p", "p", "p", "d", "p", "p", "p", "d", "p"]), + (8, ["p"] * 8 + ["d"]), +]) +def test_truth_table_against_real_scheduler(every, expected): + """The real _schedule_next_batch must produce exactly this slot sequence.""" + from freetoken.scheduler.scheduler import Scheduler + + s, calls, pf, dc = _real_scheduler(every=every) + got = [] + for _ in range(9): + b = Scheduler._schedule_next_batch(s) + got.append("p" if b is pf else ("d" if b is dc else "?")) + assert got == expected, f"every={every}: got {got}, want {expected}" + + +def test_every_one_alternates_while_both_stay_runnable(): + """every=1 is the degenerate case: one decode between every prefill.""" + from freetoken.scheduler.scheduler import Scheduler + + s, calls, pf, dc = _real_scheduler(every=1) + seq = [] + for _ in range(6): + b = Scheduler._schedule_next_batch(s) + seq.append("p" if b is pf else "d") + # slot 0 prefill (streak starts at 0), then alternation + assert seq == ["p", "d", "p", "d", "p", "d"] + + +def test_latch_releases_when_decode_becomes_available(): + """Due decode unavailable -> keep prefilling -> decode becomes available -> it fires. + + Exercises the subtle path: the policy must not lose its 'decode is due' state + while waiting, and must not fire twice once it does. + """ + from freetoken.scheduler.scheduler import Scheduler + + s, calls, pf, dc = _real_scheduler(every=2) + # phase 1: no decode runnable at all + s.decode_manager.schedule_next_batch = lambda: None + for _ in range(5): + assert Scheduler._schedule_next_batch(s) is pf + # phase 2: decode becomes available -> next slot must be decode + s.decode_manager.schedule_next_batch = lambda: dc + assert Scheduler._schedule_next_batch(s) is dc, "latch did not release" + + +def test_no_work_returns_none(): + from freetoken.scheduler.scheduler import Scheduler + + s, _, _, _ = _real_scheduler(every=8, prefill=False, decode=False) + assert Scheduler._schedule_next_batch(s) is None + + +def test_pure_decode_fallback_when_nothing_to_prefill(): + """With no prefill available every slot is decode (no policy interference).""" + from freetoken.scheduler.scheduler import Scheduler + + s, calls, pf, dc = _real_scheduler(every=2, prefill=False, decode=True) + got = [Scheduler._schedule_next_batch(s) for _ in range(5)] + assert all(b is dc for b in got) + + +def test_disabled_policy_object_not_just_missing_attribute(): + """codex: cover a policy that EXISTS but is disabled, not only an absent one.""" + from freetoken.scheduler.scheduler import Scheduler + + s, calls, pf, dc = _real_scheduler(every=None) + s._interleave = __import__( + "freetoken.scheduler.interleave", fromlist=["DecodeInterleavePolicy"] + ).DecodeInterleavePolicy(0) # present but disabled + got = [Scheduler._schedule_next_batch(s) for _ in range(6)] + assert all(b is pf for b in got), "a disabled policy must never take a slot" + + +def test_decode_manager_not_probed_when_policy_disabled(): + """When disabled, the decode manager must not be consulted at all (no extra call).""" + from freetoken.scheduler.scheduler import Scheduler + + s, calls, pf, dc = _real_scheduler(every=None) + Scheduler._schedule_next_batch(s) + assert calls["decode"] == 0, "disabled policy should not probe decode" + + +def test_config_to_scheduler_wiring(): + """codex: cover CLI/config -> Scheduler.__init__ wiring.""" + import inspect + + from freetoken.engine.config import EngineConfig + from freetoken.scheduler import scheduler as sched_mod + + assert "decode_interleave_every" in {f.name for f in __import__("dataclasses").fields(EngineConfig)} + src = inspect.getsource(sched_mod.Scheduler.__init__) + assert "_interleave" in src and "decode_interleave_every" in src, ( + "Scheduler.__init__ must build the policy from the config field" + ) + # and the CLI exposes it + from freetoken.server.args import ServerArgs + assert hasattr(ServerArgs, "decode_interleave_every") diff --git a/tests/scheduler/test_match_stop_str.py b/tests/scheduler/test_match_stop_str.py new file mode 100644 index 0000000000..e9a5c07da5 --- /dev/null +++ b/tests/scheduler/test_match_stop_str.py @@ -0,0 +1,86 @@ +"""Scheduler-side stop-string matching (Scheduler._match_stop_str) against the same +incrementally decoded text the frontend streams.""" + +from __future__ import annotations + +from types import SimpleNamespace + +import torch + +from freetoken.core import Req, SamplingParams +from freetoken.scheduler.scheduler import Scheduler + +EOS, SPECIAL = 0, 1 + + +class _ByteTok: + """Byte-level ids like Qwen's BPE fallback: a 4-byte emoji is four tokens. 0 is EOS, + 1 a special marker that skip_special_tokens drops.""" + + def decode(self, ids, skip_special_tokens=False): + def piece(t): + if t == EOS: + return b"" + if t == SPECIAL: + return b"" if skip_special_tokens else b"" + return bytes([t]) + + return b"".join(piece(t) for t in ids).decode("utf-8", errors="replace") + + +def _req(prompt: list[int], output_len: int, **sp) -> Req: + return Req( + input_ids=torch.tensor(prompt, dtype=torch.int32), + table_idx=0, + cached_len=0, + output_len=output_len, + uid=1, + sampling_params=SamplingParams(**sp), + cache_handle=None, + ) + + +def _match(req: Req) -> str | None: + stub = SimpleNamespace(tokenizer=_ByteTok(), eos_token_ids={EOS}) + return Scheduler._match_stop_str(stub, req) + + +def _feed(req: Req, ids: list[int]) -> list[str | None]: + seen = [] + for t in ids: + req.append_host(torch.tensor([t], dtype=torch.int32)) + req.complete_one() + seen.append(_match(req)) + return seen + + +def test_a_stop_spanning_more_tokens_than_characters_is_matched(): + dna = list("\N{DNA DOUBLE HELIX}".encode()) # one character, four byte tokens + req = _req([65, 66], 32, stop_strs=["\N{DNA DOUBLE HELIX}"]) + text = list(b"Here is DNA ") + dna + list(b" and more") + seen = _feed(req, text) + first = next(i for i, m in enumerate(seen) if m is not None) + assert first == len(b"Here is DNA ") + len(dna) - 1 + + +def test_an_eos_short_of_the_delivered_limit_stays_in_the_matched_text(): + # Overlap scheduling advances device_len ahead of the host: the budget is spent on the + # device (can_decode is False) while the host holds one token less than the limit. That + # EOS is not terminal, so the frontend decodes it and the scheduler must too. + req = _req([65], 3, stop_strs=["zz"], ignore_eos=True) + for t in (65, EOS): + req.append_host(torch.tensor([t], dtype=torch.int32)) + while req.can_decode: + req.complete_one() + assert req.input_ids.numel() < req.max_device_len + assert _match(req) is None + assert req.stop_decode_status.decoded_ids == [65, EOS] + + +def test_stops_match_the_text_decoded_with_the_requests_skip_special_tokens(): + # the frontend streams "ab" for a request that skips special tokens; the scheduler has + # to match the stop in that same text, not in "ab" + req = _req([65], 8, stop_strs=["ab"], skip_special_tokens=True) + assert _feed(req, [ord("a"), SPECIAL, ord("b")]) == [None, None, "ab"] + req = _req([65], 8, stop_strs=["ab"]) + assert _feed(req, [ord("a"), SPECIAL, ord("b")]) == [None, None, None] diff --git a/tests/server/test_accounting.py b/tests/server/test_accounting.py new file mode 100644 index 0000000000..fe29d414d8 --- /dev/null +++ b/tests/server/test_accounting.py @@ -0,0 +1,121 @@ +"""Failure/retry and concurrency contracts for the terminal accounting barrier.""" +import asyncio +from types import SimpleNamespace + +import pytest + +from freetoken.message import UserReply +from freetoken.server.accounting import AccountingDrainError, prepare_stop_accounting +from freetoken.server.stats import StatsTracker + + +def state_with_request(): + stats = StatsTracker() + stats.on_new_user(7) + return SimpleNamespace(stats=stats, maintenance_state="serving", instance_id="instance", + ready_at=None, config=SimpleNamespace(served_model_name="test")) + + +def terminal(state, tokens=3): + state.stats.observe(UserReply(uid=7, incremental_output="", finished=True, + completion_tokens_delta=tokens, error="aborted")) + + +def test_drain_timeout_then_late_terminal_can_retry_without_losing_tokens(): + async def run(): + state = state_with_request() + aborted = [] + + async def abort(uid): + aborted.append(uid) + state.stats.on_abort(uid) # enqueueing abort is NOT a terminal reply + + state.abort_user = abort + with pytest.raises(AccountingDrainError, match="abort barrier timed out"): + await prepare_stop_accounting(state, drain_timeout_s=0, abort_timeout_s=0) + assert state.maintenance_state == "stopping" + assert state.stats.active == 1 and not hasattr(state, "_sealed_accounting") + terminal(state, tokens=11) # last sampled-token delta arrives AFTER the timeout + sealed = await prepare_stop_accounting(state, drain_timeout_s=0, abort_timeout_s=0) + assert sealed["completion_tokens_total"] == 11 + assert sealed["drain_complete"] is True + assert aborted == [7] + sealed["completion_tokens_total"] = -1 + assert (await prepare_stop_accounting(state))["completion_tokens_total"] == 11 + asyncio.run(run()) + + +def test_concurrent_prepare_stop_serializes_abort_and_returns_independent_snapshots(): + async def run(): + state = state_with_request() + entered, release = asyncio.Event(), asyncio.Event() + aborted = [] + + async def abort(uid): + assert state.maintenance_state == "stopping" + aborted.append(uid) + entered.set() + await release.wait() + terminal(state) + + state.abort_user = abort + first = asyncio.create_task(prepare_stop_accounting(state, drain_timeout_s=0)) + await entered.wait() + second = asyncio.create_task(prepare_stop_accounting(state, drain_timeout_s=0)) + await asyncio.sleep(0) # second caller reaches the shared lock, no wall-clock race + assert not second.done() + release.set() + a, b = await asyncio.gather(first, second) + assert a == b and a is not b + assert aborted == [7] + assert a["completion_tokens_total"] == 3 + asyncio.run(run()) + + +def test_abort_transport_failure_preserves_engine_and_closes_gate(): + state = state_with_request() + + async def abort(uid): + raise ConnectionError("scheduler disconnected") + + state.abort_user = abort + with pytest.raises(AccountingDrainError, match="failed to abort active requests") as error: + asyncio.run(prepare_stop_accounting(state, drain_timeout_s=0, abort_timeout_s=0)) + assert isinstance(error.value.__cause__, ConnectionError) + assert state.maintenance_state == "stopping" and state.stats.active == 1 + assert not hasattr(state, "_sealed_accounting") + + +def test_unidentified_active_request_cannot_be_sealed(): + state = state_with_request() + state.stats = SimpleNamespace(active=1, inflight_uids=()) + with pytest.raises(AccountingDrainError, match="unidentified"): + asyncio.run(prepare_stop_accounting(state, drain_timeout_s=0)) + assert state.maintenance_state == "stopping" + assert not hasattr(state, "_sealed_accounting") + + +@pytest.mark.parametrize("maintenance", ["rebuilding", "unknown"]) +def test_invalid_prepare_transition_preserves_previous_state(maintenance): + state = state_with_request() + state.maintenance_state = maintenance + with pytest.raises(AccountingDrainError): + asyncio.run(prepare_stop_accounting(state, drain_timeout_s=0)) + assert state.maintenance_state == maintenance + assert not hasattr(state, "_sealed_accounting") + + +def test_cancelled_drain_releases_lock_without_reopening_admission(): + async def run(): + state = state_with_request() + task = asyncio.create_task(prepare_stop_accounting(state, drain_timeout_s=10)) + await asyncio.sleep(0) # admission closes before the first drain wait + assert state.maintenance_state == "stopping" + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + assert not state._accounting_prepare_lock.locked() + assert not hasattr(state, "_sealed_accounting") + terminal(state) + assert (await prepare_stop_accounting(state))["drain_complete"] is True + asyncio.run(run()) diff --git a/tests/server/test_anthropic_api.py b/tests/server/test_anthropic_api.py index f382093fc0..890339522c 100644 --- a/tests/server/test_anthropic_api.py +++ b/tests/server/test_anthropic_api.py @@ -33,6 +33,7 @@ ContentDelta, GenDone, GenResult, + GenTimings, ReasoningDelta, ToolCallsDelta, ) @@ -43,12 +44,14 @@ async def _aiter(items): yield it -def _collect_events(events, model="claude-x", uid=1, cache_report=False): +def _collect_events(events, model="claude-x", uid=1, cache_report=False, metrics=False): """Run the Anthropic event stream over neutral GenEvents; return [(type, data), ...].""" async def run(): out = [] - async for frame in A.anthropic_event_stream(_aiter(events), model, uid, cache_report=cache_report): + async for frame in A.anthropic_event_stream( + _aiter(events), model, uid, cache_report=cache_report, metrics=metrics + ): etype = None data = None for line in frame.split("\n"): @@ -428,7 +431,8 @@ def test_stream_reasoning_as_thinking_block(): "message_stop", ] blocks = [e[1] for e in collected if e[0] == "content_block_start"] - assert blocks[0]["content_block"]["type"] == "thinking" + # signature is a required string in the SDK's thinking block, present from the start. + assert blocks[0]["content_block"] == {"type": "thinking", "thinking": "", "signature": ""} assert blocks[0]["index"] == 0 assert blocks[1]["content_block"]["type"] == "text" assert blocks[1]["index"] == 1 @@ -1063,3 +1067,52 @@ def test_image_only_tool_result_keeps_an_empty_tool_message(): "role": "user", "content": [{"type": "image", "freetoken_ref": {"kind": "b64", "data": "aGk="}}], } + + +# --------------------------------------------------------------------------- # +# Per-request metrics (--enable-metrics-report) +# --------------------------------------------------------------------------- # +def _timed_result() -> GenResult: + return GenResult( + reasoning="", content="hi", tool_calls=[], finish_reason="stop", + prompt_tokens=11, completion_tokens=5, cached_tokens=8, + timings=GenTimings(ttft_ms=30.0, prefill_ms=12.0, decode_ms=40.0, total_ms=70.0), + ) + + +def test_full_response_metrics_absent_unless_asked_for(): + """Anthropic's wire has no `metrics` field, so /v1/messages must stay byte-identical for + every client that did not opt in -- absent, not null.""" + resp = A.anthropic_full_response(_timed_result(), "claude-x", uid=9) + assert "metrics" not in resp.model_dump(exclude_none=True) + + +def test_full_response_metrics_report_the_untouched_prompt_split(): + """cache_report rewrites usage.input_tokens to exclude the cached prefix; metrics must + still describe the real prefill, which is prompt_tokens - cached_tokens.""" + resp = A.anthropic_full_response( + _timed_result(), "claude-x", uid=9, cache_report=True, metrics=True + ) + body = resp.model_dump(exclude_none=True) + assert body["usage"]["input_tokens"] == 3 + assert body["metrics"]["prefill_tokens"] == 3 + assert body["metrics"]["cached_prompt_tokens"] == 8 + assert body["metrics"]["prefill_time_ms"] == 12.0 + assert body["metrics"]["ttft_ms"] == 30.0 + + +def test_stream_message_delta_carries_metrics(): + events = [ + ContentDelta("hi"), + GenDone("stop", 4, 2, cached_tokens=1, + timings=GenTimings(ttft_ms=5.0, prefill_ms=2.0, decode_ms=9.0, total_ms=14.0)), + ] + collected = _collect_events(events, metrics=True) + md = next(e[1] for e in collected if e[0] == "message_delta") + assert md["metrics"]["prefill_time_ms"] == 2.0 + assert md["metrics"]["prefill_tokens"] == 3 + # Only the terminal event carries them. + assert sum(1 for _, data in collected if isinstance(data, dict) and "metrics" in data) == 1 + + off = _collect_events(events) + assert all("metrics" not in data for _, data in off if isinstance(data, dict)) diff --git a/tests/server/test_default_thinking_mode.py b/tests/server/test_default_thinking_mode.py new file mode 100644 index 0000000000..56854641f2 --- /dev/null +++ b/tests/server/test_default_thinking_mode.py @@ -0,0 +1,123 @@ +"""Unit tests for the --default-thinking-mode server flag merge logic.""" +from __future__ import annotations + +import asyncio +from types import SimpleNamespace + +from freetoken.server import anthropic_api as A +from freetoken.server.anthropic_models import AnthropicCountTokensRequest, AnthropicMessagesRequest +from freetoken.server.api_models import ChatCompletionRequest +from freetoken.server.model_meta import apply_default_thinking_mode +from freetoken.server.openai_api import chat_request_to_genspec +from freetoken.server.responses_api import ResponsesRequest, convert_responses_to_genspec + +OFF = {"enable_thinking": False, "thinking_mode": "disabled"} +ON = {"enable_thinking": True, "thinking_mode": "enabled"} + + +def test_auto_is_noop(): + assert apply_default_thinking_mode(None, "auto") is None + assert apply_default_thinking_mode({"enable_thinking": True}, "auto") == { + "enable_thinking": True + } + assert apply_default_thinking_mode(None, None) is None + + +def test_chat_fills_unset_kwargs(): + assert apply_default_thinking_mode(None, "chat") == OFF + assert apply_default_thinking_mode({}, "chat") == OFF + assert apply_default_thinking_mode({"other": 1}, "chat") == {"other": 1, **OFF} + + +def test_thinking_fills_unset_kwargs(): + assert apply_default_thinking_mode(None, "thinking") == ON + assert apply_default_thinking_mode({}, "thinking") == ON + + +def test_explicit_request_value_wins(): + # Explicit per-request values override the server default, both directions. + assert apply_default_thinking_mode({"enable_thinking": True}, "chat") == { + "enable_thinking": True + } + assert apply_default_thinking_mode({"enable_thinking": False}, "thinking") == { + "enable_thinking": False + } + assert apply_default_thinking_mode({"thinking": False}, "chat") == {"thinking": False} + assert apply_default_thinking_mode({"thinking_mode": "thinking"}, "chat") == { + "thinking_mode": "thinking" + } + # A raw graded-effort request asks for thinking too. + assert apply_default_thinking_mode({"reasoning_effort": "high"}, "chat") == { + "reasoning_effort": "high" + } + + +def test_original_dict_not_mutated(): + original = {"other": 1} + result = apply_default_thinking_mode(original, "chat") + assert result == {"other": 1, **OFF} + assert original == {"other": 1} + + +def test_unknown_mode_is_noop(): + assert apply_default_thinking_mode({"other": 1}, "bogus") == {"other": 1} + + +def _chat(**extra): + body = {"model": "m", "messages": [{"role": "user", "content": "hi"}], **extra} + return chat_request_to_genspec( + ChatCompletionRequest.model_validate(body), {}, default_thinking_mode="chat" + ).chat_template_kwargs + + +def test_chat_default_fills_a_request_without_thinking_controls(): + assert _chat() == OFF + + +def test_chat_default_keeps_protocol_level_thinking_requests(): + assert _chat(reasoning_effort="high")["enable_thinking"] is True + assert _chat(thinking={"type": "enabled"})["enable_thinking"] is True + assert _chat(reasoning_effort="none")["enable_thinking"] is False + + +def test_responses_chat_default_keeps_reasoning_effort(): + def ctk(**extra): + req = ResponsesRequest.model_validate({"model": "m", "input": "hi", **extra}) + return convert_responses_to_genspec(req, {}, default_thinking_mode="chat").chat_template_kwargs + + assert ctk() == OFF + on = ctk(reasoning={"effort": "medium"}) + assert on["enable_thinking"] is True and on["reasoning_effort"] == "medium" + + +def _anthropic(**extra): + return AnthropicMessagesRequest.model_validate( + {"model": "claude-x", "max_tokens": 64, "messages": [{"role": "user", "content": "hi"}], **extra} + ) + + +def test_anthropic_chat_default_applies_only_without_a_thinking_block(): + def ctk(**extra): + return A.convert_anthropic_to_genspec( + _anthropic(**extra), {}, default_thinking_mode="chat" + ).chat_template_kwargs + + assert ctk() == OFF + assert ctk(thinking={"type": "enabled", "budget_tokens": 1024}) == ON + assert ctk(thinking={"type": "adaptive"}) == {} + + +def test_anthropic_count_tokens_renders_the_generation_prompt_under_the_default(monkeypatch): + captured = {} + + async def fake_count(messages, template_tools, ctk, state, **_): + captured["ctk"] = ctk + return 1 + + monkeypatch.setattr(A, "count_prompt_tokens", fake_count) + state = SimpleNamespace(config=SimpleNamespace(reasoning_parser=None, default_thinking_mode="chat")) + body = {"model": "claude-x", "messages": [{"role": "user", "content": "hi"}]} + asyncio.run(A.handle_anthropic_count_tokens(AnthropicCountTokensRequest.model_validate(body), state)) + + generation = A.convert_anthropic_to_genspec(_anthropic(), {}, default_thinking_mode="chat") + assert captured["ctk"] == generation.chat_template_kwargs == OFF diff --git a/tests/server/test_dist_port_arg.py b/tests/server/test_dist_port_arg.py new file mode 100644 index 0000000000..a49def7b65 --- /dev/null +++ b/tests/server/test_dist_port_arg.py @@ -0,0 +1,41 @@ +"""Regression test for #301: the TP rendezvous port must be independently configurable +from --port, so two instances on adjacent ports don't have one's API port collide with +the other's rendezvous port. +""" + +from __future__ import annotations + +from unittest.mock import patch + +import pytest + +from freetoken.server.args import parse_args + +ANON_PATH = "/models/anon" + + +class _Config: + def to_dict(self) -> dict: + return {"architectures": ["LlamaForCausalLM"], "torch_dtype": "bfloat16"} + + +def _parse(argv: list[str]): + with patch("freetoken.utils.cached_load_hf_config", lambda _path: _Config()): + return parse_args(["--model", ANON_PATH, *argv]) + + +def test_dist_port_defaults_to_server_port_plus_one(): + args, _ = _parse(["--port", "8081"]) + assert args.distributed_port == 8082 + assert args.distributed_addr == "tcp://127.0.0.1:8082" + + +def test_dist_port_override_is_independent_of_server_port(): + args, _ = _parse(["--port", "8082", "--dist-port", "9000"]) + assert args.distributed_port == 9000 + assert args.distributed_addr == "tcp://127.0.0.1:9000" + + +def test_dist_port_rejects_out_of_range_value(): + with pytest.raises(SystemExit): + _parse(["--dist-port", "70000"]) diff --git a/tests/server/test_function_call_parser.py b/tests/server/test_function_call_parser.py index bc22296760..8af5bc88b6 100644 --- a/tests/server/test_function_call_parser.py +++ b/tests/server/test_function_call_parser.py @@ -371,3 +371,143 @@ def test_streaming_support_flags(): # test_streaming_model_matrix.py::test_non_streaming_detector_falls_back_to_buffered_parse). for name in SUPPORTED_TOOL_CALL_PARSERS: assert FunctionCallParser(TOOLS, tool_call_parser=name).supports_streaming() is True + + +# One class per concrete detector makes format coverage visible in pytest output. +# A shared contract avoids accidentally giving a newly added format weaker checks. +class _DetectorContract: + parser_name: str + block: str + truncate_before: str + buffered_recovery = False + + def parser(self): + return FunctionCallParser(OPENCODE_TOOLS, self.parser_name) + + def test_complete_call(self): + result = self.parser().parse_non_stream(self.block) + assert result.normal_text == "" + assert len(result.calls) == 1 + assert result.calls[0].name == "read" + assert json.loads(result.calls[0].parameters) == {"filePath": "/tmp/test_calc.py"} + + def test_every_character_boundary(self): + parser = self.parser() + texts, calls = _feed(parser, self.block) + assert "".join(texts) + parser.finish_stream() == "" + assert [call.name for call in calls if call.name] == ["read"] + assert json.loads("".join(call.parameters for call in calls)) == { + "filePath": "/tmp/test_calc.py" + } + # A completed call must never be duplicated by the end-of-stream recovery. + assert parser.recover_truncated_call() == [] + + def test_truncated_mid_call_recovery(self): + parser = self.parser() + cut = self.block.rindex(self.truncate_before) + texts, calls = _feed(parser, self.block[:cut]) + recovered = parser.recover_truncated_call() + assert bool(recovered) is self.buffered_recovery + # Incremental detectors have already consumed the opener and emitted the + # call identity. Their recovery method MUST NOT invent a duplicate call: + # generation.py recovers missing arguments from their partial-parse ledger. + # Buffered M3 instead recovers a complete call from the retained markup. + all_calls = calls + recovered + assert [call.name for call in all_calls if call.name] == ["read"] + arguments = "".join(call.parameters for call in all_calls) + try: + parsed = json.loads(arguments) + except json.JSONDecodeError: + parsed = json.loads(parser.unstreamed_arguments(0)) + assert parsed == {"filePath": "/tmp/test_calc.py"} + assert parser.recover_truncated_call() == [] + assert "".join(texts) + parser.finish_stream() == "" + + +class TestQwen25Detector(_DetectorContract): + parser_name = 'qwen25' + block = '{"name": "read", "arguments": {"filePath": "/tmp/test_calc.py"}}' + truncate_before = '' + + +class TestMistralDetector(_DetectorContract): + parser_name = 'mistral' + block = '[TOOL_CALLS] [{"name": "read", "arguments": {"filePath": "/tmp/test_calc.py"}}]' + truncate_before = '}]' + + +class TestLlama32Detector(_DetectorContract): + parser_name = 'llama3' + block = '<|python_tag|>{"name": "read", "arguments": {"filePath": "/tmp/test_calc.py"}}' + truncate_before = '}' + + +class TestGlm47Detector(_DetectorContract): + parser_name = 'glm47' + block = 'readfilePath/tmp/test_calc.py' + truncate_before = '' + + +class TestDeepSeekV32Detector(_DetectorContract): + parser_name = 'deepseekv32' + block = '<｜DSML｜function_calls><｜DSML｜invoke name="read"><｜DSML｜parameter name="filePath" string="true">/tmp/test_calc.py' + truncate_before = '' + + +class TestDeepSeekV41Detector(_DetectorContract): + parser_name = 'deepseekv41' + block = '<｜DSML｜ calls><｜DSML｜ invoke name="read"><｜DSML｜ parameter name="filePath" string="true">/tmp/test_calc.py' + truncate_before = '' + + +class TestQwen3CoderDetector(_DetectorContract): + parser_name = 'qwen3_coder' + block = '/tmp/test_calc.py' + truncate_before = '' + + +class TestGemma4Detector(_DetectorContract): + parser_name = 'gemma4' + block = '<|tool_call>call:read{filePath:<|"|>/tmp/test_calc.py<|"|>}' + truncate_before = '' + + +class TestMiniMaxDetector(_DetectorContract): + parser_name = 'minimax' + block = '/tmp/test_calc.py' + truncate_before = '' + + +class TestMiniMaxM3Detector(_DetectorContract): + parser_name = 'minimax_m3' + block = ']<]minimax[>[\n]<]minimax[>[]<]minimax[>[/tmp/test_calc.py]<]minimax[>[]<]minimax[>[\n]<]minimax[>[' + truncate_before = ']<]minimax[>[' + buffered_recovery = True + + +class TestGptOssDetector(_DetectorContract): + parser_name = 'gpt_oss' + block = '<|start|>assistant<|channel|>commentary to=functions.read <|constrain|>json<|message|>{"filePath": "/tmp/test_calc.py"}<|end|>' + truncate_before = '<|end|>' + + +class TestMuseGlimmerDetector(_DetectorContract): + parser_name = 'muse_glimmer' + block = '<|start|>assistant to=read<|message|>\n\n/tmp/test_calc.py\n\n<|eot|>' + truncate_before = '' + + +def test_contract_classes_cover_every_concrete_detector(): + from freetoken.server.function_call_parser import BaseFormatDetector + + covered = { + type(cls().parser().detector) + for cls in _DetectorContract.__subclasses__() + } + concrete, todo = set(), list(BaseFormatDetector.__subclasses__()) + while todo: # a detector may extend another (DeepSeekV41Detector < DeepSeekV32Detector) + cls = todo.pop() + concrete.add(cls) + todo.extend(cls.__subclasses__()) + assert covered == concrete + assert covered == set(FunctionCallParser.ToolCallParserEnum.values()) diff --git a/tests/server/test_function_call_parser_schema_resolution.py b/tests/server/test_function_call_parser_schema_resolution.py new file mode 100644 index 0000000000..468acc3bd2 --- /dev/null +++ b/tests/server/test_function_call_parser_schema_resolution.py @@ -0,0 +1,773 @@ +"""JSON Schema resolution in schema-aware tool argument parsing +(server/function_call_parser.py): $ref chains, unions and implicit containers, +plus MiniMax-M3's recursive nesting -- one-shot and streaming. +""" + +from __future__ import annotations + +import json + +import pytest + +from freetoken.server.function_call_parser import Function, FunctionCallParser, Tool + + +# Formats that route parameter values through _convert_param_value(); the rest +# parse arguments from their own JSON grammar and ignore the declared schema. +SCHEMA_AWARE_PARSERS = [ + "glm47", + "qwen3_coder", + "minimax", + "minimax_m3", + "muse_glimmer", +] + +MINIMAX_M3_NS = "]<]minimax[>[" + + +def _tool(prop_schema: dict, root_extra: dict | None = None, *, param_name: str = "limit") -> list[Tool]: + parameters = { + "type": "object", + "properties": { + param_name: prop_schema, + }, + } + + if root_extra: + parameters.update(root_extra) + + return [ + Tool( + function=Function( + name="schema_test", + parameters=parameters, + ) + ) + ] + + +def _wire(parser_name: str, raw: str, *, param_name: str = "limit") -> str: + if parser_name == "qwen3_coder": + return ( + "" + "" + f"{raw}" + "" + "" + ) + + if parser_name == "glm47": + return ( + "schema_test" + f"{param_name}" + f"{raw}" + "" + ) + + if parser_name == "minimax": + return ( + "" + '' + f'{raw}' + "" + "" + ) + + if parser_name == "minimax_m3": + ns = MINIMAX_M3_NS + return ( + f"{ns}\n" + f'{ns}' + f"{ns}<{param_name}>{raw}{ns}" + f"{ns}\n" + f"{ns}" + ) + + if parser_name == "muse_glimmer": + return ( + "<|start|>assistant to=schema_test<|message|>" + "\n" + '\n' + f'{raw}\n' + "\n" + "" + "<|eot|>" + ) + + raise AssertionError(f"Missing wire fixture for parser {parser_name!r}") + + +def _assemble_streamed_calls(calls) -> list[tuple[str, dict]]: + """Reassemble streamed calls the way a client does: a fragment naming a tool + opens a call, the argument fragments after it concatenate into its JSON.""" + assembled: list[list] = [] + + for call in calls: + if call.name is not None: + assembled.append([call.name, []]) + + if call.parameters: + assert assembled, f"argument fragment before any call name: {call!r}" + assembled[-1][1].append(call.parameters) + + return [(name, json.loads("".join(parts) or "{}")) for name, parts in assembled] + + +def _parse_args(parser_name: str, prop_schema: dict, root_extra: dict | None, raw: str, *, streaming: bool, param_name: str = "limit") -> dict: + tools = _tool(prop_schema, root_extra, param_name=param_name) + text = _wire(parser_name, raw, param_name=param_name) + + parser = FunctionCallParser(tools, tool_call_parser=parser_name) + + if not streaming: + result = parser.parse_non_stream(text) + + assert len(result.calls) == 1, result + assert result.calls[0].name == "schema_test" + + return json.loads(result.calls[0].parameters) + + calls = [] + + # Small chunks deliberately split XML/control markers and parameter values. + for i in range(0, len(text), 7): + _, emitted = parser.parse_stream_chunk(text[i : i + 7]) + calls.extend(emitted) + + # Drain any wire-order-deferred output. + for _ in range(4): + normal, emitted = parser.parse_stream_chunk("") + calls.extend(emitted) + + if not normal and not emitted: + break + + # End-of-stream hooks in the serving layer's order (generation.py): finalize + # closes a call cut off mid-arguments, finish releases held-back text. + calls.extend(parser.finalize_stream()) + parser.finish_stream() + + assembled = _assemble_streamed_calls(calls) + + assert len(assembled) == 1, (parser_name, calls) + assert assembled[0][0] == "schema_test" + + return assembled[0][1] + + +SCHEMA_CASES = [ + # Existing behaviour: direct concrete types. + pytest.param( + {"type": "integer"}, + None, + "5", + 5, + id="direct-integer", + ), + pytest.param( + {"type": "string"}, + None, + "5", + "5", + id="direct-string-preserved", + ), + + # Single local $ref. + pytest.param( + {"$ref": "#/$defs/Limit"}, + { + "$defs": { + "Limit": {"type": "integer"}, + } + }, + "5", + 5, + id="ref-integer", + ), + + # Chained $ref. + pytest.param( + {"$ref": "#/$defs/A"}, + { + "$defs": { + "A": {"$ref": "#/$defs/B"}, + "B": {"type": "integer"}, + } + }, + "5", + 5, + id="ref-chain", + ), + + # Draft-07 style definitions should work too: the resolver accepts generic + # local JSON pointers, not only $defs. + pytest.param( + {"$ref": "#/definitions/Limit"}, + { + "definitions": { + "Limit": {"type": "integer"}, + } + }, + "5", + 5, + id="legacy-definitions-ref", + ), + + # JSON Pointer escaping: ~1 => / and ~0 => ~. + pytest.param( + {"$ref": "#/$defs/A~1B~0C"}, + { + "$defs": { + "A/B~C": {"type": "integer"}, + } + }, + "5", + 5, + id="json-pointer-escaping", + ), + + # oneOf / anyOf concrete type. + pytest.param( + { + "oneOf": [ + {"type": "integer"}, + ] + }, + None, + "5", + 5, + id="oneof-integer", + ), + pytest.param( + { + "anyOf": [ + {"type": "integer"}, + ] + }, + None, + "5", + 5, + id="anyof-integer", + ), + + # Nullable unions. + pytest.param( + { + "oneOf": [ + {"type": "integer"}, + {"type": "null"}, + ] + }, + None, + "5", + 5, + id="oneof-nullable-integer", + ), + pytest.param( + { + "anyOf": [ + {"type": "integer"}, + {"type": "null"}, + ] + }, + None, + "5", + 5, + id="anyof-nullable-integer", + ), + pytest.param( + { + "type": ["integer", "null"], + }, + None, + "5", + 5, + id="type-array-nullable-integer", + ), + + # A union member may itself be a $ref. + pytest.param( + { + "oneOf": [ + {"$ref": "#/$defs/Limit"}, + {"type": "null"}, + ] + }, + { + "$defs": { + "Limit": {"type": "integer"}, + } + }, + "5", + 5, + id="oneof-ref-nullable", + ), + + # Referenced structured values must remain JSON structures rather than + # escaped strings. + pytest.param( + {"$ref": "#/$defs/Ids"}, + { + "$defs": { + "Ids": { + "type": "array", + "items": {"type": "integer"}, + } + } + }, + "[1, 2, 3]", + [1, 2, 3], + id="ref-array", + ), + pytest.param( + {"$ref": "#/$defs/Options"}, + { + "$defs": { + "Options": { + "type": "object", + "properties": { + "limit": {"type": "integer"}, + }, + } + } + }, + '{"limit": 5}', + {"limit": 5}, + id="ref-object", + ), + + # JSON Schema does not require an explicit "type" if structure already + # establishes the effective container kind. + pytest.param( + { + "properties": { + "limit": {"type": "integer"}, + } + }, + None, + '{"limit": 5}', + {"limit": 5}, + id="implicit-object", + ), + pytest.param( + { + "items": { + "type": "integer", + } + }, + None, + "[1, 2]", + [1, 2], + id="implicit-array", + ), + + # The PR intentionally uses loose JSON parsing when no single concrete type + # can be inferred. + pytest.param( + { + "oneOf": [ + {"type": "integer"}, + {"type": "string"}, + ] + }, + None, + "5", + 5, + id="ambiguous-union-loose-json", + ), + + # An enum/const with no "type" takes its type from the literals, so a + # numeric-looking string member is not coerced out of the enum. + pytest.param( + {"enum": ["1", "2"]}, + None, + "1", + "1", + id="untyped-string-enum", + ), + pytest.param( + {"enum": [1, 2]}, + None, + "2", + 2, + id="untyped-integer-enum", + ), + pytest.param( + {"const": "10"}, + None, + "10", + "10", + id="untyped-string-const", + ), + pytest.param( + { + "anyOf": [ + {"enum": ["1", "2"]}, + {"type": "null"}, + ] + }, + None, + "1", + "1", + id="anyof-untyped-string-enum-nullable", + ), + + # Pydantic v1 / OpenAPI 3.0 wrap a $ref in a one-element allOf to attach + # siblings such as default/description. + pytest.param( + {"allOf": [{"$ref": "#/$defs/Mode"}], "default": "a"}, + { + "$defs": { + "Mode": {"type": "string", "enum": ["a", "true"]}, + } + }, + "true", + "true", + id="allof-ref-string-enum", + ), + pytest.param( + {"allOf": [{"$ref": "#/$defs/Limit"}], "description": "max rows"}, + { + "$defs": { + "Limit": {"type": "integer"}, + } + }, + "5", + 5, + id="allof-ref-integer", + ), +] + + +@pytest.mark.parametrize("parser_name", SCHEMA_AWARE_PARSERS) +@pytest.mark.parametrize("streaming", [False, True], ids=["one-shot", "streaming"]) +@pytest.mark.parametrize( + ("prop_schema", "root_extra", "raw", "expected"), + SCHEMA_CASES, +) +def test_schema_aware_parsers_resolve_indirect_schema_types( + parser_name, + streaming, + prop_schema, + root_extra, + raw, + expected, +): + args = _parse_args( + parser_name, + prop_schema, + root_extra, + raw, + streaming=streaming, + ) + + assert args == {"limit": expected} + + +# --------------------------------------------------------------------------- +# Broken / recursive references must never take the request path down. +# --------------------------------------------------------------------------- +REF_EDGE_CASES = [ + pytest.param( + {"$ref": "#/$defs/Missing"}, + None, + id="dangling-local-ref", + ), + pytest.param( + {"$ref": "https://example.invalid/schema.json#/Limit"}, + None, + id="external-ref", + ), + pytest.param( + {"$ref": "#/$defs/A"}, + { + "$defs": { + "A": "not-a-schema", + } + }, + id="ref-to-non-object-def", + ), + pytest.param( + {"$ref": "#/$defs/A/name"}, + { + "$defs": { + "A": "not-a-container", + } + }, + id="ref-through-non-object-def", + ), + pytest.param( + {"$ref": "#/$defs/A"}, + { + "$defs": { + "A": {"$ref": "#/$defs/A"}, + } + }, + id="self-referential-ref", + ), + pytest.param( + {"$ref": "#/$defs/A"}, + { + "$defs": { + "A": {"$ref": "#/$defs/B"}, + "B": {"$ref": "#/$defs/A"}, + } + }, + id="mutually-recursive-ref", + ), + pytest.param( + {"allOf": [{"$ref": "#/$defs/A"}]}, + { + "$defs": { + "A": {"allOf": [{"$ref": "#/$defs/A"}]}, + } + }, + id="self-referential-allof", + ), +] + + +@pytest.mark.parametrize("parser_name", SCHEMA_AWARE_PARSERS) +@pytest.mark.parametrize("streaming", [False, True], ids=["one-shot", "streaming"]) +@pytest.mark.parametrize( + ("prop_schema", "root_extra"), + REF_EDGE_CASES, +) +def test_schema_ref_edge_cases_terminate_without_crashing(parser_name, streaming, prop_schema, root_extra): + args = _parse_args( + parser_name, + prop_schema, + root_extra, + "5", + streaming=streaming, + ) + + # An unresolvable ref leaves the parameter untyped: the same loose JSON + # fallback as an ambiguous union, so "5" -> 5 and never a dropped parameter. + assert args == {"limit": 5} + + +# --------------------------------------------------------------------------- +# Compatibility guards +# --------------------------------------------------------------------------- +@pytest.mark.parametrize("parser_name", SCHEMA_AWARE_PARSERS) +@pytest.mark.parametrize("streaming", [False, True], ids=["one-shot", "streaming"]) +def test_explicit_string_schema_never_loose_parses_json_literals(parser_name, streaming): + args = _parse_args( + parser_name, + {"type": "string"}, + None, + "123", + streaming=streaming, + ) + + assert args == {"limit": "123"} + + +@pytest.mark.parametrize("parser_name", SCHEMA_AWARE_PARSERS) +@pytest.mark.parametrize("streaming", [False, True], ids=["one-shot", "streaming"]) +def test_nullable_string_schema_preserves_numeric_looking_string(parser_name, streaming): + args = _parse_args( + parser_name, + { + "anyOf": [ + {"type": "string"}, + {"type": "null"}, + ] + }, + None, + "123", + streaming=streaming, + ) + + assert args == {"limit": "123"} + + +# --------------------------------------------------------------------------- +# MiniMax-M3 threads the schema through nested XML +# --------------------------------------------------------------------------- +def _m3_tools(properties: dict, **root_extra) -> list[Tool]: + """One ``schema_test`` tool: ``properties`` plus any extra root keywords.""" + return [ + Tool( + function=Function( + name="schema_test", + parameters={"type": "object", "properties": properties, **root_extra}, + ) + ) + ] + + +def _m3_call(body: str) -> str: + """Wrap an invoke body in MiniMax-M3's namespace-delimited wire form.""" + ns = MINIMAX_M3_NS + + return ( + f"{ns}\n" + f'{ns}' + f"{body}" + f"{ns}\n" + f"{ns}" + ) + + +_M3_NODE_DEFS = { + "$defs": { + "Node": { + "type": "object", + "properties": { + "code": {"type": "string"}, + "count": {"type": "integer"}, + "child": {"$ref": "#/$defs/Node"}, + }, + } + } +} + + +def _minimax_m3_recursive_tools() -> list[Tool]: + return _m3_tools({"node": {"$ref": "#/$defs/Node"}}, **_M3_NODE_DEFS) + + +def _minimax_m3_recursive_wire() -> str: + ns = MINIMAX_M3_NS + + return _m3_call( + f"{ns}" + f"{ns}1{ns}" + f"{ns}2{ns}" + f"{ns}" + f"{ns}5{ns}" + f"{ns}6{ns}" + f"{ns}" + f"{ns}" + ) + + +def _parse_minimax_m3(tools: list[Tool], text: str, *, streaming: bool) -> dict: + """Parse one MiniMax-M3 call through whichever path is under test.""" + parser = FunctionCallParser(tools, tool_call_parser="minimax_m3") + + if not streaming: + result = parser.parse_non_stream(text) + + assert len(result.calls) == 1, result + assert result.calls[0].name == "schema_test" + + return json.loads(result.calls[0].parameters) + + calls = [] + + # Small chunks deliberately split NS markers and element boundaries. M3 emits + # each call in one piece at its closing marker, so no extra drain is needed. + for i in range(0, len(text), 7): + _, emitted = parser.parse_stream_chunk(text[i : i + 7]) + calls.extend(emitted) + + calls.extend(parser.finalize_stream()) + parser.finish_stream() + + assembled = _assemble_streamed_calls(calls) + + assert len(assembled) == 1, (calls,) + assert assembled[0][0] == "schema_test" + + return assembled[0][1] + + +@pytest.mark.parametrize("streaming", [False, True], ids=["one-shot", "streaming"]) +def test_minimax_m3_nested_ref_keeps_nested_schema_typing(streaming): + args = _parse_minimax_m3( + _minimax_m3_recursive_tools(), + _minimax_m3_recursive_wire(), + streaming=streaming, + ) + + assert args == { + "node": { + "code": "1", + "count": 2, + "child": { + "code": "5", + "count": 6, + }, + } + } + + +# Below the top level: an array's "items" schema and an implicitly-array nested +# property both used to lose typing and arrive as the verbatim string. +MINIMAX_M3_NESTED_CASES = [ + pytest.param( + _m3_tools( + { + "ids": { + "type": "array", + "items": {"$ref": "#/$defs/Id"}, + } + }, + **{"$defs": {"Id": {"type": "integer"}}}, + ), + _m3_call( + f"{MINIMAX_M3_NS}" + f"{MINIMAX_M3_NS}1{MINIMAX_M3_NS}" + f"{MINIMAX_M3_NS}2{MINIMAX_M3_NS}" + f"{MINIMAX_M3_NS}" + ), + {"ids": [1, 2]}, + id="array-items-ref", + ), + pytest.param( + _m3_tools( + { + "payload": { + "type": "object", + "properties": { + # No explicit "type": only the structure says array. + "cell": {"items": {"type": "integer"}}, + }, + } + } + ), + _m3_call( + f"{MINIMAX_M3_NS}" + f"{MINIMAX_M3_NS}1{MINIMAX_M3_NS}" + f"{MINIMAX_M3_NS}2{MINIMAX_M3_NS}" + f"{MINIMAX_M3_NS}" + ), + {"payload": {"cell": [1, 2]}}, + id="repeated-key-implicit-array", + ), +] + + +@pytest.mark.parametrize("streaming", [False, True], ids=["one-shot", "streaming"]) +@pytest.mark.parametrize(("tools", "text", "expected"), MINIMAX_M3_NESTED_CASES) +def test_minimax_m3_nested_schema_typing_below_the_top_level(tools, text, expected, streaming): + assert _parse_minimax_m3(tools, text, streaming=streaming) == expected + + +# --------------------------------------------------------------------------- +# Frontier: only MiniMax-M3 recurses into nested markup +# --------------------------------------------------------------------------- +# The other formats read a parameter value as opaque text; pinning that makes +# teaching them to recurse a deliberate change. +FLAT_SCHEMA_AWARE_PARSERS = [name for name in SCHEMA_AWARE_PARSERS if name != "minimax_m3"] + + +@pytest.mark.parametrize("parser_name", FLAT_SCHEMA_AWARE_PARSERS) +@pytest.mark.parametrize("streaming", [False, True], ids=["one-shot", "streaming"]) +def test_nested_markup_in_parameter_value_stays_opaque_outside_minimax_m3(parser_name, streaming): + args = _parse_args( + parser_name, + { + "type": "object", + "properties": {"code": {"type": "string"}}, + }, + None, + "1", + streaming=streaming, + param_name="node", + ) + + assert args == {"node": "1"} \ No newline at end of file diff --git a/tests/server/test_generation_accounting.py b/tests/server/test_generation_accounting.py index a6d2bdb916..17b694a562 100644 --- a/tests/server/test_generation_accounting.py +++ b/tests/server/test_generation_accounting.py @@ -22,6 +22,8 @@ from freetoken.server.generation import ( # noqa: E402 GenDone, GenSpec, + GenTimings, + build_metrics, generate_events, generate_full, ) @@ -63,13 +65,22 @@ async def wait_for_ack(self, uid: int): yield reply -def _ack(prompt: int = 0, completion: int = 0, out: str = "", finished: bool = False) -> UserReply: +def _ack( + prompt: int = 0, + completion: int = 0, + out: str = "", + finished: bool = False, + cached: int = 0, + prefill_ms: float = 0.0, +) -> UserReply: return UserReply( uid=42, incremental_output=out, finished=finished, prompt_tokens_delta=prompt, completion_tokens_delta=completion, + cached_tokens=cached, + prefill_ms=prefill_ms, finish_reason="stop" if finished else None, ) @@ -206,3 +217,103 @@ def test_ttft_mean_is_zero_without_samples(): request_ring.reset() request_ring.record_request(_row(ttft_ms=None)) assert request_ring.requests_ttft_mean_ms() == 0 + + +# ------------------------------------------------------- per-request metrics +# The timings ride the same ack stream the token totals do, so they are covered here rather +# than in each adapter's file; the adapter tests only check that the object reaches the wire. +def test_prefill_ms_from_the_engine_reaches_the_result_unchanged(): + """The scheduler measures the prefill span; the generation layer must pass it through + rather than re-deriving it from ack arrival times (which batch together).""" + st = FakeState([ + _ack(prompt=100, cached=40), + _ack(completion=1, out="a", prefill_ms=12.5), + _ack(completion=1, out="b", finished=True), + ]) + result = asyncio.run(generate_full(42, _spec(), st)) + assert result.timings.prefill_ms == 12.5 + assert result.cached_tokens == 40 + + +# A real gap between acks, so a TTFT taken at the wrong ack is distinguishable from one taken +# at the right one; FakeState yields its whole list within a single event-loop tick. +_GAP_S = 0.02 +_GAP_MS = _GAP_S * 1000 + + +class PacedState(FakeState): + async def wait_for_ack(self, uid: int): + assert uid == 42 + for index, reply in enumerate(self._replies): + if index: + await asyncio.sleep(_GAP_S) + yield reply + + +def test_ttft_is_measured_at_the_first_generated_token_not_the_admission_ack(): + """The admission ack carries prompt_tokens and no token; timing TTFT from it would report + roughly the queue time for every request.""" + st = PacedState([ + _ack(prompt=5), + _ack(completion=1, out="hi"), + _ack(completion=1, out="!", finished=True), + ]) + result = asyncio.run(generate_full(42, _spec(), st)) + # One gap to the first token (not zero, as an admission-timed TTFT would give), two to the end. + assert _GAP_MS * 0.5 < result.timings.ttft_ms < _GAP_MS * 1.9 + assert result.timings.total_ms > result.timings.ttft_ms + _GAP_MS * 0.5 + assert result.timings.decode_ms > _GAP_MS * 0.5 + + +def test_ttft_counts_a_token_the_detokenizer_held_back(): + """A token whose text is withheld (a trailing partial-stop prefix) is still a generated + token: TTFT keys off completion_tokens_delta, not off non-empty output -- otherwise it + would skip to the ack after it.""" + st = PacedState([ + _ack(prompt=5), + _ack(completion=1, out=""), + _ack(completion=1, out="done", finished=True), + ]) + result = asyncio.run(generate_full(42, _spec(), st)) + assert _GAP_MS * 0.5 < result.timings.ttft_ms < _GAP_MS * 1.9 + + +def test_stream_gendone_carries_the_same_timings(): + st = FakeState([ + _ack(prompt=8, cached=2), + _ack(completion=1, out="x", prefill_ms=7.0), + _ack(completion=1, out="y", finished=True), + ]) + + async def drain(): + return [ev async for ev in generate_events(42, _spec(), st) if isinstance(ev, GenDone)] + + (done,) = asyncio.run(drain()) + assert done.timings.prefill_ms == 7.0 + assert done.timings.ttft_ms > 0.0 + + +def test_build_metrics_splits_prompt_tokens_into_cached_and_prefilled(): + metrics = build_metrics( + prompt_tokens=8192, completion_tokens=256, cached_tokens=6144, + timings=GenTimings(ttft_ms=320.5, prefill_ms=410.0, decode_ms=5120.0, total_ms=5480.2), + ) + # The identity the issue asks for: prompt_tokens == cached_prompt_tokens + prefill_tokens. + assert metrics["cached_prompt_tokens"] + metrics["prefill_tokens"] == 8192 + assert metrics["prefill_tokens"] == 2048 + assert metrics["prefill_tokens_per_second"] == round(2048 / 0.410, 2) + # Decode divides by completion_tokens - 1: the first token came out of prefill. + assert metrics["decode_tokens_per_second"] == round(255 / 5.120, 2) + assert metrics["decode_tokens"] == 256 + assert metrics["ttft_ms"] == 320.5 + assert metrics["total_time_ms"] == 5480.2 + + +def test_build_metrics_rates_are_zero_rather_than_dividing_by_zero(): + metrics = build_metrics( + prompt_tokens=10, completion_tokens=1, cached_tokens=0, timings=GenTimings() + ) + assert metrics["prefill_tokens_per_second"] == 0.0 + # One generated token spans no decode interval, so there is no rate to report. + assert metrics["decode_tokens_per_second"] == 0.0 + assert metrics["prefill_tokens"] == 10 diff --git a/tests/server/test_launch.py b/tests/server/test_launch.py new file mode 100644 index 0000000000..defba4f307 --- /dev/null +++ b/tests/server/test_launch.py @@ -0,0 +1,37 @@ +"""The scheduler worker's own process setup, before it builds the engine.""" + +from __future__ import annotations + +import queue +import threading + +import pytest +import torch +from tqdm import tqdm + +import freetoken.gpu_select +import freetoken.scheduler +from freetoken.distributed import DistributedInfo +from freetoken.server import launch +from freetoken.server.args import ServerArgs +from freetoken.utils import progress + + +def test_the_scheduler_worker_draws_its_bars_under_a_lock_of_its_own(monkeypatch): + """Only rank 0 draws a bar, so no bar lock spans processes. tqdm's default one is a multiprocessing + semaphore, which a worker ended by a signal leaves to resource_tracker as leaked.""" + + class EngineBuilt(Exception): + pass + + def scheduler(args): + raise EngineBuilt(tqdm.get_lock()) + + monkeypatch.setattr(freetoken.scheduler, "Scheduler", scheduler) + monkeypatch.setattr(freetoken.gpu_select, "set_assigned_gpu", lambda target: None) + monkeypatch.setattr(progress, "_PROGRESS_SINK", None) + monkeypatch.setattr(tqdm, "_lock", tqdm.get_lock()) + args = ServerArgs(model_path="unused", tp_info=DistributedInfo(0, 1), dtype=torch.bfloat16) + with pytest.raises(EngineBuilt) as built: + launch._run_scheduler(args, queue.Queue()) + assert type(built.value.args[0]) is type(threading.RLock()) diff --git a/tests/server/test_launch_worker_errors.py b/tests/server/test_launch_worker_errors.py new file mode 100644 index 0000000000..8779016bdd --- /dev/null +++ b/tests/server/test_launch_worker_errors.py @@ -0,0 +1,69 @@ +"""A tokenizer/detokenizer worker that dies during load must report WHY. + +Regression test for the asymmetry that made every tokenizer-side startup failure surface as +the generic "backend worker freetoken-detokenizer-0 exited during load": ``_run_scheduler`` +pushed an ("error", reason) ack before dying, ``_run_tokenize_worker`` did not. +""" + +from __future__ import annotations + +import queue + +import pytest + +from freetoken.server.launch import _run_tokenize_worker + + +def _drain(q: "queue.Queue") -> list: + out = [] + while True: + try: + out.append(q.get_nowait()) + except queue.Empty: + return out + + +def test_tokenize_worker_reports_the_real_reason_before_dying(monkeypatch): + """A raising tokenize_worker pushes ("error", reason) so the supervisor can surface it.""" + import freetoken.tokenizer as tok_mod + + def boom(**_kwargs): + raise FileNotFoundError("tokenizer.json is missing from the checkpoint") + + monkeypatch.setattr(tok_mod, "tokenize_worker", boom) + q: "queue.Queue" = queue.Queue() + + with pytest.raises(FileNotFoundError): + _run_tokenize_worker(detach=False, ack_queue=q, tokenizer_path="/nope") + + acks = _drain(q) + assert acks, "the dying worker pushed no ack at all" + kind, reason = acks[0] + assert kind == "error" + assert "FileNotFoundError" in reason + assert "tokenizer.json is missing" in reason + + +def test_tokenize_worker_without_ack_queue_still_propagates(monkeypatch): + """No ack_queue (direct call / test harness) must not turn the failure into an AttributeError.""" + import freetoken.tokenizer as tok_mod + + def boom(**_kwargs): + raise ValueError("bad tokenizer config") + + monkeypatch.setattr(tok_mod, "tokenize_worker", boom) + + with pytest.raises(ValueError, match="bad tokenizer config"): + _run_tokenize_worker(detach=False, tokenizer_path="/nope") + + +def test_tokenize_worker_success_pushes_no_error_ack(monkeypatch): + """The happy path must stay clean — no spurious ("error", …) on the ack queue.""" + import freetoken.tokenizer as tok_mod + + monkeypatch.setattr(tok_mod, "tokenize_worker", lambda **_k: None) + q: "queue.Queue" = queue.Queue() + + _run_tokenize_worker(detach=False, ack_queue=q, tokenizer_path="/nope") + + assert not [a for a in _drain(q) if isinstance(a, tuple) and a and a[0] == "error"] diff --git a/tests/server/test_logprobs_api.py b/tests/server/test_logprobs_api.py new file mode 100644 index 0000000000..8cf0387c60 --- /dev/null +++ b/tests/server/test_logprobs_api.py @@ -0,0 +1,308 @@ +from __future__ import annotations + +import asyncio +import json +from types import SimpleNamespace + +from freetoken.message import TokenizeMsg, UserReply +from freetoken.server.openai_api import ( + ChatCompletionRequest, + CompletionRequest, + chat_request_to_genspec, + handle_chat_completion, + handle_completion, + stream_chat_completion_chunks, + stream_completion_chunks, +) + + +def run(coro): + return asyncio.run(coro) + + +class FakeState: + def __init__(self, replies: list[UserReply], reasoning_parser: str | None = None) -> None: + self.config = SimpleNamespace( + model_path="/models/unit-model", + served_model_name="unit-model", + tool_call_parser="llama3", + reasoning_parser=reasoning_parser, + ) + self.replies = replies + self.sent: TokenizeMsg | None = None + + def new_user(self) -> int: + return 42 + + async def send_one(self, msg): + self.sent = msg + + async def wait_for_ack(self, uid: int): + assert uid == 42 + for reply in self.replies: + yield reply + + +def reply(text: str, *, finished: bool = False, logprobs: dict | None = None) -> UserReply: + return UserReply( + uid=42, + incremental_output=text, + finished=finished, + prompt_tokens_delta=3 if not text else 0, + completion_tokens_delta=1 if text else 0, + cached_tokens=0, + logprobs=logprobs, + ) + + +def entry(token_id: int, token: str, logprob: float) -> dict: + return { + "token_id": token_id, + "token": token, + "bytes": list(token.encode("utf-8")), + "logprob": logprob, + "top": [ + { + "token_id": token_id, + "token": token, + "bytes": list(token.encode("utf-8")), + "logprob": logprob, + }, + { + "token_id": token_id + 1, + "token": "x", + "bytes": [120], + "logprob": logprob - 1, + }, + ], + } + + +def chat_request(**kwargs) -> ChatCompletionRequest: + payload = { + "model": "client-model", + "messages": [{"role": "user", "content": "hello"}], + "max_tokens": 8, + } + payload.update(kwargs) + return ChatCompletionRequest(**payload) + + +def parse_sse(chunks: list[bytes]) -> list[dict | str]: + events: list[dict | str] = [] + for chunk in chunks: + for line in chunk.decode().splitlines(): + if line.startswith("data: "): + data = line.removeprefix("data: ") + events.append(data if data == "[DONE]" else json.loads(data)) + return events + + +def test_chat_non_stream_logprobs() -> None: + first = entry(1, "Hello", -0.1) + second = entry(2, "!", -0.2) + result = run( + handle_chat_completion( + chat_request(logprobs=True, top_logprobs=2), + None, + FakeState([reply("Hello", logprobs=first), reply("!", finished=True, logprobs=second)]), + {}, + ) + ) + + content = result["choices"][0]["logprobs"]["content"] + assert content == [ + { + "token": "Hello", + "logprob": -0.1, + "bytes": [72, 101, 108, 108, 111], + "top_logprobs": [ + {"token": "Hello", "logprob": -0.1, "bytes": [72, 101, 108, 108, 111]}, + {"token": "x", "logprob": -1.1, "bytes": [120]}, + ], + }, + { + "token": "!", + "logprob": -0.2, + "bytes": [33], + "top_logprobs": [ + {"token": "!", "logprob": -0.2, "bytes": [33]}, + {"token": "x", "logprob": -1.2, "bytes": [120]}, + ], + }, + ] + + without_logprobs = run( + handle_chat_completion( + chat_request(), + None, + FakeState([reply("Hello", finished=True, logprobs=first)]), + {}, + ) + ) + assert without_logprobs["choices"][0].get("logprobs") is None + + +def test_chat_stream_logprobs_follow_content_deltas() -> None: + first = entry(1, "Hello", -0.1) + state = FakeState( + [ + reply("Hello", logprobs=first), + reply(" world", finished=True), + ] + ) + req = chat_request(stream=True, logprobs=True, top_logprobs=2) + spec = chat_request_to_genspec(req, {}) + events = parse_sse(run(_collect(stream_chat_completion_chunks(42, req, state, spec)))) + + content_choices = [ + event["choices"][0] + for event in events + if isinstance(event, dict) + and event["choices"] + and event["choices"][0]["delta"].get("content") + ] + assert content_choices[0]["logprobs"]["content"][0]["token"] == "Hello" + assert "logprobs" not in content_choices[1] + + +def test_completion_logprobs_non_stream_and_stream() -> None: + first = entry(1, "Hi", -0.1) + second = entry(2, "!", -0.2) + req = CompletionRequest(model="client-model", prompt="hello", logprobs=2, max_tokens=8) + + result = run( + handle_completion( + req, + None, + FakeState([reply("Hi", logprobs=first), reply("!", finished=True, logprobs=second)]), + {}, + ) + ) + assert result["choices"][0]["logprobs"] == { + "tokens": ["Hi", "!"], + "token_logprobs": [-0.1, -0.2], + "top_logprobs": [ + {"Hi": -0.1, "x": -1.1}, + {"!": -0.2, "x": -1.2}, + ], + "text_offset": [0, 2], + } + + events = parse_sse( + run( + _collect( + stream_completion_chunks( + 42, + CompletionRequest( + model="client-model", + prompt="hello", + logprobs=2, + max_tokens=8, + stream=True, + ), + FakeState([reply("Hi", finished=True, logprobs=first)]), + ) + ) + ) + ) + chunk = next(event for event in events if isinstance(event, dict) and event["choices"][0]["text"]) + assert chunk["choices"][0]["logprobs"] == { + "tokens": ["Hi"], + "token_logprobs": [-0.1], + "top_logprobs": [{"Hi": -0.1, "x": -1.1}], + "text_offset": [0], + } + + +def test_logprobs_validation_errors() -> None: + chat_top_out_of_range = run( + handle_chat_completion(chat_request(logprobs=True, top_logprobs=25), None, FakeState([]), {}) + ) + assert chat_top_out_of_range.status_code == 400 + + chat_missing_flag = run( + handle_chat_completion(chat_request(top_logprobs=1), None, FakeState([]), {}) + ) + assert chat_missing_flag.status_code == 400 + + completion_out_of_range = run( + handle_completion( + CompletionRequest(model="client-model", prompt="hello", logprobs=7), + None, + FakeState([]), + {}, + ) + ) + assert completion_out_of_range.status_code == 400 + + completion_echo = run( + handle_completion( + CompletionRequest(model="client-model", prompt="hello", echo=True, logprobs=1), + None, + FakeState([]), + {}, + ) + ) + assert completion_echo.status_code == 400 + + +def test_chat_logprobs_fail_closed_under_semantic_parsing() -> None: + # A server-side reasoning parser hides tokens from message content, so + # logprob entries cannot be aligned with it: reject up front, stream and + # non-stream alike, before any engine work is submitted. + with_parser = FakeState([], reasoning_parser="qwen3") + for stream in (False, True): + resp = run( + handle_chat_completion(chat_request(logprobs=True, stream=stream), None, with_parser, {}) + ) + assert resp.status_code == 400 + assert json.loads(resp.body)["error"]["param"] == "logprobs" + + # Tool parsing consumes tokens into tool_calls -- same conflict. + tool = {"type": "function", "function": {"name": "f", "parameters": {"type": "object"}}} + tools_resp = run( + handle_chat_completion(chat_request(logprobs=True, tools=[tool]), None, FakeState([]), {}) + ) + assert tools_resp.status_code == 400 + assert json.loads(tools_resp.body)["error"]["param"] == "logprobs" + + # tool_choice="none" disables parsing, so logprobs stay available. + ok = run( + handle_chat_completion( + chat_request(logprobs=True, tools=[tool], tool_choice="none"), + None, + FakeState([reply("Hi", finished=True, logprobs=entry(1, "Hi", -0.1))]), + {}, + ) + ) + assert ok["choices"][0]["logprobs"]["content"][0]["token"] == "Hi" + + +def test_reasoning_logprob_is_not_carried_to_content_delta() -> None: + # The generation-layer guard behind the API fail-close: even when a caller + # bypasses request validation, a hidden reasoning token's entry is dropped, + # never attached to a later visible delta (where its token string would leak). + state = FakeState( + [ + reply("thought", logprobs=entry(1, "thought", -0.1)), + reply("answer", finished=True), + ], + reasoning_parser="qwen3", + ) + req = chat_request(stream=True, logprobs=True, top_logprobs=2) + spec = chat_request_to_genspec(req, {}) + events = parse_sse(run(_collect(stream_chat_completion_chunks(42, req, state, spec)))) + + content_choice = next( + event["choices"][0] + for event in events + if isinstance(event, dict) + and event["choices"] + and event["choices"][0]["delta"].get("content") == "answer" + ) + assert "logprobs" not in content_choice + + +async def _collect(iterator): + return [chunk async for chunk in iterator] diff --git a/tests/server/test_mm_input.py b/tests/server/test_mm_input.py index 19d9f86529..33561acdcd 100644 --- a/tests/server/test_mm_input.py +++ b/tests/server/test_mm_input.py @@ -10,6 +10,7 @@ import pytest +from freetoken.mm.media import _MAX_IMAGE_BYTES as _MAX from freetoken.mm.media import collect_image_refs, fetch_image_bytes, image_reject_reason from freetoken.server.generation import GenerationError, render_messages from freetoken.server.stats import derive_model_card @@ -140,3 +141,188 @@ def test_image_token_budget_flags_land_in_the_multimodal_config(): assert parse_args(["--model", "/models/anon"])[0].mm.processor_kwargs == {} with pytest.raises(SystemExit): # argparse reports the bad pair and exits parse_args(["--model", "/models/anon", "--image-min-tokens", "2048", "--image-max-tokens", "1024"]) + + +def test_render_messages_hoists_system_to_front(): + """render_messages moves system messages to index 0 for templates that + require it (e.g. Qwen3.6 'System message must be at the beginning').""" + msgs = render_messages([ + {"role": "user", "content": "What is 2+2?"}, + {"role": "system", "content": "You are a math tutor."}, + {"role": "user", "content": "Answer briefly."}, + ]) + assert msgs[0]["role"] == "system" + assert msgs[1]["role"] == "user" + assert msgs[2]["role"] == "user" + + +def test_render_messages_preserves_system_first(): + """When system is already first, no reordering happens.""" + msgs = render_messages([ + {"role": "system", "content": "You are helpful."}, + {"role": "user", "content": "Hi"}, + ]) + assert msgs[0]["role"] == "system" + assert msgs[1]["role"] == "user" + + +def test_render_messages_no_system_unchanged(): + """No system message → order preserved.""" + msgs = render_messages([ + {"role": "user", "content": "Hello"}, + {"role": "assistant", "content": "Hi there"}, + {"role": "user", "content": "Bye"}, + ]) + assert [m["role"] for m in msgs] == ["user", "assistant", "user"] + + +def test_render_messages_keeps_a_leading_developer_message_ahead_of_system(): + """[developer, system, user]: the tokenizer maps developer -> system and merges in order, so + render_messages must not hoist system ahead of it first. A template that knows the + developer role gets both messages as sent.""" + from freetoken.tokenizer.tokenize import _map_developer_role + + msgs = [ + {"role": "developer", "content": "D"}, + {"role": "system", "content": "S"}, + {"role": "user", "content": "hi"}, + ] + rendered = render_messages(msgs) + assert [m["role"] for m in rendered] == ["developer", "system", "user"] + mapped = _map_developer_role(rendered, "{% if message.role == 'system' %}") + assert [m["role"] for m in mapped] == ["system", "user"] and mapped[0]["content"] == "D\n\nS" + assert _map_developer_role(rendered, "{% if message.role == 'developer' %}") is rendered + + +def test_render_messages_rejects_an_image_in_a_merged_system_message(): + """Merging system messages joins text; an image part must not be pasted into the prompt as + the repr of its part list (with the whole data URL).""" + with pytest.raises(ValueError, match="System message cannot contain images"): + render_messages([ + {"role": "system", "content": "A"}, + {"role": "system", "content": [ + {"type": "text", "text": "B"}, + {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{PNG}"}}, + ]}, + {"role": "user", "content": "hi"}, + ]) +# ---- error bodies and redirect policy (no network: httpx.MockTransport stands in for the remote) ---- + + +def _mock_remote(monkeypatch, routes): + """Serve ``routes`` ({url: (status, headers, body) | Exception}) through the client media.py builds; + returns the list of URLs the transport actually saw.""" + import httpx + + seen = [] + + def handler(request): + seen.append(str(request.url)) + route = routes[str(request.url)] + if isinstance(route, Exception): + raise route + status, headers, body = route + return httpx.Response(status, headers=headers, content=body, request=request) + + real = httpx.AsyncClient + monkeypatch.setattr(httpx, "AsyncClient", lambda **kw: real(transport=httpx.MockTransport(handler), **kw)) + return seen + + +def _fetch(url, config): + return asyncio.run(fetch_image_bytes([{"kind": "url", "data": url}], config)) + + +def test_local_media_error_bodies_carry_no_path_root_or_errno(tmp_path): + root = tmp_path / "media" + root.mkdir() + config = _config(allowed_local_media_path=str(root)) + outside = tmp_path / "secret.png" + outside.write_bytes(b"x") + (root / "dir").mkdir() + + cases = { + f"file://{outside}": "subpath", # escapes the root + f"file://{root / 'missing.png'}": "missing or unreadable", # ENOENT + f"file://{root / 'dir'}": "missing or unreadable", # EISDIR + } + for url, expect in cases.items(): + with pytest.raises(ValueError, match=expect) as ei: + _fetch(url, config) + body = str(ei.value) + assert str(tmp_path) not in body and "Errno" not in body, body + assert "secret.png" not in body and "missing.png" not in body, body + + +def test_remote_fetch_failures_are_classified_without_library_text(monkeypatch): + import httpx + + big = _MAX + 1 + seen = _mock_remote(monkeypatch, { + "https://cdn.example.com/404.png": (404, {}, b""), + "https://cdn.example.com/down.png": httpx.ConnectError("[Errno -2] Name or service not known"), + "https://cdn.example.com/slow.png": httpx.ReadTimeout("timed out"), + "https://cdn.example.com/declared.png": (200, {"content-length": str(big)}, b""), + "https://cdn.example.com/streamed.png": (200, {}, b"\0" * big), + }) + config = _config(allowed_media_domains="cdn.example.com") + for name, expect in [("404", "HTTP 404"), ("down", "unreachable"), ("slow", "timed out"), + ("declared", "exceeds"), ("streamed", "exceeds")]: + with pytest.raises(ValueError, match=expect) as ei: + _fetch(f"https://cdn.example.com/{name}.png", config) + assert "Errno" not in str(ei.value) + assert len(seen) == 5 + + +def test_redirects_are_rechecked_against_the_allowlist(monkeypatch): + seen = _mock_remote(monkeypatch, { + "https://cdn.example.com/a.png": (302, {"location": "https://evil.com/a.png"}, b""), + "https://cdn.example.com/b.png": (302, {"location": "/real.png"}, b""), + "https://cdn.example.com/real.png": (200, {}, b"fakepng"), + "https://evil.com/a.png": (200, {}, b"pwned"), + }) + config = _config(allowed_media_domains="cdn.example.com") + with pytest.raises(ValueError, match="redirects outside the allowed domains") as ei: + _fetch("https://cdn.example.com/a.png", config) + assert "evil.com" not in str(ei.value) + # a same-host relative redirect is followed and served + assert _fetch("https://cdn.example.com/b.png", config) == [b"fakepng"] + # the disallowed hop was never requested + assert seen == ["https://cdn.example.com/a.png", "https://cdn.example.com/b.png", "https://cdn.example.com/real.png"] + + +@pytest.mark.parametrize("target", [ + "http://127.0.0.1/x", "http://10.0.0.1/x", "http://169.254.169.254/latest/meta-data", + "http://[::1]/x", "http://[::ffff:10.0.0.1]/x", "http://localhost/x", "http://0.0.0.0/x", + "http://100.64.0.1/x", +]) +def test_redirects_to_private_or_loopback_addresses_are_refused(monkeypatch, target): + seen = _mock_remote(monkeypatch, { + "https://public.example.com/a.png": (302, {"location": target}, b""), + target: (200, {}, b"internal"), + }) + # even with an empty allowlist (any domain admitted) a redirect may not reach inside + with pytest.raises(ValueError, match="private or loopback"): + _fetch("https://public.example.com/a.png", _config()) + assert seen == ["https://public.example.com/a.png"] + + +def test_redirect_hop_count_and_scheme_are_bounded(monkeypatch): + chain = {f"https://h.example.com/{i}.png": (302, {"location": f"https://h.example.com/{i + 1}.png"}, b"") + for i in range(7)} + chain["https://h.example.com/7.png"] = (200, {}, b"fakepng") + chain["https://h.example.com/file.png"] = (302, {"location": "file:///etc/passwd"}, b"") + seen = _mock_remote(monkeypatch, chain) + with pytest.raises(ValueError, match="redirects more than"): + _fetch("https://h.example.com/0.png", _config()) + assert len(seen) == 6 # the original plus five hops; the sixth hop is never requested + with pytest.raises(ValueError, match="non-http"): + _fetch("https://h.example.com/file.png", _config()) + # a short chain that stays public is fine + assert _fetch("https://h.example.com/5.png", _config()) == [b"fakepng"] + + +def test_bad_base64_is_named_without_decoder_text(): + for ref in ({"kind": "b64", "data": "not base64!!"}, {"kind": "url", "data": "data:image/png;base64"}): + with pytest.raises(ValueError, match="not valid base64"): + asyncio.run(fetch_image_bytes([ref], _config())) diff --git a/tests/server/test_model_path_resolution.py b/tests/server/test_model_path_resolution.py new file mode 100644 index 0000000000..76016e6869 --- /dev/null +++ b/tests/server/test_model_path_resolution.py @@ -0,0 +1,87 @@ +"""``--model org/name``: a hub repo id resolves to a local checkpoint directory. + +``ft serve --model deepseek-ai/DeepSeek-V4-Flash-0731`` used to pass the repo id through +untouched as a filesystem path. Config parsing then opened it relative to cwd: + + FileNotFoundError: No DeepSeek-V4 ModelArgs JSON found under + deepseek-ai/DeepSeek-V4-Flash-0731 (looked for inference/config.json) + +Only the modelscope branch ever downloaded, and ``download_hf_weight`` (used later, by the +weight loaders) fetches ``*.safetensors`` alone -- so it would never have produced +``inference/config.json`` either. + +These pin the resolve for both sources, and pin the two things that must keep reading the +*repo id* rather than the resolved snapshot path: the served model name and the parser +cascade. Resolving too early renames the served model to a snapshot hash. +""" + +from __future__ import annotations + +import sys +import types +from unittest.mock import patch + +import pytest + +from freetoken.server.args import parse_args + +REPO_ID = "deepseek-ai/DeepSeek-V4-Flash-0731" +SNAPSHOT = "/cache/hub/models--deepseek-ai--DeepSeek-V4-Flash-0731/snapshots/7872f01" + +# Explicit dtype and parsers keep parse_args off the network: all three otherwise read the +# checkpoint config, which is exactly what cannot be loaded until the path resolves. +BASE_ARGS = ["--dtype", "bfloat16", "--tool-call-parser", "llama3", "--reasoning-parser", "off"] + + +def _parse(model_path: str, *extra: str): + return parse_args([*BASE_ARGS, "--model", model_path, *extra])[0] + + +@pytest.fixture +def hf_download(): + with patch("freetoken.utils.hf.snapshot_download", return_value=SNAPSHOT) as m: + yield m + + +def test_repo_id_resolves_to_local_snapshot(hf_download): + assert _parse(REPO_ID).model_path == SNAPSHOT + assert hf_download.call_args.args[0] == REPO_ID + + +def test_served_model_name_keeps_the_repo_id(hf_download): + """Not the snapshot hash -- this is the id clients send in the `model` field.""" + assert _parse(REPO_ID).served_model_name == "DeepSeek-V4-Flash-0731" + + +def test_local_dir_is_never_downloaded(tmp_path, hf_download): + assert _parse(str(tmp_path)).model_path == str(tmp_path) + hf_download.assert_not_called() + + +def test_formats_freetoken_cannot_load_are_skipped(hf_download): + """Many repos ship a duplicate .bin copy; the loaders only read safetensors/gguf.""" + _parse(REPO_ID) + ignored = hf_download.call_args.kwargs["ignore_patterns"] + assert "*.bin" in ignored + assert "*.safetensors" not in ignored + + +def test_dummy_weight_skips_the_tensors(hf_download): + """--dummy-weight still needs config + tokenizer, but not 156 GiB of weights.""" + _parse(REPO_ID, "--dummy-weight") + ignored = hf_download.call_args.kwargs["ignore_patterns"] + assert "*.safetensors" in ignored + assert "*.gguf" in ignored + + +def test_modelscope_source_still_uses_modelscope(monkeypatch): + """The isdir check moved out of the modelscope branch; it must still short-circuit.""" + fake = types.ModuleType("modelscope") + calls = [] + fake.snapshot_download = lambda repo, **kw: calls.append((repo, kw)) or SNAPSHOT + monkeypatch.setitem(sys.modules, "modelscope", fake) + + with patch("freetoken.utils.hf.snapshot_download") as hf: + assert _parse(REPO_ID, "--model-source", "modelscope").model_path == SNAPSHOT + hf.assert_not_called() + assert calls[0][0] == REPO_ID diff --git a/tests/server/test_openai_api.py b/tests/server/test_openai_api.py index e33018e195..fdd2b11159 100644 --- a/tests/server/test_openai_api.py +++ b/tests/server/test_openai_api.py @@ -223,6 +223,26 @@ def test_glm_reasoning_parser_honors_disabled_thinking_with_tools(): assert parser is not None and parser.detector.force_reasoning is True +def test_qwen3_continuation_of_a_final_assistant_message_is_content_not_reasoning(): + # The Qwen template closes the think block before the assistant's content, so a + # continue_final_message continuation never emits ; starting the parser inside + # reasoning filed the whole continuation under reasoning_content. + from freetoken.server.generation import _make_reasoning_parser + + state = FakeState([], reasoning_parser="qwen3") + req = chat_request( + messages=[ + {"role": "user", "content": "Count to five."}, + {"role": "assistant", "content": "One, two,"}, + ], + continue_final_message=True, + ) + parser = _make_reasoning_parser(chat_request_to_genspec(req, {}), state) + assert parser is not None and parser.detector.force_reasoning is False + plain = _make_reasoning_parser(chat_request_to_genspec(chat_request(), {}), state) + assert plain.detector.force_reasoning is True + + def test_non_stream_chat_completion_returns_openai_tool_calls_and_sends_tools(): output = '[TOOL_CALLS] [{"name":"get_weather","arguments":{"city":"Paris"}}]' state = FakeState( @@ -442,19 +462,93 @@ def test_models_route_returns_served_model_name(): assert card["max_model_len"] is None and card["context_length"] is None +def _models_card(state): + app = FastAPI() + register_openai_routes(app, lambda: state, lambda: {}) + return TestClient(app).get("/v1/models").json()["data"][0] + + def test_models_route_publishes_the_model_context_length(): """`ft launch` reads this to size each agent's context window.""" state = FakeState([]) state.config.max_seq_len = 262144 - app = FastAPI() - register_openai_routes(app, lambda: state, lambda: {}) - card = TestClient(app).get("/v1/models").json()["data"][0] + card = _models_card(state) + # No pool geometry resolved yet (still loading, or pre-("meta", …) ack): the model + # ceiling is the only answer available. assert card["max_model_len"] == 262144 assert card["context_length"] == 262144 +def test_models_route_clamps_the_context_length_to_the_kv_pool(): + """The engine clamps its own max_seq_len to the allocated pool and the scheduler admits + against that, so publishing the unclamped ceiling here made `ft launch` size each agent's + compaction window past what the server can actually hold (#448).""" + state = FakeState([]) + state.config.max_seq_len = 262144 + state.config.page_size = 1 + state.cache_pools = {"num_pages": 178176, "page_size": 1} + + card = _models_card(state) + + assert card["max_model_len"] == 178176 + assert card["context_length"] == 178176 + + +def test_models_route_keeps_the_ceiling_when_the_pool_exceeds_it(): + """A pool larger than the model's positional ceiling does not extend the context.""" + state = FakeState([]) + state.config.max_seq_len = 32768 + state.config.page_size = 1 + state.cache_pools = {"num_pages": 178176, "page_size": 1} + + card = _models_card(state) + + assert card["max_model_len"] == 32768 + assert card["context_length"] == 32768 + + +def test_models_route_prefers_the_last_rebuild_over_the_load_time_pool(): + """A rebuild moves the pool; /v1/models must follow it rather than freeze the load-time + allocation (same most-recent-truth order /v1/cache/status reports).""" + state = FakeState([]) + state.config.max_seq_len = 262144 + state.config.page_size = 1 + state.cache_pools = {"num_pages": 178176, "page_size": 1} + state.last_rebuild = {"num_pages": 40000} + + card = _models_card(state) + + assert card["max_model_len"] == 40000 + + +def test_model_retrieve_route_returns_the_served_card_or_404(): + state = FakeState([]) + state.config.max_seq_len = 262144 + app = FastAPI() + register_openai_routes(app, lambda: state, lambda: {}) + client = TestClient(app) + + card = client.get("/v1/models/unit-model").json() + assert card["id"] == "unit-model" and card["context_length"] == 262144 + + missing = client.get("/v1/models/other-model") + assert missing.status_code == 404 + assert missing.json()["error"]["code"] == "model_not_found" + + +def test_models_route_with_a_trailing_slash_still_lists(): + state = FakeState([]) + app = FastAPI() + register_openai_routes(app, lambda: state, lambda: {}) + + response = TestClient(app).get("/v1/models/") + + assert response.status_code == 200 + assert response.json()["data"][0]["id"] == "unit-model" + + async def _collect(generator): return [chunk async for chunk in generator] @@ -687,3 +781,153 @@ def test_minimax_http_non_stream_forces_implicit_reasoning_without_request_knob( message = response["choices"][0]["message"] assert message["reasoning_content"] == "private thought" assert message["content"] == "visible answer" + + +# ----------------------------------------------------- per-request metrics +def _metrics_replies() -> list[UserReply]: + return [ + UserReply(uid=42, incremental_output="", finished=False, prompt_tokens_delta=10, cached_tokens=4), + UserReply(uid=42, incremental_output="hi", finished=True, completion_tokens_delta=2, prefill_ms=8.0), + ] + + +def test_non_stream_chat_metrics_only_with_flag(): + state = FakeState(_metrics_replies()) + state.config.enable_metrics_report = True + response = run(handle_chat_completion(chat_request(tools=None), request=None, state=state, model_sampling={})) + metrics = response["metrics"] + assert metrics["prefill_time_ms"] == 8.0 + assert metrics["prefill_tokens"] == 6 and metrics["cached_prompt_tokens"] == 4 + assert metrics["decode_tokens"] == 2 + assert metrics["total_time_ms"] > 0 + + off = run(handle_chat_completion(chat_request(tools=None), request=None, state=FakeState(_metrics_replies()), model_sampling={})) + assert "metrics" not in off + + +def test_non_stream_chat_metrics_report_the_cache_hit_without_enable_cache_report(): + """--enable-cache-report governs the billing fields in `usage`. Gating the hit inside + `metrics` on it too would leave prefill_tokens_per_second computed over tokens that were + never forwarded.""" + state = FakeState(_metrics_replies()) + state.config.enable_metrics_report = True + response = run(handle_chat_completion(chat_request(tools=None), request=None, state=state, model_sampling={})) + assert "prompt_tokens_details" not in response["usage"] + assert response["metrics"]["cached_prompt_tokens"] == 4 + + +def test_stream_chat_metrics_ride_the_usage_chunk(): + state = FakeState(_metrics_replies()) + state.config.enable_metrics_report = True + req = chat_request(tools=None, stream_options={"include_usage": True}) + + async def collect(): + return [chunk async for chunk in stream_chat_completion_chunks(42, req, state)] + + events = parse_sse(run(collect())) + final = next(e for e in reversed(events) if isinstance(e, dict) and e.get("usage")) + assert final["metrics"]["prefill_time_ms"] == 8.0 + assert final["metrics"]["decode_tokens"] == 2 + # No metrics on the content chunks -- exactly one carries them. + assert sum(1 for e in events if isinstance(e, dict) and "metrics" in e) == 1 + + +def test_stream_chat_without_include_usage_has_no_metrics_chunk(): + """Metrics ride the usage chunk, so a client that opted out of usage keeps the plain + OpenAI stream rather than getting an extra trailing chunk it never asked for.""" + state = FakeState(_metrics_replies()) + state.config.enable_metrics_report = True + + async def collect(): + return [chunk async for chunk in stream_chat_completion_chunks(42, chat_request(tools=None), state)] + + events = parse_sse(run(collect())) + assert not any(isinstance(e, dict) and "metrics" in e for e in events) + + +import pytest # noqa: E402 + + +@pytest.mark.parametrize( + ("field", "value"), + [ + ("temperature", -0.5), + ("top_p", 0.0), + ], +) +def test_chat_completion_rejects_invalid_sampling(field, value): + app = FastAPI() + state = FakeState([]) + + @app.post("/v1/chat/completions") + async def chat_completion(req: ChatCompletionRequest): + return await handle_chat_completion(req, request=None, state=state, model_sampling={}) + + body = { + "model": "unit-model", + "messages": [{"role": "user", "content": "hi"}], + "stream": True, + field: value, + } + response = TestClient(app).post("/v1/chat/completions", json=body) + + assert response.status_code == 400 + error = response.json()["error"] + assert error["type"] == "invalid_request_error" + assert field in error["message"] + assert state.sent is None + + +@pytest.mark.parametrize( + "sampling", + [ + {"temperature": 0.0, "top_p": 1.0, "top_k": -1}, + {"temperature": 0.0, "top_p": 1.0, "top_k": 1}, + ], +) +def test_chat_completion_accepts_sampling_boundaries(sampling): + app = FastAPI() + state = FakeState([]) + + @app.post("/v1/chat/completions") + async def chat_completion(req: ChatCompletionRequest): + return await handle_chat_completion(req, request=None, state=state, model_sampling={}) + + response = TestClient(app).post( + "/v1/chat/completions", + json={ + "model": "unit-model", + "messages": [{"role": "user", "content": "hi"}], + "stream": True, + **sampling, + }, + ) + + assert response.status_code == 200 + assert state.sent is not None + + +@pytest.mark.parametrize("top_k", [0, -5]) +def test_chat_completion_treats_non_positive_top_k_as_disabled(top_k): + """Ollama-style clients (Open WebUI) and some SDKs send top_k=0 to mean "off"; vLLM and + the other OpenAI-compatible servers accept any top_k <= 0 as disabled, so it must not 400.""" + app = FastAPI() + state = FakeState([]) + + @app.post("/v1/chat/completions") + async def chat_completion(req: ChatCompletionRequest): + return await handle_chat_completion(req, request=None, state=state, model_sampling={}) + + response = TestClient(app).post( + "/v1/chat/completions", + json={ + "model": "unit-model", + "messages": [{"role": "user", "content": "hi"}], + "stream": True, + "top_k": top_k, + }, + ) + + assert response.status_code == 200 + assert state.sent is not None + assert state.sent.sampling_params.top_k == -1 diff --git a/tests/server/test_openai_extras.py b/tests/server/test_openai_extras.py new file mode 100644 index 0000000000..19140d8b72 --- /dev/null +++ b/tests/server/test_openai_extras.py @@ -0,0 +1,389 @@ +"""OpenAI protocol coverage: n > 1 fan-out, echo and suffix (FIM) on completions, the +sampling extras reaching SamplingParams, request_id, usage details, and the tokenize / +detokenize routes.""" + +from __future__ import annotations + +import asyncio +import json +from types import SimpleNamespace + +import pydantic +import pytest +from freetoken.message import TokenizeMsg, UserReply +from freetoken.server.api_models import ( + ChatCompletionRequest, + CompletionRequest, + DetokenizeRequest, + TokenizeRequest, +) +from freetoken.server.openai_api import ( + _usage, + chat_request_to_genspec, + handle_chat_completion, + handle_completion, + handle_detokenize, + handle_tokenize, +) +from freetoken.server.stats import prometheus_text + + +def run(coro): + return asyncio.run(coro) + + +class FakeState: + """Distinct uids per generation; every uid gets the same two-ack reply.""" + + def __init__(self, text=("hel", "lo"), reasoning_tokens=0, tokenizer=None): + self.config = SimpleNamespace( + model_path="/models/unit-model", + served_model_name="unit-model", + tool_call_parser="llama3", + reasoning_parser=None, + max_seq_len=4096, + ) + self.text = text + self.reasoning_tokens = reasoning_tokens + self._next = 100 + self.sent: list[TokenizeMsg] = [] + self.aborted: list[int] = [] + if tokenizer is not None: + manager = SimpleNamespace( + tokenizer=tokenizer, + render_prompt=lambda msg: "RENDERED:" + json.dumps(msg.text), + ) + self.frontend_tokenizer = lambda: manager + + def new_user(self) -> int: + self._next += 1 + return self._next + + async def send_one(self, msg): + self.sent.append(msg) + + async def abort_user(self, uid): + self.aborted.append(uid) + + async def wait_for_ack(self, uid: int): + yield UserReply( + uid=uid, incremental_output="", finished=False, prompt_tokens_delta=5 + ) + for i, piece in enumerate(self.text): + last = i == len(self.text) - 1 + yield UserReply( + uid=uid, + incremental_output=piece, + finished=last, + completion_tokens_delta=1, + finish_reason="stop" if last else None, + reasoning_tokens=self.reasoning_tokens if last else 0, + ) + + +def chat(**kwargs) -> ChatCompletionRequest: + body = {"model": "unit-model", "messages": [{"role": "user", "content": "hi"}]} + body.update(kwargs) + return ChatCompletionRequest(**body) + + +def parse_sse(chunks: list[bytes]) -> list: + out = [] + for chunk in chunks: + for line in chunk.decode().split("\n"): + if not line.startswith("data: "): + continue + payload = line[len("data: ") :] + out.append("[DONE]" if payload == "[DONE]" else json.loads(payload)) + return out + + +async def _collect(generator): + return [item async for item in generator] + + +# ---- n > 1 ------------------------------------------------------------------------------ + + +def test_chat_n_returns_one_choice_per_generation_and_counts_the_prompt_once(): + state = FakeState(reasoning_tokens=2) + body = run(handle_chat_completion(chat(n=2), None, state, {})) + assert [c["index"] for c in body["choices"]] == [0, 1] + assert all(c["message"]["content"] == "hello" for c in body["choices"]) + assert len(state.sent) == 2 and len({m.uid for m in state.sent}) == 2 + usage = body["usage"] + assert ( + usage["prompt_tokens"] == 5 + and usage["completion_tokens"] == 4 + and usage["total_tokens"] == 9 + ) + assert usage["completion_tokens_details"] == {"reasoning_tokens": 4} + assert body["system_fingerprint"] is None + + +def test_chat_n_stream_interleaves_choice_indices_and_ends_once(): + state = FakeState() + response = run( + handle_chat_completion( + chat(n=2, stream=True, stream_options={"include_usage": True}), + None, + state, + {}, + ) + ) + events = parse_sse(run(_collect(response.body_iterator))) + assert events.count("[DONE]") == 1 and events[-1] == "[DONE]" + indices = {c["index"] for ev in events[:-1] for c in ev.get("choices", [])} + assert indices == {0, 1} + usage_chunks = [ev for ev in events[:-1] if ev.get("usage")] + assert len(usage_chunks) == 1 + assert usage_chunks[0]["usage"]["prompt_tokens"] == 5 + assert usage_chunks[0]["usage"]["completion_tokens"] == 4 + finals = [ + c for ev in events[:-1] for c in ev.get("choices", []) if c.get("finish_reason") + ] + assert len(finals) == 2 + + +def test_n_out_of_range_is_400(): + for n in (0, 17): + body = run(handle_chat_completion(chat(n=n), None, FakeState(), {})) + assert body.status_code == 400 + assert json.loads(body.body)["error"]["param"] == "n" + + +# ---- completions: echo, suffix, n --------------------------------------------------------- + + +def test_completion_echo_prepends_the_prompt(): + state = FakeState() + req = CompletionRequest(model="unit-model", prompt="Once upon", echo=True) + body = run(handle_completion(req, None, state, {})) + assert body["choices"][0]["text"] == "Once uponhello" + + +def test_completion_echo_streams_the_prompt_first(): + state = FakeState() + req = CompletionRequest( + model="unit-model", prompt="Once upon", echo=True, stream=True + ) + response = run(handle_completion(req, None, state, {})) + events = parse_sse(run(_collect(response.body_iterator))) + assert events[0]["choices"][0]["text"] == "Once upon" + assert events[-1] == "[DONE]" + + +class _FimTokenizer: + unk_token_id = -1 + + def convert_tokens_to_ids(self, name): + return {"<|fim_prefix|>": 10, "<|fim_suffix|>": 11, "<|fim_middle|>": 12}.get( + name, -1 + ) + + +class _NoFimTokenizer: + unk_token_id = 3 + + def convert_tokens_to_ids(self, name): + return 3 + + +def test_completion_suffix_builds_a_fill_in_the_middle_prompt(): + state = FakeState(tokenizer=_FimTokenizer()) + req = CompletionRequest(model="unit-model", prompt="def f(", suffix=" return 1") + run(handle_completion(req, None, state, {})) + assert ( + state.sent[0].text + == "<|fim_prefix|>def f(<|fim_suffix|> return 1<|fim_middle|>" + ) + + +def test_completion_suffix_needs_fim_tokens(): + state = FakeState(tokenizer=_NoFimTokenizer()) + req = CompletionRequest(model="unit-model", prompt="def f(", suffix="x") + body = run(handle_completion(req, None, state, {})) + assert ( + body.status_code == 400 and json.loads(body.body)["error"]["param"] == "suffix" + ) + + +def test_completion_n_orders_choices_prompt_major(): + state = FakeState() + req = CompletionRequest(model="unit-model", prompt=["a", "b"], n=2) + body = run(handle_completion(req, None, state, {})) + assert [c["index"] for c in body["choices"]] == [0, 1, 2, 3] + assert [m.text for m in state.sent] == ["a", "a", "b", "b"] + assert body["usage"]["prompt_tokens"] == 10 # two prompts, each counted once + assert body["usage"]["completion_tokens"] == 8 + + +# ---- sampling extras -------------------------------------------------------------------- + + +def test_sampling_extras_reach_sampling_params(): + spec = chat_request_to_genspec( + chat( + min_p=0.1, + repetition_penalty=1.2, + presence_penalty=0.5, + frequency_penalty=0.3, + logit_bias={"5": 150, "7": -3}, + min_tokens=3, + stop_token_ids=[9], + include_stop_str_in_output=True, + skip_special_tokens=True, + seed=7, + user="u1", + ), + {}, + ) + sp = spec.sampling_params + assert sp.min_p == 0.1 and sp.repetition_penalty == 1.2 + assert sp.presence_penalty == 0.5 and sp.frequency_penalty == 0.3 + assert sp.logit_bias == [[5, 100.0], [7, -3.0]] # clamped to the OpenAI range + assert sp.min_tokens == 3 and sp.stop_token_ids == [9] + assert sp.include_stop_str_in_output is True and sp.skip_special_tokens is True + assert sp.needs_logits_processing + + +@pytest.mark.parametrize( + "field", + [ + {"min_p": 1.5}, + {"repetition_penalty": 0.0}, + {"logit_bias": {"x": 1.0}}, + {"min_tokens": -1}, + {"stop_token_ids": [-1]}, + ], +) +def test_bad_sampling_extras_are_400(field): + body = run(handle_chat_completion(chat(**field), None, FakeState(), {})) + assert body.status_code == 400, field + + +@pytest.mark.parametrize("field", [{"presence_penalty": 3.0}, {"frequency_penalty": -2.5}]) +def test_out_of_range_penalties_fail_request_validation(field): + # main's request models bound both penalties to [-2, 2] (FastAPI answers 422) + with pytest.raises(pydantic.ValidationError): + chat(**field) + + +def test_continue_final_message_needs_an_assistant_tail_and_reaches_the_template(): + spec = chat_request_to_genspec( + chat( + messages=[ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "Sure,"}, + ], + continue_final_message=True, + ), + {}, + ) + assert spec.chat_template_kwargs["continue_final_message"] is True + with pytest.raises(ValueError): + chat(continue_final_message=True) + + +# ---- ids and usage ------------------------------------------------------------------------ + + +def test_request_id_becomes_the_response_id(): + state = FakeState() + body = run(handle_chat_completion(chat(request_id="abc-1"), None, state, {})) + assert body["id"] == "chatcmpl-abc-1" + response = run( + handle_chat_completion(chat(request_id="abc-2", stream=True), None, state, {}) + ) + events = parse_sse(run(_collect(response.body_iterator))) + assert events[0]["id"] == "chatcmpl-abc-2" + + +def test_usage_details_appear_only_when_nonzero(): + assert "completion_tokens_details" not in _usage(10, 5) + assert _usage(10, 5, reasoning_tokens=3)["completion_tokens_details"] == { + "reasoning_tokens": 3 + } + assert _usage(10, 5, cached_tokens=4)["prompt_tokens_details"] == { + "cached_tokens": 4 + } + + +# ---- tokenize / detokenize / metrics ----------------------------------------------------- + + +class _Tokenizer: + def encode(self, text, add_special_tokens=True): + ids = [ord(c) for c in text] + return ([1] + ids) if add_special_tokens else ids + + def decode(self, ids, skip_special_tokens=False): + return "".join(chr(i) for i in ids if not (skip_special_tokens and i == 1)) + + def convert_ids_to_tokens(self, ids): + return [chr(i) for i in ids] + + +def test_tokenize_prompt_and_messages_and_detokenize(): + state = FakeState(tokenizer=_Tokenizer()) + body = run(handle_tokenize(TokenizeRequest(prompt="ab"), state)) + assert body == {"count": 3, "max_model_len": 4096, "tokens": [1, 97, 98]} + body = run( + handle_tokenize( + TokenizeRequest( + prompt="ab", add_special_tokens=False, return_token_strs=True + ), + state, + ) + ) + assert body["tokens"] == [97, 98] and body["token_strs"] == ["a", "b"] + body = run( + handle_tokenize( + TokenizeRequest(messages=[{"role": "user", "content": "x"}]), state + ) + ) + assert body["count"] == len( + "RENDERED:" + json.dumps([{"role": "user", "content": "x"}]) + ) + bad = run( + handle_tokenize( + TokenizeRequest(prompt="a", messages=[{"role": "user", "content": "x"}]), + state, + ) + ) + assert bad.status_code == 400 + body = run( + handle_detokenize( + DetokenizeRequest(tokens=[1, 104, 105], skip_special_tokens=True), state + ) + ) + assert body == {"prompt": "hi"} + + +def test_metrics_exposition_lists_the_stats_document(): + doc = { + "instance_id": "abc", + "model": {"id": "unit-model"}, + "uptime_s": 12, + "kv": {"used_pages": 3, "total_pages": 10, "page_size": 64}, + "mamba": {"used_slots": 1, "total_slots": 4}, + "swa": None, + "vram_bytes": 123, + "throughput": {"decode_tps": 99.5, "prefill_tps": 1000.0}, + "requests": { + "active": 1, + "completed": 7, + "p95_ms": 50, + "ttft_mean_ms": 20, + "prompt_tokens_total": 100, + "completion_tokens_total": 40, + }, + } + text = prometheus_text(doc) + assert 'freetoken_info{instance_id="abc",model="unit-model"} 1' in text + assert "freetoken_requests_active 1" in text + assert "freetoken_requests_completed_total 7" in text + assert "freetoken_kv_pages_total 10" in text + assert "freetoken_mamba_slots_used 1" in text + assert "freetoken_decode_tokens_per_second 99.5" in text + assert "# TYPE freetoken_prompt_tokens_total counter" in text diff --git a/tests/server/test_protocol_compat.py b/tests/server/test_protocol_compat.py new file mode 100644 index 0000000000..d91582e423 --- /dev/null +++ b/tests/server/test_protocol_compat.py @@ -0,0 +1,213 @@ +"""Real SDK model/stream validation over the real HTTP adapters and GenSpec path. + +Only the scheduler/tokenizer transport is scripted. In-process ASGI transport +keeps these deterministic and CPU-only; it does not prove TCP chunk timing, +disconnect behavior, tokenizer templates, or a real model's output quality. +""" +import asyncio +import json +from types import SimpleNamespace + +import pytest +from fastapi import FastAPI + +from freetoken.message import TokenizeMsg, UserReply +from freetoken.server.anthropic_api import register_anthropic_routes +from freetoken.server.openai_api import register_openai_routes +from freetoken.server.responses_api import register_responses_routes + +# The SDKs are optional test dependencies: skip what needs a missing one instead of +# failing collection of the whole server suite. The Anthropic half skips per test. +httpx = pytest.importorskip("httpx") +openai = pytest.importorskip("openai") + + +SCHEMA = {"type": "object", "properties": {"city": {"type": "string"}}} +TOOL = {"type": "function", "function": {"name": "get_weather", "parameters": SCHEMA}} +CALL = '{"name":"get_weather","arguments":{"city":"Paris"}}' + + +class ScriptedBackend: + def __init__(self, text, *, reasoning=False): + self.config = SimpleNamespace( + model_path="/models/test", served_model_name="test", max_seq_len=4096, + tool_call_parser="qwen25", reasoning_parser="qwen3" if reasoning else None, + ) + self.text = text + self.sent = [] + self.maintenance_state = "serving" + + def new_user(self): + return 42 + + async def send_one(self, message): + assert isinstance(message, TokenizeMsg) + self.sent.append(message) + + async def wait_for_ack(self, uid): + # One character per backend reply exercises adapter stream aggregation + # as well as the SDK SSE decoder. Usage deltas remain authoritative. + for i, char in enumerate(self.text): + yield UserReply(uid=uid, incremental_output=char, + finished=i == len(self.text) - 1, + prompt_tokens_delta=5 if i == 0 else 0, + completion_tokens_delta=1) + + async def stream_with_cancellation(self, events, request, uid): + async for event in events: + yield event + + +def client(state, vendor="openai"): + app = FastAPI() + for register in (register_openai_routes, register_anthropic_routes, register_responses_routes): + register(app, lambda: state, dict) + # Recent Anthropic SDKs use httpx2 and reject httpx client instances. + # Select the transport matching its public default client class; older SDKs + # in the supported range still use httpx. Neither path bypasses SDK validation. + http = httpx + if vendor == "anthropic": + anthropic = pytest.importorskip("anthropic") + if not issubclass(anthropic.DefaultAsyncHttpxClient, httpx.AsyncClient): + http = pytest.importorskip("httpx2") + transport = http.AsyncClient(transport=http.ASGITransport(app=app)) + cls = openai.AsyncOpenAI if vendor == "openai" else anthropic.AsyncAnthropic + return cls(base_url="http://test/v1" if vendor == "openai" else "http://test", + api_key="test", http_client=transport, max_retries=0, + _strict_response_validation=True) + + +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize("tool", [False, True]) +def test_openai_chat_sdk_parses_reasoning_text_tools_and_usage(stream, tool): + async def run(): + state = ScriptedBackend("consider" + (CALL if tool else "Hello"), reasoning=True) + async with client(state) as sdk: + response = await sdk.chat.completions.create( + model="test", messages=[{"role": "user", "content": "Hi"}], + tools=[TOOL], stream=stream, stream_options={"include_usage": True}, + ) + if stream: + chunks = [chunk async for chunk in response] + choices = [choice for chunk in chunks for choice in chunk.choices] + assert "".join(getattr(c.delta, "reasoning_content", "") or "" for c in choices) == "consider" + assert [c.finish_reason for c in choices if c.finish_reason] == ["tool_calls" if tool else "stop"] + if tool: + calls = [tc for c in choices for tc in c.delta.tool_calls or []] + assert [tc.function.name for tc in calls if tc.function.name] == ["get_weather"] + assert json.loads("".join(tc.function.arguments or "" for tc in calls)) == {"city": "Paris"} + assert all(tc.index == 0 for tc in calls) + else: + assert "".join(c.delta.content or "" for c in choices) == "Hello" + usage = chunks[-1].usage + else: + assert response.choices[0].message.reasoning_content == "consider" + assert response.choices[0].finish_reason == ("tool_calls" if tool else "stop") + if tool: + call = response.choices[0].message.tool_calls[0] + assert call.id and call.function.name == "get_weather" + assert json.loads(call.function.arguments) == {"city": "Paris"} + else: + assert response.choices[0].message.content == "Hello" + usage = response.usage + assert usage.prompt_tokens == 5 + assert usage.completion_tokens == len(state.text) + assert usage.total_tokens == 5 + len(state.text) + assert len(state.sent) == 1 + asyncio.run(run()) + + +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize("tool", [False, True]) +def test_anthropic_sdk_assembles_typed_message(stream, tool): + pytest.importorskip("anthropic") + + async def run(): + state = ScriptedBackend("consider" + (CALL if tool else "Hello"), reasoning=True) + async with client(state, "anthropic") as sdk: + kwargs = dict(model="test", max_tokens=256, + messages=[{"role": "user", "content": "Hi"}], + tools=[{"name": "get_weather", "input_schema": SCHEMA}]) + if stream: + async with sdk.messages.stream(**kwargs) as response: + events = [event async for event in response] + result = await response.get_final_message() + assert events[0].type == "message_start" + assert events[-1].type == "message_stop" + else: + result = await sdk.messages.create(**kwargs) + assert result.role == "assistant" and result.model == "test" + assert result.stop_reason == ("tool_use" if tool else "end_turn") + assert result.content[0].type == "thinking" + assert result.content[0].thinking == "consider" + if tool: + # The established full-response shape includes an empty text + # block before tools; preserving that is part of compatibility. + assert all(b.text == "" for b in result.content if b.type == "text") + block = next(b for b in result.content if b.type == "tool_use") + assert block.id + assert block.name == "get_weather" and block.input == {"city": "Paris"} + else: + block = result.content[1] + assert block.type == "text" and block.text == "Hello" + assert result.usage.input_tokens == 5 + assert result.usage.output_tokens == len(state.text) + assert len(state.sent) == 1 + asyncio.run(run()) + + +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize("tool", [False, True]) +def test_responses_sdk_parses_output_and_events(stream, tool): + async def run(): + state = ScriptedBackend(CALL if tool else "Hello") + async with client(state) as sdk: + response = await sdk.responses.create( + model="test", input="Hi", stream=stream, + tools=[{"type": "function", "name": "get_weather", "parameters": SCHEMA}], + ) + if stream: + events = [event async for event in response] + assert events[0].type == "response.created" + assert events[-1].type == "response.completed" + assert [e.sequence_number for e in events] == list(range(len(events))) + result = events[-1].response + deltas = [e for e in events if e.type == ("response.function_call_arguments.delta" if tool else "response.output_text.delta")] + text = "".join(e.delta for e in deltas) + assert json.loads(text) == {"city": "Paris"} if tool else text == "Hello" + else: + result = response + assert result.status == "completed" + if tool: + call = next(item for item in result.output if item.type == "function_call") + assert call.call_id and call.name == "get_weather" + assert json.loads(call.arguments) == {"city": "Paris"} + else: + assert result.output_text == "Hello" + assert result.usage.input_tokens == 5 + assert result.usage.output_tokens == len(state.text) + assert result.usage.total_tokens == 5 + len(state.text) + assert len(state.sent) == 1 + asyncio.run(run()) + + +@pytest.mark.parametrize("protocol", ["chat", "responses", "anthropic"]) +def test_sdk_classifies_invalid_request_as_400(protocol): + vendor = "anthropic" if protocol == "anthropic" else "openai" + sdk_module = pytest.importorskip("anthropic") if vendor == "anthropic" else openai + + async def run(): + state = ScriptedBackend("unused") + error = sdk_module.BadRequestError + async with client(state, vendor) as sdk: + with pytest.raises(error) as caught: + if protocol == "chat": + await sdk.chat.completions.create(model="test", messages=[{"role": "user", "content": "Hi"}], n=0) + elif protocol == "responses": + await sdk.responses.create(model="test", input="Hi", background=True) + else: + await sdk.messages.create(model="test", max_tokens=0, messages=[{"role": "user", "content": "Hi"}]) + assert caught.value.status_code == 400 + assert "invalid_request_error" in str(caught.value) + assert state.sent == [] + asyncio.run(run()) diff --git a/tests/server/test_qwen_semantic_special_tokens.py b/tests/server/test_qwen_semantic_special_tokens.py new file mode 100644 index 0000000000..11e35ccb3e --- /dev/null +++ b/tests/server/test_qwen_semantic_special_tokens.py @@ -0,0 +1,405 @@ +"""Qwen protocol markers are private to semantic chat-style API responses.""" + +from __future__ import annotations + +import asyncio +import json +import os +import sys +from types import SimpleNamespace + +import pytest + +_ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +_PY = os.path.join(_ROOT, "python") +if _PY not in sys.path: + sys.path.insert(0, _PY) + +from freetoken.llm.llm import LLM, RequestStatus # noqa: E402 +from freetoken.message import DetokenizeMsg, UserReply # noqa: E402 +from freetoken.server.api_server import FrontendManager # noqa: E402 +from freetoken.server.api_models import CompletionRequest # noqa: E402 +from freetoken.server.generation import ( # noqa: E402 + _QWEN_SEMANTIC_HIDDEN_TOKENS, + ContentDelta, + GenDone, + GenSpec, + ReasoningDelta, + ToolCallsDelta, + generate_events, + generate_full, +) +from freetoken.server.openai_api import handle_completion, stream_completion_chunks # noqa: E402 + + +TOOLS = [ + { + "type": "function", + "function": { + "name": "echo", + "parameters": { + "type": "object", + "properties": {"value": {"type": "string"}}, + }, + }, + } +] + + +class FakeState: + def __init__( + self, + chunks: list[str], + *, + completion_tokens: list[int] | None = None, + reasoning_parser: str | None = "qwen3", + tool_call_parser: str = "qwen3_coder", + ) -> None: + self.config = SimpleNamespace( + reasoning_parser=reasoning_parser, + tool_call_parser=tool_call_parser, + served_model_name="qwen-unit", + ) + self._chunks = chunks + self._completion_tokens = completion_tokens or [1] * len(chunks) + self.sent = None + + def new_user(self) -> int: + return 42 + + async def send_one(self, msg) -> None: + self.sent = msg + + async def wait_for_ack(self, uid: int): + assert uid == 42 + for index, (chunk, token_delta) in enumerate(zip(self._chunks, self._completion_tokens)): + last = index == len(self._chunks) - 1 + yield UserReply( + uid=uid, + incremental_output=chunk, + finished=last, + finish_reason="stop" if last else None, + prompt_tokens_delta=5 if index == 0 else 0, + completion_tokens_delta=token_delta, + ) + + +def _spec(*, tools=None) -> GenSpec: + return GenSpec( + messages=[], + sampling_params=SimpleNamespace(), + chat_template_kwargs={"enable_thinking": False}, + template_tools=tools, + parser_tools=tools, + ) + + +async def _events(chunks: list[str], *, tools=None): + return [event async for event in generate_events(42, _spec(tools=tools), FakeState(chunks))] + + +def _visible(events) -> str: + return "".join(event.text for event in events if isinstance(event, ContentDelta)) + + +@pytest.mark.parametrize("marker", _QWEN_SEMANTIC_HIDDEN_TOKENS) +def test_every_hidden_marker_is_chunk_boundary_safe(marker: str): + text = f"before{marker}after" + events = asyncio.run(_events(list(text))) + + assert _visible(events) == "beforeafter" + assert isinstance(events[-1], GenDone) + assert events[-1].completion_tokens == len(text) + + +def test_stream_and_nonstream_match_while_preserving_semantic_markers_and_literals(): + grounding = ( + "<|object_ref_start|>cat<|object_ref_end|>" + "<|box_start|>(1,2),(3,4)<|box_end|>" + "<|quad_start|>(1,2),(3,4),(5,6),(7,8)<|quad_end|>" + ) + text = ( + "<|im_start|>plananswer" + "<|vision_start|><|image_pad|><|vision_end|>" + f"{grounding} literal \"<|im_start|>\" and `<|image_pad|>`<|im_end|>" + ) + streamed = asyncio.run(_events(list(text))) + full = asyncio.run(generate_full(42, _spec(), FakeState([text], completion_tokens=[len(text)]))) + + reasoning = "".join(event.text for event in streamed if isinstance(event, ReasoningDelta)) + expected_content = f"answer{grounding} literal \"<|im_start|>\" and `<|image_pad|>`" + assert reasoning == full.reasoning == "plan" + assert _visible(streamed) == full.content == expected_content + assert full.completion_tokens == streamed[-1].completion_tokens == len(text) + + +def test_tool_and_think_markers_reach_their_parsers_before_filtering(): + block = ( + "" + 'literal <|im_start|>' + "" + ) + text = "choose tool<|im_start|>" + block + "<|im_end|>" + streamed = asyncio.run(_events(list(text), tools=TOOLS)) + full = asyncio.run(generate_full(42, _spec(tools=TOOLS), FakeState([text]))) + + stream_calls = [ + call for event in streamed if isinstance(event, ToolCallsDelta) for call in event.calls + ] + assert "".join( + event.text for event in streamed if isinstance(event, ReasoningDelta) + ) == full.reasoning == "choose tool" + assert _visible(streamed) == full.content == "" + assert len(stream_calls) == len(full.tool_calls) == 1 + assert stream_calls[0].name == full.tool_calls[0].name == "echo" + assert json.loads(stream_calls[0].parameters) == json.loads(full.tool_calls[0].parameters) == { + "value": "literal <|im_start|>" + } + + +def test_filter_state_is_isolated_per_concurrent_request(): + async def run_pair(): + return await asyncio.gather( + _events(["left<|im_", "start|>right"]), + _events(['quoted "<|im_start|>"']), + ) + + first, second = asyncio.run(run_pair()) + assert _visible(first) == "leftright" + assert _visible(second) == 'quoted "<|im_start|>"' + + +@pytest.mark.parametrize("quote", ['"', "`", "'"]) +def test_unterminated_literal_cannot_disable_filtering(quote: str): + text = f"answer {quote}oops <|im_end|>" + streamed = asyncio.run(_events(list(text))) + full = asyncio.run(generate_full(42, _spec(), FakeState([text]))) + + assert _visible(streamed) == full.content == f"answer {quote}oops " + + +@pytest.mark.parametrize("slashes", ["\\", "\\\\"]) +def test_escaped_marker_opener_does_not_bypass_unterminated_quote_filter(slashes: str): + text = f'answer "oops {slashes}<|im_end|>' + expected = f'answer "oops {slashes}' + streamed = asyncio.run(_events(list(text))) + full = asyncio.run(generate_full(42, _spec(), FakeState([text]))) + + assert _visible(streamed) == full.content == expected + + +@pytest.mark.parametrize(("opening", "closing"), [("“", "”"), ("‘", "’")]) +def test_typographic_quoted_marker_is_preserved(opening: str, closing: str): + text = f"literal {opening}<|im_start|>{closing}" + streamed = asyncio.run(_events(list(text))) + full = asyncio.run(generate_full(42, _spec(), FakeState([text]))) + + assert _visible(streamed) == full.content == text + + +def test_suppressed_marker_does_not_turn_a_possessive_into_an_open_quote(): + text = "foo<|im_end|>'s <|vision_pad|>done" + streamed = asyncio.run(_events(list(text))) + full = asyncio.run(generate_full(42, _spec(), FakeState([text]))) + + assert _visible(streamed) == full.content == "foo's done" + + +@pytest.mark.parametrize( + ("text", "expected"), + [ + ("<<|im_end|>", "<"), + ("<|im_<|im_end|>", "<|im_"), + ("x<<|vision_pad|>y", "xok" + "" + ) + chunks = ["before<|im_", block, "start|>after"] + streamed = asyncio.run(_events(chunks, tools=TOOLS)) + full = asyncio.run(generate_full(42, _spec(tools=TOOLS), FakeState(["".join(chunks)]))) + + assert _visible(streamed) == full.content == "beforeafter" + assert len([event for event in streamed if isinstance(event, ToolCallsDelta)]) == 1 + assert len(full.tool_calls) == 1 + + +def test_disproved_marker_prefix_stays_before_intervening_tool_event(): + block = ( + "ok" + "" + ) + streamed = asyncio.run(_events(["pre<|im_", block, "Xpost"], tools=TOOLS)) + + kinds = [type(event).__name__ for event in streamed] + call_start_index = kinds.index("ToolCallStart") + call_end_index = kinds.index("ToolCallsDelta") + before = "".join( + event.text + for event in streamed[:call_start_index] + if isinstance(event, ContentDelta) + ) + after = "".join( + event.text + for event in streamed[call_end_index + 1 :] + if isinstance(event, ContentDelta) + ) + + assert before == "pre<|im_" + assert after == "Xpost" + + +def test_tool_event_flushes_prior_quoted_prose_before_call(): + block = ( + "ok" + "" + ) + chunks = [ + 'before "<|im_start|> quoted', + block, + ' tail <|vision_pad|>" after<|im_end|>', + ] + streamed = asyncio.run(_events(chunks, tools=TOOLS)) + full = asyncio.run(generate_full(42, _spec(tools=TOOLS), FakeState(["".join(chunks)]))) + + kinds = [type(event).__name__ for event in streamed] + call_start_index = kinds.index("ToolCallStart") + call_end_index = kinds.index("ToolCallsDelta") + prefix = "".join( + event.text + for event in streamed[:call_start_index] + if isinstance(event, ContentDelta) + ) + suffix = "".join( + event.text + for event in streamed[call_end_index + 1 :] + if isinstance(event, ContentDelta) + ) + assert prefix == 'before "<|im_start|> quoted' + assert suffix == ' tail <|vision_pad|>" after' + assert call_start_index < call_end_index + assert _visible(streamed) == full.content == ( + 'before "<|im_start|> quoted tail <|vision_pad|>" after' + ) + + +def test_reasoning_filter_hides_transport_marker_and_preserves_fim_tokens(): + text = ( + "plan<|vision_pad|>" + "<|fim_prefix|>left<|fim_suffix|>right<|fim_middle|><|im_end|>" + ) + streamed = asyncio.run(_events(list(text))) + full = asyncio.run(generate_full(42, _spec(), FakeState([text]))) + + reasoning = "".join( + event.text for event in streamed if isinstance(event, ReasoningDelta) + ) + expected_content = "<|fim_prefix|>left<|fim_suffix|>right<|fim_middle|>" + assert reasoning == full.reasoning == "plan" + assert _visible(streamed) == full.content == expected_content + + +def test_non_qwen_semantic_dialect_is_unchanged(): + text = '<|im_start|>literal "<|vision_pad|>"<|im_end|>' + state = FakeState( + [text], + reasoning_parser=None, + tool_call_parser="llama3", + ) + + result = asyncio.run(generate_full(42, _spec(), state)) + + assert result.content == text + + +def test_legacy_qwen_parser_alias_enables_semantic_filter_without_reasoning(): + text = "<|im_start|>visible<|vision_pad|><|im_end|>" + state = FakeState([text], reasoning_parser=None, tool_call_parser="qwen") + result = asyncio.run(generate_full(42, _spec(), state)) + + assert result.content == "visible" + + +def test_legacy_completions_paths_remain_raw_for_qwen(): + raw = '<|im_start|>raw "<|vision_pad|>"<|im_end|>' + + full = asyncio.run( + handle_completion( + CompletionRequest(model="qwen-unit", prompt="continue", max_tokens=8), + request=None, + state=FakeState([raw]), + model_sampling={}, + ) + ) + assert full["choices"][0]["text"] == raw + + async def stream_text() -> str: + chunks = [ + chunk async for chunk in stream_completion_chunks( + 42, + CompletionRequest(model="qwen-unit", prompt="continue", max_tokens=8, stream=True), + FakeState([raw]), + ) + ] + frames = [] + for chunk in chunks: + for line in chunk.decode().splitlines(): + if line.startswith("data: ") and line != "data: [DONE]": + frames.append(json.loads(line[6:])) + return "".join(frame["choices"][0]["text"] for frame in frames) + + assert asyncio.run(stream_text()) == raw + + +def test_raw_generate_stream_remains_unfiltered_for_qwen(): + raw = "<|im_start|>raw<|vision_pad|><|im_end|>" + + async def stream_text() -> str: + chunks = [ + chunk async for chunk in FrontendManager.stream_generate(FakeState([raw]), 42) + ] + frames = [] + for chunk in chunks: + for line in chunk.decode().splitlines(): + if line.startswith("data: ") and line != "data: [DONE]": + frames.append(json.loads(line[6:])) + return "".join(frame["text"] for frame in frames) + + assert asyncio.run(stream_text()) == raw + + +def test_offline_generation_keeps_exact_output_token_ids(): + offline = SimpleNamespace( + status_map={7: RequestStatus(uid=7, input_ids=[], output_ids=[])}, + eos_token_ids=set(), + ) + marker_token_ids = [248045, 248044, 248046] + + LLM.offline_send_result( + offline, + [ + DetokenizeMsg(uid=7, next_token=token_id, finished=False) + for token_id in marker_token_ids + ], + ) + + assert offline.status_map[7].output_ids == marker_token_ids diff --git a/tests/server/test_reasoning_parser.py b/tests/server/test_reasoning_parser.py new file mode 100644 index 0000000000..05c207d0e1 --- /dev/null +++ b/tests/server/test_reasoning_parser.py @@ -0,0 +1,108 @@ +"""Character-boundary contracts for every reasoning detector and template mode.""" +from types import SimpleNamespace + +import pytest + +from freetoken.server.generation import GenSpec, _make_reasoning_parser +from freetoken.server.reasoning_parser import BaseReasoningParser, ReasoningParser + + +class _ReasoningContract: + name: str + opener: str + closer: str + content_opener = "" + + def parser(self): + return ReasoningParser(self.name, force_reasoning=False) + + def stream(self, text): + parser = self.parser() + parts = [parser.parse_stream_chunk(char) for char in text] + parts.append(parser.flush()) + assert parser.flush() == ("", "") + return tuple("".join(pair[i] for pair in parts) for i in (0, 1)) + + def test_complete_reasoning(self): + text = self.opener + "consider" + self.closer + self.content_opener + "answer" + assert self.parser().parse_non_stream(text) == ("consider", "answer") + + def test_every_character_boundary(self): + text = self.opener + "consider" + self.closer + self.content_opener + "answer" + assert self.stream(text) == ("consider", "answer") + + def test_truncated_reasoning_flush(self): + # Reasoning detectors expose flush(), not recover_truncated_call(). EOS + # inside an open reasoning body must retain it in the reasoning channel. + assert self.stream(self.opener + "partial thought") == ("partial thought", "") + + +class TestDeepSeekV32ReasoningParser(_ReasoningContract): + name = "deepseekv32" + opener = "" + closer = "" + + +class TestThinkReasoningParser(_ReasoningContract): + name = "qwen3" + opener = "" + closer = "" + + +class TestMiniMaxM3ReasoningParser(_ReasoningContract): + name = "minimax_m3" + opener = "" + closer = "" + + +class TestGemmaThoughtReasoningParser(_ReasoningContract): + name = "gemma4" + opener = "<|channel>thought\n" + closer = "" + + +class TestGptOssHarmonyReasoningParser(_ReasoningContract): + name = "gpt_oss" + opener = "<|channel|>analysis<|message|>" + closer = "<|end|>" + content_opener = "<|start|>assistant<|channel|>final<|message|>" + + +class TestMuseGlimmerReasoningParser(_ReasoningContract): + name = "muse_glimmer" + opener = "<|start|>assistant to=self<|message|>" + closer = "<|eom|>" + content_opener = "<|start|>assistant to=user<|message|>" + + +def test_contract_classes_cover_every_reasoning_detector(): + covered = {type(cls().parser().detector) for cls in _ReasoningContract.__subclasses__()} + assert covered == set(BaseReasoningParser.__subclasses__()) + assert covered == set(ReasoningParser.ReasoningParserEnum.values()) + + +@pytest.mark.parametrize("tools", [None, [{"type": "function", "function": {"name": "read"}}]]) +@pytest.mark.parametrize(("name", "kwargs", "forced"), [ + ("qwen3", {}, True), + ("qwen3", {"enable_thinking": True}, True), + ("qwen3", {"enable_thinking": False}, False), + ("glm", {}, True), + ("glm", {"enable_thinking": True}, True), + ("glm", {"enable_thinking": False}, False), + ("gemma4", {}, False), + ("gemma4", {"thinking_mode": "thinking"}, True), + ("gemma4", {"enable_thinking": True}, True), + ("gemma4", {"thinking": True}, True), + ("gemma4", {"thinking_mode": "chat", "enable_thinking": False}, False), + ("minimax_m3", {}, False), + ("minimax_m3", {"thinking_mode": "adaptive"}, False), + ("minimax_m3", {"thinking_mode": "disabled"}, False), + ("minimax_m3", {"thinking_mode": "enabled"}, True), +]) +def test_generation_force_reasoning_matches_template(name, kwargs, forced, tools): + spec = GenSpec(messages=[], sampling_params=None, chat_template_kwargs=kwargs, + template_tools=tools, parser_tools=tools) + parser = _make_reasoning_parser(spec, SimpleNamespace(config=SimpleNamespace(reasoning_parser=name))) + # Observe routing, not an internal flag: tools must not turn a disabled + # template's ordinary response into invisible reasoning. + assert parser.parse_non_stream("probe") == (("probe", "") if forced else ("", "probe")) diff --git a/tests/server/test_rebuild_maintenance.py b/tests/server/test_rebuild_maintenance.py index a1d10a7e56..0df3d6bfef 100644 --- a/tests/server/test_rebuild_maintenance.py +++ b/tests/server/test_rebuild_maintenance.py @@ -327,3 +327,18 @@ def test_cache_rebuild_request_rejects_unknown_mode(): assert CacheRebuildRequest(mode="if_idle").mode == "if_idle" with pytest.raises(ValidationError): CacheRebuildRequest(mode="drain") + + +from pydantic import ValidationError # noqa: E402 + +from freetoken.server.api_server import CacheRebuildRequest # noqa: E402 + + +def test_rebuild_timeout_bounds(): + import pytest + + for bad in (0, -5, 3601, float("nan")): + with pytest.raises(ValidationError): + CacheRebuildRequest(timeout=bad) + assert CacheRebuildRequest(timeout=3600).timeout == 3600 + assert CacheRebuildRequest().timeout == 300.0 diff --git a/tests/server/test_request_recorder_disconnect.py b/tests/server/test_request_recorder_disconnect.py new file mode 100644 index 0000000000..06ff17f99b --- /dev/null +++ b/tests/server/test_request_recorder_disconnect.py @@ -0,0 +1,75 @@ +"""The request recorder must not hide a client disconnect from a non-streaming handler. + +Real sockets on purpose: the failure only shows through uvicorn's receive channel (a +BaseHTTPMiddleware wrapper never yields ``http.disconnect`` while the handler computes), which +TestClient does not model. +""" + +from __future__ import annotations + +import asyncio +import socket +import threading +import time + +from fastapi import FastAPI, Request + +from freetoken.server.api_server import _RecordRequestMiddleware + + +def _free_port() -> int: + with socket.socket() as s: + s.bind(("127.0.0.1", 0)) + return s.getsockname()[1] + + +def _serve(app: FastAPI, port: int): + import uvicorn + + server = uvicorn.Server(uvicorn.Config(app, host="127.0.0.1", port=port, log_level="error")) + thread = threading.Thread(target=server.run, daemon=True) + thread.start() + deadline = time.monotonic() + 10 + while not server.started: + assert time.monotonic() < deadline, "uvicorn did not start" + time.sleep(0.02) + return server + + +def _post_and_hang_up(port: int) -> None: + body = b'{"prompt": "x"}' + with socket.create_connection(("127.0.0.1", port)) as sock: + sock.sendall( + b"POST /v1/completions HTTP/1.1\r\nHost: t\r\nContent-Type: application/json\r\n" + b"Content-Length: " + str(len(body)).encode() + b"\r\n\r\n" + body + ) + time.sleep(0.2) # the handler is now polling + + +def test_non_streaming_handler_sees_the_client_leave(): + app = FastAPI() + app.add_middleware(_RecordRequestMiddleware) + seen: dict[str, float | None] = {} + + @app.post("/v1/completions") + async def slow(request: Request): + t0 = time.monotonic() + for _ in range(60): + await asyncio.sleep(0.05) + if await request.is_disconnected(): + seen["after"] = time.monotonic() - t0 + return {"detected": True} + seen["after"] = None + return {"detected": False} + + port = _free_port() + server = _serve(app, port) + try: + _post_and_hang_up(port) + deadline = time.monotonic() + 5 + while "after" not in seen and time.monotonic() < deadline: + time.sleep(0.05) + finally: + server.should_exit = True + assert seen.get("after") is not None, "handler never saw the disconnect" + assert seen["after"] < 2.0 diff --git a/tests/server/test_responses_api.py b/tests/server/test_responses_api.py index 887e78e573..6e7721b891 100644 --- a/tests/server/test_responses_api.py +++ b/tests/server/test_responses_api.py @@ -338,10 +338,11 @@ def test_stream_tool_call_events(): # Route smoke tests # --------------------------------------------------------------------------- # class FakeState: - def __init__(self, outputs, finish_reason=None, cached_tokens=0): + def __init__(self, outputs, finish_reason=None, cached_tokens=0, prefill_ms=0.0): self._outputs = outputs self._finish_reason = finish_reason # stamped on the terminal ack self._cached_tokens = cached_tokens # stamped on the first ack (admission reply) + self._prefill_ms = prefill_ms # the scheduler stamps it on the first generated token self.maintenance_state = "serving" self.config = SimpleNamespace( mm=SimpleNamespace(text_model_only=False, disabled_encoders=frozenset()), @@ -364,7 +365,8 @@ async def wait_for_ack(self, uid): yield UserReply(uid=uid, incremental_output=text, finished=finished, finish_reason=self._finish_reason if finished else None, prompt_tokens_delta=pt, completion_tokens_delta=ct, - cached_tokens=self._cached_tokens if i == 0 else 0) + cached_tokens=self._cached_tokens if i == 0 else 0, + prefill_ms=self._prefill_ms if i == 0 else 0.0) async def stream_with_cancellation(self, gen, request, uid): async for chunk in gen: @@ -967,3 +969,38 @@ def test_convert_function_call_output_text_list_stays_a_plain_tool_message(): spec = RP.convert_responses_to_genspec(req, {}) assert [m["role"] for m in spec.messages] == ["user", "assistant", "tool"] assert spec.messages[2]["content"] == "ab" + + +# --------------------------------------------------------------------------- # +# Per-request metrics (--enable-metrics-report) +# --------------------------------------------------------------------------- # +def test_route_nonstream_metrics_only_with_flag(): + fake = FakeState([("Hello world", True, 5, 2)], cached_tokens=3, prefill_ms=6.0) + fake.config.enable_metrics_report = True + body = _client(fake).post("/v1/responses", json={"model": "gpt-x", "input": "hi"}).json() + assert body["metrics"]["prefill_time_ms"] == 6.0 + assert body["metrics"]["prefill_tokens"] == 2 and body["metrics"]["cached_prompt_tokens"] == 3 + assert body["metrics"]["total_time_ms"] > 0 + + off = _client(FakeState([("Hello world", True, 5, 2)])).post( + "/v1/responses", json={"model": "gpt-x", "input": "hi"} + ).json() + assert "metrics" not in off + + +def test_route_stream_completed_event_carries_metrics(): + fake = FakeState([("Hello world", True, 5, 2)], cached_tokens=3, prefill_ms=6.0) + fake.config.enable_metrics_report = True + r = _client(fake).post( + "/v1/responses", json={"model": "gpt-x", "input": "hi", "stream": True} + ) + completed = None + for block in r.text.split("\n\n"): + for line in block.split("\n"): + if line.startswith("data:"): + payload = json.loads(line[len("data:"):].strip()) + if payload.get("type") == "response.completed": + completed = payload + assert completed is not None + assert completed["response"]["metrics"]["prefill_time_ms"] == 6.0 + assert completed["response"]["metrics"]["cached_prompt_tokens"] == 3 diff --git a/tests/server/test_sampling_validation.py b/tests/server/test_sampling_validation.py new file mode 100644 index 0000000000..b13db921be --- /dev/null +++ b/tests/server/test_sampling_validation.py @@ -0,0 +1,56 @@ +"""resolve_sampling rejects non-finite penalties: NaN passes a `<= 0` check and would +turn every already-seen token's logit into NaN.""" + +import pytest + +from freetoken.server.generation import resolve_sampling + + +def _resolve(**extra): + return resolve_sampling( + temperature=None, top_k=None, top_p=None, max_tokens=8, ignore_eos=False, + model_sampling={}, **extra, + ) + + +@pytest.mark.parametrize("value", [float("nan"), float("inf"), 0.0, -1.0]) +def test_repetition_penalty_must_be_positive_and_finite(value): + with pytest.raises(ValueError, match="repetition_penalty"): + _resolve(repetition_penalty=value) + + +def test_a_normal_repetition_penalty_is_kept(): + assert _resolve(repetition_penalty=1.1).repetition_penalty == pytest.approx(1.1) + + +def _resolve_with(**overrides): + kwargs = dict( + temperature=None, top_k=None, top_p=None, max_tokens=8, ignore_eos=False, + model_sampling={}, + ) + kwargs.update(overrides) + return resolve_sampling(**kwargs) + + +@pytest.mark.parametrize("value", [0, -1, -5]) +def test_a_non_positive_top_k_means_disabled(value): + # vLLM's rule: any top_k <= 0 is "no top-k filter"; clients send 0 for "off" + assert _resolve_with(top_k=value).top_k == -1 + + +def test_a_checkpoint_default_top_k_of_zero_means_disabled(): + assert _resolve_with(model_sampling={"top_k": 0}).top_k == -1 + + +def test_a_positive_top_k_is_kept(): + assert _resolve_with(top_k=20).top_k == 20 + + +@pytest.mark.parametrize( + ("field", "value"), + [("temperature", float("nan")), ("temperature", float("inf")), ("temperature", -0.5), + ("top_p", 0.0), ("top_p", 1.5), ("top_p", float("nan"))], +) +def test_the_other_sampling_ranges_are_still_rejected(field, value): + with pytest.raises(ValueError, match=field): + _resolve_with(**{field: value}) diff --git a/tests/server/test_stream_cancellation.py b/tests/server/test_stream_cancellation.py new file mode 100644 index 0000000000..71d5ac8b85 --- /dev/null +++ b/tests/server/test_stream_cancellation.py @@ -0,0 +1,376 @@ +"""Cancellation-path tests for FrontendManager.stream_with_cancellation / abort_user.""" + +from __future__ import annotations + +import asyncio +from types import SimpleNamespace + +import pytest + +from freetoken.message import AbortMsg +from freetoken.server.api_server import FrontendManager + + +class _Stats: + def __init__(self): + self.aborts = [] + + def on_abort(self, uid): + self.aborts.append(uid) + + +def _state(send_impl=None): + st = SimpleNamespace( + ack_map={7: [object()]}, + event_map={7: asyncio.Event()}, + aborted_uids={}, + stats=_Stats(), + ) + sent = [] + + async def default_send(msg): + sent.append(msg) + + st.send_one = send_impl or default_send + st.sent = sent + return st + + +class _Request: + def __init__(self, disconnected=False): + self._disconnected = disconnected + + async def is_disconnected(self): + return self._disconnected + + +async def _consume(state, request, uid=7): + state.abort_user = lambda request_uid: FrontendManager.abort_user(state, request_uid) + async for _ in FrontendManager.stream_with_cancellation(state, _never(), request, uid): + pass + + +async def _never(): + await asyncio.sleep(3600) + yield b"" # pragma: no cover + + +def test_cancellation_sends_one_abort_and_cleans_maps_inline(): + async def run(): + state = _state() + task = asyncio.create_task(_consume(state, _Request())) + await asyncio.sleep(0.01) + task.cancel() + + with pytest.raises(asyncio.CancelledError): + await task + + assert asyncio.all_tasks() - {asyncio.current_task()} == set() + assert len(state.sent) == 1 + assert isinstance(state.sent[0], AbortMsg) + assert state.sent[0].uid == 7 + assert state.ack_map == {} + assert state.event_map == {} + assert state.stats.aborts == [7] + + asyncio.run(run()) + + +def test_abort_delivery_failure_preserves_cancellation(): + async def boom(msg): + raise RuntimeError("zmq down") + + async def run(): + state = _state(send_impl=boom) + task = asyncio.create_task(_consume(state, _Request())) + await asyncio.sleep(0.01) + task.cancel() + + with pytest.raises(asyncio.CancelledError): + await task + + asyncio.run(run()) + + +def test_abort_user_is_idempotent(): + async def run(): + state = _state() + + await FrontendManager.abort_user(state, 7) + assert len(state.sent) == 1 + + await FrontendManager.abort_user(state, 7) + assert len(state.sent) == 1 + assert state.stats.aborts == [7] + + asyncio.run(run()) + + +def test_cancellation_after_ack_loop_cleanup_still_aborts(): + """The production disconnect path: the CancelledError lands on wait_for_ack's + innermost await, and its finally empties both maps BEFORE the stream wrapper + gets to run. A claim keyed on the maps no-ops exactly then; the AbortMsg must + still reach the scheduler (which acks aborts for uids it no longer has).""" + + async def run(): + state = _state() + + async def drain_like_wait_for_ack(): + try: + await asyncio.sleep(3600) + yield b"" # pragma: no cover + finally: + state.ack_map.pop(7, None) + state.event_map.pop(7, None) + + async def consume(): + state.abort_user = lambda uid: FrontendManager.abort_user(state, uid) + async for _ in FrontendManager.stream_with_cancellation( + state, drain_like_wait_for_ack(), _Request(), 7 + ): + pass + + task = asyncio.create_task(consume()) + await asyncio.sleep(0.01) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + assert state.ack_map == {} and state.event_map == {} + assert [type(m) for m in state.sent] == [AbortMsg] + assert state.sent[0].uid == 7 + assert state.stats.aborts == [7] + + asyncio.run(run()) + + +def test_late_close_delivers_abort(): + """A body closed after the response started surfaces as aclose() on the stream + wrapper: GeneratorExit, which `except CancelledError` never sees. The finally + must deliver the abort for any stream that did not run to completion.""" + + async def run(): + state = _state() + state.abort_user = lambda uid: FrontendManager.abort_user(state, uid) + + async def chunks(): + while True: + yield b"data: x\n\n" + + agen = FrontendManager.stream_with_cancellation(state, chunks(), _Request(), 7) + assert await agen.__anext__() == b"data: x\n\n" + await agen.aclose() + + assert [type(m) for m in state.sent] == [AbortMsg] + assert state.sent[0].uid == 7 + assert state.stats.aborts == [7] + + asyncio.run(run()) + + +def test_concurrent_aborts_send_one_message(): + """The accounting drain gathers abort_user over every in-flight uid and can race + a disconnect abort for the same request; the claim set must collapse them.""" + + async def run(): + state = _state() + await asyncio.gather( + FrontendManager.abort_user(state, 7), + FrontendManager.abort_user(state, 7), + ) + assert len(state.sent) == 1 + assert state.stats.aborts == [7] + + asyncio.run(run()) + + +def test_normal_completion_sends_no_abort(): + async def run(): + state = _state() + + async def one(): + yield b"data: x\n\n" + + async for _ in FrontendManager.stream_with_cancellation(state, one(), _Request(), 7): + pass + assert state.sent == [] + assert 7 in state.ack_map # wait_for_ack owns normal-path cleanup, not the stream + + asyncio.run(run()) + + +# --------------------------------------------------------------------------- # +# Non-streaming handlers: a disconnected client must abort the engine request +# (same shielded abort_user as the streaming path) instead of letting it run to +# max_tokens while later requests queue behind it. +# --------------------------------------------------------------------------- # + +from freetoken.server import openai_api +from freetoken.server.api_models import ChatCompletionRequest, CompletionRequest + + +def _ack(text, finished=False): + return SimpleNamespace( + error=None, + incremental_output=text, + finished=finished, + prompt_tokens_delta=0, + completion_tokens_delta=1, + cached_tokens=0, + finish_reason="stop" if finished else None, + matched_stop=None, + logprobs=None, + ) + + +class _ApiState: + """State fake for the openai_api handlers: hangs forever when given no acks.""" + + def __init__(self, acks=None): + self.config = SimpleNamespace( + model_path="/models/unit-model", + served_model_name="unit-model", + tool_call_parser="llama3", + reasoning_parser=None, + ) + self.acks = acks + self.sent = [] + self.aborted = [] + + def new_user(self): + return 7 + + async def send_one(self, msg): + self.sent.append(msg) + + async def wait_for_ack(self, uid): + if self.acks is None: + await asyncio.sleep(3600) + for ack in self.acks or []: + yield ack + + async def abort_user(self, uid): + self.aborted.append(uid) + + +def _chat_req(**kwargs): + payload = {"model": "m", "messages": [{"role": "user", "content": "hi"}], "max_tokens": 8} + payload.update(kwargs) + return ChatCompletionRequest(**payload) + + +def test_chat_non_stream_disconnect_delivers_abort(monkeypatch): + monkeypatch.setattr(openai_api, "_DISCONNECT_POLL_SECONDS", 0.01) + state = _ApiState(acks=None) # generation never finishes on its own + + resp = asyncio.run( + openai_api.handle_chat_completion(_chat_req(), _Request(disconnected=True), state, {}) + ) + + assert resp.status_code == 499 + assert state.aborted == [7] + + +def test_completion_non_stream_disconnect_delivers_abort(monkeypatch): + monkeypatch.setattr(openai_api, "_DISCONNECT_POLL_SECONDS", 0.01) + state = _ApiState(acks=None) + req = CompletionRequest(model="m", prompt="hello", max_tokens=8) + + resp = asyncio.run( + openai_api.handle_completion(req, _Request(disconnected=True), state, {}) + ) + + assert resp.status_code == 499 + assert state.aborted == [7] + + +def test_non_stream_connected_client_gets_result_without_abort(monkeypatch): + monkeypatch.setattr(openai_api, "_DISCONNECT_POLL_SECONDS", 0.01) + state = _ApiState(acks=[_ack("Hi", finished=True)]) + + result = asyncio.run( + openai_api.handle_chat_completion(_chat_req(), _Request(disconnected=False), state, {}) + ) + + assert result["choices"][0]["message"]["content"] == "Hi" + assert state.aborted == [] + + +def test_chat_non_stream_disconnect_aborts_every_sample(monkeypatch): + # n > 1 fans out into one engine request per sample; a disconnect must stop all of them + monkeypatch.setattr(openai_api, "_DISCONNECT_POLL_SECONDS", 0.01) + + class _State(_ApiState): + def __init__(self): + super().__init__(acks=None) + self._next = 7 + + def new_user(self): + self._next += 1 + return self._next + + state = _State() + resp = asyncio.run( + openai_api.handle_chat_completion(_chat_req(n=2), _Request(disconnected=True), state, {}) + ) + + assert resp.status_code == 499 + assert sorted(state.aborted) == [8, 9] + + +def test_a_failed_choice_cancels_the_other_choices(monkeypatch): + # n > 1: one choice errors while another is still generating. gather passes the error up + # without cancelling the sibling, and the aborts drop the event it waits on. + monkeypatch.setattr(openai_api, "_DISCONNECT_POLL_SECONDS", 0.01) + + class _State(_ApiState): + def __init__(self): + super().__init__(acks=None) + self._next = 7 + + def new_user(self): + self._next += 1 + return self._next + + async def wait_for_ack(self, uid): + if uid == 8: + yield SimpleNamespace(**{**vars(_ack("")), "error": "boom"}) + return + await asyncio.sleep(3600) + yield _ack("never") + + async def scenario(): + state = _State() + resp = await openai_api.handle_chat_completion( + _chat_req(n=2), _Request(disconnected=False), state, {} + ) + await asyncio.sleep(0) + pending = [t for t in asyncio.all_tasks() if t is not asyncio.current_task()] + return resp, state, pending + + resp, state, pending = asyncio.run(scenario()) + assert resp.status_code == 400 + assert sorted(state.aborted) == [8, 9] + assert pending == [] + + +def test_messages_and_responses_non_stream_disconnect_delivers_abort(monkeypatch): + # #222 only wrapped the OpenAI handlers; the Anthropic and Responses endpoints kept + # decoding an abandoned non-streaming request to max_tokens. + from freetoken.server import anthropic_api, responses_api + from freetoken.server.anthropic_models import AnthropicMessagesRequest + from freetoken.server.responses_api import ResponsesRequest + + monkeypatch.setattr(openai_api, "_DISCONNECT_POLL_SECONDS", 0.01) + + state = _ApiState(acks=None) + req = AnthropicMessagesRequest(model="m", max_tokens=8, messages=[{"role": "user", "content": "hi"}]) + handler = anthropic_api.handle_anthropic_messages(req, _Request(disconnected=True), state, {}) + resp = asyncio.run(asyncio.wait_for(handler, timeout=5)) # unwatched: hangs, then times out + assert resp.status_code == 499 and state.aborted == [7] + + state = _ApiState(acks=None) + req = ResponsesRequest(model="m", input="hi", max_output_tokens=8) + handler = responses_api.handle_responses(req, _Request(disconnected=True), state, {}) + resp = asyncio.run(asyncio.wait_for(handler, timeout=5)) + assert resp.status_code == 499 and state.aborted == [7] diff --git a/tests/server/test_stream_keepalive.py b/tests/server/test_stream_keepalive.py new file mode 100644 index 0000000000..e86d4365fc --- /dev/null +++ b/tests/server/test_stream_keepalive.py @@ -0,0 +1,115 @@ +"""A stream that stays silent -- the prefill of a long prompt -- keeps its client, and notices when +the client leaves instead of prefilling for nobody.""" + +from __future__ import annotations + +import asyncio +from types import SimpleNamespace + +import pytest + +from freetoken.server import api_server +from freetoken.server.api_server import FrontendManager + +# Converts a stream that never ends into a failure; the streams below finish in milliseconds. +HANG_GUARD_SECONDS = 10.0 + + +class _Request: + def __init__(self) -> None: + self.gone = False + + async def is_disconnected(self) -> bool: + return self.gone + + +def _manager() -> SimpleNamespace: + async def abort_user(uid: int) -> None: + pass + + return SimpleNamespace(abort_user=abort_user) + + +def _is_sse_comment(chunk: bytes) -> bool: + return chunk.startswith(b":") and chunk.endswith(b"\n\n") + + +def test_a_silent_stream_sends_comments_until_its_next_chunk(monkeypatch): + monkeypatch.setattr(api_server, "STREAM_KEEPALIVE_SECONDS", 0.01) + released = asyncio.Event() + + async def chunks(): + yield b"data: first\n\n" + await released.wait() + yield b"data: last\n\n" + + async def consume() -> list[bytes]: + out = [] + request = _Request() + async for chunk in FrontendManager.stream_with_cancellation(_manager(), chunks(), request, 7): + out.append(chunk) + if _is_sse_comment(chunk): + released.set() + return out + + out = asyncio.run(asyncio.wait_for(consume(), HANG_GUARD_SECONDS)) + assert out[0] == b"data: first\n\n" and out[-1] == b"data: last\n\n" + assert len(out) > 2 and all(_is_sse_comment(chunk) for chunk in out[1:-1]) + + +def test_a_client_that_leaves_during_the_silence_aborts_its_request(monkeypatch): + monkeypatch.setattr(api_server, "STREAM_KEEPALIVE_SECONDS", 0.01) + prefill_done = asyncio.Event() + + async def chunks(): + yield b"data: first\n\n" + await prefill_done.wait() + yield b"data: never read\n\n" + + async def scenario() -> tuple[list[bytes], list[int]]: + aborted = asyncio.Event() + seen: list[int] = [] + + async def abort_user(uid: int) -> None: + seen.append(uid) + aborted.set() + + request = _Request() + out: list[bytes] = [] + + async def consume() -> None: + manager = SimpleNamespace(abort_user=abort_user) + async for chunk in FrontendManager.stream_with_cancellation(manager, chunks(), request, 9): + out.append(chunk) + request.gone = True + + consumer = asyncio.ensure_future(consume()) + # Only the wrapper noticing the departure can end this wait; the guard below fails it. + await aborted.wait() + with pytest.raises(asyncio.CancelledError): + await consumer + return out, seen + + out, seen = asyncio.run(asyncio.wait_for(scenario(), HANG_GUARD_SECONDS)) + assert out == [b"data: first\n\n"] + assert seen == [9] + assert not prefill_done.is_set() + + +def test_a_stream_that_never_goes_quiet_carries_no_comments(monkeypatch): + monkeypatch.setattr(api_server, "STREAM_KEEPALIVE_SECONDS", HANG_GUARD_SECONDS) + + async def chunks(): + for index in range(5): + yield f"data: {index}\n\n".encode() + + async def consume() -> list[bytes]: + return [ + chunk + async for chunk in FrontendManager.stream_with_cancellation( + _manager(), chunks(), _Request(), 3 + ) + ] + + out = asyncio.run(asyncio.wait_for(consume(), HANG_GUARD_SECONDS)) + assert out == [f"data: {index}\n\n".encode() for index in range(5)] diff --git a/tests/server/test_supervisor.py b/tests/server/test_supervisor.py index 3568d90505..9e557e09b2 100644 --- a/tests/server/test_supervisor.py +++ b/tests/server/test_supervisor.py @@ -1,5 +1,7 @@ from __future__ import annotations +import pytest + import queue from queue import Empty as _Empty @@ -211,6 +213,37 @@ def on_ready() -> None: assert "failure" not in seen # graceful stop: the death was not reported +def test_supervisor_releases_the_ack_queue_once_the_backend_is_ready(): + """Readiness is the ack queue's last read. A multiprocessing queue's semaphores are unlinked only when + it is collected, and a main process ended by a signal never collects it.""" + import multiprocessing + import weakref + + from freetoken.server.supervisor import run_backend_supervisor + + class Proc: + name = "freetoken-TP0-scheduler" + alive = True + + def is_alive(self) -> bool: + return self.alive + + proc = Proc() + ack_queue = multiprocessing.get_context("spawn").Queue() + ack_queue.put("scheduler ready") + queue_at_ready = weakref.ref(ack_queue) + handle = BackendHandle(ack_queue=ack_queue, processes=[proc], expected_acks=1) + del ack_queue + seen: dict = {} + + def on_ready() -> None: + seen["queue at ready"] = queue_at_ready() + proc.alive = False + + run_backend_supervisor(handle, LoadProgress(), on_ready=on_ready, poll=0.01, is_shutting_down=lambda: True) + assert seen == {"queue at ready": None} + + def test_supervisor_silent_on_startup_death_during_shutdown(): """A worker dying mid-load while an orderly stop is already in progress must not be reported as a load failure either.""" @@ -270,3 +303,38 @@ def test_byte_bar_emits_to_installed_sink_then_stops_after_clear(): bar.update(100) bar.close() assert seen == [] + + +@pytest.mark.parametrize("shutdown", [False, True]) +def test_shutdown_flag_is_read_at_later_worker_death(monkeypatch, shutdown): + """Readiness must not snapshot shutdown state: the flag changes on a later poll. + + Keep TP0 alive and kill a different worker, so the first process cannot mask + a dead tokenizer. No real processes or sleep timing enter the assertion. + """ + import queue + from types import SimpleNamespace + from freetoken.server import supervisor + + alive = {"scheduler": True, "tokenizer": True} + processes = [SimpleNamespace(name=name, is_alive=lambda name=name: alive[name]) + for name in alive] + acks = queue.Queue() + acks.put("scheduler ready") + acks.put("tokenizer ready") + handle = BackendHandle(ack_queue=acks, processes=processes, expected_acks=2) + stopping = False + seen = [] + + def next_poll(_): + nonlocal stopping + stopping = shutdown + alive["tokenizer"] = False + + monkeypatch.setattr(supervisor.time, "sleep", next_poll) + supervisor.run_backend_supervisor( + handle, LoadProgress(), on_ready=lambda: seen.append("ready"), + on_failure=lambda message: seen.append(message), + is_shutting_down=lambda: stopping, + ) + assert seen == (["ready"] if shutdown else ["ready", "backend worker tokenizer exited"]) diff --git a/tests/tokenizer/test_detokenize_extras.py b/tests/tokenizer/test_detokenize_extras.py new file mode 100644 index 0000000000..4717617b9a --- /dev/null +++ b/tests/tokenizer/test_detokenize_extras.py @@ -0,0 +1,97 @@ +"""Per-request detokenizer options: include_stop_str_in_output, skip_special_tokens, and +the reasoning token count reported on the finished reply.""" + +from __future__ import annotations + +from freetoken.message import DetokenizeMsg +from freetoken.tokenizer.detokenize import DetokenizeManager + +EOS, SPECIAL, THINK_END = 0, 99, 7 + + +class _Tok: + """Token ids decode to letters; 99 is a special marker, 7 the reasoning end tag.""" + + eos_token_id = EOS + unk_token_id = -1 + + def decode(self, ids, skip_special_tokens=False): + out = [] + for t in ids: + if t == SPECIAL: + if not skip_special_tokens: + out.append("") + elif t == THINK_END: + out.append("") + else: + out.append(chr(ord("a") + t)) + return "".join(out) + + def batch_decode(self, ids_list, skip_special_tokens=False): + return [self.decode(ids, skip_special_tokens) for ids in ids_list] + + +def _run(manager, uid, tokens, *, finished_kwargs=None, **msg_kwargs): + text = "" + meta = 0 + for i, t in enumerate(tokens): + last = i == len(tokens) - 1 + kwargs = dict(msg_kwargs) + if last and finished_kwargs: + kwargs.update(finished_kwargs) + strs, counts = manager.detokenize_with_meta( + [DetokenizeMsg(uid=uid, next_token=t, finished=last, **kwargs)] + ) + text += strs[0] + meta = counts[0] + return text, meta + + +def test_stop_string_is_trimmed_by_default_and_kept_on_request(): + manager = DetokenizeManager(_Tok(), frozenset({EOS})) + trimmed, _ = _run( + manager, 1, [1, 2, 3], stop_strs=["cd"], finished_kwargs={"matched_stop": "cd"} + ) + assert trimmed == "b" + kept, _ = _run( + manager, + 2, + [1, 2, 3], + stop_strs=["cd"], + keep_stop_str=True, + finished_kwargs={"matched_stop": "cd"}, + ) + assert kept == "bcd" + + +def test_skip_special_tokens_is_per_request(): + manager = DetokenizeManager(_Tok(), frozenset({EOS})) + shown, _ = _run(manager, 1, [1, SPECIAL, 2, EOS]) + hidden, _ = _run(manager, 2, [1, SPECIAL, 2, EOS], skip_special_tokens=True) + assert shown == "bc" + assert hidden == "bc" + + +def test_reasoning_tokens_count_up_to_the_end_tag(): + manager = DetokenizeManager(_Tok(), frozenset({EOS}), think_end_id=THINK_END) + text, reasoning = _run(manager, 1, [1, 2, THINK_END, 3, 4, EOS]) + assert text == "bcde" + assert reasoning == 3 # b, c and the tag itself + _, none = _run(manager, 2, [1, 2, EOS]) + assert none == 0 + + +def test_a_batch_with_one_uid_twice_advances_its_state_in_order(): + # each message must see the decode state the previous one left, and the reasoning + # counts must stay aligned with the texts + manager = DetokenizeManager(_Tok(), frozenset({EOS}), think_end_id=THINK_END) + msgs = [ + DetokenizeMsg(uid=1, next_token=1, finished=False), + DetokenizeMsg(uid=1, next_token=THINK_END, finished=False), + DetokenizeMsg(uid=2, next_token=3, finished=False), + DetokenizeMsg(uid=1, next_token=2, finished=True), + ] + texts, counts = manager.detokenize_with_meta(msgs) + assert texts == ["b", "", "d", "c"] + assert counts == [0, 0, 0, 2] + assert manager.detokenize([DetokenizeMsg(uid=2, next_token=4, finished=True)]) == ["e"] diff --git a/tests/tokenizer/test_logprobs_entry.py b/tests/tokenizer/test_logprobs_entry.py new file mode 100644 index 0000000000..89cf35cb5e --- /dev/null +++ b/tests/tokenizer/test_logprobs_entry.py @@ -0,0 +1,59 @@ +from freetoken.tokenizer.detokenize import build_logprobs_entry + + +class _FakeTokenizer: + def __init__(self, token_table: dict[int, str]) -> None: + self.token_table = token_table + + def decode(self, token_ids: list[int]) -> str: + return "".join(self.token_table[token_id] for token_id in token_ids) + + +def test_build_logprobs_entry_matches_contract() -> None: + tokenizer = _FakeTokenizer({1: "hel", 2: "lo", 3: "é"}) + entry = build_logprobs_entry( + tokenizer, + token_id=3, + chosen_logprob=-0.123, + top_ids=[3, 2, 1], + top_logprobs=[-0.1, -0.9, -1.5], + ) + + assert entry["token_id"] == 3 + assert entry["token"] == "é" + assert entry["bytes"] == list("é".encode("utf-8")) + assert entry["logprob"] == -0.123 + + top = entry["top"] + assert len(top) == 3 + assert top[0] == { + "token_id": 3, + "token": "é", + "bytes": list("é".encode("utf-8")), + "logprob": -0.1, + } + assert top[1] == { + "token_id": 2, + "token": "lo", + "bytes": list("lo".encode("utf-8")), + "logprob": -0.9, + } + assert top[2] == { + "token_id": 1, + "token": "hel", + "bytes": list("hel".encode("utf-8")), + "logprob": -1.5, + } + + +def test_build_logprobs_entry_empty_top() -> None: + tokenizer = _FakeTokenizer({7: "x", 8: "y"}) + entry = build_logprobs_entry( + tokenizer, + token_id=7, + chosen_logprob=-2.0, + top_ids=[], + top_logprobs=[], + ) + + assert entry["top"] == [] diff --git a/tests/tokenizer/test_resolve_thinking_mode.py b/tests/tokenizer/test_resolve_thinking_mode.py index 266a745620..23d69cfee6 100644 --- a/tests/tokenizer/test_resolve_thinking_mode.py +++ b/tests/tokenizer/test_resolve_thinking_mode.py @@ -26,3 +26,13 @@ def test_explicit_chat_mode(): def test_invalid_mode_falls_back_to_chat(): assert resolve_thinking_mode({"thinking_mode": "bogus"}, None) == "chat" + + +def test_explicit_disable_overrides_tools(): + """An explicit enable_thinking=False (or thinking=False) must win even + when tools are present — the caller knows their intent better than the + tools-presumed-thinking heuristic.""" + assert resolve_thinking_mode({"enable_thinking": False}, [{"type": "function"}]) == "chat" + assert resolve_thinking_mode({"thinking": False}, [{"type": "function"}]) == "chat" + assert resolve_thinking_mode({"enable_thinking": False, "thinking": True}, [{"type": "function"}]) == "chat" + assert resolve_thinking_mode({"thinking": False, "enable_thinking": True}, [{"type": "function"}]) == "chat" diff --git a/tests/tokenizer/test_tokenize.py b/tests/tokenizer/test_tokenize.py index fb73fd317a..9e441886b2 100644 --- a/tests/tokenizer/test_tokenize.py +++ b/tests/tokenizer/test_tokenize.py @@ -816,3 +816,35 @@ def encode_messages(messages, thinking_mode, reasoning_effort=None): ) assert manager.render_prompt(msg) == f"Reasoning Effort: {budget}" assert msg.chat_template_kwargs["reasoning_effort"] == budget + + +def test_developer_role_maps_to_system_unless_the_template_knows_it(): + from freetoken.tokenizer.tokenize import _map_developer_role + + msgs = [{"role": "developer", "content": "be brief"}, {"role": "user", "content": "hi"}] + mapped = _map_developer_role(msgs, "{% if message.role == 'system' %}") + assert [m["role"] for m in mapped] == ["system", "user"] + assert mapped[0]["content"] == "be brief" and msgs[0]["role"] == "developer" # input untouched + assert _map_developer_role(msgs, "{% if message.role == 'developer' %}") is msgs + assert _map_developer_role(msgs[1:], None) == msgs[1:] # nothing to map + + +def test_mapped_developer_messages_are_merged_into_one_leading_system_message(): + """The frontend hoists system messages before the tokenizer maps developer -> system; a + template that rejects a non-first system message must still get exactly one, first.""" + from freetoken.tokenizer.tokenize import _map_developer_role + + template = "{% if message.role == 'system' %}" + msgs = [ + {"role": "system", "content": "be brief"}, + {"role": "developer", "content": "answer in German"}, + {"role": "user", "content": "hi"}, + ] + mapped = _map_developer_role(msgs, template) + assert [m["role"] for m in mapped] == ["system", "user"] + assert mapped[0]["content"] == "be brief\n\nanswer in German" + late = [{"role": "user", "content": "hi"}, {"role": "developer", "content": "be brief"}] + assert [m["role"] for m in _map_developer_role(late, template)] == ["system", "user"] + image = [{"type": "text", "text": "B"}, {"type": "image", "freetoken_ref": {"kind": "url", "data": "x"}}] + with pytest.raises(ValueError, match="System message cannot contain images"): + _map_developer_role([{"role": "developer", "content": "A"}, {"role": "system", "content": image}], template) diff --git a/tests/utils/__init__.py b/tests/utils/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/utils/test_mp_coalesced.py b/tests/utils/test_mp_coalesced.py new file mode 100644 index 0000000000..c64ee1c5c3 --- /dev/null +++ b/tests/utils/test_mp_coalesced.py @@ -0,0 +1,119 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Regression tests for coalesced msgpack frames on ZMQ queues (issue #452). + +A single ZMQ frame should carry exactly one msgpack object, but a read can surface +a frame containing two packed objects back to back. ``msgpack.unpackb`` raises +``ExtraData`` there and the worker process dies. The buffered unpacker decodes the +first object and holds the remainder for the next get(). +""" + +import asyncio + +import msgpack +import pytest + +from freetoken.utils.mp import ( + ZmqAsyncPullQueue, + ZmqAsyncPushQueue, + ZmqPubQueue, + ZmqPullQueue, + ZmqPushQueue, + ZmqSubQueue, +) + + +def _identity_decoder(obj): + return obj + + +def _addr(port): + return f"tcp://127.0.0.1:{port}" + + +class TestCoalescedFrames: + """A frame carrying two packed objects must yield two get() calls, not a crash.""" + + def test_sync_pull_coalesced_frame(self): + push = ZmqPushQueue(_addr(5601), create=True, encoder=lambda o: o) + pull = ZmqPullQueue(_addr(5601), create=False, decoder=_identity_decoder) + try: + # Simulate the coalesced frame: two packb outputs concatenated. + frame = msgpack.packb({"i": 1}, use_bin_type=True) + msgpack.packb({"i": 2}, use_bin_type=True) + pull._unpacker.feed(frame) + + assert pull.get() == {"i": 1} + assert pull.get() == {"i": 2} + assert len(pull._unpacker) == 0 + finally: + push.stop() + pull.stop() + + def test_decode_coalesced_raw(self): + from freetoken.utils.mp import _CoalescedUnpacker + + unpacker = _CoalescedUnpacker() + frame = msgpack.packb("first", use_bin_type=True) + msgpack.packb("second", use_bin_type=True) + unpacker.feed(frame) + assert unpacker.take() == "first" + assert unpacker.take() == "second" + + def test_empty_still_true_with_buffered_only(self): + push = ZmqPushQueue(_addr(5602), create=True, encoder=lambda o: o) + pull = ZmqPullQueue(_addr(5602), create=False, decoder=_identity_decoder) + try: + # nothing on socket and no buffer -> empty + assert pull.empty() + # a coalesced frame arrives and is fully buffered -> socket empty but buffer has 1 left + push.put({"x": 1}) + pull._unpacker.feed(pull.get_raw()) + assert pull.get() == {"x": 1} + assert pull.empty() + finally: + push.stop() + pull.stop() + + def test_sub_queue_coalesced(self): + import msgpack + + pub = ZmqPubQueue(_addr(5603), create=True, encoder=lambda o: o) + sub = ZmqSubQueue(_addr(5603), create=False, decoder=_identity_decoder) + try: + import time + + time.sleep(0.1) # SUB subscription propagation + pub.put({"a": 1}) + pub.put({"a": 2}) + time.sleep(0.1) + # Deliver both objects even if they arrive in one frame. + assert sub.get() == {"a": 1} + assert sub.get() == {"a": 2} + finally: + pub.stop() + sub.stop() + + +def test_async_pull_coalesced(): + import sys + + from freetoken.utils.mp import ZmqAsyncPullQueue as _AsyncPull + from freetoken.utils.mp import ZmqAsyncPushQueue as _AsyncPush + + # zmq.asyncio needs a Selector loop on Windows; the project runs on Linux CI. + if sys.platform == "win32": + pytest.skip("zmq.asyncio requires SelectorEventLoop on Windows (prod runs on Linux)") + async def scenario(): + push = ZmqAsyncPushQueue(_addr(5604), create=True, encoder=lambda o: o) + pull = ZmqAsyncPullQueue(_addr(5604), create=False, decoder=_identity_decoder) + try: + await push.put({"n": 1}) + await push.put({"n": 2}) + await asyncio.sleep(0.1) + + # Even if both messages arrive in a single frame, two gets must succeed. + assert await pull.get() == {"n": 1} + assert await pull.get() == {"n": 2} + finally: + push.stop() + pull.stop() + + asyncio.run(scenario()) diff --git a/tests/utils/test_mp_large_frame.py b/tests/utils/test_mp_large_frame.py new file mode 100644 index 0000000000..f91ee5604d --- /dev/null +++ b/tests/utils/test_mp_large_frame.py @@ -0,0 +1,13 @@ +"""The coalesced-frame decoder (#466) must accept a message larger than msgpack's +Unpacker default buffer (100 MiB), which unpackb never limited.""" + +import msgpack + +from freetoken.utils.mp import _CoalescedUnpacker + + +def test_a_message_over_100_mib_decodes(): + payload = {"pixels": b"x" * (101 * 1024 * 1024)} + unpacker = _CoalescedUnpacker() + unpacker.feed(msgpack.packb(payload)) + assert len(unpacker) == 1 and len(unpacker.take()["pixels"]) == len(payload["pixels"])