diff --git a/.github/issue-title-labeler.yml b/.github/issue-title-labeler.yml new file mode 100644 index 0000000000..48826301d5 --- /dev/null +++ b/.github/issue-title-labeler.yml @@ -0,0 +1,2 @@ +amd: + - '/\b(ROCm|HIP|RCCL|RDNA\s?\d?|CDNA\s?\d?|gfx\d{3,4}[a-z]?|MI\d{3}X?|RX\s?\d{4}|Radeon(?!\s*(\d{3}M|Graphics|Vega))|AMD[\s-]*(GPUs?|graphics|cards?|Instinct))\b|\bAMD[\s-]*显卡|\(amd\)|(? 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 fb5e613f7b..ecb53e1339 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -41,6 +41,7 @@ parsers all resolve automatically from the checkpoint and the GPU. |---|---|---| | `--host` | 127.0.0.1 | Bind address | | `--port` | 1919 | Bind port | +| `--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 | @@ -72,15 +73,39 @@ 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. | Flag | Default | Meaning | |---|---|---| -| `--moe-strategy` | auto | `fused`/`offload`/`cpu`/`hybrid`; auto → offload, or hybrid with a `ft bench bw` profile. `--moe-backend` is the deprecated old spelling | +| `--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) | @@ -99,6 +124,63 @@ See [models.md](models.md#moe-strategies) for what each strategy does. | `--tool-call-parser` | auto | Tool-call format; auto-inferred from the model family | | `--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 @@ -179,9 +261,10 @@ ft checkpoint --model --out [--dtype bfloat16] [--moe-backend ``` Converts an HF safetensors checkpoint to FTW, FreeToken's self-contained -fast-load format; point `ft serve --model` at the output dir. `--moe-backend -offload` (default) packs experts into offload banks; `--moe-backend triton` -keeps them dense for resident serving. See the FTW caveats in +fast-load format; point `ft serve --model` at the output dir. MoE experts are +always packed into expert banks, which serve the fused and the offload +strategies alike (GGUF and DeepSeek-V4 experts serve offload only); +`--moe-backend` is kept for compatibility. See the FTW caveats in [models.md](models.md#notes); FTW files from older builds can be repaired with [scripts/ftw_hotfix.py](ftw-hotfix.md) instead of reconverting. @@ -203,3 +286,37 @@ profile that `ft serve --moe-strategy auto` and `--moe-hybrid-max-fetch -1` then - `--threshold` (default 2.0) sets the call: recommend hybrid when CPU bandwidth beats PCIe by that factor. +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/models.md b/docs/models.md index 548a363c67..6374558d24 100644 --- a/docs/models.md +++ b/docs/models.md @@ -5,6 +5,7 @@ for them; other checkpoints of the same architectures work too. | Model | HF checkpoints | |---|---| +| DeepSeek-V4.1-Flash | [deepseek-ai/DeepSeek-V4.1-Flash](https://huggingface.co/deepseek-ai/DeepSeek-V4.1-Flash) | | DeepSeek-V4 | [deepseek-ai/DeepSeek-V4-Flash-0731](https://huggingface.co/deepseek-ai/DeepSeek-V4-Flash-0731) | | GLM-5.3-Flash | [RedHatAI/GLM-5.3-Flash-NVFP4](https://huggingface.co/RedHatAI/GLM-5.3-Flash-NVFP4) | | GLM-5.2 | [nvidia/GLM-5.2-NVFP4](https://huggingface.co/nvidia/GLM-5.2-NVFP4) | @@ -27,6 +28,7 @@ These families accept image input by default; pass `--text-model-only` to skip t | Family | Image tokens | `--image-min-tokens` / `--image-max-tokens` | `--mm-processor-kwargs` example | | --- | --- | --- | --- | +| DeepSeek-V4.1-Flash (ViT tower, streamed under `--mm-encoder-weights host`) | one feature per 42x42 pixels, plus a learned row break per row and two delimiters; checkpoint cap 1024 total span tokens | maximum caps the entire span; minimum maps to a pixel-area floor before the cap; checkpoint minimum area is 544x544 | `{"max_image_tokens": 512, "min_pixels": 295936}` | | Qwen3.6 (both variants, every listed weight format), Qwen3.8-Flash-Next, Qwen3-VL | one token per 32x32 pixels of the resized image, dynamic resolution | pixel areas in `size.shortest_edge` / `longest_edge`; checkpoint defaults 64 to 16384 tokens | `{"size": {"longest_edge": 1048576}}` | | Gemma-4 26B-A4B, 31B (`gemma4`: ViT tower, streamed under `--mm-encoder-weights host`) | one of the soft-token budgets 70 / 140 / 280 / 560 / 1120, every image scaled to its budget as far as the aspect ratio allows | the maximum picks the largest budget within it, below 70 is refused at start-up; the minimum has no effect | `{"max_soft_tokens": 1120}` | | Gemma-4 12B (`gemma4_unified`: linear patch embedder, resident under either placement) | same budgets, one 48x48 super-patch per soft token | same as the tower releases | same | @@ -38,7 +40,8 @@ These families accept image input by default; pass `--text-model-only` to skip t `ft serve --moe-strategy {auto,fused,offload,cpu,hybrid}` (`--moe-backend` is the deprecated old spelling): -- **fused** — experts resident on GPU (needs the VRAM); never auto-selected. +- **fused** — experts resident on GPU (needs the VRAM). Serves every expert + format except GGUF and DeepSeek-V4's. - **offload** — experts live in host RAM, an LRU cache of expert slots on GPU; misses stream over PCIe. - **cpu** — misses are computed on the CPU instead of fetched. @@ -46,12 +49,38 @@ These families accept image input by default; pass `--text-model-only` to skip t CPU, overlapped. Run `ft bench bw` once per machine to calibrate the split. - **auto** — dense models always resolve to `fused`; MoE models resolve to `offload`, upgraded to `hybrid` when a cached `ft bench bw` profile - recommends it. + recommends it. Unified-memory GPUs differ, see below. + +### Unified-memory GPUs (GB10 / DGX Spark) + +The GPU and the CPU share one DRAM, so offload only copies experts between two +names for the same memory. On these GPUs `auto` resolves every MoE model to +`fused`: the default is the resident path only. GGUF and DeepSeek-V4 experts +have no resident path and stay on `offload`. Offload-only flags +(`--moe-cpu-layers`, `--moe-cache-size` / `-rate` / `-auto`, +`--moe-prefill-hit-d2d`, `--disable-moe-prefill-overlap`) are ignored with a +warning; pass `--moe-strategy offload` to use them. + +### NVLink-C2C hosts (GH200 / GB200) + +The host link is ~450 GB/s per direction instead of PCIe's ~32-64 GB/s. The +offload expert gather reads host memory zero-copy and is latency-bound, so it +needs a wider grid to keep enough loads in flight. Set +`FREETOKEN_H2D_BLOCKS_PER_BANK`: +``` + FREETOKEN_H2D_BLOCKS_PER_BANK=32 ft serve --model --moe-strategy offload +``` +On GH200 this raises the gather from ~220 to ~410 GB/s (`PCIe-gather` in +`ft bench bw`), matching the DMA ceiling. Run `ft bench bw` with the variable +set so the hybrid split is calibrated against the faster gather. Leave it +unset on PCIe GPUs: wider grids add no bandwidth there. ## Notes - `ft checkpoint` conversion is optional — it pre-converts a checkpoint into FreeToken's fast-load format, and `ft serve --model` auto-detects the result. +- An FTW converted by an older build with `--moe-backend fused` or `triton` keeps + its experts as dense weights and no longer loads; reconvert it with `ft checkpoint`. - FTW files converted by builds before the quantization refactor may fail to load; see [ftw-hotfix.md](ftw-hotfix.md) for the affected checkpoints and the repair tool. - An FTW converted before its family served images holds no vision encoder: `ft serve` @@ -59,4 +88,40 @@ These families accept image input by default; pass `--text-model-only` to skip t with `ft checkpoint`, or add the encoder in place with [scripts/ftw_hotfix.py](ftw-hotfix.md). - DeepSeek-V4 checkpoints must keep the `inference/config.json` subdir — the authoritative model args are read from there. +- DeepSeek-V4.1-Flash serves text and images from the HF checkpoint as shipped (fp8 block-32 + dense weights, fp4 experts on the offload cache, 890 B/token global KV in packed + fp4/fp8 pools). Its two 98 GiB Engram tables stream from the checkpoint shards on + demand (keep them on a fast NVMe; the 6 GiB of table scales stay in host RAM), using + bounded pinned staging buffers; an FTW conversion copies the shards that hold them next + to the checkpoint. + SWA bounded replay (DeepSeek_V41_Tech_Report.pdf, shipped in the checkpoint, §3.2.2): a replayed + segment recomputes only the SWA KV of its last `n_win` tokens and truncates each query's window to + the segment. *Encoder* replay rebuilds the encoder SWA KV behind a prefix hit from the global KV + alone (approximate; the report's fallback when the SWA KV of a hit has been evicted). *Decoder* + replay runs the decoder layers on each prompt's last `n_win` tokens only; their SWA KV is never + prefix-cached and post-training simulated it. FreeToken keeps encoder SWA KV in the radix cache (no + encoder replay); `--swa-decoder-replay bounded` (default) is the report's decoder replay, with a + prefix hit stopping at least `n_win` tokens before the prompt end so those tokens are prefilled; + `exact` runs the decoder on every prompt token (reference numerics). + Bounded output differs from exact by construction and does not depend on the prefill chunk size. + Under expert offload the prefill time is bounded by streaming each layer's experts, so decoder + replay saves decoder-layer compute, not prefill time. + Run `ft bench bw --model dsv4.1-flash` to measure local PCIe and CPU bandwidth before + choosing a MoE strategy; `--moe-strategy hybrid` splits expert misses using those + measurements. With `--moe-cache-auto`, `--kv-reserve-tokens` sets the token capacity reserved + before the expert cache is allocated, subject to the attention pool's structural minimum; size it + for the intended context and concurrency. `--max-extend-length` controls the prefill chunk size + and its workspace. Image input uses the shared multimodal processor and encoder cache, including + chunked prefill and prefix replay; `--text-model-only` skips the vision weights. DSpark + speculative decoding is not served. + `reasoning_effort` takes `low`, `high` or `max`; the 1-100 integer budget goes through the + template kwargs with thinking on, `"chat_template_kwargs": {"enable_thinking": true, + "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 8a410ffe4a..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): @@ -102,9 +112,23 @@ def create_dsv4_sparse_backend(config: ModelConfig): return DSV4SparseAttnBackend(config) +@SUPPORTED_ATTENTION_BACKENDS.register( + "dsv41_sparse", + BackendInfo(supported_types=frozenset({AttnType.DSV41})), +) +def create_dsv41_sparse_backend(config: ModelConfig): + from .dsv41_sparse import DSV41SparseAttnBackend + + return DSV41SparseAttnBackend(config) + + @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. @@ -137,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/base.py b/python/freetoken/attention/base.py index 9115b80929..846e99447a 100644 --- a/python/freetoken/attention/base.py +++ b/python/freetoken/attention/base.py @@ -20,6 +20,9 @@ class AttnType(str, Enum): MLA = "mla" # plain latent-KV MLA -> MLAKVCache DSA = "dsa" # latent-KV MLA + DSA sparse indexer -> DSAKVCache DSV4 = "dsv4" # DSV4 window+compressed sparse -> DSV4PagedKVCache + # DSV41 (DeepSeek-V4.1): window + cross-layer-shared compressed KV in packed fp8/fp4 + # rows, Full/Reindex/Reuse layer modes -> DSV41PagedKVCache + DSV41 = "dsv41" LINEAR = "linear" # GDN/mamba state layers -> LinearStatePool # GQA block-sparse (MiniMax-M3): paged GQA K/V + a per-sparse-layer index-key # slab; the indexer picks top-k 128-token blocks per query -> BSAKVCache 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/dsv41_sparse.py b/python/freetoken/attention/dsv41_sparse.py new file mode 100644 index 0000000000..baaa2cf019 --- /dev/null +++ b/python/freetoken/attention/dsv41_sparse.py @@ -0,0 +1,438 @@ +"""DeepSeek-V4.1 sparse-attention backend. + +Same contract and division of labour as ``dsv4_sparse``: the backend owns the per-forward KV +ADDRESSING over the shared page table (window ring slots, compressed rows, compress-state carry, +the decode snapshot and the CUDA-graph staging buffers) plus the selection helpers that are pure +index arithmetic (causal top-k, the hierarchical candidate pool); the model computes projections +and hands the picks back to be resolved into pool rows. + +What DSV41 adds on top of DSV4's vocabulary: + +* **cross-layer sharing** -- the compressed tiers belong to the kv-source layer; a consumer layer + addresses ``source_of(layer)``'s pools. The Top-K rows and the candidate pool the Full / Reindex + layers produce ride on the per-forward ``SharedSelection`` for the Reuse layers to read. +* **packed rows** -- writes quantize into the pool (``pool.store_*``), reads dequantize in-kernel. +* **prefill passes with a window floor** -- a ``PrefillSegment`` carries ``window_floor``: the + absolute position a query's sliding window may not reach below (Decoder SWA Bounded Replay + truncates the decoder's window at ``L - W``, the start of the prompt's last window). The metadata + carries the encoder pass (every new token) and the decoder pass (the new tokens inside the + prompt's last window, or every new token in exact mode). +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, List + +import torch +from freetoken.core import Batch, get_global_ctx + +from .base import AttentionSpec, BaseAttnBackend, BaseAttnMetadata + +if TYPE_CHECKING: + from freetoken.models import ModelConfig + + +@dataclass(frozen=True) +class PrefillSegment: + """One request's slice of a flat prefill token stream. + + ``offset``/``n`` tile the stream; ``start_pos`` is the absolute position of the first token; + ``window_floor`` bounds every query's sliding window from below (0 = the exact window; bounded + replay sets it to the start of the prompt's last window). + """ + + offset: int + n: int + table_idx: int + start_pos: int + window_floor: int = 0 + + @property + def end(self) -> int: + return self.start_pos + self.n + + +@dataclass +class SharedSelection: + """What the index-producing layers hand down the stack within ONE forward. + + ``topk_rows`` are GLOBAL main-pool rows (``-1`` masked) for the current kv source, laid out + exactly as the attention kernel consumes them; ``cmp_counts`` is the decode-only per-row live + count that bounds the kernel's loop; ``candidates`` is the Hierarchical Sparse Indexer's + compressed-position pool (``-1`` = empty slot). Layers run in order and every producer writes + before its consumers read, so one slot each suffices. + """ + + source: int | None = None + topk_rows: torch.Tensor | None = None + cmp_counts: torch.Tensor | None = None + candidates: torch.Tensor | None = None + + +@dataclass +class DSV41AttnMetadata(BaseAttnMetadata): + last_indices: torch.Tensor + # Prefill: the encoder pass tiles the batch's new tokens; the decoder pass tiles the tokens the + # decoder layers process (``decoder_rows`` gathers them out of the encoder stream; None = all of + # it). ``decoder_pad``: some request has no decoder row this chunk, and its ``last_indices`` entry + # points at one placeholder row the model appends after the decoder stream (its logits are unused). + segments: List[PrefillSegment] | None = None + decoder_segments: List[PrefillSegment] | None = None + decoder_rows: torch.Tensor | None = None + decoder_pad: bool = False + # Decode: the whole-history full-loc snapshot (captured buffer under a replay, lazy eager copy). + full_snap: torch.Tensor | None = None + table_rows: torch.Tensor | None = None + window_ar: torch.Tensor | None = None + selection: SharedSelection = field(default_factory=SharedSelection) + + def get_last_indices(self, bs: int) -> torch.Tensor: + return self.last_indices[:bs] + + @property + def stage_width(self) -> int: + assert self.full_snap is not None, "stage_width is capture/replay-only" + return self.full_snap.shape[1] + + def full_snapshot(self) -> torch.Tensor: + if self.full_snap is None: + assert self.table_rows is not None, "snapshot is decode-only" + pool = get_global_ctx().kv_cache + self.full_snap = pool.full_loc_map.index_select(0, self.table_rows).to(torch.int64) + return self.full_snap + + def window_ctx(self, pos: torch.Tensor, rows: torch.Tensor): + """Layer-invariant decode ring context ``(window_slots, prev_window_slots, + window_slots_topk [B, 1, win])``; computed fresh on every call (never cached -- a capture + must record these gathers), off the snapshot so a concurrent allocation cannot redirect + an in-flight replay.""" + snap = self.full_snapshot() + translate = get_global_ctx().kv_cache.translate_full_to_window + bs = pos.shape[0] + j = self.window_ar + assert j is not None, "window_ctx is decode-only" + win = j.shape[0] + window_slots = translate(snap[rows, pos]) + prev_window_slots = translate(snap[rows, (pos - 1).clamp_min(0)]) + p = pos[:, None] - ((pos[:, None] - j[None, :]) % win) + ws = translate(snap[rows[:, None], p.clamp(min=0)]) + window_slots_topk = torch.where((p >= 0) & (j[None, :] <= pos[:, None]), ws, -1).view(bs, 1, win) + return window_slots, prev_window_slots, window_slots_topk + + def private_window_ctx(self, pos: torch.Tensor, rows: torch.Tensor): + """The private-ring counterpart of ``window_ctx`` for request-private window layers: + ``(window_slots [B], window_slots_topk [B, 1, win])`` in the layer's ring, where a request's + row holds positions by ``pos % win`` (candidate ``j`` is the position congruent to ``j``).""" + assert self.table_rows is not None and self.window_ar is not None, "private_window_ctx is decode-only" + pool = get_global_ctx().kv_cache + bs = pos.shape[0] + j = self.window_ar + win = j.shape[0] + table = self.table_rows[rows] + window_slots = pool.ring_slots(table, pos) + p = pos[:, None] - ((pos[:, None] - j[None, :]) % win) # the position at ring index j (< 0: none yet) + window_slots_topk = pool.ring_slots(table[:, None], p).view(bs, 1, win) + return window_slots, window_slots_topk + + +@dataclass +class DSV41CaptureData: + full_snap: torch.Tensor + last_indices: torch.Tensor + table_rows: torch.Tensor # each batch row's page-table row (the private window rings key on it) + + @classmethod + def create(cls, max_bs: int, width: int, device: torch.device) -> "DSV41CaptureData": + return cls( + full_snap=torch.full((max_bs, width), -1, dtype=torch.int64, device=device), + last_indices=torch.arange(max_bs, dtype=torch.int32, device=device), + table_rows=torch.zeros(max_bs, dtype=torch.int64, device=device), + ) + + +class DSV41SparseAttnBackend(BaseAttnBackend): + def __init__(self, config: ModelConfig): + self.config = config + self.device = get_global_ctx().kv_cache.device + args = config.dsv41_args + self.window_size = args.window_size + self.swa_decoder_replay: str = args.swa_decoder_replay + self.capture: DSV41CaptureData | None = None + self.capture_bs: List[int] = [] + self.max_graph_bs = 0 + self._window_ar = torch.arange(self.window_size, device=self.device) + + @property + def pool(self): + return get_global_ctx().kv_cache + + # ----- generic contract ------------------------------------------------------------- + def forward(self, q, k, v, layer_id, batch, attn_spec: AttentionSpec | None = None): + raise NotImplementedError("DSV41 attention is driven per-tier from the model module; use DSV41SparseAttnBackend.attend().") + + def prepare_metadata(self, batch: Batch) -> None: + if not batch.is_decode: + batch.attn_metadata = self._prefill_metadata(batch) + return + last = torch.tensor([r.extend_len for r in batch.padded_reqs], dtype=torch.int32, device=self.device).cumsum_(0) - 1 + batch.attn_metadata = DSV41AttnMetadata(last_indices=last, table_rows=self._table_rows(batch), window_ar=self._window_ar) + + def _prefill_metadata(self, batch: Batch) -> DSV41AttnMetadata: + """The encoder pass runs every request's new tokens ``[cached_len, device_len)``. Exact mode runs + the decoder on the same stream. Decoder SWA Bounded Replay runs it on the new tokens inside the + prompt's last window ``[L - W, L)`` (``L = prompt_len``) with every query's window floored at + ``L - W``: the chunks of a chunked prefill each take their part of that window (earlier rows + come from the request's private ring), and a chunk that ends before it runs no decoder row.""" + segments: List[PrefillSegment] = [] + off = 0 + for r in batch.reqs: + segments.append(PrefillSegment(off, r.extend_len, r.table_idx, r.cached_len)) + off += r.extend_len + if self.swa_decoder_replay == "exact": + last = torch.tensor([s.n for s in segments], dtype=torch.int32, device=self.device).cumsum_(0) - 1 + return DSV41AttnMetadata(last_indices=last, segments=segments, decoder_segments=segments) + win = self.window_size + dec: List[PrefillSegment] = [] + rows: list[torch.Tensor] = [] + last: list[int] = [] + doff = 0 + for r, s in zip(batch.reqs, segments): + floor = max(0, r.prompt_len - win) + lo = max(s.start_pos, floor) + n = max(0, s.end - lo) + if n: + # the decoder's window KV goes to the request's private rings, never to radix-shared pages + dec.append(PrefillSegment(doff, n, s.table_idx, lo, window_floor=floor)) + rows.append(torch.arange(s.offset + lo - s.start_pos, s.offset + s.n)) + doff += n + last.append(doff - 1 if n else -1) + pad = -1 in last + last_indices = torch.tensor([doff if i < 0 else i for i in last], dtype=torch.int32, device=self.device) + decoder_rows = torch.cat(rows).to(self.device, non_blocking=True) if rows else torch.empty(0, dtype=torch.int64, device=self.device) + return DSV41AttnMetadata(last_indices=last_indices, segments=segments, decoder_segments=dec, decoder_rows=decoder_rows, decoder_pad=pad) + + def init_capture_graph(self, max_seq_len: int, bs_list: List[int]) -> None: + assert self.capture is None, "Capture already initialized." + self.max_graph_bs = max(bs_list) + self.capture = DSV41CaptureData.create(self.max_graph_bs, max_seq_len, self.device) + self.capture_bs = sorted(bs_list) + + def prepare_for_capture(self, batch: Batch) -> None: + bs = batch.size + rows = torch.full((bs,), batch.padded_reqs[0].table_idx, dtype=torch.int64, device=self.device) + self._point_to_capture(batch, bs, rows) + + def prepare_for_replay(self, batch: Batch) -> None: + self._point_to_capture(batch, batch.padded_size, self._table_rows(batch)) + + def _table_rows(self, batch: Batch) -> torch.Tensor: + assert batch.active_table_idx is not None, "decode batch is missing its page-table rows" + return batch.active_table_idx.to(torch.int64) + + def _point_to_capture(self, batch: Batch, bs: int, rows_ti: torch.Tensor) -> None: + assert self.capture is not None and bs <= self.max_graph_bs + cap = self.capture + src = self.pool.full_loc_map.index_select(0, rows_ti) + w = min(src.shape[1], cap.full_snap.shape[1]) + cap.full_snap[:bs, :w].copy_(src[:bs, :w]) + cap.table_rows[:bs].copy_(rows_ti[:bs]) + batch.attn_metadata = DSV41AttnMetadata( + last_indices=cap.last_indices[:bs], full_snap=cap.full_snap[:bs], table_rows=cap.table_rows[:bs], window_ar=self._window_ar, + ) + + # ----- metadata / selection access ------------------------------------------------- + @property + def metadata(self) -> DSV41AttnMetadata: + md = get_global_ctx().batch.attn_metadata + assert isinstance(md, DSV41AttnMetadata) + return md + + @property + def selection(self) -> SharedSelection: + return self.metadata.selection + + def snapshot(self) -> torch.Tensor: + return self.metadata.full_snapshot() + + # ----- window tier ------------------------------------------------------------------- + def window_slots_of(self, ti: int, lo: int, hi: int) -> torch.Tensor: + """SHARED window slots of positions ``[lo, hi)`` off the request's LIVE full locs (slid-out -> -1).""" + return self.pool.translate_full_to_window(self.pool.full_loc_map[ti, lo:hi]) + + def ring_slots_of(self, ti: int, lo: int, hi: int) -> torch.Tensor: + """PRIVATE-ring slots of positions ``[lo, hi)`` of page-table row ``ti`` (``hi - lo <= win``).""" + assert hi - lo <= self.window_size, f"a private ring holds one window, not {hi - lo} positions" + return self.pool.ring_slots(torch.tensor(ti, device=self.device), torch.arange(lo, hi, device=self.device)) + + def layer_window_slots_of(self, layer_id: int, ti: int, lo: int, hi: int) -> torch.Tensor: + """Where ``layer_id`` keeps the window KV of positions ``[lo, hi)``: its private ring or the shared pool.""" + if self.pool.is_private_window(layer_id): + return self.ring_slots_of(ti, lo, hi) + return self.window_slots_of(ti, lo, hi) + + def store_window(self, kv: torch.Tensor, layer_id: int, window_slots: torch.Tensor) -> None: + self.pool.store_window(kv, layer_id, window_slots) + + def window_topk_prefill(self, seg: PrefillSegment, layer_id: int | None = None) -> torch.Tensor: + """Per-query window candidates for a prefill segment as GLOBAL window slots ``[1, n, win]`` + (``-1`` where empty) in the tier ``layer_id`` reads (shared pool by default). Query at + absolute ``p`` sees ``[max(floor, first_retained, p - win + 1), p]``.""" + win, device = self.window_size, self.device + lo = max(0, seg.start_pos - win + 1, seg.window_floor) + private = layer_id is not None and self.pool.is_private_window(layer_id) + ws_pool = (self.ring_slots_of if private else self.window_slots_of)(seg.table_idx, lo, seg.end) # [end - lo] + abs_p = seg.start_pos + torch.arange(seg.n, device=device).unsqueeze(1) + cand = (abs_p - win + 1).clamp(min=lo) + torch.arange(win, device=device) + cols = torch.where(cand > abs_p, -1, cand - lo) + g = ws_pool[cols.clamp_min(0)] + return torch.where(cols < 0, -1, g).unsqueeze(0) + + # ----- compressed tiers (per kv source) ---------------------------------------------- + def source_of(self, layer_id: int) -> int: + return self.pool.source_of(layer_id) + + def compressed_rows_of(self, ti: int, group_starts: torch.Tensor, ratio: int) -> torch.Tensor: + """Main / index rows of the compressed groups whose ABSOLUTE first positions are + ``group_starts``, off the request's LIVE full locs (a page is ratio-divisible, so every + position of a group shares one row).""" + return self.pool.cmp_rows(self.pool.full_loc_map[ti, group_starts], ratio) + + def locs_prefill(self, ti: int, end: int) -> torch.Tensor: + """``[1, end]`` int32: the request's live full locs for positions ``[0, end)`` -- what the indexer + derives compressed rows from (``loc // ratio``) for the positions its queries may see.""" + return self.pool.full_loc_map[ti : ti + 1, :end] + + def store_main(self, latent: torch.Tensor, source: int, rows: torch.Tensor) -> None: + self.pool.store_main(latent, source, rows) + + def store_index(self, k: torch.Tensor, source: int, rows: torch.Tensor) -> None: + self.pool.store_index(k, source, rows) + + def decode_store_rows(self, rows: torch.Tensor, pos: torch.Tensor, ratio: int, source: int, completed: torch.Tensor) -> torch.Tensor: + """Per-row decode destination: the completed group's arithmetic row, or the row's OWN scratch + row when this step did not complete a group (graph-safe, collision-free masked store).""" + row_of_group = self.pool.cmp_rows(self.snapshot()[rows, pos], ratio) + scratch = rows + self.pool.scratch_base[source] + return torch.where(completed, row_of_group, scratch) + + # ----- compress-state ring (ratio > 1 sources) --------------------------------------- + def ring_page_base(self, window_slots: torch.Tensor, ring_size: int) -> torch.Tensor: + return torch.div(window_slots, self.window_size, rounding_mode="floor") * ring_size + + def carry_state_loc(self, window_slot: int, ring_size: int) -> torch.Tensor: + base = (window_slot // self.window_size) * ring_size + return torch.arange(base, base + ring_size, device=self.device, dtype=torch.int64) + + def read_carry(self, source: int, window_slot: int) -> torch.Tensor: + """The ``[ring_size, 2 * head_dim]`` carry block at ``window_slot``'s page (kv | score).""" + ring = self.pool.state_ring[source] + return ring.get(self.carry_state_loc(window_slot, ring.ring_size)) + + def write_carry(self, source: int, window_slot: int, kv_score: torch.Tensor) -> None: + ring = self.pool.state_ring[source] + ring.set(self.carry_state_loc(window_slot, ring.ring_size), kv_score) + + def read_carry_blocks(self, source: int, window_slots: torch.Tensor) -> torch.Tensor: + ring = self.pool.state_ring[source] + return ring.get_blocks(self.ring_page_base(window_slots, ring.ring_size)) + + def write_carry_blocks(self, source: int, window_slots: torch.Tensor, blocks: torch.Tensor) -> None: + ring = self.pool.state_ring[source] + ring.set_blocks(self.ring_page_base(window_slots, ring.ring_size), blocks) + + def write_boundary_carries(self, source: int, *, lo: int, hi: int, window_slots: torch.Tensor) -> None: + """Persist the compressor carry at every window-page boundary ``B`` in ``(lo, hi]`` so a + page-aligned radix match can resume by value. A page holds whole groups (``P % ratio == 0``), + so the carry at a boundary is the empty group -- the reset block is written so a resume never + reads a stale one. ``window_slots`` is indexed by ``pos - lo``. One batched ring write, no host syncs.""" + ring = self.pool.state_ring.get(source) + if ring is None: + return + P = self.window_size + first = (lo // P + 1) * P + if first > hi: + return + bounds = torch.arange(first, hi + 1, P, device=self.device) + page_slots = window_slots[bounds - 1 - lo].to(torch.int64) # the last slot of each page + locs = (torch.div(page_slots, P, rounding_mode="floor") * ring.ring_size)[:, None] + torch.arange(ring.ring_size, device=self.device) + empty = torch.cat([ + torch.zeros(ring.ring_size, ring.item_size, dtype=torch.float32, device=self.device), + torch.full((ring.ring_size, ring.item_size), float("-inf"), dtype=torch.float32, device=self.device), + ], dim=-1) + ring.set(locs.flatten(), empty.repeat(bounds.numel(), 1)) + + # ----- indexer ---------------------------------------------------------------------- + def indexer_logits( + self, q: torch.Tensor, weights: torch.Tensor, source: int, locs: torch.Tensor, ratio: int, live: torch.Tensor, + T: int | None = None, candidates: torch.Tensor | None = None, + ) -> torch.Tensor: + """``[B, S, T]`` full-range scores (live columns written) or ``[B, S, NC]`` candidate scores; the + row of compressed position ``t`` is ``locs[b, t * ratio] // ratio``.""" + from freetoken.kernel.triton.dsv41.indexer import indexer_logits_packed + + return indexer_logits_packed(q, weights, self.pool.idx_pool[source], self.pool.idx_fmt, locs, ratio, live, T=T, candidates=candidates) + + @staticmethod + def select_topk(scores: torch.Tensor, live: torch.Tensor, topk: int) -> torch.Tensor: + """Exact causal top-k over ``[B, S, T]`` full-range scores (columns at or past ``live`` unread): + compressed positions ``[B, S, topk]`` int32 in ascending order, ``-1`` in the tail for picks + that do not exist (fewer live or finite positions than ``topk``). Ties: lowest position wins.""" + from freetoken.kernel.triton.dsv41.topk import dsv41_topk + + b, s, t = scores.shape + return dsv41_topk(scores.reshape(b * s, t), live.reshape(b * s), topk).view(b, s, topk) + + @staticmethod + def select_topk_in_candidates(scores: torch.Tensor, candidates: torch.Tensor, topk: int) -> torch.Tensor: + """Top-k over ``[B, S, NC]`` candidate-aligned scores -> the picked compressed positions + ``[B, S, topk]`` int32. The candidate list is a sorted valid prefix (``-1`` tail) and empty / + unreachable slots score ``-inf``, so the ascending winning slots map to ascending positions.""" + from freetoken.kernel.triton.dsv41.topk import dsv41_topk + + b, s, nc = scores.shape + live = torch.full((b * s,), nc, dtype=torch.int32, device=scores.device) + slots = dsv41_topk(scores.reshape(b * s, nc), live, topk).view(b, s, topk) + pos = candidates.gather(-1, slots.clamp_min(0).to(torch.int64)).to(torch.int32) + return torch.where(slots < 0, -1, pos) + + @staticmethod + def select_candidate_blocks(scores: torch.Tensor, live: torch.Tensor, topk_blocks: int, block_size: int) -> torch.Tensor: + """Level one of the Hierarchical Sparse Indexer: the ``topk_blocks`` best live ``block_size`` + blocks by their best position, the block holding the newest live position always kept, + expanded to the compact candidate list ``[B, S, topk_blocks * block_size]`` int32 of + compressed positions (ascending, ``-1`` tail). Mirrors the reference ``select_candidate_blocks`` + (which returns the equivalent boolean mask).""" + from freetoken.kernel.triton.dsv41.topk import dsv41_candidate_blocks + + b, s, t = scores.shape + return dsv41_candidate_blocks(scores.reshape(b * s, t), live.reshape(b * s), topk_blocks, block_size).view(b, s, -1) + + @staticmethod + def positions_to_rows(positions: torch.Tensor, locs: torch.Tensor, ratio: int) -> torch.Tensor: + """Compressed positions ``[B, S, K]`` (``-1`` masked) -> GLOBAL main rows ``locs[b, p * ratio] // ratio``.""" + b = positions.shape[0] + idx = (positions.to(torch.int64) * ratio).clamp_min(0).flatten(1) + full = locs.to(torch.int64).gather(1, idx.clamp(max=locs.shape[1] - 1)).view_as(positions) + rows = torch.div(full, ratio, rounding_mode="floor") + return torch.where((positions < 0) | (full < 0), -1, rows) + + # ----- attention ----------------------------------------------------------------------- + def attend( + self, q: torch.Tensor, layer_id: int, topk_idxs: torch.Tensor, n_window: int, attn_sink: torch.Tensor, + softmax_scale: float, cmp_counts: torch.Tensor | None = None, + ) -> torch.Tensor: + """Paged sparse attention over ``[window | compressed]`` global rows; the compressed half reads + the layer's kv source (window-only layers pass ``n_window == topk``).""" + from freetoken.kernel.triton.dsv41.sparse_attn import sparse_attn_packed + + pool, args = self.pool, self.config.dsv41_args + src = args.roles[layer_id].kv_source + cmp = pool.main_pool[src] if src is not None else pool.main_pool[args.backbone_kv_sources[0]] + return sparse_attn_packed( + q, pool.window_pool[layer_id], pool.win_fmt, cmp, pool.main_fmt, attn_sink, + topk_idxs, n_window, softmax_scale, cmp_counts=cmp_counts, + ) + + +__all__ = ["DSV41SparseAttnBackend", "DSV41AttnMetadata", "PrefillSegment", "SharedSelection"] 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/__main__.py b/python/freetoken/checkpoint/__main__.py index 2b7e35f1d1..0d7c2d7ed7 100644 --- a/python/freetoken/checkpoint/__main__.py +++ b/python/freetoken/checkpoint/__main__.py @@ -37,7 +37,7 @@ def main(argv: list[str] | None = None, prog: str = "freetoken.checkpoint") -> i p.add_argument("--out", required=True, help="output FTW checkpoint dir") p.add_argument("--dtype", choices=sorted(_DTYPES), default="bfloat16") p.add_argument("--moe-backend", default="offload", - help="offload (experts -> banks) or e.g. triton (experts stay dense)") + help="kept for compatibility: MoE experts always go to banks, which every strategy reads") p.add_argument("--quant-backend", type=_parse_quant_backend, default=None, help="kernel per quantized layer type, as for ft serve; the expert banks are packed for the " "chosen MoE kernel and the server picks that kernel back up from the FTW") diff --git a/python/freetoken/checkpoint/convert.py b/python/freetoken/checkpoint/convert.py index a875d77f34..152f3871c0 100644 --- a/python/freetoken/checkpoint/convert.py +++ b/python/freetoken/checkpoint/convert.py @@ -200,7 +200,8 @@ def convert_checkpoint( cfg = EngineConfig(model_path=model_path, tp_info=DistributedInfo(tp.rank, tp.size), dtype=dtype, moe_strategy=moe_backend, quant_backend=quant_backend) mc = cfg.model_config - offload = moe_backend == "offload" and getattr(mc, "is_moe", False) + # every MoE strategy reads the routed experts from the banks, fused included + offload = getattr(mc, "is_moe", False) include_moe_experts = not offload method = None if offload: @@ -208,6 +209,7 @@ def convert_checkpoint( from freetoken.layers import set_rope_device set_rope_device(dev) + object.__setattr__(cfg, "moe_strategy", "offload") method = offload_expert_method(cfg) from freetoken.utils.progress import byte_bar, count_bar diff --git a/python/freetoken/checkpoint/ftw.py b/python/freetoken/checkpoint/ftw.py index e2bd0de24e..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) @@ -432,6 +791,45 @@ def _producer(): raise err[0] +def _split_bank_entries(reader: FTWReader, path: str, num_layers: int): + """``(all, alphas, flat regions, {base: {layer: entry}})`` of the ``experts_bank`` entries; None when there are none.""" + bank_entries = reader.entries("experts_bank") + if not bank_entries: + return None + + alpha_entries = [e for e in bank_entries if e["name"] in _ALPHA_NAMES] + row_entries = [e for e in bank_entries if e["name"] not in _ALPHA_NAMES] + + meta_layers = reader.meta("expert_bank_num_layers") + if meta_layers is not None and meta_layers != num_layers: + raise RuntimeError( + f"{path!r} was converted with {meta_layers} expert-bank layers but the " + f"model config says num_moe_layers={num_layers}; the checkpoint does not " + "match its config" + ) + + # Split row entries into the two layouts by name. + flat_entries: list[dict] = [] + per_layer_groups: dict[str, dict[int, dict]] = {} + for e in row_entries: + m = _LAYER_ENTRY_RE.match(e["name"]) + if m is None: + flat_entries.append(e) + continue + per_layer_groups.setdefault(m.group("base"), {})[int(m.group("layer"))] = e + + mixed = {e["name"] for e in flat_entries} & per_layer_groups.keys() + 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(): + 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 + + def load_ftw_banks( path: str, *, num_layers: int, workers: int = 8, chunk: int = _DEFAULT_CHUNK, layer_residency: list[str] | None = None, @@ -465,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() @@ -482,40 +887,20 @@ def _backing(layer_id: int) -> str: return "mmap" reader = FTWReader(path) - bank_entries = reader.entries("experts_bank") - if not bank_entries: + try: + split = _split_bank_entries(reader, path, num_layers) + except BaseException: reader.close() - return None - - alpha_entries = [e for e in bank_entries if e["name"] in _ALPHA_NAMES] - row_entries = [e for e in bank_entries if e["name"] not in _ALPHA_NAMES] - - meta_layers = reader.meta("expert_bank_num_layers") - if meta_layers is not None and meta_layers != num_layers: + raise + if split is None: reader.close() - raise RuntimeError( - f"{path!r} was converted with {meta_layers} expert-bank layers but the " - f"model config says num_moe_layers={num_layers}; the checkpoint does not " - "match its config" - ) + return None + bank_entries, alpha_entries, flat_entries, per_layer_groups = split # Alphas: unchanged, one flat HostBank per entry. alpha_specs = {e["name"]: (tuple(e["shape"]), _dtype_of(e["dtype"])) for e in alpha_entries} alpha_hb = alloc_banks(alpha_specs) - # Split row entries into the two layouts by name. - flat_entries: list[dict] = [] - per_layer_groups: dict[str, dict[int, dict]] = {} - for e in row_entries: - m = _LAYER_ENTRY_RE.match(e["name"]) - if m is None: - flat_entries.append(e) - continue - 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)}" - # Row banks: one padded-window HostBank per (name, layer_id) for the flat layout, plus # how to carve the real [num_experts, *row_shape] tensor out of its head; ``None`` marks # a per-layer entry (direct view, no carving needed). @@ -526,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): @@ -546,15 +943,15 @@ def _backing(layer_id: int) -> str: row_jobs.append((name, bank, win_off, win_end - win_off, layer_bytes, layer_id)) 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})" - ) row_hb[base] = [] 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) @@ -659,8 +1056,71 @@ def _read_layer(job): ) +def load_ftw_banks_to_device(path: str, *, num_layers: int, device: torch.device, workers: int = 8, chunk: int = _DEFAULT_CHUNK): + """:func:`load_ftw_banks` for resident experts: every bank lands on ``device``, one (bank, layer) at a time through a transient host buffer.""" + from freetoken.moe.expert_banks import ExpertBanks + from freetoken.moe.legacy_format import canonical_role, kind_kernel_for + from freetoken.utils.progress import byte_bar + + reader = FTWReader(path) + try: + split = _split_bank_entries(reader, path, num_layers) + if split is None: + return None + bank_entries, alpha_entries, flat_entries, per_layer_groups = split + + # (name, layer_id or None for an alpha, read offset, read length, head pad, shape, dtype) + jobs = [(e["name"], None, e["global_off"], e["nbytes"], 0, tuple(e["shape"]), _dtype_of(e["dtype"])) for e in alpha_entries] + for e in flat_entries: + total, *row_shape = e["shape"] + assert total % num_layers == 0, (e["name"], total, num_layers) + layer_bytes = total // num_layers * math.prod(row_shape) * _elsize(_dtype_of(e["dtype"])) + assert layer_bytes * num_layers == e["nbytes"], (e["name"], layer_bytes, num_layers, e["nbytes"]) + for layer_id in range(num_layers): + off = e["global_off"] + layer_id * layer_bytes + win_off = (off // ALIGN) * ALIGN + jobs.append((e["name"], layer_id, win_off, _align_up(off + layer_bytes) - win_off, off - win_off, (total // num_layers, *row_shape), _dtype_of(e["dtype"]))) + for base, by_layer in per_layer_groups.items(): + for layer_id in range(num_layers): + e = by_layer[layer_id] + jobs.append((base, layer_id, e["global_off"], e["nbytes"], 0, tuple(e["shape"]), _dtype_of(e["dtype"]))) + + def _read(job) -> torch.Tensor: + _name, _layer, off, nbytes, head_pad, shape, dtype = job + buf = _transient_buffer(nbytes) + reader.read_into(memoryview(buf), {"global_off": off, "nbytes": nbytes}, workers=workers, chunk=chunk) + count = math.prod(shape) * _elsize(dtype) + return torch.frombuffer(buf, dtype=torch.uint8, count=head_pad + count)[head_pad:].view(dtype).view(shape) + + sources: dict[str, list] = {} + alphas: dict[str, torch.Tensor] = {} + bar = byte_bar(sum(e["nbytes"] for e in bank_entries), "Loading expert banks (FTW)") + try: + # read the next (bank, layer) while the current one copies to the device + with ThreadPoolExecutor(1) as ex: + ahead = ex.submit(_read, jobs[0]) + for k, (name, layer_id, *_rest) in enumerate(jobs): + host = ahead.result() + if k + 1 < len(jobs): + ahead = ex.submit(_read, jobs[k + 1]) + if layer_id is None: + alphas[name] = host.to(device) + else: + sources.setdefault(canonical_role(name), [None] * num_layers)[layer_id] = host.to(device) + bar.update(host.numel() * host.element_size()) + del host + finally: + bar.close() + quant_format = reader.meta("quant_format") + finally: + reader.close() + + kind, kernel = kind_kernel_for(quant_format) if quant_format is not None else (None, None) + return ExpertBanks(quant_format, sources, **alphas, kind=kind, kernel=kernel) + + __all__ = [ "INDEX_NAME", "FORMAT_TAG", "FORMAT_VERSION", "ALIGN", "DEFAULT_SHARD_LIMIT", - "is_ftw_checkpoint", "ftw_tensor_names", "FTWWriter", "FTWReader", - "iter_ftw_weights", "load_ftw_banks", "layer_bank_entry_name", + "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 fd691d180c..ab5b81b0af 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,11 +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). + 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: + return bool( + self.min_p > 0.0 + or self.presence_penalty != 0.0 + or self.frequency_penalty != 0.0 + or self.repetition_penalty != 1.0 + or self.logit_bias + or self.min_tokens > 0 + ) + @dataclass(eq=False) class Req: @@ -43,6 +83,8 @@ class Req: mm_items: list | None = None mrope_positions_full: torch.Tensor | None = None # [3, prompt_len] int32, CPU mrope_delta: int = 0 + # The whole prompt's length; 0 = this request's input_ids (a ChunkedReq's end at its chunk). + prompt_len: int = 0 # --- hybrid-radix (GDN linear-state) per-request slots; None for non-hybrid models or # until allocated from LinearStatePool. Set by the scheduler (P2). --- @@ -64,11 +106,13 @@ class Req: # handler must not free resources under an in-flight forward; it sets this flag and # _process_last_data frees the request when the batch drains (after copy_done.synchronize). aborted: bool = False + stop_decode_status: DecodeStatus | None = None def __post_init__(self) -> None: assert self.input_ids.is_cpu self.device_len = len(self.input_ids) self.max_device_len = len(self.input_ids) + self.output_len + self.prompt_len = self.prompt_len or self.device_len assert 0 <= self.cached_len < self.device_len <= self.max_device_len self._alloc_ids_buf() @@ -183,6 +227,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 d522be4991..4a59c7496d 100644 --- a/python/freetoken/engine/cache_budget.py +++ b/python/freetoken/engine/cache_budget.py @@ -24,7 +24,7 @@ def expert_bytes_per_slot(sources: dict[str, "list[torch.Tensor]"]) -> int: # marlin/b12x gate_up/down alpha scales are fixed [L*E] residency (do not scale # 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_paged_pool.py. + # 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()) @@ -54,6 +54,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 +62,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 +76,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 +92,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 +129,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 +141,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 3e38cb2165..5567164573 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 @@ -34,6 +35,11 @@ class EngineConfig: quant_backend: str | None = None # PLE table backend: "disk" (default) reads rows from the checkpoint files per fill, "pinned" preloads the table into page-locked host RAM. ple_backend: str = "disk" + # 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 @@ -53,6 +59,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 @@ -68,22 +76,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 @@ -91,6 +111,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) @@ -138,8 +163,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: @@ -152,4 +187,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 e855bcd32f..5a93e139f5 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,15 +13,17 @@ 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.layers import set_rope_device +from freetoken.checkpoint.ftw import ftw_tensor_names, is_ftw_checkpoint +from freetoken.layers import OffloadMoELayer, iter_moe_layers, set_rope_device from freetoken.layers.quantization import LayerKind, QuantBackend, finalize_quant, set_quant_backend from freetoken.moe.offload_cache import iter_offload_moe_layers from freetoken.mm.config import ENCODER_SECTIONS from freetoken.models import create_model, load_weight from freetoken.models.weight import ftw_lacks_vision from freetoken.moe import is_offload_moe_strategy -from freetoken.moe.expert_banks import load_expert_banks +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.utils import align_ceil, init_logger, is_sm90_family, is_sm100_family, mem_GB, torch_dtype @@ -32,11 +35,18 @@ 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__) +# 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.""" @@ -49,6 +59,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 @@ -69,6 +85,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 @@ -115,19 +167,55 @@ 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: candidates.append(("dsv4_sparse", True)) + if AttnType.DSV41 in required: + candidates.append(("dsv41_sparse", True)) if required & {AttnType.MLA, AttnType.DSA}: candidates.append(("dsa", True)) if AttnType.BSA in required: @@ -150,10 +238,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 "" + ) + + "." ) @@ -198,6 +294,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. @@ -215,6 +328,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 " @@ -322,6 +442,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: @@ -368,6 +493,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 @@ -377,14 +503,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"): @@ -417,10 +544,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 ======================== @@ -454,6 +609,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 @@ -493,19 +663,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 = ( @@ -513,13 +707,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 @@ -531,9 +718,14 @@ def _load_weights(self, config: EngineConfig) -> None: "family served images. Reconvert it with `ft checkpoint`, add the encoder in place with " "scripts/ftw_hotfix.py (docs/ftw-hotfix.md), or start with --text-model-only" ) + resident = [layer for layer in iter_moe_layers(self.model) if not isinstance(layer, OffloadMoELayer)] + if resident and not config.use_dummy_weight and is_ftw_checkpoint(config.model_path) and not ftw_tensor_names(config.model_path, "experts_bank"): + raise ValueError(f"{config.model_path} keeps its routed experts as dense weights (an older `ft checkpoint --moe-backend fused`); reconvert it with `ft checkpoint`") with _weight_load_context(): self.model.load_state_dict(self._load_weight_state_dict(config)) finalize_quant(self.model) + if resident: + self._load_resident_experts(config, resident) def _load_weight_state_dict(self, config: EngineConfig) -> Dict[str, torch.Tensor]: model_state = self.model.state_dict() @@ -541,18 +733,32 @@ def _load_weight_state_dict(self, config: EngineConfig) -> Dict[str, torch.Tenso return _make_dummy_weight_state_dict(model_state, device=self.device) # _materialize casts each loaded tensor to its model-param dtype (model_state), so # models declaring per-tensor dtypes (e.g. DSV4's mixed fp8/fp32/bf16) are preserved; - # offload models exclude experts (served from the offload cache, not dense weights). + # MoE models exclude routed experts (they load as expert banks, not dense weights). return _materialize_loaded_weight_state_dict( model_state, load_weight( config.model_path, self.device, - include_moe_experts=not is_offload_moe_strategy(config.moe_strategy), + include_moe_experts=not getattr(config.model_config, "is_moe", False), include_vision=bool(config.active_encoders), ), device=self.device, ) + def _load_resident_experts(self, config: EngineConfig, layers) -> None: + with _weight_load_context(): + banks = load_expert_banks( + config.model_path, + config.model_config, + method=layers[0].quant_method, + device=self.device, + dtype=self.dtype, + dummy=config.use_dummy_weight, + parallel={"serial": False, "parallel": True}.get(config.expert_load, None), + resident=True, + ) + attach_resident_banks(layers, banks) + @torch.inference_mode() def _warmup_encoders(self) -> None: for item in self.mm_processor.dummy_items(self.dtype, self.device): @@ -596,6 +802,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( @@ -608,7 +815,11 @@ 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, ) @@ -683,6 +894,11 @@ def _init_offload_moe_cache(self, config: EngineConfig) -> OffloadMoeCache: raise RuntimeError(f"{exc}; {_pin_hint(self._host_tables_bytes)}") from exc if config.moe_cache_auto: size, pages, overlap = self._resolve_auto_moe_cache_size(config, banks, method) + 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: @@ -733,6 +949,13 @@ def _init_offload_moe_cache(self, config: EngineConfig) -> OffloadMoeCache: # 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 if cache.decode_target in ("cpu", "hybrid"): @@ -792,6 +1015,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) + method = sample.quant_method executor = CpuMoeExecutor( cache, top_k=sample.top_k, @@ -803,7 +1027,8 @@ def _init_cpu_moe_executor(self, config: EngineConfig, cache, layers) -> None: swiglu_alpha=float(sample.alpha), swiglu_limit=sample.limit, # FIXME: the None branch serves GGUF q4_0 banks, which have no quant method yet; drop it once GGUF joins the quant path - fmt=sample.quant_method.cpu_format if sample.quant_method is not None else None, + fmt=method.cpu_format if method is not None else None, + act_block=method.cfg.act_block if method is not None else None, ) cache.set_cpu_executor(executor) self.cpu_moe_executor = executor @@ -870,6 +1095,40 @@ def _refresh_seq_state(self, config) -> None: self.page_table[self.dummy_req.table_idx].fill_(num_tokens) self.kv_cache.attach_page_table(self.page_table) + def _resize_pools( + self, config, moe_cache_size: int | None, num_pages: int | None, + num_mamba_slots: int | None, num_swa_pages: int | None, + ) -> None: + """Resize the requested pools, shrinking ones first. + + Each pool frees its own tensors before allocating, but a pool that grows ahead of one + that shrinks needs both geometries resident at once: on a nearly full card that OOMs + a resize whose final total fits (MoE 512 -> 1024 slots with KV 262K -> 64K tokens).""" + steps = [] # (shrinks, resize) + if moe_cache_size is not None: + assert self.moe_offload_cache is not None, "no MoE offload cache to resize" + steps.append(( + moe_cache_size < self.moe_offload_cache.cache_size, + lambda: self.moe_offload_cache.rebuild(moe_cache_size), + )) + if num_pages is not None or num_swa_pages is not None: + # A window-only change (num_pages None) re-derives the window pool at the new pin + # against the CURRENT page count; _resize_kv_pool sets self.num_pages either way. + pages = num_pages if num_pages is not None else self.num_pages + steps.append(( + pages < self.num_pages, + lambda: self._resize_kv_pool(config, pages, num_swa_pages), + )) + if num_mamba_slots is not None: + # 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. + steps.append(( + num_mamba_slots + 1 < self.linear_state_pool.num_slots, + lambda: self.linear_state_pool.rebuild(num_mamba_slots + 1), + )) + for _, resize in sorted(steps, key=lambda step: not step[0]): + resize() + @torch.inference_mode() def rebuild_runtime_cache( self, @@ -905,6 +1164,14 @@ def rebuild_runtime_cache( 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( @@ -951,7 +1218,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 "" @@ -979,22 +1247,7 @@ def rebuild_runtime_cache( # 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) + self._resize_pools(config, moe_cache_size, num_pages, num_mamba_slots, num_swa_pages) # 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) @@ -1024,6 +1277,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. @@ -1035,9 +1291,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: @@ -1048,18 +1442,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) @@ -1077,6 +1471,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) @@ -1089,23 +1486,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(): @@ -1135,15 +1562,14 @@ def _is_unified_memory_gpu(index: "int | None" = None) -> bool: return False +_RESIDENT_EXPERT_QUANTS = ("none", "fp8_block", "nvfp4") + + def _fused_resident_ok(model_config) -> bool: - """Whether the resident ('fused') MoE path can hold this model's experts. - - FIXME: auto resolves to fused only for bf16 and fp8_block experts; drop this gate once the other quant formats support fused. - """ - expert_quant = getattr(model_config, "expert_quant", "none") - if expert_quant not in ("none", "fp8_block"): + """Whether the resident ('fused') MoE path can hold this model's experts; GGUF q4_0 and DeepSeek-V4 experts stay offload.""" + if getattr(model_config, "expert_quant", "none") not in _RESIDENT_EXPERT_QUANTS: return False - return getattr(model_config, "moe_weight_format", None) in (None, "bf16") + return getattr(model_config, "moe_weight_format", None) in (None, "bf16", "mxfp4") def _ensure_expandable_segments() -> None: @@ -1181,6 +1607,39 @@ def _resolve_cache_type(has_linear_attention: bool, requested: str) -> str: return requested +def _adjust_dsv41_config(config: EngineConfig, override) -> None: + """DeepSeek-V4.1 config reconciliation, the DSV4 policy with the replay knob: sync the + runtime into ``dsv41_args``, page_size = the window page P, radix -> swa_radix, decode graph + batches <= max_running_req. Unlike DSV4 the prefill chunk keeps ``max_extend_tokens`` (whole + window pages): at a 1M ceiling the window pool alone would admit ~40K-token chunks whose + activations (64 x 512 latent queries per token) do not fit next to the expert cache.""" + model_config = config.model_config + args = model_config.dsv41_args + args.max_seq_len = config.max_seq_len + args.max_batch_size = config.max_running_req + 1 # +1 dummy + args.swa_decoder_replay = config.swa_decoder_replay + P = args.window_size + override("page_size", P) + logger.info_rank0( + f"DSV41 KV pages are {P}-token window pages; page_size set to {P}; SWA decoder replay: {args.swa_decoder_replay}" + + ("" if args.swa_decoder_replay == "exact" else f" (prefix hits stop {P} tokens before the prompt end)") + ) + if getattr(config, "cache_type", "radix") != "naive": + override("cache_type", "swa_radix") + current = getattr(config, "max_extend_tokens", None) # a SchedulerConfig field + if current is not None and (extend := max(P, current // P * P)) != current: + override("max_extend_tokens", extend) + mr = config.max_running_req + if config.cuda_graph_max_bs is not None and config.cuda_graph_max_bs > mr: + logger.warning_rank0(f"cuda_graph_max_bs {config.cuda_graph_max_bs} exceeds max_running_req {mr}; clamping (larger decode batches never occur).") + override("cuda_graph_max_bs", mr) + if config.cuda_graph_bs is not None: + kept = [bs for bs in config.cuda_graph_bs if bs <= mr] + if kept != list(config.cuda_graph_bs): + logger.warning_rank0(f"dropping cuda_graph_bs entries above max_running_req {mr}: {[bs for bs in config.cuda_graph_bs if bs > mr]}") + override("cuda_graph_bs", kept) + + def _adjust_dsv4_config(config: EngineConfig, override) -> None: """DSV4 engine-config reconciliation at config-resolution time (before the pool exists). Syncs the resolved runtime config into the opaque ``dsv4_args`` payload, sets @@ -1458,6 +1917,7 @@ def override(attr: str, value: Any): # this is dangerous, use with caution model_config = config.model_config single_stream_only = getattr(model_config, "single_stream_only", False) is_dsv4 = getattr(model_config, "dsv4_args", None) is not None + is_dsv41 = getattr(model_config, "dsv41_args", None) is not None has_swa_attention = getattr(model_config, "has_swa_attention", False) has_linear_attention = getattr(model_config, "has_linear_attention", False) is_moe = getattr(model_config, "is_moe", False) @@ -1501,6 +1961,8 @@ def override(attr: str, value: Any): # this is dangerous, use with caution if is_dsv4: _adjust_dsv4_config(config, override) + if is_dsv41: + _adjust_dsv41_config(config, override) if has_swa_attention: # Both SWA cache paths use the global-paged swa pool (page_size==1 only for now). @@ -1528,6 +1990,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} @@ -1555,7 +2045,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) @@ -1594,8 +2084,8 @@ def override(attr: str, value: Any): # this is dangerous, use with caution if is_moe and config.moe_strategy == "auto": # A MoE model always defaults to the offload family: experts stream from pinned host # banks into an auto-sized GPU slot cache, which is the only default that serves a model - # bigger than the GPU. The resident 'fused' path (bf16 / block-fp8 experts, the two - # formats MoELayer can allocate) is still reachable, but only when asked for explicitly + # bigger than the GPU. The resident 'fused' path (every expert format but GGUF q4_0 and + # DeepSeek-V4's) is still reachable, but only when asked for explicitly # -- auto never picks it, because nothing here knows whether the experts would fit in # HBM and a wrong guess is a weight-load OOM rather than a slower-but-working run. default_backend = "offload" @@ -1614,6 +2104,18 @@ def override(attr: str, value: Any): # this is dangerous, use with caution "Unified-memory GPU detected; auto-selecting 'fused' MoE strategy " "(resident experts) instead of offload" ) + # the cache sizing flags get their own warning in the 'fused' block below + inert = [flag for flag, on in ( + (f"--moe-cpu-layers={config.moe_cpu_layers}", config.moe_cpu_layers), + ("--moe-prefill-hit-d2d", config.moe_prefill_hit_d2d), + ("--disable-moe-prefill-overlap", not config.moe_prefill_overlap), + ) if on] + if inert: + logger.warning_rank0( + f"Unified-memory GPU serves MoE experts resident; ignoring {', '.join(inert)} " + "(use --moe-strategy offload to keep the offload flags)" + ) + override("moe_cpu_layers", None) # Hardware-adaptive config: a cached `ft bench bw` profile can upgrade # the offload default to hybrid when this machine's CPU MoE bandwidth clears its PCIe # gather bandwidth by the bench threshold (default 2x). hybrid is VRAM-equivalent to @@ -1733,7 +2235,7 @@ def override(attr: str, value: Any): # this is dangerous, use with caution if ( is_moe - and expert_quant not in ("none", "fp8_block") + and expert_quant not in _RESIDENT_EXPERT_QUANTS and not is_offload_moe_strategy(config.moe_strategy) ): raise ValueError( @@ -1775,7 +2277,7 @@ def override(attr: str, value: Any): # this is dangerous, use with caution # DSV4 is exempt: it sizes its own table from the resolved max_seq_len (_adjust_dsv4_config). rotary = getattr(model_config, "rotary_config", None) seq_override = getattr(config, "max_seq_len_override", None) - if seq_override is not None and rotary is not None and not is_dsv4: + if seq_override is not None and rotary is not None and not (is_dsv4 or is_dsv41): if seq_override > rotary.table_positions: raise ValueError( f"--max-seq-len-override {seq_override} exceeds the model's " diff --git a/python/freetoken/engine/sample.py b/python/freetoken/engine/sample.py index 3bdae76957..89d2a6fb88 100644 --- a/python/freetoken/engine/sample.py +++ b/python/freetoken/engine/sample.py @@ -10,18 +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 + (presence / frequency / repetition penalties, logit_bias, min_tokens, min_p). + + Built on the host in ``Sampler.prepare`` from the requests' SamplingParams and + token histories, 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) + presence: torch.Tensor | None = None # [r, 1] float32 + frequency: torch.Tensor | None = None # [r, 1] float32 + repetition: torch.Tensor | None = None # [r, 1] float32 (1.0 = off) + # Generated ids per row, right-padded with vocab_size (a scratch column that is + # dropped), for the presence / frequency counts and the repetition mask. + hist_ids: torch.Tensor | None = None # [r, L] int64 + # Prompt ids per row, same padding: repetition_penalty also covers the prompt. + prompt_ids: torch.Tensor | None = None # [r, P] int64 + 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) 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) + + @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 + 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. A synchronous copy: the + histories can be long (prompt + output) and a pinned buffer must not be freed under + an in-flight non_blocking copy.""" + width = max(1, max((len(r) for r in rows), default=1)) + out = torch.full((len(rows), width), pad, dtype=torch.int64) + for i, r in enumerate(rows): + if r: + out[i, : len(r)] = torch.as_tensor(r, dtype=torch.int64) + return out.to(device) + + +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, in vLLM's order: + repetition -> presence -> frequency penalties, 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) + r = sub.shape[0] + device = sub.device + neg_inf = float("-inf") + + if plan.hist_ids is not None: + counts = torch.zeros((r, vocab_size + 1), dtype=torch.float32, device=device) + counts.scatter_add_(1, plan.hist_ids, torch.ones_like(plan.hist_ids, dtype=torch.float32)) + counts = counts[:, :vocab_size] + if plan.repetition is not None: + seen = counts > 0 + if plan.prompt_ids is not None: + in_prompt = torch.zeros((r, vocab_size + 1), dtype=torch.bool, device=device) + in_prompt.scatter_(1, plan.prompt_ids, torch.ones_like(plan.prompt_ids, dtype=torch.bool)) + seen = seen | in_prompt[:, :vocab_size] + # HF / vLLM semantics: divide positive logits, multiply negative ones. + penalized = torch.where(sub > 0, sub / plan.repetition, sub * plan.repetition) + sub = torch.where(seen, penalized, sub) + if plan.presence is not None: + sub = sub - plan.presence * (counts > 0).to(sub.dtype) + if plan.frequency is not None: + sub = sub - plan.frequency * counts + + 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, @@ -56,11 +154,106 @@ 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. + Token histories come from the host-side ``req.input_ids``; under overlap + scheduling the previous step's token may not be appended yet, so a penalty sees + the history one token late. That is the accepted cost of keeping the counts off + the device.""" + 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 + + want_hist = any( + p.presence_penalty != 0.0 or p.frequency_penalty != 0.0 or p.repetition_penalty != 1.0 + for _, p in picked + ) + if want_hist: + hist: list[list[int]] = [] + prompts: list[list[int]] = [] + want_prompt = False + for req, p in picked: + prompt_len = req.max_device_len - req.output_len + ids = req.input_ids + penalized = p.presence_penalty != 0.0 or p.frequency_penalty != 0.0 or p.repetition_penalty != 1.0 + hist.append(ids[prompt_len:].tolist() if penalized else []) + if p.repetition_penalty != 1.0: + want_prompt = True + # multimodal placeholders are pseudo-ids >= vocab_size (MM_PAD_SHIFT_VALUE); + # send them to the scratch column, not past the scatter target + prompts.append(ids[:prompt_len].clamp(max=pad).tolist()) + else: + prompts.append([]) + plan.hist_ids = _padded_rows(hist, pad, self.device) + if want_prompt: + plan.prompt_ids = _padded_rows(prompts, pad, self.device) + if any(p.presence_penalty != 0.0 for _, p in picked): + plan.presence = make_device_tensor([[p.presence_penalty] for _, p in picked], torch.float32, self.device) + if any(p.frequency_penalty != 0.0 for _, p in picked): + plan.frequency = make_device_tensor([[p.frequency_penalty] for _, p in picked], torch.float32, self.device) + if any(p.repetition_penalty != 1.0 for _, p in picked): + plan.repetition = make_device_tensor([[p.repetition_penalty] for _, p in picked], torch.float32, self.device) + + 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 prepare(self, batch: Batch) -> BatchSamplingArgs: params = [r.sampling_params for r in batch.reqs] is_greedy = [p.is_greedy for p in params] + 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) + return BatchSamplingArgs( + temperatures=None, + 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 @@ -83,11 +276,21 @@ def prepare(self, batch: Batch) -> BatchSamplingArgs: greedy_mask = ( 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) + return BatchSamplingArgs( + temperatures, + top_k=top_k, + top_p=top_p, + greedy_mask=greedy_mask, + 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.plan is not None: + logits = apply_logits_processors(logits, args.plan, self.vocab_size) if args.temperatures is None: # greedy sampling return torch.argmax(logits, dim=-1) tokens = sample_impl(logits.float(), args.temperatures, args.top_k, args.top_p) @@ -97,3 +300,59 @@ def sample(self, logits: torch.Tensor, args: BatchSamplingArgs) -> torch.Tensor: greedy_tokens = torch.argmax(logits, dim=-1).to(tokens.dtype) tokens = torch.where(args.greedy_mask, greedy_tokens, tokens) 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 5e87f8c923..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( @@ -153,8 +170,11 @@ def default_kernel_specs() -> tuple[KernelSpec, ...]: for feature_size in DEFAULT_FAST_INDEX_COPY_FEATURE_SIZES ) specs.append(_fast_index_copy_multi_spec(num_threads=1024, blocks_per_bank=8)) + 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 a8ef85d432..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 @@ -332,6 +335,17 @@ def expert_bank_row_bytes(fmt: str, hidden_size: int, moe_intermediate_size: int expert_formats=("ds_fp4",), embed_indexing=False, # plain nn.Embedding ), + AotModel( + # DSV41 pools are packed byte rows written by the dsv41 pack kernel, not store_cache; + # the routed experts share DSV4's ds_fp4 bank layout (hidden 5120, intermediate 2304). + name="deepseek-ai/DeepSeek-V4.1-Flash", + architecture="DeepseekV41ForCausalLM", + hidden_size=5120, + kv_groups=(), + top_k=6, + moe_intermediate_size=2304, + expert_formats=("ds_fp4",), + ), # ---- dense checkpoints (store/index only, no expert banks) ---- AotModel( name="Qwen/Qwen3.6-27B", @@ -395,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]]: @@ -418,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/cpu_moe_ext.cpp b/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp index e2cb0d9c33..86765013e3 100644 --- a/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp +++ b/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp @@ -49,6 +49,14 @@ 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) { @@ -1055,11 +1063,12 @@ inline void deinterleave_bf16_f32(const bf16_t* src, float* even, float* odd, in } } -// DeepSeek-V4 activation FP8 round-trip (bf16 in/out): per 128-block, -// 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) { - for (int b0 = 0; b0 < K; b0 += 128) { - const int b1 = std::min(K, b0 + 128); +// 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))); @@ -1267,6 +1276,8 @@ struct CpuMoeExecutor { // 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]; @@ -1290,13 +1301,19 @@ struct CpuMoeExecutor { std::mutex sync_mtx; std::condition_variable sync_cv; - bool stop = false; + std::atomic stop{false}; // spinning workers read it without the lock uint64_t cur_gen = 0; MoeTask* cur_task = nullptr; - std::atomic submitted{0}; - std::atomic completed{0}; + // 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}; + + // Spin-then-park waits. Off unless the Python wrapper finds spare CPUs for the spinners. + static constexpr auto kSpinWindow = std::chrono::milliseconds(50); + std::atomic spin_wait{false}; + bool coord_pinned = false; - std::atomic p1_next{0}; + alignas(64) std::atomic p1_next{0}; std::atomic p2_next{0}; std::atomic prt_next{0}; // ds_fp4 intermediate fp8 round-trip phase int64_t p1_total = 0, p2_total = 0, prt_total = 0; @@ -1803,7 +1820,7 @@ struct CpuMoeExecutor { gas_scratch.data() + (size_t)r * (I / 32)); return; } - if (fmt == WF_DSFP4) fp8_roundtrip_bf16(g, g, I); + 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); @@ -1866,19 +1883,34 @@ struct CpuMoeExecutor { } } + // Tasks arrive once per MoE layer, sub-ms apart. Spinning skips a futex wakeup per worker per layer; parking after kSpinWindow keeps an idle server off the CPU. void worker_loop(int tid) { pin_self(tid); uint64_t my_gen = 0; + using wclk = std::chrono::steady_clock; + auto last_active = wclk::now(); for (;;) { - MoeTask* t; - { + uint64_t gen = submitted.load(std::memory_order_acquire); + if (gen == my_gen && spin_wait.load(std::memory_order_relaxed)) { + unsigned polls = 0; + while (gen == my_gen) { + cpu_relax(); + if (stop.load(std::memory_order_relaxed)) return; + if ((++polls & 1023u) == 0 && wclk::now() - last_active > kSpinWindow) break; + gen = submitted.load(std::memory_order_acquire); + } + } + if (gen == my_gen) { std::unique_lock lk(task_mtx); - task_cv.wait(lk, [&] { return stop || cur_gen != my_gen; }); - if (stop) return; - my_gen = cur_gen; - t = cur_task; + task_cv.wait(lk, [&] { return stop.load() || cur_gen != my_gen; }); + if (stop.load()) return; + gen = cur_gen; } + // submit() writes cur_task before the release store to `submitted`. + my_gen = gen; + MoeTask* t = cur_task; run_task_body(t); + last_active = wclk::now(); if (done_count.fetch_add(1) + 1 == num_threads) { completed.store(my_gen, std::memory_order_release); { @@ -1928,7 +1960,7 @@ struct CpuMoeExecutor { 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); + fp8_roundtrip_bf16(src, xq, H, act_block); src = xq; } float* xe = xe_scratch.data() + (size_t)tok * (H / 2); @@ -1961,8 +1993,19 @@ struct CpuMoeExecutor { task_cv.notify_all(); } - void sync() { + // Only a coordinator on its own core may spin: host-func callbacks and eager calls share cores with the workers. + void sync(bool spin = false) { const uint64_t target = submitted.load(std::memory_order_acquire); + if (spin) { + using sclk = std::chrono::steady_clock; + const auto t0 = sclk::now(); + unsigned polls = 0; + while (completed.load(std::memory_order_acquire) < target) { + cpu_relax(); + if ((++polls & 1023u) == 0 && sclk::now() - t0 > kSpinWindow) break; + } + if (completed.load(std::memory_order_acquire) >= target) return; + } std::unique_lock lk(sync_mtx); sync_cv.wait(lk, [&] { return completed.load(std::memory_order_acquire) >= target; }); } @@ -2003,6 +2046,7 @@ struct CpuMoeExecutor { } flag_served.assign(num_slots, 0); 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) { @@ -2050,7 +2094,7 @@ struct CpuMoeExecutor { } if (t != nullptr) { submit(t); - sync(); + 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); @@ -2144,7 +2188,14 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { .def("set_input_prequant", [](CpuMoeExecutor& e, bool v) { e.input_prequant = v; }, py::arg("value")) - .def("isa_name", &CpuMoeExecutor::isa_name); + .def("set_act_block", + [](CpuMoeExecutor& e, int v) { e.act_block = v; }, + py::arg("value")) + .def("isa_name", &CpuMoeExecutor::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")); 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 1aaa1303d2..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"} @@ -149,6 +154,23 @@ def _jit_fast_index_copy_multi_module(*, num_threads: int, blocks_per_bank: int) cuda_wrappers=[("launch", f"&MultiIndexCopyKernel<{args}>::run")], ) +def default_blocks_per_bank(device: torch.device) -> int: + """H2D gather grid width per bank. The link's bandwidth-delay product sets how many + 16B loads must be in flight: PCIe (~31 GB/s) saturates at ~4K threads/bank, NVLink-C2C + (GH200, ~450 GB/s) needs ~32K (measured: 8x1024 -> 222, 32x1024 -> 412 GB/s).""" + env = os.environ.get("FREETOKEN_H2D_BLOCKS_PER_BANK") + if env is not None: + blocks_per_bank = int( + env + ) + assert blocks_per_bank > 0, ( + "FREETOKEN_H2D_BLOCKS_PER_BANK must be greater than 0, " + f"got {blocks_per_bank}" + ) + return blocks_per_bank # manual override (validate > 0) + if any(k in torch.cuda.get_device_name(device) for k in ("GH200", "GB200", "GB300")): + return 32 + return 8 def fast_index_copy_multi_jit( dst_ptrs: torch.Tensor, @@ -185,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 3f53c5fc0e..9165041753 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 @@ -18,25 +45,34 @@ def _optin_smem_bytes(device_index: int) -> int: return int(getattr(props, "shared_memory_per_block_optin", 0)) -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) @@ -47,6 +83,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, @@ -60,6 +100,8 @@ def _paged_attention_kernel( stride_kh, stride_vs, stride_vh, + stride_kss, + stride_vss, stride_ot, stride_oh, GROUP: tl.constexpr, @@ -68,6 +110,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) @@ -107,14 +151,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")) @@ -124,14 +190,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 @@ -149,6 +235,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, @@ -162,6 +252,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, @@ -179,6 +271,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) @@ -224,7 +318,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) @@ -232,24 +329,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 @@ -349,6 +478,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, @@ -364,16 +514,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 @@ -391,6 +556,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. @@ -405,6 +581,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, @@ -418,6 +598,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), @@ -435,6 +617,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, ) @@ -471,6 +655,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, @@ -485,6 +673,8 @@ def _extend_attention_kernel( stride_kh, stride_vs, stride_vh, + stride_kss, + stride_vss, stride_ot, stride_oh, GROUP: tl.constexpr, @@ -495,6 +685,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) @@ -555,15 +747,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) @@ -572,15 +782,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 @@ -602,6 +830,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, @@ -620,6 +852,8 @@ def _extend_attention_split_kernel( stride_kch, stride_vcs, stride_vch, + stride_kss, + stride_vss, stride_ot, stride_oh, GROUP: tl.constexpr, @@ -630,6 +864,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) @@ -684,15 +920,33 @@ def _extend_attention_split_kernel( if not skip_tile: slots = tl.load(kv_indices_ptr + kv_start + kv_offsets, mask=mask_n, other=0) - k = tl.load( - k_cache_ptr - + slots[None, :] * stride_kcs - + kv_head * stride_kch - + offs_d[:, None], - mask=mask_n[None, :] & mask_d[:, None], - other=0.0, - ) + 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) + elif HAS_KV_SCALE: + 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) + else: + k = tl.load( + k_cache_ptr + + slots[None, :] * stride_kcs + + kv_head * stride_kch + + 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) @@ -701,15 +955,33 @@ def _extend_attention_split_kernel( alpha = tl.exp(m_i - m_new) p = tl.exp(scores - m_new[:, None]) - v = tl.load( - v_cache_ptr - + slots[:, None] * stride_vcs - + kv_head * stride_vch - + 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_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) + elif HAS_KV_SCALE: + 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) + else: + v = tl.load( + v_cache_ptr + + slots[:, None] * stride_vcs + + kv_head * stride_vch + + 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 @@ -793,18 +1065,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 @@ -816,14 +1102,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: @@ -839,6 +1138,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, @@ -857,6 +1160,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, @@ -867,6 +1172,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, @@ -877,6 +1184,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, @@ -891,6 +1202,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, @@ -901,6 +1214,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, @@ -921,20 +1236,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 @@ -944,12 +1267,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, @@ -963,6 +1300,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, @@ -971,6 +1310,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/dsv4/fp8_linear.py b/python/freetoken/kernel/triton/dsv4/fp8_linear.py index 47e3c9f58a..fb17a35175 100644 --- a/python/freetoken/kernel/triton/dsv4/fp8_linear.py +++ b/python/freetoken/kernel/triton/dsv4/fp8_linear.py @@ -1,23 +1,25 @@ -"""Block-scaled FP8 (e4m3) linear for DeepSeek-V4, matching the reference numerics. +"""Block-scaled FP8 (e4m3) linear for DeepSeek-V4 / V4.1, matching the reference numerics. The reference (``inference/model.py`` ``linear`` + ``inference/kernel.py`` -``act_quant``/``fp8_gemm``) quantizes the *activation* to FP8 with a per-128 block +``act_quant``/``fp8_gemm``) quantizes the *activation* to FP8 with a per-``block`` power-of-two (ue8m0) scale, then runs an FP8xFP8 block-scaled GEMM against the FP8 -weight (which carries its own 128x128 ue8m0 block scale). Both operands' scales are -applied per 128-K block to a separate FP32 accumulator. This module reproduces that: +weight (which carries its own ``block x block`` ue8m0 scale). Both operands' scales are +applied per ``block``-K slab to a separate FP32 accumulator. This module reproduces that: - ``y = fp8_gemm(act_quant(x, 128, ue8m0), weight_fp8, weight_scale_e8m0)`` + ``y = fp8_gemm(act_quant(x, block, ue8m0), weight_fp8, weight_scale_e8m0)`` + +``block`` is the checkpoint's ``weight_block_size``: 128 for DeepSeek-V4, 32 for V4.1. ``act_quant`` (reference): per block ``s = 2**ceil(log2(max(|x|,1e-4)/448))`` (exact via IEEE bit ops -> matches ``fast_round_scale``), ``x_fp8 = round_e4m3(clamp(x/s, +-448))``, scale stored e8m0. The GEMM accumulates ``sum_k (A_fp8 @ B_fp8) * s_a * s_b`` -per 128-K block in FP32. +per ``block``-K slab in FP32. Also provides ``act_quant_fp8_inplace`` -- the fused FP8 quant+dequant round-trip the reference applies in-place to the window / compressor KV (``act_quant(..., 64, ..., inplace=True)``), returning BF16. -Assumes ``K % 128 == 0`` and ``N % 128 == 0`` (true for every DeepSeek-V4 projection). +Assumes ``K % block == 0`` and ``N % block == 0`` (true for every DeepSeek-V4 / V4.1 projection). """ from __future__ import annotations @@ -218,8 +220,8 @@ def fp4_act_quant_inplace(x: torch.Tensor, block: int = 32) -> torch.Tensor: def _fp8_act_gemm_kernel( a_ptr, # [M, K] float8_e4m3fn (quantized activation) w_ptr, # [N, K] float8_e4m3fn - sa_ptr, # [M, K//128] uint8 (e8m0 act codes) - sb_ptr, # [N//128, K//128] uint8 (e8m0 weight codes) + sa_ptr, # [M, K//BLOCK_K] uint8 (e8m0 act codes) + sb_ptr, # [N//BLOCK_K, K//BLOCK_K] uint8 (e8m0 weight codes; square BLOCK_K blocks) c_ptr, # [M, N] compute dtype M, N, K, stride_am, stride_ak, stride_wn, stride_wk, @@ -228,6 +230,11 @@ def _fp8_act_gemm_kernel( BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, compute_type: tl.constexpr, ): + """``BLOCK_K`` is the weight's square block edge, so each K step maps to exactly one scale + column. When the N tile IS one scale row (``BLOCK_N == BLOCK_K``, the DeepSeek-V4 block-128 + case) the weight scale is a scalar per step and the N masks vanish -- the original kernel; + otherwise (block 32) the tile spans ``BLOCK_N // BLOCK_K`` scale rows and N is masked.""" + SCALAR_SCALE: tl.constexpr = BLOCK_N == BLOCK_K pid_m = tl.program_id(0) pid_n = tl.program_id(1) offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) @@ -238,112 +245,217 @@ def _fp8_act_gemm_kernel( w_ptrs = w_ptr + offs_n[:, None] * stride_wn + offs_k[None, :] * stride_wk acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) num_k = tl.cdiv(K, BLOCK_K) - for k in range(num_k): - a = tl.load(a_ptrs, mask=m_mask[:, None], other=0.0) - w = tl.load(w_ptrs) - if e4m3_native_cx(): - p = tl.dot(a, tl.trans(w), out_dtype=tl.float32) - else: - # bf16 dot on the same e4m3 grid: operands exact in bf16, fp32 acc - p = tl.dot(a, tl.trans(e4m3_u8_to_f32(w).to(tl.bfloat16)), out_dtype=tl.float32) - sa_code = tl.load(sa_ptr + offs_m * stride_sam + k * stride_sak, mask=m_mask, other=0) - sca = tl.exp2(sa_code.to(tl.float32) - 127.0) # [BLOCK_M] - sb_code = tl.load(sb_ptr + pid_n * stride_sbn + k * stride_sbk) - scb = tl.exp2(sb_code.to(tl.float32) - 127.0) # scalar (one 128-N block) - acc += p * sca[:, None] * scb - a_ptrs += BLOCK_K * stride_ak - w_ptrs += BLOCK_K * stride_wk - c_ptrs = c_ptr + offs_m[:, None] * stride_cm + offs_n[None, :] * stride_cn - tl.store(c_ptrs, acc.to(compute_type), mask=m_mask[:, None]) + if SCALAR_SCALE: + for k in range(num_k): + a = tl.load(a_ptrs, mask=m_mask[:, None], other=0.0) + w = tl.load(w_ptrs) + if e4m3_native_cx(): + p = tl.dot(a, tl.trans(w), out_dtype=tl.float32) + else: + # bf16 dot on the same e4m3 grid: operands exact in bf16, fp32 acc + p = tl.dot(a, tl.trans(e4m3_u8_to_f32(w).to(tl.bfloat16)), out_dtype=tl.float32) + sa_code = tl.load(sa_ptr + offs_m * stride_sam + k * stride_sak, mask=m_mask, other=0) + sca = tl.exp2(sa_code.to(tl.float32) - 127.0) # [BLOCK_M] + sb_code = tl.load(sb_ptr + pid_n * stride_sbn + k * stride_sbk) + scb = tl.exp2(sb_code.to(tl.float32) - 127.0) # scalar (one N block) + acc += p * sca[:, None] * scb + a_ptrs += BLOCK_K * stride_ak + w_ptrs += BLOCK_K * stride_wk + c_ptrs = c_ptr + offs_m[:, None] * stride_cm + offs_n[None, :] * stride_cn + tl.store(c_ptrs, acc.to(compute_type), mask=m_mask[:, None]) + else: + n_mask = offs_n < N # N is a multiple of BLOCK_K, not necessarily of BLOCK_N + sn = offs_n // BLOCK_K + for k in range(num_k): + a = tl.load(a_ptrs, mask=m_mask[:, None], other=0.0) + w = tl.load(w_ptrs, mask=n_mask[:, None], other=0.0) + if e4m3_native_cx(): + p = tl.dot(a, tl.trans(w), out_dtype=tl.float32) + else: + p = tl.dot(a, tl.trans(e4m3_u8_to_f32(w).to(tl.bfloat16)), out_dtype=tl.float32) + sa_code = tl.load(sa_ptr + offs_m * stride_sam + k * stride_sak, mask=m_mask, other=0) + sca = tl.exp2(sa_code.to(tl.float32) - 127.0) # [BLOCK_M] + sb_code = tl.load(sb_ptr + sn * stride_sbn + k * stride_sbk, mask=n_mask, other=0) + scb = tl.exp2(sb_code.to(tl.float32) - 127.0) # [BLOCK_N] + acc += p * sca[:, None] * scb[None, :] + a_ptrs += BLOCK_K * stride_ak + w_ptrs += BLOCK_K * stride_wk + c_ptrs = c_ptr + offs_m[:, None] * stride_cm + offs_n[None, :] * stride_cn + tl.store(c_ptrs, acc.to(compute_type), mask=m_mask[:, None] & n_mask[None, :]) @triton.jit -def _fp8_act_gemv_splitk_kernel( - a_ptr, # [K] float8_e4m3fn - sa_ptr, # [K//128] uint8 (e8m0 act codes) - w_ptr, # [N, K] float8_e4m3fn - sb_ptr, # [N//128, K//128] uint8 (e8m0 weight codes) - part_ptr, # [SPLIT_K, N] fp32 - N, K, - stride_ak, stride_wn, stride_wk, stride_sbn, stride_sbk, stride_pk, stride_pn, - BLOCK_N: tl.constexpr, SPLIT_K: tl.constexpr, +def _fp8_gemv_kernel( + x_ptr, # [M, G, K] activation in the compute dtype (bf16); QUANT_ACT: quantized in-kernel per SLAB + w_ptr, # [G * N, K] float8_e4m3fn (or its uint8 view on pre-Ada parts) + sb_ptr, # [G * N // SLAB, K // SLAB] uint8 (e8m0 weight codes) + out_ptr, # [M, G * N] OUT dtype (SPLIT_K == 1) + part_ptr, # [SPLIT_K, M, G * N] fp32 partials (SPLIT_K > 1) + M, N, K, + stride_xm, stride_xg, stride_wn, stride_wk, stride_sbn, stride_sbk, stride_om, stride_pk, stride_pm, + BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, SLAB: tl.constexpr, SPLIT_K: tl.constexpr, + QUANT_ACT: tl.constexpr, OUT: tl.constexpr, ): + """Small-M (decode) block-scaled fp8 GEMV/GEMM: one pass over the weight at near HBM bandwidth. + + Each program owns ``BLOCK_N`` rows and ``K // SPLIT_K`` contiguous K, walked ``BLOCK_K`` (several + ``SLAB``-wide scale slabs) at a time: a wide tile keeps enough bytes in flight per SM to saturate + the memory system, which a slab-at-a-time walk (512 B per load at the V4.1 slab of 32) cannot. + Every weight tile is applied to all ``M`` (<= ``BLOCK_M``) activation rows, so a decode batch + costs one weight read like a single token does (the tensor-core GEMM path is 3-6x slower at + M = 2..8: too few tiles, tiny K steps). ``QUANT_ACT`` quantizes the activation here rather than + by a separate ``act_quant`` launch: every program recomputes the ue8m0 scale and the e4m3 rounding + of the slabs it touches (the reference formula, bit-identical to ``act_quant_fp8``) from the + L2-resident bf16 rows, so the path is one kernel (plus a small reduce when ``SPLIT_K > 1``). + Without it the activation stays bf16 (W8A16: the reference's bf16 einsum over a dequantized + weight, e.g. ``wo_a``). The third grid axis is an independent group (block-diagonal weights: + group ``g`` reads ``x[:, g]`` and rows ``[g * N, (g + 1) * N)``).""" + tl.static_assert(BLOCK_K % SLAB == 0, "the K tile holds whole scale slabs") + NSLAB: tl.constexpr = BLOCK_K // SLAB pid_n = tl.program_id(0) pid_k = tl.program_id(1) + g = tl.program_id(2).to(tl.int64) + x_ptr += g * stride_xg offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) n_mask = offs_n < N - sn = offs_n // 128 + sn = (g * N + offs_n) // SLAB + offs_n = g * N + offs_n + rows = tl.arange(0, BLOCK_M) k_per = K // SPLIT_K k_start = pid_k * k_per - acc = tl.zeros((BLOCK_N,), dtype=tl.float32) - for k0 in range(0, k_per, 128): - offs_k = k_start + k0 + tl.arange(0, 128) - a = tl.load(a_ptr + offs_k * stride_ak).to(tl.float32) + acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) + for k0 in range(0, k_per, BLOCK_K): + offs_k = k_start + k0 + tl.arange(0, BLOCK_K) + # --- the weight tile, dequantized per slab (read once for every activation row) --- w_ptrs = w_ptr + offs_n[:, None] * stride_wn + offs_k[None, :] * stride_wk if e4m3_native_cx(): w = tl.load(w_ptrs, mask=n_mask[:, None], other=0.0).to(tl.float32) else: w = e4m3_u8_to_f32(tl.load(w_ptrs, mask=n_mask[:, None], other=0)) - kb = (k_start + k0) // 128 - sb_code = tl.load(sb_ptr + sn * stride_sbn + kb * stride_sbk, mask=n_mask, other=0) - scb = tl.exp2(sb_code.to(tl.float32) - 127.0) - sa_code = tl.load(sa_ptr + kb) - sca = tl.exp2(sa_code.to(tl.float32) - 127.0) - acc += tl.sum(w * a[None, :], axis=1) * scb * sca - tl.store(part_ptr + pid_k * stride_pk + offs_n * stride_pn, acc, mask=n_mask) + kb = (k_start + k0) // SLAB + tl.arange(0, NSLAB) + sb_code = tl.load(sb_ptr + sn[:, None] * stride_sbn + kb[None, :] * stride_sbk, mask=n_mask[:, None], other=0) + scb = tl.exp2(sb_code.to(tl.float32) - 127.0) # [BLOCK_N, NSLAB] + for m in tl.static_range(BLOCK_M): + if m < M: + xs = tl.reshape(tl.load(x_ptr + m * stride_xm + offs_k).to(tl.float32), (NSLAB, SLAB)) + if QUANT_ACT: + # --- reference act_quant on this row's slabs: s = 2**ceil(log2(max(|x|, 1e-4) / 448)) --- + amax = tl.maximum(tl.max(tl.abs(xs), axis=1), 1e-4) + sa = tl.exp2(_log2_ceil(amax * (1.0 / 448.0)).to(tl.float32)) # [NSLAB] + q = tl.clamp(xs / sa[:, None], -448.0, 448.0) + if e4m3_native_cx(): + q = q.to(tl.float8e4nv).to(tl.float32) + else: + q = round_e4m3(q).to(tl.float32) + else: + sa = tl.full((NSLAB,), 1.0, tl.float32) + q = xs + a = tl.reshape(q, (BLOCK_K,)) + slab = tl.sum(tl.reshape(w * a[None, :], (BLOCK_N, NSLAB, SLAB)), axis=2) # [BLOCK_N, NSLAB] + contrib = tl.sum(slab * scb * sa[None, :], axis=1) # [BLOCK_N] + acc = tl.where(rows[:, None] == m, acc + contrib[None, :], acc) + row_mask = rows < M + if SPLIT_K == 1: + tl.store(out_ptr + rows[:, None] * stride_om + offs_n[None, :], acc.to(OUT), mask=row_mask[:, None] & n_mask[None, :]) + else: + tl.store(part_ptr + pid_k * stride_pk + rows[:, None] * stride_pm + offs_n[None, :], acc, mask=row_mask[:, None] & n_mask[None, :]) @triton.jit -def _splitk_reduce_kernel(part_ptr, out_ptr, N, SPLIT_K: tl.constexpr, - stride_pk, stride_pn, BLOCK: tl.constexpr, OUT: tl.constexpr): +def _splitk_reduce_kernel(part_ptr, out_ptr, M, N, SPLIT_K: tl.constexpr, + stride_pk, stride_pm, stride_om, BLOCK: tl.constexpr, OUT: tl.constexpr): + m = tl.program_id(1) offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) mask = offs < N acc = tl.zeros((BLOCK,), dtype=tl.float32) for k in tl.static_range(SPLIT_K): - acc += tl.load(part_ptr + k * stride_pk + offs * stride_pn, mask=mask, other=0.0) - tl.store(out_ptr + offs, acc.to(OUT), mask=mask) - - -# Swept-best decode GEMV config per (N, K) -> (BLOCK_N, SPLIT_K, num_warps). -_DECODE_FP8_CFG = { - (1024, 4096): (16, 16, 1), # wq_a - (32768, 1024): (16, 1, 2), # wq_b - (512, 4096): (16, 16, 1), # attn wkv - (4096, 8192): (16, 8, 1), # wo_b - (2048, 4096): (16, 8, 1), # shared w1 / w3 - (4096, 2048): (16, 8, 1), # shared w2 - (8192, 1024): (16, 2, 4), # indexer wq_b -} - - -def _decode_cfg(N: int, K: int) -> tuple[int, int, int]: - cfg = _DECODE_FP8_CFG.get((N, K)) - if cfg is not None: - bn, sk, nw = cfg - return bn, max(1, min(sk, K // 128)), nw - bn = 16 - n_tiles = triton.cdiv(N, bn) - sk = max(1, 1536 // n_tiles) - sk = 1 << (sk.bit_length() - 1) - return bn, max(1, min(sk, K // 128)), 1 - - -def _fp8_act_gemv(a_fp8: torch.Tensor, sa: torch.Tensor, weight: torch.Tensor, - sb: torch.Tensor, out_dtype: torch.dtype) -> torch.Tensor: - N, K = weight.shape - BLOCK_N, split_k, num_warps = _decode_cfg(N, K) - n_tiles = triton.cdiv(N, BLOCK_N) - part = torch.empty((split_k, N), dtype=torch.float32, device=a_fp8.device) - _fp8_act_gemv_splitk_kernel[(n_tiles, split_k)]( - a_fp8, sa, weight, sb, part, N, K, - a_fp8.stride(0), weight.stride(0), weight.stride(1), - sb.stride(0), sb.stride(1), part.stride(0), part.stride(1), - BLOCK_N=BLOCK_N, SPLIT_K=split_k, num_warps=num_warps, + acc += tl.load(part_ptr + k * stride_pk + m * stride_pm + offs, mask=mask, other=0.0) + tl.store(out_ptr + m * stride_om + offs, acc.to(OUT), mask=mask) + + +_GEMV_BLOCK_N = 32 +_GEMV_NARROW_BLOCK_N = 8 +def _gemv_cfg(N: int, K: int, slab: int, sm_count: int) -> tuple[int, int, int, int]: + """Choose whole scale tiles and enough independent programs to cover the device.""" + block_k = next((bk for bk in (512, 256, 128) if bk >= slab and K % bk == 0), slab) + if triton.cdiv(N, _GEMV_BLOCK_N) >= sm_count: + return _GEMV_BLOCK_N, block_k, 1, 4 + n_tiles = triton.cdiv(N, _GEMV_NARROW_BLOCK_N) + split = 1 + while n_tiles * split < sm_count and (K // block_k) % (split * 2) == 0: + split *= 2 + return _GEMV_NARROW_BLOCK_N, block_k, split, 4 + + +GEMV_MAX_M = 4 # bound the unrolled row loop's register footprint + + +def _fp8_gemv( + x: torch.Tensor, weight: torch.Tensor, sb: torch.Tensor, out_dtype: torch.dtype, block: int, *, groups: int = 1, quant_act: bool = True, +) -> torch.Tensor: + """``x [M, groups, K]`` bf16 x block-diagonal ``weight [groups * N, K]`` fp8 -> ``[M, groups * N]`` + in ``out_dtype`` (``M <= GEMV_MAX_M``). ``quant_act`` applies the reference fp8 activation quant + (W8A8); off, the activation stays bf16 (W8A16). ``groups == 1`` is the plain GEMV.""" + GN, K = weight.shape + N = GN // groups + x = x.reshape(-1, groups, K) + M = x.shape[0] + assert 1 <= M <= GEMV_MAX_M, M + block_m = triton.next_power_of_2(M) + block_n, block_k, split_k, num_warps = _gemv_cfg(GN, K, block, torch.cuda.get_device_properties(x.device).multi_processor_count) + out = torch.empty((M, GN), dtype=out_dtype, device=x.device) + part = out if split_k == 1 else torch.empty((split_k, M, GN), dtype=torch.float32, device=x.device) + _fp8_gemv_kernel[(triton.cdiv(N, block_n), split_k, groups)]( + x, weight, sb, out, part, M, N, K, + x.stride(0), x.stride(1), weight.stride(0), weight.stride(1), sb.stride(0), sb.stride(1), out.stride(0), + part.stride(0) if split_k > 1 else 0, part.stride(1) if split_k > 1 else 0, + BLOCK_M=block_m, BLOCK_N=block_n, BLOCK_K=block_k, SLAB=block, SPLIT_K=split_k, QUANT_ACT=quant_act, OUT=_TL_DTYPE[out_dtype], + num_warps=num_warps, ) - out = torch.empty(N, dtype=out_dtype, device=a_fp8.device) - _splitk_reduce_kernel[(triton.cdiv(N, 256),)]( - part, out, N, split_k, part.stride(0), part.stride(1), - BLOCK=256, OUT=_TL_DTYPE[out_dtype], num_warps=2, + if split_k > 1: + _splitk_reduce_kernel[(triton.cdiv(GN, 256), M)]( + part, out, M, GN, split_k, part.stride(0), part.stride(1), out.stride(0), + BLOCK=256, OUT=_TL_DTYPE[out_dtype], num_warps=2, + ) + return out + + +def grouped_w8a16_gemv(x: torch.Tensor, weight: torch.Tensor, scale: torch.Tensor, *, block: int) -> torch.Tensor: + """Decode form of the reference's grouped bf16 einsum ``tgd,grd->tgr`` over a block-diagonal fp8 + weight: ``x [t, G, d]`` bf16 (``t <= GEMV_MAX_M``), ``weight [G * r, d]`` fp8 with ``[G * r // + block, d // block]`` e8m0 ``scale`` -> ``[t, G * r]`` bf16. The weight is dequantized in-kernel + (exact in bf16: e4m3 x 2^k), the activation is not quantized, so this is the einsum's math with a + different fp32 summation order.""" + assert x.ndim == 3 and weight.dtype == FP8 + G = x.shape[1] + sb = (scale.view(torch.uint8) if scale.dtype == torch.float8_e8m0fnu else scale).contiguous() + return _fp8_gemv(x.contiguous(), e4m3_kernel_view(weight), sb, x.dtype, block, groups=G, quant_act=False) + + +@triton.jit +def _dequant_block_fp8_kernel(w_ptr, sb_ptr, out_ptr, N, K, stride_wn, stride_sbn, stride_on, SLAB: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr): + pid_n = tl.program_id(0) + pid_k = tl.program_id(1) + offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + offs_k = pid_k * BLOCK_K + tl.arange(0, BLOCK_K) + mask = (offs_n < N)[:, None] & (offs_k < K)[None, :] + if e4m3_native_cx(): + w = tl.load(w_ptr + offs_n[:, None] * stride_wn + offs_k[None, :], mask=mask, other=0.0).to(tl.float32) + else: + w = e4m3_u8_to_f32(tl.load(w_ptr + offs_n[:, None] * stride_wn + offs_k[None, :], mask=mask, other=0)) + code = tl.load(sb_ptr + (offs_n // SLAB)[:, None] * stride_sbn + (offs_k // SLAB)[None, :], mask=mask, other=0) + tl.store(out_ptr + offs_n[:, None] * stride_on + offs_k[None, :], (w * tl.exp2(code.to(tl.float32) - 127.0)).to(tl.bfloat16), mask=mask) + + +def dequant_block_fp8(weight: torch.Tensor, scale: torch.Tensor, *, block: int) -> torch.Tensor: + """``[N, K]`` block-scaled e4m3 -> bf16 (exact), one pass; the prefill-side companion of + :func:`grouped_w8a16_gemv` for weights the reference consumes through a bf16 einsum.""" + N, K = weight.shape + sb = (scale.view(torch.uint8) if scale.dtype == torch.float8_e8m0fnu else scale).contiguous() + out = torch.empty((N, K), dtype=torch.bfloat16, device=weight.device) + _dequant_block_fp8_kernel[(triton.cdiv(N, 64), triton.cdiv(K, 256))]( + e4m3_kernel_view(weight), sb, out, N, K, weight.stride(0), sb.stride(0), out.stride(0), + SLAB=block, BLOCK_N=64, BLOCK_K=256, num_warps=4, ) return out @@ -353,36 +465,41 @@ def block_fp8_linear( weight: torch.Tensor, scale: torch.Tensor, bias: torch.Tensor | None = None, + *, + block: int = 128, ) -> torch.Tensor: """``y = act_quant(x) @ weight^T`` (reference FP8 path). ``x``: ``[..., K]`` bf16; ``weight``: ``[N, K]`` float8_e4m3fn; ``scale``: - ``[N//128, K//128]`` float8_e8m0fnu (weight block scale). Activation is quantized - to FP8 with a per-128 ue8m0 scale; the GEMM applies both scales per 128-K block. + ``[N//block, K//block]`` float8_e8m0fnu (weight block scale). Activation is quantized + to FP8 with a per-``block`` ue8m0 scale; the GEMM applies both scales per ``block``-K slab. """ assert weight.dtype == FP8 *lead, K = x.shape N = weight.shape[0] assert weight.shape[1] == K - assert K % 128 == 0 and N % 128 == 0, (N, K) + assert K % block == 0 and N % block == 0, (N, K, block) + assert tuple(scale.shape) == (N // block, K // block), (tuple(scale.shape), N, K, block) compute_dtype = x.dtype if x.dtype in _TL_DTYPE else torch.bfloat16 sb = scale.view(torch.uint8) if scale.dtype == torch.float8_e8m0fnu else scale sb = sb.contiguous() w = e4m3_kernel_view(weight) - a_fp8, sa = act_quant_fp8(x, 128) # [M,K] fp8, [M,K//128] e8m0 codes - M = a_fp8.shape[0] - - if M == 1: - out = _fp8_act_gemv(a_fp8[0], sa[0], w, sb, compute_dtype).reshape(*lead, N) + x2d = x.reshape(-1, K) + if x2d.shape[0] == 1: + # decode: one kernel quantizes the activation rows and streams the weight once + out = _fp8_gemv(x2d.contiguous(), w, sb, compute_dtype, block).reshape(*lead, N) if bias is not None: out = out + bias.to(out.dtype) return out + a_fp8, sa = act_quant_fp8(x, block) # [M,K] fp8, [M,K//block] e8m0 codes + M = a_fp8.shape[0] + out = torch.empty((M, N), dtype=compute_dtype, device=x.device) BLOCK_M = 32 BLOCK_N = 128 - BLOCK_K = 128 + BLOCK_K = block grid = (triton.cdiv(M, BLOCK_M), triton.cdiv(N, BLOCK_N)) _fp8_act_gemm_kernel[grid]( a_fp8, w, sa, sb, out, @@ -399,4 +516,4 @@ def block_fp8_linear( return out -__all__ = ["block_fp8_linear", "act_quant_fp8", "act_quant_fp8_inplace", "fp4_act_quant_inplace"] +__all__ = ["block_fp8_linear", "act_quant_fp8", "act_quant_fp8_inplace", "fp4_act_quant_inplace", "grouped_w8a16_gemv", "dequant_block_fp8", "GEMV_MAX_M"] diff --git a/python/freetoken/kernel/triton/dsv4/fused_moe.py b/python/freetoken/kernel/triton/dsv4/fused_moe.py index 6eea606c36..0adfbdd34b 100644 --- a/python/freetoken/kernel/triton/dsv4/fused_moe.py +++ b/python/freetoken/kernel/triton/dsv4/fused_moe.py @@ -24,6 +24,9 @@ import triton import triton.language as tl +from freetoken.kernel.triton.dsv4.fp8_linear import _log2_ceil +from freetoken.kernel.triton.e4m3_compat import e4m3_native_cx, round_e4m3 + _E2M1_VALUES = [ 0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0, @@ -71,6 +74,12 @@ def _decode_dsfp4_moe_kernel( n_mask = offs_n < N slot = tl.load(topk_ids_ptr + token_id * stride_tid_m + route_k * stride_tid_k).to(tl.int64) + c_ptrs = c_ptr + token_id * stride_cm + route_k * stride_ck + offs_n * stride_cn + if slot < 0: + # an inactive route (hybrid decode computes it on the CPU): zero contribution, and no read of + # any slot -- a slot's storage may be uninitialized, and 0 * NaN is not 0 + tl.store(c_ptrs, tl.zeros((BLOCK_SIZE_N,), dtype=compute_type), mask=(route_id < total_routes) & n_mask) + return a_row = route_id if A_ROW_IS_ROUTE else token_id a_base = a_ptr + a_row * stride_am @@ -113,7 +122,6 @@ def _decode_dsfp4_moe_kernel( weight = tl.load(topk_weights_ptr + token_id * stride_tw_m + route_k * stride_tw_k) accumulator = accumulator * weight - c_ptrs = c_ptr + token_id * stride_cm + route_k * stride_ck + offs_n * stride_cn tl.store(c_ptrs, accumulator.to(compute_type), mask=(route_id < total_routes) & n_mask) @@ -124,11 +132,15 @@ def _swiglu_kernel( R, I, limit, stride_gr, stride_gi, stride_or, stride_oi, BLOCK: tl.constexpr, HAS_LIMIT: tl.constexpr, compute_type: tl.constexpr, + ACT_BLOCK: tl.constexpr, # > 0: fp8 (e4m3, ue8m0 per ACT_BLOCK) round-trip of the stored value, fused ): """Fused SwiGLU: ``out = silu(min(gate, limit)) * clamp(up, -limit, limit)``. Collapses the gate/up split + two clamps + silu + mul + the fp32 round-trip - (6 elementwise launches over [R, I]) between the two FP4 GEMVs into one pass.""" + (6 elementwise launches over [R, I]) between the two FP4 GEMVs into one pass. With + ``ACT_BLOCK`` the down projection's activation quant (``act_quant_fp8_inplace`` on the bf16 + value: ``s = 2**ceil(log2(max(|x|, 1e-4) / 448))``, e4m3 rounding, dequant) is applied to the + same tile before the store -- one launch and one memory pass fewer, bit-identical.""" row = tl.program_id(0) cb = tl.program_id(1) offs = cb * BLOCK + tl.arange(0, BLOCK) @@ -139,6 +151,17 @@ def _swiglu_kernel( g = tl.minimum(g, limit) u = tl.minimum(tl.maximum(u, -limit), limit) act = (g * tl.sigmoid(g)) * u + if ACT_BLOCK > 0: + tl.static_assert(BLOCK % ACT_BLOCK == 0, "the tile holds whole quant blocks") + xs = tl.reshape(act.to(compute_type).to(tl.float32), (BLOCK // ACT_BLOCK, ACT_BLOCK)) # the value the store would hold + amax = tl.maximum(tl.max(tl.abs(xs), axis=1), 1e-4) + sc = tl.exp2(_log2_ceil(amax * (1.0 / 448.0)).to(tl.float32)) + q = tl.clamp(xs / sc[:, None], -448.0, 448.0) + if e4m3_native_cx(): + q = q.to(tl.float8e4nv).to(tl.float32) + else: + q = round_e4m3(q) + act = tl.reshape(q * sc[:, None], (BLOCK,)) tl.store(out_ptr + row * stride_or + offs * stride_oi, act.to(compute_type), mask=mask) @@ -274,20 +297,22 @@ def _compute_type(dtype: torch.dtype): return {torch.bfloat16: tl.bfloat16, torch.float16: tl.float16, torch.float32: tl.float32}[dtype] -def fused_swiglu(gate_up: torch.Tensor, limit: float) -> torch.Tensor: - """``[..., 2I] -> [..., I]`` SwiGLU in a single kernel (see ``_swiglu_kernel``).""" +def fused_swiglu(gate_up: torch.Tensor, limit: float, *, act_block: int = 0) -> torch.Tensor: + """``[..., 2I] -> [..., I]`` SwiGLU in a single kernel (see ``_swiglu_kernel``). ``act_block > 0`` + also applies the down projection's fp8 activation round-trip (block ``act_block``) in the same pass.""" *lead, two_I = gate_up.shape I = two_I // 2 gu = gate_up.reshape(-1, two_I) R = gu.shape[0] out = torch.empty((R, I), dtype=gate_up.dtype, device=gate_up.device) BLOCK = 1024 + assert act_block == 0 or (I % act_block == 0 and BLOCK % act_block == 0), (I, act_block) grid = (R, triton.cdiv(I, BLOCK)) _swiglu_kernel[grid]( gu, out, R, I, float(limit), gu.stride(0), gu.stride(1), out.stride(0), out.stride(1), BLOCK=BLOCK, HAS_LIMIT=limit > 0, - compute_type=_compute_type(gate_up.dtype), num_warps=4, + compute_type=_compute_type(gate_up.dtype), ACT_BLOCK=act_block, num_warps=4, ) return out.reshape(*lead, I) diff --git a/python/freetoken/kernel/triton/dsv4/router.py b/python/freetoken/kernel/triton/dsv4/router.py new file mode 100644 index 0000000000..17e3506a19 --- /dev/null +++ b/python/freetoken/kernel/triton/dsv4/router.py @@ -0,0 +1,83 @@ +"""Fused MoE gate selection (DeepSeek-V4 / V4.1 ``Gate``): score activation, selection bias, top-k, +weight normalization and route scale in one program per token. + + s = act(scores) act: softmax | sigmoid | sqrt(softplus) + ids = topk(s + bias) selection is biased, the weights are not + w = s[ids] [/ (sum + 1e-20)] * route_scale + +One launch replaces the eight or so elementwise / top-k / gather launches of the torch expression -- +at decode each of those moves a few hundred floats and costs its launch floor. Ties in ``s + bias`` +go to the lowest expert id (torch.topk leaves them unspecified). +""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl +from triton.language.extra import libdevice + +_SCORE = {"softmax": 0, "sigmoid": 1, "sqrtsoftplus": 2} + + +@triton.jit +def _router_select_kernel( + scores_ptr, bias_ptr, alternate_bias_ptr, alternate_mask_ptr, w_ptr, id_ptr, E, route_scale, + stride_s, stride_w, stride_id, + SCORE: tl.constexpr, TOPK: tl.constexpr, BLOCK_K: tl.constexpr, NORMALIZE: tl.constexpr, BLOCK_E: tl.constexpr, ALTERNATE: tl.constexpr, +): + t = tl.program_id(0).to(tl.int64) + offs = tl.arange(0, BLOCK_E) + mask = offs < E + s = tl.load(scores_ptr + t * stride_s + offs, mask=mask, other=float("-inf")).to(tl.float32) + if SCORE == 0: + s = tl.exp(s - tl.max(s)) + s = s / tl.sum(tl.where(mask, s, 0.0)) + elif SCORE == 1: + s = tl.sigmoid(s) + else: + s = tl.sqrt(tl.where(s > 20.0, s, libdevice.log1p(tl.exp(s)))) # torch softplus (threshold 20) + if ALTERNATE: + bias_ptr = tl.where(tl.load(alternate_mask_ptr + t), alternate_bias_ptr, bias_ptr) + bias = tl.load(bias_ptr + offs, mask=mask, other=0.0).to(tl.float32) + key = tl.where(mask, s + bias, float("-inf")) + slots = tl.arange(0, BLOCK_K) + picked_w = tl.zeros((BLOCK_K,), dtype=tl.float32) + picked_id = tl.zeros((BLOCK_K,), dtype=tl.int32) + for k in tl.static_range(TOPK): + best = tl.max(key) + idx = tl.min(tl.where(key == best, offs, BLOCK_E)) # lowest id among ties + w = tl.sum(tl.where(offs == idx, s, 0.0)) + picked_w = tl.where(slots == k, w, picked_w) + picked_id = tl.where(slots == k, idx.to(tl.int32), picked_id) + key = tl.where(offs == idx, float("-inf"), key) + if NORMALIZE: + picked_w = picked_w / (tl.sum(picked_w) + 1e-20) + picked_w = picked_w * route_scale + tl.store(w_ptr + t * stride_w + slots, picked_w, mask=slots < TOPK) + tl.store(id_ptr + t * stride_id + slots, picked_id, mask=slots < TOPK) + + +def router_select( + scores: torch.Tensor, bias: torch.Tensor, topk: int, *, score_func: str, normalize: bool, route_scale: float, + alternate_bias: torch.Tensor | None = None, alternate_mask: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """``scores [T, E]`` fp32 (pre-activation gate logits), ``bias [E]`` -> ``(weights [T, topk] fp32, + ids [T, topk] int32)`` in selection order (best first).""" + T, E = scores.shape + assert scores.dtype == torch.float32 and scores.stride(1) == 1, (scores.dtype, scores.stride()) + weights = torch.empty((T, topk), dtype=torch.float32, device=scores.device) + ids = torch.empty((T, topk), dtype=torch.int32, device=scores.device) + if T == 0: + return weights, ids + _router_select_kernel[(T,)]( + scores, bias, alternate_bias if alternate_bias is not None else bias, + alternate_mask if alternate_mask is not None else scores, weights, ids, E, float(route_scale), + scores.stride(0), weights.stride(0), ids.stride(0), + SCORE=_SCORE[score_func], TOPK=topk, BLOCK_K=triton.next_power_of_2(topk), NORMALIZE=normalize and topk > 1, + BLOCK_E=triton.next_power_of_2(E), ALTERNATE=alternate_mask is not None, num_warps=4, + ) + return weights, ids + + +__all__ = ["router_select"] diff --git a/python/freetoken/kernel/triton/dsv41/__init__.py b/python/freetoken/kernel/triton/dsv41/__init__.py new file mode 100644 index 0000000000..00ec72daba --- /dev/null +++ b/python/freetoken/kernel/triton/dsv41/__init__.py @@ -0,0 +1,9 @@ +"""Triton kernels for DeepSeek-V4.1 Compressed Sparse Attention 2 over packed KV pools. + +``row_format`` device helpers: load a tile of packed rows (bf16 / fp8+e8m0 / fp4+e4m3 / fp4+e8m0) + and dequantize it to fp32, encode / decode e2m1 nibbles. +``pack`` host wrappers: quantize rows into a packed pool (``pack_rows``) and read them back + (``unpack_rows``), matching the reference ``act_quant`` / ``fp4_act_quant`` numerics. +``sparse_attn`` paged sparse attention gathering window (fp8) and compressed (fp4) rows. +``indexer`` Lightning-indexer logits over fp4 index keys, full-range or candidate-restricted. +""" diff --git a/python/freetoken/kernel/triton/dsv41/indexer.py b/python/freetoken/kernel/triton/dsv41/indexer.py new file mode 100644 index 0000000000..3721315837 --- /dev/null +++ b/python/freetoken/kernel/triton/dsv41/indexer.py @@ -0,0 +1,176 @@ +"""Lightning-Indexer logits over PACKED index keys (DeepSeek-V4.1, Full and Reindex modes). + + logits[b, s, j] = sum_h relu(q[b, s, h, :] . k(b, pos_j)) * weights[b, s, h] + +``q`` carries the ``index_n_heads`` fp4-baked index queries, ``k`` is the single shared index key of +a compressed position, gathered from the kv-source's packed pool and dequantized in-kernel. The +row of compressed position ``t`` is derived in-kernel from the request's full-token locs (the page +table row at prefill, the decode snapshot): ``locs[b, t * ratio] // ratio`` (negative = no row) -- +no ``[B, T]`` row table is materialized per layer. ``weights`` folds ``softmax_scale * +n_heads**-0.5`` in. + +Two position sets, one kernel: + +* **full range** (Full mode): position ``j`` is the compressed position ``j``; the output is + ``[B, S, T]`` over the staged width ``T``, but only the live tiles (columns below ``live[b, s]`` + rounded up to ``BLOCK_T``, the rest of the last tile ``-inf``) are WRITTEN. A fixed grid of worker + programs strides over the live tiles, so both the work and the launch geometry follow the live + history rather than the staged width; consumers (``dsv41/topk.py``) read below ``live`` only, so + the dead area needs no fill. +* **candidate pool** (Reindex mode, Hierarchical Sparse Indexer): position ``j`` is + ``candidates[b, s, j]`` (a compressed position, ``-1`` = empty slot); the output is ``[B, S, NC]`` + aligned with the candidate list and fully written (``-inf`` for empty / unreachable slots), so + a top-k over it indexes straight back into the list. + +Causality lives in the kernel: ``live[b, s]`` is the number of compressed positions query ``s`` may +see, and every position at or past it (or with no row) scores ``-inf``. The dot product, weighted +scores and head reduction round to bf16 at the same boundaries as the reference expression. +""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + +from freetoken.kvcache.dsv4.v41_row_format import RowFormat + +from .row_format import load_rows + +BLOCK_T = 64 + + +@triton.jit +def _score_tile( + q, w, pool_ptr, locs_row, live, pos, store_mask, out_ptrs, + n_locs, RATIO: tl.constexpr, D: tl.constexpr, ROW_BYTES: tl.constexpr, FMT: tl.constexpr, +): + """Score one tile of compressed positions ``pos`` against the query (``q [BLOCK_H, D]`` bf16, + ``w [BLOCK_H]`` fp32) and store ``[BLOCK_T]`` logits (``-inf`` where unreachable).""" + pos_ok = store_mask & (pos >= 0) & (pos < live) + full_idx = pos * RATIO + full = tl.load(locs_row + tl.maximum(full_idx, 0), mask=pos_ok & (full_idx < n_locs), other=-1) + rows = tl.where(full >= 0, full // RATIO, -1).to(tl.int32) + valid = pos_ok & (rows >= 0) + offs_d = tl.arange(0, D) + k = load_rows(pool_ptr, rows, valid, offs_d, D, ROW_BYTES, FMT).to(tl.bfloat16) # [BLOCK_T, D] + score = tl.dot(q, tl.trans(k)).to(tl.bfloat16).to(tl.float32) + score = (tl.maximum(score, 0.0) * w[:, None]).to(tl.bfloat16).to(tl.float32) + logits = tl.sum(score, axis=0).to(tl.bfloat16).to(tl.float32) + tl.store(out_ptrs, tl.where(valid, logits, float("-inf")), mask=store_mask) + + +@triton.jit +def _indexer_logits_packed_kernel( + q_ptr, w_ptr, pool_ptr, locs_ptr, live_ptr, cand_ptr, out_ptr, + T, NC, n_locs, + stride_qb, stride_qs, stride_qh, stride_qd, + stride_wb, stride_ws, stride_wh, + stride_lb, stride_lt, + stride_vb, stride_vs, + stride_cb, stride_cs, stride_cn, + stride_ob, stride_os, stride_on, + RATIO: tl.constexpr, + H: tl.constexpr, + D: tl.constexpr, + ROW_BYTES: tl.constexpr, + FMT: tl.constexpr, + BLOCK_H: tl.constexpr, + BLOCK_T: tl.constexpr, + HAS_CAND: tl.constexpr, + N_WORKERS: tl.constexpr, +): + pid_s = tl.program_id(0) + pid_b = tl.program_id(1) + pid_t = tl.program_id(2) + + live = tl.load(live_ptr + pid_b * stride_vb + pid_s * stride_vs) + locs_row = locs_ptr + pid_b * stride_lb + out_row = out_ptr + pid_b * stride_ob + pid_s * stride_os + offs_d = tl.arange(0, D) + offs_h = tl.arange(0, BLOCK_H) + h_mask = offs_h < H + q = tl.load( + q_ptr + pid_b * stride_qb + pid_s * stride_qs + offs_h[:, None] * stride_qh + offs_d[None, :] * stride_qd, + mask=h_mask[:, None], other=0.0, + ).to(tl.bfloat16) + w = tl.load(w_ptr + pid_b * stride_wb + pid_s * stride_ws + offs_h * stride_wh, mask=h_mask, other=0.0).to(tl.bfloat16).to(tl.float32) + offs = tl.arange(0, BLOCK_T) + + if HAS_CAND: + # one program per candidate tile; the list is fixed-width and fully written + offs_t = pid_t * BLOCK_T + offs + store_mask = offs_t < NC + pos = tl.load(cand_ptr + pid_b * stride_cb + pid_s * stride_cs + offs_t * stride_cn, mask=store_mask, other=-1) + _score_tile(q, w, pool_ptr, locs_row, live, pos, store_mask, out_row + offs_t * stride_on, n_locs, RATIO, D, ROW_BYTES, FMT) + else: + # worker pid_t scores live tiles pid_t, pid_t + N_WORKERS, ...; nothing past live is written + for tile in range(pid_t, tl.cdiv(tl.minimum(live, T), BLOCK_T), N_WORKERS): + offs_t = tile * BLOCK_T + offs + store_mask = offs_t < T + _score_tile(q, w, pool_ptr, locs_row, live, offs_t, store_mask, out_row + offs_t * stride_on, n_locs, RATIO, D, ROW_BYTES, FMT) + + +def indexer_logits_packed( + q: torch.Tensor, # [B, S, H, D] bf16 (fp4 round-tripped index queries) + weights: torch.Tensor, # [B, S, H] (softmax_scale * n_heads**-0.5 folded in) + k_pool: torch.Tensor, # [R, fmt.row_bytes(D)] uint8, the kv-source's packed index keys + fmt: RowFormat, + locs: torch.Tensor, # [B, W] int32/int64 full-token locs of each request (page-table row / decode snapshot) + ratio: int, # compressed position t -> full loc locs[b, t * ratio]; its row is loc // ratio + live: torch.Tensor, # [B, S] int32: compressed positions query (b, s) may see + T: int | None = None, # staged width of the full-range output (default: W // ratio) + candidates: torch.Tensor | None = None, # [B, S, NC] int32 compressed positions (-1 = empty) + out: torch.Tensor | None = None, +) -> torch.Tensor: + """Head-reduced indexer logits: ``[B, S, T]`` over the full staged range (columns below ``live`` + written, the rest untouched), or ``[B, S, NC]`` aligned with ``candidates`` (fully written). + Unreachable / empty positions are ``-inf``.""" + B, S, H, D = q.shape + assert locs.ndim == 2 and locs.shape[0] == B and live.shape == (B, S), (locs.shape, live.shape, (B, S)) + assert weights.shape == (B, S, H), (weights.shape, (B, S, H)) + assert D == triton.next_power_of_2(D), f"index_head_dim must be pow2, got {D}" + assert k_pool.dtype == torch.uint8 and k_pool.shape[1] == fmt.row_bytes(D) and k_pool.is_contiguous() + assert locs.dtype in (torch.int32, torch.int64) and locs.stride(1) == 1, (locs.dtype, locs.stride()) + n_locs = locs.shape[1] + if T is None: + T = n_locs // ratio + live = live.to(torch.int32) + has_cand = candidates is not None + if has_cand: + assert candidates.shape[:2] == (B, S), (candidates.shape, (B, S)) + cand = candidates.to(torch.int32) + NC = cand.shape[2] + cs = cand.stride() + else: + cand, NC, cs = live, 0, (0, 0, 0) + n_out = NC if has_cand else T + if out is None: + out = torch.empty((B, S, n_out), dtype=torch.float32, device=q.device) + assert out.shape == (B, S, n_out) and out.dtype == torch.float32 + if n_out == 0 or S == 0: + return out + + n_tiles = triton.cdiv(n_out, BLOCK_T) + # A device-sized worker grid strides over live history without launching a + # program for every reserved tile. Independent queries share the SM budget. + sm_count = torch.cuda.get_device_properties(q.device).multi_processor_count + workers = n_tiles if has_cand else min(n_tiles, max(1, triton.cdiv(sm_count, B * S))) + grid = (S, B, workers) + _indexer_logits_packed_kernel[grid]( + q, weights, k_pool, locs, live, cand, out, + T, NC, n_locs, + q.stride(0), q.stride(1), q.stride(2), q.stride(3), + weights.stride(0), weights.stride(1), weights.stride(2), + locs.stride(0), locs.stride(1), + live.stride(0), live.stride(1), + cs[0], cs[1], cs[2], + out.stride(0), out.stride(1), out.stride(2), + RATIO=ratio, H=H, D=D, ROW_BYTES=fmt.row_bytes(D), FMT=fmt.code, + BLOCK_H=triton.next_power_of_2(H), BLOCK_T=BLOCK_T, HAS_CAND=has_cand, N_WORKERS=workers, + num_warps=4, num_stages=2, + ) + return out + + +__all__ = ["indexer_logits_packed", "BLOCK_T"] diff --git a/python/freetoken/kernel/triton/dsv41/pack.py b/python/freetoken/kernel/triton/dsv41/pack.py new file mode 100644 index 0000000000..47a9164e6e --- /dev/null +++ b/python/freetoken/kernel/triton/dsv41/pack.py @@ -0,0 +1,183 @@ +"""Quantize rows into a packed pool and read them back (DeepSeek-V4.1 KV storage formats). + +``pack_rows`` reproduces the reference quantizers (``inference/kernel.py``) exactly: + +* fp8_e8m0_b32 ``act_quant(x, 32, "ue8m0", e8m0)``: per 32-group ``s = 2**ceil(log2(max(amax, 1e-4) / 448))``, + ``q = round_e4m3(clamp(x / s, +-448))`` +* fp4_e4m3_b16 ``fp4_act_quant(x, 16, scale_dtype=e4m3)``: ``s = e4m3(max(amax, 6 * 2^-9) / 6)`` (round to + nearest e4m3), ``q = round_e2m1(clamp(x / s, +-6))`` +* fp4_e8m0_b32 ``fp4_act_quant(x, 32)``: ``s = 2**ceil(log2(max(amax, 6 * 2^-126) / 6))``, same e2m1 rounding +* bf16 a plain byte copy + +Rows land at ``pool[row_ids[m]]`` (or ``pool[m]`` without ids), so the same launch serves the KV write +path (scatter into slots) and plain buffer conversion. ``unpack_rows`` is the inverse gather + dequant +(bf16-rounded like the reference cache); the attention / indexer kernels use the same device helper +(``row_format.load_rows``) in-line instead of materializing rows. +""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + +from freetoken.kernel.triton.dsv4.fp8_linear import _log2_ceil, _round_fp4 +from freetoken.kernel.triton.e4m3_compat import e4m3_f32_to_u8, e4m3_native_cx, round_e4m3 +from freetoken.kvcache.dsv4.v41_row_format import BF16, RowFormat + +from .row_format import fp4_f32_to_nibble, load_rows + + +@triton.jit +def _e4m3_codes(q): + """fp32 (|q| <= 448) -> e4m3 bit pattern, RNE.""" + if e4m3_native_cx(): + return q.to(tl.float8e4nv).to(tl.uint8, bitcast=True) + else: + return e4m3_f32_to_u8(round_e4m3(q)) + + +@triton.jit +def _e4m3_round(q): + """fp32 -> nearest e4m3 value, as fp32.""" + if e4m3_native_cx(): + return q.to(tl.float8e4nv).to(tl.float32) + else: + return round_e4m3(q) + + +@triton.jit +def _pack_rows_kernel( + x_ptr, out_ptr, idx_ptr, stride_xm, + D: tl.constexpr, GROUP: tl.constexpr, ROW_BYTES: tl.constexpr, FMT: tl.constexpr, HAS_IDX: tl.constexpr, +): + """One program per input row ``m``; writes packed bytes to pool row ``idx[m]`` (or ``m``).""" + m = tl.program_id(0).to(tl.int64) + if HAS_IDX: + row = tl.load(idx_ptr + m).to(tl.int64) + else: + row = m + dst = out_ptr + row * ROW_BYTES + src = x_ptr + m * stride_xm + + if FMT == 0: + offs_d = tl.arange(0, D) + bits = tl.load(src + offs_d).to(tl.bfloat16).to(tl.uint16, bitcast=True) + tl.store(dst + 2 * offs_d, (bits & 0xFF).to(tl.uint8)) + tl.store(dst + 2 * offs_d + 1, (bits >> 8).to(tl.uint8)) + elif FMT == 1: + NG: tl.constexpr = D // GROUP + offs_d = tl.arange(0, D) + offs_g = tl.arange(0, NG) + x = tl.reshape(tl.load(src + offs_d).to(tl.float32), [NG, GROUP]) + amax = tl.maximum(tl.max(tl.abs(x), axis=1), 1e-4) + e = _log2_ceil(amax * (1.0 / 448.0)) + s = tl.exp2(e.to(tl.float32)) + q = tl.clamp(tl.math.div_rn(x, s[:, None]), -448.0, 448.0) + tl.store(dst + offs_d, tl.reshape(_e4m3_codes(q), [D])) + tl.store(dst + D + offs_g, (e + 127).to(tl.uint8)) + else: + # e2m1 pairs: even channels from one strided load, odd from another, so a byte is (lo | hi << 4) + NG: tl.constexpr = D // GROUP + HALF: tl.constexpr = GROUP // 2 + offs_h = tl.arange(0, D // 2) + offs_g = tl.arange(0, NG) + xe = tl.reshape(tl.load(src + 2 * offs_h).to(tl.float32), [NG, HALF]) + xo = tl.reshape(tl.load(src + 2 * offs_h + 1).to(tl.float32), [NG, HALF]) + amax = tl.maximum(tl.max(tl.abs(xe), axis=1), tl.max(tl.abs(xo), axis=1)) + if FMT == 2: + amax = tl.maximum(amax, 6.0 * (2.0 ** -9)) + s = _e4m3_round(amax * (1.0 / 6.0)) + tl.store(dst + D // 2 + offs_g, _e4m3_codes(s)) + else: + amax = tl.maximum(amax, 6.0 * (2.0 ** -126)) + e = _log2_ceil(amax * (1.0 / 6.0)) + s = tl.exp2(e.to(tl.float32)) + tl.store(dst + D // 2 + offs_g, (e + 127).to(tl.uint8)) + # IEEE division: the reference's CUDA ``x / s`` is correctly rounded, and with an e4m3 + # (non power-of-two) scale triton's default approximate div would land a hair off the + # e2m1 midpoints and flip their round-to-even + qe = fp4_f32_to_nibble(_round_fp4(tl.clamp(tl.math.div_rn(xe, s[:, None]), -6.0, 6.0))) + qo = fp4_f32_to_nibble(_round_fp4(tl.clamp(tl.math.div_rn(xo, s[:, None]), -6.0, 6.0))) + tl.store(dst + offs_h, tl.reshape(qe | (qo << 4), [D // 2])) + + +@triton.jit +def _unpack_rows_kernel( + pool_ptr, idx_ptr, out_ptr, M, + D: tl.constexpr, ROW_BYTES: tl.constexpr, FMT: tl.constexpr, BLOCK_T: tl.constexpr, HAS_IDX: tl.constexpr, +): + pid = tl.program_id(0) + offs_t = pid * BLOCK_T + tl.arange(0, BLOCK_T) + valid = offs_t < M + if HAS_IDX: + rows = tl.load(idx_ptr + offs_t, mask=valid, other=0) + else: + rows = offs_t + offs_d = tl.arange(0, D) + vals = load_rows(pool_ptr, rows, valid, offs_d, D, ROW_BYTES, FMT) + tl.store( + out_ptr + offs_t[:, None].to(tl.int64) * D + offs_d[None, :], + vals.to(out_ptr.dtype.element_ty), mask=valid[:, None], + ) + + +def pack_rows( + x: torch.Tensor, fmt: RowFormat, pool: torch.Tensor | None = None, row_ids: torch.Tensor | None = None, +) -> torch.Tensor: + """Quantize ``x [M, D]`` (bf16 / fp16 / fp32) into packed rows. + + With ``pool [rows, row_bytes] uint8`` and ``row_ids [M]``, row ``m`` is written to ``pool[row_ids[m]]`` + (the KV write path). Without a pool a fresh ``[M, row_bytes]`` uint8 tensor is returned. + """ + assert x.dim() == 2, x.shape + m, d = x.shape + fmt.validate_dim(d) + row_bytes = fmt.row_bytes(d) + assert d & (d - 1) == 0 and d >= 16, f"pack_rows wants a power-of-two D >= 16, got {d}" + if pool is None: + assert row_ids is None + pool = torch.empty((m, row_bytes), dtype=torch.uint8, device=x.device) + else: + assert pool.dtype == torch.uint8 and pool.shape[1] == row_bytes and pool.is_contiguous(), (pool.shape, row_bytes) + assert row_ids is not None and row_ids.shape == (m,), (None if row_ids is None else row_ids.shape, m) + if m == 0: + return pool + if x.stride(1) != 1: + x = x.contiguous() + _pack_rows_kernel[(m,)]( + x, pool, row_ids if row_ids is not None else pool, x.stride(0), + D=d, GROUP=fmt.group or d, ROW_BYTES=row_bytes, FMT=fmt.code, HAS_IDX=row_ids is not None, + num_warps=4, + ) + return pool + + +def unpack_rows( + pool: torch.Tensor, fmt: RowFormat, dim: int, row_ids: torch.Tensor | None = None, + out_dtype: torch.dtype = torch.bfloat16, +) -> torch.Tensor: + """Gather ``pool[row_ids]`` (or every row) and dequantize to ``[M, dim]`` in ``out_dtype``.""" + fmt.validate_dim(dim) + assert pool.dtype == torch.uint8 and pool.shape[1] == fmt.row_bytes(dim) and pool.is_contiguous(), pool.shape + m = pool.shape[0] if row_ids is None else row_ids.numel() + out = torch.empty((m, dim), dtype=out_dtype, device=pool.device) + if m == 0: + return out + block_t = 16 + _unpack_rows_kernel[(triton.cdiv(m, block_t),)]( + pool, row_ids if row_ids is not None else pool, out, m, + D=dim, ROW_BYTES=fmt.row_bytes(dim), FMT=fmt.code, BLOCK_T=block_t, HAS_IDX=row_ids is not None, + num_warps=4, + ) + return out + + +def quant_roundtrip(x: torch.Tensor, fmt: RowFormat) -> torch.Tensor: + """``unpack(pack(x))`` -- the value the cache will hand back; the reference's ``inplace=True`` semantics.""" + if fmt is BF16: + return x.to(torch.bfloat16) + return unpack_rows(pack_rows(x, fmt), fmt, x.shape[1], out_dtype=torch.bfloat16) + + +__all__ = ["pack_rows", "unpack_rows", "quant_roundtrip"] diff --git a/python/freetoken/kernel/triton/dsv41/row_format.py b/python/freetoken/kernel/triton/dsv41/row_format.py new file mode 100644 index 0000000000..032d77e970 --- /dev/null +++ b/python/freetoken/kernel/triton/dsv41/row_format.py @@ -0,0 +1,92 @@ +"""Device-side row-format helpers shared by the DSV41 kernels. + +The Python-side descriptors live in ``freetoken.kvcache.dsv4.v41_row_format``; the ``FMT`` constexpr a +kernel receives is ``RowFormat.code``: + + 0 bf16 2 B / channel, no scale + 1 fp8_e8m0_b32 e4m3 byte / channel, ue8m0 scale byte per 32 channels + 2 fp4_e4m3_b16 e2m1 nibble / channel (even channel low nibble), e4m3 scale byte per 16 + 3 fp4_e8m0_b32 e2m1 nibble / channel, ue8m0 scale byte per 32 + +A row is ``[values | scales]``; ``ROW_BYTES`` is the pool's row stride. Every dequantized value +is rounded through bf16 before use -- the reference bakes its quant round-trip into a bf16 cache, +so this is what its attention / indexer kernels actually read. +""" + +from __future__ import annotations + +import triton +import triton.language as tl + +from freetoken.kernel.triton.e4m3_compat import e4m3_u8_to_f32 + +FMT_BF16 = 0 +FMT_FP8_E8M0_B32 = 1 +FMT_FP4_E4M3_B16 = 2 +FMT_FP4_E8M0_B32 = 3 + + +@triton.jit +def fp4_nibble_to_f32(nib): + """e2m1 code (uint8, 0..15) -> fp32: bit 3 sign, bits 2-1 exponent, bit 0 mantissa; + codes 0/1 are the subnormals 0 and 0.5, the rest ``(1 + m/2) * 2^(e-1)``.""" + c = (nib & 7).to(tl.int32) + neg = (nib & 8) != 0 + sub = c.to(tl.float32) * 0.5 + normal = (1.0 + 0.5 * (c & 1).to(tl.float32)) * tl.exp2(((c >> 1) - 1).to(tl.float32)) + mag = tl.where(c < 2, sub, normal) + return tl.where(neg, -mag, mag) + + +@triton.jit +def fp4_f32_to_nibble(q): + """fp32 value already on the e2m1 grid ({0, .5, 1, 1.5, 2, 3, 4, 6}, signed) -> e2m1 code. + Thresholds sit between grid points, so a value a hair off the grid still maps to its point.""" + a = tl.abs(q) + c = tl.where( + a < 1.75, + tl.where(a < 0.75, tl.where(a < 0.25, 0, 1), tl.where(a < 1.25, 2, 3)), + tl.where(a < 3.5, tl.where(a < 2.5, 4, 5), tl.where(a < 5.0, 6, 7)), + ) + return (c + tl.where(q < 0, 8, 0)).to(tl.uint8) + + +@triton.jit +def load_rows(pool_ptr, rows, valid, offs_d, D: tl.constexpr, ROW_BYTES: tl.constexpr, FMT: tl.constexpr): + """Gather + dequantize ``[T, D]`` fp32 from a packed pool. + + ``pool_ptr`` uint8 pool base; ``rows`` [T] row ids (any int width; masked-out rows may be -1); + ``valid`` [T] bool; ``offs_d`` [D] = ``tl.arange(0, D)``. Invalid rows read as 0. + """ + row_i64 = tl.maximum(rows, 0).to(tl.int64) + base = pool_ptr + row_i64[:, None] * ROW_BYTES + m = valid[:, None] & (offs_d >= 0)[None, :] + if FMT == 0: + lo = tl.load(base + 2 * offs_d[None, :], mask=m, other=0).to(tl.uint16) + hi = tl.load(base + 2 * offs_d[None, :] + 1, mask=m, other=0).to(tl.uint16) + val = (lo | (hi << 8)).to(tl.bfloat16, bitcast=True).to(tl.float32) + elif FMT == 1: + v = tl.load(base + offs_d[None, :], mask=m, other=0) + s = tl.load(base + D + (offs_d // 32)[None, :], mask=m, other=127) + val = e4m3_u8_to_f32(v) * tl.exp2(s.to(tl.float32) - 127.0) + else: + b = tl.load(base + (offs_d // 2)[None, :], mask=m, other=0) + nib = tl.where((offs_d % 2 == 0)[None, :], b & 0xF, b >> 4) + q = fp4_nibble_to_f32(nib) + if FMT == 2: + s = e4m3_u8_to_f32(tl.load(base + D // 2 + (offs_d // 16)[None, :], mask=m, other=0)) + else: + s = tl.exp2(tl.load(base + D // 2 + (offs_d // 32)[None, :], mask=m, other=127).to(tl.float32) - 127.0) + val = q * s + return val.to(tl.bfloat16).to(tl.float32) + + +__all__ = [ + "FMT_BF16", + "FMT_FP8_E8M0_B32", + "FMT_FP4_E4M3_B16", + "FMT_FP4_E8M0_B32", + "fp4_nibble_to_f32", + "fp4_f32_to_nibble", + "load_rows", +] diff --git a/python/freetoken/kernel/triton/dsv41/sparse_attn.py b/python/freetoken/kernel/triton/dsv41/sparse_attn.py new file mode 100644 index 0000000000..9d6b0816db --- /dev/null +++ b/python/freetoken/kernel/triton/dsv41/sparse_attn.py @@ -0,0 +1,323 @@ +"""Sparse gathered-KV flash attention over PACKED pools (DeepSeek-V4.1). + +Same contract as ``kernel/triton/dsv4/sparse_attn.py`` -- each query attends to a per-query list of +GLOBAL slots laid out ``[window part | compressed part]`` plus a per-head attention sink, K == V is one +shared latent head -- but the two pools are byte pools in (possibly different) ``RowFormat``s: the +window ring holds fp8 rows, the compressed KV holds fp4 rows, and every gathered row is dequantized +in-kernel (``row_format.load_rows``) to the bf16 value the reference cache would have held. + +Because the two pools dequantize differently, a tile never mixes them: the window part is +``n_window`` columns with ``n_window % BLOCK_T == 0`` (the caller pads it with ``-1``), so the kernel +walks the window tiles and then the compressed tiles in two loops, each with a single load site. +Both dots take bf16 operands with fp32 accumulation, like the reference ``sparse_attn`` kernel. + +Two implementations, picked from the launch shape (no knob): + * prefill (``m`` > 1): one program per (query, request, head block) + * decode (``m`` == 1): flash-decoding split over the candidate axis, merged through log-sum-exp + (the sink joins in the merge). Not bit-identical to the single-program path. + +``cmp_counts`` (device ``[b, m]`` int32) bounds the compressed columns a CUDA-graph replay visits; +without it the full ``topk_idxs`` width is walked. +""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + +from freetoken.kvcache.dsv4.v41_row_format import RowFormat + +from .row_format import load_rows + +BLOCK_H = 16 +BLOCK_T = 32 +MAX_SPLITS = 32 +# Decode splits at tile boundaries, bounded by the device's SM count below. +MIN_TILES_PER_SPLIT = 1 + + +@triton.jit +def _attend_tile(q, kv, valid, scale, m_i, l_i, acc): + """One online-softmax step over a ``[BLOCK_T, D]`` fp32 KV tile (bf16-exact values).""" + kv16 = kv.to(tl.bfloat16) + scores = tl.dot(q, tl.trans(kv16)) * scale # [BLOCK_H, BLOCK_T] fp32 + scores = tl.where(valid[None, :], scores, -float("inf")) + m_new = tl.maximum(m_i, tl.max(scores, axis=1)) + # an all-masked tile keeps m_new at -inf; short-circuit so -inf - -inf never forms + alpha = tl.where(m_new == -float("inf"), 1.0, tl.exp(m_i - m_new)) + p = tl.where(valid[None, :], tl.exp(scores - m_new[:, None]), 0.0) + l_i = l_i * alpha + tl.sum(p, axis=1) + acc = acc * alpha[:, None] + tl.dot(p.to(tl.bfloat16), kv16) + return m_new, l_i, acc + + +@triton.jit +def _sparse_attn_packed_kernel( + q_ptr, win_ptr, cmp_ptr, o_ptr, sink_ptr, idx_ptr, cnt_ptr, + scale, + H, TOPK, N_WINDOW, + stride_qb, stride_qm, stride_qh, stride_qd, + stride_ob, stride_om, stride_oh, stride_od, + stride_ib, stride_im, stride_it, + stride_nb, stride_nm, + D: tl.constexpr, + WIN_FMT: tl.constexpr, WIN_ROW_BYTES: tl.constexpr, + CMP_FMT: tl.constexpr, CMP_ROW_BYTES: tl.constexpr, + BLOCK_H: tl.constexpr, + BLOCK_T: tl.constexpr, + HAS_COUNTS: tl.constexpr, +): + pid_m = tl.program_id(0) + pid_b = tl.program_id(1) + pid_h = tl.program_id(2) + + offs_h = pid_h * BLOCK_H + tl.arange(0, BLOCK_H) + h_mask = offs_h < H + offs_d = tl.arange(0, D) + + q_ptrs = q_ptr + pid_b * stride_qb + pid_m * stride_qm + offs_h[:, None] * stride_qh + offs_d[None, :] * stride_qd + q = tl.load(q_ptrs, mask=h_mask[:, None], other=0.0).to(tl.bfloat16) # [BLOCK_H, D] + + m_i = tl.full((BLOCK_H,), -float("inf"), dtype=tl.float32) + l_i = tl.zeros((BLOCK_H,), dtype=tl.float32) + acc = tl.zeros((BLOCK_H, D), dtype=tl.float32) + + n_active = TOPK + if HAS_COUNTS: + n_active = N_WINDOW + tl.load(cnt_ptr + pid_b * stride_nb + pid_m * stride_nm) + + idx_base = idx_ptr + pid_b * stride_ib + pid_m * stride_im + for start in range(0, N_WINDOW, BLOCK_T): + offs_t = start + tl.arange(0, BLOCK_T) + idxs = tl.load(idx_base + offs_t * stride_it) + valid = idxs >= 0 + kv = load_rows(win_ptr, idxs, valid, offs_d, D, WIN_ROW_BYTES, WIN_FMT) + m_i, l_i, acc = _attend_tile(q, kv, valid, scale, m_i, l_i, acc) + for start in range(N_WINDOW, n_active, BLOCK_T): + offs_t = start + tl.arange(0, BLOCK_T) + t_mask = offs_t < n_active + idxs = tl.load(idx_base + offs_t * stride_it, mask=t_mask, other=-1) + valid = idxs >= 0 + kv = load_rows(cmp_ptr, idxs, valid, offs_d, D, CMP_ROW_BYTES, CMP_FMT) + m_i, l_i, acc = _attend_tile(q, kv, valid, scale, m_i, l_i, acc) + + sink = tl.load(sink_ptr + offs_h, mask=h_mask, other=0.0).to(tl.float32) + l_i = l_i + tl.exp(sink - m_i) + o = acc / l_i[:, None] + + o_ptrs = o_ptr + pid_b * stride_ob + pid_m * stride_om + offs_h[:, None] * stride_oh + offs_d[None, :] * stride_od + tl.store(o_ptrs, o.to(o_ptr.dtype.element_ty), mask=h_mask[:, None]) + + +@triton.jit +def _sparse_attn_packed_splitk_kernel( + q_ptr, win_ptr, cmp_ptr, mid_o_ptr, mid_lse_ptr, idx_ptr, cnt_ptr, + scale, + H, TOPK, N_WINDOW, + stride_qb, stride_qm, stride_qh, stride_qd, + stride_mb, stride_mm, stride_mh, stride_ms, stride_md, + stride_lb, stride_lm, stride_lh, stride_ls, + stride_ib, stride_im, stride_it, + stride_nb, stride_nm, + D: tl.constexpr, + WIN_FMT: tl.constexpr, WIN_ROW_BYTES: tl.constexpr, + CMP_FMT: tl.constexpr, CMP_ROW_BYTES: tl.constexpr, + BLOCK_H: tl.constexpr, + BLOCK_T: tl.constexpr, + HAS_COUNTS: tl.constexpr, + NUM_SPLITS: tl.constexpr, +): + """Stage 1: each program reduces one BLOCK_T-aligned slice of the candidate list and writes + its normalized partial output + log-sum-exp. A slice may straddle the window / compressed + boundary; it is walked as (window tiles, compressed tiles) since N_WINDOW is tile-aligned.""" + pid_ms = tl.program_id(0) + pid_m = pid_ms // NUM_SPLITS + split_id = pid_ms % NUM_SPLITS + pid_b = tl.program_id(1) + pid_h = tl.program_id(2) + + offs_h = pid_h * BLOCK_H + tl.arange(0, BLOCK_H) + h_mask = offs_h < H + offs_d = tl.arange(0, D) + + n_active = TOPK + if HAS_COUNTS: + n_active = N_WINDOW + tl.load(cnt_ptr + pid_b * stride_nb + pid_m * stride_nm) + + per_split = tl.cdiv(tl.cdiv(n_active, NUM_SPLITS), BLOCK_T) * BLOCK_T + split_start = per_split * split_id + split_end = tl.minimum(split_start + per_split, n_active) + + m_i = tl.full((BLOCK_H,), -float("inf"), dtype=tl.float32) + l_i = tl.zeros((BLOCK_H,), dtype=tl.float32) + acc = tl.zeros((BLOCK_H, D), dtype=tl.float32) + + if split_end > split_start: + q_ptrs = ( + q_ptr + pid_b * stride_qb + pid_m * stride_qm + + offs_h[:, None] * stride_qh + offs_d[None, :] * stride_qd + ) + q = tl.load(q_ptrs, mask=h_mask[:, None], other=0.0).to(tl.bfloat16) + idx_base = idx_ptr + pid_b * stride_ib + pid_m * stride_im + + win_end = tl.minimum(split_end, N_WINDOW) + for start in range(split_start, win_end, BLOCK_T): + offs_t = start + tl.arange(0, BLOCK_T) + idxs = tl.load(idx_base + offs_t * stride_it) + valid = idxs >= 0 + kv = load_rows(win_ptr, idxs, valid, offs_d, D, WIN_ROW_BYTES, WIN_FMT) + m_i, l_i, acc = _attend_tile(q, kv, valid, scale, m_i, l_i, acc) + cmp_start = tl.maximum(split_start, N_WINDOW) + for start in range(cmp_start, split_end, BLOCK_T): + offs_t = start + tl.arange(0, BLOCK_T) + t_mask = offs_t < split_end + idxs = tl.load(idx_base + offs_t * stride_it, mask=t_mask, other=-1) + valid = idxs >= 0 + kv = load_rows(cmp_ptr, idxs, valid, offs_d, D, CMP_ROW_BYTES, CMP_FMT) + m_i, l_i, acc = _attend_tile(q, kv, valid, scale, m_i, l_i, acc) + + out = tl.where(l_i[:, None] == 0.0, 0.0, acc / l_i[:, None]) + lse = tl.where(l_i == 0.0, -float("inf"), m_i + tl.log(l_i)) + + mid_base = ( + mid_o_ptr + pid_b * stride_mb + pid_m * stride_mm + + offs_h[:, None] * stride_mh + split_id * stride_ms + offs_d[None, :] * stride_md + ) + tl.store(mid_base, out, mask=h_mask[:, None]) + lse_base = ( + mid_lse_ptr + pid_b * stride_lb + pid_m * stride_lm + + offs_h * stride_lh + split_id * stride_ls + ) + tl.store(lse_base, lse, mask=h_mask) + + +@triton.jit +def _sparse_attn_splitk_merge_kernel( + mid_o_ptr, mid_lse_ptr, o_ptr, sink_ptr, + stride_mb, stride_mm, stride_mh, stride_ms, stride_md, + stride_lb, stride_lm, stride_lh, stride_ls, + stride_ob, stride_om, stride_oh, stride_od, + D: tl.constexpr, + NUM_SPLITS: tl.constexpr, +): + """Stage 2: log-sum-exp merge over the splits; the attention sink joins here once, as a null + key with logit ``attn_sink[h]`` and zero value. One program per (query, request, head).""" + pid_m = tl.program_id(0) + pid_b = tl.program_id(1) + pid_h = tl.program_id(2) + offs_d = tl.arange(0, D) + + m_i = tl.load(sink_ptr + pid_h).to(tl.float32) + l_i = 1.0 + acc = tl.zeros((D,), dtype=tl.float32) + + mid_base = mid_o_ptr + pid_b * stride_mb + pid_m * stride_mm + pid_h * stride_mh + offs_d * stride_md + lse_base = mid_lse_ptr + pid_b * stride_lb + pid_m * stride_lm + pid_h * stride_lh + + for split_id in tl.range(0, NUM_SPLITS, num_stages=2): + partial = tl.load(mid_base + split_id * stride_ms) + lse = tl.load(lse_base + split_id * stride_ls) + m_new = tl.maximum(m_i, lse) + alpha = tl.exp(m_i - m_new) + beta = tl.where(lse == -float("inf"), 0.0, tl.exp(lse - m_new)) + acc = acc * alpha + partial * beta + l_i = l_i * alpha + beta + m_i = m_new + + o = acc / l_i + o_ptrs = o_ptr + pid_b * stride_ob + pid_m * stride_om + pid_h * stride_oh + offs_d * stride_od + tl.store(o_ptrs, o.to(o_ptr.dtype.element_ty)) + + +def split_count(b: int, m: int, h: int, topk: int, device) -> int: + """How many ways to split the candidate axis; 0 means run the single-program kernel. Decode + splits until every SM has a program or every split is down to one tile.""" + if m != 1: + return 0 + sm_count = torch.cuda.get_device_properties(device).multi_processor_count + n_splits = min( + MAX_SPLITS, + triton.cdiv(topk, MIN_TILES_PER_SPLIT * BLOCK_T), + max(1, sm_count // (b * triton.cdiv(h, BLOCK_H))), + ) + return n_splits if n_splits > 1 else 0 + + +def sparse_attn_packed( + q: torch.Tensor, # [b, m, h, d] bf16 + window_pool: torch.Tensor, # [n_win_slots, win_fmt.row_bytes(d)] uint8 + win_fmt: RowFormat, + cmp_pool: torch.Tensor, # [n_cmp, cmp_fmt.row_bytes(d)] uint8 + cmp_fmt: RowFormat, + attn_sink: torch.Tensor, # [h] + topk_idxs: torch.Tensor, # [b, m, topk] int32, GLOBAL rows, layout [window | compressed] + n_window: int, # window columns (a multiple of BLOCK_T; pad with -1) + softmax_scale: float, + cmp_counts: torch.Tensor | None = None, # [b, m] int32 live compressed columns per query +) -> torch.Tensor: + """Paged sparse MLA attention over packed KV pools; see the module docstring for the contract.""" + b, m, h, d = q.shape + topk = topk_idxs.shape[-1] + assert n_window % BLOCK_T == 0 and 0 <= n_window <= topk, (n_window, topk) + assert window_pool.dtype == torch.uint8 and window_pool.shape[1] == win_fmt.row_bytes(d), (window_pool.shape, win_fmt) + assert cmp_pool.dtype == torch.uint8 and cmp_pool.shape[1] == cmp_fmt.row_bytes(d), (cmp_pool.shape, cmp_fmt) + assert window_pool.is_contiguous() and cmp_pool.is_contiguous() + q = q.contiguous() + idx = topk_idxs.contiguous().to(torch.int32) + sink = attn_sink.contiguous().to(torch.float32) + o = torch.empty_like(q) + + has_counts = cmp_counts is not None + if has_counts: + cnt = cmp_counts.contiguous().to(torch.int32).view(b, m) + stride_nb, stride_nm = cnt.stride() + else: + cnt, stride_nb, stride_nm = idx, 0, 0 + + fmt_args = dict( + WIN_FMT=win_fmt.code, WIN_ROW_BYTES=win_fmt.row_bytes(d), + CMP_FMT=cmp_fmt.code, CMP_ROW_BYTES=cmp_fmt.row_bytes(d), + ) + n_splits = split_count(b, m, h, topk, q.device) + if n_splits: + head_blocks = triton.cdiv(h, BLOCK_H) + mid_o = torch.empty((b, m, h, n_splits, d), dtype=torch.float32, device=q.device) + mid_lse = torch.empty((b, m, h, n_splits), dtype=torch.float32, device=q.device) + _sparse_attn_packed_splitk_kernel[(m * n_splits, b, head_blocks)]( + q, window_pool, cmp_pool, mid_o, mid_lse, idx, cnt, + float(softmax_scale), + h, topk, int(n_window), + q.stride(0), q.stride(1), q.stride(2), q.stride(3), + 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), idx.stride(1), idx.stride(2), + stride_nb, stride_nm, + D=d, BLOCK_H=BLOCK_H, BLOCK_T=BLOCK_T, HAS_COUNTS=has_counts, NUM_SPLITS=n_splits, + **fmt_args, num_warps=8, num_stages=2, + ) + _sparse_attn_splitk_merge_kernel[(m, b, h)]( + mid_o, mid_lse, o, sink, + 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), + o.stride(0), o.stride(1), o.stride(2), o.stride(3), + D=d, NUM_SPLITS=n_splits, num_warps=4, + ) + return o + + grid = (m, b, triton.cdiv(h, BLOCK_H)) + _sparse_attn_packed_kernel[grid]( + q, window_pool, cmp_pool, o, sink, idx, cnt, + float(softmax_scale), + h, topk, int(n_window), + q.stride(0), q.stride(1), q.stride(2), q.stride(3), + o.stride(0), o.stride(1), o.stride(2), o.stride(3), + idx.stride(0), idx.stride(1), idx.stride(2), + stride_nb, stride_nm, + D=d, BLOCK_H=BLOCK_H, BLOCK_T=BLOCK_T, HAS_COUNTS=has_counts, + **fmt_args, num_warps=8, num_stages=2, + ) + return o + + +__all__ = ["sparse_attn_packed", "split_count", "BLOCK_T"] diff --git a/python/freetoken/kernel/triton/dsv41/topk.py b/python/freetoken/kernel/triton/dsv41/topk.py new file mode 100644 index 0000000000..4d638e4d49 --- /dev/null +++ b/python/freetoken/kernel/triton/dsv41/topk.py @@ -0,0 +1,265 @@ +"""Exact top-k and block-candidate selection whose work follows the live history (DeepSeek-V4.1). + +The indexer scores a request's compressed positions into a buffer as wide as the staged history +(``max_seq_len / ratio``, fixed while a CUDA graph is captured), but only the first ``live[row]`` +columns are meaningful -- and at 1M staged positions a 4K-token request must not pay for the +other 99.6%. Every kernel here takes the live count per row on the device, keeps its output +addresses and shapes fixed, and restricts what it READS and reduces to the live prefix; nothing +past ``live`` is ever read, so the score buffer's dead area needs no initialization. + +Selection is a fixed pipeline of small kernels with explicit workspace contracts: + +* ``dsv41_topk``: exact top-k of ``[R, T]`` fp32 scores below ``live``. Level 0 gives every + ``CHUNK``-wide slice of a row its own program, which radix-selects the slice's top-k into a + ``(key, column)`` candidate workspace (a global winner beats at most k-1 columns, so it is + inside its own slice's top-k: the union is exact). Further levels apply the same program to the + candidates until one register-resident tile remains; the final program ranks it and emits the + winners in ascending column order, ``-1`` padded. Slices past ``live`` exit at once and a + slice with no more entries than ``k`` is copied through, so a short history costs a handful of + near-empty launches. The radix primitives are the QSA block top-k's (``kernel/triton/qsa/topk``). +* ``dsv41_block_scores``: per-``block_size`` block maxima over the live prefix (the Hierarchical + Sparse Indexer's level one), the block holding the newest live position pinned to ``+inf`` so it + is always kept; a fixed worker grid strides over the live blocks. +* ``dsv41_candidate_blocks``: the two above plus a fixed-width expansion into the compact candidate + list ``[R, topk_blocks * block_size]`` of compressed positions, ascending with ``-1`` in the tail. + +Tie rule (documented, deterministic): scores order descending, equal scores by ascending column, +so the lowest positions win among exact ties. ``torch.topk`` leaves ties unspecified. +""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + +from freetoken.kernel.triton.qsa.topk import _BINS, _PASSES, _RADIX, _monotone_key, _resident_prefix, _tile_ranks + +CHUNK = 4096 # register-resident slice per level program (fp32 keys + int32 columns) +FINAL_MAX = 4096 # the final program's resident candidate tile +BLOCK_SCORE_TILE = 512 # blocks per block-max program tile + + +@triton.jit +def _load_entries(src_ptr, col_ptr, base, offsets, limit, SRC_LOGITS: tl.constexpr): + """One slice of a level's input as ``(monotone key, column)``; key 0 = dead (past ``limit``, + ``-inf``, or a padding slot of the previous level).""" + live = offsets < limit + if SRC_LOGITS: + value = tl.load(src_ptr + base + offsets, mask=live, other=-float("inf")) + key = tl.where(live & (value > -float("inf")), _monotone_key(value), 0) + cols = (base + offsets).to(tl.int32) + else: + key = tl.load(src_ptr + base + offsets, mask=live, other=0).to(tl.uint32, bitcast=True) + cols = tl.load(col_ptr + base + offsets, mask=live, other=-1) + return key, cols + + +@triton.jit +def _dsv41_topk_level_kernel( + src_ptr, src_col_ptr, live_in_ptr, key_out_ptr, col_out_ptr, live_out_ptr, + stride_src_row, stride_out_row, num_columns, + TOP_K: tl.constexpr, PAD_K: tl.constexpr, CHUNK: tl.constexpr, N_SPLITS: tl.constexpr, + SRC_LOGITS: tl.constexpr, BINS: tl.constexpr, RADIX: tl.constexpr, PASSES: tl.constexpr, +): + """One program per (row, slice): the slice's top-k into candidate slot ``split * TOP_K``. Slot + ``split`` is read downstream only when ``split < cdiv(live, CHUNK)``; program 0 publishes that + live candidate count for the next level.""" + tl.static_assert(CHUNK <= 0xFFFF, "packed cumsum keeps 16 bits per half") + tl.static_assert(TOP_K <= CHUNK, "a slice must hold its own top-k") + row = tl.program_id(0) + split = tl.program_id(1) + visible = tl.maximum(tl.minimum(tl.load(live_in_ptr + row), num_columns), 0) + if split == 0: + tl.store(live_out_ptr + row, tl.minimum(tl.cdiv(visible, CHUNK), N_SPLITS) * TOP_K) + base = split * CHUNK + limit = tl.minimum(tl.maximum(visible - base, 0), CHUNK) + if limit <= 0: + return # a dead slice: nothing downstream reads its slot + slot = row.to(tl.int64) * stride_out_row + split * TOP_K + offsets = tl.arange(0, CHUNK) + key, cols = _load_entries(src_ptr + row.to(tl.int64) * stride_src_row, src_col_ptr + row.to(tl.int64) * stride_src_row, base, offsets, limit, SRC_LOGITS) + if limit <= TOP_K: + # every live entry is a winner: copy in order (dead keys stay 0 for the next level) + inside = offsets < TOP_K + tl.store(key_out_ptr + slot + offsets, key.to(tl.int32, bitcast=True), mask=inside) + tl.store(col_out_ptr + slot + offsets, cols, mask=inside) + return + prefix, ties = _resident_prefix(key, TOP_K, BINS, RADIX, PASSES) + rank, take, above, equal = _tile_ranks(key, prefix, ties, 0, 0) + write = take & (rank < TOP_K) + tl.store(key_out_ptr + slot + rank, key.to(tl.int32, bitcast=True), mask=write) + tl.store(col_out_ptr + slot + rank, cols, mask=write) + emitted = above + tl.minimum(equal, ties) + pad = tl.arange(0, PAD_K) + tl.store(key_out_ptr + slot + pad, 0, mask=(pad >= emitted) & (pad < TOP_K)) + + +@triton.jit +def _final_tile(src_row, col_row, out_row, limit, TOP_K: tl.constexpr, BLOCK: tl.constexpr, SRC_LOGITS: tl.constexpr, BINS: tl.constexpr, RADIX: tl.constexpr, PASSES: tl.constexpr): + tl.static_assert(BLOCK <= 0xFFFF, "packed cumsum keeps 16 bits per half") + offsets = tl.arange(0, BLOCK) + key, cols = _load_entries(src_row, col_row, 0, offsets, limit, SRC_LOGITS) + prefix, ties = _resident_prefix(key, tl.minimum(limit, TOP_K), BINS, RADIX, PASSES) + rank, take, above, equal = _tile_ranks(key, prefix, ties, 0, 0) + tl.store(out_row + rank, cols, mask=take & (rank < TOP_K)) + return above + tl.minimum(equal, ties) + + +@triton.jit +def _dsv41_topk_final_kernel( + src_ptr, src_col_ptr, live_ptr, out_ptr, + stride_src_row, stride_out_row, num_columns, + TOP_K: tl.constexpr, PAD_K: tl.constexpr, + BLOCK_SMALL: tl.constexpr, BLOCK_MID: tl.constexpr, BLOCK_FULL: tl.constexpr, + SRC_LOGITS: tl.constexpr, BINS: tl.constexpr, RADIX: tl.constexpr, PASSES: tl.constexpr, +): + """One program per row over the last level's live candidates (or the scores themselves when the + row fits one tile): winners at their ascending-column ranks, ``-1`` from ``emitted`` on.""" + row = tl.program_id(0) + limit = tl.maximum(tl.minimum(tl.load(live_ptr + row), num_columns), 0) + src_row = src_ptr + row.to(tl.int64) * stride_src_row + col_row = src_col_ptr + row.to(tl.int64) * stride_src_row + out_row = out_ptr + row.to(tl.int64) * stride_out_row + emitted = 0 + if limit > 0: + # register residency costs the whole tile, so a short row takes a narrower one + if limit <= BLOCK_SMALL: + emitted = _final_tile(src_row, col_row, out_row, limit, TOP_K, BLOCK_SMALL, SRC_LOGITS, BINS, RADIX, PASSES) + elif limit <= BLOCK_MID: + emitted = _final_tile(src_row, col_row, out_row, limit, TOP_K, BLOCK_MID, SRC_LOGITS, BINS, RADIX, PASSES) + else: + emitted = _final_tile(src_row, col_row, out_row, limit, TOP_K, BLOCK_FULL, SRC_LOGITS, BINS, RADIX, PASSES) + pad = tl.arange(0, PAD_K) + tl.store(out_row + pad, -1, mask=(pad >= emitted) & (pad < TOP_K)) + + +MAX_K = CHUNK // 2 # a slice must shed at least half its entries per level, or the plan cannot converge + + +def topk_plan(columns: int, k: int) -> list[tuple[int, int]]: + """``[(input width, n_splits), ...]`` of the reduction levels before the final tile; empty when + the row already fits the final tile. Derived from the buffer width alone (graph-static). Each + level maps ``width`` to ``cdiv(width, CHUNK) * k`` candidates, which shrinks the row by at least + ``CHUNK / k >= 2`` per level for ``k <= MAX_K``; larger ``k`` would stall (``k == CHUNK`` never + shrinks) and is rejected here, before any launch.""" + if not 0 < k <= MAX_K: + raise ValueError(f"dsv41_topk supports 1 <= k <= {MAX_K}, got {k}") + levels: list[tuple[int, int]] = [] + width = columns + while width > FINAL_MAX: + n_splits = -(-width // CHUNK) + levels.append((width, n_splits)) + nxt = n_splits * k + assert nxt < width, (width, k) # guaranteed by k <= CHUNK // 2 for width > FINAL_MAX == CHUNK + width = nxt + return levels + + +def dsv41_topk(scores: torch.Tensor, live: torch.Tensor, k: int, out: torch.Tensor | None = None) -> torch.Tensor: + """Exact top-``k`` columns of every ``scores`` row among its first ``live[row]`` columns. + + ``scores [R, T]`` fp32 row-contiguous (columns at or past ``live`` are never read; ``-inf`` + columns are dead), ``live [R]`` int32. Returns ``out [R, k]`` int32: the winners in ascending + column order packed at the front, ``-1`` from the ``emitted`` count on (fewer live or finite + columns than ``k`` leave a ``-1`` tail). Ties: lower column wins.""" + assert scores.ndim == 2 and scores.dtype == torch.float32 and scores.stride(1) == 1, (scores.shape, scores.dtype) + R, T = scores.shape + assert live.shape == (R,), (live.shape, R) + live = live.to(torch.int32).contiguous() + if out is None: + out = torch.empty((R, k), dtype=torch.int32, device=scores.device) + assert out.shape == (R, k) and out.dtype == torch.int32 and out.stride(1) == 1 + if R == 0 or k == 0: + return out + if T == 0: + out.fill_(-1) + return out + pad_k = triton.next_power_of_2(k) + src, src_col, src_live, width, src_logits = scores, scores, live, T, True + for level_width, n_splits in topk_plan(T, k): + keys = torch.empty((R, n_splits * k), dtype=torch.int32, device=scores.device) + cols = torch.empty((R, n_splits * k), dtype=torch.int32, device=scores.device) + live_out = torch.empty((R,), dtype=torch.int32, device=scores.device) + _dsv41_topk_level_kernel[(R, n_splits)]( + src, src_col, src_live, keys, cols, live_out, + src.stride(0), keys.stride(0), level_width, + TOP_K=k, PAD_K=pad_k, CHUNK=CHUNK, N_SPLITS=n_splits, SRC_LOGITS=src_logits, + BINS=_BINS, RADIX=_RADIX, PASSES=_PASSES, num_warps=8, num_stages=1, + ) + src, src_col, src_live, width, src_logits = keys, cols, live_out, n_splits * k, False + block = triton.next_power_of_2(max(width, 1)) + _dsv41_topk_final_kernel[(R,)]( + src, src_col, src_live, out, + src.stride(0), out.stride(0), width, + TOP_K=k, PAD_K=pad_k, + BLOCK_SMALL=min(1024, block), BLOCK_MID=min(4096, block), BLOCK_FULL=block, + SRC_LOGITS=src_logits, BINS=_BINS, RADIX=_RADIX, PASSES=_PASSES, num_warps=8, num_stages=1, + ) + return out + + +@triton.jit +def _dsv41_block_scores_kernel( + scores_ptr, live_ptr, blk_ptr, + stride_s_row, stride_b_row, num_blocks, + BLOCK_SIZE: tl.constexpr, TILE: tl.constexpr, N_WORKERS: tl.constexpr, +): + """Worker ``w`` of a row scores block tiles ``w, w + N_WORKERS, ...`` below the live block count: + the block's best live position, ``+inf`` for the block holding the newest live position.""" + row = tl.program_id(0) + worker = tl.program_id(1) + live = tl.maximum(tl.load(live_ptr + row), 0) + nb_live = tl.minimum(tl.cdiv(live, BLOCK_SIZE), num_blocks) + s_row = scores_ptr + row.to(tl.int64) * stride_s_row + b_row = blk_ptr + row.to(tl.int64) * stride_b_row + offs_b = tl.arange(0, TILE) + offs_p = tl.arange(0, BLOCK_SIZE) + for tile in range(worker, tl.cdiv(nb_live, TILE), N_WORKERS): + blocks = tile * TILE + offs_b + b_mask = blocks < nb_live + pos = blocks[:, None] * BLOCK_SIZE + offs_p[None, :] + v = tl.load(s_row + pos, mask=b_mask[:, None] & (pos < live), other=-float("inf")) + best = tl.max(v, axis=1) + best = tl.where(blocks == nb_live - 1, float("inf"), best) + tl.store(b_row + blocks, best, mask=b_mask) + + +def dsv41_block_scores(scores: torch.Tensor, live: torch.Tensor, block_size: int, out: torch.Tensor | None = None) -> tuple[torch.Tensor, torch.Tensor]: + """``(block scores [R, cdiv(T, block_size)] fp32, live blocks [R] int32)``: each live block's best + live position (``+inf`` for the newest block); blocks at or past the live count are not written.""" + R, T = scores.shape + nb = -(-T // block_size) + live = live.to(torch.int32).contiguous() + if out is None: + out = torch.empty((R, nb), dtype=torch.float32, device=scores.device) + live_blocks = torch.div(live + block_size - 1, block_size, rounding_mode="floor").clamp_(min=0).to(torch.int32) + if R == 0 or nb == 0: + return out, live_blocks + assert block_size == triton.next_power_of_2(block_size), block_size + sm_count = torch.cuda.get_device_properties(scores.device).multi_processor_count + workers = min(triton.cdiv(nb, BLOCK_SCORE_TILE), max(1, triton.cdiv(sm_count, R))) + _dsv41_block_scores_kernel[(R, workers)]( + scores, live, out, scores.stride(0), out.stride(0), nb, + BLOCK_SIZE=block_size, TILE=BLOCK_SCORE_TILE, N_WORKERS=workers, num_warps=4, + ) + return out, live_blocks + + +def dsv41_candidate_blocks(scores: torch.Tensor, live: torch.Tensor, topk_blocks: int, block_size: int) -> torch.Tensor: + """Level one of the Hierarchical Sparse Indexer over ``[R, T]`` scores: the ``topk_blocks`` best + live blocks (the newest always among them) expanded to the compact candidate list ``[R, + topk_blocks * block_size]`` of compressed positions -- ascending, ``-1`` for positions at or past + ``live`` and for unfilled blocks, all in the tail (a sorted valid prefix).""" + R, T = scores.shape + blk, live_blocks = dsv41_block_scores(scores, live, block_size) + keep = dsv41_topk(blk, live_blocks, min(topk_blocks, blk.shape[1])) # [R, KB] ascending, -1 tail + pos = keep.to(torch.int64).unsqueeze(-1) * block_size + torch.arange(block_size, device=scores.device) + dead = (keep < 0).unsqueeze(-1) | (pos >= live.to(torch.int64).view(R, 1, 1)) + out = torch.where(dead, -1, pos).flatten(1).to(torch.int32) + if out.shape[1] < topk_blocks * block_size: + out = torch.nn.functional.pad(out, (0, topk_blocks * block_size - out.shape[1]), value=-1) + return out + + +__all__ = ["dsv41_topk", "dsv41_block_scores", "dsv41_candidate_blocks", "topk_plan", "CHUNK", "FINAL_MAX", "MAX_K"] diff --git a/python/freetoken/kernel/triton/e4m3_compat.py b/python/freetoken/kernel/triton/e4m3_compat.py index 6ea0c7537d..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) @@ -109,6 +168,30 @@ def e4m3_u8_to_f16_x128(v): return h.to(tl.float16, bitcast=True) * 32768.0 +@jit +def e4m3_f32_to_u8(x): + """Encode an fp32 value that already sits on the e4m3 grid (see :func:`round_e4m3`, |x| <= 448) + into its e4m3 bit pattern (uint8). Normals: fp32 exponent rebased 127 -> 7 with the top three + mantissa bits; subnormals (|x| < 2^-6, grid 2^-9): mantissa = |x| / 2^-9, exponent field 0.""" + b = x.to(tl.uint32, bitcast=True) + sign = ((b >> 31) & 1).to(tl.uint8) << 7 + exp = ((b >> 23) & 0xFF).to(tl.int32) - 127 + 7 + man = ((b >> 20) & 0x7).to(tl.uint8) + normal = (exp.to(tl.uint8) << 3) | man + sub = (tl.abs(x) * 512.0).to(tl.uint8) # 0..7 on the 2^-9 grid + 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 @@ -126,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/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/mhc.py b/python/freetoken/kernel/triton/mhc.py index 0782efe37d..876643c164 100644 --- a/python/freetoken/kernel/triton/mhc.py +++ b/python/freetoken/kernel/triton/mhc.py @@ -27,7 +27,7 @@ def _mhc_stage1_kernel( H: tl.constexpr, N: tl.constexpr, MIX: tl.constexpr, BLK_MIX: tl.constexpr, SPLIT: tl.constexpr, # hidden elems per split (multiple of BLOCK_H) BLOCK_H: tl.constexpr, - NS: tl.constexpr, + NS: tl.constexpr, # splits per token (the grid's second axis; need not be a power of two) HAS_POST: tl.constexpr, ): """Split-K stage of the fused mHC: each program owns one hidden slice of one @@ -79,24 +79,24 @@ def _mhc_stage1_kernel( @triton.jit -def _mhc_stage2_kernel( - sq_part_ptr, mix_part_ptr, scale_ptr, base_ptr, +def _mhc_gates( + t, sq_part_ptr, mix_part_ptr, scale_ptr, base_ptr, post_out_ptr, comb_out_ptr, pre_out_ptr, rms_eps, hc_eps, post_mult, SINKHORN: tl.constexpr, H: tl.constexpr, N: tl.constexpr, MIX: tl.constexpr, BLK_MIX: tl.constexpr, - NS: tl.constexpr, + NS: tl.constexpr, BLK_NS: tl.constexpr, ): - """Reduce the split partials and run the tiny gate math (sigmoid gates, - row-softmax + in-register 4x4 Sinkhorn); emits pre gates for stage 3.""" - t = tl.program_id(0).to(tl.int64) + """Reduce token ``t``'s split partials and run the tiny gate math (sigmoid gates, + row-softmax + in-register 4x4 Sinkhorn); writes post / comb / pre for the next sublayer.""" offs_mix = tl.arange(0, BLK_MIX) mix_mask = offs_mix < MIX - offs_s = tl.arange(0, NS) + offs_s = tl.arange(0, BLK_NS) + s_mask = offs_s < NS - sqsum = tl.sum(tl.load(sq_part_ptr + t * NS + offs_s)) + sqsum = tl.sum(tl.load(sq_part_ptr + t * NS + offs_s, mask=s_mask, other=0.0)) acc = tl.sum( - tl.load(mix_part_ptr + (t * NS + offs_s)[:, None] * BLK_MIX + offs_mix[None, :]), + tl.load(mix_part_ptr + (t * NS + offs_s)[:, None] * BLK_MIX + offs_mix[None, :], mask=s_mask[:, None], other=0.0), axis=0, ) inv_rms = tl.math.rsqrt(sqsum / (N * H) + rms_eps) @@ -142,6 +142,23 @@ def _mhc_stage2_kernel( tl.store(comb_out_ptr + t * N * N + offs_n2[:, None] * N + offs_n2[None, :], comb) + +@triton.jit +def _mhc_stage2_kernel( + sq_part_ptr, mix_part_ptr, scale_ptr, base_ptr, + post_out_ptr, comb_out_ptr, pre_out_ptr, + rms_eps, hc_eps, post_mult, + SINKHORN: tl.constexpr, + H: tl.constexpr, N: tl.constexpr, MIX: tl.constexpr, BLK_MIX: tl.constexpr, + NS: tl.constexpr, BLK_NS: tl.constexpr, +): + """Gate math alone (the two-pass variant, whose stage 3 consumes these pre gates).""" + _mhc_gates( + tl.program_id(0).to(tl.int64), sq_part_ptr, mix_part_ptr, scale_ptr, base_ptr, post_out_ptr, comb_out_ptr, pre_out_ptr, + rms_eps, hc_eps, post_mult, SINKHORN, H, N, MIX, BLK_MIX, NS, BLK_NS, + ) + + @triton.jit def _mhc_stage3_kernel( res_out_ptr, pre_ptr, li_out_ptr, @@ -161,7 +178,46 @@ def _mhc_stage3_kernel( tl.store(li_out_ptr + t * H + offs_h, li.to(li_out_ptr.dtype.element_ty), mask=h_mask) -def mhc_fused_post_pre_triton( +@triton.jit +def _mhc_stage23_kernel( + res_out_ptr, pre_ptr, li_out_ptr, + sq_part_ptr, mix_part_ptr, scale_ptr, base_ptr, + post_out_ptr, comb_out_ptr, pre_out_ptr, + rms_eps, hc_eps, post_mult, + SINKHORN: tl.constexpr, + H: tl.constexpr, N: tl.constexpr, MIX: tl.constexpr, BLK_MIX: tl.constexpr, + NS: tl.constexpr, BLK_NS: tl.constexpr, BLOCK_H: tl.constexpr, +): + """Single-pass mHC: the layer input mixes with the PREVIOUS sublayer's pre gates, so the + gate math for the next sublayer (stage 2) does not gate stage 3 -- one launch does both: + every program mixes one hidden chunk, program 0 of each token also reduces the split + partials and writes the next sublayer's gates.""" + t = tl.program_id(0).to(tl.int64) + hb = tl.program_id(1) + offs_h = hb * BLOCK_H + tl.arange(0, BLOCK_H) + h_mask = offs_h < H + offs_n = tl.arange(0, N) + pre = tl.load(pre_ptr + t * N + offs_n) + li = tl.zeros([BLOCK_H], dtype=tl.float32) + for n in tl.static_range(N): + r = tl.load(res_out_ptr + (t * N + n) * H + offs_h, mask=h_mask, other=0.0).to(tl.float32) + li += tl.sum(tl.where(offs_n == n, pre, 0.0)) * r + tl.store(li_out_ptr + t * H + offs_h, li.to(li_out_ptr.dtype.element_ty), mask=h_mask) + if hb == 0: + _mhc_gates( + t, sq_part_ptr, mix_part_ptr, scale_ptr, base_ptr, post_out_ptr, comb_out_ptr, pre_out_ptr, + rms_eps, hc_eps, post_mult, SINKHORN, H, N, MIX, BLK_MIX, NS, BLK_NS, + ) + + +def _stage1_hidden_tile(h: int) -> int: + """The stage-1 hidden split. It fixes the fp32 reduction order over the hidden width, so it + depends on the width alone and results reproduce across processes, machines and chunk sizes. + 128 gives 32 / 40 programs per token at widths 4096 / 5120.""" + return min(128, triton.next_power_of_2(h)) + + +def _mhc_fused_launch( x: torch.Tensor, residual: torch.Tensor, post_mix: torch.Tensor | None, @@ -173,11 +229,12 @@ def mhc_fused_post_pre_triton( hc_eps: float, post_mult: float, sinkhorn_repeat: int, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - """Fused hc_post (skipped when ``post_mix is None``) + hc_pre. Three-stage - split-K: the GEMV/sq-sum reduction fans out over NS hidden slices. Returns - (residual_new [T,N,H] bf16, post [T,N,1] fp32, comb [T,N,N] fp32, - layer_input [T,H] bf16).""" + pre_mix_in: torch.Tensor | None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """The three stages behind both variants. ``pre_mix_in`` selects single-pass: the + layer input is mixed with it instead of with the pre gates stage 2 just computed + (those are still returned, for the next sublayer). Returns (residual_new, post, + comb, pre_new [T,N] fp32, layer_input).""" t, n, h = residual.shape mix = 2 * n + n * n assert fn.shape == (mix, n * h) and fn.dtype == torch.float32 @@ -190,12 +247,7 @@ def mhc_fused_post_pre_triton( comb_out = torch.empty(t, n, n, dtype=torch.float32, device=dev) li_out = torch.empty(t, h, dtype=residual.dtype, device=dev) - block_h = min(512, triton.next_power_of_2(h)) - # NS feeds a tl.arange in stage 2 -> keep it a power of two. - ns = 1 - while ns * 2 <= min(16, h // block_h): - ns *= 2 - split = triton.cdiv(triton.cdiv(h, ns), block_h) * block_h + block_h = split = _stage1_hidden_tile(h) ns = triton.cdiv(h, split) blk_mix = triton.next_power_of_2(mix) sq_part = torch.empty(t, ns, dtype=torch.float32, device=dev) @@ -213,20 +265,83 @@ def mhc_fused_post_pre_triton( HAS_POST=has_post, num_warps=4, num_stages=2, ) - _mhc_stage2_kernel[(t,)]( - sq_part, mix_part, hc_scale, hc_base, - post_out, comb_out, pre_out, - rms_eps, hc_eps, post_mult, - SINKHORN=sinkhorn_repeat, - H=h, N=n, MIX=mix, BLK_MIX=blk_mix, NS=ns, - num_warps=1, + block_h3 = min(1024, triton.next_power_of_2(h)) + if pre_mix_in is None: + # two-pass: stage 3 mixes with the gates stage 2 just produced + _mhc_stage2_kernel[(t,)]( + sq_part, mix_part, hc_scale, hc_base, + post_out, comb_out, pre_out, + rms_eps, hc_eps, post_mult, + SINKHORN=sinkhorn_repeat, + H=h, N=n, MIX=mix, BLK_MIX=blk_mix, NS=ns, BLK_NS=triton.next_power_of_2(ns), + num_warps=1, + ) + _mhc_stage3_kernel[(t, triton.cdiv(h, block_h3))]( + res_out, pre_out, li_out, + H=h, N=n, BLOCK_H=block_h3, + num_warps=4, + ) + else: + # single-pass: the mix uses the previous sublayer's gates, so stages 2 and 3 share a launch + assert pre_mix_in.shape == (t, n), (pre_mix_in.shape, (t, n)) + _mhc_stage23_kernel[(t, triton.cdiv(h, block_h3))]( + res_out, pre_mix_in.to(torch.float32).contiguous(), li_out, + sq_part, mix_part, hc_scale, hc_base, + post_out, comb_out, pre_out, + rms_eps, hc_eps, post_mult, + SINKHORN=sinkhorn_repeat, + H=h, N=n, MIX=mix, BLK_MIX=blk_mix, NS=ns, BLK_NS=triton.next_power_of_2(ns), BLOCK_H=block_h3, + num_warps=4, + ) + return res_out, post_out.view(t, n, 1), comb_out, pre_out, li_out + + +def mhc_fused_post_pre_triton( + x: torch.Tensor, + residual: torch.Tensor, + post_mix: torch.Tensor | None, + comb_mix: torch.Tensor | None, + fn: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + rms_eps: float, + hc_eps: float, + post_mult: float, + sinkhorn_repeat: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Fused hc_post (skipped when ``post_mix is None``) + hc_pre. Three-stage + split-K: the GEMV/sq-sum reduction fans out over NS hidden slices. Returns + (residual_new [T,N,H] bf16, post [T,N,1] fp32, comb [T,N,N] fp32, + layer_input [T,H] bf16).""" + res_out, post_out, comb_out, _pre_out, li_out = _mhc_fused_launch( + x, residual, post_mix, comb_mix, fn, hc_scale, hc_base, + rms_eps, hc_eps, post_mult, sinkhorn_repeat, pre_mix_in=None, ) - _mhc_stage3_kernel[(t, triton.cdiv(h, 1024))]( - res_out, pre_out, li_out, - H=h, N=n, BLOCK_H=min(1024, triton.next_power_of_2(h)), - num_warps=4, + return res_out, post_out, comb_out, li_out + + +def mhc_fused_post_pre_single_pass_triton( + x: torch.Tensor, + residual: torch.Tensor, + post_mix: torch.Tensor | None, + comb_mix: torch.Tensor | None, + pre_mix_in: torch.Tensor, + fn: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + rms_eps: float, + hc_eps: float, + post_mult: float, + sinkhorn_repeat: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Single-pass mHC (DeepSeek-V4.1): same fused stages, but the layer input is mixed + with the PREVIOUS sublayer's ``pre_mix_in`` and this sublayer's pre gates come back + as ``pre_new`` for the next one. Returns (residual_new, post [T,N,1], comb [T,N,N], + pre_new [T,N] fp32, layer_input).""" + return _mhc_fused_launch( + x, residual, post_mix, comb_mix, fn, hc_scale, hc_base, + rms_eps, hc_eps, post_mult, sinkhorn_repeat, pre_mix_in=pre_mix_in, ) - return res_out, post_out.view(t, n, 1), comb_out, li_out -__all__ = ["mhc_fused_post_pre_triton"] +__all__ = ["mhc_fused_post_pre_triton", "mhc_fused_post_pre_single_pass_triton"] 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 a1064bcf7d..e451bbe65b 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,7 +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.e4m3_compat import e4m3_native_cx, e4m3_to_f16_x128, e4m3_u8_to_f32 +from freetoken.kernel.triton.nvfp4_linear import _nvfp4_pair_f16 _E2M1_VALUES = [ 0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, @@ -44,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) @@ -238,7 +250,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, @@ -256,6 +267,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) @@ -277,8 +289,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 @@ -288,30 +300,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..f6e223385b 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 @@ -303,6 +433,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 +453,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 +472,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/kernel/triton/sampling.py b/python/freetoken/kernel/triton/sampling.py index a970714406..869046aa1c 100644 --- a/python/freetoken/kernel/triton/sampling.py +++ b/python/freetoken/kernel/triton/sampling.py @@ -92,7 +92,10 @@ def _sm_partial( x = tl.load(logits_ptr + base + offs, mask=mask, other=-float("inf")).to(tl.float32) * inv_t blk_max = tl.max(x, 0) new_m = tl.maximum(m, blk_max) - d = d * tl.exp(m - new_m) + tl.sum(tl.exp(x - new_m), 0) + # A fully masked prefix (min_p / min_tokens write -inf) leaves new_m at -inf, and + # exp(-inf - -inf) is NaN; shift by 0 there instead, every term is exp(-inf) = 0. + safe_m = tl.where(new_m == -float("inf"), 0.0, new_m) + d = d * tl.exp(m - safe_m) + tl.sum(tl.exp(x - safe_m), 0) m = new_m tl.store(pm_ptr + pid, m) tl.store(pl_ptr + pid, d) @@ -117,7 +120,8 @@ def _sm_finalize( pm = tl.load(pm_ptr + row * G + goff, mask=gmask, other=-float("inf")) pl = tl.load(pl_ptr + row * G + goff, mask=gmask, other=0.0) gm = tl.max(pm, 0) - gl = tl.sum(pl * tl.exp(pm - gm), 0) + # a chunk that is -inf throughout holds no mass; keep its NaN-prone exp out of the sum + gl = tl.sum(tl.where(pm == -float("inf"), 0.0, pl * tl.exp(pm - gm)), 0) inv_gl = 1.0 / gl base = row * row_stride diff --git a/python/freetoken/kvcache/__init__.py b/python/freetoken/kvcache/__init__.py index 1792da2b3c..32f7195dbb 100644 --- a/python/freetoken/kvcache/__init__.py +++ b/python/freetoken/kvcache/__init__.py @@ -34,7 +34,7 @@ def resolve_pool_class(model_config: ModelConfig) -> type[BaseKVCachePool]: specs_fn = getattr(model_config, "kv_cache_group_specs", None) if specs_fn is None: if getattr(model_config, "dsv4_args", None) is not None: - from .dsv4_paged_pool import DSV4PagedKVCache + from .dsv4.v4_pool import DSV4PagedKVCache return DSV4PagedKVCache from .mha_pool import MHAKVCache @@ -43,9 +43,13 @@ def resolve_pool_class(model_config: ModelConfig) -> type[BaseKVCachePool]: specs = list(specs_fn()) types = {spec.attn_type for spec in specs} if AttnType.DSV4 in types: - from .dsv4_paged_pool import DSV4PagedKVCache + from .dsv4.v4_pool import DSV4PagedKVCache return DSV4PagedKVCache + if AttnType.DSV41 in types: + from .dsv4.v41_pool import DSV41PagedKVCache + + return DSV41PagedKVCache if AttnType.SWA in types: from .hybrid_swa_pool import HybridSWAKVCache @@ -76,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.""" - from .dsv4_cost_model import _dsv4_pool_sizes + 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_paged_pool import DSV4PagedKVCache + 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, @@ -99,6 +120,21 @@ def create_kv_pool(config, num_pages: int, device: torch.device, dtype: torch.dt ) pool._init_paged_state(config.max_running_req, config.cache_type != "naive") return pool + if getattr(model_config, "dsv41_args", None) is not None: + from .dsv4.v41_cost_model import _dsv41_pool_sizes + from .dsv4.v41_pool import DSV41PagedKVCache + + # Same route as DSV4: the generic CacheManager over the shared page table; the pool is the + # swa_pool plug-in (window tier + per-source packed main / index pools + state rings). + pool = DSV41PagedKVCache( + sizes=_dsv41_pool_sizes(config, num_pages + 1), # +1 for dummy page + args=model_config.dsv41_args, + device=device, + dtype=dtype, + n_scratch=config.max_running_req + 1, + ) + pool._init_paged_state(config.max_running_req, config.cache_type != "naive") + return pool num_swa_tokens = None # Both the naive and radix SWA paths share the global-paged swa pool; radix sizes it by @@ -117,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, ) @@ -128,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 @@ -140,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 @@ -167,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( @@ -205,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: @@ -225,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. @@ -244,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 @@ -256,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 95669e8c81..54ad5ad13c 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,71 @@ 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) +class WindowPoolSpec: + """Window resize units and minimum usable capacity, excluding reserved pages.""" + + page_size: int + min_pages: int class BaseKVCachePool(ABC): @@ -45,6 +96,21 @@ class BaseKVCachePool(ABC): # model re-bound after a rebuild; the engine asks before it resizes. needs_rebind_on_rebuild: ClassVar[bool] = False + @classmethod + def window_spec(cls, config) -> WindowPoolSpec | None: + return None + + @property + 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 @@ -154,13 +220,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 6aa13e7f75..5b531ce265 100644 --- a/python/freetoken/kvcache/cache_status.py +++ b/python/freetoken/kvcache/cache_status.py @@ -3,13 +3,14 @@ from typing import Any, Dict +def window_pool_spec(config): + from . import resolve_pool_class + + return resolve_pool_class(config.model_config).window_spec(config) + + def _supports_swa_ratio(config) -> bool: - """Whether ``swa_full_tokens_ratio`` sizes a separate window pool for this model -- DSV4 - (always) or a radix-SWA model (Gemma). Gates the ratio in telemetry and rebuild.""" - mc = config.model_config - if mc.dsv4_args is not None: - return True - return mc.has_swa_attention and config.cache_type == "swa_radix" + return window_pool_spec(config) is not None def compute_cache_unit_bytes(engine: "Engine") -> Dict[str, int]: @@ -120,18 +121,8 @@ def _mamba() -> int: return int(_linear_pool_min_slots(config) - 1) def _swa() -> int: - # Report usable tokens, matching num_swa_pages in rebuild requests. - # DSV4's physical floor includes a dummy page; the SWA floor already excludes slot 0. - from .dsv4_cost_model import _dsv4_window_floor_pages - from .hybrid_swa_pool import _swa_pool_floor - - mc = config.model_config - if mc.dsv4_args is not None: - P = mc.dsv4_args.window_size - return int((_dsv4_window_floor_pages(config, P) - 1) * P) - if not (mc.has_swa_attention and config.cache_type == "swa_radix"): - return 0 - return int(_swa_pool_floor(config)) + spec = window_pool_spec(config) + return spec.page_size * spec.min_pages if spec is not None else 0 for key, fn in (("kv_tokens", _kv), ("moe_experts", _moe), ("mamba_slots", _mamba), ("swa_tokens", _swa)): @@ -157,19 +148,10 @@ def compute_cache_pools(engine: "Engine") -> Dict[str, int]: pools["num_pages"] = int(engine.num_pages or 0) if config is not None: pools["page_size"] = int(config.page_size or 0) - # Window pool: its own page unit (swa_page_size -- DSV4 windows are P-token pages, - # radix-SWA is token-granular page_size 1) and the concrete current size in that unit - # (num_swa_pages, usable count). Same source as the scheduler's _current_cache_geometry. - # Both 0 for models without a window pool. Lets a client denominate the swa control. - mc = config.model_config - if mc.dsv4_args is not None: - pools["swa_page_size"] = int(mc.dsv4_args.window_size or 0) - sizes = getattr(engine.kv_cache, "sizes", None) # usable = physical minus dummy - if sizes is not None: - pools["num_swa_pages"] = max(0, int(sizes.n_win_pages) - 1) - elif mc.has_swa_attention and config.cache_type == "swa_radix": - pools["swa_page_size"] = 1 # usable = pool tokens minus the slot-0 sentinel - pools["num_swa_pages"] = max(0, int(getattr(engine.kv_cache, "swa_num_tokens", 0) or 0) - 1) + spec = window_pool_spec(config) + if spec is not None: + pools["swa_page_size"] = spec.page_size + pools["num_swa_pages"] = max(0, int(engine.kv_cache.window_pages)) moe = engine.moe_offload_cache if moe is not None: pools["moe_cache_size"] = int(moe.cache_size or 0) diff --git a/python/freetoken/kvcache/dsa_pool.py b/python/freetoken/kvcache/dsa_pool.py index 13225cec06..f693227401 100644 --- a/python/freetoken/kvcache/dsa_pool.py +++ b/python/freetoken/kvcache/dsa_pool.py @@ -2,7 +2,7 @@ ``MLAKVCache`` is the MHA pool's sibling for latent-KV models: ONE slab holding the per-token latent ``ckv (kv_lora_rank) | kpe (qk_rope_head_dim)`` -- there is no -separate V (``v_cache`` aliases ``k_cache``, same convention as dsv4_paged_pool's +separate V (``v_cache`` aliases ``k_cache``, same convention as dsv4/v4_pool's single-latent tiers). ``DSAKVCache`` extends it with the DeepSeek-Sparse-Attention index-key slab: one ``index_head_dim``-wide bf16 row per token per full-indexer layer, addressed by the SAME physical rows as the latent slab (GLM-5.2, page 1; @@ -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,26 +66,52 @@ 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 ) def v_cache(self, layer_id: int) -> torch.Tensor: - # MLA: K == V (single latent); same buffer, dsv4_paged_pool precedent. + # MLA: K == V (single latent); same buffer, dsv4/v4_pool precedent. 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/__init__.py b/python/freetoken/kvcache/dsv4/__init__.py new file mode 100644 index 0000000000..7531470f0f --- /dev/null +++ b/python/freetoken/kvcache/dsv4/__init__.py @@ -0,0 +1 @@ +"""Paged KV pools for the DeepSeek-V4 family.""" diff --git a/python/freetoken/kvcache/dsv4/v41_cost_model.py b/python/freetoken/kvcache/dsv4/v41_cost_model.py new file mode 100644 index 0000000000..4db1fb18c1 --- /dev/null +++ b/python/freetoken/kvcache/dsv4/v41_cost_model.py @@ -0,0 +1,216 @@ +"""Cost model + per-tier sizing for the DSV41 paged KV pool (DeepSeek-V4.1). + +Same shape as ``v4_cost_model``: one bytes-per-P-token number for the budget division, and +independent per-tier sizing from the anchor ``full_token = num_pages * P``. The tiers: + +* window pool, every layer -- ``swa_ratio`` of the full history, packed fp8 rows +* main KV pool, per kv source -- ``full_token // ratio`` packed fp4 rows +* index-key pool, per kv source -- ``full_token // ratio`` packed fp4 rows +* compress-state ring, ratio>1 sources -- ``ring_size`` fp32 ``kv|score`` slots per WINDOW page + +Sizing reads ``ModelConfig.dsv41_args`` by attribute, like ``v4_cost_model`` reads ``dsv4_args``: +the kvcache package never imports the model package. All byte widths come from the tier row formats +below, so there is exactly one place that knows what a row costs. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field + +from .v41_row_format import FP4_E4M3_B16, FP4_E8M0_B32, FP8_E8M0_B32 +from .window_tier import reserved_window_pages + +_INT64_BYTES = 8 +_FP32_BYTES = 4 + +WIN_FMT = FP8_E8M0_B32 +MAIN_FMT = FP4_E4M3_B16 +IDX_FMT = FP4_E8M0_B32 + + +def _win_row_bytes(args) -> int: + return WIN_FMT.row_bytes(args.head_dim) + + +def _main_row_bytes(args) -> int: + return MAIN_FMT.row_bytes(args.head_dim) + + +def _idx_row_bytes(args) -> int: + return IDX_FMT.row_bytes(args.index_head_dim) + + +def _state_bytes(args) -> int: + """One compress-state ring slot: fp32 ``kv | score`` of the latent width.""" + return 2 * args.head_dim * _FP32_BYTES + + +def private_window_layer_ids(args) -> tuple[int, ...]: + """Layers whose window KV is REQUEST-PRIVATE: kept in a per-page-table-row ring of ``window`` + slots (slot = row * window + pos % window) instead of the shared, page-bound window pool, so + it never lands in radix-shared pages. The decoder layers under Decoder SWA Bounded Replay: + their KV is recomputed per request over the prompt's last window (positions before it are + never computed), so a shared page could hold rows another request never wrote. Exact mode + computes every position and keeps the decoder on the shared pool.""" + if args.swa_decoder_replay == "exact": + return () + return tuple(range(args.decoder_start_layer, args.n_layers)) + + +@dataclass +class DSV41PoolSizes: + """Per-tier row counts derived from the budget anchor ``full_token``; the global tiers are + keyed by kv-source layer id.""" + + P: int + swa_ratio: float + full_token: int + n_win_slots: int + n_win_pages: int + main_rows: dict[int, int] = field(default_factory=dict) + idx_rows: dict[int, int] = field(default_factory=dict) + state_slots: dict[int, int] = field(default_factory=dict) # ratio>1 sources only + + +def dsv41_pool_sizes( + num_pages: int, args, swa_ratio: float, P: int, n_win_pages: int | None = None +) -> DSV41PoolSizes: + full_token = num_pages * P + if n_win_pages is None: + n_win_pages = (round(swa_ratio * full_token) + P - 1) // P + n_win_pages = min(n_win_pages, num_pages) + sizes = DSV41PoolSizes(P=P, swa_ratio=swa_ratio, full_token=full_token, n_win_slots=n_win_pages * P, n_win_pages=n_win_pages) + for src in args.backbone_kv_sources: + ratio = args.compress_ratios[src] + sizes.main_rows[src] = full_token // ratio + sizes.idx_rows[src] = full_token // ratio + if ratio > 1: + sizes.state_slots[src] = n_win_pages * ratio # ring size == ratio + return sizes + + +def dsv41_pool_bytes(sizes: DSV41PoolSizes, args, n_scratch: int = 1) -> int: + """Exact bytes a ``DSV41PagedKVCache`` built from ``sizes`` allocates (mirror of ``total_bytes``).""" + n_private = len(private_window_layer_ids(args)) + total = (args.n_layers - n_private) * sizes.n_win_slots * _win_row_bytes(args) + total += n_private * n_scratch * sizes.P * _win_row_bytes(args) # per-request rings + total += (sizes.full_token + 1) * _INT64_BYTES # full_to_window (+ sentinel row) + for src in args.backbone_kv_sources: + total += (sizes.main_rows[src] + n_scratch) * _main_row_bytes(args) + total += (sizes.idx_rows[src] + n_scratch) * _idx_row_bytes(args) + if src in sizes.state_slots: + total += (sizes.state_slots[src] + 1) * _state_bytes(args) + return int(total) + + +def dsv41_cache_per_page(args, swa_ratio: float, P: int) -> int: + """Marginal bytes per P-token page across all tiers (window scaled by ``swa_ratio``).""" + n_shared = args.n_layers - len(private_window_layer_ids(args)) + total = n_shared * round(swa_ratio * P) * _win_row_bytes(args) + for src in args.backbone_kv_sources: + ratio = args.compress_ratios[src] + total += (P // ratio) * (_main_row_bytes(args) + _idx_row_bytes(args)) + if ratio > 1: + total += round(swa_ratio * ratio) * _state_bytes(args) + return int(total) + + +def dsv41_kv_unit_bytes(args, P: int) -> int: + """FULL-tier bytes per full-history token (main + index pools + the mapping); the slider's + ``kv_bytes_per_token``. 890 B/token for DeepSeek-V4.1-Flash before the mapping.""" + per_page = P * _INT64_BYTES + for src in args.backbone_kv_sources: + per_page += (P // args.compress_ratios[src]) * (_main_row_bytes(args) + _idx_row_bytes(args)) + return -(-per_page // P) + + +def dsv41_window_unit_bytes(args, P: int) -> int: + """WINDOW-tier bytes per window token (sliding KV on every shared-window layer + the state rings).""" + per_page = (args.n_layers - len(private_window_layer_ids(args))) * P * _win_row_bytes(args) + for src in args.backbone_kv_sources: + ratio = args.compress_ratios[src] + if ratio > 1: + per_page += ratio * _state_bytes(args) + return -(-per_page // P) + + +def dsv41_solve_num_pages( + available_bytes: int, args, swa_ratio: float, floor_win_pages: int, P: int, n_scratch: int = 1 +) -> DSV41PoolSizes: + """Largest budget-respecting pool; the window floor is honored in pages and the total is + byte-checked. Raises ``ValueError`` when even the minimal pool does not fit.""" + + def _sizes(num: int) -> DSV41PoolSizes: + win = max(floor_win_pages, (round(swa_ratio * num * P) + P - 1) // P) + return dsv41_pool_sizes(num, args, swa_ratio, P, n_win_pages=win) + + lo = max(floor_win_pages, 2) + if dsv41_pool_bytes(_sizes(lo), args, n_scratch) > available_bytes: + raise ValueError( + f"DSV41 KV budget {available_bytes} bytes cannot fit the minimal pool ({lo} pages incl. the " + f"window working-set floor {floor_win_pages}); raise memory_ratio or lower max_running_req/max_seq_len" + ) + hi = max(lo, available_bytes // max(1, dsv41_cache_per_page(args, 0.0, P))) + while dsv41_pool_bytes(_sizes(hi), args, n_scratch) <= available_bytes: + hi *= 2 + while lo < hi - 1: + mid = (lo + hi) // 2 + if dsv41_pool_bytes(_sizes(mid), args, n_scratch) <= available_bytes: + lo = mid + else: + hi = mid + return _sizes(lo) + + +def dsv41_auto_cost_model(args, swa_ratio: float, floor_win_pages: int, P: int, n_scratch: int = 1): + """Affine ``(cache_per_page, fixed_cache_size, min_reserve_tokens)`` for the MoE-first auto planner.""" + per_page = dsv41_cache_per_page(args, swa_ratio, P) + P * _INT64_BYTES + n0 = max(floor_win_pages, 2) + win0 = max(floor_win_pages, (round(swa_ratio * n0 * P) + P - 1) // P) + base = dsv41_pool_bytes(dsv41_pool_sizes(n0, args, swa_ratio, P, n_win_pages=win0), args, n_scratch) + # The engine combines this structural floor with the configured KV reserve. + return per_page, max(0, base - n0 * per_page), n0 * P + + +# ---- config-facing sizing (EngineConfig in, sizes out) ---- + + +def _dsv41_swa_ratio(config) -> float: + return float(config.swa_full_tokens_ratio) + + +def _dsv41_window_floor_pages(config, args) -> int: + """Minimum window pages: one prefill chunk's reach (capped at 8 pages) + the running requests' + working set (see ``reserved_window_pages``).""" + P = args.window_size + prefill_reach_pages = (config.max_seq_len + P - 1) // P + radix = config.cache_type != "naive" + return min(prefill_reach_pages, 8) + reserved_window_pages(config.max_running_req, radix) + + +def _dsv41_pool_sizes(config, num_pages: int, num_swa_pages: int | None = None) -> DSV41PoolSizes: + """Sizes for ``num_pages`` PHYSICAL pages (dummy included). Window precedence: an explicit + ``num_swa_pages`` > ``config.swa_num_pages_override`` > ``swa_ratio`` x full, floored ONCE in pages.""" + args = config.model_config.dsv41_args + P = args.window_size + swa_ratio = _dsv41_swa_ratio(config) + floor_pages = _dsv41_window_floor_pages(config, args) + target = num_swa_pages if num_swa_pages is not None else config.swa_num_pages_override + if target is not None: + win = min(num_pages, max(floor_pages, int(target) + 1)) + else: + win = max(floor_pages, (round(swa_ratio * num_pages * P) + P - 1) // P) + return dsv41_pool_sizes(num_pages, args, swa_ratio, P, n_win_pages=win) + + +__all__ = [ + "DSV41PoolSizes", + "dsv41_auto_cost_model", + "dsv41_cache_per_page", + "dsv41_kv_unit_bytes", + "dsv41_pool_bytes", + "dsv41_pool_sizes", + "dsv41_solve_num_pages", + "dsv41_window_unit_bytes", + "private_window_layer_ids", +] diff --git a/python/freetoken/kvcache/dsv4/v41_pool.py b/python/freetoken/kvcache/dsv4/v41_pool.py new file mode 100644 index 0000000000..36d3986ae5 --- /dev/null +++ b/python/freetoken/kvcache/dsv4/v41_pool.py @@ -0,0 +1,292 @@ +"""DSV41 paged KV pool (DeepSeek-V4.1): packed byte tiers over the shared page table. + +Tiers, sized from a budget (``v41_cost_model``) not from ``num_requests``: + +* ``window_pool[L]`` -- every layer; the P-sliding KV ring in ``win_fmt`` (fp8) rows, page-granular. +* ``main_pool[src]`` -- per kv-source layer; compressed KV latents in ``main_fmt`` (fp4) rows. +* ``idx_pool[src]`` -- per kv-source layer; indexer keys in ``idx_fmt`` (fp4) rows. +* ``state_ring[src]`` -- ratio>1 sources; the per-window-page compress-state ring (fp32 ``kv|score``). + +Consumer layers (Reindex / Reuse modes) alias their source's pools: ``main_pool_of(layer)``. + +Addressing is the shared ``WindowTierPagedPool``'s (``kvcache/dsv4/window_tier.py``, also under the DSV4 +pool): the page table maps ``(table_idx, pos)`` to a full-token loc; ``full_to_window`` maps that to a +window slot (page-bound by the free list); the main / index rows are pure arithmetic ``full_loc // +ratio``; the ring slot derives from the window slot. Scratch rows past every main / index pool's +arithmetic range give a graph-safe destination for the decode rows whose compressed group did not +complete this step. + +Rows are written through ``pack_rows`` (quantize + scatter) and read inside the attention / +indexer kernels; ``read_*`` helpers dequantize for tests and torch reference paths. +""" + +from __future__ import annotations + +import torch + +from freetoken.utils import init_logger + +from .v41_cost_model import ( + IDX_FMT, + MAIN_FMT, + WIN_FMT, + DSV41PoolSizes, + _idx_row_bytes, + _main_row_bytes, + _win_row_bytes, + dsv41_kv_unit_bytes, + dsv41_window_unit_bytes, + private_window_layer_ids, +) +from .window_tier import CompressStateRing, WindowTierPagedPool + +logger = init_logger(__name__) + + +class DSV41PagedKVCache(WindowTierPagedPool): + win_fmt = WIN_FMT + main_fmt = MAIN_FMT + idx_fmt = IDX_FMT + + def __init__( + self, + sizes: DSV41PoolSizes, + args, + device: torch.device, + dtype: torch.dtype = torch.bfloat16, + n_scratch: int = 1, + ) -> None: + for fmt, dim in ((self.win_fmt, args.head_dim), (self.main_fmt, args.head_dim), (self.idx_fmt, args.index_head_dim)): + fmt.validate_dim(dim) + for src in args.backbone_kv_sources: + if args.window_size % args.compress_ratios[src]: + raise ValueError(f"window {args.window_size} must be a multiple of compress ratio {args.compress_ratios[src]}") + self.args = args + self.private_window_layer_ids = private_window_layer_ids(args) + self.sizes = sizes + self._device = device + self._dtype = dtype # the model's compute dtype; the tiers themselves are packed bytes + self.P = args.window_size + self.head_dim = args.head_dim + self.index_head_dim = args.index_head_dim + self.n_scratch = int(n_scratch) + self._paged_params: tuple[int, bool] | None = None + self.full_loc_map: torch.Tensor | None = None + self._alloc_buffers() + + # ----- buffers ----- + def _alloc_buffers(self) -> None: + sizes, args, dev = self.sizes, self.args, self._device + self.full_to_window = torch.full((sizes.full_token + 1,), -1, dtype=torch.int64, device=dev) + # shared layers: the page-bound window pool; private layers: one ring of P slots per + # page-table row (n_scratch rows, the dummy included), addressed row * P + pos % P + self.window_pool: list[torch.Tensor] = [ + torch.zeros(self.n_scratch * self.P if self.is_private_window(l) else sizes.n_win_slots, _win_row_bytes(args), dtype=torch.uint8, device=dev) + for l in range(args.n_layers) + ] + self.main_pool: dict[int, torch.Tensor] = {} + self.idx_pool: dict[int, torch.Tensor] = {} + self.state_ring: dict[int, CompressStateRing] = {} + self.scratch_base: dict[int, int] = {} + for src in args.backbone_kv_sources: + rows = sizes.main_rows[src] + self.scratch_base[src] = rows + self.main_pool[src] = torch.zeros(rows + self.n_scratch, _main_row_bytes(args), dtype=torch.uint8, device=dev) + self.idx_pool[src] = torch.zeros(sizes.idx_rows[src] + self.n_scratch, _idx_row_bytes(args), dtype=torch.uint8, device=dev) + if src in sizes.state_slots: + self.state_ring[src] = CompressStateRing( + n_slots=sizes.state_slots[src], ring_size=args.compress_ratios[src], overlap=False, + head_dim=args.head_dim, device=dev, + ) + + def total_bytes(self) -> int: + n = self.full_to_window.numel() * self.full_to_window.element_size() + n += sum(t.numel() for t in self.window_pool) + n += sum(t.numel() for t in self.main_pool.values()) + n += sum(t.numel() for t in self.idx_pool.values()) + n += sum(r.buffer.numel() * r.buffer.element_size() for r in self.state_ring.values()) + return int(n) + + @property + def prefix_replay_tokens(self) -> int: + """Prompt-tail tokens a prefix hit must leave to the prefill: bounded replay runs the decoder on + the prompt's last window from that window's encoder outputs, which no cache keeps.""" + return self.P if self.private_window_layer_ids else 0 + + # ----- private window rings ----- + def is_private_window(self, layer_id: int) -> bool: + return layer_id in self.private_window_layer_ids + + def ring_slots(self, table_rows: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: + """Ring slots of ``positions`` (int64, broadcastable with ``table_rows``) for private window layers; + negative positions map to ``-1``.""" + slots = table_rows.to(torch.int64) * self.P + positions % self.P + return torch.where(positions < 0, -1, slots) + + # ----- source aliasing ----- + def source_of(self, layer_id: int) -> int: + src = self.args.roles[layer_id].kv_source + assert src is not None, f"layer {layer_id} is window-only" + return src + + def main_pool_of(self, layer_id: int) -> torch.Tensor: + return self.main_pool[self.source_of(layer_id)] + + def idx_pool_of(self, layer_id: int) -> torch.Tensor: + return self.idx_pool[self.source_of(layer_id)] + + def ratio_of(self, layer_id: int) -> int: + return self.args.roles[layer_id].ratio + + # ----- engine-facing sizing / rebuild surface ----- + @classmethod + def kv_cost(cls, config) -> tuple[int, int, int, int]: + from .v41_cost_model import _dsv41_swa_ratio, _dsv41_window_floor_pages, dsv41_auto_cost_model + + args = config.model_config.dsv41_args + P = args.window_size + per_page, fixed, min_reserve = dsv41_auto_cost_model( + args, _dsv41_swa_ratio(config), _dsv41_window_floor_pages(config, args), P, n_scratch=config.max_running_req + 1 + ) + return per_page, fixed, config.page_size, min_reserve + + @classmethod + def solve_num_pages(cls, config, available_memory: int) -> int: + from freetoken.utils import mem_GB + + from .v41_cost_model import ( + _dsv41_pool_sizes, + _dsv41_swa_ratio, + _dsv41_window_floor_pages, + dsv41_pool_bytes, + dsv41_solve_num_pages, + ) + + args = config.model_config.dsv41_args + P = args.window_size + num_pages = config.num_page_override + if num_pages is None: + sizes = dsv41_solve_num_pages( + available_memory, args, _dsv41_swa_ratio(config), _dsv41_window_floor_pages(config, args), P, + n_scratch=config.max_running_req + 1, + ) + num_pages = sizes.full_token // P - 1 # one physical page is the dummy + else: + floor = _dsv41_window_floor_pages(config, args) + if num_pages < floor: + raise ValueError( + f"--num-pages {num_pages} ({num_pages * P} tokens) is below the DSV41 window working-set " + f"floor {floor} pages ({floor * P} tokens); raise --num-pages or lower max_running_req/max_seq_len" + ) + sizes = _dsv41_pool_sizes(config, num_pages + 1) + assert num_pages > 1, "Not enough memory for KV cache, try reducing --num-pages" + real = dsv41_pool_bytes(sizes, args, config.max_running_req + 1) + logger.info( + f"Allocating {num_pages * P} tokens for DSV41 KV cache ({sizes.n_win_pages} window pages), total = {mem_GB(real)}" + ) + return num_pages + + @classmethod + def window_spec(cls, config): + from ..base import WindowPoolSpec + from .v41_cost_model import _dsv41_window_floor_pages + + args = config.model_config.dsv41_args + return WindowPoolSpec(args.window_size, _dsv41_window_floor_pages(config, args) - 1) + + @classmethod + def min_kv_tokens(cls, config) -> int: + from .v41_cost_model import _dsv41_window_floor_pages + + args = config.model_config.dsv41_args + return _dsv41_window_floor_pages(config, args) * args.window_size + + def validate_rebuild( + self, config, *, num_pages: int | None, target_moe: int, per_expert_bytes: int, + baseline_free: int, weights_bytes: int, current_num_pages: int, + extra_fixed_bytes: int = 0, extra_note: str = "", + num_swa_pages: int | None = None, **targets, + ) -> None: + from freetoken.engine.cache_budget import net_cache_budget_bytes + from freetoken.utils import mem_GB + + from ..base import CacheRebuildRejected + from .v41_cost_model import _dsv41_pool_sizes, _dsv41_window_floor_pages, dsv41_pool_bytes + + if num_pages is not None: + floor = _dsv41_window_floor_pages(config, self.args) + if num_pages < floor: + raise CacheRebuildRejected( + f"num_pages {num_pages} is below the DSV41 window working-set floor {floor} " + f"(max_running_req={config.max_running_req}); admission would deadlock" + ) + if num_pages is not None or num_swa_pages is not None: + target_pages = num_pages if num_pages is not None else current_num_pages + kv_sizes = _dsv41_pool_sizes(config, target_pages + 1, num_swa_pages=num_swa_pages) + else: + kv_sizes = self.sizes + budget = net_cache_budget_bytes(config.memory_ratio, baseline_free, weights_bytes, 0) + need = target_moe * per_expert_bytes + dsv41_pool_bytes(kv_sizes, self.args, config.max_running_req + 1) + if need > budget: + kv_part = f"kv={num_pages} P-pages" if num_pages is not None else "kv=current pool" + raise CacheRebuildRejected( + f"requested cache (moe={target_moe} slots, {kv_part}) needs {mem_GB(need)} > budget {mem_GB(budget)}; " + "old cache kept, still serving" + ) + + def rebuild_from_config(self, config, num_pages: int, *, num_swa_pages: int | None = None) -> None: + from .v41_cost_model import _dsv41_pool_sizes + + self.rebuild(_dsv41_pool_sizes(config, num_pages + 1, num_swa_pages=num_swa_pages)) + + def _drop_buffers(self) -> None: + self.window_pool = self.main_pool = self.idx_pool = self.state_ring = None # type: ignore[assignment] + + def unit_bytes(self) -> tuple[int, int]: + return dsv41_kv_unit_bytes(self.args, self.P), dsv41_window_unit_bytes(self.args, self.P) + + # ----- writes (quantize + scatter) ----- + def store_window(self, kv: torch.Tensor, layer_id: int, window_slot: torch.Tensor) -> None: + from freetoken.kernel.triton.dsv41.pack import pack_rows + + pack_rows(kv, self.win_fmt, pool=self.window_pool[layer_id], row_ids=window_slot) + + def store_main(self, latent: torch.Tensor, source: int, rows: torch.Tensor) -> None: + from freetoken.kernel.triton.dsv41.pack import pack_rows + + pack_rows(latent, self.main_fmt, pool=self.main_pool[source], row_ids=rows) + + def store_index(self, k: torch.Tensor, source: int, rows: torch.Tensor) -> None: + from freetoken.kernel.triton.dsv41.pack import pack_rows + + pack_rows(k, self.idx_fmt, pool=self.idx_pool[source], row_ids=rows) + + # ----- reads (gather + dequantize; tests and torch reference paths) ----- + def read_window(self, layer_id: int, window_slot: torch.Tensor) -> torch.Tensor: + from freetoken.kernel.triton.dsv41.pack import unpack_rows + + return unpack_rows(self.window_pool[layer_id], self.win_fmt, self.head_dim, window_slot) + + def read_main(self, source: int, rows: torch.Tensor) -> torch.Tensor: + from freetoken.kernel.triton.dsv41.pack import unpack_rows + + return unpack_rows(self.main_pool[source], self.main_fmt, self.head_dim, rows) + + def read_index(self, source: int, rows: torch.Tensor) -> torch.Tensor: + from freetoken.kernel.triton.dsv41.pack import unpack_rows + + return unpack_rows(self.idx_pool[source], self.idx_fmt, self.index_head_dim, rows) + + # ----- compress-state ring accessors ----- + def get_state(self, source: int, state_loc: torch.Tensor) -> torch.Tensor: + return self.state_ring[source].get(state_loc) + + def set_state(self, source: int, state_loc: torch.Tensor, kv_score: torch.Tensor) -> None: + self.state_ring[source].set(state_loc, kv_score) + + @property + def num_layers(self) -> int: + return self.args.n_layers + + +__all__ = ["DSV41PagedKVCache"] diff --git a/python/freetoken/kvcache/dsv4/v41_row_format.py b/python/freetoken/kvcache/dsv4/v41_row_format.py new file mode 100644 index 0000000000..f273be4f56 --- /dev/null +++ b/python/freetoken/kvcache/dsv4/v41_row_format.py @@ -0,0 +1,69 @@ +"""RowFormat: how one DeepSeek-V4.1 KV row (a head_dim-wide vector) is stored in a packed byte pool. + +DeepSeek-V4.1 keeps three cache tiers in three storage formats (tech report sec. 2.4.4 and +``inference/model.py``): + +* SWA KV -- ``fp8_e8m0_b32``: e4m3 values, one ue8m0 (power-of-two) scale per 32 channels +* main KV -- ``fp4_e4m3_b16``: e2m1 values, one e4m3 scale per 16 channels (NVFP4 without the global scale) +* indexer K -- ``fp4_e8m0_b32``: e2m1 values, one ue8m0 scale per 32 channels (OCP MXFP4) + +A row is laid out as ``[values | scales]``: ``value_bytes(dim)`` bytes of (packed) values followed by +``scale_bytes(dim)`` scale bytes, so one gather fetches both. e2m1 pairs pack the even channel into the low +nibble and the odd channel into the high nibble (``torch.float4_e2m1fn_x2`` / ``convert.py`` order). + +The Triton side (``kernel/triton/dsv41/row_format.py``) selects the dequant path from ``RowFormat.code`` +(a constexpr); the pack / unpack wrappers in ``kernel/triton/dsv41/pack.py`` reproduce the reference +``act_quant`` / ``fp4_act_quant`` numerics exactly, including the bf16 rounding of the dequantized value +the reference bakes into its cache. +""" + +from __future__ import annotations + +from dataclasses import dataclass + + +@dataclass(frozen=True) +class RowFormat: + name: str + elem_bits: int # 16 (bf16), 8 (e4m3) or 4 (e2m1) + group: int # channels per scale; 0 = unscaled + scale: str | None # None | "e8m0" | "e4m3" + code: int # Triton constexpr selector + + def value_bytes(self, dim: int) -> int: + assert (dim * self.elem_bits) % 8 == 0, (dim, self.elem_bits) + return dim * self.elem_bits // 8 + + def scale_bytes(self, dim: int) -> int: + if not self.group: + return 0 + assert dim % self.group == 0, (dim, self.group) + return dim // self.group + + def row_bytes(self, dim: int) -> int: + return self.value_bytes(dim) + self.scale_bytes(dim) + + def num_groups(self, dim: int) -> int: + return dim // self.group if self.group else 0 + + def validate_dim(self, dim: int) -> None: + if dim <= 0 or (dim * self.elem_bits) % 8 or (self.group and dim % self.group): + raise ValueError(f"row format {self.name} cannot hold a {dim}-wide row") + + +BF16 = RowFormat("bf16", elem_bits=16, group=0, scale=None, code=0) +FP8_E8M0_B32 = RowFormat("fp8_e8m0_b32", elem_bits=8, group=32, scale="e8m0", code=1) +FP4_E4M3_B16 = RowFormat("fp4_e4m3_b16", elem_bits=4, group=16, scale="e4m3", code=2) +FP4_E8M0_B32 = RowFormat("fp4_e8m0_b32", elem_bits=4, group=32, scale="e8m0", code=3) + +ROW_FORMATS: dict[str, RowFormat] = {f.name: f for f in (BF16, FP8_E8M0_B32, FP4_E4M3_B16, FP4_E8M0_B32)} + + +def row_format(name: str) -> RowFormat: + try: + return ROW_FORMATS[name] + except KeyError: + raise ValueError(f"unknown row format {name!r}; one of {sorted(ROW_FORMATS)}") from None + + +__all__ = ["RowFormat", "BF16", "FP8_E8M0_B32", "FP4_E4M3_B16", "FP4_E8M0_B32", "ROW_FORMATS", "row_format"] diff --git a/python/freetoken/kvcache/dsv4_cost_model.py b/python/freetoken/kvcache/dsv4/v4_cost_model.py similarity index 96% rename from python/freetoken/kvcache/dsv4_cost_model.py rename to python/freetoken/kvcache/dsv4/v4_cost_model.py index 2c084aff84..c6f5b7ce41 100644 --- a/python/freetoken/kvcache/dsv4_cost_model.py +++ b/python/freetoken/kvcache/dsv4/v4_cost_model.py @@ -28,14 +28,12 @@ def dsv4_reserved_window_pages(max_running_req: int, radix: bool) -> int: - """Window pages the sliding pool must always keep for the concurrent working set: each - running request's decode transients (2 per req + dummy) plus, in radix mode, PER concurrent - request one locked live-tail page AND a retained (soft-pinned) prompt-end window -- the - window is 2 pages here because the retention gap page-aligns to a whole extra page at - P==window==128, so a distinct follow-up per running request can re-lock 2 pages each. - Shared by the engine's window-floor and the manager's prefill_chunk_budget so both - reserve the same set.""" - return 2 * (max_running_req + 1) + (3 * max_running_req if radix else 0) + 1 + """Window pages the sliding pool must always keep for the concurrent working set; the + formula is shared with every sliding-window paged pool (``window_tier.reserved_window_pages``) + and by the engine's window-floor and the manager's prefill_chunk_budget.""" + from .window_tier import reserved_window_pages + + return reserved_window_pages(max_running_req, radix) def ring_size_for_ratio(ratio: int) -> int: diff --git a/python/freetoken/kvcache/dsv4_paged_pool.py b/python/freetoken/kvcache/dsv4/v4_pool.py similarity index 52% rename from python/freetoken/kvcache/dsv4_paged_pool.py rename to python/freetoken/kvcache/dsv4/v4_pool.py index 3ce23eea7d..d5eae6d080 100644 --- a/python/freetoken/kvcache/dsv4_paged_pool.py +++ b/python/freetoken/kvcache/dsv4/v4_pool.py @@ -23,135 +23,22 @@ from freetoken.utils import init_logger -from .base import BaseKVCachePool -from .dsv4_cost_model import ( +from .v4_cost_model import ( DSV4PoolSizes, dsv4_kv_unit_bytes, dsv4_window_unit_bytes, ring_size_for_ratio, ) +# The window-tier building blocks and the pool base live in window_tier.py (shared with the DSV41 +# pool); the first two are re-exported here for the DSV4 callers and tests that import them from +# this module. +from .window_tier import CompressStateRing, FreeListAllocator, WindowTierPagedPool # noqa: F401 logger = init_logger(__name__) -# LIFO free-list allocator for the window tier: one allocated unit spans ``page_unit = P`` -# slots, so a unit base is always a multiple of ``P`` (the spec's ``G*P`` page-base -# invariant). The cmp/idx tiers have no allocator -- their rows are ``full_loc // ratio`` -# pure arithmetic. -class FreeListAllocator: - def __init__( - self, - capacity: int, - device: torch.device, - page_unit: int = 1, - ) -> None: - assert capacity % page_unit == 0, ( - f"capacity {capacity} must be a multiple of page_unit {page_unit}" - ) - self._capacity = int(capacity) - self._page_unit = int(page_unit) - self._device = device - # Unit base slots: [0, page_unit, 2*page_unit, ...]. LIFO -> pop the tail. - self._free = self._fresh_free() - - def _fresh_free(self) -> torch.Tensor: - n_units = self._capacity // self._page_unit - return torch.arange(n_units, dtype=torch.int64, device=self._device) * self._page_unit - - def alloc(self, n_units: int) -> torch.Tensor: - """Return ``n_units`` unit base slots (each a multiple of ``page_unit``).""" - if n_units < 0: - raise ValueError(f"n_units must be non-negative, got {n_units}") - if n_units > self._free.numel(): - raise RuntimeError( - f"FreeListAllocator out of slots: requested {n_units} units, " - f"have {self._free.numel()} (capacity {self._capacity}, unit {self._page_unit})" - ) - if n_units == 0: - return self._free[:0].clone() - taken = self._free[-n_units:].clone() - self._free = self._free[:-n_units] - return taken - - def free(self, units: torch.Tensor) -> None: - """Return previously-allocated unit base slots to the pool (LIFO recycle).""" - if units.numel() == 0: - return - units = units.to(device=self._device, dtype=torch.int64).reshape(-1) - self._free = torch.cat([self._free, units]) - - def available(self) -> int: - """Free capacity in slots (units * page_unit).""" - return int(self._free.numel()) * self._page_unit - - @property - def capacity(self) -> int: - return self._capacity - - def reset(self) -> None: - self._free = self._fresh_free() - - - -class CompressStateRing: - """Per-layer fp32 compress-state ring: ``[n_slots + 1, 2*(1+overlap)*head_dim]``. - - Last row (index ``-1``) is a permanent scratch slot, re-cleared on every - write. Last dim is split ``kv | score``; ``set_state`` writes both halves. - """ - - def __init__( - self, - n_slots: int, - ring_size: int, - overlap: bool, - head_dim: int, - device: torch.device, - dtype: torch.dtype = torch.float32, - ) -> None: - self.ring_size = ring_size - self.head_dim = head_dim - self._item_size = (1 + int(overlap)) * head_dim # width of kv (== width of score) - last_dim = 2 * self._item_size - # +1 trailing scratch row at index -1. - self.buffer = torch.zeros((n_slots + 1, last_dim), dtype=dtype, device=device) - self.n_slots = n_slots - self._clear_scratch() - - def _clear_scratch(self) -> None: - self.buffer[-1, : self._item_size].zero_() - self.buffer[-1, self._item_size :].fill_(float("-inf")) - - @property - def item_size(self) -> int: - return self._item_size - - def get(self, state_loc: torch.Tensor) -> torch.Tensor: - """Gather ``kv_score`` rows at ``state_loc`` (``-1`` -> scratch row).""" - return self.buffer[state_loc] - - def set(self, state_loc: torch.Tensor, kv_score: torch.Tensor) -> None: - """Scatter ``kv_score`` rows to ``state_loc`` then re-clear the scratch row.""" - self.buffer[state_loc] = kv_score - self._clear_scratch() - - def get_blocks(self, page_base: torch.Tensor) -> torch.Tensor: - """Batched per-row carry-block read. ``page_base`` is ``[B]`` (the ring block base - row per row, == ``(window_slot // P) * ring_size``). Returns ``[B, ring_size, 2*item]`` - — each row's whole rolling carry block. Distinct pages -> disjoint blocks (isolation).""" - rows = page_base[:, None] + torch.arange(self.ring_size, device=page_base.device) - return self.buffer[rows] - - def set_blocks(self, page_base: torch.Tensor, blocks: torch.Tensor) -> None: - """Batched per-row carry-block write: ``buffer[base+arange(ring_size)] = blocks[row]``. - ``page_base`` ``[B]``, ``blocks`` ``[B, ring_size, 2*item]``. Re-clears scratch.""" - rows = page_base[:, None] + torch.arange(self.ring_size, device=page_base.device) - self.buffer[rows] = blocks - self._clear_scratch() - - -class DSV4PagedKVCache(BaseKVCachePool): +class DSV4PagedKVCache(WindowTierPagedPool): def __init__( self, sizes: DSV4PoolSizes, @@ -273,50 +160,11 @@ def _alloc_buffers(self) -> None: ) ) - # ----- full-loc translation (gather-only -1 safety; see field comment) ----- - def translate_full_to_window(self, full_locs: torch.Tensor) -> torch.Tensor: - # int64 gather indices: the shared page_table stores full locs as int32. - return self.full_to_window[full_locs.to(dtype=torch.int64)] - - @staticmethod - def cmp_rows(full_locs: torch.Tensor, ratio: int) -> torch.Tensor: - return torch.div(full_locs.to(dtype=torch.int64), ratio, rounding_mode="floor") - - def bind_window_pages(self, full_page_base: int, window_page_base: int) -> None: - assert full_page_base % self.P == 0 and window_page_base % self.P == 0 - self.full_to_window[full_page_base : full_page_base + self.P] = torch.arange( - window_page_base, window_page_base + self.P, dtype=torch.int64, device=self._device - ) - - def unbind_window_pages(self, full_locs: torch.Tensor) -> None: - self.full_to_window[full_locs[full_locs >= 0]] = -1 - - # ----- generic swa_pool duck-type (the CacheManager plug-in surface) ----- - # ShadowRadix layering: the shared page_table is the virtual full-token coordinate; this pool - # projects it into the physical tiers. The window tier is the managed "second currency" -- - # token-face signatures (what the generic CacheManager speaks), PAGE-ATOMIC internals (window - # pages are 1:1 page-bound to full pages; the per-page state ring requires it). alloc_swa - # receives whole ascending pages (allocate_paged's _page_to_token expansion) and free paths - # are page-complete by construction (padded finish tails, align_down frontiers, page-aligned - # tree nodes) -- asserted here, not assumed. - swa_paged = True - - @property - def sliding_window_size(self) -> int: - return self.P - - @property - def prefill_chunk_budget(self) -> int: - return self._chunk_budget - # ----- engine-facing rebuild surface ----- - # The tier buffers are bound into per-forward model scratch, invalid after a realloc. - needs_rebind_on_rebuild = True - @classmethod def kv_cost(cls, config) -> tuple[int, int, int, int]: - from .dsv4_cost_model import _dsv4_swa_ratio, _dsv4_window_floor_pages - from .dsv4_cost_model import dsv4_auto_cost_model + from .v4_cost_model import _dsv4_swa_ratio, _dsv4_window_floor_pages + from .v4_cost_model import dsv4_auto_cost_model dsv4_args = config.model_config.dsv4_args P = dsv4_args.window_size @@ -335,8 +183,8 @@ def solve_num_pages(cls, config, available_memory: int) -> int: # graceful config error, not a late OOM. from freetoken.utils import mem_GB - from .dsv4_cost_model import _dsv4_pool_sizes, _dsv4_swa_ratio, _dsv4_window_floor_pages - from .dsv4_cost_model import dsv4_pool_bytes, dsv4_solve_num_pages + from .v4_cost_model import _dsv4_pool_sizes, _dsv4_swa_ratio, _dsv4_window_floor_pages + from .v4_cost_model import dsv4_pool_bytes, dsv4_solve_num_pages dsv4_args = config.model_config.dsv4_args P = dsv4_args.window_size @@ -369,11 +217,19 @@ def solve_num_pages(cls, config, available_memory: int) -> int: ) return num_pages + @classmethod + def window_spec(cls, config): + from ..base import WindowPoolSpec + from .v4_cost_model import _dsv4_window_floor_pages + + P = config.model_config.dsv4_args.window_size + return WindowPoolSpec(P, _dsv4_window_floor_pages(config, P) - 1) + @classmethod def min_kv_tokens(cls, config) -> int: # The full anchor must cover the window working-set floor (full >= window always), so # that floor -- the value validate_rebuild enforces -- is the pool's floor in tokens. - from .dsv4_cost_model import _dsv4_window_floor_pages + from .v4_cost_model import _dsv4_window_floor_pages P = config.model_config.dsv4_args.window_size return _dsv4_window_floor_pages(config, P) * P @@ -387,9 +243,9 @@ def validate_rebuild( from freetoken.engine.cache_budget import net_cache_budget_bytes from freetoken.utils import mem_GB - from .base import CacheRebuildRejected - from .dsv4_cost_model import _dsv4_pool_sizes, _dsv4_window_floor_pages - from .dsv4_cost_model import dsv4_pool_bytes + from ..base import CacheRebuildRejected + from .v4_cost_model import _dsv4_pool_sizes, _dsv4_window_floor_pages + from .v4_cost_model import dsv4_pool_bytes dsv4_args = config.model_config.dsv4_args if num_pages is not None: @@ -426,120 +282,19 @@ def validate_rebuild( def rebuild_from_config( self, config, num_pages: int, *, num_swa_pages: int | None = None ) -> None: - from .dsv4_cost_model import _dsv4_pool_sizes + from .v4_cost_model import _dsv4_pool_sizes # +1 for the dummy page self.rebuild(_dsv4_pool_sizes(config, num_pages + 1, num_swa_pages=num_swa_pages)) - def attach_page_table(self, page_table: torch.Tensor) -> None: - # The model reads full locs through full_loc_map; under the shared route that IS the - # page table. The attention backend re-allocates its decode snapshot on re-capture. - self.full_loc_map = page_table - def unit_bytes(self) -> tuple[int, int]: # No measurable flat buffer (owned paged pool): the full (cmp/idx + mapping) and window # (sliding KV + state rings) per-token costs come from the per-tier cost model. return dsv4_kv_unit_bytes(self.args, self.P), dsv4_window_unit_bytes(self.args, self.P) - def rebuild(self, sizes) -> None: - """In-place resize to ``sizes`` (identity-preserving; free-before-alloc). The manager's - tree/page bookkeeping reset is the scheduler's generic cache_manager.rebuild; the engine - re-attaches the page table via attach_page_table afterwards.""" - import gc - - assert self._paged_params is not None, "rebuild before _init_paged_state" - self.sizes = sizes + def _drop_buffers(self) -> None: self.window_pool = self.cmp_pool = self.idx_pool = None self.state_ring = self.indexer_state_ring = None - self.full_to_window = None - gc.collect() - if torch.cuda.is_available(): - torch.cuda.empty_cache() - self._alloc_buffers() - self._init_paged_state(*self._paged_params) - - def _init_paged_state(self, max_running_req: int, radix: bool) -> None: - """Build the pool-owned window free-list + the tail dummy binding. The LAST full page and - LAST window page are the reserved dummy region: page_table's dummy row points at - ``full_token - P`` (== the generic ``fill_(num_tokens)`` convention with num_tokens = the - allocatable token count), permanently bound so graph-padded rows scatter to a real slot.""" - from .dsv4_cost_model import dsv4_reserved_window_pages - P = self.P - self._paged_params = (int(max_running_req), bool(radix)) - self.full_to_window.fill_(-1) - self._win_alloc = FreeListAllocator(self.sizes.n_win_slots - P, self._device, page_unit=P) - self.bind_window_pages(self.sizes.full_token - P, self.sizes.n_win_slots - P) - # Chunk cap: a batched prefill holds the whole chunk's window live at once (sliding frees - # only between chunks; peak ~2x the chunk), so reserve the concurrent working set and - # halve the rest -- the same formula the bespoke manager used. - n_win_pages = (self.sizes.n_win_slots // P) - 1 - reserved = dsv4_reserved_window_pages(max_running_req, radix) - self._chunk_budget = max(P, (n_win_pages - reserved) // 2 * P) - - @property - def swa_num_tokens(self) -> int: - # Allocatable window slots + 1: the generic capacity convention reserves slot 0 as a - # sentinel (cap == swa_num_tokens - 1); DSV4's reserved unit is the tail dummy page, - # already excluded from the free-list, so +1 re-encodes the same cap. - return (self.sizes.n_win_slots - self.P) + 1 - - def swa_available_size(self) -> int: - return int(self._win_alloc.available()) - - def alloc_swa(self, full_indices: torch.Tensor) -> None: - """Bind one window page per incoming FULL page. ``full_indices`` must be whole ascending - pages (the ``_page_to_token`` expansion); the in-page offsets are preserved - (``window_slot = wbase + pos % P``), which the state ring's page-block layout requires.""" - n = int(full_indices.numel()) - if n == 0: - return - P = self.P - assert n % P == 0, f"alloc_swa needs whole pages, got {n} slots" - fi = full_indices.to(device=self._device, dtype=torch.int64).view(-1, P) - fbases = fi[:, 0] - assert torch.equal(fi, fbases[:, None] + torch.arange(P, device=self._device)), ( - "alloc_swa pages must be contiguous ascending" - ) - wbases = self._win_alloc.alloc(fbases.numel()) # raises when exhausted (caller gated) - offsets = torch.arange(P, dtype=torch.int64, device=self._device) - self.full_to_window[(fbases[:, None] + offsets).flatten()] = ( - wbases[:, None] + offsets - ).flatten() - - def free_swa(self, full_indices: torch.Tensor) -> None: - """Return the window pages backing these FULL locs and unbind the mapping. Page-atomic: - the incoming locs must cover each touched page completely (guaranteed by the padded - finish tails / aligned frontiers / page-aligned tree values). Idempotent over already - unbound (slid/tombstoned) pages.""" - if full_indices.numel() == 0: - return - P = self.P - fi = full_indices.to(device=self._device, dtype=torch.int64) - fi = fi[fi >= 0] - if fi.numel() == 0: - return - fbases, counts = torch.unique( - torch.div(fi, P, rounding_mode="floor") * P, return_counts=True - ) - assert bool((counts == P).all()), ( - f"free_swa got partial pages (counts {counts[counts != P].tolist()[:4]})" - ) - ws = self.full_to_window[fbases] - live = ws[ws >= 0] - offsets = torch.arange(P, dtype=torch.int64, device=self._device) - self.full_to_window[(fbases[:, None] + offsets).flatten()] = -1 - if live.numel(): - self._win_alloc.free(torch.div(live, P, rounding_mode="floor") * P) - - def translate_loc_from_full_to_swa(self, kv_indices: torch.Tensor) -> torch.Tensor: - return self.full_to_window[kv_indices.to(dtype=torch.int64)] - - # ----- state_loc derivation (vectorized, LongTensor in/out) ----- - @staticmethod - def state_loc(window_slot: torch.Tensor, ring_size: int, P: int) -> torch.Tensor: - pages = torch.div(window_slot, P, rounding_mode="floor") - loc = pages * ring_size + (window_slot % ring_size) - return torch.where(window_slot < 0, torch.full_like(loc, -1), loc) def ring_size(self, layer_id: int) -> int: return ring_size_for_ratio(self.compress_ratios[layer_id]) @@ -586,26 +341,6 @@ def total_bytes(self) -> int: ) return int(n) - # ----- BaseKVCachePool interface ----- - def k_cache(self, index: int) -> torch.Tensor: - return self.window_pool[index] - - def v_cache(self, index: int) -> torch.Tensor: # MLA: K == V (single latent) - return self.window_pool[index] - - def store_kv(self, k, v, out_loc, layer_id) -> None: - # Thin window-write shim for ABC compat; DSV4 writes via the specialized - # setters above. K == V (single latent), so the window slot is out_loc. - self.store_window(k, layer_id, out_loc) - - @property - def device(self) -> torch.device: - return self._device - - @property - def dtype(self) -> torch.dtype: - return self._dtype - @property def num_layers(self) -> int: return self._n_layers diff --git a/python/freetoken/kvcache/dsv4/window_tier.py b/python/freetoken/kvcache/dsv4/window_tier.py new file mode 100644 index 0000000000..0a25c05872 --- /dev/null +++ b/python/freetoken/kvcache/dsv4/window_tier.py @@ -0,0 +1,331 @@ +"""Shared pieces of the sliding-window paged pools (DSV4, DSV41). + +* ``FreeListAllocator`` -- LIFO page allocator for the window tier (unit = one P-slot page). +* ``CompressStateRing`` -- per-window-page fp32 ``kv | score`` carry ring with a scratch row. +* ``reserved_window_pages`` -- the window working set every such pool keeps for its running requests. +* ``window_state_loc`` -- the derived (never stored) ring slot of a window slot. +* ``WindowTierPagedPool`` -- the pool base: the full-loc -> window-slot mapping over the shared page + table, the page-atomic window allocator behind the CacheManager's ``swa_pool`` duck-type, the + in-place rebuild, and the ABC glue every such pool shares. Subclasses allocate their own tiers + (``_alloc_buffers`` / ``_drop_buffers``), size themselves (the ``kv_cost`` family) and write rows. + +``DSV4PagedKVCache`` (the origin of this code) and ``DSV41PagedKVCache`` both build on these, so a +fix to allocation, freeing or the mapping lands in one place. +""" + +from __future__ import annotations + +import gc +from abc import abstractmethod + +import torch + +from ..base import BaseKVCachePool + + +class FreeListAllocator: + """LIFO free list over ``capacity // page_unit`` units; a unit's base slot is a multiple of + ``page_unit`` (the window tier's ``G*P`` page-base invariant).""" + + def __init__(self, capacity: int, device: torch.device, page_unit: int = 1) -> None: + assert capacity % page_unit == 0, f"capacity {capacity} must be a multiple of page_unit {page_unit}" + self._capacity = int(capacity) + self._page_unit = int(page_unit) + self._device = device + self._free = self._fresh_free() + + def _fresh_free(self) -> torch.Tensor: + n_units = self._capacity // self._page_unit + return torch.arange(n_units, dtype=torch.int64, device=self._device) * self._page_unit + + def alloc(self, n_units: int) -> torch.Tensor: + """Return ``n_units`` unit base slots (each a multiple of ``page_unit``).""" + if n_units < 0: + raise ValueError(f"n_units must be non-negative, got {n_units}") + if n_units > self._free.numel(): + raise RuntimeError( + f"FreeListAllocator out of slots: requested {n_units} units, " + f"have {self._free.numel()} (capacity {self._capacity}, unit {self._page_unit})" + ) + if n_units == 0: + return self._free[:0].clone() + taken = self._free[-n_units:].clone() + self._free = self._free[:-n_units] + return taken + + def free(self, units: torch.Tensor) -> None: + """Return previously-allocated unit base slots to the pool (LIFO recycle).""" + if units.numel() == 0: + return + units = units.to(device=self._device, dtype=torch.int64).reshape(-1) + self._free = torch.cat([self._free, units]) + + def available(self) -> int: + """Free capacity in slots (units * page_unit).""" + return int(self._free.numel()) * self._page_unit + + @property + def capacity(self) -> int: + return self._capacity + + def reset(self) -> None: + self._free = self._fresh_free() + + +class CompressStateRing: + """Per-layer fp32 compress-state ring: ``[n_slots + 1, 2*(1+overlap)*head_dim]``. + + Last row (index ``-1``) is a permanent scratch slot, re-cleared on every write. Last dim is + split ``kv | score``; ``set`` writes both halves. + """ + + def __init__( + self, + n_slots: int, + ring_size: int, + overlap: bool, + head_dim: int, + device: torch.device, + dtype: torch.dtype = torch.float32, + ) -> None: + self.ring_size = ring_size + self.head_dim = head_dim + self._item_size = (1 + int(overlap)) * head_dim # width of kv (== width of score) + last_dim = 2 * self._item_size + self.buffer = torch.zeros((n_slots + 1, last_dim), dtype=dtype, device=device) + self.n_slots = n_slots + self._clear_scratch() + + def _clear_scratch(self) -> None: + self.buffer[-1, : self._item_size].zero_() + self.buffer[-1, self._item_size :].fill_(float("-inf")) + + @property + def item_size(self) -> int: + return self._item_size + + def get(self, state_loc: torch.Tensor) -> torch.Tensor: + """Gather ``kv_score`` rows at ``state_loc`` (``-1`` -> scratch row).""" + return self.buffer[state_loc] + + def set(self, state_loc: torch.Tensor, kv_score: torch.Tensor) -> None: + """Scatter ``kv_score`` rows to ``state_loc`` then re-clear the scratch row.""" + self.buffer[state_loc] = kv_score + self._clear_scratch() + + def get_blocks(self, page_base: torch.Tensor) -> torch.Tensor: + """Batched per-row carry-block read: ``page_base`` ``[B]`` (``(window_slot // P) * ring_size``) + -> ``[B, ring_size, 2*item]``. Distinct pages -> disjoint blocks.""" + rows = page_base[:, None] + torch.arange(self.ring_size, device=page_base.device) + return self.buffer[rows] + + def set_blocks(self, page_base: torch.Tensor, blocks: torch.Tensor) -> None: + """Batched per-row carry-block write of ``[B, ring_size, 2*item]``; re-clears scratch.""" + rows = page_base[:, None] + torch.arange(self.ring_size, device=page_base.device) + self.buffer[rows] = blocks + self._clear_scratch() + + +def reserved_window_pages(max_running_req: int, radix: bool) -> int: + """Window pages the sliding pool must always keep for the concurrent working set: each + running request's decode transients (2 per req + dummy) plus, in radix mode, PER concurrent + request one locked live-tail page AND a retained (soft-pinned) prompt-end window -- the + window is 2 pages here because the retention gap page-aligns to a whole extra page at + P == window.""" + return 2 * (max_running_req + 1) + (3 * max_running_req if radix else 0) + 1 + + +def window_state_loc(window_slot: torch.Tensor, ring_size: int, P: int) -> torch.Tensor: + """Ring slot of a window slot: ``(ws // P) * ring_size + ws % ring_size``; ``-1`` stays ``-1`` + (the scratch row). ``ring_size | P`` keeps distinct pages on disjoint blocks.""" + pages = torch.div(window_slot, P, rounding_mode="floor") + loc = pages * ring_size + (window_slot % ring_size) + return torch.where(window_slot < 0, torch.full_like(loc, -1), loc) + + +class WindowTierPagedPool(BaseKVCachePool): + """A paged KV pool whose window tier is page-bound to the shared full-token page table. + + Subclasses set ``P`` (the window page), ``_device``, ``_dtype`` and ``sizes`` (with + ``full_token`` and ``n_win_slots``) before calling ``_alloc_buffers``, which must create + ``self.full_to_window`` (``[full_token + 1]`` int64, ``-1`` = unmapped, the trailing row a + permanent ``-1`` sentinel so a gather at ``-1`` is safe) and ``self.window_pool`` (one tensor + per layer, indexed by window slot). Scatter paths must never see a negative slot. + + ShadowRadix layering: the shared page table is the virtual full-token coordinate; the pool + projects it into physical tiers. The window tier is the managed "second currency" -- token-face + signatures (what the generic CacheManager speaks), PAGE-ATOMIC internals (window pages are 1:1 + page-bound to full pages; the per-page state ring requires it). ``alloc_swa`` receives whole + ascending pages and free paths are page-complete by construction -- asserted here, not assumed. + """ + + swa_paged = True + # the tier buffers are bound into per-forward model scratch, invalid after a realloc + needs_rebind_on_rebuild = True + + P: int + _device: torch.device + _dtype: torch.dtype + sizes: object + full_to_window: torch.Tensor + window_pool: list + + # ----- tiers (subclass) ----- + @abstractmethod + def _alloc_buffers(self) -> None: ... + + @abstractmethod + def _drop_buffers(self) -> None: + """Release every tier buffer before a rebuild reallocates them.""" + + # ----- full-loc translation (gather-only -1 safety; see the class docstring) ----- + def translate_full_to_window(self, full_locs: torch.Tensor) -> torch.Tensor: + # int64 gather indices: the shared page_table stores full locs as int32 + return self.full_to_window[full_locs.to(dtype=torch.int64)] + + @staticmethod + def cmp_rows(full_locs: torch.Tensor, ratio: int) -> torch.Tensor: + """The compressed row of a full loc: ``full_loc // ratio`` (negative stays negative).""" + return torch.div(full_locs.to(dtype=torch.int64), ratio, rounding_mode="floor") + + def bind_window_pages(self, full_page_base: int, window_page_base: int) -> None: + assert full_page_base % self.P == 0 and window_page_base % self.P == 0 + self.full_to_window[full_page_base : full_page_base + self.P] = torch.arange( + window_page_base, window_page_base + self.P, dtype=torch.int64, device=self._device + ) + + def unbind_window_pages(self, full_locs: torch.Tensor) -> None: + self.full_to_window[full_locs[full_locs >= 0]] = -1 + + @staticmethod + def state_loc(window_slot: torch.Tensor, ring_size: int, P: int) -> torch.Tensor: + return window_state_loc(window_slot, ring_size, P) + + # ----- generic swa_pool duck-type (the CacheManager plug-in surface) ----- + @property + def sliding_window_size(self) -> int: + return self.P + + @property + def prefill_chunk_budget(self) -> int: + return self._chunk_budget + + def _init_paged_state(self, max_running_req: int, radix: bool) -> None: + """Build the pool-owned window free list + the tail dummy binding. The LAST full page and + LAST window page are the reserved dummy region: page_table's dummy row points at + ``full_token - P`` (the generic ``fill_(num_tokens)`` convention), permanently bound so + graph-padded rows scatter to a real slot.""" + P = self.P + self._paged_params = (int(max_running_req), bool(radix)) + self.full_to_window.fill_(-1) + self._win_alloc = FreeListAllocator(self.sizes.n_win_slots - P, self._device, page_unit=P) + self.bind_window_pages(self.sizes.full_token - P, self.sizes.n_win_slots - P) + # Chunk cap: a batched prefill holds the whole chunk's window live at once (sliding frees + # only between chunks; peak ~2x the chunk), so reserve the concurrent working set and + # halve the rest. + n_win_pages = (self.sizes.n_win_slots // P) - 1 + reserved = reserved_window_pages(max_running_req, radix) + self._chunk_budget = max(P, (n_win_pages - reserved) // 2 * P) + + @property + def window_pages(self) -> int: + return self.sizes.n_win_pages - 1 + + @property + def swa_num_tokens(self) -> int: + # Allocatable window slots + 1: the generic capacity convention reserves slot 0 as a + # sentinel (cap == swa_num_tokens - 1); here the reserved unit is the tail dummy page, + # already excluded from the free list, so +1 re-encodes the same cap. + return (self.sizes.n_win_slots - self.P) + 1 + + def swa_available_size(self) -> int: + return int(self._win_alloc.available()) + + def alloc_swa(self, full_indices: torch.Tensor) -> None: + """Bind one window page per incoming FULL page. ``full_indices`` must be whole ascending + pages (the ``_page_to_token`` expansion); the in-page offsets are preserved + (``window_slot = wbase + pos % P``), which the state ring's page-block layout requires.""" + n = int(full_indices.numel()) + if n == 0: + return + P = self.P + assert n % P == 0, f"alloc_swa needs whole pages, got {n} slots" + fi = full_indices.to(device=self._device, dtype=torch.int64).view(-1, P) + fbases = fi[:, 0] + assert torch.equal(fi, fbases[:, None] + torch.arange(P, device=self._device)), "alloc_swa pages must be contiguous ascending" + wbases = self._win_alloc.alloc(fbases.numel()) # raises when exhausted (caller gated) + offsets = torch.arange(P, dtype=torch.int64, device=self._device) + self.full_to_window[(fbases[:, None] + offsets).flatten()] = (wbases[:, None] + offsets).flatten() + + def free_swa(self, full_indices: torch.Tensor) -> None: + """Return the window pages backing these FULL locs and unbind the mapping. Page-atomic: the + incoming locs must cover each touched page completely (guaranteed by the padded finish + tails / aligned frontiers / page-aligned tree values). Idempotent over already unbound + (slid / tombstoned) pages.""" + if full_indices.numel() == 0: + return + P = self.P + fi = full_indices.to(device=self._device, dtype=torch.int64) + fi = fi[fi >= 0] + if fi.numel() == 0: + return + fbases, counts = torch.unique(torch.div(fi, P, rounding_mode="floor") * P, return_counts=True) + assert bool((counts == P).all()), f"free_swa got partial pages (counts {counts[counts != P].tolist()[:4]})" + ws = self.full_to_window[fbases] + live = ws[ws >= 0] + offsets = torch.arange(P, dtype=torch.int64, device=self._device) + self.full_to_window[(fbases[:, None] + offsets).flatten()] = -1 + if live.numel(): + self._win_alloc.free(torch.div(live, P, rounding_mode="floor") * P) + + def translate_loc_from_full_to_swa(self, kv_indices: torch.Tensor) -> torch.Tensor: + return self.full_to_window[kv_indices.to(dtype=torch.int64)] + + # ----- rebuild / page table ----- + def rebuild(self, sizes) -> None: + """In-place resize to ``sizes`` (identity-preserving; free-before-alloc). The manager's + tree / page bookkeeping reset is the scheduler's generic cache_manager.rebuild; the engine + re-attaches the page table via attach_page_table afterwards.""" + assert self._paged_params is not None, "rebuild before _init_paged_state" + self.sizes = sizes + self._drop_buffers() + self.full_to_window = None # type: ignore[assignment] + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + self._alloc_buffers() + self._init_paged_state(*self._paged_params) + + def attach_page_table(self, page_table: torch.Tensor) -> None: + # the model reads full locs through full_loc_map; under the shared route that IS the page table + self.full_loc_map = page_table + + # ----- BaseKVCachePool glue: the window tier is the K == V latent cache ----- + def k_cache(self, index: int) -> torch.Tensor: + return self.window_pool[index] + + def v_cache(self, index: int) -> torch.Tensor: + return self.window_pool[index] + + def store_kv(self, k, v, out_loc, layer_id) -> None: + self.store_window(k, layer_id, out_loc) + + @abstractmethod + def store_window(self, kv: torch.Tensor, layer_id: int, window_slot: torch.Tensor) -> None: ... + + @property + def device(self) -> torch.device: + return self._device + + @property + def dtype(self) -> torch.dtype: + return self._dtype + + +__all__ = [ + "FreeListAllocator", + "CompressStateRing", + "WindowTierPagedPool", + "reserved_window_pages", + "window_state_loc", +] 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 1bd9e0b9cd..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) @@ -98,7 +116,10 @@ def insert(self, input_ids: torch.Tensor, kv_indices: torch.Tensor, node, prefix_len = self._walk(input_ids) if prefix_len != insert_len: new_node = RadixTreeNode(self.key_fn) - new_node.set_key_value(input_ids[prefix_len:], kv_indices[prefix_len:].clone()) + # A view would keep the request's full output-budget token buffer alive. + new_node.set_key_value( + input_ids[prefix_len:].clone(), kv_indices[prefix_len:].clone() + ) new_node.set_parent(node) self.full_evictable += new_node.length node = new_node @@ -156,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 @@ -177,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) @@ -230,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 @@ -252,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 41c3880e34..57a5734aa4 100644 --- a/python/freetoken/kvcache/hybrid_swa_pool.py +++ b/python/freetoken/kvcache/hybrid_swa_pool.py @@ -8,7 +8,7 @@ from freetoken.models.config import KVCacheGroupSpec from freetoken.utils import align_ceil, div_even -from .base import BaseKVCachePool +from .base import BaseKVCachePool, WindowPoolSpec @dataclass(frozen=True) @@ -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: @@ -295,6 +427,16 @@ def rebuild(self, num_full_pages: int, num_swa_tokens: int | None = None) -> Non # by CacheManager.rebuild is then consistent by construction (idle-only). self._init_swa_paged_state() + @classmethod + def window_spec(cls, config) -> WindowPoolSpec | None: + if config.cache_type != "swa_radix": + return None + return WindowPoolSpec(1, _swa_pool_floor(config)) + + @property + def window_pages(self) -> int: + return self.swa_num_tokens - 1 + @classmethod def kv_cost(cls, config, *, num_swa_pages: int | None = None) -> tuple[int, int, int, int]: """Full groups ride cache_per_page; the window pool is either a pinned absolute @@ -346,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/kvcache/radix_cache.py b/python/freetoken/kvcache/radix_cache.py index 5437286d12..ec0bf75d35 100644 --- a/python/freetoken/kvcache/radix_cache.py +++ b/python/freetoken/kvcache/radix_cache.py @@ -167,7 +167,10 @@ def insert_prefix(self, input_ids: torch.Tensor, indices: torch.Tensor) -> Inser node, prefix_len = self._tree_walk(input_ids) if prefix_len != insert_len: # NOTE: prefix_len < insert_len new_node = RadixTreeNode(self.key_fn) - new_node.set_key_value(input_ids[prefix_len:], indices[prefix_len:].clone()) + # A view would keep the request's full output-budget token buffer alive. + new_node.set_key_value( + input_ids[prefix_len:].clone(), indices[prefix_len:].clone() + ) new_node.set_parent(node) self.evictable_size += new_node.length node = new_node diff --git a/python/freetoken/kvcache/swa_radix_cache.py b/python/freetoken/kvcache/swa_radix_cache.py index 260439ad66..d2cc89c26f 100644 --- a/python/freetoken/kvcache/swa_radix_cache.py +++ b/python/freetoken/kvcache/swa_radix_cache.py @@ -215,7 +215,8 @@ def insert(self, input_ids: torch.Tensor, kv_indices: torch.Tensor, # then a live (non-tombstone) leaf for the in-window remainder (never a tombstone leaf -- # the clamp + the free_swa -page_size margin guarantee a live tail). if total < insert_len: - suffix_ids = input_ids[total:] + # Both suffix nodes can share this copy without retaining the request buffer. + suffix_ids = input_ids[total:].clone() suffix_kv = kv_indices[total:].clone() boundary = max(0, min(swa_evicted_seqlen, insert_len) - total) boundary = min(boundary, max(0, len(suffix_ids) - self.page_size)) diff --git a/python/freetoken/layers/__init__.py b/python/freetoken/layers/__init__.py index d868751a6b..a07484620e 100644 --- a/python/freetoken/layers/__init__.py +++ b/python/freetoken/layers/__init__.py @@ -16,7 +16,7 @@ LinearReplicated, LinearRowParallel, ) -from .moe import MoELayer, OffloadMoELayer, make_moe_layer +from .moe import MoELayer, OffloadMoELayer, iter_moe_layers, make_moe_layer from .norm import ( GatedRMSNorm, GemmaPlusOneRMSNorm, @@ -57,5 +57,6 @@ "LinearReplicated", "MoELayer", "OffloadMoELayer", + "iter_moe_layers", "make_moe_layer", ] 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/mhc.py b/python/freetoken/layers/mhc.py index 1cd6459d7b..dcf4915140 100644 --- a/python/freetoken/layers/mhc.py +++ b/python/freetoken/layers/mhc.py @@ -19,6 +19,16 @@ correctness baseline; the fused triton kernel (``kernel/triton/mhc.py``, dispatched by ``mhc_fused_post_pre`` below) replaces it on CUDA and is validated against this file. + +Two variants share the gate math: + +* **multi-pass** (GLM-5.3, DeepSeek-V4): ``pre`` is computed from the residual the + sublayer is about to read and applied immediately -- ``mhc_pre`` / ``mhc_fused_post_pre``. +* **single-pass** (DeepSeek-V4.1 "Single-Pass mHC", tech report sec. 2.4.1): every + sublayer consumes the ``pre`` gates the *previous* sublayer produced, so the input + mix no longer waits on this sublayer's own reduction -- ``mhc_pre_single_pass`` / + ``mhc_fused_post_pre_single_pass``. The returned ``pre_new`` is what the next + sublayer applies; the model seeds the chain with a one-hot pre on stream 0. """ from __future__ import annotations @@ -36,8 +46,8 @@ def hc_contract(x: torch.Tensor) -> torch.Tensor: return x.mean(dim=1) -def mhc_pre( - residual: torch.Tensor, # [T, n, hidden] bf16 +def mhc_gates( + residual: torch.Tensor, # [T, n, hidden] fn: torch.Tensor, # [2n + n^2, n*hidden] fp32 hc_scale: torch.Tensor, # [3] fp32 hc_base: torch.Tensor, # [2n + n^2] fp32 @@ -46,8 +56,8 @@ def mhc_pre( post_mult: float, sinkhorn_repeat: int, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """Returns (post_mix [T, n, 1] fp32, comb_mix [T, n, n] fp32, - layer_input [T, hidden] bf16).""" + """The coefficient predictor H(X): returns (pre_mix [T, n] fp32, post_mix [T, n, 1] + fp32, comb_mix [T, n, n] fp32) from one fp32 GEMM over the flattened streams.""" n, hidden = residual.shape[-2], residual.shape[-1] t = residual.shape[0] @@ -68,11 +78,46 @@ def mhc_pre( for _ in range(sinkhorn_repeat - 1): comb = comb / (comb.sum(dim=-1, keepdim=True) + hc_eps) comb = comb / (comb.sum(dim=-2, keepdim=True) + hc_eps) + return pre_mix, post_mix.view(t, n, 1), comb - layer_input = ( - (pre_mix.unsqueeze(-1) * residual.to(torch.float32)).sum(dim=1).to(residual.dtype) - ) - return post_mix.view(t, n, 1), comb, layer_input + +def mhc_mix_input(residual: torch.Tensor, pre_mix: torch.Tensor) -> torch.Tensor: + """The sublayer input ``sum_i pre_i * res_i``: [T, n, hidden] x [T, n] -> [T, hidden].""" + return (pre_mix.unsqueeze(-1) * residual.to(torch.float32)).sum(dim=1).to(residual.dtype) + + +def mhc_pre( + residual: torch.Tensor, # [T, n, hidden] bf16 + fn: torch.Tensor, # [2n + n^2, n*hidden] fp32 + hc_scale: torch.Tensor, # [3] fp32 + hc_base: torch.Tensor, # [2n + n^2] fp32 + rms_eps: float, + hc_eps: float, + post_mult: float, + sinkhorn_repeat: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Multi-pass: gates from ``residual``, pre applied to it right away. Returns + (post_mix [T, n, 1] fp32, comb_mix [T, n, n] fp32, layer_input [T, hidden] bf16).""" + pre_mix, post_mix, comb = mhc_gates(residual, fn, hc_scale, hc_base, rms_eps, hc_eps, post_mult, sinkhorn_repeat) + return post_mix, comb, mhc_mix_input(residual, pre_mix) + + +def mhc_pre_single_pass( + residual: torch.Tensor, # [T, n, hidden] + pre_mix_in: torch.Tensor, # [T, n] fp32 -- the previous sublayer's pre gates + fn: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + rms_eps: float, + hc_eps: float, + post_mult: float, + sinkhorn_repeat: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Single-pass: gates from ``residual`` for THIS sublayer's post/comb and the NEXT + sublayer's pre; the input is mixed with ``pre_mix_in``. Returns (pre_new [T, n] fp32, + post_mix [T, n, 1] fp32, comb_mix [T, n, n] fp32, layer_input [T, hidden]).""" + pre_new, post_mix, comb = mhc_gates(residual, fn, hc_scale, hc_base, rms_eps, hc_eps, post_mult, sinkhorn_repeat) + return pre_new, post_mix, comb, mhc_mix_input(residual, pre_mix_in.to(torch.float32)) def mhc_post( @@ -140,11 +185,69 @@ def mhc_fused_post_pre( ) +def mhc_fused_post_pre_single_pass_torch( + x: torch.Tensor, + residual: torch.Tensor, + post_mix: torch.Tensor | None, + comb_mix: torch.Tensor | None, + pre_mix_in: torch.Tensor, + fn: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + rms_eps: float, + hc_eps: float, + post_mult: float, + sinkhorn_repeat: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Decomposed single-pass reference: the previous sublayer's hc_post (skipped when + ``post_mix is None`` -- the model's first sublayer), then this sublayer's gates with + the input mixed by ``pre_mix_in``. Returns (residual_new, post_new [T, n, 1], + comb_new [T, n, n], pre_new [T, n], layer_input).""" + residual_new = residual if post_mix is None else mhc_post(x, residual, post_mix, comb_mix) + pre_new, post_new, comb_new, layer_input = mhc_pre_single_pass( + residual_new, pre_mix_in, fn, hc_scale, hc_base, rms_eps, hc_eps, post_mult, sinkhorn_repeat + ) + return residual_new, post_new, comb_new, pre_new, layer_input + + +def mhc_fused_post_pre_single_pass( + x: torch.Tensor, + residual: torch.Tensor, + post_mix: torch.Tensor | None, + comb_mix: torch.Tensor | None, + pre_mix_in: torch.Tensor, + fn: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + rms_eps: float, + hc_eps: float, + post_mult: float, + sinkhorn_repeat: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Single-pass sublayer boundary (DeepSeek-V4.1): fused triton on CUDA, torch elsewhere.""" + if residual.is_cuda: + from freetoken.kernel.triton.mhc import mhc_fused_post_pre_single_pass_triton + + return mhc_fused_post_pre_single_pass_triton( + x, residual, post_mix, comb_mix, pre_mix_in, fn, hc_scale, hc_base, + rms_eps, hc_eps, post_mult, sinkhorn_repeat, + ) + return mhc_fused_post_pre_single_pass_torch( + x, residual, post_mix, comb_mix, pre_mix_in, fn, hc_scale, hc_base, + rms_eps, hc_eps, post_mult, sinkhorn_repeat, + ) + + __all__ = [ "hc_expand", "hc_contract", + "mhc_gates", + "mhc_mix_input", "mhc_pre", + "mhc_pre_single_pass", "mhc_post", "mhc_fused_post_pre", "mhc_fused_post_pre_torch", + "mhc_fused_post_pre_single_pass", + "mhc_fused_post_pre_single_pass_torch", ] diff --git a/python/freetoken/layers/moe.py b/python/freetoken/layers/moe.py index c209454996..71858b54d0 100644 --- a/python/freetoken/layers/moe.py +++ b/python/freetoken/layers/moe.py @@ -1,5 +1,5 @@ import os -from typing import TYPE_CHECKING, Tuple +from typing import TYPE_CHECKING, Iterator, Tuple import torch from freetoken.core import get_global_ctx @@ -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,13 @@ 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) def forward( self, @@ -235,7 +248,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, @@ -263,6 +288,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 @@ -282,9 +308,25 @@ 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, @@ -294,6 +336,29 @@ def _decode_routed( 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,18 +373,21 @@ 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. - Each route is computed exactly once -- the GPU weights are zeroed for CPU-assigned - routes and the CPU ids are -1 for GPU-assigned routes (the C++ kernel skips id<0). + 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 + GPU-assigned routes (the C++ kernel skips id<0). """ executor = cache.cpu_executor assert executor is not None, "CPU MoE executor was not initialized" - raw = topk_ids.clone() # raw expert ids for the CPU partial - cache.ensure_experts_hybrid(self.layer_id, topk_ids) # -> slot (hit/fetched) or -1 + 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) - on_gpu = topk_ids >= 0 - - cpu_ids = torch.where(on_gpu, raw.new_full((), -1), raw).contiguous() 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 @@ -329,8 +397,10 @@ def _decode_hybrid( ) cache.copy_missing() - gpu_slots = topk_ids.clamp_min(0) # -1 -> slot 0 (zero-weighted below) - gpu_w = torch.where(on_gpu, topk_weights, topk_weights.new_zeros(())).contiguous() + 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, @@ -356,6 +426,26 @@ def _prefill_routed( pass through unmapped.""" cache = self.offload_cache assert cache is not None + if ( + 0 < hidden_states.shape[0] <= _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). 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, + self._ensure_unique(cache, topk_ids), + views=cache.bank_views(), + n=None, + alphas=cache.alphas_for_slots(self.layer_id), + is_prefill=False, + ) if cache.prefill_overlap: views = self._wait_prefill_overlap(cache) out = self._expert_gemm( @@ -383,6 +473,27 @@ def _prefill_routed( is_prefill=True, ) + 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) # 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 next layer's full-layer H2D copy, then return this layer's bank views (in @@ -435,6 +546,19 @@ def _expert_gemm( raise AssertionError(f"offload experts without a quant method only serve q4_0 banks, got {fmt!r}") +def iter_moe_layers(model) -> Iterator[MoELayer]: + if isinstance(model, MoELayer): + yield model + if not isinstance(model, BaseOP): + return + for value in model.__dict__.values(): + if isinstance(value, BaseOP): + yield from iter_moe_layers(value) + elif isinstance(value, (list, tuple)): + for item in value: + yield from iter_moe_layers(item) + + def make_moe_layer( config: "ModelConfig", *, @@ -483,10 +607,10 @@ def make_moe_layer( has_bias=has_bias, quant_config=quant_config, prefix=prefix, + layer_id=layer_id, ) if offload: assert layer_id is not None, "offload MoE backends need the layer_id" - kwargs["layer_id"] = layer_id kwargs["strategy"] = config.moe_strategy kwargs["decode_target"] = config.decode_target return layer_cls(**kwargs) diff --git a/python/freetoken/layers/quantization/configs/base.py b/python/freetoken/layers/quantization/configs/base.py index e84dcbcbba..a64136103b 100644 --- a/python/freetoken/layers/quantization/configs/base.py +++ b/python/freetoken/layers/quantization/configs/base.py @@ -89,9 +89,13 @@ def get_quant_method(self, layer: Any, prefix: str): layer_kind = layer.quant_layer_kind kind = scheme.kind if scheme else QuantKind.NONE cls = method_class(kind, layer_kind) - cfg = LAYER_CONFIGS[layer_kind].from_layer(layer, scheme) + cfg = self.layer_config(layer, layer_kind, scheme) return cls(cfg, get_quant_backend().select(layer_kind, kind)) + def layer_config(self, layer: Any, layer_kind: LayerKind, scheme: QuantScheme | None) -> LinearConfig | MoEConfig: + """The layer's Linear / MoE config; a dialect overrides it to add what only its checkpoint knows.""" + return LAYER_CONFIGS[layer_kind].from_layer(layer, scheme) + @staticmethod def from_hf( hf_config: Any, diff --git a/python/freetoken/layers/quantization/configs/fp8.py b/python/freetoken/layers/quantization/configs/fp8.py index c23cbfdee8..3e74f0c4d0 100644 --- a/python/freetoken/layers/quantization/configs/fp8.py +++ b/python/freetoken/layers/quantization/configs/fp8.py @@ -1,20 +1,25 @@ from __future__ import annotations +import dataclasses from typing import Any, ClassVar +from ..linear import LinearConfig +from ..moe import MoEConfig from ..names import is_routed_expert, name_set, substr_set -from ..registry import register_dialect +from ..registry import LayerKind, register_dialect from ..scheme import QuantKind, QuantScheme -from ..scheme import FP8_BLOCK, fp8_block_scheme, fp8_tensor_scheme, mxfp4_scheme +from ..scheme import FP8_BLOCK_SIZES, fp8_block_scheme, fp8_tensor_scheme, mxfp4_scheme from .base import QuantConfig, Stored, cfg_get @register_dialect class Fp8BlockConfig(QuantConfig): - """HF ``quant_method: fp8`` (DeepSeek-V3 style 128x128 block scales) plus the DeepSeek-V4 e8m0 / fp4-expert variant.""" + """HF ``quant_method: fp8`` (DeepSeek-V3 style 128x128 block scales) plus the DeepSeek-V4 e8m0 / fp4-expert + variant; DeepSeek-V4.1 keeps that dialect with 32x32 blocks (``weight_block_size: [32, 32]``).""" dialect = "fp8" + # the 128-block schemes every dialect table names; ``block_scheme`` carries the checkpoint's actual block edge SCHEMES: ClassVar[dict[str, QuantScheme]] = { "BLOCK": fp8_block_scheme("float"), "BLOCK_E8M0": fp8_block_scheme("e8m0"), @@ -32,15 +37,23 @@ class Fp8BlockConfig(QuantConfig): def __init__(self, q: dict[str, Any], hf_config: Any = None, *, name_map=None, unquantized=()): super().__init__(name_map, unquantized) block = tuple(int(x) for x in (q.get("weight_block_size") or ())) - if q.get("weight_per_tensor") or block != (FP8_BLOCK, FP8_BLOCK): - raise NotImplementedError(f"fp8 checkpoint with weight_block_size={block} per_tensor={q.get('weight_per_tensor')} is not supported; only 128x128 blocks are") + square = len(block) == 2 and block[0] == block[1] and block[0] in FP8_BLOCK_SIZES + if q.get("weight_per_tensor") or not square: + raise NotImplementedError( + f"fp8 checkpoint with weight_block_size={block} per_tensor={q.get('weight_per_tensor')} is not supported; " + f"only square blocks of {FP8_BLOCK_SIZES} are" + ) + self.block = block[0] # transformers skips lm_head when the checkpoint gives no list not_convert = tuple(q.get("modules_to_not_convert") or ("lm_head",)) self.not_convert = name_set(not_convert) self.not_convert_substr = substr_set(not_convert) self.convert_tables = name_set(tuple(q.get("modules_to_convert") or ())) self.e8m0 = str(q.get("scale_fmt") or "").lower() == "ue8m0" - self.expert_fp4 = str(cfg_get(hf_config, "expert_dtype") or "").lower() == "fp4" + # ``expert_dtype`` sits at the config top level (V4) or inside quantization_config (V4.1) + expert_dtype = q.get("expert_dtype") or cfg_get(hf_config, "expert_dtype") + self.expert_fp4 = str(expert_dtype or "").lower() == "fp4" + self.block_scheme = fp8_block_scheme("e8m0" if self.e8m0 else "float", self.block) def storage(self, scheme: QuantScheme) -> dict[str, Stored]: names = super().storage(scheme) @@ -55,4 +68,11 @@ def scheme_for_name(self, name: str) -> QuantScheme | None: return None if self.expert_fp4 and is_routed_expert(name): return self.SCHEMES["EXPERT_MXFP4"] - return self.SCHEMES["BLOCK_E8M0" if self.e8m0 else "BLOCK"] + return self.block_scheme + + def layer_config(self, layer: Any, layer_kind: LayerKind, scheme: QuantScheme | None) -> LinearConfig | MoEConfig: + cfg = super().layer_config(layer, layer_kind, scheme) + if layer_kind is LayerKind.MOE and scheme is not None and scheme.kind is QuantKind.MXFP4: + # the fp4 experts quantize their activations at the same block as the dense linears + cfg = dataclasses.replace(cfg, act_block=self.block) + return cfg diff --git a/python/freetoken/layers/quantization/linear/fp8_block.py b/python/freetoken/layers/quantization/linear/fp8_block.py index 3b78e9776d..3eafdfdf5b 100644 --- a/python/freetoken/layers/quantization/linear/fp8_block.py +++ b/python/freetoken/layers/quantization/linear/fp8_block.py @@ -1,4 +1,4 @@ -"""fp8 e4m3 weight with 128x128 block scales (DeepSeek-V3 style).""" +"""fp8 e4m3 weight with square block scales: 128x128 (DeepSeek-V3 style) or 32x32 (DeepSeek-V4.1).""" from __future__ import annotations @@ -7,7 +7,7 @@ import torch from ..registry import LayerKind, register_method -from ..scheme import FP8_BLOCK as BLOCK, QuantKind +from ..scheme import FP8_BLOCK, QuantKind, fp8_block_size from .base import LinearConfig, LinearKernel, LinearMethod FP8 = torch.float8_e4m3fn @@ -18,6 +18,11 @@ def _e8m0(cfg: LinearConfig) -> bool: return cfg.scheme is not None and cfg.scheme.weight.scale == "e8m0" +def layer_block_size(layer: Any) -> int: + """The block edge a finalized layer was declared with: ``K / (scale columns)``, exact by construction.""" + return layer.weight.shape[1] // layer.weight_scale_inv.shape[1] + + class Dsv4Fp8BlockLinearKernel(LinearKernel): """DeepSeek-V4's reference path: activations quantized to fp8 with power-of-two block scales, e8m0 weight scales read as codes.""" @@ -29,7 +34,7 @@ def unusable_reason(self, cfg: LinearConfig) -> str | None: def apply(self, layer: Any, x: torch.Tensor) -> torch.Tensor: from freetoken.kernel.triton.dsv4.fp8_linear import block_fp8_linear - return block_fp8_linear(x, layer.weight, layer.weight_scale_inv, layer.bias) + return block_fp8_linear(x, layer.weight, layer.weight_scale_inv, layer.bias, block=layer_block_size(layer)) class TritonFp8BlockLinearKernel(LinearKernel): @@ -38,7 +43,11 @@ class TritonFp8BlockLinearKernel(LinearKernel): name = "triton" def unusable_reason(self, cfg: LinearConfig) -> str | None: - return "reads float block scales; e8m0 codes go to the dsv4 kernel" if _e8m0(cfg) else None + if _e8m0(cfg): + return "reads float block scales; e8m0 codes go to the dsv4 kernel" + if fp8_block_size(cfg.scheme) != FP8_BLOCK: + return f"float-scale kernel serves {FP8_BLOCK}x{FP8_BLOCK} blocks only" + return None def apply(self, layer: Any, x: torch.Tensor) -> torch.Tensor: from freetoken.kernel.triton.fp8_block_linear import block_fp8_linear @@ -52,9 +61,10 @@ class Fp8BlockLinearMethod(LinearMethod): def create_weights(self, layer: Any) -> None: g = self.cfg - if g.in_features % BLOCK or any(o % BLOCK for o in g.output_sizes): - raise ValueError(f"block-fp8 needs in/out sizes divisible by {BLOCK}, got K={g.in_features} N={g.output_sizes}") + block = fp8_block_size(g.scheme) + if g.in_features % block or any(o % block for o in g.output_sizes): + raise ValueError(f"block-fp8 needs in/out sizes divisible by {block}, got K={g.in_features} N={g.output_sizes}") layer.weight = torch.empty(g.out_features, g.in_features, dtype=FP8) # e8m0 codes stay codes for the dsv4 kernel; float scales are bf16 as the readers push them today scale_dtype = E8M0 if _e8m0(g) else torch.bfloat16 - layer.weight_scale_inv = torch.empty(g.out_features // BLOCK, g.in_features // BLOCK, dtype=scale_dtype) + layer.weight_scale_inv = torch.empty(g.out_features // block, g.in_features // block, dtype=scale_dtype) diff --git a/python/freetoken/layers/quantization/moe/base.py b/python/freetoken/layers/quantization/moe/base.py index 89f13ed09d..de1ce004fe 100644 --- a/python/freetoken/layers/quantization/moe/base.py +++ b/python/freetoken/layers/quantization/moe/base.py @@ -28,10 +28,13 @@ class MoEConfig: beta: float = 0.0 limit: float | None = None interleaved: bool = False + # dynamic fp8 activation quant block of a W4A8 expert kernel, set by the quant dialect; None = no activation quant + act_block: int | None = None has_bias: bool = False apply_router_weight_on_input: bool = False strategy: str = "resident" decode_target: str = "gpu" + dtype: torch.dtype = torch.bfloat16 @classmethod def from_layer(cls, layer: Any, scheme: QuantScheme | None) -> "MoEConfig": @@ -52,6 +55,7 @@ def from_layer(cls, layer: Any, scheme: QuantScheme | None) -> "MoEConfig": apply_router_weight_on_input=bool(layer.apply_router_weight_on_input), strategy=layer.strategy, decode_target=layer.decode_target, + dtype=torch.get_default_dtype(), ) @property @@ -136,6 +140,10 @@ class MoEKernel(ABC): name: ClassVar[str] cpu_format: ClassVar[str | None] = None max_slots: ClassVar[int | None] = None + # Whether ``apply`` treats a slot of -1 as an inactive route (zero output, no read of any slot). + # Hybrid decode hands the GPU a route split with such holes; kernels without this get the holes + # pointed at slot 0 with weight 0 (slot storage is zero-filled at allocation, so that is exact). + supports_inactive_slots: ClassVar[bool] = False def unusable_reason(self, cfg: MoEConfig) -> str | None: return None @@ -147,9 +155,7 @@ def slot_limit(self, cfg: MoEConfig) -> int | None: """Most GPU cache slots the kernel can address for ``cfg``; None for no limit.""" return self.max_slots - def _common_reject(self, cfg: MoEConfig, *, resident_ok: bool, tp_ok: bool, cpu_ok: bool, plain_silu_only: bool) -> str | None: - if not resident_ok and cfg.strategy == "resident": - return "not served resident; use --moe-strategy offload or cpu" + def _common_reject(self, cfg: MoEConfig, *, tp_ok: bool, cpu_ok: bool, plain_silu_only: bool) -> str | None: if not tp_ok and cfg.tp_size > 1: return "TP > 1 is not supported for this expert format" if not cpu_ok and cfg.decode_target != "gpu": @@ -190,12 +196,15 @@ def apply(self, x: torch.Tensor, topk_weights: torch.Tensor, topk_ids: torch.Ten def cpu_format(self) -> str | None: return self.kernel.cpu_format - @abstractmethod - def create_weights(self, layer: Any) -> None: ... + def create_weights(self, layer: Any) -> None: + # resident experts stay out of the state dict; attach_resident_banks fills these after the dense weights load + layer._expert_banks = None + layer._expert_alphas = None def finalize(self, layer: Any) -> None: pass - @abstractmethod def resident_view(self, layer: Any) -> ExpertView: """Layout roles -> the resident layer's tensors, for apply().""" + assert layer._expert_banks is not None, "resident expert banks were never loaded" + return ExpertView(layer._expert_banks, n=self.cfg.num_experts, alphas=layer._expert_alphas) diff --git a/python/freetoken/layers/quantization/moe/fp8_block.py b/python/freetoken/layers/quantization/moe/fp8_block.py index 6c0167b71b..7552978d3f 100644 --- a/python/freetoken/layers/quantization/moe/fp8_block.py +++ b/python/freetoken/layers/quantization/moe/fp8_block.py @@ -5,7 +5,7 @@ import torch from ..registry import LayerKind, register_method -from ..scheme import FP8_BLOCK as BLOCK, QuantKind +from ..scheme import FP8_BLOCK as BLOCK, QuantKind, fp8_block_size from .base import BankSpec, ExpertView, fused_piece, gated_epilogue_reason, limit_or_inf, MoEConfig, MoEKernel, MoEMethod FP8 = torch.float8_e4m3fn @@ -21,9 +21,11 @@ class TritonFp8BlockMoEKernel(MoEKernel): name = "triton" def unusable_reason(self, cfg: MoEConfig) -> str | None: - reason = self._common_reject(cfg, resident_ok=True, tp_ok=False, cpu_ok=False, plain_silu_only=False) + reason = self._common_reject(cfg, tp_ok=False, cpu_ok=False, plain_silu_only=False) if reason: return reason + if cfg.scheme is not None and fp8_block_size(cfg.scheme) != BLOCK: + return f"fp8 block MoE kernel serves {BLOCK}x{BLOCK} expert blocks only" reason = gated_epilogue_reason(cfg) if reason: return f"fp8 block MoE kernel: {reason}" @@ -64,17 +66,3 @@ def apply(self, layer, x, topk_weights, topk_ids, view: ExpertView, *, is_prefil @register_method(QuantKind.FP8_BLOCK, LayerKind.MOE) class Fp8BlockMoEMethod(MoEMethod): candidates = (TritonFp8BlockMoEKernel,) - - def create_weights(self, layer) -> None: - g = self.cfg - e, i, h, b = g.num_experts, g.intermediate, g.hidden, BLOCK - layer.gate_up_proj = torch.empty(e, 2 * i, h, dtype=FP8) - layer.gate_up_scale_inv = torch.empty(e, 2 * i // b, h // b, dtype=torch.bfloat16) - layer.down_proj = torch.empty(e, h, i, dtype=FP8) - layer.down_scale_inv = torch.empty(e, h // b, i // b, dtype=torch.bfloat16) - - def resident_view(self, layer) -> ExpertView: - return ExpertView({ - "gate_up": layer.gate_up_proj, "gate_up_scale": layer.gate_up_scale_inv, - "down": layer.down_proj, "down_scale": layer.down_scale_inv, - }) diff --git a/python/freetoken/layers/quantization/moe/mxfp4.py b/python/freetoken/layers/quantization/moe/mxfp4.py index aa020ee33a..1af6d9f43e 100644 --- a/python/freetoken/layers/quantization/moe/mxfp4.py +++ b/python/freetoken/layers/quantization/moe/mxfp4.py @@ -12,17 +12,21 @@ class TritonMxfp4MoEKernel(MoEKernel): - """Standard OCP MXFP4 experts (DeepSeek-V4 ds_fp4): e2m1 pairs + e8m0 scales, no bias.""" + """Standard OCP MXFP4 experts (DeepSeek-V4 / V4.1 ds_fp4): e2m1 pairs + e8m0 scales, no bias; + W4A8 with the fp8 activation round-trips at the checkpoint's block.""" name = "triton" cpu_format = "ds_fp4" + supports_inactive_slots = True def unusable_reason(self, cfg: MoEConfig) -> str | None: if cfg.interleaved: return "standard MXFP4 kernel reads the concatenated gate|up row order" if (cfg.alpha, cfg.beta) != (1.0, 0.0): return "standard MXFP4 kernel has no alpha / beta in its swiglu" - return self._common_reject(cfg, resident_ok=False, tp_ok=False, cpu_ok=True, plain_silu_only=False) + if cfg.act_block is None: + return "W4A8 ds_fp4 experts need the fp8 activation quant block" + return self._common_reject(cfg, tp_ok=False, cpu_ok=True, plain_silu_only=False) def layout(self, cfg: MoEConfig) -> dict[str, BankSpec]: i, h = cfg.intermediate, cfg.hidden @@ -44,13 +48,14 @@ def apply(self, layer, x, topk_weights, topk_ids, view: ExpertView, *, is_prefil t = view.tensors banks = (t["gate_up"], t["gate_up_scale"], t["down"], t["down_scale"]) limit = limit_or_inf(layer) + act_block = layer.quant_method.cfg.act_block if is_prefill and view.n is not None: from freetoken.moe.fused_ds_fp4 import routed_experts_fp4_prefill - return routed_experts_fp4_prefill(x, topk_ids, topk_weights, *banks, limit, view.n) + return routed_experts_fp4_prefill(x, topk_ids, topk_weights, *banks, limit, view.n, act_block=act_block) from freetoken.moe.fused_ds_fp4 import routed_experts_fp4 - return routed_experts_fp4(x, topk_ids, topk_weights, *banks, limit) + return routed_experts_fp4(x, topk_ids, topk_weights, *banks, limit, act_block=act_block) class TritonGptossMxfp4MoEKernel(MoEKernel): @@ -106,37 +111,3 @@ def apply(self, layer, x, topk_weights, topk_ids, view: ExpertView, *, is_prefil @register_method(QuantKind.MXFP4, LayerKind.MOE) class Mxfp4MoEMethod(MoEMethod): candidates = (TritonMxfp4MoEKernel, TritonGptossMxfp4MoEKernel) - - def create_weights(self, layer) -> None: - if not self.cfg.has_bias: - raise NotImplementedError("standard MXFP4 experts are served from the offload cache, not resident") - g = self.cfg - e, i, h = g.num_experts, g.local_intermediate, g.hidden - if h % GROUP: - raise ValueError(f"MXFP4 hidden size must be divisible by {GROUP}") - layer.gate_up_proj_blocks = torch.empty(e, 2 * i, h // GROUP, 16, dtype=torch.uint8) - layer.gate_up_proj_scales = torch.empty(e, 2 * i, h // GROUP, dtype=torch.uint8) - layer.gate_up_proj_bias = torch.empty(e, 2 * i, dtype=torch.bfloat16) - layer.down_proj_blocks = torch.empty(e, h, i // GROUP, 16, dtype=torch.uint8) - layer.down_proj_scales = torch.empty(e, h, i // GROUP, dtype=torch.uint8) - layer.down_proj_bias = torch.empty(e, h, dtype=torch.bfloat16) - - def finalize(self, layer) -> None: - if getattr(layer, "gate_up_proj_blocks", None) is None: - return - from freetoken.moe.fused_mxfp4 import _transpose_mxfp4_for_decode - - # one transposed copy serves prefill and decode; the HF blocks are freed so 120B fits - layer._gu_blocks_t, layer._gu_scales_t = _transpose_mxfp4_for_decode(layer.gate_up_proj_blocks, layer.gate_up_proj_scales) - layer._dn_blocks_t, layer._dn_scales_t = _transpose_mxfp4_for_decode(layer.down_proj_blocks, layer.down_proj_scales) - layer.gate_up_proj_blocks = None - layer.gate_up_proj_scales = None - layer.down_proj_blocks = None - layer.down_proj_scales = None - torch.cuda.empty_cache() - - def resident_view(self, layer) -> ExpertView: - return ExpertView({ - "gate_up": layer._gu_blocks_t, "gate_up_scale": layer._gu_scales_t, "gate_up_bias": layer.gate_up_proj_bias, - "down": layer._dn_blocks_t, "down_scale": layer._dn_scales_t, "down_bias": layer.down_proj_bias, - }) diff --git a/python/freetoken/layers/quantization/moe/nvfp4.py b/python/freetoken/layers/quantization/moe/nvfp4.py index bb341cdfb7..102dbff96d 100644 --- a/python/freetoken/layers/quantization/moe/nvfp4.py +++ b/python/freetoken/layers/quantization/moe/nvfp4.py @@ -1,4 +1,4 @@ -"""NVFP4 experts: served from the offload cache through the Triton inline-dequant kernels, vLLM's Marlin or flashinfer's b12x. +"""NVFP4 experts: served resident or from the offload cache through the Triton inline-dequant kernels, vLLM's Marlin or flashinfer's b12x. FreeToken stores MoE experts as ModelOpt NVFP4 (packed e2m1 codes + fp8-e4m3 per-16 block scales + per-tensor global scale) in pinned host banks and gathers the routed @@ -41,14 +41,20 @@ class TritonNvfp4MoEKernel(MoEKernel): cpu_format = "nvfp4" def unusable_reason(self, cfg: MoEConfig) -> str | None: - reason = self._common_reject(cfg, resident_ok=False, 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)) @@ -230,7 +236,7 @@ class MarlinNvfp4MoEKernel(MoEKernel): def unusable_reason(self, cfg: MoEConfig) -> str | None: if not backend.is_vllm_installed(): return "vLLM is not installed" - reason = self._common_reject(cfg, resident_ok=False, tp_ok=False, cpu_ok=False, plain_silu_only=True) + reason = self._common_reject(cfg, tp_ok=False, cpu_ok=False, plain_silu_only=True) if reason: return reason if not _marlin_symbols_ok(): @@ -241,7 +247,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), @@ -252,7 +258,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) @@ -496,14 +502,14 @@ def unusable_reason(self, cfg: MoEConfig) -> str | None: return f"b12x requires sm_120+, got sm_{cc[0]}{cc[1]}" if not backend.is_flashinfer_installed(): return "flashinfer is not installed" - reason = self._common_reject(cfg, resident_ok=False, tp_ok=False, cpu_ok=False, plain_silu_only=True) + reason = self._common_reject(cfg, tp_ok=False, cpu_ok=False, plain_silu_only=True) if reason: return reason return _b12x_unusable_reason(cc) 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 @@ -512,7 +518,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), @@ -525,7 +531,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() @@ -567,9 +573,3 @@ def apply(self, layer, x, topk_weights, topk_ids, view: ExpertView, *, is_prefil @register_method(QuantKind.NVFP4, LayerKind.MOE) class Nvfp4MoEMethod(MoEMethod): candidates = (TritonNvfp4MoEKernel, MarlinNvfp4MoEKernel, B12xNvfp4MoEKernel) - - def create_weights(self, layer) -> None: - raise NotImplementedError("NVFP4 experts are served from the offload cache, not resident") - - def resident_view(self, layer) -> ExpertView: - raise NotImplementedError("NVFP4 experts are not resident") diff --git a/python/freetoken/layers/quantization/moe/unquantized.py b/python/freetoken/layers/quantization/moe/unquantized.py index ef435686c5..ded7d8135c 100644 --- a/python/freetoken/layers/quantization/moe/unquantized.py +++ b/python/freetoken/layers/quantization/moe/unquantized.py @@ -19,7 +19,9 @@ def unusable_reason(self, cfg: MoEConfig) -> str | None: def layout(self, cfg: MoEConfig) -> dict[str, BankSpec]: i = cfg.local_intermediate - return {"gate_up": BankSpec((2 * i, cfg.hidden), torch.bfloat16), "down": BankSpec((cfg.hidden, i), torch.bfloat16)} + # resident experts compute in the model dtype; offload banks stay bf16, the FTW and CPU executor format + dt = cfg.dtype if cfg.strategy == "resident" else torch.bfloat16 + return {"gate_up": BankSpec((2 * i, cfg.hidden), dt), "down": BankSpec((cfg.hidden, i), dt)} def pack(self, pieces, cfg: MoEConfig, out): out["gate_up"].copy_(fused_piece(pieces, "gate_up")) @@ -37,12 +39,3 @@ def apply(self, layer, x, topk_weights, topk_ids, view: ExpertView, *, is_prefil @register_method(QuantKind.NONE, LayerKind.MOE) class UnquantizedMoEMethod(MoEMethod): candidates = (FusedMoEKernel,) - - def create_weights(self, layer) -> None: - g = self.cfg - i = g.local_intermediate - layer.gate_up_proj = torch.empty(g.num_experts, 2 * i, g.hidden) - layer.down_proj = torch.empty(g.num_experts, g.hidden, i) - - def resident_view(self, layer) -> ExpertView: - return ExpertView({"gate_up": layer.gate_up_proj, "down": layer.down_proj}) diff --git a/python/freetoken/layers/quantization/scheme.py b/python/freetoken/layers/quantization/scheme.py index 668eb6fb37..89f968afc8 100644 --- a/python/freetoken/layers/quantization/scheme.py +++ b/python/freetoken/layers/quantization/scheme.py @@ -25,7 +25,8 @@ def __str__(self) -> str: return self.value -FP8_BLOCK = 128 +FP8_BLOCK = 128 # DeepSeek-V3 style square block; V4.1 exports 32x32 +FP8_BLOCK_SIZES = (32, 128) NVFP4_GROUP = 16 MX_GROUP = 32 @@ -80,8 +81,17 @@ def fp8_tensor_scheme(scale: str, *, per_row: bool = False, input_scale: bool = return QuantScheme(QuantKind.FP8_TENSOR, WeightDesc("e4m3", (1, -1) if per_row else (-1, -1), scale), roles) -def fp8_block_scheme(scale: str) -> QuantScheme: - return QuantScheme(QuantKind.FP8_BLOCK, WeightDesc("e4m3", (FP8_BLOCK, FP8_BLOCK), scale), {"weight", "weight_scale_inv"}) +def fp8_block_scheme(scale: str, block: int = FP8_BLOCK) -> QuantScheme: + """Square ``block x block`` e4m3 weight blocks, one scale each; ``block`` is the checkpoint's ``weight_block_size``.""" + if block not in FP8_BLOCK_SIZES: + raise NotImplementedError(f"fp8 block size {block} is not supported; one of {FP8_BLOCK_SIZES}") + return QuantScheme(QuantKind.FP8_BLOCK, WeightDesc("e4m3", (block, block), scale), {"weight", "weight_scale_inv"}) + + +def fp8_block_size(scheme: QuantScheme) -> int: + """The square block edge of an FP8_BLOCK scheme (128 for DeepSeek-V3 style exports, 32 for DeepSeek-V4.1).""" + assert scheme.kind is QuantKind.FP8_BLOCK, scheme + return scheme.weight.group[0] def mxfp8_scheme() -> QuantScheme: 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 9442ee059e..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 @@ -73,6 +87,7 @@ class TokenizeMsg(BaseTokenizerMsg): chat_template_kwargs: Dict[str, Any] | None = None tools: List[Dict[str, Any]] | None = None images: List[bytes] | None = None + inline_system_policy: str | None = None @dataclass diff --git a/python/freetoken/mm/processors/deepseek_v41.py b/python/freetoken/mm/processors/deepseek_v41.py new file mode 100644 index 0000000000..e40c961569 --- /dev/null +++ b/python/freetoken/mm/processors/deepseek_v41.py @@ -0,0 +1,88 @@ +"""DeepSeek-V4.1 image resize/patch layout and learned image-span delimiters.""" + +from __future__ import annotations + +import math +import struct + +import numpy as np +import torch +from PIL import ImageOps + +from freetoken.message import MMItem +from freetoken.mm import mm_pad_value +from freetoken.mm.processor import MMProcessor, PromptReplacement, content_hash + + +def image_grid(width, height, patch, ratio, min_pixels, max_tokens, max_wh_ratio=None): + if max_wh_ratio is not None: + width = min(width, height * max_wh_ratio) + if width * height < min_pixels: + scale = (min_pixels / (width * height)) ** 0.5 + width, height = int(width * scale), int(height * scale) + w, h = math.ceil(width / patch) * patch, math.ceil(height / patch) * patch + if math.ceil(h / (patch * ratio)) * (math.ceil(w / (patch * ratio)) + 1) + 2 > max_tokens: + aspect, cell = height / width, patch * ratio + wf = math.sqrt((max_tokens - 2) / aspect + 0.25) - 0.5 + hf = wf * aspect + if wf < 1: + h, w = (max_tokens - 2) // 2 * cell, cell + elif hf < 1: + h, w = cell, (max_tokens - 3) * cell + else: + scale = min(math.floor(wf) * cell / width, math.floor(hf) * cell / height) + h, w = math.floor(height * scale / patch) * patch, math.floor(width * scale / patch) * patch + return h, w + + +class DeepseekV41MMProcessor(MMProcessor): + def __init__(self, hf_config, model_path, mm): + super().__init__(model_path, mm) + self.vc = hf_config.vision_config + self.placeholder = [hf_config.image_token_id] + self.min_pixels = self.vc.min_pixels + self.max_tokens = mm.image_max_tokens if mm.image_max_tokens is not None else self.vc.max_image_tokens + if mm.image_min_tokens is not None: + # Delimiters and row breaks count against the total span budget. + self.min_pixels = max(1, mm.image_min_tokens - 2) * (self.vc.patch_size * self.vc.downsample_ratio)**2 + allowed = {"min_pixels", "max_image_tokens", "max_wh_ratio"} + unknown = set(mm.processor_kwargs) - allowed + if unknown: + raise ValueError(f"unsupported DeepSeek image processor options: {sorted(unknown)}") + self.min_pixels = mm.processor_kwargs.get("min_pixels", self.min_pixels) + self.max_tokens = mm.processor_kwargs.get("max_image_tokens", self.max_tokens) + self.max_wh_ratio = mm.processor_kwargs.get("max_wh_ratio", self.vc.max_wh_ratio) + if self.max_tokens < 4 or self.min_pixels < 0: + raise ValueError("image token budget must be >= 4 and min_pixels must be nonnegative") + if self.max_wh_ratio is not None and self.max_wh_ratio <= 0: + raise ValueError("max_wh_ratio must be positive") + + def process(self, images): + p, r = self.vc.patch_size, self.vc.downsample_ratio + items = [] + for image in images: + h, w = image_grid(image.width, image.height, p, r, self.min_pixels, self.max_tokens, self.max_wh_ratio) + if self.max_wh_ratio is not None and image.width >= self.max_wh_ratio * image.height: + image = image.resize((w, h)) + else: + image = ImageOps.pad(image, (w, h), color=(127, 127, 127)) + x = torch.from_numpy(np.asarray(image, dtype=np.float32)).permute(2, 0, 1) / 255 + x = ((x - 0.5) / 0.5).to(torch.bfloat16) + nh, nw = h // p, w // p + patches = x.reshape(3, nh, p, nw, p).permute(1, 3, 0, 2, 4).reshape(nh * nw, 3, p, p).contiguous() + item_hash = content_hash(patches, struct.pack("<3i", 1, nh, nw)) + items.append(MMItem(modality="image", hash=item_hash, pad_value=mm_pad_value(item_hash), offsets=[], feature=patches, + model_specific_data={"grid_thw": [1, nh, nw]})) + return items + + def prompt_replacement(self, item): + _, h, w = item.grid_thw + r = self.vc.downsample_ratio + n = ((h + r - 1) // r) * ((w + r - 1) // r + 1) + 2 + # Every span row, including delimiters/newlines, has a learned visual embedding. + return PromptReplacement([self.placeholder[0]] * n) + + def dummy_items(self, dtype, device): + p = self.vc.patch_size + return [MMItem(modality="image", hash=0, pad_value=0, offsets=[[0, 4]], + feature=torch.zeros(1, 3, p, p, dtype=dtype, device=device), model_specific_data={"grid_thw": [1, 1, 1]})] diff --git a/python/freetoken/models/config.py b/python/freetoken/models/config.py index c5943dce1b..2370d7b82c 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 @@ -205,7 +224,8 @@ class DSV4AttentionGroupConfig(BaseAttentionGroupConfig): """DSV4 sparse attention (window + compressed tiers + Lightning Indexer). Standalone class on purpose: subclassing SWAAttentionGroupConfig would flip has_swa_attention and reroute DSV4 through the SWA gates. Geometry beyond the latent width lives in - dsv4_args; the pool prices itself from there, not from this spec.""" + dsv4_args (``dsv41_args`` for the ``v41`` variant); the pool prices itself from there, + not from this spec.""" kind: ClassVar[Literal["dsv4"]] = "dsv4" cache_kind: ClassVar[Literal["dsv4_paged"]] = "dsv4_paged" @@ -213,6 +233,8 @@ class on purpose: subclassing SWAAttentionGroupConfig would flip has_swa_attenti num_kv_heads: int head_dim: int sliding_window: int # the P-token window page + # "v41": DeepSeek-V4.1's cross-layer-shared packed tiers (AttnType.DSV41, its own pool and backend) + variant: Literal["v4", "v41"] = "v4" AttentionGroupConfig: TypeAlias = ( @@ -327,6 +349,9 @@ class ModelConfig: # CSA/HCA compressors, Lightning Indexer, manifold-constrained Hyper-Connections, # hash routing). Opaque to model-agnostic engine code; None for non-DSV4 models. dsv4_args: Any | None = None + # DeepSeek-V4.1 payload (DeepseekV41Args): DSV41 layer roles, Engram, replay policy. Opaque + # to model-agnostic engine code; None for every other model. + dsv41_args: Any | None = None # GLM-5.2 (glm_moe_dsa) MLA/DSA payload (GlmMoeDsaArgs): the MLA low-rank dims and the # DSA indexer geometry the model module needs. Opaque to model-agnostic engine code; # None for every other model. @@ -448,7 +473,7 @@ def attn_type_for_layer(self, layer_id: int) -> AttnType: if isinstance(group, SWAAttentionGroupConfig): return AttnType.SWA if isinstance(group, DSV4AttentionGroupConfig): - return AttnType.DSV4 + return AttnType.DSV41 if group.variant == "v41" else AttnType.DSV4 return _full_group_attn_type(group) def kv_cache_group_specs(self) -> Tuple[KVCacheGroupSpec, ...]: @@ -502,7 +527,7 @@ def kv_cache_group_specs(self) -> Tuple[KVCacheGroupSpec, ...]: num_kv_heads=group.num_kv_heads, head_dim=group.head_dim, sliding_window=group.sliding_window, - attn_type=AttnType.DSV4, + attn_type=AttnType.DSV41 if group.variant == "v41" else AttnType.DSV4, ) ) return tuple(specs) diff --git a/python/freetoken/models/deepseek_v4/__init__.py b/python/freetoken/models/deepseek_v4/__init__.py index 37f5076ba7..e19438d1ba 100644 --- a/python/freetoken/models/deepseek_v4/__init__.py +++ b/python/freetoken/models/deepseek_v4/__init__.py @@ -10,7 +10,7 @@ DeepSeek-V4-Flash is a first-class registered model on the shared paged-KV engine: its window / compressed-attention / compressed-index KV live in DSV4-owned pools addressed by page tables, and sparse attention is a physical-slot gather (see -:mod:`freetoken.attention.dsv4_sparse` and :mod:`freetoken.kvcache.dsv4_paged_pool`). +:mod:`freetoken.attention.dsv4_sparse` and :mod:`freetoken.kvcache.dsv4.v4_pool`). """ from .args import DeepseekV4Args, load_args diff --git a/python/freetoken/models/deepseek_v4/model.py b/python/freetoken/models/deepseek_v4/model.py index a2d4387ade..9faefb83fe 100644 --- a/python/freetoken/models/deepseek_v4/model.py +++ b/python/freetoken/models/deepseek_v4/model.py @@ -4,7 +4,7 @@ + stateful KV compressors / Lightning Indexer, manifold-constrained Hyper-Connections, sqrtsoftplus / hash MoE), wired onto FreeToken's shared paged engine: - - KV lives in DSV4-owned paged pools (:class:`~freetoken.kvcache.dsv4_paged_pool.DSV4PagedKVCache`): + - KV lives in DSV4-owned paged pools (:class:`~freetoken.kvcache.dsv4.v4_pool.DSV4PagedKVCache`): each layer's window-ring + compressed KV is a region of a shared global pool, addressed by per-layer slot maps. Sparse attention is a PAGED physical-slot gather: ``sparse_attn_paged`` reads KV directly from the two global pools (window / compressed) at diff --git a/python/freetoken/models/deepseek_v4/moe.py b/python/freetoken/models/deepseek_v4/moe.py index 6b6245674d..83e6585403 100644 --- a/python/freetoken/models/deepseek_v4/moe.py +++ b/python/freetoken/models/deepseek_v4/moe.py @@ -106,7 +106,7 @@ def _prefill_routed( 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) diff --git a/python/freetoken/models/deepseek_v41/__init__.py b/python/freetoken/models/deepseek_v41/__init__.py new file mode 100644 index 0000000000..bcdef078c7 --- /dev/null +++ b/python/freetoken/models/deepseek_v41/__init__.py @@ -0,0 +1,14 @@ +"""DeepSeek-V4.1-Flash: Causal Encoder-Decoder MoE with Compressed Sparse Attention 2. + +40 layers (20 encoder + 20 decoder), every layer a sliding window (128) plus -- from layer 2 -- global +sparse attention over cross-layer-shared compressed KV in packed fp4 (Full / Reindex / Reuse +modes, a hierarchical candidate pool in the decoder), single-pass mHC residual streams, Engram n-gram +memory at layers 1 and 14 streamed from disk, one shared + 384 routed MXFP4 experts per layer served +from the offload cache. Images use the DeepSeek ViT; DSpark speculative decoding is not served. +""" + +from .config import parse_config +from .model import DeepseekV41ForCausalLM +from .weight import ftw_side_files, iter_expert_pieces, iter_weights, iter_vision_weights + +__all__ = ["DeepseekV41ForCausalLM", "parse_config", "iter_weights", "iter_vision_weights", "iter_expert_pieces", "ftw_side_files"] diff --git a/python/freetoken/models/deepseek_v41/args.py b/python/freetoken/models/deepseek_v41/args.py new file mode 100644 index 0000000000..867dcb21a9 --- /dev/null +++ b/python/freetoken/models/deepseek_v41/args.py @@ -0,0 +1,225 @@ +"""DeepSeek-V4.1-Flash hyperparameters and the static per-layer attention roles. + +``DeepseekV41Args`` mirrors the reference ``ModelArgs`` (``inference/model.py``) field for field, +read from the HF ``config.json``'s ``text_config`` (the ``inference/config.json`` dialect is also +accepted). Everything the engine reconciles at runtime (``max_seq_len``, the replay policy) is +attached by ``parse_config`` / the engine, not read from the checkpoint. + +``LayerRole`` is the attention mode table (tech report sec. 2.3.1): every compressing layer is Full +(owns the compressed KV + indexer K), Reindex (own indexer over a shared K) or Reuse (shared +Top-K); the candidate pool of the Hierarchical Sparse Indexer (sec. 2.3.2) is produced by the +first decoder Full layer and consumed by the decoder Reindex layers after it. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from enum import Enum +from typing import Any, Literal + +# text_config key -> ModelArgs key, where the two dialects differ +_HF_TO_ARGS = { + "hidden_size": "dim", + "moe_intermediate_size": "moe_inter_dim", + "num_hidden_layers": "n_layers", + "num_nextn_predict_layers": "n_mtp_layers", + "num_attention_heads": "n_heads", + "num_experts_per_tok": "n_activated_experts", + "scoring_func": "score_func", + "routed_scaling_factor": "route_scale", + "qk_rope_head_dim": "rope_head_dim", + "rms_norm_eps": "norm_eps", + "sliding_window": "window_size", + "kv_source_layer_ids": "kv_source_layers", + "index_source_layer_ids": "index_source_layers", + "candidate_source_layer_id": "candidate_source_layer", + "engram_pad_token_id": "engram_pad_id", + "dspark_num_experts_per_tok": "dspark_n_activated_experts", +} + + +class Mode(str, Enum): + WINDOW = "window" # sliding window only (compress_ratio == 0) + FULL = "full" # computes main KV + indexer K, runs its indexer + REINDEX = "reindex" # reuses main KV + indexer K, runs its own indexer + REUSE = "reuse" # reuses main KV and the latest Top-K + + +@dataclass(frozen=True) +class LayerRole: + layer_id: int + ratio: int + mode: Mode + kv_source: int | None # the layer whose main KV / indexer K this layer reads (itself when Full) + is_candidate_source: bool # builds the hierarchical candidate pool + uses_candidates: bool # indexes only inside the candidate pool + + @property + def compresses(self) -> bool: + return self.ratio > 0 + + @property + def has_indexer(self) -> bool: + return self.mode in (Mode.FULL, Mode.REINDEX) + + +@dataclass +class DeepseekV41Args: + vocab_size: int = 129280 + dim: int = 5120 + moe_inter_dim: int = 2304 + n_layers: int = 40 + n_mtp_layers: int = 3 + n_heads: int = 64 + # moe + n_routed_experts: int = 384 + n_shared_experts: int = 1 + n_activated_experts: int = 6 + score_func: Literal["softmax", "sigmoid", "sqrtsoftplus"] = "sqrtsoftplus" + norm_topk_prob: bool = True + route_scale: float = 1.5 + swiglu_limit: float = 10.0 + # attention: latent q, one shared latent kv head, grouped low-rank output projection + q_lora_rank: int = 1280 + head_dim: int = 512 + rope_head_dim: int = 64 + norm_eps: float = 1e-20 + o_groups: int = 8 + o_lora_rank: int = 1024 + # sparse attention + window_size: int = 128 + compress_ratios: tuple[int, ...] = () # one per layer, MTP layers included + kv_source_layers: tuple[int, ...] = () + index_source_layers: tuple[int, ...] = () + compress_rope_theta: float = 160000.0 + original_seq_len: int = 65536 + rope_theta: float = 10000.0 + rope_factor: float = 16.0 + beta_fast: int = 32 + beta_slow: int = 1 + index_n_heads: int = 32 + index_head_dim: int = 128 + index_topk: int = 512 + candidate_source_layer: int = -1 + candidate_topk_blocks: int = 0 + candidate_block_size: int = 0 + # hyper-connections (single-pass mHC) + hc_mult: int = 4 + hc_sinkhorn_iters: int = 20 + hc_eps: float = 1e-6 + # engram + engram_layer_ids: tuple[int, ...] = () + engram_num_embeddings: tuple[int, ...] = () + engram_max_ngram_size: int = 4 + engram_vocab_size: int = 0 + engram_n_heads: int = 0 + engram_head_dim: int = 0 + engram_pad_id: int = 2 + engram_compressed_vocab_size: int = 0 + # dspark (weights are not served in v1; kept so the config round-trips) + dspark_block_size: int = 0 + dspark_target_layer_ids: tuple[int, ...] = () + # ---- runtime (engine-reconciled, never read from the checkpoint) ---- + max_seq_len: int = 4096 + max_batch_size: int = 1 + # "bounded": the decoder runs on the prompt's last window_size tokens with the sliding window + # truncated there (Decoder SWA Bounded Replay, tech report sec. 3.2.2); "exact": on every token. + swa_decoder_replay: Literal["bounded", "exact"] = "bounded" + roles: tuple[LayerRole, ...] = field(default_factory=tuple, repr=False) + + def __post_init__(self) -> None: + for name in ("compress_ratios", "kv_source_layers", "index_source_layers", "engram_layer_ids", + "engram_num_embeddings", "dspark_target_layer_ids"): + setattr(self, name, tuple(int(v) for v in getattr(self, name))) + if len(self.compress_ratios) < self.n_layers: + raise ValueError(f"compress_ratios has {len(self.compress_ratios)} entries for {self.n_layers} layers") + if self.n_shared_experts != 1: + raise ValueError("DeepSeek-V4.1 has exactly one shared expert") + if self.hc_mult < 1: + raise ValueError("hc_mult must be positive") + self.roles = tuple(self._role(layer) for layer in range(self.n_layers)) + if len(self.engram_layer_ids) != len(self.engram_num_embeddings): + raise ValueError("engram_layer_ids and engram_num_embeddings differ in length") + + # ---- derived ---- + def _role(self, layer: int) -> LayerRole: + ratio = self.compress_ratios[layer] + if ratio == 0: + return LayerRole(layer, 0, Mode.WINDOW, None, False, False) + sources = [s for s in self.kv_source_layers if s <= layer] + if not sources: + raise ValueError(f"layer {layer} compresses but no kv source precedes it") + source = max(sources) + if self.compress_ratios[source] != ratio: + raise ValueError(f"layer {layer} (ratio {ratio}) reads kv source {source} (ratio {self.compress_ratios[source]})") + if layer == source: + mode = Mode.FULL + elif layer in self.index_source_layers: + mode = Mode.REINDEX + else: + mode = Mode.REUSE + cand = self.candidate_source_layer + is_cand_src = cand >= 0 and layer == cand + uses = cand >= 0 and cand < layer and mode is Mode.REINDEX + if uses and self.kv_source_layers and max(s for s in self.kv_source_layers if s <= cand) != source: + raise ValueError(f"layer {layer} would index the candidate pool of layer {cand} over a different kv source") + return LayerRole(layer, ratio, mode, source, is_cand_src, uses) + + @property + def nope_head_dim(self) -> int: + return self.head_dim - self.rope_head_dim + + @property + def backbone_kv_sources(self) -> tuple[int, ...]: + return tuple(s for s in self.kv_source_layers if s < self.n_layers) + + @property + def decoder_start_layer(self) -> int: + """The first decoder layer: the last kv source (CED projects the decoder's global KV from + its input, the final encoder hidden state).""" + return max(self.backbone_kv_sources) if self.backbone_kv_sources else self.n_layers + + @property + def engram_hash_cols(self) -> int: + return (self.engram_max_ngram_size - 1) * self.engram_n_heads + + @property + def index_softmax_scale(self) -> float: + return self.index_head_dim**-0.5 + + def freqs_params(self, layer: int) -> tuple[int, int, float, float, int, int]: + """``(rotary_dim, original_seq_len, theta, factor, beta_fast, beta_slow)``: compressing layers + rotate at the compressed theta under YaRN, window-only layers at the base theta without it.""" + if self.compress_ratios[layer]: + return self.rope_head_dim, self.original_seq_len, self.compress_rope_theta, self.rope_factor, self.beta_fast, self.beta_slow + return self.rope_head_dim, 0, self.rope_theta, self.rope_factor, self.beta_fast, self.beta_slow + + @classmethod + def from_hf(cls, hf_config: Any) -> "DeepseekV41Args": + """From the HF config (``text_config`` of the multimodal wrapper, or the text config itself, + or the ``inference/config.json`` dialect).""" + text = getattr(hf_config, "text_config", None) or hf_config + if isinstance(text, dict): + raw = text + elif hasattr(text, "to_dict"): # PretrainedConfig / RawConfigShim + raw = text.to_dict() + else: + raw = vars(text) + names = {f.name for f in cls.__dataclass_fields__.values()} - {"roles", "max_seq_len", "max_batch_size", "swa_decoder_replay"} + kwargs: dict[str, Any] = {} + for key, value in raw.items(): + key = _HF_TO_ARGS.get(key, key) + if key in names: + kwargs[key] = value + # HF nests the YaRN parameters; the inference dialect keeps them flat + scaling = raw.get("rope_scaling") or raw.get("rope_parameters") + if isinstance(scaling, dict): + if "factor" in scaling: + kwargs["rope_factor"] = scaling["factor"] + for src, dst in (("beta_fast", "beta_fast"), ("beta_slow", "beta_slow"), ("original_max_position_embeddings", "original_seq_len")): + if src in scaling: + kwargs[dst] = scaling[src] + return cls(**kwargs) + + +__all__ = ["DeepseekV41Args", "LayerRole", "Mode"] diff --git a/python/freetoken/models/deepseek_v41/attention.py b/python/freetoken/models/deepseek_v41/attention.py new file mode 100644 index 0000000000..1643ddb6d1 --- /dev/null +++ b/python/freetoken/models/deepseek_v41/attention.py @@ -0,0 +1,305 @@ +"""DSV41 attention layer (reference ``Attention``): latent MLA over a sliding window plus, on +compressing layers, ``index_topk`` compressed positions -- one paged sparse-attention call. + +Per layer mode (``LayerRole``): + +* Full -- ``publish_*`` compresses its input into the kv source's main KV (fp4) and index keys + (fp4), then its indexer selects fresh Top-K rows and publishes them on the forward's + ``SharedSelection``; the decoder's first Full layer also builds the candidate pool. +* Reindex -- selects fresh Top-K rows with its own indexer over the shared keys (inside the + candidate pool in the decoder) and publishes them. +* Reuse -- attends over the latest published Top-K rows. +* Window -- the sliding window only. + +All modes compute their own query and sliding-window KV (fp8 into the layer's window ring). The +attention output is rotated back (K == V share one rotated latent) and projected through the +grouped low-rank ``wo_a`` (block-diagonal over ``o_groups``; fp8 payload kept, W8A16 at decode) and ``wo_b``. + +Prefill runs ragged over a flat token stream tiled by ``PrefillSegment``s; the compressor / indexer +state is per request, so those run per segment. Decode is batched (one token per row) and +CUDA-graph safe: every gather reads the backend's snapshot, never the live page table. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import List + +import torch +from freetoken.attention.dsv41_sparse import PrefillSegment +from freetoken.core import get_global_ctx +from freetoken.kernel.triton.dsv4.fp8_linear import GEMV_MAX_M, dequant_block_fp8, grouped_w8a16_gemv +from freetoken.layers import BaseOP, LinearReplicated, RMSNorm +from freetoken.layers.quantization import QuantKind +from freetoken.layers.quantization.scheme import fp8_block_size + +from .args import DeepseekV41Args, LayerRole, Mode +from .compressor import Compressor +from .indexer import Indexer +from .rope import apply_rotary_emb, apply_rotary_emb_decode, get_freqs_cis + + +class DSV41Attention(BaseOP): + def __init__(self, args: DeepseekV41Args, role: LayerRole, *, quant_config=None, prefix: str = ""): + self.args = args + self.role = role + self.layer_id = role.layer_id + self.n_heads = args.n_heads + self.head_dim = args.head_dim + self.rope_dim = args.rope_head_dim + self.n_groups = args.o_groups + self.o_lora_rank = args.o_lora_rank + self.window = args.window_size + self.softmax_scale = args.head_dim**-0.5 + + self.wq_a = LinearReplicated(args.dim, args.q_lora_rank, has_bias=False, quant_config=quant_config, prefix=f"{prefix}.wq_a") + self.q_norm = RMSNorm(args.q_lora_rank, args.norm_eps) + self.wq_b = LinearReplicated(args.q_lora_rank, args.n_heads * args.head_dim, has_bias=False, quant_config=quant_config, prefix=f"{prefix}.wq_b") + self.wkv = LinearReplicated(args.dim, args.head_dim, has_bias=False, quant_config=quant_config, prefix=f"{prefix}.wkv") + self.kv_norm = RMSNorm(args.head_dim, args.norm_eps) + # Block-diagonal over groups. The reference dequantizes it to bf16 and runs a grouped einsum + # (the activation stays bf16); a quantized checkpoint keeps the e4m3 payload + e8m0 scale + # here and dequantizes in-kernel -- the same math, half the bytes per decode step, and no + # bf16 copy in VRAM. An unquantized checkpoint ships it bf16. + wo_a_scheme = quant_config.scheme_for(f"{prefix}.wo_a") if quant_config is not None else None + self._wo_a_block = fp8_block_size(wo_a_scheme) if wo_a_scheme is not None and wo_a_scheme.kind is QuantKind.FP8_BLOCK else None + rows, cols = args.o_groups * args.o_lora_rank, args.n_heads * args.head_dim // args.o_groups + if self._wo_a_block is not None: + self.wo_a = torch.empty(rows, cols, dtype=torch.float8_e4m3fn) + self.wo_a_scale = torch.empty(rows // self._wo_a_block, cols // self._wo_a_block, dtype=torch.float8_e8m0fnu) + else: + self.wo_a = torch.empty(rows, cols, dtype=torch.bfloat16) + self.wo_b = LinearReplicated(args.o_groups * args.o_lora_rank, args.dim, has_bias=False, quant_config=quant_config, prefix=f"{prefix}.wo_b") + self.attn_sink = torch.empty(args.n_heads, dtype=torch.float32) + + self.compressor = Compressor(args, self.layer_id, role.ratio, quant_config=quant_config, prefix=f"{prefix}.compressor") if role.mode is Mode.FULL else None + self.indexer = Indexer(args, role, quant_config=quant_config, prefix=f"{prefix}.indexer") if role.has_indexer else None + self._freqs_params = args.freqs_params(self.layer_id) + self._freqs: torch.Tensor | None = None + + # addressing lives on the live backend, pool buffers on the live pool (both read per access) + @property + def attn(self): + return get_global_ctx().attn_backend + + @property + def freqs_key(self) -> tuple: + """The RoPE table this layer rotates with (window layers: base theta; compressing layers: the + compressed theta under YaRN). Layers with the same key share one table.""" + return self._freqs_params + + def bind(self, device: torch.device, tables: dict | None = None) -> None: + """Attach the RoPE table; ``tables`` (key -> table) lets the model share one table per key + instead of every layer holding its own ``[max_seq_len, rd // 2]`` copy.""" + if tables is not None and self.freqs_key in tables: + self._freqs = tables[self.freqs_key] + return + rd, original_seq_len, theta, factor, beta_fast, beta_slow = self._freqs_params + self._freqs = get_freqs_cis(rd, self.args.max_seq_len, original_seq_len, theta, factor, beta_fast, beta_slow, device) + if tables is not None: + tables[self.freqs_key] = self._freqs + + @property + def freqs(self) -> torch.Tensor: + assert self._freqs is not None, "attention layer is not bound (call bind first)" + return self._freqs + + # ----- shared pieces ------------------------------------------------------------------- + def _q(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + qr = self.q_norm.forward(self.wq_a.forward(x)) + q = self.wq_b.forward(qr).unflatten(-1, (self.n_heads, self.head_dim)) + return qr, q + + def _kv(self, x: torch.Tensor) -> torch.Tensor: + return self.kv_norm.forward(self.wkv.forward(x)) + + def _wo(self, o: torch.Tensor) -> torch.Tensor: + """``wo_b(wo_a(o))``: up to ``GEMV_MAX_M`` rows the grouped W8A16 GEMV over the fp8 payload, else the + reference's grouped einsum over the bf16 (or dequantized) ``wo_a``.""" + t = o.shape[0] + o = o.reshape(t, self.n_groups, -1) + if self._wo_a_block is not None and t <= GEMV_MAX_M: + return self.wo_b.forward(grouped_w8a16_gemv(o, self.wo_a, self.wo_a_scale, block=self._wo_a_block)) + wo_a = self.wo_a if self._wo_a_block is None else dequant_block_fp8(self.wo_a, self.wo_a_scale, block=self._wo_a_block) + return self.wo_b.forward(torch.einsum("tgd,grd->tgr", o, wo_a.view(self.n_groups, self.o_lora_rank, -1)).flatten(1)) + + def load_state_dict(self, state_dict, *, prefix: str = "", _internal: bool = False) -> None: + key = f"{prefix}.wo_a" if prefix else "wo_a" + scale_key = f"{key}_scale" + if self._wo_a_block is None and scale_key in state_dict: + # a bf16 layer (no fp8 scheme) fed a quantized checkpoint dequantizes wo_a on the way in + from .weight import dequant_fp8_block + + state_dict[key] = dequant_fp8_block(state_dict[key], state_dict.pop(scale_key)) + elif self._wo_a_block is not None and scale_key not in state_dict: + # an FTW conversion made before wo_a kept its fp8 payload carries a dequantized bf16 + # tensor and no scale: say so instead of failing on the missing key + raise RuntimeError( + f"{key}: the checkpoint has no {scale_key} for the fp8 wo_a layout; an FTW checkpoint " + "converted by an earlier build stores wo_a dequantized -- reconvert it with `ft checkpoint`" + ) + super().load_state_dict(state_dict, prefix=prefix, _internal=_internal) + + # ----- prefill ----------------------------------------------------------------------------- + def publish_prefill(self, x: torch.Tensor, segments: List[PrefillSegment]) -> None: + """Full mode: compress this layer's input for every segment into the source's main KV and + index keys. Runs on EVERY new token (under bounded replay the decoder source publishes for + the whole prompt even though its attention runs on the prompt's last window only).""" + assert self.compressor is not None and self.indexer is not None + attn, ratio, src = self.attn, self.role.ratio, self.layer_id + for seg in segments: + slots = attn.window_slots_of(seg.table_idx, seg.start_pos, seg.end) + tail = None + if self.compressor.needs_tail_carry(seg.start_pos): + tail = int(attn.window_slots_of(seg.table_idx, seg.start_pos - 1, seg.start_pos).item()) + latent, starts = self.compressor.forward_prefill(x[seg.offset : seg.offset + seg.n], seg.start_pos, slots, tail) + if not latent.shape[0]: + continue + rows = attn.compressed_rows_of(seg.table_idx, starts, ratio) + freqs = self.freqs.index_select(0, starts) + attn.store_index(self.indexer.index_keys(latent, freqs), src, rows) + apply_rotary_emb(latent[..., -self.rope_dim :], freqs) + attn.store_main(latent, src, rows) + + def forward_prefill(self, x: torch.Tensor, segments: List[PrefillSegment], positions: torch.Tensor) -> torch.Tensor: + """Ragged prefill over ``x [T, dim]`` tiled by ``segments`` (``positions [T]`` absolute).""" + attn, role = self.attn, self.role + freqs = self.freqs.index_select(0, positions) + qr, q = self._q(x) + apply_rotary_emb(q[..., -self.rope_dim :], freqs) + kv = self._kv(x) + apply_rotary_emb(kv[..., -self.rope_dim :], freqs) + + sel = attn.selection + win_parts, cmp_parts, pools = [], [], [] + for seg in segments: + lo, hi = seg.offset, seg.offset + seg.n + attn.store_window(kv[lo:hi], self.layer_id, attn.layer_window_slots_of(self.layer_id, seg.table_idx, seg.start_pos, seg.end)) + win_parts.append(attn.window_topk_prefill(seg, self.layer_id)) + if not role.compresses: + continue + if role.has_indexer: + locs = attn.locs_prefill(seg.table_idx, seg.end // role.ratio * role.ratio) + cand = sel.candidates[:, lo:hi] if role.uses_candidates else None + rows, pool = self.indexer.select_prefill( + x[lo:hi], qr[lo:hi], freqs[lo:hi], start_pos=seg.start_pos, ratio=role.ratio, locs=locs, candidates=cand, + ) + cmp_parts.append(rows) + if pool is not None: + pools.append(pool) + else: + assert sel.source == role.kv_source, "Reuse layer without a published selection for its source" + cmp_parts.append(sel.topk_rows[:, lo:hi]) + if role.has_indexer: + sel.source, sel.topk_rows = role.kv_source, torch.cat(cmp_parts, dim=1) + if pools: + sel.candidates = torch.cat(pools, dim=1) + win = torch.cat(win_parts, dim=1) + topk = torch.cat([win, torch.cat(cmp_parts, dim=1)], dim=-1) if role.compresses else win + + o = attn.attend(q.unsqueeze(0), self.layer_id, topk.to(torch.int32), self.window, self.attn_sink, self.softmax_scale)[0] + apply_rotary_emb(o[..., -self.rope_dim :], freqs, inverse=True) + return self._wo(o) + + # ----- decode ------------------------------------------------------------------------------ + def forward_decode(self, x: torch.Tensor, pos: torch.Tensor, rows: torch.Tensor, dctx: "DecodeStepContext", cmp_stage_cap: int) -> torch.Tensor: + """Batched single-token attention. ``pos [B]`` device positions, ``rows [B]`` local rows into + the decode snapshot, ``dctx`` the step's layer-invariant context (ring slots, window + candidates, RoPE rows), ``cmp_stage_cap`` the static compressed staging bound (max position + any row reaches; the capture width under a graph).""" + attn, role = self.attn, self.role + B = x.shape[0] + freqs_t = dctx.freqs[self.freqs_key] + qr, q = self._q(x) + apply_rotary_emb_decode(q[..., -self.rope_dim :].unsqueeze(1), freqs_t) + kv = self._kv(x) + apply_rotary_emb_decode(kv[..., -self.rope_dim :].unsqueeze(1), freqs_t) + # the shared ring context feeds the compressor's carry ring on every layer; a request-private + # window layer stores and reads its own KV through its ring context instead + window_slots, prev_window_slots = dctx.window_slots, dctx.prev_window_slots + private = dctx.private is not None and attn.pool.is_private_window(self.layer_id) + tier = dctx.private if private else dctx.shared + attn.store_window(kv, self.layer_id, tier.slots if private else window_slots) + + cmp_counts = None + if role.compresses: + ratio, src = role.ratio, role.kv_source + n_stage = (cmp_stage_cap + 1) // ratio + sel = attn.selection + if role.mode is Mode.FULL: + latent, completed = self.compressor.forward_decode(x, pos, prev_window_slots, window_slots) + dst = attn.decode_store_rows(rows, pos, ratio, src, completed) + freqs_g = self.freqs.index_select(0, (pos + 1 - ratio).clamp_min(0)) + k = self.indexer.k_norm.forward(self.indexer.wk.forward(latent)) + apply_rotary_emb_decode(k[..., -self.rope_dim :].unsqueeze(1), freqs_g) + attn.store_index(k, src, dst) + apply_rotary_emb_decode(latent[..., -self.rope_dim :].unsqueeze(1), freqs_g) + attn.store_main(latent, src, dst) + if role.has_indexer: + cand = sel.candidates if role.uses_candidates else None + cmp_rows, counts, pool = self.indexer.select_decode(x, qr, freqs_t, pos=pos, ratio=ratio, locs=attn.snapshot(), T=n_stage, candidates=cand) + sel.source, sel.topk_rows, sel.cmp_counts = src, cmp_rows, counts + if pool is not None: + sel.candidates = pool + # the step's [window | compressed] candidate buffer of this tier takes the new rows once; + # the Reuse layers that follow read them in place (no per-layer cat / cast) + tier.topk[..., self.window :].copy_(cmp_rows) + else: + assert sel.source == src, "Reuse layer without a published selection for its source" + counts = sel.cmp_counts + topk = tier.topk + cmp_counts = counts + else: + topk = tier.window_topk + + o = attn.attend(q.view(B, 1, self.n_heads, self.head_dim), self.layer_id, topk, self.window, self.attn_sink, self.softmax_scale, cmp_counts=cmp_counts) + apply_rotary_emb_decode(o[..., -self.rope_dim :], freqs_t, inverse=True) + return self._wo(o.view(B, self.n_heads, self.head_dim)) + + +@dataclass +class TierCandidates: + """One window tier's per-step decode candidates: the request's store slot and the int32 + ``[B, 1, win]`` window candidates, plus the ``[B, 1, win + topk]`` ``[window | compressed]`` + buffer whose compressed half the index layers of that tier fill as they publish.""" + + slots: torch.Tensor + window_topk: torch.Tensor + topk: torch.Tensor + + +@dataclass +class DecodeStepContext: + """Everything a decode step resolves once for all layers (a capture records it once): the shared + pool's ring slots (the compressor's carry ring keys on them), the shared / private tier + candidates, and each RoPE table gathered at the step's positions.""" + + window_slots: torch.Tensor + prev_window_slots: torch.Tensor + shared: TierCandidates + private: TierCandidates | None + freqs: dict + + @classmethod + def build(cls, layers, md, pos: torch.Tensor, rows: torch.Tensor, window: int, topk: int) -> "DecodeStepContext": + B = pos.shape[0] + + def tier(slots, window_topk): + win32 = window_topk.to(torch.int32) + buf = torch.empty((B, 1, window + topk), dtype=torch.int32, device=pos.device) + buf[..., :window].copy_(window_topk) + return TierCandidates(slots, win32, buf) + + window_slots, prev_window_slots, window_topk = md.window_ctx(pos, rows) + private = None + if get_global_ctx().kv_cache.private_window_layer_ids: + private = tier(*md.private_window_ctx(pos, rows)) + tables: dict = {} + for layer in layers: + key = layer.freqs_key + if key not in tables: + tables[key] = layer.freqs.index_select(0, pos) + return cls(window_slots, prev_window_slots, tier(window_slots, window_topk), private, tables) + + +__all__ = ["DSV41Attention", "DecodeStepContext", "TierCandidates"] diff --git a/python/freetoken/models/deepseek_v41/compressor.py b/python/freetoken/models/deepseek_v41/compressor.py new file mode 100644 index 0000000000..5d5d07fe5b --- /dev/null +++ b/python/freetoken/models/deepseek_v41/compressor.py @@ -0,0 +1,131 @@ +"""DSV41 KV compressor: ``compress_ratio`` consecutive tokens -> one latent (reference ``Compressor``). + +* ratio 1 (the CED decoder source): a plain bf16 projection + RMSNorm, one latent per token. +* ratio > 1 (the encoder sources): ``wkv`` / ``wgate`` promoted to fp32, a softmax over the group's + gate scores pools the group's ``wkv`` rows; the trailing partial group is carried across chunks + and decode steps in the pool's compress-state ring (per window page, so a page-aligned radix hit + resumes it by value -- see ``DSV41SparseAttnBackend`` for the addressing). + +Returns the latent BEFORE RoPE: the indexer key is projected from the unrotated latent, the +attention layer rotates and quantizes it afterwards. +""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F +from freetoken.core import get_global_ctx +from freetoken.layers import BaseOP, LinearReplicated, RMSNorm + +from .args import DeepseekV41Args + + +class Compressor(BaseOP): + def __init__(self, args: DeepseekV41Args, layer_id: int, ratio: int, *, quant_config=None, prefix: str = ""): + assert ratio >= 1 + self.layer_id = layer_id + self.ratio = ratio + self.dim = args.dim + self.head_dim = args.head_dim + self.window = args.window_size + if ratio > 1: + # fp32 like the reference (the softmax pooling runs in fp32); the checkpoint ships bf16 + self.wkv = torch.empty(args.head_dim, args.dim, dtype=torch.float32) + self.wgate = torch.empty(args.head_dim, args.dim, dtype=torch.float32) + else: + self.wkv = LinearReplicated(args.dim, args.head_dim, has_bias=False, quant_config=quant_config, prefix=f"{prefix}.wkv") + self.norm = RMSNorm(args.head_dim, args.norm_eps) + + @property + def attn(self): + return get_global_ctx().attn_backend + + def _project(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """fp32 ``wkv`` / ``wgate`` projections.""" + xf = x.float() + return F.linear(xf, self.wkv), F.linear(xf, self.wgate) + + # ----- prefill ----------------------------------------------------------------------- + def needs_tail_carry(self, start_pos: int) -> bool: + """Whether a prefill from ``start_pos`` resumes a partial group (never at a page-aligned start).""" + return self.ratio > 1 and start_pos % self.ratio != 0 + + def forward_prefill( + self, x: torch.Tensor, start_pos: int, window_slots: torch.Tensor, tail_window_slot: int | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Compress the tokens ``[start_pos, start_pos + n)`` of one request. + + ``window_slots`` are the tokens' window slots (the ring block of each window page carries the + partial group); ``tail_window_slot`` is the previous token's slot, needed only when + ``needs_tail_carry(start_pos)``. Returns ``(latent [G, head_dim] bf16 pre-RoPE, group_starts [G] + int64)`` for every group that completed in this call. + """ + n = x.shape[0] + device = x.device + if self.ratio == 1: + latent = self.norm.forward(self.wkv.forward(x)) + return latent, torch.arange(start_pos, start_pos + n, device=device) + + ratio, d = self.ratio, self.head_dim + kv, score = self._project(x) + # the partial group in progress at start_pos: its rows live in the previous token's ring block + offset = start_pos % ratio + if offset: + assert tail_window_slot is not None, "resuming mid-group needs the previous token's window slot" + carry = self.attn.read_carry(self.layer_id, tail_window_slot) # [ratio, 2d] + kv = torch.cat([carry[:offset, :d], kv], dim=0) + score = torch.cat([carry[:offset, d:], score], dim=0) + total = kv.shape[0] + groups = total // ratio + remainder = total % ratio + cut = groups * ratio + if groups: + kv_g = kv[:cut].view(groups, ratio, d) + score_g = score[:cut].view(groups, ratio, d) + pooled = (kv_g * score_g.softmax(dim=1)).sum(dim=1) + latent = self.norm.forward(pooled.to(x.dtype)) + else: + latent = x.new_empty(0, d) + group_starts = (start_pos - offset) + ratio * torch.arange(groups, device=device) + # ring: the reset block at every page boundary crossed (a page holds whole groups), and the + # in-progress tail (the remainder rows) on the last page so decode / the next chunk resume + end = start_pos + n + self.attn.write_boundary_carries(self.layer_id, lo=start_pos, hi=end, window_slots=window_slots) + if end % self.window != 0: + block = self._carry_block(kv[cut:], score[cut:], remainder) + self.attn.write_carry(self.layer_id, int(window_slots[-1].item()), block) + return latent, group_starts + + def _carry_block(self, kv_tail: torch.Tensor, score_tail: torch.Tensor, filled: int) -> torch.Tensor: + ratio, d = self.ratio, self.head_dim + ks = kv_tail.new_zeros(ratio, d) + ss = kv_tail.new_full((ratio, d), float("-inf")) + if filled: + ks[:filled] = kv_tail + ss[:filled] = score_tail + return torch.cat([ks, ss], dim=-1) + + # ----- decode ------------------------------------------------------------------------ + def forward_decode(self, x: torch.Tensor, pos: torch.Tensor, prev_window_slots: torch.Tensor, window_slots: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """One token per row. Returns ``(latent [B, head_dim] bf16 pre-RoPE, completed [B] bool)``; + a row whose group did not complete carries a meaningless latent the caller discards.""" + B = x.shape[0] + if self.ratio == 1: + latent = self.norm.forward(self.wkv.forward(x)) + return latent, torch.ones(B, dtype=torch.bool, device=x.device) + + from freetoken.kernel.triton.dsv4.compress import gated_pool + + ratio, d = self.ratio, self.head_dim + kv, score = self._project(x) + idx = (pos % ratio).view(B, 1, 1).expand(B, 1, d) + block = self.attn.read_carry_blocks(self.layer_id, prev_window_slots) # [B, ratio, 2d] + ks = block[..., :d].clone().scatter_(1, idx, kv.unsqueeze(1)) + ss = block[..., d:].clone().scatter_(1, idx, score.unsqueeze(1)) + self.attn.write_carry_blocks(self.layer_id, window_slots, torch.cat([ks, ss], dim=-1)) + pooled = gated_pool(ks, ss, x.dtype).view(B, d) + completed = (pos + 1) % ratio == 0 + return self.norm.forward(pooled), completed + + +__all__ = ["Compressor"] diff --git a/python/freetoken/models/deepseek_v41/config.py b/python/freetoken/models/deepseek_v41/config.py new file mode 100644 index 0000000000..9e738939c8 --- /dev/null +++ b/python/freetoken/models/deepseek_v41/config.py @@ -0,0 +1,95 @@ +"""Engine-facing config for DeepSeek-V4.1-Flash. + +``parse_config`` maps the standard transformer fields into :class:`ModelConfig`, carries the full +:class:`DeepseekV41Args` in ``ModelConfig.dsv41_args`` for the model module and the KV pool (which +sizes its tiers from it), and declares the ``v41`` variant of the DSV4 attention group. +The engine reconciles the runtime knobs (``max_seq_len``, ``swa_decoder_replay``) onto ``dsv41_args`` +at config resolution. +""" + +from __future__ import annotations + +from dataclasses import dataclass, fields +from typing import Any + +from freetoken.models.config import DSV4AttentionGroupConfig, ModelConfig, RotaryConfig + +from .args import DeepseekV41Args + + +@dataclass(frozen=True) +class VisionConfig: + num_hidden_layers: int = 32 + hidden_size: int = 1024 + num_attention_heads: int = 16 + intermediate_size: int = 2816 + patch_size: int = 14 + rope_theta: float = 10000.0 + downsample_ratio: int = 3 + max_image_tokens: int = 1024 + min_pixels: int = 544 * 544 + max_wh_ratio: float | None = None + + +def parse_config(hf_config: Any) -> ModelConfig: + args = DeepseekV41Args.from_hf(hf_config) + text = getattr(hf_config, "text_config", None) or hf_config + max_position = int(getattr(text, "max_position_embeddings", 0) or getattr(hf_config, "max_position_embeddings", 0) or 0) + if max_position <= 0: + max_position = int(args.original_seq_len * args.rope_factor) + vc = getattr(hf_config, "vision_config", None) + vision = VisionConfig(**{f.name: getattr(vc, f.name, f.default) for f in fields(VisionConfig)}) if vc is not None else None + + return ModelConfig( + num_layers=args.n_layers, + num_qo_heads=args.n_heads, + num_kv_heads=1, # one shared latent KV head (K == V) + head_dim=args.head_dim, + hidden_size=args.dim, + vocab_size=args.vocab_size, + intermediate_size=args.moe_inter_dim, + hidden_act="silu", + rms_norm_eps=args.norm_eps, + tie_word_embeddings=False, + rotary_config=RotaryConfig( + head_dim=args.head_dim, + rotary_dim=args.rope_head_dim, + max_position=max_position, + base=args.rope_theta, + scaling={ + "rope_type": "yarn", + "factor": args.rope_factor, + "beta_fast": args.beta_fast, + "beta_slow": args.beta_slow, + "original_max_position_embeddings": args.original_seq_len, + }, + ), + num_experts=args.n_routed_experts, + num_experts_per_tok=args.n_activated_experts, + moe_intermediate_size=args.moe_inter_dim, + norm_topk_prob=args.norm_topk_prob, + n_shared_experts=args.n_shared_experts, + routed_scaling_factor=args.route_scale, + swiglu_limit=args.swiglu_limit, + model_type="deepseek_v41", + architectures=["DeepseekV41ForCausalLM"], + moe_enabled=True, + expert_quant="ds_fp4", + attn_sm_scale=args.head_dim**-0.5, + dsv41_args=args, + vision_config=vision, + image_token_id=getattr(hf_config, "image_token_id", None), + attention_groups=( + DSV4AttentionGroupConfig( + name="dsv41", + layer_ids=tuple(range(args.n_layers)), + num_kv_heads=1, + head_dim=args.head_dim, + sliding_window=args.window_size, + variant="v41", + ), + ), + ) + + +__all__ = ["parse_config"] diff --git a/python/freetoken/models/deepseek_v41/engram.py b/python/freetoken/models/deepseek_v41/engram.py new file mode 100644 index 0000000000..1c997a7cfd --- /dev/null +++ b/python/freetoken/models/deepseek_v41/engram.py @@ -0,0 +1,262 @@ +"""Engram conditional memory (reference ``inference/engram.py`` + ``Engram``): n-gram hash lookups +added into the residual streams before layers 1 and 14. + +Hashing (host side, ``EngramHash``): token ids are first mapped through a compressed vocabulary +(tokens that normalize alike collapse together), then each position is hashed with the +``max_ngram_size - 1`` tokens before it -- a rolling XOR of ``id * multiplier`` over the lookback, +one bucket range (a distinct prime modulus) per (n-gram order, head). Lookback stops at the sequence +start (pad id) and never crosses an image span. Addresses depend +only on the token ids, so the rows can be fetched before the forward starts. + +Lookup (``EngramTable`` protocol): the 384M-row fp8 tables live on disk; a backend gathers the +``n_hash_cols`` rows of every token into a device tensor ``[T, n_hash_cols * head_dim]`` in bf16. + +Gate (``EngramLayer``): ``wkv`` (fp8 linear) turns the gathered rows into one key per residual +stream plus a shared value; the gate is a normalized dot product of stream and key, signed-sqrt +then sigmoid; ``stream += gate * value``. +""" + +from __future__ import annotations + +from typing import Protocol + +import numpy as np +import torch +from freetoken.layers import BaseOP, LinearReplicated + +from .args import DeepseekV41Args + + +# ----- hashing --------------------------------------------------------------------------------- + + +def _is_prime(n: int) -> bool: + """Deterministic Miller-Rabin for n < 3.3e24 (the bases below are a proven set), matching + ``sympy.isprime`` on every 64-bit input.""" + if n < 2: + return False + small = (2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41) + for p in small: + if n % p == 0: + return n == p + d, s = n - 1, 0 + while d % 2 == 0: + d //= 2 + s += 1 + for a in small: + x = pow(a, d, n) + if x in (1, n - 1): + continue + for _ in range(s - 1): + x = pow(x, 2, n) + if x == n - 1: + break + else: + return False + return True + + +def find_next_prime(start: int, seen: set[int]) -> int: + """The smallest prime above ``start`` not handed out yet (reference ``find_next_prime``).""" + candidate = start + 1 + while not _is_prime(candidate) or candidate in seen: + candidate += 1 + return candidate + + +def bucket_primes(args: DeepseekV41Args) -> list[list[list[int]]]: + """``[layer][ngram order - 2][head]`` bucket moduli: distinct primes drawn in order from + ``engram_vocab_size - 1`` upward, never reused across layers / orders / heads.""" + primes: list[list[list[int]]] = [] + seen: set[int] = set() + for _ in args.engram_layer_ids: + per_ngram = [] + for _ in range(args.engram_max_ngram_size - 1): + sizes, current = [], args.engram_vocab_size - 1 + for _ in range(args.engram_n_heads): + current = find_next_prime(current, seen) + seen.add(current) + sizes.append(current) + per_ngram.append(sizes) + primes.append(per_ngram) + return primes + + +def hash_multipliers(layer_ids: tuple[int, ...], max_ngram_size: int, compressed_vocab_size: int) -> torch.Tensor: + """``[n_layers, max_ngram_size]`` odd int64 multipliers from a per-layer RNG (reference + ``compute_hash_multipliers``), bounded so ``id * multiplier`` cannot overflow int64.""" + max_long = np.iinfo(np.int64).max + bound = max(1, (max_long // compressed_vocab_size) // 2) + rows = [] + for layer_id in layer_ids: + gen = np.random.default_rng(10007 * layer_id) + values = gen.integers(low=0, high=bound, size=(max_ngram_size,), dtype=np.int64) + rows.append(torch.tensor(values * 2 + 1)) + return torch.stack(rows) + + +def build_compressed_token_map(tokenizer) -> tuple[list[int], int]: + """Token id -> compressed id: tokens that normalize alike (NFKC, accents stripped, lowercased, + whitespace collapsed) share one id. Reference ``build_compressed_token_map``; the size feeds the + multipliers, so it is asserted against ``engram_compressed_vocab_size``.""" + from tokenizers import Regex, normalizers + + sentinel = "\ue000" + normalizer = normalizers.Sequence([ + normalizers.NFKC(), + normalizers.NFD(), + normalizers.StripAccents(), + normalizers.Lowercase(), + normalizers.Replace(Regex(r"[ \t\r\n]+"), " "), + normalizers.Replace(Regex(r"^ $"), sentinel), + normalizers.Strip(), + normalizers.Replace(sentinel, " "), + ]) + backend = tokenizer.backend_tokenizer + key_to_new: dict[str, int] = {} + lookup = [0] * len(tokenizer) + for token_id in range(len(tokenizer)): + text = backend.decode([token_id], skip_special_tokens=False) + if "\ufffd" in text: + key = backend.id_to_token(token_id) + else: + normalized = normalizer.normalize_str(text) + key = normalized if normalized else text + new_id = key_to_new.get(key) + if new_id is None: + new_id = len(key_to_new) + key_to_new[key] = new_id + lookup[token_id] = new_id + return lookup, len(key_to_new) + + +class EngramHash: + """Maps token runs to table rows: ``row_ids(tokens) -> [T, n_layers, n_hash_cols]`` int64. + + ``token_map`` is the compressed vocabulary (``build_compressed_token_map``); ``pad_id`` is the + compressed id of ``engram_pad_id``. Runs on the host (numpy / CPU torch) since the addresses + are needed before the forward launches. + """ + + def __init__(self, args: DeepseekV41Args, token_map: list[int] | torch.Tensor, compressed_vocab_size: int): + if compressed_vocab_size != args.engram_compressed_vocab_size: + raise ValueError( + f"compressed vocab has {compressed_vocab_size} ids, config says {args.engram_compressed_vocab_size}; " + "the hash multipliers derive from it, so the tokenizer does not match the checkpoint" + ) + self.args = args + self.max_ngram = args.engram_max_ngram_size + self.token_map = torch.as_tensor(token_map, dtype=torch.int64) + self.pad_id = int(self.token_map[args.engram_pad_id]) + primes = bucket_primes(args) + flat = [[p for per_ngram in layer for p in per_ngram] for layer in primes] + self.primes = torch.tensor(primes, dtype=torch.int64) # [L, ngram-1, heads] + self.offsets = torch.tensor(np.array([np.cumsum([0, *sizes[:-1]]) for sizes in flat]), dtype=torch.int64) # [L, cols] + self.multipliers = hash_multipliers(args.engram_layer_ids, self.max_ngram, compressed_vocab_size) # [L, ngram] + self.rows_per_layer = [int(o[-1] + s[-1]) for o, s in zip(self.offsets.tolist(), [[p for pn in layer for p in pn] for layer in primes])] + + @property + def n_cols(self) -> int: + return (self.max_ngram - 1) * self.args.engram_n_heads + + def row_ids(self, tokens: torch.Tensor, context: torch.Tensor | None = None) -> torch.Tensor: + """Rows for ``tokens [T]`` (raw ids, one contiguous run). ``context`` holds the + ``max_ngram_size - 1`` raw ids immediately before the run (fewer -> the run starts a + sequence and the missing lookback is the pad id). Returns ``[T, n_layers, n_hash_cols]``.""" + tokens = tokens.to(torch.int64).cpu() + ctx_len = self.max_ngram - 1 + if context is None: + context = tokens.new_empty(0) + context = context.to(torch.int64).cpu()[-ctx_len:] if ctx_len else tokens.new_empty(0) + raw = torch.cat([context, tokens]) + # Content-pad ids identify image spans; n-grams stop at them, including across chunks. + dead = (raw < 0) | (raw >= self.token_map.numel()) + compressed = self.token_map[raw.clamp(0, self.token_map.numel() - 1)] + n_ctx = context.numel() + T = tokens.numel() + positions = torch.arange(n_ctx, n_ctx + T) + cols = [] + blocked = torch.zeros(T, dtype=torch.bool) + for shift in range(self.max_ngram): + src = positions - shift + blocked |= (src < 0) | dead[src.clamp_min(0)] + cols.append(torch.where(blocked, torch.full_like(src, self.pad_id), compressed[src.clamp_min(0)])) + stack = torch.stack(cols, dim=-1) # [T, max_ngram]; column s is the token s places back + products = stack.unsqueeze(1) * self.multipliers # [T, L, max_ngram] + rolling, hashes = products[..., 0], [] + for i in range(1, self.max_ngram): + rolling = torch.bitwise_xor(rolling, products[..., i]) + hashes.append(rolling.unsqueeze(-1) % self.primes[:, i - 1]) # [T, L, heads] + return torch.cat(hashes, dim=-1) + self.offsets # [T, L, cols] + + +# ----- lookup ------------------------------------------------------------------------------------ + + +class EngramTable(Protocol): + """Row store of one Engram layer: ``lookup`` returns the dequantized rows of the current + forward's tokens as ``[T, n_hash_cols * head_dim]`` bf16 (staged by the host before launch).""" + + def lookup(self, num_tokens: int) -> torch.Tensor: ... + + +class ZeroEngramTable: + """Stand-in when no table is attached (dummy weights / unit tests): every lookup reads zeros.""" + + def __init__(self, width: int, device: torch.device) -> None: + self.width = width + self.device = device + + def lookup(self, num_tokens: int) -> torch.Tensor: + return torch.zeros(num_tokens, self.width, dtype=torch.bfloat16, device=self.device) + + +# ----- the layer ----------------------------------------------------------------------------------- + + +class EngramLayer(BaseOP): + """``streams [T, hc, dim] += gate * value`` with key / value from the gathered n-gram rows.""" + + def __init__(self, args: DeepseekV41Args, layer_id: int, *, quant_config=None, prefix: str = ""): + self.layer_id = layer_id + self.dim = args.dim + self.hc_mult = args.hc_mult + self.eps = args.norm_eps + self.clamp_value = 1e-6 + self.width = args.engram_hash_cols * args.engram_head_dim + self.wkv = LinearReplicated(self.width, args.dim * (args.hc_mult + 1), has_bias=False, quant_config=quant_config, prefix=f"{prefix}.wkv") + self.q_weight = torch.empty(args.hc_mult, args.dim, dtype=torch.float32) + self.k_weight = torch.empty(args.hc_mult, args.dim, dtype=torch.float32) + self._table: EngramTable | None = None + + def attach_table(self, table: EngramTable) -> None: + self._table = table + + def forward(self, streams: torch.Tensor, rows: torch.Tensor | None = None, *, image_mask: torch.Tensor | None = None) -> torch.Tensor: + """``streams [T, hc, dim]``; ``rows`` overrides the table lookup (tests / reference parity).""" + T = streams.shape[0] + if rows is None: + assert self._table is not None, "Engram table was never attached" + rows = self._table.lookup(T) + kv = self.wkv.forward(rows.to(streams.dtype)) + key, value = kv.split([self.hc_mult * self.dim, self.dim], dim=-1) + key = key.float().view(T, self.hc_mult, self.dim) + weight = self.q_weight * self.k_weight # only ever used as a product + h = streams.float() + # normalized per (token, stream) over dim, not jointly over the streams + rstd = torch.rsqrt(h.square().mean(-1) + self.eps) * torch.rsqrt(key.square().mean(-1) + self.eps) + dot = (h * weight * key).sum(-1) * rstd * self.dim**-0.5 + gate = torch.sigmoid(torch.copysign(dot.abs().clamp_min(self.clamp_value).sqrt(), dot)) + out = (h + gate.unsqueeze(-1) * value.float().unsqueeze(1)).to(streams.dtype) + return out if image_mask is None else torch.where(image_mask[:, None, None], streams, out) + + +__all__ = [ + "EngramHash", + "EngramLayer", + "EngramTable", + "ZeroEngramTable", + "build_compressed_token_map", + "bucket_primes", + "hash_multipliers", +] diff --git a/python/freetoken/models/deepseek_v41/engram_table.py b/python/freetoken/models/deepseek_v41/engram_table.py new file mode 100644 index 0000000000..55450bbd80 --- /dev/null +++ b/python/freetoken/models/deepseek_v41/engram_table.py @@ -0,0 +1,296 @@ +"""Engram tables streamed directly from checkpoint shards. + +Each Engram layer owns a ``[384M, 256]`` fp8 table (98 GB) plus a ``[384M, 8]`` ue8m0 scale +tensor (3 GB). The tables stay where the checkpoint put them: ``EngramDiskTable`` maps the shard +tensor as one ``RowStore`` extent and reads the 24 rows a token needs per fill (io_uring, O_DIRECT, +duplicates deduplicated), while the scales are host-resident and gathered next to the values, so +the staged row is exactly the pool row format ``fp8_e8m0_b32`` and the device dequantizes it with +the same kernel the KV pools use. + +Protocol (the Qwen3.8 PLE disk table's, ``models/qwen4_exp/ple_disk.py``): the engine wraps every +dispatch in ``EngramHost.forward_host_ctx``. Eager: hash + stage + read before the launch. CUDA-graph +decode: the sampled token lives on the device under overlap scheduling, so the host reads it back, +fills the graph's pinned staging and releases a per-table flag the captured ``lookup`` WAITs on (or, +without stream memops, fills before the replay). The fill runs on a worker thread that is started +BEFORE the graph is launched: with a graph this large (40 layers) the driver's launch call itself +blocks once the GPU sits on the WAIT with the rest of the graph queued behind it, so a fill scheduled +after the launch returns would deadlock. The worker works from an immutable dispatch snapshot plus +the readback event (which the GPU completes on its own), and the dispatch context waits for it after +the launch so its failure surfaces before the step's output is consumed. +""" + +from __future__ import annotations + +import json +import os +import struct +from concurrent.futures import Future, ThreadPoolExecutor +from contextlib import contextmanager +from dataclasses import dataclass +from typing import Sequence + +import torch +from freetoken.core import Batch +from freetoken.kernel.pinned import alloc_pinned_tensor +from freetoken.kvcache.dsv4.v41_row_format import FP8_E8M0_B32 +from freetoken.utils import init_logger + +from .engram import EngramHash + +logger = init_logger(__name__) + + +@dataclass(frozen=True) +class _TensorExtent: + """Where a safetensors tensor's bytes sit in its shard: a contiguous ``[offset, offset + nbytes)``.""" + + path: str + offset: int + nbytes: int + shape: tuple[int, ...] + dtype: str # the safetensors dtype string (e.g. "F8_E4M3") + + +def _tensor_extent(folder: str, name: str) -> _TensorExtent: + """Locate ``name`` in the checkpoint directory without reading it (the table is streamed in place).""" + from freetoken.models.loader import safetensors_weight_map + + path = os.path.join(folder, safetensors_weight_map(folder)[name]) + with open(path, "rb") as fh: + n = struct.unpack(" EngramRowSource: + """Locate a layer's table in whatever safetensors files the checkpoint directory holds: the HF + shards (through the index) or the shards an FTW conversion copied next to it (through their headers).""" + from safetensors import safe_open + + weight = _tensor_extent(folder, f"layers.{layer_id}.engram.embed.weight") + if weight.dtype != "F8_E4M3": + raise ValueError(f"engram table of layer {layer_id} is {weight.dtype}, expected F8_E4M3") + rows, head_dim = weight.shape + scale = _tensor_extent(folder, f"layers.{layer_id}.engram.embed.scale") + with safe_open(scale.path, framework="pt", device="cpu") as f: + scales = f.get_tensor(f"layers.{layer_id}.engram.embed.scale").view(torch.uint8).contiguous() + if tuple(scales.shape) != (rows, head_dim // 32): + raise ValueError(f"engram scales of layer {layer_id} are {tuple(scales.shape)}, expected {(rows, head_dim // 32)}") + return EngramRowSource(weight.path, weight.offset, rows, head_dim, scales) + + +class EngramDiskTable: + """One layer's table: stages ``[T, n_cols]`` rows into pinned memory on the host, dequantizes on + the device. Implements the model's ``EngramTable`` protocol (``lookup``).""" + + def __init__( + self, source: EngramRowSource, n_cols: int, device: torch.device, *, max_graph_rows: int, max_extend_tokens: int, + use_io_uring: bool = True, + ) -> None: + from freetoken.kernel.row_store import RowStore + + self.head_dim = source.head_dim + self.n_cols = n_cols + self.num_rows = source.num_rows + self.row_bytes = FP8_E8M0_B32.row_bytes(source.head_dim) # values | scales + self.scales = source.scales + self.device = device + self.store = RowStore( + paths=[source.path], extent_file=[0], extent_base=[source.base], rows_per_extent=source.num_rows, + row_bytes=source.head_dim, row_stride=source.head_dim, use_io_uring=use_io_uring, + ) + # the graph staging is allocated up front (pinned alloc inside capture is illegal) and + # outlives any one capture so a re-capture keeps the same pointers; the eager staging + # grows to the largest prefill chunk seen (never inside a capture) + self.max_graph_rows = max_graph_rows + self._graph_pinned = alloc_pinned_tensor(max_graph_rows * n_cols * self.row_bytes, dtype=torch.uint8) + self._graph_pinned.zero_() + self._graph_dev = torch.empty(max_graph_rows * n_cols * self.row_bytes, dtype=torch.uint8, device=device) + # completes once the last graph launch that read ``_graph_pinned`` has run (recorded by the + # host coordinator after the launch: a record inside the capture would not be observable) + self.graph_consumed = torch.cuda.Event() + # the overlap scheduler stages 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 + # that last read it has run (``_eager_read[i]``, recorded right after that copy) + self._eager_pinned: list[torch.Tensor] = [] + self._eager_read = [torch.cuda.Event(), torch.cuda.Event()] + self._eager_slot = 0 + self._eager_dev = None + self._grow_eager(max_extend_tokens) + self.flag = alloc_pinned_tensor(1, dtype=torch.int64) + self.flag.zero_() + self.wait_sync = False # set by the host coordinator after its probe + + def _grow_eager(self, num_tokens: int) -> None: + nbytes = num_tokens * self.n_cols * self.row_bytes + if self._eager_pinned and self._eager_pinned[0].numel() >= nbytes: + return + for event in self._eager_read: + event.synchronize() + self._eager_pinned = [alloc_pinned_tensor(nbytes, dtype=torch.uint8) for _ in self._eager_read] + for buf in self._eager_pinned: + buf.zero_() # a warmup prefill stages nothing and reads whatever sits here + self._eager_dev = torch.empty(nbytes, dtype=torch.uint8, device=self.device) + + # ----- host ----- + def stage(self, row_ids: torch.Tensor, *, graph: bool) -> None: + """Queue ``row_ids [T, n_cols]`` (int64, host): values through the store, scales from the resident tensor.""" + n = row_ids.numel() + if graph: + self.graph_consumed.synchronize() + pinned = self._graph_pinned + assert n <= self.max_graph_rows * self.n_cols, f"graph staging holds {self.max_graph_rows} rows, batch has {n // self.n_cols}" + else: + self._grow_eager(n // self.n_cols) + self._eager_slot ^= 1 + self._eager_read[self._eager_slot].synchronize() + pinned = self._eager_pinned[self._eager_slot] + ids = row_ids.reshape(-1).to(torch.int64).contiguous() + rows = pinned[: n * self.row_bytes].view(n, self.row_bytes) + self.store.stage_rows(ids.data_ptr(), n, rows.data_ptr(), self.row_bytes) + rows[:, self.head_dim :].copy_(self.scales.index_select(0, ids)) + + def flush(self, *, signal: bool) -> None: + self.store.flush(self.flag.data_ptr() if signal else 0) + + # ----- device (EngramTable protocol) ----- + def lookup(self, num_tokens: int) -> torch.Tensor: + from freetoken.kernel.row_store import wait_reset + from freetoken.kernel.triton.dsv41.pack import unpack_rows + + stream = torch.cuda.current_stream(self.device) + capturing = torch.cuda.is_current_stream_capturing() + if capturing and self.wait_sync: + wait_reset(stream, self.flag) + pinned, dev = (self._graph_pinned, self._graph_dev) if capturing else (self._eager_pinned[self._eager_slot], self._eager_dev) + n = num_tokens * self.n_cols + nbytes = n * self.row_bytes + dev[:nbytes].copy_(pinned[:nbytes], non_blocking=True) + if not capturing: + self._eager_read[self._eager_slot].record(stream) + values = unpack_rows(dev[:nbytes].view(n, self.row_bytes), FP8_E8M0_B32, self.head_dim) + return values.view(num_tokens, self.n_cols * self.head_dim) + + +def _context(ids: torch.Tensor, position: int, width: int) -> torch.Tensor: + """The up-to-``width`` raw ids before ``position`` (fewer at the sequence start).""" + return ids[max(0, position - width) : position] + + +class EngramHost: + """The per-model coordinator: hashes each batch's token runs and fills every layer's table. + + The token range a prefill hashes is the attention metadata's encoder segments -- the one + execution plan the scheduler, the attention layers and the model already agreed on. A graph + decode's fill runs on a worker thread from an immutable dispatch snapshot (position + the + ``max_ngram - 1`` preceding ids per row) plus the readback of the sampled token; the request + objects themselves advance on the main thread right after dispatch and are never read late. + """ + + def __init__(self, hash: EngramHash, tables: Sequence[EngramDiskTable], device: torch.device, *, sync_mode: str = "auto") -> None: + from freetoken.kernel.row_store import probe_wait_sync + + self.hash = hash + self.tables = list(tables) + self.device = device + self.wait_sync = probe_wait_sync(sync_mode, device) + for t in self.tables: + t.wait_sync = self.wait_sync + self._token_readback = alloc_pinned_tensor(min(t.max_graph_rows for t in self.tables), dtype=torch.int32) + self._readback_event = torch.cuda.Event() + self._worker = ThreadPoolExecutor(max_workers=1, thread_name_prefix="engram-fill") + self._pending: Future | None = None + logger.info_rank0(f"Engram disk backend: {self.tables[0].store.io_backend()}, {'wait-sync' if self.wait_sync else 'launch-gating'}") + + def close(self) -> None: + self._worker.shutdown(wait=True) + + def __del__(self) -> None: + worker = getattr(self, "_worker", None) + if worker is not None: + worker.shutdown(wait=False) + + def fill(self, runs: Sequence[tuple[torch.Tensor, torch.Tensor]], *, graph: bool) -> None: + """``runs``: per request ``(context ids, token ids)`` in batch order; stages every layer's rows.""" + ids = torch.cat([self.hash.row_ids(tokens, ctx) for ctx, tokens in runs]) if runs else torch.empty(0, len(self.tables), self.hash.n_cols, dtype=torch.int64) + for l, table in enumerate(self.tables): + table.stage(ids[:, l], graph=graph) + table.flush(signal=graph and self.wait_sync) + + @torch.inference_mode() + def _deferred_fill(self, bs: int, histories: list[tuple[int, ...]]) -> None: + """Worker side of a graph decode: wait for the sampled tokens, hash them against the dispatch + snapshot, stage, and release the tables' flags. Reads nothing that the main thread mutates. + Inference mode is thread-local; staging buffers were created under it by the server.""" + try: + self._readback_event.synchronize() + tokens = self._token_readback[:bs].to(torch.int64) + runs = [(torch.tensor(hist, dtype=torch.int64), tokens[i : i + 1]) for i, hist in enumerate(histories)] + self.fill(runs, graph=True) + except BaseException: + from freetoken.kernel.row_store import signal + + for t in self.tables: # unblock the stream before surfacing; the step's output is discarded + signal(t.flag) + raise + + def _await_pending(self) -> None: + """Surface the deferred fill's outcome (it depends only on GPU work already enqueued, so this + cannot wait on the host).""" + if self._pending is not None: + pending, self._pending = self._pending, None + pending.result() + + def host_fill_batch(self, batch: Batch, use_graph: bool) -> None: + """Stage this batch's rows: eagerly, or -- a graph decode under flag-sync -- on the worker + thread once the sampled tokens have been read back.""" + self._await_pending() + width = self.hash.max_ngram - 1 + if batch.is_decode: + reqs = list(batch.padded_reqs) + # the dispatch snapshot: each row's position context, taken before anything advances + histories = [tuple(_context(r.input_ids, r.device_len - 1, width).tolist()) 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)) + self._pending = self._worker.submit(self._deferred_fill, bs, histories) + return + tokens = batch.input_ids.to("cpu").to(torch.int64) + runs = [(torch.tensor(hist, dtype=torch.int64), tokens[i : i + 1]) for i, hist in enumerate(histories)] + self.fill(runs, graph=use_graph) + return + # prefill: the encoder segments are the token range every consumer of this forward executes + segments = batch.attn_metadata.segments + assert segments is not None and len(segments) == len(batch.reqs) + runs = [(_context(r.input_ids, seg.start_pos, width), r.input_ids[seg.start_pos : seg.end]) for r, seg in zip(batch.reqs, segments)] + self.fill(runs, graph=False) + + @contextmanager + def forward_host_ctx(self, batch: Batch, use_graph: bool): + self.host_fill_batch(batch, use_graph) + yield + # the launch has been enqueued: wait for the deferred fill here so a failure surfaces before + # this step's output is consumed (the fill only waits on the GPU readback, never on the host) + self._await_pending() + if use_graph: + # only after this step's fill: the next fill waits on this launch, not the fill on it + stream = torch.cuda.current_stream(self.device) + for t in self.tables: + t.graph_consumed.record(stream) + + +__all__ = ["EngramDiskTable", "EngramHost", "EngramRowSource", "engram_row_source"] diff --git a/python/freetoken/models/deepseek_v41/indexer.py b/python/freetoken/models/deepseek_v41/indexer.py new file mode 100644 index 0000000000..4ab3da2bd4 --- /dev/null +++ b/python/freetoken/models/deepseek_v41/indexer.py @@ -0,0 +1,127 @@ +"""DSV41 Lightning Indexer (reference ``Indexer``): scores the compressed positions and keeps the +``index_topk`` best per query. + +A Full-mode layer owns the index keys: ``k = k_norm(wk(latent))`` from its compressor's unrotated +latent, RoPE'd at the group position, fp4-quantized (ue8m0 per 32) into the kv source's index pool. +Every Full / Reindex layer scores its own fp4-baked queries ``wq_b(qr)`` against those keys, +``sum_h relu(q_h . k) * w_h`` with ``w = weights_proj(x) * softmax_scale * n_heads**-0.5``. + +The Hierarchical Sparse Indexer (sec. 2.3.2): the candidate-source layer (the decoder's first Full +layer) additionally keeps the ``candidate_topk_blocks`` best ``candidate_block_size``-position blocks +as a candidate pool; the decoder's Reindex layers score only that pool. Logits are produced in query +blocks bounded by ``LOGITS_BUDGET_BYTES`` so a long-context prefill never materializes ``[T, 1M]``. +""" + +from __future__ import annotations + +import torch +from freetoken.core import get_global_ctx +from freetoken.kernel.triton.dsv4.fp8_linear import fp4_act_quant_inplace +from freetoken.layers import BaseOP, LinearReplicated, RMSNorm + +from .args import DeepseekV41Args, LayerRole, Mode +from .rope import apply_rotary_emb, apply_rotary_emb_decode + +LOGITS_BUDGET_BYTES = 512 << 20 + + +class Indexer(BaseOP): + def __init__(self, args: DeepseekV41Args, role: LayerRole, *, quant_config=None, prefix: str = ""): + self.role = role + self.n_heads = args.index_n_heads + self.head_dim = args.index_head_dim + self.rope_dim = args.rope_head_dim + self.topk = args.index_topk + self.weights_scale = args.index_softmax_scale * args.index_n_heads**-0.5 + self.candidate_topk_blocks = args.candidate_topk_blocks + self.candidate_block_size = args.candidate_block_size + self.wq_b = LinearReplicated(args.q_lora_rank, self.n_heads * self.head_dim, has_bias=False, quant_config=quant_config, prefix=f"{prefix}.wq_b") + self.weights_proj = LinearReplicated(args.dim, self.n_heads, has_bias=False, quant_config=quant_config, prefix=f"{prefix}.weights_proj") + if role.mode is Mode.FULL: + self.wk = LinearReplicated(args.head_dim, self.head_dim, has_bias=False, quant_config=quant_config, prefix=f"{prefix}.wk") + self.k_norm = RMSNorm(self.head_dim, args.norm_eps) + + @property + def attn(self): + return get_global_ctx().attn_backend + + # ----- keys (Full mode) ---------------------------------------------------------------- + def index_keys(self, latent: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor: + """``[G, index_head_dim]`` keys from the unrotated latents, rotated at their group positions. The + fp4 quantization is the pool's packing.""" + k = self.k_norm.forward(self.wk.forward(latent)) + apply_rotary_emb(k[..., -self.rope_dim :], freqs) + return k + + # ----- queries --------------------------------------------------------------------------- + def _queries(self, x: torch.Tensor, qr: torch.Tensor, rotate) -> tuple[torch.Tensor, torch.Tensor]: + q = self.wq_b.forward(qr).unflatten(-1, (self.n_heads, self.head_dim)) + rotate(q[..., -self.rope_dim :]) + fp4_act_quant_inplace(q, 32) + weights = self.weights_proj.forward(x) * self.weights_scale + return q, weights + + # ----- selection ------------------------------------------------------------------------- + def select_prefill( + self, x: torch.Tensor, qr: torch.Tensor, freqs: torch.Tensor, *, start_pos: int, ratio: int, locs: torch.Tensor, + candidates: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + """Per-query picks for one request's ``n`` new tokens at ``[start_pos, start_pos + n)``. + + ``locs [1, W]`` are the request's live full locs for the positions its queries may see (the + source's row of compressed position ``t`` is ``locs[t * ratio] // ratio``); ``candidates [1, n, + NC]`` restricts a Reindex layer to the pool. Returns ``(rows [1, n, topk] int32 global main rows, + -1 tail; candidates [1, n, NC] | None -- the pool this layer built as candidate source)``. + """ + n = x.shape[0] + T = locs.shape[1] // ratio + if T == 0: # a prompt shorter than one compression block: no source row exists yet + rows = torch.full((1, n, self.topk), -1, dtype=torch.int32, device=x.device) + pool = None + if self.role.is_candidate_source: + pool = torch.full((1, n, self.candidate_topk_blocks * self.candidate_block_size), -1, dtype=torch.int32, device=x.device) + return rows, pool + q, w = self._queries(x, qr, lambda t: apply_rotary_emb(t, freqs)) + live = ((start_pos + torch.arange(1, n + 1, device=x.device)) // ratio).to(torch.int32).view(1, n) + width = candidates.shape[-1] if candidates is not None else T + qb = max(1, LOGITS_BUDGET_BYTES // max(1, width * 4)) + picks, pools = [], [] + for s0 in range(0, n, qb): + s1 = min(n, s0 + qb) + cand = candidates[:, s0:s1] if candidates is not None else None + logits = self.attn.indexer_logits(q[None, s0:s1], w[None, s0:s1], self.role.kv_source, locs, ratio, live[:, s0:s1], T=T, candidates=cand) + if cand is not None: + pos = self.attn.select_topk_in_candidates(logits, cand, self.topk) + else: + pos = self.attn.select_topk(logits, live[:, s0:s1], self.topk) + if self.role.is_candidate_source: + pools.append(self.attn.select_candidate_blocks(logits, live[:, s0:s1], self.candidate_topk_blocks, self.candidate_block_size)) + picks.append(self.attn.positions_to_rows(pos, locs, ratio)) + rows = torch.cat(picks, dim=1) + pool = torch.cat(pools, dim=1) if pools else None + return rows.to(torch.int32), pool + + def select_decode( + self, x: torch.Tensor, qr: torch.Tensor, freqs: torch.Tensor, *, pos: torch.Tensor, ratio: int, locs: torch.Tensor, + T: int, candidates: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: + """Batched one-token selection. ``locs [B, W]`` is the decode snapshot (graph-static width), + ``T`` the staged compressed width, ``pos [B]``. Returns ``(rows [B, 1, topk] int32, counts [B, 1] + int32 valid picks, candidates [B, 1, NC] | None)``.""" + B = x.shape[0] + q, w = self._queries(x, qr, lambda t: apply_rotary_emb_decode(t.unsqueeze(1), freqs)) + live = ((pos + 1) // ratio).to(torch.int32).view(B, 1) + logits = self.attn.indexer_logits(q.view(B, 1, self.n_heads, self.head_dim), w.view(B, 1, -1), self.role.kv_source, locs, ratio, live, T=T, candidates=candidates) + pool = None + if candidates is not None: + picks = self.attn.select_topk_in_candidates(logits, candidates, self.topk) + else: + picks = self.attn.select_topk(logits, live, self.topk) + if self.role.is_candidate_source: + pool = self.attn.select_candidate_blocks(logits, live, self.candidate_topk_blocks, self.candidate_block_size) + rows = self.attn.positions_to_rows(picks, locs, ratio) + counts = (picks >= 0).sum(dim=-1, dtype=torch.int32) # valid picks form a prefix + return rows.to(torch.int32), counts, pool + + +__all__ = ["Indexer", "LOGITS_BUDGET_BYTES"] diff --git a/python/freetoken/models/deepseek_v41/model.py b/python/freetoken/models/deepseek_v41/model.py new file mode 100644 index 0000000000..8112022f76 --- /dev/null +++ b/python/freetoken/models/deepseek_v41/model.py @@ -0,0 +1,289 @@ +"""DeepSeek-V4.1-Flash model (engine-native port of ``inference/model.py``). + + embed -> hc_mult residual streams -> [Engram at 1, 14] -> 40 blocks (single-pass mHC around + DSV41 attention and MoE) -> collapse with the last block's pre gates -> norm -> head + +Single-pass mHC (``layers/mhc.py``): each sublayer boundary applies the PREVIOUS sublayer's post / +comb to the streams, predicts this sublayer's (pre, post, comb) from the updated streams, and mixes +the sublayer input with the pre gates the previous sublayer predicted. The chain state +``(y, residual, post, comb, pre)`` threads through the layer loop; layer 0's attention starts it +with a one-hot pre on stream 0 and no pending post. + +Causal Encoder-Decoder: the decoder's kv source (layer 20) compresses ITS INPUT -- the final encoder +hidden state -- into the global KV every decoder layer reads. Under Decoder SWA Bounded Replay the +decoder layers run on each request's last ``window_size`` prompt tokens only (their sliding window +truncated at the start of that window, see the backend; a prefix hit stops before it, so those tokens +are always prefilled); layer 20 still publishes global KV for every prompt token. +``--swa-decoder-replay exact`` runs the decoder on every token (the reference numerics). + +KV addressing is the attention backend's; pool buffers are read off the live pool per access, so a +runtime cache rebuild needs no unbind. Decode is batched and CUDA-graph safe (the DSV4 precedent). +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, List, NamedTuple + +import torch +from freetoken.attention.dsv41_sparse import DSV41AttnMetadata, PrefillSegment +from freetoken.core import get_global_ctx +from freetoken.layers import BaseOP, OPList, ParallelLMHead, RMSNorm, VocabParallelEmbedding +from freetoken.layers.mhc import mhc_fused_post_pre_single_pass, mhc_mix_input, mhc_post +from freetoken.models.blocks import BaseLLMModel, embed_input_ids + +from .args import DeepseekV41Args +from .attention import DSV41Attention, DecodeStepContext +from .engram import EngramLayer +from .moe import MoE + +if TYPE_CHECKING: + from freetoken.models.config import ModelConfig + +POST_MULT = 2.0 # post gates are 2 * sigmoid (reference hc_split_sinkhorn) + + +class Streams(NamedTuple): + """The single-pass mHC chain state between sublayers: the last sublayer output ``y`` with its + not-yet-applied ``post`` / ``comb``, the residual streams ``[T, hc, dim]`` they apply to, and + the ``pre`` gates ``[T, hc]`` the next sublayer mixes its input with.""" + + y: torch.Tensor | None + residual: torch.Tensor + post: torch.Tensor | None + comb: torch.Tensor | None + pre: torch.Tensor + + def materialize(self) -> torch.Tensor: + """The streams with the pending post applied.""" + if self.post is None: + return self.residual + return mhc_post(self.y, self.residual, self.post, self.comb) + + +class Block(BaseOP): + def __init__(self, config: ModelConfig, layer_id: int, *, prefix: str = ""): + args: DeepseekV41Args = config.dsv41_args + self.layer_id = layer_id + self.role = args.roles[layer_id] + self.norm_eps = args.norm_eps + self.hc_eps = args.hc_eps + self.sinkhorn = args.hc_sinkhorn_iters + hc, dim = args.hc_mult, args.dim + mix = (2 + hc) * hc + self.attn = DSV41Attention(args, self.role, quant_config=config.quant, prefix=f"{prefix}.attn") + self.ffn = MoE(config, layer_id, prefix=f"{prefix}.ffn") + self.attn_norm = RMSNorm(dim, args.norm_eps) + self.ffn_norm = RMSNorm(dim, args.norm_eps) + self.engram = EngramLayer(args, layer_id, quant_config=config.quant, prefix=f"{prefix}.engram") if layer_id in args.engram_layer_ids else None + self.hc_attn_fn = torch.empty(mix, hc * dim, dtype=torch.float32) + self.hc_ffn_fn = torch.empty(mix, hc * dim, dtype=torch.float32) + self.hc_attn_base = torch.empty(mix, dtype=torch.float32) + self.hc_ffn_base = torch.empty(mix, dtype=torch.float32) + self.hc_attn_scale = torch.empty(3, dtype=torch.float32) + self.hc_ffn_scale = torch.empty(3, dtype=torch.float32) + + def _boundary(self, s: Streams, fn, scale, base): + residual, post, comb, pre_next, x = mhc_fused_post_pre_single_pass( + s.y if s.y is not None else s.residual[:, 0], s.residual, s.post, s.comb, s.pre, fn, scale, base, + self.norm_eps, self.hc_eps, POST_MULT, self.sinkhorn, + ) + return residual, post, comb, pre_next, x + + def attn_input(self, s: Streams, *, image_mask: torch.Tensor | None = None): + """Apply the pending post, predict the attention gates, mix and norm the attention input. + Returns ``(residual, post, comb, pre_next, x)``.""" + if self.engram is not None: + s = Streams(None, self.engram.forward(s.materialize(), image_mask=image_mask), None, None, s.pre) + residual, post, comb, pre_next, x = self._boundary(s, self.hc_attn_fn, self.hc_attn_scale, self.hc_attn_base) + return residual, post, comb, pre_next, self.attn_norm.forward(x) + + def ffn_step(self, y: torch.Tensor, residual, post, comb, pre, *, image_mask: torch.Tensor | None = None) -> Streams: + """The FFN sublayer after the attention output ``y``: boundary, norm, MoE; returns the chain state.""" + residual, post, comb, pre_next, x = self._boundary(Streams(y, residual, post, comb, pre), self.hc_ffn_fn, self.hc_ffn_scale, self.hc_ffn_base) + return Streams(self.ffn.forward(self.ffn_norm.forward(x), image_mask=image_mask), residual, post, comb, pre_next) + + def forward_prefill(self, s: Streams, segments: List[PrefillSegment], positions: torch.Tensor, image_mask: torch.Tensor | None = None) -> Streams: + residual, post, comb, pre_next, x = self.attn_input(s, image_mask=image_mask) + if self.attn.compressor is not None: + self.attn.publish_prefill(x, segments) + y = self.attn.forward_prefill(x, segments, positions) + return self.ffn_step(y, residual, post, comb, pre_next, image_mask=image_mask) + + def forward_prefill_narrowed( + self, s: Streams, segments: List[PrefillSegment], rows: torch.Tensor, query_segments: List[PrefillSegment], + query_positions: torch.Tensor, image_mask: torch.Tensor | None = None, query_image_mask: torch.Tensor | None = None, + ) -> Streams | None: + """``forward_prefill`` with a row cut after the global KV is published: the attention input and + ``publish_prefill`` cover every token of ``segments``, the attention and the FFN only ``rows`` + (tiled by ``query_segments``) -- the CED decoder's kv source under bounded replay. None when no + query segment is left (every chunk in the batch ends before its prompt's last window).""" + residual, post, comb, pre_next, x = self.attn_input(s, image_mask=image_mask) + if self.attn.compressor is not None: + self.attn.publish_prefill(x, segments) + if not query_segments: + return None + y = self.attn.forward_prefill(x[rows], query_segments, query_positions) + return self.ffn_step(y, residual[rows], post[rows], comb[rows], pre_next[rows], image_mask=query_image_mask) + + def forward_decode(self, s: Streams, pos: torch.Tensor, rows: torch.Tensor, dctx, cmp_stage_cap: int) -> Streams: + residual, post, comb, pre_next, x = self.attn_input(s) + y = self.attn.forward_decode(x, pos, rows, dctx, cmp_stage_cap) + return self.ffn_step(y, residual, post, comb, pre_next) + + +class Transformer(BaseOP): + def __init__(self, config: ModelConfig, *, prefix: str = "model"): + args: DeepseekV41Args = config.dsv41_args + self.args = args + self.embed = VocabParallelEmbedding(num_embeddings=args.vocab_size, embedding_dim=args.dim) + self.layers = OPList([Block(config, i, prefix=f"{prefix}.layers.{i}") for i in range(args.n_layers)]) + self.norm = RMSNorm(args.dim, args.norm_eps) + + def bind(self, device: torch.device) -> None: + tables: dict = {} # one RoPE table per distinct parameter set, shared by the layers using it + for block in self.layers.op_list: + block.attn.bind(device, tables) + + def _entry(self, input_ids: torch.Tensor) -> Streams: + h = embed_input_ids(self.embed, input_ids, get_global_ctx().batch) + residual = h.unsqueeze(1).expand(-1, self.args.hc_mult, -1).contiguous() + pre = torch.zeros(h.shape[0], self.args.hc_mult, dtype=torch.float32, device=h.device) + pre[:, 0] = 1.0 + return Streams(None, residual, None, None, pre) + + def _exit(self, s: Streams) -> torch.Tensor: + return self.norm.forward(mhc_mix_input(s.materialize(), s.pre)) + + def prefill(self, input_ids: torch.Tensor, positions: torch.Tensor, md: DSV41AttnMetadata) -> torch.Tensor: + """Ragged prefill over the flat new-token stream; returns the hidden states the head reads (the + decoder stream under bounded replay, where ``md.last_indices`` also point, plus the placeholder + row ``md.decoder_pad`` asks for).""" + s = self._entry(input_ids) + dtype = s.residual.dtype + image_mask = input_ids >= self.args.vocab_size if get_global_ctx().batch.mm_embeds is not None else None + segments, dec_segments, rows = md.segments, md.decoder_segments, md.decoder_rows + decoder_start = self.args.decoder_start_layer + if rows is not None: + dec_positions = positions[rows] + dec_image_mask = None if image_mask is None else image_mask[rows] + for block in self.layers.op_list: + if rows is None or block.layer_id < decoder_start: + s = block.forward_prefill(s, segments, positions, image_mask) + elif block.layer_id == decoder_start: + s = block.forward_prefill_narrowed(s, segments, rows, dec_segments, dec_positions, image_mask, dec_image_mask) + if s is None: + break + else: + s = block.forward_prefill(s, dec_segments, dec_positions, dec_image_mask) + pad = [torch.zeros(1, self.args.dim, dtype=dtype, device=input_ids.device)] if md.decoder_pad else [] + return torch.cat([self._exit(s)] + pad) if s is not None else pad[0] + + def decode(self, input_ids: torch.Tensor, pos: torch.Tensor, md: DSV41AttnMetadata, cmp_stage_cap: int) -> torch.Tensor: + B = input_ids.shape[0] + rows = torch.arange(B, device=input_ids.device) + s = self._entry(input_ids) + # the layer-invariant step context, resolved once (a capture records it once): ring slots and + # window candidates of the shared pool -- and of the private rings when bounded replay keeps + # the decoder's window KV per request -- plus each RoPE table at the step's positions + dctx = DecodeStepContext.build([b.attn for b in self.layers.op_list], md, pos, rows, self.args.window_size, self.args.index_topk) + for block in self.layers.op_list: + s = block.forward_decode(s, pos, rows, dctx, cmp_stage_cap) + return self._exit(s) + + +class DeepseekV41ForCausalLM(BaseLLMModel): + def __init__(self, config: ModelConfig): + self._config = config + self._args: DeepseekV41Args = config.dsv41_args + self.model = Transformer(config) + self.head = ParallelLMHead(num_embeddings=self._args.vocab_size, embedding_dim=self._args.dim, quant_config=config.quant, prefix="head") + self._bound = False + if config.is_multimodal: + from .vision import DeepseekV41Vision + + self.visual = DeepseekV41Vision(config.vision_config, config.hidden_size, + num_layers=self._args.n_layers, num_experts=self._args.n_routed_experts) + + def encode(self, item): + return self.visual.forward(item) + + def place_encoder_weights(self, mode: str) -> None: + self.visual.place_weights(mode) + + def _ensure_bound(self) -> None: + if not self._bound: + self.model.bind(get_global_ctx().kv_cache.device) + if self._config.is_multimodal: + for i, block in enumerate(self.model.layers.op_list): + block.ffn.gate._bias_vl = self.visual.routing_bias[i] + self._bound = True + + def mark_for_rebind(self) -> None: + """A runtime pool rebuild changes nothing the model caches (rope tables are pool-independent), + but the engine asks; re-bind is cheap and keeps the contract with the DSV4 sibling.""" + self._bound = False + + def engram_layers(self) -> List[EngramLayer]: + return [b.engram for b in self.model.layers.op_list if b.engram is not None] + + def load_host_tables(self, engine_config) -> int: + """Attach the Engram tables (streamed from the checkpoint shards; zeros for dummy weights) + and take over ``forward_host_ctx`` so every dispatch stages its rows. Returns the pinned + host bytes the engine reserves from its pin budget.""" + import os + + layers = self.engram_layers() + if not layers: + return 0 + device = torch.device("cuda", torch.cuda.current_device()) + if getattr(engine_config, "use_dummy_weight", False): + from .engram import ZeroEngramTable + + for layer in layers: + layer.attach_table(ZeroEngramTable(layer.width, device)) + return 0 + from freetoken.utils import download_hf_weight + from transformers import AutoTokenizer + + from .engram import EngramHash, build_compressed_token_map + from .engram_table import EngramDiskTable, EngramHost, engram_row_source + + folder = download_hf_weight(engine_config.model_path) + token_map, compressed_vocab = build_compressed_token_map(AutoTokenizer.from_pretrained(folder, trust_remote_code=True)) + hash = EngramHash(self._args, token_map, compressed_vocab) + graph_rows = max(engine_config.max_running_req, engine_config.cuda_graph_max_bs or 0, 1) + tables = [ + EngramDiskTable( + engram_row_source(folder, layer.layer_id), hash.n_cols, device, + max_graph_rows=graph_rows, + max_extend_tokens=engine_config.max_extend_tokens, + use_io_uring=os.getenv("FREETOKEN_ENGRAM_IO_URING", "1") != "0", + ) + for layer in layers + ] + for layer, table in zip(layers, tables): + layer.attach_table(table) + self._engram_host = EngramHost(hash, tables, device, sync_mode=os.getenv("FREETOKEN_ENGRAM_SYNC", "auto")) + self.forward_host_ctx = self._engram_host.forward_host_ctx + return sum(t._graph_pinned.numel() + sum(b.numel() for b in t._eager_pinned) for t in tables) + + def forward(self) -> torch.Tensor: + self._ensure_bound() + batch = get_global_ctx().batch + md = batch.attn_metadata + assert isinstance(md, DSV41AttnMetadata) + if batch.is_prefill: + hidden = self.model.prefill(batch.input_ids, batch.positions.long(), md) + else: + input_ids = batch.input_ids + B = batch.padded_size + pos = batch.positions.long().view(-1)[:B] + if torch.cuda.is_current_stream_capturing(): + cmp_stage_cap = md.stage_width - 1 + else: + cmp_stage_cap = int(pos.max().item()) + hidden = self.model.decode(input_ids.view(B), pos, md, cmp_stage_cap) + return self.head.forward(hidden) + + +__all__ = ["DeepseekV41ForCausalLM", "Transformer", "Block", "Streams"] diff --git a/python/freetoken/models/deepseek_v41/moe.py b/python/freetoken/models/deepseek_v41/moe.py new file mode 100644 index 0000000000..91c62ecb01 --- /dev/null +++ b/python/freetoken/models/deepseek_v41/moe.py @@ -0,0 +1,90 @@ +"""DeepSeek-V4.1 MoE block: sqrtsoftplus router with a selection-only correction bias, one shared +clamped-SwiGLU expert, and the routed MXFP4 experts on the shared offload cache. + +Router (reference ``Gate``): ``scores = sqrt(softplus(x @ W^T))`` in fp32; the bias picks the experts, +the unbiased scores weight them (renormalized, ``+1e-20`` as in training), times ``route_scale``. +Image-span tokens select experts using the checkpoint's separate ``bias_vl``. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Tuple + +import torch +from freetoken.kernel.triton.dsv4.bf16_linear import bf16_linear_fp32 +from freetoken.kernel.triton.dsv4.router import router_select +from freetoken.kernel.triton.dsv4.swiglu import fused_swiglu +from freetoken.layers import BaseOP, LinearReplicated, make_moe_layer + +from .args import DeepseekV41Args + +if TYPE_CHECKING: + from freetoken.models.config import ModelConfig + +TopK = Tuple[torch.Tensor, torch.Tensor] + + +class Router(BaseOP): + def __init__(self, args: DeepseekV41Args): + self.topk = args.n_activated_experts + self.score_func = args.score_func + self.norm_topk_prob = args.norm_topk_prob + self.route_scale = args.route_scale + self.weight = torch.empty(args.n_routed_experts, args.dim, dtype=torch.bfloat16) + self.bias = torch.empty(args.n_routed_experts, dtype=torch.float32) + self._bias_vl = None # bound to the visual stack's per-layer routing bias after loading + + def forward(self, x: torch.Tensor, *, image_mask: torch.Tensor | None = None) -> TopK: + if image_mask is not None: + assert self._bias_vl is not None, "image routing bias was not loaded" + scores = bf16_linear_fp32(x, self.weight) + return router_select( + scores.view(-1, scores.shape[-1]), self.bias, self.topk, + score_func=self.score_func, normalize=self.norm_topk_prob, route_scale=self.route_scale, + alternate_bias=self._bias_vl, alternate_mask=image_mask, + ) + + +class SharedExpert(BaseOP): + """The dense clamped-SwiGLU expert every token passes through (fp8 block-32 linears).""" + + def __init__(self, args: DeepseekV41Args, *, quant_config=None, prefix: str = ""): + self.w1 = LinearReplicated(args.dim, args.moe_inter_dim, has_bias=False, quant_config=quant_config, prefix=f"{prefix}.w1") + self.w2 = LinearReplicated(args.moe_inter_dim, args.dim, has_bias=False, quant_config=quant_config, prefix=f"{prefix}.w2") + self.w3 = LinearReplicated(args.dim, args.moe_inter_dim, has_bias=False, quant_config=quant_config, prefix=f"{prefix}.w3") + self.limit = args.swiglu_limit + + def forward(self, x: torch.Tensor) -> torch.Tensor: + h = fused_swiglu(self.w1.forward(x), self.w3.forward(x), self.limit, x.dtype) + return self.w2.forward(h) + + +class MoE(BaseOP): + def __init__(self, config: ModelConfig, layer_id: int, *, prefix: str = ""): + args: DeepseekV41Args = config.dsv41_args + self.dim = args.dim + self.gate = Router(args) + self.shared_experts = SharedExpert(args, quant_config=config.quant, prefix=f"{prefix}.shared_experts") + # every layer is MoE: the offload cache indexes experts by layer id directly + self.experts = make_moe_layer( + config, + layer_id=layer_id, + activation="swiglu_clamp", + renormalize=args.norm_topk_prob, + limit=args.swiglu_limit, + quant_config=config.quant, + prefix=f"{prefix}.experts", + ) + + def forward(self, x: torch.Tensor, *, image_mask: torch.Tensor | None = None) -> torch.Tensor: + shape = x.shape + x = x.view(-1, self.dim) + weights, indices = self.gate.forward(x, image_mask=image_mask) + # the shared expert runs before the routed experts (same stream, so it does not overlap the + # hybrid decode's CPU work, whose submit follows it -- reordering that is follow-up work) + shared = self.shared_experts.forward(x) + routed = self.experts.routed_forward(x, weights, indices) + return (routed + shared).view(shape) + + +__all__ = ["MoE", "Router", "SharedExpert"] diff --git a/python/freetoken/models/deepseek_v41/rope.py b/python/freetoken/models/deepseek_v41/rope.py new file mode 100644 index 0000000000..6d7df9f8d3 --- /dev/null +++ b/python/freetoken/models/deepseek_v41/rope.py @@ -0,0 +1,58 @@ +"""Rotary embeddings for DeepSeek-V4.1 (verbatim reference math, ``inference/model.py``). + +Interleaved complex rotation on the last ``rope_head_dim`` channels, with YaRN when +``original_seq_len > 0``. Two regimes share two cached tables: window-only layers at the base +theta without YaRN, compressing layers at ``compress_rope_theta`` under YaRN. A compressed latent +stands for the first token of its group, so group ``j`` rotates at position ``j * ratio``; the +attention output is rotated back (``inverse=True``) because K == V share one rotated latent. +""" + +from __future__ import annotations + +import math +from functools import lru_cache + +import torch + + +def precompute_freqs_cis(dim: int, seqlen: int, original_seq_len: int, base: float, factor: float, beta_fast: int, beta_slow: int) -> torch.Tensor: + freqs = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim)) + if original_seq_len > 0: + + def corrected_dim(rotations): + return dim * math.log(original_seq_len / (rotations * 2 * math.pi)) / (2 * math.log(base)) + + low = max(math.floor(corrected_dim(beta_fast)), 0) + high = min(math.ceil(corrected_dim(beta_slow)), dim - 1) + ramp = ((torch.arange(dim // 2, dtype=torch.float32) - low) / max(high - low, 1e-3)).clamp(0, 1) + smooth = 1 - ramp + freqs = freqs / factor * (1 - smooth) + freqs * smooth + freqs = torch.outer(torch.arange(seqlen), freqs) + return torch.polar(torch.ones_like(freqs), freqs) + + +@lru_cache(2) # the model's two rope regimes; a 1M table is ~256 MB, so never one per layer +def get_freqs_cis(dim: int, seqlen: int, original_seq_len: int, base: float, factor: float, beta_fast: int, beta_slow: int, device: torch.device) -> torch.Tensor: + return precompute_freqs_cis(dim, seqlen, original_seq_len, base, factor, beta_fast, beta_slow).to(device) + + +def apply_rotary_emb(x: torch.Tensor, freqs_cis: torch.Tensor, inverse: bool = False) -> torch.Tensor: + """In-place rotation of ``x[..., :]`` (``[T, rd]`` or ``[T, H, rd]``) by per-token ``freqs_cis [T, rd//2]``.""" + xc = torch.view_as_complex(x.float().unflatten(-1, (-1, 2))) + if inverse: + freqs_cis = freqs_cis.conj() + if xc.ndim == 3: + freqs_cis = freqs_cis.unsqueeze(1) + x.copy_(torch.view_as_real(xc * freqs_cis).flatten(-2)) + return x + + +def apply_rotary_emb_decode(x: torch.Tensor, freqs_cis: torch.Tensor, inverse: bool = False) -> torch.Tensor: + """Batched decode rotation with PER-ROW freqs: ``x [B, 1, ..., rd]``, ``freqs_cis [B, rd//2]`` + (the fused dsv4 triton kernel; bit-identical to ``apply_rotary_emb``).""" + from freetoken.kernel.triton.dsv4.rope import rope_decode_inplace + + return rope_decode_inplace(x, freqs_cis, inverse) + + +__all__ = ["precompute_freqs_cis", "get_freqs_cis", "apply_rotary_emb", "apply_rotary_emb_decode"] diff --git a/python/freetoken/models/deepseek_v41/vision.py b/python/freetoken/models/deepseek_v41/vision.py new file mode 100644 index 0000000000..f5fea969ea --- /dev/null +++ b/python/freetoken/models/deepseek_v41/vision.py @@ -0,0 +1,139 @@ +"""DeepSeek ViT and 3x3 pixel-unshuffle aligner, with the shared encoder weight streamer.""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F +from freetoken.layers import BaseOP, LinearReplicated, OPList +from freetoken.models.weight_stream import BlockWeightStreamer + + +class VisionNorm(BaseOP): + def __init__(self, dim: int): + self.weight = torch.empty(dim, dtype=torch.float32) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + y = x.float() + return (y * torch.rsqrt(y.square().mean(-1, keepdim=True) + 1e-6) * self.weight).to(x.dtype) + + +class VisionAttention(BaseOP): + def __init__(self, vc): + self.n_heads = vc.num_attention_heads + self.head_dim = vc.hidden_size // self.n_heads + self.wqkv = LinearReplicated(vc.hidden_size, 3 * vc.hidden_size, has_bias=True) + self.wo = LinearReplicated(vc.hidden_size, vc.hidden_size, has_bias=True) + + def forward(self, x, cos, sin): + n = x.shape[0] + q, k, v = (t.reshape(n, self.n_heads, self.head_dim) for t in self.wqkv.forward(x).chunk(3, -1)) + def rotary(t): + a, b = t.float().chunk(2, -1) + return torch.cat((a * cos - b * sin, b * cos + a * sin), -1).to(t.dtype) + q, k = rotary(q), rotary(k) + q, k, v = (t.transpose(0, 1).unsqueeze(0) for t in (q, k, v)) + y = F.scaled_dot_product_attention(q, k, v) + return self.wo.forward(y[0].transpose(0, 1).reshape(n, -1)) + + +class VisionMLP(BaseOP): + def __init__(self, vc): + self.w1 = LinearReplicated(vc.hidden_size, 2 * vc.intermediate_size, has_bias=False) + self.w2 = LinearReplicated(vc.intermediate_size, vc.hidden_size, has_bias=False) + + def forward(self, x): + gate, up = self.w1.forward(x).chunk(2, -1) + return self.w2.forward(F.silu(gate) * up) + + +class VisionBlock(BaseOP): + def __init__(self, vc): + self.norm1 = VisionNorm(vc.hidden_size) + self.attn = VisionAttention(vc) + self.norm2 = VisionNorm(vc.hidden_size) + self.mlp = VisionMLP(vc) + + def forward(self, x, cos, sin): + x = x + self.attn.forward(self.norm1.forward(x), cos, sin) + return x + self.mlp.forward(self.norm2.forward(x)) + + +class PatchEmbed(BaseOP): + def __init__(self, vc): + self.proj = LinearReplicated(3 * vc.patch_size**2, vc.hidden_size, has_bias=True) + + def forward(self, x): + return self.proj.forward(x.flatten(1)) + + +class VisionTransformer(BaseOP): + def __init__(self, vc): + self.patch_embed = PatchEmbed(vc) + self.blocks = OPList([VisionBlock(vc) for _ in range(vc.num_hidden_layers)]) + self.norm = VisionNorm(vc.hidden_size) + self._vc = vc + self._streamer = None + + def place_weights(self, mode: str) -> None: + if mode == "host" and self._streamer is None: + self._streamer = BlockWeightStreamer(self.blocks.op_list, self.patch_embed.proj.weight.device) + elif mode == "gpu" and self._streamer is not None: + self._streamer.unstream() + self._streamer = None + elif mode not in ("gpu", "host"): + raise ValueError(f"unknown vision weight placement {mode!r}") + + def forward(self, patches, h: int, w: int): + x = self.patch_embed.forward(patches) + dim = self._vc.hidden_size // self._vc.num_attention_heads // 2 + inv = 1.0 / (self._vc.rope_theta ** (torch.arange(0, dim, 2, dtype=torch.float32, device=x.device) / dim)) + rows = torch.arange(h, device=x.device)[:, None].expand(h, w) + cols = torch.arange(w, device=x.device)[None, :].expand(h, w) + angles = (torch.stack((rows, cols), -1).reshape(-1, 2, 1).float() * inv).flatten(1) + cos, sin = angles.cos().unsqueeze(1), angles.sin().unsqueeze(1) + blocks = enumerate(self.blocks.op_list) if self._streamer is None else self._streamer.blocks(self.blocks.op_list) + for _, block in blocks: + x = block.forward(x, cos, sin) + return self.norm.forward(x) + + +class VisionAligner(BaseOP): + def __init__(self, vc, dim: int): + self.ratio = vc.downsample_ratio + self.w1 = LinearReplicated(vc.hidden_size * self.ratio**2, dim, has_bias=True) + self.w2 = LinearReplicated(dim, dim, has_bias=True) + + def forward(self, x, h: int, w: int): + r = self.ratio + x = x.view(h, w, -1).permute(2, 0, 1) + x = F.pad(x, (0, -w % r, 0, -h % r)) + x = F.unfold(x.unsqueeze(0), r, stride=r).squeeze(0).T + return self.w2.forward(F.gelu(self.w1.forward(x))) + + +class DeepseekV41Vision(BaseOP): + def __init__(self, vc, dim: int, *, num_layers: int, num_experts: int): + self.vision = VisionTransformer(vc) + self.aligner = VisionAligner(vc, dim) + self.image_start = torch.empty(dim) + self.image_end = torch.empty(dim) + self.image_newline = torch.empty(dim) + # Keep all image-only parameters under visual, including FTW/text-only filtering. + self.routing_bias = torch.empty(num_layers, num_experts, dtype=torch.float32) + + def place_weights(self, mode: str) -> None: + self.vision.place_weights(mode) + + @torch.inference_mode() + def forward(self, item): + _, h, w = item.grid_thw + patches = item.feature.to(device=self.image_start.device, dtype=self.image_start.dtype) + x = self.aligner.forward(self.vision.forward(patches, h, w), h, w) + r = self.aligner.ratio + nh, nw = (h + r - 1) // r, (w + r - 1) // r + out = x.new_empty(nh * (nw + 1) + 2, x.shape[-1]) + out[0], out[-1] = self.image_start, self.image_end + rows = out[1:-1].view(nh, nw + 1, -1) + rows[:, :nw] = x.view(nh, nw, -1) + rows[:, nw] = self.image_newline + return out diff --git a/python/freetoken/models/deepseek_v41/weight.py b/python/freetoken/models/deepseek_v41/weight.py new file mode 100644 index 0000000000..1ccd23eaef --- /dev/null +++ b/python/freetoken/models/deepseek_v41/weight.py @@ -0,0 +1,242 @@ +"""Weight loading for DeepSeek-V4.1-Flash. + +* :func:`iter_weights` streams the resident (non-expert) tensors keyed by the model's attribute + paths: checkpoint ``layers.N.*`` -> ``model.layers.N.*``, ``embed`` / ``norm`` -> ``model.*``, + ``head`` -> ``head``. fp8 linears keep their e4m3 payload and e8m0 ``scale`` (declared by the quant + method as ``weight_scale_inv``); ``wo_a`` keeps its e4m3 payload and ``wo_a_scale`` too (the layer + dequantizes in-kernel for the reference's bf16 grouped einsum). The engine casts everything else to + the declared parameter dtype (fp32 compressor weights, hc gates, engram gate weights). +* :func:`iter_expert_pieces` streams the routed MXFP4 experts (e2m1 pairs + e8m0 per-32 scales) per + expert for the offload banks. + +Dropped: the DSpark ``mtp.*`` stage; vision weights and ``gate.bias_vl`` when serving text-only. +The Engram tables (``layers.{1,14}.engram.embed.*``) are +not state-dict tensors: the disk row store maps them in place (``engram_table.py``); an FTW conversion +carries them as side files (:func:`ftw_side_files`). +""" + +from __future__ import annotations + +from contextlib import ExitStack +import json +import os +import re + +import safetensors +import torch +from freetoken.distributed import get_tp_info +from freetoken.layers.quantization import QuantKind +from freetoken.models.loader import drop_page_cache +from tqdm import tqdm + +from .args import DeepseekV41Args, Mode + + +class _ShardReader: + def __init__(self, folder: str, weight_map: dict, device): + self._folder = folder + self._weight_map = weight_map + self._device = str(device) + self._handles: dict[str, object] = {} + self._stack = ExitStack() + + def has(self, name: str) -> bool: + return name in self._weight_map + + def get(self, name: str) -> torch.Tensor: + shard = self._weight_map[name] + handle = self._handles.get(shard) + if handle is None: + handle = self._stack.enter_context(safetensors.safe_open(os.path.join(self._folder, shard), framework="pt", device=self._device)) + self._handles[shard] = handle + return handle.get_tensor(name) + + def close(self) -> None: + try: + self._stack.close() + finally: + for shard in self._handles: + drop_page_cache(os.path.join(self._folder, shard)) + self._handles.clear() + + +def weight_map(model_path: str) -> dict: + with open(os.path.join(model_path, "model.safetensors.index.json")) as f: + return json.load(f)["weight_map"] + + +def load_args(model_path: str) -> DeepseekV41Args: + from freetoken.utils.hf import cached_load_hf_config + + return DeepseekV41Args.from_hf(cached_load_hf_config(model_path)) + + +def dequant_fp8_block(weight: torch.Tensor, scale: torch.Tensor) -> torch.Tensor: + """Square block-scaled e4m3 -> bf16; ``scale`` holds e8m0 codes (``2^(code-127)``), the block + edge follows from the shapes.""" + n, k = weight.shape + block = n // scale.shape[0] + assert (n // block, k // block) == tuple(scale.shape), (weight.shape, scale.shape) + s = torch.exp2(scale.view(torch.uint8).to(torch.float32) - 127.0) + s = s.repeat_interleave(block, dim=0).repeat_interleave(block, dim=1) + return (weight.to(torch.float32) * s).to(torch.bfloat16) + + +def iter_weights(model_path: str, device, *, include_moe_experts: bool = True, include_non_moe: bool = True, include_vision: bool = True): + args = load_args(model_path) + reader = _ShardReader(model_path, weight_map(model_path), device) + get = reader.get + def linear(src: str, dst: str): + yield f"{dst}.weight", get(f"{src}.weight") + if reader.has(f"{src}.scale"): + yield f"{dst}.weight_scale_inv", get(f"{src}.scale") + + try: + if include_moe_experts: + # Resident experts support unquantized checkpoints; FP4 uses the offload cache. + if reader.has("layers.0.ffn.experts.0.w1.scale"): + raise ValueError( + "DeepSeek-V4.1's fp4 routed experts are served from the offload cache; run with " + "--moe-strategy offload (include_moe_experts must be False)." + ) + for L in range(args.n_layers): + p = f"layers.{L}.ffn.experts" + gate_up = torch.stack([torch.cat([get(f"{p}.{e}.w1.weight"), get(f"{p}.{e}.w3.weight")]) for e in range(args.n_routed_experts)]) + down = torch.stack([get(f"{p}.{e}.w2.weight") for e in range(args.n_routed_experts)]) + yield f"model.layers.{L}.ffn.experts.gate_up_proj", gate_up + yield f"model.layers.{L}.ffn.experts.down_proj", down + if not include_non_moe: + return + if include_vision and reader.has("vision.patch_embed.proj.weight"): + yield from iter_vision_weights(model_path, device) + yield "model.embed.weight", get("embed.weight") + yield "model.norm.weight", get("norm.weight") + yield "head.weight", get("head.weight") + for L in range(args.n_layers): + role = args.roles[L] + a, m = f"layers.{L}.attn", f"model.layers.{L}.attn" + yield from linear(f"{a}.wq_a", f"{m}.wq_a") + yield f"{m}.q_norm.weight", get(f"{a}.q_norm.weight") + yield from linear(f"{a}.wq_b", f"{m}.wq_b") + yield from linear(f"{a}.wkv", f"{m}.wkv") + yield f"{m}.kv_norm.weight", get(f"{a}.kv_norm.weight") + yield f"{m}.wo_a", get(f"{a}.wo_a.weight") + if reader.has(f"{a}.wo_a.scale"): # the fp8 payload + scale; the layer dequantizes in-kernel + yield f"{m}.wo_a_scale", get(f"{a}.wo_a.scale") + yield from linear(f"{a}.wo_b", f"{m}.wo_b") + yield f"{m}.attn_sink", get(f"{a}.attn_sink") + if role.mode is Mode.FULL: + c = f"{a}.compressor" + yield f"{m}.compressor.norm.weight", get(f"{c}.norm.weight") + if role.ratio > 1: + yield f"{m}.compressor.wkv", get(f"{c}.wkv.weight") + yield f"{m}.compressor.wgate", get(f"{c}.wgate.weight") + else: + yield from linear(f"{c}.wkv", f"{m}.compressor.wkv") + yield from linear(f"{a}.indexer.wk", f"{m}.indexer.wk") + yield f"{m}.indexer.k_norm.weight", get(f"{a}.indexer.k_norm.weight") + if role.has_indexer: + yield from linear(f"{a}.indexer.wq_b", f"{m}.indexer.wq_b") + yield from linear(f"{a}.indexer.weights_proj", f"{m}.indexer.weights_proj") + + p = f"layers.{L}" + yield f"model.{p}.attn_norm.weight", get(f"{p}.attn_norm.weight") + yield f"model.{p}.ffn_norm.weight", get(f"{p}.ffn_norm.weight") + yield f"model.{p}.ffn.gate.weight", get(f"{p}.ffn.gate.weight") + yield f"model.{p}.ffn.gate.bias", get(f"{p}.ffn.gate.bias") + for proj in ("w1", "w2", "w3"): + yield from linear(f"{p}.ffn.shared_experts.{proj}", f"model.{p}.ffn.shared_experts.{proj}") + for nm in ("hc_attn_fn", "hc_ffn_fn", "hc_attn_base", "hc_ffn_base", "hc_attn_scale", "hc_ffn_scale"): + yield f"model.{p}.{nm}", get(f"{p}.{nm}") + if L in args.engram_layer_ids: + e = f"{p}.engram" + yield from linear(f"{e}.wkv", f"model.{e}.wkv") + yield f"model.{e}.q_weight", get(f"{e}.q_weight") + yield f"model.{e}.k_weight", get(f"{e}.k_weight") + finally: + reader.close() + + +def iter_vision_weights(model_path: str, device): + names = weight_map(model_path) + reader = _ShardReader(model_path, names, device) + try: + for name in names: + if name.startswith(("vision.", "aligner.")) or name in ("image_start", "image_end", "image_newline"): + yield f"visual.{name}", reader.get(name) + args = load_args(model_path) + yield "visual.routing_bias", torch.stack([reader.get(f"layers.{i}.ffn.gate.bias_vl") for i in range(args.n_layers)]) + finally: + reader.close() + + +# ----- FTW side files: the Engram tables ---------------------------------------------------------- +def engram_table_names(layer_id: int) -> tuple[str, str]: + return f"layers.{layer_id}.engram.embed.weight", f"layers.{layer_id}.engram.embed.scale" + + +def ftw_side_files(model_path: str, out_dir: str) -> list[str]: + """Copy the checkpoint shards holding the Engram tables (fp8 rows + ue8m0 scales) next to an FTW + checkpoint. The tables are never FTW entries: ``load_host_tables`` maps them in place from whatever + safetensors file holds them (``engram_table.engram_row_source`` reads the headers), and the FTW + loader never reads ``.safetensors`` as weights, so the shard's other tensors ride along unused.""" + import shutil + + from freetoken.models.loader import safetensors_weight_map + from freetoken.utils import download_hf_weight + + folder = download_hf_weight(model_path) + args = load_args(folder) + shard_of = safetensors_weight_map(folder) + shards = sorted({shard_of[n] for layer_id in args.engram_layer_ids if layer_id < args.n_layers for n in engram_table_names(layer_id)}) + for shard in shards: + shutil.copyfile(os.path.join(folder, shard), os.path.join(out_dir, shard)) + return shards + + +# ----- routed MXFP4 expert pieces --------------------------------------------------------------- +_EXPERT_RE = re.compile(r"^layers\.(?P\d+)\.ffn\.experts\.(?P\d+)\.(?Pw1|w2|w3)\.(?Pweight|scale)$") +_PROJ_ROLE = {"w1": "gate", "w3": "up", "w2": "down"} +_KIND_SUFFIX = {"weight": "", "scale": "_scale"} + + +def iter_expert_pieces(model_path: str, config, kind: QuantKind, *, parallel: bool | None = False, workers: int = 8, chunk: int = 8 << 20): + """One piece per expert: ``{gate, up, down}`` e2m1 pairs and their e8m0 ``_scale`` companions. + The DSpark stage's experts (``mtp.*``) never match.""" + if kind is not QuantKind.MXFP4: + return None + if get_tp_info().size > 1: + raise NotImplementedError("DeepSeek-V4.1 expert banks support TP=1 only") + from freetoken.models.weight import iter_expert_tensors_parallel + from freetoken.moe.expert_pieces import per_expert_pieces + + args = load_args(model_path) + L, E = args.n_layers, args.n_routed_experts + + def locate(raw_name: str): + m = _EXPERT_RE.match(raw_name) + if m is None or int(m["layer"]) >= L: + return None + return int(m["layer"]), int(m["expert"]), _PROJ_ROLE[m["proj"]] + _KIND_SUFFIX[m["kind"]] + + if parallel: + tensors = iter_expert_tensors_parallel(model_path, lambda n: locate(n) is not None, workers=workers, chunk=chunk) + return per_expert_pieces(tensors, locate, tensors_per_expert=6) + + def _serial(): + reader = _ShardReader(model_path, weight_map(model_path), torch.device("cpu")) + try: + for li in tqdm(range(L), desc="Loading DeepSeek-V4.1 experts (serial)", disable=not get_tp_info().is_primary()): + for e in range(E): + base = f"layers.{li}.ffn.experts.{e}" + for proj in ("w1", "w3", "w2"): + for kind_ in ("weight", "scale"): + name = f"{base}.{proj}.{kind_}" + yield name, reader.get(name) + finally: + reader.close() + + return per_expert_pieces(_serial(), locate, tensors_per_expert=6) + + +__all__ = ["iter_weights", "iter_vision_weights", "iter_expert_pieces", "ftw_side_files", "load_args", "weight_map", "dequant_fp8_block"] 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/qwen3_5_moe/weight.py b/python/freetoken/models/qwen3_5_moe/weight.py index 1bac8227b0..de5b9e8983 100644 --- a/python/freetoken/models/qwen3_5_moe/weight.py +++ b/python/freetoken/models/qwen3_5_moe/weight.py @@ -1,6 +1,6 @@ """Qwen3.5 / 3.6 / 3.8 checkpoint reader. -The dense pass reads every Linear module under the scheme the checkpoint's QuantConfig gives it, the same answer the model built its buffers from, so bf16, block-fp8, ModelOpt and llm-compressor exports in any mix all land as the model's state dict. Routed experts are read by the offload cache (``nvfp4_expert_spec`` / ``iter_expert_pieces``); only bf16 stacked experts and resident block-fp8 experts come from here. +The dense pass reads every Linear module under the scheme the checkpoint's QuantConfig gives it, the same answer the model built its buffers from, so bf16, block-fp8, ModelOpt and llm-compressor exports in any mix all land as the model's state dict. Routed experts are read by the expert-bank loader (``nvfp4_expert_spec`` / ``iter_expert_pieces``); only the bf16 stacked experts it reads come from here. """ from __future__ import annotations @@ -241,9 +241,9 @@ def iter_weights( include_non_moe: bool, include_vision: bool = True, ) -> Iterator[tuple[str, torch.Tensor]]: - """Yield the dense weights fused to the model's buffers, and the routed experts only where a resident path takes them from here: bf16 stacked experts as stored, block-fp8 experts restacked per layer. + """Yield the dense weights fused to the model's buffers, and with ``include_moe_experts`` the bf16 stacked experts as stored. - Per-expert NVFP4 experts always come from the offload cache's expert reader. + Block-fp8 and NVFP4 experts always come from the expert-bank reader. """ if get_tp_info().size > 1: raise NotImplementedError("qwen3_5_moe weight loading supports TP=1 only") @@ -253,8 +253,6 @@ def iter_weights( if include_non_moe or stacked: reader = _DenseReader(get_quant_config(), get_model_spec(hf_config.architectures[0])) if include_non_moe else None yield from _iter_shards(model_path, device, reader, stacked=stacked, include_vision=include_vision) - if include_moe_experts and config.is_moe and config.expert_quant == "fp8_block": - yield from _resident_fp8_experts(model_path, config) def _iter_shards(model_path: str, device: torch.device, reader: _DenseReader | None, *, stacked: bool, include_vision: bool): @@ -327,7 +325,7 @@ def _is_expert(raw_name: str) -> bool: # ====================================================================================== -# Block-FP8 routed experts (Qwen3.5-35B-A3B-FP8): offload expert pieces and resident stacks. +# Block-FP8 routed experts (Qwen3.5-35B-A3B-FP8): expert pieces. # ====================================================================================== # Routed-expert checkpoint key (per-expert, un-fused). ``mtp.layers...`` is excluded by the @@ -338,35 +336,6 @@ def _is_expert(raw_name: str) -> bool: ) -def _resident_fp8_experts(model_path, config): - from freetoken.kernel.triton.fp8_block_linear import FP8 - - B = 128 - L, E, H, I, dense = _moe_dims(config) - shapes = { - "gate_up_proj": ((E, 2 * I, H), FP8), - "gate_up_scale_inv": ((E, 2 * I // B, H // B), torch.bfloat16), - "down_proj": ((E, H, I), FP8), - "down_scale_inv": ((E, H // B, I // B), torch.bfloat16), - } - layers: dict[int, dict[str, torch.Tensor]] = {} - placed = [0] * L - for li, e0, e1, piece in iter_expert_pieces(model_path, config, QuantKind.FP8_BLOCK, parallel=None): - stack = layers.setdefault(li, {n: torch.empty(shape, dtype=dt) for n, (shape, dt) in shapes.items()}) - stack["gate_up_proj"][e0:e1, :I] = piece["gate"] - stack["gate_up_proj"][e0:e1, I:] = piece["up"] - stack["gate_up_scale_inv"][e0:e1, : I // B] = piece["gate_scale"] - stack["gate_up_scale_inv"][e0:e1, I // B :] = piece["up_scale"] - stack["down_proj"][e0:e1] = piece["down"] - stack["down_scale_inv"][e0:e1] = piece["down_scale"] - placed[li] += e1 - e0 - if placed[li] == E: - pre = f"model.layers.{dense + li}.mlp.experts" - for name, tensor in layers.pop(li).items(): - yield f"{pre}.{name}", tensor - assert not layers, f"incomplete resident fp8 experts for layers {sorted(layers)}" - - def _moe_dims(model_config): L = model_config.num_moe_layers return ( 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/model.py b/python/freetoken/models/qwen4_exp/model.py index 258badac42..7f4f826807 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, ) @@ -99,6 +101,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) @@ -131,14 +134,23 @@ 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..91c8c1bd17 100644 --- a/python/freetoken/models/qwen4_exp/weight.py +++ b/python/freetoken/models/qwen4_exp/weight.py @@ -24,17 +24,19 @@ 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.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 +logger = init_logger(__name__) + # 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. @@ -188,6 +190,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 +306,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 +356,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 +417,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. diff --git a/python/freetoken/models/register.py b/python/freetoken/models/register.py index 2a4ced10af..fe150ddc29 100644 --- a/python/freetoken/models/register.py +++ b/python/freetoken/models/register.py @@ -156,6 +156,20 @@ class ModelSpec: # the head, the KV compressors and the indexer's scorer ship bf16; the fp8 config has no modules_to_not_convert unquantized_modules=("head", "*.compressor.wkv", "*.compressor.wgate", "*.indexer.weights_proj"), ), + # DeepSeek-V4.1-Flash (model_type deepseek_v41): multimodal wrapper config (text tower in + # text_config). The checkpoint has no ``model.`` root; the compressors, the indexer's key projection / + # scorer, the head and the vision tower + aligner ship bf16 (the fp8 config has no modules_to_not_convert). + "DeepseekV41ForCausalLM": ModelSpec( + "freetoken.models.deepseek_v41", + "DeepseekV41ForCausalLM", + mm_processor="freetoken.mm.processors.deepseek_v41:DeepseekV41MMProcessor", + encoders=(EncoderSpec("vision", "vision_config", ("image",)),), + checkpoint_roots=(("model.layers", "layers"), ("model.embed", "embed"), ("model.norm", "norm")), + packed_modules_mapping=_EXPERTS_W123_PACKED, + unquantized_modules=( + "head", "*.compressor.wkv", "*.compressor.wgate", "*.indexer.wk", "*.indexer.weights_proj", "vision.*", "aligner.*", + ), + ), "Qwen3_5MoeForConditionalGeneration": ModelSpec( "freetoken.models.qwen3_5_moe", "Qwen3_5MoeForConditionalGeneration", 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 83b3d19665..e9d6ffe54e 100644 --- a/python/freetoken/moe/benchbw.py +++ b/python/freetoken/moe/benchbw.py @@ -108,6 +108,7 @@ class Workload: activation: str = "silu" swiglu_alpha: float = 1.702 swiglu_limit: float | None = None + act_block: int | None = None # ds_fp4 W4A8 activation quant block # Preset workloads. Dims from the model configs / benchmarks/bench_offload_cache_copy.py. @@ -120,7 +121,8 @@ class Workload: activation="gpt_oss_swiglu", swiglu_limit=7.0), "gpt-oss-20b": Workload("gpt-oss-20b", 2880, 2880, 32, 4, ("mxfp4_triton",), activation="gpt_oss_swiglu", swiglu_limit=7.0), - "dsv4": Workload("dsv4", 4096, 2048, 256, 6, ("ds_fp4",), swiglu_limit=7.0), + "dsv4": Workload("dsv4", 4096, 2048, 256, 6, ("ds_fp4",), swiglu_limit=7.0, act_block=128), + "dsv4.1-flash": Workload("dsv4.1-flash", 5120, 2304, 384, 6, ("ds_fp4",), swiglu_limit=10.0, act_block=32), "glm4.7-nvfp4": Workload("glm4.7-nvfp4", 5120, 1536, 160, 8, ("nvfp4",)), "glm5.3-flash-nvfp4": Workload( "glm5.3-flash-nvfp4", 4096, 2048, 288, 8, ("nvfp4",), @@ -142,7 +144,7 @@ class Workload: "fp8_block": Workload("dtype:fp8", 2048, 512, 128, 8, ("fp8_block",)), "mxfp4_triton": Workload("dtype:mxfp4", 2880, 2880, 128, 4, ("mxfp4_triton",), activation="gpt_oss_swiglu", swiglu_limit=7.0), - "ds_fp4": Workload("dtype:ds_fp4", 4096, 2048, 128, 6, ("ds_fp4",), swiglu_limit=7.0), + "ds_fp4": Workload("dtype:ds_fp4", 4096, 2048, 128, 6, ("ds_fp4",), swiglu_limit=7.0, act_block=128), } @@ -460,7 +462,7 @@ def _build_cpu_moe_executor(fmt: str, wl: Workload, banks: dict, num_threads: in cache, top_k=wl.top_k, activation=wl.activation, apply_router_weight_on_input=False, num_threads=num_threads, max_tokens=1, device=torch.device("cuda"), swiglu_alpha=wl.swiglu_alpha, - swiglu_limit=wl.swiglu_limit, + swiglu_limit=wl.swiglu_limit, act_block=wl.act_block, ) diff --git a/python/freetoken/moe/cpu_executor.py b/python/freetoken/moe/cpu_executor.py index e2e48d1461..c4809075c8 100644 --- a/python/freetoken/moe/cpu_executor.py +++ b/python/freetoken/moe/cpu_executor.py @@ -53,6 +53,8 @@ # 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" # Activation ids must match ActKind in csrc/cpu_moe/cpu_moe_ext.cpp. Id 3 is the # clamped (up + 1) swiglu: "swigluoai" runs it in the generic GEMV epilogue, @@ -116,6 +118,17 @@ def physical_core_cpus() -> list[int]: return reps or allowed or [0] +def _spin_wait_ok(core_ids: list[int], coord_core: int) -> bool: + """Spin only if each pool thread owns a distinct CPU and two CPUs stay free for the main and CUDA callback threads.""" + if not _SPIN_WAIT or not hasattr(os, "sched_getaffinity"): + return False + pool = list(core_ids) + ([coord_core] if coord_core >= 0 else []) + if len(set(pool)) != len(pool): + return False + spare = set(os.sched_getaffinity(0)) - set(pool) + return len(spare) >= 2 + + def resolve_threads_and_affinity(requested: int) -> tuple[int, list[int]]: """Return (num_threads, core_ids) for the worker pool. @@ -158,6 +171,7 @@ def __init__( swiglu_alpha: float = 1.702, swiglu_limit: float | None = None, fmt: str | None = None, + act_block: int | None = None, ) -> None: from freetoken.kernel import _cpu_moe from freetoken.moe.legacy_format import canonical_role @@ -246,6 +260,8 @@ def __init__( self.num_threads = nthreads self.core_ids = core_ids self.isa = self._ext.isa_name() + self.spin_wait = _spin_wait_ok(core_ids, coord_core) + self._ext.set_spin_wait(self.spin_wait) spare = len(physical_core_cpus()) - nthreads - (1 if coord_core >= 0 else 0) - 1 clamp = max(1, min(torch.get_num_threads(), spare)) @@ -308,6 +324,11 @@ def __init__( # 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( @@ -319,7 +340,8 @@ def __init__( 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"top_k={self.top_k} act={activation} max_tokens={self.max_tokens} " + f"spin_wait={self.spin_wait}" ) def _make_table(self, layers: list[torch.Tensor]) -> torch.Tensor: @@ -466,7 +488,7 @@ def _resolve_dsfp4_banks(self, banks: dict) -> tuple[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 - activations (block 128) to match DSV4's W4A8 reference, hence the %128 dims.""" + activations (the checkpoint's fp8 block) to match the DeepSeek W4A8 reference.""" gup, gus = banks["gate_up"], banks["gate_up_scale"] dnp, dns = banks["down"], banks["down_scale"] assert gup[0].dtype == torch.uint8 and dnp[0].dtype == torch.uint8, (gup[0].dtype, dnp[0].dtype) @@ -474,7 +496,7 @@ def _resolve_dsfp4_banks(self, banks: dict) -> tuple[dict, tuple[int, int]]: I = int(gup[0].shape[1] // 2) H = int(gup[0].shape[2] * 2) assert gup[0].shape[1] == 2 * I - assert H % 128 == 0 and I % 128 == 0, (H, I) # FP8 activation round-trip block=128 + assert H % 32 == 0 and I % 32 == 0, (H, I) # e8m0 scales per 32 along K 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) @@ -564,7 +586,7 @@ def decode_submit( # pre-quantized activations and skips its serial scalar pass. from freetoken.kernel.triton.dsv4.fp8_linear import act_quant_fp8_roundtrip - hidden_states = act_quant_fp8_roundtrip(hidden_states, block=128) + hidden_states = act_quant_fp8_roundtrip(hidden_states, block=self._act_block) # D2H: ship this step's activations + routing to pinned host memory. io["x"].copy_(hidden_states, non_blocking=True) diff --git a/python/freetoken/moe/expert_banks.py b/python/freetoken/moe/expert_banks.py index e3a768b6d0..7500174ae4 100644 --- a/python/freetoken/moe/expert_banks.py +++ b/python/freetoken/moe/expert_banks.py @@ -75,6 +75,7 @@ def build_expert_banks( device: torch.device, layer_sink=None, dummy: bool = False, + resident: bool = False, ) -> ExpertBanks: """Fill host banks in the kernel's layout from a stream of expert pieces. @@ -83,6 +84,7 @@ def build_expert_banks( complete once its ``num_experts`` rows have arrived: with ``layer_sink=None`` its banks are pinned in the background, otherwise the sink receives them (converter). ``dummy`` skips the pieces and fills the banks with finite random contents. + ``resident`` packs into banks on ``device`` instead of host banks (the resident 'fused' experts). """ from freetoken.moe.host_banks import LayerCompletionTracker, PinPipeline, pin_banks from freetoken.moe.legacy_format import legacy_format_for @@ -91,8 +93,13 @@ def build_expert_banks( layout = method.layout() E = method.cfg.num_experts specs = {role: ((E, *spec.shape), spec.dtype) for role, spec in layout.items() if not spec.resident} - hb = alloc_layer_banks(specs, num_layers) - banks = {role: [b.tensor for b in hb[role]] for role in specs} + if resident: + assert layer_sink is None + hb = None + banks = {role: [torch.empty(shape, dtype=dtype, device=device) for _ in range(num_layers)] for role, (shape, dtype) in specs.items()} + else: + hb = alloc_layer_banks(specs, num_layers) + banks = {role: [b.tensor for b in hb[role]] for role in specs} alphas = { role: torch.empty(num_layers * E, dtype=spec.dtype, device=device) for role, spec in layout.items() if spec.resident @@ -104,7 +111,7 @@ def build_expert_banks( _dummy_fill(role, tensor) for alpha in alphas.values(): alpha.fill_(1.0) - if torch.cuda.is_available(): + if hb is not None and torch.cuda.is_available(): pin_banks(hb) return ExpertBanks( legacy_format_for(method.kind, kernel.name), banks, @@ -136,7 +143,7 @@ def _fill(sink) -> None: if layer_sink is not None: _fill(layer_sink) - elif torch.cuda.is_available(): + elif torch.cuda.is_available() and not resident: with PinPipeline() as pins: _fill(pins) else: @@ -192,16 +199,16 @@ def _legacy_expert_banks(model_path, model_config, device, dtype, dummy, paralle ) -def _method_expert_banks(model_path, model_config, method, device, dummy, parallel, workers, chunk, layer_sink=None) -> ExpertBanks: +def _method_expert_banks(model_path, model_config, method, device, dummy, parallel, workers, chunk, layer_sink=None, resident=False) -> ExpertBanks: from freetoken.moe.expert_pieces import iter_expert_pieces num_layers = model_config.num_moe_layers if dummy: - return build_expert_banks(method, num_layers, None, device=device, dummy=True) + return build_expert_banks(method, num_layers, None, device=device, dummy=True, resident=resident) pieces = iter_expert_pieces( model_path, model_config, method.kind, parallel=parallel, workers=workers, chunk=chunk ) - return build_expert_banks(method, num_layers, pieces, device=device, layer_sink=layer_sink) + return build_expert_banks(method, num_layers, pieces, device=device, layer_sink=layer_sink, resident=resident) def _host_ram_fits_parallel(model_path: str) -> bool: @@ -286,6 +293,7 @@ def load_expert_banks( decode_target: str = "gpu", layer_sink=None, layer_residency: list[str] | None = None, + resident: bool = False, ) -> ExpertBanks: """Load (or fabricate, with ``dummy=True``) the expert banks. Two paths, both returning the same normalized ``ExpertBanks`` and both pinning after fill: @@ -310,14 +318,21 @@ def load_expert_banks( ``layer_residency``: per-layer ``HostResidency`` labels applied at settle time -- explicitly on the FTW fast path, ambiently (``requested_residency``) in the slow-path providers. Applied labels are echoed on ``ExpertBanks.layer_residency``; a loader that settles some other way leaves it ``None`` (CPU-layer decode still works on pinned banks, it just saves no pin quota). + + ``resident`` (with ``method``): the banks land on ``device`` for the resident 'fused' experts; no host bank outlives the load. """ - from freetoken.checkpoint.ftw import is_ftw_checkpoint, load_ftw_banks + from freetoken.checkpoint.ftw import is_ftw_checkpoint, load_ftw_banks, load_ftw_banks_to_device + if resident: + assert method is not None and layer_sink is None and layer_residency is None if model_path and is_ftw_checkpoint(model_path) and not dummy: - banks = load_ftw_banks( - model_path, num_layers=model_config.num_moe_layers, workers=workers, chunk=chunk, - layer_residency=layer_residency, - ) + if resident: + banks = load_ftw_banks_to_device(model_path, num_layers=model_config.num_moe_layers, device=device, workers=workers, chunk=chunk) + else: + banks = load_ftw_banks( + model_path, num_layers=model_config.num_moe_layers, workers=workers, chunk=chunk, + layer_residency=layer_residency, + ) if banks is not None: logger.info_rank0(f"expert banks: FTW fast path (FTW checkpoint {model_path})") return banks @@ -354,7 +369,7 @@ def load_expert_banks( def _build(par: bool) -> ExpertBanks: if method is not None: - return _method_expert_banks(model_path, model_config, method, device, dummy, par, workers, chunk, layer_sink) + return _method_expert_banks(model_path, model_config, method, device, dummy, par, workers, chunk, layer_sink, resident) return _legacy_expert_banks(model_path, model_config, device, dtype, dummy, par, workers, chunk, decode_target, layer_sink) with requested_residency(layer_residency) as residency_plan: @@ -393,3 +408,33 @@ def _echo_residency(banks: ExpertBanks, requested, plan) -> ExpertBanks: "but saves no pinned quota" ) return banks + + +def attach_resident_banks(layers, banks: ExpertBanks) -> None: + """Give each resident MoE layer its rows of the device ``banks``; bank layer ``i`` goes to the layer whose ``layer_id`` is ``i``.""" + ids = sorted(layer.layer_id for layer in layers if layer.layer_id is not None) + num_layers = len(next(iter(banks.sources.values()))) + if ids != list(range(num_layers)) or len(ids) != len(layers): + raise ValueError(f"resident MoE layer ids {ids} do not cover the {num_layers} expert bank layers") + for layer in layers: + i, e = layer.layer_id, layer.num_experts + layout = layer.quant_method.layout() + rows = {role: (e, *spec.shape) for role, spec in layout.items() if not spec.resident} + if set(banks.sources) != set(rows): + raise ValueError(f"expert banks hold {sorted(banks.sources)}, the {layer.quant_method.kernel.name} kernel reads {sorted(rows)}") + for role, shape in rows.items(): + bank = banks.sources[role][i] + if tuple(bank.shape) != shape or bank.dtype != layout[role].dtype: + raise ValueError( + f"bank {role!r} of layer {i} is {tuple(bank.shape)} {bank.dtype} but the {layer.quant_method.kernel.name} kernel " + f"wants {shape} {layout[role].dtype}; the banks were packed for another kernel or model" + ) + alphas = [] + # apply() reads view.alphas as (gate_up, down) + for role in (role for role in ("gate_up_alpha", "down_alpha") if role in layout): + a = getattr(banks, role) + if a is None or a.numel() != num_layers * e or a.dtype != layout[role].dtype: + raise ValueError(f"the {layer.quant_method.kernel.name} kernel needs {num_layers * e} {layout[role].dtype} {role} values, the expert banks carry {None if a is None else (a.numel(), a.dtype)}") + alphas.append(a) + layer._expert_banks = {role: per_layer[i] for role, per_layer in banks.sources.items()} + layer._expert_alphas = tuple(a[i * e : (i + 1) * e] for a in alphas) or None diff --git a/python/freetoken/moe/fused_ds_fp4.py b/python/freetoken/moe/fused_ds_fp4.py index 65fffb5cc9..b4dc5cd0d0 100644 --- a/python/freetoken/moe/fused_ds_fp4.py +++ b/python/freetoken/moe/fused_ds_fp4.py @@ -12,10 +12,7 @@ import torch import triton -from freetoken.kernel.triton.dsv4.fp8_linear import ( - act_quant_fp8_inplace, - act_quant_fp8_roundtrip, -) +from freetoken.kernel.triton.dsv4.fp8_linear import act_quant_fp8_roundtrip from freetoken.kernel.triton.dsv4.fused_moe import ( _decode_dsfp4_moe_kernel, _e2m1_lut, @@ -50,7 +47,8 @@ def _grouped_decode( """Grouped per-route GEMM ``c[t,r,:] = a[row] @ dequant(W[slot])^T``. Returns ``[T, top_k, N]``. Used for both gate_up (a_row=token) and down - (a_row=route). FP4 weights are dequantized inline from the slot cache. + (a_row=route). FP4 weights are dequantized inline from the slot cache. A slot + of ``-1`` is an inactive route: its output is zero and its storage is never read. """ T, top_k = slots.shape N = packed_cache.shape[1] @@ -62,14 +60,13 @@ def _grouped_decode( topk_weights = out.new_empty((1, 1), dtype=torch.float32) scale_u8 = scale_cache.view(torch.uint8) - # Decode is a per-route GEMV, bound by the inline FP4 dequant (LUT gather + block - # scale), not weight HBM bandwidth -- so it tops out ~29% of HBM peak. BN=16/ - # BKB=128/1 warp maximizes memory-level parallelism (many single-warp CTAs, no - # cross-warp reduction) -> ~950 GB/s on H100 (was ~830 at BN=8/BKB=256/2). K_BYTES - # must be a multiple of BLOCK_SIZE_KB (gate_up 2048, down 1024 -- both /128). + # One fixed launch geometry, never tuned per machine: the row / warp split changes the per-row + # K reduction layout (and so the output bits) on some Triton versions. 16 rows on 2 warps is at + # or near the fastest split for the V4 and V4.1 decode shapes. Each K tile contains whole + # packed FP4 scale groups. BLOCK_SIZE_N = 16 BLOCK_SIZE_KB = 128 - _NW = 1 + _NW = 2 assert (K // 2) % BLOCK_SIZE_KB == 0, (K, BLOCK_SIZE_KB) grid = (total_routes, triton.cdiv(N, BLOCK_SIZE_N)) _decode_dsfp4_moe_kernel[grid]( @@ -103,28 +100,28 @@ def routed_experts_fp4( down_packed: torch.Tensor, # [S, H, I//2] uint8 down_scale: torch.Tensor, # [S, H, I//32] e8m0 swiglu_limit: float, + act_block: int = 128, ) -> torch.Tensor: """Full routed-expert output (summed over the top-k routes), excludes shared expert. - Precision matches the reference ``fp4_gemm(act_quant(x, 128), W_fp4)``: the gate_up - and down activations are FP8-round-tripped (block 128, ue8m0) before each GEMM. Since - an fp8 value x pow2 scale is exact in bf16, the round-tripped activation entering the - bf16 decode kernel is bit-identical to the reference's dequantized FP8 activation - (validated max diff = 0 vs the tilelang ``fp4_gemm`` reference).""" + Precision matches the reference ``fp4_gemm(act_quant(x, act_block), W_fp4)``: the gate_up + and down activations are FP8-round-tripped (block ``act_block`` -- the checkpoint's fp8 block, + 128 on V4 and 32 on V4.1 -- ue8m0) before each GEMM. Since an fp8 value x pow2 scale is exact + in bf16, the round-tripped activation entering the bf16 decode kernel is bit-identical to the + reference's dequantized FP8 activation (validated max diff = 0 vs the tilelang ``fp4_gemm`` + reference). The routing weight scales each route's down output in the GEMM epilogue.""" T, top_k = slots.shape H = x.shape[1] two_I = gate_up_packed.shape[1] I = two_I // 2 - x = act_quant_fp8_roundtrip(x, 128) # gate_up activation -> FP8 round-trip (no clone) + x = act_quant_fp8_roundtrip(x, act_block) # gate_up activation -> FP8 round-trip (no clone) gate_up = _grouped_decode( x, gate_up_packed, gate_up_scale, slots, None, a_row_is_route=False, mul_routed_weight=False, ) # [T, top_k, 2I] - act = fused_swiglu(gate_up, swiglu_limit) # [T, top_k, I] - - act = act.reshape(T * top_k, I) - act_quant_fp8_inplace(act, 128) # down activation -> FP8 round-trip + # [T, top_k, I] with the down activation's FP8 round-trip in the same pass + act = fused_swiglu(gate_up, swiglu_limit, act_block=act_block).reshape(T * top_k, I) down = _grouped_decode( act, down_packed, down_scale, slots, topk_weights, a_row_is_route=True, mul_routed_weight=True, @@ -191,6 +188,7 @@ def routed_experts_fp4_prefill( down_scale: torch.Tensor, # [S, H, I//32] e8m0 swiglu_limit: float, num_rows: int, + act_block: int = 128, ) -> torch.Tensor: """Grouped-GEMM counterpart of :func:`routed_experts_fp4` for dense prefill chunks: one moe_align sort shared by both GEMMs, each expert's weights @@ -201,7 +199,7 @@ def routed_experts_fp4_prefill( if T * top_k < _GROUPED_MIN_ROUTES: return routed_experts_fp4( x, slots, topk_weights, - gate_up_packed, gate_up_scale, down_packed, down_scale, swiglu_limit, + gate_up_packed, gate_up_scale, down_packed, down_scale, swiglu_limit, act_block=act_block, ) H = x.shape[1] two_I = gate_up_packed.shape[1] @@ -217,16 +215,14 @@ def routed_experts_fp4_prefill( sorted_ids, expert_ids, ntpp = moe_align_block_size(slots, cfg["BLOCK_SIZE_M"], num_rows) tw = topk_weights.reshape(-1).contiguous() - x = act_quant_fp8_roundtrip(x, 128) # gate_up activation -> FP8 round-trip (no clone) + x = act_quant_fp8_roundtrip(x, act_block) # gate_up activation -> FP8 round-trip (no clone) gate_up = torch.empty((T, top_k, two_I), dtype=x.dtype, device=x.device) _grouped_prefill( x, gate_up_packed, gate_up_scale, gate_up, tw, sorted_ids, expert_ids, ntpp, routes, top_k, False, cfg, ) - act = fused_swiglu(gate_up, swiglu_limit) # [T, top_k, I] - - act = act.reshape(routes, I) - act_quant_fp8_inplace(act, 128) # down activation -> FP8 round-trip + # [T, top_k, I] with the down activation's FP8 round-trip in the same pass + act = fused_swiglu(gate_up, swiglu_limit, act_block=act_block).reshape(routes, I) down = torch.empty((T, top_k, H), dtype=x.dtype, device=x.device) _grouped_prefill( act, down_packed, down_scale, down, tw, diff --git a/python/freetoken/moe/fused_nvfp4.py b/python/freetoken/moe/fused_nvfp4.py index 7988fa9efc..8334bbbc4c 100644 --- a/python/freetoken/moe/fused_nvfp4.py +++ b/python/freetoken/moe/fused_nvfp4.py @@ -265,7 +265,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), @@ -276,6 +275,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 32e0f3cc1e..246b6b248f 100644 --- a/python/freetoken/moe/offload_cache.py +++ b/python/freetoken/moe/offload_cache.py @@ -2,6 +2,8 @@ import math import os +import weakref +from contextlib import contextmanager from dataclasses import dataclass from typing import Iterator @@ -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, @@ -25,6 +30,8 @@ # entry the batch sees is >= this size. _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 logger = init_logger(__name__) @@ -144,11 +151,27 @@ 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 + # 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.layout is None: assert self.quant_format in _BANK_SCHEMAS, f"unknown quant_format {self.quant_format!r}" # Attached by the engine for decode_target == "cpu" (CpuMoeExecutor); None @@ -165,7 +188,11 @@ def __post_init__(self) -> None: "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, @@ -199,6 +226,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 @@ -257,6 +285,7 @@ def __post_init__(self) -> None: # Source pointers are per layer (_copy_src_ptrs[layer_id] -> [num_banks] device # tensor); dst/feat are layer-invariant. self._copy_fused_ok = False + self._copy_blocks_per_bank = 8 self._copy_dst_ptrs: torch.Tensor | None = None self._copy_src_ptrs: list[torch.Tensor] | None = None self._copy_feat_bytes: torch.Tensor | None = None @@ -266,6 +295,9 @@ def __post_init__(self) -> None: # _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 + # 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 # 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] = [] @@ -288,6 +320,151 @@ def __post_init__(self) -> None: self._batch_memcpy = None self.prefill_hit_rows = 0 self.prefill_total_rows = 0 + # 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() + + 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 _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, @@ -361,6 +538,10 @@ def set_bank_sources( ) self.banks = [(self.bank_sources[n], self.bank_caches[n]) for n in self.bank_schema] 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() @@ -438,6 +619,7 @@ def _build_fused_copy_plan(self) -> None: 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() + self._copy_blocks_per_bank = default_blocks_per_bank(self.device) self._copy_fused_ok = True def validate_rebuild(self, cache_size: int) -> None: @@ -497,6 +679,8 @@ def rebuild(self, cache_size: int) -> None: ) 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 + if self.verify is not None: + self.verify.bind(self) # 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) @@ -509,6 +693,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_() @@ -520,6 +706,10 @@ 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 @@ -632,15 +822,22 @@ 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: + def _invalidate_prefill_buffer(self, buffer_id: int, layer_id: int = -1) -> 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_() + if self.audit is not None: + from freetoken.moe.slot_audit import INVALIDATE + + self.audit.record_clear(self, INVALIDATE, layer_id, slot_start, self.num_experts) + # One fixed-shape launch. The previous 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 + + invalidate_prefill_slots( + self.id_of_slot, self.slot_for_id, self.usage, slot_start, self.num_experts + ) def begin_prefill(self) -> None: if not self.prefill_overlap: @@ -682,9 +879,13 @@ def prefetch_prefill_layer(self, layer_id: int) -> None: ) def copy() -> None: - self._invalidate_prefill_buffer(buffer_id) + self._invalidate_prefill_buffer(buffer_id, layer_id) for (per_layer, _), buffer in zip(self.banks, self.prefill_bank_buffers): buffer[buffer_id].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) @@ -784,11 +985,16 @@ def _prefetch_split(self, layer_id: int, buffer_id: int) -> None: 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] @@ -813,6 +1019,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, ...]: @@ -840,8 +1051,21 @@ def release_prefill_layer(self, layer_id: int) -> None: self._prefill_buffer_has_release_event[buffer_id] = True self._prefill_buffer_released[buffer_id] = True - def ensure_experts(self, layer_id: int, expert_ids: torch.Tensor) -> None: - from freetoken.moe.offload_kernels import ensure_experts + 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, + ) -> 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: # ``expert_ids`` still holds raw expert ids here (the kernel rewrites them to @@ -850,7 +1074,49 @@ def ensure_experts(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(self, layer_id, expert_ids) + 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) + 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 + + 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) -> None: """Capped-fetch LRU for the hybrid backend. @@ -878,15 +1144,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 @@ -979,10 +1266,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) @@ -992,7 +1279,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] @@ -1004,13 +1291,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: @@ -1023,51 +1313,76 @@ def copy_missing(self) -> None: # 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]) + 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 + 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) + + 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, + ) -> 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.""" + self._copy_rows(layer_id, dst_slots, src_rows, num, slim=slim) + 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 + ) -> 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, - ) + # 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). + if slim: + fast_index_copy_multi_slim_jit( + self._copy_dst_ptrs, self._copy_src_ptrs[layer_id], self._copy_feat_bytes, dst_slots, src_rows, num, + ) + else: + fast_index_copy_multi_jit( + self._copy_dst_ptrs, self._copy_src_ptrs[layer_id], self._copy_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, - ) + 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. -def iter_offload_moe_layers(model) -> Iterator: - from freetoken.layers import BaseOP, OffloadMoELayer + 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]) - if isinstance(model, OffloadMoELayer): - yield model - if not isinstance(model, BaseOP): - return +def iter_offload_moe_layers(model) -> Iterator: + from freetoken.layers import OffloadMoELayer, iter_moe_layers - for value in model.__dict__.values(): - if isinstance(value, BaseOP): - yield from iter_offload_moe_layers(value) - elif isinstance(value, (list, tuple)): - for item in value: - yield from iter_offload_moe_layers(item) + return (layer for layer in iter_moe_layers(model) if isinstance(layer, OffloadMoELayer)) def attach_offload_moe_cache(model, cache: OffloadMoeCache) -> list: diff --git a/python/freetoken/moe/offload_kernels.py b/python/freetoken/moe/offload_kernels.py index cf513f52d1..5993738d10 100644 --- a/python/freetoken/moe/offload_kernels.py +++ b/python/freetoken/moe/offload_kernels.py @@ -40,6 +40,81 @@ def ensure_experts(cache, layer_id: int, expert_ids: torch.Tensor) -> None: ) +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, + cache.usage, + cache.step, + expert_ids, + cache.src_indices, + cache.evict_slots, + cache.num_indices, + 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 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, max_fetch: int, fetch_fraction: float = 0.0 ) -> None: 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..981239758f --- /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, 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..1f096c9c3a --- /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, 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 8f77be9330..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 @@ -66,6 +91,11 @@ def __init__(self, num_pages: int, page_size: int, page_table: torch.Tensor, typ supports_runtime_rebuild = True prefill_chunk_budget = None # generic shared page pool: no per-model prefill chunk cap + @property + def prefix_replay_tokens(self) -> int: + """Prompt-tail tokens a prefix hit must leave to the prefill (a pool whose model recomputes them).""" + return getattr(self.swa_pool, "prefix_replay_tokens", 0) if self.swa_pool is not None else 0 + @property def prefill_chunk_align(self) -> int: """Granularity a non-final prefill chunk should end on. A hybrid snapshot is donated only @@ -90,10 +120,10 @@ def _make_prefix_cache(self, device, page_size, type): return SWARadixCache(device, page_size, self.sliding_window_size) return create_prefix_cache(device=device, type=type, page_size=page_size) - def match_req(self, req: PendingReq) -> MatchResult: + def match_req(self, req: PendingReq, max_len: int | None = None) -> MatchResult: input_len = req.input_len assert input_len > 0, "Input length must be greater than 0." - ids = req.input_ids[: input_len - 1] + ids = req.input_ids[: input_len - 1 if max_len is None else max(0, min(input_len - 1, max_len))] if self.is_swa: from freetoken.kvcache.swa_radix_cache import SWACacheHandle m = self.prefix_cache.match_prefix(ids) @@ -101,6 +131,16 @@ def match_req(self, req: PendingReq) -> 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) @@ -548,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. @@ -561,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 bd8129aae8..9019604707 100644 --- a/python/freetoken/scheduler/prefill.py +++ b/python/freetoken/scheduler/prefill.py @@ -71,8 +71,9 @@ def _try_allocate_one(self, req: PendingReq): if self.table_manager.available_size == 0: return None - # TODO: consider host cache match case - mr = self.cache_manager.match_req(req) + # TODO: consider host cache match case (it needs the same replay cap) + replay = self.cache_manager.prefix_replay_tokens + mr = self.cache_manager.match_req(req, max_len=max(0, req.input_len - replay) if replay else None) handle = mr.cuda_handle cached_len = handle.cached_len # TODO: better estimate policy @@ -142,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) @@ -214,6 +216,7 @@ def _add_one_req( uid=pending_req.uid, cache_handle=cache_handle, sampling_params=pending_req.sampling_params, + prompt_len=pending_req.input_len, ) req.mm_items = pending_req.mm_items req.mrope_positions_full = pending_req.mrope_positions_full @@ -223,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 @@ -240,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 fa0d96a714..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 ahead of the token delivered to the host. + # 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). @@ -737,17 +816,10 @@ def _current_cache_geometry(self) -> dict: the CONCRETE current window (usable pages) so a rollback restores it byte-for-byte, whether it was pinned or ratio-derived.""" eng = self.engine - config = self.config - mc = config.model_config - num_swa_pages = None - if getattr(mc, "dsv4_args", None) is not None: - sizes = getattr(eng.kv_cache, "sizes", None) - if sizes is not None: # usable window pages = physical n_win_pages minus the dummy page - num_swa_pages = max(0, sizes.n_win_pages - 1) - elif getattr(mc, "has_swa_attention", False) and ( - getattr(config, "cache_type", None) == "swa_radix" - ): # usable window tokens = pool tokens minus the slot-0 sentinel - num_swa_pages = max(0, int(getattr(eng.kv_cache, "swa_num_tokens", 0) or 0) - 1) + from freetoken.kvcache.cache_status import window_pool_spec + + spec = window_pool_spec(self.config) + num_swa_pages = eng.kv_cache.window_pages if spec is not None else None return dict( num_pages=eng.num_pages, moe_cache_size=eng.moe_offload_cache.cache_size if eng.moe_offload_cache is not None else None, @@ -868,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) @@ -887,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 0d60940c97..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, @@ -120,26 +123,36 @@ async def handle_anthropic_messages( default_max_tokens=( 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)) @@ -153,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)) @@ -164,7 +179,10 @@ async def handle_anthropic_count_tokens(req: AnthropicCountTokensRequest, state: 400, "invalid_request_error", "messages: no tokenizable content" ) try: - n_tokens = await count_prompt_tokens(messages, template_tools, ctk, state) + n_tokens = await count_prompt_tokens( + messages, template_tools, ctk, state, + inline_system_policy=getattr(state.config, "anthropic_inline_system", "auto"), + ) except GenerationError as exc: # The chat template could not render this conversation (bad role ordering, an unmatched # tool_result, ...) — a client error, exactly as /v1/messages classifies the same failure. @@ -182,13 +200,13 @@ 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 a counted prompt is exactly the prompt a generation would tokenize.""" - # Collect all system content (top-level `system` + any system-role messages - # Claude Code interleaves in the array) and emit ONE system message at the - # front: strict chat templates (e.g. Qwen3.5) require system at the beginning. + # Hoisting turn-specific reminders rewrites the cached prompt prefix. + # Preserve their position; the tokenizer adapts them to the model's roles. system_texts: list[str] = [] if req.system: if isinstance(req.system, str): @@ -201,7 +219,11 @@ def convert_anthropic_prompt( other: list[dict[str, Any]] = [] for msg in req.messages: if msg.role == "system": - system_texts.append(_content_text(msg.content)) + text = _content_text(msg.content) + if not other: + system_texts.append(text) + elif text: + other.append({"role": "system", "content": text}) continue if isinstance(msg.content, str): @@ -291,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: @@ -299,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( @@ -308,9 +335,11 @@ def convert_anthropic_to_genspec( model_sampling: dict[str, Any], 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, @@ -327,6 +356,7 @@ def convert_anthropic_to_genspec( chat_template_kwargs=ctk, template_tools=template_tools, parser_tools=parser_tools, + inline_system_policy=inline_system_policy, ) @@ -375,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: @@ -402,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, ) @@ -420,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. @@ -444,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]: @@ -576,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 ffd7172802..bd45e2d09d 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 = 0.0 frequency_penalty: float = 0.0 + # 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 5592a2b0d5..828fc6c775 100644 --- a/python/freetoken/server/api_server.py +++ b/python/freetoken/server/api_server.py @@ -57,6 +57,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: @@ -137,6 +141,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())) @@ -358,24 +369,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)) @@ -446,38 +489,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): @@ -563,8 +625,12 @@ def _resolve_num_swa_pages(state: FrontendManager, req: CacheRebuildRequest) -> or pools.get("num_pages", 0) or 0 ) page_size = int(pools.get("page_size", 0) or getattr(config, "page_size", 1) or 1) - is_dsv4 = getattr(getattr(config, "model_config", None), "dsv4_args", None) is not None - swa_page_size = page_size if is_dsv4 else 1 + from freetoken.kvcache.cache_status import window_pool_spec + + spec = window_pool_spec(config) + if spec is None: + raise ValueError("swa_full_tokens_ratio requested but this model has no resizable window pool") + swa_page_size = spec.page_size window_tokens = int(round(req.swa_full_tokens_ratio * num_pages * page_size)) return max(1, -(-window_tokens // swa_page_size)) # ceil-div to the pool's page unit @@ -604,12 +670,16 @@ async def cache_rebuild(req: CacheRebuildRequest): {"status": "failed", "error": "swa_full_tokens_ratio must be in (0, 1]"}, status_code=422, ) + try: + num_swa_pages = _resolve_num_swa_pages(state, req) + except ValueError as exc: + return JSONResponse({"status": "failed", "error": str(exc)}, status_code=422) result = await dispatch_rebuild( state, moe_cache_size=req.moe_cache_size, num_pages=req.num_pages, num_mamba_slots=req.num_mamba_slots, - num_swa_pages=_resolve_num_swa_pages(state, req), + num_swa_pages=num_swa_pages, mode=req.mode, timeout=req.timeout, ) @@ -710,6 +780,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 @@ -725,11 +819,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( @@ -739,11 +832,11 @@ def cache_geometry(state: Any) -> dict: # the internal currency), else the load-time meta value. 0.0 for models without a window pool. swa_full_tokens_ratio = float(getattr(state, "swa_full_tokens_ratio", 0.0) or 0.0) last_swa_pages = last.get("num_swa_pages") - if last_swa_pages: - is_dsv4 = getattr(getattr(config, "model_config", None), "dsv4_args", None) is not None - full = num_pages if is_dsv4 else num_pages * page_size + swa_page_size = int(pools.get("swa_page_size", 0) or 0) + if last_swa_pages and swa_page_size: + full = num_pages * page_size if full > 0: - swa_full_tokens_ratio = min(1.0, last_swa_pages / full) + swa_full_tokens_ratio = min(1.0, last_swa_pages * swa_page_size / full) try: model_config = config.model_config num_experts = int(getattr(model_config, "num_experts", 0) or 0) @@ -785,9 +878,8 @@ def cache_geometry(state: Any) -> dict: "moe_cache_policy": getattr(config, "moe_cache_policy", None), "unit_bytes": unit_bytes, "swa_full_tokens_ratio": swa_full_tokens_ratio, - # The window pool's page unit for num_swa_pages: DSV4 = P (== page_size), radix-SWA = 1 - # token, 0 for models without a window pool. Lets a client denominate the swa control. - "swa_page_size": int(pools.get("swa_page_size", 0) or 0), + # The pool's window page unit; 0 for models without a resizable window pool. + "swa_page_size": swa_page_size, # Concrete current window size in that unit (usable pages): the last rebuild's value if it # pinned/derived one, else the load-time pool size. 0 for models without a window pool. "num_swa_pages": int(last.get("num_swa_pages") or pools.get("num_swa_pages", 0) or 0), diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py index 23a62f8419..3e693bbe78 100644 --- a/python/freetoken/server/args.py +++ b/python/freetoken/server/args.py @@ -49,6 +49,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" @@ -59,6 +65,11 @@ class ServerArgs(SchedulerConfig): # prompt_tokens_details.cached_tokens, Anthropic cache_read_input_tokens, Responses # 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. @@ -99,10 +110,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: @@ -159,6 +166,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 @@ -203,6 +219,8 @@ def _infer_tool_call_parser(model_path: str) -> str: return "qwen3_coder" if "qwen" in marker: return "qwen25" + if "deepseek_v41" in marker or "deepseekv41" in marker: + return "deepseekv41" if "deepseek" in marker and ("v4" in marker or "deepseek_v4" in marker): return "deepseekv32" if "deepseek" in marker and ("v3.2" in marker or "v32" in marker): @@ -365,6 +383,18 @@ def _infer_reasoning_parser(model_path: str) -> str | None: help="The port number for the server to listen on.", ) + 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", @@ -418,6 +448,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, @@ -425,6 +469,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", @@ -535,6 +596,30 @@ def _infer_reasoning_parser(model_path: str) -> str | None: ), ) + parser.add_argument( + "--anthropic-inline-system", + choices=("auto", "preserve", "fold"), + default=ServerArgs.anthropic_inline_system, + help="Preserve inline system instructions when supported by the renderer, " + "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, @@ -567,6 +652,7 @@ def _infer_reasoning_parser(model_path: str) -> str | None: "qwen3_coder", "mistral", "deepseekv32", + "deepseekv41", "gemma4", "glm47", "minimax", @@ -594,14 +680,29 @@ 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, choices=["auto", *MOE_STRATEGIES], help=( "How the routed experts are served. 'auto' resolves a MoE model to the offload family " - "(offload, or hybrid when a `ft bench bw` profile recommends it); resident " - "'fused' experts must be requested explicitly." + "(offload, or hybrid when a `ft bench bw` profile recommends it), and to resident " + "'fused' experts on unified-memory GPUs (GB10 / DGX Spark)." ), ) @@ -636,6 +737,27 @@ def _infer_reasoning_parser(model_path: str) -> str | None: ), ) + parser.add_argument( + "--swa-decoder-replay", + default=ServerArgs.swa_decoder_replay, + choices=["bounded", "exact"], + help=( + "DeepSeek-V4.1 Decoder SWA Bounded Replay. 'bounded' (default) runs the 20 decoder layers " + "on each prompt's last 128 tokens with their sliding window truncated there, as in the " + "tech report; 'exact' runs them on every prompt token (the reference numerics)." + ), + ) + + 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, @@ -685,14 +807,46 @@ 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( + "--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( @@ -735,6 +889,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", @@ -819,6 +984,9 @@ def _infer_reasoning_parser(model_path: str) -> str | None: if entry: kwargs["quant_backend"] = entry + 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"]) @@ -856,16 +1024,26 @@ def _infer_reasoning_parser(model_path: str) -> str | None: 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 0804dcc3f7..88cc65b367 100644 --- a/python/freetoken/server/function_call_parser.py +++ b/python/freetoken/server/function_call_parser.py @@ -78,6 +78,8 @@ def _should_forward_unknown_tool(name: Any) -> bool: "<｜DSML｜function_calls>", "<｜DSML｜tool_calls>", "<｜DSML｜invoke", + "<｜DSML｜ calls>", + "<｜DSML｜ invoke", "", ] @@ -347,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": @@ -392,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: @@ -1563,37 +1686,45 @@ class DeepSeekV32Detector(BaseFormatDetector): Reference: https://huggingface.co/deepseek-ai/DeepSeek-V3.2 """ + # the DSML tag names: (block, alternative block), invoke, parameter + block_tags = ("function_calls", "tool_calls") + invoke_tag = "invoke" + param_tag = "parameter" + def __init__(self): super().__init__() self.dsml_token = "｜DSML｜" - self.bot_token = f"<{self.dsml_token}function_calls>" - self.eot_token = f"" - self.alt_bot_token = f"<{self.dsml_token}tool_calls>" - self.alt_eot_token = f"" - self.invoke_start_prefix = f"<{self.dsml_token}invoke" - self.invoke_end_token = f"" - self.param_end_token = f"" + block, alt_block = self.block_tags + self.bot_token = f"<{self.dsml_token}{block}>" + self.eot_token = f"" + self.alt_bot_token = f"<{self.dsml_token}{alt_block}>" + self.alt_eot_token = f"" + self.invoke_start_prefix = f"<{self.dsml_token}{self.invoke_tag}" + self.invoke_end_token = f"" + self.param_start_prefix = f"<{self.dsml_token}{self.param_tag}" + self.param_end_token = f"" # Regex for complete invoke extraction - _de = re.escape(self.dsml_token) + _inv = re.escape(self.dsml_token + self.invoke_tag) + _par = re.escape(self.dsml_token + self.param_tag) self.invoke_regex = re.compile( - rf'<{_de}invoke\s+name="([^"]+)"\s*>(.*?)', + rf'<{_inv}\s+name="([^"]+)"\s*>(.*?)', re.DOTALL, ) # Regex for parameter extraction self.param_regex = re.compile( - rf'<{_de}parameter\s+name="([^"]+)"(?:\s+string="(true|false)")?\s*>(.*?)', + rf'<{_par}\s+name="([^"]+)"(?:\s+string="(true|false)")?\s*>(.*?)', re.DOTALL, ) # Regex for partial invoke (name known, body still streaming) self.partial_invoke_regex = re.compile( - rf'<{_de}invoke\s+name="([^"]+)"\s*>(.*)', + rf'<{_inv}\s+name="([^"]+)"\s*>(.*)', re.DOTALL, ) # Streaming state machine tag regexes (anchored matches over the buffer). - self.invoke_open_regex = re.compile(rf'<{_de}invoke\s+name="([^"]+)"\s*>') + self.invoke_open_regex = re.compile(rf'<{_inv}\s+name="([^"]+)"\s*>') self.param_open_regex = re.compile( - rf'<{_de}parameter\s+name="([^"]+)"(?:\s+string="(true|false)")?\s*>' + rf'<{_par}\s+name="([^"]+)"(?:\s+string="(true|false)")?\s*>' ) self._last_arguments = "" @@ -1687,7 +1818,7 @@ def parse_streaming_increment(self, new_text: str, tools: List[Tool]) -> Streami normal_parts: List[str] = [] calls: List[ToolCallItem] = [] - param_open_token = f"<{self.dsml_token}parameter" + param_open_token = self.param_start_prefix def _emit_args(fragment: str) -> None: if fragment: @@ -2591,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 @@ -2612,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 @@ -2632,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] @@ -2643,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] @@ -3512,6 +3653,16 @@ def detect_and_parse(self, text: str, tools: List[Tool]) -> StreamingParseResult return StreamingParseResult(normal_text="".join(normal_parts).strip(), calls=calls) +class DeepSeekV41Detector(DeepSeekV32Detector): + """DeepSeek-V4.1's DSML: the V3.2 grammar with every tag name led by a space and the block + named ``calls`` (``<｜DSML｜ calls>`` / ``<｜DSML｜ invoke name=...>`` / ``<｜DSML｜ parameter + ...>``, the checkpoint's ``encoding/encoding.py``).""" + + block_tags = (" calls", " calls") + invoke_tag = " invoke" + param_tag = " parameter" + + class FunctionCallParser: """ Parser for function/tool calls in model outputs. @@ -3523,6 +3674,7 @@ class FunctionCallParser: ToolCallParserEnum: Dict[str, Type[BaseFormatDetector]] = { "deepseekv32": DeepSeekV32Detector, + "deepseekv41": DeepSeekV41Detector, "gemma4": Gemma4Detector, "gpt-oss": GptOssDetector, "gpt_oss": GptOssDetector, diff --git a/python/freetoken/server/generation.py b/python/freetoken/server/generation.py index 91056e8d1f..e5a215658b 100644 --- a/python/freetoken/server/generation.py +++ b/python/freetoken/server/generation.py @@ -14,6 +14,7 @@ from __future__ import annotations import asyncio +import math import json import time from collections.abc import AsyncIterator @@ -24,7 +25,8 @@ 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: # Chat templates render through jinja2 (a transformers dependency): a TemplateError means @@ -43,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 @@ -73,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 @@ -107,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 @@ -114,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 @@ -129,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 @@ -142,6 +419,7 @@ class GenSpec: chat_template_kwargs: dict[str, Any] = field(default_factory=dict) template_tools: list[dict[str, Any]] | None = None # tools the model sees (TokenizeMsg.tools) parser_tools: list[dict[str, Any]] | None = None # tools for FunctionCallParser; None disables parsing + inline_system_policy: str | None = None @property def parse_tools(self) -> bool: @@ -167,9 +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) @@ -181,21 +469,79 @@ 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_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), stop_strs=[s for s in stop_list if s], # drop empty strings (would match everything) + 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]: @@ -302,6 +648,7 @@ async def submit_generation(spec: GenSpec, state: Any) -> int: chat_template_kwargs=spec.chat_template_kwargs, tools=spec.template_tools, images=images, + inline_system_policy=spec.inline_system_policy, ) ) return uid @@ -312,6 +659,8 @@ async def count_prompt_tokens( tools: list[dict[str, Any]] | None, chat_template_kwargs: dict[str, Any], state: Any, + *, + inline_system_policy: str | None = None, ) -> int: """Token count of an already-converted (messages, tools, chat_template_kwargs) prompt, using the frontend's own tokenizer (``state.frontend_tokenizer()``) so the count equals the @@ -334,10 +683,13 @@ class the generation path maps to a 400. A tokenizer *initialization* failure (m chat_template_kwargs=chat_template_kwargs, tools=tools, images=images, + inline_system_policy=inline_system_policy, ) manager = await asyncio.to_thread(state.frontend_tokenizer) # init failure -> server fault try: (user_msg,) = await asyncio.to_thread(manager.tokenize, [msg]) + except InlineSystemError as exc: + raise GenerationError(str(exc)) from exc except _TemplateError as exc: raise GenerationError(str(exc)) from exc return int(user_msg.input_ids.numel()) @@ -361,6 +713,7 @@ async def prerender_error(spec: GenSpec, state: Any) -> GenerationError | None: sampling_params=SamplingParams(), chat_template_kwargs=spec.chat_template_kwargs, tools=spec.template_tools, + inline_system_policy=spec.inline_system_policy, ) try: manager = await asyncio.to_thread(build) @@ -382,7 +735,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. @@ -491,6 +849,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, @@ -532,27 +969,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, ) @@ -561,25 +995,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. @@ -594,10 +1032,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: @@ -609,6 +1063,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. @@ -619,6 +1081,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: @@ -634,8 +1118,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]: @@ -649,12 +1131,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 ( @@ -663,7 +1141,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, @@ -671,14 +1149,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 @@ -686,14 +1166,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: @@ -703,7 +1187,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) @@ -714,7 +1200,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: @@ -724,7 +1210,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" @@ -738,58 +1226,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) @@ -806,13 +1328,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..ef4c40cb3c 100644 --- a/python/freetoken/server/launch.py +++ b/python/freetoken/server/launch.py @@ -52,7 +52,18 @@ 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: 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 dc2f73a979..7f94efc250 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") @@ -60,26 +75,24 @@ 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, - ), + 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), @@ -131,6 +144,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() @@ -146,6 +169,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, @@ -155,15 +298,19 @@ 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): return create_error_response( "response_format json_object/json_schema is not supported (no constrained decoding)", 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 @@ -185,7 +332,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)) @@ -196,45 +348,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( @@ -242,21 +435,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") @@ -277,17 +482,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, @@ -305,7 +507,7 @@ async def stream_chat_completion_chunks( "type": "function", "function": {"name": ev.name, "arguments": ""}, }]}, - "index": 0, "finish_reason": None, + "index": index, "finish_reason": None, }], ) ) @@ -322,7 +524,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, }], ) ) @@ -346,7 +548,7 @@ async def stream_chat_completion_chunks( "index": open_tool["index"], "function": {"arguments": remainder}, }]}, - "index": 0, "finish_reason": None, + "index": index, "finish_reason": None, }], ) ) @@ -358,7 +560,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 @@ -366,21 +568,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" @@ -394,47 +602,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) @@ -442,80 +667,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, @@ -541,7 +931,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, @@ -550,7 +940,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: @@ -613,10 +1023,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, } @@ -630,7 +1041,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, @@ -640,6 +1053,10 @@ 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 @@ -651,14 +1068,9 @@ def _response_format_unsupported(response_format: dict[str, Any] | None) -> bool 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)" return None @@ -701,10 +1113,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 1857d3d7c0..c9d0212c32 100644 --- a/python/freetoken/server/stats.py +++ b/python/freetoken/server/stats.py @@ -120,12 +120,10 @@ def derive_model_card(config: Any) -> dict: def _swa_page_size(config: Any) -> int: - """The window pool's own page unit: P (window_size) for DSV4, 1 token for radix-SWA. - Mirrors compute_cache_pools' swa_page_size.""" - dsv4 = getattr(getattr(config, "model_config", None), "dsv4_args", None) - if dsv4 is not None: - return int(getattr(dsv4, "window_size", 0) or 1) - return 1 + from freetoken.kvcache.cache_status import window_pool_spec + + spec = window_pool_spec(config) + return spec.page_size if spec is not None else 1 def build_stats(state: Any, p95_ms: int, ttft_mean_ms: int) -> dict: @@ -138,9 +136,10 @@ def build_stats(state: Any, p95_ms: int, ttft_mean_ms: int) -> dict: config = state.config ready_at = getattr(state, "ready_at", None) uptime_s = max(0, int(time.monotonic() - ready_at)) if ready_at is not None else 0 + pools = getattr(state, "cache_pools", None) or {} kv = ( {"used_pages": tr.kv_used_pages, "total_pages": tr.kv_total_pages, - "page_size": getattr(config, "page_size", 1)} + "page_size": pools.get("page_size") or getattr(config, "page_size", 1)} if tr.kv_total_pages > 0 else None ) mamba = ( @@ -175,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/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/effort.py b/python/freetoken/tokenizer/effort.py index c81cd7bef1..2e9e3cd21f 100644 --- a/python/freetoken/tokenizer/effort.py +++ b/python/freetoken/tokenizer/effort.py @@ -84,7 +84,8 @@ class EffortProfile: template interpolates anything". ``strength_dialect`` True means the template reads the ``reasoning_strength`` spelling (muse-glimmer's family, whose documented ladder tops out at xhigh) rather than only - ``reasoning_effort``. + ``reasoning_effort``. ``integer_range`` records a validated 1-100 budget + accepted independently of the named effort vocabulary. """ supported: frozenset[str] @@ -92,6 +93,7 @@ class EffortProfile: consumes_effort: bool validates: bool = False strength_dialect: bool = False + integer_range: tuple[int, int] | None = None @dataclass(frozen=True) @@ -118,7 +120,7 @@ class ThinkingProfile: _MAX_QUANTIZE_DISTANCE = 0.15 -def quantize_effort(value: Any, profile: EffortProfile) -> str | None: +def quantize_effort(value: Any, profile: EffortProfile) -> str | int | None: """Map a client's effort onto ``profile``; ``None`` means "send nothing". In-vocabulary values pass through untouched. Other named levels land on the @@ -128,6 +130,11 @@ def quantize_effort(value: Any, profile: EffortProfile) -> str | None: the template default. With "max" excluded the remaining scale values are unique, so quantization is deterministic across processes. """ + if profile.integer_range is not None and isinstance(value, (int, float)): + low, high = profile.integer_range + if type(value) is not int or not low <= value <= high: + raise ValueError(f"reasoning_effort must be an integer in [{low}, {high}]") + return value supported = effective_efforts(profile) if not supported: return None @@ -181,6 +188,7 @@ def probe_effort_profile( diverged: set[str] = set() matches_baseline: dict[str, bool] = {name: True for name in KNOWN_REASONING_EFFORTS} ran_rounds = 0 + integer_range = None for base_kwargs, tools in _PROBE_ROUNDS: try: @@ -198,6 +206,10 @@ def probe_effort_profile( if _renderings_differ(rendering, baseline): diverged.add(name) matches_baseline[name] = False + # Only a validated budget range is a numeric capability; arbitrary Jinja + # interpolation must not advertise support for numeric reasoning effort. + if _accepts_integer_budget(render, base_kwargs, tools): + integer_range = (1, 100) if ran_rounds == 0: # Nothing learnable: sending no effort is the only safe rendering. @@ -214,7 +226,7 @@ def probe_effort_profile( strength_dialect = False supported = frozenset(name for name in KNOWN_REASONING_EFFORTS if name not in rejected) - consumes = bool(rejected or diverged) + consumes = bool(rejected or diverged or integer_range) default = None if consumes: defaults = [name for name in supported if matches_baseline[name]] @@ -226,9 +238,26 @@ def probe_effort_profile( consumes_effort=consumes, validates=bool(rejected), strength_dialect=strength_dialect, + integer_range=integer_range, ) +def _accepts_integer_budget(render, base_kwargs, tools) -> bool: + for value in (0, 101): + try: + render({**base_kwargs, "reasoning_effort": value}, tools) + except Exception: # noqa: BLE001 -- the encoder owns validation + continue + return False + outputs = [] + for value in range(1, 101): + try: + outputs.append(render({**base_kwargs, "reasoning_effort": value}, tools)) + except Exception: # noqa: BLE001 + return False + return _renderings_differ(outputs[0], outputs[-1]) + + def probe_thinking_profile( render: Callable[[dict[str, Any], list[dict[str, Any]] | None], Any], efforts: EffortProfile, diff --git a/python/freetoken/tokenizer/inline_system.py b/python/freetoken/tokenizer/inline_system.py new file mode 100644 index 0000000000..7e6f2c06f9 --- /dev/null +++ b/python/freetoken/tokenizer/inline_system.py @@ -0,0 +1,173 @@ +# SPDX-License-Identifier: Apache-2.0 +# Renderer-probe approach adapted from vLLM PR #58772. +"""Preserve or fold late system instructions without hoisting the prompt head.""" +from __future__ import annotations + +from copy import deepcopy +from typing import Any, Callable + + +# Allow small role-framing differences, not a repeated environment header. +_INLINE_ROLE_SLACK_TOKENS = 8 + + +class InlineSystemError(ValueError): + """An input instruction cannot be placed without corrupting tool ordering.""" + + +def has_inline_system(messages: list[dict]) -> bool: + started = False + for message in messages: + if message.get("role") == "system": + if started: + return True + else: + started = True + return False + + +def _add_text(message: dict, text: str, *, prepend: bool = False) -> dict: + result = dict(message) + content = result.get("content") + def join(a, b): + return "\n\n".join(part for part in (a, b) if part) + if content is None or isinstance(content, str): + result["content"] = join(text, content) if prepend else join(content, text) + elif isinstance(content, list): + parts = deepcopy(content) + index = 0 if prepend else -1 + if parts and parts[index].get("type") == "text": + previous = parts[index].get("text") or "" + parts[index]["text"] = join(text, previous) if prepend else join(previous, text) + else: + parts.insert(0 if prepend else len(parts), {"type": "text", "text": text}) + result["content"] = parts + else: + raise InlineSystemError("inline system folding requires text or content blocks") + return result + + +def _after_tool_results(messages: list[dict]) -> list[dict]: + out, held = [], [] + pending_ids: set[str] = set() + for message in messages: + role = message.get("role") + if role == "system" and pending_ids: + held.append(message) + continue + if role == "tool": + out.append(message) + pending_ids.discard(message.get("tool_call_id", "")) + if not pending_ids: + out.extend(held) + held.clear() + continue + if pending_ids and held: + raise InlineSystemError("inline system instruction interrupts incomplete tool results") + out.append(message) + pending_ids = {call.get("id", "") for call in message.get("tool_calls", [])} if role == "assistant" else set() + if held: + raise InlineSystemError("inline system instruction is waiting for missing tool results") + return out + + +def normalize_inline_system(messages: list[dict], mode: str) -> list[dict]: + if mode not in ("preserve", "fold"): + raise InlineSystemError(f"invalid inline system mode: {mode}") + if not has_inline_system(messages): + return messages + ordered = _after_tool_results(deepcopy(messages)) + # Without a stable leading system, some renderers attach tools to the first + # late system message. Folding avoids moving tools when that message changes. + if mode == "preserve" and ordered[0].get("role") == "system": + return ordered + out, pending = [], [] + started = False + for message in ordered: + role = message.get("role") + if role == "system" and started: + text = message.get("content") or "" + if not isinstance(text, str): + raise InlineSystemError("inline system instructions must be text") + if text: + if out and out[-1].get("role") in ("user", "tool") and not pending: + out[-1] = _add_text(out[-1], text) + else: + pending.append(text) + continue + if pending: + text = "\n\n".join(pending) + pending.clear() + if role == "user": + message = _add_text(message, text, prepend=True) + else: + out.append({"role": "user", "content": text}) + out.append(message) + started = started or role != "system" + if pending: + out.append({"role": "user", "content": "\n\n".join(pending)}) + return out + + +def _insertion(base: list[int], new: list[int]) -> tuple[int, list[int]] | None: + start = 0 + while start < min(len(base), len(new)) and base[start] == new[start]: + start += 1 + suffix = 0 + while suffix < min(len(base), len(new)) - start and base[-suffix-1] == new[-suffix-1]: + suffix += 1 + if start + suffix != len(base): + return None + return start, new[start:len(new)-suffix] + + +def probe_inline_system(render: Callable[[list[dict]], str], tokenizer: Any) -> tuple[str, str]: + """Conservatively require a distinct, prefix-stable system segment.""" + note, reasoning = "Echo inline instruction", "Foxtrot previous reasoning" + system = {"role": "system", "content": "Alpha stable instructions"} + user = {"role": "user", "content": "Bravo first question"} + cases = [ + ("middle", [system, user, {"role": "assistant", "content": "Charlie first answer"}, + {"role": "user", "content": "Delta second question"}], 3), + ("trailing", [system, user], 2), + ("tool result", [system, user, + {"role": "assistant", "content": "", "reasoning_content": reasoning, "thinking": reasoning, + "reasoning": reasoning, "tool_calls": [{"id": "probe_call", "type": "function", + "function": {"name": "read", "arguments": {}}}]}, + {"role": "tool", "tool_call_id": "probe_call", "content": "Golf tool output"}], 4), + ] + max_system_tokens = None + try: + special_ids = set(tokenizer.all_special_ids) | set(tokenizer.get_added_vocab().values()) + if not special_ids: + return "fold", "renderer has no identifiable special tokens" + for name, messages, index in cases: + base_text = render(deepcopy(messages)) + base = tokenizer.encode(base_text, add_special_tokens=False) + with_note = messages[:index] + [{"role": "system", "content": note}] + messages[index:] + new_text = render(deepcopy(with_note)) + new = tokenizer.encode(new_text, add_special_tokens=False) + found = _insertion(base, new) + if found is None: + return "fold", f"{name}: other prompt tokens change" + position, segment = found + if index == len(messages) and position >= len(base): + return "fold", f"{name}: instruction follows the generation opening" + if tokenizer.decode(segment).count(note) != 1 or not special_ids.intersection(segment): + return "fold", f"{name}: no distinct role-marked instruction" + if (reasoning in base_text) != (reasoning in new_text): + return "fold", f"{name}: reasoning retention changes" + if name == "middle": + as_user = messages[:index] + [{"role": "user", "content": note}] + messages[index:] + user_tokens = tokenizer.encode(render(deepcopy(as_user)), add_special_tokens=False) + user_segment = _insertion(base, user_tokens) + if user_segment is None: + return "fold", "middle: cannot isolate a user instruction" + if user_segment[1] == segment: + return "fold", "system instructions render as ordinary user turns" + max_system_tokens = len(user_segment[1]) + _INLINE_ROLE_SLACK_TOKENS + if max_system_tokens is not None and len(segment) > max_system_tokens: + return "fold", f"{name}: system insertion overhead ({len(segment)} > {max_system_tokens} tokens)" + except Exception as exc: + return "fold", f"probe could not render: {type(exc).__name__}" + return "preserve", "renderer preserves distinct system turns and existing prompt tokens" 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 758ffd4dcf..31b74bdac2 100644 --- a/python/freetoken/tokenizer/tokenize.py +++ b/python/freetoken/tokenizer/tokenize.py @@ -4,6 +4,7 @@ import json import os import threading +from copy import deepcopy from types import ModuleType from typing import TYPE_CHECKING, Any, List @@ -22,6 +23,7 @@ probe_thinking_profile, quantize_effort, ) +from .inline_system import has_inline_system, normalize_inline_system, probe_inline_system logger = init_logger(__name__) @@ -36,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" @@ -49,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 @@ -57,7 +95,9 @@ def __init__(self, tokenizer: PreTrainedTokenizerBase, mm_processor: MMProcessor self._effort_profile: EffortProfile | None = None self._thinking_profile: ThinkingProfile | None = None self._effort_lock = threading.Lock() - self._logged_effort_maps: set[tuple[Any, str | None]] = set() + self._logged_effort_maps: set[tuple[Any, str | int | None]] = set() + self._inline_system_lock = threading.Lock() + self._inline_system_modes: dict[str, str] = {} def tokenize(self, msgs: List[TokenizeMsg]) -> List[UserMsg]: results: List[UserMsg] = [] @@ -103,21 +143,43 @@ def render_prompt(self, msg: TokenizeMsg) -> str: validation, count_tokens) must quantize identically.""" if not isinstance(msg.text, list): return msg.text - return self._render( - msg.text, msg.tools, self._sanitize_effort(msg.chat_template_kwargs or {}) - ) + kwargs = self._sanitize_effort(msg.chat_template_kwargs or {}) + messages = msg.text + if msg.inline_system_policy is not None: + if msg.inline_system_policy not in ("auto", "preserve", "fold"): + raise ValueError(f"invalid inline system policy: {msg.inline_system_policy}") + if has_inline_system(messages): + mode = msg.inline_system_policy + if mode == "auto": + mode = self.inline_system_mode(msg.tools, kwargs) + messages = normalize_inline_system(messages, mode) + return self._render(messages, msg.tools, kwargs, images=msg.images) + + def inline_system_mode(self, tools: list[dict] | None, kwargs: dict[str, Any]) -> str: + key = json.dumps([tools, kwargs], sort_keys=True) + with self._inline_system_lock: + if key not in self._inline_system_modes: + mode, reason = probe_inline_system( + lambda messages: self._render(messages, deepcopy(tools), deepcopy(kwargs)), self.tokenizer + ) + if len(self._inline_system_modes) >= 128: + self._inline_system_modes.pop(next(iter(self._inline_system_modes))) + self._inline_system_modes[key] = mode + logger.info("Anthropic inline system policy: %s (%s)", mode, reason) + return self._inline_system_modes[key] def _render( self, messages: list[dict[str, Any]], tools: list[dict[str, Any]] | None, chat_template_kwargs: dict[str, Any], + *, images: list[bytes] | None = None, ) -> str: """Raw render, no effort sanitation — the probe needs unsupported values to actually reach the template so rejection is observable.""" if self._dsv4_encoder is not None: return _apply_dsv4_chat_encoder( - self._dsv4_encoder, messages, tools, chat_template_kwargs + self._dsv4_encoder, messages, tools, chat_template_kwargs, images=images ) # Broadcast the effort in every spelling the ecosystem's templates read # (muse-glimmer grades ``reasoning_strength``; Jinja ignores undeclared @@ -130,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) @@ -197,8 +266,12 @@ def _load_dsv4_encoder_if_needed(tokenizer: PreTrainedTokenizerBase) -> ModuleTy model_path = getattr(tokenizer, "name_or_path", None) or getattr(tokenizer, "_name_or_path", "") if not model_path: return None - encoder_path = os.path.join(str(model_path), "encoding", "encoding_dsv4.py") - if not os.path.isfile(encoder_path): + # DeepSeek-V4 ships encoding/encoding_dsv4.py, V4.1 encoding/encoding.py; same encode_messages contract + encoder_path = next( + (p for p in (os.path.join(str(model_path), "encoding", name) for name in ("encoding_dsv4.py", "encoding.py")) if os.path.isfile(p)), + None, + ) + if encoder_path is None: return None spec = importlib.util.spec_from_file_location("encoding_dsv4", encoder_path) if spec is None or spec.loader is None: @@ -215,9 +288,28 @@ def _apply_dsv4_chat_encoder( messages: list[dict], tools: list[dict] | None, chat_template_kwargs: dict, + *, images: list[bytes] | None = None, ) -> str: rendered_messages = [dict(message) for message in messages] + image_data = iter(images) if images is not None else None for message in rendered_messages: + if isinstance(message.get("content"), list): + content = [] + for part in message["content"]: + if isinstance(part, dict) and part.get("type") == "image": + # Prerender still has refs; the tokenizer worker receives fetched bytes. + # The checkpoint encoder validates sources even when rendering only text. + ref = part.get("freetoken_ref") + if image_data is not None: + try: + part = {"type": "image", "data": next(image_data)} + except StopIteration: + raise ValueError("image blocks outnumber fetched images") from None + elif ref is not None: + key = "data" if ref["kind"] == "b64" else "url" + part = {"type": "image", key: ref["data"]} + content.append(part) + message["content"] = content if message.get("tool_calls"): message["tool_calls"] = _dsv4_tool_calls(message["tool_calls"]) if tools: 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/scripts/ftw_hotfix.py b/scripts/ftw_hotfix.py index f9cb62d1a8..9640c7b7f7 100644 --- a/scripts/ftw_hotfix.py +++ b/scripts/ftw_hotfix.py @@ -263,7 +263,9 @@ def is_checkpoint_tower_name(name: str) -> bool: head = name.split(".") if head[0] == "model" and len(head) > 1: head = head[1:] - return "vision" in head[0] or "visual" in head[0] or head[0] in _TOWER_SEGMENTS + return ("vision" in head[0] or "visual" in head[0] or head[0] in _TOWER_SEGMENTS + or head[0] in ("aligner", "image_start", "image_end", "image_newline") + or name.endswith(".ffn.gate.bias_vl")) def read_tower(source: TensorSource, ftw_dir: str, checkpoint_names: list[str]) -> dict[str, torch.Tensor]: diff --git a/tests/README.md b/tests/README.md index 26dba63391..47c7455c35 100644 --- a/tests/README.md +++ b/tests/README.md @@ -53,12 +53,17 @@ checkpoint is set: | `FREETOKEN_GEMMA4_MODEL` | `models/test_gemma4_vision.py`, `tokenizer/test_mm_tokenize.py` — local Gemma-4 `gemma4` checkpoint (tower parity, boi/eoi expansion; the tokenizer test also takes the 12B `gemma4_unified` checkpoint) | | `FREETOKEN_GEMMA4_UNIFIED_MODEL` | `models/test_gemma4_vision.py` — local gemma-4-12B `gemma4_unified` checkpoint (embedder parity) | | `FREETOKEN_GLM53_MODEL` | `tokenizer/test_mm_tokenize.py` — local GLM-5.3-Flash checkpoint (image-token expansion) | -| `FREETOKEN_MUSE_MODEL` | `models/test_muse_glimmer_vision.py`, `tokenizer/test_mm_tokenize.py` — local Muse-Glimmer-30B checkpoint (tower parity, image_start/image_end expansion) | +| `FREETOKEN_DSV4_MODEL` | `tokenizer/test_tokenize.py` - local DSV4 checkpoint with its native `encoding/encoding_dsv4.py`; prefix stability and reminder retention, tokenizer only (no GPU or model weights loaded) | +| `FREETOKEN_INLINE_GLM_MODEL` | `tokenizer/test_tokenize.py` - local GLM-5.3 tokenizer/template; inline-system preserve policy, prefix stability and reasoning retention (no weights loaded) | +| `FREETOKEN_INLINE_QWEN_MODEL` | `tokenizer/test_tokenize.py` - local Qwen3.6 tokenizer/template; inline-system fold policy, prefix stability and reasoning retention (no weights loaded) | +| `FREETOKEN_MUSE_MODEL` | `models/test_muse_glimmer_vision.py`, `tokenizer/test_mm_tokenize.py`, `tokenizer/test_tokenize.py` - local Muse-Glimmer-30B checkpoint (vision tests; tokenizer-only inline-system regressions) | | `FREETOKEN_MINIMAX_M3_MODEL` | `tokenizer/test_mm_tokenize.py` — local MiniMax-M3 checkpoint (start/end expansion, token budget) | | `FREETOKEN_TEST_MOE_CACHE_SIZE` | `e2e/test_aime.py` — >0 switches to the offload MoE backend with this cache size | | `FREETOKEN_TEST_MEM_RATIO` | `e2e/test_aime.py` — offload-mode memory_ratio (default `0.9`) | | `FREETOKEN_REBUILD_TEST_MODEL` | `e2e/test_cache_rebuild.py` — a SMALL local model dir; boots a real server (falls back to `FREETOKEN_TEST_MODEL`) | | `FREETOKEN_GEMMA4_GGUF_GLOB` | `models/test_gemma4_gguf_rope.py` — glob matching a local gemma-4 GGUF file | +| `FREETOKEN_TEST_MODELS_ROOT` | `models/test_quant_config.py` — directory whose checkpoints the stored-tensor scheme scan checks (default `/mnt/nvme/models`) | +| `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/attention/test_dsv41_backend.py b/tests/attention/test_dsv41_backend.py new file mode 100644 index 0000000000..387742eb63 --- /dev/null +++ b/tests/attention/test_dsv41_backend.py @@ -0,0 +1,208 @@ +"""DSV41 attention backend: addressing and selection contracts, CPU-only (kernels live in tests/kernels). + +* prefill metadata: encoder segments, the decoder pass under exact vs bounded replay; +* window candidates with a floor (bounded replay truncates the window at the prompt's last window); +* decode snapshot staging and the layer-invariant ring context; +* the selection helpers' reshape, candidate gather and position-to-row mapping. +""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch + +from freetoken.core import Batch, Context, Req, SamplingParams, get_global_ctx, set_global_ctx +from freetoken.kvcache.dsv4.v41_cost_model import dsv41_pool_sizes +from freetoken.kvcache.dsv4.v41_pool import DSV41PagedKVCache +from freetoken.models.deepseek_v41.args import DeepseekV41Args + +P, MRR, DEVICE = 128, 4, torch.device("cpu") +RATIOS = (0, 2, 2, 1, 1) +SOURCES = (1, 3) + + +def _ctx(pool): + try: + ctx = get_global_ctx() + except AssertionError: + ctx = Context(page_size=P) + set_global_ctx(ctx) + ctx.kv_cache = pool + return ctx + + +def _stack(swa_decoder_replay="exact", num_pages=32, max_seq_len=8192): + from freetoken.attention.dsv41_sparse import DSV41SparseAttnBackend + + args = DeepseekV41Args(n_layers=len(RATIOS), head_dim=512, index_head_dim=128, window_size=P, compress_ratios=RATIOS, + kv_source_layers=SOURCES, swa_decoder_replay=swa_decoder_replay) + pool = DSV41PagedKVCache(dsv41_pool_sizes(num_pages + 1, args, 1.0, P), args, DEVICE, n_scratch=MRR + 1) + pool._init_paged_state(MRR, True) + pt = torch.zeros(MRR + 1, max_seq_len, dtype=torch.int32) + pt[MRR].fill_(num_pages * P) + pt[2, :300] = torch.arange(300, dtype=torch.int32) + for page in range(3): + pool.bind_window_pages(page * P, page * P) + pool.full_loc_map = pt + _ctx(pool) + backend = DSV41SparseAttnBackend(SimpleNamespace(dsv41_args=args)) + return backend, pool, pt + + +def _req(table_idx, cached_len, n_new, output_len=16, uid=0, prompt_len=0): + return Req(input_ids=torch.zeros(cached_len + n_new, dtype=torch.int32), table_idx=table_idx, cached_len=cached_len, + output_len=output_len, uid=uid, sampling_params=SamplingParams(), cache_handle=None, prompt_len=prompt_len) + + +def _prefill_batch(reqs): + batch = Batch(reqs=reqs, phase="prefill") + batch.padded_reqs = reqs + return batch + + +def test_prefill_metadata_exact_mode_runs_the_decoder_on_every_token(): + backend, pool, _ = _stack("exact") + assert pool.prefix_replay_tokens == 0 + batch = _prefill_batch([_req(0, 0, 300), _req(2, 256, 44, uid=1)]) + backend.prepare_metadata(batch) + md = batch.attn_metadata + assert [(s.offset, s.n, s.table_idx, s.start_pos, s.window_floor) for s in md.segments] == [(0, 300, 0, 0, 0), (300, 44, 2, 256, 0)] + assert md.decoder_segments is md.segments and md.decoder_rows is None + assert md.last_indices.tolist() == [299, 343] + + +def test_prefill_metadata_bounded_mode_runs_the_decoder_on_the_last_window(): + """The encoder runs every new token; the decoder runs each request's part of its prompt's last + window ``[L - P, L)``, floored at ``L - P``. A chunk ending before that window runs no decoder row + and its ``last_indices`` entry points at the placeholder row the model appends.""" + backend, pool, _ = _stack("bounded") + assert pool.prefix_replay_tokens == P and pool.sliding_window_size == P + cold = _req(0, 0, 300) # a whole prompt: L - P = 172 + early = _req(1, 256, 128, uid=1, prompt_len=600) # chunk [256, 384) of a 600-token prompt: before 472 + straddle = _req(2, 384, 128, uid=2, prompt_len=600) # chunk [384, 512): its rows [472, 512) are in the window + batch = _prefill_batch([cold, early, straddle]) + backend.prepare_metadata(batch) + md = batch.attn_metadata + assert [(s.offset, s.n, s.start_pos, s.window_floor) for s in md.segments] == [(0, 300, 0, 0), (300, 128, 256, 0), (428, 128, 384, 0)] + assert [(d.offset, d.n, d.table_idx, d.start_pos, d.window_floor) for d in md.decoder_segments] == [(0, 128, 0, 172, 172), (128, 40, 2, 472, 472)] + assert md.decoder_rows.tolist() == list(range(172, 300)) + list(range(428 + 88, 556)) + assert md.decoder_pad and md.last_indices.tolist() == [127, 168, 167] # 168: the placeholder after 168 decoder rows + + +def test_prefill_metadata_bounded_mode_without_any_decoder_row(): + backend, _, _ = _stack("bounded") + batch = _prefill_batch([_req(0, 0, 256, prompt_len=1024)]) + backend.prepare_metadata(batch) + md = batch.attn_metadata + assert md.decoder_segments == [] and md.decoder_rows.numel() == 0 + assert md.decoder_pad and md.last_indices.tolist() == [0] + + +def test_prefill_metadata_short_prompt_runs_the_decoder_everywhere(): + """A prompt shorter than the window floors the decoder's window at 0: it sees what exact mode sees.""" + backend, _, _ = _stack("bounded") + batch = _prefill_batch([_req(0, 0, 50)]) + backend.prepare_metadata(batch) + md = batch.attn_metadata + assert [(d.offset, d.n, d.start_pos, d.window_floor) for d in md.decoder_segments] == [(0, 50, 0, 0)] + assert not md.decoder_pad and md.last_indices.tolist() == [49] + + +def test_private_window_layers_address_per_request_rings(): + """Under bounded replay the decoder layers keep their window KV in per-request rings (slot = + table row * P + pos % P), never in the radix-shared window pages; the encoder layers stay on + the shared pool.""" + from freetoken.attention.dsv41_sparse import PrefillSegment + + backend, pool, _ = _stack("bounded") + assert pool.private_window_layer_ids == (3, 4) and not any(pool.is_private_window(l) for l in (0, 1, 2)) + assert pool.window_pool[3].shape[0] == (MRR + 1) * P and pool.window_pool[0].shape[0] == pool.sizes.n_win_slots + # prefill: a decoder segment's slots are its ring row; an encoder segment's are page-bound + seg = PrefillSegment(0, 44, 2, 256, window_floor=256) + assert backend.layer_window_slots_of(3, 2, 256, 300).tolist() == [2 * P + p % P for p in range(256, 300)] + assert backend.layer_window_slots_of(0, 2, 256, 300).tolist() == list(range(256, 300)) + g = backend.window_topk_prefill(seg, layer_id=3)[0] + assert g.shape == (44, P) and g[0, 0].item() == 2 * P + 0 and (g[0, 1:] == -1).all() # position 256 sees itself only (floor) + assert set(g[43].tolist()) - {-1} == {2 * P + p % P for p in range(256, 300)} + # decode: the private ring context keys on the batch row's page-table row + batch = _decode_batch([2], [300]) + backend.prepare_metadata(batch) + md = batch.attn_metadata + slots, topk = md.private_window_ctx(torch.tensor([300]), torch.arange(1)) + assert slots.tolist() == [2 * P + 300 % P] and topk.shape == (1, 1, P) + assert set(topk[0, 0].tolist()) == {2 * P + j for j in range(P)} # every ring index holds a position <= 300 + + +def test_window_candidates_honor_the_floor_and_the_retained_prefix(): + from freetoken.attention.dsv41_sparse import PrefillSegment + + backend, pool, _ = _stack() + # a cold segment at positions [0, 10): query p sees [0, p] + g = backend.window_topk_prefill(PrefillSegment(0, 10, 2, 0)) + assert g.shape == (1, 10, P) + assert g[0, 3, :4].tolist() == [0, 1, 2, 3] and (g[0, 3, 4:] == -1).all() + # an extend at [256, 300) sees back into the retained prefix, but not below the floor + g = backend.window_topk_prefill(PrefillSegment(0, 44, 2, 256, window_floor=200)) + row = g[0, 0] # query at 256 -> [200, 256] + assert row[0].item() == 200 and row[56].item() == 256 and (row[57:] == -1).all() + # without a floor the query at 256 sees [129, 256] + g = backend.window_topk_prefill(PrefillSegment(0, 44, 2, 256)) + assert g[0, 0, 0].item() == 129 and g[0, 0, 127].item() == 256 + + +def _decode_batch(rows, positions): + reqs = [_req(int(t), 0, 1, uid=i) for i, t in enumerate(rows)] + batch = Batch(reqs=reqs, phase="decode") + batch.padded_reqs = reqs + batch.active_table_idx = torch.tensor(rows, dtype=torch.int64) + batch.positions = torch.tensor(positions, dtype=torch.int64) + return batch + + +def test_decode_snapshot_and_ring_context(): + backend, pool, pt = _stack() + batch = _decode_batch([2, MRR], [259, 0]) + backend.prepare_metadata(batch) + md = batch.attn_metadata + assert md.full_snap is None + snap = md.full_snapshot() + assert torch.equal(snap[0, :300], torch.arange(300)) + pt[2, :10] = 7 # a later mutation must not reach the snapshot + assert torch.equal(snap[0, :300], torch.arange(300)) + pos, rows = batch.positions, torch.arange(2) + ws, prev, topk = md.window_ctx(pos, rows) + dummy_ws = pool.sizes.n_win_slots - P # the dummy row's full loc binds to the dummy window page + assert ws.tolist() == [259, dummy_ws] and prev.tolist() == [258, dummy_ws] + assert topk.shape == (2, 1, P) and (topk[0] >= 0).all() + assert set(topk[0, 0].tolist()) == set(range(259 - 127, 260)) + assert topk[1, 0, 0].item() == dummy_ws and (topk[1, 0, 1:] == -1).all() + # the indexer derives compressed rows from the live locs: prefill off the page table row, + # decode off the snapshot (position t -> full_loc(t * ratio) // ratio) + assert torch.equal(backend.locs_prefill(2, 20), pt[2:3, :20]) + assert snap[0, :16].tolist() == list(range(16)) # taken before the mutation above + # capture staging reuses one buffer + backend.init_capture_graph(512, [1, 2]) + backend.prepare_for_replay(batch) + assert batch.attn_metadata.full_snap.data_ptr() == backend.capture.full_snap.data_ptr() + assert batch.attn_metadata.stage_width == 512 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="the selection kernels need CUDA") +def test_topk_helpers(): + from freetoken.attention.dsv41_sparse import DSV41SparseAttnBackend as B + + scores = torch.tensor([[[0.1, 5.0, 3.0, 7.0, 9.0]]], device="cuda") # columns past live are never read + live = torch.tensor([[3]], device="cuda", dtype=torch.int32) + assert B.select_topk(scores, live, 4).tolist() == [[[0, 1, 2, -1]]] + assert B.select_topk(scores, live, 2).tolist() == [[[1, 2]]] + # a candidate list is a sorted valid prefix; empty slots score -inf + cand = torch.tensor([[[2, 7, 9, -1]]], device="cuda", dtype=torch.int32) + cscores = torch.tensor([[[1.0, 4.0, -torch.inf, -torch.inf]]], device="cuda") + assert B.select_topk_in_candidates(cscores, cand, 3).tolist() == [[[2, 7, -1]]] + # rows come from the live locs: compressed position p -> locs[p * ratio] // ratio + locs = torch.arange(20, 40, device="cuda", dtype=torch.int32).view(1, 20) + locs[0, 14:16] = -1 # position 7 (ratio 2) has slid out + rows = B.positions_to_rows(torch.tensor([[[2, 7, 9, -1]]], device="cuda", dtype=torch.int32), locs, 2) + assert rows.tolist() == [[[12, -1, 19, -1]]] 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_ftw_hotfix_vision.py b/tests/checkpoint/test_ftw_hotfix_vision.py index d4a3eac68e..eb0c592f35 100644 --- a/tests/checkpoint/test_ftw_hotfix_vision.py +++ b/tests/checkpoint/test_ftw_hotfix_vision.py @@ -14,7 +14,7 @@ # env var -> local checkpoint, as tests/README.md lists them _CHECKPOINTS = ("FREETOKEN_QWEN3VL_MODEL", "FREETOKEN_GEMMA4_MODEL", "FREETOKEN_GEMMA4_UNIFIED_MODEL", - "FREETOKEN_GLM53_MODEL", "FREETOKEN_MUSE_MODEL", "FREETOKEN_MINIMAX_M3_MODEL") + "FREETOKEN_GLM53_MODEL", "FREETOKEN_MUSE_MODEL", "FREETOKEN_MINIMAX_M3_MODEL", "FREETOKEN_DSV41_MODEL") @pytest.fixture(scope="module") @@ -34,6 +34,10 @@ def hotfix(): ("vision_tower.vision_model.embeddings.patch_embedding.weight", True), ("multi_modal_projector.linear_1.weight", True), ("patch_merge_mlp.linear_2.bias", True), + ("vision.blocks.0.norm1.weight", True), + ("aligner.w1.weight", True), + ("image_newline", True), + ("layers.3.ffn.gate.bias_vl", True), ("model.embed_audio.embedding_projection.weight", False), ("model.language_model.layers.0.self_attn.q_proj.weight", False), ("model.layers.3.mlp.experts.0.gate_proj.weight", False), 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/dsv4/test_dsv4_attn_backend.py b/tests/dsv4/test_dsv4_attn_backend.py index 8ffc6fd4d9..7bc72210f2 100644 --- a/tests/dsv4/test_dsv4_attn_backend.py +++ b/tests/dsv4/test_dsv4_attn_backend.py @@ -11,8 +11,8 @@ import torch from freetoken.core import Batch, Context, Req, SamplingParams, get_global_ctx, set_global_ctx -from freetoken.kvcache.dsv4_paged_pool import DSV4PagedKVCache -from freetoken.kvcache.dsv4_cost_model import dsv4_pool_sizes +from freetoken.kvcache.dsv4.v4_pool import DSV4PagedKVCache +from freetoken.kvcache.dsv4.v4_cost_model import dsv4_pool_sizes from freetoken.models.deepseek_v4.args import DeepseekV4Args P, MRR, DEVICE = 128, 4, torch.device("cpu") diff --git a/tests/engine/test_attention_backend_arch.py b/tests/engine/test_attention_backend_arch.py index 9271bd4254..44f9f87777 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 073112bb5a..69ab3f2c40 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,30 +371,71 @@ 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: (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), # the floor rule compute_cache_floors documents above. @@ -367,7 +465,7 @@ def _offload_engine_config(**overrides): num_layers=10, num_moe_layers=10, num_experts=8, - expert_quant="nvfp4", # quantized experts -> must resolve to an offload backend + expert_quant="nvfp4", moe_strategy="auto", ), ) @@ -389,7 +487,7 @@ def test_guard_raises_actionable_error_when_too_small(): assert "128" in msg and "moe-cache" in msg -def test_adjust_config_defaults_moe_cache_auto_for_auto_resolved_offload_backend(): +def test_adjust_config_defaults_moe_cache_auto_for_auto_resolved_offload_backend(monkeypatch): """Bare `ft serve `: no --moe-backend, no --moe-cache-* flags at all. args.py's parse-time default only fires when the backend is *already* @@ -400,9 +498,12 @@ def test_adjust_config_defaults_moe_cache_auto_for_auto_resolved_offload_backend moe_cache_auto=True, so _init_offload_moe_cache's _require_offload_cache_size guard is never reached with moe_cache_size still 0. """ + import freetoken.engine.engine as engine_module from freetoken.engine.engine import _adjust_config from freetoken.moe import is_offload_moe_strategy + # a unified-memory host (GB10) resolves auto to fused instead + monkeypatch.setattr(engine_module, "_is_unified_memory_gpu", lambda index=None: False) config = _offload_engine_config() _adjust_config(config) @@ -520,3 +621,19 @@ 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_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..a215961ef7 --- /dev/null +++ b/tests/engine/test_logits_processors.py @@ -0,0 +1,207 @@ +"""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, Sampler, apply_logits_processors + +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_presence_and_frequency_penalties_count_generated_tokens(): + logits = torch.zeros((2, V)) + plan = _plan( + [1], + hist_ids=_ids([[3, 3, 5]]), + presence=torch.tensor([[0.5]]), + frequency=torch.tensor([[0.25]]), + ) + out = apply_logits_processors(logits, plan, V) + assert torch.equal(out[0], logits[0]) # row outside the plan untouched + assert out[1, 3].item() == -(0.5 + 2 * 0.25) + assert out[1, 5].item() == -(0.5 + 0.25) + assert out[1, 0].item() == 0.0 + + +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]]) + plan = _plan( + [0], + hist_ids=_ids([[0, 1]]), # generated: 0 (positive), 1 (negative) + prompt_ids=_ids([[2, 3]]), # prompt: 2 (positive), 3 (negative) + repetition=torch.tensor([[2.0]]), + ) + out = apply_logits_processors(logits, plan, V) + 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 + + +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), + ) + + +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 plan.hist_ids[0].tolist()[:2] == [3, 3] # row 1's generated tokens + assert plan.prompt_ids[0].tolist()[:2] == [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() == -1.0 + ) # presence on the repeated 3 (repetition on 0 is a no-op) + 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_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_logits_processors(torch.ones(1, V), args.plan, V) + 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 diff --git a/tests/engine/test_sample_logprobs.py b/tests/engine/test_sample_logprobs.py new file mode 100644 index 0000000000..6bd204604d --- /dev/null +++ b/tests/engine/test_sample_logprobs.py @@ -0,0 +1,167 @@ +from types import SimpleNamespace + +import torch + +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=SimpleNamespace(needs_logits_processing=False, logprobs=True, top_logprobs=17, is_greedy=True)), + SimpleNamespace(sampling_params=SimpleNamespace(needs_logits_processing=False, logprobs=False, top_logprobs=0, is_greedy=True)), + ] + ) + 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_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_dsv41_indexer.py b/tests/kernels/test_dsv41_indexer.py new file mode 100644 index 0000000000..19b45ab8b9 --- /dev/null +++ b/tests/kernels/test_dsv41_indexer.py @@ -0,0 +1,139 @@ +"""DSV41 indexer logits over packed fp4 index keys: full-range (Full mode) and candidate-restricted +(Reindex mode) against a torch reference on the dequantized keys, with causality from ``live``. The +row of a compressed position comes from the request's full-token locs (``locs[t * ratio] // ratio``); +in Full mode only the live columns are written.""" + +from __future__ import annotations + +import pytest +import torch + +from freetoken.kvcache.dsv4.v41_row_format import FP4_E8M0_B32 + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") + +D, H = 128, 32 +ROWS = 4200 + + +RATIO = 2 + + +def _setup(b, s, t, seed): + from freetoken.kernel.triton.dsv41.pack import pack_rows, unpack_rows + + torch.manual_seed(seed) + keys = torch.randn(ROWS, D, device="cuda", dtype=torch.bfloat16) + pool = pack_rows(keys, FP4_E8M0_B32) + deq = unpack_rows(pool, FP4_E8M0_B32, D) + q = torch.randn(b, s, H, D, device="cuda", dtype=torch.bfloat16) + w = (torch.randn(b, s, H, device="cuda") * 0.3).bfloat16() + # each request owns a random permutation of pool rows for its compressed positions; a few holes. + # The kernel derives rows from full-token locs: compressed position p -> loc = row * RATIO (+ any + # in-group offset) at locs[p * RATIO]; a hole is a negative loc. + k_rows = torch.stack([torch.randperm(ROWS, device="cuda")[:t] for _ in range(b)]).to(torch.int32) + k_rows[:, 5::11] = -1 + locs = torch.full((b, t * RATIO), -1, device="cuda", dtype=torch.int32) + for i in range(b): + for p in range(t): + r = int(k_rows[i, p]) + if r >= 0: + locs[i, p * RATIO : (p + 1) * RATIO] = torch.arange(r * RATIO, (r + 1) * RATIO, device="cuda", dtype=torch.int32) + return pool, deq, q, w, k_rows, locs + + +def _reference(q, w, deq, k_rows, live, positions): + """positions: [B, S, N] compressed positions (-1 empty) -> logits [B, S, N].""" + b, s, n = positions.shape + out = torch.full((b, s, n), float("-inf"), device="cuda") + for i in range(b): + for j in range(s): + for c in range(n): + p = int(positions[i, j, c]) + if p < 0 or p >= int(live[i, j]): + continue + r = int(k_rows[i, p]) + if r < 0: + continue + score = (q[i, j] @ deq[r]).relu() * w[i, j] + out[i, j, c] = score.sum() + return out + + +def test_full_range_matches_reference_with_causal_live(): + from freetoken.kernel.triton.dsv41.indexer import indexer_logits_packed + + b, s, t = 2, 5, 300 + pool, deq, q, w, k_rows, locs = _setup(b, s, t, seed=1) + live = torch.tensor([[10, 64, 65, 200, 300], [0, 1, 128, 129, 250]], device="cuda", dtype=torch.int32) + out = torch.full((b, s, t), 12345.0, device="cuda") # a sentinel: the dead area must stay untouched + got = indexer_logits_packed(q, w, pool, FP4_E8M0_B32, locs, RATIO, live, out=out) + positions = torch.arange(t, device="cuda").view(1, 1, t).expand(b, s, t) + want = _reference(q, w, deq, k_rows, live, positions) + assert got.shape == (b, s, t) + written = positions < live.unsqueeze(-1) + touched = positions < (torch.div(live + 63, 64, rounding_mode="floor") * 64).unsqueeze(-1) # whole live tiles + assert (got[~touched] == 12345.0).all() and torch.isneginf(got[touched & ~written]).all() + assert torch.equal(torch.isinf(got[written]), torch.isinf(want[written])) + finite = written & ~torch.isinf(want) + torch.testing.assert_close(got[finite], want[finite], atol=2e-2, rtol=2e-2) + + +def test_quantized_queries_preserve_reference_score_rounding(): + from freetoken.kernel.triton.dsv41.indexer import indexer_logits_packed + from freetoken.kernel.triton.dsv41.pack import pack_rows, unpack_rows + + torch.manual_seed(23) + t, s, topk = 4096, 32, 512 + pool = pack_rows(torch.randn(t, D, device="cuda", dtype=torch.bfloat16), FP4_E8M0_B32) + keys = unpack_rows(pool, FP4_E8M0_B32, D) + q = torch.randn(s * H, D, device="cuda", dtype=torch.bfloat16) + q = unpack_rows(pack_rows(q, FP4_E8M0_B32), FP4_E8M0_B32, D).view(1, s, H, D) + w = torch.randn(1, s, H, device="cuda", dtype=torch.bfloat16) * (D * H) ** -0.5 + locs = torch.arange(t, device="cuda", dtype=torch.int32).view(1, t) + live = torch.full((1, s), t, device="cuda", dtype=torch.int32) + got = indexer_logits_packed(q, w, pool, FP4_E8M0_B32, locs, 1, live) + reference = torch.einsum("bshd,btd->bsht", q, keys.unsqueeze(0)) + reference = (reference.relu() * w.unsqueeze(-1)).sum(dim=2).float() + torch.testing.assert_close(got, reference, atol=0, rtol=0) + # BF16 scores can tie at the cutoff; compare selected scores, not arbitrary tie order. + torch.testing.assert_close(got.topk(topk).values, reference.topk(topk).values, atol=0, rtol=0) + + +def test_full_range_worker_grid_covers_wide_histories(monkeypatch): + """More live tiles than worker programs: each worker strides over its share.""" + from freetoken.kernel.triton.dsv41 import indexer as mod + + b, s, t = 1, 2, 4096 # 64 tiles + pool, deq, q, w, k_rows, locs = _setup(b, s, t, seed=4) + live = torch.tensor([[4096, 4000]], device="cuda", dtype=torch.int32) + from types import SimpleNamespace + + properties = torch.cuda.get_device_properties(q.device) + with monkeypatch.context() as patch: + patch.setattr(torch.cuda, "get_device_properties", lambda _: SimpleNamespace( + multi_processor_count=5, major=properties.major, minor=properties.minor, + )) + got = mod.indexer_logits_packed(q, w, pool, FP4_E8M0_B32, locs, RATIO, live) + positions = torch.arange(t, device="cuda").view(1, 1, t).expand(b, s, t) + want = _reference(q, w, deq, k_rows, live, positions) + written = positions < live.unsqueeze(-1) + assert torch.equal(torch.isinf(got[written]), torch.isinf(want[written])) + finite = written & ~torch.isinf(want) + torch.testing.assert_close(got[finite], want[finite], atol=2e-2, rtol=2e-2) + + +def test_candidate_pool_matches_reference(): + from freetoken.kernel.triton.dsv41.indexer import indexer_logits_packed + + b, s, t, nc = 2, 3, 400, 96 + pool, deq, q, w, k_rows, locs = _setup(b, s, t, seed=2) + live = torch.tensor([[400, 350, 33], [1, 400, 200]], device="cuda", dtype=torch.int32) + cand = torch.randint(0, t, (b, s, nc), device="cuda", dtype=torch.int32) + cand[..., -10:] = -1 # empty tail of the pool + got = indexer_logits_packed(q, w, pool, FP4_E8M0_B32, locs, RATIO, live, candidates=cand) + want = _reference(q, w, deq, k_rows, live, cand) + assert got.shape == (b, s, nc) + assert torch.equal(torch.isinf(got), torch.isinf(want)) + finite = ~torch.isinf(want) + torch.testing.assert_close(got[finite], want[finite], atol=2e-2, rtol=2e-2) diff --git a/tests/kernels/test_dsv41_pack.py b/tests/kernels/test_dsv41_pack.py new file mode 100644 index 0000000000..176889dc7f --- /dev/null +++ b/tests/kernels/test_dsv41_pack.py @@ -0,0 +1,112 @@ +"""Packed KV row formats (``kvcache/dsv4/v41_row_format.py`` + ``kernel/triton/dsv41/pack.py``) against a torch +transcription of the reference quantizers in DeepSeek-V4.1's ``inference/kernel.py``: the round-tripped +value the pool hands back must equal what the reference bakes into its bf16 cache, bit for bit.""" + +from __future__ import annotations + +import pytest +import torch + +from freetoken.kvcache.dsv4.v41_row_format import BF16, FP4_E4M3_B16, FP4_E8M0_B32, FP8_E8M0_B32, RowFormat + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") + +FP4_GRID = torch.tensor([0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0]) + + +def _round_fp4(x: torch.Tensor) -> torch.Tensor: + """Round-to-nearest-even onto the signed e2m1 grid (the ``float4_e2m1fn`` cast).""" + a = x.abs() + grid = FP4_GRID.to(x.device) + # nearest grid point; ties go to the even-mantissa neighbour (0, 1, 2, 4 are even) + dist = (a[..., None] - grid).abs() + lo = torch.searchsorted(grid, a.reshape(-1).contiguous(), right=True).reshape(a.shape).clamp(1, 7) + below, above = grid[lo - 1], grid[lo] + mid = (below + above) / 2 + even_above = torch.isin(above, torch.tensor([0.0, 1.0, 2.0, 4.0], device=x.device)) + pick_above = (a > mid) | ((a == mid) & even_above) + r = torch.where(pick_above, above, below) + r = torch.where(a >= 6.0, torch.full_like(r, 6.0), r) + del dist + return torch.copysign(r, x) + + +def reference_roundtrip(x: torch.Tensor, fmt: RowFormat) -> torch.Tensor: + """``act_quant`` / ``fp4_act_quant`` with ``inplace=True``: quant + dequant back to bf16.""" + m, d = x.shape + if fmt is BF16: + return x.to(torch.bfloat16) + g = x.float().view(m, d // fmt.group, fmt.group) + amax = g.abs().amax(dim=-1) + if fmt is FP8_E8M0_B32: + s = torch.exp2(torch.ceil(torch.log2(amax.clamp_min(1e-4) / 448.0))) + q = (g / s[..., None]).clamp(-448.0, 448.0).to(torch.float8_e4m3fn).float() + elif fmt is FP4_E8M0_B32: + s = torch.exp2(torch.ceil(torch.log2(amax.clamp_min(6.0 * 2.0**-126) / 6.0))) + q = _round_fp4((g / s[..., None]).clamp(-6.0, 6.0)) + elif fmt is FP4_E4M3_B16: + s = (amax.clamp_min(6.0 * 2.0**-9) / 6.0).to(torch.float8_e4m3fn).float() + q = _round_fp4((g / s[..., None]).clamp(-6.0, 6.0)) + else: + raise AssertionError(fmt) + return (q * s[..., None]).view(m, d).to(torch.bfloat16) + + +@pytest.mark.parametrize("fmt", [BF16, FP8_E8M0_B32, FP4_E4M3_B16, FP4_E8M0_B32], ids=lambda f: f.name) +@pytest.mark.parametrize("dim", [128, 512]) +def test_roundtrip_matches_the_reference_quantizer(fmt: RowFormat, dim: int): + from freetoken.kernel.triton.dsv41.pack import pack_rows, unpack_rows + + torch.manual_seed(0) + m = 300 + # mixed magnitudes, exact zeros, a tiny group and a huge one + x = torch.randn(m, dim, device="cuda") * torch.exp2(torch.randint(-8, 6, (m, 1), device="cuda").float()) + x[7] = 0 + x[11, : fmt.group or 32] *= 1e-6 + x[13, -32:] = 400.0 + x = x.to(torch.bfloat16) + + packed = pack_rows(x, fmt) + assert packed.shape == (m, fmt.row_bytes(dim)) and packed.dtype == torch.uint8 + got = unpack_rows(packed, fmt, dim) + want = reference_roundtrip(x, fmt) + assert torch.equal(got, want), f"{fmt.name}: {(got.float() - want.float()).abs().max().item()}" + + +def test_scatter_into_pool_rows_and_gather_back(): + from freetoken.kernel.triton.dsv41.pack import pack_rows, unpack_rows + + torch.manual_seed(1) + fmt, dim, rows = FP4_E4M3_B16, 512, 64 + pool = torch.zeros((rows, fmt.row_bytes(dim)), dtype=torch.uint8, device="cuda") + x = torch.randn(5, dim, device="cuda", dtype=torch.bfloat16) + ids = torch.tensor([3, 60, 0, 17, 42], device="cuda", dtype=torch.int64) + assert pack_rows(x, fmt, pool=pool, row_ids=ids) is pool + got = unpack_rows(pool, fmt, dim, row_ids=ids.to(torch.int32)) + assert torch.equal(got, reference_roundtrip(x, fmt)) + untouched = torch.ones(rows, dtype=torch.bool, device="cuda") + untouched[ids] = False + assert int(pool[untouched].sum()) == 0 + + +def test_e4m3_encode_helper_matches_torch_bits(): + """``e4m3_f32_to_u8`` (the pre-sm_89 scale encoder) on every finite e4m3 code.""" + import triton + import triton.language as tl + + from freetoken.kernel.triton.e4m3_compat import e4m3_f32_to_u8 + + @triton.jit + def _enc(x_ptr, o_ptr, N: tl.constexpr): + offs = tl.arange(0, N) + tl.store(o_ptr + offs, e4m3_f32_to_u8(tl.load(x_ptr + offs))) + + codes = torch.arange(256, dtype=torch.uint8, device="cuda") + finite = (codes & 0x7F) != 0x7F # skip the NaN codes + vals = codes.view(torch.float8_e4m3fn).float() + vals = torch.where(finite, vals, torch.zeros_like(vals)) + out = torch.empty(256, dtype=torch.uint8, device="cuda") + _enc[(1,)](vals, out, N=256) + want = torch.where(finite, codes, torch.zeros_like(codes)) + # -0.0 encodes to 0x80 in torch and 0x80 here; +0.0 -> 0x00 + assert torch.equal(out, want) diff --git a/tests/kernels/test_dsv41_sparse_attn.py b/tests/kernels/test_dsv41_sparse_attn.py new file mode 100644 index 0000000000..8a4f5f2a4f --- /dev/null +++ b/tests/kernels/test_dsv41_sparse_attn.py @@ -0,0 +1,107 @@ +"""DSV41 sparse attention over packed pools against a torch reference on the DEQUANTIZED pools. + +``m == 1`` takes the split-k decode path, ``m > 1`` the single-program prefill one. Tolerance follows +the dsv4 sparse-attention tests (5e-2 on fp32-cast bf16): the online-softmax accumulation order differs +from the reference and split-k merges through log-sum-exp on top of that.""" + +from __future__ import annotations + +import pytest +import torch + +from freetoken.kvcache.dsv4.v41_row_format import BF16, FP4_E4M3_B16, FP8_E8M0_B32 + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") + +D, H = 128, 8 +N_WINDOW = 128 # a multiple of BLOCK_T +N_WIN_SLOTS, N_CMP = 512, 1024 +TOL = dict(atol=5e-2, rtol=5e-2) + + +def _reference(q, win, cmp, sink, idx, n_window, scale, counts): + """Per (request, query, head) softmax over the live columns; ``sink`` is a null key.""" + b, m, h, d = q.shape + out = torch.zeros_like(q, dtype=torch.float32) + for i in range(b): + for j in range(m): + width = idx.shape[-1] if counts is None else n_window + int(counts[i, j]) + kv = [] + for c in range(min(width, idx.shape[-1])): + slot = int(idx[i, j, c]) + if slot >= 0: + kv.append((win if c < n_window else cmp)[slot].float()) + if not kv: + continue + kv_t = torch.stack(kv) + logits = q[i, j].float() @ kv_t.T * scale + mx = torch.maximum(logits.max(dim=-1).values, sink.float()) + probs = torch.exp(logits - mx[:, None]) + denom = probs.sum(dim=-1) + torch.exp(sink.float() - mx) + out[i, j] = (probs @ kv_t) / denom[:, None] + return out + + +def _pools(win_fmt, cmp_fmt, seed): + from freetoken.kernel.triton.dsv41.pack import pack_rows, unpack_rows + + torch.manual_seed(seed) + win_vals = torch.randn(N_WIN_SLOTS, D, device="cuda", dtype=torch.bfloat16) + cmp_vals = torch.randn(N_CMP, D, device="cuda", dtype=torch.bfloat16) + win_pool = pack_rows(win_vals, win_fmt) + cmp_pool = pack_rows(cmp_vals, cmp_fmt) + return win_pool, cmp_pool, unpack_rows(win_pool, win_fmt, D), unpack_rows(cmp_pool, cmp_fmt, D) + + +def _topk(b, m, n_cmp_cols, seed, ring_fill=None): + torch.manual_seed(seed) + win = torch.randint(0, N_WIN_SLOTS, (b, m, N_WINDOW), device="cuda", dtype=torch.int32) + if ring_fill is not None: # a ring still filling: trailing window columns empty + win[..., ring_fill:] = -1 + cmp = torch.randint(0, N_CMP, (b, m, n_cmp_cols), device="cuda", dtype=torch.int32) + cmp[..., ::7] = -1 # scattered unreachable picks + return torch.cat([win, cmp], dim=-1) + + +@pytest.mark.parametrize("m", [1, 3]) +@pytest.mark.parametrize("fmts", [(FP8_E8M0_B32, FP4_E4M3_B16), (BF16, BF16)], ids=["fp8+fp4", "bf16"]) +def test_matches_reference_on_dequantized_pools(m, fmts): + from freetoken.kernel.triton.dsv41.sparse_attn import sparse_attn_packed + + win_fmt, cmp_fmt = fmts + b, n_cmp_cols = 2, 256 + win_pool, cmp_pool, win_deq, cmp_deq = _pools(win_fmt, cmp_fmt, seed=1) + q = torch.randn(b, m, H, D, device="cuda", dtype=torch.bfloat16) + sink = torch.randn(H, device="cuda") * 0.5 + idx = _topk(b, m, n_cmp_cols, seed=2, ring_fill=100) + scale = D**-0.5 + + got = sparse_attn_packed(q, win_pool, win_fmt, cmp_pool, cmp_fmt, sink, idx, N_WINDOW, scale) + want = _reference(q, win_deq, cmp_deq, sink, idx, N_WINDOW, scale, None) + torch.testing.assert_close(got.float(), want, **TOL) + + +def test_cmp_counts_bound_the_visited_columns(): + from freetoken.kernel.triton.dsv41.sparse_attn import sparse_attn_packed + + b, m, n_cmp_cols = 3, 1, 512 + win_pool, cmp_pool, win_deq, cmp_deq = _pools(FP8_E8M0_B32, FP4_E4M3_B16, seed=3) + q = torch.randn(b, m, H, D, device="cuda", dtype=torch.bfloat16) + sink = torch.zeros(H, device="cuda") + idx = _topk(b, m, n_cmp_cols, seed=4) + counts = torch.tensor([[0], [37], [512]], device="cuda", dtype=torch.int32) + scale = D**-0.5 + + got = sparse_attn_packed(q, win_pool, FP8_E8M0_B32, cmp_pool, FP4_E4M3_B16, sink, idx, N_WINDOW, scale, cmp_counts=counts) + want = _reference(q, win_deq, cmp_deq, sink, idx, N_WINDOW, scale, counts) + torch.testing.assert_close(got.float(), want, **TOL) + + +def test_all_masked_query_yields_zeros(): + from freetoken.kernel.triton.dsv41.sparse_attn import sparse_attn_packed + + win_pool, cmp_pool, _, _ = _pools(FP8_E8M0_B32, FP4_E4M3_B16, seed=5) + q = torch.randn(1, 2, H, D, device="cuda", dtype=torch.bfloat16) + idx = torch.full((1, 2, N_WINDOW + 64), -1, device="cuda", dtype=torch.int32) + out = sparse_attn_packed(q, win_pool, FP8_E8M0_B32, cmp_pool, FP4_E4M3_B16, torch.zeros(H, device="cuda"), idx, N_WINDOW, 0.1) + assert torch.equal(out, torch.zeros_like(out)) diff --git a/tests/kernels/test_dsv41_topk.py b/tests/kernels/test_dsv41_topk.py new file mode 100644 index 0000000000..e7fc66f8d8 --- /dev/null +++ b/tests/kernels/test_dsv41_topk.py @@ -0,0 +1,159 @@ +"""dsv41_topk / dsv41_candidate_blocks against an independent torch reference with the documented tie +rule, over the contract's edge cases: empty and short histories, unequal live counts in one batch, +-inf holes, widths crossing the slice and level boundaries, exact ties, newest-block inclusion, +the sorted-valid-prefix / -1-tail output layout, and CUDA-graph replay across changing live lengths +(stale workspace reads would surface as wrong picks).""" + +from __future__ import annotations + +import pytest +import torch + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") + + +def ref_topk(scores: torch.Tensor, live: torch.Tensor, k: int) -> torch.Tensor: + """(score desc, column asc) over the live, finite columns; ascending columns, -1 padded.""" + R, T = scores.shape + out = torch.full((R, k), -1, dtype=torch.int32) + for r in range(R): + n = int(live[r]) + s = scores[r, :n].double().cpu() + cols = torch.arange(n) + alive = torch.isfinite(s) | (s == float("inf")) + cand = [(float(-s[c]), int(c)) for c in cols[alive]] + cand.sort() + picked = sorted(c for _, c in cand[:k]) + out[r, : len(picked)] = torch.tensor(picked, dtype=torch.int32) + return out.to(scores.device) + + +def ref_candidates(scores, live, topk_blocks, block_size): + R, T = scores.shape + nb = -(-T // block_size) + out = torch.full((R, topk_blocks * block_size), -1, dtype=torch.int32) + for r in range(R): + n = int(live[r]) + if n == 0: + continue + nbl = -(-n // block_size) + s = torch.nn.functional.pad(scores[r, :n].float().cpu(), (0, nbl * block_size - n), value=float("-inf")).view(nbl, block_size).amax(-1) + s[nbl - 1] = float("inf") + cand = sorted((float(-s[b]), b) for b in range(nbl) if s[b] > float("-inf")) + keep = sorted(b for _, b in cand[:topk_blocks]) + pos = [p for b in keep for p in range(b * block_size, (b + 1) * block_size) if p < n] + out[r, : len(pos)] = torch.tensor(pos, dtype=torch.int32) + return out.to(scores.device) + + +def _scores(R, T, seed=0, ties=False): + g = torch.Generator(device="cpu").manual_seed(seed) + s = torch.randn(R, T, generator=g) + if ties: + s = torch.randint(0, 5, (R, T), generator=g).float() # heavy exact ties + return s.cuda() + + +@pytest.mark.parametrize("T,k", [(64, 8), (4096, 512), (4097, 512), (8192, 512), (100_000, 512), (262_144, 512), (1_048_576, 512), (131_072, 2048)]) +def test_topk_matches_reference_across_widths(T, k): + from freetoken.kernel.triton.dsv41.topk import dsv41_topk + + R = 3 + scores = _scores(R, T, seed=T) + live = torch.tensor([0, min(T, 5000), T], dtype=torch.int32, device="cuda")[:R] + # holes inside the live range + scores[2, ::7] = float("-inf") + got = dsv41_topk(scores, live, k) + want = ref_topk(scores, live, k) + assert torch.equal(got, want), (got[:, :8], want[:, :8]) + + +def test_topk_short_and_unequal_rows_and_exact_ties(): + from freetoken.kernel.triton.dsv41.topk import dsv41_topk + + T, k = 20_000, 512 + scores = _scores(5, T, seed=3, ties=True) + live = torch.tensor([0, 1, 300, 4096, 20_000], dtype=torch.int32, device="cuda") + got = dsv41_topk(scores, live, k) + want = ref_topk(scores, live, k) + assert torch.equal(got, want) + # the layout contract: a sorted valid prefix, then -1 + for r in range(5): + row = got[r].tolist() + n = sum(1 for v in row if v >= 0) + assert row[n:] == [-1] * (k - n) and row[:n] == sorted(row[:n]) and n == min(int(live[r]), k) + + +def test_topk_all_dead_row_and_fewer_finite_than_k(): + from freetoken.kernel.triton.dsv41.topk import dsv41_topk + + scores = torch.full((2, 10_000), float("-inf"), device="cuda") + scores[1, [10, 4095, 4096, 9000]] = torch.tensor([1.0, 2.0, 3.0, 0.5], device="cuda") + live = torch.tensor([10_000, 9001], dtype=torch.int32, device="cuda") + got = dsv41_topk(scores, live, 8) + assert (got[0] == -1).all() + assert got[1].tolist() == [10, 4095, 4096, 9000, -1, -1, -1, -1] + + +@pytest.mark.parametrize("T", [4000, 200_000]) +def test_candidate_blocks_match_reference_and_pin_the_newest_block(T): + from freetoken.kernel.triton.dsv41.topk import dsv41_candidate_blocks + + R, KB, BS = 4, 2048, 8 + scores = _scores(R, T, seed=T + 1) + scores[0, :] = -1000.0 # the newest block is kept even when every score is poor + live = torch.tensor([13, 8, min(T, 16_000), T], dtype=torch.int32, device="cuda") + got = dsv41_candidate_blocks(scores, live, KB, BS) + want = ref_candidates(scores, live, KB, BS) + assert got.shape == (R, KB * BS) + assert torch.equal(got, want) + # short histories: every live position is a candidate (the implicit fast path) + assert got[0].tolist()[:13] == list(range(13)) and (got[0, 13:] == -1).all() + assert got[1].tolist()[:8] == list(range(8)) and (got[1, 8:] == -1).all() + + +def test_topk_under_graph_replay_across_live_lengths(): + """Capture once at the staged width, replay with live lengths that move up and down: every replay + must equal the reference (a stale candidate slot or live count from a previous replay would not).""" + from freetoken.kernel.triton.dsv41.topk import dsv41_candidate_blocks, dsv41_topk + + T, k, R = 300_000, 512, 2 + scores = torch.empty((R, T), device="cuda") + live = torch.empty((R,), dtype=torch.int32, device="cuda") + out = torch.empty((R, k), dtype=torch.int32, device="cuda") + cand = torch.empty((R, 2048 * 8), dtype=torch.int32, device="cuda") + stream = torch.cuda.Stream() + with torch.cuda.stream(stream): + scores.copy_(_scores(R, T, seed=9)) + live.fill_(T) + dsv41_topk(scores, live, k, out) # warm-up compiles + cand.copy_(dsv41_candidate_blocks(scores, live, 2048, 8)) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + dsv41_topk(scores, live, k, out) + cand.copy_(dsv41_candidate_blocks(scores, live, 2048, 8)) + torch.cuda.synchronize() + for step, (l0, l1) in enumerate([(T, 7), (4096, 4097), (0, 1), (250_000, 12_000), (3, T), (16_384, 16_385)]): + scores.copy_(_scores(R, T, seed=100 + step)) + live.copy_(torch.tensor([l0, l1], dtype=torch.int32)) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(out, ref_topk(scores, live, k)), (step, l0, l1) + assert torch.equal(cand, ref_candidates(scores, live, 2048, 8)), (step, l0, l1) + + +def test_topk_plan_terminates_or_rejects(): + """Every accepted k converges (each level at least halves the row); k past CHUNK // 2 -- which + would stall the planner (k == CHUNK never shrinks, 3072 stalls at 6144) -- is rejected up front.""" + from freetoken.kernel.triton.dsv41.topk import CHUNK, FINAL_MAX, MAX_K, topk_plan + + for k in (1, 8, 512, 2048, MAX_K): + for columns in (1, 4096, 4097, 1 << 20, (1 << 20) + 3): + levels = topk_plan(columns, k) + widths = [w for w, _ in levels] + [(levels[-1][1] * k) if levels else columns] + assert all(a > b for a, b in zip(widths, widths[1:])) and widths[-1] <= FINAL_MAX + assert len(levels) <= 12 + for k in (MAX_K + 1, 2049, 3072, CHUNK, 0, -1): + with pytest.raises(ValueError): + topk_plan(1 << 20, k) diff --git a/tests/kernels/test_dsv4_fp8_block_linear.py b/tests/kernels/test_dsv4_fp8_block_linear.py new file mode 100644 index 0000000000..2194a65e1b --- /dev/null +++ b/tests/kernels/test_dsv4_fp8_block_linear.py @@ -0,0 +1,101 @@ +"""``kernel/triton/dsv4/fp8_linear.block_fp8_linear`` against a torch transcription of the reference +``act_quant`` + ``fp8_gemm`` (ue8m0 activation scales, ``block x block`` e8m0 weight scales), for the +DeepSeek-V4 block (128) and the DeepSeek-V4.1 block (32), on the GEMV (M=1) and GEMM (M>1) paths.""" + +from __future__ import annotations + +import pytest +import torch + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") + +FP8 = torch.float8_e4m3fn +E8M0 = torch.float8_e8m0fnu + + +def _pow2_ceil_scale(amax: torch.Tensor, qmax: float) -> torch.Tensor: + """``2 ** ceil(log2(amax / qmax))`` -- the reference ``fast_round_scale``.""" + return torch.exp2(torch.ceil(torch.log2(amax / qmax))) + + +def quantize_weight(w: torch.Tensor, block: int) -> tuple[torch.Tensor, torch.Tensor]: + """Per ``block x block`` ue8m0-scaled e4m3 weight, scale as e8m0 codes.""" + n, k = w.shape + blocks = w.float().view(n // block, block, k // block, block).permute(0, 2, 1, 3) + amax = blocks.abs().amax(dim=(-1, -2)).clamp_min(1e-4) + scale = _pow2_ceil_scale(amax, 448.0) # [n//block, k//block] + q = (blocks / scale[..., None, None]).clamp(-448.0, 448.0).to(FP8) + w_fp8 = q.permute(0, 2, 1, 3).reshape(n, k) + return w_fp8, scale.to(E8M0) + + +def act_quant_reference(x: torch.Tensor, block: int) -> tuple[torch.Tensor, torch.Tensor]: + m, k = x.shape + groups = x.float().view(m, k // block, block) + amax = groups.abs().amax(dim=-1).clamp_min(1e-4) + scale = _pow2_ceil_scale(amax, 448.0) + q = (groups / scale[..., None]).clamp(-448.0, 448.0).to(FP8) + return q.view(m, k), scale + + +def linear_reference(x: torch.Tensor, w_fp8: torch.Tensor, w_scale: torch.Tensor, block: int) -> torch.Tensor: + m, k = x.shape + n = w_fp8.shape[0] + a_fp8, a_scale = act_quant_reference(x, block) + a = a_fp8.float().view(m, k // block, block) + w = w_fp8.float().view(n // block, block, k // block, block) + ws = w_scale.float() # [n//block, k//block] + out = torch.zeros(m, n, dtype=torch.float32, device=x.device) + for kb in range(k // block): + wk = w[:, :, kb, :].reshape(n, block) # [n, block] + partial = a[:, kb, :] @ wk.t() # [m, n] + out += partial * a_scale[:, kb, None] * ws[:, kb].repeat_interleave(block)[None, :] + return out + + +@pytest.mark.parametrize("block", [32, 128]) +@pytest.mark.parametrize("m", [1, 2, 5, 8, 9, 64]) +def test_block_fp8_linear_matches_reference(block: int, m: int): + from freetoken.kernel.triton.dsv4.fp8_linear import block_fp8_linear + + torch.manual_seed(0) + # K // block = 5 slabs exercises the GEMV split that must divide the slab count; N spans two GEMM N-tiles + k, n = 5 * block, 256 if block == 128 else 224 + x = torch.randn(m, k, device="cuda", dtype=torch.bfloat16) * 3 + w_fp8, w_scale = quantize_weight(torch.randn(n, k, device="cuda", dtype=torch.bfloat16), block) + + got = block_fp8_linear(x, w_fp8, w_scale, block=block).float() + want = linear_reference(x, w_fp8, w_scale, block) + # bf16 output; the fp32 reduction order inside the MMA may differ from the slab loop + torch.testing.assert_close(got, want, rtol=2e-2, atol=2e-2 * want.abs().max().item()) + + +def test_block_argument_must_match_the_scale_shape(): + from freetoken.kernel.triton.dsv4.fp8_linear import block_fp8_linear + + x = torch.randn(4, 256, device="cuda", dtype=torch.bfloat16) + w_fp8, w_scale = quantize_weight(torch.randn(256, 256, device="cuda", dtype=torch.bfloat16), 32) + with pytest.raises(AssertionError): + block_fp8_linear(x, w_fp8, w_scale, block=128) + + +@pytest.mark.parametrize("block", [32, 128]) +def test_grouped_w8a16_gemv_matches_the_dequantized_einsum(block): + """``wo_a`` at decode: the fp8 payload dequantized in-kernel against a bf16 activation must equal + the reference's bf16 einsum over the dequantized weight (fp32 accumulation, order aside).""" + from freetoken.kernel.triton.dsv4.fp8_linear import dequant_block_fp8, grouped_w8a16_gemv + + G, r, d = 4, 256, 1024 + torch.manual_seed(block) + w = torch.randn(G * r, d, device="cuda") * 0.05 + w_fp8, scale = quantize_weight(w, block) + x = torch.randn(G, d, device="cuda", dtype=torch.bfloat16) + deq = dequant_block_fp8(w_fp8, scale, block=block) + ref_deq = (w_fp8.float() * torch.exp2(scale.view(torch.uint8).float() - 127.0).repeat_interleave(block, 0).repeat_interleave(block, 1)).to(torch.bfloat16) + assert torch.equal(deq, ref_deq) + for t in (1, 3): + xt = torch.randn(t, G, d, device="cuda", dtype=torch.bfloat16) if t > 1 else x.unsqueeze(0) + want = torch.einsum("tgd,grd->tgr", xt, ref_deq.view(G, r, d)).flatten(1) + got = grouped_w8a16_gemv(xt, w_fp8, scale, block=block) + assert got.shape == (t, G * r) and got.dtype == torch.bfloat16 + torch.testing.assert_close(got.float(), want.float(), atol=2e-2, rtol=2e-2) diff --git a/tests/kernels/test_dsv4_router_select.py b/tests/kernels/test_dsv4_router_select.py new file mode 100644 index 0000000000..9829ed038a --- /dev/null +++ b/tests/kernels/test_dsv4_router_select.py @@ -0,0 +1,52 @@ +"""Fused gate selection (``kernel/triton/dsv4/router``) against the torch expression it replaces, for +every score function, with and without top-k weight normalization, at decode and prefill widths.""" + +from __future__ import annotations + +import pytest +import torch +import torch.nn.functional as F + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") + + +def reference(scores, bias, topk, score_func, normalize, route_scale): + if score_func == "softmax": + s = scores.softmax(dim=-1) + elif score_func == "sigmoid": + s = scores.sigmoid() + else: + s = F.softplus(scores).sqrt() + ids = (s + bias).topk(topk, dim=-1)[1] + w = s.gather(1, ids) + if normalize and topk > 1: + w = w / (w.sum(dim=-1, keepdim=True) + 1e-20) + return w * route_scale, ids + + +@pytest.mark.parametrize("score_func", ["sqrtsoftplus", "sigmoid", "softmax"]) +@pytest.mark.parametrize("normalize", [True, False]) +@pytest.mark.parametrize("T", [1, 3, 300]) +def test_router_select_matches_torch(score_func, normalize, T): + from freetoken.kernel.triton.dsv4.router import router_select + + g = torch.Generator(device="cpu").manual_seed(T) + E, K = 384, 6 + scores = (torch.randn(T, E, generator=g) * 3).cuda() + bias = (torch.randn(E, generator=g) * 0.1).cuda() + w, ids = router_select(scores, bias, K, score_func=score_func, normalize=normalize, route_scale=2.5) + w_ref, ids_ref = reference(scores, bias, K, score_func, normalize, 2.5) + assert torch.equal(ids.sort(dim=-1).values, ids_ref.to(torch.int32).sort(dim=-1).values) + # same expert set; compare the weights per expert (selection order may differ only on exact ties) + order = ids.argsort(dim=-1) + order_ref = ids_ref.argsort(dim=-1) + torch.testing.assert_close(w.gather(1, order), w_ref.gather(1, order_ref), atol=1e-6, rtol=1e-5) + + +def test_router_select_breaks_ties_toward_the_lowest_expert(): + from freetoken.kernel.triton.dsv4.router import router_select + + scores = torch.zeros(1, 384, device="cuda") # every expert ties + bias = torch.zeros(384, device="cuda") + _, ids = router_select(scores, bias, 6, score_func="sigmoid", normalize=True, route_scale=1.0) + assert ids.tolist() == [[0, 1, 2, 3, 4, 5]] 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_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..591c6f1d4e --- /dev/null +++ b/tests/kernels/test_qsa_nvfp4.py @@ -0,0 +1,141 @@ +"""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) 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/kernels/test_triton_sampling.py b/tests/kernels/test_triton_sampling.py new file mode 100644 index 0000000000..0766ea9343 --- /dev/null +++ b/tests/kernels/test_triton_sampling.py @@ -0,0 +1,23 @@ +"""The Triton sampling softmax, which samples when flashinfer is not installed.""" + +import pytest +import torch + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") + + +def test_softmax_survives_a_fully_masked_prefix(): + """min_p / min_tokens write -inf over most of the vocabulary. A column chunk or leading + block that is -inf throughout gave exp(-inf - -inf) = NaN and a NaN row.""" + from freetoken.kernel.triton.sampling import softmax + + V = 32768 # several column chunks per row + torch.manual_seed(0) + logits = torch.full((2, V), float("-inf"), device="cuda") + logits[0, [20000, 30000, 32000]] = torch.tensor([1.0, 2.0, 0.5], device="cuda") + logits[1] = torch.randn(V, device="cuda") + logits[1, :8192] = float("-inf") + temps = torch.tensor([0.7, 1.3], device="cuda") + probs = softmax(logits, temps) + assert not probs.isnan().any() + torch.testing.assert_close(probs, torch.softmax(logits / temps[:, None], dim=-1)) diff --git a/tests/kvcache/test_cache_unit_bytes.py b/tests/kvcache/test_cache_unit_bytes.py index 496622f985..cce5df6246 100644 --- a/tests/kvcache/test_cache_unit_bytes.py +++ b/tests/kvcache/test_cache_unit_bytes.py @@ -16,7 +16,7 @@ compute_cache_status_meta, compute_cache_unit_bytes, ) -from freetoken.kvcache.dsv4_cost_model import _dsv4_window_floor_pages +from freetoken.kvcache.dsv4.v4_cost_model import _dsv4_window_floor_pages def _init_tp() -> None: @@ -204,7 +204,7 @@ def test_floors_dsv4_reports_real_window_floor(): dsv4_args=SimpleNamespace(window_size=P), has_swa_attention=False ), ) - from freetoken.kvcache.dsv4_paged_pool import DSV4PagedKVCache + from freetoken.kvcache.dsv4.v4_pool import DSV4PagedKVCache eng = SimpleNamespace(config=cfg, moe_offload_cache=None, linear_state_pool=None) eng.kv_cache = object.__new__(DSV4PagedKVCache) # min_kv_tokens is a classmethod 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_dsv41_pool.py b/tests/kvcache/test_dsv41_pool.py new file mode 100644 index 0000000000..f41df85174 --- /dev/null +++ b/tests/kvcache/test_dsv41_pool.py @@ -0,0 +1,205 @@ +"""DSV41 paged KV pool + cost model (CPU, no model): args validation, per-source tiers, sizing / +byte accounting, the window free-list duck-type, full-loc translation; packed writes on CUDA.""" + +from __future__ import annotations + +import pytest +import torch + +from freetoken.kvcache.dsv4.v41_cost_model import ( + dsv41_cache_per_page, + dsv41_kv_unit_bytes, + dsv41_pool_bytes, + dsv41_pool_sizes, + dsv41_solve_num_pages, + dsv41_window_unit_bytes, +) +from freetoken.kvcache.dsv4.v41_pool import DSV41PagedKVCache +from freetoken.kvcache.dsv4.v41_row_format import FP4_E4M3_B16, FP4_E8M0_B32, FP8_E8M0_B32 +from freetoken.models.deepseek_v41.args import DeepseekV41Args + +DEVICE = torch.device("cpu") +P = 128 +# a 10-layer miniature of V4.1's layout: 2 window-only, ratio-2 encoder (source 2, source 5), ratio-1 decoder (source 7) +RATIOS = (0, 0, 2, 2, 2, 2, 2, 1, 1, 1) +SOURCES = (2, 5, 7) + + +def _args(**over) -> DeepseekV41Args: + base = dict(n_layers=10, head_dim=512, index_head_dim=128, window_size=P, compress_ratios=RATIOS, + kv_source_layers=SOURCES, swa_decoder_replay="exact") + base.update(over) + return DeepseekV41Args(**base) + + +def _pool(num_pages=8, swa_ratio=0.5, n_scratch=1, **over): + args = _args(**over) + sizes = dsv41_pool_sizes(num_pages, args, swa_ratio, P) + return DSV41PagedKVCache(sizes, args, DEVICE, n_scratch=n_scratch), sizes, args + + +def test_window_control_uses_pool_units_and_restores_concrete_capacity(): + from types import SimpleNamespace + + from freetoken.attention import AttnType + from freetoken.kvcache.cache_status import ( + _supports_swa_ratio, compute_cache_floors, compute_cache_pools, + ) + from freetoken.kvcache.dsv4.v41_cost_model import _dsv41_pool_sizes + from freetoken.scheduler.scheduler import Scheduler + from freetoken.server.api_server import CacheRebuildRequest, _resolve_num_swa_pages, cache_geometry + from freetoken.server.stats import _swa_page_size + + pool, _, args = _pool(num_pages=64) + config = SimpleNamespace( + page_size=P, max_running_req=1, max_seq_len=1024, cache_type="swa_radix", + swa_full_tokens_ratio=0.5, swa_num_pages_override=None, + model_config=SimpleNamespace( + dsv4_args=None, dsv41_args=_args(), has_swa_attention=False, + kv_cache_group_specs=lambda: [SimpleNamespace(attn_type=AttnType.DSV41)], + ), + ) + engine = SimpleNamespace( + config=config, kv_cache=pool, num_pages=63, + moe_offload_cache=None, linear_state_pool=None, + ) + assert _supports_swa_ratio(config) + assert compute_cache_pools(engine)["swa_page_size"] == P + assert compute_cache_pools(engine)["num_swa_pages"] == 31 + floor = _dsv41_pool_sizes(config, 64, num_swa_pages=1).n_win_pages - 1 + assert compute_cache_floors(engine)["swa_tokens"] == floor * P + prior = Scheduler._current_cache_geometry(SimpleNamespace(engine=engine, config=config)) + assert prior["num_swa_pages"] == 31 + state = SimpleNamespace(config=config, cache_pools=compute_cache_pools(engine)) + req = CacheRebuildRequest(num_pages=64, swa_full_tokens_ratio=0.5) + assert _resolve_num_swa_pages(state, req) == 32 + assert _swa_page_size(config) == P + state.stats = SimpleNamespace(kv_total_pages=63, mamba_total_slots=0) + state.last_rebuild = {"num_pages": 64, "num_swa_pages": 32} + assert cache_geometry(state)["swa_full_tokens_ratio"] == 0.5 + pool._init_paged_state(1, True) + pool.rebuild_from_config(config, 63, num_swa_pages=24) + assert pool.window_pages == 24 + pool.rebuild_from_config(config, prior["num_pages"], num_swa_pages=prior["num_swa_pages"]) + assert pool.window_pages == 31 + + +def test_pool_reads_sources_and_rings_from_args(): + pool, _, _ = _pool() + assert [pool.source_of(l) for l in range(2, 10)] == [2, 2, 2, 5, 5, 7, 7, 7] + assert set(pool.state_ring) == {2, 5} and pool.state_ring[2].ring_size == 2 + with pytest.raises(ValueError): # the window is not a multiple of a compress ratio + _pool(window_size=P - 1) + with pytest.raises(ValueError): # an indexer key the fp4 row format cannot hold + _pool(index_head_dim=48) + + +@pytest.mark.parametrize("replay", ["bounded", "exact"]) +def test_v41_flash_row_bytes_and_private_window_layers(replay): + """The tech report's headline: 3 ratio-2 encoder sources + 1 ratio-1 decoder source = 890 B/token.""" + ratios = (0, 0) + (2,) * 18 + (1,) * 20 + args = DeepseekV41Args(n_layers=40, head_dim=512, index_head_dim=128, window_size=128, compress_ratios=ratios, + kv_source_layers=(2, 8, 14, 20), swa_decoder_replay=replay) + pool = DSV41PagedKVCache(dsv41_pool_sizes(4, args, 1.0, 128), args, DEVICE) + win, main, idx = pool.window_pool[0].shape[1], pool.main_pool[2].shape[1], pool.idx_pool[2].shape[1] + assert (win, main, idx) == (528, 288, 68) + assert pool.state_ring[2].buffer.shape[1] * 4 == 4096 + assert sum((main + idx) // pool.ratio_of(s) for s in args.backbone_kv_sources) == 890 + assert dsv41_kv_unit_bytes(args, 128) == 890 + 8 # + the int64 full->window map slot + assert pool.private_window_layer_ids == (tuple(range(20, 40)) if replay == "bounded" else ()) + + +def test_pool_tiers_per_source_and_aliasing(): + pool, sizes, args = _pool() + assert len(pool.window_pool) == 10 and pool.window_pool[0].shape == (sizes.n_win_slots, 528) + assert set(pool.main_pool) == set(SOURCES) == set(pool.idx_pool) + assert set(pool.state_ring) == {2, 5} # the ratio-1 decoder source carries no partial group + assert pool.main_pool[2].shape == (sizes.full_token // 2 + 1, 288) + assert pool.main_pool[7].shape == (sizes.full_token + 1, 288) + assert pool.idx_pool[5].shape == (sizes.full_token // 2 + 1, 68) + assert pool.main_pool_of(4) is pool.main_pool[2] and pool.idx_pool_of(9) is pool.idx_pool[7] + assert pool.scratch_base[2] == sizes.full_token // 2 + with pytest.raises(AssertionError): + pool.main_pool_of(0) + ring = pool.state_ring[2] + assert ring.ring_size == 2 and ring.item_size == 512 and ring.buffer.shape == (sizes.state_slots[2] + 1, 1024) + assert torch.all(ring.buffer[-1, :512] == 0) and torch.all(torch.isneginf(ring.buffer[-1, 512:])) + + +def test_pool_bytes_match_allocation_and_solver_respects_budget(): + pool, sizes, args = _pool(num_pages=16, swa_ratio=0.25) + assert pool.total_bytes() == dsv41_pool_bytes(sizes, args, n_scratch=1) + assert dsv41_cache_per_page(args, 0.25, P) > 0 + assert dsv41_window_unit_bytes(args, P) == -(-(10 * P * 528 + 2 * 2 * 4096) // P) + budget = 64 << 20 + solved = dsv41_solve_num_pages(budget, args, 0.25, floor_win_pages=4, P=P, n_scratch=3) + assert dsv41_pool_bytes(solved, args, 3) <= budget + # one more page, sized the way the solver sizes (window = max(floor, ceil(ratio * pages))), overflows + more = solved.full_token // P + 1 + bigger = dsv41_pool_sizes(more, args, 0.25, P, n_win_pages=max(4, (round(0.25 * more * P) + P - 1) // P)) + assert dsv41_pool_bytes(bigger, args, 3) > budget + assert solved.n_win_pages >= 4 + with pytest.raises(ValueError): + dsv41_solve_num_pages(1 << 10, args, 0.25, floor_win_pages=4, P=P) + + +def test_translation_state_loc_and_cmp_rows(): + pool, sizes, args = _pool(num_pages=16) + pool.bind_window_pages(full_page_base=0, window_page_base=2 * P) + pool.bind_window_pages(full_page_base=3 * P, window_page_base=0) + assert pool.translate_full_to_window(torch.tensor([0, 1, 127])).tolist() == [2 * P, 2 * P + 1, 2 * P + 127] + assert pool.translate_full_to_window(torch.tensor([3 * P + 5])).item() == 5 + assert pool.translate_full_to_window(torch.tensor([P + 7, -1])).tolist() == [-1, -1] + full = torch.tensor([0, 1, 2, 127, 128, 4 * P - 1]) + assert pool.cmp_rows(full, 2).tolist() == [0, 0, 1, 63, 64, 2 * P - 1] + assert pool.cmp_rows(full, 1).tolist() == full.tolist() + assert pool.cmp_rows(torch.tensor([-1]), 2).item() < 0 + top = pool.cmp_rows(torch.tensor([sizes.full_token - 1]), 2).item() + assert top < pool.scratch_base[2] + ws = pool.translate_full_to_window(torch.tensor([127, 3 * P])) + assert DSV41PagedKVCache.state_loc(ws, 2, P).tolist() == [2 * 2 + 1, 0] + assert DSV41PagedKVCache.state_loc(torch.tensor([-1]), 2, P).item() == -1 + + +def _expand(bases): + return (torch.tensor(bases, dtype=torch.int64)[:, None] + torch.arange(P)).flatten() + + +def test_swa_duck_type_alloc_free_and_dummy(): + pool, sizes, _ = _pool(num_pages=8) + pool._init_paged_state(max_running_req=2, radix=True) + cap = sizes.n_win_slots - P + assert pool.swa_available_size() == cap and pool.swa_num_tokens - 1 == cap + assert pool.sliding_window_size == P and pool.prefill_chunk_budget >= P + pool.alloc_swa(_expand([0, 2 * P])) + assert pool.swa_available_size() == cap - 2 * P + ws = pool.translate_loc_from_full_to_swa(torch.arange(2 * P, 3 * P)) + assert (ws >= 0).all() and int(ws[0]) % P == 0 and torch.equal(ws - ws[0], torch.arange(P)) + pool.free_swa(_expand([0])) + pool.free_swa(_expand([0])) # idempotent + assert pool.swa_available_size() == cap - P + with pytest.raises(AssertionError): + pool.free_swa(torch.arange(2 * P, 2 * P + 5)) + dummy = pool.translate_loc_from_full_to_swa(torch.arange(sizes.full_token - P, sizes.full_token)) + assert int(dummy[0]) == sizes.n_win_slots - P + assert pool.k_cache(3) is pool.window_pool[3] and pool.num_layers == 10 and pool.unit_bytes()[0] > 0 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="packed writes run through triton") +def test_packed_writes_round_trip_on_cuda(): + from kernels.test_dsv41_pack import reference_roundtrip # tests/ is on sys.path under pytest + + args = _args() + sizes = dsv41_pool_sizes(4, args, 0.5, P) + pool = DSV41PagedKVCache(sizes, args, torch.device("cuda"), n_scratch=2) + kv = torch.randn(3, 512, device="cuda", dtype=torch.bfloat16) + slots = torch.tensor([0, 130, 5], device="cuda") + pool.store_window(kv, 4, slots) + assert torch.equal(pool.read_window(4, slots), reference_roundtrip(kv, FP8_E8M0_B32)) + pool.store_main(kv, 2, slots) + assert torch.equal(pool.read_main(2, slots), reference_roundtrip(kv, FP4_E4M3_B16)) + k = torch.randn(3, 128, device="cuda", dtype=torch.bfloat16) + pool.store_index(k, 7, slots) + assert torch.equal(pool.read_index(7, slots), reference_roundtrip(k, FP4_E8M0_B32)) + pool.store_kv(kv, kv, slots, 0) + assert torch.equal(pool.read_window(0, slots), reference_roundtrip(kv, FP8_E8M0_B32)) diff --git a/tests/kvcache/test_dsv4_free_list.py b/tests/kvcache/test_dsv4_free_list.py index a3ae8d0cd7..7b0dfb5566 100644 --- a/tests/kvcache/test_dsv4_free_list.py +++ b/tests/kvcache/test_dsv4_free_list.py @@ -5,7 +5,7 @@ import pytest import torch -from freetoken.kvcache.dsv4_paged_pool import FreeListAllocator +from freetoken.kvcache.dsv4.v4_pool import FreeListAllocator DEVICE = torch.device("cpu") P = 128 diff --git a/tests/kvcache/test_dsv4_pool.py b/tests/kvcache/test_dsv4_pool.py index c345b67c82..c8e21997ed 100644 --- a/tests/kvcache/test_dsv4_pool.py +++ b/tests/kvcache/test_dsv4_pool.py @@ -10,8 +10,8 @@ import pytest import torch -from freetoken.kvcache.dsv4_cost_model import dsv4_pool_sizes -from freetoken.kvcache.dsv4_paged_pool import CompressStateRing, DSV4PagedKVCache +from freetoken.kvcache.dsv4.v4_cost_model import dsv4_pool_sizes +from freetoken.kvcache.dsv4.v4_pool import CompressStateRing, DSV4PagedKVCache from freetoken.models.deepseek_v4.args import DeepseekV4Args DEVICE = torch.device("cpu") 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 a5ec94a084..45772fe4e5 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 @@ -178,13 +185,13 @@ def test_dsv4_rebuild_from_config_builds_the_pool_sizes_and_attaches_the_table() the engine used to do -- and re-points full_loc_map at the shared page table.""" from types import SimpleNamespace - from freetoken.kvcache.dsv4_cost_model import _dsv4_pool_sizes - from freetoken.kvcache.dsv4_cost_model import ( + from freetoken.kvcache.dsv4.v4_cost_model import _dsv4_pool_sizes + from freetoken.kvcache.dsv4.v4_cost_model import ( dsv4_kv_unit_bytes, dsv4_pool_sizes, dsv4_window_unit_bytes, ) - from freetoken.kvcache.dsv4_paged_pool import DSV4PagedKVCache + from freetoken.kvcache.dsv4.v4_pool import DSV4PagedKVCache from freetoken.models.deepseek_v4.args import DeepseekV4Args P, mrr = 128, 1 @@ -222,7 +229,7 @@ def test_dsv4_refresh_seq_state_tracks_page_table_width(): from freetoken.engine.engine import Engine from freetoken.utils import align_ceil - from freetoken.kvcache.dsv4_paged_pool import DSV4PagedKVCache + from freetoken.kvcache.dsv4.v4_pool import DSV4PagedKVCache from freetoken.scheduler.table import TableManager P, mrr = 128, 4 @@ -264,7 +271,7 @@ def test_every_kv_pool_answers_the_sizing_surface(): the model resolved to. DSV4 must carry its OWN overrides (getattr would resolve to the base class and silently misprice the window floor if an override vanished).""" from freetoken.kvcache.dsa_pool import DSAKVCache, MLAKVCache - from freetoken.kvcache.dsv4_paged_pool import DSV4PagedKVCache + from freetoken.kvcache.dsv4.v4_pool import DSV4PagedKVCache from freetoken.kvcache.hybrid_swa_pool import HybridSWAKVCache from freetoken.kvcache.mha_pool import MHAKVCache @@ -279,7 +286,7 @@ def test_every_kv_pool_answers_the_rebuild_surface(): """The engine drives every pool through this one surface; a pool that skips an implementation must fail loudly at import/instantiation, not at rebuild time.""" from freetoken.kvcache.dsa_pool import DSAKVCache, MLAKVCache - from freetoken.kvcache.dsv4_paged_pool import DSV4PagedKVCache + from freetoken.kvcache.dsv4.v4_pool import DSV4PagedKVCache from freetoken.kvcache.hybrid_swa_pool import HybridSWAKVCache from freetoken.kvcache.mha_pool import MHAKVCache 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_pool_sizing_surface.py b/tests/kvcache/test_pool_sizing_surface.py index 516090b099..33f27983de 100644 --- a/tests/kvcache/test_pool_sizing_surface.py +++ b/tests/kvcache/test_pool_sizing_surface.py @@ -30,7 +30,7 @@ def _model_config(specs, **attrs): def test_resolve_pool_class_follows_attn_type(): from freetoken.kvcache.dsa_pool import DSAKVCache, MLAKVCache - from freetoken.kvcache.dsv4_paged_pool import DSV4PagedKVCache + from freetoken.kvcache.dsv4.v4_pool import DSV4PagedKVCache from freetoken.kvcache.hybrid_swa_pool import HybridSWAKVCache from freetoken.kvcache.mha_pool import MHAKVCache @@ -91,11 +91,11 @@ def _dsv4_config(num_page_override=None): def test_dsv4_kv_cost_and_floor_parity(): - from freetoken.kvcache.dsv4_cost_model import _dsv4_swa_ratio, _dsv4_window_floor_pages - from freetoken.kvcache.dsv4_paged_pool import DSV4PagedKVCache + from freetoken.kvcache.dsv4.v4_cost_model import _dsv4_swa_ratio, _dsv4_window_floor_pages + from freetoken.kvcache.dsv4.v4_pool import DSV4PagedKVCache - pytest.importorskip("freetoken.kvcache.dsv4_cost_model") - from freetoken.kvcache.dsv4_cost_model import dsv4_auto_cost_model + pytest.importorskip("freetoken.kvcache.dsv4.v4_cost_model") + from freetoken.kvcache.dsv4.v4_cost_model import dsv4_auto_cost_model config = _dsv4_config() P = 128 @@ -147,8 +147,8 @@ def check(pages, baseline): def test_dsv4_validate_rebuild_floor(): from freetoken.kvcache.base import CacheRebuildRejected - from freetoken.kvcache.dsv4_cost_model import _dsv4_window_floor_pages - from freetoken.kvcache.dsv4_paged_pool import DSV4PagedKVCache + from freetoken.kvcache.dsv4.v4_cost_model import _dsv4_window_floor_pages + from freetoken.kvcache.dsv4.v4_pool import DSV4PagedKVCache config = _dsv4_config() pool = object.__new__(DSV4PagedKVCache) # floor path reads no instance state @@ -165,7 +165,7 @@ def test_rebuild_from_config_explicit_num_swa_pages_wins(): # The explicit parameter must take precedence over config.swa_num_pages_override -- # the engine's override write would otherwise mask a dropped parameter forever. from freetoken.distributed import set_tp_info, try_get_tp_info - from freetoken.kvcache.dsv4_cost_model import _dsv4_pool_sizes + from freetoken.kvcache.dsv4.v4_cost_model import _dsv4_pool_sizes from freetoken.kvcache.hybrid_swa_pool import _swa_paged_num_tokens, _swa_pool_floor if try_get_tp_info() is None: @@ -183,7 +183,7 @@ def test_rebuild_from_config_explicit_num_swa_pages_wins(): assert _swa_paged_num_tokens(swa_cfg, 11, num_swa_pages=explicit) == explicit + 1 assert _swa_paged_num_tokens(swa_cfg, 11) == 999_999 + 1 # override path unchanged - from freetoken.kvcache.dsv4_cost_model import _dsv4_window_floor_pages + from freetoken.kvcache.dsv4.v4_cost_model import _dsv4_window_floor_pages d_cfg = _dsv4_config() d_floor = _dsv4_window_floor_pages(d_cfg, d_cfg.page_size) 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/layers/test_mhc.py b/tests/layers/test_mhc.py index 4cf144e720..0359bed9c9 100644 --- a/tests/layers/test_mhc.py +++ b/tests/layers/test_mhc.py @@ -152,6 +152,75 @@ def test_triton_fused_respects_input_dtype(dtype, tol): assert err < tol, f"{name} [{dtype}]: max abs err {err}" +def test_single_pass_applies_the_previous_pre_and_hands_on_its_own(): + """Single-pass mHC (DeepSeek-V4.1): the input mix uses ``pre_mix_in``, the gates + computed here are the next sublayer's pre; post/comb are unchanged from multi-pass.""" + from freetoken.layers.mhc import ( + mhc_fused_post_pre_single_pass, + mhc_gates, + mhc_mix_input, + mhc_pre_single_pass, + ) + + fn, scale, base = _weights(seed=21) + res = _residual(seed=22) + x = torch.randn(T, HIDDEN, dtype=torch.bfloat16) + pre_in = torch.rand(T, N, dtype=torch.float32) + post0, comb0, _ = mhc_pre(res, fn, scale, base, RMS_EPS, EPS, POST_MULT, SINKHORN) + + r1, p1, c1, pre1, li1 = mhc_fused_post_pre_single_pass( + x, res, post0, comb0, pre_in, fn, scale, base, RMS_EPS, EPS, POST_MULT, SINKHORN + ) + r_ref = mhc_post(x, res, post0, comb0) + pre_ref, p_ref, c_ref = mhc_gates(r_ref, fn, scale, base, RMS_EPS, EPS, POST_MULT, SINKHORN) + assert torch.equal(r1, r_ref) and torch.equal(p1, p_ref) and torch.equal(c1, c_ref) + assert torch.equal(pre1, pre_ref) + assert torch.equal(li1, mhc_mix_input(r_ref, pre_in)) + # the multi-pass input would have used pre_ref, not pre_in + assert not torch.equal(li1, mhc_mix_input(r_ref, pre_ref)) + # a one-hot pre picks one stream (the model's seed for the chain) + one_hot = torch.zeros(T, N) + one_hot[:, 0] = 1.0 + _, _, _, li0 = mhc_pre_single_pass(res, one_hot, fn, scale, base, RMS_EPS, EPS, POST_MULT, SINKHORN) + assert torch.equal(li0, res[:, 0]) + # no post on the first sublayer: the streams pass through untouched + r0, *_ = mhc_fused_post_pre_single_pass( + x, res, None, None, one_hot, fn, scale, base, RMS_EPS, EPS, POST_MULT, SINKHORN + ) + assert torch.equal(r0, res) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +@pytest.mark.parametrize("t,hidden,with_post", [(1, 64, True), (9, 64, True), (3, 5120, True), (3, 5120, False)]) +def test_triton_single_pass_matches_torch(t, hidden, with_post): + from freetoken.layers.mhc import mhc_fused_post_pre_single_pass_torch + from freetoken.kernel.triton.mhc import mhc_fused_post_pre_single_pass_triton + + torch.manual_seed(23) + mix = 2 * N + N * N + fn = torch.randn(mix, N * hidden, dtype=torch.float32, device="cuda") * 0.05 + scale = torch.rand(3, dtype=torch.float32, device="cuda") + 0.5 + base = torch.randn(mix, dtype=torch.float32, device="cuda") * 0.3 + res = torch.randn(t, N, hidden, dtype=torch.bfloat16, device="cuda") + x = torch.randn(t, hidden, dtype=torch.bfloat16, device="cuda") + pre_in = torch.rand(t, N, dtype=torch.float32, device="cuda") + post0 = torch.rand(t, N, 1, dtype=torch.float32, device="cuda") * POST_MULT if with_post else None + comb0 = torch.softmax(torch.randn(t, N, N, device="cuda"), dim=-1) if with_post else None + + ref = mhc_fused_post_pre_single_pass_torch( + x, res, post0, comb0, pre_in, fn, scale, base, RMS_EPS, EPS, POST_MULT, SINKHORN + ) + got = mhc_fused_post_pre_single_pass_triton( + x, res, post0, comb0, pre_in, fn, scale, base, RMS_EPS, EPS, POST_MULT, SINKHORN + ) + names = ("residual", "post", "comb", "pre_new", "layer_input") + tols = (2e-2, 2e-3, 2e-3, 2e-3, 2e-2) + for name, r, g, tol in zip(names, ref, got, tols): + assert g.shape == r.shape, (name, g.shape, r.shape) + err = (g.float() - r.float()).abs().max().item() + assert err < tol, f"{name}: max abs err {err}" + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") def test_triton_pre_only_matches_torch(): """HAS_POST=False path (layer 0's standalone hc_pre through the fused kernel).""" @@ -176,3 +245,31 @@ def test_triton_pre_only_matches_torch(): assert (got_post.float() - ref_post.float()).abs().max().item() < 2e-3 assert (got_comb.float() - ref_comb.float()).abs().max().item() < 2e-3 assert (got_li.float() - ref_li.float()).abs().max().item() < 2e-2 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +def test_mhc_split_preserves_chunked_results_and_graph_replay(): + from freetoken.kernel.triton.mhc import mhc_fused_post_pre_single_pass_triton as fused + + torch.manual_seed(26) + hidden, tokens = 5120, 7 + res = torch.randn(tokens, N, hidden, device="cuda", dtype=torch.bfloat16) + x = torch.randn(tokens, hidden, device="cuda", dtype=torch.bfloat16) + post = torch.rand(tokens, N, 1, device="cuda") + comb = torch.softmax(torch.randn(tokens, N, N, device="cuda"), -1) + pre = torch.rand(tokens, N, device="cuda") + fn = torch.randn(2 * N + N * N, N * hidden, device="cuda") * 0.05 + scale, base = torch.ones(3, device="cuda"), torch.zeros(2 * N + N * N, device="cuda") + def run(lo, hi): + return fused(x[lo:hi], res[lo:hi], post[lo:hi], comb[lo:hi], pre[lo:hi], fn, scale, base, + RMS_EPS, EPS, POST_MULT, SINKHORN) + whole = run(0, tokens) + chunks = [run(0, 2), run(2, tokens)] + for i, expected in enumerate(whole): + assert torch.equal(expected, torch.cat([chunk[i] for chunk in chunks])) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = run(0, tokens) + graph.replay() + for got, expected in zip(captured, whole): + assert torch.equal(got, expected) diff --git a/tests/models/deepseek_v41/__init__.py b/tests/models/deepseek_v41/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/models/deepseek_v41/common.py b/tests/models/deepseek_v41/common.py new file mode 100644 index 0000000000..d0f8efeef9 --- /dev/null +++ b/tests/models/deepseek_v41/common.py @@ -0,0 +1,176 @@ +"""Shared fixtures for the deepseek_v41 tests: a tiny V4.1-shaped config (2 window layers, a ratio-2 +encoder group with one Full + one Reuse layer, a ratio-1 decoder group with a Full candidate source and +a Reindex layer) and synthetic bf16 and FP8/MXFP4 checkpoints in the real tensor naming. Holds no tests itself.""" + +from __future__ import annotations + +import json +import os +from types import SimpleNamespace + +import pytest +import torch + +requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU") + +VOCAB = 512 +DIM = 256 +N_LAYERS = 6 +RATIOS = [0, 0, 2, 2, 1, 1, 0] # + one trailing MTP layer the model ignores +KV_SOURCES = [2, 4] +INDEX_SOURCES = [2, 4, 5] +ENGRAM_LAYERS = [1] +ENGRAM_VOCAB = 500 # bucket prime start; 3 orders x 2 heads ~= 3000 rows +ENGRAM_ROWS = 4096 + + +def tiny_text_config(**over) -> dict: + text = dict( + model_type="deepseek_v41_text", + vocab_size=VOCAB, hidden_size=DIM, moe_intermediate_size=64, num_hidden_layers=N_LAYERS, num_attention_heads=4, + num_key_value_heads=1, head_dim=64, qk_rope_head_dim=16, q_lora_rank=64, o_lora_rank=32, o_groups=2, + hidden_act="silu", swiglu_limit=10.0, rms_norm_eps=1e-20, max_position_embeddings=16384, + rope_theta=10000, rope_scaling=dict(rope_type="yarn", factor=16, beta_fast=32, beta_slow=1, original_max_position_embeddings=1024), + n_routed_experts=8, n_shared_experts=1, num_experts_per_tok=2, scoring_func="sqrtsoftplus", topk_method="noaux_tc", + norm_topk_prob=True, routed_scaling_factor=1.5, sliding_window=128, + compress_ratios=list(RATIOS), compress_rope_theta=160000, kv_source_layer_ids=list(KV_SOURCES), + index_source_layer_ids=list(INDEX_SOURCES), index_n_heads=2, index_head_dim=32, index_topk=8, + candidate_source_layer_id=4, candidate_topk_blocks=2, candidate_block_size=4, + hc_mult=4, hc_sinkhorn_iters=20, hc_eps=1e-6, + engram_layer_ids=list(ENGRAM_LAYERS), engram_num_embeddings=[ENGRAM_ROWS], engram_max_ngram_size=4, + engram_vocab_size=ENGRAM_VOCAB, engram_n_heads=2, engram_head_dim=32, engram_pad_token_id=2, + engram_compressed_vocab_size=VOCAB, # the synthetic tokenizer maps every id to itself + num_nextn_predict_layers=1, dspark_block_size=0, dspark_target_layer_ids=[], + ) + text.update(over) + return text + + +def tiny_hf_config(**over) -> SimpleNamespace: + return SimpleNamespace(architectures=["DeepseekV41ForCausalLM"], model_type="deepseek_v41", text_config=tiny_text_config(**over)) + + +def checkpoint_tensors(text: dict, seed: int = 0) -> dict[str, torch.Tensor]: + """Random bf16 weights under the checkpoint's names (the reference ``convert.py`` dialect without + quantization: no ``.scale`` tensors). hc / sink / router-bias / engram gates in fp32 as shipped.""" + g = torch.Generator().manual_seed(seed) + + def rnd(*shape, scale=0.02, dtype=torch.bfloat16): + return (torch.randn(*shape, generator=g) * scale).to(dtype) + + dim, heads, hd, rd = text["hidden_size"], text["num_attention_heads"], text["head_dim"], text["qk_rope_head_dim"] + qr, ol, og = text["q_lora_rank"], text["o_lora_rank"], text["o_groups"] + inter, E, hc = text["moe_intermediate_size"], text["n_routed_experts"], text["hc_mult"] + ih, ihd = text["index_n_heads"], text["index_head_dim"] + mix = (2 + hc) * hc + t: dict[str, torch.Tensor] = { + "embed.weight": rnd(text["vocab_size"], dim, scale=1.0), + "norm.weight": torch.ones(dim, dtype=torch.bfloat16), + "head.weight": rnd(text["vocab_size"], dim), + } + for L in range(text["num_hidden_layers"]): + a = f"layers.{L}.attn" + t[f"{a}.wq_a.weight"] = rnd(qr, dim) + t[f"{a}.q_norm.weight"] = torch.ones(qr, dtype=torch.bfloat16) + t[f"{a}.wq_b.weight"] = rnd(heads * hd, qr) + t[f"{a}.wkv.weight"] = rnd(hd, dim) + t[f"{a}.kv_norm.weight"] = torch.ones(hd, dtype=torch.bfloat16) + t[f"{a}.wo_a.weight"] = rnd(og * ol, heads * hd // og) + t[f"{a}.wo_b.weight"] = rnd(dim, og * ol) + t[f"{a}.attn_sink"] = rnd(heads, scale=0.5, dtype=torch.float32) + ratio = text["compress_ratios"][L] + if L in text["kv_source_layer_ids"]: + t[f"{a}.compressor.wkv.weight"] = rnd(hd, dim) + if ratio > 1: + t[f"{a}.compressor.wgate.weight"] = rnd(hd, dim) + t[f"{a}.compressor.norm.weight"] = torch.ones(hd, dtype=torch.bfloat16) + t[f"{a}.indexer.wk.weight"] = rnd(ihd, hd) + t[f"{a}.indexer.k_norm.weight"] = torch.ones(ihd, dtype=torch.bfloat16) + if L in text["index_source_layer_ids"]: + t[f"{a}.indexer.wq_b.weight"] = rnd(ih * ihd, qr) + t[f"{a}.indexer.weights_proj.weight"] = rnd(ih, dim) + p = f"layers.{L}" + t[f"{p}.attn_norm.weight"] = torch.ones(dim, dtype=torch.bfloat16) + t[f"{p}.ffn_norm.weight"] = torch.ones(dim, dtype=torch.bfloat16) + t[f"{p}.ffn.gate.weight"] = rnd(E, dim) + t[f"{p}.ffn.gate.bias"] = rnd(E, scale=0.1, dtype=torch.float32) + t[f"{p}.ffn.gate.bias_vl"] = rnd(E, scale=0.1, dtype=torch.float32) + for w, shape in (("w1", (inter, dim)), ("w2", (dim, inter)), ("w3", (inter, dim))): + t[f"{p}.ffn.shared_experts.{w}.weight"] = rnd(*shape) + for e in range(E): + t[f"{p}.ffn.experts.{e}.{w}.weight"] = rnd(*shape) + for nm in ("hc_attn_fn", "hc_ffn_fn"): + t[f"{p}.{nm}"] = rnd(mix, hc * dim, scale=0.05, dtype=torch.float32) + for nm in ("hc_attn_base", "hc_ffn_base"): + t[f"{p}.{nm}"] = rnd(mix, scale=0.3, dtype=torch.float32) + for nm in ("hc_attn_scale", "hc_ffn_scale"): + t[f"{p}.{nm}"] = (torch.rand(3, generator=g) + 0.5).to(torch.float32) + if L in text["engram_layer_ids"]: + e = f"{p}.engram" + cols = (text["engram_max_ngram_size"] - 1) * text["engram_n_heads"] + t[f"{e}.wkv.weight"] = rnd(dim * (hc + 1), cols * text["engram_head_dim"]) + t[f"{e}.q_weight"] = torch.ones(hc, dim, dtype=torch.bfloat16) + t[f"{e}.k_weight"] = torch.ones(hc, dim, dtype=torch.bfloat16) + rows = text["engram_num_embeddings"][text["engram_layer_ids"].index(L)] + t[f"{e}.embed.weight"] = rnd(rows, text["engram_head_dim"], scale=1.0).to(torch.float8_e4m3fn) + t[f"{e}.embed.scale"] = torch.full((rows, text["engram_head_dim"] // 32), 127, dtype=torch.uint8).view(torch.float8_e8m0fnu) + # the DSpark stage and the vision tower: present in a real checkpoint, dropped by the loader + t["mtp.0.attn_norm.weight"] = torch.ones(dim, dtype=torch.bfloat16) + t["vision.norm.weight"] = torch.ones(8, dtype=torch.bfloat16) + return t + + +def write_tiny_checkpoint(folder: str, seed: int = 0, *, quantized: bool = False, **over) -> dict[str, torch.Tensor]: + """A loadable checkpoint directory: HF-style config.json + one safetensors shard + index.""" + from safetensors.torch import save_file + + if quantized: + over.setdefault("moe_intermediate_size", 256) + text = tiny_text_config(**over) + os.makedirs(folder, exist_ok=True) + with open(os.path.join(folder, "config.json"), "w") as f: + config = {"architectures": ["DeepseekV41ForCausalLM"], "model_type": "deepseek_v41", "text_config": text} + if quantized: + config["quantization_config"] = { + "quant_method": "fp8", "activation_scheme": "dynamic", + "weight_block_size": [32, 32], "scale_fmt": "ue8m0", "expert_dtype": "fp4", + } + json.dump(config, f) + tensors = checkpoint_tensors(text, seed) + if quantized: + tensors = _quantize_checkpoint(tensors) + save_file(tensors, os.path.join(folder, "model-00001-of-00001.safetensors")) + with open(os.path.join(folder, "model.safetensors.index.json"), "w") as f: + json.dump({"metadata": {}, "weight_map": {k: "model-00001-of-00001.safetensors" for k in tensors}}, f) + return tensors + + +def _quantize_checkpoint(tensors): + from .reference.kernel import _round_fp4 + + out = dict(tensors) + fp8_suffixes = ( + ".attn.wq_a", ".attn.wq_b", ".attn.wkv", ".attn.wo_a", ".attn.wo_b", + ".indexer.wq_b", ".shared_experts.w1", ".shared_experts.w2", ".shared_experts.w3", ".engram.wkv", + ) + grid = torch.tensor([0., .5, 1., 1.5, 2., 3., 4., 6.]) + for name, weight in tensors.items(): + if not name.endswith(".weight"): + continue + stem = name.removesuffix(".weight") + if ".ffn.experts." in name: + groups = weight.float().unflatten(-1, (-1, 32)) + scale = torch.exp2(torch.ceil(torch.log2(groups.abs().amax(-1).clamp_min(1e-12) / 6))) + q = _round_fp4(groups / scale.unsqueeze(-1)).flatten(-2) + codes = (q.abs().unsqueeze(-1) == grid).to(torch.int32).argmax(-1).to(torch.uint8) + codes |= torch.signbit(q).to(torch.uint8) * 8 + out[name] = codes[:, 0::2] | (codes[:, 1::2] << 4) + out[stem + ".scale"] = scale.to(torch.float8_e8m0fnu) + elif stem.endswith(fp8_suffixes): + n, k = weight.shape + blocks = weight.float().reshape(n // 32, 32, k // 32, 32).permute(0, 2, 1, 3) + scale = torch.exp2(torch.ceil(torch.log2(blocks.abs().amax((-1, -2)).clamp_min(1e-4) / 448))) + expanded = scale.repeat_interleave(32, 0).repeat_interleave(32, 1) + out[name] = (weight.float() / expanded).to(torch.float8_e4m3fn) + out[stem + ".scale"] = scale.to(torch.float8_e8m0fnu) + return out diff --git a/tests/models/deepseek_v41/harness.py b/tests/models/deepseek_v41/harness.py new file mode 100644 index 0000000000..98b19e2de4 --- /dev/null +++ b/tests/models/deepseek_v41/harness.py @@ -0,0 +1,182 @@ +"""A minimal in-process engine for the tiny DeepSeek-V4.1 model: builds the DSV41 pool, backend and +context by hand (no scheduler), loads a synthetic bf16 or FP8/MXFP4 checkpoint with resident or offloaded experts, and drives +ragged prefill / batched decode batches through ``model.forward()``. Shared by the forward smoke test +and the reference-parity test.""" + +from __future__ import annotations + +from dataclasses import replace + +import torch + +from freetoken.core import Batch, Context, Req, SamplingParams, get_global_ctx, set_global_ctx +from freetoken.distributed.info import set_tp_info, try_get_tp_info +from freetoken.engine.engine import _materialize_loaded_weight_state_dict +from freetoken.kvcache.dsv4.v41_cost_model import dsv41_pool_sizes +from freetoken.kvcache.dsv4.v41_pool import DSV41PagedKVCache +from freetoken.layers import set_rope_device +from freetoken.layers.quantization import NoQuantConfig +from freetoken.layers.quantization.method import finalize_quant +from freetoken.models import create_model +from freetoken.models.deepseek_v41.config import parse_config +from freetoken.models.deepseek_v41.engram import ZeroEngramTable +from freetoken.models.deepseek_v41.weight import iter_weights +from freetoken.utils.hf import cached_load_hf_config +from freetoken.utils.torch_utils import torch_dtype + +P = 128 + + +class TinyEngine: + def __init__(self, checkpoint: str, *, max_seq_len: int = 1024, max_running_req: int = 2, swa_decoder_replay: str = "exact", engram_table=None, quantized: bool = False): + self.device = torch.device("cuda") + self.checkpoint = checkpoint + if try_get_tp_info() is None: + set_tp_info(0, 1) + set_rope_device(self.device) + hf_config = cached_load_hf_config(checkpoint) + mc = parse_config(hf_config) + args = mc.dsv41_args + args.max_seq_len = max_seq_len + args.max_batch_size = max_running_req + 1 + args.swa_decoder_replay = swa_decoder_replay + self.args = args + quant = NoQuantConfig() + if quantized: + from freetoken.models.register import checkpoint_quant_config, get_model_spec + + quant = checkpoint_quant_config(checkpoint, hf_config, get_model_spec("DeepseekV41ForCausalLM")) + self.config = replace(mc, moe_strategy="offload" if quantized else "fused", decode_target="gpu", quant=quant) + mc = self.config + with torch.device("meta"), torch_dtype(torch.bfloat16): + self.model = create_model(self.config) + state = _materialize_loaded_weight_state_dict( + self.model.state_dict(), iter_weights(checkpoint, self.device, include_moe_experts=False), device=self.device, + ) + self.model.load_state_dict(state) + finalize_quant(self.model) + from freetoken.moe.expert_banks import load_expert_banks + + method = self.model.model.layers.op_list[0].ffn.experts.quant_method + if quantized: + from freetoken.moe.offload_cache import OffloadMoeCache, attach_offload_moe_cache + + self.banks = load_expert_banks( + checkpoint, mc, method=method, device=self.device, dtype=torch.bfloat16, parallel=False, + ) + self.expert_cache = OffloadMoeCache( + num_layers=mc.num_moe_layers, num_experts=mc.num_experts, + cache_size=mc.num_moe_layers * mc.num_experts, device=self.device, + quant_format=self.banks.quant_format, layout=method.layout(), + ) + self.expert_cache.set_bank_sources(self.banks.sources, layer_residency=self.banks.layer_residency) + attach_offload_moe_cache(self.model, self.expert_cache) + else: # resident experts load as banks, not through the state dict (as in Engine._load_resident_experts) + from freetoken.layers import iter_moe_layers + from freetoken.moe.expert_banks import attach_resident_banks + + attach_resident_banks(list(iter_moe_layers(self.model)), load_expert_banks( + checkpoint, mc, method=method, device=self.device, dtype=torch.bfloat16, parallel=False, resident=True, + )) + for layer in self.model.engram_layers(): + layer.attach_table(engram_table if engram_table is not None else ZeroEngramTable(layer.width, self.device)) + + self.max_running_req = max_running_req + num_pages = max_seq_len // P + self.pool = DSV41PagedKVCache(dsv41_pool_sizes(num_pages + 1, args, 1.0, P), args, self.device, n_scratch=max_running_req + 1) + self.pool._init_paged_state(max_running_req, radix=False) + self.page_table = torch.zeros(max_running_req + 1, max_seq_len, dtype=torch.int32, device=self.device) + self.page_table[max_running_req].fill_(num_pages * P) # the dummy row -> the reserved tail page + self.pool.attach_page_table(self.page_table) + try: + ctx = get_global_ctx() + except AssertionError: + ctx = Context(page_size=P) + set_global_ctx(ctx) + from freetoken.attention.dsv41_sparse import DSV41SparseAttnBackend + + self.ctx = ctx + ctx.page_table = self.page_table + ctx.kv_cache = self.pool # the backend reads the pool's device from the context + self.backend = DSV41SparseAttnBackend(self.config) + self._next_page = 0 + self._bind() + + def _bind(self) -> None: + """Point the process-wide ``Context`` at this engine's page table, pool and backend. The + context is a singleton the model reads through ``get_global_ctx()``, so a test that holds two + engines must rebind before each forward; ``prefill`` / ``decode`` do this, and the model + re-resolves its backend when the binding changes.""" + ctx = self.ctx + if ( + getattr(ctx, "attn_backend", None) is self.backend + and getattr(ctx, "kv_cache", None) is self.pool + and getattr(ctx, "page_table", None) is self.page_table + ): + return + ctx.page_table = self.page_table + ctx.kv_cache = self.pool + ctx.attn_backend = self.backend + self.model.mark_for_rebind() + + # ----- requests ----------------------------------------------------------------------- + def new_request(self, table_idx: int, tokens: list[int], max_new: int = 64) -> Req: + """A request with pages for ``len(tokens) + max_new`` positions bound on row ``table_idx``.""" + n_pages = -(-(len(tokens) + max_new) // P) + first = self._next_page + self._next_page += n_pages + locs = torch.arange(first * P, (first + n_pages) * P, device=self.device) + self.page_table[table_idx, : n_pages * P] = locs.to(torch.int32) + self.pool.alloc_swa(locs) + return Req(input_ids=torch.tensor(tokens, dtype=torch.int32), table_idx=table_idx, cached_len=0, output_len=max_new, + uid=table_idx, sampling_params=SamplingParams(), cache_handle=None) + + def new_request_on_prefix(self, table_idx: int, donor: Req, tokens: list[int], max_new: int = 64) -> Req: + """A request whose first ``hit`` positions alias ``donor``'s pages (what a radix prefix hit does + through the page table) and whose remaining positions get pages of their own. ``hit`` is what the + scheduler admits: the donor's prompt, capped by ``match_req`` (the last token, and the pool's + ``prefix_replay_tokens`` under bounded replay) and page-aligned. The caller prefills the rest.""" + limit = len(tokens) - 1 + if self.pool.prefix_replay_tokens: + limit = min(limit, max(0, len(tokens) - self.pool.prefix_replay_tokens)) + hit = min(donor.device_len, limit) // P * P + assert tokens[:hit] == donor.input_ids[:hit].tolist(), "the prefix must match the donor's prompt" + self.page_table[table_idx, :hit] = self.page_table[donor.table_idx, :hit] + n_pages = -(-(len(tokens) + max_new - hit) // P) + first = self._next_page + self._next_page += n_pages + locs = torch.arange(first * P, (first + n_pages) * P, device=self.device) + self.page_table[table_idx, hit : hit + n_pages * P] = locs.to(torch.int32) + self.pool.alloc_swa(locs) + return Req(input_ids=torch.tensor(tokens, dtype=torch.int32), table_idx=table_idx, cached_len=hit, output_len=max_new, + uid=table_idx, sampling_params=SamplingParams(), cache_handle=None) + + def prefill(self, reqs: list[Req]) -> torch.Tensor: + """Run each request's ``[cached_len, device_len)`` tokens; returns ``[B, vocab]`` logits.""" + batch = Batch(reqs=reqs, phase="prefill") + batch.padded_reqs = list(reqs) + batch.input_ids = torch.cat([r.input_ids[r.cached_len : r.device_len] for r in reqs]).to(self.device) + batch.positions = torch.cat([torch.arange(r.cached_len, r.device_len) for r in reqs]).to(self.device) + self._bind() + self.backend.prepare_metadata(batch) + with self.ctx.forward_batch(batch), self.model.forward_host_ctx(batch, False): + return self.model.forward().float() + + def decode(self, reqs: list[Req], tokens: list[int]) -> torch.Tensor: + """Feed one new token per request (appended at ``device_len - 1``); returns ``[B, vocab]``.""" + for r, t in zip(reqs, tokens): + r.append_host(torch.tensor([t], dtype=torch.int32)) + r.complete_one() + batch = Batch(reqs=reqs, phase="decode") + batch.padded_reqs = list(reqs) + batch.input_ids = torch.tensor(tokens, dtype=torch.int32, device=self.device) + batch.positions = torch.tensor([r.device_len - 1 for r in reqs], device=self.device) + batch.active_table_idx = torch.tensor([r.table_idx for r in reqs], device=self.device) + self._bind() + self.backend.prepare_metadata(batch) + with self.ctx.forward_batch(batch), self.model.forward_host_ctx(batch, False): + return self.model.forward().float() + + def finish_prefill(self, reqs: list[Req]) -> None: + for r in reqs: + r.cached_len = r.device_len diff --git a/tests/models/deepseek_v41/reference/LICENSE b/tests/models/deepseek_v41/reference/LICENSE new file mode 100644 index 0000000000..d42fae903e --- /dev/null +++ b/tests/models/deepseek_v41/reference/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2023 DeepSeek + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. \ No newline at end of file diff --git a/tests/models/deepseek_v41/reference/__init__.py b/tests/models/deepseek_v41/reference/__init__.py new file mode 100644 index 0000000000..79d0c33108 --- /dev/null +++ b/tests/models/deepseek_v41/reference/__init__.py @@ -0,0 +1,2 @@ +"""The DeepSeek-V4.1 reference implementation (MIT, vendored from the HF repo's ``inference/``) with its +tilelang kernels transcribed to torch -- the numerical oracle FreeToken's port is tested against.""" diff --git a/tests/models/deepseek_v41/reference/engram.py b/tests/models/deepseek_v41/reference/engram.py new file mode 100644 index 0000000000..625553113f --- /dev/null +++ b/tests/models/deepseek_v41/reference/engram.py @@ -0,0 +1,186 @@ +# Vendored from https://huggingface.co/deepseek-ai/DeepSeek-V4.1-Flash (inference/engram.py), MIT License +# (see LICENSE next to this file); unmodified. +from dataclasses import dataclass + +import numpy as np +import torch +from sympy import isprime +from torch import nn + + +def find_next_prime(start: int, seen_primes: set[int]) -> int: + """The smallest prime above `start` that has not been handed out yet.""" + candidate = start + 1 + while not isprime(candidate) or candidate in seen_primes: + candidate += 1 + return candidate + + +def build_compressed_token_map(tokenizer) -> tuple[list[int], int]: + """Map every token id onto a smaller id space where tokens that normalize alike collapse together. + + N-grams are hashed over these compressed ids, so " The", "the" and "THE" all hash the same way. + Returns the lookup plus the size of the compressed vocab -- and that size matters beyond bounds + checking, because every hash multiplier is derived from it. + """ + from tokenizers import Regex, normalizers + + # a private-use char, so a token that is exactly one space survives Strip() instead of + # collapsing to the empty string and merging with unrelated tokens + sentinel = "\ue000" + normalizer = normalizers.Sequence( + [ + normalizers.NFKC(), + normalizers.NFD(), + normalizers.StripAccents(), + normalizers.Lowercase(), + normalizers.Replace(Regex(r"[ \t\r\n]+"), " "), + normalizers.Replace(Regex(r"^ $"), sentinel), + normalizers.Strip(), + normalizers.Replace(sentinel, " "), + ] + ) + + # the raw Rust tokenizer, matching what training decodes with (no clean_up_tokenization_spaces) + backend = tokenizer.backend_tokenizer + key_to_new: dict[str, int] = {} + lookup = [0] * len(tokenizer) + for token_id in range(len(tokenizer)): + text = backend.decode([token_id], skip_special_tokens=False) + if "\ufffd" in text: + # a partial UTF-8 byte token: nothing to normalize, so key it by its raw form + key = backend.id_to_token(token_id) + else: + normalized = normalizer.normalize_str(text) + key = normalized if normalized else text + + new_id = key_to_new.get(key) + if new_id is None: + new_id = len(key_to_new) + key_to_new[key] = new_id + lookup[token_id] = new_id + + return lookup, len(key_to_new) + + +def compute_hash_multipliers( + layer_ids: tuple[int, ...], max_ngram_size: int, tokenizer_vocab_size: int +) -> torch.Tensor: + """One multiplier per (layer, lookback), from a per-layer RNG so layers hash differently. + + Kept odd, and bounded so that `token_id * multiplier` cannot overflow int64. + """ + max_long = np.iinfo(np.int64).max + multiplier_bound = max(1, (max_long // tokenizer_vocab_size) // 2) + rows = [] + for layer_id in layer_ids: + generator = np.random.default_rng(10007 * layer_id) + values = generator.integers( + low=0, + high=multiplier_bound, + size=(max_ngram_size,), + dtype=np.int64, + ) + rows.append(torch.tensor(values * 2 + 1)) + return torch.stack(rows) + + +@dataclass(frozen=True) +class EngramLayout: + """Bucket layout of the n-gram hash tables. + + A position is hashed as `max_ngram_size - 1` n-grams (2-gram .. max_ngram_size-gram), each split + over `n_heads` heads. Every (n-gram size, head) pair owns its own prime-sized bucket range in the + layer's table; the primes are drawn in order and never reused, which keeps the ranges disjoint. + """ + + max_ngram_size: int + layer_ids: tuple[int, ...] + num_embeddings: tuple[int, ...] # table rows, per engram layer + primes: tuple[tuple[tuple[int, ...], ...], ...] # [layer][n-gram size][head] bucket modulus + n_heads: int + head_dim: int + + @classmethod + def from_args(cls, args) -> "EngramLayout | None": + layer_ids = tuple(args.engram_layer_ids) + if not layer_ids: + return None + max_ngram_size, n_heads = args.engram_max_ngram_size, args.engram_n_heads + primes, seen = [], set() + for _ in layer_ids: + per_ngram = [] + for _ in range(max_ngram_size - 1): + sizes, current = [], args.engram_vocab_size - 1 + for _ in range(n_heads): + current = find_next_prime(current, seen) + seen.add(current) + sizes.append(current) + per_ngram.append(tuple(sizes)) + primes.append(tuple(per_ngram)) + return cls( + max_ngram_size=max_ngram_size, + layer_ids=layer_ids, + num_embeddings=tuple(args.engram_num_embeddings), + primes=tuple(primes), + n_heads=n_heads, + head_dim=args.engram_head_dim, + ) + + +class NgramHashState(nn.Module): + """Maps each position to the hash ids of the n-grams ending there. + + Ids go through the compressed table, then each position is hashed with the `max_ngram_size - 1` + tokens before it. Look-back stops at the start of the sequence and at any dead token (an image + span, cached as DEAD), so an n-gram never spans one. The cache carries all of this across the + prefill/decode split. + """ + + DEAD = -1 + + def __init__(self, args, layout: EngramLayout, tokenizer): + super().__init__() + self.layout = layout + # every hash multiplier derives from the compressed vocab size, so a mismatch there would + # silently rehash the whole table + token_map, vocab_size = build_compressed_token_map(tokenizer) + assert vocab_size == args.engram_compressed_vocab_size, (vocab_size, args.engram_compressed_vocab_size) + self.pad_id = token_map[args.engram_pad_id] + flat = [[p for per_ngram in layer for p in per_ngram] for layer in layout.primes] + offsets = [np.cumsum([0, *sizes[:-1]]) for sizes in flat] + multipliers = compute_hash_multipliers(layout.layer_ids, layout.max_ngram_size, vocab_size) + self.register_buffer("primes", torch.tensor(layout.primes), persistent=False) + self.register_buffer("offsets", torch.tensor(np.array(offsets)), persistent=False) + self.register_buffer("multipliers", multipliers, persistent=False) + self.register_buffer("token_map", torch.tensor(token_map), persistent=False) + self.register_buffer( + "cache", torch.empty(args.max_batch_size, args.max_seq_len, dtype=torch.int64), persistent=False + ) + + @torch.inference_mode() + def forward(self, input_ids: torch.Tensor, start_pos: int, token_mask: torch.Tensor | None = None) -> torch.Tensor: + """token_mask: [B, L], False for tokens that take no part in an n-gram (image spans). + Returns the hash ids, shaped [B, L, n_engram_layers, n_hash_cols].""" + batch, seqlen = input_ids.shape + compressed = self.token_map[input_ids] + if token_mask is not None: + compressed = torch.where(token_mask, compressed, self.DEAD) + self.cache[:batch, start_pos : start_pos + seqlen] = compressed + + positions = torch.arange(start_pos, start_pos + seqlen, device=input_ids.device).expand(batch, seqlen) + tokens, blocked = [], torch.zeros_like(positions, dtype=torch.bool) + for shift in range(self.layout.max_ngram_size): + source = self.cache[:batch].gather(1, (positions - shift).clamp_min(0)) + blocked = blocked | (positions < shift) | (source == self.DEAD) + tokens.append(torch.where(blocked, self.pad_id, source)) + tokens = torch.stack(tokens, dim=-1) # [B, L, max_ngram_size] + + # XOR the multiplied ids together one lookback at a time, so the running value after step i + # is the hash of the (i+1)-gram; each lands in its own prime-sized bucket range + products = tokens.unsqueeze(2) * self.multipliers # [B, L, n_engram_layers, max_ngram_size] + rolling, hashes = products[..., 0], [] + for i in range(1, self.layout.max_ngram_size): + rolling = torch.bitwise_xor(rolling, products[..., i]) + hashes.append(rolling.unsqueeze(-1) % self.primes[:, i - 1]) + return torch.cat(hashes, dim=-1) + self.offsets diff --git a/tests/models/deepseek_v41/reference/kernel.py b/tests/models/deepseek_v41/reference/kernel.py new file mode 100644 index 0000000000..fab0ae08e2 --- /dev/null +++ b/tests/models/deepseek_v41/reference/kernel.py @@ -0,0 +1,113 @@ +"""Pure-torch transcriptions of the reference tilelang kernels (inference/kernel.py), for the test oracle. + +Same numerics as the tilelang code: fp32 math, power-of-two (ue8m0) or e4m3 scales rounded the way +the kernels round them, the e2m1 grid with the hardware's round-to-nearest-even, and the fp8/fp4 +in-place round trips that write the dequantized value back in the input dtype. Quantized GEMMs +dequantize independently in torch and round each projection's output to bf16. +""" + +from __future__ import annotations + +import torch + +_FP4_GRID = torch.tensor([0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0]) + + +def _pow2_ceil(x: torch.Tensor) -> torch.Tensor: + """``fast_round_scale``: 2 ** ceil(log2(x)) for x > 0.""" + return torch.exp2(torch.ceil(torch.log2(x))) + + +def _round_fp4(x: torch.Tensor) -> torch.Tensor: + """Round-to-nearest-even onto the signed e2m1 grid (the ``float4_e2m1fn`` cast).""" + a = x.abs() + grid = _FP4_GRID.to(x.device) + hi = torch.searchsorted(grid, a.reshape(-1).contiguous(), right=True).reshape(a.shape).clamp(1, 7) + below, above = grid[hi - 1], grid[hi] + mid = (below + above) / 2 + even_above = torch.isin(above, torch.tensor([0.0, 1.0, 2.0, 4.0], device=x.device)) + r = torch.where((a > mid) | ((a == mid) & even_above), above, below) + r = torch.where(a >= 6.0, torch.full_like(r, 6.0), r) + return torch.copysign(r, x) + + +def act_quant(x: torch.Tensor, block_size: int = 128, scale_fmt=None, scale_dtype=torch.float32, inplace: bool = False): + """Block-wise fp8 quantization; ``inplace=True`` writes the quant+dequant round trip back into ``x``.""" + n = x.size(-1) + assert n % block_size == 0 + g = x.float().unflatten(-1, (-1, block_size)) + amax = g.abs().amax(dim=-1).clamp_min(1e-4) + s = _pow2_ceil(amax / 448.0) if scale_fmt is not None else amax / 448.0 + q = (g / s.unsqueeze(-1)).clamp(-448.0, 448.0).to(torch.float8_e4m3fn) + if inplace: + x.copy_((q.float() * s.unsqueeze(-1)).flatten(-2).to(x.dtype)) + return x + return q.flatten(-2), s.to(scale_dtype) + + +def fp4_act_quant(x: torch.Tensor, block_size: int = 32, inplace: bool = False, scale_dtype=torch.float8_e8m0fnu): + """Block-wise fp4 with ue8m0 (pow2) or e4m3 scales; ``inplace=True`` writes the round trip back.""" + n = x.size(-1) + assert n % block_size == 0 + g = x.float().unflatten(-1, (-1, block_size)) + amax = g.abs().amax(dim=-1) + if scale_dtype == torch.float8_e4m3fn: + amax = amax.clamp_min(6 * 2**-9) + s = (amax / 6.0).to(torch.float8_e4m3fn).float() + else: + amax = amax.clamp_min(6 * 2**-126) + s = _pow2_ceil(amax / 6.0) + q = _round_fp4((g / s.unsqueeze(-1)).clamp(-6.0, 6.0)) + if inplace: + x.copy_((q * s.unsqueeze(-1)).flatten(-2).to(x.dtype)) + return x + return q.flatten(-2), s.to(scale_dtype) + + +def fp8_gemm(a, a_s, b, b_s, scale_dtype=torch.float32, block_size=128): + """``C = (A * s_a) @ (B * s_b)^T`` over dequantized operands.""" + k = a.size(-1) + a_deq = (a.float().unflatten(-1, (-1, block_size)) * a_s.float().unsqueeze(-1)).flatten(-2) + n = b.size(0) + bs = b_s.float().repeat_interleave(block_size, 0)[:n].repeat_interleave(block_size, 1)[:, :k] + return (a_deq @ (b.float() * bs).t()).bfloat16() + + +def fp4_gemm(a, a_s, b, b_s, scale_dtype=torch.float8_e8m0fnu, act_block_size=32): + packed = b.view(torch.uint8) + codes = torch.stack((packed & 15, packed >> 4), dim=-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=b.device) + weight = grid[codes] * b_s.float().repeat_interleave(32, -1) + x = (a.float().unflatten(-1, (-1, act_block_size)) * a_s.float().unsqueeze(-1)).flatten(-2) + return (x @ weight.t()).bfloat16() + + +def sparse_attn(q: torch.Tensor, kv: torch.Tensor, attn_sink: torch.Tensor, topk_idxs: torch.Tensor, softmax_scale: float) -> torch.Tensor: + """Per (batch, query): softmax over the gathered ``topk_idxs`` (-1 = absent) plus the sink null key. + + ``q [b, m, h, d]``, ``kv [b, n, d]``, ``topk_idxs [b, m, topk]`` -> ``o [b, m, h, d]`` in q's dtype. + A row with no valid index yields zeros (the finite ``-1e30`` running max of the kernel).""" + b, m, h, d = q.shape + valid = topk_idxs >= 0 + gathered = kv[torch.arange(b, device=q.device)[:, None, None], topk_idxs.clamp_min(0).long()] # [b, m, topk, d] + scores = torch.einsum("bmhd,bmtd->bmht", q.float(), gathered.float()) * softmax_scale + scores = scores.masked_fill(~valid[:, :, None, :], float("-inf")) + mx = scores.amax(dim=-1).clamp_min(-1e30) + p = torch.exp(scores - mx.unsqueeze(-1)) + denom = p.sum(-1) + torch.exp(attn_sink.float()[None, None, :] - mx) + o = torch.einsum("bmht,bmtd->bmhd", p, gathered.float()) / denom.unsqueeze(-1) + return o.to(q.dtype) + + +def hc_split_sinkhorn(mixes: torch.Tensor, hc_scale: torch.Tensor, hc_base: torch.Tensor, hc_mult: int = 4, sinkhorn_iters: int = 20, eps: float = 1e-6): + hc = hc_mult + mixes, sc, base = mixes.float(), hc_scale.float(), hc_base.float() + pre = torch.sigmoid(mixes[..., :hc] * sc[0] + base[:hc]) + eps + post = 2 * torch.sigmoid(mixes[..., hc : 2 * hc] * sc[1] + base[hc : 2 * hc]) + comb = (mixes[..., 2 * hc :] * sc[2] + base[2 * hc :]).unflatten(-1, (hc, hc)) + comb = comb.softmax(dim=-1) + eps + comb = comb / (comb.sum(dim=-2, keepdim=True) + eps) + for _ in range(sinkhorn_iters - 1): + comb = comb / (comb.sum(dim=-1, keepdim=True) + eps) + comb = comb / (comb.sum(dim=-2, keepdim=True) + eps) + return pre, post, comb diff --git a/tests/models/deepseek_v41/reference/model.py b/tests/models/deepseek_v41/reference/model.py new file mode 100644 index 0000000000..593c7518c4 --- /dev/null +++ b/tests/models/deepseek_v41/reference/model.py @@ -0,0 +1,1458 @@ +# Vendored from https://huggingface.co/deepseek-ai/DeepSeek-V4.1-Flash (inference/model.py), MIT License +# (see LICENSE next to this file). FreeToken keeps this copy as the numerical ORACLE its port is tested +# against. Edits, all marked "FreeToken:", are limited to (1) local imports with the tilelang kernels +# replaced by the pure-torch transcriptions in kernel.py, and (2) forward_bounded(), a Decoder SWA +# Bounded Replay oracle built from the unmodified layer math. +import math +from contextlib import contextmanager +from dataclasses import dataclass +from functools import lru_cache +from typing import Literal + +import torch +import torch.distributed as dist +import torch.nn.functional as F +from torch import nn + +# FreeToken: local imports. +from .engram import EngramLayout, NgramHashState +from .kernel import ( + act_quant, + fp4_act_quant, + fp4_gemm, + fp8_gemm, + hc_split_sinkhorn, + sparse_attn, +) + +from .vision import ViT, Aligner + +IMAGE_START, IMAGE, IMAGE_NEW_LINE, IMAGE_END = range(4) # FreeToken: image_processor constants + +# Set once by Transformer.__init__; one model per process, so layers just read them. +world_size = 1 +rank = 0 +default_dtype = torch.float8_e4m3fn # storage dtype for Linear weights, from ModelArgs.dtype +fp8_block_size = 32 # one fp8 scale per 32x32 weight block / 32 activations +fp4_block_size = 32 # one fp4 scale per 32 elements along K +scale_fmt = "ue8m0" +scale_dtype = torch.float8_e8m0fnu + + +@contextmanager +def set_dtype(dtype): + """Temporarily override torch's default dtype, restoring it even if the body raises.""" + prev = torch.get_default_dtype() + torch.set_default_dtype(dtype) + try: + yield + finally: + torch.set_default_dtype(prev) + + +@dataclass +class ModelArgs: + """Field names are exactly the config JSON keys. The defaults are a small model that + `python model.py` can run, not the released shapes -- though the scale-independent + values (norm_eps, score_func, hc_*, engram_*) do match it.""" + + # runtime limits rather than model shape: they size the KV caches + max_batch_size: int = 4 + max_seq_len: int = 4096 + temperature: float = 1 + dtype: Literal["bf16", "fp8"] = "fp8" + expert_dtype: Literal["fp4"] | None = "fp4" + vocab_size: int = 129280 + dim: int = 1024 + moe_inter_dim: int = 1024 + n_layers: int = 5 + n_mtp_layers: int = 1 # extra draft layers appended after the backbone, indices n_layers.. + n_heads: int = 16 + # moe + n_routed_experts: int = 8 + n_shared_experts: int = 1 + n_activated_experts: int = 2 + score_func: Literal["softmax", "sigmoid", "sqrtsoftplus"] = "sqrtsoftplus" + gate_temp: float = 1.0 + norm_topk_prob: bool = True + route_scale: float = 1.0 + swiglu_limit: float = 0.0 + # attention: latent q/kv projections, plus a LoRA-factorised output projection over o_groups + q_lora_rank: int = 256 + head_dim: int = 128 + rope_head_dim: int = 32 + norm_eps: float = 1e-20 + o_groups: int = 8 + o_lora_rank: int = 256 + # sparse attention: every layer attends over a sliding window, and may add compressed KV on top + window_size: int = 128 + # one entry per layer, MTP layers included: 0 = sliding window only, r = KV compressed r-to-1 + compress_ratios: tuple[int, ...] = (0, 2, 2, 1, 1, 0) + # layers sharing a ratio also share one compressed KV and one indexer, produced by the first + kv_source_layers: tuple[int, ...] = (1, 3) + index_source_layers: tuple[int, ...] = (1, 3) + # rope, with YaRN extrapolation when original_seq_len > 0. Compressed KV rotates at its own + # theta because one latent stands for compress_ratio tokens, so its positions are further apart. + compress_rope_theta: float = 40000.0 + original_seq_len: int = 0 + rope_theta: float = 10000.0 + rope_factor: float = 40 + beta_fast: int = 32 + beta_slow: int = 1 + # the indexer: a small extra attention that scores compressed positions, so each query can keep + # just `index_topk` of them. Names match DeepSeek-V3.2-Exp, where this mechanism first appeared. + index_n_heads: int = 16 + index_head_dim: int = 64 + index_topk: int = 64 + # candidate pre-filtering: candidate_source_layer < 0 turns it off and the other two are unused + candidate_source_layer: int = -1 + candidate_topk_blocks: int = 0 + candidate_block_size: int = 0 + # hyper-connections: the residual stream is carried as hc_mult parallel copies + hc_mult: int = 4 + hc_sinkhorn_iters: int = 20 + hc_eps: float = 1e-6 + # engram: n-gram hash lookups added into the residual stream at a few layers + engram_layer_ids: tuple[int, ...] = () + engram_num_embeddings: tuple[int, ...] = () # unpadded table rows; each rank allocates ceil(rows / world_size) + engram_max_ngram_size: int = 1 + engram_vocab_size: int = 0 # bucket size each (n-gram size, head) starts searching primes from + engram_n_heads: int = 0 + engram_head_dim: int = 0 + engram_pad_id: int = 2 # token that fills n-gram slots with no history; matches training + # size of the compressed tokenizer vocab; every hash multiplier is derived from it + engram_compressed_vocab_size: int = 0 + # vision (VL); vision_n_layers == 0 disables the vision path + vision_n_layers: int = 0 + vision_dim: int = 1024 + vision_n_heads: int = 16 + vision_inter_dim: int = 2816 + vision_patch_size: int = 14 + vision_rope_theta: float = 10000.0 + vision_downsample_ratio: int = 3 + vision_max_n_token: int = 1024 + vision_min_pixels: int = 544 * 544 + vision_max_wh_ratio: int | None = None + # raw id of <｜deepseek_image｜>; every position of an image span carries this id in input_ids + image_token_id: int = 129264 + # dspark draft head. Only the forward pass is implemented here -- nothing calls forward_spec, + # so these are read but the speculative-decoding loop itself is out of scope for this repo. + dspark_block_size: int = 0 + dspark_noise_token_id: int = 0 + dspark_target_layer_ids: tuple[int, ...] = () + dspark_markov_rank: int = 256 + dspark_n_routed_experts: int = 0 + dspark_n_activated_experts: int = 0 + + @property + def vision_enabled(self) -> bool: + return self.vision_n_layers > 0 + + def get_moe_config(self, layer_id: int) -> tuple[int, int]: + """Return the routed/activated expert counts for a given layer.""" + if layer_id < self.n_layers: + return self.n_routed_experts, self.n_activated_experts + return ( + self.dspark_n_routed_experts or self.n_routed_experts, + self.dspark_n_activated_experts or self.n_activated_experts, + ) + + +class ParallelEmbedding(nn.Module): + """Embedding sharded along the vocab dimension. Each rank holds vocab_size // world_size rows. + Out-of-range indices are zero-masked before all_reduce to combine partial embeddings.""" + + def __init__(self, vocab_size: int, dim: int): + super().__init__() + self.vocab_size = vocab_size + self.dim = dim + assert vocab_size % world_size == 0, ( + f"Vocabulary size must be divisible by world size (world_size={world_size})" + ) + self.part_vocab_size = vocab_size // world_size + self.vocab_start_idx = rank * self.part_vocab_size + self.vocab_end_idx = self.vocab_start_idx + self.part_vocab_size + self.weight = nn.Parameter(torch.empty(self.part_vocab_size, self.dim)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if world_size > 1: + # ids off this rank read row 0 then get zeroed, so the all_reduce sums one real row + mask = (x < self.vocab_start_idx) | (x >= self.vocab_end_idx) + x = x - self.vocab_start_idx + x[mask] = 0 + y = F.embedding(x, self.weight) + if world_size > 1: + y[mask] = 0 + dist.all_reduce(y) + return y + + +def linear(x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor | None = None) -> torch.Tensor: + """Pick a GEMM from the weight dtype. Quantized weights need a quantized activation, and both + fp4 and fp8 weights take an fp8 one -- for fp4 the kernel handles the mixed precision.""" + assert bias is None + + if weight.dtype == torch.float4_e2m1fn_x2: + x, s = act_quant(x, fp8_block_size, scale_fmt, scale_dtype) + return fp4_gemm( + x, + s, + weight, + weight.scale, + scale_dtype, + act_block_size=fp8_block_size, + ) + elif weight.dtype == torch.float8_e4m3fn: + x, s = act_quant(x, fp8_block_size, scale_fmt, scale_dtype) + return fp8_gemm( + x, + s, + weight, + weight.scale, + scale_dtype, + block_size=fp8_block_size, + ) + else: + return F.linear(x, weight) + + +class Linear(nn.Module): + """bf16, fp8 or fp4 weights. Quantized ones get a `scale`, also attached to `.weight` so that + `linear()` can reach it.""" + + def __init__(self, in_features: int, out_features: int, bias: bool = False, dtype=None): + super().__init__() + self.in_features = in_features + self.out_features = out_features + dtype = dtype or default_dtype + if dtype == torch.float4_e2m1fn_x2: + # two values per byte: [out, in] logically, [out, in//2] stored + self.weight = nn.Parameter(torch.empty(out_features, in_features // 2, dtype=torch.float4_e2m1fn_x2)) + self.weight.scale = self.scale = nn.Parameter( + torch.empty(out_features, in_features // fp4_block_size, dtype=torch.float8_e8m0fnu) + ) + elif dtype == torch.float8_e4m3fn: + self.weight = nn.Parameter(torch.empty(out_features, in_features, dtype=dtype)) + self.weight.scale = self.scale = nn.Parameter( + torch.empty( + (out_features + fp8_block_size - 1) // fp8_block_size, + (in_features + fp8_block_size - 1) // fp8_block_size, + dtype=torch.float8_e8m0fnu, + ) + ) + else: + self.weight = nn.Parameter(torch.empty(out_features, in_features, dtype=dtype)) + self.register_parameter("scale", None) + if bias: + self.bias = nn.Parameter(torch.empty(out_features)) + else: + self.register_parameter("bias", None) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return linear(x, self.weight, self.bias) + + +class ColumnParallelLinear(Linear): + """Splits the output dim across ranks; each rank's slice of the output is already complete.""" + + def __init__(self, in_features: int, out_features: int, bias: bool = False, dtype=None): + assert out_features % world_size == 0, ( + f"Output features must be divisible by world size (world_size={world_size})" + ) + self.part_out_features = out_features // world_size + super().__init__(in_features, self.part_out_features, bias, dtype) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return linear(x, self.weight, self.bias) + + +class RowParallelLinear(Linear): + """Splits the reduction dim, so each rank holds a partial sum: hence the fp32 all_reduce, with + the bias added only after it.""" + + def __init__(self, in_features: int, out_features: int, bias: bool = False, dtype=None): + assert in_features % world_size == 0, ( + f"Input features must be divisible by world size (world_size={world_size})" + ) + self.part_in_features = in_features // world_size + super().__init__(self.part_in_features, out_features, bias, dtype) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + y = linear(x, self.weight, None) + if world_size > 1: + y = y.float() + dist.all_reduce(y) + if self.bias is not None: + y += self.bias + return y.type_as(x) + + +class RMSNorm(nn.Module): + def __init__(self, dim: int, eps: float = 1e-6): + super().__init__() + self.dim = dim + self.eps = eps + self.weight = nn.Parameter(torch.ones(dim)) + + def forward(self, x: torch.Tensor): + dtype = x.dtype + x = x.float() + var = x.square().mean(-1, keepdim=True) + x = x * torch.rsqrt(var + self.eps) + return (self.weight * x).to(dtype) + + +class ParallelEngramEmbedding(nn.Module): + """The n-gram hash table, sharded over its rows. Stays fp8: rows are dequantized on lookup.""" + + def __init__(self, num_embeddings: int, dim: int): + super().__init__() + self.num_embeddings = num_embeddings + self.dim = dim + self.part_num_embeddings = (num_embeddings + world_size - 1) // world_size + self.vocab_start_idx = rank * self.part_num_embeddings + self.vocab_end_idx = self.vocab_start_idx + self.part_num_embeddings + + self.block_size = fp8_block_size + # the table stays fp8 as stored: rows are dequantized with `scale` on lookup + self.weight = nn.Parameter(torch.empty(self.part_num_embeddings, dim, dtype=torch.float8_e4m3fn)) + self.scale = nn.Parameter(torch.empty(self.part_num_embeddings, dim // self.block_size, dtype=scale_dtype)) + + def forward(self, indices: torch.Tensor) -> torch.Tensor: + mask = (indices < self.vocab_start_idx) | (indices >= self.vocab_end_idx) + local_indices = indices - self.vocab_start_idx + local_indices = local_indices.masked_fill(mask, 0) + + values = F.embedding(local_indices, self.weight) + scales = F.embedding(local_indices, self.scale) + values = values.float().unflatten(-1, (-1, self.block_size)) * scales.float().unsqueeze(-1) + values = values.flatten(-2).to(torch.bfloat16) + values = values.masked_fill(mask.unsqueeze(-1), 0) + + if world_size > 1: + dist.all_reduce(values) + return values + + +class Engram(nn.Module): + """Writes an n-gram lookup into the residual stream, gated by how well it matches that stream. + + The hash ids fetch `n_hash_cols` rows; `wkv` turns them into one key per hc copy plus a shared + value. The gate is a normalized dot product of stream against key. + """ + + def __init__(self, args: ModelArgs, layer_id: int, layout: EngramLayout): + super().__init__() + self.layer_id = layer_id + self.layer_hash_index = layout.layer_ids.index(layer_id) + self.dim = args.dim + self.hc_mult = args.hc_mult + self.clamp_value = 1e-6 + + self.embed = ParallelEngramEmbedding(layout.num_embeddings[self.layer_hash_index], layout.head_dim) + n_hash_cols = (layout.max_ngram_size - 1) * layout.n_heads + self.wkv = Linear(n_hash_cols * layout.head_dim, args.dim * (args.hc_mult + 1)) + self.eps = args.norm_eps + self.q_weight = nn.Parameter(torch.ones(args.hc_mult, args.dim)) + self.k_weight = nn.Parameter(torch.ones(args.hc_mult, args.dim)) + + def forward(self, x: torch.Tensor, hash_ids: torch.Tensor, token_mask: torch.Tensor | None = None) -> torch.Tensor: + """x: [B, L, hc_mult, dim]; hash_ids: [B, L, n_hash_cols]; token_mask: [B, L], False shuts + the gate so those positions pass through untouched.""" + kv = self.wkv(self.embed(hash_ids).flatten(-2)) + key, value = kv.split([self.hc_mult * self.dim, self.dim], dim=-1) + key = key.float().unflatten(-1, (self.hc_mult, self.dim)) + weight = self.q_weight.float() * self.k_weight.float() # only ever used as a product + h, eps = x.float(), self.eps + # normalized per (token, hc copy) over `dim`, NOT jointly over the copies + rstd = torch.rsqrt(h.square().mean(-1) + eps) * torch.rsqrt(key.square().mean(-1) + eps) + dot = (h * weight * key).sum(-1) * rstd * self.dim**-0.5 + # signed sqrt before the sigmoid, matching the training kernel + gate = torch.sigmoid(torch.copysign(dot.abs().clamp_min(self.clamp_value).sqrt(), dot)) + if token_mask is not None: + gate = gate.masked_fill(~token_mask.unsqueeze(-1), 0) + return (h + gate.unsqueeze(-1) * value.float().unsqueeze(-2)).to(x.dtype) + + +@lru_cache(2) +def precompute_freqs_cis(dim, seqlen, original_seq_len, base, factor, beta_fast, beta_slow) -> torch.Tensor: + """Rotary frequencies as complex exponentials, one row per position. + + With original_seq_len > 0 this applies YaRN: dimensions whose wavelength already fits inside the + training context keep their frequency, those far beyond it are divided by `factor`, and the + `beta_fast`..`beta_slow` band in between is faded across with a linear ramp. + """ + freqs = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim)) + if original_seq_len > 0: + # the dim whose wavelength completes `rotations` turns over the training context + def corrected_dim(rotations): + return dim * math.log(original_seq_len / (rotations * 2 * math.pi)) / (2 * math.log(base)) + + low = max(math.floor(corrected_dim(beta_fast)), 0) + high = min(math.ceil(corrected_dim(beta_slow)), dim - 1) + ramp = ((torch.arange(dim // 2, dtype=torch.float32) - low) / max(high - low, 1e-3)).clamp(0, 1) + smooth = 1 - ramp + freqs = freqs / factor * (1 - smooth) + freqs * smooth + + freqs = torch.outer(torch.arange(seqlen), freqs) + return torch.polar(torch.ones_like(freqs), freqs) + + +def apply_rotary_emb(x: torch.Tensor, freqs_cis: torch.Tensor, inverse: bool = False) -> torch.Tensor: + """Rotate `x` in place, taking adjacent element pairs as complex numbers. Accepts [b, s, d] and + [b, s, h, d]; `inverse` conjugates the rotation, which is how the attention output gets the + query's rotation removed again so the cache can stay in one shared rotated form.""" + y = x + x = torch.view_as_complex(x.float().unflatten(-1, (-1, 2))) + if inverse: + freqs_cis = freqs_cis.conj() + if x.ndim == 3: + freqs_cis = freqs_cis.view(1, x.size(1), x.size(-1)) + else: + freqs_cis = freqs_cis.view(1, x.size(1), 1, x.size(-1)) + x = torch.view_as_real(x * freqs_cis).flatten(-2) + y.copy_(x) + return y + + +@lru_cache(1) +def get_window_topk_idxs(window_size: int, bsz: int, seqlen: int, start_pos: int): + """Which sliding-window cache slots each query attends to; -1 marks a slot holding nothing. + + The cache is a ring of `window_size` slots. Prefill needs one row per query, each seeing its own + causal window. A decode step has a single query that sees the whole ring, listed oldest first. + Order within a row does not matter to `sparse_attn`, which handles every slot independently. + """ + if start_pos == 0: + end = torch.arange(seqlen).unsqueeze(1) + idxs = (end - window_size + 1).clamp(0) + torch.arange(min(seqlen, window_size)) + idxs = torch.where(idxs > end, -1, idxs) # before the sequence started + else: + oldest = start_pos % window_size + 1 + idxs = torch.cat([torch.arange(oldest, window_size), torch.arange(oldest)]) + idxs = torch.where(idxs > start_pos, -1, idxs) # ring still filling + # sparse_attn needs real [b, m, topk] int32 memory, hence the materializing expand + return idxs.int().unsqueeze(0).expand(bsz, -1, -1).contiguous() + + +class Compressor(nn.Module): + """Pools `compress_ratio` consecutive tokens into one KV latent with a learned softmax gate. + + Returns the latent before RoPE, or None while a group is still filling up -- so during decode it + only yields every `compress_ratio` steps, holding the partial group in `kv_state`/`score_state`. + Pre-RoPE is deliberate: the indexer needs the unrotated form, so Attention rotates afterwards. + """ + + def __init__(self, args: ModelArgs, layer_id: int): + super().__init__() + compress_ratio = args.compress_ratios[layer_id] + head_dim = args.head_dim + self.compress_ratio = compress_ratio + self.head_dim = head_dim + self.norm = RMSNorm(head_dim, args.norm_eps) + # ratio 1 is a plain projection, so it stays in the checkpoint's bf16; the softmax pooling + # above ratio 1 runs in fp32, so those weights are promoted to fp32 to match + self.wkv = Linear(args.dim, head_dim, dtype=torch.float32 if compress_ratio > 1 else torch.bfloat16) + if compress_ratio > 1: + self.wgate = Linear(args.dim, head_dim, dtype=torch.float32) + # tail of an incomplete group, carried across decode steps + self.kv_state: torch.Tensor + self.score_state: torch.Tensor + state_shape = (args.max_batch_size, compress_ratio, head_dim) + self.register_buffer("kv_state", torch.zeros(state_shape, dtype=torch.float32), persistent=False) + self.register_buffer( + "score_state", torch.full(state_shape, -torch.inf, dtype=torch.float32), persistent=False + ) + + def forward(self, x: torch.Tensor, start_pos: int) -> torch.Tensor | None: + bsz, seqlen, _ = x.size() + ratio, dtype = self.compress_ratio, x.dtype + if ratio == 1: # one token per group: nothing to pool, so no gate and no fp32 + return self.norm(self.wkv(x)) + + x = x.float() + kv, score = self.wkv(x), self.wgate(x) + if start_pos == 0: + should_compress = seqlen >= ratio + remainder = seqlen % ratio + cutoff = seqlen - remainder + if remainder: # trailing partial group waits in the state + kv, self.kv_state[:bsz, :remainder] = kv.split([cutoff, remainder], dim=1) + score, self.score_state[:bsz, :remainder] = score.split([cutoff, remainder], dim=1) + kv = kv.unflatten(1, (-1, ratio)) + score = score.unflatten(1, (-1, ratio)) + kv = (kv * score.softmax(dim=2)).sum(dim=2) + else: # one token per step: fill a slot, and pool only when the group just completed + should_compress = (start_pos + 1) % ratio == 0 + slot = start_pos % ratio + self.kv_state[:bsz, slot] = kv.squeeze(1) + self.score_state[:bsz, slot] = score.squeeze(1) + if should_compress: + kv = (self.kv_state[:bsz] * self.score_state[:bsz].softmax(dim=1)).sum(dim=1, keepdim=True) + if not should_compress: + return None + return self.norm(kv.to(dtype)) + + +class Indexer(torch.nn.Module): + """Keeps the `index_topk` best compressed positions per query. + + A small side attention: fp4 query heads against one shared key per compressed position, scores + rectified then combined by `weights_proj`. With a candidate source this is the second of two + levels; `select_candidate_blocks` is the first. + """ + + def __init__(self, args: ModelArgs, layer_id: int): + super().__init__() + # the index keys are derived from the compressor's latent, so only a layer that compresses + # its own KV can produce them; every other indexer reads them from that layer's cache + self.owns_k = layer_id in args.kv_source_layers + self.compress_ratio = args.compress_ratios[layer_id] + self.is_candidate_source = layer_id == args.candidate_source_layer + self.uses_candidates = 0 <= args.candidate_source_layer < layer_id + self.candidate_topk_blocks = args.candidate_topk_blocks + self.candidate_block_size = args.candidate_block_size + self.dim = args.dim + self.n_heads = args.index_n_heads + self.n_local_heads = args.index_n_heads // world_size + self.index_head_dim = args.index_head_dim + self.rope_head_dim = args.rope_head_dim + self.index_topk = args.index_topk + self.q_lora_rank = args.q_lora_rank + self.softmax_scale = self.index_head_dim**-0.5 + self.wq_b = ColumnParallelLinear(self.q_lora_rank, self.n_heads * self.index_head_dim) + self.weights_proj = ColumnParallelLinear(self.dim, self.n_heads, dtype=torch.bfloat16) + self.freqs_cis: torch.Tensor | None = None + if self.owns_k: + self.wk = Linear(args.head_dim, self.index_head_dim, dtype=torch.bfloat16) + self.k_norm = RMSNorm(self.index_head_dim, args.norm_eps) + self.k_cache: torch.Tensor + self.register_buffer( + "k_cache", + torch.zeros(args.max_batch_size, args.max_seq_len // self.compress_ratio, args.index_head_dim), + persistent=False, + ) + + def forward(self, x: torch.Tensor, qr: torch.Tensor, latent: torch.Tensor, start_pos: int, offset: int): + """`latent` is this layer's RoPE-free compressed latent, None when this layer does not + compress or when its current group is still incomplete. An index-key owner turns it into + index keys here, which has to happen before Attention overwrites that same storage with + the RoPE'd, quantized values.""" + assert self.freqs_cis is not None + bsz, seqlen, _ = x.size() + ratio, rd, end_pos = self.compress_ratio, self.rope_head_dim, start_pos + seqlen + + # latent is None while a group is still filling up, so there is nothing to publish yet + if self.owns_k and latent is not None: + # a latent stands for the first token of its group, so group j takes position j * ratio + freqs = ( + self.freqs_cis[: seqlen - seqlen % ratio : ratio] + if start_pos == 0 + else self.freqs_cis[start_pos + 1 - ratio].unsqueeze(0) + ) + k = self.k_norm(self.wk(latent)) + apply_rotary_emb(k[..., -rd:], freqs) + fp4_act_quant(k, fp4_block_size, True) + self.k_cache[:bsz, start_pos // ratio : start_pos // ratio + k.size(1)] = k + shared_attn.index_k = self.k_cache + + q = self.wq_b(qr).unflatten(-1, (self.n_local_heads, self.index_head_dim)) + apply_rotary_emb(q[..., -rd:], self.freqs_cis[start_pos:end_pos]) + fp4_act_quant(q, fp4_block_size, True) + + index_k = shared_attn.index_k[:bsz, : end_pos // ratio] + weights = self.weights_proj(x) * (self.softmax_scale * self.n_heads**-0.5) + index_score = torch.einsum("bshd,btd->bsht", q, index_k) + index_score = (index_score.relu_() * weights.unsqueeze(-1)).sum(dim=2) + if world_size > 1: + dist.all_reduce(index_score) + + # how many compressed positions each query can see: a block becomes visible once the query + # has passed its last token. One query per decode step, so there it is just a number. + if start_pos == 0: + compress_lens = (torch.arange(1, seqlen + 1, device=x.device) // ratio).unsqueeze(-1) + index_score.masked_fill_(torch.arange(seqlen // ratio, device=x.device) >= compress_lens, -torch.inf) + else: + compress_lens = end_pos // ratio + + if self.is_candidate_source: + shared_attn.candidates = select_candidate_blocks( + index_score, compress_lens, self.candidate_topk_blocks, self.candidate_block_size + ) + elif self.uses_candidates: + # level two: score with our own weights, but only inside the source's candidate blocks + index_score = index_score.masked_fill(~shared_attn.candidates, -torch.inf) + + # top-k by score, re-sorted into position order; unreachable -> -1, rest shifted by offset + topk = min(self.index_topk, end_pos // ratio) + idxs = index_score.topk(topk, dim=-1, sorted=False).indices.sort(dim=-1).values + return torch.where(idxs < compress_lens, idxs + offset, -1).int() + + # FreeToken: Decoder SWA Bounded Replay oracle -- the indexer for a TAIL prefill: queries are the + # positions [s, s + n) of a prompt of L_all tokens; the keys (a source's) were published for all + # L_all tokens by forward()'s start_pos == 0 path. Same scoring / masking / top-k as forward(). + def forward_tail(self, x, qr, s: int, n: int, offset: int): + assert self.freqs_cis is not None + bsz = x.size(0) + ratio, rd = self.compress_ratio, self.rope_head_dim + q = self.wq_b(qr).unflatten(-1, (self.n_local_heads, self.index_head_dim)) + apply_rotary_emb(q[..., -rd:], self.freqs_cis[s : s + n]) + fp4_act_quant(q, fp4_block_size, True) + end_pos = s + n + index_k = shared_attn.index_k[:bsz, : end_pos // ratio] + weights = self.weights_proj(x) * (self.softmax_scale * self.n_heads**-0.5) + index_score = torch.einsum("bshd,btd->bsht", q, index_k) + index_score = (index_score.relu_() * weights.unsqueeze(-1)).sum(dim=2) + compress_lens = ((s + torch.arange(1, n + 1, device=x.device)) // ratio).unsqueeze(-1) + index_score.masked_fill_(torch.arange(end_pos // ratio, device=x.device) >= compress_lens, -torch.inf) + if self.is_candidate_source: + shared_attn.candidates = select_candidate_blocks( + index_score, compress_lens, self.candidate_topk_blocks, self.candidate_block_size + ) + elif self.uses_candidates: + index_score = index_score.masked_fill(~shared_attn.candidates, -torch.inf) + topk = min(self.index_topk, end_pos // ratio) + idxs = index_score.topk(topk, dim=-1, sorted=False).indices.sort(dim=-1).values + return torch.where(idxs < compress_lens, idxs + offset, -1).int() + + +def select_candidate_blocks( + logits: torch.Tensor, + compress_lens: torch.Tensor | int, + topk_blocks: int, + block_size: int, +) -> torch.Tensor: + """Level one of the two-level top-k: keep the `topk_blocks` highest-scoring blocks per query. + + `logits` is [..., n_positions] with positions the query cannot reach already at -inf, which is + what makes a block score of -inf mean "not reachable yet". `compress_lens` is a plain int during + decode, or broadcasts against logits' leading dims during prefill. Returns a bool mask shaped + like `logits`, so the layers consuming it just mask and never think about blocks again. + """ + width = logits.size(-1) + # score each block by its best position; -inf pads the last one out to block_size + scores = F.pad(logits, (0, -width % block_size), value=-torch.inf) + scores = scores.unflatten(-1, (-1, block_size)).amax(dim=-1) + num_blocks = scores.size(-1) + + # the block with this query's newest position is only partly filled, so pin it in: it holds the + # most recent tokens but could otherwise be outscored by an older, full block + last = (compress_lens - 1) // block_size + scores = scores.masked_fill(torch.arange(num_blocks, device=logits.device) == last, torch.inf) + + top = scores.topk(min(topk_blocks, num_blocks), dim=-1) + # fewer reachable blocks than topk_blocks means leftover picks came back -inf: drop them + keep = torch.zeros_like(scores, dtype=torch.bool).scatter_(-1, top.indices, top.values > -torch.inf) + return keep.repeat_interleave(block_size, dim=-1)[..., :width] + + +class Attention(nn.Module): + """Latent attention over two KV sources at once, concatenated into one `sparse_attn` call: a + sliding window of raw KV, plus -- when compress_ratio > 0 -- `index_topk` compressed positions + reaching further back. Q and the output projection are both low-rank, the latter grouped. + + compress_ratio > 0 does not mean the layer compresses its own KV: only kv_source_layers do, + the rest read that same cache. + """ + + def __init__(self, layer_id: int, args: ModelArgs): + super().__init__() + self.layer_id = layer_id + self.dim = args.dim + self.n_heads = args.n_heads + self.n_local_heads = args.n_heads // world_size + self.q_lora_rank = args.q_lora_rank + self.o_lora_rank = args.o_lora_rank + self.head_dim = args.head_dim + self.rope_head_dim = args.rope_head_dim + self.nope_head_dim = args.head_dim - args.rope_head_dim + self.n_groups = args.o_groups + self.n_local_groups = self.n_groups // world_size + self.window_size = args.window_size + self.compress_ratio = args.compress_ratios[layer_id] + self.eps = args.norm_eps + + self.attn_sink = nn.Parameter(torch.empty(self.n_local_heads, dtype=torch.float32)) + self.wq_a = Linear(self.dim, self.q_lora_rank) + self.q_norm = RMSNorm(self.q_lora_rank, self.eps) + self.wq_b = ColumnParallelLinear(self.q_lora_rank, self.n_heads * self.head_dim) + self.wkv = Linear(self.dim, self.head_dim) + self.kv_norm = RMSNorm(self.head_dim, self.eps) + self.wo_a = ColumnParallelLinear( + self.n_heads * self.head_dim // self.n_groups, + self.n_groups * args.o_lora_rank, + dtype=torch.bfloat16, + ) + self.wo_b = RowParallelLinear(self.n_groups * args.o_lora_rank, self.dim) + self.softmax_scale = self.head_dim**-0.5 + + is_backbone = layer_id < args.n_layers + self.is_kv_source = is_backbone and layer_id in args.kv_source_layers + self.is_index_source = is_backbone and layer_id in args.index_source_layers + self.compressor: Compressor | None = None + self.indexer: Indexer | None = None + if self.is_kv_source: + self.compressor = Compressor(args, layer_id) + if self.is_index_source: + self.indexer = Indexer(args, layer_id) + + self.window_kv_cache: torch.Tensor + self.register_buffer( + "window_kv_cache", + torch.zeros(args.max_batch_size, args.window_size, self.head_dim), + persistent=False, + ) + if self.is_kv_source: + self.compress_kv_cache: torch.Tensor + self.register_buffer( + "compress_kv_cache", + torch.zeros( + args.max_batch_size, + args.max_seq_len // self.compress_ratio, + self.head_dim, + ), + persistent=False, + ) + if self.compress_ratio: + original_seq_len, rope_theta = ( + args.original_seq_len, + args.compress_rope_theta, + ) + else: + # disable YaRN and use base rope_theta in pure sliding-window attention + original_seq_len, rope_theta = 0, args.rope_theta + freqs_cis = precompute_freqs_cis( + self.rope_head_dim, + args.max_seq_len, + original_seq_len, + rope_theta, + args.rope_factor, + args.beta_fast, + args.beta_slow, + ) + self.freqs_cis: torch.Tensor + self.register_buffer("freqs_cis", freqs_cis, persistent=False) + + def _window_kv(self, x, freqs_cis, start_pos): + """This layer's sliding-window K and the window positions every query may attend to. The K + stays fp8, quantized over the whole post-RoPE vector, RoPE tail included.""" + bsz, seqlen, _ = x.size() + win = self.window_size + kv = self.kv_norm(self.wkv(x)) + apply_rotary_emb(kv[..., -self.rope_head_dim :], freqs_cis) + act_quant(kv, fp8_block_size, scale_fmt, scale_dtype, True) + if start_pos == 0: # prefill: attend over this chunk, seeding the ring buffer for decode + if seqlen <= win: + self.window_kv_cache[:bsz, :seqlen] = kv + else: + cutoff = seqlen % win + self.window_kv_cache[:bsz, cutoff:win], self.window_kv_cache[:bsz, :cutoff] = kv[:, -win:].split( + [win - cutoff, cutoff], dim=1 + ) + window_kv = kv + else: # decode: one token into the ring buffer, attend over the whole window + self.window_kv_cache[:bsz, start_pos % win] = kv.squeeze(1) + window_kv = self.window_kv_cache[:bsz] + return window_kv, get_window_topk_idxs(win, bsz, seqlen, start_pos) + + def _compress_topk_idxs(self, x, qr, latent, start_pos, offset, compress_len): + """Which compressed positions each query attends to. Index sources run their own indexer; + the layers in between reuse the result their source published.""" + if not self.is_index_source: + return shared_attn.topk_idxs + + bsz, seqlen, _ = x.size() + if compress_len == 0: + idxs = torch.empty(bsz, seqlen, 0, dtype=torch.int32, device=x.device) + else: + assert self.indexer is not None + if self.indexer.freqs_cis is None: + self.indexer.freqs_cis = self.freqs_cis + idxs = self.indexer(x, qr, latent, start_pos, offset) + shared_attn.topk_idxs = idxs + return idxs + + def _compress_kv(self, x, qr, start_pos, offset): + """The shared compressed KV and the compressed positions every query may attend to. This + layer compresses its own KV only when it is a source; otherwise it just reads the cache.""" + bsz, seqlen, _ = x.size() + ratio = self.compress_ratio + compress_len = (start_pos + seqlen) // ratio + latent = None + if self.is_kv_source: + latent = self.compressor(x, start_pos) + shared_attn.compress_kv = self.compress_kv_cache + # the indexer needs the latent before RoPE, so it runs before the cache is written + idxs = self._compress_topk_idxs(x, qr, latent, start_pos, offset, compress_len) + if latent is not None: + # a latent stands for the first token of its group, so group j takes position j * ratio + freqs = ( + self.freqs_cis[: seqlen - seqlen % ratio : ratio] + if start_pos == 0 + else self.freqs_cis[start_pos + 1 - ratio].unsqueeze(0) + ) + apply_rotary_emb(latent[..., -self.rope_head_dim :], freqs) + # Compressed KV uses groups of 16 with E4M3 scales; the indexer uses 32 with E8M0. + fp4_act_quant(latent, 16, True, scale_dtype=torch.float8_e4m3fn) + self.compress_kv_cache[:bsz, start_pos // ratio : start_pos // ratio + latent.size(1)] = latent + # read after the write, so this does not depend on the slice aliasing the cache + return shared_attn.compress_kv[:bsz, :compress_len], idxs + + def forward(self, x: torch.Tensor, start_pos: int): + bsz, seqlen, _ = x.size() + freqs_cis = self.freqs_cis[start_pos : start_pos + seqlen] + rd = self.rope_head_dim + + qr = self.q_norm(self.wq_a(x)) + q = self.wq_b(qr).unflatten(-1, (self.n_local_heads, self.head_dim)) + apply_rotary_emb(q[..., -rd:], freqs_cis) + + kv, topk_idxs = self._window_kv(x, freqs_cis, start_pos) + if self.compress_ratio: + compress_kv, compress_idxs = self._compress_kv(x, qr, start_pos, kv.size(1)) + kv = torch.cat([kv, compress_kv], dim=1) + topk_idxs = torch.cat([topk_idxs, compress_idxs], dim=-1) + + o = sparse_attn(q, kv, self.attn_sink, topk_idxs, self.softmax_scale) + apply_rotary_emb(o[..., -rd:], freqs_cis, True) + + # wo_a is block-diagonal over groups (each projects only its own heads), hence einsum not + # Linear. convert.py dequantizes it to bf16; an fp8 grouped GEMM would halve the memory. + o = o.view(bsz, seqlen, self.n_local_groups, -1) + wo_a = self.wo_a.weight.view(self.n_local_groups, self.o_lora_rank, -1) + o = torch.einsum("bsgd,grd->bsgr", o, wo_a) + x = self.wo_b(o.flatten(2)) + return x + + # FreeToken: Decoder SWA Bounded Replay oracle. Prefill of the prompt TAIL [s, s + n) as its own + # window segment (a query sees SWA keys in [max(s, i - win + 1), i]) while the global branch + # covers the whole prompt: a kv source compresses ``x_all`` (every token, forward()'s start_pos + # == 0 path), the indexer scores the tail queries at their absolute positions. + def forward_tail(self, x: torch.Tensor, s: int, x_all: torch.Tensor | None = None): + bsz, n, _ = x.size() + win, rd, ratio = self.window_size, self.rope_head_dim, self.compress_ratio + freqs_cis = self.freqs_cis[s : s + n] + + qr = self.q_norm(self.wq_a(x)) + q = self.wq_b(qr).unflatten(-1, (self.n_local_heads, self.head_dim)) + apply_rotary_emb(q[..., -rd:], freqs_cis) + + kv = self.kv_norm(self.wkv(x)) + apply_rotary_emb(kv[..., -rd:], freqs_cis) + act_quant(kv, fp8_block_size, scale_fmt, scale_dtype, True) + # seed the decode ring with the tail's keys at their absolute slots + for i in range(max(0, n - win), n): + self.window_kv_cache[:bsz, (s + i) % win] = kv[:, i] + topk_idxs = get_window_topk_idxs(win, bsz, n, 0).to(x.device) + + if ratio: + end = s + n + compress_len = end // ratio + latent = None + if self.is_kv_source: + assert x_all is not None, "the kv source compresses every prompt token" + latent = self.compressor(x_all, 0) + shared_attn.compress_kv = self.compress_kv_cache + if self.indexer is not None: + if self.indexer.freqs_cis is None: + self.indexer.freqs_cis = self.freqs_cis + # publish the index keys for every group (forward()'s start_pos == 0 path) + L = x_all.size(1) + k = self.indexer.k_norm(self.indexer.wk(latent)) + apply_rotary_emb(k[..., -rd:], self.freqs_cis[: L - L % ratio : ratio]) + fp4_act_quant(k, fp4_block_size, True) + self.indexer.k_cache[:bsz, : k.size(1)] = k + shared_attn.index_k = self.indexer.k_cache + if not self.is_index_source: + idxs = shared_attn.topk_idxs + elif compress_len == 0: + idxs = torch.empty(bsz, n, 0, dtype=torch.int32, device=x.device) + else: + if self.indexer.freqs_cis is None: + self.indexer.freqs_cis = self.freqs_cis + idxs = self.indexer.forward_tail(x, qr, s, n, kv.size(1)) + shared_attn.topk_idxs = idxs + if latent is not None: + L = x_all.size(1) + apply_rotary_emb(latent[..., -rd:], self.freqs_cis[: L - L % ratio : ratio]) + fp4_act_quant(latent, 16, True, scale_dtype=torch.float8_e4m3fn) + self.compress_kv_cache[:bsz, : latent.size(1)] = latent + kv = torch.cat([kv, shared_attn.compress_kv[:bsz, :compress_len]], dim=1) + topk_idxs = torch.cat([topk_idxs, idxs], dim=-1) + + o = sparse_attn(q, kv, self.attn_sink, topk_idxs, self.softmax_scale) + apply_rotary_emb(o[..., -rd:], freqs_cis, True) + o = o.view(bsz, n, self.n_local_groups, -1) + wo_a = self.wo_a.weight.view(self.n_local_groups, self.o_lora_rank, -1) + o = torch.einsum("bsgd,grd->bsgr", o, wo_a) + return self.wo_b(o.flatten(2)) + + +class Gate(nn.Module): + """MoE gating. The correction bias steers expert selection only; the routing weights come from the + unbiased scores. Image-span tokens use a separate bias (training `noaux_tc_for_vl`).""" + + def __init__(self, layer_id: int, args: ModelArgs): + super().__init__() + n_routed_experts, n_activated_experts = args.get_moe_config(layer_id) + self.dim = args.dim + self.topk = n_activated_experts + self.score_func = args.score_func + self.gate_temp = args.gate_temp + self.norm_topk_prob = args.norm_topk_prob + self.route_scale = args.route_scale + self.weight = nn.Parameter(torch.empty(n_routed_experts, args.dim)) + self.bias = nn.Parameter(torch.empty(n_routed_experts, dtype=torch.float32)) + self.bias_vl = nn.Parameter(torch.empty(n_routed_experts, dtype=torch.float32)) if args.vision_enabled else None + + def forward(self, x: torch.Tensor, image_mask: torch.Tensor | None = None) -> tuple[torch.Tensor, torch.Tensor]: + """x: [n, dim]; image_mask: [n] bool, True for tokens inside an image span.""" + scores = linear(x.float(), self.weight.float()) / self.gate_temp + if self.score_func == "softmax": + scores = scores.softmax(dim=-1) + elif self.score_func == "sigmoid": + scores = scores.sigmoid() + else: + scores = F.softplus(scores).sqrt() + bias = self.bias + if image_mask is not None and self.bias_vl is not None: + bias = torch.where(image_mask.unsqueeze(-1), self.bias_vl, bias) + # the bias picks experts but does not scale them: weights come from the raw scores + indices = (scores + bias).topk(self.topk, dim=-1)[1] + weights = scores.gather(1, indices) + if self.norm_topk_prob and self.topk > 1: + weights /= weights.sum(dim=-1, keepdim=True) + 1e-20 # not norm_eps, matches training + weights *= self.route_scale + return weights, indices + + +class Expert(nn.Module): + """One SwiGLU FFN. The clamps come straight from training, where they keep fp8/fp4 activations in + range: the up branch is clamped on both sides, the gate branch only from above.""" + + def __init__(self, dim: int, inter_dim: int, dtype=None, swiglu_limit=0.0): + super().__init__() + self.w1 = Linear(dim, inter_dim, dtype=dtype) + self.w2 = Linear(inter_dim, dim, dtype=dtype) + self.w3 = Linear(dim, inter_dim, dtype=dtype) + self.swiglu_limit = swiglu_limit + + def forward(self, x: torch.Tensor, weights: torch.Tensor | None = None) -> torch.Tensor: + dtype = x.dtype + gate = self.w1(x).float() + up = self.w3(x).float() + if self.swiglu_limit > 0: + up = torch.clamp(up, min=-self.swiglu_limit, max=self.swiglu_limit) + gate = torch.clamp(gate, max=self.swiglu_limit) + x = F.silu(gate) * up + if weights is not None: + x = weights * x + return self.w2(x.to(dtype)) + + +class MoE(nn.Module): + """Top-k routed experts plus one shared expert every token goes through. Experts are split + across ranks, so `self.experts` is None for those another rank owns.""" + + def __init__(self, layer_id: int, args: ModelArgs): + super().__init__() + n_routed_experts, n_activated_experts = args.get_moe_config(layer_id) + self.layer_id = layer_id + self.dim = args.dim + assert n_routed_experts % world_size == 0, ( + f"Number of experts must be divisible by world size (world_size={world_size})" + ) + self.n_routed_experts = n_routed_experts + self.n_local_experts = n_routed_experts // world_size + self.n_activated_experts = n_activated_experts + self.experts_start_idx = rank * self.n_local_experts + self.experts_end_idx = self.experts_start_idx + self.n_local_experts + self.gate = Gate(layer_id, args) + expert_dtype = torch.float4_e2m1fn_x2 if args.expert_dtype == "fp4" else None + self.experts = nn.ModuleList( + [ + Expert( + args.dim, + args.moe_inter_dim, + dtype=expert_dtype, + swiglu_limit=args.swiglu_limit, + ) + if self.experts_start_idx <= i < self.experts_end_idx + else None + for i in range(self.n_routed_experts) + ] + ) + assert args.n_shared_experts == 1 + self.shared_experts = Expert(args.dim, args.moe_inter_dim, swiglu_limit=args.swiglu_limit) + + def forward(self, x: torch.Tensor, image_mask: torch.Tensor | None = None) -> torch.Tensor: + shape = x.size() + x = x.view(-1, self.dim) + weights, indices = self.gate(x, None if image_mask is None else image_mask.flatten()) + y = torch.zeros_like(x, dtype=torch.float32) + counts = torch.bincount(indices.flatten(), minlength=self.n_routed_experts).tolist() + for i in range(self.experts_start_idx, self.experts_end_idx): + if counts[i] == 0: + continue + expert = self.experts[i] + idx, top = torch.where(indices == i) + y[idx] += expert(x[idx], weights[idx, top, None]) + if world_size > 1: + dist.all_reduce(y) + y += self.shared_experts(x) + return y.type_as(x).view(shape) + + +class Block(nn.Module): + """A block whose residual stream is `hc_mult` parallel copies (Hyper-Connections). + + Attention and FFN each sit between `hc_pre` (collapse the copies into one sublayer input) and + `hc_post` (expand back out, mixing the residual in through `comb`). `hc_mixes` derives all three + coefficient sets from the stream itself, `comb` made doubly stochastic by Sinkhorn. + + The coefficients a sublayer computes are used by the *next* one -- see `forward`. + """ + + attention_cls = Attention + + def __init__( + self, + layer_id: int, + args: ModelArgs, + engram_layout: EngramLayout | None = None, + ): + super().__init__() + self.layer_id = layer_id + self.norm_eps = args.norm_eps + self.attn = self.attention_cls(layer_id, args) + self.ffn = MoE(layer_id, args) + self.engram = None + if engram_layout is not None and layer_id in engram_layout.layer_ids: + self.engram = Engram(args, layer_id, engram_layout) + self.attn_norm = RMSNorm(args.dim, self.norm_eps) + self.ffn_norm = RMSNorm(args.dim, self.norm_eps) + self.hc_mult = hc_mult = args.hc_mult + self.hc_sinkhorn_iters = args.hc_sinkhorn_iters + self.hc_eps = args.hc_eps + mix_hc = (2 + hc_mult) * hc_mult + hc_dim = hc_mult * args.dim + with set_dtype(torch.float32): + self.hc_attn_fn = nn.Parameter(torch.empty(mix_hc, hc_dim)) + self.hc_ffn_fn = nn.Parameter(torch.empty(mix_hc, hc_dim)) + self.hc_attn_base = nn.Parameter(torch.empty(mix_hc)) + self.hc_ffn_base = nn.Parameter(torch.empty(mix_hc)) + self.hc_attn_scale = nn.Parameter(torch.empty(3)) + self.hc_ffn_scale = nn.Parameter(torch.empty(3)) + + def hc_mixes(self, x: torch.Tensor, hc_fn: torch.Tensor, hc_scale: torch.Tensor, hc_base: torch.Tensor): + """x: [b,s,hc,d], hc_fn: [mix_hc, hc*d], hc_scale: [3], hc_base: [mix_hc]. Returns the + pre / post / comb coefficients, split out of one projection of the flattened stream.""" + # normalized over the whole flattened hc*d stream, one statistic per token + x = x.flatten(2).float() + rsqrt = torch.rsqrt(x.square().mean(-1, keepdim=True) + self.norm_eps) + mixes = F.linear(x, hc_fn) * rsqrt + return hc_split_sinkhorn(mixes, hc_scale, hc_base, self.hc_mult, self.hc_sinkhorn_iters, self.hc_eps) + + def hc_pre(self, x: torch.Tensor, pre_mix: torch.Tensor): + """Collapse the hc copies into one, weighted by pre_mix. [b,s,hc,d] x [b,s,hc] -> [b,s,d]""" + y = torch.sum(pre_mix.unsqueeze(-1) * x.float(), dim=2) + return y.to(x.dtype) + + def hc_post(self, x: torch.Tensor, residual: torch.Tensor, post: torch.Tensor, comb: torch.Tensor): + """Expand the sublayer output back to hc copies and mix the residual in through `comb`. + x: [b,s,d], residual: [b,s,hc,d], post: [b,s,hc], comb: [b,s,hc,hc] -> [b,s,hc,d]""" + y = post.unsqueeze(-1) * x.unsqueeze(-2) + torch.sum(comb.unsqueeze(-1) * residual.unsqueeze(-2), dim=2) + return y.type_as(x) + + def forward( + self, + x: torch.Tensor, + start_pos: int, + pre_mix: torch.Tensor, + image_mask: torch.Tensor | None, + *attn_args, + ) -> tuple[torch.Tensor, torch.Tensor]: + """`pre_mix` collapses the hc_mult copies down to one input for this block's attention. Each + sub-block's own `hc_mixes` produces the mix for the *next* one, so attention uses what the + previous layer's FFN produced and the FFN uses what this attention produced. + + image_mask: [b, s] bool, True inside image spans (selects the VL routing bias).""" + residual = x + attn_pre, attn_post, attn_comb = self.hc_mixes(x, self.hc_attn_fn, self.hc_attn_scale, self.hc_attn_base) + x = self.hc_pre(x, pre_mix) + x = self.attn_norm(x) + x = self.attn(x, start_pos, *attn_args) + x = self.hc_post(x, residual, attn_post, attn_comb) + + residual = x + ffn_pre, ffn_post, ffn_comb = self.hc_mixes(x, self.hc_ffn_fn, self.hc_ffn_scale, self.hc_ffn_base) + x = self.hc_pre(x, attn_pre) + x = self.ffn_norm(x) + x = self.ffn(x, image_mask) + x = self.hc_post(x, residual, ffn_post, ffn_comb) + return x, ffn_pre + + # FreeToken: Decoder SWA Bounded Replay oracle. ``h`` / ``pre_mix`` are the tail [s, N) streams, + # or -- at the decoder's kv source -- ``h_all`` / ``pre_mix_all`` cover every token so the + # source compresses them all before the block narrows to the tail. + def forward_tail(self, h, s: int, pre_mix, h_all=None, pre_mix_all=None): + if h_all is not None: + attn_pre, attn_post, attn_comb = self.hc_mixes(h_all, self.hc_attn_fn, self.hc_attn_scale, self.hc_attn_base) + x_all = self.attn_norm(self.hc_pre(h_all, pre_mix_all)) + residual = h_all[:, s:] + attn_pre, attn_post, attn_comb = attn_pre[:, s:], attn_post[:, s:], attn_comb[:, s:] + x = self.attn.forward_tail(x_all[:, s:], s, x_all) + else: + residual = h + attn_pre, attn_post, attn_comb = self.hc_mixes(h, self.hc_attn_fn, self.hc_attn_scale, self.hc_attn_base) + x = self.attn_norm(self.hc_pre(h, pre_mix)) + x = self.attn.forward_tail(x, s) + x = self.hc_post(x, residual, attn_post, attn_comb) + + residual = x + ffn_pre, ffn_post, ffn_comb = self.hc_mixes(x, self.hc_ffn_fn, self.hc_ffn_scale, self.hc_ffn_base) + x = self.hc_pre(x, attn_pre) + x = self.ffn_norm(x) + x = self.ffn(x, None) + x = self.hc_post(x, residual, ffn_post, ffn_comb) + return x, ffn_pre + + +class ParallelHead(nn.Module): + def __init__(self, vocab_size: int, dim: int, norm_eps: float = 1e-6, hc_eps: float = 1e-6): + super().__init__() + self.vocab_size = vocab_size + self.dim = dim + self.norm_eps = norm_eps + self.hc_eps = hc_eps + self.part_vocab_size = vocab_size // world_size + # bf16 in the checkpoint, kept as fp32 here so the logits come out in fp32 directly + self.weight = nn.Parameter(torch.empty(self.part_vocab_size, self.dim, dtype=torch.float32)) + + def forward(self, x: torch.Tensor, full_logits=False): + """x: [b, s, d]. Generation only needs the last position, so that is the default.""" + if not full_logits: + x = x[:, -1] + logits = F.linear(x.float(), self.weight) + if world_size > 1: + all_logits = [torch.empty_like(logits) for _ in range(world_size)] + dist.all_gather(all_logits, logits) + logits = torch.cat(all_logits, dim=-1) + return logits + + +@lru_cache(1) +def get_dspark_topk_idxs(window_size: int, bsz: int, block_size: int, start_pos: int): + assert start_pos > 0 + matrix = torch.cat( + [ + torch.arange(min(window_size, start_pos + 1)), + window_size + torch.arange(block_size), + ] + ) + return matrix.int().view(1, 1, -1).expand(bsz, block_size, -1).contiguous() + + +class DSparkAttention(Attention): + def forward(self, x: torch.Tensor, start_pos: int, main_x: torch.Tensor): + assert self.compress_ratio == 0 + bsz, seqlen, _ = main_x.size() + win = self.window_size + rd = self.rope_head_dim + + main_freqs_cis = self.freqs_cis[start_pos : start_pos + seqlen] + main_kv = self.kv_norm(self.wkv(main_x)) + apply_rotary_emb(main_kv[..., -rd:], main_freqs_cis) + act_quant(main_kv, fp8_block_size, scale_fmt, scale_dtype, True) + + if start_pos == 0: + if seqlen <= win: + self.window_kv_cache[:bsz, :seqlen] = main_kv + else: + cutoff = seqlen % win + self.window_kv_cache[:bsz, cutoff:win], self.window_kv_cache[:bsz, :cutoff] = main_kv[:, -win:].split( + [win - cutoff, cutoff], dim=1 + ) + return x + + bsz, block_size, _ = x.size() + freqs_cis = self.freqs_cis[start_pos + seqlen : start_pos + seqlen + block_size] + + qr = self.q_norm(self.wq_a(x)) + q = self.wq_b(qr).unflatten(-1, (self.n_local_heads, self.head_dim)) + apply_rotary_emb(q[..., -rd:], freqs_cis) + kv = self.kv_norm(self.wkv(x)) + apply_rotary_emb(kv[..., -rd:], freqs_cis) + act_quant(kv, fp8_block_size, scale_fmt, scale_dtype, True) + + topk_idxs = get_dspark_topk_idxs(win, bsz, block_size, start_pos) + self.window_kv_cache[:bsz, start_pos % win] = main_kv.squeeze(1) + kv = torch.cat([self.window_kv_cache[:bsz], kv], dim=1) + o = sparse_attn(q, kv, self.attn_sink, topk_idxs, self.softmax_scale) + apply_rotary_emb(o[..., -rd:], freqs_cis, True) + + o = o.view(bsz, block_size, self.n_local_groups, -1) + wo_a = self.wo_a.weight.view(self.n_local_groups, self.o_lora_rank, -1) + o = torch.einsum("bsgd,grd->bsgr", o, wo_a) + x = self.wo_b(o.flatten(2)) + return x + + +class DSparkMarkovHead(nn.Module): + def __init__(self, vocab_size: int, dspark_markov_rank: int): + super().__init__() + self.embed = ParallelEmbedding(vocab_size, dspark_markov_rank) + self.head = ParallelHead(vocab_size, dspark_markov_rank) + + def forward(self, token_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + embed = self.embed(token_ids) + logits = self.head(embed, full_logits=True) + return logits, embed + + +class DSparkConfidenceHead(nn.Module): + def __init__(self, input_dim: int): + super().__init__() + # proj in the checkpoint is stored in bf16, while the parameter here is stored in fp32 for fp32 confidence score. + self.proj = Linear(input_dim, 1, dtype=torch.float32) + + def forward(self, hidden: torch.Tensor, markov_embed: torch.Tensor): + hidden = torch.cat([hidden, markov_embed], dim=-1) + return self.proj(hidden.float()).squeeze(-1) + + +class DSparkBlock(Block): + """DSpark stage stored under the mtp.* checkpoint namespace.""" + + attention_cls = DSparkAttention + + def __init__(self, layer_id: int, args: ModelArgs): + super().__init__(layer_id, args) + stage_id = layer_id - args.n_layers + self.block_size = args.dspark_block_size + self.noise_token_id = args.dspark_noise_token_id + self.temperature = args.temperature + if stage_id == 0: + assert len(args.dspark_target_layer_ids) > 0, "DSpark needs target layers" + self.main_proj = Linear(args.dim * len(args.dspark_target_layer_ids), args.dim) + self.main_norm = RMSNorm(args.dim, args.norm_eps) + if stage_id == args.n_mtp_layers - 1: + self.norm = RMSNorm(args.dim, args.norm_eps) + self.markov_head = DSparkMarkovHead(args.vocab_size, args.dspark_markov_rank) + self.confidence_head = DSparkConfidenceHead(args.dim + args.dspark_markov_rank) + self.embed: ParallelEmbedding | None = None + self.head: ParallelHead | None = None + + def forward(self, x: torch.Tensor, start_pos: int, pre_mix: torch.Tensor, main_x: torch.Tensor): + if start_pos == 0: + self.attn(x, start_pos, main_x) # prefill only seeds the window KV cache + return x, pre_mix + return super().forward(x, start_pos, pre_mix, None, main_x) # drafts are text: no VL bias + + def forward_embed(self, main_hidden: torch.Tensor, input_ids: torch.Tensor): + assert self.embed is not None + main_x = self.main_norm(self.main_proj(main_hidden)) + draft_input_ids = input_ids.new_full([input_ids.size(0), self.block_size], self.noise_token_id) + draft_input_ids[:, 0] = input_ids + x = self.embed(draft_input_ids) + x = x.unsqueeze(2).repeat(1, 1, self.hc_mult, 1) + return x, main_x + + def forward_head( + self, + x: torch.Tensor, + pre_mix: torch.Tensor, + input_ids: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + assert self.head is not None + x = self.hc_pre(x, pre_mix) + logits = self.head(self.norm(x), full_logits=True) + output_ids = input_ids.new_empty(input_ids.size(0), self.block_size + 1) + output_ids[:, 0] = input_ids + markov_embeds = [] + for i in range(self.block_size): + logits_bias, markov_embed = self.markov_head(output_ids[:, i]) + logits[:, i].add_(logits_bias) + markov_embeds.append(markov_embed) + output_ids[:, i + 1] = sample(logits[:, i], self.temperature) + markov_embed = torch.stack(markov_embeds, dim=1) + confidence = self.confidence_head(x, markov_embed) + return output_ids, logits, confidence + + +def make_identity_pre_mix(x: torch.Tensor, hc_mult: int) -> torch.Tensor: + """initial one-hot mix""" + pre_mix = x.new_zeros(x.size(0), x.size(1), hc_mult, dtype=torch.float32) + pre_mix[:, :, 0] = 1.0 + return pre_mix + + +class SharedAttentionRuntime: + """What attention layers hand down the stack instead of recomputing. Layers run in order and + every source writes before its consumers read, so one slot each is enough and nothing needs + resetting between forwards. Sources: compress_kv and index_k from kv_source_layers, topk_idxs + from index_source_layers, candidates from candidate_source_layer.""" + + def __init__(self): + self.compress_kv: torch.Tensor | None = None + self.index_k: torch.Tensor | None = None + self.topk_idxs: torch.Tensor | None = None + self.candidates: torch.Tensor | None = None + + +# Only ever one model per process, same as world_size / rank / default_dtype above. +shared_attn = SharedAttentionRuntime() + + +class Transformer(nn.Module): + """DeepSeek-V4.1: embed -> expand to hc_mult copies -> blocks -> collapse -> logits. Building + this sets the globals at the top of the file. The tokenizer only feeds the engram token map.""" + + def __init__(self, args: ModelArgs, tokenizer=None): + global world_size, rank, default_dtype + world_size = dist.get_world_size() if dist.is_initialized() else 1 + rank = dist.get_rank() if dist.is_initialized() else 0 + default_dtype = torch.float8_e4m3fn if args.dtype == "fp8" else torch.bfloat16 + super().__init__() + self.max_seq_len = args.max_seq_len + self.temperature = args.temperature + self.norm_eps = args.norm_eps + self.hc_eps = args.hc_eps + self.engram_layout = EngramLayout.from_args(args) + self.engram_hash = ( + NgramHashState(args, self.engram_layout, tokenizer) if self.engram_layout is not None else None + ) + self.embed = ParallelEmbedding(args.vocab_size, args.dim) + self.layers = torch.nn.ModuleList() + for layer_id in range(args.n_layers): + self.layers.append(Block(layer_id, args, self.engram_layout)) + self.norm = RMSNorm(args.dim, self.norm_eps) + self.head = ParallelHead(args.vocab_size, args.dim, self.norm_eps, self.hc_eps) + self.mtp = torch.nn.ModuleList() + self.target_layer_ids = args.dspark_target_layer_ids + if args.dspark_block_size: + for layer_id in range(args.n_mtp_layers): + self.mtp.append(DSparkBlock(args.n_layers + layer_id, args)) + self.mtp[-1].embed = self.embed + self.mtp[-1].head = self.head + self.hc_mult = args.hc_mult + self.vision = None + if args.vision_enabled: + self.vision = ViT(args) + self.aligner = Aligner(args) + # learned embeddings for the image span delimiters + self.image_start = nn.Parameter(torch.empty(args.dim)) + self.image_end = nn.Parameter(torch.empty(args.dim)) + self.image_newline = nn.Parameter(torch.empty(args.dim)) + + @torch.inference_mode() + def encode_image(self, patches: torch.Tensor, n_vit_h: int, n_vit_w: int) -> torch.Tensor: + return self.aligner(self.vision(patches, n_vit_h, n_vit_w), n_vit_h, n_vit_w) + + def merge_image_embeddings(self, images, h: torch.Tensor): + """Overwrite each image's token span in h with its ViT/aligner features. The IMAGE slots take + the aligner rows in row-major order; the span delimiters take learned embeddings.""" + for i, sample in enumerate(images): + for img in sample or (): + types = img.types.to(h.device) + span = h[i, img.start : img.start + types.numel()] + span[types == IMAGE_START] = self.image_start.to(h.dtype) + span[types == IMAGE_END] = self.image_end.to(h.dtype) + span[types == IMAGE_NEW_LINE] = self.image_newline.to(h.dtype) + embeds = self.encode_image(img.patches.to(h.device), img.n_vit_h, img.n_vit_w) + span[types == IMAGE] = embeds.to(h.dtype) + + @torch.inference_mode() + def forward( + self, input_ids: torch.Tensor, start_pos: int = 0, images=None, token_types: torch.Tensor | None = None + ): + """input_ids: [b, s], every entry a real token id -- generate.py only ever passes positions it + has already filled, so the padding it uses internally never reaches here. token_types / + images carry the VL inputs built by image_processor.prepare_vl_inputs; image spans must lie + inside the first (start_pos 0) chunk.""" + image_mask = None if token_types is None else token_types >= 0 # TEXT is -1 + # image tokens take no part in an n-gram and get no engram contribution; text-only needs no mask + engram_mask = None if image_mask is None else ~image_mask + engram_hashes = self.engram_hash(input_ids, start_pos, engram_mask) if self.engram_hash is not None else None + h = self.embed(input_ids) + if images is not None: + assert start_pos == 0, "image spans must be prefilled in a single chunk" + self.merge_image_embeddings(images, h) + # Expand to hc_mult copies for Hyper-Connections + h = h.unsqueeze(2).repeat(1, 1, self.hc_mult, 1) + main_hiddens = [] + pre_mix = make_identity_pre_mix(h, self.hc_mult) + for i, layer in enumerate(self.layers): + if layer.engram is not None: + h = layer.engram(h, engram_hashes[:, :, layer.engram.layer_hash_index, :], engram_mask) + # the MTP head reads the attention input of its target layers, not their output + if i in self.target_layer_ids: + main_hiddens.append(h.mean(dim=2)) + h, pre_mix = layer(h, start_pos, pre_mix, image_mask) + h = layer.hc_pre(h, pre_mix) + logits = self.head(self.norm(h)) + output_ids = sample(logits, self.temperature) + main_hidden = torch.cat(main_hiddens, dim=-1) if main_hiddens else None + return output_ids, logits, main_hidden + + # FreeToken: Decoder SWA Bounded Replay oracle (tech report sec. 3.2.2): the encoder runs on + # every prompt token, the decoder on the last ``n_win`` with its sliding window truncated + # there; the decoder's kv source (the last kv_source layer) compresses every token. Returns + # the last position's logits and leaves the caches ready for decode from position seqlen. + @torch.inference_mode() + def forward_bounded(self, input_ids: torch.Tensor, n_win: int): + engram_hashes = self.engram_hash(input_ids, 0, None) if self.engram_hash is not None else None + h = self.embed(input_ids).unsqueeze(2).repeat(1, 1, self.hc_mult, 1) + pre_mix = make_identity_pre_mix(h, self.hc_mult) + N = input_ids.size(1) + s = max(0, N - n_win) + dec_start = max(i for i, layer in enumerate(self.layers) if layer.attn.is_kv_source) + for i, layer in enumerate(self.layers): + if layer.engram is not None: + assert i < dec_start, "Engram sits in the encoder" + h = layer.engram(h, engram_hashes[:, :, layer.engram.layer_hash_index, :], None) + if i < dec_start: + h, pre_mix = layer(h, 0, pre_mix, None) + elif i == dec_start: + h, pre_mix = layer.forward_tail(None, s, None, h_all=h, pre_mix_all=pre_mix) + else: + h, pre_mix = layer.forward_tail(h, s, pre_mix) + h = layer.hc_pre(h, pre_mix) + return self.head(self.norm(h)) + + @torch.inference_mode() + def forward_spec(self, input_ids: torch.Tensor, main_hidden: torch.Tensor, start_pos: int = 0): + h, main_x = self.mtp[0].forward_embed(main_hidden, input_ids) + pre_mix = make_identity_pre_mix(h, self.hc_mult) + for layer in self.mtp: + h, pre_mix = layer(h, start_pos, pre_mix, main_x) + if start_pos == 0: + return None + return self.mtp[-1].forward_head(h, pre_mix, input_ids) + + +def sample(logits, temperature: float = 1.0): + """Gumbel-max trick: equivalent to multinomial sampling but faster on GPU, + since it avoids the GPU-to-CPU sync in torch.multinomial.""" + if temperature == 0: + return logits.argmax(dim=-1) + logits = logits / max(temperature, 1e-5) + probs = torch.softmax(logits, dim=-1, dtype=torch.float32) + return probs.div_(torch.empty_like(probs).exponential_(1)).argmax(dim=-1) + + +if __name__ == "__main__": + torch.set_default_dtype(torch.bfloat16) + torch.set_default_device("cuda") + torch.manual_seed(0) + args = ModelArgs(dspark_block_size=6, dspark_target_layer_ids=(3, 4)) + x = torch.randint(0, args.vocab_size, (2, 150)) + model = Transformer(args) + + output_ids, logits, main_hidden = model(x[:, :128]) + model.forward_spec(output_ids, main_hidden) + for i in range(128, 150): + output_ids, logits, main_hidden = model(x[:, i : i + 1], i) + result = model.forward_spec(output_ids, main_hidden, i) + assert result is not None + output_ids, logits, confidence = result diff --git a/tests/models/deepseek_v41/reference/vision.py b/tests/models/deepseek_v41/reference/vision.py new file mode 100644 index 0000000000..2fd86160f7 --- /dev/null +++ b/tests/models/deepseek_v41/reference/vision.py @@ -0,0 +1,121 @@ +# Vendored from deepseek-ai/DeepSeek-V4.1-Flash, revision dba1be0a, inference/vision.py. +# MIT License; see LICENSE in this directory. +from functools import lru_cache + +import torch +import torch.nn.functional as F +from torch import nn + + +@lru_cache(8) +def get_vision_cos_sin(n_h: int, n_w: int, dim: int, theta: float): + inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim)) + hpos = torch.arange(n_h).unsqueeze(1).expand(n_h, n_w) + wpos = torch.arange(n_w).unsqueeze(0).expand(n_h, n_w) + freqs = torch.stack([hpos, wpos], dim=-1).reshape(-1, 2, 1).float() * inv_freq + freqs = freqs.flatten(1) + return freqs.cos().unsqueeze(1), freqs.sin().unsqueeze(1) + + +def apply_rotary(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: + dtype = x.dtype + x1, x2 = x.float().chunk(2, dim=-1) + return torch.cat([x1 * cos - x2 * sin, x2 * cos + x1 * sin], dim=-1).to(dtype) + + +class RMSNorm(nn.Module): + def __init__(self, dim: int, eps: float = 1e-6): + super().__init__() + self.eps = eps + self.weight = nn.Parameter(torch.ones(dim, dtype=torch.float32)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + dtype = x.dtype + x = x.float() + x = x * torch.rsqrt(x.square().mean(-1, keepdim=True) + self.eps) + return (self.weight * x).to(dtype) + + +class PatchEmbed(nn.Module): + def __init__(self, args): + super().__init__() + self.proj = nn.Linear(3 * args.vision_patch_size**2, args.vision_dim) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.proj(x.flatten(1)) + + +class Attention(nn.Module): + def __init__(self, args): + super().__init__() + self.n_heads = args.vision_n_heads + self.head_dim = args.vision_dim // args.vision_n_heads + self.wqkv = nn.Linear(args.vision_dim, 3 * args.vision_dim) + self.wo = nn.Linear(args.vision_dim, args.vision_dim) + + def forward(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: + n = x.size(0) + q, k, v = (t.view(n, self.n_heads, self.head_dim) for t in self.wqkv(x).chunk(3, dim=-1)) + q = apply_rotary(q, cos, sin) + k = apply_rotary(k, cos, sin) + o = F.scaled_dot_product_attention(q.transpose(0, 1), k.transpose(0, 1), v.transpose(0, 1)) + return self.wo(o.transpose(0, 1).reshape(n, -1)) + + +class MLP(nn.Module): + def __init__(self, args): + super().__init__() + self.w1 = nn.Linear(args.vision_dim, 2 * args.vision_inter_dim, bias=False) + self.w2 = nn.Linear(args.vision_inter_dim, args.vision_dim, bias=False) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + gate, up = self.w1(x).chunk(2, dim=-1) + return self.w2(F.silu(gate) * up) + + +class Block(nn.Module): + def __init__(self, args): + super().__init__() + self.norm1 = RMSNorm(args.vision_dim) + self.attn = Attention(args) + self.norm2 = RMSNorm(args.vision_dim) + self.mlp = MLP(args) + + def forward(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: + x = x + self.attn(self.norm1(x), cos, sin) + return x + self.mlp(self.norm2(x)) + + +class ViT(nn.Module): + """DeepSeek ViT: full bidirectional attention over one image with 2D RoPE.""" + + def __init__(self, args): + super().__init__() + self.rope_dim = args.vision_dim // args.vision_n_heads // 2 + self.rope_theta = args.vision_rope_theta + self.patch_embed = PatchEmbed(args) + self.blocks = nn.ModuleList([Block(args) for _ in range(args.vision_n_layers)]) + self.norm = RMSNorm(args.vision_dim) + + def forward(self, patches: torch.Tensor, n_h: int, n_w: int) -> torch.Tensor: + x = self.patch_embed(patches) + cos, sin = get_vision_cos_sin(n_h, n_w, self.rope_dim, self.rope_theta) + for block in self.blocks: + x = block(x, cos, sin) + return self.norm(x) + + +class Aligner(nn.Module): + def __init__(self, args): + super().__init__() + self.downsample_ratio = args.vision_downsample_ratio + in_dim = args.vision_dim * self.downsample_ratio**2 + self.w1 = nn.Linear(in_dim, args.dim) + self.w2 = nn.Linear(args.dim, args.dim) + + def forward(self, x: torch.Tensor, n_h: int, n_w: int) -> torch.Tensor: + r = self.downsample_ratio + x = x.view(n_h, n_w, -1).permute(2, 0, 1) + x = F.pad(x, (0, -n_w % r, 0, -n_h % r)) + x = F.unfold(x.unsqueeze(0), r, stride=r).squeeze(0).transpose(0, 1) + return self.w2(F.gelu(self.w1(x))) diff --git a/tests/models/deepseek_v41/test_config.py b/tests/models/deepseek_v41/test_config.py new file mode 100644 index 0000000000..fab385fb93 --- /dev/null +++ b/tests/models/deepseek_v41/test_config.py @@ -0,0 +1,169 @@ +"""Config parsing, the DSV41 layer-role table, registry dispatch and loader / model key parity.""" + +from __future__ import annotations + +from dataclasses import replace + +import pytest +import torch + +from freetoken.layers import set_rope_device +from freetoken.layers.quantization import NoQuantConfig +from freetoken.distributed.info import set_tp_info, try_get_tp_info +from freetoken.models import create_model +from freetoken.models.deepseek_v41.args import DeepseekV41Args, Mode +from freetoken.models.deepseek_v41.config import parse_config +from freetoken.models.register import get_model_spec +from freetoken.utils.torch_utils import torch_dtype + +from .common import ENGRAM_LAYERS, INDEX_SOURCES, KV_SOURCES, N_LAYERS, RATIOS, tiny_hf_config, write_tiny_checkpoint + + +def test_roles_follow_the_dsv41_mode_table(): + args = parse_config(tiny_hf_config()).dsv41_args + modes = [r.mode for r in args.roles] + assert modes == [Mode.WINDOW, Mode.WINDOW, Mode.FULL, Mode.REUSE, Mode.FULL, Mode.REINDEX] + assert [r.kv_source for r in args.roles] == [None, None, 2, 2, 4, 4] + assert args.roles[4].is_candidate_source and args.roles[5].uses_candidates and not args.roles[2].uses_candidates + assert args.decoder_start_layer == 4 and args.backbone_kv_sources == (2, 4) + assert args.freqs_params(0)[1] == 0 and args.freqs_params(2)[1] == 1024 # yarn only on compressing layers + assert args.swa_decoder_replay == "bounded" + + +def test_shipping_layout_roles(): + """The real 40-layer layout: 3 encoder Full layers (ratio 2), the decoder Full at 20 building the + candidate pool, Reindex at 24/28/32/36 inside it, Reuse elsewhere.""" + text = tiny_hf_config().text_config + text.update( + num_hidden_layers=40, compress_ratios=[0, 0] + [2] * 18 + [1] * 20 + [0, 0, 0], + kv_source_layer_ids=[2, 8, 14, 20], index_source_layer_ids=[2, 8, 14, 20, 24, 28, 32, 36], + candidate_source_layer_id=20, num_nextn_predict_layers=3, engram_layer_ids=[1, 14], engram_num_embeddings=[10, 10], + ) + args = DeepseekV41Args.from_hf(tiny_hf_config(**text)) + full = [r.layer_id for r in args.roles if r.mode is Mode.FULL] + reindex = [r.layer_id for r in args.roles if r.mode is Mode.REINDEX] + assert full == [2, 8, 14, 20] and reindex == [24, 28, 32, 36] + assert sum(r.mode is Mode.REUSE for r in args.roles) == 40 - 2 - 4 - 4 + assert [r.layer_id for r in args.roles if r.uses_candidates] == reindex + assert args.decoder_start_layer == 20 + + +def test_bad_layouts_are_rejected(): + with pytest.raises(ValueError, match="no kv source"): + DeepseekV41Args.from_hf(tiny_hf_config(kv_source_layer_ids=[4])) + with pytest.raises(ValueError, match="reads kv source"): + DeepseekV41Args.from_hf(tiny_hf_config(compress_ratios=[0, 0, 2, 1, 1, 1, 0])) + with pytest.raises(ValueError, match="different kv source"): + DeepseekV41Args.from_hf(tiny_hf_config(compress_ratios=[0, 0, 2, 2, 2, 2, 0], kv_source_layer_ids=[2, 4], candidate_source_layer_id=2, index_source_layer_ids=[2, 4, 5])) + + +def test_registry_and_kv_sources(): + spec = get_model_spec("DeepseekV41ForCausalLM") + assert spec.module == "freetoken.models.deepseek_v41" + from freetoken.attention import AttnType + from freetoken.models.config import DSV4AttentionGroupConfig + + mc = parse_config(tiny_hf_config()) + args = mc.dsv41_args + assert args.backbone_kv_sources == (2, 4) and tuple(r.ratio for r in args.roles) == tuple(RATIOS[:N_LAYERS]) + (group,) = mc.attention_groups + assert type(group) is DSV4AttentionGroupConfig and group.variant == "v41" + assert mc.kv_cache_group_specs()[0].attn_type is AttnType.DSV41 and mc.attn_type_for_layer(0) is AttnType.DSV41 + assert mc.num_layers == N_LAYERS and mc.dsv41_args is not None + + +def _meta_model(mc): + if try_get_tp_info() is None: + set_tp_info(0, 1) + set_rope_device(torch.device("cpu")) + mc = replace(mc, moe_strategy="offload", decode_target="gpu", quant=NoQuantConfig()) + with torch.device("meta"), torch_dtype(torch.bfloat16): + return create_model(mc) + + +def test_loader_names_every_model_parameter(tmp_path): + from freetoken.models.deepseek_v41.weight import iter_weights + + write_tiny_checkpoint(str(tmp_path)) + model = _meta_model(parse_config(tiny_hf_config())) + params = model.state_dict() + loaded = dict(iter_weights(str(tmp_path), torch.device("cpu"), include_moe_experts=False)) + assert set(loaded) == set(params), (set(loaded) ^ set(params)) + for name, tensor in loaded.items(): + assert tuple(tensor.shape) == tuple(params[name].shape), name + # routed experts, the DSpark stage, the vision tower and the image routing bias are not resident weights + assert not any(".experts." in k or k.startswith(("mtp.", "vision.")) or k.endswith("bias_vl") for k in loaded) + # per-role structure + assert model.model.layers.op_list[3].attn.compressor is None and model.model.layers.op_list[3].attn.indexer is None + assert model.model.layers.op_list[4].attn.compressor.ratio == 1 and model.model.layers.op_list[2].attn.compressor.ratio == 2 + assert [b.engram is not None for b in model.model.layers.op_list] == [L in ENGRAM_LAYERS for L in range(N_LAYERS)] + assert all(L in INDEX_SOURCES for L in range(N_LAYERS) if model.model.layers.op_list[L].attn.indexer is not None) + assert all(L in KV_SOURCES for L in range(N_LAYERS) if model.model.layers.op_list[L].attn.compressor is not None) + + +@pytest.mark.parametrize("missing_weight", [False, True]) +def test_loader_closes_shards_when_expert_loading_stops(tmp_path, monkeypatch, missing_weight): + from freetoken.models.deepseek_v41 import weight + + write_tiny_checkpoint(str(tmp_path)) + closed = [] + close, get = weight._ShardReader.close, weight._ShardReader.get + + def close_reader(reader): + close(reader) + closed.append(True) + + def read_weight(reader, name): + if missing_weight and name.endswith(".w3.weight"): + raise KeyError(name) + return get(reader, name) + + monkeypatch.setattr(weight._ShardReader, "close", close_reader) + monkeypatch.setattr(weight._ShardReader, "get", read_weight) + tensors = weight.iter_weights(str(tmp_path), torch.device("cpu")) + if missing_weight: + with pytest.raises(KeyError): + next(tensors) + else: + next(tensors) + tensors.close() + assert closed == [True] + + +def test_fp8_wo_a_layer_rejects_a_dequantized_ftw_layout_clearly(): + """An fp8-declared wo_a fed a state dict without ``wo_a_scale`` (an FTW converted before wo_a kept + its payload) fails with a reconvert message, not a raw missing-key error; a bf16-declared layer fed + the fp8 payload + scale dequantizes it on the way in.""" + from types import SimpleNamespace + + from freetoken.layers.quantization.scheme import fp8_block_scheme + from freetoken.models.deepseek_v41.attention import DSV41Attention + + if try_get_tp_info() is None: + set_tp_info(0, 1) + args = parse_config(tiny_hf_config()).dsv41_args + args.max_seq_len = 512 + # wo_a is fp8; the LinearReplicated modules stay unquantized (the bf16 quant method) + plain = NoQuantConfig() + quant = SimpleNamespace( + scheme_for=lambda name: fp8_block_scheme("ue8m0", 32) if name.endswith(".wo_a") else None, + get_quant_method=lambda layer, prefix: plain.get_quant_method(layer, prefix), + ) + with torch.device("meta"), torch_dtype(torch.bfloat16): + fp8_layer = DSV41Attention(args, args.roles[2], quant_config=quant, prefix="attn") + bf16_layer = DSV41Attention(args, args.roles[2], quant_config=None, prefix="attn") + assert fp8_layer.wo_a.dtype == torch.float8_e4m3fn and bf16_layer.wo_a.dtype == torch.bfloat16 + rows, cols = fp8_layer.wo_a.shape + with torch.device("meta"): + state = {k: torch.empty_like(v) for k, v in fp8_layer.state_dict(prefix="attn").items()} + del state["attn.wo_a_scale"] + state["attn.wo_a"] = torch.empty(rows, cols, dtype=torch.bfloat16, device="meta") + with pytest.raises(RuntimeError, match="reconvert"): + fp8_layer.load_state_dict(dict(state), prefix="attn") + # the bf16 layer given the fp8 payload + scale dequantizes (real tensors: the dequant runs) + w = (torch.randn(rows, cols) * 0.05).to(torch.float8_e4m3fn) + scale = torch.full((rows // 32, cols // 32), 127, dtype=torch.uint8).view(torch.float8_e8m0fnu) + state2 = {k: torch.zeros(v.shape, dtype=v.dtype) for k, v in bf16_layer.state_dict(prefix="attn").items()} + state2["attn.wo_a"], state2["attn.wo_a_scale"] = w, scale + bf16_layer.load_state_dict(state2, prefix="attn") + assert torch.equal(bf16_layer.wo_a, w.to(torch.bfloat16)) diff --git a/tests/models/deepseek_v41/test_engine_config.py b/tests/models/deepseek_v41/test_engine_config.py new file mode 100644 index 0000000000..116209ea51 --- /dev/null +++ b/tests/models/deepseek_v41/test_engine_config.py @@ -0,0 +1,81 @@ +"""Engine config resolution for DeepSeek-V4.1 on the tiny checkpoint: the DSV41 backend and pool are +picked, the page size is the window page, the replay knob lands on the model args, and the fp8 +dialect (block 32, fp4 experts) reaches the quant layer with the matching activation block.""" + +from __future__ import annotations + +import json +import os + +import pytest +import torch + +from freetoken.distributed import DistributedInfo +from freetoken.scheduler.config import SchedulerConfig +from freetoken.engine.engine import _adjust_config +from freetoken.kvcache import resolve_pool_class +from freetoken.kvcache.dsv4.v41_pool import DSV41PagedKVCache + +from .common import write_tiny_checkpoint + + +def _engine_config(path, **over): + return SchedulerConfig(model_path=path, tp_info=DistributedInfo(rank=0, size=1), dtype=torch.bfloat16, **over) + + +@pytest.mark.parametrize("replay", ["bounded", "exact"]) +def test_resolution_picks_dsv41_and_the_replay_knob(tmp_path, monkeypatch, replay): + from freetoken.engine import engine + + monkeypatch.setattr(engine, "is_sm100_family", lambda: False) + monkeypatch.setattr(engine, "is_sm90_family", lambda: True) + write_tiny_checkpoint(str(tmp_path)) + config = _engine_config(str(tmp_path), attention_backend="auto", moe_strategy="offload", swa_decoder_replay=replay, max_seq_len_override=2048) + _adjust_config(config) + assert config.attention_backend == "dsv41_sparse" + assert config.page_size == 128 and config.cache_type == "swa_radix" + assert resolve_pool_class(config.model_config) is DSV41PagedKVCache + args = config.model_config.dsv41_args + assert args.swa_decoder_replay == replay and args.max_seq_len == 2048 and args.max_batch_size == config.max_running_req + 1 + assert config.max_extend_tokens == 8192 # the prefill chunk stays bounded (whole window pages) + # the cache contract follows the replay mode: a bounded-mode prefix hit must leave the prompt's last + # window to the prefill (the pool's prefix_replay_tokens); the history a resume reads stays one window + from freetoken.kvcache.dsv4.v41_cost_model import dsv41_pool_sizes + + spec = config.model_config.kv_cache_group_specs()[0] + pool = DSV41PagedKVCache(dsv41_pool_sizes(16, args, 1.0, 128), args, torch.device("cpu")) + assert pool.prefix_replay_tokens == (128 if replay == "bounded" else 0) + assert pool.sliding_window_size == 128 and spec.sliding_window == 128 and not spec.is_swa + # bounded replay keeps the decoder's per-request window KV in private rings, off the shared pages + assert pool.private_window_layer_ids == (tuple(range(args.decoder_start_layer, args.n_layers)) if replay == "bounded" else ()) + assert DSV41PagedKVCache.min_kv_tokens(config) // 128 == 8 + 3 * config.max_running_req + 2 * (config.max_running_req + 1) + 1 + + +def test_fp8_block32_dialect_reaches_the_quant_layer(tmp_path): + """A V4.1-style quantization_config on the tiny checkpoint: dense linears get the 32-block + e8m0 scheme, routed experts the MXFP4 scheme with a 32-wide activation block.""" + from freetoken.distributed.info import set_tp_info, try_get_tp_info + from freetoken.layers import OffloadMoELayer + from freetoken.layers.quantization import Fp8BlockConfig, QuantKind + from freetoken.layers.quantization.scheme import fp8_block_size + + write_tiny_checkpoint(str(tmp_path)) + cfg_path = os.path.join(str(tmp_path), "config.json") + with open(cfg_path) as f: + cfg = json.load(f) + cfg["quantization_config"] = {"quant_method": "fp8", "activation_scheme": "dynamic", "weight_block_size": [32, 32], "scale_fmt": "ue8m0", "expert_dtype": "fp4"} + with open(cfg_path, "w") as f: + json.dump(cfg, f) + config = _engine_config(str(tmp_path), moe_strategy="offload") + quant = config.model_config.quant + assert type(quant) is Fp8BlockConfig and quant.block == 32 + dense = quant.scheme_for("model.layers.3.attn.wq_a") + assert dense.kind is QuantKind.FP8_BLOCK and fp8_block_size(dense) == 32 + if try_get_tp_info() is None: + set_tp_info(0, 1) + experts = OffloadMoELayer(3, num_experts=4, top_k=2, hidden_size=64, intermediate_size=64, limit=10.0, quant_config=quant, prefix="model.layers.3.ffn.experts") + assert experts.quant_method.scheme.kind is QuantKind.MXFP4 and experts.quant_method.cfg.act_block == 32 + # the family's bf16 modules stay unquantized although the fp8 config lists no exceptions + for name in ("model.layers.2.attn.compressor.wkv", "model.layers.2.attn.indexer.wk", "model.layers.5.attn.indexer.weights_proj", "head"): + assert quant.scheme_for(name) is None, name + assert quant.scheme_for("model.layers.1.engram.wkv").kind is QuantKind.FP8_BLOCK diff --git a/tests/models/deepseek_v41/test_engram.py b/tests/models/deepseek_v41/test_engram.py new file mode 100644 index 0000000000..dec852ec03 --- /dev/null +++ b/tests/models/deepseek_v41/test_engram.py @@ -0,0 +1,281 @@ +"""Engram: the hash against a transcription of the reference ``NgramHashState``, the prime buckets, +the disk table against a torch oracle over the synthetic checkpoint (eager, decode with context, +CUDA-graph flag sync), and the layer's gate math against the reference ``Engram.forward``.""" + +from __future__ import annotations + +import numpy as np +import pytest +import torch + +from freetoken.models.deepseek_v41.args import DeepseekV41Args +from freetoken.models.deepseek_v41.engram import EngramHash, bucket_primes, hash_multipliers + +from .common import VOCAB, requires_cuda, tiny_hf_config, write_tiny_checkpoint + + +def _args(**over) -> DeepseekV41Args: + return DeepseekV41Args.from_hf(tiny_hf_config(**over)) + + +def _reference_hashes(args: DeepseekV41Args, ids: torch.Tensor, start_pos: int, cache: torch.Tensor) -> torch.Tensor: + """``NgramHashState.forward`` for one sequence with an identity token map; ``cache`` is the + sequence's compressed-id history buffer (mutated like the reference's).""" + layout_primes = bucket_primes(args) + flat = [[p for per_ngram in layer for p in per_ngram] for layer in layout_primes] + offsets = torch.tensor(np.array([np.cumsum([0, *sizes[:-1]]) for sizes in flat])) + primes = torch.tensor(layout_primes) + multipliers = hash_multipliers(args.engram_layer_ids, args.engram_max_ngram_size, args.engram_compressed_vocab_size) + pad_id = args.engram_pad_id + seqlen = ids.numel() + cache[start_pos : start_pos + seqlen] = ids + positions = torch.arange(start_pos, start_pos + seqlen) + tokens, blocked = [], torch.zeros_like(positions, dtype=torch.bool) + for shift in range(args.engram_max_ngram_size): + source = cache[(positions - shift).clamp_min(0)] + blocked = blocked | (positions < shift) + tokens.append(torch.where(blocked, pad_id, source)) + tokens = torch.stack(tokens, dim=-1) + products = tokens.unsqueeze(1) * multipliers + rolling, hashes = products[..., 0], [] + for i in range(1, args.engram_max_ngram_size): + rolling = torch.bitwise_xor(rolling, products[..., i]) + hashes.append(rolling.unsqueeze(-1) % primes[:, i - 1]) + return torch.cat(hashes, dim=-1) + offsets + + +def test_primes_are_distinct_and_ordered(): + args = _args() + primes = bucket_primes(args) + flat = [p for layer in primes for per in layer for p in per] + assert len(flat) == len(set(flat)) and flat == sorted(flat) and min(flat) > args.engram_vocab_size - 1 + from sympy import isprime + + assert all(isprime(p) for p in flat) + + +def test_hash_matches_the_reference_across_a_prefill_decode_split(): + args = _args() + hash = EngramHash(args, list(range(VOCAB)), VOCAB) + ids = torch.randint(3, VOCAB, (37,), generator=torch.Generator().manual_seed(0)) + cache = torch.empty(64, dtype=torch.int64) + want = _reference_hashes(args, ids[:30], 0, cache) + got = hash.row_ids(ids[:30]) + assert torch.equal(got, want) + # decode continues with the three preceding tokens as context + for step in range(30, 37): + want = _reference_hashes(args, ids[step : step + 1], step, cache) + got = hash.row_ids(ids[step : step + 1], context=ids[:step]) + assert torch.equal(got, want), step + # a run resumed mid-sequence (chunked prefill) hashes like the uninterrupted one + whole = hash.row_ids(ids) + assert torch.equal(torch.cat([hash.row_ids(ids[:20]), hash.row_ids(ids[20:], context=ids[:20])]), whole) + assert whole.shape == (37, 1, hash.n_cols) and int(whole.max()) < hash.rows_per_layer[0] + with pytest.raises(ValueError, match="compressed vocab"): + EngramHash(args, list(range(VOCAB)), VOCAB + 1) + + +@requires_cuda +def test_disk_table_matches_the_dequantized_shard(tmp_path): + from freetoken.models.deepseek_v41.engram_table import EngramDiskTable, engram_row_source + + tensors = write_tiny_checkpoint(str(tmp_path)) + args = _args() + hash = EngramHash(args, list(range(VOCAB)), VOCAB) + src = engram_row_source(str(tmp_path), 1) + assert src.num_rows == tensors["layers.1.engram.embed.weight"].shape[0] and src.head_dim == 32 + table = EngramDiskTable(src, hash.n_cols, torch.device("cuda"), max_graph_rows=4, max_extend_tokens=16) + fp8 = tensors["layers.1.engram.embed.weight"] + scale = torch.exp2(tensors["layers.1.engram.embed.scale"].view(torch.uint8).float() - 127.0) + + def oracle(rows: torch.Tensor) -> torch.Tensor: + vals = fp8[rows.flatten()].float().view(-1, 32 // 32, 32) * scale[rows.flatten()].unsqueeze(-1) + return vals.view(rows.shape[0], -1).to(torch.bfloat16).cuda() + + ids = torch.randint(3, VOCAB, (40,), generator=torch.Generator().manual_seed(1)) + rows = hash.row_ids(ids)[:, 0] + table.stage(rows, graph=False) # larger than the initial eager staging: grows + table.flush(signal=False) + assert torch.equal(table.lookup(40), oracle(rows)) + # the graph staging + flag protocol: capture a lookup, fill after the launch, replay + from freetoken.kernel.row_store import probe_wait_sync + + table.wait_sync = probe_wait_sync("auto", torch.device("cuda")) + dec_rows = hash.row_ids(ids[40 - 1 :], context=ids[: 40 - 1])[:, 0] + out = torch.empty(1, hash.n_cols * 32, dtype=torch.bfloat16, device="cuda") + graph = torch.cuda.CUDAGraph() + stream = torch.cuda.Stream() + with torch.cuda.stream(stream): + if not table.wait_sync: + table.stage(dec_rows, graph=True) + table.flush(signal=False) + with torch.cuda.graph(graph, stream=stream): + out.copy_(table.lookup(1)) + torch.cuda.synchronize() + graph.replay() # under wait-sync the replay blocks until the host signals + table.stage(dec_rows, graph=True) + table.flush(signal=table.wait_sync) + torch.cuda.synchronize() + assert torch.equal(out, oracle(dec_rows)) + + +@requires_cuda +def test_eager_staging_does_not_overwrite_an_in_flight_lookup(tmp_path): + """The overlap scheduler stages batch k+1 while batch k's lookup copy is still queued behind + the layers before the Engram layer (a sleep kernel here): batch k must read its own rows.""" + from freetoken.models.deepseek_v41.engram_table import EngramDiskTable, engram_row_source + + tensors = write_tiny_checkpoint(str(tmp_path)) + hash = EngramHash(_args(), list(range(VOCAB)), VOCAB) + table = EngramDiskTable(engram_row_source(str(tmp_path), 1), hash.n_cols, torch.device("cuda"), max_graph_rows=4, max_extend_tokens=64) + fp8 = tensors["layers.1.engram.embed.weight"] + scale = torch.exp2(tensors["layers.1.engram.embed.scale"].view(torch.uint8).float() - 127.0) + + def oracle(rows): + vals = fp8[rows.flatten()].float().view(-1, 1, 32) * scale[rows.flatten()].unsqueeze(-1) + return vals.view(rows.shape[0], -1).to(torch.bfloat16).cuda() + + gen = torch.Generator().manual_seed(0) + batches = [hash.row_ids(torch.randint(3, VOCAB, (40,), generator=gen))[:, 0] for _ in range(4)] + got = [] + for rows in batches: + table.stage(rows, graph=False) + table.flush(signal=False) + torch.cuda._sleep(50_000_000) + got.append(table.lookup(rows.shape[0])) + torch.cuda.synchronize() + for i, (rows, out) in enumerate(zip(batches, got)): + assert torch.equal(out, oracle(rows)), f"batch {i} read another batch's rows" + + +@requires_cuda +@torch.inference_mode() +def test_host_fills_a_graph_decode_step_from_a_worker_thread(tmp_path): + """The deferred graph-decode fill must not wait for the launch call to return: a 40-layer graph + blocks cuGraphLaunch on the host while the GPU sits on the WAIT, so the host coordinator reads the + tokens back and signals from a worker thread. The worker hashes against the DISPATCH snapshot of + the request state: advancing the request right after dispatch (what the engine does) must not + change which rows are fetched.""" + from types import SimpleNamespace + + from freetoken.models.deepseek_v41.engram_table import EngramDiskTable, EngramHost, engram_row_source + + tensors = write_tiny_checkpoint(str(tmp_path)) + args = _args() + hash = EngramHash(args, list(range(VOCAB)), VOCAB) + table = EngramDiskTable(engram_row_source(str(tmp_path), 1), hash.n_cols, torch.device("cuda"), max_graph_rows=4, max_extend_tokens=16) + host = EngramHost(hash, [table], torch.device("cuda")) + if not host.wait_sync: + pytest.skip("stream memops unavailable: launch-gating mode needs no worker") + ids = torch.randint(3, VOCAB, (20,), generator=torch.Generator().manual_seed(7)) + req = SimpleNamespace(input_ids=ids.clone(), device_len=20, cached_len=19) + batch = SimpleNamespace(is_decode=True, is_prefill=False, padded_reqs=[req], padded_size=1, input_ids=ids[19:20].to(torch.int32).cuda()) + out = torch.empty(1, hash.n_cols * 32, dtype=torch.bfloat16, device="cuda") + stream = torch.cuda.Stream() + graph = torch.cuda.CUDAGraph() + with torch.cuda.stream(stream): + with torch.cuda.graph(graph, stream=stream): + out.copy_(table.lookup(1)) + torch.cuda.synchronize() + with torch.cuda.stream(stream): + # slow the readback so the request advances (complete_one + append) before the worker runs + torch.cuda._sleep(200_000_000) + with host.forward_host_ctx(batch, use_graph=True): + graph.replay() + req.device_len += 1 + req.input_ids = torch.cat([req.input_ids, torch.tensor([5], dtype=ids.dtype)]) + torch.cuda.synchronize() + fp8 = tensors["layers.1.engram.embed.weight"] + rows = hash.row_ids(ids[19:20], context=ids[:19])[:, 0].flatten() + want = fp8[rows].float().view(1, -1).to(torch.bfloat16).cuda() # scales are all 1.0 in the synthetic table + assert torch.equal(out, want) + host.close() + + +@requires_cuda +def test_host_surfaces_a_failed_deferred_fill_at_dispatch_exit(tmp_path): + from types import SimpleNamespace + + from freetoken.models.deepseek_v41.engram_table import EngramDiskTable, EngramHost, engram_row_source + + write_tiny_checkpoint(str(tmp_path)) + args = _args() + hash = EngramHash(args, list(range(VOCAB)), VOCAB) + table = EngramDiskTable(engram_row_source(str(tmp_path), 1), hash.n_cols, torch.device("cuda"), max_graph_rows=4, max_extend_tokens=16) + host = EngramHost(hash, [table], torch.device("cuda")) + if not host.wait_sync: + pytest.skip("stream memops unavailable") + ids = torch.randint(3, VOCAB, (8,), generator=torch.Generator().manual_seed(8)) + req = SimpleNamespace(input_ids=ids, device_len=8, cached_len=7) + batch = SimpleNamespace(is_decode=True, is_prefill=False, padded_reqs=[req], padded_size=1, input_ids=ids[7:8].to(torch.int32).cuda()) + stream = torch.cuda.Stream() + graph = torch.cuda.CUDAGraph() + out = torch.empty(1, hash.n_cols * 32, dtype=torch.bfloat16, device="cuda") + with torch.cuda.stream(stream): + with torch.cuda.graph(graph, stream=stream): + out.copy_(table.lookup(1)) + torch.cuda.synchronize() + table.stage = lambda *a, **k: (_ for _ in ()).throw(RuntimeError("disk gone")) # the fill fails + with pytest.raises(RuntimeError, match="disk gone"): + with torch.cuda.stream(stream): + with host.forward_host_ctx(batch, use_graph=True): + graph.replay() + torch.cuda.synchronize() # the failing worker released the flag: the stream is not stuck + host.close() + + +def test_ftw_side_files_carry_the_tables_and_reopen(tmp_path): + """An FTW conversion copies the shards holding the Engram tables next to it (the converter's + ``ftw_side_files`` hook, which skips the HF index); the row source resolves the copies through + their headers exactly like the HF shards, byte for byte.""" + from freetoken.kernel.row_store import RowStore + from freetoken.models.deepseek_v41.engram_table import engram_row_source + from freetoken.models.deepseek_v41.weight import ftw_side_files + + src, out = tmp_path / "hf", tmp_path / "ftw" + out.mkdir() + tensors = write_tiny_checkpoint(str(src)) + (out / "config.json").write_bytes((src / "config.json").read_bytes()) # the converter copies the metadata + shard = "model-00001-of-00001.safetensors" # the tiny checkpoint's single shard holds the table + assert ftw_side_files(str(src), str(out)) == [shard] + assert (out / shard).read_bytes() == (src / shard).read_bytes() and not (out / "model.safetensors.index.json").exists() + source = engram_row_source(str(out), 1) + assert source.path == str(out / shard) and source.num_rows == 4096 and source.head_dim == 32 + assert torch.equal(source.scales, tensors["layers.1.engram.embed.scale"].view(torch.uint8)) + store = RowStore(paths=[source.path], extent_file=[0], extent_base=[source.base], rows_per_extent=source.num_rows, + row_bytes=source.head_dim, row_stride=source.head_dim, use_io_uring=False) + ids = torch.tensor([0, 7, 4095], dtype=torch.int64) + buf = torch.zeros(3 * 32, dtype=torch.uint8) + store.stage_rows(ids.data_ptr(), 3, buf.data_ptr(), 0) + store.flush(0) + assert torch.equal(buf.view(3, 32), tensors["layers.1.engram.embed.weight"][ids].view(torch.uint8)) + + +def test_layer_gate_matches_the_reference(): + """``EngramLayer.forward`` against the reference ``Engram.forward`` (bf16 wkv, fp32 gate math).""" + from freetoken.distributed.info import set_tp_info, try_get_tp_info + from freetoken.models.deepseek_v41.engram import EngramLayer + + if try_get_tp_info() is None: + set_tp_info(0, 1) + args = _args() + torch.manual_seed(0) + layer = EngramLayer(args, 1) + T, hc, dim = 5, args.hc_mult, args.dim + wkv = torch.randn(dim * (hc + 1), layer.width, dtype=torch.bfloat16) * 0.05 + layer.wkv.weight = wkv + layer.q_weight = torch.rand(hc, dim) + 0.5 + layer.k_weight = torch.rand(hc, dim) + 0.5 + x = torch.randn(T, hc, dim, dtype=torch.bfloat16) + rows = torch.randn(T, layer.width, dtype=torch.bfloat16) + + kv = torch.nn.functional.linear(rows, wkv) + key, value = kv.split([hc * dim, dim], dim=-1) + key = key.float().view(T, hc, dim) + weight = layer.q_weight * layer.k_weight + h = x.float() + rstd = torch.rsqrt(h.square().mean(-1) + args.norm_eps) * torch.rsqrt(key.square().mean(-1) + args.norm_eps) + dot = (h * weight * key).sum(-1) * rstd * dim**-0.5 + gate = torch.sigmoid(torch.copysign(dot.abs().clamp_min(1e-6).sqrt(), dot)) + want = (h + gate.unsqueeze(-1) * value.float().unsqueeze(-2)).to(x.dtype) + assert torch.equal(layer.forward(x, rows), want) diff --git a/tests/models/deepseek_v41/test_forward_smoke.py b/tests/models/deepseek_v41/test_forward_smoke.py new file mode 100644 index 0000000000..e0861e8874 --- /dev/null +++ b/tests/models/deepseek_v41/test_forward_smoke.py @@ -0,0 +1,184 @@ +"""End-to-end forward smoke on the tiny model: cold prefill, batched decode, a second request, chunked +prefill and bounded replay all produce finite logits of the right shape (numerics: test_reference_parity).""" + +from __future__ import annotations + +import pytest +import torch + +from .common import VOCAB, requires_cuda, write_tiny_checkpoint + +pytestmark = requires_cuda + + +@pytest.fixture(scope="module") +def checkpoint(tmp_path_factory): + folder = tmp_path_factory.mktemp("dsv41-tiny") + write_tiny_checkpoint(str(folder)) + return str(folder) + + +def _tokens(n: int, seed: int) -> list[int]: + g = torch.Generator().manual_seed(seed) + return torch.randint(3, VOCAB, (n,), generator=g).tolist() + + +def test_prefill_and_decode(checkpoint): + from .harness import TinyEngine + + eng = TinyEngine(checkpoint, max_seq_len=1024, max_running_req=2, swa_decoder_replay="exact") + r0 = eng.new_request(0, _tokens(300, 1)) + logits = eng.prefill([r0]) + assert logits.shape == (1, VOCAB) and torch.isfinite(logits).all() + eng.finish_prefill([r0]) + r1 = eng.new_request(1, _tokens(50, 2)) + assert torch.isfinite(eng.prefill([r1])).all() + eng.finish_prefill([r1]) + for step in range(5): # batched decode over two requests at different positions + out = eng.decode([r0, r1], [10 + step, 20 + step]) + assert out.shape == (2, VOCAB) and torch.isfinite(out).all() + + +def test_chunked_prefill_matches_single_shot(checkpoint): + """Two 128-aligned chunks must reproduce the single-shot prefill's last-token logits: the + compressor carry, the window ring and the compressed rows are all resumed through the pool. The + chunks run their projections at other row counts, so the match is to rounding, not bitwise.""" + from .harness import TinyEngine + from .test_reference_parity import _compare + + eng = TinyEngine(checkpoint, max_seq_len=1024, max_running_req=2, swa_decoder_replay="exact") + toks = _tokens(300, 3) + whole = eng.new_request(0, toks) + ref = eng.prefill([whole]) + chunked = eng.new_request(1, toks) + chunked.device_len = 256 # first chunk [0, 256) + eng.prefill([chunked]) + chunked.cached_len, chunked.device_len = 256, 300 + got = eng.prefill([chunked]) + _compare("chunked prefill", got, ref) + + +def test_bounded_replay_is_exact_within_one_window(checkpoint): + """A prompt no longer than the window replays every token with no floor: bounded == exact.""" + from .harness import TinyEngine + from .test_reference_parity import _compare + + toks = _tokens(100, 5) + exact = TinyEngine(checkpoint, max_seq_len=1024, max_running_req=2, swa_decoder_replay="exact") + want = exact.prefill([exact.new_request(0, toks)]) + bounded = TinyEngine(checkpoint, max_seq_len=1024, max_running_req=2, swa_decoder_replay="bounded") + got = bounded.prefill([bounded.new_request(0, toks)]) + _compare("bounded within one window", got, want) + + +@pytest.mark.parametrize("cut", [128, 256]) +def test_bounded_chunked_prefill_matches_unchunked(checkpoint, cut): + """A chunked bounded prefill gives each chunk its part of the prompt's last window ``[172, 300)``: + a first chunk that ends before it (cut 128) runs no decoder row and hands the head a placeholder; + one that reaches into it (cut 256) leaves rows in the private ring for the last chunk, which is + shorter than the window (44 tokens). Both match the unchunked prefill.""" + from .harness import TinyEngine + from .test_reference_parity import _compare + + eng = TinyEngine(checkpoint, max_seq_len=1024, max_running_req=2, swa_decoder_replay="bounded") + toks = _tokens(300, 6) + ref = eng.prefill([eng.new_request(0, toks)]) + chunked = eng.new_request(1, toks) + chunked.device_len = cut + first = eng.prefill([chunked]) + assert first.shape == (1, VOCAB) and torch.isfinite(first).all() + chunked.cached_len, chunked.device_len = cut, 300 + _compare(f"bounded chunks cut at {cut}", eng.prefill([chunked]), ref) + + +def test_bounded_batch_mixes_a_decoderless_chunk_with_a_whole_prompt(checkpoint): + """One batch: a chunk that ends before its prompt's last window (no decoder row, a placeholder for + the head) next to a whole prompt; the whole prompt's logits are those of its own prefill.""" + from .harness import TinyEngine + from .test_reference_parity import _compare + + eng = TinyEngine(checkpoint, max_seq_len=2048, max_running_req=2, swa_decoder_replay="bounded") + whole = _tokens(300, 8) + want = eng.prefill([eng.new_request(0, whole)]) + chunk = eng.new_request(0, _tokens(600, 9)) + chunk.device_len = 256 + other = eng.new_request(1, whole) + got = eng.prefill([chunk, other]) + assert got.shape == (2, VOCAB) and torch.isfinite(got).all() + _compare("whole prompt next to a decoderless chunk", got[1:], want) + + +def test_bounded_prefix_hit_matches_cold_prefill(checkpoint): + """A prefix hit that would leave fewer than a window of new tokens (768 cached + 20 new) is capped + one window before the prompt end, so the request prefills ``[640, 788)`` -- the decoder's whole + window among it -- and produces the cold prefill's logits.""" + from .harness import TinyEngine + from .test_reference_parity import _compare + + eng = TinyEngine(checkpoint, max_seq_len=2048, max_running_req=2, swa_decoder_replay="bounded") + prefix = _tokens(768, 11) + donor = eng.new_request(0, prefix + _tokens(20, 12)) + cold = eng.prefill([donor]) + eng.finish_prefill([donor]) + hit = eng.new_request_on_prefix(1, donor, prefix + _tokens(20, 12)) # the same prompt over the shared pages + assert hit.cached_len == 640 and hit.extend_len == 148 + got = eng.prefill([hit]) + _compare("bounded prefix hit", got, cold) + + +def test_shared_prefix_replays_do_not_disturb_each_other(checkpoint): + """Two requests sharing a 768-token prefix with different suffix lengths: B's hit is capped at 640 + (a window before its prompt end), its decoder writes only B's private ring and it reads the shared + encoder history, so A's pending decode sees the logits it would have seen without B, and neither + the shared pages nor A's ring change.""" + from .harness import TinyEngine + + prefix = _tokens(768, 13) + a_toks, b_toks = prefix + _tokens(20, 14), prefix + _tokens(40, 15) + + def run(interleave: bool): + eng = TinyEngine(checkpoint, max_seq_len=2048, max_running_req=2, swa_decoder_replay="bounded") + a = eng.new_request(0, a_toks) + eng.prefill([a]) + eng.finish_prefill([a]) + dec_layer, enc_layer = eng.args.decoder_start_layer, 2 + shared_slots = eng.pool.translate_full_to_window(eng.page_table[0, 640:768].long()) + enc_before = eng.pool.window_pool[enc_layer][shared_slots].clone() + P = eng.args.window_size + ring_before = eng.pool.window_pool[dec_layer][0 * P : 1 * P].clone() # A's ring row + if interleave: + b = eng.new_request_on_prefix(1, a, b_toks) + eng.prefill([b]) + eng.finish_prefill([b]) + assert torch.equal(eng.pool.window_pool[enc_layer][shared_slots], enc_before) + assert torch.equal(eng.pool.window_pool[dec_layer][0 * P : 1 * P], ring_before) + return eng.decode([a], [7]) + + assert torch.equal(run(True), run(False)) + + +def test_commit_dedup_onto_a_longer_prompts_pages_keeps_the_decoder_state(checkpoint): + """The radix commit may replace a request's freshly written pages with an existing node's (the + reviewer's donor case): D prefilled 256 tokens, so under bounded replay its decoder never computed + positions [0, 128); A (exactly D's first 128 tokens) is repointed onto D's pages at commit. A's + decoder KV lives in A's own ring, so its next decode is unchanged by the repoint (D's shared rows + came from a longer prefill, so to rounding).""" + from .harness import TinyEngine + from .test_reference_parity import _compare + + donor_toks = _tokens(256, 21) + a_toks = donor_toks[:128] + + def run(repoint: bool): + eng = TinyEngine(checkpoint, max_seq_len=1024, max_running_req=2, swa_decoder_replay="bounded") + d = eng.new_request(0, donor_toks) + eng.prefill([d]) + eng.finish_prefill([d]) + a = eng.new_request(1, a_toks) + eng.prefill([a]) + eng.finish_prefill([a]) + if repoint: # what CacheManager.cache_req does when the tree already holds the prefix + eng.page_table[1, :128] = eng.page_table[0, :128] + return eng.decode([a], [3]) + + _compare("decode after a commit repoint", run(True), run(False)) diff --git a/tests/models/deepseek_v41/test_reference_parity.py b/tests/models/deepseek_v41/test_reference_parity.py new file mode 100644 index 0000000000..839cdafb82 --- /dev/null +++ b/tests/models/deepseek_v41/test_reference_parity.py @@ -0,0 +1,219 @@ +"""Logits parity against the vendored reference implementation on the tiny synthetic checkpoint. + +The reference runs bf16 or FP8/FP4 weights with its own torch quantizers, sparse attention, engram +hash and single-pass mHC; FreeToken runs the same checkpoint through its kernels and packed pools. +Compared: exact-mode prefill, greedy decode steps, a two-request batch, and the bounded-replay path +against the reference's ``forward_bounded`` oracle. +""" + +from __future__ import annotations + +from dataclasses import asdict, fields + +import pytest +import torch + +from .common import VOCAB, requires_cuda, tiny_text_config, write_tiny_checkpoint + +pytestmark = requires_cuda + +# The sparse-attention implementations accumulate in different orders. +ATOL = RTOL = 1e-2 +# The fp4 experts also round the routed sum to bf16 before the shared-expert add and weight each +# expert's down output; the reference keeps that sum fp32 and weights the intermediate before its fp8 +# round-trip (10 seeds x 16 checks: max |err| 0.019, against 0.007 in the reference's order). +QUANT_TOL = 2e-2 + + +@pytest.fixture(scope="module") +def checkpoint(tmp_path_factory): + folder = tmp_path_factory.mktemp("dsv41-parity") + tensors = write_tiny_checkpoint(str(folder), seed=11) + return str(folder), tensors + + +def _tokens(n: int, seed: int) -> torch.Tensor: + return torch.randint(3, VOCAB, (n,), generator=torch.Generator().manual_seed(seed)) + + +class Reference: + """The vendored ``Transformer`` on CUDA, loaded from the synthetic checkpoint (identity token map).""" + + def __init__(self, tensors: dict, text: dict, *, max_seq_len: int, max_batch_size: int, quantized: bool = False, vision_config: dict | None = None): + from freetoken.models.deepseek_v41.args import DeepseekV41Args + from types import SimpleNamespace + + from . import reference + from .reference import engram as ref_engram + from .reference import model as ref_model + + args = DeepseekV41Args.from_hf(SimpleNamespace(text_config=text)) + names = {f.name for f in fields(ref_model.ModelArgs)} + kwargs = {k: v for k, v in asdict(args).items() if k in names} + kwargs.update(dtype="fp8" if quantized else "bf16", expert_dtype="fp4" if quantized else None, + max_seq_len=max_seq_len, max_batch_size=max_batch_size, vision_n_layers=0, temperature=0.0) + if vision_config: + kwargs.update(vision_n_layers=vision_config["num_hidden_layers"], vision_dim=vision_config["hidden_size"], + vision_n_heads=vision_config["num_attention_heads"], vision_inter_dim=vision_config["intermediate_size"], + vision_patch_size=vision_config["patch_size"], vision_downsample_ratio=vision_config["downsample_ratio"]) + self.margs = ref_model.ModelArgs(**kwargs) + # the synthetic checkpoint's tokenizer is the identity over VOCAB ids + ref_engram.build_compressed_token_map = lambda tokenizer: (list(range(VOCAB)), VOCAB) + torch.set_default_dtype(torch.bfloat16) + try: + with torch.device("cuda"): + self.model = ref_model.Transformer(self.margs, tokenizer=object()) + finally: + torch.set_default_dtype(torch.float32) + state = {k: v for k, v in tensors.items() if not k.startswith("mtp.") + and (vision_config or (not k.startswith("vision.") and not k.endswith("bias_vl")))} + if quantized: + state = {k: v.view(torch.float4_e2m1fn_x2) if ".ffn.experts." in k and k.endswith(".weight") else v + for k, v in state.items()} + # The reference's grouped output projection holds dequantized bf16 weights. + for name in list(state): + if name.endswith(".wo_a.scale"): + scale = state.pop(name).float().repeat_interleave(32, 0).repeat_interleave(32, 1) + weight_name = name.removesuffix(".scale") + ".weight" + state[weight_name] = (state[weight_name].float() * scale).bfloat16() + missing, unexpected = self.model.load_state_dict({k: v.cuda() for k, v in state.items()}, strict=False) + assert not unexpected, unexpected + assert not [m for m in missing if "freqs_cis" not in m], missing + self.reference = reference + + # the reference builds its index tensors on the default device (generate.py sets it to cuda) + def prefill(self, ids: torch.Tensor) -> torch.Tensor: + with torch.device("cuda"): + return self.model(ids.view(1, -1).cuda(), 0)[1].float() + + def prefill_batch(self, ids: torch.Tensor) -> torch.Tensor: + with torch.device("cuda"): + return self.model(ids.cuda(), 0)[1].float() + + def prefill_bounded(self, ids: torch.Tensor) -> torch.Tensor: + with torch.device("cuda"): + return self.model.forward_bounded(ids.reshape(-1, ids.shape[-1]).cuda(), self.margs.window_size).float() + + def decode(self, tokens: torch.Tensor, pos: int) -> torch.Tensor: + with torch.device("cuda"): + return self.model(tokens.view(-1, 1).cuda(), pos)[1].float() + + +def _engram_table_for(engine, tensors, text): + """Our engine reads the Engram rows straight from the synthetic shard (the disk table), hashed + with the same identity token map the reference oracle uses.""" + from freetoken.models.deepseek_v41.engram import EngramHash + from freetoken.models.deepseek_v41.engram_table import EngramDiskTable, EngramHost, engram_row_source + + hash = EngramHash(engine.args, list(range(VOCAB)), VOCAB) + tables = [ + EngramDiskTable(engram_row_source(engine.checkpoint, layer.layer_id), hash.n_cols, engine.device, max_graph_rows=8, max_extend_tokens=64) + for layer in engine.model.engram_layers() + ] + for layer, table in zip(engine.model.engram_layers(), tables): + layer.attach_table(table) + host = EngramHost(hash, tables, engine.device) + engine.model.forward_host_ctx = host.forward_host_ctx + engine.engram_host = host + + +def _compare(name, got, want, *, tol=ATOL): + err = (got - want).abs().max().item() + assert torch.equal(got.argmax(-1), want.argmax(-1)), f"{name}: argmax differs (max abs err {err:.4f})" + torch.testing.assert_close(got, want, atol=tol, rtol=tol, msg=lambda m: f"{name}: {m}") + + +@pytest.mark.parametrize("mode", ["exact", "bounded"]) +@pytest.mark.parametrize("batch_size", [1, 2]) +def test_quantized_offload_prefill_and_decode_match_reference(tmp_path, mode, batch_size): + from .harness import TinyEngine + + tensors = write_tiny_checkpoint(str(tmp_path), seed=11, quantized=True) + text = tiny_text_config(moe_intermediate_size=256) + ref = Reference(tensors, text, max_seq_len=1024, max_batch_size=3, quantized=True) + eng = TinyEngine(str(tmp_path), swa_decoder_replay=mode, quantized=True) + _engram_table_for(eng, tensors, text) + ids = torch.stack([_tokens(300, 1 + i) for i in range(batch_size)]) + want = ref.prefill_batch(ids) if mode == "exact" else ref.prefill_bounded(ids) + reqs = [eng.new_request(i, row.tolist()) for i, row in enumerate(ids)] + got = eng.prefill(reqs) + _compare(f"quantized {mode} prefill", got, want, tol=QUANT_TOL) + eng.finish_prefill(reqs) + for step in range(3): + token = want.argmax(-1) + want = ref.decode(token, ids.shape[-1] + step) + got = eng.decode(reqs, token.tolist()) + _compare(f"quantized {mode} decode {step}", got, want, tol=QUANT_TOL) + + +def test_exact_prefill_and_greedy_decode_match_the_reference(checkpoint): + from .harness import TinyEngine + + folder, tensors = checkpoint + text = tiny_text_config() + ref = Reference(tensors, text, max_seq_len=1024, max_batch_size=3) + eng = TinyEngine(folder, max_seq_len=1024, max_running_req=2, swa_decoder_replay="exact") + _engram_table_for(eng, tensors, text) + + ids = _tokens(300, 1) + want = ref.prefill(ids) + req = eng.new_request(0, ids.tolist()) + got = eng.prefill([req]) + _compare("prefill", got, want) + eng.finish_prefill([req]) + # greedy decode from the reference's own picks keeps both sides on one trajectory + pos = 300 + for step in range(12): + nxt = int(want.argmax(-1).item()) + want = ref.decode(torch.tensor([nxt]), pos) + got = eng.decode([req], [nxt]) + _compare(f"decode step {step}", got, want) + pos += 1 + + +def test_batched_prefill_and_decode_match_the_reference(checkpoint): + from .harness import TinyEngine + + folder, tensors = checkpoint + text = tiny_text_config() + ref = Reference(tensors, text, max_seq_len=1024, max_batch_size=3) + eng = TinyEngine(folder, max_seq_len=1024, max_running_req=2, swa_decoder_replay="exact") + _engram_table_for(eng, tensors, text) + + ids = torch.stack([_tokens(200, 2), _tokens(200, 3)]) + want = ref.prefill_batch(ids) + reqs = [eng.new_request(i, ids[i].tolist()) for i in range(2)] + got = eng.prefill(reqs) + _compare("batched prefill", got, want) + eng.finish_prefill(reqs) + pos = 200 + for step in range(6): + nxt = want.argmax(-1).cpu() + want = ref.decode(nxt, pos) + got = eng.decode(reqs, nxt.tolist()) + _compare(f"batched decode step {step}", got, want) + pos += 1 + + +def test_bounded_replay_matches_the_reference_oracle(checkpoint): + from .harness import TinyEngine + + folder, tensors = checkpoint + text = tiny_text_config() + ref = Reference(tensors, text, max_seq_len=1024, max_batch_size=3) + eng = TinyEngine(folder, max_seq_len=1024, max_running_req=2, swa_decoder_replay="bounded") + _engram_table_for(eng, tensors, text) + + ids = _tokens(300, 4) + want = ref.prefill_bounded(ids) + req = eng.new_request(0, ids.tolist()) + got = eng.prefill([req]) + _compare("bounded prefill", got, want) + eng.finish_prefill([req]) + pos = 300 + for step in range(6): + nxt = int(want.argmax(-1).item()) + want = ref.decode(torch.tensor([nxt]), pos) + got = eng.decode([req], [nxt]) + _compare(f"decode after bounded prefill, step {step}", got, want) + pos += 1 diff --git a/tests/models/deepseek_v41/test_vision.py b/tests/models/deepseek_v41/test_vision.py new file mode 100644 index 0000000000..0e1d4a4830 --- /dev/null +++ b/tests/models/deepseek_v41/test_vision.py @@ -0,0 +1,214 @@ +"""Image layout, official ViT/logits parity and image replay through the shared encoder cache.""" + +from __future__ import annotations + +import io +import json +from dataclasses import asdict +from types import SimpleNamespace + +import numpy as np +import pytest +import torch +from PIL import Image +from safetensors.torch import save_file + +from freetoken.core import Batch +from freetoken.engine.engine import Engine +from freetoken.mm.config import MultimodalConfig +from freetoken.mm.encoder_cache import EncoderCache +from freetoken.mm.processors.deepseek_v41 import DeepseekV41MMProcessor +from freetoken.models.deepseek_v41.config import VisionConfig +from freetoken.scheduler.mm import mm_rows_after, plan_mm_batch +from freetoken.utils.torch_utils import torch_dtype + +from .common import DIM, VOCAB, requires_cuda, tiny_text_config, write_tiny_checkpoint +from .reference.vision import Aligner, ViT + +VC = asdict(VisionConfig(num_hidden_layers=2, hidden_size=32, num_attention_heads=4, intermediate_size=48, + patch_size=2, min_pixels=0, max_image_tokens=64)) + + +def _args(): + return SimpleNamespace(vision_n_layers=2, vision_dim=32, vision_n_heads=4, vision_inter_dim=48, + vision_patch_size=2, vision_downsample_ratio=3, vision_rope_theta=10000.0, dim=DIM) + + +def _checkpoint(path, quantized=False): + tensors = write_tiny_checkpoint(str(path), seed=11, quantized=quantized) + torch.manual_seed(12) + with torch_dtype(torch.bfloat16): + vit, aligner = ViT(_args()), Aligner(_args()) + # Replace the unused dummy vision weight from the text-only fixture. + tensors = {k: v for k, v in tensors.items() if not k.startswith("vision.")} + for prefix, module in (("vision", vit), ("aligner", aligner)): + tensors.update({f"{prefix}.{k}": v.detach().bfloat16().contiguous() for k, v in module.state_dict().items()}) + tensors.update({k: torch.randn(DIM, dtype=torch.bfloat16) for k in ("image_start", "image_end", "image_newline")}) + save_file(tensors, path / 'model-00001-of-00001.safetensors') + (path / 'model.safetensors.index.json').write_text(json.dumps({'weight_map': {k: 'model-00001-of-00001.safetensors' for k in tensors}})) + config = json.loads((path / 'config.json').read_text()) + config.update(vision_config=VC, image_token_id=VOCAB - 1) + (path / 'config.json').write_text(json.dumps(config)) + return tensors + + +def _processor(): + cfg = SimpleNamespace(vision_config=SimpleNamespace(**VC), image_token_id=VOCAB - 1) + return DeepseekV41MMProcessor(cfg, '.', MultimodalConfig()) + + +def _image(color): + buf = io.BytesIO() + Image.new('RGB', (14, 10), color).save(buf, format='PNG') + return buf.getvalue() + + +def _inputs(length=151): + ids = [3] * 17 + [VOCAB - 1] + [4] * (length - 19) + [VOCAB - 1] + return _processor().apply(torch.tensor(ids, dtype=torch.int32), [_image('red'), _image('blue')]) + + +def _prefill(eng, reqs): + if not hasattr(eng, 'encoder_cache'): + eng.encoder_cache = EncoderCache() + eng.dtype = torch.bfloat16 + eng._mm_registered = set() + for req in reqs: + if id(req) not in eng._mm_registered: + for item in req.mm_items or (): # what the scheduler claims at admission + eng.encoder_cache.register(item.hash, req.uid, mm_rows_after(item, req.cached_len)) + eng._mm_registered.add(id(req)) + batch = Batch(reqs=reqs, phase='prefill') + batch.padded_reqs = reqs + batch.input_ids = torch.cat([r.input_ids[r.cached_len:r.device_len] for r in reqs]).cuda() + batch.positions = torch.cat([torch.arange(r.cached_len, r.device_len) for r in reqs]).cuda() + jobs, plan, rows, _ = plan_mm_batch(reqs, eng.encoder_cache) + batch.mm_encoder_jobs, batch.mm_gather_plan = jobs, plan + if plan: + batch.mm_rows = torch.tensor(rows, dtype=torch.long, device=eng.device) + Engine._run_mm_encoder(eng, batch) + eng._bind() + eng.backend.prepare_metadata(batch) + with eng.ctx.forward_batch(batch), eng.model.forward_host_ctx(batch, False): + return eng.model.forward().float() + + +def test_processor_patch_order_and_all_span_rows_are_embeddings(): + pixels = np.arange(10 * 14 * 3, dtype=np.uint8).reshape(10, 14, 3) + item = _processor().process([Image.fromarray(pixels)])[0] + assert item.grid_thw == [1, 5, 7] + expected = ((torch.tensor(pixels[:2, :2]).permute(2, 0, 1).float() / 255 - .5) / .5).bfloat16() + assert torch.equal(item.feature[0], expected) + assert len(_processor().prompt_replacement(item).full) == 10 + result = _inputs() + assert [i.offsets for i in result.mm_items] == [[[17, 27]], [[159, 169]]] + assert result.mrope_positions is None + assert result.mm_items[0].hash != result.mm_items[1].hash + + +def test_text_only_skips_every_visual_parameter(tmp_path): + from freetoken.models.deepseek_v41.weight import iter_weights + from freetoken.mm.processor import get_mm_processor + from .test_engine_config import _engine_config + _checkpoint(tmp_path) + enabled = _engine_config(str(tmp_path), moe_strategy='offload') + assert enabled.model_config.is_multimodal + disabled = MultimodalConfig(disabled_encoders=frozenset({'vision', 'audio'})) + text = _engine_config(str(tmp_path), moe_strategy='offload', mm=disabled) + assert not text.model_config.is_multimodal and not text.active_encoders + assert get_mm_processor(str(tmp_path), disabled) is None + loaded = dict(iter_weights(str(tmp_path), 'cpu', include_moe_experts=False, include_vision=False)) + assert not any(k.startswith('visual.') for k in loaded) + assert get_mm_processor(str(tmp_path)).placeholder == [VOCAB - 1] + + +@requires_cuda +@pytest.mark.parametrize('placement', ['gpu', 'host']) +def test_vision_encoder_matches_official_with_padded_grid(tmp_path, placement): + from .harness import TinyEngine + tensors = _checkpoint(tmp_path) + eng = TinyEngine(str(tmp_path)) + eng.model.place_encoder_weights(placement) + item = _processor().process([Image.new('RGB', (14, 10), 'red')])[0] + with torch_dtype(torch.bfloat16), torch.device('cuda'): + vit, aligner = ViT(_args()), Aligner(_args()) + for prefix, module in (('vision', vit), ('aligner', aligner)): + module.load_state_dict({k.removeprefix(prefix + '.'): v.cuda() for k, v in tensors.items() if k.startswith(prefix + '.')}) + with torch.inference_mode(), torch.device('cuda'): + x = aligner(vit(item.feature.cuda(), 5, 7), 5, 7) + got = eng.model.encode(item) + torch.testing.assert_close(got[1:-1].view(2, 4, DIM)[:, :3].reshape(-1, DIM), x, atol=1e-3, rtol=1e-3) + assert torch.equal(got[0], tensors['image_start'].cuda()) + assert torch.equal(got[-1], tensors['image_end'].cuda()) + assert torch.equal(got[[4, 8]], tensors['image_newline'].cuda().expand(2, -1)) + eng.model.place_encoder_weights('gpu') + + +@requires_cuda +@pytest.mark.parametrize('quantized', [False, True]) +def test_image_logits_and_decode_match_reference(tmp_path, quantized): + from .harness import TinyEngine + from .test_reference_parity import ATOL, QUANT_TOL, Reference, _compare, _engram_table_for + tol = QUANT_TOL if quantized else ATOL + tensors = _checkpoint(tmp_path, quantized) + text = tiny_text_config(moe_intermediate_size=256) if quantized else tiny_text_config() + ref = Reference(tensors, text, max_seq_len=512, max_batch_size=2, quantized=quantized, vision_config=VC) + eng = TinyEngine(str(tmp_path), max_seq_len=512, swa_decoder_replay='exact', quantized=quantized) + _engram_table_for(eng, tensors, text) + result = _inputs() + ids = result.input_ids.clone() + types = torch.full_like(ids, -1) + images = [] + for item in result.mm_items: + lo, hi = item.offsets[0] + ids[lo:hi] = VOCAB - 1 + typ = torch.tensor([0] + ([1] * 3 + [2]) * 2 + [3]) + types[lo:hi] = typ + images.append(SimpleNamespace(start=lo, patches=item.feature, n_vit_h=5, n_vit_w=7, types=typ)) + with torch.device('cuda'): + want = ref.model(ids.cuda()[None], images=[images], token_types=types.cuda()[None])[1] + req = eng.new_request(0, result.input_ids.tolist()) + req.mm_items = result.mm_items + got = _prefill(eng, [req]) + _compare('image prefill', got, want, tol=tol) + eng.finish_prefill([req]) + for step in range(3): + token = want.argmax(-1) + want = ref.decode(token, len(ids) + step) + got = eng.decode([req], token.tolist()) + _compare(f'image decode {step}', got, want, tol=tol) + + +@requires_cuda +@pytest.mark.parametrize('mode', ['exact', 'bounded']) +def test_images_survive_chunking_and_prefix_replay(tmp_path, mode): + from .harness import TinyEngine + from .test_reference_parity import _compare, _engram_table_for + tensors = _checkpoint(tmp_path) + eng = TinyEngine(str(tmp_path), max_seq_len=2048, swa_decoder_replay=mode) + _engram_table_for(eng, tensors, tiny_text_config()) + # bounded mode caps a hit a window before the prompt end: trailing text keeps that cap at 128 + tail = 150 if mode == 'bounded' else 0 + def inputs(): + ids = [3] * 80 + [VOCAB - 1] + [4] * 37 + [VOCAB - 1] + [4] * 31 + [VOCAB - 1] + [4] * tail + return _processor().apply(torch.tensor(ids, dtype=torch.int32), [_image(c) for c in ('red', 'blue', 'green')]) + whole = inputs() + assert [i.offsets for i in whole.mm_items] == [[[80, 90]], [[127, 137]], [[168, 178]]] + a = eng.new_request(0, whole.input_ids.tolist()) + a.mm_items = whole.mm_items + want = _prefill(eng, [a]) + chunk = inputs() + b = eng.new_request(1, chunk.input_ids.tolist()) + b.mm_items = chunk.mm_items + b.device_len = 128 + _prefill(eng, [b]) + b.cached_len, b.device_len = 128, len(chunk.input_ids) + got = _prefill(eng, [b]) + _compare('chunked image prefill', got, want) + prefix = inputs() + # The second image straddles the chunk / prefix boundary at 128 (in bounded mode the first chunk + # also ends before the prompt's last window and runs no decoder row). + hit = eng.new_request_on_prefix(1, a, prefix.input_ids.tolist()) + assert hit.cached_len == 128 + hit.mm_items = prefix.mm_items + _compare('image prefix hit', _prefill(eng, [hit]), want) 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 ac598e65df..2b47748f53 100644 --- a/tests/models/qwen4_exp/test_skeleton.py +++ b/tests/models/qwen4_exp/test_skeleton.py @@ -16,7 +16,7 @@ import torch import torch.nn.functional as F -from freetoken.layers import BaseOP, LinearReplicated +from freetoken.layers import BaseOP, LinearReplicated, iter_moe_layers from freetoken.models.config import ModelConfig from freetoken.models.qwen4_exp.config import parse_config from freetoken.models.qwen4_exp.hc import GatedResidual @@ -30,12 +30,17 @@ def _config(num_layers: int = 4) -> ModelConfig: def _fill(op, gen: torch.Generator, scale: float = 0.05) -> None: - """Random floats / zeroed ints for every state-dict tensor of an op tree.""" + """Random floats / zeroed ints for every state-dict tensor of an op tree, and random expert banks for its resident MoE layers.""" for tensor in op.state_dict().values(): if tensor.is_floating_point(): tensor.normal_(0.0, scale, generator=gen) else: tensor.zero_() + for layer in iter_moe_layers(op): + layer._expert_banks = { + role: torch.randn((layer.num_experts, *spec.shape), generator=gen, device=gen.device, dtype=spec.dtype) * scale + for role, spec in layer.quant_method.layout().items() + } def _group_norm(x, weight, eps, groups): @@ -427,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): + streams["slim"].add(torch.cuda.current_stream().cuda_stream) + return real_slim(*args) + + def spy_multi(*args): + streams["demand"].add(torch.cuda.current_stream().cuda_stream) + return real_multi(*args) + + 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/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/models/test_quant_config.py b/tests/models/test_quant_config.py index 6ac1b84b17..bfcfb87664 100644 --- a/tests/models/test_quant_config.py +++ b/tests/models/test_quant_config.py @@ -1,6 +1,7 @@ """What each checkpoint's quant config resolves to: the dialect, the method behind every layer, and agreement with the stored tensors. -Rows need the checkpoint's config.json locally (/mnt/nvme/models or the HF cache); absent ones skip. +Rows need the checkpoint's config.json locally (/mnt/nvme/models or the HF cache); absent ones skip. The +stored-tensor scan also reads $FREETOKEN_TEST_MODELS_ROOT. """ from __future__ import annotations @@ -47,6 +48,8 @@ from freetoken.utils.torch_utils import torch_dtype MODELS = "/mnt/nvme/models" +# the stored-tensor scan takes every checkpoint under this root (the case table above matches by name) +SCAN_ROOT = os.environ.get("FREETOKEN_TEST_MODELS_ROOT", MODELS) HF_CACHE = os.path.expanduser("~/.cache/huggingface/hub") BF16, FP8B, FP8T, MXFP8, NVFP4 = UnquantizedLinearMethod, Fp8BlockLinearMethod, Fp8TensorLinearMethod, Mxfp8LinearMethod, Nvfp4LinearMethod @@ -121,14 +124,6 @@ class Case: check: Callable[[Any, dict], None] | None = None # family extras beyond the method table -def _installed(module: str) -> bool: - try: - __import__(module) - return True - except ImportError: - return False - - def _qwen35_fp8(model, idx): assert model.model.layers.op_list[idx["lin"]].linear_attn._split_in_proj @@ -162,13 +157,18 @@ def _dsv4(model, idx): assert "model.head.weight" in keys and not any(".ffn.experts." in k for k in keys) +def _banked_experts(model, idx): + experts = op_at(model, f"model.layers.{idx['lin']}.mlp.experts") + assert not experts.state_dict() + + def _gpt_oss(cls): def check(model, idx): experts = op_at(model, "model.layers.0.mlp.experts") assert type(experts) is cls assert experts.interleaved and experts.alpha == pytest.approx(1.702) and experts.limit == 7.0 if cls is GptOssMoELayer: - assert experts.gate_up_proj_blocks.dtype is torch.uint8 and experts.gate_up_proj_blocks.dim() == 4 + assert not experts.state_dict() return check @@ -198,9 +198,8 @@ def _gemma_tied_head(model, idx): "model.layers.{lin}.mlp.shared_expert.down_proj": NVFP4, "model.layers.{lin}.mlp.experts": Nvfp4MoEMethod, "lm_head": NVFP4, }, check=_qwen35_nvfp4), - # resident NVFP4 experts have no kernel without vLLM's Marlin - Case("Qwen3.6-35B-A3B-NVFP4", ModelOptConfig, {}, strategy="fused", raises=KernelSelectionError, - skip_if=lambda: not _installed("flashinfer") or _installed("vllm")), + # resident NVFP4 experts load through the expert-bank loader, not the state dict + Case("Qwen3.6-35B-A3B-NVFP4", ModelOptConfig, {"model.layers.{lin}.mlp.experts": (Nvfp4MoEMethod, "triton")}, strategy="fused", check=_banked_experts), Case("Qwen3.8-27B-NVFP4", ModelOptConfig, { "model.layers.{lin}.linear_attn.in_proj_qkvz": FP8T, "model.layers.{full}.self_attn.o_proj": FP8T, "model.layers.{lin}.mlp.gate_up_proj": NVFP4, "model.layers.{lin}.mlp.down_proj": NVFP4, "lm_head": NVFP4, @@ -439,6 +438,25 @@ def test_modelopt_nvfp4_reads_the_activation_quantizer_from_config_groups(extra, assert scheme.kind is QuantKind.NVFP4 and scheme.has("input_scale") is has_input_scale +def test_fp8_dialect_reads_the_block_size_and_the_nested_expert_dtype(): + """DeepSeek-V4.1 keeps the ``fp8`` dialect with 32x32 blocks and puts ``expert_dtype`` inside quantization_config.""" + from freetoken.layers.quantization.scheme import fp8_block_size + + q = {"quant_method": "fp8", "activation_scheme": "dynamic", "weight_block_size": [32, 32], "scale_fmt": "ue8m0", "expert_dtype": "fp4"} + quant = QuantConfig.from_hf(SimpleNamespace(quantization_config=q)) + assert type(quant) is Fp8BlockConfig and quant.block == 32 + dense = quant.scheme_for("layers.3.attn.wq_a") + assert dense.kind is QuantKind.FP8_BLOCK and fp8_block_size(dense) == 32 and dense.weight.scale == "e8m0" + assert quant.storage(dense)["weight_scale_inv"].name == "scale" + assert quant.scheme_for("layers.3.ffn.experts.7.w1").kind is QuantKind.MXFP4 + # the V4 128-block export still resolves to the 128 scheme + q128 = {"quant_method": "fp8", "activation_scheme": "dynamic", "weight_block_size": [128, 128], "scale_fmt": "ue8m0"} + assert fp8_block_size(QuantConfig.from_hf(SimpleNamespace(quantization_config=q128)).scheme_for("layers.3.attn.wq_a")) == 128 + for bad in ([64, 64], [32, 128], [128]): + with pytest.raises(NotImplementedError): + QuantConfig.from_hf(SimpleNamespace(quantization_config={**q, "weight_block_size": bad})) + + def test_every_dialect_names_the_tensors_behind_its_schemes(): from freetoken.layers.quantization.registry import dialects @@ -495,7 +513,7 @@ def _weight_map(p: Path) -> dict[str, str] | None: def _candidate_dirs() -> list[Path]: - roots = glob.glob(os.path.join(HF_CACHE, "models--*/snapshots/*/")) + glob.glob(os.path.join(MODELS, "*/")) + roots = glob.glob(os.path.join(HF_CACHE, "models--*/snapshots/*/")) + glob.glob(os.path.join(SCAN_ROOT, "*/")) return [Path(d) for d in roots if (Path(d) / "config.json").exists() and _weight_map(Path(d)) is not None] 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 f496de3142..547c94c60a 100644 --- a/tests/moe/test_cpu_moe.py +++ b/tests/moe/test_cpu_moe.py @@ -391,6 +391,7 @@ def test_cpu_decode_dsfp4_matches_gpu(bs): max_tokens=bs, device=dev, swiglu_limit=swiglu_limit, + act_block=128, ) hidden = torch.randn(bs, H, device=dev, dtype=torch.bfloat16) @@ -585,7 +586,7 @@ def test_cpu_moe_decode_cuda_graph_replay_dsfp4(): ex = CpuMoeExecutor( cache, top_k=top_k, activation="silu", apply_router_weight_on_input=False, num_threads=8, max_tokens=bs, device=dev, - swiglu_limit=limit, + swiglu_limit=limit, act_block=128, ) b = cache.bank_sources layer_banks = [b[n][layer].to(dev) for n in ( diff --git a/tests/moe/test_dsfp4_prequant.py b/tests/moe/test_dsfp4_prequant.py index 6d6b49d395..11519502e2 100644 --- a/tests/moe/test_dsfp4_prequant.py +++ b/tests/moe/test_dsfp4_prequant.py @@ -46,7 +46,7 @@ def test_dsfp4_gpu_prequant_bit_parity(): def run(cpu_roundtrip: bool) -> torch.Tensor: 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, + 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 diff --git a/tests/moe/test_dsfp4_reference_semantics.py b/tests/moe/test_dsfp4_reference_semantics.py new file mode 100644 index 0000000000..2f89cf334d --- /dev/null +++ b/tests/moe/test_dsfp4_reference_semantics.py @@ -0,0 +1,170 @@ +"""ds_fp4 (DeepSeek W4A8) experts against an INDEPENDENT torch transcription over +``fp4_gemm(act_quant(x, block), W)`` at both fp8 activation blocks (128 on V4, 32 on V4.1): + + x_q = fp8_roundtrip(x, block) # act_quant(..., inplace) + gate = bf16(x_q @ deq(W1)^T), up = bf16(x_q @ deq(W3)^T) # fp4_gemm outputs bf16 + h = bf16(silu(min(gate, L)) * clamp(up, -L, L)) + y_r = bf16(w_r * (fp8_roundtrip(h, block) @ deq(W2)^T)) # routing weight in the down epilogue + y = sum_r y_r + +The routing weight scales the down output (the placement main serves V4 with); the reference +``Expert.forward`` scales the intermediate before its fp8 quant instead. The GEMV (decode), the +grouped GEMM (prefill) and the CPU executor all follow this transcription. +""" + +from __future__ import annotations + +import pytest +import torch + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") + +E, H, I, TOP_K, LIMIT = 16, 512, 256, 4, 7.0 +E2M1 = torch.tensor([0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0]) + + +def _banks(device, seed=0): + g = torch.Generator(device="cpu").manual_seed(seed) + + def u8(*shape, low=0, high=256): + return torch.randint(low, high, shape, dtype=torch.uint8, generator=g).to(device) + + return (u8(E, 2 * I, H // 2), u8(E, 2 * I, H // 32, low=120, high=130), u8(E, H, I // 2), u8(E, H, I // 32, low=120, high=130)) + + +def _dequant(packed: torch.Tensor, scale: torch.Tensor) -> torch.Tensor: + """[E, N, K//2] e2m1 pairs (even channel low nibble) + [E, N, K//32] e8m0 -> [E, N, K] fp32.""" + lut = E2M1.to(packed.device) + lo, hi = lut[(packed & 0xF).long()], lut[(packed >> 4).long()] + vals = torch.stack([lo, hi], dim=-1).flatten(-2) # [E, N, K] + s = torch.exp2(scale.view(torch.uint8).float() - 127.0).repeat_interleave(32, dim=-1) + return vals * s + + +def _fp8_roundtrip(x: torch.Tensor, block: int) -> torch.Tensor: + g = x.float().unflatten(-1, (-1, block)) + s = torch.exp2(torch.ceil(torch.log2(g.abs().amax(-1).clamp_min(1e-4) / 448.0))) + q = (g / s.unsqueeze(-1)).clamp(-448.0, 448.0).to(torch.float8_e4m3fn).float() + return (q * s.unsqueeze(-1)).flatten(-2).to(torch.bfloat16) + + +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).""" + 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) + for t in range(T): + for r in range(TOP_K): + e = int(slots[t, r]) + gu = (xq[t] @ W13[e].T).to(torch.bfloat16).float() + 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 + + +def _inputs(T, device, seed=1): + g = torch.Generator(device="cpu").manual_seed(seed) + x = (torch.randn(T, H, dtype=torch.bfloat16, generator=g) * 0.5).to(device) + slots = torch.stack([torch.randperm(E, generator=g)[:TOP_K] for _ in range(T)]).to(device=device, dtype=torch.int32).contiguous() + w = torch.rand(T, TOP_K, generator=g) + w = (1.5 * w / w.sum(-1, keepdim=True)).float().to(device).contiguous() # route_scale-style weights > 1 too + return x, slots, w + + +@pytest.mark.parametrize("block", [32, 128]) +def test_gemv_path_matches_the_transcription(block): + from freetoken.moe.fused_ds_fp4 import routed_experts_fp4 + + 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) + + +@pytest.mark.parametrize("block", [32, 128]) +def test_cpu_executor_matches_the_transcription(block): + """The ``ds_fp4`` CPU executor (cpu / hybrid decode) implements the same contract as the GPU path.""" + from types import SimpleNamespace + + from freetoken.kernel.pinned import alloc_pinned_tensor + from freetoken.moe.cpu_executor import CpuMoeExecutor + + gu_p, gu_s, dn_p, dn_s = _banks("cuda") + pinned = {} + for name, t in (("gate_up_packed", gu_p), ("gate_up_scale", gu_s), ("down_packed", dn_p), ("down_scale", dn_s)): + buf = alloc_pinned_tensor(*t.shape, dtype=torch.uint8) + buf.copy_(t.cpu()) + pinned[name] = buf + cache = SimpleNamespace(quant_format="ds_fp4", bank_sources={k: [v] for k, v in pinned.items()}, num_layers=1, num_experts=E) + dev = torch.device("cuda", 0) + x, slots, w = _inputs(2, dev) + ex = CpuMoeExecutor(cache, top_k=TOP_K, activation="silu", apply_router_weight_on_input=False, num_threads=4, max_tokens=2, device=dev, swiglu_limit=LIMIT, act_block=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) + + +def test_grouped_prefill_path_matches_the_transcription(): + from freetoken.moe import fused_ds_fp4 + + 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) + + +@pytest.mark.parametrize("block", [32, 128]) +def test_inactive_routes_contribute_zero_without_reading_their_slots(block): + """Hybrid decode hands the GPU kernel slot -1 (weight 0) for the routes the CPU computes. The kernel + must produce exactly zero for them without touching any slot: every slot the active routes do not + use is filled with 0xFF (an e8m0 scale of NaN), and the result still equals the reference over the + active routes and the CPU executor's output for the same split.""" + from types import SimpleNamespace + + from freetoken.kernel.pinned import alloc_pinned_tensor + from freetoken.moe.cpu_executor import CpuMoeExecutor + from freetoken.moe.fused_ds_fp4 import routed_experts_fp4 + + gu_p, gu_s, dn_p, dn_s = (t.clone() for t in _banks("cuda")) + x, slots, w = _inputs(2, "cuda") + # routes: token 0 keeps routes 0..1 on the GPU, token 1 keeps route 3; the rest go to the CPU + on_gpu = torch.zeros_like(slots, dtype=torch.bool) + on_gpu[0, :2] = True + on_gpu[1, 3] = True + gpu_slots = torch.where(on_gpu, slots, -1) + gpu_w = torch.where(on_gpu, w, 0.0) + used = set(slots[on_gpu].tolist()) + poison = [e for e in range(E) if e not in used] + for t in (gu_p, gu_s, dn_p, dn_s): + t[poison] = 0xFF + got = routed_experts_fp4(x, gpu_slots, gpu_w, gu_p, gu_s, dn_p, dn_s, LIMIT, act_block=block) + assert torch.isfinite(got).all() + clean = _banks("cuda") + # the poisoned-bank result equals the clean-bank result with the same holes (neither reads them), + # 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) + # 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])): + buf = alloc_pinned_tensor(*t.shape, dtype=torch.uint8) + buf.copy_(t.cpu()) + pinned[name] = buf + cache = SimpleNamespace(quant_format="ds_fp4", bank_sources={k: [v] for k, v in pinned.items()}, num_layers=1, num_experts=E) + ex = CpuMoeExecutor(cache, top_k=TOP_K, activation="silu", apply_router_weight_on_input=False, num_threads=4, max_tokens=2, device=torch.device("cuda", 0), swiglu_limit=LIMIT, act_block=block) + cpu_ids = torch.where(on_gpu, -1, slots) + 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) 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_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 0546fbcd6f..7bdbc95681 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 @@ -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_swigluoai_matches_dequant_reference(): """MiniMax-M3's swigluoai routed experts through the Triton prefill grouped GEMM diff --git a/tests/moe/test_offload.py b/tests/moe/test_offload.py index cd5f5dc773..19e25eafe3 100644 --- a/tests/moe/test_offload.py +++ b/tests/moe/test_offload.py @@ -277,6 +277,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 @@ -343,7 +387,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)) @@ -383,6 +427,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 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..6bfa84a5ed --- /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, "bank_views", lambda: ()) + 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..8b7b9d802c --- /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": None, "is_prefill": False} + 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": None, "is_prefill": False} + + +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_cache_rebuild.py b/tests/scheduler/test_cache_rebuild.py index 0c0580708f..5e280ae7fd 100644 --- a/tests/scheduler/test_cache_rebuild.py +++ b/tests/scheduler/test_cache_rebuild.py @@ -168,3 +168,41 @@ def test_rebuild_cache_refreshes_prefill_budget(monkeypatch): cache_manager.prefill_chunk_budget = 1000 # the (stubbed) engine rebuild shrank the pool Scheduler.rebuild_cache(sched, num_pages=16) assert sched.prefill_budget == 1000 # tracks the shrunk cap, not the stale 5000 + + +def _pool_engine(calls, *, moe_slots=512, num_pages=4096, mamba_slots=9): + """An engine carrying only the three pools _resize_pools touches, each logging its resize.""" + from types import SimpleNamespace + + engine = SimpleNamespace(num_pages=num_pages) + engine.moe_offload_cache = SimpleNamespace( + cache_size=moe_slots, rebuild=lambda n: calls.append(("moe", n))) + engine.linear_state_pool = SimpleNamespace( + num_slots=mamba_slots, rebuild=lambda n: calls.append(("mamba", n))) + engine._resize_kv_pool = lambda config, pages, swa: calls.append(("kv", pages, swa)) + return engine + + +def test_resize_pools_frees_shrinking_pools_before_growing_ones(): + """#5: MoE 512 -> 1024 slots with KV 4096 -> 1024 pages OOMed on a full card because the MoE + pool was reallocated while the KV pool still held its old pages.""" + from freetoken.engine.engine import Engine + + calls = [] + Engine._resize_pools(_pool_engine(calls), None, moe_cache_size=1024, num_pages=1024, + num_mamba_slots=4, num_swa_pages=None) + assert calls == [("kv", 1024, None), ("mamba", 5), ("moe", 1024)] + + calls.clear() # the rollback of that request: MoE shrinks back first, then KV and mamba grow + Engine._resize_pools(_pool_engine(calls, moe_slots=1024, num_pages=1024, mamba_slots=5), None, + moe_cache_size=512, num_pages=4096, num_mamba_slots=8, num_swa_pages=None) + assert calls == [("moe", 512), ("kv", 4096, None), ("mamba", 9)] + + +def test_resize_pools_keeps_request_order_among_growers_and_window_only_uses_current_pages(): + from freetoken.engine.engine import Engine + + calls = [] + Engine._resize_pools(_pool_engine(calls), None, moe_cache_size=600, num_pages=None, + num_mamba_slots=None, num_swa_pages=7) + assert calls == [("moe", 600), ("kv", 4096, 7)] 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_dsv4_generic_manager.py b/tests/scheduler/test_dsv4_generic_manager.py index ad69b89ec5..1bca10bc8a 100644 --- a/tests/scheduler/test_dsv4_generic_manager.py +++ b/tests/scheduler/test_dsv4_generic_manager.py @@ -14,8 +14,8 @@ import torch from freetoken.core import Req, SamplingParams -from freetoken.kvcache.dsv4_cost_model import dsv4_pool_sizes -from freetoken.kvcache.dsv4_paged_pool import DSV4PagedKVCache +from freetoken.kvcache.dsv4.v4_cost_model import dsv4_pool_sizes +from freetoken.kvcache.dsv4.v4_pool import DSV4PagedKVCache from freetoken.models.deepseek_v4.args import DeepseekV4Args from freetoken.scheduler.cache import CacheManager @@ -191,6 +191,33 @@ def _pending(uid, n): cm.check_integrity() +@pytest.mark.parametrize("replay, new_tokens, cached", [(0, 2, 256), (P, 2, 128), (P, 200, 256)]) +def test_admission_caps_the_hit_by_the_pools_prefix_replay(replay, new_tokens, cached): + """A pool whose model recomputes the prompt's tail (DeepSeek-V4.1's bounded decoder replay) holds + back ``prefix_replay_tokens`` at admission: the match ends at or before ``L - replay`` (page-aligned + by the radix), a hit already that far from the prompt end is untouched, and the request's + ``prompt_len`` is the whole prompt.""" + from freetoken.scheduler.decode import DecodeManager + from freetoken.scheduler.prefill import PrefillManager + from freetoken.scheduler.table import TableManager + from freetoken.scheduler.utils import PendingReq + + cm, pool, pt = _stack(num_pages=48) + pool.prefix_replay_tokens = replay + prompt = torch.arange(1, 301, dtype=torch.int32) + _lifecycle(cm, _req(MRR - 1, prompt, n_decode=310), total_len=310) + pm = PrefillManager(cm, TableManager(max_running_reqs=MRR, page_table=pt), DecodeManager(page_size=P)) + tokens = torch.cat([prompt, torch.arange(1000, 1000 + new_tokens, dtype=torch.int32)]) + pm.pending_list = [PendingReq(uid=1, input_ids=tokens, sampling_params=SamplingParams(max_tokens=1))] + batch = pm.schedule_next_batch(1024) + req = batch.reqs[0] + assert req.cached_len == cached and req.extend_len == len(tokens) - cached and req.prompt_len == len(tokens) + cm.allocate_paged(batch.reqs) + req.complete_one() + cm.cache_req(req, finished=True) + cm.check_integrity() + + def test_abort_anywhere_fuzz_conserves_and_isolates(): """admit / extend / finish-or-abort at arbitrary points with shared prefixes + eviction pressure, on a deliberately small window tier. The generic conservation check (free + tree 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/scheduler/test_scheduler_kv_usage.py b/tests/scheduler/test_scheduler_kv_usage.py index 6ab949705e..bf0e825c31 100644 --- a/tests/scheduler/test_scheduler_kv_usage.py +++ b/tests/scheduler/test_scheduler_kv_usage.py @@ -2,6 +2,8 @@ from types import SimpleNamespace +from freetoken.attention import AttnType + def test_kv_usage_pages_excludes_evictable_prefix_cache(): # page_usage now lives on the CacheManagerLike interface (polymorphic vs DSV4); test the @@ -53,10 +55,15 @@ def _fake_engine(swa=True, moe=True, mamba=True): config=SimpleNamespace( page_size=16, cache_type="swa_radix", - model_config=SimpleNamespace(dsv4_args=None, has_swa_attention=swa), + max_running_req=1, + model_config=SimpleNamespace( + dsv4_args=None, has_swa_attention=swa, + kv_cache_group_specs=lambda: [SimpleNamespace(attn_type=AttnType.SWA, is_swa=True, sliding_window=128)] if swa else [], + swa_attention_group=lambda: SimpleNamespace(sliding_window=128), + ), ), kv_cache=SimpleNamespace( - swa_num_tokens=513, unit_bytes=lambda: (1 << 20, 1 << 21) + window_pages=512, swa_num_tokens=513, unit_bytes=lambda: (1 << 20, 1 << 21) ), moe_offload_cache=SimpleNamespace( cache_size=24, num_layers=8, num_experts=16, 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 327bfcd87d..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"): @@ -231,10 +234,7 @@ def test_convert_thinking_replay_in_tool_loop(): assert spec.messages[2]["role"] == "tool" -def test_convert_hoists_and_merges_system_messages(): - # Claude Code interleaves system messages mid-array; strict chat templates - # (e.g. Qwen3.5: "System message must be at the beginning") require ONE system - # message at the front. Merge top-level system + in-array system, hoist to front. +def test_convert_preserves_mid_conversation_system_reminders(): req = AnthropicMessagesRequest.model_validate( { "model": "claude-x", @@ -247,11 +247,69 @@ def test_convert_hoists_and_merges_system_messages(): ], } ) + original = req.model_dump() spec = A.convert_anthropic_to_genspec(req, {}) - assert [m["role"] for m in spec.messages] == ["system", "user", "assistant"] - assert sum(1 for m in spec.messages if m["role"] == "system") == 1 - assert "top-level sys" in spec.messages[0]["content"] - assert "mid-stream sys" in spec.messages[0]["content"] + assert spec.messages == [ + {"role": "system", "content": "top-level sys"}, + {"role": "user", "content": "hello"}, + {"role": "system", "content": "mid-stream sys"}, + {"role": "assistant", "content": "hi"}, + ] + assert req.model_dump() == original + + +def test_convert_merges_only_leading_system_messages(): + req = AnthropicMessagesRequest.model_validate({ + "model": "claude-x", "max_tokens": 64, + "system": [{"type": "text", "text": "top-level sys"}], + "messages": [ + {"role": "system", "content": "leading sys"}, + {"role": "system", "content": [{"type": "text", "text": "second sys"}]}, + {"role": "user", "content": "hello"}, + {"role": "system", "content": [{"type": "text", "text": "turn reminder"}]}, + {"role": "assistant", "content": "hi"}, + {"role": "system", "content": ""}, + {"role": "system", "content": "next reminder"}, + ], + }) + original = req.model_dump() + spec = A.convert_anthropic_to_genspec(req, {}) + assert spec.messages == [ + {"role": "system", "content": "top-level sys\n\nleading sys\n\nsecond sys"}, + {"role": "user", "content": "hello"}, + {"role": "system", "content": "turn reminder"}, + {"role": "assistant", "content": "hi"}, + {"role": "system", "content": "next reminder"}, + ] + assert req.model_dump() == original + + +@pytest.mark.parametrize("system", [None, "stable instructions"]) +@pytest.mark.parametrize("count_tokens", [False, True]) +def test_appending_system_reminder_preserves_converted_prefix(system, count_tokens): + from freetoken.server.anthropic_models import AnthropicCountTokensRequest + + data = { + "model": "claude-x", + "messages": [ + {"role": "user", "content": "inspect the source"}, + {"role": "assistant", "content": "I found the function"}, + ], + "tools": [{"name": "read", "input_schema": {"type": "object", "properties": {}}}], + } + if system is not None: + data["system"] = system + if not count_tokens: + data["max_tokens"] = 64 + request_type = AnthropicCountTokensRequest if count_tokens else AnthropicMessagesRequest + first, first_tools, _, _ = A.convert_anthropic_prompt(request_type.model_validate(data)) + data["messages"].append({"role": "system", "content": "1000 tokens left"}) + second, second_tools, _, _ = A.convert_anthropic_prompt(request_type.model_validate(data)) + assert second[:-1] == first + assert second[-1] == { + "role": "system", "content": "1000 tokens left" + } + assert second_tools == first_tools # --------------------------------------------------------------------------- # @@ -373,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 @@ -664,6 +723,7 @@ def test_count_tokens_route(): assert msg.text == spec.messages assert msg.tools == spec.template_tools assert msg.chat_template_kwargs == spec.chat_template_kwargs + assert msg.inline_system_policy == spec.inline_system_policy == "auto" def test_count_tokens_empty_messages_400(): @@ -673,6 +733,28 @@ def test_count_tokens_empty_messages_400(): assert r.json()["error"]["type"] == "invalid_request_error" +@pytest.mark.parametrize("policy", ["auto", "preserve", "fold"]) +def test_inline_system_override_reaches_counting_and_generation(policy): + manager = _FakeTokenizeManager() + fake = FakeState([("ok", True, 3, 1)]) + fake._count_manager = manager + fake.config.anthropic_inline_system = policy + submitted = [] + + async def record(msg): + submitted.append(msg) + + fake.send_one = record + client = _client(fake) + body = {**_COUNT_BODY, "messages": [ + {"role": "user", "content": "question"}, {"role": "system", "content": "instruction"}]} + assert client.post("/v1/messages/count_tokens", json=body).status_code == 200 + assert client.post("/v1/messages", json={**body, "max_tokens": 32}).status_code == 200 + assert manager.msgs[0].inline_system_policy == submitted[0].inline_system_policy == policy + assert manager.msgs[0].text == submitted[0].text + assert submitted[0].text[0]["content"] == body["system"] + + def test_count_tokens_works_while_not_serving(): # Counting never touches the engine, so it must not 503 during load/rebuild. client, _ = _count_client(_FakeTokenizeManager(), maintenance_state="rebuilding") @@ -719,6 +801,15 @@ def test_count_tokens_template_render_error_400(): assert r.json()["error"]["type"] == "invalid_request_error" +def test_count_tokens_inline_system_placement_error_is_400(): + from freetoken.tokenizer.inline_system import InlineSystemError + + client, _ = _count_client(_RaisingManager(InlineSystemError("missing tool results"))) + r = client.post("/v1/messages/count_tokens", json=_COUNT_BODY) + assert r.status_code == 400 + assert r.json()["error"]["type"] == "invalid_request_error" + + def test_count_tokens_excluded_from_request_ring(): # count_tokens never enters generation accounting; its latency (incl. first-touch tokenizer # load) must not land in the /v1/requests ring or the /v1/stats p95. @@ -976,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_effort_dialect.py b/tests/server/test_effort_dialect.py index 6a68a4ece3..a25de8e728 100644 --- a/tests/server/test_effort_dialect.py +++ b/tests/server/test_effort_dialect.py @@ -159,6 +159,22 @@ def test_empty_effort_is_treated_as_absent(): assert state.sent.chat_template_kwargs == {} +def test_numeric_effort_rides_template_kwargs_not_the_wire_field(): + # The wire field takes named tiers only (as vLLM and SGLang type it); a numeric budget goes + # through chat_template_kwargs, where the checkpoint's encoder validates it at render time. + from pydantic import ValidationError + import pytest + + for effort in (37, True, 25.5): + with pytest.raises(ValidationError): + chat_request(reasoning_effort=effort) + ctk = {"enable_thinking": True, "reasoning_effort": 37} + state = FakeState(reasoning_parser="qwen3") + response = run(handle_chat_completion(chat_request(chat_template_kwargs=ctk), None, state, {})) + assert not isinstance(response, JSONResponse) + assert state.sent.chat_template_kwargs == ctk + + def test_foreign_thinking_shapes_stay_ignored(): # extra="allow" swallowed any thinking shape before the field existed; # a bare string, a bool, or a typeless dict must keep working unchanged. diff --git a/tests/server/test_function_call_parser.py b/tests/server/test_function_call_parser.py index 34c1c04d10..8af5bc88b6 100644 --- a/tests/server/test_function_call_parser.py +++ b/tests/server/test_function_call_parser.py @@ -337,9 +337,177 @@ def test_dsv32_streaming_multi_param_args_prefix_stable(): } +V41_BLOCK = ( + "<｜DSML｜ calls>\n" + '<｜DSML｜ invoke name="glob">\n' + '<｜DSML｜ parameter name="pattern" string="true">*.py\n' + '<｜DSML｜ parameter name="path" string="true">/src\n' + "\n" + "" +) + + +def test_dsv41_dsml_tags_lead_with_a_space(): + """DeepSeek-V4.1 leads every DSML tag name with a space (its encoder's ``<｜DSML｜ calls>``): + the V3.2 detector sees no tool call there, the V4.1 one parses it whole and streamed.""" + text = "Let me look.\n\n" + V41_BLOCK + assert not FunctionCallParser(OPENCODE_TOOLS, tool_call_parser="deepseekv32").has_tool_call(text) + parser = FunctionCallParser(OPENCODE_TOOLS, tool_call_parser="deepseekv41") + assert parser.has_tool_call(text) + result = parser.parse_non_stream(text) + assert result.normal_text == "Let me look." + assert [(c.name, json.loads(c.parameters)) for c in result.calls] == [("glob", {"pattern": "*.py", "path": "/src"})] + + streamed = FunctionCallParser(OPENCODE_TOOLS, tool_call_parser="deepseekv41") + texts, calls = _feed(streamed, [text[i : i + 5] for i in range(0, len(text), 5)]) + assert "".join(texts).strip() + streamed.finish_stream() == "Let me look." + assert [c.name for c in calls if c.name] == ["glob"] + assert json.loads("".join(c.parameters for c in calls if c.name is None)) == {"pattern": "*.py", "path": "/src"} + + def test_streaming_support_flags(): # Every registered detector is incremental-safe (buffered fallback remains as # the escape hatch for future formats, exercised via monkeypatch in # 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 8276163b0c..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, ) @@ -91,6 +102,19 @@ def _last_row() -> dict: return rows[-1] +@pytest.mark.parametrize("pools, expected", [(None, 1), ({"page_size": 64}, 64)]) +def test_stats_reports_backend_page_size(pools, expected): + from freetoken.server.stats import StatsTracker, build_stats + + tracker = StatsTracker() + tracker.kv_total_pages = 130 + config = SimpleNamespace(served_model_name="unit-model", max_seq_len=262144, + page_size=1, served_modalities=(), model_config=SimpleNamespace()) + state = SimpleNamespace(stats=tracker, config=config, cache_pools=pools) + doc = build_stats(state, 0, 0) + assert doc["kv"] == {"used_pages": 0, "total_pages": 130, "page_size": expected} + + def test_non_stream_records_the_request_with_real_token_totals(): request_ring.reset() st = FakeState([_ack(prompt=5, completion=1, out="a"), _ack(completion=2, out="bc", finished=True)]) @@ -193,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_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_message_wire.py b/tests/server/test_message_wire.py index e228667154..fdc095921d 100644 --- a/tests/server/test_message_wire.py +++ b/tests/server/test_message_wire.py @@ -28,6 +28,18 @@ from freetoken.core import SamplingParams +def test_inline_system_policy_tokenizer_roundtrip(): + for policy in (None, "auto", "preserve", "fold"): + msg = TokenizeMsg(uid=1, text=[{"role": "system", "content": "instruction"}], + sampling_params=SamplingParams(), inline_system_policy=policy) + decoded = BaseTokenizerMsg.decoder(BaseTokenizerMsg.encoder(msg)) + assert decoded.inline_system_policy == policy + assert decoded.text == msg.text + legacy = BaseTokenizerMsg.encoder(msg) + legacy.pop("inline_system_policy") + assert BaseTokenizerMsg.decoder(legacy).inline_system_policy is None + + def test_cache_rebuild_msg_roundtrip(): msg = CacheRebuildMsg(request_id="abc", moe_cache_size=8, num_pages=1024, mode="if_idle") out = BaseTokenizerMsg.decoder(BaseTokenizerMsg.encoder(msg)) diff --git a/tests/server/test_mm_input.py b/tests/server/test_mm_input.py index 19d9f86529..a694a50a84 100644 --- a/tests/server/test_mm_input.py +++ b/tests/server/test_mm_input.py @@ -140,3 +140,68 @@ 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"}, + ]) 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..7e6d4ed30e 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,65 @@ 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) diff --git a/tests/server/test_openai_extras.py b/tests/server/test_openai_extras.py new file mode 100644 index 0000000000..115112b9d0 --- /dev/null +++ b/tests/server/test_openai_extras.py @@ -0,0 +1,383 @@ +"""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 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}, + {"presence_penalty": 3.0}, + {"frequency_penalty": -2.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 + + +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_parser_auto_selection.py b/tests/server/test_parser_auto_selection.py index 78b1663e8d..e1eb87982b 100644 --- a/tests/server/test_parser_auto_selection.py +++ b/tests/server/test_parser_auto_selection.py @@ -87,6 +87,13 @@ def test_qwen3_5_is_not_shadowed_by_the_generic_qwen_branch(): assert _inferred("Qwen3MoeForCausalLM")[0] == "qwen25" +def test_deepseek_v41_is_not_shadowed_by_the_v4_branch(): + """V4.1's DSML leads every tag name with a space; the V4 arm (``"v4"`` is a substring of + ``"v41"``) would hand it a tool parser that never sees a call.""" + assert _inferred("DeepseekV41ForCausalLM") == ("deepseekv41", "deepseekv32") + assert _inferred("DeepseekV4ForCausalLM")[0] == "deepseekv32" + + def test_an_explicit_choice_beats_inference(): config = _Config({"architectures": ["DeepseekV4ForCausalLM"], "torch_dtype": "bfloat16"}) with patch("freetoken.utils.cached_load_hf_config", lambda _path: config): 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_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..3b2e1d6473 --- /dev/null +++ b/tests/server/test_sampling_validation.py @@ -0,0 +1,23 @@ +"""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) 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_streaming_model_matrix.py b/tests/server/test_streaming_model_matrix.py index 43922021a2..eb6e3a2ed2 100644 --- a/tests/server/test_streaming_model_matrix.py +++ b/tests/server/test_streaming_model_matrix.py @@ -67,6 +67,11 @@ '<｜DSML｜parameter name="filePath" string="true">/tmp/test_calc.py' "" ), + "deepseekv41": ( + '<｜DSML｜ calls>\n<｜DSML｜ invoke name="read">\n' + '<｜DSML｜ parameter name="filePath" string="true">/tmp/test_calc.py\n' + "\n" + ), "qwen25": '\n{"name": "read", "arguments": {"filePath": "/tmp/test_calc.py"}}\n', "qwen3_coder": ( "/tmp/test_calc.py" @@ -104,6 +109,7 @@ # Substrings that must never leak into user-visible content. MARKUP_MARKERS = { "deepseekv32": ["｜DSML｜"], + "deepseekv41": ["｜DSML｜"], "qwen25": ["", ""], "qwen3_coder": ["", "", "", ""], @@ -121,6 +127,7 @@ # opens the block implicitly (model emits only the closing marker). REASONING_FAMILIES = { "dsv4": ("deepseekv32", "deepseekv32", "", ""), + "dsv4.1": ("deepseekv41", "deepseekv32", "", ""), "qwen3.5": ("qwen3_coder", "qwen3", "", ""), "qwen": ("qwen25", "qwen3", "", ""), "glm4.7": ("glm47", "glm", "", ""), diff --git a/tests/server/test_supervisor.py b/tests/server/test_supervisor.py index 3568d90505..c4772603b0 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 @@ -270,3 +272,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_effort.py b/tests/tokenizer/test_effort.py index e56df5fe68..e0cf7a81f6 100644 --- a/tests/tokenizer/test_effort.py +++ b/tests/tokenizer/test_effort.py @@ -1,6 +1,8 @@ """Unit tests for the reasoning-effort dialect layer (tokenizer/effort.py).""" from __future__ import annotations +import pytest + from freetoken.tokenizer.effort import ( EFFORT_SCALE, EffortProfile, @@ -131,3 +133,37 @@ def render(kwargs, tools): profile = probe_effort_profile(render) assert profile.supported == frozenset({"xhigh", "medium", "low"}) assert profile.consumes_effort + + +def _numeric_render(kwargs, tools): + if not tools and not kwargs.get("enable_thinking"): + return "chat" + effort = kwargs.get("reasoning_effort", "high") + if isinstance(effort, str): + effort = {"low": 50, "high": 75, "max": 100}[effort] + if type(effort) is not int or not 1 <= effort <= 100: + raise ValueError("invalid budget") + return f"Reasoning Effort: {effort}" + + +def test_probe_preserves_validated_numeric_budgets(): + profile = probe_effort_profile(_numeric_render) + assert profile.integer_range == (1, 100) + assert profile.default == "high" + for budget in range(1, 101): + assert quantize_effort(budget, profile) == budget + for bad in (0, 101, -1, True, False, 1.0, 1.5): + with pytest.raises(ValueError, match="integer"): + quantize_effort(bad, profile) + + +def test_numeric_capability_requires_validation_and_the_whole_range(): + assert probe_effort_profile(_dsv4_render).integer_range is None + assert probe_effort_profile(lambda kw, tools: str(kw)).integer_range is None + + def holes(kwargs, tools): + if kwargs.get("reasoning_effort") == 42: + raise ValueError("unsupported") + return _numeric_render(kwargs, tools) + + assert probe_effort_profile(holes).integer_range is None 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 369576a18b..9e441886b2 100644 --- a/tests/tokenizer/test_tokenize.py +++ b/tests/tokenizer/test_tokenize.py @@ -1,6 +1,8 @@ from __future__ import annotations import json +import os +from copy import deepcopy import pytest import torch @@ -8,14 +10,303 @@ from freetoken.core import SamplingParams from freetoken.message import TokenizeMsg from freetoken.tokenizer.tokenize import TokenizeManager, _dsv4_arguments_str +from freetoken.tokenizer.inline_system import InlineSystemError, normalize_inline_system, probe_inline_system + + +def test_fold_keeps_leading_system_and_appends_text_without_mutation(): + messages = [ + {"role": "system", "content": "stable"}, + {"role": "user", "content": [{"type": "image_url", "image_url": {"url": "image"}}, + {"type": "text", "text": "question"}]}, + {"role": "system", "content": "instruction 1"}, + {"role": "system", "content": "instruction 2"}, + ] + original = deepcopy(messages) + result = normalize_inline_system(messages, "fold") + assert result == [messages[0], {"role": "user", "content": [ + messages[1]["content"][0], {"type": "text", "text": "question\n\ninstruction 1\n\ninstruction 2"}]}] + assert messages == original + + +@pytest.mark.parametrize("mode", ["preserve", "fold"]) +def test_inline_system_waits_for_all_parallel_tool_results(mode): + messages = [ + {"role": "system", "content": "stable"}, {"role": "user", "content": "question"}, + {"role": "assistant", "reasoning_content": "reason", "tool_calls": [{"id": "a"}, {"id": "b"}]}, + {"role": "system", "content": "before results"}, + {"role": "tool", "tool_call_id": "a", "content": "result a"}, + {"role": "system", "content": "between results"}, + {"role": "tool", "tool_call_id": "b", "content": "result b"}, + {"role": "system", "content": "after results"}, + ] + original = deepcopy(messages) + result = normalize_inline_system(messages, mode) + assert [m["tool_call_id"] for m in result if m["role"] == "tool"] == ["a", "b"] + assert result[2] == messages[2] + assert [m["role"] for m in result[3:5]] == ["tool", "tool"] + if mode == "fold": + assert len(result) == 5 + assert result[-1]["content"] == "result b\n\nbefore results\n\nbetween results\n\nafter results" + else: + assert [m["content"] for m in result[5:]] == ["before results", "between results", "after results"] + assert messages == original + + +@pytest.mark.parametrize("mode", ["preserve", "fold"]) +def test_inline_system_does_not_hide_missing_tool_results(mode): + messages = [{"role": "assistant", "tool_calls": [{"id": "missing"}]}, + {"role": "system", "content": "instruction"}] + with pytest.raises(InlineSystemError, match="missing tool results"): + normalize_inline_system(messages, mode) + + +def test_fold_after_assistant_prepends_next_user_or_supplies_a_user_turn(): + messages = [{"role": "user", "content": "first"}, {"role": "assistant", "content": "answer"}, + {"role": "system", "content": "instruction"}] + assert normalize_inline_system(messages, "fold")[-1] == {"role": "user", "content": "instruction"} + result = normalize_inline_system(messages + [{"role": "user", "content": "next"}], "fold") + assert result == messages[:2] + [{"role": "user", "content": "instruction\n\nnext"}] + + +def test_preserve_keeps_instruction_role_and_only_folds_without_a_stable_system(): + messages = [{"role": "system", "content": "stable"}, {"role": "user", "content": "question"}, + {"role": "system", "content": "late"}] + assert normalize_inline_system(messages, "preserve") == messages + assert normalize_inline_system(messages[1:], "preserve") == [{"role": "user", "content": "question\n\nlate"}] + + +class ProbeTokenizer: + vocabulary = {text: 10000+i for i, text in enumerate( + ["", "", "", "", "", "", + "", "", "", ""])} + all_special_ids = list(vocabulary.values()) + + def get_added_vocab(self): + return self.vocabulary + + def encode(self, text, **kwargs): + import re + parts = re.split("(" + "|".join(re.escape(x) for x in self.vocabulary) + ")", text) + return [token for part in parts for token in ([self.vocabulary[part]] if part in self.vocabulary else map(ord, part))] + + def decode(self, ids): + inverse = {v: k for k, v in self.vocabulary.items()} + return "".join(inverse[x] if x in inverse else chr(x) for x in ids) + + +@pytest.mark.parametrize("behavior,expected", [("preserve", "preserve"), ("reject", "fold"), + ("drop", "fold"), ("bare", "fold"), ("user", "fold"), ("rewrite", "fold"), + ("after_generation", "fold"), ("drop_reasoning", "fold")]) +def test_renderer_probe_checks_more_than_template_acceptance(behavior, expected): + def render(messages): + inline = any(m["role"] == "system" for m in messages[1:]) + output, late = "", "" + for i, message in enumerate(messages): + role, content = message["role"], message.get("content", "") + if role == "system" and i: + if behavior == "reject": + raise ValueError("system first only") + if behavior == "drop": + continue + if behavior == "bare": + output += content + continue + if behavior == "user": + role = "user" + if behavior == "after_generation": + late += content + continue + if behavior == "rewrite" and inline and i == 0: + content = "changed head" + reason = message.get("reasoning_content") + if reason and not (behavior == "drop_reasoning" and inline): + content = "" + reason + "" + content + output += f"<{role}>{content}" + return output + "" + late + assert probe_inline_system(render, ProbeTokenizer())[0] == expected + + +@pytest.mark.parametrize("position", ["middle", "trailing", "tool result"]) +@pytest.mark.parametrize("metadata", ["Reasoning strength: high. Valid recipients: self, user.", "tool schema " * 1000], + ids=["metadata", "tools"]) +def test_inline_probe_rejects_repeated_system_metadata(position, metadata): + def render(messages): + output = "" + for i, message in enumerate(messages): + role = message["role"] + content = message.get("content", "") + if role == "system" and i: + if i < len(messages)-1: + location = "middle" + else: + location = "tool result" if messages[i-1]["role"] == "tool" else "trailing" + if location == position: + content += metadata + content = message.get("reasoning_content", "") + content + output += f"<{role}>{content}" + return output + "" + + mode, reason = probe_inline_system(render, ProbeTokenizer()) + assert mode == "fold" + assert "overhead" in reason + + +def test_inline_probe_allows_small_role_framing_overhead(): + def render(messages): + return "".join(f'<{m["role"]}>' + ("role: " if m["role"] == "system" else "") + + m.get("reasoning_content", "") + m.get("content", "") + + f'' for m in messages) + "" + + assert probe_inline_system(render, ProbeTokenizer())[0] == "preserve" + + +def test_inline_probe_folds_when_user_insertion_cannot_be_measured(): + def render(messages): + rewrite = any(m["role"] == "user" and m.get("content") == "Echo inline instruction" for m in messages) + return ("changed head" if rewrite else "stable head") + "".join( + f'<{m["role"]}>' + m.get("reasoning_content", "") + m.get("content", "") + + f'' for m in messages) + "" + + mode, reason = probe_inline_system(render, ProbeTokenizer()) + assert mode == "fold" + assert "cannot isolate" in reason + + +def test_inline_probe_cache_is_scoped_to_tools_and_render_kwargs(monkeypatch): + import freetoken.tokenizer.tokenize as module + calls = [] + monkeypatch.setattr(module, "probe_inline_system", lambda *_: (calls.append(1) or ("fold", "test"))) + manager = TokenizeManager(FakeTokenizer()) + for _ in range(3): + assert manager.inline_system_mode(None, {}) == "fold" + manager.inline_system_mode([], {}) + manager.inline_system_mode(None, {"enable_thinking": True}) + assert len(calls) == 3 + + +def test_inline_policy_is_opt_in_and_not_forwarded_to_templates(): + tokenizer = FakeTokenizer() + manager = TokenizeManager(tokenizer) + messages = [{"role": "user", "content": "question"}, {"role": "system", "content": "instruction"}] + msg = TokenizeMsg(uid=1, text=messages, sampling_params=SamplingParams()) + manager.render_prompt(msg) + assert tokenizer.chat_template_messages == messages + msg.inline_system_policy = "fold" + manager.render_prompt(msg) + assert tokenizer.chat_template_messages == [{"role": "user", "content": "question\n\ninstruction"}] + assert "inline_system_policy" not in tokenizer.chat_template_kwargs + + +@pytest.mark.needs_weights +@pytest.mark.parametrize("variable,expected_mode", [ + ("FREETOKEN_DSV4_MODEL", "fold"), + ("FREETOKEN_INLINE_GLM_MODEL", "preserve"), + ("FREETOKEN_INLINE_QWEN_MODEL", "fold"), +]) +def test_native_model_inline_system_tool_history(variable, expected_mode): + from transformers import AutoTokenizer + from freetoken.server.anthropic_api import convert_anthropic_to_genspec + from freetoken.server.anthropic_models import AnthropicMessagesRequest + + path = os.environ.get(variable) + if not path: + pytest.skip(f"set {variable} to a local checkpoint") + manager = TokenizeManager(AutoTokenizer.from_pretrained(path, local_files_only=True)) + data = { + "model": "local", "max_tokens": 64, "system": "Stable project instructions.", + "thinking": {"type": "enabled"}, + "tools": [{"name": "read", "input_schema": {"type": "object", "properties": {}}}], + "messages": [ + {"role": "user", "content": "Inspect source:\n" + "value = 1\n" * 512}, + {"role": "system", "content": "Keep all existing tests."}, + ], + } + ids, prompts = [], [] + for turn in (1, 2): + data["messages"].extend([ + {"role": "assistant", "content": [ + {"type": "thinking", "thinking": f"Retained thought {turn}.", "signature": ""}, + {"type": "tool_use", "id": f"read_{turn}", "name": "read", "input": {}}, + ]}, + {"role": "system", "content": f"During tool call {turn}."}, + {"role": "user", "content": [ + {"type": "tool_result", "tool_use_id": f"read_{turn}", "content": f"Source result {turn}."}, + ]}, + {"role": "system", "content": f"Remaining budget: {1000-turn}."}, + ]) + req = AnthropicMessagesRequest.model_validate(data) + spec = convert_anthropic_to_genspec(req, {}) + assert manager.inline_system_mode(spec.template_tools, spec.chat_template_kwargs) == expected_mode + original = deepcopy(spec.messages) + msg = TokenizeMsg(uid=turn, text=spec.messages, tools=spec.template_tools, + sampling_params=spec.sampling_params, chat_template_kwargs=spec.chat_template_kwargs, + inline_system_policy=spec.inline_system_policy) + prompts.append(manager.render_prompt(msg)) + ids.append(manager.tokenize([msg])[0].input_ids.tolist()) + assert spec.messages == original + assert prompts[-1].count("Keep all existing tests.") == 1 + assert "Retained thought 1." in prompts[-1] + for previous in range(1, turn+1): + assert prompts[-1].count(f"During tool call {previous}.") == 1 + assert prompts[-1].count(f"Remaining budget: {1000-previous}.") == 1 + common = next((i for i, pair in enumerate(zip(*ids)) if pair[0] != pair[1]), min(map(len, ids))) + assert common >= len(ids[0]) - 64 + assert common > 2048 + + +@pytest.mark.needs_weights +@pytest.mark.parametrize("tool_count", [0, 20]) +def test_native_muse_inline_system_does_not_repeat_environment(tool_count): + from transformers import AutoTokenizer + + path = os.environ.get("FREETOKEN_MUSE_MODEL") + if not path: + pytest.skip("set FREETOKEN_MUSE_MODEL to a local checkpoint") + manager = TokenizeManager(AutoTokenizer.from_pretrained(path, local_files_only=True)) + tools = [{"type": "function", "function": {"name": f"read_{i}", + "description": "Inspect the repository source. " * 100, + "parameters": {"type": "object", "properties": {}}}} for i in range(tool_count)] or None + messages = [{"role": "system", "content": "Stable instructions."}, + {"role": "user", "content": "Inspect this file.\n" + "value = 1\n" * 256}] + assert manager.inline_system_mode(tools, {}) == "fold" + for turn in range(3): + messages.extend([{"role": "assistant", "content": f"Answer {turn}."}, + {"role": "user", "content": f"Question {turn}."}, + {"role": "system", "content": f"Budget reminder {turn}."}]) + original = deepcopy(messages) + msg = TokenizeMsg(uid=turn, text=messages, tools=tools, sampling_params=SamplingParams(), + inline_system_policy="auto") + prompt = manager.render_prompt(msg) + assert prompt.count("Reasoning strength:") == 1 + assert prompt.count("# Valid recipients:") == 1 + assert prompt.count("// Function schemas") == bool(tools) + for previous in range(turn+1): + assert prompt.count(f"Budget reminder {previous}.") == 1 + ids = manager.tokenize([msg])[0].input_ids.tolist() + msg.inline_system_policy = "fold" + assert ids == manager.tokenize([msg])[0].input_ids.tolist() + assert messages == original + + latest = [m for i, m in enumerate(messages) if m["role"] != "system" or i in (0, len(messages)-1)] + updated = deepcopy(latest) + updated[-1]["content"] = "Budget reminder 900." + tokens = [] + for conversation in (latest, updated): + msg = TokenizeMsg(uid=0, text=conversation, tools=tools, sampling_params=SamplingParams(), + inline_system_policy="auto") + tokens.append(manager.tokenize([msg])[0].input_ids.tolist()) + common = next((i for i, pair in enumerate(zip(*tokens)) if pair[0] != pair[1]), min(map(len, tokens))) + assert common >= min(map(len, tokens)) - 16 class FakeTokenizer: def __init__(self) -> None: self.chat_template_kwargs = None + self.chat_template_messages = None def apply_chat_template(self, messages, **kwargs): self.chat_template_kwargs = kwargs + self.chat_template_messages = deepcopy(messages) return "rendered prompt" def encode(self, prompt, return_tensors=None, add_special_tokens=True): @@ -74,6 +365,31 @@ def test_tokenize_manager_passes_tools_to_chat_template(): assert input_ids.tolist() == [1, 2, 3] +@pytest.mark.parametrize("initial_system", [[], [{"role": "system", "content": "stable instructions"}]]) +def test_generic_template_folds_system_messages_without_mutating_input(initial_system): + tokenizer = FakeTokenizer() + manager = TokenizeManager(tokenizer) + messages = initial_system + [ + {"role": "user", "content": "inspect files"}, + {"role": "system", "content": "first reminder"}, + {"role": "assistant", "content": "found the bug"}, + {"role": "system", "content": "second reminder"}, + ] + original = deepcopy(messages) + tools = [{"type": "function", "function": {"name": "read", "parameters": {"type": "object"}}}] + original_tools = deepcopy(tools) + manager.tokenize([TokenizeMsg(uid=1, text=messages, tools=tools, sampling_params=SamplingParams(), + inline_system_policy="auto")]) + assert tokenizer.chat_template_messages == initial_system + [ + {"role": "user", "content": "inspect files\n\nfirst reminder"}, + {"role": "assistant", "content": "found the bug"}, + {"role": "user", "content": "second reminder"}, + ] + assert tokenizer.chat_template_kwargs["tools"] == original_tools + assert messages == original + assert tools == original_tools + + class FakeDsv4Tokenizer: chat_template = None @@ -89,6 +405,36 @@ def encode(self, prompt, return_tensors=None, add_special_tokens=True): return torch.tensor([[4, 5, 6]], dtype=torch.long) +def test_dsv41_encoder_gets_image_sources_before_and_after_fetch(tmp_path): + import copy + from freetoken.mm.media import collect_image_refs + encoding_dir = tmp_path / "encoding" + encoding_dir.mkdir() + (encoding_dir / "encoding_dsv4.py").write_text(''' +def encode_messages(messages, thinking_mode, reasoning_effort=None): + parts = messages[0]["content"] + assert parts[1] == {"type": "text", "text": "compare"} + assert parts[0].get("url") == "https://example.com/image.png" or parts[0].get("data") == b"first image bytes" + assert parts[2].get("data") in ("aW1hZ2U=", b"second image bytes") + return "two images" +''') + messages = [{"role": "user", "content": [ + {"type": "image", "freetoken_ref": {"kind": "url", "data": "https://example.com/image.png"}}, + {"type": "text", "text": "compare"}, + {"type": "image", "freetoken_ref": {"kind": "b64", "data": "aW1hZ2U="}}, + ]}] + manager = TokenizeManager(FakeDsv4Tokenizer(tmp_path)) + msg = TokenizeMsg(uid=1, text=messages, sampling_params=SamplingParams()) + original = copy.deepcopy(messages) + assert manager.render_prompt(msg) == "two images" + assert messages == original + assert len(collect_image_refs(messages)) == 2 + msg.images = [b"first image bytes", b"second image bytes"] + stripped = copy.deepcopy(messages) + assert manager.render_prompt(msg) == "two images" + assert messages == stripped + + def test_tokenize_manager_uses_dsv4_encoder_when_chat_template_is_missing(tmp_path): encoding_dir = tmp_path / "encoding" encoding_dir.mkdir() @@ -126,6 +472,166 @@ def encode_messages(messages, thinking_mode, reasoning_effort=None): assert input_ids.tolist() == [4, 5, 6] +@pytest.mark.parametrize("initial_system", [[], [{"role": "system", "content": "stable instructions"}]]) +def test_dsv4_fold_keeps_instructions_and_tools_on_stable_prefix(tmp_path, initial_system): + encoding_dir = tmp_path / "encoding" + encoding_dir.mkdir() + (encoding_dir / "encoding_dsv4.py").write_text( + "import json\n" + "def encode_messages(messages, thinking_mode, reasoning_effort=None):\n" + " return json.dumps(messages)\n" + ) + tokenizer = FakeDsv4Tokenizer(tmp_path) + manager = TokenizeManager(tokenizer) + tools = [{"type": "function", "function": {"name": "read", "parameters": {"type": "object"}}}] + messages = initial_system + [ + {"role": "user", "content": "inspect files"}, + {"role": "system", "content": "first reminder"}, + {"role": "assistant", "content": "found the bug"}, + ] + original, original_tools = deepcopy(messages), deepcopy(tools) + manager.tokenize([TokenizeMsg(uid=1, text=messages, tools=tools, sampling_params=SamplingParams(), + inline_system_policy="fold")]) + first = json.loads(tokenizer.prompt) + extended = messages + [{"role": "system", "content": "next reminder"}] + original_extended = deepcopy(extended) + manager.tokenize([TokenizeMsg(uid=2, text=extended, tools=tools, sampling_params=SamplingParams(), + inline_system_policy="fold")]) + second = json.loads(tokenizer.prompt) + assert [m["role"] for m in first] == ["system", "user", "assistant"] + assert first[0]["content"] == (initial_system[0]["content"] if initial_system else "") + assert first[0]["tools"] == original_tools + assert first[1]["content"] == "inspect files\n\nfirst reminder" + assert second[:-1] == first + assert second[-1] == {"role": "user", "content": "next reminder"} + assert all("tools" not in m for m in second[1:]) + assert messages == original and extended == original_extended and tools == original_tools + + +@pytest.fixture(scope="module") +def native_dsv4_manager(): + model_path = os.environ.get("FREETOKEN_DSV4_MODEL") + if not model_path: + pytest.skip("set FREETOKEN_DSV4_MODEL to a local DSV4 checkpoint") + from transformers import AutoTokenizer + + manager = TokenizeManager(AutoTokenizer.from_pretrained(model_path, local_files_only=True)) + assert manager._dsv4_encoder is not None, "checkpoint must supply the native DSV4 encoder" + return manager + + +def test_dsv4_no_tools_retains_reminders_when_encoder_drops_developer_turns(tmp_path): + encoding_dir = tmp_path / "encoding" + encoding_dir.mkdir() + (encoding_dir / "encoding_dsv4.py").write_text( + "import json\n" + "def encode_messages(messages, thinking_mode, reasoning_effort=None):\n" + " if thinking_mode == 'thinking':\n" + " messages = [m for m in messages if m['role'] != 'developer']\n" + " return json.dumps(messages)\n" + ) + manager = TokenizeManager(FakeDsv4Tokenizer(tmp_path)) + messages = [ + {"role": "user", "content": "First task."}, + {"role": "system", "content": "Keep this earlier instruction."}, + {"role": "assistant", "content": "First answer."}, + {"role": "user", "content": "Second task."}, + ] + original = deepcopy(messages) + msg = TokenizeMsg(uid=1, text=messages, sampling_params=SamplingParams(), + chat_template_kwargs={"enable_thinking": True}, inline_system_policy="fold") + rendered = json.loads(manager.render_prompt(msg)) + assert all(m["role"] in ("system", "user", "assistant") for m in rendered) + assert sum(m["content"].count("Keep this earlier instruction.") for m in rendered) == 1 + assert messages == original + + +@pytest.mark.needs_weights +@pytest.mark.parametrize("system", [None, "You are a coding assistant."]) +@pytest.mark.parametrize("change", ["update", "append"]) +def test_native_dsv4_reminder_update_preserves_token_prefix(native_dsv4_manager, system, change): + from freetoken.server.anthropic_api import convert_anthropic_to_genspec + from freetoken.server.anthropic_models import AnthropicMessagesRequest + + data = { + "model": "dsv4", "max_tokens": 64, + "tools": [{"name": "read", "input_schema": {"type": "object", "properties": {}}}], + "messages": [ + {"role": "user", "content": "Inspect this file:\n" + "value = 1\n" * 512}, + {"role": "system", "content": "Keep the changes minimal."}, + {"role": "assistant", "content": [ + {"type": "thinking", "thinking": "I should read the source.", "signature": ""}, + {"type": "tool_use", "id": "call_read", "name": "read", "input": {}}, + ]}, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "call_read", "content": "value = 1"}]}, + {"role": "system", "content": "Remaining token budget: 1000"}, + ], + } + if system is not None: + data["system"] = system + tokens = [] + for budget in (1000, 900): + if budget == 900 and change == "append": + data["messages"].extend([ + {"role": "assistant", "content": [ + {"type": "thinking", "thinking": "The file is simple.", "signature": ""}, + {"type": "text", "text": "I inspected the file."}, + ]}, + {"role": "system", "content": "Remaining token budget: 900"}, + ]) + else: + data["messages"][-1]["content"] = f"Remaining token budget: {budget}" + req = AnthropicMessagesRequest.model_validate(data) + original = req.model_dump() + spec = convert_anthropic_to_genspec(req, {}) + msg = TokenizeMsg(uid=budget, text=spec.messages, tools=spec.template_tools, + sampling_params=spec.sampling_params, chat_template_kwargs=spec.chat_template_kwargs, + inline_system_policy=spec.inline_system_policy) + prompt = native_dsv4_manager.render_prompt(msg) + assert prompt.count("Keep the changes minimal.") == 1 + assert prompt.index("Inspect this file:") < prompt.index(f"Remaining token budget: {budget}") + tokens.append(native_dsv4_manager.tokenize([msg])[0].input_ids.tolist()) + assert req.model_dump() == original + common = next((i for i, pair in enumerate(zip(*tokens)) if pair[0] != pair[1]), min(map(len, tokens))) + assert common > 2048 + assert common >= min(map(len, tokens)) - 64 + + +@pytest.mark.needs_weights +@pytest.mark.parametrize("disable_declared_tools", [False, True]) +@pytest.mark.parametrize("thinking_enabled", [False, True]) +def test_native_dsv4_no_tools_keeps_earlier_reminder( + native_dsv4_manager, disable_declared_tools, thinking_enabled +): + from freetoken.server.anthropic_api import convert_anthropic_to_genspec + from freetoken.server.anthropic_models import AnthropicMessagesRequest + + data = { + "model": "dsv4", "max_tokens": 64, "system": "Stable instructions.", + "thinking": {"type": "enabled" if thinking_enabled else "disabled"}, + "messages": [ + {"role": "user", "content": "First task."}, + {"role": "system", "content": "Keep this earlier instruction."}, + {"role": "assistant", "content": "First answer."}, + {"role": "user", "content": "Second task."}, + ], + } + if disable_declared_tools: + data["tools"] = [{"name": "read", "input_schema": {"type": "object", "properties": {}}}] + data["tool_choice"] = {"type": "none"} + req = AnthropicMessagesRequest.model_validate(data) + original = req.model_dump() + spec = convert_anthropic_to_genspec(req, {}) + assert not spec.template_tools + msg = TokenizeMsg(uid=1, text=spec.messages, tools=spec.template_tools, + sampling_params=spec.sampling_params, chat_template_kwargs=spec.chat_template_kwargs, + inline_system_policy=spec.inline_system_policy) + prompt = native_dsv4_manager.render_prompt(msg) + assert prompt.count("Keep this earlier instruction.") == 1 + assert "Second task." in prompt + assert req.model_dump() == original + + def test_dsv4_encoder_gets_tool_call_arguments_as_json_string(tmp_path): """Regression: render_messages hands the template dict arguments; the dsv4 encoder contract is a JSON-object STRING -- a dict trips its fallback that @@ -285,3 +791,60 @@ def test_tokenize_survives_an_unhashable_effort(): manager.tokenize([msg]) assert "reasoning_effort" not in tokenizer.chat_template_kwargs + + +@pytest.mark.parametrize("budget", [1, 25, 50, 75, 100]) +def test_numeric_encoder_effort_survives_prompt_rendering(tmp_path, budget): + encoding_dir = tmp_path / "encoding" + encoding_dir.mkdir() + (encoding_dir / "encoding.py").write_text(''' +def encode_messages(messages, thinking_mode, reasoning_effort=None): + if thinking_mode != "thinking": + return "chat" + effort = "high" if reasoning_effort is None else reasoning_effort + if isinstance(effort, str): + effort = {"low": 50, "high": 75, "max": 100}[effort] + if type(effort) is not int or not 1 <= effort <= 100: + raise ValueError("invalid budget") + return f"Reasoning Effort: {effort}" +''') + manager = TokenizeManager(FakeDsv4Tokenizer(tmp_path)) + msg = TokenizeMsg( + uid=1, text=[{"role": "user", "content": "hello"}], + sampling_params=SamplingParams(), + chat_template_kwargs={"enable_thinking": True, "reasoning_effort": budget}, + ) + 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..e599db7265 --- /dev/null +++ b/tests/utils/test_mp_coalesced.py @@ -0,0 +1,116 @@ +# 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() + + +async 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)") + 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() \ No newline at end of file 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"])