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_load_weight_generic.py b/benchmarks/bench_load_weight_generic.py index 4e4c5fe012..62807bc8fe 100644 --- a/benchmarks/bench_load_weight_generic.py +++ b/benchmarks/bench_load_weight_generic.py @@ -41,6 +41,8 @@ import torch +from freetoken.memory import effective_memory_available + GiB = float(1 << 30) ALL_MODES = ("parallel", "ftw") @@ -54,14 +56,6 @@ def _default_ftw_dir(model_path: str) -> str: # ---------------- memory sampling + checksum ---------------- -def _meminfo_available_bytes() -> int: - with open("/proc/meminfo") as f: - for line in f: - if line.startswith("MemAvailable:"): - return int(line.split()[1]) * 1024 - return 0 - - def _status_kb(key: str) -> int: with open("/proc/self/status") as f: for line in f: @@ -71,7 +65,7 @@ def _status_kb(key: str) -> int: class MemSampler(threading.Thread): - """Background sampler: peak process RSS + system MemAvailable low-water-mark.""" + """Background sampler: peak RSS + effective available-memory low-water-mark.""" def __init__(self, interval: float = 0.2): super().__init__(daemon=True) @@ -82,7 +76,9 @@ def __init__(self, interval: float = 0.2): def run(self) -> None: while not self._stop_evt.wait(self.interval): - self.min_avail = min(self.min_avail, _meminfo_available_bytes()) + available = effective_memory_available() + if available is not None: + self.min_avail = min(self.min_avail, available) self.max_rss = max(self.max_rss, _status_kb("VmRSS:")) def stop(self) -> None: @@ -179,7 +175,7 @@ def _bench_load(mode: str, model_path: str, *, parallel: bool, workers: int, chu t = time.perf_counter() try: banks = load_expert_banks(model_path, mc, device=torch.device("cuda", torch.cuda.current_device()), - dtype=torch.bfloat16, parallel=parallel, + parallel=parallel, workers=workers, chunk=chunk) except NotImplementedError as e: s.stop() 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 0c33943987..d9ec62a0d0 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -11,7 +11,7 @@ ft [args] | `ft ctl` | Query and manage a running server over HTTP | | `ft launch` | Configure and launch a coding agent against a server | | `ft checkpoint` | Convert an HF checkpoint to the FTW fast-load format | -| `ft bench bw` | Benchmark CPU vs PCIe bandwidth to calibrate the MoE backend | +| `ft bench bw` | Benchmark CPU vs PCIe bandwidth to calibrate hybrid MoE decode | `ft --version` prints the installed version (torch-free; nightly wheels carry a `+g` build stamp, tagged releases a bare version). Every command supports @@ -33,6 +33,7 @@ parsers all resolve automatically from the checkpoint and the GPU. |---|---|---| | `--model-path`, `--model` | required | Local dir, HF repo id, or an FTW dir (auto-detected) | | `--served-model-name` | basename of `--model` | Model id reported by `/v1/models` | +| `--hf-overrides` | — | JSON object applied to the checkpoint's config as vLLM's `--hf-overrides`: a nested config section updates key by key, any other value is replaced whole. A YaRN `rope_parameters` override serves `original_max_position_embeddings * factor` positions | ### Server & runtime @@ -40,6 +41,8 @@ parsers all resolve automatically from the checkpoint and the GPU. |---|---|---| | `--host` | 127.0.0.1 | Bind address | | `--port` | 1919 | Bind port | +| `--api-key` | `$FREETOKEN_API_KEY`, else off | Require the key on every route except `/health`, as `Authorization: Bearer ` or `x-api-key: `; other requests get a 401 | +| `--dist-port` | `--port` + 1 | Internal TP rendezvous port (loopback-only regardless of `--host`) | | `--gpu` | GPU 0 | GPU to run on: a UUID from `nvidia-smi -L` or an `nvidia-smi` index; see [below](#choosing-a-gpu) | | `--max-running-requests` | 4 | Max concurrently running requests | | `--max-output-tokens` | 32768 | Default output budget for requests that omit one | @@ -71,22 +74,51 @@ 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` | +| `--logprob-read-snapshot-donate` | all | hybrid GDN models: snapshots a finished `max_tokens=1` + `logprobs` request donates to the prefix cache: `all` (x64 prefill boundary and final state), `boundary` (only the former) or `none` | +| `--logprob-read-isolate` | off | a `max_tokens=1` + `logprobs` request prefills in a pass of its own and neither resumes from nor stores prefix-cache entries, so its logprobs do not depend on other traffic; costs read throughput under concurrency and prefix reuse for reads | +| `--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; fused on unified-memory GPUs (GB10), see [models.md](models.md#unified-memory-gpus-gb10--dgx-spark). `--moe-backend` is the deprecated old spelling | | `--quant-backend` | auto | Kernel per quantized layer type, `layer[.kind]=name` entries: `linear=marlin,moe=b12x` or `moe.nvfp4=triton`. A layer-level entry applies to every kind whose table lists the name | | `--nvfp4-backend` | — | Deprecated: stands in for `--quant-backend moe.nvfp4=` (`flashinfer` means b12x); cannot be combined with `--quant-backend` | -| `--moe-cache-size` / `--moe-cache-rate` / `--moe-cache-auto` | auto | GPU expert-cache size as slots / fraction of all experts / sized from free VRAM (mutually exclusive; auto is enabled by default for offload-family strategies) | +| `--moe-cache-size` / `--moe-cache-rate` / `--moe-cache-mem` / `--moe-cache-auto` | auto | GPU expert-cache size as slots / fraction of all experts / memory such as `12GiB` / sized from free VRAM (mutually exclusive; auto is enabled by default for offload-family strategies) | +| `--moe-cache-scope` | unified | `unified`: one LRU pool of slots per expert quant (quant type and bank layout), so all MoE layers share one pool when they share a quant format; `per-layer`: each layer has its own pool. `per-layer` does not combine with `--moe-strategy cpu` or `--moe-cpu-layers` | +| `--moe-cache-split` | equal | How the pools share the slots: `equal` (in proportion to their layer counts), or `profile:`: every pool keeps `top_k` slots and the rest go to the most used experts in the profile's per-layer expert counts. With `offload`, a pool smaller than the expert count lowers `--max-running-requests` to its slots / top_k | | `--kv-reserve-tokens` | 8192 | KV token floor reserved before `--moe-cache-auto` fills experts | | `--moe-cpu-threads` | physical cores | CPU worker threads for the cpu/hybrid executor | -| `--moe-cpu-layers` | all on GPU | With `offload`: which MoE layers decode on CPU (`3,7,11`, a count, a fraction, or `auto`). `auto` is for Windows/WSL only, where CUDA pinned memory is capped; every value needs an expert format the CPU executor serves (bf16, nvfp4, mxfp4), so fp8 experts cannot use it | +| `--moe-cpu-layers` | all on GPU | With `offload`: which MoE layers decode on CPU (`3,7,11`, a count, a fraction, or `auto`). `auto` is for Windows/WSL only, where CUDA pinned memory is capped; every value needs an expert format the CPU executor serves (bf16, fp8_block, nvfp4, mxfp4, GGUF on x86 CPUs) | | `--moe-hybrid-max-fetch` | auto | With `hybrid`: max experts fetched over PCIe per layer per step; rest computed on CPU | +| `--moe-prefill-demand-tokens` | 0 | Prefill: a chunk of at most this many tokens loads only the experts it routes to, through the expert cache as decode does (GPU decode only); 0 streams whole layers | | `--moe-prefill-hit-d2d` | off | Prefill: copy cache-hit experts device-side, stream only misses (CUDA >= 13) | | `--disable-moe-prefill-overlap` | overlap on | Disable the two-buffer prefill copy overlap | @@ -98,6 +130,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 @@ -118,7 +207,7 @@ so a client can gate its attachment controls without reading the checkpoint conf | `--image-min-tokens`, `--image-max-tokens` | processor defaults | Per-image token budget: the image processor resizes every image to take between these many tokens, converted to the family's own units by its processor. A family with fixed budgets honors the maximum only and refuses one below its smallest budget at start-up | | `--mm-processor-kwargs` | none | JSON object of extra keyword arguments for the checkpoint's image processor call, for knobs the token budget does not cover; applied after the budget, so an explicit key wins | | `--mm-embed-cache-device` | cpu | Where encoded image embeddings live between prefill chunks. `cpu` keeps them out of the VRAM budget; `cuda` skips the copy back | -| `--allowed-media-domains` | any | Comma-separated hostname allowlist for image URLs; requests for other domains are rejected with a 400. Empty allows any domain | +| `--allowed-media-domains` | any | Comma-separated hostname allowlist for image URLs; requests for other domains are rejected with a 400. Empty allows any domain. Redirects are re-checked against the list and may not land on a private or loopback address | | `--allowed-local-media-path` | off | Directory `file://` image refs may be read from; unset rejects local files | ## ft shell @@ -129,14 +218,17 @@ ft shell --model ~/models/Qwen3.6-35B-A3B # serve + chat in one process ``` - Attach mode talks to `--server URL` (default `http://127.0.0.1:1919`) +- `--api-key KEY` (default `$FREETOKEN_API_KEY`) attaches to a server started with `--api-key`; with `--model` the shell uses that server's own key - `/help` inside the shell lists the commands (`/think`, `/cache`, `/reset`). ## ft ctl ```bash -ft ctl [--base-url http://127.0.0.1:1919] [--timeout 10] [--json] +ft ctl [--base-url http://127.0.0.1:1919] [--api-key KEY] [--timeout 10] [--json] ``` +`--api-key` defaults to `$FREETOKEN_API_KEY`. + | Subcommand | Endpoint | Purpose | |---|---|---| | `health` | `GET /health` | Server status, model, load progress | @@ -164,6 +256,7 @@ Code and Hermes need no declaration. | Flag | Meaning | |---|---| | `--server URL` | Server to point the agent at (default `http://127.0.0.1:1919`) | +| `--api-key KEY` | Key of a server started with `--api-key` (default `$FREETOKEN_API_KEY`). The agent gets it in its environment; Hermes and OpenClaw keep it in their config files | | `--dry-run` | Print the planned config changes and command, touch nothing | | `-y`, `--yes` | Approve install/config prompts | | `--config` | Configure without launching | @@ -178,27 +271,66 @@ 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 (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. ## ft bench bw ```bash -ft bench bw # once per GPU -ft bench bw --dtype nvfp4,bf16 # only the formats you serve -ft bench bw --gpu 1 # a specific GPU (UUID or nvidia-smi index, as for ft serve) +ft bench bw --model /path/to/model.gguf # every MoE layer kind of a checkpoint (GGUF or HF) +ft bench bw # canonical geometry per expert format +ft bench bw --gpu 1 # a specific GPU (UUID or nvidia-smi index, as for ft serve) ``` -Measures host-RAM vs PCIe bandwidth with the real cpu/offload MoE kernels and writes a -profile that `ft serve --moe-strategy auto` and `--moe-hybrid-max-fetch -1` then read. - -- One profile per GPU, at `~/.cache/freetoken/benchbw/.json`. -- Keyed on expert format + GPU, so a profile from other hardware is ignored rather than - misapplied. An older single `benchbw.json` still counts if its GPU name matches. -- What to measure: `--dtype`, `--model`, `--formats`, `--isa`. -- `--threshold` (default 2.0) sets the call: recommend hybrid when CPU bandwidth beats PCIe - by that factor. - +Measures host-RAM vs PCIe bandwidth with the real cpu/offload MoE kernels and writes one +row per MoE layer kind, which `ft serve --moe-strategy hybrid --moe-hybrid-max-fetch -1` reads +to split each layer's misses. + +- One profile per GPU, at `~/.cache/freetoken/benchbw/.json`; a profile from other + hardware or another profile version is ignored. +- A row's key is the layer's CPU format and options, expert shape, top_k and CPU worker + threads; a layer without an exact row keeps a fetch cap of 1. Bench with the same + `--cpu-threads` as the server's `--moe-cpu-threads`. +- Each run replaces the rows it benched and keeps the others. +- What to measure: `--model` (checkpoint paths or presets), `--dtype`, `--formats`, `--isa`. +- `--threshold` (default 2.0) sets each row's `recommended`: 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/ftw-hotfix.md b/docs/ftw-hotfix.md index 299d522d6e..0c6b6b9feb 100644 --- a/docs/ftw-hotfix.md +++ b/docs/ftw-hotfix.md @@ -32,6 +32,8 @@ byte range, never the whole checkpoint; `--revision` pins the Hub revision to re | `KeyError: 'model.embed.weight'` | deepseek-ai/DeepSeek-V4-Flash-0731 | rename the index entries | none | | `RuntimeError: Unexpected keys ... .weight_scale` | nvidia/GLM-5.2-NVFP4 | dequantize the old runtime-fp8 weights back to bf16 | none | | `PLE shard indices are not contiguous 0..N-1: []` | Qwen3.8-Flash-Next (FTWs converted before #420) | write the PLE table as `ple-table-*.safetensors` | 47.7 GiB | +| `ValueError: ... stores NVFP4 global scales as fp16` | NVFP4 FTWs converted by older builds | none: the fp16 rounding cannot be undone, so the script refuses; reconvert with `ft checkpoint` | - | +| none: loads, but the output is wrong | Qwen3.5 / Qwen3.6 / Qwen3.8-27B FTWs converted by older builds (norm weights stored as bf16(1 + w)) | none: the script does not detect it; reconvert with `ft checkpoint` | - | | `ValueError: ... holds no vision encoder tensors` | image-input families (Qwen3-VL, Qwen3.6, Qwen3.8-Flash-Next, Gemma-4, GLM-5.3-Flash, Muse-Glimmer, MiniMax-M3) converted before the family served images | add the vision encoder: the tower's checkpoint tensors are fetched and the family's own weight reader emits them as the converter would | 1 to 4 GiB | The script decides by itself: it builds the current model from the FTW's `config.json`, compares diff --git a/docs/install.md b/docs/install.md index 9438c21a67..f474216bf0 100644 --- a/docs/install.md +++ b/docs/install.md @@ -1,5 +1,7 @@ # Install +AMD users: see the [AMD ROCm installation guide (WIP)](install_amd.md). + ## Requirements - Linux x86_64, NVIDIA GPU, driver r580+ (CUDA 13) diff --git a/docs/install_amd.md b/docs/install_amd.md new file mode 100644 index 0000000000..ed934ad657 --- /dev/null +++ b/docs/install_amd.md @@ -0,0 +1,55 @@ +# AMD ROCm installation (WIP) + +AMD support is experimental and remains a work in progress. See the +[AMD Support roadmap](https://github.com/FlashML-org/FreeToken/issues/541) for +the current integration and qualification status. + +## Requirements + +- Linux x86_64 +- AMD RDNA3/RDNA4 GPU (`gfx1100`-`gfx1103`, `gfx1200`, or `gfx1201`) +- ROCm 7.14 +- Python >= 3.10 + +## Install from source + +Use an official ROCm PyTorch image whose PyTorch version satisfies the project's +`torch>=2.11,<2.12` constraint. For RDNA4, the matching ROCm 7.14 image is: + +```bash +VIDEO_GID="$(getent group video | cut -d: -f3)" +RENDER_GID="$(getent group render | cut -d: -f3)" +docker run --rm -it \ + --device=/dev/kfd --device=/dev/dri \ + --group-add="$VIDEO_GID" --group-add="$RENDER_GID" --ipc=host \ + --cap-add=SYS_PTRACE --security-opt seccomp=unconfined \ + -e PYTORCH_ROCM_ARCH=gfx1201 -e FREETOKEN_ROCM_ARCH=gfx1201 \ + -v "$PWD:/workspace/FreeToken" -w /workspace/FreeToken \ + rocm/pytorch:rocm7.14_ubuntu24.04_py3.12_pytorch_release_2.11.0 bash +``` + +Inside the container, preserve the ROCm-enabled PyTorch already supplied by the +image and disable build isolation so it is also used to compile the extensions: + +```bash +python -m pip install --no-build-isolation -e . +``` + +Set both architecture variables to `gfx1200` for RX 9060 family GPUs, or to the +actual target reported by `rocminfo`. + +## Optional kernel backends + +FreeToken's FlashInfer, `sgl_kernel`, and vLLM kernel integrations are CUDA-only. +On ROCm, these paths use the built-in fallbacks; do not install the `[accel]`, +`[fi]`, or `[sgl]` extras. Forcing `--attention-backend fi`, `fa`, or `trtllm`, +or NVFP4 Marlin/b12x (e.g. `--quant-backend moe.nvfp4=marlin` or +`--quant-backend moe.nvfp4=b12x`), fails with a ROCm-specific error. + +## CPU and hybrid MoE + +`--moe-strategy cpu`, `hybrid` and `--moe-cpu-layers` sync the GPU with the CPU +expert pool through HIP stream memory operations. ROCm 10 records them into CUDA +graphs itself; on ROCm 7.14, which does not, FreeToken adds the graph nodes +explicitly. When the startup probe of that path fails, CUDA graph capture stops +with an error; run those strategies with `--cuda-graph-max-bs 0`. diff --git a/docs/models.md b/docs/models.md index 548a363c67..4d3743278f 100644 --- a/docs/models.md +++ b/docs/models.md @@ -5,7 +5,8 @@ for them; other checkpoints of the same architectures work too. | Model | HF checkpoints | |---|---| -| DeepSeek-V4 | [deepseek-ai/DeepSeek-V4-Flash-0731](https://huggingface.co/deepseek-ai/DeepSeek-V4-Flash-0731) | +| 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), [deepseek-ai/DeepSeek-V4-Flash-Vision-Exp](https://huggingface.co/deepseek-ai/DeepSeek-V4-Flash-Vision-Exp) | | 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) | | GLM-4.7 | [nvidia/GLM-4.7-NVFP4](https://huggingface.co/nvidia/GLM-4.7-NVFP4) | @@ -27,36 +28,109 @@ 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 | | GLM-5.3-Flash (`glm5_next`: ViT tower, streamed under `--mm-encoder-weights host`) | one token per 28x28 pixels of the resized image, dynamic resolution on a canvas zero-padded to a 28-multiple | token counts, passed through as the processor's `min_image_tokens` / `max_image_tokens`; checkpoint defaults 16 to 8000 tokens | `{"max_image_tokens": 2048}` | | Muse-Glimmer-30B (windowed ViT tower, streamed under `--mm-encoder-weights host`) | one token per 28x28 pixels of the resized image, aspect ratio kept under a token cap; checkpoint default 4096 tokens | the maximum is the cap (`max_image_tokens`); the minimum has no effect | `{"max_image_tokens": 1024}` | | MiniMax-M3 (`minimax_m3`: CLIP-style ViT tower, streamed under `--mm-encoder-weights host`) | one token per 28x28 pixels of the resized image, dynamic resolution | pixel areas in `size.shortest_edge` / `longest_edge`; checkpoint defaults 4 to 576 tokens | `{"size": {"longest_edge": 1048576}}` | +| DeepSeek-V4-Flash-Vision (`deepseek_v4`: native-resolution ViT tower, streamed under `--mm-encoder-weights host`) | one token per 3x3 grid of 14x14 patches, against a per-image cap; the block carries sentinel tokens and lead pads that align it to a 4-token compression stride, so its length depends on where it lands in the prompt | the maximum is a cap on the block (`vision_max_n_token`, checkpoint default 384) and lowers it, the minimum is read as the checkpoint's upscale floor and has no effect | not read | ## MoE strategies `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 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. - **hybrid** — per step, fetches some misses over PCIe and computes the rest on - CPU, overlapped. Run `ft bench bw` once per machine to calibrate the split. + CPU, overlapped. Run `ft bench bw --model ` once per machine and + model 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. + `offload`. 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. 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 +- A checkpoint that fails to load: from a source checkout, run `PYTHONPATH=python python scripts/checkpoint_diagnose/diagnose.py `. It reads the headers, not the weights, says whether FreeToken or the checkpoint is at fault, and for a local checkpoint gives the commands that fix a non-standard one; see [its README](../scripts/checkpoint_diagnose/README.md). GGUF files are checked with [scripts/gguf_check](../scripts/gguf_check/README.md). - `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. +- Muse-Glimmer keeps the attention gate as its own linear next to `qkv_proj`; an FTW converted + when the gate was fused into `qkvg_proj` no longer loads, reconvert it with `ft checkpoint`. +- NVFP4 global scales are kept in fp32. An NVFP4 FTW that stores them as fp16 (converted by an + older build) is refused at load; reconvert it with `ft checkpoint`. +- Qwen3.5 / Qwen3.6 / Qwen3.8-27B load their zero-centred norm weights as stored and add the 1 in + fp32 at runtime. An FTW of these models converted by an older build stores the norms with the 1 + already added: it loads without an error but gives wrong output; reconvert it with `ft checkpoint`. - An FTW converted before its family served images holds no vision encoder: `ft serve` refuses it unless started with `--text-model-only` (or `--mm-disable vision`); reconvert it 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 ` 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..ec795b0823 --- /dev/null +++ b/docs/openai_api.md @@ -0,0 +1,111 @@ +# 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` | honoured (`top_logprobs` 0..20; message content only, under the supported reasoning and tool parsers) | 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 | +| `allowed_token_ids`, `return_tokens_as_token_ids`, `logprob_token_ids`, `prompt_logprobs`, `guided_*`, `structured_outputs`, `bad_words` | 400 when set (`null` / `false` / empty pass) | honoured | partly | partly | +| `best_of`, `use_beam_search`, `length_penalty`, `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, a list of strings, a list of token ids or a list of token-id lists +(each id in `[0, len(tokenizer))`; not with `echo` or `suffix`); 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` | honoured, 0..32: raw pre-temperature logprobs of the sampled token and its top alternatives; -inf is sent as -9999 | honoured (`--max-logprobs`, default 20) | honoured | `n_probs` | +| `allowed_token_ids` | honoured, 1..1024 ids below `len(tokenizer)`: sampling is masked to them at every position. With `logprobs` the values are renormalized over the set (vLLM's `processed_logprobs` at temperature 0), so `temperature > 0`, penalties, `min_tokens` and `logit_bias` are a 400 there; entries past the set are left out | honoured (raw logprobs unless `--logprobs-mode processed_logprobs`) | no | no | + +| `return_tokens_as_token_ids` | honoured: `tokens` and the `top_logprobs` keys are `token_id:N`, so ids that decode to the same text stay apart | honoured | n/a | no | + +Set values of the vLLM fields this server does not implement (`logprob_token_ids`, +`prompt_logprobs`, `guided_*`, `structured_outputs`, `bad_words`, ...) are a 400 rather than +ignored; `null` / `false` / empty pass. Empty prompts and more than 256 prompts x `n` are a 400. +An over-long prompt's error carries `code: "context_length_exceeded"`. + +## 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/freetoken-kernel-cache/pyproject.toml b/freetoken-kernel-cache/pyproject.toml index ec37efbfa6..bc75e1b2f7 100644 --- a/freetoken-kernel-cache/pyproject.toml +++ b/freetoken-kernel-cache/pyproject.toml @@ -8,7 +8,7 @@ requires = [ # 77 is the first release that understands the PEP 639 `license` SPDX string below. "setuptools>=77", "wheel", - "torch>=2.11,<2.12", + "torch>=2.14,<2.15", "apache-tvm-ffi==0.1.13.post3", ] build-backend = "build_backend" @@ -22,7 +22,7 @@ requires-python = ">=3.10" license = "Apache-2.0" dependencies = ["apache-tvm-ffi==0.1.13.post3"] -# Provenance pin: PyPI's torch 2.11.0 is itself the cu130 build, but this index serves +# Provenance pin: PyPI's torch 2.14.1 is itself the cu130 build, but this index serves # ONLY cu130 wheels, so an isolated build can never resolve a different-CUDA torch and # stamp the wheel +cu126/+cu128. uv applies these to build-system.requires when this # directory is the build root. Mirrors ../pyproject.toml. diff --git a/install.sh b/install.sh index 2a9a3dec30..68738cc695 100755 --- a/install.sh +++ b/install.sh @@ -20,7 +20,8 @@ # Defaults to $DEFAULT_KERNEL_CACHE_WHEEL_URL; if # unset and FREETOKEN_WHEEL is local, the script # auto-detects a sibling freetoken_kernel_cache-*.whl. -# FREETOKEN_HOME install root (default: ~/.freetoken); venv at $FREETOKEN_HOME/venv +# FREETOKEN_HOME install root (default: ~/.freetoken); venv at $FREETOKEN_HOME/venv; +# a successful install writes it to environment.d # FREETOKEN_PY_VERSION python for the venv (default: 3.12 — must match the wheel tag) # FREETOKEN_BIN_DIR where to symlink `ft` (default: ~/.local/bin) # FREETOKEN_ENV_DIR environment.d dir (default: ~/.config/environment.d) @@ -205,7 +206,7 @@ mkdir -p "$FT_HOME" # re-install can't inherit a stale/mismatched torch (e.g. an old cu128 venv after a cu130 bump). "$UV" venv "$VENV" --python "$PY_VERSION" --clear -# PyPI's torch 2.11.0 is the same cu130 build the pytorch index +# PyPI's torch 2.14.1 is the same cu130 build the pytorch index # serves; the explicit index pins provenance to the cu130 channel. `unsafe-best-match` # is needed because the pytorch index also mirrors stale copies of common deps (e.g. # packaging<=24.1) that would shadow PyPI under uv's first-index strategy; all indexes @@ -248,8 +249,8 @@ ln -sf "$FT_BIN" "$BIN_DIR/ft" say "symlinked $BIN_DIR/ft -> $FT_BIN" mkdir -p "$ENV_DIR" -printf 'FREETOKEN_FT_BIN=%s\n' "$FT_BIN" > "$ENV_DIR/50-freetoken.conf" -say "wrote $ENV_DIR/50-freetoken.conf (FREETOKEN_FT_BIN) — GUI picks it up after next login" +printf 'FREETOKEN_HOME=%s\n' "$(cd "$FT_HOME" && pwd)" >"$ENV_DIR/50-freetoken.conf" +say "wrote $ENV_DIR/50-freetoken.conf (FREETOKEN_HOME) - FreeToken Desktop reads the file, other programs see the variable after next login" # --- 5. Self-check --------------------------------------------------------- if "$FT_BIN" --help >/dev/null 2>&1; then @@ -264,10 +265,9 @@ ${C_GREEN}FreeToken engine installed.${C_RESET} ft binary $FT_BIN on PATH as $BIN_DIR/ft (ensure $BIN_DIR is on PATH) - Desktop env FREETOKEN_FT_BIN via environment.d (re-login to apply) + Desktop env FREETOKEN_HOME in $ENV_DIR/50-freetoken.conf -Run in this shell without re-login: - export FREETOKEN_FT_BIN="$FT_BIN" +Run: ft serve --model --port 1919 EOF diff --git a/pyproject.toml b/pyproject.toml index 25cb204bb4..384b905f2c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -3,7 +3,7 @@ # own, and a mismatch links the C++ extensions against the wrong libtorch. # setuptools floor: 77 is the first release that understands the PEP 639 `license` # SPDX string and `license-files` below. -requires = ["setuptools>=77", "torch>=2.11,<2.12", "wheel"] +requires = ["setuptools>=77", "torch>=2.14,<2.15", "wheel"] build-backend = "setuptools.build_meta" [project] @@ -23,6 +23,7 @@ classifiers = [ "Intended Audience :: Science/Research", "Operating System :: POSIX :: Linux", "Environment :: GPU :: NVIDIA CUDA", + "Environment :: GPU :: AMD ROCm", "Programming Language :: Python :: 3", "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", @@ -38,13 +39,11 @@ dependencies = [ "einops>=0.8,<1", "fastapi>=0.115,<1", # slot_cache: the device-side LRU admission kernel behind the MoE expert cache. - "flashlib==0.3.0", # pin it - "gguf>=0.19,<1", + "flashlib==0.4.0", # pin it "huggingface_hub>=1.5,<2", "msgpack>=1.1,<2", "modelscope>=1.37,<2", - # ceiling: numba (via flashlib) needs numpy<2.5 - "numpy>=2.0,<2.5", + "numpy>=2.0,<3", "openai>=2.0,<3", "partial-json-parser>=0.2,<1", "pillow>=10,<13", @@ -52,14 +51,17 @@ dependencies = [ "pydantic>=2.9,<3", "pyzmq>=27,<28", "safetensors>=0.6,<1", - # floor+ceiling: sglang-kernel 0.4.5 links libtorch symbols only 2.11 has. - # PyPI's torch 2.11.0 wheel is itself the cu130 build, so plain pip resolves + # floor+ceiling: sglang-kernel 0.4.9 requires torch==2.14.1. + # PyPI's torch 2.14.1 wheel is itself the cu130 build, so plain pip resolves # correctly from PyPI alone; uv additionally pins the index below. - "torch>=2.11,<2.12", - "torchvision>=0.26,<0.27", + "torch>=2.14,<2.15", + "torchvision>=0.29,<0.30", "tqdm>=4.66,<5", - "transformers>=5.16,<5.17", - "triton==3.6.0; platform_system == 'Linux'", + "transformers>=5.18,<5.19", + # CUDA torch 2.14 resolves Triton 3.8; AMD's ROCm 7.14 image supplies its + # gfx1201-enabled Triton 3.7 build. Keep both supported without replacing the + # runtime-specific wheel selected by the PyTorch distribution. + "triton>=3.6,<3.9; platform_system == 'Linux'", "uvicorn>=0.30,<1", ] @@ -78,7 +80,7 @@ dev = ["pytest>=6.0"] # `freetoken[accel]` for the full native fast path. fi = ["flashinfer-python[cu13]==0.6.18.post1"] # renamed from sgl-kernel at 0.4; still imports as `sgl_kernel`, so never co-install both -sgl = ["sglang-kernel==0.4.5"] +sgl = ["sglang-kernel==0.4.9"] accel = ["freetoken[fi,sgl]"] # NOTE: the Marlin W4A16 NVFP4 expert-GEMM path (sm_80-99) borrows vLLM's AOT wheel # (vllm>=0.14,<0.15), which pins transformers>=4.56,<5 and so is INCOMPATIBLE with the @@ -87,7 +89,7 @@ accel = ["freetoken[fi,sgl]"] # in a dedicated environment when you specifically need the Marlin path. # uv-only provenance pins — pip ignores this section and resolves from PyPI, which -# works: PyPI's torch 2.11.0 is the same cu130 build this index serves. +# works: PyPI's torch 2.14.1 is the same cu130 build this index serves. # `explicit = true` scopes the index to the one package that needs it # (the torch index also mirrors stale copies of common deps, e.g. packaging<=24.1, # which would otherwise shadow PyPI under uv's first-index strategy). @@ -110,7 +112,7 @@ version = {attr = "freetoken.version.__version__"} where = ["python"] [tool.setuptools.package-data] -"*" = ["csrc/**/*", "moe/configs/**/*.json"] +"*" = ["csrc/**/*", "moe/configs/**/*.json", "third_party/*/LICENSE", "third_party/*/VERSION", "third_party/*/freetoken.patch"] [tool.pytest.ini_options] testpaths = ["tests"] 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 b07c952323..9d2a3bc8b0 100644 --- a/python/freetoken/attention/qsa_sparse.py +++ b/python/freetoken/attention/qsa_sparse.py @@ -50,6 +50,7 @@ # Block-score transient budget (vLLM's number): the fp32 [rows, n_blocks] logits tile is # 256 KB per row at a 1M-token context, so a long prefill must be scored in row chunks. _LOGITS_WORKSPACE_BYTES = 128 << 20 +_SCORE_ROWS = 16 # prefill query rows per score program (qsa_mqa_paged_rows) TORCH_TOPK_ENV = "FREETOKEN_QSA_TORCH_TOPK" @@ -92,6 +93,12 @@ 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 + # Prefill scoring: row-tile maps per (start, end) row chunk, built by the first QSA layer. + score_tiles: dict | None = None # fmt: on def get_last_indices(self, bs: int) -> torch.Tensor: @@ -112,7 +119,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 +159,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] = [] @@ -167,23 +190,22 @@ def _cmp_pages(self, slot: int) -> torch.Tensor: def _index_rope_cache(self) -> torch.Tensor: """cos/sin table of the indexer rope: same rotary_dim and frequencies as the main - attention, ``head_size`` 128 instead of 256, so it is a separate get_rope instance. + attention with ``head_size`` 128 instead of 256, so the attention's own table. The table itself (not RotaryEmbedding.forward) because the indexer's norm+rope is one fused kernel and the compressed keys rope at their group's position, not the query's.""" if self._index_cos_sin is None: - from freetoken.layers.rotary import get_rope + from freetoken.layers.rotary import rope_cos_sin_table rotary = self.rotary_config with torch.device(self.device): - rope = get_rope( - head_dim=self.index_head_dim, - rotary_dim=rotary.rotary_dim, - max_position=rotary.max_position, - base=rotary.base, - rope_scaling=tuple(rotary.scaling.items()) if rotary.scaling else None, + self._index_cos_sin = rope_cos_sin_table( + self.index_head_dim, + rotary.rotary_dim, + rotary.max_position, + rotary.base, + tuple(rotary.scaling.items()) if rotary.scaling else None, ) - self._index_cos_sin = rope._cos_sin_cache.to(self.device) return self._index_cos_sin def _index_rope_rows(self, positions: torch.Tensor) -> torch.Tensor: @@ -288,7 +310,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: @@ -298,7 +327,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), @@ -307,8 +343,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.""" @@ -389,6 +467,8 @@ def _select(self, index, md: QSASparseMetadata, slot: int) -> torch.Tensor: expand_qsa_block_indices, qsa_index_norm_rope, qsa_mqa_paged, + qsa_mqa_paged_rows, + qsa_score_tile_map, ) rows = index.q.shape[0] @@ -413,22 +493,48 @@ def _select(self, index, md: QSASparseMetadata, slot: int) -> torch.Tensor: columns = md.block_table.shape[1] * self.cmp_page_size indices = self._scratch("indices", rows, self.select_width, dtype=torch.int32) rows_per_chunk = max(1, _LOGITS_WORKSPACE_BYTES // max(columns * 4, 1)) + row_tiles = not md.is_decode and self.index_heads == 4 for start in range(0, rows, rows_per_chunk): end = min(start + rows_per_chunk, rows) chunk = slice(start, end) logits = self._scratch("logits", end - start, columns, dtype=torch.float32) visible = self._scratch("visible", end - start, dtype=torch.int32) - qsa_mqa_paged( - q_index[chunk], - cmp_pages, - md.block_table, - md.token_to_req[chunk], - positions[chunk], - md.seq_lens, - self.ratio, - logits, - visible, - ) + if row_tiles: + # Prefill: several query rows of one request per program (decode, which a + # graph may capture, keeps the one-row kernel). Same scores bit for bit. + if md.score_tiles is None: + md.score_tiles = {} + tile_map = md.score_tiles.get((start, end)) + if tile_map is None: + tile_map = qsa_score_tile_map( + md.qo_indptr_cpu, start, end, _SCORE_ROWS, self.device + ) + md.score_tiles[(start, end)] = tile_map + qsa_mqa_paged_rows( + q_index[chunk], + cmp_pages, + md.block_table, + tile_map, + md.token_to_req[chunk], + positions[chunk], + md.seq_lens, + self.ratio, + logits, + visible, + rows_per_tile=_SCORE_ROWS, + ) + else: + qsa_mqa_paged( + q_index[chunk], + cmp_pages, + md.block_table, + md.token_to_req[chunk], + positions[chunk], + md.seq_lens, + self.ratio, + logits, + visible, + ) blocks = self._scratch("blocks", end - start, self.block_topk, dtype=torch.int32) self._top_blocks(logits, visible, blocks) expand_qsa_block_indices( @@ -518,6 +624,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..004e00fe4a 100644 --- a/python/freetoken/attention/triton.py +++ b/python/freetoken/attention/triton.py @@ -5,6 +5,7 @@ import torch from freetoken.core import Batch, get_global_ctx +from freetoken.kernel.triton.attention import decode_kv_split_cap from .base import AttentionSpec, BaseAttnBackend, BaseAttnMetadata from .utils import BaseCaptureData @@ -16,9 +17,9 @@ @dataclass class TritonCaptureData(BaseCaptureData): q_to_req: torch.Tensor + # flat decode split scratch: each graph batch size views it with its own split count attn_logits: torch.Tensor attn_lse: torch.Tensor - num_kv_splits: torch.Tensor swa_page_table: torch.Tensor | None = None @classmethod @@ -30,7 +31,7 @@ def create( *, num_q_heads: int, max_head_dim: int, - max_kv_splits: int, + scratch_rows: int, **kwargs, ): return cls( @@ -40,22 +41,8 @@ def create( cu_seqlens_q=torch.arange(0, max_bs + 1, dtype=torch.int32, device=device), page_table=torch.zeros((max_bs, max_seq_len), dtype=torch.int32, device=device), q_to_req=torch.arange(max_bs, dtype=torch.int32, device=device), - attn_logits=torch.empty( - (max_bs, num_q_heads, max_kv_splits, max_head_dim), - dtype=torch.float32, - device=device, - ), - attn_lse=torch.empty( - (max_bs, num_q_heads, max_kv_splits), - dtype=torch.float32, - device=device, - ), - num_kv_splits=torch.full( - (max_bs,), - max_kv_splits, - dtype=torch.int32, - device=device, - ), + attn_logits=torch.empty((scratch_rows * num_q_heads * max_head_dim,), dtype=torch.float32, device=device), + attn_lse=torch.empty((scratch_rows * num_q_heads,), dtype=torch.float32, device=device), **kwargs, ) @@ -72,7 +59,6 @@ class TritonMetadata(BaseAttnMetadata): max_q_len: int attn_logits: torch.Tensor | None = None attn_lse: torch.Tensor | None = None - num_kv_splits: torch.Tensor | None = None swa_indices: torch.Tensor | None = None def get_last_indices(self, bs: int) -> torch.Tensor: @@ -87,7 +73,6 @@ def __init__(self, config: ModelConfig): self.capture: TritonCaptureData | None = None self.capture_bs: List[int] = [] self.max_graph_bs = 0 - self.max_kv_splits = 8 self.prefill_tile_min_q = 128 self.num_q_heads = int(getattr(config, "num_qo_heads", 1)) kv_groups = getattr(config, "kv_cache_group_specs", lambda: ())() @@ -95,6 +80,11 @@ def __init__(self, config: ModelConfig): (group.head_dim for group in kv_groups), default=int(getattr(config, "head_dim", 1)), ) + self.kv_heads = sorted({group.num_kv_heads for group in kv_groups}) or [int(getattr(config, "num_kv_heads", 1))] + + def _split_cap(self, bs: int) -> int: + """The most decode KV splits any layer group takes at this batch size.""" + return max(decode_kv_split_cap(bs, max(1, self.num_q_heads), heads, self.device.index) for heads in self.kv_heads) def _ensure_decode_scratch( self, @@ -103,32 +93,27 @@ def _ensure_decode_scratch( num_q_heads: int, head_dim: int, ) -> None: + splits = self._split_cap(bs) if ( metadata.attn_logits is not None and metadata.attn_lse is not None - and metadata.num_kv_splits is not None and metadata.attn_logits.shape[0] >= bs and metadata.attn_logits.shape[1] >= num_q_heads + and metadata.attn_logits.shape[2] >= splits and metadata.attn_logits.shape[3] >= head_dim ): return scratch_head_dim = max(self.max_head_dim, head_dim) metadata.attn_logits = torch.empty( - (bs, num_q_heads, self.max_kv_splits, scratch_head_dim), + (bs, num_q_heads, splits, scratch_head_dim), dtype=torch.float32, device=self.device, ) metadata.attn_lse = torch.empty( - (bs, num_q_heads, self.max_kv_splits), + (bs, num_q_heads, splits), dtype=torch.float32, device=self.device, ) - metadata.num_kv_splits = torch.full( - (bs,), - self.max_kv_splits, - dtype=torch.int32, - device=self.device, - ) def forward( self, @@ -151,10 +136,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 @@ -167,7 +161,6 @@ def forward( self._ensure_decode_scratch(metadata, bs, q.shape[1], q.shape[-1]) assert metadata.attn_logits is not None assert metadata.attn_lse is not None - assert metadata.num_kv_splits is not None return decode_paged_attention( q=q, k_cache=k_cache, @@ -177,11 +170,14 @@ def forward( q_positions=metadata.q_positions, attn_logits=metadata.attn_logits[:bs], attn_lse=metadata.attn_lse[:bs], - num_kv_splits=metadata.num_kv_splits[:bs], - max_kv_splits=self.max_kv_splits, 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 +198,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 +218,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: @@ -230,17 +236,16 @@ def prepare_metadata(self, batch: Batch) -> None: cached_lens = [req.cached_len for req in reqs] num_query_tokens = sum(seqlens_q) is_decode = max(seqlens_q) == 1 - prefix_lens = torch.tensor(cached_lens, dtype=torch.int32, device=device) + CPU_KWARGS = {"device": "cpu", "dtype": torch.int32, "pin_memory": device.type == "cuda"} + prefix_lens = torch.tensor(cached_lens, **CPU_KWARGS).to(device, non_blocking=True) - indptr = torch.tensor([0] + seqlens_k, dtype=torch.int32, device=device).cumsum_(0) + indptr = torch.tensor([0] + seqlens_k, **CPU_KWARGS).cumsum_(0).to(device, non_blocking=True) if is_decode: cu_seqlens_q_gpu = torch.arange(0, padded_size + 1, device=device, dtype=torch.int32) elif all(l == 0 for l in cached_lens): cu_seqlens_q_gpu = indptr else: - cu_seqlens_q_gpu = torch.tensor( - [0] + seqlens_q, dtype=torch.int32, device=device - ).cumsum_(0) + cu_seqlens_q_gpu = torch.tensor([0] + seqlens_q, **CPU_KWARGS).cumsum_(0).to(device, non_blocking=True) indices = torch.cat([page_table[req.table_idx, : req.device_len] for req in reqs]) swa_indices = None if getattr(self.kvcache, "swa_paged", False): @@ -281,7 +286,7 @@ def init_capture_graph(self, max_seq_len: int, bs_list: List[int]) -> None: self.device, num_q_heads=max(1, self.num_q_heads), max_head_dim=max(1, self.max_head_dim), - max_kv_splits=self.max_kv_splits, + scratch_rows=max(bs * self._split_cap(bs) for bs in bs_list), ) if self._swa_capture_enabled(): self.capture.swa_page_table = torch.zeros( @@ -325,14 +330,20 @@ def _point_to_capture(self, metadata: TritonMetadata, bs: int) -> None: metadata.indices = indices metadata.q_to_req = self.capture.q_to_req[: metadata.q_to_req.numel()] metadata.q_positions = self.capture.positions[:q_tokens] - metadata.attn_logits = self.capture.attn_logits[:bs] - metadata.attn_lse = self.capture.attn_lse[:bs] - metadata.num_kv_splits = self.capture.num_kv_splits[:bs] + metadata.attn_logits, metadata.attn_lse = self._capture_scratch(bs) + + def _capture_scratch(self, bs: int) -> tuple[torch.Tensor, torch.Tensor]: + assert self.capture is not None + splits, heads, dim = self._split_cap(bs), max(1, self.num_q_heads), max(1, self.max_head_dim) + logits = self.capture.attn_logits[: bs * heads * splits * dim].view(bs, heads, splits, dim) + lse = self.capture.attn_lse[: bs * heads * splits].view(bs, heads, splits) + return logits, lse def prepare_for_capture(self, batch: Batch) -> None: bs = batch.size assert bs in self.capture_bs and self.capture is not None capture = self.capture + attn_logits, attn_lse = self._capture_scratch(bs) batch.attn_metadata = TritonMetadata( cu_seqlens_q_gpu=capture.cu_seqlens_q[: bs + 1], indptr=capture.cu_seqlens_k[: bs + 1], @@ -342,9 +353,8 @@ def prepare_for_capture(self, batch: Batch) -> None: is_decode=True, prefix_lens=capture.positions[:bs], max_q_len=1, - attn_logits=capture.attn_logits[:bs], - attn_lse=capture.attn_lse[:bs], - num_kv_splits=capture.num_kv_splits[:bs], + attn_logits=attn_logits, + attn_lse=attn_lse, swa_indices=self._capture_swa_indices(), ) 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..da5aa629f9 100644 --- a/python/freetoken/checkpoint/convert.py +++ b/python/freetoken/checkpoint/convert.py @@ -103,6 +103,25 @@ def _copy_metadata(model_path: str, out_dir: str) -> list[str]: return copied +def _refuse_occupied_output(out_dir: str) -> None: + """Refuse an output dir that already holds FTW output: the index of a finished conversion, + or shards left by an interrupted one. There is no resume; the writer would open shard + 00000 with "wb" and a stale index could then describe the wrong bytes.""" + from .ftw import INDEX_NAME, is_ftw_checkpoint + + if is_ftw_checkpoint(out_dir): + found = INDEX_NAME + else: + shards = glob.glob(os.path.join(glob.escape(out_dir), "freetoken-*.ftw")) + if not shards: + return + found = os.path.basename(min(shards)) + raise SystemExit( + f"{out_dir} already holds an FTW checkpoint ({found}); " + f"convert into an empty directory or remove the old one first" + ) + + class _ConvertSink: """Layer-completion sink for ``load_expert_banks(layer_sink=...)``: writes each completed layer's banks as their own FTW entries immediately (name @@ -184,6 +203,7 @@ def convert_checkpoint( if is_ftw_checkpoint(model_path): raise SystemExit(f"{model_path} is already an FTW checkpoint") + _refuse_occupied_output(out_dir) tp = try_get_tp_info() if tp is None: set_tp_info(rank=0, size=1) @@ -200,15 +220,17 @@ 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 + method = layer_methods = None if offload: - from freetoken.engine.engine import offload_expert_method + from freetoken.engine.engine import offload_expert_methods from freetoken.layers import set_rope_device set_rope_device(dev) - method = offload_expert_method(cfg) + object.__setattr__(cfg, "moe_strategy", "offload") + method, layer_methods = offload_expert_methods(cfg) from freetoken.utils.progress import byte_bar, count_bar @@ -237,6 +259,7 @@ def convert_checkpoint( # 2) offload expert banks (post-repack) + alpha scales (slow path auto-picks parallel/serial) quant_format = None + layer_quant_formats = None num_layers = None if offload: # every method-packed format streams each layer to its own FTW entry as it completes (via the sink); the GGUF provider reports through ExpertBanks.streamed whether it engaged the sink or materialized the whole bank set first @@ -245,9 +268,11 @@ def convert_checkpoint( names = legacy_bank_names(legacy_format_for(method.kind, method.kernel.name)) if method is not None else {} sink = _ConvertSink(writer, names=names) banks = load_expert_banks( - model_path, mc, method=method, device=dev, dtype=dtype, layer_sink=sink + model_path, mc, method=method, device=dev, layer_sink=sink, layer_methods=layer_methods ) quant_format = banks.quant_format + layer_quant_formats = banks.layer_formats + assert layer_quant_formats is None or banks.streamed, "experts of several formats always stream to the sink" if banks.streamed: sink.close() num_layers = sink.num_layers # however many distinct layers the sink actually saw @@ -311,6 +336,9 @@ def convert_checkpoint( # read back at load (ftw.load_ftw_banks). dtype/moe_backend were dropped: each # tensor already carries its own dtype, and nothing reads a model-level backend. "quant_format": quant_format, + # Experts of several formats: each MoE layer's quant_format (quant_format above is the most common). A layer only + # stores the banks its kernel reads; the reader gives it zero-byte placeholders for the others. + "layer_quant_formats": layer_quant_formats, # The reader takes num_layers from the model config (copied into this # checkpoint); recording it here too gives load_ftw_banks a cross-check that # the banks match the config they ship with. None for non-offload checkpoints. diff --git a/python/freetoken/checkpoint/ftw.py b/python/freetoken/checkpoint/ftw.py index e2bd0de24e..c7068a587c 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,330 @@ 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_nvfp4_global(name: str) -> bool: + m = _LAYER_ENTRY_RE.match(name) + return (m["base"] if m else name).endswith(("gate_up_global", "down_global", ".weight_global")) + + +# fork-only: FTW dirs already warned about fp16 NVFP4 globals (one reader is opened per load stage) +_FP16_GLOBALS_WARNED: set[str] = set() + + 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 +466,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 +510,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 +522,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 +550,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 +566,31 @@ 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} + # fork-only: an FTW converted before #739 keeps NVFP4 globals in fp16; the read paths upcast them to fp32 + self.fp16_globals = {n for n, t in self.tensors.items() if t["dtype"] == "float16" and _is_nvfp4_global(n)} + if self.fp16_globals and os.path.realpath(path) not in _FP16_GLOBALS_WARNED: + _FP16_GLOBALS_WARNED.add(os.path.realpath(path)) + logger.warning( + f"{path!r} stores {len(self.fp16_globals)} NVFP4 global scale entries as fp16 " + f"(e.g. {min(self.fp16_globals)!r}); upcasting them to fp32 at load, the fp16 rounding stays. " + "Reconvert with `ft checkpoint` for exact globals" + ) 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 +729,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 +774,16 @@ 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 e["name"] in reader.fp16_globals: + t = t.float() + if not _put((e["name"], t, buf, e["nbytes"])): return except BaseException as ex: # surface to consumer err.append(ex) @@ -432,6 +811,54 @@ 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)}") + # experts of several formats: a layer stores only the banks its kernel reads + several = reader.meta("layer_quant_formats") is not None + if several and flat_entries: + raise FTWFormatError("experts of several formats are stored per layer") + for base, by_layer in per_layer_groups.items(): + layers = sorted(by_layer) + if not ((set(layers) <= set(range(num_layers))) if several else layers == list(range(num_layers))): + raise FTWFormatError( + f"FTW bank {base!r} has per-layer entries for layers {layers}, " + f"expected exactly range({num_layers})" + ) + if several: + covered = set().union(*(by_layer.keys() for by_layer in per_layer_groups.values())) + if covered != set(range(num_layers)): + raise FTWFormatError(f"FTW expert banks have no entries for layers {sorted(set(range(num_layers)) - covered)}") + 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 +892,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 +916,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,36 +940,62 @@ 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] = [] + upcast = name in reader.fp16_globals for layer_id in range(num_layers): off = e["global_off"] + layer_id * layer_bytes win_off = (off // ALIGN) * ALIGN win_end = _align_up(off + layer_bytes) head_pad = off - win_off - bank = HostBank((win_end - win_off,), torch.uint8, backing=_backing(layer_id)) + if upcast: + # read the fp16 window into a transient buffer, keep an fp32 bank + bank = HostBank((num_experts, *row_shape), torch.float32, backing=_backing(layer_id)) + row_view_args[name].append(None) + else: + bank = HostBank((win_end - win_off,), torch.uint8, backing=_backing(layer_id)) + row_view_args[name].append((head_pad, layer_bytes, num_experts, tuple(row_shape), dtype)) row_hb[name].append(bank) - row_view_args[name].append((head_pad, layer_bytes, num_experts, tuple(row_shape), dtype)) - row_jobs.append((name, bank, win_off, win_end - win_off, layer_bytes, layer_id)) + row_jobs.append((name, bank, win_off, win_end - win_off, layer_bytes, layer_id, + (head_pad, (num_experts, *row_shape)) if upcast else None)) + num_experts = next((e["shape"][0] for by_layer in per_layer_groups.values() for e in by_layer.values()), None) 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 - bank = HostBank(tuple(e["shape"]), _dtype_of(e["dtype"]), backing=_backing(layer_id)) + e = by_layer.get(layer_id) + if e is None: # experts of several formats: this layer's kernel reads no such bank + row_hb[base].append(None) + row_view_args[base].append("placeholder") + continue + if e["global_off"] % ALIGN: + reader.close() + raise FTWFormatError( + f"FTW bank {base!r} layer {layer_id} is not {ALIGN}-aligned" + ) + upcast = e["name"] in reader.fp16_globals + dtype = torch.float32 if upcast else _dtype_of(e["dtype"]) + bank = HostBank(tuple(e["shape"]), dtype, backing=_backing(layer_id)) row_hb[base].append(bank) row_view_args[base].append(None) layer_jobs.append((base, bank, e, layer_id)) @@ -576,15 +1016,30 @@ def _read_alpha(e): bar.update(e["nbytes"]) def _read_row(job): - _name, bank, win_off, win_len, layer_bytes, layer_id = job - reader.read_into(bank.memoryview(), {"global_off": win_off, "nbytes": win_len}, - workers=workers, chunk=chunk) + _name, bank, win_off, win_len, layer_bytes, layer_id, upcast = job + window = {"global_off": win_off, "nbytes": win_len} + if upcast is None: + reader.read_into(bank.memoryview(), window, workers=workers, chunk=chunk) + else: + head_pad, shape = upcast + buf = _transient_buffer(win_len) + reader.read_into(memoryview(buf), window, workers=workers, chunk=chunk) + raw = torch.frombuffer(buf, dtype=torch.uint8, count=head_pad + layer_bytes)[head_pad:] + with torch.inference_mode(): # the banks are inference tensors; this thread does not inherit the mode + bank.tensor.copy_(raw.view(torch.float16).view(shape)) pins.submit(bank, residency[layer_id]) bar.update(layer_bytes) def _read_layer(job): _name, bank, entry, layer_id = job - reader.read_into(bank.memoryview(), entry, workers=workers, chunk=chunk) + if entry["name"] in reader.fp16_globals: + buf = _transient_buffer(entry["nbytes"]) + reader.read_into(memoryview(buf), entry, workers=workers, chunk=chunk) + raw = torch.frombuffer(buf, dtype=torch.float16, count=entry["nbytes"] // 2) + with torch.inference_mode(): # as in _read_row + bank.tensor.copy_(raw.view(bank.tensor.shape)) + else: + reader.read_into(bank.memoryview(), entry, workers=workers, chunk=chunk) pins.submit(bank, residency[layer_id]) bar.update(entry["nbytes"]) @@ -598,10 +1053,15 @@ def _read_layer(job): bar.close() reader.close() + from freetoken.moe.offload_cache import placeholder_bank + sources: dict[str, list] = {} for name, banks in row_hb.items(): views = [] for bank, view_args in zip(banks, row_view_args[name]): + if view_args == "placeholder": + views.append(placeholder_bank(num_experts)) + continue if view_args is None: # per-layer entry: already shaped [num_experts, ...] views.append(bank.tensor) continue @@ -622,6 +1082,8 @@ def _read_layer(job): applied = list(residency) for banks in row_hb.values(): for layer_id, bank in enumerate(banks): + if bank is None: + continue if (applied[layer_id] == HostResidency.LOCKED.value and bank.residency is not HostResidency.LOCKED): applied[layer_id] = HostResidency.PAGEABLE.value @@ -630,7 +1092,7 @@ def _read_layer(job): by_layer = [0] * num_layers for name, banks in row_hb.items(): for layer_id, bank in enumerate(banks): - by_layer[layer_id] += bank.nbytes + by_layer[layer_id] += bank.nbytes if bank is not None else 0 locked = [i for i in unpinned if applied[i] == HostResidency.LOCKED.value] pageable = [i for i in unpinned if i not in set(locked)] pinned_b = sum(b for i, b in enumerate(by_layer) if i not in set(unpinned)) @@ -653,14 +1115,83 @@ def _read_layer(job): # alphas are the small per-expert scale vectors, distinguished by their reserved names # (not a separate kind); everything else under experts_bank is a weight source. alpha_kw = {n: alpha_hb[n].tensor for n in alpha_hb} + layer_formats = reader.meta("layer_quant_formats") return ExpertBanks( quant_format, sources, **alpha_kw, layer_residency=applied, kind=kind, kernel=kernel, + layer_tags=layer_formats, layer_formats=layer_formats, ) +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 + if reader.meta("layer_quant_formats") is not None: + raise ValueError(f"{path!r} holds experts of several formats, which are served from the offload slot cache only") + 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 host.dtype == torch.float16 and _is_nvfp4_global(name): + host = host.float() + 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/control_cli.py b/python/freetoken/control_cli.py index 50a32aae6e..e33eb4acb8 100644 --- a/python/freetoken/control_cli.py +++ b/python/freetoken/control_cli.py @@ -10,6 +10,8 @@ from collections.abc import Sequence from typing import Any +from freetoken.launch import resolve_api_key + DEFAULT_BASE_URL = "http://127.0.0.1:1919" # (CLI dest, pool, rebuild body key). moe/mamba are slot counts and pass straight through; @@ -74,6 +76,7 @@ def _request_json( body: dict[str, Any] | None = None, query: dict[str, Any] | None = None, timeout: float = 10.0, + api_key: str | None, ) -> dict[str, Any]: with _open_request( method, @@ -83,6 +86,7 @@ def _request_json( query=query, accept="application/json", timeout=timeout, + api_key=api_key, ) as response: raw = response.read() @@ -104,6 +108,7 @@ def _request_sse_text( *, body: dict[str, Any] | None = None, timeout: float = 10.0, + api_key: str | None, ) -> str: with _open_request( method, @@ -112,6 +117,7 @@ def _request_sse_text( body=body, accept="text/event-stream", timeout=timeout, + api_key=api_key, ) as response: return _read_sse_text(response) @@ -125,9 +131,12 @@ def _open_request( query: dict[str, Any] | None = None, accept: str, timeout: float, + api_key: str | None, ) -> Any: data = None headers = {"Accept": accept} + if api_key: + headers["Authorization"] = f"Bearer {api_key}" if body is not None: data = json.dumps(body).encode("utf-8") headers["Content-Type"] = "application/json" @@ -284,6 +293,10 @@ def _build_parser(prog: str) -> argparse.ArgumentParser: default=DEFAULT_BASE_URL, help=f"FreeToken server URL (default: {DEFAULT_BASE_URL})", ) + parser.add_argument( + "--api-key", + help="API key of a server started with --api-key (default: $FREETOKEN_API_KEY)", + ) parser.add_argument("--timeout", type=float, default=10.0, help="HTTP timeout in seconds") parser.add_argument("--json", action="store_true", help="Print raw JSON response") @@ -355,14 +368,19 @@ def main(argv: Sequence[str] | None = None, *, prog: str = "ft ctl") -> int: args = parser.parse_args(list(argv) if argv is not None else None) except SystemExit as exc: return int(exc.code) if isinstance(exc.code, int) else 2 + args.api_key = resolve_api_key(args.api_key) try: if args.command == "health": - doc = _request_json("GET", args.base_url, "/health", timeout=args.timeout) + doc = _request_json( + "GET", args.base_url, "/health", timeout=args.timeout, api_key=args.api_key + ) _print_doc(doc, _format_health, raw_json=args.json) return 0 if args.command == "stats": - doc = _request_json("GET", args.base_url, "/v1/stats", timeout=args.timeout) + doc = _request_json( + "GET", args.base_url, "/v1/stats", timeout=args.timeout, api_key=args.api_key + ) _print_doc(doc, _format_stats, raw_json=args.json) return 0 if args.command == "generate": @@ -376,6 +394,7 @@ def main(argv: Sequence[str] | None = None, *, prog: str = "ft ctl") -> int: "ignore_eos": args.ignore_eos, }, timeout=args.timeout, + api_key=args.api_key, ) if args.json: print(json.dumps({"text": text}, ensure_ascii=False, indent=2, sort_keys=True)) @@ -390,13 +409,16 @@ def main(argv: Sequence[str] | None = None, *, prog: str = "ft ctl") -> int: "/v1/requests", query=query, timeout=args.timeout, + api_key=args.api_key, ) _print_doc(doc, _format_requests, raw_json=args.json) return 0 if args.command == "cache" and args.cache_command in (None, "status"): if args.cache_command is None and _cache_targets(args): return _run_cache_rebuild(args) - doc = _request_json("GET", args.base_url, "/v1/cache/status", timeout=args.timeout) + doc = _request_json( + "GET", args.base_url, "/v1/cache/status", timeout=args.timeout, api_key=args.api_key + ) _print_doc(doc, _format_cache_status, raw_json=args.json) return 0 if args.command == "cache" and args.cache_command == "rebuild": @@ -444,7 +466,9 @@ def _run_cache_rebuild(args: argparse.Namespace) -> int: print("error: cache rebuild requires at least one cache target", file=sys.stderr) return 2 # The token targets can only be converted against the live geometry, so read it first. - status = _request_json("GET", args.base_url, "/v1/cache/status", timeout=args.timeout) + status = _request_json( + "GET", args.base_url, "/v1/cache/status", timeout=args.timeout, api_key=args.api_key + ) body = _rebuild_body(targets, status.get("geometry") or {}) body["timeout"] = args.wait doc = _request_json( @@ -453,6 +477,7 @@ def _run_cache_rebuild(args: argparse.Namespace) -> int: "/v1/cache/rebuild", body=body, timeout=args.wait + args.timeout, + api_key=args.api_key, ) if args.json or doc.get("status") != "ok": _print_doc(doc, _format_rebuild, raw_json=args.json) @@ -462,7 +487,9 @@ def _run_cache_rebuild(args: argparse.Namespace) -> int: # /cache prints. Best-effort -- the rebuild already succeeded. print("status=ok") try: - print(_format_cache_status(_request_json("GET", args.base_url, "/v1/cache/status"))) + print(_format_cache_status( + _request_json("GET", args.base_url, "/v1/cache/status", api_key=args.api_key) + )) except ControlCliError: pass return 0 diff --git a/python/freetoken/core.py b/python/freetoken/core.py index fd691d180c..a375dc9dcf 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,71 @@ class SamplingParams: # Stop strings (OpenAI `stop` / Anthropic `stop_sequences`). Generation finishes when one # appears in the decoded output; the matched substring (and anything after) is trimmed. stop_strs: list[str] = field(default_factory=list) + # ---- logits processors (engine/sample.py), all off by default ---- + # min_p: drop tokens whose probability is below min_p x the top probability. + min_p: float = 0.0 + # OpenAI presence/frequency penalties over the generated tokens; repetition_penalty is the + # HF/vLLM multiplicative penalty over prompt + generated tokens (1.0 = off). The sampler + # keeps their state on the device (Req.output_token_counts, Req.prompt_token_mask). + presence_penalty: float = 0.0 + frequency_penalty: float = 0.0 + repetition_penalty: float = 1.0 + # OpenAI logit_bias as [[token id, bias], ...] (clamped to [-100, 100] by the API). A list, + # not a dict: the messages cross process boundaries as msgpack, which refuses int map keys. + logit_bias: list[list[float]] | None = None + # No EOS / stop token before this many generated tokens (vLLM min_tokens, SGLang + # min_new_tokens). The scheduler fills min_tokens_stop_ids with the ids to mask. + min_tokens: int = 0 + min_tokens_stop_ids: list[int] = field(default_factory=list) + # Extra token ids that end generation with finish_reason "stop" (vLLM stop_token_ids). + stop_token_ids: list[int] = field(default_factory=list) + # Keep the matched stop string in the output instead of trimming it (vLLM + # include_stop_str_in_output, SGLang no_stop_trim). + include_stop_str_in_output: bool = False + # Decode with skip_special_tokens. Off by default here: the reasoning and tool parsers + # consume the /tool tags from the decoded text. + skip_special_tokens: bool = False + # Sampled-token logprobs (OpenAI `logprobs`/`top_logprobs`): when on, the sampler + # reports the chosen token's raw (pre-temperature) logprob and top-k alternatives. + logprobs: bool = False + top_logprobs: int = 0 + # vLLM allowed_token_ids: sampling is masked to these ids (sorted, unique). A request that + # also asks for logprobs gets them renormalized over this set (Sampler.compute_logprobs). + allowed_token_ids: list[int] | None = None + # Set by the tokenizer worker when the rendered prompt leaves a block open, so + # the detokenizer counts an unclosed generation as reasoning tokens. + starts_in_reasoning: bool = False + + def client_token_ids(self) -> list[tuple[str, list[int]]]: + """The token ids a client supplied, by request field (empty ones left out): each + must be checked against the tokenizer before the request reaches the engine.""" + fields = [ + ("allowed_token_ids", self.allowed_token_ids or []), + ("logit_bias", [int(tid) for tid, _ in self.logit_bias or ()]), + ("stop_token_ids", self.stop_token_ids), + ] + return [(name, ids) for name, ids in fields if ids] + + @property + def is_logprob_read(self) -> bool: + """A scoring read (e.g. AnyJev label logprobs): one token whose logprobs are the answer.""" + return bool(self.logprobs) and self.max_tokens == 1 @property def is_greedy(self) -> bool: return self.temperature <= 0.0 or self.top_k == 1 + @property + def needs_logits_processing(self) -> bool: + # the penalties are not LogitsPlan work: the sampler applies them from its own + # device-side state (Sampler.prepare's penalties). + return bool( + self.min_p > 0.0 + or self.logit_bias + or self.min_tokens > 0 + or self.allowed_token_ids + ) + @dataclass(eq=False) class Req: @@ -43,6 +105,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 +128,16 @@ 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 + output_token_counts: torch.Tensor | None = field(default=None, init=False, repr=False) + stop_decode_status: DecodeStatus | None = None + # [vocab] bool on the device: the prompt's ids, for repetition_penalty. Built once. + prompt_token_mask: torch.Tensor | None = field(default=None, init=False, repr=False) def __post_init__(self) -> None: assert self.input_ids.is_cpu 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 +252,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/daemon/README.md b/python/freetoken/daemon/README.md index ea3b016eaf..ebe4a17008 100644 --- a/python/freetoken/daemon/README.md +++ b/python/freetoken/daemon/README.md @@ -60,9 +60,9 @@ Target a non-default daemon with `--url http://host:1900` (or `$FREETOKEN_DAEMON | Method / path | Notes | | --- | --- | | `GET /health` | Daemon self-health; always answers, never gated by `--token`. | -| `POST /engine/start` `{model,port,args[]}` | Idempotent on the full `(model,port,args)`; a differing config on the same port → `409`. | +| `POST /engine/start` `{model,port,args[],apiKey?}` | Idempotent on the full `(model,port,args,apiKey)`; a differing config on the same port → `409`. | | `POST /engine/stop` `{force?:false}` | Close admission, drain/abort, durably enqueue the final-accounting receipt, then `SIGTERM`→grace→`SIGKILL`. A prepare/outbox failure preserves the engine. | -| `POST /engine/switch` `{model,port,args[],force?:false}` | One serialized stop-accounting-start transaction. | +| `POST /engine/switch` `{model,port,args[],apiKey?,force?:false}` | One serialized stop-accounting-start transaction. | | `GET /engine/status` | `{running,pid,model,port,uptimeS,lastExitCode,…}`; outlives any single serve. | | `GET /engine/logs?since=` | SSE, ANSI-stripped, tqdm-`\r` collapsed, ring replay, `id:`, `Last-Event-ID` resume. | | `GET /engine/metrics` | `{ramBytes,vramBytes}` — the serve tree's own footprint only. | @@ -75,6 +75,13 @@ Target a non-default daemon with `--url http://host:1900` (or `$FREETOKEN_DAEMON Set `--token` (or `$FREETOKEN_DAEMON_TOKEN`) to require an `X-FT-Token` header on everything except `/health`. +`apiKey` is the serve's own `--api-key`. The daemon hands it to the serve as `FREETOKEN_API_KEY` in +its environment (never on its command line), sends it on its own calls to the serve, and keeps it in +the owner-only state file so a re-adopted serve stays reachable. The daemon's own +`FREETOKEN_API_KEY` is never passed down: a serve requires a key only when its start request carries +one, and an `--api-key` inside `args` is rejected with a 400 because the daemon would not know it. +`ft daemon start|switch --api-key KEY` (default `$FREETOKEN_API_KEY`) sets it. + The daemon reaches the serve's destructive `POST /v1/admin/prepare-stop` endpoint only over loopback; the serve rejects non-loopback callers even when its inference API is bound to `0.0.0.0`. `force` is never implicit: it is an explicit recovery choice that may record null diff --git a/python/freetoken/daemon/app.py b/python/freetoken/daemon/app.py index d7a0a53c60..dd506f78fc 100644 --- a/python/freetoken/daemon/app.py +++ b/python/freetoken/daemon/app.py @@ -30,6 +30,7 @@ class StartBody(BaseModel): model: str port: int | None = None args: list[str] = [] + apiKey: str | None = None class StopBody(BaseModel): @@ -161,6 +162,15 @@ def resolve_port(explicit: int | None) -> int: st = manager.status() return st.get("port") or default_serve_port + def reject_api_key_arg(args: list[str]) -> None: + # the daemon must know the key to reach the serve, so the key only travels as apiKey + # argparse also accepts any unambiguous prefix, and no other serve flag starts with --ap + flags = (a.split("=", 1)[0] for a in args) + if any(len(f) >= 4 and "--api-key".startswith(f) for f in flags): + raise HTTPException( + status_code=400, detail="pass the serve's API key as apiKey, not in args" + ) + def accounting_error(exc: Exception) -> JSONResponse: code = ( "accounting_outbox_failed" @@ -192,9 +202,17 @@ async def health(): @app.post("/engine/start", dependencies=auth) async def engine_start(body: StartBody): + reject_api_key_arg(body.args) port = resolve_port(body.port) try: - return await run(lifecycle_pool, manager.start, body.model, port, list(body.args)) + return await run( + lifecycle_pool, + manager.start, + body.model, + port, + list(body.args), + body.apiKey or None, + ) except Conflict as exc: st = manager.status() return JSONResponse( @@ -236,6 +254,7 @@ async def shutdown_daemon(request: Request, body: StopBody | None = None): @app.post("/engine/switch", dependencies=auth) async def engine_switch(body: SwitchBody): + reject_api_key_arg(body.args) port = resolve_port(body.port) try: return await run( @@ -245,6 +264,7 @@ async def engine_switch(body: SwitchBody): port, list(body.args), body.force, + body.apiKey or None, ) except (AccountingPrepareError, AccountingOutboxError) as exc: return accounting_error(exc) diff --git a/python/freetoken/daemon/client.py b/python/freetoken/daemon/client.py index 6616778e72..9ded96276e 100644 --- a/python/freetoken/daemon/client.py +++ b/python/freetoken/daemon/client.py @@ -16,6 +16,8 @@ from collections.abc import Sequence from typing import Any +from freetoken.launch import resolve_api_key + DEFAULT_URL = "http://127.0.0.1:1900" DEFAULT_TIMEOUT = 10.0 # prepare-stop (15s transport budget) + default SIGTERM grace (10s) + reap wait (10s), @@ -144,6 +146,10 @@ def _build_parser(prog: str) -> argparse.ArgumentParser: sp = sub.add_parser(name, parents=[common], help=f"POST /engine/{name}") sp.add_argument("model", help="Model path/id") sp.add_argument("--port", type=int, default=None, help="Serve port") + sp.add_argument( + "--api-key", + help="API key the serve requires, set in its environment (default: $FREETOKEN_API_KEY)", + ) if name == "switch": sp.add_argument( "--force", @@ -181,6 +187,8 @@ def main(argv: Sequence[str] | None = None, *, prog: str = "ft daemon") -> int: body: dict[str, Any] = {"model": args.model, "args": list(args.serve_args)} if args.port is not None: body["port"] = args.port + if api_key := resolve_api_key(args.api_key): + body["apiKey"] = api_key if args.verb == "switch" and args.force: body["force"] = True method, path = "POST", f"/engine/{args.verb}" diff --git a/python/freetoken/daemon/pidfile.py b/python/freetoken/daemon/pidfile.py index 89ffc2b4df..b9865a83d3 100644 --- a/python/freetoken/daemon/pidfile.py +++ b/python/freetoken/daemon/pidfile.py @@ -4,15 +4,16 @@ open for the daemon's whole life. flock (not a bare port-bind) avoids TIME_WAIT races and is released automatically if the daemon dies, so a crashed daemon never wedges its own restart. -``ServeStateStore`` persists ``{model, port, pid, args, starttime, logPath}`` as JSON on every +``ServeStateStore`` persists ``{model, port, pid, args, starttime, log_path, api_key}`` as JSON on every lifecycle change, so a restarted daemon can re-adopt a still-running serve. ``starttime`` + -``args`` are what make re-adoption PID-reuse-safe and config-exact.""" +``args`` are what make re-adoption PID-reuse-safe and config-exact. +The file is owner-only because it holds the serve's API key.""" from __future__ import annotations import json import os -from dataclasses import asdict, dataclass +from dataclasses import asdict, dataclass, field class AlreadyRunning(RuntimeError): @@ -63,6 +64,7 @@ class ServeState: args: list[str] starttime: int | None = None log_path: str | None = None + api_key: str | None = field(default=None, repr=False) class ServeStateStore: @@ -72,7 +74,12 @@ def __init__(self, path: str) -> None: def save(self, state: ServeState) -> None: os.makedirs(os.path.dirname(self.path) or ".", exist_ok=True) tmp = f"{self.path}.tmp" - with open(tmp, "w") as fh: + # unlink first: O_CREAT applies the 0o600 mode only to a file it creates + try: + os.remove(tmp) + except FileNotFoundError: + pass + with os.fdopen(os.open(tmp, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600), "w") as fh: json.dump(asdict(state), fh) fh.flush() os.fsync(fh.fileno()) @@ -102,6 +109,7 @@ def load(self) -> ServeState | None: args=list(doc.get("args") or []), starttime=doc.get("starttime"), log_path=doc.get("log_path"), + api_key=doc.get("api_key"), ) except (KeyError, TypeError, ValueError): return None diff --git a/python/freetoken/daemon/proxy.py b/python/freetoken/daemon/proxy.py index 9e382a4d72..c4f0ba2864 100644 --- a/python/freetoken/daemon/proxy.py +++ b/python/freetoken/daemon/proxy.py @@ -49,6 +49,7 @@ def __init__( now: Callable[[], float] = time.monotonic, opener: Callable[[str, float], dict] | None = None, prepare_opener: Callable[[str, float], dict] | None = None, + api_key: Callable[[], str | None] = lambda: None, ) -> None: self._host = host self._ttl = ttl_s @@ -57,6 +58,7 @@ def __init__( self._now = now self._opener = opener or self._urlopen self._prepare_opener = prepare_opener or self._urlopen_prepare + self._api_key = api_key self._lock = threading.Lock() self._cache: dict[tuple[str, int], tuple[float, dict]] = {} @@ -109,19 +111,23 @@ def _fetch(self, path: str, port: int) -> dict: result["reachable"] = True return result - @staticmethod - def _urlopen(url: str, timeout: float) -> dict: - req = urllib.request.Request(url, headers={"Accept": "application/json"}, method="GET") + def _headers(self) -> dict[str, str]: + headers = {"Accept": "application/json"} + if key := self._api_key(): + headers["Authorization"] = f"Bearer {key}" + return headers + + def _urlopen(self, url: str, timeout: float) -> dict: + req = urllib.request.Request(url, headers=self._headers(), method="GET") with urllib.request.urlopen(req, timeout=timeout) as resp: raw = resp.read() return json.loads(raw.decode("utf-8")) - @staticmethod - def _urlopen_prepare(url: str, timeout: float) -> dict: + def _urlopen_prepare(self, url: str, timeout: float) -> dict: req = urllib.request.Request( url, data=b"{}", - headers={"Accept": "application/json", "Content-Type": "application/json"}, + headers={**self._headers(), "Content-Type": "application/json"}, method="POST", ) with urllib.request.urlopen(req, timeout=timeout) as resp: diff --git a/python/freetoken/daemon/serve_manager.py b/python/freetoken/daemon/serve_manager.py index 681bdef497..9d5e51fe71 100644 --- a/python/freetoken/daemon/serve_manager.py +++ b/python/freetoken/daemon/serve_manager.py @@ -26,6 +26,8 @@ from dataclasses import dataclass from typing import Any, Callable +from freetoken.launch import API_KEY_ENV + from . import osproc from .accounting import ( AccountingOutbox, @@ -37,7 +39,7 @@ class Conflict(RuntimeError): - """A different serve (model/port/args) is already running; the client should switch().""" + """A different serve (model/port/args/key) is already running; the client should switch().""" @dataclass @@ -128,7 +130,14 @@ def build_serve_command( return argv, log_path -def spawn_serve(argv: list[str], log_path: str) -> PopenChild: +def serve_env(api_key: str | None) -> dict[str, str]: + env = {k: v for k, v in os.environ.items() if k != API_KEY_ENV} + if api_key: + env[API_KEY_ENV] = api_key + return env + + +def spawn_serve(argv: list[str], log_path: str, api_key: str | None = None) -> PopenChild: """Spawn the serve as a session leader with its stdout+stderr going to a real logfile fd (never a PIPE): the file survives daemon death, so a re-adopting daemon just resumes tailing it and the serve never writes to a dead pipe. ``start_new_session`` makes it a process @@ -138,6 +147,7 @@ def spawn_serve(argv: list[str], log_path: str) -> PopenChild: try: proc = subprocess.Popen( argv, + env=serve_env(api_key), stdout=logf, stderr=subprocess.STDOUT, stdin=subprocess.DEVNULL, @@ -158,7 +168,7 @@ def __init__( ring, state_store, *, - spawn_fn: Callable[[str, int, list[str]], object] | None = None, + spawn_fn: Callable[[str, int, list[str], str | None], object] | None = None, adopt_fn: Callable[[object], object] | None = None, tailer_factory: Callable[[object], object] | None = None, signal_fn: Callable[[int, int], None] | None = None, @@ -217,6 +227,7 @@ def __init__( self._model: str | None = None self._port: int | None = None self._args: list[str] = [] + self._api_key: str | None = None self._started_at: float | None = None self._last_exit: ExitInfo | None = None self._adopted = False @@ -232,11 +243,11 @@ def __init__( # ---- default factories (overridable for tests) ---- - def _default_spawn(self, model: str, port: int, args: list[str]): + def _default_spawn(self, model: str, port: int, args: list[str], api_key: str | None): argv, log_path = build_serve_command( model, port, args, python=self._python, log_dir=self._log_dir ) - return spawn_serve(argv, log_path) + return spawn_serve(argv, log_path, api_key) def _default_adopt(self, state): if not osproc.is_ft_serve_on_port(state.pid, state.port, starttime=state.starttime): @@ -256,13 +267,25 @@ def _default_adopt(self, state): # ---- lifecycle ---- def start( - self, model: str, port: int, args: list[str] | None = None, *, _auto: bool = False + self, + model: str, + port: int, + args: list[str] | None = None, + api_key: str | None = None, + *, + _auto: bool = False, ) -> dict: with self._lifecycle: - return self._start(model, port, args, _auto=_auto) + return self._start(model, port, args, api_key, _auto=_auto) def _start( - self, model: str, port: int, args: list[str] | None = None, *, _auto: bool = False + self, + model: str, + port: int, + args: list[str] | None = None, + api_key: str | None = None, + *, + _auto: bool = False, ) -> dict: args = list(args or []) with self._cond: @@ -278,10 +301,13 @@ def _start( self._stop_requested = False # an explicit client start clears a prior stop intent while True: if self._child is not None and not self._stopping: - if (self._model, self._port, self._args) == (model, port, args): + if (self._model, self._port, self._args, self._api_key) == ( + model, port, args, api_key + ): return {"pid": self._child.pid, "idempotent": True} raise Conflict( - f"serve already running (model={self._model!r} port={self._port}); " + f"serve already running with a different config " + f"(model={self._model!r} port={self._port}); " f"use switch to replace it" ) if self._starting: @@ -299,7 +325,7 @@ def _start( child = None try: - child = self._spawn_fn(model, port, args) + child = self._spawn_fn(model, port, args, api_key) finally: if child is None: with self._cond: @@ -311,6 +337,7 @@ def _start( self._model = model self._port = port self._args = args + self._api_key = api_key self._started_at = self._now() self._last_exit = None self._adopted = getattr(child, "adopted", False) @@ -318,7 +345,7 @@ def _start( self._starting = False self._cond.notify_all() - self._persist(child, model, port, args) + self._persist(child, model, port, args, api_key) self._maybe_apply_oom(child.pid) self._begin_watch(child) self._emit(f"serve started (pid={child.pid} model={model} port={port})") @@ -407,10 +434,11 @@ def switch( port: int, args: list[str] | None = None, force: bool = False, + api_key: str | None = None, ) -> dict: with self._lifecycle: stopped = self._stop(force=force) - started = self._start(model, port, args) + started = self._start(model, port, args, api_key) return {**started, "accounting": stopped["accounting"]} def pending_accounting(self) -> list[dict[str, Any]]: @@ -734,6 +762,7 @@ def readopt(self) -> bool: self._model = state.model self._port = state.port self._args = list(state.args) + self._api_key = state.api_key self._started_at = self._now() # unknown true start; report uptime since adoption self._last_exit = None self._adopted = True @@ -824,13 +853,13 @@ def _exit_line(self, child, info: ExitInfo, was_stopping: bool) -> str: return f"serve exited with code {info.code} (pid={child.pid})" def _restart_async(self, dead_child) -> None: - model, port, args = self._model, self._port, list(self._args) + model, port, args, api_key = self._model, self._port, list(self._args), self._api_key if model is None or port is None: return def _run(): try: - self.start(model, port, args, _auto=True) + self.start(model, port, args, api_key, _auto=True) except Exception as exc: # noqa: BLE001 self._emit(f"auto-restart failed: {exc}") @@ -867,6 +896,10 @@ def serve_args(self) -> list[str]: with self._cond: return list(self._args) + def serve_api_key(self) -> str | None: + with self._cond: + return self._api_key + def status(self) -> dict: with self._cond: child = self._child @@ -892,7 +925,9 @@ def status(self) -> dict: # ---- helpers ---- - def _persist(self, child, model: str, port: int, args: list[str]) -> None: + def _persist( + self, child, model: str, port: int, args: list[str], api_key: str | None + ) -> None: from .pidfile import ServeState self._store.save( @@ -903,6 +938,7 @@ def _persist(self, child, model: str, port: int, args: list[str]) -> None: args=list(args), starttime=getattr(child, "starttime", None), log_path=getattr(child, "log_path", None), + api_key=api_key, ) ) diff --git a/python/freetoken/daemon/server.py b/python/freetoken/daemon/server.py index d4e742b0be..76e7e789c9 100644 --- a/python/freetoken/daemon/server.py +++ b/python/freetoken/daemon/server.py @@ -132,7 +132,8 @@ def main(argv: Sequence[str] | None = None, *, prog: str = "ft daemon") -> int: ring = LogRing(capacity=args.log_capacity) store = ServeStateStore(os.path.join(state_dir, "serve.json")) - probe = ServeProbe() + # read at call time: the running serve's key, from the manager created below + probe = ServeProbe(api_key=lambda: manager.serve_api_key()) def tailer_factory(child): log_path = getattr(child, "log_path", None) diff --git a/python/freetoken/engine/cache_budget.py b/python/freetoken/engine/cache_budget.py index d522be4991..9e824f972b 100644 --- a/python/freetoken/engine/cache_budget.py +++ b/python/freetoken/engine/cache_budget.py @@ -14,8 +14,8 @@ import torch -def expert_bytes_per_slot(sources: dict[str, "list[torch.Tensor]"]) -> int: - """Bytes one expert slot occupies on GPU: summed row bytes over all banks. +def layer_slot_bytes(sources: dict[str, "list[torch.Tensor]"]) -> list[int]: + """Bytes one expert slot of each layer occupies on GPU: its row bytes summed over all banks. Each bank source is per-layer ``[num_experts, *row_shape]`` tensors and is already TP-sharded upstream, so the per-row byte count is the per-rank slot @@ -24,8 +24,15 @@ 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. - return sum(t[0][0].numel() * t[0].element_size() for t in sources.values()) + # slot-byte idiom in kvcache/linear_state_pool.py and kvcache/dsv4/v4_pool.py. + num_layers = len(next(iter(sources.values()))) + return [sum(t[l][0].numel() * t[l].element_size() for t in sources.values()) for l in range(num_layers)] + + +def expert_bytes_per_slot(sources: dict[str, "list[torch.Tensor]"]) -> int: + """Bytes one expert slot occupies on GPU, rounded up to the mean over layers when their bank layouts differ.""" + per_layer = layer_slot_bytes(sources) + return -(-sum(per_layer) // len(per_layer)) def net_cache_budget_bytes( @@ -54,6 +61,7 @@ def plan_cache_budget( prefill_overlap: bool, kv_reserve_pages: int, max_slots: int, + page_size: int = 1, ) -> tuple[int, int, bool]: """Split ``budget_bytes`` MoE-first into (moe_cache_size, num_pages, prefill_overlap). @@ -61,7 +69,8 @@ def plan_cache_budget( weights + fixed_cache_size; the (1-memory_ratio) remainder is the graph headroom). Experts greedily fill the budget after reserving ``kv_reserve_pages`` for KV, clamped to ``[floor, min(total_experts, max_slots)]`` (floor is ``2*num_experts`` when prefill - overlap is feasible else ``num_experts``); KV pages take whatever remains. + overlap is feasible, in slots and in bytes, else ``num_experts``); KV pages take whatever + remains. ``page_size`` (tokens per page) only feeds the error message. """ assert per_expert_bytes > 0, "per_expert_bytes must be positive" assert cache_per_page > 0, "cache_per_page must be positive (owned-KV models unsupported here)" @@ -74,6 +83,10 @@ def plan_cache_budget( assert hi >= lo, f"slot cap {hi} below the minimum {lo} slots" kv_reserve_bytes = kv_reserve_pages * cache_per_page + # Overlap is a prefill speedup, not a requirement: when its 2*num_experts floor leaves no + # room for the KV reserve (and the 2 pages asserted below), plan without it. + if overlap and 2 * num_experts * per_expert_bytes + max(kv_reserve_pages, 2) * cache_per_page > budget_bytes: + overlap, lo = False, num_experts # MoE-priority: reserve KV first, then experts greedily take the remaining budget. raw = (budget_bytes - kv_reserve_bytes) // per_expert_bytes moe_cache_size = max(lo, min(raw, hi)) @@ -86,10 +99,13 @@ def plan_cache_budget( # the reserve (or negative), yielding a plan that exceeds budget_bytes. Reject here so # --moe-cache-auto fails in arithmetic instead of OOMing in a later CUDA allocation. total = moe_cache_size * per_expert_bytes + num_pages * cache_per_page + fit_pages = max(0, (budget_bytes - moe_cache_size * per_expert_bytes) // cache_per_page) assert total <= budget_bytes, ( f"cache budget too small: minimum plan (moe={moe_cache_size} slots, " - f"kv={num_pages} pages) needs {total} B > budget {budget_bytes} B " - "(raise memory_ratio, lower kv_reserve_tokens, or free GPU memory)" + f"kv={num_pages} pages) needs {total / 2**30:.2f} GiB > budget {budget_bytes / 2**30:.2f} GiB; " + f"beside {moe_cache_size} slots at most {fit_pages} KV pages ({fit_pages * page_size} tokens) fit. " + "Lower --kv-reserve-tokens, use a smaller --kv-cache-dtype, raise --memory-ratio, " + "or set --moe-cache-size and --max-seq-len-override explicitly" ) assert num_pages > 1, "not enough memory for KV cache after MoE allocation" return moe_cache_size, num_pages, overlap @@ -120,7 +136,9 @@ def resolve_moe_cache_auto( """ budget_bytes = net_cache_budget_bytes(memory_ratio, baseline_free, weights_bytes, fixed_cache_size) max_slots = total_experts if max_slots is None else min(max_slots, total_experts) - kv_reserve_pages = div_ceil(kv_reserve_tokens, page_size) + # Every pool keeps page 0 as an unreachable dummy/sentinel. The CLI floor is expressed in + # usable tokens, so reserve that internal page in addition to the user-visible capacity. + kv_reserve_pages = div_ceil(kv_reserve_tokens, page_size) + 1 return plan_cache_budget( budget_bytes=budget_bytes, per_expert_bytes=per_expert_bytes, @@ -130,4 +148,5 @@ def resolve_moe_cache_auto( prefill_overlap=prefill_overlap, kv_reserve_pages=kv_reserve_pages, max_slots=max_slots, + page_size=page_size, ) diff --git a/python/freetoken/engine/config.py b/python/freetoken/engine/config.py index 44c5a38e73..53b28c3ef1 100644 --- a/python/freetoken/engine/config.py +++ b/python/freetoken/engine/config.py @@ -1,9 +1,10 @@ from __future__ import annotations import copy +import os from dataclasses import dataclass, field, replace from functools import cached_property -from typing import TYPE_CHECKING, List +from typing import TYPE_CHECKING, Any, List, Mapping import torch from freetoken.distributed import DistributedInfo @@ -23,6 +24,8 @@ class EngineConfig: model_path: str tp_info: DistributedInfo dtype: torch.dtype + # --hf-overrides: applied to the checkpoint config the model is built from (cached_load_hf_config) + hf_overrides: Mapping[str, Any] = field(default_factory=dict) max_running_req: int = 4 attention_backend: str = "auto" moe_strategy: str = "auto" @@ -32,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 @@ -40,8 +48,15 @@ class EngineConfig: moe_cache_size: int = 0 moe_cache_rate: float | None = None moe_cache_auto: bool = False + # --moe-cache-mem: slot-cache budget in bytes, turned into a slot count once the expert banks are loaded + moe_cache_mem: int | None = None + # --moe-cache-scope: unified (one LRU pool per expert quant) or per-layer (one pool per layer) + moe_cache_scope: str = "unified" + # --moe-cache-split: how pools share the slots; equal or profile: (by expert use from a profile file) + moe_cache_split: str = "equal" kv_reserve_tokens: int = 8192 # KV floor for --moe-cache-auto; small by design (MoE-priority) moe_cache_policy: str = "lru" + moe_prefill_demand_tokens: int = 0 moe_prefill_overlap: bool = True # Prefill hit/miss split: serve cache-resident experts D2D during prefill # prefetch instead of re-streaming the full layer over PCIe. Needs CUDA >= 12.8 @@ -51,6 +66,8 @@ class EngineConfig: # CPU MoE backend (--moe-strategy cpu): number of CPU worker threads computing # the decode experts. 0 = auto (physical cores). Ignored by other backends. moe_cpu_threads: int = 0 + # skip the startup prefill warmup forwards; first requests then pay the kernel module loads (and the full JIT compile on a cold start) + prefill_warmup: bool = True # Hybrid CPU/GPU decode (--moe-strategy offload only): which MoE layers decode on # the CPU executor instead of the GPU offload/PCIe path. Spec is an explicit id # list ("3,7,11"), a count ("8" -> 8 layers evenly strided across depth), or a @@ -66,22 +83,34 @@ class EngineConfig: cuda_graph_bs: List[int] | None = None cuda_graph_max_bs: int | None = None page_size: int = 1 + # KV-cache storage quantization: "none" stores the compute dtype, "fp8" stores e4m3 + # codes plus one fp32 scale per (token, slab, layer, kv head) -- about 2x the tokens + # per GiB, at a small accuracy cost. --kv-cache-dtype; resolved from "auto" by + # _adjust_config, which also refuses it on a pool family or attention backend that + # cannot read the scales. + kv_quant: str = "none" memory_ratio: float = 0.9 # Hybrid GDN models default to the HybridRadixCache (cross-request GDN-state prefix reuse); # `--cache-type naive` opts out. linear_state_cache_ratio sizes the GDN snapshot cache as # ceil(ratio * max_running_req) extra slots. - linear_state_cache_ratio: float = 2.0 + linear_state_cache_ratio: float = float(os.environ.get("FT_LINEAR_STATE_CACHE_RATIO", "2.0")) # Window/full ratio for the SWA radix cache (`--cache-type radix` on SWA models) and the DSV4 # window tier: the DEFAULT window-pool size = max(working-set floor, ratio x full-pool tokens). # < 1.0 trades retained window-prefix capacity for memory savings; must be in (0, 1]. It is the # DSV4 window/full ratio directly. Used only when swa_num_pages_override is None (a runtime # rebuild can pin an absolute window instead). swa_full_tokens_ratio: float = 0.2 + # Force one decode step after this many consecutive prefill steps, so a long chunked + # prefill cannot starve in-flight decodes. None/0 keeps the historical prefill-first + # order (the scheduler's own TODO names this: "support other policies: e.g. DECODE + # first"). At 8 the cost is ~1% of prefill wall time. + decode_interleave_every: int | None = None # Absolute window-pool size in the pool's own pages (usable, dummy excluded); None -> use the # ratio default above. A runtime cache rebuild sets this (num_swa_pages) to pin the window # regardless of the full anchor; the ratio is the startup default and the fallback. swa_num_pages_override: int | None = None - distributed_timeout: float = 60.0 + distributed_timeout: float = 1800.0 # ranks reach the first collective minutes apart on a 100+ GiB offload load + distributed_port: int = 2333 use_dummy_weight: bool = False use_pynccl: bool = True max_seq_len_override: int | None = None @@ -89,6 +118,11 @@ class EngineConfig: # KV capacity in tokens; resolved into num_page_override by _adjust_config once page_size # is final. Mutually exclusive with num_page_override. num_token_override: int | None = None + # Extra KV pages mirrored to pinned host RAM (QSA paged pools, e.g. Qwen3.8-Flash-Next): + # the GPU pool becomes an LRU cache over the logical page space (scheduler, page table and + # radix tree all see num_pages + kv_host_pages pages), extending context past VRAM + # capacity. 0 (default) = off. + kv_host_pages: int = 0 # Runtime knobs of the multimodal path; the architecture side (vision_config, mrope) lives in ModelConfig. mm: MultimodalConfig = field(default_factory=MultimodalConfig) @@ -103,7 +137,7 @@ def __post_init__(self): @cached_property def hf_config(self): - return cached_load_hf_config(self.model_path) + return cached_load_hf_config(self.model_path, self.hf_overrides) @cached_property def model_spec(self) -> ModelSpec: @@ -130,19 +164,29 @@ def model_config(self) -> ModelConfig: hf_config = copy.copy(self.hf_config) built = {e.config_key for e in self.active_encoders} for key in set(ENCODER_SECTIONS) | {e.config_key for e in self.model_spec.encoders}: - if key not in built: + if key not in built and getattr(hf_config, key, None) is not None: setattr(hf_config, key, None) spec = self.model_spec 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: return self.max_seq_len_override - return self.model_config.rotary_config.max_position + return self.model_config.rotary_config.table_positions @property def max_forward_len(self) -> int: @@ -150,4 +194,4 @@ def max_forward_len(self) -> int: @property def distributed_addr(self) -> str: - return "tcp://127.0.0.1:2333" + return f"tcp://127.0.0.1:{self.distributed_port}" diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index cc14cf5339..db06a1e136 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -1,10 +1,12 @@ from __future__ import annotations +import collections import contextlib import errno import gc import math import os +import socket from datetime import timedelta from typing import Any, Dict, Iterable, NamedTuple, Tuple @@ -12,17 +14,20 @@ 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.layers.quantization import LayerKind, QuantBackend, finalize_quant, set_quant_backend +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, QuantKind, 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.moe.offload_cache import OffloadMoeCache, attach_offload_moe_cache, bank_layouts +from freetoken.moe.pool_policy import group_layers, make_pools, pool_of_layer, resolve_profile, split_slots from freetoken.utils import align_ceil, init_logger, is_sm90_family, is_sm100_family, mem_GB, torch_dtype from .config import EngineConfig @@ -32,11 +37,35 @@ from freetoken.kvcache.base import CacheRebuildRejected from freetoken.kvcache.cache_status import _supports_swa_ratio from freetoken.kvcache.linear_state_pool import ( - _linear_pool_min_slots, _linear_pool_num_slots, state_pool_bytes, + _linear_pool_min_slots, _linear_pool_num_slots, gdn_prefill_workspace_bytes, + state_pool_bytes, ) logger = init_logger(__name__) +def _clamp_decode_batch_to_pools(config: EngineConfig, pool_sizes: list[int]) -> None: + """Clamp the decode batch so every GPU-decoded layer's routed experts fit in its pool.""" + num_experts = config.model_config.num_experts + small = [s for s in pool_sizes if s < num_experts] + if not small: + return + limit = min(small) // config.model_config.num_experts_per_tok + if config.max_running_req > limit: + logger.warning_rank0( + f"the smallest MoE slot pool has {min(small)} slots; lowering max_running_req " + f"{config.max_running_req} -> {limit} so a decode step fits in it" + ) + object.__setattr__(config, "max_running_req", limit) + if config.cuda_graph_max_bs is None or config.cuda_graph_max_bs > limit: + object.__setattr__(config, "cuda_graph_max_bs", limit) + if config.cuda_graph_bs is not None: + object.__setattr__(config, "cuda_graph_bs", [bs for bs in config.cuda_graph_bs if bs <= limit]) +# Decode steps between --moe-collect-stats reports. Reading the counters costs a host sync, +# so the window is coarse; it is short enough that a normal-length reply still produces a +# few reports rather than one at the very end. +MOE_STATS_INTERVAL = 256 + + def _require_offload_cache_size(cache_size: int, num_experts: int) -> None: """The offload MoE cache needs at least one slot per expert per layer. A too-small size (e.g. a bare offload run with moe_cache_size unset and auto disabled) must fail loudly.""" @@ -49,6 +78,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 @@ -57,6 +92,10 @@ def _flashinfer_available() -> bool: def _sgl_flash_attn_available() -> bool: + from freetoken.kernel.backend import is_rocm + + if is_rocm(): + return False try: from sgl_kernel.flash_attn import flash_attn_with_kvcache # noqa: F401 except Exception as exc: @@ -69,6 +108,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 +190,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 +261,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 "" + ) + + "." ) @@ -164,6 +283,7 @@ def _validate_attention_backend_choice(config, override, required: frozenset[Att where a DSV4 or MLA checkpoint rejects a generic backend before weights load, and where a generic model rejects dsa/dsv4_sparse.""" from freetoken.attention import validate_attn_backend + from freetoken.kernel.backend import is_rocm # Name membership first (ArgumentTypeError listing the supported names): the CLI already # ran this, but the programmatic EngineConfig path reaches here unvalidated and would @@ -198,23 +318,57 @@ 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. for part in backend_parts: info = attention_backend_info(part) if info.requires_flashinfer and not _flashinfer_available(): + if is_rocm(): + raise RuntimeError( + f"Attention backend {part!r} uses CUDA-only FlashInfer kernels " + "unavailable on ROCm; use --attention-backend auto or triton." + ) raise RuntimeError( f"Attention backend {config.attention_backend!r} requires flashinfer, which is " "not installed. Install it with `pip install 'freetoken[fi]'` (or " "'freetoken[accel]'), or use --attention-backend triton." ) if info.requires_sgl_kernel and not _sgl_flash_attn_available(): + if is_rocm(): + raise RuntimeError( + f"Attention backend {part!r} uses CUDA-only sgl_kernel kernels " + "unavailable on ROCm; use --attention-backend auto or triton." + ) raise RuntimeError( f"Attention backend {config.attention_backend!r} requires sgl_kernel, which is " "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 " @@ -266,7 +420,7 @@ def _make_dummy_weight_state_dict( # MXFP8 e8m0 exponent codes: 127 encodes scale 1.0; zeros would collapse # every scale to 2^-127 and zero the model. Scoped BY NAME: other uint8 # buffers are packed payloads whose bytes mean something else entirely - # (GGUF qweight blocks embed fp16 scales -- 0x7F7F is fp16 NaN), so they + # (GGUF blocks embed fp16 scales -- 0x7F7F is fp16 NaN), so they # keep the benign all-zeros fill below. state_dict[key] = torch.full(param.shape, 127, dtype=param.dtype, device=device) else: @@ -322,6 +476,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: @@ -354,6 +513,9 @@ def __init__(self, config: EngineConfig): logger.info_rank0(f"Free memory before loading model: {mem_GB(init_free_memory)}") # ======================= Model initialization ======================== + from freetoken.kernel.triton.autotune_cache import bound_autotune_flush_buffer + + bound_autotune_flush_buffer(self.device.index) set_rope_device(self.device) with torch.device("meta"), torch_dtype(config.dtype): self.model = create_model(config.model_config) @@ -368,6 +530,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 +540,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 +581,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 +646,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 +700,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 +744,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 +755,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 +770,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, + dummy=config.use_dummy_weight, + parallel={"serial": False, "parallel": True}.get(config.expert_load, None), + resident=True, + layer_methods=[layer.quant_method for layer in sorted(layers, key=lambda layer: layer.layer_id)], + ) + 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): @@ -586,7 +829,7 @@ def _run_mm_encoder(self, batch: Batch) -> None: if parts: batch.mm_embeds = torch.cat([p.to(self.dtype) for p in parts], dim=0) - def _resolve_auto_moe_cache_size(self, config: EngineConfig, banks, method=None) -> tuple[int, int, bool]: + def _resolve_auto_moe_cache_size(self, config: EngineConfig, banks, method=None, max_slots=None) -> tuple[int, int, bool]: """Resolve --moe-cache-auto into (moe_cache_size, num_pages, prefill_overlap). Pure glue over the Phase-1 budget policy; isolated here so it is unit-testable @@ -596,6 +839,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,16 +852,31 @@ def _resolve_auto_moe_cache_size(self, config: EngineConfig, banks, method=None) num_experts=num_experts, total_experts=total_experts, prefill_overlap=config.moe_prefill_overlap, - kv_reserve_tokens=max(config.kv_reserve_tokens, min_reserve), + kv_reserve_tokens=max( + config.kv_reserve_tokens, + min_reserve, + (config.num_page_override or 0) * page_tokens, + ), page_size=page_tokens, - max_slots=method.slot_limit() if method is not None else None, + max_slots=max_slots if max_slots is not None or method is None else method.slot_limit(), ) def _init_offload_moe_cache(self, config: EngineConfig) -> OffloadMoeCache: + from freetoken.moe.expert_banks import is_mixed + method = shared_offload_method(self.model) num_moe_layers = config.model_config.num_moe_layers + # experts of several formats: one method per MoE layer, one slot pool per format + layer_methods = offload_layer_methods(self.model) + if not is_mixed(layer_methods): + layer_methods = None + elif config.moe_strategy == "cpu" or config.moe_cpu_layers: + raise ValueError( + "this model's experts come in several formats, which are served from per-format GPU slot pools; " + "they do not combine with --moe-strategy cpu or --moe-cpu-layers (use offload or hybrid)" + ) cpu_layer_ids = _resolve_cpu_layers(config, num_moe_layers, reserved=self._host_tables_bytes, method=method) - _check_pin_budget(config, reserved=self._host_tables_bytes, method=method) + _check_pin_budget(config, reserved=self._host_tables_bytes, method=method, layer_methods=layer_methods) # the kernels were picked for model_config.decode_target; --moe-cpu-layers auto may still find that every bank fits the pin budget decode_target = config.model_config.decode_target if decode_target == "cpu" and not cpu_layer_ids: @@ -666,6 +925,11 @@ def _init_offload_moe_cache(self, config: EngineConfig) -> OffloadMoeCache: else HostResidency.PINNED.value for i in range(config.model_config.num_moe_layers) ] + # read before the expert load, so a bad profile fails in seconds + expert_counts = resolve_profile( + config.moe_cache_split, config.model_path, + config.model_config.num_moe_layers, config.model_config.num_experts, + ) try: with _weight_load_context(): banks = load_expert_banks( @@ -673,16 +937,39 @@ def _init_offload_moe_cache(self, config: EngineConfig) -> OffloadMoeCache: config.model_config, method=method, device=self.device, - dtype=self.dtype, dummy=config.use_dummy_weight, parallel=expert_parallel, - decode_target=("cpu" if decode_target in ("cpu", "hybrid") else "gpu"), layer_residency=requested_residency, + layer_methods=layer_methods, ) except PinFailed as exc: raise RuntimeError(f"{exc}; {_pin_hint(self._host_tables_bytes)}") from exc + from freetoken.engine.cache_budget import expert_bytes_per_slot, layer_slot_bytes + + self._moe_layer_slot_bytes = layer_slot_bytes(banks.sources) + self._moe_pool_groups = group_layers(config.moe_cache_scope, bank_layouts(banks.sources, banks.layer_tags), self._moe_layer_slot_bytes) + if len(self._moe_pool_groups) > 1 and cpu_layer_ids: + raise ValueError( + "several MoE slot pools (--moe-cache-scope per-layer, or experts in several quant formats) need every MoE " + "layer to decode from the GPU slot cache; they do not combine with --moe-strategy cpu or --moe-cpu-layers" + ) + self._moe_expert_counts = expert_counts if len(self._moe_pool_groups) > 1 else None + if len(self._moe_pool_groups) == 1 and config.moe_cache_split != "equal": + logger.warning_rank0( + f"--moe-cache-split {config.moe_cache_split} is ignored: all MoE layers share one slot pool " + f"(--moe-cache-scope per-layer gives each layer its own)" + ) + self._moe_slot_limit = _slot_limit(layer_methods) if layer_methods is not None else (method.slot_limit() if method is not None else None) + # with GPU decode, a pool below num_experts slots bounds the batch; the startup split keeps room for one token per step + pools = self._moe_pools(1) + size_cap = sum(p.max_num_slots for p in pools) if config.moe_cache_auto: - size, pages, overlap = self._resolve_auto_moe_cache_size(config, banks, method) + size, pages, overlap = self._resolve_auto_moe_cache_size(config, banks, method, max_slots=size_cap) + if config.moe_prefill_overlap and not overlap: + logger.info_rank0( + f"--moe-cache-auto: prefill overlap disabled, its {2 * config.model_config.num_experts} " + "slot floor does not fit beside the KV reserve" + ) object.__setattr__(config, "moe_cache_size", size) object.__setattr__(config, "moe_prefill_overlap", overlap) if config.num_page_override is None: @@ -698,15 +985,46 @@ def _init_offload_moe_cache(self, config: EngineConfig) -> OffloadMoeCache: f"--moe-cache-auto resolved moe_cache_size={size} " f"num_pages={pages} (prefill_overlap={overlap})" ) + elif config.moe_cache_mem is not None: + per_slot = expert_bytes_per_slot(banks.sources) + size = min(config.moe_cache_mem // per_slot, size_cap) + # per_slot is a mean over layers, so with pools of different slot sizes the split can exceed the budget: shrink until it fits + while size > 0 and self._moe_pool_bytes(split_slots(size, pools, self._moe_expert_counts)) > config.moe_cache_mem: + size -= 1 + overlap = config.moe_prefill_overlap and size >= 2 * config.model_config.num_experts + object.__setattr__(config, "moe_cache_size", size) + object.__setattr__(config, "moe_prefill_overlap", overlap) + logger.info_rank0( + f"--moe-cache-mem {mem_GB(config.moe_cache_mem)} resolved moe_cache_size={size} " + f"({per_slot / 2**20:.2f} MiB per slot, prefill_overlap={overlap})" + ) _require_offload_cache_size(config.moe_cache_size, config.model_config.num_experts) - layout = max_slots = None - if method is not None: + pool_sizes = split_slots(config.moe_cache_size, pools, self._moe_expert_counts) + if len(pools) > 1: + logger.info_rank0( + f"MoE slot pools: {config.moe_cache_scope}, {len(pools)} pools of {min(pool_sizes)}-{max(pool_sizes)} " + f"slots ({config.moe_cache_split})" + ) + if decode_target == "gpu": + _clamp_decode_batch_to_pools(config, pool_sizes) + layout = layer_layouts = max_slots = None + if layer_methods is not None: + _check_layer_formats(banks, layer_methods) + layer_layouts = [m.layout() for m in layer_methods] + max_slots = self._moe_slot_limit + elif method is not None: + if banks.layer_formats is not None: + raise ValueError( + f"expert banks hold experts of several formats ({sorted(set(banks.layer_formats))}) but the model " + f"binds {method.kind} / {method.kernel.name} on every MoE layer; reconvert the checkpoint" + ) if banks.kind is not None and (banks.kind, banks.kernel) != (method.kind, method.kernel.name): raise ValueError( f"expert banks were packed for {banks.kind} / {banks.kernel} but the model " f"binds {method.kind} / {method.kernel.name}; reconvert or select that kernel" ) - layout = method.layout() + # GGUF layers keep their own ggml types, so the cache reads each layer's layout off its banks + layout = None if method.kind is QuantKind.GGUF else method.layout() max_slots = method.slot_limit() cache = OffloadMoeCache( # Models with leading dense layers (GLM-4) only have experts on the MoE @@ -716,69 +1034,109 @@ def _init_offload_moe_cache(self, config: EngineConfig) -> OffloadMoeCache: cache_size=config.moe_cache_size, device=self.device, cache_policy=config.moe_cache_policy, + prefill_demand_tokens=config.moe_prefill_demand_tokens, prefill_overlap=config.moe_prefill_overlap, prefill_hit_d2d=config.moe_prefill_hit_d2d, quant_format=banks.quant_format, decode_target=decode_target, hybrid_max_fetch=config.moe_hybrid_max_fetch, layout=layout, + layer_layouts=layer_layouts, max_slots=max_slots, + pool_of_layer=pool_of_layer(pools, config.model_config.num_moe_layers), + pool_sizes=pool_sizes, ) # before set_bank_sources: the residency validation and the copy plan's skip of non-pinned layers key on the CPU-layer set cache.cpu_layer_ids = cpu_layer_ids - cache.set_bank_sources(banks.sources, layer_residency=banks.layer_residency) + cache.set_bank_sources(banks.sources, layer_residency=banks.layer_residency, layer_tags=banks.layer_tags, row_tail_bytes=banks.row_tail_bytes) cache.set_alphas(banks.gate_up_alpha, banks.down_alpha) - if decode_target == "hybrid": - self._resolve_hybrid_fetch(config, cache) # Must be set before CUDA graph capture so the (device-side) accumulation ops are # captured and re-run on every decode replay. cache.collect_stats = config.moe_collect_stats + # The routing histogram rides the same switch: on its own the miss rate says how + # often we fetch, but not whether a smarter policy could have avoided the fetch. + # decode_routing_stats turns it into routing-skew numbers, including the hit rate of + # the best fixed expert set per layer (not a ceiling: LRU's temporal locality can beat it). + cache.collect_decode_freq = config.moe_collect_stats + # attach_offload_moe_cache walks for OffloadMoELayers, or defers to a model's + # _iter_offload_moe_layers() hook when its MoE blocks are bespoke nn.Modules (DSV4). layers = attach_offload_moe_cache(self.model, cache) assert len(layers) == config.model_config.num_moe_layers if cache.decode_target in ("cpu", "hybrid"): self._init_cpu_moe_executor(config, cache, layers) + if decode_target == "hybrid": + self._resolve_hybrid_fetch(config, cache) self.ctx.moe_offload_cache = cache self.moe_offload_cache = cache return cache + def _moe_pool_bytes(self, pool_sizes: list[int]) -> int: + return sum(s * self._moe_layer_slot_bytes[g[0]] for g, s in zip(self._moe_pool_groups, pool_sizes)) + + def _moe_pools(self, batch: int) -> list: + """The slot pools, each with room for the routed experts of ``batch`` decode tokens.""" + model_config = self.config.model_config + num_experts = model_config.num_experts + min_num_slots = min(num_experts, model_config.num_experts_per_tok * batch) + return make_pools(self._moe_pool_groups, num_experts, min_num_slots, self._moe_slot_limit) + def _resolve_hybrid_fetch(self, config: EngineConfig, cache) -> None: """Resolve --moe-hybrid-max-fetch -1 (auto) into a bandwidth-matched fetch fraction. Perfect fetch/compute overlap wants fetched : cpu-computed misses = pcie_bw : (cpu_bw - pcie_bw), i.e. fetching a pcie_bw / cpu_bw fraction of each decode step's misses -- both sides then finish together instead of one idling. The - achieved bandwidths come from the cached `ft bench bw` profile (the same one the - auto backend pick reads); without a usable profile the old fixed cap of 1 applies. + achieved bandwidths come from the cached `ft bench bw` profile, from the row of each + MoE layer's kind (format, options, expert shape, top_k, CPU worker threads); a layer + without a row keeps the fixed cap of 1. """ if config.moe_hybrid_max_fetch >= 0: return # explicit fixed cap - from freetoken.moe.bench_profile import load_hybrid_fetch_fraction + from freetoken.moe.bench_profile import fetch_fraction, layer_kind_key, load_layer_kinds + from freetoken.moe.cpu_executor import _format_label + cache.hybrid_max_fetch = 1 + methods = offload_layer_methods(self.model) + if methods is None: + logger.warning_rank0("--moe-hybrid-max-fetch auto: the MoE layers bind no expert method to look up in the `ft bench bw` profile; using a fixed fetch cap of 1") + return gpu_name, gpu_uuid = _profile_gpu(self.device.index) - fraction = load_hybrid_fetch_fraction( - cache.quant_format, gpu_name=gpu_name, gpu_uuid=gpu_uuid - ) - if fraction is None: - cache.hybrid_max_fetch = 1 + rows = load_layer_kinds(gpu_name=gpu_name, gpu_uuid=gpu_uuid) or {} + executor = self.cpu_moe_executor + keys, labels = [], [] + for m in methods: + fmt, opts = m.cpu_spec() + c = m.cfg + keys.append(layer_kind_key(fmt, opts, c.hidden, c.local_intermediate, c.top_k, executor.num_threads)) + labels.append(f"{_format_label(fmt, tuple(sorted(opts.items())))} {c.hidden}x{c.local_intermediate}") + fractions = [fetch_fraction(rows.get(key)) for key in keys] + cache.hybrid_fetch_fraction = [f or 0.0 for f in fractions] + groups = collections.Counter(zip(labels, fractions)) + fetched = [f"{f:.1%} on {n} {label} layers" for (label, f), n in groups.items() if f is not None] + if fetched: + logger.info_rank0( + "--moe-hybrid-max-fetch auto: each decode step fetches this share of a layer's expert misses over PCIe " + f"(benched PCIe/CPU bandwidth ratio) and computes the rest on the CPU: {'; '.join(fetched)}" + ) + if missing := sorted({label for label, f in groups if f is None}): + capped = sum(n for (_, f), n in groups.items() if f is None) logger.warning_rank0( - "--moe-hybrid-max-fetch auto: no usable `ft bench bw` profile for " - f"{cache.quant_format!r} experts; using a fixed fetch cap of 1" + f"--moe-hybrid-max-fetch auto: no `ft bench bw` row for {capped} MoE layers ({', '.join(missing)}) at " + f"{executor.num_threads} CPU threads; using a fixed fetch cap of 1 there. `ft bench bw --model ` benches each layer kind" + ) + benched_isas = {rows[key].get("cpu_moe_isa") for key, f in zip(keys, fractions) if f is not None} - {None} + if other := sorted(benched_isas - set(executor.isa.split(","))): + logger.warning_rank0( + f"--moe-hybrid-max-fetch auto: some `ft bench bw` rows were benched on CPU kernels {other}, but this executor runs " + f"{executor.isa}; using them anyway" ) - return - cache.hybrid_max_fetch = cache.num_experts # inert: the fraction is the cap - cache.hybrid_fetch_fraction = fraction - logger.info_rank0( - f"--moe-hybrid-max-fetch auto: fetching {fraction:.1%} of each decode step's " - "expert misses over PCIe (benched PCIe/CPU bandwidth ratio), the rest on the CPU" - ) def _init_cpu_moe_executor(self, config: EngineConfig, cache, layers) -> None: """Build the persistent CPU MoE executor (decode-time expert compute). Must run before CUDA graph capture: the worker pool has to be live for the - eager warmup forward, and the pinned IO buffers / host-func task pointers - must be stable for the captured nodes. Buffers/tasks themselves are - allocated lazily on the first (eager) forward at each batch size. + eager warmup forward, and the pinned IO block must be stable for the captured + nodes. """ from freetoken.moe.cpu_executor import CpuMoeExecutor @@ -791,7 +1149,9 @@ def _init_cpu_moe_executor(self, config: EngineConfig, cache, layers) -> None: ) # Decode batches never exceed max_running_req, but CUDA-graph padding can # round a batch up to the largest captured size; cover both. - max_tokens = max(config.max_running_req, config.cuda_graph_max_bs or 0, 1) + max_tokens = max(config.max_running_req, config.cuda_graph_max_bs or 0, *(config.cuda_graph_bs or []), 1) + by_id = sorted(layers, key=lambda layer: layer.layer_id) + layer_formats = [layer.quant_method.cpu_spec() for layer in by_id] executor = CpuMoeExecutor( cache, top_k=sample.top_k, @@ -802,8 +1162,7 @@ def _init_cpu_moe_executor(self, config: EngineConfig, cache, layers) -> None: device=self.device, 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, + layer_formats=layer_formats, ) cache.set_cpu_executor(executor) self.cpu_moe_executor = executor @@ -829,19 +1188,18 @@ def _sync_get_memory(self) -> Tuple[int, int]: return min_free_memory, max_free_memory - def _target_moe_and_expert_bytes(self, moe_cache_size: int | None) -> tuple[int, int]: - from freetoken.engine.cache_budget import expert_bytes_per_slot - - target_moe = ( - moe_cache_size - if moe_cache_size is not None - else (self.moe_offload_cache.cache_size if self.moe_offload_cache else 0) - ) - per_expert_bytes = ( - expert_bytes_per_slot(self.moe_offload_cache.bank_sources) - if self.moe_offload_cache is not None else 0 - ) - return target_moe, per_expert_bytes + def _target_moe_and_expert_bytes(self, moe_cache_size: int | None, pool_sizes: list[int] | None) -> tuple[int, int]: + cache = self.moe_offload_cache + if cache is None: + return moe_cache_size or 0, 0 + if moe_cache_size is None: + target_moe = cache.cache_size + nbytes = sum(c.numel() * c.element_size() for _, c in cache.banks) + else: + target_moe = moe_cache_size + nbytes = self._moe_pool_bytes(pool_sizes) + # pools of several bank layouts have different slot sizes; round up so target * bytes covers them + return target_moe, -(-nbytes // target_moe) def _resize_kv_pool(self, config, num_pages: int, num_swa_pages: int | None) -> None: # IN-PLACE, identity-preserving: the CacheManager's swa_pool reference, ctx.kv_cache and @@ -900,11 +1258,22 @@ def rebuild_runtime_cache( "moe_cache_size requested but this model has no MoE offload cache" ) try: - self.moe_offload_cache.validate_rebuild(moe_cache_size) + # a rebuild does not lower max_running_req again, so with GPU decode every pool must hold a full batch; hybrid caps its fetches instead + batch = config.max_running_req if self.moe_offload_cache.decode_target == "gpu" else 1 + moe_pool_sizes = split_slots(moe_cache_size, self._moe_pools(batch), self._moe_expert_counts) + self.moe_offload_cache.validate_rebuild(moe_cache_size, moe_pool_sizes) except ValueError as e: raise CacheRebuildRejected(str(e)) from e if num_pages is not None and num_pages <= 0: raise CacheRebuildRejected(f"num_pages must be positive, got {num_pages}") + if num_pages is not None and self.kv_offloader is not None: + # The logical/physical split (host mirror, LRU maps, QSA index slab sized to the + # logical space) is built once at startup; a runtime resize would have to rebuild + # all of it. Reject recoverably instead. + raise CacheRebuildRejected( + "num_pages rebuild is not supported with --kv-host-pages (restart with a " + "different --kv-host-pages / --memory-ratio instead)" + ) if num_mamba_slots is not None: if self.linear_state_pool is None: raise CacheRebuildRejected( @@ -937,7 +1306,9 @@ def rebuild_runtime_cache( # must reject (recoverable) so the old caches stay intact and serving continues, # rather than freeing and then OOMing into permanent failure. The engine supplies # the memory account; the pool answers whether its target geometry fits. - target_moe, per_expert_bytes = self._target_moe_and_expert_bytes(moe_cache_size) + target_moe, per_expert_bytes = self._target_moe_and_expert_bytes( + moe_cache_size, moe_pool_sizes if moe_cache_size is not None else None + ) # Price the sibling GDN state pool at ITS target (physical slots = usable + padding # sink) and hand the bytes in -- the KV pool only budgets its own tiers. target_mamba = ( @@ -951,7 +1322,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 "" @@ -972,29 +1344,54 @@ def rebuild_runtime_cache( # 1. Tear down CUDA graphs + backend capture scratch (free-before-alloc). self.attn_backend.reset_capture() self.graph_runner.destroy_cuda_graphs() - # 2. Resize caches in place (each frees its old GPU tensors before allocating). + # 2. Resize caches in place (each frees its old GPU tensors before allocating), + # shrinking resizes before growing ones: a pool's grow must not run while a + # still-large sibling is resident, or a target that passed the fit check above can + # OOM mid-rebuild (#643). With shrinks first the transient peak never exceeds + # max(current, target) totals. Footprints are compared BEFORE the window-override + # write below (the old side must price the pre-override config). + kv_old_bytes = type(self.kv_cache).rebuild_footprint(config, self.num_pages, None) + kv_new_bytes = type(self.kv_cache).rebuild_footprint( + config, num_pages if num_pages is not None else self.num_pages, num_swa_pages + ) + + def _resize_moe() -> None: + assert self.moe_offload_cache is not None, "no MoE offload cache to resize" + self.moe_offload_cache.rebuild(moe_cache_size, moe_pool_sizes) + + def _resize_kv() -> None: + # sets self.num_pages (rebuilds KV + window); a window-only change passes the + # CURRENT page count so only the window pool is re-derived + self._resize_kv_pool( + config, num_pages if num_pages is not None else self.num_pages, num_swa_pages + ) + + def _resize_mamba() -> None: + # Like every resize here it must sit between graph teardown and re-capture so + # the recaptured graphs bind the new state tensors. +1 for the reserved + # padding sink: num_mamba_slots is the usable count. + self.linear_state_pool.rebuild(num_mamba_slots + 1) + + resize_ops = [ + (moe_cache_size is not None, + moe_cache_size is None or moe_cache_size <= self.moe_offload_cache.cache_size, + _resize_moe), + (num_pages is not None or num_swa_pages is not None, + kv_new_bytes <= kv_old_bytes, _resize_kv), + (num_mamba_slots is not None, + num_mamba_slots is None or num_mamba_slots + 1 <= self.linear_state_pool.num_slots, + _resize_mamba), + ] # Pin the new window first (validated above) so any KV-pool rebuild below sizes the window # to it (_dsv4_pool_sizes / _swa_paged_num_tokens read config.swa_num_pages_override). # frozen EngineConfig — mutate in place like the moe_cache_size path; `config.x = y` raises # FrozenInstanceError, which here aborts the rebuild after the CUDA graphs are gone (→ 503). if num_swa_pages is not None: object.__setattr__(config, "swa_num_pages_override", num_swa_pages) - if moe_cache_size is not None: - assert self.moe_offload_cache is not None, "no MoE offload cache to resize" - self.moe_offload_cache.rebuild(moe_cache_size) - if num_pages is not None: - # sets self.num_pages (rebuilds KV + window) - self._resize_kv_pool(config, num_pages, num_swa_pages) - elif num_swa_pages is not None: - # Window-only change: no page-count change, but re-derive the window pool at the new - # pin against the CURRENT page count. This re-allocs the same-size full pool and - # the resized window, both inside the pool's own rebuild_from_config. - self._resize_kv_pool(config, self.num_pages, num_swa_pages) - if num_mamba_slots is not None: - # Reallocate the GDN state pool (frees old tensors first). Must sit between graph - # teardown and re-capture so the recaptured graphs bind the new state tensors. - # +1 for the reserved padding sink: num_mamba_slots is the usable count. - self.linear_state_pool.rebuild(num_mamba_slots + 1) + for shrinking_phase in (True, False): + for requested, shrinks, op in resize_ops: + if requested and shrinks == shrinking_phase: + op() # 3. Refresh max_seq_len (+ generic page table) for the new token budget. self._refresh_seq_state(config) aligned_max_seq_len = _page_table_width(self.max_seq_len, config.page_size) @@ -1024,6 +1421,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 +1435,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 +1586,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 +1615,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 +1630,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 +1706,14 @@ def _is_unified_memory_gpu(index: "int | None" = None) -> bool: return False +_RESIDENT_EXPERT_QUANTS = ("none", "fp8_block", "nvfp4", "gguf") + + 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; 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 +1751,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 @@ -1294,10 +1897,10 @@ def _cpu_moe_executor_viable(model_config) -> bool: """Whether an automatic CPU-decode decision may target the CPU MoE executor. A default boot must degrade to GPU offload instead of crashing in CpuMoeExecutor after the whole load; explicit cpu/hybrid/--moe-cpu-layers picks still fail loudly.""" - from freetoken.moe.cpu_executor import _WFMT_IDS, compiled_extension_supports + from freetoken.moe.cpu_executor import compiled_extension_supports try: - from freetoken.kernel import _cpu_moe # noqa: F401 + from freetoken.kernel import _cpu_moe except ImportError: return False act = getattr(model_config, "hidden_act", "silu") @@ -1308,7 +1911,11 @@ def _cpu_moe_executor_viable(model_config) -> bool: return False expert_quant = getattr(model_config, "expert_quant", "none") fmt = expert_quant if expert_quant != "none" else (moe_wfmt or "bf16") - return fmt == "mxfp4" or fmt in _WFMT_IDS + if fmt == "gguf": + from freetoken.kernel import ggml_cpu + + return fmt in _cpu_moe.weight_formats() and ggml_cpu.select() is not None + return fmt == "mxfp4" or fmt in _cpu_moe.weight_formats() def _pin_budget_bytes(reserved: int = 0) -> int | None: @@ -1324,10 +1931,10 @@ def _pin_budget_bytes(reserved: int = 0) -> int | None: return max(0, cap - reserved) -def _bank_bytes(config: EngineConfig, method=None) -> int | None: +def _bank_bytes(config: EngineConfig, method=None, layer_methods=None) -> int | None: from freetoken.moe.expert_banks import bank_bytes_estimate, ftw_bank_bytes - return ftw_bank_bytes(config.model_path) or bank_bytes_estimate(config.model_config, method) + return ftw_bank_bytes(config.model_path) or bank_bytes_estimate(config.model_config, method, layer_methods) def _pin_hint(reserved: int) -> str: @@ -1339,12 +1946,12 @@ def _pin_hint(reserved: int) -> str: ) -def _check_pin_budget(config: EngineConfig, *, reserved: int, method=None) -> None: +def _check_pin_budget(config: EngineConfig, *, reserved: int, method=None, layer_methods=None) -> None: """Stop a plain offload boot whose banks exceed a known pin budget before any bank is read.""" if config.moe_cpu_layers or config.moe_strategy not in ("offload", "hybrid"): return budget = _pin_budget_bytes(reserved) - bank_bytes = _bank_bytes(config, method) if budget is not None else None + bank_bytes = _bank_bytes(config, method, layer_methods) if budget is not None else None if bank_bytes and bank_bytes > budget: raise ValueError( f"expert banks need {bank_bytes / 2**30:.1f} GiB of pinned host RAM but the pin budget is " @@ -1385,9 +1992,13 @@ def _auto_cpu_layers(config: EngineConfig, num_moe_layers: int, *, reserved: int "moe_cache_size": 0, "moe_cache_rate": None, "moe_cache_auto": False, + "moe_cache_mem": None, + "moe_cache_scope": "unified", + "moe_cache_split": "equal", "moe_cpu_layers": None, "moe_cpu_threads": 0, "moe_hybrid_max_fetch": -1, + "moe_prefill_demand_tokens": 0, "moe_prefill_overlap": True, "moe_prefill_hit_d2d": False, "expert_load": "auto", @@ -1420,35 +2031,82 @@ def _adjust_ftw_quant_backend(model_path: str, quant_backend: QuantBackend) -> Q return QuantBackend(quant_backend.items + (((LayerKind.MOE, kind), kernel),)) -def shared_offload_method(model): - """The expert method every offload MoE layer of ``model`` uses, or None for models whose MoE layers carry none (GGUF). - - The offload cache holds one bank layout, so the layers must agree on (kind, kernel).""" +def offload_layer_methods(model) -> list | None: + """The expert method of each offload MoE layer of ``model``, by MoE-layer id; None for models whose MoE layers carry none (GGUF).""" layers = [l for l in iter_offload_moe_layers(model) if getattr(l, "quant_method", None) is not None] if not layers: return None - keys = {(layer.quant_method.kind, layer.quant_method.kernel.name) for layer in layers} - if len(keys) != 1: - found = sorted(f"{kind} / {kernel}" for kind, kernel in keys) - raise ValueError(f"expert layers disagree on format / kernel: {found}; per-layer mixed expert formats are not supported") - kind, kernel = keys.pop() - logger.info_rank0(f"MoE experts: {kind} via {kernel}") - return layers[0].quant_method + layers.sort(key=lambda layer: layer.layer_id) + assert [layer.layer_id for layer in layers] == list(range(len(layers))), [layer.layer_id for layer in layers] + return [layer.quant_method for layer in layers] -def offload_expert_method(config: EngineConfig): - """The offload expert method of ``config``'s model, from a meta-device build. +def shared_offload_method(model): + """The expert method the offload MoE layers of ``model`` use, or None for models whose MoE layers carry none (GGUF). - For tools that load expert banks without an engine (the FTW converter); they pack for GPU decode.""" + With experts of several formats (a checkpoint that keeps some layers unquantized) this is the format of the most + layers; the cache then takes every layer's method (``offload_layer_methods``) and keeps one slot pool per format.""" + from freetoken.moe.expert_banks import method_groups + + methods = offload_layer_methods(model) + if methods is None: + return None + groups = method_groups(methods) + found = ", ".join(f"{m.kind} via {m.kernel.name}" + (f" on {len(ids)} layers" if len(groups) > 1 else "") for m, ids in groups) + logger.info_rank0(f"MoE experts: {found}") + return groups[0][0] + + +def _slot_limit(layer_methods) -> int | None: + """The tightest slot cap of the layers' expert kernels (a pool never mixes kernels, but the limit is checked per pool against the largest).""" + limits = [limit for limit in (m.slot_limit() for m in layer_methods) if limit is not None] + return min(limits) if limits else None + + +def _check_layer_formats(banks, layer_methods) -> None: + from freetoken.moe.legacy_format import legacy_format_for + + want = [legacy_format_for(m.kind, m.kernel.name) for m in layer_methods] + have = banks.layer_formats or [banks.quant_format] * len(want) + bad = [layer for layer, (w, h) in enumerate(zip(want, have)) if w != h] + if bad: + raise ValueError( + f"expert banks of MoE layers {bad[:8]} were packed as {[have[l] for l in bad[:8]]} but the model binds " + f"{[want[l] for l in bad[:8]]}; reconvert the checkpoint or select those kernels" + ) + + +def _meta_model(config: EngineConfig, decode_target: str): from freetoken.models import create_model from freetoken.utils.torch_utils import torch_dtype set_quant_backend(_adjust_ftw_quant_backend(config.model_path, QuantBackend.parse(config.quant_backend))) object.__setattr__(config.model_config, "moe_strategy", config.moe_strategy) - object.__setattr__(config.model_config, "decode_target", "gpu") + object.__setattr__(config.model_config, "decode_target", decode_target) with torch.device("meta"), torch_dtype(config.dtype): - model = create_model(config.model_config) - return shared_offload_method(model) + return create_model(config.model_config) + + +def offload_expert_methods(config: EngineConfig) -> tuple: + """``(method, layer_methods)`` of ``config``'s model, from a meta-device build: the most common expert method and, + for experts of several formats, one per MoE layer (else None). + + For tools that load expert banks without an engine (the FTW converter); they pack for GPU decode.""" + from freetoken.moe.expert_banks import is_mixed + + model = _meta_model(config, "gpu") + layer_methods = offload_layer_methods(model) + return shared_offload_method(model), (layer_methods if is_mixed(layer_methods) else None) + + +def meta_offload_layer_methods(config: EngineConfig, decode_target: str) -> list | None: + """The expert method of each offload MoE layer of ``config``'s model, from a meta-device build for ``decode_target``; see ``offload_layer_methods``.""" + return offload_layer_methods(_meta_model(config, decode_target)) + + +def offload_expert_method(config: EngineConfig): + """The (most common) offload expert method of ``config``'s model; see ``offload_expert_methods``.""" + return offload_expert_methods(config)[0] def _adjust_config(config: EngineConfig): @@ -1458,6 +2116,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 +2160,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 +2189,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} @@ -1552,10 +2241,13 @@ def override(attr: str, value: Any): # this is dangerous, use with caution "--dtype float16 with MXFP8 resident weights is unsupported (the " "W8A16 fold is only validated exact in bfloat16); use bfloat16." ) + if _dtype == torch.float16 and getattr(getattr(model_config, "quant", None), "weight_scales", False): + # GGUF NVFP4 results are stored in the activation dtype before weight_scale_2 (~1e-4) multiplies them, which overflows fp16. + raise ValueError("--dtype float16 with NVFP4 GGUF weight scales is unsupported (the unscaled results overflow float16); use bfloat16.") 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) @@ -1568,7 +2260,7 @@ def override(attr: str, value: Any): # this is dangerous, use with caution # swigluoai (csrc ActKind; "gpt_oss_swiglu" rides inside the mxfp4 kernel and # swigluoai the generic GEMV epilogue). A model with any other expert # activation cannot decode on the CPU: reject an explicit cpu/hybrid pick at - # config time, and keep auto from upgrading offload -> hybrid off the profile. + # config time. # hidden_act (the dense activation) stands proxy for the expert activation -- # true for every in-tree model. mxfp4 experts pass regardless: their act runs # inside the mxfp4 kernel, not the generic epilogue. @@ -1594,8 +2286,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 + # 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" @@ -1604,8 +2296,7 @@ def override(attr: str, value: Any): # this is dangerous, use with caution # GB10 this added a measured ~130 s stall to every request, #369). Resident # experts are the safe default here, not the risky one: the model and its banks # page from the same pool, so the "wrong guess = weight-load OOM" rationale above - # does not apply. The benchbw hybrid upgrade is skipped too: CPU execution adds - # no bandwidth when both sides share one memory. Only formats the resident path + # does not apply. Only formats the resident path # can actually hold take this branch; the rest stay on offload as before. unified_memory = _is_unified_memory_gpu() if unified_memory and _fused_resident_ok(model_config): @@ -1614,50 +2305,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" ) - # 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 - # offload -- same auto-sized GPU slot cache (_resolve_auto_moe_cache_size), plus a - # host-RAM CPU executor -- so this never raises the OOM risk; with no profile (or one - # from different hardware) it stays offload. offload remains the always-safe fallback. - # Key the lookup on the real expert format: mxfp4/q4_0 live in moe_weight_format when - # expert_quant is "none", and "none" with no weight format means plain bf16 experts. - moe_wfmt = getattr(model_config, "moe_weight_format", None) - bench_fmt = expert_quant if expert_quant != "none" else (moe_wfmt or "bf16") - from freetoken.moe.bench_profile import load_backend_recommendation - - gpu_name, gpu_uuid = _profile_gpu() - if ( - default_backend == "offload" - and not unified_memory - and load_backend_recommendation( - bench_fmt, gpu_name=gpu_name, gpu_uuid=gpu_uuid - ) - == "hybrid" - ): - from freetoken.moe.cpu_executor import compiled_extension_supports - - _act = getattr(model_config, "hidden_act", "silu") - if not _cpu_moe_act_ok: - logger.info_rank0( - f"benchbw profile recommends hybrid, but the CPU MoE executor does not " - f"support this model's expert activation " - f"{getattr(model_config, 'hidden_act', None)!r}; staying on offload" - ) - elif moe_wfmt != "mxfp4" and not compiled_extension_supports(_act): - # Stale prebuilt _cpu_moe.so: an explicit cpu/hybrid pick still - # hard-fails in the executor, but a default must not turn into a - # post-load crash -- degrade to offload. - logger.info_rank0( - f"benchbw profile recommends hybrid, but the compiled _cpu_moe " - f"extension predates activation {_act!r} (rebuild with " - f"`python setup.py build_ext --inplace`); staying on offload" - ) - else: - default_backend = "hybrid" - logger.info_rank0( - f"benchbw profile recommends hybrid for {bench_fmt!r} experts on this GPU" + # 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) override("moe_strategy", default_backend) logger.info_rank0(f"Auto-selected MoE strategy: {config.moe_strategy}") @@ -1665,6 +2324,7 @@ def override(attr: str, value: Any): # this is dangerous, use with caution is_offload_moe_strategy(config.moe_strategy) and config.moe_cache_size <= 0 and config.moe_cache_rate is None + and config.moe_cache_mem is None and not getattr(config, "moe_cache_auto", False) ): # args.py's "no sizing flag -> default --moe-cache-auto" only fires when the @@ -1698,11 +2358,13 @@ def override(attr: str, value: Any): # this is dangerous, use with caution # An explicit 'fused' keeps the experts resident, so there is no slot cache to size. The # sizing flags no longer redirect the backend, so ignore them here and say so -- the # geometry the user asked for is what runs. Report the flag actually passed: --moe-cache- - # rate was already folded into moe_cache_size above, and the three are mutually exclusive. + # rate was already folded into moe_cache_size above, and the sizing flags are mutually exclusive. if config.moe_cache_rate: inert = f"--moe-cache-rate={config.moe_cache_rate}" elif config.moe_cache_size: inert = f"--moe-cache-size={config.moe_cache_size}" + elif config.moe_cache_mem is not None: + inert = f"--moe-cache-mem={config.moe_cache_mem}" elif getattr(config, "moe_cache_auto", False): inert = "--moe-cache-auto" else: @@ -1714,6 +2376,7 @@ def override(attr: str, value: Any): # this is dangerous, use with caution ) override("moe_cache_size", 0) override("moe_cache_rate", None) + override("moe_cache_mem", None) override("moe_cache_auto", False) if is_moe and config.moe_strategy == "cpu": @@ -1724,6 +2387,7 @@ def override(attr: str, value: Any): # this is dangerous, use with caution num_experts = config.model_config.num_experts if getattr(config, "moe_cache_auto", False): override("moe_cache_auto", False) + override("moe_cache_mem", None) override("moe_cache_size", 2 * num_experts) override("moe_prefill_overlap", True) logger.info_rank0( @@ -1733,7 +2397,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( @@ -1770,18 +2434,17 @@ def override(attr: str, value: Any): # this is dangerous, use with caution ) override("num_page_override", config.num_token_override // config.page_size) - # The rope cos/sin table is baked to rotary_config.max_position, and neither rope kernel + # The rope cos/sin table covers rotary_config.table_positions, and neither rope kernel # bounds-checks the position it gathers with -- a longer ceiling reads past the table. # 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 > rotary.max_position: + 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 " - f"rope table ({rotary.max_position} positions). Serving past it would read " - "out of bounds; extend the checkpoint's rope_scaling / " - "max_position_embeddings in config.json instead." + f"rope table ({rotary.table_positions} positions). Serving past it would read " + "out of bounds; extend the rope with --hf-overrides (YaRN rope_parameters) instead." ) # The startup ServerArgs dump is the *requested* config, printed in the frontend process diff --git a/python/freetoken/engine/graph.py b/python/freetoken/engine/graph.py index eb9684f609..0d79bad93f 100644 --- a/python/freetoken/engine/graph.py +++ b/python/freetoken/engine/graph.py @@ -1,12 +1,16 @@ from __future__ import annotations +import ctypes +import functools import gc +import sys from dataclasses import dataclass from typing import TYPE_CHECKING, Dict, List import torch from freetoken.core import Batch, Req, get_global_ctx from freetoken.distributed import get_tp_info +from freetoken.env import ENV from freetoken.utils import init_logger, mem_GB from freetoken.utils.progress import emit_progress from tqdm import tqdm @@ -102,6 +106,30 @@ def get_free_memory(device: torch.device) -> int: return torch.cuda.mem_get_info(device)[0] +@functools.cache +def _cu_graph_upload(): + if sys.platform != "linux" or torch.version.hip is not None: + return None + fn = ctypes.CDLL("libcuda.so.1").cuGraphUpload + fn.argtypes = [ctypes.c_void_p, ctypes.c_void_p] + fn.restype = ctypes.c_int + return fn + + +def _upload_graph(graph: torch.cuda.CUDAGraph, stream: torch.cuda.Stream) -> None: + """Upload the graph at capture time. + + Otherwise the first replay pushes every node behind the graph's own stream WAIT, and a large graph then blocks cudaGraphLaunch until the WAIT is released. PLE wait-sync releases it only after the launch returns, so that deadlocks. + """ + upload = _cu_graph_upload() if ENV.GRAPH_UPLOAD else None + if upload is None: + return + rc = upload(graph.raw_cuda_graph_exec(), stream.cuda_stream) + if rc != 0: + raise RuntimeError(f"cuGraphUpload failed with CUresult {rc}") + stream.synchronize() + + class GraphRunner: def __init__( self, @@ -192,6 +220,7 @@ def _capture_graphs(self, max_seq_len: int, vocab_size: int, model: BaseLLMModel # CUDA graph capture to replay cold-cache expert copies. with torch.cuda.graph(graph, pool=pool, stream=self.stream): self.buffer.logits[:bs] = model.forward() + _upload_graph(graph, self.stream) self._reset_moe_offload_cache() if pool is None: pool = graph.pool() # reuse cuda graph handle to reduce memory diff --git a/python/freetoken/engine/sample.py b/python/freetoken/engine/sample.py index 3bdae76957..f07121ba9f 100644 --- a/python/freetoken/engine/sample.py +++ b/python/freetoken/engine/sample.py @@ -1,7 +1,7 @@ from __future__ import annotations -from dataclasses import dataclass -from typing import TYPE_CHECKING, List +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, List, NamedTuple import torch from freetoken.utils import is_sm90_supported, nvtx_annotate @@ -10,18 +10,137 @@ from freetoken.core import Batch +@dataclass +class LogitsPlan: + """Inputs of the logits processors for the rows of one batch that asked for any + (logit_bias, min_tokens, min_p). The penalties are not here: the sampler applies + them from per-request device state (``RowPenalty``). + + Built on the host in ``Sampler.prepare`` from the requests' SamplingParams, applied + on the device by ``apply_logits_processors`` before the sampling kernel. ``None`` + when no request in the batch uses a processor, so the default path stays exactly + as it was. Row indices inside the plan are LOCAL (0..r-1, in the order of ``rows``); + ``rows`` maps them back to the batch. + """ + + rows: torch.Tensor # [r] int64 batch rows + temps: torch.Tensor # [r, 1] float32 temperatures (min_p works on probabilities) + bias_rows: torch.Tensor | None = None # [b] int64 local rows + bias_ids: torch.Tensor | None = None # [b] int64 + bias_vals: torch.Tensor | None = None # [b] float32 + # Rows still under min_tokens and the ids (EOS + stop_token_ids, padded with + # vocab_size, a scratch column that is dropped) they must not sample yet. + min_rows: torch.Tensor | None = None # [m] int64 local rows + min_ids: torch.Tensor | None = None # [m, S] int64 + min_p: torch.Tensor | None = None # [r, 1] float32 (0 = off for that row) + # Rows with allowed_token_ids and their ids, padded with vocab_size like min_ids. + allow_rows: torch.Tensor | None = None # [a] int64 local rows + allow_ids: torch.Tensor | None = None # [a, A] int64 + + +class RowPenalty(NamedTuple): + """One batch row's presence / frequency / repetition penalty. ``counts`` are the + request's generated-token counts ([V] int32), updated on the device after each draw so + overlap scheduling never sees them late; ``prompt_mask`` ([V] bool) marks the prompt's + ids for the repetition penalty and is None when that is off.""" + + row: int + counts: torch.Tensor + presence: float + frequency: float + repetition: float = 1.0 + prompt_mask: torch.Tensor | None = None + + @dataclass class BatchSamplingArgs: temperatures: torch.Tensor | None top_k: torch.Tensor | None = None top_p: torch.Tensor | None = None greedy_mask: torch.Tensor | None = None + penalties: list[RowPenalty] = field(default_factory=list) + plan: LogitsPlan | None = None + # Rows that asked for logprobs, as an int64 device index built on the host (no + # nonzero sync). None with logprob_all_rows set when every row asks: no gather then. + logprob_row_idx: torch.Tensor | None = None + logprob_all_rows: bool = False + max_top_logprobs: int = 0 + # Logprob rows restricted to their allowed_token_ids, indexing the logprob rows above. + restrict_rows: torch.Tensor | None = None # [q] int64 + restrict_ids: torch.Tensor | None = None # [q, A] int64, padded with vocab_size def make_device_tensor(data: List, dtype: torch.dtype, device: torch.device) -> torch.Tensor: + if device.type != "cuda": # the unit tests build plans on the CPU + return torch.tensor(data, dtype=dtype, device=device) return torch.tensor(data, dtype=dtype, pin_memory=True).to(device, non_blocking=True) +def _padded_rows(rows: list[list[int]], pad: int, device: torch.device) -> torch.Tensor: + """[len(rows), max_len] int64 with ``pad`` on the right.""" + width = max(1, max((len(r) for r in rows), default=1)) + return make_device_tensor([r + [pad] * (width - len(r)) for r in rows], torch.int64, device) + + +def disallowed_mask(ids: torch.Tensor, vocab_size: int) -> torch.Tensor: + """[a, V] bool, True outside each row's ids (the vocab_size padding column is dropped).""" + block = torch.ones((ids.shape[0], vocab_size + 1), dtype=torch.bool, device=ids.device) + block.scatter_(1, ids, False) + return block[:, :vocab_size] + + +def apply_penalties(logits: torch.Tensor, penalties: list[RowPenalty]) -> torch.Tensor: + """A float32 copy of ``logits`` with each row's penalties applied in vLLM's order: + repetition (over prompt + generated ids), then frequency and presence (generated).""" + logits = logits.float().clone() + for row, counts, presence, frequency, repetition, prompt_mask in penalties: + if repetition != 1.0: + x = logits[row] + seen = counts > 0 + if prompt_mask is not None: + seen = seen | prompt_mask + # HF / vLLM semantics: divide positive logits, multiply negative ones. + penalized = torch.where(x > 0, x / repetition, x * repetition) + logits[row] = torch.where(seen, penalized, x) + if presence or frequency: + logits[row] -= frequency * counts + presence * (counts > 0) + return logits + + +def apply_logits_processors( + logits: torch.Tensor, plan: LogitsPlan, vocab_size: int +) -> torch.Tensor: + """A float32 copy of ``logits`` with the plan applied to its rows: logit_bias, the + allowed_token_ids and min_tokens masks, then min_p. Rows outside the plan are copied + unchanged.""" + out = logits.to(torch.float32, copy=True) + sub = out[plan.rows] # [r, V] (advanced indexing copies) + device = sub.device + neg_inf = float("-inf") + + if plan.bias_ids is not None: + sub.index_put_((plan.bias_rows, plan.bias_ids), plan.bias_vals, accumulate=True) + + if plan.allow_ids is not None: + rows = sub[plan.allow_rows] + sub[plan.allow_rows] = rows.masked_fill(disallowed_mask(plan.allow_ids, vocab_size), neg_inf) + + 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 +175,134 @@ class Sampler: device: torch.device vocab_size: int + def _plan(self, batch: Batch, params) -> LogitsPlan | None: + """The LogitsPlan for this batch, or None when no request asks for a processor.""" + MIN_T = 1e-6 + rows = [i for i, p in enumerate(params) if p.needs_logits_processing] + if not rows: + return None + picked = [(batch.reqs[i], params[i]) for i in rows] + temps = [[max(0.0 if p.is_greedy else p.temperature, MIN_T)] for _, p in picked] + plan = LogitsPlan( + rows=make_device_tensor(rows, torch.int64, self.device), + temps=make_device_tensor(temps, torch.float32, self.device), + ) + pad = self.vocab_size + + bias_rows: list[int] = [] + bias_ids: list[int] = [] + bias_vals: list[float] = [] + for local, (_, p) in enumerate(picked): + for tid, val in p.logit_bias or (): + if 0 <= int(tid) < self.vocab_size: + bias_rows.append(local) + bias_ids.append(int(tid)) + bias_vals.append(float(val)) + if bias_ids: + plan.bias_rows = make_device_tensor(bias_rows, torch.int64, self.device) + plan.bias_ids = make_device_tensor(bias_ids, torch.int64, self.device) + plan.bias_vals = make_device_tensor(bias_vals, torch.float32, self.device) + + min_rows: list[int] = [] + min_ids: list[list[int]] = [] + for local, (req, p) in enumerate(picked): + if p.min_tokens > 0 and p.min_tokens_stop_ids: + # device_len, not the host ids: under overlap scheduling the previous token + # reaches req.input_ids only after this batch is prepared + generated = req.device_len - (req.max_device_len - req.output_len) + if generated < p.min_tokens: + min_rows.append(local) + min_ids.append([t for t in p.min_tokens_stop_ids if 0 <= t < self.vocab_size]) + if min_rows: + plan.min_rows = make_device_tensor(min_rows, torch.int64, self.device) + plan.min_ids = _padded_rows(min_ids, pad, self.device) + + allow = [(local, self._allowed(p)) for local, (_, p) in enumerate(picked)] + allow = [(local, ids) for local, ids in allow if ids] + if allow: + plan.allow_rows = make_device_tensor([local for local, _ in allow], torch.int64, self.device) + plan.allow_ids = _padded_rows([ids for _, ids in allow], 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 _allowed(self, p) -> list[int]: + # an empty list after the filter leaves the row unmasked rather than all -inf + return [t for t in p.allowed_token_ids or () if 0 <= t < self.vocab_size] + + def _prompt_mask(self, req) -> torch.Tensor: + """[V] bool marking the request's prompt ids, uploaded once per request.""" + prompt_len = req.max_device_len - req.output_len + # multimodal placeholders are pseudo-ids >= vocab_size (MM_PAD_SHIFT_VALUE); + # send them to the scratch column, not past the scatter target + ids = req.input_ids[:prompt_len].to(torch.int64).clamp(max=self.vocab_size) + if self.device.type == "cuda": + ids = ids.pin_memory().to(self.device, non_blocking=True) + mask = torch.zeros(self.vocab_size + 1, dtype=torch.bool, device=self.device) + mask[ids] = True + return mask[: self.vocab_size] + def prepare(self, batch: Batch) -> BatchSamplingArgs: params = [r.sampling_params for r in batch.reqs] is_greedy = [p.is_greedy for p in params] + penalties = [] + for row, req in enumerate(batch.reqs): + p = req.sampling_params + repetition = p.repetition_penalty + if not (p.presence_penalty or p.frequency_penalty or repetition != 1.0): + continue + if not req.can_decode: + continue + if req.output_token_counts is None: + req.output_token_counts = torch.zeros( + self.vocab_size, dtype=torch.int32, device=self.device + ) + prompt_mask = None + if repetition != 1.0: + if req.prompt_token_mask is None: + req.prompt_token_mask = self._prompt_mask(req) + prompt_mask = req.prompt_token_mask + penalties.append( + RowPenalty( + row, + req.output_token_counts, + p.presence_penalty, + p.frequency_penalty, + repetition, + prompt_mask, + ) + ) + plan = self._plan(batch, params) + want_logprobs = [p.logprobs for p in params] + logprob_all_rows = bool(want_logprobs) and all(want_logprobs) + logprob_row_idx = ( + make_device_tensor( + [i for i, want in enumerate(want_logprobs) if want], torch.int64, self.device + ) + if any(want_logprobs) and not logprob_all_rows + else None + ) + max_top_logprobs = min( + max((p.top_logprobs for p in params if p.logprobs), default=0), self.vocab_size + ) + restrict = [(i, self._allowed(p)) for i, p in enumerate(p for p in params if p.logprobs)] + restrict = [(i, ids) for i, ids in restrict if ids] + restrict_rows = restrict_ids = None + if restrict: + restrict_rows = make_device_tensor([i for i, _ in restrict], torch.int64, self.device) + restrict_ids = _padded_rows([ids for _, ids in restrict], self.vocab_size, self.device) if all(is_greedy): - return BatchSamplingArgs(temperatures=None) + return BatchSamplingArgs( + temperatures=None, + penalties=penalties, + plan=plan, + logprob_row_idx=logprob_row_idx, + logprob_all_rows=logprob_all_rows, + max_top_logprobs=max_top_logprobs, + restrict_rows=restrict_rows, + restrict_ids=restrict_ids, + ) MIN_P = MIN_T = 1e-6 # Greedy outputs are selected explicitly in sample(); use neutral sampling @@ -83,17 +325,108 @@ 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, + penalties=penalties, + plan=plan, + logprob_row_idx=logprob_row_idx, + logprob_all_rows=logprob_all_rows, + max_top_logprobs=max_top_logprobs, + restrict_rows=restrict_rows, + restrict_ids=restrict_ids, + ) @nvtx_annotate("Sampler") def sample(self, logits: torch.Tensor, args: BatchSamplingArgs) -> torch.Tensor: with torch.cuda.nvtx.range("Sampler"): + if args.penalties: + logits = apply_penalties(logits, args.penalties) + # after the penalties: min_p must see the final distribution + if args.plan is not None: + logits = apply_logits_processors(logits, args.plan, self.vocab_size) if args.temperatures is None: # greedy sampling - return torch.argmax(logits, dim=-1) - tokens = sample_impl(logits.float(), args.temperatures, args.top_k, args.top_p) - if args.greedy_mask is not None: - # Mixed batches still run probability sampling for all rows, but - # greedy rows must follow argmax's deterministic tie-breaking. - greedy_tokens = torch.argmax(logits, dim=-1).to(tokens.dtype) - tokens = torch.where(args.greedy_mask, greedy_tokens, tokens) + tokens = torch.argmax(logits, dim=-1) + else: + tokens = sample_impl(logits.float(), args.temperatures, args.top_k, args.top_p) + if args.greedy_mask is not None: + # Mixed batches still run probability sampling for all rows, but + # greedy rows must follow argmax's deterministic tie-breaking. + greedy_tokens = torch.argmax(logits, dim=-1).to(tokens.dtype) + tokens = torch.where(args.greedy_mask, greedy_tokens, tokens) + # Update on the sampling stream: overlapped scheduling can prepare the next + # batch before the previous token reaches Req.input_ids on the CPU. + for row, counts, *_ in args.penalties: + counts.scatter_add_( + 0, tokens[row : row + 1].long(), counts.new_ones(1) + ) return tokens + + def compute_logprobs( + self, + logits: torch.Tensor, + sampled_tokens: torch.Tensor, + args: BatchSamplingArgs, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None: + rows = args.logprob_row_idx + if not args.logprob_all_rows and (rows is None or rows.numel() == 0): + return None + + request_logits = (logits if rows is None else logits.index_select(0, rows)).float() + if args.restrict_rows is not None: + # not in place: on an fp32 batch request_logits can be the engine's logits view + masked = request_logits[args.restrict_rows].masked_fill( + disallowed_mask(args.restrict_ids, self.vocab_size), float("-inf") + ) + request_logits = request_logits.index_put((args.restrict_rows,), masked) + # Reported values are raw model logprobs (pre-temperature log_softmax over logits), + # renormalized over allowed_token_ids on the restricted rows. + request_logprobs = torch.log_softmax(request_logits, dim=-1) + + request_tokens = sampled_tokens.to(dtype=torch.long, device=logits.device) + if rows is not None: + request_tokens = request_tokens.index_select(0, rows) + request_row_idx = torch.arange(request_logprobs.shape[0], device=logits.device) + request_chosen_logprobs = request_logprobs[request_row_idx, request_tokens] + + if rows is None: + chosen_logprobs = request_chosen_logprobs + else: + chosen_logprobs = torch.full( + (logits.shape[0],), float("nan"), dtype=torch.float32, device=logits.device + ) + chosen_logprobs.index_copy_(0, 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 + ) + if rows is None: + top_ids = request_top_ids.to(torch.int32) + top_logprobs = request_top_logprobs + else: + 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[rows] = request_top_ids.to(torch.int32) + top_logprobs[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..03e9388024 100644 --- a/python/freetoken/env.py +++ b/python/freetoken/env.py @@ -75,6 +75,27 @@ 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) + # cuGraphUpload each decode graph at capture (#697); off keeps the first-launch upload, which + # replays ~0.2 ms/step faster on Qwen3.8-Flash / RTX 5080 but can block the first launch on a stream WAIT + GRAPH_UPLOAD = EnvBool(True) + # 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/__main__.py b/python/freetoken/kernel/__main__.py index 7be541a67c..5b66d484df 100644 --- a/python/freetoken/kernel/__main__.py +++ b/python/freetoken/kernel/__main__.py @@ -6,27 +6,32 @@ def generate_clangd(): import subprocess from freetoken.kernel.utils import DEFAULT_INCLUDE - from freetoken.utils import init_logger + from freetoken.utils import get_rocm_gfx_arch, init_logger, is_rocm from tvm_ffi.libinfo import find_dlpack_include_path, find_include_path logger = init_logger(__name__) logger.info("Generating .clangd file...") include_paths = [find_include_path(), find_dlpack_include_path()] + DEFAULT_INCLUDE - status = subprocess.run( - args=["nvidia-smi", "--query-gpu=compute_cap", "--format=csv,noheader"], - capture_output=True, - check=True, - ) - compute_cap = status.stdout.decode("utf-8").strip().split("\n")[0] - major, minor = compute_cap.split(".") + + # TODO(ROCm): hiprtc JIT cache should be separate from nvcc JIT cache to avoid stale binaries. + if is_rocm(): + arch_flags = ["-xhip", f"--offload-arch={get_rocm_gfx_arch() or 'gfx1201'}"] + else: + try: + status = subprocess.run( + args=["nvidia-smi", "--query-gpu=compute_cap", "--format=csv,noheader"], + capture_output=True, + check=True, + ) + compute_cap = status.stdout.decode("utf-8").strip().split("\n")[0] + major, minor = compute_cap.split(".") + except (subprocess.CalledProcessError, FileNotFoundError, ValueError): + import torch + + major, minor = torch.cuda.get_device_capability() + arch_flags = ["-xcuda", f"--cuda-gpu-arch=sm_{major}{minor}"] compile_flags = ",\n ".join( - [ - "-xcuda", - f"--cuda-gpu-arch=sm_{major}{minor}", - "-std=c++20", - "-Wall", - "-Wextra", - ] + arch_flags + ["-std=c++20", "-Wall", "-Wextra"] + [f"-isystem{path}" for path in include_paths] ) clangd_content = f""" diff --git a/python/freetoken/kernel/_toolchain.py b/python/freetoken/kernel/_toolchain.py index b49cebbb43..93d0322151 100644 --- a/python/freetoken/kernel/_toolchain.py +++ b/python/freetoken/kernel/_toolchain.py @@ -1,4 +1,4 @@ -"""CUDA toolchain/torch consistency checks. +"""CUDA/HIP toolchain/torch consistency checks. Standalone on purpose: setup.py and the kernel-cache build backend load this file by path, so it must not import the freetoken package. @@ -16,6 +16,11 @@ _TRUE_VALUES = {"1", "true", "yes", "on"} +def _is_rocm() -> bool: + import torch + return getattr(torch.version, "hip", None) is not None + + def _nvcc_path() -> str | None: from torch.utils.cpp_extension import CUDA_HOME @@ -49,6 +54,8 @@ def check_nvcc_matches_torch() -> None: nvcc-built binaries link libcudart.so.; at runtime only the torch wheel's own CUDA runtime is guaranteed to be loadable. """ + if _is_rocm(): + return # ROCm uses hipcc, not nvcc if os.getenv(ALLOW_MISMATCH_ENV, "").strip().lower() in _TRUE_VALUES: return torch_major = torch_cuda_major() diff --git a/python/freetoken/kernel/aot.py b/python/freetoken/kernel/aot.py index 5e87f8c923..c9e378d5b6 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( @@ -143,6 +160,15 @@ def build(build_directory: pathlib.Path) -> object: return KernelSpec(name=_make_name("radix"), build=build) +def _ggml_spec() -> KernelSpec: + def build(build_directory: pathlib.Path) -> object: + from .ggml import load_module + + return load_module(str(build_directory)) + + return KernelSpec(name=_make_name("ggml"), build=build) + + def default_kernel_specs() -> tuple[KernelSpec, ...]: specs: list[KernelSpec] = [] specs.extend(_store_spec(element_size) for element_size in DEFAULT_STORE_ELEMENT_SIZES) @@ -153,10 +179,14 @@ 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()) + specs.append(_ggml_spec()) names = [spec.name for spec in specs] duplicates = {name for name in names if names.count(name) > 1} diff --git a/python/freetoken/kernel/aot_models.py b/python/freetoken/kernel/aot_models.py index a8ef85d432..1c7f15c84a 100644 --- a/python/freetoken/kernel/aot_models.py +++ b/python/freetoken/kernel/aot_models.py @@ -11,12 +11,14 @@ (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. + DSV4 (plain nn.Embedding) and GGUF embeddings (ggml get_rows) bypass it. The whole table targets the shipped serving configuration: TP=1 (TP>1 shards kv heads, shrinking the store row) and a 2-byte compute dtype (``--dtype @@ -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 @@ -91,11 +94,8 @@ def expert_bank_row_bytes(fmt: str, hidden_size: int, moe_intermediate_size: int "down": H * I, "down_scale": (H // B) * fp8_block_scale_pad(H // B, I // B) * 2, } - if fmt == "q4_0": - # gemma4/gguf.py _q4_0_expert_specs: GGML Q4_0 rows, 32 elems -> 18 bytes - return {"gate_up": 2 * I * (H // 32 * 18), "down": H * (I // 32 * 18)} if fmt in ("nvfp4", "nvfp4_marlin", "nvfp4_b12x"): - # models/nvfp4_banks.py: packed e2m1 pairs + per-16 fp8-e4m3 scales + fp16 + # models/nvfp4_banks.py: packed e2m1 pairs + per-16 fp8-e4m3 scales + fp32 # per-row globals; marlin/b12x repacks are byte-identical with the globals # folded into GPU-resident alphas (layers/quantization/moe/nvfp4.py), so no global banks. banks = { @@ -105,8 +105,8 @@ def expert_bank_row_bytes(fmt: str, hidden_size: int, moe_intermediate_size: int "down_scale": H * (I // 16), } if fmt == "nvfp4": - banks["gate_up_global"] = 2 * I * 2 - banks["down_global"] = H * 2 + banks["gate_up_global"] = 2 * I * 4 + banks["down_global"] = H * 4 return banks if fmt == "mxfp4_triton": # gpt-oss mxfp4 expert banks: transposed split-K blocks/scales + bf16 bias @@ -145,7 +145,7 @@ def expert_bank_row_bytes(fmt: str, hidden_size: int, moe_intermediate_size: int moe_intermediate_size=768, expert_formats=("bf16",), aliases=("Qwen/Qwen3-30B-A3B-Thinking-2507",), - arch_aliases=("Qwen3VLMoeForConditionalGeneration",), # Qwen3-VL-30B-A3B: same text tower + arch_aliases=("Qwen3VLMoeForConditionalGeneration", "Qwen3MoeGGUFForCausalLM"), # Qwen3-VL-30B-A3B: same text tower ), AotModel( name="Qwen/Qwen3.5-35B-A3B", @@ -155,7 +155,7 @@ def expert_bank_row_bytes(fmt: str, hidden_size: int, moe_intermediate_size: int top_k=8, moe_intermediate_size=512, expert_formats=("bf16",), - arch_aliases=("Qwen3_5MoeForCausalLM",), # text-only release of the same tower + arch_aliases=("Qwen3_5MoeForCausalLM", "Qwen3_5MoeGGUFForCausalLM"), # text-only release of the same tower ), AotModel( name="Qwen/Qwen3.5-35B-A3B-FP8", @@ -205,6 +205,8 @@ def expert_bank_row_bytes(fmt: str, hidden_size: int, moe_intermediate_size: int top_k=10, moe_intermediate_size=640, expert_formats=(*_NVFP4_FORMATS, "fp8_block"), + # GGUF expert rows differ per layer and reach the fused copy as runtime sizes + arch_aliases=("Qwen4ExpGGUFForCausalLM",), ), AotModel( name="google/gemma-4-26B-A4B-it", @@ -215,9 +217,8 @@ def expert_bank_row_bytes(fmt: str, hidden_size: int, moe_intermediate_size: int kv_groups=((2, 512), (8, 256)), top_k=8, moe_intermediate_size=704, - # q4_0 = the Q4_0 GGUF release (Gemma4GGUFForCausalLM); its GGUFEmbedding - # path skips indexing(), but the bf16 safetensors path needs the variant. - expert_formats=("bf16", "q4_0"), + expert_formats=("bf16",), + # GGUF expert rows differ per layer and reach the fused copy as runtime sizes arch_aliases=("Gemma4ForCausalLM", "Gemma4GGUFForCausalLM"), ), AotModel( @@ -255,6 +256,7 @@ def expert_bank_row_bytes(fmt: str, hidden_size: int, moe_intermediate_size: int top_k=8, moe_intermediate_size=1536, expert_formats=("bf16",), + arch_aliases=("Glm4MoeGGUFForCausalLM",), ), AotModel( name="nvidia/GLM-4.7-NVFP4", @@ -282,7 +284,8 @@ def expert_bank_row_bytes(fmt: str, hidden_size: int, moe_intermediate_size: int # KDA conv/recurrent state lives in the LinearStatePool, not paged KV. name="RedHatAI/GLM-5.3-Flash-NVFP4", architecture="Glm5NextForCausalLM", - arch_aliases=("Glm5NextForConditionalGeneration",), + # GGUF expert rows differ per layer and reach the fused copy as runtime sizes + arch_aliases=("Glm5NextForConditionalGeneration", "Glm5NextGGUFForCausalLM"), hidden_size=4096, kv_groups=(), top_k=8, @@ -302,6 +305,7 @@ def expert_bank_row_bytes(fmt: str, hidden_size: int, moe_intermediate_size: int moe_intermediate_size=1536, # experts reuse the dense intermediate_size expert_formats=_NVFP4_FORMATS, aliases=("MiniMaxAI/MiniMax-M2.5",), + arch_aliases=("MiniMaxM2GGUFForCausalLM",), ), AotModel( # GQA block-sparse (BSA): the BSAKVCache stores K/V through store_cache @@ -320,7 +324,7 @@ def expert_bank_row_bytes(fmt: str, hidden_size: int, moe_intermediate_size: int moe_intermediate_size=3072, expert_formats=("nvfp4",), aliases=("MiniMaxAI/MiniMax-M3",), - arch_aliases=("MiniMaxM3SparseForCausalLM",), + arch_aliases=("MiniMaxM3SparseForCausalLM", "MiniMaxM3GGUFForCausalLM"), ), AotModel( name="deepseek-ai/DeepSeek-V4-Flash", @@ -332,6 +336,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", @@ -339,7 +354,7 @@ def expert_bank_row_bytes(fmt: str, hidden_size: int, moe_intermediate_size: int hidden_size=5120, kv_groups=((4, 256),), aliases=("Qwen/Qwen3.6-27B-FP8", "nvidia/Qwen3.6-27B-NVFP4"), - arch_aliases=("Qwen3_5ForCausalLM",), # text-only release of the same tower + arch_aliases=("Qwen3_5ForCausalLM", "Qwen3_5GGUFForCausalLM"), # text-only release of the same tower ), AotModel( name="google/gemma-4-12B-it", @@ -360,25 +375,28 @@ def expert_bank_row_bytes(fmt: str, hidden_size: int, moe_intermediate_size: int hidden_size=6656, kv_groups=((2, 128),), # sliding and full (NoPE) layers share the same geometry aliases=("RedHatAI/Muse-Glimmer-30B-NVFP4",), + arch_aliases=("MuseGlimmerGGUFForCausalLM",), ), AotModel( name="meta-llama/Llama-3.1-8B-Instruct", architecture="LlamaForCausalLM", hidden_size=4096, kv_groups=((8, 128),), + arch_aliases=("LlamaGGUFForCausalLM",), ), AotModel( name="Qwen/Qwen2-7B", architecture="Qwen2ForCausalLM", hidden_size=3584, kv_groups=((4, 128),), + arch_aliases=("Qwen2GGUFForCausalLM",), ), AotModel( name="Qwen/Qwen3-8B", architecture="Qwen3ForCausalLM", hidden_size=4096, kv_groups=((8, 128),), - arch_aliases=("Qwen3VLForConditionalGeneration",), # Qwen3-VL-8B: same text tower + arch_aliases=("Qwen3VLForConditionalGeneration", "Qwen3GGUFForCausalLM"), # Qwen3-VL-8B: same text tower ), AotModel( name="mistralai/Mistral-7B-Instruct-v0.3", @@ -395,8 +413,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 +437,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/backend.py b/python/freetoken/kernel/backend.py index 2a886cf77d..73a7dc6f36 100644 --- a/python/freetoken/kernel/backend.py +++ b/python/freetoken/kernel/backend.py @@ -4,6 +4,7 @@ ops; otherwise they fall back to the pure-Triton kernels in ``freetoken.kernel.triton``. ``find_spec`` only checks that the package is importable (no import side effects), and the result is cached. +On ROCm, the optional-package probes return False so callers use the fallbacks. """ from __future__ import annotations @@ -25,16 +26,22 @@ def _importable(name: str) -> bool: @functools.cache def is_flashinfer_installed() -> bool: + if is_rocm(): + return False return _importable("flashinfer") @functools.cache def is_sgl_kernel_installed() -> bool: + if is_rocm(): + return False return _importable("sgl_kernel") @functools.cache def is_vllm_installed() -> bool: + if is_rocm(): + return False return _importable("vllm") @@ -47,6 +54,13 @@ def device_capability() -> tuple[int, int]: return (int(major), int(minor)) +@functools.cache +def is_rocm() -> bool: + """True when torch is built for ROCm (AMD GPU).""" + import torch + return getattr(torch.version, "hip", None) is not None + + @functools.cache def driver_cuda_version() -> int | None: """Max CUDA version the installed NVIDIA driver supports (``13000`` == CUDA 13.0), @@ -55,6 +69,8 @@ def driver_cuda_version() -> int | None: toolkit version. Resolved through the ``_pinned_tensor`` extension's link-time cudart, so it works wherever the extension builds (including Windows) -- no dlopen by soname.""" + if is_rocm(): + return None try: from freetoken.kernel.pinned import _load_pinned_extension diff --git a/python/freetoken/kernel/batch_memcpy.py b/python/freetoken/kernel/batch_memcpy.py index b39e5cde24..67e2bf9870 100644 --- a/python/freetoken/kernel/batch_memcpy.py +++ b/python/freetoken/kernel/batch_memcpy.py @@ -28,6 +28,12 @@ def _probe(fn) -> None: src = torch.arange(16, dtype=torch.uint8).pin_memory() dst = torch.zeros(16, dtype=torch.uint8, device="cuda") stream = torch.cuda.Stream() + # torch.zeros filled dst on the CURRENT stream; the copy below runs on this one. + # Without the join the fill can land after the copy and clobber it, leaving the + # probe to read back zeros -- a false "copied wrong bytes" on hardware that + # supports the API. A cold process hides this (the first torch.zeros and + # torch.cuda.Stream() each synchronize the device); a warmed-up server does not. + stream.wait_stream(torch.cuda.current_stream()) fn( torch.tensor([dst.data_ptr()]), torch.tensor([src.data_ptr()]), diff --git a/python/freetoken/kernel/causal_conv1d.py b/python/freetoken/kernel/causal_conv1d.py index 81f1fab7b4..c196658648 100644 --- a/python/freetoken/kernel/causal_conv1d.py +++ b/python/freetoken/kernel/causal_conv1d.py @@ -22,9 +22,19 @@ def causal_conv1d_varlen( cu_seqlens: torch.Tensor, # [batch+1] int32 prefix sums of per-request lengths cache_indices: torch.Tensor, # [batch] int32 slot id per request has_initial_state: torch.Tensor, # [batch] bool (carry conv state across chunks) + max_seq_len: int | None = None, # host-known longest extend_len; None derives it on device ) -> torch.Tensor: - """Varlen (prefill) depthwise causal conv with silu; writes silu(conv) into ``x`` - in place and refreshes ``conv_states[cache_indices]`` with each request's tail.""" + """Varlen (prefill) depthwise causal conv with silu; returns silu(conv) and refreshes + ``conv_states[cache_indices]`` with each request's tail. + + Whether ``x`` survives depends on the backend: sgl_kernel convolves ``x`` in place and + returns it, while the triton fallback leaves ``x`` untouched and returns a new tensor. + Callers must use the returned tensor and must not rely on ``x`` afterwards (every + current call site builds ``x`` fresh and never reads it again). + + ``max_seq_len`` only sizes the triton launch grid. Deriving it from ``cu_seqlens`` costs a + D2H sync per prefill, which is illegal under CUDA graph capture, so the scheduler passes the + host value it already has.""" from freetoken.kernel.backend import is_sgl_kernel_installed if not is_sgl_kernel_installed(): @@ -33,7 +43,8 @@ def causal_conv1d_varlen( ) return triton_causal_conv1d_varlen( - x, weight, conv_states, cu_seqlens, cache_indices, has_initial_state + x, weight, conv_states, cu_seqlens, cache_indices, has_initial_state, + max_seq_len=max_seq_len, ) from sgl_kernel import causal_conv1d_fwd diff --git a/python/freetoken/kernel/csrc/cpu_moe/common.h b/python/freetoken/kernel/csrc/cpu_moe/common.h new file mode 100644 index 0000000000..0edee0150a --- /dev/null +++ b/python/freetoken/kernel/csrc/cpu_moe/common.h @@ -0,0 +1,208 @@ +// Numeric helpers and ISA detection shared by the CPU MoE formats. +#pragma once + +#include +#include +#include +#include +#include +#include +#include + +#if defined(__x86_64__) || defined(__i386__) +#include +#define CPU_MOE_X86 1 +#else +#define CPU_MOE_X86 0 +#endif + +// Each SIMD kernel is compiled for its own ISA and picked at runtime (pick_isa), so the build sets no global -march. +#define CPU_MOE_TARGET(...) __attribute__((target(__VA_ARGS__))) + +#if CPU_MOE_X86 && ((defined(__GNUC__) && __GNUC__ >= 10) || defined(__clang__)) +#define CPU_MOE_HAS_AVX512BF16 1 +#define CPU_MOE_HAS_AVX512VNNI 1 +#endif + +namespace freetoken::cpu_moe { + +// Scratch rows are read with 64-byte SIMD loads, so their base is cache-line aligned whatever the heap state; std::vector alone guarantees 16 bytes. +template +struct CacheLineAllocator { + using value_type = T; + CacheLineAllocator() = default; + template + CacheLineAllocator(const CacheLineAllocator&) {} + T* allocate(size_t n) { return static_cast(::operator new(n * sizeof(T), std::align_val_t(64))); } + void deallocate(T* p, size_t) { ::operator delete(p, std::align_val_t(64)); } + template + bool operator==(const CacheLineAllocator&) const { return true; } +}; + +template +using AlignedVector = std::vector>; + +using bf16_t = uint16_t; + +inline float bf16_to_f32(bf16_t v) { + uint32_t u = static_cast(v) << 16; + float f; + std::memcpy(&f, &u, sizeof(f)); + return f; +} + +inline bf16_t f32_to_bf16(float f) { + uint32_t u; + std::memcpy(&u, &f, sizeof(u)); + // round-to-nearest-even + const uint32_t lsb = (u >> 16) & 1u; + u += 0x7fffu + lsb; + return static_cast(u >> 16); +} + +// ACT_SWIGLUOAI is the clamped (up + 1) swiglu (gpt-oss "swigluoai" / +// MiniMax-M3): gate/up are combined jointly with the runtime alpha/limit +// scalars, so it is handled in the gated_rows epilogue (act_apply never sees it; +// the mxfp4 kernel additionally fuses its own copy of the same math). +// ACT_SWIGLU_CLAMP (GLM-5.3 "swiglu_limit") is the same clamped form WITHOUT +// the (up + 1) bias: clamp(gate, max=lim) * sigmoid(alpha*gate) * clamp(up, +-lim). +enum ActKind { + ACT_SILU = 0, + ACT_GELU = 1, + ACT_GELU_TANH = 2, + ACT_SWIGLUOAI = 3, + ACT_SWIGLU_CLAMP = 4, +}; + +inline float act_apply(int act, float x) { + if (act == ACT_SILU) return x / (1.0f + std::exp(-x)); + if (act == ACT_GELU) + return 0.5f * x * (1.0f + std::erf(x * 0.70710678118654752440f)); + // gelu_tanh + const float k0 = 0.79788456080286535588f; // sqrt(2/pi) + const float inner = k0 * (x + 0.044715f * x * x * x); + return 0.5f * x * (1.0f + std::tanh(inner)); +} + +// Software prefetch distance (bytes) ahead of the current weight row stream. The +// weight stream is the bandwidth bottleneck (read once, never reused); nudging the +// HW prefetcher with a few cache lines of lookahead raises sustained throughput. +constexpr int PF_AHEAD = 512; + +const float kE2M1[16] = {0.0f, 0.5f, 1.0f, 1.5f, 2.0f, 3.0f, 4.0f, 6.0f, + -0.0f, -0.5f, -1.0f, -1.5f, -2.0f, -3.0f, -4.0f, -6.0f}; + +// e4m3 (OCP "fn": finite, max-normal 448, exp bias 7). Decoded into a 256-entry LUT. +inline float e4m3_decode(uint8_t v) { + const float sign = (v & 0x80u) ? -1.0f : 1.0f; + const uint32_t exp = (v >> 3) & 0xFu; + const uint32_t man = v & 0x7u; + if (exp == 0) return sign * (man / 8.0f) * 0.015625f; // 2^(1-7) = 2^-6 + return sign * (1.0f + man / 8.0f) * std::ldexp(1.0f, (int)exp - 7); +} + +#if CPU_MOE_X86 +CPU_MOE_TARGET("avx2,fma") +inline float hsum256(__m256 v) { + __m128 lo = _mm256_castps256_ps128(v); + lo = _mm_add_ps(lo, _mm256_extractf128_ps(v, 1)); + lo = _mm_add_ps(lo, _mm_movehl_ps(lo, lo)); + lo = _mm_add_ss(lo, _mm_shuffle_ps(lo, lo, 0x55)); + return _mm_cvtss_f32(lo); +} + +// AVX2 e2m1 nibble decode: codes (0..15) in 8 int32 lanes -> fp32. AVX2 vpermps is +// only 8-wide, so instead of a 16-entry LUT we use the e2m1 sign/magnitude symmetry: +// value = (code&8 ? - : +) * mag8[code&7], mag8 = e2m1[0..7]. The sign is bit 3 of +// the code shifted into the fp32 sign bit (bit 31). Bit-identical to the e2m1 LUT. +CPU_MOE_TARGET("avx2,fma") +inline __m256 e2m1_decode8(__m256i codes, __m256 mag8) { + __m256 mag = _mm256_permutevar8x32_ps(mag8, _mm256_and_si256(codes, _mm256_set1_epi32(7))); + __m256i sgn = _mm256_slli_epi32(_mm256_and_si256(codes, _mm256_set1_epi32(8)), 28); + return _mm256_xor_ps(mag, _mm256_castsi256_ps(sgn)); +} +#endif // CPU_MOE_X86 + +// SIMD tiers, ascending. Each format picks the highest tier <= the one chosen by +// pick_isa() that it implements (fp4 and fp8_block formats have no bf16-specific tier, so the +// avx512bf16 tier maps to their avx512 kernel). +enum IsaTier { ISA_SCALAR = 0, ISA_AVX2 = 1, ISA_AVX512 = 2, ISA_AVX512BF16 = 3 }; + +// Best tier the CPU+build supports, optionally capped DOWN by +// FREETOKEN_CPU_MOE_ISA={scalar,avx2,avx512,avx512bf16} (A/B testing on a machine +// that supports more). FREETOKEN_CPU_MOE_SCALAR=1 forces scalar (legacy alias). +inline IsaTier pick_isa() { +#if CPU_MOE_X86 + if (getenv("FREETOKEN_CPU_MOE_SCALAR")) return ISA_SCALAR; + IsaTier best = ISA_SCALAR; + if (__builtin_cpu_supports("avx2") && __builtin_cpu_supports("fma")) best = ISA_AVX2; + if (best >= ISA_AVX2 && __builtin_cpu_supports("avx512f")) best = ISA_AVX512; +#ifdef CPU_MOE_HAS_AVX512BF16 + if (best >= ISA_AVX512 && __builtin_cpu_supports("avx512bf16")) best = ISA_AVX512BF16; +#endif + if (const char* f = getenv("FREETOKEN_CPU_MOE_ISA")) { + IsaTier want = best; + if (!std::strcmp(f, "scalar")) want = ISA_SCALAR; + else if (!std::strcmp(f, "avx2")) want = ISA_AVX2; + else if (!std::strcmp(f, "avx512")) want = ISA_AVX512; + else if (!std::strcmp(f, "avx512bf16")) want = ISA_AVX512BF16; + if (want < best) best = want; // cap downward; never force above hw/build support + } + return best; +#else + return ISA_SCALAR; +#endif +} + +inline const char* tier_name(IsaTier t) { +#if CPU_MOE_X86 +#ifdef CPU_MOE_HAS_AVX512BF16 + if (t >= ISA_AVX512BF16) return "avx512bf16"; +#endif + if (t >= ISA_AVX512) return "avx512f"; + if (t >= ISA_AVX2) return "avx2"; +#endif + (void)t; + return "scalar"; +} + +// AVX-VNNI (VEX-256 VPDPBUSD) availability: Alder/Raptor Lake, Sapphire Rapids+, Zen5. +// Distinct from AVX-512 VNNI. Opt out with FREETOKEN_CPU_MOE_NO_VNNI=1 (A/B the W4A8 path). +inline bool cpu_has_avxvnni() { +#if CPU_MOE_X86 + const char* no = getenv("FREETOKEN_CPU_MOE_NO_VNNI"); + if (no && no[0] && no[0] != '0') return false; // ignore unset/empty/"0" + return __builtin_cpu_supports("avx2") && __builtin_cpu_supports("avxvnni"); +#else + return false; +#endif +} + +// AVX-512 VNNI (512-bit VPDPBUSD): Cascade Lake+, Ice Lake, Sapphire/Emerald Rapids, +// Zen4+. 2x the 256-bit AVX-VNNI width. FREETOKEN_CPU_MOE_NO_AVX512VNNI=1 forces the +// 256-bit path (A/B the two W4A8 kernels on the same box); FREETOKEN_CPU_MOE_NO_VNNI=1 +// still disables the whole W4A8 family (back to the faithful fp32 nvdot), so it is +// honored here too. +inline bool cpu_has_avx512vnni() { +#if CPU_MOE_X86 && defined(CPU_MOE_HAS_AVX512VNNI) + const char* no = getenv("FREETOKEN_CPU_MOE_NO_AVX512VNNI"); + if (no && no[0] && no[0] != '0') return false; + const char* no_vnni = getenv("FREETOKEN_CPU_MOE_NO_VNNI"); + if (no_vnni && no_vnni[0] && no_vnni[0] != '0') return false; + return __builtin_cpu_supports("avx512vnni"); +#else + return false; +#endif +} + +// Split an interleaved bf16 row into fp32 even/odd halves (even[m]=src[2m]). +// bf16->fp32 is exact, so this only reorders -- done once per token/route and +// reused across every output row of the GEMV. +inline void deinterleave_bf16_f32(const bf16_t* src, float* even, float* odd, int K) { + for (int m = 0; m < K / 2; ++m) { + even[m] = bf16_to_f32(src[2 * m]); + odd[m] = bf16_to_f32(src[2 * m + 1]); + } +} + +} // namespace freetoken::cpu_moe diff --git a/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp b/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp index 8349123bd6..8162138ceb 100644 --- a/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp +++ b/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp @@ -1,35 +1,31 @@ -// CPU-compute MoE executor for the "cpu" offload backend. +// CPU-compute MoE executor for the cpu and hybrid MoE strategies, and its Python binding. // // Decode ships activations to the CPU, computes the routed experts here (reading -// the pinned host expert banks at full RAM bandwidth), and ships the results -// back. To keep the whole decode path inside a single CUDA graph we expose -// submit/sync as host nodes via cudaLaunchHostFunc -- the callbacks only touch a -// CPU worker pool + pinned host buffers and never call any CUDA API. +// the pinned host expert banks at full RAM bandwidth), and ships the results back. +// One task (one MoE layer) is in flight at a time: submit() wakes the persistent worker +// pool, the pool runs pass 1 (gate_up) and pass 2 (down), and sync() waits for it. +// The GPU rings a doorbell to submit and waits on it to sync (GpuSync, gpu_sync.h), so the +// whole decode step stays in one CUDA graph; cudaLaunchHostFunc nodes are the fallback. // -// One task is in flight at a time (per MoE layer): submit() wakes the pool, -// sync() blocks the host-func thread until the pool drains. The heavy GEMV runs -// on the persistent worker threads, not the host-func thread. -// -// Weight formats: bf16, NVFP4, MXFP4, ds_fp4 and Q4_0 expert banks (see WFmt and -// the per-format bank schemas). Compute is FP32-accumulate; the intermediate is -// stored bf16 to match the GPU decode path. ISA is chosen once at construction -// (AVX-512-BF16 dpbf16 -> AVX-512F widening -> AVX2+FMA -> scalar). +// Weight formats live in kernels/, one file each behind the MoeFormat interface (format.h). +// Compute is FP32-accumulate; the intermediate is stored bf16 to match the GPU decode path. #include #include #include -#include #include #include -#include #include +#include #include #include #include #include +#include +#include #include -#include +#include #include #if defined(__linux__) @@ -40,1501 +36,175 @@ #define CPU_MOE_HAS_AFFINITY 0 #endif -#if defined(__x86_64__) || defined(__i386__) -#include -#define CPU_MOE_X86 1 -#else -#define CPU_MOE_X86 0 -#endif - -namespace { - -using bf16_t = uint16_t; - -inline float bf16_to_f32(bf16_t v) { - uint32_t u = static_cast(v) << 16; - float f; - std::memcpy(&f, &u, sizeof(f)); - return f; -} - -inline bf16_t f32_to_bf16(float f) { - uint32_t u; - std::memcpy(&u, &f, sizeof(u)); - // round-to-nearest-even - const uint32_t lsb = (u >> 16) & 1u; - u += 0x7fffu + lsb; - return static_cast(u >> 16); -} - -// ACT_SWIGLUOAI is the clamped (up + 1) swiglu (gpt-oss "swigluoai" / -// MiniMax-M3): gate/up are combined jointly with the runtime alpha/limit -// scalars, so it is handled in the do_pass1 epilogue (act_apply never sees it; -// the mxfp4 kernel additionally fuses its own copy of the same math). -// ACT_SWIGLU_CLAMP (GLM-5.3 "swiglu_limit") is the same clamped form WITHOUT -// the (up + 1) bias: clamp(gate, max=lim) * sigmoid(alpha*gate) * clamp(up, +-lim). -enum ActKind { - ACT_SILU = 0, - ACT_GELU = 1, - ACT_GELU_TANH = 2, - ACT_SWIGLUOAI = 3, - ACT_SWIGLU_CLAMP = 4, -}; - -inline float act_apply(int act, float x) { - if (act == ACT_SILU) return x / (1.0f + std::exp(-x)); - if (act == ACT_GELU) - return 0.5f * x * (1.0f + std::erf(x * 0.70710678118654752440f)); - // gelu_tanh - const float k0 = 0.79788456080286535588f; // sqrt(2/pi) - const float inner = k0 * (x + 0.044715f * x * x * x); - return 0.5f * x * (1.0f + std::tanh(inner)); -} - -// ------------------------------- dot products ------------------------------- -// dot(weight[bf16], act[bf16], n) -> fp32. The selected impl is a function -// pointer chosen at runtime; the per-row call overhead is negligible vs n. - -using dot_fn = float (*)(const bf16_t*, const bf16_t*, int); - -float dot_scalar(const bf16_t* w, const bf16_t* x, int n) { - float acc = 0.0f; - for (int i = 0; i < n; ++i) acc += bf16_to_f32(w[i]) * bf16_to_f32(x[i]); - return acc; -} - -// Software prefetch distance (bytes) ahead of the current weight row stream. The -// weight stream is the bandwidth bottleneck (read once, never reused); nudging the -// HW prefetcher with a few cache lines of lookahead raises sustained throughput. -constexpr int PF_AHEAD = 512; - -#if CPU_MOE_X86 -__attribute__((target("avx512f"))) -float dot_avx512f(const bf16_t* w, const bf16_t* x, int n) { - // 4 independent accumulators -> more in-flight loads (memory-level parallelism), - // which is what lifts a bandwidth-bound GEMV toward peak. - __m512 a0 = _mm512_setzero_ps(), a1 = _mm512_setzero_ps(); - __m512 a2 = _mm512_setzero_ps(), a3 = _mm512_setzero_ps(); - int i = 0; - for (; i + 64 <= n; i += 64) { - _mm_prefetch(reinterpret_cast(w + i) + PF_AHEAD, _MM_HINT_T0); - for (int j = 0; j < 64; j += 16) { - __m256i wi = _mm256_loadu_si256(reinterpret_cast(w + i + j)); - __m256i xi = _mm256_loadu_si256(reinterpret_cast(x + i + j)); - __m512 wf = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(wi), 16)); - __m512 xf = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(xi), 16)); - __m512& acc = (j == 0) ? a0 : (j == 16) ? a1 : (j == 32) ? a2 : a3; - acc = _mm512_fmadd_ps(wf, xf, acc); - } - } - for (; i + 16 <= n; i += 16) { - __m256i wi = _mm256_loadu_si256(reinterpret_cast(w + i)); - __m256i xi = _mm256_loadu_si256(reinterpret_cast(x + i)); - __m512 wf = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(wi), 16)); - __m512 xf = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(xi), 16)); - a0 = _mm512_fmadd_ps(wf, xf, a0); - } - float s = _mm512_reduce_add_ps(_mm512_add_ps(_mm512_add_ps(a0, a1), _mm512_add_ps(a2, a3))); - for (; i < n; ++i) s += bf16_to_f32(w[i]) * bf16_to_f32(x[i]); - return s; -} - -#if (defined(__GNUC__) && __GNUC__ >= 10) || defined(__clang__) -#define CPU_MOE_HAS_AVX512BF16 1 -__attribute__((target("avx512bf16,avx512f"))) -static inline __m512bh load_bh(const bf16_t* p) { - __m512i raw = _mm512_loadu_si512(reinterpret_cast(p)); - __m512bh out; - std::memcpy(&out, &raw, sizeof(out)); - return out; -} - -__attribute__((target("avx512bf16,avx512f"))) -float dot_avx512bf16(const bf16_t* w, const bf16_t* x, int n) { - // 4 accumulators (128 bf16/iter) for memory-level parallelism + a prefetch nudge. - __m512 a0 = _mm512_setzero_ps(), a1 = _mm512_setzero_ps(); - __m512 a2 = _mm512_setzero_ps(), a3 = _mm512_setzero_ps(); - int i = 0; - for (; i + 128 <= n; i += 128) { - _mm_prefetch(reinterpret_cast(w + i) + PF_AHEAD, _MM_HINT_T0); - a0 = _mm512_dpbf16_ps(a0, load_bh(w + i), load_bh(x + i)); - a1 = _mm512_dpbf16_ps(a1, load_bh(w + i + 32), load_bh(x + i + 32)); - a2 = _mm512_dpbf16_ps(a2, load_bh(w + i + 64), load_bh(x + i + 64)); - a3 = _mm512_dpbf16_ps(a3, load_bh(w + i + 96), load_bh(x + i + 96)); - } - for (; i + 32 <= n; i += 32) { - a0 = _mm512_dpbf16_ps(a0, load_bh(w + i), load_bh(x + i)); - } - float s = _mm512_reduce_add_ps(_mm512_add_ps(_mm512_add_ps(a0, a1), _mm512_add_ps(a2, a3))); - for (; i < n; ++i) s += bf16_to_f32(w[i]) * bf16_to_f32(x[i]); - return s; -} -#endif // avx512bf16 available - -// --------- AVX2 + FMA (256-bit fallback for CPUs without AVX-512) ---------- -// Covers Intel 12-14th gen / Arrow Lake (AVX-512 fused off) and AMD Zen<4. The -// bf16->fp32 widen is a zero-extend + <<16; with 4 independent accumulators the -// GEMV is memory-bandwidth bound, same as the AVX-512 path (just half the width). -__attribute__((target("avx2,fma"))) -inline float hsum256(__m256 v) { - __m128 lo = _mm256_castps256_ps128(v); - lo = _mm_add_ps(lo, _mm256_extractf128_ps(v, 1)); - lo = _mm_add_ps(lo, _mm_movehl_ps(lo, lo)); - lo = _mm_add_ss(lo, _mm_shuffle_ps(lo, lo, 0x55)); - return _mm_cvtss_f32(lo); -} - -__attribute__((target("avx2,fma"))) -float dot_avx2(const bf16_t* w, const bf16_t* x, int n) { - __m256 a0 = _mm256_setzero_ps(), a1 = _mm256_setzero_ps(); - __m256 a2 = _mm256_setzero_ps(), a3 = _mm256_setzero_ps(); - int i = 0; - for (; i + 32 <= n; i += 32) { - _mm_prefetch(reinterpret_cast(w + i) + PF_AHEAD, _MM_HINT_T0); - for (int j = 0; j < 32; j += 8) { - __m128i wi = _mm_loadu_si128(reinterpret_cast(w + i + j)); - __m128i xi = _mm_loadu_si128(reinterpret_cast(x + i + j)); - __m256 wf = _mm256_castsi256_ps(_mm256_slli_epi32(_mm256_cvtepu16_epi32(wi), 16)); - __m256 xf = _mm256_castsi256_ps(_mm256_slli_epi32(_mm256_cvtepu16_epi32(xi), 16)); - __m256& acc = (j == 0) ? a0 : (j == 8) ? a1 : (j == 16) ? a2 : a3; - acc = _mm256_fmadd_ps(wf, xf, acc); - } - } - for (; i + 8 <= n; i += 8) { - __m128i wi = _mm_loadu_si128(reinterpret_cast(w + i)); - __m128i xi = _mm_loadu_si128(reinterpret_cast(x + i)); - __m256 wf = _mm256_castsi256_ps(_mm256_slli_epi32(_mm256_cvtepu16_epi32(wi), 16)); - __m256 xf = _mm256_castsi256_ps(_mm256_slli_epi32(_mm256_cvtepu16_epi32(xi), 16)); - a0 = _mm256_fmadd_ps(wf, xf, a0); - } - float s = hsum256(_mm256_add_ps(_mm256_add_ps(a0, a1), _mm256_add_ps(a2, a3))); - for (; i < n; ++i) s += bf16_to_f32(w[i]) * bf16_to_f32(x[i]); - return s; -} -#endif // CPU_MOE_X86 - -// --------------------------- NVFP4 (W4A16) dequant --------------------------- -// Weights: e2m1 4-bit codes (2/byte, low nibble first), per-16 block scale in -// fp8-e4m3, per-output-row global scale in fp16. Dequant matches the GPU kernels -// (freetoken/kernel/triton/nvfp4_dequant.py): w = E2M1[code] * e4m3(scale) * global. -// Activations stay bf16 (W4A16); the GEMV dequantizes weights inside the K-loop. - -const float kE2M1[16] = {0.0f, 0.5f, 1.0f, 1.5f, 2.0f, 3.0f, 4.0f, 6.0f, - -0.0f, -0.5f, -1.0f, -1.5f, -2.0f, -3.0f, -4.0f, -6.0f}; - -// e2m1 * 2 as exact int8 (all e2m1 values are multiples of 0.5). Used by the AVX-VNNI -// W4A8 path: nibble -> int8 weight via PSHUFB LUT, then VPDPBUSD against int8 activations; -// the *2 is undone by a 0.5 folded into the final scale. Mirrors ggml's kvalues_mxfp4. -alignas(16) const int8_t kE2M1x2[16] = {0, 1, 2, 3, 4, 6, 8, 12, 0, -1, -2, -3, -4, -6, -8, -12}; - -inline float fp16_to_f32(uint16_t h) { - const uint32_t sign = (uint32_t)(h & 0x8000u) << 16; - uint32_t exp = (h >> 10) & 0x1Fu; - uint32_t man = h & 0x3FFu; - uint32_t f; - if (exp == 0) { - if (man == 0) { - f = sign; - } else { - exp = 127 - 15 + 1; - while ((man & 0x400u) == 0) { - man <<= 1; - --exp; - } - man &= 0x3FFu; - f = sign | (exp << 23) | (man << 13); - } - } else if (exp == 0x1Fu) { - f = sign | 0x7F800000u | (man << 13); - } else { - f = sign | ((exp + (127 - 15)) << 23) | (man << 13); - } - float out; - std::memcpy(&out, &f, sizeof(out)); - return out; -} - -// e4m3 (OCP "fn": finite, max-normal 448, exp bias 7). Decoded into a 256-entry LUT. -inline float e4m3_decode(uint8_t v) { - const float sign = (v & 0x80u) ? -1.0f : 1.0f; - const uint32_t exp = (v >> 3) & 0xFu; - const uint32_t man = v & 0x7u; - if (exp == 0) return sign * (man / 8.0f) * 0.015625f; // 2^(1-7) = 2^-6 - return sign * (1.0f + man / 8.0f) * std::ldexp(1.0f, (int)exp - 7); -} - -// Activations pre-deinterleaved to fp32 (xe[m]=x[2m], xo[m]=x[2m+1]); see the -// ds_fp4 dot below for why (drops the hot loop to ~1.5 shuffle ops / 16 weights). -using nvdot_fn = float (*)(const uint8_t*, const uint8_t*, float, const float*, const float*, - int, const float*, const float*); - -float dot_nvfp4_scalar(const uint8_t* packed, const uint8_t* scale, float global, - const float* xe, const float* xo, int K, const float* e2m1, - const float* e4m3) { - float acc = 0.0f; - const int nb = K / 16; - for (int b = 0; b < nb; ++b) { - const float bs = e4m3[scale[b]]; - const uint8_t* pk = packed + (size_t)b * 8; - const float* xeb = xe + (size_t)b * 8; // 16 K -> 8 even + 8 odd - const float* xob = xo + (size_t)b * 8; - float bsum = 0.0f; - for (int j = 0; j < 8; ++j) { - const uint8_t byte = pk[j]; - bsum += e2m1[byte & 0xF] * xeb[j]; - bsum += e2m1[byte >> 4] * xob[j]; - } - acc += bs * bsum; - } - return acc * global; -} - -// ---- NVFP4 W4A8 (int8 activations) dot: nibble->int8 via LUT, per-16 act scale ---- -// asi8: int8 activations laid out per-16 block as [even(8), odd(8)]; asb[b] = per-block -// activation scale (absmax/127). Result folds the e2m1*2 -> *0.5 into the scale. -using nvi8dot_fn = float (*)(const uint8_t*, const uint8_t*, float, const int8_t*, int, - const float*, const float*); - -[[maybe_unused]] float dot_nvfp4_i8_scalar(const uint8_t* packed, const uint8_t* scale, - float global, const int8_t* asi8, int K, const float* e4m3, - const float* asb) { - float acc = 0.0f; - const int nb = K / 16; - for (int b = 0; b < nb; ++b) { - const uint8_t* pk = packed + (size_t)b * 8; - const int8_t* ae = asi8 + (size_t)b * 16; // even(8) - const int8_t* ao = ae + 8; // odd(8) - int isum = 0; - for (int j = 0; j < 8; ++j) { - isum += (int)kE2M1x2[pk[j] & 0xF] * (int)ae[j]; - isum += (int)kE2M1x2[pk[j] >> 4] * (int)ao[j]; - } - acc += (e4m3[scale[b]] * asb[b]) * (float)isum; - } - return acc * (0.5f * global); -} - -#if CPU_MOE_X86 -// AVX2 e2m1 nibble decode: codes (0..15) in 8 int32 lanes -> fp32. AVX2 vpermps is -// only 8-wide, so instead of a 16-entry LUT we use the e2m1 sign/magnitude symmetry: -// value = (code&8 ? - : +) * mag8[code&7], mag8 = e2m1[0..7]. The sign is bit 3 of -// the code shifted into the fp32 sign bit (bit 31). Bit-identical to the e2m1 LUT. -__attribute__((target("avx2,fma"))) -inline __m256 e2m1_decode8(__m256i codes, __m256 mag8) { - __m256 mag = _mm256_permutevar8x32_ps(mag8, _mm256_and_si256(codes, _mm256_set1_epi32(7))); - __m256i sgn = _mm256_slli_epi32(_mm256_and_si256(codes, _mm256_set1_epi32(8)), 28); - return _mm256_xor_ps(mag, _mm256_castsi256_ps(sgn)); -} - -// Two 16-K blocks (16 packed bytes) per iter: lo nibbles -> even-K, hi -> odd-K, -// gathered via two vpermps. The per-16 e4m3 scale differs across the two blocks, so -// it is applied per lane (low 8 lanes = block b, high 8 = block b+1). -__attribute__((target("avx512f"))) -inline __m512 nvfp4_blk2(const uint8_t* pk, const float* xeb, const float* xob, __m512 lut, - __m512i loma, float s0, float s1) { - __m512i wi = _mm512_cvtepu8_epi32(_mm_loadu_si128(reinterpret_cast(pk))); - __m512 vlo = _mm512_permutexvar_ps(_mm512_and_si512(wi, loma), lut); - __m512 vhi = _mm512_permutexvar_ps(_mm512_and_si512(_mm512_srli_epi32(wi, 4), loma), lut); - __m512 prod = _mm512_fmadd_ps(vlo, _mm512_loadu_ps(xeb), _mm512_mul_ps(vhi, _mm512_loadu_ps(xob))); - // lanes 0-7 (block b) -> s0, lanes 8-15 (block b+1) -> s1 (pure AVX512F mask move) - __m512 scv = _mm512_mask_mov_ps(_mm512_set1_ps(s0), 0xFF00, _mm512_set1_ps(s1)); - return _mm512_mul_ps(prod, scv); -} - -__attribute__((target("avx512f"))) -float dot_nvfp4_avx512(const uint8_t* packed, const uint8_t* scale, float global, - const float* xe, const float* xo, int K, const float* e2m1, - const float* e4m3) { - const __m512 lut = _mm512_loadu_ps(e2m1); // 16 e2m1 values - const __m512i loma = _mm512_set1_epi32(0xF); - __m512 acc0 = _mm512_setzero_ps(), acc1 = _mm512_setzero_ps(); - const int nb = K / 16; // 8 packed bytes + one e4m3 scale per block - int b = 0; - for (; b + 4 <= nb; b += 4) { // two blk2 calls -> 4 blocks, 2 accumulators - acc0 = _mm512_add_ps(acc0, nvfp4_blk2(packed + (size_t)b * 8, xe + (size_t)b * 8, - xo + (size_t)b * 8, lut, loma, e4m3[scale[b]], e4m3[scale[b + 1]])); - acc1 = _mm512_add_ps(acc1, nvfp4_blk2(packed + (size_t)(b + 2) * 8, xe + (size_t)(b + 2) * 8, - xo + (size_t)(b + 2) * 8, lut, loma, e4m3[scale[b + 2]], e4m3[scale[b + 3]])); - } - for (; b + 2 <= nb; b += 2) - acc0 = _mm512_add_ps(acc0, nvfp4_blk2(packed + (size_t)b * 8, xe + (size_t)b * 8, - xo + (size_t)b * 8, lut, loma, e4m3[scale[b]], e4m3[scale[b + 1]])); - float s = _mm512_reduce_add_ps(_mm512_add_ps(acc0, acc1)); - for (; b < nb; ++b) { // odd final 16-K block - const uint8_t* pk = packed + (size_t)b * 8; - const float* xeb = xe + (size_t)b * 8; - const float* xob = xo + (size_t)b * 8; - float bsum = 0.0f; - for (int j = 0; j < 8; ++j) { - bsum += e2m1[pk[j] & 0xF] * xeb[j] + e2m1[pk[j] >> 4] * xob[j]; - } - s += e4m3[scale[b]] * bsum; - } - return s * global; -} - -// AVX2: one 16-K block (8 packed bytes) per call, 8 even + 8 odd lanes. -__attribute__((target("avx2,fma"))) -inline __m256 nvfp4_blk_avx2(const uint8_t* pk, const float* xeb, const float* xob, - __m256 mag8, float sc) { - __m256i wi = _mm256_cvtepu8_epi32(_mm_loadl_epi64(reinterpret_cast(pk))); - __m256 vlo = e2m1_decode8(_mm256_and_si256(wi, _mm256_set1_epi32(0xF)), mag8); - __m256 vhi = e2m1_decode8(_mm256_srli_epi32(wi, 4), mag8); - __m256 prod = _mm256_fmadd_ps(vlo, _mm256_loadu_ps(xeb), _mm256_mul_ps(vhi, _mm256_loadu_ps(xob))); - return _mm256_mul_ps(prod, _mm256_set1_ps(sc)); -} - -__attribute__((target("avx2,fma"))) -float dot_nvfp4_avx2(const uint8_t* packed, const uint8_t* scale, float global, - const float* xe, const float* xo, int K, const float* e2m1, - const float* e4m3) { - const __m256 mag8 = _mm256_loadu_ps(e2m1); // e2m1[0..7] magnitudes - __m256 acc0 = _mm256_setzero_ps(), acc1 = _mm256_setzero_ps(); - const int nb = K / 16; - int b = 0; - for (; b + 2 <= nb; b += 2) { - acc0 = _mm256_add_ps(acc0, nvfp4_blk_avx2(packed + (size_t)b * 8, xe + (size_t)b * 8, - xo + (size_t)b * 8, mag8, e4m3[scale[b]])); - acc1 = _mm256_add_ps(acc1, nvfp4_blk_avx2(packed + (size_t)(b + 1) * 8, xe + (size_t)(b + 1) * 8, - xo + (size_t)(b + 1) * 8, mag8, e4m3[scale[b + 1]])); - } - for (; b < nb; ++b) - acc0 = _mm256_add_ps(acc0, nvfp4_blk_avx2(packed + (size_t)b * 8, xe + (size_t)b * 8, - xo + (size_t)b * 8, mag8, e4m3[scale[b]])); - return hsum256(_mm256_add_ps(acc0, acc1)) * global; -} - -// AVX-VNNI W4A8: decode 8 packed bytes (16 nibbles) of one 16-block to int8 [lo(8),hi(8)] -// via PSHUFB against the e2m1*2 LUT (replaces the 2 vpermps fp32 expands -- ~4x less -// port-5 traffic). lo=even-K weights, hi=odd-K, matching the [even(8),odd(8)] act layout. -__attribute__((target("avx2,avxvnni,fma"))) -inline __m128i nvfp4_decode_block_i8(const uint8_t* pk, __m128i lut) { - __m128i b = _mm_loadl_epi64(reinterpret_cast(pk)); // 8 bytes - __m128i lo = _mm_and_si128(b, _mm_set1_epi8(0x0F)); - __m128i hi = _mm_and_si128(_mm_srli_epi16(b, 4), _mm_set1_epi8(0x0F)); - return _mm_shuffle_epi8(lut, _mm_unpacklo_epi64(lo, hi)); // [lo(8),hi(8)] -> int8 -} - -// Two 16-blocks per VPDPBUSD (32 int8). Sign trick (ggml mul_add_epi8): |w|*(sign(w)*a)=w*a, -// so u8*s8 needs no offset/correction term. Per-block scale (e4m3 * act-scale) folded in fp32 -// (lanes 0-3 -> block b, 4-7 -> block b+1). Bit-faithful weight; only the int8 activation quant -// (W4A8) differs from the bf16 reference. -__attribute__((target("avx2,avxvnni,fma"))) -float dot_nvfp4_i8_vnni(const uint8_t* packed, const uint8_t* scale, float global, - const int8_t* asi8, int K, const float* e4m3, const float* asb) { - const __m128i lut = _mm_loadu_si128(reinterpret_cast(kE2M1x2)); - __m256 accF = _mm256_setzero_ps(); - const int nb = K / 16; - int b = 0; - for (; b + 2 <= nb; b += 2) { - __m128i wb = nvfp4_decode_block_i8(packed + (size_t)b * 8, lut); - __m128i wb1 = nvfp4_decode_block_i8(packed + (size_t)(b + 1) * 8, lut); - __m256i w = _mm256_set_m128i(wb1, wb); // [blk b | blk b+1] - __m256i a = _mm256_loadu_si256(reinterpret_cast(asi8 + (size_t)b * 16)); - __m256i aw = _mm256_sign_epi8(w, w); // |w| (u8 operand) - __m256i sa = _mm256_sign_epi8(a, w); // sign(w)*a (s8 operand) - __m256i di = _mm256_dpbusd_avx_epi32(_mm256_setzero_si256(), aw, sa); - __m256 scv = _mm256_blend_ps(_mm256_set1_ps(e4m3[scale[b]] * asb[b]), - _mm256_set1_ps(e4m3[scale[b + 1]] * asb[b + 1]), 0xF0); - accF = _mm256_fmadd_ps(_mm256_cvtepi32_ps(di), scv, accF); - } - float s = hsum256(accF); - for (; b < nb; ++b) { // tail (odd block count) - const uint8_t* pk = packed + (size_t)b * 8; - const int8_t* ae = asi8 + (size_t)b * 16; const int8_t* ao = ae + 8; - int isum = 0; - for (int j = 0; j < 8; ++j) - isum += (int)kE2M1x2[pk[j] & 0xF] * (int)ae[j] + (int)kE2M1x2[pk[j] >> 4] * (int)ao[j]; - s += (e4m3[scale[b]] * asb[b]) * (float)isum; - } - return s * (0.5f * global); -} - -#if (defined(__GNUC__) && __GNUC__ >= 10) || defined(__clang__) -#define CPU_MOE_HAS_AVX512VNNI 1 - -// Software-prefetch distance for the W4A8 weight stream, in 16-K blocks (8 packed -// bytes each). Returns -1 when FREETOKEN_CPU_MOE_PF_BLOCKS is unset: the kernel then -// uses the built-in default min(512 blocks = 4 KB, 2 rows) -- 4 KB is the empirical -// optimum on large-row machines (Emerald Rapids sweep), while the 2-row cap keeps a small-row -// model's overshoot bounded (the executor works in 32-row tiles, so a fixed byte -// distance otherwise prefetches another worker's tile: duplicated DRAM traffic that -// regresses at the bandwidth ceiling). An EXPLICIT env value is honored verbatim -// (no clamp; 0 disables): the per-machine optimum can sit past the safe default (+20% -// at 4 KB on a 24-thread Ice Lake with 256B rows), so the escape hatch must reach it. -// Prefetch never faults, so overshooting a row/bank tail is safe. -static int nvfp4_pf_blocks() { - static const int v = [] { - const char* s = getenv("FREETOKEN_CPU_MOE_PF_BLOCKS"); - return (s && s[0]) ? atoi(s) : -1; - }(); - return v; -} - -// AVX-512 VNNI W4A8: FOUR 16-K blocks per VPDPBUSD (64 int8) -- 2x the AVX-VNNI -// (256-bit) path. Decode 32 packed bytes -> 64 int8 with a single _mm512_shuffle_epi8 -// (e2m1*2 LUT replicated to all 4 128-bit lanes). AVX-512 has no _mm512_sign_epi8, so -// the u8*s8 sign trick (|w| as u8, sign(w)*a as s8) uses abs_epi8 + a masked negate. -// Bit-faithful weight; only the int8 activation quant (W4A8) differs from bf16. -// One 4-block group (64 int8) -> scaled fp32 partial (16 lanes). Isolated as a helper so -// the caller can run several independent chains into separate accumulators (the decode -> -// dpbusd -> scale chain is long, so a single accumulator leaves the core latency-bound). -__attribute__((target("avx512f,avx512bw,avx512vnni,avx2"))) -static inline __m512 nvfp4_i8_grp4(const uint8_t* packed, const uint8_t* scale, - const int8_t* asi8, const float* e4m3, const float* asb, - int b, __m512i lut, __m512i idx, __m512i mask0F, - __m512i idxsc) { - const __mmask64 hi_half = 0xFF00FF00FF00FF00ULL; // bytes 8-15 of each 128b lane -> hi nibbles - __m256i raw = _mm256_loadu_si256(reinterpret_cast(packed + (size_t)b * 8)); - __m512i src = _mm512_permutexvar_epi64(idx, _mm512_castsi256_si512(raw)); - __m512i lo = _mm512_and_si512(src, mask0F); - __m512i hi = _mm512_and_si512(_mm512_srli_epi16(src, 4), mask0F); - __m512i comb = _mm512_mask_blend_epi8(hi_half, lo, hi); - __m512i w = _mm512_shuffle_epi8(lut, comb); // 64 int8 weights (e2m1*2) - // u8*s8 sign trick without _mm512_sign_epi8: aw=|w|, sa = (w<0 ? -a : a) (a moot at w==0) - __m512i a = _mm512_loadu_si512(reinterpret_cast(asi8 + (size_t)b * 16)); - __m512i aw = _mm512_abs_epi8(w); - __mmask64 neg = _mm512_movepi8_mask(w); - __m512i sa = _mm512_mask_sub_epi8(a, neg, _mm512_setzero_si512(), a); - __m512i di = _mm512_dpbusd_epi32(_mm512_setzero_si512(), aw, sa); // 16 int32 (groups of 4) - // int32[0..3]->blk b, [4..7]->b+1, [8..11]->b+2, [12..15]->b+3. The 4 block scales - // (e4m3 LUT x per-block act scale) are computed vectorized: a 4-byte load + epu8->epi32 - // widen + one 4-lane LUT gather + one mul replaces 8 scalar loads + 4 scalar muls + - // a set_ps assembly, which otherwise dominates the per-group op count. - int sc_raw; - memcpy(&sc_raw, scale + b, 4); - __m128i sc4 = _mm_cvtepu8_epi32(_mm_cvtsi32_si128(sc_raw)); - __m128 s4 = _mm_mul_ps(_mm_i32gather_ps(e4m3, sc4, 4), _mm_loadu_ps(asb + b)); - __m512 scv = _mm512_permutexvar_ps(idxsc, _mm512_castps128_ps512(s4)); - return _mm512_mul_ps(_mm512_cvtepi32_ps(di), scv); -} - -__attribute__((target("avx512f,avx512bw,avx512vnni,avx2"))) -float dot_nvfp4_i8_avx512vnni(const uint8_t* packed, const uint8_t* scale, float global, - const int8_t* asi8, int K, const float* e4m3, const float* asb) { - const __m512i lut = _mm512_broadcast_i32x4( - _mm_loadu_si128(reinterpret_cast(kE2M1x2))); - const __m512i idx = _mm512_set_epi64(3, 3, 2, 2, 1, 1, 0, 0); // block i -> 128b lane i - const __m512i mask0F = _mm512_set1_epi8(0x0F); - const __m512i idxsc = _mm512_set_epi32(3, 3, 3, 3, 2, 2, 2, 2, 1, 1, 1, 1, 0, 0, 0, 0); - __m512 acc0 = _mm512_setzero_ps(), acc1 = _mm512_setzero_ps(); - __m512 acc2 = _mm512_setzero_ps(), acc3 = _mm512_setzero_ps(); - const int nb = K / 16; - // Built-in default: min(4 KB, 2 rows) -- see nvfp4_pf_blocks(). An explicitly set - // env value is used verbatim so a per-machine sweep can reach operating points past - // the conservative default. - const int pfb = nvfp4_pf_blocks(); - const int pf = (pfb < 0) ? std::min(512, 2 * nb) : pfb; - int b = 0; - // Four independent 4-block groups per iter -> four accumulator chains hide the ~10-op - // decode->dpbusd->scale latency (1-2 chains leave the loop latency-bound: the per-core - // rate sat at ~60% of the core's achievable DRAM stream rate). - for (; b + 16 <= nb; b += 16) { - // Prefetch the weight stream ahead: the interleaved decode lowers the L1-miss - // concurrency the HW prefetcher sustains on its own (134 -> 167 GB/s on Emerald Rapids). - if (pf > 0) { - _mm_prefetch(reinterpret_cast(packed + ((size_t)b + (size_t)pf) * 8), - _MM_HINT_T0); - _mm_prefetch(reinterpret_cast(packed + ((size_t)b + (size_t)pf) * 8 + 64), - _MM_HINT_T0); - } - acc0 = _mm512_add_ps(acc0, nvfp4_i8_grp4(packed, scale, asi8, e4m3, asb, b, lut, idx, - mask0F, idxsc)); - acc1 = _mm512_add_ps(acc1, nvfp4_i8_grp4(packed, scale, asi8, e4m3, asb, b + 4, lut, idx, - mask0F, idxsc)); - acc2 = _mm512_add_ps(acc2, nvfp4_i8_grp4(packed, scale, asi8, e4m3, asb, b + 8, lut, idx, - mask0F, idxsc)); - acc3 = _mm512_add_ps(acc3, nvfp4_i8_grp4(packed, scale, asi8, e4m3, asb, b + 12, lut, idx, - mask0F, idxsc)); - } - for (; b + 4 <= nb; b += 4) - acc0 = _mm512_add_ps(acc0, nvfp4_i8_grp4(packed, scale, asi8, e4m3, asb, b, lut, idx, - mask0F, idxsc)); - float s = _mm512_reduce_add_ps( - _mm512_add_ps(_mm512_add_ps(acc0, acc1), _mm512_add_ps(acc2, acc3))); - for (; b < nb; ++b) { // tail (<4 remaining 16-K blocks) - const uint8_t* pk = packed + (size_t)b * 8; - const int8_t* ae = asi8 + (size_t)b * 16; const int8_t* ao = ae + 8; - int isum = 0; - for (int j = 0; j < 8; ++j) - isum += (int)kE2M1x2[pk[j] & 0xF] * (int)ae[j] + (int)kE2M1x2[pk[j] >> 4] * (int)ao[j]; - s += (e4m3[scale[b]] * asb[b]) * (float)isum; - } - return s * (0.5f * global); -} -#endif // avx512vnni available -#endif - -// ===================================================================================== -// CUDA stream memory operations (driver API, resolved via dlopen -- no link-time or -// toolchain dependence). The GPU side of the flag handshake: submit = WRITE_VALUE -// (done[slot]=0 then ready[slot]=1), sync = WAIT_VALUE(done[slot] >= 1). The wait is -// executed by the GPU front-end (no SM-resident kernel), so GPU "utilization" stays -// truthful during CPU compute windows -- a resident spin kernel pinned it at 99%, -// which laptop CPU/GPU dynamic power schedulers answered by clamping the CPU's max -// frequency (GEMV workers -1.5x: the reported edge regression). Availability is -// probed functionally at startup (memops_probe); anything unsupported (Windows WDDM, -// vGPU, old drivers) falls back to the cudaLaunchHostFunc path. -#if defined(_WIN32) -#include -static void* cumemop_dlopen() { return (void*)::LoadLibraryA("nvcuda.dll"); } -static void* cumemop_dlsym(void* h, const char* n) { - return (void*)::GetProcAddress((HMODULE)h, n); -} -#else -#include -static void* cumemop_dlopen() { - void* h = dlopen("libcuda.so.1", RTLD_LAZY | RTLD_LOCAL); - if (h == nullptr) h = dlopen("libcuda.so", RTLD_LAZY | RTLD_LOCAL); - return h; -} -static void* cumemop_dlsym(void* h, const char* n) { return dlsym(h, n); } -#endif - -using cuMemOp64_fn = int (*)(void* stream, unsigned long long addr, unsigned long long value, - unsigned int flags); -static cuMemOp64_fn g_cu_write64 = nullptr; -static cuMemOp64_fn g_cu_wait64 = nullptr; -static constexpr unsigned int kCuWaitValueGeq = 0x0; // CU_STREAM_WAIT_VALUE_GEQ -static constexpr unsigned int kCuWriteDefault = 0x0; // CU_STREAM_WRITE_VALUE_DEFAULT - -static bool cumemop_resolve() { - static bool resolved = [] { - void* h = cumemop_dlopen(); - if (h == nullptr) return false; - // 11.7+ made the v2 entry points the default; older drivers export only the v1 - // names with the same signature. - g_cu_write64 = reinterpret_cast(cumemop_dlsym(h, "cuStreamWriteValue64_v2")); - if (g_cu_write64 == nullptr) - g_cu_write64 = reinterpret_cast(cumemop_dlsym(h, "cuStreamWriteValue64")); - g_cu_wait64 = reinterpret_cast(cumemop_dlsym(h, "cuStreamWaitValue64_v2")); - if (g_cu_wait64 == nullptr) - g_cu_wait64 = reinterpret_cast(cumemop_dlsym(h, "cuStreamWaitValue64")); - return g_cu_write64 != nullptr && g_cu_wait64 != nullptr; - }(); - return resolved; -} - -// Functional probe on a scratch pinned int64: enqueue WRITE(7) + WAIT(>=7) + sync. -// Returns true only if the whole memop path works on THIS stream/device/driver. -static bool cumemops_probe(uintptr_t stream, uintptr_t scratch_addr) { - if (!cumemop_resolve()) return false; - auto* s = reinterpret_cast(stream); - if (g_cu_write64(s, (unsigned long long)scratch_addr, 7ULL, kCuWriteDefault) != 0) return false; - if (g_cu_wait64(s, (unsigned long long)scratch_addr, 7ULL, kCuWaitValueGeq) != 0) return false; - return cudaStreamSynchronize(reinterpret_cast(stream)) == cudaSuccess; -} - -// GPU side of the flag handshake (see the block comment above): enqueued on the -// caller's (possibly capturing) stream; the WAIT immediate is the constant 1, -// replay-safe under CUDA graphs. -// The startup probe validates EAGER memops; a driver could still reject them at graph -// capture time. Those enqueue errors would otherwise be swallowed here and surface only -// as a later EndCapture failure -- log the first CUresult so triage is one step. -static void cumemop_check(int rc, const char* what) { - static std::atomic warned{false}; - if (rc != 0 && !warned.exchange(true)) { - std::fprintf(stderr, - "[freetoken/cpu_moe] %s failed with CUresult=%d (first occurrence; " - "subsequent errors are not repeated). If this happened during CUDA " - "graph capture, the driver lacks capture support for stream memops -- " - "set FREETOKEN_CPU_MOE_FLAG_SYNC=0.\n", - what, rc); - } -} - -static void cumemop_submit(uintptr_t stream, uintptr_t done_addr, uintptr_t ready_addr, - int64_t slot) { - auto* s = reinterpret_cast(stream); - // Order matters and is preserved by the front end: reset done BEFORE raising ready, - // so the coordinator's completion write for THIS step can never be wiped. - cumemop_check(g_cu_write64(s, (unsigned long long)(done_addr + (size_t)slot * 8), 0ULL, - kCuWriteDefault), - "cuStreamWriteValue64(done)"); - cumemop_check(g_cu_write64(s, (unsigned long long)(ready_addr + (size_t)slot * 8), 1ULL, - kCuWriteDefault), - "cuStreamWriteValue64(ready)"); -} - -static void cumemop_sync(uintptr_t stream, uintptr_t done_addr, int64_t slot) { - cumemop_check(g_cu_wait64(reinterpret_cast(stream), - (unsigned long long)(done_addr + (size_t)slot * 8), 1ULL, - kCuWaitValueGeq), - "cuStreamWaitValue64(done)"); -} - -struct DotChoice { - dot_fn fn; - const char* name; -}; +#include "format.h" +#include "gpu_sync.h" +#include "gpu_sync_io.h" -// SIMD tiers, ascending. Each format picks the highest tier <= the one chosen by -// pick_isa() that it implements (fp4 formats have no bf16-specific tier, so the -// avx512bf16 tier maps to their avx512 kernel). -enum IsaTier { ISA_SCALAR = 0, ISA_AVX2 = 1, ISA_AVX512 = 2, ISA_AVX512BF16 = 3 }; +namespace freetoken::cpu_moe { -// Best tier the CPU+build supports, optionally capped DOWN by -// FREETOKEN_CPU_MOE_ISA={scalar,avx2,avx512,avx512bf16} (A/B testing on a machine -// that supports more). FREETOKEN_CPU_MOE_SCALAR=1 forces scalar (legacy alias). -inline IsaTier pick_isa() { -#if CPU_MOE_X86 - if (getenv("FREETOKEN_CPU_MOE_SCALAR")) return ISA_SCALAR; - IsaTier best = ISA_SCALAR; - if (__builtin_cpu_supports("avx2") && __builtin_cpu_supports("fma")) best = ISA_AVX2; - if (best >= ISA_AVX2 && __builtin_cpu_supports("avx512f")) best = ISA_AVX512; -#ifdef CPU_MOE_HAS_AVX512BF16 - if (best >= ISA_AVX512 && __builtin_cpu_supports("avx512bf16")) best = ISA_AVX512BF16; -#endif - if (const char* f = getenv("FREETOKEN_CPU_MOE_ISA")) { - IsaTier want = best; - if (!std::strcmp(f, "scalar")) want = ISA_SCALAR; - else if (!std::strcmp(f, "avx2")) want = ISA_AVX2; - else if (!std::strcmp(f, "avx512")) want = ISA_AVX512; - else if (!std::strcmp(f, "avx512bf16")) want = ISA_AVX512BF16; - if (want < best) best = want; // cap downward; never force above hw/build support - } - return best; -#else - return ISA_SCALAR; -#endif -} - -DotChoice select_dot() { - const IsaTier t = pick_isa(); -#if CPU_MOE_X86 -#ifdef CPU_MOE_HAS_AVX512BF16 - if (t >= ISA_AVX512BF16) return {dot_avx512bf16, "avx512bf16"}; -#endif - if (t >= ISA_AVX512) return {dot_avx512f, "avx512f"}; - if (t >= ISA_AVX2) return {dot_avx2, "avx2"}; -#endif - (void)t; - return {dot_scalar, "scalar"}; -} - -nvdot_fn select_nvdot() { - const IsaTier t = pick_isa(); -#if CPU_MOE_X86 - if (t >= ISA_AVX512) return dot_nvfp4_avx512; - if (t >= ISA_AVX2) return dot_nvfp4_avx2; -#endif - (void)t; - return dot_nvfp4_scalar; -} +std::unique_ptr make_bf16_format(const MoeParams& p, const FormatArgs& a); +std::unique_ptr make_nvfp4_format(const MoeParams& p, const FormatArgs& a); +std::unique_ptr make_mxfp4_format(const MoeParams& p, const FormatArgs& a); +std::unique_ptr make_ds_fp4_format(const MoeParams& p, const FormatArgs& a); +std::unique_ptr make_fp8_block_format(const MoeParams& p, const FormatArgs& a); +std::unique_ptr make_gguf_format(const MoeParams& p, const FormatArgs& a); +std::unique_ptr make_gguf_scaled_format(const MoeParams& p, const FormatArgs& a); +void set_ggml_cpu(uintptr_t addr); -// AVX-VNNI (VEX-256 VPDPBUSD) availability: Alder/Raptor Lake, Sapphire Rapids+, Zen5. -// Distinct from AVX-512 VNNI. Opt out with FREETOKEN_CPU_MOE_NO_VNNI=1 (A/B the W4A8 path). -inline bool cpu_has_avxvnni() { -#if CPU_MOE_X86 - const char* no = getenv("FREETOKEN_CPU_MOE_NO_VNNI"); - if (no && no[0] && no[0] != '0') return false; // ignore unset/empty/"0" - return __builtin_cpu_supports("avx2") && __builtin_cpu_supports("avxvnni"); -#else - return false; -#endif -} +namespace { -// AVX-512 VNNI (512-bit VPDPBUSD): Cascade Lake+, Ice Lake, Sapphire/Emerald Rapids, -// Zen4+. 2x the 256-bit AVX-VNNI width. FREETOKEN_CPU_MOE_NO_AVX512VNNI=1 forces the -// 256-bit path (A/B the two W4A8 kernels on the same box); FREETOKEN_CPU_MOE_NO_VNNI=1 -// still disables the whole W4A8 family (back to the faithful fp32 nvdot), so it is -// honored here too. -inline bool cpu_has_avx512vnni() { -#if CPU_MOE_X86 && defined(CPU_MOE_HAS_AVX512VNNI) - const char* no = getenv("FREETOKEN_CPU_MOE_NO_AVX512VNNI"); - if (no && no[0] && no[0] != '0') return false; - const char* no_vnni = getenv("FREETOKEN_CPU_MOE_NO_VNNI"); - if (no_vnni && no_vnni[0] && no_vnni[0] != '0') return false; - return __builtin_cpu_supports("avx512vnni"); +inline void pin_this_thread(int cpu) { +#if CPU_MOE_HAS_AFFINITY + cpu_set_t set; + CPU_ZERO(&set); + CPU_SET(cpu, &set); + pthread_setaffinity_np(pthread_self(), sizeof(set), &set); #else - return false; -#endif -} - -// Best W4A8 (int8-activation) nvfp4 dot, or nullptr if no SIMD VNNI (caller keeps the -// faithful fp32 nvdot path). The scalar i8 dot exists only as a correctness reference. -nvi8dot_fn select_nvi8dot() { -#if CPU_MOE_X86 -#if defined(CPU_MOE_HAS_AVX512VNNI) - if (cpu_has_avx512vnni()) return dot_nvfp4_i8_avx512vnni; -#endif - if (cpu_has_avxvnni()) return dot_nvfp4_i8_vnni; + (void)cpu; #endif - return nullptr; -} - -// ----------------------- DeepSeek-V4 ds_fp4 (W4A8) --------------------------- -// Row-major e2m1 (2/byte, low nibble first) + e8m0 per-32 block scale, no global -// (w = E2M1[code] * 2^(e8m0-127)); activations are FP8-e4m3 round-tripped (per-128 -// block, ue8m0 scale) before each GEMM. Matches kernel/triton/dsv4 (fused_moe + -// fp8_linear): silu(clamp(gate,max=lim)) * clamp(up,-lim,lim), router weight on the -// down output. - -// Activations are pre-deinterleaved to fp32 (xe[m]=x[2m], xo[m]=x[2m+1]) once per -// token/route and reused across all output rows. This drops the hot dot to one -// vpmovzxbd + two vpermps per 32 weights (1.5 shuffle ops / 16 vs 4 for a bf16, -// dup-permute, per-element gather), so the row-major fp4 GEMV stops being port-5 -// bound and approaches the bf16 memory-bandwidth ceiling. -using dsdot_fn = float (*)(const uint8_t*, const uint8_t*, const float*, const float*, int, - const float*, const float*); - -float dot_dsfp4_scalar(const uint8_t* packed, const uint8_t* scale, const float* xe, - const float* xo, int K, const float* e2m1, const float* e8m0) { - float acc = 0.0f; - const int nb = K / 32; - for (int b = 0; b < nb; ++b) { - const float sc = e8m0[scale[b]]; - const uint8_t* pk = packed + (size_t)b * 16; // 16 bytes = 32 codes - const float* xeb = xe + (size_t)b * 16; - const float* xob = xo + (size_t)b * 16; - float bsum = 0.0f; - for (int j = 0; j < 16; ++j) { - const uint8_t byte = pk[j]; - bsum += e2m1[byte & 0xF] * xeb[j]; // low nibble -> even-K activation - bsum += e2m1[byte >> 4] * xob[j]; // high nibble -> odd-K activation - } - acc += sc * bsum; - } - return acc; -} - -#if CPU_MOE_X86 -// One 32-block: 16 bytes -> 16 low + 16 high nibble values via two vpermps, times -// the pre-split even/odd fp32 activations, folded by the per-32 e8m0 scale. -__attribute__((target("avx512f"))) -inline __m512 dsfp4_blk(const uint8_t* pk, const float* xeb, const float* xob, __m512 lut, - __m512i loma, float sc) { - __m512i wi = _mm512_cvtepu8_epi32(_mm_loadu_si128(reinterpret_cast(pk))); - __m512 vlo = _mm512_permutexvar_ps(_mm512_and_si512(wi, loma), lut); - __m512 vhi = _mm512_permutexvar_ps(_mm512_and_si512(_mm512_srli_epi32(wi, 4), loma), lut); - __m512 prod = _mm512_fmadd_ps(vlo, _mm512_loadu_ps(xeb), _mm512_mul_ps(vhi, _mm512_loadu_ps(xob))); - return _mm512_mul_ps(prod, _mm512_set1_ps(sc)); -} - -__attribute__((target("avx512f"))) -float dot_dsfp4_avx512(const uint8_t* packed, const uint8_t* scale, const float* xe, - const float* xo, int K, const float* e2m1, const float* e8m0) { - const __m512 lut = _mm512_loadu_ps(e2m1); - const __m512i loma = _mm512_set1_epi32(0xF); - __m512 acc0 = _mm512_setzero_ps(), acc1 = _mm512_setzero_ps(); - const int nb = K / 32; // 16 packed bytes + one e8m0 scale per block - int b = 0; - for (; b + 2 <= nb; b += 2) { // two independent accumulators hide FMA latency - acc0 = _mm512_add_ps(acc0, dsfp4_blk(packed + (size_t)b * 16, xe + (size_t)b * 16, - xo + (size_t)b * 16, lut, loma, e8m0[scale[b]])); - acc1 = _mm512_add_ps(acc1, dsfp4_blk(packed + (size_t)(b + 1) * 16, xe + (size_t)(b + 1) * 16, - xo + (size_t)(b + 1) * 16, lut, loma, e8m0[scale[b + 1]])); - } - for (; b < nb; ++b) - acc0 = _mm512_add_ps(acc0, dsfp4_blk(packed + (size_t)b * 16, xe + (size_t)b * 16, - xo + (size_t)b * 16, lut, loma, e8m0[scale[b]])); - return _mm512_reduce_add_ps(_mm512_add_ps(acc0, acc1)); -} - -// AVX2: a 32-K block is 16 bytes -> two 8-lane halves (8 even + 8 odd each). -__attribute__((target("avx2,fma"))) -inline __m256 dsfp4_half_avx2(const uint8_t* pk, const float* xeb, const float* xob, __m256 mag8) { - __m256i wi = _mm256_cvtepu8_epi32(_mm_loadl_epi64(reinterpret_cast(pk))); - __m256 vlo = e2m1_decode8(_mm256_and_si256(wi, _mm256_set1_epi32(0xF)), mag8); - __m256 vhi = e2m1_decode8(_mm256_srli_epi32(wi, 4), mag8); - return _mm256_fmadd_ps(vlo, _mm256_loadu_ps(xeb), _mm256_mul_ps(vhi, _mm256_loadu_ps(xob))); -} - -__attribute__((target("avx2,fma"))) -float dot_dsfp4_avx2(const uint8_t* packed, const uint8_t* scale, const float* xe, - const float* xo, int K, const float* e2m1, const float* e8m0) { - const __m256 mag8 = _mm256_loadu_ps(e2m1); - __m256 acc0 = _mm256_setzero_ps(), acc1 = _mm256_setzero_ps(); - const int nb = K / 32; - for (int b = 0; b < nb; ++b) { - const uint8_t* pk = packed + (size_t)b * 16; - const float* xeb = xe + (size_t)b * 16; - const float* xob = xo + (size_t)b * 16; - const __m256 sc = _mm256_set1_ps(e8m0[scale[b]]); - acc0 = _mm256_fmadd_ps(dsfp4_half_avx2(pk, xeb, xob, mag8), sc, acc0); - acc1 = _mm256_fmadd_ps(dsfp4_half_avx2(pk + 8, xeb + 8, xob + 8, mag8), sc, acc1); - } - return hsum256(_mm256_add_ps(acc0, acc1)); -} -#endif - -dsdot_fn select_dsdot() { - const IsaTier t = pick_isa(); -#if CPU_MOE_X86 - if (t >= ISA_AVX512) return dot_dsfp4_avx512; - if (t >= ISA_AVX2) return dot_dsfp4_avx2; -#endif - (void)t; - return dot_dsfp4_scalar; -} - -// ------------------------- mxfp4 (gpt-oss) GEMV ----------------------------- -// Transposed split-K layout: blk[Kpairs, N2] (N innermost), scl[Kpairs/16, N2] -// e8m0 per 32-K. Computes out[c] = sum_kb (E2M1[lo]*x[2kb] + E2M1[hi]*x[2kb+1]) -// * 2^(e8m0-127) for a contiguous column tile (blk/scl already offset to col 0 of -// the tile). Vectorized over N (16 columns / __m512), K stays the outer (cache- -// sequential) loop. Used by both gate_up (K=H) and down (K=I). -using mxgemv_fn = void (*)(float*, const uint8_t*, const uint8_t*, const bf16_t*, int, int, - int, const float*, const float*); - -void mxfp4_gemv_scalar(float* out, const uint8_t* blk, const uint8_t* scl, const bf16_t* x, - int Kpairs, int N2, int ncol, const float* e2m1, const float* e8m0) { - for (int c = 0; c < ncol; ++c) out[c] = 0.0f; - for (int kb = 0; kb < Kpairs; ++kb) { - const uint8_t* w = blk + (size_t)kb * N2; - const uint8_t* s = scl + (size_t)(kb >> 4) * N2; - const float xl = bf16_to_f32(x[2 * kb]); - const float xh = bf16_to_f32(x[2 * kb + 1]); - for (int c = 0; c < ncol; ++c) { - const uint8_t byte = w[c]; - out[c] += (e2m1[byte & 0xF] * xl + e2m1[byte >> 4] * xh) * e8m0[s[c]]; - } - } -} - -#if CPU_MOE_X86 -__attribute__((target("avx512f"))) -void mxfp4_gemv_avx512(float* out, const uint8_t* blk, const uint8_t* scl, const bf16_t* x, - int Kpairs, int N2, int ncol, const float* e2m1, const float* e8m0) { - (void)e8m0; // e8m0[c]=2^(c-127) computed via bit construction (no gather) - const __m512 lut = _mm512_loadu_ps(e2m1); - const __m512i loma = _mm512_set1_epi32(0xF); - // K-outer / N-inner: each kb cache line is read once and all live column chunks - // (up to 4 -> 64 cols) accumulate from registers, so DRAM/L2 stream the tile once. - int c0 = 0; - for (; c0 + 16 <= ncol; c0 += 64) { - const int nchunk = std::min(4, (ncol - c0) / 16); - __m512 acc[4]; - for (int ci = 0; ci < nchunk; ++ci) acc[ci] = _mm512_setzero_ps(); - for (int kblk = 0; kblk < Kpairs; kblk += 16) { // 16 K-pairs = 32 K = one scale row - __m512 sc[4]; - for (int ci = 0; ci < nchunk; ++ci) { - __m128i sraw = _mm_loadu_si128(reinterpret_cast( - scl + (size_t)(kblk >> 4) * N2 + c0 + ci * 16)); - sc[ci] = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu8_epi32(sraw), 23)); - } - __m512 blk_acc[4]; - for (int ci = 0; ci < nchunk; ++ci) blk_acc[ci] = _mm512_setzero_ps(); - for (int kk = 0; kk < 16; ++kk) { - const int kb = kblk + kk; - const uint8_t* wbase = blk + (size_t)kb * N2 + c0; - // The transposed layout strides K by N2 bytes; prefetch ahead so the strided - // reads are not exposed to DRAM latency (the HW streamer misses big strides). - constexpr int PFD = 8; - if (kb + PFD < Kpairs) - _mm_prefetch(reinterpret_cast(blk + (size_t)(kb + PFD) * N2 + c0), - _MM_HINT_T0); - const __m512 xl = _mm512_set1_ps(bf16_to_f32(x[2 * kb])); - const __m512 xh = _mm512_set1_ps(bf16_to_f32(x[2 * kb + 1])); - for (int ci = 0; ci < nchunk; ++ci) { - __m512i wi = _mm512_cvtepu8_epi32( - _mm_loadu_si128(reinterpret_cast(wbase + ci * 16))); - __m512 vlo = _mm512_permutexvar_ps(_mm512_and_si512(wi, loma), lut); - __m512 vhi = _mm512_permutexvar_ps(_mm512_and_si512(_mm512_srli_epi32(wi, 4), loma), lut); - blk_acc[ci] = _mm512_fmadd_ps(vlo, xl, blk_acc[ci]); - blk_acc[ci] = _mm512_fmadd_ps(vhi, xh, blk_acc[ci]); - } - } - for (int ci = 0; ci < nchunk; ++ci) acc[ci] = _mm512_fmadd_ps(blk_acc[ci], sc[ci], acc[ci]); - } - for (int ci = 0; ci < nchunk; ++ci) _mm512_storeu_ps(out + c0 + ci * 16, acc[ci]); - } - for (int c = c0; c < ncol; ++c) { // tail columns (< 16) - float o = 0.0f; - for (int kb = 0; kb < Kpairs; ++kb) { - const uint8_t byte = blk[(size_t)kb * N2 + c]; - uint32_t bits = (uint32_t)scl[(size_t)(kb >> 4) * N2 + c] << 23; - float sc; - std::memcpy(&sc, &bits, 4); - o += (e2m1[byte & 0xF] * bf16_to_f32(x[2 * kb]) + - e2m1[byte >> 4] * bf16_to_f32(x[2 * kb + 1])) * sc; - } - out[c] = o; - } } -__attribute__((target("avx2,fma"))) -void mxfp4_gemv_avx2(float* out, const uint8_t* blk, const uint8_t* scl, const bf16_t* x, - int Kpairs, int N2, int ncol, const float* e2m1, const float* e8m0) { - (void)e8m0; // e8m0[s]=2^(s-127) built via s<<23 (no gather) - const __m256 mag8 = _mm256_loadu_ps(e2m1); - int c0 = 0; - for (; c0 + 8 <= ncol; c0 += 32) { // up to 4 chunks of 8 = 32 cols - const int nchunk = std::min(4, (ncol - c0) / 8); - __m256 acc[4]; - for (int ci = 0; ci < nchunk; ++ci) acc[ci] = _mm256_setzero_ps(); - for (int kblk = 0; kblk < Kpairs; kblk += 16) { // 16 K-pairs = one scale row - __m256 sc[4]; - for (int ci = 0; ci < nchunk; ++ci) { - __m128i sraw = _mm_loadl_epi64(reinterpret_cast( - scl + (size_t)(kblk >> 4) * N2 + c0 + ci * 8)); - sc[ci] = _mm256_castsi256_ps(_mm256_slli_epi32(_mm256_cvtepu8_epi32(sraw), 23)); - } - __m256 blk_acc[4]; - for (int ci = 0; ci < nchunk; ++ci) blk_acc[ci] = _mm256_setzero_ps(); - for (int kk = 0; kk < 16; ++kk) { - const int kb = kblk + kk; - const uint8_t* wbase = blk + (size_t)kb * N2 + c0; - constexpr int PFD = 8; - if (kb + PFD < Kpairs) - _mm_prefetch(reinterpret_cast(blk + (size_t)(kb + PFD) * N2 + c0), - _MM_HINT_T0); - const __m256 xl = _mm256_set1_ps(bf16_to_f32(x[2 * kb])); - const __m256 xh = _mm256_set1_ps(bf16_to_f32(x[2 * kb + 1])); - for (int ci = 0; ci < nchunk; ++ci) { - __m256i wi = _mm256_cvtepu8_epi32( - _mm_loadl_epi64(reinterpret_cast(wbase + ci * 8))); - __m256 vlo = e2m1_decode8(_mm256_and_si256(wi, _mm256_set1_epi32(0xF)), mag8); - __m256 vhi = e2m1_decode8(_mm256_srli_epi32(wi, 4), mag8); - blk_acc[ci] = _mm256_fmadd_ps(vlo, xl, blk_acc[ci]); - blk_acc[ci] = _mm256_fmadd_ps(vhi, xh, blk_acc[ci]); - } - } - for (int ci = 0; ci < nchunk; ++ci) acc[ci] = _mm256_fmadd_ps(blk_acc[ci], sc[ci], acc[ci]); - } - for (int ci = 0; ci < nchunk; ++ci) _mm256_storeu_ps(out + c0 + ci * 8, acc[ci]); - } - for (int c = c0; c < ncol; ++c) { // tail columns (< 8); none when ncol%8==0 - float o = 0.0f; - for (int kb = 0; kb < Kpairs; ++kb) { - const uint8_t byte = blk[(size_t)kb * N2 + c]; - uint32_t bits = (uint32_t)scl[(size_t)(kb >> 4) * N2 + c] << 23; - float sc; - std::memcpy(&sc, &bits, 4); - o += (e2m1[byte & 0xF] * bf16_to_f32(x[2 * kb]) + - e2m1[byte >> 4] * bf16_to_f32(x[2 * kb + 1])) * sc; - } - out[c] = o; - } -} -#endif - -mxgemv_fn select_mxgemv() { - const IsaTier t = pick_isa(); +inline void cpu_relax() { #if CPU_MOE_X86 - if (t >= ISA_AVX512) return mxfp4_gemv_avx512; - if (t >= ISA_AVX2) return mxfp4_gemv_avx2; + _mm_pause(); +#elif defined(__aarch64__) + asm volatile("yield" ::: "memory"); #endif - (void)t; - return mxfp4_gemv_scalar; -} - -// Round a clamped |x|<=448 to nearest float8-e4m3 (RNE), back to fp32. Matches -// torch.float8_e4m3fn / triton .to(float8e4nv). -inline float e4m3_round(float x) { - const float sign = x < 0.0f ? -1.0f : 1.0f; - const float a = std::fabs(x); - if (a == 0.0f) return 0.0f; - if (a >= 448.0f) return sign * 448.0f; - int e; - std::frexp(a, &e); // a in [2^(e-1), 2^e) - float step = std::ldexp(1.0f, e - 4); - const float min_step = std::ldexp(1.0f, -9); // e4m3 subnormal step (2^-9) - if (step < min_step) step = min_step; - float r = std::nearbyint(a / step) * step; - if (r > 448.0f) r = 448.0f; - return sign * r; -} - -// IEEE ceil(log2(v)) for v>0 (matches dsv4 _log2_ceil / fast_round_scale). -inline int ceil_log2_pos(float v) { - uint32_t bits; - std::memcpy(&bits, &v, sizeof(bits)); - const int exp = (int)((bits >> 23) & 0xFF); - const int man = (int)(bits & 0x7FFFFF); - return exp - 127 + (man != 0 ? 1 : 0); } -// Split an interleaved bf16 row into fp32 even/odd halves (even[m]=src[2m]). -// bf16->fp32 is exact, so this only reorders -- done once per token/route and -// reused across every output row of the GEMV. -inline void deinterleave_bf16_f32(const bf16_t* src, float* even, float* odd, int K) { - for (int m = 0; m < K / 2; ++m) { - even[m] = bf16_to_f32(src[2 * m]); - odd[m] = bf16_to_f32(src[2 * m + 1]); - } -} - -// DeepSeek-V4 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); - float amax = 1e-4f; - for (int i = b0; i < b1; ++i) amax = std::max(amax, std::fabs(bf16_to_f32(src[i]))); - const float s = std::ldexp(1.0f, ceil_log2_pos(amax * (1.0f / 448.0f))); - const float inv_s = 1.0f / s; - for (int i = b0; i < b1; ++i) { - float q = bf16_to_f32(src[i]) * inv_s; - q = std::min(448.0f, std::max(-448.0f, q)); - dst[i] = f32_to_bf16(e4m3_round(q) * s); - } - } -} - -// --------------------------------- executor --------------------------------- - -struct CpuMoeExecutor; +constexpr int kMaxTopK = 64; +// One MoE layer on the first num_tokens rows of the pinned IO. +// It is also the doorbell word the GPU writes, a constant baked into each captured graph; num_tokens >= 1 keeps it nonzero, since 0 means no doorbell. struct MoeTask { - CpuMoeExecutor* exec; - int layer_id; - int num_tokens; - const bf16_t* x; // [num_tokens, H] - const int32_t* ids; // [num_tokens, top_k] (raw expert ids; <0 = skip) - const float* w; // [num_tokens, top_k] - bf16_t* y; // [num_tokens, H] + uint64_t layer : 32; + uint64_t num_tokens : 32; }; +static_assert(sizeof(MoeTask) == sizeof(uint64_t), "MoeTask must pack into one 64-bit word"); -// Output-row tiling. Small enough to give every worker independent work even at -// batch size 1; large enough to amortize the atomic work-grab. -// -// Bandwidth notes (Sapphire Rapids 8480+, 13 cores): the two passes already read -// every expert weight byte exactly once per token (each output row block is owned -// by one worker), and x stays hot in L1 across a (token,expert)'s rows -- so the -// kernel is single-read bandwidth-optimal at bs=1 (~205 GB/s vs ~55 GB/s PCIe). -// One worker per *physical* core, pinned, is the sweet spot; SMT oversubscription -// thrashes the spin-barrier. Deferred (not worth it here / for this workload): -// - AMX-bf16: a GEMM tile engine; decode is M=1 GEMV so tiles sit idle. It would -// only pay off in a grouped/batched (dedup) path. -// - expert dedup for bs>1: read each distinct expert once and GEMM its tokens. -// Helps locality+bytes when bs is large; decode batches here are tiny (<=4). -// - NUMA: a single node is assumed. Multi-socket machines would split each -// expert's K dimension per node (banks are already per-row contiguous). -constexpr int IBLK = 32; -constexpr int HBLK = 32; - -// -------------------------------- Q4_0 (W4A8) -------------------------------- -// Native GGUF Q4_0 experts (gemma4 GGUF): per-32 block = fp16 scale d + 16 packed -// bytes; byte j holds element j in its low nibble and j+16 in its high nibble, so a -// block's storage order is [lo0..lo15, hi0..hi15] and w = (nibble - 8) * d. Matches -// the reference dequant (models/gguf/dequant.py) and the packed banks the GPU offload -// path streams. -// -// llama.cpp ggml_vec_dot_q4_0_q8_0: W4A8. The activation is pre-quantized to Q8_0 -// (per-32-block int8 ``aq`` + fp32 scale ``asb``); each block unpacks its 16 bytes to -// 32 int8 weights in [-8,7] (bytes_from_nibbles_32: low nibbles -> elems 0..15, high -// -> 16..31) and runs an integer block dot -- VPDPBUSD (AVX-VNNI) or VPMADDUBSW+VPMADDWD -// (AVX2) with the ggml sign trick |w|*(sign(w)*a)=w*a, or a scalar int loop -- then -// scales the block sum by wd*xd in fp32. No fp weight dequant / shuffle chain. The GPU -// offload path (ggml_moe_a8_vec / MMVQ) is also W4A8, so cpu and hybrid stay close. -using q4dot_fn = float (*)(const uint8_t*, const int8_t*, const float*, int); - -float q4_0_dot_i8_scalar(const uint8_t* w, const int8_t* aq, const float* asb, int K) { - float acc = 0.0f; - const int nb = K / 32; - for (int b = 0; b < nb; ++b) { - const uint8_t* blk = w + (size_t)b * 18; - uint16_t dh; - std::memcpy(&dh, blk, sizeof(dh)); - const uint8_t* q = blk + 2; // 16 nibble bytes - const int8_t* a = aq + (size_t)b * 32; - int isum = 0; - for (int j = 0; j < 16; ++j) { - isum += ((int)(q[j] & 0x0F) - 8) * (int)a[j]; // elem j - isum += ((int)(q[j] >> 4) - 8) * (int)a[16 + j]; // elem 16+j - } - acc += fp16_to_f32(dh) * asb[b] * (float)isum; - } - return acc; -} - -#if CPU_MOE_X86 -// fp16 block scale -> fp32 via HW F16C (single value in lane 0). -__attribute__((target("f16c"))) -static inline float q4_scale(uint16_t h) { - return _mm_cvtss_f32(_mm_cvtph_ps(_mm_cvtsi32_si128((int)h))); +inline uint64_t to_word(MoeTask t) { + uint64_t w; + std::memcpy(&w, &t, sizeof(w)); + return w; } -// Unpack one Q4_0 block's 16 bytes -> 32 int8 weights in [-8,7] (elems 0..15 = low -// nibbles, 16..31 = high nibbles). ``eight`` = _mm256_set1_epi8(8). -__attribute__((target("avx2"))) -static inline __m256i q4_unpack32(const uint8_t* blk, __m128i mask, __m256i eight) { - const __m128i qb = _mm_loadu_si128(reinterpret_cast(blk + 2)); - const __m128i lo = _mm_and_si128(qb, mask); - const __m128i hi = _mm_and_si128(_mm_srli_epi16(qb, 4), mask); - return _mm256_sub_epi8(_mm256_set_m128i(hi, lo), eight); +inline MoeTask from_word(uint64_t w) { + MoeTask t; + std::memcpy(&t, &w, sizeof(t)); + return t; } -// AVX2 W4A8 (llama.cpp non-VNNI mul_sum_i8_pairs): integer block dot via VPMADDUBSW + -// VPMADDWD (sign trick), scaled by wd*xd. |aw*sa| pair sums <= 8*127*2 < 32767 -> no -// int16 saturation. This is the fast path on AVX2 CPUs without AVX-VNNI (and the -// avx512-tier fallback, since the block dot is 256-bit either way). -__attribute__((target("avx2,fma,f16c"))) -float q4_0_dot_i8_avx2(const uint8_t* w, const int8_t* aq, const float* asb, int K) { - const __m128i mask = _mm_set1_epi8(0x0F); - const __m256i eight = _mm256_set1_epi8(8); - const __m256i ones16 = _mm256_set1_epi16(1); - __m256 accF = _mm256_setzero_ps(); - const int nb = K / 32; - for (int b = 0; b < nb; ++b) { - const uint8_t* blk = w + (size_t)b * 18; - _mm_prefetch(reinterpret_cast(blk) + 512, _MM_HINT_T0); - uint16_t dh; - std::memcpy(&dh, blk, sizeof(dh)); - __m256i wq = q4_unpack32(blk, mask, eight); - __m256i a = _mm256_loadu_si256(reinterpret_cast(aq + (size_t)b * 32)); - __m256i aw = _mm256_sign_epi8(wq, wq); // |wq| (unsigned operand) - __m256i sa = _mm256_sign_epi8(a, wq); // sign(wq) * a (signed operand) - __m256i d32 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, sa), ones16); // 8 int32 - accF = _mm256_fmadd_ps(_mm256_cvtepi32_ps(d32), _mm256_set1_ps(q4_scale(dh) * asb[b]), accF); - } - return hsum256(accF); -} - -// AVX-VNNI W4A8: one VPDPBUSD per block (the fast path on modern CPUs). -__attribute__((target("avx2,avxvnni,fma,f16c"))) -float q4_0_dot_i8_vnni(const uint8_t* w, const int8_t* aq, const float* asb, int K) { - const __m128i mask = _mm_set1_epi8(0x0F); - const __m256i eight = _mm256_set1_epi8(8); - __m256 accF = _mm256_setzero_ps(); - const int nb = K / 32; - for (int b = 0; b < nb; ++b) { - const uint8_t* blk = w + (size_t)b * 18; - _mm_prefetch(reinterpret_cast(blk) + 512, _MM_HINT_T0); - uint16_t dh; - std::memcpy(&dh, blk, sizeof(dh)); - __m256i wq = q4_unpack32(blk, mask, eight); - __m256i a = _mm256_loadu_si256(reinterpret_cast(aq + (size_t)b * 32)); - __m256i aw = _mm256_sign_epi8(wq, wq); // |wq| (unsigned operand) - __m256i sa = _mm256_sign_epi8(a, wq); // sign(wq) * a (signed operand) - __m256i di = _mm256_dpbusd_avx_epi32(_mm256_setzero_si256(), aw, sa); - // All 32 elems of the block share wd*xd; distribute over di's 8 partial sums and - // reduce at the end (equivalent to scale * block_total). - accF = _mm256_fmadd_ps(_mm256_cvtepi32_ps(di), _mm256_set1_ps(q4_scale(dh) * asb[b]), accF); - } - return hsum256(accF); -} -#endif // CPU_MOE_X86 +// Every weight format this module serves, keyed by the name cpu_executor.py passes; weight_formats() exports the names. +const std::pair (*)(const MoeParams&, const FormatArgs&)> kFormats[] = { + {"bf16", make_bf16_format}, {"nvfp4", make_nvfp4_format}, {"mxfp4_triton", make_mxfp4_format}, + {"ds_fp4", make_ds_fp4_format}, {"fp8_block", make_fp8_block_format}, {"gguf", make_gguf_format}, + {"gguf_scaled", make_gguf_scaled_format}, +}; -// All tiers are W4A8 (int8 activations pre-quantized to Q8_0). AVX-VNNI is orthogonal to -// the ISA tier (gated by cpu_has_avxvnni() / FREETOKEN_CPU_MOE_NO_VNNI), so it wins when -// present; otherwise the 256-bit VPMADDUBSW kernel covers both the avx2 and avx512 tiers. -q4dot_fn select_q4dot() { - const IsaTier t = pick_isa(); -#if CPU_MOE_X86 - if (cpu_has_avxvnni()) return q4_0_dot_i8_vnni; - if (t >= ISA_AVX2) return q4_0_dot_i8_avx2; -#endif - (void)t; - return q4_0_dot_i8_scalar; +std::unique_ptr make_format(const std::string& name, const MoeParams& p, const FormatArgs& a) { + for (const auto& [n, make] : kFormats) + if (name == n) return make(p, a); + throw std::invalid_argument("unknown CPU MoE weight format " + name); } -enum WFmt { WF_BF16 = 0, WF_NVFP4 = 1, WF_MXFP4 = 2, WF_DSFP4 = 3, WF_Q4_0 = 4 }; - -// Each ctor pointer arg is the address of a CPU int64 array of length -// num_layers (one base address per layer, built by cpu_executor.py's -// _make_table), not a single flat bank. tbl_at resolves -// tbl[layer_id] once per task/pass; a null table (bank unused by this fmt, ptr -// arg 0) resolves to nullptr without dereferencing. -inline const void* tbl_at(const uint64_t* tbl, int layer_id) { - return tbl ? reinterpret_cast(tbl[layer_id]) : nullptr; -} +// One format instance as cpu_executor.py passes it: name, bank tables by role, integer options. +using FormatSpec = std::tuple, std::map>; struct CpuMoeExecutor { + // A host-function callback gets one pointer, so it carries the executor with the task. + struct HostTask { + CpuMoeExecutor* exec; + MoeTask task; + }; + + const int num_layers, num_experts; + const bool apply_on_input; + MoeParams params; // the formats keep a reference; declared before fmts + std::vector> fmts; + std::vector layer_fmt; // [num_layers], index into fmts + int n_iblk = 0, n_hblk = 0; // kRowBlock blocks per intermediate row and per output row + + // Pinned IO shared by every layer and batch size (one task in flight): [max_tokens, H] x and y, [max_tokens, top_k] ids and w. + const bf16_t* io_x; + const int32_t* io_ids; + const float* io_w; + bf16_t* io_y; + AlignedVector g_scratch; // [max_tokens * top_k * I] intermediate + + // Worker pool. int num_threads; - int num_layers, num_experts, top_k; - int H, I; - int act, apply_on_input; - int fmt; // WFmt - bool needs_di = false; // pre-deinterleave activations to fp32 (nvfp4/ds_fp4) - // Per-layer pointer tables (one base address per layer, see tbl_at). gate_up_tbl - // doubles as the bf16 gate_up table and the nvfp4/mxfp4/q4_0/ds_fp4 packed-gate_up - // table (down_tbl likewise for down); which reinterpretation applies is picked by - // fmt at each resolve site (see gemm1_dot/gemm2_dot/do_pass1_mxfp4/do_pass1_dsfp4). - const uint64_t* gate_up_tbl; // bf16: [E,2I,H] rows; else: packed e2m1/mxfp4-blocks - const uint64_t* down_tbl; // bf16: [E,H,I] rows; else: packed e2m1/mxfp4-blocks - const uint64_t* gu_scale_tbl; // nvfp4/mxfp4/ds_fp4: [E,2I,*] block scales - const uint64_t* gu_global_tbl; // nvfp4: [E,2I] fp16 row globals - const uint64_t* dn_scale_tbl; // nvfp4/mxfp4/ds_fp4: [E,H,*] block scales - const uint64_t* dn_global_tbl; // nvfp4: [E,H] fp16 row globals - const uint64_t* gu_bias_tbl; // mxfp4: [E,2I] bf16 biases - const uint64_t* dn_bias_tbl; // mxfp4: [E,H] bf16 biases - float swiglu_alpha; - float swiglu_limit; // +inf == no clamp - dot_fn dot; - nvdot_fn nvdot; - nvi8dot_fn nvi8dot = nullptr; // AVX-VNNI W4A8 nvfp4 dot (nullptr -> use fp32 nvdot) - bool use_vnni = false; // nvfp4 + AVX-VNNI: decode via int8 VPDPBUSD (W4A8) - bool use_q4a8 = false; // q4_0: always W4A8 (llama.cpp Q4_0 x Q8_0); int8 pre-quant - dsdot_fn dsdot; - mxgemv_fn mxgemv; - q4dot_fn q4dot; - // ds_fp4: the caller already FP8-round-tripped the input activations on the GPU - // (same reference grid), so submit() must not repeat it on the host-callback - // thread. That scalar per-element pass is single-threaded ON THE DECODE CRITICAL - // PATH (~0.3ms/layer at H=4096, every worker and the GPU waiting on it); moving - // it to a captured GPU elementwise kernel removes it while keeping the official - // W4A8 numerics bit-exact. Set via set_input_prequant (see cpu_executor.py). - bool input_prequant = false; - // Q4_0 packed-row byte strides (H/32*18 for gate_up over K=H, I/32*18 for down over K=I). - int q4_gu_row_bytes = 0, q4_dn_row_bytes = 0; - float e2m1_lut[16]; - float e4m3_lut[256]; - float e8m0_lut[256]; // mxfp4 block scale: 2^(s-127), s clamped to [0,254] - const char* isa; - - std::vector g_scratch; // [max_tokens * top_k * I] intermediate - std::vector xq_scratch; // [max_tokens * H] ds_fp4 fp8-roundtripped input - // ds_fp4 activations pre-deinterleaved to fp32 (even/odd K) for the row-major dot. - std::vector xe_scratch, xo_scratch; // [max_tokens * H/2] (input) - std::vector ge_scratch, go_scratch; // [max_tokens*top_k*I/2] (intermediate) - // AVX-VNNI W4A8: per-16-block int8 activations [even(8),odd(8)] + per-block scale. - std::vector xi8_scratch, gi8_scratch; // [max_tokens*H], [max_tokens*top_k*I] - std::vector xas_scratch, gas_scratch; // [max_tokens*H/16], [..*top_k*I/16] - std::string isa_str; - + std::vector core_ids; // worker tid -> logical CPU to pin to (may be empty) std::vector workers; std::mutex task_mtx; std::condition_variable task_cv; 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}; - - std::atomic p1_next{0}; + MoeTask cur_task{}; + // Separate cache lines: idle workers poll `submitted` while busy workers hit p1_next/p2_next. + alignas(64) std::atomic submitted{0}; + alignas(64) std::atomic completed{0}; + + // 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; + + // Work split of the current task. + alignas(64) std::atomic p1_next{0}; std::atomic p2_next{0}; - std::atomic prt_next{0}; // ds_fp4 intermediate fp8 round-trip phase + std::atomic prt_next{0}; // intermediate-row prep phase (prep_inter) int64_t p1_total = 0, p2_total = 0, prt_total = 0; - int n_iblk = 0, n_hblk = 0; std::atomic done_count{0}; std::atomic bar_count{0}; std::atomic bar_sense{0}; - std::vector owned_tasks; // persistent task descriptors (graph-stable) - std::vector core_ids; // worker tid -> logical CPU to pin to (may be empty) - // ---- Flag-based GPU<->CPU handshake (replaces the per-layer cudaLaunchHostFunc pair) ---- - // A tiny GPU kernel bumps ready_flags[slot] at submit; this coordinator thread busy-polls - // it, runs the slot's task on the worker pool, and sets done_flags[slot], which a GPU - // spin-wait kernel polls at sync. This removes the ~2x30-50us host-func dispatch round - // trips per MoE layer per decode step that otherwise idle the GPU (~6 ms/step on a - // 75-layer model). One slot per (layer, decode batch size) pair -- the Python side - // allocates slots as tasks are created. Flags live in mapped-pinned host memory (UVA: - // the same pointers are used by the GPU kernels and by this thread). + // The GPU rings the doorbell with a task word at submit; this coordinator thread busy-polls it, runs the task on the worker pool, and completes it, which the GPU waits on at sync. + // This removes the ~2x30-50us host-func dispatch round trips per MoE layer per decode step that otherwise idle the GPU (~6 ms/step on a 75-layer model). + std::unique_ptr gpu_sync; + std::unique_ptr kernel_io; // replaces the copies around the doorbell once enable_kernel_io() runs std::thread coord_thread; std::atomic coord_stop{false}; - volatile int64_t* ready_flags = nullptr; // GPU increments, this thread polls - volatile int64_t* done_flags = nullptr; // this thread sets, GPU spin-waits - int coord_num_slots = 0; - std::vector flag_task; // slot -> task (registered lazily) - std::vector flag_served; // slot -> completed dispatch count (tests/debug) - std::mutex flag_task_mtx; - - // Portable ordering for the flag handshake: "ready observed => the DMA'd inputs that - // preceded the bump are visible" and "y stores are visible before done". Plain - // volatile loads lean on x86 TSO; acquire/release makes it hold on aarch64 too - // (GH200/Jetson) at zero x86 cost. (MSVC branch is x86-only today: compiler barrier - // + TSO.) - static int64_t flag_load_acquire(const volatile int64_t* p) { -#if defined(_MSC_VER) - const int64_t v = *p; - _ReadWriteBarrier(); - return v; -#else - return __atomic_load_n(const_cast(p), __ATOMIC_ACQUIRE); -#endif - } - - static void flag_store_release(volatile int64_t* p, int64_t v) { -#if defined(_MSC_VER) - _ReadWriteBarrier(); - *p = v; -#else - __atomic_store_n(const_cast(p), v, __ATOMIC_RELEASE); -#endif - } + std::atomic flag_served{0}; // completed doorbells (tests, watchdog) + std::map host_tasks; // host-function userData, keyed by doorbell word (graph-stable) CpuMoeExecutor(int num_threads_, int num_layers_, int num_experts_, int top_k_, - int hidden_size, int inter_size, int max_tokens, int activation_id, - int apply_router_weight_on_input, int weight_format, - uintptr_t gate_up_ptr, uintptr_t down_ptr, uintptr_t gate_up_scale_ptr, - uintptr_t gate_up_global_ptr, uintptr_t down_scale_ptr, - uintptr_t down_global_ptr, uintptr_t gate_up_bias_ptr, - uintptr_t down_bias_ptr, double swiglu_alpha_, double swiglu_limit_, - std::vector core_ids_) - : num_threads(num_threads_ > 0 ? num_threads_ : 1), - num_layers(num_layers_), + int hidden_size, int inter_size, int max_tokens_, int activation_id, + int apply_router_weight_on_input, std::vector formats, + std::vector layer_formats, double swiglu_alpha_, double swiglu_limit_, + std::vector core_ids_, uintptr_t x_ptr, uintptr_t ids_ptr, uintptr_t w_ptr, + uintptr_t y_ptr) + : num_layers(num_layers_), num_experts(num_experts_), - top_k(top_k_), - H(hidden_size), - I(inter_size), - act(activation_id), - apply_on_input(apply_router_weight_on_input), - fmt(weight_format), - gate_up_tbl(reinterpret_cast(gate_up_ptr)), - down_tbl(reinterpret_cast(down_ptr)), - gu_scale_tbl(reinterpret_cast(gate_up_scale_ptr)), - gu_global_tbl(reinterpret_cast(gate_up_global_ptr)), - dn_scale_tbl(reinterpret_cast(down_scale_ptr)), - dn_global_tbl(reinterpret_cast(down_global_ptr)), - gu_bias_tbl(reinterpret_cast(gate_up_bias_ptr)), - dn_bias_tbl(reinterpret_cast(down_bias_ptr)), - swiglu_alpha(static_cast(swiglu_alpha_)), - swiglu_limit(static_cast(swiglu_limit_)), + apply_on_input(apply_router_weight_on_input != 0), + params{top_k_, max_tokens_, hidden_size, inter_size, activation_id, static_cast(swiglu_alpha_), + static_cast(swiglu_limit_)}, + io_x(reinterpret_cast(x_ptr)), + io_ids(reinterpret_cast(ids_ptr)), + io_w(reinterpret_cast(w_ptr)), + io_y(reinterpret_cast(y_ptr)), + num_threads(num_threads_ > 0 ? num_threads_ : 1), core_ids(std::move(core_ids_)) { - DotChoice c = select_dot(); - dot = c.fn; - nvdot = select_nvdot(); - dsdot = select_dsdot(); - mxgemv = select_mxgemv(); - q4dot = select_q4dot(); - if (weight_format == WF_Q4_0) { - if (H % 32 != 0 || I % 32 != 0) - throw std::runtime_error("Q4_0 CPU MoE requires H and I to be multiples of 32"); - q4_gu_row_bytes = (H / 32) * 18; // K = H (gate_up rows) - q4_dn_row_bytes = (I / 32) * 18; // K = I (down rows) - } - isa = c.name; - // nvfp4 (AVX-VNNI only): W4A8 int8 decode when the CPU supports it. q4_0 is always - // W4A8 (activations pre-quantized to Q8_0); select_q4dot picks VPDPBUSD / VPMADDUBSW - // / scalar for the tier, so the tag reflects which of those q4dot resolved to. - nvi8dot = select_nvi8dot(); - use_vnni = (weight_format == WF_NVFP4) && (nvi8dot != nullptr); - use_q4a8 = (weight_format == WF_Q4_0); - const char* q4tag = use_q4a8 ? (cpu_has_avxvnni() ? "+vnni(q4_0-w4a8)" : "+q4_0-w4a8") : ""; - const char* vnni_tag = - cpu_has_avx512vnni() ? "+avx512vnni(nvfp4-w4a8)" : "+vnni(nvfp4-w4a8)"; - isa_str = std::string(c.name) + (use_vnni ? vnni_tag : "") + q4tag; - isa = isa_str.c_str(); - for (int i = 0; i < 16; ++i) e2m1_lut[i] = kE2M1[i]; - for (int i = 0; i < 256; ++i) e4m3_lut[i] = e4m3_decode((uint8_t)i); - // e8m0 (mxfp4 block scale) = 2^(s-127); the GPU GEMV clamps s to [0,254]. - for (int i = 0; i < 256; ++i) e8m0_lut[i] = std::ldexp(1.0f, std::min(i, 254) - 127); - g_scratch.assign(static_cast(max_tokens) * top_k * I, 0); - // Row-major fp4 (nvfp4/ds_fp4) pre-deinterleaves activations to fp32 even/odd. - needs_di = (fmt == WF_NVFP4 || fmt == WF_DSFP4); - if (needs_di) { - if (fmt == WF_DSFP4) xq_scratch.assign(static_cast(max_tokens) * H, 0); - xe_scratch.assign(static_cast(max_tokens) * (H / 2), 0); - xo_scratch.assign(static_cast(max_tokens) * (H / 2), 0); - ge_scratch.assign(static_cast(max_tokens) * top_k * (I / 2), 0); - go_scratch.assign(static_cast(max_tokens) * top_k * (I / 2), 0); - if (use_vnni) { - xi8_scratch.assign(static_cast(max_tokens) * H, 0); - xas_scratch.assign(static_cast(max_tokens) * (H / 16), 0); - gi8_scratch.assign(static_cast(max_tokens) * top_k * I, 0); - gas_scratch.assign(static_cast(max_tokens) * top_k * (I / 16), 0); - } - } - // q4_0 W4A8: per-32-block Q8_0 activations (int8 + fp32 scale) for input + intermediate. - if (use_q4a8) { - xi8_scratch.assign(static_cast(max_tokens) * H, 0); - xas_scratch.assign(static_cast(max_tokens) * (H / 32), 0); - gi8_scratch.assign(static_cast(max_tokens) * top_k * I, 0); - gas_scratch.assign(static_cast(max_tokens) * top_k * (I / 32), 0); - } + if (params.top_k > kMaxTopK) throw std::invalid_argument("CPU MoE supports top_k <= " + std::to_string(kMaxTopK)); + for (auto& [name, banks, options] : formats) + fmts.push_back(make_format(name, params, FormatArgs{name, std::move(banks), std::move(options)})); + if (static_cast(layer_formats.size()) != num_layers) + throw std::invalid_argument("CPU MoE needs the format of every layer"); + for (const int f : layer_formats) + if (f < 0 || f >= static_cast(fmts.size())) throw std::out_of_range("CPU MoE layer format index out of range"); + layer_fmt = std::move(layer_formats); + n_iblk = (params.I + kRowBlock - 1) / kRowBlock; + n_hblk = (params.H + kRowBlock - 1) / kRowBlock; + g_scratch.assign(static_cast(params.max_tokens) * params.top_k * params.I, 0); for (int t = 0; t < num_threads; ++t) workers.emplace_back([this, t] { worker_loop(t); }); } - // Quantize a bf16 activation row to Q8_0 (llama.cpp): per-32-block symmetric int8 in - // natural order + one fp32 scale (amax/127) per block. Done once per token/route, - // amortized over every output row that the q4_0 W4A8 GEMV reads. K % 32 == 0. - void quant_q8_0(const bf16_t* x, int K, int8_t* aq, float* asb) { - const int nb = K / 32; - for (int b = 0; b < nb; ++b) { - const bf16_t* xb = x + (size_t)b * 32; - float xf[32], amax = 0.0f; - for (int j = 0; j < 32; ++j) { - xf[j] = bf16_to_f32(xb[j]); - amax = std::max(amax, std::fabs(xf[j])); - } - const float d = amax > 0.0f ? amax / 127.0f : 1.0f; - asb[b] = d; - const float inv = amax > 0.0f ? 1.0f / d : 0.0f; - int8_t* o = aq + (size_t)b * 32; - for (int j = 0; j < 32; ++j) - o[j] = (int8_t)std::max(-127, std::min(127, (int)std::lround(xf[j] * inv))); - } - } - - // Quantize the pre-deinterleaved fp32 even/odd activations to per-16-block int8 in the - // [even(8),odd(8)] layout the VNNI dot expects. Done once per token/route (amortized over - // every output row), so a scalar pass is fine relative to the GEMV. - void quant_i8_pg16(const float* xe, const float* xo, int K, int8_t* asi8, float* asb) { - const int nb = K / 16; - for (int b = 0; b < nb; ++b) { - const float* xeb = xe + (size_t)b * 8; - const float* xob = xo + (size_t)b * 8; - float amax = 0.0f; - for (int j = 0; j < 8; ++j) - amax = std::max(amax, std::max(std::fabs(xeb[j]), std::fabs(xob[j]))); - const float s = amax > 0.0f ? amax / 127.0f : 1.0f; - asb[b] = s; - const float inv = 1.0f / s; - int8_t* ae = asi8 + (size_t)b * 16; - for (int j = 0; j < 8; ++j) { - int qe = (int)std::lround(xeb[j] * inv), qo = (int)std::lround(xob[j] * inv); - ae[j] = (int8_t)std::max(-127, std::min(127, qe)); - ae[8 + j] = (int8_t)std::max(-127, std::min(127, qo)); - } - } - } - - // gate_up output row `row` (in [0, 2I)) dotted with activation over K = H. ``e`` is - // the layer-local expert row (0..num_experts); the layer bases (already resolved - // once per task/pass by the caller via tbl_at) pick the layer's own tensors. - // bf16 uses the interleaved bf16 row; nvfp4 uses the pre-split fp32 even/odd halves - // (or, with AVX-VNNI, the per-16-block int8 activations). - inline float gemm1_dot(const bf16_t* gate_up_l, const uint8_t* gu_packed_l, - const uint8_t* gu_scale_l, const uint16_t* gu_global_l, int e, int row, - const bf16_t* x, const float* xe, const float* xo, const int8_t* xi8, - const float* xas) { - if (fmt == WF_BF16) { - const bf16_t* w = gate_up_l + ((size_t)e * (2 * I) + row) * H; - return dot(w, x, H); - } - if (fmt == WF_Q4_0) { - const uint8_t* w = - gu_packed_l + ((size_t)e * (2 * I) + row) * (size_t)q4_gu_row_bytes; - return q4dot(w, xi8, xas, H); // W4A8: int8 activations (Q8_0), scale in xas - } - const size_t r = (size_t)e * (2 * I) + row; - if (use_vnni) - return nvi8dot(gu_packed_l + r * (size_t)(H / 2), gu_scale_l + r * (size_t)(H / 16), - fp16_to_f32(gu_global_l[r]), xi8, H, e4m3_lut, xas); - return nvdot(gu_packed_l + r * (size_t)(H / 2), gu_scale_l + r * (size_t)(H / 16), - fp16_to_f32(gu_global_l[r]), xe, xo, H, e2m1_lut, e4m3_lut); - } - - // down output row `row` (in [0, H)) dotted with the intermediate over K = I. Same - // layer-local-base convention as gemm1_dot. - inline float gemm2_dot(const bf16_t* down_l, const uint8_t* dn_packed_l, - const uint8_t* dn_scale_l, const uint16_t* dn_global_l, int e, int row, - const bf16_t* g, const float* ge, const float* go, const int8_t* gi8, - const float* gas) { - if (fmt == WF_BF16) { - const bf16_t* w = down_l + ((size_t)e * H + row) * I; - return dot(w, g, I); - } - if (fmt == WF_Q4_0) { - const uint8_t* w = dn_packed_l + ((size_t)e * H + row) * (size_t)q4_dn_row_bytes; - return q4dot(w, gi8, gas, I); // W4A8: int8 activations (Q8_0), scale in gas - } - const size_t r = (size_t)e * H + row; - if (use_vnni) - return nvi8dot(dn_packed_l + r * (size_t)(I / 2), dn_scale_l + r * (size_t)(I / 16), - fp16_to_f32(dn_global_l[r]), gi8, I, e4m3_lut, gas); - return nvdot(dn_packed_l + r * (size_t)(I / 2), dn_scale_l + r * (size_t)(I / 16), - fp16_to_f32(dn_global_l[r]), ge, go, I, e2m1_lut, e4m3_lut); - } - void pin_self(int tid) { -#if CPU_MOE_HAS_AFFINITY - if (core_ids.empty()) return; - const int cpu = core_ids[tid % static_cast(core_ids.size())]; - cpu_set_t set; - CPU_ZERO(&set); - CPU_SET(cpu, &set); - pthread_setaffinity_np(pthread_self(), sizeof(set), &set); -#else - (void)tid; -#endif + if (!core_ids.empty()) pin_this_thread(core_ids[tid % static_cast(core_ids.size())]); } ~CpuMoeExecutor() { @@ -1547,23 +217,33 @@ struct CpuMoeExecutor { task_cv.notify_all(); for (auto& th : workers) if (th.joinable()) th.join(); - for (MoeTask* t : owned_tasks) delete t; } - uintptr_t create_task(int layer_id, int num_tokens, uintptr_t x_ptr, - uintptr_t ids_ptr, uintptr_t w_ptr, uintptr_t y_ptr) { - MoeTask* t = new MoeTask{this, - layer_id, - num_tokens, - reinterpret_cast(x_ptr), - reinterpret_cast(ids_ptr), - reinterpret_cast(w_ptr), - reinterpret_cast(y_ptr)}; - owned_tasks.push_back(t); - return reinterpret_cast(t); + MoeTask task(int layer_id, int num_tokens) const { + if (layer_id < 0 || layer_id >= num_layers || num_tokens < 1 || num_tokens > params.max_tokens) + throw std::out_of_range("CPU MoE task (layer " + std::to_string(layer_id) + ", " + + std::to_string(num_tokens) + " tokens) is outside the executor's layers or IO"); + return MoeTask{static_cast(layer_id), static_cast(num_tokens)}; + } + + std::string isa_name() const { + std::vector names; + for (const auto& f : fmts) + if (const std::string n = f->isa_name(); std::find(names.begin(), names.end(), n) == names.end()) names.push_back(n); + std::string out; + for (const auto& n : names) out += (out.empty() ? "" : ",") + n; + return out; } - const char* isa_name() const { return isa; } + const MoeFormat& format_of(const MoeTask& t) const { return *fmts[layer_fmt[t.layer]]; } + + // Whether any route of the task names an expert; do_pass1 and do_pass2 skip the others. + bool has_routes(const MoeTask& t) const { + const int64_t n = static_cast(t.num_tokens) * params.top_k; + for (int64_t i = 0; i < n; ++i) + if (io_ids[i] >= 0 && io_ids[i] < num_experts) return true; + return false; + } void barrier(int& local_sense) { local_sense ^= 1; @@ -1571,276 +251,60 @@ struct CpuMoeExecutor { bar_count.store(0); bar_sense.store(local_sense); } else { - while (bar_sense.load() != local_sense) { -#if CPU_MOE_X86 - _mm_pause(); -#endif - } - } - } - - void do_pass1(const MoeTask* t, int64_t p) { - if (fmt == WF_MXFP4) { - do_pass1_mxfp4(t, p); - return; - } - if (fmt == WF_DSFP4) { - do_pass1_dsfp4(t, p); - return; - } - const int64_t ib = p % n_iblk; - const int64_t tk = p / n_iblk; - const int k = static_cast(tk % top_k); - const int tok = static_cast(tk / top_k); - const int e = t->ids[static_cast(tok) * top_k + k]; - if (e < 0 || e >= num_experts) return; - const float w_in = apply_on_input ? t->w[static_cast(tok) * top_k + k] : 1.0f; - // Resolve this task's layer base once; row indexing below is layer-local (e). - const bf16_t* gate_up_l = reinterpret_cast(tbl_at(gate_up_tbl, t->layer_id)); - const uint8_t* gu_packed_l = reinterpret_cast(tbl_at(gate_up_tbl, t->layer_id)); - const uint8_t* gu_scale_l = reinterpret_cast(tbl_at(gu_scale_tbl, t->layer_id)); - const uint16_t* gu_global_l = - reinterpret_cast(tbl_at(gu_global_tbl, t->layer_id)); - const bf16_t* x_row = t->x + (size_t)tok * H; - const float* xe = needs_di ? xe_scratch.data() + (size_t)tok * (H / 2) : nullptr; - const float* xo = needs_di ? xo_scratch.data() + (size_t)tok * (H / 2) : nullptr; - const int8_t* xi8 = - (use_vnni || use_q4a8) ? xi8_scratch.data() + (size_t)tok * H : nullptr; - const float* xas = use_vnni ? xas_scratch.data() + (size_t)tok * (H / 16) - : use_q4a8 ? xas_scratch.data() + (size_t)tok * (H / 32) - : nullptr; - bf16_t* g_row = g_scratch.data() + ((size_t)tok * top_k + k) * I; - const int i0 = static_cast(ib) * IBLK; - const int i1 = std::min(I, i0 + IBLK); - const bool clamped = act == ACT_SWIGLUOAI || act == ACT_SWIGLU_CLAMP; - const float up_bias = act == ACT_SWIGLUOAI ? 1.0f : 0.0f; - const float lim = swiglu_limit, alpha = swiglu_alpha; - for (int i = i0; i < i1; ++i) { - // gate = row i, up = row I+i - float gate = - gemm1_dot(gate_up_l, gu_packed_l, gu_scale_l, gu_global_l, e, i, x_row, xe, xo, xi8, xas) * w_in; - float up = gemm1_dot(gate_up_l, gu_packed_l, gu_scale_l, gu_global_l, e, I + i, x_row, - xe, xo, xi8, xas) * w_in; - if (clamped) { - // clamp(gate, max=lim) * sigmoid(alpha * gate) * (clamp(up, +-lim) + up_bias) - // -- swigluoai carries the +1 up bias (gpt-oss/MiniMax); swiglu_clamp - // (GLM-5.3) does not. lim == +inf: no clamp. - if (gate > lim) gate = lim; - if (up > lim) up = lim; - else if (up < -lim) up = -lim; - const float glu = gate / (1.0f + std::exp(-gate * alpha)); - g_row[i] = f32_to_bf16(glu * (up + up_bias)); - } else { - g_row[i] = f32_to_bf16(act_apply(act, gate) * up); - } - } - } - - void do_pass2(const MoeTask* t, int64_t p) { - if (fmt == WF_MXFP4) { - do_pass2_mxfp4(t, p); - return; - } - if (fmt == WF_DSFP4) { - do_pass2_dsfp4(t, p); - return; - } - const int64_t hb = p % n_hblk; - const int tok = static_cast(p / n_hblk); - const int h0 = static_cast(hb) * HBLK; - const int h1 = std::min(H, h0 + HBLK); - // Resolve this task's layer base once; row indexing below is layer-local (e). - const bf16_t* down_l = reinterpret_cast(tbl_at(down_tbl, t->layer_id)); - const uint8_t* dn_packed_l = reinterpret_cast(tbl_at(down_tbl, t->layer_id)); - const uint8_t* dn_scale_l = reinterpret_cast(tbl_at(dn_scale_tbl, t->layer_id)); - const uint16_t* dn_global_l = - reinterpret_cast(tbl_at(dn_global_tbl, t->layer_id)); - bf16_t* y_row = t->y + (size_t)tok * H; - for (int h = h0; h < h1; ++h) { - float acc = 0.0f; - for (int k = 0; k < top_k; ++k) { - const int e = t->ids[static_cast(tok) * top_k + k]; - if (e < 0 || e >= num_experts) continue; - const float w_out = apply_on_input ? 1.0f : t->w[static_cast(tok) * top_k + k]; - const size_t gr = (size_t)tok * top_k + k; - const bf16_t* g_row = g_scratch.data() + gr * I; - const float* ge = needs_di ? ge_scratch.data() + gr * (I / 2) : nullptr; - const float* go = needs_di ? go_scratch.data() + gr * (I / 2) : nullptr; - const int8_t* gi8 = (use_vnni || use_q4a8) ? gi8_scratch.data() + gr * I : nullptr; - const float* gas = use_vnni ? gas_scratch.data() + gr * (I / 16) - : use_q4a8 ? gas_scratch.data() + gr * (I / 32) - : nullptr; - acc += gemm2_dot(down_l, dn_packed_l, dn_scale_l, dn_global_l, e, h, g_row, ge, go, gi8, - gas) * w_out; - } - y_row[h] = f32_to_bf16(acc); + while (bar_sense.load() != local_sense) cpu_relax(); } } - // ----------------------------- mxfp4 (gpt-oss) ----------------------------- - // Transposed split-K layout (N innermost), so the GEMV streams K and accumulates - // a contiguous block of N output columns -> cache-line-efficient, no repack, no - // extra host memory. Pass 1 fuses gate_up + clamped-swiglu(+bias); pass 2 fuses - // down(+bias) * router-weight, summed over the token's routes. - // - // Dequant: w = E2M1[code] * 2^(e8m0_scale - 127); two codes per byte (low nibble - // first), one e8m0 scale per 32 contiguous K. Matches kernel/triton/mxfp4_moe.py. - - void do_pass1_mxfp4(const MoeTask* t, int64_t p) { + // Pass 1 item p = one kRowBlock block of one (token, route)'s intermediate row. + void do_pass1(const MoeTask& t, int64_t p) { + const int top_k = params.top_k, H = params.H, I = params.I; const int64_t ib = p % n_iblk; const int64_t tk = p / n_iblk; const int k = static_cast(tk % top_k); const int tok = static_cast(tk / top_k); - const int e = t->ids[static_cast(tok) * top_k + k]; + const int e = io_ids[static_cast(tok) * top_k + k]; if (e < 0 || e >= num_experts) return; - // Resolve this task's layer base once; row indexing below is layer-local (e). - const uint8_t* gu_packed_l = reinterpret_cast(tbl_at(gate_up_tbl, t->layer_id)); - const uint8_t* gu_scale_l = reinterpret_cast(tbl_at(gu_scale_tbl, t->layer_id)); - const bf16_t* gu_bias_l = reinterpret_cast(tbl_at(gu_bias_tbl, t->layer_id)); - const int N2 = 2 * I; // gate_up output width (gate/up interleaved) - const int Hh = H / 2; // packed-K rows - const int i0 = static_cast(ib) * IBLK; - const int i1 = std::min(I, i0 + IBLK); - const int nunit = i1 - i0; // intermediate units owned by this tile - const int col0 = 2 * i0; // first gate_up column - const int ncol = 2 * nunit; // gate_up columns owned by this tile - const bf16_t* x_row = t->x + (size_t)tok * H; - const uint8_t* blk_e = gu_packed_l + (size_t)e * Hh * N2; - const uint8_t* scl_e = gu_scale_l + (size_t)e * (size_t)(H / 32) * N2; - float gu[2 * IBLK]; - mxgemv(gu, blk_e + col0, scl_e + col0, x_row, Hh, N2, ncol, e2m1_lut, e8m0_lut); - const bf16_t* bias_e = gu_bias_l + (size_t)e * N2 + col0; - bf16_t* g_row = g_scratch.data() + ((size_t)tok * top_k + k) * I; - const float lim = swiglu_limit, alpha = swiglu_alpha; - for (int j = 0; j < nunit; ++j) { - float gate = gu[2 * j] + bf16_to_f32(bias_e[2 * j]); - float up = gu[2 * j + 1] + bf16_to_f32(bias_e[2 * j + 1]); - if (gate > lim) gate = lim; - if (up > lim) up = lim; - else if (up < -lim) up = -lim; - const float glu = gate / (1.0f + std::exp(-gate * alpha)); // gate * sigmoid(alpha*gate) - g_row[i0 + j] = f32_to_bf16(glu * (up + 1.0f)); - } + const float w_in = apply_on_input ? io_w[static_cast(tok) * top_k + k] : 1.0f; + const int i0 = static_cast(ib) * kRowBlock; + const int i1 = std::min(I, i0 + kRowBlock); + format_of(t).gate_up(static_cast(t.layer), e, tok, i0, i1, io_x + (size_t)tok * H, w_in, + g_scratch.data() + ((size_t)tok * top_k + k) * I); } - void do_pass2_mxfp4(const MoeTask* t, int64_t p) { + // Pass 2 item p = one kRowBlock block of one token's output row, summed over its routes. + void do_pass2(const MoeTask& t, int64_t p) { + const int top_k = params.top_k, H = params.H, I = params.I; const int64_t hb = p % n_hblk; const int tok = static_cast(p / n_hblk); - const int Ih = I / 2; // packed-K rows - const int h0 = static_cast(hb) * HBLK; - const int h1 = std::min(H, h0 + HBLK); - const int nh = h1 - h0; - // Resolve this task's layer base once; row indexing below is layer-local (e). - const uint8_t* dn_packed_l = reinterpret_cast(tbl_at(down_tbl, t->layer_id)); - const uint8_t* dn_scale_l = reinterpret_cast(tbl_at(dn_scale_tbl, t->layer_id)); - const bf16_t* dn_bias_l = reinterpret_cast(tbl_at(dn_bias_tbl, t->layer_id)); - float acc[HBLK]; - for (int c = 0; c < nh; ++c) acc[c] = 0.0f; + const int h0 = static_cast(hb) * kRowBlock; + const int h1 = std::min(H, h0 + kRowBlock); + DownRoute routes[kMaxTopK]; + float w_out[kMaxTopK]; + int n = 0; for (int k = 0; k < top_k; ++k) { - const int e = t->ids[static_cast(tok) * top_k + k]; + const int e = io_ids[static_cast(tok) * top_k + k]; if (e < 0 || e >= num_experts) continue; - const float wt = t->w[static_cast(tok) * top_k + k]; - const bf16_t* g_row = g_scratch.data() + ((size_t)tok * top_k + k) * I; - const uint8_t* blk_e = dn_packed_l + (size_t)e * Ih * H; - const uint8_t* scl_e = dn_scale_l + (size_t)e * (size_t)(I / 32) * H; - float part[HBLK]; - mxgemv(part, blk_e + h0, scl_e + h0, g_row, Ih, H, nh, e2m1_lut, e8m0_lut); - const bf16_t* bias_e = dn_bias_l + (size_t)e * H + h0; - for (int c = 0; c < nh; ++c) acc[c] += (part[c] + bf16_to_f32(bias_e[c])) * wt; - } - bf16_t* y_row = t->y + (size_t)tok * H; - for (int c = 0; c < nh; ++c) y_row[h0 + c] = f32_to_bf16(acc[c]); - } - - // ----------------------------- ds_fp4 (DSV4) ------------------------------- - // Row-major e2m1 + e8m0/32 (no global); silu-swiglu with clamp; FP8-roundtripped - // activations (x once in submit -> xq_scratch; the intermediate g in a dedicated - // round-trip phase between the two passes). Router weight applies on the down output. - - void do_pass1_dsfp4(const MoeTask* t, int64_t p) { - const int64_t ib = p % n_iblk; - const int64_t tk = p / n_iblk; - const int k = static_cast(tk % top_k); - const int tok = static_cast(tk / top_k); - const int e = t->ids[static_cast(tok) * top_k + k]; - if (e < 0 || e >= num_experts) return; - // Resolve this task's layer base once; row indexing below is layer-local (e). - const uint8_t* gu_packed_l = reinterpret_cast(tbl_at(gate_up_tbl, t->layer_id)); - const uint8_t* gu_scale_l = reinterpret_cast(tbl_at(gu_scale_tbl, t->layer_id)); - const int N2 = 2 * I, Hh = H / 2, Hs = H / 32; - // fp8-roundtripped input, pre-deinterleaved to fp32 even/odd halves. - const float* xe = xe_scratch.data() + (size_t)tok * (H / 2); - const float* xo = xo_scratch.data() + (size_t)tok * (H / 2); - const uint8_t* gp = gu_packed_l + (size_t)e * N2 * Hh; - const uint8_t* gs = gu_scale_l + (size_t)e * N2 * Hs; - bf16_t* g_row = g_scratch.data() + ((size_t)tok * top_k + k) * I; - const int i0 = static_cast(ib) * IBLK; - const int i1 = std::min(I, i0 + IBLK); - const float lim = swiglu_limit; - for (int i = i0; i < i1; ++i) { - // gate_up is stored bf16 by the reference GEMV before swiglu; round to match. - float gate = bf16_to_f32(f32_to_bf16( - dsdot(gp + (size_t)i * Hh, gs + (size_t)i * Hs, xe, xo, H, e2m1_lut, e8m0_lut))); - float up = bf16_to_f32(f32_to_bf16(dsdot( - gp + (size_t)(I + i) * Hh, gs + (size_t)(I + i) * Hs, xe, xo, H, e2m1_lut, e8m0_lut))); - if (lim > 0.0f) { - if (gate > lim) gate = lim; - if (up > lim) up = lim; - else if (up < -lim) up = -lim; - } - const float glu = gate / (1.0f + std::exp(-gate)); // silu(gate) - g_row[i] = f32_to_bf16(glu * up); - } - } - - // Prepare one intermediate row (token,route) for the down GEMV: ds_fp4 first FP8 - // round-trips it (DSV4 act_quant), then both formats deinterleave to fp32 even/odd - // (reused across every down output row). - void prep_g_row(int64_t r) { - bf16_t* g = g_scratch.data() + (size_t)r * I; - if (use_q4a8) { // q4_0 W4A8: Q8_0-quantize the intermediate row for the down GEMV. - quant_q8_0(g, I, gi8_scratch.data() + (size_t)r * I, - gas_scratch.data() + (size_t)r * (I / 32)); - return; - } - if (fmt == WF_DSFP4) fp8_roundtrip_bf16(g, g, I); - float* ge = ge_scratch.data() + (size_t)r * (I / 2); - float* go = go_scratch.data() + (size_t)r * (I / 2); - deinterleave_bf16_f32(g, ge, go, I); - if (use_vnni) - quant_i8_pg16(ge, go, I, gi8_scratch.data() + (size_t)r * I, - gas_scratch.data() + (size_t)r * (I / 16)); - } - - void do_pass2_dsfp4(const MoeTask* t, int64_t p) { - const int64_t hb = p % n_hblk; - const int tok = static_cast(p / n_hblk); - const int Ih = I / 2, Is = I / 32; - const int h0 = static_cast(hb) * HBLK; - const int h1 = std::min(H, h0 + HBLK); - // Resolve this task's layer base once; row indexing below is layer-local (e). - const uint8_t* dn_packed_l = reinterpret_cast(tbl_at(down_tbl, t->layer_id)); - const uint8_t* dn_scale_l = reinterpret_cast(tbl_at(dn_scale_tbl, t->layer_id)); - bf16_t* y_row = t->y + (size_t)tok * H; + const size_t r = (size_t)tok * top_k + k; + routes[n] = {e, static_cast(r), g_scratch.data() + r * I}; + w_out[n] = apply_on_input ? 1.0f : io_w[r]; + ++n; + } + float part[kMaxTopK * kRowBlock]; + const MoeFormat& f = format_of(t); + if (n > 0) f.down(static_cast(t.layer), h0, h1, routes, n, part); + const bool rounds_routes = f.rounds_routes(); + bf16_t* y_row = io_y + (size_t)tok * H; for (int h = h0; h < h1; ++h) { float acc = 0.0f; - for (int k = 0; k < top_k; ++k) { - const int e = t->ids[static_cast(tok) * top_k + k]; - if (e < 0 || e >= num_experts) continue; - const float wt = t->w[static_cast(tok) * top_k + k]; - const float* ge = ge_scratch.data() + ((size_t)tok * top_k + k) * (I / 2); - const float* go = go_scratch.data() + ((size_t)tok * top_k + k) * (I / 2); - const uint8_t* dp = dn_packed_l + (size_t)e * (size_t)H * Ih + (size_t)h * Ih; - const uint8_t* ds = dn_scale_l + (size_t)e * (size_t)H * Is + (size_t)h * Is; - // The reference rounds each route's weighted down output to bf16 before the - // fp32 sum over routes (down [T, top_k, H] bf16 -> .sum(dim=1)). - acc += bf16_to_f32(f32_to_bf16(dsdot(dp, ds, ge, go, I, e2m1_lut, e8m0_lut) * wt)); + for (int j = 0; j < n; ++j) { + const float v = part[j * kRowBlock + h - h0] * w_out[j]; + acc += rounds_routes ? bf16_to_f32(f32_to_bf16(v)) : v; } y_row[h] = f32_to_bf16(acc); } } - void run_task_body(const MoeTask* t) { + void run_task_body(const MoeTask& t) { int local_sense = 0; for (;;) { int64_t p = p1_next.fetch_add(1, std::memory_order_relaxed); @@ -1848,14 +312,14 @@ struct CpuMoeExecutor { do_pass1(t, p); } barrier(local_sense); - // Row-major fp4: prepare the intermediate rows (per token,route) before the down - // GEMV -- ds_fp4 FP8 round-trips (DSV4 act_quant), both deinterleave to fp32; q4_0 - // W4A8 Q8_0-quantizes. Needs all of pass1 done (a full row spans every iblk). - if (needs_di || use_q4a8) { + // Prepare the intermediate rows (per token,route) for the down GEMV, see prep_inter. + // Needs all of pass1 done (a full row spans every iblk). + if (prt_total > 0) { + MoeFormat& f = *fmts[layer_fmt[t.layer]]; for (;;) { int64_t r = prt_next.fetch_add(1, std::memory_order_relaxed); if (r >= prt_total) break; - prep_g_row(r); + f.prep_inter(r, g_scratch.data() + (size_t)r * params.I); } barrier(local_sense); } @@ -1866,19 +330,33 @@ 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; } - run_task_body(t); + // submit() writes cur_task before the release store to `submitted`. + my_gen = gen; + run_task_body(cur_task); + last_active = wclk::now(); if (done_count.fetch_add(1) + 1 == num_threads) { completed.store(my_gen, std::memory_order_release); { @@ -1889,69 +367,21 @@ struct CpuMoeExecutor { } } - void submit(MoeTask* t) { - n_iblk = (I + IBLK - 1) / IBLK; - n_hblk = (H + HBLK - 1) / HBLK; - // Grow the per-token intermediate scratch if a larger batch shows up than the - // construction-time hint (CUDA-graph capture warms the largest bs first, so - // this happens at most once, before any capture, while the pool is idle). - const size_t need = static_cast(t->num_tokens) * top_k * I; - if (need > g_scratch.size()) g_scratch.resize(need); - p1_total = static_cast(t->num_tokens) * top_k * n_iblk; - p2_total = static_cast(t->num_tokens) * n_hblk; - prt_total = (needs_di || use_q4a8) ? static_cast(t->num_tokens) * top_k : 0; + void submit(MoeTask t) { + const int num_tokens = static_cast(t.num_tokens); + MoeFormat& fmt = *fmts[layer_fmt[t.layer]]; + p1_total = static_cast(num_tokens) * params.top_k * n_iblk; + p2_total = static_cast(num_tokens) * n_hblk; + prt_total = fmt.preps_inter() ? static_cast(num_tokens) * params.top_k : 0; p1_next.store(0, std::memory_order_relaxed); p2_next.store(0, std::memory_order_relaxed); prt_next.store(0, std::memory_order_relaxed); done_count.store(0, std::memory_order_relaxed); bar_count.store(0, std::memory_order_relaxed); bar_sense.store(0, std::memory_order_relaxed); - // ds_fp4: FP8 round-trip the per-token input once, up front (single-threaded; - // tiny for decode, and done before the workers are woken below). - if (needs_di) { - const size_t xn = static_cast(t->num_tokens) * H; - if (xn / 2 > xe_scratch.size()) { - xe_scratch.resize(xn / 2); - xo_scratch.resize(xn / 2); - ge_scratch.resize(static_cast(t->num_tokens) * top_k * (I / 2)); - go_scratch.resize(static_cast(t->num_tokens) * top_k * (I / 2)); - if (use_vnni) { - xi8_scratch.resize(xn); - xas_scratch.resize(xn / 16); - gi8_scratch.resize(static_cast(t->num_tokens) * top_k * I); - gas_scratch.resize(static_cast(t->num_tokens) * top_k * (I / 16)); - } - } - const bool ds = (fmt == WF_DSFP4) && !input_prequant; - if (ds && xn > xq_scratch.size()) xq_scratch.resize(xn); - for (int tok = 0; tok < t->num_tokens; ++tok) { - const bf16_t* src = t->x + (size_t)tok * H; - if (ds) { // DSV4 FP8-round-trips the input before the gate_up GEMV - bf16_t* xq = xq_scratch.data() + (size_t)tok * H; - fp8_roundtrip_bf16(src, xq, H); - src = xq; - } - float* xe = xe_scratch.data() + (size_t)tok * (H / 2); - float* xo = xo_scratch.data() + (size_t)tok * (H / 2); - deinterleave_bf16_f32(src, xe, xo, H); - if (use_vnni) - quant_i8_pg16(xe, xo, H, xi8_scratch.data() + (size_t)tok * H, - xas_scratch.data() + (size_t)tok * (H / 16)); - } - } - // q4_0 W4A8: Q8_0-quantize the per-token input once (single-threaded, tiny for decode). - if (use_q4a8) { - const size_t xn = static_cast(t->num_tokens) * H; - if (xn > xi8_scratch.size()) { - xi8_scratch.resize(xn); - xas_scratch.resize(xn / 32); - gi8_scratch.resize(static_cast(t->num_tokens) * top_k * I); - gas_scratch.resize(static_cast(t->num_tokens) * top_k * (I / 32)); - } - for (int tok = 0; tok < t->num_tokens; ++tok) - quant_q8_0(t->x + (size_t)tok * H, H, xi8_scratch.data() + (size_t)tok * H, - xas_scratch.data() + (size_t)tok * (H / 32)); - } + // Prepare the per-token input once, up front. + // Single-threaded: tiny for decode, and done before the workers are woken below. + for (int tok = 0; tok < num_tokens; ++tok) fmt.prep_input(tok, io_x + (size_t)tok * params.H); { std::lock_guard lk(task_mtx); cur_task = t; @@ -1961,67 +391,73 @@ 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; }); } - void submit_with_cuda_stream(uintptr_t stream, uintptr_t task) { - cudaLaunchHostFunc(reinterpret_cast(stream), &CpuMoeExecutor::submit_cb, - reinterpret_cast(task)); + void submit_with_cuda_stream(uintptr_t stream, int layer_id, int num_tokens) { + const MoeTask t = task(layer_id, num_tokens); + HostTask& h = host_tasks.try_emplace(to_word(t), HostTask{this, t}).first->second; + cudaLaunchHostFunc(reinterpret_cast(stream), &CpuMoeExecutor::submit_cb, &h); } - void sync_with_cuda_stream(uintptr_t stream, uintptr_t task) { - cudaLaunchHostFunc(reinterpret_cast(stream), &CpuMoeExecutor::sync_cb, - reinterpret_cast(task)); + void sync_with_cuda_stream(uintptr_t stream) { + cudaLaunchHostFunc(reinterpret_cast(stream), &CpuMoeExecutor::sync_cb, this); } - // Register a (layer, batch-size) slot's task so the coordinator can dispatch it on a - // flag bump. - void register_flag_task(int slot, uintptr_t task) { - std::lock_guard lk(flag_task_mtx); - if (static_cast(flag_task.size()) <= slot) flag_task.resize(slot + 1, nullptr); - flag_task[slot] = reinterpret_cast(task); - } - - int64_t flag_served_count(int slot) const { - return (slot >= 0 && slot < static_cast(flag_served.size())) ? flag_served[slot] : 0; - } - - // Start the busy-poll coordinator over the mapped-pinned flag arrays. ``pin_core`` >= 0 + // Start the busy-poll coordinator over the doorbell. ``pin_core`` >= 0 // pins the coordinator to that logical CPU (the worker auto-sizing reserves it), so its // polling never migrates onto / contends with a GEMV worker's core. - void start_flag_coordinator(uintptr_t ready_ptr, uintptr_t done_ptr, int num_slots, - int pin_core) { - ready_flags = reinterpret_cast(ready_ptr); - done_flags = reinterpret_cast(done_ptr); - coord_num_slots = num_slots; - { - std::lock_guard lk(flag_task_mtx); - if (static_cast(flag_task.size()) < num_slots) flag_task.resize(num_slots, nullptr); - } - flag_served.assign(num_slots, 0); + void start_flag_coordinator(int pin_core) { + gpu_sync = std::make_unique(); coord_stop.store(false); + coord_pinned = CPU_MOE_HAS_AFFINITY && pin_core >= 0; coord_thread = std::thread([this, pin_core] { -#if CPU_MOE_HAS_AFFINITY - if (pin_core >= 0) { - cpu_set_t set; - CPU_ZERO(&set); - CPU_SET(pin_core, &set); - pthread_setaffinity_np(pthread_self(), sizeof(set), &set); - } -#endif + if (pin_core >= 0) pin_this_thread(pin_core); coordinator_loop(); }); } + void flag_signal(uintptr_t stream, int layer_id, int num_tokens) { + gpu_sync->gpu_submit(stream, to_word(task(layer_id, num_tokens))); + } + void flag_wait(uintptr_t stream) { gpu_sync->gpu_wait(stream); } + void enable_kernel_io() { + kernel_io = std::make_unique(gpu_sync->ready_word(), gpu_sync->done_word(), io_x, io_ids, io_w, io_y, params.H, + params.top_k); + } + void flag_publish(uintptr_t stream, uintptr_t x, uintptr_t ids, uintptr_t w, int layer_id, int num_tokens) { + kernel_io->publish(stream, reinterpret_cast(x), reinterpret_cast(ids), + reinterpret_cast(w), num_tokens, to_word(task(layer_id, num_tokens))); + } + void flag_combine(uintptr_t stream, uintptr_t out, int num_tokens, bool accumulate) { + kernel_io->combine(stream, reinterpret_cast(out), num_tokens, accumulate); + } + bool flag_pending() const { return gpu_sync && gpu_sync->pending(); } + // Complete a doorbell the coordinator never answered so the stream can make progress. + void poison_flag() { + if (gpu_sync) gpu_sync->cpu_done(); + } + void coordinator_loop() { // Idle backoff (edge-device friendliness): spin hot only while decode traffic is - // flowing, with a TIME-based hot window (an iteration count varies 2-6x with slot - // count and pause cost, so some configs dozed every token). 50 ms since the last - // flag comfortably covers intra- and inter-token gaps -- including host-func-only - // stretches (prefill bursts, --moe-cpu-layers subsets) -- so steady decode never + // flowing, with a TIME-based hot window (an iteration count varies 2-6x with pause + // cost, so some configs dozed every token). 50 ms since the last + // flag comfortably covers intra- and inter-token gaps -- including prefill bursts and + // --moe-cpu-layers gaps -- so steady decode never // sleeps and keeps the ~us-level wakeup. Past it, i.e. the engine is actually idle, // escalate to timed sleeps capped at 2 ms: a dozing coordinator costs <0.1% of a // core instead of 100%, and the only price is a <=2 ms discovery delay on the FIRST @@ -2035,30 +471,19 @@ struct CpuMoeExecutor { int64_t sleep_us = 100; bool dozing = false; while (!coord_stop.load(std::memory_order_relaxed)) { - bool any = false; - for (int L = 0; L < coord_num_slots; ++L) { - // Binary handshake (memop-compatible: the GPU-side WAIT compares against an - // immediate baked at graph capture, so the protocol resets per step instead of - // counting). Acquire: everything the GPU made visible before setting ready -- - // the D2H input copies -- is visible to the worker pool after this read. - if (flag_load_acquire(&ready_flags[L]) != 0) { - flag_store_release(&ready_flags[L], 0); // consume this step's doorbell - MoeTask* t; - { - std::lock_guard lk(flag_task_mtx); - t = (L < static_cast(flag_task.size())) ? flag_task[L] : nullptr; - } - if (t != nullptr) { - submit(t); - sync(); - } - // Release: the workers' y stores are visible before the GPU sees done. - flag_store_release(&done_flags[L], 1); - if (L < static_cast(flag_served.size())) ++flag_served[L]; - any = true; + const uint64_t word = gpu_sync->cpu_poll(); + if (word != 0) { + const MoeTask t = from_word(word); + if (has_routes(t)) { + submit(t); + sync(coord_pinned && spin_wait.load(std::memory_order_relaxed)); + } else { + // What the pool would write for a task without routes, minus its fork and join. + std::memset(io_y, 0, sizeof(bf16_t) * t.num_tokens * params.H); } - } - if (any) { + // Release: the workers' y stores are visible before the GPU sees done. + gpu_sync->cpu_done(); + flag_served.fetch_add(1, std::memory_order_relaxed); last_active = coord_clock::now(); empty_polls = 0; sleep_us = 100; @@ -2067,15 +492,13 @@ struct CpuMoeExecutor { } // Hot phase: pause-spin, consulting the clock only every 1024 empty polls (the // amortization must gate the CLOCK, not the sleep -- gating the sleep decision - // let 1023/1024 idle iterations hot-scan all slots, ~20% of a core at 992 slots). + // let 1023/1024 idle iterations hot-poll). // Doze phase: sleep on EVERY iteration until traffic returns; one wake-up scan // (~0.5 us) per 2 ms sleep is ~0.03% duty. if (!dozing) { if ((++empty_polls & 1023u) != 0 || coord_clock::now() - last_active < kHotWindow) { -#if CPU_MOE_X86 - _mm_pause(); -#endif + cpu_relax(); continue; } dozing = true; @@ -2083,78 +506,81 @@ struct CpuMoeExecutor { std::this_thread::sleep_for(std::chrono::microseconds(sleep_us)); sleep_us = std::min(sleep_us * 2, kSleepCapUs); } - // Teardown: release any in-flight (or future) spin-wait immediately so a replay - // caught mid-shutdown exits its sync kernel now instead of owning the watchdog - // stall. Runs before the destructor's join() returns, while the flag arrays are - // still alive on the Python side. - for (int L = 0; L < coord_num_slots; ++L) { - flag_store_release(&done_flags[L], INT64_MAX); - } + // Teardown: complete any in-flight wait so a replay caught mid-shutdown does not stall + // until the watchdog fires. Runs before the destructor's join() returns, while the doorbell is still alive. + gpu_sync->cpu_done(); } // Eager (non-graph) path: run one task to completion on the pool. - void run_task(uintptr_t task) { - MoeTask* t = reinterpret_cast(task); - submit(t); + void run(int layer_id, int num_tokens) { + submit(task(layer_id, num_tokens)); sync(); } static void CUDART_CB submit_cb(void* ud) { - MoeTask* t = reinterpret_cast(ud); - t->exec->submit(t); - } - static void CUDART_CB sync_cb(void* ud) { - MoeTask* t = reinterpret_cast(ud); - t->exec->sync(); + HostTask* h = reinterpret_cast(ud); + h->exec->submit(h->task); } + static void CUDART_CB sync_cb(void* ud) { reinterpret_cast(ud)->sync(); } }; } // namespace +} // namespace freetoken::cpu_moe PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { namespace py = pybind11; + using freetoken::cpu_moe::CpuMoeExecutor; + using freetoken::cpu_moe::GpuSync; py::class_(m, "CpuMoeExecutor") - .def(py::init>(), + .def(py::init, + std::vector, double, double, std::vector, uintptr_t, uintptr_t, uintptr_t, + uintptr_t>(), py::arg("num_threads"), py::arg("num_layers"), py::arg("num_experts"), py::arg("top_k"), py::arg("hidden_size"), py::arg("inter_size"), py::arg("max_tokens"), py::arg("activation_id"), - py::arg("apply_router_weight_on_input"), py::arg("weight_format"), - py::arg("gate_up_ptr"), py::arg("down_ptr"), py::arg("gate_up_scale_ptr"), - py::arg("gate_up_global_ptr"), py::arg("down_scale_ptr"), - py::arg("down_global_ptr"), py::arg("gate_up_bias_ptr"), - py::arg("down_bias_ptr"), py::arg("swiglu_alpha"), py::arg("swiglu_limit"), - py::arg("core_ids")) - .def("create_task", &CpuMoeExecutor::create_task, py::arg("layer_id"), - py::arg("num_tokens"), py::arg("x_ptr"), py::arg("ids_ptr"), py::arg("w_ptr"), + py::arg("apply_router_weight_on_input"), py::arg("formats"), py::arg("layer_formats"), + py::arg("swiglu_alpha"), py::arg("swiglu_limit"), + py::arg("core_ids"), py::arg("x_ptr"), py::arg("ids_ptr"), py::arg("w_ptr"), py::arg("y_ptr")) .def("submit_with_cuda_stream", &CpuMoeExecutor::submit_with_cuda_stream, - py::arg("stream"), py::arg("task"), py::call_guard()) + py::arg("stream"), py::arg("layer_id"), py::arg("num_tokens"), + py::call_guard()) .def("sync_with_cuda_stream", &CpuMoeExecutor::sync_with_cuda_stream, - py::arg("stream"), py::arg("task"), py::call_guard()) - .def("run_task", &CpuMoeExecutor::run_task, py::arg("task"), + py::arg("stream"), py::call_guard()) + .def("run", &CpuMoeExecutor::run, py::arg("layer_id"), py::arg("num_tokens"), py::call_guard()) - .def("register_flag_task", &CpuMoeExecutor::register_flag_task, - py::arg("slot"), py::arg("task")) - .def("flag_served_count", &CpuMoeExecutor::flag_served_count, py::arg("slot")) - .def("start_flag_coordinator", &CpuMoeExecutor::start_flag_coordinator, - py::arg("ready_ptr"), py::arg("done_ptr"), py::arg("num_slots"), - py::arg("pin_core")) - .def("set_input_prequant", - [](CpuMoeExecutor& e, bool v) { e.input_prequant = v; }, - py::arg("value")) - .def("isa_name", &CpuMoeExecutor::isa_name); - m.def("memops_probe", &cumemops_probe, py::arg("stream"), py::arg("scratch_addr")); - m.def("memop_submit", &cumemop_submit, py::arg("stream"), py::arg("done_addr"), - py::arg("ready_addr"), py::arg("slot")); - m.def("memop_sync", &cumemop_sync, py::arg("stream"), py::arg("done_addr"), - py::arg("slot")); + .def("flag_served", [](const CpuMoeExecutor& e) { return e.flag_served.load(); }) + .def("start_flag_coordinator", &CpuMoeExecutor::start_flag_coordinator, py::arg("pin_core")) + .def("flag_signal", &CpuMoeExecutor::flag_signal, py::arg("stream"), py::arg("layer_id"), + py::arg("num_tokens")) + .def("flag_wait", &CpuMoeExecutor::flag_wait, py::arg("stream")) + .def("flag_pending", &CpuMoeExecutor::flag_pending) + .def("enable_kernel_io", &CpuMoeExecutor::enable_kernel_io) + .def("flag_publish", &CpuMoeExecutor::flag_publish, py::arg("stream"), py::arg("x"), py::arg("ids"), py::arg("w"), + py::arg("layer_id"), py::arg("num_tokens")) + .def("flag_combine", &CpuMoeExecutor::flag_combine, py::arg("stream"), py::arg("out"), py::arg("num_tokens"), + py::arg("accumulate")) + .def("poison_flag", &CpuMoeExecutor::poison_flag) + .def("isa_name", &CpuMoeExecutor::isa_name) + .def("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", &GpuSync::probe, py::arg("stream")); + m.def("kernel_io_probe", &freetoken::cpu_moe::KernelIo::probe, py::arg("stream")); + m.def("host_func_capture_safe", &GpuSync::host_func_capture_safe); // ABI capability marker: the highest ActKind this build implements in the // GENERIC epilogue. CpuMoeExecutor.__init__ probes it before requesting an act // id the epilogue must handle -- a prebuilt .so from before ACT_SWIGLUOAI // accepts id 3 without error and silently computes the wrong activation // (act_apply falls through to gelu_tanh); the probe turns a stale extension // into a loud rebuild instruction instead of wrong model outputs. - m.def("max_generic_act_id", []() { return static_cast(ACT_SWIGLU_CLAMP); }); + m.def("max_generic_act_id", []() { return static_cast(freetoken::cpu_moe::ACT_SWIGLU_CLAMP); }); + // The ggml-cpu variant the gguf format calls: the kernel table address that freetoken.kernel.ggml_cpu.select() returns. + m.def("set_ggml_cpu", &freetoken::cpu_moe::set_ggml_cpu, py::arg("api")); + m.def("weight_formats", []() { + py::list names; + for (const auto& f : freetoken::cpu_moe::kFormats) names.append(f.first); + return names; + }); } diff --git a/python/freetoken/kernel/csrc/cpu_moe/format.h b/python/freetoken/kernel/csrc/cpu_moe/format.h new file mode 100644 index 0000000000..286ae110f7 --- /dev/null +++ b/python/freetoken/kernel/csrc/cpu_moe/format.h @@ -0,0 +1,131 @@ +// The interface between the CPU MoE executor and one expert weight format. +#pragma once + +#include +#include +#include +#include + +#include "common.h" + +namespace freetoken::cpu_moe { + +// Output-row tiling. Small enough to give every worker independent work even at +// batch size 1; large enough to amortize the atomic work-grab. +// +// Bandwidth notes (Sapphire Rapids 8480+, 13 cores): the two passes already read +// every expert weight byte exactly once per token (each output row block is owned +// by one worker), and x stays hot in L1 across a (token,expert)'s rows -- so the +// kernel is single-read bandwidth-optimal at bs=1 (~205 GB/s vs ~55 GB/s PCIe). +// One worker per *physical* core, pinned, is the sweet spot; SMT oversubscription +// thrashes the spin-barrier. Deferred (not worth it here / for this workload): +// - AMX-bf16: a GEMM tile engine; decode is M=1 GEMV so tiles sit idle. It would +// only pay off in a grouped/batched (dedup) path. +// - expert dedup for bs>1: read each distinct expert once and GEMM its tokens. +// Helps locality+bytes when bs is large; decode batches here are tiny (<=4). +// - NUMA: a single node is assumed. Multi-socket machines would split each +// expert's K dimension per node (banks are already per-row contiguous). +constexpr int kRowBlock = 32; + +struct MoeParams { + int top_k = 0, max_tokens = 0; + int H = 0, I = 0; + int act = ACT_SILU; + float swiglu_alpha = 1.0f; + float swiglu_limit = INFINITY; // +inf == no clamp +}; + +// What cpu_executor.py passes for one format, by name: banks map canonical roles (gate_up, gate_up_scale, ...) to per-layer tables, options hold small integers. +// A format reads what it needs once, in its constructor. +struct FormatArgs { + std::string format; + std::map banks; + std::map options; + + const uint64_t* bank(const std::string& role) const { + const auto it = banks.find(role); + if (it == banks.end() || it->second == 0) + throw std::invalid_argument("CPU MoE " + format + " experts need the " + role + " bank"); + return reinterpret_cast(it->second); + } + int64_t option(const std::string& name) const { + const auto it = options.find(name); + if (it == options.end()) throw std::invalid_argument("CPU MoE " + format + " experts need the " + name + " option"); + return it->second; + } + int64_t option(const std::string& name, int64_t fallback) const { + const auto it = options.find(name); + return it == options.end() ? fallback : it->second; + } +}; + +// A bank table is the address of a CPU int64 array of length +// num_layers (one base address per layer, built by cpu_executor.py's +// _make_table), not a single flat bank. tbl_at resolves +// tbl[layer_id] once per task/pass. +inline const void* tbl_at(const uint64_t* tbl, int layer_id) { + return reinterpret_cast(tbl[layer_id]); +} + +// One route of a pass 2 item: expert e on intermediate row r = tok * top_k + k, whose bf16 values are g. +struct DownRoute { + int e; + int64_t r; + const bf16_t* g; +}; + +// g[i] for i in [i0, i1) from the gate row i and the up row I + i; dot(row) is the row's GEMV. +// ACT_SWIGLUOAI and ACT_SWIGLU_CLAMP clamp both inputs; the other kinds go through act_apply. +template +inline void gated_rows(const MoeParams& p, int i0, int i1, float w_in, bf16_t* g, RowDot dot) { + const bool clamped = p.act == ACT_SWIGLUOAI || p.act == ACT_SWIGLU_CLAMP; + const float up_bias = p.act == ACT_SWIGLUOAI ? 1.0f : 0.0f; + const float lim = p.swiglu_limit, alpha = p.swiglu_alpha; + for (int i = i0; i < i1; ++i) { + // gate = row i, up = row I+i + float gate = dot(i) * w_in; + float up = dot(p.I + i) * w_in; + if (clamped) { + // clamp(gate, max=lim) * sigmoid(alpha * gate) * (clamp(up, +-lim) + up_bias) + // -- swigluoai carries the +1 up bias (gpt-oss/MiniMax); swiglu_clamp + // (GLM-5.3) does not. lim == +inf: no clamp. + if (gate > lim) gate = lim; + if (up > lim) up = lim; + else if (up < -lim) up = -lim; + const float glu = gate / (1.0f + std::exp(-gate * alpha)); + g[i] = f32_to_bf16(glu * (up + up_bias)); + } else { + g[i] = f32_to_bf16(act_apply(p.act, gate) * up); + } + } +} + +// One expert weight format: its kernels, its activation layout and the scratch for it. +// The executor owns the threads and the work split. +// Scratch is sized once from max_tokens; gate_up, down and prep_inter run on the worker threads, prep_input before they wake. +class MoeFormat { + public: + explicit MoeFormat(const MoeParams& p) : p(p) {} + virtual ~MoeFormat() = default; + + virtual std::string isa_name() const = 0; + // Once per input row x[tok] before pass 1. + virtual void prep_input(int /*tok*/, const bf16_t* /*x*/) {} + // Whether pass 2 needs prep_inter on every intermediate row. + virtual bool preps_inter() const { return false; } + virtual void prep_inter(int64_t /*r*/, bf16_t* /*g*/) {} + // Pass 1: g[i] for i in [i0, i1) (at most kRowBlock rows), expert e of `layer` on token tok. + // w_in scales the gate and up dots: the router weight when it applies on the input, else 1. + virtual void gate_up(int layer, int e, int tok, int i0, int i1, const bf16_t* x, float w_in, + bf16_t* g) const = 0; + // Pass 2: out[j * kRowBlock + h - h0] for h in [h0, h1) (at most kRowBlock rows) and each of the n routes, before the router weight. + // Visiting the routes' rows h together keeps n weight streams in flight, which the down GEMV needs to reach DRAM bandwidth. + virtual void down(int layer, int h0, int h1, const DownRoute* routes, int n, float* out) const = 0; + // Round each router-weighted down output to bf16 before the sum over routes. + virtual bool rounds_routes() const { return false; } + + protected: + const MoeParams& p; +}; + +} // namespace freetoken::cpu_moe diff --git a/python/freetoken/kernel/csrc/cpu_moe/gpu_sync.h b/python/freetoken/kernel/csrc/cpu_moe/gpu_sync.h new file mode 100644 index 0000000000..829a8f34c3 --- /dev/null +++ b/python/freetoken/kernel/csrc/cpu_moe/gpu_sync.h @@ -0,0 +1,78 @@ +// GPU <-> CPU handshake for the CPU MoE coordinator, one implementation per GPU runtime (gpu_sync_cuda.cpp, gpu_sync_hip.cpp). +// +// A backend keeps two contracts, in eager mode and in graph replay: +// 1. submit, GPU -> CPU: after gpu_submit(stream, task), cpu_poll() returns task once, and every earlier write of the stream is visible to the CPU. +// 2. wait, CPU -> GPU: gpu_wait(stream) holds later work of the stream until cpu_done(), and that work sees every CPU write made before cpu_done(). +// task is a nonzero word, opaque here; the executor encodes the layer and token count in it, as a constant baked into each captured graph. +// One task is in flight at a time: decode runs on one stream and every layer waits before the next one signals. +// KernelIo (gpu_sync_io.h) may replace the copies around submit and wait: its kernel copies the inputs into the IO rows and does gpu_submit's writes itself. +// It skips a step without a CPU route: done = 1 and no ready, so cpu_poll() never sees the step and gpu_wait() passes at once. +#pragma once + +#include +#include + +#if defined(_MSC_VER) +#include +#endif + +namespace freetoken::cpu_moe { + +// Portable ordering for the flag handshake: "ready observed => the DMA'd inputs that +// preceded the bump are visible" and "y stores are visible before done". Plain +// volatile loads lean on x86 TSO; acquire/release makes it hold on aarch64 too +// (GH200/Jetson) at zero x86 cost. (MSVC branch is x86-only today: compiler barrier +// + TSO.) +inline int64_t flag_load_acquire(const volatile int64_t* p) { +#if defined(_MSC_VER) + const int64_t v = *p; + _ReadWriteBarrier(); + return v; +#else + return __atomic_load_n(const_cast(p), __ATOMIC_ACQUIRE); +#endif +} + +inline void flag_store_release(volatile int64_t* p, int64_t v) { +#if defined(_MSC_VER) + _ReadWriteBarrier(); + *p = v; +#else + __atomic_store_n(const_cast(p), v, __ATOMIC_RELEASE); +#endif +} + +class GpuSync { + public: + // Allocates the handshake memory; throws if this runtime cannot. + GpuSync(); + ~GpuSync(); + GpuSync(const GpuSync&) = delete; + GpuSync& operator=(const GpuSync&) = delete; + + // Whether both contracts hold on this device and stream, including graph replay where the runtime needs a real check. + static bool probe(uintptr_t stream); + // Whether the cudaLaunchHostFunc fallback is replayed correctly from a captured graph. + static bool host_func_capture_safe(); + + // 1. submit. gpu_submit() is enqueued on `stream`, which may be capturing; cpu_poll() returns 0 when nothing is submitted. + void gpu_submit(uintptr_t stream, uint64_t task); + uint64_t cpu_poll(); + + // 2. wait. cpu_done() also unblocks the stream on teardown and from the watchdog. + void gpu_wait(uintptr_t stream); + void cpu_done(); + + // Submitted but not polled yet (watchdog only). + bool pending() const; + + // Device addresses of the two words, for KernelIo. + uintptr_t ready_word() const; + uintptr_t done_word() const; + + private: + struct Impl; + std::unique_ptr impl_; +}; + +} // namespace freetoken::cpu_moe diff --git a/python/freetoken/kernel/csrc/cpu_moe/gpu_sync_cuda.cpp b/python/freetoken/kernel/csrc/cpu_moe/gpu_sync_cuda.cpp new file mode 100644 index 0000000000..a2a089bf67 --- /dev/null +++ b/python/freetoken/kernel/csrc/cpu_moe/gpu_sync_cuda.cpp @@ -0,0 +1,166 @@ +#include +#include +#include +#include +#include +#include + +#if defined(_WIN32) +#include +#else +#include +#endif + +#include + +#include "gpu_sync.h" + +namespace freetoken::cpu_moe { +namespace { + +// ===================================================================================== +// CUDA stream memory operations (driver API, resolved via dlopen -- no link-time or +// toolchain dependence). The GPU side of the flag handshake: submit = WRITE_VALUE +// (done=0 then ready=task word), sync = WAIT_VALUE(done >= 1). The wait is +// executed by the GPU front-end (no SM-resident kernel), so GPU "utilization" stays +// truthful during CPU compute windows -- a resident spin kernel pinned it at 99%, +// which laptop CPU/GPU dynamic power schedulers answered by clamping the CPU's max +// frequency (GEMV workers -1.5x: the reported edge regression). Availability is +// probed functionally at startup (memops_probe); anything unsupported (some driver models, +// vGPU, old drivers) falls back to the cudaLaunchHostFunc path. +#if defined(_WIN32) +static void* cumemop_dlopen() { return (void*)::LoadLibraryA("nvcuda.dll"); } +static void* cumemop_dlsym(void* h, const char* n) { + return (void*)::GetProcAddress((HMODULE)h, n); +} +#else +static void* cumemop_dlopen() { + void* h = dlopen("libcuda.so.1", RTLD_LAZY | RTLD_LOCAL); + if (h == nullptr) h = dlopen("libcuda.so", RTLD_LAZY | RTLD_LOCAL); + return h; +} +static void* cumemop_dlsym(void* h, const char* n) { return dlsym(h, n); } +#endif + +using cuMemOp64_fn = int (*)(void* stream, unsigned long long addr, unsigned long long value, + unsigned int flags); +static cuMemOp64_fn g_cu_write64 = nullptr; +static cuMemOp64_fn g_cu_wait64 = nullptr; +static constexpr unsigned int kCuWaitValueGeq = 0x0; // CU_STREAM_WAIT_VALUE_GEQ +static constexpr unsigned int kCuWriteDefault = 0x0; // CU_STREAM_WRITE_VALUE_DEFAULT + +static bool cumemop_resolve() { + static bool resolved = [] { + void* h = cumemop_dlopen(); + if (h == nullptr) return false; + // 11.7+ made the v2 entry points the default; older drivers export only the v1 + // names with the same signature. + g_cu_write64 = reinterpret_cast(cumemop_dlsym(h, "cuStreamWriteValue64_v2")); + if (g_cu_write64 == nullptr) + g_cu_write64 = reinterpret_cast(cumemop_dlsym(h, "cuStreamWriteValue64")); + g_cu_wait64 = reinterpret_cast(cumemop_dlsym(h, "cuStreamWaitValue64_v2")); + if (g_cu_wait64 == nullptr) + g_cu_wait64 = reinterpret_cast(cumemop_dlsym(h, "cuStreamWaitValue64")); + return g_cu_write64 != nullptr && g_cu_wait64 != nullptr; + }(); + return resolved; +} + +// GPU side of the flag handshake (see the block comment above): enqueued on the +// caller's (possibly capturing) stream; the WAIT immediate is the constant 1, +// replay-safe under CUDA graphs. +// The startup probe validates EAGER memops; a driver could still reject them at graph +// capture time. Those enqueue errors would otherwise be swallowed here and surface only +// as a later EndCapture failure -- log the first CUresult so triage is one step. +static void cumemop_check(int rc, const char* what) { + static std::atomic warned{false}; + if (rc != 0 && !warned.exchange(true)) { + std::fprintf(stderr, + "[freetoken/cpu_moe] %s failed with CUresult=%d (first occurrence; " + "subsequent errors are not repeated). If this happened during CUDA " + "graph capture, the driver lacks capture support for stream memops -- " + "set FREETOKEN_CPU_MOE_FLAG_SYNC=0.\n", + what, rc); + } +} + +} // namespace + +// ready and done live on separate cache lines of one mapped-pinned block; with UVA the GPU memops and the coordinator use the same addresses. +struct GpuSync::Impl { + void* block = nullptr; + void* device_block = nullptr; + volatile int64_t* ready = nullptr; + volatile int64_t* done = nullptr; +}; + +GpuSync::GpuSync() : impl_(std::make_unique()) { + if (!cumemop_resolve()) throw std::runtime_error("CPU MoE flag sync: CUDA stream memops are unavailable"); + const cudaError_t rc = cudaHostAlloc(&impl_->block, 128, cudaHostAllocPortable | cudaHostAllocMapped); + if (rc != cudaSuccess) + throw std::runtime_error(std::string("CPU MoE flag allocation failed: ") + cudaGetErrorString(rc)); + std::memset(impl_->block, 0, 128); + impl_->ready = static_cast(impl_->block); + impl_->done = impl_->ready + 8; + const cudaError_t dev_rc = cudaHostGetDevicePointer(&impl_->device_block, impl_->block, 0); + if (dev_rc != cudaSuccess) { + cudaFreeHost(impl_->block); + throw std::runtime_error(std::string("CPU MoE flag device address: ") + cudaGetErrorString(dev_rc)); + } +} + +GpuSync::~GpuSync() { cudaFreeHost(impl_->block); } + +// Functional probe on a scratch pinned int64: enqueue WRITE(7) + WAIT(>=7) + sync. +// Returns true only if the whole memop path works on THIS stream/device/driver. +bool GpuSync::probe(uintptr_t stream) { + if (!cumemop_resolve()) return false; + void* scratch = nullptr; + if (cudaHostAlloc(&scratch, sizeof(int64_t), cudaHostAllocPortable | cudaHostAllocMapped) != cudaSuccess) + return false; + *static_cast(scratch) = 0; + auto* s = reinterpret_cast(stream); + const auto addr = reinterpret_cast(scratch); + const bool ok = g_cu_write64(s, addr, 7ULL, kCuWriteDefault) == 0 && g_cu_wait64(s, addr, 7ULL, kCuWaitValueGeq) == 0 && + cudaStreamSynchronize(reinterpret_cast(stream)) == cudaSuccess; + // a failed enqueue can leave the WRITE pending on the scratch word + cudaStreamSynchronize(reinterpret_cast(stream)); + cudaFreeHost(scratch); + return ok; +} + +bool GpuSync::host_func_capture_safe() { return true; } + +void GpuSync::gpu_submit(uintptr_t stream, uint64_t task) { + auto* s = reinterpret_cast(stream); + // Order matters and is preserved by the front end: reset done BEFORE raising ready, + // so the coordinator's completion write for THIS step can never be wiped. + cumemop_check(g_cu_write64(s, reinterpret_cast(impl_->done), 0ULL, kCuWriteDefault), + "cuStreamWriteValue64(done)"); + cumemop_check(g_cu_write64(s, reinterpret_cast(impl_->ready), task, kCuWriteDefault), + "cuStreamWriteValue64(ready)"); +} + +void GpuSync::gpu_wait(uintptr_t stream) { + cumemop_check(g_cu_wait64(reinterpret_cast(stream), reinterpret_cast(impl_->done), 1ULL, + kCuWaitValueGeq), + "cuStreamWaitValue64(done)"); +} + +uint64_t GpuSync::cpu_poll() { + // Acquire: everything the GPU made visible before writing ready -- the D2H input copies -- is visible to the worker pool after this read. + const int64_t task = flag_load_acquire(impl_->ready); + if (task == 0) return 0; + flag_store_release(impl_->ready, 0); // consume this step's doorbell + return static_cast(task); +} + +void GpuSync::cpu_done() { flag_store_release(impl_->done, 1); } + +bool GpuSync::pending() const { return flag_load_acquire(impl_->ready) != 0 && flag_load_acquire(impl_->done) == 0; } + +uintptr_t GpuSync::ready_word() const { return reinterpret_cast(impl_->device_block); } + +uintptr_t GpuSync::done_word() const { return reinterpret_cast(impl_->device_block) + 8 * sizeof(int64_t); } + +} // namespace freetoken::cpu_moe diff --git a/python/freetoken/kernel/csrc/cpu_moe/gpu_sync_hip.cpp b/python/freetoken/kernel/csrc/cpu_moe/gpu_sync_hip.cpp new file mode 100644 index 0000000000..31ff9be539 --- /dev/null +++ b/python/freetoken/kernel/csrc/cpu_moe/gpu_sync_hip.cpp @@ -0,0 +1,275 @@ +// HIP stream memory operations for the CPU MoE flag handshake. +// +// HIP WAIT addresses must come from hipMallocSignalMemory, one 64-bit word per allocation. +// +// ROCm 10 records hipStreamWrite/WaitValue64 into the graph under stream capture, one node per call. +// ROCm 7.14 runs them while the stream is being captured but records nothing, so there we add explicit batch-memop nodes and splice them into the stream's dependency set. +// Which of the two applies is measured once at runtime (records_native_memops), not read from the version number. +// The probe then does a real capture + instantiate + replay through the chosen path. +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +#include "gpu_sync.h" + +namespace freetoken::cpu_moe { + +namespace { + +void hip_check(hipError_t rc, const char* what) { + if (rc != hipSuccess) + throw std::runtime_error(std::string("CPU MoE HIP flag sync: ") + what + " failed: " + hipGetErrorString(rc)); +} + +// Capability failures must not poison eager fallback or hide a caller's earlier error. +struct ProbeErrorScope { + ProbeErrorScope() { hip_check(hipPeekAtLastError(), "pending error before probe"); } + ~ProbeErrorScope() { (void)hipGetLastError(); } +}; + +hipError_t alloc_signal(uint64_t** out) { + void* storage = nullptr; + const hipError_t rc = hipExtMallocWithFlags(&storage, sizeof(uint64_t), hipMallocSignalMemory); + if (rc == hipSuccess) { + *out = static_cast(storage); + __atomic_store_n(*out, 0ULL, __ATOMIC_RELEASE); + } + return rc; +} + +void free_signal(uint64_t* signal) { + if (signal != nullptr) (void)hipFree(signal); +} + +volatile int64_t* as_flag(uint64_t* signal) { return reinterpret_cast(signal); } + +bool capturing(hipStream_t stream) { + hipStreamCaptureStatus status = hipStreamCaptureStatusNone; + hip_check(hipStreamIsCapturing(stream, &status), "capture status query"); + return status == hipStreamCaptureStatusActive; +} + +// Whether this runtime records a native hipStreamWriteValue64 issued under capture as a graph node. +bool records_native_memops() { + static const bool native = [] { + ProbeErrorScope errors; + uint64_t* word = nullptr; + hipStream_t stream = nullptr; + hipGraph_t graph = nullptr; + hipError_t rc = alloc_signal(&word); + if (rc == hipSuccess) rc = hipStreamCreateWithFlags(&stream, hipStreamNonBlocking); + if (rc == hipSuccess) rc = hipStreamBeginCapture(stream, hipStreamCaptureModeRelaxed); + const bool capture_started = rc == hipSuccess; + if (rc == hipSuccess) rc = hipStreamWriteValue64(stream, word, 1, hipStreamWriteValueDefault); + if (capture_started) { + const hipError_t end_rc = hipStreamEndCapture(stream, &graph); + if (rc == hipSuccess) rc = end_rc; + } + size_t nodes = 0; + if (rc == hipSuccess) rc = hipGraphGetNodes(graph, nullptr, &nodes); + if (graph != nullptr) (void)hipGraphDestroy(graph); + if (stream != nullptr) { + (void)hipStreamSynchronize(stream); + (void)hipStreamDestroy(stream); + } + free_signal(word); + return rc == hipSuccess && nodes == 1; + }(); + return native; +} + +// Appends one batch-memop node to the capture of `stream`; the node keeps a pointer to `ops`. +hipError_t add_capture_memop_node(hipStream_t stream, hipStreamBatchMemOpParams* ops, unsigned int count) { + hipStreamCaptureStatus status = hipStreamCaptureStatusNone; + hipGraph_t graph = nullptr; + const hipGraphNode_t* dependencies = nullptr; + size_t dependency_count = 0; + hipError_t rc = hipStreamGetCaptureInfo_v2(stream, &status, nullptr, &graph, &dependencies, &dependency_count); + if (rc != hipSuccess) return rc; + if (status != hipStreamCaptureStatusActive || graph == nullptr) return hipErrorIllegalState; + + hipCtx_t context = nullptr; +#if defined(__GNUC__) +#pragma GCC diagnostic push +#pragma GCC diagnostic ignored "-Wdeprecated-declarations" +#endif + rc = hipCtxGetCurrent(&context); +#if defined(__GNUC__) +#pragma GCC diagnostic pop +#endif + if (rc != hipSuccess) return rc; + + hipBatchMemOpNodeParams params{}; + params.ctx = context; + params.count = count; + params.paramArray = ops; + hipGraphNode_t node = nullptr; + rc = hipGraphAddBatchMemOpNode(&node, graph, dependencies, dependency_count, ¶ms); + if (rc != hipSuccess) return rc; + rc = hipStreamUpdateCaptureDependencies(stream, &node, 1, hipStreamSetCaptureDependencies); + // EndCapture rejects a node that was added but never joined to the stream. + if (rc != hipSuccess && node != nullptr) (void)hipGraphDestroyNode(node); + return rc; +} + +hipStreamBatchMemOpParams write_op(uint64_t* address, uint64_t value) { + hipStreamBatchMemOpParams op{}; + op.writeValue.operation = hipStreamMemOpWriteValue64; + op.writeValue.address = reinterpret_cast(address); + op.writeValue.value64 = value; + return op; +} + +hipStreamBatchMemOpParams wait_op(uint64_t* address) { + hipStreamBatchMemOpParams op{}; + op.waitValue.operation = hipStreamMemOpWaitValue64; + op.waitValue.address = reinterpret_cast(address); + op.waitValue.value64 = 1; + op.waitValue.flags = hipStreamWaitValueGte; + return op; +} + +bool needs_explicit_nodes(hipStream_t stream) { return capturing(stream) && !records_native_memops(); } + +// node_op is null for a native call, else the parameter slot of an explicit node, which the node keeps a pointer to. +// The ops of one batch node run in parallel, so each write gets its own node and consecutive writes stay in order. +hipError_t write_value(hipStream_t stream, uint64_t* address, uint64_t value, hipStreamBatchMemOpParams* node_op) { + if (node_op == nullptr) return hipStreamWriteValue64(stream, address, value, hipStreamWriteValueDefault); + *node_op = write_op(address, value); + return add_capture_memop_node(stream, node_op, 1); +} + +hipError_t wait_done(hipStream_t stream, uint64_t* done, hipStreamBatchMemOpParams* node_op) { + if (node_op == nullptr) return hipStreamWaitValue64(stream, done, 1, hipStreamWaitValueGte, UINT64_MAX); + *node_op = wait_op(done); + return add_capture_memop_node(stream, node_op, 1); +} + +} // namespace + +struct GpuSync::Impl { + uint64_t* ready = nullptr; // signal memory + uint64_t* done = nullptr; + // Explicit nodes keep these addresses, so one parameter pair per distinct task, never moved (std::map nodes are stable). + std::map> submit_ops; // done = 0, ready = task + hipStreamBatchMemOpParams wait_op{}; +}; + +GpuSync::GpuSync() : impl_(std::make_unique()) { + hipError_t rc = alloc_signal(&impl_->ready); + if (rc == hipSuccess) rc = alloc_signal(&impl_->done); + if (rc != hipSuccess) { + free_signal(impl_->ready); + free_signal(impl_->done); + hip_check(rc, "signal allocation"); + } + // Measure now: the measurement captures and allocates, which a graph capture in progress would reject. + (void)records_native_memops(); +} + +GpuSync::~GpuSync() { + free_signal(impl_->ready); + free_signal(impl_->done); +} + +bool GpuSync::probe(uintptr_t /*stream*/) { + ProbeErrorScope errors; + int device = 0; + int can_wait = 0; + if (hipGetDevice(&device) != hipSuccess || + hipDeviceGetAttribute(&can_wait, hipDeviceAttributeCanUseStreamWaitValue, device) != hipSuccess || can_wait == 0) + return false; + const bool explicit_nodes = !records_native_memops(); + + uint64_t* ready = nullptr; + uint64_t* done = nullptr; + hipStream_t stream = nullptr; + hipGraph_t graph = nullptr; + hipGraphExec_t graph_exec = nullptr; + std::array node_ops{}; + hipError_t rc = alloc_signal(&ready); + if (rc == hipSuccess) rc = alloc_signal(&done); + if (rc == hipSuccess) rc = hipStreamCreateWithFlags(&stream, hipStreamNonBlocking); + // A runtime that runs the WAIT during capture instead of recording it must not block here. + if (rc == hipSuccess) __atomic_store_n(done, 1ULL, __ATOMIC_RELEASE); + if (rc == hipSuccess) rc = hipStreamBeginCapture(stream, hipStreamCaptureModeRelaxed); + const bool capture_started = rc == hipSuccess; + if (rc == hipSuccess) rc = write_value(stream, ready, 1, explicit_nodes ? &node_ops[0] : nullptr); + if (rc == hipSuccess) rc = wait_done(stream, done, explicit_nodes ? &node_ops[1] : nullptr); + // End the capture even on failure so its stream and graph are not left active during cleanup. + if (capture_started) { + const hipError_t end_rc = hipStreamEndCapture(stream, &graph); + if (rc == hipSuccess) rc = end_rc; + } + if (rc == hipSuccess) rc = hipStreamSynchronize(stream); + size_t nodes = 0; + if (rc == hipSuccess) rc = hipGraphGetNodes(graph, nullptr, &nodes); + if (rc == hipSuccess && nodes != 2) rc = hipErrorNotSupported; + if (rc == hipSuccess) rc = hipGraphInstantiate(&graph_exec, graph, nullptr, nullptr, 0); + + bool served = false; + if (rc == hipSuccess) { + __atomic_store_n(ready, 0ULL, __ATOMIC_RELEASE); + __atomic_store_n(done, 0ULL, __ATOMIC_RELEASE); + std::thread responder([&] { + const auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(2); + while (__atomic_load_n(ready, __ATOMIC_ACQUIRE) == 0 && std::chrono::steady_clock::now() < deadline) + std::this_thread::yield(); + served = __atomic_load_n(ready, __ATOMIC_ACQUIRE) == 1; + // The probe graph never resets done, so this also releases a launch that arrives after the deadline. + __atomic_store_n(done, 1ULL, __ATOMIC_RELEASE); + }); + rc = hipGraphLaunch(graph_exec, stream); + if (rc == hipSuccess) rc = hipStreamSynchronize(stream); + responder.join(); + } + const bool ok = rc == hipSuccess && served; + + if (graph_exec != nullptr) (void)hipGraphExecDestroy(graph_exec); + if (graph != nullptr) (void)hipGraphDestroy(graph); + if (stream != nullptr) (void)hipStreamDestroy(stream); + free_signal(ready); + free_signal(done); + return ok; +} + +void GpuSync::gpu_submit(uintptr_t stream, uint64_t task) { + auto s = reinterpret_cast(stream); + hipStreamBatchMemOpParams* ops = needs_explicit_nodes(s) ? impl_->submit_ops[task].data() : nullptr; + // Reset done before raising ready, so the coordinator's completion for this step can never be wiped. + hip_check(write_value(s, impl_->done, 0, ops), "write done"); + hip_check(write_value(s, impl_->ready, task, ops ? ops + 1 : nullptr), "write ready"); +} + +void GpuSync::gpu_wait(uintptr_t stream) { + auto s = reinterpret_cast(stream); + hip_check(wait_done(s, impl_->done, needs_explicit_nodes(s) ? &impl_->wait_op : nullptr), "wait done"); +} + +bool GpuSync::host_func_capture_safe() { return false; } + +uint64_t GpuSync::cpu_poll() { + const int64_t task = flag_load_acquire(as_flag(impl_->ready)); + if (task == 0) return 0; + flag_store_release(as_flag(impl_->ready), 0); + return static_cast(task); +} + +void GpuSync::cpu_done() { flag_store_release(as_flag(impl_->done), 1); } + +bool GpuSync::pending() const { + return flag_load_acquire(as_flag(impl_->ready)) != 0 && flag_load_acquire(as_flag(impl_->done)) == 0; +} + +uintptr_t GpuSync::ready_word() const { return reinterpret_cast(impl_->ready); } + +uintptr_t GpuSync::done_word() const { return reinterpret_cast(impl_->done); } + +} // namespace freetoken::cpu_moe diff --git a/python/freetoken/kernel/csrc/cpu_moe/gpu_sync_io.cu b/python/freetoken/kernel/csrc/cpu_moe/gpu_sync_io.cu new file mode 100644 index 0000000000..b41961f0bf --- /dev/null +++ b/python/freetoken/kernel/csrc/cpu_moe/gpu_sync_io.cu @@ -0,0 +1,199 @@ +#include +#include +#include + +#include + +#include "gpu_sync_io.h" + +namespace freetoken::cpu_moe { +namespace { + +__device__ inline float bf16_bits_to_float(uint16_t b) { return __uint_as_float(static_cast(b) << 16); } + +// Round to nearest even, the same rounding as a torch bf16 add. +__device__ inline uint16_t float_to_bf16_bits(float f) { + const uint32_t u = __float_as_uint(f); + if ((u & 0x7fffffffu) > 0x7f800000u) return 0x7fc0; + return static_cast((u + 0x7fffu + ((u >> 16) & 1u)) >> 16); +} + +// One block. Copies x/ids/w into the IO rows, then done = 0 and ready = task. +// A step without a CPU route writes done = 1 and rings nothing: the coordinator stays asleep and the GPU wait passes at once. +__global__ void publish_kernel(const uint16_t* __restrict__ x, const int32_t* __restrict__ ids, const float* __restrict__ w, + int64_t n_x, int64_t n_ids, uint16_t* io_x, int32_t* io_ids, float* io_w, + volatile int64_t* ready, volatile int64_t* done, int64_t task, int32_t* has_cpu) { + __shared__ int any; + if (threadIdx.x == 0) any = 0; + __syncthreads(); + for (int64_t i = threadIdx.x; i < n_ids; i += blockDim.x) + if (ids[i] >= 0) any = 1; + __syncthreads(); + if (!any) { + if (threadIdx.x == 0) { + *has_cpu = 0; + *done = 1; + __threadfence_system(); + } + return; + } + const bool vec = (reinterpret_cast(x) | reinterpret_cast(io_x)) % 16 == 0 && n_x % 8 == 0; + if (vec) { + for (int64_t i = threadIdx.x; i < n_x / 8; i += blockDim.x) + reinterpret_cast(io_x)[i] = reinterpret_cast(x)[i]; + } else { + for (int64_t i = threadIdx.x; i < n_x; i += blockDim.x) io_x[i] = x[i]; + } + for (int64_t i = threadIdx.x; i < n_ids; i += blockDim.x) { + io_ids[i] = ids[i]; + io_w[i] = w[i]; + } + __threadfence_system(); + __syncthreads(); + if (threadIdx.x == 0) { + *has_cpu = 1; + *done = 0; + __threadfence_system(); + *ready = task; + __threadfence_system(); + } +} + +// out = CPU rows, or out += CPU rows with accumulate; a step publish skipped contributes zeros. +// io_y is volatile: the same host rows hold a new step on every call. +__global__ void combine_kernel(uint16_t* out, const volatile uint16_t* io_y, const int32_t* has_cpu, int64_t n, bool accumulate) { + const bool on = *has_cpu != 0; + if (accumulate && !on) return; + for (int64_t i = blockIdx.x * static_cast(blockDim.x) + threadIdx.x; i < n; i += static_cast(gridDim.x) * blockDim.x) { + float v = accumulate ? bf16_bits_to_float(out[i]) : 0.0f; + if (on) v += bf16_bits_to_float(io_y[i]); + out[i] = float_to_bf16_bits(v); + } +} + +cudaError_t launch_publish(cudaStream_t stream, const uint16_t* x, const int32_t* ids, const float* w, int64_t n_x, int64_t n_ids, + uint16_t* io_x, int32_t* io_ids, float* io_w, volatile int64_t* ready, volatile int64_t* done, int64_t task, + int32_t* has_cpu) { + publish_kernel<<<1, 256, 0, stream>>>(x, ids, w, n_x, n_ids, io_x, io_ids, io_w, ready, done, task, has_cpu); + return cudaGetLastError(); +} + +cudaError_t launch_combine(cudaStream_t stream, uint16_t* out, const uint16_t* io_y, const int32_t* has_cpu, int64_t n, bool accumulate) { + if (n == 0) return cudaSuccess; + const auto blocks = static_cast(std::min((n + 255) / 256, 1024)); + combine_kernel<<>>(out, io_y, has_cpu, n, accumulate); + return cudaGetLastError(); +} + +void check(cudaError_t rc, const char* what) { + if (rc != cudaSuccess) throw std::runtime_error(std::string("CPU MoE kernel IO: ") + what + ": " + cudaGetErrorString(rc)); +} + +template +T* device_address(const void* host) { + void* dev = nullptr; + check(cudaHostGetDevicePointer(&dev, const_cast(host), 0), "device address of the IO rows"); + return static_cast(dev); +} + +} // namespace + +KernelIo::KernelIo(uintptr_t ready, uintptr_t done, const void* x, const int32_t* ids, const float* w, const void* y, int hidden, + int top_k) + : ready_(reinterpret_cast(ready)), + done_(reinterpret_cast(done)), + x_(device_address(x)), + ids_(device_address(ids)), + w_(device_address(w)), + y_(device_address(y)), + hidden_(hidden), + top_k_(top_k) { + check(cudaMalloc(reinterpret_cast(&has_cpu_), sizeof(int32_t)), "has_cpu allocation"); + const cudaError_t rc = cudaMemset(has_cpu_, 0, sizeof(int32_t)); + if (rc != cudaSuccess) { + (void)cudaFree(has_cpu_); + check(rc, "has_cpu reset"); + } +} + +KernelIo::~KernelIo() { (void)cudaFree(has_cpu_); } + +void KernelIo::publish(uintptr_t stream, const void* x, const int32_t* ids, const float* w, int num_tokens, uint64_t task) { + check(launch_publish(reinterpret_cast(stream), static_cast(x), ids, w, + static_cast(num_tokens) * hidden_, static_cast(num_tokens) * top_k_, x_, ids_, w_, ready_, done_, + static_cast(task), has_cpu_), + "publish launch"); +} + +void KernelIo::combine(uintptr_t stream, void* out, int num_tokens, bool accumulate) { + check(launch_combine(reinterpret_cast(stream), static_cast(out), y_, has_cpu_, + static_cast(num_tokens) * hidden_, accumulate), + "combine launch"); +} + +// A ringing step, its combine and a skipped step on scratch mapped memory, checking every word they write. +std::string KernelIo::probe(uintptr_t stream_handle) { + struct alignas(16) Area { + int64_t ready, done; + int32_t has_cpu, pad[3]; + uint16_t x[8], io_x[8], io_y[8], out[8]; + int32_t ids[2], io_ids[2]; + float w[2], io_w[2]; + }; + const auto stream = reinterpret_cast(stream_handle); + void* host = nullptr; + void* dev = nullptr; + if (cudaHostAlloc(&host, sizeof(Area), cudaHostAllocPortable | cudaHostAllocMapped) != cudaSuccess) + return "pinned scratch allocation failed"; + std::string reason; + auto ok = [&](cudaError_t rc, const char* what) { + if (rc != cudaSuccess && reason.empty()) reason = std::string(what) + ": " + cudaGetErrorString(rc); + return rc == cudaSuccess; + }; + auto fail = [&](const char* why) { + if (reason.empty()) reason = why; + }; + if (ok(cudaHostGetDevicePointer(&dev, host, 0), "device address of the scratch")) { + auto* h = static_cast(host); + auto* d = static_cast(dev); + std::memset(h, 0, sizeof(Area)); + for (int i = 0; i < 8; ++i) { + h->x[i] = static_cast(0x3f80 + i); + h->io_y[i] = 0x3f80; // 1.0 + h->out[i] = 0x4000; // 2.0 + } + h->ids[0] = 3; + h->ids[1] = -1; + h->w[0] = 0.25f; + h->w[1] = 0.75f; + auto publish = [&] { + return launch_publish(stream, d->x, d->ids, d->w, 8, 2, d->io_x, d->io_ids, d->io_w, &d->ready, &d->done, 5, &d->has_cpu); + }; + + if (ok(publish(), "publish launch") && ok(cudaStreamSynchronize(stream), "publish")) { + if (h->ready != 5 || h->done != 0 || h->has_cpu != 1) fail("publish wrote wrong handshake words"); + if (std::memcmp(h->io_x, h->x, sizeof(h->x)) || std::memcmp(h->io_ids, h->ids, sizeof(h->ids)) || + std::memcmp(h->io_w, h->w, sizeof(h->w))) + fail("publish wrote wrong IO rows"); + } + if (reason.empty() && ok(launch_combine(stream, d->out, d->io_y, &d->has_cpu, 8, true), "combine launch") && + ok(cudaStreamSynchronize(stream), "combine")) { + for (int i = 0; i < 8; ++i) + if (h->out[i] != 0x4040) fail("combine returned a wrong sum"); // 3.0 + } + if (reason.empty()) { + h->ready = 0; + h->ids[0] = -1; + if (ok(publish(), "skip launch") && ok(launch_combine(stream, d->out, d->io_y, &d->has_cpu, 8, false), "skip combine launch") && + ok(cudaStreamSynchronize(stream), "skipped step")) { + bool zeros = true; + for (int i = 0; i < 8; ++i) zeros = zeros && h->out[i] == 0; + if (h->ready != 0 || h->done != 1 || h->has_cpu != 0 || !zeros) fail("a skipped step wrote wrong words"); + } + } + } + (void)cudaFreeHost(host); + return reason; +} + +} // namespace freetoken::cpu_moe diff --git a/python/freetoken/kernel/csrc/cpu_moe/gpu_sync_io.h b/python/freetoken/kernel/csrc/cpu_moe/gpu_sync_io.h new file mode 100644 index 0000000000..6eacb3936b --- /dev/null +++ b/python/freetoken/kernel/csrc/cpu_moe/gpu_sync_io.h @@ -0,0 +1,39 @@ +// Kernel IO for the GPU <-> CPU handshake (gpu_sync.h): kernels on the executor's mapped pinned IO rows in place of the copies around submit and wait. +// One source for both runtimes: nvcc builds gpu_sync_io.cu on CUDA, hipcc on ROCm through freetoken/hip_compat.h. +#pragma once + +#include +#include + +namespace freetoken::cpu_moe { + +class KernelIo { + public: + // ready, done: device addresses of GpuSync's words. x, ids, w, y: host addresses of the pinned IO rows. Throws if setup fails. + KernelIo(uintptr_t ready, uintptr_t done, const void* x, const int32_t* ids, const float* w, const void* y, int hidden, int top_k); + ~KernelIo(); + KernelIo(const KernelIo&) = delete; + KernelIo& operator=(const KernelIo&) = delete; + + // Copies the step's device inputs into the IO rows and does gpu_submit's writes, or skips the step when no route goes to the CPU. + // x [n, hidden] bf16, ids [n, top_k] int32 (< 0 = not a CPU route), w [n, top_k] fp32. + void publish(uintptr_t stream, const void* x, const int32_t* ids, const float* w, int num_tokens, uint64_t task); + // Follows gpu_wait(): writes the CPU rows to out [n, hidden] bf16, or adds them to it; a skipped step gives zeros. + void combine(uintptr_t stream, void* out, int num_tokens, bool accumulate); + + // Whether both kernels run and write the right words on this device: "" or the reason. + static std::string probe(uintptr_t stream); + + private: + volatile int64_t* ready_; + volatile int64_t* done_; + uint16_t* x_; + int32_t* ids_; + float* w_; + const uint16_t* y_; + int hidden_; + int top_k_; + int32_t* has_cpu_ = nullptr; // device word: whether the last published step has CPU routes +}; + +} // namespace freetoken::cpu_moe diff --git a/python/freetoken/kernel/csrc/cpu_moe/kernels/bf16.cpp b/python/freetoken/kernel/csrc/cpu_moe/kernels/bf16.cpp new file mode 100644 index 0000000000..273c394609 --- /dev/null +++ b/python/freetoken/kernel/csrc/cpu_moe/kernels/bf16.cpp @@ -0,0 +1,159 @@ +#include "../format.h" + +namespace freetoken::cpu_moe { +namespace { + +// ------------------------------- dot products ------------------------------- +// dot(weight[bf16], act[bf16], n) -> fp32. The selected impl is a function +// pointer chosen at runtime; the per-row call overhead is negligible vs n. + +using dot_fn = float (*)(const bf16_t*, const bf16_t*, int); + +float dot_scalar(const bf16_t* w, const bf16_t* x, int n) { + float acc = 0.0f; + for (int i = 0; i < n; ++i) acc += bf16_to_f32(w[i]) * bf16_to_f32(x[i]); + return acc; +} + +#if CPU_MOE_X86 +CPU_MOE_TARGET("avx512f") +float dot_avx512f(const bf16_t* w, const bf16_t* x, int n) { + // 4 independent accumulators -> more in-flight loads (memory-level parallelism), + // which is what lifts a bandwidth-bound GEMV toward peak. + __m512 a0 = _mm512_setzero_ps(), a1 = _mm512_setzero_ps(); + __m512 a2 = _mm512_setzero_ps(), a3 = _mm512_setzero_ps(); + int i = 0; + for (; i + 64 <= n; i += 64) { + _mm_prefetch(reinterpret_cast(w + i) + PF_AHEAD, _MM_HINT_T0); + for (int j = 0; j < 64; j += 16) { + __m256i wi = _mm256_loadu_si256(reinterpret_cast(w + i + j)); + __m256i xi = _mm256_loadu_si256(reinterpret_cast(x + i + j)); + __m512 wf = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(wi), 16)); + __m512 xf = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(xi), 16)); + __m512& acc = (j == 0) ? a0 : (j == 16) ? a1 : (j == 32) ? a2 : a3; + acc = _mm512_fmadd_ps(wf, xf, acc); + } + } + for (; i + 16 <= n; i += 16) { + __m256i wi = _mm256_loadu_si256(reinterpret_cast(w + i)); + __m256i xi = _mm256_loadu_si256(reinterpret_cast(x + i)); + __m512 wf = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(wi), 16)); + __m512 xf = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(xi), 16)); + a0 = _mm512_fmadd_ps(wf, xf, a0); + } + float s = _mm512_reduce_add_ps(_mm512_add_ps(_mm512_add_ps(a0, a1), _mm512_add_ps(a2, a3))); + for (; i < n; ++i) s += bf16_to_f32(w[i]) * bf16_to_f32(x[i]); + return s; +} + +#ifdef CPU_MOE_HAS_AVX512BF16 +CPU_MOE_TARGET("avx512bf16,avx512f") +static inline __m512bh load_bh(const bf16_t* p) { + __m512i raw = _mm512_loadu_si512(reinterpret_cast(p)); + __m512bh out; + std::memcpy(&out, &raw, sizeof(out)); + return out; +} + +CPU_MOE_TARGET("avx512bf16,avx512f") +float dot_avx512bf16(const bf16_t* w, const bf16_t* x, int n) { + // 4 accumulators (128 bf16/iter) for memory-level parallelism + a prefetch nudge. + __m512 a0 = _mm512_setzero_ps(), a1 = _mm512_setzero_ps(); + __m512 a2 = _mm512_setzero_ps(), a3 = _mm512_setzero_ps(); + int i = 0; + for (; i + 128 <= n; i += 128) { + _mm_prefetch(reinterpret_cast(w + i) + PF_AHEAD, _MM_HINT_T0); + a0 = _mm512_dpbf16_ps(a0, load_bh(w + i), load_bh(x + i)); + a1 = _mm512_dpbf16_ps(a1, load_bh(w + i + 32), load_bh(x + i + 32)); + a2 = _mm512_dpbf16_ps(a2, load_bh(w + i + 64), load_bh(x + i + 64)); + a3 = _mm512_dpbf16_ps(a3, load_bh(w + i + 96), load_bh(x + i + 96)); + } + for (; i + 32 <= n; i += 32) { + a0 = _mm512_dpbf16_ps(a0, load_bh(w + i), load_bh(x + i)); + } + float s = _mm512_reduce_add_ps(_mm512_add_ps(_mm512_add_ps(a0, a1), _mm512_add_ps(a2, a3))); + for (; i < n; ++i) s += bf16_to_f32(w[i]) * bf16_to_f32(x[i]); + return s; +} +#endif // avx512bf16 available + +// --------- AVX2 + FMA (256-bit fallback for CPUs without AVX-512) ---------- +// Covers Intel 12-14th gen / Arrow Lake (AVX-512 fused off) and AMD Zen<4. The +// bf16->fp32 widen is a zero-extend + <<16; with 4 independent accumulators the +// GEMV is memory-bandwidth bound, same as the AVX-512 path (just half the width). +CPU_MOE_TARGET("avx2,fma") +float dot_avx2(const bf16_t* w, const bf16_t* x, int n) { + __m256 a0 = _mm256_setzero_ps(), a1 = _mm256_setzero_ps(); + __m256 a2 = _mm256_setzero_ps(), a3 = _mm256_setzero_ps(); + int i = 0; + for (; i + 32 <= n; i += 32) { + _mm_prefetch(reinterpret_cast(w + i) + PF_AHEAD, _MM_HINT_T0); + for (int j = 0; j < 32; j += 8) { + __m128i wi = _mm_loadu_si128(reinterpret_cast(w + i + j)); + __m128i xi = _mm_loadu_si128(reinterpret_cast(x + i + j)); + __m256 wf = _mm256_castsi256_ps(_mm256_slli_epi32(_mm256_cvtepu16_epi32(wi), 16)); + __m256 xf = _mm256_castsi256_ps(_mm256_slli_epi32(_mm256_cvtepu16_epi32(xi), 16)); + __m256& acc = (j == 0) ? a0 : (j == 8) ? a1 : (j == 16) ? a2 : a3; + acc = _mm256_fmadd_ps(wf, xf, acc); + } + } + for (; i + 8 <= n; i += 8) { + __m128i wi = _mm_loadu_si128(reinterpret_cast(w + i)); + __m128i xi = _mm_loadu_si128(reinterpret_cast(x + i)); + __m256 wf = _mm256_castsi256_ps(_mm256_slli_epi32(_mm256_cvtepu16_epi32(wi), 16)); + __m256 xf = _mm256_castsi256_ps(_mm256_slli_epi32(_mm256_cvtepu16_epi32(xi), 16)); + a0 = _mm256_fmadd_ps(wf, xf, a0); + } + float s = hsum256(_mm256_add_ps(_mm256_add_ps(a0, a1), _mm256_add_ps(a2, a3))); + for (; i < n; ++i) s += bf16_to_f32(w[i]) * bf16_to_f32(x[i]); + return s; +} +#endif // CPU_MOE_X86 + +dot_fn select_dot() { + const IsaTier t = pick_isa(); +#if CPU_MOE_X86 +#ifdef CPU_MOE_HAS_AVX512BF16 + if (t >= ISA_AVX512BF16) return dot_avx512bf16; +#endif + if (t >= ISA_AVX512) return dot_avx512f; + if (t >= ISA_AVX2) return dot_avx2; +#endif + (void)t; + return dot_scalar; +} + +class Bf16Format final : public MoeFormat { + public: + Bf16Format(const MoeParams& p, const FormatArgs& a) + : MoeFormat(p), dot_(select_dot()), gate_up_(a.bank("gate_up")), down_(a.bank("down")) {} + + std::string isa_name() const override { return tier_name(pick_isa()); } + + void gate_up(int layer, int e, int /*tok*/, int i0, int i1, const bf16_t* x, float w_in, + bf16_t* g) const override { + const bf16_t* w = static_cast(tbl_at(gate_up_, layer)) + (size_t)e * (2 * p.I) * p.H; + gated_rows(p, i0, i1, w_in, g, [&](int row) { return dot_(w + (size_t)row * p.H, x, p.H); }); + } + + void down(int layer, int h0, int h1, const DownRoute* routes, int n, float* out) const override { + const bf16_t* w = static_cast(tbl_at(down_, layer)); + for (int h = h0; h < h1; ++h) + for (int j = 0; j < n; ++j) + out[j * kRowBlock + h - h0] = dot_(w + ((size_t)routes[j].e * p.H + h) * p.I, routes[j].g, p.I); + } + + private: + dot_fn dot_; + // bf16 rows: gate_up [E, 2I, H], down [E, H, I]. + const uint64_t* gate_up_; + const uint64_t* down_; +}; + +} // namespace + +std::unique_ptr make_bf16_format(const MoeParams& p, const FormatArgs& a) { + return std::make_unique(p, a); +} + +} // namespace freetoken::cpu_moe diff --git a/python/freetoken/kernel/csrc/cpu_moe/kernels/ds_fp4.cpp b/python/freetoken/kernel/csrc/cpu_moe/kernels/ds_fp4.cpp new file mode 100644 index 0000000000..a7afe9804d --- /dev/null +++ b/python/freetoken/kernel/csrc/cpu_moe/kernels/ds_fp4.cpp @@ -0,0 +1,269 @@ +#include + +#include "../format.h" + +namespace freetoken::cpu_moe { +namespace { + +// ----------------------- DeepSeek-V4 ds_fp4 (W4A8) --------------------------- +// Row-major e2m1 (2/byte, low nibble first) + e8m0 per-32 block scale, no global +// (w = E2M1[code] * 2^(e8m0-127)); activations are FP8-e4m3 round-tripped (per-128 +// block, ue8m0 scale) before each GEMM. Matches kernel/triton/dsv4 (fused_moe + +// fp8_linear): silu(clamp(gate,max=lim)) * clamp(up,-lim,lim), router weight on the +// down output. + +// Activations are pre-deinterleaved to fp32 (xe[m]=x[2m], xo[m]=x[2m+1]) once per +// token/route and reused across all output rows. This drops the hot dot to one +// vpmovzxbd + two vpermps per 32 weights (1.5 shuffle ops / 16 vs 4 for a bf16, +// dup-permute, per-element gather), so the row-major fp4 GEMV stops being port-5 +// bound and approaches the bf16 memory-bandwidth ceiling. +using dsdot_fn = float (*)(const uint8_t*, const uint8_t*, const float*, const float*, int, + const float*, const float*); + +float dot_dsfp4_scalar(const uint8_t* packed, const uint8_t* scale, const float* xe, + const float* xo, int K, const float* e2m1, const float* e8m0) { + float acc = 0.0f; + const int nb = K / 32; + for (int b = 0; b < nb; ++b) { + const float sc = e8m0[scale[b]]; + const uint8_t* pk = packed + (size_t)b * 16; // 16 bytes = 32 codes + const float* xeb = xe + (size_t)b * 16; + const float* xob = xo + (size_t)b * 16; + float bsum = 0.0f; + for (int j = 0; j < 16; ++j) { + const uint8_t byte = pk[j]; + bsum += e2m1[byte & 0xF] * xeb[j]; // low nibble -> even-K activation + bsum += e2m1[byte >> 4] * xob[j]; // high nibble -> odd-K activation + } + acc += sc * bsum; + } + return acc; +} + +#if CPU_MOE_X86 +// One 32-block: 16 bytes -> 16 low + 16 high nibble values via two vpermps, times +// the pre-split even/odd fp32 activations, folded by the per-32 e8m0 scale. +CPU_MOE_TARGET("avx512f") +inline __m512 dsfp4_blk(const uint8_t* pk, const float* xeb, const float* xob, __m512 lut, + __m512i loma, float sc) { + __m512i wi = _mm512_cvtepu8_epi32(_mm_loadu_si128(reinterpret_cast(pk))); + __m512 vlo = _mm512_permutexvar_ps(_mm512_and_si512(wi, loma), lut); + __m512 vhi = _mm512_permutexvar_ps(_mm512_and_si512(_mm512_srli_epi32(wi, 4), loma), lut); + __m512 prod = _mm512_fmadd_ps(vlo, _mm512_loadu_ps(xeb), _mm512_mul_ps(vhi, _mm512_loadu_ps(xob))); + return _mm512_mul_ps(prod, _mm512_set1_ps(sc)); +} + +CPU_MOE_TARGET("avx512f") +float dot_dsfp4_avx512(const uint8_t* packed, const uint8_t* scale, const float* xe, + const float* xo, int K, const float* e2m1, const float* e8m0) { + const __m512 lut = _mm512_loadu_ps(e2m1); + const __m512i loma = _mm512_set1_epi32(0xF); + __m512 acc0 = _mm512_setzero_ps(), acc1 = _mm512_setzero_ps(); + const int nb = K / 32; // 16 packed bytes + one e8m0 scale per block + int b = 0; + for (; b + 2 <= nb; b += 2) { // two independent accumulators hide FMA latency + acc0 = _mm512_add_ps(acc0, dsfp4_blk(packed + (size_t)b * 16, xe + (size_t)b * 16, + xo + (size_t)b * 16, lut, loma, e8m0[scale[b]])); + acc1 = _mm512_add_ps(acc1, dsfp4_blk(packed + (size_t)(b + 1) * 16, xe + (size_t)(b + 1) * 16, + xo + (size_t)(b + 1) * 16, lut, loma, e8m0[scale[b + 1]])); + } + for (; b < nb; ++b) + acc0 = _mm512_add_ps(acc0, dsfp4_blk(packed + (size_t)b * 16, xe + (size_t)b * 16, + xo + (size_t)b * 16, lut, loma, e8m0[scale[b]])); + return _mm512_reduce_add_ps(_mm512_add_ps(acc0, acc1)); +} + +// AVX2: a 32-K block is 16 bytes -> two 8-lane halves (8 even + 8 odd each). +CPU_MOE_TARGET("avx2,fma") +inline __m256 dsfp4_half_avx2(const uint8_t* pk, const float* xeb, const float* xob, __m256 mag8) { + __m256i wi = _mm256_cvtepu8_epi32(_mm_loadl_epi64(reinterpret_cast(pk))); + __m256 vlo = e2m1_decode8(_mm256_and_si256(wi, _mm256_set1_epi32(0xF)), mag8); + __m256 vhi = e2m1_decode8(_mm256_srli_epi32(wi, 4), mag8); + return _mm256_fmadd_ps(vlo, _mm256_loadu_ps(xeb), _mm256_mul_ps(vhi, _mm256_loadu_ps(xob))); +} + +CPU_MOE_TARGET("avx2,fma") +float dot_dsfp4_avx2(const uint8_t* packed, const uint8_t* scale, const float* xe, + const float* xo, int K, const float* e2m1, const float* e8m0) { + const __m256 mag8 = _mm256_loadu_ps(e2m1); + __m256 acc0 = _mm256_setzero_ps(), acc1 = _mm256_setzero_ps(); + const int nb = K / 32; + for (int b = 0; b < nb; ++b) { + const uint8_t* pk = packed + (size_t)b * 16; + const float* xeb = xe + (size_t)b * 16; + const float* xob = xo + (size_t)b * 16; + const __m256 sc = _mm256_set1_ps(e8m0[scale[b]]); + acc0 = _mm256_fmadd_ps(dsfp4_half_avx2(pk, xeb, xob, mag8), sc, acc0); + acc1 = _mm256_fmadd_ps(dsfp4_half_avx2(pk + 8, xeb + 8, xob + 8, mag8), sc, acc1); + } + return hsum256(_mm256_add_ps(acc0, acc1)); +} +#endif + +dsdot_fn select_dsdot() { + const IsaTier t = pick_isa(); +#if CPU_MOE_X86 + if (t >= ISA_AVX512) return dot_dsfp4_avx512; + if (t >= ISA_AVX2) return dot_dsfp4_avx2; +#endif + (void)t; + return dot_dsfp4_scalar; +} + +// Round a clamped |x|<=448 to nearest float8-e4m3 (RNE), back to fp32. Matches +// torch.float8_e4m3fn / triton .to(float8e4nv). +inline float e4m3_round(float x) { + const float sign = x < 0.0f ? -1.0f : 1.0f; + const float a = std::fabs(x); + if (a == 0.0f) return 0.0f; + if (a >= 448.0f) return sign * 448.0f; + int e; + std::frexp(a, &e); // a in [2^(e-1), 2^e) + float step = std::ldexp(1.0f, e - 4); + const float min_step = std::ldexp(1.0f, -9); // e4m3 subnormal step (2^-9) + if (step < min_step) step = min_step; + float r = std::nearbyint(a / step) * step; + if (r > 448.0f) r = 448.0f; + return sign * r; +} + +// IEEE ceil(log2(v)) for v>0 (matches dsv4 _log2_ceil / fast_round_scale). +inline int ceil_log2_pos(float v) { + uint32_t bits; + std::memcpy(&bits, &v, sizeof(bits)); + const int exp = (int)((bits >> 23) & 0xFF); + const int man = (int)(bits & 0x7FFFFF); + return exp - 127 + (man != 0 ? 1 : 0); +} + +// DeepSeek-V4 / V4.1 activation FP8 round-trip (bf16 in/out): per `block` (the checkpoint's +// fp8 block: 128 on V4, 32 on V4.1), s = 2^ceil(log2(max(|x|,1e-4)/448)); +// y = round_e4m3(clamp(x/s,+-448)) * s. +void fp8_roundtrip_bf16(const bf16_t* src, bf16_t* dst, int K, int block) { + for (int b0 = 0; b0 < K; b0 += block) { + const int b1 = std::min(K, b0 + block); + float amax = 1e-4f; + for (int i = b0; i < b1; ++i) amax = std::max(amax, std::fabs(bf16_to_f32(src[i]))); + const float s = std::ldexp(1.0f, ceil_log2_pos(amax * (1.0f / 448.0f))); + const float inv_s = 1.0f / s; + for (int i = b0; i < b1; ++i) { + float q = bf16_to_f32(src[i]) * inv_s; + q = std::min(448.0f, std::max(-448.0f, q)); + dst[i] = f32_to_bf16(e4m3_round(q) * s); + } + } +} + +class DsFp4Format final : public MoeFormat { + public: + DsFp4Format(const MoeParams& p, const FormatArgs& a) : MoeFormat(p), dot_(select_dsdot()) { + gate_up_ = a.bank("gate_up"); + gu_scale_ = a.bank("gate_up_scale"); + down_ = a.bank("down"); + dn_scale_ = a.bank("down_scale"); + act_block_ = static_cast(a.option("act_block")); + input_prequant_ = a.option("input_prequant", 0) != 0; + const size_t tokens = p.max_tokens, routes = (size_t)p.max_tokens * p.top_k; + xq_.resize(tokens * p.H); + xe_.resize(tokens * (p.H / 2)); + xo_.resize(tokens * (p.H / 2)); + ge_.resize(routes * (p.I / 2)); + go_.resize(routes * (p.I / 2)); + for (int i = 0; i < 16; ++i) e2m1_lut_[i] = kE2M1[i]; + for (int i = 0; i < 256; ++i) e8m0_lut_[i] = std::ldexp(1.0f, std::min(i, 254) - 127); + } + + std::string isa_name() const override { return tier_name(pick_isa()); } + + + // DSV4 FP8-round-trips the input before the gate_up GEMV, unless the GPU already did. + void prep_input(int tok, const bf16_t* x) override { + const bf16_t* src = x; + if (!input_prequant_) { + bf16_t* xq = xq_.data() + (size_t)tok * p.H; + fp8_roundtrip_bf16(x, xq, p.H, act_block_); + src = xq; + } + deinterleave_bf16_f32(src, xe_.data() + (size_t)tok * (p.H / 2), xo_.data() + (size_t)tok * (p.H / 2), p.H); + } + + bool preps_inter() const override { return true; } + + void prep_inter(int64_t r, bf16_t* g) override { + fp8_roundtrip_bf16(g, g, p.I, act_block_); + deinterleave_bf16_f32(g, ge_.data() + (size_t)r * (p.I / 2), go_.data() + (size_t)r * (p.I / 2), p.I); + } + + void gate_up(int layer, int e, int tok, int i0, int i1, const bf16_t* /*x*/, float /*w_in*/, + bf16_t* g) const override { + const int H = p.H, I = p.I; + const uint8_t* gu_packed_l = static_cast(tbl_at(gate_up_, layer)); + const uint8_t* gu_scale_l = static_cast(tbl_at(gu_scale_, layer)); + const int N2 = 2 * I, Hh = H / 2, Hs = H / 32; + const float* xe = xe_.data() + (size_t)tok * (H / 2); + const float* xo = xo_.data() + (size_t)tok * (H / 2); + const uint8_t* gp = gu_packed_l + (size_t)e * N2 * Hh; + const uint8_t* gs = gu_scale_l + (size_t)e * N2 * Hs; + const float lim = p.swiglu_limit; + for (int i = i0; i < i1; ++i) { + // gate_up is stored bf16 by the reference GEMV before swiglu; round to match. + float gate = bf16_to_f32(f32_to_bf16( + dot_(gp + (size_t)i * Hh, gs + (size_t)i * Hs, xe, xo, H, e2m1_lut_, e8m0_lut_))); + float up = bf16_to_f32(f32_to_bf16(dot_( + gp + (size_t)(I + i) * Hh, gs + (size_t)(I + i) * Hs, xe, xo, H, e2m1_lut_, e8m0_lut_))); + if (lim > 0.0f) { + if (gate > lim) gate = lim; + if (up > lim) up = lim; + else if (up < -lim) up = -lim; + } + const float glu = gate / (1.0f + std::exp(-gate)); // silu(gate) + g[i] = f32_to_bf16(glu * up); + } + } + + void down(int layer, int h0, int h1, const DownRoute* routes, int n, float* out) const override { + const int H = p.H, I = p.I, Ih = I / 2, Is = I / 32; + const uint8_t* dn_packed_l = static_cast(tbl_at(down_, layer)); + const uint8_t* dn_scale_l = static_cast(tbl_at(dn_scale_, layer)); + for (int h = h0; h < h1; ++h) { + for (int j = 0; j < n; ++j) { + const size_t e = (size_t)routes[j].e, r = (size_t)routes[j].r; + const uint8_t* dp = dn_packed_l + e * (size_t)H * Ih + (size_t)h * Ih; + const uint8_t* ds = dn_scale_l + e * (size_t)H * Is + (size_t)h * Is; + out[j * kRowBlock + h - h0] = dot_(dp, ds, ge_.data() + r * Ih, go_.data() + r * Ih, I, e2m1_lut_, e8m0_lut_); + } + } + } + + // The reference rounds each route's weighted down output to bf16 before the fp32 sum over routes. + bool rounds_routes() const override { return true; } + + private: + dsdot_fn dot_; + // e2m1 rows [E,2I,H/2] with e8m0 scales per 32 along K [E,2I,H/32]; down likewise with K = I. + const uint64_t* gate_up_ = nullptr; + const uint64_t* gu_scale_ = nullptr; + const uint64_t* down_ = nullptr; + const uint64_t* dn_scale_ = nullptr; + // Activation FP8 round-trip block, the checkpoint's fp8 block. + int act_block_ = 128; + // The caller already FP8-round-tripped the input activations on the GPU + // (same reference grid), so prep_input must not repeat it on the host-callback + // thread. That scalar per-element pass is single-threaded ON THE DECODE CRITICAL + // PATH (~0.3ms/layer at H=4096, every worker and the GPU waiting on it); moving + // it to a captured GPU elementwise kernel removes it while keeping the official + // W4A8 numerics bit-exact (see cpu_executor.py). + bool input_prequant_ = false; + float e2m1_lut_[16]; + float e8m0_lut_[256]; + AlignedVector xq_; + AlignedVector xe_, xo_, ge_, go_; +}; + +} // namespace + +std::unique_ptr make_ds_fp4_format(const MoeParams& p, const FormatArgs& a) { + return std::make_unique(p, a); +} + +} // namespace freetoken::cpu_moe diff --git a/python/freetoken/kernel/csrc/cpu_moe/kernels/fp8_block.cpp b/python/freetoken/kernel/csrc/cpu_moe/kernels/fp8_block.cpp new file mode 100644 index 0000000000..d27fb1868e --- /dev/null +++ b/python/freetoken/kernel/csrc/cpu_moe/kernels/fp8_block.cpp @@ -0,0 +1,203 @@ +#include +#include + +#include "../format.h" + +namespace freetoken::cpu_moe { +namespace { + +// ------------------------ block-FP8 (W8A16) dequant ------------------------ +// fp8-e4m3 weights row-major [rows, K] plus one bf16 scale per 128x128 block. The scale tensor is named weight_scale_inv but multiplies, as in kernel/triton/fp8_blockscale_moe.py. +// Whole 128-blocks read the activations as fp32 even/odd halves (xe[m] = x[2m], xo[m] = x[2m+1]); a ragged tail reads the bf16 row x. +// 0x7F/0xFF (NaN in torch float8_e4m3fn) decode to +-480 in every kernel, as e4m3_decode does. +using fp8_block_dot_fn = float (*)(const uint8_t*, const float*, const float*, const bf16_t*, + const bf16_t*, int, const float*); + +constexpr int FP8_BLK = 128; // K elements covered by one weight scale + +inline float fp8_block_tail(const uint8_t* w, const bf16_t* x, int k0, int K, const float* e4m3) { + float blk = 0.0f; + for (int k = k0; k < K; ++k) blk += e4m3[w[k]] * bf16_to_f32(x[k]); + return blk; +} + +float dot_fp8_block_scalar(const uint8_t* w, const float* xe, const float* xo, const bf16_t* x, + const bf16_t* s, int K, const float* e4m3) { + float acc = 0.0f; + int b = 0, k0 = 0; + for (; k0 + FP8_BLK <= K; k0 += FP8_BLK, ++b) { + // 8 independent sums: a single running sum is bound by the add latency + float a[8] = {}; + const uint8_t* wb = w + k0; + const float* pe = xe + k0 / 2; + const float* po = xo + k0 / 2; + for (int j = 0; j < FP8_BLK / 2; j += 4) + for (int l = 0; l < 4; ++l) { + a[2 * l] += e4m3[wb[2 * (j + l)]] * pe[j + l]; + a[2 * l + 1] += e4m3[wb[2 * (j + l) + 1]] * po[j + l]; + } + acc += (((a[0] + a[1]) + (a[2] + a[3])) + ((a[4] + a[5]) + (a[6] + a[7]))) * + bf16_to_f32(s[b]); + } + if (k0 < K) acc += fp8_block_tail(w, x, k0, K, e4m3) * bf16_to_f32(s[b]); + return acc; +} + +#if CPU_MOE_X86 +// A byte b read as fp16 bits (b & 0x80) << 8 | (b & 0x7F) << 7 is its e4m3 value times 2^-8, subnormals included, so the block scale carries the 256. +// In a 16-bit lane holding bytes (2m, 2m+1), srai 1 then & 0xBF80 builds that for byte 2m+1; a left shift by 8 first does byte 2m. Every AVX2 CPU has F16C. +CPU_MOE_TARGET("avx2,fma,f16c") +float dot_fp8_block_avx2(const uint8_t* w, const float* xe, const float* xo, const bf16_t* x, + const bf16_t* s, int K, const float* e4m3) { + const __m256i keep = _mm256_set1_epi16((short)0xBF80); + __m256 total = _mm256_setzero_ps(); + int b = 0, k0 = 0; + for (; k0 + FP8_BLK <= K; k0 += FP8_BLK, ++b) { + _mm_prefetch(reinterpret_cast(w + k0) + PF_AHEAD, _MM_HINT_T0); + __m256 a0 = _mm256_setzero_ps(), a1 = _mm256_setzero_ps(); + __m256 a2 = _mm256_setzero_ps(), a3 = _mm256_setzero_ps(); + for (int j = 0; j < FP8_BLK; j += 32) { + const __m256i v = _mm256_loadu_si256(reinterpret_cast(w + k0 + j)); + const __m256i odd = _mm256_and_si256(_mm256_srai_epi16(v, 1), keep); + const __m256i even = _mm256_and_si256(_mm256_srai_epi16(_mm256_slli_epi16(v, 8), 1), keep); + const float* pe = xe + (k0 + j) / 2; + const float* po = xo + (k0 + j) / 2; + a0 = _mm256_fmadd_ps(_mm256_cvtph_ps(_mm256_castsi256_si128(even)), _mm256_loadu_ps(pe), a0); + a1 = _mm256_fmadd_ps(_mm256_cvtph_ps(_mm256_extracti128_si256(even, 1)), _mm256_loadu_ps(pe + 8), a1); + a2 = _mm256_fmadd_ps(_mm256_cvtph_ps(_mm256_castsi256_si128(odd)), _mm256_loadu_ps(po), a2); + a3 = _mm256_fmadd_ps(_mm256_cvtph_ps(_mm256_extracti128_si256(odd, 1)), _mm256_loadu_ps(po + 8), a3); + } + const __m256 blk = _mm256_add_ps(_mm256_add_ps(a0, a1), _mm256_add_ps(a2, a3)); + total = _mm256_fmadd_ps(blk, _mm256_set1_ps(bf16_to_f32(s[b]) * 256.0f), total); + } + float acc = hsum256(total); + if (k0 < K) acc += fp8_block_tail(w, x, k0, K, e4m3) * bf16_to_f32(s[b]); + return acc; +} + +CPU_MOE_TARGET("avx512f,avx2,fma,f16c") +float dot_fp8_block_avx512f(const uint8_t* w, const float* xe, const float* xo, const bf16_t* x, + const bf16_t* s, int K, const float* e4m3) { + const __m256i keep = _mm256_set1_epi16((short)0xBF80); + __m512 total = _mm512_setzero_ps(); + int b = 0, k0 = 0; + for (; k0 + FP8_BLK <= K; k0 += FP8_BLK, ++b) { + _mm_prefetch(reinterpret_cast(w + k0) + PF_AHEAD, _MM_HINT_T0); + __m512 a0 = _mm512_setzero_ps(), a1 = _mm512_setzero_ps(); + __m512 a2 = _mm512_setzero_ps(), a3 = _mm512_setzero_ps(); + for (int j = 0; j < FP8_BLK; j += 64) { + const __m256i v0 = _mm256_loadu_si256(reinterpret_cast(w + k0 + j)); + const __m256i v1 = _mm256_loadu_si256(reinterpret_cast(w + k0 + j + 32)); + const __m256i o0 = _mm256_and_si256(_mm256_srai_epi16(v0, 1), keep); + const __m256i e0 = _mm256_and_si256(_mm256_srai_epi16(_mm256_slli_epi16(v0, 8), 1), keep); + const __m256i o1 = _mm256_and_si256(_mm256_srai_epi16(v1, 1), keep); + const __m256i e1 = _mm256_and_si256(_mm256_srai_epi16(_mm256_slli_epi16(v1, 8), 1), keep); + const float* pe = xe + (k0 + j) / 2; + const float* po = xo + (k0 + j) / 2; + a0 = _mm512_fmadd_ps(_mm512_cvtph_ps(e0), _mm512_loadu_ps(pe), a0); + a1 = _mm512_fmadd_ps(_mm512_cvtph_ps(o0), _mm512_loadu_ps(po), a1); + a2 = _mm512_fmadd_ps(_mm512_cvtph_ps(e1), _mm512_loadu_ps(pe + 16), a2); + a3 = _mm512_fmadd_ps(_mm512_cvtph_ps(o1), _mm512_loadu_ps(po + 16), a3); + } + const __m512 blk = _mm512_add_ps(_mm512_add_ps(a0, a1), _mm512_add_ps(a2, a3)); + total = _mm512_fmadd_ps(blk, _mm512_set1_ps(bf16_to_f32(s[b]) * 256.0f), total); + } + float acc = _mm512_reduce_add_ps(total); + if (k0 < K) acc += fp8_block_tail(w, x, k0, K, e4m3) * bf16_to_f32(s[b]); + return acc; +} +#endif // CPU_MOE_X86 + +fp8_block_dot_fn select_fp8_block_dot() { + const IsaTier t = pick_isa(); +#if CPU_MOE_X86 + // the avx512bf16 tier keeps the fp32 kernel: on Sapphire Rapids it is faster than widening to bf16 for dpbf16 + if (t >= ISA_AVX512) return dot_fp8_block_avx512f; + if (t >= ISA_AVX2) return dot_fp8_block_avx2; +#endif + (void)t; + return dot_fp8_block_scalar; +} + +class Fp8BlockFormat final : public MoeFormat { + public: + Fp8BlockFormat(const MoeParams& p, const FormatArgs& a) : MoeFormat(p), dot_(select_fp8_block_dot()) { + gate_up_ = a.bank("gate_up"); + gu_scale_ = a.bank("gate_up_scale"); + down_ = a.bank("down"); + dn_scale_ = a.bank("down_scale"); + gu_scale_stride_ = a.option("gu_scale_stride"); + dn_scale_stride_ = a.option("dn_scale_stride"); + const size_t tokens = p.max_tokens, routes = (size_t)p.max_tokens * p.top_k; + xe_.resize(tokens * (p.H / 2)); + xo_.resize(tokens * (p.H / 2)); + ge_.resize(routes * (p.I / 2)); + go_.resize(routes * (p.I / 2)); + for (int i = 0; i < 256; ++i) e4m3_lut_[i] = e4m3_decode((uint8_t)i); + } + + std::string isa_name() const override { return tier_name(pick_isa()); } + + + void prep_input(int tok, const bf16_t* x) override { + deinterleave_bf16_f32(x, xe_.data() + (size_t)tok * (p.H / 2), xo_.data() + (size_t)tok * (p.H / 2), p.H); + } + + bool preps_inter() const override { return true; } + + void prep_inter(int64_t r, bf16_t* g) override { + deinterleave_bf16_f32(g, ge_.data() + (size_t)r * (p.I / 2), go_.data() + (size_t)r * (p.I / 2), p.I); + } + + // The halves cover whole 128-blocks; the bf16 row covers a ragged tail. + void gate_up(int layer, int e, int tok, int i0, int i1, const bf16_t* x, float w_in, + bf16_t* g) const override { + const int H = p.H, I = p.I; + const uint8_t* packed = static_cast(tbl_at(gate_up_, layer)); + const bf16_t* scale = static_cast(tbl_at(gu_scale_, layer)); + const float* xe = xe_.data() + (size_t)tok * (H / 2); + const float* xo = xo_.data() + (size_t)tok * (H / 2); + gated_rows(p, i0, i1, w_in, g, [&](int row) { + const uint8_t* w = packed + ((size_t)e * (2 * I) + row) * (size_t)H; + const bf16_t* sc = scale + ((size_t)e * ((2 * I + FP8_BLK - 1) / FP8_BLK) + row / FP8_BLK) * + (size_t)gu_scale_stride_; + return dot_(w, xe, xo, x, sc, H, e4m3_lut_); + }); + } + + void down(int layer, int h0, int h1, const DownRoute* routes, int n, float* out) const override { + const int H = p.H, I = p.I; + const uint8_t* packed = static_cast(tbl_at(down_, layer)); + const bf16_t* scale = static_cast(tbl_at(dn_scale_, layer)); + for (int h = h0; h < h1; ++h) { + for (int j = 0; j < n; ++j) { + const size_t e = (size_t)routes[j].e, r = (size_t)routes[j].r; + const uint8_t* w = packed + (e * H + h) * (size_t)I; + const bf16_t* sc = scale + (e * ((H + FP8_BLK - 1) / FP8_BLK) + h / FP8_BLK) * + (size_t)dn_scale_stride_; + out[j * kRowBlock + h - h0] = + dot_(w, ge_.data() + r * (I / 2), go_.data() + r * (I / 2), routes[j].g, sc, I, e4m3_lut_); + } + } + } + + private: + fp8_block_dot_fn dot_; + // e4m3 rows [E,2I,H] and [E,H,I] with one bf16 scale per 128x128 block. + const uint64_t* gate_up_ = nullptr; + const uint64_t* gu_scale_ = nullptr; + const uint64_t* down_ = nullptr; + const uint64_t* dn_scale_ = nullptr; + // Row strides of the scale tensors, including any padding of the host bank layout. + int64_t gu_scale_stride_ = 0, dn_scale_stride_ = 0; + float e4m3_lut_[256]; + AlignedVector xe_, xo_, ge_, go_; +}; + +} // namespace + +std::unique_ptr make_fp8_block_format(const MoeParams& p, const FormatArgs& a) { + return std::make_unique(p, a); +} + +} // namespace freetoken::cpu_moe diff --git a/python/freetoken/kernel/csrc/cpu_moe/kernels/gguf.cpp b/python/freetoken/kernel/csrc/cpu_moe/kernels/gguf.cpp new file mode 100644 index 0000000000..63c8107dd3 --- /dev/null +++ b/python/freetoken/kernel/csrc/cpu_moe/kernels/gguf.cpp @@ -0,0 +1,124 @@ +#include +#include +#include +#include +#include + +#include "../format.h" +#include "cpu_kernels.h" + +namespace freetoken::cpu_moe { +namespace { + +std::atomic g_ggml{nullptr}; + +const ft_ggml_cpu_api& ggml_api() { + const ft_ggml_cpu_api* api = g_ggml.load(); + if (api == nullptr) throw std::logic_error("CPU MoE gguf experts need a ggml-cpu variant; call _cpu_moe.set_ggml_cpu first"); + return *api; +} + +} // namespace + +void set_ggml_cpu(uintptr_t addr) { + const auto* api = reinterpret_cast(addr); + if (api == nullptr || api->abi != FT_GGML_CPU_ABI) throw std::invalid_argument("set_ggml_cpu: not a ggml-cpu kernel table of this build"); + const ft_ggml_cpu_api* prev = nullptr; + if (!g_ggml.compare_exchange_strong(prev, api) && prev != api) + throw std::invalid_argument(std::string("set_ggml_cpu: this process already uses the ") + prev->name + " ggml-cpu variant"); +} + +namespace { + +// Packed ggml rows of the layers whose experts have this gate_up type and this down type, dotted by ggml-cpu against the activation converted to each type's vec_dot input (q8_0, q8_1 or q8_K, or the float type itself). +// With scales, each expert's (gate, up, down, 1) NVFP4 weight_scale_2 multiplies its gate and up rows before the glu and its down rows before the routing weight, as llama.cpp's build_moe_ffn does. +class GgufFormat final : public MoeFormat { + public: + GgufFormat(const MoeParams& p, const FormatArgs& a, bool scaled) + : MoeFormat(p), ggml_(ggml_api()), gate_up_(a.bank("gate_up")), down_(a.bank("down")), scales_(scaled ? a.bank("global") : nullptr) { + static std::once_flag tables; + std::call_once(tables, ggml_.init); + gu_ = kernel(static_cast(a.option("gate_up_type")), p.H, "gate_up"); + dn_ = kernel(static_cast(a.option("down_type")), p.I, "down"); + xf_.resize(p.H); + xq_.resize(static_cast(p.max_tokens) * gu_.act_row_bytes); + gq_.resize(static_cast(p.max_tokens) * p.top_k * dn_.act_row_bytes); + } + + std::string isa_name() const override { return std::string(ggml_.name) + "(ggml-cpu)"; } + + void prep_input(int tok, const bf16_t* x) override { + for (int i = 0; i < p.H; ++i) xf_[i] = bf16_to_f32(x[i]); + gu_.from_float(xf_.data(), xq_.data() + static_cast(tok) * gu_.act_row_bytes, p.H); + } + + bool preps_inter() const override { return true; } + + void prep_inter(int64_t r, bf16_t* g) override { + thread_local std::vector gf; + gf.resize(p.I); + for (int i = 0; i < p.I; ++i) gf[i] = bf16_to_f32(g[i]); + dn_.from_float(gf.data(), gq_.data() + static_cast(r) * dn_.act_row_bytes, p.I); + } + + void gate_up(int layer, int e, int tok, int i0, int i1, const bf16_t* /*x*/, float w_in, bf16_t* g) const override { + const uint8_t* w = static_cast(tbl_at(gate_up_, layer)) + static_cast(e) * 2 * p.I * gu_.row_bytes; + const uint8_t* xq = xq_.data() + static_cast(tok) * gu_.act_row_bytes; + const float* sc = expert_scales(layer, e); + gated_rows(p, i0, i1, w_in, g, [&](int row) { + float s; + gu_.vec_dot(p.H, &s, 0, w + static_cast(row) * gu_.row_bytes, 0, xq, 0, 1); + return sc == nullptr ? s : s * sc[row < p.I ? 0 : 1]; + }); + } + + void down(int layer, int h0, int h1, const DownRoute* routes, int n, float* out) const override { + const uint8_t* w = static_cast(tbl_at(down_, layer)); + for (int h = h0; h < h1; ++h) + for (int j = 0; j < n; ++j) + dn_.vec_dot(p.I, &out[j * kRowBlock + h - h0], 0, w + (static_cast(routes[j].e) * p.H + h) * dn_.row_bytes, 0, + gq_.data() + static_cast(routes[j].r) * dn_.act_row_bytes, 0, 1); + if (scales_ == nullptr) return; + for (int j = 0; j < n; ++j) { + const float s = expert_scales(layer, routes[j].e)[2]; + for (int h = h0; h < h1; ++h) out[j * kRowBlock + h - h0] *= s; + } + } + + private: + // the expert's (gate, up, down, 1) row of the layer's [E, 4] fp32 scale bank, or nullptr without scales + const float* expert_scales(int layer, int e) const { + return scales_ == nullptr ? nullptr : static_cast(tbl_at(scales_, layer)) + static_cast(e) * 4; + } + + ft_ggml_row_kernel kernel(int type, int k, const char* role) const { + ft_ggml_row_kernel out{}; + const int rc = ggml_.row_kernel(type, k, &out); + const std::string name = ggml_.type_name(type); + if (rc == FT_GGML_NO_KERNEL) throw std::invalid_argument("ggml-cpu has no dot product for " + name + " experts"); + if (rc != FT_GGML_OK) + throw std::invalid_argument("CPU MoE " + name + " " + role + " rows of " + std::to_string(k) + " values are not whole blocks"); + return out; + } + + const ft_ggml_cpu_api& ggml_; + const uint64_t* gate_up_; + const uint64_t* down_; + const uint64_t* scales_; + ft_ggml_row_kernel gu_{}, dn_{}; + std::vector xf_; + // the activation in the vec_dot input type: x per token, g per route + AlignedVector xq_, gq_; +}; + +} // namespace + +std::unique_ptr make_gguf_format(const MoeParams& p, const FormatArgs& a) { + return std::make_unique(p, a, false); +} + +std::unique_ptr make_gguf_scaled_format(const MoeParams& p, const FormatArgs& a) { + return std::make_unique(p, a, true); +} + +} // namespace freetoken::cpu_moe diff --git a/python/freetoken/kernel/csrc/cpu_moe/kernels/mxfp4.cpp b/python/freetoken/kernel/csrc/cpu_moe/kernels/mxfp4.cpp new file mode 100644 index 0000000000..484d6dc001 --- /dev/null +++ b/python/freetoken/kernel/csrc/cpu_moe/kernels/mxfp4.cpp @@ -0,0 +1,247 @@ +#include "../format.h" + +namespace freetoken::cpu_moe { +namespace { + +// ------------------------- mxfp4 (gpt-oss) GEMV ----------------------------- +// Transposed split-K layout: blk[Kpairs, N2] (N innermost), scl[Kpairs/16, N2] +// e8m0 per 32-K. Computes out[c] = sum_kb (E2M1[lo]*x[2kb] + E2M1[hi]*x[2kb+1]) +// * 2^(e8m0-127) for a contiguous column tile (blk/scl already offset to col 0 of +// the tile). Vectorized over N (16 columns / __m512), K stays the outer (cache- +// sequential) loop. Used by both gate_up (K=H) and down (K=I). +using mxgemv_fn = void (*)(float*, const uint8_t*, const uint8_t*, const bf16_t*, int, int, + int, const float*, const float*); + +void mxfp4_gemv_scalar(float* out, const uint8_t* blk, const uint8_t* scl, const bf16_t* x, + int Kpairs, int N2, int ncol, const float* e2m1, const float* e8m0) { + for (int c = 0; c < ncol; ++c) out[c] = 0.0f; + for (int kb = 0; kb < Kpairs; ++kb) { + const uint8_t* w = blk + (size_t)kb * N2; + const uint8_t* s = scl + (size_t)(kb >> 4) * N2; + const float xl = bf16_to_f32(x[2 * kb]); + const float xh = bf16_to_f32(x[2 * kb + 1]); + for (int c = 0; c < ncol; ++c) { + const uint8_t byte = w[c]; + out[c] += (e2m1[byte & 0xF] * xl + e2m1[byte >> 4] * xh) * e8m0[s[c]]; + } + } +} + +#if CPU_MOE_X86 +CPU_MOE_TARGET("avx512f") +void mxfp4_gemv_avx512(float* out, const uint8_t* blk, const uint8_t* scl, const bf16_t* x, + int Kpairs, int N2, int ncol, const float* e2m1, const float* e8m0) { + (void)e8m0; // e8m0[c]=2^(c-127) computed via bit construction (no gather) + const __m512 lut = _mm512_loadu_ps(e2m1); + const __m512i loma = _mm512_set1_epi32(0xF); + // K-outer / N-inner: each kb cache line is read once and all live column chunks + // (up to 4 -> 64 cols) accumulate from registers, so DRAM/L2 stream the tile once. + int c0 = 0; + for (; c0 + 16 <= ncol; c0 += 64) { + const int nchunk = std::min(4, (ncol - c0) / 16); + __m512 acc[4]; + for (int ci = 0; ci < nchunk; ++ci) acc[ci] = _mm512_setzero_ps(); + for (int kblk = 0; kblk < Kpairs; kblk += 16) { // 16 K-pairs = 32 K = one scale row + __m512 sc[4]; + for (int ci = 0; ci < nchunk; ++ci) { + __m128i sraw = _mm_loadu_si128(reinterpret_cast( + scl + (size_t)(kblk >> 4) * N2 + c0 + ci * 16)); + sc[ci] = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu8_epi32(sraw), 23)); + } + __m512 blk_acc[4]; + for (int ci = 0; ci < nchunk; ++ci) blk_acc[ci] = _mm512_setzero_ps(); + for (int kk = 0; kk < 16; ++kk) { + const int kb = kblk + kk; + const uint8_t* wbase = blk + (size_t)kb * N2 + c0; + // The transposed layout strides K by N2 bytes; prefetch ahead so the strided + // reads are not exposed to DRAM latency (the HW streamer misses big strides). + constexpr int PFD = 8; + if (kb + PFD < Kpairs) + _mm_prefetch(reinterpret_cast(blk + (size_t)(kb + PFD) * N2 + c0), + _MM_HINT_T0); + const __m512 xl = _mm512_set1_ps(bf16_to_f32(x[2 * kb])); + const __m512 xh = _mm512_set1_ps(bf16_to_f32(x[2 * kb + 1])); + for (int ci = 0; ci < nchunk; ++ci) { + __m512i wi = _mm512_cvtepu8_epi32( + _mm_loadu_si128(reinterpret_cast(wbase + ci * 16))); + __m512 vlo = _mm512_permutexvar_ps(_mm512_and_si512(wi, loma), lut); + __m512 vhi = _mm512_permutexvar_ps(_mm512_and_si512(_mm512_srli_epi32(wi, 4), loma), lut); + blk_acc[ci] = _mm512_fmadd_ps(vlo, xl, blk_acc[ci]); + blk_acc[ci] = _mm512_fmadd_ps(vhi, xh, blk_acc[ci]); + } + } + for (int ci = 0; ci < nchunk; ++ci) acc[ci] = _mm512_fmadd_ps(blk_acc[ci], sc[ci], acc[ci]); + } + for (int ci = 0; ci < nchunk; ++ci) _mm512_storeu_ps(out + c0 + ci * 16, acc[ci]); + } + for (int c = c0; c < ncol; ++c) { // tail columns (< 16) + float o = 0.0f; + for (int kb = 0; kb < Kpairs; ++kb) { + const uint8_t byte = blk[(size_t)kb * N2 + c]; + uint32_t bits = (uint32_t)scl[(size_t)(kb >> 4) * N2 + c] << 23; + float sc; + std::memcpy(&sc, &bits, 4); + o += (e2m1[byte & 0xF] * bf16_to_f32(x[2 * kb]) + + e2m1[byte >> 4] * bf16_to_f32(x[2 * kb + 1])) * sc; + } + out[c] = o; + } +} + +CPU_MOE_TARGET("avx2,fma") +void mxfp4_gemv_avx2(float* out, const uint8_t* blk, const uint8_t* scl, const bf16_t* x, + int Kpairs, int N2, int ncol, const float* e2m1, const float* e8m0) { + (void)e8m0; // e8m0[s]=2^(s-127) built via s<<23 (no gather) + const __m256 mag8 = _mm256_loadu_ps(e2m1); + int c0 = 0; + for (; c0 + 8 <= ncol; c0 += 32) { // up to 4 chunks of 8 = 32 cols + const int nchunk = std::min(4, (ncol - c0) / 8); + __m256 acc[4]; + for (int ci = 0; ci < nchunk; ++ci) acc[ci] = _mm256_setzero_ps(); + for (int kblk = 0; kblk < Kpairs; kblk += 16) { // 16 K-pairs = one scale row + __m256 sc[4]; + for (int ci = 0; ci < nchunk; ++ci) { + __m128i sraw = _mm_loadl_epi64(reinterpret_cast( + scl + (size_t)(kblk >> 4) * N2 + c0 + ci * 8)); + sc[ci] = _mm256_castsi256_ps(_mm256_slli_epi32(_mm256_cvtepu8_epi32(sraw), 23)); + } + __m256 blk_acc[4]; + for (int ci = 0; ci < nchunk; ++ci) blk_acc[ci] = _mm256_setzero_ps(); + for (int kk = 0; kk < 16; ++kk) { + const int kb = kblk + kk; + const uint8_t* wbase = blk + (size_t)kb * N2 + c0; + constexpr int PFD = 8; + if (kb + PFD < Kpairs) + _mm_prefetch(reinterpret_cast(blk + (size_t)(kb + PFD) * N2 + c0), + _MM_HINT_T0); + const __m256 xl = _mm256_set1_ps(bf16_to_f32(x[2 * kb])); + const __m256 xh = _mm256_set1_ps(bf16_to_f32(x[2 * kb + 1])); + for (int ci = 0; ci < nchunk; ++ci) { + __m256i wi = _mm256_cvtepu8_epi32( + _mm_loadl_epi64(reinterpret_cast(wbase + ci * 8))); + __m256 vlo = e2m1_decode8(_mm256_and_si256(wi, _mm256_set1_epi32(0xF)), mag8); + __m256 vhi = e2m1_decode8(_mm256_srli_epi32(wi, 4), mag8); + blk_acc[ci] = _mm256_fmadd_ps(vlo, xl, blk_acc[ci]); + blk_acc[ci] = _mm256_fmadd_ps(vhi, xh, blk_acc[ci]); + } + } + for (int ci = 0; ci < nchunk; ++ci) acc[ci] = _mm256_fmadd_ps(blk_acc[ci], sc[ci], acc[ci]); + } + for (int ci = 0; ci < nchunk; ++ci) _mm256_storeu_ps(out + c0 + ci * 8, acc[ci]); + } + for (int c = c0; c < ncol; ++c) { // tail columns (< 8); none when ncol%8==0 + float o = 0.0f; + for (int kb = 0; kb < Kpairs; ++kb) { + const uint8_t byte = blk[(size_t)kb * N2 + c]; + uint32_t bits = (uint32_t)scl[(size_t)(kb >> 4) * N2 + c] << 23; + float sc; + std::memcpy(&sc, &bits, 4); + o += (e2m1[byte & 0xF] * bf16_to_f32(x[2 * kb]) + + e2m1[byte >> 4] * bf16_to_f32(x[2 * kb + 1])) * sc; + } + out[c] = o; + } +} +#endif + +mxgemv_fn select_mxgemv() { + const IsaTier t = pick_isa(); +#if CPU_MOE_X86 + if (t >= ISA_AVX512) return mxfp4_gemv_avx512; + if (t >= ISA_AVX2) return mxfp4_gemv_avx2; +#endif + (void)t; + return mxfp4_gemv_scalar; +} + +// ----------------------------- mxfp4 (gpt-oss) ----------------------------- +// Transposed split-K layout (N innermost), so the GEMV streams K and accumulates +// a contiguous block of N output columns -> cache-line-efficient, no repack, no +// extra host memory. Pass 1 fuses gate_up + clamped-swiglu(+bias); pass 2 adds the +// down bias, and the executor applies the router weight and sums the routes. +// +// Dequant: w = E2M1[code] * 2^(e8m0_scale - 127); two codes per byte (low nibble +// first), one e8m0 scale per 32 contiguous K. Matches kernel/triton/mxfp4_moe.py. +class Mxfp4Format final : public MoeFormat { + public: + Mxfp4Format(const MoeParams& p, const FormatArgs& a) : MoeFormat(p), gemv_(select_mxgemv()) { + gate_up_ = a.bank("gate_up"); + gu_scale_ = a.bank("gate_up_scale"); + gu_bias_ = a.bank("gate_up_bias"); + down_ = a.bank("down"); + dn_scale_ = a.bank("down_scale"); + dn_bias_ = a.bank("down_bias"); + for (int i = 0; i < 16; ++i) e2m1_lut_[i] = kE2M1[i]; + // e8m0 (mxfp4 block scale) = 2^(s-127); the GPU GEMV clamps s to [0,254]. + for (int i = 0; i < 256; ++i) e8m0_lut_[i] = std::ldexp(1.0f, std::min(i, 254) - 127); + } + + std::string isa_name() const override { return tier_name(pick_isa()); } + + void gate_up(int layer, int e, int /*tok*/, int i0, int i1, const bf16_t* x, float /*w_in*/, + bf16_t* g) const override { + const int H = p.H, I = p.I; + const uint8_t* gu_packed_l = static_cast(tbl_at(gate_up_, layer)); + const uint8_t* gu_scale_l = static_cast(tbl_at(gu_scale_, layer)); + const bf16_t* gu_bias_l = static_cast(tbl_at(gu_bias_, layer)); + const int N2 = 2 * I; // gate_up output width (gate/up interleaved) + const int Hh = H / 2; // packed-K rows + const int nunit = i1 - i0; // intermediate units owned by this tile + const int col0 = 2 * i0; // first gate_up column + const int ncol = 2 * nunit; // gate_up columns owned by this tile + const uint8_t* blk_e = gu_packed_l + (size_t)e * Hh * N2; + const uint8_t* scl_e = gu_scale_l + (size_t)e * (size_t)(H / 32) * N2; + float gu[2 * kRowBlock]; + gemv_(gu, blk_e + col0, scl_e + col0, x, Hh, N2, ncol, e2m1_lut_, e8m0_lut_); + const bf16_t* bias_e = gu_bias_l + (size_t)e * N2 + col0; + const float lim = p.swiglu_limit, alpha = p.swiglu_alpha; + for (int j = 0; j < nunit; ++j) { + float gate = gu[2 * j] + bf16_to_f32(bias_e[2 * j]); + float up = gu[2 * j + 1] + bf16_to_f32(bias_e[2 * j + 1]); + if (gate > lim) gate = lim; + if (up > lim) up = lim; + else if (up < -lim) up = -lim; + const float glu = gate / (1.0f + std::exp(-gate * alpha)); // gate * sigmoid(alpha*gate) + g[i0 + j] = f32_to_bf16(glu * (up + 1.0f)); + } + } + + // The transposed layout already streams a whole column block per call, so the routes go one after another. + void down(int layer, int h0, int h1, const DownRoute* routes, int n, float* out) const override { + const int H = p.H, I = p.I; + const int Ih = I / 2; // packed-K rows + const int nh = h1 - h0; + const uint8_t* dn_packed_l = static_cast(tbl_at(down_, layer)); + const uint8_t* dn_scale_l = static_cast(tbl_at(dn_scale_, layer)); + const bf16_t* dn_bias_l = static_cast(tbl_at(dn_bias_, layer)); + for (int j = 0; j < n; ++j) { + const size_t e = (size_t)routes[j].e; + const uint8_t* blk_e = dn_packed_l + e * Ih * H; + const uint8_t* scl_e = dn_scale_l + e * (size_t)(I / 32) * H; + float* part = out + j * kRowBlock; + gemv_(part, blk_e + h0, scl_e + h0, routes[j].g, Ih, H, nh, e2m1_lut_, e8m0_lut_); + const bf16_t* bias_e = dn_bias_l + e * H + h0; + for (int c = 0; c < nh; ++c) part[c] += bf16_to_f32(bias_e[c]); + } + } + + private: + mxgemv_fn gemv_; + // gpt-oss split-K layout, N innermost: blocks [E,H/2,2I], e8m0 scales [E,H/32,2I], bf16 bias [E,2I]; down likewise with K = I, N = H. + const uint64_t* gate_up_ = nullptr; + const uint64_t* gu_scale_ = nullptr; + const uint64_t* gu_bias_ = nullptr; + const uint64_t* down_ = nullptr; + const uint64_t* dn_scale_ = nullptr; + const uint64_t* dn_bias_ = nullptr; + float e2m1_lut_[16]; + float e8m0_lut_[256]; +}; + +} // namespace + +std::unique_ptr make_mxfp4_format(const MoeParams& p, const FormatArgs& a) { + return std::make_unique(p, a); +} + +} // namespace freetoken::cpu_moe diff --git a/python/freetoken/kernel/csrc/cpu_moe/kernels/nvfp4.cpp b/python/freetoken/kernel/csrc/cpu_moe/kernels/nvfp4.cpp new file mode 100644 index 0000000000..3c5a56d51f --- /dev/null +++ b/python/freetoken/kernel/csrc/cpu_moe/kernels/nvfp4.cpp @@ -0,0 +1,466 @@ +#include + +#include "../format.h" + +namespace freetoken::cpu_moe { +namespace { + +// --------------------------- NVFP4 (W4A16) dequant --------------------------- +// Weights: e2m1 4-bit codes (2/byte, low nibble first), per-16 block scale in +// fp8-e4m3, per-output-row global scale in fp32. Dequant matches the GPU kernels +// (freetoken/kernel/triton/nvfp4_dequant.py): w = E2M1[code] * e4m3(scale) * global. +// Activations stay bf16 (W4A16); the GEMV dequantizes weights inside the K-loop. + + +// e2m1 * 2 as exact int8 (all e2m1 values are multiples of 0.5). Used by the AVX-VNNI +// W4A8 path: nibble -> int8 weight via PSHUFB LUT, then VPDPBUSD against int8 activations; +// the *2 is undone by a 0.5 folded into the final scale. Mirrors ggml's kvalues_mxfp4. +alignas(16) const int8_t kE2M1x2[16] = {0, 1, 2, 3, 4, 6, 8, 12, 0, -1, -2, -3, -4, -6, -8, -12}; + +// Activations pre-deinterleaved to fp32 (xe[m]=x[2m], xo[m]=x[2m+1]); see the +// ds_fp4 dot in ds_fp4.cpp for why (drops the hot loop to ~1.5 shuffle ops / 16 weights). +using nvdot_fn = float (*)(const uint8_t*, const uint8_t*, float, const float*, const float*, + int, const float*, const float*); + +float dot_nvfp4_scalar(const uint8_t* packed, const uint8_t* scale, float global, + const float* xe, const float* xo, int K, const float* e2m1, + const float* e4m3) { + float acc = 0.0f; + const int nb = K / 16; + for (int b = 0; b < nb; ++b) { + const float bs = e4m3[scale[b]]; + const uint8_t* pk = packed + (size_t)b * 8; + const float* xeb = xe + (size_t)b * 8; // 16 K -> 8 even + 8 odd + const float* xob = xo + (size_t)b * 8; + float bsum = 0.0f; + for (int j = 0; j < 8; ++j) { + const uint8_t byte = pk[j]; + bsum += e2m1[byte & 0xF] * xeb[j]; + bsum += e2m1[byte >> 4] * xob[j]; + } + acc += bs * bsum; + } + return acc * global; +} + +// ---- NVFP4 W4A8 (int8 activations) dot: nibble->int8 via LUT, per-16 act scale ---- +// asi8: int8 activations laid out per-16 block as [even(8), odd(8)]; asb[b] = per-block +// activation scale (absmax/127). Result folds the e2m1*2 -> *0.5 into the scale. +using nvi8dot_fn = float (*)(const uint8_t*, const uint8_t*, float, const int8_t*, int, + const float*, const float*); + +[[maybe_unused]] float dot_nvfp4_i8_scalar(const uint8_t* packed, const uint8_t* scale, + float global, const int8_t* asi8, int K, const float* e4m3, + const float* asb) { + float acc = 0.0f; + const int nb = K / 16; + for (int b = 0; b < nb; ++b) { + const uint8_t* pk = packed + (size_t)b * 8; + const int8_t* ae = asi8 + (size_t)b * 16; // even(8) + const int8_t* ao = ae + 8; // odd(8) + int isum = 0; + for (int j = 0; j < 8; ++j) { + isum += (int)kE2M1x2[pk[j] & 0xF] * (int)ae[j]; + isum += (int)kE2M1x2[pk[j] >> 4] * (int)ao[j]; + } + acc += (e4m3[scale[b]] * asb[b]) * (float)isum; + } + return acc * (0.5f * global); +} + +#if CPU_MOE_X86 +// Two 16-K blocks (16 packed bytes) per iter: lo nibbles -> even-K, hi -> odd-K, +// gathered via two vpermps. The per-16 e4m3 scale differs across the two blocks, so +// it is applied per lane (low 8 lanes = block b, high 8 = block b+1). +CPU_MOE_TARGET("avx512f") +inline __m512 nvfp4_blk2(const uint8_t* pk, const float* xeb, const float* xob, __m512 lut, + __m512i loma, float s0, float s1) { + __m512i wi = _mm512_cvtepu8_epi32(_mm_loadu_si128(reinterpret_cast(pk))); + __m512 vlo = _mm512_permutexvar_ps(_mm512_and_si512(wi, loma), lut); + __m512 vhi = _mm512_permutexvar_ps(_mm512_and_si512(_mm512_srli_epi32(wi, 4), loma), lut); + __m512 prod = _mm512_fmadd_ps(vlo, _mm512_loadu_ps(xeb), _mm512_mul_ps(vhi, _mm512_loadu_ps(xob))); + // lanes 0-7 (block b) -> s0, lanes 8-15 (block b+1) -> s1 (pure AVX512F mask move) + __m512 scv = _mm512_mask_mov_ps(_mm512_set1_ps(s0), 0xFF00, _mm512_set1_ps(s1)); + return _mm512_mul_ps(prod, scv); +} + +CPU_MOE_TARGET("avx512f") +float dot_nvfp4_avx512(const uint8_t* packed, const uint8_t* scale, float global, + const float* xe, const float* xo, int K, const float* e2m1, + const float* e4m3) { + const __m512 lut = _mm512_loadu_ps(e2m1); // 16 e2m1 values + const __m512i loma = _mm512_set1_epi32(0xF); + __m512 acc0 = _mm512_setzero_ps(), acc1 = _mm512_setzero_ps(); + const int nb = K / 16; // 8 packed bytes + one e4m3 scale per block + int b = 0; + for (; b + 4 <= nb; b += 4) { // two blk2 calls -> 4 blocks, 2 accumulators + acc0 = _mm512_add_ps(acc0, nvfp4_blk2(packed + (size_t)b * 8, xe + (size_t)b * 8, + xo + (size_t)b * 8, lut, loma, e4m3[scale[b]], e4m3[scale[b + 1]])); + acc1 = _mm512_add_ps(acc1, nvfp4_blk2(packed + (size_t)(b + 2) * 8, xe + (size_t)(b + 2) * 8, + xo + (size_t)(b + 2) * 8, lut, loma, e4m3[scale[b + 2]], e4m3[scale[b + 3]])); + } + for (; b + 2 <= nb; b += 2) + acc0 = _mm512_add_ps(acc0, nvfp4_blk2(packed + (size_t)b * 8, xe + (size_t)b * 8, + xo + (size_t)b * 8, lut, loma, e4m3[scale[b]], e4m3[scale[b + 1]])); + float s = _mm512_reduce_add_ps(_mm512_add_ps(acc0, acc1)); + for (; b < nb; ++b) { // odd final 16-K block + const uint8_t* pk = packed + (size_t)b * 8; + const float* xeb = xe + (size_t)b * 8; + const float* xob = xo + (size_t)b * 8; + float bsum = 0.0f; + for (int j = 0; j < 8; ++j) { + bsum += e2m1[pk[j] & 0xF] * xeb[j] + e2m1[pk[j] >> 4] * xob[j]; + } + s += e4m3[scale[b]] * bsum; + } + return s * global; +} + +// AVX2: one 16-K block (8 packed bytes) per call, 8 even + 8 odd lanes. +CPU_MOE_TARGET("avx2,fma") +inline __m256 nvfp4_blk_avx2(const uint8_t* pk, const float* xeb, const float* xob, + __m256 mag8, float sc) { + __m256i wi = _mm256_cvtepu8_epi32(_mm_loadl_epi64(reinterpret_cast(pk))); + __m256 vlo = e2m1_decode8(_mm256_and_si256(wi, _mm256_set1_epi32(0xF)), mag8); + __m256 vhi = e2m1_decode8(_mm256_srli_epi32(wi, 4), mag8); + __m256 prod = _mm256_fmadd_ps(vlo, _mm256_loadu_ps(xeb), _mm256_mul_ps(vhi, _mm256_loadu_ps(xob))); + return _mm256_mul_ps(prod, _mm256_set1_ps(sc)); +} + +CPU_MOE_TARGET("avx2,fma") +float dot_nvfp4_avx2(const uint8_t* packed, const uint8_t* scale, float global, + const float* xe, const float* xo, int K, const float* e2m1, + const float* e4m3) { + const __m256 mag8 = _mm256_loadu_ps(e2m1); // e2m1[0..7] magnitudes + __m256 acc0 = _mm256_setzero_ps(), acc1 = _mm256_setzero_ps(); + const int nb = K / 16; + int b = 0; + for (; b + 2 <= nb; b += 2) { + acc0 = _mm256_add_ps(acc0, nvfp4_blk_avx2(packed + (size_t)b * 8, xe + (size_t)b * 8, + xo + (size_t)b * 8, mag8, e4m3[scale[b]])); + acc1 = _mm256_add_ps(acc1, nvfp4_blk_avx2(packed + (size_t)(b + 1) * 8, xe + (size_t)(b + 1) * 8, + xo + (size_t)(b + 1) * 8, mag8, e4m3[scale[b + 1]])); + } + for (; b < nb; ++b) + acc0 = _mm256_add_ps(acc0, nvfp4_blk_avx2(packed + (size_t)b * 8, xe + (size_t)b * 8, + xo + (size_t)b * 8, mag8, e4m3[scale[b]])); + return hsum256(_mm256_add_ps(acc0, acc1)) * global; +} + +// AVX-VNNI W4A8: decode 8 packed bytes (16 nibbles) of one 16-block to int8 [lo(8),hi(8)] +// via PSHUFB against the e2m1*2 LUT (replaces the 2 vpermps fp32 expands -- ~4x less +// port-5 traffic). lo=even-K weights, hi=odd-K, matching the [even(8),odd(8)] act layout. +CPU_MOE_TARGET("avx2,avxvnni,fma") +inline __m128i nvfp4_decode_block_i8(const uint8_t* pk, __m128i lut) { + __m128i b = _mm_loadl_epi64(reinterpret_cast(pk)); // 8 bytes + __m128i lo = _mm_and_si128(b, _mm_set1_epi8(0x0F)); + __m128i hi = _mm_and_si128(_mm_srli_epi16(b, 4), _mm_set1_epi8(0x0F)); + return _mm_shuffle_epi8(lut, _mm_unpacklo_epi64(lo, hi)); // [lo(8),hi(8)] -> int8 +} + +// Two 16-blocks per VPDPBUSD (32 int8). Sign trick (ggml mul_add_epi8): |w|*(sign(w)*a)=w*a, +// so u8*s8 needs no offset/correction term. Per-block scale (e4m3 * act-scale) folded in fp32 +// (lanes 0-3 -> block b, 4-7 -> block b+1). Bit-faithful weight; only the int8 activation quant +// (W4A8) differs from the bf16 reference. +CPU_MOE_TARGET("avx2,avxvnni,fma") +float dot_nvfp4_i8_vnni(const uint8_t* packed, const uint8_t* scale, float global, + const int8_t* asi8, int K, const float* e4m3, const float* asb) { + const __m128i lut = _mm_loadu_si128(reinterpret_cast(kE2M1x2)); + __m256 accF = _mm256_setzero_ps(); + const int nb = K / 16; + int b = 0; + for (; b + 2 <= nb; b += 2) { + __m128i wb = nvfp4_decode_block_i8(packed + (size_t)b * 8, lut); + __m128i wb1 = nvfp4_decode_block_i8(packed + (size_t)(b + 1) * 8, lut); + __m256i w = _mm256_set_m128i(wb1, wb); // [blk b | blk b+1] + __m256i a = _mm256_loadu_si256(reinterpret_cast(asi8 + (size_t)b * 16)); + __m256i aw = _mm256_sign_epi8(w, w); // |w| (u8 operand) + __m256i sa = _mm256_sign_epi8(a, w); // sign(w)*a (s8 operand) + __m256i di = _mm256_dpbusd_avx_epi32(_mm256_setzero_si256(), aw, sa); + __m256 scv = _mm256_blend_ps(_mm256_set1_ps(e4m3[scale[b]] * asb[b]), + _mm256_set1_ps(e4m3[scale[b + 1]] * asb[b + 1]), 0xF0); + accF = _mm256_fmadd_ps(_mm256_cvtepi32_ps(di), scv, accF); + } + float s = hsum256(accF); + for (; b < nb; ++b) { // tail (odd block count) + const uint8_t* pk = packed + (size_t)b * 8; + const int8_t* ae = asi8 + (size_t)b * 16; const int8_t* ao = ae + 8; + int isum = 0; + for (int j = 0; j < 8; ++j) + isum += (int)kE2M1x2[pk[j] & 0xF] * (int)ae[j] + (int)kE2M1x2[pk[j] >> 4] * (int)ao[j]; + s += (e4m3[scale[b]] * asb[b]) * (float)isum; + } + return s * (0.5f * global); +} + +#ifdef CPU_MOE_HAS_AVX512VNNI + +// Software-prefetch distance for the W4A8 weight stream, in 16-K blocks (8 packed +// bytes each). Returns -1 when FREETOKEN_CPU_MOE_PF_BLOCKS is unset: the kernel then +// uses the built-in default min(512 blocks = 4 KB, 2 rows) -- 4 KB is the empirical +// optimum on large-row machines (Emerald Rapids sweep), while the 2-row cap keeps a small-row +// model's overshoot bounded (the executor works in 32-row tiles, so a fixed byte +// distance otherwise prefetches another worker's tile: duplicated DRAM traffic that +// regresses at the bandwidth ceiling). An EXPLICIT env value is honored verbatim +// (no clamp; 0 disables): the per-machine optimum can sit past the safe default (+20% +// at 4 KB on a 24-thread Ice Lake with 256B rows), so the escape hatch must reach it. +// Prefetch never faults, so overshooting a row/bank tail is safe. +static int nvfp4_pf_blocks() { + static const int v = [] { + const char* s = getenv("FREETOKEN_CPU_MOE_PF_BLOCKS"); + return (s && s[0]) ? atoi(s) : -1; + }(); + return v; +} + +// AVX-512 VNNI W4A8: FOUR 16-K blocks per VPDPBUSD (64 int8) -- 2x the AVX-VNNI +// (256-bit) path. Decode 32 packed bytes -> 64 int8 with a single _mm512_shuffle_epi8 +// (e2m1*2 LUT replicated to all 4 128-bit lanes). AVX-512 has no _mm512_sign_epi8, so +// the u8*s8 sign trick (|w| as u8, sign(w)*a as s8) uses abs_epi8 + a masked negate. +// Bit-faithful weight; only the int8 activation quant (W4A8) differs from bf16. +// One 4-block group (64 int8) -> scaled fp32 partial (16 lanes). Isolated as a helper so +// the caller can run several independent chains into separate accumulators (the decode -> +// dpbusd -> scale chain is long, so a single accumulator leaves the core latency-bound). +CPU_MOE_TARGET("avx512f,avx512bw,avx512vnni,avx2") +static inline __m512 nvfp4_i8_grp4(const uint8_t* packed, const uint8_t* scale, + const int8_t* asi8, const float* e4m3, const float* asb, + int b, __m512i lut, __m512i idx, __m512i mask0F, + __m512i idxsc) { + const __mmask64 hi_half = 0xFF00FF00FF00FF00ULL; // bytes 8-15 of each 128b lane -> hi nibbles + __m256i raw = _mm256_loadu_si256(reinterpret_cast(packed + (size_t)b * 8)); + __m512i src = _mm512_permutexvar_epi64(idx, _mm512_castsi256_si512(raw)); + __m512i lo = _mm512_and_si512(src, mask0F); + __m512i hi = _mm512_and_si512(_mm512_srli_epi16(src, 4), mask0F); + __m512i comb = _mm512_mask_blend_epi8(hi_half, lo, hi); + __m512i w = _mm512_shuffle_epi8(lut, comb); // 64 int8 weights (e2m1*2) + // u8*s8 sign trick without _mm512_sign_epi8: aw=|w|, sa = (w<0 ? -a : a) (a moot at w==0) + __m512i a = _mm512_loadu_si512(reinterpret_cast(asi8 + (size_t)b * 16)); + __m512i aw = _mm512_abs_epi8(w); + __mmask64 neg = _mm512_movepi8_mask(w); + __m512i sa = _mm512_mask_sub_epi8(a, neg, _mm512_setzero_si512(), a); + __m512i di = _mm512_dpbusd_epi32(_mm512_setzero_si512(), aw, sa); // 16 int32 (groups of 4) + // int32[0..3]->blk b, [4..7]->b+1, [8..11]->b+2, [12..15]->b+3. The 4 block scales + // (e4m3 LUT x per-block act scale) are computed vectorized: a 4-byte load + epu8->epi32 + // widen + one 4-lane LUT gather + one mul replaces 8 scalar loads + 4 scalar muls + + // a set_ps assembly, which otherwise dominates the per-group op count. + int sc_raw; + memcpy(&sc_raw, scale + b, 4); + __m128i sc4 = _mm_cvtepu8_epi32(_mm_cvtsi32_si128(sc_raw)); + __m128 s4 = _mm_mul_ps(_mm_i32gather_ps(e4m3, sc4, 4), _mm_loadu_ps(asb + b)); + __m512 scv = _mm512_permutexvar_ps(idxsc, _mm512_castps128_ps512(s4)); + return _mm512_mul_ps(_mm512_cvtepi32_ps(di), scv); +} + +CPU_MOE_TARGET("avx512f,avx512bw,avx512vnni,avx2") +float dot_nvfp4_i8_avx512vnni(const uint8_t* packed, const uint8_t* scale, float global, + const int8_t* asi8, int K, const float* e4m3, const float* asb) { + const __m512i lut = _mm512_broadcast_i32x4( + _mm_loadu_si128(reinterpret_cast(kE2M1x2))); + const __m512i idx = _mm512_set_epi64(3, 3, 2, 2, 1, 1, 0, 0); // block i -> 128b lane i + const __m512i mask0F = _mm512_set1_epi8(0x0F); + const __m512i idxsc = _mm512_set_epi32(3, 3, 3, 3, 2, 2, 2, 2, 1, 1, 1, 1, 0, 0, 0, 0); + __m512 acc0 = _mm512_setzero_ps(), acc1 = _mm512_setzero_ps(); + __m512 acc2 = _mm512_setzero_ps(), acc3 = _mm512_setzero_ps(); + const int nb = K / 16; + // Built-in default: min(4 KB, 2 rows) -- see nvfp4_pf_blocks(). An explicitly set + // env value is used verbatim so a per-machine sweep can reach operating points past + // the conservative default. + const int pfb = nvfp4_pf_blocks(); + const int pf = (pfb < 0) ? std::min(512, 2 * nb) : pfb; + int b = 0; + // Four independent 4-block groups per iter -> four accumulator chains hide the ~10-op + // decode->dpbusd->scale latency (1-2 chains leave the loop latency-bound: the per-core + // rate sat at ~60% of the core's achievable DRAM stream rate). + for (; b + 16 <= nb; b += 16) { + // Prefetch the weight stream ahead: the interleaved decode lowers the L1-miss + // concurrency the HW prefetcher sustains on its own (134 -> 167 GB/s on Emerald Rapids). + if (pf > 0) { + _mm_prefetch(reinterpret_cast(packed + ((size_t)b + (size_t)pf) * 8), + _MM_HINT_T0); + _mm_prefetch(reinterpret_cast(packed + ((size_t)b + (size_t)pf) * 8 + 64), + _MM_HINT_T0); + } + acc0 = _mm512_add_ps(acc0, nvfp4_i8_grp4(packed, scale, asi8, e4m3, asb, b, lut, idx, + mask0F, idxsc)); + acc1 = _mm512_add_ps(acc1, nvfp4_i8_grp4(packed, scale, asi8, e4m3, asb, b + 4, lut, idx, + mask0F, idxsc)); + acc2 = _mm512_add_ps(acc2, nvfp4_i8_grp4(packed, scale, asi8, e4m3, asb, b + 8, lut, idx, + mask0F, idxsc)); + acc3 = _mm512_add_ps(acc3, nvfp4_i8_grp4(packed, scale, asi8, e4m3, asb, b + 12, lut, idx, + mask0F, idxsc)); + } + for (; b + 4 <= nb; b += 4) + acc0 = _mm512_add_ps(acc0, nvfp4_i8_grp4(packed, scale, asi8, e4m3, asb, b, lut, idx, + mask0F, idxsc)); + float s = _mm512_reduce_add_ps( + _mm512_add_ps(_mm512_add_ps(acc0, acc1), _mm512_add_ps(acc2, acc3))); + for (; b < nb; ++b) { // tail (<4 remaining 16-K blocks) + const uint8_t* pk = packed + (size_t)b * 8; + const int8_t* ae = asi8 + (size_t)b * 16; const int8_t* ao = ae + 8; + int isum = 0; + for (int j = 0; j < 8; ++j) + isum += (int)kE2M1x2[pk[j] & 0xF] * (int)ae[j] + (int)kE2M1x2[pk[j] >> 4] * (int)ao[j]; + s += (e4m3[scale[b]] * asb[b]) * (float)isum; + } + return s * (0.5f * global); +} +#endif // avx512vnni available +#endif + +nvdot_fn select_nvdot() { + const IsaTier t = pick_isa(); +#if CPU_MOE_X86 + if (t >= ISA_AVX512) return dot_nvfp4_avx512; + if (t >= ISA_AVX2) return dot_nvfp4_avx2; +#endif + (void)t; + return dot_nvfp4_scalar; +} + +// Best W4A8 (int8-activation) nvfp4 dot, or nullptr if no SIMD VNNI (caller keeps the +// faithful fp32 nvdot path). The scalar i8 dot exists only as a correctness reference. +nvi8dot_fn select_nvi8dot() { +#if CPU_MOE_X86 +#if defined(CPU_MOE_HAS_AVX512VNNI) + if (cpu_has_avx512vnni()) return dot_nvfp4_i8_avx512vnni; +#endif + if (cpu_has_avxvnni()) return dot_nvfp4_i8_vnni; +#endif + return nullptr; +} + +// Quantize the pre-deinterleaved fp32 even/odd activations to per-16-block int8 in the +// [even(8),odd(8)] layout the VNNI dot expects. Done once per token/route (amortized over +// every output row), so a scalar pass is fine relative to the GEMV. +void quant_i8_pg16(const float* xe, const float* xo, int K, int8_t* asi8, float* asb) { + const int nb = K / 16; + for (int b = 0; b < nb; ++b) { + const float* xeb = xe + (size_t)b * 8; + const float* xob = xo + (size_t)b * 8; + float amax = 0.0f; + for (int j = 0; j < 8; ++j) + amax = std::max(amax, std::max(std::fabs(xeb[j]), std::fabs(xob[j]))); + const float s = amax > 0.0f ? amax / 127.0f : 1.0f; + asb[b] = s; + const float inv = 1.0f / s; + int8_t* ae = asi8 + (size_t)b * 16; + for (int j = 0; j < 8; ++j) { + int qe = (int)std::lround(xeb[j] * inv), qo = (int)std::lround(xob[j] * inv); + ae[j] = (int8_t)std::max(-127, std::min(127, qe)); + ae[8 + j] = (int8_t)std::max(-127, std::min(127, qo)); + } + } +} + +class Nvfp4Format final : public MoeFormat { + public: + Nvfp4Format(const MoeParams& p, const FormatArgs& a) + : MoeFormat(p), nvdot_(select_nvdot()), nvi8dot_(select_nvi8dot()) { + gate_up_ = a.bank("gate_up"); + gu_scale_ = a.bank("gate_up_scale"); + gu_global_ = a.bank("gate_up_global"); + down_ = a.bank("down"); + dn_scale_ = a.bank("down_scale"); + dn_global_ = a.bank("down_global"); + const size_t tokens = p.max_tokens, routes = (size_t)p.max_tokens * p.top_k; + xe_.resize(tokens * (p.H / 2)); + xo_.resize(tokens * (p.H / 2)); + ge_.resize(routes * (p.I / 2)); + go_.resize(routes * (p.I / 2)); + if (nvi8dot_ != nullptr) { + xi8_.resize(tokens * p.H); + xas_.resize(tokens * (p.H / 16)); + gi8_.resize(routes * p.I); + gas_.resize(routes * (p.I / 16)); + } + isa_ = tier_name(pick_isa()); + if (nvi8dot_ != nullptr) isa_ += cpu_has_avx512vnni() ? "+avx512vnni(nvfp4-w4a8)" : "+vnni(nvfp4-w4a8)"; + for (int i = 0; i < 16; ++i) e2m1_lut_[i] = kE2M1[i]; + for (int i = 0; i < 256; ++i) e4m3_lut_[i] = e4m3_decode((uint8_t)i); + } + + std::string isa_name() const override { return isa_; } + + + void prep_input(int tok, const bf16_t* x) override { + float* xe = xe_.data() + (size_t)tok * (p.H / 2); + float* xo = xo_.data() + (size_t)tok * (p.H / 2); + deinterleave_bf16_f32(x, xe, xo, p.H); + if (nvi8dot_ != nullptr) + quant_i8_pg16(xe, xo, p.H, xi8_.data() + (size_t)tok * p.H, xas_.data() + (size_t)tok * (p.H / 16)); + } + + bool preps_inter() const override { return true; } + + void prep_inter(int64_t r, bf16_t* g) override { + float* ge = ge_.data() + (size_t)r * (p.I / 2); + float* go = go_.data() + (size_t)r * (p.I / 2); + deinterleave_bf16_f32(g, ge, go, p.I); + if (nvi8dot_ != nullptr) + quant_i8_pg16(ge, go, p.I, gi8_.data() + (size_t)r * p.I, gas_.data() + (size_t)r * (p.I / 16)); + } + + void gate_up(int layer, int e, int tok, int i0, int i1, const bf16_t* /*x*/, float w_in, + bf16_t* g) const override { + const int H = p.H; + const uint8_t* packed = static_cast(tbl_at(gate_up_, layer)); + const uint8_t* scale = static_cast(tbl_at(gu_scale_, layer)); + const float* global = static_cast(tbl_at(gu_global_, layer)); + const float* xe = xe_.data() + (size_t)tok * (H / 2); + const float* xo = xo_.data() + (size_t)tok * (H / 2); + const int8_t* xi8 = nvi8dot_ != nullptr ? xi8_.data() + (size_t)tok * H : nullptr; + const float* xas = nvi8dot_ != nullptr ? xas_.data() + (size_t)tok * (H / 16) : nullptr; + gated_rows(p, i0, i1, w_in, g, [&](int row) { + const size_t r = (size_t)e * (2 * p.I) + row; + if (nvi8dot_ != nullptr) + return nvi8dot_(packed + r * (size_t)(H / 2), scale + r * (size_t)(H / 16), + global[r], xi8, H, e4m3_lut_, xas); + return nvdot_(packed + r * (size_t)(H / 2), scale + r * (size_t)(H / 16), + global[r], xe, xo, H, e2m1_lut_, e4m3_lut_); + }); + } + + void down(int layer, int h0, int h1, const DownRoute* routes, int n, float* out) const override { + const int I = p.I; + const uint8_t* packed = static_cast(tbl_at(down_, layer)); + const uint8_t* scale = static_cast(tbl_at(dn_scale_, layer)); + const float* global = static_cast(tbl_at(dn_global_, layer)); + for (int h = h0; h < h1; ++h) { + for (int j = 0; j < n; ++j) { + const size_t r = (size_t)routes[j].r; + const size_t row = (size_t)routes[j].e * p.H + h; + out[j * kRowBlock + h - h0] = nvi8dot_ != nullptr + ? nvi8dot_(packed + row * (size_t)(I / 2), scale + row * (size_t)(I / 16), global[row], + gi8_.data() + r * I, I, e4m3_lut_, gas_.data() + r * (I / 16)) + : nvdot_(packed + row * (size_t)(I / 2), scale + row * (size_t)(I / 16), global[row], + ge_.data() + r * (I / 2), go_.data() + r * (I / 2), I, e2m1_lut_, e4m3_lut_); + } + } + } + + private: + nvdot_fn nvdot_; + nvi8dot_fn nvi8dot_; // AVX-VNNI W4A8 nvfp4 dot (nullptr -> use fp32 nvdot) + // ModelOpt rows: packed e2m1 [E,2I,H/2], e4m3 scales per 16 along K [E,2I,H/16], fp32 row globals [E,2I]; down likewise with K = I. + const uint64_t* gate_up_ = nullptr; + const uint64_t* gu_scale_ = nullptr; + const uint64_t* gu_global_ = nullptr; + const uint64_t* down_ = nullptr; + const uint64_t* dn_scale_ = nullptr; + const uint64_t* dn_global_ = nullptr; + std::string isa_; + float e2m1_lut_[16]; + float e4m3_lut_[256]; + // Activations pre-deinterleaved to fp32 even/odd for the row-major dot; with VNNI also per-16-block int8 + scale. + AlignedVector xe_, xo_, ge_, go_; + AlignedVector xi8_, gi8_; + AlignedVector xas_, gas_; +}; + +} // namespace + +std::unique_ptr make_nvfp4_format(const MoeParams& p, const FormatArgs& a) { + return std::make_unique(p, a); +} + +} // namespace freetoken::cpu_moe diff --git a/python/freetoken/kernel/csrc/ggml/LICENSE b/python/freetoken/kernel/csrc/ggml/LICENSE new file mode 100644 index 0000000000..e7dca554bc --- /dev/null +++ b/python/freetoken/kernel/csrc/ggml/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2023-2026 The ggml authors + +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. diff --git a/python/freetoken/kernel/csrc/ggml/VERSION b/python/freetoken/kernel/csrc/ggml/VERSION new file mode 100644 index 0000000000..f7b5c01705 --- /dev/null +++ b/python/freetoken/kernel/csrc/ggml/VERSION @@ -0,0 +1,15 @@ +llama.cpp 8c2134a8755d8a9792af513dab20282affe903e4 (2026-10-14) + +Only the ggml files that FreeToken compiles or includes: +- the GGUF ops (freetoken/ggml_ops.cu, built by kernel/ggml.py): ggml/include -> include, ggml/src -> src, ggml/src/ggml-cuda -> ggml-cuda. +- the CPU MoE row kernels (freetoken/cpu_kernels.c and freetoken/variant_module.cpp, one freetoken.kernel.ggml_cpu module per x86 variant, built by setup.py): ggml/src/{ggml-quants.c,ggml-quants.h,ggml-backend-impl.h} -> src, ggml/include/ggml-cpu.h -> include, ggml/src/ggml-cpu/{quants.c,quants.h,vec.cpp,vec.h,ggml-cpu-impl.h,simd-mappings.h,arch-fallback.h,arch/x86/quants.c,arch/x86/cpu-feats.cpp} -> ggml-cpu. These have no local changes. + +Local changes, kept in freetoken.patch (already applied): +- quantize.cu/.cuh: the q8_1 quantize kernels read fp32, bf16 or fp16 activations. +- mmq.cu/.cuh: MMQ takes bf16 or fp16 activations, except on the native fp4 path. Its kernels are built for fp32 and bf16 results (dst_t); a bf16 result that needs the stream-k fixup is computed in fp32 first, and so is an fp16 result unless GGML_MMQ_F16_DST builds fp16 kernels too. +- mmvq.cu/.cuh: MMVQ takes bf16 or fp16 activations, can write bf16 or fp16 results, and can reuse activations already quantized to q8_1 (src1_q8_1_pre). +- mmq.cu/.cuh, mmid.cu/.cuh: MoE MMQ runs a flat grid over only the (expert, column tile) pairs that hold tokens, listed by mm_ids_tile_map; ncols_max > 0 keeps the regular grid. +- common.cuh: the native fp4 MMA (BLACKWELL_MMA_AVAILABLE, blackwell_mma_available) is off unless GGML_CUDA_FP4_MMA is defined; it needs an sm_120a/sm_120f build. +- mmvq.cu, common.cuh: the NVFP4 scale fusion (x_scale, gate_scale) reads strided scale vectors, e.g. columns of a per-expert scale bank. + +To update: copy the same files from a newer llama.cpp, apply freetoken.patch, and update this file. diff --git a/python/freetoken/kernel/csrc/ggml/freetoken.patch b/python/freetoken/kernel/csrc/ggml/freetoken.patch new file mode 100644 index 0000000000..595498dd15 --- /dev/null +++ b/python/freetoken/kernel/csrc/ggml/freetoken.patch @@ -0,0 +1,1058 @@ +--- a/ggml/src/ggml-cuda/common.cuh ++++ b/ggml/src/ggml-cuda/common.cuh +@@ -293,7 +293,8 @@ + #define AMPERE_MMA_AVAILABLE + #endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) && __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE + +-#if !defined(GGML_USE_HIP) && __CUDA_ARCH__ >= GGML_CUDA_CC_BLACKWELL && __CUDA_ARCH__ < GGML_CUDA_CC_RUBIN ++// FreeToken: the native fp4 MMA needs an sm_120a/sm_120f build, so it is off unless GGML_CUDA_FP4_MMA is defined; FP4 weights then take Q8_1 activations ++#if !defined(GGML_USE_HIP) && __CUDA_ARCH__ >= GGML_CUDA_CC_BLACKWELL && __CUDA_ARCH__ < GGML_CUDA_CC_RUBIN && defined(GGML_CUDA_FP4_MMA) + # define BLACKWELL_MMA_AVAILABLE + #endif // !defined(GGML_USE_HIP) && __CUDA_ARCH__ >= GGML_CUDA_CC_BLACKWELL + +@@ -374,8 +375,13 @@ + } + + static bool blackwell_mma_available(const int cc) { ++#ifdef GGML_CUDA_FP4_MMA + return GGML_CUDA_CC_IS_NVIDIA(cc) && ggml_cuda_highest_compiled_arch(cc) >= GGML_CUDA_CC_BLACKWELL && + ggml_cuda_highest_compiled_arch(cc) < GGML_CUDA_CC_RUBIN; ++#else ++ GGML_UNUSED(cc); ++ return false; ++#endif // GGML_CUDA_FP4_MMA + } + + // Checks whether the tensor's base data pointer and higher-dimensional strides are byte-aligned to `alignment` bytes. +@@ -1582,6 +1588,9 @@ + const void * gate_bias = nullptr; + const void * x_scale = nullptr; + const void * gate_scale = nullptr; ++ // FreeToken: element strides of x_scale / gate_scale, so they can be columns of a per-expert scale bank ++ uint32_t x_scale_stride = 1; ++ uint32_t gate_scale_stride = 1; + ggml_glu_op glu_op; + float glu_limit = 0.0f; + const void * shared_up = nullptr; +--- a/ggml/src/ggml-cuda/mmid.cu ++++ b/ggml/src/ggml-cuda/mmid.cu +@@ -177,3 +177,45 @@ + break; + } + } ++ ++// Lists the (expert, column tile) pairs that hold columns for a flat MoE grid, in expert order. ++// The other ntiles grid columns get tile jt_none of expert 0, which starts past its columns, so mul_mat_q skips them. ++template ++__launch_bounds__(nthreads, 1) ++static __global__ void mm_ids_tile_map( ++ const int32_t * __restrict__ expert_bounds, int2 * __restrict__ tile_map, const int n_experts, const int J, const int ntiles, const int jt_none) { ++ __shared__ int scan[nthreads]; ++ ++ int start = 0; // First tile of the current chunk of experts. ++ for (int e0 = 0; e0 < n_experts; e0 += nthreads) { ++ const int e = e0 + threadIdx.x; ++ const int n = e < n_experts ? (expert_bounds[e + 1] - expert_bounds[e] + J - 1) / J : 0; ++ ++ scan[threadIdx.x] = n; ++ __syncthreads(); ++#pragma unroll ++ for (int offset = 1; offset < nthreads; offset *= 2) { ++ const int add = threadIdx.x >= offset ? scan[threadIdx.x - offset] : 0; ++ __syncthreads(); ++ scan[threadIdx.x] += add; ++ __syncthreads(); ++ } ++ ++ const int first = start + scan[threadIdx.x] - n; ++ for (int jt = 0; jt < n; ++jt) { ++ tile_map[first + jt] = make_int2(e, jt); ++ } ++ start += scan[nthreads - 1]; ++ __syncthreads(); ++ } ++ ++ for (int t = start + threadIdx.x; t < ntiles; t += nthreads) { ++ tile_map[t] = make_int2(0, jt_none); ++ } ++} ++ ++void ggml_cuda_launch_mm_ids_tile_map( ++ const int32_t * expert_bounds, int2 * tile_map, const int n_experts, const int J, const int ntiles, const int jt_none, cudaStream_t stream) { ++ constexpr int nthreads = 512; ++ mm_ids_tile_map<<<1, nthreads, 0, stream>>>(expert_bounds, tile_map, n_experts, J, ntiles, jt_none); ++} +--- a/ggml/src/ggml-cuda/mmid.cuh ++++ b/ggml/src/ggml-cuda/mmid.cuh +@@ -3,3 +3,6 @@ + void ggml_cuda_launch_mm_ids_helper( + const int32_t * ids, int32_t * ids_src1, int32_t * ids_dst, int32_t * expert_bounds, + int n_experts, int n_tokens, int n_expert_used, int nchannels_y, int si1, int sis1, bool write_inverse, cudaStream_t stream); ++ ++void ggml_cuda_launch_mm_ids_tile_map( ++ const int32_t * expert_bounds, int2 * tile_map, int n_experts, int J, int ntiles, int jt_none, cudaStream_t stream); +--- a/ggml/src/ggml-cuda/mmq.cu ++++ b/ggml/src/ggml-cuda/mmq.cu +@@ -132,9 +132,11 @@ + } + + void ggml_cuda_mul_mat_q( +- ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, ggml_tensor * dst) { +- GGML_ASSERT( src1->type == GGML_TYPE_F32); +- GGML_ASSERT( dst->type == GGML_TYPE_F32); ++ ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, ggml_tensor * dst, ++ int64_t ncols_max) { ++ // FreeToken: bf16 and fp16 activations are quantized directly, except for the native fp4 paths; dst is fp32, bf16 or fp16 ++ GGML_ASSERT( src1->type == GGML_TYPE_F32 || src1->type == GGML_TYPE_BF16 || src1->type == GGML_TYPE_F16); ++ GGML_ASSERT( dst->type == GGML_TYPE_F32 || dst->type == GGML_TYPE_BF16 || dst->type == GGML_TYPE_F16); + GGML_ASSERT(!ids || ids->type == GGML_TYPE_I32); // Optional, used for batched GGML_MUL_MAT_ID. + + GGML_TENSOR_BINARY_OP_LOCALS; +@@ -156,7 +158,9 @@ + + const char * src0_d = (const char *) src0->data; + const float * src1_d = (const float *) src1->data; +- float * dst_d = (float *) dst->data; ++ const mmq_dst dst_d = {dst->data, dst->type}; ++ const nv_bfloat16 * src1_bf16 = src1->type == GGML_TYPE_BF16 ? (const nv_bfloat16 *) src1->data : nullptr; ++ const half * src1_f16 = src1->type == GGML_TYPE_F16 ? (const half *) src1->data : nullptr; + + // If src0 is a temporary compute buffer, clear any potential padding. + if (ggml_backend_buffer_get_usage(src0->buffer) == GGML_BACKEND_BUFFER_USAGE_COMPUTE) { +@@ -183,6 +187,7 @@ + const ggml_prec prec_src1 = ggml_cuda_mmq_get_prec_src1(src0, dst, cc); + + const bool use_native_fp4 = prec_src1 == GGML_PREC_Q4; ++ GGML_ASSERT(!use_native_fp4 || src1->type == GGML_TYPE_F32); + const size_t y_block_size = use_native_fp4 ? sizeof(block_fp4_mmq) : sizeof(block_q8_1_mmq); + const size_t y_values_per_block = use_native_fp4 ? QK_FP4_MMQ : QK8_1_MMQ; + +@@ -192,7 +197,8 @@ + int64_t ncols_opt = ne11; + if (ids) { + const int64_t n_expert_used = ids->ne[0]; +- ncols_opt = ne12; ++ // a flat grid runs only the tiles that hold tokens, so it sizes them for the average expert ++ ncols_opt = ncols_max > 0 ? ncols_max : (ne12*n_expert_used + ne02 - 1) / ne02; + + // Each expert only sees ne12*n_expert_used/ne02 tokens on average. + // On RDNA3 and RDNA4 it is faster to pick the tile size against this value instead of ne12. +@@ -249,6 +255,12 @@ + quantize_mmq_fp4_cuda(src1_d, nullptr, src1_q8_1.get(), src1_scale.ptr, src0->type, use_aligned_float8, ne10, s11, s12, s13, ne10_padded, + ne11, ne12, ne13, stream); + ++ } else if (src1_bf16) { ++ quantize_mmq_q8_1_cuda(src1_bf16, nullptr, src1_q8_1.get(), src0->type, ne10, s11, s12, s13, ne10_padded, ++ ne11, ne12, ne13, stream); ++ } else if (src1_f16) { ++ quantize_mmq_q8_1_cuda(src1_f16, nullptr, src1_q8_1.get(), src0->type, ne10, s11, s12, s13, ne10_padded, ++ ne11, ne12, ne13, stream); + } else { + quantize_mmq_q8_1_cuda(src1_d, nullptr, src1_q8_1.get(), src0->type, ne10, s11, s12, s13, ne10_padded, + ne11, ne12, ne13, stream); +@@ -299,6 +311,17 @@ + CUDA_CHECK(cudaGetLastError()); + } + ++ // Every expert's column tiles fit in ntx_flat grid columns, so a flat grid needs no host sync to size it. ++ ggml_cuda_pool_alloc tile_map(ctx.pool()); ++ int ntx_flat = 0; ++ if (ncols_max <= 0) { ++ const int jt_none = (ne_get_rows + J_best - 1) / J_best; ++ ntx_flat = jt_none + std::min(ne02, ne_get_rows); ++ tile_map.alloc(ntx_flat); ++ ggml_cuda_launch_mm_ids_tile_map(expert_bounds.get(), tile_map.get(), ne02, J_best, ntx_flat, jt_none, stream); ++ CUDA_CHECK(cudaGetLastError()); ++ } ++ + const size_t nbytes_src1_q8_1 = ne12*n_expert_used*ne10_padded * y_block_size/y_values_per_block + src1_q8_1_padding; + ggml_cuda_pool_alloc src1_q8_1(ctx.pool(), nbytes_src1_q8_1); + ggml_cuda_pool_alloc src1_scale(ctx.pool()); +@@ -325,9 +348,21 @@ + quantize_mmq_fp4_cuda(src1_d, ids_src1.get(), src1_q8_1.get(), src1_scale.ptr, src0->type, use_aligned_float8, ne10, s11, s12, s13, + ne10_padded, ne11_flat, ne12_flat, ne13_flat, stream); + } ++ } else if (dedup_bcast && src1_bf16) { ++ quantize_scatter_mmq_q8_1_cuda(src1_bf16, ids_src1.get(), src1_q8_1.get(), src0->type, ne10, ++ /*stride_token=*/s12, ne10_padded, ne12, ne11_flat, n_expert_used, stream); ++ } else if (dedup_bcast && src1_f16) { ++ quantize_scatter_mmq_q8_1_cuda(src1_f16, ids_src1.get(), src1_q8_1.get(), src0->type, ne10, ++ /*stride_token=*/s12, ne10_padded, ne12, ne11_flat, n_expert_used, stream); + } else if (dedup_bcast) { + quantize_scatter_mmq_q8_1_cuda(src1_d, ids_src1.get(), src1_q8_1.get(), src0->type, ne10, + /*stride_token=*/s12, ne10_padded, ne12, ne11_flat, n_expert_used, stream); ++ } else if (src1_bf16) { ++ quantize_mmq_q8_1_cuda(src1_bf16, ids_src1.get(), src1_q8_1.get(), src0->type, ne10, s11, s12, s13, ++ ne10_padded, ne11_flat, ne12_flat, ne13_flat, stream); ++ } else if (src1_f16) { ++ quantize_mmq_q8_1_cuda(src1_f16, ids_src1.get(), src1_q8_1.get(), src0->type, ne10, s11, s12, s13, ++ ne10_padded, ne11_flat, ne12_flat, ne13_flat, stream); + } else { + quantize_mmq_q8_1_cuda(src1_d, ids_src1.get(), src1_q8_1.get(), src0->type, ne10, s11, s12, s13, + ne10_padded, ne11_flat, ne12_flat, ne13_flat, stream); +@@ -347,7 +382,7 @@ + ne00, ne01, ne_get_rows, s01, ne_get_rows, s1, + ne02, ne02, s02, s12, s2, + ne03, ne13, s03, s13, s3, +- ne12, J_best}; ++ ncols_max > 0 ? ncols_max : ne12, J_best, tile_map.ptr, ntx_flat}; + + ggml_cuda_mul_mat_q_switch_type(ctx, args, stream, prec_src1); + } +--- a/ggml/src/ggml-cuda/mmq.cuh ++++ b/ggml/src/ggml-cuda/mmq.cuh +@@ -10,6 +10,17 @@ + #define MMQ_ITER_K_FP4 512 + #define MMQ_NWARPS 8 + ++// FreeToken: MMQ writes fp32, bf16 or fp16 results (dst_t); the sums stay fp32 until the store ++template static __device__ __forceinline__ T mmq_cast(float v); ++template <> __device__ __forceinline__ float mmq_cast(float v) { return v; } ++template <> __device__ __forceinline__ nv_bfloat16 mmq_cast(float v) { return __float2bfloat16(v); } ++template <> __device__ __forceinline__ half mmq_cast(float v) { return __float2half(v); } ++ ++struct mmq_dst { ++ void * p; ++ ggml_type type; ++}; ++ + typedef void (*ggml_cuda_mmq_load_tiles_t)(const char * __restrict__ x, int * x_tile, const int kbx0, const int i_max, const int stride); + typedef void (*ggml_cuda_mmq_vec_dot_t)(const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00); + typedef void (*ggml_cuda_mmq_write_back_t)(const float * __restrict__ sum, const int32_t * __restrict__ get_rows_to_sorted, +@@ -429,8 +440,8 @@ + #include "mmq-load-tiles.cuh" + #include "mmq-vec-dot.cuh" + +-template static __device__ __forceinline__ void ggml_cuda_mmq_write_back_dp4a( +- const float * __restrict__ sum, const int32_t * __restrict__ ids_dst, float * __restrict__ dst, ++template static __device__ __forceinline__ void ggml_cuda_mmq_write_back_dp4a( ++ const float * __restrict__ sum, const int32_t * __restrict__ ids_dst, dst_t * __restrict__ dst, + const float * __restrict__ y_scale, const int stride, const int i_max, const int j_max) { + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback, prec_src1) / warp_size; +@@ -456,21 +467,21 @@ + + if constexpr (type == GGML_TYPE_NVFP4) { + if (y_scale_used) { +- dst[ids_dst[j]*stride + i] = y_scale[j] * sum[(j0/nwarps) * (I/warp_size) + i0/warp_size]; ++ dst[ids_dst[j]*stride + i] = mmq_cast(y_scale[j] * sum[(j0/nwarps) * (I/warp_size) + i0/warp_size]); + } else { +- dst[ids_dst[j]*stride + i] = sum[(j0/nwarps) * (I/warp_size) + i0/warp_size]; ++ dst[ids_dst[j]*stride + i] = mmq_cast(sum[(j0/nwarps) * (I/warp_size) + i0/warp_size]); + } + } else { +- dst[ids_dst[j]*stride + i] = sum[(j0/nwarps) * (I/warp_size) + i0/warp_size]; ++ dst[ids_dst[j]*stride + i] = mmq_cast(sum[(j0/nwarps) * (I/warp_size) + i0/warp_size]); + GGML_UNUSED(y_scale_used); + } + } + } + } + +-template ++template + static __device__ __forceinline__ void ggml_cuda_mmq_write_back_mma( +- const float * __restrict__ sum, const int * __restrict__ ids_dst, float * __restrict__ dst, ++ const float * __restrict__ sum, const int * __restrict__ ids_dst, dst_t * __restrict__ dst, + const float * __restrict__ y_scale, const int stride, const int i_max, const int j_max) { + + #if defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) +@@ -506,12 +517,12 @@ + + if constexpr (type == GGML_TYPE_NVFP4) { + if (y_scale_used) { +- dst[ids_dst[j]*stride + i] = y_scale[j] * sum[(j0/tile_C::J + n)*tile_C::ne + l]; ++ dst[ids_dst[j]*stride + i] = mmq_cast(y_scale[j] * sum[(j0/tile_C::J + n)*tile_C::ne + l]); + } else { +- dst[ids_dst[j]*stride + i] = sum[(j0/tile_C::J + n)*tile_C::ne + l]; ++ dst[ids_dst[j]*stride + i] = mmq_cast(sum[(j0/tile_C::J + n)*tile_C::ne + l]); + } + } else { +- dst[ids_dst[j]*stride + i] = sum[(j0/tile_C::J + n)*tile_C::ne + l]; ++ dst[ids_dst[j]*stride + i] = mmq_cast(sum[(j0/tile_C::J + n)*tile_C::ne + l]); + GGML_UNUSED(y_scale_used); + } + } +@@ -869,10 +880,10 @@ + + // --------------------------------------------------------------------------------------------- + +-template ++template + static __device__ __forceinline__ void mul_mat_q_process_tile( + const char * __restrict__ x, const int offset_x, const int * __restrict__ y, +- const int * __restrict__ ids_dst, float * __restrict__ dst, float * __restrict__ tmp_fixup, ++ const int * __restrict__ ids_dst, dst_t * __restrict__ dst, float * __restrict__ tmp_fixup, + const float * __restrict__ y_scale, + const int stride_row_x, const int ncols_y, const int stride_col_dst, + const int tile_x_max_i, const int tile_y_max_j, const int kb0_start, const int kb0_stop) { +@@ -942,19 +953,32 @@ + + if (fixup) { + write_back(sum, ids_dst, tmp_fixup + blockIdx.x*(J*I), y_scale, I, I, J); +- } else { ++ } else if constexpr (std::is_same_v) { + write_back(sum, ids_dst, dst, y_scale, stride_col_dst, tile_x_max_i, tile_y_max_j); ++ } else if constexpr (ggml_cuda_mmq_get_config(type, J, fallback, prec_src1).use_mma_data_layout()) { ++ ggml_cuda_mmq_write_back_mma(sum, ids_dst, dst, y_scale, stride_col_dst, tile_x_max_i, tile_y_max_j); ++ } else { ++ ggml_cuda_mmq_write_back_dp4a(sum, ids_dst, dst, y_scale, stride_col_dst, tile_x_max_i, tile_y_max_j); + } + } + + ++// A flat MoE grid runs only the (expert, column tile) pairs listed in tile_map, see ggml_cuda_launch_mm_ids_tile_map. ++static __device__ __forceinline__ void mmq_flat_tile(const int2 * __restrict__ tile_map, int & zt, int & jt) { ++ if (tile_map) { ++ const int2 t = tile_map[jt]; ++ zt = t.x; ++ jt = t.y; ++ } ++} ++ + // The mul_mat_q kernel implements "stream-k" work partitioning as described in https://arxiv.org/abs/2301.03598 + +-template ++template + __launch_bounds__(ggml_cuda_mmq_get_nthreads(type, J, fallback, prec_src1), ggml_cuda_mmq_get_occupancy(type, J, fallback, prec_src1)) + static __global__ void mul_mat_q( + const char * __restrict__ x, const int * __restrict__ y, const int32_t * __restrict__ ids_dst, +- const int32_t * __restrict__ expert_bounds, float * __restrict__ dst, float * __restrict__ tmp_fixup, ++ const int32_t * __restrict__ expert_bounds, const int2 * __restrict__ tile_map, dst_t * __restrict__ dst, float * __restrict__ tmp_fixup, + const float * __restrict__ y_scale, + const uint3 blocks_per_ne00, const int nrows_x, const int ncols_dst, const int stride_row_x, const int ncols_y, const int stride_col_dst, + const uint3 channel_ratio, const uint3 nchannels_y, const int stride_channel_x, const int stride_channel_y, const int stride_channel_dst, +@@ -993,9 +1017,10 @@ + if constexpr (!ggml_cuda_mmq_get_stream_k(type, J, fallback, prec_src1)) { + const uint2 tmp2 = fast_div_modulo(blockIdx.z, nchannels_y); + const int wt = tmp2.x; +- const int zt = tmp2.y; +- const int jt = blockIdx.y; ++ int zt = tmp2.y; ++ int jt = blockIdx.y; + const int it = blockIdx.x; ++ mmq_flat_tile(tile_map, zt, jt); + + // Defaults for regular matrix multiplication: + int col_low = 0; +@@ -1076,14 +1101,15 @@ + while (kbc < kbc_stop && kb0_stop == int(blocks_per_ne00.z)) { + int tmp = fastdiv(kbc, blocks_per_ne00); + uint2 tmp2 = fast_div_modulo(tmp, ntx); +- const int jt = tmp2.y; ++ int jt = tmp2.y; + tmp = tmp2.x; + tmp2 = fast_div_modulo(tmp, nchannels_y); +- const int zt = tmp2.y; ++ int zt = tmp2.y; + tmp = tmp2.x; + tmp2 = fast_div_modulo(tmp, nsamples_y); + const int wt = tmp2.y; + const int it = tmp2.x; ++ mmq_flat_tile(tile_map, zt, jt); + + // Defaults for regular matrix multiplication: + int col_low = 0; +@@ -1165,14 +1191,15 @@ + + int tmp = fastdiv(kbc, blocks_per_ne00); + uint2 tmp2 = fast_div_modulo(tmp, ntx); +- const int jt = tmp2.y; ++ int jt = tmp2.y; + tmp = tmp2.x; + tmp2 = fast_div_modulo(tmp, nchannels_y); +- const int zt = tmp2.y; ++ int zt = tmp2.y; + tmp = tmp2.x; + tmp2 = fast_div_modulo(tmp, nsamples_y); + const int wt = tmp2.y; + const int it = tmp2.x; ++ mmq_flat_tile(tile_map, zt, jt); + + // Defaults for regular matrix multiplication: + int col_low = 0; +@@ -1375,14 +1402,26 @@ + } + } + ++// FreeToken: copies an fp32 result with dst's layout into dst ++template ++static __global__ void mmq_store_dst(const float * __restrict__ src, dst_t * __restrict__ dst, const int nrows, const int stride) { ++ const int64_t col = blockIdx.x; ++ const int row = blockIdx.y*blockDim.x + threadIdx.x; ++ if (row < nrows) { ++ dst[col*stride + row] = mmq_cast(src[col*stride + row]); ++ } ++} ++ + struct mmq_args { +- const char * x; ggml_type type_x; const int * y; const int32_t * ids_dst; const int32_t * expert_bounds; float * dst; ++ const char * x; ggml_type type_x; const int * y; const int32_t * ids_dst; const int32_t * expert_bounds; mmq_dst dst; + const float * y_scale; + int64_t ncols_x; int64_t nrows_x; int64_t ncols_dst; int64_t stride_row_x; int64_t ncols_y; int64_t nrows_dst; + int64_t nchannels_x; int64_t nchannels_y; int64_t stride_channel_x; int64_t stride_channel_y; int64_t stride_channel_dst; + int64_t nsamples_x; int64_t nsamples_y; int64_t stride_sample_x; int64_t stride_sample_y; int64_t stride_sample_dst; + int64_t ncols_max; + int J_best; // Tile width in ne11(dense)/ne12(MoE) direction to use for optimal performance. ++ const int2 * tile_map = nullptr; // Flat MoE grid: the (expert, column tile) of each of the ntx_flat grid columns. ++ int ntx_flat = 0; + }; + + static size_t mmq_get_nbytes_shared(const ggml_cuda_mmq_config & config, const int cc) { +@@ -1406,12 +1445,11 @@ + + const dim3 block_dims(warp_size, nwarps, 1); + +- CUDA_SET_SHARED_MEMORY_LIMIT((mul_mat_q), nbytes_shared); +- CUDA_SET_SHARED_MEMORY_LIMIT((mul_mat_q), nbytes_shared); +- ++ // a flat MoE grid puts all experts into one channel of ntx_flat column tiles ++ const int64_t nchannels_grid = args.tile_map ? 1 : args.nchannels_y; + const int nty = (args.nrows_x + config.I - 1) / config.I; +- const int ntx = (args.ncols_max + config.J - 1) / config.J; +- const int ntzw = args.nchannels_y * args.nsamples_y; ++ const int ntx = args.tile_map ? args.ntx_flat : (args.ncols_max + config.J - 1) / config.J; ++ const int ntzw = nchannels_grid * args.nsamples_y; + const dim3 block_nums_xy_tiling(nty, ntx, ntzw); + + GGML_ASSERT(args.nchannels_y % args.nchannels_x == 0); +@@ -1421,18 +1459,71 @@ + + const uint3 blocks_per_ne00_fd = init_fastdiv_values(args.ncols_x / ggml_cuda_type_traits::qk); + const uint3 ntx_fd = init_fastdiv_values(ntx); +- const uint3 nchannels_y_fd = init_fastdiv_values(args.nchannels_y); ++ const uint3 nchannels_y_fd = init_fastdiv_values(nchannels_grid); + const uint3 nsamples_y_fd = init_fastdiv_values(args.nsamples_y); + const uint3 channel_ratio_fd = init_fastdiv_values(channel_ratio); + const uint3 sample_ratio_fd = init_fastdiv_values(sample_ratio); + +- if (!config.stream_k) { +- mul_mat_q<<>> +- (args.x, args.y, args.ids_dst, args.expert_bounds, args.dst, nullptr, args.y_scale, ++ const auto run = [&](auto * dst, const dim3 & block_nums, float * tmp_fixup) { ++ using dst_t = std::remove_pointer_t; ++ CUDA_SET_SHARED_MEMORY_LIMIT((mul_mat_q), nbytes_shared); ++ CUDA_SET_SHARED_MEMORY_LIMIT((mul_mat_q), nbytes_shared); ++ mul_mat_q<<>> ++ (args.x, args.y, args.ids_dst, args.expert_bounds, args.tile_map, dst, tmp_fixup, args.y_scale, + blocks_per_ne00_fd, args.nrows_x, args.ncols_dst, args.stride_row_x, args.ncols_y, args.nrows_dst, + channel_ratio_fd, nchannels_y_fd, args.stride_channel_x, args.stride_channel_y, args.stride_channel_dst, + sample_ratio_fd, nsamples_y_fd, args.stride_sample_x, args.stride_sample_y, args.stride_sample_dst, + ntx_fd); ++ }; ++ ++ // FreeToken: the fixup kernel adds partial sums into an fp32 dst, so a bf16 dst that needs it is computed in fp32 first. ++ // An fp16 dst always is, unless GGML_MMQ_F16_DST builds the fp16 kernels. ++ const auto launch = [&](const dim3 & block_nums, const bool fixup, float * tmp_fixup) { ++ bool via_f32 = fixup && args.dst.type != GGML_TYPE_F32; ++#ifndef GGML_MMQ_F16_DST ++ via_f32 = via_f32 || args.dst.type == GGML_TYPE_F16; ++#endif // GGML_MMQ_F16_DST ++ ggml_cuda_pool_alloc dst_f32(ctx.pool(id)); ++ float * dst_float = args.dst.type == GGML_TYPE_F32 ? (float *) args.dst.p : nullptr; ++ if (via_f32) { ++ GGML_ASSERT(args.ids_dst || (args.nchannels_y == 1 && args.nsamples_y == 1)); ++ dst_float = dst_f32.alloc(args.ncols_dst * args.nrows_dst); ++ } ++ ++ if (dst_float) { ++ run(dst_float, block_nums, tmp_fixup); ++ } else if (args.dst.type == GGML_TYPE_BF16) { ++ run((nv_bfloat16 *) args.dst.p, block_nums, tmp_fixup); ++ } ++#ifdef GGML_MMQ_F16_DST ++ else { ++ run((half *) args.dst.p, block_nums, tmp_fixup); ++ } ++#endif // GGML_MMQ_F16_DST ++ ++ if (fixup) { ++ CUDA_CHECK(cudaGetLastError()); ++ const dim3 block_nums_fixup(block_nums.x, config.I/warp_size, 1); ++ const dim3 block_dims_fixup(block_dims.x, block_dims.y/2, block_dims.z); ++ mul_mat_q_stream_k_fixup<<>> ++ (args.ids_dst, args.expert_bounds, dst_float, tmp_fixup, blocks_per_ne00_fd, args.nrows_x, args.ncols_dst, ++ args.nrows_dst, nchannels_y_fd, args.stride_channel_dst, nsamples_y_fd, args.stride_sample_dst, ++ ntx_fd); ++ } ++ ++ if (via_f32) { ++ CUDA_CHECK(cudaGetLastError()); ++ const dim3 block_nums_store(args.ncols_dst, (args.nrows_x + 255) / 256, 1); ++ if (args.dst.type == GGML_TYPE_BF16) { ++ mmq_store_dst<<>>(dst_float, (nv_bfloat16 *) args.dst.p, args.nrows_x, args.nrows_dst); ++ } else { ++ mmq_store_dst<<>>(dst_float, (half *) args.dst.p, args.nrows_x, args.nrows_dst); ++ } ++ } ++ }; ++ ++ if (!config.stream_k) { ++ launch(block_nums_xy_tiling, false, nullptr); + return; + } + +@@ -1441,7 +1532,8 @@ + const int ntiles_dst = ntx * nty * ntzw; + const int tiles_nwaves = (ntiles_dst + nsm - 1) / nsm; + const int tiles_efficiency_percent = 100 * ntiles_dst / (nsm*tiles_nwaves); +- const dim3 block_nums_stream_k(GGML_CUDA_CC_IS_NVIDIA(cc) && tiles_efficiency_percent >= 90 ? ntiles_dst : nsm, 1, 1); ++ // The fixup kernel does not read tile_map, so a flat MoE grid always runs one tile per block. ++ const dim3 block_nums_stream_k(args.tile_map || (GGML_CUDA_CC_IS_NVIDIA(cc) && tiles_efficiency_percent >= 90) ? ntiles_dst : nsm, 1, 1); + + GGML_ASSERT(ntiles_dst * blocks_per_ne00_fd.z < (1 << 30)); // Assert that variable kbc will not overflow. + +@@ -1453,25 +1545,7 @@ + tmp_fixup.alloc(block_nums_stream_k.x * config.J*config.I); + } + +- const dim3 block_nums_fixup(block_nums_stream_k.x, config.I/warp_size, 1); +- const dim3 block_dims_fixup(block_dims.x, block_dims.y/2, block_dims.z); +- +- mul_mat_q<<>> +- (args.x, args.y, args.ids_dst, args.expert_bounds, args.dst, tmp_fixup.ptr, args.y_scale, +- blocks_per_ne00_fd, args.nrows_x, args.ncols_dst, args.stride_row_x, args.ncols_y, args.nrows_dst, +- channel_ratio_fd, nchannels_y_fd, args.stride_channel_x, args.stride_channel_y, args.stride_channel_dst, +- sample_ratio_fd, nsamples_y_fd, args.stride_sample_x, args.stride_sample_y, args.stride_sample_dst, +- ntx_fd); +- +- if (!fixup_needed) { +- return; +- } +- +- CUDA_CHECK(cudaGetLastError()); +- mul_mat_q_stream_k_fixup<<>> +- (args.ids_dst, args.expert_bounds, args.dst, tmp_fixup.ptr, blocks_per_ne00_fd, args.nrows_x, args.ncols_dst, +- args.nrows_dst, nchannels_y_fd, args.stride_channel_dst, nsamples_y_fd, args.stride_sample_dst, +- ntx_fd); ++ launch(block_nums_stream_k, fixup_needed, tmp_fixup.ptr); + } + + template +@@ -1579,7 +1653,9 @@ + + // ------------------------------------------------------------------------------------------------------------------------- + ++// FreeToken: under ids, ncols_max > 0 bounds the tokens any one expert gets and sizes the tile grid; otherwise the grid is flat, see ggml_cuda_launch_mm_ids_tile_map + void ggml_cuda_mul_mat_q( +- ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, ggml_tensor * dst); ++ ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, ggml_tensor * dst, ++ int64_t ncols_max = -1); + + bool ggml_cuda_should_use_mmq(enum ggml_type type, int cc, int64_t ne11, int64_t n_experts); +--- a/ggml/src/ggml-cuda/mmvq.cu ++++ b/ggml/src/ggml-cuda/mmvq.cu +@@ -1,4 +1,5 @@ + #include "mmvq.cuh" ++#include "convert.cuh" + #include "quantize.cuh" + #include "unary.cuh" + #include "vecdotq.cuh" +@@ -596,10 +597,11 @@ + return 1; + } + +-template ++// FreeToken: dst_t is float, nv_bfloat16 or half ++template + __launch_bounds__(calc_nwarps(type, ncols_dst, get_device_table_id(), small_k, halve_iters)*ggml_cuda_get_physical_warp_size(), 1) + static __global__ void mul_mat_vec_q( +- const void * vx_ptr, const void * vy_ptr, const int32_t * ids_ptr, const ggml_cuda_mm_fusion_args_device fusion, float * dst_ptr, ++ const void * vx_ptr, const void * vy_ptr, const int32_t * ids_ptr, const ggml_cuda_mm_fusion_args_device fusion, dst_t * dst_ptr, + const uint32_t ncols_x, const uint3 nchannels_y, const uint32_t stride_row_x, const uint32_t stride_col_y, + uint32_t stride_col_dst, const uint3 channel_ratio, const uint32_t stride_channel_x, + const uint32_t stride_channel_y, const uint32_t stride_channel_dst, const uint3 sample_ratio, +@@ -608,7 +610,7 @@ + const void * GGML_CUDA_RESTRICT vx = vx_ptr; + const void * GGML_CUDA_RESTRICT vy = vy_ptr; + const int32_t * GGML_CUDA_RESTRICT ids = ids_ptr; +- float * GGML_CUDA_RESTRICT dst = dst_ptr; ++ dst_t * GGML_CUDA_RESTRICT dst = dst_ptr; + + constexpr int qk = ggml_cuda_type_traits::qk; + constexpr int qi = ggml_cuda_type_traits::qi; +@@ -629,7 +631,9 @@ + const uint32_t channel_dst = shared_expert ? 0 : blockIdx.y; + if (shared_expert) { + vx = fusion.shared_up; +- dst = fusion.shared_dst; ++ if constexpr (std::is_same_v) { ++ dst = fusion.shared_dst; ++ } + stride_col_dst = fusion.shared_stride_col_dst; + } + +@@ -702,10 +706,10 @@ + } + if constexpr (type == GGML_TYPE_NVFP4) { + if (use_scale) { +- x_scales = x_scale[ids ? channel_x : 0]; ++ x_scales = x_scale[(ids ? channel_x : 0) * fusion.x_scale_stride]; + } + if (use_gate_scale) { +- gate_scales = gate_scale[ids ? channel_x : 0]; ++ gate_scales = gate_scale[(ids ? channel_x : 0) * fusion.gate_scale_stride]; + } + } + } +@@ -837,7 +841,7 @@ + } + } + } +- dst[j*stride_col_dst + i] = result; ++ dst[j*stride_col_dst + i] = ggml_cuda_cast(result); + } + } + } +@@ -854,11 +858,11 @@ + // Grid: (ceil(nrows_x / c_rows_per_block), nchannels_dst) + // Block: (warp_size, ncols_dst) - each warp handles one token independently. + // No shared memory reduction needed since each warp works alone. +-template ++template + __launch_bounds__(get_mmvq_mmid_max_batch_for_device()*ggml_cuda_get_physical_warp_size(), 1) + static __global__ void mul_mat_vec_q_moe( + const void * vx_ptr, const void * vy_ptr, const int32_t * ids_ptr, const ggml_cuda_mm_fusion_args_device fusion, +- float * dst_ptr, ++ dst_t * dst_ptr, + const uint32_t ncols_x, const uint3 nchannels_y, const uint32_t nrows_x, + const uint32_t stride_row_x, const uint32_t stride_col_y, uint32_t stride_col_dst, + const uint32_t stride_channel_x, const uint32_t stride_channel_y, const uint32_t stride_channel_dst, +@@ -866,7 +870,7 @@ + const void * GGML_CUDA_RESTRICT vx = vx_ptr; + const void * GGML_CUDA_RESTRICT vy = vy_ptr; + const int32_t * GGML_CUDA_RESTRICT ids = ids_ptr; +- float * GGML_CUDA_RESTRICT dst = dst_ptr; ++ dst_t * GGML_CUDA_RESTRICT dst = dst_ptr; + + constexpr int qk = ggml_cuda_type_traits::qk; + constexpr int qi = ggml_cuda_type_traits::qi; +@@ -878,7 +882,9 @@ + const bool shared_expert = has_fusion && fusion.shared_up && blockIdx.y == gridDim.y - 1; + if (shared_expert) { + vx = fusion.shared_up; +- dst = fusion.shared_dst; ++ if constexpr (std::is_same_v) { ++ dst = fusion.shared_dst; ++ } + stride_col_dst = fusion.shared_stride_col_dst; + } + +@@ -963,7 +969,7 @@ + + if constexpr (type == GGML_TYPE_NVFP4) { + if (x_scale) { +- result *= x_scale[channel_x]; ++ result *= x_scale[channel_x * fusion.x_scale_stride]; + } + } + if (x_bias) { +@@ -973,7 +979,7 @@ + float gate_value = tmp_gate[threadIdx.x]; + if constexpr (type == GGML_TYPE_NVFP4) { + if (gate_scale) { +- gate_value *= gate_scale[channel_x]; ++ gate_value *= gate_scale[channel_x * fusion.gate_scale_stride]; + } + } + if (gate_bias) { +@@ -998,7 +1004,7 @@ + } + } + } +- dst[channel_dst*stride_channel_dst + token_idx*stride_col_dst + row0 + threadIdx.x] = result; ++ dst[channel_dst*stride_channel_dst + token_idx*stride_col_dst + row0 + threadIdx.x] = ggml_cuda_cast(result); + } + + if constexpr (!has_fusion) { +@@ -1020,9 +1026,9 @@ + return {block_nums, block_dims}; + } + +-template ++template + static void mul_mat_vec_q_switch_fusion( +- const void * vx, const void * vy, const int32_t * ids, const ggml_cuda_mm_fusion_args_device fusion, float * dst, ++ const void * vx, const void * vy, const int32_t * ids, const ggml_cuda_mm_fusion_args_device fusion, dst_t * dst, + const uint32_t ncols_x, const uint3 nchannels_y, const uint32_t stride_row_x, const uint32_t stride_col_y, + const uint32_t stride_col_dst, const uint3 channel_ratio, const uint32_t stride_channel_x, + const uint32_t stride_channel_y, const uint32_t stride_channel_dst, const uint3 sample_ratio, +@@ -1035,7 +1041,7 @@ + if constexpr (c_ncols_dst == 1) { + if (has_fusion) { + const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(block_nums, block_dims, nbytes_shared, stream); +- ggml_cuda_kernel_launch(mul_mat_vec_q, launch_params, ++ ggml_cuda_kernel_launch(mul_mat_vec_q, launch_params, + vx, vy, ids, fusion, dst, ncols_x, nchannels_y, stride_row_x, stride_col_y, stride_col_dst, + channel_ratio, stride_channel_x, stride_channel_y, stride_channel_dst, + sample_ratio, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride); +@@ -1046,15 +1052,15 @@ + GGML_ASSERT(!has_fusion && "fusion only supported for ncols_dst=1"); + + const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(block_nums, block_dims, nbytes_shared, stream); +- ggml_cuda_kernel_launch(mul_mat_vec_q, launch_params, ++ ggml_cuda_kernel_launch(mul_mat_vec_q, launch_params, + vx, vy, ids, fusion, dst, ncols_x, nchannels_y, stride_row_x, stride_col_y, stride_col_dst, + channel_ratio, stride_channel_x, stride_channel_y, stride_channel_dst, + sample_ratio, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride); + } + +-template ++template + static void mul_mat_vec_q_moe_launch( +- const void * vx, const void * vy, const int32_t * ids, const ggml_cuda_mm_fusion_args_device fusion, float * dst, ++ const void * vx, const void * vy, const int32_t * ids, const ggml_cuda_mm_fusion_args_device fusion, dst_t * dst, + const uint32_t ncols_x, const uint3 nchannels_y, const uint32_t nrows_x, + const uint32_t stride_row_x, const uint32_t stride_col_y, const uint32_t stride_col_dst, + const uint32_t stride_channel_x, const uint32_t stride_channel_y, const uint32_t stride_channel_dst, +@@ -1071,13 +1077,13 @@ + fusion.x_scale != nullptr || fusion.gate_scale != nullptr; + + if (has_fusion) { +- ggml_cuda_kernel_launch(mul_mat_vec_q_moe, launch_params, ++ ggml_cuda_kernel_launch(mul_mat_vec_q_moe, launch_params, + vx, vy, ids, fusion, dst, ncols_x, nchannels_y, nrows_x, + stride_row_x, stride_col_y, stride_col_dst, + stride_channel_x, stride_channel_y, stride_channel_dst, + ncols_dst, ids_stride); + } else { +- ggml_cuda_kernel_launch(mul_mat_vec_q_moe, launch_params, ++ ggml_cuda_kernel_launch(mul_mat_vec_q_moe, launch_params, + vx, vy, ids, fusion, dst, ncols_x, nchannels_y, nrows_x, + stride_row_x, stride_col_y, stride_col_dst, + stride_channel_x, stride_channel_y, stride_channel_dst, +@@ -1085,9 +1091,9 @@ + } + } + +-template ++template + static void mul_mat_vec_q_switch_ncols_dst( +- const void * vx, const void * vy, const int32_t * ids, const ggml_cuda_mm_fusion_args_device fusion, float * dst, ++ const void * vx, const void * vy, const int32_t * ids, const ggml_cuda_mm_fusion_args_device fusion, dst_t * dst, + const int ncols_x, const int nrows_x, const int ncols_dst, + const int stride_row_x, const int stride_col_y, const int stride_col_dst, + const int nchannels_x, const int nchannels_y, const int nchannels_dst, +@@ -1277,8 +1283,10 @@ + GGML_ABORT("fatal error"); + } + } ++ ++template + static void mul_mat_vec_q_switch_type( +- const void * vx, const ggml_type type_x, const void * vy, const int32_t * ids, const ggml_cuda_mm_fusion_args_device fusion, float * dst, ++ const void * vx, const ggml_type type_x, const void * vy, const int32_t * ids, const ggml_cuda_mm_fusion_args_device fusion, dst_t * dst, + const int ncols_x, const int nrows_x, const int ncols_dst, + const int stride_row_x, const int stride_col_y, const int stride_col_dst, + const int nchannels_x, const int nchannels_y, const int nchannels_dst, +@@ -1431,9 +1439,10 @@ + + void ggml_cuda_mul_mat_vec_q( + ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, ggml_tensor * dst, +- const ggml_cuda_mm_fusion_args_host * fusion) { +- GGML_ASSERT( src1->type == GGML_TYPE_F32); +- GGML_ASSERT( dst->type == GGML_TYPE_F32); ++ const ggml_cuda_mm_fusion_args_host * fusion, const void * src1_q8_1_pre) { ++ // FreeToken: activations and results may also be bf16 or fp16 ++ GGML_ASSERT( src1->type == GGML_TYPE_F32 || src1->type == GGML_TYPE_BF16 || src1->type == GGML_TYPE_F16); ++ GGML_ASSERT( dst->type == GGML_TYPE_F32 || dst->type == GGML_TYPE_BF16 || dst->type == GGML_TYPE_F16); + GGML_ASSERT(!ids || ids->type == GGML_TYPE_I32); // Optional, used for batched GGML_MUL_MAT_ID. + + GGML_TENSOR_BINARY_OP_LOCALS; +@@ -1498,17 +1507,20 @@ + GGML_ASSERT(!ids || fusion->gate_bias->ne[1] == src0->ne[2]); + fusion_local.gate_bias = fusion->gate_bias->data; + } ++ // FreeToken: a scale may be a strided vector, read with its element stride + if (fusion->x_scale) { + GGML_ASSERT(fusion->x_scale->type == GGML_TYPE_F32); +- GGML_ASSERT(ggml_is_contiguous(fusion->x_scale)); ++ GGML_ASSERT(ggml_nelements(fusion->x_scale) == fusion->x_scale->ne[0] && fusion->x_scale->nb[0] % sizeof(float) == 0); + GGML_ASSERT(ggml_nelements(fusion->x_scale) == (ids ? src0->ne[2] : 1)); + fusion_local.x_scale = fusion->x_scale->data; ++ fusion_local.x_scale_stride = fusion->x_scale->nb[0] / sizeof(float); + } + if (fusion->gate_scale) { + GGML_ASSERT(fusion->gate_scale->type == GGML_TYPE_F32); +- GGML_ASSERT(ggml_is_contiguous(fusion->gate_scale)); ++ GGML_ASSERT(ggml_nelements(fusion->gate_scale) == fusion->gate_scale->ne[0] && fusion->gate_scale->nb[0] % sizeof(float) == 0); + GGML_ASSERT(ggml_nelements(fusion->gate_scale) == (ids ? src0->ne[2] : 1)); + fusion_local.gate_scale = fusion->gate_scale->data; ++ fusion_local.gate_scale_stride = fusion->gate_scale->nb[0] / sizeof(float); + } + fusion_local.glu_op = fusion->glu_op; + fusion_local.glu_limit = fusion->glu_limit; +@@ -1526,12 +1538,22 @@ + } + + const int64_t ne10_padded = GGML_PAD(ne10, MATRIX_ROW_PADDING); +- ggml_cuda_pool_alloc src1_q8_1(ctx.pool(), ne13*ne12 * ne11*ne10_padded * sizeof(block_q8_1)/QK8_1); +- { ++ // FreeToken: src1_q8_1_pre is src1 already quantized by quantize_row_q8_1_cuda with this padding and layout ++ ggml_cuda_pool_alloc src1_q8_1(ctx.pool()); ++ const void * src1_q8 = src1_q8_1_pre; ++ if (src1_q8 == nullptr) { ++ src1_q8_1.alloc(ne13*ne12 * ne11*ne10_padded * sizeof(block_q8_1)/QK8_1); ++ src1_q8 = src1_q8_1.get(); + const int64_t s11 = src1->nb[1] / ts_src1; + const int64_t s12 = src1->nb[2] / ts_src1; + const int64_t s13 = src1->nb[3] / ts_src1; +- quantize_row_q8_1_cuda(src1_d, nullptr, src1_q8_1.get(), src0->type, ne10, s11, s12, s13, ne10_padded, ne11, ne12, ne13, stream); ++ if (src1->type == GGML_TYPE_BF16) { ++ quantize_row_q8_1_cuda((const nv_bfloat16 *) src1->data, nullptr, src1_q8_1.get(), src0->type, ne10, s11, s12, s13, ne10_padded, ne11, ne12, ne13, stream); ++ } else if (src1->type == GGML_TYPE_F16) { ++ quantize_row_q8_1_cuda((const half *) src1->data, nullptr, src1_q8_1.get(), src0->type, ne10, s11, s12, s13, ne10_padded, ne11, ne12, ne13, stream); ++ } else { ++ quantize_row_q8_1_cuda(src1_d, nullptr, src1_q8_1.get(), src0->type, ne10, s11, s12, s13, ne10_padded, ne11, ne12, ne13, stream); ++ } + } + + const int64_t s01 = src0->nb[1] / ts_src0; +@@ -1556,11 +1578,23 @@ + + const int64_t ids_stride = ids ? ids->nb[1] / ggml_type_size(ids->type) : 0; + +- mul_mat_vec_q_switch_type( +- src0->data, src0->type, src1_q8_1.get(), ids_d, fusion_local, dst_d, ne00, +- ne01, ncols_dst, s01, stride_col_y, stride_col_dst, +- ne02, nchannels_y, nchannels_dst, s02, stride_channel_y, stride_channel_dst, +- ne03, ne3, s03, s13, s3, ids_stride, stream); ++ const auto launch = [&](auto * dst_ptr) { ++ mul_mat_vec_q_switch_type( ++ src0->data, src0->type, src1_q8, ids_d, fusion_local, dst_ptr, ne00, ++ ne01, ncols_dst, s01, stride_col_y, stride_col_dst, ++ ne02, nchannels_y, nchannels_dst, s02, stride_channel_y, stride_channel_dst, ++ ne03, ne3, s03, s13, s3, ids_stride, stream); ++ }; ++ if (dst->type == GGML_TYPE_BF16) { ++ // FreeToken: the fused shared expert writes an fp32 shared_dst ++ GGML_ASSERT(!fusion || !fusion->shared_up); ++ launch((nv_bfloat16 *) dst->data); ++ } else if (dst->type == GGML_TYPE_F16) { ++ GGML_ASSERT(!fusion || !fusion->shared_up); ++ launch((half *) dst->data); ++ } else { ++ launch(dst_d); ++ } + } + + void ggml_cuda_op_mul_mat_vec_q( +--- a/ggml/src/ggml-cuda/mmvq.cuh ++++ b/ggml/src/ggml-cuda/mmvq.cuh +@@ -8,8 +8,10 @@ + // based on the quantization type and GPU architecture (compute capability). + int get_mmvq_mmid_max_batch(ggml_type type, int cc); + ++// FreeToken: src1_q8_1_pre skips the internal quantization of src1 + void ggml_cuda_mul_mat_vec_q(ggml_backend_cuda_context & ctx, +- const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, ggml_tensor * dst, const ggml_cuda_mm_fusion_args_host * fusion = nullptr); ++ const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, ggml_tensor * dst, const ggml_cuda_mm_fusion_args_host * fusion = nullptr, ++ const void * src1_q8_1_pre = nullptr); + + void ggml_cuda_op_mul_mat_vec_q( + ggml_backend_cuda_context & ctx, +--- a/ggml/src/ggml-cuda/quantize.cu ++++ b/ggml/src/ggml-cuda/quantize.cu +@@ -1,4 +1,5 @@ + #include "quantize.cuh" ++#include "convert.cuh" + #include + + #if defined(BLACKWELL_MMA_AVAILABLE) +@@ -50,13 +51,14 @@ + #endif // CUDART_VERSION >= 12080 + #endif // defined(BLACKWELL_MMA_AVAILABLE) + ++template + __launch_bounds__(CUDA_QUANTIZE_BLOCK_SIZE, 1) + static __global__ void quantize_q8_1( +- const float * x_ptr, void * vy_ptr, ++ const src_t * x_ptr, void * vy_ptr, + const int64_t ne00, const int64_t s01, const int64_t s02, const int64_t s03, + const int64_t ne0, const uint32_t ne1, const uint3 ne2) { + ggml_cuda_pdl_lc(); +- const float * GGML_CUDA_RESTRICT x = x_ptr; ++ const src_t * GGML_CUDA_RESTRICT x = x_ptr; + void * GGML_CUDA_RESTRICT vy = vy_ptr; + const int64_t i0 = (int64_t)blockDim.x*blockIdx.x + threadIdx.x; + +@@ -81,7 +83,7 @@ + const int64_t iqs = i_cont % QK8_1; // quant index + + ggml_cuda_pdl_sync(); +- const float xi = i0 < ne00 ? x[i03*s03 + i02*s02 + i01*s01 + i00] : 0.0f; ++ const float xi = i0 < ne00 ? ggml_cuda_cast(x[i03*s03 + i02*s02 + i01*s01 + i00]) : 0.0f; + float amax = fabsf(xi); + float sum = xi; + +@@ -452,10 +454,29 @@ + GGML_UNUSED(n_expert_used); + } + ++// FreeToken: 4 consecutive activations as float, from fp32, bf16 or fp16 ++static __device__ __forceinline__ float4 load_float4(const float * x) { ++ return *(const float4 *) x; ++} ++ ++static __device__ __forceinline__ float4 load_float4(const nv_bfloat16 * x) { ++ const uint2 v = *(const uint2 *) x; ++ const float2 lo = __bfloat1622float2(*(const nv_bfloat162 *) &v.x); ++ const float2 hi = __bfloat1622float2(*(const nv_bfloat162 *) &v.y); ++ return make_float4(lo.x, lo.y, hi.x, hi.y); ++} ++ ++static __device__ __forceinline__ float4 load_float4(const half * x) { ++ const uint2 v = *(const uint2 *) x; ++ const float2 lo = __half22float2(*(const half2 *) &v.x); ++ const float2 hi = __half22float2(*(const half2 *) &v.y); ++ return make_float4(lo.x, lo.y, hi.x, hi.y); ++} ++ + // scatter: grid over tokens, quantize once, write to all the token's compact rows +-template ++template + static __global__ void quantize_mmq_q8_1( +- const float * __restrict__ x, const int32_t * __restrict__ ids, void * __restrict__ vy, ++ const src_t * __restrict__ x, const int32_t * __restrict__ ids, void * __restrict__ vy, + const int64_t ne00, const int64_t s01, const int64_t s02, const int64_t s03, + const int64_t ne0, const int ne1, const int ne2, const int n_expert_used) { + +@@ -481,14 +502,13 @@ + base_idx = i3*s03 + i2*s02 + i01*s01; + } + +- const float4 * x4 = (const float4 *) x; + block_q8_1_mmq * y = (block_q8_1_mmq *) vy; + + const int64_t k_block = i0 / QK8_1_MMQ; // column block in the channel + const int64_t iqs = i0 % QK8_1_MMQ; // quant index in block + + // Load 4 floats per thread and calculate max. abs. value between them: +- const float4 xi = i0 < ne00 ? x4[(base_idx + i00)/4] : make_float4(0.0f, 0.0f, 0.0f, 0.0f); ++ const float4 xi = i0 < ne00 ? load_float4(x + base_idx + i00) : make_float4(0.0f, 0.0f, 0.0f, 0.0f); + float amax = fabsf(xi.x); + amax = fmaxf(amax, fabsf(xi.y)); + amax = fmaxf(amax, fabsf(xi.z)); +@@ -554,8 +574,9 @@ + GGML_UNUSED(n_expert_used); + } + ++template + void quantize_row_q8_1_cuda( +- const float * x, const int32_t * ids, void * vy, const ggml_type type_src0, ++ const src_t * x, const int32_t * ids, void * vy, const ggml_type type_src0, + const int64_t ne00, const int64_t s01, const int64_t s02, const int64_t s03, + const int64_t ne0, const int64_t ne1, const int64_t ne2, const int64_t ne3, cudaStream_t stream) { + GGML_ASSERT(!ids); +@@ -567,12 +588,13 @@ + const dim3 num_blocks(block_num_x, ne1, ne2*ne3); + const dim3 block_size(CUDA_QUANTIZE_BLOCK_SIZE, 1, 1); + const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(num_blocks, block_size, 0, stream); +- ggml_cuda_kernel_launch(quantize_q8_1, launch_params, x, vy, ne00, s01, s02, s03, ne0, ne1, ne2_fastdiv); ++ ggml_cuda_kernel_launch(quantize_q8_1, launch_params, x, vy, ne00, s01, s02, s03, ne0, ne1, ne2_fastdiv); + GGML_UNUSED(type_src0); + } + ++template + void quantize_mmq_q8_1_cuda( +- const float * x, const int32_t * ids, void * vy, const ggml_type type_src0, ++ const src_t * x, const int32_t * ids, void * vy, const ggml_type type_src0, + const int64_t ne00, const int64_t s01, const int64_t s02, const int64_t s03, + const int64_t ne0, const int64_t ne1, const int64_t ne2, const int64_t ne3, cudaStream_t stream) { + GGML_ASSERT(ne00 % 4 == 0); +@@ -584,15 +606,15 @@ + const dim3 block_size(CUDA_QUANTIZE_BLOCK_SIZE_MMQ, 1, 1); + switch (mmq_get_q8_1_ds_layout(type_src0)) { + case MMQ_Q8_1_DS_LAYOUT_D4: +- quantize_mmq_q8_1 ++ quantize_mmq_q8_1 + <<>>(x, ids, vy, ne00, s01, s02, s03, ne0, ne1, ne2, /*n_expert_used=*/0); + break; + case MMQ_Q8_1_DS_LAYOUT_DS4: +- quantize_mmq_q8_1 ++ quantize_mmq_q8_1 + <<>>(x, ids, vy, ne00, s01, s02, s03, ne0, ne1, ne2, /*n_expert_used=*/0); + break; + case MMQ_Q8_1_DS_LAYOUT_D2S6: +- quantize_mmq_q8_1 ++ quantize_mmq_q8_1 + <<>>(x, ids, vy, ne00, s01, s02, s03, ne0, ne1, ne2, /*n_expert_used=*/0); + break; + default: +@@ -601,8 +623,9 @@ + } + + // scatter=true reuses the quant kernel: grid over tokens, ids = inverse map (token slot -> compact row) ++template + void quantize_scatter_mmq_q8_1_cuda( +- const float * x, const int32_t * ids_src1_inv, void * vy, const ggml_type type_src0, ++ const src_t * x, const int32_t * ids_src1_inv, void * vy, const ggml_type type_src0, + const int64_t ne00, const int64_t stride_token, const int64_t ne0, + const int64_t n_tokens, const int64_t nrows_dst, const int n_expert_used, cudaStream_t stream) { + GGML_ASSERT(ne00 % 4 == 0); +@@ -613,15 +636,15 @@ + const dim3 block_size(CUDA_QUANTIZE_BLOCK_SIZE_MMQ, 1, 1); + switch (mmq_get_q8_1_ds_layout(type_src0)) { + case MMQ_Q8_1_DS_LAYOUT_D4: +- quantize_mmq_q8_1<<>>( ++ quantize_mmq_q8_1<<>>( + x, ids_src1_inv, vy, ne00, /*s01=*/0, /*s02=*/stride_token, /*s03=*/0, ne0, /*ne1=*/(int) nrows_dst, /*ne2=*/1, n_expert_used); + break; + case MMQ_Q8_1_DS_LAYOUT_DS4: +- quantize_mmq_q8_1<<>>( ++ quantize_mmq_q8_1<<>>( + x, ids_src1_inv, vy, ne00, /*s01=*/0, /*s02=*/stride_token, /*s03=*/0, ne0, /*ne1=*/(int) nrows_dst, /*ne2=*/1, n_expert_used); + break; + case MMQ_Q8_1_DS_LAYOUT_D2S6: +- quantize_mmq_q8_1<<>>( ++ quantize_mmq_q8_1<<>>( + x, ids_src1_inv, vy, ne00, /*s01=*/0, /*s02=*/stride_token, /*s03=*/0, ne0, /*ne1=*/(int) nrows_dst, /*ne2=*/1, n_expert_used); + break; + default: +@@ -692,3 +715,15 @@ + quantize_mmq_mxfp4<<>>(x, ids, vy, ne00, s01, s02, s03, ne0, ne1, ne2, /*n_expert_used=*/0); + } + } ++ ++#define FREETOKEN_QUANTIZE_INSTANCES(src_t) \ ++ template void quantize_row_q8_1_cuda(const src_t *, const int32_t *, void *, ggml_type, int64_t, int64_t, \ ++ int64_t, int64_t, int64_t, int64_t, int64_t, int64_t, cudaStream_t); \ ++ template void quantize_mmq_q8_1_cuda(const src_t *, const int32_t *, void *, ggml_type, int64_t, int64_t, \ ++ int64_t, int64_t, int64_t, int64_t, int64_t, int64_t, cudaStream_t); \ ++ template void quantize_scatter_mmq_q8_1_cuda(const src_t *, const int32_t *, void *, ggml_type, int64_t, \ ++ int64_t, int64_t, int64_t, int64_t, int, cudaStream_t); ++ ++FREETOKEN_QUANTIZE_INSTANCES(float) ++FREETOKEN_QUANTIZE_INSTANCES(nv_bfloat16) ++FREETOKEN_QUANTIZE_INSTANCES(half) +--- a/ggml/src/ggml-cuda/quantize.cuh ++++ b/ggml/src/ggml-cuda/quantize.cuh +@@ -16,13 +16,16 @@ + ggml_type type_src0, int64_t ne00, int64_t s01, int64_t s02, int64_t s03, + int64_t ne0, int64_t ne1, int64_t ne2, int64_t ne3, cudaStream_t stream); + ++// FreeToken: src_t is float, nv_bfloat16 or half ++template + void quantize_row_q8_1_cuda( +- const float * x, const int32_t * ids, void * vy, ++ const src_t * x, const int32_t * ids, void * vy, + ggml_type type_src0, int64_t ne00, int64_t s01, int64_t s02, int64_t s03, + int64_t ne0, int64_t ne1, int64_t ne2, int64_t ne3, cudaStream_t stream); + ++template + void quantize_mmq_q8_1_cuda( +- const float * x, const int32_t * ids, void * vy, ++ const src_t * x, const int32_t * ids, void * vy, + ggml_type type_src0, int64_t ne00, int64_t s01, int64_t s02, int64_t s03, + int64_t ne0, int64_t ne1, int64_t ne2, int64_t ne3, cudaStream_t stream); + +@@ -57,7 +60,8 @@ + int n_expert_used, + cudaStream_t stream); + +-void quantize_scatter_mmq_q8_1_cuda(const float * x, ++template ++void quantize_scatter_mmq_q8_1_cuda(const src_t * x, + const int32_t * ids_src1_inv, + void * vy, + ggml_type type_src0, diff --git a/python/freetoken/kernel/csrc/ggml/freetoken/cpu_kernels.c b/python/freetoken/kernel/csrc/ggml/freetoken/cpu_kernels.c new file mode 100644 index 0000000000..79a036383a --- /dev/null +++ b/python/freetoken/kernel/csrc/ggml/freetoken/cpu_kernels.c @@ -0,0 +1,172 @@ +// The ggml-base symbols and lookup tables that the vendored ggml-cpu quant kernels reference, and the row kernel of each weight type. + +#include "cpu_kernels.h" + +#include "ggml-impl.h" +#include "ggml-cpu/ggml-cpu-impl.h" +#include "ggml-cpu/quants.h" +#include "ggml-cpu/simd-mappings.h" +#include "ggml-cpu/vec.h" + +#include +#include +#include +#include + +float ggml_table_f32_f16[1 << 16]; +float ggml_table_f32_e8m0_half[1 << 8]; +float ggml_table_f32_ue4m3[1 << 8]; + +void ggml_abort(const char * file, int line, const char * fmt, ...) { + va_list args; + va_start(args, fmt); + fprintf(stderr, "ggml: %s:%d: ", file, line); + vfprintf(stderr, fmt, args); + va_end(args); + fputc('\n', stderr); + fflush(stderr); + abort(); +} + +struct type_info { + const char * name; + int64_t blck_size; + size_t type_size; +}; + +static const struct type_info k_types[GGML_TYPE_COUNT] = { + [GGML_TYPE_F32] = {"f32", 1, sizeof(float)}, + [GGML_TYPE_F16] = {"f16", 1, sizeof(ggml_fp16_t)}, + [GGML_TYPE_BF16] = {"bf16", 1, sizeof(ggml_bf16_t)}, + [GGML_TYPE_Q1_0] = {"q1_0", QK1_0, sizeof(block_q1_0)}, + [GGML_TYPE_Q2_0] = {"q2_0", QK2_0, sizeof(block_q2_0)}, + [GGML_TYPE_Q4_0] = {"q4_0", QK4_0, sizeof(block_q4_0)}, + [GGML_TYPE_Q4_1] = {"q4_1", QK4_1, sizeof(block_q4_1)}, + [GGML_TYPE_Q5_0] = {"q5_0", QK5_0, sizeof(block_q5_0)}, + [GGML_TYPE_Q5_1] = {"q5_1", QK5_1, sizeof(block_q5_1)}, + [GGML_TYPE_Q8_0] = {"q8_0", QK8_0, sizeof(block_q8_0)}, + [GGML_TYPE_Q8_1] = {"q8_1", QK8_1, sizeof(block_q8_1)}, + [GGML_TYPE_MXFP4] = {"mxfp4", QK_MXFP4, sizeof(block_mxfp4)}, + [GGML_TYPE_NVFP4] = {"nvfp4", QK_NVFP4, sizeof(block_nvfp4)}, + [GGML_TYPE_Q2_K] = {"q2_K", QK_K, sizeof(block_q2_K)}, + [GGML_TYPE_Q3_K] = {"q3_K", QK_K, sizeof(block_q3_K)}, + [GGML_TYPE_Q4_K] = {"q4_K", QK_K, sizeof(block_q4_K)}, + [GGML_TYPE_Q5_K] = {"q5_K", QK_K, sizeof(block_q5_K)}, + [GGML_TYPE_Q6_K] = {"q6_K", QK_K, sizeof(block_q6_K)}, + [GGML_TYPE_Q8_K] = {"q8_K", QK_K, sizeof(block_q8_K)}, + [GGML_TYPE_TQ1_0] = {"tq1_0", QK_K, sizeof(block_tq1_0)}, + [GGML_TYPE_TQ2_0] = {"tq2_0", QK_K, sizeof(block_tq2_0)}, + [GGML_TYPE_IQ2_XXS] = {"iq2_xxs", QK_K, sizeof(block_iq2_xxs)}, + [GGML_TYPE_IQ2_XS] = {"iq2_xs", QK_K, sizeof(block_iq2_xs)}, + [GGML_TYPE_IQ2_S] = {"iq2_s", QK_K, sizeof(block_iq2_s)}, + [GGML_TYPE_IQ3_XXS] = {"iq3_xxs", QK_K, sizeof(block_iq3_xxs)}, + [GGML_TYPE_IQ3_S] = {"iq3_s", QK_K, sizeof(block_iq3_s)}, + [GGML_TYPE_IQ1_S] = {"iq1_s", QK_K, sizeof(block_iq1_s)}, + [GGML_TYPE_IQ1_M] = {"iq1_m", QK_K, sizeof(block_iq1_m)}, + [GGML_TYPE_IQ4_NL] = {"iq4_nl", QK4_NL, sizeof(block_iq4_nl)}, + [GGML_TYPE_IQ4_XS] = {"iq4_xs", QK_K, sizeof(block_iq4_xs)}, +}; + +static int known_type(int type) { return type >= 0 && type < GGML_TYPE_COUNT && k_types[type].name != NULL; } + +int64_t ggml_blck_size(enum ggml_type type) { return k_types[type].blck_size; } + +size_t ggml_type_size(enum ggml_type type) { return k_types[type].type_size; } + +size_t ggml_row_size(enum ggml_type type, int64_t ne) { return k_types[type].type_size * ne / k_types[type].blck_size; } + +const char * ggml_type_name(enum ggml_type type) { return known_type(type) ? k_types[type].name : "unknown"; } + +const char * ft_ggml_type_name(int type) { return ggml_type_name((enum ggml_type) type); } + +struct dot_info { + ft_ggml_vec_dot_t vec_dot; + int act_type; +}; + +// The weight types of type_traits_cpu in ggml-cpu.c, with their vec_dot and vec_dot_type. +static const struct dot_info k_dots[GGML_TYPE_COUNT] = { + [GGML_TYPE_F32] = {(ft_ggml_vec_dot_t) ggml_vec_dot_f32, GGML_TYPE_F32}, + [GGML_TYPE_F16] = {(ft_ggml_vec_dot_t) ggml_vec_dot_f16, GGML_TYPE_F16}, + [GGML_TYPE_BF16] = {(ft_ggml_vec_dot_t) ggml_vec_dot_bf16, GGML_TYPE_BF16}, + [GGML_TYPE_Q1_0] = {ggml_vec_dot_q1_0_q8_0, GGML_TYPE_Q8_0}, + [GGML_TYPE_Q2_0] = {ggml_vec_dot_q2_0_q8_0, GGML_TYPE_Q8_0}, + [GGML_TYPE_Q4_0] = {ggml_vec_dot_q4_0_q8_0, GGML_TYPE_Q8_0}, + [GGML_TYPE_Q4_1] = {ggml_vec_dot_q4_1_q8_1, GGML_TYPE_Q8_1}, + [GGML_TYPE_Q5_0] = {ggml_vec_dot_q5_0_q8_0, GGML_TYPE_Q8_0}, + [GGML_TYPE_Q5_1] = {ggml_vec_dot_q5_1_q8_1, GGML_TYPE_Q8_1}, + [GGML_TYPE_Q8_0] = {ggml_vec_dot_q8_0_q8_0, GGML_TYPE_Q8_0}, + [GGML_TYPE_MXFP4] = {ggml_vec_dot_mxfp4_q8_0, GGML_TYPE_Q8_0}, + [GGML_TYPE_NVFP4] = {ggml_vec_dot_nvfp4_q8_0, GGML_TYPE_Q8_0}, + [GGML_TYPE_Q2_K] = {ggml_vec_dot_q2_K_q8_K, GGML_TYPE_Q8_K}, + [GGML_TYPE_Q3_K] = {ggml_vec_dot_q3_K_q8_K, GGML_TYPE_Q8_K}, + [GGML_TYPE_Q4_K] = {ggml_vec_dot_q4_K_q8_K, GGML_TYPE_Q8_K}, + [GGML_TYPE_Q5_K] = {ggml_vec_dot_q5_K_q8_K, GGML_TYPE_Q8_K}, + [GGML_TYPE_Q6_K] = {ggml_vec_dot_q6_K_q8_K, GGML_TYPE_Q8_K}, + [GGML_TYPE_TQ1_0] = {ggml_vec_dot_tq1_0_q8_K, GGML_TYPE_Q8_K}, + [GGML_TYPE_TQ2_0] = {ggml_vec_dot_tq2_0_q8_K, GGML_TYPE_Q8_K}, + [GGML_TYPE_IQ2_XXS] = {ggml_vec_dot_iq2_xxs_q8_K, GGML_TYPE_Q8_K}, + [GGML_TYPE_IQ2_XS] = {ggml_vec_dot_iq2_xs_q8_K, GGML_TYPE_Q8_K}, + [GGML_TYPE_IQ2_S] = {ggml_vec_dot_iq2_s_q8_K, GGML_TYPE_Q8_K}, + [GGML_TYPE_IQ3_XXS] = {ggml_vec_dot_iq3_xxs_q8_K, GGML_TYPE_Q8_K}, + [GGML_TYPE_IQ3_S] = {ggml_vec_dot_iq3_s_q8_K, GGML_TYPE_Q8_K}, + [GGML_TYPE_IQ1_S] = {ggml_vec_dot_iq1_s_q8_K, GGML_TYPE_Q8_K}, + [GGML_TYPE_IQ1_M] = {ggml_vec_dot_iq1_m_q8_K, GGML_TYPE_Q8_K}, + [GGML_TYPE_IQ4_NL] = {ggml_vec_dot_iq4_nl_q8_0, GGML_TYPE_Q8_0}, + [GGML_TYPE_IQ4_XS] = {ggml_vec_dot_iq4_xs_q8_K, GGML_TYPE_Q8_K}, +}; + +static void fp32_to_fp32_row(const float * x, void * y, int64_t k) { memcpy(y, x, (size_t) k * sizeof(float)); } + +static void fp32_to_fp16_row(const float * x, void * y, int64_t k) { + ggml_fp16_t * out = y; + for (int64_t i = 0; i < k; ++i) { + out[i] = GGML_CPU_FP32_TO_FP16(x[i]); + } +} + +static void fp32_to_bf16_row(const float * x, void * y, int64_t k) { + ggml_bf16_t * out = y; + for (int64_t i = 0; i < k; ++i) { + out[i] = GGML_FP32_TO_BF16(x[i]); + } +} + +static ft_ggml_from_float_t act_quantizer(int act_type) { + switch (act_type) { + case GGML_TYPE_F32: return fp32_to_fp32_row; + case GGML_TYPE_F16: return fp32_to_fp16_row; + case GGML_TYPE_BF16: return fp32_to_bf16_row; + case GGML_TYPE_Q8_0: return quantize_row_q8_0; + case GGML_TYPE_Q8_1: return quantize_row_q8_1; + case GGML_TYPE_Q8_K: return quantize_row_q8_K; + default: return NULL; + } +} + +// The tables ggml_cpu_init fills for the quant kernels. +void ft_ggml_cpu_init(void) { + for (int i = 0; i < (1 << 16); ++i) { + ggml_table_f32_f16[i] = GGML_COMPUTE_FP16_TO_FP32((ggml_fp16_t) i); + } + for (int i = 0; i < (1 << 8); ++i) { + ggml_table_f32_e8m0_half[i] = GGML_E8M0_TO_FP32_HALF(i); + ggml_table_f32_ue4m3[i] = ggml_ue4m3_to_fp32((uint8_t) i); + } +} + +int ft_ggml_cpu_row_kernel(int type, int64_t k, struct ft_ggml_row_kernel * out) { + if (!known_type(type) || k_dots[type].vec_dot == NULL) { + return FT_GGML_NO_KERNEL; + } + const int act = k_dots[type].act_type; + if (k % k_types[type].blck_size != 0 || k % k_types[act].blck_size != 0) { + return FT_GGML_PARTIAL_BLOCK; + } + out->vec_dot = k_dots[type].vec_dot; + out->from_float = act_quantizer(act); + out->act_type = act; + out->row_bytes = ggml_row_size((enum ggml_type) type, k); + out->act_row_bytes = ggml_row_size((enum ggml_type) act, k); + return FT_GGML_OK; +} diff --git a/python/freetoken/kernel/csrc/ggml/freetoken/cpu_kernels.h b/python/freetoken/kernel/csrc/ggml/freetoken/cpu_kernels.h new file mode 100644 index 0000000000..d58c174854 --- /dev/null +++ b/python/freetoken/kernel/csrc/ggml/freetoken/cpu_kernels.h @@ -0,0 +1,43 @@ +// ggml-cpu row kernels for the CPU MoE executor. setup.py builds them once per x86 variant into freetoken.kernel.ggml_cpu._; call them only through the table of a variant whose score() is not 0. +#pragma once + +#include +#include + +#ifdef __cplusplus +extern "C" { +#endif + +typedef void (*ft_ggml_vec_dot_t)(int n, float * s, size_t bs, const void * x, size_t bx, const void * y, size_t by, int nrc); +typedef void (*ft_ggml_from_float_t)(const float * x, void * y, int64_t k); + +struct ft_ggml_row_kernel { + ft_ggml_vec_dot_t vec_dot; // one weight row . one activation row in act_type + ft_ggml_from_float_t from_float; // fp32 activation row -> act_type + int act_type; + size_t row_bytes; // one weight row of k values + size_t act_row_bytes; // one activation row of k values in act_type +}; + +enum { FT_GGML_OK = 0, FT_GGML_NO_KERNEL = 1, FT_GGML_PARTIAL_BLOCK = 2 }; + +void ft_ggml_cpu_init(void); + +// The kernels for weight rows of ggml type `type` over k values. +int ft_ggml_cpu_row_kernel(int type, int64_t k, struct ft_ggml_row_kernel * out); + +const char * ft_ggml_type_name(int type); + +// One variant's functions, as its module's api() hands them to _cpu_moe.set_ggml_cpu. +#define FT_GGML_CPU_ABI 1 +struct ft_ggml_cpu_api { + int abi; // FT_GGML_CPU_ABI of the module that built it + const char * name; // the variant, e.g. "haswell" + void (*init)(void); + int (*row_kernel)(int type, int64_t k, struct ft_ggml_row_kernel * out); + const char * (*type_name)(int type); +}; + +#ifdef __cplusplus +} +#endif diff --git a/python/freetoken/kernel/csrc/ggml/freetoken/ggml_ops.cu b/python/freetoken/kernel/csrc/ggml/freetoken/ggml_ops.cu new file mode 100644 index 0000000000..170cbcac71 --- /dev/null +++ b/python/freetoken/kernel/csrc/ggml/freetoken/ggml_ops.cu @@ -0,0 +1,469 @@ +// tvm-ffi ops over the vendored ggml-cuda kernels. Weights stay in their packed ggml blocks. +// Activations go through the same q8_1 quantization and kernel choice as llama.cpp's mul_mat / mul_mat_id. +// Activations are fp32, bf16 or fp16 with contiguous rows. Results have the activation dtype and are accumulated in fp32; the caller allocates every output. + +#include + +#include "common.cuh" +#include "convert.cuh" +#include "getrows.cuh" +#include "mmq.cuh" +#include "mmvq.cuh" +#include "quantize.cuh" + +#include +#include +#include +#include +#include + +#include +#include +#include +#include + +using tvm::ffi::Optional; +using tvm::ffi::String; +using tvm::ffi::TensorView; + +namespace { + +using host::RuntimeCheck; + +struct DeviceGuard { + int prev, device; + + explicit DeviceGuard(DLDevice d) : prev(ggml_cuda_get_device()), device(d.device_id) { +#ifdef GGML_USE_HIP + RuntimeCheck(d.device_type == kDLROCM, "ggml ops take ROCm tensors"); +#else + RuntimeCheck(d.device_type == kDLCUDA, "ggml ops take CUDA tensors"); +#endif + ggml_cuda_set_device(device); + } + + ~DeviceGuard() { + if (prev != device) { + cudaSetDevice(prev); + } + } +}; + +// ggml treats a null stream as unset and would create its own; name the default stream explicitly +cudaStream_t current_stream(DLDevice d) { + auto stream = static_cast(TVMFFIEnvGetStream(d.device_type, d.device_id)); +#ifdef GGML_USE_HIP + return stream == nullptr ? hipStreamLegacy : stream; +#else + return stream == nullptr ? cudaStreamLegacy : stream; +#endif +} + +ggml_backend_cuda_context & context(DLDevice d) { + static std::unique_ptr ctxs[GGML_CUDA_MAX_DEVICES]; + const int dev = d.device_id; + if (!ctxs[dev]) { + ctxs[dev] = std::make_unique(dev); + } + ctxs[dev]->curr_stream_no = 0; + ctxs[dev]->streams[dev][0] = current_stream(d); + return *ctxs[dev]; +} + +ggml_tensor describe(ggml_type type, const void * data, std::array ne, std::array nb) { + ggml_tensor t{}; + t.type = type; + t.data = const_cast(data); + for (int i = 0; i < 4; ++i) { + t.ne[i] = ne[i]; + t.nb[i] = nb[i]; + } + return t; +} + +int device_cc() { return ggml_cuda_info().devices[ggml_cuda_get_device()].cc; } + +bool is_dtype(DLDataType dt, DLDataTypeCode code, int bits) { return dt.code == code && dt.bits == bits && dt.lanes == 1; } + +ggml_type float_type(TensorView t) { + const DLDataType dt = t.dtype(); + if (is_dtype(dt, kDLFloat, 32)) { + return GGML_TYPE_F32; + } + if (is_dtype(dt, kDLBfloat, 16)) { + return GGML_TYPE_BF16; + } + RuntimeCheck(is_dtype(dt, kDLFloat, 16), "ggml tensors must be fp32, bf16 or fp16"); + return GGML_TYPE_F16; +} + +// the activation dtype, which the result must share +ggml_type io_type(TensorView x, TensorView out) { + const ggml_type type = float_type(x); + RuntimeCheck(float_type(out) == type, "the result must have the activation dtype"); + return type; +} + +int64_t check_weight(TensorView w, ggml_type type, int64_t k) { + RuntimeCheck(is_dtype(w.dtype(), kDLUInt, 8), "ggml weight must be a uint8 tensor"); + RuntimeCheck(w.stride(w.ndim() - 1) == 1, "ggml weight rows must be contiguous"); + const int64_t bs = ggml_blck_size(type); + RuntimeCheck(k % bs == 0, "K=", k, " is not a multiple of the ", ggml_type_name(type), " block size ", bs); + const int64_t row_bytes = k / bs * (int64_t) ggml_type_size(type); + RuntimeCheck(w.size(w.ndim() - 1) == row_bytes, "ggml weight row is ", w.size(w.ndim() - 1), " bytes, ", ggml_type_name(type), + " K=", k, " needs ", row_bytes); + return row_bytes; +} + +// one fp32 scale per expert, e.g. a column of a [E, n] scale bank; the MMVQ epilogue multiplies NVFP4 results by it +ggml_tensor describe_scale(TensorView s, ggml_type type, int64_t e) { + RuntimeCheck(type == GGML_TYPE_NVFP4, "per-expert scales apply to NVFP4 weights, not ", ggml_type_name(type)); + RuntimeCheck(is_dtype(s.dtype(), kDLFloat, 32) && s.ndim() == 1 && s.size(0) == e && s.stride(0) > 0, "a per-expert scale must be an fp32 [E] view"); + const size_t nb0 = sizeof(float) * s.stride(0); + return describe(GGML_TYPE_F32, s.data_ptr(), {e, 1, 1, 1}, {nb0, nb0 * e, nb0 * e, nb0 * e}); +} + +int64_t q8_1_row_bytes(int64_t k) { return GGML_PAD(k, MATRIX_ROW_PADDING) / QK8_1 * (int64_t) sizeof(block_q8_1); } + +// MMQ tiles read up to MATRIX_ROW_PADDING values past the last row; the caller keeps this many bytes after it +int64_t mmq_tail_bytes(int64_t type_id, int64_t k) { + const int64_t rem = k % MATRIX_ROW_PADDING; + return rem == 0 ? 0 : (int64_t) ggml_row_size((ggml_type) type_id, MATRIX_ROW_PADDING - rem); +} + +// MMQ / MMVQ index weight blocks from the data pointer with 32-bit ints, MMQ tail reads included +void check_block_offsets(TensorView w, ggml_type type, int64_t k) { + if (w.numel() == 0) { + return; + } + int64_t bytes = 1 + mmq_tail_bytes(type, k); + for (int d = 0; d < w.ndim(); ++d) { + bytes += (w.size(d) - 1) * w.stride(d); + } + const int64_t blocks = bytes / (int64_t) ggml_type_size(type); + RuntimeCheck(blocks <= INT32_MAX, "ggml weight spans ", blocks, " ", ggml_type_name(type), " blocks; MMQ / MMVQ address at most ", INT32_MAX); +} + +struct Activation { + const void * data; + ggml_type type; + size_t esz; +}; + +// The native fp4 MMQ path quantizes fp32 activations only; it is built only with GGML_CUDA_FP4_MMA. +bool mmq_reads_16bit(ggml_type type) { +#ifdef GGML_CUDA_FP4_MMA + return type != GGML_TYPE_MXFP4 && type != GGML_TYPE_NVFP4; +#else + GGML_UNUSED(type); + return true; +#endif +} + +// x as the kernel reads it: x itself, or an fp32 copy in ``tmp`` for a kernel that reads fp32 only +Activation activation(TensorView x, bool f16_ok, ggml_cuda_pool_alloc & tmp, cudaStream_t stream) { + RuntimeCheck(x.IsContiguous(), "ggml activations must be contiguous"); + const ggml_type type = float_type(x); + if (type == GGML_TYPE_F32) { + return {x.data_ptr(), type, 4}; + } + if (f16_ok) { + return {x.data_ptr(), type, 2}; + } + float * f32 = tmp.alloc(x.numel()); + ggml_get_to_fp32_cuda(type)(x.data_ptr(), f32, x.numel(), stream); + return {f32, GGML_TYPE_F32, 4}; +} + +} // namespace + +// out[M, q8_1_row_bytes(K)] = x[M, K] as MMVQ's q8_1 blocks, for GEMVs that share one input +void op_quantize_q8_1(TensorView x, TensorView out) { + RuntimeCheck(x.ndim() == 2 && x.stride(1) == 1, "quantize_q8_1 takes x[M, K] with contiguous rows"); + const DeviceGuard guard(x.device()); + const int64_t m = x.size(0), k = x.size(1), k_padded = GGML_PAD(k, MATRIX_ROW_PADDING); + RuntimeCheck(is_dtype(out.dtype(), kDLUInt, 8) && out.IsContiguous() && out.ndim() == 2 && out.size(0) == m && + out.size(1) == q8_1_row_bytes(k), + "out must be a contiguous uint8 [M, q8_1_row_bytes(K)] tensor"); + if (m == 0) { + return; + } + const int64_t s1 = x.stride(0); + const cudaStream_t stream = current_stream(x.device()); + const ggml_type type = float_type(x); + if (type == GGML_TYPE_BF16) { + quantize_row_q8_1_cuda((const nv_bfloat16 *) x.data_ptr(), nullptr, out.data_ptr(), GGML_TYPE_Q8_0, k, s1, s1 * m, s1 * m, k_padded, m, 1, 1, stream); + } else if (type == GGML_TYPE_F16) { + quantize_row_q8_1_cuda((const half *) x.data_ptr(), nullptr, out.data_ptr(), GGML_TYPE_Q8_0, k, s1, s1 * m, s1 * m, k_padded, m, 1, 1, stream); + } else { + quantize_row_q8_1_cuda((const float *) x.data_ptr(), nullptr, out.data_ptr(), GGML_TYPE_Q8_0, k, s1, s1 * m, s1 * m, k_padded, m, 1, 1, stream); + } +} + +// out[M, N] = x[M, K] @ W[N, K]^T; out may be a view with strided rows; MMVQ reads ``xq`` (op_quantize_q8_1 of x) instead of quantizing x again +void op_mul_mat(TensorView w, int64_t type_id, TensorView x, TensorView out, Optional xq) { + RuntimeCheck(x.ndim() == 2 && w.ndim() == 2, "mul_mat takes w[N, row_bytes] and x[M, K]"); + const DeviceGuard guard(x.device()); + const ggml_type type = (ggml_type) type_id; + const int64_t m = x.size(0), k = x.size(1), n = w.size(0); + check_weight(w, type, k); + check_block_offsets(w, type, k); + RuntimeCheck(out.ndim() == 2 && out.size(0) == m && out.size(1) == n && out.stride(1) == 1, "mul_mat out must be a [M, N] view with contiguous rows"); + const ggml_type otype = io_type(x, out); + if (m == 0) { + return; + } + const size_t s1 = w.stride(0); + ggml_tensor src0 = describe(type, w.data_ptr(), {k, n, 1, 1}, {ggml_type_size(type), s1, s1 * n, s1 * n}); + const size_t oe = ggml_type_size(otype), os = out.stride(0) * oe; + ggml_tensor dst = describe(otype, out.data_ptr(), {n, m, 1, 1}, {oe, os, os * m, os * m}); + auto & ctx = context(x.device()); + const cudaStream_t stream = ctx.stream(); + const int cc = device_cc(); + ggml_cuda_pool_alloc x32(ctx.pool()); + if (ggml_cuda_should_use_mmvq(type, cc, m)) { + const Activation a = activation(x, true, x32, stream); + ggml_tensor src1 = describe(a.type, a.data, {k, m, 1, 1}, {a.esz, a.esz * k, a.esz * k * m, a.esz * k * m}); + const void * q8 = nullptr; + if (xq.has_value()) { + const TensorView q = xq.value(); + RuntimeCheck(is_dtype(q.dtype(), kDLUInt, 8) && q.IsContiguous() && q.ndim() == 2 && q.size(0) == m && q.size(1) == q8_1_row_bytes(k), + "xq must be quantize_q8_1 of x"); + q8 = q.data_ptr(); + } + ggml_cuda_mul_mat_vec_q(ctx, &src0, &src1, nullptr, &dst, nullptr, q8); + } else if (ggml_cuda_should_use_mmq(type, cc, m, 0)) { + const Activation a = activation(x, mmq_reads_16bit(type), x32, stream); + ggml_tensor src1 = describe(a.type, a.data, {k, m, 1, 1}, {a.esz, a.esz * k, a.esz * k * m, a.esz * k * m}); + ggml_cuda_mul_mat_q(ctx, &src0, &src1, nullptr, &dst); + } else { + RuntimeCheck(false, "no ggml kernel for ", ggml_type_name(type), " with M=", m); + } +} + +// out[T, k, N]: row (t, j) multiplies x[t, j or 0] by expert ids[t, j] of W[E, N, K]; the matrix kernel runs a flat grid over the tiles that hold routes +// scale[ids[t, j]] multiplies row (t, j); only the vector kernel takes it +void op_mul_mat_id(TensorView w, int64_t type_id, TensorView x, TensorView ids, TensorView out, Optional scale) { + RuntimeCheck(w.ndim() == 3, "mul_mat_id takes w[E, N, row_bytes]"); + RuntimeCheck(is_dtype(ids.dtype(), kDLInt, 32) && ids.ndim() == 2 && ids.IsContiguous(), "ids must be a contiguous int32 [T, k] tensor"); + const DeviceGuard guard(x.device()); + const ggml_type type = (ggml_type) type_id; + const int64_t t = ids.size(0), top_k = ids.size(1); + RuntimeCheck(x.ndim() == 3 && x.size(0) == t && (x.size(1) == 1 || x.size(1) == top_k), "x must be [T, 1, K] or [T, k, K]"); + const int64_t ky = x.size(1), k = x.size(2), e = w.size(0), n = w.size(1); + check_weight(w, type, k); + check_block_offsets(w, type, k); + RuntimeCheck(out.ndim() == 3 && out.size(0) == t && out.size(1) == top_k && out.size(2) == n && out.IsContiguous(), + "out must be a contiguous [T, k, N] tensor"); + const ggml_type otype = io_type(x, out); + if (t == 0) { + return; + } + const size_t s1 = w.stride(1), s2 = w.stride(0); + ggml_tensor src0 = describe(type, w.data_ptr(), {k, n, e, 1}, {ggml_type_size(type), s1, s2, s2 * e}); + const size_t ids_stride = sizeof(int32_t) * top_k, ids_bytes = ids_stride * t; + ggml_tensor tids = describe(GGML_TYPE_I32, ids.data_ptr(), {top_k, t, 1, 1}, {sizeof(int32_t), ids_stride, ids_bytes, ids_bytes}); + auto & ctx = context(x.device()); + const cudaStream_t stream = ctx.stream(); + const int cc = device_cc(); + const bool mmvq = t <= MMVQ_MAX_BATCH_SIZE && t <= get_mmvq_mmid_max_batch(type, cc); + RuntimeCheck(mmvq || ggml_cuda_should_use_mmq(type, cc, t, e), "no ggml MoE kernel for ", ggml_type_name(type), " with T=", t); + ggml_cuda_pool_alloc x32(ctx.pool()); + const Activation a = activation(x, mmvq || mmq_reads_16bit(type), x32, stream); + ggml_tensor src1 = describe(a.type, a.data, {k, ky, t, 1}, {a.esz, a.esz * k, a.esz * k * ky, a.esz * k * ky * t}); + const size_t oe = ggml_type_size(otype); + ggml_tensor dst = describe(otype, out.data_ptr(), {n, top_k, t, 1}, {oe, oe * n, oe * n * top_k, oe * n * top_k * t}); + RuntimeCheck(mmvq || !scale.has_value(), "mul_mat_id takes a per-expert scale only on the vector kernel, T=", t); + if (mmvq) { + ggml_cuda_mm_fusion_args_host fusion{}; + ggml_tensor tscale; + if (scale.has_value()) { + tscale = describe_scale(scale.value(), type, e); + fusion.x_scale = &tscale; + } + ggml_cuda_mul_mat_vec_q(ctx, &src0, &src1, &tids, &dst, scale.has_value() ? &fusion : nullptr); + } else { + ggml_cuda_mul_mat_q(ctx, &src0, &src1, &tids, &dst, 0); + } +} + +namespace { + +// The GLUs the MMVQ epilogue computes; any other ggml_glu_op falls through to a plain gate * up there. +ggml_glu_op fused_glu_op(const String & name) { + if (name == "swiglu") { + return GGML_GLU_OP_SWIGLU; + } + if (name == "geglu") { + return GGML_GLU_OP_GEGLU; + } + if (name == "swiglu_oai") { + return GGML_GLU_OP_SWIGLU_OAI; + } + RuntimeCheck(name == "swiglu_clamp", "MMVQ fuses swiglu, geglu, swiglu_oai or swiglu_clamp, not ", std::string(name)); + return GGML_GLU_OP_SWIGLU_CLAMP; +} + +} // namespace + +// out[T, k, I] = glu(x @ gate^T, x @ up^T) for small T in one MMVQ pass; w[E, 2I, row_bytes] holds each expert's gate rows then its up rows; limit is swiglu_clamp's +// gate_scale / up_scale multiply the expert's gate / up results before the glu +void op_mul_mat_id_glu(TensorView w, int64_t type_id, TensorView x, TensorView ids, TensorView out, String glu, double limit, + Optional gate_scale, Optional up_scale) { + const ggml_glu_op glu_op = fused_glu_op(glu); + RuntimeCheck(w.ndim() == 3 && w.size(1) % 2 == 0, "mul_mat_id_glu takes w[E, 2I, row_bytes]"); + RuntimeCheck(is_dtype(ids.dtype(), kDLInt, 32) && ids.ndim() == 2 && ids.IsContiguous(), "ids must be a contiguous int32 [T, k] tensor"); + RuntimeCheck(x.ndim() == 2 && x.size(0) == ids.size(0), "x must be [T, K]"); + const DeviceGuard guard(x.device()); + const ggml_type type = (ggml_type) type_id; + const int64_t t = ids.size(0), top_k = ids.size(1), k = x.size(1), e = w.size(0), n = w.size(1) / 2; + check_weight(w, type, k); + check_block_offsets(w, type, k); + RuntimeCheck(out.IsContiguous() && out.ndim() == 3 && out.size(0) == t && out.size(1) == top_k && out.size(2) == n, + "out must be a contiguous [T, k, I] tensor"); + const ggml_type otype = io_type(x, out); + const int cc = device_cc(); + RuntimeCheck(t <= MMVQ_MAX_BATCH_SIZE && t <= get_mmvq_mmid_max_batch(type, cc), "mul_mat_id_glu is the small-batch path, T=", t); + if (t == 0) { + return; + } + auto & ctx = context(x.device()); + ggml_cuda_pool_alloc x32(ctx.pool()); + const Activation a = activation(x, true, x32, ctx.stream()); + const size_t s1 = w.stride(1), s2 = w.stride(0); + const char * base = (const char *) w.data_ptr(); + ggml_tensor up = describe(type, base + n * s1, {k, n, e, 1}, {ggml_type_size(type), s1, s2, s2 * e}); + ggml_tensor gate = describe(type, base, {k, n, e, 1}, {ggml_type_size(type), s1, s2, s2 * e}); + ggml_tensor src1 = describe(a.type, a.data, {k, 1, t, 1}, {a.esz, a.esz * k, a.esz * k, a.esz * k * t}); + const size_t ids_stride = sizeof(int32_t) * top_k, ids_bytes = ids_stride * t; + ggml_tensor tids = describe(GGML_TYPE_I32, ids.data_ptr(), {top_k, t, 1, 1}, {sizeof(int32_t), ids_stride, ids_bytes, ids_bytes}); + const size_t oe = ggml_type_size(otype); + ggml_tensor dst = describe(otype, out.data_ptr(), {n, top_k, t, 1}, {oe, oe * n, oe * n * top_k, oe * n * top_k * t}); + ggml_cuda_mm_fusion_args_host fusion{}; + fusion.gate = &gate; + fusion.glu_op = glu_op; + fusion.glu_limit = (float) limit; + ggml_tensor tgate_scale, tup_scale; + if (gate_scale.has_value()) { + tgate_scale = describe_scale(gate_scale.value(), type, e); + fusion.gate_scale = &tgate_scale; + } + if (up_scale.has_value()) { + tup_scale = describe_scale(up_scale.value(), type, e); + fusion.x_scale = &tup_scale; + } + ggml_cuda_mul_mat_vec_q(ctx, &up, &src1, &tids, &dst, &fusion); +} + +namespace { + +template +__global__ void moe_weighted_sum_kernel(const T * __restrict__ y, const float * __restrict__ w, T * __restrict__ out, + const int64_t n, const int top_k) { + const int64_t token = blockIdx.x; + const int64_t col = (int64_t) blockIdx.y * blockDim.x + threadIdx.x; + if (col >= n) { + return; + } + const int64_t row0 = token * top_k; + float sum = 0.0f; + for (int j = 0; j < top_k; ++j) { + sum += ggml_cuda_cast(y[(row0 + j) * n + col]) * w[row0 + j]; + } + out[token * n + col] = ggml_cuda_cast(sum); +} + +} // namespace + +// out[T, N] = sum_j weights[t, j] * y[t, j, :], summed in fp32 +void op_moe_weighted_sum(TensorView y, TensorView weights, TensorView out) { + RuntimeCheck(y.ndim() == 3 && y.IsContiguous(), "y must be a contiguous [T, k, N] tensor"); + const int64_t t = y.size(0), top_k = y.size(1), n = y.size(2); + RuntimeCheck(float_type(weights) == GGML_TYPE_F32 && weights.IsContiguous() && weights.numel() == t * top_k, "weights must be a contiguous fp32 [T, k] tensor"); + RuntimeCheck(out.ndim() == 2 && out.size(0) == t && out.size(1) == n && out.IsContiguous(), "out must be a contiguous [T, N] tensor"); + const ggml_type otype = io_type(y, out); + const DeviceGuard guard(y.device()); + if (t == 0) { + return; + } + constexpr int threads = 256; + const dim3 blocks(t, (n + threads - 1) / threads); + const cudaStream_t stream = current_stream(y.device()); + const float * wp = (const float *) weights.data_ptr(); + if (otype == GGML_TYPE_BF16) { + moe_weighted_sum_kernel<<>>((const nv_bfloat16 *) y.data_ptr(), wp, (nv_bfloat16 *) out.data_ptr(), n, top_k); + } else if (otype == GGML_TYPE_F16) { + moe_weighted_sum_kernel<<>>((const half *) y.data_ptr(), wp, (half *) out.data_ptr(), n, top_k); + } else { + moe_weighted_sum_kernel<<>>((const float *) y.data_ptr(), wp, (float *) out.data_ptr(), n, top_k); + } + CUDA_CHECK(cudaGetLastError()); +} + +// out[rows, K] = W[rows, row_bytes] dequantized to out's dtype +void op_dequantize(TensorView w, int64_t type_id, int64_t k, TensorView out) { + const DeviceGuard guard(w.device()); + const ggml_type type = (ggml_type) type_id; + check_weight(w, type, k); + RuntimeCheck(w.IsContiguous(), "dequantize needs a contiguous weight"); + const int64_t rows = w.numel() / w.size(w.ndim() - 1); + RuntimeCheck(out.IsContiguous() && out.numel() == rows * k, "out must be a contiguous [rows, K] tensor"); + const ggml_type otype = float_type(out); + if (rows == 0) { + return; + } + const cudaStream_t stream = current_stream(w.device()); + if (type == otype) { + CUDA_CHECK(cudaMemcpyAsync(out.data_ptr(), w.data_ptr(), rows * k * ggml_type_size(type), cudaMemcpyDeviceToDevice, stream)); + return; + } + if (otype == GGML_TYPE_BF16) { + const auto convert = ggml_get_to_bf16_cuda(type); + RuntimeCheck(convert != nullptr, "no ggml bf16 conversion for ", ggml_type_name(type)); + convert(w.data_ptr(), (nv_bfloat16 *) out.data_ptr(), rows * k, stream); + } else if (otype == GGML_TYPE_F16) { + const auto convert = ggml_get_to_fp16_cuda(type); + RuntimeCheck(convert != nullptr, "no ggml fp16 conversion for ", ggml_type_name(type)); + convert(w.data_ptr(), (half *) out.data_ptr(), rows * k, stream); + } else { + const auto convert = ggml_get_to_fp32_cuda(type); + RuntimeCheck(convert != nullptr, "no ggml fp32 conversion for ", ggml_type_name(type)); + convert(w.data_ptr(), (float *) out.data_ptr(), rows * k, stream); + } +} + +// out[i] = dequant(W[ids[i]]) for a table W[rows, row_bytes] +void op_get_rows(TensorView w, int64_t type_id, int64_t k, TensorView ids, TensorView out) { + RuntimeCheck(is_dtype(ids.dtype(), kDLInt, 32) && ids.ndim() == 1 && ids.IsContiguous(), "ids must be a contiguous int32 vector"); + RuntimeCheck(w.ndim() == 2, "get_rows takes w[rows, row_bytes]"); + const DeviceGuard guard(w.device()); + const ggml_type type = (ggml_type) type_id; + check_weight(w, type, k); + const int64_t n = ids.size(0); + RuntimeCheck(out.IsContiguous() && out.ndim() == 2 && out.size(0) == n && out.size(1) == k, "out must be a contiguous [n, K] tensor"); + const ggml_type otype = float_type(out); + if (n == 0) { + return; + } + const size_t nb01 = w.stride(0), esz = otype == GGML_TYPE_F32 ? 4 : 2; + get_rows_cuda(w.data_ptr(), type, (const int32_t *) ids.data_ptr(), out.data_ptr(), otype, k, nb01, nb01 * w.size(0), nb01 * w.size(0), + n, 1, 1, 4, 4 * n, 4 * n, k * esz, k * esz * n, k * esz * n, current_stream(w.device())); +} + +bool op_uses_mmvq(int64_t type_id, int64_t m) { return ggml_cuda_should_use_mmvq((ggml_type) type_id, device_cc(), m); } + +int64_t op_mmvq_max_tokens(int64_t type_id) { + return std::min(MMVQ_MAX_BATCH_SIZE, get_mmvq_mmid_max_batch((ggml_type) type_id, device_cc())); +} + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(quantize_q8_1, op_quantize_q8_1); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(q8_1_row_bytes, q8_1_row_bytes); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(mmq_tail_bytes, mmq_tail_bytes); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(mul_mat, op_mul_mat); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(mul_mat_id, op_mul_mat_id); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(mul_mat_id_glu, op_mul_mat_id_glu); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(moe_weighted_sum, op_moe_weighted_sum); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(dequantize, op_dequantize); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(get_rows, op_get_rows); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(uses_mmvq, op_uses_mmvq); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(mmvq_max_tokens, op_mmvq_max_tokens); diff --git a/python/freetoken/kernel/csrc/ggml/freetoken/ggml_shim.cu b/python/freetoken/kernel/csrc/ggml/freetoken/ggml_shim.cu new file mode 100644 index 0000000000..80ef72855c --- /dev/null +++ b/python/freetoken/kernel/csrc/ggml/freetoken/ggml_shim.cu @@ -0,0 +1,255 @@ +// The ggml-base and ggml-cuda.cu host symbols that the vendored ggml-cuda kernels reference. +// Tensor descriptors built by ggml_ops.cu never have a backend buffer, and temporary buffers come from tvm-ffi's env allocator (torch's caching allocator under torch). + +#include "common.cuh" + +#include +#include + +#include +#include +#include +#include +#include + +namespace { + +struct TypeInfo { + const char * name = nullptr; + int64_t blck_size = 0; + size_t type_size = 0; + bool quantized = false; +}; + +const TypeInfo * type_infos() { + static TypeInfo t[GGML_TYPE_COUNT]; + static const bool init = [] { + auto set = [](ggml_type type, const char * name, int64_t blck, size_t size, bool quantized) { + t[type] = {name, blck, size, quantized}; + }; + set(GGML_TYPE_F32, "f32", 1, sizeof(float), false); + set(GGML_TYPE_F16, "f16", 1, sizeof(ggml_fp16_t), false); + set(GGML_TYPE_BF16, "bf16", 1, sizeof(ggml_bf16_t), false); + set(GGML_TYPE_F64, "f64", 1, sizeof(double), false); + set(GGML_TYPE_I8, "i8", 1, sizeof(int8_t), false); + set(GGML_TYPE_I16, "i16", 1, sizeof(int16_t), false); + set(GGML_TYPE_I32, "i32", 1, sizeof(int32_t), false); + set(GGML_TYPE_I64, "i64", 1, sizeof(int64_t), false); + set(GGML_TYPE_Q1_0, "q1_0", QK1_0, sizeof(block_q1_0), true); + set(GGML_TYPE_Q2_0, "q2_0", QK2_0, sizeof(block_q2_0), true); + set(GGML_TYPE_Q4_0, "q4_0", QK4_0, sizeof(block_q4_0), true); + set(GGML_TYPE_Q4_1, "q4_1", QK4_1, sizeof(block_q4_1), true); + set(GGML_TYPE_Q5_0, "q5_0", QK5_0, sizeof(block_q5_0), true); + set(GGML_TYPE_Q5_1, "q5_1", QK5_1, sizeof(block_q5_1), true); + set(GGML_TYPE_Q8_0, "q8_0", QK8_0, sizeof(block_q8_0), true); + set(GGML_TYPE_Q8_1, "q8_1", QK8_1, sizeof(block_q8_1), true); + set(GGML_TYPE_MXFP4, "mxfp4", QK_MXFP4, sizeof(block_mxfp4), true); + set(GGML_TYPE_NVFP4, "nvfp4", QK_NVFP4, sizeof(block_nvfp4), true); + set(GGML_TYPE_TQ1_0, "tq1_0", QK_K, sizeof(block_tq1_0), true); + set(GGML_TYPE_TQ2_0, "tq2_0", QK_K, sizeof(block_tq2_0), true); + set(GGML_TYPE_Q2_K, "q2_K", QK_K, sizeof(block_q2_K), true); + set(GGML_TYPE_Q3_K, "q3_K", QK_K, sizeof(block_q3_K), true); + set(GGML_TYPE_Q4_K, "q4_K", QK_K, sizeof(block_q4_K), true); + set(GGML_TYPE_Q5_K, "q5_K", QK_K, sizeof(block_q5_K), true); + set(GGML_TYPE_Q6_K, "q6_K", QK_K, sizeof(block_q6_K), true); + set(GGML_TYPE_Q8_K, "q8_K", QK_K, sizeof(block_q8_K), true); + set(GGML_TYPE_IQ2_XXS, "iq2_xxs", QK_K, sizeof(block_iq2_xxs), true); + set(GGML_TYPE_IQ2_XS, "iq2_xs", QK_K, sizeof(block_iq2_xs), true); + set(GGML_TYPE_IQ2_S, "iq2_s", QK_K, sizeof(block_iq2_s), true); + set(GGML_TYPE_IQ3_XXS, "iq3_xxs", QK_K, sizeof(block_iq3_xxs), true); + set(GGML_TYPE_IQ3_S, "iq3_s", QK_K, sizeof(block_iq3_s), true); + set(GGML_TYPE_IQ1_S, "iq1_s", QK_K, sizeof(block_iq1_s), true); + set(GGML_TYPE_IQ1_M, "iq1_m", QK_K, sizeof(block_iq1_m), true); + set(GGML_TYPE_IQ4_NL, "iq4_nl", QK4_NL, sizeof(block_iq4_nl), true); + set(GGML_TYPE_IQ4_XS, "iq4_xs", QK_K, sizeof(block_iq4_xs), true); + return true; + }(); + (void) init; + return t; +} + +const TypeInfo & type_info(ggml_type type) { + if (type < 0 || type >= GGML_TYPE_COUNT || type_infos()[type].name == nullptr) { + throw std::invalid_argument("unknown ggml type " + std::to_string((int) type)); + } + return type_infos()[type]; +} + +struct EnvPool : ggml_cuda_pool { + explicit EnvPool(int device) : device(device) {} + + void * alloc(size_t size, size_t * actual_size) override { + *actual_size = size; + if (size == 0) { + return nullptr; + } +#ifdef GGML_USE_HIP + constexpr DLDeviceType device_type = kDLROCM; +#else + constexpr DLDeviceType device_type = kDLCUDA; +#endif + auto buf = tvm::ffi::Tensor::FromEnvAlloc(TVMFFIEnvTensorAlloc, tvm::ffi::Shape({(int64_t) size}), DLDataType{kDLUInt, 8, 1}, + DLDevice{device_type, device}); + void * ptr = buf.data_ptr(); + live.emplace(ptr, std::move(buf)); + return ptr; + } + + void free(void * ptr, size_t) override { live.erase(ptr); } + + int device; + std::unordered_map live; +}; + +} // namespace + +void ggml_abort(const char * file, int line, const char * fmt, ...) { + char msg[1024]; + va_list args; + va_start(args, fmt); + vsnprintf(msg, sizeof(msg), fmt, args); + va_end(args); + throw std::runtime_error(std::string("ggml: ") + file + ":" + std::to_string(line) + ": " + msg); +} + +void ggml_log_internal(enum ggml_log_level level, const char * format, ...) { + if (level < GGML_LOG_LEVEL_WARN || level == GGML_LOG_LEVEL_CONT) { + return; + } + va_list args; + va_start(args, format); + vfprintf(stderr, format, args); + va_end(args); +} + +int64_t ggml_blck_size(enum ggml_type type) { return type_info(type).blck_size; } + +size_t ggml_type_size(enum ggml_type type) { return type_info(type).type_size; } + +const char * ggml_type_name(enum ggml_type type) { return type_info(type).name; } + +bool ggml_is_quantized(enum ggml_type type) { return type_info(type).quantized; } + +size_t ggml_row_size(enum ggml_type type, int64_t ne) { return ggml_type_size(type) * ne / ggml_blck_size(type); } + +size_t ggml_element_size(const struct ggml_tensor * tensor) { return ggml_type_size(tensor->type); } + +int64_t ggml_nelements(const struct ggml_tensor * tensor) { + return tensor->ne[0] * tensor->ne[1] * tensor->ne[2] * tensor->ne[3]; +} + +size_t ggml_nbytes(const struct ggml_tensor * tensor) { + for (int i = 0; i < GGML_MAX_DIMS; ++i) { + if (tensor->ne[i] <= 0) { + return 0; + } + } + size_t nbytes; + const size_t blck_size = ggml_blck_size(tensor->type); + if (blck_size == 1) { + nbytes = ggml_type_size(tensor->type); + for (int i = 0; i < GGML_MAX_DIMS; ++i) { + nbytes += (tensor->ne[i] - 1) * tensor->nb[i]; + } + } else { + nbytes = tensor->ne[0] * tensor->nb[0] / blck_size; + for (int i = 1; i < GGML_MAX_DIMS; ++i) { + nbytes += (tensor->ne[i] - 1) * tensor->nb[i]; + } + } + return nbytes; +} + +bool ggml_is_matrix(const struct ggml_tensor * tensor) { return tensor->ne[2] == 1 && tensor->ne[3] == 1; } + +bool ggml_is_contiguous(const struct ggml_tensor * tensor) { + size_t next_nb = ggml_type_size(tensor->type); + if (tensor->ne[0] != ggml_blck_size(tensor->type) && tensor->nb[0] != next_nb) { + return false; + } + next_nb *= tensor->ne[0] / ggml_blck_size(tensor->type); + for (int i = 1; i < GGML_MAX_DIMS; i++) { + if (tensor->ne[i] != 1 && tensor->nb[i] != next_nb) { + return false; + } + next_nb *= tensor->ne[i]; + } + return true; +} + +bool ggml_is_contiguously_allocated(const struct ggml_tensor * tensor) { + return ggml_nbytes(tensor) == ggml_nelements(tensor) * ggml_type_size(tensor->type) / ggml_blck_size(tensor->type); +} + +bool ggml_are_same_shape(const struct ggml_tensor * t0, const struct ggml_tensor * t1) { + return t0->ne[0] == t1->ne[0] && t0->ne[1] == t1->ne[1] && t0->ne[2] == t1->ne[2] && t0->ne[3] == t1->ne[3]; +} + +bool ggml_are_same_stride(const struct ggml_tensor * t0, const struct ggml_tensor * t1) { + return t0->nb[0] == t1->nb[0] && t0->nb[1] == t1->nb[1] && t0->nb[2] == t1->nb[2] && t0->nb[3] == t1->nb[3]; +} + +enum ggml_backend_buffer_usage ggml_backend_buffer_get_usage(ggml_backend_buffer_t) { + return GGML_BACKEND_BUFFER_USAGE_WEIGHTS; +} + +size_t ggml_backend_buffer_get_alloc_size(ggml_backend_buffer_t, const struct ggml_tensor * tensor) { + return ggml_nbytes(tensor); +} + +[[noreturn]] void ggml_cuda_error(const char * stmt, const char * func, const char * file, int line, const char * msg) { + throw std::runtime_error(std::string("ggml-cuda: ") + msg + " in " + func + " at " + file + ":" + std::to_string(line) + ": " + stmt); +} + +int ggml_cuda_get_device() { + int id; + CUDA_CHECK(cudaGetDevice(&id)); + return id; +} + +void ggml_cuda_set_device(int device) { + int current; + CUDA_CHECK(cudaGetDevice(¤t)); + if (current != device) { + CUDA_CHECK(cudaSetDevice(device)); + } +} + +const ggml_cuda_device_info & ggml_cuda_info() { + static const ggml_cuda_device_info info = [] { + ggml_cuda_device_info in = {}; + CUDA_CHECK(cudaGetDeviceCount(&in.physical_device_count)); + GGML_ASSERT(in.physical_device_count <= GGML_CUDA_MAX_DEVICES); + in.device_count = in.physical_device_count; + for (int id = 0; id < in.device_count; ++id) { + cudaDeviceProp prop; + CUDA_CHECK(cudaGetDeviceProperties(&prop, id)); + auto & d = in.devices[id]; +#ifdef GGML_USE_HIP + GGML_ASSERT(std::string(prop.gcnArchName).starts_with("gfx")); + d.cc = GGML_CUDA_CC_OFFSET_AMD + static_cast(std::strtol(prop.gcnArchName + 3, nullptr, 16)); +#else + d.cc = 100 * prop.major + 10 * prop.minor; +#endif + d.nsm = prop.multiProcessorCount; + d.smpb = prop.sharedMemPerBlock; + d.smpbo = prop.sharedMemPerBlockOptin; + d.warp_size = prop.warpSize; + d.total_vram = prop.totalGlobalMem; + d.physical_device = id; + d.physical_share_count = 1; + int coop = 0; + CUDA_CHECK(cudaDeviceGetAttribute(&coop, cudaDevAttrCooperativeLaunch, id)); + d.supports_cooperative_launch = coop != 0; + } + return in; + }(); + return info; +} + +// The streams belong to the caller, so there is nothing to release. +ggml_backend_cuda_context::~ggml_backend_cuda_context() {} + +std::unique_ptr ggml_backend_cuda_context::new_pool_for_device(int device, int) { + return std::make_unique(device); +} diff --git a/python/freetoken/kernel/csrc/ggml/freetoken/variant_module.cpp b/python/freetoken/kernel/csrc/ggml/freetoken/variant_module.cpp new file mode 100644 index 0000000000..8f3d3e5246 --- /dev/null +++ b/python/freetoken/kernel/csrc/ggml/freetoken/variant_module.cpp @@ -0,0 +1,32 @@ +// One ggml-cpu variant as a Python module: score() is 0 unless this CPU has every ISA extension the variant is built for, and api() is the address of its kernel table. +// setup.py builds this file without the variant's ISA flags, like upstream's -feats target, so importing the module and calling score() is safe on any x86 CPU. +#define PY_SSIZE_T_CLEAN +#include + +#include "cpu_kernels.h" +#include "ggml-cpu/arch/x86/cpu-feats.cpp" + +#define FT_CAT_(a, b) a##b +#define FT_CAT(a, b) FT_CAT_(a, b) +#define FT_STR_(a) #a +#define FT_STR(a) FT_STR_(a) + +namespace { + +const ft_ggml_cpu_api kApi = {FT_GGML_CPU_ABI, FT_STR(FT_GGML_VARIANT), ft_ggml_cpu_init, ft_ggml_cpu_row_kernel, ft_ggml_type_name}; + +PyObject* score(PyObject*, PyObject*) { return PyLong_FromLong(ggml_backend_cpu_x86_score()); } + +PyObject* api(PyObject*, PyObject*) { return PyLong_FromVoidPtr(const_cast(&kApi)); } + +PyMethodDef kMethods[] = { + {"score", score, METH_NOARGS, "0 if this CPU lacks an ISA extension the variant is built for, else upstream's variant score."}, + {"api", api, METH_NOARGS, "The address of the variant's kernel table, for _cpu_moe.set_ggml_cpu."}, + {nullptr, nullptr, 0, nullptr}, +}; + +PyModuleDef kModule = {PyModuleDef_HEAD_INIT, "_" FT_STR(FT_GGML_VARIANT), nullptr, -1, kMethods}; + +} // namespace + +PyMODINIT_FUNC FT_CAT(PyInit__, FT_GGML_VARIANT)(void) { return PyModule_Create(&kModule); } diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cpu/arch-fallback.h b/python/freetoken/kernel/csrc/ggml/ggml-cpu/arch-fallback.h new file mode 100644 index 0000000000..4e59d0ae0a --- /dev/null +++ b/python/freetoken/kernel/csrc/ggml/ggml-cpu/arch-fallback.h @@ -0,0 +1,375 @@ + +#pragma once + +// Rename `_generic` functions if no native implementation is available. +// This effectively selects the generic implementation. + +#if defined(GGML_CPU_GENERIC) +// quants.c +#define quantize_row_q8_0_generic quantize_row_q8_0 +#define quantize_row_q8_1_generic quantize_row_q8_1 +#define quantize_row_q8_K_generic quantize_row_q8_K +#define ggml_vec_dot_q4_0_q8_0_generic ggml_vec_dot_q4_0_q8_0 +#define ggml_vec_dot_q4_1_q8_1_generic ggml_vec_dot_q4_1_q8_1 +#define ggml_vec_dot_q5_0_q8_0_generic ggml_vec_dot_q5_0_q8_0 +#define ggml_vec_dot_q5_1_q8_1_generic ggml_vec_dot_q5_1_q8_1 +#define ggml_vec_dot_q8_0_q8_0_generic ggml_vec_dot_q8_0_q8_0 +#define ggml_vec_dot_mxfp4_q8_0_generic ggml_vec_dot_mxfp4_q8_0 +#define ggml_vec_dot_nvfp4_q8_0_generic ggml_vec_dot_nvfp4_q8_0 +#define ggml_vec_dot_q1_0_q8_0_generic ggml_vec_dot_q1_0_q8_0 +#define ggml_vec_dot_q2_0_q8_0_generic ggml_vec_dot_q2_0_q8_0 +#define ggml_vec_dot_tq1_0_q8_K_generic ggml_vec_dot_tq1_0_q8_K +#define ggml_vec_dot_tq2_0_q8_K_generic ggml_vec_dot_tq2_0_q8_K +#define ggml_vec_dot_q2_K_q8_K_generic ggml_vec_dot_q2_K_q8_K +#define ggml_vec_dot_q3_K_q8_K_generic ggml_vec_dot_q3_K_q8_K +#define ggml_vec_dot_q4_K_q8_K_generic ggml_vec_dot_q4_K_q8_K +#define ggml_vec_dot_q5_K_q8_K_generic ggml_vec_dot_q5_K_q8_K +#define ggml_vec_dot_q6_K_q8_K_generic ggml_vec_dot_q6_K_q8_K +#define ggml_vec_dot_iq2_xxs_q8_K_generic ggml_vec_dot_iq2_xxs_q8_K +#define ggml_vec_dot_iq2_xs_q8_K_generic ggml_vec_dot_iq2_xs_q8_K +#define ggml_vec_dot_iq2_s_q8_K_generic ggml_vec_dot_iq2_s_q8_K +#define ggml_vec_dot_iq3_xxs_q8_K_generic ggml_vec_dot_iq3_xxs_q8_K +#define ggml_vec_dot_iq3_s_q8_K_generic ggml_vec_dot_iq3_s_q8_K +#define ggml_vec_dot_iq1_s_q8_K_generic ggml_vec_dot_iq1_s_q8_K +#define ggml_vec_dot_iq1_m_q8_K_generic ggml_vec_dot_iq1_m_q8_K +#define ggml_vec_dot_iq4_nl_q8_0_generic ggml_vec_dot_iq4_nl_q8_0 +#define ggml_vec_dot_iq4_xs_q8_K_generic ggml_vec_dot_iq4_xs_q8_K +// repack.cpp +#define ggml_quantize_mat_q8_0_4x4_generic ggml_quantize_mat_q8_0_4x4 +#define ggml_quantize_mat_q8_0_4x8_generic ggml_quantize_mat_q8_0_4x8 +#define ggml_quantize_mat_q8_K_4x4_generic ggml_quantize_mat_q8_K_4x4 +#define ggml_quantize_mat_q8_K_4x8_generic ggml_quantize_mat_q8_K_4x8 +#define ggml_gemv_q1_0_4x4_q8_0_generic ggml_gemv_q1_0_4x4_q8_0 +#define ggml_gemv_q1_0_4x8_q8_0_generic ggml_gemv_q1_0_4x8_q8_0 +#define ggml_gemv_q4_0_4x4_q8_0_generic ggml_gemv_q4_0_4x4_q8_0 +#define ggml_gemv_q4_0_4x8_q8_0_generic ggml_gemv_q4_0_4x8_q8_0 +#define ggml_gemv_q4_0_8x8_q8_0_generic ggml_gemv_q4_0_8x8_q8_0 +#define ggml_gemv_q2_K_8x8_q8_K_generic ggml_gemv_q2_K_8x8_q8_K +#define ggml_gemv_q4_K_8x4_q8_K_generic ggml_gemv_q4_K_8x4_q8_K +#define ggml_gemv_q4_K_8x8_q8_K_generic ggml_gemv_q4_K_8x8_q8_K +#define ggml_gemv_q5_K_8x4_q8_K_generic ggml_gemv_q5_K_8x4_q8_K +#define ggml_gemv_q5_K_8x8_q8_K_generic ggml_gemv_q5_K_8x8_q8_K +#define ggml_gemv_q6_K_8x4_q8_K_generic ggml_gemv_q6_K_8x4_q8_K +#define ggml_gemv_q6_K_8x8_q8_K_generic ggml_gemv_q6_K_8x8_q8_K +#define ggml_gemv_iq4_nl_4x4_q8_0_generic ggml_gemv_iq4_nl_4x4_q8_0 +#define ggml_gemv_iq4_nl_8x8_q8_0_generic ggml_gemv_iq4_nl_8x8_q8_0 +#define ggml_gemv_mxfp4_4x4_q8_0_generic ggml_gemv_mxfp4_4x4_q8_0 +#define ggml_gemv_mxfp4_8x8_q8_0_generic ggml_gemv_mxfp4_8x8_q8_0 +#define ggml_gemv_q8_0_4x4_q8_0_generic ggml_gemv_q8_0_4x4_q8_0 +#define ggml_gemv_q8_0_4x8_q8_0_generic ggml_gemv_q8_0_4x8_q8_0 +#define ggml_gemm_q1_0_4x4_q8_0_generic ggml_gemm_q1_0_4x4_q8_0 +#define ggml_gemm_q1_0_4x8_q8_0_generic ggml_gemm_q1_0_4x8_q8_0 +#define ggml_gemm_q4_0_4x4_q8_0_generic ggml_gemm_q4_0_4x4_q8_0 +#define ggml_gemm_q4_0_4x8_q8_0_generic ggml_gemm_q4_0_4x8_q8_0 +#define ggml_gemm_q4_0_8x8_q8_0_generic ggml_gemm_q4_0_8x8_q8_0 +#define ggml_gemm_q2_K_8x8_q8_K_generic ggml_gemm_q2_K_8x8_q8_K +#define ggml_gemm_q4_K_8x4_q8_K_generic ggml_gemm_q4_K_8x4_q8_K +#define ggml_gemm_q4_K_8x8_q8_K_generic ggml_gemm_q4_K_8x8_q8_K +#define ggml_gemm_q5_K_8x4_q8_K_generic ggml_gemm_q5_K_8x4_q8_K +#define ggml_gemm_q5_K_8x8_q8_K_generic ggml_gemm_q5_K_8x8_q8_K +#define ggml_gemm_q6_K_8x4_q8_K_generic ggml_gemm_q6_K_8x4_q8_K +#define ggml_gemm_q6_K_8x8_q8_K_generic ggml_gemm_q6_K_8x8_q8_K +#define ggml_gemm_iq4_nl_4x4_q8_0_generic ggml_gemm_iq4_nl_4x4_q8_0 +#define ggml_gemm_iq4_nl_8x8_q8_0_generic ggml_gemm_iq4_nl_8x8_q8_0 +#define ggml_gemm_mxfp4_4x4_q8_0_generic ggml_gemm_mxfp4_4x4_q8_0 +#define ggml_gemm_mxfp4_8x8_q8_0_generic ggml_gemm_mxfp4_8x8_q8_0 +#define ggml_gemm_q8_0_4x4_q8_0_generic ggml_gemm_q8_0_4x4_q8_0 +#define ggml_gemm_q8_0_4x8_q8_0_generic ggml_gemm_q8_0_4x8_q8_0 +#elif defined(__aarch64__) || defined(__arm__) || defined(_M_ARM) || defined(_M_ARM64) || defined(_M_ARM64EC) +// repack.cpp +#define ggml_quantize_mat_q8_K_4x4_generic ggml_quantize_mat_q8_K_4x4 +#define ggml_quantize_mat_q8_K_4x8_generic ggml_quantize_mat_q8_K_4x8 +#define ggml_gemv_iq4_nl_8x8_q8_0_generic ggml_gemv_iq4_nl_8x8_q8_0 +#define ggml_gemv_mxfp4_8x8_q8_0_generic ggml_gemv_mxfp4_8x8_q8_0 +#define ggml_gemv_q2_K_8x8_q8_K_generic ggml_gemv_q2_K_8x8_q8_K +#define ggml_gemm_iq4_nl_8x8_q8_0_generic ggml_gemm_iq4_nl_8x8_q8_0 +#define ggml_gemm_mxfp4_8x8_q8_0_generic ggml_gemm_mxfp4_8x8_q8_0 +#define ggml_gemm_q2_K_8x8_q8_K_generic ggml_gemm_q2_K_8x8_q8_K +#elif defined(__x86_64__) || defined(__i386__) || defined(_M_IX86) || defined(_M_X64) +// repack.cpp +#define ggml_quantize_mat_q8_0_4x4_generic ggml_quantize_mat_q8_0_4x4 +#define ggml_quantize_mat_q8_K_4x4_generic ggml_quantize_mat_q8_K_4x4 +#define ggml_gemv_q1_0_4x4_q8_0_generic ggml_gemv_q1_0_4x4_q8_0 +#define ggml_gemv_q1_0_4x8_q8_0_generic ggml_gemv_q1_0_4x8_q8_0 +#define ggml_gemv_q4_0_4x4_q8_0_generic ggml_gemv_q4_0_4x4_q8_0 +#define ggml_gemv_q4_0_4x8_q8_0_generic ggml_gemv_q4_0_4x8_q8_0 +#define ggml_gemv_q4_K_8x4_q8_K_generic ggml_gemv_q4_K_8x4_q8_K +#define ggml_gemv_q5_K_8x4_q8_K_generic ggml_gemv_q5_K_8x4_q8_K +#define ggml_gemv_q5_K_8x8_q8_K_generic ggml_gemv_q5_K_8x8_q8_K +#define ggml_gemv_q6_K_8x4_q8_K_generic ggml_gemv_q6_K_8x4_q8_K +#define ggml_gemv_q6_K_8x8_q8_K_generic ggml_gemv_q6_K_8x8_q8_K +#define ggml_gemv_iq4_nl_4x4_q8_0_generic ggml_gemv_iq4_nl_4x4_q8_0 +#define ggml_gemv_mxfp4_4x4_q8_0_generic ggml_gemv_mxfp4_4x4_q8_0 +#define ggml_gemv_q8_0_4x4_q8_0_generic ggml_gemv_q8_0_4x4_q8_0 +#define ggml_gemv_q8_0_4x8_q8_0_generic ggml_gemv_q8_0_4x8_q8_0 +#define ggml_gemm_q1_0_4x4_q8_0_generic ggml_gemm_q1_0_4x4_q8_0 +#define ggml_gemm_q1_0_4x8_q8_0_generic ggml_gemm_q1_0_4x8_q8_0 +#define ggml_gemm_q4_0_4x4_q8_0_generic ggml_gemm_q4_0_4x4_q8_0 +#define ggml_gemm_q4_0_4x8_q8_0_generic ggml_gemm_q4_0_4x8_q8_0 +#define ggml_gemm_q4_K_8x4_q8_K_generic ggml_gemm_q4_K_8x4_q8_K +#define ggml_gemm_q5_K_8x4_q8_K_generic ggml_gemm_q5_K_8x4_q8_K +#define ggml_gemm_q5_K_8x8_q8_K_generic ggml_gemm_q5_K_8x8_q8_K +#define ggml_gemm_q6_K_8x4_q8_K_generic ggml_gemm_q6_K_8x4_q8_K +#define ggml_gemm_q6_K_8x8_q8_K_generic ggml_gemm_q6_K_8x8_q8_K +#define ggml_gemm_iq4_nl_4x4_q8_0_generic ggml_gemm_iq4_nl_4x4_q8_0 +#define ggml_gemm_mxfp4_4x4_q8_0_generic ggml_gemm_mxfp4_4x4_q8_0 +#define ggml_gemm_q8_0_4x4_q8_0_generic ggml_gemm_q8_0_4x4_q8_0 +#define ggml_gemm_q8_0_4x8_q8_0_generic ggml_gemm_q8_0_4x8_q8_0 +#elif defined(__POWERPC__) || defined(__powerpc__) +// ref: https://github.com/ggml-org/llama.cpp/pull/14146#issuecomment-2972561679 +// quants.c +#define quantize_row_q8_K_generic quantize_row_q8_K +#define ggml_vec_dot_nvfp4_q8_0_generic ggml_vec_dot_nvfp4_q8_0 +#define ggml_vec_dot_q1_0_q8_0_generic ggml_vec_dot_q1_0_q8_0 +#define ggml_vec_dot_q2_0_q8_0_generic ggml_vec_dot_q2_0_q8_0 +#define ggml_vec_dot_tq1_0_q8_K_generic ggml_vec_dot_tq1_0_q8_K +#define ggml_vec_dot_tq2_0_q8_K_generic ggml_vec_dot_tq2_0_q8_K +#define ggml_vec_dot_iq1_m_q8_K_generic ggml_vec_dot_iq1_m_q8_K +// repack.cpp +#define ggml_quantize_mat_q8_0_4x4_generic ggml_quantize_mat_q8_0_4x4 +#define ggml_quantize_mat_q8_0_4x8_generic ggml_quantize_mat_q8_0_4x8 +#define ggml_quantize_mat_q8_K_4x4_generic ggml_quantize_mat_q8_K_4x4 +#define ggml_quantize_mat_q8_K_4x8_generic ggml_quantize_mat_q8_K_4x8 +#define ggml_gemv_q1_0_4x4_q8_0_generic ggml_gemv_q1_0_4x4_q8_0 +#define ggml_gemv_q1_0_4x8_q8_0_generic ggml_gemv_q1_0_4x8_q8_0 +#define ggml_gemv_q4_0_4x4_q8_0_generic ggml_gemv_q4_0_4x4_q8_0 +#define ggml_gemv_q4_0_4x8_q8_0_generic ggml_gemv_q4_0_4x8_q8_0 +#define ggml_gemv_q4_0_8x8_q8_0_generic ggml_gemv_q4_0_8x8_q8_0 +#define ggml_gemv_q2_K_8x8_q8_K_generic ggml_gemv_q2_K_8x8_q8_K +#define ggml_gemv_q4_K_8x4_q8_K_generic ggml_gemv_q4_K_8x4_q8_K +#define ggml_gemv_q4_K_8x8_q8_K_generic ggml_gemv_q4_K_8x8_q8_K +#define ggml_gemv_q5_K_8x4_q8_K_generic ggml_gemv_q5_K_8x4_q8_K +#define ggml_gemv_q5_K_8x8_q8_K_generic ggml_gemv_q5_K_8x8_q8_K +#define ggml_gemv_q6_K_8x4_q8_K_generic ggml_gemv_q6_K_8x4_q8_K +#define ggml_gemv_q6_K_8x8_q8_K_generic ggml_gemv_q6_K_8x8_q8_K +#define ggml_gemv_iq4_nl_4x4_q8_0_generic ggml_gemv_iq4_nl_4x4_q8_0 +#define ggml_gemv_iq4_nl_8x8_q8_0_generic ggml_gemv_iq4_nl_8x8_q8_0 +#define ggml_gemv_mxfp4_4x4_q8_0_generic ggml_gemv_mxfp4_4x4_q8_0 +#define ggml_gemv_mxfp4_8x8_q8_0_generic ggml_gemv_mxfp4_8x8_q8_0 +#define ggml_gemv_q8_0_4x4_q8_0_generic ggml_gemv_q8_0_4x4_q8_0 +#define ggml_gemv_q8_0_4x8_q8_0_generic ggml_gemv_q8_0_4x8_q8_0 +#define ggml_gemm_q1_0_4x4_q8_0_generic ggml_gemm_q1_0_4x4_q8_0 +#define ggml_gemm_q1_0_4x8_q8_0_generic ggml_gemm_q1_0_4x8_q8_0 +#define ggml_gemm_q4_0_4x4_q8_0_generic ggml_gemm_q4_0_4x4_q8_0 +#define ggml_gemm_q4_0_4x8_q8_0_generic ggml_gemm_q4_0_4x8_q8_0 +#define ggml_gemm_q4_0_8x8_q8_0_generic ggml_gemm_q4_0_8x8_q8_0 +#define ggml_gemm_q2_K_8x8_q8_K_generic ggml_gemm_q2_K_8x8_q8_K +#define ggml_gemm_q4_K_8x4_q8_K_generic ggml_gemm_q4_K_8x4_q8_K +#define ggml_gemm_q4_K_8x8_q8_K_generic ggml_gemm_q4_K_8x8_q8_K +#define ggml_gemm_q5_K_8x4_q8_K_generic ggml_gemm_q5_K_8x4_q8_K +#define ggml_gemm_q5_K_8x8_q8_K_generic ggml_gemm_q5_K_8x8_q8_K +#define ggml_gemm_q6_K_8x4_q8_K_generic ggml_gemm_q6_K_8x4_q8_K +#define ggml_gemm_q6_K_8x8_q8_K_generic ggml_gemm_q6_K_8x8_q8_K +#define ggml_gemm_iq4_nl_4x4_q8_0_generic ggml_gemm_iq4_nl_4x4_q8_0 +#define ggml_gemm_iq4_nl_8x8_q8_0_generic ggml_gemm_iq4_nl_8x8_q8_0 +#define ggml_gemm_mxfp4_4x4_q8_0_generic ggml_gemm_mxfp4_4x4_q8_0 +#define ggml_gemm_mxfp4_8x8_q8_0_generic ggml_gemm_mxfp4_8x8_q8_0 +#define ggml_gemm_q8_0_4x4_q8_0_generic ggml_gemm_q8_0_4x4_q8_0 +#define ggml_gemm_q8_0_4x8_q8_0_generic ggml_gemm_q8_0_4x8_q8_0 +#elif defined(__loongarch64) +// quants.c +#define quantize_row_q8_K_generic quantize_row_q8_K +#define ggml_vec_dot_tq1_0_q8_K_generic ggml_vec_dot_tq1_0_q8_K +#define ggml_vec_dot_tq2_0_q8_K_generic ggml_vec_dot_tq2_0_q8_K +#define ggml_vec_dot_iq1_m_q8_K_generic ggml_vec_dot_iq1_m_q8_K +#define ggml_vec_dot_mxfp4_q8_0_generic ggml_vec_dot_mxfp4_q8_0 +#define ggml_vec_dot_nvfp4_q8_0_generic ggml_vec_dot_nvfp4_q8_0 +#define ggml_vec_dot_q1_0_q8_0_generic ggml_vec_dot_q1_0_q8_0 +#define ggml_vec_dot_q2_0_q8_0_generic ggml_vec_dot_q2_0_q8_0 +// repack.cpp +#define ggml_quantize_mat_q8_0_4x4_generic ggml_quantize_mat_q8_0_4x4 +#define ggml_quantize_mat_q8_0_4x8_generic ggml_quantize_mat_q8_0_4x8 +#define ggml_quantize_mat_q8_K_4x4_generic ggml_quantize_mat_q8_K_4x4 +#define ggml_quantize_mat_q8_K_4x8_generic ggml_quantize_mat_q8_K_4x8 +#define ggml_gemv_q1_0_4x4_q8_0_generic ggml_gemv_q1_0_4x4_q8_0 +#define ggml_gemv_q1_0_4x8_q8_0_generic ggml_gemv_q1_0_4x8_q8_0 +#define ggml_gemv_q4_0_4x4_q8_0_generic ggml_gemv_q4_0_4x4_q8_0 +#define ggml_gemv_q4_0_4x8_q8_0_generic ggml_gemv_q4_0_4x8_q8_0 +#define ggml_gemv_q4_0_8x8_q8_0_generic ggml_gemv_q4_0_8x8_q8_0 +#define ggml_gemv_q2_K_8x8_q8_K_generic ggml_gemv_q2_K_8x8_q8_K +#define ggml_gemv_q4_K_8x4_q8_K_generic ggml_gemv_q4_K_8x4_q8_K +#define ggml_gemv_q4_K_8x8_q8_K_generic ggml_gemv_q4_K_8x8_q8_K +#define ggml_gemv_q5_K_8x4_q8_K_generic ggml_gemv_q5_K_8x4_q8_K +#define ggml_gemv_q5_K_8x8_q8_K_generic ggml_gemv_q5_K_8x8_q8_K +#define ggml_gemv_q6_K_8x4_q8_K_generic ggml_gemv_q6_K_8x4_q8_K +#define ggml_gemv_q6_K_8x8_q8_K_generic ggml_gemv_q6_K_8x8_q8_K +#define ggml_gemv_iq4_nl_4x4_q8_0_generic ggml_gemv_iq4_nl_4x4_q8_0 +#define ggml_gemv_iq4_nl_8x8_q8_0_generic ggml_gemv_iq4_nl_8x8_q8_0 +#define ggml_gemv_mxfp4_4x4_q8_0_generic ggml_gemv_mxfp4_4x4_q8_0 +#define ggml_gemv_mxfp4_8x8_q8_0_generic ggml_gemv_mxfp4_8x8_q8_0 +#define ggml_gemv_q8_0_4x4_q8_0_generic ggml_gemv_q8_0_4x4_q8_0 +#define ggml_gemv_q8_0_4x8_q8_0_generic ggml_gemv_q8_0_4x8_q8_0 +#define ggml_gemm_q1_0_4x4_q8_0_generic ggml_gemm_q1_0_4x4_q8_0 +#define ggml_gemm_q1_0_4x8_q8_0_generic ggml_gemm_q1_0_4x8_q8_0 +#define ggml_gemm_q4_0_4x4_q8_0_generic ggml_gemm_q4_0_4x4_q8_0 +#define ggml_gemm_q4_0_4x8_q8_0_generic ggml_gemm_q4_0_4x8_q8_0 +#define ggml_gemm_q4_0_8x8_q8_0_generic ggml_gemm_q4_0_8x8_q8_0 +#define ggml_gemm_q2_K_8x8_q8_K_generic ggml_gemm_q2_K_8x8_q8_K +#define ggml_gemm_q4_K_8x4_q8_K_generic ggml_gemm_q4_K_8x4_q8_K +#define ggml_gemm_q4_K_8x8_q8_K_generic ggml_gemm_q4_K_8x8_q8_K +#define ggml_gemm_q5_K_8x4_q8_K_generic ggml_gemm_q5_K_8x4_q8_K +#define ggml_gemm_q5_K_8x8_q8_K_generic ggml_gemm_q5_K_8x8_q8_K +#define ggml_gemm_q6_K_8x4_q8_K_generic ggml_gemm_q6_K_8x4_q8_K +#define ggml_gemm_q6_K_8x8_q8_K_generic ggml_gemm_q6_K_8x8_q8_K +#define ggml_gemm_iq4_nl_4x4_q8_0_generic ggml_gemm_iq4_nl_4x4_q8_0 +#define ggml_gemm_iq4_nl_8x8_q8_0_generic ggml_gemm_iq4_nl_8x8_q8_0 +#define ggml_gemm_mxfp4_4x4_q8_0_generic ggml_gemm_mxfp4_4x4_q8_0 +#define ggml_gemm_mxfp4_8x8_q8_0_generic ggml_gemm_mxfp4_8x8_q8_0 +#define ggml_gemm_q8_0_4x4_q8_0_generic ggml_gemm_q8_0_4x4_q8_0 +#define ggml_gemm_q8_0_4x8_q8_0_generic ggml_gemm_q8_0_4x8_q8_0 +#elif defined(__riscv) +// quants.c +#define ggml_vec_dot_nvfp4_q8_0_generic ggml_vec_dot_nvfp4_q8_0 +#define ggml_vec_dot_q2_0_q8_0_generic ggml_vec_dot_q2_0_q8_0 +// repack.cpp +#define ggml_quantize_mat_q8_0_4x1_generic ggml_quantize_mat_q8_0_4x1 +#define ggml_quantize_mat_q8_0_4x4_generic ggml_quantize_mat_q8_0_4x4 +#define ggml_quantize_mat_q8_K_4x1_generic ggml_quantize_mat_q8_K_4x1 +#define ggml_quantize_mat_q8_K_4x4_generic ggml_quantize_mat_q8_K_4x4 +#define ggml_quantize_mat_q8_K_4x8_generic ggml_quantize_mat_q8_K_4x8 +#define ggml_gemv_q1_0_4x4_q8_0_generic ggml_gemv_q1_0_4x4_q8_0 +#define ggml_gemv_q1_0_4x8_q8_0_generic ggml_gemv_q1_0_4x8_q8_0 +#define ggml_gemv_q4_0_4x4_q8_0_generic ggml_gemv_q4_0_4x4_q8_0 +#define ggml_gemv_q4_0_4x8_q8_0_generic ggml_gemv_q4_0_4x8_q8_0 +#define ggml_gemv_q2_K_8x8_q8_K_generic ggml_gemv_q2_K_8x8_q8_K +#define ggml_gemv_q4_K_8x4_q8_K_generic ggml_gemv_q4_K_8x4_q8_K +#define ggml_gemv_q4_K_8x8_q8_K_generic ggml_gemv_q4_K_8x8_q8_K +#define ggml_gemv_q5_K_8x4_q8_K_generic ggml_gemv_q5_K_8x4_q8_K +#define ggml_gemv_q5_K_8x8_q8_K_generic ggml_gemv_q5_K_8x8_q8_K +#define ggml_gemv_q6_K_8x4_q8_K_generic ggml_gemv_q6_K_8x4_q8_K +#define ggml_gemv_q6_K_8x8_q8_K_generic ggml_gemv_q6_K_8x8_q8_K +#define ggml_gemv_iq4_nl_4x4_q8_0_generic ggml_gemv_iq4_nl_4x4_q8_0 +#define ggml_gemv_iq4_nl_8x8_q8_0_generic ggml_gemv_iq4_nl_8x8_q8_0 +#define ggml_gemv_mxfp4_4x4_q8_0_generic ggml_gemv_mxfp4_4x4_q8_0 +#define ggml_gemv_mxfp4_8x8_q8_0_generic ggml_gemv_mxfp4_8x8_q8_0 +#define ggml_gemv_q8_0_4x4_q8_0_generic ggml_gemv_q8_0_4x4_q8_0 +#define ggml_gemv_q8_0_4x8_q8_0_generic ggml_gemv_q8_0_4x8_q8_0 +#define ggml_gemm_q1_0_4x4_q8_0_generic ggml_gemm_q1_0_4x4_q8_0 +#define ggml_gemm_q1_0_4x8_q8_0_generic ggml_gemm_q1_0_4x8_q8_0 +#define ggml_gemm_q4_0_4x4_q8_0_generic ggml_gemm_q4_0_4x4_q8_0 +#define ggml_gemm_q4_0_4x8_q8_0_generic ggml_gemm_q4_0_4x8_q8_0 +#define ggml_gemm_q2_K_8x8_q8_K_generic ggml_gemm_q2_K_8x8_q8_K +#define ggml_gemm_q4_K_8x4_q8_K_generic ggml_gemm_q4_K_8x4_q8_K +#define ggml_gemm_q4_K_8x8_q8_K_generic ggml_gemm_q4_K_8x8_q8_K +#define ggml_gemm_q5_K_8x4_q8_K_generic ggml_gemm_q5_K_8x4_q8_K +#define ggml_gemm_q5_K_8x8_q8_K_generic ggml_gemm_q5_K_8x8_q8_K +#define ggml_gemm_q6_K_8x4_q8_K_generic ggml_gemm_q6_K_8x4_q8_K +#define ggml_gemm_q6_K_8x8_q8_K_generic ggml_gemm_q6_K_8x8_q8_K +#define ggml_gemm_iq4_nl_4x4_q8_0_generic ggml_gemm_iq4_nl_4x4_q8_0 +#define ggml_gemm_iq4_nl_8x8_q8_0_generic ggml_gemm_iq4_nl_8x8_q8_0 +#define ggml_gemm_mxfp4_4x4_q8_0_generic ggml_gemm_mxfp4_4x4_q8_0 +#define ggml_gemm_mxfp4_8x8_q8_0_generic ggml_gemm_mxfp4_8x8_q8_0 +#define ggml_gemm_q8_0_4x4_q8_0_generic ggml_gemm_q8_0_4x4_q8_0 +#define ggml_gemm_q8_0_4x8_q8_0_generic ggml_gemm_q8_0_4x8_q8_0 +#elif defined(__s390x__) +// quants.c +#define quantize_row_q8_K_generic quantize_row_q8_K +#define ggml_vec_dot_nvfp4_q8_0_generic ggml_vec_dot_nvfp4_q8_0 +#define ggml_vec_dot_q2_0_q8_0_generic ggml_vec_dot_q2_0_q8_0 +#define ggml_vec_dot_tq1_0_q8_K_generic ggml_vec_dot_tq1_0_q8_K +#define ggml_vec_dot_tq2_0_q8_K_generic ggml_vec_dot_tq2_0_q8_K +#define ggml_vec_dot_q2_K_q8_K_generic ggml_vec_dot_q2_K_q8_K +#define ggml_vec_dot_iq2_xxs_q8_K_generic ggml_vec_dot_iq2_xxs_q8_K +#define ggml_vec_dot_iq2_xs_q8_K_generic ggml_vec_dot_iq2_xs_q8_K +#define ggml_vec_dot_iq2_s_q8_K_generic ggml_vec_dot_iq2_s_q8_K +#define ggml_vec_dot_iq3_xxs_q8_K_generic ggml_vec_dot_iq3_xxs_q8_K +#define ggml_vec_dot_iq3_s_q8_K_generic ggml_vec_dot_iq3_s_q8_K +#define ggml_vec_dot_iq1_s_q8_K_generic ggml_vec_dot_iq1_s_q8_K +#define ggml_vec_dot_iq1_m_q8_K_generic ggml_vec_dot_iq1_m_q8_K +// repack.cpp +#define ggml_quantize_mat_q8_0_4x8_generic ggml_quantize_mat_q8_0_4x8 +#define ggml_quantize_mat_q8_K_4x4_generic ggml_quantize_mat_q8_K_4x4 +#define ggml_quantize_mat_q8_K_4x8_generic ggml_quantize_mat_q8_K_4x8 +#define ggml_gemv_q1_0_4x4_q8_0_generic ggml_gemv_q1_0_4x4_q8_0 +#define ggml_gemv_q1_0_4x8_q8_0_generic ggml_gemv_q1_0_4x8_q8_0 +#define ggml_gemv_q4_0_4x8_q8_0_generic ggml_gemv_q4_0_4x8_q8_0 +#define ggml_gemv_q4_0_8x8_q8_0_generic ggml_gemv_q4_0_8x8_q8_0 +#define ggml_gemv_q2_K_8x8_q8_K_generic ggml_gemv_q2_K_8x8_q8_K +#define ggml_gemv_q4_K_8x4_q8_K_generic ggml_gemv_q4_K_8x4_q8_K +#define ggml_gemv_q4_K_8x8_q8_K_generic ggml_gemv_q4_K_8x8_q8_K +#define ggml_gemv_q5_K_8x4_q8_K_generic ggml_gemv_q5_K_8x4_q8_K +#define ggml_gemv_q5_K_8x8_q8_K_generic ggml_gemv_q5_K_8x8_q8_K +#define ggml_gemv_q6_K_8x4_q8_K_generic ggml_gemv_q6_K_8x4_q8_K +#define ggml_gemv_q6_K_8x8_q8_K_generic ggml_gemv_q6_K_8x8_q8_K +#define ggml_gemv_iq4_nl_4x4_q8_0_generic ggml_gemv_iq4_nl_4x4_q8_0 +#define ggml_gemv_iq4_nl_8x8_q8_0_generic ggml_gemv_iq4_nl_8x8_q8_0 +#define ggml_gemv_mxfp4_4x4_q8_0_generic ggml_gemv_mxfp4_4x4_q8_0 +#define ggml_gemv_mxfp4_8x8_q8_0_generic ggml_gemv_mxfp4_8x8_q8_0 +#define ggml_gemv_q8_0_4x4_q8_0_generic ggml_gemv_q8_0_4x4_q8_0 +#define ggml_gemv_q8_0_4x8_q8_0_generic ggml_gemv_q8_0_4x8_q8_0 +#define ggml_gemm_q1_0_4x4_q8_0_generic ggml_gemm_q1_0_4x4_q8_0 +#define ggml_gemm_q1_0_4x8_q8_0_generic ggml_gemm_q1_0_4x8_q8_0 +#define ggml_gemm_q4_0_4x8_q8_0_generic ggml_gemm_q4_0_4x8_q8_0 +#define ggml_gemm_q4_0_8x8_q8_0_generic ggml_gemm_q4_0_8x8_q8_0 +#define ggml_gemm_q2_K_8x8_q8_K_generic ggml_gemm_q2_K_8x8_q8_K +#define ggml_gemm_q4_K_8x4_q8_K_generic ggml_gemm_q4_K_8x4_q8_K +#define ggml_gemm_q4_K_8x8_q8_K_generic ggml_gemm_q4_K_8x8_q8_K +#define ggml_gemm_q5_K_8x4_q8_K_generic ggml_gemm_q5_K_8x4_q8_K +#define ggml_gemm_q5_K_8x8_q8_K_generic ggml_gemm_q5_K_8x8_q8_K +#define ggml_gemm_q6_K_8x4_q8_K_generic ggml_gemm_q6_K_8x4_q8_K +#define ggml_gemm_q6_K_8x8_q8_K_generic ggml_gemm_q6_K_8x8_q8_K +#define ggml_gemm_iq4_nl_4x4_q8_0_generic ggml_gemm_iq4_nl_4x4_q8_0 +#define ggml_gemm_iq4_nl_8x8_q8_0_generic ggml_gemm_iq4_nl_8x8_q8_0 +#define ggml_gemm_mxfp4_4x4_q8_0_generic ggml_gemm_mxfp4_4x4_q8_0 +#define ggml_gemm_mxfp4_8x8_q8_0_generic ggml_gemm_mxfp4_8x8_q8_0 +#define ggml_gemm_q8_0_4x4_q8_0_generic ggml_gemm_q8_0_4x4_q8_0 +#define ggml_gemm_q8_0_4x8_q8_0_generic ggml_gemm_q8_0_4x8_q8_0 +#elif defined(__wasm__) +// quants.c +#define ggml_vec_dot_tq1_0_q8_K_generic ggml_vec_dot_tq1_0_q8_K +#define ggml_vec_dot_tq2_0_q8_K_generic ggml_vec_dot_tq2_0_q8_K +#define ggml_vec_dot_iq2_xxs_q8_K_generic ggml_vec_dot_iq2_xxs_q8_K +#define ggml_vec_dot_iq2_xs_q8_K_generic ggml_vec_dot_iq2_xs_q8_K +#define ggml_vec_dot_iq2_s_q8_K_generic ggml_vec_dot_iq2_s_q8_K +#define ggml_vec_dot_iq3_xxs_q8_K_generic ggml_vec_dot_iq3_xxs_q8_K +#define ggml_vec_dot_iq3_s_q8_K_generic ggml_vec_dot_iq3_s_q8_K +#define ggml_vec_dot_iq1_s_q8_K_generic ggml_vec_dot_iq1_s_q8_K +#define ggml_vec_dot_iq1_m_q8_K_generic ggml_vec_dot_iq1_m_q8_K +#define ggml_vec_dot_iq4_nl_q8_0_generic ggml_vec_dot_iq4_nl_q8_0 +#define ggml_vec_dot_iq4_xs_q8_K_generic ggml_vec_dot_iq4_xs_q8_K +#define ggml_vec_dot_mxfp4_q8_0_generic ggml_vec_dot_mxfp4_q8_0 +#define ggml_vec_dot_nvfp4_q8_0_generic ggml_vec_dot_nvfp4_q8_0 +#define ggml_vec_dot_q1_0_q8_0_generic ggml_vec_dot_q1_0_q8_0 +#define ggml_vec_dot_q2_0_q8_0_generic ggml_vec_dot_q2_0_q8_0 +// repack.cpp +#define ggml_quantize_mat_q8_0_4x4_generic ggml_quantize_mat_q8_0_4x4 +#define ggml_quantize_mat_q8_0_4x8_generic ggml_quantize_mat_q8_0_4x8 +#define ggml_quantize_mat_q8_K_4x4_generic ggml_quantize_mat_q8_K_4x4 +#define ggml_quantize_mat_q8_K_4x8_generic ggml_quantize_mat_q8_K_4x8 +#define ggml_gemv_q1_0_4x4_q8_0_generic ggml_gemv_q1_0_4x4_q8_0 +#define ggml_gemv_q1_0_4x8_q8_0_generic ggml_gemv_q1_0_4x8_q8_0 +#define ggml_gemv_q4_0_4x4_q8_0_generic ggml_gemv_q4_0_4x4_q8_0 +#define ggml_gemv_q4_0_4x8_q8_0_generic ggml_gemv_q4_0_4x8_q8_0 +#define ggml_gemv_q4_0_8x8_q8_0_generic ggml_gemv_q4_0_8x8_q8_0 +#define ggml_gemv_q2_K_8x8_q8_K_generic ggml_gemv_q2_K_8x8_q8_K +#define ggml_gemv_q4_K_8x4_q8_K_generic ggml_gemv_q4_K_8x4_q8_K +#define ggml_gemv_q4_K_8x8_q8_K_generic ggml_gemv_q4_K_8x8_q8_K +#define ggml_gemv_q5_K_8x4_q8_K_generic ggml_gemv_q5_K_8x4_q8_K +#define ggml_gemv_q5_K_8x8_q8_K_generic ggml_gemv_q5_K_8x8_q8_K +#define ggml_gemv_q6_K_8x4_q8_K_generic ggml_gemv_q6_K_8x4_q8_K +#define ggml_gemv_q6_K_8x8_q8_K_generic ggml_gemv_q6_K_8x8_q8_K +#define ggml_gemv_iq4_nl_4x4_q8_0_generic ggml_gemv_iq4_nl_4x4_q8_0 +#define ggml_gemv_iq4_nl_8x8_q8_0_generic ggml_gemv_iq4_nl_8x8_q8_0 +#define ggml_gemv_mxfp4_4x4_q8_0_generic ggml_gemv_mxfp4_4x4_q8_0 +#define ggml_gemv_mxfp4_8x8_q8_0_generic ggml_gemv_mxfp4_8x8_q8_0 +#define ggml_gemv_q8_0_4x4_q8_0_generic ggml_gemv_q8_0_4x4_q8_0 +#define ggml_gemv_q8_0_4x8_q8_0_generic ggml_gemv_q8_0_4x8_q8_0 +#define ggml_gemm_q1_0_4x4_q8_0_generic ggml_gemm_q1_0_4x4_q8_0 +#define ggml_gemm_q1_0_4x8_q8_0_generic ggml_gemm_q1_0_4x8_q8_0 +#define ggml_gemm_q4_0_4x4_q8_0_generic ggml_gemm_q4_0_4x4_q8_0 +#define ggml_gemm_q4_0_4x8_q8_0_generic ggml_gemm_q4_0_4x8_q8_0 +#define ggml_gemm_q4_0_8x8_q8_0_generic ggml_gemm_q4_0_8x8_q8_0 +#define ggml_gemm_q2_K_8x8_q8_K_generic ggml_gemm_q2_K_8x8_q8_K +#define ggml_gemm_q4_K_8x4_q8_K_generic ggml_gemm_q4_K_8x4_q8_K +#define ggml_gemm_q4_K_8x8_q8_K_generic ggml_gemm_q4_K_8x8_q8_K +#define ggml_gemm_q5_K_8x4_q8_K_generic ggml_gemm_q5_K_8x4_q8_K +#define ggml_gemm_q5_K_8x8_q8_K_generic ggml_gemm_q5_K_8x8_q8_K +#define ggml_gemm_q6_K_8x4_q8_K_generic ggml_gemm_q6_K_8x4_q8_K +#define ggml_gemm_q6_K_8x8_q8_K_generic ggml_gemm_q6_K_8x8_q8_K +#define ggml_gemm_iq4_nl_4x4_q8_0_generic ggml_gemm_iq4_nl_4x4_q8_0 +#define ggml_gemm_iq4_nl_8x8_q8_0_generic ggml_gemm_iq4_nl_8x8_q8_0 +#define ggml_gemm_mxfp4_4x4_q8_0_generic ggml_gemm_mxfp4_4x4_q8_0 +#define ggml_gemm_mxfp4_8x8_q8_0_generic ggml_gemm_mxfp4_8x8_q8_0 +#define ggml_gemm_q8_0_4x4_q8_0_generic ggml_gemm_q8_0_4x4_q8_0 +#define ggml_gemm_q8_0_4x8_q8_0_generic ggml_gemm_q8_0_4x8_q8_0 +#endif diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cpu/arch/x86/cpu-feats.cpp b/python/freetoken/kernel/csrc/ggml/ggml-cpu/arch/x86/cpu-feats.cpp new file mode 100644 index 0000000000..d775a03638 --- /dev/null +++ b/python/freetoken/kernel/csrc/ggml/ggml-cpu/arch/x86/cpu-feats.cpp @@ -0,0 +1,327 @@ +#include "ggml-backend-impl.h" + +#if defined(__x86_64__) || (defined(_MSC_VER) && defined(_M_AMD64)) + +#ifdef _MSC_VER +#include +#endif + +#include +#include +#include +#include +#include + +// ref: https://cdrdv2-public.intel.com/782156/325383-sdm-vol-2abcd.pdf +struct cpuid_x86 { + bool SSE3(void) { return f_1_ecx[0]; } + bool PCLMULQDQ(void) { return f_1_ecx[1]; } + bool MONITOR(void) { return f_1_ecx[3]; } + bool SSSE3(void) { return f_1_ecx[9]; } + bool FMA(void) { return f_1_ecx[12]; } + bool CMPXCHG16B(void) { return f_1_ecx[13]; } + bool SSE41(void) { return f_1_ecx[19]; } + bool SSE42(void) { return f_1_ecx[20]; } + bool MOVBE(void) { return f_1_ecx[22]; } + bool POPCNT(void) { return f_1_ecx[23]; } + bool AES(void) { return f_1_ecx[25]; } + bool XSAVE(void) { return f_1_ecx[26]; } + bool OSXSAVE(void) { return f_1_ecx[27]; } + bool AVX(void) { return f_1_ecx[28]; } + bool F16C(void) { return f_1_ecx[29]; } + bool RDRAND(void) { return f_1_ecx[30]; } + + bool MSR(void) { return f_1_edx[5]; } + bool CX8(void) { return f_1_edx[8]; } + bool SEP(void) { return f_1_edx[11]; } + bool CMOV(void) { return f_1_edx[15]; } + bool CLFSH(void) { return f_1_edx[19]; } + bool MMX(void) { return f_1_edx[23]; } + bool FXSR(void) { return f_1_edx[24]; } + bool SSE(void) { return f_1_edx[25]; } + bool SSE2(void) { return f_1_edx[26]; } + + bool FSGSBASE(void) { return f_7_ebx[0]; } + bool BMI1(void) { return f_7_ebx[3]; } + bool HLE(void) { return is_intel && f_7_ebx[4]; } + bool AVX2(void) { return f_7_ebx[5]; } + bool BMI2(void) { return f_7_ebx[8]; } + bool ERMS(void) { return f_7_ebx[9]; } + bool INVPCID(void) { return f_7_ebx[10]; } + bool RTM(void) { return is_intel && f_7_ebx[11]; } + bool AVX512F(void) { return f_7_ebx[16]; } + bool AVX512DQ(void) { return f_7_ebx[17]; } + bool RDSEED(void) { return f_7_ebx[18]; } + bool ADX(void) { return f_7_ebx[19]; } + bool AVX512PF(void) { return f_7_ebx[26]; } + bool AVX512ER(void) { return f_7_ebx[27]; } + bool AVX512CD(void) { return f_7_ebx[28]; } + bool AVX512BW(void) { return f_7_ebx[30]; } + bool AVX512VL(void) { return f_7_ebx[31]; } + + bool SHA(void) { return f_7_ebx[29]; } + + bool PREFETCHWT1(void) { return f_7_ecx[0]; } + + bool LAHF(void) { return f_81_ecx[0]; } + bool LZCNT(void) { return is_intel && f_81_ecx[5]; } + bool ABM(void) { return is_amd && f_81_ecx[5]; } + bool SSE4a(void) { return is_amd && f_81_ecx[6]; } + bool XOP(void) { return is_amd && f_81_ecx[11]; } + bool TBM(void) { return is_amd && f_81_ecx[21]; } + + bool SYSCALL(void) { return is_intel && f_81_edx[11]; } + bool MMXEXT(void) { return is_amd && f_81_edx[22]; } + bool RDTSCP(void) { return is_intel && f_81_edx[27]; } + bool _3DNOWEXT(void) { return is_amd && f_81_edx[30]; } + bool _3DNOW(void) { return is_amd && f_81_edx[31]; } + + bool AVX512_VBMI(void) { return f_7_ecx[1]; } + bool AVX512_VNNI(void) { return f_7_ecx[11]; } + bool AVX512_FP16(void) { return f_7_edx[23]; } + bool AVX512_BF16(void) { return f_7_1_eax[5]; } + bool AVX_VNNI(void) { return f_7_1_eax[4]; } + + bool AMX_TILE(void) { return f_7_edx[24]; } + bool AMX_INT8(void) { return f_7_edx[25]; } + bool AMX_FP16(void) { return f_7_1_eax[21]; } + bool AMX_BF16(void) { return f_7_edx[22]; } + +#ifdef _MSC_VER + static void cpuid(int cpu_info[4], int eax) { + __cpuid(cpu_info, eax); + } + static void cpuidex(int cpu_info[4], int eax, int ecx) { + __cpuidex(cpu_info, eax, ecx); + } +#else + static void cpuid(int cpu_info[4], int eax) { + __asm__ __volatile__( + "cpuid" + : "=a"(cpu_info[0]), "=b"(cpu_info[1]), "=c"(cpu_info[2]), "=d"(cpu_info[3]) + : "a"(eax), "c"(0)); + } + static void cpuidex(int cpu_info[4], int eax, int ecx) { + __asm__ __volatile__( + "cpuid" + : "=a"(cpu_info[0]), "=b"(cpu_info[1]), "=c"(cpu_info[2]), "=d"(cpu_info[3]) + : "a"(eax), "c"(ecx)); + } +#endif + + cpuid_x86() { + std::array cpui; + std::vector> data; + + // calling __cpuid with 0x0 as the function_id argument + // gets the number of the highest valid function ID. + cpuid(cpui.data(), 0); + int n_ids = cpui[0]; + + for (int i = 0; i <= n_ids; ++i) { + cpuidex(cpui.data(), i, 0); + data.push_back(cpui); + } + + // capture vendor string + char vendor[0x20] = {}; + *reinterpret_cast(vendor) = data[0][1]; + *reinterpret_cast(vendor + 4) = data[0][3]; + *reinterpret_cast(vendor + 8) = data[0][2]; + this->vendor = vendor; + if (this->vendor == "GenuineIntel") { + is_intel = true; + } else if (this->vendor == "AuthenticAMD") { + is_amd = true; + } + + // load bitset with flags for function 0x00000001 + if (n_ids >= 1) { + f_1_ecx = data[1][2]; + f_1_edx = data[1][3]; + } + + // load bitset with flags for function 0x00000007 + if (n_ids >= 7) { + f_7_ebx = data[7][1]; + f_7_ecx = data[7][2]; + f_7_edx = data[7][3]; + cpuidex(cpui.data(), 7, 1); + f_7_1_eax = cpui[0]; + } + + // calling __cpuid with 0x80000000 as the function_id argument + // gets the number of the highest valid extended ID. + cpuid(cpui.data(), 0x80000000); + unsigned int n_ex_ids = cpui[0]; + + std::vector> ext_data; + for (unsigned int i = 0x80000000; i <= n_ex_ids; ++i) { + cpuidex(cpui.data(), i, 0); + ext_data.push_back(cpui); + } + + // load bitset with flags for function 0x80000001 + if (n_ex_ids >= 0x80000001) { + f_81_ecx = ext_data[1][2]; + f_81_edx = ext_data[1][3]; + } + + // interpret CPU brand string if reported + char brand[0x40] = {}; + if (n_ex_ids >= 0x80000004) { + std::memcpy(brand, ext_data[2].data(), sizeof(cpui)); + std::memcpy(brand + 16, ext_data[3].data(), sizeof(cpui)); + std::memcpy(brand + 32, ext_data[4].data(), sizeof(cpui)); + this->brand = brand; + } + } + + bool is_intel = false; + bool is_amd = false; + std::string vendor; + std::string brand; + std::bitset<32> f_1_ecx; + std::bitset<32> f_1_edx; + std::bitset<32> f_7_ebx; + std::bitset<32> f_7_ecx; + std::bitset<32> f_7_edx; + std::bitset<32> f_7_1_eax; + std::bitset<32> f_81_ecx; + std::bitset<32> f_81_edx; +}; + +#if 0 +void test_x86_is() { + cpuid_x86 is; + printf("CPU Vendor: %s\n", is.vendor.c_str()); + printf("Brand: %s\n", is.brand.c_str()); + printf("is_intel: %d\n", is.is_intel); + printf("is_amd: %d\n", is.is_amd); + printf("sse3: %d\n", is.SSE3()); + printf("pclmulqdq: %d\n", is.PCLMULQDQ()); + printf("ssse3: %d\n", is.SSSE3()); + printf("fma: %d\n", is.FMA()); + printf("cmpxchg16b: %d\n", is.CMPXCHG16B()); + printf("sse41: %d\n", is.SSE41()); + printf("sse42: %d\n", is.SSE42()); + printf("movbe: %d\n", is.MOVBE()); + printf("popcnt: %d\n", is.POPCNT()); + printf("aes: %d\n", is.AES()); + printf("xsave: %d\n", is.XSAVE()); + printf("osxsave: %d\n", is.OSXSAVE()); + printf("avx: %d\n", is.AVX()); + printf("f16c: %d\n", is.F16C()); + printf("rdrand: %d\n", is.RDRAND()); + printf("msr: %d\n", is.MSR()); + printf("cx8: %d\n", is.CX8()); + printf("sep: %d\n", is.SEP()); + printf("cmov: %d\n", is.CMOV()); + printf("clflush: %d\n", is.CLFSH()); + printf("mmx: %d\n", is.MMX()); + printf("fxsr: %d\n", is.FXSR()); + printf("sse: %d\n", is.SSE()); + printf("sse2: %d\n", is.SSE2()); + printf("fsgsbase: %d\n", is.FSGSBASE()); + printf("bmi1: %d\n", is.BMI1()); + printf("hle: %d\n", is.HLE()); + printf("avx2: %d\n", is.AVX2()); + printf("bmi2: %d\n", is.BMI2()); + printf("erms: %d\n", is.ERMS()); + printf("invpcid: %d\n", is.INVPCID()); + printf("rtm: %d\n", is.RTM()); + printf("avx512f: %d\n", is.AVX512F()); + printf("rdseed: %d\n", is.RDSEED()); + printf("adx: %d\n", is.ADX()); + printf("avx512pf: %d\n", is.AVX512PF()); + printf("avx512er: %d\n", is.AVX512ER()); + printf("avx512cd: %d\n", is.AVX512CD()); + printf("sha: %d\n", is.SHA()); + printf("prefetchwt1: %d\n", is.PREFETCHWT1()); + printf("lahf: %d\n", is.LAHF()); + printf("lzcnt: %d\n", is.LZCNT()); + printf("abm: %d\n", is.ABM()); + printf("sse4a: %d\n", is.SSE4a()); + printf("xop: %d\n", is.XOP()); + printf("tbm: %d\n", is.TBM()); + printf("syscall: %d\n", is.SYSCALL()); + printf("mmxext: %d\n", is.MMXEXT()); + printf("rdtscp: %d\n", is.RDTSCP()); + printf("3dnowext: %d\n", is._3DNOWEXT()); + printf("3dnow: %d\n", is._3DNOW()); + printf("avx512_vbmi: %d\n", is.AVX512_VBMI()); + printf("avx512_vnni: %d\n", is.AVX512_VNNI()); + printf("avx512_fp16: %d\n", is.AVX512_FP16()); + printf("avx512_bf16: %d\n", is.AVX512_BF16()); + printf("amx_tile: %d\n", is.AMX_TILE()); + printf("amx_int8: %d\n", is.AMX_INT8()); + printf("amx_fp16: %d\n", is.AMX_FP16()); + printf("amx_bf16: %d\n", is.AMX_BF16()); +} +#endif + +static int ggml_backend_cpu_x86_score() { + // FIXME: this does not check for OS support + + int score = 1; + cpuid_x86 is; + +#ifdef GGML_FMA + if (!is.FMA()) { return 0; } + score += 1; +#endif +#ifdef GGML_F16C + if (!is.F16C()) { return 0; } + score += 1<<1; +#endif +#ifdef GGML_SSE42 + if (!is.SSE42()) { return 0; } + score += 1<<2; +#endif +#ifdef GGML_BMI2 + if (!is.BMI2()) { return 0; } + score += 1<<3; +#endif +#ifdef GGML_AVX + if (!is.AVX()) { return 0; } + score += 1<<4; +#endif +#ifdef GGML_AVX2 + if (!is.AVX2()) { return 0; } + score += 1<<5; +#endif +#ifdef GGML_AVX_VNNI + if (!is.AVX_VNNI()) { return 0; } + score += 1<<6; +#endif +#ifdef GGML_AVX512 + if (!is.AVX512F()) { return 0; } + if (!is.AVX512CD()) { return 0; } + if (!is.AVX512VL()) { return 0; } + if (!is.AVX512DQ()) { return 0; } + if (!is.AVX512BW()) { return 0; } + score += 1<<7; +#endif +#ifdef GGML_AVX512_VBMI + if (!is.AVX512_VBMI()) { return 0; } + score += 1<<8; +#endif +#ifdef GGML_AVX512_BF16 + if (!is.AVX512_BF16()) { return 0; } + score += 1<<9; +#endif +#ifdef GGML_AVX512_VNNI + if (!is.AVX512_VNNI()) { return 0; } + score += 1<<10; +#endif +#ifdef GGML_AMX_INT8 + if (!is.AMX_INT8()) { return 0; } + score += 1<<11; +#endif + + return score; +} + +GGML_BACKEND_DL_SCORE_IMPL(ggml_backend_cpu_x86_score) + +#endif // defined(__x86_64__) || (defined(_MSC_VER) && defined(_M_AMD64)) diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cpu/arch/x86/quants.c b/python/freetoken/kernel/csrc/ggml/ggml-cpu/arch/x86/quants.c new file mode 100644 index 0000000000..21cd6312a1 --- /dev/null +++ b/python/freetoken/kernel/csrc/ggml/ggml-cpu/arch/x86/quants.c @@ -0,0 +1,4172 @@ +#define GGML_COMMON_IMPL_C +#include "ggml-common.h" +#include "ggml-quants.h" +#include "ggml-impl.h" +#include "ggml-cpu.h" +#include "simd-mappings.h" + +#include "../../quants.h" +#include "../../ggml-cpu-impl.h" + +#include +#include +#include +#include // for qsort +#include // for GGML_ASSERT + +#define GROUP_MAX_EPS 1e-15f +#define GROUP_MAX_EPS_IQ3_XXS 1e-8f +#define GROUP_MAX_EPS_IQ2_S 1e-8f +#define GROUP_MAX_EPS_IQ1_M 1e-7f +#define GROUP_MAX_EPS_IQ1_S 1e-12f + +#define UNUSED GGML_UNUSED + +// some compilers don't provide _mm256_set_m128i, e.g. gcc 7 +#define MM256_SET_M128I(a, b) _mm256_insertf128_si256(_mm256_castsi128_si256(b), (a), 1) + +#if defined(__AVX__) || defined(__AVX2__) || defined(__AVX512F__) || defined(__SSSE3__) +// multiply int8_t, add results pairwise twice +static inline __m128i mul_sum_i8_pairs(const __m128i x, const __m128i y) { + // Get absolute values of x vectors + const __m128i ax = _mm_sign_epi8(x, x); + // Sign the values of the y vectors + const __m128i sy = _mm_sign_epi8(y, x); + // Perform multiplication and create 16-bit values + const __m128i dot = _mm_maddubs_epi16(ax, sy); + const __m128i ones = _mm_set1_epi16(1); + return _mm_madd_epi16(ones, dot); +} + +#if __AVX__ || __AVX2__ || __AVX512F__ +// horizontally add 8 floats +static inline float hsum_float_8(const __m256 x) { + __m128 res = _mm256_extractf128_ps(x, 1); + res = _mm_add_ps(res, _mm256_castps256_ps128(x)); + res = _mm_add_ps(res, _mm_movehl_ps(res, res)); + res = _mm_add_ss(res, _mm_movehdup_ps(res)); + return _mm_cvtss_f32(res); +} + +// horizontally add 8 int32_t +static inline int hsum_i32_8(const __m256i a) { + const __m128i sum128 = _mm_add_epi32(_mm256_castsi256_si128(a), _mm256_extractf128_si256(a, 1)); + const __m128i hi64 = _mm_unpackhi_epi64(sum128, sum128); + const __m128i sum64 = _mm_add_epi32(hi64, sum128); + const __m128i hi32 = _mm_shuffle_epi32(sum64, _MM_SHUFFLE(2, 3, 0, 1)); + return _mm_cvtsi128_si32(_mm_add_epi32(sum64, hi32)); +} + +// horizontally add 4 int32_t +static inline int hsum_i32_4(const __m128i a) { + const __m128i hi64 = _mm_unpackhi_epi64(a, a); + const __m128i sum64 = _mm_add_epi32(hi64, a); + const __m128i hi32 = _mm_shuffle_epi32(sum64, _MM_SHUFFLE(2, 3, 0, 1)); + return _mm_cvtsi128_si32(_mm_add_epi32(sum64, hi32)); +} + +#if defined(__AVX2__) || defined(__AVX512F__) +static inline __m256i mul_add_epi8(const __m256i x, const __m256i y) { + const __m256i ax = _mm256_sign_epi8(x, x); + const __m256i sy = _mm256_sign_epi8(y, x); + return _mm256_maddubs_epi16(ax, sy); +} + +// spread 32 bits to 32 bytes { 0x00, 0xFF } +static inline __m256i bytes_from_bits_32(const uint8_t * x) { + uint32_t x32; + memcpy(&x32, x, sizeof(uint32_t)); + const __m256i shuf_mask = _mm256_set_epi64x( + 0x0303030303030303, 0x0202020202020202, + 0x0101010101010101, 0x0000000000000000); + __m256i bytes = _mm256_shuffle_epi8(_mm256_set1_epi32(x32), shuf_mask); + const __m256i bit_mask = _mm256_set1_epi64x(0x7fbfdfeff7fbfdfe); + bytes = _mm256_or_si256(bytes, bit_mask); + return _mm256_cmpeq_epi8(bytes, _mm256_set1_epi64x(-1)); +} + +// Unpack 32 4-bit fields into 32 bytes +// The output vector contains 32 bytes, each one in [ 0 .. 15 ] interval +static inline __m256i bytes_from_nibbles_32(const uint8_t * rsi) +{ + const __m128i tmp = _mm_loadu_si128((const __m128i *)rsi); + const __m256i bytes = MM256_SET_M128I(_mm_srli_epi16(tmp, 4), tmp); + const __m256i lowMask = _mm256_set1_epi8( 0xF ); + return _mm256_and_si256(lowMask, bytes); +} + +// add int16_t pairwise and return as float vector +static inline __m256 sum_i16_pairs_float(const __m256i x) { + const __m256i ones = _mm256_set1_epi16(1); + const __m256i summed_pairs = _mm256_madd_epi16(ones, x); + return _mm256_cvtepi32_ps(summed_pairs); +} + +static inline __m256 mul_sum_us8_pairs_float(const __m256i ax, const __m256i sy) { +#if defined(__AVX512VNNI__) && defined(__AVX512VL__) + const __m256i zero = _mm256_setzero_si256(); + const __m256i summed_pairs = _mm256_dpbusd_epi32(zero, ax, sy); + return _mm256_cvtepi32_ps(summed_pairs); +#elif defined(__AVXVNNI__) + const __m256i zero = _mm256_setzero_si256(); + const __m256i summed_pairs = _mm256_dpbusd_avx_epi32(zero, ax, sy); + return _mm256_cvtepi32_ps(summed_pairs); +#else + // Perform multiplication and create 16-bit values + const __m256i dot = _mm256_maddubs_epi16(ax, sy); + return sum_i16_pairs_float(dot); +#endif +} + +// multiply int8_t, add results pairwise twice and return as float vector +static inline __m256 mul_sum_i8_pairs_float(const __m256i x, const __m256i y) { +#if __AVXVNNIINT8__ + const __m256i zero = _mm256_setzero_si256(); + const __m256i summed_pairs = _mm256_dpbssd_epi32(zero, x, y); + return _mm256_cvtepi32_ps(summed_pairs); +#else + // Get absolute values of x vectors + const __m256i ax = _mm256_sign_epi8(x, x); + // Sign the values of the y vectors + const __m256i sy = _mm256_sign_epi8(y, x); + return mul_sum_us8_pairs_float(ax, sy); +#endif +} + +static inline __m128i packNibbles( __m256i bytes ) +{ + // Move bits within 16-bit lanes from 0000_abcd_0000_efgh into 0000_0000_abcd_efgh +#if __AVX512F__ + const __m256i bytes_srli_4 = _mm256_srli_epi16(bytes, 4); // 0000_0000_abcd_0000 + bytes = _mm256_or_si256(bytes, bytes_srli_4); // 0000_abcd_abcd_efgh + return _mm256_cvtepi16_epi8(bytes); // abcd_efgh +#else + const __m256i lowByte = _mm256_set1_epi16( 0xFF ); + __m256i high = _mm256_andnot_si256( lowByte, bytes ); + __m256i low = _mm256_and_si256( lowByte, bytes ); + high = _mm256_srli_epi16( high, 4 ); + bytes = _mm256_or_si256( low, high ); + + // Compress uint16_t lanes into bytes + __m128i r0 = _mm256_castsi256_si128( bytes ); + __m128i r1 = _mm256_extracti128_si256( bytes, 1 ); + return _mm_packus_epi16( r0, r1 ); +#endif +} +#elif defined(__AVX__) +static inline __m128i packNibbles( __m128i bytes1, __m128i bytes2 ) +{ + // Move bits within 16-bit lanes from 0000_abcd_0000_efgh into 0000_0000_abcd_efgh + const __m128i lowByte = _mm_set1_epi16( 0xFF ); + __m128i high = _mm_andnot_si128( lowByte, bytes1 ); + __m128i low = _mm_and_si128( lowByte, bytes1 ); + high = _mm_srli_epi16( high, 4 ); + bytes1 = _mm_or_si128( low, high ); + high = _mm_andnot_si128( lowByte, bytes2 ); + low = _mm_and_si128( lowByte, bytes2 ); + high = _mm_srli_epi16( high, 4 ); + bytes2 = _mm_or_si128( low, high ); + + return _mm_packus_epi16( bytes1, bytes2); +} + +static inline __m128i mul_add_epi8_sse(const __m128i x, const __m128i y) { + const __m128i ax = _mm_sign_epi8(x, x); + const __m128i sy = _mm_sign_epi8(y, x); + return _mm_maddubs_epi16(ax, sy); +} + +// spread 32 bits to 32 bytes { 0x00, 0xFF } +static inline __m256i bytes_from_bits_32(const uint8_t * x) { + uint32_t x32; + memcpy(&x32, x, sizeof(uint32_t)); + const __m128i shuf_maskl = _mm_set_epi64x(0x0101010101010101, 0x0000000000000000); + const __m128i shuf_maskh = _mm_set_epi64x(0x0303030303030303, 0x0202020202020202); + __m128i bytesl = _mm_shuffle_epi8(_mm_set1_epi32(x32), shuf_maskl); + __m128i bytesh = _mm_shuffle_epi8(_mm_set1_epi32(x32), shuf_maskh); + const __m128i bit_mask = _mm_set1_epi64x(0x7fbfdfeff7fbfdfe); + bytesl = _mm_or_si128(bytesl, bit_mask); + bytesh = _mm_or_si128(bytesh, bit_mask); + bytesl = _mm_cmpeq_epi8(bytesl, _mm_set1_epi64x(-1)); + bytesh = _mm_cmpeq_epi8(bytesh, _mm_set1_epi64x(-1)); + return MM256_SET_M128I(bytesh, bytesl); +} + +// Unpack 32 4-bit fields into 32 bytes +// The output vector contains 32 bytes, each one in [ 0 .. 15 ] interval +static inline __m256i bytes_from_nibbles_32(const uint8_t * rsi) +{ + // Load 16 bytes from memory + __m128i tmpl = _mm_loadu_si128((const __m128i *)rsi); + __m128i tmph = _mm_srli_epi16(tmpl, 4); + const __m128i lowMask = _mm_set1_epi8(0xF); + tmpl = _mm_and_si128(lowMask, tmpl); + tmph = _mm_and_si128(lowMask, tmph); + return MM256_SET_M128I(tmph, tmpl); +} + +// add int16_t pairwise and return as float vector +static inline __m256 sum_i16_pairs_float(const __m128i xh, const __m128i xl) { + const __m128i ones = _mm_set1_epi16(1); + const __m128i summed_pairsl = _mm_madd_epi16(ones, xl); + const __m128i summed_pairsh = _mm_madd_epi16(ones, xh); + const __m256i summed_pairs = MM256_SET_M128I(summed_pairsh, summed_pairsl); + return _mm256_cvtepi32_ps(summed_pairs); +} + +static inline __m256 mul_sum_us8_pairs_float(const __m256i ax, const __m256i sy) { + const __m128i axl = _mm256_castsi256_si128(ax); + const __m128i axh = _mm256_extractf128_si256(ax, 1); + const __m128i syl = _mm256_castsi256_si128(sy); + const __m128i syh = _mm256_extractf128_si256(sy, 1); + // Perform multiplication and create 16-bit values + const __m128i dotl = _mm_maddubs_epi16(axl, syl); + const __m128i doth = _mm_maddubs_epi16(axh, syh); + return sum_i16_pairs_float(doth, dotl); +} + +// multiply int8_t, add results pairwise twice and return as float vector +static inline __m256 mul_sum_i8_pairs_float(const __m256i x, const __m256i y) { + const __m128i xl = _mm256_castsi256_si128(x); + const __m128i xh = _mm256_extractf128_si256(x, 1); + const __m128i yl = _mm256_castsi256_si128(y); + const __m128i yh = _mm256_extractf128_si256(y, 1); + // Get absolute values of x vectors + const __m128i axl = _mm_sign_epi8(xl, xl); + const __m128i axh = _mm_sign_epi8(xh, xh); + // Sign the values of the y vectors + const __m128i syl = _mm_sign_epi8(yl, xl); + const __m128i syh = _mm_sign_epi8(yh, xh); + // Perform multiplication and create 16-bit values + const __m128i dotl = _mm_maddubs_epi16(axl, syl); + const __m128i doth = _mm_maddubs_epi16(axh, syh); + return sum_i16_pairs_float(doth, dotl); +} + +// larger version of mul_sum_i8_pairs_float where x and y are each represented by four 128-bit vectors +static inline __m256 mul_sum_i8_quad_float(const __m128i x_1_0, const __m128i x_1_1, const __m128i x_2_0, const __m128i x_2_1, + const __m128i y_1_0, const __m128i y_1_1, const __m128i y_2_0, const __m128i y_2_1) { + const __m128i mone = _mm_set1_epi16(1); + + const __m128i p16_1_0 = mul_add_epi8_sse(x_1_0, y_1_0); + const __m128i p16_1_1 = mul_add_epi8_sse(x_1_1, y_1_1); + const __m128i p16_2_0 = mul_add_epi8_sse(x_2_0, y_2_0); + const __m128i p16_2_1 = mul_add_epi8_sse(x_2_1, y_2_1); + const __m128i p_1_0 = _mm_madd_epi16(p16_1_0, mone); + const __m128i p_1_1 = _mm_madd_epi16(p16_1_1, mone); + const __m128i p_2_0 = _mm_madd_epi16(p16_2_0, mone); + const __m128i p_2_1 = _mm_madd_epi16(p16_2_1, mone); + const __m128i p_1 = _mm_add_epi32(p_1_0, p_1_1); + const __m128i p_2 = _mm_add_epi32(p_2_0, p_2_1); + return _mm256_cvtepi32_ps(MM256_SET_M128I(p_2, p_1)); +} + +// quad fp16 delta calculation +static inline __m256 quad_fp16_delta_float(const float x0, const float y0, const float x1, const float y1) { + // GGML_CPU_FP16_TO_FP32 is faster than Intel F16C + return _mm256_set_m128(_mm_set1_ps(GGML_CPU_FP16_TO_FP32(x1) * GGML_CPU_FP16_TO_FP32(y1)), + _mm_set1_ps(GGML_CPU_FP16_TO_FP32(x0) * GGML_CPU_FP16_TO_FP32(y0))); +} + +static inline __m256 quad_mx_delta_float(const uint8_t x0, const float y0, const uint8_t x1, const float y1) { + return _mm256_set_m128(_mm_set1_ps(GGML_CPU_E8M0_TO_FP32_HALF(x1) * GGML_CPU_FP16_TO_FP32(y1)), + _mm_set1_ps(GGML_CPU_E8M0_TO_FP32_HALF(x0) * GGML_CPU_FP16_TO_FP32(y0))); +} +#endif +#elif defined(__SSSE3__) +static inline __m128i bytes_from_bits_16(const uint8_t * x) { + uint16_t x16; + memcpy(&x16, x, sizeof(uint16_t)); + + const __m128i shuf_mask = _mm_set_epi64x(0x0101010101010101, 0x0000000000000000); + __m128i bytes = _mm_shuffle_epi8(_mm_set1_epi16((short) x16), shuf_mask); + const __m128i bit_mask = _mm_set_epi64x(0x7fbfdfeff7fbfdfe, 0x7fbfdfeff7fbfdfe); + bytes = _mm_or_si128(bytes, bit_mask); + + return _mm_cmpeq_epi8(bytes, _mm_set1_epi64x(-1)); +} + +// horizontally add 4x4 floats +static inline float hsum_float_4x4(const __m128 a, const __m128 b, const __m128 c, const __m128 d) { + __m128 res_0 =_mm_hadd_ps(a, b); + __m128 res_1 =_mm_hadd_ps(c, d); + __m128 res =_mm_hadd_ps(res_0, res_1); + res =_mm_hadd_ps(res, res); + res =_mm_hadd_ps(res, res); + + return _mm_cvtss_f32(res); +} +#endif // __AVX__ || __AVX2__ || __AVX512F__ +#endif // defined(__AVX__) || defined(__AVX2__) || defined(__AVX512F__) || defined(__SSSE3__) + +void quantize_row_q8_0(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k) { + assert(QK8_0 == 32); + assert(k % QK8_0 == 0); + const int nb = k / QK8_0; + + block_q8_0 * GGML_RESTRICT y = vy; + +#if defined(__AVX2__) || defined(__AVX__) + for (int i = 0; i < nb; i++) { + // Load elements into 4 AVX vectors + __m256 v0 = _mm256_loadu_ps( x ); + __m256 v1 = _mm256_loadu_ps( x + 8 ); + __m256 v2 = _mm256_loadu_ps( x + 16 ); + __m256 v3 = _mm256_loadu_ps( x + 24 ); + x += 32; + + // Compute max(abs(e)) for the block + const __m256 signBit = _mm256_set1_ps( -0.0f ); + __m256 maxAbs = _mm256_andnot_ps( signBit, v0 ); + maxAbs = _mm256_max_ps( maxAbs, _mm256_andnot_ps( signBit, v1 ) ); + maxAbs = _mm256_max_ps( maxAbs, _mm256_andnot_ps( signBit, v2 ) ); + maxAbs = _mm256_max_ps( maxAbs, _mm256_andnot_ps( signBit, v3 ) ); + + __m128 max4 = _mm_max_ps( _mm256_extractf128_ps( maxAbs, 1 ), _mm256_castps256_ps128( maxAbs ) ); + max4 = _mm_max_ps( max4, _mm_movehl_ps( max4, max4 ) ); + max4 = _mm_max_ss( max4, _mm_movehdup_ps( max4 ) ); + const float maxScalar = _mm_cvtss_f32( max4 ); + + // Quantize these floats + const float d = maxScalar / 127.f; + y[i].d = GGML_CPU_FP32_TO_FP16(d); + const float id = ( maxScalar != 0.0f ) ? 127.f / maxScalar : 0.0f; + const __m256 mul = _mm256_set1_ps( id ); + + // Apply the multiplier + v0 = _mm256_mul_ps( v0, mul ); + v1 = _mm256_mul_ps( v1, mul ); + v2 = _mm256_mul_ps( v2, mul ); + v3 = _mm256_mul_ps( v3, mul ); + + // Round to nearest integer + v0 = _mm256_round_ps( v0, _MM_ROUND_NEAREST ); + v1 = _mm256_round_ps( v1, _MM_ROUND_NEAREST ); + v2 = _mm256_round_ps( v2, _MM_ROUND_NEAREST ); + v3 = _mm256_round_ps( v3, _MM_ROUND_NEAREST ); + + // Convert floats to integers + __m256i i0 = _mm256_cvtps_epi32( v0 ); + __m256i i1 = _mm256_cvtps_epi32( v1 ); + __m256i i2 = _mm256_cvtps_epi32( v2 ); + __m256i i3 = _mm256_cvtps_epi32( v3 ); + +#if defined(__AVX2__) + // Convert int32 to int16 + i0 = _mm256_packs_epi32( i0, i1 ); // 0, 1, 2, 3, 8, 9, 10, 11, 4, 5, 6, 7, 12, 13, 14, 15 + i2 = _mm256_packs_epi32( i2, i3 ); // 16, 17, 18, 19, 24, 25, 26, 27, 20, 21, 22, 23, 28, 29, 30, 31 + // Convert int16 to int8 + i0 = _mm256_packs_epi16( i0, i2 ); // 0, 1, 2, 3, 8, 9, 10, 11, 16, 17, 18, 19, 24, 25, 26, 27, 4, 5, 6, 7, 12, 13, 14, 15, 20, 21, 22, 23, 28, 29, 30, 31 + + // We got our precious signed bytes, but the order is now wrong + // These AVX2 pack instructions process 16-byte pieces independently + // The following instruction is fixing the order + const __m256i perm = _mm256_setr_epi32( 0, 4, 1, 5, 2, 6, 3, 7 ); + i0 = _mm256_permutevar8x32_epi32( i0, perm ); + + _mm256_storeu_si256((__m256i *)y[i].qs, i0); +#else + // Since we don't have in AVX some necessary functions, + // we split the registers in half and call AVX2 analogs from SSE + __m128i ni0 = _mm256_castsi256_si128( i0 ); + __m128i ni1 = _mm256_extractf128_si256( i0, 1); + __m128i ni2 = _mm256_castsi256_si128( i1 ); + __m128i ni3 = _mm256_extractf128_si256( i1, 1); + __m128i ni4 = _mm256_castsi256_si128( i2 ); + __m128i ni5 = _mm256_extractf128_si256( i2, 1); + __m128i ni6 = _mm256_castsi256_si128( i3 ); + __m128i ni7 = _mm256_extractf128_si256( i3, 1); + + // Convert int32 to int16 + ni0 = _mm_packs_epi32( ni0, ni1 ); + ni2 = _mm_packs_epi32( ni2, ni3 ); + ni4 = _mm_packs_epi32( ni4, ni5 ); + ni6 = _mm_packs_epi32( ni6, ni7 ); + // Convert int16 to int8 + ni0 = _mm_packs_epi16( ni0, ni2 ); + ni4 = _mm_packs_epi16( ni4, ni6 ); + + _mm_storeu_si128((__m128i *)(y[i].qs + 0), ni0); + _mm_storeu_si128((__m128i *)(y[i].qs + 16), ni4); +#endif + } +#else + GGML_UNUSED(nb); + // scalar + quantize_row_q8_0_ref(x, y, k); +#endif +} + +void quantize_row_q8_1(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k) { + assert(k % QK8_1 == 0); + const int nb = k / QK8_1; + + block_q8_1 * GGML_RESTRICT y = vy; +#if defined(__AVX2__) || defined(__AVX__) + for (int i = 0; i < nb; i++) { + // Load elements into 4 AVX vectors + __m256 v0 = _mm256_loadu_ps( x ); + __m256 v1 = _mm256_loadu_ps( x + 8 ); + __m256 v2 = _mm256_loadu_ps( x + 16 ); + __m256 v3 = _mm256_loadu_ps( x + 24 ); + x += 32; + + // Compute max(abs(e)) for the block + const __m256 signBit = _mm256_set1_ps( -0.0f ); + __m256 maxAbs = _mm256_andnot_ps( signBit, v0 ); + maxAbs = _mm256_max_ps( maxAbs, _mm256_andnot_ps( signBit, v1 ) ); + maxAbs = _mm256_max_ps( maxAbs, _mm256_andnot_ps( signBit, v2 ) ); + maxAbs = _mm256_max_ps( maxAbs, _mm256_andnot_ps( signBit, v3 ) ); + + __m128 max4 = _mm_max_ps( _mm256_extractf128_ps( maxAbs, 1 ), _mm256_castps256_ps128( maxAbs ) ); + max4 = _mm_max_ps( max4, _mm_movehl_ps( max4, max4 ) ); + max4 = _mm_max_ss( max4, _mm_movehdup_ps( max4 ) ); + const float max_scalar = _mm_cvtss_f32( max4 ); + + // Quantize these floats + const float d = max_scalar / 127.f; + y[i].d = GGML_CPU_FP32_TO_FP16(d); + const float id = ( max_scalar != 0.0f ) ? 127.f / max_scalar : 0.0f; + const __m256 mul = _mm256_set1_ps( id ); + + // Apply the multiplier + v0 = _mm256_mul_ps( v0, mul ); + v1 = _mm256_mul_ps( v1, mul ); + v2 = _mm256_mul_ps( v2, mul ); + v3 = _mm256_mul_ps( v3, mul ); + + // Round to nearest integer + v0 = _mm256_round_ps( v0, _MM_ROUND_NEAREST ); + v1 = _mm256_round_ps( v1, _MM_ROUND_NEAREST ); + v2 = _mm256_round_ps( v2, _MM_ROUND_NEAREST ); + v3 = _mm256_round_ps( v3, _MM_ROUND_NEAREST ); + + // Convert floats to integers + __m256i i0 = _mm256_cvtps_epi32( v0 ); + __m256i i1 = _mm256_cvtps_epi32( v1 ); + __m256i i2 = _mm256_cvtps_epi32( v2 ); + __m256i i3 = _mm256_cvtps_epi32( v3 ); + +#if defined(__AVX2__) + // Compute the sum of the quants and set y[i].s + y[i].s = GGML_CPU_FP32_TO_FP16(d * hsum_i32_8(_mm256_add_epi32(_mm256_add_epi32(i0, i1), _mm256_add_epi32(i2, i3)))); + + // Convert int32 to int16 + i0 = _mm256_packs_epi32( i0, i1 ); // 0, 1, 2, 3, 8, 9, 10, 11, 4, 5, 6, 7, 12, 13, 14, 15 + i2 = _mm256_packs_epi32( i2, i3 ); // 16, 17, 18, 19, 24, 25, 26, 27, 20, 21, 22, 23, 28, 29, 30, 31 + // Convert int16 to int8 + i0 = _mm256_packs_epi16( i0, i2 ); // 0, 1, 2, 3, 8, 9, 10, 11, 16, 17, 18, 19, 24, 25, 26, 27, 4, 5, 6, 7, 12, 13, 14, 15, 20, 21, 22, 23, 28, 29, 30, 31 + + // We got our precious signed bytes, but the order is now wrong + // These AVX2 pack instructions process 16-byte pieces independently + // The following instruction is fixing the order + const __m256i perm = _mm256_setr_epi32( 0, 4, 1, 5, 2, 6, 3, 7 ); + i0 = _mm256_permutevar8x32_epi32( i0, perm ); + + _mm256_storeu_si256((__m256i *)y[i].qs, i0); +#else + // Since we don't have in AVX some necessary functions, + // we split the registers in half and call AVX2 analogs from SSE + __m128i ni0 = _mm256_castsi256_si128( i0 ); + __m128i ni1 = _mm256_extractf128_si256( i0, 1); + __m128i ni2 = _mm256_castsi256_si128( i1 ); + __m128i ni3 = _mm256_extractf128_si256( i1, 1); + __m128i ni4 = _mm256_castsi256_si128( i2 ); + __m128i ni5 = _mm256_extractf128_si256( i2, 1); + __m128i ni6 = _mm256_castsi256_si128( i3 ); + __m128i ni7 = _mm256_extractf128_si256( i3, 1); + + // Compute the sum of the quants and set y[i].s + const __m128i s0 = _mm_add_epi32(_mm_add_epi32(ni0, ni1), _mm_add_epi32(ni2, ni3)); + const __m128i s1 = _mm_add_epi32(_mm_add_epi32(ni4, ni5), _mm_add_epi32(ni6, ni7)); + y[i].s = GGML_CPU_FP32_TO_FP16(d * hsum_i32_4(_mm_add_epi32(s0, s1))); + + // Convert int32 to int16 + ni0 = _mm_packs_epi32( ni0, ni1 ); + ni2 = _mm_packs_epi32( ni2, ni3 ); + ni4 = _mm_packs_epi32( ni4, ni5 ); + ni6 = _mm_packs_epi32( ni6, ni7 ); + // Convert int16 to int8 + ni0 = _mm_packs_epi16( ni0, ni2 ); + ni4 = _mm_packs_epi16( ni4, ni6 ); + + _mm_storeu_si128((__m128i *)(y[i].qs + 0), ni0); + _mm_storeu_si128((__m128i *)(y[i].qs + 16), ni4); +#endif + } +#else + GGML_UNUSED(nb); + // scalar + quantize_row_q8_1_ref(x, y, k); +#endif +} + +// placeholder implementation for Apple targets +void quantize_row_q8_K(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k) { + quantize_row_q8_K_ref(x, y, k); +} + +//===================================== Dot products ================================= + +// +// Helper functions +// + +#if __AVX__ || __AVX2__ || __AVX512F__ + +// shuffles to pick the required scales in dot products +static inline __m256i get_scale_shuffle_q3k(int i) { + static const uint8_t k_shuffle[128] = { + 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 2, 3, 2, 3, 2, 3, 2, 3, 2, 3, 2, 3, 2, 3, 2, 3, + 4, 5, 4, 5, 4, 5, 4, 5, 4, 5, 4, 5, 4, 5, 4, 5, 6, 7, 6, 7, 6, 7, 6, 7, 6, 7, 6, 7, 6, 7, 6, 7, + 8, 9, 8, 9, 8, 9, 8, 9, 8, 9, 8, 9, 8, 9, 8, 9, 10,11,10,11,10,11,10,11,10,11,10,11,10,11,10,11, + 12,13,12,13,12,13,12,13,12,13,12,13,12,13,12,13, 14,15,14,15,14,15,14,15,14,15,14,15,14,15,14,15, + }; + return _mm256_loadu_si256((const __m256i*)k_shuffle + i); +} +static inline __m256i get_scale_shuffle_k4(int i) { + static const uint8_t k_shuffle[256] = { + 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, + 2, 3, 2, 3, 2, 3, 2, 3, 2, 3, 2, 3, 2, 3, 2, 3, 2, 3, 2, 3, 2, 3, 2, 3, 2, 3, 2, 3, 2, 3, 2, 3, + 4, 5, 4, 5, 4, 5, 4, 5, 4, 5, 4, 5, 4, 5, 4, 5, 4, 5, 4, 5, 4, 5, 4, 5, 4, 5, 4, 5, 4, 5, 4, 5, + 6, 7, 6, 7, 6, 7, 6, 7, 6, 7, 6, 7, 6, 7, 6, 7, 6, 7, 6, 7, 6, 7, 6, 7, 6, 7, 6, 7, 6, 7, 6, 7, + 8, 9, 8, 9, 8, 9, 8, 9, 8, 9, 8, 9, 8, 9, 8, 9, 8, 9, 8, 9, 8, 9, 8, 9, 8, 9, 8, 9, 8, 9, 8, 9, + 10,11,10,11,10,11,10,11,10,11,10,11,10,11,10,11,10,11,10,11,10,11,10,11,10,11,10,11,10,11,10,11, + 12,13,12,13,12,13,12,13,12,13,12,13,12,13,12,13,12,13,12,13,12,13,12,13,12,13,12,13,12,13,12,13, + 14,15,14,15,14,15,14,15,14,15,14,15,14,15,14,15,14,15,14,15,14,15,14,15,14,15,14,15,14,15,14,15 + }; + return _mm256_loadu_si256((const __m256i*)k_shuffle + i); +} +static inline __m128i get_scale_shuffle(int i) { + static const uint8_t k_shuffle[128] = { + 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, + 2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3, + 4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5, + 6, 6, 6, 6, 6, 6, 6, 6, 7, 7, 7, 7, 7, 7, 7, 7, + 8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9, + 10,10,10,10,10,10,10,10, 11,11,11,11,11,11,11,11, + 12,12,12,12,12,12,12,12, 13,13,13,13,13,13,13,13, + 14,14,14,14,14,14,14,14, 15,15,15,15,15,15,15,15 + }; + return _mm_loadu_si128((const __m128i*)k_shuffle + i); +} +#endif + +void ggml_vec_dot_q1_0_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) { + const int qk = QK1_0; + const int nb = n / qk; + + assert(n % qk == 0); + assert(nrc == 1); + UNUSED(nrc); + UNUSED(bx); + UNUSED(by); + UNUSED(bs); + + const block_q1_0 * GGML_RESTRICT x = vx; + const block_q8_0 * GGML_RESTRICT y = vy; + +#if defined(__AVX2__) + const __m256i ones_8 = _mm256_set1_epi8(1); + const __m256i ones_16 = _mm256_set1_epi16(1); + const __m256i byte_shuf = _mm256_setr_epi8( + 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, + 2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3); + const __m256i bit_masks = _mm256_setr_epi8( + 1, 2, 4, 8, 16, 32, 64, (char) -128, 1, 2, 4, 8, 16, 32, 64, (char) -128, + 1, 2, 4, 8, 16, 32, 64, (char) -128, 1, 2, 4, 8, 16, 32, 64, (char) -128); + const __m256i zero = _mm256_setzero_si256(); + __m256 acc = _mm256_setzero_ps(); + + for (int ib = 0; ib < nb; ++ib) { + const float d0 = GGML_CPU_FP16_TO_FP32(x[ib].d); + const uint32_t * GGML_RESTRICT qs32 = (const uint32_t *) x[ib].qs; + const block_q8_0 * GGML_RESTRICT y_ptr = &y[ib * 4]; + + __m256 acc_block; + { + const __m256i qy = _mm256_loadu_si256((const __m256i *) y_ptr[0].qs); + const __m256i sm = _mm256_cmpeq_epi8( + _mm256_and_si256(_mm256_shuffle_epi8(_mm256_set1_epi32((int) qs32[0]), byte_shuf), bit_masks), zero); + const __m256i sy = _mm256_sub_epi8(_mm256_xor_si256(qy, sm), sm); + const __m256i s32 = _mm256_madd_epi16(_mm256_maddubs_epi16(ones_8, sy), ones_16); + acc_block = _mm256_mul_ps(_mm256_set1_ps(GGML_CPU_FP16_TO_FP32(y_ptr[0].d)), _mm256_cvtepi32_ps(s32)); + } + for (int K = 1; K < 4; ++K) { + const __m256i qy = _mm256_loadu_si256((const __m256i *) y_ptr[K].qs); + const __m256i sm = _mm256_cmpeq_epi8( + _mm256_and_si256(_mm256_shuffle_epi8(_mm256_set1_epi32((int) qs32[K]), byte_shuf), bit_masks), zero); + const __m256i sy = _mm256_sub_epi8(_mm256_xor_si256(qy, sm), sm); + const __m256i s32 = _mm256_madd_epi16(_mm256_maddubs_epi16(ones_8, sy), ones_16); + acc_block = _mm256_fmadd_ps(_mm256_set1_ps(GGML_CPU_FP16_TO_FP32(y_ptr[K].d)), _mm256_cvtepi32_ps(s32), acc_block); + } + acc = _mm256_fmadd_ps(_mm256_set1_ps(d0), acc_block, acc); + } + + *s = hsum_float_8(acc); +#elif defined(__AVX__) + const __m128i ones_8 = _mm_set1_epi8(1); + const __m128i ones_16 = _mm_set1_epi16(1); + const __m128i zero = _mm_setzero_si128(); + __m256 acc = _mm256_setzero_ps(); + + for (int ib = 0; ib < nb; ++ib) { + const float d0 = GGML_CPU_FP16_TO_FP32(x[ib].d); + const block_q8_0 * GGML_RESTRICT y_ptr = &y[ib * 4]; + __m256 acc_block; + { + const __m256i bit_mask = bytes_from_bits_32(&x[ib].qs[0]); + const __m128i bit_mask_0 = _mm256_castsi256_si128(bit_mask); + const __m128i bit_mask_1 = _mm256_extractf128_si256(bit_mask, 1); + const __m128i qy_0 = _mm_loadu_si128((const __m128i *) &y_ptr[0].qs[0]); + const __m128i qy_1 = _mm_loadu_si128((const __m128i *) &y_ptr[0].qs[16]); + const __m128i sign_mask_0 = _mm_cmpeq_epi8(bit_mask_0, zero); + const __m128i sign_mask_1 = _mm_cmpeq_epi8(bit_mask_1, zero); + const __m128i sy_0 = _mm_sub_epi8(_mm_xor_si128(qy_0, sign_mask_0), sign_mask_0); + const __m128i sy_1 = _mm_sub_epi8(_mm_xor_si128(qy_1, sign_mask_1), sign_mask_1); + const __m128i sum16_0 = _mm_maddubs_epi16(ones_8, sy_0); + const __m128i sum16_1 = _mm_maddubs_epi16(ones_8, sy_1); + const __m128i sum32_0 = _mm_madd_epi16(sum16_0, ones_16); + const __m128i sum32_1 = _mm_madd_epi16(sum16_1, ones_16); + const __m256 q = _mm256_cvtepi32_ps(MM256_SET_M128I(sum32_1, sum32_0)); + acc_block = _mm256_mul_ps(_mm256_set1_ps(GGML_CPU_FP16_TO_FP32(y_ptr[0].d)), q); + } + for(int K = 1; K < 4; ++K) { + const __m256i bit_mask = bytes_from_bits_32(&x[ib].qs[(K) * 4]); + const __m128i bit_mask_0 = _mm256_castsi256_si128(bit_mask); + const __m128i bit_mask_1 = _mm256_extractf128_si256(bit_mask, 1); + const __m128i qy_0 = _mm_loadu_si128((const __m128i *) &y_ptr[(K)].qs[0]); + const __m128i qy_1 = _mm_loadu_si128((const __m128i *) &y_ptr[(K)].qs[16]); + const __m128i sign_mask_0 = _mm_cmpeq_epi8(bit_mask_0, zero); + const __m128i sign_mask_1 = _mm_cmpeq_epi8(bit_mask_1, zero); + const __m128i sy_0 = _mm_sub_epi8(_mm_xor_si128(qy_0, sign_mask_0), sign_mask_0); + const __m128i sy_1 = _mm_sub_epi8(_mm_xor_si128(qy_1, sign_mask_1), sign_mask_1); + const __m128i sum16_0 = _mm_maddubs_epi16(ones_8, sy_0); + const __m128i sum16_1 = _mm_maddubs_epi16(ones_8, sy_1); + const __m128i sum32_0 = _mm_madd_epi16(sum16_0, ones_16); + const __m128i sum32_1 = _mm_madd_epi16(sum16_1, ones_16); + const __m256 q = _mm256_cvtepi32_ps(MM256_SET_M128I(sum32_1, sum32_0)); + acc_block = _mm256_add_ps(acc_block, _mm256_mul_ps(_mm256_set1_ps(GGML_CPU_FP16_TO_FP32(y_ptr[(K)].d)), q)); + } +#undef Q1_AVX_BLOCK + + acc = _mm256_add_ps(acc, _mm256_mul_ps(_mm256_set1_ps(d0), acc_block)); + } + + *s = hsum_float_8(acc); +#elif defined(__SSSE3__) + const __m128i ones_8 = _mm_set1_epi8(1); + const __m128i ones_16 = _mm_set1_epi16(1); + const __m128i zero = _mm_setzero_si128(); + __m128 acc_0 = _mm_setzero_ps(); + __m128 acc_1 = _mm_setzero_ps(); + __m128 acc_2 = _mm_setzero_ps(); + __m128 acc_3 = _mm_setzero_ps(); + + for (int ib = 0; ib < nb; ++ib) { + const __m128 d0 = _mm_set1_ps(GGML_CPU_FP16_TO_FP32(x[ib].d)); + const block_q8_0 * GGML_RESTRICT y_ptr = &y[ib * 4]; + +#define Q1_SSSE3_BLOCK(QS_OFF, Y_IDX, ACC) \ + { \ + const __m128i bit_mask_0 = bytes_from_bits_16(&x[ib].qs[(QS_OFF) + 0]); \ + const __m128i bit_mask_1 = bytes_from_bits_16(&x[ib].qs[(QS_OFF) + 2]); \ + const __m128i qy_0 = _mm_loadu_si128((const __m128i *) &y_ptr[(Y_IDX)].qs[0]); \ + const __m128i qy_1 = _mm_loadu_si128((const __m128i *) &y_ptr[(Y_IDX)].qs[16]); \ + const __m128i sign_mask_0 = _mm_cmpeq_epi8(bit_mask_0, zero); \ + const __m128i sign_mask_1 = _mm_cmpeq_epi8(bit_mask_1, zero); \ + const __m128i sy_0 = _mm_sub_epi8(_mm_xor_si128(qy_0, sign_mask_0), sign_mask_0); \ + const __m128i sy_1 = _mm_sub_epi8(_mm_xor_si128(qy_1, sign_mask_1), sign_mask_1); \ + const __m128i sum_0 = _mm_madd_epi16(_mm_maddubs_epi16(ones_8, sy_0), ones_16); \ + const __m128i sum_1 = _mm_madd_epi16(_mm_maddubs_epi16(ones_8, sy_1), ones_16); \ + const __m128 q = _mm_cvtepi32_ps(_mm_add_epi32(sum_0, sum_1)); \ + (ACC) = _mm_add_ps((ACC), _mm_mul_ps(_mm_mul_ps(d0, _mm_set1_ps(GGML_CPU_FP16_TO_FP32(y_ptr[(Y_IDX)].d))), q)); \ + } + Q1_SSSE3_BLOCK(0, 0, acc_0) + Q1_SSSE3_BLOCK(4, 1, acc_1) + Q1_SSSE3_BLOCK(8, 2, acc_2) + Q1_SSSE3_BLOCK(12, 3, acc_3) +#undef Q1_SSSE3_BLOCK + } + + *s = hsum_float_4x4(acc_0, acc_1, acc_2, acc_3); +#else + UNUSED(nb); + UNUSED(x); + UNUSED(y); + ggml_vec_dot_q1_0_q8_0_generic(n, s, bs, vx, bx, vy, by, nrc); +#endif +} + +#if defined(__AVX2__) +// four consecutive Q2_0 weights per byte: byte b becomes b & 3, (b >> 2) & 3, (b >> 4) & 3, (b >> 6) & 3 (OR, never add: the overlapping copies only touch masked-off bits) +static inline __m256i q2_0_unpack_32_avx2(const uint8_t * qs) { + __m256i v = _mm256_cvtepu8_epi32(_mm_loadl_epi64((const __m128i *) qs)); + v = _mm256_or_si256(v, _mm256_slli_epi32(v, 12)); + v = _mm256_or_si256(v, _mm256_slli_epi32(v, 6)); + return _mm256_and_si256(v, _mm256_set1_epi32(0x03030303)); +} + +// 32 weights against one Q8_0 block: sum((v - 1) * y) = sum(v * y) - sum(y), exact for every int8 y +static inline __m256 q2_0_dot_32_avx2(const uint8_t * qs, const int8_t * qy) { + const __m256i v = q2_0_unpack_32_avx2(qs); + const __m256i y = _mm256_loadu_si256((const __m256i *) qy); + const __m256i p = _mm256_sub_epi16(_mm256_maddubs_epi16(v, y), _mm256_maddubs_epi16(_mm256_set1_epi8(1), y)); + return _mm256_cvtepi32_ps(_mm256_madd_epi16(p, _mm256_set1_epi16(1))); +} + +static inline __m256 q2_0_acc_block_avx2(const block_q2_0 * x, const block_q8_0 * y, __m256 acc) { + const float dx = GGML_CPU_FP16_TO_FP32(x->d); + const __m256 d0 = _mm256_set1_ps(dx * GGML_CPU_FP16_TO_FP32(y[0].d)); + const __m256 d1 = _mm256_set1_ps(dx * GGML_CPU_FP16_TO_FP32(y[1].d)); + const __m256 p0 = q2_0_dot_32_avx2(x->qs + 0, y[0].qs); + const __m256 p1 = q2_0_dot_32_avx2(x->qs + 8, y[1].qs); +#if defined(__FMA__) + acc = _mm256_fmadd_ps(d0, p0, acc); + return _mm256_fmadd_ps(d1, p1, acc); +#else + acc = _mm256_add_ps(acc, _mm256_mul_ps(d0, p0)); + return _mm256_add_ps(acc, _mm256_mul_ps(d1, p1)); +#endif +} +#endif + +void ggml_vec_dot_q2_0_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) { +#if defined(__AVX2__) + assert(n % QK2_0 == 0); + assert(nrc == 1); + UNUSED(nrc); + UNUSED(bx); + UNUSED(by); + UNUSED(bs); + + const block_q2_0 * GGML_RESTRICT x = vx; + const block_q8_0 * GGML_RESTRICT y = vy; + const int nb = n / QK2_0; + + // two Q2_0 blocks per step, each into its own accumulator (two FMA chains) + __m256 acc0 = _mm256_setzero_ps(); + __m256 acc1 = _mm256_setzero_ps(); + int i = 0; + for (; i + 1 < nb; i += 2) { + acc0 = q2_0_acc_block_avx2(x + i + 0, y + 2*i + 0, acc0); + acc1 = q2_0_acc_block_avx2(x + i + 1, y + 2*i + 2, acc1); + } + if (i < nb) { + acc0 = q2_0_acc_block_avx2(x + i, y + 2*i, acc0); + } + + *s = hsum_float_8(_mm256_add_ps(acc0, acc1)); +#else + ggml_vec_dot_q2_0_q8_0_generic(n, s, bs, vx, bx, vy, by, nrc); +#endif +} + +void ggml_vec_dot_q4_0_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) { + const int qk = QK8_0; + const int nb = n / qk; + + assert(n % qk == 0); + assert(nrc == 1); + UNUSED(nrc); + UNUSED(bx); + UNUSED(by); + UNUSED(bs); + + const block_q4_0 * GGML_RESTRICT x = vx; + const block_q8_0 * GGML_RESTRICT y = vy; + + int ib = 0; + float sumf = 0; + +#if defined(__AVX2__) + // Initialize accumulator with zeros + __m256 acc = _mm256_setzero_ps(); + + // Main loop + for (; ib < nb; ++ib) { + /* Compute combined scale for the block */ + const __m256 d = _mm256_set1_ps( GGML_CPU_FP16_TO_FP32(x[ib].d) * GGML_CPU_FP16_TO_FP32(y[ib].d) ); + + __m256i qx = bytes_from_nibbles_32(x[ib].qs); + + // Now we have a vector with bytes in [ 0 .. 15 ] interval. Offset them into [ -8 .. +7 ] interval. + const __m256i off = _mm256_set1_epi8( 8 ); + qx = _mm256_sub_epi8( qx, off ); + + __m256i qy = _mm256_loadu_si256((const __m256i *)y[ib].qs); + + const __m256 q = mul_sum_i8_pairs_float(qx, qy); + + /* Multiply q with scale and accumulate */ + acc = _mm256_fmadd_ps( d, q, acc ); + } + + sumf = hsum_float_8(acc); +#elif defined(__AVX__) + __m256 accum = _mm256_setzero_ps(); + for (; ib + 1 < nb; ib += 2) { + const __m128i q4bits_1 = _mm_loadu_si128((const __m128i *)x[ib + 0].qs); + const __m128i q4bits_2 = _mm_loadu_si128((const __m128i *)x[ib + 1].qs); + const __m128i q8b_1_0 = _mm_loadu_si128((const __m128i *)y[ib + 0].qs); + const __m128i q8b_1_1 = _mm_loadu_si128((const __m128i *)y[ib + 0].qs + 1); + const __m128i q8b_2_0 = _mm_loadu_si128((const __m128i *)y[ib + 1].qs); + const __m128i q8b_2_1 = _mm_loadu_si128((const __m128i *)y[ib + 1].qs + 1); + + const __m128i q4b_1_0 = _mm_sub_epi8(_mm_and_si128(_mm_set1_epi8(15), q4bits_1), _mm_set1_epi8(8)); + const __m128i q4b_1_1 = _mm_sub_epi8(_mm_and_si128(_mm_set1_epi8(15), _mm_srli_epi16(q4bits_1, 4)), _mm_set1_epi8(8)); + const __m128i q4b_2_0 = _mm_sub_epi8(_mm_and_si128(_mm_set1_epi8(15), q4bits_2), _mm_set1_epi8(8)); + const __m128i q4b_2_1 = _mm_sub_epi8(_mm_and_si128(_mm_set1_epi8(15), _mm_srli_epi16(q4bits_2, 4)), _mm_set1_epi8(8)); + + const __m128i p16_1_0 = mul_add_epi8_sse(q4b_1_0, q8b_1_0); + const __m128i p16_1_1 = mul_add_epi8_sse(q4b_1_1, q8b_1_1); + const __m128i p16_2_0 = mul_add_epi8_sse(q4b_2_0, q8b_2_0); + const __m128i p16_2_1 = mul_add_epi8_sse(q4b_2_1, q8b_2_1); + const __m128i p_1 = _mm_add_epi16(p16_1_0, p16_1_1); + const __m128i p_2 = _mm_add_epi16(p16_2_0, p16_2_1); + const __m256 p = sum_i16_pairs_float(p_2, p_1); + + const __m256 deltas = quad_fp16_delta_float(x[ib].d, y[ib].d, x[ib + 1].d, y[ib + 1].d); + accum = _mm256_add_ps(_mm256_mul_ps(deltas, p), accum); + } + + sumf = hsum_float_8(accum); +#elif defined(__SSSE3__) + // set constants + const __m128i lowMask = _mm_set1_epi8(0xF); + const __m128i off = _mm_set1_epi8(8); + + // Initialize accumulator with zeros + __m128 acc_0 = _mm_setzero_ps(); + __m128 acc_1 = _mm_setzero_ps(); + __m128 acc_2 = _mm_setzero_ps(); + __m128 acc_3 = _mm_setzero_ps(); + + for (; ib + 1 < nb; ib += 2) { + _mm_prefetch(&x[ib] + sizeof(block_q4_0), _MM_HINT_T0); + _mm_prefetch(&y[ib] + sizeof(block_q8_0), _MM_HINT_T0); + + // Compute combined scale for the block 0 and 1 + const __m128 d_0_1 = _mm_set1_ps( GGML_CPU_FP16_TO_FP32(x[ib].d) * GGML_CPU_FP16_TO_FP32(y[ib].d) ); + + const __m128i tmp_0_1 = _mm_loadu_si128((const __m128i *)x[ib].qs); + + __m128i bx_0 = _mm_and_si128(lowMask, tmp_0_1); + __m128i by_0 = _mm_loadu_si128((const __m128i *)y[ib].qs); + bx_0 = _mm_sub_epi8(bx_0, off); + const __m128i i32_0 = mul_sum_i8_pairs(bx_0, by_0); + + __m128i bx_1 = _mm_and_si128(lowMask, _mm_srli_epi64(tmp_0_1, 4)); + __m128i by_1 = _mm_loadu_si128((const __m128i *)(y[ib].qs + 16)); + bx_1 = _mm_sub_epi8(bx_1, off); + const __m128i i32_1 = mul_sum_i8_pairs(bx_1, by_1); + + _mm_prefetch(&x[ib] + 2 * sizeof(block_q4_0), _MM_HINT_T0); + _mm_prefetch(&y[ib] + 2 * sizeof(block_q8_0), _MM_HINT_T0); + + // Compute combined scale for the block 2 and 3 + const __m128 d_2_3 = _mm_set1_ps( GGML_CPU_FP16_TO_FP32(x[ib + 1].d) * GGML_CPU_FP16_TO_FP32(y[ib + 1].d) ); + + const __m128i tmp_2_3 = _mm_loadu_si128((const __m128i *)x[ib + 1].qs); + + __m128i bx_2 = _mm_and_si128(lowMask, tmp_2_3); + __m128i by_2 = _mm_loadu_si128((const __m128i *)y[ib + 1].qs); + bx_2 = _mm_sub_epi8(bx_2, off); + const __m128i i32_2 = mul_sum_i8_pairs(bx_2, by_2); + + __m128i bx_3 = _mm_and_si128(lowMask, _mm_srli_epi64(tmp_2_3, 4)); + __m128i by_3 = _mm_loadu_si128((const __m128i *)(y[ib + 1].qs + 16)); + bx_3 = _mm_sub_epi8(bx_3, off); + const __m128i i32_3 = mul_sum_i8_pairs(bx_3, by_3); + + // Convert int32_t to float + __m128 p0 = _mm_cvtepi32_ps(i32_0); + __m128 p1 = _mm_cvtepi32_ps(i32_1); + __m128 p2 = _mm_cvtepi32_ps(i32_2); + __m128 p3 = _mm_cvtepi32_ps(i32_3); + + // Apply the scale + __m128 p0_d = _mm_mul_ps( d_0_1, p0 ); + __m128 p1_d = _mm_mul_ps( d_0_1, p1 ); + __m128 p2_d = _mm_mul_ps( d_2_3, p2 ); + __m128 p3_d = _mm_mul_ps( d_2_3, p3 ); + + // Acummulate + acc_0 = _mm_add_ps(p0_d, acc_0); + acc_1 = _mm_add_ps(p1_d, acc_1); + acc_2 = _mm_add_ps(p2_d, acc_2); + acc_3 = _mm_add_ps(p3_d, acc_3); + } + + sumf = hsum_float_4x4(acc_0, acc_1, acc_2, acc_3); + +#endif + for (; ib < nb; ++ib) { + int sumi0 = 0; + int sumi1 = 0; + + for (int j = 0; j < qk/2; ++j) { + const int v0 = (x[ib].qs[j] & 0x0F) - 8; + const int v1 = (x[ib].qs[j] >> 4) - 8; + + sumi0 += (v0 * y[ib].qs[j]); + sumi1 += (v1 * y[ib].qs[j + qk/2]); + } + + int sumi = sumi0 + sumi1; + sumf += sumi*GGML_CPU_FP16_TO_FP32(x[ib].d)*GGML_CPU_FP16_TO_FP32(y[ib].d); + } + + *s = sumf; +} + +void ggml_vec_dot_q4_1_q8_1(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) { + const int qk = QK8_1; + const int nb = n / qk; + + assert(n % qk == 0); + assert(nrc == 1); + UNUSED(nrc); + UNUSED(bx); + UNUSED(by); + UNUSED(bs); + + const block_q4_1 * GGML_RESTRICT x = vx; + const block_q8_1 * GGML_RESTRICT y = vy; + + int ib = 0; + +#if defined(__AVX2__) || defined(__AVX__) + // Initialize accumulator with zeros + __m256 acc = _mm256_setzero_ps(); + + float summs = 0; + + // Main loop + for (; ib < nb; ++ib) { + const float d0 = GGML_CPU_FP16_TO_FP32(x[ib].d); + const float d1 = GGML_CPU_FP16_TO_FP32(y[ib].d); + + summs += GGML_CPU_FP16_TO_FP32(x[ib].m) * GGML_CPU_FP16_TO_FP32(y[ib].s); + + const __m256 d0v = _mm256_set1_ps( d0 ); + const __m256 d1v = _mm256_set1_ps( d1 ); + + // Compute combined scales + const __m256 d0d1 = _mm256_mul_ps( d0v, d1v ); + + // Load 16 bytes, and unpack 4 bit fields into bytes, making 32 bytes + const __m256i qx = bytes_from_nibbles_32(x[ib].qs); + const __m256i qy = _mm256_loadu_si256( (const __m256i *)y[ib].qs ); + + const __m256 xy = mul_sum_us8_pairs_float(qx, qy); + + // Accumulate d0*d1*x*y +#if defined(__AVX2__) + acc = _mm256_fmadd_ps( d0d1, xy, acc ); +#else + acc = _mm256_add_ps( _mm256_mul_ps( d0d1, xy ), acc ); +#endif + } + + *s = hsum_float_8(acc) + summs; +#else + UNUSED(nb); + UNUSED(x); + UNUSED(y); + UNUSED(ib); + ggml_vec_dot_q4_1_q8_1_generic(n, s, bs, vx, bx, vy, by, nrc); +#endif +} + +void ggml_vec_dot_mxfp4_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) { + assert(nrc == 1); + UNUSED(nrc); + UNUSED(bx); + UNUSED(by); + UNUSED(bs); + assert(n % QK_MXFP4 == 0); + static_assert(QK_MXFP4 == QK8_0, "QK_MXFP4 and QK8_0 must be the same"); + + const block_mxfp4 * GGML_RESTRICT x = vx; + const block_q8_0 * GGML_RESTRICT y = vy; + + const int nb = n / QK_MXFP4; + + int ib = 0; + float sumf = 0; + +#if defined __AVX2__ + + const __m128i values128 = _mm_loadu_si128((const __m128i*)kvalues_fp4); + const __m128i m4b = _mm_set1_epi8(0x0f); + const __m256i mone = _mm256_set1_epi16(1); + + __m256 accum1 = _mm256_setzero_ps(); + __m256 accum2 = _mm256_setzero_ps(); + + for (; ib + 1 < nb; ib += 2) { + const __m128i q4bits_1 = _mm_loadu_si128((const __m128i*)x[ib + 0].qs); + const __m128i q4bits_2 = _mm_loadu_si128((const __m128i*)x[ib + 1].qs); + const __m256i q8b_1 = _mm256_loadu_si256((const __m256i *)y[ib + 0].qs); + const __m256i q8b_2 = _mm256_loadu_si256((const __m256i *)y[ib + 1].qs); + const __m256i q4b_1 = MM256_SET_M128I(_mm_shuffle_epi8(values128, _mm_and_si128(_mm_srli_epi16(q4bits_1, 4), m4b)), + _mm_shuffle_epi8(values128, _mm_and_si128(q4bits_1, m4b))); + const __m256i q4b_2 = MM256_SET_M128I(_mm_shuffle_epi8(values128, _mm_and_si128(_mm_srli_epi16(q4bits_2, 4), m4b)), + _mm_shuffle_epi8(values128, _mm_and_si128(q4bits_2, m4b))); + const __m256i p16_1 = mul_add_epi8(q4b_1, q8b_1); + const __m256i p16_2 = mul_add_epi8(q4b_2, q8b_2); + const __m256i p_1 = _mm256_madd_epi16(p16_1, mone); + const __m256i p_2 = _mm256_madd_epi16(p16_2, mone); + const __m256 scale0 = _mm256_set1_ps(GGML_CPU_FP16_TO_FP32(y[ib + 0].d)*GGML_CPU_E8M0_TO_FP32_HALF(x[ib + 0].e)); + const __m256 scale1 = _mm256_set1_ps(GGML_CPU_FP16_TO_FP32(y[ib + 1].d)*GGML_CPU_E8M0_TO_FP32_HALF(x[ib + 1].e)); + accum1 = _mm256_fmadd_ps(scale0, _mm256_cvtepi32_ps(p_1), accum1); + accum2 = _mm256_fmadd_ps(scale1, _mm256_cvtepi32_ps(p_2), accum2); + } + + sumf = hsum_float_8(_mm256_add_ps(accum1, accum2)); + +#elif defined __AVX__ + const __m128i values128 = _mm_loadu_si128((const __m128i*)kvalues_fp4); + const __m128i m4b = _mm_set1_epi8(0x0f); + + __m256 accum = _mm256_setzero_ps(); + for (; ib + 1 < nb; ib += 2) { + const __m128i q4bits_1 = _mm_loadu_si128((const __m128i *)x[ib + 0].qs); + const __m128i q4bits_2 = _mm_loadu_si128((const __m128i *)x[ib + 1].qs); + const __m128i q8b_1_0 = _mm_loadu_si128((const __m128i *)y[ib + 0].qs); + const __m128i q8b_1_1 = _mm_loadu_si128((const __m128i *)y[ib + 0].qs + 1); + const __m128i q8b_2_0 = _mm_loadu_si128((const __m128i *)y[ib + 1].qs); + const __m128i q8b_2_1 = _mm_loadu_si128((const __m128i *)y[ib + 1].qs + 1); + + const __m128i q4b_1_0 = _mm_shuffle_epi8(values128, _mm_and_si128(q4bits_1, m4b)); + const __m128i q4b_1_1 = _mm_shuffle_epi8(values128, _mm_and_si128(_mm_srli_epi16(q4bits_1, 4), m4b)); + const __m128i q4b_2_0 = _mm_shuffle_epi8(values128, _mm_and_si128(q4bits_2, m4b)); + const __m128i q4b_2_1 = _mm_shuffle_epi8(values128, _mm_and_si128(_mm_srli_epi16(q4bits_2, 4), m4b)); + + const __m256 p = mul_sum_i8_quad_float(q4b_1_0, q4b_1_1, q4b_2_0, q4b_2_1, q8b_1_0, q8b_1_1, q8b_2_0, q8b_2_1); + const __m256 deltas = quad_mx_delta_float(x[ib].e, y[ib].d, x[ib + 1].e, y[ib + 1].d); + accum = _mm256_add_ps(_mm256_mul_ps(deltas, p), accum); + } + + sumf = hsum_float_8(accum); + +#endif + for (; ib < nb; ++ib) { + const float d = GGML_CPU_FP16_TO_FP32(y[ib].d)*GGML_CPU_E8M0_TO_FP32_HALF(x[ib].e); + int sumi1 = 0; + int sumi2 = 0; + for (int j = 0; j < QK_MXFP4/2; ++j) { + sumi1 += y[ib].qs[j + 0] * kvalues_fp4[x[ib].qs[j] & 0xf]; + sumi2 += y[ib].qs[j + QK_MXFP4/2] * kvalues_fp4[x[ib].qs[j] >> 4]; + } + sumf += d * (sumi1 + sumi2); + } + *s = sumf; +} + +void ggml_vec_dot_nvfp4_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) { + assert(nrc == 1); + UNUSED(nrc); + UNUSED(bx); + UNUSED(by); + UNUSED(bs); + assert(n % QK_NVFP4 == 0); + + const block_nvfp4 * GGML_RESTRICT x = vx; + const block_q8_0 * GGML_RESTRICT y = vy; + + const int nb = n / QK_NVFP4; + int ib = 0; + float sumf = 0; + +#if defined(__AVX2__) + + const __m128i values128 = _mm_loadu_si128((const __m128i*)kvalues_fp4); + const __m128i m4b = _mm_set1_epi8(0x0f); + const __m256i mone = _mm256_set1_epi16(1); + + __m256 accum = _mm256_setzero_ps(); + for(; ib < nb; ib++){ + + const __m128i q4bits_01 = _mm_loadu_si128((const __m128i *)(x[ib].qs + 0)); + const __m128i q4bits_23 = _mm_loadu_si128((const __m128i *)(x[ib].qs + 16)); + + const __m256i q8_01 = _mm256_loadu_si256((const __m256i *)y[2*ib + 0].qs); + const __m256i q8_23 = _mm256_loadu_si256((const __m256i *)y[2*ib + 1].qs); + + const __m128i q4_01_lo = _mm_shuffle_epi8(values128, _mm_and_si128(q4bits_01, m4b)); + const __m128i q4_01_hi = _mm_shuffle_epi8(values128, _mm_and_si128(_mm_srli_epi16(q4bits_01, 4), m4b)); + const __m128i q4_23_lo = _mm_shuffle_epi8(values128, _mm_and_si128(q4bits_23, m4b)); + const __m128i q4_23_hi = _mm_shuffle_epi8(values128, _mm_and_si128(_mm_srli_epi16(q4bits_23, 4), m4b)); + + //reordering + const __m256i q4_01 = MM256_SET_M128I(_mm_unpackhi_epi64(q4_01_lo,q4_01_hi), _mm_unpacklo_epi64(q4_01_lo,q4_01_hi)); + const __m256i q4_23 = MM256_SET_M128I(_mm_unpackhi_epi64(q4_23_lo,q4_23_hi),_mm_unpacklo_epi64(q4_23_lo,q4_23_hi)); + + const __m256i p01 = mul_add_epi8(q4_01,q8_01); + const __m256i p_1 = _mm256_madd_epi16(p01, mone); + + const __m256i p23 = mul_add_epi8(q4_23,q8_23); + const __m256i p_2 = _mm256_madd_epi16(p23, mone); + + const float dy0 = GGML_CPU_FP16_TO_FP32(y[2*ib].d); + const float dy1 = GGML_CPU_FP16_TO_FP32(y[2*ib+1].d); + + const float s0 = GGML_CPU_UE4M3_TO_FP32(x[ib].d[0]) * dy0; + const float s1 = GGML_CPU_UE4M3_TO_FP32(x[ib].d[1]) * dy0; + const float s2 = GGML_CPU_UE4M3_TO_FP32(x[ib].d[2]) * dy1; + const float s3 = GGML_CPU_UE4M3_TO_FP32(x[ib].d[3]) * dy1; + + const __m256 scales01 = _mm256_set_m128(_mm_set1_ps(s1), _mm_set1_ps(s0)); + const __m256 scales23 = _mm256_set_m128(_mm_set1_ps(s3), _mm_set1_ps(s2)); + + accum = _mm256_fmadd_ps(scales01, _mm256_cvtepi32_ps(p_1), accum); + accum = _mm256_fmadd_ps(scales23, _mm256_cvtepi32_ps(p_2), accum); + } + sumf = hsum_float_8(accum); + +#elif defined(__AVX__) + + const __m128i values128 = _mm_loadu_si128((const __m128i*)kvalues_fp4); + const __m128i m4b = _mm_set1_epi8(0x0f); + + __m256 accum = _mm256_setzero_ps(); + for(; ib < nb; ib++){ + + const __m128i q4bits_01 = _mm_loadu_si128((const __m128i *)(x[ib].qs + 0)); + const __m128i q4bits_23 = _mm_loadu_si128((const __m128i *)(x[ib].qs + 16)); + + const __m128i q8_0 = _mm_loadu_si128((const __m128i *)(y[2*ib + 0].qs + 0)); + const __m128i q8_1 = _mm_loadu_si128((const __m128i *)(y[2*ib + 0].qs + 16)); + const __m128i q8_2 = _mm_loadu_si128((const __m128i *)(y[2*ib + 1].qs + 0)); + const __m128i q8_3 = _mm_loadu_si128((const __m128i *)(y[2*ib + 1].qs + 16)); + + const __m128i q4_01_lo = _mm_shuffle_epi8(values128, _mm_and_si128(q4bits_01, m4b)); + const __m128i q4_01_hi = _mm_shuffle_epi8(values128, _mm_and_si128(_mm_srli_epi16(q4bits_01, 4), m4b)); + const __m128i q4_23_lo = _mm_shuffle_epi8(values128, _mm_and_si128(q4bits_23, m4b)); + const __m128i q4_23_hi = _mm_shuffle_epi8(values128, _mm_and_si128(_mm_srli_epi16(q4bits_23, 4), m4b)); + + const __m128i q4_0 = _mm_unpacklo_epi64(q4_01_lo, q4_01_hi); + const __m128i q4_1 = _mm_unpackhi_epi64(q4_01_lo, q4_01_hi); + const __m128i q4_2 = _mm_unpacklo_epi64(q4_23_lo, q4_23_hi); + const __m128i q4_3 = _mm_unpackhi_epi64(q4_23_lo, q4_23_hi); + + const __m128i p0_i32 = mul_sum_i8_pairs(q4_0, q8_0); + const __m128i p1_i32 = mul_sum_i8_pairs(q4_1, q8_1); + const __m128i p2_i32 = mul_sum_i8_pairs(q4_2, q8_2); + const __m128i p3_i32 = mul_sum_i8_pairs(q4_3, q8_3); + + const __m128 p0 = _mm_cvtepi32_ps(p0_i32); + const __m128 p1 = _mm_cvtepi32_ps(p1_i32); + const __m128 p2 = _mm_cvtepi32_ps(p2_i32); + const __m128 p3 = _mm_cvtepi32_ps(p3_i32); + + const __m256 p01 = _mm256_set_m128(p1, p0); + const __m256 p23 = _mm256_set_m128(p3, p2); + + const float dy0 = GGML_CPU_FP16_TO_FP32(y[2*ib].d); + const float dy1 = GGML_CPU_FP16_TO_FP32(y[2*ib+1].d); + + const float s0 = GGML_CPU_UE4M3_TO_FP32(x[ib].d[0]) * dy0; + const float s1 = GGML_CPU_UE4M3_TO_FP32(x[ib].d[1]) * dy0; + const float s2 = GGML_CPU_UE4M3_TO_FP32(x[ib].d[2]) * dy1; + const float s3 = GGML_CPU_UE4M3_TO_FP32(x[ib].d[3]) * dy1; + + const __m256 scales01 = _mm256_set_m128(_mm_set1_ps(s1), _mm_set1_ps(s0)); + const __m256 scales23 = _mm256_set_m128(_mm_set1_ps(s3), _mm_set1_ps(s2)); + + accum = _mm256_add_ps(accum, _mm256_mul_ps(p01, scales01)); + accum = _mm256_add_ps(accum, _mm256_mul_ps(p23, scales23)); + } + sumf = hsum_float_8(accum); + +#endif + + for (;ib < nb; ++ib) { + for (int s_idx = 0; s_idx < 4; ++s_idx) { + const float d = GGML_CPU_UE4M3_TO_FP32(x[ib].d[s_idx]); + const int q8_block = s_idx / 2; + const int q8_off = (s_idx % 2) * QK_NVFP4_SUB; + const float dy = GGML_CPU_FP16_TO_FP32(y[2*ib + q8_block].d); + + int sumi_lo = 0, sumi_hi = 0; + for (int j = 0; j < QK_NVFP4_SUB/2; ++j) { + const uint8_t qv = x[ib].qs[s_idx*(QK_NVFP4_SUB/2) + j]; + sumi_lo += y[2*ib + q8_block].qs[q8_off + j + 0] * kvalues_fp4[qv & 0xf]; + sumi_hi += y[2*ib + q8_block].qs[q8_off + j + QK_NVFP4_SUB/2] * kvalues_fp4[qv >> 4]; + } + + sumf += dy * d * (sumi_lo + sumi_hi); + } + } + *s = sumf; +} + +void ggml_vec_dot_q5_0_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) { + const int qk = QK8_0; + const int nb = n / qk; + + int ib = 0; + + assert(n % qk == 0); + assert(qk == QK5_0); + assert(nrc == 1); + UNUSED(nrc); + UNUSED(bx); + UNUSED(by); + UNUSED(bs); + + const block_q5_0 * GGML_RESTRICT x = vx; + const block_q8_0 * GGML_RESTRICT y = vy; + +#if defined(__AVX2__) + // Initialize accumulator with zeros + __m256 acc = _mm256_setzero_ps(); + + // Main loop + for (; ib < nb; ++ib) { + /* Compute combined scale for the block */ + const __m256 d = _mm256_set1_ps(GGML_CPU_FP16_TO_FP32(x[ib].d) * GGML_CPU_FP16_TO_FP32(y[ib].d)); + + __m256i qx = bytes_from_nibbles_32(x[ib].qs); + __m256i bxhi = bytes_from_bits_32(x[ib].qh); + bxhi = _mm256_andnot_si256(bxhi, _mm256_set1_epi8((char)0xF0)); + qx = _mm256_or_si256(qx, bxhi); + + __m256i qy = _mm256_loadu_si256((const __m256i *)y[ib].qs); + + const __m256 q = mul_sum_i8_pairs_float(qx, qy); + + /* Multiply q with scale and accumulate */ + acc = _mm256_fmadd_ps(d, q, acc); + } + + *s = hsum_float_8(acc); +#elif defined(__AVX__) + // Initialize accumulator with zeros + __m256 acc = _mm256_setzero_ps(); + __m128i mask = _mm_set1_epi8((char)0xF0); + + // Main loop + for (; ib < nb; ++ib) { + /* Compute combined scale for the block */ + const __m256 d = _mm256_set1_ps(GGML_CPU_FP16_TO_FP32(x[ib].d) * GGML_CPU_FP16_TO_FP32(y[ib].d)); + + __m256i bx_0 = bytes_from_nibbles_32(x[ib].qs); + const __m256i bxhi = bytes_from_bits_32(x[ib].qh); + __m128i bxhil = _mm256_castsi256_si128(bxhi); + __m128i bxhih = _mm256_extractf128_si256(bxhi, 1); + bxhil = _mm_andnot_si128(bxhil, mask); + bxhih = _mm_andnot_si128(bxhih, mask); + __m128i bxl = _mm256_castsi256_si128(bx_0); + __m128i bxh = _mm256_extractf128_si256(bx_0, 1); + bxl = _mm_or_si128(bxl, bxhil); + bxh = _mm_or_si128(bxh, bxhih); + bx_0 = MM256_SET_M128I(bxh, bxl); + + const __m256i by_0 = _mm256_loadu_si256((const __m256i *)y[ib].qs); + + const __m256 q = mul_sum_i8_pairs_float(bx_0, by_0); + + /* Multiply q with scale and accumulate */ + acc = _mm256_add_ps(_mm256_mul_ps(d, q), acc); + } + + *s = hsum_float_8(acc); +#else + UNUSED(nb); + UNUSED(ib); + UNUSED(x); + UNUSED(y); + ggml_vec_dot_q5_0_q8_0_generic(n, s, bs, vx, bx, vy, by, nrc); +#endif +} + +void ggml_vec_dot_q5_1_q8_1(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) { + const int qk = QK8_1; + const int nb = n / qk; + + int ib = 0; + + assert(n % qk == 0); + assert(qk == QK5_1); + assert(nrc == 1); + UNUSED(nrc); + UNUSED(bx); + UNUSED(by); + UNUSED(bs); + + const block_q5_1 * GGML_RESTRICT x = vx; + const block_q8_1 * GGML_RESTRICT y = vy; + +#if defined(__AVX2__) + // Initialize accumulator with zeros + __m256 acc = _mm256_setzero_ps(); + + float summs = 0.0f; + + // Main loop + for (; ib < nb; ++ib) { + const __m256 dx = _mm256_set1_ps(GGML_CPU_FP16_TO_FP32(x[ib].d)); + + summs += GGML_CPU_FP16_TO_FP32(x[ib].m) * GGML_CPU_FP16_TO_FP32(y[ib].s); + + __m256i qx = bytes_from_nibbles_32(x[ib].qs); + __m256i bxhi = bytes_from_bits_32(x[ib].qh); + bxhi = _mm256_and_si256(bxhi, _mm256_set1_epi8(0x10)); + qx = _mm256_or_si256(qx, bxhi); + + const __m256 dy = _mm256_set1_ps(GGML_CPU_FP16_TO_FP32(y[ib].d)); + const __m256i qy = _mm256_loadu_si256((const __m256i *)y[ib].qs); + + const __m256 q = mul_sum_us8_pairs_float(qx, qy); + + acc = _mm256_fmadd_ps(q, _mm256_mul_ps(dx, dy), acc); + } + + *s = hsum_float_8(acc) + summs; +#elif defined(__AVX__) + // Initialize accumulator with zeros + __m256 acc = _mm256_setzero_ps(); + __m128i mask = _mm_set1_epi8(0x10); + + float summs = 0.0f; + + // Main loop + for (; ib < nb; ++ib) { + const __m256 dx = _mm256_set1_ps(GGML_CPU_FP16_TO_FP32(x[ib].d)); + + summs += GGML_CPU_FP16_TO_FP32(x[ib].m) * GGML_CPU_FP16_TO_FP32(y[ib].s); + + __m256i bx_0 = bytes_from_nibbles_32(x[ib].qs); + const __m256i bxhi = bytes_from_bits_32(x[ib].qh); + __m128i bxhil = _mm256_castsi256_si128(bxhi); + __m128i bxhih = _mm256_extractf128_si256(bxhi, 1); + bxhil = _mm_and_si128(bxhil, mask); + bxhih = _mm_and_si128(bxhih, mask); + __m128i bxl = _mm256_castsi256_si128(bx_0); + __m128i bxh = _mm256_extractf128_si256(bx_0, 1); + bxl = _mm_or_si128(bxl, bxhil); + bxh = _mm_or_si128(bxh, bxhih); + bx_0 = MM256_SET_M128I(bxh, bxl); + + const __m256 dy = _mm256_set1_ps(GGML_CPU_FP16_TO_FP32(y[ib].d)); + const __m256i by_0 = _mm256_loadu_si256((const __m256i *)y[ib].qs); + + const __m256 q = mul_sum_us8_pairs_float(bx_0, by_0); + + acc = _mm256_add_ps(_mm256_mul_ps(q, _mm256_mul_ps(dx, dy)), acc); + } + + *s = hsum_float_8(acc) + summs; +#else + UNUSED(nb); + UNUSED(ib); + UNUSED(x); + UNUSED(y); + ggml_vec_dot_q5_1_q8_1_generic(n, s, bs, vx, bx, vy, by, nrc); +#endif +} + +void ggml_vec_dot_q8_0_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) { + const int qk = QK8_0; + const int nb = n / qk; + + assert(n % qk == 0); + assert(nrc == 1); + UNUSED(nrc); + UNUSED(bx); + UNUSED(by); + UNUSED(bs); + + const block_q8_0 * GGML_RESTRICT x = vx; + const block_q8_0 * GGML_RESTRICT y = vy; + + int ib = 0; + float sumf = 0; + +#if defined(__AVX2__) + // Initialize accumulator with zeros + __m256 acc = _mm256_setzero_ps(); + + // Main loop + for (; ib < nb; ++ib) { + // Compute combined scale for the block + const __m256 d = _mm256_set1_ps(GGML_CPU_FP16_TO_FP32(x[ib].d) * GGML_CPU_FP16_TO_FP32(y[ib].d)); + __m256i qx = _mm256_loadu_si256((const __m256i *)x[ib].qs); + __m256i qy = _mm256_loadu_si256((const __m256i *)y[ib].qs); + + const __m256 q = mul_sum_i8_pairs_float(qx, qy); + + // Multiply q with scale and accumulate + acc = _mm256_fmadd_ps( d, q, acc ); + } + + sumf = hsum_float_8(acc); +#elif defined(__AVX__) + __m256 accum = _mm256_setzero_ps(); + + for (; ib + 1 < nb; ib += 2) { + const __m128i qx_1_0 = _mm_loadu_si128((const __m128i *)x[ib].qs); + const __m128i qx_1_1 = _mm_loadu_si128((const __m128i *)x[ib].qs + 1); + const __m128i qx_2_0 = _mm_loadu_si128((const __m128i *)x[ib + 1].qs); + const __m128i qx_2_1 = _mm_loadu_si128((const __m128i *)x[ib + 1].qs + 1); + const __m128i qy_1_0 = _mm_loadu_si128((const __m128i *)y[ib].qs); + const __m128i qy_1_1 = _mm_loadu_si128((const __m128i *)y[ib].qs + 1); + const __m128i qy_2_0 = _mm_loadu_si128((const __m128i *)y[ib + 1].qs); + const __m128i qy_2_1 = _mm_loadu_si128((const __m128i *)y[ib + 1].qs + 1); + + const __m256 p = mul_sum_i8_quad_float(qx_1_0, qx_1_1, qx_2_0, qx_2_1, qy_1_0, qy_1_1, qy_2_0, qy_2_1); + const __m256 deltas = quad_fp16_delta_float(x[ib].d, y[ib].d, x[ib + 1].d, y[ib + 1].d); + accum = _mm256_add_ps(_mm256_mul_ps(deltas, p), accum); + } + + sumf = hsum_float_8(accum); +#endif + for (; ib < nb; ++ib) { + int sumi = 0; + + for (int j = 0; j < qk; j++) { + sumi += x[ib].qs[j]*y[ib].qs[j]; + } + + sumf += sumi*(GGML_CPU_FP16_TO_FP32(x[ib].d)*GGML_CPU_FP16_TO_FP32(y[ib].d)); + } + + *s = sumf; +} + +void ggml_vec_dot_tq1_0_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) { + assert(nrc == 1); + UNUSED(nrc); + UNUSED(bx); + UNUSED(by); + UNUSED(bs); + + const block_tq1_0 * GGML_RESTRICT x = vx; + const block_q8_K * GGML_RESTRICT y = vy; + + const int nb = n / QK_K; + +#if defined(__AVX2__) + __m256 sumf = _mm256_setzero_ps(); + + for (int i = 0; i < nb; ++i) { + // 16-bit sums + __m256i sumi0 = _mm256_setzero_si256(); + __m256i sumi1 = _mm256_setzero_si256(); + __m256i sumi2 = _mm256_setzero_si256(); + + // first 32 bytes of 5 elements + { + __m256i qx0 = _mm256_loadu_si256((const __m256i *) (x[i].qs)); + // 8-bit multiplies with shifts, masks and adds + __m256i qx1 = _mm256_add_epi8(qx0, _mm256_add_epi8(qx0, qx0)); // 1 * 3 + __m256i qx2 = _mm256_add_epi8(_mm256_and_si256(_mm256_slli_epi16(qx0, 3), _mm256_set1_epi8(-8)), qx0); // 1 * 9 + __m256i qx3 = _mm256_add_epi8(_mm256_and_si256(_mm256_slli_epi16(qx1, 3), _mm256_set1_epi8(-8)), qx1); // 3 * 9 + __m256i qx4 = _mm256_add_epi8(_mm256_and_si256(_mm256_slli_epi16(qx2, 3), _mm256_set1_epi8(-8)), qx2); // 9 * 9 + + // TODO: can _mm256_mulhi_epu16 be faster even if 16-bits? + + // Cancel the +1 from avg so that it behaves like a halving add + qx0 = _mm256_subs_epu8(qx0, _mm256_set1_epi8(1)); + qx1 = _mm256_subs_epu8(qx1, _mm256_set1_epi8(1)); + qx2 = _mm256_subs_epu8(qx2, _mm256_set1_epi8(1)); + qx3 = _mm256_subs_epu8(qx3, _mm256_set1_epi8(1)); + qx4 = _mm256_subs_epu8(qx4, _mm256_set1_epi8(1)); + // Multiply by 3 and get the top 2 bits + qx0 = _mm256_avg_epu8(qx0, _mm256_avg_epu8(qx0, _mm256_setzero_si256())); + qx1 = _mm256_avg_epu8(qx1, _mm256_avg_epu8(qx1, _mm256_setzero_si256())); + qx2 = _mm256_avg_epu8(qx2, _mm256_avg_epu8(qx2, _mm256_setzero_si256())); + qx3 = _mm256_avg_epu8(qx3, _mm256_avg_epu8(qx3, _mm256_setzero_si256())); + qx4 = _mm256_avg_epu8(qx4, _mm256_avg_epu8(qx4, _mm256_setzero_si256())); + qx0 = _mm256_and_si256(_mm256_srli_epi16(qx0, 6), _mm256_set1_epi8(3)); + qx1 = _mm256_and_si256(_mm256_srli_epi16(qx1, 6), _mm256_set1_epi8(3)); + qx2 = _mm256_and_si256(_mm256_srli_epi16(qx2, 6), _mm256_set1_epi8(3)); + qx3 = _mm256_and_si256(_mm256_srli_epi16(qx3, 6), _mm256_set1_epi8(3)); + qx4 = _mm256_and_si256(_mm256_srli_epi16(qx4, 6), _mm256_set1_epi8(3)); + + const __m256i qy0 = _mm256_loadu_si256((const __m256i *) (y[i].qs + 0)); + const __m256i qy1 = _mm256_loadu_si256((const __m256i *) (y[i].qs + 32)); + const __m256i qy2 = _mm256_loadu_si256((const __m256i *) (y[i].qs + 64)); + const __m256i qy3 = _mm256_loadu_si256((const __m256i *) (y[i].qs + 96)); + const __m256i qy4 = _mm256_loadu_si256((const __m256i *) (y[i].qs + 128)); + + qx0 = _mm256_maddubs_epi16(qx0, qy0); + qx1 = _mm256_maddubs_epi16(qx1, qy1); + qx2 = _mm256_maddubs_epi16(qx2, qy2); + qx3 = _mm256_maddubs_epi16(qx3, qy3); + qx4 = _mm256_maddubs_epi16(qx4, qy4); + + sumi0 = _mm256_add_epi16(sumi0, _mm256_add_epi16(qx0, qx1)); + sumi1 = _mm256_add_epi16(sumi1, _mm256_add_epi16(qx2, qx3)); + sumi2 = _mm256_add_epi16(sumi2, qx4); + } + + // last 16 bytes of 5-element, along with the 4 bytes of 4 elements + { + __m128i qx0 = _mm_loadu_si128((const __m128i *) (x[i].qs + 32)); + uint32_t qh; + memcpy(&qh, x[i].qh, sizeof(qh)); // potentially unaligned + __m256i qx5_l = _mm256_cvtepu8_epi16(_mm_set1_epi32(qh)); + __m128i qx1 = _mm_add_epi8(qx0, _mm_add_epi8(qx0, qx0)); // 1 * 3 + __m128i qx2 = _mm_add_epi8(_mm_and_si128(_mm_slli_epi16(qx0, 3), _mm_set1_epi8(-8)), qx0); // 1 * 9 + __m128i qx3 = _mm_add_epi8(_mm_and_si128(_mm_slli_epi16(qx1, 3), _mm_set1_epi8(-8)), qx1); // 3 * 9 + __m128i qx4 = _mm_add_epi8(_mm_and_si128(_mm_slli_epi16(qx2, 3), _mm_set1_epi8(-8)), qx2); // 9 * 9 + __m256i qx01 = MM256_SET_M128I(qx1, qx0); + __m256i qx23 = MM256_SET_M128I(qx3, qx2); + + // avx2 does not have 8-bit multiplies, so 16-bit it is. + qx5_l = _mm256_mullo_epi16(qx5_l, _mm256_set_epi16(27, 27, 27, 27, 9, 9, 9, 9, 3, 3, 3, 3, 1, 1, 1, 1)); + qx5_l = _mm256_and_si256(qx5_l, _mm256_set1_epi16(0xFF)); + __m128i qx5 = _mm_packus_epi16(_mm256_castsi256_si128(qx5_l), _mm256_extracti128_si256(qx5_l, 1)); + + __m256i qx45 = MM256_SET_M128I(qx5, qx4); + + // Cancel the +1 from avg so that it behaves like a halving add + qx01 = _mm256_subs_epu8(qx01, _mm256_set1_epi8(1)); + qx23 = _mm256_subs_epu8(qx23, _mm256_set1_epi8(1)); + qx45 = _mm256_subs_epu8(qx45, _mm256_set1_epi8(1)); + // Multiply by 3 and get the top 2 bits + qx01 = _mm256_avg_epu8(qx01, _mm256_avg_epu8(qx01, _mm256_setzero_si256())); + qx23 = _mm256_avg_epu8(qx23, _mm256_avg_epu8(qx23, _mm256_setzero_si256())); + qx45 = _mm256_avg_epu8(qx45, _mm256_avg_epu8(qx45, _mm256_setzero_si256())); + qx01 = _mm256_and_si256(_mm256_srli_epi16(qx01, 6), _mm256_set1_epi8(3)); + qx23 = _mm256_and_si256(_mm256_srli_epi16(qx23, 6), _mm256_set1_epi8(3)); + qx45 = _mm256_and_si256(_mm256_srli_epi16(qx45, 6), _mm256_set1_epi8(3)); + + const __m256i qy01 = _mm256_loadu_si256((const __m256i *) (y[i].qs + 160)); + const __m256i qy23 = _mm256_loadu_si256((const __m256i *) (y[i].qs + 192)); + const __m256i qy45 = _mm256_loadu_si256((const __m256i *) (y[i].qs + 224)); + + qx01 = _mm256_maddubs_epi16(qx01, qy01); + qx23 = _mm256_maddubs_epi16(qx23, qy23); + qx45 = _mm256_maddubs_epi16(qx45, qy45); + + sumi0 = _mm256_add_epi16(sumi0, qx01); + sumi1 = _mm256_add_epi16(sumi1, qx23); + sumi2 = _mm256_add_epi16(sumi2, qx45); + } + + const __m256i ysum = _mm256_loadu_si256((const __m256i *) y[i].bsums); + const __m256 d = _mm256_set1_ps(y[i].d * GGML_CPU_FP16_TO_FP32(x[i].d)); + + sumi0 = _mm256_sub_epi16(sumi0, ysum); + sumi0 = _mm256_add_epi16(sumi0, _mm256_add_epi16(sumi1, sumi2)); + sumi0 = _mm256_madd_epi16(sumi0, _mm256_set1_epi16(1)); + + sumf = _mm256_add_ps(_mm256_mul_ps(_mm256_cvtepi32_ps(sumi0), d), sumf); + } + + *s = hsum_float_8(sumf); + +#else + UNUSED(x); + UNUSED(y); + UNUSED(nb); + ggml_vec_dot_tq1_0_q8_K_generic(n, s, bs, vx, bx, vy, by, nrc); +#endif +} + +void ggml_vec_dot_tq2_0_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) { + assert(nrc == 1); + UNUSED(nrc); + UNUSED(bx); + UNUSED(by); + UNUSED(bs); + + const block_tq2_0 * GGML_RESTRICT x = vx; + const block_q8_K * GGML_RESTRICT y = vy; + + const int nb = n / QK_K; + +#if defined(__AVX2__) + __m256 sumf = _mm256_setzero_ps(); + + for (int i = 0; i < nb; ++i) { + // 16-bit sums, because 256*127 still fits + __m256i sumi0 = _mm256_setzero_si256(); + __m256i sumi1 = _mm256_setzero_si256(); + + for (size_t j = 0; j < sizeof(x->qs); j += 32) { + __m256i qx0 = _mm256_loadu_si256((const __m256i *) (x[i].qs + j)); + __m256i qx1 = _mm256_srli_epi16(qx0, 2); + __m256i qx2 = _mm256_srli_epi16(qx0, 4); + __m256i qx3 = _mm256_srli_epi16(qx0, 6); + + // 0, 1, 2 (should not be 3) + qx0 = _mm256_and_si256(qx0, _mm256_set1_epi8(3)); + qx1 = _mm256_and_si256(qx1, _mm256_set1_epi8(3)); + qx2 = _mm256_and_si256(qx2, _mm256_set1_epi8(3)); + qx3 = _mm256_and_si256(qx3, _mm256_set1_epi8(3)); + + const __m256i qy0 = _mm256_loadu_si256((const __m256i *) (y[i].qs + j*4 + 0)); + const __m256i qy1 = _mm256_loadu_si256((const __m256i *) (y[i].qs + j*4 + 32)); + const __m256i qy2 = _mm256_loadu_si256((const __m256i *) (y[i].qs + j*4 + 64)); + const __m256i qy3 = _mm256_loadu_si256((const __m256i *) (y[i].qs + j*4 + 96)); + + qx0 = _mm256_maddubs_epi16(qx0, qy0); + qx1 = _mm256_maddubs_epi16(qx1, qy1); + qx2 = _mm256_maddubs_epi16(qx2, qy2); + qx3 = _mm256_maddubs_epi16(qx3, qy3); + + sumi0 = _mm256_add_epi16(sumi0, _mm256_add_epi16(qx0, qx1)); + sumi1 = _mm256_add_epi16(sumi1, _mm256_add_epi16(qx2, qx3)); + } + + const __m256i ysum = _mm256_loadu_si256((const __m256i *) y[i].bsums); + const __m256 d = _mm256_set1_ps(y[i].d * GGML_CPU_FP16_TO_FP32(x[i].d)); + + sumi0 = _mm256_add_epi16(sumi0, sumi1); + sumi0 = _mm256_sub_epi16(sumi0, ysum); + sumi0 = _mm256_madd_epi16(sumi0, _mm256_set1_epi16(1)); + + sumf = _mm256_add_ps(_mm256_mul_ps(_mm256_cvtepi32_ps(sumi0), d), sumf); + } + + *s = hsum_float_8(sumf); + +#else + UNUSED(x); + UNUSED(y); + UNUSED(nb); + ggml_vec_dot_tq2_0_q8_K_generic(n, s, bs, vx, bx, vy, by, nrc); +#endif +} + +void ggml_vec_dot_q2_K_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) { + assert(nrc == 1); + UNUSED(nrc); + UNUSED(bx); + UNUSED(by); + UNUSED(bs); + + const block_q2_K * GGML_RESTRICT x = vx; + const block_q8_K * GGML_RESTRICT y = vy; + + const int nb = n / QK_K; + +#if defined __AVX2__ + + const __m256i m3 = _mm256_set1_epi8(3); + const __m128i m4 = _mm_set1_epi8(0xF); + + __m256 acc = _mm256_setzero_ps(); + + for (int i = 0; i < nb; ++i) { + + const float d = y[i].d * GGML_CPU_FP16_TO_FP32(x[i].d); + const float dmin = -y[i].d * GGML_CPU_FP16_TO_FP32(x[i].dmin); + + const uint8_t * GGML_RESTRICT q2 = x[i].qs; + const int8_t * GGML_RESTRICT q8 = y[i].qs; + + const __m128i mins_and_scales = _mm_loadu_si128((const __m128i*)x[i].scales); + const __m128i scales8 = _mm_and_si128(mins_and_scales, m4); + const __m128i mins8 = _mm_and_si128(_mm_srli_epi16(mins_and_scales, 4), m4); + const __m256i mins = _mm256_cvtepi8_epi16(mins8); + const __m256i prod = _mm256_madd_epi16(mins, _mm256_loadu_si256((const __m256i*)y[i].bsums)); + + acc = _mm256_fmadd_ps(_mm256_broadcast_ss(&dmin), _mm256_cvtepi32_ps(prod), acc); + + const __m256i all_scales = _mm256_cvtepi8_epi16(scales8); + const __m128i l_scales = _mm256_extracti128_si256(all_scales, 0); + const __m128i h_scales = _mm256_extracti128_si256(all_scales, 1); + const __m256i scales[2] = {MM256_SET_M128I(l_scales, l_scales), MM256_SET_M128I(h_scales, h_scales)}; + + __m256i sumi = _mm256_setzero_si256(); + + for (int j = 0; j < QK_K/128; ++j) { + + const __m256i q2bits = _mm256_loadu_si256((const __m256i*)q2); q2 += 32; + + const __m256i q8_0 = _mm256_loadu_si256((const __m256i*)q8); q8 += 32; + const __m256i q8_1 = _mm256_loadu_si256((const __m256i*)q8); q8 += 32; + const __m256i q8_2 = _mm256_loadu_si256((const __m256i*)q8); q8 += 32; + const __m256i q8_3 = _mm256_loadu_si256((const __m256i*)q8); q8 += 32; + + const __m256i q2_0 = _mm256_and_si256(q2bits, m3); + const __m256i q2_1 = _mm256_and_si256(_mm256_srli_epi16(q2bits, 2), m3); + const __m256i q2_2 = _mm256_and_si256(_mm256_srli_epi16(q2bits, 4), m3); + const __m256i q2_3 = _mm256_and_si256(_mm256_srli_epi16(q2bits, 6), m3); + + __m256i p0 = _mm256_maddubs_epi16(q2_0, q8_0); + __m256i p1 = _mm256_maddubs_epi16(q2_1, q8_1); + __m256i p2 = _mm256_maddubs_epi16(q2_2, q8_2); + __m256i p3 = _mm256_maddubs_epi16(q2_3, q8_3); + + p0 = _mm256_madd_epi16(_mm256_shuffle_epi8(scales[j], get_scale_shuffle_q3k(0)), p0); + p1 = _mm256_madd_epi16(_mm256_shuffle_epi8(scales[j], get_scale_shuffle_q3k(1)), p1); + p2 = _mm256_madd_epi16(_mm256_shuffle_epi8(scales[j], get_scale_shuffle_q3k(2)), p2); + p3 = _mm256_madd_epi16(_mm256_shuffle_epi8(scales[j], get_scale_shuffle_q3k(3)), p3); + + p0 = _mm256_add_epi32(p0, p1); + p2 = _mm256_add_epi32(p2, p3); + + sumi = _mm256_add_epi32(sumi, _mm256_add_epi32(p0, p2)); + } + + acc = _mm256_fmadd_ps(_mm256_broadcast_ss(&d), _mm256_cvtepi32_ps(sumi), acc); + + } + + *s = hsum_float_8(acc); + +#elif defined __AVX__ + + const __m128i m3 = _mm_set1_epi8(0x3); + const __m128i m4 = _mm_set1_epi8(0xF); + const __m128i m2 = _mm_set1_epi8(0x2); + + __m256 acc = _mm256_setzero_ps(); + + for (int i = 0; i < nb; ++i) { + + const float dall = y[i].d * GGML_CPU_FP16_TO_FP32(x[i].d); + const float dmin = -y[i].d * GGML_CPU_FP16_TO_FP32(x[i].dmin); + + const uint8_t * GGML_RESTRICT q2 = x[i].qs; + const int8_t * GGML_RESTRICT q8 = y[i].qs; + + // load mins and scales from block_q2_K.scales[QK_K/16] + const __m128i mins_and_scales = _mm_loadu_si128((const __m128i*)x[i].scales); + const __m128i scales16 = _mm_and_si128(mins_and_scales, m4); + const __m128i mins16 = _mm_and_si128(_mm_srli_epi16(mins_and_scales, 4), m4); + const __m128i mins_0 = _mm_cvtepi8_epi16(mins16); + const __m128i mins_1 = _mm_cvtepi8_epi16(_mm_unpackhi_epi64(mins16, mins16)); + + // summs = y[i].bsums * (x[i].scales >> 4) in 16bits*8*2 to 32bits*4*2 + const __m128i summs_0 = _mm_madd_epi16(mins_0, _mm_loadu_si128((const __m128i*)&y[i].bsums[0])); + const __m128i summs_1 = _mm_madd_epi16(mins_1, _mm_loadu_si128((const __m128i*)&y[i].bsums[8])); + + // sumf += -dmin * summs in 32bits*8 + acc = _mm256_add_ps(_mm256_mul_ps(_mm256_broadcast_ss(&dmin), _mm256_cvtepi32_ps(MM256_SET_M128I(summs_1, summs_0))), acc); + + const __m128i scales_0 = _mm_cvtepi8_epi16(scales16); + const __m128i scales_1 = _mm_cvtepi8_epi16(_mm_unpackhi_epi64(scales16, scales16)); + const __m128i scales[2] = { scales_0, scales_1 }; + + __m128i sumi_0 = _mm_setzero_si128(); + __m128i sumi_1 = _mm_setzero_si128(); + + for (int j = 0; j < QK_K/128; ++j) { + + // load Q8 quants int8*16*8 from block_q8_K.qs[QK_K] + const __m128i q8_0 = _mm_loadu_si128((const __m128i*)q8); q8 += 16; + const __m128i q8_1 = _mm_loadu_si128((const __m128i*)q8); q8 += 16; + const __m128i q8_2 = _mm_loadu_si128((const __m128i*)q8); q8 += 16; + const __m128i q8_3 = _mm_loadu_si128((const __m128i*)q8); q8 += 16; + const __m128i q8_4 = _mm_loadu_si128((const __m128i*)q8); q8 += 16; + const __m128i q8_5 = _mm_loadu_si128((const __m128i*)q8); q8 += 16; + const __m128i q8_6 = _mm_loadu_si128((const __m128i*)q8); q8 += 16; + const __m128i q8_7 = _mm_loadu_si128((const __m128i*)q8); q8 += 16; + + // load 2bits*16*8 from block_q2_K.qs[QK_K/4] + __m128i q2bits = _mm_loadu_si128((const __m128i*)q2); q2 += 16; + const __m128i q2_0 = _mm_and_si128(q2bits, m3); + const __m128i q2_2 = _mm_and_si128(_mm_srli_epi16(q2bits, 2), m3); + const __m128i q2_4 = _mm_and_si128(_mm_srli_epi16(q2bits, 4), m3); + const __m128i q2_6 = _mm_and_si128(_mm_srli_epi16(q2bits, 6), m3); + q2bits = _mm_loadu_si128((const __m128i*)q2); q2 += 16; + const __m128i q2_1 = _mm_and_si128(q2bits, m3); + const __m128i q2_3 = _mm_and_si128(_mm_srli_epi16(q2bits, 2), m3); + const __m128i q2_5 = _mm_and_si128(_mm_srli_epi16(q2bits, 4), m3); + const __m128i q2_7 = _mm_and_si128(_mm_srli_epi16(q2bits, 6), m3); + + // isuml = q8[l] * ((q2[l] >> shift) & 3) in 8bits*16*8 to 16bits*8*8 + __m128i p0 = _mm_maddubs_epi16(q2_0, q8_0); + __m128i p1 = _mm_maddubs_epi16(q2_1, q8_1); + __m128i p2 = _mm_maddubs_epi16(q2_2, q8_2); + __m128i p3 = _mm_maddubs_epi16(q2_3, q8_3); + __m128i p4 = _mm_maddubs_epi16(q2_4, q8_4); + __m128i p5 = _mm_maddubs_epi16(q2_5, q8_5); + __m128i p6 = _mm_maddubs_epi16(q2_6, q8_6); + __m128i p7 = _mm_maddubs_epi16(q2_7, q8_7); + + // isum += (x[i].scales[is++] & 0xF) * isuml in 16bits*8*8 to 32bits*4*8 + __m128i shuffle = _mm_set1_epi16(0x0100); + p0 = _mm_madd_epi16(_mm_shuffle_epi8(scales[j], shuffle), p0); + shuffle = _mm_add_epi16(shuffle, m2); + p1 = _mm_madd_epi16(_mm_shuffle_epi8(scales[j], shuffle), p1); + shuffle = _mm_add_epi16(shuffle, m2); + p2 = _mm_madd_epi16(_mm_shuffle_epi8(scales[j], shuffle), p2); + shuffle = _mm_add_epi16(shuffle, m2); + p3 = _mm_madd_epi16(_mm_shuffle_epi8(scales[j], shuffle), p3); + shuffle = _mm_add_epi16(shuffle, m2); + p4 = _mm_madd_epi16(_mm_shuffle_epi8(scales[j], shuffle), p4); + shuffle = _mm_add_epi16(shuffle, m2); + p5 = _mm_madd_epi16(_mm_shuffle_epi8(scales[j], shuffle), p5); + shuffle = _mm_add_epi16(shuffle, m2); + p6 = _mm_madd_epi16(_mm_shuffle_epi8(scales[j], shuffle), p6); + shuffle = _mm_add_epi16(shuffle, m2); + p7 = _mm_madd_epi16(_mm_shuffle_epi8(scales[j], shuffle), p7); + + p0 = _mm_add_epi32(p0, p1); + p2 = _mm_add_epi32(p2, p3); + p4 = _mm_add_epi32(p4, p5); + p6 = _mm_add_epi32(p6, p7); + + // isum in 32bits*4*2 + sumi_0 = _mm_add_epi32(sumi_0, _mm_add_epi32(p0, p2)); + sumi_1 = _mm_add_epi32(sumi_1, _mm_add_epi32(p4, p6)); + } + + // sumf += dall * isum - dmin * summs in 32bits + __m256i sumi = MM256_SET_M128I(sumi_1, sumi_0); + acc = _mm256_add_ps(_mm256_mul_ps(_mm256_broadcast_ss(&dall), _mm256_cvtepi32_ps(sumi)), acc); + } + + *s = hsum_float_8(acc); + +#else + UNUSED(x); + UNUSED(y); + UNUSED(nb); + ggml_vec_dot_q2_K_q8_K_generic(n, s, bs, vx, bx, vy, by, nrc); +#endif +} + +void ggml_vec_dot_q3_K_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) { + assert(n % QK_K == 0); + assert(nrc == 1); + UNUSED(nrc); + UNUSED(bx); + UNUSED(by); + UNUSED(bs); + + const uint32_t kmask1 = 0x03030303; + const uint32_t kmask2 = 0x0f0f0f0f; + + const block_q3_K * GGML_RESTRICT x = vx; + const block_q8_K * GGML_RESTRICT y = vy; + + const int nb = n / QK_K; + +#if defined __AVX2__ + + const __m256i m3 = _mm256_set1_epi8(3); + const __m256i mone = _mm256_set1_epi8(1); + const __m128i m32 = _mm_set1_epi8(32); + + __m256 acc = _mm256_setzero_ps(); + + uint32_t aux[3]; + + for (int i = 0; i < nb; ++i) { + + const float d = y[i].d * GGML_CPU_FP16_TO_FP32(x[i].d); + + const uint8_t * GGML_RESTRICT q3 = x[i].qs; + const int8_t * GGML_RESTRICT q8 = y[i].qs; + + // Set up scales + memcpy(aux, x[i].scales, 12); + __m128i scales128 = _mm_set_epi32( + ((aux[1] >> 4) & kmask2) | (((aux[2] >> 6) & kmask1) << 4), + ((aux[0] >> 4) & kmask2) | (((aux[2] >> 4) & kmask1) << 4), + (aux[1] & kmask2) | (((aux[2] >> 2) & kmask1) << 4), + (aux[0] & kmask2) | (((aux[2] >> 0) & kmask1) << 4)); + scales128 = _mm_sub_epi8(scales128, m32); + const __m256i all_scales = _mm256_cvtepi8_epi16(scales128); + const __m128i l_scales = _mm256_extracti128_si256(all_scales, 0); + const __m128i h_scales = _mm256_extracti128_si256(all_scales, 1); + const __m256i scales[2] = {MM256_SET_M128I(l_scales, l_scales), MM256_SET_M128I(h_scales, h_scales)}; + + // high bit + const __m256i hbits = _mm256_loadu_si256((const __m256i*)x[i].hmask); + + // integer accumulator + __m256i sumi = _mm256_setzero_si256(); + + int bit = 0; + int is = 0; + + for (int j = 0; j < QK_K/128; ++j) { + // load low 2 bits + const __m256i q3bits = _mm256_loadu_si256((const __m256i*)q3); q3 += 32; + + // prepare low and high bits + const __m256i q3l_0 = _mm256_and_si256(q3bits, m3); + const __m256i q3h_0 = _mm256_slli_epi16(_mm256_srli_epi16(_mm256_andnot_si256(hbits, _mm256_slli_epi16(mone, bit)), bit), 2); + ++bit; + + const __m256i q3l_1 = _mm256_and_si256(_mm256_srli_epi16(q3bits, 2), m3); + const __m256i q3h_1 = _mm256_slli_epi16(_mm256_srli_epi16(_mm256_andnot_si256(hbits, _mm256_slli_epi16(mone, bit)), bit), 2); + ++bit; + + const __m256i q3l_2 = _mm256_and_si256(_mm256_srli_epi16(q3bits, 4), m3); + const __m256i q3h_2 = _mm256_slli_epi16(_mm256_srli_epi16(_mm256_andnot_si256(hbits, _mm256_slli_epi16(mone, bit)), bit), 2); + ++bit; + + const __m256i q3l_3 = _mm256_and_si256(_mm256_srli_epi16(q3bits, 6), m3); + const __m256i q3h_3 = _mm256_slli_epi16(_mm256_srli_epi16(_mm256_andnot_si256(hbits, _mm256_slli_epi16(mone, bit)), bit), 2); + ++bit; + + // load Q8 quants + const __m256i q8_0 = _mm256_loadu_si256((const __m256i*)q8); q8 += 32; + const __m256i q8_1 = _mm256_loadu_si256((const __m256i*)q8); q8 += 32; + const __m256i q8_2 = _mm256_loadu_si256((const __m256i*)q8); q8 += 32; + const __m256i q8_3 = _mm256_loadu_si256((const __m256i*)q8); q8 += 32; + + // Dot product: we multiply the 2 low bits and 1 high bit part separately, so we can use _mm256_maddubs_epi16, + // and then subtract. The high bit part has the 2 already subtracted (and so, it is zero if the high bit was not set, + // and 2 if the high bit was set) + __m256i q8s_0 = _mm256_maddubs_epi16(q3h_0, q8_0); + __m256i q8s_1 = _mm256_maddubs_epi16(q3h_1, q8_1); + __m256i q8s_2 = _mm256_maddubs_epi16(q3h_2, q8_2); + __m256i q8s_3 = _mm256_maddubs_epi16(q3h_3, q8_3); + + __m256i p16_0 = _mm256_maddubs_epi16(q3l_0, q8_0); + __m256i p16_1 = _mm256_maddubs_epi16(q3l_1, q8_1); + __m256i p16_2 = _mm256_maddubs_epi16(q3l_2, q8_2); + __m256i p16_3 = _mm256_maddubs_epi16(q3l_3, q8_3); + + p16_0 = _mm256_sub_epi16(p16_0, q8s_0); + p16_1 = _mm256_sub_epi16(p16_1, q8s_1); + p16_2 = _mm256_sub_epi16(p16_2, q8s_2); + p16_3 = _mm256_sub_epi16(p16_3, q8s_3); + + // multiply with scales + p16_0 = _mm256_madd_epi16(_mm256_shuffle_epi8(scales[j], get_scale_shuffle_q3k(is + 0)), p16_0); + p16_1 = _mm256_madd_epi16(_mm256_shuffle_epi8(scales[j], get_scale_shuffle_q3k(is + 1)), p16_1); + p16_2 = _mm256_madd_epi16(_mm256_shuffle_epi8(scales[j], get_scale_shuffle_q3k(is + 2)), p16_2); + p16_3 = _mm256_madd_epi16(_mm256_shuffle_epi8(scales[j], get_scale_shuffle_q3k(is + 3)), p16_3); + + // accumulate + p16_0 = _mm256_add_epi32(p16_0, p16_1); + p16_2 = _mm256_add_epi32(p16_2, p16_3); + sumi = _mm256_add_epi32(sumi, _mm256_add_epi32(p16_0, p16_2)); + + } + + // multiply with block scale and accumulate + acc = _mm256_fmadd_ps(_mm256_broadcast_ss(&d), _mm256_cvtepi32_ps(sumi), acc); + + } + + *s = hsum_float_8(acc); + +#elif defined __AVX__ + + const __m128i m3 = _mm_set1_epi8(3); + const __m128i mone = _mm_set1_epi8(1); + const __m128i m32 = _mm_set1_epi8(32); + const __m128i m2 = _mm_set1_epi8(2); + + __m256 acc = _mm256_setzero_ps(); + + const uint32_t *aux; + + for (int i = 0; i < nb; ++i) { + + const float d = y[i].d * GGML_CPU_FP16_TO_FP32(x[i].d); + + const uint8_t * GGML_RESTRICT q3 = x[i].qs; + const int8_t * GGML_RESTRICT q8 = y[i].qs; + + // Set up scales + aux = (const uint32_t *)x[i].scales; + __m128i scales128 = _mm_set_epi32( + ((aux[1] >> 4) & kmask2) | (((aux[2] >> 6) & kmask1) << 4), + ((aux[0] >> 4) & kmask2) | (((aux[2] >> 4) & kmask1) << 4), + (aux[1] & kmask2) | (((aux[2] >> 2) & kmask1) << 4), + (aux[0] & kmask2) | (((aux[2] >> 0) & kmask1) << 4)); + scales128 = _mm_sub_epi8(scales128, m32); + const __m128i scales_0 = _mm_cvtepi8_epi16(scales128); + const __m128i scales_1 = _mm_cvtepi8_epi16(_mm_unpackhi_epi64(scales128, scales128)); + const __m128i scales[2] = { scales_0, scales_1 }; + + // high bit *128*2 from block_q3_K.hmask[QK_K/8] + const __m128i hbits_0 = _mm_loadu_si128((const __m128i*)&x[i].hmask[0]); + const __m128i hbits_1 = _mm_loadu_si128((const __m128i*)&x[i].hmask[16]); + + // integer accumulator + __m128i sumi_0 = _mm_setzero_si128(); + __m128i sumi_1 = _mm_setzero_si128(); + + for (int j = 0; j < QK_K/128; ++j) { + // load low 2 bits *64*2 from block_q3_K.qs[QK_K/4] + const __m128i q3bits_0 = _mm_loadu_si128((const __m128i*)q3); q3 += 16; + const __m128i q3bits_1 = _mm_loadu_si128((const __m128i*)q3); q3 += 16; + + // prepare low and high bits + const int bit = j << 2; + + const __m128i q3l_0 = _mm_and_si128(q3bits_0, m3); + const __m128i q3l_1 = _mm_and_si128(q3bits_1, m3); + const __m128i q3h_0 = _mm_slli_epi16(_mm_srli_epi16(_mm_andnot_si128(hbits_0, _mm_slli_epi16(mone, bit)), bit), 2); + const __m128i q3h_1 = _mm_slli_epi16(_mm_srli_epi16(_mm_andnot_si128(hbits_1, _mm_slli_epi16(mone, bit)), bit), 2); + + const __m128i q3l_2 = _mm_and_si128(_mm_srli_epi16(q3bits_0, 2), m3); + const __m128i q3l_3 = _mm_and_si128(_mm_srli_epi16(q3bits_1, 2), m3); + const __m128i q3h_2 = _mm_slli_epi16(_mm_srli_epi16(_mm_andnot_si128(hbits_0, _mm_slli_epi16(mone, bit+1)), bit+1), 2); + const __m128i q3h_3 = _mm_slli_epi16(_mm_srli_epi16(_mm_andnot_si128(hbits_1, _mm_slli_epi16(mone, bit+1)), bit+1), 2); + + const __m128i q3l_4 = _mm_and_si128(_mm_srli_epi16(q3bits_0, 4), m3); + const __m128i q3l_5 = _mm_and_si128(_mm_srli_epi16(q3bits_1, 4), m3); + const __m128i q3h_4 = _mm_slli_epi16(_mm_srli_epi16(_mm_andnot_si128(hbits_0, _mm_slli_epi16(mone, bit+2)), bit+2), 2); + const __m128i q3h_5 = _mm_slli_epi16(_mm_srli_epi16(_mm_andnot_si128(hbits_1, _mm_slli_epi16(mone, bit+2)), bit+2), 2); + + const __m128i q3l_6 = _mm_and_si128(_mm_srli_epi16(q3bits_0, 6), m3); + const __m128i q3l_7 = _mm_and_si128(_mm_srli_epi16(q3bits_1, 6), m3); + const __m128i q3h_6 = _mm_slli_epi16(_mm_srli_epi16(_mm_andnot_si128(hbits_0, _mm_slli_epi16(mone, bit+3)), bit+3), 2); + const __m128i q3h_7 = _mm_slli_epi16(_mm_srli_epi16(_mm_andnot_si128(hbits_1, _mm_slli_epi16(mone, bit+3)), bit+3), 2); + + // load Q8 quants from block_q8_K.qs[QK_K] + const __m128i q8_0 = _mm_loadu_si128((const __m128i*)q8); q8 += 16; + const __m128i q8_1 = _mm_loadu_si128((const __m128i*)q8); q8 += 16; + const __m128i q8_2 = _mm_loadu_si128((const __m128i*)q8); q8 += 16; + const __m128i q8_3 = _mm_loadu_si128((const __m128i*)q8); q8 += 16; + const __m128i q8_4 = _mm_loadu_si128((const __m128i*)q8); q8 += 16; + const __m128i q8_5 = _mm_loadu_si128((const __m128i*)q8); q8 += 16; + const __m128i q8_6 = _mm_loadu_si128((const __m128i*)q8); q8 += 16; + const __m128i q8_7 = _mm_loadu_si128((const __m128i*)q8); q8 += 16; + + // Dot product: we multiply the 2 low bits and 1 high bit part separately, so we can use _mm256_maddubs_epi16, + // and then subtract. The high bit part has the 2 already subtracted (and so, it is zero if the high bit was not set, + // and 2 if the high bit was set) + __m128i q8s_0 = _mm_maddubs_epi16(q3h_0, q8_0); + __m128i q8s_1 = _mm_maddubs_epi16(q3h_1, q8_1); + __m128i q8s_2 = _mm_maddubs_epi16(q3h_2, q8_2); + __m128i q8s_3 = _mm_maddubs_epi16(q3h_3, q8_3); + __m128i q8s_4 = _mm_maddubs_epi16(q3h_4, q8_4); + __m128i q8s_5 = _mm_maddubs_epi16(q3h_5, q8_5); + __m128i q8s_6 = _mm_maddubs_epi16(q3h_6, q8_6); + __m128i q8s_7 = _mm_maddubs_epi16(q3h_7, q8_7); + + __m128i p16_0 = _mm_maddubs_epi16(q3l_0, q8_0); + __m128i p16_1 = _mm_maddubs_epi16(q3l_1, q8_1); + __m128i p16_2 = _mm_maddubs_epi16(q3l_2, q8_2); + __m128i p16_3 = _mm_maddubs_epi16(q3l_3, q8_3); + __m128i p16_4 = _mm_maddubs_epi16(q3l_4, q8_4); + __m128i p16_5 = _mm_maddubs_epi16(q3l_5, q8_5); + __m128i p16_6 = _mm_maddubs_epi16(q3l_6, q8_6); + __m128i p16_7 = _mm_maddubs_epi16(q3l_7, q8_7); + + p16_0 = _mm_sub_epi16(p16_0, q8s_0); + p16_1 = _mm_sub_epi16(p16_1, q8s_1); + p16_2 = _mm_sub_epi16(p16_2, q8s_2); + p16_3 = _mm_sub_epi16(p16_3, q8s_3); + p16_4 = _mm_sub_epi16(p16_4, q8s_4); + p16_5 = _mm_sub_epi16(p16_5, q8s_5); + p16_6 = _mm_sub_epi16(p16_6, q8s_6); + p16_7 = _mm_sub_epi16(p16_7, q8s_7); + + // multiply with scales + __m128i shuffle = _mm_set1_epi16(0x0100); + p16_0 = _mm_madd_epi16(_mm_shuffle_epi8(scales[j], shuffle), p16_0); + shuffle = _mm_add_epi16(shuffle, m2); + p16_1 = _mm_madd_epi16(_mm_shuffle_epi8(scales[j], shuffle), p16_1); + shuffle = _mm_add_epi16(shuffle, m2); + p16_2 = _mm_madd_epi16(_mm_shuffle_epi8(scales[j], shuffle), p16_2); + shuffle = _mm_add_epi16(shuffle, m2); + p16_3 = _mm_madd_epi16(_mm_shuffle_epi8(scales[j], shuffle), p16_3); + shuffle = _mm_add_epi16(shuffle, m2); + p16_4 = _mm_madd_epi16(_mm_shuffle_epi8(scales[j], shuffle), p16_4); + shuffle = _mm_add_epi16(shuffle, m2); + p16_5 = _mm_madd_epi16(_mm_shuffle_epi8(scales[j], shuffle), p16_5); + shuffle = _mm_add_epi16(shuffle, m2); + p16_6 = _mm_madd_epi16(_mm_shuffle_epi8(scales[j], shuffle), p16_6); + shuffle = _mm_add_epi16(shuffle, m2); + p16_7 = _mm_madd_epi16(_mm_shuffle_epi8(scales[j], shuffle), p16_7); + + // accumulate + p16_0 = _mm_add_epi32(p16_0, p16_1); + p16_2 = _mm_add_epi32(p16_2, p16_3); + p16_4 = _mm_add_epi32(p16_4, p16_5); + p16_6 = _mm_add_epi32(p16_6, p16_7); + sumi_0 = _mm_add_epi32(sumi_0, _mm_add_epi32(p16_0, p16_2)); + sumi_1 = _mm_add_epi32(sumi_1, _mm_add_epi32(p16_4, p16_6)); + + } + + // multiply with block scale and accumulate + __m256i sumi = MM256_SET_M128I(sumi_1, sumi_0); + acc = _mm256_add_ps(_mm256_mul_ps(_mm256_broadcast_ss(&d), _mm256_cvtepi32_ps(sumi)), acc); + + } + + *s = hsum_float_8(acc); + +#else + UNUSED(kmask1); + UNUSED(kmask2); + UNUSED(x); + UNUSED(y); + UNUSED(nb); + ggml_vec_dot_q3_K_q8_K_generic(n, s, bs, vx, bx, vy, by, nrc); +#endif +} + +void ggml_vec_dot_q4_K_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) { + assert(n % QK_K == 0); + assert(nrc == 1); + UNUSED(nrc); + UNUSED(bx); + UNUSED(by); + UNUSED(bs); + + const block_q4_K * GGML_RESTRICT x = vx; + const block_q8_K * GGML_RESTRICT y = vy; + + const int nb = n / QK_K; + + static const uint32_t kmask1 = 0x3f3f3f3f; + static const uint32_t kmask2 = 0x0f0f0f0f; + static const uint32_t kmask3 = 0x03030303; + + uint32_t utmp[4]; + +#if defined __AVX2__ + + const __m256i m4 = _mm256_set1_epi8(0xF); + + __m256 acc = _mm256_setzero_ps(); + __m128 acc_m = _mm_setzero_ps(); + + for (int i = 0; i < nb; ++i) { + + const float d = y[i].d * GGML_CPU_FP16_TO_FP32(x[i].d); + const float dmin = -y[i].d * GGML_CPU_FP16_TO_FP32(x[i].dmin); + + memcpy(utmp, x[i].scales, 12); + utmp[3] = ((utmp[2] >> 4) & kmask2) | (((utmp[1] >> 6) & kmask3) << 4); + const uint32_t uaux = utmp[1] & kmask1; + utmp[1] = (utmp[2] & kmask2) | (((utmp[0] >> 6) & kmask3) << 4); + utmp[2] = uaux; + utmp[0] &= kmask1; + + const uint8_t * GGML_RESTRICT q4 = x[i].qs; + const int8_t * GGML_RESTRICT q8 = y[i].qs; + + const __m256i mins_and_scales = _mm256_cvtepu8_epi16(_mm_set_epi32(utmp[3], utmp[2], utmp[1], utmp[0])); + + const __m256i q8sums = _mm256_loadu_si256((const __m256i*)y[i].bsums); + const __m128i q8s = _mm_hadd_epi16(_mm256_extracti128_si256(q8sums, 0), _mm256_extracti128_si256(q8sums, 1)); + const __m128i prod = _mm_madd_epi16(_mm256_extracti128_si256(mins_and_scales, 1), q8s); + acc_m = _mm_fmadd_ps(_mm_set1_ps(dmin), _mm_cvtepi32_ps(prod), acc_m); + + const __m128i sc128 = _mm256_extracti128_si256(mins_and_scales, 0); + const __m256i scales = MM256_SET_M128I(sc128, sc128); + + __m256i sumi = _mm256_setzero_si256(); + + for (int j = 0; j < QK_K/64; ++j) { + + const __m256i scale_l = _mm256_shuffle_epi8(scales, get_scale_shuffle_k4(2*j+0)); + const __m256i scale_h = _mm256_shuffle_epi8(scales, get_scale_shuffle_k4(2*j+1)); + + const __m256i q4bits = _mm256_loadu_si256((const __m256i*)q4); q4 += 32; + const __m256i q4l = _mm256_and_si256(q4bits, m4); + const __m256i q4h = _mm256_and_si256(_mm256_srli_epi16(q4bits, 4), m4); + + const __m256i q8l = _mm256_loadu_si256((const __m256i*)q8); q8 += 32; + __m256i p16l = _mm256_maddubs_epi16(q4l, q8l); + p16l = _mm256_madd_epi16(scale_l, p16l); + + const __m256i q8h = _mm256_loadu_si256((const __m256i*)q8); q8 += 32; + __m256i p16h = _mm256_maddubs_epi16(q4h, q8h); + p16h = _mm256_madd_epi16(scale_h, p16h); + const __m256i sumj = _mm256_add_epi32(p16l, p16h); + + sumi = _mm256_add_epi32(sumi, sumj); + } + + __m256 vd = _mm256_set1_ps(d); + acc = _mm256_fmadd_ps(vd, _mm256_cvtepi32_ps(sumi), acc); + + } + + acc_m = _mm_add_ps(acc_m, _mm_movehl_ps(acc_m, acc_m)); + acc_m = _mm_add_ss(acc_m, _mm_movehdup_ps(acc_m)); + + *s = hsum_float_8(acc) + _mm_cvtss_f32(acc_m); + +#elif defined __AVX__ + + const __m128i m4 = _mm_set1_epi8(0xF); + const __m128i m2 = _mm_set1_epi8(0x2); + + __m256 acc = _mm256_setzero_ps(); + __m128 acc_m = _mm_setzero_ps(); + + for (int i = 0; i < nb; ++i) { + + const float d = y[i].d * GGML_CPU_FP16_TO_FP32(x[i].d); + const float dmin = -y[i].d * GGML_CPU_FP16_TO_FP32(x[i].dmin); + + const uint8_t * GGML_RESTRICT q4 = x[i].qs; + const int8_t * GGML_RESTRICT q8 = y[i].qs; + + memcpy(utmp, x[i].scales, 12); + utmp[3] = ((utmp[2] >> 4) & kmask2) | (((utmp[1] >> 6) & kmask3) << 4); + const uint32_t uaux = utmp[1] & kmask1; + utmp[1] = (utmp[2] & kmask2) | (((utmp[0] >> 6) & kmask3) << 4); + utmp[2] = uaux; + utmp[0] &= kmask1; + + const __m128i utmps = _mm_set_epi32(utmp[3], utmp[2], utmp[1], utmp[0]); + const __m128i scales = _mm_cvtepu8_epi16(utmps); + const __m128i mins = _mm_cvtepu8_epi16(_mm_unpackhi_epi64(utmps, utmps)); + + const __m128i q8sums_0 = _mm_loadu_si128((const __m128i*)&y[i].bsums[0]); + const __m128i q8sums_1 = _mm_loadu_si128((const __m128i*)&y[i].bsums[8]); + const __m128i q8s = _mm_hadd_epi16(q8sums_0, q8sums_1); + const __m128i prod = _mm_madd_epi16(mins, q8s); + acc_m = _mm_add_ps(_mm_mul_ps(_mm_set1_ps(dmin), _mm_cvtepi32_ps(prod)), acc_m); + + __m128i sumi_0 = _mm_setzero_si128(); + __m128i sumi_1 = _mm_setzero_si128(); + + __m128i shuffle = _mm_set1_epi16(0x0100); + for (int j = 0; j < QK_K/64; ++j) { + + const __m128i scale_l = _mm_shuffle_epi8(scales, shuffle); + shuffle = _mm_add_epi16(shuffle, m2); + const __m128i scale_h = _mm_shuffle_epi8(scales, shuffle); + shuffle = _mm_add_epi16(shuffle, m2); + + __m128i q4bits = _mm_loadu_si128((const __m128i*)q4); q4 += 16; + const __m128i q4l_0 = _mm_and_si128(q4bits, m4); + const __m128i q4h_0 = _mm_and_si128(_mm_srli_epi16(q4bits, 4), m4); + q4bits = _mm_loadu_si128((const __m128i*)q4); q4 += 16; + const __m128i q4l_1 = _mm_and_si128(q4bits, m4); + const __m128i q4h_1 = _mm_and_si128(_mm_srli_epi16(q4bits, 4), m4); + + const __m128i q8l_0 = _mm_loadu_si128((const __m128i*)q8); q8 += 16; + __m128i p16l = _mm_maddubs_epi16(q4l_0, q8l_0); + p16l = _mm_madd_epi16(scale_l, p16l); + sumi_0 = _mm_add_epi32(sumi_0, p16l); + const __m128i q8l_1 = _mm_loadu_si128((const __m128i*)q8); q8 += 16; + p16l = _mm_maddubs_epi16(q4l_1, q8l_1); + p16l = _mm_madd_epi16(scale_l, p16l); + sumi_1 = _mm_add_epi32(sumi_1, p16l); + + const __m128i q8h_0 = _mm_loadu_si128((const __m128i*)q8); q8 += 16; + __m128i p16h = _mm_maddubs_epi16(q4h_0, q8h_0); + p16h = _mm_madd_epi16(scale_h, p16h); + sumi_0 = _mm_add_epi32(sumi_0, p16h); + const __m128i q8h_1 = _mm_loadu_si128((const __m128i*)q8); q8 += 16; + p16h = _mm_maddubs_epi16(q4h_1, q8h_1); + p16h = _mm_madd_epi16(scale_h, p16h); + sumi_1 = _mm_add_epi32(sumi_1, p16h); + + } + + __m256 vd = _mm256_set1_ps(d); + __m256i sumi = MM256_SET_M128I(sumi_1, sumi_0); + acc = _mm256_add_ps(_mm256_mul_ps(vd, _mm256_cvtepi32_ps(sumi)), acc); + + } + + acc_m = _mm_add_ps(acc_m, _mm_movehl_ps(acc_m, acc_m)); + acc_m = _mm_add_ss(acc_m, _mm_movehdup_ps(acc_m)); + + *s = hsum_float_8(acc) + _mm_cvtss_f32(acc_m); + +#else + UNUSED(x); + UNUSED(y); + UNUSED(nb); + UNUSED(kmask1); + UNUSED(kmask2); + UNUSED(kmask3); + UNUSED(utmp); + ggml_vec_dot_q4_K_q8_K_generic(n, s, bs, vx, bx, vy, by, nrc); +#endif +} + +void ggml_vec_dot_q5_K_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) { + assert(n % QK_K == 0); + assert(nrc == 1); + UNUSED(nrc); + UNUSED(bx); + UNUSED(by); + UNUSED(bs); + + const block_q5_K * GGML_RESTRICT x = vx; + const block_q8_K * GGML_RESTRICT y = vy; + + const int nb = n / QK_K; + + static const uint32_t kmask1 = 0x3f3f3f3f; + static const uint32_t kmask2 = 0x0f0f0f0f; + static const uint32_t kmask3 = 0x03030303; + + uint32_t utmp[4]; + +#if defined __AVX2__ + + const __m256i m4 = _mm256_set1_epi8(0xF); + const __m128i mzero = _mm_setzero_si128(); + const __m256i mone = _mm256_set1_epi8(1); + + __m256 acc = _mm256_setzero_ps(); + + float summs = 0.f; + + for (int i = 0; i < nb; ++i) { + const uint8_t * GGML_RESTRICT q5 = x[i].qs; + const int8_t * GGML_RESTRICT q8 = y[i].qs; + + const float d = y[i].d * GGML_CPU_FP16_TO_FP32(x[i].d); + const float dmin = -y[i].d * GGML_CPU_FP16_TO_FP32(x[i].dmin); + + memcpy(utmp, x[i].scales, 12); + utmp[3] = ((utmp[2] >> 4) & kmask2) | (((utmp[1] >> 6) & kmask3) << 4); + const uint32_t uaux = utmp[1] & kmask1; + utmp[1] = (utmp[2] & kmask2) | (((utmp[0] >> 6) & kmask3) << 4); + utmp[2] = uaux; + utmp[0] &= kmask1; + + const __m256i mins_and_scales = _mm256_cvtepu8_epi16(_mm_set_epi32(utmp[3], utmp[2], utmp[1], utmp[0])); + + const __m256i q8sums = _mm256_loadu_si256((const __m256i*)y[i].bsums); + const __m128i q8s = _mm_hadd_epi16(_mm256_extracti128_si256(q8sums, 0), _mm256_extracti128_si256(q8sums, 1)); + const __m128i prod = _mm_madd_epi16(_mm256_extracti128_si256(mins_and_scales, 1), q8s); + const __m128i hsum = _mm_hadd_epi32(_mm_hadd_epi32(prod, mzero), mzero); + summs += dmin * _mm_extract_epi32(hsum, 0); + + const __m128i sc128 = _mm256_extracti128_si256(mins_and_scales, 0); + const __m256i scales = MM256_SET_M128I(sc128, sc128); + + const __m256i hbits = _mm256_loadu_si256((const __m256i*)x[i].qh); + __m256i hmask = mone; + + __m256i sumi = _mm256_setzero_si256(); + + int bit = 0; + + for (int j = 0; j < QK_K/64; ++j) { + + const __m256i scale_0 = _mm256_shuffle_epi8(scales, get_scale_shuffle_k4(2*j+0)); + const __m256i scale_1 = _mm256_shuffle_epi8(scales, get_scale_shuffle_k4(2*j+1)); + + const __m256i q5bits = _mm256_loadu_si256((const __m256i*)q5); q5 += 32; + + const __m256i q5l_0 = _mm256_and_si256(q5bits, m4); + const __m256i q5h_0 = _mm256_slli_epi16(_mm256_srli_epi16(_mm256_and_si256(hbits, hmask), bit++), 4); + const __m256i q5_0 = _mm256_add_epi8(q5l_0, q5h_0); + hmask = _mm256_slli_epi16(hmask, 1); + + const __m256i q5l_1 = _mm256_and_si256(_mm256_srli_epi16(q5bits, 4), m4); + const __m256i q5h_1 = _mm256_slli_epi16(_mm256_srli_epi16(_mm256_and_si256(hbits, hmask), bit++), 4); + const __m256i q5_1 = _mm256_add_epi8(q5l_1, q5h_1); + hmask = _mm256_slli_epi16(hmask, 1); + + const __m256i q8_0 = _mm256_loadu_si256((const __m256i*)q8); q8 += 32; + const __m256i q8_1 = _mm256_loadu_si256((const __m256i*)q8); q8 += 32; + + __m256i p16_0 = _mm256_maddubs_epi16(q5_0, q8_0); + __m256i p16_1 = _mm256_maddubs_epi16(q5_1, q8_1); + + p16_0 = _mm256_madd_epi16(scale_0, p16_0); + p16_1 = _mm256_madd_epi16(scale_1, p16_1); + + sumi = _mm256_add_epi32(sumi, _mm256_add_epi32(p16_0, p16_1)); + + } + + __m256 vd = _mm256_set1_ps(d); + acc = _mm256_fmadd_ps(vd, _mm256_cvtepi32_ps(sumi), acc); + + } + + *s = hsum_float_8(acc) + summs; + +#elif defined __AVX__ + + const __m128i m4 = _mm_set1_epi8(0xF); + const __m128i mzero = _mm_setzero_si128(); + const __m128i mone = _mm_set1_epi8(1); + const __m128i m2 = _mm_set1_epi8(2); + + __m256 acc = _mm256_setzero_ps(); + + float summs = 0.f; + + for (int i = 0; i < nb; ++i) { + + const float d = y[i].d * GGML_CPU_FP16_TO_FP32(x[i].d); + const float dmin = -y[i].d * GGML_CPU_FP16_TO_FP32(x[i].dmin); + + const uint8_t * GGML_RESTRICT q5 = x[i].qs; + const int8_t * GGML_RESTRICT q8 = y[i].qs; + + memcpy(utmp, x[i].scales, 12); + utmp[3] = ((utmp[2] >> 4) & kmask2) | (((utmp[1] >> 6) & kmask3) << 4); + const uint32_t uaux = utmp[1] & kmask1; + utmp[1] = (utmp[2] & kmask2) | (((utmp[0] >> 6) & kmask3) << 4); + utmp[2] = uaux; + utmp[0] &= kmask1; + + const __m128i utmps = _mm_set_epi32(utmp[3], utmp[2], utmp[1], utmp[0]); + const __m128i scales = _mm_cvtepu8_epi16(utmps); + const __m128i mins = _mm_cvtepu8_epi16(_mm_unpackhi_epi64(utmps, utmps)); + + const __m128i q8sums_0 = _mm_loadu_si128((const __m128i*)&y[i].bsums[0]); + const __m128i q8sums_1 = _mm_loadu_si128((const __m128i*)&y[i].bsums[8]); + const __m128i q8s = _mm_hadd_epi16(q8sums_0, q8sums_1); + const __m128i prod = _mm_madd_epi16(mins, q8s); + const __m128i hsum = _mm_hadd_epi32(_mm_hadd_epi32(prod, mzero), mzero); + summs += dmin * _mm_extract_epi32(hsum, 0); + + const __m128i hbits_0 = _mm_loadu_si128((const __m128i*)&x[i].qh[0]); + const __m128i hbits_1 = _mm_loadu_si128((const __m128i*)&x[i].qh[16]); + __m128i hmask = mone; + + __m128i sumi_0 = _mm_setzero_si128(); + __m128i sumi_1 = _mm_setzero_si128(); + + int bit = 0; + + __m128i shuffle = _mm_set1_epi16(0x0100); + for (int j = 0; j < QK_K/64; ++j) { + + const __m128i scale_0 = _mm_shuffle_epi8(scales, shuffle); + shuffle = _mm_add_epi16(shuffle, m2); + const __m128i scale_1 = _mm_shuffle_epi8(scales, shuffle); + shuffle = _mm_add_epi16(shuffle, m2); + + const __m128i q5bits_0 = _mm_loadu_si128((const __m128i*)q5); q5 += 16; + const __m128i q5bits_1 = _mm_loadu_si128((const __m128i*)q5); q5 += 16; + + __m128i q5l_0 = _mm_and_si128(q5bits_0, m4); + __m128i q5l_1 = _mm_and_si128(q5bits_1, m4); + __m128i q5h_0 = _mm_slli_epi16(_mm_srli_epi16(_mm_and_si128(hbits_0, hmask), bit), 4); + __m128i q5h_1 = _mm_slli_epi16(_mm_srli_epi16(_mm_and_si128(hbits_1, hmask), bit++), 4); + __m128i q5_0 = _mm_add_epi8(q5l_0, q5h_0); + __m128i q5_1 = _mm_add_epi8(q5l_1, q5h_1); + hmask = _mm_slli_epi16(hmask, 1); + + __m128i q8_0 = _mm_loadu_si128((const __m128i*)q8); q8 += 16; + __m128i q8_1 = _mm_loadu_si128((const __m128i*)q8); q8 += 16; + __m128i p16_0 = _mm_maddubs_epi16(q5_0, q8_0); + __m128i p16_1 = _mm_maddubs_epi16(q5_1, q8_1); + p16_0 = _mm_madd_epi16(scale_0, p16_0); + p16_1 = _mm_madd_epi16(scale_0, p16_1); + + q5l_0 = _mm_and_si128(_mm_srli_epi16(q5bits_0, 4), m4); + q5l_1 = _mm_and_si128(_mm_srli_epi16(q5bits_1, 4), m4); + q5h_0 = _mm_slli_epi16(_mm_srli_epi16(_mm_and_si128(hbits_0, hmask), bit), 4); + q5h_1 = _mm_slli_epi16(_mm_srli_epi16(_mm_and_si128(hbits_1, hmask), bit++), 4); + q5_0 = _mm_add_epi8(q5l_0, q5h_0); + q5_1 = _mm_add_epi8(q5l_1, q5h_1); + hmask = _mm_slli_epi16(hmask, 1); + + q8_0 = _mm_loadu_si128((const __m128i*)q8); q8 += 16; + q8_1 = _mm_loadu_si128((const __m128i*)q8); q8 += 16; + __m128i p16_2 = _mm_maddubs_epi16(q5_0, q8_0); + __m128i p16_3 = _mm_maddubs_epi16(q5_1, q8_1); + p16_2 = _mm_madd_epi16(scale_1, p16_2); + p16_3 = _mm_madd_epi16(scale_1, p16_3); + + sumi_0 = _mm_add_epi32(sumi_0, _mm_add_epi32(p16_0, p16_2)); + sumi_1 = _mm_add_epi32(sumi_1, _mm_add_epi32(p16_1, p16_3)); + + } + + __m256 vd = _mm256_set1_ps(d); + __m256i sumi = MM256_SET_M128I(sumi_1, sumi_0); + acc = _mm256_add_ps(_mm256_mul_ps(vd, _mm256_cvtepi32_ps(sumi)), acc); + + } + + *s = hsum_float_8(acc) + summs; + +#else + UNUSED(x); + UNUSED(y); + UNUSED(nb); + UNUSED(kmask1); + UNUSED(kmask2); + UNUSED(kmask3); + UNUSED(utmp); + ggml_vec_dot_q5_K_q8_K_generic(n, s, bs, vx, bx, vy, by, nrc); +#endif +} + +void ggml_vec_dot_q6_K_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) { + assert(n % QK_K == 0); + assert(nrc == 1); + UNUSED(nrc); + UNUSED(bx); + UNUSED(by); + UNUSED(bs); + + const block_q6_K * GGML_RESTRICT x = vx; + const block_q8_K * GGML_RESTRICT y = vy; + + const int nb = n / QK_K; + +#if defined __AVX2__ + + const __m256i m3 = _mm256_set1_epi8(3); + const __m256i m15 = _mm256_set1_epi8(15); + + __m256 acc = _mm256_setzero_ps(); + + for (int i = 0; i < nb; ++i) { + + const float d = y[i].d * GGML_CPU_FP16_TO_FP32(x[i].d); + + const uint8_t * GGML_RESTRICT q4 = x[i].ql; + const uint8_t * GGML_RESTRICT qh = x[i].qh; + const int8_t * GGML_RESTRICT q8 = y[i].qs; + + const __m256i q8sums = _mm256_loadu_si256((const __m256i*)y[i].bsums); + const __m128i scales = _mm_loadu_si128((const __m128i*)x[i].scales); + const __m256i scales_16 = _mm256_cvtepi8_epi16(scales); + const __m256i q8sclsub = _mm256_slli_epi32(_mm256_madd_epi16(q8sums, scales_16), 5); + + __m256i sumi = _mm256_setzero_si256(); + + int is = 0; + + for (int j = 0; j < QK_K/128; ++j) { + const __m256i q4bits1 = _mm256_loadu_si256((const __m256i*)q4); q4 += 32; + const __m256i q4bits2 = _mm256_loadu_si256((const __m256i*)q4); q4 += 32; + const __m256i q4bitsH = _mm256_loadu_si256((const __m256i*)qh); qh += 32; + + const __m256i q4h_0 = _mm256_slli_epi16(_mm256_and_si256(q4bitsH, m3), 4); + const __m256i q4h_1 = _mm256_slli_epi16(_mm256_and_si256(q4bitsH, _mm256_set1_epi8(12)), 2); + const __m256i q4h_2 = _mm256_and_si256(q4bitsH, _mm256_set1_epi8(48)); + const __m256i q4h_3 = _mm256_srli_epi16(_mm256_and_si256(q4bitsH, _mm256_set1_epi8(-64)), 2); + + const __m256i q4_0 = _mm256_or_si256(_mm256_and_si256(q4bits1, m15), q4h_0); + const __m256i q4_1 = _mm256_or_si256(_mm256_and_si256(q4bits2, m15), q4h_1); + const __m256i q4_2 = _mm256_or_si256(_mm256_and_si256(_mm256_srli_epi16(q4bits1, 4), m15), q4h_2); + const __m256i q4_3 = _mm256_or_si256(_mm256_and_si256(_mm256_srli_epi16(q4bits2, 4), m15), q4h_3); + + const __m256i q8_0 = _mm256_loadu_si256((const __m256i*)q8); q8 += 32; + const __m256i q8_1 = _mm256_loadu_si256((const __m256i*)q8); q8 += 32; + const __m256i q8_2 = _mm256_loadu_si256((const __m256i*)q8); q8 += 32; + const __m256i q8_3 = _mm256_loadu_si256((const __m256i*)q8); q8 += 32; + + __m256i p16_0 = _mm256_maddubs_epi16(q4_0, q8_0); + __m256i p16_1 = _mm256_maddubs_epi16(q4_1, q8_1); + __m256i p16_2 = _mm256_maddubs_epi16(q4_2, q8_2); + __m256i p16_3 = _mm256_maddubs_epi16(q4_3, q8_3); + + const __m128i scale_0 = _mm_shuffle_epi8(scales, get_scale_shuffle(is + 0)); + const __m128i scale_1 = _mm_shuffle_epi8(scales, get_scale_shuffle(is + 1)); + const __m128i scale_2 = _mm_shuffle_epi8(scales, get_scale_shuffle(is + 2)); + const __m128i scale_3 = _mm_shuffle_epi8(scales, get_scale_shuffle(is + 3)); + is += 4; + + p16_0 = _mm256_madd_epi16(_mm256_cvtepi8_epi16(scale_0), p16_0); + p16_1 = _mm256_madd_epi16(_mm256_cvtepi8_epi16(scale_1), p16_1); + p16_2 = _mm256_madd_epi16(_mm256_cvtepi8_epi16(scale_2), p16_2); + p16_3 = _mm256_madd_epi16(_mm256_cvtepi8_epi16(scale_3), p16_3); + + sumi = _mm256_add_epi32(sumi, _mm256_add_epi32(p16_0, p16_1)); + sumi = _mm256_add_epi32(sumi, _mm256_add_epi32(p16_2, p16_3)); + + } + + sumi = _mm256_sub_epi32(sumi, q8sclsub); + acc = _mm256_fmadd_ps(_mm256_broadcast_ss(&d), _mm256_cvtepi32_ps(sumi), acc); + } + + *s = hsum_float_8(acc); + +#elif defined __AVX__ + + const __m128i m3 = _mm_set1_epi8(3); + const __m128i m15 = _mm_set1_epi8(15); + + __m256 acc = _mm256_setzero_ps(); + + for (int i = 0; i < nb; ++i) { + + const float d = y[i].d * GGML_CPU_FP16_TO_FP32(x[i].d); + + const uint8_t * GGML_RESTRICT q4 = x[i].ql; + const uint8_t * GGML_RESTRICT qh = x[i].qh; + const int8_t * GGML_RESTRICT q8 = y[i].qs; + + // handle the q6_k -32 offset separately using bsums + const __m128i q8sums_0 = _mm_loadu_si128((const __m128i*)y[i].bsums); + const __m128i q8sums_1 = _mm_loadu_si128((const __m128i*)y[i].bsums + 1); + const __m128i scales = _mm_loadu_si128((const __m128i*)x[i].scales); + const __m128i scales_16_0 = _mm_cvtepi8_epi16(scales); + const __m128i scales_16_1 = _mm_cvtepi8_epi16(_mm_bsrli_si128(scales, 8)); + const __m128i q8sclsub_0 = _mm_slli_epi32(_mm_madd_epi16(q8sums_0, scales_16_0), 5); + const __m128i q8sclsub_1 = _mm_slli_epi32(_mm_madd_epi16(q8sums_1, scales_16_1), 5); + + __m128i sumi_0 = _mm_setzero_si128(); + __m128i sumi_1 = _mm_setzero_si128(); + + int is = 0; + + for (int j = 0; j < QK_K/128; ++j) { + + const __m128i q4bitsH_0 = _mm_loadu_si128((const __m128i*)qh); qh += 16; + const __m128i q4bitsH_1 = _mm_loadu_si128((const __m128i*)qh); qh += 16; + + const __m128i q4h_0 = _mm_slli_epi16(_mm_and_si128(q4bitsH_0, m3), 4); + const __m128i q4h_1 = _mm_slli_epi16(_mm_and_si128(q4bitsH_1, m3), 4); + const __m128i q4h_2 = _mm_slli_epi16(_mm_and_si128(q4bitsH_0, _mm_set1_epi8(12)), 2); + const __m128i q4h_3 = _mm_slli_epi16(_mm_and_si128(q4bitsH_1, _mm_set1_epi8(12)), 2); + const __m128i q4h_4 = _mm_and_si128(q4bitsH_0, _mm_set1_epi8(48)); + const __m128i q4h_5 = _mm_and_si128(q4bitsH_1, _mm_set1_epi8(48)); + const __m128i q4h_6 = _mm_srli_epi16(_mm_and_si128(q4bitsH_0, _mm_set1_epi8(-64)), 2); + const __m128i q4h_7 = _mm_srli_epi16(_mm_and_si128(q4bitsH_1, _mm_set1_epi8(-64)), 2); + + const __m128i q4bits1_0 = _mm_loadu_si128((const __m128i*)q4); q4 += 16; + const __m128i q4bits1_1 = _mm_loadu_si128((const __m128i*)q4); q4 += 16; + const __m128i q4bits2_0 = _mm_loadu_si128((const __m128i*)q4); q4 += 16; + const __m128i q4bits2_1 = _mm_loadu_si128((const __m128i*)q4); q4 += 16; + + const __m128i q4_0 = _mm_or_si128(_mm_and_si128(q4bits1_0, m15), q4h_0); + const __m128i q4_1 = _mm_or_si128(_mm_and_si128(q4bits1_1, m15), q4h_1); + const __m128i q4_2 = _mm_or_si128(_mm_and_si128(q4bits2_0, m15), q4h_2); + const __m128i q4_3 = _mm_or_si128(_mm_and_si128(q4bits2_1, m15), q4h_3); + const __m128i q4_4 = _mm_or_si128(_mm_and_si128(_mm_srli_epi16(q4bits1_0, 4), m15), q4h_4); + const __m128i q4_5 = _mm_or_si128(_mm_and_si128(_mm_srli_epi16(q4bits1_1, 4), m15), q4h_5); + const __m128i q4_6 = _mm_or_si128(_mm_and_si128(_mm_srli_epi16(q4bits2_0, 4), m15), q4h_6); + const __m128i q4_7 = _mm_or_si128(_mm_and_si128(_mm_srli_epi16(q4bits2_1, 4), m15), q4h_7); + + const __m128i q8_0 = _mm_loadu_si128((const __m128i*)q8); q8 += 16; + const __m128i q8_1 = _mm_loadu_si128((const __m128i*)q8); q8 += 16; + const __m128i q8_2 = _mm_loadu_si128((const __m128i*)q8); q8 += 16; + const __m128i q8_3 = _mm_loadu_si128((const __m128i*)q8); q8 += 16; + const __m128i q8_4 = _mm_loadu_si128((const __m128i*)q8); q8 += 16; + const __m128i q8_5 = _mm_loadu_si128((const __m128i*)q8); q8 += 16; + const __m128i q8_6 = _mm_loadu_si128((const __m128i*)q8); q8 += 16; + const __m128i q8_7 = _mm_loadu_si128((const __m128i*)q8); q8 += 16; + + __m128i p16_0 = _mm_maddubs_epi16(q4_0, q8_0); + __m128i p16_1 = _mm_maddubs_epi16(q4_1, q8_1); + __m128i p16_2 = _mm_maddubs_epi16(q4_2, q8_2); + __m128i p16_3 = _mm_maddubs_epi16(q4_3, q8_3); + __m128i p16_4 = _mm_maddubs_epi16(q4_4, q8_4); + __m128i p16_5 = _mm_maddubs_epi16(q4_5, q8_5); + __m128i p16_6 = _mm_maddubs_epi16(q4_6, q8_6); + __m128i p16_7 = _mm_maddubs_epi16(q4_7, q8_7); + + const __m128i scale_0 = _mm_shuffle_epi8(scales, get_scale_shuffle(is + 0)); + const __m128i scale_1 = _mm_shuffle_epi8(scales, get_scale_shuffle(is + 1)); + const __m128i scale_2 = _mm_shuffle_epi8(scales, get_scale_shuffle(is + 2)); + const __m128i scale_3 = _mm_shuffle_epi8(scales, get_scale_shuffle(is + 3)); + is += 4; + + p16_0 = _mm_madd_epi16(_mm_cvtepi8_epi16(scale_0), p16_0); + p16_1 = _mm_madd_epi16(_mm_cvtepi8_epi16(_mm_bsrli_si128(scale_0, 8)), p16_1); + p16_2 = _mm_madd_epi16(_mm_cvtepi8_epi16(scale_1), p16_2); + p16_3 = _mm_madd_epi16(_mm_cvtepi8_epi16(_mm_bsrli_si128(scale_1, 8)), p16_3); + p16_4 = _mm_madd_epi16(_mm_cvtepi8_epi16(scale_2), p16_4); + p16_5 = _mm_madd_epi16(_mm_cvtepi8_epi16(_mm_bsrli_si128(scale_2, 8)), p16_5); + p16_6 = _mm_madd_epi16(_mm_cvtepi8_epi16(scale_3), p16_6); + p16_7 = _mm_madd_epi16(_mm_cvtepi8_epi16(_mm_bsrli_si128(scale_3, 8)), p16_7); + + sumi_0 = _mm_add_epi32(sumi_0, _mm_add_epi32(p16_0, p16_2)); + sumi_1 = _mm_add_epi32(sumi_1, _mm_add_epi32(p16_1, p16_3)); + sumi_0 = _mm_add_epi32(sumi_0, _mm_add_epi32(p16_4, p16_6)); + sumi_1 = _mm_add_epi32(sumi_1, _mm_add_epi32(p16_5, p16_7)); + + } + + sumi_0 = _mm_sub_epi32(sumi_0, q8sclsub_0); + sumi_1 = _mm_sub_epi32(sumi_1, q8sclsub_1); + const __m256i sumi = MM256_SET_M128I(sumi_1, sumi_0); + acc = _mm256_add_ps(_mm256_mul_ps(_mm256_set1_ps(d), _mm256_cvtepi32_ps(sumi)), acc); + } + + *s = hsum_float_8(acc); + +#else + UNUSED(x); + UNUSED(y); + UNUSED(nb); + ggml_vec_dot_q6_K_q8_K_generic(n, s, bs, vx, bx, vy, by, nrc); +#endif +} + +#if defined (__AVX__) || defined (__AVX2__) +static const int8_t keven_signs_q2xs[1024] = { + 1, 1, 1, 1, 1, 1, 1, 1, -1, 1, 1, 1, 1, 1, 1, -1, 1, -1, 1, 1, 1, 1, 1, -1, -1, -1, 1, 1, 1, 1, 1, 1, + 1, 1, -1, 1, 1, 1, 1, -1, -1, 1, -1, 1, 1, 1, 1, 1, 1, -1, -1, 1, 1, 1, 1, 1, -1, -1, -1, 1, 1, 1, 1, -1, + 1, 1, 1, -1, 1, 1, 1, -1, -1, 1, 1, -1, 1, 1, 1, 1, 1, -1, 1, -1, 1, 1, 1, 1, -1, -1, 1, -1, 1, 1, 1, -1, + 1, 1, -1, -1, 1, 1, 1, 1, -1, 1, -1, -1, 1, 1, 1, -1, 1, -1, -1, -1, 1, 1, 1, -1, -1, -1, -1, -1, 1, 1, 1, 1, + 1, 1, 1, 1, -1, 1, 1, -1, -1, 1, 1, 1, -1, 1, 1, 1, 1, -1, 1, 1, -1, 1, 1, 1, -1, -1, 1, 1, -1, 1, 1, -1, + 1, 1, -1, 1, -1, 1, 1, 1, -1, 1, -1, 1, -1, 1, 1, -1, 1, -1, -1, 1, -1, 1, 1, -1, -1, -1, -1, 1, -1, 1, 1, 1, + 1, 1, 1, -1, -1, 1, 1, 1, -1, 1, 1, -1, -1, 1, 1, -1, 1, -1, 1, -1, -1, 1, 1, -1, -1, -1, 1, -1, -1, 1, 1, 1, + 1, 1, -1, -1, -1, 1, 1, -1, -1, 1, -1, -1, -1, 1, 1, 1, 1, -1, -1, -1, -1, 1, 1, 1, -1, -1, -1, -1, -1, 1, 1, -1, + 1, 1, 1, 1, 1, -1, 1, -1, -1, 1, 1, 1, 1, -1, 1, 1, 1, -1, 1, 1, 1, -1, 1, 1, -1, -1, 1, 1, 1, -1, 1, -1, + 1, 1, -1, 1, 1, -1, 1, 1, -1, 1, -1, 1, 1, -1, 1, -1, 1, -1, -1, 1, 1, -1, 1, -1, -1, -1, -1, 1, 1, -1, 1, 1, + 1, 1, 1, -1, 1, -1, 1, 1, -1, 1, 1, -1, 1, -1, 1, -1, 1, -1, 1, -1, 1, -1, 1, -1, -1, -1, 1, -1, 1, -1, 1, 1, + 1, 1, -1, -1, 1, -1, 1, -1, -1, 1, -1, -1, 1, -1, 1, 1, 1, -1, -1, -1, 1, -1, 1, 1, -1, -1, -1, -1, 1, -1, 1, -1, + 1, 1, 1, 1, -1, -1, 1, 1, -1, 1, 1, 1, -1, -1, 1, -1, 1, -1, 1, 1, -1, -1, 1, -1, -1, -1, 1, 1, -1, -1, 1, 1, + 1, 1, -1, 1, -1, -1, 1, -1, -1, 1, -1, 1, -1, -1, 1, 1, 1, -1, -1, 1, -1, -1, 1, 1, -1, -1, -1, 1, -1, -1, 1, -1, + 1, 1, 1, -1, -1, -1, 1, -1, -1, 1, 1, -1, -1, -1, 1, 1, 1, -1, 1, -1, -1, -1, 1, 1, -1, -1, 1, -1, -1, -1, 1, -1, + 1, 1, -1, -1, -1, -1, 1, 1, -1, 1, -1, -1, -1, -1, 1, -1, 1, -1, -1, -1, -1, -1, 1, -1, -1, -1, -1, -1, -1, -1, 1, 1, + 1, 1, 1, 1, 1, 1, -1, -1, -1, 1, 1, 1, 1, 1, -1, 1, 1, -1, 1, 1, 1, 1, -1, 1, -1, -1, 1, 1, 1, 1, -1, -1, + 1, 1, -1, 1, 1, 1, -1, 1, -1, 1, -1, 1, 1, 1, -1, -1, 1, -1, -1, 1, 1, 1, -1, -1, -1, -1, -1, 1, 1, 1, -1, 1, + 1, 1, 1, -1, 1, 1, -1, 1, -1, 1, 1, -1, 1, 1, -1, -1, 1, -1, 1, -1, 1, 1, -1, -1, -1, -1, 1, -1, 1, 1, -1, 1, + 1, 1, -1, -1, 1, 1, -1, -1, -1, 1, -1, -1, 1, 1, -1, 1, 1, -1, -1, -1, 1, 1, -1, 1, -1, -1, -1, -1, 1, 1, -1, -1, + 1, 1, 1, 1, -1, 1, -1, 1, -1, 1, 1, 1, -1, 1, -1, -1, 1, -1, 1, 1, -1, 1, -1, -1, -1, -1, 1, 1, -1, 1, -1, 1, + 1, 1, -1, 1, -1, 1, -1, -1, -1, 1, -1, 1, -1, 1, -1, 1, 1, -1, -1, 1, -1, 1, -1, 1, -1, -1, -1, 1, -1, 1, -1, -1, + 1, 1, 1, -1, -1, 1, -1, -1, -1, 1, 1, -1, -1, 1, -1, 1, 1, -1, 1, -1, -1, 1, -1, 1, -1, -1, 1, -1, -1, 1, -1, -1, + 1, 1, -1, -1, -1, 1, -1, 1, -1, 1, -1, -1, -1, 1, -1, -1, 1, -1, -1, -1, -1, 1, -1, -1, -1, -1, -1, -1, -1, 1, -1, 1, + 1, 1, 1, 1, 1, -1, -1, 1, -1, 1, 1, 1, 1, -1, -1, -1, 1, -1, 1, 1, 1, -1, -1, -1, -1, -1, 1, 1, 1, -1, -1, 1, + 1, 1, -1, 1, 1, -1, -1, -1, -1, 1, -1, 1, 1, -1, -1, 1, 1, -1, -1, 1, 1, -1, -1, 1, -1, -1, -1, 1, 1, -1, -1, -1, + 1, 1, 1, -1, 1, -1, -1, -1, -1, 1, 1, -1, 1, -1, -1, 1, 1, -1, 1, -1, 1, -1, -1, 1, -1, -1, 1, -1, 1, -1, -1, -1, + 1, 1, -1, -1, 1, -1, -1, 1, -1, 1, -1, -1, 1, -1, -1, -1, 1, -1, -1, -1, 1, -1, -1, -1, -1, -1, -1, -1, 1, -1, -1, 1, + 1, 1, 1, 1, -1, -1, -1, -1, -1, 1, 1, 1, -1, -1, -1, 1, 1, -1, 1, 1, -1, -1, -1, 1, -1, -1, 1, 1, -1, -1, -1, -1, + 1, 1, -1, 1, -1, -1, -1, 1, -1, 1, -1, 1, -1, -1, -1, -1, 1, -1, -1, 1, -1, -1, -1, -1, -1, -1, -1, 1, -1, -1, -1, 1, + 1, 1, 1, -1, -1, -1, -1, 1, -1, 1, 1, -1, -1, -1, -1, -1, 1, -1, 1, -1, -1, -1, -1, -1, -1, -1, 1, -1, -1, -1, -1, 1, + 1, 1, -1, -1, -1, -1, -1, -1, -1, 1, -1, -1, -1, -1, -1, 1, 1, -1, -1, -1, -1, -1, -1, 1, -1, -1, -1, -1, -1, -1, -1, -1, +}; +#endif + +void ggml_vec_dot_iq2_xxs_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) { + assert(n % QK_K == 0); + assert(nrc == 1); + UNUSED(nrc); + UNUSED(bx); + UNUSED(by); + UNUSED(bs); + + const block_iq2_xxs * GGML_RESTRICT x = vx; + const block_q8_K * GGML_RESTRICT y = vy; + + const int nb = n / QK_K; + +#if defined(__AVX2__) + + const uint64_t * signs64 = (const uint64_t *)keven_signs_q2xs; + + uint32_t aux32[4]; + const uint8_t * aux8 = (const uint8_t *)aux32; + + __m256 accumf = _mm256_setzero_ps(); + for (int i = 0; i < nb; ++i) { + const float d = GGML_CPU_FP16_TO_FP32(x[i].d) * y[i].d; + const uint16_t * GGML_RESTRICT q2 = x[i].qs; + const int8_t * GGML_RESTRICT q8 = y[i].qs; + __m256i sumi1 = _mm256_setzero_si256(); + __m256i sumi2 = _mm256_setzero_si256(); + for (int ib32 = 0; ib32 < QK_K/32; ib32 += 2) { + const __m256i q8_1 = _mm256_loadu_si256((const __m256i *)q8); q8 += 32; + const __m256i q8_2 = _mm256_loadu_si256((const __m256i *)q8); q8 += 32; + memcpy(aux32, q2, 4*sizeof(uint32_t)); q2 += 8; + const __m256i q2_1 = _mm256_set_epi64x(iq2xxs_grid[aux8[ 3]], iq2xxs_grid[aux8[ 2]], iq2xxs_grid[aux8[1]], iq2xxs_grid[aux8[0]]); + const __m256i q2_2 = _mm256_set_epi64x(iq2xxs_grid[aux8[11]], iq2xxs_grid[aux8[10]], iq2xxs_grid[aux8[9]], iq2xxs_grid[aux8[8]]); + const __m256i s2_1 = _mm256_set_epi64x(signs64[(aux32[1] >> 21) & 127], signs64[(aux32[1] >> 14) & 127], + signs64[(aux32[1] >> 7) & 127], signs64[(aux32[1] >> 0) & 127]); + const __m256i s2_2 = _mm256_set_epi64x(signs64[(aux32[3] >> 21) & 127], signs64[(aux32[3] >> 14) & 127], + signs64[(aux32[3] >> 7) & 127], signs64[(aux32[3] >> 0) & 127]); + const __m256i q8s_1 = _mm256_sign_epi8(q8_1, s2_1); + const __m256i q8s_2 = _mm256_sign_epi8(q8_2, s2_2); + const __m256i dot1 = _mm256_maddubs_epi16(q2_1, q8s_1); + const __m256i dot2 = _mm256_maddubs_epi16(q2_2, q8s_2); + const uint16_t ls1 = aux32[1] >> 28; + const uint16_t ls2 = aux32[3] >> 28; + const __m256i p1 = _mm256_madd_epi16(dot1, _mm256_set1_epi16(2*ls1+1)); + const __m256i p2 = _mm256_madd_epi16(dot2, _mm256_set1_epi16(2*ls2+1)); + sumi1 = _mm256_add_epi32(sumi1, p1); + sumi2 = _mm256_add_epi32(sumi2, p2); + } + + accumf = _mm256_fmadd_ps(_mm256_set1_ps(d), _mm256_cvtepi32_ps(_mm256_add_epi32(sumi1, sumi2)), accumf); + + } + + *s = 0.125f * hsum_float_8(accumf); + +#elif defined(__AVX__) + const uint64_t * signs64 = (const uint64_t *)keven_signs_q2xs; + + uint32_t aux32[4]; + const uint8_t * aux8 = (const uint8_t *)aux32; + + __m256 accumf = _mm256_setzero_ps(); + for (int i = 0; i < nb; ++i) { + const float d = GGML_CPU_FP16_TO_FP32(x[i].d) * y[i].d; + const uint16_t * GGML_RESTRICT q2 = x[i].qs; + const int8_t * GGML_RESTRICT q8 = y[i].qs; + __m128i sumi1_0 = _mm_setzero_si128(); + __m128i sumi1_1 = _mm_setzero_si128(); + __m128i sumi2_0 = _mm_setzero_si128(); + __m128i sumi2_1 = _mm_setzero_si128(); + for (int ib32 = 0; ib32 < QK_K/32; ib32 += 2) { + const __m128i q8_1_0 = _mm_loadu_si128((const __m128i *)q8); q8 += 16; + const __m128i q8_1_1 = _mm_loadu_si128((const __m128i *)q8); q8 += 16; + const __m128i q8_2_0 = _mm_loadu_si128((const __m128i *)q8); q8 += 16; + const __m128i q8_2_1 = _mm_loadu_si128((const __m128i *)q8); q8 += 16; + memcpy(aux32, q2, 4*sizeof(uint32_t)); q2 += 8; + const __m128i q2_1_0 = _mm_set_epi64x(iq2xxs_grid[aux8[1]], iq2xxs_grid[aux8[0]]); + const __m128i q2_1_1 = _mm_set_epi64x(iq2xxs_grid[aux8[3]], iq2xxs_grid[aux8[2]]); + const __m128i q2_2_0 = _mm_set_epi64x(iq2xxs_grid[aux8[9]], iq2xxs_grid[aux8[8]]); + const __m128i q2_2_1 = _mm_set_epi64x(iq2xxs_grid[aux8[11]], iq2xxs_grid[aux8[10]]); + const __m128i s2_1_0 = _mm_set_epi64x(signs64[(aux32[1] >> 7) & 127], signs64[(aux32[1] >> 0) & 127]); + const __m128i s2_1_1 = _mm_set_epi64x(signs64[(aux32[1] >> 21) & 127], signs64[(aux32[1] >> 14) & 127]); + const __m128i s2_2_0 = _mm_set_epi64x(signs64[(aux32[3] >> 7) & 127], signs64[(aux32[3] >> 0) & 127]); + const __m128i s2_2_1 = _mm_set_epi64x(signs64[(aux32[3] >> 21) & 127], signs64[(aux32[3] >> 14) & 127]); + const __m128i q8s_1_0 = _mm_sign_epi8(q8_1_0, s2_1_0); + const __m128i q8s_1_1 = _mm_sign_epi8(q8_1_1, s2_1_1); + const __m128i q8s_2_0 = _mm_sign_epi8(q8_2_0, s2_2_0); + const __m128i q8s_2_1 = _mm_sign_epi8(q8_2_1, s2_2_1); + const __m128i dot1_0 = _mm_maddubs_epi16(q2_1_0, q8s_1_0); + const __m128i dot1_1 = _mm_maddubs_epi16(q2_1_1, q8s_1_1); + const __m128i dot2_0 = _mm_maddubs_epi16(q2_2_0, q8s_2_0); + const __m128i dot2_1 = _mm_maddubs_epi16(q2_2_1, q8s_2_1); + const uint16_t ls1 = aux32[1] >> 28; + const uint16_t ls2 = aux32[3] >> 28; + const __m128i p1_0 = _mm_madd_epi16(dot1_0, _mm_set1_epi16(2*ls1+1)); + const __m128i p1_1 = _mm_madd_epi16(dot1_1, _mm_set1_epi16(2*ls1+1)); + const __m128i p2_0 = _mm_madd_epi16(dot2_0, _mm_set1_epi16(2*ls2+1)); + const __m128i p2_1 = _mm_madd_epi16(dot2_1, _mm_set1_epi16(2*ls2+1)); + sumi1_0 = _mm_add_epi32(sumi1_0, p1_0); + sumi1_1 = _mm_add_epi32(sumi1_1, p1_1); + sumi2_0 = _mm_add_epi32(sumi2_0, p2_0); + sumi2_1 = _mm_add_epi32(sumi2_1, p2_1); + } + + accumf = _mm256_add_ps(_mm256_mul_ps(_mm256_set1_ps(d), _mm256_cvtepi32_ps(MM256_SET_M128I(_mm_add_epi32(sumi1_1, sumi2_1), _mm_add_epi32(sumi1_0, sumi2_0)))), accumf); + + } + + *s = 0.125f * hsum_float_8(accumf); + +#else + UNUSED(x); + UNUSED(y); + UNUSED(nb); + ggml_vec_dot_iq2_xxs_q8_K_generic(n, s, bs, vx, bx, vy, by, nrc); +#endif +} + +void ggml_vec_dot_iq2_xs_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) { + assert(n % QK_K == 0); + assert(nrc == 1); + UNUSED(nrc); + UNUSED(bx); + UNUSED(by); + UNUSED(bs); + + const block_iq2_xs * GGML_RESTRICT x = vx; + const block_q8_K * GGML_RESTRICT y = vy; + + const int nb = n / QK_K; + +#if defined(__AVX2__) + + const __m256i mone = _mm256_set1_epi8(1); + static const char block_sign_shuffle_mask_1[32] = { + 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x02, 0x02, 0x02, 0x02, 0x02, 0x02, 0x02, 0x02, + 0x04, 0x04, 0x04, 0x04, 0x04, 0x04, 0x04, 0x04, 0x06, 0x06, 0x06, 0x06, 0x06, 0x06, 0x06, 0x06, + }; + static const char block_sign_shuffle_mask_2[32] = { + 0x08, 0x08, 0x08, 0x08, 0x08, 0x08, 0x08, 0x08, 0x0a, 0x0a, 0x0a, 0x0a, 0x0a, 0x0a, 0x0a, 0x0a, + 0x0c, 0x0c, 0x0c, 0x0c, 0x0c, 0x0c, 0x0c, 0x0c, 0x0e, 0x0e, 0x0e, 0x0e, 0x0e, 0x0e, 0x0e, 0x0e, + }; + static const uint8_t bit_selector_mask_bytes[32] = { + 0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80, 0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80, + 0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80, 0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80, + }; + + const __m256i bit_selector_mask = _mm256_loadu_si256((const __m256i*)bit_selector_mask_bytes); + const __m256i block_sign_shuffle_1 = _mm256_loadu_si256((const __m256i*)block_sign_shuffle_mask_1); + const __m256i block_sign_shuffle_2 = _mm256_loadu_si256((const __m256i*)block_sign_shuffle_mask_2); + + static const uint8_t k_bit_helper[32] = { + 0x00, 0x80, 0x80, 0x00, 0x80, 0x00, 0x00, 0x80, 0x80, 0x00, 0x00, 0x80, 0x00, 0x80, 0x80, 0x00, + 0x00, 0x80, 0x80, 0x00, 0x80, 0x00, 0x00, 0x80, 0x80, 0x00, 0x00, 0x80, 0x00, 0x80, 0x80, 0x00, + }; + const __m256i bit_helper = _mm256_loadu_si256((const __m256i*)k_bit_helper); + const __m256i m511 = _mm256_set1_epi16(511); + const __m128i m4 = _mm_set1_epi8(0xf); + const __m128i m1 = _mm_set1_epi8(1); + + uint64_t aux64; + + // somewhat hacky, but gives a significant boost in performance + __m256i aux_gindex; + const uint16_t * gindex = (const uint16_t *)&aux_gindex; + + __m256 accumf = _mm256_setzero_ps(); + for (int i = 0; i < nb; ++i) { + const float d = GGML_CPU_FP16_TO_FP32(x[i].d) * y[i].d; + const uint16_t * GGML_RESTRICT q2 = x[i].qs; + const int8_t * GGML_RESTRICT q8 = y[i].qs; + + memcpy(&aux64, x[i].scales, 8); + __m128i stmp = _mm_set1_epi64x(aux64); + stmp = _mm_unpacklo_epi8(_mm_and_si128(stmp, m4), _mm_and_si128(_mm_srli_epi16(stmp, 4), m4)); + const __m128i scales = _mm_add_epi8(_mm_slli_epi16(stmp, 1), m1); + + __m256i sumi1 = _mm256_setzero_si256(); + __m256i sumi2 = _mm256_setzero_si256(); + for (int ib32 = 0; ib32 < QK_K/32; ib32 += 4) { + + const __m256i q2_data = _mm256_loadu_si256((const __m256i*)q2); q2 += 16; + aux_gindex = _mm256_and_si256(q2_data, m511); + + const __m256i partial_sign_bits = _mm256_srli_epi16(q2_data, 9); + const __m256i partial_sign_bits_upper = _mm256_srli_epi16(q2_data, 13); + const __m256i partial_sign_bits_for_counting = _mm256_xor_si256(partial_sign_bits, partial_sign_bits_upper); + + const __m256i odd_bits = _mm256_shuffle_epi8(bit_helper, partial_sign_bits_for_counting); + const __m256i full_sign_bits = _mm256_or_si256(partial_sign_bits, odd_bits); + + const __m256i q8_1 = _mm256_loadu_si256((const __m256i *)q8); q8 += 32; + const __m256i q8_2 = _mm256_loadu_si256((const __m256i *)q8); q8 += 32; + const __m256i q8_3 = _mm256_loadu_si256((const __m256i *)q8); q8 += 32; + const __m256i q8_4 = _mm256_loadu_si256((const __m256i *)q8); q8 += 32; + + const __m256i q2_1 = _mm256_set_epi64x(iq2xs_grid[gindex[ 3]], iq2xs_grid[gindex[ 2]], + iq2xs_grid[gindex[ 1]], iq2xs_grid[gindex[ 0]]); + const __m256i q2_2 = _mm256_set_epi64x(iq2xs_grid[gindex[ 7]], iq2xs_grid[gindex[ 6]], + iq2xs_grid[gindex[ 5]], iq2xs_grid[gindex[ 4]]); + const __m256i q2_3 = _mm256_set_epi64x(iq2xs_grid[gindex[11]], iq2xs_grid[gindex[10]], + iq2xs_grid[gindex[ 9]], iq2xs_grid[gindex[ 8]]); + const __m256i q2_4 = _mm256_set_epi64x(iq2xs_grid[gindex[15]], iq2xs_grid[gindex[14]], + iq2xs_grid[gindex[13]], iq2xs_grid[gindex[12]]); + + const __m128i full_signs_l = _mm256_castsi256_si128(full_sign_bits); + const __m128i full_signs_h = _mm256_extractf128_si256(full_sign_bits, 1); + const __m256i full_signs_1 = MM256_SET_M128I(full_signs_l, full_signs_l); + const __m256i full_signs_2 = MM256_SET_M128I(full_signs_h, full_signs_h); + + __m256i signs; + signs = _mm256_shuffle_epi8(full_signs_1, block_sign_shuffle_1); + signs = _mm256_cmpeq_epi8(_mm256_and_si256(signs, bit_selector_mask), bit_selector_mask); + const __m256i q8s_1 = _mm256_sign_epi8(q8_1, _mm256_or_si256(signs, mone)); + + signs = _mm256_shuffle_epi8(full_signs_1, block_sign_shuffle_2); + signs = _mm256_cmpeq_epi8(_mm256_and_si256(signs, bit_selector_mask), bit_selector_mask); + const __m256i q8s_2 = _mm256_sign_epi8(q8_2, _mm256_or_si256(signs, mone)); + + signs = _mm256_shuffle_epi8(full_signs_2, block_sign_shuffle_1); + signs = _mm256_cmpeq_epi8(_mm256_and_si256(signs, bit_selector_mask), bit_selector_mask); + const __m256i q8s_3 = _mm256_sign_epi8(q8_3, _mm256_or_si256(signs, mone)); + + signs = _mm256_shuffle_epi8(full_signs_2, block_sign_shuffle_2); + signs = _mm256_cmpeq_epi8(_mm256_and_si256(signs, bit_selector_mask), bit_selector_mask); + const __m256i q8s_4 = _mm256_sign_epi8(q8_4, _mm256_or_si256(signs, mone)); + + const __m256i dot1 = _mm256_maddubs_epi16(q2_1, q8s_1); + const __m256i dot2 = _mm256_maddubs_epi16(q2_2, q8s_2); + const __m256i dot3 = _mm256_maddubs_epi16(q2_3, q8s_3); + const __m256i dot4 = _mm256_maddubs_epi16(q2_4, q8s_4); + + const __m256i sc1 = _mm256_cvtepi8_epi16(_mm_shuffle_epi8(scales, get_scale_shuffle(ib32+0))); + const __m256i sc2 = _mm256_cvtepi8_epi16(_mm_shuffle_epi8(scales, get_scale_shuffle(ib32+1))); + const __m256i sc3 = _mm256_cvtepi8_epi16(_mm_shuffle_epi8(scales, get_scale_shuffle(ib32+2))); + const __m256i sc4 = _mm256_cvtepi8_epi16(_mm_shuffle_epi8(scales, get_scale_shuffle(ib32+3))); + + sumi1 = _mm256_add_epi32(sumi1, _mm256_madd_epi16(dot1, sc1)); + sumi2 = _mm256_add_epi32(sumi2, _mm256_madd_epi16(dot2, sc2)); + sumi1 = _mm256_add_epi32(sumi1, _mm256_madd_epi16(dot3, sc3)); + sumi2 = _mm256_add_epi32(sumi2, _mm256_madd_epi16(dot4, sc4)); + } + + accumf = _mm256_fmadd_ps(_mm256_set1_ps(d), _mm256_cvtepi32_ps(_mm256_add_epi32(sumi1, sumi2)), accumf); + + } + + *s = 0.125f * hsum_float_8(accumf); + +#elif defined(__AVX__) + const __m128i mone = _mm_set1_epi8(1); + static const char block_sign_shuffle_mask_1[32] = { + 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x02, 0x02, 0x02, 0x02, 0x02, 0x02, 0x02, 0x02, + 0x04, 0x04, 0x04, 0x04, 0x04, 0x04, 0x04, 0x04, 0x06, 0x06, 0x06, 0x06, 0x06, 0x06, 0x06, 0x06, + }; + static const char block_sign_shuffle_mask_2[32] = { + 0x08, 0x08, 0x08, 0x08, 0x08, 0x08, 0x08, 0x08, 0x0a, 0x0a, 0x0a, 0x0a, 0x0a, 0x0a, 0x0a, 0x0a, + 0x0c, 0x0c, 0x0c, 0x0c, 0x0c, 0x0c, 0x0c, 0x0c, 0x0e, 0x0e, 0x0e, 0x0e, 0x0e, 0x0e, 0x0e, 0x0e, + }; + static const uint8_t bit_selector_mask_bytes[32] = { + 0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80, 0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80, + 0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80, 0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80, + }; + + const __m128i bit_selector_mask_0 = _mm_loadu_si128((const __m128i*)bit_selector_mask_bytes); + const __m128i bit_selector_mask_1 = _mm_loadu_si128((const __m128i*)bit_selector_mask_bytes + 1); + const __m128i block_sign_shuffle_1_0 = _mm_loadu_si128((const __m128i*)block_sign_shuffle_mask_1); + const __m128i block_sign_shuffle_1_1 = _mm_loadu_si128((const __m128i*)block_sign_shuffle_mask_1 + 1); + const __m128i block_sign_shuffle_2_0 = _mm_loadu_si128((const __m128i*)block_sign_shuffle_mask_2); + const __m128i block_sign_shuffle_2_1 = _mm_loadu_si128((const __m128i*)block_sign_shuffle_mask_2 + 1); + + static const uint8_t k_bit_helper[32] = { + 0x00, 0x80, 0x80, 0x00, 0x80, 0x00, 0x00, 0x80, 0x80, 0x00, 0x00, 0x80, 0x00, 0x80, 0x80, 0x00, + 0x00, 0x80, 0x80, 0x00, 0x80, 0x00, 0x00, 0x80, 0x80, 0x00, 0x00, 0x80, 0x00, 0x80, 0x80, 0x00, + }; + const __m128i bit_helper_0 = _mm_loadu_si128((const __m128i*)k_bit_helper); + const __m128i bit_helper_1 = _mm_loadu_si128((const __m128i*)k_bit_helper + 1); + const __m128i m511 = _mm_set1_epi16(511); + const __m128i m4 = _mm_set1_epi8(0xf); + const __m128i m1 = _mm_set1_epi8(1); + + uint64_t aux64; + + // somewhat hacky, but gives a significant boost in performance + __m256i aux_gindex; + const uint16_t * gindex = (const uint16_t *)&aux_gindex; + + __m256 accumf = _mm256_setzero_ps(); + for (int i = 0; i < nb; ++i) { + const float d = GGML_CPU_FP16_TO_FP32(x[i].d) * y[i].d; + const uint16_t * GGML_RESTRICT q2 = x[i].qs; + const int8_t * GGML_RESTRICT q8 = y[i].qs; + + memcpy(&aux64, x[i].scales, 8); + __m128i stmp = _mm_set1_epi64x(aux64); + stmp = _mm_unpacklo_epi8(_mm_and_si128(stmp, m4), _mm_and_si128(_mm_srli_epi16(stmp, 4), m4)); + const __m128i scales = _mm_add_epi8(_mm_slli_epi16(stmp, 1), m1); + + __m128i sumi1_0 = _mm_setzero_si128(); + __m128i sumi1_1 = _mm_setzero_si128(); + __m128i sumi2_0 = _mm_setzero_si128(); + __m128i sumi2_1 = _mm_setzero_si128(); + for (int ib32 = 0; ib32 < QK_K/32; ib32 += 4) { + + const __m128i q2_data_0 = _mm_loadu_si128((const __m128i*)q2); + const __m128i q2_data_1 = _mm_loadu_si128((const __m128i*)q2 + 1); q2 += 16; + aux_gindex = MM256_SET_M128I(_mm_and_si128(q2_data_1, m511), _mm_and_si128(q2_data_0, m511)); + + const __m128i partial_sign_bits_0 = _mm_srli_epi16(q2_data_0, 9); + const __m128i partial_sign_bits_1 = _mm_srli_epi16(q2_data_1, 9); + const __m128i partial_sign_bits_upper_0 = _mm_srli_epi16(q2_data_0, 13); + const __m128i partial_sign_bits_upper_1 = _mm_srli_epi16(q2_data_1, 13); + const __m128i partial_sign_bits_for_counting_0 = _mm_xor_si128(partial_sign_bits_0, partial_sign_bits_upper_0); + const __m128i partial_sign_bits_for_counting_1 = _mm_xor_si128(partial_sign_bits_1, partial_sign_bits_upper_1); + + const __m128i odd_bits_0 = _mm_shuffle_epi8(bit_helper_0, partial_sign_bits_for_counting_0); + const __m128i odd_bits_1 = _mm_shuffle_epi8(bit_helper_1, partial_sign_bits_for_counting_1); + const __m128i full_sign_bits_0 = _mm_or_si128(partial_sign_bits_0, odd_bits_0); + const __m128i full_sign_bits_1 = _mm_or_si128(partial_sign_bits_1, odd_bits_1); + + const __m128i q8_1_0 = _mm_loadu_si128((const __m128i *)q8); q8 += 16; + const __m128i q8_1_1 = _mm_loadu_si128((const __m128i *)q8); q8 += 16; + const __m128i q8_2_0 = _mm_loadu_si128((const __m128i *)q8); q8 += 16; + const __m128i q8_2_1 = _mm_loadu_si128((const __m128i *)q8); q8 += 16; + const __m128i q8_3_0 = _mm_loadu_si128((const __m128i *)q8); q8 += 16; + const __m128i q8_3_1 = _mm_loadu_si128((const __m128i *)q8); q8 += 16; + const __m128i q8_4_0 = _mm_loadu_si128((const __m128i *)q8); q8 += 16; + const __m128i q8_4_1 = _mm_loadu_si128((const __m128i *)q8); q8 += 16; + + const __m128i q2_1_0 = _mm_set_epi64x(iq2xs_grid[gindex[1]], iq2xs_grid[gindex[0]]); + const __m128i q2_1_1 = _mm_set_epi64x(iq2xs_grid[gindex[3]], iq2xs_grid[gindex[2]]); + const __m128i q2_2_0 = _mm_set_epi64x(iq2xs_grid[gindex[5]], iq2xs_grid[gindex[4]]); + const __m128i q2_2_1 = _mm_set_epi64x(iq2xs_grid[gindex[7]], iq2xs_grid[gindex[6]]); + const __m128i q2_3_0 = _mm_set_epi64x(iq2xs_grid[gindex[9]], iq2xs_grid[gindex[8]]); + const __m128i q2_3_1 = _mm_set_epi64x(iq2xs_grid[gindex[11]], iq2xs_grid[gindex[10]]); + const __m128i q2_4_0 = _mm_set_epi64x(iq2xs_grid[gindex[13]], iq2xs_grid[gindex[12]]); + const __m128i q2_4_1 = _mm_set_epi64x(iq2xs_grid[gindex[15]], iq2xs_grid[gindex[14]]); + + // AVX2 full_signs_1 is full_sign_bits_0 here + // AVX2 full_signs_2 is full_sign_bits_1 here + __m128i signs_0, signs_1; + signs_0 = _mm_shuffle_epi8(full_sign_bits_0, block_sign_shuffle_1_0); + signs_1 = _mm_shuffle_epi8(full_sign_bits_0, block_sign_shuffle_1_1); + signs_0 = _mm_cmpeq_epi8(_mm_and_si128(signs_0, bit_selector_mask_0), bit_selector_mask_0); + signs_1 = _mm_cmpeq_epi8(_mm_and_si128(signs_1, bit_selector_mask_1), bit_selector_mask_1); + const __m128i q8s_1_0 = _mm_sign_epi8(q8_1_0, _mm_or_si128(signs_0, mone)); + const __m128i q8s_1_1 = _mm_sign_epi8(q8_1_1, _mm_or_si128(signs_1, mone)); + + signs_0 = _mm_shuffle_epi8(full_sign_bits_0, block_sign_shuffle_2_0); + signs_1 = _mm_shuffle_epi8(full_sign_bits_0, block_sign_shuffle_2_1); + signs_0 = _mm_cmpeq_epi8(_mm_and_si128(signs_0, bit_selector_mask_0), bit_selector_mask_0); + signs_1 = _mm_cmpeq_epi8(_mm_and_si128(signs_1, bit_selector_mask_1), bit_selector_mask_1); + const __m128i q8s_2_0 = _mm_sign_epi8(q8_2_0, _mm_or_si128(signs_0, mone)); + const __m128i q8s_2_1 = _mm_sign_epi8(q8_2_1, _mm_or_si128(signs_1, mone)); + + signs_0 = _mm_shuffle_epi8(full_sign_bits_1, block_sign_shuffle_1_0); + signs_1 = _mm_shuffle_epi8(full_sign_bits_1, block_sign_shuffle_1_1); + signs_0 = _mm_cmpeq_epi8(_mm_and_si128(signs_0, bit_selector_mask_0), bit_selector_mask_0); + signs_1 = _mm_cmpeq_epi8(_mm_and_si128(signs_1, bit_selector_mask_1), bit_selector_mask_1); + const __m128i q8s_3_0 = _mm_sign_epi8(q8_3_0, _mm_or_si128(signs_0, mone)); + const __m128i q8s_3_1 = _mm_sign_epi8(q8_3_1, _mm_or_si128(signs_1, mone)); + + signs_0 = _mm_shuffle_epi8(full_sign_bits_1, block_sign_shuffle_2_0); + signs_1 = _mm_shuffle_epi8(full_sign_bits_1, block_sign_shuffle_2_1); + signs_0 = _mm_cmpeq_epi8(_mm_and_si128(signs_0, bit_selector_mask_0), bit_selector_mask_0); + signs_1 = _mm_cmpeq_epi8(_mm_and_si128(signs_1, bit_selector_mask_1), bit_selector_mask_1); + const __m128i q8s_4_0 = _mm_sign_epi8(q8_4_0, _mm_or_si128(signs_0, mone)); + const __m128i q8s_4_1 = _mm_sign_epi8(q8_4_1, _mm_or_si128(signs_1, mone)); + + const __m128i dot1_0 = _mm_maddubs_epi16(q2_1_0, q8s_1_0); + const __m128i dot1_1 = _mm_maddubs_epi16(q2_1_1, q8s_1_1); + const __m128i dot2_0 = _mm_maddubs_epi16(q2_2_0, q8s_2_0); + const __m128i dot2_1 = _mm_maddubs_epi16(q2_2_1, q8s_2_1); + const __m128i dot3_0 = _mm_maddubs_epi16(q2_3_0, q8s_3_0); + const __m128i dot3_1 = _mm_maddubs_epi16(q2_3_1, q8s_3_1); + const __m128i dot4_0 = _mm_maddubs_epi16(q2_4_0, q8s_4_0); + const __m128i dot4_1 = _mm_maddubs_epi16(q2_4_1, q8s_4_1); + + __m128i sc_tmp = _mm_shuffle_epi8(scales, get_scale_shuffle(ib32+0)); + const __m128i sc1_0 = _mm_cvtepi8_epi16(sc_tmp); + const __m128i sc1_1 = _mm_cvtepi8_epi16(_mm_srli_si128(sc_tmp, 8)); + sc_tmp = _mm_shuffle_epi8(scales, get_scale_shuffle(ib32+1)); + const __m128i sc2_0 = _mm_cvtepi8_epi16(sc_tmp); + const __m128i sc2_1 = _mm_cvtepi8_epi16(_mm_srli_si128(sc_tmp, 8)); + sc_tmp = _mm_shuffle_epi8(scales, get_scale_shuffle(ib32+2)); + const __m128i sc3_0 = _mm_cvtepi8_epi16(sc_tmp); + const __m128i sc3_1 = _mm_cvtepi8_epi16(_mm_srli_si128(sc_tmp, 8)); + sc_tmp = _mm_shuffle_epi8(scales, get_scale_shuffle(ib32+3)); + const __m128i sc4_0 = _mm_cvtepi8_epi16(sc_tmp); + const __m128i sc4_1 = _mm_cvtepi8_epi16(_mm_srli_si128(sc_tmp, 8)); + + sumi1_0 = _mm_add_epi32(sumi1_0, _mm_madd_epi16(dot1_0, sc1_0)); + sumi1_1 = _mm_add_epi32(sumi1_1, _mm_madd_epi16(dot1_1, sc1_1)); + sumi2_0 = _mm_add_epi32(sumi2_0, _mm_madd_epi16(dot2_0, sc2_0)); + sumi2_1 = _mm_add_epi32(sumi2_1, _mm_madd_epi16(dot2_1, sc2_1)); + sumi1_0 = _mm_add_epi32(sumi1_0, _mm_madd_epi16(dot3_0, sc3_0)); + sumi1_1 = _mm_add_epi32(sumi1_1, _mm_madd_epi16(dot3_1, sc3_1)); + sumi2_0 = _mm_add_epi32(sumi2_0, _mm_madd_epi16(dot4_0, sc4_0)); + sumi2_1 = _mm_add_epi32(sumi2_1, _mm_madd_epi16(dot4_1, sc4_1)); + } + + accumf = _mm256_add_ps(_mm256_mul_ps(_mm256_set1_ps(d), _mm256_cvtepi32_ps(MM256_SET_M128I(_mm_add_epi32(sumi1_1, sumi2_1), _mm_add_epi32(sumi1_0, sumi2_0)))), accumf); + + } + + *s = 0.125f * hsum_float_8(accumf); + +#else + UNUSED(x); + UNUSED(y); + UNUSED(nb); + ggml_vec_dot_iq2_xs_q8_K_generic(n, s, bs, vx, bx, vy, by, nrc); +#endif +} + +void ggml_vec_dot_iq2_s_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) { + assert(n % QK_K == 0); + assert(nrc == 1); + UNUSED(nrc); + UNUSED(bx); + UNUSED(by); + UNUSED(bs); + + const block_iq2_s * GGML_RESTRICT x = vx; + const block_q8_K * GGML_RESTRICT y = vy; + + const int nb = n / QK_K; + +#if defined(__AVX2__) + + static const uint8_t k_mask1[32] = {0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01, + 0x02, 0x02, 0x02, 0x02, 0x02, 0x02, 0x02, 0x02, 0x03, 0x03, 0x03, 0x03, 0x03, 0x03, 0x03, 0x03 + }; + + static const uint8_t k_mask2[32] = {0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80, 0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80, + 0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80, 0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80, + }; + + const __m128i m4 = _mm_set1_epi8(0xf); + const __m128i m1 = _mm_set1_epi8(1); + + const __m256i mask1 = _mm256_loadu_si256((const __m256i*)k_mask1); + const __m256i mask2 = _mm256_loadu_si256((const __m256i*)k_mask2); + + uint64_t aux64; + + __m256 accumf = _mm256_setzero_ps(); + for (int i = 0; i < nb; ++i) { + const float d = GGML_CPU_FP16_TO_FP32(x[i].d) * y[i].d; + const uint8_t * GGML_RESTRICT qs = x[i].qs; + const uint8_t * GGML_RESTRICT qh = x[i].qh; + const uint16_t * GGML_RESTRICT signs = (const uint16_t *)(x[i].qs + QK_K/8); + const int8_t * GGML_RESTRICT q8 = y[i].qs; + + memcpy(&aux64, x[i].scales, 8); + const __m128i scales8 = _mm_add_epi8(_mm_slli_epi16(_mm_and_si128(_mm_set_epi64x(aux64 >> 4, aux64), m4), 1), m1); + const __m256i scales16 = _mm256_cvtepi8_epi16(scales8); // 0 2 4 6 8 10 12 14 1 3 5 7 9 11 13 15 + + __m256i sumi1 = _mm256_setzero_si256(); + __m256i sumi2 = _mm256_setzero_si256(); + for (int ib32 = 0; ib32 < QK_K/32; ib32 += 2) { + const __m256i q8_1 = _mm256_loadu_si256((const __m256i *)q8); q8 += 32; + const __m256i q8_2 = _mm256_loadu_si256((const __m256i *)q8); q8 += 32; + const __m256i q2_1 = _mm256_set_epi64x(iq2s_grid[qs[3] | ((qh[ib32+0] << 2) & 0x300)], + iq2s_grid[qs[2] | ((qh[ib32+0] << 4) & 0x300)], + iq2s_grid[qs[1] | ((qh[ib32+0] << 6) & 0x300)], + iq2s_grid[qs[0] | ((qh[ib32+0] << 8) & 0x300)]); + const __m256i q2_2 = _mm256_set_epi64x(iq2s_grid[qs[7] | ((qh[ib32+1] << 2) & 0x300)], + iq2s_grid[qs[6] | ((qh[ib32+1] << 4) & 0x300)], + iq2s_grid[qs[5] | ((qh[ib32+1] << 6) & 0x300)], + iq2s_grid[qs[4] | ((qh[ib32+1] << 8) & 0x300)]); + qs += 8; + + __m256i aux256 = _mm256_set1_epi32(signs[0] | ((uint32_t) signs[1] << 16)); + aux256 = _mm256_and_si256(_mm256_shuffle_epi8(aux256,mask1), mask2); + const __m256i s2_1 = _mm256_cmpeq_epi8(aux256, mask2); + const __m256i q8s_1 = _mm256_sub_epi8(_mm256_xor_si256(s2_1, q8_1), s2_1); + + aux256 = _mm256_set1_epi32(signs[2] | ((uint32_t) signs[3] << 16)); + aux256 = _mm256_and_si256(_mm256_shuffle_epi8(aux256,mask1), mask2); + const __m256i s2_2 = _mm256_cmpeq_epi8(aux256, mask2); + const __m256i q8s_2 = _mm256_sub_epi8(_mm256_xor_si256(s2_2, q8_2), s2_2); + + signs += 4; + + const __m256i dot1 = _mm256_maddubs_epi16(q2_1, q8s_1); // blocks 2*ib32+0, 2*ib32+1 + const __m256i dot2 = _mm256_maddubs_epi16(q2_2, q8s_2); // blocks 2*ib32+2, 2*ib32+3 + + const __m256i p1 = _mm256_madd_epi16(dot1, _mm256_shuffle_epi8(scales16, get_scale_shuffle_k4(ib32+0))); + const __m256i p2 = _mm256_madd_epi16(dot2, _mm256_shuffle_epi8(scales16, get_scale_shuffle_k4(ib32+1))); + sumi1 = _mm256_add_epi32(sumi1, p1); + sumi2 = _mm256_add_epi32(sumi2, p2); + } + + accumf = _mm256_fmadd_ps(_mm256_set1_ps(d), _mm256_cvtepi32_ps(_mm256_add_epi32(sumi1, sumi2)), accumf); + + } + + *s = 0.125f * hsum_float_8(accumf); + +#elif defined(__AVX__) + static const uint8_t k_mask1[32] = {0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01, + 0x02, 0x02, 0x02, 0x02, 0x02, 0x02, 0x02, 0x02, 0x03, 0x03, 0x03, 0x03, 0x03, 0x03, 0x03, 0x03 + }; + + static const uint8_t k_mask2[32] = {0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80, 0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80, + 0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80, 0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80, + }; + + const __m128i m4 = _mm_set1_epi8(0xf); + const __m128i m1 = _mm_set1_epi8(1); + + const __m128i mask1_0 = _mm_loadu_si128((const __m128i*)k_mask1); + const __m128i mask1_1 = _mm_loadu_si128((const __m128i*)k_mask1 + 1); + const __m128i mask2_0 = _mm_loadu_si128((const __m128i*)k_mask2); + const __m128i mask2_1 = _mm_loadu_si128((const __m128i*)k_mask2 + 1); + + uint64_t aux64; + + __m256 accumf = _mm256_setzero_ps(); + for (int i = 0; i < nb; ++i) { + const float d = GGML_CPU_FP16_TO_FP32(x[i].d) * y[i].d; + const uint8_t * GGML_RESTRICT qs = x[i].qs; + const uint8_t * GGML_RESTRICT qh = x[i].qh; + const uint16_t * GGML_RESTRICT signs = (const uint16_t *)(x[i].qs + QK_K/8); + const int8_t * GGML_RESTRICT q8 = y[i].qs; + + memcpy(&aux64, x[i].scales, 8); + const __m128i scales8 = _mm_add_epi8(_mm_slli_epi16(_mm_and_si128(_mm_set_epi64x(aux64 >> 4, aux64), m4), 1), m1); + const __m128i scales16_0 = _mm_cvtepi8_epi16(scales8); + const __m128i scales16_1 = _mm_cvtepi8_epi16(_mm_srli_si128(scales8, 8)); + + __m128i sumi1_0 = _mm_setzero_si128(); + __m128i sumi1_1 = _mm_setzero_si128(); + __m128i sumi2_0 = _mm_setzero_si128(); + __m128i sumi2_1 = _mm_setzero_si128(); + for (int ib32 = 0; ib32 < QK_K/32; ib32 += 2) { + const __m128i q8_1_0 = _mm_loadu_si128((const __m128i *)q8); q8 += 16; + const __m128i q8_1_1 = _mm_loadu_si128((const __m128i *)q8); q8 += 16; + const __m128i q8_2_0 = _mm_loadu_si128((const __m128i *)q8); q8 += 16; + const __m128i q8_2_1 = _mm_loadu_si128((const __m128i *)q8); q8 += 16; + const __m128i q2_1_0 = _mm_set_epi64x(iq2s_grid[qs[1] | ((qh[ib32+0] << 6) & 0x300)], + iq2s_grid[qs[0] | ((qh[ib32+0] << 8) & 0x300)]); + const __m128i q2_1_1 = _mm_set_epi64x(iq2s_grid[qs[3] | ((qh[ib32+0] << 2) & 0x300)], + iq2s_grid[qs[2] | ((qh[ib32+0] << 4) & 0x300)]); + const __m128i q2_2_0 = _mm_set_epi64x(iq2s_grid[qs[5] | ((qh[ib32+1] << 6) & 0x300)], + iq2s_grid[qs[4] | ((qh[ib32+1] << 8) & 0x300)]); + const __m128i q2_2_1 = _mm_set_epi64x(iq2s_grid[qs[7] | ((qh[ib32+1] << 2) & 0x300)], + iq2s_grid[qs[6] | ((qh[ib32+1] << 4) & 0x300)]); + qs += 8; + + __m128i aux128_0 = _mm_set1_epi32(signs[0] | ((uint32_t) signs[1] << 16)); + __m128i aux128_1 = aux128_0; + aux128_0 = _mm_and_si128(_mm_shuffle_epi8(aux128_0,mask1_0), mask2_0); + aux128_1 = _mm_and_si128(_mm_shuffle_epi8(aux128_1,mask1_1), mask2_1); + const __m128i s2_1_0 = _mm_cmpeq_epi8(aux128_0, mask2_0); + const __m128i s2_1_1 = _mm_cmpeq_epi8(aux128_1, mask2_1); + const __m128i q8s_1_0 = _mm_sub_epi8(_mm_xor_si128(s2_1_0, q8_1_0), s2_1_0); + const __m128i q8s_1_1 = _mm_sub_epi8(_mm_xor_si128(s2_1_1, q8_1_1), s2_1_1); + + aux128_0 = _mm_set1_epi32(signs[2] | ((uint32_t) signs[3] << 16)); + aux128_1 = aux128_0; + aux128_0 = _mm_and_si128(_mm_shuffle_epi8(aux128_0,mask1_0), mask2_0); + aux128_1 = _mm_and_si128(_mm_shuffle_epi8(aux128_1,mask1_1), mask2_1); + const __m128i s2_2_0 = _mm_cmpeq_epi8(aux128_0, mask2_0); + const __m128i s2_2_1 = _mm_cmpeq_epi8(aux128_1, mask2_1); + const __m128i q8s_2_0 = _mm_sub_epi8(_mm_xor_si128(s2_2_0, q8_2_0), s2_2_0); + const __m128i q8s_2_1 = _mm_sub_epi8(_mm_xor_si128(s2_2_1, q8_2_1), s2_2_1); + + signs += 4; + + const __m128i dot1_0 = _mm_maddubs_epi16(q2_1_0, q8s_1_0); + const __m128i dot1_1 = _mm_maddubs_epi16(q2_1_1, q8s_1_1); + const __m128i dot2_0 = _mm_maddubs_epi16(q2_2_0, q8s_2_0); + const __m128i dot2_1 = _mm_maddubs_epi16(q2_2_1, q8s_2_1); + + const __m128i p1_0 = _mm_madd_epi16(dot1_0, _mm_shuffle_epi8(scales16_0, _mm256_extractf128_si256(get_scale_shuffle_k4(ib32+0), 0))); + const __m128i p1_1 = _mm_madd_epi16(dot1_1, _mm_shuffle_epi8(scales16_1, _mm256_extractf128_si256(get_scale_shuffle_k4(ib32+0), 1))); + const __m128i p2_0 = _mm_madd_epi16(dot2_0, _mm_shuffle_epi8(scales16_0, _mm256_extractf128_si256(get_scale_shuffle_k4(ib32+1), 0))); + const __m128i p2_1 = _mm_madd_epi16(dot2_1, _mm_shuffle_epi8(scales16_1, _mm256_extractf128_si256(get_scale_shuffle_k4(ib32+1), 1))); + sumi1_0 = _mm_add_epi32(sumi1_0, p1_0); + sumi1_1 = _mm_add_epi32(sumi1_1, p1_1); + sumi2_0 = _mm_add_epi32(sumi2_0, p2_0); + sumi2_1 = _mm_add_epi32(sumi2_1, p2_1); + } + + accumf = _mm256_add_ps(_mm256_mul_ps(_mm256_set1_ps(d), _mm256_cvtepi32_ps(MM256_SET_M128I(_mm_add_epi32(sumi1_1, sumi2_1), _mm_add_epi32(sumi1_0, sumi2_0)))), accumf); + + } + + *s = 0.125f * hsum_float_8(accumf); + +#else + UNUSED(x); + UNUSED(y); + UNUSED(nb); + ggml_vec_dot_iq2_s_q8_K_generic(n, s, bs, vx, bx, vy, by, nrc); +#endif +} + +void ggml_vec_dot_iq3_xxs_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) { + assert(n % QK_K == 0); + assert(nrc == 1); + UNUSED(nrc); + UNUSED(bx); + UNUSED(by); + UNUSED(bs); + + const block_iq3_xxs * GGML_RESTRICT x = vx; + const block_q8_K * GGML_RESTRICT y = vy; + + const int nb = n / QK_K; + +#if defined(__AVX2__) + + const uint64_t * signs64 = (const uint64_t *)keven_signs_q2xs; + + uint32_t aux32[2]; + + __m256 accumf = _mm256_setzero_ps(); + for (int i = 0; i < nb; ++i) { + const float d = GGML_CPU_FP16_TO_FP32(x[i].d) * y[i].d; + const uint8_t * GGML_RESTRICT q3 = x[i].qs; + const uint8_t * GGML_RESTRICT gas = x[i].qs + QK_K/4; + const int8_t * GGML_RESTRICT q8 = y[i].qs; + __m256i sumi1 = _mm256_setzero_si256(); + __m256i sumi2 = _mm256_setzero_si256(); + for (int ib32 = 0; ib32 < QK_K/32; ib32 += 2) { + const __m256i q8_1 = _mm256_loadu_si256((const __m256i *)q8); q8 += 32; + const __m256i q8_2 = _mm256_loadu_si256((const __m256i *)q8); q8 += 32; + const __m256i q2_1 = _mm256_set_epi32(iq3xxs_grid[q3[7]], iq3xxs_grid[q3[6]], iq3xxs_grid[q3[5]], iq3xxs_grid[q3[4]], + iq3xxs_grid[q3[3]], iq3xxs_grid[q3[2]], iq3xxs_grid[q3[1]], iq3xxs_grid[q3[0]]); + q3 += 8; + const __m256i q2_2 = _mm256_set_epi32(iq3xxs_grid[q3[7]], iq3xxs_grid[q3[6]], iq3xxs_grid[q3[5]], iq3xxs_grid[q3[4]], + iq3xxs_grid[q3[3]], iq3xxs_grid[q3[2]], iq3xxs_grid[q3[1]], iq3xxs_grid[q3[0]]); + q3 += 8; + memcpy(aux32, gas, 8); gas += 8; + const __m256i s2_1 = _mm256_set_epi64x(signs64[(aux32[0] >> 21) & 127], signs64[(aux32[0] >> 14) & 127], + signs64[(aux32[0] >> 7) & 127], signs64[(aux32[0] >> 0) & 127]); + const __m256i s2_2 = _mm256_set_epi64x(signs64[(aux32[1] >> 21) & 127], signs64[(aux32[1] >> 14) & 127], + signs64[(aux32[1] >> 7) & 127], signs64[(aux32[1] >> 0) & 127]); + const __m256i q8s_1 = _mm256_sign_epi8(q8_1, s2_1); + const __m256i q8s_2 = _mm256_sign_epi8(q8_2, s2_2); + const __m256i dot1 = _mm256_maddubs_epi16(q2_1, q8s_1); + const __m256i dot2 = _mm256_maddubs_epi16(q2_2, q8s_2); + const uint16_t ls1 = aux32[0] >> 28; + const uint16_t ls2 = aux32[1] >> 28; + const __m256i p1 = _mm256_madd_epi16(dot1, _mm256_set1_epi16(2*ls1+1)); + const __m256i p2 = _mm256_madd_epi16(dot2, _mm256_set1_epi16(2*ls2+1)); + sumi1 = _mm256_add_epi32(sumi1, p1); + sumi2 = _mm256_add_epi32(sumi2, p2); + } + + accumf = _mm256_fmadd_ps(_mm256_set1_ps(d), _mm256_cvtepi32_ps(_mm256_add_epi32(sumi1, sumi2)), accumf); + + } + + *s = 0.25f * hsum_float_8(accumf); + +#elif defined(__AVX__) + const uint64_t * signs64 = (const uint64_t *)keven_signs_q2xs; + + uint32_t aux32[2]; + + __m256 accumf = _mm256_setzero_ps(); + for (int i = 0; i < nb; ++i) { + const float d = GGML_CPU_FP16_TO_FP32(x[i].d) * y[i].d; + const uint8_t * GGML_RESTRICT q3 = x[i].qs; + const uint8_t * GGML_RESTRICT gas = x[i].qs + QK_K/4; + const int8_t * GGML_RESTRICT q8 = y[i].qs; + __m128i sumi1_0 = _mm_setzero_si128(); + __m128i sumi1_1 = _mm_setzero_si128(); + __m128i sumi2_0 = _mm_setzero_si128(); + __m128i sumi2_1 = _mm_setzero_si128(); + for (int ib32 = 0; ib32 < QK_K/32; ib32 += 2) { + const __m128i q8_1_0 = _mm_loadu_si128((const __m128i *)q8); q8 += 16; + const __m128i q8_1_1 = _mm_loadu_si128((const __m128i *)q8); q8 += 16; + const __m128i q8_2_0 = _mm_loadu_si128((const __m128i *)q8); q8 += 16; + const __m128i q8_2_1 = _mm_loadu_si128((const __m128i *)q8); q8 += 16; + const __m128i q2_1_0 = _mm_set_epi32(iq3xxs_grid[q3[3]], iq3xxs_grid[q3[2]], iq3xxs_grid[q3[1]], iq3xxs_grid[q3[0]]); + const __m128i q2_1_1 = _mm_set_epi32(iq3xxs_grid[q3[7]], iq3xxs_grid[q3[6]], iq3xxs_grid[q3[5]], iq3xxs_grid[q3[4]]); + q3 += 8; + const __m128i q2_2_0 = _mm_set_epi32(iq3xxs_grid[q3[3]], iq3xxs_grid[q3[2]], iq3xxs_grid[q3[1]], iq3xxs_grid[q3[0]]); + const __m128i q2_2_1 = _mm_set_epi32(iq3xxs_grid[q3[7]], iq3xxs_grid[q3[6]], iq3xxs_grid[q3[5]], iq3xxs_grid[q3[4]]); + q3 += 8; + memcpy(aux32, gas, 8); gas += 8; + const __m128i s2_1_0 = _mm_set_epi64x(signs64[(aux32[0] >> 7) & 127], signs64[(aux32[0] >> 0) & 127]); + const __m128i s2_1_1 = _mm_set_epi64x(signs64[(aux32[0] >> 21) & 127], signs64[(aux32[0] >> 14) & 127]); + const __m128i s2_2_0 = _mm_set_epi64x(signs64[(aux32[1] >> 7) & 127], signs64[(aux32[1] >> 0) & 127]); + const __m128i s2_2_1 = _mm_set_epi64x(signs64[(aux32[1] >> 21) & 127], signs64[(aux32[1] >> 14) & 127]); + const __m128i q8s_1_0 = _mm_sign_epi8(q8_1_0, s2_1_0); + const __m128i q8s_1_1 = _mm_sign_epi8(q8_1_1, s2_1_1); + const __m128i q8s_2_0 = _mm_sign_epi8(q8_2_0, s2_2_0); + const __m128i q8s_2_1 = _mm_sign_epi8(q8_2_1, s2_2_1); + const __m128i dot1_0 = _mm_maddubs_epi16(q2_1_0, q8s_1_0); + const __m128i dot1_1 = _mm_maddubs_epi16(q2_1_1, q8s_1_1); + const __m128i dot2_0 = _mm_maddubs_epi16(q2_2_0, q8s_2_0); + const __m128i dot2_1 = _mm_maddubs_epi16(q2_2_1, q8s_2_1); + const uint16_t ls1 = aux32[0] >> 28; + const uint16_t ls2 = aux32[1] >> 28; + const __m128i p1_0 = _mm_madd_epi16(dot1_0, _mm_set1_epi16(2*ls1+1)); + const __m128i p1_1 = _mm_madd_epi16(dot1_1, _mm_set1_epi16(2*ls1+1)); + const __m128i p2_0 = _mm_madd_epi16(dot2_0, _mm_set1_epi16(2*ls2+1)); + const __m128i p2_1 = _mm_madd_epi16(dot2_1, _mm_set1_epi16(2*ls2+1)); + sumi1_0 = _mm_add_epi32(sumi1_0, p1_0); + sumi1_1 = _mm_add_epi32(sumi1_1, p1_1); + sumi2_0 = _mm_add_epi32(sumi2_0, p2_0); + sumi2_1 = _mm_add_epi32(sumi2_1, p2_1); + } + + accumf = _mm256_add_ps(_mm256_mul_ps(_mm256_set1_ps(d), _mm256_cvtepi32_ps(MM256_SET_M128I(_mm_add_epi32(sumi1_1, sumi2_1), _mm_add_epi32(sumi1_0, sumi2_0)))), accumf); + + } + + *s = 0.25f * hsum_float_8(accumf); + +#else + UNUSED(x); + UNUSED(y); + UNUSED(nb); + ggml_vec_dot_iq3_xxs_q8_K_generic(n, s, bs, vx, bx, vy, by, nrc); +#endif +} + +void ggml_vec_dot_iq3_s_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) { + assert(n % QK_K == 0); + assert(nrc == 1); + UNUSED(nrc); + UNUSED(bx); + UNUSED(by); + UNUSED(bs); + + const block_iq3_s * GGML_RESTRICT x = vx; + const block_q8_K * GGML_RESTRICT y = vy; + + const int nb = n / QK_K; + +#if defined(__AVX2__) + + static const uint8_t k_mask1[32] = {0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01, + 0x02, 0x02, 0x02, 0x02, 0x02, 0x02, 0x02, 0x02, 0x03, 0x03, 0x03, 0x03, 0x03, 0x03, 0x03, 0x03 + }; + + static const uint8_t k_mask2[32] = {0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80, 0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80, + 0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80, 0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80, + }; + + const __m256i mask1 = _mm256_loadu_si256((const __m256i*)k_mask1); + const __m256i mask2 = _mm256_loadu_si256((const __m256i*)k_mask2); + + const __m256i idx_shift = _mm256_set_epi32(1, 2, 3, 4, 5, 6, 7, 8); + const __m256i idx_mask = _mm256_set1_epi32(256); + + typedef union { + __m256i vec[2]; + uint32_t index[16]; + } index_t; + + index_t idx; + + __m256 accumf = _mm256_setzero_ps(); + for (int i = 0; i < nb; ++i) { + const float d = GGML_CPU_FP16_TO_FP32(x[i].d) * y[i].d; + const uint8_t * GGML_RESTRICT qs = x[i].qs; + const uint8_t * GGML_RESTRICT qh = x[i].qh; + const uint16_t * GGML_RESTRICT signs = (const uint16_t *)x[i].signs; + const int8_t * GGML_RESTRICT q8 = y[i].qs; + __m256i sumi1 = _mm256_setzero_si256(); + __m256i sumi2 = _mm256_setzero_si256(); + for (int ib32 = 0; ib32 < QK_K/32; ib32 += 2) { + const __m256i q8_1 = _mm256_loadu_si256((const __m256i *)q8); q8 += 32; + const __m256i q8_2 = _mm256_loadu_si256((const __m256i *)q8); q8 += 32; + const __m256i idx_l = _mm256_cvtepu8_epi16(_mm_loadu_si128((const __m128i *)qs)); qs += 16; + idx.vec[0] = _mm256_set1_epi32(qh[ib32+0]); + idx.vec[1] = _mm256_set1_epi32(qh[ib32+1]); + idx.vec[0] = _mm256_and_si256(_mm256_sllv_epi32(idx.vec[0], idx_shift), idx_mask); + idx.vec[1] = _mm256_and_si256(_mm256_sllv_epi32(idx.vec[1], idx_shift), idx_mask); + idx.vec[0] = _mm256_or_si256(idx.vec[0], _mm256_cvtepi16_epi32(_mm256_castsi256_si128(idx_l))); + idx.vec[1] = _mm256_or_si256(idx.vec[1], _mm256_cvtepi16_epi32(_mm256_extractf128_si256(idx_l, 1))); + + // At leat on my CPU (Ryzen 7950X), using _mm256_i32gather_epi32 is slower than _mm256_set_epi32. Strange. + //const __m256i q2_1 = _mm256_i32gather_epi32((const int *)iq3s_grid, idx.vec[0], 4); + //const __m256i q2_2 = _mm256_i32gather_epi32((const int *)iq3s_grid, idx.vec[1], 4); + const __m256i q2_1 = _mm256_set_epi32( + iq3s_grid[idx.index[7]], iq3s_grid[idx.index[6]], iq3s_grid[idx.index[5]], iq3s_grid[idx.index[4]], + iq3s_grid[idx.index[3]], iq3s_grid[idx.index[2]], iq3s_grid[idx.index[1]], iq3s_grid[idx.index[0]] + ); + const __m256i q2_2 = _mm256_set_epi32( + iq3s_grid[idx.index[15]], iq3s_grid[idx.index[14]], iq3s_grid[idx.index[13]], iq3s_grid[idx.index[12]], + iq3s_grid[idx.index[11]], iq3s_grid[idx.index[10]], iq3s_grid[idx.index[ 9]], iq3s_grid[idx.index[ 8]] + ); + + __m256i aux256 = _mm256_set1_epi32(signs[0] | (signs[1] << 16)); + aux256 = _mm256_and_si256(_mm256_shuffle_epi8(aux256,mask1), mask2); + const __m256i s2_1 = _mm256_cmpeq_epi8(aux256, mask2); + const __m256i q8s_1 = _mm256_sub_epi8(_mm256_xor_si256(s2_1, q8_1), s2_1); + + aux256 = _mm256_set1_epi32(signs[2] | (signs[3] << 16)); + aux256 = _mm256_and_si256(_mm256_shuffle_epi8(aux256,mask1), mask2); + const __m256i s2_2 = _mm256_cmpeq_epi8(aux256, mask2); + const __m256i q8s_2 = _mm256_sub_epi8(_mm256_xor_si256(s2_2, q8_2), s2_2); + + signs += 4; + + const __m256i dot1 = _mm256_maddubs_epi16(q2_1, q8s_1); + const __m256i dot2 = _mm256_maddubs_epi16(q2_2, q8s_2); + const uint16_t ls1 = x[i].scales[ib32/2] & 0xf; + const uint16_t ls2 = x[i].scales[ib32/2] >> 4; + const __m256i p1 = _mm256_madd_epi16(dot1, _mm256_set1_epi16(2*ls1+1)); + const __m256i p2 = _mm256_madd_epi16(dot2, _mm256_set1_epi16(2*ls2+1)); + sumi1 = _mm256_add_epi32(sumi1, p1); + sumi2 = _mm256_add_epi32(sumi2, p2); + } + + accumf = _mm256_fmadd_ps(_mm256_set1_ps(d), _mm256_cvtepi32_ps(_mm256_add_epi32(sumi1, sumi2)), accumf); + + } + + *s = hsum_float_8(accumf); + +#elif defined(__AVX__) + static const uint8_t k_mask1[32] = {0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01, + 0x02, 0x02, 0x02, 0x02, 0x02, 0x02, 0x02, 0x02, 0x03, 0x03, 0x03, 0x03, 0x03, 0x03, 0x03, 0x03 + }; + + static const uint8_t k_mask2[32] = {0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80, 0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80, + 0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80, 0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80, + }; + + const __m128i mask1_0 = _mm_loadu_si128((const __m128i*)k_mask1); + const __m128i mask1_1 = _mm_loadu_si128((const __m128i*)k_mask1 + 1); + const __m128i mask2_0 = _mm_loadu_si128((const __m128i*)k_mask2); + const __m128i mask2_1 = _mm_loadu_si128((const __m128i*)k_mask2 + 1); + + const __m128i idx_mul_0 = _mm_set_epi32(32, 64, 128, 256); + const __m128i idx_mul_1 = _mm_set_epi32(2, 4, 8, 16); + const __m128i idx_mask = _mm_set1_epi32(256); + + typedef union { + __m128i vec[4]; + uint32_t index[16]; + } index_t; + + index_t idx; + + __m256 accumf = _mm256_setzero_ps(); + for (int i = 0; i < nb; ++i) { + const float d = GGML_CPU_FP16_TO_FP32(x[i].d) * y[i].d; + const uint8_t * GGML_RESTRICT qs = x[i].qs; + const uint8_t * GGML_RESTRICT qh = x[i].qh; + const uint16_t * GGML_RESTRICT signs = (const uint16_t *)x[i].signs; + const int8_t * GGML_RESTRICT q8 = y[i].qs; + __m128i sumi1_0 = _mm_setzero_si128(); + __m128i sumi1_1 = _mm_setzero_si128(); + __m128i sumi2_0 = _mm_setzero_si128(); + __m128i sumi2_1 = _mm_setzero_si128(); + for (int ib32 = 0; ib32 < QK_K/32; ib32 += 2) { + const __m128i q8_1_0 = _mm_loadu_si128((const __m128i *)q8); q8 += 16; + const __m128i q8_1_1 = _mm_loadu_si128((const __m128i *)q8); q8 += 16; + const __m128i q8_2_0 = _mm_loadu_si128((const __m128i *)q8); q8 += 16; + const __m128i q8_2_1 = _mm_loadu_si128((const __m128i *)q8); q8 += 16; + const __m128i qs_tmp = _mm_loadu_si128((const __m128i *)qs); + const __m128i idx_l_0 = _mm_cvtepu8_epi16(qs_tmp); + const __m128i idx_l_1 = _mm_cvtepu8_epi16(_mm_srli_si128(qs_tmp, 8)); qs += 16; + idx.vec[0] = _mm_set1_epi32(qh[ib32+0]); + idx.vec[1] = idx.vec[0]; + idx.vec[2] = _mm_set1_epi32(qh[ib32+1]); + idx.vec[3] = idx.vec[2]; + + idx.vec[0] = _mm_and_si128(_mm_mullo_epi32(idx.vec[0], idx_mul_0), idx_mask); + idx.vec[1] = _mm_and_si128(_mm_mullo_epi32(idx.vec[1], idx_mul_1), idx_mask); + idx.vec[2] = _mm_and_si128(_mm_mullo_epi32(idx.vec[2], idx_mul_0), idx_mask); + idx.vec[3] = _mm_and_si128(_mm_mullo_epi32(idx.vec[3], idx_mul_1), idx_mask); + + idx.vec[0] = _mm_or_si128(idx.vec[0], _mm_cvtepi16_epi32(idx_l_0)); + idx.vec[1] = _mm_or_si128(idx.vec[1], _mm_cvtepi16_epi32(_mm_srli_si128(idx_l_0, 8))); + idx.vec[2] = _mm_or_si128(idx.vec[2], _mm_cvtepi16_epi32(idx_l_1)); + idx.vec[3] = _mm_or_si128(idx.vec[3], _mm_cvtepi16_epi32(_mm_srli_si128(idx_l_1, 8))); + + const __m128i q2_1_0 = _mm_set_epi32(iq3s_grid[idx.index[3]], iq3s_grid[idx.index[2]], iq3s_grid[idx.index[1]], iq3s_grid[idx.index[0]]); + const __m128i q2_1_1 = _mm_set_epi32(iq3s_grid[idx.index[7]], iq3s_grid[idx.index[6]], iq3s_grid[idx.index[5]], iq3s_grid[idx.index[4]]); + const __m128i q2_2_0 = _mm_set_epi32(iq3s_grid[idx.index[11]], iq3s_grid[idx.index[10]], iq3s_grid[idx.index[9]], iq3s_grid[idx.index[8]]); + const __m128i q2_2_1 = _mm_set_epi32(iq3s_grid[idx.index[15]], iq3s_grid[idx.index[14]], iq3s_grid[idx.index[13]], iq3s_grid[idx.index[12]]); + + __m128i aux128_0 = _mm_set1_epi32(signs[0] | (signs[1] << 16)); + __m128i aux128_1 = aux128_0; + aux128_0 = _mm_and_si128(_mm_shuffle_epi8(aux128_0,mask1_0), mask2_0); + aux128_1 = _mm_and_si128(_mm_shuffle_epi8(aux128_1,mask1_1), mask2_1); + const __m128i s2_1_0 = _mm_cmpeq_epi8(aux128_0, mask2_0); + const __m128i s2_1_1 = _mm_cmpeq_epi8(aux128_1, mask2_1); + const __m128i q8s_1_0 = _mm_sub_epi8(_mm_xor_si128(s2_1_0, q8_1_0), s2_1_0); + const __m128i q8s_1_1 = _mm_sub_epi8(_mm_xor_si128(s2_1_1, q8_1_1), s2_1_1); + + aux128_0 = _mm_set1_epi32(signs[2] | (signs[3] << 16)); + aux128_1 = aux128_0; + aux128_0 = _mm_and_si128(_mm_shuffle_epi8(aux128_0,mask1_0), mask2_0); + aux128_1 = _mm_and_si128(_mm_shuffle_epi8(aux128_1,mask1_1), mask2_1); + const __m128i s2_2_0 = _mm_cmpeq_epi8(aux128_0, mask2_0); + const __m128i s2_2_1 = _mm_cmpeq_epi8(aux128_1, mask2_1); + const __m128i q8s_2_0 = _mm_sub_epi8(_mm_xor_si128(s2_2_0, q8_2_0), s2_2_0); + const __m128i q8s_2_1 = _mm_sub_epi8(_mm_xor_si128(s2_2_1, q8_2_1), s2_2_1); + + signs += 4; + + const __m128i dot1_0 = _mm_maddubs_epi16(q2_1_0, q8s_1_0); + const __m128i dot1_1 = _mm_maddubs_epi16(q2_1_1, q8s_1_1); + const __m128i dot2_0 = _mm_maddubs_epi16(q2_2_0, q8s_2_0); + const __m128i dot2_1 = _mm_maddubs_epi16(q2_2_1, q8s_2_1); + const uint16_t ls1 = x[i].scales[ib32/2] & 0xf; + const uint16_t ls2 = x[i].scales[ib32/2] >> 4; + const __m128i p1_0 = _mm_madd_epi16(dot1_0, _mm_set1_epi16(2*ls1+1)); + const __m128i p1_1 = _mm_madd_epi16(dot1_1, _mm_set1_epi16(2*ls1+1)); + const __m128i p2_0 = _mm_madd_epi16(dot2_0, _mm_set1_epi16(2*ls2+1)); + const __m128i p2_1 = _mm_madd_epi16(dot2_1, _mm_set1_epi16(2*ls2+1)); + sumi1_0 = _mm_add_epi32(sumi1_0, p1_0); + sumi1_1 = _mm_add_epi32(sumi1_1, p1_1); + sumi2_0 = _mm_add_epi32(sumi2_0, p2_0); + sumi2_1 = _mm_add_epi32(sumi2_1, p2_1); + } + + accumf = _mm256_add_ps(_mm256_mul_ps(_mm256_set1_ps(d), _mm256_cvtepi32_ps(MM256_SET_M128I(_mm_add_epi32(sumi1_1, sumi2_1), _mm_add_epi32(sumi1_0, sumi2_0)))), accumf); + + } + + *s = hsum_float_8(accumf); + +#else + UNUSED(x); + UNUSED(y); + UNUSED(nb); + ggml_vec_dot_iq3_s_q8_K_generic(n, s, bs, vx, bx, vy, by, nrc); +#endif +} + +void ggml_vec_dot_iq1_s_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) { + assert(n % QK_K == 0); + assert(nrc == 1); + UNUSED(nrc); + UNUSED(bx); + UNUSED(by); + UNUSED(bs); + + const block_iq1_s * GGML_RESTRICT x = vx; + const block_q8_K * GGML_RESTRICT y = vy; + + const int nb = n / QK_K; + +#if defined __AVX2__ + + __m256 accum = _mm256_setzero_ps(); + float accum1 = 0; + for (int i = 0; i < nb; ++i) { + + const int8_t * q8 = y[i].qs; + const uint8_t * qs = x[i].qs; + const uint16_t * qh = x[i].qh; + + __m256i sumi = _mm256_setzero_si256(); + int sumi1 = 0; + for (int ib = 0; ib < QK_K/32; ib += 2) { +#ifdef __BMI2__ + const uint64_t packed_idx1 = _pdep_u64(*(const uint32_t *)qs, 0x00ff00ff00ff00ffULL) | _pdep_u64(qh[ib], 0x700070007000700ULL); + const uint64_t packed_idx2 = _pdep_u64(*(const uint32_t *)(qs + 4), 0x00ff00ff00ff00ffULL) | _pdep_u64(qh[ib + 1], 0x700070007000700ULL); + const uint16_t *idx1 = (const uint16_t *)(&packed_idx1); + const uint16_t *idx2 = (const uint16_t *)(&packed_idx2); + const __m256i q1b_1 = _mm256_set_epi64x(iq1s_grid[idx1[3]], iq1s_grid[idx1[2]], iq1s_grid[idx1[1]], iq1s_grid[idx1[0]]); + const __m256i q1b_2 = _mm256_set_epi64x(iq1s_grid[idx2[3]], iq1s_grid[idx2[2]], iq1s_grid[idx2[1]], iq1s_grid[idx2[0]]); +#else + const __m256i q1b_1 = _mm256_set_epi64x(iq1s_grid[qs[3] | ((qh[ib+0] >> 1) & 0x700)], iq1s_grid[qs[2] | ((qh[ib+0] << 2) & 0x700)], + iq1s_grid[qs[1] | ((qh[ib+0] << 5) & 0x700)], iq1s_grid[qs[0] | ((qh[ib+0] << 8) & 0x700)]); + const __m256i q1b_2 = _mm256_set_epi64x(iq1s_grid[qs[7] | ((qh[ib+1] >> 1) & 0x700)], iq1s_grid[qs[6] | ((qh[ib+1] << 2) & 0x700)], + iq1s_grid[qs[5] | ((qh[ib+1] << 5) & 0x700)], iq1s_grid[qs[4] | ((qh[ib+1] << 8) & 0x700)]); +#endif + qs += 8; + const __m256i q8b_1 = _mm256_loadu_si256((const __m256i*)q8); q8 += 32; + const __m256i q8b_2 = _mm256_loadu_si256((const __m256i*)q8); q8 += 32; + + const __m256i dot1 = mul_add_epi8(q1b_1, q8b_1); + const __m256i dot2 = mul_add_epi8(q1b_2, q8b_2); + const int16_t ls1 = 2*((qh[ib+0] >> 12) & 7) + 1; + const int16_t ls2 = 2*((qh[ib+1] >> 12) & 7) + 1; + const __m256i p1 = _mm256_madd_epi16(dot1, _mm256_set1_epi16(ls1)); + const __m256i p2 = _mm256_madd_epi16(dot2, _mm256_set1_epi16(ls2)); + + sumi = _mm256_add_epi32(sumi, _mm256_add_epi32(p1, p2)); + sumi1 += (y[i].bsums[2*ib+0] + y[i].bsums[2*ib+1]) * (qh[ib+0] & 0x8000 ? -1 : 1) * ls1 + + (y[i].bsums[2*ib+2] + y[i].bsums[2*ib+3]) * (qh[ib+1] & 0x8000 ? -1 : 1) * ls2; + } + + const float d = y[i].d * GGML_CPU_FP16_TO_FP32(x[i].d); + accum = _mm256_fmadd_ps(_mm256_set1_ps(d), _mm256_cvtepi32_ps(sumi), accum); + accum1 += d * sumi1; + + } + + *s = hsum_float_8(accum) + IQ1S_DELTA * accum1; + +#elif defined __AVX__ + __m256 accum = _mm256_setzero_ps(); + float accum1 = 0; + for (int i = 0; i < nb; ++i) { + + const int8_t * q8 = y[i].qs; + const uint8_t * qs = x[i].qs; + const uint16_t * qh = x[i].qh; + + __m128i sumi1_0 = _mm_setzero_si128(); + __m128i sumi1_1 = _mm_setzero_si128(); + int sumi1 = 0; + for (int ib = 0; ib < QK_K/32; ib += 2) { + const __m128i q1b_1_0 = _mm_set_epi64x(iq1s_grid[qs[1] | ((qh[ib+0] << 5) & 0x700)], iq1s_grid[qs[0] | ((qh[ib+0] << 8) & 0x700)]); + const __m128i q1b_1_1 = _mm_set_epi64x(iq1s_grid[qs[3] | ((qh[ib+0] >> 1) & 0x700)], iq1s_grid[qs[2] | ((qh[ib+0] << 2) & 0x700)]); + const __m128i q1b_2_0 = _mm_set_epi64x(iq1s_grid[qs[5] | ((qh[ib+1] << 5) & 0x700)], iq1s_grid[qs[4] | ((qh[ib+1] << 8) & 0x700)]); + const __m128i q1b_2_1 = _mm_set_epi64x(iq1s_grid[qs[7] | ((qh[ib+1] >> 1) & 0x700)], iq1s_grid[qs[6] | ((qh[ib+1] << 2) & 0x700)]); + qs += 8; + const __m128i q8b_1_0 = _mm_loadu_si128((const __m128i *)q8); q8 += 16; + const __m128i q8b_1_1 = _mm_loadu_si128((const __m128i *)q8); q8 += 16; + const __m128i q8b_2_0 = _mm_loadu_si128((const __m128i *)q8); q8 += 16; + const __m128i q8b_2_1 = _mm_loadu_si128((const __m128i *)q8); q8 += 16; + + const __m128i dot1_0 = mul_add_epi8_sse(q1b_1_0, q8b_1_0); + const __m128i dot1_1 = mul_add_epi8_sse(q1b_1_1, q8b_1_1); + const __m128i dot2_0 = mul_add_epi8_sse(q1b_2_0, q8b_2_0); + const __m128i dot2_1 = mul_add_epi8_sse(q1b_2_1, q8b_2_1); + const int16_t ls1 = 2*((qh[ib+0] >> 12) & 7) + 1; + const int16_t ls2 = 2*((qh[ib+1] >> 12) & 7) + 1; + const __m128i p1_0 = _mm_madd_epi16(dot1_0, _mm_set1_epi16(ls1)); + const __m128i p1_1 = _mm_madd_epi16(dot1_1, _mm_set1_epi16(ls1)); + const __m128i p2_0 = _mm_madd_epi16(dot2_0, _mm_set1_epi16(ls2)); + const __m128i p2_1 = _mm_madd_epi16(dot2_1, _mm_set1_epi16(ls2)); + + sumi1_0 = _mm_add_epi32(sumi1_0, _mm_add_epi32(p1_0, p2_0)); + sumi1_1 = _mm_add_epi32(sumi1_1, _mm_add_epi32(p1_1, p2_1)); + sumi1 += (y[i].bsums[2*ib+0] + y[i].bsums[2*ib+1]) * (qh[ib+0] & 0x8000 ? -1 : 1) * ls1 + + (y[i].bsums[2*ib+2] + y[i].bsums[2*ib+3]) * (qh[ib+1] & 0x8000 ? -1 : 1) * ls2; + } + + const float d = y[i].d * GGML_CPU_FP16_TO_FP32(x[i].d); + accum = _mm256_add_ps(_mm256_mul_ps(_mm256_set1_ps(d), _mm256_cvtepi32_ps(MM256_SET_M128I(sumi1_1, sumi1_0))), accum); + accum1 += d * sumi1; + + } + + *s = hsum_float_8(accum) + IQ1S_DELTA * accum1; + +#else + UNUSED(x); + UNUSED(y); + UNUSED(nb); + ggml_vec_dot_iq1_s_q8_K_generic(n, s, bs, vx, bx, vy, by, nrc); +#endif +} + +void ggml_vec_dot_iq1_m_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) { + assert(n % QK_K == 0); + assert(nrc == 1); + UNUSED(nrc); + UNUSED(bx); + UNUSED(by); + UNUSED(bs); + + const block_iq1_m * GGML_RESTRICT x = vx; + const block_q8_K * GGML_RESTRICT y = vy; + + const int nb = n / QK_K; + + iq1m_scale_t scale; + +#if defined __AVX2__ + + const __m256i mask = _mm256_set1_epi16(0x7); + const __m256i mone = _mm256_set1_epi16(1); + const __m256i mone8 = _mm256_set1_epi8(1); + const __m256i mtwo8 = _mm256_set1_epi8(2); + // VPSHUFB cannot cross 128-bit lanes so odd shifts go to upper half. + const __m256i scales_shift = _mm256_set_epi64x(9, 3, 6, 0); + + __m256 accum1 = _mm256_setzero_ps(); + __m256 accum2 = _mm256_setzero_ps(); + for (int i = 0; i < nb; ++i) { + + const int8_t * q8 = y[i].qs; + const uint8_t * qs = x[i].qs; + const uint8_t * qh = x[i].qh; + const uint16_t * sc = (const uint16_t *)x[i].scales; + + scale.u16 = (sc[0] >> 12) | ((sc[1] >> 8) & 0x00f0) | ((sc[2] >> 4) & 0x0f00) | (sc[3] & 0xf000); + // Extract 3-bit scales (16 values) + __m256i scales = _mm256_set1_epi64x(*(const uint64_t*)sc); + scales = _mm256_srlv_epi64(scales, scales_shift); + scales = _mm256_add_epi16(_mm256_slli_epi16(_mm256_and_si256(scales, mask), 1), mone); + + // Indices to repeat each scale 8 times. + __m256i scales_idx1 = _mm256_set1_epi16(0x0100); + __m256i scales_idx2 = _mm256_add_epi8(scales_idx1, _mm256_set1_epi8(8)); + + __m256i sumi1 = _mm256_setzero_si256(); + __m256i sumi2 = _mm256_setzero_si256(); + for (int ib = 0; ib < QK_K/32; ib += 2) { +#ifdef __BMI2__ + const uint64_t packed_idx1 = _pdep_u64(*(const uint32_t *)qs, 0x00ff00ff00ff00ffULL) + | _pdep_u64(*(const uint16_t*)(qh) & 0x7777, 0xf000f000f000f00ULL); + const uint64_t packed_idx2 = _pdep_u64(*(const uint32_t *)(qs + 4), 0x00ff00ff00ff00ffULL) + | _pdep_u64(*(const uint16_t*)(qh + 2) & 0x7777, 0xf000f000f000f00ULL); + const uint16_t *idx1 = (const uint16_t *)(&packed_idx1); + const uint16_t *idx2 = (const uint16_t *)(&packed_idx2); + const __m256i q1b_1 = _mm256_set_epi64x(iq1s_grid[idx1[3]], iq1s_grid[idx1[2]], iq1s_grid[idx1[1]], iq1s_grid[idx1[0]]); + const __m256i q1b_2 = _mm256_set_epi64x(iq1s_grid[idx2[3]], iq1s_grid[idx2[2]], iq1s_grid[idx2[1]], iq1s_grid[idx2[0]]); + + // Convert signs to bytes 0x81 (negative) or 0x01 (positive) + const uint64_t delta_sign = _pdep_u64(*(const uint32_t*)(qh) & 0x88888888, 0xf0f0f0f0f0f0f0f0ULL); + const __m256i delta1 = _mm256_or_si256(mone8, _mm256_cvtepi8_epi64(_mm_set1_epi32(delta_sign))); + const __m256i delta2 = _mm256_or_si256(mone8, _mm256_cvtepi8_epi64(_mm_set1_epi32(delta_sign >> 32))); +#else + const __m256i q1b_1 = _mm256_set_epi64x( + iq1s_grid[qs[3] | (((uint16_t)qh[1] << 4) & 0x700)], iq1s_grid[qs[2] | (((uint16_t)qh[1] << 8) & 0x700)], + iq1s_grid[qs[1] | (((uint16_t)qh[0] << 4) & 0x700)], iq1s_grid[qs[0] | (((uint16_t)qh[0] << 8) & 0x700)] + ); + const __m256i q1b_2 = _mm256_set_epi64x( + iq1s_grid[qs[7] | (((uint16_t)qh[3] << 4) & 0x700)], iq1s_grid[qs[6] | (((uint16_t)qh[3] << 8) & 0x700)], + iq1s_grid[qs[5] | (((uint16_t)qh[2] << 4) & 0x700)], iq1s_grid[qs[4] | (((uint16_t)qh[2] << 8) & 0x700)] + ); + + const __m256i delta1 = _mm256_set_epi64x(qh[1] & 0x80 ? 0xffffffffffffffff : 0x0101010101010101, + qh[1] & 0x08 ? 0xffffffffffffffff : 0x0101010101010101, + qh[0] & 0x80 ? 0xffffffffffffffff : 0x0101010101010101, + qh[0] & 0x08 ? 0xffffffffffffffff : 0x0101010101010101); + const __m256i delta2 = _mm256_set_epi64x(qh[3] & 0x80 ? 0xffffffffffffffff : 0x0101010101010101, + qh[3] & 0x08 ? 0xffffffffffffffff : 0x0101010101010101, + qh[2] & 0x80 ? 0xffffffffffffffff : 0x0101010101010101, + qh[2] & 0x08 ? 0xffffffffffffffff : 0x0101010101010101); +#endif + const __m256i q8b_1 = _mm256_loadu_si256((const __m256i*)q8); q8 += 32; + const __m256i q8b_2 = _mm256_loadu_si256((const __m256i*)q8); q8 += 32; + + const __m256i dot1 = mul_add_epi8(q1b_1, q8b_1); + const __m256i dot2 = mul_add_epi8(q1b_2, q8b_2); + const __m256i dot3 = _mm256_maddubs_epi16(mone8, _mm256_sign_epi8(q8b_1, delta1)); + const __m256i dot4 = _mm256_maddubs_epi16(mone8, _mm256_sign_epi8(q8b_2, delta2)); + + __m256i scale1 = _mm256_shuffle_epi8(scales, scales_idx1); + __m256i scale2 = _mm256_shuffle_epi8(scales, scales_idx2); + + scales_idx1 = _mm256_add_epi8(scales_idx1, mtwo8); + scales_idx2 = _mm256_add_epi8(scales_idx2, mtwo8); + + const __m256i p1 = _mm256_madd_epi16(dot1, scale1); + const __m256i p2 = _mm256_madd_epi16(dot2, scale2); + const __m256i p3 = _mm256_madd_epi16(dot3, scale1); + const __m256i p4 = _mm256_madd_epi16(dot4, scale2); + + sumi1 = _mm256_add_epi32(sumi1, _mm256_add_epi32(p1, p2)); + sumi2 = _mm256_add_epi32(sumi2, _mm256_add_epi32(p3, p4)); + + qs += 8; qh += 4; + } + + const __m256 d = _mm256_set1_ps(y[i].d * GGML_CPU_FP16_TO_FP32(scale.f16)); + + accum1 = _mm256_fmadd_ps(d, _mm256_cvtepi32_ps(sumi1), accum1); + accum2 = _mm256_fmadd_ps(d, _mm256_cvtepi32_ps(sumi2), accum2); + } + + *s = hsum_float_8(accum1) + IQ1M_DELTA * hsum_float_8(accum2); + +#elif defined __AVX__ + const __m128i mask = _mm_set1_epi16(0x7); + const __m128i mone = _mm_set1_epi16(1); + + __m256 accum1 = _mm256_setzero_ps(); + __m256 accum2 = _mm256_setzero_ps(); + for (int i = 0; i < nb; ++i) { + + const int8_t * q8 = y[i].qs; + const uint8_t * qs = x[i].qs; + const uint8_t * qh = x[i].qh; + const uint16_t * sc = (const uint16_t *)x[i].scales; + + scale.u16 = (sc[0] >> 12) | ((sc[1] >> 8) & 0x00f0) | ((sc[2] >> 4) & 0x0f00) | (sc[3] & 0xf000); + + __m128i sumi1_0 = _mm_setzero_si128(); + __m128i sumi1_1 = _mm_setzero_si128(); + __m128i sumi2_0 = _mm_setzero_si128(); + __m128i sumi2_1 = _mm_setzero_si128(); + for (int ib = 0; ib < QK_K/32; ib += 2) { + const __m128i q1b_1_0 = _mm_set_epi64x( + iq1s_grid[qs[1] | (((uint16_t)qh[0] << 4) & 0x700)], iq1s_grid[qs[0] | (((uint16_t)qh[0] << 8) & 0x700)]); + const __m128i q1b_1_1 = _mm_set_epi64x( + iq1s_grid[qs[3] | (((uint16_t)qh[1] << 4) & 0x700)], iq1s_grid[qs[2] | (((uint16_t)qh[1] << 8) & 0x700)]); + const __m128i q1b_2_0 = _mm_set_epi64x( + iq1s_grid[qs[5] | (((uint16_t)qh[2] << 4) & 0x700)], iq1s_grid[qs[4] | (((uint16_t)qh[2] << 8) & 0x700)]); + const __m128i q1b_2_1 = _mm_set_epi64x( + iq1s_grid[qs[7] | (((uint16_t)qh[3] << 4) & 0x700)], iq1s_grid[qs[6] | (((uint16_t)qh[3] << 8) & 0x700)]); + const __m128i q8b_1_0 = _mm_loadu_si128((const __m128i *)q8); q8 += 16; + const __m128i q8b_1_1 = _mm_loadu_si128((const __m128i *)q8); q8 += 16; + const __m128i q8b_2_0 = _mm_loadu_si128((const __m128i *)q8); q8 += 16; + const __m128i q8b_2_1 = _mm_loadu_si128((const __m128i *)q8); q8 += 16; + + const __m128i dot1_0 = mul_add_epi8_sse(q1b_1_0, q8b_1_0); + const __m128i dot1_1 = mul_add_epi8_sse(q1b_1_1, q8b_1_1); + const __m128i dot2_0 = mul_add_epi8_sse(q1b_2_0, q8b_2_0); + const __m128i dot2_1 = mul_add_epi8_sse(q1b_2_1, q8b_2_1); + + const __m128i delta1_0 = _mm_set_epi64x(qh[0] & 0x80 ? 0xffffffffffffffff : 0x0101010101010101, + qh[0] & 0x08 ? 0xffffffffffffffff : 0x0101010101010101); + const __m128i delta1_1 = _mm_set_epi64x(qh[1] & 0x80 ? 0xffffffffffffffff : 0x0101010101010101, + qh[1] & 0x08 ? 0xffffffffffffffff : 0x0101010101010101); + const __m128i delta2_0 = _mm_set_epi64x(qh[2] & 0x80 ? 0xffffffffffffffff : 0x0101010101010101, + qh[2] & 0x08 ? 0xffffffffffffffff : 0x0101010101010101); + const __m128i delta2_1 = _mm_set_epi64x(qh[3] & 0x80 ? 0xffffffffffffffff : 0x0101010101010101, + qh[3] & 0x08 ? 0xffffffffffffffff : 0x0101010101010101); + + const __m128i dot3_0 = mul_add_epi8_sse(delta1_0, q8b_1_0); + const __m128i dot3_1 = mul_add_epi8_sse(delta1_1, q8b_1_1); + const __m128i dot4_0 = mul_add_epi8_sse(delta2_0, q8b_2_0); + const __m128i dot4_1 = mul_add_epi8_sse(delta2_1, q8b_2_1); + + __m128i scale1_0 = _mm_set1_epi16(sc[ib/2] >> 0); + __m128i scale1_1 = _mm_set1_epi16(sc[ib/2] >> 3); + __m128i scale2_0 = _mm_set1_epi16(sc[ib/2] >> 6); + __m128i scale2_1 = _mm_set1_epi16(sc[ib/2] >> 9); + + scale1_0 = _mm_add_epi16(_mm_slli_epi16(_mm_and_si128(scale1_0, mask), 1), mone); + scale1_1 = _mm_add_epi16(_mm_slli_epi16(_mm_and_si128(scale1_1, mask), 1), mone); + scale2_0 = _mm_add_epi16(_mm_slli_epi16(_mm_and_si128(scale2_0, mask), 1), mone); + scale2_1 = _mm_add_epi16(_mm_slli_epi16(_mm_and_si128(scale2_1, mask), 1), mone); + const __m128i p1_0 = _mm_madd_epi16(dot1_0, scale1_0); + const __m128i p1_1 = _mm_madd_epi16(dot1_1, scale1_1); + const __m128i p2_0 = _mm_madd_epi16(dot2_0, scale2_0); + const __m128i p2_1 = _mm_madd_epi16(dot2_1, scale2_1); + const __m128i p3_0 = _mm_madd_epi16(dot3_0, scale1_0); + const __m128i p3_1 = _mm_madd_epi16(dot3_1, scale1_1); + const __m128i p4_0 = _mm_madd_epi16(dot4_0, scale2_0); + const __m128i p4_1 = _mm_madd_epi16(dot4_1, scale2_1); + + sumi1_0 = _mm_add_epi32(sumi1_0, _mm_add_epi32(p1_0, p2_0)); + sumi1_1 = _mm_add_epi32(sumi1_1, _mm_add_epi32(p1_1, p2_1)); + sumi2_0 = _mm_add_epi32(sumi2_0, _mm_add_epi32(p3_0, p4_0)); + sumi2_1 = _mm_add_epi32(sumi2_1, _mm_add_epi32(p3_1, p4_1)); + + qs += 8; qh += 4; + } + + const __m256 d = _mm256_set1_ps(y[i].d * GGML_CPU_FP16_TO_FP32(scale.f16)); + + accum1 = _mm256_add_ps(_mm256_mul_ps(d, _mm256_cvtepi32_ps(MM256_SET_M128I(sumi1_1, sumi1_0))), accum1); + accum2 = _mm256_add_ps(_mm256_mul_ps(d, _mm256_cvtepi32_ps(MM256_SET_M128I(sumi2_1, sumi2_0))), accum2); + } + + *s = hsum_float_8(accum1) + IQ1M_DELTA * hsum_float_8(accum2); + +#else + UNUSED(x); + UNUSED(y); + UNUSED(nb); + UNUSED(scale); + ggml_vec_dot_iq1_m_q8_K_generic(n, s, bs, vx, bx, vy, by, nrc); +#endif +} + +void ggml_vec_dot_iq4_nl_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) { + assert(nrc == 1); + UNUSED(nrc); + UNUSED(bx); + UNUSED(by); + UNUSED(bs); + assert(n % QK4_NL == 0); + static_assert(QK4_NL == QK8_0, "QK4_NL and QK8_0 must be the same"); + + const block_iq4_nl * GGML_RESTRICT x = vx; + const block_q8_0 * GGML_RESTRICT y = vy; + + const int nb = n / QK4_NL; + + int ib = 0; + float sumf = 0; + +#if defined __AVX2__ + + const __m128i values128 = _mm_loadu_si128((const __m128i*)kvalues_iq4nl); + const __m128i m4b = _mm_set1_epi8(0x0f); + const __m256i mone = _mm256_set1_epi16(1); + + __m256 accum1 = _mm256_setzero_ps(); + __m256 accum2 = _mm256_setzero_ps(); + for (; ib + 1 < nb; ib += 2) { + const __m128i q4bits_1 = _mm_loadu_si128((const __m128i*)x[ib + 0].qs); + const __m128i q4bits_2 = _mm_loadu_si128((const __m128i*)x[ib + 1].qs); + const __m256i q8b_1 = _mm256_loadu_si256((const __m256i *)y[ib + 0].qs); + const __m256i q8b_2 = _mm256_loadu_si256((const __m256i *)y[ib + 1].qs); + const __m256i q4b_1 = MM256_SET_M128I(_mm_shuffle_epi8(values128, _mm_and_si128(_mm_srli_epi16(q4bits_1, 4), m4b)), + _mm_shuffle_epi8(values128, _mm_and_si128(q4bits_1, m4b))); + const __m256i q4b_2 = MM256_SET_M128I(_mm_shuffle_epi8(values128, _mm_and_si128(_mm_srli_epi16(q4bits_2, 4), m4b)), + _mm_shuffle_epi8(values128, _mm_and_si128(q4bits_2, m4b))); + const __m256i p16_1 = mul_add_epi8(q4b_1, q8b_1); + const __m256i p16_2 = mul_add_epi8(q4b_2, q8b_2); + const __m256i p_1 = _mm256_madd_epi16(p16_1, mone); + const __m256i p_2 = _mm256_madd_epi16(p16_2, mone); + accum1 = _mm256_fmadd_ps(_mm256_set1_ps(GGML_CPU_FP16_TO_FP32(y[ib + 0].d)*GGML_CPU_FP16_TO_FP32(x[ib + 0].d)), + _mm256_cvtepi32_ps(p_1), accum1); + accum2 = _mm256_fmadd_ps(_mm256_set1_ps(GGML_CPU_FP16_TO_FP32(y[ib + 1].d)*GGML_CPU_FP16_TO_FP32(x[ib + 1].d)), + _mm256_cvtepi32_ps(p_2), accum2); + } + + sumf = hsum_float_8(_mm256_add_ps(accum1, accum2)); + +#elif defined __AVX__ + const __m128i values128 = _mm_loadu_si128((const __m128i*)kvalues_iq4nl); + const __m128i m4b = _mm_set1_epi8(0x0f); + + __m256 accum = _mm256_setzero_ps(); + for (; ib + 1 < nb; ib += 2) { + const __m128i q4bits_1 = _mm_loadu_si128((const __m128i *)x[ib + 0].qs); + const __m128i q4bits_2 = _mm_loadu_si128((const __m128i *)x[ib + 1].qs); + const __m128i q8b_1_0 = _mm_loadu_si128((const __m128i *)y[ib + 0].qs); + const __m128i q8b_1_1 = _mm_loadu_si128((const __m128i *)y[ib + 0].qs + 1); + const __m128i q8b_2_0 = _mm_loadu_si128((const __m128i *)y[ib + 1].qs); + const __m128i q8b_2_1 = _mm_loadu_si128((const __m128i *)y[ib + 1].qs + 1); + + const __m128i q4b_1_0 = _mm_shuffle_epi8(values128, _mm_and_si128(q4bits_1, m4b)); + const __m128i q4b_1_1 = _mm_shuffle_epi8(values128, _mm_and_si128(_mm_srli_epi16(q4bits_1, 4), m4b)); + const __m128i q4b_2_0 = _mm_shuffle_epi8(values128, _mm_and_si128(q4bits_2, m4b)); + const __m128i q4b_2_1 = _mm_shuffle_epi8(values128, _mm_and_si128(_mm_srli_epi16(q4bits_2, 4), m4b)); + + const __m256 p = mul_sum_i8_quad_float(q4b_1_0, q4b_1_1, q4b_2_0, q4b_2_1, q8b_1_0, q8b_1_1, q8b_2_0, q8b_2_1); + const __m256 deltas = quad_fp16_delta_float(x[ib].d, y[ib].d, x[ib + 1].d, y[ib + 1].d); + accum = _mm256_add_ps(_mm256_mul_ps(deltas, p), accum); + } + + sumf = hsum_float_8(accum); + +#endif + for (; ib < nb; ++ib) { + const float d = GGML_CPU_FP16_TO_FP32(y[ib].d)*GGML_CPU_FP16_TO_FP32(x[ib].d); + int sumi1 = 0, sumi2 = 0; + for (int j = 0; j < QK4_NL/2; ++j) { + sumi1 += y[ib].qs[j+ 0] * kvalues_iq4nl[x[ib].qs[j] & 0xf]; + sumi2 += y[ib].qs[j+QK4_NL/2] * kvalues_iq4nl[x[ib].qs[j] >> 4]; + } + sumf += d * (sumi1 + sumi2); + } + *s = sumf; +} + +void ggml_vec_dot_iq4_xs_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) { + assert(nrc == 1); + UNUSED(nrc); + UNUSED(bx); + UNUSED(by); + UNUSED(bs); + assert(n % QK_K == 0); + + const block_iq4_xs * GGML_RESTRICT x = vx; + const block_q8_K * GGML_RESTRICT y = vy; + + const int nb = n / QK_K; + +#if defined __AVX2__ + + const __m128i values128 = _mm_loadu_si128((const __m128i*)kvalues_iq4nl); + const __m128i m4b = _mm_set1_epi8(0x0f); + + __m256 accum = _mm256_setzero_ps(); + for (int ibl = 0; ibl < nb; ++ibl) { + const uint8_t * qs = x[ibl].qs; + const int8_t * q8 = y[ibl].qs; + uint16_t sh = x[ibl].scales_h; + __m256i sumi1 = _mm256_setzero_si256(); + __m256i sumi2 = _mm256_setzero_si256(); + for (int ib = 0; ib < QK_K/32; ib += 2) { + const __m128i q4bits_1 = _mm_loadu_si128((const __m128i*)qs); qs += 16; + const __m128i q4bits_2 = _mm_loadu_si128((const __m128i*)qs); qs += 16; + const __m256i q8b_1 = _mm256_loadu_si256((const __m256i *)q8); q8 += 32; + const __m256i q8b_2 = _mm256_loadu_si256((const __m256i *)q8); q8 += 32; + const __m256i q4b_1 = MM256_SET_M128I(_mm_shuffle_epi8(values128, _mm_and_si128(_mm_srli_epi16(q4bits_1, 4), m4b)), + _mm_shuffle_epi8(values128, _mm_and_si128(q4bits_1, m4b))); + const __m256i q4b_2 = MM256_SET_M128I(_mm_shuffle_epi8(values128, _mm_and_si128(_mm_srli_epi16(q4bits_2, 4), m4b)), + _mm_shuffle_epi8(values128, _mm_and_si128(q4bits_2, m4b))); + const __m256i p16_1 = mul_add_epi8(q4b_1, q8b_1); + const __m256i p16_2 = mul_add_epi8(q4b_2, q8b_2); + const int16_t ls1 = ((x[ibl].scales_l[ib/2] & 0xf) | ((sh << 4) & 0x30)) - 32; + const int16_t ls2 = ((x[ibl].scales_l[ib/2] >> 4) | ((sh << 2) & 0x30)) - 32; + sh >>= 4; + const __m256i p_1 = _mm256_madd_epi16(p16_1, _mm256_set1_epi16(ls1)); + const __m256i p_2 = _mm256_madd_epi16(p16_2, _mm256_set1_epi16(ls2)); + sumi1 = _mm256_add_epi32(p_1, sumi1); + sumi2 = _mm256_add_epi32(p_2, sumi2); + } + accum = _mm256_fmadd_ps(_mm256_set1_ps(GGML_CPU_FP16_TO_FP32(x[ibl].d)*y[ibl].d), + _mm256_cvtepi32_ps(_mm256_add_epi32(sumi1, sumi2)), accum); + } + + *s = hsum_float_8(accum); + +#elif defined __AVX__ + const __m128i values128 = _mm_loadu_si128((const __m128i*)kvalues_iq4nl); + const __m128i m4b = _mm_set1_epi8(0x0f); + + __m256 accum = _mm256_setzero_ps(); + for (int ibl = 0; ibl < nb; ++ibl) { + const uint8_t * qs = x[ibl].qs; + const int8_t * q8 = y[ibl].qs; + uint16_t sh = x[ibl].scales_h; + __m128i sumi1_0 = _mm_setzero_si128(); + __m128i sumi1_1 = _mm_setzero_si128(); + __m128i sumi2_0 = _mm_setzero_si128(); + __m128i sumi2_1 = _mm_setzero_si128(); + for (int ib = 0; ib < QK_K/32; ib += 2) { + const __m128i q4bits_1 = _mm_loadu_si128((const __m128i *)qs); qs += 16; + const __m128i q4bits_2 = _mm_loadu_si128((const __m128i *)qs); qs += 16; + const __m128i q8b_1_0 = _mm_loadu_si128((const __m128i *)q8); q8 += 16; + const __m128i q8b_1_1 = _mm_loadu_si128((const __m128i *)q8); q8 += 16; + const __m128i q8b_2_0 = _mm_loadu_si128((const __m128i *)q8); q8 += 16; + const __m128i q8b_2_1 = _mm_loadu_si128((const __m128i *)q8); q8 += 16; + const __m128i q4b_1_0 = _mm_shuffle_epi8(values128, _mm_and_si128(q4bits_1, m4b)); + const __m128i q4b_1_1 = _mm_shuffle_epi8(values128, _mm_and_si128(_mm_srli_epi16(q4bits_1, 4), m4b)); + const __m128i q4b_2_0 = _mm_shuffle_epi8(values128, _mm_and_si128(q4bits_2, m4b)); + const __m128i q4b_2_1 = _mm_shuffle_epi8(values128, _mm_and_si128(_mm_srli_epi16(q4bits_2, 4), m4b)); + const __m128i p16_1_0 = mul_add_epi8_sse(q4b_1_0, q8b_1_0); + const __m128i p16_1_1 = mul_add_epi8_sse(q4b_1_1, q8b_1_1); + const __m128i p16_2_0 = mul_add_epi8_sse(q4b_2_0, q8b_2_0); + const __m128i p16_2_1 = mul_add_epi8_sse(q4b_2_1, q8b_2_1); + const int16_t ls1 = ((x[ibl].scales_l[ib/2] & 0xf) | ((sh << 4) & 0x30)) - 32; + const int16_t ls2 = ((x[ibl].scales_l[ib/2] >> 4) | ((sh << 2) & 0x30)) - 32; + sh >>= 4; + const __m128i p_1_0 = _mm_madd_epi16(p16_1_0, _mm_set1_epi16(ls1)); + const __m128i p_1_1 = _mm_madd_epi16(p16_1_1, _mm_set1_epi16(ls1)); + const __m128i p_2_0 = _mm_madd_epi16(p16_2_0, _mm_set1_epi16(ls2)); + const __m128i p_2_1 = _mm_madd_epi16(p16_2_1, _mm_set1_epi16(ls2)); + sumi1_0 = _mm_add_epi32(p_1_0, sumi1_0); + sumi1_1 = _mm_add_epi32(p_1_1, sumi1_1); + sumi2_0 = _mm_add_epi32(p_2_0, sumi2_0); + sumi2_1 = _mm_add_epi32(p_2_1, sumi2_1); + } + __m128i sumi12_0 = _mm_add_epi32(sumi1_0, sumi2_0); + __m128i sumi12_1 = _mm_add_epi32(sumi1_1, sumi2_1); + accum = _mm256_add_ps(_mm256_mul_ps(_mm256_set1_ps(GGML_CPU_FP16_TO_FP32(x[ibl].d)*y[ibl].d), + _mm256_cvtepi32_ps(MM256_SET_M128I(sumi12_1, sumi12_0))), accum); + } + + *s = hsum_float_8(accum); + +#else + UNUSED(x); + UNUSED(y); + UNUSED(nb); + ggml_vec_dot_iq4_xs_q8_K_generic(n, s, bs, vx, bx, vy, by, nrc); +#endif +} diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cpu/ggml-cpu-impl.h b/python/freetoken/kernel/csrc/ggml/ggml-cpu/ggml-cpu-impl.h new file mode 100644 index 0000000000..82345b2063 --- /dev/null +++ b/python/freetoken/kernel/csrc/ggml/ggml-cpu/ggml-cpu-impl.h @@ -0,0 +1,538 @@ +#pragma once + +// GGML CPU internal header + +#include "ggml.h" +#include "ggml-impl.h" + +#include // load `stdlib.h` before other headers to work around MinGW bug: https://sourceforge.net/p/mingw-w64/bugs/192/ +//#include +#include +#include // memcpy +#include // fabsf + +#ifdef __cplusplus +extern "C" { +#endif + +struct ggml_compute_params { + // ith = thread index, nth = number of threads + int ith, nth; + + // work buffer for all threads + size_t wsize; + void * wdata; + + struct ggml_threadpool * threadpool; + + // use reference implementation + bool use_ref; +}; + + +#if defined(_MSC_VER) + +#define m512bh(p) p +#define m512i(p) p + +#else + +#define m512bh(p) (__m512bh)(p) +#define m512i(p) (__m512i)(p) + +#endif + +// __FMA__ and __F16C__ are not defined in MSVC, however they are implied with AVX2/AVX512 +#if defined(_MSC_VER) && (defined(__AVX2__) || defined(__AVX512F__)) +#ifndef __FMA__ +#define __FMA__ +#endif +#ifndef __F16C__ +#define __F16C__ +#endif +#endif + +// __SSE3__ and __SSSE3__ are not defined in MSVC, but SSE3/SSSE3 are present when AVX/AVX2/AVX512 are available +#if defined(_MSC_VER) && (defined(__AVX__) || defined(__AVX2__) || defined(__AVX512F__)) +#ifndef __SSE3__ +#define __SSE3__ +#endif +#ifndef __SSSE3__ +#define __SSSE3__ +#endif +#endif + +#if defined(__s390x__) && defined(__VEC__) +#ifndef __VXE__ +#define __VXE__ +#endif // __VXE__ +#ifndef __VXE2__ +#define __VXE2__ +#endif // __VXE2__ +#endif // __s390x__ && __VEC__ + +#if defined(__ARM_FEATURE_SVE) && defined(__linux__) +#include +#endif + +#if defined(__ARM_NEON) + +// ref: https://github.com/ggml-org/llama.cpp/pull/5404 +#if defined(_MSC_VER) && !defined(__clang__) +#define ggml_vld1q_u32(w,x,y,z) { ((w) + ((uint64_t)(x) << 32)), ((y) + ((uint64_t)(z) << 32)) } +#else +#define ggml_vld1q_u32(w,x,y,z) { (w), (x), (y), (z) } +#endif // _MSC_VER + +#if !defined(__aarch64__) + +// 32-bit ARM compatibility + +// vaddlvq_s16 +// vpaddq_s16 +// vpaddq_s32 +// vaddvq_s32 +// vaddvq_f32 +// vmaxvq_f32 +// vcvtnq_s32_f32 +// vzip1_u8 +// vzip2_u8 + +inline static int32_t vaddlvq_s16(int16x8_t v) { + int32x4_t v0 = vreinterpretq_s32_s64(vpaddlq_s32(vpaddlq_s16(v))); + return vgetq_lane_s32(v0, 0) + vgetq_lane_s32(v0, 2); +} + +inline static int16x8_t vpaddq_s16(int16x8_t a, int16x8_t b) { + int16x4_t a0 = vpadd_s16(vget_low_s16(a), vget_high_s16(a)); + int16x4_t b0 = vpadd_s16(vget_low_s16(b), vget_high_s16(b)); + return vcombine_s16(a0, b0); +} + +inline static int32x4_t vpaddq_s32(int32x4_t a, int32x4_t b) { + int32x2_t a0 = vpadd_s32(vget_low_s32(a), vget_high_s32(a)); + int32x2_t b0 = vpadd_s32(vget_low_s32(b), vget_high_s32(b)); + return vcombine_s32(a0, b0); +} + +inline static int32_t vaddvq_s32(int32x4_t v) { + return vgetq_lane_s32(v, 0) + vgetq_lane_s32(v, 1) + vgetq_lane_s32(v, 2) + vgetq_lane_s32(v, 3); +} + +inline static float vaddvq_f32(float32x4_t v) { + return vgetq_lane_f32(v, 0) + vgetq_lane_f32(v, 1) + vgetq_lane_f32(v, 2) + vgetq_lane_f32(v, 3); +} + +inline static float vmaxvq_f32(float32x4_t v) { + return + MAX(MAX(vgetq_lane_f32(v, 0), vgetq_lane_f32(v, 1)), + MAX(vgetq_lane_f32(v, 2), vgetq_lane_f32(v, 3))); +} + +inline static int32x4_t vcvtnq_s32_f32(float32x4_t v) { + int32x4_t res; + + res[0] = roundf(vgetq_lane_f32(v, 0)); + res[1] = roundf(vgetq_lane_f32(v, 1)); + res[2] = roundf(vgetq_lane_f32(v, 2)); + res[3] = roundf(vgetq_lane_f32(v, 3)); + + return res; +} + +inline static uint8x8_t vzip1_u8(uint8x8_t a, uint8x8_t b) { + uint8x8_t res; + + res[0] = a[0]; res[1] = b[0]; + res[2] = a[1]; res[3] = b[1]; + res[4] = a[2]; res[5] = b[2]; + res[6] = a[3]; res[7] = b[3]; + + return res; +} + +inline static uint8x8_t vzip2_u8(uint8x8_t a, uint8x8_t b) { + uint8x8_t res; + + res[0] = a[4]; res[1] = b[4]; + res[2] = a[5]; res[3] = b[5]; + res[4] = a[6]; res[5] = b[6]; + res[6] = a[7]; res[7] = b[7]; + + return res; +} + +// vld1q_s16_x2 +// vld1q_u8_x2 +// vld1q_u8_x4 +// vld1q_s8_x2 +// vld1q_s8_x4 +// TODO: double-check these work correctly + +typedef struct ggml_int16x8x2_t { + int16x8_t val[2]; +} ggml_int16x8x2_t; + +inline static ggml_int16x8x2_t ggml_vld1q_s16_x2(const int16_t * ptr) { + ggml_int16x8x2_t res; + + res.val[0] = vld1q_s16(ptr + 0); + res.val[1] = vld1q_s16(ptr + 8); + + return res; +} + +typedef struct ggml_uint8x16x2_t { + uint8x16_t val[2]; +} ggml_uint8x16x2_t; + +inline static ggml_uint8x16x2_t ggml_vld1q_u8_x2(const uint8_t * ptr) { + ggml_uint8x16x2_t res; + + res.val[0] = vld1q_u8(ptr + 0); + res.val[1] = vld1q_u8(ptr + 16); + + return res; +} + +typedef struct ggml_uint8x16x4_t { + uint8x16_t val[4]; +} ggml_uint8x16x4_t; + +inline static ggml_uint8x16x4_t ggml_vld1q_u8_x4(const uint8_t * ptr) { + ggml_uint8x16x4_t res; + + res.val[0] = vld1q_u8(ptr + 0); + res.val[1] = vld1q_u8(ptr + 16); + res.val[2] = vld1q_u8(ptr + 32); + res.val[3] = vld1q_u8(ptr + 48); + + return res; +} + +typedef struct ggml_int8x16x2_t { + int8x16_t val[2]; +} ggml_int8x16x2_t; + +inline static ggml_int8x16x2_t ggml_vld1q_s8_x2(const int8_t * ptr) { + ggml_int8x16x2_t res; + + res.val[0] = vld1q_s8(ptr + 0); + res.val[1] = vld1q_s8(ptr + 16); + + return res; +} + +typedef struct ggml_int8x16x4_t { + int8x16_t val[4]; +} ggml_int8x16x4_t; + +inline static ggml_int8x16x4_t ggml_vld1q_s8_x4(const int8_t * ptr) { + ggml_int8x16x4_t res; + + res.val[0] = vld1q_s8(ptr + 0); + res.val[1] = vld1q_s8(ptr + 16); + res.val[2] = vld1q_s8(ptr + 32); + res.val[3] = vld1q_s8(ptr + 48); + + return res; +} + +// NOTE: not tested +inline static int8x16_t ggml_vqtbl1q_s8(int8x16_t a, uint8x16_t b) { + int8x16_t res; + + res[ 0] = a[b[ 0]]; + res[ 1] = a[b[ 1]]; + res[ 2] = a[b[ 2]]; + res[ 3] = a[b[ 3]]; + res[ 4] = a[b[ 4]]; + res[ 5] = a[b[ 5]]; + res[ 6] = a[b[ 6]]; + res[ 7] = a[b[ 7]]; + res[ 8] = a[b[ 8]]; + res[ 9] = a[b[ 9]]; + res[10] = a[b[10]]; + res[11] = a[b[11]]; + res[12] = a[b[12]]; + res[13] = a[b[13]]; + res[14] = a[b[14]]; + res[15] = a[b[15]]; + + return res; +} + +// NOTE: not tested +inline static uint8x16_t ggml_vqtbl1q_u8(uint8x16_t a, uint8x16_t b) { + uint8x16_t res; + + res[ 0] = a[b[ 0]]; + res[ 1] = a[b[ 1]]; + res[ 2] = a[b[ 2]]; + res[ 3] = a[b[ 3]]; + res[ 4] = a[b[ 4]]; + res[ 5] = a[b[ 5]]; + res[ 6] = a[b[ 6]]; + res[ 7] = a[b[ 7]]; + res[ 8] = a[b[ 8]]; + res[ 9] = a[b[ 9]]; + res[10] = a[b[10]]; + res[11] = a[b[11]]; + res[12] = a[b[12]]; + res[13] = a[b[13]]; + res[14] = a[b[14]]; + res[15] = a[b[15]]; + + return res; +} + +#else + +#define ggml_int16x8x2_t int16x8x2_t +#define ggml_uint8x16x2_t uint8x16x2_t +#define ggml_uint8x16x4_t uint8x16x4_t +#define ggml_int8x16x2_t int8x16x2_t +#define ggml_int8x16x4_t int8x16x4_t + +#define ggml_vld1q_s16_x2 vld1q_s16_x2 +#define ggml_vld1q_u8_x2 vld1q_u8_x2 +#define ggml_vld1q_u8_x4 vld1q_u8_x4 +#define ggml_vld1q_s8_x2 vld1q_s8_x2 +#define ggml_vld1q_s8_x4 vld1q_s8_x4 +#define ggml_vqtbl1q_s8 vqtbl1q_s8 +#define ggml_vqtbl1q_u8 vqtbl1q_u8 + +#endif // !defined(__aarch64__) + +#if !defined(__ARM_FEATURE_DOTPROD) + +// NOTE: this fallback produces the same total sum as native vdotq_s32 but with different per-lane grouping — do not use when individual lane values matter. +inline static int32x4_t ggml_vdotq_s32(int32x4_t acc, int8x16_t a, int8x16_t b) { + const int16x8_t p0 = vmull_s8(vget_low_s8 (a), vget_low_s8 (b)); + const int16x8_t p1 = vmull_s8(vget_high_s8(a), vget_high_s8(b)); + + return vaddq_s32(acc, vaddq_s32(vpaddlq_s16(p0), vpaddlq_s16(p1))); +} + +#else + +#define ggml_vdotq_s32(a, b, c) vdotq_s32(a, b, c) + +#endif // !defined(__ARM_FEATURE_DOTPROD) + +static inline int32x4_t ggml_nvfp4_dot8(const int8x8_t q4_lo, const int8x8_t q8_lo, + const int8x8_t q4_hi, const int8x8_t q8_hi) { + const int16x8_t p_lo = vmull_s8(q4_lo, q8_lo); + const int16x8_t p_hi = vmull_s8(q4_hi, q8_hi); + const int32x4_t sum_lo = vpaddlq_s16(p_lo); + const int32x4_t sum_hi = vpaddlq_s16(p_hi); + return vaddq_s32(sum_lo, sum_hi); +} + +#endif // defined(__ARM_NEON) + +#ifdef __wasm_simd128__ +#include +#endif + +#ifdef __POWER9_VECTOR__ +#include +#endif + +#if defined(_MSC_VER) || defined(__MINGW32__) +#include +#elif defined(__SSE__) || defined(__SSE3__) || defined(__SSSE3__) || defined(__AVX__) || defined(__F16C__) || defined(__AVX2__) || defined(__AVX512F__) || defined(__AVX512BF16__) +#include +#endif + +#ifdef __riscv_v_intrinsic +#include +#endif + +#if defined(__loongarch64) +#if defined(__loongarch_asx) +#include +#endif +#if defined(__loongarch_sx) +#include +#endif +#endif + +#if defined(__VXE__) || defined(__VXE2__) +#include + +#define vec_neg(a) (-(a)) // Vector Negate +#define vec_add(a, b) ((a) + (b)) // Vector Add +#define vec_sub(a, b) ((a) - (b)) // Vector Subtract +#define vec_mul(a, b) ((a) * (b)) // Vector Multiply +#define vec_div(a, b) ((a) / (b)) // Vector Divide +#define vec_sl(a, b) ((a) << (b)) // Vector Shift Left +#define vec_sra(a, b) ((a) >> (b)) // Vector Shift Right +#define vec_sr(a, b) ((a) >> (b)) // Vector Shift Right Algebraic +#define vec_slo(a, b) vec_slb(a, (b) << 64) // Vector Shift Left by Octet +#define vec_sro(a, b) vec_srb(a, (b) << 64) // Vector Shift Right by Octet + +#ifndef vec_and +#define vec_and(a, b) ((a) & (b)) // Vector AND +#endif + +#ifndef vec_or +#define vec_or(a, b) ((a) | (b)) // Vector OR +#endif + +#ifndef vec_xor +#define vec_xor(a, b) ((a) ^ (b)) // Vector XOR +#endif + +typedef signed char char8x16_t __attribute__((vector_size(16))); +typedef unsigned char uchar8x16_t __attribute__((vector_size(16))); + +typedef int8_t int8x16_t __attribute__((vector_size(16))); +typedef int16_t int16x8_t __attribute__((vector_size(16))); +typedef int32_t int32x4_t __attribute__((vector_size(16))); +typedef int64_t int64x2_t __attribute__((vector_size(16))); + +typedef uint8_t uint8x16_t __attribute__((vector_size(16))); +typedef uint16_t uint16x8_t __attribute__((vector_size(16))); +typedef uint32_t uint32x4_t __attribute__((vector_size(16))); +typedef uint64_t uint64x2_t __attribute__((vector_size(16))); + +typedef float float32x4_t __attribute__((vector_size(16))); +typedef double double64x2_t __attribute__((vector_size(16))); + +typedef struct ggml_uint8x16x2_t { + uint8x16_t val[2]; +} ggml_uint8x16x2_t; + +inline static ggml_uint8x16x2_t ggml_vec_xl_u8x2(const uint8_t * ptr) { + ggml_uint8x16x2_t res; + + res.val[0] = vec_xl( 0, ptr); + res.val[1] = vec_xl(16, ptr); + + return res; +} + +typedef struct ggml_uint8x16x4_t { + uint8x16_t val[4]; +} ggml_uint8x16x4_t; + +inline static ggml_uint8x16x4_t ggml_vec_xl_u8x4(const uint8_t * ptr) { + ggml_uint8x16x4_t res; + + res.val[0] = vec_xl( 0, ptr); + res.val[1] = vec_xl(16, ptr); + res.val[2] = vec_xl(32, ptr); + res.val[3] = vec_xl(48, ptr); + + return res; +} + +typedef struct ggml_int8x16x4_t { + int8x16_t val[4]; +} ggml_int8x16x4_t; + +inline static ggml_int8x16x4_t ggml_vec_xl_s8x4(const int8_t * ptr) { + ggml_int8x16x4_t res; + + res.val[0] = vec_xl( 0, ptr); + res.val[1] = vec_xl(16, ptr); + res.val[2] = vec_xl(32, ptr); + res.val[3] = vec_xl(48, ptr); + + return res; +} + +typedef struct ggml_int16x8x2_t { + int16x8_t val[2]; +} ggml_int16x8x2_t; + +inline static ggml_int16x8x2_t ggml_vec_xl_s16x2(const int16_t * ptr) { + ggml_int16x8x2_t res; + + res.val[0] = vec_xl( 0, ptr); + res.val[1] = vec_xl(16, ptr); + + return res; +} + +/* + ! WARNING: Very slow. Use vec_perm if possible. Refer to iq4_xs + ! or iq4_nl for example implementation. +*/ +inline static int8x16_t ggml_vec_tbl(int8x16_t a, uint8x16_t b) { + int8x16_t res; + + res[ 0] = a[b[ 0]]; + res[ 1] = a[b[ 1]]; + res[ 2] = a[b[ 2]]; + res[ 3] = a[b[ 3]]; + res[ 4] = a[b[ 4]]; + res[ 5] = a[b[ 5]]; + res[ 6] = a[b[ 6]]; + res[ 7] = a[b[ 7]]; + res[ 8] = a[b[ 8]]; + res[ 9] = a[b[ 9]]; + res[10] = a[b[10]]; + res[11] = a[b[11]]; + res[12] = a[b[12]]; + res[13] = a[b[13]]; + res[14] = a[b[14]]; + res[15] = a[b[15]]; + + return res; +} + +inline static int16x8_t vec_padd_s16(int16x8_t a, int16x8_t b) { + const uchar8x16_t v_maske = { 0, 1, 4, 5, 8, 9, 12, 13, + 16, 17, 20, 21, 24, 25, 28, 29 }; + + const int16x8_t v_abo = vec_pack((int32x4_t)a, (int32x4_t)b); + const int16x8_t v_abe = vec_perm(a, b, v_maske); + return v_abo + v_abe; +} + +/** + * @see https://github.com/ggml-org/llama.cpp/pull/14037 + */ +inline static float vec_hsum_f32x4(float32x4_t v) { + float32x4_t v_temp = v + vec_reve(v); + return v_temp[0] + v_temp[1]; +} + +inline static int32_t vec_hsum_i32x4(int32x4_t v) { + int32x4_t v_temp = v + vec_reve(v); + return v_temp[0] + v_temp[1]; +} + +inline static int32x4_t ggml_vec_dot(int32x4_t acc, int8x16_t a, int8x16_t b) { + const int16x8_t p = vec_mule(a, b) + vec_mulo(a, b); + return acc + (vec_unpackh(p) + vec_unpackl(p)); +} + +#endif + +#if defined(__loongarch_sx) +/* float type data load instructions */ +static __m128 __lsx_vreplfr2vr_s(const float val) { + v4f32 res = {val, val, val, val}; + return (__m128)res; +} +#endif + +#if defined(__loongarch_asx) +static __m256 __lasx_xvreplfr2vr_s(const float val) { + v8f32 res = {val, val, val, val, val, val, val, val}; + return (__m256)res; +} +#endif + +// TODO: move to ggml-threading +void ggml_barrier(struct ggml_threadpool * tp); + +void ggml_threadpool_chunk_set(struct ggml_threadpool * tp, int value); +int ggml_threadpool_chunk_add(struct ggml_threadpool * tp, int value); + +#ifdef __cplusplus +} +#endif diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cpu/quants.c b/python/freetoken/kernel/csrc/ggml/ggml-cpu/quants.c new file mode 100644 index 0000000000..5e36459f8c --- /dev/null +++ b/python/freetoken/kernel/csrc/ggml/ggml-cpu/quants.c @@ -0,0 +1,1339 @@ +#define GGML_COMMON_IMPL_C +#include "ggml-common.h" + +#include "ggml-cpu-impl.h" +#include "simd-mappings.h" +#include "ggml-quants.h" +#include "quants.h" + +#include "arch-fallback.h" + +#include +#include +#include +#include // for qsort +#include // for GGML_ASSERT + +#define GROUP_MAX_EPS 1e-15f +#define GROUP_MAX_EPS_IQ3_XXS 1e-8f +#define GROUP_MAX_EPS_IQ2_S 1e-8f +#define GROUP_MAX_EPS_IQ1_M 1e-7f +#define GROUP_MAX_EPS_IQ1_S 1e-12f + +#define UNUSED GGML_UNUSED + +void quantize_row_q1_0(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k) { + quantize_row_q1_0_ref(x, y, k); +} + +void quantize_row_q2_0(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k) { + quantize_row_q2_0_ref(x, y, k); +} + +void quantize_row_q4_0(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k) { + quantize_row_q4_0_ref(x, y, k); +} + +void quantize_row_q4_1(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k) { + quantize_row_q4_1_ref(x, y, k); +} + +void quantize_row_q5_0(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k) { + quantize_row_q5_0_ref(x, y, k); +} + +void quantize_row_q5_1(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k) { + quantize_row_q5_1_ref(x, y, k); +} + +void quantize_row_q8_0_generic(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k) { + quantize_row_q8_0_ref(x, y, k); +} + +void quantize_row_q8_1_generic(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k) { + quantize_row_q8_1_ref(x, y, k); +} + +void quantize_row_mxfp4(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k) { + quantize_row_mxfp4_ref(x, y, k); +} + +void quantize_row_nvfp4(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k) { + quantize_row_nvfp4_ref(x, y, k); +} + +// +// 2-6 bit quantization in super-blocks +// + +//========================- 2-bit (de)-quantization + +void quantize_row_q2_K(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k) { + quantize_row_q2_K_ref(x, vy, k); +} + +//========================= 3-bit (de)-quantization + +void quantize_row_q3_K(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k) { + quantize_row_q3_K_ref(x, vy, k); +} + +// ====================== 4-bit (de)-quantization + +void quantize_row_q4_K(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k) { + assert(k % QK_K == 0); + block_q4_K * GGML_RESTRICT y = vy; + quantize_row_q4_K_ref(x, y, k); +} + +// ====================== 5-bit (de)-quantization + +void quantize_row_q5_K(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k) { + assert(k % QK_K == 0); + block_q5_K * GGML_RESTRICT y = vy; + quantize_row_q5_K_ref(x, y, k); +} + +// ====================== 6-bit (de)-quantization + +void quantize_row_q6_K(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k) { + assert(k % QK_K == 0); + block_q6_K * GGML_RESTRICT y = vy; + quantize_row_q6_K_ref(x, y, k); +} + +// ====================== Ternary (de)-quantization (BitNet b1.58 and TriLMs) + +void quantize_row_tq1_0(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k) { + assert(k % QK_K == 0); + block_tq1_0 * GGML_RESTRICT y = vy; + quantize_row_tq1_0_ref(x, y, k); +} + +void quantize_row_tq2_0(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k) { + assert(k % QK_K == 0); + block_tq2_0 * GGML_RESTRICT y = vy; + quantize_row_tq2_0_ref(x, y, k); +} + +//===================================== Q8_K ============================================== + +void quantize_row_q8_K_generic(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k) { + quantize_row_q8_K_ref(x, y, k); +} + +//===================================== Dot products ================================= + +void ggml_vec_dot_q1_0_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) { + const int qk = QK1_0; + const int nb = n / qk; + + assert(n % qk == 0); + assert(nrc == 1); + UNUSED(nrc); + UNUSED(bx); + UNUSED(by); + UNUSED(bs); + + const block_q1_0 * GGML_RESTRICT x = vx; + const block_q8_0 * GGML_RESTRICT y = vy; + + float sumf = 0.0; + + for (int i = 0; i < nb; i++) { + const float d0 = GGML_CPU_FP16_TO_FP32(x[i].d); + + float sumi = 0.0f; + + for (int k = 0; k < 4; k++) { + const block_q8_0 * GGML_RESTRICT yb = &y[i * 4 + k]; + const float d1 = GGML_CPU_FP16_TO_FP32(yb->d); + int sumi_block = 0; + + const uint8_t * GGML_RESTRICT bits = &x[i].qs[k * 4]; + const int8_t * GGML_RESTRICT qy = yb->qs; + + for (int b = 0; b < 4; ++b, qy += 8) { + const unsigned mask = bits[b]; + sumi_block += ((mask & 0x01) ? qy[0] : -qy[0]) + + ((mask & 0x02) ? qy[1] : -qy[1]) + + ((mask & 0x04) ? qy[2] : -qy[2]) + + ((mask & 0x08) ? qy[3] : -qy[3]) + + ((mask & 0x10) ? qy[4] : -qy[4]) + + ((mask & 0x20) ? qy[5] : -qy[5]) + + ((mask & 0x40) ? qy[6] : -qy[6]) + + ((mask & 0x80) ? qy[7] : -qy[7]); + } + + sumi += d1 * sumi_block; + } + + sumf += d0 * sumi; + } + + *s = sumf; +} + +void ggml_vec_dot_q2_0_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) { + const int qk = QK2_0; + const int nb = n / qk; + + assert(n % qk == 0); + assert(nrc == 1); + UNUSED(nrc); + UNUSED(bx); + UNUSED(by); + UNUSED(bs); + + const block_q2_0 * GGML_RESTRICT x = vx; + const block_q8_0 * GGML_RESTRICT y = vy; + + float sumf = 0.0f; + + for (int i = 0; i < nb; i++) { + const float d0 = GGML_CPU_FP16_TO_FP32(x[i].d); + + float sumi = 0.0f; + + // group 64: one Q2_0 block (64 weights) maps to two Q8_0 blocks (2 * 32 = 64) + for (int k = 0; k < 2; k++) { + const block_q8_0 * GGML_RESTRICT yb = &y[i * 2 + k]; + const float d1 = GGML_CPU_FP16_TO_FP32(yb->d); + int sumi_block = 0; + + const uint8_t * GGML_RESTRICT qs = &x[i].qs[k * 8]; + const int8_t * GGML_RESTRICT qy = yb->qs; + + for (int b = 0; b < 8; ++b) { + const uint8_t byte = qs[b]; + // Extract 4 two-bit values, map {0,1,2,3} -> {-1,0,1,2} + sumi_block += ((int)((byte >> 0) & 3) - 1) * qy[b*4 + 0]; + sumi_block += ((int)((byte >> 2) & 3) - 1) * qy[b*4 + 1]; + sumi_block += ((int)((byte >> 4) & 3) - 1) * qy[b*4 + 2]; + sumi_block += ((int)((byte >> 6) & 3) - 1) * qy[b*4 + 3]; + } + + sumi += d1 * sumi_block; + } + + sumf += d0 * sumi; + } + + *s = sumf; +} + +void ggml_vec_dot_q4_0_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) { + const int qk = QK8_0; + const int nb = n / qk; + + assert(n % qk == 0); + assert(nrc == 1); + UNUSED(nrc); + UNUSED(bx); + UNUSED(by); + UNUSED(bs); + + const block_q4_0 * GGML_RESTRICT x = vx; + const block_q8_0 * GGML_RESTRICT y = vy; + + int ib = 0; + float sumf = 0; + + for (; ib < nb; ++ib) { + int sumi0 = 0; + int sumi1 = 0; + + for (int j = 0; j < qk/2; ++j) { + const int v0 = (x[ib].qs[j] & 0x0F) - 8; + const int v1 = (x[ib].qs[j] >> 4) - 8; + + sumi0 += (v0 * y[ib].qs[j]); + sumi1 += (v1 * y[ib].qs[j + qk/2]); + } + + int sumi = sumi0 + sumi1; + sumf += sumi*GGML_CPU_FP16_TO_FP32(x[ib].d)*GGML_CPU_FP16_TO_FP32(y[ib].d); + } + + *s = sumf; +} + +// TODO: add WASM SIMD +void ggml_vec_dot_q4_1_q8_1_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) { + const int qk = QK8_1; + const int nb = n / qk; + + assert(n % qk == 0); + assert(nrc == 1); + UNUSED(nrc); + UNUSED(bx); + UNUSED(by); + UNUSED(bs); + + const block_q4_1 * GGML_RESTRICT x = vx; + const block_q8_1 * GGML_RESTRICT y = vy; + + int ib = 0; + float sumf = 0; + + for (; ib < nb; ++ib) { + int sumi0 = 0; + int sumi1 = 0; + + for (int j = 0; j < qk/2; ++j) { + const int v0 = (x[ib].qs[j] & 0x0F); + const int v1 = (x[ib].qs[j] >> 4); + + sumi0 += (v0 * y[ib].qs[j]); + sumi1 += (v1 * y[ib].qs[j + qk/2]); + } + + int sumi = sumi0 + sumi1; + sumf += (GGML_CPU_FP16_TO_FP32(x[ib].d)*GGML_CPU_FP16_TO_FP32(y[ib].d))*sumi + GGML_CPU_FP16_TO_FP32(x[ib].m)*GGML_CPU_FP16_TO_FP32(y[ib].s); + } + + *s = sumf; +} + +void ggml_vec_dot_mxfp4_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) { + assert(nrc == 1); + UNUSED(nrc); + UNUSED(bx); + UNUSED(by); + UNUSED(bs); + assert(n % QK_MXFP4 == 0); + static_assert(QK_MXFP4 == QK8_0, "QK_MXFP4 and QK8_0 must be the same"); + + const block_mxfp4 * GGML_RESTRICT x = vx; + const block_q8_0 * GGML_RESTRICT y = vy; + + const int nb = n / QK_MXFP4; + + int ib = 0; + float sumf = 0; + + for (; ib < nb; ++ib) { + const float d = GGML_CPU_FP16_TO_FP32(y[ib].d)*GGML_E8M0_TO_FP32_HALF(x[ib].e); + + int sumi1 = 0; + int sumi2 = 0; + for (int j = 0; j < QK_MXFP4/2; ++j) { + sumi1 += y[ib].qs[j + 0] * kvalues_mxfp4[x[ib].qs[j] & 0xf]; + sumi2 += y[ib].qs[j + QK_MXFP4/2] * kvalues_mxfp4[x[ib].qs[j] >> 4]; + } + sumf += d * (sumi1 + sumi2); + } + *s = sumf; +} + +// NVFP4: super-block of 64 elements = 4 sub-blocks of 16 = 2 q8_0 blocks +void ggml_vec_dot_nvfp4_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) { + assert(nrc == 1); + UNUSED(nrc); + UNUSED(bx); + UNUSED(by); + UNUSED(bs); + assert(n % QK_NVFP4 == 0); + + const block_nvfp4 * GGML_RESTRICT x = vx; + const block_q8_0 * GGML_RESTRICT y = vy; + + const int nb = n / QK_NVFP4; + + float sumf = 0; + + for (int ib = 0; ib < nb; ++ib) { + for (int s_idx = 0; s_idx < 4; ++s_idx) { + const float d = ggml_ue4m3_to_fp32(x[ib].d[s_idx]); + const int q8_block = s_idx / 2; + const int q8_off = (s_idx % 2) * QK_NVFP4_SUB; + const float dy = GGML_CPU_FP16_TO_FP32(y[2*ib + q8_block].d); + + int sumi_lo = 0, sumi_hi = 0; + for (int j = 0; j < QK_NVFP4_SUB/2; ++j) { + const uint8_t qv = x[ib].qs[s_idx*(QK_NVFP4_SUB/2) + j]; + sumi_lo += y[2*ib + q8_block].qs[q8_off + j + 0] * kvalues_mxfp4[qv & 0xf]; + sumi_hi += y[2*ib + q8_block].qs[q8_off + j + QK_NVFP4_SUB/2] * kvalues_mxfp4[qv >> 4]; + } + + sumf += dy * d * (sumi_lo + sumi_hi); + } + } + *s = sumf; +} + +void ggml_vec_dot_q5_0_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) { + const int qk = QK8_0; + const int nb = n / qk; + + int ib = 0; + float sumf = 0; + + assert(n % qk == 0); + assert(qk == QK5_0); + assert(nrc == 1); + UNUSED(nrc); + UNUSED(bx); + UNUSED(by); + UNUSED(bs); + + const block_q5_0 * GGML_RESTRICT x = vx; + const block_q8_0 * GGML_RESTRICT y = vy; + + for (; ib < nb; ++ib) { + uint32_t qh; + memcpy(&qh, x[ib].qh, sizeof(qh)); + + int sumi0 = 0; + int sumi1 = 0; + + for (int j = 0; j < qk/2; ++j) { + const uint8_t xh_0 = ((qh & (1u << (j + 0 ))) >> (j + 0 )) << 4; + const uint8_t xh_1 = ((qh & (1u << (j + 16))) >> (j + 12)); + + const int32_t x0 = (int8_t)(((x[ib].qs[j] & 0x0F) | xh_0) - 16); + const int32_t x1 = (int8_t)(((x[ib].qs[j] >> 4) | xh_1) - 16); + + sumi0 += (x0 * y[ib].qs[j]); + sumi1 += (x1 * y[ib].qs[j + qk/2]); + } + + int sumi = sumi0 + sumi1; + sumf += (GGML_CPU_FP16_TO_FP32(x[ib].d)*GGML_CPU_FP16_TO_FP32(y[ib].d)) * sumi; + } + + *s = sumf; +} + +void ggml_vec_dot_q5_1_q8_1_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) { + const int qk = QK8_1; + const int nb = n / qk; + + int ib = 0; + float sumf = 0; + + assert(n % qk == 0); + assert(qk == QK5_1); + assert(nrc == 1); + UNUSED(nrc); + UNUSED(bx); + UNUSED(by); + UNUSED(bs); + + const block_q5_1 * GGML_RESTRICT x = vx; + const block_q8_1 * GGML_RESTRICT y = vy; + + for (; ib < nb; ++ib) { + uint32_t qh; + memcpy(&qh, x[ib].qh, sizeof(qh)); + + int sumi0 = 0; + int sumi1 = 0; + + for (int j = 0; j < qk/2; ++j) { + const uint8_t xh_0 = ((qh >> (j + 0)) << 4) & 0x10; + const uint8_t xh_1 = ((qh >> (j + 12)) ) & 0x10; + + const int32_t x0 = (x[ib].qs[j] & 0xF) | xh_0; + const int32_t x1 = (x[ib].qs[j] >> 4) | xh_1; + + sumi0 += (x0 * y[ib].qs[j]); + sumi1 += (x1 * y[ib].qs[j + qk/2]); + } + + int sumi = sumi0 + sumi1; + sumf += (GGML_CPU_FP16_TO_FP32(x[ib].d)*GGML_CPU_FP16_TO_FP32(y[ib].d))*sumi + GGML_CPU_FP16_TO_FP32(x[ib].m)*GGML_CPU_FP16_TO_FP32(y[ib].s); + } + + *s = sumf; +} + +void ggml_vec_dot_q8_0_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) { + const int qk = QK8_0; + const int nb = n / qk; + + assert(n % qk == 0); + assert(nrc == 1); + UNUSED(nrc); + UNUSED(bx); + UNUSED(by); + UNUSED(bs); + + const block_q8_0 * GGML_RESTRICT x = vx; + const block_q8_0 * GGML_RESTRICT y = vy; + + int ib = 0; + float sumf = 0; + + for (; ib < nb; ++ib) { + int sumi = 0; + + for (int j = 0; j < qk; j++) { + sumi += x[ib].qs[j]*y[ib].qs[j]; + } + + sumf += sumi*(GGML_CPU_FP16_TO_FP32(x[ib].d)*GGML_CPU_FP16_TO_FP32(y[ib].d)); + } + + *s = sumf; +} + +void ggml_vec_dot_tq1_0_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) { + assert(nrc == 1); + UNUSED(nrc); + UNUSED(bx); + UNUSED(by); + UNUSED(bs); + + const block_tq1_0 * GGML_RESTRICT x = vx; + const block_q8_K * GGML_RESTRICT y = vy; + + const int nb = n / QK_K; + + const uint8_t pow3[6] = {1, 3, 9, 27, 81, 243}; + + float sumf = 0.0f; + + for (int i = 0; i < nb; ++i) { + int sum = 0; + + for (size_t j = 0; j < sizeof(x->qs) - sizeof(x->qs) % 32; j += 32) { + for (size_t l = 0; l < 5; ++l) { + for (size_t m = 0; m < 32; ++m) { + uint8_t q = x[i].qs[j + m] * pow3[l]; + uint16_t xi = ((uint16_t) q * 3) >> 8; + sum += (xi - 1) * y[i].qs[j*5 + l*32 + m]; + } + } + } + for (size_t j = sizeof(x->qs) - sizeof(x->qs) % 32; j < sizeof(x->qs); j += 16) { + for (size_t l = 0; l < 5; ++l) { + for (size_t m = 0; m < 16; ++m) { + uint8_t q = x[i].qs[j + m] * pow3[l]; + uint16_t xi = ((uint16_t) q * 3) >> 8; + sum += (xi - 1) * y[i].qs[j*5 + l*16 + m]; + } + } + } + + for (size_t l = 0; l < 4; ++l) { + for (size_t j = 0; j < sizeof(x->qh); ++j) { + uint8_t q = x[i].qh[j] * pow3[l]; + uint16_t xi = ((uint16_t) q * 3) >> 8; + sum += (xi - 1) * y[i].qs[sizeof(x->qs)*5 + l*sizeof(x->qh) + j]; + } + } + + sumf += (float) sum * (GGML_CPU_FP16_TO_FP32(x[i].d) * y[i].d); + } + + *s = sumf; +} + +void ggml_vec_dot_tq2_0_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) { + assert(nrc == 1); + UNUSED(nrc); + UNUSED(bx); + UNUSED(by); + UNUSED(bs); + + const block_tq2_0 * GGML_RESTRICT x = vx; + const block_q8_K * GGML_RESTRICT y = vy; + + const int nb = n / QK_K; + float sumf = 0.0f; + + for (int i = 0; i < nb; ++i) { + int32_t sumi = 0; + + for (size_t j = 0; j < sizeof(x->qs); j += 32) { + for (size_t l = 0; l < 4; ++l) { + for (size_t k = 0; k < 32; ++k) { + sumi += y[i].qs[j*4 + l*32 + k] * (((x[i].qs[j + k] >> (l*2)) & 3) - 1); + } + } + } + + const float d = y[i].d * GGML_CPU_FP16_TO_FP32(x[i].d); + + sumf += (float) sumi * d; + } + + *s = sumf; +} + +void ggml_vec_dot_q2_K_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) { + assert(nrc == 1); + UNUSED(nrc); + UNUSED(bx); + UNUSED(by); + UNUSED(bs); + + const block_q2_K * GGML_RESTRICT x = vx; + const block_q8_K * GGML_RESTRICT y = vy; + + const int nb = n / QK_K; + + float sumf = 0; + + for (int i = 0; i < nb; ++i) { + + const uint8_t * q2 = x[i].qs; + const int8_t * q8 = y[i].qs; + const uint8_t * sc = x[i].scales; + + int summs = 0; + for (int j = 0; j < 16; ++j) { + summs += y[i].bsums[j] * (sc[j] >> 4); + } + + const float dall = y[i].d * GGML_CPU_FP16_TO_FP32(x[i].d); + const float dmin = y[i].d * GGML_CPU_FP16_TO_FP32(x[i].dmin); + + int isum = 0; + int is = 0; + int d; + for (int k = 0; k < QK_K/128; ++k) { + int shift = 0; + for (int j = 0; j < 4; ++j) { + d = sc[is++] & 0xF; + int isuml = 0; + for (int l = 0; l < 16; ++l) isuml += q8[l] * ((q2[l] >> shift) & 3); + isum += d * isuml; + d = sc[is++] & 0xF; + isuml = 0; + for (int l = 16; l < 32; ++l) isuml += q8[l] * ((q2[l] >> shift) & 3); + isum += d * isuml; + shift += 2; + q8 += 32; + } + q2 += 32; + } + sumf += dall * isum - dmin * summs; + } + *s = sumf; +} + +void ggml_vec_dot_q3_K_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) { + assert(n % QK_K == 0); + assert(nrc == 1); + UNUSED(nrc); + UNUSED(bx); + UNUSED(by); + UNUSED(bs); + + const uint32_t kmask1 = 0x03030303; + const uint32_t kmask2 = 0x0f0f0f0f; + + const block_q3_K * GGML_RESTRICT x = vx; + const block_q8_K * GGML_RESTRICT y = vy; + + const int nb = n / QK_K; + + // scalar version + // This function is written like this so the compiler can manage to vectorize most of it + // Using -Ofast, GCC and clang manage to produce code that is within a factor of 2 or so from the + // manually vectorized version above. Every other version I tried would run at least 4 times slower. + // The ideal situation would be if we could just write the code once, and the compiler would + // automatically produce the best possible set of machine instructions, instead of us having to manually + // write vectorized versions for AVX, ARM_NEON, etc. + + int8_t aux8[QK_K]; + int16_t aux16[8]; + float sums [8]; + int32_t aux32[8]; + memset(sums, 0, 8*sizeof(float)); + + uint32_t auxs[4]; + const int8_t * scales = (const int8_t*)auxs; + + float sumf = 0; + for (int i = 0; i < nb; ++i) { + const uint8_t * GGML_RESTRICT q3 = x[i].qs; + const uint8_t * GGML_RESTRICT hm = x[i].hmask; + const int8_t * GGML_RESTRICT q8 = y[i].qs; + memset(aux32, 0, 8*sizeof(int32_t)); + int8_t * GGML_RESTRICT a = aux8; + uint8_t m = 1; + for (int j = 0; j < QK_K; j += 128) { + for (int l = 0; l < 32; ++l) a[l] = q3[l] & 3; + for (int l = 0; l < 32; ++l) a[l] -= (hm[l] & m ? 0 : 4); + a += 32; m <<= 1; + for (int l = 0; l < 32; ++l) a[l] = (q3[l] >> 2) & 3; + for (int l = 0; l < 32; ++l) a[l] -= (hm[l] & m ? 0 : 4); + a += 32; m <<= 1; + for (int l = 0; l < 32; ++l) a[l] = (q3[l] >> 4) & 3; + for (int l = 0; l < 32; ++l) a[l] -= (hm[l] & m ? 0 : 4); + a += 32; m <<= 1; + for (int l = 0; l < 32; ++l) a[l] = (q3[l] >> 6) & 3; + for (int l = 0; l < 32; ++l) a[l] -= (hm[l] & m ? 0 : 4); + a += 32; m <<= 1; + q3 += 32; + } + a = aux8; + + memcpy(auxs, x[i].scales, 12); + uint32_t tmp = auxs[2]; + auxs[2] = ((auxs[0] >> 4) & kmask2) | (((tmp >> 4) & kmask1) << 4); + auxs[3] = ((auxs[1] >> 4) & kmask2) | (((tmp >> 6) & kmask1) << 4); + auxs[0] = (auxs[0] & kmask2) | (((tmp >> 0) & kmask1) << 4); + auxs[1] = (auxs[1] & kmask2) | (((tmp >> 2) & kmask1) << 4); + for (int j = 0; j < QK_K/16; ++j) { + for (int l = 0; l < 8; ++l) aux16[l] = q8[l] * a[l]; + for (int l = 0; l < 8; ++l) aux32[l] += (scales[j] - 32) * aux16[l]; + q8 += 8; a += 8; + for (int l = 0; l < 8; ++l) aux16[l] = q8[l] * a[l]; + for (int l = 0; l < 8; ++l) aux32[l] += (scales[j] - 32) * aux16[l]; + q8 += 8; a += 8; + } + const float d = GGML_CPU_FP16_TO_FP32(x[i].d) * y[i].d; + for (int l = 0; l < 8; ++l) sums[l] += d * aux32[l]; + } + for (int l = 0; l < 8; ++l) sumf += sums[l]; + *s = sumf; +} + +void ggml_vec_dot_q4_K_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) { + assert(n % QK_K == 0); + assert(nrc == 1); + UNUSED(nrc); + UNUSED(bx); + UNUSED(by); + UNUSED(bs); + + const block_q4_K * GGML_RESTRICT x = vx; + const block_q8_K * GGML_RESTRICT y = vy; + + const int nb = n / QK_K; + + static const uint32_t kmask1 = 0x3f3f3f3f; + static const uint32_t kmask2 = 0x0f0f0f0f; + static const uint32_t kmask3 = 0x03030303; + + uint32_t utmp[4]; + + const uint8_t * scales = (const uint8_t*)&utmp[0]; + const uint8_t * mins = (const uint8_t*)&utmp[2]; + + int8_t aux8[QK_K]; + int16_t aux16[8]; + float sums [8]; + int32_t aux32[8]; + memset(sums, 0, 8*sizeof(float)); + + float sumf = 0; + for (int i = 0; i < nb; ++i) { + const uint8_t * GGML_RESTRICT q4 = x[i].qs; + const int8_t * GGML_RESTRICT q8 = y[i].qs; + memset(aux32, 0, 8*sizeof(int32_t)); + int8_t * GGML_RESTRICT a = aux8; + for (int j = 0; j < QK_K/64; ++j) { + for (int l = 0; l < 32; ++l) a[l] = (int8_t)(q4[l] & 0xF); + a += 32; + for (int l = 0; l < 32; ++l) a[l] = (int8_t)(q4[l] >> 4); + a += 32; q4 += 32; + } + memcpy(utmp, x[i].scales, 12); + utmp[3] = ((utmp[2] >> 4) & kmask2) | (((utmp[1] >> 6) & kmask3) << 4); + const uint32_t uaux = utmp[1] & kmask1; + utmp[1] = (utmp[2] & kmask2) | (((utmp[0] >> 6) & kmask3) << 4); + utmp[2] = uaux; + utmp[0] &= kmask1; + + int sumi = 0; + for (int j = 0; j < QK_K/16; ++j) sumi += y[i].bsums[j] * mins[j/2]; + a = aux8; + int is = 0; + for (int j = 0; j < QK_K/32; ++j) { + int32_t scale = scales[is++]; + for (int l = 0; l < 8; ++l) aux16[l] = q8[l] * a[l]; + for (int l = 0; l < 8; ++l) aux32[l] += scale * aux16[l]; + q8 += 8; a += 8; + for (int l = 0; l < 8; ++l) aux16[l] = q8[l] * a[l]; + for (int l = 0; l < 8; ++l) aux32[l] += scale * aux16[l]; + q8 += 8; a += 8; + for (int l = 0; l < 8; ++l) aux16[l] = q8[l] * a[l]; + for (int l = 0; l < 8; ++l) aux32[l] += scale * aux16[l]; + q8 += 8; a += 8; + for (int l = 0; l < 8; ++l) aux16[l] = q8[l] * a[l]; + for (int l = 0; l < 8; ++l) aux32[l] += scale * aux16[l]; + q8 += 8; a += 8; + } + const float d = GGML_CPU_FP16_TO_FP32(x[i].d) * y[i].d; + for (int l = 0; l < 8; ++l) sums[l] += d * aux32[l]; + const float dmin = GGML_CPU_FP16_TO_FP32(x[i].dmin) * y[i].d; + sumf -= dmin * sumi; + } + for (int l = 0; l < 8; ++l) sumf += sums[l]; + *s = sumf; +} + +void ggml_vec_dot_q5_K_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) { + assert(n % QK_K == 0); + assert(nrc == 1); + UNUSED(nrc); + UNUSED(bx); + UNUSED(by); + UNUSED(bs); + + const block_q5_K * GGML_RESTRICT x = vx; + const block_q8_K * GGML_RESTRICT y = vy; + + const int nb = n / QK_K; + + static const uint32_t kmask1 = 0x3f3f3f3f; + static const uint32_t kmask2 = 0x0f0f0f0f; + static const uint32_t kmask3 = 0x03030303; + + uint32_t utmp[4]; + + const uint8_t * scales = (const uint8_t*)&utmp[0]; + const uint8_t * mins = (const uint8_t*)&utmp[2]; + + int8_t aux8[QK_K]; + int16_t aux16[8]; + float sums [8]; + int32_t aux32[8]; + memset(sums, 0, 8*sizeof(float)); + + float sumf = 0; + for (int i = 0; i < nb; ++i) { + const uint8_t * GGML_RESTRICT q4 = x[i].qs; + const uint8_t * GGML_RESTRICT hm = x[i].qh; + const int8_t * GGML_RESTRICT q8 = y[i].qs; + memset(aux32, 0, 8*sizeof(int32_t)); + int8_t * GGML_RESTRICT a = aux8; + uint8_t m = 1; + for (int j = 0; j < QK_K/64; ++j) { + for (int l = 0; l < 32; ++l) a[l] = (int8_t)(q4[l] & 0xF); + for (int l = 0; l < 32; ++l) a[l] += (hm[l] & m ? 16 : 0); + a += 32; m <<= 1; + for (int l = 0; l < 32; ++l) a[l] = (int8_t)(q4[l] >> 4); + for (int l = 0; l < 32; ++l) a[l] += (hm[l] & m ? 16 : 0); + a += 32; m <<= 1; + q4 += 32; + } + memcpy(utmp, x[i].scales, 12); + utmp[3] = ((utmp[2] >> 4) & kmask2) | (((utmp[1] >> 6) & kmask3) << 4); + const uint32_t uaux = utmp[1] & kmask1; + utmp[1] = (utmp[2] & kmask2) | (((utmp[0] >> 6) & kmask3) << 4); + utmp[2] = uaux; + utmp[0] &= kmask1; + + int sumi = 0; + for (int j = 0; j < QK_K/16; ++j) sumi += y[i].bsums[j] * mins[j/2]; + a = aux8; + int is = 0; + for (int j = 0; j < QK_K/32; ++j) { + int32_t scale = scales[is++]; + for (int l = 0; l < 8; ++l) aux16[l] = q8[l] * a[l]; + for (int l = 0; l < 8; ++l) aux32[l] += scale * aux16[l]; + q8 += 8; a += 8; + for (int l = 0; l < 8; ++l) aux16[l] = q8[l] * a[l]; + for (int l = 0; l < 8; ++l) aux32[l] += scale * aux16[l]; + q8 += 8; a += 8; + for (int l = 0; l < 8; ++l) aux16[l] = q8[l] * a[l]; + for (int l = 0; l < 8; ++l) aux32[l] += scale * aux16[l]; + q8 += 8; a += 8; + for (int l = 0; l < 8; ++l) aux16[l] = q8[l] * a[l]; + for (int l = 0; l < 8; ++l) aux32[l] += scale * aux16[l]; + q8 += 8; a += 8; + } + const float d = GGML_CPU_FP16_TO_FP32(x[i].d) * y[i].d; + for (int l = 0; l < 8; ++l) sums[l] += d * aux32[l]; + const float dmin = GGML_CPU_FP16_TO_FP32(x[i].dmin) * y[i].d; + sumf -= dmin * sumi; + } + for (int l = 0; l < 8; ++l) sumf += sums[l]; + *s = sumf; +} + +void ggml_vec_dot_q6_K_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) { + assert(n % QK_K == 0); + assert(nrc == 1); + UNUSED(nrc); + UNUSED(bx); + UNUSED(by); + UNUSED(bs); + + const block_q6_K * GGML_RESTRICT x = vx; + const block_q8_K * GGML_RESTRICT y = vy; + + const int nb = n / QK_K; + + int8_t aux8[QK_K]; + int16_t aux16[8]; + float sums [8]; + int32_t aux32[8]; + memset(sums, 0, 8*sizeof(float)); + + float sumf = 0; + for (int i = 0; i < nb; ++i) { + const uint8_t * GGML_RESTRICT q4 = x[i].ql; + const uint8_t * GGML_RESTRICT qh = x[i].qh; + const int8_t * GGML_RESTRICT q8 = y[i].qs; + memset(aux32, 0, 8*sizeof(int32_t)); + int8_t * GGML_RESTRICT a = aux8; + for (int j = 0; j < QK_K; j += 128) { + for (int l = 0; l < 32; ++l) { + a[l + 0] = (int8_t)((q4[l + 0] & 0xF) | (((qh[l] >> 0) & 3) << 4)) - 32; + a[l + 32] = (int8_t)((q4[l + 32] & 0xF) | (((qh[l] >> 2) & 3) << 4)) - 32; + a[l + 64] = (int8_t)((q4[l + 0] >> 4) | (((qh[l] >> 4) & 3) << 4)) - 32; + a[l + 96] = (int8_t)((q4[l + 32] >> 4) | (((qh[l] >> 6) & 3) << 4)) - 32; + } + a += 128; + q4 += 64; + qh += 32; + } + a = aux8; + int is = 0; + for (int j = 0; j < QK_K/16; ++j) { + int scale = x[i].scales[is++]; + for (int l = 0; l < 8; ++l) aux16[l] = q8[l] * a[l]; + for (int l = 0; l < 8; ++l) aux32[l] += scale * aux16[l]; + q8 += 8; a += 8; + for (int l = 0; l < 8; ++l) aux16[l] = q8[l] * a[l]; + for (int l = 0; l < 8; ++l) aux32[l] += scale * aux16[l]; + q8 += 8; a += 8; + } + const float d = GGML_CPU_FP16_TO_FP32(x[i].d) * y[i].d; + for (int l = 0; l < 8; ++l) sums[l] += d * aux32[l]; + } + for (int l = 0; l < 8; ++l) sumf += sums[l]; + *s = sumf; +} + +void ggml_vec_dot_iq2_xxs_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) { + assert(n % QK_K == 0); + assert(nrc == 1); + UNUSED(nrc); + UNUSED(bx); + UNUSED(by); + UNUSED(bs); + + const block_iq2_xxs * GGML_RESTRICT x = vx; + const block_q8_K * GGML_RESTRICT y = vy; + + const int nb = n / QK_K; + + uint32_t aux32[2]; + const uint8_t * aux8 = (const uint8_t *)aux32; + + float sumf = 0.f; + for (int i = 0; i < nb; ++i) { + const float d = GGML_CPU_FP16_TO_FP32(x[i].d) * y[i].d; + const uint16_t * GGML_RESTRICT q2 = x[i].qs; + const int8_t * GGML_RESTRICT q8 = y[i].qs; + int32_t bsum = 0; + for (int ib32 = 0; ib32 < QK_K/32; ++ib32) { + memcpy(aux32, q2, 2*sizeof(uint32_t)); + q2 += 4; + const uint32_t ls = 2*(aux32[1] >> 28) + 1; + int32_t sumi = 0; + for (int l = 0; l < 4; ++l) { + const uint8_t * grid = (const uint8_t *)(iq2xxs_grid + aux8[l]); + const uint8_t signs = ksigns_iq2xs[(aux32[1] >> 7*l) & 127]; + for (int j = 0; j < 8; ++j) { + sumi += grid[j] * q8[j] * (signs & kmask_iq2xs[j] ? -1 : 1); + } + q8 += 8; + } + bsum += sumi * ls; + } + sumf += d * bsum; + } + *s = 0.125f * sumf; +} + +void ggml_vec_dot_iq2_xs_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) { + assert(n % QK_K == 0); + assert(nrc == 1); + UNUSED(nrc); + UNUSED(bx); + UNUSED(by); + UNUSED(bs); + + const block_iq2_xs * GGML_RESTRICT x = vx; + const block_q8_K * GGML_RESTRICT y = vy; + + const int nb = n / QK_K; + + float sumf = 0.f; + for (int i = 0; i < nb; ++i) { + const float d = GGML_CPU_FP16_TO_FP32(x[i].d) * y[i].d; + const uint16_t * GGML_RESTRICT q2 = x[i].qs; + const uint8_t * GGML_RESTRICT sc = x[i].scales; + const int8_t * GGML_RESTRICT q8 = y[i].qs; + int32_t bsum = 0; + for (int ib32 = 0; ib32 < QK_K/32; ++ib32) { + const uint16_t ls1 = 2*(sc[ib32] & 0xf) + 1; + const uint16_t ls2 = 2*(sc[ib32] >> 4) + 1; + int32_t sumi = 0; + for (int l = 0; l < 2; ++l) { + const uint8_t * grid = (const uint8_t *)(iq2xs_grid + (q2[l] & 511)); + const uint8_t signs = ksigns_iq2xs[q2[l] >> 9]; + for (int j = 0; j < 8; ++j) { + sumi += grid[j] * q8[j] * (signs & kmask_iq2xs[j] ? -1 : 1); + } + q8 += 8; + } + bsum += sumi * ls1; + sumi = 0; + for (int l = 2; l < 4; ++l) { + const uint8_t * grid = (const uint8_t *)(iq2xs_grid + (q2[l] & 511)); + const uint8_t signs = ksigns_iq2xs[q2[l] >> 9]; + for (int j = 0; j < 8; ++j) { + sumi += grid[j] * q8[j] * (signs & kmask_iq2xs[j] ? -1 : 1); + } + q8 += 8; + } + bsum += sumi * ls2; + q2 += 4; + } + sumf += d * bsum; + } + *s = 0.125f * sumf; +} + +void ggml_vec_dot_iq2_s_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) { + assert(n % QK_K == 0); + assert(nrc == 1); + UNUSED(nrc); + UNUSED(bx); + UNUSED(by); + UNUSED(bs); + + const block_iq2_s * GGML_RESTRICT x = vx; + const block_q8_K * GGML_RESTRICT y = vy; + + const int nb = n / QK_K; + + float sumf = 0; + for (int i = 0; i < nb; i++) { + + const float d = GGML_CPU_FP16_TO_FP32(x[i].d) * y[i].d; + const int8_t * q8 = y[i].qs; + const uint8_t * qs = x[i].qs; + const uint8_t * qh = x[i].qh; + const uint8_t * signs = qs + QK_K/8; + + int bsum = 0; + for (int ib32 = 0; ib32 < QK_K/32; ++ib32) { + int ls1 = 1 + 2*(x[i].scales[ib32] & 0xf); + int ls2 = 1 + 2*(x[i].scales[ib32] >> 4); + int sumi1 = 0, sumi2 = 0; + for (int l = 0; l < 2; ++l) { + const uint8_t * grid = (const uint8_t *)(iq2s_grid + (qs[l] | (qh[ib32] << (8-2*l) & 0x300))); + for (int j = 0; j < 8; ++j) { + sumi1 += q8[j] * grid[j] * (signs[l] & kmask_iq2xs[j] ? -1 : 1); + } + q8 += 8; + } + for (int l = 2; l < 4; ++l) { + const uint8_t * grid = (const uint8_t *)(iq2s_grid + (qs[l] | (qh[ib32] << (8-2*l) & 0x300))); + for (int j = 0; j < 8; ++j) { + sumi2 += q8[j] * grid[j] * (signs[l] & kmask_iq2xs[j] ? -1 : 1); + } + q8 += 8; + } + bsum += ls1 * sumi1 + ls2 * sumi2; + qs += 4; + signs += 4; + } + + sumf += d * bsum; + } + + *s = 0.125f * sumf; +} + +void ggml_vec_dot_iq3_xxs_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) { + assert(n % QK_K == 0); + assert(nrc == 1); + UNUSED(nrc); + UNUSED(bx); + UNUSED(by); + UNUSED(bs); + + const block_iq3_xxs * GGML_RESTRICT x = vx; + const block_q8_K * GGML_RESTRICT y = vy; + + const int nb = n / QK_K; + + uint32_t aux32; + + float sumf = 0.f; + for (int i = 0; i < nb; ++i) { + const float d = GGML_CPU_FP16_TO_FP32(x[i].d) * y[i].d; + const uint8_t * GGML_RESTRICT q3 = x[i].qs; + const uint8_t * GGML_RESTRICT gas = x[i].qs + QK_K/4; + const int8_t * GGML_RESTRICT q8 = y[i].qs; + int32_t bsum = 0; + for (int ib32 = 0; ib32 < QK_K/32; ++ib32) { + memcpy(&aux32, gas, sizeof(uint32_t)); gas += sizeof(uint32_t); + const uint32_t ls = 2*(aux32 >> 28) + 1; + int32_t sumi = 0; + for (int l = 0; l < 4; ++l) { + const uint8_t * grid1 = (const uint8_t *)(iq3xxs_grid + q3[2*l+0]); + const uint8_t * grid2 = (const uint8_t *)(iq3xxs_grid + q3[2*l+1]); + const uint8_t signs = ksigns_iq2xs[(aux32 >> 7*l) & 127]; + for (int j = 0; j < 4; ++j) { + sumi += grid1[j] * q8[j+0] * (signs & kmask_iq2xs[j+0] ? -1 : 1); + sumi += grid2[j] * q8[j+4] * (signs & kmask_iq2xs[j+4] ? -1 : 1); + } + q8 += 8; + } + q3 += 8; + bsum += sumi * ls; + } + sumf += d * bsum; + } + *s = 0.25f * sumf; +} + +void ggml_vec_dot_iq3_s_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) { + assert(n % QK_K == 0); + assert(nrc == 1); + UNUSED(nrc); + UNUSED(bx); + UNUSED(by); + UNUSED(bs); + + const block_iq3_s * GGML_RESTRICT x = vx; + const block_q8_K * GGML_RESTRICT y = vy; + + const int nb = n / QK_K; + + float sumf = 0.f; + for (int i = 0; i < nb; ++i) { + const float d = GGML_CPU_FP16_TO_FP32(x[i].d) * y[i].d; + const uint8_t * GGML_RESTRICT qs = x[i].qs; + const uint8_t * GGML_RESTRICT qh = x[i].qh; + const uint8_t * GGML_RESTRICT signs = x[i].signs; + const int8_t * GGML_RESTRICT q8 = y[i].qs; + int32_t bsum = 0; + for (int ib32 = 0; ib32 < QK_K/32; ib32 += 2) { + const uint32_t ls1 = 2*(x[i].scales[ib32/2] & 0xf) + 1; + const uint32_t ls2 = 2*(x[i].scales[ib32/2] >> 4) + 1; + int32_t sumi = 0; + for (int l = 0; l < 4; ++l) { + const uint8_t * grid1 = (const uint8_t *)(iq3s_grid + (qs[2*l+0] | ((qh[ib32+0] << (8-2*l)) & 256))); + const uint8_t * grid2 = (const uint8_t *)(iq3s_grid + (qs[2*l+1] | ((qh[ib32+0] << (7-2*l)) & 256))); + for (int j = 0; j < 4; ++j) { + sumi += grid1[j] * q8[j+0] * (signs[l] & kmask_iq2xs[j+0] ? -1 : 1); + sumi += grid2[j] * q8[j+4] * (signs[l] & kmask_iq2xs[j+4] ? -1 : 1); + } + q8 += 8; + } + qs += 8; + signs += 4; + bsum += sumi * ls1; + sumi = 0; + for (int l = 0; l < 4; ++l) { + const uint8_t * grid1 = (const uint8_t *)(iq3s_grid + (qs[2*l+0] | ((qh[ib32+1] << (8-2*l)) & 256))); + const uint8_t * grid2 = (const uint8_t *)(iq3s_grid + (qs[2*l+1] | ((qh[ib32+1] << (7-2*l)) & 256))); + for (int j = 0; j < 4; ++j) { + sumi += grid1[j] * q8[j+0] * (signs[l] & kmask_iq2xs[j+0] ? -1 : 1); + sumi += grid2[j] * q8[j+4] * (signs[l] & kmask_iq2xs[j+4] ? -1 : 1); + } + q8 += 8; + } + qs += 8; + signs += 4; + bsum += sumi * ls2; + } + sumf += d * bsum; + } + *s = sumf; +} + +void ggml_vec_dot_iq1_s_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) { + assert(n % QK_K == 0); + assert(nrc == 1); + UNUSED(nrc); + UNUSED(bx); + UNUSED(by); + UNUSED(bs); + + const block_iq1_s * GGML_RESTRICT x = vx; + const block_q8_K * GGML_RESTRICT y = vy; + + const int nb = n / QK_K; + + float sumf = 0; + for (int i = 0; i < nb; i++) { + + const int8_t * q8 = y[i].qs; + const uint8_t * qs = x[i].qs; + const uint16_t * qh = x[i].qh; + + int sumi = 0, sumi1 = 0; + for (int ib = 0; ib < QK_K/32; ++ib) { + const int ls = 2*((qh[ib] >> 12) & 7) + 1; + const int delta = qh[ib] & 0x8000 ? -1 : 1; + int lsum = 0; + for (int l = 0; l < 4; ++l) { + const int8_t * grid = (const int8_t *)(iq1s_grid + (qs[l] | (((qh[ib] >> 3*l) & 7) << 8))); + for (int j = 0; j < 8; ++j) { + lsum += q8[j] * grid[j]; + } + q8 += 8; + } + sumi += ls * lsum; + sumi1 += ls * delta * (y[i].bsums[2*ib+0] + y[i].bsums[2*ib+1]); + qs += 4; + } + + sumf += GGML_CPU_FP16_TO_FP32(x[i].d) * y[i].d * (sumi + IQ1S_DELTA * sumi1); + } + + *s = sumf; +} + +void ggml_vec_dot_iq1_m_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) { + assert(n % QK_K == 0); + assert(nrc == 1); + UNUSED(nrc); + UNUSED(bx); + UNUSED(by); + UNUSED(bs); + + const block_iq1_m * GGML_RESTRICT x = vx; + const block_q8_K * GGML_RESTRICT y = vy; + + const int nb = n / QK_K; + + iq1m_scale_t scale; + + int sum1[2], sum2[2], delta[4]; + + float sumf = 0; + for (int i = 0; i < nb; i++) { + + const int8_t * q8 = y[i].qs; + const uint8_t * qs = x[i].qs; + const uint8_t * qh = x[i].qh; + const uint16_t * sc = (const uint16_t *)x[i].scales; + + scale.u16 = (sc[0] >> 12) | ((sc[1] >> 8) & 0x00f0) | ((sc[2] >> 4) & 0x0f00) | (sc[3] & 0xf000); + + int sumi1 = 0, sumi2 = 0; + for (int ib = 0; ib < QK_K/32; ++ib) { + delta[0] = qh[0] & 0x08 ? -1 : 1; + delta[1] = qh[0] & 0x80 ? -1 : 1; + delta[2] = qh[1] & 0x08 ? -1 : 1; + delta[3] = qh[1] & 0x80 ? -1 : 1; + sum1[0] = sum1[1] = sum2[0] = sum2[1] = 0; + for (int l = 0; l < 4; ++l) { + const int8_t * grid = (const int8_t *)(iq1s_grid + (qs[l] | (((uint16_t)qh[l/2] << (8 - 4*(l%2))) & 0x700))); + int lsum1 = 0, lsum2 = 0; + for (int j = 0; j < 8; ++j) { + lsum1 += q8[j] * grid[j]; + lsum2 += q8[j]; + } + q8 += 8; + sum1[l/2] += lsum1; + sum2[l/2] += lsum2*delta[l]; + } + + const int ls1 = 2*((sc[ib/2] >> (6*(ib%2)+0)) & 0x7) + 1; + const int ls2 = 2*((sc[ib/2] >> (6*(ib%2)+3)) & 0x7) + 1; + + sumi1 += sum1[0] * ls1 + sum1[1] * ls2; + sumi2 += sum2[0] * ls1 + sum2[1] * ls2; + qs += 4; + qh += 2; + } + + sumf += GGML_CPU_FP16_TO_FP32(scale.f16) * y[i].d * (sumi1 + IQ1M_DELTA * sumi2); + } + + *s = sumf; +} + +void ggml_vec_dot_iq4_nl_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) { + assert(nrc == 1); + UNUSED(nrc); + UNUSED(bx); + UNUSED(by); + UNUSED(bs); + assert(n % QK4_NL == 0); + static_assert(QK4_NL == QK8_0, "QK4_NL and QK8_0 must be the same"); + + const block_iq4_nl * GGML_RESTRICT x = vx; + const block_q8_0 * GGML_RESTRICT y = vy; + + const int nb = n / QK4_NL; + + int ib = 0; + float sumf = 0; + + for (; ib < nb; ++ib) { + const float d = GGML_CPU_FP16_TO_FP32(y[ib].d)*GGML_CPU_FP16_TO_FP32(x[ib].d); + int sumi1 = 0, sumi2 = 0; + for (int j = 0; j < QK4_NL/2; ++j) { + sumi1 += y[ib].qs[j+ 0] * kvalues_iq4nl[x[ib].qs[j] & 0xf]; + sumi2 += y[ib].qs[j+QK4_NL/2] * kvalues_iq4nl[x[ib].qs[j] >> 4]; + } + sumf += d * (sumi1 + sumi2); + } + *s = sumf; +} + +void ggml_vec_dot_iq4_xs_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) { + assert(nrc == 1); + UNUSED(nrc); + UNUSED(bx); + UNUSED(by); + UNUSED(bs); + assert(n % QK_K == 0); + + const block_iq4_xs * GGML_RESTRICT x = vx; + const block_q8_K * GGML_RESTRICT y = vy; + + const int nb = n / QK_K; + + float sumf = 0; + for (int ibl = 0; ibl < nb; ++ibl) { + const float d4d8 = GGML_CPU_FP16_TO_FP32(x[ibl].d) * y[ibl].d; + uint16_t h = x[ibl].scales_h; + const uint8_t * qs = x[ibl].qs; + const int8_t * q8 = y[ibl].qs; + for (int ib = 0; ib < QK_K/32; ib += 2) { + const uint8_t ls1 = (x[ibl].scales_l[ib/2] & 0xf) | ((h << 4) & 0x30); + const uint8_t ls2 = (x[ibl].scales_l[ib/2] >> 4) | ((h << 2) & 0x30); + h >>= 4; + const float d1 = d4d8*(ls1 - 32); + const float d2 = d4d8*(ls2 - 32); + int sumi1 = 0, sumi2 = 0; + for (int j = 0; j < 16; ++j) { + sumi1 += q8[j+ 0] * kvalues_iq4nl[qs[j] & 0xf]; + sumi2 += q8[j+16] * kvalues_iq4nl[qs[j] >> 4]; + } + sumf += d1 * (sumi1 + sumi2); + qs += 16; + q8 += 32; + sumi1 = sumi2 = 0; + for (int j = 0; j < 16; ++j) { + sumi1 += q8[j+ 0] * kvalues_iq4nl[qs[j] & 0xf]; + sumi2 += q8[j+16] * kvalues_iq4nl[qs[j] >> 4]; + } + sumf += d2 * (sumi1 + sumi2); + qs += 16; + q8 += 32; + } + } + *s = sumf; +} + +// ============================ 4-bit non-linear quants + +void quantize_row_iq4_nl(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k) { + assert(k % QK4_NL == 0); + quantize_row_iq4_nl_ref(x, y, k); +} + +void quantize_row_iq4_xs(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k) { + assert(k % QK_K == 0); + quantize_iq4_xs(x, y, 1, k, NULL); +} diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cpu/quants.h b/python/freetoken/kernel/csrc/ggml/ggml-cpu/quants.h new file mode 100644 index 0000000000..93ea7eeffe --- /dev/null +++ b/python/freetoken/kernel/csrc/ggml/ggml-cpu/quants.h @@ -0,0 +1,106 @@ +#pragma once + +#define GGML_COMMON_DECL_C +#include "ggml-common.h" + +#include "ggml.h" + +// GGML CPU internal header + +#ifdef __cplusplus +extern "C" { +#endif + +// Quantization +void quantize_row_q1_0(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +void quantize_row_q2_0(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +void quantize_row_q4_0(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +void quantize_row_q4_1(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +void quantize_row_q5_0(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +void quantize_row_q5_1(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +void quantize_row_q8_0(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +void quantize_row_q8_1(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); + +void quantize_row_mxfp4(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +void quantize_row_nvfp4(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); + +void quantize_row_q2_K(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +void quantize_row_q3_K(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +void quantize_row_q4_K(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +void quantize_row_q5_K(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +void quantize_row_q6_K(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +void quantize_row_q8_K(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); + +void quantize_row_tq1_0(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +void quantize_row_tq2_0(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); + +void quantize_row_iq4_nl (const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +void quantize_row_iq4_xs (const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); + +// Dot product +void ggml_vec_dot_q1_0_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); +void ggml_vec_dot_q2_0_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); +void ggml_vec_dot_q4_0_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); +void ggml_vec_dot_q4_1_q8_1(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); +void ggml_vec_dot_q5_0_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); +void ggml_vec_dot_q5_1_q8_1(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); +void ggml_vec_dot_q8_0_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); + +void ggml_vec_dot_mxfp4_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); +void ggml_vec_dot_nvfp4_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); + +void ggml_vec_dot_q2_K_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); +void ggml_vec_dot_q3_K_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); +void ggml_vec_dot_q4_K_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); +void ggml_vec_dot_q5_K_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); +void ggml_vec_dot_q6_K_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); + +void ggml_vec_dot_tq1_0_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); +void ggml_vec_dot_tq2_0_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); + +void ggml_vec_dot_iq2_xxs_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); +void ggml_vec_dot_iq2_xs_q8_K (int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); +void ggml_vec_dot_iq2_s_q8_K (int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); +void ggml_vec_dot_iq3_xxs_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); +void ggml_vec_dot_iq1_s_q8_K (int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); +void ggml_vec_dot_iq1_m_q8_K (int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); +void ggml_vec_dot_iq4_nl_q8_0 (int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); +void ggml_vec_dot_iq4_xs_q8_K (int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); +void ggml_vec_dot_iq3_s_q8_K (int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); + +// Generic implementation +void quantize_row_q8_0_generic(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k); +void quantize_row_q8_1_generic(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k); +void quantize_row_q8_K_generic(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +void ggml_vec_dot_q1_0_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); +void ggml_vec_dot_q2_0_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); +void ggml_vec_dot_q4_0_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); +void ggml_vec_dot_q4_1_q8_1_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); +void ggml_vec_dot_q5_0_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); +void ggml_vec_dot_q5_1_q8_1_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); +void ggml_vec_dot_q8_0_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); + +void ggml_vec_dot_mxfp4_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); +void ggml_vec_dot_nvfp4_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); + +void ggml_vec_dot_tq1_0_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); +void ggml_vec_dot_tq2_0_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); + +void ggml_vec_dot_q2_K_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); +void ggml_vec_dot_q3_K_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); +void ggml_vec_dot_q4_K_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); +void ggml_vec_dot_q5_K_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); +void ggml_vec_dot_q6_K_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); +void ggml_vec_dot_iq2_xxs_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); +void ggml_vec_dot_iq2_xs_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); +void ggml_vec_dot_iq2_s_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); +void ggml_vec_dot_iq3_xxs_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); +void ggml_vec_dot_iq3_s_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); +void ggml_vec_dot_iq1_s_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); +void ggml_vec_dot_iq1_m_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); +void ggml_vec_dot_iq4_nl_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); +void ggml_vec_dot_iq4_xs_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); + +#ifdef __cplusplus +} +#endif diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cpu/simd-mappings.h b/python/freetoken/kernel/csrc/ggml/ggml-cpu/simd-mappings.h new file mode 100644 index 0000000000..fb53ecdf42 --- /dev/null +++ b/python/freetoken/kernel/csrc/ggml/ggml-cpu/simd-mappings.h @@ -0,0 +1,1366 @@ +#pragma once + +#include "ggml-cpu-impl.h" + +#ifdef __ARM_FEATURE_SVE +#include +#endif // __ARM_FEATURE_SVE + +#if defined(__ARM_NEON) && !defined(__CUDACC__) && !defined(__MUSACC__) +// if YCM cannot find , make a symbolic link to it, for example: +// +// $ ln -sfn /Library/Developer/CommandLineTools/usr/lib/clang/13.1.6/include/arm_neon.h ./src/ +// +#include +#endif + +#if defined(__riscv_v_intrinsic) +#include +#endif + +#ifdef __cplusplus +extern "C" { +#endif + +// +// simd mappings +// + +// FP16 to FP32 conversion + +// 16-bit float +// on Arm, we use __fp16, which requires the IEEE fp16 format: implied on +// AArch64, selected by -mfp16-format=ieee on 32 bit Arm, where the compiler +// may otherwise reject the type +// on x86, we use uint16_t +// +// for old CUDA compilers (<= 11), we use uint16_t: ref https://github.com/ggml-org/llama.cpp/pull/10616 +// for MUSA compilers , we use uint16_t: ref https://github.com/ggml-org/llama.cpp/pull/11843 +// +#if defined(__ARM_NEON) && defined(__ARM_FP16_FORMAT_IEEE) && !(defined(__CUDACC__) && __CUDACC_VER_MAJOR__ <= 11) && !defined(__MUSACC__) + #define GGML_CPU_COMPUTE_FP16_TO_FP32(x) neon_compute_fp16_to_fp32(x) + #define GGML_CPU_COMPUTE_FP32_TO_FP16(x) neon_compute_fp32_to_fp16(x) + + #define GGML_CPU_FP16_TO_FP32(x) GGML_CPU_COMPUTE_FP16_TO_FP32(x) + + static inline float neon_compute_fp16_to_fp32(ggml_fp16_t h) { + __fp16 tmp; + memcpy(&tmp, &h, sizeof(ggml_fp16_t)); + return (float)tmp; + } + + static inline ggml_fp16_t neon_compute_fp32_to_fp16(float f) { + ggml_fp16_t res; + __fp16 tmp = f; + memcpy(&res, &tmp, sizeof(ggml_fp16_t)); + return res; + } +#elif defined(__F16C__) + #ifdef _MSC_VER + #define GGML_CPU_COMPUTE_FP16_TO_FP32(x) _mm_cvtss_f32(_mm_cvtph_ps(_mm_cvtsi32_si128(x))) + #define GGML_CPU_COMPUTE_FP32_TO_FP16(x) _mm_extract_epi16(_mm_cvtps_ph(_mm_set_ss(x), 0), 0) + #else + #define GGML_CPU_COMPUTE_FP16_TO_FP32(x) _cvtsh_ss(x) + #define GGML_CPU_COMPUTE_FP32_TO_FP16(x) _cvtss_sh(x, 0) + #endif +#elif defined(__POWER9_VECTOR__) + #define GGML_CPU_COMPUTE_FP16_TO_FP32(x) power_compute_fp16_to_fp32(x) + #define GGML_CPU_COMPUTE_FP32_TO_FP16(x) power_compute_fp32_to_fp16(x) + /* the inline asm below is about 12% faster than the lookup method */ + #define GGML_CPU_FP16_TO_FP32(x) GGML_CPU_COMPUTE_FP16_TO_FP32(x) + #define GGML_CPU_FP32_TO_FP16(x) GGML_CPU_COMPUTE_FP32_TO_FP16(x) + + static inline float power_compute_fp16_to_fp32(ggml_fp16_t h) { + float f; + double d; + __asm__( + "mtfprd %0,%2\n" + "xscvhpdp %0,%0\n" + "frsp %1,%0\n" : + /* temp */ "=d"(d), + /* out */ "=f"(f): + /* in */ "r"(h)); + return f; + } + + static inline ggml_fp16_t power_compute_fp32_to_fp16(float f) { + double d; + ggml_fp16_t r; + __asm__( /* xscvdphp can work on double or single precision */ + "xscvdphp %0,%2\n" + "mffprd %1,%0\n" : + /* temp */ "=d"(d), + /* out */ "=r"(r): + /* in */ "f"(f)); + return r; + } +#elif defined(__riscv) && defined(__riscv_zfhmin) + static inline float riscv_compute_fp16_to_fp32(ggml_fp16_t h) { + _Float16 hf; + memcpy(&hf, &h, sizeof(ggml_fp16_t)); + return hf; + } + + static inline ggml_fp16_t riscv_compute_fp32_to_fp16(float f) { + ggml_fp16_t res; + _Float16 hf = (_Float16)f; + memcpy(&res, &hf, sizeof(ggml_fp16_t)); + return res; + } + + #define GGML_CPU_COMPUTE_FP16_TO_FP32(x) riscv_compute_fp16_to_fp32(x) + #define GGML_CPU_COMPUTE_FP32_TO_FP16(x) riscv_compute_fp32_to_fp16(x) + #define GGML_CPU_FP16_TO_FP32(x) GGML_CPU_COMPUTE_FP16_TO_FP32(x) + #define GGML_CPU_FP32_TO_FP16(x) GGML_CPU_COMPUTE_FP32_TO_FP16(x) +#endif + +// precomputed f32 table for f16 (256 KB) +// defined in ggml-cpu.c, initialized in ggml_cpu_init() +extern float ggml_table_f32_f16[1 << 16]; + +// precomputed f32 table for e8m0 half (1 KB) +// defined in ggml-cpu.c, initialized in ggml_cpu_init() +extern float ggml_table_f32_e8m0_half[1 << 8]; + +// precomputed f32 table for ue4m3 (1 KB) +// defined in ggml-cpu.c, initialized in ggml_cpu_init() +extern float ggml_table_f32_ue4m3[1 << 8]; + +// Use lookup table for E8M0 on x86 (faster than bit manipulation) +#if defined(__AVX__) || defined(__AVX2__) || defined(__AVX512F__) +#define GGML_CPU_E8M0_TO_FP32_HALF(x) ggml_table_f32_e8m0_half[(uint8_t)(x)] +#else +#define GGML_CPU_E8M0_TO_FP32_HALF(x) GGML_E8M0_TO_FP32_HALF(x) +#endif + +// Use lookup table for UE4M3 on x86 and ARM (faster than bit manipulation) +#if defined(__AVX__) || defined(__AVX2__) || defined(__AVX512F__) || defined(__ARM_NEON) +#define GGML_CPU_UE4M3_TO_FP32(x) ggml_table_f32_ue4m3[(uint8_t)(x)] +#else +#define GGML_CPU_UE4M3_TO_FP32(x) ggml_ue4m3_to_fp32(x) +#endif + +// Pre AVX512 CPUs (AVX2 with F16C, such as i5-8th gen or i7-4790) +// do not produce codegen as checks below don't test for __F16C__ +// but this check does. +#if defined(__F16C__) +#define GGML_CPU_FP16_TO_FP32(x) GGML_CPU_COMPUTE_FP16_TO_FP32(x) +#define GGML_CPU_FP32_TO_FP16(x) GGML_CPU_COMPUTE_FP32_TO_FP16(x) +#endif + +// On ARM NEON, it's quicker to directly convert x -> x instead of calling into ggml_lookup_fp16_to_fp32, +// so we define GGML_CPU_FP16_TO_FP32 and GGML_CPU_FP32_TO_FP16 elsewhere for NEON. +// This is also true for POWER9. +#if !defined(GGML_CPU_FP16_TO_FP32) +inline static float ggml_lookup_fp16_to_fp32(ggml_fp16_t f) { + uint16_t s; + memcpy(&s, &f, sizeof(uint16_t)); + return ggml_table_f32_f16[s]; +} + +#define GGML_CPU_FP16_TO_FP32(x) ggml_lookup_fp16_to_fp32(x) +#endif + +#if !defined(GGML_CPU_FP32_TO_FP16) +#define GGML_CPU_FP32_TO_FP16(x) GGML_COMPUTE_FP32_TO_FP16(x) +#endif + + +// we define a common set of C macros which map to specific intrinsics based on the current architecture +// we then implement the fundamental computation operations below using only these macros +// adding support for new architectures requires to define the corresponding SIMD macros +// +// GGML_F32_STEP / GGML_F16_STEP +// number of elements to process in a single step +// +// GGML_F32_EPR / GGML_F16_EPR +// number of elements to fit in a single register +// + +#if defined(__ARM_FEATURE_SVE) && defined(__ARM_FEATURE_FMA) + +#define GGML_SIMD + +// F32 SVE +#define GGML_F32_EPR 8 +#define DEFAULT_PG svptrue_b32() + +#define GGML_F32xt svfloat32_t +#define GGML_F32xt_ZERO svdup_n_f32(0.0f) +#define GGML_F32xt_SET1(x) svdup_n_f32(x) +#define GGML_F32xt_LOAD_IMPL(pg, a) svld1_f32(pg, a) +#define GGML_F32xt_LOAD(a) GGML_F32xt_LOAD_IMPL(DEFAULT_PG, a) +#define GGML_F32xt_STORE_IMPL(pg, a, b) svst1_f32(pg, a, b) +#define GGML_F32xt_STORE(a, b) GGML_F32xt_STORE_IMPL(DEFAULT_PG, a, b) +#define GGML_F32xt_FMA_IMPL(pg, a, b, c) svmad_f32_m(pg, b, c, a) +#define GGML_F32xt_FMA(a, b, c) GGML_F32xt_FMA_IMPL(DEFAULT_PG, a, b, c) +#define GGML_F32xt_ADD_IMPL(pg, a, b) svadd_f32_m(pg, a, b) +#define GGML_F32xt_ADD(a, b) GGML_F32xt_ADD_IMPL(DEFAULT_PG, a, b) +#define GGML_F32xt_MUL_IMPL(pg, a, b) svmul_f32_m(pg, a, b) +#define GGML_F32xt_MUL(a, b) GGML_F32xt_MUL_IMPL(DEFAULT_PG, a, b) +#define GGML_F32xt_REDUCE_ONE_IMPL(pg, a) svaddv(pg, a) +#define GGML_F32xt_REDUCE_ONE(a) GGML_F32xt_REDUCE_ONE_IMPL(DEFAULT_PG, a) +#define GGML_F32xt_REDUCE_IMPL(pg, res, sum1, sum2, sum3, sum4, sum5, sum6, sum7, sum8) \ +{ \ + sum1 = svadd_f32_m(DEFAULT_PG, sum1, sum2); \ + sum3 = svadd_f32_m(DEFAULT_PG, sum3, sum4); \ + sum5 = svadd_f32_m(DEFAULT_PG, sum5, sum6); \ + sum7 = svadd_f32_m(DEFAULT_PG, sum7, sum8); \ + sum1 = svadd_f32_m(DEFAULT_PG, sum1, sum3); \ + sum5 = svadd_f32_m(DEFAULT_PG, sum5, sum7); \ + sum1 = svadd_f32_m(DEFAULT_PG, sum1, sum5); \ + (res) = (ggml_float) GGML_F32xt_REDUCE_ONE(sum1); \ +} +#define GGML_F32xt_REDUCE(res, sum1, sum2, sum3, sum4, sum5, sum6, sum7, sum8) \ + GGML_F32xt_REDUCE_IMPL(DEFAULT_PG, res, sum1, sum2, sum3, sum4, sum5, sum6, sum7, sum8) + +#define GGML_F32_VEC GGML_F32xt +#define GGML_F32_VEC_ZERO GGML_F32xt_ZERO +#define GGML_F32_VEC_SET1 GGML_F32xt_SET1 +#define GGML_F32_VEC_LOAD GGML_F32xt_LOAD +#define GGML_F32_VEC_STORE GGML_F32xt_STORE +#define GGML_F32_VEC_FMA GGML_F32xt_FMA +#define GGML_F32_VEC_ADD GGML_F32xt_ADD +#define GGML_F32_VEC_MUL GGML_F32xt_MUL +#define GGML_F32_VEC_REDUCE GGML_F32xt_REDUCE + +// F16 SVE +#define DEFAULT_PG32 svptrue_b32() +#define DEFAULT_PG16 svptrue_b16() + +#define GGML_F32Cxt svfloat16_t +#define GGML_F32Cxt_ZERO svdup_n_f16(0.0f) +#define GGML_F32Cxt_SET1(x) svdup_n_f16(x) +#define GGML_F32Cxt_LOAD(p) svld1_f16(DEFAULT_PG16, (const __fp16 *)(p)) +#define GGML_F32Cxt_STORE(dst_ptr, src_vec) svst1_f16(DEFAULT_PG16, (__fp16 *)(dst_ptr), (src_vec)) + +#define GGML_F32Cxt_FMA_IMPL(pg, a, b, c) svmad_f16_x(pg, b, c, a) +#define GGML_F32Cxt_FMA(a, b, c) GGML_F32Cxt_FMA_IMPL(DEFAULT_PG16, a, b, c) +#define GGML_F32Cxt_ADD_IMPL(pg, a, b) svadd_f16_x(pg, a, b) +#define GGML_F32Cxt_ADD(a, b) GGML_F32Cxt_ADD_IMPL(DEFAULT_PG16, a, b) +#define GGML_F32Cxt_MUL_IMPL(pg, a, b) svmul_f16_x(pg, a, b) +#define GGML_F32Cxt_MUL(a, b) GGML_F32Cxt_MUL_IMPL(DEFAULT_PG16, a, b) +#define GGML_F32Cxt_REDUCE GGML_F16xt_REDUCE_MIXED + +#define GGML_F16x_VEC GGML_F32Cxt +#define GGML_F16x_VEC_ZERO GGML_F32Cxt_ZERO +#define GGML_F16x_VEC_SET1 GGML_F32Cxt_SET1 +#define GGML_F16x_VEC_LOAD(p, i) GGML_F32Cxt_LOAD(p) +#define GGML_F16x_VEC_STORE(p, r, i) GGML_F32Cxt_STORE((__fp16 *)(p), r) +#define GGML_F16x_VEC_FMA GGML_F32Cxt_FMA +#define GGML_F16x_VEC_ADD GGML_F32Cxt_ADD +#define GGML_F16x_VEC_MUL GGML_F32Cxt_MUL +#define GGML_F16x_VEC_REDUCE GGML_F32Cxt_REDUCE + +#define GGML_F16xt_REDUCE_ONE_IMPL(pg, a) svaddv_f16(pg, a) +#define GGML_F16xt_REDUCE_ONE(a) GGML_F16xt_REDUCE_ONE_IMPL(DEFAULT_PG16, a) + +#define GGML_F16xt_REDUCE_MIXED_IMPL(pg16, res, sum1, sum2, sum3, sum4) \ +{ \ + sum1 = svadd_f16_x(pg16, sum1, sum2); \ + sum3 = svadd_f16_x(pg16, sum3, sum4); \ + sum1 = svadd_f16_x(pg16, sum1, sum3); \ + __fp16 sum_f16 = svaddv_f16(pg16, sum1); \ + (res) = (ggml_float) sum_f16; \ +} +#define GGML_F16xt_REDUCE_MIXED(res, sum1, sum2, sum3, sum4) \ + GGML_F16xt_REDUCE_MIXED_IMPL(DEFAULT_PG16, res, sum1, sum2, sum3, sum4) + +// F16 NEON + +#if defined(__ARM_FEATURE_FP16_VECTOR_ARITHMETIC) + #define GGML_F16_STEP 32 + #define GGML_F16_EPR 8 + + #define GGML_F16x8 float16x8_t + #define GGML_F16x8_ZERO vdupq_n_f16(0.0f) + #define GGML_F16x8_SET1(x) vdupq_n_f16(x) + #define GGML_F16x8_LOAD(x) vld1q_f16((const __fp16 *)(x)) + #define GGML_F16x8_STORE vst1q_f16 + #define GGML_F16x8_FMA(a, b, c) vfmaq_f16(a, b, c) + #define GGML_F16x8_ADD vaddq_f16 + #define GGML_F16x8_MUL vmulq_f16 + #define GGML_F16x8_REDUCE(res, x) \ + do { \ + int offset = GGML_F16_ARR >> 1; \ + for (int i = 0; i < offset; ++i) { \ + (x)[i] = vaddq_f16((x)[i], (x)[offset+i]); \ + } \ + offset >>= 1; \ + for (int i = 0; i < offset; ++i) { \ + (x)[i] = vaddq_f16((x)[i], (x)[offset+i]); \ + } \ + offset >>= 1; \ + for (int i = 0; i < offset; ++i) { \ + (x)[i] = vaddq_f16((x)[i], (x)[offset+i]); \ + } \ + const float32x4_t t0 = vcvt_f32_f16(vget_low_f16 ((x)[0])); \ + const float32x4_t t1 = vcvt_f32_f16(vget_high_f16((x)[0])); \ + (res) = (ggml_float) vaddvq_f32(vaddq_f32(t0, t1)); \ + } while (0) + + #define GGML_F16_VEC GGML_F16x8 + #define GGML_F16_VEC_ZERO GGML_F16x8_ZERO + #define GGML_F16_VEC_SET1 GGML_F16x8_SET1 + #define GGML_F16_VEC_LOAD(p, i) GGML_F16x8_LOAD(p) + #define GGML_F16_VEC_STORE(p, r, i) GGML_F16x8_STORE((__fp16 *)(p), (r)[i]) + #define GGML_F16_VEC_FMA GGML_F16x8_FMA + #define GGML_F16_VEC_ADD GGML_F16x8_ADD + #define GGML_F16_VEC_MUL GGML_F16x8_MUL + #define GGML_F16_VEC_REDUCE GGML_F16x8_REDUCE +#else + // if FP16 vector arithmetic is not supported, we use FP32 instead + // and take advantage of the vcvt_ functions to convert to/from FP16 + + #define GGML_F16_STEP 16 + #define GGML_F16_EPR 4 + + #define GGML_F32Cx4 float32x4_t + #define GGML_F32Cx4_ZERO vdupq_n_f32(0.0f) + #define GGML_F32Cx4_SET1(x) vdupq_n_f32(x) + #define GGML_F32Cx4_LOAD(x) vcvt_f32_f16(vld1_f16((const __fp16 *)(x))) + #define GGML_F32Cx4_STORE(x, y) vst1_f16(x, vcvt_f16_f32(y)) + #define GGML_F32Cx4_FMA(a, b, c) vfmaq_f32(a, b, c) + #define GGML_F32Cx4_ADD vaddq_f32 + #define GGML_F32Cx4_MUL vmulq_f32 + #define GGML_F32Cx4_REDUCE GGML_F32x4_REDUCE + + #define GGML_F16_VEC GGML_F32Cx4 + #define GGML_F16_VEC_ZERO GGML_F32Cx4_ZERO + #define GGML_F16_VEC_SET1 GGML_F32Cx4_SET1 + #define GGML_F16_VEC_LOAD(p, i) GGML_F32Cx4_LOAD(p) + #define GGML_F16_VEC_STORE(p, r, i) GGML_F32Cx4_STORE((__fp16 *)(p), r[i]) + #define GGML_F16_VEC_FMA GGML_F32Cx4_FMA + #define GGML_F16_VEC_ADD GGML_F32Cx4_ADD + #define GGML_F16_VEC_MUL GGML_F32Cx4_MUL + #define GGML_F16_VEC_REDUCE GGML_F32Cx4_REDUCE +#endif + +#elif defined(__ARM_NEON) && defined(__ARM_FEATURE_FMA) && defined(__ARM_FP16_FORMAT_IEEE) + +#define GGML_SIMD + +// F32 NEON + +#define GGML_F32_STEP 16 +#define GGML_F32_EPR 4 + +#define GGML_F32x4 float32x4_t +#define GGML_F32x4_ZERO vdupq_n_f32(0.0f) +#define GGML_F32x4_SET1(x) vdupq_n_f32(x) +#define GGML_F32x4_LOAD vld1q_f32 +#define GGML_F32x4_STORE vst1q_f32 +#define GGML_F32x4_FMA(a, b, c) vfmaq_f32(a, b, c) +#define GGML_F32x4_ADD vaddq_f32 +#define GGML_F32x4_MUL vmulq_f32 +#define GGML_F32x4_REDUCE_ONE(x) vaddvq_f32(x) +#define GGML_F32x4_REDUCE(res, x) \ +{ \ + int offset = GGML_F32_ARR >> 1; \ + for (int i = 0; i < offset; ++i) { \ + (x)[i] = vaddq_f32((x)[i], (x)[offset+i]); \ + } \ + offset >>= 1; \ + for (int i = 0; i < offset; ++i) { \ + (x)[i] = vaddq_f32((x)[i], (x)[offset+i]); \ + } \ + offset >>= 1; \ + for (int i = 0; i < offset; ++i) { \ + (x)[i] = vaddq_f32((x)[i], (x)[offset+i]); \ + } \ + (res) = (ggml_float) GGML_F32x4_REDUCE_ONE((x)[0]); \ +} + +#define GGML_F32_VEC GGML_F32x4 +#define GGML_F32_VEC_ZERO GGML_F32x4_ZERO +#define GGML_F32_VEC_SET1 GGML_F32x4_SET1 +#define GGML_F32_VEC_LOAD GGML_F32x4_LOAD +#define GGML_F32_VEC_STORE GGML_F32x4_STORE +#define GGML_F32_VEC_FMA GGML_F32x4_FMA +#define GGML_F32_VEC_ADD GGML_F32x4_ADD +#define GGML_F32_VEC_MUL GGML_F32x4_MUL +#define GGML_F32_VEC_REDUCE GGML_F32x4_REDUCE + +// F16 NEON + +#if defined(__ARM_FEATURE_FP16_VECTOR_ARITHMETIC) + #define GGML_F16_STEP 32 + #define GGML_F16_EPR 8 + + #define GGML_F16x8 float16x8_t + #define GGML_F16x8_ZERO vdupq_n_f16(0.0f) + #define GGML_F16x8_SET1(x) vdupq_n_f16(x) + #define GGML_F16x8_LOAD(x) vld1q_f16((const __fp16 *)(x)) + #define GGML_F16x8_STORE vst1q_f16 + #define GGML_F16x8_FMA(a, b, c) vfmaq_f16(a, b, c) + #define GGML_F16x8_ADD vaddq_f16 + #define GGML_F16x8_MUL vmulq_f16 + #define GGML_F16x8_REDUCE(res, x) \ + do { \ + int offset = GGML_F16_ARR >> 1; \ + for (int i = 0; i < offset; ++i) { \ + (x)[i] = vaddq_f16((x)[i], (x)[offset+i]); \ + } \ + offset >>= 1; \ + for (int i = 0; i < offset; ++i) { \ + (x)[i] = vaddq_f16((x)[i], (x)[offset+i]); \ + } \ + offset >>= 1; \ + for (int i = 0; i < offset; ++i) { \ + (x)[i] = vaddq_f16((x)[i], (x)[offset+i]); \ + } \ + const float32x4_t t0 = vcvt_f32_f16(vget_low_f16 ((x)[0])); \ + const float32x4_t t1 = vcvt_f32_f16(vget_high_f16((x)[0])); \ + (res) = (ggml_float) vaddvq_f32(vaddq_f32(t0, t1)); \ + } while (0) + + #define GGML_F16_VEC GGML_F16x8 + #define GGML_F16_VEC_ZERO GGML_F16x8_ZERO + #define GGML_F16_VEC_SET1 GGML_F16x8_SET1 + #define GGML_F16_VEC_LOAD(p, i) GGML_F16x8_LOAD(p) + #define GGML_F16_VEC_STORE(p, r, i) GGML_F16x8_STORE((__fp16 *)(p), (r)[i]) + #define GGML_F16_VEC_FMA GGML_F16x8_FMA + #define GGML_F16_VEC_ADD GGML_F16x8_ADD + #define GGML_F16_VEC_MUL GGML_F16x8_MUL + #define GGML_F16_VEC_REDUCE GGML_F16x8_REDUCE +#else + // if FP16 vector arithmetic is not supported, we use FP32 instead + // and take advantage of the vcvt_ functions to convert to/from FP16 + + #define GGML_F16_STEP 16 + #define GGML_F16_EPR 4 + + #define GGML_F32Cx4 float32x4_t + #define GGML_F32Cx4_ZERO vdupq_n_f32(0.0f) + #define GGML_F32Cx4_SET1(x) vdupq_n_f32(x) + #define GGML_F32Cx4_LOAD(x) vcvt_f32_f16(vld1_f16((const __fp16 *)(x))) + #define GGML_F32Cx4_STORE(x, y) vst1_f16(x, vcvt_f16_f32(y)) + #define GGML_F32Cx4_FMA(a, b, c) vfmaq_f32(a, b, c) + #define GGML_F32Cx4_ADD vaddq_f32 + #define GGML_F32Cx4_MUL vmulq_f32 + #define GGML_F32Cx4_REDUCE GGML_F32x4_REDUCE + + #define GGML_F16_VEC GGML_F32Cx4 + #define GGML_F16_VEC_ZERO GGML_F32Cx4_ZERO + #define GGML_F16_VEC_SET1 GGML_F32Cx4_SET1 + #define GGML_F16_VEC_LOAD(p, i) GGML_F32Cx4_LOAD(p) + #define GGML_F16_VEC_STORE(p, r, i) GGML_F32Cx4_STORE((__fp16 *)(p), r[i]) + #define GGML_F16_VEC_FMA GGML_F32Cx4_FMA + #define GGML_F16_VEC_ADD GGML_F32Cx4_ADD + #define GGML_F16_VEC_MUL GGML_F32Cx4_MUL + #define GGML_F16_VEC_REDUCE GGML_F32Cx4_REDUCE +#endif + +#elif defined(__AVX512F__) + +#define GGML_SIMD + +// F32 AVX512 + +#define GGML_F32_STEP 64 +#define GGML_F32_EPR 16 + +#define GGML_F32x16 __m512 +#define GGML_F32x16_ZERO _mm512_setzero_ps() +#define GGML_F32x16_SET1(x) _mm512_set1_ps(x) +#define GGML_F32x16_LOAD _mm512_loadu_ps +#define GGML_F32x16_STORE _mm512_storeu_ps +// _mm512_fmadd_ps is defined in AVX512F so no guard is required +#define GGML_F32x16_FMA(a, b, c) _mm512_fmadd_ps(b, c, a) +#define GGML_F32x16_ADD _mm512_add_ps +#define GGML_F32x16_MUL _mm512_mul_ps +#define GGML_F32x16_REDUCE(res, x) \ +do { \ + int offset = GGML_F32_ARR >> 1; \ + for (int i = 0; i < offset; ++i) { \ + x[i] = _mm512_add_ps(x[i], x[offset+i]); \ + } \ + offset >>= 1; \ + for (int i = 0; i < offset; ++i) { \ + x[i] = _mm512_add_ps(x[i], x[offset+i]); \ + } \ + offset >>= 1; \ + for (int i = 0; i < offset; ++i) { \ + x[i] = _mm512_add_ps(x[i], x[offset+i]); \ + } \ + res = (ggml_float) _mm512_reduce_add_ps(x[0]); \ +} while (0) + +// TODO: is this optimal ? + +#define GGML_F32_VEC GGML_F32x16 +#define GGML_F32_VEC_ZERO GGML_F32x16_ZERO +#define GGML_F32_VEC_SET1 GGML_F32x16_SET1 +#define GGML_F32_VEC_LOAD GGML_F32x16_LOAD +#define GGML_F32_VEC_STORE GGML_F32x16_STORE +#define GGML_F32_VEC_FMA GGML_F32x16_FMA +#define GGML_F32_VEC_ADD GGML_F32x16_ADD +#define GGML_F32_VEC_MUL GGML_F32x16_MUL +#define GGML_F32_VEC_REDUCE GGML_F32x16_REDUCE + +// F16 AVX512 + +#if defined(__AVX512FP16__) + +#define GGML_F16_STEP 128 +#define GGML_F16_EPR 32 + +#define GGML_F16x32 __m512h +#define GGML_F16x32_ZERO _mm512_setzero_ph() +#define GGML_F16x32_SET1(x) _mm512_set1_ph(__extension__(_Float16)(x)) +#define GGML_F16x32_LOAD(x) _mm512_loadu_ph(x) +#define GGML_F16x32_STORE(x, y) _mm512_storeu_ph(x, y) +#define GGML_F16x32_FMA(a, b, c) _mm512_fmadd_ph(b, c, a) +#define GGML_F16x32_ADD _mm512_add_ph +#define GGML_F16x32_MUL _mm512_mul_ph +#define GGML_F16x32_REDUCE(res, x) \ +do { \ + int offset = GGML_F16_ARR >> 1; \ + for (int i = 0; i < offset; ++i) { \ + x[i] = _mm512_add_ph(x[i], x[offset+i]); \ + } \ + offset >>= 1; \ + for (int i = 0; i < offset; ++i) { \ + x[i] = _mm512_add_ph(x[i], x[offset+i]); \ + } \ + offset >>= 1; \ + for (int i = 0; i < offset; ++i) { \ + x[i] = _mm512_add_ph(x[i], x[offset+i]); \ + } \ + res = (ggml_float) _mm512_reduce_add_ph(x[0]); \ +} while (0) + +#define GGML_F16_VEC GGML_F16x32 +#define GGML_F16_VEC_ZERO GGML_F16x32_ZERO +#define GGML_F16_VEC_SET1 GGML_F16x32_SET1 +#define GGML_F16_VEC_LOAD(p, i) GGML_F16x32_LOAD(p) +#define GGML_F16_VEC_STORE(p, r, i) GGML_F16x32_STORE(p, r[i]) +#define GGML_F16_VEC_FMA GGML_F16x32_FMA +#define GGML_F16_VEC_ADD GGML_F16x32_ADD +#define GGML_F16_VEC_MUL GGML_F16x32_MUL +#define GGML_F16_VEC_REDUCE GGML_F16x32_REDUCE + +#else // Fallback FP16 <-> FP32 + +#define GGML_F16_STEP 64 +#define GGML_F16_EPR 16 + +#define GGML_F32Cx16 __m512 +#define GGML_F32Cx16_ZERO _mm512_setzero_ps() +#define GGML_F32Cx16_SET1(x) _mm512_set1_ps(x) + +// unlike _mm256_cvt intrinsics that require F16C, _mm512_cvt is defined in AVX512F +// so F16C guard isn't required +#define GGML_F32Cx16_LOAD(x) _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)(x))) +#define GGML_F32Cx16_STORE(x, y) _mm256_storeu_si256((__m256i *)(x), _mm512_cvtps_ph(y, 0)) + +#define GGML_F32Cx16_FMA(a, b, c) _mm512_fmadd_ps(b, c, a) +#define GGML_F32Cx16_ADD _mm512_add_ps +#define GGML_F32Cx16_MUL _mm512_mul_ps +#define GGML_F32Cx16_REDUCE(res, x) \ +do { \ + int offset = GGML_F32_ARR >> 1; \ + for (int i = 0; i < offset; ++i) { \ + x[i] = _mm512_add_ps(x[i], x[offset+i]); \ + } \ + offset >>= 1; \ + for (int i = 0; i < offset; ++i) { \ + x[i] = _mm512_add_ps(x[i], x[offset+i]); \ + } \ + offset >>= 1; \ + for (int i = 0; i < offset; ++i) { \ + x[i] = _mm512_add_ps(x[i], x[offset+i]); \ + } \ + res = (ggml_float) _mm512_reduce_add_ps(x[0]); \ +} while (0) + +#define GGML_F16_VEC GGML_F32Cx16 +#define GGML_F16_VEC_ZERO GGML_F32Cx16_ZERO +#define GGML_F16_VEC_SET1 GGML_F32Cx16_SET1 +#define GGML_F16_VEC_LOAD(p, i) GGML_F32Cx16_LOAD(p) +#define GGML_F16_VEC_STORE(p, r, i) GGML_F32Cx16_STORE(p, r[i]) +#define GGML_F16_VEC_FMA GGML_F32Cx16_FMA +#define GGML_F16_VEC_ADD GGML_F32Cx16_ADD +#define GGML_F16_VEC_MUL GGML_F32Cx16_MUL + +#define GGML_F16_VEC_REDUCE GGML_F32Cx16_REDUCE + +#endif // __AVX512FP16__ +#elif defined(__AVX__) + +#define GGML_SIMD + +// F32 AVX + +#define GGML_F32_STEP 32 +#define GGML_F32_EPR 8 + +#define GGML_F32x8 __m256 +#define GGML_F32x8_ZERO _mm256_setzero_ps() +#define GGML_F32x8_SET1(x) _mm256_set1_ps(x) +#define GGML_F32x8_LOAD _mm256_loadu_ps +#define GGML_F32x8_STORE _mm256_storeu_ps +#if defined(__FMA__) + #define GGML_F32x8_FMA(a, b, c) _mm256_fmadd_ps(b, c, a) +#else + #define GGML_F32x8_FMA(a, b, c) _mm256_add_ps(_mm256_mul_ps(b, c), a) +#endif +#define GGML_F32x8_ADD _mm256_add_ps +#define GGML_F32x8_MUL _mm256_mul_ps +#define GGML_F32x8_REDUCE(res, x) \ +do { \ + int offset = GGML_F32_ARR >> 1; \ + for (int i = 0; i < offset; ++i) { \ + x[i] = _mm256_add_ps(x[i], x[offset+i]); \ + } \ + offset >>= 1; \ + for (int i = 0; i < offset; ++i) { \ + x[i] = _mm256_add_ps(x[i], x[offset+i]); \ + } \ + offset >>= 1; \ + for (int i = 0; i < offset; ++i) { \ + x[i] = _mm256_add_ps(x[i], x[offset+i]); \ + } \ + const __m128 t0 = _mm_add_ps(_mm256_castps256_ps128(x[0]), \ + _mm256_extractf128_ps(x[0], 1)); \ + const __m128 t1 = _mm_hadd_ps(t0, t0); \ + res = (ggml_float) _mm_cvtss_f32(_mm_hadd_ps(t1, t1)); \ +} while (0) +// TODO: is this optimal ? + +#define GGML_F32_VEC GGML_F32x8 +#define GGML_F32_VEC_ZERO GGML_F32x8_ZERO +#define GGML_F32_VEC_SET1 GGML_F32x8_SET1 +#define GGML_F32_VEC_LOAD GGML_F32x8_LOAD +#define GGML_F32_VEC_STORE GGML_F32x8_STORE +#define GGML_F32_VEC_FMA GGML_F32x8_FMA +#define GGML_F32_VEC_ADD GGML_F32x8_ADD +#define GGML_F32_VEC_MUL GGML_F32x8_MUL +#define GGML_F32_VEC_REDUCE GGML_F32x8_REDUCE + +// F16 AVX + +#define GGML_F16_STEP 32 +#define GGML_F16_EPR 8 + +// F16 arithmetic is not supported by AVX, so we use F32 instead + +#define GGML_F32Cx8 __m256 +#define GGML_F32Cx8_ZERO _mm256_setzero_ps() +#define GGML_F32Cx8_SET1(x) _mm256_set1_ps(x) + +#if defined(__F16C__) +// the _mm256_cvt intrinsics require F16C +#define GGML_F32Cx8_LOAD(x) _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)(x))) +#define GGML_F32Cx8_STORE(x, y) _mm_storeu_si128((__m128i *)(x), _mm256_cvtps_ph(y, 0)) +#else +static inline __m256 __avx_f32cx8_load(const ggml_fp16_t * x) { + float tmp[8]; + + for (int i = 0; i < 8; i++) { + tmp[i] = GGML_CPU_FP16_TO_FP32(x[i]); + } + + return _mm256_loadu_ps(tmp); +} +static inline void __avx_f32cx8_store(ggml_fp16_t *x, __m256 y) { + float arr[8]; + + _mm256_storeu_ps(arr, y); + + for (int i = 0; i < 8; i++) + x[i] = GGML_CPU_FP32_TO_FP16(arr[i]); +} +#define GGML_F32Cx8_LOAD(x) __avx_f32cx8_load(x) +#define GGML_F32Cx8_STORE(x, y) __avx_f32cx8_store(x, y) +#endif + +#define GGML_F32Cx8_FMA GGML_F32x8_FMA +#define GGML_F32Cx8_ADD _mm256_add_ps +#define GGML_F32Cx8_MUL _mm256_mul_ps +#define GGML_F32Cx8_REDUCE GGML_F32x8_REDUCE + +#define GGML_F16_VEC GGML_F32Cx8 +#define GGML_F16_VEC_ZERO GGML_F32Cx8_ZERO +#define GGML_F16_VEC_SET1 GGML_F32Cx8_SET1 +#define GGML_F16_VEC_LOAD(p, i) GGML_F32Cx8_LOAD(p) +#define GGML_F16_VEC_STORE(p, r, i) GGML_F32Cx8_STORE(p, r[i]) +#define GGML_F16_VEC_FMA GGML_F32Cx8_FMA +#define GGML_F16_VEC_ADD GGML_F32Cx8_ADD +#define GGML_F16_VEC_MUL GGML_F32Cx8_MUL +#define GGML_F16_VEC_REDUCE GGML_F32Cx8_REDUCE + +#elif defined(__POWER9_VECTOR__) + +#define GGML_SIMD + +// F32 POWER9 + +#define GGML_F32_STEP 32 +#define GGML_F32_EPR 4 + +#define GGML_F32x4 vector float +#define GGML_F32x4_ZERO {0.0f} +#define GGML_F32x4_SET1 vec_splats +#define GGML_F32x4_LOAD(p) vec_xl(0, p) +#define GGML_F32x4_STORE(p, r) vec_xst(r, 0, p) +#define GGML_F32x4_FMA(a, b, c) vec_madd(b, c, a) +#define GGML_F32x4_ADD vec_add +#define GGML_F32x4_MUL vec_mul +#define GGML_F32x4_REDUCE(res, x) \ +{ \ + int offset = GGML_F32_ARR >> 1; \ + for (int i = 0; i < offset; ++i) { \ + x[i] = vec_add(x[i], x[offset+i]); \ + } \ + offset >>= 1; \ + for (int i = 0; i < offset; ++i) { \ + x[i] = vec_add(x[i], x[offset+i]); \ + } \ + offset >>= 1; \ + for (int i = 0; i < offset; ++i) { \ + x[i] = vec_add(x[i], x[offset+i]); \ + } \ + res = vec_extract(x[0], 0) + \ + vec_extract(x[0], 1) + \ + vec_extract(x[0], 2) + \ + vec_extract(x[0], 3); \ +} +#define GGML_F32x4_REDUCE_4(res, s0, s1, s2, s3) \ +{ \ + vector float v = vec_add(vec_add(s0, s1), \ + vec_add(s2, s3)); \ + v = vec_add(v, vec_sld(v, v, 8)); \ + v = vec_add(v, vec_sld(v, v, 4)); \ + res += (ggml_float) vec_extract(v, 0); \ +} + +#define GGML_F32_VEC GGML_F32x4 +#define GGML_F32_VEC_ZERO GGML_F32x4_ZERO +#define GGML_F32_VEC_SET1 GGML_F32x4_SET1 +#define GGML_F32_VEC_LOAD GGML_F32x4_LOAD +#define GGML_F32_VEC_STORE GGML_F32x4_STORE +#define GGML_F32_VEC_FMA GGML_F32x4_FMA +#define GGML_F32_VEC_ADD GGML_F32x4_ADD +#define GGML_F32_VEC_MUL GGML_F32x4_MUL +#define GGML_F32_VEC_REDUCE GGML_F32x4_REDUCE + +// F16 POWER9 +#define GGML_F16_STEP GGML_F32_STEP +#define GGML_F16_EPR GGML_F32_EPR +#define GGML_F16_VEC GGML_F32x4 +#define GGML_F16_VEC_ZERO GGML_F32x4_ZERO +#define GGML_F16_VEC_SET1 GGML_F32x4_SET1 +#define GGML_F16_VEC_FMA GGML_F32x4_FMA +#define GGML_F16_VEC_ADD GGML_F32x4_ADD +#define GGML_F16_VEC_MUL GGML_F32x4_MUL +#define GGML_F16_VEC_REDUCE GGML_F32x4_REDUCE +// Use vec_xl, not vec_ld, in case the load address is not aligned. +#define GGML_F16_VEC_LOAD(p, i) (i & 0x1) ? \ + vec_extract_fp32_from_shorth(vec_xl(0, p - GGML_F16_EPR)) : \ + vec_extract_fp32_from_shortl(vec_xl(0, p)) +static inline unsigned char ggml_endian_byte(int i) { + uint16_t tmp_val = 1; + return ((unsigned char *)&tmp_val)[i]; +} +#define GGML_ENDIAN_BYTE(i) ggml_endian_byte(i) +#define GGML_F16_VEC_STORE(p, r, i) \ + if (i & 0x1) \ + vec_xst(vec_pack_to_short_fp32(r[i - GGML_ENDIAN_BYTE(1)], \ + r[i - GGML_ENDIAN_BYTE(0)]), \ + 0, p - GGML_F16_EPR) + +//BF16 POWER9 +#define GGML_BF16_STEP 16 +#define GGML_BF16_EPR 8 + +#define GGML_BF16x8 vector unsigned short +#define GGML_BF16x8_ZERO vec_splats((unsigned short)0) +#define GGML_BF16x8_LOAD(p) vec_xl(0, (const unsigned short *)(p)) + +#define GGML_BF16_VEC GGML_BF16x8 +#define GGML_BF16_VEC_ZERO GGML_BF16x8_ZERO +#define GGML_BF16_VEC_LOAD GGML_BF16x8_LOAD +#if defined(__LITTLE_ENDIAN__) +#define GGML_BF16_TO_F32_LO(v) ((vector float) vec_mergel(GGML_BF16_VEC_ZERO, (v))) +#define GGML_BF16_TO_F32_HI(v) ((vector float) vec_mergeh(GGML_BF16_VEC_ZERO, (v))) +#else +#define GGML_BF16_TO_F32_LO(v) ((vector float) vec_mergel((v), GGML_BF16_VEC_ZERO)) +#define GGML_BF16_TO_F32_HI(v) ((vector float) vec_mergeh((v), GGML_BF16_VEC_ZERO)) +#endif +#define GGML_BF16_FMA_LO(acc, x, y) \ + (acc) = GGML_F32x4_FMA((acc), GGML_BF16_TO_F32_LO(x), GGML_BF16_TO_F32_LO(y)) +#define GGML_BF16_FMA_HI(acc, x, y) \ + (acc) = GGML_F32x4_FMA((acc), GGML_BF16_TO_F32_HI(x), GGML_BF16_TO_F32_HI(y)) + +#elif defined(__wasm_simd128__) + +#define GGML_SIMD + +// F32 WASM + +#define GGML_F32_STEP 16 +#define GGML_F32_EPR 4 + +#define GGML_F32x4 v128_t +#define GGML_F32x4_ZERO wasm_f32x4_splat(0.0f) +#define GGML_F32x4_SET1(x) wasm_f32x4_splat(x) +#define GGML_F32x4_LOAD wasm_v128_load +#define GGML_F32x4_STORE wasm_v128_store +#define GGML_F32x4_FMA(a, b, c) wasm_f32x4_add(wasm_f32x4_mul(b, c), a) +#define GGML_F32x4_ADD wasm_f32x4_add +#define GGML_F32x4_MUL wasm_f32x4_mul +#define GGML_F32x4_REDUCE(res, x) \ +{ \ + int offset = GGML_F32_ARR >> 1; \ + for (int i = 0; i < offset; ++i) { \ + x[i] = wasm_f32x4_add(x[i], x[offset+i]); \ + } \ + offset >>= 1; \ + for (int i = 0; i < offset; ++i) { \ + x[i] = wasm_f32x4_add(x[i], x[offset+i]); \ + } \ + offset >>= 1; \ + for (int i = 0; i < offset; ++i) { \ + x[i] = wasm_f32x4_add(x[i], x[offset+i]); \ + } \ + res = wasm_f32x4_extract_lane(x[0], 0) + \ + wasm_f32x4_extract_lane(x[0], 1) + \ + wasm_f32x4_extract_lane(x[0], 2) + \ + wasm_f32x4_extract_lane(x[0], 3); \ +} + +#define GGML_F32_VEC GGML_F32x4 +#define GGML_F32_VEC_ZERO GGML_F32x4_ZERO +#define GGML_F32_VEC_SET1 GGML_F32x4_SET1 +#define GGML_F32_VEC_LOAD GGML_F32x4_LOAD +#define GGML_F32_VEC_STORE GGML_F32x4_STORE +#define GGML_F32_VEC_FMA GGML_F32x4_FMA +#define GGML_F32_VEC_ADD GGML_F32x4_ADD +#define GGML_F32_VEC_MUL GGML_F32x4_MUL +#define GGML_F32_VEC_REDUCE GGML_F32x4_REDUCE + +// F16 WASM + +#define GGML_F16_STEP 16 +#define GGML_F16_EPR 4 + +inline static v128_t __wasm_f16x4_load(const ggml_fp16_t * p) { + float tmp[4]; + + tmp[0] = GGML_CPU_FP16_TO_FP32(p[0]); + tmp[1] = GGML_CPU_FP16_TO_FP32(p[1]); + tmp[2] = GGML_CPU_FP16_TO_FP32(p[2]); + tmp[3] = GGML_CPU_FP16_TO_FP32(p[3]); + + return wasm_v128_load(tmp); +} + +inline static void __wasm_f16x4_store(ggml_fp16_t * p, v128_t x) { + float tmp[4]; + + wasm_v128_store(tmp, x); + + p[0] = GGML_CPU_FP32_TO_FP16(tmp[0]); + p[1] = GGML_CPU_FP32_TO_FP16(tmp[1]); + p[2] = GGML_CPU_FP32_TO_FP16(tmp[2]); + p[3] = GGML_CPU_FP32_TO_FP16(tmp[3]); +} + +#define GGML_F16x4 v128_t +#define GGML_F16x4_ZERO wasm_f32x4_splat(0.0f) +#define GGML_F16x4_SET1(x) wasm_f32x4_splat(x) +#define GGML_F16x4_LOAD(x) __wasm_f16x4_load(x) +#define GGML_F16x4_STORE(x, y) __wasm_f16x4_store(x, y) +#define GGML_F16x4_FMA GGML_F32x4_FMA +#define GGML_F16x4_ADD wasm_f32x4_add +#define GGML_F16x4_MUL wasm_f32x4_mul +#define GGML_F16x4_REDUCE(res, x) \ +{ \ + int offset = GGML_F16_ARR >> 1; \ + for (int i = 0; i < offset; ++i) { \ + x[i] = wasm_f32x4_add(x[i], x[offset+i]); \ + } \ + offset >>= 1; \ + for (int i = 0; i < offset; ++i) { \ + x[i] = wasm_f32x4_add(x[i], x[offset+i]); \ + } \ + offset >>= 1; \ + for (int i = 0; i < offset; ++i) { \ + x[i] = wasm_f32x4_add(x[i], x[offset+i]); \ + } \ + res = (ggml_float) (wasm_f32x4_extract_lane(x[0], 0) + \ + wasm_f32x4_extract_lane(x[0], 1) + \ + wasm_f32x4_extract_lane(x[0], 2) + \ + wasm_f32x4_extract_lane(x[0], 3)); \ +} + +#define GGML_F16_VEC GGML_F16x4 +#define GGML_F16_VEC_ZERO GGML_F16x4_ZERO +#define GGML_F16_VEC_SET1 GGML_F16x4_SET1 +#define GGML_F16_VEC_LOAD(p, i) GGML_F16x4_LOAD(p) +#define GGML_F16_VEC_STORE(p, r, i) GGML_F16x4_STORE(p, r[i]) +#define GGML_F16_VEC_FMA GGML_F16x4_FMA +#define GGML_F16_VEC_ADD GGML_F16x4_ADD +#define GGML_F16_VEC_MUL GGML_F16x4_MUL +#define GGML_F16_VEC_REDUCE GGML_F16x4_REDUCE + +#elif defined(__SSE3__) + +#define GGML_SIMD + +// F32 SSE + +#define GGML_F32_STEP 32 +#define GGML_F32_EPR 4 + +#define GGML_F32x4 __m128 +#define GGML_F32x4_ZERO _mm_setzero_ps() +#define GGML_F32x4_SET1(x) _mm_set1_ps(x) +#define GGML_F32x4_LOAD _mm_loadu_ps +#define GGML_F32x4_STORE _mm_storeu_ps +#if defined(__FMA__) + // TODO: Does this work? + #define GGML_F32x4_FMA(a, b, c) _mm_fmadd_ps(b, c, a) +#else + #define GGML_F32x4_FMA(a, b, c) _mm_add_ps(_mm_mul_ps(b, c), a) +#endif +#define GGML_F32x4_ADD _mm_add_ps +#define GGML_F32x4_MUL _mm_mul_ps +#define GGML_F32x4_REDUCE(res, x) \ +{ \ + int offset = GGML_F32_ARR >> 1; \ + for (int i = 0; i < offset; ++i) { \ + x[i] = _mm_add_ps(x[i], x[offset+i]); \ + } \ + offset >>= 1; \ + for (int i = 0; i < offset; ++i) { \ + x[i] = _mm_add_ps(x[i], x[offset+i]); \ + } \ + offset >>= 1; \ + for (int i = 0; i < offset; ++i) { \ + x[i] = _mm_add_ps(x[i], x[offset+i]); \ + } \ + const __m128 t0 = _mm_hadd_ps(x[0], x[0]); \ + res = (ggml_float) _mm_cvtss_f32(_mm_hadd_ps(t0, t0)); \ +} +// TODO: is this optimal ? + +#define GGML_F32_VEC GGML_F32x4 +#define GGML_F32_VEC_ZERO GGML_F32x4_ZERO +#define GGML_F32_VEC_SET1 GGML_F32x4_SET1 +#define GGML_F32_VEC_LOAD GGML_F32x4_LOAD +#define GGML_F32_VEC_STORE GGML_F32x4_STORE +#define GGML_F32_VEC_FMA GGML_F32x4_FMA +#define GGML_F32_VEC_ADD GGML_F32x4_ADD +#define GGML_F32_VEC_MUL GGML_F32x4_MUL +#define GGML_F32_VEC_REDUCE GGML_F32x4_REDUCE + +// F16 SSE + +#define GGML_F16_STEP 32 +#define GGML_F16_EPR 4 + +static inline __m128 __sse_f16x4_load(const ggml_fp16_t * x) { + float tmp[4]; + + tmp[0] = GGML_CPU_FP16_TO_FP32(x[0]); + tmp[1] = GGML_CPU_FP16_TO_FP32(x[1]); + tmp[2] = GGML_CPU_FP16_TO_FP32(x[2]); + tmp[3] = GGML_CPU_FP16_TO_FP32(x[3]); + + return _mm_loadu_ps(tmp); +} + +static inline void __sse_f16x4_store(ggml_fp16_t * x, __m128 y) { + float arr[4]; + + _mm_storeu_ps(arr, y); + + x[0] = GGML_CPU_FP32_TO_FP16(arr[0]); + x[1] = GGML_CPU_FP32_TO_FP16(arr[1]); + x[2] = GGML_CPU_FP32_TO_FP16(arr[2]); + x[3] = GGML_CPU_FP32_TO_FP16(arr[3]); +} + +#define GGML_F32Cx4 __m128 +#define GGML_F32Cx4_ZERO _mm_setzero_ps() +#define GGML_F32Cx4_SET1(x) _mm_set1_ps(x) +#define GGML_F32Cx4_LOAD(x) __sse_f16x4_load(x) +#define GGML_F32Cx4_STORE(x, y) __sse_f16x4_store(x, y) +#define GGML_F32Cx4_FMA GGML_F32x4_FMA +#define GGML_F32Cx4_ADD _mm_add_ps +#define GGML_F32Cx4_MUL _mm_mul_ps +#define GGML_F32Cx4_REDUCE GGML_F32x4_REDUCE + +#define GGML_F16_VEC GGML_F32Cx4 +#define GGML_F16_VEC_ZERO GGML_F32Cx4_ZERO +#define GGML_F16_VEC_SET1 GGML_F32Cx4_SET1 +#define GGML_F16_VEC_LOAD(p, i) GGML_F32Cx4_LOAD(p) +#define GGML_F16_VEC_STORE(p, r, i) GGML_F32Cx4_STORE(p, r[i]) +#define GGML_F16_VEC_FMA GGML_F32Cx4_FMA +#define GGML_F16_VEC_ADD GGML_F32Cx4_ADD +#define GGML_F16_VEC_MUL GGML_F32Cx4_MUL +#define GGML_F16_VEC_REDUCE GGML_F32Cx4_REDUCE + +#elif defined(__loongarch_asx) + +#define GGML_SIMD + +// F32 LASX +#define GGML_F32_STEP 32 +#define GGML_F32_EPR 8 + +#define GGML_F32x8 __m256 +#define GGML_F32x8_ZERO (__m256)__lasx_xvldi(0) +#define GGML_F32x8_SET1(x) (__m256)__lasx_xvreplfr2vr_s((x)) +#define GGML_F32x8_LOAD(x) (__m256)__lasx_xvld((x), 0) +#define GGML_F32x8_STORE(x,y) __lasx_xvst((y), (x), 0) +#define GGML_F32x8_FMA(a, b, c) __lasx_xvfmadd_s(b, c, a) +#define GGML_F32x8_ADD __lasx_xvfadd_s +#define GGML_F32x8_MUL __lasx_xvfmul_s +#define GGML_F32x8_REDUCE(res, x) \ +do { \ + int offset = GGML_F32_ARR >> 1; \ + for (int i = 0; i < offset; ++i) { \ + x[i] = __lasx_xvfadd_s(x[i], x[offset+i]); \ + } \ + offset >>= 1; \ + for (int i = 0; i < offset; ++i) { \ + x[i] = __lasx_xvfadd_s(x[i], x[offset+i]); \ + } \ + offset >>= 1; \ + for (int i = 0; i < offset; ++i) { \ + x[i] = __lasx_xvfadd_s(x[i], x[offset+i]); \ + } \ + float *tmp_p = (float *)&x[0]; \ + res = tmp_p[0] + tmp_p[1] + tmp_p[2] + tmp_p[3] + tmp_p[4] + tmp_p[5] + tmp_p[6] + tmp_p[7]; \ +} while (0) +// TODO: is this optimal ? + +#define GGML_F32_VEC GGML_F32x8 +#define GGML_F32_VEC_ZERO GGML_F32x8_ZERO +#define GGML_F32_VEC_SET1 GGML_F32x8_SET1 +#define GGML_F32_VEC_LOAD GGML_F32x8_LOAD +#define GGML_F32_VEC_STORE GGML_F32x8_STORE +#define GGML_F32_VEC_FMA GGML_F32x8_FMA +#define GGML_F32_VEC_ADD GGML_F32x8_ADD +#define GGML_F32_VEC_MUL GGML_F32x8_MUL +#define GGML_F32_VEC_REDUCE GGML_F32x8_REDUCE + +// F16 LASX + +#define GGML_F16_STEP 32 +#define GGML_F16_EPR 8 + +// F16 arithmetic is not supported by LASX, so we use F32 instead + +#define GGML_F32Cx8 __m256 +#define GGML_F32Cx8_ZERO (__m256)__lasx_xvldi(0) +#define GGML_F32Cx8_SET1(x) (__m256)__lasx_xvreplfr2vr_s((x)) + +static inline __m256 __lasx_f32cx8_load(const ggml_fp16_t * x) { + __m256i a; + memcpy(&a, x, sizeof(ggml_fp16_t) * 8); + a = __lasx_xvpermi_d(a, 0 | (1 << 4)); + return __lasx_xvfcvtl_s_h(a); +} + +static inline void __lasx_f32cx8_store(ggml_fp16_t * x, __m256 y) { + __m256i a = __lasx_xvfcvt_h_s(y, y); + a = __lasx_xvpermi_d(a, 0 | (2 << 2)); + memcpy(x, &a, sizeof(ggml_fp16_t) * 8); +} +#define GGML_F32Cx8_LOAD(x) __lasx_f32cx8_load(x) +#define GGML_F32Cx8_STORE(x, y) __lasx_f32cx8_store(x, y) + +#define GGML_F32Cx8_FMA GGML_F32x8_FMA +#define GGML_F32Cx8_ADD __lasx_xvfadd_s +#define GGML_F32Cx8_MUL __lasx_xvfmul_s +#define GGML_F32Cx8_REDUCE GGML_F32x8_REDUCE + +#define GGML_F16_VEC GGML_F32Cx8 +#define GGML_F16_VEC_ZERO GGML_F32Cx8_ZERO +#define GGML_F16_VEC_SET1 GGML_F32Cx8_SET1 +#define GGML_F16_VEC_LOAD(p, i) GGML_F32Cx8_LOAD(p) +#define GGML_F16_VEC_STORE(p, r, i) GGML_F32Cx8_STORE(p, r[i]) +#define GGML_F16_VEC_FMA GGML_F32Cx8_FMA +#define GGML_F16_VEC_ADD GGML_F32Cx8_ADD +#define GGML_F16_VEC_MUL GGML_F32Cx8_MUL +#define GGML_F16_VEC_REDUCE GGML_F32Cx8_REDUCE + +#elif defined(__loongarch_sx) + +#define GGML_SIMD + +// F32 LSX + +#define GGML_F32_STEP 32 +#define GGML_F32_EPR 4 + +#define GGML_F32x4 __m128 +#define GGML_F32x4_ZERO (__m128)__lsx_vldi(0) +#define GGML_F32x4_SET1(x) (__m128)__lsx_vreplfr2vr_s((x)) +#define GGML_F32x4_LOAD(x) (__m128)__lsx_vld((x), 0) +#define GGML_F32x4_STORE(x, y) __lsx_vst(y, x, 0) +#define GGML_F32x4_FMA(a, b, c) __lsx_vfmadd_s(b, c, a) +#define GGML_F32x4_ADD __lsx_vfadd_s +#define GGML_F32x4_MUL __lsx_vfmul_s + +#define GGML_F32x4_REDUCE(res, x) \ +{ \ + int offset = GGML_F32_ARR >> 1; \ + for (int i = 0; i < offset; ++i) { \ + x[i] = __lsx_vfadd_s(x[i], x[offset+i]); \ + } \ + offset >>= 1; \ + for (int i = 0; i < offset; ++i) { \ + x[i] = __lsx_vfadd_s(x[i], x[offset+i]); \ + } \ + offset >>= 1; \ + for (int i = 0; i < offset; ++i) { \ + x[i] = __lsx_vfadd_s(x[i], x[offset+i]); \ + } \ + __m128i t0 = __lsx_vpickev_w((__m128i)x[0], (__m128i)x[0]); \ + __m128i t1 = __lsx_vpickod_w((__m128i)x[0], (__m128i)x[0]); \ + __m128 t2 = __lsx_vfadd_s((__m128)t0, (__m128)t1); \ + __m128i t3 = __lsx_vpickev_w((__m128i)t2, (__m128i)t2); \ + __m128i t4 = __lsx_vpickod_w((__m128i)t2, (__m128i)t2); \ + __m128 t5 = __lsx_vfadd_s((__m128)t3, (__m128)t4); \ + res = (ggml_float) ((v4f32)t5)[0]; \ +} + +#define GGML_F32_VEC GGML_F32x4 +#define GGML_F32_VEC_ZERO GGML_F32x4_ZERO +#define GGML_F32_VEC_SET1 GGML_F32x4_SET1 +#define GGML_F32_VEC_LOAD GGML_F32x4_LOAD +#define GGML_F32_VEC_STORE GGML_F32x4_STORE +#define GGML_F32_VEC_FMA GGML_F32x4_FMA +#define GGML_F32_VEC_ADD GGML_F32x4_ADD +#define GGML_F32_VEC_MUL GGML_F32x4_MUL +#define GGML_F32_VEC_REDUCE GGML_F32x4_REDUCE + +// F16 LSX + +#define GGML_F16_STEP 32 +#define GGML_F16_EPR 4 + +static inline __m128 __lsx_f16x4_load(const ggml_fp16_t * x) { + return __lsx_vfcvtl_s_h(__lsx_vld((const void *)x, 0)); +} + +static inline void __lsx_f16x4_store(ggml_fp16_t * x, __m128 y) { + __m128i a = __lsx_vfcvt_h_s(y, y); + memcpy(x, &a, sizeof(ggml_fp16_t) * 4); +} + +#define GGML_F32Cx4 __m128 +#define GGML_F32Cx4_ZERO (__m128)__lsx_vldi(0) +#define GGML_F32Cx4_SET1(x) (__m128)__lsx_vreplfr2vr_s((x)) +#define GGML_F32Cx4_LOAD(x) (__m128)__lsx_f16x4_load(x) +#define GGML_F32Cx4_STORE(x, y) __lsx_f16x4_store(x, y) +#define GGML_F32Cx4_FMA GGML_F32x4_FMA +#define GGML_F32Cx4_ADD __lsx_vfadd_s +#define GGML_F32Cx4_MUL __lsx_vfmul_s +#define GGML_F32Cx4_REDUCE GGML_F32x4_REDUCE + +#define GGML_F16_VEC GGML_F32Cx4 +#define GGML_F16_VEC_ZERO GGML_F32Cx4_ZERO +#define GGML_F16_VEC_SET1 GGML_F32Cx4_SET1 +#define GGML_F16_VEC_LOAD(p, i) GGML_F32Cx4_LOAD(p) +#define GGML_F16_VEC_STORE(p, r, i) GGML_F32Cx4_STORE(p, r[i]) +#define GGML_F16_VEC_FMA GGML_F32Cx4_FMA +#define GGML_F16_VEC_ADD GGML_F32Cx4_ADD +#define GGML_F16_VEC_MUL GGML_F32Cx4_MUL +#define GGML_F16_VEC_REDUCE GGML_F32Cx4_REDUCE + +#elif defined(__VXE__) || defined(__VXE2__) + +#define GGML_SIMD + +// F32 s390x + +#define GGML_F32_STEP 32 +#define GGML_F32_EPR 4 + +#define GGML_F32x4 float32x4_t +#define GGML_F32x4_ZERO vec_splats(0.0f) +#define GGML_F32x4_SET1 vec_splats +#define GGML_F32x4_LOAD(p) vec_xl(0, p) +#define GGML_F32x4_STORE(p, r) vec_xst(r, 0, p) +#define GGML_F32x4_FMA(a, b, c) vec_madd(b, c, a) +#define GGML_F32x4_ADD vec_add +#define GGML_F32x4_MUL vec_mul +#define GGML_F32x4_REDUCE(res, x) \ +{ \ + int offset = GGML_F32_ARR >> 1; \ + for (int i = 0; i < offset; ++i) { \ + x[i] = vec_add(x[i], x[offset + i]); \ + } \ + offset >>= 1; \ + for (int i = 0; i < offset; ++i) { \ + x[i] = vec_add(x[i], x[offset + i]); \ + } \ + offset >>= 1; \ + for (int i = 0; i < offset; ++i) { \ + x[i] = vec_add(x[i], x[offset + i]); \ + } \ + float32x4_t tmp = x[0] + vec_reve(x[0]); \ + res = tmp[0] + tmp[1]; \ +} +#define GGML_F32x4_REDUCE_4(res, s0, s1, s2, s3) \ +{ \ + float32x4_t v = vec_add(vec_add(s0, s1), \ + vec_add(s2, s3)); \ + v = vec_add(v, vec_sld(v, v, 8)); \ + v = vec_add(v, vec_sld(v, v, 4)); \ + res += (ggml_float)vec_extract(v, 0); \ +} + +#define GGML_F32_VEC GGML_F32x4 +#define GGML_F32_VEC_ZERO GGML_F32x4_ZERO +#define GGML_F32_VEC_SET1 GGML_F32x4_SET1 +#define GGML_F32_VEC_LOAD GGML_F32x4_LOAD +#define GGML_F32_VEC_STORE GGML_F32x4_STORE +#define GGML_F32_VEC_FMA GGML_F32x4_FMA +#define GGML_F32_VEC_ADD GGML_F32x4_ADD +#define GGML_F32_VEC_MUL GGML_F32x4_MUL +#define GGML_F32_VEC_REDUCE GGML_F32x4_REDUCE + +// F16 s390x +#define GGML_F16_STEP GGML_F32_STEP +#define GGML_F16_EPR GGML_F32_EPR + +static inline uint32x4_t __lzs_f32cx4_to_f16(float32x4_t v_f) { + float32x4_t v_base = vec_mul(vec_mul(vec_abs(v_f), vec_splats(0x1.0p+112f)), vec_splats(0x1.0p-110f)); + + const uint32x4_t v_w = (uint32x4_t)v_f; + const uint32x4_t v_shl1_w = vec_add(v_w, v_w); + const uint32x4_t v_sign = vec_and(v_w, vec_splats(UINT32_C(0x80000000))); + const uint32x4_t v_bias = vec_max(vec_and(v_shl1_w, vec_splats(UINT32_C(0xFF000000))), vec_splats(UINT32_C(0x71000000))); + + v_base = vec_add((float32x4_t)vec_add(vec_sr(v_bias, 1), vec_splats(UINT32_C(0x07800000))), v_base); + + const uint32x4_t v_bits = (uint32x4_t)v_base; + const uint32x4_t v_nonsign = vec_add(vec_and(vec_sr(v_bits, 13), vec_splats(UINT32_C(0x00007C00))), + vec_and(v_bits, vec_splats(UINT32_C(0x00000FFF)))); + const uint32x4_t v_is_nan = (uint32x4_t)vec_cmpgt(v_shl1_w, vec_splats(UINT32_C(0xFF000000))); + + return vec_or(vec_sr(v_sign, 16), vec_sel(v_nonsign, vec_splats(UINT32_C(0x7E00)), v_is_nan)); +} + +static inline float32x4_t __lzs_f16cx4_load(const ggml_fp16_t * x) { + float tmp[4]; + + for (int i = 0; i < 4; i++) { + tmp[i] = GGML_CPU_FP16_TO_FP32(x[i]); + } + + // note: keep type-cast here to prevent compiler bugs + // see: https://github.com/ggml-org/llama.cpp/issues/12846 + return vec_xl(0, (const float *)(tmp)); +} + +static inline void __lzs_f16cx4_store(ggml_fp16_t * x, float32x4_t v_y) { + const uint32x4_t v_h = __lzs_f32cx4_to_f16(v_y); + const uint64_t tmp = ((uint64x2_t)vec_pack(v_h, v_h))[0]; + memcpy(x, &tmp, sizeof(tmp)); +} + +#define GGML_F16_VEC GGML_F32x4 +#define GGML_F16_VEC_ZERO GGML_F32x4_ZERO +#define GGML_F16_VEC_SET1 GGML_F32x4_SET1 +#define GGML_F16_VEC_LOAD(p, i) __lzs_f16cx4_load(p) +#define GGML_F16_VEC_STORE(p, r, i) __lzs_f16cx4_store(p, r[i]) +#define GGML_F16_VEC_FMA GGML_F32x4_FMA +#define GGML_F16_VEC_ADD GGML_F32x4_ADD +#define GGML_F16_VEC_MUL GGML_F32x4_MUL +#define GGML_F16_VEC_REDUCE GGML_F32x4_REDUCE + +// BF16 s390x +#define GGML_BF16_STEP 16 +#define GGML_BF16_EPR 8 + +#define GGML_BF16x8 __vector unsigned short +#define GGML_BF16x8_ZERO vec_splats((unsigned short)0) +#define GGML_BF16x8_LOAD(p) vec_xl(0, (const unsigned short *)(p)) + +#define GGML_BF16_VEC GGML_BF16x8 +#define GGML_BF16_VEC_ZERO GGML_BF16x8_ZERO +#define GGML_BF16_VEC_LOAD GGML_BF16x8_LOAD +#define GGML_BF16_TO_F32_LO(v) ((float32x4_t) vec_mergel((v), GGML_BF16_VEC_ZERO)) +#define GGML_BF16_TO_F32_HI(v) ((float32x4_t) vec_mergeh((v), GGML_BF16_VEC_ZERO)) +#define GGML_BF16_FMA_LO(acc, x, y) \ + (acc) = GGML_F32x4_FMA((acc), GGML_BF16_TO_F32_LO(x), GGML_BF16_TO_F32_LO(y)) +#define GGML_BF16_FMA_HI(acc, x, y) \ + (acc) = GGML_F32x4_FMA((acc), GGML_BF16_TO_F32_HI(x), GGML_BF16_TO_F32_HI(y)) + +#elif defined(__riscv_v_intrinsic) + +// compatible with vlen >= 128 + +#define GGML_SIMD + +// F32 + +#define GGML_F32_STEP 16 +#define GGML_F32_EPR 4 + +#define GGML_F32x4 vfloat32m1_t +#define GGML_F32x4_ZERO __riscv_vfmv_v_f_f32m1(0.0f, GGML_F32_EPR) +#define GGML_F32x4_SET1(x) __riscv_vfmv_v_f_f32m1(x, GGML_F32_EPR) +#define GGML_F32x4_LOAD(x) __riscv_vle32_v_f32m1(x, GGML_F32_EPR) +#define GGML_F32x4_STORE(b, v) __riscv_vse32_v_f32m1(b, v, GGML_F32_EPR) +#define GGML_F32x4_FMA(a, b, c) __riscv_vfmacc_vv_f32m1(a, b, c, GGML_F32_EPR) +#define GGML_F32x4_ADD(a, b) __riscv_vfadd_vv_f32m1(a, b, GGML_F32_EPR) +#define GGML_F32x4_MUL(a, b) __riscv_vfmul_vv_f32m1(a, b, GGML_F32_EPR) + +#define GGML_F32_VEC GGML_F32x4 +#define GGML_F32_VEC_ZERO GGML_F32x4_ZERO +#define GGML_F32_VEC_SET1 GGML_F32x4_SET1 +#define GGML_F32_VEC_LOAD GGML_F32x4_LOAD +#define GGML_F32_VEC_STORE GGML_F32x4_STORE +#define GGML_F32_VEC_FMA GGML_F32x4_FMA +#define GGML_F32_VEC_ADD GGML_F32x4_ADD +#define GGML_F32_VEC_MUL GGML_F32x4_MUL +#define GGML_F32_VEC_REDUCE GGML_F32x4_REDUCE + +#endif + +// GGML_F32_ARR / GGML_F16_ARR +// number of registers to use per step +#ifdef GGML_SIMD +#define GGML_F32_ARR (GGML_F32_STEP/GGML_F32_EPR) +#define GGML_F16_ARR (GGML_F16_STEP/GGML_F16_EPR) +#endif + +// GGML_F16_DOT_* +// like GGML_F16_* but for dot products which need F32 accumulation on AVX512-FP16 + +#if defined(__AVX512FP16__) + +#define GGML_F16_DOT_STEP GGML_F32_STEP +#define GGML_F16_DOT_EPR GGML_F32_EPR +#define GGML_F16_DOT_ARR GGML_F32_ARR +#define GGML_F16_DOT_VEC GGML_F32x16 +#define GGML_F16_DOT_VEC_ZERO GGML_F32x16_ZERO +#define GGML_F16_DOT_VEC_LOAD(p, i) _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)(p))) +#define GGML_F16_DOT_VEC_FMA GGML_F32x16_FMA +#define GGML_F16_DOT_VEC_REDUCE GGML_F32x16_REDUCE + +#else + +#define GGML_F16_DOT_STEP GGML_F16_STEP +#define GGML_F16_DOT_EPR GGML_F16_EPR +#define GGML_F16_DOT_ARR GGML_F16_ARR +#define GGML_F16_DOT_VEC GGML_F16_VEC +#define GGML_F16_DOT_VEC_ZERO GGML_F16_VEC_ZERO +#define GGML_F16_DOT_VEC_LOAD GGML_F16_VEC_LOAD +#define GGML_F16_DOT_VEC_FMA GGML_F16_VEC_FMA +#define GGML_F16_DOT_VEC_REDUCE GGML_F16_VEC_REDUCE + +#endif // defined(__AVX512FP16__) + +#ifdef __cplusplus +} +#endif diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cpu/vec.cpp b/python/freetoken/kernel/csrc/ggml/ggml-cpu/vec.cpp new file mode 100644 index 0000000000..3918a6f4d9 --- /dev/null +++ b/python/freetoken/kernel/csrc/ggml/ggml-cpu/vec.cpp @@ -0,0 +1,613 @@ +#include "vec.h" + +#include + +// precomputed gelu table for f16 (128 KB) +ggml_fp16_t ggml_table_gelu_f16[1 << 16]; + +// precomputed quick gelu table for f16 (128 KB) +ggml_fp16_t ggml_table_gelu_quick_f16[1 << 16]; + +void ggml_vec_dot_f32(int n, float * GGML_RESTRICT s, size_t bs, const float * GGML_RESTRICT x, size_t bx, const float * GGML_RESTRICT y, size_t by, int nrc) { + assert(nrc == 1); + GGML_UNUSED(nrc); + GGML_UNUSED(bx); + GGML_UNUSED(by); + GGML_UNUSED(bs); + +#if defined(GGML_SIMD) + float sumf = 0.0f; + + #if defined(__ARM_FEATURE_SVE) + const int sve_register_length = ggml_cpu_get_sve_cnt() * 8; + const int ggml_f32_epr = sve_register_length / 32;//8;//svcntw(); // SVE128:4, SVE256:8, SVE512:16 + const int ggml_f32_step = 8 * ggml_f32_epr; // choose 8 SVE registers + + const int np = (n & ~(ggml_f32_step - 1)); + svfloat32_t sum1 = svdup_n_f32(0.0f); + svfloat32_t sum2 = svdup_n_f32(0.0f); + svfloat32_t sum3 = svdup_n_f32(0.0f); + svfloat32_t sum4 = svdup_n_f32(0.0f); + svfloat32_t sum5 = svdup_n_f32(0.0f); + svfloat32_t sum6 = svdup_n_f32(0.0f); + svfloat32_t sum7 = svdup_n_f32(0.0f); + svfloat32_t sum8 = svdup_n_f32(0.0f); + svfloat32_t ax1,ax2,ax3,ax4,ax5,ax6,ax7,ax8; + svfloat32_t ay1,ay2,ay3,ay4,ay5,ay6,ay7,ay8; + for (int i = 0; i < np; i += ggml_f32_step) { + ax1 = GGML_F32_VEC_LOAD(x + i); + ay1 = GGML_F32_VEC_LOAD(y + i); + sum1 = GGML_F32_VEC_FMA(sum1, ax1, ay1); + + ax2 = GGML_F32_VEC_LOAD(x + i + 1*ggml_f32_epr); + ay2 = GGML_F32_VEC_LOAD(y + i + 1*ggml_f32_epr); + sum2 = GGML_F32_VEC_FMA(sum2, ax2, ay2); + + ax3 = GGML_F32_VEC_LOAD(x + i + 2*ggml_f32_epr); + ay3 = GGML_F32_VEC_LOAD(y + i + 2*ggml_f32_epr); + sum3 = GGML_F32_VEC_FMA(sum3, ax3, ay3); + + ax4 = GGML_F32_VEC_LOAD(x + i + 3*ggml_f32_epr); + ay4 = GGML_F32_VEC_LOAD(y + i + 3*ggml_f32_epr); + sum4 = GGML_F32_VEC_FMA(sum4, ax4, ay4); + + ax5 = GGML_F32_VEC_LOAD(x + i + 4*ggml_f32_epr); + ay5 = GGML_F32_VEC_LOAD(y + i + 4*ggml_f32_epr); + sum5 = GGML_F32_VEC_FMA(sum5, ax5, ay5); + + ax6 = GGML_F32_VEC_LOAD(x + i + 5*ggml_f32_epr); + ay6 = GGML_F32_VEC_LOAD(y + i + 5*ggml_f32_epr); + sum6 = GGML_F32_VEC_FMA(sum6, ax6, ay6); + + ax7 = GGML_F32_VEC_LOAD(x + i + 6*ggml_f32_epr); + ay7 = GGML_F32_VEC_LOAD(y + i + 6*ggml_f32_epr); + sum7 = GGML_F32_VEC_FMA(sum7, ax7, ay7); + + ax8 = GGML_F32_VEC_LOAD(x + i + 7*ggml_f32_epr); + ay8 = GGML_F32_VEC_LOAD(y + i + 7*ggml_f32_epr); + sum8 = GGML_F32_VEC_FMA(sum8, ax8, ay8); + } + // leftovers + // Since 8 unrolls are done in above loop, leftovers lie in range [0, ggml_f32_step] which is handled in below loop + const int np2 = (n & ~(ggml_f32_epr - 1)); + for (int i = np; i < np2; i += ggml_f32_epr) { + ax1 = GGML_F32_VEC_LOAD(x + i); + ay1 = GGML_F32_VEC_LOAD(y + i); + sum1 = GGML_F32_VEC_FMA(sum1, ax1, ay1); + } + // maximum number of leftover elements will be less that ggml_f32_epr. Apply predicated svmla on available elements only + if (np2 < n) { + svbool_t pg = svwhilelt_b32(np2, n); + ax1 = svld1_f32(pg, x + np2); + ay1 = svld1_f32(pg, y + np2); + sum1 = svmla_f32_m(pg, sum1, ax1, ay1); + } + // reduce sum1,sum2 to sum1 + GGML_F32_VEC_REDUCE(sumf, sum1, sum2, sum3, sum4, sum5, sum6, sum7, sum8); + #elif defined(__riscv_v_intrinsic) + int vl = __riscv_vsetvlmax_e32m8(); + vfloat32m1_t vs = __riscv_vfmv_v_f_f32m1(0.0f, 1); + vfloat32m8_t vsum; + vfloat32m8_t ax; + vfloat32m8_t ay; + vsum = __riscv_vfmv_v_f_f32m8_tu(vsum, 0.0f, vl); + for (int i = 0; i < n; i += vl) { + vl = __riscv_vsetvl_e32m8(n - i); + ax = __riscv_vle32_v_f32m8_tu(ax, &x[i], vl); + ay = __riscv_vle32_v_f32m8_tu(ay, &y[i], vl); + vsum = __riscv_vfmacc_vv_f32m8_tu(vsum, ax, ay, vl); + } + vl = __riscv_vsetvlmax_e32m8(); + vs = __riscv_vfredusum_vs_f32m8_f32m1(vsum, vs, vl); + sumf += __riscv_vfmv_f_s_f32m1_f32(vs); + #else + const int np = (n & ~(GGML_F32_STEP - 1)); + + GGML_F32_VEC sum[GGML_F32_ARR] = { GGML_F32_VEC_ZERO }; + + GGML_F32_VEC ax[GGML_F32_ARR]; + GGML_F32_VEC ay[GGML_F32_ARR]; + + for (int i = 0; i < np; i += GGML_F32_STEP) { + for (int j = 0; j < GGML_F32_ARR; j++) { + ax[j] = GGML_F32_VEC_LOAD(x + i + j*GGML_F32_EPR); + ay[j] = GGML_F32_VEC_LOAD(y + i + j*GGML_F32_EPR); + + sum[j] = GGML_F32_VEC_FMA(sum[j], ax[j], ay[j]); + } + } + + // reduce sum0..sum3 to sum0 + GGML_F32_VEC_REDUCE(sumf, sum); + + // leftovers + for (int i = np; i < n; ++i) { + sumf += x[i]*y[i]; + } + #endif +#else + // scalar + ggml_float sumf = 0.0; + for (int i = 0; i < n; ++i) { + sumf += (ggml_float)(x[i]*y[i]); + } +#endif + + *s = sumf; +} + +void ggml_vec_dot_bf16(int n, float * GGML_RESTRICT s, size_t bs, ggml_bf16_t * GGML_RESTRICT x, size_t bx, ggml_bf16_t * GGML_RESTRICT y, size_t by, int nrc) { + assert(nrc == 1); + GGML_UNUSED(nrc); + GGML_UNUSED(bx); + GGML_UNUSED(by); + GGML_UNUSED(bs); + int i = 0; + ggml_float sumf = 0; + +#if defined(__AVX512BF16__) + __m512 c1 = _mm512_setzero_ps(); + __m512 c2 = _mm512_setzero_ps(); + for (; i + 64 <= n; i += 64) { + c1 = _mm512_dpbf16_ps(c1, m512bh(_mm512_loadu_si512((x + i))), + m512bh(_mm512_loadu_si512((y + i)))); + c2 = _mm512_dpbf16_ps(c2, m512bh(_mm512_loadu_si512((x + i + 32))), + m512bh(_mm512_loadu_si512((y + i + 32)))); + } + sumf += (ggml_float)_mm512_reduce_add_ps(c1); + sumf += (ggml_float)_mm512_reduce_add_ps(c2); + +#elif defined(__AVX512F__) +#define LOAD(p) _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(_mm256_loadu_si256((const __m256i *)(p))), 16)) + __m512 c1 = _mm512_setzero_ps(); + __m512 c2 = _mm512_setzero_ps(); + for (; i + 32 <= n; i += 32) { + c1 = _mm512_add_ps(_mm512_mul_ps(LOAD(x + i), LOAD(y + i)), c1); + c2 = _mm512_add_ps(_mm512_mul_ps(LOAD(x + i + 16), LOAD(y + i + 16)), c2); + } + sumf += (ggml_float)_mm512_reduce_add_ps(c1); + sumf += (ggml_float)_mm512_reduce_add_ps(c2); + +#undef LOAD +#elif defined(__AVX2__) || defined(__AVX__) +#if defined(__AVX2__) +#define LOAD(p) _mm256_castsi256_ps(_mm256_slli_epi32(_mm256_cvtepu16_epi32(_mm_loadu_si128((const __m128i *)(p))), 16)) +#else +#define LOAD(p) _mm256_castsi256_ps(_mm256_insertf128_si256(_mm256_castsi128_si256(_mm_slli_epi32(_mm_cvtepu16_epi32(_mm_loadu_si128((const __m128i *)(p))), 16)), (_mm_slli_epi32(_mm_cvtepu16_epi32(_mm_bsrli_si128(_mm_loadu_si128((const __m128i *)(p)), 8)), 16)), 1)) +#endif + __m256 c1 = _mm256_setzero_ps(); + __m256 c2 = _mm256_setzero_ps(); + __m256 c3 = _mm256_setzero_ps(); + __m256 c4 = _mm256_setzero_ps(); + for (; i + 32 <= n; i += 32) { + c1 = _mm256_add_ps(_mm256_mul_ps(LOAD(x + i), LOAD(y + i)), c1); + c2 = _mm256_add_ps(_mm256_mul_ps(LOAD(x + i + 8), LOAD(y + i + 8)), c2); + c3 = _mm256_add_ps(_mm256_mul_ps(LOAD(x + i + 16), LOAD(y + i + 16)), c3); + c4 = _mm256_add_ps(_mm256_mul_ps(LOAD(x + i + 24), LOAD(y + i + 24)), c4); + } + __m128 g; + c1 = _mm256_add_ps(_mm256_add_ps(c1, c3), + _mm256_add_ps(c2, c4)); + g = _mm_add_ps(_mm256_extractf128_ps(c1, 1), + _mm256_castps256_ps128(c1)); + g = _mm_add_ps(g, _mm_movehl_ps(g, g)); + g = _mm_add_ss(g, _mm_movehdup_ps(g)); + sumf += (ggml_float)_mm_cvtss_f32(g); + +#undef LOAD +#elif defined(__riscv_v_intrinsic) && defined(__riscv_zvfbfwma) + size_t vl = __riscv_vsetvlmax_e32m4(); + + // initialize accumulators to all zeroes + vfloat32m4_t vsum0 = __riscv_vfmv_v_f_f32m4(0.0f, vl); + vfloat32m4_t vsum1 = __riscv_vfmv_v_f_f32m4(0.0f, vl); + + // calculate step size + const size_t epr = __riscv_vsetvlmax_e16m2(); + const size_t step = epr * 2; + const int np = (n & ~(step - 1)); + + // unroll by 2 + for (; i < np; i += step) { + vbfloat16m2_t ax0 = __riscv_vle16_v_bf16m2((const __bf16 *)&x[i], epr); + vbfloat16m2_t ay0 = __riscv_vle16_v_bf16m2((const __bf16 *)&y[i], epr); + vsum0 = __riscv_vfwmaccbf16_vv_f32m4(vsum0, ax0, ay0, epr); + __asm__ __volatile__ ("" ::: "memory"); + + vbfloat16m2_t ax1 = __riscv_vle16_v_bf16m2((const __bf16 *)&x[i + epr], epr); + vbfloat16m2_t ay1 = __riscv_vle16_v_bf16m2((const __bf16 *)&y[i + epr], epr); + vsum1 = __riscv_vfwmaccbf16_vv_f32m4(vsum1, ax1, ay1, epr); + __asm__ __volatile__ ("" ::: "memory"); + } + + // accumulate in 1 register + vsum0 = __riscv_vfadd_vv_f32m4(vsum0, vsum1, vl); + + // leftovers + for (i = np; i < n; i += vl) { + vl = __riscv_vsetvl_e16m2(n - i); + vbfloat16m2_t ax0 = __riscv_vle16_v_bf16m2((const __bf16 *)&x[i], vl); + vbfloat16m2_t ay0 = __riscv_vle16_v_bf16m2((const __bf16 *)&y[i], vl); + vsum0 = __riscv_vfwmaccbf16_vv_f32m4(vsum0, ax0, ay0, vl); + } + + // reduce + vl = __riscv_vsetvlmax_e32m4(); + vfloat32m1_t redsum = __riscv_vfredusum_vs_f32m4_f32m1(vsum0, __riscv_vfmv_v_f_f32m1(0.0f, 1), vl); + sumf += __riscv_vfmv_f_s_f32m1_f32(redsum); + +#elif defined(__POWER9_VECTOR__) || defined(__VXE__) || defined(__VXE2__) + const int np = (n & ~(GGML_BF16_STEP - 1)); + if (np > 0) { + GGML_F32_VEC sum[4] = {GGML_F32_VEC_ZERO}; + for (; i < np; i += GGML_BF16_STEP) { + GGML_BF16_VEC vx0 = GGML_BF16_VEC_LOAD(x + i); + GGML_BF16_VEC vx1 = GGML_BF16_VEC_LOAD(x + i + 8); + GGML_BF16_VEC vy0 = GGML_BF16_VEC_LOAD(y + i); + GGML_BF16_VEC vy1 = GGML_BF16_VEC_LOAD(y + i + 8); + GGML_BF16_FMA_LO(sum[0], vx0, vy0); + GGML_BF16_FMA_HI(sum[1], vx0, vy0); + GGML_BF16_FMA_LO(sum[2], vx1, vy1); + GGML_BF16_FMA_HI(sum[3], vx1, vy1); + } + GGML_F32x4_REDUCE_4(sumf, sum[0], sum[1], sum[2], sum[3]); + } +#endif + + for (; i < n; ++i) { + sumf += (ggml_float)(GGML_BF16_TO_FP32(x[i]) * + GGML_BF16_TO_FP32(y[i])); + } + *s = sumf; +} + +void ggml_vec_dot_f16(int n, float * GGML_RESTRICT s, size_t bs, ggml_fp16_t * GGML_RESTRICT x, size_t bx, ggml_fp16_t * GGML_RESTRICT y, size_t by, int nrc) { + assert(nrc == 1); + GGML_UNUSED(nrc); + GGML_UNUSED(bx); + GGML_UNUSED(by); + GGML_UNUSED(bs); + + ggml_float sumf = 0.0; + + +#if defined(GGML_SIMD) + #if defined(__ARM_FEATURE_SVE) + const int ggml_f16_epr = svcnth(); + const int ggml_f16_step = 8 * ggml_f16_epr; + const int np = n - (n % ggml_f16_step); + const int np2 = n - (n % ggml_f16_epr); + + svfloat32_t sum1_lo = svdup_n_f32(0.0f); + svfloat32_t sum1_hi = svdup_n_f32(0.0f); + svfloat32_t sum2_lo = svdup_n_f32(0.0f); + svfloat32_t sum2_hi = svdup_n_f32(0.0f); + svfloat32_t sum3_lo = svdup_n_f32(0.0f); + svfloat32_t sum3_hi = svdup_n_f32(0.0f); + svfloat32_t sum4_lo = svdup_n_f32(0.0f); + svfloat32_t sum4_hi = svdup_n_f32(0.0f); + + for (int i = 0; i < np; i += ggml_f16_step) { + ggml_sve_f16_fma_widened(&sum1_lo, &sum1_hi, GGML_F16x_VEC_LOAD(x + i + 0 * ggml_f16_epr, 0), GGML_F16x_VEC_LOAD(y + i + 0 * ggml_f16_epr, 0)); + ggml_sve_f16_fma_widened(&sum2_lo, &sum2_hi, GGML_F16x_VEC_LOAD(x + i + 1 * ggml_f16_epr, 1), GGML_F16x_VEC_LOAD(y + i + 1 * ggml_f16_epr, 1)); + ggml_sve_f16_fma_widened(&sum3_lo, &sum3_hi, GGML_F16x_VEC_LOAD(x + i + 2 * ggml_f16_epr, 2), GGML_F16x_VEC_LOAD(y + i + 2 * ggml_f16_epr, 2)); + ggml_sve_f16_fma_widened(&sum4_lo, &sum4_hi, GGML_F16x_VEC_LOAD(x + i + 3 * ggml_f16_epr, 3), GGML_F16x_VEC_LOAD(y + i + 3 * ggml_f16_epr, 3)); + ggml_sve_f16_fma_widened(&sum1_lo, &sum1_hi, GGML_F16x_VEC_LOAD(x + i + 4 * ggml_f16_epr, 4), GGML_F16x_VEC_LOAD(y + i + 4 * ggml_f16_epr, 4)); + ggml_sve_f16_fma_widened(&sum2_lo, &sum2_hi, GGML_F16x_VEC_LOAD(x + i + 5 * ggml_f16_epr, 5), GGML_F16x_VEC_LOAD(y + i + 5 * ggml_f16_epr, 5)); + ggml_sve_f16_fma_widened(&sum3_lo, &sum3_hi, GGML_F16x_VEC_LOAD(x + i + 6 * ggml_f16_epr, 6), GGML_F16x_VEC_LOAD(y + i + 6 * ggml_f16_epr, 6)); + ggml_sve_f16_fma_widened(&sum4_lo, &sum4_hi, GGML_F16x_VEC_LOAD(x + i + 7 * ggml_f16_epr, 7), GGML_F16x_VEC_LOAD(y + i + 7 * ggml_f16_epr, 7)); + } + + for (int i = np; i < np2; i += ggml_f16_epr) { + ggml_sve_f16_fma_widened(&sum1_lo, &sum1_hi, GGML_F16x_VEC_LOAD(x + i, 0), GGML_F16x_VEC_LOAD(y + i, 0)); + } + + if (np2 < n) { + const svbool_t pg = svwhilelt_b16(np2, n); + const svfloat16_t rx = svld1_f16(pg, (const __fp16 *)(x + np2)); + const svfloat16_t ry = svld1_f16(pg, (const __fp16 *)(y + np2)); + + ggml_sve_f16_fma_widened(&sum1_lo, &sum1_hi, rx, ry); + } + + sum1_lo = svadd_f32_m(DEFAULT_PG32, sum1_lo, sum2_lo); + sum1_hi = svadd_f32_m(DEFAULT_PG32, sum1_hi, sum2_hi); + sum3_lo = svadd_f32_m(DEFAULT_PG32, sum3_lo, sum4_lo); + sum3_hi = svadd_f32_m(DEFAULT_PG32, sum3_hi, sum4_hi); + sum1_lo = svadd_f32_m(DEFAULT_PG32, sum1_lo, sum3_lo); + sum1_hi = svadd_f32_m(DEFAULT_PG32, sum1_hi, sum3_hi); + + sumf = ggml_sve_sum_f32x2(sum1_lo, sum1_hi); + #elif defined(__riscv_v_intrinsic) + #if defined(__riscv_zvfh) + int vl = __riscv_vsetvlmax_e32m2(); + vfloat32m1_t vs = __riscv_vfmv_v_f_f32m1(0.0f, 1); + vfloat32m2_t vsum; + vfloat16m1_t ax; + vfloat16m1_t ay; + vsum = __riscv_vreinterpret_v_u32m2_f32m2(__riscv_vmv_v_x_u32m2(0, vl)); + for (int i = 0; i < n; i += vl) { + vl = __riscv_vsetvl_e16m1(n - i); + ax = __riscv_vle16_v_f16m1_tu(ax, (const _Float16 *)&x[i], vl); + ay = __riscv_vle16_v_f16m1_tu(ay, (const _Float16 *)&y[i], vl); + vsum = __riscv_vfwmacc_vv_f32m2_tu(vsum, ax, ay, vl); + } + vl = __riscv_vsetvlmax_e32m1(); + vfloat32m1_t ac0 = __riscv_vfadd_vv_f32m1(__riscv_vget_v_f32m2_f32m1(vsum, 0), __riscv_vget_v_f32m2_f32m1(vsum, 1), vl); + vs = __riscv_vfredusum_vs_f32m1_f32m1(ac0, vs, vl); + sumf += __riscv_vfmv_f_s_f32m1_f32(vs); + #else + for (int i = 0; i < n; ++i) { + sumf += (ggml_float)(GGML_CPU_FP16_TO_FP32(x[i])*GGML_CPU_FP16_TO_FP32(y[i])); + } + #endif // __riscv_zvfh + #else + const int np = (n & ~(GGML_F16_DOT_STEP - 1)); + + GGML_F16_DOT_VEC sum[GGML_F16_DOT_ARR] = { GGML_F16_DOT_VEC_ZERO }; + + GGML_F16_DOT_VEC ax[GGML_F16_DOT_ARR]; + GGML_F16_DOT_VEC ay[GGML_F16_DOT_ARR]; + + for (int i = 0; i < np; i += GGML_F16_DOT_STEP) { + for (int j = 0; j < GGML_F16_DOT_ARR; j++) { + ax[j] = GGML_F16_DOT_VEC_LOAD(x + i + j*GGML_F16_DOT_EPR, j); + ay[j] = GGML_F16_DOT_VEC_LOAD(y + i + j*GGML_F16_DOT_EPR, j); + + sum[j] = GGML_F16_DOT_VEC_FMA(sum[j], ax[j], ay[j]); + } + } + + // reduce sum0..sum3 to sum0 + GGML_F16_DOT_VEC_REDUCE(sumf, sum); + + // leftovers + for (int i = np; i < n; ++i) { + sumf += (ggml_float)(GGML_CPU_FP16_TO_FP32(x[i])*GGML_CPU_FP16_TO_FP32(y[i])); + } + // if you hit this, you are likely running outside the FP range + assert(!isnan(sumf) && !isinf(sumf)); + #endif +#else + for (int i = 0; i < n; ++i) { + sumf += (ggml_float)(GGML_CPU_FP16_TO_FP32(x[i])*GGML_CPU_FP16_TO_FP32(y[i])); + } +#endif // GGML_SIMD + + *s = sumf; +} + +void ggml_vec_silu_f32(const int n, float * y, const float * x) { + int i = 0; +#if defined(__AVX512F__) && defined(__AVX512DQ__) + for (; i + 15 < n; i += 16) { + _mm512_storeu_ps(y + i, ggml_v_silu(_mm512_loadu_ps(x + i))); + } +#elif defined(__AVX2__) && defined(__FMA__) + for (; i + 7 < n; i += 8) { + _mm256_storeu_ps(y + i, ggml_v_silu(_mm256_loadu_ps(x + i))); + } +#elif defined(__SSE2__) + for (; i + 3 < n; i += 4) { + _mm_storeu_ps(y + i, ggml_v_silu(_mm_loadu_ps(x + i))); + } +#elif defined(__ARM_FEATURE_SVE) && defined(__aarch64__) + const int vlen = svcntw(); + for (; i < n; i += vlen) { + const svbool_t pg = svwhilelt_b32_s32(i, n); + svst1_f32(pg, y + i, ggml_v_silu(pg, svld1_f32(pg, x + i))); + } +#elif defined(__ARM_NEON) && defined(__aarch64__) + for (; i + 3 < n; i += 4) { + vst1q_f32(y + i, ggml_v_silu(vld1q_f32(x + i))); + } +#elif defined(__riscv_v_intrinsic) + for (int vl; i < n; i += vl) { + vl = __riscv_vsetvl_e32m2(n - i); + vfloat32m2_t vx = __riscv_vle32_v_f32m2(&x[i], vl); + vfloat32m2_t vy = ggml_v_silu_m2(vx, vl); + __riscv_vse32_v_f32m2(&y[i], vy, vl); + } +#endif + for (; i < n; ++i) { + y[i] = ggml_silu_f32(x[i]); + } +} + +void ggml_vec_swiglu_f32(const int n, float * y, const float * x, const float * g) { + int i = 0; +#if defined(__AVX512F__) && defined(__AVX512DQ__) + for (; i + 15 < n; i += 16) { + _mm512_storeu_ps(y + i, _mm512_mul_ps(ggml_v_silu(_mm512_loadu_ps(x + i)), _mm512_loadu_ps(g + i))); + } +#elif defined(__AVX2__) && defined(__FMA__) + for (; i + 7 < n; i += 8) { + _mm256_storeu_ps(y + i, _mm256_mul_ps(ggml_v_silu(_mm256_loadu_ps(x + i)), _mm256_loadu_ps(g + i))); + } +#elif defined(__SSE2__) + for (; i + 3 < n; i += 4) { + _mm_storeu_ps(y + i, _mm_mul_ps(ggml_v_silu(_mm_loadu_ps(x + i)), _mm_loadu_ps(g + i))); + } +#elif defined(__ARM_FEATURE_SVE) && defined(__aarch64__) + const int vlen = svcntw(); + for (; i < n; i += vlen) { + const svbool_t pg = svwhilelt_b32_s32(i, n); + svst1_f32(pg, y + i, svmul_f32_x(pg, ggml_v_silu(pg, svld1_f32(pg, x + i)), svld1_f32(pg, g + i))); + } +#elif defined(__ARM_NEON) && defined(__aarch64__) + for (; i + 3 < n; i += 4) { + vst1q_f32(y + i, vmulq_f32(ggml_v_silu(vld1q_f32(x + i)), vld1q_f32(g + i))); + } +#elif defined(__riscv_v_intrinsic) + for (int vl; i < n; i += vl) { + vl = __riscv_vsetvl_e32m2(n - i); + vfloat32m2_t vx = __riscv_vle32_v_f32m2(&x[i], vl); + vfloat32m2_t vg = __riscv_vle32_v_f32m2(&g[i], vl); + vfloat32m2_t vy = __riscv_vfmul_vv_f32m2(ggml_v_silu_m2(vx, vl), vg, vl); + __riscv_vse32_v_f32m2(&y[i], vy, vl); + } +#endif + for (; i < n; ++i) { + y[i] = ggml_silu_f32(x[i]) * g[i]; + } +} + +ggml_float ggml_vec_cvar_f32(const int n, float * y, const float * x, const float mean) { + int i = 0; + ggml_float sum = 0; +// TODO: optimize to process the remaining elements in groups using the smaller vector sizes from AVX2 and SSE +// ref: https://github.com/ggml-org/llama.cpp/pull/15953#pullrequestreview-3310928344 +#if defined(__AVX512F__) && defined(__AVX512DQ__) + for (; i + 15 < n; i += 16) { + __m512 val = _mm512_sub_ps(_mm512_loadu_ps(x + i), + _mm512_set1_ps(mean)); + _mm512_storeu_ps(y + i, val); + sum += (ggml_float)_mm512_reduce_add_ps(_mm512_mul_ps(val, val)); + } +#elif defined(__AVX2__) && defined(__FMA__) + for (; i + 7 < n; i += 8) { + __m256 val = _mm256_sub_ps(_mm256_loadu_ps(x + i), + _mm256_set1_ps(mean)); + _mm256_storeu_ps(y + i, val); + val = _mm256_mul_ps(val,val); + __m128 val2 = _mm_add_ps(_mm256_extractf128_ps(val, 1), + _mm256_castps256_ps128(val)); + val2 = _mm_add_ps(val2, _mm_movehl_ps(val2, val2)); + val2 = _mm_add_ss(val2, _mm_movehdup_ps(val2)); + sum += (ggml_float)_mm_cvtss_f32(val2); + } +#elif defined(__SSE2__) + for (; i + 3 < n; i += 4) { + __m128 val = _mm_sub_ps(_mm_loadu_ps(x + i), + _mm_set1_ps(mean)); + _mm_storeu_ps(y + i, val); + val = _mm_mul_ps(val, val); +#if defined(__AVX__) || defined(__AVX2__) || defined(__AVX512F__) + val = _mm_add_ps(val, _mm_movehl_ps(val, val)); + val = _mm_add_ss(val, _mm_movehdup_ps(val)); +#else + __m128 tmp = _mm_shuffle_ps(val, val, _MM_SHUFFLE(2, 3, 0, 1)); + val = _mm_add_ps(val, tmp); + tmp = _mm_movehl_ps(tmp, val); + val = _mm_add_ss(val, tmp); +#endif // __AVX__ || __AVX2__ || __AVX512F__ + sum += (ggml_float)_mm_cvtss_f32(val); + } +#elif defined(__ARM_NEON) && defined(__aarch64__) + for (; i + 3 < n; i += 4) { + float32x4_t val = vsubq_f32(vld1q_f32(x + i), + vdupq_n_f32(mean)); + vst1q_f32(y + i, val); + val = vmulq_f32(val, val); + sum += (ggml_float)vaddvq_f32(val); + } +#elif defined(__VXE__) || defined(__VXE2__) + for (; i + 3 < n; i += 4) { + float32x4_t val = vec_sub(vec_xl(0, x + i), vec_splats(mean)); + vec_xst(val, 0, y + i); + val = vec_mul(val, val); + sum += (ggml_float)vec_hsum_f32x4(val); + } +#elif defined(__riscv_v_intrinsic) + vfloat64m1_t vsum = __riscv_vfmv_v_f_f64m1(0, 1); + for (int vl; i < n; i += vl) { + vl = __riscv_vsetvl_e32m2(n - i); + vfloat32m2_t val = __riscv_vfsub_vf_f32m2(__riscv_vle32_v_f32m2(&x[i], vl), mean, vl); + __riscv_vse32_v_f32m2(&y[i], val, vl); + val = __riscv_vfmul_vv_f32m2(val, val, vl); + vsum = __riscv_vfwredusum_vs_f32m2_f64m1(val, vsum, vl); + } + sum = (ggml_float)__riscv_vfmv_f_s_f64m1_f64(vsum); +#endif + for (; i < n; ++i) { + float val = x[i] - mean; + y[i] = val; + val *= val; + sum += (ggml_float)val; + } + return sum/n; +} + +ggml_float ggml_vec_soft_max_f32(const int n, float * y, const float * x, float max) { + int i = 0; + ggml_float sum = 0; +#if defined(__AVX512F__) && defined(__AVX512DQ__) + for (; i + 15 < n; i += 16) { + __m512 val = ggml_v_expf(_mm512_sub_ps(_mm512_loadu_ps(x + i), + _mm512_set1_ps(max))); + _mm512_storeu_ps(y + i, val); + sum += (ggml_float)_mm512_reduce_add_ps(val); + } +#elif defined(__AVX2__) && defined(__FMA__) + for (; i + 7 < n; i += 8) { + __m256 val = ggml_v_expf(_mm256_sub_ps(_mm256_loadu_ps(x + i), + _mm256_set1_ps(max))); + _mm256_storeu_ps(y + i, val); + __m128 val2 = _mm_add_ps(_mm256_extractf128_ps(val, 1), + _mm256_castps256_ps128(val)); + val2 = _mm_add_ps(val2, _mm_movehl_ps(val2, val2)); + val2 = _mm_add_ss(val2, _mm_movehdup_ps(val2)); + sum += (ggml_float)_mm_cvtss_f32(val2); + } +#elif defined(__SSE2__) + for (; i + 3 < n; i += 4) { + __m128 val = ggml_v_expf(_mm_sub_ps(_mm_loadu_ps(x + i), + _mm_set1_ps(max))); + _mm_storeu_ps(y + i, val); +#if defined(__AVX__) || defined(__AVX2__) || defined(__AVX512F__) + val = _mm_add_ps(val, _mm_movehl_ps(val, val)); + val = _mm_add_ss(val, _mm_movehdup_ps(val)); +#else + __m128 tmp = _mm_shuffle_ps(val, val, _MM_SHUFFLE(2, 3, 0, 1)); + val = _mm_add_ps(val, tmp); + tmp = _mm_movehl_ps(tmp, val); + val = _mm_add_ss(val, tmp); +#endif + sum += (ggml_float)_mm_cvtss_f32(val); + } +#elif defined(__ARM_FEATURE_SVE) && defined(__aarch64__) + const int vlen = svcntw(); + for (; i < n; i += vlen) { + const svbool_t pg = svwhilelt_b32_s32(i, n); + svfloat32_t val = ggml_v_expf(pg, svsub_f32_x(pg, svld1_f32(pg, x + i), + svdup_n_f32_x(pg, max))); + svst1_f32(pg, y + i, val); + sum += (ggml_float)svaddv_f32(pg, val); + } +#elif defined(__ARM_NEON) && defined(__aarch64__) + for (; i + 3 < n; i += 4) { + float32x4_t val = ggml_v_expf(vsubq_f32(vld1q_f32(x + i), + vdupq_n_f32(max))); + vst1q_f32(y + i, val); + sum += (ggml_float)vaddvq_f32(val); + } +#elif defined(__riscv_v_intrinsic) + vfloat64m1_t vsum = __riscv_vfmv_v_f_f64m1(0, 1); + for (int avl; i < n; i += avl) { + avl = __riscv_vsetvl_e32m2(n - i); + vfloat32m2_t val = ggml_v_expf_m2(__riscv_vfsub_vf_f32m2(__riscv_vle32_v_f32m2(&x[i], avl), max, avl), avl); + __riscv_vse32_v_f32m2(&y[i], val, avl); + vsum = __riscv_vfwredusum_vs_f32m2_f64m1(val, vsum, avl); + } + return (ggml_float)__riscv_vfmv_f_s_f64m1_f64(vsum); +#endif + for (; i < n; ++i) { + float val = expf(x[i] - max); + sum += (ggml_float)val; + y[i] = val; + } + return sum; +} + +ggml_float ggml_vec_log_soft_max_f32(const int n, float * y, const float * x, float max) { + // log(soft_max) = log(soft_max_i / soft_max_sum) = log(soft_max_i) - log(soft_max_sum) = (logit_i - max) - log(soft_max_i) + + int i = 0; + ggml_float sum = 0; + for (; i < n; ++i) { + float val = x[i] - max; + y[i] = val; + sum += (ggml_float)expf(val); + } + return sum = (ggml_float)logf(sum); +} diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cpu/vec.h b/python/freetoken/kernel/csrc/ggml/ggml-cpu/vec.h new file mode 100644 index 0000000000..59fa88f2c9 --- /dev/null +++ b/python/freetoken/kernel/csrc/ggml/ggml-cpu/vec.h @@ -0,0 +1,1630 @@ +// Vectorized functions for fundamental operations + +#pragma once + +#include "ggml-impl.h" +#include "simd-mappings.h" +#include "ggml.h" +#include "ggml-cpu.h" + +#if defined(GGML_USE_ACCELERATE) +#include +#endif + +// floating point type used to accumulate sums +typedef double ggml_float; + +#if defined(__ARM_FEATURE_SVE) +inline static void ggml_sve_f16_fma_widened( + svfloat32_t * acc_lo, + svfloat32_t * acc_hi, + svfloat16_t x, + svfloat16_t y) { +#if defined(__ARM_FEATURE_SVE2) + *acc_lo = svmlalb_f32(*acc_lo, x, y); + *acc_hi = svmlalt_f32(*acc_hi, x, y); +#else + // Plain SVE fallback path if SVE2 instructions not available + svfloat16_t x_even = svtrn1_f16(x, x); + svfloat16_t x_odd = svtrn2_f16(x, x); + + svfloat16_t y_even = svtrn1_f16(y, y); + svfloat16_t y_odd = svtrn2_f16(y, y); + + svbool_t pg = svptrue_b32(); + + *acc_lo = svmla_f32_x(pg, *acc_lo, svcvt_f32_f16_x(pg, x_even), svcvt_f32_f16_x(pg, y_even)); + *acc_hi = svmla_f32_x(pg, *acc_hi, svcvt_f32_f16_x(pg, x_odd), svcvt_f32_f16_x(pg, y_odd)); +#endif +} + +inline static ggml_float ggml_sve_sum_f32x2(svfloat32_t sum_lo, svfloat32_t sum_hi) { + return (ggml_float) (svaddv_f32(svptrue_b32(), sum_lo) + svaddv_f32(svptrue_b32(), sum_hi)); +} +#endif + +#define GGML_GELU_FP16 +#define GGML_GELU_QUICK_FP16 + +#define GGML_SOFT_MAX_UNROLL 4 +#define GGML_VEC_DOT_UNROLL 2 +#define GGML_VEC_MAD_UNROLL 32 + +#ifdef __cplusplus +extern "C" { +#endif + +// +// global data +// + +// precomputed gelu table for f16 (128 KB) +extern ggml_fp16_t ggml_table_gelu_f16[1 << 16]; + +// precomputed quick gelu table for f16 (128 KB) +extern ggml_fp16_t ggml_table_gelu_quick_f16[1 << 16]; + +// +// fundamental operations +// + +void ggml_vec_dot_f32(int n, float * GGML_RESTRICT s, size_t bs, const float * GGML_RESTRICT x, size_t bx, const float * GGML_RESTRICT y, size_t by, int nrc); +void ggml_vec_dot_bf16(int n, float * GGML_RESTRICT s, size_t bs, ggml_bf16_t * GGML_RESTRICT x, size_t bx, ggml_bf16_t * GGML_RESTRICT y, size_t by, int nrc); +void ggml_vec_dot_f16(int n, float * GGML_RESTRICT s, size_t bs, ggml_fp16_t * GGML_RESTRICT x, size_t bx, ggml_fp16_t * GGML_RESTRICT y, size_t by, int nrc); + +void ggml_vec_silu_f32(const int n, float * y, const float * x); +ggml_float ggml_vec_cvar_f32(const int n, float * y, const float * x, const float mean); //it will also center y ( y = y - mean ) +ggml_float ggml_vec_soft_max_f32(const int n, float * y, const float * x, float max); +ggml_float ggml_vec_log_soft_max_f32(const int n, float * y, const float * x, float max); + +inline static void ggml_vec_set_i8(const int n, int8_t * x, const int8_t v) { for (int i = 0; i < n; ++i) x[i] = v; } +inline static void ggml_vec_set_i16(const int n, int16_t * x, const int16_t v) { for (int i = 0; i < n; ++i) x[i] = v; } + +inline static void ggml_vec_set_i32(const int n, int32_t * x, const int32_t v) { for (int i = 0; i < n; ++i) x[i] = v; } +inline static void ggml_vec_cpy_i32(const int n, int32_t * y, const int32_t * x) { for (int i = 0; i < n; ++i) y[i] = x[i]; } + +inline static void ggml_vec_set_f16(const int n, ggml_fp16_t * x, const ggml_fp16_t v) { for (int i = 0; i < n; ++i) x[i] = v; } +inline static void ggml_vec_set_bf16(const int n, ggml_bf16_t * x, const ggml_bf16_t v) { for (int i = 0; i < n; ++i) x[i] = v; } + +inline static void ggml_vec_add_f32 (const int n, float * z, const float * x, const float * y) { + int i = 0; +#if defined(__AVX2__) + for (; i + 7 < n; i += 8) { + __m256 vx = _mm256_loadu_ps(x + i); + __m256 vy = _mm256_loadu_ps(y + i); + __m256 vz = _mm256_add_ps(vx, vy); + _mm256_storeu_ps(z + i, vz); + } +#endif + for (; i < n; ++i) { + z[i] = x[i] + y[i]; + } +} + +inline static void ggml_vec_add_f16 (const int n, ggml_fp16_t * z, const ggml_fp16_t * x, const ggml_fp16_t * y) { + for (int i = 0; i < n; ++i) { + z[i] = GGML_CPU_FP32_TO_FP16(GGML_CPU_FP16_TO_FP32(x[i]) + GGML_CPU_FP16_TO_FP32(y[i])); + } +} +inline static void ggml_vec_add1_f32(const int n, float * z, const float * x, const float v) { for (int i = 0; i < n; ++i) z[i] = x[i] + v; } +inline static void ggml_vec_acc_f32 (const int n, float * y, const float * x) { for (int i = 0; i < n; ++i) y[i] += x[i]; } +inline static void ggml_vec_acc1_f32(const int n, float * y, const float v) { for (int i = 0; i < n; ++i) y[i] += v; } +inline static void ggml_vec_sub_f32 (const int n, float * z, const float * x, const float * y) { for (int i = 0; i < n; ++i) z[i] = x[i] - y[i]; } +inline static void ggml_vec_sub_f16 (const int n, ggml_fp16_t * z, const ggml_fp16_t * x, const ggml_fp16_t * y) { + for (int i = 0; i < n; ++i) { + z[i] = GGML_CPU_FP32_TO_FP16(GGML_CPU_FP16_TO_FP32(x[i]) - GGML_CPU_FP16_TO_FP32(y[i])); + } +} +inline static void ggml_vec_set_f32 (const int n, float * x, const float v) { for (int i = 0; i < n; ++i) x[i] = v; } +inline static void ggml_vec_cpy_f32 (const int n, float * y, const float * x) { for (int i = 0; i < n; ++i) y[i] = x[i]; } +inline static void ggml_vec_neg_f32 (const int n, float * y, const float * x) { for (int i = 0; i < n; ++i) y[i] = -x[i]; } +inline static void ggml_vec_neg_f16 (const int n, ggml_fp16_t * y, const ggml_fp16_t * x) { + for (int i = 0; i < n; ++i) { + y[i] = GGML_CPU_FP32_TO_FP16(-GGML_CPU_FP16_TO_FP32(x[i])); + } +} + +inline static void ggml_vec_mul_f32 (const int n, float * z, const float * x, const float * y) { for (int i = 0; i < n; ++i) z[i] = x[i]*y[i]; } +inline static void ggml_vec_mul_f16 (const int n, ggml_fp16_t * z, const ggml_fp16_t * x, const ggml_fp16_t * y) { + for (int i = 0; i < n; ++i) { + z[i] = GGML_CPU_FP32_TO_FP16(GGML_CPU_FP16_TO_FP32(x[i]) * GGML_CPU_FP16_TO_FP32(y[i])); + } +} +inline static void ggml_vec_div_f32 (const int n, float * z, const float * x, const float * y) { for (int i = 0; i < n; ++i) z[i] = x[i]/y[i]; } +inline static void ggml_vec_div_f16 (const int n, ggml_fp16_t * z, const ggml_fp16_t * x, const ggml_fp16_t * y) { + for (int i = 0; i < n; ++i) { + z[i] = GGML_CPU_FP32_TO_FP16(GGML_CPU_FP16_TO_FP32(x[i]) / GGML_CPU_FP16_TO_FP32(y[i])); + } +} + +// compute GGML_VEC_DOT_UNROLL dot products at once +// xs - x row stride in bytes +inline static void ggml_vec_dot_f16_unroll(const int n, const int xs, float * GGML_RESTRICT s, void * GGML_RESTRICT xv, ggml_fp16_t * GGML_RESTRICT y) { + ggml_float sumf[GGML_VEC_DOT_UNROLL] = { 0.0 }; + + ggml_fp16_t * GGML_RESTRICT x[GGML_VEC_DOT_UNROLL]; + + for (int i = 0; i < GGML_VEC_DOT_UNROLL; ++i) { + x[i] = (ggml_fp16_t *) ((char *) xv + i*xs); + } + +#if defined(GGML_SIMD) + #if defined(__ARM_FEATURE_SVE) + + const int ggml_f16_epr = svcnth(); + const int ggml_f16_step = 2 * ggml_f16_epr; + int np = n - (n % ggml_f16_step); + int np2 = n - (n % ggml_f16_epr); + + svfloat32_t sum_0_0_lo = svdup_n_f32(0.0f); + svfloat32_t sum_0_0_hi = svdup_n_f32(0.0f); + svfloat32_t sum_0_1_lo = svdup_n_f32(0.0f); + svfloat32_t sum_0_1_hi = svdup_n_f32(0.0f); + svfloat32_t sum_1_0_lo = svdup_n_f32(0.0f); + svfloat32_t sum_1_0_hi = svdup_n_f32(0.0f); + svfloat32_t sum_1_1_lo = svdup_n_f32(0.0f); + svfloat32_t sum_1_1_hi = svdup_n_f32(0.0f); + + for (int i = 0; i < np; i += ggml_f16_step) { + const svfloat16_t ay0 = GGML_F16x_VEC_LOAD(y + i, 0); + const svfloat16_t ax00 = GGML_F16x_VEC_LOAD(x[0] + i, 0); + const svfloat16_t ax01 = GGML_F16x_VEC_LOAD(x[1] + i, 0); + + ggml_sve_f16_fma_widened(&sum_0_0_lo, &sum_0_0_hi, ax00, ay0); + ggml_sve_f16_fma_widened(&sum_1_0_lo, &sum_1_0_hi, ax01, ay0); + + const svfloat16_t ay1 = GGML_F16x_VEC_LOAD(y + i + 1 * ggml_f16_epr, 0); + const svfloat16_t ax10 = GGML_F16x_VEC_LOAD(x[0] + i + 1 * ggml_f16_epr, 0); + const svfloat16_t ax11 = GGML_F16x_VEC_LOAD(x[1] + i + 1 * ggml_f16_epr, 0); + + ggml_sve_f16_fma_widened(&sum_0_1_lo, &sum_0_1_hi, ax10, ay1); + ggml_sve_f16_fma_widened(&sum_1_1_lo, &sum_1_1_hi, ax11, ay1); + } + + for (int i = np; i < np2; i += ggml_f16_epr) { + const svfloat16_t ry = GGML_F16x_VEC_LOAD(y + i, 0); + const svfloat16_t rx0 = GGML_F16x_VEC_LOAD(x[0] + i, 0); + const svfloat16_t rx1 = GGML_F16x_VEC_LOAD(x[1] + i, 0); + + ggml_sve_f16_fma_widened(&sum_0_0_lo, &sum_0_0_hi, rx0, ry); + ggml_sve_f16_fma_widened(&sum_1_0_lo, &sum_1_0_hi, rx1, ry); + } + + if (np2 < n) { + const svbool_t pg = svwhilelt_b16(np2, n); + const svfloat16_t ay = svld1_f16(pg, (const __fp16 *)(y + np2)); + const svfloat16_t ax0 = svld1_f16(pg, (const __fp16 *)(x[0] + np2)); + const svfloat16_t ax1 = svld1_f16(pg, (const __fp16 *)(x[1] + np2)); + + ggml_sve_f16_fma_widened(&sum_0_0_lo, &sum_0_0_hi, ax0, ay); + ggml_sve_f16_fma_widened(&sum_1_0_lo, &sum_1_0_hi, ax1, ay); + } + + svfloat32_t sum_0_lo = svadd_f32_x(DEFAULT_PG32, sum_0_0_lo, sum_0_1_lo); + svfloat32_t sum_0_hi = svadd_f32_x(DEFAULT_PG32, sum_0_0_hi, sum_0_1_hi); + svfloat32_t sum_1_lo = svadd_f32_x(DEFAULT_PG32, sum_1_0_lo, sum_1_1_lo); + svfloat32_t sum_1_hi = svadd_f32_x(DEFAULT_PG32, sum_1_0_hi, sum_1_1_hi); + sumf[0] = ggml_sve_sum_f32x2(sum_0_lo, sum_0_hi); + sumf[1] = ggml_sve_sum_f32x2(sum_1_lo, sum_1_hi); + np = n; + #elif defined(__riscv_v_intrinsic) + #if defined(__riscv_zvfh) + size_t vl = __riscv_vsetvlmax_e32m4(); + + // initialize accumulators to all zeroes + vfloat32m4_t vsum0_0 = __riscv_vfmv_v_f_f32m4(0.0f, vl); + vfloat32m4_t vsum0_1 = __riscv_vfmv_v_f_f32m4(0.0f, vl); + vfloat32m4_t vsum1_0 = __riscv_vfmv_v_f_f32m4(0.0f, vl); + vfloat32m4_t vsum1_1 = __riscv_vfmv_v_f_f32m4(0.0f, vl); + + // calculate step size + const size_t epr = __riscv_vsetvlmax_e16m2(); + const size_t step = epr * 2; + int np = (n & ~(step - 1)); + + // unroll by 2 along the row dimension + for (int i = 0; i < np; i += step) { + vfloat16m2_t ay0 = __riscv_vle16_v_f16m2((const _Float16 *)(y + i), epr); + vfloat16m2_t ax0_0 = __riscv_vle16_v_f16m2((const _Float16 *)(x[0] + i), epr); + vfloat16m2_t ax1_0 = __riscv_vle16_v_f16m2((const _Float16 *)(x[1] + i), epr); + vsum0_0 = __riscv_vfwmacc_vv_f32m4(vsum0_0, ax0_0, ay0, epr); + vsum1_0 = __riscv_vfwmacc_vv_f32m4(vsum1_0, ax1_0, ay0, epr); + + vfloat16m2_t ay1 = __riscv_vle16_v_f16m2((const _Float16 *)(y + i + epr), epr); + vfloat16m2_t ax0_1 = __riscv_vle16_v_f16m2((const _Float16 *)(x[0] + i + epr), epr); + vfloat16m2_t ax1_1 = __riscv_vle16_v_f16m2((const _Float16 *)(x[1] + i + epr), epr); + vsum0_1 = __riscv_vfwmacc_vv_f32m4(vsum0_1, ax0_1, ay1, epr); + vsum1_1 = __riscv_vfwmacc_vv_f32m4(vsum1_1, ax1_1, ay1, epr); + } + + vfloat32m4_t vsum0 = __riscv_vfadd_vv_f32m4(vsum0_0, vsum0_1, vl); + vfloat32m4_t vsum1 = __riscv_vfadd_vv_f32m4(vsum1_0, vsum1_1, vl); + + // leftovers + for (int i = np; i < n; i += vl) { + vl = __riscv_vsetvl_e16m2(n - i); + vfloat16m2_t ay = __riscv_vle16_v_f16m2((const _Float16 *)(y + i), vl); + vfloat16m2_t ax0 = __riscv_vle16_v_f16m2((const _Float16 *)(x[0] + i), vl); + vfloat16m2_t ax1 = __riscv_vle16_v_f16m2((const _Float16 *)(x[1] + i), vl); + + vsum0 = __riscv_vfwmacc_vv_f32m4(vsum0, ax0, ay, vl); + vsum1 = __riscv_vfwmacc_vv_f32m4(vsum1, ax1, ay, vl); + } + + // reduce + vl = __riscv_vsetvlmax_e32m2(); + vfloat32m2_t acc0_0 = __riscv_vfadd_vv_f32m2(__riscv_vget_v_f32m4_f32m2(vsum0, 0), + __riscv_vget_v_f32m4_f32m2(vsum0, 1), vl); + vl = __riscv_vsetvlmax_e32m1(); + vfloat32m1_t acc0_1 = __riscv_vfadd_vv_f32m1(__riscv_vget_v_f32m2_f32m1(acc0_0, 0), + __riscv_vget_v_f32m2_f32m1(acc0_0, 1), vl); + vfloat32m1_t redsum0 = __riscv_vfredusum_vs_f32m1_f32m1( + acc0_1, __riscv_vfmv_v_f_f32m1(0.0f, 1), vl); + + vl = __riscv_vsetvlmax_e32m2(); + vfloat32m2_t acc1_0 = __riscv_vfadd_vv_f32m2(__riscv_vget_v_f32m4_f32m2(vsum1, 0), + __riscv_vget_v_f32m4_f32m2(vsum1, 1), vl); + vl = __riscv_vsetvlmax_e32m1(); + vfloat32m1_t acc1_1 = __riscv_vfadd_vv_f32m1(__riscv_vget_v_f32m2_f32m1(acc1_0, 0), + __riscv_vget_v_f32m2_f32m1(acc1_0, 1), vl); + vfloat32m1_t redsum1 = __riscv_vfredusum_vs_f32m1_f32m1( + acc1_1, __riscv_vfmv_v_f_f32m1(0.0f, 1), vl); + sumf[0] = __riscv_vfmv_f_s_f32m1_f32(redsum0); + sumf[1] = __riscv_vfmv_f_s_f32m1_f32(redsum1); + np = n; + #else + const int np = 0; + #endif + #else + const int np = (n & ~(GGML_F16_DOT_STEP - 1)); + + GGML_F16_DOT_VEC sum[GGML_VEC_DOT_UNROLL][GGML_F16_DOT_ARR] = { { GGML_F16_DOT_VEC_ZERO } }; + + GGML_F16_DOT_VEC ax[GGML_F16_DOT_ARR]; + GGML_F16_DOT_VEC ay[GGML_F16_DOT_ARR]; + + for (int i = 0; i < np; i += GGML_F16_DOT_STEP) { + for (int j = 0; j < GGML_F16_DOT_ARR; j++) { + ay[j] = GGML_F16_DOT_VEC_LOAD(y + i + j*GGML_F16_DOT_EPR, j); + + for (int k = 0; k < GGML_VEC_DOT_UNROLL; ++k) { + ax[j] = GGML_F16_DOT_VEC_LOAD(x[k] + i + j*GGML_F16_DOT_EPR, j); + + sum[k][j] = GGML_F16_DOT_VEC_FMA(sum[k][j], ax[j], ay[j]); + } + } + } + + // reduce sum0..sum3 to sum0 + for (int k = 0; k < GGML_VEC_DOT_UNROLL; ++k) { + GGML_F16_DOT_VEC_REDUCE(sumf[k], sum[k]); + } + #endif +#else + // scalar path + const int np = 0; +#endif + // scalar and leftovers + for (int i = np; i < n; ++i) { + for (int j = 0; j < GGML_VEC_DOT_UNROLL; ++j) { + sumf[j] += (ggml_float)(GGML_CPU_FP16_TO_FP32(x[j][i])*GGML_CPU_FP16_TO_FP32(y[i])); + } + } + + for (int i = 0; i < GGML_VEC_DOT_UNROLL; ++i) { + s[i] = (float)sumf[i]; + } +} + +inline static void ggml_vec_mad_f32(const int n, float * GGML_RESTRICT y, const float * GGML_RESTRICT x, const float v) { +#if defined(GGML_SIMD) + #if defined(__ARM_FEATURE_SVE) + + const int sve_register_length = ggml_cpu_get_sve_cnt() * 8; + const int ggml_f32_epr = sve_register_length / 32;//8;//svcntw(); // SVE128:4, SVE256:8, SVE512:16 + const int ggml_f32_step = 8 * ggml_f32_epr; // choose 8 SVE registers + GGML_F32_VEC vx = GGML_F32_VEC_SET1(v); + + const int np = (n & ~(ggml_f32_step - 1)); + svfloat32_t ax1, ax2, ax3, ax4, ax5, ax6, ax7, ax8; + svfloat32_t ay1, ay2, ay3, ay4, ay5, ay6, ay7, ay8; + for (int i = 0; i < np; i += ggml_f32_step) { + + ax1 = GGML_F32_VEC_LOAD(x + i); + ay1 = GGML_F32_VEC_LOAD(y + i); + ay1 = GGML_F32_VEC_FMA(ay1, ax1, vx); + + GGML_F32_VEC_STORE(y + i, ay1); + + ax2 = GGML_F32_VEC_LOAD(x + i + 1*ggml_f32_epr); + ay2 = GGML_F32_VEC_LOAD(y + i + 1*ggml_f32_epr); + ay2 = GGML_F32_VEC_FMA(ay2, ax2, vx); + + GGML_F32_VEC_STORE(y + i + 1*ggml_f32_epr, ay2); + + ax3 = GGML_F32_VEC_LOAD(x + i + 2*ggml_f32_epr); + ay3 = GGML_F32_VEC_LOAD(y + i + 2*ggml_f32_epr); + ay3 = GGML_F32_VEC_FMA(ay3, ax3, vx); + + GGML_F32_VEC_STORE(y + i + 2*ggml_f32_epr, ay3); + + ax4 = GGML_F32_VEC_LOAD(x + i + 3*ggml_f32_epr); + ay4 = GGML_F32_VEC_LOAD(y + i + 3*ggml_f32_epr); + ay4 = GGML_F32_VEC_FMA(ay4, ax4, vx); + + GGML_F32_VEC_STORE(y + i + 3*ggml_f32_epr, ay4); + + ax5 = GGML_F32_VEC_LOAD(x + i + 4*ggml_f32_epr); + ay5 = GGML_F32_VEC_LOAD(y + i + 4*ggml_f32_epr); + ay5 = GGML_F32_VEC_FMA(ay5, ax5, vx); + + GGML_F32_VEC_STORE(y + i + 4*ggml_f32_epr, ay5); + + ax6 = GGML_F32_VEC_LOAD(x + i + 5*ggml_f32_epr); + ay6 = GGML_F32_VEC_LOAD(y + i + 5*ggml_f32_epr); + ay6 = GGML_F32_VEC_FMA(ay6, ax6, vx); + + GGML_F32_VEC_STORE(y + i + 5*ggml_f32_epr, ay6); + + ax7 = GGML_F32_VEC_LOAD(x + i + 6*ggml_f32_epr); + ay7 = GGML_F32_VEC_LOAD(y + i + 6*ggml_f32_epr); + ay7 = GGML_F32_VEC_FMA(ay7, ax7, vx); + + GGML_F32_VEC_STORE(y + i + 6*ggml_f32_epr, ay7); + + ax8 = GGML_F32_VEC_LOAD(x + i + 7*ggml_f32_epr); + ay8 = GGML_F32_VEC_LOAD(y + i + 7*ggml_f32_epr); + ay8 = GGML_F32_VEC_FMA(ay8, ax8, vx); + + GGML_F32_VEC_STORE(y + i + 7*ggml_f32_epr, ay8); + } + // leftovers + // Since 8 unrolls are done in above loop, leftovers lie in range [0, ggml_f32_step] which is handled in below loop + const int np2 = (n & ~(ggml_f32_epr - 1)); + for (int i = np; i < np2; i += ggml_f32_epr) { + ax1 = GGML_F32_VEC_LOAD(x + i); + ay1 = GGML_F32_VEC_LOAD(y + i); + ay1 = GGML_F32_VEC_FMA(ay1, ax1, vx); + + GGML_F32_VEC_STORE(y + i, ay1); + } + // maximum number of leftover elements will be less that ggml_f32_epr. Apply predicated svmad on available elements only + if (np2 < n) { + svbool_t pg =svwhilelt_b32(np2, n); + ax1 = svld1_f32(pg, x + np2); + ay1 = svld1_f32(pg, y + np2); + ay1 = svmad_f32_m(pg, ax1, vx, ay1); + + svst1_f32(pg, y + np2, ay1); + } + #elif defined(__riscv_v_intrinsic) + for (int i = 0, avl; i < n; i += avl) { + avl = __riscv_vsetvl_e32m8(n - i); + vfloat32m8_t ax = __riscv_vle32_v_f32m8(&x[i], avl); + vfloat32m8_t ay = __riscv_vle32_v_f32m8(&y[i], avl); + vfloat32m8_t ny = __riscv_vfmadd_vf_f32m8(ax, v, ay, avl); + __riscv_vse32_v_f32m8(&y[i], ny, avl); + } + #else + const int np = (n & ~(GGML_F32_STEP - 1)); + + GGML_F32_VEC vx = GGML_F32_VEC_SET1(v); + + GGML_F32_VEC ax[GGML_F32_ARR]; + GGML_F32_VEC ay[GGML_F32_ARR]; + + for (int i = 0; i < np; i += GGML_F32_STEP) { + for (int j = 0; j < GGML_F32_ARR; j++) { + ax[j] = GGML_F32_VEC_LOAD(x + i + j*GGML_F32_EPR); + ay[j] = GGML_F32_VEC_LOAD(y + i + j*GGML_F32_EPR); + ay[j] = GGML_F32_VEC_FMA(ay[j], ax[j], vx); + + GGML_F32_VEC_STORE(y + i + j*GGML_F32_EPR, ay[j]); + } + } + + // leftovers + for (int i = np; i < n; ++i) { + y[i] += x[i]*v; + } + #endif +#else + // scalar + for (int i = 0; i < n; ++i) { + y[i] += x[i]*v; + } +#endif +} + +inline static void ggml_vec_mad_f16(const int n, ggml_fp16_t * GGML_RESTRICT y, const ggml_fp16_t * GGML_RESTRICT x, const float v) { +#if defined(GGML_SIMD) && defined(__ARM_FEATURE_SVE) + const int sve_register_length = svcntb() * 8; + const int ggml_f16_epr = sve_register_length / 16; + const int ggml_f16_step = 8 * ggml_f16_epr; + + GGML_F16x_VEC vx = GGML_F16x_VEC_SET1(v); + + int np = (n & ~(ggml_f16_step - 1)); + + svfloat16_t ax1, ax2, ax3, ax4, ax5, ax6, ax7, ax8; + svfloat16_t ay1, ay2, ay3, ay4, ay5, ay6, ay7, ay8; + for (int i = 0; i < np; i += ggml_f16_step) { + ax1 = GGML_F16x_VEC_LOAD(x + i + 0 * ggml_f16_epr, 0); + ay1 = GGML_F16x_VEC_LOAD(y + i + 0 * ggml_f16_epr, 0); + ay1 = GGML_F16x_VEC_FMA(ay1, ax1, vx); + + GGML_F16x_VEC_STORE(y + i + 0 * ggml_f16_epr, ay1, 0); + + ax2 = GGML_F16x_VEC_LOAD(x + i + 1 * ggml_f16_epr, 1); + ay2 = GGML_F16x_VEC_LOAD(y + i + 1 * ggml_f16_epr, 1); + ay2 = GGML_F16x_VEC_FMA(ay2, ax2, vx); + + GGML_F16x_VEC_STORE(y + i + 1 * ggml_f16_epr, ay2, 1); + + ax3 = GGML_F16x_VEC_LOAD(x + i + 2 * ggml_f16_epr, 2); + ay3 = GGML_F16x_VEC_LOAD(y + i + 2 * ggml_f16_epr, 2); + ay3 = GGML_F16x_VEC_FMA(ay3, ax3, vx); + + GGML_F16x_VEC_STORE(y + i + 2 * ggml_f16_epr, ay3, 2); + + ax4 = GGML_F16x_VEC_LOAD(x + i + 3 * ggml_f16_epr, 3); + ay4 = GGML_F16x_VEC_LOAD(y + i + 3 * ggml_f16_epr, 3); + ay4 = GGML_F16x_VEC_FMA(ay4, ax4, vx); + + GGML_F16x_VEC_STORE(y + i + 3 * ggml_f16_epr, ay4, 3); + + ax5 = GGML_F16x_VEC_LOAD(x + i + 4 * ggml_f16_epr, 4); + ay5 = GGML_F16x_VEC_LOAD(y + i + 4 * ggml_f16_epr, 4); + ay5 = GGML_F16x_VEC_FMA(ay5, ax5, vx); + + GGML_F16x_VEC_STORE(y + i + 4 * ggml_f16_epr, ay5, 4); + + ax6 = GGML_F16x_VEC_LOAD(x + i + 5 * ggml_f16_epr, 5); + ay6 = GGML_F16x_VEC_LOAD(y + i + 5 * ggml_f16_epr, 5); + ay6 = GGML_F16x_VEC_FMA(ay6, ax6, vx); + + GGML_F16x_VEC_STORE(y + i + 5 * ggml_f16_epr, ay6, 5); + + ax7 = GGML_F16x_VEC_LOAD(x + i + 6 * ggml_f16_epr, 6); + ay7 = GGML_F16x_VEC_LOAD(y + i + 6 * ggml_f16_epr, 6); + ay7 = GGML_F16x_VEC_FMA(ay7, ax7, vx); + + GGML_F16x_VEC_STORE(y + i + 6 * ggml_f16_epr, ay7, 6); + + ax8 = GGML_F16x_VEC_LOAD(x + i + 7 * ggml_f16_epr, 7); + ay8 = GGML_F16x_VEC_LOAD(y + i + 7 * ggml_f16_epr, 7); + ay8 = GGML_F16x_VEC_FMA(ay8, ax8, vx); + + GGML_F16x_VEC_STORE(y + i + 7 * ggml_f16_epr, ay8, 7); + } + const int np2 = (n & ~(ggml_f16_epr - 1)); + for (int k = np; k < np2; k += ggml_f16_epr) { + svfloat16_t rx = GGML_F16x_VEC_LOAD(x + k, 0); + svfloat16_t ry = GGML_F16x_VEC_LOAD(y + k, 0); + ry = GGML_F16x_VEC_FMA(ry, rx, vx); + + GGML_F16x_VEC_STORE(y + k, ry, 0); + } + + if (np2 < n) { + svbool_t pg = svwhilelt_b16(np2, n); + svfloat16_t hx = svld1_f16(pg, (const __fp16 *)(x + np2)); + svfloat16_t hy = svld1_f16(pg, (const __fp16 *)(y + np2)); + hy = svmad_f16_x(pg, hx, vx, hy); + svst1_f16(pg, (__fp16 *)(y + np2), hy); + } + np = n; +#elif defined(__riscv_v_intrinsic) // implies __riscv_v_intrinsic + #if defined (__riscv_zvfh) + const ggml_fp16_t s = GGML_CPU_FP32_TO_FP16(v); + const _Float16 scale = *(const _Float16*)(&s); + + // calculate step size + const int epr = __riscv_vsetvlmax_e16m4(); + const int step = epr * 2; + int np = (n & ~(step - 1)); + + // unroll by 2 + for (int i = 0; i < np; i += step) { + vfloat16m4_t ax0 = __riscv_vle16_v_f16m4((const _Float16*)x + i, epr); + vfloat16m4_t ay0 = __riscv_vle16_v_f16m4((const _Float16*)y + i, epr); + ay0 = __riscv_vfmacc_vf_f16m4(ay0, scale, ax0, epr); + __riscv_vse16_v_f16m4((_Float16*)y + i, ay0, epr); + __asm__ __volatile__ ("" ::: "memory"); + + vfloat16m4_t ax1 = __riscv_vle16_v_f16m4((const _Float16*)x + i + epr, epr); + vfloat16m4_t ay1 = __riscv_vle16_v_f16m4((const _Float16*)y + i + epr, epr); + ay1 = __riscv_vfmacc_vf_f16m4(ay1, scale, ax1, epr); + __riscv_vse16_v_f16m4((_Float16*)y + i + epr, ay1, epr); + __asm__ __volatile__ ("" ::: "memory"); + } + + // leftovers + int vl; + for (int i = np; i < n; i += vl) { + vl = __riscv_vsetvl_e16m4(n - i); + vfloat16m4_t ax0 = __riscv_vle16_v_f16m4((const _Float16*)x + i, vl); + vfloat16m4_t ay0 = __riscv_vle16_v_f16m4((const _Float16*)y + i, vl); + ay0 = __riscv_vfmacc_vf_f16m4(ay0, scale, ax0, vl); + __riscv_vse16_v_f16m4((_Float16*)y + i, ay0, vl); + } + np = n; + #else + // fall to scalar path + const int np = 0; + #endif +#elif defined(GGML_SIMD) + const int np = (n & ~(GGML_F16_STEP - 1)); + + GGML_F16_VEC vx = GGML_F16_VEC_SET1(v); + + GGML_F16_VEC ax[GGML_F16_ARR]; + GGML_F16_VEC ay[GGML_F16_ARR]; + + for (int i = 0; i < np; i += GGML_F16_STEP) { + for (int j = 0; j < GGML_F16_ARR; j++) { + ax[j] = GGML_F16_VEC_LOAD(x + i + j*GGML_F16_EPR, j); + ay[j] = GGML_F16_VEC_LOAD(y + i + j*GGML_F16_EPR, j); + ay[j] = GGML_F16_VEC_FMA(ay[j], ax[j], vx); + + GGML_F16_VEC_STORE(y + i + j*GGML_F16_EPR, ay, j); + } + } +#else + // scalar path + const int np = 0; +#endif + + // scalar and leftovers + for (int i = np; i < n; ++i) { + y[i] = GGML_CPU_FP32_TO_FP16(GGML_CPU_FP16_TO_FP32(y[i]) + GGML_CPU_FP16_TO_FP32(x[i])*v); + } +} + +// xs and vs are byte strides of x and v +inline static void ggml_vec_mad_f32_unroll(const int n, const int xs, const int vs, float * GGML_RESTRICT y, const float * GGML_RESTRICT xv, const float * GGML_RESTRICT vv) { + + const float * GGML_RESTRICT x[GGML_VEC_MAD_UNROLL]; + const float * GGML_RESTRICT v[GGML_VEC_MAD_UNROLL]; + + for (int i = 0; i < GGML_VEC_MAD_UNROLL; ++i) { + x[i] = (const float *) ((const char *) xv + i*xs); + v[i] = (const float *) ((const char *) vv + i*vs); + } + +#if defined(GGML_SIMD) + #if defined(__ARM_FEATURE_SVE) + // scalar Route to scalar implementation //TODO: Write SVE code + for (int k = 0; k < GGML_VEC_MAD_UNROLL; ++k) { + for (int i = 0; i < n; ++i) { + y[i] += x[k][i]*v[k][0]; + } + } + #elif defined(__riscv_v_intrinsic) + for (int i = 0, avl; i < n; i += avl) { + avl = __riscv_vsetvl_e32m8(n - i); + vfloat32m8_t ay = __riscv_vle32_v_f32m8(&y[i], avl); + for (int k = 0; k < GGML_VEC_MAD_UNROLL; k++) { + vfloat32m8_t ax = __riscv_vle32_v_f32m8(&x[k][i], avl); + ay = __riscv_vfmadd_vf_f32m8(ax, v[k][0], ay, avl); + } + __riscv_vse32_v_f32m8(&y[i], ay, avl); + } + #else + const int np = (n & ~(GGML_F32_STEP - 1)); + + GGML_F32_VEC vx[GGML_VEC_MAD_UNROLL]; + + for (int k = 0; k < GGML_VEC_MAD_UNROLL; ++k) { + vx[k] = GGML_F32_VEC_SET1(v[k][0]); + } + + GGML_F32_VEC ax[GGML_VEC_MAD_UNROLL][GGML_F32_ARR]; + GGML_F32_VEC ay[GGML_F32_ARR]; + + for (int i = 0; i < np; i += GGML_F32_STEP) { + for (int j = 0; j < GGML_F32_ARR; j++) { + ay[j] = GGML_F32_VEC_LOAD(y + i + j*GGML_F32_EPR); + + for (int k = 0; k < GGML_VEC_MAD_UNROLL; ++k) { + ax[k][j] = GGML_F32_VEC_LOAD(x[k] + i + j*GGML_F32_EPR); + ay[j] = GGML_F32_VEC_FMA(ay[j], ax[k][j], vx[k]); + } + + GGML_F32_VEC_STORE(y + i + j*GGML_F32_EPR, ay[j]); + } + } + + // leftovers + for (int k = 0; k < GGML_VEC_MAD_UNROLL; ++k) { + for (int i = np; i < n; ++i) { + y[i] += x[k][i]*v[k][0]; + } + } + #endif +#else + // scalar + for (int k = 0; k < GGML_VEC_MAD_UNROLL; ++k) { + for (int i = 0; i < n; ++i) { + y[i] += x[k][i]*v[k][0]; + } + } +#endif +} + +inline static void ggml_vec_mad1_f32(const int n, float * y, const float * x, const float s, const float b) { +#if defined(GGML_USE_ACCELERATE) + vDSP_vsmsa(x, 1, &s, &b, y, 1, n); +#elif defined(GGML_SIMD) + #if defined(__ARM_FEATURE_SVE) + // scalar ; TODO: Write SVE code + for (int i = 0; i < n; ++i) { + y[i] = x[i]*s + b; + } + #elif defined(__riscv_v_intrinsic) + for (int i = 0, avl; i < n; i += avl) { + avl = __riscv_vsetvl_e32m8(n - i); + vfloat32m8_t ax = __riscv_vle32_v_f32m8(&x[i], avl); + vfloat32m8_t vb = __riscv_vfmv_v_f_f32m8(b, avl); + vfloat32m8_t ny = __riscv_vfmadd_vf_f32m8(ax, s, vb, avl); + __riscv_vse32_v_f32m8(&y[i], ny, avl); + } + #else + const int np = (n & ~(GGML_F32_STEP - 1)); + + GGML_F32_VEC vs = GGML_F32_VEC_SET1(s); + GGML_F32_VEC vb = GGML_F32_VEC_SET1(b); + + GGML_F32_VEC ay[GGML_F32_ARR]; + + for (int i = 0; i < np; i += GGML_F32_STEP) { + for (int j = 0; j < GGML_F32_ARR; j++) { + ay[j] = GGML_F32_VEC_LOAD(x + i + j*GGML_F32_EPR); + ay[j] = GGML_F32_VEC_FMA(vb, ay[j], vs); + + GGML_F32_VEC_STORE(y + i + j*GGML_F32_EPR, ay[j]); + } + } + + // leftovers + for (int i = np; i < n; ++i) { + y[i] = x[i]*s + b; + } + #endif +#else + // scalar + for (int i = 0; i < n; ++i) { + y[i] = x[i]*s + b; + } +#endif +} + +//inline static void ggml_vec_scale_f32(const int n, float * y, const float v) { for (int i = 0; i < n; ++i) y[i] *= v; } +inline static void ggml_vec_scale_f32(const int n, float * y, const float v) { +#if defined(GGML_USE_ACCELERATE) + vDSP_vsmul(y, 1, &v, y, 1, n); +#elif defined(GGML_SIMD) + #if defined(__ARM_FEATURE_SVE) + const int sve_register_length = ggml_cpu_get_sve_cnt() * 8; + const int ggml_f32_epr = sve_register_length / 32;//8;//svcntw(); // SVE128:4, SVE256:8, SVE512:16 + const int ggml_f32_step = 2 * ggml_f32_epr; + + GGML_F32_VEC vx = GGML_F32_VEC_SET1(v); + const int np = (n & ~(ggml_f32_step - 1)); + svfloat32_t ay1; + svfloat32_t ay2; + for (int i = 0; i < np; i += ggml_f32_step) { + ay1 = GGML_F32_VEC_LOAD(y + i); + ay1 = GGML_F32_VEC_MUL(ay1, vx); + GGML_F32_VEC_STORE(y + i, ay1); + + ay2 = GGML_F32_VEC_LOAD(y + i + 1*ggml_f32_epr); + ay2 = GGML_F32_VEC_MUL(ay2, vx); + GGML_F32_VEC_STORE(y + i + 1*ggml_f32_epr, ay2); + } + // leftovers + // maximum number of leftover elements will be less that ggml_f32_epr. Apply predicated svmad on available elements only + for (int i = np; i < n; i += ggml_f32_epr) { + svbool_t pg = svwhilelt_b32(i, n); + ay1 = svld1_f32(pg, y + i); + ay1 = svmul_f32_m(pg, ay1, vx); + svst1_f32(pg, y + i, ay1); + } + #elif defined(__riscv_v_intrinsic) + for (int i = 0, avl; i < n; i += avl) { + avl = __riscv_vsetvl_e32m8(n - i); + vfloat32m8_t ay = __riscv_vle32_v_f32m8(&y[i], avl); + vfloat32m8_t ny = __riscv_vfmul_vf_f32m8(ay, v, avl); + __riscv_vse32_v_f32m8(&y[i], ny, avl); + } + #else + const int np = (n & ~(GGML_F32_STEP - 1)); + + GGML_F32_VEC vx = GGML_F32_VEC_SET1(v); + + GGML_F32_VEC ay[GGML_F32_ARR]; + + for (int i = 0; i < np; i += GGML_F32_STEP) { + for (int j = 0; j < GGML_F32_ARR; j++) { + ay[j] = GGML_F32_VEC_LOAD(y + i + j*GGML_F32_EPR); + ay[j] = GGML_F32_VEC_MUL(ay[j], vx); + + GGML_F32_VEC_STORE(y + i + j*GGML_F32_EPR, ay[j]); + } + } + + // leftovers + for (int i = np; i < n; ++i) { + y[i] *= v; + } + #endif +#else + // scalar + for (int i = 0; i < n; ++i) { + y[i] *= v; + } +#endif +} + +inline static void ggml_vec_scale_f16(const int n, ggml_fp16_t * y, const float v) { +#if defined(GGML_SIMD) && defined(__ARM_FEATURE_SVE) + const int sve_register_length = svcntb() * 8; + const int ggml_f16_epr = sve_register_length / 16; + const int ggml_f16_step = 2 * ggml_f16_epr; + + GGML_F16x_VEC vx = GGML_F16x_VEC_SET1(v); + int np = (n & ~(ggml_f16_step - 1)); + svfloat16_t ay1, ay2; + + for (int i = 0; i < np; i += ggml_f16_step) { + ay1 = GGML_F16x_VEC_LOAD(y + i + 0*ggml_f16_epr, 0); + ay1 = GGML_F16x_VEC_MUL(ay1, vx); + GGML_F16x_VEC_STORE(y + i + 0*ggml_f16_epr, ay1, 0); + + ay2 = GGML_F16x_VEC_LOAD(y + i + 1*ggml_f16_epr, 1); + ay2 = GGML_F16x_VEC_MUL(ay2, vx); + GGML_F16x_VEC_STORE(y + i + 1*ggml_f16_epr, ay2, 1); + } + // leftovers + // maximum number of leftover elements will be less that ggmlF_16x_epr. Apply predicated svmad on available elements only + if (np < n) { + svbool_t pg = svwhilelt_b16(np, n); + svfloat16_t hy = svld1_f16(pg, (__fp16 *)(y + np)); + svfloat16_t out = svmul_f16_m(pg, hy, vx); + svst1_f16(pg, (__fp16 *)(y + np), out); + } + np = n; +#elif defined(__riscv_v_intrinsic) + #if defined(__riscv_zvfh) + const ggml_fp16_t s = GGML_CPU_FP32_TO_FP16(v); + const _Float16 scale = *(const _Float16*)(&s); + + // calculate step size + const int epr = __riscv_vsetvlmax_e16m4(); + const int step = epr * 2; + int np = (n & ~(step - 1)); + + // unroll by 2 + for (int i = 0; i < np; i += step) { + vfloat16m4_t ay0 = __riscv_vle16_v_f16m4((const _Float16*)y + i, epr); + ay0 = __riscv_vfmul_vf_f16m4(ay0, scale, epr); + __riscv_vse16_v_f16m4((_Float16*)y + i, ay0, epr); + __asm__ __volatile__ ("" ::: "memory"); + + vfloat16m4_t ay1 = __riscv_vle16_v_f16m4((const _Float16*)y + i + epr, epr); + ay1 = __riscv_vfmul_vf_f16m4(ay1, scale, epr); + __riscv_vse16_v_f16m4((_Float16*)y + i + epr, ay1, epr); + __asm__ __volatile__ ("" ::: "memory"); + } + + // leftovers + int vl; + for (int i = np; i < n; i += vl) { + vl = __riscv_vsetvl_e16m4(n - i); + vfloat16m4_t ay0 = __riscv_vle16_v_f16m4((const _Float16*)y + i, vl); + ay0 = __riscv_vfmul_vf_f16m4(ay0, scale, vl); + __riscv_vse16_v_f16m4((_Float16*)y + i, ay0, vl); + } + np = n; + #else + // fall to scalar path + const int np = 0; + #endif +#elif defined(GGML_SIMD) + const int np = (n & ~(GGML_F16_STEP - 1)); + + GGML_F16_VEC vx = GGML_F16_VEC_SET1(v); + + GGML_F16_VEC ay[GGML_F16_ARR]; + + for (int i = 0; i < np; i += GGML_F16_STEP) { + for (int j = 0; j < GGML_F16_ARR; j++) { + ay[j] = GGML_F16_VEC_LOAD(y + i + j*GGML_F16_EPR, j); + ay[j] = GGML_F16_VEC_MUL(ay[j], vx); + + GGML_F16_VEC_STORE(y + i + j*GGML_F16_EPR, ay, j); + } + } +#else + // scalar path + const int np = 0; +#endif + // scalar and leftovers + for (int i = np; i < n; ++i) { + y[i] = GGML_CPU_FP32_TO_FP16(GGML_CPU_FP16_TO_FP32(y[i])*v); + } +} + +inline static void ggml_vec_norm_f32 (const int n, float * s, const float * x) { ggml_vec_dot_f32(n, s, 0, x, 0, x, 0, 1); *s = sqrtf(*s); } +inline static void ggml_vec_sqr_f32 (const int n, float * y, const float * x) { for (int i = 0; i < n; ++i) y[i] = x[i]*x[i]; } +inline static void ggml_vec_sqr_f16 (const int n, ggml_fp16_t * y, const ggml_fp16_t * x) { + for (int i = 0; i < n; ++i) { + float v = GGML_CPU_FP16_TO_FP32(x[i]); + y[i] = GGML_CPU_FP32_TO_FP16(v*v); + } +} +inline static void ggml_vec_sqrt_f32 (const int n, float * y, const float * x) { for (int i = 0; i < n; ++i) y[i] = sqrtf(x[i]); } +inline static void ggml_vec_sqrt_f16 (const int n, ggml_fp16_t * y, const ggml_fp16_t * x) { + for (int i = 0; i < n; ++i) { + y[i] = GGML_CPU_FP32_TO_FP16(sqrtf(GGML_CPU_FP16_TO_FP32(x[i]))); + } +} +inline static void ggml_vec_log_f32 (const int n, float * y, const float * x) { for (int i = 0; i < n; ++i) y[i] = logf(x[i]); } +inline static void ggml_vec_log_f16 (const int n, ggml_fp16_t * y, const ggml_fp16_t * x) { + for (int i = 0; i < n; ++i) { + y[i] = GGML_CPU_FP32_TO_FP16(logf(GGML_CPU_FP16_TO_FP32(x[i]))); + } +} +inline static void ggml_vec_sin_f32 (const int n, float * y, const float * x) { for (int i = 0; i < n; ++i) y[i] = sinf(x[i]); } +inline static void ggml_vec_sin_f16 (const int n, ggml_fp16_t * y, const ggml_fp16_t * x) { + for (int i = 0; i < n; ++i) { + y[i] = GGML_CPU_FP32_TO_FP16(sinf(GGML_CPU_FP16_TO_FP32(x[i]))); + } +} +inline static void ggml_vec_cos_f32 (const int n, float * y, const float * x) { for (int i = 0; i < n; ++i) y[i] = cosf(x[i]); } +inline static void ggml_vec_cos_f16 (const int n, ggml_fp16_t * y, const ggml_fp16_t * x) { + for (int i = 0; i < n; ++i) { + y[i] = GGML_CPU_FP32_TO_FP16(cosf(GGML_CPU_FP16_TO_FP32(x[i]))); + } +} +inline static void ggml_vec_abs_f32 (const int n, float * y, const float * x) { for (int i = 0; i < n; ++i) y[i] = fabsf(x[i]); } +inline static void ggml_vec_abs_f16 (const int n, ggml_fp16_t * y, const ggml_fp16_t * x) { + for (int i = 0; i < n; ++i) { + y[i] = GGML_CPU_FP32_TO_FP16(fabsf(GGML_CPU_FP16_TO_FP32(x[i]))); + } +} +inline static void ggml_vec_sgn_f32 (const int n, float * y, const float * x) { for (int i = 0; i < n; ++i) y[i] = (x[i] > 0.f) ? 1.f : ((x[i] < 0.f) ? -1.f : 0.f); } +inline static void ggml_vec_sgn_f16 (const int n, ggml_fp16_t * y, const ggml_fp16_t * x) { + for (int i = 0; i < n; ++i) { + float v = GGML_CPU_FP16_TO_FP32(x[i]); + y[i] = GGML_CPU_FP32_TO_FP16((v > 0.f) ? 1.f : ((v < 0.f) ? -1.f : 0.f)); + } +} +inline static void ggml_vec_step_f32 (const int n, float * y, const float * x) { for (int i = 0; i < n; ++i) y[i] = (x[i] > 0.f) ? 1.f : 0.f; } +inline static void ggml_vec_step_f16 (const int n, ggml_fp16_t * y, const ggml_fp16_t * x) { + for (int i = 0; i < n; ++i) { + y[i] = GGML_CPU_FP32_TO_FP16((GGML_CPU_FP16_TO_FP32(x[i]) > 0.f) ? 1.f : 0.f); + } +} +inline static void ggml_vec_tanh_f32 (const int n, float * y, const float * x) { for (int i = 0; i < n; ++i) y[i] = tanhf(x[i]); } +inline static void ggml_vec_tanh_f16 (const int n, ggml_fp16_t * y, const ggml_fp16_t * x) { + for (int i = 0; i < n; ++i) { + y[i] = GGML_CPU_FP32_TO_FP16(tanhf(GGML_CPU_FP16_TO_FP32(x[i]))); + } +} +inline static void ggml_vec_elu_f32 (const int n, float * y, const float * x) { for (int i = 0; i < n; ++i) y[i] = (x[i] > 0.f) ? x[i] : expm1f(x[i]); } +inline static void ggml_vec_elu_f16 (const int n, ggml_fp16_t * y, const ggml_fp16_t * x) { + for (int i = 0; i < n; ++i) { + const float v = GGML_CPU_FP16_TO_FP32(x[i]); + y[i] = GGML_CPU_FP32_TO_FP16((v > 0.f) ? v : expm1f(v)); + } +} +inline static void ggml_vec_relu_f32 (const int n, float * y, const float * x) { for (int i = 0; i < n; ++i) y[i] = (x[i] > 0.f) ? x[i] : 0.f; } +inline static void ggml_vec_relu_f16 (const int n, ggml_fp16_t * y, const ggml_fp16_t * x) { + for (int i = 0; i < n; ++i) { + float v = GGML_CPU_FP16_TO_FP32(x[i]); + y[i] = GGML_CPU_FP32_TO_FP16((v > 0.f) ? v : 0.f); + } +} +inline static void ggml_vec_leaky_relu_f32 (const int n, float * y, const float * x, const float ns) { for (int i = 0; i < n; ++i) y[i] = ((x[i] > 0.f) ? x[i] : 0.f) + ns * ((x[i] < 0.0f) ? x[i] : 0.f); } +inline static void ggml_vec_leaky_relu_f16 (const int n, ggml_fp16_t * y, const ggml_fp16_t * x, const float ns) { + for (int i = 0; i < n; ++i) { + float v = GGML_CPU_FP16_TO_FP32(x[i]); + y[i] = GGML_CPU_FP32_TO_FP16(((v > 0.f) ? v : 0.f) + ns * ((v < 0.0f) ? v : 0.f)); + } +} +inline static void ggml_vec_sigmoid_f32 (const int n, float * y, const float * x) { for (int i = 0; i < n; ++i) y[i] = 1.f / (1.f + expf(-x[i])); } +inline static void ggml_vec_sigmoid_f16 (const int n, ggml_fp16_t * y, const ggml_fp16_t * x) { + for (int i = 0; i < n; ++i) { + y[i] = GGML_CPU_FP32_TO_FP16(1.f / (1.f + expf(-GGML_CPU_FP16_TO_FP32(x[i])))); + } +} +// TODO: optimize performance +inline static void ggml_vec_hardswish_f32 (const int n, float * y, const float * x) { for (int i = 0; i < n; ++i) y[i] = x[i] * fminf(1.0f, fmaxf(0.0f, (x[i] + 3.0f) / 6.0f)); } +inline static void ggml_vec_hardswish_f16 (const int n, ggml_fp16_t * y, const ggml_fp16_t * x) { + for (int i = 0; i < n; ++i) { + float v = GGML_CPU_FP16_TO_FP32(x[i]); + y[i] = GGML_CPU_FP32_TO_FP16(v * fminf(1.0f, fmaxf(0.0f, (v + 3.0f) / 6.0f))); + } +} +inline static void ggml_vec_hardsigmoid_f32 (const int n, float * y, const float * x) { for (int i = 0; i < n; ++i) y[i] = fminf(1.0f, fmaxf(0.0f, (x[i] + 3.0f) / 6.0f)); } +inline static void ggml_vec_hardsigmoid_f16 (const int n, ggml_fp16_t * y, const ggml_fp16_t * x) { + for (int i = 0; i < n; ++i) { + y[i] = GGML_CPU_FP32_TO_FP16(fminf(1.0f, fmaxf(0.0f, (GGML_CPU_FP16_TO_FP32(x[i]) + 3.0f) / 6.0f))); + } +} +inline static void ggml_vec_exp_f32 (const int n, float * y, const float * x) { for (int i = 0; i < n; ++i) y[i] = expf(x[i]); } +inline static void ggml_vec_exp_f16 (const int n, ggml_fp16_t * y, const ggml_fp16_t * x) { + for (int i = 0; i < n; ++i) { + y[i] = GGML_CPU_FP32_TO_FP16(expf(GGML_CPU_FP16_TO_FP32(x[i]))); + } +} + +static const float GELU_COEF_A = 0.044715f; +static const float GELU_QUICK_COEF = -1.702f; +static const float SQRT_2_OVER_PI = 0.79788456080286535587989211986876f; +static const float SQRT_2_INV = 0.70710678118654752440084436210484f; + +inline static float ggml_gelu_f32(float x) { + return 0.5f*x*(1.0f + tanhf(SQRT_2_OVER_PI*x*(1.0f + GELU_COEF_A*x*x))); +} + +inline static void ggml_vec_gelu_f16(const int n, ggml_fp16_t * y, const ggml_fp16_t * x) { + const uint16_t * i16 = (const uint16_t *) x; + for (int i = 0; i < n; ++i) { + y[i] = ggml_table_gelu_f16[i16[i]]; + } +} + +inline static void ggml_vec_gelu_bf16(const int n, ggml_bf16_t * y, const ggml_bf16_t * x) { + for (int i = 0; i < n; ++i) { + y[i] = GGML_FP32_TO_BF16(ggml_gelu_f32(GGML_BF16_TO_FP32(x[i]))); + } +} + +inline static void ggml_vec_gelu_erf_f16(const int n, ggml_fp16_t * y, const ggml_fp16_t * x) { + for (int i = 0; i < n; ++i) { + float xi = GGML_CPU_FP16_TO_FP32(x[i]); + float res = 0.5f*xi*(1.0f + erff(xi*SQRT_2_INV)); + y[i] = GGML_CPU_FP32_TO_FP16(res); + } +} + +inline static void ggml_vec_gelu_erf_bf16(const int n, ggml_bf16_t * y, const ggml_bf16_t * x) { + for (int i = 0; i < n; ++i) { + float xi = GGML_BF16_TO_FP32(x[i]); + y[i] = GGML_FP32_TO_BF16(0.5f*xi*(1.0f + erff(xi*SQRT_2_INV))); + } +} + +#ifdef GGML_GELU_FP16 +inline static void ggml_vec_gelu_f32(const int n, float * y, const float * x) { + uint16_t t; + for (int i = 0; i < n; ++i) { + if (x[i] <= -10.0f) { + y[i] = 0.0f; + } else if (x[i] >= 10.0f) { + y[i] = x[i]; + } else { + ggml_fp16_t fp16 = GGML_CPU_FP32_TO_FP16(x[i]); + memcpy(&t, &fp16, sizeof(uint16_t)); + y[i] = GGML_CPU_FP16_TO_FP32(ggml_table_gelu_f16[t]); + } + } +} +#else +inline static void ggml_vec_gelu_f32(const int n, float * y, const float * x) { + for (int i = 0; i < n; ++i) { + y[i] = ggml_gelu_f32(x[i]); + } +} +#endif + +inline static void ggml_vec_gelu_erf_f32(const int n, float * y, const float * x) { + for (int i = 0; i < n; ++i) { + float xi = x[i]; + y[i] = 0.5f*xi*(1.0f + erff(xi*SQRT_2_INV)); + } +} + +inline static float ggml_gelu_quick_f32(float x) { + return x*(1.0f/(1.0f+expf(GELU_QUICK_COEF*x))); +} + +inline static void ggml_vec_gelu_quick_f16(const int n, ggml_fp16_t * y, const ggml_fp16_t * x) { + const uint16_t * i16 = (const uint16_t *) x; + for (int i = 0; i < n; ++i) { + y[i] = ggml_table_gelu_quick_f16[i16[i]]; + } +} + +inline static void ggml_vec_gelu_quick_bf16(const int n, ggml_bf16_t * y, const ggml_bf16_t * x) { + for (int i = 0; i < n; ++i) { + y[i] = GGML_FP32_TO_BF16(ggml_gelu_quick_f32(GGML_BF16_TO_FP32(x[i]))); + } +} + +#ifdef GGML_GELU_QUICK_FP16 +inline static void ggml_vec_gelu_quick_f32(const int n, float * y, const float * x) { + uint16_t t; + for (int i = 0; i < n; ++i) { + ggml_fp16_t fp16 = GGML_CPU_FP32_TO_FP16(x[i]); + memcpy(&t, &fp16, sizeof(uint16_t)); + y[i] = GGML_CPU_FP16_TO_FP32(ggml_table_gelu_quick_f16[t]); + } +} +#else +inline static void ggml_vec_gelu_quick_f32(const int n, float * y, const float * x) { + for (int i = 0; i < n; ++i) { + y[i] = ggml_gelu_quick_f32(x[i]); + } +} +#endif + +// Sigmoid Linear Unit (SiLU) function +inline static float ggml_silu_f32(float x) { + return x/(1.0f + expf(-x)); +} +inline static ggml_fp16_t ggml_silu_f16(ggml_fp16_t x) { + float v = GGML_CPU_FP16_TO_FP32(x); + return GGML_CPU_FP32_TO_FP16(v/(1.0f + expf(-v))); +} + +#if __FINITE_MATH_ONLY__ +#error "some routines in ggml.c require non-finite math arithmetics -- pass -fno-finite-math-only to the compiler to fix" +#error "ref: https://github.com/ggml-org/llama.cpp/pull/7154#issuecomment-2143844461" +#endif + +/* Below function was borrowed from the GitHub repository: +https://github.com/openvinotoolkit/openvino/blob/master/src/plugins/intel_cpu/src/nodes/kernels/scaled_attn/common.hpp */ +#if defined(__ARM_FEATURE_SVE) && defined(__aarch64__) + inline static svfloat32_t exp_ps_sve(svbool_t pg, svfloat32_t src) { + // Constants + const svfloat32_t log2_e = svdup_n_f32(1.4426950409f); + const svfloat32_t ln2 = svdup_n_f32(0.6931473921f); + const svfloat32_t half_ln2_sq = svdup_n_f32(0.2413862043f); + const svuint32_t not_mask17 = svdup_n_u32(~((1u << 17) - 1)); + const svfloat32_t one = svdup_n_f32(1.0f); + const svfloat32_t inactive1 = svdup_n_f32(0.0f); + const svint32_t inactive2 = svdup_n_s32(0); + + // Algorithm starts here + svfloat32_t t0 = svmul_f32_m(pg, src, log2_e); // y = x * log2(e) + svfloat32_t t1 = svrintm_f32_m(inactive1, pg, t0); // rount to int (float) + svint32_t t2 = svcvt_s32_f32_m(inactive2, pg, t1); // n + + t1 = svsub_f32_m(pg, t0, t1); // a = y - floor(y) + t1 = svadd_f32_m(pg, t1, one); // b = a + 1 + + svuint32_t t3 = svlsr_n_u32_m(pg, svreinterpret_u32_f32(t1), 17); // v = b >> 17 (u32) + svfloat32_t t4 = svexpa_f32(t3); // c = fexpa(v) + t4 = svscale_f32_m(pg, t4, t2); // fexpa(v) * 2^(n) + + // and_(t2.d, t1.d, not_mask17.d) + svfloat32_t t5 = svreinterpret_f32_u32(svand_u32_m(pg, svreinterpret_u32_f32(t1), not_mask17)); + t5 = svsub_f32_m(pg, t1, t5); // z + t0 = svmla_f32_m(pg, ln2, t5, half_ln2_sq); // ln2 + half_ln2_sq * z + t0 = svmla_f32_m(pg, one, t5, t0); // 1 + (ln2 * z) + (half_ln2_sq * z * z) + t0 = svmul_f32_m(pg, t0, t4); // Final result + + return t0; + } +#endif + +#if defined(__ARM_FEATURE_SVE) && defined(__aarch64__) + +inline static svfloat32_t ggml_v_expf(svbool_t pg, svfloat32_t x) { + const svfloat32_t r = svdup_n_f32_x(pg, 0x1.8p23f); + const svfloat32_t z = svmla_n_f32_x(pg, r, x, 0x1.715476p+0f); + const svfloat32_t n = svsub_f32_x(pg, z, r); + const svfloat32_t b = svmls_n_f32_x(pg, svmls_n_f32_x(pg, x, n, 0x1.62e4p-1f), n, 0x1.7f7d1cp-20f); + const svuint32_t e = svlsl_n_u32_x(pg, svreinterpret_u32_f32(z), 23); + const svfloat32_t k = svreinterpret_f32_u32(svadd_u32_x(pg, e, svreinterpret_u32_f32(svdup_n_f32_x(pg, 1)))); + const svbool_t c = svacgt_n_f32(pg, n, 126); + const svfloat32_t u = svmul_f32_x(pg, b, b); + const svfloat32_t j = svmla_f32_x(pg, + svmul_n_f32_x(pg, b, 0x1.ffffecp-1f), + svmla_f32_x(pg, svmla_f32_x(pg, svdup_n_f32_x(pg, 0x1.fffdb6p-2f), svdup_n_f32_x(pg, 0x1.555e66p-3f), b), + svmla_f32_x(pg, svdup_n_f32_x(pg, 0x1.573e2ep-5f), svdup_n_f32_x(pg, 0x1.0e4020p-7f), b), u), u); + const svuint32_t d = svdup_n_u32_z(svcmple_n_f32(pg, n, 0.0), 0x82000000); + const svfloat32_t s1 = svreinterpret_f32_u32(svadd_n_u32_x(pg, d, 0x7f000000)); + const svfloat32_t s2 = svreinterpret_f32_u32(svsub_u32_x(pg, e, d)); + return svsel_f32(svacgt_f32(pg, n, svdup_n_f32_x(pg, 192)), svmul_f32_x(pg, s1, s1), + svsel_f32(c, svmul_f32_x(pg, svmla_f32_x(pg, s2, s2, j), s1), svmla_f32_x(pg, k, k, j))); +} + +// computes silu x/(1+exp(-x)) in single precision vector +inline static svfloat32_t ggml_v_silu(svbool_t pg, svfloat32_t x) { + const svfloat32_t one = svdup_n_f32_x(pg, 1.0f); + const svfloat32_t zero = svdup_n_f32_x(pg, 0.0f); + const svfloat32_t neg_x = svsub_f32_x(pg, zero, x); + const svfloat32_t exp_neg_x = ggml_v_expf(pg, neg_x); + const svfloat32_t one_plus_exp_neg_x = svadd_f32_x(pg, one, exp_neg_x); + return svdiv_f32_x(pg, x, one_plus_exp_neg_x); +} + +#elif defined(__ARM_NEON) && defined(__aarch64__) + +// adapted from arm limited optimized routine +// the maximum error is 1.45358 plus 0.5 ulps +// numbers above 88.38 will flush to infinity +// numbers beneath -103.97 will flush to zero +inline static float32x4_t ggml_v_expf(float32x4_t x) { + const float32x4_t r = vdupq_n_f32(0x1.8p23f); + const float32x4_t z = vfmaq_f32(r, x, vdupq_n_f32(0x1.715476p+0f)); + const float32x4_t n = vsubq_f32(z, r); + const float32x4_t b = vfmsq_f32(vfmsq_f32(x, n, vdupq_n_f32(0x1.62e4p-1f)), n, + vdupq_n_f32(0x1.7f7d1cp-20f)); + const uint32x4_t e = vshlq_n_u32(vreinterpretq_u32_f32(z), 23); + const float32x4_t k = vreinterpretq_f32_u32(vaddq_u32(e, vreinterpretq_u32_f32(vdupq_n_f32(1)))); + const uint32x4_t c = vcagtq_f32(n, vdupq_n_f32(126)); + const float32x4_t u = vmulq_f32(b, b); + const float32x4_t j = vfmaq_f32( + vmulq_f32(vdupq_n_f32(0x1.ffffecp-1f), b), + vfmaq_f32(vfmaq_f32(vdupq_n_f32(0x1.fffdb6p-2f), vdupq_n_f32(0x1.555e66p-3f), b), + vfmaq_f32(vdupq_n_f32(0x1.573e2ep-5f), vdupq_n_f32(0x1.0e4020p-7f), b), u), u); + if (!vpaddd_u64(vreinterpretq_u64_u32(c))) + return vfmaq_f32(k, j, k); + const uint32x4_t d = vandq_u32(vclezq_f32(n), vdupq_n_u32(0x82000000)); + const float32x4_t s1 = vreinterpretq_f32_u32(vaddq_u32(d, vdupq_n_u32(0x7f000000))); + const float32x4_t s2 = vreinterpretq_f32_u32(vsubq_u32(e, d)); + return vbslq_f32(vcagtq_f32(n, vdupq_n_f32(192)), vmulq_f32(s1, s1), + vbslq_f32(c, vmulq_f32(vfmaq_f32(s2, s2, j), s1), vfmaq_f32(k, k, j))); +} + +// computes silu x/(1+exp(-x)) in single precision vector +inline static float32x4_t ggml_v_silu(float32x4_t x) { + const float32x4_t one = vdupq_n_f32(1.0f); + const float32x4_t zero = vdupq_n_f32(0.0f); + const float32x4_t neg_x = vsubq_f32(zero, x); + const float32x4_t exp_neg_x = ggml_v_expf(neg_x); + const float32x4_t one_plus_exp_neg_x = vaddq_f32(one, exp_neg_x); + return vdivq_f32(x, one_plus_exp_neg_x); +} + +#elif defined(__AVX512F__) && defined(__AVX512DQ__) + +// adapted from arm limited optimized routine +// the maximum error is 1.45358 plus 0.5 ulps +// numbers above 88.38 will flush to infinity +// numbers beneath -103.97 will flush to zero +inline static __m512 ggml_v_expf(__m512 x) { + const __m512 r = _mm512_set1_ps(0x1.8p23f); + const __m512 z = _mm512_fmadd_ps(x, _mm512_set1_ps(0x1.715476p+0f), r); + const __m512 n = _mm512_sub_ps(z, r); + const __m512 b = + _mm512_fnmadd_ps(n, _mm512_set1_ps(0x1.7f7d1cp-20f), + _mm512_fnmadd_ps(n, _mm512_set1_ps(0x1.62e4p-1f), x)); + const __mmask16 d = + _mm512_cmp_ps_mask(_mm512_abs_ps(n), _mm512_set1_ps(192), _CMP_GT_OQ); + const __m512 u = _mm512_mul_ps(b, b); + const __m512 j = _mm512_fmadd_ps( + _mm512_fmadd_ps(_mm512_fmadd_ps(_mm512_set1_ps(0x1.0e4020p-7f), b, + _mm512_set1_ps(0x1.573e2ep-5f)), + u, + _mm512_fmadd_ps(_mm512_set1_ps(0x1.555e66p-3f), b, + _mm512_set1_ps(0x1.fffdb6p-2f))), + u, + _mm512_fmadd_ps(_mm512_set1_ps(0x1.ffffecp-1f), b, _mm512_set1_ps(1.0F))); + const __m512 res = _mm512_scalef_ps(j, n); + if (_mm512_kortestz(d, d)) + return res; + const __m512 zero = _mm512_setzero_ps(); + const __m512 alt = _mm512_mask_blend_ps( + _mm512_cmp_ps_mask(n, zero, _CMP_LE_OQ), _mm512_set1_ps(INFINITY), zero); + return _mm512_mask_blend_ps(d, res, alt); +} + +// computes silu x/(1+exp(-x)) in single precision vector +inline static __m512 ggml_v_silu(__m512 x) { + const __m512 one = _mm512_set1_ps(1); + const __m512 zero = _mm512_setzero_ps(); + const __m512 neg_x = _mm512_sub_ps(zero, x); + const __m512 exp_neg_x = ggml_v_expf(neg_x); + const __m512 one_plus_exp_neg_x = _mm512_add_ps(one, exp_neg_x); + return _mm512_div_ps(x, one_plus_exp_neg_x); +} + +#elif defined(__AVX2__) && defined(__FMA__) + +// adapted from arm limited optimized routine +// the maximum error is 1.45358 plus 0.5 ulps +// numbers above 88.38 will flush to infinity +// numbers beneath -103.97 will flush to zero +inline static __m256 ggml_v_expf(__m256 x) { + const __m256 r = _mm256_set1_ps(0x1.8p23f); + const __m256 z = _mm256_fmadd_ps(x, _mm256_set1_ps(0x1.715476p+0f), r); + const __m256 n = _mm256_sub_ps(z, r); + const __m256 b = _mm256_fnmadd_ps(n, _mm256_set1_ps(0x1.7f7d1cp-20f), + _mm256_fnmadd_ps(n, _mm256_set1_ps(0x1.62e4p-1f), x)); + const __m256i e = _mm256_slli_epi32(_mm256_castps_si256(z), 23); + const __m256 k = _mm256_castsi256_ps( + _mm256_add_epi32(e, _mm256_castps_si256(_mm256_set1_ps(1)))); + const __m256i c = _mm256_castps_si256( + _mm256_cmp_ps(_mm256_andnot_ps(_mm256_set1_ps(-0.f), n), + _mm256_set1_ps(126), _CMP_GT_OQ)); + const __m256 u = _mm256_mul_ps(b, b); + const __m256 j = _mm256_fmadd_ps(_mm256_fmadd_ps(_mm256_fmadd_ps(_mm256_set1_ps(0x1.0e4020p-7f), b, + _mm256_set1_ps(0x1.573e2ep-5f)), u, + _mm256_fmadd_ps(_mm256_set1_ps(0x1.555e66p-3f), b, + _mm256_set1_ps(0x1.fffdb6p-2f))), + u, _mm256_mul_ps(_mm256_set1_ps(0x1.ffffecp-1f), b)); + if (!_mm256_movemask_ps(_mm256_castsi256_ps(c))) + return _mm256_fmadd_ps(j, k, k); + const __m256i g = _mm256_and_si256( + _mm256_castps_si256(_mm256_cmp_ps(n, _mm256_setzero_ps(), _CMP_LE_OQ)), + _mm256_set1_epi32(0x82000000u)); + const __m256 s1 = + _mm256_castsi256_ps(_mm256_add_epi32(g, _mm256_set1_epi32(0x7f000000u))); + const __m256 s2 = _mm256_castsi256_ps(_mm256_sub_epi32(e, g)); + const __m256i d = _mm256_castps_si256( + _mm256_cmp_ps(_mm256_andnot_ps(_mm256_set1_ps(-0.f), n), + _mm256_set1_ps(192), _CMP_GT_OQ)); + return _mm256_or_ps( + _mm256_and_ps(_mm256_castsi256_ps(d), _mm256_mul_ps(s1, s1)), + _mm256_andnot_ps( + _mm256_castsi256_ps(d), + _mm256_or_ps( + _mm256_and_ps(_mm256_castsi256_ps(c), + _mm256_mul_ps(_mm256_fmadd_ps(s2, j, s2), s1)), + _mm256_andnot_ps(_mm256_castsi256_ps(c), _mm256_fmadd_ps(k, j, k))))); +} + +// computes silu x/(1+exp(-x)) in single precision vector +inline static __m256 ggml_v_silu(__m256 x) { + const __m256 one = _mm256_set1_ps(1); + const __m256 zero = _mm256_setzero_ps(); + const __m256 neg_x = _mm256_sub_ps(zero, x); + const __m256 exp_neg_x = ggml_v_expf(neg_x); + const __m256 one_plus_exp_neg_x = _mm256_add_ps(one, exp_neg_x); + return _mm256_div_ps(x, one_plus_exp_neg_x); +} + +#elif defined(__SSE2__) // __AVX2__ / __ARM_NEON + +#if defined(__FMA__) +#define MADD128(x, y, z) _mm_fmadd_ps(x, y, z) +#define NMADD128(x, y, z) _mm_fnmadd_ps(x, y, z) +#else +#define MADD128(x, y, z) _mm_add_ps(_mm_mul_ps(x, y), z) +#define NMADD128(x, y, z) _mm_sub_ps(z, _mm_mul_ps(x, y)) +#endif + +// adapted from arm limited optimized routine +// the maximum error is 1.45358 plus 0.5 ulps +// numbers above 88.38 will flush to infinity +// numbers beneath -103.97 will flush to zero +inline static __m128 ggml_v_expf(__m128 x) { + const __m128 r = _mm_set1_ps(0x1.8p23f); + const __m128 z = MADD128(x, _mm_set1_ps(0x1.715476p+0f), r); + const __m128 n = _mm_sub_ps(z, r); + const __m128 b = + NMADD128(n, _mm_set1_ps(0x1.7f7d1cp-20f), NMADD128(n, _mm_set1_ps(0x1.62e4p-1f), x)); + const __m128i e = _mm_slli_epi32(_mm_castps_si128(z), 23); + const __m128 k = _mm_castsi128_ps(_mm_add_epi32(e, _mm_castps_si128(_mm_set1_ps(1)))); + const __m128i c = + _mm_castps_si128(_mm_cmpgt_ps(_mm_andnot_ps(_mm_set1_ps(-0.f), n), _mm_set1_ps(126))); + const __m128 u = _mm_mul_ps(b, b); + const __m128 j = + MADD128(MADD128(MADD128(_mm_set1_ps(0x1.0e4020p-7f), b, _mm_set1_ps(0x1.573e2ep-5f)), u, + MADD128(_mm_set1_ps(0x1.555e66p-3f), b, _mm_set1_ps(0x1.fffdb6p-2f))), + u, _mm_mul_ps(_mm_set1_ps(0x1.ffffecp-1f), b)); + if (!_mm_movemask_epi8(c)) + return MADD128(j, k, k); + const __m128i g = _mm_and_si128(_mm_castps_si128(_mm_cmple_ps(n, _mm_setzero_ps())), + _mm_set1_epi32(0x82000000u)); + const __m128 s1 = _mm_castsi128_ps(_mm_add_epi32(g, _mm_set1_epi32(0x7f000000u))); + const __m128 s2 = _mm_castsi128_ps(_mm_sub_epi32(e, g)); + const __m128i d = + _mm_castps_si128(_mm_cmpgt_ps(_mm_andnot_ps(_mm_set1_ps(-0.f), n), _mm_set1_ps(192))); + return _mm_or_ps( + _mm_and_ps(_mm_castsi128_ps(d), _mm_mul_ps(s1, s1)), + _mm_andnot_ps(_mm_castsi128_ps(d), + _mm_or_ps(_mm_and_ps(_mm_castsi128_ps(c), _mm_mul_ps(MADD128(s2, j, s2), s1)), + _mm_andnot_ps(_mm_castsi128_ps(c), MADD128(k, j, k))))); +} + +// computes silu x/(1+exp(-x)) in single precision vector +inline static __m128 ggml_v_silu(__m128 x) { + const __m128 one = _mm_set1_ps(1); + const __m128 zero = _mm_setzero_ps(); + const __m128 neg_x = _mm_sub_ps(zero, x); + const __m128 exp_neg_x = ggml_v_expf(neg_x); + const __m128 one_plus_exp_neg_x = _mm_add_ps(one, exp_neg_x); + return _mm_div_ps(x, one_plus_exp_neg_x); +} + +#elif defined(__riscv_v_intrinsic) + +// adapted from arm limited optimized routine +// the maximum error is 1.45358 plus 0.5 ulps +// numbers above 88.38 will flush to infinity +// numbers beneath -103.97 will flush to zero +inline static vfloat32m2_t ggml_v_expf_m2(vfloat32m2_t x, int vl) { + const vfloat32m2_t r = __riscv_vfmv_v_f_f32m2(0x1.8p23f, vl); +#ifdef __riscv_xtheadvector + // workaround for compiler bug (gcc 14.3.0: Error: unrecognized opcode `th.vmv1r.v v2,v4') + vfloat32m2_t z = __riscv_vfadd_vf_f32m2(r, 0.0f, vl); + z = __riscv_vfmacc_vf_f32m2(z, 0x1.715476p+0f, x, vl); +#else + const vfloat32m2_t z = __riscv_vfmacc_vf_f32m2(r, 0x1.715476p+0f, x, vl); +#endif + const vfloat32m2_t n = __riscv_vfsub_vv_f32m2(z, r, vl); + const vfloat32m2_t b = __riscv_vfnmsac_vf_f32m2(__riscv_vfnmsac_vf_f32m2(x, 0x1.62e4p-1f, n, vl), + 0x1.7f7d1cp-20f, n, vl); + const vuint32m2_t e = __riscv_vsll_vx_u32m2(__riscv_vreinterpret_v_f32m2_u32m2(z), 23, vl); + const vfloat32m2_t k = __riscv_vreinterpret_v_u32m2_f32m2(__riscv_vadd_vx_u32m2(e, 0x3f800000, vl)); // 1.0f + const vbool16_t c = __riscv_vmfgt_vf_f32m2_b16(__riscv_vfabs_v_f32m2(n, vl), 126.0f, vl); + const vfloat32m2_t u = __riscv_vfmul_vv_f32m2(b, b, vl); + const vfloat32m2_t j = __riscv_vfmacc_vv_f32m2( + __riscv_vfmul_vf_f32m2(b, 0x1.ffffecp-1f, vl), + __riscv_vfmacc_vv_f32m2( + __riscv_vfmacc_vf_f32m2(__riscv_vfmv_v_f_f32m2(0x1.fffdb6p-2f, vl), 0x1.555e66p-3f, b, vl), + __riscv_vfmacc_vf_f32m2(__riscv_vfmv_v_f_f32m2(0x1.573e2ep-5f, vl), 0x1.0e4020p-7f, b, vl), + u, vl), u, vl); + if (!__riscv_vcpop_m_b16(c, vl)) + return __riscv_vfmacc_vv_f32m2(k, j, k, vl); + const vbool16_t dm = __riscv_vmfle_vf_f32m2_b16(n, 0.0f, vl); + const vuint32m2_t d = __riscv_vmerge_vxm_u32m2(__riscv_vmv_v_x_u32m2(0, vl), 0x82000000, dm, vl); + const vfloat32m2_t s1 = __riscv_vreinterpret_v_u32m2_f32m2(__riscv_vadd_vx_u32m2(d, 0x7f000000, vl)); + const vfloat32m2_t s2 = __riscv_vreinterpret_v_u32m2_f32m2(__riscv_vsub_vv_u32m2(e, d, vl)); + const vfloat32m2_t r1 = __riscv_vmerge_vvm_f32m2( + __riscv_vfmacc_vv_f32m2(k, k, j, vl), + __riscv_vfmul_vv_f32m2(__riscv_vfmacc_vv_f32m2(s2, s2, j, vl), s1, vl), + c, vl); + return __riscv_vmerge_vvm_f32m2( + r1, __riscv_vfmul_vv_f32m2(s1, s1, vl), + __riscv_vmfgt_vf_f32m2_b16(__riscv_vfabs_v_f32m2(n, vl), 192.0f, vl), + vl); +} + +// computes silu x/(1+exp(-x)) in single precision vector +inline static vfloat32m2_t ggml_v_silu_m2(vfloat32m2_t x, int vl) { + const vfloat32m2_t neg_x = __riscv_vfneg_v_f32m2(x, vl); + const vfloat32m2_t exp_neg_x = ggml_v_expf_m2(neg_x, vl); + const vfloat32m2_t one_plus_exp_neg_x = __riscv_vfadd_vf_f32m2(exp_neg_x, 1.0f, vl); + return __riscv_vfdiv_vv_f32m2(x, one_plus_exp_neg_x, vl); +} + +#endif // __ARM_NEON / __AVX2__ / __SSE2__ / __riscv_v_intrinsic + +inline static void ggml_vec_silu_f16(const int n, ggml_fp16_t * y, const ggml_fp16_t * x) { + for (int i = 0; i < n; ++i) { + y[i] = ggml_silu_f16(x[i]); + } +} + +inline static void ggml_vec_silu_bf16(const int n, ggml_bf16_t * y, const ggml_bf16_t * x) { + for (int i = 0; i < n; ++i) { + y[i] = GGML_FP32_TO_BF16(ggml_silu_f32(GGML_BF16_TO_FP32(x[i]))); + } +} + +inline static float ggml_silu_backward_f32(float x, float dy) { + const float s = 1.0f/(1.0f + expf(-x)); + return dy*s*(1.0f + x*(1.0f - s)); +} + +inline static ggml_fp16_t ggml_silu_backward_f16(ggml_fp16_t x, ggml_fp16_t dy) { + const float v = GGML_CPU_FP16_TO_FP32(x); + const float s = 1.0f/(1.0f + expf(-v)); + return GGML_CPU_FP32_TO_FP16(GGML_CPU_FP16_TO_FP32(dy)*s*(1.0f + v*(1.0f - s))); +} + +inline static void ggml_vec_silu_backward_f32(const int n, float * dx, const float * x, const float * dy) { + for (int i = 0; i < n; ++i) { + dx[i] = ggml_silu_backward_f32(x[i], dy[i]); + } +} + +inline static void ggml_vec_silu_backward_f16(const int n, ggml_fp16_t * dx, const ggml_fp16_t * x, const ggml_fp16_t * dy) { + for (int i = 0; i < n; ++i) { + dx[i] = ggml_silu_backward_f16(x[i], dy[i]); + } +} + +inline static void ggml_vec_reglu_f32 (const int n, float * y, const float * x, const float * g) { + for (int i = 0; i < n; ++i) { + y[i] = (x[i] > 0.f) ? x[i] * g[i] : 0.f; + } +} + +inline static void ggml_vec_reglu_f16 (const int n, ggml_fp16_t * y, const ggml_fp16_t * x, const ggml_fp16_t * g) { + for (int i = 0; i < n; ++i) { + float v = GGML_CPU_FP16_TO_FP32(x[i]); + y[i] = GGML_CPU_FP32_TO_FP16((v > 0.f) ? v * GGML_CPU_FP16_TO_FP32(g[i]) : 0.f); + } +} + +inline static void ggml_vec_reglu_bf16(const int n, ggml_bf16_t * y, const ggml_bf16_t * x, const ggml_bf16_t * g) { + for (int i = 0; i < n; ++i) { + float v = GGML_BF16_TO_FP32(x[i]); + y[i] = GGML_FP32_TO_BF16((v > 0.f) ? v * GGML_BF16_TO_FP32(g[i]) : 0.f); + } +} + +#ifdef GGML_GELU_FP16 +inline static void ggml_vec_geglu_f32(const int n, float * y, const float * x, const float * g) { + uint16_t t; + for (int i = 0; i < n; ++i) { + if (x[i] <= -10.0f) { + y[i] = 0.0f; + } else if (x[i] >= 10.0f) { + y[i] = x[i] * g[i]; + } else { + ggml_fp16_t fp16 = GGML_CPU_FP32_TO_FP16(x[i]); + memcpy(&t, &fp16, sizeof(uint16_t)); + y[i] = GGML_CPU_FP16_TO_FP32(ggml_table_gelu_f16[t]) * g[i]; + } + } +} +#else +inline static void ggml_vec_geglu_f32(const int n, float * y, const float * x, const float * g) { + for (int i = 0; i < n; ++i) { + y[i] = ggml_gelu_f32(x[i]) * g[i]; + } +} +#endif + +inline static void ggml_vec_geglu_f16(const int n, ggml_fp16_t * y, const ggml_fp16_t * x, const ggml_fp16_t * g) { + const uint16_t * i16 = (const uint16_t *) x; + for (int i = 0; i < n; ++i) { + float v = GGML_CPU_FP16_TO_FP32(g[i]); + y[i] = GGML_CPU_FP32_TO_FP16(GGML_CPU_FP16_TO_FP32(ggml_table_gelu_f16[i16[i]]) * v); + } +} + +inline static void ggml_vec_geglu_bf16(const int n, ggml_bf16_t * y, const ggml_bf16_t * x, const ggml_bf16_t * g) { + for (int i = 0; i < n; ++i) { + y[i] = GGML_FP32_TO_BF16(ggml_gelu_f32(GGML_BF16_TO_FP32(x[i])) * GGML_BF16_TO_FP32(g[i])); + } +} + +void ggml_vec_swiglu_f32(const int n, float * y, const float * x, const float * g); + +inline static void ggml_vec_swiglu_f16(const int n, ggml_fp16_t * y, const ggml_fp16_t * x, const ggml_fp16_t * g) { + for (int i = 0; i < n; ++i) { + float xi = GGML_CPU_FP16_TO_FP32(x[i]); + float gi = GGML_CPU_FP16_TO_FP32(g[i]); + y[i] = GGML_CPU_FP32_TO_FP16((xi/(1.0f + expf(-xi))) * gi); + } +} + +inline static void ggml_vec_swiglu_bf16(const int n, ggml_bf16_t * y, const ggml_bf16_t * x, const ggml_bf16_t * g) { + for (int i = 0; i < n; ++i) { + float xi = GGML_BF16_TO_FP32(x[i]); + float gi = GGML_BF16_TO_FP32(g[i]); + y[i] = GGML_FP32_TO_BF16((xi/(1.0f + expf(-xi))) * gi); + } +} + +inline static void ggml_vec_geglu_erf_f32(const int n, float * y, const float * x, const float * g) { + for (int i = 0; i < n; ++i) { + float xi = x[i]; + y[i] = 0.5f * xi * (1.0f + erff(xi*SQRT_2_INV)) * g[i]; + } +} + +inline static void ggml_vec_geglu_erf_f16(const int n, ggml_fp16_t * y, const ggml_fp16_t * x, const ggml_fp16_t * g) { + for (int i = 0; i < n; ++i) { + float xi = GGML_CPU_FP16_TO_FP32(x[i]); + float gi = GGML_CPU_FP16_TO_FP32(g[i]); + y[i] = GGML_CPU_FP32_TO_FP16(0.5f * xi * (1.0f + erff(xi*SQRT_2_INV)) * gi); + } +} + +inline static void ggml_vec_geglu_erf_bf16(const int n, ggml_bf16_t * y, const ggml_bf16_t * x, const ggml_bf16_t * g) { + for (int i = 0; i < n; ++i) { + float xi = GGML_BF16_TO_FP32(x[i]); + float gi = GGML_BF16_TO_FP32(g[i]); + y[i] = GGML_FP32_TO_BF16(0.5f * xi * (1.0f + erff(xi*SQRT_2_INV)) * gi); + } +} + +#ifdef GGML_GELU_QUICK_FP16 +inline static void ggml_vec_geglu_quick_f32(const int n, float * y, const float * x, const float * g) { + uint16_t t; + for (int i = 0; i < n; ++i) { + ggml_fp16_t fp16 = GGML_CPU_FP32_TO_FP16(x[i]); + memcpy(&t, &fp16, sizeof(uint16_t)); + y[i] = GGML_CPU_FP16_TO_FP32(ggml_table_gelu_quick_f16[t]) * g[i]; + } +} +#else +inline static void ggml_vec_geglu_quick_f32(const int n, float * y, const float * x, const float * g) { + for (int i = 0; i < n; ++i) { + y[i] = ggml_gelu_quick_f32(x[i]) * g[i]; + } +} +#endif + +inline static void ggml_vec_geglu_quick_f16(const int n, ggml_fp16_t * y, const ggml_fp16_t * x, const ggml_fp16_t * g) { + const uint16_t * i16 = (const uint16_t *) x; + for (int i = 0; i < n; ++i) { + float v = GGML_CPU_FP16_TO_FP32(g[i]); + y[i] = GGML_CPU_FP32_TO_FP16(GGML_CPU_FP16_TO_FP32(ggml_table_gelu_quick_f16[i16[i]]) * v); + } +} + +inline static void ggml_vec_geglu_quick_bf16(const int n, ggml_bf16_t * y, const ggml_bf16_t * x, const ggml_bf16_t * g) { + for (int i = 0; i < n; ++i) { + y[i] = GGML_FP32_TO_BF16(ggml_gelu_quick_f32(GGML_BF16_TO_FP32(x[i])) * GGML_BF16_TO_FP32(g[i])); + } +} + +inline static void ggml_vec_sum_f32(const int n, float * s, const float * x) { +#ifndef GGML_USE_ACCELERATE + ggml_float sum = 0.0; + for (int i = 0; i < n; ++i) { + sum += (ggml_float)x[i]; + } + *s = (float)sum; +#else + vDSP_sve(x, 1, s, n); +#endif +} + +inline static void ggml_vec_cumsum_f32(const int n, float * y, const float * x) { + for (int i = 0; i < n; ++i) { + if (i == 0) { + y[i] = x[i]; + } else { + y[i] = y[i - 1] + x[i]; + } + } +} + +inline static void ggml_vec_sum_f32_ggf(const int n, ggml_float * s, const float * x) { + ggml_float sum = 0.0; + for (int i = 0; i < n; ++i) { + sum += (ggml_float)x[i]; + } + *s = sum; +} + +inline static void ggml_vec_sum_f16_ggf(const int n, float * s, const ggml_fp16_t * x) { + float sum = 0.0f; + for (int i = 0; i < n; ++i) { + sum += GGML_CPU_FP16_TO_FP32(x[i]); + } + *s = sum; +} + +inline static void ggml_vec_sum_bf16_ggf(const int n, float * s, const ggml_bf16_t * x) { + float sum = 0.0f; + for (int i = 0; i < n; ++i) { + sum += GGML_BF16_TO_FP32(x[i]); + } + *s = sum; +} + +inline static void ggml_vec_max_f32(const int n, float * s, const float * x) { +#ifndef GGML_USE_ACCELERATE + float max = -INFINITY; + for (int i = 0; i < n; ++i) { + max = MAX(max, x[i]); + } + *s = max; +#else + vDSP_maxv(x, 1, s, n); +#endif +} + +inline static void ggml_vec_norm_inv_f32(const int n, float * s, const float * x) { + ggml_vec_norm_f32(n, s, x); + *s = 1.f/(*s); +} + +inline static void ggml_vec_argmax_f32(const int n, int * s, const float * x) { + float max = -INFINITY; + int idx = 0; + for (int i = 0; i < n; ++i) { + max = MAX(max, x[i]); + if (max == x[i]) { idx = i; } + } + *s = idx; +} + +#ifdef __cplusplus +} +#endif diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/common.cuh b/python/freetoken/kernel/csrc/ggml/ggml-cuda/common.cuh new file mode 100644 index 0000000000..a767c688f5 --- /dev/null +++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/common.cuh @@ -0,0 +1,1725 @@ +#pragma once + +#include "ggml.h" +#include "ggml-impl.h" +#include "ggml-cuda.h" + +#include +#include +#include +#include + +#if defined(GGML_USE_HIP) +#define GGML_COMMON_DECL_HIP +#define GGML_COMMON_IMPL_HIP +#else +#define GGML_COMMON_DECL_CUDA +#define GGML_COMMON_IMPL_CUDA +#if defined(GGML_USE_MUSA) +#define GGML_COMMON_DECL_MUSA +#define GGML_COMMON_IMPL_MUSA +#endif +#endif +#include "ggml-common.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#if defined(GGML_USE_HIP) +#include "vendors/hip.h" +#elif defined(GGML_USE_MUSA) +#include "vendors/musa.h" +#else +#include "vendors/cuda.h" +#endif // defined(GGML_USE_HIP) + +#define STRINGIZE_IMPL(...) #__VA_ARGS__ +#define STRINGIZE(...) STRINGIZE_IMPL(__VA_ARGS__) + +#define WARP_SIZE 32 +#define CUDART_HMAX 11070 // CUDA 11.7, min. ver. for which __hmax and __hmax2 are known to work (may be higher than needed) +#define CUDART_HMASK 12000 // CUDA 12.0, min. ver. for half2 -> uint mask comparisons + +#define GGML_CUDA_CC_PASCAL 600 +#define GGML_CUDA_CC_DP4A 610 // minimum compute capability for __dp4a, an intrinsic for byte-wise dot products +#define GGML_CUDA_CC_VOLTA 700 +#define GGML_CUDA_CC_TURING 750 +#define GGML_CUDA_CC_AMPERE 800 +#define GGML_CUDA_CC_ORIN 870 +#define GGML_CUDA_CC_ADA_LOVELACE 890 +#define GGML_CUDA_CC_HOPPER 900 +// While BW spans CC 1000, 1100 & 1200, we are integrating Tensor Core instructions available to 1200 family, see +// https://docs.nvidia.com/cutlass/media/docs/cpp/blackwell_functionality.html#blackwell-sm120-gemms +#define GGML_CUDA_CC_BLACKWELL 1200 +#define GGML_CUDA_CC_DGX_SPARK 1210 +#define GGML_CUDA_CC_RUBIN 1300 +#define GGML_CUDA_CC_OFFSET_AMD 0x1000000 +#define GGML_CUDA_CC_OFFSET_MTHREADS 0x0100000 +#define GGML_CUDA_CC_IS_NVIDIA(cc) (cc < GGML_CUDA_CC_OFFSET_MTHREADS) + +// AMD +// GCN/CDNA, wave size is 64 +#define GGML_CUDA_CC_GCN4 (GGML_CUDA_CC_OFFSET_AMD + 0x803) // Tonga, Fiji, Polaris, minimum for fast fp16 +#define GGML_CUDA_CC_VEGA (GGML_CUDA_CC_OFFSET_AMD + 0x900) // Vega56/64, minimum for fp16 dual issue +#define GGML_CUDA_CC_VEGA20 (GGML_CUDA_CC_OFFSET_AMD + 0x906) // MI50/Radeon VII, minimum for dp4a +#define GGML_CUDA_CC_GFX909 (GGML_CUDA_CC_OFFSET_AMD + 0x909) // GCN APU +#define GGML_CUDA_CC_GFX90C (GGML_CUDA_CC_OFFSET_AMD + 0x90c) // GCN APU +#define GGML_CUDA_CC_CDNA1 (GGML_CUDA_CC_OFFSET_AMD + 0x908) // MI100, minimum for MFMA, acc registers +#define GGML_CUDA_CC_CDNA2 (GGML_CUDA_CC_OFFSET_AMD + 0x90a) // MI210 (gfx90a), minimum acc register renaming +#define GGML_CUDA_CC_CDNA3 (GGML_CUDA_CC_OFFSET_AMD + 0x942) // MI300 +#define GGML_CUDA_CC_CDNA4 (GGML_CUDA_CC_OFFSET_AMD + 0x950) // MI350X/MI355X + +// RDNA removes MFMA, dp4a, xnack, acc registers, wave size is 32 +#define GGML_CUDA_CC_RDNA1 (GGML_CUDA_CC_OFFSET_AMD + 0x1010) // RX 5000 +#define GGML_CUDA_CC_RDNA2 (GGML_CUDA_CC_OFFSET_AMD + 0x1030) // RX 6000, minimum for dp4a +#define GGML_CUDA_CC_RDNA3 (GGML_CUDA_CC_OFFSET_AMD + 0x1100) // RX 7000, minimum for WMMA +#define GGML_CUDA_CC_RDNA3_5 (GGML_CUDA_CC_OFFSET_AMD + 0x1150) // AI 370, AI Max 395 laptops. +#define GGML_CUDA_CC_RDNA4 (GGML_CUDA_CC_OFFSET_AMD + 0x1200) // RX 9000 + +#define GGML_CUDA_CC_IS_AMD(cc) (cc >= GGML_CUDA_CC_OFFSET_AMD) +#define GGML_CUDA_CC_IS_RDNA(cc) (cc >= GGML_CUDA_CC_RDNA1) +#define GGML_CUDA_CC_IS_RDNA1(cc) (cc >= GGML_CUDA_CC_RDNA1 && cc < GGML_CUDA_CC_RDNA2) +#define GGML_CUDA_CC_IS_RDNA2(cc) (cc >= GGML_CUDA_CC_RDNA2 && cc < GGML_CUDA_CC_RDNA3) +#define GGML_CUDA_CC_IS_RDNA3_0(cc) (cc >= GGML_CUDA_CC_RDNA3 && cc < GGML_CUDA_CC_RDNA3_5) +#define GGML_CUDA_CC_IS_RDNA3_5(cc) (cc >= GGML_CUDA_CC_RDNA3_5 && cc < GGML_CUDA_CC_RDNA4) +#define GGML_CUDA_CC_IS_RDNA3(cc) (GGML_CUDA_CC_IS_RDNA3_0(cc) || GGML_CUDA_CC_IS_RDNA3_5(cc)) +#define GGML_CUDA_CC_IS_RDNA4(cc) (cc >= GGML_CUDA_CC_RDNA4) +#define GGML_CUDA_CC_IS_GCN_APU(cc) ((cc) == GGML_CUDA_CC_GFX909 || (cc) == GGML_CUDA_CC_GFX90C) +#define GGML_CUDA_CC_IS_GCN(cc) ((cc > GGML_CUDA_CC_OFFSET_AMD && cc < GGML_CUDA_CC_CDNA1) || GGML_CUDA_CC_IS_GCN_APU(cc)) +#define GGML_CUDA_CC_IS_CDNA(cc) (!GGML_CUDA_CC_IS_GCN_APU(cc) && cc >= GGML_CUDA_CC_CDNA1 && cc < GGML_CUDA_CC_RDNA1) +#define GGML_CUDA_CC_IS_CDNA1(cc) (GGML_CUDA_CC_IS_CDNA(cc) && cc >= GGML_CUDA_CC_CDNA1 && cc < GGML_CUDA_CC_CDNA2) +#define GGML_CUDA_CC_IS_CDNA2(cc) (GGML_CUDA_CC_IS_CDNA(cc) && cc >= GGML_CUDA_CC_CDNA2 && cc < GGML_CUDA_CC_CDNA3) +#define GGML_CUDA_CC_IS_CDNA3(cc) (GGML_CUDA_CC_IS_CDNA(cc) && cc >= GGML_CUDA_CC_CDNA3 && cc < GGML_CUDA_CC_CDNA4) +#define GGML_CUDA_CC_IS_CDNA4(cc) (GGML_CUDA_CC_IS_CDNA(cc) && cc >= GGML_CUDA_CC_CDNA4 && cc < GGML_CUDA_CC_RDNA1) + +// Moore Threads +#define MUSART_HMASK 40300 // MUSA rc4.3, min. ver. for half2 -> uint mask comparisons + +#define GGML_CUDA_CC_QY1 (GGML_CUDA_CC_OFFSET_MTHREADS + 0x210) // MTT S80, MTT S3000 +#define GGML_CUDA_CC_QY2 (GGML_CUDA_CC_OFFSET_MTHREADS + 0x220) // MTT S4000 +#define GGML_CUDA_CC_PH1 (GGML_CUDA_CC_OFFSET_MTHREADS + 0x310) // MTT S5000 + +#define GGML_CUDA_CC_IS_MTHREADS(cc) (cc >= GGML_CUDA_CC_OFFSET_MTHREADS && cc < GGML_CUDA_CC_OFFSET_AMD) +#define GGML_CUDA_CC_IS_QY1(cc) (cc >= GGML_CUDA_CC_QY1 && cc < GGML_CUDA_CC_QY2) +#define GGML_CUDA_CC_IS_QY2(cc) (cc >= GGML_CUDA_CC_QY2 && cc < GGML_CUDA_CC_PH1) +#define GGML_CUDA_CC_IS_PH1(cc) (cc >= GGML_CUDA_CC_PH1) + +#if !defined(GGML_USE_HIP) && (defined(GGML_USE_MUSA) || CUDART_VERSION >= 11070) +# define GGML_CUDA_USE_CUB +#endif // !defined(GGML_USE_HIP) && (defined(GGML_USE_MUSA) || CUDART_VERSION >= 11070) + +// PDL host-side support (cudaLaunchKernelEx) requires CUDART >= 11.8. +// However, this has been bugged in CTK < 12.3 for MSVC builds, see +// https://github.com/ggml-org/llama.cpp/pull/22522#discussion_r3302393293 +// __CUDA_ARCH__ is undefined in host passes; GPU arch check happens in device-side code. +#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) && \ + (CUDART_VERSION >= 12030 || (!(defined(_MSC_VER) && !defined(__clang__)) && CUDART_VERSION >= 11080)) +# define GGML_CUDA_USE_PDL +#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) && (CUDART_VERSION >= 12030 || (!(defined(_MSC_VER) && !defined(__clang__)) && CUDART_VERSION >= 11080)) + +static __device__ __forceinline__ void ggml_cuda_syncwarp() { +#ifndef GGML_USE_HIP + __syncwarp(); +#endif // GGML_USE_HIP +} + +static __device__ __forceinline__ void ggml_cuda_pdl_sync() { +#if defined(GGML_CUDA_USE_PDL) && __CUDA_ARCH__ >= GGML_CUDA_CC_HOPPER + cudaGridDependencySynchronize(); +#endif // defined(GGML_CUDA_USE_PDL) && __CUDA_ARCH__ >= GGML_CUDA_CC_HOPPER +} + +static __device__ __forceinline__ void ggml_cuda_pdl_lc() { +#if defined(GGML_CUDA_USE_PDL) && __CUDA_ARCH__ >= GGML_CUDA_CC_HOPPER + cudaTriggerProgrammaticLaunchCompletion(); +#endif // defined(GGML_CUDA_USE_PDL) && __CUDA_ARCH__ >= GGML_CUDA_CC_HOPPER +} + +#ifdef __CUDA_ARCH_LIST__ +constexpr bool ggml_cuda_has_arch_impl(int) { + return false; +} + +template +constexpr bool ggml_cuda_has_arch_impl(const int arch, const int first, Archs... rest) { + return arch == first || ggml_cuda_has_arch_impl(arch, rest...); +} + +constexpr bool ggml_cuda_has_arch(const int arch) { + return ggml_cuda_has_arch_impl(arch, __CUDA_ARCH_LIST__); +} + +constexpr int ggml_cuda_highest_compiled_arch_impl(const int /*arch*/, const int cur) { + if (cur == 0) { + return -1; + } + return cur; +} + +template +constexpr int ggml_cuda_highest_compiled_arch_impl(const int arch, const int cur, const int first, Archs... rest) { + if (first <= arch && first > cur) { + return ggml_cuda_highest_compiled_arch_impl(arch, first, rest...); + } else { + return ggml_cuda_highest_compiled_arch_impl(arch, cur, rest...); + } +} + +constexpr int ggml_cuda_highest_compiled_arch(const int arch) { + return ggml_cuda_highest_compiled_arch_impl(arch, 0, __CUDA_ARCH_LIST__); +} +#else +static int ggml_cuda_highest_compiled_arch(const int arch) { + return arch; +} +#endif // __CUDA_ARCH_LIST__ + +// --------------------------------------------------------------------------------------------------------- + +#define MATRIX_ROW_PADDING 512 // last row of quant. matrices is a multiple of this to avoid out-of-bounds memory accesses + +#define GGML_CUDA_MAX_STREAMS 8 + +[[noreturn]] +void ggml_cuda_error(const char * stmt, const char * func, const char * file, int line, const char * msg); + +#define CUDA_CHECK_GEN(err, success, error_fn) \ + do { \ + auto err_ = (err); \ + if (err_ != (success)) { \ + ggml_cuda_error(#err, __func__, __FILE__, __LINE__, error_fn(err_)); \ + } \ + } while (0) + +#define CUDA_CHECK(err) CUDA_CHECK_GEN(err, cudaSuccess, cudaGetErrorString) + + +#if CUDART_VERSION >= 12000 || defined(GGML_USE_MUSA) + static const char * cublas_get_error_str(const cublasStatus_t err) { + return cublasGetStatusString(err); + } +#else + static const char * cublas_get_error_str(const cublasStatus_t err) { + switch (err) { + case CUBLAS_STATUS_SUCCESS: return "CUBLAS_STATUS_SUCCESS"; + case CUBLAS_STATUS_NOT_INITIALIZED: return "CUBLAS_STATUS_NOT_INITIALIZED"; + case CUBLAS_STATUS_ALLOC_FAILED: return "CUBLAS_STATUS_ALLOC_FAILED"; + case CUBLAS_STATUS_INVALID_VALUE: return "CUBLAS_STATUS_INVALID_VALUE"; + case CUBLAS_STATUS_ARCH_MISMATCH: return "CUBLAS_STATUS_ARCH_MISMATCH"; + case CUBLAS_STATUS_MAPPING_ERROR: return "CUBLAS_STATUS_MAPPING_ERROR"; + case CUBLAS_STATUS_EXECUTION_FAILED: return "CUBLAS_STATUS_EXECUTION_FAILED"; + case CUBLAS_STATUS_INTERNAL_ERROR: return "CUBLAS_STATUS_INTERNAL_ERROR"; + case CUBLAS_STATUS_NOT_SUPPORTED: return "CUBLAS_STATUS_NOT_SUPPORTED"; + default: return "unknown error"; + } + } +#endif // CUDART_VERSION >= 12000 + +#define CUBLAS_CHECK(err) CUDA_CHECK_GEN(err, CUBLAS_STATUS_SUCCESS, cublas_get_error_str) + +#ifdef GGML_USE_NCCL +#define NCCL_CHECK(err) CUDA_CHECK_GEN(err, ncclSuccess, ncclGetErrorString) +#endif // GGML_USE_NCCL + +#if !defined(GGML_USE_HIP) && !defined(GGML_CUDA_NO_VMM) +static const char * cu_get_error_str(CUresult err) { + const char * err_str; + cuGetErrorString(err, &err_str); + return err_str; +} +#define CU_CHECK(err) CUDA_CHECK_GEN(err, CUDA_SUCCESS, cu_get_error_str) +#endif + +#if !defined(GGML_USE_HIP) +# define CUDA_SET_SHARED_MEMORY_LIMIT(kernel, nbytes) \ + do { \ + static bool shared_memory_limit_raised[GGML_CUDA_MAX_DEVICES] = { false }; \ + const int id = ggml_cuda_get_device(); \ + if (!shared_memory_limit_raised[id]) { \ + CUDA_CHECK(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, nbytes)); \ + shared_memory_limit_raised[id] = true; \ + } \ + } while (0) +#else +# define CUDA_SET_SHARED_MEMORY_LIMIT(kernel, nbytes) \ + do { \ + GGML_UNUSED(nbytes); \ + } while (0) +#endif // !defined(GGML_USE_HIP) + +#if CUDART_VERSION >= 11010 || defined(GGML_USE_MUSA) +#define GGML_CUDA_ASSUME(x) __builtin_assume(x) +#else +#define GGML_CUDA_ASSUME(x) +#endif // CUDART_VERSION >= 11010 + +#if (!defined(GGML_USE_HIP) && !defined(GGML_CUDA_NO_VMM)) || (defined(GGML_USE_HIP) && !defined(GGML_HIP_NO_VMM)) +#define GGML_USE_VMM +#endif // (!defined(GGML_USE_HIP) && !defined(GGML_CUDA_NO_VMM)) || (defined(GGML_USE_HIP) && !defined(GGML_HIP_NO_VMM)) + +#if defined(GGML_USE_HIP) || defined(GGML_USE_MUSA) || __CUDA_ARCH__ >= GGML_CUDA_CC_PASCAL +#define FP16_AVAILABLE +#endif // defined(GGML_USE_HIP) || defined(GGML_USE_MUSA) || __CUDA_ARCH__ >= GGML_CUDA_CC_PASCAL + +#if defined(FP16_AVAILABLE) && __CUDA_ARCH__ != 610 +#define FAST_FP16_AVAILABLE +#endif // defined(FP16_AVAILABLE) && __CUDA_ARCH__ != 610 + +#if defined(GGML_USE_HIP) && defined(CDNA) && !defined(GGML_HIP_NO_MMQ_MFMA) +#define AMD_MFMA_AVAILABLE +#endif // defined(GGML_USE_HIP) && defined(CDNA) && !defined(GGML_HIP_NO_MMQ_MFMA) + +#if defined(GGML_USE_HIP) && (defined(RDNA4) || defined(RDNA3)) +#define AMD_WMMA_AVAILABLE +#endif // defined(GGML_USE_HIP) && defined(RDNA4) + +// The Volta instructions are in principle available on Turing or newer but they are effectively unusable: +#if !defined(GGML_USE_HIP) && __CUDA_ARCH__ == GGML_CUDA_CC_VOLTA +#define VOLTA_MMA_AVAILABLE +#endif // !defined(GGML_USE_HIP) && __CUDA_ARCH__ == GGML_CUDA_CC_VOLTA + +#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) && __CUDA_ARCH__ >= GGML_CUDA_CC_TURING +#define TURING_MMA_AVAILABLE +#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) && __CUDA_ARCH__ >= GGML_CUDA_CC_TURING + +#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) && __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE +#define AMPERE_MMA_AVAILABLE +#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) && __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE + +// FreeToken: the native fp4 MMA needs an sm_120a/sm_120f build, so it is off unless GGML_CUDA_FP4_MMA is defined; FP4 weights then take Q8_1 activations +#if !defined(GGML_USE_HIP) && __CUDA_ARCH__ >= GGML_CUDA_CC_BLACKWELL && __CUDA_ARCH__ < GGML_CUDA_CC_RUBIN && defined(GGML_CUDA_FP4_MMA) +# define BLACKWELL_MMA_AVAILABLE +#endif // !defined(GGML_USE_HIP) && __CUDA_ARCH__ >= GGML_CUDA_CC_BLACKWELL + +#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) && __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE +#define CP_ASYNC_AVAILABLE +#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) && __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE + +#if !defined(GGML_CUDA_NO_FA) && !(defined(GGML_USE_MUSA) && __MUSA_ARCH__ < 220) +#define FLASH_ATTN_AVAILABLE +#endif // !defined(GGML_CUDA_NO_FA) && !(defined(GGML_USE_MUSA) && __MUSA_ARCH__ < 220) + +static bool fp16_available(const int cc) { + return ggml_cuda_highest_compiled_arch(cc) >= GGML_CUDA_CC_PASCAL || + (GGML_CUDA_CC_IS_MTHREADS(cc) && cc >= GGML_CUDA_CC_PH1); +} + +static bool fast_fp16_available(const int cc) { + return GGML_CUDA_CC_IS_AMD(cc) || + (GGML_CUDA_CC_IS_NVIDIA(cc) && fp16_available(cc) && ggml_cuda_highest_compiled_arch(cc) != 610) || + (GGML_CUDA_CC_IS_MTHREADS(cc) && fp16_available(cc)); +} + +// To be used for feature selection of external libraries, e.g. cuBLAS. +static bool fast_fp16_hardware_available(const int cc) { + return (GGML_CUDA_CC_IS_NVIDIA(cc) && cc >= GGML_CUDA_CC_PASCAL && cc != 610) || GGML_CUDA_CC_IS_AMD(cc) || + (GGML_CUDA_CC_IS_MTHREADS(cc) && cc >= GGML_CUDA_CC_QY2); +} + +// To be used for feature selection of external libraries, e.g. cuBLAS. +static bool fp16_mma_hardware_available(const int cc) { + return (GGML_CUDA_CC_IS_NVIDIA(cc) && cc >= GGML_CUDA_CC_VOLTA) || + GGML_CUDA_CC_IS_CDNA(cc) || GGML_CUDA_CC_IS_RDNA3(cc) || GGML_CUDA_CC_IS_RDNA4(cc) || + (GGML_CUDA_CC_IS_MTHREADS(cc) && cc >= GGML_CUDA_CC_QY2); +} + +// To be used for feature selection of external libraries, e.g. cuBLAS. +static bool fast_bf16_hardware_available(const int cc) { + return (GGML_CUDA_CC_IS_AMD(cc) && (cc >= GGML_CUDA_CC_RDNA3 || GGML_CUDA_CC_IS_CDNA(cc))) + || (GGML_CUDA_CC_IS_NVIDIA(cc) && cc >= GGML_CUDA_CC_AMPERE); +} + +static bool bf16_mma_hardware_available(const int cc) { + return (GGML_CUDA_CC_IS_NVIDIA(cc) && cc >= GGML_CUDA_CC_AMPERE) || + GGML_CUDA_CC_IS_CDNA(cc) || cc >= GGML_CUDA_CC_RDNA3 || + (GGML_CUDA_CC_IS_MTHREADS(cc) && cc >= GGML_CUDA_CC_PH1); +} + +static bool fp32_mma_hardware_available(const int cc) { + return GGML_CUDA_CC_IS_CDNA(cc); +} + +static bool amd_mfma_available(const int cc) { +#if !defined(GGML_HIP_NO_MMQ_MFMA) + return GGML_CUDA_CC_IS_CDNA(cc); +#else + return false; +#endif //!defined(GGML_HIP_NO_MMQ_MFMA) +} + +static bool amd_wmma_available(const int cc) { + return (GGML_CUDA_CC_IS_RDNA4(cc) || GGML_CUDA_CC_IS_RDNA3(cc)); +} + +static bool volta_mma_available(const int cc) { + return GGML_CUDA_CC_IS_NVIDIA(cc) && ggml_cuda_highest_compiled_arch(cc) == GGML_CUDA_CC_VOLTA; +} + +static bool turing_mma_available(const int cc) { + return GGML_CUDA_CC_IS_NVIDIA(cc) && ggml_cuda_highest_compiled_arch(cc) >= GGML_CUDA_CC_TURING; +} + +static bool ampere_mma_available(const int cc) { + return GGML_CUDA_CC_IS_NVIDIA(cc) && ggml_cuda_highest_compiled_arch(cc) >= GGML_CUDA_CC_AMPERE; +} + +static bool cp_async_available(const int cc) { + return GGML_CUDA_CC_IS_NVIDIA(cc) && ggml_cuda_highest_compiled_arch(cc) >= GGML_CUDA_CC_AMPERE; +} + +static bool blackwell_mma_available(const int cc) { +#ifdef GGML_CUDA_FP4_MMA + return GGML_CUDA_CC_IS_NVIDIA(cc) && ggml_cuda_highest_compiled_arch(cc) >= GGML_CUDA_CC_BLACKWELL && + ggml_cuda_highest_compiled_arch(cc) < GGML_CUDA_CC_RUBIN; +#else + GGML_UNUSED(cc); + return false; +#endif // GGML_CUDA_FP4_MMA +} + +// Checks whether the tensor's base data pointer and higher-dimensional strides are byte-aligned to `alignment` bytes. +static bool ggml_cuda_is_aligned(const ggml_tensor * tensor, const size_t alignment) { + GGML_ASSERT(tensor != nullptr); + return (reinterpret_cast(tensor->data) % alignment) == 0 && + tensor->nb[1] % alignment == 0 && + tensor->nb[2] % alignment == 0 && + tensor->nb[3] % alignment == 0; +} + +static constexpr __device__ int ggml_cuda_get_physical_warp_size() { +#if defined(GGML_USE_HIP) && (defined(__GFX9__) || defined(__GFX8__)) + return 64; +#else + return 32; +#endif // defined(GGML_USE_HIP) && (defined(__GFX9__) || defined(__GFX8__)) +} + +// Maximum number of bytes that can be copied in a single instruction. +static constexpr __device__ int ggml_cuda_get_max_cpy_bytes() { +#if defined(GGML_USE_HIP) || defined(GGML_USE_MUSA) + return 16; +#else +#if __CUDA_ARCH__ >= GGML_CUDA_CC_VOLTA + return 16; +#else + return 8; +#endif // __CUDA_ARCH__ >= GGML_CUDA_CC_VOLTA +#endif // defined(GGML_USE_HIP) || defined(GGML_USE_MUSA) +} + + +[[noreturn]] +static __device__ void no_device_code( + const char * file_name, const int line, const char * function_name, const int arch, const char * arch_list) { + +#if defined(GGML_USE_HIP) + printf("%s:%d: ERROR: HIP kernel %s has no device code compatible with HIP arch %d.\n", + file_name, line, function_name, arch); + GGML_UNUSED(arch_list); +#else + printf("%s:%d: ERROR: CUDA kernel %s has no device code compatible with CUDA arch %d. ggml-cuda.cu was compiled for: %s\n", + file_name, line, function_name, arch, arch_list); +#endif // defined(GGML_USE_HIP) + __trap(); + + GGML_UNUSED(no_device_code); // suppress unused function warning +} + +#ifdef __CUDA_ARCH__ +#define NO_DEVICE_CODE no_device_code(__FILE__, __LINE__, __FUNCTION__, __CUDA_ARCH__, STRINGIZE(__CUDA_ARCH_LIST__)) +#else +#define NO_DEVICE_CODE //GGML_ABORT("NO_DEVICE_CODE not valid in host code.") +#endif // __CUDA_ARCH__ + +// The compiler is always able to unroll loops if they contain continue expressions. +// In such cases loop unrolling can still be achieved via recursion: +template +struct ggml_cuda_unroll { + template + __device__ void operator()(const Func & f, Args... args) const { + f(n - 1, args...); + ggml_cuda_unroll{}(f, args...); + } +}; + +template <> +struct ggml_cuda_unroll<1> { + template + __device__ void operator()(const Func & f, Args... args) const { + f(0, args...); + } +}; + +template +static __device__ __forceinline__ int warp_reduce_sum(int x) { +#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) && __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE + return __reduce_add_sync(0xffffffff, x); +#else +#pragma unroll + for (int offset = width/2; offset > 0; offset >>= 1) { + x += __shfl_xor_sync(0xffffffff, x, offset, width); + } + return x; +#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) && __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE +} + +template +static __device__ __forceinline__ float warp_reduce_sum(float x) { +#pragma unroll + for (int offset = width/2; offset > 0; offset >>= 1) { + x += __shfl_xor_sync(0xffffffff, x, offset, width); + } + return x; +} + +template +static __device__ __forceinline__ float2 warp_reduce_sum(float2 a) { +#pragma unroll + for (int offset = width/2; offset > 0; offset >>= 1) { + a.x += __shfl_xor_sync(0xffffffff, a.x, offset, width); + a.y += __shfl_xor_sync(0xffffffff, a.y, offset, width); + } + return a; +} + +template +static __device__ __forceinline__ half2 warp_reduce_sum(half2 a) { +#ifdef FP16_AVAILABLE +#pragma unroll + for (int offset = width/2; offset > 0; offset >>= 1) { + a = __hadd2(a, __shfl_xor_sync(0xffffffff, a, offset, width)); + } + return a; + +#else + NO_DEVICE_CODE; + return a; +#endif // FP16_AVAILABLE +} + +template +static __device__ __forceinline__ int warp_reduce_all(int x) { + if (width == ggml_cuda_get_physical_warp_size()) { + return __all_sync(0xffffffff, x); + } else { +#pragma unroll + for (int offset = width/2; offset > 0; offset >>= 1) { + x = __shfl_xor_sync(0xffffffff, x, offset, width) && x; + } + return x; + } +} + +template +static __device__ __forceinline__ int warp_reduce_any(int x) { + if (width == ggml_cuda_get_physical_warp_size()) { + return __any_sync(0xffffffff, x); + } else { +#pragma unroll + for (int offset = width/2; offset > 0; offset >>= 1) { + x = __shfl_xor_sync(0xffffffff, x, offset, width) || x; + } + return x; + } +} + +template +static __device__ __forceinline__ float warp_reduce_max(float x) { +#pragma unroll + for (int offset = width/2; offset > 0; offset >>= 1) { + x = fmaxf(x, __shfl_xor_sync(0xffffffff, x, offset, width)); + } + return x; +} + +template +static __device__ __forceinline__ T warp_prefix_inclusive_sum(T x) { + const int lane_id = threadIdx.x % width; +#pragma unroll + for (int offset = 1; offset < width; offset <<= 1) { + const T t = __shfl_up_sync(0xffffffff, x, offset, width); + if (lane_id >= offset) { + x += t; + } + } + return x; +} + +template +static __device__ __forceinline__ float2 warp_prefix_inclusive_sum(float2 a) { + const int lane_id = threadIdx.x % width; +#pragma unroll + for (int offset = 1; offset < width; offset <<= 1) { + const float t_x = __shfl_up_sync(0xffffffff, a.x, offset, width); + const float t_y = __shfl_up_sync(0xffffffff, a.y, offset, width); + if (lane_id >= offset) { + a.x += t_x; + a.y += t_y; + } + } + return a; +} + +template +static __device__ __forceinline__ half2 warp_prefix_inclusive_sum(half2 a) { +#ifdef FP16_AVAILABLE + const int lane_id = threadIdx.x % width; +#pragma unroll + for (int offset = 1; offset < width; offset <<= 1) { + const half2 t = __shfl_up_sync(0xffffffff, a, offset, width); + if (lane_id >= offset) { + a = __hadd2(a, t); + } + } + return a; + +#else + NO_DEVICE_CODE; + return a; +#endif // FP16_AVAILABLE +} + +enum class block_reduce_method { + MAX, + SUM, +}; + +template +struct block_reduce_policy; + +template +inline constexpr bool is_any = (std::is_same_v || ...); + +template +inline constexpr bool ggml_cuda_dependent_false_v = false; + +template struct block_reduce_policy { + static __device__ T reduce(T val) { + if constexpr(is_any) { + return warp_reduce_sum(val); + } else { + static_assert(ggml_cuda_dependent_false_v, "Unsupported type for block reduce sum"); + } + } + + static __device__ T sentinel() { + if constexpr (std::is_same_v) { + return 0.0f; + } else if constexpr (std::is_same_v) { + return make_float2(0.0f, 0.0f); + } else if constexpr (std::is_same_v) { + return make_half2(0.0f, 0.0f); + } else if constexpr (std::is_same_v) { + return 0; + } else { + static_assert(ggml_cuda_dependent_false_v, "Unsupported type for block reduce sum"); + } + } +}; + +template struct block_reduce_policy { + static __device__ T reduce(T val) { + if constexpr (is_any) { + return warp_reduce_max(val); + } else { + static_assert(ggml_cuda_dependent_false_v, "Unsupported type for block reduce max"); + } + } + + static __device__ T sentinel() { + if constexpr (std::is_same_v) { + return -INFINITY; + } else if constexpr (std::is_same_v) { + return make_half2(-INFINITY, -INFINITY); + } else { + static_assert(ggml_cuda_dependent_false_v, "Unsupported type for block reduce max"); + } + } +}; + +template +static __device__ T block_reduce(T val, [[maybe_unused]] T * shared_vals) { + // for multi-warp reductions, callers must not reuse shared_vals until all reads from this invocation have completed + val = block_reduce_policy::reduce(val); + const unsigned int block_size = block_size_template == 0 ? blockDim.x : block_size_template; + if (block_size > WARP_SIZE) { + assert((block_size <= 1024) && (block_size % WARP_SIZE) == 0); + const int warp_id = threadIdx.x / WARP_SIZE; + const int lane_id = threadIdx.x % WARP_SIZE; + if (lane_id == 0) { + shared_vals[warp_id] = val; + } + __syncthreads(); + val = block_reduce_policy::sentinel(); + if (lane_id < (static_cast(block_size) / WARP_SIZE)) { + val = shared_vals[lane_id]; + } + return block_reduce_policy::reduce(val); + } + + return val; +} + +static __device__ __forceinline__ half ggml_cuda_hmax(const half a, const half b) { +#ifdef FP16_AVAILABLE + +#if !defined(GGML_USE_HIP) && CUDART_VERSION < CUDART_HMAX + return __float2half(fmaxf(__half2float(a), __half2float(b))); +#else + return __hmax(a, b); +#endif // !defined(GGML_USE_HIP) && CUDART_VERSION < CUDART_HMAX + +#else + NO_DEVICE_CODE; + GGML_UNUSED(b); + return a; +#endif // FP16_AVAILABLE +} + +static __device__ __forceinline__ half2 ggml_cuda_hmax2(const half2 a, const half2 b) { +#if defined(GGML_USE_HIP) + return half2(__hmax(a.x, b.x), __hmax(a.y, b.y)); +#elif CUDART_VERSION >= CUDART_HMAX + return __hmax2(a, b); +#else + half2 ret; + reinterpret_cast(ret.x) = __float2half(fmaxf( __low2float(a), __low2float(b))); + reinterpret_cast(ret.y) = __float2half(fmaxf(__high2float(a), __high2float(b))); + return ret; +#endif +} + +template +static __device__ __forceinline__ half2 warp_reduce_max(half2 x) { +#pragma unroll + for (int offset = width/2; offset > 0; offset >>= 1) { + x = ggml_cuda_hmax2(x, __shfl_xor_sync(0xffffffff, x, offset, width)); + } + return x; +} + +#if (defined(CUDART_VERSION) && CUDART_VERSION < CUDART_HMASK) || defined(GGML_USE_HIP) || \ + (defined(MUSART_VERSION) && MUSART_VERSION < MUSART_HMASK) +static __device__ __forceinline__ uint32_t __hgt2_mask(const half2 a, const half2 b) { + const uint32_t mask_low = 0x0000FFFF * (float( __low2half(a)) > float( __low2half(b))); + const uint32_t mask_high = 0xFFFF0000 * (float(__high2half(a)) > float(__high2half(b))); + return mask_low | mask_high; +} +#endif // (defined(CUDART_VERSION) && CUDART_VERSION < CUDART_HMASK) || defined(GGML_USE_HIP) || (defined(MUSART_VERSION) && MUSART_VERSION < MUSART_HMASK) + +static __device__ __forceinline__ int ggml_cuda_dp4a(const int a, const int b, int c) { +#if defined(GGML_USE_HIP) +#if defined(CDNA) || defined(RDNA2) || defined(__gfx906__) + c = __builtin_amdgcn_sdot4(a, b, c, false); +#elif defined(RDNA3) || defined(RDNA4) + c = __builtin_amdgcn_sudot4( true, a, true, b, c, false); +#elif defined(RDNA1) || defined(__gfx900__) + int tmp1; + int tmp2; + asm("\n \ + v_mul_i32_i24 %1, sext(%3), sext(%4) dst_sel:DWORD dst_unused:UNUSED_PAD src0_sel:BYTE_0 src1_sel:BYTE_0 \n \ + v_mul_i32_i24 %2, sext(%3), sext(%4) dst_sel:DWORD dst_unused:UNUSED_PAD src0_sel:BYTE_1 src1_sel:BYTE_1 \n \ + v_add3_u32 %0, %1, %2, %0 \n \ + v_mul_i32_i24 %1, sext(%3), sext(%4) dst_sel:DWORD dst_unused:UNUSED_PAD src0_sel:BYTE_2 src1_sel:BYTE_2 \n \ + v_mul_i32_i24 %2, sext(%3), sext(%4) dst_sel:DWORD dst_unused:UNUSED_PAD src0_sel:BYTE_3 src1_sel:BYTE_3 \n \ + v_add3_u32 %0, %1, %2, %0 \n \ + " + : "+v"(c), "=&v"(tmp1), "=&v"(tmp2) + : "v"(a), "v"(b) + ); +#else + const int8x4_t va = reinterpret_cast(a); + const int8x4_t vb = reinterpret_cast(b); + c += va[0] * vb[0] + va[1] * vb[1] + va[2] * vb[2] + va[3] * vb[3]; +#endif + return c; + +#else // defined(GGML_USE_HIP) + +#if __CUDA_ARCH__ >= GGML_CUDA_CC_DP4A || defined(GGML_USE_MUSA) + return __dp4a(a, b, c); +#else // __CUDA_ARCH__ >= GGML_CUDA_CC_DP4A || defined(GGML_USE_MUSA) + const int8_t * a8 = (const int8_t *) &a; + const int8_t * b8 = (const int8_t *) &b; + return c + a8[0]*b8[0] + a8[1]*b8[1] + a8[2]*b8[2] + a8[3]*b8[3]; +#endif // __CUDA_ARCH__ >= GGML_CUDA_CC_DP4A || defined(GGML_USE_MUSA) + +#endif // defined(GGML_USE_HIP) +} + +static __device__ __forceinline__ void ggml_cuda_mad(float & acc, const float v, const float u) { + acc += v*u; +} + +static __device__ __forceinline__ void ggml_cuda_mad(float & acc, const float2 v, const float2 u) { + acc += v.x*u.x; + acc += v.y*u.y; +} + +#if defined(GGML_USE_HIP) && (defined(RDNA2) || defined(RDNA3) || defined(RDNA4) || defined(__gfx906__) || defined(CDNA)) +#define V_DOT2_F32_F16_AVAILABLE +#endif // defined(GGML_USE_HIP) && (defined(RDNA2) || defined(RDNA3) || defined(RDNA4) || defined(__gfx906__) || defined(CDNA)) + +static __device__ __forceinline__ void ggml_cuda_mad(float & acc, const half2 v, const half2 u) { +#ifdef V_DOT2_F32_F16_AVAILABLE + asm volatile("v_dot2_f32_f16 %0, %1, %2, %0" : "+v"(acc) : "v"(v), "v"(u)); +#else +#ifdef FAST_FP16_AVAILABLE + const float2 tmp = __half22float2(v*u); + acc += tmp.x + tmp.y; +#else + const float2 tmpv = __half22float2(v); + const float2 tmpu = __half22float2(u); + acc += tmpv.x * tmpu.x; + acc += tmpv.y * tmpu.y; +#endif // FAST_FP16_AVAILABLE +#endif // V_DOT2_F32_F16_AVAILABLE +} + +static __device__ __forceinline__ void ggml_cuda_mad(half2 & acc, const half2 v, const half2 u) { +#ifdef FAST_FP16_AVAILABLE + acc += v*u; +#else + const float2 tmpv = __half22float2(v); + const float2 tmpu = __half22float2(u); + float2 tmpacc = __half22float2(acc); + tmpacc.x += tmpv.x * tmpu.x; + tmpacc.y += tmpv.y * tmpu.y; + acc = make_half2(tmpacc.x, tmpacc.y); +#endif // FAST_FP16_AVAILABLE +} + +// Aligned memory transfers of 8/16 bytes can be faster than 2 transfers with 4 bytes, especially on AMD. +// Important: do not use this function if dst and src both point at registers. +// Due to the strict aliasing rule the compiler can do incorrect optimizations if src and dst have different types. +// The function is intended for copies between registers and SRAM/VRAM to make the compiler emit the right instructions. +// If dst and src point at different address spaces then they are guaranteed to not be aliased. +template +static __device__ __forceinline__ void ggml_cuda_memcpy_1(void * __restrict__ dst, const void * __restrict__ src) { + static_assert( + nbytes <= ggml_cuda_get_max_cpy_bytes() || alignment == 0, + "You are misusing the alignment parameter for ggml_cuda_memcpy_1. " + "The intent is for the parameter is only as a workaround if either one of the pointers is not properly aligned. " + "If you use it to do more bytes per copy than ggml_cuda_max_cpy_bytes() the reads and writes may not be coalesced. " + "Call ggml_cuda_memcpy_1 in a loop instead."); + if constexpr (alignment != 0) { + static_assert(nbytes % alignment == 0, "bad alignment"); + } + constexpr int nb_per_cpy = alignment == 0 ? nbytes : alignment; + +#pragma unroll + for (int i = 0; i < nbytes/nb_per_cpy; ++i) { + if constexpr (nb_per_cpy == 1) { + ((char *) dst)[i] = ((const char *) src)[i]; + } else if constexpr (nb_per_cpy == 2) { + ((short *) dst)[i] = ((const short *) src)[i]; + } else if constexpr (nb_per_cpy == 4) { + ((int *) dst)[i] = ((const int *) src)[i]; + } else if constexpr (nb_per_cpy == 8) { + ((int2 *) dst)[i] = ((const int2 *) src)[i]; + } else if constexpr (nb_per_cpy == 16) { + ((int4 *) dst)[i] = ((const int4 *) src)[i]; + } else { + static_assert(nbytes == 0 && nbytes == -1, "bad nbytes"); + } + } +} + +static __device__ __forceinline__ float ggml_cuda_e8m0_to_fp32(uint8_t x) { +#if CUDART_VERSION >= 12080 + const nv_bfloat16 e = __nv_cvt_e8m0_to_bf16raw(x); + return (float) e; +#else + uint32_t bits; + if (x == 0) { + bits = 0x00400000; + } else { + bits = (uint32_t) x << 23; + } + + float result; + memcpy(&result, &bits, sizeof(float)); + return result; +#endif // CUDART_VERSION >= 12050 +} + +static __device__ __forceinline__ float ggml_cuda_ue4m3_to_fp32(uint8_t x) { +#if defined(GGML_USE_HIP) && defined(CDNA3) && defined(FP8_AVAILABLE) && HIP_VERSION >= 60200000 + // ROCm does not support fp8 in software on devices with fp8 hardware, + // but CDNA3 supports only e4m3_fnuz (no inf). + const uint32_t bits = x * (x != 0x7F && x != 0xFF); // Convert NaN to 0.0f to match CPU implementation. + const __hip_fp8_e4m3_fnuz xf = *reinterpret_cast(&bits); + return static_cast(xf) / 2; +#else +#if defined(FP8_AVAILABLE) && !defined(GGML_USE_HIP) + const uint32_t bits = x * (x != 0x7F && x != 0xFF); // Convert NaN to 0.0f to match CPU implementation. + const __nv_fp8_e4m3 xf = *reinterpret_cast(&bits); + return static_cast(xf) / 2; +#else + if (x == 0 || (x == 0x7F && x != 0xFF)) { // Convert NaN to 0.0f + return 0.0f; + } + const int exp = (x >> 3) & 0xF; + const int man = x & 0x7; + float raw; + if (exp == 0) { + raw = ldexpf((float) man, -9); + } else { + raw = ldexpf(1.0f + (float) man / 8.0f, exp - 7); + } + return static_cast(raw / 2); +#endif // defined(FP8_AVAILABLE) && !defined(GGML_USE_HIP) +#endif // defined(GGML_USE_HIP) && defined(CDNA3) && defined(FP8_AVAILABLE) && HIP_VERSION >= 60200000 +} + +static __device__ __forceinline__ uint8_t ggml_cuda_fp32_to_ue4m3(float x) { +#if defined(BLACKWELL_MMA_AVAILABLE) // This is used for NVFP4 subblock scale quantizations only + if (!(x > 0.0f)) { + return 0; + } + const __nv_fp8_e4m3 xf(x); + return xf.__x; +#else + GGML_UNUSED(x); + NO_DEVICE_CODE; // Used only for NVFP4 Scales for Activations, only for Blackwell +#endif // defined(BLACKWELL_MMA_AVAILABLE) +} + +__device__ __forceinline__ uint8_t ggml_cuda_float_to_fp4_e2m1(float x, float e) { + const uint8_t sign_bit = (x < 0.0f) << 3; + float ax = fabsf(x) * e; + + // Positive LUT + static constexpr float pos_lut[8] = { 0.0f, 0.5f, 1.0f, 1.5f, 2.0f, 3.0f, 4.0f, 6.0f }; + + int best_i = 0; + float best_err = fabsf(ax - pos_lut[0]); + +#pragma unroll + for (int i = 1; i < 8; ++i) { + const float err = fabsf(ax - pos_lut[i]); + if (err < best_err) { + best_err = err; + best_i = i; + } + } + + return static_cast(best_i | sign_bit); +} + +// See https://gmplib.org/~tege/divcnst-pldi94.pdf figure 4.1. +// Precompute mp (m' in the paper) and L such that division +// can be computed using a multiply (high 32b of 64b result) +// and a shift: +// +// n/d = (mulhi(n, mp) + n) >> L; +static const uint3 init_fastdiv_values(uint64_t d_64) { + GGML_ASSERT(d_64 != 0); + GGML_ASSERT(d_64 <= std::numeric_limits::max()); + + uint32_t d = (uint32_t)d_64; + + // compute L = ceil(log2(d)); + uint32_t L = 0; + while (L < 32 && (uint32_t{ 1 } << L) < d) { + L++; + } + + uint32_t mp = (uint32_t) ((uint64_t{ 1 } << 32) * ((uint64_t{ 1 } << L) - d) / d + 1); + // pack divisor as well to reduce error surface + return make_uint3(mp, L, d); +} + +static __device__ __forceinline__ uint32_t fastdiv(uint32_t n, const uint3 fastdiv_values) { + // expects fastdiv_values to contain in + // fastdiv_values.z is unused and optimized away by the compiler. + // Compute high 32 bits of n * mp + const uint32_t hi = __umulhi(n, fastdiv_values.x); + // add n, apply bit shift + return (hi + n) >> fastdiv_values.y; +} + +static __device__ __forceinline__ uint32_t fastmodulo(uint32_t n, const uint3 fastdiv_values) { + // expects fastdiv_values to contain in (see init_fastdiv_values) + return n - fastdiv(n, fastdiv_values) * fastdiv_values.z; +} + +// Calculate both division and modulo at once, returns +static __device__ __forceinline__ uint2 fast_div_modulo(uint32_t n, const uint3 fastdiv_values) { + // expects fastdiv_values to contain in (see init_fastdiv_values) + const uint32_t div_val = fastdiv(n, fastdiv_values); + const uint32_t mod_val = n - div_val * fastdiv_values.z; + return make_uint2(div_val, mod_val); +} + +typedef void (*dequantize_kernel_t)(const void * vx, const int64_t ib, const int iqs, float2 & v); + +template +using dequantize_kq_t = void (*)(const void * vx, const int64_t ib, dst_t * y, const int tid); + +static __device__ __forceinline__ float get_alibi_slope( + const float max_bias, const uint32_t h, const uint32_t n_head_log2, const float m0, const float m1 +) { + if (max_bias <= 0.0f) { + return 1.0f; + } + const float base = h < n_head_log2 ? m0 : m1; + const int exph = h < n_head_log2 ? h + 1 : 2*(h - n_head_log2) + 1; + + return powf(base, exph); +} + +template +struct ggml_cuda_type_traits; + +template<> +struct ggml_cuda_type_traits { + static constexpr int qk = 1; + static constexpr int qr = 1; + static constexpr int bs = sizeof(ggml_half); +}; + +template<> +struct ggml_cuda_type_traits { + static constexpr int qk = QK1_0; + static constexpr int qr = QR1_0; + static constexpr int qi = QI1_0; + static constexpr int bs = sizeof(block_q1_0); +}; + +template<> +struct ggml_cuda_type_traits { + static constexpr int qk = QK2_0; + static constexpr int qr = QR2_0; + static constexpr int qi = QI2_0; + static constexpr int bs = sizeof(block_q2_0); +}; + +template<> +struct ggml_cuda_type_traits { + static constexpr int qk = QK4_0; + static constexpr int qr = QR4_0; + static constexpr int qi = QI4_0; + static constexpr int bs = sizeof(block_q4_0); +}; + +template<> +struct ggml_cuda_type_traits { + static constexpr int qk = QK4_1; + static constexpr int qr = QR4_1; + static constexpr int qi = QI4_1; + static constexpr int bs = sizeof(block_q4_1); +}; + +template<> +struct ggml_cuda_type_traits { + static constexpr int qk = QK5_0; + static constexpr int qr = QR5_0; + static constexpr int qi = QI5_0; + static constexpr int bs = sizeof(block_q5_0); +}; + +template<> +struct ggml_cuda_type_traits { + static constexpr int qk = QK5_1; + static constexpr int qr = QR5_1; + static constexpr int qi = QI5_1; + static constexpr int bs = sizeof(block_q5_1); +}; + +template<> +struct ggml_cuda_type_traits { + static constexpr int qk = QK8_0; + static constexpr int qr = QR8_0; + static constexpr int qi = QI8_0; + static constexpr int bs = sizeof(block_q8_0); +}; + +template<> +struct ggml_cuda_type_traits { + static constexpr int qk = QK_MXFP4; + static constexpr int qr = QR_MXFP4; + static constexpr int qi = QI_MXFP4; + static constexpr int bs = sizeof(block_mxfp4); +}; + +template<> +struct ggml_cuda_type_traits { + static constexpr int qk = QK_NVFP4; + static constexpr int qr = QR_NVFP4; + static constexpr int qi = QI_NVFP4; + static constexpr int bs = sizeof(block_nvfp4); +}; + +template<> +struct ggml_cuda_type_traits { + static constexpr int qk = QK_K; + static constexpr int qr = QR2_K; + static constexpr int qi = QI2_K; + static constexpr int bs = sizeof(block_q2_K); +}; + +template<> +struct ggml_cuda_type_traits { + static constexpr int qk = QK_K; + static constexpr int qr = QR3_K; + static constexpr int qi = QI3_K; + static constexpr int bs = sizeof(block_q3_K); +}; + +template<> +struct ggml_cuda_type_traits { + static constexpr int qk = QK_K; + static constexpr int qr = QR4_K; + static constexpr int qi = QI4_K; + static constexpr int bs = sizeof(block_q4_K); +}; + +template<> +struct ggml_cuda_type_traits { + static constexpr int qk = QK_K; + static constexpr int qr = QR5_K; + static constexpr int qi = QI5_K; + static constexpr int bs = sizeof(block_q5_K); +}; + +template<> +struct ggml_cuda_type_traits { + static constexpr int qk = QK_K; + static constexpr int qr = QR6_K; + static constexpr int qi = QI6_K; + static constexpr int bs = sizeof(block_q6_K); +}; + +template<> +struct ggml_cuda_type_traits { + static constexpr int qk = QK_K; + static constexpr int qr = QR2_XXS; + static constexpr int qi = QI2_XXS; + static constexpr int bs = sizeof(block_iq2_xxs); +}; + +template<> +struct ggml_cuda_type_traits { + static constexpr int qk = QK_K; + static constexpr int qr = QR2_XS; + static constexpr int qi = QI2_XS; + static constexpr int bs = sizeof(block_iq2_xs); +}; + +template<> +struct ggml_cuda_type_traits { + static constexpr int qk = QK_K; + static constexpr int qr = QR2_S; + static constexpr int qi = QI2_S; + static constexpr int bs = sizeof(block_iq2_s); +}; + +template<> +struct ggml_cuda_type_traits { + static constexpr int qk = QK_K; + static constexpr int qr = QR3_XXS; + static constexpr int qi = QI3_XXS; + static constexpr int bs = sizeof(block_iq3_xxs); +}; + +template<> +struct ggml_cuda_type_traits { + static constexpr int qk = QK_K; + static constexpr int qr = QR1_S; + static constexpr int qi = QI1_S; + static constexpr int bs = sizeof(block_iq1_s); +}; + +template<> +struct ggml_cuda_type_traits { + static constexpr int qk = QK_K; + static constexpr int qr = QR1_M; + static constexpr int qi = QI1_M; + static constexpr int bs = sizeof(block_iq1_m); +}; + +template<> +struct ggml_cuda_type_traits { + static constexpr int qk = QK4_NL; + static constexpr int qr = QR4_NL; + static constexpr int qi = QI4_NL; + static constexpr int bs = sizeof(block_iq4_nl); +}; + +template<> +struct ggml_cuda_type_traits { + static constexpr int qk = QK_K; + static constexpr int qr = QR4_XS; + static constexpr int qi = QI4_XS; + static constexpr int bs = sizeof(block_iq4_xs); +}; + +template<> +struct ggml_cuda_type_traits { + static constexpr int qk = QK_K; + static constexpr int qr = QR3_S; + static constexpr int qi = QI3_S; + static constexpr int bs = sizeof(block_iq3_s); +}; + +////////////////////// + +struct ggml_cuda_device_info { + int device_count; // number of (possibly virtual) devices exposed to the rest of ggml + int physical_device_count; // number of physical CUDA devices actually present + + struct cuda_device_info { + int cc; // compute capability + int nsm; // number of streaming multiprocessors + size_t smpb; // max. shared memory per block + size_t smpbo; // max. shared memory per block (with opt-in) + bool integrated; // Device is integrated as opposed to discrete + bool vmm; // virtual memory support + size_t vmm_granularity; // granularity of virtual memory + size_t total_vram; + int warp_size; // Number of threads in a dispatch + bool supports_cooperative_launch; // whether cooperative launch is supported + int physical_device; // backing physical CUDA device for this (virtual) device + int physical_share_count; // number of (virtual) devices sharing this device's physical GPU + int virtual_index; // index of this (virtual) device among those sharing its physical GPU + }; + + cuda_device_info devices[GGML_CUDA_MAX_DEVICES] = {}; + + std::array default_tensor_split = {}; +}; + +const ggml_cuda_device_info & ggml_cuda_info(); + +void ggml_cuda_set_device(int device); +int ggml_cuda_get_device(); + +struct ggml_cuda_pool { + virtual ~ggml_cuda_pool() = default; + + virtual void * alloc(size_t size, size_t * actual_size) = 0; + virtual void free(void * ptr, size_t size) = 0; +}; + +template +struct ggml_cuda_pool_alloc { + ggml_cuda_pool * pool = nullptr; + T * ptr = nullptr; + size_t actual_size = 0; + + ggml_cuda_pool_alloc() = default; + + explicit ggml_cuda_pool_alloc(ggml_cuda_pool & pool) : pool(&pool) { + } + + ggml_cuda_pool_alloc(ggml_cuda_pool & pool, size_t size) : pool(&pool) { + alloc(size); + } + + ~ggml_cuda_pool_alloc() { + if (ptr != nullptr) { + pool->free(ptr, actual_size); + } + } + + // size is in number of elements + T * alloc(size_t size) { + GGML_ASSERT(pool != nullptr); + GGML_ASSERT(ptr == nullptr); + ptr = (T *) pool->alloc(size * sizeof(T), &this->actual_size); + return ptr; + } + + T * alloc(ggml_cuda_pool & pool, size_t size) { + this->pool = &pool; + return alloc(size); + } + + T * get() { + return ptr; + } + + ggml_cuda_pool_alloc(const ggml_cuda_pool_alloc &) = delete; + ggml_cuda_pool_alloc(ggml_cuda_pool_alloc &&) = delete; + ggml_cuda_pool_alloc& operator=(const ggml_cuda_pool_alloc &) = delete; + ggml_cuda_pool_alloc& operator=(ggml_cuda_pool_alloc &&) = delete; +}; + + +// backend interface + +struct ggml_tensor_extra_gpu { + void * data_device[GGML_CUDA_MAX_DEVICES]; // 1 pointer for each device for split tensors + cudaEvent_t events[GGML_CUDA_MAX_DEVICES][GGML_CUDA_MAX_STREAMS]; // events for synchronizing multiple GPUs +}; + + +#if (defined(GGML_CUDA_USE_GRAPHS) || defined(GGML_HIP_GRAPHS)) || defined(GGML_MUSA_GRAPHS) +#define USE_CUDA_GRAPH +#endif + +struct ggml_cuda_graph { +#ifdef USE_CUDA_GRAPH + ~ggml_cuda_graph() { + if (instance != nullptr) { + CUDA_CHECK(cudaGraphExecDestroy(instance)); + } + if (graph != nullptr) { + CUDA_CHECK(cudaGraphDestroy(graph)); + } + } + cudaGraph_t graph = nullptr; + cudaGraphExec_t instance = nullptr; + size_t num_nodes = 0; + std::vector nodes; + bool disable_due_to_gpu_arch = false; + bool warmup_complete = false; + uint64_t uid = 0; + int64_t last_used_time = 0; + struct node_properties { + ggml_tensor node; + void * node_src_data_ptrs[GGML_MAX_SRC]; + int64_t node_src_ne[GGML_MAX_SRC][GGML_MAX_DIMS]; + size_t node_src_nb[GGML_MAX_SRC][GGML_MAX_DIMS]; + }; + std::vector node_props; + + bool is_enabled() const { + static const bool disable_cuda_graphs_due_to_env = (getenv("GGML_CUDA_DISABLE_GRAPHS") != nullptr); + return !(disable_due_to_gpu_arch || disable_cuda_graphs_due_to_env); + } +#endif +}; + +struct ggml_cuda_concurrent_event { + std::vector join_events; + cudaEvent_t fork_event = nullptr; + + int n_streams = 0; + std::unordered_map stream_mapping; + + // Original order of nodes in this concurrent region (before interleaving) + // Used to restore grouping for fusion within streams + std::vector original_order; + + const ggml_tensor * join_node; + + ggml_cuda_concurrent_event() = default; + + ggml_cuda_concurrent_event(const ggml_cuda_concurrent_event &) = delete; + ggml_cuda_concurrent_event & operator=(const ggml_cuda_concurrent_event &) = delete; + + explicit ggml_cuda_concurrent_event(int n_streams) : n_streams(n_streams) { + join_events.resize(n_streams); + + for (size_t i = 0; i < join_events.size(); ++i) { + CUDA_CHECK(cudaEventCreateWithFlags(&join_events[i], cudaEventDisableTiming)); + } + + CUDA_CHECK(cudaEventCreateWithFlags(&fork_event, cudaEventDisableTiming)); + } + + ggml_cuda_concurrent_event(ggml_cuda_concurrent_event && other) noexcept + : join_events(std::move(other.join_events)) + , fork_event(other.fork_event) + , n_streams(other.n_streams) + , stream_mapping(std::move(other.stream_mapping)) + , original_order(std::move(other.original_order)) + , join_node(other.join_node) { + other.fork_event = nullptr; + } + + // 1. check if any branches write to overlapping memory ranges (except the join node) + // 2. check whether all srcs are either within the branch or outside the nodes covered by ggml_cuda_concurrent_event + // we assume all nodes have the same buffer + bool is_valid() const { + std::vector>> write_ranges; + write_ranges.resize(n_streams); + + // get join_node's memory range to exclude from overlap checking. + // multiple nodes can use join_node's buffer; we synchronize on the join node. + const ggml_tensor * join_t = join_node->view_src ? join_node->view_src : join_node; + const int64_t join_start = (int64_t) join_t->data; + const int64_t join_end = join_start + ggml_nbytes(join_t); + + for (const auto & [tensor, stream] : stream_mapping) { + const ggml_tensor * t = tensor->view_src ? tensor->view_src : tensor; + const int64_t t_start = (int64_t) t->data; + const int64_t t_end = t_start + ggml_nbytes(t); + + // skip tensors that overlap with join_node's buffer. + if ((t_start <= join_start && join_start < t_end) || (join_start <= t_start && t_start < join_end)) { + continue; + } + + // concurrent streams begin from 1 + write_ranges[stream - 1].emplace_back(t_start, t_end); + } + + for (int i = 0; i < n_streams; ++i) { + // sorts first by start then by end of write range + std::sort(write_ranges[i].begin(), write_ranges[i].end()); + } + + bool writes_overlap = false; + bool dependent_srcs = false; + for (const auto & [tensor, stream] : stream_mapping) { + const ggml_tensor * t = tensor->view_src ? tensor->view_src : tensor; + const int64_t t_start = (int64_t) t->data; + const int64_t t_end = t_start + ggml_nbytes(t); + + // skip tensors that overlap with join_node's buffer + if ((t_start <= join_start && join_start < t_end) || (join_start <= t_start && t_start < join_end)) { + continue; + } + + // check if this buffer's write data overlaps with another stream's + std::pair data_range = std::make_pair(t_start, t_end); + for (int i = 0; i < n_streams; ++i) { + if (i == stream - 1) { + continue; + } + auto it = std::lower_bound(write_ranges[i].begin(), write_ranges[i].end(), data_range); + + if (it != write_ranges[i].end()) { + const std::pair & other = *it; + + // std::lower_bound returns the first element where other >= data_range (lexicographically). + // This guarantees other.first >= data_range.first. + // Therefore, overlap occurs iff other.first < data_range.second + // (i.e., the other range starts before this range ends). + if (other.first < data_range.second) { + GGML_LOG_DEBUG("Writes overlap for %s", tensor->name); + writes_overlap = true; + break; + } + } + } + + //check if all srcs are either in branch or don't have a branch + for (int i = 0; i < GGML_MAX_SRC; ++i) { + if (!tensor->src[i]) { + continue; + } + + auto it = stream_mapping.find(tensor->src[i]); + + if (it == stream_mapping.end()) { + continue; + } + + if (it->second != stream) { + dependent_srcs = true; + break; + } + } + + if (dependent_srcs || writes_overlap) { + break; + } + } + + return !writes_overlap && !dependent_srcs; + } + + ~ggml_cuda_concurrent_event() { + if (fork_event != nullptr) { + CUDA_CHECK(cudaEventDestroy(fork_event)); + } + for (cudaEvent_t e : join_events) { + if (e != nullptr) { + CUDA_CHECK(cudaEventDestroy(e)); + } + } + } +}; + +struct ggml_cuda_stream_context { + std::unordered_map concurrent_events; + + void reset() { + concurrent_events.clear(); + } +}; + +struct ggml_backend_cuda_context { + int device; + std::string name; + cudaEvent_t copy_event = nullptr; + + cudaStream_t streams[GGML_CUDA_MAX_DEVICES][GGML_CUDA_MAX_STREAMS] = {}; + cublasHandle_t cublas_handles[GGML_CUDA_MAX_DEVICES][GGML_CUDA_MAX_STREAMS] = {}; + void * cublas_workspaces[GGML_CUDA_MAX_DEVICES][GGML_CUDA_MAX_STREAMS] = {}; + size_t cublas_workspace_sizes[GGML_CUDA_MAX_DEVICES] = {}; + + int curr_stream_no = 0; + +#ifdef USE_CUDA_GRAPH + // Map from first_node_ptr to cuda_graph - allows multiple graphs per context + // when the computation is split across CPU/GPU (e.g., with --n-cpu-moe) + std::unordered_map> cuda_graphs; + + int64_t last_graph_eviction_sweep = 0; + + ggml_cuda_graph * cuda_graph(const void * first_node_ptr) { + const int64_t time_now = ggml_time_us(); + + // sweep every 5s, evicting cuda graphs unused for >=10s + if (time_now - last_graph_eviction_sweep >= 5'000'000) { + last_graph_eviction_sweep = time_now; + for (auto it = cuda_graphs.begin(); it != cuda_graphs.end(); ) { + if (time_now - it->second->last_used_time >= 10'000'000) { + it = cuda_graphs.erase(it); + } else { + ++it; + } + } + } + + auto it = cuda_graphs.find(first_node_ptr); + if (it == cuda_graphs.end()) { + it = cuda_graphs.emplace(first_node_ptr, std::make_unique()).first; + } + it->second->last_used_time = time_now; + return it->second.get(); + } + + // Check if any CUDA graph is enabled for this context (used by kernels that need to know + // if graphs are in use without having access to the specific graph key) + bool any_cuda_graph_enabled() const { + for (const auto & [key, graph] : cuda_graphs) { + if (graph && graph->is_enabled()) { + return true; + } + } + return false; + } + + // Check if any CUDA graph has an instance for this context + bool any_cuda_graph_has_instance() const { + for (const auto & [key, graph] : cuda_graphs) { + if (graph && graph->instance != nullptr) { + return true; + } + } + return false; + } +#endif // USE_CUDA_GRAPH + + explicit ggml_backend_cuda_context(int device) : + device(device), + name(GGML_CUDA_NAME + std::to_string(device)) { + } + + ggml_cuda_stream_context concurrent_stream_context; + + ~ggml_backend_cuda_context(); + + cudaStream_t stream(int device, int stream) { + if (streams[device][stream] == nullptr) { + ggml_cuda_set_device(device); + CUDA_CHECK(cudaStreamCreateWithFlags(&streams[device][stream], cudaStreamNonBlocking)); + } + return streams[device][stream]; + } + + cudaStream_t stream() { return stream(device, curr_stream_no); } + + ggml_cuda_stream_context & stream_context() { return concurrent_stream_context; } + + cublasHandle_t cublas_handle() { + if (cublas_handles[device][curr_stream_no] == nullptr) { + ggml_cuda_set_device(device); + CUBLAS_CHECK(cublasCreate(&cublas_handles[device][curr_stream_no])); + CUBLAS_CHECK(cublasSetMathMode(cublas_handles[device][curr_stream_no], CUBLAS_TF32_TENSOR_OP_MATH)); + CUBLAS_CHECK(cublasSetStream(cublas_handles[device][curr_stream_no], stream())); +#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) && (CUBLAS_VER_MAJOR > 11 || (CUBLAS_VER_MAJOR == 11 && CUBLAS_VER_MINOR >= 2)) + if (cublas_workspace_sizes[device] == 0) { + const int cc = ggml_cuda_info().devices[device].cc; + cublas_workspace_sizes[device] = (cc >= GGML_CUDA_CC_HOPPER) ? 32 * 1024 * 1024 : 4 * 1024 * 1024; + } + CUDA_CHECK(cudaMalloc(&cublas_workspaces[device][curr_stream_no], cublas_workspace_sizes[device])); + CUBLAS_CHECK(cublasSetWorkspace(cublas_handles[device][curr_stream_no], cublas_workspaces[device][curr_stream_no], cublas_workspace_sizes[device])); +#endif + } + return cublas_handles[device][curr_stream_no]; + } + + // pool + std::unique_ptr pools[GGML_CUDA_MAX_DEVICES][GGML_CUDA_MAX_STREAMS]; + + static std::unique_ptr new_pool_for_device(int device, int stream_no); + + ggml_cuda_pool & pool(int device) { + if (pools[device][curr_stream_no] == nullptr) { + pools[device][curr_stream_no] = new_pool_for_device(device, curr_stream_no); + } + return *pools[device][curr_stream_no]; + } + + ggml_cuda_pool & pool() { + return pool(device); + } +}; + +struct ggml_cuda_mm_fusion_args_host { + const ggml_tensor * x_bias = nullptr; + const ggml_tensor * gate = nullptr; + const ggml_tensor * gate_bias = nullptr; + const ggml_tensor * x_scale = nullptr; + const ggml_tensor * gate_scale = nullptr; + ggml_glu_op glu_op; + float glu_limit = 0.0f; + const ggml_tensor * shared_up = nullptr; + const ggml_tensor * shared_gate = nullptr; + ggml_tensor * shared_dst = nullptr; +}; +struct ggml_cuda_mm_fusion_args_device { + const void * x_bias = nullptr; + const void * gate = nullptr; + const void * gate_bias = nullptr; + const void * x_scale = nullptr; + const void * gate_scale = nullptr; + // FreeToken: element strides of x_scale / gate_scale, so they can be columns of a per-expert scale bank + uint32_t x_scale_stride = 1; + uint32_t gate_scale_stride = 1; + ggml_glu_op glu_op; + float glu_limit = 0.0f; + const void * shared_up = nullptr; + const void * shared_gate = nullptr; + float * shared_dst = nullptr; + uint32_t shared_stride_col_dst = 0; +}; + +struct ggml_cuda_kernel_launch_params { + dim3 block_nums; + dim3 block_dims; + size_t shmem; + cudaStream_t stream; + + // size_t shmem + ggml_cuda_kernel_launch_params(const dim3& block_nums_, const dim3& block_dims_, const size_t shmem_, const cudaStream_t stream_) + : block_nums(block_nums_), block_dims(block_dims_), shmem(shmem_), stream(stream_) {} + + // Some call sites pass ints instead of the required size_t. This 2nd constructor casts int->size_t to avoid these -Wnarrowing warnings. + ggml_cuda_kernel_launch_params(const dim3& block_nums_, const dim3& block_dims_, const int shmem_, const cudaStream_t stream_) + : block_nums(block_nums_), block_dims(block_dims_), shmem((size_t)shmem_), stream(stream_) {} +}; + +#if defined(GGML_CUDA_USE_PDL) +struct ggml_cuda_pdl_config { + cudaLaunchAttribute attr; + cudaLaunchConfig_t cfg; + + ggml_cuda_pdl_config(const ggml_cuda_kernel_launch_params & params) { + attr.id = cudaLaunchAttributeProgrammaticStreamSerialization; + attr.val.programmaticStreamSerializationAllowed = 1; + + cfg = {}; + cfg.gridDim = params.block_nums; + cfg.blockDim = params.block_dims; + cfg.dynamicSmemBytes = params.shmem; + cfg.stream = params.stream; + cfg.attrs = &attr; + cfg.numAttrs = 1; + } + + // Delete due to &attr + ggml_cuda_pdl_config(const ggml_cuda_pdl_config&) = delete; + ggml_cuda_pdl_config& operator=(const ggml_cuda_pdl_config&) = delete; + ggml_cuda_pdl_config& operator=(ggml_cuda_pdl_config&&) = delete; + +}; + +static bool ggml_cuda_kernel_can_use_pdl(const void * kernel) { + const int device = ggml_cuda_get_device(); + + struct cache_key { + int device; + const void * kernel; + + bool operator==(const cache_key & other) const { return device == other.device && kernel == other.kernel; } + }; + + struct cache_key_hash { + // MurmurHash3 mixing function for better hash distribution (vs. just std::hash which in some implementations simply returns the identity) + static size_t hash_mix(size_t x) { + std::uint64_t y = x; + const std::uint64_t m = 0xe9846af9b1a615d; + + y ^= y >> 32; + y *= m; + y ^= y >> 32; + y *= m; + y ^= y >> 28; + + return static_cast(y); + } + + size_t operator()(const cache_key & key) const { + // Use a nonzero seed to avoid mapping all-zero keys to zero + size_t h = 42; + h = hash_mix(h + key.device); + h = hash_mix(h + reinterpret_cast(key.kernel)); + return h; + } + }; + + static std::mutex cache_mutex; + static std::unordered_map cache; + + const cache_key key = { device, kernel }; + std::lock_guard lock(cache_mutex); + const auto it = cache.find(key); + if (it != cache.end()) { + return it->second; + } + + cudaFuncAttributes attr = {}; + CUDA_CHECK(cudaFuncGetAttributes(&attr, kernel)); + + // PDL device-side primitives are emitted only for PTX versions >= 90. + // We have to guard on a loaded kernel's PTX version so a kernel forward-JIT'ed + // from pre-Hopper PTX to a Hopper-or-newer GPU does not opt into PDL. + const bool can_use_pdl = attr.ptxVersion >= 90; + cache.emplace(key, can_use_pdl); + return can_use_pdl; +} + +#endif //defined(GGML_CUDA_USE_PDL) + +// PDL and __restrict__ need to be mutually exclusive, see https://github.com/ggml-org/llama.cpp/pull/24030 +# if (defined(GGML_CUDA_USE_PDL) && __CUDA_ARCH__ >= GGML_CUDA_CC_HOPPER) +# define GGML_CUDA_RESTRICT +# else +# define GGML_CUDA_RESTRICT __restrict__ +# endif // defined(GGML_CUDA_USE_PDL) && __CUDA_ARCH__ >= GGML_CUDA_CC_HOPPER + +template +static __inline__ void ggml_cuda_kernel_launch(Kernel kernel, const ggml_cuda_kernel_launch_params & launch_params, Args&&... args) { +#if defined(GGML_CUDA_USE_PDL) + + static const bool env_pdl_enabled = []() { + const char * env = getenv("GGML_CUDA_PDL"); + return env == nullptr || std::atoi(env) != 0; + }(); + + if (env_pdl_enabled && ggml_cuda_kernel_can_use_pdl(reinterpret_cast(kernel))) { + auto pdl_cfg = ggml_cuda_pdl_config(launch_params); + + CUDA_CHECK(cudaLaunchKernelEx(&pdl_cfg.cfg, kernel, std::forward(args)... )); + return; + } +#endif //defined(GGML_CUDA_USE_PDL) + + kernel<<>>(std::forward(args)... ); + CUDA_CHECK(cudaGetLastError()); +} diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/convert.cu b/python/freetoken/kernel/csrc/ggml/ggml-cuda/convert.cu new file mode 100644 index 0000000000..77fd288945 --- /dev/null +++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/convert.cu @@ -0,0 +1,743 @@ +#include "convert.cuh" +#include "dequantize.cuh" + +#include + +#define CUDA_Q8_0_NE_ALIGN 2048 + +template +static __global__ void dequantize_block(const void * __restrict__ vx, dst_t * __restrict__ y, + const int64_t ne00, const int64_t ne01, + const int64_t ne0203, const uint3 ne02, + const int64_t s01, const int64_t s02, const int64_t s03) { + const int64_t i00 = 2 * (int64_t(blockDim.x)*blockIdx.x + threadIdx.x); + + if (i00 >= ne00) { + return; + } + + for (int64_t i01 = blockIdx.y; i01 < ne01; i01 += gridDim.y) { + for (int64_t i0203 = blockIdx.z; i0203 < ne0203; i0203 += gridDim.z) { + const uint2 dm = fast_div_modulo((uint32_t)i0203, ne02); + const int64_t i02 = dm.y; + const int64_t i03 = dm.x; + + const int64_t ibx0 = i03*s03 + i02*s02 + i01*s01; + + const int64_t ib = ibx0 + i00/qk; // block index + const int64_t iqs = (i00%qk)/qr; // quant index + const int64_t iybs = i00 - i00%qk; // y block start index + const int64_t y_offset = qr == 1 ? 1 : qk/2; + + // dequantize + float2 v; + dequantize_kernel(vx, ib, iqs, v); + + const int64_t iy0 = (i0203*ne01 + i01)*ne00 + iybs + iqs; + y[iy0 + 0] = ggml_cuda_cast(v.x); + y[iy0 + y_offset] = ggml_cuda_cast(v.y); + } + } +} + +template +static __global__ void dequantize_block_q8_0_f16(const void * __restrict__ vx, half * __restrict__ y, const int64_t k) { +#if __CUDA_ARCH__ >= GGML_CUDA_CC_PASCAL + constexpr int nint = CUDA_Q8_0_NE_ALIGN/sizeof(int) + WARP_SIZE; + + const int64_t i0 = CUDA_Q8_0_NE_ALIGN*blockIdx.x; + const int * x0 = ((int *) vx) + blockIdx.x * nint; + half2 * y2 = (half2 *) (y + i0); + + __shared__ int vals[nint]; + +#pragma unroll + for (int ix0 = 0; ix0 < nint; ix0 += WARP_SIZE) { + if (need_check && i0*sizeof(block_q8_0)/QK8_0 + sizeof(int)*(ix0 + threadIdx.x) >= k*sizeof(block_q8_0)/QK8_0) { + break; + } + + const int ix = ix0 + threadIdx.x; + vals[ix] = x0[ix]; + } + + __syncthreads(); + +#pragma unroll + for (int iy = 0; iy < CUDA_Q8_0_NE_ALIGN; iy += 2*WARP_SIZE) { + if (need_check && i0 + iy + 2*threadIdx.x >= k) { + return; + } + + const half * b0 = ((const half *) vals) + (sizeof(block_q8_0)/sizeof(half)) * ((iy + 2*threadIdx.x)/QK8_0); + const half d = *b0; + const char2 qs = ((const char2 *) (b0 + 1))[threadIdx.x % (QK8_0/2)]; + + y2[iy/2 + threadIdx.x] = __hmul2(make_half2(qs.x, qs.y), __half2half2(d)); + } +#else + GGML_UNUSED_VARS(vx, y, k); + NO_DEVICE_CODE; +#endif // __CUDA_ARCH__ >= GGML_CUDA_CC_PASCAL +} + +template +static __global__ void dequantize_block_q4_0(const void * __restrict__ vx, dst_t * __restrict__ yy, int nb32) { + + const int64_t i = blockIdx.x; + + // assume 32 threads + const int64_t tid = threadIdx.x; + const int64_t il = tid/8; + const int64_t ir = tid%8; + const int64_t ib = 8*i + ir; + if (ib >= nb32) { + return; + } + + dst_t * y = yy + 256*i + 32*ir + 4*il; + + const block_q4_0 * x = (const block_q4_0 *)vx + ib; + const float d = __half2float(x->d); + const float dm = -8*d; + + const uint8_t * q = x->qs + 4*il; + + for (int l = 0; l < 4; ++l) { + y[l+ 0] = ggml_cuda_cast(d * (q[l] & 0xF) + dm); + y[l+16] = ggml_cuda_cast(d * (q[l] >> 4) + dm); + } +} + +template +static __global__ void dequantize_block_q4_1(const void * __restrict__ vx, dst_t * __restrict__ yy, int nb32) { + + const int64_t i = blockIdx.x; + + // assume 32 threads + const int64_t tid = threadIdx.x; + const int64_t il = tid/8; + const int64_t ir = tid%8; + const int64_t ib = 8*i + ir; + if (ib >= nb32) { + return; + } + + dst_t * y = yy + 256*i + 32*ir + 4*il; + + const block_q4_1 * x = (const block_q4_1 *)vx + ib; + const float2 d = __half22float2(x->dm); + + const uint8_t * q = x->qs + 4*il; + + for (int l = 0; l < 4; ++l) { + y[l+ 0] = ggml_cuda_cast(d.x * (q[l] & 0xF) + d.y); + y[l+16] = ggml_cuda_cast(d.x * (q[l] >> 4) + d.y); + } +} + +//================================== k-quants + +template +static __global__ void dequantize_block_q2_K(const void * __restrict__ vx, dst_t * __restrict__ yy) { + const int64_t i = blockIdx.x; + + dequantize_q2_K(vx, i, yy + i*QK_K, threadIdx.x); +} + +template +static __global__ void dequantize_block_q3_K(const void * __restrict__ vx, dst_t * __restrict__ yy) { + const int64_t i = blockIdx.x; + + dequantize_q3_K(vx, i, yy + i*QK_K, threadIdx.x); +} + +template +static __global__ void dequantize_block_q4_K(const void * __restrict__ vx, dst_t * __restrict__ yy) { + const int64_t i = blockIdx.x; + + dequantize_q4_K(vx, i, yy + i*QK_K, threadIdx.x); +} + +template +static __global__ void dequantize_block_q5_K(const void * __restrict__ vx, dst_t * __restrict__ yy) { + const int64_t i = blockIdx.x; + + dequantize_q5_K(vx, i, yy + i*QK_K, threadIdx.x); +} + +template +static __global__ void dequantize_block_q6_K(const void * __restrict__ vx, dst_t * __restrict__ yy) { + const int64_t i = blockIdx.x; + + dequantize_q6_K(vx, i, yy + i*QK_K, threadIdx.x); +} + +template +static __global__ void dequantize_block_iq2_xxs(const void * __restrict__ vx, dst_t * __restrict__ yy) { + const int64_t i = blockIdx.x; + + dequantize_iq2_xxs(vx, i, yy + i*QK_K, threadIdx.x); +} + +template +static __global__ void dequantize_block_iq2_xs(const void * __restrict__ vx, dst_t * __restrict__ yy) { + const int64_t i = blockIdx.x; + + dequantize_iq2_xs(vx, i, yy + i*QK_K, threadIdx.x); +} + +template +static __global__ void dequantize_block_iq2_s(const void * __restrict__ vx, dst_t * __restrict__ yy) { + const int64_t i = blockIdx.x; + + dequantize_iq2_s(vx, i, yy + i*QK_K, threadIdx.x); +} + +template +static __global__ void dequantize_block_iq3_xxs(const void * __restrict__ vx, dst_t * __restrict__ yy) { + const int64_t i = blockIdx.x; + + dequantize_iq3_xxs(vx, i, yy + i*QK_K, threadIdx.x); +} + +template +static __global__ void dequantize_block_iq3_s(const void * __restrict__ vx, dst_t * __restrict__ yy) { + const int64_t i = blockIdx.x; + + dequantize_iq3_s(vx, i, yy + i*QK_K, threadIdx.x); +} + +template +static __global__ void dequantize_block_iq1_s(const void * __restrict__ vx, dst_t * __restrict__ yy) { + const int64_t i = blockIdx.x; + + dequantize_iq1_s(vx, i, yy + i*QK_K, threadIdx.x); +} + +template +static __global__ void dequantize_block_iq1_m(const void * __restrict__ vx, dst_t * __restrict__ yy) { + const int64_t i = blockIdx.x; + + dequantize_iq1_m(vx, i, yy + i*QK_K, threadIdx.x); +} + +template +static __global__ void dequantize_block_iq4_nl(const void * __restrict__ vx, dst_t * __restrict__ yy, int nb32) { + const int64_t i = blockIdx.x; + + const int64_t ib = 8*i + threadIdx.x%8; + if (ib >= nb32) { + return; + } + + dequantize_iq4_nl(vx, i, yy + i*QK_K, threadIdx.x); +} + +template +static __global__ void dequantize_block_iq4_xs(const void * __restrict__ vx, dst_t * __restrict__ yy) { + const int64_t i = blockIdx.x; + + dequantize_iq4_xs(vx, i, yy + i*QK_K, threadIdx.x); +} + +template +static __global__ void dequantize_block_mxfp4(const void * __restrict__ vx, dst_t * __restrict__ yy) { + const int64_t i = blockIdx.x; + + dequantize_mxfp4(vx, i, yy + i*QK_K, threadIdx.x); +} + +template +static void dequantize_block_cuda(const void * vx, dst_t * y, + const int64_t ne00, const int64_t ne01, const int64_t ne02, const int64_t ne03, + const int64_t s01, const int64_t s02, const int64_t s03, cudaStream_t stream) { + const int64_t ne0203 = ne02*ne03; + const uint3 ne02_fdv = init_fastdiv_values(ne02); + const dim3 num_blocks((ne00 + 2*CUDA_DEQUANTIZE_BLOCK_SIZE - 1) / (2*CUDA_DEQUANTIZE_BLOCK_SIZE), (int)std::min(ne01, (int64_t)65535), (int)std::min(ne0203, (int64_t)65535)); + dequantize_block<<>> + (vx, y, ne00, ne01, ne0203, ne02_fdv, s01, s02, s03); +} + +template +static void dequantize_block_cont_cuda(const void * __restrict__ vx, dst_t * __restrict__ y, const int64_t k, cudaStream_t stream) { + dequantize_block_cuda(vx, y, k, 1, 1, 1, k/qk, k/qk, k/qk, stream); +} + +static void dequantize_block_q8_0_f16_cuda(const void * __restrict__ vx, half * __restrict__ y, const int64_t k, cudaStream_t stream) { + const int num_blocks = (k + CUDA_Q8_0_NE_ALIGN - 1) / CUDA_Q8_0_NE_ALIGN; + if (k % CUDA_Q8_0_NE_ALIGN == 0) { + const bool need_check = false; + dequantize_block_q8_0_f16<<>>(vx, y, k); + } else { + const bool need_check = true; + dequantize_block_q8_0_f16<<>>(vx, y, k); + } +} + +template +static void dequantize_row_q2_K_cuda(const void * vx, dst_t * y, const int64_t k, cudaStream_t stream) { + const int nb = k / QK_K; + dequantize_block_q2_K<<>>(vx, y); +} + +template +static void dequantize_row_q3_K_cuda(const void * vx, dst_t * y, const int64_t k, cudaStream_t stream) { + const int nb = k / QK_K; + dequantize_block_q3_K<<>>(vx, y); +} + +template +static void dequantize_row_q4_0_cuda(const void * vx, dst_t * y, const int64_t k, cudaStream_t stream) { + const int nb32 = k / 32; + const int nb = (k + 255) / 256; + dequantize_block_q4_0<<>>(vx, y, nb32); +} + +template +static void dequantize_row_q4_1_cuda(const void * vx, dst_t * y, const int64_t k, cudaStream_t stream) { + const int nb32 = k / 32; + const int nb = (k + 255) / 256; + dequantize_block_q4_1<<>>(vx, y, nb32); +} + +template +static void dequantize_row_q4_K_cuda(const void * vx, dst_t * y, const int64_t k, cudaStream_t stream) { + const int nb = k / QK_K; + dequantize_block_q4_K<<>>(vx, y); +} + +template +static void dequantize_row_q5_K_cuda(const void * vx, dst_t * y, const int64_t k, cudaStream_t stream) { + const int nb = k / QK_K; + dequantize_block_q5_K<<>>(vx, y); +} + +template +static void dequantize_row_q6_K_cuda(const void * vx, dst_t * y, const int64_t k, cudaStream_t stream) { + const int nb = k / QK_K; + dequantize_block_q6_K<<>>(vx, y); +} + +template +static void dequantize_row_iq2_xxs_cuda(const void * vx, dst_t * y, const int64_t k, cudaStream_t stream) { + const int nb = k / QK_K; + dequantize_block_iq2_xxs<<>>(vx, y); +} + +template +static void dequantize_row_iq2_xs_cuda(const void * vx, dst_t * y, const int64_t k, cudaStream_t stream) { + const int nb = k / QK_K; + dequantize_block_iq2_xs<<>>(vx, y); +} + +template +static void dequantize_row_iq2_s_cuda(const void * vx, dst_t * y, const int64_t k, cudaStream_t stream) { + const int nb = k / QK_K; + dequantize_block_iq2_s<<>>(vx, y); +} + +template +static void dequantize_row_iq3_xxs_cuda(const void * vx, dst_t * y, const int64_t k, cudaStream_t stream) { + const int nb = k / QK_K; + dequantize_block_iq3_xxs<<>>(vx, y); +} + +template +static void dequantize_row_iq3_s_cuda(const void * vx, dst_t * y, const int64_t k, cudaStream_t stream) { + const int nb = k / QK_K; + dequantize_block_iq3_s<<>>(vx, y); +} + +template +static void dequantize_row_iq1_s_cuda(const void * vx, dst_t * y, const int64_t k, cudaStream_t stream) { + const int nb = k / QK_K; + dequantize_block_iq1_s<<>>(vx, y); +} + +template +static void dequantize_row_iq4_nl_cuda(const void * vx, dst_t * y, const int64_t k, cudaStream_t stream) { + const int nb32 = k / QK4_NL; + const int nb = (k + QK_K - 1) / QK_K; + dequantize_block_iq4_nl<<>>(vx, y, nb32); +} + +template +static void dequantize_row_iq1_m_cuda(const void * vx, dst_t * y, const int64_t k, cudaStream_t stream) { + const int nb = k / QK_K; + dequantize_block_iq1_m<<>>(vx, y); +} + +template +static void dequantize_row_iq4_xs_cuda(const void * vx, dst_t * y, const int64_t k, cudaStream_t stream) { + const int nb = (k + QK_K - 1) / QK_K; + dequantize_block_iq4_xs<<>>(vx, y); +} + +template +static void dequantize_row_mxfp4_cuda(const void * vx, dst_t * y, const int64_t k, cudaStream_t stream) { + const int nb = (k + QK_K - 1) / QK_K; + dequantize_block_mxfp4<<>>(vx, y); +} + +template +static __global__ void dequantize_block_nvfp4( + const void * __restrict__ vx, + dst_t * __restrict__ yy, + const int64_t ne) { + const int64_t i = blockIdx.x; + const int tid = threadIdx.x; + + const int64_t base = i * QK_NVFP4; + if (base >= ne) { + return; + } + + const block_nvfp4 * x = (const block_nvfp4 *) vx; + const block_nvfp4 & xb = x[i]; + + const int sub = tid / (QK_NVFP4_SUB / 2); + const int j = tid % (QK_NVFP4_SUB / 2); + + const float d = ggml_cuda_ue4m3_to_fp32(xb.d[sub]); + const uint8_t q = xb.qs[sub * (QK_NVFP4_SUB / 2) + j]; + + const int64_t y0 = base + sub * QK_NVFP4_SUB + j; + const int64_t y1 = y0 + QK_NVFP4_SUB / 2; + + yy[y0] = ggml_cuda_cast(d * kvalues_mxfp4[q & 0x0F]); + yy[y1] = ggml_cuda_cast(d * kvalues_mxfp4[q >> 4]); +} + +template +static void dequantize_row_nvfp4_cuda( + const void * vx, + dst_t * y, + const int64_t k, + cudaStream_t stream) { + GGML_ASSERT(k % QK_NVFP4 == 0); + const int nb = k / QK_NVFP4; + dequantize_block_nvfp4<<>>(vx, y, k); +} +template +static __global__ void convert_unary( + const void * __restrict__ vx, dst_t * __restrict__ y, const int64_t ne00, const int64_t ne01, + const int64_t ne0203, const uint3 ne02, + const int64_t s01, const int64_t s02, const int64_t s03) { + const int64_t i00 = (int64_t)blockDim.x*blockIdx.x + threadIdx.x; + + if (i00 >= ne00) { + return; + } + + const src_t * x = (const src_t *) vx; + + for (int64_t i01 = blockIdx.y; i01 < ne01; i01 += gridDim.y) { + for (int64_t i0203 = blockIdx.z; i0203 < ne0203; i0203 += gridDim.z) { + const uint2 dm = fast_div_modulo((uint32_t)i0203, ne02); + const int64_t i02 = dm.y; + const int64_t i03 = dm.x; + + const int64_t ix = i03*s03 + i02*s02 + i01*s01 + i00; + const int64_t iy = (i0203*ne01 + i01)*ne00 + i00; + y[iy] = ggml_cuda_cast(x[ix]); + } + } +} + +template struct alignas(sizeof(T)*4) cvt_vec4 { T v[4]; }; + +// four elements per thread, so a warp moves 512B (RDNA) / 1k (CDNA) per load +template +static __global__ void convert_unary_cont_vec4( + const void * __restrict__ vx, dst_t * __restrict__ y, const int64_t k4) { + const int64_t i = (int64_t)blockDim.x*blockIdx.x + threadIdx.x; + + if (i >= k4) { + return; + } + + const cvt_vec4 xv = ((const cvt_vec4 *) vx)[i]; + + cvt_vec4 yv; +#pragma unroll + for (int j = 0; j < 4; ++j) { + yv.v[j] = ggml_cuda_cast(xv.v[j]); + } + + ((cvt_vec4 *) y)[i] = yv; +} + +template +static void convert_unary_cuda(const void * vx, dst_t * y, + const int64_t ne00, const int64_t ne01, const int64_t ne02, const int64_t ne03, + const int64_t s01, const int64_t s02, const int64_t s03, cudaStream_t stream) { + const int64_t ne0203 = ne02*ne03; + const uint3 ne02_fdv = init_fastdiv_values(ne02); + const dim3 num_blocks((ne00 + CUDA_DEQUANTIZE_BLOCK_SIZE - 1) / CUDA_DEQUANTIZE_BLOCK_SIZE, (int)std::min(ne01, (int64_t)65535), (int)std::min(ne0203, (int64_t)65535)); + convert_unary<<>> + (vx, y, ne00, ne01, ne0203, ne02_fdv, s01, s02, s03); +} + +template +static void convert_unary_cont_cuda(const void * vx, dst_t * y, const int64_t k, cudaStream_t stream) { + if (k % 4 == 0 && + (uintptr_t) vx % alignof(cvt_vec4) == 0 && + (uintptr_t) y % alignof(cvt_vec4) == 0) { + const int64_t k4 = k/4; + const int64_t num_blocks = (k4 + CUDA_DEQUANTIZE_BLOCK_SIZE - 1) / CUDA_DEQUANTIZE_BLOCK_SIZE; + convert_unary_cont_vec4<<>>(vx, y, k4); + return; + } + + convert_unary_cuda(vx, y, k, 1, 1, 1, k, k, k, stream); +} + +to_bf16_cuda_t ggml_get_to_bf16_cuda(ggml_type type) { + switch (type) { + case GGML_TYPE_Q1_0: + return dequantize_block_cont_cuda; + case GGML_TYPE_Q2_0: + return dequantize_block_cont_cuda; + case GGML_TYPE_Q4_0: + return dequantize_row_q4_0_cuda; + case GGML_TYPE_Q4_1: + return dequantize_row_q4_1_cuda; + case GGML_TYPE_Q5_0: + return dequantize_block_cont_cuda; + case GGML_TYPE_Q5_1: + return dequantize_block_cont_cuda; + case GGML_TYPE_Q8_0: + return dequantize_block_cont_cuda; + case GGML_TYPE_Q2_K: + return dequantize_row_q2_K_cuda; + case GGML_TYPE_Q3_K: + return dequantize_row_q3_K_cuda; + case GGML_TYPE_Q4_K: + return dequantize_row_q4_K_cuda; + case GGML_TYPE_Q5_K: + return dequantize_row_q5_K_cuda; + case GGML_TYPE_Q6_K: + return dequantize_row_q6_K_cuda; + case GGML_TYPE_IQ2_XXS: + return dequantize_row_iq2_xxs_cuda; + case GGML_TYPE_IQ2_XS: + return dequantize_row_iq2_xs_cuda; + case GGML_TYPE_IQ2_S: + return dequantize_row_iq2_s_cuda; + case GGML_TYPE_IQ3_XXS: + return dequantize_row_iq3_xxs_cuda; + case GGML_TYPE_IQ1_S: + return dequantize_row_iq1_s_cuda; + case GGML_TYPE_IQ1_M: + return dequantize_row_iq1_m_cuda; + case GGML_TYPE_IQ4_NL: + return dequantize_row_iq4_nl_cuda; + case GGML_TYPE_IQ4_XS: + return dequantize_row_iq4_xs_cuda; + case GGML_TYPE_IQ3_S: + return dequantize_row_iq3_s_cuda; + case GGML_TYPE_MXFP4: + return dequantize_row_mxfp4_cuda; + case GGML_TYPE_NVFP4: + return dequantize_row_nvfp4_cuda; + case GGML_TYPE_F32: + return convert_unary_cont_cuda; + case GGML_TYPE_F16: + return convert_unary_cont_cuda; + default: + return nullptr; + } +} + +to_fp16_cuda_t ggml_get_to_fp16_cuda(ggml_type type) { + switch (type) { + case GGML_TYPE_Q1_0: + return dequantize_block_cont_cuda; + case GGML_TYPE_Q2_0: + return dequantize_block_cont_cuda; + case GGML_TYPE_Q4_0: + return dequantize_row_q4_0_cuda; + case GGML_TYPE_Q4_1: + return dequantize_row_q4_1_cuda; + case GGML_TYPE_Q5_0: + return dequantize_block_cont_cuda; + case GGML_TYPE_Q5_1: + return dequantize_block_cont_cuda; + case GGML_TYPE_Q8_0: + if (fp16_available(ggml_cuda_info().devices[ggml_cuda_get_device()].cc)) { + return dequantize_block_q8_0_f16_cuda; + } + return dequantize_block_cont_cuda; + case GGML_TYPE_Q2_K: + return dequantize_row_q2_K_cuda; + case GGML_TYPE_Q3_K: + return dequantize_row_q3_K_cuda; + case GGML_TYPE_Q4_K: + return dequantize_row_q4_K_cuda; + case GGML_TYPE_Q5_K: + return dequantize_row_q5_K_cuda; + case GGML_TYPE_Q6_K: + return dequantize_row_q6_K_cuda; + case GGML_TYPE_IQ2_XXS: + return dequantize_row_iq2_xxs_cuda; + case GGML_TYPE_IQ2_XS: + return dequantize_row_iq2_xs_cuda; + case GGML_TYPE_IQ2_S: + return dequantize_row_iq2_s_cuda; + case GGML_TYPE_IQ3_XXS: + return dequantize_row_iq3_xxs_cuda; + case GGML_TYPE_IQ1_S: + return dequantize_row_iq1_s_cuda; + case GGML_TYPE_IQ1_M: + return dequantize_row_iq1_m_cuda; + case GGML_TYPE_IQ4_NL: + return dequantize_row_iq4_nl_cuda; + case GGML_TYPE_IQ4_XS: + return dequantize_row_iq4_xs_cuda; + case GGML_TYPE_IQ3_S: + return dequantize_row_iq3_s_cuda; + case GGML_TYPE_MXFP4: + return dequantize_row_mxfp4_cuda; + case GGML_TYPE_NVFP4: + return dequantize_row_nvfp4_cuda; + case GGML_TYPE_F32: + return convert_unary_cont_cuda; + case GGML_TYPE_BF16: + return convert_unary_cont_cuda; + default: + return nullptr; + } +} + +to_fp32_cuda_t ggml_get_to_fp32_cuda(ggml_type type) { + switch (type) { + case GGML_TYPE_Q1_0: + return dequantize_block_cont_cuda; + case GGML_TYPE_Q2_0: + return dequantize_block_cont_cuda; + case GGML_TYPE_Q4_0: + return dequantize_row_q4_0_cuda; + case GGML_TYPE_Q4_1: + return dequantize_row_q4_1_cuda; + case GGML_TYPE_Q5_0: + return dequantize_block_cont_cuda; + case GGML_TYPE_Q5_1: + return dequantize_block_cont_cuda; + case GGML_TYPE_Q8_0: + return dequantize_block_cont_cuda; + case GGML_TYPE_Q2_K: + return dequantize_row_q2_K_cuda; + case GGML_TYPE_Q3_K: + return dequantize_row_q3_K_cuda; + case GGML_TYPE_Q4_K: + return dequantize_row_q4_K_cuda; + case GGML_TYPE_Q5_K: + return dequantize_row_q5_K_cuda; + case GGML_TYPE_Q6_K: + return dequantize_row_q6_K_cuda; + case GGML_TYPE_IQ2_XXS: + return dequantize_row_iq2_xxs_cuda; + case GGML_TYPE_IQ2_XS: + return dequantize_row_iq2_xs_cuda; + case GGML_TYPE_IQ2_S: + return dequantize_row_iq2_s_cuda; + case GGML_TYPE_IQ3_XXS: + return dequantize_row_iq3_xxs_cuda; + case GGML_TYPE_IQ1_S: + return dequantize_row_iq1_s_cuda; + case GGML_TYPE_IQ1_M: + return dequantize_row_iq1_m_cuda; + case GGML_TYPE_IQ4_NL: + return dequantize_row_iq4_nl_cuda; + case GGML_TYPE_IQ4_XS: + return dequantize_row_iq4_xs_cuda; + case GGML_TYPE_IQ3_S: + return dequantize_row_iq3_s_cuda; + case GGML_TYPE_MXFP4: + return dequantize_row_mxfp4_cuda; + case GGML_TYPE_NVFP4: + return dequantize_row_nvfp4_cuda; + case GGML_TYPE_F16: + return convert_unary_cont_cuda; + case GGML_TYPE_BF16: + return convert_unary_cont_cuda; + default: + return nullptr; + } +} + +to_fp16_nc_cuda_t ggml_get_to_fp16_nc_cuda(ggml_type type) { + switch (type) { + case GGML_TYPE_F32: + return convert_unary_cuda; + case GGML_TYPE_Q1_0: + return dequantize_block_cuda; + case GGML_TYPE_Q2_0: + return dequantize_block_cuda; + case GGML_TYPE_Q4_0: + return dequantize_block_cuda; + case GGML_TYPE_Q4_1: + return dequantize_block_cuda; + case GGML_TYPE_Q5_0: + return dequantize_block_cuda; + case GGML_TYPE_Q5_1: + return dequantize_block_cuda; + case GGML_TYPE_Q8_0: + return dequantize_block_cuda; + case GGML_TYPE_BF16: + return convert_unary_cuda; + default: + return nullptr; + } +} + +to_bf16_nc_cuda_t ggml_get_to_bf16_nc_cuda(ggml_type type) { + switch (type) { + case GGML_TYPE_F32: + return convert_unary_cuda; + case GGML_TYPE_Q1_0: + return dequantize_block_cuda; + case GGML_TYPE_Q2_0: + return dequantize_block_cuda; + case GGML_TYPE_Q4_0: + return dequantize_block_cuda; + case GGML_TYPE_Q4_1: + return dequantize_block_cuda; + case GGML_TYPE_Q5_0: + return dequantize_block_cuda; + case GGML_TYPE_Q5_1: + return dequantize_block_cuda; + case GGML_TYPE_Q8_0: + return dequantize_block_cuda; + case GGML_TYPE_F16: + return convert_unary_cuda; + default: + return nullptr; + } +} + +to_fp32_nc_cuda_t ggml_get_to_fp32_nc_cuda(ggml_type type) { + switch (type) { + case GGML_TYPE_F16: + return convert_unary_cuda; + case GGML_TYPE_Q1_0: + return dequantize_block_cuda; + case GGML_TYPE_Q2_0: + return dequantize_block_cuda; + case GGML_TYPE_Q4_0: + return dequantize_block_cuda; + case GGML_TYPE_Q4_1: + return dequantize_block_cuda; + case GGML_TYPE_Q5_0: + return dequantize_block_cuda; + case GGML_TYPE_Q5_1: + return dequantize_block_cuda; + case GGML_TYPE_Q8_0: + return dequantize_block_cuda; + case GGML_TYPE_BF16: + return convert_unary_cuda; + default: + return nullptr; + } +} diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/convert.cuh b/python/freetoken/kernel/csrc/ggml/ggml-cuda/convert.cuh new file mode 100644 index 0000000000..91255ccf31 --- /dev/null +++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/convert.cuh @@ -0,0 +1,66 @@ +#pragma once +#include "common.cuh" + +#define CUDA_DEQUANTIZE_BLOCK_SIZE 256 + +template +using to_t_cuda_t = void (*)(const void * x, T * y, int64_t k, cudaStream_t stream); + +typedef to_t_cuda_t to_fp32_cuda_t; +typedef to_t_cuda_t to_fp16_cuda_t; +typedef to_t_cuda_t to_bf16_cuda_t; + +to_fp16_cuda_t ggml_get_to_fp16_cuda(ggml_type type); + +to_bf16_cuda_t ggml_get_to_bf16_cuda(ggml_type type); + +to_fp32_cuda_t ggml_get_to_fp32_cuda(ggml_type type); + +// TODO more general support for non-contiguous inputs + +template +using to_t_nc_cuda_t = void (*)(const void * x, T * y, + int64_t ne00, int64_t ne01, int64_t ne02, int64_t ne03, + int64_t s01, int64_t s02, int64_t s03, cudaStream_t stream); + +typedef to_t_nc_cuda_t to_fp32_nc_cuda_t; +typedef to_t_nc_cuda_t to_fp16_nc_cuda_t; +typedef to_t_nc_cuda_t to_bf16_nc_cuda_t; + +to_fp32_nc_cuda_t ggml_get_to_fp32_nc_cuda(ggml_type type); +to_fp16_nc_cuda_t ggml_get_to_fp16_nc_cuda(ggml_type type); +to_bf16_nc_cuda_t ggml_get_to_bf16_nc_cuda(ggml_type type); + +template + __host__ __device__ inline dst_t ggml_cuda_cast(src_t x) { + if constexpr (std::is_same_v) { + return x; + } else if constexpr(std::is_same_v) { + return __float2bfloat16(float(x)); + } else if constexpr(std::is_same_v) { + return __bfloat162float(x); + } else if constexpr(std::is_same_v && std::is_same_v) { + return __float22half2_rn(x); + } else if constexpr(std::is_same_v && std::is_same_v) { +#ifdef GGML_USE_HIP + return make_float2(__bfloat162float(__low2bfloat16(x)), __bfloat162float(__high2bfloat16(x))); +#else +#if !defined(GGML_USE_MUSA) && __CUDA_ARCH__ >= 800 + return __bfloat1622float2(x); +#else + return make_float2(__bfloat162float(x.x), __bfloat162float(x.y)); +#endif // !defined(GGML_USE_MUSA) && __CUDA_ARCH__ >= 800 +#endif // GGML_USE_HIP + } else if constexpr(std::is_same_v && std::is_same_v) { + // bypass compile error on cuda 12.0.1 +#ifdef GGML_USE_HIP + return __float22bfloat162_rn(x); +#else + return {x.x, x.y}; +#endif // GGML_USE_HIP + } else if constexpr(std::is_same_v) { + return int32_t(x); + } else { + return float(x); + } +} diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/dequantize.cuh b/python/freetoken/kernel/csrc/ggml/ggml-cuda/dequantize.cuh new file mode 100644 index 0000000000..651524be7b --- /dev/null +++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/dequantize.cuh @@ -0,0 +1,452 @@ +#include "common.cuh" +#include "convert.cuh" + +static __device__ __forceinline__ void dequantize_q1_0(const void * vx, const int64_t ib, const int iqs, float2 & v){ + const block_q1_0 * x = (const block_q1_0 *) vx; + + const float d = x[ib].d; + + const int bit_index_0 = iqs; + const int bit_index_1 = iqs + 1; + + const int byte_index_0 = bit_index_0 / 8; + const int bit_offset_0 = bit_index_0 % 8; + + const int byte_index_1 = bit_index_1 / 8; + const int bit_offset_1 = bit_index_1 % 8; + + // Extract bits: 1 = +d, 0 = -d (branchless) + const int bit_0 = (x[ib].qs[byte_index_0] >> bit_offset_0) & 1; + const int bit_1 = (x[ib].qs[byte_index_1] >> bit_offset_1) & 1; + + v.x = (2*bit_0 - 1) * d; + v.y = (2*bit_1 - 1) * d; +} + +static __device__ __forceinline__ void dequantize_q2_0(const void * vx, const int64_t ib, const int iqs, float2 & v){ + const block_q2_0 * x = (const block_q2_0 *) vx; + + const float d = x[ib].d; + + // Q2_0: 2 bits per element, 4 elements per byte. + // Stored code c in {0,1,2,3} maps to symbol s = c - 1 in {-1, 0, +1, +2}. + const int byte_index_0 = iqs / 4; + const int bit_offset_0 = (iqs % 4) * 2; + + const int byte_index_1 = (iqs + 1) / 4; + const int bit_offset_1 = ((iqs + 1) % 4) * 2; + + const int c0 = (x[ib].qs[byte_index_0] >> bit_offset_0) & 0x3; + const int c1 = (x[ib].qs[byte_index_1] >> bit_offset_1) & 0x3; + + v.x = (c0 - 1) * d; + v.y = (c1 - 1) * d; +} + +static __device__ __forceinline__ void dequantize_q4_0(const void * vx, const int64_t ib, const int iqs, float2 & v){ + const block_q4_0 * x = (const block_q4_0 *) vx; + + const float d = x[ib].d; + + const int vui = x[ib].qs[iqs]; + + v.x = vui & 0xF; + v.y = vui >> 4; + + v.x = (v.x - 8.0f) * d; + v.y = (v.y - 8.0f) * d; +} + +static __device__ __forceinline__ void dequantize_q4_1(const void * vx, const int64_t ib, const int iqs, float2 & v){ + const block_q4_1 * x = (const block_q4_1 *) vx; + + const float2 dm = __half22float2(x[ib].dm); + + const int vui = x[ib].qs[iqs]; + + v.x = vui & 0xF; + v.y = vui >> 4; + + v.x = (v.x * dm.x) + dm.y; + v.y = (v.y * dm.x) + dm.y; +} + +static __device__ __forceinline__ void dequantize_q5_0(const void * vx, const int64_t ib, const int iqs, float2 & v){ + const block_q5_0 * x = (const block_q5_0 *) vx; + + const float d = x[ib].d; + + uint32_t qh; + memcpy(&qh, x[ib].qh, sizeof(qh)); + + const int xh_0 = ((qh >> (iqs + 0)) << 4) & 0x10; + const int xh_1 = ((qh >> (iqs + 12)) ) & 0x10; + + v.x = ((x[ib].qs[iqs] & 0xf) | xh_0); + v.y = ((x[ib].qs[iqs] >> 4) | xh_1); + + v.x = (v.x - 16.0f) * d; + v.y = (v.y - 16.0f) * d; +} + +static __device__ __forceinline__ void dequantize_q5_1(const void * vx, const int64_t ib, const int iqs, float2 & v){ + const block_q5_1 * x = (const block_q5_1 *) vx; + + const float2 dm = __half22float2(x[ib].dm); + + uint32_t qh; + memcpy(&qh, x[ib].qh, sizeof(qh)); + + const int xh_0 = ((qh >> (iqs + 0)) << 4) & 0x10; + const int xh_1 = ((qh >> (iqs + 12)) ) & 0x10; + + v.x = ((x[ib].qs[iqs] & 0xf) | xh_0); + v.y = ((x[ib].qs[iqs] >> 4) | xh_1); + + v.x = (v.x * dm.x) + dm.y; + v.y = (v.y * dm.x) + dm.y; +} + +static __device__ __forceinline__ void dequantize_q8_0(const void * vx, const int64_t ib, const int iqs, float2 & v){ + const block_q8_0 * x = (const block_q8_0 *) vx; + + const float d = x[ib].d; + + v.x = x[ib].qs[iqs + 0]; + v.y = x[ib].qs[iqs + 1]; + + v.x *= d; + v.y *= d; +} + +//================================== k-quants + +// Each call dequantizes one super-block of QK_K values into y using the +// thread layout of the caller: 32 threads for q4_K, 64 threads otherwise. + +template +static __device__ __forceinline__ void dequantize_q2_K(const void * vx, const int64_t ib, dst_t * yy, const int tid) { + const block_q2_K * x = (const block_q2_K *) vx; + + const int64_t n = tid/32; + const int64_t l = tid - 32*n; + const int64_t is = 8*n + l/16; + + const uint8_t q = x[ib].qs[32*n + l]; + dst_t * y = yy + 128*n; + + float dall = __low2half(x[ib].dm); + float dmin = __high2half(x[ib].dm); + y[l+ 0] = ggml_cuda_cast(dall * (x[ib].scales[is+0] & 0xF) * ((q >> 0) & 3) - dmin * (x[ib].scales[is+0] >> 4)); + y[l+32] = ggml_cuda_cast(dall * (x[ib].scales[is+2] & 0xF) * ((q >> 2) & 3) - dmin * (x[ib].scales[is+2] >> 4)); + y[l+64] = ggml_cuda_cast(dall * (x[ib].scales[is+4] & 0xF) * ((q >> 4) & 3) - dmin * (x[ib].scales[is+4] >> 4)); + y[l+96] = ggml_cuda_cast(dall * (x[ib].scales[is+6] & 0xF) * ((q >> 6) & 3) - dmin * (x[ib].scales[is+6] >> 4)); +} + +template +static __device__ __forceinline__ void dequantize_q3_K(const void * vx, const int64_t ib, dst_t * yy, const int tid) { + const block_q3_K * x = (const block_q3_K *) vx; + + const int64_t r = tid/4; + const int64_t t = r/2; + const int64_t is0 = r%2; + const int64_t l0 = 16*is0 + 4*(tid%4); + const int64_t n = t / 4; + const int64_t j = t - 4*n; + + uint8_t m = 1 << (4*n + j); + int64_t is = 8*n + 2*j + is0; + int shift = 2*j; + + int8_t us = is < 4 ? (x[ib].scales[is-0] & 0xF) | (((x[ib].scales[is+8] >> 0) & 3) << 4) : + is < 8 ? (x[ib].scales[is-0] & 0xF) | (((x[ib].scales[is+4] >> 2) & 3) << 4) : + is < 12 ? (x[ib].scales[is-8] >> 4) | (((x[ib].scales[is+0] >> 4) & 3) << 4) : + (x[ib].scales[is-8] >> 4) | (((x[ib].scales[is-4] >> 6) & 3) << 4); + float d_all = x[ib].d; + float dl = d_all * (us - 32); + + dst_t * y = yy + 128*n + 32*j; + const uint8_t * q = x[ib].qs + 32*n; + const uint8_t * hm = x[ib].hmask; + + for (int l = l0; l < l0+4; ++l) { + y[l] = ggml_cuda_cast(dl * ((int8_t)((q[l] >> shift) & 3) - ((hm[l] & m) ? 0 : 4))); + } +} + +static inline __device__ void get_scale_min_k4(int j, const uint8_t * q, uint8_t & d, uint8_t & m) { + if (j < 4) { + d = q[j] & 63; m = q[j + 4] & 63; + } else { + d = (q[j+4] & 0xF) | ((q[j-4] >> 6) << 4); + m = (q[j+4] >> 4) | ((q[j-0] >> 6) << 4); + } +} + +template +static __device__ __forceinline__ void dequantize_q4_K(const void * vx, const int64_t ib, dst_t * yy, const int tid) { + const block_q4_K * x = (const block_q4_K *) vx; + + // assume 32 threads + const int64_t il = tid/8; + const int64_t ir = tid%8; + const int64_t is = 2*il; + const int64_t n = 4; + + dst_t * y = yy + 64*il + n*ir; + + const float dall = __low2half(x[ib].dm); + const float dmin = __high2half(x[ib].dm); + + const uint8_t * q = x[ib].qs + 32*il + n*ir; + + uint8_t sc, m; + get_scale_min_k4(is + 0, x[ib].scales, sc, m); + const float d1 = dall * sc; const float m1 = dmin * m; + get_scale_min_k4(is + 1, x[ib].scales, sc, m); + const float d2 = dall * sc; const float m2 = dmin * m; + for (int l = 0; l < n; ++l) { + y[l + 0] = ggml_cuda_cast(d1 * (q[l] & 0xF) - m1); + y[l +32] = ggml_cuda_cast(d2 * (q[l] >> 4) - m2); + } +} + +template +static __device__ __forceinline__ void dequantize_q5_K(const void * vx, const int64_t ib, dst_t * yy, const int tid) { + const block_q5_K * x = (const block_q5_K *) vx; + + // assume 64 threads - this is very slightly better than the one below + const int64_t il = tid/16; // il is in 0...3 + const int64_t ir = tid%16; // ir is in 0...15 + const int64_t is = 2*il; // is is in 0...6 + + dst_t * y = yy + 64*il + 2*ir; + + const float dall = __low2half(x[ib].dm); + const float dmin = __high2half(x[ib].dm); + + const uint8_t * ql = x[ib].qs + 32*il + 2*ir; + const uint8_t * qh = x[ib].qh + 2*ir; + + uint8_t sc, m; + get_scale_min_k4(is + 0, x[ib].scales, sc, m); + const float d1 = dall * sc; const float m1 = dmin * m; + get_scale_min_k4(is + 1, x[ib].scales, sc, m); + const float d2 = dall * sc; const float m2 = dmin * m; + + uint8_t hm = 1 << (2*il); + y[ 0] = ggml_cuda_cast(d1 * ((ql[ 0] & 0xF) + (qh[ 0] & hm ? 16 : 0)) - m1); + y[ 1] = ggml_cuda_cast(d1 * ((ql[ 1] & 0xF) + (qh[ 1] & hm ? 16 : 0)) - m1); + hm <<= 1; + y[32] = ggml_cuda_cast(d2 * ((ql[ 0] >> 4) + (qh[ 0] & hm ? 16 : 0)) - m2); + y[33] = ggml_cuda_cast(d2 * ((ql[ 1] >> 4) + (qh[ 1] & hm ? 16 : 0)) - m2); +} + +template +static __device__ __forceinline__ void dequantize_q6_K(const void * vx, const int64_t ib, dst_t * yy, const int tid) { + const block_q6_K * x = (const block_q6_K *) vx; + + // assume 64 threads - this is very slightly better than the one below + const int64_t ip = tid/32; // ip is 0 or 1 + const int64_t il = tid - 32*ip; // 0...32 + const int64_t is = 8*ip + il/16; + + dst_t * y = yy + 128*ip + il; + + const float d = x[ib].d; + + const uint8_t * ql = x[ib].ql + 64*ip + il; + const uint8_t qh = x[ib].qh[32*ip + il]; + const int8_t * sc = x[ib].scales + is; + + y[ 0] = ggml_cuda_cast(d * sc[0] * ((int8_t)((ql[ 0] & 0xF) | (((qh >> 0) & 3) << 4)) - 32)); + y[32] = ggml_cuda_cast(d * sc[2] * ((int8_t)((ql[32] & 0xF) | (((qh >> 2) & 3) << 4)) - 32)); + y[64] = ggml_cuda_cast(d * sc[4] * ((int8_t)((ql[ 0] >> 4) | (((qh >> 4) & 3) << 4)) - 32)); + y[96] = ggml_cuda_cast(d * sc[6] * ((int8_t)((ql[32] >> 4) | (((qh >> 6) & 3) << 4)) - 32)); +} + +//================================== i-quants + +// Each call dequantizes one super-block of QK_K values into y with 32 +// threads; iq4_nl packs QK_K/QK4_NL sub-blocks per super-block. + +template +static __device__ __forceinline__ void dequantize_iq2_xxs(const void * vx, const int64_t ibs, dst_t * yy, const int tid) { + + const block_iq2_xxs * x = (const block_iq2_xxs *) vx; + + const int64_t il = tid/8; // 0...3 + const int64_t ib = tid%8; // 0...7 + dst_t * y = yy + 32*ib + 8*il; + const uint16_t * q2 = x[ibs].qs + 4*ib; + const uint8_t * aux8 = (const uint8_t *)q2; + const uint8_t * grid = (const uint8_t *)(iq2xxs_grid + aux8[il]); + const uint32_t aux32 = q2[2] | (q2[3] << 16); + const float d = (float)x[ibs].d * (0.5f + (aux32 >> 28)) * 0.25f; + const uint8_t signs = ksigns_iq2xs[(aux32 >> 7*il) & 127]; + for (int j = 0; j < 8; ++j) { + y[j] = ggml_cuda_cast(d * grid[j] * (signs & kmask_iq2xs[j] ? -1.f : 1.f)); + } +} + +template +static __device__ __forceinline__ void dequantize_iq2_xs(const void * vx, const int64_t ibs, dst_t * yy, const int tid) { + + const block_iq2_xs * x = (const block_iq2_xs *) vx; + + const int64_t il = tid/8; // 0...3 + const int64_t ib = tid%8; // 0...7 + dst_t * y = yy + 32*ib + 8*il; + const uint16_t * q2 = x[ibs].qs + 4*ib; + const uint8_t * grid = (const uint8_t *)(iq2xs_grid + (q2[il] & 511)); + const float d = (float)x[ibs].d * (0.5f + ((x[ibs].scales[ib] >> 4*(il/2)) & 0xf)) * 0.25f; + const uint8_t signs = ksigns_iq2xs[q2[il] >> 9]; + for (int j = 0; j < 8; ++j) { + y[j] = ggml_cuda_cast(d * grid[j] * (signs & kmask_iq2xs[j] ? -1.f : 1.f)); + } +} + +template +static __device__ __forceinline__ void dequantize_iq2_s(const void * vx, const int64_t ibs, dst_t * yy, const int tid) { + + const block_iq2_s * x = (const block_iq2_s *) vx; + + const int64_t il = tid/8; // 0...3 + const int64_t ib = tid%8; // 0...7 + dst_t * y = yy + 32*ib + 8*il; + const uint8_t * grid = (const uint8_t *)(iq2s_grid + (x[ibs].qs[4*ib+il] | ((x[ibs].qh[ib] << (8-2*il)) & 0x300))); + const float d = (float)x[ibs].d * (0.5f + ((x[ibs].scales[ib] >> 4*(il/2)) & 0xf)) * 0.25f; + const uint8_t signs = x[ibs].qs[QK_K/8+4*ib+il]; + for (int j = 0; j < 8; ++j) { + y[j] = ggml_cuda_cast(d * grid[j] * (signs & kmask_iq2xs[j] ? -1.f : 1.f)); + } +} + +template +static __device__ __forceinline__ void dequantize_iq3_xxs(const void * vx, const int64_t ibs, dst_t * yy, const int tid) { + + const block_iq3_xxs * x = (const block_iq3_xxs *) vx; + + const int64_t il = tid/8; // 0...3 + const int64_t ib = tid%8; // 0...7 + dst_t * y = yy + 32*ib + 8*il; + const uint8_t * q3 = x[ibs].qs + 8*ib; + const uint16_t * gas = (const uint16_t *)(x[ibs].qs + QK_K/4) + 2*ib; + const uint8_t * grid1 = (const uint8_t *)(iq3xxs_grid + q3[2*il+0]); + const uint8_t * grid2 = (const uint8_t *)(iq3xxs_grid + q3[2*il+1]); + const uint32_t aux32 = gas[0] | (gas[1] << 16); + const float d = (float)x[ibs].d * (0.5f + (aux32 >> 28)) * 0.5f; + const uint8_t signs = ksigns_iq2xs[(aux32 >> 7*il) & 127]; + for (int j = 0; j < 4; ++j) { + y[j+0] = ggml_cuda_cast(d * grid1[j] * (signs & kmask_iq2xs[j+0] ? -1.f : 1.f)); + y[j+4] = ggml_cuda_cast(d * grid2[j] * (signs & kmask_iq2xs[j+4] ? -1.f : 1.f)); + } +} + +template +static __device__ __forceinline__ void dequantize_iq3_s(const void * vx, const int64_t ibs, dst_t * yy, const int tid) { + + const block_iq3_s * x = (const block_iq3_s *) vx; + + const int64_t il = tid/8; // 0...3 + const int64_t ib = tid%8; // 0...7 + dst_t * y = yy + 32*ib + 8*il; + const uint8_t * qs = x[ibs].qs + 8*ib; + const uint8_t * grid1 = (const uint8_t *)(iq3s_grid + (qs[2*il+0] | ((x[ibs].qh[ib] << (8-2*il)) & 256))); + const uint8_t * grid2 = (const uint8_t *)(iq3s_grid + (qs[2*il+1] | ((x[ibs].qh[ib] << (7-2*il)) & 256))); + const float d = (float)x[ibs].d * (1 + 2*((x[ibs].scales[ib/2] >> 4*(ib%2)) & 0xf)); + const uint8_t signs = x[ibs].signs[4*ib + il]; + for (int j = 0; j < 4; ++j) { + y[j+0] = ggml_cuda_cast(d * grid1[j] * (signs & kmask_iq2xs[j+0] ? -1.f : 1.f)); + y[j+4] = ggml_cuda_cast(d * grid2[j] * (signs & kmask_iq2xs[j+4] ? -1.f : 1.f)); + } +} + +template +static __device__ __forceinline__ void dequantize_iq1_s(const void * vx, const int64_t ibs, dst_t * yy, const int tid) { + + const block_iq1_s * x = (const block_iq1_s *) vx; + + const int64_t il = tid/8; // 0...3 + const int64_t ib = tid%8; // 0...7 + dst_t * y = yy + 32*ib + 8*il; + const float delta = x[ibs].qh[ib] & 0x8000 ? -1 - IQ1S_DELTA : -1 + IQ1S_DELTA; + const float d = (float)x[ibs].d * (2*((x[ibs].qh[ib] >> 12) & 7) + 1); + uint32_t grid32[2]; const int8_t * q = (const int8_t *)grid32; + grid32[0] = iq1s_grid_gpu[x[ibs].qs[4*ib+il] | (((x[ibs].qh[ib] >> 3*il) & 7) << 8)]; + grid32[1] = (grid32[0] >> 4) & 0x0f0f0f0f; + grid32[0] &= 0x0f0f0f0f; + for (int j = 0; j < 8; ++j) { + y[j] = ggml_cuda_cast(d * (q[j] + delta)); + } +} + +template +static __device__ __forceinline__ void dequantize_iq1_m(const void * vx, const int64_t ibs, dst_t * yy, const int tid) { + + const block_iq1_m * x = (const block_iq1_m *) vx; + + const int64_t il = tid/8; // 0...3 + const int64_t ib = tid%8; // 0...7 + dst_t * y = yy + 32*ib + 8*il; + const uint16_t * sc = (const uint16_t *)x[ibs].scales; + iq1m_scale_t scale; + scale.u16 = (sc[0] >> 12) | ((sc[1] >> 8) & 0x00f0) | ((sc[2] >> 4) & 0x0f00) | (sc[3] & 0xf000); + const int64_t ib16 = 2*ib + il/2; // sc[ib16/4] >> 3*(ib16%4) -> sc[ib/2] >> 3*((2*ib+il/2)%4); + const float d = (float)scale.f16 * (2*((sc[ib16/4] >> 3*(ib16%4)) & 0x7) + 1); + const float delta = x[ibs].qh[2*ib+il/2] & (0x08 << 4*(il%2)) ? -1 - IQ1M_DELTA : -1 + IQ1M_DELTA; + uint32_t grid32[2]; const int8_t * q = (const int8_t *)grid32; + grid32[0] = iq1s_grid_gpu[x[ibs].qs[4*ib+il] | (((x[ibs].qh[2*ib+il/2] >> 4*(il%2)) & 7) << 8)]; + grid32[1] = (grid32[0] >> 4) & 0x0f0f0f0f; + grid32[0] &= 0x0f0f0f0f; + for (int j = 0; j < 8; ++j) { + y[j] = ggml_cuda_cast(d * (q[j] + delta)); + } +} + +template +static __device__ __forceinline__ void dequantize_iq4_nl(const void * vx, const int64_t ibs, dst_t * yy, const int tid) { + + const block_iq4_nl * x = (const block_iq4_nl *) vx + ibs*(QK_K/QK4_NL); + + const int64_t il = tid/8; // 0...3 + const int64_t ib = tid%8; // 0...7 + dst_t * y = yy + 32*ib + 4*il; + const uint8_t * q4 = x[ib].qs + 4*il; + const float d = (float)x[ib].d; + for (int j = 0; j < 4; ++j) { + y[j+ 0] = ggml_cuda_cast(d * kvalues_iq4nl[q4[j] & 0xf]); + y[j+16] = ggml_cuda_cast(d * kvalues_iq4nl[q4[j] >> 4]); + } +} + +template +static __device__ __forceinline__ void dequantize_iq4_xs(const void * vx, const int64_t ibs, dst_t * yy, const int tid) { + const block_iq4_xs * x = (const block_iq4_xs *)vx; + + const int64_t il = tid/8; // 0...3 + const int64_t ib = tid%8; // 0...7 + dst_t * y = yy + 32*ib + 4*il; + const uint8_t * q4 = x[ibs].qs + 16*ib + 4*il; + const float d = (float)x[ibs].d * ((((x[ibs].scales_l[ib/2] >> 4*(ib%2)) & 0xf) | (((x[ibs].scales_h >> 2*ib) & 3) << 4)) - 32); + for (int j = 0; j < 4; ++j) { + y[j+ 0] = ggml_cuda_cast(d * kvalues_iq4nl[q4[j] & 0xf]); + y[j+16] = ggml_cuda_cast(d * kvalues_iq4nl[q4[j] >> 4]); + } +} + +template +static __device__ __forceinline__ void dequantize_mxfp4(const void * vx, const int64_t ibs, dst_t * yy, const int tid) { + + const block_mxfp4 * x = (const block_mxfp4 *) vx + ibs*(QK_K/QK_MXFP4); + + const int64_t il = tid/8; // 0...3 + const int64_t ib = tid%8; // 0...7 + dst_t * y = yy + 32*ib + 4*il; + const uint8_t * q4 = x[ib].qs + 4*il; + const float d = ggml_cuda_e8m0_to_fp32(x[ib].e); + for (int j = 0; j < 4; ++j) { + y[j+ 0] = ggml_cuda_cast(d * kvalues_mxfp4[q4[j] & 0xf]*0.5f); + y[j+16] = ggml_cuda_cast(d * kvalues_mxfp4[q4[j] >> 4]*0.5f); + } +} diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/getrows.cu b/python/freetoken/kernel/csrc/ggml/ggml-cuda/getrows.cu new file mode 100644 index 0000000000..d97e8fdb0f --- /dev/null +++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/getrows.cu @@ -0,0 +1,488 @@ +#include "getrows.cuh" +#include "dequantize.cuh" +#include "convert.cuh" + +template +static __global__ void k_get_rows( + const void * __restrict__ src0, const int32_t * __restrict__ src1, dst_t * __restrict__ dst, + const int64_t ne00, /*const int64_t ne01, const int64_t ne02, const int64_t ne03,*/ + /*const int64_t ne10,*/ const int64_t ne11, const uint3 ne12_fdv, /*const int64_t ne13,*/ + /*const size_t s0,*/ const size_t s1, const size_t s2, const size_t s3, + /*const size_t nb00,*/ const size_t nb01, const size_t nb02, const size_t nb03, + const size_t s10, const size_t s11, const size_t s12/*, const size_t s13*/) { + + ggml_cuda_pdl_sync(); + for (int64_t z = blockIdx.z; z < ne11*(int64_t)ne12_fdv.z; z += gridDim.z) { + for (int64_t i00 = 2*(blockIdx.y*blockDim.x + threadIdx.x); i00 < ne00; i00 += gridDim.y*blockDim.x) { + // The x and y dimensions of the grid are swapped because the maximum allowed grid size for x is higher. + const int i10 = blockIdx.x; + const uint2 dm = fast_div_modulo((uint32_t)z, ne12_fdv); + const int i11 = dm.x; + const int i12 = dm.y; + + const int i01 = src1[i10*s10 + i11*s11 + i12*s12]; + + dst_t * dst_row = dst + i10*s1 + i11*s2 + i12*s3; + const void * src0_row = (const char *) src0 + i01*nb01 + i11*nb02 + i12*nb03; + + const int ib = i00/qk; // block index + const int iqs = (i00%qk)/qr; // quant index + const int iybs = i00 - i00%qk; // dst block start index + const int y_offset = qr == 1 ? 1 : qk/2; + + // dequantize + float2 v; + dequantize_kernel(src0_row, ib, iqs, v); + + dst_row[iybs + iqs + 0] = ggml_cuda_cast(v.x); + dst_row[iybs + iqs + y_offset] = ggml_cuda_cast(v.y); + } + } +} + +template dequantize_kq> +static __global__ void k_get_rows_kq( + const void * __restrict__ src0, const int32_t * __restrict__ src1, dst_t * __restrict__ dst, + const int64_t ne00, /*const int64_t ne01, const int64_t ne02, const int64_t ne03,*/ + /*const int64_t ne10,*/ const int64_t ne11, const uint3 ne12_fdv, /*const int64_t ne13,*/ + /*const size_t s0,*/ const size_t s1, const size_t s2, const size_t s3, + /*const size_t nb00,*/ const size_t nb01, const size_t nb02, const size_t nb03, + const size_t s10, const size_t s11, const size_t s12/*, const size_t s13*/) { + + ggml_cuda_pdl_sync(); + const int64_t nsb = ne00/QK_K; // super-blocks per row + for (int64_t z = blockIdx.z; z < ne11*(int64_t)ne12_fdv.z; z += gridDim.z) { + // The x and y dimensions of the grid are swapped because the maximum allowed grid size for x is higher. + const int i10 = blockIdx.x; + const uint2 dm = fast_div_modulo((uint32_t)z, ne12_fdv); + const int i11 = dm.x; + const int i12 = dm.y; + + const int i01 = src1[i10*s10 + i11*s11 + i12*s12]; + + dst_t * dst_row = dst + i10*s1 + i11*s2 + i12*s3; + const void * src0_row = (const char *) src0 + i01*nb01 + i11*nb02 + i12*nb03; + + for (int64_t ib = blockIdx.y; ib < nsb; ib += gridDim.y) { + dequantize_kq(src0_row, ib, dst_row + ib*QK_K, threadIdx.x); + } + } +} + +template +static __global__ void k_get_rows_float( + const src0_t * src0_ptr, const int32_t * src1_ptr, dst_t * dst_ptr, + const int64_t ne00, /*const int64_t ne01, const int64_t ne02, const int64_t ne03,*/ + /*const int64_t ne10,*/ const int64_t ne11, const uint3 ne12_fdv, /*const int64_t ne13,*/ + /*const size_t s0,*/ const size_t s1, const size_t s2, const size_t s3, + /*const size_t nb00,*/ const size_t nb01, const size_t nb02, const size_t nb03, + const size_t s10, const size_t s11, const size_t s12/*, const size_t s13*/) { + + ggml_cuda_pdl_lc(); + const src0_t * GGML_CUDA_RESTRICT src0 = src0_ptr; + const int32_t * GGML_CUDA_RESTRICT src1 = src1_ptr; + dst_t * GGML_CUDA_RESTRICT dst = dst_ptr; + ggml_cuda_pdl_sync(); + for (int64_t z = blockIdx.z; z < ne11*(int64_t)ne12_fdv.z; z += gridDim.z) { + // The x and y dimensions of the grid are swapped because the maximum allowed grid size for x is higher. + const int i10 = blockIdx.x; + const uint2 dm = fast_div_modulo((uint32_t)z, ne12_fdv); + const int i11 = dm.x; + const int i12 = dm.y; + + const int i01 = src1[i10*s10 + i11*s11 + i12*s12]; + + dst_t * GGML_CUDA_RESTRICT dst_row = dst + i10*s1 + i11*s2 + i12*s3; + const src0_t * GGML_CUDA_RESTRICT src0_row = (const src0_t *)((const char *) src0 + i01*nb01 + i11*nb02 + i12*nb03); + + for (int64_t i00 = blockIdx.y*blockDim.x + threadIdx.x; i00 < ne00; i00 += gridDim.y*blockDim.x) { + dst_row[i00] = ggml_cuda_cast(src0_row[i00]); + } + } +} + +template +static __global__ void k_get_rows_float_vec( + const dst_t * src0_ptr, const int32_t * src1_ptr, dst_t * dst_ptr, + const int64_t ne00v, + const int64_t ne11, const uint3 ne12_fdv, + const size_t s1, const size_t s2, const size_t s3, + const size_t nb01, const size_t nb02, const size_t nb03, + const size_t s10, const size_t s11, const size_t s12) { + + ggml_cuda_pdl_lc(); + ggml_cuda_pdl_sync(); + for (int64_t z = blockIdx.z; z < ne11*(int64_t)ne12_fdv.z; z += gridDim.z) { + const int i10 = blockIdx.x; + const uint2 dm = fast_div_modulo((uint32_t)z, ne12_fdv); + const int i11 = dm.x; + const int i12 = dm.y; + + const int i01 = src1_ptr[i10*s10 + i11*s11 + i12*s12]; + + int4 * GGML_CUDA_RESTRICT dst_row = (int4 *) (dst_ptr + i10*s1 + i11*s2 + i12*s3); + const int4 * GGML_CUDA_RESTRICT src0_row = (const int4 *)((const char *) src0_ptr + i01*nb01 + i11*nb02 + i12*nb03); + + for (int64_t i = blockIdx.y*blockDim.x + threadIdx.x; i < ne00v; i += gridDim.y*blockDim.x) { + dst_row[i] = src0_row[i]; + } + } +} + +template +static __global__ void k_get_rows_back_float( + const grad_t * __restrict__ grad, const int32_t * __restrict__ rows, dst_t * __restrict__ dst, + const int64_t ncols, const int64_t nrows_grad, const int64_t nrows_dst) { + const int col = blockIdx.x*blockDim.x + threadIdx.x; + + if (col >= ncols) { + return; + } + + ggml_cuda_pdl_sync(); + + // grid.y is clamped to the CUDA grid limit, so stride over the destination rows + for (int64_t dst_row = blockIdx.y; dst_row < nrows_dst; dst_row += gridDim.y) { + float sum = 0.0f; + + for (int64_t i = 0; i < nrows_grad; ++i) { + if (rows[i] != dst_row) { + continue; + } + sum += grad[i*ncols + col]; + } + + dst[dst_row*ncols + col] = sum; + } +} + +template +static void get_rows_cuda_q( + const void * src0_d, const int32_t * src1_d, dst_t * dst_d, + const int64_t ne00, const size_t nb01, const size_t nb02, const size_t nb03, + const int64_t ne10, const int64_t ne11, const int64_t ne12, const size_t nb10, const size_t nb11, const size_t nb12, + const size_t nb1, const size_t nb2, const size_t nb3, + cudaStream_t stream) { + const dim3 block_dims(CUDA_GET_ROWS_BLOCK_SIZE, 1, 1); + const int block_num_y = (ne00 + 2*CUDA_GET_ROWS_BLOCK_SIZE - 1) / (2*CUDA_GET_ROWS_BLOCK_SIZE); + const dim3 block_nums(ne10, MIN(block_num_y, UINT16_MAX), MIN(ne11*ne12, UINT16_MAX)); + + // strides in elements + // const size_t s0 = nb0 / sizeof(dst_t); + const size_t s1 = nb1 / sizeof(dst_t); + const size_t s2 = nb2 / sizeof(dst_t); + const size_t s3 = nb3 / sizeof(dst_t); + + const size_t s10 = nb10 / sizeof(int32_t); + const size_t s11 = nb11 / sizeof(int32_t); + const size_t s12 = nb12 / sizeof(int32_t); + // const size_t s13 = nb13 / sizeof(int32_t); + + GGML_ASSERT(ne00 % 2 == 0); + + GGML_ASSERT(ne12 > 0); + GGML_ASSERT(ne11 <= std::numeric_limits::max() / ne12); + const uint3 ne12_fdv = init_fastdiv_values(ne12); + + k_get_rows<<>>( + src0_d, src1_d, dst_d, + ne00, /*ne01, ne02, ne03,*/ + /*ne10,*/ ne11, ne12_fdv, /*ne13,*/ + /* s0,*/ s1, s2, s3, + /* nb00,*/ nb01, nb02, nb03, + s10, s11, s12/*, s13*/); +} + +template dequantize_kq> +static void get_rows_cuda_kq( + const void * src0_d, const int32_t * src1_d, dst_t * dst_d, + const int64_t ne00, const size_t nb01, const size_t nb02, const size_t nb03, + const int64_t ne10, const int64_t ne11, const int64_t ne12, const size_t nb10, const size_t nb11, const size_t nb12, + const size_t nb1, const size_t nb2, const size_t nb3, + cudaStream_t stream) { + GGML_ASSERT(ne00 % QK_K == 0); + const int64_t nsb = ne00/QK_K; + + const dim3 block_dims(block_dim, 1, 1); + const dim3 block_nums(ne10, MIN(nsb, UINT16_MAX), MIN(ne11*ne12, UINT16_MAX)); + + // strides in elements + // const size_t s0 = nb0 / sizeof(dst_t); + const size_t s1 = nb1 / sizeof(dst_t); + const size_t s2 = nb2 / sizeof(dst_t); + const size_t s3 = nb3 / sizeof(dst_t); + + const size_t s10 = nb10 / sizeof(int32_t); + const size_t s11 = nb11 / sizeof(int32_t); + const size_t s12 = nb12 / sizeof(int32_t); + // const size_t s13 = nb13 / sizeof(int32_t); + + GGML_ASSERT(ne12 > 0); + GGML_ASSERT(ne11 <= std::numeric_limits::max() / ne12); + const uint3 ne12_fdv = init_fastdiv_values(ne12); + + k_get_rows_kq<<>>( + src0_d, src1_d, dst_d, + ne00, /*ne01, ne02, ne03,*/ + /*ne10,*/ ne11, ne12_fdv, /*ne13,*/ + /* s0,*/ s1, s2, s3, + /* nb00,*/ nb01, nb02, nb03, + s10, s11, s12/*, s13*/); +} + +template +static void get_rows_cuda_float( + const src0_t * src0_d, const int32_t * src1_d, dst_t * dst_d, + const int64_t ne00, const size_t nb01, const size_t nb02, const size_t nb03, + const int64_t ne10, const int64_t ne11, const int64_t ne12, const size_t nb10, const size_t nb11, const size_t nb12, + const size_t nb1, const size_t nb2, const size_t nb3, + cudaStream_t stream) { + const dim3 block_dims(CUDA_GET_ROWS_BLOCK_SIZE, 1, 1); + + // strides in elements + // const size_t s0 = nb0 / sizeof(dst_t); + const size_t s1 = nb1 / sizeof(dst_t); + const size_t s2 = nb2 / sizeof(dst_t); + const size_t s3 = nb3 / sizeof(dst_t); + + const size_t s10 = nb10 / sizeof(int32_t); + const size_t s11 = nb11 / sizeof(int32_t); + const size_t s12 = nb12 / sizeof(int32_t); + // const size_t s13 = nb13 / sizeof(int32_t); + + GGML_ASSERT(ne12 > 0); + GGML_ASSERT(ne11 <= std::numeric_limits::max() / ne12); + const uint3 ne12_fdv = init_fastdiv_values(ne12); + + if constexpr (std::is_same::value) { + constexpr int VEC = 16 / sizeof(dst_t); + const int64_t ne00v = ne00 / VEC; + const int64_t vec_block_num_y = (ne00v + CUDA_GET_ROWS_BLOCK_SIZE - 1) / CUDA_GET_ROWS_BLOCK_SIZE; + const bool enough_blocks = vec_block_num_y * ne10 * ne11 * ne12 >= 128; + const bool can_vec = VEC > 1 && enough_blocks && + (ne00 % VEC == 0) && + (nb01 % 16 == 0) && (nb02 % 16 == 0) && (nb03 % 16 == 0) && + (nb1 % 16 == 0) && (nb2 % 16 == 0) && (nb3 % 16 == 0) && + (((uintptr_t) src0_d) % 16 == 0) && (((uintptr_t) dst_d) % 16 == 0); + + if (can_vec) { + const int block_num_y = vec_block_num_y; + const dim3 block_nums(ne10, MIN(block_num_y, UINT16_MAX), MIN(ne11*ne12, UINT16_MAX)); + const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params{block_nums, block_dims, 0, stream}; + ggml_cuda_kernel_launch(k_get_rows_float_vec, launch_params, + (const dst_t *) src0_d, src1_d, dst_d, + ne00v, ne11, ne12_fdv, + s1, s2, s3, + nb01, nb02, nb03, + s10, s11, s12); + return; + } + } + + const int block_num_y = (ne00 + CUDA_GET_ROWS_BLOCK_SIZE - 1) / CUDA_GET_ROWS_BLOCK_SIZE; + const dim3 block_nums(ne10, MIN(block_num_y, UINT16_MAX), MIN(ne11*ne12, UINT16_MAX)); + + const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params{block_nums, block_dims, 0, stream}; + ggml_cuda_kernel_launch(k_get_rows_float, launch_params, + src0_d, src1_d, dst_d, + ne00, /*ne01, ne02, ne03,*/ + /*ne10,*/ ne11, ne12_fdv, /*ne13,*/ + /* s0,*/ s1, s2, s3, + /* nb00,*/ nb01, nb02, nb03, + s10, s11, s12/*, s13*/); +} + +template +static void ggml_cuda_get_rows_switch_src0_type( + const void * src0_d, const ggml_type src0_type, const int32_t * src1_d, dst_t * dst_d, + const int64_t ne00, const size_t nb01, const size_t nb02, const size_t nb03, + const int64_t ne10, const int64_t ne11, const int64_t ne12, const size_t nb10, const size_t nb11, const size_t nb12, + const size_t nb1, const size_t nb2, const size_t nb3, + cudaStream_t stream) { + switch (src0_type) { + case GGML_TYPE_F16: + get_rows_cuda_float((const half *) src0_d, src1_d, dst_d, + ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream); + break; + case GGML_TYPE_F32: + get_rows_cuda_float((const float *) src0_d, src1_d, dst_d, + ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream); + break; + case GGML_TYPE_I32: + get_rows_cuda_float((const int32_t *) src0_d, src1_d, dst_d, + ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream); + break; + case GGML_TYPE_BF16: + get_rows_cuda_float((const nv_bfloat16 *) src0_d, src1_d, dst_d, + ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream); + break; + case GGML_TYPE_Q1_0: + get_rows_cuda_q(src0_d, src1_d, dst_d, + ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream); + break; + case GGML_TYPE_Q2_0: + get_rows_cuda_q(src0_d, src1_d, dst_d, + ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream); + break; + case GGML_TYPE_Q4_0: + get_rows_cuda_q(src0_d, src1_d, dst_d, + ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream); + break; + case GGML_TYPE_Q4_1: + get_rows_cuda_q(src0_d, src1_d, dst_d, + ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream); + break; + case GGML_TYPE_Q5_0: + get_rows_cuda_q(src0_d, src1_d, dst_d, + ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream); + break; + case GGML_TYPE_Q5_1: + get_rows_cuda_q(src0_d, src1_d, dst_d, + ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream); + break; + case GGML_TYPE_Q8_0: + get_rows_cuda_q(src0_d, src1_d, dst_d, + ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream); + break; + case GGML_TYPE_Q2_K: + get_rows_cuda_kq<64, dst_t, dequantize_q2_K>(src0_d, src1_d, dst_d, + ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream); + break; + case GGML_TYPE_Q3_K: + get_rows_cuda_kq<64, dst_t, dequantize_q3_K>(src0_d, src1_d, dst_d, + ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream); + break; + case GGML_TYPE_Q4_K: + get_rows_cuda_kq<32, dst_t, dequantize_q4_K>(src0_d, src1_d, dst_d, + ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream); + break; + case GGML_TYPE_Q5_K: + get_rows_cuda_kq<64, dst_t, dequantize_q5_K>(src0_d, src1_d, dst_d, + ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream); + break; + case GGML_TYPE_Q6_K: + get_rows_cuda_kq<64, dst_t, dequantize_q6_K>(src0_d, src1_d, dst_d, + ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream); + break; + case GGML_TYPE_IQ2_XXS: + get_rows_cuda_kq<32, dst_t, dequantize_iq2_xxs>(src0_d, src1_d, dst_d, + ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream); + break; + case GGML_TYPE_IQ2_XS: + get_rows_cuda_kq<32, dst_t, dequantize_iq2_xs>(src0_d, src1_d, dst_d, + ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream); + break; + case GGML_TYPE_IQ2_S: + get_rows_cuda_kq<32, dst_t, dequantize_iq2_s>(src0_d, src1_d, dst_d, + ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream); + break; + case GGML_TYPE_IQ3_XXS: + get_rows_cuda_kq<32, dst_t, dequantize_iq3_xxs>(src0_d, src1_d, dst_d, + ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream); + break; + case GGML_TYPE_IQ3_S: + get_rows_cuda_kq<32, dst_t, dequantize_iq3_s>(src0_d, src1_d, dst_d, + ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream); + break; + case GGML_TYPE_IQ1_S: + get_rows_cuda_kq<32, dst_t, dequantize_iq1_s>(src0_d, src1_d, dst_d, + ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream); + break; + case GGML_TYPE_IQ1_M: + get_rows_cuda_kq<32, dst_t, dequantize_iq1_m>(src0_d, src1_d, dst_d, + ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream); + break; + case GGML_TYPE_IQ4_NL: + get_rows_cuda_kq<32, dst_t, dequantize_iq4_nl>(src0_d, src1_d, dst_d, + ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream); + break; + case GGML_TYPE_IQ4_XS: + get_rows_cuda_kq<32, dst_t, dequantize_iq4_xs>(src0_d, src1_d, dst_d, + ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream); + break; + case GGML_TYPE_MXFP4: + get_rows_cuda_kq<32, dst_t, dequantize_mxfp4>(src0_d, src1_d, dst_d, + ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream); + break; + default: + GGML_ABORT("%s: unsupported src0 type: %s\n", __func__, ggml_type_name(src0_type)); + } +} + +void get_rows_cuda( + const void * src0_d, ggml_type src0_type, const int32_t * src1_d, void * dst_d, ggml_type dst_type, + int64_t ne00, size_t nb01, size_t nb02, size_t nb03, + int64_t ne10, int64_t ne11, int64_t ne12, size_t nb10, size_t nb11, size_t nb12, + size_t nb1, size_t nb2, size_t nb3, + cudaStream_t stream) { + switch (dst_type) { + case GGML_TYPE_F32: + ggml_cuda_get_rows_switch_src0_type(src0_d, src0_type, src1_d, (float *) dst_d, + ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream); + break; + case GGML_TYPE_I32: + ggml_cuda_get_rows_switch_src0_type(src0_d, src0_type, src1_d, (int32_t *) dst_d, + ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream); + break; + case GGML_TYPE_F16: + ggml_cuda_get_rows_switch_src0_type(src0_d, src0_type, src1_d, (half *) dst_d, + ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream); + break; + case GGML_TYPE_BF16: + ggml_cuda_get_rows_switch_src0_type(src0_d, src0_type, src1_d, (nv_bfloat16 *) dst_d, + ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream); + break; + default: + GGML_ABORT("%s: unsupported dst type: %s\n", __func__, ggml_type_name(dst_type)); + } +} + +void ggml_cuda_op_get_rows(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { + const ggml_tensor * src0 = dst->src[0]; + const ggml_tensor * src1 = dst->src[1]; + + cudaStream_t stream = ctx.stream(); + + GGML_TENSOR_BINARY_OP_LOCALS + + GGML_ASSERT(src1->type == GGML_TYPE_I32); + GGML_ASSERT(ne13 == 1); + + GGML_ASSERT(src0->nb[0] == ggml_type_size(src0->type)); + GGML_ASSERT(src1->nb[0] == ggml_type_size(src1->type)); + GGML_ASSERT(dst->nb[0] == ggml_type_size(dst->type)); + + get_rows_cuda(src0->data, src0->type, (const int32_t *) src1->data, dst->data, dst->type, + ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream); +} + +void ggml_cuda_op_get_rows_back(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { + const ggml_tensor * src0 = dst->src[0]; // gradients of forward pass output + const ggml_tensor * src1 = dst->src[1]; // src1 in forward pass + + GGML_TENSOR_BINARY_OP_LOCALS + + const float * src0_d = (const float *) src0->data; + const int32_t * src1_d = (const int32_t *) src1->data; + float * dst_d = (float *) dst->data; + + cudaStream_t stream = ctx.stream(); + + GGML_ASSERT(src0->type == GGML_TYPE_F32); + GGML_ASSERT(src1->type == GGML_TYPE_I32); + GGML_ASSERT(dst->type == GGML_TYPE_F32); + + GGML_ASSERT(ggml_is_contiguous(src0)); + GGML_ASSERT(ggml_is_contiguous(src1)); + GGML_ASSERT(ggml_is_contiguous(dst)); + + GGML_ASSERT(ne02*ne03 == 1); + GGML_ASSERT(ne12*ne13 == 1); + GGML_ASSERT(ne2*ne3 == 1); + + const dim3 block_dims(CUDA_GET_ROWS_BACK_BLOCK_SIZE, 1, 1); + const int block_num_x = (ne00 + CUDA_GET_ROWS_BACK_BLOCK_SIZE - 1) / CUDA_GET_ROWS_BACK_BLOCK_SIZE; + const dim3 block_nums(block_num_x, MIN(ne1, (int64_t)UINT16_MAX), 1); + + k_get_rows_back_float<<>>(src0_d, src1_d, dst_d, ne00, ne10, ne1); +} diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/getrows.cuh b/python/freetoken/kernel/csrc/ggml/ggml-cuda/getrows.cuh new file mode 100644 index 0000000000..3c5bea5f48 --- /dev/null +++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/getrows.cuh @@ -0,0 +1,15 @@ +#include "common.cuh" + +#define CUDA_GET_ROWS_BLOCK_SIZE 256 +#define CUDA_GET_ROWS_BACK_BLOCK_SIZE 256 + +void get_rows_cuda( + const void * src0_d, ggml_type src0_type, const int32_t * src1_d, void * dst_d, ggml_type dst_type, + int64_t ne00, size_t nb01, size_t nb02, size_t nb03, + int64_t ne10, int64_t ne11, int64_t ne12, size_t nb10, size_t nb11, size_t nb12, + size_t nb1, size_t nb2, size_t nb3, + cudaStream_t stream); + +void ggml_cuda_op_get_rows(ggml_backend_cuda_context & ctx, ggml_tensor * dst); + +void ggml_cuda_op_get_rows_back(ggml_backend_cuda_context & ctx, ggml_tensor * dst); diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/mma.cuh b/python/freetoken/kernel/csrc/ggml/ggml-cuda/mma.cuh new file mode 100644 index 0000000000..93a51aa0be --- /dev/null +++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/mma.cuh @@ -0,0 +1,1594 @@ +#pragma once +// This file contains primitives that expose the tensor core PTX instructions for CUDA code. +// The primitives can be used in a similar way as the nvcuda::wmma interface but with a well-defined memory layout. +// The documentation for the PTX instructions can be found under: +// https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#matrix-multiply-accumulate-operation-using-mma-instruction +// +// Like with nvcuda::wmma there are three types of matrix tiles: A, B, and C with A @ B = C. +// A is a row-major matrix with shape M x K. +// B is a column-major matrix with shape K x N. +// C is a column-major matrix with shape M x N. +// A, B, and C are represented using the same fundamental data type: a row-major matrix with I rows and J columns. +// Note that J is measured in physical 32 bit elements instead of logical elements. +// The methods get_i and get_j can be used to get the physical 32 bit index of the lth element of a thread within a tile. +// All matrix tiles have ne physical 32 bit elements per warp. +// +// As described in the PTX documentation, all pointers for load_ldmatrix must be to shared memory and aligned to 16 bytes. +// The API in this file also assumes that the pointers for load_generic are aligned to 16 bytes, unaligned pointers are considered undefined behavior. + +#include "common.cuh" + +// On Volta each warp is doing 4 8x8 mma operations in parallel. +// The basic memory layout for a 32x8 output tile is to stack 4 input tiles in I direction and to mirror the B tile. +// However, the i indices in this file are by default permuted to simplify the index calculations. +// #define GGML_CUDA_MMA_NO_VOLTA_PERM + +#if CUDART_VERSION >= 11080 + +static __device__ __forceinline__ int ggml_cuda_movmatrix(const int x) { + int ret = 0; + +#ifdef TURING_MMA_AVAILABLE + asm("movmatrix.sync.aligned.m8n8.trans.b16 %0, %1;" + : "=r"(ret) : "r"(x)); +#else + GGML_UNUSED(x); + NO_DEVICE_CODE; +#endif // defined(TURING_MMA_AVAILABLE) + return ret; +} + +#else + +static __device__ __forceinline__ int ggml_cuda_movmatrix(const int x) { + // Imagine transposing row-major matrix to column-major matrix. + const int src_i_low = 2 * (threadIdx.x % 4); + const int src_i_high = src_i_low + 1; + const int src_j = threadIdx.x / 4; + + const int src_laneid_low = src_i_low * 4 + src_j / 2; + const int src_laneid_high = src_i_high * 4 + src_j / 2; + + const int shift_low = ((src_j + 0) % 2) * 16; + const int shift_high = ((src_j + 1) % 2) * 16; + + const int ret_low = (__shfl_sync(0xFFFFFFFF, x, src_laneid_low, WARP_SIZE) >> shift_low) & 0x0000FFFF; + const int ret_high = (__shfl_sync(0xFFFFFFFF, x, src_laneid_high, WARP_SIZE) << shift_high) & 0xFFFF0000; + + return ret_low | ret_high; +} + +#endif // CUDART_VERSION >= 11080 + +static __device__ __forceinline__ half2 ggml_cuda_movmatrix(const half2 x) { + half2 ret; + *((int *) &ret) = ggml_cuda_movmatrix(*((const int *) &x)); + return ret; +} + +namespace ggml_cuda_mma { + + // Some architectures like Volta or CDNA3 perform multiple matrix multiplications per warp in parallel, + // effectively the warp is being split into subgroups of threads that each perform a single mma instruction. + // In those cases the data can be split in different ways across the warp. + enum data_layout { + // By default the data uses the I direction as its major dimension and the J direction as its minor dimension. + // For the A/C matrices this means I major == row major, J major == column major. + // For the B matrix this means I major == column major, J major == row major. + // MIRRORED == Each data value is held exactly once per thread subgroup. + DATA_LAYOUT_I_MAJOR = 0, // Always used for Turing, Ampere, Ada Lovelace, consumer Blackwell, matrix A&B for RDNA4 and CDNA. + DATA_LAYOUT_J_MAJOR = 10, // Matrix C for CDNA and RDNA4, int and float matrix C for RDNA3. + DATA_LAYOUT_I_MAJOR_MIRRORED = 20, // Volta, matrix A&B for RDNA3. + DATA_LAYOUT_J_MAJOR_MIRRORED = 30, + DATA_LAYOUT_I_MAJOR_SCRAMBLED = 40, // Scrambled matrix C for faster transposition (RDNA4/CDNA), convert to float to unscramble. + }; + // Implemented mma combinations are: + // - (I_MAJOR, I_MAJOR) -> I_MAJOR + // - (I_MAJOR, I_MAJOR_MIRRORED) -> I_MAJOR + // - (I_MAJOR, J_MAJOR_MIRRORED) -> I_MAJOR + + static constexpr __device__ data_layout get_input_data_layout() { +#if defined(RDNA3) || defined(VOLTA_MMA_AVAILABLE) + return DATA_LAYOUT_I_MAJOR_MIRRORED; +#else + return DATA_LAYOUT_I_MAJOR; +#endif // defined(RDNA3) || defined(VOLTA_MMA_AVAILABLE) + } + + template + struct tile {}; + + template + struct tile { + static constexpr int I = I_; + static constexpr int J = J_; + static constexpr data_layout dl = DATA_LAYOUT_I_MAJOR; + +#if defined(AMD_MFMA_AVAILABLE) + static constexpr int ne = I * J / 64; + T x[ne] = {0}; + + static constexpr __device__ bool supported() { + if (I == 16 && J == 8) return true; + if (I == 32 && J == 4) return true; + if (I == 16 && J == 16) return true; + if (I == 32 && J == 32) return true; + return false; + } + + static __device__ __forceinline__ int get_i(const int l) { + if constexpr (I == 16 && J == 4) { + return threadIdx.x % 16; + } else if constexpr (I == 16 && J == 8) { + return threadIdx.x % 16; + } else if constexpr (I == 32 && J == 4) { + return threadIdx.x % 32; + } else if constexpr (I == 16 && J == 16) { + return threadIdx.x % 16; + } else if constexpr (I == 32 && J == 32) { + return threadIdx.x % 32; + } else { + NO_DEVICE_CODE; + return -1; + } + } + + static __device__ __forceinline__ int get_j(const int l) { + if constexpr (I == 16 && J == 4) { + return threadIdx.x / 16; + } else if constexpr (I == 16 && J == 8) { + return 2 * (threadIdx.x / 16) + l; + } else if constexpr (I == 32 && J == 4) { + return 2 * (threadIdx.x / 32) + l; + } else if constexpr (I == 16 && J == 16) { + return 4 * (threadIdx.x / 16) + l; + } else if constexpr (I == 32 && J == 32) { + return 4 * (threadIdx.x / 32) + 8 * (l / 4) + (l % 4); + } else { + NO_DEVICE_CODE; + return -1; + } + } +#elif defined(VOLTA_MMA_AVAILABLE) + static constexpr int ne = I * J / 32; + T x[ne] = {0}; + + static constexpr __device__ bool supported() { + if (I == 32 && J == 8) return true; + return false; + } + + static __device__ __forceinline__ int get_i(const int l) { + if constexpr (I == 32 && J == 8) { +#ifdef GGML_CUDA_MMA_NO_VOLTA_PERM + return (((threadIdx.x % 16) / 4) * 8) + ((threadIdx.x / 16) * 4) + (l & 2) + (threadIdx.x % 2); +#else + return (l & 2) + (threadIdx.x & ~2); +#endif // GGML_CUDA_MMA_NO_VOLTA_PERM + } else { + NO_DEVICE_CODE; + return -1; + } + } + + static __device__ __forceinline__ int get_j(const int l) { + if constexpr (I == 32 && J == 8) { + return (threadIdx.x & 2) + (l & (4 + 1)); + } else { + NO_DEVICE_CODE; + return -1; + } + } +#elif defined(AMD_WMMA_AVAILABLE) + static constexpr int ne = I * J / 32; + T x[ne] = {0}; + + static constexpr __device__ bool supported() { + if (I == 16 && J == 16) return true; + if (I == 16 && J == 8) return true; + if (I == 16 && J == 4) return true; + return false; + } + + static __device__ __forceinline__ int get_i(const int l) { + if constexpr (supported()) { + return threadIdx.x % 16; + } else { + NO_DEVICE_CODE; + return -1; + } + } + + static __device__ __forceinline__ int get_j(const int l) { + if constexpr (I == 16 && J == 16) { +#if defined(RDNA3) + if constexpr (std::is_same_v || std::is_same_v) { + // matrix C + return 2 * l + (threadIdx.x / 16); + } else { + // matrix A&B + return l; + } +#else + // matrix C is the transposed matrix A&B on RDNA4 + return ne * (threadIdx.x / 16) + l; +#endif // defined(RDNA3) + } else if constexpr (I == 16 && J == 8) { + // mmq input for RDNA4 + return ne * (threadIdx.x / 16) + l; + } else if constexpr (I == 16 && J == 4) { + return ne * (threadIdx.x / 16) + l; + } else { + NO_DEVICE_CODE; + return -1; + } + } +#else + static constexpr int ne = I * J / 32; + T x[ne] = {0}; + + static constexpr __device__ bool supported() { + if (I == 8 && J == 4) return true; + if (I == 8 && J == 8) return true; + if (I == 16 && J == 8) return true; + if (I == 16 && J == 16) return true; + if (I == 32 && J == 8) return true; + return false; + } + + static __device__ __forceinline__ int get_i(const int l) { + if constexpr (I == 8 && J == 4) { + return threadIdx.x / 4; + } else if constexpr (I == 8 && J == 8) { + return threadIdx.x / 4; + } else if constexpr (I == 16 && J == 8) { + return ((l / 2) * 8) + (threadIdx.x / 4); + } else if constexpr (I == 16 && J == 16) { + return (((l / 2) % 2) * 8) + (threadIdx.x / 4); + } else if constexpr (I == 32 && J == 8) { + return tile<16, 8, T>::get_i(l); // Memory layout simply repeated with same pattern in i direction. + } else { + NO_DEVICE_CODE; + return -1; + } + } + + static __device__ __forceinline__ int get_j(const int l) { + if constexpr (I == 8 && J == 4) { + return threadIdx.x % 4; + } else if constexpr (I == 8 && J == 8) { + return (l * 4) + (threadIdx.x % 4); + } else if constexpr (I == 16 && J == 8) { + return ((threadIdx.x % 4) * 2) + (l % 2); + } else if constexpr (I == 16 && J == 16) { + return ((l / 4) * 8) + ((threadIdx.x % 4) * 2) + (l % 2); + } else if constexpr (I == 32 && J == 8) { + return tile<16, 8, T>::get_j(l); // Memory layout simply repeated with same pattern in i direction. + } else { + NO_DEVICE_CODE; + return -1; + } + } +#endif // defined(GGML_USE_HIP) + }; + + template + struct tile { + static constexpr int I = I_; + static constexpr int J = J_; + static constexpr data_layout dl = DATA_LAYOUT_I_MAJOR; + +#if defined(VOLTA_MMA_AVAILABLE) + static constexpr int ne = I * J / WARP_SIZE; + half2 x[ne] = {{0.0f, 0.0f}}; + + static constexpr __device__ bool supported() { + if (I == 32 && J == 4) return true; + return false; + } + + static __device__ __forceinline__ int get_i(const int l) { + if constexpr (I == 32 && J == 4) { +#ifdef GGML_CUDA_MMA_NO_VOLTA_PERM + return (((threadIdx.x % 16) / 4) * 8) + ((threadIdx.x / 16) * 4) + (threadIdx.x % 4); +#else + return threadIdx.x; +#endif // GGML_CUDA_MMA_NO_VOLTA_PERM + } else { + NO_DEVICE_CODE; + return -1; + } + } + + static __device__ __forceinline__ int get_j(const int l) { + if constexpr (I == 32 && J == 4) { + return l; + } else { + NO_DEVICE_CODE; + return -1; + } + } +#elif defined(AMD_WMMA_AVAILABLE) + static constexpr int ne = I * J / 32; + half2 x[ne] = {{0.0f, 0.0f}}; + + static constexpr __device__ bool supported() { + if (I == 16 && J == 8) return true; + if (I == 16 && J == 16) return true; + if (I == 32 && J == 8) return true; + return false; + } + + static __device__ __forceinline__ int get_i(const int l) { + if constexpr (I == 16 && J == 8) { + return threadIdx.x % 16; + } else if constexpr (I == 16 && J == 16) { + return threadIdx.x % 16; + } else if constexpr (I == 32 && J == 8) { + return (threadIdx.x % 16) * 2 + l / (ne/2); + } else { + NO_DEVICE_CODE; + return -1; + } + } + + static __device__ __forceinline__ int get_j(const int l) { + if constexpr (I == 16 && J == 8) { + return (threadIdx.x / 16) * ne + l; + } else if constexpr (I == 16 && J == 16) { +#ifdef RDNA3 + return l*2 + (threadIdx.x / 16); +#else + return (threadIdx.x / 16) * ne + l; +#endif // RDNA3 + } else if constexpr (I == 32 && J == 8) { + return (threadIdx.x / 16) * (ne/2) + l % (ne/2); + } else { + NO_DEVICE_CODE; + return -1; + } + } +#elif defined(AMD_MFMA_AVAILABLE) + static constexpr int ne = I * J / 64; + half2 x[ne] = {{0.0f, 0.0f}}; + + static constexpr __device__ bool supported() { + if (I == 16 && J == 8) return true; + if (I == 16 && J == 16) return true; + if (I == 32 && J == 8) return true; + return false; + } + + static __device__ __forceinline__ int get_i(const int l) { + if constexpr (I == 16 && J == 8) { + return threadIdx.x % 16; + } else if constexpr (I == 16 && J == 16) { + return threadIdx.x % 16; + } else if constexpr (I == 32 && J == 8) { + return (threadIdx.x % 16) * 2 + l / (ne/2); + } else { + NO_DEVICE_CODE; + return -1; + } + } + + static __device__ __forceinline__ int get_j(const int l) { + if constexpr (I == 16 && J == 8) { + return (threadIdx.x / 16) * ne + l; + } else if constexpr (I == 16 && J == 16) { + return (threadIdx.x / 16) * ne + l; + } else if constexpr (I == 32 && J == 8) { + return (threadIdx.x / 16) * (ne/2) + l % (ne/2); + } else { + NO_DEVICE_CODE; + return -1; + } + } +#else + static constexpr int ne = I * J / WARP_SIZE; + half2 x[ne] = {{0.0f, 0.0f}}; + + static constexpr __device__ bool supported() { + if (I == 8 && J == 4) return true; + if (I == 8 && J == 8) return true; + if (I == 16 && J == 8) return true; + if (I == 16 && J == 16) return true; + if (I == 32 && J == 8) return true; + return false; + } + + static __device__ __forceinline__ int get_i(const int l) { + if constexpr (I == 8 && J == 8) { + return threadIdx.x / 4; + } else if constexpr (I == 16 && J == 4) { + return (l * 8) + (threadIdx.x / 4); + } else if constexpr (I == 16 && J == 8) { + return ((l % 2) * 8) + (threadIdx.x / 4); + } else if constexpr (I == 32 && J == 8) { + return ((l / 4) * 16) + ((l % 2) * 8) + (threadIdx.x / 4); + } else { + NO_DEVICE_CODE; + return -1; + } + } + + static __device__ __forceinline__ int get_j(const int l) { + if constexpr (I == 8 && J == 8) { + return (l * 4) + (threadIdx.x % 4); + } else if constexpr (I == 16 && J == 4) { + return threadIdx.x % 4; + } else if constexpr (I == 16 && J == 8) { + return ((l / 2) * 4) + (threadIdx.x % 4); + } else if constexpr (I == 32 && J == 8) { + return ((l & 2) * 2) + (threadIdx.x % 4); + } else { + NO_DEVICE_CODE; + return -1; + } + } +#endif // defined(VOLTA_MMA_AVAILABLE) + }; + + template + struct tile { + static constexpr int I = I_; + static constexpr int J = J_; + static constexpr data_layout dl = DATA_LAYOUT_I_MAJOR; + +#if defined(AMD_WMMA_AVAILABLE) + static constexpr int ne = tile::ne; + nv_bfloat162 x[ne] = {{0.0f, 0.0f}}; + + static constexpr __device__ bool supported() { + return tile::supported(); + } + + static __device__ __forceinline__ int get_i(const int l) { + return tile::get_i(l); + } + + static __device__ __forceinline__ int get_j(const int l) { + return tile::get_j(l); + } +#elif defined(AMD_MFMA_AVAILABLE) + static constexpr int ne = tile::ne; + nv_bfloat162 x[ne] = {{0.0f, 0.0f}}; + + static constexpr __device__ bool supported() { + return tile::supported(); + } + + static __device__ __forceinline__ int get_i(const int l) { + return tile::get_i(l); + } + + static __device__ __forceinline__ int get_j(const int l) { + return tile::get_j(l); + } +#else + static constexpr int ne = I * J / WARP_SIZE; + nv_bfloat162 x[ne] = {{0.0f, 0.0f}}; + + static constexpr __device__ bool supported() { + if (I == 8 && J == 8) return true; + if (I == 16 && J == 4) return true; + if (I == 16 && J == 8) return true; + return false; + } + + static __device__ __forceinline__ int get_i(const int l) { + if constexpr (I == 8 && J == 8) { + return threadIdx.x / 4; + } else if constexpr (I == 16 && J == 4) { + return (l * 8) + (threadIdx.x / 4); + } else if constexpr (I == 16 && J == 8) { + return ((l % 2) * 8) + (threadIdx.x / 4); + } else { + NO_DEVICE_CODE; + return -1; + } + } + + static __device__ __forceinline__ int get_j(const int l) { + if constexpr (I == 8 && J == 8) { + return (l * 4) + (threadIdx.x % 4); + } else if constexpr (I == 16 && J == 4) { + return threadIdx.x % 4; + } else if constexpr (I == 16 && J == 8) { + return ((l / 2) * 4) + (threadIdx.x % 4); + } else { + NO_DEVICE_CODE; + return -1; + } + } +#endif // defined(AMD_WMMA_AVAILABLE) + }; + + template + struct tile { + static constexpr int I = I_; + static constexpr int J = J_; + static constexpr data_layout dl = DATA_LAYOUT_J_MAJOR; + + static constexpr int ne = tile::ne; + T x[ne] = {0}; + + static constexpr __device__ bool supported() { + return tile::supported(); + } + + static __device__ __forceinline__ int get_i(const int l) { + return tile::get_j(l); + } + + static __device__ __forceinline__ int get_j(const int l) { + return tile::get_i(l); + } + }; + + template + struct tile { + static constexpr int I = I_; + static constexpr int J = J_; + static constexpr data_layout dl = DATA_LAYOUT_I_MAJOR_MIRRORED; + + // RDNA3 + static constexpr int ne = I * J / 32 * 2; + + T x[ne] = {0}; + + static constexpr __device__ bool supported() { + if (I == 16 && J == 16) return true; + if (I == 16 && J == 8) return true; + if (I == 16 && J == 4) return true; + if (I == 32 && J == 8) return true; + return false; + } + + static __device__ __forceinline__ int get_i(const int l) { + if constexpr (I == 16) { + return threadIdx.x % 16; + } else if constexpr (I == 32) { + return (threadIdx.x % 16) * 2 + l / (ne/2); + } else { + NO_DEVICE_CODE; + return -1; + } + } + + static __device__ __forceinline__ int get_j(const int l) { + if constexpr (I == 16) { + return l; + } else if constexpr (I == 32) { + return l % (ne/2); + } else { + NO_DEVICE_CODE; + return -1; + } + } + }; + + template + struct tile { + static constexpr int I = I_; + static constexpr int J = J_; + static constexpr data_layout dl = DATA_LAYOUT_I_MAJOR_MIRRORED; +#if defined(RDNA3) + static constexpr int ne = tile::ne; + + half2 x[ne] = {{0.0f, 0.0f}}; + + static constexpr __device__ bool supported() { + return tile::supported(); + } + + static __device__ __forceinline__ int get_i(const int l) { + return tile::get_i(l); + } + + static __device__ __forceinline__ int get_j(const int l) { + return tile::get_j(l); + } +#else // Volta + static constexpr int ne = I * J / (WARP_SIZE/4); + + half2 x[ne] = {{0.0f, 0.0f}}; + + static constexpr __device__ bool supported() { + if (I == 8 && J == 4) return true; + return false; + } + + static __device__ __forceinline__ int get_i(const int /*l*/) { + if constexpr (I == 8 && J == 4) { + return ((threadIdx.x / 16) * 4) + (threadIdx.x % 4); + } else { + NO_DEVICE_CODE; + return -1; + } + } + + static __device__ __forceinline__ int get_j(const int l) { + if constexpr (I == 8 && J == 4) { + return l; + } else { + NO_DEVICE_CODE; + return -1; + } + } +#endif // defined(RDNA3) + }; + + template + struct tile { + static constexpr int I = I_; + static constexpr int J = J_; + static constexpr data_layout dl = DATA_LAYOUT_I_MAJOR_MIRRORED; + static constexpr int ne = tile::ne; + + nv_bfloat162 x[ne] = {{0.0f, 0.0f}}; + + static constexpr __device__ bool supported() { + return tile::supported(); + } + + static __device__ __forceinline__ int get_i(const int l) { + return tile::get_i(l); + } + + static __device__ __forceinline__ int get_j(const int l) { + return tile::get_j(l); + } + }; + + template + struct tile { + static constexpr int I = I_; + static constexpr int J = J_; + static constexpr data_layout dl = DATA_LAYOUT_J_MAJOR_MIRRORED; + static constexpr int ne = I * J / (WARP_SIZE/4); + + half2 x[ne] = {{0.0f, 0.0f}}; + + static constexpr __device__ bool supported() { + if (I == 8 && J == 4) return true; + return false; + } + + static __device__ __forceinline__ int get_i(const int l) { + if constexpr (I == 8 && J == 4) { + return ((l / 2) * 4) + (threadIdx.x % 4); + } else { + NO_DEVICE_CODE; + return -1; + } + } + + static __device__ __forceinline__ int get_j(const int l) { + if constexpr (I == 8 && J == 4) { + return ((threadIdx.x / 16) * 2) + (l % 2); + } else { + NO_DEVICE_CODE; + return -1; + } + } + }; + + template + struct tile { + static constexpr int I = I_; + static constexpr int J = J_; + static constexpr data_layout dl = DATA_LAYOUT_I_MAJOR_SCRAMBLED; + + static constexpr int ne = I * J / ggml_cuda_get_physical_warp_size(); + half2 x[ne] = {{0.0f, 0.0f}}; + + static constexpr __device__ bool supported() { + if (I == 16 && J == 16) return true; + return false; + } + + static __device__ __forceinline__ int get_i(const int l) { + return tile::get_i(l); + } + }; + + static __device__ __forceinline__ tile<16, 16, half2, DATA_LAYOUT_I_MAJOR> unscramble(const tile<16, 16, half2, DATA_LAYOUT_I_MAJOR_SCRAMBLED> & t) { +#if defined(AMD_MFMA_AVAILABLE) || (defined(AMD_WMMA_AVAILABLE) && defined(RDNA4)) + tile<16, 16, half2, DATA_LAYOUT_I_MAJOR> ret; +#pragma unroll + for (int l0 = 0; l0 < t.ne/2; ++l0) { + ret.x[2*l0 + 0] = __lows2half2(t.x[l0], t.x[l0 + t.ne/2]); + ret.x[2*l0 + 1] = __highs2half2(t.x[l0], t.x[l0 + t.ne/2]); + } + return ret; +#else + NO_DEVICE_CODE; + GGML_UNUSED(t); +#endif // defined(AMD_MFMA_AVAILABLE) || (defined(AMD_WMMA_AVAILABLE) && defined(RDNA4)) + } + +#if defined(TURING_MMA_AVAILABLE) + template + static __device__ __forceinline__ tile get_half2(const tile & tile_float) { + tile ret; +#pragma unroll + for (int l0 = 0; l0 < tile_float.ne; l0 += 2) { + ret.x[l0/2] = make_half2(tile_float.x[l0 + 0], tile_float.x[l0 + 1]); + } + return ret; + } + + static __device__ __forceinline__ tile<8, 8, half2> get_transposed(const tile<16, 4, half2> & t) { + tile<8, 8, half2> ret; + ret.x[0] = ggml_cuda_movmatrix(t.x[0]); + ret.x[1] = ggml_cuda_movmatrix(t.x[1]); + + return ret; + } +#elif defined(AMD_WMMA_AVAILABLE) && defined(RDNA3) + static __device__ __forceinline__ tile<16, 8, half2, DATA_LAYOUT_I_MAJOR_MIRRORED> get_half2( + const tile<16, 16, float, DATA_LAYOUT_I_MAJOR> & tile_float) { + tile<16, 8, half2, DATA_LAYOUT_I_MAJOR_MIRRORED> ret; +#pragma unroll + for (int l = 0; l < tile_float.ne; ++l) { + float tmp[2]; + int i = threadIdx.x / 16; + tmp[i] = tile_float.x[l]; + i ^= 1; + tmp[i] = __shfl_xor_sync(0xFFFFFFFF, tile_float.x[l], 16, WARP_SIZE); + ret.x[l] = make_half2(tmp[0], tmp[1]); + } + return ret; + } +#elif defined(AMD_WMMA_AVAILABLE) || defined(AMD_MFMA_AVAILABLE) + template + static __device__ __forceinline__ tile get_half2(const tile & tile_float) { + tile ret; +#pragma unroll + for (int l0 = 0; l0 < tile_float.ne; l0 += 2) { + ret.x[l0/2] = make_half2(tile_float.x[l0 + 0], tile_float.x[l0 + 1]); + } + return ret; + } + + static __device__ __forceinline__ tile<8, 8, half2> get_transposed(const tile<16, 4, half2> & t) { + NO_DEVICE_CODE; + return tile<8, 8, half2>{}; + } +#else // Volta + template + static __device__ __forceinline__ tile get_half2(const tile & tile_float) { + tile ret; +#pragma unroll + for (int l0 = 0; l0 < tile_float.ne; l0 += 4) { + ret.x[l0/2 + 0] = make_half2(tile_float.x[l0 + 0], tile_float.x[l0 + 1]); + ret.x[l0/2 + 1] = make_half2(tile_float.x[l0 + 2], tile_float.x[l0 + 3]); + + // On Volta FP16 and FP32 tiles have a different memory layout, + // for the conversion threads with an offset of 2 need to exchange half their values: + ret.x[l0/2 + (((threadIdx.x % 4) / 2) ^ 1)] = __shfl_xor_sync( + 0xFFFFFFFF, ret.x[l0/2 + (((threadIdx.x % 4) / 2) ^ 1)], 2, WARP_SIZE); + } + return ret; + } +#endif // defined(TURING_MMA_AVAILABLE) + + template + static __device__ __forceinline__ void load_generic(tile & t, const T * __restrict__ xs0, const int stride) { +#pragma unroll + for (int l = 0; l < t.ne; ++l) { + t.x[l] = xs0[t.get_i(l)*stride + t.get_j(l)]; + } + } + + template + static __device__ __forceinline__ uint32_t swizzle(const uint32_t offset, const uint32_t i) { + static_assert(sizeof(T) <= 4, "unsupported type size"); + constexpr int stride_bytes = stride*sizeof(T); + static_assert(stride_bytes % 16 == 0, "bad stride"); + constexpr uint32_t shift = sizeof(T) == 1 ? 4 : (sizeof(T) == 2 ? 3 : 2); + if (stride_bytes % 32 != 0) { + return offset; // Equivalent to padding with 16 bytes. + } + if (stride_bytes % 64 != 0) { + return offset ^ (((i / 4) % 2) << shift); + } + if (stride_bytes % 128 != 0) { + return offset ^ (((i / 2) % 4) << shift); + } + return offset ^ ((i % 8) << shift); + } + + template + static __device__ __forceinline__ T * swizzle(T * ptr, const uint32_t offset, const uint32_t i) { + return ptr + swizzle(offset, i); + } + + template + static __device__ __forceinline__ void load_ldmatrix( + tile<8, 8, T> & t, const T * __restrict__ xs0, const int stride) { +#ifdef TURING_MMA_AVAILABLE + int * xi = (int *) t.x; + const int * xs = (const int *) xs0 + (threadIdx.x % t.I) * stride + ((threadIdx.x / t.I) * (t.J / 2)) % t.J; + asm volatile("ldmatrix.sync.aligned.m8n8.x2.b16 {%0, %1}, [%2];" + : "=r"(xi[0]), "=r"(xi[1]) + : "l"(xs)); +#else + GGML_UNUSED_VARS(t, xs0, stride); + NO_DEVICE_CODE; +#endif // TURING_MMA_AVAILABLE + } + + template + static __device__ __forceinline__ void load_ldmatrix( + tile<16, 4, T, dl> & t, const T * __restrict__ xs0, const int stride) { +#ifdef TURING_MMA_AVAILABLE + int * xi = (int *) t.x; + const int * xs = (const int *) xs0 + (threadIdx.x % t.I) * stride; + asm volatile("ldmatrix.sync.aligned.m8n8.x2.b16 {%0, %1}, [%2];" + : "=r"(xi[0]), "=r"(xi[1]) + : "l"(xs)); +#elif defined(AMD_WMMA_AVAILABLE) +#ifdef RDNA3 + static_assert(dl == DATA_LAYOUT_I_MAJOR_MIRRORED, "bad data layout"); + static_assert(sizeof(t.x) == 16, "bad ne"); + ggml_cuda_memcpy_1<8>(t.x + 0, xs0 + t.get_i(0)*stride + 0); + ggml_cuda_memcpy_1<8>(t.x + 2, xs0 + t.get_i(0)*stride + 2); +#else + static_assert(dl == DATA_LAYOUT_I_MAJOR, "bad data layout"); + static_assert(sizeof(t.x) == 8, "bad ne"); + ggml_cuda_memcpy_1<8>(t.x, xs0 + t.get_i(0)*stride + t.get_j(0)); +#endif // RDNA3 +#elif defined(AMD_MFMA_AVAILABLE) + static_assert(sizeof(t.x) == 4, "bad ne"); + ggml_cuda_memcpy_1<4>(t.x, xs0 + t.get_i(0)*stride + t.get_j(0)); +#else + GGML_UNUSED_VARS(t, xs0, stride); + NO_DEVICE_CODE; +#endif // TURING_MMA_AVAILABLE + } + + template + static __device__ __forceinline__ void load_ldmatrix( + tile<16, 8, T, dl> & t, const T * __restrict__ xs0, const int stride) { +#if defined(TURING_MMA_AVAILABLE) + int * xi = (int * ) t.x; + const int * xs = (const int *) xs0 + (threadIdx.x % t.I) * stride + (threadIdx.x / t.I) * (t.J / 2); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.b16 {%0, %1, %2, %3}, [%4];" + : "=r"(xi[0]), "=r"(xi[1]), "=r"(xi[2]), "=r"(xi[3]) + : "l"(xs)); +#elif defined(VOLTA_MMA_AVAILABLE) + ggml_cuda_memcpy_1<4*sizeof(T)>(t.x + 0, xs0 + t.get_i(0)*stride + 0); + ggml_cuda_memcpy_1<4*sizeof(T)>(t.x + 4, xs0 + t.get_i(4)*stride + 4); +#elif defined(AMD_WMMA_AVAILABLE) +#ifdef RDNA3 + static_assert(dl == DATA_LAYOUT_I_MAJOR_MIRRORED, "bad data layout"); + static_assert(sizeof(t.x) == 32, "bad ne"); + ggml_cuda_memcpy_1<16>(t.x + 0, xs0 + t.get_i(0)*stride + 0); + ggml_cuda_memcpy_1<16>(t.x + 4, xs0 + t.get_i(0)*stride + 4); +#else + static_assert(dl == DATA_LAYOUT_I_MAJOR, "bad data layout"); + static_assert(sizeof(t.x) == 16, "bad ne"); + ggml_cuda_memcpy_1<16>(t.x, xs0 + t.get_i(0)*stride + t.get_j(0)); +#endif // RDNA3 +#elif defined(AMD_MFMA_AVAILABLE) + static_assert(sizeof(t.x) == 8, "bad ne"); + ggml_cuda_memcpy_1<8>(t.x, xs0 + t.get_i(0)*stride + t.get_j(0)); +#else + GGML_UNUSED_VARS(t, xs0, stride); + NO_DEVICE_CODE; +#endif // TURING_MMA_AVAILABLE + } + + static __device__ __forceinline__ void load_ldmatrix( + tile<8, 4, half2, DATA_LAYOUT_I_MAJOR_MIRRORED> & t, const half2 * __restrict__ xs0, const int stride) { + ggml_cuda_memcpy_1<4*sizeof(half2)>(t.x, xs0 + t.get_i(0)*stride); + } + + static __device__ __forceinline__ void load_ldmatrix( + tile<8, 4, half2, DATA_LAYOUT_J_MAJOR_MIRRORED> & t, const half2 * __restrict__ xs0, const int stride) { +#ifdef VOLTA_MMA_AVAILABLE +#pragma unroll + for (int l0 = 0; l0 < t.ne; l0 += 2) { + ggml_cuda_memcpy_1<2*sizeof(half2)>(t.x + l0, xs0 + t.get_i(l0)*stride + t.get_j(l0)); + } +#else + GGML_UNUSED_VARS(t, xs0, stride); + NO_DEVICE_CODE; +#endif // VOLTA_MMA_AVAILABLE + } + + static __device__ __forceinline__ void load_ldmatrix( + tile<32, 4, half2> & t, const half2 * __restrict__ xs0, const int stride) { +#if defined(VOLTA_MMA_AVAILABLE) + ggml_cuda_memcpy_1<4*sizeof(half2)>(t.x, xs0 + t.get_i(0)*stride); +#else + GGML_UNUSED_VARS(t, xs0, stride); + NO_DEVICE_CODE; +#endif // defined(VOLTA_MMA_AVAILABLE) + } + + template + static __device__ __forceinline__ void load_ldmatrix_trans( + tile & t, const T * __restrict__ xs0, const int stride) { +#ifdef TURING_MMA_AVAILABLE + static_assert(I == 16, "bad tile width"); + static_assert(dl == DATA_LAYOUT_I_MAJOR, "bad data layout"); + int * xi = (int *) t.x; + const int * xs = (const int *) xs0 + (threadIdx.x % t.I) * stride + (threadIdx.x / t.I) * (t.J / 2); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.trans.b16 {%0, %1, %2, %3}, [%4];" + : "=r"(xi[0]), "=r"(xi[2]), "=r"(xi[1]), "=r"(xi[3]) + : "l"(xs)); +#elif defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + static_assert(dl == DATA_LAYOUT_I_MAJOR || dl == DATA_LAYOUT_I_MAJOR_MIRRORED, "bad data layout"); + if constexpr (I == 32) { +#pragma unroll + for (int l0 = 0; l0 < t.ne/2; ++l0) { + const half2 tmp0 = xs0[(2*t.get_j(l0) + 0)*stride + t.get_i(l0)/2]; + const half2 tmp1 = xs0[(2*t.get_j(l0) + 1)*stride + t.get_i(l0)/2]; + + t.x[l0] = __lows2half2(tmp0, tmp1); + t.x[l0 + t.ne/2] = __highs2half2(tmp0, tmp1); + } + } else { + half * xh = (half *) t.x; +#pragma unroll + for (int l = 0; l < t.ne; ++l) { + xh[2*l + 0] = ((const half *) xs0)[(2*t.get_j(l) + 0)*(2*stride) + t.get_i(l)]; + xh[2*l + 1] = ((const half *) xs0)[(2*t.get_j(l) + 1)*(2*stride) + t.get_i(l)]; + } + } +#else + GGML_UNUSED_VARS(t, xs0, stride); + NO_DEVICE_CODE; +#endif // TURING_MMA_AVAILABLE + } + + template + static __device__ __forceinline__ void load_ldmatrix_swizzled( + tile & t, const T * __restrict__ xs0, const int offset) { +#if defined(TURING_MMA_AVAILABLE) + static_assert(I == 16, "bad tile width"); + static_assert(J == 8, "bad tile height"); + const int i = threadIdx.x % t.I; + const int j = (threadIdx.x / t.I) * (t.J / 2); + int offset_ij = offset + i * stride + j; + offset_ij = swizzle(offset_ij, i); + int * xi = (int *) t.x; + asm volatile("ldmatrix.sync.aligned.m8n8.x4.b16 {%0, %1, %2, %3}, [%4];" + : "=r"(xi[0]), "=r"(xi[1]), "=r"(xi[2]), "=r"(xi[3]) + : "l"(xs0 + offset_ij)); +#elif defined(VOLTA_MMA_AVAILABLE) +#pragma unroll + for (int o = 0; o < t.ne; o += 4) { + const int offset_ij = offset + t.get_i(o) * stride + o; + ggml_cuda_memcpy_1<4*sizeof(T)>(t.x + o, swizzle(xs0, offset_ij, t.get_i(o))); + } +#elif defined(AMD_WMMA_AVAILABLE) +#ifdef RDNA3 + static_assert(dl == DATA_LAYOUT_I_MAJOR_MIRRORED, "bad data layout"); + static_assert(sizeof(t.x) == 32, "bad ne"); + static_assert(I == 16, "bad tile width"); + static_assert(J == 8, "bad tile height"); +#pragma unroll + for (int o = 0; o < 8; o += 4) { + const int offset_ij = offset + t.get_i(0) * stride + o; + ggml_cuda_memcpy_1<16>(t.x + o, swizzle(xs0, offset_ij, t.get_i(0))); + } +#else + static_assert(dl == DATA_LAYOUT_I_MAJOR, "bad data layout"); + static_assert(sizeof(t.x) == 16, "bad ne"); + const int offset_ij = offset + t.get_i(0)*stride + t.get_j(0); + ggml_cuda_memcpy_1<16>(t.x, swizzle(xs0, offset_ij, t.get_i(0))); +#endif // RDNA3 +#elif defined(AMD_MFMA_AVAILABLE) + static_assert(sizeof(t.x) == 8, "bad ne"); + const int offset_ij = offset + t.get_i(0)*stride + t.get_j(0); + ggml_cuda_memcpy_1<8>(t.x, swizzle(xs0, offset_ij, t.get_i(0))); +#else + GGML_UNUSED_VARS(t, xs0, offset); + NO_DEVICE_CODE; +#endif // defined(TURING_MMA_AVAILABLE) + } + + template + static __device__ __forceinline__ void load_ldmatrix_swizzled( + tile<8, 4, half2, DATA_LAYOUT_J_MAJOR_MIRRORED> & t, const half2 * __restrict__ xs0, const int offset) { +#ifdef VOLTA_MMA_AVAILABLE +#pragma unroll + for (int l0 = 0; l0 < t.ne; l0 += 2) { + const int offset_ij = offset + t.get_i(l0)*stride + t.get_j(l0); + ggml_cuda_memcpy_1<2*sizeof(half2)>(t.x + l0, swizzle(xs0, offset_ij, t.get_i(l0))); + } +#else + GGML_UNUSED_VARS(t, xs0, offset); + NO_DEVICE_CODE; +#endif // VOLTA_MMA_AVAILABLE + } + + template + static __device__ __forceinline__ void load_ldmatrix_trans_swizzled( + tile & t, const T * __restrict__ xs0, const int offset) { +#if defined(TURING_MMA_AVAILABLE) + static_assert(I == 16, "bad tile width"); + static_assert(dl == DATA_LAYOUT_I_MAJOR, "bad data layout"); + const int i = threadIdx.x % t.I; + const int j = (threadIdx.x / t.I) * (t.J / 2); + int offset_ij = offset + i * stride + j; + offset_ij = swizzle(offset_ij, i); + int * xi = (int *) t.x; + asm volatile("ldmatrix.sync.aligned.m8n8.x4.trans.b16 {%0, %1, %2, %3}, [%4];" + : "=r"(xi[0]), "=r"(xi[2]), "=r"(xi[1]), "=r"(xi[3]) + : "l"(xs0 + offset_ij)); +#elif defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + static_assert(dl == DATA_LAYOUT_I_MAJOR || dl == DATA_LAYOUT_I_MAJOR_MIRRORED, "bad data layout"); + if constexpr (I == 32) { +#pragma unroll + for (int l0 = 0; l0 < t.ne/2; ++l0) { + half2 tmp[2]; +#pragma unroll + for (int o = 0; o < 2; ++o) { + const int j = 2*t.get_j(l0) + o; + int offset_ij = offset + j*stride + t.get_i(l0)/2; + offset_ij = swizzle(offset_ij, j); + tmp[o] = xs0[offset_ij]; + } + + t.x[l0] = __lows2half2(tmp[0], tmp[1]); + t.x[l0 + t.ne/2] = __highs2half2(tmp[0], tmp[1]); + } + } else { + half * xh = (half *) t.x; +#pragma unroll + for (int l = 0; l < t.ne; ++l) { +#pragma unroll + for (int o = 0; o < 2; ++o) { + const int j = 2*t.get_j(l) + o; + xh[2*l + o] = ((const half *) xs0)[swizzle<2*stride, half>(2*offset + j*(2*stride) + t.get_i(l), j)]; + } + } + } +#else + GGML_UNUSED_VARS(t, xs0, offset); + NO_DEVICE_CODE; +#endif // defined(TURING_MMA_AVAILABLE) + } + + static __device__ __forceinline__ void mma( + tile<16, 8, int> & D, const tile<16, 4, int> & A, const tile<8, 4, int> & B) { +#ifdef TURING_MMA_AVAILABLE +#if __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE + asm("mma.sync.aligned.m16n8k16.row.col.s32.s8.s8.s32 {%0, %1, %2, %3}, {%4, %5}, {%6}, {%0, %1, %2, %3};" + : "+r"(D.x[0]), "+r"(D.x[1]), "+r"(D.x[2]), "+r"(D.x[3]) + : "r"(A.x[0]), "r"(A.x[1]), "r"(B.x[0])); +#else + // On Turing m16n8k16 mma is not available, use 2x m8n8k16 mma instead: + asm("mma.sync.aligned.m8n8k16.row.col.s32.s8.s8.s32 {%0, %1}, {%2}, {%3}, {%0, %1};" + : "+r"(D.x[0]), "+r"(D.x[1]) + : "r"(A.x[0]), "r"(B.x[0])); + asm("mma.sync.aligned.m8n8k16.row.col.s32.s8.s8.s32 {%0, %1}, {%2}, {%3}, {%0, %1};" + : "+r"(D.x[2]), "+r"(D.x[3]) + : "r"(A.x[1]), "r"(B.x[0])); +#endif // __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE +#else + GGML_UNUSED_VARS(D, A, B); + NO_DEVICE_CODE; +#endif // TURING_MMA_AVAILABLE + } + + static __device__ __forceinline__ void mma( + tile<16, 8, int> & D, const tile<16, 8, int> & A, const tile<8, 8, int> & B) { +#ifdef TURING_MMA_AVAILABLE +#if __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE + asm("mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};" + : "+r"(D.x[0]), "+r"(D.x[1]), "+r"(D.x[2]), "+r"(D.x[3]) + : "r"(A.x[0]), "r"(A.x[1]), "r"(A.x[2]), "r"(A.x[3]), "r"(B.x[0]), "r"(B.x[1])); +#else + // On Turing m16n8k32 mma is not available, use 4x m8n8k16 mma instead: + asm("mma.sync.aligned.m8n8k16.row.col.s32.s8.s8.s32 {%0, %1}, {%2}, {%3}, {%0, %1};" + : "+r"(D.x[0]), "+r"(D.x[1]) + : "r"(A.x[0]), "r"(B.x[0])); + asm("mma.sync.aligned.m8n8k16.row.col.s32.s8.s8.s32 {%0, %1}, {%2}, {%3}, {%0, %1};" + : "+r"(D.x[2]), "+r"(D.x[3]) + : "r"(A.x[1]), "r"(B.x[0])); + asm("mma.sync.aligned.m8n8k16.row.col.s32.s8.s8.s32 {%0, %1}, {%2}, {%3}, {%0, %1};" + : "+r"(D.x[0]), "+r"(D.x[1]) + : "r"(A.x[2]), "r"(B.x[1])); + asm("mma.sync.aligned.m8n8k16.row.col.s32.s8.s8.s32 {%0, %1}, {%2}, {%3}, {%0, %1};" + : "+r"(D.x[2]), "+r"(D.x[3]) + : "r"(A.x[3]), "r"(B.x[1])); +#endif // __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE +#else + GGML_UNUSED_VARS(D, A, B); + NO_DEVICE_CODE; +#endif // TURING_MMA_AVAILABLE + } + + static __device__ __forceinline__ void mma( + tile<16, 4, half2> & D, const tile<16, 8, half2> & A, const tile<8, 8, half2> & B) { +#ifdef TURING_MMA_AVAILABLE + const int * Axi = (const int *) A.x; + const int * Bxi = (const int *) B.x; + int * Dxi = (int *) D.x; +#if __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE + asm("mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 {%0, %1}, {%2, %3, %4, %5}, {%6, %7}, {%0, %1};" + : "+r"(Dxi[0]), "+r"(Dxi[1]) + : "r"(Axi[0]), "r"(Axi[1]), "r"(Axi[2]), "r"(Axi[3]), "r"(Bxi[0]), "r"(Bxi[1])); +#else + // On Turing m16n8k16 mma is not available, use 2x m8n8k8 mma instead: + asm("mma.sync.aligned.m16n8k8.row.col.f16.f16.f16.f16 {%0, %1}, {%2, %3}, {%4}, {%0, %1};" + : "+r"(Dxi[0]), "+r"(Dxi[1]) + : "r"(Axi[0]), "r"(Axi[1]), "r"(Bxi[0])); + asm("mma.sync.aligned.m16n8k8.row.col.f16.f16.f16.f16 {%0, %1}, {%2, %3}, {%4}, {%0, %1};" + : "+r"(Dxi[0]), "+r"(Dxi[1]) + : "r"(Axi[2]), "r"(Axi[3]), "r"(Bxi[1])); +#endif // __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE +#else + GGML_UNUSED_VARS(D, A, B); + NO_DEVICE_CODE; +#endif // TURING_MMA_AVAILABLE + } + + static __device__ __forceinline__ void mma( + tile<16, 8, half2> & D, const tile<16, 8, half2> & A, const tile<16, 8, half2> & B) { +#ifdef TURING_MMA_AVAILABLE + const int * Axi = (const int *) A.x; + const int * Bxi = (const int *) B.x; + int * Dxi = (int *) D.x; +#if __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE + asm("mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 {%0, %1}, {%2, %3, %4, %5}, {%6, %7}, {%0, %1};" + : "+r"(Dxi[0]), "+r"(Dxi[1]) + : "r"(Axi[0]), "r"(Axi[1]), "r"(Axi[2]), "r"(Axi[3]), "r"(Bxi[0]), "r"(Bxi[2])); + asm("mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 {%0, %1}, {%2, %3, %4, %5}, {%6, %7}, {%0, %1};" + : "+r"(Dxi[2]), "+r"(Dxi[3]) + : "r"(Axi[0]), "r"(Axi[1]), "r"(Axi[2]), "r"(Axi[3]), "r"(Bxi[1]), "r"(Bxi[3])); +#else + // On Turing m16n8k16 mma is not available, use 4x m8n8k8 mma instead: + asm("mma.sync.aligned.m16n8k8.row.col.f16.f16.f16.f16 {%0, %1}, {%2, %3}, {%4}, {%0, %1};" + : "+r"(Dxi[0]), "+r"(Dxi[1]) + : "r"(Axi[0]), "r"(Axi[1]), "r"(Bxi[0])); + asm("mma.sync.aligned.m16n8k8.row.col.f16.f16.f16.f16 {%0, %1}, {%2, %3}, {%4}, {%0, %1};" + : "+r"(Dxi[0]), "+r"(Dxi[1]) + : "r"(Axi[2]), "r"(Axi[3]), "r"(Bxi[2])); + asm("mma.sync.aligned.m16n8k8.row.col.f16.f16.f16.f16 {%0, %1}, {%2, %3}, {%4}, {%0, %1};" + : "+r"(Dxi[2]), "+r"(Dxi[3]) + : "r"(Axi[0]), "r"(Axi[1]), "r"(Bxi[1])); + asm("mma.sync.aligned.m16n8k8.row.col.f16.f16.f16.f16 {%0, %1}, {%2, %3}, {%4}, {%0, %1};" + : "+r"(Dxi[2]), "+r"(Dxi[3]) + : "r"(Axi[2]), "r"(Axi[3]), "r"(Bxi[3])); +#endif // __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE +#elif defined(AMD_WMMA_AVAILABLE) +#if defined(RDNA4) + using halfx8_t = __attribute__((ext_vector_type(8))) _Float16; + halfx8_t& acc_frag = reinterpret_cast(D.x[0]); + const halfx8_t& a_frag = reinterpret_cast(A.x[0]); + const halfx8_t& b_frag = reinterpret_cast(B.x[0]); + acc_frag = __builtin_amdgcn_wmma_f16_16x16x16_f16_w32_gfx12(a_frag, b_frag, acc_frag); +#else + GGML_UNUSED_VARS(D, A, B); + NO_DEVICE_CODE; +#endif // defined(RDNA4) +#elif defined(AMD_MFMA_AVAILABLE) + // MFMA: FP16 input, FP32 accumulate, convert back to half2. + using halfx4_t = __attribute__((ext_vector_type(4))) _Float16; + using floatx4_t = __attribute__((ext_vector_type(4))) float; + + // Convert existing half2 accumulator to float for MFMA: + floatx4_t acc_f32; + { + const halfx4_t acc_h = reinterpret_cast(D.x[0]); +#pragma unroll + for (int i = 0; i < 4; ++i) { + acc_f32[i] = (float)acc_h[i]; + } + } + + const halfx4_t& a_frag = reinterpret_cast(A.x[0]); + const halfx4_t& b_frag = reinterpret_cast(B.x[0]); + acc_f32 = __builtin_amdgcn_mfma_f32_16x16x16f16(a_frag, b_frag, acc_f32, 0, 0, 0); + + // Convert back to half2: + { + halfx4_t result_h; +#pragma unroll + for (int i = 0; i < 4; ++i) { + result_h[i] = (_Float16)acc_f32[i]; + } + reinterpret_cast(D.x[0]) = result_h; + } +#else + GGML_UNUSED_VARS(D, A, B); + NO_DEVICE_CODE; +#endif // TURING_MMA_AVAILABLE + } + + static __device__ __forceinline__ void mma( + tile<16, 16, half2, DATA_LAYOUT_I_MAJOR_SCRAMBLED> & D, const tile<32, 8, half2, DATA_LAYOUT_I_MAJOR> & A, + const tile<16, 8, half2, DATA_LAYOUT_I_MAJOR> & B) { +#if defined(AMD_MFMA_AVAILABLE) || (defined(AMD_WMMA_AVAILABLE) && defined(RDNA4)) + tile<16, 8, half2> * D16 = (tile<16, 8, half2> *) &D; + const tile<16, 8, half2> * A16 = (const tile<16, 8, half2> *) &A; + mma(D16[0], A16[0], B); + mma(D16[1], A16[1], B); +#else + GGML_UNUSED_VARS(D, A, B); + NO_DEVICE_CODE; +#endif // defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) && defined(RDNA4) + } + + template + static __device__ __forceinline__ void mma( + tile<16, 8, float, dl_d> & D, const tile<16, 8, float, dl_ab> & A, const tile<8, 8, float, dl_ab> & B) { +#ifdef AMPERE_MMA_AVAILABLE + const int * Axi = (const int *) A.x; + const int * Bxi = (const int *) B.x; + int * Dxi = (int *) D.x; + asm("mma.sync.aligned.m16n8k8.row.col.f32.tf32.tf32.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};" + : "+r"(Dxi[0]), "+r"(Dxi[1]), "+r"(Dxi[2]), "+r"(Dxi[3]) + : "r"(Axi[0]), "r"(Axi[1]), "r"(Axi[2]), "r"(Axi[3]), "r"(Bxi[0]), "r"(Bxi[1])); +#else + GGML_UNUSED_VARS(D, A, B); + NO_DEVICE_CODE; +#endif // AMPERE_MMA_AVAILABLE + } + + template + static __device__ __forceinline__ void mma( + tile<16, 16, float, dl_d> & D, const tile<16, 8, float, dl_ab> & A, const tile<16, 8, float, dl_ab> & B) { +#ifdef AMD_MFMA_AVAILABLE + using floatx4_t = __attribute__((ext_vector_type(4))) float; + floatx4_t& acc_frag = reinterpret_cast(D.x[0]); +#if defined(CDNA3) + using floatx2_t = __attribute__((ext_vector_type(2))) float; + const floatx2_t& a_frag = reinterpret_cast(A.x[0]); + const floatx2_t& b_frag = reinterpret_cast(B.x[0]); + acc_frag = __builtin_amdgcn_mfma_f32_16x16x8_xf32(a_frag, b_frag, acc_frag, 0, 0, 0); +#elif defined(CDNA4) || defined(CDNA2) || defined(CDNA1) + // CDNA4 (gfx950) does not support xf32 MFMA, use f32 path like CDNA2/CDNA1 +#pragma unroll + for (int i = 0; i < 2; ++i) { + acc_frag = __builtin_amdgcn_mfma_f32_16x16x4f32(A.x[i], B.x[i], acc_frag, 0, 0, 0); + } +#else + GGML_UNUSED_VARS(D, A, B); + NO_DEVICE_CODE; +#endif // defined(CDNA3) +#else + GGML_UNUSED_VARS(D, A, B); + NO_DEVICE_CODE; +#endif // AMD_MFMA_AVAILABLE + } + + template + static __device__ __forceinline__ void mma_block_scaled_fp4(tile<16, 8, float> & D, + const tile<16, 8, int> & A, + const tile<8, 8, int> & B, + uint32_t a_scale, + uint32_t b_scale) { +#ifdef BLACKWELL_MMA_AVAILABLE + const int * Axi = (const int *) A.x; + const int * Bxi = (const int *) B.x; + float * Dxi = (float *) D.x; + + if constexpr (type == GGML_TYPE_MXFP4) { + asm volatile( + "mma.sync.aligned.kind::mxf4.block_scale.scale_vec::2X.m16n8k64.row.col.f32.e2m1.e2m1.f32.ue8m0 " + "{%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3}, " + "%10, {0, 0}, %11, {0, 0};" + : "+f"(Dxi[0]), "+f"(Dxi[1]), "+f"(Dxi[2]), "+f"(Dxi[3]) + : "r"(Axi[0]), "r"(Axi[1]), "r"(Axi[2]), "r"(Axi[3]), "r"(Bxi[0]), "r"(Bxi[1]), "r"(a_scale), "r"(b_scale)); + } else { + asm volatile( + "mma.sync.aligned.kind::mxf4nvf4.block_scale.scale_vec::4X.m16n8k64.row.col.f32.e2m1.e2m1.f32.ue4m3 " + "{%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3}, " + "%10, {0, 0}, %11, {0, 0};" + : "+f"(Dxi[0]), "+f"(Dxi[1]), "+f"(Dxi[2]), "+f"(Dxi[3]) + : "r"(Axi[0]), "r"(Axi[1]), "r"(Axi[2]), "r"(Axi[3]), "r"(Bxi[0]), "r"(Bxi[1]), "r"(a_scale), "r"(b_scale)); + } +#else + GGML_UNUSED_VARS(D, A, B, a_scale, b_scale); +#endif // BLACKWELL_MMA_AVAILABLE + } + + static __device__ __forceinline__ void mma( + tile<16, 8, float> & D, const tile<16, 8, half2> & A, const tile<8, 8, half2> & B) { +#ifdef TURING_MMA_AVAILABLE + const int * Axi = (const int *) A.x; + const int * Bxi = (const int *) B.x; + int * Dxi = (int *) D.x; +#if __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE + asm("mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};" + : "+r"(Dxi[0]), "+r"(Dxi[1]), "+r"(Dxi[2]), "+r"(Dxi[3]) + : "r"(Axi[0]), "r"(Axi[1]), "r"(Axi[2]), "r"(Axi[3]), "r"(Bxi[0]), "r"(Bxi[1])); +#else + // On Turing m16n8k16 mma is not available, use 2x m8n8k8 mma instead: + asm("mma.sync.aligned.m16n8k8.row.col.f32.f16.f16.f32 {%0, %1, %2, %3}, {%4, %5}, {%6}, {%0, %1, %2, %3};" + : "+r"(Dxi[0]), "+r"(Dxi[1]), "+r"(Dxi[2]), "+r"(Dxi[3]) + : "r"(Axi[0]), "r"(Axi[1]), "r"(Bxi[0])); + asm("mma.sync.aligned.m16n8k8.row.col.f32.f16.f16.f32 {%0, %1, %2, %3}, {%4, %5}, {%6}, {%0, %1, %2, %3};" + : "+r"(Dxi[0]), "+r"(Dxi[1]), "+r"(Dxi[2]), "+r"(Dxi[3]) + : "r"(Axi[2]), "r"(Axi[3]), "r"(Bxi[1])); +#endif // __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE +#else + GGML_UNUSED_VARS(D, A, B); + NO_DEVICE_CODE; +#endif // TURING_MMA_AVAILABLE + } + + static __device__ __forceinline__ void mma( + tile<16, 8, float> & D, const tile<16, 8, nv_bfloat162> & A, const tile<8, 8, nv_bfloat162> & B) { +#ifdef AMPERE_MMA_AVAILABLE + const int * Axi = (const int *) A.x; + const int * Bxi = (const int *) B.x; + int * Dxi = (int *) D.x; + asm("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};" + : "+r"(Dxi[0]), "+r"(Dxi[1]), "+r"(Dxi[2]), "+r"(Dxi[3]) + : "r"(Axi[0]), "r"(Axi[1]), "r"(Axi[2]), "r"(Axi[3]), "r"(Bxi[0]), "r"(Bxi[1])); +#else + GGML_UNUSED_VARS(D, A, B); + NO_DEVICE_CODE; +#endif // AMPERE_MMA_AVAILABLE + } + + template + static __device__ __forceinline__ void mma( + tile<16, 16, float, dl_d> & D, const tile<16, 8, half2, dl_ab> & A, const tile<16, 8, half2, dl_ab> & B) { +#ifdef TURING_MMA_AVAILABLE + const int * Axi = (const int *) A.x; + const int * Bxi = (const int *) B.x; + int * Dxi = (int *) D.x; +#if __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE + asm("mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};" + : "+r"(Dxi[0]), "+r"(Dxi[1]), "+r"(Dxi[2]), "+r"(Dxi[3]) + : "r"(Axi[0]), "r"(Axi[1]), "r"(Axi[2]), "r"(Axi[3]), "r"(Bxi[0]), "r"(Bxi[2])); + asm("mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};" + : "+r"(Dxi[4]), "+r"(Dxi[5]), "+r"(Dxi[6]), "+r"(Dxi[7]) + : "r"(Axi[0]), "r"(Axi[1]), "r"(Axi[2]), "r"(Axi[3]), "r"(Bxi[1]), "r"(Bxi[3])); +#else + // On Turing m16n8k16 mma is not available, use 4x m8n8k8 mma instead: + asm("mma.sync.aligned.m16n8k8.row.col.f32.f16.f16.f32 {%0, %1, %2, %3}, {%4, %5}, {%6}, {%0, %1, %2, %3};" + : "+r"(Dxi[0]), "+r"(Dxi[1]), "+r"(Dxi[2]), "+r"(Dxi[3]) + : "r"(Axi[0]), "r"(Axi[1]), "r"(Bxi[0])); + asm("mma.sync.aligned.m16n8k8.row.col.f32.f16.f16.f32 {%0, %1, %2, %3}, {%4, %5}, {%6}, {%0, %1, %2, %3};" + : "+r"(Dxi[0]), "+r"(Dxi[1]), "+r"(Dxi[2]), "+r"(Dxi[3]) + : "r"(Axi[2]), "r"(Axi[3]), "r"(Bxi[2])); + asm("mma.sync.aligned.m16n8k8.row.col.f32.f16.f16.f32 {%0, %1, %2, %3}, {%4, %5}, {%6}, {%0, %1, %2, %3};" + : "+r"(Dxi[4]), "+r"(Dxi[5]), "+r"(Dxi[6]), "+r"(Dxi[7]) + : "r"(Axi[0]), "r"(Axi[1]), "r"(Bxi[1])); + asm("mma.sync.aligned.m16n8k8.row.col.f32.f16.f16.f32 {%0, %1, %2, %3}, {%4, %5}, {%6}, {%0, %1, %2, %3};" + : "+r"(Dxi[4]), "+r"(Dxi[5]), "+r"(Dxi[6]), "+r"(Dxi[7]) + : "r"(Axi[2]), "r"(Axi[3]), "r"(Bxi[3])); +#endif // __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE +#elif defined(AMD_WMMA_AVAILABLE) +#if defined(RDNA4) + using halfx8_t = __attribute__((ext_vector_type(8))) _Float16; + using floatx8_t = __attribute__((ext_vector_type(8))) float; + floatx8_t& acc_frag = reinterpret_cast(D.x[0]); + const halfx8_t& a_frag = reinterpret_cast(A.x[0]); + const halfx8_t& b_frag = reinterpret_cast(B.x[0]); + acc_frag = __builtin_amdgcn_wmma_f32_16x16x16_f16_w32_gfx12(a_frag, b_frag, acc_frag); +#elif defined(RDNA3) + using halfx16_t = __attribute__((ext_vector_type(16))) _Float16; + using floatx8_t = __attribute__((ext_vector_type(8))) float; + floatx8_t& acc_frag = reinterpret_cast(D.x[0]); + const halfx16_t& a_frag = reinterpret_cast(A.x[0]); + const halfx16_t& b_frag = reinterpret_cast(B.x[0]); + acc_frag = __builtin_amdgcn_wmma_f32_16x16x16_f16_w32(a_frag, b_frag, acc_frag); +#else + GGML_UNUSED_VARS(D, A, B); + NO_DEVICE_CODE; +#endif // RDNA4 +#elif defined(AMD_MFMA_AVAILABLE) + using halfx4_t = __attribute__((ext_vector_type(4))) _Float16; + using floatx4_t = __attribute__((ext_vector_type(4))) float; + floatx4_t& acc_frag = reinterpret_cast(D.x[0]); + const halfx4_t& a_frag = reinterpret_cast(A.x[0]); + const halfx4_t& b_frag = reinterpret_cast(B.x[0]); + acc_frag = __builtin_amdgcn_mfma_f32_16x16x16f16(a_frag, b_frag, acc_frag, 0, 0, 0); +#else + GGML_UNUSED_VARS(D, A, B); + NO_DEVICE_CODE; +#endif // TURING_MMA_AVAILABLE + } + + template + static __device__ __forceinline__ void mma( + tile<16, 16, float, dl_d> & D, const tile<16, 8, nv_bfloat162, dl_ab> & A, const tile<16, 8, nv_bfloat162, dl_ab> & B) { +#if defined(AMD_WMMA_AVAILABLE) +#if defined(RDNA4) + using bf16x8_t = __attribute__((ext_vector_type(8))) __bf16; + using floatx8_t = __attribute__((ext_vector_type(8))) float; + floatx8_t& acc_frag = reinterpret_cast(D.x[0]); + const bf16x8_t& a_frag = reinterpret_cast(A.x[0]); + const bf16x8_t& b_frag = reinterpret_cast(B.x[0]); + acc_frag = __builtin_amdgcn_wmma_f32_16x16x16_bf16_w32_gfx12(a_frag, b_frag, acc_frag); +#elif defined(RDNA3) + using bf16x16_t = __attribute__((ext_vector_type(16))) __bf16; + using floatx8_t = __attribute__((ext_vector_type(8))) float; + floatx8_t& acc_frag = reinterpret_cast(D.x[0]); + const bf16x16_t& a_frag = reinterpret_cast(A.x[0]); + const bf16x16_t& b_frag = reinterpret_cast(B.x[0]); + acc_frag = __builtin_amdgcn_wmma_f32_16x16x16_bf16_w32(a_frag, b_frag, acc_frag); +#else + GGML_UNUSED_VARS(D, A, B); + NO_DEVICE_CODE; +#endif // defined(RDNA4) +#elif defined(AMD_MFMA_AVAILABLE) + using floatx4_t = __attribute__((ext_vector_type(4))) float; + floatx4_t& acc_frag = reinterpret_cast(D.x[0]); +#if defined(CDNA4) || defined(CDNA3) || defined(CDNA2) + using bf16x4_t = __attribute__((ext_vector_type(4))) __bf16; + const bf16x4_t& a_frag = reinterpret_cast(A.x[0]); + const bf16x4_t& b_frag = reinterpret_cast(B.x[0]); + acc_frag = __builtin_amdgcn_mfma_f32_16x16x16bf16_1k(a_frag, b_frag, acc_frag, 0, 0, 0); +#elif defined(CDNA1) +#pragma unroll + for (int i = 0; i < 2; ++i) { + using bf16x2_t = __attribute__((ext_vector_type(2))) __bf16; + const bf16x2_t& a_frag = reinterpret_cast(A.x[i]); + const bf16x2_t& b_frag = reinterpret_cast(B.x[i]); + acc_frag = __builtin_amdgcn_mfma_f32_16x16x8bf16(a_frag, b_frag, acc_frag, 0, 0, 0); + } +#else + GGML_UNUSED_VARS(D, A, B); + NO_DEVICE_CODE; +#endif // defined(CDNA3) || defined(CDNA2) +#else + GGML_UNUSED_VARS(D, A, B); + NO_DEVICE_CODE; +#endif // defined(AMD_WMMA_AVAILABLE) + } + + template + static __device__ __forceinline__ void mma( + tile<16, 16, int, dl_d> & D, const tile<16, 8, int, dl_ab> & A, const tile<16, 8, int, dl_ab> & B) { +#if defined(AMD_MFMA_AVAILABLE) + using int32x4_t = __attribute__((__vector_size__(4 * sizeof(int)))) int; + int32x4_t * acc = (int32x4_t *) D.x; +#if defined(CDNA4) || defined(CDNA3) + acc[0] = __builtin_amdgcn_mfma_i32_16x16x32_i8(((int64_t *) A.x)[0], ((int64_t *) B.x)[0], acc[0], 0, 0, 0); +#elif defined(CDNA2) || defined(CDNA1) + acc[0] = __builtin_amdgcn_mfma_i32_16x16x16i8(A.x[0], B.x[0], acc[0], 0, 0, 0); + acc[0] = __builtin_amdgcn_mfma_i32_16x16x16i8(A.x[1], B.x[1], acc[0], 0, 0, 0); +#endif // defined(CDNA4) || defined(CDNA3) +#elif defined(AMD_WMMA_AVAILABLE) + using int32x8_t = __attribute__((__vector_size__(8 * sizeof(int)))) int; + int32x8_t * acc = (int32x8_t *) D.x; +#if defined(RDNA4) + using int32x2_t = __attribute__((__vector_size__(2 * sizeof(int)))) int; + int32x2_t * a_vec = (int32x2_t *) A.x; + int32x2_t * b_vec = (int32x2_t *) B.x; + acc[0] = __builtin_amdgcn_wmma_i32_16x16x16_iu8_w32_gfx12(true, a_vec[0], true, b_vec[0], acc[0], true); + acc[0] = __builtin_amdgcn_wmma_i32_16x16x16_iu8_w32_gfx12(true, a_vec[1], true, b_vec[1], acc[0], true); +#elif defined(RDNA3) + using int32x4_t = __attribute__((__vector_size__(4 * sizeof(int)))) int; + int32x4_t * a_vec = (int32x4_t *) A.x; + int32x4_t * b_vec = (int32x4_t *) B.x; + acc[0] = __builtin_amdgcn_wmma_i32_16x16x16_iu8_w32(true, a_vec[0], true, b_vec[0], acc[0], true); + acc[0] = __builtin_amdgcn_wmma_i32_16x16x16_iu8_w32(true, a_vec[1], true, b_vec[1], acc[0], true); +#endif // RDNA4 +#else + GGML_UNUSED_VARS(D, A, B); + NO_DEVICE_CODE; +#endif // AMD_MFMA_AVAILABLE + } + + static __device__ __forceinline__ void mma( + tile<32, 32, int> & D, const tile<32, 4, int> & A, const tile<32, 4, int> & B) { +#if defined(AMD_MFMA_AVAILABLE) + using int32x16_t = __attribute__((__vector_size__(16 * sizeof(int)))) int; + int32x16_t * acc = (int32x16_t *) D.x; +#if defined(CDNA4) || defined(CDNA3) + acc[0] = __builtin_amdgcn_mfma_i32_32x32x16_i8(((int64_t *) A.x)[0], ((int64_t *) B.x)[0], acc[0], 0, 0, 0); +#elif defined(CDNA2) || defined(CDNA1) + acc[0] = __builtin_amdgcn_mfma_i32_32x32x8i8(A.x[0], B.x[0], acc[0], 0, 0, 0); + acc[0] = __builtin_amdgcn_mfma_i32_32x32x8i8(A.x[1], B.x[1], acc[0], 0, 0, 0); +#endif // defined(CDNA4) || defined(CDNA3) + +#else + GGML_UNUSED_VARS(D, A, B); + NO_DEVICE_CODE; +#endif // AMD_MFMA_AVAILABLE + } + + template + static __device__ __forceinline__ void mma( + tile<32, J, T1> & D, const tile<32, K, T2> & A, const tile & B) { + tile <16, J, T1> * D16 = reinterpret_cast< tile<16, J, T1> *>(&D); + const tile<16, K, T2> * A16 = reinterpret_cast *>(&A); + mma(D16[0], A16[0], B); + mma(D16[1], A16[1], B); + } + + static __device__ __forceinline__ void mma( + tile<32, 8, float> & D, const tile<32, 4, half2> & A, const tile<8, 4, half2, DATA_LAYOUT_I_MAJOR_MIRRORED> & B) { +#if defined(VOLTA_MMA_AVAILABLE) + const int * Axi = (const int *) A.x; + const int * Bxi = (const int *) B.x; + int * Dxi = (int *) D.x; + asm("mma.sync.aligned.m8n8k4.row.col.f32.f16.f16.f32 " + "{%0, %1, %2, %3, %4, %5, %6, %7}, {%8, %9}, {%10, %11}, {%0, %1, %2, %3, %4, %5, %6, %7};" + : "+r"(Dxi[0]), "+r"(Dxi[1]), "+r"(Dxi[2]), "+r"(Dxi[3]), "+r"(Dxi[4]), "+r"(Dxi[5]), "+r"(Dxi[6]), "+r"(Dxi[7]) + : "r"(Axi[0]), "r"(Axi[1]), "r"(Bxi[0]), "r"(Bxi[1])); + asm("mma.sync.aligned.m8n8k4.row.col.f32.f16.f16.f32 " + "{%0, %1, %2, %3, %4, %5, %6, %7}, {%8, %9}, {%10, %11}, {%0, %1, %2, %3, %4, %5, %6, %7};" + : "+r"(Dxi[0]), "+r"(Dxi[1]), "+r"(Dxi[2]), "+r"(Dxi[3]), "+r"(Dxi[4]), "+r"(Dxi[5]), "+r"(Dxi[6]), "+r"(Dxi[7]) + : "r"(Axi[2]), "r"(Axi[3]), "r"(Bxi[2]), "r"(Bxi[3])); +#else + GGML_UNUSED_VARS(D, A, B); + NO_DEVICE_CODE; +#endif // defined(VOLTA_MMA_AVAILABLE) + } + + static __device__ __forceinline__ void mma( + tile<32, 4, half2> & D, const tile<32, 4, half2> & A, const tile<8, 4, half2, DATA_LAYOUT_J_MAJOR_MIRRORED> & B) { +#if defined(VOLTA_MMA_AVAILABLE) + const int * Axi = (const int *) A.x; + const int * Bxi = (const int *) B.x; + int * Dxi = (int *) D.x; + asm("mma.sync.aligned.m8n8k4.row.row.f16.f16.f16.f16 " + "{%0, %1, %2, %3}, {%4, %5}, {%6, %7}, {%0, %1, %2, %3};" + : "+r"(Dxi[0]), "+r"(Dxi[1]), "+r"(Dxi[2]), "+r"(Dxi[3]) + : "r"(Axi[0]), "r"(Axi[1]), "r"(Bxi[0]), "r"(Bxi[1])); + asm("mma.sync.aligned.m8n8k4.row.row.f16.f16.f16.f16 " + "{%0, %1, %2, %3}, {%4, %5}, {%6, %7}, {%0, %1, %2, %3};" + : "+r"(Dxi[0]), "+r"(Dxi[1]), "+r"(Dxi[2]), "+r"(Dxi[3]) + : "r"(Axi[2]), "r"(Axi[3]), "r"(Bxi[2]), "r"(Bxi[3])); +#else + GGML_UNUSED_VARS(D, A, B); + NO_DEVICE_CODE; +#endif // defined(VOLTA_MMA_AVAILABLE) + } + + static __device__ __forceinline__ void mma( + tile<16, 16, half2, DATA_LAYOUT_I_MAJOR> & D, const tile<32, 8, half2, DATA_LAYOUT_I_MAJOR_MIRRORED> & A, + const tile<16, 8, half2, DATA_LAYOUT_I_MAJOR_MIRRORED> & B) { +#if defined(AMD_WMMA_AVAILABLE) && defined(RDNA3) + using halfx16_t = __attribute__((ext_vector_type(16))) _Float16; + halfx16_t * xD = (halfx16_t *) D.x; + const halfx16_t * xA = (const halfx16_t *) A.x; + const halfx16_t * xB = (const halfx16_t *) B.x; + xD[0] = __builtin_amdgcn_wmma_f16_16x16x16_f16_w32(xA[0], xB[0], xD[0], /*opsel =*/ 0); + xD[0] = __builtin_amdgcn_wmma_f16_16x16x16_f16_w32(xA[1], xB[0], xD[0], /*opsel =*/ 1); +#else + GGML_UNUSED_VARS(D, A, B); + NO_DEVICE_CODE; +#endif // TURING_MMA_AVAILABLE + } + + template + static __device__ __forceinline__ void mma( + tile<16, 16, int, dl_d> & D, const tile<16, 4, int, dl_ab> & A, const tile<16, 4, int, dl_ab> & B) { +#if defined(AMD_MFMA_AVAILABLE) + using int32x4_t = __attribute__((__vector_size__(4 * sizeof(int)))) int; + int32x4_t * acc = (int32x4_t *) D.x; +#if defined(CDNA4) || defined(CDNA3) + const int64_t xA = uint32_t(A.x[0]); + const int64_t xB = uint32_t(B.x[0]); + acc[0] = __builtin_amdgcn_mfma_i32_16x16x32_i8(xA, xB, acc[0], 0, 0, 0); +#elif defined(CDNA2) || defined(CDNA1) + acc[0] = __builtin_amdgcn_mfma_i32_16x16x16i8(A.x[0], B.x[0], acc[0], 0, 0, 0); +#endif // defined(CDNA4) || defined(CDNA3) +#elif defined(AMD_WMMA_AVAILABLE) + using int32x8_t = __attribute__((__vector_size__(8 * sizeof(int)))) int; + int32x8_t * acc = (int32x8_t *) D.x; +#if defined(RDNA4) + using int32x2_t = __attribute__((__vector_size__(2 * sizeof(int)))) int; + int32x2_t * a_vec = (int32x2_t *) A.x; + int32x2_t * b_vec = (int32x2_t *) B.x; + acc[0] = __builtin_amdgcn_wmma_i32_16x16x16_iu8_w32_gfx12(true, a_vec[0], true, b_vec[0], acc[0], false); +#elif defined(RDNA3) + using int32x4_t = __attribute__((__vector_size__(4 * sizeof(int)))) int; + int32x4_t * a_vec = (int32x4_t *) A.x; + int32x4_t * b_vec = (int32x4_t *) B.x; + acc[0] = __builtin_amdgcn_wmma_i32_16x16x16_iu8_w32(true, a_vec[0], true, b_vec[0], acc[0], false); +#endif // RDNA4 +#else + GGML_UNUSED(D); + GGML_UNUSED(A); + GGML_UNUSED(B); + NO_DEVICE_CODE; +#endif // AMD_WMMA_AVAILABLE + } +} diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmid.cu b/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmid.cu new file mode 100644 index 0000000000..345b39eb5f --- /dev/null +++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmid.cu @@ -0,0 +1,221 @@ +#include "common.cuh" +#include "mmid.cuh" + +// To reduce shared memory use, store "it" and "iex_used" with 22/10 bits each. +struct mm_ids_helper_store { + uint32_t data; + + __device__ mm_ids_helper_store(const uint32_t it, const uint32_t iex_used) { + data = (it & 0x003FFFFF) | (iex_used << 22); + } + + __device__ uint32_t it() const { + return data & 0x003FFFFF; + } + + __device__ uint32_t iex_used() const { + return data >> 22; + } +}; +static_assert(sizeof(mm_ids_helper_store) == 4, "unexpected size for mm_ids_helper_store"); + +// the generic path passes 0, which needs no padding since it never groups lanes by token +template struct mm_ids_pow2 { static constexpr int value = 2*mm_ids_pow2<(n + 1)/2>::value; }; +template <> struct mm_ids_pow2<1> { static constexpr int value = 1; }; +template <> struct mm_ids_pow2<0> { static constexpr int value = 1; }; + +// Helper function for mul_mat_id, converts ids to a more convenient format. +// ids_src1 describes how to permute the flattened column indices of src1 in order to get a compact src1 tensor sorted by expert. +// ids_dst describes the same mapping but for the dst tensor. +// The upper and lower bounds for the ith expert in the compact src1 tensor are stored in expert_bounds[i:i+1]. +template +__launch_bounds__(ggml_cuda_get_physical_warp_size(), 1) +static __global__ void mm_ids_helper( + const int32_t * __restrict__ ids, int32_t * __restrict__ ids_src1, int32_t * __restrict__ ids_dst, int32_t * __restrict__ expert_bounds, + const int n_tokens, const int n_expert_used_var, const int nchannels_y, const int si1, const int sis1, const bool write_inverse) { + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + const int n_expert_used = n_expert_used_template == 0 ? n_expert_used_var : n_expert_used_template; + const int expert = blockIdx.x; + + extern __shared__ char data_mm_ids_helper[]; + mm_ids_helper_store * store = (mm_ids_helper_store *) data_mm_ids_helper; + + int nex_prev = 0; // Number of columns for experts with a lower index. + int it_compact = 0; // Running index for the compact slice of this expert. + + if constexpr (n_expert_used_template == 0) { + // Generic implementation: + for (int it = 0; it < n_tokens; ++it) { + int iex_used = -1; // The index at which the expert is used, if any. + for (int iex = threadIdx.x; iex < n_expert_used; iex += warp_size) { + const int expert_used = ids[it*si1 + iex]; + nex_prev += expert_used < expert; + if (expert_used == expert) { + iex_used = iex; + } + } + + if (iex_used != -1) { + store[it_compact] = mm_ids_helper_store(it, iex_used); + } + + if (warp_reduce_any(iex_used != -1)) { + it_compact++; + } + } + } else { + // Implementation optimized for specific numbers of experts used: + // a warp holds a whole number of token slots, so the slot count is padded to a power of 2 + constexpr int neu_padded = mm_ids_pow2::value; + static_assert(neu_padded <= warp_size && warp_size % neu_padded == 0, "bad n_expert_used"); + for (int it0 = 0; it0 < n_tokens; it0 += warp_size/neu_padded) { + const int it = it0 + threadIdx.x / neu_padded; + + const int iex = threadIdx.x % neu_padded; // The index at which the expert is used, if any. + const int expert_used = (neu_padded == n_expert_used || iex < n_expert_used) && it < n_tokens ? + ids[it*si1 + iex] : INT_MAX; + const int iex_used = expert_used == expert ? iex : -1; + nex_prev += expert_used < expert; + + // Whether the threads at this token position have used the expert: + const int it_compact_add_self = warp_reduce_any(iex_used != -1); + + // Do a scan over threads at lower token positions in warp to get the correct index for writing data: + int it_compact_add_lower = 0; +#pragma unroll + for (int offset = neu_padded; offset < warp_size; offset += neu_padded) { + const int tmp = __shfl_up_sync(0xFFFFFFFF, it_compact_add_self, offset, warp_size); + if (threadIdx.x >= static_cast(offset)) { + it_compact_add_lower += tmp; + } + } + + if (iex_used != -1) { + store[it_compact + it_compact_add_lower] = mm_ids_helper_store(it, iex_used); + } + + // The thread with the highest index in the warp always has the sum over the whole warp, use it to increment all threads: + it_compact += __shfl_sync(0xFFFFFFFF, it_compact_add_lower + it_compact_add_self, warp_size - 1, warp_size); + } + } + nex_prev = warp_reduce_sum(nex_prev); + ggml_cuda_syncwarp(); + + for (int itc = threadIdx.x; itc < it_compact; itc += warp_size) { + const mm_ids_helper_store store_it = store[itc]; + const int it = store_it.it(); + const int iex_used = store_it.iex_used(); + ids_dst[nex_prev + itc] = it*n_expert_used + iex_used; + // ids_src1 holds the forward map, or the inverse map (token slot -> compact row) for quant dedup + if (write_inverse) { + ids_src1[it*n_expert_used + iex_used] = nex_prev + itc; + } else { + ids_src1[nex_prev + itc] = it*sis1 + iex_used % nchannels_y; + } + } + + if (threadIdx.x != 0) { + return; + } + + expert_bounds[expert] = nex_prev; + + if (expert < static_cast(gridDim.x) - 1) { + return; + } + + expert_bounds[gridDim.x] = nex_prev + it_compact; +} + +template +static void launch_mm_ids_helper( + const int32_t * __restrict__ ids, int32_t * __restrict__ ids_src1, int32_t * __restrict__ ids_dst, int32_t * __restrict__ expert_bounds, + const int n_experts, const int n_tokens, const int n_expert_used_var, const int nchannels_y, const int si1, const int sis1, const bool write_inverse, cudaStream_t stream) { + GGML_ASSERT(n_tokens < (1 << 22) && "too few bits in mm_ids_helper_store"); + GGML_ASSERT(n_expert_used_var < (1 << 10) && "too few bits in mm_ids_helper_store"); + + const int id = ggml_cuda_get_device(); + const int warp_size = ggml_cuda_info().devices[id].warp_size; + const size_t smpbo = ggml_cuda_info().devices[id].smpbo; + CUDA_SET_SHARED_MEMORY_LIMIT(mm_ids_helper, smpbo); + + const dim3 num_blocks(n_experts, 1, 1); + const dim3 block_size(warp_size, 1, 1); + const size_t nbytes_shared = n_tokens*sizeof(mm_ids_helper_store); + GGML_ASSERT(nbytes_shared <= smpbo); + mm_ids_helper<<>> + (ids, ids_src1, ids_dst, expert_bounds, n_tokens, n_expert_used_var, nchannels_y, si1, sis1, write_inverse); +} + +void ggml_cuda_launch_mm_ids_helper( + const int32_t * __restrict__ ids, int32_t * __restrict__ ids_src1, int32_t * __restrict__ ids_dst, int32_t * __restrict__ expert_bounds, + const int n_experts, const int n_tokens, const int n_expert_used, const int nchannels_y, const int si1, const int sis1, const bool write_inverse, cudaStream_t stream) { + switch (n_expert_used) { + case 2: + launch_mm_ids_helper< 2>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, write_inverse, stream); + break; + case 4: + launch_mm_ids_helper< 4>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, write_inverse, stream); + break; + case 6: + launch_mm_ids_helper< 6>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, write_inverse, stream); + break; + case 8: + launch_mm_ids_helper< 8>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, write_inverse, stream); + break; + case 10: + launch_mm_ids_helper<10>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, write_inverse, stream); + break; + case 16: + launch_mm_ids_helper<16>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, write_inverse, stream); + break; + case 32: + launch_mm_ids_helper<32>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, write_inverse, stream); + break; + default: + launch_mm_ids_helper< 0>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, write_inverse, stream); + break; + } +} + +// Lists the (expert, column tile) pairs that hold columns for a flat MoE grid, in expert order. +// The other ntiles grid columns get tile jt_none of expert 0, which starts past its columns, so mul_mat_q skips them. +template +__launch_bounds__(nthreads, 1) +static __global__ void mm_ids_tile_map( + const int32_t * __restrict__ expert_bounds, int2 * __restrict__ tile_map, const int n_experts, const int J, const int ntiles, const int jt_none) { + __shared__ int scan[nthreads]; + + int start = 0; // First tile of the current chunk of experts. + for (int e0 = 0; e0 < n_experts; e0 += nthreads) { + const int e = e0 + threadIdx.x; + const int n = e < n_experts ? (expert_bounds[e + 1] - expert_bounds[e] + J - 1) / J : 0; + + scan[threadIdx.x] = n; + __syncthreads(); +#pragma unroll + for (int offset = 1; offset < nthreads; offset *= 2) { + const int add = threadIdx.x >= offset ? scan[threadIdx.x - offset] : 0; + __syncthreads(); + scan[threadIdx.x] += add; + __syncthreads(); + } + + const int first = start + scan[threadIdx.x] - n; + for (int jt = 0; jt < n; ++jt) { + tile_map[first + jt] = make_int2(e, jt); + } + start += scan[nthreads - 1]; + __syncthreads(); + } + + for (int t = start + threadIdx.x; t < ntiles; t += nthreads) { + tile_map[t] = make_int2(0, jt_none); + } +} + +void ggml_cuda_launch_mm_ids_tile_map( + const int32_t * expert_bounds, int2 * tile_map, const int n_experts, const int J, const int ntiles, const int jt_none, cudaStream_t stream) { + constexpr int nthreads = 512; + mm_ids_tile_map<<<1, nthreads, 0, stream>>>(expert_bounds, tile_map, n_experts, J, ntiles, jt_none); +} diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmid.cuh b/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmid.cuh new file mode 100644 index 0000000000..b0e8edb88e --- /dev/null +++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmid.cuh @@ -0,0 +1,8 @@ +#pragma once + +void ggml_cuda_launch_mm_ids_helper( + const int32_t * ids, int32_t * ids_src1, int32_t * ids_dst, int32_t * expert_bounds, + int n_experts, int n_tokens, int n_expert_used, int nchannels_y, int si1, int sis1, bool write_inverse, cudaStream_t stream); + +void ggml_cuda_launch_mm_ids_tile_map( + const int32_t * expert_bounds, int2 * tile_map, int n_experts, int J, int ntiles, int jt_none, cudaStream_t stream); diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq-config-ampere.cuh b/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq-config-ampere.cuh new file mode 100644 index 0000000000..9f9fd19738 --- /dev/null +++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq-config-ampere.cuh @@ -0,0 +1,383 @@ +static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config_ampere(ggml_type type, int J, bool fallback) { + CASE(GGML_TYPE_Q1_0, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q1_0, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q1_0, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q1_0, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q1_0, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q1_0, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q1_0, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q1_0, 256, 1, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q1_0, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q1_0, 256, 1, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q1_0, 256, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q1_0, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q1_0, 256, 1, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q1_0, 256, 1, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q1_0, 256, 1, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q1_0, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + + CASE(GGML_TYPE_Q2_0, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q2_0, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q2_0, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q2_0, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q2_0, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q2_0, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q2_0, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q2_0, 256, 1, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q2_0, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q2_0, 256, 1, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q2_0, 256, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q2_0, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q2_0, 256, 1, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q2_0, 256, 1, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q2_0, 256, 1, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q2_0, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + + CASE(GGML_TYPE_Q4_0, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q4_0, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q4_0, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q4_0, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q4_0, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q4_0, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q4_0, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q4_0, 256, 1, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q4_0, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q4_0, 256, 1, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q4_0, 256, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q4_0, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q4_0, 256, 1, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q4_0, 256, 1, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q4_0, 256, 1, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q4_0, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + + CASE(GGML_TYPE_Q4_1, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q4_1, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q4_1, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q4_1, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q4_1, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q4_1, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q4_1, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q4_1, 256, 1, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q4_1, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q4_1, 256, 1, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q4_1, 256, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q4_1, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q4_1, 256, 1, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q4_1, 256, 1, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q4_1, 256, 1, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q4_1, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + + CASE(GGML_TYPE_Q5_0, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q5_0, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q5_0, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q5_0, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q5_0, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q5_0, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q5_0, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q5_0, 256, 1, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q5_0, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q5_0, 256, 1, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q5_0, 256, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q5_0, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q5_0, 256, 1, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q5_0, 256, 1, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q5_0, 256, 1, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q5_0, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + + CASE(GGML_TYPE_Q5_1, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q5_1, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q5_1, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q5_1, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q5_1, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q5_1, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q5_1, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q5_1, 256, 1, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q5_1, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q5_1, 256, 1, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q5_1, 256, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q5_1, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q5_1, 256, 1, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q5_1, 256, 1, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q5_1, 256, 1, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q5_1, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + + CASE(GGML_TYPE_Q8_0, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q8_0, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q8_0, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q8_0, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q8_0, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q8_0, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q8_0, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q8_0, 256, 1, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q8_0, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q8_0, 256, 1, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q8_0, 256, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q8_0, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q8_0, 256, 1, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q8_0, 256, 1, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q8_0, 256, 1, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q8_0, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + +// --------------------------------------------------------------------------------------------- + + CASE(GGML_TYPE_Q2_K, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q2_K, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q2_K, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q2_K, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q2_K, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q2_K, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q2_K, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q2_K, 256, 1, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q2_K, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q2_K, 256, 1, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q2_K, 256, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q2_K, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q2_K, 256, 1, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q2_K, 256, 1, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q2_K, 256, 1, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q2_K, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, true, false); + + CASE(GGML_TYPE_Q3_K, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q3_K, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q3_K, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q3_K, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q3_K, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q3_K, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q3_K, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q3_K, 256, 1, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q3_K, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q3_K, 256, 1, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q3_K, 256, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q3_K, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q3_K, 256, 1, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q3_K, 256, 1, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q3_K, 256, 1, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q3_K, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false); + + CASE(GGML_TYPE_Q4_K, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q4_K, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q4_K, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q4_K, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q4_K, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q4_K, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q4_K, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q4_K, 256, 1, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q4_K, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q4_K, 256, 1, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q4_K, 256, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q4_K, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q4_K, 256, 1, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q4_K, 256, 1, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q4_K, 256, 1, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q4_K, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + + CASE(GGML_TYPE_Q5_K, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q5_K, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q5_K, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q5_K, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q5_K, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q5_K, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q5_K, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q5_K, 256, 1, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q5_K, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q5_K, 256, 1, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q5_K, 256, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q5_K, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q5_K, 256, 1, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q5_K, 256, 1, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q5_K, 256, 1, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q5_K, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + + CASE(GGML_TYPE_Q6_K, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q6_K, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q6_K, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q6_K, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q6_K, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q6_K, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q6_K, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q6_K, 256, 1, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q6_K, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q6_K, 256, 1, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q6_K, 256, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q6_K, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q6_K, 256, 1, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q6_K, 256, 1, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q6_K, 256, 1, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q6_K, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, true, false); + +// --------------------------------------------------------------------------------------------- + + CASE(GGML_TYPE_IQ1_S, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_IQ1_S, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_IQ1_S, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_IQ1_S, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_IQ1_S, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_IQ1_S, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ1_S, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ1_S, 256, 1, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ1_S, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ1_S, 256, 1, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ1_S, 256, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ1_S, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ1_S, 256, 1, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ1_S, 256, 1, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ1_S, 256, 1, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ1_S, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + + CASE(GGML_TYPE_IQ2_XXS, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_IQ2_XXS, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_IQ2_XXS, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_IQ2_XXS, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_IQ2_XXS, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_IQ2_XXS, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ2_XXS, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ2_XXS, 256, 1, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ2_XXS, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ2_XXS, 256, 1, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ2_XXS, 256, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ2_XXS, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ2_XXS, 256, 1, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ2_XXS, 256, 1, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ2_XXS, 256, 1, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ2_XXS, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + + CASE(GGML_TYPE_IQ2_XS, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_IQ2_XS, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_IQ2_XS, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_IQ2_XS, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_IQ2_XS, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_IQ2_XS, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ2_XS, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ2_XS, 256, 1, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ2_XS, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ2_XS, 256, 1, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ2_XS, 256, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ2_XS, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ2_XS, 256, 1, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ2_XS, 256, 1, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ2_XS, 256, 1, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ2_XS, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false); + + CASE(GGML_TYPE_IQ2_S, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_IQ2_S, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_IQ2_S, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_IQ2_S, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_IQ2_S, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_IQ2_S, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ2_S, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ2_S, 256, 1, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ2_S, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ2_S, 256, 1, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ2_S, 256, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ2_S, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ2_S, 256, 1, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ2_S, 256, 1, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ2_S, 256, 1, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ2_S, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false); + + CASE(GGML_TYPE_IQ3_XXS, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_IQ3_XXS, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_IQ3_XXS, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_IQ3_XXS, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_IQ3_XXS, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_IQ3_XXS, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ3_XXS, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ3_XXS, 256, 1, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ3_XXS, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ3_XXS, 256, 1, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ3_XXS, 256, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ3_XXS, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ3_XXS, 256, 1, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ3_XXS, 256, 1, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ3_XXS, 256, 1, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ3_XXS, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + + CASE(GGML_TYPE_IQ3_S, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_IQ3_S, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_IQ3_S, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_IQ3_S, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_IQ3_S, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_IQ3_S, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ3_S, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ3_S, 256, 1, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ3_S, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ3_S, 256, 1, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ3_S, 256, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ3_S, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ3_S, 256, 1, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ3_S, 256, 1, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ3_S, 256, 1, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ3_S, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + + CASE(GGML_TYPE_IQ4_XS, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_IQ4_XS, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_IQ4_XS, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_IQ4_XS, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_IQ4_XS, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_IQ4_XS, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ4_XS, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ4_XS, 256, 1, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ4_XS, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ4_XS, 256, 1, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ4_XS, 256, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ4_XS, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ4_XS, 256, 1, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ4_XS, 256, 1, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ4_XS, 256, 1, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ4_XS, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + + CASE(GGML_TYPE_IQ4_NL, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_IQ4_NL, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_IQ4_NL, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_IQ4_NL, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_IQ4_NL, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_IQ4_NL, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ4_NL, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ4_NL, 256, 1, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ4_NL, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ4_NL, 256, 1, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ4_NL, 256, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ4_NL, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ4_NL, 256, 1, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ4_NL, 256, 1, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ4_NL, 256, 1, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ4_NL, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + +// --------------------------------------------------------------------------------------------- + + CASE(GGML_TYPE_MXFP4, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_MXFP4, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_MXFP4, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_MXFP4, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_MXFP4, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_MXFP4, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_MXFP4, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_MXFP4, 256, 1, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_MXFP4, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_MXFP4, 256, 1, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_MXFP4, 256, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_MXFP4, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_MXFP4, 256, 1, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_MXFP4, 256, 1, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_MXFP4, 256, 1, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_MXFP4, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + + CASE(GGML_TYPE_NVFP4, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_NVFP4, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_NVFP4, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_NVFP4, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_NVFP4, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_NVFP4, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_NVFP4, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_NVFP4, 256, 1, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_NVFP4, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_NVFP4, 256, 1, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_NVFP4, 256, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_NVFP4, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_NVFP4, 256, 1, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_NVFP4, 256, 1, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_NVFP4, 256, 1, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_NVFP4, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, true, false); + + return ggml_cuda_mmq_config(GGML_TYPE_COUNT, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, 256, false, true); +} diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq-config-blackwell.cuh b/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq-config-blackwell.cuh new file mode 100644 index 0000000000..9fbe32b697 --- /dev/null +++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq-config-blackwell.cuh @@ -0,0 +1,37 @@ +static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config_blackwell(ggml_type type, int J, bool fallback) { + CASE(GGML_TYPE_MXFP4, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, true); + CASE(GGML_TYPE_MXFP4, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, true); + CASE(GGML_TYPE_MXFP4, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, true); + CASE(GGML_TYPE_MXFP4, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, true); + CASE(GGML_TYPE_MXFP4, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, true); + CASE(GGML_TYPE_MXFP4, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, false); + CASE(GGML_TYPE_MXFP4, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, false); + CASE(GGML_TYPE_MXFP4, 256, 1, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, false); + CASE(GGML_TYPE_MXFP4, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, false); + CASE(GGML_TYPE_MXFP4, 256, 1, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, false); + CASE(GGML_TYPE_MXFP4, 256, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, false); + CASE(GGML_TYPE_MXFP4, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, false); + CASE(GGML_TYPE_MXFP4, 256, 1, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, false); + CASE(GGML_TYPE_MXFP4, 256, 1, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, false); + CASE(GGML_TYPE_MXFP4, 256, 1, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, false); + CASE(GGML_TYPE_MXFP4, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, false); + + CASE(GGML_TYPE_NVFP4, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, true); + CASE(GGML_TYPE_NVFP4, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, true); + CASE(GGML_TYPE_NVFP4, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, true); + CASE(GGML_TYPE_NVFP4, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, true); + CASE(GGML_TYPE_NVFP4, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, true); + CASE(GGML_TYPE_NVFP4, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, false); + CASE(GGML_TYPE_NVFP4, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, false); + CASE(GGML_TYPE_NVFP4, 256, 1, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, false); + CASE(GGML_TYPE_NVFP4, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, false); + CASE(GGML_TYPE_NVFP4, 256, 1, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, false); + CASE(GGML_TYPE_NVFP4, 256, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, false); + CASE(GGML_TYPE_NVFP4, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, false); + CASE(GGML_TYPE_NVFP4, 256, 1, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, false); + CASE(GGML_TYPE_NVFP4, 256, 1, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, false); + CASE(GGML_TYPE_NVFP4, 256, 1, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, false); + CASE(GGML_TYPE_NVFP4, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, false); + + return ggml_cuda_mmq_get_config_ampere(type, J, fallback); +} diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq-config-cdna.cuh b/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq-config-cdna.cuh new file mode 100644 index 0000000000..4a8d89f720 --- /dev/null +++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq-config-cdna.cuh @@ -0,0 +1,185 @@ +static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config_cdna(ggml_type type, int J, bool fallback) { + CASE(GGML_TYPE_Q1_0, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q1_0, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q1_0, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q1_0, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q1_0, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q1_0, 512, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q1_0, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + + CASE(GGML_TYPE_Q2_0, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q2_0, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q2_0, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q2_0, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q2_0, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q2_0, 512, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q2_0, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + + CASE(GGML_TYPE_Q4_0, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q4_0, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q4_0, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q4_0, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q4_0, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q4_0, 512, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q4_0, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + + CASE(GGML_TYPE_Q4_1, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q4_1, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q4_1, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q4_1, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q4_1, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q4_1, 512, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q4_1, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + + CASE(GGML_TYPE_Q5_0, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q5_0, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q5_0, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q5_0, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q5_0, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q5_0, 512, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q5_0, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + + CASE(GGML_TYPE_Q5_1, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q5_1, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q5_1, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q5_1, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q5_1, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q5_1, 512, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q5_1, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + + CASE(GGML_TYPE_Q8_0, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q8_0, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q8_0, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q8_0, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q8_0, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q8_0, 512, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q8_0, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + +// --------------------------------------------------------------------------------------------- + + CASE(GGML_TYPE_Q2_K, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q2_K, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q2_K, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q2_K, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q2_K, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q2_K, 512, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q2_K, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, true, false); + + CASE(GGML_TYPE_Q3_K, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q3_K, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q3_K, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q3_K, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q3_K, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q3_K, 512, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q3_K, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false); + + CASE(GGML_TYPE_Q4_K, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q4_K, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q4_K, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q4_K, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q4_K, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q4_K, 512, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q4_K, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + + CASE(GGML_TYPE_Q5_K, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q5_K, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q5_K, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q5_K, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q5_K, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q5_K, 512, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q5_K, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + + CASE(GGML_TYPE_Q6_K, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q6_K, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q6_K, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q6_K, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q6_K, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q6_K, 512, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q6_K, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, true, false); + +// --------------------------------------------------------------------------------------------- + + CASE(GGML_TYPE_IQ1_S, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_IQ1_S, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_IQ1_S, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_IQ1_S, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ1_S, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ1_S, 512, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ1_S, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + + CASE(GGML_TYPE_IQ2_XXS, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_IQ2_XXS, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_IQ2_XXS, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_IQ2_XXS, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ2_XXS, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ2_XXS, 512, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ2_XXS, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + + CASE(GGML_TYPE_IQ2_XS, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_IQ2_XS, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_IQ2_XS, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_IQ2_XS, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ2_XS, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ2_XS, 512, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ2_XS, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false); + + CASE(GGML_TYPE_IQ2_S, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_IQ2_S, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_IQ2_S, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_IQ2_S, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ2_S, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ2_S, 512, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ2_S, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false); + + CASE(GGML_TYPE_IQ3_XXS, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_IQ3_XXS, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_IQ3_XXS, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_IQ3_XXS, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ3_XXS, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ3_XXS, 512, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ3_XXS, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + + CASE(GGML_TYPE_IQ3_S, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_IQ3_S, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_IQ3_S, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_IQ3_S, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ3_S, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ3_S, 512, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ3_S, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + + CASE(GGML_TYPE_IQ4_XS, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_IQ4_XS, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_IQ4_XS, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_IQ4_XS, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ4_XS, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ4_XS, 512, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ4_XS, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + + CASE(GGML_TYPE_IQ4_NL, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_IQ4_NL, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_IQ4_NL, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_IQ4_NL, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ4_NL, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ4_NL, 512, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ4_NL, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + +// --------------------------------------------------------------------------------------------- + + CASE(GGML_TYPE_MXFP4, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_MXFP4, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_MXFP4, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_MXFP4, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_MXFP4, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_MXFP4, 512, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_MXFP4, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + + CASE(GGML_TYPE_NVFP4, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_NVFP4, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_NVFP4, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_NVFP4, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_NVFP4, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_NVFP4, 512, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_NVFP4, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, true, false); + + return ggml_cuda_mmq_config(GGML_TYPE_COUNT, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, 256, false, true); +} diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq-config-gcn.cuh b/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq-config-gcn.cuh new file mode 100644 index 0000000000..24af2ef2b0 --- /dev/null +++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq-config-gcn.cuh @@ -0,0 +1,281 @@ +static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config_gcn(ggml_type type, int J, bool fallback) { + CASE(GGML_TYPE_Q1_0, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q1_0, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q1_0, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q1_0, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q1_0, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q1_0, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q1_0, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q1_0, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q1_0, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q1_0, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q1_0, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q2_0, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q2_0, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q2_0, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q2_0, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q2_0, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_0, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_0, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_0, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_0, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_0, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_0, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q4_0, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_0, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_0, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_0, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_0, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_0, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_0, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_0, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_0, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_0, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_0, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q4_1, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_1, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_1, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_1, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_1, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_1, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_1, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_1, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_1, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_1, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_1, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q5_0, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_0, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_0, 512, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_0, 512, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_0, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_0, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_0, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_0, 512, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_0, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_0, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_0, 512, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q5_1, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_1, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_1, 512, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_1, 512, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_1, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_1, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_1, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_1, 512, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_1, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_1, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_1, 512, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q8_0, 512, 2, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q8_0, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q8_0, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q8_0, 512, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q8_0, 512, 2, 64, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q8_0, 512, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q8_0, 512, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q8_0, 512, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q8_0, 512, 2, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q8_0, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q8_0, 512, 2, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q8_0, 512, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q8_0, 512, 2, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q8_0, 512, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q8_0, 512, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q8_0, 512, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q8_0, 512, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q8_0, 512, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q8_0, 512, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + +// --------------------------------------------------------------------------------------------- + + CASE(GGML_TYPE_Q2_K, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q2_K, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q2_K, 512, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q2_K, 512, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q2_K, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_K, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_K, 512, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_K, 512, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_K, 512, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_K, 512, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_K, 512, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q3_K, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q3_K, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q3_K, 512, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q3_K, 512, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q3_K, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q3_K, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q3_K, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q3_K, 512, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q3_K, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q3_K, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q3_K, 512, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q4_K, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_K, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_K, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_K, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_K, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_K, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_K, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_K, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_K, 256, 2, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_K, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_K, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q5_K, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_K, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_K, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_K, 512, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_K, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_K, 256, 3, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_K, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_K, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_K, 256, 2, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_K, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_K, 512, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q6_K, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q6_K, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q6_K, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q6_K, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q6_K, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q6_K, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q6_K, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q6_K, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q6_K, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q6_K, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q6_K, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); + +// --------------------------------------------------------------------------------------------- + + CASE(GGML_TYPE_IQ1_S, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ1_S, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ1_S, 512, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ1_S, 512, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ1_S, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ1_S, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ1_S, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ1_S, 512, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ1_S, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ1_S, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ1_S, 512, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_IQ2_XXS, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_XXS, 512, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_XXS, 512, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_XXS, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XXS, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XXS, 512, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XXS, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XXS, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XXS, 512, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_IQ2_XS, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_XS, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_XS, 512, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_XS, 512, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_XS, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XS, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XS, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XS, 512, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XS, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XS, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XS, 512, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_IQ2_S, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_S, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_S, 512, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_S, 512, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_S, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_S, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_S, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_S, 512, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_S, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_S, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_S, 512, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_IQ3_XXS, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ3_XXS, 512, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ3_XXS, 512, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ3_XXS, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_XXS, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_XXS, 512, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_XXS, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_XXS, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_XXS, 512, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_IQ3_S, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ3_S, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ3_S, 512, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ3_S, 512, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ3_S, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_S, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_S, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_S, 512, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_S, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_S, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_S, 512, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_IQ4_XS, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ4_XS, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ4_XS, 512, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ4_XS, 512, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ4_XS, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_XS, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_XS, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_XS, 512, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_XS, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_XS, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_XS, 512, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_IQ4_NL, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ4_NL, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ4_NL, 512, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ4_NL, 512, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ4_NL, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_NL, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_NL, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_NL, 512, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_NL, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_NL, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_NL, 512, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + +// --------------------------------------------------------------------------------------------- + + CASE(GGML_TYPE_MXFP4, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_MXFP4, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_MXFP4, 512, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_MXFP4, 512, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_MXFP4, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_MXFP4, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_MXFP4, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_MXFP4, 512, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_MXFP4, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_MXFP4, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_MXFP4, 512, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_NVFP4, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_NVFP4, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_NVFP4, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_NVFP4, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_NVFP4, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_NVFP4, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_NVFP4, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_NVFP4, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_NVFP4, 256, 2, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_NVFP4, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_NVFP4, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false); + + return ggml_cuda_mmq_config(GGML_TYPE_COUNT, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, 256, false, true); +} diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq-config-pascal-dp4a.cuh b/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq-config-pascal-dp4a.cuh new file mode 100644 index 0000000000..83eb7c146e --- /dev/null +++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq-config-pascal-dp4a.cuh @@ -0,0 +1,273 @@ +static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config_pascal_dp4a(ggml_type type, int J, bool fallback) { + CASE(GGML_TYPE_Q1_0, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q1_0, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q1_0, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q1_0, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q1_0, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q1_0, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q1_0, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q1_0, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q1_0, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q1_0, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q1_0, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q2_0, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q2_0, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q2_0, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q2_0, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q2_0, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_0, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_0, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_0, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_0, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_0, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_0, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q4_0, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_0, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_0, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_0, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_0, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_0, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_0, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_0, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_0, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_0, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_0, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q4_1, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_1, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_1, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_1, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_1, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_1, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_1, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_1, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_1, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_1, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_1, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q5_0, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_0, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_0, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_0, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_0, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_0, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_0, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_0, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_0, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_0, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_0, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q5_1, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_1, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_1, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_1, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_1, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_1, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_1, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_1, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_1, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_1, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_1, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q8_0, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q8_0, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q8_0, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q8_0, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q8_0, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q8_0, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q8_0, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q8_0, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q8_0, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q8_0, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q8_0, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + +// --------------------------------------------------------------------------------------------- + + CASE(GGML_TYPE_Q2_K, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q2_K, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q2_K, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q2_K, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q2_K, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_K, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_K, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_K, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_K, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_K, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_K, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q3_K, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q3_K, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q3_K, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q3_K, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q3_K, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q3_K, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q3_K, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q3_K, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q3_K, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q3_K, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q3_K, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q4_K, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_K, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_K, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_K, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_K, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_K, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_K, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_K, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_K, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_K, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_K, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q5_K, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_K, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_K, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_K, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_K, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_K, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_K, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_K, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_K, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_K, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_K, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q6_K, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q6_K, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q6_K, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q6_K, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q6_K, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q6_K, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q6_K, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q6_K, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q6_K, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q6_K, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q6_K, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); + +// --------------------------------------------------------------------------------------------- + + CASE(GGML_TYPE_IQ1_S, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ1_S, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ1_S, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ1_S, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ1_S, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ1_S, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ1_S, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ1_S, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ1_S, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ1_S, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ1_S, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_IQ2_XXS, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_XXS, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_XXS, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_XXS, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_XXS, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XXS, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XXS, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XXS, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XXS, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XXS, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XXS, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_IQ2_XS, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_XS, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_XS, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_XS, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_XS, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XS, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XS, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XS, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XS, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XS, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XS, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_IQ2_S, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_S, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_S, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_S, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_S, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_S, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_S, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_S, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_S, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_S, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_S, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_IQ3_XXS, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ3_XXS, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ3_XXS, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ3_XXS, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ3_XXS, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_XXS, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_XXS, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_XXS, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_XXS, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_XXS, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_XXS, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_IQ3_S, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ3_S, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ3_S, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ3_S, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ3_S, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_S, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_S, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_S, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_S, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_S, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_S, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_IQ4_XS, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ4_XS, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ4_XS, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ4_XS, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ4_XS, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_XS, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_XS, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_XS, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_XS, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_XS, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_XS, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_IQ4_NL, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ4_NL, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ4_NL, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ4_NL, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ4_NL, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_NL, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_NL, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_NL, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_NL, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_NL, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_NL, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + +// --------------------------------------------------------------------------------------------- + + CASE(GGML_TYPE_MXFP4, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_MXFP4, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_MXFP4, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_MXFP4, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_MXFP4, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_MXFP4, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_MXFP4, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_MXFP4, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_MXFP4, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_MXFP4, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_MXFP4, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_NVFP4, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_NVFP4, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_NVFP4, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_NVFP4, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_NVFP4, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_NVFP4, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_NVFP4, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_NVFP4, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_NVFP4, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_NVFP4, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_NVFP4, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false); + + return ggml_cuda_mmq_config(GGML_TYPE_COUNT, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, 256, false, true); +} diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq-config-pascal-older.cuh b/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq-config-pascal-older.cuh new file mode 100644 index 0000000000..2a8dc9e1a9 --- /dev/null +++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq-config-pascal-older.cuh @@ -0,0 +1,273 @@ +static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config_pascal_older(ggml_type type, int J, bool fallback) { + CASE(GGML_TYPE_Q1_0, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q1_0, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q1_0, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q1_0, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q1_0, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q1_0, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q1_0, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q1_0, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q1_0, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q1_0, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q1_0, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q2_0, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q2_0, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q2_0, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q2_0, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q2_0, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_0, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_0, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_0, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_0, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_0, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_0, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q4_0, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_0, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_0, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_0, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_0, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_0, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_0, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_0, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_0, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_0, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_0, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q4_1, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_1, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_1, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_1, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_1, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_1, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_1, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_1, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_1, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_1, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_1, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q5_0, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_0, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_0, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_0, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_0, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_0, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_0, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_0, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_0, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_0, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_0, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q5_1, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_1, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_1, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_1, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_1, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_1, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_1, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_1, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_1, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_1, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_1, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q8_0, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q8_0, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q8_0, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q8_0, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q8_0, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q8_0, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q8_0, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q8_0, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q8_0, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q8_0, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q8_0, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + +// --------------------------------------------------------------------------------------------- + + CASE(GGML_TYPE_Q2_K, 256, 1, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q2_K, 256, 1, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q2_K, 256, 1, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q2_K, 256, 1, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q2_K, 256, 1, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_K, 256, 1, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_K, 256, 1, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_K, 256, 1, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_K, 256, 1, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_K, 256, 1, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_K, 256, 1, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q3_K, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q3_K, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q3_K, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q3_K, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q3_K, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q3_K, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q3_K, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q3_K, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q3_K, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q3_K, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q3_K, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q4_K, 256, 1, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_K, 256, 1, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_K, 256, 1, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_K, 256, 1, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_K, 256, 1, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_K, 256, 1, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_K, 256, 1, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_K, 256, 1, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_K, 256, 1, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_K, 256, 1, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_K, 256, 1, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q5_K, 256, 1, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_K, 256, 1, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_K, 256, 1, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_K, 256, 1, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_K, 256, 1, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_K, 256, 1, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_K, 256, 1, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_K, 256, 1, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_K, 256, 1, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_K, 256, 1, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_K, 256, 1, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q6_K, 256, 1, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q6_K, 256, 1, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q6_K, 256, 1, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q6_K, 256, 1, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q6_K, 256, 1, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q6_K, 256, 1, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q6_K, 256, 1, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q6_K, 256, 1, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q6_K, 256, 1, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q6_K, 256, 1, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q6_K, 256, 1, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); + +// --------------------------------------------------------------------------------------------- + + CASE(GGML_TYPE_IQ1_S, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ1_S, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ1_S, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ1_S, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ1_S, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ1_S, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ1_S, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ1_S, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ1_S, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ1_S, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ1_S, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_IQ2_XXS, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_XXS, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_XXS, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_XXS, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_XXS, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XXS, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XXS, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XXS, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XXS, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XXS, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XXS, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_IQ2_XS, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_XS, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_XS, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_XS, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_XS, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XS, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XS, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XS, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XS, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XS, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XS, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_IQ2_S, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_S, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_S, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_S, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_S, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_S, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_S, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_S, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_S, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_S, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_S, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_IQ3_XXS, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ3_XXS, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ3_XXS, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ3_XXS, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ3_XXS, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_XXS, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_XXS, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_XXS, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_XXS, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_XXS, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_XXS, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_IQ3_S, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ3_S, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ3_S, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ3_S, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ3_S, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_S, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_S, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_S, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_S, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_S, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_S, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_IQ4_XS, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ4_XS, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ4_XS, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ4_XS, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ4_XS, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_XS, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_XS, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_XS, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_XS, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_XS, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_XS, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_IQ4_NL, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ4_NL, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ4_NL, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ4_NL, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ4_NL, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_NL, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_NL, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_NL, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_NL, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_NL, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_NL, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + +// --------------------------------------------------------------------------------------------- + + CASE(GGML_TYPE_MXFP4, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_MXFP4, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_MXFP4, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_MXFP4, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_MXFP4, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_MXFP4, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_MXFP4, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_MXFP4, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_MXFP4, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_MXFP4, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_MXFP4, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_NVFP4, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_NVFP4, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_NVFP4, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_NVFP4, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_NVFP4, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_NVFP4, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_NVFP4, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_NVFP4, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_NVFP4, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_NVFP4, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_NVFP4, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false); + + return ggml_cuda_mmq_config(GGML_TYPE_COUNT, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, 256, false, true); +} diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq-config-rdna2.cuh b/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq-config-rdna2.cuh new file mode 100644 index 0000000000..8324d9e1a8 --- /dev/null +++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq-config-rdna2.cuh @@ -0,0 +1,273 @@ +static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config_rdna2(ggml_type type, int J, bool fallback) { + CASE(GGML_TYPE_Q1_0, 256, 2, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q1_0, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q1_0, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q1_0, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q1_0, 256, 2, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q1_0, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q1_0, 256, 2, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q1_0, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q1_0, 256, 2, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q1_0, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q1_0, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q2_0, 256, 2, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q2_0, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q2_0, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q2_0, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q2_0, 256, 2, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_0, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_0, 256, 2, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_0, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_0, 256, 2, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_0, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_0, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q4_0, 256, 2, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_0, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_0, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_0, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_0, 256, 2, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_0, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_0, 256, 2, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_0, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_0, 256, 2, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_0, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_0, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q4_1, 256, 2, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_1, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_1, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_1, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_1, 256, 2, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_1, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_1, 256, 2, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_1, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_1, 256, 2, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_1, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_1, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q5_0, 256, 2, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_0, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_0, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_0, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_0, 256, 2, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_0, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_0, 256, 2, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_0, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_0, 256, 2, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_0, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_0, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q5_1, 256, 2, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_1, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_1, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_1, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_1, 256, 2, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_1, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_1, 256, 2, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_1, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_1, 256, 2, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_1, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_1, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q8_0, 256, 2, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q8_0, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q8_0, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q8_0, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q8_0, 256, 2, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q8_0, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q8_0, 256, 2, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q8_0, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q8_0, 256, 2, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q8_0, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q8_0, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + +// --------------------------------------------------------------------------------------------- + + CASE(GGML_TYPE_Q2_K, 256, 2, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q2_K, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q2_K, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q2_K, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q2_K, 256, 2, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_K, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_K, 256, 2, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_K, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_K, 256, 2, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_K, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_K, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q3_K, 256, 2, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q3_K, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q3_K, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q3_K, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q3_K, 256, 2, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q3_K, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q3_K, 256, 2, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q3_K, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q3_K, 256, 2, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q3_K, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q3_K, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q4_K, 256, 2, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_K, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_K, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_K, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_K, 256, 2, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_K, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_K, 256, 2, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_K, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_K, 256, 2, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_K, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_K, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q5_K, 256, 2, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_K, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_K, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_K, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_K, 256, 2, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_K, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_K, 256, 2, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_K, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_K, 256, 2, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_K, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_K, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q6_K, 256, 2, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q6_K, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q6_K, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q6_K, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q6_K, 256, 2, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q6_K, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q6_K, 256, 2, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q6_K, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q6_K, 256, 2, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q6_K, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q6_K, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); + +// --------------------------------------------------------------------------------------------- + + CASE(GGML_TYPE_IQ1_S, 256, 2, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ1_S, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ1_S, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ1_S, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ1_S, 256, 2, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ1_S, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ1_S, 256, 2, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ1_S, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ1_S, 256, 2, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ1_S, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ1_S, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_IQ2_XS, 256, 2, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_XS, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_XS, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_XS, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_XS, 256, 2, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XS, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XS, 256, 2, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XS, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XS, 256, 2, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XS, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XS, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_IQ2_S, 256, 2, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_S, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_S, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_S, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_S, 256, 2, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_S, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_S, 256, 2, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_S, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_S, 256, 2, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_S, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_S, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_IQ3_S, 256, 2, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ3_S, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ3_S, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ3_S, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ3_S, 256, 2, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_S, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_S, 256, 2, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_S, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_S, 256, 2, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_S, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_S, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_IQ4_XS, 256, 2, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ4_XS, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ4_XS, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ4_XS, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ4_XS, 256, 2, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_XS, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_XS, 256, 2, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_XS, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_XS, 256, 2, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_XS, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_XS, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_IQ4_NL, 256, 2, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ4_NL, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ4_NL, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ4_NL, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ4_NL, 256, 2, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_NL, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_NL, 256, 2, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_NL, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_NL, 256, 2, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_NL, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_NL, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + +// --------------------------------------------------------------------------------------------- + + CASE(GGML_TYPE_MXFP4, 256, 2, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_MXFP4, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_MXFP4, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_MXFP4, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_MXFP4, 256, 2, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_MXFP4, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_MXFP4, 256, 2, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_MXFP4, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_MXFP4, 256, 2, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_MXFP4, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_MXFP4, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_NVFP4, 256, 2, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_NVFP4, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_NVFP4, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_NVFP4, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_NVFP4, 256, 2, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_NVFP4, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_NVFP4, 256, 2, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_NVFP4, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_NVFP4, 256, 2, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_NVFP4, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_NVFP4, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false); + + return ggml_cuda_mmq_config(GGML_TYPE_COUNT, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, 256, false, true); +} diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq-config-rdna3-5.cuh b/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq-config-rdna3-5.cuh new file mode 100644 index 0000000000..ead3825d21 --- /dev/null +++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq-config-rdna3-5.cuh @@ -0,0 +1,290 @@ +static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config_rdna3_5(ggml_type type, int J, bool fallback) { + CASE(GGML_TYPE_Q1_0, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q1_0, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q1_0, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q1_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q1_0, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q1_0, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q1_0, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q1_0, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q1_0, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q1_0, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q1_0, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q1_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_0, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q2_0, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q2_0, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q2_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q2_0, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_0, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_0, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_0, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_0, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_0, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_0, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q4_0, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_0, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_0, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_0, 128, 2, 64, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_0, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_0, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_0, 128, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_0, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_0, 128, 2, 64, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_0, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_0, 128, 2, 64, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_0, 128, 2, 64, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q4_1, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_1, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_1, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_1, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_1, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_1, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_1, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_1, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_1, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_1, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_1, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_1, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q5_0, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_0, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_0, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_0, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_0, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_0, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_0, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_0, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_0, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_0, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q5_1, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_1, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_1, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_1, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_1, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_1, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_1, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_1, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_1, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_1, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_1, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_1, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q8_0, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q8_0, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q8_0, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q8_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q8_0, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q8_0, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q8_0, 128, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q8_0, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q8_0, 128, 2, 64, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q8_0, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q8_0, 128, 2, 64, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q8_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + +// --------------------------------------------------------------------------------------------- + + CASE(GGML_TYPE_Q2_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q2_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q2_K, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q2_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_K, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_K, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_K, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q3_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q3_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q3_K, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q3_K, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q3_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q3_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q3_K, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q3_K, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q3_K, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q3_K, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q3_K, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q3_K, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q4_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_K, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_K, 128, 2, 64, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_K, 128, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_K, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_K, 128, 2, 64, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_K, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_K, 128, 2, 64, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_K, 128, 2, 64, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q5_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_K, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_K, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_K, 128, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_K, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_K, 128, 2, 64, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_K, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_K, 128, 2, 64, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_K, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q6_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q6_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q6_K, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q6_K, 128, 2, 64, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q6_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q6_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q6_K, 128, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q6_K, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q6_K, 128, 2, 64, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q6_K, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q6_K, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q6_K, 128, 2, 64, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); + +// --------------------------------------------------------------------------------------------- + + CASE(GGML_TYPE_IQ1_S, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ1_S, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ1_S, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ1_S, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ1_S, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ1_S, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ1_S, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ1_S, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ1_S, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ1_S, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ1_S, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ1_S, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_IQ2_XXS, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_XXS, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_XXS, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XXS, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_IQ2_XS, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_XS, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_XS, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_XS, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_XS, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XS, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XS, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XS, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XS, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XS, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XS, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XS, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_IQ2_S, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_S, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_S, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_S, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_S, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_S, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_S, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_S, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_S, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_S, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_S, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_S, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_IQ3_XXS, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ3_XXS, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ3_XXS, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_XXS, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_IQ3_S, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ3_S, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ3_S, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ3_S, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ3_S, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_S, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_S, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_S, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_S, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_S, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_S, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_S, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_IQ4_XS, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ4_XS, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ4_XS, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ4_XS, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ4_XS, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_XS, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_XS, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_XS, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_XS, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_XS, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_XS, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_XS, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_IQ4_NL, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ4_NL, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ4_NL, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ4_NL, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ4_NL, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_NL, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_NL, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_NL, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_NL, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_NL, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_NL, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_NL, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + +// --------------------------------------------------------------------------------------------- + + CASE(GGML_TYPE_MXFP4, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_MXFP4, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_MXFP4, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_MXFP4, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_MXFP4, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_MXFP4, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_MXFP4, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_MXFP4, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_MXFP4, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_MXFP4, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_MXFP4, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_MXFP4, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_NVFP4, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_NVFP4, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_NVFP4, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_NVFP4, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_NVFP4, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_NVFP4, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_NVFP4, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_NVFP4, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_NVFP4, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_NVFP4, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_NVFP4, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_NVFP4, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false); + + return ggml_cuda_mmq_config(GGML_TYPE_COUNT, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, 256, false, true); +} diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq-config-rdna3.cuh b/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq-config-rdna3.cuh new file mode 100644 index 0000000000..3a3ef7bd9c --- /dev/null +++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq-config-rdna3.cuh @@ -0,0 +1,274 @@ +static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config_rdna3(ggml_type type, int J, bool fallback) { + CASE(GGML_TYPE_Q1_0, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q1_0, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q1_0, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q1_0, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q1_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q1_0, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q1_0, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q1_0, 128, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q1_0, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q1_0, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q1_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q2_0, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q2_0, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q2_0, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q2_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q2_0, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_0, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_0, 128, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_0, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_0, 128, 2, 64, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_0, 128, 2, 64, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_0, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q4_0, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_0, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_0, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_0, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_0, 128, 1, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_0, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_0, 128, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_0, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_0, 128, 4, 64, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_0, 128, 4, 64, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q4_1, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_1, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_1, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_1, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_1, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_1, 128, 1, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_1, 128, 1, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_1, 128, 1, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_1, 128, 1, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_1, 128, 1, 64, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_1, 128, 1, 64, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q5_0, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_0, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_0, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_0, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_0, 128, 1, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_0, 128, 4, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_0, 128, 1, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_0, 128, 1, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_0, 128, 1, 64, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_0, 128, 4, 64, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q5_1, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_1, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_1, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_1, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_1, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_1, 128, 4, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_1, 128, 1, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_1, 128, 1, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_1, 128, 1, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_1, 128, 1, 64, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_1, 128, 4, 64, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q8_0, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q8_0, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q8_0, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q8_0, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q8_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q8_0, 128, 1, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q8_0, 128, 1, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q8_0, 128, 1, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q8_0, 128, 1, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q8_0, 128, 1, 64, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q8_0, 128, 1, 64, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + +// --------------------------------------------------------------------------------------------- + + CASE(GGML_TYPE_Q2_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q2_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q2_K, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q2_K, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q2_K, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q2_K, 128, 1, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_K, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_K, 128, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_K, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_K, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_K, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q3_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q3_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q3_K, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q3_K, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q3_K, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q3_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q3_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q3_K, 128, 1, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q3_K, 128, 1, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q3_K, 128, 1, 64, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q3_K, 128, 1, 64, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q4_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_K, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_K, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_K, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_K, 128, 1, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_K, 128, 1, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_K, 128, 1, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_K, 128, 1, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_K, 128, 1, 64, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_K, 128, 1, 64, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q5_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_K, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_K, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_K, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_K, 128, 1, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_K, 128, 1, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_K, 128, 1, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_K, 128, 4, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_K, 128, 4, 64, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_K, 128, 4, 64, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q6_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q6_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q6_K, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q6_K, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q6_K, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q6_K, 128, 4, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q6_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q6_K, 128, 1, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q6_K, 128, 1, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q6_K, 128, 1, 64, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q6_K, 128, 1, 64, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); + +// --------------------------------------------------------------------------------------------- + + CASE(GGML_TYPE_IQ1_S, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ1_S, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ1_S, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ1_S, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ1_S, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ1_S, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ1_S, 128, 1, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ1_S, 128, 4, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ1_S, 128, 4, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ1_S, 128, 1, 64, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ1_S, 128, 1, 64, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_IQ2_XXS, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_XXS, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_XXS, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_XXS, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XXS, 128, 1, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XXS, 128, 4, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XXS, 128, 4, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XXS, 128, 2, 64, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XXS, 128, 4, 64, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_IQ2_XS, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_XS, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_XS, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_XS, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_XS, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_XS, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XS, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XS, 128, 4, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XS, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XS, 128, 4, 64, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XS, 128, 4, 64, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_IQ2_S, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_S, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_S, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_S, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_S, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_S, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_S, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_S, 128, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_S, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_S, 128, 4, 64, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_S, 128, 4, 64, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_IQ3_XXS, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ3_XXS, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ3_XXS, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ3_XXS, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_XXS, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_XXS, 128, 4, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_XXS, 128, 4, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_XXS, 128, 2, 64, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_XXS, 128, 1, 64, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_IQ3_S, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ3_S, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ3_S, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ3_S, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ3_S, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ3_S, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_S, 128, 4, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_S, 128, 4, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_S, 128, 1, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_S, 128, 2, 64, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_S, 128, 2, 64, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_IQ4_XS, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ4_XS, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ4_XS, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ4_XS, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ4_XS, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ4_XS, 128, 4, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_XS, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_XS, 128, 1, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_XS, 128, 1, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_XS, 128, 2, 64, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_XS, 128, 2, 64, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_IQ4_NL, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ4_NL, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ4_NL, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ4_NL, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ4_NL, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ4_NL, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_NL, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_NL, 128, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_NL, 128, 1, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_NL, 128, 2, 64, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_NL, 128, 2, 64, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + +// --------------------------------------------------------------------------------------------- + + CASE(GGML_TYPE_MXFP4, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_MXFP4, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_MXFP4, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_MXFP4, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_MXFP4, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_MXFP4, 128, 1, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_MXFP4, 128, 1, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_MXFP4, 128, 1, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_MXFP4, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_MXFP4, 128, 2, 64, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_MXFP4, 128, 2, 64, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_NVFP4, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_NVFP4, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_NVFP4, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_NVFP4, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_NVFP4, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_NVFP4, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_NVFP4, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_NVFP4, 128, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_NVFP4, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_NVFP4, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_NVFP4, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false); + + return ggml_cuda_mmq_config(GGML_TYPE_COUNT, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, 256, false, true); +} diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq-config-rdna4.cuh b/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq-config-rdna4.cuh new file mode 100644 index 0000000000..9293d9d558 --- /dev/null +++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq-config-rdna4.cuh @@ -0,0 +1,290 @@ +static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config_rdna4(ggml_type type, int J, bool fallback) { + CASE(GGML_TYPE_Q1_0, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q1_0, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q1_0, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q1_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q1_0, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q1_0, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q1_0, 128, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q1_0, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q1_0, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q1_0, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q1_0, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q1_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_0, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q2_0, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q2_0, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q2_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q2_0, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_0, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_0, 128, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_0, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_0, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_0, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_0, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q4_0, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_0, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_0, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_0, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_0, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_0, 128, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_0, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_0, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_0, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_0, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q4_1, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_1, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_1, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_1, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_1, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_1, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_1, 128, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_1, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_1, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_1, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_1, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_1, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q5_0, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_0, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_0, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_0, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_0, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_0, 128, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_0, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_0, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_0, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_0, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q5_1, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_1, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_1, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_1, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_1, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_1, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_1, 128, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_1, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_1, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_1, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_1, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_1, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q8_0, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q8_0, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q8_0, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q8_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q8_0, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q8_0, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q8_0, 128, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q8_0, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q8_0, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q8_0, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q8_0, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q8_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + +// --------------------------------------------------------------------------------------------- + + CASE(GGML_TYPE_Q2_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q2_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q2_K, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q2_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_K, 128, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_K, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_K, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q3_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q3_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q3_K, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q3_K, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q3_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q3_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q3_K, 128, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q3_K, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q3_K, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q3_K, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q3_K, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q3_K, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q4_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_K, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_K, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_K, 128, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_K, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_K, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_K, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_K, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_K, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q5_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_K, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_K, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_K, 128, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_K, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_K, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_K, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_K, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_K, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q6_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q6_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q6_K, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q6_K, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q6_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q6_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q6_K, 128, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q6_K, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q6_K, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q6_K, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q6_K, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q6_K, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); + +// --------------------------------------------------------------------------------------------- + + CASE(GGML_TYPE_IQ1_S, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ1_S, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ1_S, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ1_S, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ1_S, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ1_S, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ1_S, 128, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ1_S, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ1_S, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ1_S, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ1_S, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ1_S, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_IQ2_XXS, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_XXS, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_XXS, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_XXS, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XXS, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XXS, 128, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XXS, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_IQ2_XS, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_XS, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_XS, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_XS, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_XS, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XS, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XS, 128, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XS, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XS, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XS, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XS, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XS, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_IQ2_S, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_S, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_S, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_S, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_S, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_S, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_S, 128, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_S, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_S, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_S, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_S, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_S, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_IQ3_XXS, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ3_XXS, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ3_XXS, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ3_XXS, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_XXS, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_XXS, 128, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_XXS, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_IQ3_S, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ3_S, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ3_S, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ3_S, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ3_S, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_S, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_S, 128, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_S, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_S, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_S, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_S, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_S, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_IQ4_XS, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ4_XS, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ4_XS, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ4_XS, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ4_XS, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_XS, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_XS, 128, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_XS, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_XS, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_XS, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_XS, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_XS, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_IQ4_NL, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ4_NL, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ4_NL, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ4_NL, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ4_NL, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_NL, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_NL, 128, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_NL, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_NL, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_NL, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_NL, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_NL, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + +// --------------------------------------------------------------------------------------------- + + CASE(GGML_TYPE_MXFP4, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_MXFP4, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_MXFP4, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_MXFP4, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_MXFP4, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_MXFP4, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_MXFP4, 128, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_MXFP4, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_MXFP4, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_MXFP4, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_MXFP4, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_MXFP4, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_NVFP4, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_NVFP4, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_NVFP4, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_NVFP4, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_NVFP4, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_NVFP4, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_NVFP4, 128, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_NVFP4, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_NVFP4, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_NVFP4, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_NVFP4, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_NVFP4, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false); + + return ggml_cuda_mmq_config(GGML_TYPE_COUNT, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, 256, false, true); +} diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq-load-tiles.cuh b/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq-load-tiles.cuh new file mode 100644 index 0000000000..017b93b546 --- /dev/null +++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq-load-tiles.cuh @@ -0,0 +1,1768 @@ +#pragma once + +#include "vecdotq.cuh" + +#include "mmq.cuh" + +template static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q1_0( + const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) { + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size; + constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8); + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8); + int * x_qs = (int *) x_tile; + float * x_df = (float *) (x_qs + 2*MMQ_TILE_NE_K); +#else + constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q8_0, I); + int * x_qs = (int *) x_tile; + float * x_df = (float *) (x_qs + txs.qs); +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + + constexpr int blocks_per_iter = MMQ_ITER_K / QK1_0; + constexpr int threads_per_row = blocks_per_iter * QI1_0; + constexpr int nrows = warp_size / threads_per_row; + constexpr int scale_entries_per_block = QK1_0 / QK8_1; + constexpr int scale_entries_per_row = blocks_per_iter * scale_entries_per_block; + + const int txi = threadIdx.x % threads_per_row; + const int kbx = txi / QI1_0; + const int kqsx = txi % QI1_0; + +#pragma unroll + for (int i0 = 0; i0 < I; i0 += nrows*nwarps) { + int i = i0 + threadIdx.y*nrows + threadIdx.x/threads_per_row; + + if (fallback) { + i = min(i, i_max); + } + + const block_q1_0 * bxi = (const block_q1_0 *) x + kbx0 + i*stride + kbx; + const int16_t * qxi = (const int16_t *) bxi->qs + kqsx * 2; + + const int dst_offset = kbx*(scale_entries_per_block*QI8_0) + kqsx*QI8_0; +#pragma unroll + for (int j = 0; j < 2; ++j) { + const int q = qxi[j]; + + // unpack crumbs into nibble indices + const int n0 = __byte_perm(0x11100100, 0x11100100, q >> 0); // [0, 1, 4, 5] [ 8, 9, 12, 13] + const int n1 = __byte_perm(0x11100100, 0x11100100, q >> 2); // [2, 3, 6, 7] [10, 11, 14, 15] + // unpack nibbles into byte values + const int s0 = __byte_perm(0x01FF, 0x01FF, n0 >> 0); + const int s1 = __byte_perm(0x01FF, 0x01FF, n1 >> 0); + const int s2 = __byte_perm(0x01FF, 0x01FF, n0 >> 16); + const int s3 = __byte_perm(0x01FF, 0x01FF, n1 >> 16); + // unshuffle values + const int v0 = __byte_perm(s0, s1, 0x5410); + const int v1 = __byte_perm(s0, s1, 0x7632); + const int v2 = __byte_perm(s2, s3, 0x5410); + const int v3 = __byte_perm(s2, s3, 0x7632); + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + x_qs[i*sram_stride + dst_offset + j*4+0] = v0; + x_qs[i*sram_stride + dst_offset + j*4+1] = v1; + x_qs[i*sram_stride + dst_offset + j*4+2] = v2; + x_qs[i*sram_stride + dst_offset + j*4+3] = v3; +#else + x_qs[i*(2*MMQ_TILE_NE_K + 1) + dst_offset + j*4+0] = v0; + x_qs[i*(2*MMQ_TILE_NE_K + 1) + dst_offset + j*4+1] = v1; + x_qs[i*(2*MMQ_TILE_NE_K + 1) + dst_offset + j*4+2] = v2; + x_qs[i*(2*MMQ_TILE_NE_K + 1) + dst_offset + j*4+3] = v3; +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + } + } + + const int ksx = threadIdx.x % scale_entries_per_row; + const int scale_block = ksx / scale_entries_per_block; + +#pragma unroll + for (int i0 = 0; i0 < I; i0 += nwarps) { + int i = i0 + threadIdx.y; + + if (fallback) { + i = min(i, i_max); + } + + const block_q1_0 * bxi = (const block_q1_0 *) x + kbx0 + i*stride + scale_block; + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + x_df[i*sram_stride + ksx] = bxi->d; +#else + x_df[i*(2*MMQ_TILE_NE_K/QI8_0) + i/(QI8_0/2) + ksx] = bxi->d; +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + } +} + +template static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q2_0( + const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) { + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size; + constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8); + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8); + int * x_qs = (int *) x_tile; + float * x_df = (float *) (x_qs + 2*MMQ_TILE_NE_K); +#else + constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q8_0, I); + int * x_qs = (int *) x_tile; + float * x_df = (float *) (x_qs + txs.qs); +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + + constexpr int blocks_per_iter = MMQ_ITER_K / QK2_0; + constexpr int threads_per_row = blocks_per_iter * QI2_0; + constexpr int nrows = warp_size / threads_per_row; + constexpr int scale_entries_per_block = QK2_0 / QK8_1; + constexpr int scale_entries_per_row = blocks_per_iter * scale_entries_per_block; + + const int txi = threadIdx.x % threads_per_row; + const int kbx = txi / QI2_0; + const int kqsx = txi % QI2_0; + +#pragma unroll + for (int i0 = 0; i0 < I; i0 += nrows*nwarps) { + int i = i0 + threadIdx.y*nrows + threadIdx.x/threads_per_row; + + if (fallback) { + i = min(i, i_max); + } + + const block_q2_0 * bxi = (const block_q2_0 *) x + kbx0 + i*stride + kbx; + const int16_t * qxi = (const int16_t *) bxi->qs + kqsx * 4; + + const int dst_offset = kbx*(scale_entries_per_block*QI8_0) + kqsx*QI8_0; + +#pragma unroll + for (int j = 0; j < 4; ++j) { + const int q = qxi[j]; + +#if defined(GGML_USE_HIP) + const uint32_t qx_indices = (q & 0x03) | ((q & 0x0C) << 6) | ((q & 0x30) << 12) | ((q & 0xC0) << 18); + const uint32_t qy_bits = q >> 8; + const uint32_t qy_indices = (qy_bits & 0x03) | ((qy_bits & 0x0C) << 6) | ((qy_bits & 0x30) << 12) | ((qy_bits & 0xC0) << 18); + const int qx = __builtin_amdgcn_perm(0x020100FF, 0x020100FF, qx_indices); + const int qy = __builtin_amdgcn_perm(0x020100FF, 0x020100FF, qy_indices); +#else + // unpack even and odd crumbs into byte values + const int qe = __byte_perm(0x020100FF, 0x020100FF, q >> 0); + const int qo = __byte_perm(0x020100FF, 0x020100FF, q >> 2); + // unshuffle values + const int qx = __byte_perm(qe, qo, 0x5140); + const int qy = __byte_perm(qe, qo, 0x7362); +#endif // defined(GGML_USE_HIP) + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + x_qs[i*sram_stride + dst_offset + j*2+0] = qx; + x_qs[i*sram_stride + dst_offset + j*2+1] = qy; +#else + x_qs[i*(2*MMQ_TILE_NE_K + 1) + dst_offset + j*2+0] = qx; + x_qs[i*(2*MMQ_TILE_NE_K + 1) + dst_offset + j*2+1] = qy; +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + } + } + + const int ksx = threadIdx.x % scale_entries_per_row; + const int scale_block = ksx / scale_entries_per_block; + +#pragma unroll + for (int i0 = 0; i0 < I; i0 += nwarps) { + int i = i0 + threadIdx.y; + + if (fallback) { + i = min(i, i_max); + } + + const block_q2_0 * bxi = (const block_q2_0 *) x + kbx0 + i*stride + scale_block; + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + x_df[i*sram_stride + ksx] = bxi->d; +#else + x_df[i*(2*MMQ_TILE_NE_K/QI8_0) + i/(QI8_0/2) + ksx] = bxi->d; +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + } +} + +template static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q4_0( + const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) { + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size; + constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8); + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8); + int * x_qs = (int *) x_tile; + float * x_df = (float *) (x_qs + 2*MMQ_TILE_NE_K); +#else + constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q4_0, I); + int * x_qs = (int *) x_tile; + float * x_df = (float *) (x_qs + txs.qs); +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + + constexpr int threads_per_row = MMQ_ITER_K / (4 * QR4_0); + constexpr int nrows = warp_size / threads_per_row; + const int txi = warp_size > threads_per_row ? threadIdx.x % threads_per_row : threadIdx.x; + const int kbx = txi / QI4_0; + const int kqsx = txi % QI4_0; + +#pragma unroll + for (int i0 = 0; i0 < I; i0 += nrows*nwarps) { + int i = i0 + (nrows == 1 ? threadIdx.y : threadIdx.y*nrows + threadIdx.x/threads_per_row); + + if (fallback) { + i = min(i, i_max); + } + + const block_q4_0 * bxi = (const block_q4_0 *) x + kbx0 + i*stride + kbx; + const int qs0 = get_int_b2(bxi->qs, kqsx); + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + x_qs[i*sram_stride + kbx*(2*QI4_0) + kqsx + 0] = __vsub4((qs0 >> 0) & 0x0F0F0F0F, 0x08080808); + x_qs[i*sram_stride + kbx*(2*QI4_0) + kqsx + QI4_0] = __vsub4((qs0 >> 4) & 0x0F0F0F0F, 0x08080808); +#else + x_qs[i*(MMQ_TILE_NE_K + 1) + txi] = qs0; +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) + } + + constexpr int blocks_per_tile_x_row = MMQ_TILE_NE_K / QI4_0; + constexpr int rows_per_warp = warp_size / blocks_per_tile_x_row; + const int kbxd = threadIdx.x % blocks_per_tile_x_row; + +#pragma unroll + for (int i0 = 0; i0 < I; i0 += nwarps * rows_per_warp) { + int i = i0 + threadIdx.y * rows_per_warp + threadIdx.x / blocks_per_tile_x_row; + + if (fallback) { + i = min(i, i_max); + } + + const block_q4_0 * bxi = (const block_q4_0 *) x + kbx0 + i*stride + kbxd; + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + x_df[i*sram_stride + kbxd] = bxi->d; +#else + x_df[i*(MMQ_TILE_NE_K/QI4_0) + i/QI4_0 + kbxd] = bxi->d; +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + } +} + +template static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q4_1( + const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) { + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size; + constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8); + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8); + int * x_qs = (int *) x_tile; + half2 * x_dm = (half2 *) (x_qs + 2*MMQ_TILE_NE_K); +#else + constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q4_1, I); + int * x_qs = (int *) x_tile; + half2 * x_dm = (half2 *) (x_qs + txs.qs); +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + + constexpr int threads_per_row = MMQ_ITER_K / (4 * QR4_1); + constexpr int nrows = warp_size / threads_per_row; + const int txi = warp_size > threads_per_row ? threadIdx.x % threads_per_row : threadIdx.x; + const int kbx = txi / QI4_1; + const int kqsx = txi % QI4_1; + +#pragma unroll + for (int i0 = 0; i0 < I; i0 += nrows*nwarps) { + int i = i0 + (nrows == 1 ? threadIdx.y : threadIdx.y*nrows + threadIdx.x/threads_per_row); + + if (fallback) { + i = min(i, i_max); + } + + const block_q4_1 * bxi = (const block_q4_1 *) x + kbx0 + i*stride + kbx; + const int qs0 = get_int_b4(bxi->qs, kqsx); + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + x_qs[i*sram_stride + kbx*(2*QI4_1) + kqsx + 0] = (qs0 >> 0) & 0x0F0F0F0F; + x_qs[i*sram_stride + kbx*(2*QI4_1) + kqsx + QI4_1] = (qs0 >> 4) & 0x0F0F0F0F; +#else + x_qs[i*(MMQ_TILE_NE_K + 1) + txi] = qs0; +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + } + + constexpr int blocks_per_tile_x_row = MMQ_TILE_NE_K / QI4_1; + constexpr int rows_per_warp = warp_size / blocks_per_tile_x_row; + const int kbxd = threadIdx.x % blocks_per_tile_x_row; + +#pragma unroll + for (int i0 = 0; i0 < I; i0 += nwarps * rows_per_warp) { + int i = i0 + threadIdx.y * rows_per_warp + threadIdx.x / blocks_per_tile_x_row; + + if (fallback) { + i = min(i, i_max); + } + + const block_q4_1 * bxi = (const block_q4_1 *) x + kbx0 + i*stride + kbxd; + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + x_dm[i*sram_stride + kbxd] = bxi->dm; +#else + x_dm[i*(MMQ_TILE_NE_K/QI4_1) + i/QI4_1 + kbxd] = bxi->dm; +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + } +} + +template static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q5_0( + const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) { + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size; + constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8); + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8); + int * x_qs = (int *) x_tile; + float * x_df = (float *) (x_qs + MMQ_TILE_NE_K*2); +#else + constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q5_0, I); + int * x_qs = (int *) x_tile; + float * x_df = (float *) (x_qs + txs.qs); +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + + constexpr int threads_per_row = MMQ_ITER_K / (4 * QR5_0); + constexpr int nrows = warp_size / threads_per_row; + const int txi = warp_size > threads_per_row ? threadIdx.x % threads_per_row : threadIdx.x; + const int kbx = txi / QI5_0; + const int kqsx = txi % QI5_0; + +#pragma unroll + for (int i0 = 0; i0 < I; i0 += nrows*nwarps) { + int i = i0 + (nrows == 1 ? threadIdx.y : threadIdx.y*nrows + threadIdx.x/threads_per_row); + + if (fallback) { + i = min(i, i_max); + } + + const block_q5_0 * bxi = (const block_q5_0 *) x + kbx0 + i*stride + kbx; + + const int ql = get_int_b2(bxi->qs, kqsx); + const int qh = get_int_b2(bxi->qh, 0) >> (4 * kqsx); + + int qs0 = (ql >> 0) & 0x0F0F0F0F; + qs0 |= (qh << 4) & 0x00000010; // 0 -> 4 + qs0 |= (qh << 11) & 0x00001000; // 1 -> 12 + qs0 |= (qh << 18) & 0x00100000; // 2 -> 20 + qs0 |= (qh << 25) & 0x10000000; // 3 -> 28 + qs0 = __vsub4(qs0, 0x10101010); // subtract 16 + + int qs1 = (ql >> 4) & 0x0F0F0F0F; + qs1 |= (qh >> 12) & 0x00000010; // 16 -> 4 + qs1 |= (qh >> 5) & 0x00001000; // 17 -> 12 + qs1 |= (qh << 2) & 0x00100000; // 18 -> 20 + qs1 |= (qh << 9) & 0x10000000; // 19 -> 28 + qs1 = __vsub4(qs1, 0x10101010); // subtract 16 + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + x_qs[i*sram_stride + kbx*(2*QI5_0) + kqsx + 0] = qs0; + x_qs[i*sram_stride + kbx*(2*QI5_0) + kqsx + QI5_0] = qs1; +#else + x_qs[i*(2*MMQ_TILE_NE_K + 1) + kbx*(2*QI5_0) + kqsx + 0] = qs0; + x_qs[i*(2*MMQ_TILE_NE_K + 1) + kbx*(2*QI5_0) + kqsx + QI5_0] = qs1; +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + } + + constexpr int blocks_per_tile_x_row = MMQ_TILE_NE_K / QI5_0; + constexpr int rows_per_warp = warp_size / blocks_per_tile_x_row; + const int kbxd = threadIdx.x % blocks_per_tile_x_row; + +#pragma unroll + for (int i0 = 0; i0 < I; i0 += nwarps * rows_per_warp) { + int i = i0 + threadIdx.y * rows_per_warp + threadIdx.x / blocks_per_tile_x_row; + + if (fallback) { + i = min(i, i_max); + } + + const block_q5_0 * bxi = (const block_q5_0 *) x + kbx0 + i*stride + kbxd; + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + x_df[i*sram_stride + kbxd] = bxi->d; +#else + x_df[i*(MMQ_TILE_NE_K/QI5_0) + i/QI5_0 + kbxd] = bxi->d; +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + } +} + +template static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q5_1( + const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) { + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size; + constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8); + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8); + int * x_qs = (int *) x_tile; + half2 * x_dm = (half2 *) (x_qs + 2*MMQ_TILE_NE_K); +#else + constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q5_1, I); + int * x_qs = (int *) x_tile; + half2 * x_dm = (half2 *) (x_qs + txs.qs); +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + + constexpr int threads_per_row = MMQ_ITER_K / (4 * QR5_1); + constexpr int nrows = warp_size / threads_per_row; + const int txi = warp_size > threads_per_row ? threadIdx.x % threads_per_row : threadIdx.x; + const int kbx = txi / QI5_1; + const int kqsx = txi % QI5_1; + +#pragma unroll + for (int i0 = 0; i0 < I; i0 += nrows*nwarps) { + int i = i0 + (nrows == 1 ? threadIdx.y : threadIdx.y*nrows + threadIdx.x/threads_per_row); + + if (fallback) { + i = min(i, i_max); + } + + const block_q5_1 * bxi = (const block_q5_1 *) x + kbx0 + i*stride + kbx; + + const int ql = get_int_b4(bxi->qs, kqsx); + const int qh = get_int_b4(bxi->qh, 0) >> (4 * kqsx); + + int qs0 = (ql >> 0) & 0x0F0F0F0F; + qs0 |= (qh << 4) & 0x00000010; // 0 -> 4 + qs0 |= (qh << 11) & 0x00001000; // 1 -> 12 + qs0 |= (qh << 18) & 0x00100000; // 2 -> 20 + qs0 |= (qh << 25) & 0x10000000; // 3 -> 28 + + int qs1 = (ql >> 4) & 0x0F0F0F0F; + qs1 |= (qh >> 12) & 0x00000010; // 16 -> 4 + qs1 |= (qh >> 5) & 0x00001000; // 17 -> 12 + qs1 |= (qh << 2) & 0x00100000; // 18 -> 20 + qs1 |= (qh << 9) & 0x10000000; // 19 -> 28 + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + x_qs[i*sram_stride + kbx*(2*QI5_1) + kqsx + 0] = qs0; + x_qs[i*sram_stride + kbx*(2*QI5_1) + kqsx + QI5_1] = qs1; +#else + x_qs[i*(2*MMQ_TILE_NE_K + 1) + kbx*(2*QI5_1) + kqsx + 0] = qs0; + x_qs[i*(2*MMQ_TILE_NE_K + 1) + kbx*(2*QI5_1) + kqsx + QI5_1] = qs1; +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + } + + constexpr int blocks_per_tile_x_row = MMQ_TILE_NE_K / QI5_1; + constexpr int rows_per_warp = warp_size / blocks_per_tile_x_row; + const int kbxd = threadIdx.x % blocks_per_tile_x_row; + +#pragma unroll + for (int i0 = 0; i0 < I; i0 += nwarps * rows_per_warp) { + int i = i0 + threadIdx.y * rows_per_warp + threadIdx.x / blocks_per_tile_x_row; + + if (fallback) { + i = min(i, i_max); + } + + const block_q5_1 * bxi = (const block_q5_1 *) x + kbx0 + i*stride + kbxd; + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + x_dm[i*sram_stride + kbxd] = bxi->dm; +#else + x_dm[i*(MMQ_TILE_NE_K/QI5_1) + i/QI5_1 + kbxd] = bxi->dm; +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + } +} + +template static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q8_0( + const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) { + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size; + constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8); + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8); + int * x_qs = (int *) x_tile; + float * x_df = (float *) (x_tile + 2*MMQ_TILE_NE_K); +#else + constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q8_0, I); + int * x_qs = (int *) x_tile; + float * x_df = (float *) (x_qs + txs.qs); +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + + // MMQ_ITER_K / (4 * QR8_0) == 64 required. but NV has only 32 threads per warp + constexpr int threads_per_row = 32; + constexpr int nrows = warp_size / threads_per_row; + const int txi = warp_size > threads_per_row ? threadIdx.x % threads_per_row : threadIdx.x; + const int kbx = txi / QI8_0; + const int kqsx = txi % QI8_0; + +#pragma unroll + for (int i0 = 0; i0 < I; i0 += nrows*nwarps) { + int i = i0 + (nrows == 1 ? threadIdx.y : threadIdx.y*nrows + threadIdx.x/threads_per_row); + + if (fallback) { + i = min(i, i_max); + } + + const block_q8_0 * bxi = (const block_q8_0 *) x + kbx0 + i*stride + kbx; + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + x_qs[i*sram_stride + 0 + txi] = get_int_b2(bxi[0].qs, kqsx); + x_qs[i*sram_stride + MMQ_TILE_NE_K + txi] = get_int_b2(bxi[MMQ_TILE_NE_K/QI8_0].qs, kqsx); +#else + x_qs[i*(2*MMQ_TILE_NE_K + 1) + 0 + txi] = get_int_b2(bxi[0].qs, kqsx); + x_qs[i*(2*MMQ_TILE_NE_K + 1) + MMQ_TILE_NE_K + txi] = get_int_b2(bxi[MMQ_TILE_NE_K/QI8_0].qs, kqsx); +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + } + + constexpr int blocks_per_tile_x_row = 2*MMQ_TILE_NE_K / QI8_0; + constexpr int rows_per_warp = warp_size / blocks_per_tile_x_row; + const int kbxd = threadIdx.x % blocks_per_tile_x_row; + +#pragma unroll + for (int i0 = 0; i0 < I; i0 += nwarps * rows_per_warp) { + int i = i0 + threadIdx.y * rows_per_warp + threadIdx.x / blocks_per_tile_x_row; + + if (fallback) { + i = min(i, i_max); + } + + const block_q8_0 * bxi = (const block_q8_0 *) x + kbx0 + i*stride + kbxd; + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + x_df[i*sram_stride + kbxd] = bxi->d; +#else + x_df[i*(2*MMQ_TILE_NE_K/QI8_0) + i/(QI8_0/2) + kbxd] = bxi->d; +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + } +} + +// --------------------------------------------------------------------------------------------- + +template static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q2_K( + const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) { + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size; + constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8); + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8); + int * x_qs = (int *) x_tile; + half2 * x_dm = (half2 *) (x_qs + 2*MMQ_TILE_NE_K); +#else + constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q2_K, I); + int * x_qs = (int *) x_tile; + half2 * x_dm = (half2 *) (x_qs + txs.qs); +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + + constexpr int threads_per_row = MMQ_ITER_K / (4 * QR2_K); + constexpr int nrows = ggml_cuda_get_physical_warp_size() / threads_per_row; + const int kqsx = threadIdx.x % threads_per_row; + +#pragma unroll + for (int i0 = 0; i0 < I; i0 += nrows*nwarps) { + int i = i0 + threadIdx.y*nrows + threadIdx.x/threads_per_row; + + if (fallback) { + i = min(i, i_max); + } + + const block_q2_K * bxi = (const block_q2_K *) x + kbx0 + i*stride; + + const int x_ql_0 = get_int_b2(bxi->qs, kqsx); + +#pragma unroll + for (int l = 0; l < QR2_K; ++l) { + const int k = (kqsx/8)*32 + l*8 + kqsx % 8; + + const int x_qs_k = (x_ql_0 >> (2*l)) & 0x03030303; + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + x_qs[i*sram_stride + k] = x_qs_k; +#else + x_qs[i*(2*MMQ_TILE_NE_K + 1) + k] = x_qs_k; +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + } + + const int sc_m = bxi->scales[kqsx]; +#ifdef FAST_FP16_AVAILABLE + const half2 x_dm_ik = __hmul2(bxi->dm, make_half2(sc_m & 0x0F, sc_m >> 4)); +#else + const float2 bxi_dmf = __half22float2(bxi->dm); + const half2 x_dm_ik = make_half2(bxi_dmf.x*(sc_m & 0x0F), bxi_dmf.y*(sc_m >> 4)); +#endif // FAST_FP16_AVAILABLE + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + x_dm[i*sram_stride + kqsx] = x_dm_ik; +#else + x_dm[i*(MMQ_TILE_NE_K + 1) + kqsx] = x_dm_ik; +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + } +} + +template static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q3_K( + const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) { + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size; + constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8); + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8); + int * x_qs = (int *) x_tile; + float * x_df = (float *) (x_qs + MMQ_TILE_NE_K*2); +#else + constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q3_K, I); + int * x_qs = (int *) x_tile; + float * x_df = (float *) (x_qs + txs.qs); + int * x_sc = (int *) (x_df + txs.dm); +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) + + constexpr int threads_per_row = MMQ_ITER_K / (4 * QR3_K); + constexpr int nrows = warp_size / threads_per_row; + const int kqsx = threadIdx.x % threads_per_row; + +#pragma unroll + for (int i0 = 0; i0 < I; i0 += nrows*nwarps) { + int i = i0 + threadIdx.y*nrows + threadIdx.x/threads_per_row; + + if (fallback) { + i = min(i, i_max); + } + + const block_q3_K * bxi = (const block_q3_K *) x + kbx0 + i*stride; + + const int x_ql_0 = get_int_b2(bxi->qs, kqsx); + const int x_qh_0 = get_int_b2(bxi->hmask, kqsx % (QI3_K/2)) >> (4 * (kqsx / (QI3_K/2))); + +#pragma unroll + for (int l = 0; l < QR3_K; ++l) { + const int k = (kqsx/8)*32 + l*8 + kqsx % 8; + + const int x_ql_k = (x_ql_0 >> (2*l)) & 0x03030303; + const int x_qh_k = ((x_qh_0 >> l) << 2) & 0x04040404; + + const int x_qs_k = __vsub4(x_ql_k | x_qh_k, 0x04040404); + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + x_qs[i*sram_stride + k] = x_qs_k; +#else + x_qs[i*(2*MMQ_TILE_NE_K + 1) + k] = x_qs_k; +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + } + } + + constexpr int rows_per_warp = warp_size / 4; +#pragma unroll + for (int i0 = 0; i0 < I; i0 += nwarps*rows_per_warp) { + int i = i0 + threadIdx.y*rows_per_warp + threadIdx.x/4; + + if (fallback) { + i = min(i, i_max); + } + + const block_q3_K * bxi = (const block_q3_K *) x + kbx0 + i*stride; + + const int ksc = threadIdx.x % 4; + + const int ksc_low = ksc % (QI3_K/8); + const int shift_low = 4 * (ksc / (QI3_K/8)); + const int sc_low = (get_int_b2(bxi->scales, ksc_low) >> shift_low) & 0x0F0F0F0F; + + const int ksc_high = QI3_K/8; + const int shift_high = 2 * ksc; + const int sc_high = ((get_int_b2(bxi->scales, ksc_high) >> shift_high) << 4) & 0x30303030; + + const int sc = __vsub4(sc_low | sc_high, 0x20202020); + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + const int8_t * sc8 = (const int8_t *) ≻ + const float d = bxi->d; + +#pragma unroll + for (int l = 0; l < int(sizeof(int)); ++l) { + x_df[i*sram_stride + sizeof(int)*ksc + l] = d*sc8[l]; + } +#else + x_sc[i*(MMQ_TILE_NE_K/8) + i/8 + ksc] = sc; +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + } + +#if !(defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)) +#pragma unroll + for (int i0 = 0; i0 < I; i0 += nwarps*warp_size) { + int i = (i0 + threadIdx.y*warp_size + threadIdx.x) % I; + + if (fallback) { + i = min(i, i_max); + } + + const block_q3_K * bxi = (const block_q3_K *) x + kbx0 + i*stride; + + x_df[i] = bxi->d; + } +#endif // !(defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE)) || defined(AMD_WMMA_AVAILABLE) +} + +static __device__ __forceinline__ int unpack_scales_q45_K(const int * scales, const int ksc) { + // scale arrangement after the following two lines: + // - ksc == 0: sc0, sc1, sc2, sc3 + // - ksc == 1: sc4, sc5, sc6, sc7 + // - ksc == 2: m0, m1, m2, m3 + // - ksc == 3: m4, m5, m6, m7 + return ((scales[(ksc%2) + (ksc!=0)] >> (4 * (ksc & (ksc/2)))) & 0x0F0F0F0F) | // lower 4 bits + ((scales[ksc/2] >> (2 * (ksc % 2))) & 0x30303030); // upper 2 bits +} + +template static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q4_K( + const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) { + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size; + constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8); + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8); + int * x_qs = (int *) x_tile; + half2 * x_dm = (half2 *) (x_qs + 2*MMQ_TILE_NE_K); +#else + constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q4_K, I); + int * x_qs = (int *) x_tile; + half2 * x_dm = (half2 *) (x_qs + txs.qs); + int * x_sc = (int *) (x_dm + txs.dm); +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + + constexpr int threads_per_row = MMQ_ITER_K / (4 * QR4_K); + constexpr int nrows = warp_size / threads_per_row; + const int txi = warp_size > threads_per_row ? threadIdx.x % threads_per_row : threadIdx.x; + +#pragma unroll + for (int i0 = 0; i0 < I; i0 += nrows*nwarps) { + int i = i0 + (nrows == 1 ? threadIdx.y : threadIdx.y*nrows + threadIdx.x/threads_per_row); + + if (fallback) { + i = min(i, i_max); + } + + const block_q4_K * bxi = (const block_q4_K *) x + kbx0 + i*stride; + const int qs0 = get_int_b4(bxi->qs, txi); + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + x_qs[i*sram_stride + 16*(txi/8) + txi % 8 + 0] = (qs0 >> 0) & 0x0F0F0F0F; + x_qs[i*sram_stride + 16*(txi/8) + txi % 8 + 8] = (qs0 >> 4) & 0x0F0F0F0F; +#else + x_qs[i*(MMQ_TILE_NE_K + 1) + txi] = qs0; +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + } + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + constexpr int rows_per_warp = warp_size / 2; +#pragma unroll + for (int i0 = 0; i0 < I; i0 += nwarps*rows_per_warp) { +#if defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + // Need if on AMD instead of % because warp_size == 64 + // This causes double work and throughput loss (MI300X) + // H100 loses about 100 t/s with 'if' condition over '%' + int i = i0 + threadIdx.y*rows_per_warp + threadIdx.x/2; + if (i < I) { +#else + int i = (i0 + threadIdx.y*rows_per_warp + threadIdx.x/2) % I; + { +#endif // defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + if (fallback) { + i = min(i, i_max); + } + + const block_q4_K * bxi = (const block_q4_K *) x + kbx0 + i*stride; + + const int * scales = (const int *) bxi->scales; + const int ksc = threadIdx.x % 2; + + const int sc32 = unpack_scales_q45_K(scales, ksc + 0); + const int m32 = unpack_scales_q45_K(scales, ksc + 2); + + const uint8_t * sc8 = (const uint8_t *) &sc32; + const uint8_t * m8 = (const uint8_t *) &m32; + + const half2 dm = bxi->dm * make_half2(1.0f, -1.0f); + + #pragma unroll + for (int l = 0; l < int(sizeof(int)); ++l) { + x_dm[i*sram_stride + sizeof(int)*ksc + l] = dm*make_half2(sc8[l], m8[l]); + } + } + } +#else +#pragma unroll + for (int i0 = 0; i0 < I; i0 += nwarps*warp_size) { + int i = (i0 + threadIdx.y*warp_size + threadIdx.x) % I; + + if (fallback) { + i = min(i, i_max); + } + + const block_q4_K * bxi = (const block_q4_K *) x + kbx0 + i*stride; + + x_dm[i] = bxi->dm; + } + constexpr int rows_per_warp = warp_size / 4; +#pragma unroll + for (int i0 = 0; i0 < I; i0 += nwarps*rows_per_warp) { + int i = (i0 + threadIdx.y*rows_per_warp + threadIdx.x/(MMQ_TILE_NE_K/8)) % I; + + if (fallback) { + i = min(i, i_max); + } + + const block_q4_K * bxi = (const block_q4_K *) x + kbx0 + i*stride + (threadIdx.x % (MMQ_TILE_NE_K/8)) / (QI4_K/8); + + const int * scales = (const int *) bxi->scales; + + const int ksc = threadIdx.x % (MMQ_TILE_NE_K/8); + const int scales8 = unpack_scales_q45_K(scales, ksc); + + x_sc[i*(MMQ_TILE_NE_K/8) + i/8 + ksc] = scales8; + } +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) +} + +template static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q5_K( + const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) { + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size; + constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8); + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8); + int * x_qs = (int *) x_tile; + half2 * x_dm = (half2 *) (x_qs + MMQ_TILE_NE_K*2); +#else + constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q5_K, I); + int * x_qs = (int *) x_tile; + half2 * x_dm = (half2 *) (x_qs + txs.qs); + int * x_sc = (int *) (x_dm + txs.dm); +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) + + constexpr int threads_per_row = MMQ_ITER_K / (4 * QR5_K); + constexpr int nrows = warp_size / threads_per_row; + const int txi = warp_size > threads_per_row ? threadIdx.x % threads_per_row : threadIdx.x; + +#pragma unroll + for (int i0 = 0; i0 < I; i0 += nrows*nwarps) { + int i = i0 + (nrows == 1 ? threadIdx.y : threadIdx.y*nrows + threadIdx.x/threads_per_row); + + if (fallback) { + i = min(i, i_max); + } + + const block_q5_K * bxi = (const block_q5_K *) x + kbx0 + i*stride; + const int ky = QR5_K*txi; + + const int ql = get_int_b4(bxi->qs, txi); + const int ql0 = (ql >> 0) & 0x0F0F0F0F; + const int ql1 = (ql >> 4) & 0x0F0F0F0F; + + const int qh = get_int_b4(bxi->qh, txi % (QI5_K/4)); + const int qh0 = ((qh >> (2 * (txi / (QI5_K/4)) + 0)) << 4) & 0x10101010; + const int qh1 = ((qh >> (2 * (txi / (QI5_K/4)) + 1)) << 4) & 0x10101010; + + const int kq0 = ky - ky % (QI5_K/2) + txi % (QI5_K/4) + 0; + const int kq1 = ky - ky % (QI5_K/2) + txi % (QI5_K/4) + QI5_K/4; + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + x_qs[i*sram_stride + kq0] = ql0 | qh0; + x_qs[i*sram_stride + kq1] = ql1 | qh1; +#else + x_qs[i*(2*MMQ_TILE_NE_K + 1) + kq0] = ql0 | qh0; + x_qs[i*(2*MMQ_TILE_NE_K + 1) + kq1] = ql1 | qh1; +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + } + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + constexpr int rows_per_warp = warp_size / 2; +#pragma unroll + for (int i0 = 0; i0 < I; i0 += nwarps*rows_per_warp) { +#if defined(AMD_MFMA_AVAILABLE) + // Need if on AMD instead of % because warp_size == 64 + // This causes double work and throughput loss (MI300X) + // H100 loses about 100 t/s with 'if' condition over '%' + int i = i0 + threadIdx.y*rows_per_warp + threadIdx.x/2; + if (i < I) { +#else + int i = (i0 + threadIdx.y*rows_per_warp + threadIdx.x/2) % I; + { +#endif // defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + if (fallback) { + i = min(i, i_max); + } + + const block_q5_K * bxi = (const block_q5_K *) x + kbx0 + i*stride; + + const int * scales = (const int *) bxi->scales; + const int ksc = threadIdx.x % 2; + + const int sc32 = unpack_scales_q45_K(scales, ksc + 0); + const int m32 = unpack_scales_q45_K(scales, ksc + 2); + + const uint8_t * sc8 = (const uint8_t *) &sc32; + const uint8_t * m8 = (const uint8_t *) &m32; + + const half2 dm = bxi->dm * make_half2(1.0f, -1.0f); + +#pragma unroll + for (int l = 0; l < int(sizeof(int)); ++l) { + x_dm[i*sram_stride + sizeof(int)*ksc + l] = dm*make_half2(sc8[l], m8[l]); + } + } + } +#else +#pragma unroll + for (int i0 = 0; i0 < I; i0 += nwarps*warp_size) { + int i = (i0 + threadIdx.y*warp_size + threadIdx.x) % I; + + if (fallback) { + i = min(i, i_max); + } + + const block_q5_K * bxi = (const block_q5_K *) x + kbx0 + i*stride; + + x_dm[i] = bxi->dm; + } + + constexpr int rows_per_warp = warp_size / 4; +#pragma unroll + for (int i0 = 0; i0 < I; i0 += nwarps*rows_per_warp) { + int i = (i0 + threadIdx.y*rows_per_warp + threadIdx.x/(MMQ_TILE_NE_K/8)) % I; + + if (fallback) { + i = min(i, i_max); + } + + const block_q5_K * bxi = (const block_q5_K *) x + kbx0 + i*stride; + + const int * scales = (const int *) bxi->scales; + + const int ksc = threadIdx.x % (MMQ_TILE_NE_K/8); + const int scales8 = unpack_scales_q45_K(scales, ksc); + + x_sc[i*(MMQ_TILE_NE_K/8) + i/8 + ksc] = scales8; + } +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) +} + +template static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q6_K( + const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) { + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size; + constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8); + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8); + int * x_qs = (int *) x_tile; + float * x_df = (float *) (x_qs + MMQ_TILE_NE_K*2); + int * x_sc = (int *) (x_df + MMQ_TILE_NE_K/QI6_K); +#else + constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q6_K, I); + int * x_qs = (int *) x_tile; + float * x_df = (float *) (x_qs + txs.qs); + int * x_sc = (int *) (x_df + txs.dm); +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + + constexpr int threads_per_row = MMQ_ITER_K / (4 * QR6_K); + constexpr int nrows = warp_size / threads_per_row; + const int txi = warp_size > threads_per_row ? threadIdx.x % threads_per_row : threadIdx.x; + +#pragma unroll + for (int i0 = 0; i0 < I; i0 += nrows*nwarps) { + int i = i0 + (nrows == 1 ? threadIdx.y : threadIdx.y*nrows + threadIdx.x/threads_per_row); + + if (fallback) { + i = min(i, i_max); + } + + const block_q6_K * bxi = (const block_q6_K *) x + kbx0 + i*stride; + + const int ql = get_int_b2(bxi->ql, txi); + const int ql0 = (ql >> 0) & 0x0F0F0F0F; + const int ql1 = (ql >> 4) & 0x0F0F0F0F; + + const int qh = get_int_b2(bxi->qh, (QI6_K/4) * (txi / (QI6_K/2)) + txi % (QI6_K/4)); + const int qh0 = ((qh >> ((txi & 0x08) >> 2)) << 4) & 0x30303030; + const int qh1 = (qh >> ((txi & 0x08) >> 2)) & 0x30303030; + + const int kq0 = 2*txi - txi % (QI6_K/2) + 0; + const int kq1 = 2*txi - txi % (QI6_K/2) + QI6_K/2; + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + x_qs[i*sram_stride + kq0] = __vsub4(ql0 | qh0, 0x20202020); + x_qs[i*sram_stride + kq1] = __vsub4(ql1 | qh1, 0x20202020); +#else + x_qs[i*(2*MMQ_TILE_NE_K + 1) + kq0] = __vsub4(ql0 | qh0, 0x20202020); + x_qs[i*(2*MMQ_TILE_NE_K + 1) + kq1] = __vsub4(ql1 | qh1, 0x20202020); +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + } + +#pragma unroll + for (int i0 = 0; i0 < I; i0 += nwarps*warp_size) { + int i = (i0 + threadIdx.y*warp_size + threadIdx.x) % I; + + if (fallback) { + i = min(i, i_max); + } + + const block_q6_K * bxi = (const block_q6_K *) x + kbx0 + i*stride; + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + x_df[i*sram_stride] = bxi->d; +#else + x_df[i*(MMQ_TILE_NE_K/QI6_K) + i/QI6_K] = bxi->d; +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + } + + constexpr int rows_per_warp = warp_size / 4; +#pragma unroll + for (int i0 = 0; i0 < I; i0 += nwarps*rows_per_warp) { + int i = (i0 + threadIdx.y*rows_per_warp + threadIdx.x/(MMQ_TILE_NE_K/8)) % I; + + if (fallback) { + i = min(i, i_max); + } + + const block_q6_K * bxi = (const block_q6_K *) x + kbx0 + i*stride + (threadIdx.x % (MMQ_TILE_NE_K/8)) / 4; + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + x_sc[i*sram_stride + threadIdx.x%4] = get_int_b2(bxi->scales, threadIdx.x % (MMQ_TILE_NE_K/8)); +#else + x_sc[i*(MMQ_TILE_NE_K/8) + i/8 + threadIdx.x%(MMQ_TILE_NE_K/8)] = get_int_b2(bxi->scales, threadIdx.x%(QI6_K/8)); +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + } +} + +// --------------------------------------------------------------------------------------------- + +template static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_iq1_s( + const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) { + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size; + constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8); + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8); + int * x_qs = (int *) x_tile; + half2 * x_ds = (half2 *) (x_qs + MMQ_TILE_NE_K*2); +#else + constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_IQ3_S, I); + int * x_qs = (int *) x_tile; + half2 * x_ds = (half2 *) (x_qs + txs.qs); +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + + constexpr int threads_per_row = MMQ_ITER_K / (4 * QR1_S); + constexpr int nrows = warp_size / threads_per_row; + const int kqsx = threadIdx.x % threads_per_row; + +#pragma unroll + for (int i0 = 0; i0 < I; i0 += nwarps * nrows) { + int i = i0 + threadIdx.y*nrows + threadIdx.x/threads_per_row; + + if (fallback) { + i = min(i, i_max); + } + + const block_iq1_s * bxi = (const block_iq1_s *) x + kbx0 + i*stride; + + const int qs_packed = get_int_b2(bxi->qs, kqsx); + const uint8_t * qs = (const uint8_t *) &qs_packed; + + const int qh = bxi->qh[kqsx]; + + #pragma unroll + for (int l = 0; l < QR1_S/2; ++l) { + const int grid = iq1s_grid_gpu[qs[l] | (((qh >> (3*l)) & 0x07) << 8)]; + + const int grid0 = (grid >> 0) & 0x0F0F0F0F; + const int grid1 = (grid >> 4) & 0x0F0F0F0F; + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + x_qs[i*sram_stride + 8*kqsx + (2*l+0)] = grid0; + x_qs[i*sram_stride + 8*kqsx + (2*l+1)] = grid1; +#else + x_qs[i*(2*MMQ_TILE_NE_K + 1) + 8*kqsx + (2*l+0)] = grid0; + x_qs[i*(2*MMQ_TILE_NE_K + 1) + 8*kqsx + (2*l+1)] = grid1; +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + } + + const float d1q = __half2float(bxi->d) * (((qh >> 11) & 0x0E) + 1); + const float delta = -1.0f + IQ1S_DELTA - (qh & 0x8000) * (2.0f*IQ1S_DELTA/0x8000); + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + x_ds[i*sram_stride + kqsx] = make_half2(d1q, d1q*delta); +#else + x_ds[i*(MMQ_TILE_NE_K/4) + i/4 + kqsx] = make_half2(d1q, d1q*delta); +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + } +} + +template static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_iq2_xxs( + const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) { + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size; + constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8); + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8); + int * x_qs = (int *) x_tile; + float * x_df = (float *) (x_qs + MMQ_TILE_NE_K*2); +#else + constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_IQ2_XXS, I); + int * x_qs = (int *) x_tile; + float * x_df = (float *) (x_qs + txs.qs); +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + + constexpr int threads_per_row = (MMQ_ITER_K / (4 * QR2_XXS)) / 2; + constexpr int nrows = warp_size / threads_per_row; + const int kqsx = warp_size > threads_per_row ? threadIdx.x % threads_per_row : threadIdx.x; + +#pragma unroll + for (int i0 = 0; i0 < I; i0 += nwarps * nrows) { + int i = i0 + threadIdx.y*nrows + threadIdx.x/threads_per_row; + + if (fallback) { + i = min(i, i_max); + } + + const block_iq2_xxs * bxi = (const block_iq2_xxs *) x + kbx0 + i*stride; + + const int q2 = get_int_b2(bxi->qs, 2*kqsx+0); + const uint8_t * aux8 = (const uint8_t *) &q2; + const uint32_t aux32 = get_int_b2(bxi->qs, 2*kqsx+1); + +#pragma unroll + for (int l = 0; l < QR2_XXS; ++l) { + const uint2 grid_pos = ((const uint2*)iq2xxs_grid)[aux8[l]]; + const uint32_t signs = unpack_ksigns(aux32 >> (7 * l)); + + const int signs0 = __vcmpne4(signs & 0x08040201, 0); + const int grid0 = __vsub4(grid_pos.x ^ signs0, signs0); + + const int signs1 = __vcmpne4(signs & 0x80402010, 0); + const int grid1 = __vsub4(grid_pos.y ^ signs1, signs1); + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + x_qs[i*sram_stride + 8*kqsx + (2*l + 0)] = grid0; + x_qs[i*sram_stride + 8*kqsx + (2*l + 1)] = grid1; +#else + x_qs[i*(2*MMQ_TILE_NE_K + 1) + 8*kqsx + (2*l + 0)] = grid0; + x_qs[i*(2*MMQ_TILE_NE_K + 1) + 8*kqsx + (2*l + 1)] = grid1; +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + } + + const int ls = aux32 >> 27 | 1; // (scale * 2 + 1) + const float d = bxi->d; +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + x_df[i*sram_stride + kqsx] = d * ls / 8; // (d * scale + d / 2) / 4 +#else + x_df[i*(MMQ_TILE_NE_K/4) + i/4 + kqsx] = d * ls / 8; // (d * scale + d / 2) / 4 +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + } +} + +template static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_iq2_xs( + const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) { + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size; + constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8); + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8); + int * x_qs = (int *) x_tile; + float * x_df = (float *) (x_qs + MMQ_TILE_NE_K*2); +#else + constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_IQ2_XS, I); + int * x_qs = (int *) x_tile; + float * x_df = (float *) (x_qs + txs.qs); +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + + constexpr int threads_per_row = (MMQ_ITER_K / (4 * QR2_XS)) / 2; + constexpr int nrows = warp_size / threads_per_row; + const int kqsx = threadIdx.x % threads_per_row; + +#pragma unroll + for (int i0 = 0; i0 < I; i0 += nwarps * nrows) { + int i = i0 + threadIdx.y*nrows + threadIdx.x/threads_per_row; + + if (fallback) { + i = min(i, i_max); + } + + const block_iq2_xs * bxi = (const block_iq2_xs *) x + kbx0 + i*stride; + + const int2 q2_packed = make_int2(get_int_b2(bxi->qs, 2*kqsx+0), get_int_b2(bxi->qs, 2*kqsx+1)); + const uint16_t * q2 = (const uint16_t *) &q2_packed; + + #pragma unroll + for (int l = 0; l < QR2_XS; ++l) { + const uint2 grid_pos = ((const uint2*)iq2xs_grid)[q2[l] & 0x1FF]; + const uint32_t signs = unpack_ksigns(q2[l] >> 9); + + const int signs0 = __vcmpne4(signs & 0x08040201, 0); + const int grid_l = __vsub4(grid_pos.x ^ signs0, signs0); + + const int signs1 = __vcmpne4(signs & 0x80402010, 0); + const int grid_h = __vsub4(grid_pos.y ^ signs1, signs1); + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + x_qs[i*sram_stride + 8*kqsx + (2*l + 0)] = grid_l; + x_qs[i*sram_stride + 8*kqsx + (2*l + 1)] = grid_h; +#else + x_qs[i*(2*MMQ_TILE_NE_K + 1) + 8*kqsx + (2*l + 0)] = grid_l; + x_qs[i*(2*MMQ_TILE_NE_K + 1) + 8*kqsx + (2*l + 1)] = grid_h; +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + } + + const int ls = bxi->scales[kqsx]; + const float d = bxi->d; +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + x_df[i*sram_stride + 2*kqsx+0] = ((ls & 0x0F)*d + d/2)/4; + x_df[i*sram_stride + 2*kqsx+1] = ((ls >> 4)*d + d/2)/4; +#else + x_df[i*(2*MMQ_TILE_NE_K*2/QI8_0) + i/(QI8_0/4) + 2*kqsx+0] = ((ls & 0x0F)*d + d/2)/4; + x_df[i*(2*MMQ_TILE_NE_K*2/QI8_0) + i/(QI8_0/4) + 2*kqsx+1] = ((ls >> 4)*d + d/2)/4; +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + } +} + +template static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_iq2_s( + const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) { + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size; + constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8); + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8); + int * x_qs = (int *) x_tile; + float * x_df = (float *) (x_qs + MMQ_TILE_NE_K*2); +#else + constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_IQ2_S, I); + int * x_qs = (int *) x_tile; + float * x_df = (float *) (x_qs + txs.qs); +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + constexpr int threads_per_row = (MMQ_ITER_K / (4 * QR2_S)) / 2; + constexpr int nrows = warp_size / threads_per_row; + const int kqsx = threadIdx.x % threads_per_row; + +#pragma unroll + for (int i0 = 0; i0 < I; i0 += nwarps * nrows) { + int i = i0 + threadIdx.y*nrows + threadIdx.x/threads_per_row; + + if (fallback) { + i = min(i, i_max); + } + + const block_iq2_s * bxi = (const block_iq2_s *) x + kbx0 + i*stride; + + const int qs_packed = get_int_b2(bxi->qs, kqsx); + const uint8_t * qs = (const uint8_t *) &qs_packed; + + const int qh = bxi->qh[kqsx]; + + const int signs_packed_32 = get_int_b2(bxi->qs, QK_K/32 + kqsx); + const uint8_t * signs_packed_8 = (const uint8_t *) &signs_packed_32; + +#pragma unroll + for (int l = 0; l < QR2_S; ++l) { + const int * grid_pos = (const int *)(iq2s_grid + (qs[l] | ((qh << (8-2*l)) & 0x300))); + + const int signs0 = __vcmpne4(((signs_packed_8[l] & 0x03) << 7) | ((signs_packed_8[l] & 0x0C) << 21), 0x00000000); + const int signs1 = __vcmpne4(((signs_packed_8[l] & 0x30) << 3) | ((signs_packed_8[l] & 0xC0) << 17), 0x00000000); + + const int grid_l = __vsub4(grid_pos[0] ^ signs0, signs0); + const int grid_h = __vsub4(grid_pos[1] ^ signs1, signs1); + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + x_qs[i*sram_stride + 8*kqsx + (2*l + 0)] = grid_l; + x_qs[i*sram_stride + 8*kqsx + (2*l + 1)] = grid_h; +#else + x_qs[i*(2*MMQ_TILE_NE_K + 1) + 8*kqsx + (2*l + 0)] = grid_l; + x_qs[i*(2*MMQ_TILE_NE_K + 1) + 8*kqsx + (2*l + 1)] = grid_h; +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + } + + const int ls = bxi->scales[kqsx]; + const float d = bxi->d; +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + x_df[i*sram_stride + 2*kqsx+0] = ((ls & 0x0F)*d + d/2)/4; + x_df[i*sram_stride + 2*kqsx+1] = ((ls >> 4)*d + d/2)/4; +#else + x_df[i*(2*MMQ_TILE_NE_K*2/QI8_0) + i/(QI8_0/4) + 2*kqsx+0] = ((ls & 0x0F)*d + d/2)/4; + x_df[i*(2*MMQ_TILE_NE_K*2/QI8_0) + i/(QI8_0/4) + 2*kqsx+1] = ((ls >> 4)*d + d/2)/4; +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + } +} + +template static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_iq3_xxs( + const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) { + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size; + constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8); + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8); + int * x_qs = (int *) x_tile; + float * x_df = (float *) (x_qs + MMQ_TILE_NE_K*2); +#else + constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_IQ3_XXS, I); + int * x_qs = (int *) x_tile; + float * x_df = (float *) (x_qs + txs.qs); +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + + constexpr int threads_per_row = (MMQ_ITER_K / (4 * QR3_XXS)) / 2; + constexpr int nrows = warp_size / threads_per_row; + const int kqsx = threadIdx.x % threads_per_row; + +#pragma unroll + for (int i0 = 0; i0 < I; i0 += nwarps * nrows) { + int i = i0 + threadIdx.y*nrows + threadIdx.x/threads_per_row; + + if (fallback) { + i = min(i, i_max); + } + + const block_iq3_xxs * bxi = (const block_iq3_xxs *) x + kbx0 + i*stride; + + const int2 q3_packed = make_int2(get_int_b2(bxi->qs, 2*kqsx+0), get_int_b2(bxi->qs, 2*kqsx+1)); + const uint8_t * q3 = (const uint8_t *) &q3_packed; + const uint32_t aux32 = get_int_b2(bxi->qs, QK_K/16 + kqsx); + +#pragma unroll + for (int l = 0; l < QR3_XXS; ++l) { + const int2 grid_pos = make_int2(iq3xxs_grid[q3[2*l+0]], iq3xxs_grid[q3[2*l+1]]); + const uint32_t signs = unpack_ksigns(aux32 >> (7*l)); + + const int signs0 = __vcmpne4(signs & 0x08040201, 0); + const int grid_l = __vsub4(grid_pos.x ^ signs0, signs0); + + const int signs1 = __vcmpne4(signs & 0x80402010, 0); + const int grid_h = __vsub4(grid_pos.y ^ signs1, signs1); + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + x_qs[i*sram_stride + 8*kqsx + (2*l + 0)] = grid_l; + x_qs[i*sram_stride + 8*kqsx + (2*l + 1)] = grid_h; +#else + x_qs[i*(2*MMQ_TILE_NE_K + 1) + 8*kqsx + (2*l + 0)] = grid_l; + x_qs[i*(2*MMQ_TILE_NE_K + 1) + 8*kqsx + (2*l + 1)] = grid_h; +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + } + + const int ls = aux32 >> 28; + const float d = bxi->d; +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + x_df[i*sram_stride + kqsx] = (ls*d + d/2)/2; +#else + x_df[i*(MMQ_TILE_NE_K/4) + i/4 + kqsx] = (ls*d + d/2)/2; +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + } +} + +template static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_iq3_s( + const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) { + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size; + constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8); + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8); + int * x_qs = (int *) x_tile; + float * x_df = (float *) (x_qs + MMQ_TILE_NE_K*2); +#else + constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_IQ3_S, I); + int * x_qs = (int *) x_tile; + float * x_df = (float *) (x_qs + txs.qs); +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + + constexpr int threads_per_row = (MMQ_ITER_K / (4 * QR3_S)) / 2; + constexpr int nrows = warp_size / threads_per_row; + const int kqsx = threadIdx.x % threads_per_row; + +#pragma unroll + for (int i0 = 0; i0 < I; i0 += nwarps * nrows) { + int i = i0 + threadIdx.y*nrows + threadIdx.x/threads_per_row; + + if (fallback) { + i = min(i, i_max); + } + + const block_iq3_s * bxi = (const block_iq3_s *) x + kbx0 + i*stride; + + const int2 qs_packed = make_int2(get_int_b2(bxi->qs, 2*kqsx+0), get_int_b2(bxi->qs, 2*kqsx+1)); + const uint8_t * qs = (const uint8_t *) &qs_packed; + + const int qh = bxi->qh[kqsx]; + + const int signs_packed_32 = get_int_b2(bxi->signs, kqsx); + const uint8_t * signs_packed_8 = (const uint8_t *) &signs_packed_32; + +#pragma unroll + for (int l = 0; l < QR3_S; ++l) { + const int2 grid_pos = make_int2( + iq3s_grid[qs[2*l+0] | ((qh << (8 - 2*l)) & 0x100)], + iq3s_grid[qs[2*l+1] | ((qh << (7 - 2*l)) & 0x100)]); + + const int signs0 = __vcmpne4(((signs_packed_8[l] & 0x03) << 7) | ((signs_packed_8[l] & 0x0C) << 21), 0x00000000); + const int signs1 = __vcmpne4(((signs_packed_8[l] & 0x30) << 3) | ((signs_packed_8[l] & 0xC0) << 17), 0x00000000); + + const int grid_l = __vsub4(grid_pos.x ^ signs0, signs0); + const int grid_h = __vsub4(grid_pos.y ^ signs1, signs1); + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + x_qs[i*sram_stride + 8*kqsx + (2*l+0)] = grid_l; + x_qs[i*sram_stride + 8*kqsx + (2*l+1)] = grid_h; +#else + x_qs[i*(2*MMQ_TILE_NE_K + 1) + 8*kqsx + (2*l+0)] = grid_l; + x_qs[i*(2*MMQ_TILE_NE_K + 1) + 8*kqsx + (2*l+1)] = grid_h; +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + } + + const int ls = 1 + 2*((bxi->scales[kqsx/2] >> (((2*kqsx) << 1) & 0x04)) & 0x0F); + const float d = bxi->d; +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + x_df[i*sram_stride + kqsx] = ls*d; +#else + x_df[i*(MMQ_TILE_NE_K/4) + i/4 + kqsx] = ls*d; +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + } +} + +template static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_iq4_xs( + const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) { + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size; + constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8); + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8); + int * x_qs = (int *) x_tile; + float * x_df = (float *) (x_qs + MMQ_TILE_NE_K*2); +#else + constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_IQ4_XS, I); + int * x_qs = (int *) x_tile; + float * x_df = (float *) (x_qs + txs.qs); +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + + constexpr int threads_per_row = MMQ_ITER_K / (4 * QR4_XS); + constexpr int nrows = warp_size / threads_per_row; + const int kqsx = threadIdx.x % threads_per_row; + +#pragma unroll + for (int i0 = 0; i0 < I; i0 += nrows*nwarps) { + int i = i0 + (nrows == 1 ? threadIdx.y : threadIdx.y*nrows + threadIdx.x/threads_per_row); + + if (fallback) { + i = min(i, i_max); + } + + const block_iq4_xs * bxi = (const block_iq4_xs *) x + kbx0 + i*stride; + + const int aux_q4 = get_int_b4(bxi->qs, kqsx); + const int2 v = get_int_from_table_16(aux_q4, kvalues_iq4nl); + const int k0 = 8 * (kqsx / 4) + kqsx % 4; + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + x_qs[i*sram_stride + k0 + 0] = v.x; + x_qs[i*sram_stride + k0 + 4] = v.y; +#else + x_qs[i*(2*MMQ_TILE_NE_K + 1) + k0 + 0] = v.x; + x_qs[i*(2*MMQ_TILE_NE_K + 1) + k0 + 4] = v.y; +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + } + + constexpr int rows_per_warp = warp_size / 8; +#pragma unroll + for (int i0 = 0; i0 < I; i0 += nwarps * rows_per_warp) { + int i = i0 + threadIdx.y * rows_per_warp + threadIdx.x / (MMQ_TILE_NE_K/4); + + if (fallback) { + i = min(i, i_max); + } + + const block_iq4_xs * bxi = (const block_iq4_xs *) x + kbx0 + i*stride; + + const float d = __half2float(bxi->d); + + const int ls = ((bxi->scales_l[(threadIdx.x % 8)/2] >> (4*(threadIdx.x % 2))) & 0x0F) + | (((bxi->scales_h >> (2*(threadIdx.x % 8))) & 0x03) << 4); + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + x_df[i*sram_stride + threadIdx.x % 8] = d * (ls - 32); +#else + x_df[i*(MMQ_TILE_NE_K/4) + i/4 + threadIdx.x % 8] = d * (ls - 32); +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + } +} + +template static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_iq4_nl( + const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) { + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size; + constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8); + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8); + int * x_qs = (int *) x_tile; + float * x_df = (float *) (x_qs + MMQ_TILE_NE_K*2); +#else + constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_IQ4_NL, I); + int * x_qs = (int *) x_tile; + float * x_df = (float *) (x_qs + txs.qs); +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + + constexpr int threads_per_row = MMQ_ITER_K / (4 * QR4_NL); + constexpr int nrows = warp_size / threads_per_row; + const int txi = warp_size > threads_per_row ? threadIdx.x % threads_per_row : threadIdx.x; + const int kbx = txi / QI4_NL; + const int kqsx = txi % QI4_NL; + +#pragma unroll + for (int i0 = 0; i0 < I; i0 += nrows*nwarps) { + int i = i0 + (nrows == 1 ? threadIdx.y : threadIdx.y*nrows + threadIdx.x/threads_per_row); + + if (fallback) { + i = min(i, i_max); + } + + const block_iq4_nl * bxi = (const block_iq4_nl *) x + kbx0 + i*stride + kbx; + + const int aux_q4 = get_int_b2(bxi->qs, kqsx); + const int2 v = get_int_from_table_16(aux_q4, kvalues_iq4nl); + const int k0 = kbx * (2 * QI4_NL) + kqsx; + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + x_qs[i*sram_stride + k0 + 0] = v.x; + x_qs[i*sram_stride + k0 + QI4_NL] = v.y; +#else + x_qs[i*(2*MMQ_TILE_NE_K + 1) + k0 + 0] = v.x; + x_qs[i*(2*MMQ_TILE_NE_K + 1) + k0 + QI4_NL] = v.y; +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + } + + constexpr int blocks_per_tile_x_row = MMQ_TILE_NE_K / QI4_NL; + constexpr int rows_per_warp = warp_size / blocks_per_tile_x_row; + const int kbxd = threadIdx.x % blocks_per_tile_x_row; + +#pragma unroll + for (int i0 = 0; i0 < I; i0 += nwarps * rows_per_warp) { + int i = i0 + threadIdx.y * rows_per_warp + threadIdx.x / blocks_per_tile_x_row; + + if (fallback) { + i = min(i, i_max); + } + + const block_iq4_nl * bxi = (const block_iq4_nl *) x + kbx0 + i*stride + kbxd; + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + x_df[i*sram_stride + kbxd] = __half2float(bxi->d); +#else + x_df[i*(MMQ_TILE_NE_K/QI4_NL) + i/QI4_NL + kbxd] = __half2float(bxi->d); +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + } +} + +// --------------------------------------------------------------------------------------------- + +template static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_mxfp4( + const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) { + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size; + constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8); + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8); + int * x_qs = (int *) x_tile; + float * x_df = (float *) (x_qs + MMQ_TILE_NE_K*2); +#else + constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_MXFP4, I); + int * x_qs = (int *) x_tile; + float * x_df = (float *) (x_qs + txs.qs); +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + + constexpr int threads_per_row = MMQ_ITER_K / (4 * QR_MXFP4); + constexpr int nrows = warp_size / threads_per_row; + const int txi = warp_size > threads_per_row ? threadIdx.x % threads_per_row : threadIdx.x; + const int kbx = txi / QI_MXFP4; + const int kqsx = txi % QI_MXFP4; + +#pragma unroll + for (int i0 = 0; i0 < I; i0 += nrows*nwarps) { + int i = i0 + (nrows == 1 ? threadIdx.y : threadIdx.y*nrows + threadIdx.x/threads_per_row); + + if (fallback) { + i = min(i, i_max); + } + + const block_mxfp4 * bxi = (const block_mxfp4 *) x + kbx0 + i*stride + kbx; + + const int aux_q4 = get_int_b1(bxi->qs, kqsx); + const int2 v = get_int_from_table_16(aux_q4, kvalues_mxfp4); + const int k0 = kbx * (2 * QI_MXFP4) + kqsx; + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + x_qs[i*sram_stride + k0 + 0] = v.x; + x_qs[i*sram_stride + k0 + QI_MXFP4] = v.y; +#else + x_qs[i*(2*MMQ_TILE_NE_K + 1) + k0 + 0] = v.x; + x_qs[i*(2*MMQ_TILE_NE_K + 1) + k0 + QI_MXFP4] = v.y; +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + } + + constexpr int blocks_per_tile_x_row = MMQ_TILE_NE_K / QI_MXFP4; + constexpr int rows_per_warp = warp_size / blocks_per_tile_x_row; + const int kbxd = threadIdx.x % blocks_per_tile_x_row; + +#pragma unroll + for (int i0 = 0; i0 < I; i0 += nwarps * rows_per_warp) { + int i = i0 + threadIdx.y * rows_per_warp + threadIdx.x / blocks_per_tile_x_row; + + if (fallback) { + i = min(i, i_max); + } + + const block_mxfp4 * bxi = (const block_mxfp4 *) x + kbx0 + i*stride + kbxd; + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + x_df[i*sram_stride + kbxd] = ggml_cuda_e8m0_to_fp32(bxi->e)*0.5f; +#else + x_df[i*(MMQ_TILE_NE_K/QI_MXFP4) + i/QI_MXFP4 + kbxd] = ggml_cuda_e8m0_to_fp32(bxi->e)*0.5f; +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + } +} + +template static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_mxfp4_fp4( + const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) { + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q4) / warp_size; + constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q4); + constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q4); + + int * x_qs = (int *) x_tile; + uint32_t * x_sc = (uint32_t *) (x_qs + 2 * MMQ_TILE_NE_K); + + const int txi = threadIdx.x; + + constexpr int iter_k = ggml_cuda_mmq_get_K_vram(type, J, fallback, GGML_PREC_Q4); + + constexpr int threads_per_row = iter_k / QK_MXFP4; // each thread processes 1 block + constexpr int rows_per_warp = warp_size / threads_per_row; + const int kbx = txi % threads_per_row; + const int row_in_warp = txi / threads_per_row; + +#pragma unroll + for (int i0 = 0; i0 < I; i0 += rows_per_warp * nwarps) { + int i = i0 + threadIdx.y * rows_per_warp + row_in_warp; + + if constexpr (fallback) { + i = min(i, i_max); + } + + const block_mxfp4 * bxi = (const block_mxfp4 *) x + kbx0 + i * stride + kbx; + + // quantize_mxfp4_mmq permutes nibbles to match the quantized format + const int k0 = kbx * 4; + memcpy(x_qs + i*sram_stride + k0, bxi->qs, 16); + + // Load E8M0 scales: pack 2 consecutive scales into one uint32 + if (kbx % 2 == 0) { + uint32_t e = bxi->e; + e |= ((bxi + 1)->e << 8); + x_sc[i*sram_stride + kbx / 2] = e; + } + } +} + +template static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_nvfp4( + const char * __restrict__ x, int * __restrict__ x_tile, const int kb0, const int i_max, const int stride) { + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback, prec_src1) / warp_size; + constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback, prec_src1); + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, prec_src1); + int * x_qs = (int *) x_tile; + float * x_df = (float *) (x_qs + MMQ_TILE_NE_K*2); +#else + constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_NVFP4, I); + int * x_qs = (int *) x_tile; + float * x_df = (float *) (x_qs + txs.qs); +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + + constexpr int threads_per_row = MMQ_ITER_K / QK_NVFP4; + constexpr int rows_per_warp = warp_size / threads_per_row; + const int kbx = threadIdx.x % threads_per_row; + const int row_in_warp = threadIdx.x / threads_per_row; + +#pragma unroll + for (int i0 = 0; i0 < I; i0 += rows_per_warp * nwarps) { + int i = i0 + threadIdx.y * rows_per_warp + row_in_warp; + + if constexpr (fallback) { + i = min(i, i_max); + } + + const block_nvfp4 * bxi = (const block_nvfp4 *) x + kb0 + i * stride + kbx; + const uint32_t * __restrict__ src_qs = reinterpret_cast(bxi->qs); + const int kqs = 16 * kbx; + const int ksc = 4 * kbx; + +#pragma unroll + for (int sub = 0; sub < QK_NVFP4 / QK_NVFP4_SUB; ++sub) { + const int2 q0 = get_int_from_table_16(src_qs[2 * sub + 0], kvalues_mxfp4); + const int2 q1 = get_int_from_table_16(src_qs[2 * sub + 1], kvalues_mxfp4); + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + x_qs[i*sram_stride + kqs + 4 * sub + 0] = q0.x; + x_qs[i*sram_stride + kqs + 4 * sub + 1] = q1.x; + x_qs[i*sram_stride + kqs + 4 * sub + 2] = q0.y; + x_qs[i*sram_stride + kqs + 4 * sub + 3] = q1.y; + x_df[i*sram_stride + ksc + sub] = ggml_cuda_ue4m3_to_fp32(bxi->d[sub]); +#else + x_qs[i * (2 * MMQ_TILE_NE_K + 1) + kqs + 4 * sub + 0] = q0.x; + x_qs[i * (2 * MMQ_TILE_NE_K + 1) + kqs + 4 * sub + 1] = q1.x; + x_qs[i * (2 * MMQ_TILE_NE_K + 1) + kqs + 4 * sub + 2] = q0.y; + x_qs[i * (2 * MMQ_TILE_NE_K + 1) + kqs + 4 * sub + 3] = q1.y; + x_df[i * (2 * MMQ_TILE_NE_K * 2 / QI_NVFP4) + i / (QK_NVFP4_SUB / QI_NVFP4) + ksc + sub] = ggml_cuda_ue4m3_to_fp32(bxi->d[sub]); +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + } + } +} + +template static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_nvfp4_nvfp4( + const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) { + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q4) / warp_size; + constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q4); + constexpr int iter_k = ggml_cuda_mmq_get_K_vram(type, J, fallback, GGML_PREC_Q4); + constexpr int threads_per_row = iter_k / QK_NVFP4; // each thread processes 1 block + constexpr int rows_per_warp = warp_size / threads_per_row; + constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q4); + + uint32_t * x_u32 = (uint32_t *) x_tile; + + const int txi = threadIdx.x; + const int kbx = txi % threads_per_row; + const int row_in_warp = txi / threads_per_row; + + const block_nvfp4 * bxi_base = (const block_nvfp4 *) x + kbx0 + kbx; + uint32_t * x_u32_scale = x_u32 + 64 + kbx; + +#pragma unroll + for (int i0 = 0; i0 < I; i0 += rows_per_warp * nwarps) { + int i = i0 + threadIdx.y * rows_per_warp + row_in_warp; + + if constexpr (fallback) { + i = min(i, i_max); + } + + const block_nvfp4 * bxi = bxi_base + i * stride; + + const uint32_t * src_qs = reinterpret_cast(bxi->qs); + +#pragma unroll + for (int sub = 0; sub < QK_NVFP4 / QK_NVFP4_SUB; ++sub) { + x_u32[i*sram_stride + 8*kbx + 2 * sub + 0] = src_qs[2 * sub + 0]; + x_u32[i*sram_stride + 8*kbx + 2 * sub + 1] = src_qs[2 * sub + 1]; + } + + x_u32_scale[i*sram_stride] = get_int_b4(bxi->d, 0); + } +} diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq-vec-dot.cuh b/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq-vec-dot.cuh new file mode 100644 index 0000000000..a75dd046e5 --- /dev/null +++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq-vec-dot.cuh @@ -0,0 +1,1237 @@ +#pragma once + +#include "vecdotq.cuh" +#include "mma.cuh" + +using namespace ggml_cuda_mma; + +#include "mmq.cuh" + +template static __device__ __forceinline__ void ggml_cuda_mmq_vec_dot_q4_0_q8_1_dp4a( + const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) { + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size; + constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8); + + constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q4_0, I); + const int * x_qs = (const int *) x; + const float * x_df = (const float *) x_qs + txs.qs; + const int * y_qs = (const int *) y + 4; + const half2 * y_ds = (const half2 *) y; + +// #pragma unroll + for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QR4_0*VDR_Q4_0_Q8_1_MMQ) { + const int k0 = k00 + k01; + +#pragma unroll + for (int j0 = 0; j0 < J; j0 += nwarps) { + const int j = j0 + threadIdx.y; + +#pragma unroll + for (int i0 = 0; i0 < I; i0 += warp_size) { + const int i = i0 + threadIdx.x; + const int kyqs = QI8_1 * ((k01/2) / (QI8_1/2)) + (k01/2) % (QI8_1/2); + + int u[2*VDR_Q4_0_Q8_1_MMQ]; + + constexpr int max_cpy = ggml_cuda_get_max_cpy_bytes(); + constexpr int mcpy_int = max_cpy / sizeof(int); + static_assert(VDR_Q4_0_Q8_1_MMQ == 4, "bad VDR_Q4_0_Q8_1_MMQ"); + + int tmp0[4], tmp1[4]; + + #pragma unroll + for (int l0 = 0; l0 < 4 / mcpy_int; ++l0) { + ggml_cuda_memcpy_1(tmp0 + l0 * mcpy_int, &y_qs[j*MMQ_TILE_Y_K + kyqs + l0 * mcpy_int] ); + ggml_cuda_memcpy_1(tmp1 + l0 * mcpy_int, &y_qs[j*MMQ_TILE_Y_K + kyqs + QI4_0 + l0 * mcpy_int]); + } + + u[0]=tmp0[0]; u[2]=tmp0[1]; u[4]=tmp0[2]; u[6]=tmp0[3]; + u[1]=tmp1[0]; u[3]=tmp1[1]; u[5]=tmp1[2]; u[7]=tmp1[3]; + + sum[j0/nwarps*I/warp_size + i0/warp_size] += vec_dot_q4_0_q8_1_impl + (&x_qs[i*(MMQ_TILE_NE_K + 1) + k0/QR4_0], u, + x_df[i*(MMQ_TILE_NE_K/QI4_0) + i/QI4_0 + k0/(QR4_0*QI4_0)], y_ds[j*MMQ_TILE_Y_K + k01/QI8_1]); + } + } + } +} + +template static __device__ __forceinline__ void ggml_cuda_mmq_vec_dot_q4_1_q8_1_dp4a( + const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) { + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size; + constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8); + + constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q4_1, I); + const int * x_qs = (const int *) x; + const half2 * x_dm = (const half2 *) x_qs + txs.qs; + const int * y_qs = (const int *) y + 4; + const half2 * y_ds = (const half2 *) y; + +// #pragma unroll + for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QR4_1*VDR_Q4_1_Q8_1_MMQ) { + const int k0 = k00 + k01; + +#pragma unroll + for (int j0 = 0; j0 < J; j0 += nwarps) { + const int j = j0 + threadIdx.y; + +#pragma unroll + for (int i0 = 0; i0 < I; i0 += warp_size) { + const int i = i0 + threadIdx.x; + const int kyqs = QI8_1 * ((k01/2) / (QI8_1/2)) + (k01/2) % (QI8_1/2); + + int u[2*VDR_Q4_1_Q8_1_MMQ]; + + constexpr int max_cpy = ggml_cuda_get_max_cpy_bytes(); + constexpr int mcpy_int = max_cpy / sizeof(int); + static_assert(VDR_Q4_0_Q8_1_MMQ == 4, "bad VDR_Q4_0_Q8_1_MMQ"); + + int tmp0[4], tmp1[4]; + + #pragma unroll + for (int l0 = 0; l0 < 4 / mcpy_int; ++l0) { + ggml_cuda_memcpy_1(tmp0 + l0 * mcpy_int, &y_qs[j*MMQ_TILE_Y_K + kyqs + l0 * mcpy_int] ); + ggml_cuda_memcpy_1(tmp1 + l0 * mcpy_int, &y_qs[j*MMQ_TILE_Y_K + kyqs + QI4_1 + l0 * mcpy_int]); + } + + u[0]=tmp0[0]; u[2]=tmp0[1]; u[4]=tmp0[2]; u[6]=tmp0[3]; + u[1]=tmp1[0]; u[3]=tmp1[1]; u[5]=tmp1[2]; u[7]=tmp1[3]; + + sum[j0/nwarps*I/warp_size + i0/warp_size] += vec_dot_q4_1_q8_1_impl + (&x_qs[i*(MMQ_TILE_NE_K + 1) + k0/QR4_1], u, + x_dm[i*(MMQ_TILE_NE_K/QI4_1) + i/QI4_1 + k0/(QR4_1*QI4_1)], y_ds[j*MMQ_TILE_Y_K + k01/QI8_1]); + } + } + } +} + +template static __device__ __forceinline__ void ggml_cuda_mmq_vec_dot_q8_0_q8_1_dp4a( + const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) { + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size; + constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8); + + constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q8_0, I); + const int * x_qs = (const int *) x; + const float * x_df = (const float *) x_qs + txs.qs; + const int * y_qs = (const int *) y + 4; + const float * y_df = (const float *) y; + +// #pragma unroll + for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += VDR_Q8_0_Q8_1_MMQ) { + const int k0 = k00 + k01; + +#pragma unroll + for (int j0 = 0; j0 < J; j0 += nwarps) { + const int j = j0 + threadIdx.y; + +#pragma unroll + for (int i0 = 0; i0 < I; i0 += warp_size) { + const int i = i0 + threadIdx.x; + + sum[j0/nwarps*I/warp_size + i0/warp_size] += vec_dot_q8_0_q8_1_impl + (&x_qs[i*(2*MMQ_TILE_NE_K + 1) + k0], &y_qs[j*MMQ_TILE_Y_K + k0 % MMQ_TILE_NE_K], + x_df[i*(2*MMQ_TILE_NE_K/QI8_0) + i/(QI8_0/2) + k0/QI8_0], y_df[j*MMQ_TILE_Y_K + (k0/QI8_1) % (MMQ_TILE_NE_K/QI8_1)]); + } + } + } +} + +template +static __device__ __forceinline__ void ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma( + const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) { +#if defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + constexpr data_layout input_layout = get_input_data_layout(); + typedef tile<16, 8, int, input_layout> tile_A; + typedef tile<16, 8, int, input_layout> tile_B; + typedef tile<16, 16, int, DATA_LAYOUT_J_MAJOR> tile_C; + + constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8); + constexpr int rows_per_warp = ggml_cuda_mmq_get_rows_per_warp(type, J, fallback, GGML_PREC_Q8); + constexpr int ntx = rows_per_warp/tile_C::I; // Number of x minitiles per warp. + + y += (threadIdx.y % ntx) * (tile_C::J*MMQ_TILE_Y_K); + + const int * x_qs = (const int *) x; + const float * x_df = (const float *) x_qs + 2*MMQ_TILE_NE_K; + const int * y_qs = (const int *) y + 4; + const float * y_df = (const float *) y; + const half2 * y_ds = (const half2 *) y; + + const int i0 = (threadIdx.y / ntx) * rows_per_warp; + + for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QI8_0) { + const int k0 = k00 + k01; + + tile_A A[ntx]; +#pragma unroll + for (int n = 0; n < ntx; ++n) { + load_ldmatrix(A[n], x_qs + (i0 + n*tile_A::I)*sram_stride + k0, sram_stride); + } + +#pragma unroll + for (int j0 = 0; j0 < J; j0 += ntx*tile_C::J) { + tile_B B; + load_ldmatrix(B, y_qs + j0*MMQ_TILE_Y_K + k01, MMQ_TILE_Y_K); + + float dB; + const int j = j0 + tile_C::get_j(0); + if (ds_layout == MMQ_Q8_1_DS_LAYOUT_D4) { + dB = y_df[j*MMQ_TILE_Y_K + k01/QI8_1]; + } else { + dB = __low2float(y_ds[j*MMQ_TILE_Y_K + k01/QI8_1]); + } + +#pragma unroll + for (int n = 0; n < ntx; ++n) { + tile_C C; + mma(C, A[n], B); + +#pragma unroll + for (int l = 0; l < tile_C::ne; ++l) { + const int i = i0 + n*tile_A::I + tile_C::get_i(l); + const float dA = x_df[i*sram_stride + k0/QI8_0]; + sum[(j0/tile_C::J + n)*tile_C::ne + l] += C.x[l]*dA*dB; + } + } + } + } +#else + typedef tile<16, 8, int> tile_A; + typedef tile< 8, 8, int> tile_B; + typedef tile<16, 8, int> tile_C; + + constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8); + constexpr int rows_per_warp = ggml_cuda_mmq_get_rows_per_warp(type, J, fallback, GGML_PREC_Q8); + constexpr int ntx = rows_per_warp/tile_C::I; // Number of x minitiles per warp. + + y += (threadIdx.y % ntx) * (tile_C::J*MMQ_TILE_Y_K); + + const int * x_qs = (const int *) x; + const float * x_df = (const float *) x_qs + 2*MMQ_TILE_NE_K; + const int * y_qs = (const int *) y + 4; + const float * y_df = (const float *) y; + const half2 * y_ds = (const half2 *) y; + + tile_A A[ntx][MMQ_TILE_NE_K/QI8_0]; + float dA[ntx][tile_C::ne/2][MMQ_TILE_NE_K/QI8_0]; + + const int i0 = (threadIdx.y/ntx)*rows_per_warp; + +#pragma unroll + for (int n = 0; n < ntx; ++n) { +#pragma unroll + for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QI8_0) { + const int k0 = k00 + k01; + + load_ldmatrix(A[n][k01/QI8_0], x_qs + (i0 + n*tile_A::I)*sram_stride + k0, sram_stride); + } + +#pragma unroll + for (int l = 0; l < tile_C::ne/2; ++l) { + const int i = i0 + n*tile_A::I + tile_C::get_i(2*l); + +#pragma unroll + for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QI8_0) { + const int k0 = k00 + k01; + + dA[n][l][k01/QI8_0] = x_df[i*sram_stride + k0/QI8_0]; + } + } + } + +#pragma unroll + for (int j0 = 0; j0 < J; j0 += ntx*tile_C::J) { +#pragma unroll + for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QI8_0) { + tile_B B; + float dB[tile_C::ne/2]; + + load_generic(B, y_qs + j0*MMQ_TILE_Y_K + k01, MMQ_TILE_Y_K); // faster than load_ldmatrix + +#pragma unroll + for (int l = 0; l < tile_C::ne/2; ++l) { + const int j = j0 + tile_C::get_j(l); + + if (ds_layout == MMQ_Q8_1_DS_LAYOUT_D4) { + dB[l] = y_df[j*MMQ_TILE_Y_K + k01/QI8_1]; + } else { + dB[l] = __low2float(y_ds[j*MMQ_TILE_Y_K + k01/QI8_1]); + } + } + +#pragma unroll + for (int n = 0; n < ntx; ++n) { + tile_C C; + mma(C, A[n][k01/QI8_0], B); + +#pragma unroll + for (int l = 0; l < tile_C::ne; ++l) { + sum[(j0/tile_C::J + n)*tile_C::ne + l] += C.x[l]*dA[n][l/2][k01/QI8_0]*dB[l%2]; + } + } + } + } +#endif // defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) +} + + +template static __device__ __forceinline__ void ggml_cuda_mmq_vec_dot_q8_1_q8_1_dp4a( + const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) { + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size; + constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8); + + constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q5_1, I); + const int * x_qs = (const int *) x; + const half2 * x_dm = (const half2 *) x_qs + txs.qs; + const int * y_qs = (const int *) y + 4; + const half2 * y_ds = (const half2 *) y; + +// #pragma unroll + for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += VDR_Q8_0_Q8_1_MMQ) { + const int k0 = k00 + k01; + +#pragma unroll + for (int j0 = 0; j0 < J; j0 += nwarps) { + const int j = j0 + threadIdx.y; + +#pragma unroll + for (int i0 = 0; i0 < I; i0 += warp_size) { + const int i = i0 + threadIdx.x; + + sum[j0/nwarps*I/warp_size + i0/warp_size] += vec_dot_q8_1_q8_1_impl + (&x_qs[i*(2*MMQ_TILE_NE_K + 1) + k0], &y_qs[j*MMQ_TILE_Y_K + k01], + x_dm[i*(MMQ_TILE_NE_K/QI5_1) + i/QI5_1 + k0/QI8_1], y_ds[j*MMQ_TILE_Y_K + k01/QI8_1]); + } + } + } +} + +template static __device__ __forceinline__ void ggml_cuda_mmq_vec_dot_q8_1_q8_1_mma( + const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) { +#if defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + constexpr data_layout input_layout = get_input_data_layout(); + typedef tile<16, 8, int, input_layout> tile_A; + typedef tile<16, 8, int, input_layout> tile_B; + typedef tile<16, 16, int, DATA_LAYOUT_J_MAJOR> tile_C; + + constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8); + constexpr int rows_per_warp = ggml_cuda_mmq_get_rows_per_warp(type, J, fallback, GGML_PREC_Q8); + constexpr int ntx = rows_per_warp/tile_C::I; // Number of x minitiles per warp. + + y += (threadIdx.y % ntx) * (tile_C::J*MMQ_TILE_Y_K); + + const int * x_qs = (const int *) x; + const half2 * x_dm = (const half2 *) x_qs + 2*MMQ_TILE_NE_K; + const int * y_qs = (const int *) y + 4; + const half2 * y_dm = (const half2 *) y; + + const int i0 = (threadIdx.y / ntx) * rows_per_warp; + + for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QI8_1) { + const int k0 = k00 + k01; + + tile_A A[ntx]; +#pragma unroll + for (int n = 0; n < ntx; ++n) { + load_ldmatrix(A[n], x_qs + (i0 + n*tile_A::I)*sram_stride + k0, sram_stride); + } + +#pragma unroll + for (int j0 = 0; j0 < J; j0 += ntx*tile_C::J) { + tile_B B; + load_ldmatrix(B, y_qs + j0*MMQ_TILE_Y_K + k01, MMQ_TILE_Y_K); + + const int j = j0 + tile_C::get_j(0); + const float2 dsB = __half22float2(y_dm[j*MMQ_TILE_Y_K + k01/QI8_1]); + +#pragma unroll + for (int n = 0; n < ntx; ++n) { + tile_C C; + mma(C, A[n], B); + +#pragma unroll + for (int l = 0; l < tile_C::ne; ++l) { + const int i = i0 + n*tile_A::I + tile_C::get_i(l); + float2 dmA = __half22float2(x_dm[i*sram_stride + k0/QI8_1]); + sum[(j0/tile_C::J + n)*tile_C::ne + l] += dmA.x*dsB.x*C.x[l]; + sum[(j0/tile_C::J + n)*tile_C::ne + l] += dmA.y*dsB.y; + } + } + } + } +#else + typedef tile<16, 8, int> tile_A; + typedef tile< 8, 8, int> tile_B; + typedef tile<16, 8, int> tile_C; + + constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8); + constexpr int rows_per_warp = ggml_cuda_mmq_get_rows_per_warp(type, J, fallback, GGML_PREC_Q8); + constexpr int ntx = rows_per_warp/tile_C::I; // Number of x minitiles per warp. + + y += (threadIdx.y % ntx) * (tile_C::J*MMQ_TILE_Y_K); + + const int * x_qs = (const int *) x; + const half2 * x_dm = (const half2 *) x_qs + 2*MMQ_TILE_NE_K; + const int * y_qs = (const int *) y + 4; + const half2 * y_dm = (const half2 *) y; + + tile_A A[ntx][MMQ_TILE_NE_K/QI8_1]; + float2 dmA[ntx][tile_C::ne/2][MMQ_TILE_NE_K/QI8_1]; + + const int i0 = (threadIdx.y/ntx)*rows_per_warp; + +#pragma unroll + for (int n = 0; n < ntx; ++n) { +#pragma unroll + for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QI8_1) { + const int k0 = k00 + k01; + + load_ldmatrix(A[n][k01/QI8_1], x_qs + (i0 + n*tile_A::I)*sram_stride + k0, sram_stride); + } + +#pragma unroll + for (int l = 0; l < tile_C::ne/2; ++l) { + const int i = i0 + n*tile_A::I + tile_C::get_i(2*l); + +#pragma unroll + for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QI8_1) { + const int k0 = k00 + k01; + + dmA[n][l][k01/QI8_1] = __half22float2(x_dm[i*sram_stride + k0/QI8_1]); + } + } + } + +#pragma unroll + for (int j0 = 0; j0 < J; j0 += ntx*tile_C::J) { +#pragma unroll + for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QI8_1) { + tile_B B; + float2 dsB[tile_C::ne/2]; + + load_generic(B, y_qs + j0*MMQ_TILE_Y_K + k01, MMQ_TILE_Y_K); // faster than load_ldmatrix + +#pragma unroll + for (int l = 0; l < tile_C::ne/2; ++l) { + const int j = j0 + tile_C::get_j(l); + + dsB[l] = __half22float2(y_dm[j*MMQ_TILE_Y_K + k01/QI8_1]); + } + +#pragma unroll + for (int n = 0; n < ntx; ++n) { + tile_C C; + mma(C, A[n][k01/QI8_1], B); + +#pragma unroll + for (int l = 0; l < tile_C::ne; ++l) { + sum[(j0/tile_C::J + n)*tile_C::ne + l] += dmA[n][l/2][k01/QI8_1].x*dsB[l%2].x*C.x[l]; + sum[(j0/tile_C::J + n)*tile_C::ne + l] += dmA[n][l/2][k01/QI8_1].y*dsB[l%2].y; + } + } + } + } +#endif // defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) +} + +// Used for NVFP4, Q3_K, IQ2_S, and IQ2_XS +template static __device__ __forceinline__ void ggml_cuda_mmq_vec_dot_q8_0_16_q8_1_dp4a( + const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) { + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size; + constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8); + + constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(type, I); + const int * x_qs = (const int *) x; + const float * x_df = (const float *) x_qs + txs.qs; + const int * y_qs = (const int *) y + 4; + const float * y_df = (const float *) y; + +// #pragma unroll + for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QI8_0) { + const int k0 = k00 + k01; + +#pragma unroll + for (int j0 = 0; j0 < J; j0 += nwarps) { + const int j = j0 + threadIdx.y; + +#pragma unroll + for (int i0 = 0; i0 < I; i0 += warp_size) { + const int i = i0 + threadIdx.x; + + sum[j0/nwarps*I/warp_size + i0/warp_size] += vec_dot_q8_0_16_q8_1_impl( + &x_qs[i*(2*MMQ_TILE_NE_K + 1) + k0], + &y_qs[j*MMQ_TILE_Y_K + k01], + &x_df[i*(2*MMQ_TILE_NE_K*2/QI8_0) + i/(QI8_0/4) + k0/(QI8_0/2)], + y_df[j*MMQ_TILE_Y_K + k01/QI8_1]); + } + } + } +} + +// Used for Q3_K, IQ2_S, and IQ2_XS: +template static __device__ __forceinline__ void ggml_cuda_mmq_vec_dot_q8_0_16_q8_1_mma( + const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) { +#if defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + constexpr data_layout input_layout = get_input_data_layout(); + typedef tile<16, 4, int, input_layout> tile_A; + typedef tile<16, 4, int, input_layout> tile_B; + typedef tile<16, 16, int, DATA_LAYOUT_J_MAJOR> tile_C; + + constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, prec_src1); + constexpr int rows_per_warp = ggml_cuda_mmq_get_rows_per_warp(type, J, fallback, prec_src1); + constexpr int ntx = rows_per_warp/tile_C::I; // Number of x minitiles per warp. + + y += (threadIdx.y % ntx) * (tile_C::J*MMQ_TILE_Y_K); + + const int * x_qs = (const int *) x; + const float * x_df = (const float *) x_qs + MMQ_TILE_NE_K*2; + const int * y_qs = (const int *) y + 4; + const float * y_df = (const float *) y; + + const int i0 = (threadIdx.y / ntx) * rows_per_warp; + + for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += 4) { + const int k0 = k00 + k01; + + tile_A A[ntx]; +#pragma unroll + for (int n = 0; n < ntx; ++n) { + load_ldmatrix(A[n], x_qs + (i0 + n*tile_A::I)*sram_stride + k0, sram_stride); + } + +#pragma unroll + for (int j0 = 0; j0 < J; j0 += ntx*tile_C::J) { + tile_B B; + load_ldmatrix(B, y_qs + j0*MMQ_TILE_Y_K + k01, MMQ_TILE_Y_K); + + const int j = j0 + tile_C::get_j(0); + const float dB = y_df[j*MMQ_TILE_Y_K + k01/QI8_1]; + +#pragma unroll + for (int n = 0; n < ntx; ++n) { + tile_C C; + mma(C, A[n], B); + +#pragma unroll + for (int l = 0; l < tile_C::ne; ++l) { + const int i = i0 + n*tile_C::I + tile_C::get_i(l); + sum[(j0/tile_C::J + n)*tile_C::ne + l] += C.x[l] * x_df[i*sram_stride + k0/4] * dB; + } + } + } + } +#elif defined(TURING_MMA_AVAILABLE) + + typedef tile<16, 4, int> tile_A; + typedef tile<16, 8, int> tile_A_8; + typedef tile< 8, 4, int> tile_B; + typedef tile<16, 8, int> tile_C; + + constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, prec_src1); + constexpr int rows_per_warp = ggml_cuda_mmq_get_rows_per_warp(type, J, fallback, prec_src1); + constexpr int ntx = rows_per_warp/tile_C::I; // Number of x minitiles per warp. + + y += (threadIdx.y % ntx) * (tile_C::J*MMQ_TILE_Y_K); + + const int * x_qs = (const int *) x; + const float * x_df = (const float *) x_qs + MMQ_TILE_NE_K*2; + const int * y_qs = (const int *) y + 4; + const float * y_df = (const float *) y; + + const int i0 = (threadIdx.y / ntx) * (ntx*tile_A::I); + + tile_A A[ntx][8]; + float dA[ntx][tile_C::ne/2][8]; + +#pragma unroll + for (int n = 0; n < ntx; ++n) { +#pragma unroll + for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += 8) { + const int k0 = k00 + k01; + + load_ldmatrix(((tile_A_8 *) A[n])[k01/8], x_qs + (i0 + n*tile_A::I)*sram_stride + k0, sram_stride); + } + +#pragma unroll + for (int l = 0; l < tile_C::ne/2; ++l) { + const int i = i0 + n*tile_C::I + tile_C::get_i(2*l); + +#pragma unroll + for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += 4) { + const int k0 = k00 + k01; + + dA[n][l][k01/4] = x_df[i*sram_stride + k0/4]; + } + } + } + +#pragma unroll + for (int j0 = 0; j0 < J; j0 += ntx*tile_C::J) { +#pragma unroll + for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QR3_K*VDR_Q3_K_Q8_1_MMQ) { + tile_B B[2]; + float dB[tile_C::ne/2]; + + // Here load_generic is faster than load_ldmatrix. + load_generic(B[0], y_qs + j0*MMQ_TILE_Y_K + (k01 + 0), MMQ_TILE_Y_K); + load_generic(B[1], y_qs + j0*MMQ_TILE_Y_K + (k01 + tile_B::J), MMQ_TILE_Y_K); + +#pragma unroll + for (int l = 0; l < tile_C::ne/2; ++l) { + const int j = j0 + tile_C::get_j(l); + + dB[l] = y_df[j*MMQ_TILE_Y_K + k01/QI8_1]; + } + +#pragma unroll + for (int n = 0; n < ntx; ++n) { + tile_C C[2]; + mma(C[0], A[n][k01/4 + 0], B[0]); + mma(C[1], A[n][k01/4 + 1], B[1]); + +#pragma unroll + for (int l = 0; l < tile_C::ne; ++l) { + sum[(j0/tile_C::J + n)*tile_C::ne + l] += dB[l%2]*(C[0].x[l]*dA[n][l/2][k01/4 + 0] + C[1].x[l]*dA[n][l/2][k01/4 + 1]); + } + } + } + } +#else + GGML_UNUSED_VARS(x, y, sum, k00); + NO_DEVICE_CODE; +#endif // AMD_MFMA_AVAILABLE || AMD_WMMA_AVAILABLE +} + +template static __device__ __forceinline__ void ggml_cuda_mmq_vec_dot_q2_K_q8_1_dp4a( + const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) { + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size; + constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8); + + constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q2_K, I); + const int * x_qs = (const int *) x; + const half2 * x_dm = (const half2 *) x_qs + txs.qs; + const int * y_qs = (const int *) y + 4; + const half2 * y_ds = (const half2 *) y; + + float2 y_df[J/nwarps]; +#pragma unroll + for (int j0 = 0; j0 < J; j0 += nwarps) { + const int j = j0 + threadIdx.y; + + y_df[j0/nwarps] = __half22float2(y_ds[j*MMQ_TILE_Y_K]); + } + +#pragma unroll + for (int k01 = 0; k01 < MMQ_TILE_NE_K/2; k01 += QR2_K*VDR_Q2_K_Q8_1_MMQ) { + const int k0 = k00 + k01; + +#pragma unroll + for (int j0 = 0; j0 < J; j0 += nwarps) { + const int j = j0 + threadIdx.y; + +#pragma unroll + for (int i0 = 0; i0 < I; i0 += warp_size) { + const int i = i0 + threadIdx.x; + + constexpr int ns = 2; + sum[j0/nwarps*I/warp_size + i0/warp_size] += vec_dot_q2_K_q8_1_impl_mmq( + &x_qs[i*(2*MMQ_TILE_NE_K + 1) + k0], &y_qs[j*MMQ_TILE_Y_K + k01], + &x_dm[i*(MMQ_TILE_NE_K + 1) + k0/4], k01 < MMQ_TILE_NE_K/2 ? y_df[j0/nwarps].x : y_df[j0/nwarps].y, + &y_ds[j*MMQ_TILE_Y_K + (1 + k01/QI8_1)]); + } + } + } + + // Some compilers fail to unroll the loop over k01 if there is a conditional statement for ns in the inner loop. + // As a workaround 2 separate loops are used instead. +#pragma unroll + for (int k01 = MMQ_TILE_NE_K/2; k01 < MMQ_TILE_NE_K; k01 += QR2_K*VDR_Q2_K_Q8_1_MMQ) { + const int k0 = k00 + k01; + +#pragma unroll + for (int j0 = 0; j0 < J; j0 += nwarps) { + const int j = j0 + threadIdx.y; + +#pragma unroll + for (int i0 = 0; i0 < I; i0 += warp_size) { + const int i = i0 + threadIdx.x; + + constexpr int ns = 1; + sum[j0/nwarps*I/warp_size + i0/warp_size] += vec_dot_q2_K_q8_1_impl_mmq( + &x_qs[i*(2*MMQ_TILE_NE_K + 1) + k0], &y_qs[j*MMQ_TILE_Y_K + k01], + &x_dm[i*(MMQ_TILE_NE_K + 1) + k0/4], k01 < MMQ_TILE_NE_K/2 ? y_df[j0/nwarps].x : y_df[j0/nwarps].y, + &y_ds[j*MMQ_TILE_Y_K + (1 + k01/QI8_1)]); + } + } + } +} + +template static __device__ __forceinline__ void ggml_cuda_mmq_vec_dot_q2_K_q8_1_mma( + const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) { +#if defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + constexpr data_layout input_layout = get_input_data_layout(); + typedef tile<16, 4, int, input_layout> tile_A; + typedef tile<16, 4, int, input_layout> tile_B; + typedef tile<16, 16, int, DATA_LAYOUT_J_MAJOR> tile_C; + + constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8); + constexpr int rows_per_warp = ggml_cuda_mmq_get_rows_per_warp(type, J, fallback, GGML_PREC_Q8); + constexpr int ntx = rows_per_warp/tile_C::I; // Number of x minitiles per warp. + + y += (threadIdx.y % ntx) * (tile_C::J*MMQ_TILE_Y_K); + + const int * x_qs = (const int *) x; + const half2 * x_dm = (const half2 *) x_qs + MMQ_TILE_NE_K*2; + const int * y_qs = (const int *) y + 4; + const half2 * y_ds = (const half2 *) y; + + const int i0 = (threadIdx.y / ntx) * rows_per_warp; + + for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += 4) { + const int k0 = k00 + k01; + + tile_A A[ntx]; +#pragma unroll + for (int n = 0; n < ntx; ++n) { + load_ldmatrix(A[n], x_qs + (i0 + n*tile_A::I)*sram_stride + k0, sram_stride); + } + +#pragma unroll + for (int j0 = 0; j0 < J; j0 += ntx*tile_C::J) { + tile_B B; + load_ldmatrix(B, y_qs + j0*MMQ_TILE_Y_K + k01, MMQ_TILE_Y_K); + + const int j = j0 + tile_C::get_j(0); + const float dB = (k01 < MMQ_TILE_NE_K/2) ? __half22float2(y_ds[j*MMQ_TILE_Y_K]).x : __half22float2(y_ds[j*MMQ_TILE_Y_K]).y; + const float sB = (k01 >= MMQ_TILE_NE_K * 3/4) ? 0 + : (((k01/4)%2) ? __half22float2(y_ds[j*MMQ_TILE_Y_K + (1 + k01/QI8_1)]).y + : __half22float2(y_ds[j*MMQ_TILE_Y_K + (1 + k01/QI8_1)]).x); + + tile_C Cm; + if (k01 >= MMQ_TILE_NE_K * 3/4) { + tile_A A1; +#pragma unroll + for (int l = 0; l < tile_A::ne; ++l) { + A1.x[l] = 0x01010101; + } + mma(Cm, A1, B); + } + +#pragma unroll + for (int n = 0; n < ntx; ++n) { + tile_C Cd; + mma(Cd, A[n], B); + +#pragma unroll + for (int l = 0; l < tile_C::ne; ++l) { + const int i = i0 + n*tile_C::I + tile_C::get_i(l); + const float2 dm = __half22float2(x_dm[i*sram_stride + k0/4]); + float tmp = Cd.x[l]*dm.x; + if (k01 >= MMQ_TILE_NE_K * 3/4) { + tmp -= Cm.x[l]*dm.y; + } + sum[(j0/tile_C::J + n)*tile_C::ne + l] += tmp*dB; + sum[(j0/tile_C::J + n)*tile_C::ne + l] -= dm.y*sB; + } + } + } + } +#elif defined(TURING_MMA_AVAILABLE) + + typedef tile<16, 4, int> tile_A; + typedef tile<16, 8, int> tile_A_8; + typedef tile< 8, 4, int> tile_B; + typedef tile<16, 8, int> tile_C; + + constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8); + constexpr int rows_per_warp = ggml_cuda_mmq_get_rows_per_warp(type, J, fallback, GGML_PREC_Q8); + constexpr int ntx = rows_per_warp/tile_C::I; // Number of x minitiles per warp. + + y += (threadIdx.y % ntx) * (tile_C::J*MMQ_TILE_Y_K); + + const int * x_qs = (const int *) x; + const half2 * x_dm = (const half2 *) x_qs + MMQ_TILE_NE_K*2; + const int * y_qs = (const int *) y + 4; + const half2 * y_ds = (const half2 *) y; + + const int i0 = (threadIdx.y / ntx) * (ntx*tile_A::I); + + tile_A A[ntx][8]; + float dA[ntx][tile_C::ne/2][8]; + float mA[ntx][tile_C::ne/2][8]; + +#pragma unroll + for (int n = 0; n < ntx; ++n) { +#pragma unroll + for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QI8_1) { + const int k0 = k00 + k01; + + load_ldmatrix(((tile_A_8 *) A[n])[k01/QI8_1], x_qs + (i0 + n*tile_A::I)*sram_stride + k0, sram_stride); + } + } + +#pragma unroll + for (int n = 0; n < ntx; ++n) { +#pragma unroll + for (int l = 0; l < tile_C::ne/2; ++l) { + const int i = i0 + n*tile_C::I + tile_C::get_i(2*l); + +#pragma unroll + for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QI8_1/2) { + const int k0 = k00 + k01; + + const float2 dm = __half22float2(x_dm[i*sram_stride + k0/(QI8_1/2)]); + + dA[n][l][k01/(QI8_1/2)] = dm.x; + mA[n][l][k01/(QI8_1/2)] = dm.y; + } + } + } + +#pragma unroll + for (int j0 = 0; j0 < J; j0 += ntx*tile_C::J) { + float2 dB[tile_C::ne/2]; + +#pragma unroll + for (int l = 0; l < tile_C::ne/2; ++l) { + const int j = j0 + tile_C::get_j(l); + + dB[l] = __half22float2(y_ds[j*MMQ_TILE_Y_K]); + } + +#pragma unroll + for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QI8_1) { + tile_B B[2]; + + // Here load_generic is faster than load_ldmatrix. + load_generic(B[0], y_qs + j0*MMQ_TILE_Y_K + (k01 + 0), MMQ_TILE_Y_K); + load_generic(B[1], y_qs + j0*MMQ_TILE_Y_K + (k01 + tile_B::J), MMQ_TILE_Y_K); + + tile_C Cm[2]; + if (k01 >= MMQ_TILE_NE_K * 3/4) { + tile_A A1; + A1.x[0] = 0x01010101; + A1.x[1] = 0x01010101; + mma(Cm[0], A1, B[0]); + mma(Cm[1], A1, B[1]); + } + +#pragma unroll + for (int n = 0; n < ntx; ++n) { + tile_C Cd[2]; + + mma(Cd[0], A[n][k01/4 + 0], B[0]); + mma(Cd[1], A[n][k01/4 + 1], B[1]); + +#pragma unroll + for (int l = 0; l < tile_C::ne; ++l) { + float tmp = Cd[0].x[l]*dA[n][l/2][k01/4 + 0] + Cd[1].x[l]*dA[n][l/2][k01/4 + 1]; + if (k01 >= MMQ_TILE_NE_K * 3/4) { + tmp -= Cm[0].x[l]*mA[n][l/2][k01/4 + 0] + Cm[1].x[l]*mA[n][l/2][k01/4 + 1]; + } + sum[(j0/tile_C::J + n)*tile_C::ne + l] += tmp*(k01 < MMQ_TILE_NE_K/2 ? dB[l%2].x : dB[l%2].y); + } + } + } + +#pragma unroll + for (int k01 = 0; k01 < MMQ_TILE_NE_K * 3/4; k01 += QI8_1) { + float2 sB[tile_C::ne/2]; + +#pragma unroll + for (int l = 0; l < tile_C::ne/2; ++l) { + const int j = j0 + tile_C::get_j(l); + + sB[l] = __half22float2(y_ds[j*MMQ_TILE_Y_K + (1 + k01/QI8_1)]); + } + +#pragma unroll + for (int n = 0; n < ntx; ++n) { +#pragma unroll + for (int l = 0; l < tile_C::ne; ++l) { + sum[(j0/tile_C::J + n)*tile_C::ne + l] -= mA[n][l/2][k01/4 + 0]*sB[l%2].x; + sum[(j0/tile_C::J + n)*tile_C::ne + l] -= mA[n][l/2][k01/4 + 1]*sB[l%2].y; + } + } + } + } +#else + GGML_UNUSED_VARS(x, y, sum, k00); + NO_DEVICE_CODE; +#endif // AMD_MFMA_AVAILABLE || AMD_WMMA_AVAILABLE +} + +template static __device__ __forceinline__ void ggml_cuda_mmq_vec_dot_q3_K_q8_1_dp4a( + const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) { + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size; + constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8); + + constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q3_K, I); + const int * x_qs = (const int *) x; + const float * x_df = (const float *) x_qs + txs.qs; + const int * x_sc = (const int *) x_df + txs.dm; + const int * y_qs = (const int *) y + 4; + const float * y_df = (const float *) y; + +// #pragma unroll + for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QR3_K*VDR_Q3_K_Q8_1_MMQ) { + const int k0 = k00 + k01; + +#pragma unroll + for (int j0 = 0; j0 < J; j0 += nwarps) { + const int j = j0 + threadIdx.y; + +#pragma unroll + for (int i0 = 0; i0 < I; i0 += warp_size) { + const int i = i0 + threadIdx.x; + + const int8_t * scales = ((const int8_t *) (x_sc + i*(MMQ_TILE_NE_K/8) + i/8)) + k0/4; + + sum[j0/nwarps*I/warp_size + i0/warp_size] += vec_dot_q3_K_q8_1_impl_mmq( + &x_qs[i*(2*MMQ_TILE_NE_K + 1) + k0], &y_qs[j*MMQ_TILE_Y_K + k01], scales, + x_df[i], y_df[j*MMQ_TILE_Y_K + k01/QI8_1]); + } + } + } +} + +template static __device__ __forceinline__ void ggml_cuda_mmq_vec_dot_q4_K_q8_1_dp4a( + const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) { + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size; + constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8); + + constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q4_K, I); + const int * x_qs = (const int *) x; + const half2 * x_dm = (const half2 *) x_qs + txs.qs; + const int * x_sc = (const int *) x_dm + txs.dm; + const int * y_qs = (const int *) y + 4; + const half2 * y_ds = (const half2 *) y; + +// #pragma unroll + for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QR4_K*VDR_Q4_K_Q8_1_MMQ) { + const int k0 = k00 + k01; + +#pragma unroll + for (int j0 = 0; j0 < J; j0 += nwarps) { + const int j = j0 + threadIdx.y; + +#pragma unroll + for (int i0 = 0; i0 < I; i0 += warp_size) { + const int i = i0 + threadIdx.x; + + const uint8_t * sc = (const uint8_t *) &x_sc[i * (MMQ_TILE_NE_K/8) + i/8 + k0/32] + 2*(k01/16); + + sum[j0/nwarps*I/warp_size + i0/warp_size] += vec_dot_q4_K_q8_1_impl_mmq( + &x_qs[i*(MMQ_TILE_NE_K + 1) + k0/2], &y_qs[j*MMQ_TILE_Y_K + k01], sc, sc+8, + x_dm[i], &y_ds[j*MMQ_TILE_Y_K + k01/QI8_1]); + } + } + } +} + +template static __device__ __forceinline__ void ggml_cuda_mmq_vec_dot_q5_K_q8_1_dp4a( + const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) { + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size; + constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8); + + constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q5_K, I); + const int * x_qs = (const int *) x; + const half2 * x_dm = (const half2 *) x_qs + txs.qs; + const int * x_sc = (const int *) x_dm + txs.dm; + const int * y_qs = (const int *) y + 4; + const half2 * y_ds = (const half2 *) y; + +// #pragma unroll + for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QR5_K*VDR_Q5_K_Q8_1_MMQ) { + const int k0 = k00 + k01; + +#pragma unroll + for (int j0 = 0; j0 < J; j0 += nwarps) { + const int j = j0 + threadIdx.y; + +#pragma unroll + for (int i0 = 0; i0 < I; i0 += warp_size) { + const int i = i0 + threadIdx.x; + + const uint8_t * sc = ((const uint8_t *) &x_sc[i * (MMQ_TILE_NE_K/8) + i/8 + k00/32]) + 2*(k01/16); + + sum[j0/nwarps*I/warp_size + i0/warp_size] += vec_dot_q5_K_q8_1_impl_mmq( + &x_qs[i*(QR5_K*MMQ_TILE_NE_K + 1) + k0], &y_qs[j*MMQ_TILE_Y_K + k01], sc, sc+8, + x_dm[i], &y_ds[j*MMQ_TILE_Y_K + k01/QI8_1]); + } + } + } +} + +template static __device__ __forceinline__ void ggml_cuda_mmq_vec_dot_q6_K_q8_1_dp4a( + const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) { + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size; + constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8); + + constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q6_K, I); + const int * x_qs = (const int *) x; + const float * x_df = (const float *) x_qs + txs.qs; + const int * x_sc = (const int *) x_df + txs.dm; + const int * y_qs = (const int *) y + 4; + const float * y_df = (const float *) y; + +// #pragma unroll + for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QR6_K*VDR_Q6_K_Q8_1_MMQ) { + const int k0 = k00 + k01; + +#pragma unroll + for (int j0 = 0; j0 < J; j0 += nwarps) { + const int j = j0 + threadIdx.y; + +#pragma unroll + for (int i0 = 0; i0 < I; i0 += warp_size) { + const int i = i0 + threadIdx.x; + + const int8_t * sc = ((const int8_t *) &x_sc[i * (MMQ_TILE_NE_K/8) + i/8 + k0/16]); + + sum[j0/nwarps*I/warp_size + i0/warp_size] += vec_dot_q6_K_q8_1_impl_mmq( + &x_qs[i*(QR6_K*MMQ_TILE_NE_K + 1) + k0], &y_qs[j*MMQ_TILE_Y_K + k01], sc, + x_df[i*(MMQ_TILE_NE_K/QI6_K) + i/QI6_K], &y_df[j*MMQ_TILE_Y_K + k01/QI8_1]); + } + } + } +} + +template static __device__ __forceinline__ void ggml_cuda_mmq_vec_dot_q6_K_q8_1_mma( + const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) { +#if defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + constexpr data_layout input_layout = get_input_data_layout(); + typedef tile<16, 4, int, input_layout> tile_A; + typedef tile<16, 4, int, input_layout> tile_B; + typedef tile<16, 16, int, DATA_LAYOUT_J_MAJOR> tile_C; + + constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8); + constexpr int rows_per_warp = ggml_cuda_mmq_get_rows_per_warp(type, J, fallback, GGML_PREC_Q8); + constexpr int ntx = rows_per_warp/tile_C::I; // Number of x minitiles per warp. + + y += (threadIdx.y % ntx) * (tile_C::J*MMQ_TILE_Y_K); + + const int * x_qs = (const int *) x; + const float * x_df = (const float *) x_qs + MMQ_TILE_NE_K*2; + const int * x_sc = (const int *) x_df + MMQ_TILE_NE_K/QI6_K; + const int * y_qs = (const int *) y + 4; + const float * y_df = (const float *) y; + + const int i0 = (threadIdx.y / ntx) * rows_per_warp; + + for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += 4) { + const int k0 = k00 + k01; + + tile_A A[ntx]; +#pragma unroll + for (int n = 0; n < ntx; ++n) { + load_ldmatrix(A[n], x_qs + (i0 + n*tile_A::I)*sram_stride + k0, sram_stride); + } + +#pragma unroll + for (int j0 = 0; j0 < J; j0 += ntx*tile_C::J) { + tile_B B; + load_ldmatrix(B, y_qs + j0*MMQ_TILE_Y_K + k01, MMQ_TILE_Y_K); + + const int j = j0 + tile_C::get_j(0); + const float dB = y_df[j*MMQ_TILE_Y_K + k01/QI8_1]; + +#pragma unroll + for (int n = 0; n < ntx; ++n) { + tile_C C; + mma(C, A[n], B); + +#pragma unroll + for (int l = 0; l < tile_C::ne; ++l) { + const int i = i0 + n*tile_C::I + tile_C::get_i(l); + const int8_t * sc = (const int8_t *) (x_sc + i*sram_stride + k00/16); + sum[(j0/tile_C::J + n)*tile_C::ne + l] += C.x[l] * sc[k01/4] * x_df[i*sram_stride] * dB; + } + } + } + } +#elif defined(TURING_MMA_AVAILABLE) + + typedef tile<16, 4, int> tile_A; + typedef tile< 8, 4, int> tile_B; + typedef tile<16, 8, int> tile_C; + + constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8); + constexpr int rows_per_warp = ggml_cuda_mmq_get_rows_per_warp(type, J, fallback, GGML_PREC_Q8); + constexpr int ntx = rows_per_warp/tile_C::I; // Number of x minitiles per warp. + + y += (threadIdx.y % ntx) * (tile_C::J*MMQ_TILE_Y_K); + + const int * x_qs = (const int *) x; + const float * x_df = (const float *) x_qs + MMQ_TILE_NE_K*2; + const int * x_sc = (const int *) x_df + MMQ_TILE_NE_K/QI6_K; + const int * y_qs = (const int *) y + 4; + const float * y_df = (const float *) y; + + const int i0 = (threadIdx.y / ntx) * (ntx*tile_A::I); + + tile_A A[ntx][8]; + int scA[ntx][tile_C::ne/2][8]; + float dA[ntx][tile_C::ne/2]; + +#pragma unroll + for (int n = 0; n < ntx; ++n) { +#pragma unroll + for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += 8) { + const int k0 = k00 + k01; + + load_ldmatrix(A[n][k01/4 + 0], x_qs + (i0 + n*tile_A::I)*sram_stride + (k0 + 0), sram_stride); + load_ldmatrix(A[n][k01/4 + 1], x_qs + (i0 + n*tile_A::I)*sram_stride + (k0 + tile_A::J), sram_stride); + } + +#pragma unroll + for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += 16) { + const int k0 = k00 + k01; + +#pragma unroll + for (int l = 0; l < tile_C::ne/2; ++l) { + const int i = i0 + n*tile_C::I + tile_C::get_i(2*l); + + const int sc_packed = x_sc[i*sram_stride + k0/16]; + const int8_t * sc = (const int8_t *) &sc_packed; + +#pragma unroll + for (int ksc = 0; ksc < sizeof(int); ++ksc) { + scA[n][l][k01/4 + ksc] = sc[ksc]; + } + } + } + +#pragma unroll + for (int l = 0; l < tile_C::ne/2; ++l) { + const int i = i0 + n*tile_C::I + tile_C::get_i(2*l); + + dA[n][l] = x_df[i*sram_stride]; + } + } + +#pragma unroll + for (int j0 = 0; j0 < J; j0 += ntx*tile_C::J) { + float tmp[ntx][tile_C::ne] = {{0.0f}}; + +#pragma unroll + for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += 8) { + tile_B B[2]; + float dB[tile_C::ne/2]; + + // Here load_generic is faster than load_ldmatrix. + load_generic(B[0], y_qs + j0*MMQ_TILE_Y_K + 0 + k01, MMQ_TILE_Y_K); + load_generic(B[1], y_qs + j0*MMQ_TILE_Y_K + tile_B::J + k01, MMQ_TILE_Y_K); + +#pragma unroll + for (int l = 0; l < tile_C::ne/2; ++l) { + const int j = j0 + tile_C::get_j(l); + + dB[l] = y_df[j*MMQ_TILE_Y_K + k01/QI8_1]; + } + +#pragma unroll + for (int n = 0; n < ntx; ++n) { + tile_C C[2]; + mma(C[0], A[n][k01/4 + 0], B[0]); + mma(C[1], A[n][k01/4 + 1], B[1]); + +#pragma unroll + for (int l = 0; l < tile_C::ne; ++l) { + tmp[n][l] += (C[0].x[l]*scA[n][l/2][k01/4 + 0] + C[1].x[l]*scA[n][l/2][k01/4 + 1])*dB[l%2]; + } + } + } + +#pragma unroll + for (int n = 0; n < ntx; ++n) { +#pragma unroll + for (int l = 0; l < tile_C::ne; ++l) { + sum[(j0/tile_C::J + n)*tile_C::ne + l] += tmp[n][l]*dA[n][l/2]; + } + } + } +#else + GGML_UNUSED_VARS(x, y, sum, k00); + NO_DEVICE_CODE; +#endif // AMD_MFMA_AVAILABLE || AMD_WMMA_AVAILABLE +} + +// --------------------------------------------------------------------------------------------- + +// Shared MMA kernel for MXFP4 and NVFP4 on Blackwell. +// Both quantizations encode values as e2m1 (FP4) and produce one uint32 scale per +// m16n8k64 MMA call; only the PTX kind (scale_vec::2X ue8m0 vs scale_vec::4X ue4m3) +// and the per-type stride constant differ. +template static __device__ __forceinline__ void ggml_cuda_mmq_vec_dot_fp4_fp4_mma( + const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) { + + typedef tile<16, 8, int> tile_A; + typedef tile<8, 8, int> tile_B; + typedef tile<16, 8, float> tile_C; + + constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q4); + constexpr int rows_per_warp = ggml_cuda_mmq_get_rows_per_warp(type, J, fallback, GGML_PREC_Q4); + constexpr int ntx = rows_per_warp / tile_C::I; + constexpr int nfrags = MMQ_TILE_NE_K / tile_A::J; + + y += (threadIdx.y % ntx) * (tile_C::J * MMQ_TILE_Y_K); + + const int * x_qs = (const int *) x; + const uint32_t * x_sc = (const uint32_t *) (x_qs + 2 * MMQ_TILE_NE_K); + const int * y_qs = (const int *) y + 4; + const uint32_t * y_sc = (const uint32_t *) y; + + // 2 threads per quad supply the packed scale register to the block_scale MMA, + // see https://docs.nvidia.com/cuda/parallel-thread-execution/#warp-level-block-scaling + const int tidx_A = threadIdx.x / 4 + (threadIdx.x % 2) * 8; + const int tidx_B = threadIdx.x / 4; + const int i0 = (threadIdx.y / ntx) * rows_per_warp; + + tile_A A[ntx][nfrags]; + uint32_t scaleA[ntx][nfrags]; + +#pragma unroll + for (int n = 0; n < ntx; ++n) { +#pragma unroll + for (int frag = 0; frag < nfrags; ++frag) { + const int k0 = k00 + frag * tile_A::J; + load_ldmatrix(A[n][frag], x_qs + (i0 + n * tile_A::I) * sram_stride + k0, sram_stride); + scaleA[n][frag] = x_sc[(i0 + n * tile_A::I + tidx_A) * sram_stride + k0 / tile_A::J]; + } + } + +#pragma unroll + for (int j0 = 0; j0 < J; j0 += ntx * tile_C::J) { + tile_B B[nfrags]; + uint32_t scaleB[nfrags]; + +#pragma unroll + for (int frag = 0; frag < nfrags; ++frag) { + const int k0 = frag * tile_B::J; + load_generic(B[frag], y_qs + j0 * MMQ_TILE_Y_K + k0, MMQ_TILE_Y_K); + scaleB[frag] = y_sc[(j0 + tidx_B) * MMQ_TILE_Y_K + frag]; + } + +#pragma unroll + for (int n = 0; n < ntx; ++n) { + // accumulate in place into the output sum array + tile_C & C = *reinterpret_cast(sum + (j0 / tile_C::J + n) * tile_C::ne); +#pragma unroll + for (int frag = 0; frag < nfrags; ++frag) { + mma_block_scaled_fp4(C, A[n][frag], B[frag], scaleA[n][frag], scaleB[frag]); + } + } + } +} + diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq.cu b/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq.cu new file mode 100644 index 0000000000..f754837dd3 --- /dev/null +++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq.cu @@ -0,0 +1,522 @@ +#include "common.cuh" +#include "mmq.cuh" +#include "quantize.cuh" +#include "mmid.cuh" + +#include + +static void ggml_cuda_mul_mat_q_switch_type(ggml_backend_cuda_context & ctx, const mmq_args & args, cudaStream_t stream, const ggml_prec prec_src1) { + switch (args.type_x) { + case GGML_TYPE_Q1_0: + mul_mat_q_case(ctx, args, stream); + break; + case GGML_TYPE_Q2_0: + mul_mat_q_case(ctx, args, stream); + break; + case GGML_TYPE_Q4_0: + mul_mat_q_case(ctx, args, stream); + break; + case GGML_TYPE_Q4_1: + mul_mat_q_case(ctx, args, stream); + break; + case GGML_TYPE_Q5_0: + mul_mat_q_case(ctx, args, stream); + break; + case GGML_TYPE_Q5_1: + mul_mat_q_case(ctx, args, stream); + break; + case GGML_TYPE_Q8_0: + mul_mat_q_case(ctx, args, stream); + break; +// ----------------------------------------------------------------------- + case GGML_TYPE_Q2_K: + mul_mat_q_case(ctx, args, stream); + break; + case GGML_TYPE_Q3_K: + mul_mat_q_case(ctx, args, stream); + break; + case GGML_TYPE_Q4_K: + mul_mat_q_case(ctx, args, stream); + break; + case GGML_TYPE_Q5_K: + mul_mat_q_case(ctx, args, stream); + break; + case GGML_TYPE_Q6_K: + mul_mat_q_case(ctx, args, stream); + break; +// ----------------------------------------------------------------------- + case GGML_TYPE_IQ1_S: + mul_mat_q_case(ctx, args, stream); + break; + case GGML_TYPE_IQ2_XXS: + mul_mat_q_case(ctx, args, stream); + break; + case GGML_TYPE_IQ2_XS: + mul_mat_q_case(ctx, args, stream); + break; + case GGML_TYPE_IQ2_S: + mul_mat_q_case(ctx, args, stream); + break; + case GGML_TYPE_IQ3_XXS: + mul_mat_q_case(ctx, args, stream); + break; + case GGML_TYPE_IQ3_S: + mul_mat_q_case(ctx, args, stream); + break; + case GGML_TYPE_IQ4_XS: + mul_mat_q_case(ctx, args, stream); + break; + case GGML_TYPE_IQ4_NL: + mul_mat_q_case(ctx, args, stream); + break; +// ----------------------------------------------------------------------- + case GGML_TYPE_MXFP4: + // src1 at Q4 uses the native FP4 instructions, which are Blackwell-only + if (prec_src1 == GGML_PREC_Q4) { + mul_mat_q_case(ctx, args, stream); + break; + } + mul_mat_q_case(ctx, args, stream); + break; + case GGML_TYPE_NVFP4: + if (prec_src1 == GGML_PREC_Q4) { + mul_mat_q_case(ctx, args, stream); + break; + } + mul_mat_q_case(ctx, args, stream); + break; + default: + GGML_ABORT("fatal error"); + } +} + +// overrides the src1 precision requested by the graph, "auto" keeps the requested one +static ggml_prec ggml_cuda_mmq_get_prec_env() { + const char * env_c = getenv("GGML_CUDA_MMQ_PREC"); + if (env_c == nullptr) { + return GGML_PREC_UNDEFINED; + } + std::string env_cpp = env_c; + for (char & c : env_cpp) { + c = std::tolower(c); + } + if (env_cpp == "q4") { + return GGML_PREC_Q4; + } + if (env_cpp == "q8") { + return GGML_PREC_Q8; + } + if (env_cpp != "auto") { + GGML_LOG_WARN("%s: Unknown value for GGML_CUDA_MMQ_PREC: '%s'. Available: 'q4', 'q8', 'auto'.\n", __func__, env_cpp.c_str()); + } + return GGML_PREC_UNDEFINED; +} + +// src1 is quantized to Q8_1 unless the FP4 types can use 4-bit activations, in which case they +// default to the native W4A4 instructions on Blackwell. +static ggml_prec ggml_cuda_mmq_get_prec_src1(const ggml_tensor * src0, const ggml_tensor * dst, const int cc) { + static const ggml_prec prec_env = ggml_cuda_mmq_get_prec_env(); + + ggml_prec prec = prec_env; + if (prec == GGML_PREC_UNDEFINED) { + prec = (ggml_prec) ggml_get_op_params_i32(dst, 3); + } + + // Q4 only for the FP4 types on Blackwell + GGML_ASSERT(prec == GGML_PREC_UNDEFINED || prec == GGML_PREC_Q8 || prec == GGML_PREC_Q4); + const bool can_use_q4 = (src0->type == GGML_TYPE_NVFP4 || src0->type == GGML_TYPE_MXFP4) && blackwell_mma_available(cc); + if (prec == GGML_PREC_Q8 || !can_use_q4) { + return GGML_PREC_Q8; + } + return GGML_PREC_Q4; +} + +void ggml_cuda_mul_mat_q( + ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, ggml_tensor * dst, + int64_t ncols_max) { + // FreeToken: bf16 and fp16 activations are quantized directly, except for the native fp4 paths; dst is fp32, bf16 or fp16 + GGML_ASSERT( src1->type == GGML_TYPE_F32 || src1->type == GGML_TYPE_BF16 || src1->type == GGML_TYPE_F16); + GGML_ASSERT( dst->type == GGML_TYPE_F32 || dst->type == GGML_TYPE_BF16 || dst->type == GGML_TYPE_F16); + GGML_ASSERT(!ids || ids->type == GGML_TYPE_I32); // Optional, used for batched GGML_MUL_MAT_ID. + + GGML_TENSOR_BINARY_OP_LOCALS; + + cudaStream_t stream = ctx.stream(); + + const int id = ggml_cuda_get_device(); + const int cc = ggml_cuda_info().devices[id].cc; + const size_t smpbo = ggml_cuda_info().devices[id].smpbo; + + const size_t ts_src0 = ggml_type_size(src0->type); + const size_t ts_src1 = ggml_type_size(src1->type); + const size_t ts_dst = ggml_type_size(dst->type); + + GGML_ASSERT( nb00 == ts_src0); + GGML_ASSERT( nb10 == ts_src1); + GGML_ASSERT( nb0 == ts_dst); + GGML_ASSERT(!ids || ids->nb[0] == ggml_type_size(ids->type)); + + const char * src0_d = (const char *) src0->data; + const float * src1_d = (const float *) src1->data; + const mmq_dst dst_d = {dst->data, dst->type}; + const nv_bfloat16 * src1_bf16 = src1->type == GGML_TYPE_BF16 ? (const nv_bfloat16 *) src1->data : nullptr; + const half * src1_f16 = src1->type == GGML_TYPE_F16 ? (const half *) src1->data : nullptr; + + // If src0 is a temporary compute buffer, clear any potential padding. + if (ggml_backend_buffer_get_usage(src0->buffer) == GGML_BACKEND_BUFFER_USAGE_COMPUTE) { + const size_t size_data = ggml_nbytes(src0); + const size_t size_alloc = ggml_backend_buffer_get_alloc_size(src0->buffer, src0); + if (size_alloc > size_data) { + GGML_ASSERT(ggml_is_contiguously_allocated(src0)); + GGML_ASSERT(!src0->view_src); + CUDA_CHECK(cudaMemsetAsync((char *) src0->data + size_data, 0, size_alloc - size_data, stream)); + } + } + + const int64_t ne10_padded = GGML_PAD(ne10, MATRIX_ROW_PADDING); + + const int64_t s01 = src0->nb[1] / ts_src0; + const int64_t s1 = dst->nb[1] / ts_dst; + const int64_t s02 = src0->nb[2] / ts_src0; + const int64_t s2 = dst->nb[2] / ts_dst; + const int64_t s03 = src0->nb[3] / ts_src0; + const int64_t s3 = dst->nb[3] / ts_dst; + + const bool fallback = ggml_cuda_mmq_needs_fallback(ne01); + + const ggml_prec prec_src1 = ggml_cuda_mmq_get_prec_src1(src0, dst, cc); + + const bool use_native_fp4 = prec_src1 == GGML_PREC_Q4; + GGML_ASSERT(!use_native_fp4 || src1->type == GGML_TYPE_F32); + const size_t y_block_size = use_native_fp4 ? sizeof(block_fp4_mmq) : sizeof(block_q8_1_mmq); + const size_t y_values_per_block = use_native_fp4 ? QK_FP4_MMQ : QK8_1_MMQ; + + int J_best = 0; + int nthreads_best = 0; + { + int64_t ncols_opt = ne11; + if (ids) { + const int64_t n_expert_used = ids->ne[0]; + // a flat grid runs only the tiles that hold tokens, so it sizes them for the average expert + ncols_opt = ncols_max > 0 ? ncols_max : (ne12*n_expert_used + ne02 - 1) / ne02; + + // Each expert only sees ne12*n_expert_used/ne02 tokens on average. + // On RDNA3 and RDNA4 it is faster to pick the tile size against this value instead of ne12. + if (GGML_CUDA_CC_IS_RDNA3(cc) || GGML_CUDA_CC_IS_RDNA4(cc)) { + ncols_opt = (ne12*n_expert_used + ne02 - 1) / ne02; + } + } + + int ntiles_J_best = INT_MAX; + + for (int J = 8; J <= 128 && ntiles_J_best > 1; J += 8) { + const ggml_cuda_mmq_config config = ggml_cuda_mmq_get_config(src0->type, J, fallback, cc, prec_src1); + if (config.type == GGML_TYPE_COUNT) { + continue; + } + + if (mmq_get_nbytes_shared(config, cc) > smpbo) { + continue; + } + + const int ntiles_x = (ncols_opt + config.J - 1) / config.J; + + if (ntiles_x < ntiles_J_best) { + J_best = J; + nthreads_best = config.nthreads; + ntiles_J_best = ntiles_x; + } + } + } + GGML_ASSERT(J_best > 0); + + // A tile of size J can read in at most J - 1 extra columns. + // For simplicity, round up the padding of a full tile to a multiple of the number of bytes that nthreads can load in parallel. + const size_t src1_load_chunk_size = nthreads_best * sizeof(int); + const size_t src1_q8_1_padding = ((J_best * sizeof(block_q8_1_mmq) + src1_load_chunk_size - 1) / src1_load_chunk_size) + * src1_load_chunk_size; + + if (!ids) { + const size_t nbytes_src1_q8_1 = ne13*ne12 * ne11*ne10_padded * y_block_size/y_values_per_block + src1_q8_1_padding; + ggml_cuda_pool_alloc src1_q8_1(ctx.pool(), nbytes_src1_q8_1); + ggml_cuda_pool_alloc src1_scale(ctx.pool()); + if (src0->type == GGML_TYPE_NVFP4 && use_native_fp4) { + src1_scale.alloc(ne13*ne12*ne11); + } + + { + const int64_t s11 = src1->nb[1] / ts_src1; + const int64_t s12 = src1->nb[2] / ts_src1; + const int64_t s13 = src1->nb[3] / ts_src1; + if (use_native_fp4) { + static constexpr size_t align_float8 = 32; + const bool use_aligned_float8 = ggml_cuda_is_aligned(src1, align_float8); + static_assert(sizeof(block_fp4_mmq) == 4 * sizeof(block_q8_1)); + quantize_mmq_fp4_cuda(src1_d, nullptr, src1_q8_1.get(), src1_scale.ptr, src0->type, use_aligned_float8, ne10, s11, s12, s13, ne10_padded, + ne11, ne12, ne13, stream); + + } else if (src1_bf16) { + quantize_mmq_q8_1_cuda(src1_bf16, nullptr, src1_q8_1.get(), src0->type, ne10, s11, s12, s13, ne10_padded, + ne11, ne12, ne13, stream); + } else if (src1_f16) { + quantize_mmq_q8_1_cuda(src1_f16, nullptr, src1_q8_1.get(), src0->type, ne10, s11, s12, s13, ne10_padded, + ne11, ne12, ne13, stream); + } else { + quantize_mmq_q8_1_cuda(src1_d, nullptr, src1_q8_1.get(), src0->type, ne10, s11, s12, s13, ne10_padded, + ne11, ne12, ne13, stream); + } + CUDA_CHECK(cudaGetLastError()); + } + + // Stride depends on quantization format + const int64_t s12 = use_native_fp4 ? + ne11 * ne10_padded * sizeof(block_fp4_mmq) / (QK_FP4_MMQ * sizeof(int)) : + ne11 * ne10_padded * sizeof(block_q8_1) / (QK8_1 * sizeof(int)); + const int64_t s13 = ne12*s12; + + const mmq_args args = { + src0_d, src0->type, (const int *) src1_q8_1.ptr, nullptr, nullptr, dst_d, + src0->type == GGML_TYPE_NVFP4 && use_native_fp4 ? src1_scale.ptr : nullptr, + ne00, ne01, ne1, s01, ne11, s1, + ne02, ne12, s02, s12, s2, + ne03, ne13, s03, s13, s3, + ne1, J_best}; + ggml_cuda_mul_mat_q_switch_type(ctx, args, stream, prec_src1); + return; + } + + GGML_ASSERT(ne13 == 1); + GGML_ASSERT(nb12 % nb11 == 0); + GGML_ASSERT(nb2 % nb1 == 0); + + const int64_t n_expert_used = ids->ne[0]; + const int64_t ne_get_rows = ne12 * n_expert_used; + GGML_ASSERT(ne1 == n_expert_used); + + ggml_cuda_pool_alloc ids_src1(ctx.pool(), ne_get_rows); + ggml_cuda_pool_alloc ids_dst(ctx.pool(), ne_get_rows + J_best-1); // Needs to be padded for unconditional memory access. + ggml_cuda_pool_alloc expert_bounds(ctx.pool(), ne02 + 1); + + // gate/up activations are broadcast across experts (ne11 == 1): quantize each token once and + // scatter to its slots. ids_src1 then holds the inverse map (token slot -> compact row). + const bool dedup_bcast = ne11 == 1 && n_expert_used > 1; + + { + GGML_ASSERT(ids->nb[0] == ggml_element_size(ids)); + const int si1 = ids->nb[1] / ggml_element_size(ids); + const int sis1 = nb12 / nb11; + + ggml_cuda_launch_mm_ids_helper((const int32_t *) ids->data, ids_src1.get(), ids_dst.get(), expert_bounds.get(), + ne02, ne12, n_expert_used, ne11, si1, sis1, /*write_inverse =*/ dedup_bcast, stream); + CUDA_CHECK(cudaGetLastError()); + } + + // Every expert's column tiles fit in ntx_flat grid columns, so a flat grid needs no host sync to size it. + ggml_cuda_pool_alloc tile_map(ctx.pool()); + int ntx_flat = 0; + if (ncols_max <= 0) { + const int jt_none = (ne_get_rows + J_best - 1) / J_best; + ntx_flat = jt_none + std::min(ne02, ne_get_rows); + tile_map.alloc(ntx_flat); + ggml_cuda_launch_mm_ids_tile_map(expert_bounds.get(), tile_map.get(), ne02, J_best, ntx_flat, jt_none, stream); + CUDA_CHECK(cudaGetLastError()); + } + + const size_t nbytes_src1_q8_1 = ne12*n_expert_used*ne10_padded * y_block_size/y_values_per_block + src1_q8_1_padding; + ggml_cuda_pool_alloc src1_q8_1(ctx.pool(), nbytes_src1_q8_1); + ggml_cuda_pool_alloc src1_scale(ctx.pool()); + if (src0->type == GGML_TYPE_NVFP4 && use_native_fp4) { + src1_scale.alloc(ne12*n_expert_used); + } + + const int64_t ne11_flat = ne12*n_expert_used; + const int64_t ne12_flat = 1; + const int64_t ne13_flat = 1; + + { + const int64_t s11 = src1->nb[1] / ts_src1; + const int64_t s12 = src1->nb[2] / ts_src1; + const int64_t s13 = src1->nb[3] / ts_src1; + + if (use_native_fp4) { + static constexpr size_t align_float8 = 32; + const bool use_aligned_float8 = ggml_cuda_is_aligned(src1, align_float8); + if (dedup_bcast) { + quantize_scatter_mmq_fp4_cuda(src1_d, ids_src1.get(), src1_q8_1.get(), src1_scale.ptr, src0->type, use_aligned_float8, ne10, + /*stride_token=*/s12, ne10_padded, ne12, ne11_flat, n_expert_used, stream); + } else { + quantize_mmq_fp4_cuda(src1_d, ids_src1.get(), src1_q8_1.get(), src1_scale.ptr, src0->type, use_aligned_float8, ne10, s11, s12, s13, + ne10_padded, ne11_flat, ne12_flat, ne13_flat, stream); + } + } else if (dedup_bcast && src1_bf16) { + quantize_scatter_mmq_q8_1_cuda(src1_bf16, ids_src1.get(), src1_q8_1.get(), src0->type, ne10, + /*stride_token=*/s12, ne10_padded, ne12, ne11_flat, n_expert_used, stream); + } else if (dedup_bcast && src1_f16) { + quantize_scatter_mmq_q8_1_cuda(src1_f16, ids_src1.get(), src1_q8_1.get(), src0->type, ne10, + /*stride_token=*/s12, ne10_padded, ne12, ne11_flat, n_expert_used, stream); + } else if (dedup_bcast) { + quantize_scatter_mmq_q8_1_cuda(src1_d, ids_src1.get(), src1_q8_1.get(), src0->type, ne10, + /*stride_token=*/s12, ne10_padded, ne12, ne11_flat, n_expert_used, stream); + } else if (src1_bf16) { + quantize_mmq_q8_1_cuda(src1_bf16, ids_src1.get(), src1_q8_1.get(), src0->type, ne10, s11, s12, s13, + ne10_padded, ne11_flat, ne12_flat, ne13_flat, stream); + } else if (src1_f16) { + quantize_mmq_q8_1_cuda(src1_f16, ids_src1.get(), src1_q8_1.get(), src0->type, ne10, s11, s12, s13, + ne10_padded, ne11_flat, ne12_flat, ne13_flat, stream); + } else { + quantize_mmq_q8_1_cuda(src1_d, ids_src1.get(), src1_q8_1.get(), src0->type, ne10, s11, s12, s13, + ne10_padded, ne11_flat, ne12_flat, ne13_flat, stream); + } + CUDA_CHECK(cudaGetLastError()); + } + + static_assert(QK_FP4_MMQ == 8 * QK_MXFP4, "QK_FP4_MMQ needs to be 8 * QK_MXFP4"); + const int64_t s12 = use_native_fp4 ? ne11 * ne10_padded * sizeof(block_fp4_mmq) / (QK_FP4_MMQ * sizeof(int)) : + ne11 * ne10_padded * sizeof(block_q8_1) / (QK8_1 * sizeof(int)); + const int64_t s13 = ne12*s12; + + // Note that ne02 is used instead of ne12 because the number of y channels determines the z dimension of the CUDA grid. + const mmq_args args = { + src0_d, src0->type, (const int *) src1_q8_1.get(), ids_dst.get(), expert_bounds.get(), dst_d, + src1_scale.ptr, + ne00, ne01, ne_get_rows, s01, ne_get_rows, s1, + ne02, ne02, s02, s12, s2, + ne03, ne13, s03, s13, s3, + ncols_max > 0 ? ncols_max : ne12, J_best, tile_map.ptr, ntx_flat}; + + ggml_cuda_mul_mat_q_switch_type(ctx, args, stream, prec_src1); +} + +bool ggml_cuda_should_use_mmq(enum ggml_type type, int cc, int64_t ne11, int64_t n_experts) { +#ifdef GGML_CUDA_FORCE_CUBLAS + return false; +#endif // GGML_CUDA_FORCE_CUBLAS + + bool mmq_supported; + + switch (type) { + case GGML_TYPE_Q1_0: + case GGML_TYPE_Q2_0: + case GGML_TYPE_Q4_0: + case GGML_TYPE_Q4_1: + case GGML_TYPE_Q5_0: + case GGML_TYPE_Q5_1: + case GGML_TYPE_Q8_0: +// ------------------------------------------------- + case GGML_TYPE_Q2_K: + case GGML_TYPE_Q3_K: + case GGML_TYPE_Q4_K: + case GGML_TYPE_Q5_K: + case GGML_TYPE_Q6_K: +// ------------------------------------------------- + case GGML_TYPE_IQ1_S: + case GGML_TYPE_IQ2_XXS: + case GGML_TYPE_IQ2_XS: + case GGML_TYPE_IQ2_S: + case GGML_TYPE_IQ3_XXS: + case GGML_TYPE_IQ3_S: + case GGML_TYPE_IQ4_XS: + case GGML_TYPE_IQ4_NL: +// ------------------------------------------------- + case GGML_TYPE_MXFP4: + case GGML_TYPE_NVFP4: + mmq_supported = true; + break; + default: + mmq_supported = false; + break; + } + + if (!mmq_supported) { + return false; + } + + // MMQ tiles require at least 48 KiB per-block shared memory; fall back to BLAS otherwise. + { + const int id = ggml_cuda_get_device(); + const size_t smpbo = ggml_cuda_info().devices[id].smpbo; + if (smpbo < 48 * 1024) { + return false; + } + } + + if (turing_mma_available(cc)) { + return true; + } + + if (ggml_cuda_highest_compiled_arch(cc) < GGML_CUDA_CC_DP4A) { + // for MoE, mmq is faster even without native dp4a + // TODO: check if cards older than pascal might benefit from this as well + return cc >= GGML_CUDA_CC_PASCAL && n_experts > 0; + } + +#ifdef GGML_CUDA_FORCE_MMQ + return true; +#endif //GGML_CUDA_FORCE_MMQ + + if (GGML_CUDA_CC_IS_NVIDIA(cc)) { + return !fp16_mma_hardware_available(cc) || ne11 < MMQ_DP4A_MAX_BATCH_SIZE; + } + + if (amd_mfma_available(cc)) { + // As of ROCM 7.0 rocblas/tensile performs very poorly on CDNA3 and hipblaslt (via ROCBLAS_USE_HIPBLASLT) + // performs better but is currently suffering from a crash on this architecture. + // TODO: Revisit when hipblaslt is fixed on CDNA3 + if (GGML_CUDA_CC_IS_CDNA3(cc)) { + return true; + } + if (n_experts > 64 || ne11 <= 128) { + return true; + } + if (type == GGML_TYPE_Q4_0 || type == GGML_TYPE_Q4_1 || type == GGML_TYPE_Q5_0 || type == GGML_TYPE_Q5_1) { + return true; + } + if (ne11 <= 256 && (type == GGML_TYPE_Q4_K || type == GGML_TYPE_Q5_K)) { + return true; + } + return false; + } + + if (amd_wmma_available(cc)) { + if (GGML_CUDA_CC_IS_RDNA3(cc)) { + // High expert counts are almost always better on MMQ due to + // the synchronization overhead in the cuBLAS/hipBLAS path: + // https://github.com/ggml-org/llama.cpp/pull/18202 + if (n_experts >= 64) { + return true; + } + + // For some quantization types MMQ can have lower peak TOPS than hipBLAS + // so it's only faster for sufficiently small batch sizes: + switch (type) { + case GGML_TYPE_Q2_K: + return ne11 <= 128; + case GGML_TYPE_Q6_K: + return ne11 <= (GGML_CUDA_CC_IS_RDNA3_0(cc) ? 128 : 256); + case GGML_TYPE_IQ2_XS: + case GGML_TYPE_IQ2_S: + return GGML_CUDA_CC_IS_RDNA3_5(cc) || ne11 <= 128; + default: + return true; + } + } + + // For RDNA4 MMQ is consistently faster than dequantization + hipBLAS: + // https://github.com/ggml-org/llama.cpp/pull/18537#issuecomment-3706422301 + return true; + } + + // gfx900 (Vega 10), gfx909, and gfx90c lack native dp4a, losing to dequant + hipBLAS + // for dense matrices; keep MMQ only for MoE, where the + // hipBLAS path is much slower. + if (cc == GGML_CUDA_CC_VEGA || GGML_CUDA_CC_IS_GCN_APU(cc)) { + return n_experts > 0; + } + + // MUSA: the MMQ kernels compute wrong values on PH1 (MTT S5000). + if (cc == GGML_CUDA_CC_PH1) { + return false; + } + + return (!GGML_CUDA_CC_IS_CDNA(cc)) || ne11 < MMQ_DP4A_MAX_BATCH_SIZE; +} diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq.cuh b/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq.cuh new file mode 100644 index 0000000000..596a54eabe --- /dev/null +++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq.cuh @@ -0,0 +1,1661 @@ +#pragma once + +#include "common.cuh" + +#include +#include + +#define MMQ_DP4A_MAX_BATCH_SIZE 64 // Max. batch size to use for dp4a MMQ kernels when FP16 tensor cores are available. +#define MMQ_ITER_K 256 +#define MMQ_ITER_K_FP4 512 +#define MMQ_NWARPS 8 + +// FreeToken: MMQ writes fp32, bf16 or fp16 results (dst_t); the sums stay fp32 until the store +template static __device__ __forceinline__ T mmq_cast(float v); +template <> __device__ __forceinline__ float mmq_cast(float v) { return v; } +template <> __device__ __forceinline__ nv_bfloat16 mmq_cast(float v) { return __float2bfloat16(v); } +template <> __device__ __forceinline__ half mmq_cast(float v) { return __float2half(v); } + +struct mmq_dst { + void * p; + ggml_type type; +}; + +typedef void (*ggml_cuda_mmq_load_tiles_t)(const char * __restrict__ x, int * x_tile, const int kbx0, const int i_max, const int stride); +typedef void (*ggml_cuda_mmq_vec_dot_t)(const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00); +typedef void (*ggml_cuda_mmq_write_back_t)(const float * __restrict__ sum, const int32_t * __restrict__ get_rows_to_sorted, + float * __restrict__ dst, const float * __restrict__ y_scale, const int stride, const int i_max, const int j_max); + +enum mmq_q8_1_ds_layout { + MMQ_Q8_1_DS_LAYOUT_D4, + MMQ_Q8_1_DS_LAYOUT_DS4, + MMQ_Q8_1_DS_LAYOUT_D2S6, +}; + +static constexpr int QK8_1_MMQ = 4*QK8_1; +static constexpr int QK_FP4_MMQ = 2*QK8_1_MMQ; + +struct block_q8_1_mmq { + // The y float data is converted to a data layout that can simply be copied to shared memory as a contiguous block. + // The y float data is first grouped as blocks of 128 values. + // These blocks are then treated as individual data values and transposed. + // + // To avoid shared memory bank conflicts each block is padded with 16 bytes. + // This padding is also used to store block scales/partial sums. + // The scales multiplied with the quantized data are equal to the unquantized values. + // The partial sums are obtained by summing up a subgroup of the contained values (prior to quantization) + // and are only needed for performance reasons. + // + // The exact data stored depends on the x data type. + union { + float d4[4]; // 1 32 bit scale per 32 values, stored as d0,d1,d2,d3 + half2 ds4[4]; // 1 16 bit scale + 1 16 bit partial sum per 32 values, stored as d0,s0,d1,s1,d2,s2,d3,s3 + half d2s6[8]; // 1 16 bit scale per 64 values + 1 16 bit partial sum per 16 values for the first 96 values, + // stored as d0,d1,s1,s2,s3,s4,s5 + }; + int8_t qs[QK8_1_MMQ]; +}; + +// this struct is used for fp4 data types (currently only used for Blackwell) +// mxfp4 has block size 32, each int32 of d4 contains 2 e8m0 scales in the lower 16 bits +// nvfp4 has block size 16, each int32 of d4 contains 4 ue4m3 scales +struct block_fp4_mmq { + uint32_t d4[4]; + int8_t qs[QK_FP4_MMQ / 2]; +}; + +static_assert(sizeof(block_q8_1_mmq) == QK8_1_MMQ + 4*sizeof(half2), "Unexpected block_q8_1_mmq size"); +static_assert(sizeof(block_q8_1_mmq) == 4*sizeof(block_q8_1), "Unexpected block_q8_1_mmq size"); +static_assert(sizeof(block_fp4_mmq) == sizeof(block_q8_1_mmq), "Unexpected block_fp4_mmq size"); + +static mmq_q8_1_ds_layout mmq_get_q8_1_ds_layout(const ggml_type type_x) { + switch (type_x) { + case GGML_TYPE_Q1_0: + case GGML_TYPE_Q2_0: + return MMQ_Q8_1_DS_LAYOUT_D4; + case GGML_TYPE_Q4_0: + case GGML_TYPE_Q4_1: + return MMQ_Q8_1_DS_LAYOUT_DS4; + case GGML_TYPE_Q5_0: + return MMQ_Q8_1_DS_LAYOUT_D4; + case GGML_TYPE_Q5_1: + return MMQ_Q8_1_DS_LAYOUT_DS4; + case GGML_TYPE_Q8_0: + return MMQ_Q8_1_DS_LAYOUT_D4; + case GGML_TYPE_MXFP4: + return MMQ_Q8_1_DS_LAYOUT_D4; + case GGML_TYPE_NVFP4: + return MMQ_Q8_1_DS_LAYOUT_D4; + case GGML_TYPE_Q2_K: + return MMQ_Q8_1_DS_LAYOUT_D2S6; + case GGML_TYPE_Q3_K: + return MMQ_Q8_1_DS_LAYOUT_D4; + case GGML_TYPE_Q4_K: + case GGML_TYPE_Q5_K: + return MMQ_Q8_1_DS_LAYOUT_DS4; + case GGML_TYPE_Q6_K: + case GGML_TYPE_IQ2_XXS: + case GGML_TYPE_IQ2_XS: + case GGML_TYPE_IQ2_S: + case GGML_TYPE_IQ3_XXS: + case GGML_TYPE_IQ3_S: + return MMQ_Q8_1_DS_LAYOUT_D4; + case GGML_TYPE_IQ1_S: + return MMQ_Q8_1_DS_LAYOUT_DS4; + case GGML_TYPE_IQ4_XS: + case GGML_TYPE_IQ4_NL: + return MMQ_Q8_1_DS_LAYOUT_D4; + default: + GGML_ABORT("fatal error"); + } +} + +struct tile_x_sizes { + int qs; + int dm; + int sc; +}; + +// Decouple shared memory tile sizes from WARP_SIZE to allow for different warp sizes. +// The K dimension of the tiles has either, +// 1*MMQ_TILE_NE_K==32 (always for TILE_Y_K) or 2*MMQ_TILE_NE_K==64 (typically for TILE_X_K), +// 32 bit elements for the quantized data (does not include scales). +// In other words, the size of the quantized data in the K dimension is a multiple of MMQ_TILE_NE_K. +// The final tile size in K direction is padded to avoid shared memory bank conflicts, +// in terms of 32 bit elements that means K % 2 == 1 for dp4a or K % 8 == 4 for mma. +#define MMQ_TILE_NE_K 32 + +// block_q8_1_mmq has (128 8-bit ints == 32 32-bit ints + 4 32-bit scales) +#define MMQ_TILE_Y_K (MMQ_TILE_NE_K + MMQ_TILE_NE_K / QI8_1) +#define MMQ_TILE_Y_FP4_K MMQ_TILE_Y_K + +enum ggml_cuda_mmq_sram_layout { + GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, + GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, + GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, + GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, + GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, + GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, // MXFP4 and NVFP4 on Blackwell. + GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, // Generic NVFP4 +}; + +static constexpr __host__ __device__ int ggml_cuda_mmq_get_sram_stride(ggml_cuda_mmq_sram_layout sram_layout) { + switch (sram_layout) { + case GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0: + return 2*MMQ_TILE_NE_K + 2*MMQ_TILE_NE_K/QI8_0 + 4; + case GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1: + return 2*MMQ_TILE_NE_K + 2*MMQ_TILE_NE_K/QI8_1 + 4; + case GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K: + return 2*MMQ_TILE_NE_K + MMQ_TILE_NE_K + 4; + case GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K: + return 2*MMQ_TILE_NE_K + MMQ_TILE_NE_K/2 + 4; + case GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K: + return 2*MMQ_TILE_NE_K + MMQ_TILE_NE_K/QI6_K + MMQ_TILE_NE_K/8 + 7; + case GGML_CUDA_MMQ_SRAM_LAYOUT_FP4: + return 2*MMQ_TILE_NE_K + 8 + 4; + case GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4: + return 2*MMQ_TILE_NE_K + MMQ_TILE_NE_K/2 + 4; + default: + return -1; + } +} + +static_assert(ggml_cuda_mmq_get_sram_stride(GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0) % 8 == 4, "Wrong padding."); +static_assert(ggml_cuda_mmq_get_sram_stride(GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1) % 8 == 4, "Wrong padding."); +static_assert(ggml_cuda_mmq_get_sram_stride(GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K) % 8 == 4, "Wrong padding."); +static_assert(ggml_cuda_mmq_get_sram_stride(GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K) % 8 == 4, "Wrong padding."); +static_assert(ggml_cuda_mmq_get_sram_stride(GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K) % 8 == 4, "Wrong padding."); +static_assert(ggml_cuda_mmq_get_sram_stride(GGML_CUDA_MMQ_SRAM_LAYOUT_FP4) % 8 == 4, "Wrong padding."); +static_assert(ggml_cuda_mmq_get_sram_stride(GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4) % 8 == 4, "Wrong padding."); + +static_assert(ggml_cuda_mmq_get_sram_stride(GGML_CUDA_MMQ_SRAM_LAYOUT_FP4) == ggml_cuda_mmq_get_sram_stride(GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1), "Wrong tile size for MXFP4"); + +// Config options for the MMQ kernel. +// Should not affect results, only speed/register pressure/shared memory use. +struct ggml_cuda_mmq_config { + ggml_type type; // src0->type + int nthreads; // Number of threads per CUDA block. + int occupancy; // Targeted occupancy for the MMA kernel. + int I; // SRAM tile width in src0->ne[1]/dst->ne[0] direction. + int J; // SRAM tile width in src1->ne[1]/dst->ne[1] direction. + ggml_cuda_mmq_sram_layout sram_layout; // SRAM tile length in src0->ne[0]/src1->ne[0] direction (physical 32 bit elements). + int K_vram; // VRAM tile length in src0->ne[0]/src1->ne[0] direction (logical elements). + bool stream_k; // Whether or not to use stream-k decomposition. + bool fallback; // Whether a fallback for out-of-bounds check in src0->ne[1] direction is needed. + + constexpr __host__ __device__ ggml_cuda_mmq_config( + ggml_type type, int nthreads, int occupancy, int I, int J, ggml_cuda_mmq_sram_layout sram_layout, int K_vram, bool stream_k, bool fallback) : + type(type), nthreads(nthreads), occupancy(occupancy), I(I), J(J), sram_layout(sram_layout), K_vram(K_vram), stream_k(stream_k), fallback(fallback) {} + + constexpr __device__ int rows_per_warp() const { +#if defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + return 16; +#else + return J >= 48 && J % 16 == 0 ? 32 : 16; +#endif // defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + } + + // TODO transition all combinations of GPUs and quantizations to the MMA data layout. + __host__ int use_mma_data_layout(const int cc) const { + if (amd_mfma_available(cc) || amd_wmma_available(cc) || turing_mma_available(cc)) { + return true; + } + return false; + } + + constexpr __device__ bool use_mma_data_layout() const { +#if defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) + return true; +#else + return false; +#endif // defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) + } + +}; + +#define CASE(type_, nthreads_, occupancy_, I_, J_, sram_layout_, K_vram_, stream_k_, fallback_) \ + if (type == (type_) && J == (J_) && fallback == (fallback_)) { \ + static_assert((nthreads_) % 32 == 0 && (nthreads_) <= 512, "bad nthreads"); \ + static_assert( (occupancy_) <= 8, "bad occupancy"); \ + static_assert((I_) % 32 == 0, "bad I"); \ + static_assert((J_) % 8 == 0 && (J_) <= 128, "bad J"); \ + static_assert((K_vram_) % 256 == 0, "bad K_vram"); \ + return ggml_cuda_mmq_config((type_), (nthreads_), (occupancy_), (I_), (J_), (sram_layout_), (K_vram_), (stream_k_), (fallback_)); \ + } \ + +#include "mmq-config-pascal-older.cuh" +#include "mmq-config-pascal-dp4a.cuh" +#include "mmq-config-ampere.cuh" +#include "mmq-config-blackwell.cuh" + +#include "mmq-config-gcn.cuh" +#include "mmq-config-cdna.cuh" +#include "mmq-config-rdna2.cuh" +#include "mmq-config-rdna3.cuh" +#include "mmq-config-rdna3-5.cuh" +#include "mmq-config-rdna4.cuh" + +#undef CASE + +static __host__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config(const ggml_type type, const int J, const bool fallback, const int cc, const ggml_prec prec_src1) { + if (GGML_CUDA_CC_IS_AMD(cc)) { + if (GGML_CUDA_CC_IS_GCN(cc)) { + return ggml_cuda_mmq_get_config_gcn(type, J, fallback); + } + if (GGML_CUDA_CC_IS_CDNA(cc)) { + return ggml_cuda_mmq_get_config_cdna(type, J, fallback); + } + if (GGML_CUDA_CC_IS_RDNA4(cc)) { + return ggml_cuda_mmq_get_config_rdna4(type, J, fallback); + } + if (GGML_CUDA_CC_IS_RDNA3_5(cc)) { + return ggml_cuda_mmq_get_config_rdna3_5(type, J, fallback); + } + if (GGML_CUDA_CC_IS_RDNA3(cc)) { // covers RDNA 3.0 + return ggml_cuda_mmq_get_config_rdna3(type, J, fallback); + } + return ggml_cuda_mmq_get_config_rdna2(type, J, fallback); + } + if (blackwell_mma_available(cc)) { + // only src1 at Q4 uses the native FP4 config, higher precisions keep src1 at Q8_1 + if (prec_src1 != GGML_PREC_Q4 && (type == GGML_TYPE_NVFP4 || type == GGML_TYPE_MXFP4)) { + return ggml_cuda_mmq_get_config_ampere(type, J, fallback); + } + return ggml_cuda_mmq_get_config_blackwell(type, J, fallback); + } + if (ggml_cuda_highest_compiled_arch(cc) >= GGML_CUDA_CC_VOLTA) { + return ggml_cuda_mmq_get_config_ampere(type, J, fallback); + } + if (ggml_cuda_highest_compiled_arch(cc) >= GGML_CUDA_CC_DP4A) { + return ggml_cuda_mmq_get_config_pascal_dp4a(type, J, fallback); + } + return ggml_cuda_mmq_get_config_pascal_older(type, J, fallback); +} + +static constexpr __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config(ggml_type type, int J, bool fallback, ggml_prec prec_src1) { +#ifdef GGML_USE_HIP +#ifdef GCN + return ggml_cuda_mmq_get_config_gcn(type, J, fallback); +#elif defined(CDNA) + return ggml_cuda_mmq_get_config_cdna(type, J, fallback); +#elif defined(RDNA4) + return ggml_cuda_mmq_get_config_rdna4(type, J, fallback); +#elif defined(RDNA3_5) + return ggml_cuda_mmq_get_config_rdna3_5(type, J, fallback); +#elif defined(RDNA3) + return ggml_cuda_mmq_get_config_rdna3(type, J, fallback); +#else + return ggml_cuda_mmq_get_config_rdna2(type, J, fallback); +#endif // CDNA +#else +#ifdef BLACKWELL_MMA_AVAILABLE + // only src1 at Q4 uses the native FP4 config, higher precisions keep src1 at Q8_1 + if (prec_src1 != GGML_PREC_Q4 && (type == GGML_TYPE_NVFP4 || type == GGML_TYPE_MXFP4)) { + return ggml_cuda_mmq_get_config_ampere(type, J, fallback); + } + return ggml_cuda_mmq_get_config_blackwell(type, J, fallback); +#elif !defined(GGML_USE_MUSA) && __CUDA_ARCH__ >= GGML_CUDA_CC_VOLTA + return ggml_cuda_mmq_get_config_ampere(type, J, fallback); +#elif !defined(GGML_USE_MUSA) && __CUDA_ARCH__ >= GGML_CUDA_CC_DP4A + return ggml_cuda_mmq_get_config_pascal_dp4a(type, J, fallback); +#else + return ggml_cuda_mmq_get_config_pascal_older(type, J, fallback); +#endif // BLACKWELL_MMA_AVAILABLE +#endif // GGML_USE_HIP + GGML_UNUSED_VARS(type, J, fallback, prec_src1); +} + +static __host__ int ggml_cuda_mmq_get_type(const ggml_type type, const int J, const bool fallback, const int cc, const ggml_prec prec_src1) { + return ggml_cuda_mmq_get_config(type, J, fallback, cc, prec_src1).type; +} + +static constexpr __device__ int ggml_cuda_mmq_get_type(ggml_type type, int J, bool fallback, ggml_prec prec_src1) { + return ggml_cuda_mmq_get_config(type, J, fallback, prec_src1).type; +} + +static constexpr __device__ int ggml_cuda_mmq_get_nthreads(ggml_type type, int J, bool fallback, ggml_prec prec_src1) { + return ggml_cuda_mmq_get_config(type, J, fallback, prec_src1).nthreads; +} + +static constexpr __device__ int ggml_cuda_mmq_get_occupancy(ggml_type type, int J, bool fallback, ggml_prec prec_src1) { + return ggml_cuda_mmq_get_config(type, J, fallback, prec_src1).occupancy; +} + +static __host__ int ggml_cuda_mmq_get_I(const ggml_type type, const int J, const bool fallback, const int cc, const ggml_prec prec_src1) { + return ggml_cuda_mmq_get_config(type, J, fallback, cc, prec_src1).I; +} + +static constexpr __device__ int ggml_cuda_mmq_get_I(ggml_type type, int J, bool fallback, ggml_prec prec_src1) { + return ggml_cuda_mmq_get_config(type, J, fallback, prec_src1).I; +} + +static __host__ int ggml_cuda_mmq_get_J(const ggml_type type, const int J, const bool fallback, const int cc, const ggml_prec prec_src1) { + return ggml_cuda_mmq_get_config(type, J, fallback, cc, prec_src1).J; +} + +static constexpr __device__ int ggml_cuda_mmq_get_J(ggml_type type, int J, bool fallback, ggml_prec prec_src1) { + return ggml_cuda_mmq_get_config(type, J, fallback, prec_src1).J; +} + +static __host__ ggml_cuda_mmq_sram_layout ggml_cuda_mmq_get_sram_layout(const ggml_type type, const int J, const bool fallback, const int cc, const ggml_prec prec_src1) { + return ggml_cuda_mmq_get_config(type, J, fallback, cc, prec_src1).sram_layout; +} + +static constexpr __device__ ggml_cuda_mmq_sram_layout ggml_cuda_mmq_get_sram_layout(ggml_type type, int J, bool fallback, ggml_prec prec_src1) { + return ggml_cuda_mmq_get_config(type, J, fallback, prec_src1).sram_layout; +} + +static __host__ int ggml_cuda_mmq_get_K_vram(const ggml_type type, const int J, const bool fallback, const int cc, const ggml_prec prec_src1) { + return ggml_cuda_mmq_get_config(type, J, fallback, cc, prec_src1).K_vram; +} + +static constexpr __device__ int ggml_cuda_mmq_get_K_vram(ggml_type type, int J, bool fallback, ggml_prec prec_src1) { + return ggml_cuda_mmq_get_config(type, J, fallback, prec_src1).K_vram; +} + +static __host__ bool ggml_cuda_mmq_get_stream_k(const ggml_type type, const int J, const bool fallback, const int cc, const ggml_prec prec_src1) { + return ggml_cuda_mmq_get_config(type, J, fallback, cc, prec_src1).stream_k; +} + +static constexpr __device__ bool ggml_cuda_mmq_get_stream_k(ggml_type type, int J, bool fallback, ggml_prec prec_src1) { + return ggml_cuda_mmq_get_config(type, J, fallback, prec_src1).stream_k; +} + +static __host__ int ggml_cuda_mmq_get_fallback(const ggml_type type, const int J, const bool fallback, const int cc, const ggml_prec prec_src1) { + return ggml_cuda_mmq_get_config(type, J, fallback, cc, prec_src1).fallback; +} + +static constexpr __device__ int ggml_cuda_mmq_get_fallback(ggml_type type, int J, bool fallback, ggml_prec prec_src1) { + return ggml_cuda_mmq_get_config(type, J, fallback, prec_src1).fallback; +} + +// --------------------------------------------------------------------------------------------- + +static __host__ int ggml_cuda_mmq_get_sram_stride(const ggml_type type, const int J, const bool fallback, const int cc, const ggml_prec prec_src1) { + return ggml_cuda_mmq_get_sram_stride(ggml_cuda_mmq_get_sram_layout(type, J, fallback, cc, prec_src1)); +} + +static constexpr __device__ int ggml_cuda_mmq_get_sram_stride(ggml_type type, int J, bool fallback, ggml_prec prec_src1) { + return ggml_cuda_mmq_get_sram_stride(ggml_cuda_mmq_get_sram_layout(type, J, fallback, prec_src1)); +} + +static __host__ bool ggml_cuda_mmq_needs_fallback(const int64_t nrows_x) { + return nrows_x % 128 != 0; +} + +static constexpr __device__ int ggml_cuda_mmq_get_rows_per_warp(ggml_type type, int J, bool fallback, ggml_prec prec_src1) { + return ggml_cuda_mmq_get_config(type, J, fallback, prec_src1).rows_per_warp(); +} + +#define MMQ_DP4A_TXS_Q4_0 tile_x_sizes{I*MMQ_TILE_NE_K + I, I*MMQ_TILE_NE_K/QI4_0 + I/QI4_0, 0} +#define MMQ_DP4A_TXS_Q4_1 tile_x_sizes{I*MMQ_TILE_NE_K + I, I*MMQ_TILE_NE_K/QI4_1 + I/QI4_1, 0} +#define MMQ_DP4A_TXS_Q8_0 tile_x_sizes{I*MMQ_TILE_NE_K*2 + I, I*MMQ_TILE_NE_K*2/QI8_0 + I/(QI8_0/2), 0} +#define MMQ_DP4A_TXS_Q8_0_16 tile_x_sizes{I*MMQ_TILE_NE_K*2 + I, I*MMQ_TILE_NE_K*4/QI8_0 + I/(QI8_0/4), 0} +#define MMQ_DP4A_TXS_Q8_1 tile_x_sizes{I*MMQ_TILE_NE_K*2 + I, I*MMQ_TILE_NE_K*2/QI8_1 + I/(QI8_1/2), 0} +#define MMQ_DP4A_TXS_Q2_K tile_x_sizes{I*MMQ_TILE_NE_K*2 + I, I*MMQ_TILE_NE_K + I, 0} +#define MMQ_DP4A_TXS_Q3_K tile_x_sizes{I*MMQ_TILE_NE_K*2 + I, I, I*MMQ_TILE_NE_K/8 + I/8} +#define MMQ_DP4A_TXS_Q4_K tile_x_sizes{I*MMQ_TILE_NE_K + I, I*MMQ_TILE_NE_K/QI4_K, I*MMQ_TILE_NE_K/8 + I/8} +#define MMQ_DP4A_TXS_Q5_K tile_x_sizes{I*MMQ_TILE_NE_K*2 + I, I*MMQ_TILE_NE_K/QI5_K + I/QI5_K, I*MMQ_TILE_NE_K/8 + I/8} +#define MMQ_DP4A_TXS_Q6_K tile_x_sizes{I*MMQ_TILE_NE_K*2 + I, I*MMQ_TILE_NE_K/QI6_K + I/QI6_K, I*MMQ_TILE_NE_K/8 + I/8} + +static constexpr __host__ __device__ tile_x_sizes mmq_get_dp4a_tile_x_sizes(ggml_type type, int I) { + switch (type) { + case GGML_TYPE_Q1_0: return MMQ_DP4A_TXS_Q8_0; + case GGML_TYPE_Q2_0: return MMQ_DP4A_TXS_Q8_0; + case GGML_TYPE_Q4_0: return MMQ_DP4A_TXS_Q4_0; + case GGML_TYPE_Q4_1: return MMQ_DP4A_TXS_Q4_1; + case GGML_TYPE_Q5_0: return MMQ_DP4A_TXS_Q8_0; + case GGML_TYPE_Q5_1: return MMQ_DP4A_TXS_Q8_1; + case GGML_TYPE_Q8_0: return MMQ_DP4A_TXS_Q8_0; + case GGML_TYPE_MXFP4: return MMQ_DP4A_TXS_Q8_1; + case GGML_TYPE_NVFP4: return MMQ_DP4A_TXS_Q8_0_16; + case GGML_TYPE_Q2_K: return MMQ_DP4A_TXS_Q2_K; + case GGML_TYPE_Q3_K: return MMQ_DP4A_TXS_Q3_K; + case GGML_TYPE_Q4_K: return MMQ_DP4A_TXS_Q4_K; + case GGML_TYPE_Q5_K: return MMQ_DP4A_TXS_Q5_K; + case GGML_TYPE_Q6_K: return MMQ_DP4A_TXS_Q6_K; + case GGML_TYPE_IQ2_XXS: return MMQ_DP4A_TXS_Q8_0; + case GGML_TYPE_IQ2_XS: return MMQ_DP4A_TXS_Q8_0_16; + case GGML_TYPE_IQ2_S: return MMQ_DP4A_TXS_Q8_0_16; + case GGML_TYPE_IQ3_XXS: return MMQ_DP4A_TXS_Q8_0; + case GGML_TYPE_IQ3_S: return MMQ_DP4A_TXS_Q8_0; + case GGML_TYPE_IQ1_S: return MMQ_DP4A_TXS_Q8_0; + case GGML_TYPE_IQ4_XS: return MMQ_DP4A_TXS_Q8_0; + case GGML_TYPE_IQ4_NL: return MMQ_DP4A_TXS_Q8_0; + default: return tile_x_sizes{0, 0, 0}; + } +} + +// FIXME temporary until all combinations of data types and GPUs can use the MMA data layout +static __host__ int ggml_cuda_mmq_get_nbytes_shared_x(const ggml_cuda_mmq_config & config, const int cc) { + if (config.use_mma_data_layout(cc)) { + return config.I * ggml_cuda_mmq_get_sram_stride(config.sram_layout) * 4; + } + const tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(config.type, config.I); + return (txs.qs + txs.dm + txs.sc) * 4; +} + +// ------------------------------------------------------------ + +#include "mmq-load-tiles.cuh" +#include "mmq-vec-dot.cuh" + +template static __device__ __forceinline__ void ggml_cuda_mmq_write_back_dp4a( + const float * __restrict__ sum, const int32_t * __restrict__ ids_dst, dst_t * __restrict__ dst, + const float * __restrict__ y_scale, const int stride, const int i_max, const int j_max) { + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback, prec_src1) / warp_size; + constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback, prec_src1); + + const bool y_scale_used = y_scale != nullptr; + +#pragma unroll + for (int j0 = 0; j0 < J; j0 += nwarps) { + const int j = j0 + threadIdx.y; + + if (j > j_max) { + return; + } + +#pragma unroll + for (int i0 = 0; i0 < I; i0 += warp_size) { + const int i = i0 + threadIdx.x; + + if (fallback && i > i_max) { + continue; + } + + if constexpr (type == GGML_TYPE_NVFP4) { + if (y_scale_used) { + dst[ids_dst[j]*stride + i] = mmq_cast(y_scale[j] * sum[(j0/nwarps) * (I/warp_size) + i0/warp_size]); + } else { + dst[ids_dst[j]*stride + i] = mmq_cast(sum[(j0/nwarps) * (I/warp_size) + i0/warp_size]); + } + } else { + dst[ids_dst[j]*stride + i] = mmq_cast(sum[(j0/nwarps) * (I/warp_size) + i0/warp_size]); + GGML_UNUSED(y_scale_used); + } + } + } +} + +template +static __device__ __forceinline__ void ggml_cuda_mmq_write_back_mma( + const float * __restrict__ sum, const int * __restrict__ ids_dst, dst_t * __restrict__ dst, + const float * __restrict__ y_scale, const int stride, const int i_max, const int j_max) { + +#if defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + typedef tile<16, 16, int, DATA_LAYOUT_J_MAJOR> tile_C; +#else + typedef tile<16, 8, int> tile_C; +#endif // defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + + constexpr int rows_per_warp = ggml_cuda_mmq_get_rows_per_warp(type, J, fallback, prec_src1); + constexpr int ntx = rows_per_warp/tile_C::I; // Number of x minitiles per warp. + + const int i0 = (threadIdx.y / ntx) * (ntx*tile_C::I); + + const bool y_scale_used = y_scale != nullptr; + +#pragma unroll + for (int j0 = 0; j0 < J; j0 += ntx*tile_C::J) { +#pragma unroll + for (int n = 0; n < ntx; ++n) { +#pragma unroll + for (int l = 0; l < tile_C::ne; ++l) { + const int j = j0 + (threadIdx.y % ntx) * tile_C::J + tile_C::get_j(l); + + if (j > j_max) { + continue; + } + + const int i = i0 + n*tile_C::I + tile_C::get_i(l); + + if (fallback && i > i_max) { + continue; + } + + if constexpr (type == GGML_TYPE_NVFP4) { + if (y_scale_used) { + dst[ids_dst[j]*stride + i] = mmq_cast(y_scale[j] * sum[(j0/tile_C::J + n)*tile_C::ne + l]); + } else { + dst[ids_dst[j]*stride + i] = mmq_cast(sum[(j0/tile_C::J + n)*tile_C::ne + l]); + } + } else { + dst[ids_dst[j]*stride + i] = mmq_cast(sum[(j0/tile_C::J + n)*tile_C::ne + l]); + GGML_UNUSED(y_scale_used); + } + } + } + } +} + +// ------------------------------------------------------------------------------------------------------------------------------------- + +// TODO remove this struct and use ggml_cuda_mmq_sram_layout instead. +struct ggml_cuda_mmq_util_funcs { + int vdr; + ggml_cuda_mmq_load_tiles_t load_tiles; + ggml_cuda_mmq_vec_dot_t vec_dot; + ggml_cuda_mmq_write_back_t write_back; + + constexpr __host__ __device__ ggml_cuda_mmq_util_funcs( + int vdr, ggml_cuda_mmq_load_tiles_t load_tiles, ggml_cuda_mmq_vec_dot_t vec_dot, ggml_cuda_mmq_write_back_t write_back) : + vdr(vdr), load_tiles(load_tiles), vec_dot(vec_dot), write_back(write_back) {} +}; + +template +static constexpr __device__ ggml_cuda_mmq_util_funcs ggml_cuda_mmq_get_util_funcs() { + if (!ggml_cuda_mmq_get_config(type, J, fallback, prec_src1).use_mma_data_layout()) { + switch (type) { + case GGML_TYPE_Q1_0: + return ggml_cuda_mmq_util_funcs( + VDR_Q1_0_Q8_1_MMQ, + ggml_cuda_mmq_load_tiles_q1_0, + ggml_cuda_mmq_vec_dot_q8_0_q8_1_dp4a, + ggml_cuda_mmq_write_back_dp4a); + case GGML_TYPE_Q2_0: + return ggml_cuda_mmq_util_funcs( + VDR_Q2_0_Q8_1_MMQ, + ggml_cuda_mmq_load_tiles_q2_0, + ggml_cuda_mmq_vec_dot_q8_0_q8_1_dp4a, + ggml_cuda_mmq_write_back_dp4a); + case GGML_TYPE_Q4_0: + return ggml_cuda_mmq_util_funcs( + VDR_Q4_0_Q8_1_MMQ, + ggml_cuda_mmq_load_tiles_q4_0, + ggml_cuda_mmq_vec_dot_q4_0_q8_1_dp4a, + ggml_cuda_mmq_write_back_dp4a); + case GGML_TYPE_Q4_1: + return ggml_cuda_mmq_util_funcs( + VDR_Q4_1_Q8_1_MMQ, + ggml_cuda_mmq_load_tiles_q4_1, + ggml_cuda_mmq_vec_dot_q4_1_q8_1_dp4a, + ggml_cuda_mmq_write_back_dp4a); + case GGML_TYPE_Q5_0: + return ggml_cuda_mmq_util_funcs( + VDR_Q5_0_Q8_1_MMQ, + ggml_cuda_mmq_load_tiles_q5_0, + ggml_cuda_mmq_vec_dot_q8_0_q8_1_dp4a, + ggml_cuda_mmq_write_back_dp4a); + case GGML_TYPE_Q5_1: + return ggml_cuda_mmq_util_funcs( + VDR_Q5_1_Q8_1_MMQ, + ggml_cuda_mmq_load_tiles_q5_1, + ggml_cuda_mmq_vec_dot_q8_1_q8_1_dp4a, + ggml_cuda_mmq_write_back_dp4a); + case GGML_TYPE_Q8_0: + return ggml_cuda_mmq_util_funcs( + VDR_Q8_0_Q8_1_MMQ, + ggml_cuda_mmq_load_tiles_q8_0, + ggml_cuda_mmq_vec_dot_q8_0_q8_1_dp4a, + ggml_cuda_mmq_write_back_dp4a); +// --------------------------------------------------------------------------------------------- + case GGML_TYPE_Q2_K: + return ggml_cuda_mmq_util_funcs( + VDR_Q2_K_Q8_1_MMQ, + ggml_cuda_mmq_load_tiles_q2_K, + ggml_cuda_mmq_vec_dot_q2_K_q8_1_dp4a, + ggml_cuda_mmq_write_back_dp4a); + case GGML_TYPE_Q3_K: + return ggml_cuda_mmq_util_funcs( + VDR_Q3_K_Q8_1_MMQ, + ggml_cuda_mmq_load_tiles_q3_K, + ggml_cuda_mmq_vec_dot_q3_K_q8_1_dp4a, + ggml_cuda_mmq_write_back_dp4a); + case GGML_TYPE_Q4_K: + return ggml_cuda_mmq_util_funcs( + VDR_Q4_K_Q8_1_MMQ, + ggml_cuda_mmq_load_tiles_q4_K, + ggml_cuda_mmq_vec_dot_q4_K_q8_1_dp4a, + ggml_cuda_mmq_write_back_dp4a); + case GGML_TYPE_Q5_K: + return ggml_cuda_mmq_util_funcs( + VDR_Q5_K_Q8_1_MMQ, + ggml_cuda_mmq_load_tiles_q5_K, + ggml_cuda_mmq_vec_dot_q5_K_q8_1_dp4a, + ggml_cuda_mmq_write_back_dp4a); + case GGML_TYPE_Q6_K: + return ggml_cuda_mmq_util_funcs( + VDR_Q6_K_Q8_1_MMQ, + ggml_cuda_mmq_load_tiles_q6_K, + ggml_cuda_mmq_vec_dot_q6_K_q8_1_dp4a, + ggml_cuda_mmq_write_back_dp4a); +// --------------------------------------------------------------------------------------------- + case GGML_TYPE_IQ1_S: + return ggml_cuda_mmq_util_funcs( + VDR_IQ1_S_Q8_1_MMQ, + ggml_cuda_mmq_load_tiles_iq1_s, + ggml_cuda_mmq_vec_dot_q8_1_q8_1_dp4a, + ggml_cuda_mmq_write_back_dp4a); + case GGML_TYPE_IQ2_XXS: + return ggml_cuda_mmq_util_funcs( + VDR_IQ2_XXS_Q8_1_MMQ, + ggml_cuda_mmq_load_tiles_iq2_xxs, + ggml_cuda_mmq_vec_dot_q8_0_q8_1_dp4a, + ggml_cuda_mmq_write_back_dp4a); + case GGML_TYPE_IQ2_XS: + return ggml_cuda_mmq_util_funcs( + VDR_IQ2_XS_Q8_1_MMQ, + ggml_cuda_mmq_load_tiles_iq2_xs, + ggml_cuda_mmq_vec_dot_q8_0_16_q8_1_dp4a, + ggml_cuda_mmq_write_back_dp4a); + case GGML_TYPE_IQ2_S: + return ggml_cuda_mmq_util_funcs( + VDR_IQ2_S_Q8_1_MMQ, + ggml_cuda_mmq_load_tiles_iq2_s, + ggml_cuda_mmq_vec_dot_q8_0_16_q8_1_dp4a, + ggml_cuda_mmq_write_back_dp4a); + case GGML_TYPE_IQ3_XXS: + return ggml_cuda_mmq_util_funcs( + VDR_IQ3_XXS_Q8_1_MMQ, + ggml_cuda_mmq_load_tiles_iq3_xxs, + ggml_cuda_mmq_vec_dot_q8_0_q8_1_dp4a, + ggml_cuda_mmq_write_back_dp4a); + case GGML_TYPE_IQ3_S: + return ggml_cuda_mmq_util_funcs( + VDR_IQ3_S_Q8_1_MMQ, + ggml_cuda_mmq_load_tiles_iq3_s, + ggml_cuda_mmq_vec_dot_q8_0_q8_1_dp4a, + ggml_cuda_mmq_write_back_dp4a); + case GGML_TYPE_IQ4_XS: + return ggml_cuda_mmq_util_funcs( + VDR_IQ4_XS_Q8_1_MMQ, + ggml_cuda_mmq_load_tiles_iq4_xs, + ggml_cuda_mmq_vec_dot_q8_0_q8_1_dp4a, + ggml_cuda_mmq_write_back_dp4a); + case GGML_TYPE_IQ4_NL: + return ggml_cuda_mmq_util_funcs( + VDR_IQ4_NL_Q8_1_MMQ, + ggml_cuda_mmq_load_tiles_iq4_nl, + ggml_cuda_mmq_vec_dot_q8_0_q8_1_dp4a, + ggml_cuda_mmq_write_back_dp4a); +// --------------------------------------------------------------------------------------------- + case GGML_TYPE_MXFP4: + return ggml_cuda_mmq_util_funcs( + VDR_MXFP4_Q8_1_MMQ, + ggml_cuda_mmq_load_tiles_mxfp4, + ggml_cuda_mmq_vec_dot_q8_0_q8_1_dp4a, + ggml_cuda_mmq_write_back_dp4a); + case GGML_TYPE_NVFP4: + return ggml_cuda_mmq_util_funcs( + VDR_NVFP4_Q8_1_MMQ, + ggml_cuda_mmq_load_tiles_nvfp4, + ggml_cuda_mmq_vec_dot_q8_0_16_q8_1_dp4a, + ggml_cuda_mmq_write_back_dp4a); + default: + return ggml_cuda_mmq_util_funcs(1, nullptr, nullptr, nullptr); + } + } + +// --------------------------------------------------------------------------------------------- + +#ifdef BLACKWELL_MMA_AVAILABLE + switch (type) { + case GGML_TYPE_MXFP4: + if (prec_src1 == GGML_PREC_Q4) { + return ggml_cuda_mmq_util_funcs( + -1, + ggml_cuda_mmq_load_tiles_mxfp4_fp4, + ggml_cuda_mmq_vec_dot_fp4_fp4_mma, + ggml_cuda_mmq_write_back_mma); + } + break; + case GGML_TYPE_NVFP4: + if (prec_src1 == GGML_PREC_Q4) { + return ggml_cuda_mmq_util_funcs( + -1, + ggml_cuda_mmq_load_tiles_nvfp4_nvfp4, + ggml_cuda_mmq_vec_dot_fp4_fp4_mma, + ggml_cuda_mmq_write_back_mma); + } + break; + default: + break; + } +#endif // BLACKWELL_MMA_AVAILABLE + +// --------------------------------------------------------------------------------------------- + + switch (type) { + case GGML_TYPE_Q1_0: + return ggml_cuda_mmq_util_funcs( + -1, + ggml_cuda_mmq_load_tiles_q1_0, + ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma, + ggml_cuda_mmq_write_back_mma); + case GGML_TYPE_Q2_0: + return ggml_cuda_mmq_util_funcs( + -1, + ggml_cuda_mmq_load_tiles_q2_0, + ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma, + ggml_cuda_mmq_write_back_mma); + case GGML_TYPE_Q4_0: + return ggml_cuda_mmq_util_funcs( + -1, + ggml_cuda_mmq_load_tiles_q4_0, + ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma, + ggml_cuda_mmq_write_back_mma); + case GGML_TYPE_Q4_1: + return ggml_cuda_mmq_util_funcs( + -1, + ggml_cuda_mmq_load_tiles_q4_1, + ggml_cuda_mmq_vec_dot_q8_1_q8_1_mma, + ggml_cuda_mmq_write_back_mma); + case GGML_TYPE_Q5_0: + return ggml_cuda_mmq_util_funcs( + -1, + ggml_cuda_mmq_load_tiles_q5_0, + ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma, + ggml_cuda_mmq_write_back_mma); + case GGML_TYPE_Q5_1: + return ggml_cuda_mmq_util_funcs( + -1, + ggml_cuda_mmq_load_tiles_q5_1, + ggml_cuda_mmq_vec_dot_q8_1_q8_1_mma, + ggml_cuda_mmq_write_back_mma); + case GGML_TYPE_Q8_0: + return ggml_cuda_mmq_util_funcs( + -1, + ggml_cuda_mmq_load_tiles_q8_0, + ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma, + ggml_cuda_mmq_write_back_mma); +// --------------------------------------------------------------------------------------------- + case GGML_TYPE_Q2_K: + return ggml_cuda_mmq_util_funcs( + -1, + ggml_cuda_mmq_load_tiles_q2_K, + ggml_cuda_mmq_vec_dot_q2_K_q8_1_mma, + ggml_cuda_mmq_write_back_mma); + case GGML_TYPE_Q3_K: + return ggml_cuda_mmq_util_funcs( + -1, + ggml_cuda_mmq_load_tiles_q3_K, + ggml_cuda_mmq_vec_dot_q8_0_16_q8_1_mma, + ggml_cuda_mmq_write_back_mma); + case GGML_TYPE_Q4_K: + return ggml_cuda_mmq_util_funcs( + -1, + ggml_cuda_mmq_load_tiles_q4_K, + ggml_cuda_mmq_vec_dot_q8_1_q8_1_mma, + ggml_cuda_mmq_write_back_mma); + case GGML_TYPE_Q5_K: + return ggml_cuda_mmq_util_funcs( + -1, + ggml_cuda_mmq_load_tiles_q5_K, + ggml_cuda_mmq_vec_dot_q8_1_q8_1_mma, + ggml_cuda_mmq_write_back_mma); + case GGML_TYPE_Q6_K: + return ggml_cuda_mmq_util_funcs( + -1, + ggml_cuda_mmq_load_tiles_q6_K, + ggml_cuda_mmq_vec_dot_q6_K_q8_1_mma, + ggml_cuda_mmq_write_back_mma); +// --------------------------------------------------------------------------------------------- + case GGML_TYPE_IQ1_S: + return ggml_cuda_mmq_util_funcs( + -1, + ggml_cuda_mmq_load_tiles_iq1_s, + ggml_cuda_mmq_vec_dot_q8_1_q8_1_mma, + ggml_cuda_mmq_write_back_mma); + case GGML_TYPE_IQ2_XXS: + return ggml_cuda_mmq_util_funcs( + -1, + ggml_cuda_mmq_load_tiles_iq2_xxs, + ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma, + ggml_cuda_mmq_write_back_mma); + case GGML_TYPE_IQ2_XS: + return ggml_cuda_mmq_util_funcs( + -1, + ggml_cuda_mmq_load_tiles_iq2_xs, + ggml_cuda_mmq_vec_dot_q8_0_16_q8_1_mma, + ggml_cuda_mmq_write_back_mma); + case GGML_TYPE_IQ2_S: + return ggml_cuda_mmq_util_funcs( + -1, + ggml_cuda_mmq_load_tiles_iq2_s, + ggml_cuda_mmq_vec_dot_q8_0_16_q8_1_mma, + ggml_cuda_mmq_write_back_mma); + case GGML_TYPE_IQ3_XXS: + return ggml_cuda_mmq_util_funcs( + -1, + ggml_cuda_mmq_load_tiles_iq3_xxs, + ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma, + ggml_cuda_mmq_write_back_mma); + case GGML_TYPE_IQ3_S: + return ggml_cuda_mmq_util_funcs( + -1, + ggml_cuda_mmq_load_tiles_iq3_s, + ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma, + ggml_cuda_mmq_write_back_mma); + case GGML_TYPE_IQ4_XS: + return ggml_cuda_mmq_util_funcs( + -1, + ggml_cuda_mmq_load_tiles_iq4_xs, + ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma, + ggml_cuda_mmq_write_back_mma); + case GGML_TYPE_IQ4_NL: + return ggml_cuda_mmq_util_funcs( + -1, + ggml_cuda_mmq_load_tiles_iq4_nl, + ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma, + ggml_cuda_mmq_write_back_mma); +// --------------------------------------------------------------------------------------------- + case GGML_TYPE_MXFP4: + return ggml_cuda_mmq_util_funcs( + -1, + ggml_cuda_mmq_load_tiles_mxfp4, + ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma, + ggml_cuda_mmq_write_back_mma); + case GGML_TYPE_NVFP4: + return ggml_cuda_mmq_util_funcs( + -1, + ggml_cuda_mmq_load_tiles_nvfp4, + ggml_cuda_mmq_vec_dot_q8_0_16_q8_1_mma, + ggml_cuda_mmq_write_back_mma); + default: + return ggml_cuda_mmq_util_funcs(1, nullptr, nullptr, nullptr); + } +} + +template +static constexpr __device__ int ggml_cuda_mmq_get_vdr() { + return ggml_cuda_mmq_get_util_funcs().vdr; +} + +template +static constexpr __device__ ggml_cuda_mmq_load_tiles_t ggml_cuda_mmq_get_load_tiles() { + return ggml_cuda_mmq_get_util_funcs().load_tiles; +} + +template +static constexpr __device__ ggml_cuda_mmq_vec_dot_t ggml_cuda_mmq_get_vec_dot() { + return ggml_cuda_mmq_get_util_funcs().vec_dot; +} + +template +static constexpr __device__ ggml_cuda_mmq_write_back_t ggml_cuda_mmq_get_write_back() { + return ggml_cuda_mmq_get_util_funcs().write_back; +} + +// --------------------------------------------------------------------------------------------- + +template +static __device__ __forceinline__ void mul_mat_q_process_tile( + const char * __restrict__ x, const int offset_x, const int * __restrict__ y, + const int * __restrict__ ids_dst, dst_t * __restrict__ dst, float * __restrict__ tmp_fixup, + const float * __restrict__ y_scale, + const int stride_row_x, const int ncols_y, const int stride_col_dst, + const int tile_x_max_i, const int tile_y_max_j, const int kb0_start, const int kb0_stop) { + + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback, prec_src1) / warp_size; + constexpr int qk = ggml_cuda_type_traits::qk; + constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback, prec_src1); + constexpr ggml_cuda_mmq_load_tiles_t load_tiles = ggml_cuda_mmq_get_load_tiles(); + constexpr ggml_cuda_mmq_vec_dot_t vec_dot = ggml_cuda_mmq_get_vec_dot(); + constexpr ggml_cuda_mmq_write_back_t write_back = ggml_cuda_mmq_get_write_back(); + + extern __shared__ int data_mul_mat_q[]; + int * tile_y = data_mul_mat_q + J; + int * tile_x = tile_y + GGML_PAD(J*MMQ_TILE_Y_K, nwarps*warp_size); + +#if defined(BLACKWELL_MMA_AVAILABLE) + // FP4 tile stores 8 blocks. src1 above Q4 uses the generic + // Q8_1 tile layout instead of the packed FP4 tile. + constexpr int ne_block = ((type == GGML_TYPE_MXFP4 || type == GGML_TYPE_NVFP4) && prec_src1 == GGML_PREC_Q4) ? + QK_FP4_MMQ : QK8_1_MMQ; +#else + constexpr int ne_block = QK8_1_MMQ; +#endif // defined(BLACKWELL_MMA_AVAILABLE) + + constexpr int ITER_K = ggml_cuda_mmq_get_K_vram(type, J, fallback, prec_src1); + constexpr int blocks_per_iter = ITER_K / qk; + + float sum[J*I / (nwarps*warp_size)] = {0.0f}; + + constexpr int sz = sizeof(block_q8_1_mmq) / sizeof(int); + + for (int kb0 = kb0_start; kb0 < kb0_stop; kb0 += blocks_per_iter) { + load_tiles(x, tile_x, offset_x + kb0, tile_x_max_i, stride_row_x); + { + const int * by0 = y + ncols_y * (kb0 * qk / ne_block) * sz; +#pragma unroll + for (int l0 = 0; l0 < J * MMQ_TILE_Y_K; l0 += nwarps * warp_size) { + int l = l0 + threadIdx.y*warp_size + threadIdx.x; + + tile_y[l] = by0[l]; + } + } + + __syncthreads(); + + vec_dot(tile_x, tile_y, sum, 0); + + __syncthreads(); + + { + const int * by0 = y + ncols_y * ((kb0 * qk / ne_block) * sz + sz); +#pragma unroll + for (int l0 = 0; l0 < J * MMQ_TILE_Y_K; l0 += nwarps * warp_size) { + int l = l0 + threadIdx.y*warp_size + threadIdx.x; + + tile_y[l] = by0[l]; + } + } + + __syncthreads(); + + vec_dot(tile_x, tile_y, sum, MMQ_TILE_NE_K); + + __syncthreads(); + } + + if (fixup) { + write_back(sum, ids_dst, tmp_fixup + blockIdx.x*(J*I), y_scale, I, I, J); + } else if constexpr (std::is_same_v) { + write_back(sum, ids_dst, dst, y_scale, stride_col_dst, tile_x_max_i, tile_y_max_j); + } else if constexpr (ggml_cuda_mmq_get_config(type, J, fallback, prec_src1).use_mma_data_layout()) { + ggml_cuda_mmq_write_back_mma(sum, ids_dst, dst, y_scale, stride_col_dst, tile_x_max_i, tile_y_max_j); + } else { + ggml_cuda_mmq_write_back_dp4a(sum, ids_dst, dst, y_scale, stride_col_dst, tile_x_max_i, tile_y_max_j); + } +} + + +// A flat MoE grid runs only the (expert, column tile) pairs listed in tile_map, see ggml_cuda_launch_mm_ids_tile_map. +static __device__ __forceinline__ void mmq_flat_tile(const int2 * __restrict__ tile_map, int & zt, int & jt) { + if (tile_map) { + const int2 t = tile_map[jt]; + zt = t.x; + jt = t.y; + } +} + +// The mul_mat_q kernel implements "stream-k" work partitioning as described in https://arxiv.org/abs/2301.03598 + +template +__launch_bounds__(ggml_cuda_mmq_get_nthreads(type, J, fallback, prec_src1), ggml_cuda_mmq_get_occupancy(type, J, fallback, prec_src1)) +static __global__ void mul_mat_q( + const char * __restrict__ x, const int * __restrict__ y, const int32_t * __restrict__ ids_dst, + const int32_t * __restrict__ expert_bounds, const int2 * __restrict__ tile_map, dst_t * __restrict__ dst, float * __restrict__ tmp_fixup, + const float * __restrict__ y_scale, + const uint3 blocks_per_ne00, const int nrows_x, const int ncols_dst, const int stride_row_x, const int ncols_y, const int stride_col_dst, + const uint3 channel_ratio, const uint3 nchannels_y, const int stride_channel_x, const int stride_channel_y, const int stride_channel_dst, + const uint3 sample_ratio, const uint3 nsamples_y, const int stride_sample_x, const int stride_sample_y, const int stride_sample_dst, + const uint3 ntx) { + + // Skip unused template specializations for faster compilation: + if (ggml_cuda_mmq_get_config(type, J, fallback, prec_src1).type == GGML_TYPE_COUNT) { + NO_DEVICE_CODE; + return; + } + + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback, prec_src1) / warp_size; + constexpr int qk = ggml_cuda_type_traits::qk; + constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback, prec_src1); + + const uint32_t nty = (nrows_x + I - 1) / I; // Number of tiles y + + // Initialize the ids for writing back data with just the index. + // For regular matrix multiplications this is never changed. + // For MoE the correct indices are loaded from ids_dst. + extern __shared__ int ids_dst_shared[]; // Stored at beginning of shared memory. +#pragma unroll + for (int j0 = 0; j0 < J; j0 += nwarps*warp_size) { + const int j = j0 + threadIdx.y*warp_size + threadIdx.x; + + if (j0 + nwarps*warp_size > J && j >= J) { + break; + } + + ids_dst_shared[j] = j; + } + __syncthreads(); + + if constexpr (!ggml_cuda_mmq_get_stream_k(type, J, fallback, prec_src1)) { + const uint2 tmp2 = fast_div_modulo(blockIdx.z, nchannels_y); + const int wt = tmp2.x; + int zt = tmp2.y; + int jt = blockIdx.y; + const int it = blockIdx.x; + mmq_flat_tile(tile_map, zt, jt); + + // Defaults for regular matrix multiplication: + int col_low = 0; + int col_high = ncols_dst; + int col_diff = ncols_dst; + int offset_y = wt*stride_sample_y + zt*stride_channel_y; + int offset_dst = wt*stride_sample_dst + zt*stride_channel_dst + jt*J*stride_col_dst; + int offset_y_scale; + if constexpr (type == GGML_TYPE_NVFP4) { + offset_y_scale = wt*nchannels_y.z*ncols_y + zt*ncols_y; + } else { + GGML_UNUSED(offset_y_scale); + } + + if (ids_dst) { + col_low = expert_bounds[zt + 0]; + col_high = expert_bounds[zt + 1]; + col_diff = col_high - col_low; + + offset_y = 0; + offset_dst = 0; + if constexpr (type == GGML_TYPE_NVFP4) { + offset_y_scale = 0; + } + + if (jt*J >= col_diff) { + return; + } + + // __syncthreads(); // There is no previous tile that could cause a race condition. +#pragma unroll + for (int j0 = 0; j0 < J; j0 += nwarps*warp_size) { + const int j = j0 + threadIdx.y*warp_size + threadIdx.x; + + if (j0 + nwarps*warp_size > J && j >= J) { + break; + } + + ids_dst_shared[j] = ids_dst[col_low + jt*J + j]; + } + __syncthreads(); + } + + offset_y += (col_low + jt*J)*(sizeof(block_q8_1_mmq)/sizeof(int)); + offset_dst += it*I; + const float * y_scale_tile = nullptr; + if constexpr (type == GGML_TYPE_NVFP4) { + offset_y_scale += col_low + jt*J; + y_scale_tile = y_scale ? y_scale + offset_y_scale : nullptr; + } + + const int tile_x_max_i = nrows_x - it*I - 1; + const int tile_y_max_j = col_diff - jt*J - 1; + + const int offset_x = fastdiv(wt, sample_ratio)*stride_sample_x + fastdiv(zt, channel_ratio)*stride_channel_x + it*I*stride_row_x; + + constexpr bool fixup = false; + mul_mat_q_process_tile + (x, offset_x, y + offset_y, ids_dst_shared, dst + offset_dst, tmp_fixup, y_scale_tile, + stride_row_x, ncols_y, stride_col_dst, + tile_x_max_i, tile_y_max_j, 0, blocks_per_ne00.z); + return; + } + + constexpr int ITER_K = ggml_cuda_mmq_get_K_vram(type, J, fallback, prec_src1); + constexpr int blocks_per_iter = ITER_K / qk; + + // kbc == k block continuous, current index in continuous ijk space. + int kbc = int64_t(blockIdx.x) *(nsamples_y.z*nchannels_y.z*ntx.z*nty*blocks_per_ne00.z) / gridDim.x; + int kbc_stop = int64_t(blockIdx.x + 1)*(nsamples_y.z*nchannels_y.z*ntx.z*nty*blocks_per_ne00.z) / gridDim.x; + + kbc -= fastmodulo(kbc, blocks_per_ne00) % blocks_per_iter; + kbc_stop -= fastmodulo(kbc_stop, blocks_per_ne00) % blocks_per_iter; + + // kb0 == k index when doing the matrix multiplication for an output tile. + int kb0_start = fastmodulo(kbc, blocks_per_ne00); + int kb0_stop = min(blocks_per_ne00.z, uint32_t(kb0_start + kbc_stop - kbc)); + while (kbc < kbc_stop && kb0_stop == int(blocks_per_ne00.z)) { + int tmp = fastdiv(kbc, blocks_per_ne00); + uint2 tmp2 = fast_div_modulo(tmp, ntx); + int jt = tmp2.y; + tmp = tmp2.x; + tmp2 = fast_div_modulo(tmp, nchannels_y); + int zt = tmp2.y; + tmp = tmp2.x; + tmp2 = fast_div_modulo(tmp, nsamples_y); + const int wt = tmp2.y; + const int it = tmp2.x; + mmq_flat_tile(tile_map, zt, jt); + + // Defaults for regular matrix multiplication: + int col_low = 0; + int col_high = ncols_dst; + int col_diff = ncols_dst; + int offset_y = wt*stride_sample_y + zt*stride_channel_y; + int offset_dst = wt*stride_sample_dst + zt*stride_channel_dst + jt*J*stride_col_dst; + int offset_y_scale; + if constexpr (type == GGML_TYPE_NVFP4) { + offset_y_scale = wt*nchannels_y.z*ncols_y + zt*ncols_y; + } else { + GGML_UNUSED(offset_y_scale); + } + + if (ids_dst) { + col_low = expert_bounds[zt + 0]; + col_high = expert_bounds[zt + 1]; + col_diff = col_high - col_low; + + offset_y = 0; + offset_dst = 0; + if constexpr (type == GGML_TYPE_NVFP4) { + offset_y_scale = 0; + } + + if (jt*J >= col_diff) { + kbc += blocks_per_ne00.z; + kbc -= fastmodulo(kbc, blocks_per_ne00); + + kb0_start = 0; + kb0_stop = min(blocks_per_ne00.z, uint32_t(kbc_stop - kbc)); + + continue; + } + + __syncthreads(); +#pragma unroll + for (int j0 = 0; j0 < J; j0 += nwarps*warp_size) { + const int j = j0 + threadIdx.y*warp_size + threadIdx.x; + + if (j0 + nwarps*warp_size > J && j >= J) { + break; + } + + ids_dst_shared[j] = ids_dst[col_low + jt*J + j]; + } + __syncthreads(); + } + + offset_y += (col_low + jt * J) * (sizeof(block_q8_1_mmq) / sizeof(int)); + offset_dst += it*I; + const float * y_scale_tile = nullptr; + if constexpr (type == GGML_TYPE_NVFP4) { + offset_y_scale += col_low + jt * J; + y_scale_tile = y_scale ? y_scale + offset_y_scale : nullptr; + } + + const int tile_x_max_i = nrows_x - it*I - 1; + const int tile_y_max_j = col_diff - jt*J - 1; + + const int offset_x = fastdiv(wt, sample_ratio)*stride_sample_x + fastdiv(zt, channel_ratio)*stride_channel_x + it*I*stride_row_x; + + constexpr bool fixup = false; // All but (potentially) the last iterations write their data to dst rather than the fixup buffer. + mul_mat_q_process_tile + (x, offset_x, y + offset_y, ids_dst_shared, dst + offset_dst, tmp_fixup, y_scale_tile, + stride_row_x, ncols_y, stride_col_dst, + tile_x_max_i, tile_y_max_j, kb0_start, kb0_stop); + + kbc += blocks_per_ne00.z; + kbc -= fastmodulo(kbc, blocks_per_ne00); + + kb0_start = 0; + kb0_stop = min(blocks_per_ne00.z, uint32_t(kbc_stop - kbc)); + } + + if (kbc >= kbc_stop) { + return; + } + + int tmp = fastdiv(kbc, blocks_per_ne00); + uint2 tmp2 = fast_div_modulo(tmp, ntx); + int jt = tmp2.y; + tmp = tmp2.x; + tmp2 = fast_div_modulo(tmp, nchannels_y); + int zt = tmp2.y; + tmp = tmp2.x; + tmp2 = fast_div_modulo(tmp, nsamples_y); + const int wt = tmp2.y; + const int it = tmp2.x; + mmq_flat_tile(tile_map, zt, jt); + + // Defaults for regular matrix multiplication: + int col_low = 0; + int col_high = ncols_dst; + int col_diff = ncols_dst; + int offset_y = wt*stride_sample_y + zt*stride_channel_y; + int offset_dst = wt*stride_sample_dst + zt*stride_channel_dst + jt*J*stride_col_dst; + int offset_y_scale; + if constexpr (type == GGML_TYPE_NVFP4) { + offset_y_scale = wt*nchannels_y.z*ncols_y + zt*ncols_y; + } else { + GGML_UNUSED(offset_y_scale); + } + + if (ids_dst) { + col_low = expert_bounds[zt + 0]; + col_high = expert_bounds[zt + 1]; + col_diff = col_high - col_low; + + offset_y = 0; + offset_dst = 0; + if constexpr (type == GGML_TYPE_NVFP4) { + offset_y_scale = 0; + } + + if (jt*J >= col_diff) { + return; + } + + // The memory layout for the fixup buffer is always contiguous, therefore reset ids: + __syncthreads(); +#pragma unroll + for (int j0 = 0; j0 < J; j0 += nwarps*warp_size) { + const int j = j0 + threadIdx.y*warp_size + threadIdx.x; + + if (j0 + nwarps*warp_size > J && j >= J) { + break; + } + + ids_dst_shared[j] = j; + } + __syncthreads(); + } + + offset_y += (col_low + jt * J) * (sizeof(block_q8_1_mmq) / sizeof(int)); + offset_dst += it*I; + const float * y_scale_tile = nullptr; + if constexpr (type == GGML_TYPE_NVFP4) { + offset_y_scale += col_low + jt * J; + y_scale_tile = y_scale ? y_scale + offset_y_scale : nullptr; + } + + const int tile_x_max_i = nrows_x - it*I - 1; + const int tile_y_max_j = col_diff - jt*J - 1; + + const int offset_x = fastdiv(wt, sample_ratio)*stride_sample_x + fastdiv(zt, channel_ratio)*stride_channel_x + it*I*stride_row_x; + + constexpr bool fixup = true; // Last index writes its data to fixup buffer to avoid data races with other blocks. + mul_mat_q_process_tile + (x, offset_x, y + offset_y, ids_dst_shared, dst + offset_dst, tmp_fixup, y_scale_tile, + stride_row_x, ncols_y, stride_col_dst, + tile_x_max_i, tile_y_max_j, kb0_start, kb0_stop); +} + +template +__launch_bounds__(ggml_cuda_mmq_get_nthreads(type, J, fallback, prec_src1)/2, 1) +static __global__ void mul_mat_q_stream_k_fixup( + const int32_t * __restrict__ ids_dst, const int32_t * __restrict__ expert_bounds, float * __restrict__ dst, + float * __restrict__ tmp_last_tile, const uint3 blocks_per_ne00, const int nrows_x, const int ncols_dst, + const int stride_col_dst, const uint3 nchannels_y, const int stride_channel_dst, const uint3 nsamples_y, + const int stride_sample_dst, const uint3 ntx) { + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + constexpr int nwarps = (ggml_cuda_mmq_get_nthreads(type, J, fallback, prec_src1) / 2) / warp_size; + constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback, prec_src1); + constexpr int qk = ggml_cuda_type_traits::qk; + constexpr int ITER_K = ggml_cuda_mmq_get_K_vram(type, J, fallback, prec_src1); + constexpr int blocks_per_iter = ITER_K / qk; + + float sum[J / nwarps] = {0.0f}; + const int i = blockIdx.y*warp_size + threadIdx.x; + + const int nty = (nrows_x + I - 1) / I; + + const int bidx0 = blockIdx.x; + + // kbc == k block continuous, current index in continuous ijk space. + int kbc0 = int64_t(blockIdx.x) *(nsamples_y.z*nchannels_y.z*ntx.z*nty*blocks_per_ne00.z) / gridDim.x; + int kbc0_stop = int64_t(blockIdx.x + 1)*(nsamples_y.z*nchannels_y.z*ntx.z*nty*blocks_per_ne00.z) / gridDim.x; + + kbc0 -= fastmodulo(kbc0, blocks_per_ne00) % blocks_per_iter; + kbc0_stop -= fastmodulo(kbc0_stop, blocks_per_ne00) % blocks_per_iter; + + const bool did_not_have_any_data = kbc0 == kbc0_stop; + const bool wrote_beginning_of_tile = fastmodulo(kbc0, blocks_per_ne00) == 0; + const bool did_not_write_last = fastdiv(kbc0, blocks_per_ne00) == fastdiv(kbc0_stop, blocks_per_ne00) && fastmodulo(kbc0_stop, blocks_per_ne00) != 0; + if (did_not_have_any_data || wrote_beginning_of_tile || did_not_write_last) { + return; + } + + bool any_fixup = false; + + // Iterate over previous blocks and sum up partial sums written to fixup buffer. + // All CUDA blocks that get here must have a previous block that needs a fixup. + int bidx = bidx0 - 1; + int kbc_stop = kbc0; + while(true) { + int kbc = int64_t(bidx)*(nsamples_y.z*nchannels_y.z*ntx.z*nty*blocks_per_ne00.z) / gridDim.x; + kbc -= fastmodulo(kbc, blocks_per_ne00) % blocks_per_iter; + + if (kbc == kbc_stop) { // Did not have any data. + bidx--; + kbc_stop = kbc; + continue; + } + + any_fixup = true; + + +#pragma unroll + for (int j0 = 0; j0 < J; j0 += nwarps) { + const int j = j0 + threadIdx.y; + + sum[j0/nwarps] += tmp_last_tile[bidx*(J*I) + j*I + i]; + } + + // If this block started in a previous tile we are done and don't need to combine additional partial results. + if (fastmodulo(kbc, blocks_per_ne00) == 0 || fastdiv(kbc, blocks_per_ne00) < fastdiv(kbc0, blocks_per_ne00)) { + break; + } + bidx--; + kbc_stop = kbc; + } + + if (!any_fixup) { + return; + } + + int tmp = fastdiv(kbc0, blocks_per_ne00); + uint2 tmp2 = fast_div_modulo(tmp, ntx); + const int jt = tmp2.y; + tmp = tmp2.x; + tmp2 = fast_div_modulo(tmp, nchannels_y); + const int zt = tmp2.y; + tmp = tmp2.x; + tmp2 = fast_div_modulo(tmp, nsamples_y); + const int wt = tmp2.y; + const int it = tmp2.x; + + if (!ids_dst) { + const int offset_dst = wt*stride_sample_dst + zt*stride_channel_dst + jt*J*stride_col_dst + it*I; + dst += offset_dst; + + const int i_max = nrows_x - it*I - 1; + const int j_max = ncols_dst - jt*J - 1; + if (fallback && i > i_max) { + return; + } + +#pragma unroll + for (int j0 = 0; j0 < J; j0 += nwarps) { + const int j = j0 + threadIdx.y; + + if (j > j_max) { + return; + } + + dst[j*stride_col_dst + i] += sum[j0/nwarps]; + } + return; + } + + __shared__ int ids_dst_shared[J]; + const int col_low = expert_bounds[zt + 0]; + const int col_high = expert_bounds[zt + 1]; + const int col_diff = col_high - col_low; + + for (int j = threadIdx.y*warp_size + threadIdx.x; j < J; j += nwarps*warp_size) { + ids_dst_shared[j] = ids_dst[col_low + jt*J + j]; + } + __syncthreads(); + + const int offset_dst = it*I; + dst += offset_dst; + + const int i_max = nrows_x - it*I - 1; + const int j_max = col_diff - jt*J - 1; + if (fallback && i > i_max) { + return; + } + +#pragma unroll + for (int j0 = 0; j0 < J; j0 += nwarps) { + const int j = j0 + threadIdx.y; + + if (j > j_max) { + return; + } + + dst[ids_dst_shared[j]*stride_col_dst + i] += sum[j0/nwarps]; + } +} + +// FreeToken: copies an fp32 result with dst's layout into dst +template +static __global__ void mmq_store_dst(const float * __restrict__ src, dst_t * __restrict__ dst, const int nrows, const int stride) { + const int64_t col = blockIdx.x; + const int row = blockIdx.y*blockDim.x + threadIdx.x; + if (row < nrows) { + dst[col*stride + row] = mmq_cast(src[col*stride + row]); + } +} + +struct mmq_args { + const char * x; ggml_type type_x; const int * y; const int32_t * ids_dst; const int32_t * expert_bounds; mmq_dst dst; + const float * y_scale; + int64_t ncols_x; int64_t nrows_x; int64_t ncols_dst; int64_t stride_row_x; int64_t ncols_y; int64_t nrows_dst; + int64_t nchannels_x; int64_t nchannels_y; int64_t stride_channel_x; int64_t stride_channel_y; int64_t stride_channel_dst; + int64_t nsamples_x; int64_t nsamples_y; int64_t stride_sample_x; int64_t stride_sample_y; int64_t stride_sample_dst; + int64_t ncols_max; + int J_best; // Tile width in ne11(dense)/ne12(MoE) direction to use for optimal performance. + const int2 * tile_map = nullptr; // Flat MoE grid: the (expert, column tile) of each of the ntx_flat grid columns. + int ntx_flat = 0; +}; + +static size_t mmq_get_nbytes_shared(const ggml_cuda_mmq_config & config, const int cc) { + const size_t nbs_ids = config.J*sizeof(int); + const size_t nbs_x = ggml_cuda_mmq_get_nbytes_shared_x(config, cc); + const size_t nbs_y = config.J * (sizeof(block_q8_1_mmq)); + return nbs_ids + nbs_x + GGML_PAD(nbs_y, config.nthreads*sizeof(int)); +} + +template +static void launch_mul_mat_q(ggml_backend_cuda_context & ctx, const mmq_args & args, cudaStream_t stream) { + const int id = ggml_cuda_get_device(); + const int cc = ggml_cuda_info().devices[id].cc; + const int nsm = ggml_cuda_info().devices[id].nsm; + const int warp_size = ggml_cuda_info().devices[id].warp_size; + + const ggml_cuda_mmq_config config = ggml_cuda_mmq_get_config(type, J, fallback, cc, prec_src1); + GGML_ASSERT(config.nthreads % warp_size == 0); + const int nwarps = config.nthreads / warp_size; + const int nbytes_shared = mmq_get_nbytes_shared(config, cc); + + const dim3 block_dims(warp_size, nwarps, 1); + + // a flat MoE grid puts all experts into one channel of ntx_flat column tiles + const int64_t nchannels_grid = args.tile_map ? 1 : args.nchannels_y; + const int nty = (args.nrows_x + config.I - 1) / config.I; + const int ntx = args.tile_map ? args.ntx_flat : (args.ncols_max + config.J - 1) / config.J; + const int ntzw = nchannels_grid * args.nsamples_y; + const dim3 block_nums_xy_tiling(nty, ntx, ntzw); + + GGML_ASSERT(args.nchannels_y % args.nchannels_x == 0); + GGML_ASSERT(args.nsamples_y % args.nsamples_x == 0); + const int channel_ratio = args.nchannels_y / args.nchannels_x; + const int sample_ratio = args.nsamples_y / args.nsamples_x; + + const uint3 blocks_per_ne00_fd = init_fastdiv_values(args.ncols_x / ggml_cuda_type_traits::qk); + const uint3 ntx_fd = init_fastdiv_values(ntx); + const uint3 nchannels_y_fd = init_fastdiv_values(nchannels_grid); + const uint3 nsamples_y_fd = init_fastdiv_values(args.nsamples_y); + const uint3 channel_ratio_fd = init_fastdiv_values(channel_ratio); + const uint3 sample_ratio_fd = init_fastdiv_values(sample_ratio); + + const auto run = [&](auto * dst, const dim3 & block_nums, float * tmp_fixup) { + using dst_t = std::remove_pointer_t; + CUDA_SET_SHARED_MEMORY_LIMIT((mul_mat_q), nbytes_shared); + CUDA_SET_SHARED_MEMORY_LIMIT((mul_mat_q), nbytes_shared); + mul_mat_q<<>> + (args.x, args.y, args.ids_dst, args.expert_bounds, args.tile_map, dst, tmp_fixup, args.y_scale, + blocks_per_ne00_fd, args.nrows_x, args.ncols_dst, args.stride_row_x, args.ncols_y, args.nrows_dst, + channel_ratio_fd, nchannels_y_fd, args.stride_channel_x, args.stride_channel_y, args.stride_channel_dst, + sample_ratio_fd, nsamples_y_fd, args.stride_sample_x, args.stride_sample_y, args.stride_sample_dst, + ntx_fd); + }; + + // FreeToken: the fixup kernel adds partial sums into an fp32 dst, so a bf16 dst that needs it is computed in fp32 first. + // An fp16 dst always is, unless GGML_MMQ_F16_DST builds the fp16 kernels. + const auto launch = [&](const dim3 & block_nums, const bool fixup, float * tmp_fixup) { + bool via_f32 = fixup && args.dst.type != GGML_TYPE_F32; +#ifndef GGML_MMQ_F16_DST + via_f32 = via_f32 || args.dst.type == GGML_TYPE_F16; +#endif // GGML_MMQ_F16_DST + ggml_cuda_pool_alloc dst_f32(ctx.pool(id)); + float * dst_float = args.dst.type == GGML_TYPE_F32 ? (float *) args.dst.p : nullptr; + if (via_f32) { + GGML_ASSERT(args.ids_dst || (args.nchannels_y == 1 && args.nsamples_y == 1)); + dst_float = dst_f32.alloc(args.ncols_dst * args.nrows_dst); + } + + if (dst_float) { + run(dst_float, block_nums, tmp_fixup); + } else if (args.dst.type == GGML_TYPE_BF16) { + run((nv_bfloat16 *) args.dst.p, block_nums, tmp_fixup); + } +#ifdef GGML_MMQ_F16_DST + else { + run((half *) args.dst.p, block_nums, tmp_fixup); + } +#endif // GGML_MMQ_F16_DST + + if (fixup) { + CUDA_CHECK(cudaGetLastError()); + const dim3 block_nums_fixup(block_nums.x, config.I/warp_size, 1); + const dim3 block_dims_fixup(block_dims.x, block_dims.y/2, block_dims.z); + mul_mat_q_stream_k_fixup<<>> + (args.ids_dst, args.expert_bounds, dst_float, tmp_fixup, blocks_per_ne00_fd, args.nrows_x, args.ncols_dst, + args.nrows_dst, nchannels_y_fd, args.stride_channel_dst, nsamples_y_fd, args.stride_sample_dst, + ntx_fd); + } + + if (via_f32) { + CUDA_CHECK(cudaGetLastError()); + const dim3 block_nums_store(args.ncols_dst, (args.nrows_x + 255) / 256, 1); + if (args.dst.type == GGML_TYPE_BF16) { + mmq_store_dst<<>>(dst_float, (nv_bfloat16 *) args.dst.p, args.nrows_x, args.nrows_dst); + } else { + mmq_store_dst<<>>(dst_float, (half *) args.dst.p, args.nrows_x, args.nrows_dst); + } + } + }; + + if (!config.stream_k) { + launch(block_nums_xy_tiling, false, nullptr); + return; + } + + // For the stream-k kernel it is possible to run it with tiling by setting the number of CUDA blocks equal to the number of tiles. + // This is worthwhile if the efficiency of tiling is high and skipping the fixup kernel is more important. + const int ntiles_dst = ntx * nty * ntzw; + const int tiles_nwaves = (ntiles_dst + nsm - 1) / nsm; + const int tiles_efficiency_percent = 100 * ntiles_dst / (nsm*tiles_nwaves); + // The fixup kernel does not read tile_map, so a flat MoE grid always runs one tile per block. + const dim3 block_nums_stream_k(args.tile_map || (GGML_CUDA_CC_IS_NVIDIA(cc) && tiles_efficiency_percent >= 90) ? ntiles_dst : nsm, 1, 1); + + GGML_ASSERT(ntiles_dst * blocks_per_ne00_fd.z < (1 << 30)); // Assert that variable kbc will not overflow. + + const bool fixup_needed = ntiles_dst % block_nums_stream_k.x != 0; + + ggml_cuda_pool & pool = ctx.pool(id); + ggml_cuda_pool_alloc tmp_fixup(pool); + if (fixup_needed) { + tmp_fixup.alloc(block_nums_stream_k.x * config.J*config.I); + } + + launch(block_nums_stream_k, fixup_needed, tmp_fixup.ptr); +} + +template +void mul_mat_q_switch_J(ggml_backend_cuda_context & ctx, const mmq_args & args, cudaStream_t stream) { + switch (args.J_best) { + case 8: + launch_mul_mat_q(ctx, args, stream); + break; + case 16: + launch_mul_mat_q(ctx, args, stream); + break; + case 24: + launch_mul_mat_q(ctx, args, stream); + break; + case 32: + launch_mul_mat_q(ctx, args, stream); + break; + case 40: + launch_mul_mat_q(ctx, args, stream); + break; + case 48: + launch_mul_mat_q(ctx, args, stream); + break; + case 56: + launch_mul_mat_q(ctx, args, stream); + break; + case 64: + launch_mul_mat_q(ctx, args, stream); + break; + case 72: + launch_mul_mat_q(ctx, args, stream); + break; + case 80: + launch_mul_mat_q(ctx, args, stream); + break; + case 88: + launch_mul_mat_q(ctx, args, stream); + break; + case 96: + launch_mul_mat_q(ctx, args, stream); + break; + case 104: + launch_mul_mat_q(ctx, args, stream); + break; + case 112: + launch_mul_mat_q(ctx, args, stream); + break; + case 120: + launch_mul_mat_q(ctx, args, stream); + break; + case 128: + launch_mul_mat_q(ctx, args, stream); + break; + default: + fprintf(stderr, "J_best=%d\n", args.J_best); + GGML_ABORT("fatal error"); + } +} + +template +void mul_mat_q_case(ggml_backend_cuda_context & ctx, const mmq_args & args, cudaStream_t stream) { + if (ggml_cuda_mmq_needs_fallback(args.nrows_x)) { + constexpr bool fallback = true; + mul_mat_q_switch_J(ctx, args, stream); + } else { + constexpr bool fallback = false; + mul_mat_q_switch_J(ctx, args, stream); + } +} + +#define DECL_MMQ_CASE(type) \ + template void mul_mat_q_case(ggml_backend_cuda_context & ctx, const mmq_args & args, cudaStream_t stream) \ + +// FP4 variant: uses native FP4 MMA instead of keeping src1 at Q8_1. +#define DECL_MMQ_CASE_W4A4(type) \ + template void mul_mat_q_case(ggml_backend_cuda_context & ctx, const mmq_args & args, cudaStream_t stream) \ + +extern DECL_MMQ_CASE(GGML_TYPE_Q1_0); +extern DECL_MMQ_CASE(GGML_TYPE_Q2_0); +extern DECL_MMQ_CASE(GGML_TYPE_Q4_0); +extern DECL_MMQ_CASE(GGML_TYPE_Q4_1); +extern DECL_MMQ_CASE(GGML_TYPE_Q5_0); +extern DECL_MMQ_CASE(GGML_TYPE_Q5_1); +extern DECL_MMQ_CASE(GGML_TYPE_Q8_0); +// ----------------------------------------- +extern DECL_MMQ_CASE(GGML_TYPE_Q2_K); +extern DECL_MMQ_CASE(GGML_TYPE_Q3_K); +extern DECL_MMQ_CASE(GGML_TYPE_Q4_K); +extern DECL_MMQ_CASE(GGML_TYPE_Q5_K); +extern DECL_MMQ_CASE(GGML_TYPE_Q6_K); +// ----------------------------------------- +extern DECL_MMQ_CASE(GGML_TYPE_IQ1_S); +extern DECL_MMQ_CASE(GGML_TYPE_IQ2_XXS); +extern DECL_MMQ_CASE(GGML_TYPE_IQ2_XS); +extern DECL_MMQ_CASE(GGML_TYPE_IQ2_S); +extern DECL_MMQ_CASE(GGML_TYPE_IQ3_XXS); +extern DECL_MMQ_CASE(GGML_TYPE_IQ3_S); +extern DECL_MMQ_CASE(GGML_TYPE_IQ4_NL); +extern DECL_MMQ_CASE(GGML_TYPE_IQ4_XS); +// ----------------------------------------- +extern DECL_MMQ_CASE(GGML_TYPE_MXFP4); +extern DECL_MMQ_CASE(GGML_TYPE_NVFP4); +extern DECL_MMQ_CASE_W4A4(GGML_TYPE_MXFP4); +extern DECL_MMQ_CASE_W4A4(GGML_TYPE_NVFP4); + +// ------------------------------------------------------------------------------------------------------------------------- + +// FreeToken: under ids, ncols_max > 0 bounds the tokens any one expert gets and sizes the tile grid; otherwise the grid is flat, see ggml_cuda_launch_mm_ids_tile_map +void ggml_cuda_mul_mat_q( + ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, ggml_tensor * dst, + int64_t ncols_max = -1); + +bool ggml_cuda_should_use_mmq(enum ggml_type type, int cc, int64_t ne11, int64_t n_experts); diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmvq.cu b/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmvq.cu new file mode 100644 index 0000000000..8487a035c1 --- /dev/null +++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmvq.cu @@ -0,0 +1,1629 @@ +#include "mmvq.cuh" +#include "convert.cuh" +#include "quantize.cuh" +#include "unary.cuh" +#include "vecdotq.cuh" + +#include +#include + +// only enabled on DGX Spark, where it is a gain on every type below. On the higher-bandwidth parts the kernel +// has little exposed latency left to hide and the extra requests cost more than they save. +// For perf data, see https://github.com/ggml-org/llama.cpp/pull/26705#issuecomment-5569335031 +#if __CUDA_ARCH__ == GGML_CUDA_CC_DGX_SPARK +// returns true only for those quants that benefit from prefetch and false otherwise +static constexpr __host__ __device__ bool mmvq_should_prefetch(ggml_type type) { + switch (type) { + case GGML_TYPE_Q4_0: + case GGML_TYPE_Q5_0: + case GGML_TYPE_Q8_0: + case GGML_TYPE_MXFP4: + case GGML_TYPE_Q3_K: + case GGML_TYPE_Q4_K: + case GGML_TYPE_Q5_K: + case GGML_TYPE_Q6_K: + case GGML_TYPE_IQ1_M: + case GGML_TYPE_IQ4_NL: + case GGML_TYPE_IQ4_XS: + return true; + default: + return false; + } +} + +static __device__ __forceinline__ void mmvq_prefetch_l2(const void * p) { + asm volatile("prefetch.global.L2 [%0];" :: "l"(p)); +} +#endif + +typedef float (*vec_dot_q_cuda_t)(const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs); + +static constexpr __device__ vec_dot_q_cuda_t get_vec_dot_q_cuda(ggml_type type) { + switch (type) { + case GGML_TYPE_Q1_0: return vec_dot_q1_0_q8_1; + case GGML_TYPE_Q2_0: return vec_dot_q2_0_q8_1; + case GGML_TYPE_Q4_0: return vec_dot_q4_0_q8_1; + case GGML_TYPE_Q4_1: return vec_dot_q4_1_q8_1; + case GGML_TYPE_Q5_0: return vec_dot_q5_0_q8_1; + case GGML_TYPE_Q5_1: return vec_dot_q5_1_q8_1; + case GGML_TYPE_Q8_0: return vec_dot_q8_0_q8_1; + case GGML_TYPE_MXFP4: return vec_dot_mxfp4_q8_1; + case GGML_TYPE_NVFP4: return vec_dot_nvfp4_q8_1; + case GGML_TYPE_Q2_K: return vec_dot_q2_K_q8_1; + case GGML_TYPE_Q3_K: return vec_dot_q3_K_q8_1; + case GGML_TYPE_Q4_K: return vec_dot_q4_K_q8_1; + case GGML_TYPE_Q5_K: return vec_dot_q5_K_q8_1; + case GGML_TYPE_Q6_K: return vec_dot_q6_K_q8_1; + case GGML_TYPE_IQ2_XXS: return vec_dot_iq2_xxs_q8_1; + case GGML_TYPE_IQ2_XS: return vec_dot_iq2_xs_q8_1; + case GGML_TYPE_IQ2_S: return vec_dot_iq2_s_q8_1; + case GGML_TYPE_IQ3_XXS: return vec_dot_iq3_xxs_q8_1; + case GGML_TYPE_IQ1_S: return vec_dot_iq1_s_q8_1; + case GGML_TYPE_IQ1_M: return vec_dot_iq1_m_q8_1; + case GGML_TYPE_IQ4_NL: return vec_dot_iq4_nl_q8_1; + case GGML_TYPE_IQ4_XS: return vec_dot_iq4_xs_q8_1; + case GGML_TYPE_IQ3_S: return vec_dot_iq3_s_q8_1; + default: return nullptr; + } +} + +static constexpr __host__ __device__ int get_vdr_mmvq(ggml_type type) { + switch (type) { + case GGML_TYPE_Q1_0: return VDR_Q1_0_Q8_1_MMVQ; + case GGML_TYPE_Q2_0: return VDR_Q2_0_Q8_1_MMVQ; + case GGML_TYPE_Q4_0: return VDR_Q4_0_Q8_1_MMVQ; + case GGML_TYPE_Q4_1: return VDR_Q4_1_Q8_1_MMVQ; + case GGML_TYPE_Q5_0: return VDR_Q5_0_Q8_1_MMVQ; + case GGML_TYPE_Q5_1: return VDR_Q5_1_Q8_1_MMVQ; + case GGML_TYPE_Q8_0: return VDR_Q8_0_Q8_1_MMVQ; + case GGML_TYPE_MXFP4: return VDR_MXFP4_Q8_1_MMVQ; + case GGML_TYPE_NVFP4: return VDR_NVFP4_Q8_1_MMVQ; + case GGML_TYPE_Q2_K: return VDR_Q2_K_Q8_1_MMVQ; + case GGML_TYPE_Q3_K: return VDR_Q3_K_Q8_1_MMVQ; + case GGML_TYPE_Q4_K: return VDR_Q4_K_Q8_1_MMVQ; + case GGML_TYPE_Q5_K: return VDR_Q5_K_Q8_1_MMVQ; + case GGML_TYPE_Q6_K: return VDR_Q6_K_Q8_1_MMVQ; + case GGML_TYPE_IQ2_XXS: return VDR_IQ2_XXS_Q8_1_MMVQ; + case GGML_TYPE_IQ2_XS: return VDR_IQ2_XS_Q8_1_MMVQ; + case GGML_TYPE_IQ2_S: return VDR_IQ2_S_Q8_1_MMVQ; + case GGML_TYPE_IQ3_XXS: return VDR_IQ3_XXS_Q8_1_MMVQ; + case GGML_TYPE_IQ3_S: return VDR_IQ3_S_Q8_1_MMVQ; + case GGML_TYPE_IQ4_NL: return VDR_IQ4_NL_Q8_1_MMVQ; + case GGML_TYPE_IQ4_XS: return VDR_IQ4_XS_Q8_1_MMVQ; + default: return 1; + } +} + +enum mmvq_parameter_table_id { + MMVQ_PARAMETERS_GENERIC = 0, + MMVQ_PARAMETERS_TURING, + MMVQ_PARAMETERS_GCN, + MMVQ_PARAMETERS_RDNA2, + MMVQ_PARAMETERS_RDNA3_0, + MMVQ_PARAMETERS_RDNA4, + MMVQ_PARAMETERS_GB10 +}; + +static constexpr __device__ mmvq_parameter_table_id get_device_table_id() { +#if defined(RDNA4) + return MMVQ_PARAMETERS_RDNA4; +#elif defined(RDNA3_0) + return MMVQ_PARAMETERS_RDNA3_0; +#elif defined(RDNA2) || defined(RDNA3_5) + return MMVQ_PARAMETERS_RDNA2; +#elif defined(GCN) || defined(CDNA) + return MMVQ_PARAMETERS_GCN; +#elif __CUDA_ARCH__ >= GGML_CUDA_CC_VOLTA && __CUDA_ARCH__ < GGML_CUDA_CC_AMPERE + return MMVQ_PARAMETERS_TURING; +#elif __CUDA_ARCH__ == GGML_CUDA_CC_DGX_SPARK + return MMVQ_PARAMETERS_GB10; +#else + return MMVQ_PARAMETERS_GENERIC; +#endif +} + +static __host__ mmvq_parameter_table_id get_device_table_id(int cc) { + if (GGML_CUDA_CC_IS_RDNA4(cc)) { + return MMVQ_PARAMETERS_RDNA4; + } + if (GGML_CUDA_CC_IS_RDNA3_0(cc)) { + return MMVQ_PARAMETERS_RDNA3_0; + } + if (GGML_CUDA_CC_IS_RDNA2(cc) || GGML_CUDA_CC_IS_RDNA3_5(cc)) { + return MMVQ_PARAMETERS_RDNA2; + } + if (GGML_CUDA_CC_IS_GCN(cc) || GGML_CUDA_CC_IS_CDNA(cc)) { + return MMVQ_PARAMETERS_GCN; + } + if (GGML_CUDA_CC_IS_NVIDIA(cc) && ggml_cuda_highest_compiled_arch(cc) >= GGML_CUDA_CC_VOLTA && ggml_cuda_highest_compiled_arch(cc) < GGML_CUDA_CC_AMPERE) { + return MMVQ_PARAMETERS_TURING; + } + if (GGML_CUDA_CC_IS_NVIDIA(cc) && ggml_cuda_highest_compiled_arch(cc) == GGML_CUDA_CC_DGX_SPARK) { + return MMVQ_PARAMETERS_GB10; + } + return MMVQ_PARAMETERS_GENERIC; +} + +// Per-architecture maximum batch size for which MMVQ should be used for MUL_MAT_ID. +// Returns a value <= MMVQ_MAX_BATCH_SIZE. Default is MMVQ_MAX_BATCH_SIZE. +// Check https://github.com/ggml-org/llama.cpp/pull/20905#issuecomment-4145835627 for details + +static constexpr __host__ __device__ int get_mmvq_mmid_max_batch_pascal_older(ggml_type type) { + switch (type) { + case GGML_TYPE_IQ1_S: return 6; + case GGML_TYPE_IQ1_M: return 6; + case GGML_TYPE_IQ2_S: return 4; + case GGML_TYPE_IQ2_XS: return 5; + case GGML_TYPE_IQ2_XXS: return 5; + case GGML_TYPE_IQ3_S: return 4; + case GGML_TYPE_IQ3_XXS: return 4; + case GGML_TYPE_IQ4_NL: return 6; + case GGML_TYPE_IQ4_XS: return 5; + case GGML_TYPE_MXFP4: return 4; + case GGML_TYPE_NVFP4: return 4; + case GGML_TYPE_Q2_K: return 4; + case GGML_TYPE_Q3_K: return 4; + case GGML_TYPE_Q4_0: return 6; + case GGML_TYPE_Q4_1: return 6; + case GGML_TYPE_Q4_K: return 5; + case GGML_TYPE_Q5_0: return 6; + case GGML_TYPE_Q5_1: return 6; + case GGML_TYPE_Q5_K: return 5; + case GGML_TYPE_Q6_K: return 4; + case GGML_TYPE_Q8_0: return 4; + default: return MMVQ_MAX_BATCH_SIZE; + } +} + +static constexpr __host__ __device__ int get_mmvq_mmid_max_batch_turing_plus(ggml_type type) { + switch (type) { + case GGML_TYPE_IQ2_S: return 7; + case GGML_TYPE_IQ3_S: return 6; + case GGML_TYPE_IQ3_XXS: return 7; + case GGML_TYPE_MXFP4: return 7; + case GGML_TYPE_NVFP4: return 8; + case GGML_TYPE_Q2_K: return 7; + case GGML_TYPE_Q3_K: return 5; + default: return MMVQ_MAX_BATCH_SIZE; + } +} + +static constexpr __host__ __device__ int get_mmvq_mmid_max_batch_gcn(ggml_type type) { + switch (type) { + case GGML_TYPE_IQ1_S: return 5; + case GGML_TYPE_IQ1_M: return 5; + case GGML_TYPE_IQ2_S: return 4; + case GGML_TYPE_IQ2_XS: return 4; + case GGML_TYPE_IQ2_XXS: return 4; + case GGML_TYPE_IQ3_S: return 4; + case GGML_TYPE_IQ3_XXS: return 4; + case GGML_TYPE_IQ4_NL: return 6; + case GGML_TYPE_IQ4_XS: return 4; + case GGML_TYPE_Q2_K: return 4; + case GGML_TYPE_Q3_K: return 4; + case GGML_TYPE_Q4_0: return 5; + case GGML_TYPE_Q4_1: return 5; + case GGML_TYPE_Q4_K: return 4; + case GGML_TYPE_Q5_K: return 4; + case GGML_TYPE_Q6_K: return 4; + case GGML_TYPE_Q8_0: return 4; + default: return MMVQ_MAX_BATCH_SIZE; + } +} + +static constexpr __host__ __device__ int get_mmvq_mmid_max_batch_cdna(ggml_type type) { + switch (type) { + case GGML_TYPE_IQ2_S: return 5; + case GGML_TYPE_IQ2_XS: return 5; + case GGML_TYPE_IQ2_XXS: return 5; + case GGML_TYPE_IQ3_S: return 4; + case GGML_TYPE_IQ3_XXS: return 5; + default: return MMVQ_MAX_BATCH_SIZE; + } +} + +static constexpr __host__ __device__ int get_mmvq_mmid_max_batch_rdna1_rdna2(ggml_type type) { + switch (type) { + case GGML_TYPE_IQ2_S: return 4; + case GGML_TYPE_IQ2_XS: return 4; + case GGML_TYPE_IQ2_XXS: return 4; + case GGML_TYPE_IQ3_S: return 4; + case GGML_TYPE_IQ3_XXS: return 4; + case GGML_TYPE_Q2_K: return 7; + case GGML_TYPE_Q3_K: return 4; + case GGML_TYPE_Q4_K: return 5; + case GGML_TYPE_Q5_K: return 6; + case GGML_TYPE_Q6_K: return 5; + default: return MMVQ_MAX_BATCH_SIZE; + } +} + +static constexpr __host__ __device__ int get_mmvq_mmid_max_batch_rdna3(ggml_type type) { + switch (type) { + case GGML_TYPE_IQ1_S: return 6; + case GGML_TYPE_IQ1_M: return 6; + case GGML_TYPE_IQ2_S: return 4; + case GGML_TYPE_IQ2_XS: return 4; + case GGML_TYPE_IQ2_XXS: return 4; + case GGML_TYPE_IQ3_S: return 4; + case GGML_TYPE_IQ3_XXS: return 4; + case GGML_TYPE_IQ4_NL: return 6; + case GGML_TYPE_IQ4_XS: return 6; + case GGML_TYPE_Q4_K: return 4; + case GGML_TYPE_Q5_K: return 4; + case GGML_TYPE_Q6_K: return 4; + default: return MMVQ_MAX_BATCH_SIZE; + } +} + +static constexpr __host__ __device__ int get_mmvq_mmid_max_batch_rdna4(ggml_type type) { + switch (type) { + case GGML_TYPE_IQ1_S: return 7; + case GGML_TYPE_IQ1_M: return 7; + case GGML_TYPE_IQ2_S: return 4; + case GGML_TYPE_IQ2_XS: return 4; + case GGML_TYPE_IQ2_XXS: return 4; + case GGML_TYPE_IQ3_S: return 4; + case GGML_TYPE_IQ3_XXS: return 4; + case GGML_TYPE_IQ4_NL: return 7; + case GGML_TYPE_IQ4_XS: return 5; + case GGML_TYPE_MXFP4: return 5; + case GGML_TYPE_NVFP4: return 5; + case GGML_TYPE_Q3_K: return 4; + case GGML_TYPE_Q4_0: return 7; + case GGML_TYPE_Q4_1: return 7; + case GGML_TYPE_Q4_K: return 4; + case GGML_TYPE_Q5_0: return 7; + case GGML_TYPE_Q5_1: return 7; + case GGML_TYPE_Q5_K: return 5; + case GGML_TYPE_Q6_K: return 5; + case GGML_TYPE_Q8_0: return 7; + default: return MMVQ_MAX_BATCH_SIZE; + } +} + +// Host function: returns the max batch size for the current arch+type at runtime. +int get_mmvq_mmid_max_batch(ggml_type type, int cc) { + // NVIDIA: P100, Volta, Ada Lovelace, and Blackwell always use MMVQ for MUL_MAT_ID. + if (GGML_CUDA_CC_IS_NVIDIA(cc)) { + if (cc == GGML_CUDA_CC_PASCAL || cc == GGML_CUDA_CC_VOLTA || cc >= GGML_CUDA_CC_ADA_LOVELACE) { + return MMVQ_MAX_BATCH_SIZE; + } + if (cc >= GGML_CUDA_CC_TURING) { + return get_mmvq_mmid_max_batch_turing_plus(type); + } + return get_mmvq_mmid_max_batch_pascal_older(type); + } + + // AMD + if (GGML_CUDA_CC_IS_AMD(cc)) { + if (GGML_CUDA_CC_IS_RDNA4(cc)) { + return get_mmvq_mmid_max_batch_rdna4(type); + } + if (GGML_CUDA_CC_IS_RDNA3(cc)) { + return get_mmvq_mmid_max_batch_rdna3(type); + } + if (GGML_CUDA_CC_IS_RDNA1(cc) || GGML_CUDA_CC_IS_RDNA2(cc)) { + return get_mmvq_mmid_max_batch_rdna1_rdna2(type); + } + if (GGML_CUDA_CC_IS_CDNA(cc)) { + return get_mmvq_mmid_max_batch_cdna(type); + } + if (GGML_CUDA_CC_IS_GCN(cc)) { + return get_mmvq_mmid_max_batch_gcn(type); + } + } + return MMVQ_MAX_BATCH_SIZE; +} + +bool ggml_cuda_should_use_mmvq(enum ggml_type type, int cc, int64_t ne11) { + if (!ggml_is_quantized(type)) { + return false; + } + // k-quants cost more to decode and mvq redoes that per column, so MMQ wins sooner. + // Only list quant-types MMQ supports, others would fall back to cuBLAS. + if (GGML_CUDA_CC_IS_NVIDIA(cc) && cc == GGML_CUDA_CC_ADA_LOVELACE) { + switch (type) { // tuned on RTX 4090 + case GGML_TYPE_Q2_K: + return ne11 <= 4; + case GGML_TYPE_Q3_K: + return ne11 <= 6; + default: + return ne11 <= MMVQ_MAX_BATCH_SIZE; + } + } + if (GGML_CUDA_CC_IS_NVIDIA(cc) && cc == GGML_CUDA_CC_BLACKWELL) { + switch (type) { // tuned on RTX 5090 + case GGML_TYPE_Q2_K: + case GGML_TYPE_Q3_K: + case GGML_TYPE_Q4_K: + return ne11 <= 5; + case GGML_TYPE_Q5_K: + return ne11 <= 6; + case GGML_TYPE_Q6_K: + return ne11 <= 7; + default: + return ne11 <= MMVQ_MAX_BATCH_SIZE; + } + } + if (GGML_CUDA_CC_IS_NVIDIA(cc) && cc == GGML_CUDA_CC_DGX_SPARK) { + switch (type) { // tuned on DGX Spark GB10 + case GGML_TYPE_Q2_K: + return ne11 <= 6; + default: + return ne11 <= MMVQ_MAX_BATCH_SIZE; + } + } + if (GGML_CUDA_CC_IS_NVIDIA(cc) && cc == GGML_CUDA_CC_ORIN) { + switch (type) { // tuned for Jetson Orin + case GGML_TYPE_Q2_K: + case GGML_TYPE_Q3_K: + case GGML_TYPE_Q4_K: + case GGML_TYPE_Q5_K: + case GGML_TYPE_Q6_K: + return ne11 <= 1; + default: + return ne11 <= MMVQ_MAX_BATCH_SIZE; + } + } + if (GGML_CUDA_CC_IS_NVIDIA(cc) && cc == GGML_CUDA_CC_VOLTA) { + switch (type) { + case GGML_TYPE_Q2_K: + return ne11 <= 4; + case GGML_TYPE_Q3_K: + return ne11 <= 6; + case GGML_TYPE_Q4_K: + return ne11 <= 5; + case GGML_TYPE_Q5_K: + return ne11 <= 6; + case GGML_TYPE_Q6_K: + return ne11 <= 7; + default: + return ne11 <= MMVQ_MAX_BATCH_SIZE; + } + } + if (GGML_CUDA_CC_IS_CDNA(cc)) { + if (GGML_CUDA_CC_IS_CDNA1(cc)) { + switch (type) { + case GGML_TYPE_Q4_0: + case GGML_TYPE_Q4_1: + return ne11 <= 7; + case GGML_TYPE_Q5_1: + return ne11 <= 7; + case GGML_TYPE_Q8_0: + return ne11 <= 6; + case GGML_TYPE_Q2_K: + return ne11 <= 4; + case GGML_TYPE_Q3_K: + return ne11 <= 3; + case GGML_TYPE_Q4_K: + return ne11 <= 2; + case GGML_TYPE_Q5_K: + return ne11 <= 3; + case GGML_TYPE_Q6_K: + return ne11 <= 4; + case GGML_TYPE_IQ1_S: + return ne11 <= 5; + case GGML_TYPE_IQ2_XXS: + case GGML_TYPE_IQ3_S: + case GGML_TYPE_IQ4_XS: + return ne11 <= 6; + default: + return ne11 <= MMVQ_MAX_BATCH_SIZE; + } + } + switch (type) { // tuned for CDNA2 + case GGML_TYPE_Q2_K: + return ne11 <= 5; + case GGML_TYPE_Q3_K: + case GGML_TYPE_Q4_K: + case GGML_TYPE_Q5_K: + return ne11 <= 3; + case GGML_TYPE_Q6_K: + return ne11 <= 5; + default: + return ne11 <= MMVQ_MAX_BATCH_SIZE; + } + } + return ne11 <= MMVQ_MAX_BATCH_SIZE; +} + +// Device constexpr: returns the max batch size for the current arch+type at compile time. +template +static constexpr __device__ int get_mmvq_mmid_max_batch_for_device() { +#if defined(RDNA4) + return get_mmvq_mmid_max_batch_rdna4(type); +#elif defined(RDNA3) + return get_mmvq_mmid_max_batch_rdna3(type); +#elif defined(RDNA2) || defined(RDNA1) + return get_mmvq_mmid_max_batch_rdna1_rdna2(type); +#elif defined(CDNA) + return get_mmvq_mmid_max_batch_cdna(type); +#elif defined(GCN) + return get_mmvq_mmid_max_batch_gcn(type); +#elif !defined(GGML_USE_MUSA) && (__CUDA_ARCH__ == GGML_CUDA_CC_PASCAL || __CUDA_ARCH__ == GGML_CUDA_CC_VOLTA || __CUDA_ARCH__ >= GGML_CUDA_CC_ADA_LOVELACE) + return MMVQ_MAX_BATCH_SIZE; +#elif !defined(GGML_USE_MUSA) && __CUDA_ARCH__ >= GGML_CUDA_CC_TURING + return get_mmvq_mmid_max_batch_turing_plus(type); +#else + return get_mmvq_mmid_max_batch_pascal_older(type); +#endif +} + +static constexpr __host__ __device__ int calc_nwarps(ggml_type type, int ncols_dst, mmvq_parameter_table_id table_id, bool small_k = false, bool halve_iters = false) { + if (table_id == MMVQ_PARAMETERS_GENERIC) { + switch (ncols_dst) { + case 1: + case 2: + case 3: + case 4: + return 4; + case 5: + case 6: + case 7: + case 8: + return 2; + default: + return 1; + } + } else if (table_id == MMVQ_PARAMETERS_GCN) { + switch (ncols_dst) { + case 1: + case 2: + case 3: + case 4: + return 2; + case 5: + case 6: + case 7: + case 8: + default: + return 1; + } + } + if (table_id == MMVQ_PARAMETERS_RDNA4) { + // nwarps=8 benefits types with simple vec_dot on RDNA4 (ncols_dst=1). + // Types with complex vec_dot (Q3_K, IQ2_*, IQ3_*) regress due to register + // pressure and lookup table contention at higher thread counts. + if (ncols_dst == 1) { + switch (type) { + case GGML_TYPE_Q4_0: + case GGML_TYPE_Q4_1: + case GGML_TYPE_Q5_0: + case GGML_TYPE_Q5_1: + case GGML_TYPE_Q8_0: + case GGML_TYPE_Q2_K: + case GGML_TYPE_Q4_K: + case GGML_TYPE_Q5_K: + case GGML_TYPE_Q6_K: + case GGML_TYPE_IQ4_NL: + case GGML_TYPE_IQ4_XS: + return 8; + default: + return 1; + } + } + return 1; + } + if (table_id == MMVQ_PARAMETERS_RDNA3_0) { + // RDNA3 (W7900): stricter whitelist than RDNA4. + // Q2_K / Q5_K / IQ4_XS regress in full quant sweeps. + if (ncols_dst == 1) { + switch (type) { + case GGML_TYPE_Q4_0: + case GGML_TYPE_Q4_1: + case GGML_TYPE_Q5_0: + case GGML_TYPE_Q5_1: + case GGML_TYPE_Q8_0: + return 8; + case GGML_TYPE_Q6_K: + return 2; + case GGML_TYPE_IQ4_NL: + return 8; + default: + return 1; + } + } + return 1; + } + if (table_id == MMVQ_PARAMETERS_TURING) { + if (ncols_dst == 1) { + switch (type) { + case GGML_TYPE_Q2_K: + case GGML_TYPE_Q3_K: + case GGML_TYPE_Q4_K: + case GGML_TYPE_Q5_K: + case GGML_TYPE_Q6_K: + return 2; + default: + return 4; + } + } + switch (ncols_dst) { + case 2: + case 3: + case 4: + return 4; + case 5: + case 6: + case 7: + case 8: + return 2; + default: + return 1; + } + } + if (table_id == MMVQ_PARAMETERS_GB10) { + const int generic = calc_nwarps(type, ncols_dst, MMVQ_PARAMETERS_GENERIC); + // Only worth the wider block when it actually retires the K loop in half the trips (Observation) + if (ncols_dst == 1 && !small_k && halve_iters) { + switch (type) { + case GGML_TYPE_Q4_0: + case GGML_TYPE_Q4_1: + case GGML_TYPE_Q5_0: + case GGML_TYPE_Q5_1: + case GGML_TYPE_Q8_0: + case GGML_TYPE_Q4_K: + case GGML_TYPE_Q5_K: + case GGML_TYPE_Q6_K: + case GGML_TYPE_IQ4_NL: + return 2 * generic; + default: + break; + } + } + return generic; + } + return 1; +} + +static constexpr __host__ __device__ int calc_rows_per_block(int ncols_dst, int table_id, bool small_k = false, int nwarps = 1) { + if (table_id == MMVQ_PARAMETERS_GENERIC || table_id == MMVQ_PARAMETERS_GCN || table_id == MMVQ_PARAMETERS_TURING || table_id == MMVQ_PARAMETERS_GB10) { + switch (ncols_dst) { + case 1: + return small_k ? nwarps : 1; + case 2: + case 3: + case 4: + case 5: + case 6: + case 7: + case 8: + return 2; + default: + return 1; + } + } + return 1; +} + +// FreeToken: dst_t is float, nv_bfloat16 or half +template +__launch_bounds__(calc_nwarps(type, ncols_dst, get_device_table_id(), small_k, halve_iters)*ggml_cuda_get_physical_warp_size(), 1) +static __global__ void mul_mat_vec_q( + const void * vx_ptr, const void * vy_ptr, const int32_t * ids_ptr, const ggml_cuda_mm_fusion_args_device fusion, dst_t * dst_ptr, + const uint32_t ncols_x, const uint3 nchannels_y, const uint32_t stride_row_x, const uint32_t stride_col_y, + uint32_t stride_col_dst, const uint3 channel_ratio, const uint32_t stride_channel_x, + const uint32_t stride_channel_y, const uint32_t stride_channel_dst, const uint3 sample_ratio, + const uint32_t stride_sample_x, const uint32_t stride_sample_y, const uint32_t stride_sample_dst, + const uint32_t ids_stride) { + const void * GGML_CUDA_RESTRICT vx = vx_ptr; + const void * GGML_CUDA_RESTRICT vy = vy_ptr; + const int32_t * GGML_CUDA_RESTRICT ids = ids_ptr; + dst_t * GGML_CUDA_RESTRICT dst = dst_ptr; + + constexpr int qk = ggml_cuda_type_traits::qk; + constexpr int qi = ggml_cuda_type_traits::qi; + constexpr int vdr = get_vdr_mmvq(type); + constexpr mmvq_parameter_table_id table_id = get_device_table_id(); + constexpr int nwarps = calc_nwarps(type, ncols_dst, table_id, small_k, halve_iters); + constexpr int rows_per_cuda_block = calc_rows_per_block(ncols_dst, table_id, small_k, nwarps); + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + + constexpr vec_dot_q_cuda_t vec_dot_q_cuda = get_vec_dot_q_cuda(type); + + const int tid = warp_size*threadIdx.y + threadIdx.x; + const int row0 = rows_per_cuda_block*blockIdx.x; + const int blocks_per_row_x = ncols_x / qk; + constexpr int blocks_per_iter = vdr * nwarps*warp_size / qi; + + const bool shared_expert = has_fusion && fusion.shared_up && blockIdx.y == gridDim.y - 1; + const uint32_t channel_dst = shared_expert ? 0 : blockIdx.y; + if (shared_expert) { + vx = fusion.shared_up; + if constexpr (std::is_same_v) { + dst = fusion.shared_dst; + } + stride_col_dst = fusion.shared_stride_col_dst; + } + + uint32_t channel_x; + uint32_t channel_y; + uint32_t sample_dst; + + ggml_cuda_pdl_sync(); + channel_x = shared_expert ? 0 : ncols_dst == 1 && ids ? ids[channel_dst] : fastdiv(channel_dst, channel_ratio); + channel_y = ncols_dst == 1 && ids ? fastmodulo(channel_dst, nchannels_y) : channel_dst; + sample_dst = blockIdx.z; + + const uint32_t sample_x = fastdiv(sample_dst, sample_ratio); + const uint32_t sample_y = sample_dst; + + bool use_gate = false; + bool use_bias = false; + bool use_gate_bias = false; + bool use_scale = false; + bool use_gate_scale = false; + [[maybe_unused]] const void * vgate = nullptr; + const float * x_bias = nullptr; + const float * gate_bias = nullptr; + const float * x_scale = nullptr; + const float * gate_scale = nullptr; + ggml_glu_op active_glu; + float glu_limit = 0.0f; + + if constexpr (has_fusion) { + use_gate = fusion.gate != nullptr; + use_bias = fusion.x_bias != nullptr; + use_gate_bias = fusion.gate_bias != nullptr && use_gate; + vgate = shared_expert ? fusion.shared_gate : fusion.gate; + x_bias = (const float *) fusion.x_bias; + gate_bias = (const float *) fusion.gate_bias; + active_glu = fusion.glu_op; + glu_limit = fusion.glu_limit; + if constexpr (type == GGML_TYPE_NVFP4) { + use_scale = fusion.x_scale != nullptr; + use_gate_scale = fusion.gate_scale != nullptr && use_gate; + x_scale = (const float *) fusion.x_scale; + gate_scale = (const float *) fusion.gate_scale; + } + } + + + [[maybe_unused]] float x_biases[ncols_dst] = { 0.0f }; + [[maybe_unused]] float gate_biases[ncols_dst] = { 0.0f }; + [[maybe_unused]] float x_scales = 1.0f; + [[maybe_unused]] float gate_scales = 1.0f; + if constexpr (has_fusion) { + // 1. Hide latency by prefetching bias, gates and scales here + // 2. load only on threads that won't die after partial sum calculation + const uint32_t channel_bias = ids ? channel_x : channel_dst; + if (threadIdx.x < rows_per_cuda_block && threadIdx.y == 0 && + (rows_per_cuda_block == 1 || uint32_t(row0 + threadIdx.x) < stride_col_dst)) { + if (use_bias) { + x_bias = x_bias + sample_dst * stride_sample_dst + channel_bias * stride_channel_dst + row0; +#pragma unroll + for (int j = 0; j < ncols_dst; ++j) { + x_biases[j] = x_bias[j * stride_col_dst + threadIdx.x]; + } + } + if (use_gate_bias) { + gate_bias = gate_bias + sample_dst * stride_sample_dst + channel_bias * stride_channel_dst + row0; +#pragma unroll + for (int j = 0; j < ncols_dst; ++j) { + gate_biases[j] = gate_bias[j * stride_col_dst + threadIdx.x]; + } + } + if constexpr (type == GGML_TYPE_NVFP4) { + if (use_scale) { + x_scales = x_scale[(ids ? channel_x : 0) * fusion.x_scale_stride]; + } + if (use_gate_scale) { + gate_scales = gate_scale[(ids ? channel_x : 0) * fusion.gate_scale_stride]; + } + } + } + } + + // partial sum for each thread + float tmp[ncols_dst][rows_per_cuda_block] = {{0.0f}}; + float tmp_gate[ncols_dst][rows_per_cuda_block] = {{0.0f}}; + + const block_q8_1 * y = ((const block_q8_1 *) vy) + sample_y*stride_sample_y + channel_y*stride_channel_y; + const int kbx_offset = sample_x*stride_sample_x + channel_x*stride_channel_x + row0*stride_row_x; + + for (int kbx = tid / (qi/vdr); kbx < blocks_per_row_x; kbx += blocks_per_iter) { + const int kby = kbx * (qk/QK8_1); // y block index that aligns with kbx + + // x block quant index when casting the quants to int + const int kqs = vdr * (tid % (qi/vdr)); + +#if __CUDA_ARCH__ == GGML_CUDA_CC_DGX_SPARK + // start the next iterations' weight loads early + if constexpr (mmvq_should_prefetch(type)) { + constexpr int pf_dist = 2; // loop iterations, not blocks + const int kbx_pf = kbx + pf_dist*blocks_per_iter; + if (kbx_pf < blocks_per_row_x) { +#pragma unroll + for (int i = 0; i < rows_per_cuda_block; ++i) { + const size_t off = (size_t)(kbx_offset + i*stride_row_x + kbx_pf) * ggml_cuda_type_traits::bs; + mmvq_prefetch_l2((const char *) vx + off); + if constexpr (has_fusion) { + if (use_gate) { + mmvq_prefetch_l2((const char *) vgate + off); + } + } + } + } + } +#endif + +#pragma unroll + for (int j = 0; j < ncols_dst; ++j) { +#pragma unroll + for (int i = 0; i < rows_per_cuda_block; ++i) { + tmp[j][i] += vec_dot_q_cuda( + vx, &y[j*stride_col_y + kby], kbx_offset + i*stride_row_x + kbx, kqs); + if constexpr (has_fusion) { + if (use_gate) { + tmp_gate[j][i] += vec_dot_q_cuda( + vgate, &y[j*stride_col_y + kby], kbx_offset + i*stride_row_x + kbx, kqs); + } + } + } + } + } + + __shared__ float tmp_shared[nwarps-1 > 0 ? nwarps-1 : 1][ncols_dst][rows_per_cuda_block][warp_size]; + [[maybe_unused]] __shared__ float tmp_shared_gate[(has_fusion && (nwarps-1 > 0)) ? nwarps-1 : 1][ncols_dst][rows_per_cuda_block][warp_size]; + + if (threadIdx.y > 0) { +#pragma unroll + for (int j = 0; j < ncols_dst; ++j) { +#pragma unroll + for (int i = 0; i < rows_per_cuda_block; ++i) { + tmp_shared[threadIdx.y-1][j][i][threadIdx.x] = tmp[j][i]; + if constexpr (has_fusion) { + if (use_gate) { + tmp_shared_gate[threadIdx.y-1][j][i][threadIdx.x] = tmp_gate[j][i]; + } + } + } + } + } + __syncthreads(); + if (threadIdx.y > 0) { + return; + } + + dst += sample_dst*stride_sample_dst + channel_dst*stride_channel_dst + row0; + + // sum up partial sums and write back result +#pragma unroll + for (int j = 0; j < ncols_dst; ++j) { +#pragma unroll + for (int i = 0; i < rows_per_cuda_block; ++i) { +#pragma unroll + for (int l = 0; l < nwarps-1; ++l) { + tmp[j][i] += tmp_shared[l][j][i][threadIdx.x]; + if constexpr (has_fusion) { + if (use_gate) { + tmp_gate[j][i] += tmp_shared_gate[l][j][i][threadIdx.x]; + } + } + } + tmp[j][i] = warp_reduce_sum(tmp[j][i]); + if constexpr (has_fusion) { + if (use_gate) { + tmp_gate[j][i] = warp_reduce_sum(tmp_gate[j][i]); + } + } + + if (threadIdx.x == i && (rows_per_cuda_block == 1 || uint32_t(row0 + i) < stride_col_dst)) { + float result = tmp[j][i]; + if constexpr (has_fusion) { + if constexpr (type == GGML_TYPE_NVFP4) { + result *= x_scales; + } + result += x_biases[j]; + if (use_gate) { + float gate_value = tmp_gate[j][i]; + if constexpr (type == GGML_TYPE_NVFP4) { + gate_value *= gate_scales; + } + gate_value += gate_biases[j]; + switch (active_glu) { + case GGML_GLU_OP_SWIGLU: + result *= ggml_cuda_op_silu_single(gate_value); + break; + case GGML_GLU_OP_GEGLU: + result *= ggml_cuda_op_gelu_single(gate_value); + break; + case GGML_GLU_OP_SWIGLU_OAI: + result = ggml_cuda_op_swiglu_oai_single(gate_value, result); + break; + case GGML_GLU_OP_SWIGLU_CLAMP: + result = ggml_cuda_op_swiglu_clamp_single(gate_value, result, glu_limit); + break; + default: + result = result * gate_value; + break; + } + } + } + dst[j*stride_col_dst + i] = ggml_cuda_cast(result); + } + } + } + + if constexpr (!has_fusion) { + GGML_UNUSED_VARS(use_gate, use_bias, use_gate_bias, use_scale, use_gate_scale, active_glu, glu_limit, gate_bias, x_bias, x_scale, gate_scale, tmp_gate); + } + if constexpr (type != GGML_TYPE_NVFP4) { + GGML_UNUSED_VARS(use_scale, use_gate_scale, x_scale, gate_scale, x_scales, gate_scales); + } +} + +// Dedicated MoE multi-token kernel. +// Grid: (ceil(nrows_x / c_rows_per_block), nchannels_dst) +// Block: (warp_size, ncols_dst) - each warp handles one token independently. +// No shared memory reduction needed since each warp works alone. +template +__launch_bounds__(get_mmvq_mmid_max_batch_for_device()*ggml_cuda_get_physical_warp_size(), 1) +static __global__ void mul_mat_vec_q_moe( + const void * vx_ptr, const void * vy_ptr, const int32_t * ids_ptr, const ggml_cuda_mm_fusion_args_device fusion, + dst_t * dst_ptr, + const uint32_t ncols_x, const uint3 nchannels_y, const uint32_t nrows_x, + const uint32_t stride_row_x, const uint32_t stride_col_y, uint32_t stride_col_dst, + const uint32_t stride_channel_x, const uint32_t stride_channel_y, const uint32_t stride_channel_dst, + const uint32_t ncols_dst, const uint32_t ids_stride) { + const void * GGML_CUDA_RESTRICT vx = vx_ptr; + const void * GGML_CUDA_RESTRICT vy = vy_ptr; + const int32_t * GGML_CUDA_RESTRICT ids = ids_ptr; + dst_t * GGML_CUDA_RESTRICT dst = dst_ptr; + + constexpr int qk = ggml_cuda_type_traits::qk; + constexpr int qi = ggml_cuda_type_traits::qi; + constexpr int vdr = get_vdr_mmvq(type); + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + + constexpr vec_dot_q_cuda_t vec_dot_q_cuda = get_vec_dot_q_cuda(type); + + const bool shared_expert = has_fusion && fusion.shared_up && blockIdx.y == gridDim.y - 1; + if (shared_expert) { + vx = fusion.shared_up; + if constexpr (std::is_same_v) { + dst = fusion.shared_dst; + } + stride_col_dst = fusion.shared_stride_col_dst; + } + + // fuse gate, bias, scales, and glu_op into the up projection + bool use_gate = false; + const void * vgate = nullptr; + const float * x_bias = nullptr; + const float * gate_bias = nullptr; + const float * x_scale = nullptr; + const float * gate_scale = nullptr; + ggml_glu_op active_glu = GGML_GLU_OP_SWIGLU; + float glu_limit = 0.0f; + + if constexpr (has_fusion) { + use_gate = fusion.gate != nullptr; + vgate = shared_expert ? fusion.shared_gate : fusion.gate; + x_bias = (const float *) fusion.x_bias; + gate_bias = (const float *) fusion.gate_bias; + active_glu = fusion.glu_op; + glu_limit = fusion.glu_limit; + if constexpr (type == GGML_TYPE_NVFP4) { + x_scale = (const float *) fusion.x_scale; + gate_scale = (const float *) fusion.gate_scale; + } + } + + const uint32_t token_idx = threadIdx.y; + const int row0 = c_rows_per_block*blockIdx.x; + const int blocks_per_row_x = ncols_x / qk; + constexpr int blocks_per_iter = vdr * warp_size / qi; + + const uint32_t channel_dst = shared_expert ? 0 : blockIdx.y; + + if (token_idx >= ncols_dst) { + return; + } + + ggml_cuda_pdl_sync(); + const uint32_t channel_x = shared_expert ? 0 : ids[channel_dst + token_idx * ids_stride]; + const uint32_t channel_y = fastmodulo(channel_dst, nchannels_y); + + const block_q8_1 * y = ((const block_q8_1 *) vy) + channel_y*stride_channel_y + token_idx*stride_col_y; + const int kbx_offset = channel_x*stride_channel_x + row0*stride_row_x; + + // partial sum for each thread + float tmp[c_rows_per_block] = {0.0f}; + float tmp_gate[c_rows_per_block] = {0.0f}; + + for (int kbx = threadIdx.x / (qi/vdr); kbx < blocks_per_row_x; kbx += blocks_per_iter) { + const int kby = kbx * (qk/QK8_1); + const int kqs = vdr * (threadIdx.x % (qi/vdr)); + +#pragma unroll + for (int i = 0; i < c_rows_per_block; ++i) { + tmp[i] += vec_dot_q_cuda(vx, &y[kby], kbx_offset + i*stride_row_x + kbx, kqs); + if constexpr (has_fusion) { + if (use_gate) { + tmp_gate[i] += vec_dot_q_cuda(vgate, &y[kby], kbx_offset + i*stride_row_x + kbx, kqs); + } + } + } + } + + ggml_cuda_pdl_lc(); + + // Warp-level reduction only - no shared memory needed +#pragma unroll + for (int i = 0; i < c_rows_per_block; ++i) { + tmp[i] = warp_reduce_sum(tmp[i]); + if constexpr (has_fusion) { + if (use_gate) { + tmp_gate[i] = warp_reduce_sum(tmp_gate[i]); + } + } + } + + // Write results + if (threadIdx.x < c_rows_per_block && (c_rows_per_block == 1 || uint32_t(row0 + threadIdx.x) < nrows_x)) { + float result = tmp[threadIdx.x]; + if constexpr (has_fusion) { + const uint32_t bias_idx = channel_x*stride_channel_dst + row0 + threadIdx.x; + + if constexpr (type == GGML_TYPE_NVFP4) { + if (x_scale) { + result *= x_scale[channel_x * fusion.x_scale_stride]; + } + } + if (x_bias) { + result += x_bias[bias_idx]; + } + if (use_gate) { + float gate_value = tmp_gate[threadIdx.x]; + if constexpr (type == GGML_TYPE_NVFP4) { + if (gate_scale) { + gate_value *= gate_scale[channel_x * fusion.gate_scale_stride]; + } + } + if (gate_bias) { + gate_value += gate_bias[bias_idx]; + } + switch (active_glu) { + case GGML_GLU_OP_SWIGLU: + result *= ggml_cuda_op_silu_single(gate_value); + break; + case GGML_GLU_OP_GEGLU: + result *= ggml_cuda_op_gelu_single(gate_value); + break; + case GGML_GLU_OP_SWIGLU_OAI: + result = ggml_cuda_op_swiglu_oai_single(gate_value, result); + break; + case GGML_GLU_OP_SWIGLU_CLAMP: + result = ggml_cuda_op_swiglu_clamp_single(gate_value, result, glu_limit); + break; + default: + result = result * gate_value; + break; + } + } + } + dst[channel_dst*stride_channel_dst + token_idx*stride_col_dst + row0 + threadIdx.x] = ggml_cuda_cast(result); + } + + if constexpr (!has_fusion) { + GGML_UNUSED_VARS(use_gate, tmp_gate, vgate, x_bias, gate_bias, active_glu, glu_limit, x_scale, gate_scale); + } else if constexpr (type != GGML_TYPE_NVFP4) { + GGML_UNUSED_VARS(x_scale, gate_scale); + } +} + +template +static std::pair calc_launch_params( + const int ncols_dst, const int nrows_x, const int nchannels_dst, const int nsamples_or_ntokens, + const int warp_size, const mmvq_parameter_table_id table_id, const bool small_k = false, const bool halve_iters = false) { + const int nwarps = calc_nwarps(type, ncols_dst, table_id, small_k, halve_iters); + const int rpb = calc_rows_per_block(ncols_dst, table_id, small_k, nwarps); + const int64_t nblocks = (nrows_x + rpb - 1) / rpb; + const dim3 block_nums(nblocks, nchannels_dst, nsamples_or_ntokens); + const dim3 block_dims(warp_size, nwarps, 1); + return {block_nums, block_dims}; +} + +template +static void mul_mat_vec_q_switch_fusion( + const void * vx, const void * vy, const int32_t * ids, const ggml_cuda_mm_fusion_args_device fusion, dst_t * dst, + const uint32_t ncols_x, const uint3 nchannels_y, const uint32_t stride_row_x, const uint32_t stride_col_y, + const uint32_t stride_col_dst, const uint3 channel_ratio, const uint32_t stride_channel_x, + const uint32_t stride_channel_y, const uint32_t stride_channel_dst, const uint3 sample_ratio, + const uint32_t stride_sample_x, const uint32_t stride_sample_y, const uint32_t stride_sample_dst, + const dim3 & block_nums, const dim3 & block_dims, const int nbytes_shared, + const uint32_t ids_stride, cudaStream_t stream) { + + const bool has_fusion = fusion.gate != nullptr || fusion.x_bias != nullptr || fusion.gate_bias != nullptr || + fusion.x_scale != nullptr || fusion.gate_scale != nullptr; + if constexpr (c_ncols_dst == 1) { + if (has_fusion) { + const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(block_nums, block_dims, nbytes_shared, stream); + ggml_cuda_kernel_launch(mul_mat_vec_q, launch_params, + vx, vy, ids, fusion, dst, ncols_x, nchannels_y, stride_row_x, stride_col_y, stride_col_dst, + channel_ratio, stride_channel_x, stride_channel_y, stride_channel_dst, + sample_ratio, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride); + return; + } + } + + GGML_ASSERT(!has_fusion && "fusion only supported for ncols_dst=1"); + + const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(block_nums, block_dims, nbytes_shared, stream); + ggml_cuda_kernel_launch(mul_mat_vec_q, launch_params, + vx, vy, ids, fusion, dst, ncols_x, nchannels_y, stride_row_x, stride_col_y, stride_col_dst, + channel_ratio, stride_channel_x, stride_channel_y, stride_channel_dst, + sample_ratio, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride); +} + +template +static void mul_mat_vec_q_moe_launch( + const void * vx, const void * vy, const int32_t * ids, const ggml_cuda_mm_fusion_args_device fusion, dst_t * dst, + const uint32_t ncols_x, const uint3 nchannels_y, const uint32_t nrows_x, + const uint32_t stride_row_x, const uint32_t stride_col_y, const uint32_t stride_col_dst, + const uint32_t stride_channel_x, const uint32_t stride_channel_y, const uint32_t stride_channel_dst, + const uint32_t ncols_dst, const uint32_t ids_stride, + const int warp_size, const int nchannels_dst, cudaStream_t stream) { + + constexpr int rows_per_block = 2; // 2 gives best perf based on tuning + const int64_t nblocks_rows = (nrows_x + rows_per_block - 1) / rows_per_block; + const dim3 block_nums(nblocks_rows, nchannels_dst + (fusion.shared_up != nullptr)); + const dim3 block_dims(warp_size, ncols_dst); + const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(block_nums, block_dims, 0, stream); + + const bool has_fusion = fusion.gate != nullptr || fusion.x_bias != nullptr || fusion.gate_bias != nullptr || + fusion.x_scale != nullptr || fusion.gate_scale != nullptr; + + if (has_fusion) { + ggml_cuda_kernel_launch(mul_mat_vec_q_moe, launch_params, + vx, vy, ids, fusion, dst, ncols_x, nchannels_y, nrows_x, + stride_row_x, stride_col_y, stride_col_dst, + stride_channel_x, stride_channel_y, stride_channel_dst, + ncols_dst, ids_stride); + } else { + ggml_cuda_kernel_launch(mul_mat_vec_q_moe, launch_params, + vx, vy, ids, fusion, dst, ncols_x, nchannels_y, nrows_x, + stride_row_x, stride_col_y, stride_col_dst, + stride_channel_x, stride_channel_y, stride_channel_dst, + ncols_dst, ids_stride); + } +} + +template +static void mul_mat_vec_q_switch_ncols_dst( + const void * vx, const void * vy, const int32_t * ids, const ggml_cuda_mm_fusion_args_device fusion, dst_t * dst, + const int ncols_x, const int nrows_x, const int ncols_dst, + const int stride_row_x, const int stride_col_y, const int stride_col_dst, + const int nchannels_x, const int nchannels_y, const int nchannels_dst, + const int stride_channel_x, const int stride_channel_y, const int stride_channel_dst, + const int nsamples_x, const int nsamples_dst, const int stride_sample_x, const int stride_sample_y, const int stride_sample_dst, + const int ids_stride, cudaStream_t stream) { + + GGML_ASSERT(ncols_x % ggml_blck_size(type) == 0); + GGML_ASSERT(ncols_dst <= MMVQ_MAX_BATCH_SIZE); + + const uint3 nchannels_y_fd = ids ? init_fastdiv_values(nchannels_y) : make_uint3(0, 0, 0); + const uint3 channel_ratio_fd = ids ? make_uint3(0, 0, 0) : init_fastdiv_values(nchannels_dst / nchannels_x); + const uint3 sample_ratio_fd = init_fastdiv_values(nsamples_dst / nsamples_x); + + const int device = ggml_cuda_get_device(); + const int cc = ggml_cuda_info().devices[device].cc; + const int warp_size = ggml_cuda_info().devices[device].warp_size; + const mmvq_parameter_table_id table_id = get_device_table_id(cc); + + const bool has_ids = ids != nullptr; + + // How the K loop divides up at the baseline block width, both decisions below use these. + constexpr int qk = ggml_cuda_type_traits::qk; + constexpr int qi = ggml_cuda_type_traits::qi; + constexpr int vdr = get_vdr_mmvq(type); + const int blocks_per_row_x = ncols_x / qk; + const int blocks_per_iter_1warp = vdr * warp_size / qi; + + const auto should_use_small_k = [&](int c_ncols_dst) { + // When K is small, increase rows_per_block to match nwarps so each warp has more work to do + // Trigger when the full thread block covers all K blocks in a single loop iteration and few threads remain idle. + const int nwarps = calc_nwarps(type, c_ncols_dst, table_id); + bool use = nwarps > 1 && blocks_per_row_x < nwarps * blocks_per_iter_1warp; + + constexpr std::array iq_slow_turing = { + GGML_TYPE_IQ3_XXS, + GGML_TYPE_IQ3_S, + }; + constexpr std::array iq_slow_other = { + GGML_TYPE_IQ1_S, GGML_TYPE_IQ1_M, GGML_TYPE_IQ2_XXS, GGML_TYPE_IQ2_XS, + GGML_TYPE_IQ2_S, GGML_TYPE_IQ3_XXS, GGML_TYPE_IQ3_S, GGML_TYPE_IQ4_XS, + }; + constexpr std::array slow_pascal = { + GGML_TYPE_IQ3_S, + GGML_TYPE_Q2_K, + GGML_TYPE_Q3_K, + }; + + const bool is_nvidia_turing_plus = GGML_CUDA_CC_IS_NVIDIA(cc) && cc >= GGML_CUDA_CC_TURING; + const bool is_nvidia_pascal_older = GGML_CUDA_CC_IS_NVIDIA(cc) && cc < GGML_CUDA_CC_VOLTA; + + if (is_nvidia_turing_plus) { + if (ncols_dst == 1 && + std::find(iq_slow_turing.begin(), iq_slow_turing.end(), type) != iq_slow_turing.end()) { + use = false; + } + } else if ((ncols_dst == 1 && std::find(iq_slow_other.begin(), iq_slow_other.end(), type) != iq_slow_other.end()) || + (is_nvidia_pascal_older && std::find(slow_pascal.begin(), slow_pascal.end(), type) != slow_pascal.end()) || + GGML_CUDA_CC_IS_RDNA(cc)) { + use = false; + } + + return use; + }; + + // Whether doubling nwarps pays off on the ncols_dst == 1 path, where K sets the K loop trip count. + const auto should_halve_iters = [&] { + if (table_id != MMVQ_PARAMETERS_GB10) { + return false; + } + + // Expert rows are gathered per token, so a wider block adds reduction work without reuse. + if (has_ids) { + return false; + } + + const int blocks_per_iter = calc_nwarps(type, 1, table_id) * blocks_per_iter_1warp; + const int iters = (blocks_per_row_x + blocks_per_iter - 1) / blocks_per_iter; + const int iters_wide = (blocks_per_row_x + blocks_per_iter * 2 - 1) / (blocks_per_iter * 2); + + // An odd trip count leaves half the wider block idle for its last iteration, that tail is + // only affordable once the loop is long enough to dilute it to an eighth of the work (observation). + const int idle = iters_wide * 2 - iters; + + return idle * 8 <= iters_wide * 2; + }; + + if (has_ids && ncols_dst > 1) { + // Multi-token MUL_MAT_ID path - dedicated MoE kernel + mul_mat_vec_q_moe_launch( + vx, vy, ids, fusion, dst, ncols_x, nchannels_y_fd, nrows_x, + stride_row_x, stride_col_y, stride_col_dst, + stride_channel_x, stride_channel_y, stride_channel_dst, + ncols_dst, ids_stride, warp_size, nchannels_dst, stream); + return; + } + + switch (ncols_dst) { + case 1: { + // static, else MSVC lambda capture breaks the constexpr uses below + static constexpr int c_ncols_dst = 1; + + // Tag types keep the flags compile-time, so __launch_bounds__ matches what is launched. + const auto launch = [&](auto small_k_tag, auto halve_iters_tag) { + constexpr bool c_small_k = decltype(small_k_tag)::value; + // Types the table does not promote would compile a second, identical kernel. + constexpr bool c_promoted = + calc_nwarps(type, c_ncols_dst, MMVQ_PARAMETERS_GB10, false, true) != + calc_nwarps(type, c_ncols_dst, MMVQ_PARAMETERS_GB10, false, false); + + constexpr bool c_halve_iters = decltype(halve_iters_tag)::value && c_promoted; + + const std::pair dims = calc_launch_params(c_ncols_dst, nrows_x, nchannels_dst + (fusion.shared_up != nullptr), + nsamples_dst, warp_size, table_id, c_small_k, c_halve_iters); + mul_mat_vec_q_switch_fusion( + vx, vy, ids, fusion, dst, ncols_x, nchannels_y_fd, stride_row_x, stride_col_y, stride_col_dst, + channel_ratio_fd, stride_channel_x, stride_channel_y, stride_channel_dst, sample_ratio_fd, + stride_sample_x, stride_sample_y, stride_sample_dst, dims.first, dims.second, 0, ids_stride, + stream); + }; + + if (should_use_small_k(c_ncols_dst)) { + launch(std::true_type{}, std::false_type{}); + } else if (should_halve_iters()) { + launch(std::false_type{}, std::true_type{}); + } else { + launch(std::false_type{}, std::false_type{}); + } + } break; + case 2: { + constexpr int c_ncols_dst = 2; + std::pair dims = calc_launch_params(c_ncols_dst, nrows_x, nchannels_dst, nsamples_dst, warp_size, table_id); + mul_mat_vec_q_switch_fusion(vx, vy, ids, fusion, dst, ncols_x, nchannels_y_fd, stride_row_x, stride_col_y, stride_col_dst, + channel_ratio_fd, stride_channel_x, stride_channel_y, stride_channel_dst, + sample_ratio_fd, stride_sample_x, stride_sample_y, stride_sample_dst, + dims.first, dims.second, 0, ids_stride, stream); + } break; + case 3: { + constexpr int c_ncols_dst = 3; + std::pair dims = calc_launch_params(c_ncols_dst, nrows_x, nchannels_dst, nsamples_dst, warp_size, table_id); + mul_mat_vec_q_switch_fusion(vx, vy, ids, fusion, dst, ncols_x, nchannels_y_fd, stride_row_x, stride_col_y, stride_col_dst, + channel_ratio_fd, stride_channel_x, stride_channel_y, stride_channel_dst, + sample_ratio_fd, stride_sample_x, stride_sample_y, stride_sample_dst, + dims.first, dims.second, 0, ids_stride, stream); + } break; + case 4: { + constexpr int c_ncols_dst = 4; + std::pair dims = calc_launch_params(c_ncols_dst, nrows_x, nchannels_dst, nsamples_dst, warp_size, table_id); + mul_mat_vec_q_switch_fusion(vx, vy, ids, fusion, dst, ncols_x, nchannels_y_fd, stride_row_x, stride_col_y, stride_col_dst, + channel_ratio_fd, stride_channel_x, stride_channel_y, stride_channel_dst, + sample_ratio_fd, stride_sample_x, stride_sample_y, stride_sample_dst, + dims.first, dims.second, 0, ids_stride, stream); + } break; + case 5: { + constexpr int c_ncols_dst = 5; + std::pair dims = calc_launch_params(c_ncols_dst, nrows_x, nchannels_dst, nsamples_dst, warp_size, table_id); + mul_mat_vec_q_switch_fusion(vx, vy, ids, fusion, dst, ncols_x, nchannels_y_fd, stride_row_x, stride_col_y, stride_col_dst, + channel_ratio_fd, stride_channel_x, stride_channel_y, stride_channel_dst, + sample_ratio_fd, stride_sample_x, stride_sample_y, stride_sample_dst, + dims.first, dims.second, 0, ids_stride, stream); + } break; + case 6: { + constexpr int c_ncols_dst = 6; + std::pair dims = calc_launch_params(c_ncols_dst, nrows_x, nchannels_dst, nsamples_dst, warp_size, table_id); + mul_mat_vec_q_switch_fusion(vx, vy, ids, fusion, dst, ncols_x, nchannels_y_fd, stride_row_x, stride_col_y, stride_col_dst, + channel_ratio_fd, stride_channel_x, stride_channel_y, stride_channel_dst, + sample_ratio_fd, stride_sample_x, stride_sample_y, stride_sample_dst, + dims.first, dims.second, 0, ids_stride, stream); + } break; + case 7: { + constexpr int c_ncols_dst = 7; + std::pair dims = calc_launch_params(c_ncols_dst, nrows_x, nchannels_dst, nsamples_dst, warp_size, table_id); + mul_mat_vec_q_switch_fusion(vx, vy, ids, fusion, dst, ncols_x, nchannels_y_fd, stride_row_x, stride_col_y, stride_col_dst, + channel_ratio_fd, stride_channel_x, stride_channel_y, stride_channel_dst, + sample_ratio_fd, stride_sample_x, stride_sample_y, stride_sample_dst, + dims.first, dims.second, 0, ids_stride, stream); + } break; + case 8: { + constexpr int c_ncols_dst = 8; + std::pair dims = calc_launch_params(c_ncols_dst, nrows_x, nchannels_dst, nsamples_dst, warp_size, table_id); + mul_mat_vec_q_switch_fusion(vx, vy, ids, fusion, dst, ncols_x, nchannels_y_fd, stride_row_x, stride_col_y, stride_col_dst, + channel_ratio_fd, stride_channel_x, stride_channel_y, stride_channel_dst, + sample_ratio_fd, stride_sample_x, stride_sample_y, stride_sample_dst, + dims.first, dims.second, 0, ids_stride, stream); + } break; + default: + GGML_ABORT("fatal error"); + } +} + +template +static void mul_mat_vec_q_switch_type( + const void * vx, const ggml_type type_x, const void * vy, const int32_t * ids, const ggml_cuda_mm_fusion_args_device fusion, dst_t * dst, + const int ncols_x, const int nrows_x, const int ncols_dst, + const int stride_row_x, const int stride_col_y, const int stride_col_dst, + const int nchannels_x, const int nchannels_y, const int nchannels_dst, + const int stride_channel_x, const int stride_channel_y, const int stride_channel_dst, + const int nsamples_x, const int nsamples_dst, const int stride_sample_x, const int stride_sample_y, const int stride_sample_dst, + const int ids_stride, cudaStream_t stream) { + switch (type_x) { + case GGML_TYPE_Q1_0: + mul_mat_vec_q_switch_ncols_dst + (vx, vy, ids, fusion, dst, ncols_x, nrows_x, ncols_dst, stride_row_x, stride_col_y, stride_col_dst, + nchannels_x, nchannels_y, nchannels_dst, stride_channel_x, stride_channel_y, stride_channel_dst, + nsamples_x, nsamples_dst, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride, stream); + break; + case GGML_TYPE_Q2_0: + mul_mat_vec_q_switch_ncols_dst + (vx, vy, ids, fusion, dst, ncols_x, nrows_x, ncols_dst, stride_row_x, stride_col_y, stride_col_dst, + nchannels_x, nchannels_y, nchannels_dst, stride_channel_x, stride_channel_y, stride_channel_dst, + nsamples_x, nsamples_dst, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride, stream); + break; + case GGML_TYPE_Q4_0: + mul_mat_vec_q_switch_ncols_dst + (vx, vy, ids, fusion, dst, ncols_x, nrows_x, ncols_dst, stride_row_x, stride_col_y, stride_col_dst, + nchannels_x, nchannels_y, nchannels_dst, stride_channel_x, stride_channel_y, stride_channel_dst, + nsamples_x, nsamples_dst, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride, stream); + break; + case GGML_TYPE_Q4_1: + mul_mat_vec_q_switch_ncols_dst + (vx, vy, ids, fusion, dst, ncols_x, nrows_x, ncols_dst, stride_row_x, stride_col_y, stride_col_dst, + nchannels_x, nchannels_y, nchannels_dst, stride_channel_x, stride_channel_y, stride_channel_dst, + nsamples_x, nsamples_dst, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride, stream); + break; + case GGML_TYPE_Q5_0: + mul_mat_vec_q_switch_ncols_dst + (vx, vy, ids, fusion, dst, ncols_x, nrows_x, ncols_dst, stride_row_x, stride_col_y, stride_col_dst, + nchannels_x, nchannels_y, nchannels_dst, stride_channel_x, stride_channel_y, stride_channel_dst, + nsamples_x, nsamples_dst, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride, stream); + break; + case GGML_TYPE_Q5_1: + mul_mat_vec_q_switch_ncols_dst + (vx, vy, ids, fusion, dst, ncols_x, nrows_x, ncols_dst, stride_row_x, stride_col_y, stride_col_dst, + nchannels_x, nchannels_y, nchannels_dst, stride_channel_x, stride_channel_y, stride_channel_dst, + nsamples_x, nsamples_dst, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride, stream); + break; + case GGML_TYPE_Q8_0: + mul_mat_vec_q_switch_ncols_dst + (vx, vy, ids, fusion, dst, ncols_x, nrows_x, ncols_dst, stride_row_x, stride_col_y, stride_col_dst, + nchannels_x, nchannels_y, nchannels_dst, stride_channel_x, stride_channel_y, stride_channel_dst, + nsamples_x, nsamples_dst, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride, stream); + break; + case GGML_TYPE_MXFP4: + mul_mat_vec_q_switch_ncols_dst + (vx, vy, ids, fusion, dst, ncols_x, nrows_x, ncols_dst, stride_row_x, stride_col_y, stride_col_dst, + nchannels_x, nchannels_y, nchannels_dst, stride_channel_x, stride_channel_y, stride_channel_dst, + nsamples_x, nsamples_dst, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride, stream); + break; + case GGML_TYPE_NVFP4: + mul_mat_vec_q_switch_ncols_dst + (vx, vy, ids, fusion, dst, ncols_x, nrows_x, ncols_dst, stride_row_x, stride_col_y, stride_col_dst, + nchannels_x, nchannels_y, nchannels_dst, stride_channel_x, stride_channel_y, stride_channel_dst, + nsamples_x, nsamples_dst, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride, stream); + break; + case GGML_TYPE_Q2_K: + mul_mat_vec_q_switch_ncols_dst + (vx, vy, ids, fusion, dst, ncols_x, nrows_x, ncols_dst, stride_row_x, stride_col_y, stride_col_dst, + nchannels_x, nchannels_y, nchannels_dst, stride_channel_x, stride_channel_y, stride_channel_dst, + nsamples_x, nsamples_dst, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride, stream); + break; + case GGML_TYPE_Q3_K: + mul_mat_vec_q_switch_ncols_dst + (vx, vy, ids, fusion, dst, ncols_x, nrows_x, ncols_dst, stride_row_x, stride_col_y, stride_col_dst, + nchannels_x, nchannels_y, nchannels_dst, stride_channel_x, stride_channel_y, stride_channel_dst, + nsamples_x, nsamples_dst, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride, stream); + break; + case GGML_TYPE_Q4_K: + mul_mat_vec_q_switch_ncols_dst + (vx, vy, ids, fusion, dst, ncols_x, nrows_x, ncols_dst, stride_row_x, stride_col_y, stride_col_dst, + nchannels_x, nchannels_y, nchannels_dst, stride_channel_x, stride_channel_y, stride_channel_dst, + nsamples_x, nsamples_dst, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride, stream); + break; + case GGML_TYPE_Q5_K: + mul_mat_vec_q_switch_ncols_dst + (vx, vy, ids, fusion, dst, ncols_x, nrows_x, ncols_dst, stride_row_x, stride_col_y, stride_col_dst, + nchannels_x, nchannels_y, nchannels_dst, stride_channel_x, stride_channel_y, stride_channel_dst, + nsamples_x, nsamples_dst, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride, stream); + break; + case GGML_TYPE_Q6_K: + mul_mat_vec_q_switch_ncols_dst + (vx, vy, ids, fusion, dst, ncols_x, nrows_x, ncols_dst, stride_row_x, stride_col_y, stride_col_dst, + nchannels_x, nchannels_y, nchannels_dst, stride_channel_x, stride_channel_y, stride_channel_dst, + nsamples_x, nsamples_dst, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride, stream); + break; + case GGML_TYPE_IQ2_XXS: + mul_mat_vec_q_switch_ncols_dst + (vx, vy, ids, fusion, dst, ncols_x, nrows_x, ncols_dst, stride_row_x, stride_col_y, stride_col_dst, + nchannels_x, nchannels_y, nchannels_dst, stride_channel_x, stride_channel_y, stride_channel_dst, + nsamples_x, nsamples_dst, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride, stream); + break; + case GGML_TYPE_IQ2_XS: + mul_mat_vec_q_switch_ncols_dst + (vx, vy, ids, fusion, dst, ncols_x, nrows_x, ncols_dst, stride_row_x, stride_col_y, stride_col_dst, + nchannels_x, nchannels_y, nchannels_dst, stride_channel_x, stride_channel_y, stride_channel_dst, + nsamples_x, nsamples_dst, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride, stream); + break; + case GGML_TYPE_IQ2_S: + mul_mat_vec_q_switch_ncols_dst + (vx, vy, ids, fusion, dst, ncols_x, nrows_x, ncols_dst, stride_row_x, stride_col_y, stride_col_dst, + nchannels_x, nchannels_y, nchannels_dst, stride_channel_x, stride_channel_y, stride_channel_dst, + nsamples_x, nsamples_dst, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride, stream); + break; + case GGML_TYPE_IQ3_XXS: + mul_mat_vec_q_switch_ncols_dst + (vx, vy, ids, fusion, dst, ncols_x, nrows_x, ncols_dst, stride_row_x, stride_col_y, stride_col_dst, + nchannels_x, nchannels_y, nchannels_dst, stride_channel_x, stride_channel_y, stride_channel_dst, + nsamples_x, nsamples_dst, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride, stream); + break; + case GGML_TYPE_IQ1_S: + mul_mat_vec_q_switch_ncols_dst + (vx, vy, ids, fusion, dst, ncols_x, nrows_x, ncols_dst, stride_row_x, stride_col_y, stride_col_dst, + nchannels_x, nchannels_y, nchannels_dst, stride_channel_x, stride_channel_y, stride_channel_dst, + nsamples_x, nsamples_dst, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride, stream); + break; + case GGML_TYPE_IQ1_M: + mul_mat_vec_q_switch_ncols_dst + (vx, vy, ids, fusion, dst, ncols_x, nrows_x, ncols_dst, stride_row_x, stride_col_y, stride_col_dst, + nchannels_x, nchannels_y, nchannels_dst, stride_channel_x, stride_channel_y, stride_channel_dst, + nsamples_x, nsamples_dst, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride, stream); + break; + case GGML_TYPE_IQ4_NL: + mul_mat_vec_q_switch_ncols_dst + (vx, vy, ids, fusion, dst, ncols_x, nrows_x, ncols_dst, stride_row_x, stride_col_y, stride_col_dst, + nchannels_x, nchannels_y, nchannels_dst, stride_channel_x, stride_channel_y, stride_channel_dst, + nsamples_x, nsamples_dst, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride, stream); + break; + case GGML_TYPE_IQ4_XS: + mul_mat_vec_q_switch_ncols_dst + (vx, vy, ids, fusion, dst, ncols_x, nrows_x, ncols_dst, stride_row_x, stride_col_y, stride_col_dst, + nchannels_x, nchannels_y, nchannels_dst, stride_channel_x, stride_channel_y, stride_channel_dst, + nsamples_x, nsamples_dst, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride, stream); + break; + case GGML_TYPE_IQ3_S: + mul_mat_vec_q_switch_ncols_dst + (vx, vy, ids, fusion, dst, ncols_x, nrows_x, ncols_dst, stride_row_x, stride_col_y, stride_col_dst, + nchannels_x, nchannels_y, nchannels_dst, stride_channel_x, stride_channel_y, stride_channel_dst, + nsamples_x, nsamples_dst, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride, stream); + break; + default: + GGML_ABORT("fatal error"); + } +} + +void ggml_cuda_mul_mat_vec_q( + ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, ggml_tensor * dst, + const ggml_cuda_mm_fusion_args_host * fusion, const void * src1_q8_1_pre) { + // FreeToken: activations and results may also be bf16 or fp16 + GGML_ASSERT( src1->type == GGML_TYPE_F32 || src1->type == GGML_TYPE_BF16 || src1->type == GGML_TYPE_F16); + GGML_ASSERT( dst->type == GGML_TYPE_F32 || dst->type == GGML_TYPE_BF16 || dst->type == GGML_TYPE_F16); + GGML_ASSERT(!ids || ids->type == GGML_TYPE_I32); // Optional, used for batched GGML_MUL_MAT_ID. + + GGML_TENSOR_BINARY_OP_LOCALS; + + cudaStream_t stream = ctx.stream(); + + const size_t ts_src0 = ggml_type_size(src0->type); + const size_t ts_src1 = ggml_type_size(src1->type); + const size_t ts_dst = ggml_type_size(dst->type); + + GGML_ASSERT( nb00 == ts_src0); + GGML_ASSERT( nb10 == ts_src1); + GGML_ASSERT( nb0 == ts_dst); + GGML_ASSERT(!ids || ids->nb[0] == ggml_type_size(ids->type)); + + GGML_ASSERT(!ids || ne12 <= MMVQ_MAX_BATCH_SIZE); + + const float * src1_d = (const float *) src1->data; + const int32_t * ids_d = ids ? (const int32_t *) ids->data : nullptr; + float * dst_d = (float *) dst->data; + + ggml_cuda_mm_fusion_args_device fusion_local{}; + + if (fusion) { + const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; + GGML_ASSERT( !ids || dst->ne[2] <= get_mmvq_mmid_max_batch(src0->type, cc)); + GGML_ASSERT( ids || dst->ne[1] == 1); + // Scale fusion is only allowed for NVFP4 currently as the cost of checking this at run-time in the prologue is + // non-negligible for some models such as gpt-oss-20b + GGML_ASSERT((fusion->x_scale == nullptr && fusion->gate_scale == nullptr) || src0->type == GGML_TYPE_NVFP4); + + if (fusion->shared_up) { + GGML_ASSERT(ids && fusion->gate && fusion->shared_gate && fusion->shared_dst); + GGML_ASSERT(!fusion->x_bias && !fusion->gate_bias && !fusion->x_scale && !fusion->gate_scale); + GGML_ASSERT(ne11 == 1 && ne03 == 1 && ne13 == 1); + GGML_ASSERT(fusion->shared_up->type == src0->type && fusion->shared_gate->type == src0->type); + GGML_ASSERT(ggml_are_same_shape(fusion->shared_up, fusion->shared_gate)); + GGML_ASSERT(ggml_is_contiguous(fusion->shared_up) && ggml_is_contiguous(fusion->shared_gate)); + GGML_ASSERT(fusion->shared_up->ne[0] == ne00 && fusion->shared_up->ne[1] == ne01); + GGML_ASSERT(fusion->shared_up->nb[1] == nb01 && ggml_is_matrix(fusion->shared_up)); + GGML_ASSERT(fusion->shared_dst->type == GGML_TYPE_F32 && ggml_is_contiguous(fusion->shared_dst)); + GGML_ASSERT(fusion->shared_dst->ne[0] == ne0 && fusion->shared_dst->ne[1] == ne2); + fusion_local.shared_up = fusion->shared_up->data; + fusion_local.shared_gate = fusion->shared_gate->data; + fusion_local.shared_dst = (float *) fusion->shared_dst->data; + fusion_local.shared_stride_col_dst = fusion->shared_dst->nb[1] / ts_dst; + } + + if (fusion->x_bias) { + GGML_ASSERT(fusion->x_bias->type == GGML_TYPE_F32); + GGML_ASSERT(fusion->x_bias->ne[0] == dst->ne[0]); + GGML_ASSERT(!ids || fusion->x_bias->ne[1] == src0->ne[2]); + fusion_local.x_bias = fusion->x_bias->data; + } + if (fusion->gate) { + GGML_ASSERT(fusion->gate->type == src0->type && ggml_are_same_stride(fusion->gate, src0)); + fusion_local.gate = fusion->gate->data; + } + if (fusion->gate_bias) { + GGML_ASSERT(fusion->gate_bias->type == GGML_TYPE_F32); + GGML_ASSERT(fusion->gate_bias->ne[0] == dst->ne[0]); + GGML_ASSERT(!ids || fusion->gate_bias->ne[1] == src0->ne[2]); + fusion_local.gate_bias = fusion->gate_bias->data; + } + // FreeToken: a scale may be a strided vector, read with its element stride + if (fusion->x_scale) { + GGML_ASSERT(fusion->x_scale->type == GGML_TYPE_F32); + GGML_ASSERT(ggml_nelements(fusion->x_scale) == fusion->x_scale->ne[0] && fusion->x_scale->nb[0] % sizeof(float) == 0); + GGML_ASSERT(ggml_nelements(fusion->x_scale) == (ids ? src0->ne[2] : 1)); + fusion_local.x_scale = fusion->x_scale->data; + fusion_local.x_scale_stride = fusion->x_scale->nb[0] / sizeof(float); + } + if (fusion->gate_scale) { + GGML_ASSERT(fusion->gate_scale->type == GGML_TYPE_F32); + GGML_ASSERT(ggml_nelements(fusion->gate_scale) == fusion->gate_scale->ne[0] && fusion->gate_scale->nb[0] % sizeof(float) == 0); + GGML_ASSERT(ggml_nelements(fusion->gate_scale) == (ids ? src0->ne[2] : 1)); + fusion_local.gate_scale = fusion->gate_scale->data; + fusion_local.gate_scale_stride = fusion->gate_scale->nb[0] / sizeof(float); + } + fusion_local.glu_op = fusion->glu_op; + fusion_local.glu_limit = fusion->glu_limit; + } + + // If src0 is a temporary compute buffer, clear any potential padding. + if (ggml_backend_buffer_get_usage(src0->buffer) == GGML_BACKEND_BUFFER_USAGE_COMPUTE) { + const size_t size_data = ggml_nbytes(src0); + const size_t size_alloc = ggml_backend_buffer_get_alloc_size(src0->buffer, src0); + if (size_alloc > size_data) { + GGML_ASSERT(ggml_is_contiguously_allocated(src0)); + GGML_ASSERT(!src0->view_src); + CUDA_CHECK(cudaMemsetAsync((char *) src0->data + size_data, 0, size_alloc - size_data, stream)); + } + } + + const int64_t ne10_padded = GGML_PAD(ne10, MATRIX_ROW_PADDING); + // FreeToken: src1_q8_1_pre is src1 already quantized by quantize_row_q8_1_cuda with this padding and layout + ggml_cuda_pool_alloc src1_q8_1(ctx.pool()); + const void * src1_q8 = src1_q8_1_pre; + if (src1_q8 == nullptr) { + src1_q8_1.alloc(ne13*ne12 * ne11*ne10_padded * sizeof(block_q8_1)/QK8_1); + src1_q8 = src1_q8_1.get(); + const int64_t s11 = src1->nb[1] / ts_src1; + const int64_t s12 = src1->nb[2] / ts_src1; + const int64_t s13 = src1->nb[3] / ts_src1; + if (src1->type == GGML_TYPE_BF16) { + quantize_row_q8_1_cuda((const nv_bfloat16 *) src1->data, nullptr, src1_q8_1.get(), src0->type, ne10, s11, s12, s13, ne10_padded, ne11, ne12, ne13, stream); + } else if (src1->type == GGML_TYPE_F16) { + quantize_row_q8_1_cuda((const half *) src1->data, nullptr, src1_q8_1.get(), src0->type, ne10, s11, s12, s13, ne10_padded, ne11, ne12, ne13, stream); + } else { + quantize_row_q8_1_cuda(src1_d, nullptr, src1_q8_1.get(), src0->type, ne10, s11, s12, s13, ne10_padded, ne11, ne12, ne13, stream); + } + } + + const int64_t s01 = src0->nb[1] / ts_src0; + const int64_t s11 = ne10_padded / QK8_1; + const int64_t s1 = dst->nb[1] / ts_dst; + const int64_t s02 = src0->nb[2] / ts_src0; + const int64_t s2 = dst->nb[2] / ts_dst; + const int64_t s03 = src0->nb[3] / ts_src0; + const int64_t s3 = dst->nb[3] / ts_dst; + + const int64_t s12 = ne11*s11; + const int64_t s13 = ne12*s12; + + // For MUL_MAT_ID the memory layout is different than for MUL_MAT: + const int64_t ncols_dst = ids ? ne2 : ne1; + const int64_t nchannels_y = ids ? ne11 : ne12; + const int64_t nchannels_dst = ids ? ne1 : ne2; + const int64_t stride_col_dst = ids ? s2 : s1; + const int64_t stride_col_y = ids ? s12 : s11; + const int64_t stride_channel_dst = ids ? s1 : s2; + const int64_t stride_channel_y = ids ? s11 : s12; + + const int64_t ids_stride = ids ? ids->nb[1] / ggml_type_size(ids->type) : 0; + + const auto launch = [&](auto * dst_ptr) { + mul_mat_vec_q_switch_type( + src0->data, src0->type, src1_q8, ids_d, fusion_local, dst_ptr, ne00, + ne01, ncols_dst, s01, stride_col_y, stride_col_dst, + ne02, nchannels_y, nchannels_dst, s02, stride_channel_y, stride_channel_dst, + ne03, ne3, s03, s13, s3, ids_stride, stream); + }; + if (dst->type == GGML_TYPE_BF16) { + // FreeToken: the fused shared expert writes an fp32 shared_dst + GGML_ASSERT(!fusion || !fusion->shared_up); + launch((nv_bfloat16 *) dst->data); + } else if (dst->type == GGML_TYPE_F16) { + GGML_ASSERT(!fusion || !fusion->shared_up); + launch((half *) dst->data); + } else { + launch(dst_d); + } +} + +void ggml_cuda_op_mul_mat_vec_q( + ggml_backend_cuda_context & ctx, + const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, const char * src0_dd_i, const float * src1_ddf_i, + const char * src1_ddq_i, float * dst_dd_i, const int64_t row_low, const int64_t row_high, const int64_t src1_ncols, + const int64_t src1_padded_row_size, cudaStream_t stream) { + + const int64_t ne00 = src0->ne[0]; + const int64_t row_diff = row_high - row_low; + + const int64_t ne10 = src1->ne[0]; + GGML_ASSERT(ne10 % QK8_1 == 0); + + const int64_t ne0 = dst->ne[0]; + + int id = ggml_cuda_get_device(); + + // the main device has a larger memory buffer to hold the results from all GPUs + // nrows_dst == nrows of the matrix that the kernel writes into + const int64_t nrows_dst = id == ctx.device ? ne0 : row_diff; + + const int stride_row_x = ne00 / ggml_blck_size(src0->type); + const int stride_col_y = src1_padded_row_size / QK8_1; + + ggml_cuda_mm_fusion_args_device fusion_local{}; + mul_mat_vec_q_switch_type( + src0_dd_i, src0->type, src1_ddq_i, nullptr, fusion_local, dst_dd_i, ne00, row_diff, src1_ncols, stride_row_x, stride_col_y, nrows_dst, + 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, stream); + + GGML_UNUSED_VARS(src1, dst, src1_ddf_i, src1_ncols, src1_padded_row_size); +} diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmvq.cuh b/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmvq.cuh new file mode 100644 index 0000000000..6973780c73 --- /dev/null +++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmvq.cuh @@ -0,0 +1,20 @@ +#include "common.cuh" + +#define MMVQ_MAX_BATCH_SIZE 8 // Max. batch size for which to use MMVQ kernels. + +bool ggml_cuda_should_use_mmvq(enum ggml_type type, int cc, int64_t ne11); + +// Returns the maximum batch size for which MMVQ should be used for MUL_MAT_ID, +// based on the quantization type and GPU architecture (compute capability). +int get_mmvq_mmid_max_batch(ggml_type type, int cc); + +// FreeToken: src1_q8_1_pre skips the internal quantization of src1 +void ggml_cuda_mul_mat_vec_q(ggml_backend_cuda_context & ctx, + const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, ggml_tensor * dst, const ggml_cuda_mm_fusion_args_host * fusion = nullptr, + const void * src1_q8_1_pre = nullptr); + +void ggml_cuda_op_mul_mat_vec_q( + ggml_backend_cuda_context & ctx, + const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, const char * src0_dd_i, const float * src1_ddf_i, + const char * src1_ddq_i, float * dst_dd_i, const int64_t row_low, const int64_t row_high, const int64_t src1_ncols, + const int64_t src1_padded_row_size, cudaStream_t stream); diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/quantize.cu b/python/freetoken/kernel/csrc/ggml/ggml-cuda/quantize.cu new file mode 100644 index 0000000000..46c9f861a1 --- /dev/null +++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/quantize.cu @@ -0,0 +1,729 @@ +#include "quantize.cuh" +#include "convert.cuh" +#include + +#if defined(BLACKWELL_MMA_AVAILABLE) +// this maps to 256-bit loads in PTX on supported devices, +// and otherwise falls back to 2 128-bit loads +struct __builtin_align__(32) float8 { + float x; float y; float z; float w; + float p; float q; float r; float s; +}; + +#if CUDART_VERSION >= 12080 +static __device__ __forceinline__ float nvfp4_native_scale_error( + const float vals[QK_NVFP4_SUB], const float inv_col_scale, const float inv_scale, const float scale) { + const float scale_dequant = 2.0f * scale; + float err = 0.0f; + +#pragma unroll + for (int k = 0; k < QK_NVFP4_SUB; k += 4) { + const float v0 = vals[k + 0] * inv_col_scale; + const float v1 = vals[k + 1] * inv_col_scale; + const float v2 = vals[k + 2] * inv_col_scale; + const float v3 = vals[k + 3] * inv_col_scale; + + const __nv_fp4x4_e2m1 q(make_float4(v0 * inv_scale, v1 * inv_scale, v2 * inv_scale, v3 * inv_scale)); + const __nv_fp4x4_storage_t q_storage = q.__x; + const __nv_fp4x2_storage_t q_lo = static_cast<__nv_fp4x2_storage_t>(q_storage); + const __nv_fp4x2_storage_t q_hi = static_cast<__nv_fp4x2_storage_t>(q_storage >> 8U); + + const __half2_raw hraw2_lo = __nv_cvt_fp4x2_to_halfraw2(q_lo, __NV_E2M1); + const __half2_raw hraw2_hi = __nv_cvt_fp4x2_to_halfraw2(q_hi, __NV_E2M1); + const __half2 h2_lo = static_cast<__half2>(hraw2_lo); + const __half2 h2_hi = static_cast<__half2>(hraw2_hi); + const float2 dq_lo = __half22float2(h2_lo); + const float2 dq_hi = __half22float2(h2_hi); + + const float err0 = fabsf(v0) - fabsf(dq_lo.x) * scale_dequant; + const float err1 = fabsf(v1) - fabsf(dq_lo.y) * scale_dequant; + const float err2 = fabsf(v2) - fabsf(dq_hi.x) * scale_dequant; + const float err3 = fabsf(v3) - fabsf(dq_hi.y) * scale_dequant; + + err = fmaf(err0, err0, err); + err = fmaf(err1, err1, err); + err = fmaf(err2, err2, err); + err = fmaf(err3, err3, err); + } + + return err; +} +#endif // CUDART_VERSION >= 12080 +#endif // defined(BLACKWELL_MMA_AVAILABLE) + +template +__launch_bounds__(CUDA_QUANTIZE_BLOCK_SIZE, 1) +static __global__ void quantize_q8_1( + const src_t * x_ptr, void * vy_ptr, + const int64_t ne00, const int64_t s01, const int64_t s02, const int64_t s03, + const int64_t ne0, const uint32_t ne1, const uint3 ne2) { + ggml_cuda_pdl_lc(); + const src_t * GGML_CUDA_RESTRICT x = x_ptr; + void * GGML_CUDA_RESTRICT vy = vy_ptr; + const int64_t i0 = (int64_t)blockDim.x*blockIdx.x + threadIdx.x; + + if (i0 >= ne0) { + return; + } + + const int64_t i3 = fastdiv(blockIdx.z, ne2); + const int64_t i2 = blockIdx.z - i3*ne2.z; + const int64_t i1 = blockIdx.y; + + const int64_t & i00 = i0; + const int64_t & i01 = i1; + const int64_t & i02 = i2; + const int64_t & i03 = i3; + + const int64_t i_cont = ((i3*ne2.z + i2) * ne1 + i1) * ne0 + i0; + + block_q8_1 * y = (block_q8_1 *) vy; + + const int64_t ib = i_cont / QK8_1; // block index + const int64_t iqs = i_cont % QK8_1; // quant index + + ggml_cuda_pdl_sync(); + const float xi = i0 < ne00 ? ggml_cuda_cast(x[i03*s03 + i02*s02 + i01*s01 + i00]) : 0.0f; + float amax = fabsf(xi); + float sum = xi; + + amax = warp_reduce_max(amax); + sum = warp_reduce_sum(sum); + + const float d = amax / 127.0f; + const int8_t q = amax == 0.0f ? 0 : roundf(xi / d); + + y[ib].qs[iqs] = q; + + if (iqs > 0) { + return; + } + + y[ib].ds = make_half2(d, sum); +} + +__device__ __forceinline__ uint8_t compute_e8m0_scale(float amax) { + if (!(amax > 0.0f)) { + return 0; + } + + // FP4 E2M1: max exponent (unbiased) is 2. + constexpr int FP4_E2M1_EMAX = 2; + + const float e = log2f(amax); + + // "even" -> round-to-nearest integer, ties-to-even + const int e_int = __float2int_rn(e); + + const int shared_exp = e_int - FP4_E2M1_EMAX; + + int biased = shared_exp + 127; + + biased = max(biased, 0); + biased = min(biased, 254); + + return static_cast(biased); +} + +// scatter: grid over tokens, quantize once, write to all the token's compact rows +template +static __global__ void quantize_mmq_nvfp4( + const float * __restrict__ x, const int32_t * __restrict__ ids, void * __restrict__ vy, float * __restrict__ scale, + const int64_t ne00, const int64_t s01, const int64_t s02, const int64_t s03, + const int64_t ne0, const int64_t ne1, const int64_t ne2, const int n_expert_used) { +#if defined(BLACKWELL_MMA_AVAILABLE) + + int64_t base_idx; + if constexpr (scatter) { + base_idx = (int64_t) blockIdx.x * s02; // one physical row per token + } else { + const int64_t i2 = blockIdx.y % ne2; + const int64_t i3 = blockIdx.y / ne2; + const int64_t i01 = ids ? ids[blockIdx.x] : blockIdx.x; + base_idx = i3 * s03 + i2 * s02 + i01 * s01; + } + const float * __restrict__ x_row = x + base_idx; + + float amax = 0.0f; + if constexpr (use_aligned_float8) { + for (int64_t i0 = 8 * threadIdx.x; i0 < ne00; i0 += 8 * blockDim.x) { + const float * x_base = x_row + i0; + const float8 v = reinterpret_cast(x_base)[0]; + amax = fmaxf(amax, fabsf(v.x)); + amax = fmaxf(amax, fabsf(v.y)); + amax = fmaxf(amax, fabsf(v.z)); + amax = fmaxf(amax, fabsf(v.w)); + amax = fmaxf(amax, fabsf(v.p)); + amax = fmaxf(amax, fabsf(v.q)); + amax = fmaxf(amax, fabsf(v.r)); + amax = fmaxf(amax, fabsf(v.s)); + } + } else { + for (int64_t i0 = threadIdx.x; i0 < ne00; i0 += blockDim.x) { + amax = fmaxf(amax, fabsf(x_row[i0])); + } + } + + amax = warp_reduce_max(amax); + + __shared__ float warp_amax[CUDA_QUANTIZE_BLOCK_SIZE_MMQ / WARP_SIZE]; + const int lane = threadIdx.x % WARP_SIZE; + const int warp = threadIdx.x / WARP_SIZE; + + if (lane == 0) { + warp_amax[warp] = amax; + } + __syncthreads(); + + if (warp == 0) { + amax = threadIdx.x < int(CUDA_QUANTIZE_BLOCK_SIZE_MMQ / WARP_SIZE) ? warp_amax[lane] : 0.0f; + amax = warp_reduce_max(amax); + if (lane == 0) { + warp_amax[0] = amax / (6.0f * 448.0f); + if constexpr (scatter) { +#pragma unroll + for (int slot = 0; slot < n_expert_used; ++slot) { + const int64_t i = ids[(int64_t) blockIdx.x * n_expert_used + slot]; + scale[i] = warp_amax[0]; + } + } else { + scale[blockIdx.y * ne1 + blockIdx.x] = warp_amax[0]; + } + } + } + __syncthreads(); + + block_fp4_mmq * y = (block_fp4_mmq *) vy; + const int64_t n_subblocks = (ne0 + QK_NVFP4_SUB - 1) / QK_NVFP4_SUB; + + for (int64_t isb = threadIdx.x; isb < n_subblocks; isb += blockDim.x) { + const int64_t i0_base = isb * QK_NVFP4_SUB; + const int64_t k_block = i0_base / QK_FP4_MMQ; + const int sub = (i0_base % QK_FP4_MMQ) / QK_NVFP4_SUB; + + const float row_scale = warp_amax[0]; + const float inv_col_scale = row_scale > 0.0f ? 1.0f / row_scale : 0.0f; + + float vals[QK_NVFP4_SUB]; + if constexpr (use_aligned_float8) { + const float * x_base = x_row + i0_base; + const float8 v0 = i0_base + 7 < ne00 ? reinterpret_cast(x_base)[0] : float8{0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f}; + const float8 v1 = i0_base + 15 < ne00 ? reinterpret_cast(x_base + 8)[0] : float8{0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f}; + vals[0] = v0.x; vals[1] = v0.y; vals[2] = v0.z; vals[3] = v0.w; + vals[4] = v0.p; vals[5] = v0.q; vals[6] = v0.r; vals[7] = v0.s; + vals[8] = v1.x; vals[9] = v1.y; vals[10] = v1.z; vals[11] = v1.w; + vals[12] = v1.p; vals[13] = v1.q; vals[14] = v1.r; vals[15] = v1.s; + } else { +#pragma unroll + for (int k = 0; k < QK_NVFP4_SUB; ++k) { + const int64_t i00 = i0_base + k; + vals[k] = i00 < ne00 ? x_row[i00] : 0.0f; + } + } + + uint32_t q0 = 0; + uint32_t q1 = 0; + + float amax_sub = 0.0f; +#pragma unroll + for (int k = 0; k < QK_NVFP4_SUB; ++k) { + amax_sub = fmaxf(amax_sub, fabsf(vals[k] * inv_col_scale)); + } + + static constexpr int test_offsets[5] = { 0, -1, 1, -2, 2 }; + const int first_fp8_code = (int) ggml_cuda_fp32_to_ue4m3(amax_sub / 6.0f); + + uint8_t fp8_code = (uint8_t) first_fp8_code; + float subblock_scale = ggml_cuda_ue4m3_to_fp32(fp8_code); + float inv_scale_err = subblock_scale > 0.0f ? 0.5f / subblock_scale : 0.0f; +#if CUDART_VERSION >= 12080 + float best_err = nvfp4_native_scale_error(vals, inv_col_scale, inv_scale_err, subblock_scale); +#else + float best_err = 0.0f; +#pragma unroll + for (int k = 0; k < QK_NVFP4_SUB; ++k) { + const float v = vals[k] * inv_col_scale; + const uint8_t q = ggml_cuda_float_to_fp4_e2m1(v, inv_scale_err); + const float err_diff = fabsf(v) - fabsf(kvalues_fp4[q & 0x7]) * subblock_scale; + best_err = fmaf(err_diff, err_diff, best_err); + } +#endif // CUDART_VERSION >= 12080 + +#pragma unroll + for (int i = 1; i < 5; ++i) { + const int test_code = first_fp8_code + test_offsets[i]; + if (test_code < 0 || test_code > 0x7e) { + continue; + } + + const float test_scale = ggml_cuda_ue4m3_to_fp32((uint8_t) test_code); + const float test_inv_scale = test_scale > 0.0f ? 0.5f / test_scale : 0.0f; +#if CUDART_VERSION >= 12080 + const float cur_err = nvfp4_native_scale_error(vals, inv_col_scale, test_inv_scale, test_scale); +#else + float cur_err = 0.0f; +#pragma unroll + for (int k = 0; k < QK_NVFP4_SUB; ++k) { + const float v = vals[k] * inv_col_scale; + const uint8_t q = ggml_cuda_float_to_fp4_e2m1(v, test_inv_scale); + const float err_diff = fabsf(v) - fabsf(kvalues_fp4[q & 0x7]) * test_scale; + cur_err = fmaf(err_diff, err_diff, cur_err); + } +#endif // CUDART_VERSION >= 12080 + + if (cur_err < best_err) { + best_err = cur_err; + fp8_code = (uint8_t) test_code; + subblock_scale = test_scale; + } + } +#if CUDART_VERSION >= 12080 + const float inv_scale = subblock_scale > 0.0f ? 0.5f / subblock_scale : 0.0f; + const float s = inv_col_scale * inv_scale; + + __nv_fp4x4_e2m1 q0_lo(make_float4(vals[0] * s, vals[8] * s, vals[1] * s, vals[9] * s)); + __nv_fp4x4_e2m1 q0_hi(make_float4(vals[2] * s, vals[10] * s, vals[3] * s, vals[11] * s)); + __nv_fp4x4_e2m1 q1_lo(make_float4(vals[4] * s, vals[12] * s, vals[5] * s, vals[13] * s)); + __nv_fp4x4_e2m1 q1_hi(make_float4(vals[6] * s, vals[14] * s, vals[7] * s, vals[15] * s)); + + const char2 q0_lo_c = *reinterpret_cast(&q0_lo); + const char2 q0_hi_c = *reinterpret_cast(&q0_hi); + const char2 q1_lo_c = *reinterpret_cast(&q1_lo); + const char2 q1_hi_c = *reinterpret_cast(&q1_hi); + + q0 = uint32_t(uint8_t(q0_lo_c.x)) | (uint32_t(uint8_t(q0_lo_c.y)) << 8) | + (uint32_t(uint8_t(q0_hi_c.x)) << 16) | (uint32_t(uint8_t(q0_hi_c.y)) << 24); + q1 = uint32_t(uint8_t(q1_lo_c.x)) | (uint32_t(uint8_t(q1_lo_c.y)) << 8) | + (uint32_t(uint8_t(q1_hi_c.x)) << 16) | (uint32_t(uint8_t(q1_hi_c.y)) << 24); +#else + const float inv_scale = subblock_scale > 0.0f ? 0.5f / subblock_scale : 0.0f; +#pragma unroll + for (int k = 0; k < QK_NVFP4_SUB / 4; ++k) { + q0 |= uint32_t(ggml_cuda_float_to_fp4_e2m1(vals[k + 0] * inv_col_scale, inv_scale)) << (8 * k); + q0 |= uint32_t(ggml_cuda_float_to_fp4_e2m1(vals[k + 8] * inv_col_scale, inv_scale)) << (8 * k + 4); + q1 |= uint32_t(ggml_cuda_float_to_fp4_e2m1(vals[k + 4] * inv_col_scale, inv_scale)) << (8 * k); + q1 |= uint32_t(ggml_cuda_float_to_fp4_e2m1(vals[k + 12] * inv_col_scale, inv_scale)) << (8 * k + 4); + } +#endif // CUDART_VERSION >= 12080 + + if constexpr (scatter) { +#pragma unroll + for (int slot = 0; slot < n_expert_used; ++slot) { + const int64_t i = ids[(int64_t) blockIdx.x * n_expert_used + slot]; + block_fp4_mmq * yb = y + (k_block * ne1 + i); + uint32_t * yqs = reinterpret_cast(yb->qs); + yqs[2 * sub + 0] = q0; + yqs[2 * sub + 1] = q1; + reinterpret_cast(yb->d4)[sub] = fp8_code; + } + } else { + const int64_t blocks_per_col = (ne0 + QK_FP4_MMQ - 1) / QK_FP4_MMQ; + block_fp4_mmq * yb = y + (blockIdx.y * (blocks_per_col * ne1) + k_block * ne1 + blockIdx.x); + uint32_t * yqs = reinterpret_cast(yb->qs); + yqs[2 * sub + 0] = q0; + yqs[2 * sub + 1] = q1; + reinterpret_cast(yb->d4)[sub] = fp8_code; + } + } +#else + GGML_UNUSED_VARS(x, ids, vy, scale, ne00, s01, s02, s03, ne0, ne1, ne2, n_expert_used); + NO_DEVICE_CODE; // This is for Blackwell NVFP4 activations only. +#endif // defined(BLACKWELL_MMA_AVAILABLE) + +} + +// quantize values in the format mxfp4 is stored which is interleaved nibbles +// i.e. a block a0-a31 is represented as a0a16,a1a17 ...a15a31 +// scatter: grid over tokens, quantize once, write to all the token's compact rows +template +static __global__ void quantize_mmq_mxfp4(const float * __restrict__ x, + const int32_t * __restrict__ ids, + void * __restrict__ vy, + const int64_t ne00, + const int64_t s01, + const int64_t s02, + const int64_t s03, + const int64_t ne0, + const int ne1, + const int ne2, + const int n_expert_used) { + constexpr int vals_per_scale = 32; + constexpr int vals_per_warp = 2 * vals_per_scale; // Each warp processes 2 blocks of 32 = 64 values + + const int warp_id = threadIdx.y; + const int lane_id_32 = threadIdx.x; + + const int nwarps = blockDim.y; + + const int64_t warp_start_offset = (blockIdx.y * nwarps + warp_id) * vals_per_warp; + + if (warp_start_offset >= ne0) { + return; + } + + const int64_t block_fp4_mmq_size = QK_FP4_MMQ; + const int64_t k_block = warp_start_offset / block_fp4_mmq_size; + const int64_t quad_idx_in_block = (warp_start_offset % block_fp4_mmq_size) / vals_per_warp; + + const int group_id = lane_id_32 / 4; + const int lane_in_group = lane_id_32 % 4; + const int base = group_id * 2; + + ggml_cuda_pdl_sync(); + int64_t base_pos; + if constexpr (scatter) { + base_pos = (int64_t) blockIdx.x * s02; // one physical row per token + } else { + const int64_t i2 = blockIdx.z % ne2; + const int64_t i3 = blockIdx.z / ne2; + const int64_t i01 = ids ? ids[blockIdx.x] : blockIdx.x; + base_pos = i3 * s03 + i2 * s02 + i01 * s01; + } + + uint8_t scales[2]; + char2 packed[2]; + +#pragma unroll + for (int b = 0; b < 2; ++b) { + const int64_t i0 = warp_start_offset + b * vals_per_scale + lane_id_32; + const float xi = (i0 < ne00) ? x[base_pos + i0] : 0.0f; + + float amax = fabsf(xi); +#pragma unroll + for (int mask = 16; mask > 0; mask >>= 1) { + amax = fmaxf(amax, __shfl_xor_sync(0xFFFFFFFF, amax, mask, WARP_SIZE)); + } + + const uint8_t e = compute_e8m0_scale(amax); + scales[b] = e; + const float inv_s = (amax == 0.0f) ? 0.0f : __frcp_rn(ggml_cuda_e8m0_to_fp32(e)); + +#if CUDART_VERSION >= 12080 + const float scaled_val = xi * inv_s; + + const float val0 = __shfl_sync(0xFFFFFFFF, scaled_val, base, WARP_SIZE); + const float val1 = __shfl_sync(0xFFFFFFFF, scaled_val, base + 16, WARP_SIZE); + const float val2 = __shfl_sync(0xFFFFFFFF, scaled_val, base + 1, WARP_SIZE); + const float val3 = __shfl_sync(0xFFFFFFFF, scaled_val, base + 17, WARP_SIZE); + + __nv_fp4x4_e2m1 fp4_packed(make_float4(val0, val1, val2, val3)); + packed[b] = *(char2 *) &fp4_packed; +#else + // Fallback: manual FP4 conversion using LUT + const uint8_t q_val = ggml_cuda_float_to_fp4_e2m1(xi, inv_s); + + const uint8_t q_lo_0 = __shfl_sync(0xFFFFFFFF, q_val, base, WARP_SIZE); + const uint8_t q_lo_1 = __shfl_sync(0xFFFFFFFF, q_val, base + 1, WARP_SIZE); + const uint8_t q_hi_0 = __shfl_sync(0xFFFFFFFF, q_val, base + 16, WARP_SIZE); + const uint8_t q_hi_1 = __shfl_sync(0xFFFFFFFF, q_val, base + 17, WARP_SIZE); + + char2 q; + q.x = (q_hi_0 << 4) | q_lo_0; + q.y = (q_hi_1 << 4) | q_lo_1; + packed[b] = q; +#endif // CUDART_VERSION >= 12080 + } + + block_fp4_mmq * y = (block_fp4_mmq *) vy; + if constexpr (scatter) { +#pragma unroll + for (int slot = 0; slot < n_expert_used; ++slot) { + const int64_t i = ids[(int64_t) blockIdx.x * n_expert_used + slot]; + block_fp4_mmq * yb = y + (k_block * ne1 + i); + char2 * yqs2 = (char2 *) yb->qs; + if (lane_in_group == 0) { + yqs2[quad_idx_in_block * 16 + 0 * 8 + group_id] = packed[0]; + yqs2[quad_idx_in_block * 16 + 1 * 8 + group_id] = packed[1]; + } + if (lane_id_32 == 0) { + yb->d4[quad_idx_in_block] = (scales[1] << 8) | scales[0]; + } + } + } else { + const int64_t ib0 = blockIdx.z * ((int64_t) ne1 * (ne0 / block_fp4_mmq_size)); + block_fp4_mmq * yb = y + (ib0 + k_block * ne1 + blockIdx.x); + char2 * yqs2 = (char2 *) yb->qs; + if (lane_in_group == 0) { + yqs2[quad_idx_in_block * 16 + 0 * 8 + group_id] = packed[0]; + yqs2[quad_idx_in_block * 16 + 1 * 8 + group_id] = packed[1]; + } + if (lane_id_32 == 0) { + yb->d4[quad_idx_in_block] = (scales[1] << 8) | scales[0]; + } + } + GGML_UNUSED(n_expert_used); +} + +// FreeToken: 4 consecutive activations as float, from fp32, bf16 or fp16 +static __device__ __forceinline__ float4 load_float4(const float * x) { + return *(const float4 *) x; +} + +static __device__ __forceinline__ float4 load_float4(const nv_bfloat16 * x) { + const uint2 v = *(const uint2 *) x; + const float2 lo = __bfloat1622float2(*(const nv_bfloat162 *) &v.x); + const float2 hi = __bfloat1622float2(*(const nv_bfloat162 *) &v.y); + return make_float4(lo.x, lo.y, hi.x, hi.y); +} + +static __device__ __forceinline__ float4 load_float4(const half * x) { + const uint2 v = *(const uint2 *) x; + const float2 lo = __half22float2(*(const half2 *) &v.x); + const float2 hi = __half22float2(*(const half2 *) &v.y); + return make_float4(lo.x, lo.y, hi.x, hi.y); +} + +// scatter: grid over tokens, quantize once, write to all the token's compact rows +template +static __global__ void quantize_mmq_q8_1( + const src_t * __restrict__ x, const int32_t * __restrict__ ids, void * __restrict__ vy, + const int64_t ne00, const int64_t s01, const int64_t s02, const int64_t s03, + const int64_t ne0, const int ne1, const int ne2, const int n_expert_used) { + + constexpr int vals_per_scale = ds_layout == MMQ_Q8_1_DS_LAYOUT_D2S6 ? 64 : 32; + constexpr int vals_per_sum = ds_layout == MMQ_Q8_1_DS_LAYOUT_D2S6 ? 16 : 32; + + const int64_t i0 = ((int64_t)blockDim.x*blockIdx.y + threadIdx.x)*4; + + if (i0 >= ne0) { + return; + } + + const int64_t i00 = i0; + ggml_cuda_pdl_sync(); + + int64_t base_idx; + if constexpr (scatter) { + base_idx = (int64_t) blockIdx.x * s02; // one physical row per token + } else { + const int64_t i2 = blockIdx.z % ne2; + const int64_t i3 = blockIdx.z / ne2; + const int64_t i01 = ids ? ids[blockIdx.x] : blockIdx.x; + base_idx = i3*s03 + i2*s02 + i01*s01; + } + + block_q8_1_mmq * y = (block_q8_1_mmq *) vy; + + const int64_t k_block = i0 / QK8_1_MMQ; // column block in the channel + const int64_t iqs = i0 % QK8_1_MMQ; // quant index in block + + // Load 4 floats per thread and calculate max. abs. value between them: + const float4 xi = i0 < ne00 ? load_float4(x + base_idx + i00) : make_float4(0.0f, 0.0f, 0.0f, 0.0f); + float amax = fabsf(xi.x); + amax = fmaxf(amax, fabsf(xi.y)); + amax = fmaxf(amax, fabsf(xi.z)); + amax = fmaxf(amax, fabsf(xi.w)); + + // Exchange max. abs. value between vals_per_scale/4 threads. +#pragma unroll + for (int offset = vals_per_scale/8; offset > 0; offset >>= 1) { + amax = fmaxf(amax, __shfl_xor_sync(0xFFFFFFFF, amax, offset, WARP_SIZE)); + } + + float sum; + if (ds_layout != MMQ_Q8_1_DS_LAYOUT_D4) { + sum = xi.x + xi.y + xi.z + xi.w; + + // Calculate sums across vals_per_sum/4 threads. +#pragma unroll + for (int offset = vals_per_sum/8; offset > 0; offset >>= 1) { + sum += __shfl_xor_sync(0xFFFFFFFF, sum, offset, WARP_SIZE); + } + } + + const float d_inv = 127.0f / amax; + char4 q; + q.x = roundf(xi.x*d_inv); + q.y = roundf(xi.y*d_inv); + q.z = roundf(xi.z*d_inv); + q.w = roundf(xi.w*d_inv); + const float d = 1.0f / d_inv; + + // write the block once (normal) or to each of the token's compact rows (scatter) + const int nwrite = scatter ? n_expert_used : 1; +#pragma unroll + for (int slot = 0; slot < nwrite; ++slot) { + int64_t ib; + if constexpr (scatter) { + const int64_t i = ids[(int64_t) blockIdx.x * n_expert_used + slot]; + ib = k_block*ne1 + i; + } else { + const int64_t ib0 = blockIdx.z*((int64_t)gridDim.x*gridDim.y*blockDim.x/QK8_1); // first block of channel + ib = ib0 + k_block*ne1 + blockIdx.x; + } + + // Write back 4 int8 values as a single 32 bit value for better memory bandwidth: + char4 * yqs4 = (char4 *) y[ib].qs; + yqs4[iqs/4] = q; + + if (ds_layout == MMQ_Q8_1_DS_LAYOUT_D2S6) { + if (iqs % 16 == 0 && iqs < 96) { + y[ib].d2s6[2 + iqs/16] = sum; + if (iqs % 64 == 0) { + y[ib].d2s6[iqs/64] = d; + } + } + } else if (iqs % 32 == 0) { + if (ds_layout == MMQ_Q8_1_DS_LAYOUT_DS4) { + y[ib].ds4[iqs/32] = make_half2(d, sum); + } else { + y[ib].d4[iqs/32] = d; + } + } + } + GGML_UNUSED(n_expert_used); +} + +template +void quantize_row_q8_1_cuda( + const src_t * x, const int32_t * ids, void * vy, const ggml_type type_src0, + const int64_t ne00, const int64_t s01, const int64_t s02, const int64_t s03, + const int64_t ne0, const int64_t ne1, const int64_t ne2, const int64_t ne3, cudaStream_t stream) { + GGML_ASSERT(!ids); + GGML_ASSERT(ne0 % QK8_1 == 0); + + const uint3 ne2_fastdiv = init_fastdiv_values(ne2); + + const int64_t block_num_x = (ne0 + CUDA_QUANTIZE_BLOCK_SIZE - 1) / CUDA_QUANTIZE_BLOCK_SIZE; + const dim3 num_blocks(block_num_x, ne1, ne2*ne3); + const dim3 block_size(CUDA_QUANTIZE_BLOCK_SIZE, 1, 1); + const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(num_blocks, block_size, 0, stream); + ggml_cuda_kernel_launch(quantize_q8_1, launch_params, x, vy, ne00, s01, s02, s03, ne0, ne1, ne2_fastdiv); + GGML_UNUSED(type_src0); +} + +template +void quantize_mmq_q8_1_cuda( + const src_t * x, const int32_t * ids, void * vy, const ggml_type type_src0, + const int64_t ne00, const int64_t s01, const int64_t s02, const int64_t s03, + const int64_t ne0, const int64_t ne1, const int64_t ne2, const int64_t ne3, cudaStream_t stream) { + GGML_ASSERT(ne00 % 4 == 0); + GGML_ASSERT(ne0 % QK8_1_MMQ == 0); + + // ne1 tends to assume the highest values, therefore use it as the "x" dimension of the CUDA grid: + const int64_t block_num_y = (ne0 + 4*CUDA_QUANTIZE_BLOCK_SIZE_MMQ - 1) / (4*CUDA_QUANTIZE_BLOCK_SIZE_MMQ); + const dim3 num_blocks(ne1, block_num_y, ne2*ne3); + const dim3 block_size(CUDA_QUANTIZE_BLOCK_SIZE_MMQ, 1, 1); + switch (mmq_get_q8_1_ds_layout(type_src0)) { + case MMQ_Q8_1_DS_LAYOUT_D4: + quantize_mmq_q8_1 + <<>>(x, ids, vy, ne00, s01, s02, s03, ne0, ne1, ne2, /*n_expert_used=*/0); + break; + case MMQ_Q8_1_DS_LAYOUT_DS4: + quantize_mmq_q8_1 + <<>>(x, ids, vy, ne00, s01, s02, s03, ne0, ne1, ne2, /*n_expert_used=*/0); + break; + case MMQ_Q8_1_DS_LAYOUT_D2S6: + quantize_mmq_q8_1 + <<>>(x, ids, vy, ne00, s01, s02, s03, ne0, ne1, ne2, /*n_expert_used=*/0); + break; + default: + GGML_ABORT("fatal error"); + } +} + +// scatter=true reuses the quant kernel: grid over tokens, ids = inverse map (token slot -> compact row) +template +void quantize_scatter_mmq_q8_1_cuda( + const src_t * x, const int32_t * ids_src1_inv, void * vy, const ggml_type type_src0, + const int64_t ne00, const int64_t stride_token, const int64_t ne0, + const int64_t n_tokens, const int64_t nrows_dst, const int n_expert_used, cudaStream_t stream) { + GGML_ASSERT(ne00 % 4 == 0); + GGML_ASSERT(ne0 % QK8_1_MMQ == 0); + + const int64_t block_num_y = (ne0 + 4*CUDA_QUANTIZE_BLOCK_SIZE_MMQ - 1) / (4*CUDA_QUANTIZE_BLOCK_SIZE_MMQ); + const dim3 num_blocks(n_tokens, block_num_y, 1); + const dim3 block_size(CUDA_QUANTIZE_BLOCK_SIZE_MMQ, 1, 1); + switch (mmq_get_q8_1_ds_layout(type_src0)) { + case MMQ_Q8_1_DS_LAYOUT_D4: + quantize_mmq_q8_1<<>>( + x, ids_src1_inv, vy, ne00, /*s01=*/0, /*s02=*/stride_token, /*s03=*/0, ne0, /*ne1=*/(int) nrows_dst, /*ne2=*/1, n_expert_used); + break; + case MMQ_Q8_1_DS_LAYOUT_DS4: + quantize_mmq_q8_1<<>>( + x, ids_src1_inv, vy, ne00, /*s01=*/0, /*s02=*/stride_token, /*s03=*/0, ne0, /*ne1=*/(int) nrows_dst, /*ne2=*/1, n_expert_used); + break; + case MMQ_Q8_1_DS_LAYOUT_D2S6: + quantize_mmq_q8_1<<>>( + x, ids_src1_inv, vy, ne00, /*s01=*/0, /*s02=*/stride_token, /*s03=*/0, ne0, /*ne1=*/(int) nrows_dst, /*ne2=*/1, n_expert_used); + break; + default: + GGML_ABORT("fatal error"); + } +} + +// scatter=true reuses the quant kernels: grid over tokens, ids = inverse map (token slot -> compact row) +void quantize_scatter_mmq_fp4_cuda( + const float * x, const int32_t * ids_src1_inv, void * vy, float * scale, const ggml_type type_src0, const bool use_aligned_float8, + const int64_t ne00, const int64_t stride_token, const int64_t ne0, + const int64_t n_tokens, const int64_t nrows_dst, const int n_expert_used, cudaStream_t stream) { + GGML_ASSERT(ne0 > 0); + if (type_src0 == GGML_TYPE_NVFP4) { + GGML_ASSERT(scale); + GGML_ASSERT(ne00 % QK_NVFP4 == 0); + const dim3 block_size(CUDA_QUANTIZE_BLOCK_SIZE_MMQ, 1, 1); + const dim3 num_blocks(n_tokens, 1, 1); + if (use_aligned_float8) { + quantize_mmq_nvfp4<<>>( + x, ids_src1_inv, vy, scale, ne00, /*s01=*/0, /*s02=*/stride_token, /*s03=*/0, ne0, /*ne1=*/nrows_dst, /*ne2=*/1, n_expert_used); + } else { + quantize_mmq_nvfp4<<>>( + x, ids_src1_inv, vy, scale, ne00, /*s01=*/0, /*s02=*/stride_token, /*s03=*/0, ne0, /*ne1=*/nrows_dst, /*ne2=*/1, n_expert_used); + } + } else { + GGML_ASSERT(type_src0 == GGML_TYPE_MXFP4); + constexpr int nwarps = 8; + constexpr int vals_per_block = nwarps * 2 * QK_MXFP4; + const int64_t block_num_y = (ne0 + vals_per_block - 1) / vals_per_block; + const dim3 block_size(WARP_SIZE, nwarps, 1); + const dim3 num_blocks(n_tokens, block_num_y, 1); + quantize_mmq_mxfp4<<>>( + x, ids_src1_inv, vy, ne00, /*s01=*/0, /*s02=*/stride_token, /*s03=*/0, ne0, /*ne1=*/(int) nrows_dst, /*ne2=*/1, n_expert_used); + } +} + +void quantize_mmq_fp4_cuda( + const float * x, const int32_t * ids, void * vy, float * scale, const ggml_type type_src0, const bool use_aligned_float8, + const int64_t ne00, const int64_t s01, const int64_t s02, const int64_t s03, + const int64_t ne0, const int64_t ne1, const int64_t ne2, const int64_t ne3, cudaStream_t stream) { + GGML_ASSERT(type_src0 == GGML_TYPE_MXFP4 || type_src0 == GGML_TYPE_NVFP4); + GGML_ASSERT(ne0 > 0); + + if (type_src0 == GGML_TYPE_NVFP4) { + GGML_ASSERT(scale); + GGML_ASSERT(ne00 % QK_NVFP4 == 0); + const dim3 block_size(CUDA_QUANTIZE_BLOCK_SIZE_MMQ, 1, 1); + const dim3 num_blocks(ne1, ne2 * ne3, 1); + if (use_aligned_float8) { + quantize_mmq_nvfp4<<>>( + x, ids, vy, scale, ne00, s01, s02, s03, ne0, ne1, ne2, /*n_expert_used=*/0); + } else { + quantize_mmq_nvfp4<<>>( + x, ids, vy, scale, ne00, s01, s02, s03, ne0, ne1, ne2, /*n_expert_used=*/0); + } + } else { + GGML_ASSERT(ne0 % (2 * QK_MXFP4) == 0); + + constexpr int nwarps = 8; + constexpr int vals_per_warp = 2 * QK_MXFP4; + constexpr int vals_per_block = nwarps * vals_per_warp; + + const int64_t block_num_y = (ne0 + vals_per_block - 1) / vals_per_block; + const dim3 num_blocks(ne1, block_num_y, ne2 * ne3); + const dim3 block_size(WARP_SIZE, nwarps, 1); + + quantize_mmq_mxfp4<<>>(x, ids, vy, ne00, s01, s02, s03, ne0, ne1, ne2, /*n_expert_used=*/0); + } +} + +#define FREETOKEN_QUANTIZE_INSTANCES(src_t) \ + template void quantize_row_q8_1_cuda(const src_t *, const int32_t *, void *, ggml_type, int64_t, int64_t, \ + int64_t, int64_t, int64_t, int64_t, int64_t, int64_t, cudaStream_t); \ + template void quantize_mmq_q8_1_cuda(const src_t *, const int32_t *, void *, ggml_type, int64_t, int64_t, \ + int64_t, int64_t, int64_t, int64_t, int64_t, int64_t, cudaStream_t); \ + template void quantize_scatter_mmq_q8_1_cuda(const src_t *, const int32_t *, void *, ggml_type, int64_t, \ + int64_t, int64_t, int64_t, int64_t, int, cudaStream_t); + +FREETOKEN_QUANTIZE_INSTANCES(float) +FREETOKEN_QUANTIZE_INSTANCES(nv_bfloat16) +FREETOKEN_QUANTIZE_INSTANCES(half) diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/quantize.cuh b/python/freetoken/kernel/csrc/ggml/ggml-cuda/quantize.cuh new file mode 100644 index 0000000000..99283ea6e1 --- /dev/null +++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/quantize.cuh @@ -0,0 +1,74 @@ +#pragma once + +#include "common.cuh" +#include "mmq.cuh" + +#include + +#define CUDA_QUANTIZE_BLOCK_SIZE 256 +#define CUDA_QUANTIZE_BLOCK_SIZE_MMQ 128 + +static_assert(MATRIX_ROW_PADDING % CUDA_QUANTIZE_BLOCK_SIZE == 0, "Risk of out-of-bounds access."); +static_assert(MATRIX_ROW_PADDING % (4*CUDA_QUANTIZE_BLOCK_SIZE_MMQ) == 0, "Risk of out-of-bounds access."); + +typedef void (*quantize_cuda_t)( + const float * x, const int32_t * ids, void * vy, + ggml_type type_src0, int64_t ne00, int64_t s01, int64_t s02, int64_t s03, + int64_t ne0, int64_t ne1, int64_t ne2, int64_t ne3, cudaStream_t stream); + +// FreeToken: src_t is float, nv_bfloat16 or half +template +void quantize_row_q8_1_cuda( + const src_t * x, const int32_t * ids, void * vy, + ggml_type type_src0, int64_t ne00, int64_t s01, int64_t s02, int64_t s03, + int64_t ne0, int64_t ne1, int64_t ne2, int64_t ne3, cudaStream_t stream); + +template +void quantize_mmq_q8_1_cuda( + const src_t * x, const int32_t * ids, void * vy, + ggml_type type_src0, int64_t ne00, int64_t s01, int64_t s02, int64_t s03, + int64_t ne0, int64_t ne1, int64_t ne2, int64_t ne3, cudaStream_t stream); + +void quantize_mmq_fp4_cuda(const float * x, + const int32_t * ids, + void * vy, + float * scale, + ggml_type type_src0, + bool use_aligned_float8, + int64_t ne00, + int64_t s01, + int64_t s02, + int64_t s03, + int64_t ne0, + int64_t ne1, + int64_t ne2, + int64_t ne3, + cudaStream_t stream); + +// quantize each token once and scatter the block to its compact rows (via the inverse map) +void quantize_scatter_mmq_fp4_cuda(const float * x, + const int32_t * ids_src1_inv, + void * vy, + float * scale, + ggml_type type_src0, + bool use_aligned_float8, + int64_t ne00, + int64_t stride_token, + int64_t ne0, + int64_t n_tokens, + int64_t nrows_dst, + int n_expert_used, + cudaStream_t stream); + +template +void quantize_scatter_mmq_q8_1_cuda(const src_t * x, + const int32_t * ids_src1_inv, + void * vy, + ggml_type type_src0, + int64_t ne00, + int64_t stride_token, + int64_t ne0, + int64_t n_tokens, + int64_t nrows_dst, + int n_expert_used, + cudaStream_t stream); diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-iq1_s.cu b/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-iq1_s.cu new file mode 100644 index 0000000000..84ec850294 --- /dev/null +++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-iq1_s.cu @@ -0,0 +1,5 @@ +// This file has been autogenerated by generate_cu_files.py, do not edit manually. + +#include "../mmq.cuh" + +DECL_MMQ_CASE(GGML_TYPE_IQ1_S); diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-iq2_s.cu b/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-iq2_s.cu new file mode 100644 index 0000000000..583c4e5a51 --- /dev/null +++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-iq2_s.cu @@ -0,0 +1,5 @@ +// This file has been autogenerated by generate_cu_files.py, do not edit manually. + +#include "../mmq.cuh" + +DECL_MMQ_CASE(GGML_TYPE_IQ2_S); diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-iq2_xs.cu b/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-iq2_xs.cu new file mode 100644 index 0000000000..edaf1560de --- /dev/null +++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-iq2_xs.cu @@ -0,0 +1,5 @@ +// This file has been autogenerated by generate_cu_files.py, do not edit manually. + +#include "../mmq.cuh" + +DECL_MMQ_CASE(GGML_TYPE_IQ2_XS); diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-iq2_xxs.cu b/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-iq2_xxs.cu new file mode 100644 index 0000000000..233d9342c9 --- /dev/null +++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-iq2_xxs.cu @@ -0,0 +1,5 @@ +// This file has been autogenerated by generate_cu_files.py, do not edit manually. + +#include "../mmq.cuh" + +DECL_MMQ_CASE(GGML_TYPE_IQ2_XXS); diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-iq3_s.cu b/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-iq3_s.cu new file mode 100644 index 0000000000..6092dc7136 --- /dev/null +++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-iq3_s.cu @@ -0,0 +1,5 @@ +// This file has been autogenerated by generate_cu_files.py, do not edit manually. + +#include "../mmq.cuh" + +DECL_MMQ_CASE(GGML_TYPE_IQ3_S); diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-iq3_xxs.cu b/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-iq3_xxs.cu new file mode 100644 index 0000000000..1d5bd201fb --- /dev/null +++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-iq3_xxs.cu @@ -0,0 +1,5 @@ +// This file has been autogenerated by generate_cu_files.py, do not edit manually. + +#include "../mmq.cuh" + +DECL_MMQ_CASE(GGML_TYPE_IQ3_XXS); diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-iq4_nl.cu b/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-iq4_nl.cu new file mode 100644 index 0000000000..eb02fab002 --- /dev/null +++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-iq4_nl.cu @@ -0,0 +1,5 @@ +// This file has been autogenerated by generate_cu_files.py, do not edit manually. + +#include "../mmq.cuh" + +DECL_MMQ_CASE(GGML_TYPE_IQ4_NL); diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-iq4_xs.cu b/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-iq4_xs.cu new file mode 100644 index 0000000000..1eb3b74307 --- /dev/null +++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-iq4_xs.cu @@ -0,0 +1,5 @@ +// This file has been autogenerated by generate_cu_files.py, do not edit manually. + +#include "../mmq.cuh" + +DECL_MMQ_CASE(GGML_TYPE_IQ4_XS); diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-mxfp4.cu b/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-mxfp4.cu new file mode 100644 index 0000000000..0c1c7f291b --- /dev/null +++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-mxfp4.cu @@ -0,0 +1,7 @@ +// This file has been autogenerated by generate_cu_files.py, do not edit manually. + +#include "../mmq.cuh" + +DECL_MMQ_CASE(GGML_TYPE_MXFP4); + +DECL_MMQ_CASE_W4A4(GGML_TYPE_MXFP4); diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-nvfp4.cu b/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-nvfp4.cu new file mode 100644 index 0000000000..bed8c8e2e4 --- /dev/null +++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-nvfp4.cu @@ -0,0 +1,7 @@ +// This file has been autogenerated by generate_cu_files.py, do not edit manually. + +#include "../mmq.cuh" + +DECL_MMQ_CASE(GGML_TYPE_NVFP4); + +DECL_MMQ_CASE_W4A4(GGML_TYPE_NVFP4); diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-q1_0.cu b/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-q1_0.cu new file mode 100644 index 0000000000..f0686b0d0d --- /dev/null +++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-q1_0.cu @@ -0,0 +1,5 @@ +// This file has been autogenerated by generate_cu_files.py, do not edit manually. + +#include "../mmq.cuh" + +DECL_MMQ_CASE(GGML_TYPE_Q1_0); diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-q2_0.cu b/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-q2_0.cu new file mode 100644 index 0000000000..750180e330 --- /dev/null +++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-q2_0.cu @@ -0,0 +1,5 @@ +// This file has been autogenerated by generate_cu_files.py, do not edit manually. + +#include "../mmq.cuh" + +DECL_MMQ_CASE(GGML_TYPE_Q2_0); diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-q2_k.cu b/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-q2_k.cu new file mode 100644 index 0000000000..6415369dc1 --- /dev/null +++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-q2_k.cu @@ -0,0 +1,5 @@ +// This file has been autogenerated by generate_cu_files.py, do not edit manually. + +#include "../mmq.cuh" + +DECL_MMQ_CASE(GGML_TYPE_Q2_K); diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-q3_k.cu b/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-q3_k.cu new file mode 100644 index 0000000000..ffb6213af8 --- /dev/null +++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-q3_k.cu @@ -0,0 +1,5 @@ +// This file has been autogenerated by generate_cu_files.py, do not edit manually. + +#include "../mmq.cuh" + +DECL_MMQ_CASE(GGML_TYPE_Q3_K); diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-q4_0.cu b/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-q4_0.cu new file mode 100644 index 0000000000..0c0b0c8a8e --- /dev/null +++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-q4_0.cu @@ -0,0 +1,5 @@ +// This file has been autogenerated by generate_cu_files.py, do not edit manually. + +#include "../mmq.cuh" + +DECL_MMQ_CASE(GGML_TYPE_Q4_0); diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-q4_1.cu b/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-q4_1.cu new file mode 100644 index 0000000000..ee67f6942a --- /dev/null +++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-q4_1.cu @@ -0,0 +1,5 @@ +// This file has been autogenerated by generate_cu_files.py, do not edit manually. + +#include "../mmq.cuh" + +DECL_MMQ_CASE(GGML_TYPE_Q4_1); diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-q4_k.cu b/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-q4_k.cu new file mode 100644 index 0000000000..9eeb3cd7f3 --- /dev/null +++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-q4_k.cu @@ -0,0 +1,5 @@ +// This file has been autogenerated by generate_cu_files.py, do not edit manually. + +#include "../mmq.cuh" + +DECL_MMQ_CASE(GGML_TYPE_Q4_K); diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-q5_0.cu b/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-q5_0.cu new file mode 100644 index 0000000000..cc57fb9753 --- /dev/null +++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-q5_0.cu @@ -0,0 +1,5 @@ +// This file has been autogenerated by generate_cu_files.py, do not edit manually. + +#include "../mmq.cuh" + +DECL_MMQ_CASE(GGML_TYPE_Q5_0); diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-q5_1.cu b/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-q5_1.cu new file mode 100644 index 0000000000..721ac790c4 --- /dev/null +++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-q5_1.cu @@ -0,0 +1,5 @@ +// This file has been autogenerated by generate_cu_files.py, do not edit manually. + +#include "../mmq.cuh" + +DECL_MMQ_CASE(GGML_TYPE_Q5_1); diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-q5_k.cu b/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-q5_k.cu new file mode 100644 index 0000000000..a2e90ffd5d --- /dev/null +++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-q5_k.cu @@ -0,0 +1,5 @@ +// This file has been autogenerated by generate_cu_files.py, do not edit manually. + +#include "../mmq.cuh" + +DECL_MMQ_CASE(GGML_TYPE_Q5_K); diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-q6_k.cu b/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-q6_k.cu new file mode 100644 index 0000000000..470938fef8 --- /dev/null +++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-q6_k.cu @@ -0,0 +1,5 @@ +// This file has been autogenerated by generate_cu_files.py, do not edit manually. + +#include "../mmq.cuh" + +DECL_MMQ_CASE(GGML_TYPE_Q6_K); diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-q8_0.cu b/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-q8_0.cu new file mode 100644 index 0000000000..974477bbb7 --- /dev/null +++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/template-instances/mmq-instance-q8_0.cu @@ -0,0 +1,5 @@ +// This file has been autogenerated by generate_cu_files.py, do not edit manually. + +#include "../mmq.cuh" + +DECL_MMQ_CASE(GGML_TYPE_Q8_0); diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/unary.cuh b/python/freetoken/kernel/csrc/ggml/ggml-cuda/unary.cuh new file mode 100644 index 0000000000..04f3af6443 --- /dev/null +++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/unary.cuh @@ -0,0 +1,123 @@ +#pragma once +#include "common.cuh" + +#define CUDA_NEG_BLOCK_SIZE 256 +#define CUDA_STEP_BLOCK_SIZE 256 +#define CUDA_GELU_BLOCK_SIZE 256 +#define CUDA_SILU_BLOCK_SIZE 256 +#define CUDA_SILU_BACK_BLOCK_SIZE 256 +#define CUDA_TANH_BLOCK_SIZE 256 +#define CUDA_RELU_BLOCK_SIZE 256 +#define CUDA_SIGMOID_BLOCK_SIZE 256 +#define CUDA_HARDSIGMOID_BLOCK_SIZE 256 +#define CUDA_EXP_BLOCK_SIZE 256 +#define CUDA_HARDSWISH_BLOCK_SIZE 256 +#define CUDA_SQR_BLOCK_SIZE 256 +#define CUDA_SQRT_BLOCK_SIZE 256 +#define CUDA_SIN_BLOCK_SIZE 256 +#define CUDA_COS_BLOCK_SIZE 256 +#define CUDA_GLU_BLOCK_SIZE 256 +#define CUDA_XIELU_BLOCK_SIZE 256 + +void ggml_cuda_op_abs(ggml_backend_cuda_context & ctx, ggml_tensor * dst); + +void ggml_cuda_op_sgn(ggml_backend_cuda_context & ctx, ggml_tensor * dst); + +void ggml_cuda_op_neg(ggml_backend_cuda_context & ctx, ggml_tensor * dst); + +void ggml_cuda_op_step(ggml_backend_cuda_context & ctx, ggml_tensor * dst); + +void ggml_cuda_op_gelu(ggml_backend_cuda_context & ctx, ggml_tensor * dst); + +void ggml_cuda_op_silu(ggml_backend_cuda_context & ctx, ggml_tensor * dst); + +void ggml_cuda_op_silu_back(ggml_backend_cuda_context & ctx, ggml_tensor * dst); + +void ggml_cuda_op_gelu_erf(ggml_backend_cuda_context & ctx, ggml_tensor * dst); + +void ggml_cuda_op_gelu_quick(ggml_backend_cuda_context & ctx, ggml_tensor * dst); + +void ggml_cuda_op_tanh(ggml_backend_cuda_context & ctx, ggml_tensor * dst); + +void ggml_cuda_op_relu(ggml_backend_cuda_context & ctx, ggml_tensor * dst); + +void ggml_cuda_op_sigmoid(ggml_backend_cuda_context & ctx, ggml_tensor * dst); + +void ggml_cuda_op_hardsigmoid(ggml_backend_cuda_context & ctx, ggml_tensor * dst); + +void ggml_cuda_op_exp(ggml_backend_cuda_context & ctx, ggml_tensor * dst); + +void ggml_cuda_op_hardswish(ggml_backend_cuda_context & ctx, ggml_tensor * dst); + +void ggml_cuda_op_leaky_relu(ggml_backend_cuda_context & ctx, ggml_tensor * dst); + +void ggml_cuda_op_sqr(ggml_backend_cuda_context & ctx, ggml_tensor * dst); + +void ggml_cuda_op_sqrt(ggml_backend_cuda_context & ctx, ggml_tensor * dst); + +void ggml_cuda_op_sin(ggml_backend_cuda_context & ctx, ggml_tensor * dst); + +void ggml_cuda_op_cos(ggml_backend_cuda_context & ctx, ggml_tensor * dst); + +void ggml_cuda_op_log(ggml_backend_cuda_context & ctx, ggml_tensor * dst); + +void ggml_cuda_op_expm1(ggml_backend_cuda_context & ctx, ggml_tensor * dst); + +void ggml_cuda_op_softplus(ggml_backend_cuda_context & ctx, ggml_tensor * dst); + +void ggml_cuda_op_elu(ggml_backend_cuda_context & ctx, ggml_tensor * dst); + +void ggml_cuda_op_floor(ggml_backend_cuda_context & ctx, ggml_tensor * dst); + +void ggml_cuda_op_ceil(ggml_backend_cuda_context & ctx, ggml_tensor * dst); + +void ggml_cuda_op_round(ggml_backend_cuda_context & ctx, ggml_tensor * dst); + +void ggml_cuda_op_trunc(ggml_backend_cuda_context & ctx, ggml_tensor * dst); + +void ggml_cuda_op_reglu(ggml_backend_cuda_context & ctx, ggml_tensor * dst); + +void ggml_cuda_op_geglu(ggml_backend_cuda_context & ctx, ggml_tensor * dst); + +void ggml_cuda_op_swiglu(ggml_backend_cuda_context & ctx, ggml_tensor * dst); + +void ggml_cuda_op_swiglu_oai(ggml_backend_cuda_context & ctx, ggml_tensor * dst); + +void ggml_cuda_op_swiglu_clamp(ggml_backend_cuda_context & ctx, ggml_tensor * dst); + +void ggml_cuda_op_geglu_erf(ggml_backend_cuda_context & ctx, ggml_tensor * dst); + +void ggml_cuda_op_geglu_quick(ggml_backend_cuda_context & ctx, ggml_tensor * dst); + +void ggml_cuda_op_xielu(ggml_backend_cuda_context & ctx, ggml_tensor * dst); + +void ggml_cuda_op_unary_mul(ggml_backend_cuda_context & ctx, ggml_tensor * unary_node, ggml_tensor * mul_node); + +void ggml_cuda_op_relu_sqr(ggml_backend_cuda_context & ctx, ggml_tensor * relu_node, ggml_tensor * sqr_node); + +__device__ __forceinline__ float ggml_cuda_op_silu_single(float x) { + return x / (1.0f + expf(-x)); +} + +__device__ __forceinline__ float ggml_cuda_op_gelu_single(float x) { + const float GELU_COEF_A = 0.044715f; + const float SQRT_2_OVER_PI = 0.79788456080286535587989211986876f; + + return 0.5f * x * (1.0f + tanhf(SQRT_2_OVER_PI * x * (1.0f + GELU_COEF_A * x * x))); +} + +__device__ __forceinline__ float ggml_cuda_op_swiglu_oai_single(float x, float g, float alpha = 1.702f, float limit = 7.0f) { + x = fminf(x, limit); + g = fmaxf(fminf(g, limit), -limit); + + float out_glu = x / (1.0f + expf(-x * alpha)); + out_glu = out_glu * (1.0f + g); + return out_glu; +} + +__device__ __forceinline__ float ggml_cuda_op_swiglu_clamp_single(float gate, float up, float limit) { + gate = fminf(gate, limit); + up = fmaxf(fminf(up, limit), -limit); + + return ggml_cuda_op_silu_single(gate) * up; +} diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/vecdotq.cuh b/python/freetoken/kernel/csrc/ggml/ggml-cuda/vecdotq.cuh new file mode 100644 index 0000000000..f65fe530eb --- /dev/null +++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/vecdotq.cuh @@ -0,0 +1,1380 @@ +#pragma once + +#include "common.cuh" + +#include + +static __device__ __forceinline__ int get_int_b1(const void * x, const int & i32) { + const uint8_t * x8 = (const uint8_t *) x; + + int x32 = x8[4*i32 + 0] << 0; + x32 |= x8[4*i32 + 1] << 8; + x32 |= x8[4*i32 + 2] << 16; + x32 |= x8[4*i32 + 3] << 24; + + return x32; +} + +static __device__ __forceinline__ int get_int_b2(const void * x, const int & i32) { + const uint16_t * x16 = (const uint16_t *) x; // assume at least 2 byte alignment + + int x32 = x16[2*i32 + 0] << 0; + x32 |= x16[2*i32 + 1] << 16; + + return x32; +} + +static __device__ __forceinline__ int get_int_b4(const void * x, const int & i32) { + return ((const int *) x)[i32]; // assume at least 4 byte alignment +} + +// q4 contains 8 indices with 4 bit each. +// This function selects those bytes from table that are at those indices and returns them as int2. +// The first int contains the bytes with even indices in q4, the second int contains the bytes with odd indices in q4. +static __device__ __forceinline__ int2 get_int_from_table_16(const int & q4, const int8_t * table) { +#if defined(GGML_USE_HIP) + // Load the 16-byte table into four 32-bit unsigned integers. + const uint32_t *values = (const uint32_t *)table; + + const uint32_t q_even = q4; + const uint32_t q_odd = (q4 >> 4); + + // Perform lookups in the lower half of the table (indices 0-7). + uint32_t v_even_low = __builtin_amdgcn_perm(values[1], values[0], q_even & 0x07070707); + uint32_t v_odd_low = __builtin_amdgcn_perm(values[1], values[0], q_odd & 0x07070707); + + // Perform lookups in the upper half of the table (indices 8-15). + uint32_t v_even_high = __builtin_amdgcn_perm(values[3], values[2], q_even & 0x07070707); + uint32_t v_odd_high = __builtin_amdgcn_perm(values[3], values[2], q_odd & 0x07070707); + + // Select between the low and high results based on the MSB of each index nibble. + uint32_t mask_even = 0x03020100 | ((q_even & 0x08080808) >> 1); + uint32_t res_x = __builtin_amdgcn_perm(v_even_high, v_even_low, mask_even); + uint32_t mask_odd = 0x03020100 | ((q_odd & 0x08080808) >> 1); + uint32_t res_y = __builtin_amdgcn_perm(v_odd_high, v_odd_low, mask_odd); + + return make_int2(res_x, res_y); +#elif !defined(GGML_USE_MUSA) + // CUDA does not have an instruction for selecting bytes with 4 bit indices. + // However, __byte_perm is an instruction that selects bytes with 3 bit indices that can be used instead. + const uint32_t * table32 = (const uint32_t *) table; + + // __byte_perm selects bytes based on the lower 16 bits in its third argument. + // Therefore, do 2 iterations over the 32 bits in q4 with 0 and 16 shift. + // To handle the fourth bit, first call _byte_perm both for the low and the high 64 bit of table, using the low 3 bits. + // Then, call __byte_perm again to select from the low and high bytes based on the fourth bit. + uint32_t tmp[2]; + const uint32_t low_high_selection_indices = (0x32103210 | ((q4 & 0x88888888) >> 1)); +#pragma unroll + for (uint32_t i = 0; i < 2; ++i) { + const uint32_t shift = 16 * i; + + const uint32_t low = __byte_perm(table32[0], table32[1], q4 >> shift); + const uint32_t high = __byte_perm(table32[2], table32[3], q4 >> shift); + tmp[i] = __byte_perm(low, high, low_high_selection_indices >> shift); + } + + // tmp contains the bytes from tyble in the same order as the 4 bit indices in q4. + // However, for the result we need ints with all even/odd 4 bit indices in q4. + // Therefore, 2 more calls to __byte_perm to put the bytes in the correct order. + return make_int2(__byte_perm(tmp[0], tmp[1], 0x6420), __byte_perm(tmp[0], tmp[1], 0x7531)); +#else + // Generic implementation. + const int q0_32 = (q4 >> 0) & 0x0F0F0F0F; + const int8_t * q0_8 = (const int8_t *) &q0_32; + const char4 val0_8 = make_char4( + table[q0_8[0]], table[q0_8[1]], table[q0_8[2]], table[q0_8[3]]); + + const int q1_32 = (q4 >> 4) & 0x0F0F0F0F; + const int8_t * q1_8 = (const int8_t *) &q1_32; + const char4 val1_8 = make_char4( + table[q1_8[0]], table[q1_8[1]], table[q1_8[2]], table[q1_8[3]]); + + return make_int2(*((const int *) &val0_8), *((const int *) &val1_8)); +#endif +} + +static __device__ __forceinline__ uint32_t unpack_ksigns(const uint8_t v) { + // v is a 7 bit int, with the 8th sign being encodable as popcnt + // with xor we can "correct" the bit instead of having to mask + const uint32_t p = __popc(v) & 1; + const uint32_t s = v ^ p << 7; + // broadcast over uint to allow for 0x08040201 / 0x80402010 as selectors + return s * 0x01010101; +} + +// VDR = vec dot ratio, how many contiguous integers each thread processes when the vec dot kernel is called +// MMVQ = mul_mat_vec_q, MMQ = mul_mat_q + +#define VDR_Q1_0_Q8_1_MMVQ 1 // Process one 32-element chunk at a time for parallelism +#define VDR_Q1_0_Q8_1_MMQ 4 // Q1_0 has 128 bits (4 ints) per block + +#define VDR_Q2_0_Q8_1_MMVQ 1 // Process one 32-element chunk at a time for parallelism +#define VDR_Q2_0_Q8_1_MMQ 2 // Q2_0 group 64: 128 bits (4 ints) per block, 2 32-element chunks + +#define VDR_Q4_0_Q8_1_MMVQ 2 +#define VDR_Q4_0_Q8_1_MMQ 4 + +template static __device__ __forceinline__ float vec_dot_q4_0_q8_1_impl( + const int * v, const int * u, const float & d4, const half2 & ds8) { + + int sumi = 0; + +#pragma unroll + for (int i = 0; i < vdr; ++i) { + const int vi0 = (v[i] >> 0) & 0x0F0F0F0F; + const int vi1 = (v[i] >> 4) & 0x0F0F0F0F; + + // SIMD dot product of quantized values + sumi = ggml_cuda_dp4a(vi0, u[2*i+0], sumi); + sumi = ggml_cuda_dp4a(vi1, u[2*i+1], sumi); + } + + const float2 ds8f = __half22float2(ds8); + + // second part effectively subtracts 8 from each quant value + return d4 * (sumi * ds8f.x - (8*vdr/QI4_0) * ds8f.y); +} + +#define VDR_Q4_1_Q8_1_MMVQ 2 +#define VDR_Q4_1_Q8_1_MMQ 4 + +template static __device__ __forceinline__ float vec_dot_q4_1_q8_1_impl( + const int * v, const int * u, const half2 & dm4, const half2 & ds8) { + + int sumi = 0; + +#pragma unroll + for (int i = 0; i < vdr; ++i) { + const int vi0 = (v[i] >> 0) & 0x0F0F0F0F; + const int vi1 = (v[i] >> 4) & 0x0F0F0F0F; + + // SIMD dot product of quantized values + sumi = ggml_cuda_dp4a(vi0, u[2*i+0], sumi); + sumi = ggml_cuda_dp4a(vi1, u[2*i+1], sumi); + } + +#ifdef FAST_FP16_AVAILABLE + const float2 tmp = __half22float2(__hmul2(dm4, ds8)); + const float d4d8 = tmp.x; + const float m4s8 = tmp.y; +#else + const float2 dm4f = __half22float2(dm4); + const float2 ds8f = __half22float2(ds8); + const float d4d8 = dm4f.x * ds8f.x; + const float m4s8 = dm4f.y * ds8f.y; +#endif // FAST_FP16_AVAILABLE + + // scale second part of sum by QI8_1/(vdr * QR4_1) to compensate for multiple threads adding it + return sumi * d4d8 + m4s8 / (QI8_1 / (vdr * QR4_1)); +} + +#define VDR_Q5_0_Q8_1_MMVQ 2 +#define VDR_Q5_0_Q8_1_MMQ 4 + +template static __device__ __forceinline__ float vec_dot_q5_0_q8_1_impl( + const int * vl, const int * vh, const int * u, const float & d5, const half2 & ds8) { + + int sumi = 0; + +#pragma unroll + for (int i = 0; i < vdr; ++i) { + int vi0 = (vl[i] >> 0) & 0x0F0F0F0F; // lower 4 qs bits, still need qh as 5th bits + vi0 |= (vh[i] << 4) & 0x00000010; // 0 -> 4 + vi0 |= (vh[i] << 11) & 0x00001000; // 1 -> 12 + vi0 |= (vh[i] << 18) & 0x00100000; // 2 -> 20 + vi0 |= (vh[i] << 25) & 0x10000000; // 3 -> 28 + sumi = ggml_cuda_dp4a(vi0, u[2*i+0], sumi); // SIMD dot product of quantized values + + int vi1 = (vl[i] >> 4) & 0x0F0F0F0F; // upper 4 qs bits, still need qh as 5th bits + vi1 |= (vh[i] >> 12) & 0x00000010; // 16 -> 4 + vi1 |= (vh[i] >> 5) & 0x00001000; // 17 -> 12 + vi1 |= (vh[i] << 2) & 0x00100000; // 18 -> 20 + vi1 |= (vh[i] << 9) & 0x10000000; // 19 -> 28 + sumi = ggml_cuda_dp4a(vi1, u[2*i+1], sumi); // SIMD dot product of quantized values + } + + const float2 ds8f = __half22float2(ds8); + + // second part effectively subtracts 16 from each quant value + return d5 * (sumi * ds8f.x - (16*vdr/QI5_0) * ds8f.y); +} + +#define VDR_Q5_1_Q8_1_MMVQ 2 +#define VDR_Q5_1_Q8_1_MMQ 4 + +template static __device__ __forceinline__ float vec_dot_q5_1_q8_1_impl( + const int * vl, const int * vh, const int * u, const half2 & dm5, const half2 & ds8) { + + int sumi = 0; + +#pragma unroll + for (int i = 0; i < vdr; ++i) { + int vi0 = (vl[i] >> 0) & 0x0F0F0F0F; // lower 4 qs bits, still need qh as 5th bits + vi0 |= (vh[i] << 4) & 0x00000010; // 0 -> 4 + vi0 |= (vh[i] << 11) & 0x00001000; // 1 -> 12 + vi0 |= (vh[i] << 18) & 0x00100000; // 2 -> 20 + vi0 |= (vh[i] << 25) & 0x10000000; // 3 -> 28 + sumi = ggml_cuda_dp4a(vi0, u[2*i+0], sumi); // SIMD dot product of quantized values + + int vi1 = (vl[i] >> 4) & 0x0F0F0F0F; // upper 4 qs bits, still need qh as 5th bits + vi1 |= (vh[i] >> 12) & 0x00000010; // 16 -> 4 + vi1 |= (vh[i] >> 5) & 0x00001000; // 17 -> 12 + vi1 |= (vh[i] << 2) & 0x00100000; // 18 -> 20 + vi1 |= (vh[i] << 9) & 0x10000000; // 19 -> 28 + sumi = ggml_cuda_dp4a(vi1, u[2*i+1], sumi); // SIMD dot product of quantized values + } + +#ifdef FAST_FP16_AVAILABLE + const float2 tmp = __half22float2(__hmul2(dm5, ds8)); + const float d5d8 = tmp.x; + const float m5s8 = tmp.y; +#else + const float2 dm5f = __half22float2(dm5); + const float2 ds8f = __half22float2(ds8); + const float d5d8 = dm5f.x * ds8f.x; + const float m5s8 = dm5f.y * ds8f.y; +#endif // FAST_FP16_AVAILABLE + + // scale second part of sum by QI5_1 / vdr to compensate for multiple threads adding it + return sumi*d5d8 + m5s8 / (QI5_1 / vdr); +} + +#define VDR_Q8_0_Q8_1_MMVQ 2 +#define VDR_Q8_0_Q8_1_MMQ 8 + +template static __device__ __forceinline__ T vec_dot_q8_0_q8_1_impl( + const int * v, const int * u, const T & d8_0, const T & d8_1) { + + int sumi = 0; + +#pragma unroll + for (int i = 0; i < vdr; ++i) { + // SIMD dot product of quantized values + sumi = ggml_cuda_dp4a(v[i], u[i], sumi); + } + + return d8_0*d8_1 * ((T) sumi); +} + +template static __device__ __forceinline__ float vec_dot_q8_1_q8_1_impl( + const int * v, const int * u, const half2 & dm8, const half2 & ds8) { + + int sumi = 0; + +#pragma unroll + for (int i = 0; i < vdr; ++i) { + // SIMD dot product of quantized values + sumi = ggml_cuda_dp4a(v[i], u[i], sumi); + } + +#ifdef FAST_FP16_AVAILABLE + const float2 tmp = __half22float2(__hmul2(dm8, ds8)); + const float d8d8 = tmp.x; + const float m8s8 = tmp.y; +#else + const float2 dm8f = __half22float2(dm8); + const float2 ds8f = __half22float2(ds8); + const float d8d8 = dm8f.x * ds8f.x; + const float m8s8 = dm8f.y * ds8f.y; +#endif // FAST_FP16_AVAILABLE + + // scale second part of sum by QI8_1/ vdr to compensate for multiple threads adding it + return sumi*d8d8 + m8s8 / (QI8_1 / vdr); +} + +template static __device__ __forceinline__ float vec_dot_q8_0_16_q8_1_impl( + const int * v, const int * u, const float * d8_0, const float & d8_1) { + + float sumf = 0.0f; + +#pragma unroll + for (int i0 = 0; i0 < vdr; i0 += QI8_0/2) { + int sumi = 0; + +#pragma unroll + for (int i = i0; i < i0 + QI8_0/2; ++i) { + // SIMD dot product of quantized values + sumi = ggml_cuda_dp4a(v[i], u[i], sumi); + } + + sumf += d8_0[i0/(QI8_0/2)]*sumi; + } + + return d8_1*sumf; +} + +#define VDR_MXFP4_Q8_1_MMVQ 2 +#define VDR_MXFP4_Q8_1_MMQ 4 + +static __device__ __forceinline__ float vec_dot_mxfp4_q8_1( + const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs) { + + const block_mxfp4 * bq4 = (const block_mxfp4 *) vbq + kbx; + + const int * q8 = (const int *) bq8_1->qs + iqs; + + int sumi = 0; +#pragma unroll + for (int l = 0; l < VDR_MXFP4_Q8_1_MMVQ; ++l) { + const int aux_q4 = get_int_b1(bq4->qs, iqs + l); + const int2 v = get_int_from_table_16(aux_q4, kvalues_mxfp4); + + sumi = ggml_cuda_dp4a(v.x, q8[l + 0], sumi); + sumi = ggml_cuda_dp4a(v.y, q8[l + 4], sumi); + } + + const float d = ggml_cuda_e8m0_to_fp32(bq4->e) * 0.5f * __low2float(bq8_1->ds); + return d * sumi; +} + +#define VDR_NVFP4_Q8_1_MMVQ 4 +#define VDR_NVFP4_Q8_1_MMQ 8 + +static __device__ __forceinline__ float vec_dot_nvfp4_q8_1( + const void * __restrict__ vbq, + const block_q8_1 * __restrict__ bq8_1, + const int32_t & kbx, + const int32_t & iqs) { + + const block_nvfp4 * bq4 = (const block_nvfp4 *) vbq + kbx; + float sum = 0.0f; +#pragma unroll + for (int i = 0; i < VDR_NVFP4_Q8_1_MMVQ/2; i++) { + const int32_t iqs0 = iqs + 2*i; + const int32_t iqs1 = iqs0 + 1; + const int32_t is = iqs0 >> 1; + const int2 v0 = get_int_from_table_16(get_int_b4(bq4->qs, iqs0), kvalues_mxfp4); + const int2 v1 = get_int_from_table_16(get_int_b4(bq4->qs, iqs1), kvalues_mxfp4); + const block_q8_1 * bq8 = bq8_1 + (is >> 1); + const int32_t i8 = ((is & 1) << 2); + + int sumi = ggml_cuda_dp4a(v0.x, get_int_b4(bq8->qs, i8 + 0), 0); + sumi = ggml_cuda_dp4a(v0.y, get_int_b4(bq8->qs, i8 + 2), sumi); + sumi = ggml_cuda_dp4a(v1.x, get_int_b4(bq8->qs, i8 + 1), sumi); + sumi = ggml_cuda_dp4a(v1.y, get_int_b4(bq8->qs, i8 + 3), sumi); + + const float d = ggml_cuda_ue4m3_to_fp32(bq4->d[is]) * __low2float(bq8->ds); + sum += d * float(sumi); + } + + return sum; +} +#define VDR_Q2_K_Q8_1_MMVQ 1 +#define VDR_Q2_K_Q8_1_MMQ 4 + +// contiguous v/x values +static __device__ __forceinline__ float vec_dot_q2_K_q8_1_impl_mmvq( + const int & v, const int * __restrict__ u, const uint8_t * __restrict__ scales, + const half2 & dm2, const float * __restrict__ d8) { + + float sumf_d = 0.0f; + float sumf_m = 0.0f; + +#pragma unroll + for (int i = 0; i < QR2_K; ++i) { + const int sc = scales[2*i]; + + const int vi = (v >> (2*i)) & 0x03030303; + + sumf_d += d8[i] * (ggml_cuda_dp4a(vi, u[i], 0) * (sc & 0xF)); // SIMD dot product + + // fill int with 4x m + int m = sc >> 4; + m |= m << 8; + m |= m << 16; + sumf_m += d8[i] * ggml_cuda_dp4a(m, u[i], 0); // multiply constant q2_K part with sum of q8_1 values + } + + const float2 dm2f = __half22float2(dm2); + + return dm2f.x*sumf_d - dm2f.y*sumf_m; +} + +// contiguous v/x + u/y values +template +static __device__ __forceinline__ float vec_dot_q2_K_q8_1_impl_mmq( + const int * __restrict__ v, const int * __restrict__ u, const half2 * dm2, const float & d8, const half2 * s8) { + + float sumf = 0.0f; + float sumf_d8 = 0.0f; + +#pragma unroll + for (int i0 = 0; i0 < QR2_K*VDR_Q2_K_Q8_1_MMQ; i0 += QI8_1) { + const float2 dm2f0 = __half22float2(dm2[i0/(QI8_1/2) + 0]); + int sumi_d0 = 0; + + const float2 dm2f1 = __half22float2(dm2[i0/(QI8_1/2) + 1]); + int sumi_d1 = 0; + +#pragma unroll + for (int i = i0; i < i0 + QI8_1/2; ++i) { + sumi_d0 = ggml_cuda_dp4a(v[i], u[i], sumi_d0); + } + sumf_d8 += dm2f0.x * sumi_d0; + +#pragma unroll + for (int i = i0 + QI8_1/2; i < i0 + QI8_1; ++i) { + sumi_d1 = ggml_cuda_dp4a(v[i], u[i], sumi_d1); + } + sumf_d8 += dm2f1.x * sumi_d1; + + if (i0/QI8_1 < ns8) { + const float2 s8f = __half22float2(s8[i0/QI8_1]); + sumf -= dm2f0.y*s8f.x; + sumf -= dm2f1.y*s8f.y; + } else { + int sumi_m0 = 0; +#pragma unroll + for (int i = i0; i < i0 + QI8_1/2; ++i) { + sumi_m0 = ggml_cuda_dp4a(0x01010101, u[i], sumi_m0); + } + sumf_d8 -= dm2f0.y * sumi_m0; + + int sumi_m1 = 0; +#pragma unroll + for (int i = i0 + QI8_1/2; i < i0 + QI8_1; ++i) { + sumi_m1 = ggml_cuda_dp4a(0x01010101, u[i], sumi_m1); + } + sumf_d8 -= dm2f1.y * sumi_m1; + } + } + + return sumf + d8*sumf_d8; +} + +#define VDR_Q3_K_Q8_1_MMVQ 1 +#define VDR_Q3_K_Q8_1_MMQ 2 + +// contiguous v/x values +static __device__ __forceinline__ float vec_dot_q3_K_q8_1_impl_mmvq( + const int & vl, const int & vh, const int * __restrict__ u, const uint8_t * __restrict__ scales, + const int & scale_offset, const float & d3, const float * __restrict__ d8) { + + float sumf = 0.0f; + +#pragma unroll + for (int i = 0; i < QR3_K; ++i) { + const int isc = scale_offset + 2*i; + + const int isc_low = isc % (QK_K/32); + const int sc_shift_low = 4 * (isc / (QK_K/32)); + const int sc_low = (scales[isc_low] >> sc_shift_low) & 0xF; + + const int isc_high = isc % (QK_K/64); + const int sc_shift_high = 2 * (isc / (QK_K/64)); + const int sc_high = ((scales[(QK_K/32) + isc_high] >> sc_shift_high) & 3) << 4; + + const int sc = (sc_low | sc_high) - 32; + + const int vil = (vl >> (2*i)) & 0x03030303; + + const int vih = ((vh >> i) << 2) & 0x04040404; + + const int vi = __vsub4(vil, vih); + + sumf += d8[i] * (ggml_cuda_dp4a(vi, u[i], 0) * sc); // SIMD dot product + } + + return d3 * sumf; +} + +// contiguous v/x + u/y values +static __device__ __forceinline__ float vec_dot_q3_K_q8_1_impl_mmq( + const int * __restrict__ v, const int * __restrict__ u, const int8_t * __restrict__ scales, + const float & d3, const float & d8) { + + int sumi = 0; + +#pragma unroll + for (int i0 = 0; i0 < QR3_K*VDR_Q3_K_Q8_1_MMQ; i0 += QI8_1/2) { + int sumi_sc = 0; + +#pragma unroll + for (int i = i0; i < i0 + QI8_1/2; ++i) { + sumi_sc = ggml_cuda_dp4a(v[i], u[i], sumi_sc); // SIMD dot product + } + + sumi += sumi_sc * scales[i0 / (QI8_1/2)]; + } + + return d3*d8 * sumi; +} + +#define VDR_Q4_K_Q8_1_MMVQ 2 +#define VDR_Q4_K_Q8_1_MMQ 8 + +// contiguous v/x values +static __device__ __forceinline__ float vec_dot_q4_K_q8_1_impl_vmmq( + const int * __restrict__ v, const int * __restrict__ u, const uint8_t * __restrict__ sc, + const uint8_t * __restrict__ m, const half2 & dm4, const float * __restrict__ d8) { + + float sumf_d = 0.0f; + float sumf_m = 0.0f; + +#pragma unroll + for (int i = 0; i < QR4_K; ++i) { + const int v0i = (v[0] >> (4*i)) & 0x0F0F0F0F; + const int v1i = (v[1] >> (4*i)) & 0x0F0F0F0F; + + const int dot1 = ggml_cuda_dp4a(v1i, u[2*i+1], ggml_cuda_dp4a(v0i, u[2*i+0], 0)); // SIMD dot product + const int dot2 = ggml_cuda_dp4a(0x01010101, u[2*i+1], ggml_cuda_dp4a(0x01010101, u[2*i+0], 0)); // sum of u + + sumf_d += d8[i] * (dot1 * sc[i]); + sumf_m += d8[i] * (dot2 * m[i]); // multiply constant part of q4_K with sum of q8_1 values + } + + const float2 dm4f = __half22float2(dm4); + + return dm4f.x*sumf_d - dm4f.y*sumf_m; +} + +// contiguous v/x + u/y values +static __device__ __forceinline__ float vec_dot_q4_K_q8_1_impl_mmq( + const int * __restrict__ v, const int * __restrict__ u, const uint8_t * __restrict__ sc, + const uint8_t * __restrict__ m, const half2 & dm4, const half2 * __restrict__ ds8) { + + float sumf_d = 0.0f; + float sumf_m = 0.0f; + +#pragma unroll + for (int i = 0; i < QR4_K*VDR_Q4_K_Q8_1_MMQ/QI8_1; ++i) { + int sumi_d = 0; + +#pragma unroll + for (int j = 0; j < QI8_1; ++j) { + sumi_d = ggml_cuda_dp4a((v[j] >> (4*i)) & 0x0F0F0F0F, u[i*QI8_1 + j], sumi_d); // SIMD dot product + } + + const float2 ds8f = __half22float2(ds8[i]); + + sumf_d += ds8f.x * (sc[i] * sumi_d); + sumf_m += ds8f.y * m[i]; // sum of q8_1 block * q4_K min val + } + + const float2 dm4f = __half22float2(dm4); + + return dm4f.x*sumf_d - dm4f.y*sumf_m; +} + +#define VDR_Q5_K_Q8_1_MMVQ 2 +#define VDR_Q5_K_Q8_1_MMQ 8 + +// contiguous v/x values +static __device__ __forceinline__ float vec_dot_q5_K_q8_1_impl_vmmq( + const int * __restrict__ vl, const int * __restrict__ vh, const int * __restrict__ u, const uint8_t * __restrict__ sc, + const uint8_t * __restrict__ m, const half2 & dm5, const float * __restrict__ d8) { + + float sumf_d = 0.0f; + float sumf_m = 0.0f; + +#pragma unroll + for (int i = 0; i < QR5_K; ++i) { + const int vl0i = (vl[0] >> (4*i)) & 0x0F0F0F0F; + const int vl1i = (vl[1] >> (4*i)) & 0x0F0F0F0F; + + const int vh0i = ((vh[0] >> i) << 4) & 0x10101010; + const int vh1i = ((vh[1] >> i) << 4) & 0x10101010; + + const int v0i = vl0i | vh0i; + const int v1i = vl1i | vh1i; + + const int dot1 = ggml_cuda_dp4a(v0i, u[2*i+0], ggml_cuda_dp4a(v1i, u[2*i+1], 0)); // SIMD dot product + const int dot2 = ggml_cuda_dp4a(0x01010101, u[2*i+0], ggml_cuda_dp4a(0x01010101, u[2*i+1], 0)); // sum of u + + sumf_d += d8[i] * (dot1 * sc[i]); + sumf_m += d8[i] * (dot2 * m[i]); + + } + + const float2 dm5f = __half22float2(dm5); + + return dm5f.x*sumf_d - dm5f.y*sumf_m; +} + +// contiguous v/x + u/y values +static __device__ __forceinline__ float vec_dot_q5_K_q8_1_impl_mmq( + const int * __restrict__ v, const int * __restrict__ u, const uint8_t * __restrict__ sc, + const uint8_t * __restrict__ m, const half2 & dm4, const half2 * __restrict__ ds8) { + + float sumf_d = 0.0f; + float sumf_m = 0.0f; + +#pragma unroll + for (int i = 0; i < QR5_K*VDR_Q5_K_Q8_1_MMQ/QI8_1; ++i) { + int sumi_d = 0; + +#pragma unroll + for (int j = 0; j < QI8_1; ++j) { + sumi_d = ggml_cuda_dp4a(v[i*QI8_1 + j], u[i*QI8_1 + j], sumi_d); // SIMD dot product + } + + const float2 ds8f = __half22float2(ds8[i]); + + sumf_d += ds8f.x * (sc[i] * sumi_d); + sumf_m += ds8f.y * m[i]; // sum of q8_1 block * q4_K min val + } + + const float2 dm4f = __half22float2(dm4); + + return dm4f.x*sumf_d - dm4f.y*sumf_m; +} + +#define VDR_Q6_K_Q8_1_MMVQ 1 +#define VDR_Q6_K_Q8_1_MMQ 8 + +// contiguous v/x values +static __device__ __forceinline__ float vec_dot_q6_K_q8_1_impl_mmvq( + const int & vl, const int & vh, const int * __restrict__ u, const int8_t * __restrict__ scales, + const float & d, const float * __restrict__ d8) { + + float sumf = 0.0f; + +#pragma unroll + for (int i = 0; i < QR6_K; ++i) { + const int sc = scales[4*i]; + + const int vil = (vl >> (4*i)) & 0x0F0F0F0F; + + const int vih = ((vh >> (4*i)) << 4) & 0x30303030; + + const int vi = __vsub4((vil | vih), 0x20202020); // vi = (vil | vih) - 32 + + sumf += d8[i] * (ggml_cuda_dp4a(vi, u[i], 0) * sc); // SIMD dot product + } + + return d*sumf; +} + +// contiguous v/x + u/y values +static __device__ __forceinline__ float vec_dot_q6_K_q8_1_impl_mmq( + const int * __restrict__ v, const int * __restrict__ u, const int8_t * __restrict__ sc, + const float & d6, const float * __restrict__ d8) { + + float sumf_d = 0.0f; + + const int sc_packed = get_int_b4(sc, 0); + const int8_t * sc_reg = (const int8_t *) &sc_packed; + +#pragma unroll + for (int i0 = 0; i0 < VDR_Q6_K_Q8_1_MMQ; i0 += 4) { + int2 sumi_d = {0, 0}; // 2 q6_K scales per q8_1 scale + +#pragma unroll + for (int i = i0; i < i0 + 2; ++i) { + sumi_d.x = ggml_cuda_dp4a(v[2*i+0], u[2*i+0], sumi_d.x); // SIMD dot product + sumi_d.x = ggml_cuda_dp4a(v[2*i+1], u[2*i+1], sumi_d.x); // SIMD dot product + + sumi_d.y = ggml_cuda_dp4a(v[2*i+4], u[2*i+4], sumi_d.y); // SIMD dot product + sumi_d.y = ggml_cuda_dp4a(v[2*i+5], u[2*i+5], sumi_d.y); // SIMD dot product + } + + sumf_d += d8[i0/4] * (sc_reg[i0/2+0]*sumi_d.x + sc_reg[i0/2+1]*sumi_d.y); + } + + return d6 * sumf_d; +} + +static __device__ __forceinline__ float vec_dot_q1_0_q8_1( + const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs) { + + const block_q1_0 * bq1_0 = (const block_q1_0 *) vbq + kbx; + + // Q1_0: 128 elements with ONE scale + // Q8_1: 32 elements per block with individual scales + // iqs selects which of the 4 chunks of 32 elements to process (0-3) + + const float d1 = bq1_0->d; + const int16_t * qs = (const int16_t *) bq1_0->qs + iqs * 2; + + // Process only the chunk specified by iqs + const block_q8_1 * bq8_1_chunk = bq8_1 + iqs; + + int sumi = 0; +#pragma unroll + for (int j = 0; j < 2; ++j) { + const int q = qs[j]; + + const int u0 = get_int_b4(bq8_1_chunk->qs, j*4+0); + const int u1 = get_int_b4(bq8_1_chunk->qs, j*4+1); + const int u2 = get_int_b4(bq8_1_chunk->qs, j*4+2); + const int u3 = get_int_b4(bq8_1_chunk->qs, j*4+3); + + // unpack crumbs into nibble indices + const int n0 = __byte_perm(0x11100100, 0x11100100, q >> 0); // [0, 1, 4, 5] [ 8, 9, 12, 13] + const int n1 = __byte_perm(0x11100100, 0x11100100, q >> 2); // [2, 3, 6, 7] [10, 11, 14, 15] + // unpack nibbles into byte values + const int s0 = __byte_perm(0x01FF, 0x01FF, n0 >> 0); + const int s1 = __byte_perm(0x01FF, 0x01FF, n1 >> 0); + const int s2 = __byte_perm(0x01FF, 0x01FF, n0 >> 16); + const int s3 = __byte_perm(0x01FF, 0x01FF, n1 >> 16); + // unshuffle values + const int v0 = __byte_perm(s0, s1, 0x5410); + const int v1 = __byte_perm(s0, s1, 0x7632); + const int v2 = __byte_perm(s2, s3, 0x5410); + const int v3 = __byte_perm(s2, s3, 0x7632); + + sumi = ggml_cuda_dp4a(v0, u0, sumi); + sumi = ggml_cuda_dp4a(v1, u1, sumi); + sumi = ggml_cuda_dp4a(v2, u2, sumi); + sumi = ggml_cuda_dp4a(v3, u3, sumi); + } + + // Apply Q1_0's single scale and this chunk's Q8_1 scale + const float d8 = __low2float(bq8_1_chunk->ds); + return d1 * d8 * sumi; +} + +static __device__ __forceinline__ float vec_dot_q2_0_q8_1( + const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs) { + + const block_q2_0 * bq2_0 = (const block_q2_0 *) vbq + kbx; + + // Q2_0 (group 64): 64 elements with ONE scale, 2 bits per element (4 elements per byte) + // Q8_1: 32 elements per block with individual scales + // iqs selects which of the 2 chunks of 32 elements to process (0-1) + + const float d2 = bq2_0->d; + const int16_t * qs = (const int16_t *) bq2_0->qs + iqs * 4; + + // Process only the chunk specified by iqs + const block_q8_1 * bq8_1_chunk = bq8_1 + iqs; + + int sumi = 0; +#pragma unroll + for (int j = 0; j < 4; ++j) { + const int q = qs[j]; + const int u = get_int_b4(bq8_1_chunk->qs, j*2+0); + const int v = get_int_b4(bq8_1_chunk->qs, j*2+1); + +#if defined(GGML_USE_HIP) + const uint32_t qx_indices = (q & 0x03) | ((q & 0x0C) << 6) | ((q & 0x30) << 12) | ((q & 0xC0) << 18); + const uint32_t qy_bits = q >> 8; + const uint32_t qy_indices = (qy_bits & 0x03) | ((qy_bits & 0x0C) << 6) | ((qy_bits & 0x30) << 12) | ((qy_bits & 0xC0) << 18); + const int qx = __builtin_amdgcn_perm(0x020100FF, 0x020100FF, qx_indices); + const int qy = __builtin_amdgcn_perm(0x020100FF, 0x020100FF, qy_indices); +#else + // unpack even and odd crumbs into byte values + const int qe = __byte_perm(0x020100FF, 0x020100FF, q >> 0); + const int qo = __byte_perm(0x020100FF, 0x020100FF, q >> 2); + // unshuffle values + const int qx = __byte_perm(qe, qo, 0x5140); + const int qy = __byte_perm(qe, qo, 0x7362); +#endif // defined(GGML_USE_HIP) + + sumi = ggml_cuda_dp4a(u, qx, sumi); + sumi = ggml_cuda_dp4a(v, qy, sumi); + } + + // Apply Q2_0's single scale and this chunk's Q8_1 scale + const float d8 = __low2float(bq8_1_chunk->ds); + return d2 * d8 * sumi; +} + +static __device__ __forceinline__ float vec_dot_q4_0_q8_1( + const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs) { + + const block_q4_0 * bq4_0 = (const block_q4_0 *) vbq + kbx; + + int v[VDR_Q4_0_Q8_1_MMVQ]; + int u[2*VDR_Q4_0_Q8_1_MMVQ]; + +#pragma unroll + for (int i = 0; i < VDR_Q4_0_Q8_1_MMVQ; ++i) { + v[i] = get_int_b2(bq4_0->qs, iqs + i); + u[2*i+0] = get_int_b4(bq8_1->qs, iqs + i); + u[2*i+1] = get_int_b4(bq8_1->qs, iqs + i + QI4_0); + } + + return vec_dot_q4_0_q8_1_impl(v, u, bq4_0->d, bq8_1->ds); +} + + +static __device__ __forceinline__ float vec_dot_q4_1_q8_1( + const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs) { + + const block_q4_1 * bq4_1 = (const block_q4_1 *) vbq + kbx; + + int v[VDR_Q4_1_Q8_1_MMVQ]; + int u[2*VDR_Q4_1_Q8_1_MMVQ]; + +#pragma unroll + for (int i = 0; i < VDR_Q4_1_Q8_1_MMVQ; ++i) { + v[i] = get_int_b4(bq4_1->qs, iqs + i); + u[2*i+0] = get_int_b4(bq8_1->qs, iqs + i); + u[2*i+1] = get_int_b4(bq8_1->qs, iqs + i + QI4_1); + } + + return vec_dot_q4_1_q8_1_impl(v, u, bq4_1->dm, bq8_1->ds); +} + +static __device__ __forceinline__ float vec_dot_q5_0_q8_1( + const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs) { + + const block_q5_0 * bq5_0 = (const block_q5_0 *) vbq + kbx; + + int vl[VDR_Q5_0_Q8_1_MMVQ]; + int vh[VDR_Q5_0_Q8_1_MMVQ]; + int u[2*VDR_Q5_0_Q8_1_MMVQ]; + +#pragma unroll + for (int i = 0; i < VDR_Q5_0_Q8_1_MMVQ; ++i) { + vl[i] = get_int_b2(bq5_0->qs, iqs + i); + vh[i] = get_int_b2(bq5_0->qh, 0) >> (4 * (iqs + i)); + u[2*i+0] = get_int_b4(bq8_1->qs, iqs + i); + u[2*i+1] = get_int_b4(bq8_1->qs, iqs + i + QI5_0); + } + + return vec_dot_q5_0_q8_1_impl(vl, vh, u, bq5_0->d, bq8_1->ds); +} + +static __device__ __forceinline__ float vec_dot_q5_1_q8_1( + const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs) { + + const block_q5_1 * bq5_1 = (const block_q5_1 *) vbq + kbx; + + int vl[VDR_Q5_1_Q8_1_MMVQ]; + int vh[VDR_Q5_1_Q8_1_MMVQ]; + int u[2*VDR_Q5_1_Q8_1_MMVQ]; + +#pragma unroll + for (int i = 0; i < VDR_Q5_1_Q8_1_MMVQ; ++i) { + vl[i] = get_int_b4(bq5_1->qs, iqs + i); + vh[i] = get_int_b4(bq5_1->qh, 0) >> (4 * (iqs + i)); + u[2*i+0] = get_int_b4(bq8_1->qs, iqs + i); + u[2*i+1] = get_int_b4(bq8_1->qs, iqs + i + QI5_1); + } + + return vec_dot_q5_1_q8_1_impl(vl, vh, u, bq5_1->dm, bq8_1->ds); +} + +static __device__ __forceinline__ float vec_dot_q8_0_q8_1( + const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs) { + + const block_q8_0 * bq8_0 = (const block_q8_0 *) vbq + kbx; + + int v[VDR_Q8_0_Q8_1_MMVQ]; + int u[VDR_Q8_0_Q8_1_MMVQ]; + +#pragma unroll + for (int i = 0; i < VDR_Q8_0_Q8_1_MMVQ; ++i) { + v[i] = get_int_b2(bq8_0->qs, iqs + i); + u[i] = get_int_b4(bq8_1->qs, iqs + i); + } + + return vec_dot_q8_0_q8_1_impl(v, u, bq8_0->d, __low2half(bq8_1->ds)); +} + +static __device__ __forceinline__ float vec_dot_q2_K_q8_1( + const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs) { + + const block_q2_K * bq2_K = (const block_q2_K *) vbq + kbx; + + const int bq8_offset = QR2_K * (iqs / QI8_1); + const int scale_offset = iqs - iqs % QI8_1 + (iqs % QI8_1) / (QI8_1/2); + + const uint8_t * scales = bq2_K->scales + scale_offset; + + const int v = get_int_b4(bq2_K->qs, iqs); + int u[QR2_K]; + float d8[QR2_K]; + +#pragma unroll + for (int i = 0; i < QR2_K; ++ i) { + u[i] = get_int_b4(bq8_1[bq8_offset + i].qs, iqs % QI8_1); + d8[i] = __low2float(bq8_1[bq8_offset + i].ds); + } + + return vec_dot_q2_K_q8_1_impl_mmvq(v, u, scales, bq2_K->dm, d8); +} + +static __device__ __forceinline__ float vec_dot_q3_K_q8_1( + const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs) { + + const block_q3_K * bq3_K = (const block_q3_K *) vbq + kbx; + + const int bq8_offset = QR3_K * (iqs / (QI3_K/2)); + const int scale_offset = iqs - iqs % QI8_1 + (iqs % QI8_1) / (QI8_1/2); + + const float d = bq3_K->d; + + const int vl = get_int_b2(bq3_K->qs, iqs); + + // invert the mask with ~ so that a 0/1 results in 4/0 being subtracted + const int vh = ~get_int_b2(bq3_K->hmask, iqs % (QI3_K/2)) >> bq8_offset; + + int u[QR3_K]; + float d8[QR3_K]; + +#pragma unroll + for (int i = 0; i < QR3_K; ++i) { + u[i] = get_int_b4(bq8_1[bq8_offset + i].qs, iqs % QI8_1); + d8[i] = __low2float(bq8_1[bq8_offset + i].ds); + } + + return vec_dot_q3_K_q8_1_impl_mmvq(vl, vh, u, bq3_K->scales, scale_offset, d, d8); +} + +static __device__ __forceinline__ float vec_dot_q4_K_q8_1( + const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs) { + + const block_q4_K * bq4_K = (const block_q4_K *) vbq + kbx; + + int v[2]; + int u[2*QR4_K]; + float d8[QR4_K]; + + // iqs is in 0,2..30. bq8_offset = iqs/4 -> bq8_offset = 0, 2, 4, 6 + const int bq8_offset = QR4_K * ((iqs/2) / (QI8_1/2)); + + // iqs = 0....3 -> bq8_offset = 0, want q4_offset = 0, 4, 8, 12 + // iqs = 4....7 -> bq8_offset = 2, want q4_offset = 32, 36, 40, 44 + // iqs = 8...11 -> bq8_offset = 4, want q4_offset = 64, 68, 72, 76 + // iqs = 12..15 -> bq8_offset = 6, want q4_offset = 96, 100, 104, 108 + + const int * q4 = (const int *)(bq4_K->qs + 16 * bq8_offset + 4 * ((iqs/2)%4)); + v[0] = q4[0]; + v[1] = q4[4]; + + // branchless so nvcc can hoist this out of the ncols_dst loop + const uint16_t * scales = (const uint16_t *)bq4_K->scales; + const int j = bq8_offset/2; + const int jm = j & 1; + + const uint32_t s0 = scales[jm + 0]; + const uint32_t s2 = scales[jm + 2]; + const uint32_t s4 = scales[jm + 4]; + + const uint32_t hi = (uint32_t) -(int32_t) (j >= 2); + + uint16_t aux[2]; + aux[0] = (uint16_t) (((s0 & 0x3f3f) & ~hi) | ((((s4 >> 0) & 0x0f0f) | ((s0 & 0xc0c0) >> 2)) & hi)); + aux[1] = (uint16_t) (((s2 & 0x3f3f) & ~hi) | ((((s4 >> 4) & 0x0f0f) | ((s2 & 0xc0c0) >> 2)) & hi)); + const uint8_t * sc = (const uint8_t *)aux; + const uint8_t * m = sc + 2; + + for (int i = 0; i < QR4_K; ++i) { + const block_q8_1 * bq8i = bq8_1 + bq8_offset + i; + d8[i] = __low2float(bq8i->ds); + + const int * q8 = (const int *)bq8i->qs + ((iqs/2)%4); + u[2*i+0] = q8[0]; + u[2*i+1] = q8[4]; + } + + return vec_dot_q4_K_q8_1_impl_vmmq(v, u, sc, m, bq4_K->dm, d8); +} + +static __device__ __forceinline__ float vec_dot_q5_K_q8_1( + const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs) { + + const block_q5_K * bq5_K = (const block_q5_K *) vbq + kbx; + + int vl[2]; + int vh[2]; + int u[2*QR5_K]; + float d8[QR5_K]; + + const int bq8_offset = QR5_K * ((iqs/2) / (QI8_1/2)); + const int * ql = (const int *)(bq5_K->qs + 16 * bq8_offset + 4 * ((iqs/2)%4)); + const int * qh = (const int *)(bq5_K->qh + 4 * ((iqs/2)%4)); + + vl[0] = ql[0]; + vl[1] = ql[4]; + + vh[0] = qh[0] >> bq8_offset; + vh[1] = qh[4] >> bq8_offset; + + // same as q4_K + const uint16_t * scales = (const uint16_t *)bq5_K->scales; + const int j = bq8_offset/2; + const int jm = j & 1; + + const uint32_t s0 = scales[jm + 0]; + const uint32_t s2 = scales[jm + 2]; + const uint32_t s4 = scales[jm + 4]; + + const uint32_t hi = (uint32_t) -(int32_t) (j >= 2); + + uint16_t aux[2]; + aux[0] = (uint16_t) (((s0 & 0x3f3f) & ~hi) | ((((s4 >> 0) & 0x0f0f) | ((s0 & 0xc0c0) >> 2)) & hi)); + aux[1] = (uint16_t) (((s2 & 0x3f3f) & ~hi) | ((((s4 >> 4) & 0x0f0f) | ((s2 & 0xc0c0) >> 2)) & hi)); + + const uint8_t * sc = (const uint8_t *)aux; + const uint8_t * m = sc + 2; + +#pragma unroll + for (int i = 0; i < QR5_K; ++i) { + const block_q8_1 * bq8i = bq8_1 + bq8_offset + i; + d8[i] = __low2float(bq8i->ds); + + const int * q8 = (const int *)bq8i->qs + ((iqs/2)%4); + u[2*i+0] = q8[0]; + u[2*i+1] = q8[4]; + } + + return vec_dot_q5_K_q8_1_impl_vmmq(vl, vh, u, sc, m, bq5_K->dm, d8); +} + +static __device__ __forceinline__ float vec_dot_q6_K_q8_1( + const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs) { + + const block_q6_K * bq6_K = (const block_q6_K *) vbq + kbx; + + const int bq8_offset = 2 * QR6_K * (iqs / (QI6_K/2)) + (iqs % (QI6_K/2)) / (QI6_K/4); + const int scale_offset = (QI6_K/4) * (iqs / (QI6_K/2)) + (iqs % (QI6_K/2)) / (QI6_K/8); + const int vh_shift = 2 * ((iqs % (QI6_K/2)) / (QI6_K/4)); + + const int vl = get_int_b2(bq6_K->ql, iqs); + const int vh = get_int_b2(bq6_K->qh, (QI6_K/4) * (iqs / (QI6_K/2)) + iqs % (QI6_K/4)) >> vh_shift; + + const int8_t * scales = bq6_K->scales + scale_offset; + + int u[QR6_K]; + float d8[QR6_K]; + +#pragma unroll + for (int i = 0; i < QR6_K; ++i) { + u[i] = get_int_b4(bq8_1[bq8_offset + 2*i].qs, iqs % QI8_1); + d8[i] = __low2float(bq8_1[bq8_offset + 2*i].ds); + } + + return vec_dot_q6_K_q8_1_impl_mmvq(vl, vh, u, scales, bq6_K->d, d8); +} + +#define VDR_IQ2_XXS_Q8_1_MMVQ 2 +#define VDR_IQ2_XXS_Q8_1_MMQ 2 + +static __device__ __forceinline__ float vec_dot_iq2_xxs_q8_1( + const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs) { + + const block_iq2_xxs * bq2 = (const block_iq2_xxs *) vbq + kbx; + + const int q2 = get_int_b2(bq2->qs, iqs); + const uint8_t * aux8 = (const uint8_t *) &q2; + const uint32_t aux32 = get_int_b2(bq2->qs, iqs + 1); + + int sumi = 0; +#pragma unroll + for (int k0 = 0; k0 < 8; k0 += 2) { + const uint2 grid_pos = ((const uint2*)iq2xxs_grid)[aux8[k0/2]]; + const uint32_t signs = unpack_ksigns(aux32 >> (7 * k0 / 2)); + + const int signs0 = __vcmpne4(signs & 0x08040201, 0); + const int grid0 = __vsub4(grid_pos.x ^ signs0, signs0); + const int u0 = get_int_b4(bq8_1[iqs/2].qs, k0 + 0); + sumi = ggml_cuda_dp4a(grid0, u0, sumi); + + const int signs1 = __vcmpne4(signs & 0x80402010, 0); + const int grid1 = __vsub4(grid_pos.y ^ signs1, signs1); + const int u1 = get_int_b4(bq8_1[iqs/2].qs, k0 + 1); + sumi = ggml_cuda_dp4a(grid1, u1, sumi); + } + + const int ls = aux32 >> 27 | 1; // (scale * 2 + 1) + sumi = sumi * ls / 8; // (sumi * scale + sumi / 2) / 4 + const float d = __half2float(bq2->d) * __low2float(bq8_1[iqs/2].ds); + return d * sumi; +} + +#define VDR_IQ2_XS_Q8_1_MMVQ 2 +#define VDR_IQ2_XS_Q8_1_MMQ 2 + +static __device__ __forceinline__ float vec_dot_iq2_xs_q8_1( + const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs) { + + const block_iq2_xs * bq2 = (const block_iq2_xs *) vbq + kbx; + + const int2 q2_packed = make_int2(get_int_b2(bq2->qs, iqs + 0), get_int_b2(bq2->qs, iqs + 1)); + const uint16_t * q2 = (const uint16_t *) &q2_packed; + const int ls0 = bq2->scales[iqs/2] & 0x0F; + const int ls1 = bq2->scales[iqs/2] >> 4; + + int sumi0 = 0; + int sumi1 = 0; +#pragma unroll + for (int l0 = 0; l0 < 8; l0 += 2) { + const uint2 grid_pos = ((const uint2*)iq2xs_grid)[q2[l0/2] & 0x1FF]; + const uint32_t signs = unpack_ksigns(q2[l0/2] >> 9); + + const int signs0 = __vcmpne4(signs & 0x08040201, 0); + const int grid_l = __vsub4(grid_pos.x ^ signs0, signs0); + const int u0 = get_int_b4(bq8_1[iqs/2].qs, l0 + 0); + + const int signs1 = __vcmpne4(signs & 0x80402010, 0); + const int grid_h = __vsub4(grid_pos.y ^ signs1, signs1); + const int u1 = get_int_b4(bq8_1[iqs/2].qs, l0 + 1); + + if (l0 < 4) { + sumi0 = ggml_cuda_dp4a(grid_l, u0, sumi0); + sumi0 = ggml_cuda_dp4a(grid_h, u1, sumi0); + } else { + sumi1 = ggml_cuda_dp4a(grid_l, u0, sumi1); + sumi1 = ggml_cuda_dp4a(grid_h, u1, sumi1); + } + } + const int sumi = (sumi0*ls0 + sumi1*ls1 + (sumi0 + sumi1)/2)/4; + const float d = __half2float(bq2->d) * __low2float(bq8_1[iqs/2].ds); + return d * sumi; +} + +#define VDR_IQ2_S_Q8_1_MMVQ 2 +#define VDR_IQ2_S_Q8_1_MMQ 2 + +static __device__ __forceinline__ float vec_dot_iq2_s_q8_1( + const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs) { + + const block_iq2_s * bq2 = (const block_iq2_s *) vbq + kbx; + + const int qs_packed = get_int_b2(bq2->qs, iqs/2); + const uint8_t * qs = (const uint8_t *) &qs_packed; + + const int qh = bq2->qh[iqs/2]; + + const int signs_packed_32 = get_int_b2(bq2->qs, QK_K/32 + iqs/2); + const uint8_t * signs_packed_8 = (const uint8_t *) &signs_packed_32; + + const int ls0 = bq2->scales[iqs/2] & 0x0F; + const int ls1 = bq2->scales[iqs/2] >> 4; + + int sumi0 = 0; + int sumi1 = 0; +#pragma unroll + for (int l0 = 0; l0 < 8; l0 += 2) { + const int * grid_pos = (const int *)(iq2s_grid + (qs[l0/2] | ((qh << (8-l0)) & 0x300))); + + const int signs0 = __vcmpne4(((signs_packed_8[l0/2] & 0x03) << 7) | ((signs_packed_8[l0/2] & 0x0C) << 21), 0x00000000); + const int signs1 = __vcmpne4(((signs_packed_8[l0/2] & 0x30) << 3) | ((signs_packed_8[l0/2] & 0xC0) << 17), 0x00000000); + + const int grid_l = __vsub4(grid_pos[0] ^ signs0, signs0); + const int grid_h = __vsub4(grid_pos[1] ^ signs1, signs1); + + const int u0 = get_int_b4(bq8_1[iqs/2].qs, l0 + 0); + const int u1 = get_int_b4(bq8_1[iqs/2].qs, l0 + 1); + + if (l0 < 4) { + sumi0 = ggml_cuda_dp4a(grid_l, u0, sumi0); + sumi0 = ggml_cuda_dp4a(grid_h, u1, sumi0); + } else { + sumi1 = ggml_cuda_dp4a(grid_l, u0, sumi1); + sumi1 = ggml_cuda_dp4a(grid_h, u1, sumi1); + } + } + const int sumi = (sumi0*ls0 + sumi1*ls1 + (sumi0 + sumi1)/2)/4; + + const float d = __half2float(bq2->d) * __low2float(bq8_1[iqs/2].ds); + return d * sumi; +} + +#define VDR_IQ3_XXS_Q8_1_MMVQ 2 +#define VDR_IQ3_XXS_Q8_1_MMQ 2 + +static __device__ __forceinline__ float vec_dot_iq3_xxs_q8_1( + const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs) { + + const block_iq3_xxs * bq3 = (const block_iq3_xxs *) vbq + kbx; + + const int2 q3_packed = make_int2(get_int_b2(bq3->qs, iqs), get_int_b2(bq3->qs, iqs+1)); + const uint8_t * q3 = (const uint8_t *) &q3_packed; + const uint32_t aux32 = get_int_b2(bq3->qs, QK_K/16 + iqs/2); + + int sumi = 0; +#pragma unroll + for (int l0 = 0; l0 < 8; l0 += 2) { + const int2 grid_pos = make_int2(iq3xxs_grid[q3[l0 + 0]], iq3xxs_grid[q3[l0 + 1]]); + const uint32_t signs = unpack_ksigns(aux32 >> (7*l0/2)); + + const int signs0 = __vcmpne4(signs & 0x08040201, 0); + const int grid_l = __vsub4(grid_pos.x ^ signs0, signs0); + + const int u0 = get_int_b4(bq8_1[iqs/2].qs, l0 + 0); + + const int signs1 = __vcmpne4(signs & 0x80402010, 0); + const int grid_h = __vsub4(grid_pos.y ^ signs1, signs1); + + const int u1 = get_int_b4(bq8_1[iqs/2].qs, l0 + 1); + + sumi = ggml_cuda_dp4a(grid_l, u0, sumi); + sumi = ggml_cuda_dp4a(grid_h, u1, sumi); + } + + const int ls = aux32 >> 28; + sumi = (ls*sumi + sumi/2)/2; + const float d = __half2float(bq3->d) * __low2float(bq8_1[iqs/2].ds); + return d * sumi; +} + +#define VDR_IQ3_S_Q8_1_MMVQ 2 +#define VDR_IQ3_S_Q8_1_MMQ 2 + +// TODO: don't use lookup table for signs +static __device__ __forceinline__ float vec_dot_iq3_s_q8_1( + const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs) { + + const block_iq3_s * bq3 = (const block_iq3_s *) vbq + kbx; + + const int2 qs_packed = make_int2(get_int_b2(bq3->qs, iqs + 0), get_int_b2(bq3->qs, iqs + 1)); + const uint8_t * qs = (const uint8_t *) &qs_packed; + + const int qh = bq3->qh[iqs/2]; + + const int signs_packed_32 = get_int_b2(bq3->signs, iqs/2); + const uint8_t * signs_packed_8 = (const uint8_t *) &signs_packed_32; + + int sumi = 0; +#pragma unroll + for (int l0 = 0; l0 < 8; l0 += 2) { + const int2 grid_pos = make_int2( + iq3s_grid[qs[l0 + 0] | ((qh << (8 - l0)) & 0x100)], + iq3s_grid[qs[l0 + 1] | ((qh << (7 - l0)) & 0x100)]); + + const int signs0 = __vcmpne4(((signs_packed_8[l0/2] & 0x03) << 7) | ((signs_packed_8[l0/2] & 0x0C) << 21), 0x00000000); + const int signs1 = __vcmpne4(((signs_packed_8[l0/2] & 0x30) << 3) | ((signs_packed_8[l0/2] & 0xC0) << 17), 0x00000000); + + const int grid_l = __vsub4(grid_pos.x ^ signs0, signs0); + const int grid_h = __vsub4(grid_pos.y ^ signs1, signs1); + + const int u0 = get_int_b4(bq8_1[iqs/2].qs, l0 + 0); + const int u1 = get_int_b4(bq8_1[iqs/2].qs, l0 + 1); + + sumi = ggml_cuda_dp4a(grid_l, u0, sumi); + sumi = ggml_cuda_dp4a(grid_h, u1, sumi); + } + + sumi *= 1 + 2*((bq3->scales[iqs/4] >> ((iqs << 1) & 0x04)) & 0x0F); + + const float d = __half2float(bq3->d) * __low2float(bq8_1[iqs/2].ds); + return d * sumi; +} + +#define VDR_IQ1_S_Q8_1_MMVQ 1 +#define VDR_IQ1_S_Q8_1_MMQ 1 + +static __device__ __forceinline__ float vec_dot_iq1_s_q8_1( + const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs) { + const block_iq1_s * bq1 = (const block_iq1_s *) vbq + kbx; + + const int qs_packed = get_int_b2(bq1->qs, iqs); + const uint8_t * qs = (const uint8_t *) &qs_packed; + + const int qh = bq1->qh[iqs]; + + int sumi = 0; +#pragma unroll + for (int l0 = 0; l0 < 8; l0 += 2) { + const int grid = iq1s_grid_gpu[qs[l0/2] | (((qh >> 3*(l0/2)) & 0x07) << 8)]; + + const int grid0 = (grid >> 0) & 0x0F0F0F0F; + const int grid1 = (grid >> 4) & 0x0F0F0F0F; + + const int u0 = get_int_b4(bq8_1[iqs].qs, l0 + 0); + const int u1 = get_int_b4(bq8_1[iqs].qs, l0 + 1); + + sumi = ggml_cuda_dp4a(grid0, u0, sumi); + sumi = ggml_cuda_dp4a(grid1, u1, sumi); + } + + const float d1q = __half2float(bq1->d) * (((qh >> 11) & 0x0E) + 1); + const float delta = -1.0f + IQ1S_DELTA - (qh & 0x8000) * (2.0f*IQ1S_DELTA/0x8000); + const float2 ds = __half22float2(bq8_1[iqs].ds); + return d1q * (ds.x*sumi + ds.y*delta); +} + +#define VDR_IQ1_M_Q8_1_MMVQ 1 +#define VDR_IQ1_M_Q8_1_MMQ 1 + +static __device__ __forceinline__ float vec_dot_iq1_m_q8_1( + const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs) { + + const block_iq1_m * bq1 = (const block_iq1_m *) vbq + kbx; + + const int qs_packed = get_int_b4(bq1->qs, iqs); + const uint8_t * qs = (const uint8_t *) &qs_packed; + + int sumi[2] = {0}; + float sumf[2] = {0.0f}; +#pragma unroll + for (int l0 = 0; l0 < 8; l0 += 2) { + const int qhl = bq1->qh[2*iqs + l0/4] >> (4 * ((l0/2) % 2)); + + const int grid = iq1s_grid_gpu[qs[l0/2] | ((qhl & 0x07) << 8)]; + + const int grid0 = (grid >> 0) & 0x0F0F0F0F; + const int grid1 = (grid >> 4) & 0x0F0F0F0F; + + const int u0 = get_int_b4(bq8_1[iqs].qs, l0 + 0); + const int u1 = get_int_b4(bq8_1[iqs].qs, l0 + 1); + + sumi[l0/4] = ggml_cuda_dp4a(grid0, u0, sumi[l0/4]); + sumi[l0/4] = ggml_cuda_dp4a(grid1, u1, sumi[l0/4]); + + const float delta = -1.0f + IQ1M_DELTA - (qhl & 0x08) * (2.0f*IQ1M_DELTA/0x08); + int sumy = 0; + sumy = ggml_cuda_dp4a(u0, 0x01010101, sumy); + sumy = ggml_cuda_dp4a(u1, 0x01010101, sumy); + sumf[l0/4] += delta*sumy; + } + + const uint16_t * sc = (const uint16_t *) bq1->scales; + + iq1m_scale_t scale; + scale.u16 = (sc[0] >> 12) | ((sc[1] >> 8) & 0x00F0) | ((sc[2] >> 4) & 0x0F00) | (sc[3] & 0xF000); + const float d = __half2float(scale.f16) * __low2float(bq8_1[iqs].ds); + + const int tmp = sc[iqs/2] >> (6*(iqs%2)); + const int sc0 = 2*((tmp >> 0) & 0x07) + 1; + const int sc1 = 2*((tmp >> 3) & 0x07) + 1; + return d * ((sumi[0] + sumf[0]) * sc0 + (sumi[1] + sumf[1]) * sc1); +} + +#define VDR_IQ4_NL_Q8_1_MMVQ 2 +#define VDR_IQ4_NL_Q8_1_MMQ 4 + +static __device__ __forceinline__ float vec_dot_iq4_nl_q8_1( + const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs) { + + const block_iq4_nl * bq4 = (const block_iq4_nl *) vbq + kbx; + + const int * q8 = (const int *) bq8_1->qs + iqs; + + int sumi = 0; +#pragma unroll + for (int l = 0; l < VDR_Q4_0_Q8_1_MMVQ; ++l) { + const int aux_q4 = get_int_b2(bq4->qs, iqs + l); + const int2 v = get_int_from_table_16(aux_q4, kvalues_iq4nl); + + sumi = ggml_cuda_dp4a(v.x, q8[l + 0], sumi); + sumi = ggml_cuda_dp4a(v.y, q8[l + 4], sumi); + } + + const float d = __half2float(bq4->d) * __low2float(bq8_1->ds); + return d * sumi; +} + +#define VDR_IQ4_XS_Q8_1_MMVQ 4 +#define VDR_IQ4_XS_Q8_1_MMQ 4 + +static __device__ __forceinline__ float vec_dot_iq4_xs_q8_1( + const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs) { + + const block_iq4_xs * bq4 = (const block_iq4_xs *) vbq + kbx; + + int sumi = 0; +#pragma unroll + for (int j = 0; j < 4; ++j) { + const int aux_q4 = get_int_b4(bq4->qs, iqs + j); + const int2 v = get_int_from_table_16(aux_q4, kvalues_iq4nl); + + const int u0 = get_int_b4(bq8_1[iqs/4].qs, j + 0); + const int u1 = get_int_b4(bq8_1[iqs/4].qs, j + 4); + + sumi = ggml_cuda_dp4a(v.x, u0, sumi); + sumi = ggml_cuda_dp4a(v.y, u1, sumi); + } + + const int ls = ((bq4->scales_l[iqs/8] >> (iqs & 0x04)) & 0x0F) | (((bq4->scales_h >> (iqs/2)) & 0x03) << 4); + sumi *= ls - 32; + + const float d = __half2float(bq4->d) * __low2float(bq8_1[iqs/4].ds); + return d * sumi; +} diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/vendors/cuda.h b/python/freetoken/kernel/csrc/ggml/ggml-cuda/vendors/cuda.h new file mode 100644 index 0000000000..323c980193 --- /dev/null +++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/vendors/cuda.h @@ -0,0 +1,28 @@ +#pragma once + +#include +#include +#include +#include +#include + +#ifdef GGML_USE_NCCL +#include +#endif // GGML_USE_NCCL + +#if CUDART_VERSION >= 11080 +#include +#define FP8_AVAILABLE +#endif // CUDART_VERSION >= 11080 + +#if CUDART_VERSION >= 12080 +#include +#endif // CUDART_VERSION >= 12080 + +#if CUDART_VERSION < 11020 +#define CU_DEVICE_ATTRIBUTE_VIRTUAL_MEMORY_MANAGEMENT_SUPPORTED CU_DEVICE_ATTRIBUTE_VIRTUAL_ADDRESS_MANAGEMENT_SUPPORTED +#define CUBLAS_TF32_TENSOR_OP_MATH CUBLAS_TENSOR_OP_MATH +#define CUBLAS_COMPUTE_16F CUDA_R_16F +#define CUBLAS_COMPUTE_32F CUDA_R_32F +#define cublasComputeType_t cudaDataType_t +#endif // CUDART_VERSION < 11020 diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/vendors/hip.h b/python/freetoken/kernel/csrc/ggml/ggml-cuda/vendors/hip.h new file mode 100644 index 0000000000..da7e3da0b9 --- /dev/null +++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/vendors/hip.h @@ -0,0 +1,313 @@ +#pragma once + +#define HIP_DISABLE_WARP_SYNC_BUILTINS 1 +#include +#include +#include +#include + +#ifdef GGML_USE_NCCL +#include +#endif // GGML_USE_NCCL + + +#define CUBLAS_GEMM_DEFAULT HIPBLAS_GEMM_DEFAULT +#define CUBLAS_GEMM_DEFAULT_TENSOR_OP HIPBLAS_GEMM_DEFAULT +#define CUBLAS_OP_N HIPBLAS_OP_N +#define CUBLAS_OP_T HIPBLAS_OP_T +#define CUBLAS_STATUS_SUCCESS HIPBLAS_STATUS_SUCCESS +#define CUBLAS_TF32_TENSOR_OP_MATH 0 +#define CUDA_R_16F HIPBLAS_R_16F +#define CUDA_R_16BF HIPBLAS_R_16B +#define CUDA_R_32F HIPBLAS_R_32F +#define CUBLAS_SIDE_RIGHT HIPBLAS_SIDE_RIGHT +#define CUBLAS_FILL_MODE_UPPER HIPBLAS_FILL_MODE_UPPER +#define CUBLAS_DIAG_NON_UNIT HIPBLAS_DIAG_NON_UNIT +#define CU_DEVICE_ATTRIBUTE_VIRTUAL_MEMORY_MANAGEMENT_SUPPORTED hipDeviceAttributeVirtualMemoryManagementSupported +#define CU_MEM_ALLOC_GRANULARITY_RECOMMENDED hipMemAllocationGranularityRecommended +#define CU_MEM_ALLOCATION_TYPE_PINNED hipMemAllocationTypePinned +#define CU_MEM_LOCATION_TYPE_DEVICE hipMemLocationTypeDevice +#define CU_MEM_ACCESS_FLAGS_PROT_READWRITE hipMemAccessFlagsProtReadWrite +#define CU_CHECK(fn) {hipError_t err = fn; if(err != hipSuccess) { GGML_ABORT("HipVMM Failure: %s\n", hipGetErrorString(err)); }} +#define __shfl_sync(mask, var, laneMask, width) __shfl(var, laneMask, width) +#define __shfl_up_sync(mask, var, laneMask, width) __shfl_up(var, laneMask, width) +#define __shfl_xor_sync(mask, var, laneMask, width) __shfl_xor(var, laneMask, width) +#define __all_sync(mask, var) __all(var) +#define __any_sync(mask, var) __any(var) +#define cublasStrsmBatched hipblasStrsmBatched +#define cublasCreate hipblasCreate +#define cublasDestroy hipblasDestroy +#define cublasGemmEx hipblasGemmEx +#define cublasGemmBatchedEx hipblasGemmBatchedEx +#define cublasGemmStridedBatchedEx hipblasGemmStridedBatchedEx +#define cublasHandle_t hipblasHandle_t +#define cublasSetMathMode(handle, mode) CUBLAS_STATUS_SUCCESS +#define cublasSetStream hipblasSetStream +#define cublasSgemm hipblasSgemm +#define cublasSgemmBatched hipblasSgemmBatched +#define cublasSgemmStridedBatched hipblasSgemmStridedBatched +#define cublasStatus_t hipblasStatus_t +#define cublasOperation_t hipblasOperation_t +#define cudaDevAttrCooperativeLaunch hipDeviceAttributeCooperativeLaunch +#define cudaDeviceCanAccessPeer hipDeviceCanAccessPeer +#define cudaDeviceDisablePeerAccess hipDeviceDisablePeerAccess +#define cudaDeviceEnablePeerAccess hipDeviceEnablePeerAccess +#define cudaDeviceGetAttribute hipDeviceGetAttribute +#define cudaDeviceGetPCIBusId hipDeviceGetPCIBusId +#define cudaDeviceProp hipDeviceProp_t +#define cudaDeviceSynchronize hipDeviceSynchronize +#define cudaError_t hipError_t +#define cudaErrorMemoryAllocation hipErrorOutOfMemory +#define cudaErrorPeerAccessAlreadyEnabled hipErrorPeerAccessAlreadyEnabled +#define cudaErrorPeerAccessNotEnabled hipErrorPeerAccessNotEnabled +#define cudaEventCreateWithFlags hipEventCreateWithFlags +#define cudaEventDisableTiming hipEventDisableTiming +#define cudaEventRecord hipEventRecord +#define cudaEventSynchronize hipEventSynchronize +#define cudaEvent_t hipEvent_t +#define cudaEventDestroy hipEventDestroy +#define cudaFree hipFree +#define cudaFreeHost hipHostFree +#define cudaGetDevice hipGetDevice +#define cudaGetDeviceCount hipGetDeviceCount +#define cudaGetDeviceProperties hipGetDeviceProperties +#define cudaGetErrorString hipGetErrorString +#define cudaGetLastError hipGetLastError +#define cudaHostAlloc hipHostMalloc +#define cudaHostAllocPortable hipHostMallocPortable +#define cudaHostAllocMapped hipHostMallocMapped +#define cudaHostGetDevicePointer hipHostGetDevicePointer +#define cudaHostRegister hipHostRegister +#define cudaHostRegisterPortable hipHostRegisterPortable +#define cudaHostRegisterReadOnly hipHostRegisterReadOnly +#define cudaHostUnregister hipHostUnregister +#define cudaLaunchCooperativeKernel hipLaunchCooperativeKernel +#define cudaLaunchHostFunc hipLaunchHostFunc +#define cudaMalloc hipMalloc +#define cudaMallocHost(ptr, size) hipHostMalloc(ptr, size, hipHostMallocDefault) +#define cudaMallocManaged hipMallocManaged +#define cudaMemAdvise hipMemAdvise +#define cudaMemcpy hipMemcpy +#define cudaMemcpyAsync hipMemcpyAsync +#define cudaMemcpyPeerAsync hipMemcpyPeerAsync +#define cudaMemcpy2DAsync hipMemcpy2DAsync +#define cudaMemcpyDeviceToDevice hipMemcpyDeviceToDevice +#define cudaMemcpyDeviceToHost hipMemcpyDeviceToHost +#define cudaMemcpyHostToDevice hipMemcpyHostToDevice +#define cudaMemcpyKind hipMemcpyKind +#define cudaMemset hipMemset +#define cudaMemsetAsync hipMemsetAsync +#define cudaMemGetInfo hipMemGetInfo +#define cudaOccupancyMaxPotentialBlockSize hipOccupancyMaxPotentialBlockSize +#define cudaSetDevice hipSetDevice +#define cuDeviceGet hipDeviceGet +#define CUdevice hipDevice_t +#define CUdeviceptr hipDeviceptr_t +#define cuMemUnmap hipMemUnmap +#define CUmemAccessDesc hipMemAccessDesc +#define cuMemAddressFree hipMemAddressFree +#define cuMemRelease hipMemRelease +#define CUmemGenericAllocationHandle hipMemGenericAllocationHandle_t +#define cuMemCreate hipMemCreate +#define cuMemAddressReserve hipMemAddressReserve +#define cuMemMap hipMemMap +#define cuMemSetAccess hipMemSetAccess +#define cuMemGetAllocationGranularity hipMemGetAllocationGranularity +#define CUmemAllocationProp hipMemAllocationProp +#define cuDeviceGetAttribute hipDeviceGetAttribute +#define cudaStreamCreateWithFlags hipStreamCreateWithFlags +#define cudaStreamDestroy hipStreamDestroy +#define cudaStreamFireAndForget hipStreamFireAndForget +#define cudaStreamNonBlocking hipStreamNonBlocking +#define cudaStreamPerThread hipStreamPerThread +#define cudaStreamSynchronize hipStreamSynchronize +#define cudaStreamWaitEvent hipStreamWaitEvent +#define cudaGraphExec_t hipGraphExec_t +#define cudaGraphNode_t hipGraphNode_t +#define cudaKernelNodeParams hipKernelNodeParams +#define cudaKernelNodeParams hipKernelNodeParams +#define cudaGraphExecDestroy hipGraphExecDestroy +#define cudaGraphLaunch hipGraphLaunch +#define cudaErrorGraphExecUpdateFailure hipErrorGraphExecUpdateFailure +#define cudaGraphExecUpdateResult hipGraphExecUpdateResult +#define cudaGraphNodeType hipGraphNodeType +#define cudaGraphNodeTypeKernel hipGraphNodeTypeKernel +#define cudaGraphInstantiate hipGraphInstantiate +#define cudaStreamEndCapture hipStreamEndCapture +#define cudaGraphDestroy hipGraphDestroy +#define cudaGraphKernelNodeSetParams hipGraphKernelNodeSetParams +#define cudaErrorInvalidDeviceFunction hipErrorInvalidDeviceFunction +#define cudaGraphKernelNodeGetParams hipGraphKernelNodeGetParams +#define cudaGraphNodeGetType hipGraphNodeGetType +#define cudaGraphGetNodes hipGraphGetNodes +#define cudaGraphExecUpdate hipGraphExecUpdate +#define cudaStreamCaptureModeRelaxed hipStreamCaptureModeRelaxed +#define cudaStreamBeginCapture hipStreamBeginCapture +#define cudaGraph_t hipGraph_t +#define cudaStream_t hipStream_t +#define cudaSuccess hipSuccess +#define cudaOccupancyMaxActiveBlocksPerMultiprocessor hipOccupancyMaxActiveBlocksPerMultiprocessor +#define cudaFuncSetAttribute hipFuncSetAttribute +#define cudaFuncAttributeMaxDynamicSharedMemorySize hipFuncAttributeMaxDynamicSharedMemorySize +#define __trap() do { abort(); __builtin_unreachable(); } while(0) +#define CUBLAS_STATUS_SUCCESS HIPBLAS_STATUS_SUCCESS +#define CUBLAS_STATUS_NOT_INITIALIZED HIPBLAS_STATUS_NOT_INITIALIZED +#define CUBLAS_STATUS_ALLOC_FAILED HIPBLAS_STATUS_ALLOC_FAILED +#define CUBLAS_STATUS_INVALID_VALUE HIPBLAS_STATUS_INVALID_VALUE +#define CUBLAS_STATUS_ARCH_MISMATCH HIPBLAS_STATUS_ARCH_MISMATCH +#define CUBLAS_STATUS_MAPPING_ERROR HIPBLAS_STATUS_MAPPING_ERROR +#define CUBLAS_STATUS_EXECUTION_FAILED HIPBLAS_STATUS_EXECUTION_FAILED +#define CUBLAS_STATUS_INTERNAL_ERROR HIPBLAS_STATUS_INTERNAL_ERROR +#define CUBLAS_STATUS_NOT_SUPPORTED HIPBLAS_STATUS_NOT_SUPPORTED + +#if HIP_VERSION >= 60500000 +#define CUBLAS_COMPUTE_16F HIPBLAS_COMPUTE_16F +#define CUBLAS_COMPUTE_32F HIPBLAS_COMPUTE_32F +#define CUBLAS_COMPUTE_32F_FAST_16F HIPBLAS_COMPUTE_32F_FAST_16F +#define cublasComputeType_t hipblasComputeType_t +#define cudaDataType_t hipDataType +#else +#define CUBLAS_COMPUTE_16F HIPBLAS_R_16F +#define CUBLAS_COMPUTE_32F HIPBLAS_R_32F +#define CUBLAS_COMPUTE_32F_FAST_16F HIPBLAS_R_32F +#define cublasComputeType_t hipblasDatatype_t +#define cudaDataType_t hipblasDatatype_t +#endif // HIP_VERSION >= 6050000 + +#if !defined(__HIP_PLATFORM_AMD__) +#error "The HIP backend supports only AMD targets" +#endif // !defined(__HIP_PLATFORM_AMD__) + +#define __CUDA_ARCH__ 1300 + +#if defined(__gfx900__) || defined(__gfx906__) || defined(__gfx909__) || defined(__gfx90c__) +#define GCN5 +#endif // defined(__gfx900__) || defined(__gfx906__) || defined(__gfx909__) || defined(__gfx90c__) + +#if defined(__gfx803__) +#define GCN4 +#endif // defined(__gfx803__) + +#if defined(GCN5) || defined(GCN4) +#define GCN +#endif // defined(GCN5) || defined(GCN4) + +#if defined(__gfx950__) +#define CDNA4 +#endif // defined(__gfx950__) + +#if defined(__gfx942__) +#define CDNA3 +#endif // defined(__gfx942__) + +#if defined(__gfx90a__) +#define CDNA2 +#endif // defined(__gfx90a__) + +#if defined(__gfx908__) +#define CDNA1 +#endif // defined(__gfx908__) + +#if defined(CDNA4) || defined(CDNA3) || defined(CDNA2) || defined(CDNA1) +#define CDNA // For the entire family +#endif // defined(CDNA4) || defined(CDNA3) || defined(CDNA2) || defined(CDNA1) + +#if defined(__GFX12__) +#define RDNA4 +#endif // defined(__GFX12__) + +#if defined(__GFX11__) +#define RDNA3 +#endif // defined(__GFX11__) + +#if defined(__gfx1150__) || defined(__gfx1151__) || defined(__gfx1152__) || defined(__gfx1153__) +#define RDNA3_5 +#endif // defined(__gfx1150__) || defined(__gfx1151__) || defined(__gfx1152__) || defined(__gfx1153__) + +#if defined(RDNA3) && !defined(RDNA3_5) +#define RDNA3_0 +#endif // defined(RDNA3) && !defined(RDNA3_5) + +#if defined(__gfx1030__) || defined(__gfx1031__) || defined(__gfx1032__) || defined(__gfx1033__) || \ + defined(__gfx1034__) || defined(__gfx1035__) || defined(__gfx1036__) || defined(__gfx1037__) +#define RDNA2 +#endif + +#if defined(__gfx1010__) || defined(__gfx1012__) +#define RDNA1 +#endif // defined(__gfx1010__) || defined(__gfx1012__) + +#if defined(RDNA4) || defined(RDNA3) || defined(RDNA2) || defined(RDNA1) +#define RDNA // For the entire family +#endif // defined(RDNA4) || defined(RDNA3) || defined(RDNA2) || defined(RDNA1) + +#ifndef __has_builtin + #define __has_builtin(x) 0 +#endif + +typedef __hip_bfloat16 nv_bfloat16; +typedef __hip_bfloat162 nv_bfloat162; + +#if HIP_VERSION >= 60300000 +#include +typedef __hip_fp8_e4m3 __nv_fp8_e4m3; +#define FP8_AVAILABLE +#endif // HIP_VERSION >= 60300000 + +typedef int8_t int8x4_t __attribute__((ext_vector_type(4))); +typedef uint8_t uint8x4_t __attribute__((ext_vector_type(4))); +static __device__ __forceinline__ int __vsubss4(const int a, const int b) { + const int8x4_t va = reinterpret_cast(a); + const int8x4_t vb = reinterpret_cast(b); +#if __has_builtin(__builtin_elementwise_sub_sat) + const int8x4_t c = __builtin_elementwise_sub_sat(va, vb); + return reinterpret_cast(c); +#else + int8x4_t c; + int16_t tmp; +#pragma unroll + for (int i = 0; i < 4; i++) { + tmp = va[i] - vb[i]; + if(tmp > std::numeric_limits::max()) tmp = std::numeric_limits::max(); + if(tmp < std::numeric_limits::min()) tmp = std::numeric_limits::min(); + c[i] = tmp; + } + return reinterpret_cast(c); +#endif // __has_builtin(__builtin_elementwise_sub_sat) +} + +static __device__ __forceinline__ int __vsub4(const int a, const int b) { + // do some small modifications to a and b to make the subtraction not underflow + const unsigned int a_large = a | 0x80808080; + const unsigned int b_small = b & 0x7f7f7f7f; + const unsigned int result_low_7bits = a_large - b_small; + + // if two ops share the same high bit, we should flip the high bit of the result + const unsigned int should_flip_high_1bit = (a ^ ~b) & 0x80808080; + + return result_low_7bits ^ should_flip_high_1bit; +} + +static __device__ __forceinline__ unsigned int __vcmpeq4(unsigned int a, unsigned int b) { + const uint8x4_t& va = reinterpret_cast(a); + const uint8x4_t& vb = reinterpret_cast(b); + unsigned int c; + uint8x4_t& vc = reinterpret_cast(c); +#pragma unroll + for (int i = 0; i < 4; ++i) { + vc[i] = va[i] == vb[i] ? 0xff : 0x00; + } + return c; +} + +static __device__ __forceinline__ unsigned int __vcmpne4(unsigned int a, unsigned int b) { + const unsigned int x = a ^ b; + + // any non-equal bit in a byte will set the high bit of that byte here + // the addition will not overflow in the byte as op1 and op2 are both less than 0x80 + const unsigned int ne_low_7bits = ((x & 0x7f7f7f7f) + 0x7f7f7f7f) & 0x80808080; + const unsigned int ne_high_1bit = x & 0x80808080; + const unsigned int ne_any_bit = ne_low_7bits | ne_high_1bit; + + return (ne_any_bit >> 7) * 0xff; +} diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/vendors/musa.h b/python/freetoken/kernel/csrc/ggml/ggml-cuda/vendors/musa.h new file mode 100644 index 0000000000..22126decfb --- /dev/null +++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/vendors/musa.h @@ -0,0 +1,161 @@ +#pragma once + +#include +#include +#include +#include +#include + +#ifdef __MUSA_ARCH__ +#define __CUDA_ARCH__ 1300 // GGML_CUDA_CC_RUBIN +#endif // __MUSA_ARCH__ + +#define CUBLAS_COMPUTE_16F CUDA_R_16F +#define CUBLAS_COMPUTE_32F CUDA_R_32F +#define CUBLAS_COMPUTE_32F_FAST_16F MUBLAS_COMPUTE_32F_FAST_16F +#define CUBLAS_GEMM_DEFAULT MUBLAS_GEMM_DEFAULT +#define CUBLAS_GEMM_DEFAULT_TENSOR_OP MUBLAS_GEMM_DEFAULT +#define CUBLAS_OP_N MUBLAS_OP_N +#define CUBLAS_OP_T MUBLAS_OP_T +#define CUBLAS_DEFAULT_MATH MUBLAS_DEFAULT_MATH +#define CUBLAS_SIDE_RIGHT MUBLAS_SIDE_RIGHT +#define CUBLAS_FILL_MODE_UPPER MUBLAS_FILL_MODE_UPPER +#define CUBLAS_DIAG_NON_UNIT MUBLAS_DIAG_NON_UNIT +#define CUBLAS_STATUS_SUCCESS MUBLAS_STATUS_SUCCESS +#define CUBLAS_TF32_TENSOR_OP_MATH MUBLAS_TENSOR_OP_MATH +#define CUDA_R_16F MUSA_R_16F +#define CUDA_R_16BF MUSA_R_16BF +#define CUDA_R_32F MUSA_R_32F +#define cublasStrsmBatched mublasStrsmBatched +#define cublasComputeType_t cudaDataType_t +#define cublasCreate mublasCreate +#define cublasDestroy mublasDestroy +#define cublasGemmEx mublasGemmEx +#define cublasGemmBatchedEx mublasGemmBatchedEx +#define cublasGemmStridedBatchedEx mublasGemmStridedBatchedEx +#define cublasHandle_t mublasHandle_t +#define cublasSetMathMode mublasSetMathMode +#define cublasSetStream mublasSetStream +#define cublasSgemm mublasSgemm +#define cublasSgemmBatched mublasSgemmBatched +#define cublasSgemmStridedBatched mublasSgemmStridedBatched +#define cublasStatus_t mublasStatus_t +#define cublasOperation_t mublasOperation_t +#define cublasGetStatusString mublasGetStatusString +#define cudaDataType_t musaDataType_t +#define cudaDeviceCanAccessPeer musaDeviceCanAccessPeer +#define cudaDeviceDisablePeerAccess musaDeviceDisablePeerAccess +#define cudaDeviceEnablePeerAccess musaDeviceEnablePeerAccess +#define cudaDeviceGetPCIBusId musaDeviceGetPCIBusId +#define cudaDeviceProp musaDeviceProp +#define cudaDeviceSynchronize musaDeviceSynchronize +#define cudaDeviceGetAttribute musaDeviceGetAttribute +#define cudaError_t musaError_t +#define cudaErrorMemoryAllocation musaErrorMemoryAllocation +#define cudaErrorPeerAccessAlreadyEnabled musaErrorPeerAccessAlreadyEnabled +#define cudaErrorPeerAccessNotEnabled musaErrorPeerAccessNotEnabled +#define cudaEventCreateWithFlags musaEventCreateWithFlags +#define cudaEventDisableTiming musaEventDisableTiming +#define cudaEventRecord musaEventRecord +#define cudaEventSynchronize musaEventSynchronize +#define cudaEvent_t musaEvent_t +#define cudaEventDestroy musaEventDestroy +#define cudaFree musaFree +#define cudaFreeHost musaFreeHost +#define cudaGetDevice musaGetDevice +#define cudaGetDeviceCount musaGetDeviceCount +#define cudaGetDeviceProperties musaGetDeviceProperties +#define cudaGetErrorString musaGetErrorString +#define cudaGetLastError musaGetLastError +#define cudaHostRegister musaHostRegister +#define cudaHostRegisterPortable musaHostRegisterPortable +#define cudaHostRegisterReadOnly musaHostRegisterReadOnly +#define cudaHostUnregister musaHostUnregister +#define cudaLaunchCooperativeKernel musaLaunchCooperativeKernel +#define cudaLaunchHostFunc musaLaunchHostFunc +#define cudaMalloc musaMalloc +#define cudaMallocHost musaMallocHost +#define cudaMallocManaged musaMallocManaged +#define cudaMemcpy musaMemcpy +#define cudaMemcpyAsync musaMemcpyAsync +#define cudaMemcpyPeerAsync musaMemcpyPeerAsync +#define cudaMemcpy2DAsync musaMemcpy2DAsync +#define cudaMemcpyDeviceToDevice musaMemcpyDeviceToDevice +#define cudaMemcpyDeviceToHost musaMemcpyDeviceToHost +#define cudaMemcpyHostToDevice musaMemcpyHostToDevice +#define cudaMemcpyKind musaMemcpyKind +#define cudaMemset musaMemset +#define cudaMemsetAsync musaMemsetAsync +#define cudaMemGetInfo musaMemGetInfo +#define cudaOccupancyMaxPotentialBlockSize musaOccupancyMaxPotentialBlockSize +#define cudaSetDevice musaSetDevice +#define cudaStreamCreateWithFlags musaStreamCreateWithFlags +#define cudaStreamDestroy musaStreamDestroy +#define cudaStreamFireAndForget musaStreamFireAndForget +#define cudaStreamNonBlocking musaStreamNonBlocking +#define cudaStreamPerThread musaStreamPerThread +#define cudaStreamSynchronize musaStreamSynchronize +#define cudaStreamWaitEvent musaStreamWaitEvent +#define cudaStream_t musaStream_t +#define cudaSuccess musaSuccess + +// Additional mappings for MUSA virtual memory pool +#define CU_DEVICE_ATTRIBUTE_VIRTUAL_MEMORY_MANAGEMENT_SUPPORTED MU_DEVICE_ATTRIBUTE_VIRTUAL_ADDRESS_MANAGEMENT_SUPPORTED +#define CU_MEM_ACCESS_FLAGS_PROT_READWRITE MU_MEM_ACCESS_FLAGS_PROT_READWRITE +#define CU_MEM_ALLOC_GRANULARITY_RECOMMENDED MU_MEM_ALLOC_GRANULARITY_RECOMMENDED +#define CU_MEM_ALLOCATION_TYPE_PINNED MU_MEM_ALLOCATION_TYPE_PINNED +#define CU_MEM_LOCATION_TYPE_DEVICE MU_MEM_LOCATION_TYPE_DEVICE +#define CUdevice MUdevice +#define CUdeviceptr MUdeviceptr +#define CUmemAccessDesc MUmemAccessDesc +#define CUmemAllocationProp MUmemAllocationProp +#define CUmemGenericAllocationHandle MUmemGenericAllocationHandle +#define cuDeviceGet muDeviceGet +#define cuDeviceGetAttribute muDeviceGetAttribute +#define cuMemAddressFree muMemAddressFree +#define cuMemAddressReserve muMemAddressReserve +#define cuMemCreate muMemCreate +#define cuMemGetAllocationGranularity muMemGetAllocationGranularity +#define cuMemMap muMemMap +#define cuMemRelease muMemRelease +#define cuMemSetAccess muMemSetAccess +#define cuMemUnmap muMemUnmap +#define cudaDevAttrCooperativeLaunch musaDevAttrCooperativeLaunch +#define cudaFuncAttributeMaxDynamicSharedMemorySize musaFuncAttributeMaxDynamicSharedMemorySize +#define cudaFuncSetAttribute musaFuncSetAttribute +#define cudaMemcpy3DPeerParms musaMemcpy3DPeerParms +#define make_cudaExtent make_musaExtent +#define make_cudaPitchedPtr make_musaPitchedPtr + +// Additional mappings for MUSA graphs +#define CUDA_SUCCESS MUSA_SUCCESS +#define CUresult MUresult +#define cuGetErrorString muGetErrorString +#define cudaErrorGraphExecUpdateFailure musaErrorGraphExecUpdateFailure +#define cudaErrorInvalidDeviceFunction musaErrorInvalidDeviceFunction +#define cudaGraphDestroy musaGraphDestroy +#define cudaGraphExecDestroy musaGraphExecDestroy +#define cudaGraphExec_t musaGraphExec_t +#define cudaGraphExecUpdate musaGraphExecUpdate +#define cudaGraphExecUpdateResult musaGraphExecUpdateResult +#define cudaGraphGetNodes musaGraphGetNodes +#define cudaGraphInstantiate musaGraphInstantiate +#define cudaGraphKernelNodeGetParams musaGraphKernelNodeGetParams +#define cudaGraphKernelNodeSetParams musaGraphKernelNodeSetParams +#define cudaGraphLaunch musaGraphLaunch +#define cudaGraphNodeGetType musaGraphNodeGetType +#define cudaGraphNode_t musaGraphNode_t +#define cudaGraphNodeType musaGraphNodeType +#define cudaGraphNodeTypeKernel musaGraphNodeTypeKernel +#define cudaGraph_t musaGraph_t +#define cudaKernelNodeParams musaKernelNodeParams +#define cudaStreamCaptureModeRelaxed musaStreamCaptureModeRelaxed +#define cudaStreamBeginCapture musaStreamBeginCapture +#define cudaStreamEndCapture musaStreamEndCapture +#define cudaStreamCaptureStatus musaStreamCaptureStatus +#define cudaStreamCaptureStatusNone musaStreamCaptureStatusNone +#define cudaStreamIsCapturing musaStreamIsCapturing +#define cudaOccupancyMaxActiveBlocksPerMultiprocessor musaOccupancyMaxActiveBlocksPerMultiprocessor + +typedef __mt_bfloat16 nv_bfloat16; +typedef __mt_bfloat162 nv_bfloat162; diff --git a/python/freetoken/kernel/csrc/ggml/include/ggml-alloc.h b/python/freetoken/kernel/csrc/ggml/include/ggml-alloc.h new file mode 100644 index 0000000000..81cce2e2a6 --- /dev/null +++ b/python/freetoken/kernel/csrc/ggml/include/ggml-alloc.h @@ -0,0 +1,87 @@ +#pragma once + +#include "ggml.h" + +#ifdef __cplusplus +extern "C" { +#endif + +typedef struct ggml_backend_buffer_type * ggml_backend_buffer_type_t; +typedef struct ggml_backend_buffer * ggml_backend_buffer_t; +typedef struct ggml_backend * ggml_backend_t; + +// Tensor allocator +struct ggml_tallocr { + ggml_backend_buffer_t buffer; + void * base; + size_t alignment; + size_t offset; +}; + +GGML_API struct ggml_tallocr ggml_tallocr_new(ggml_backend_buffer_t buffer); +GGML_API enum ggml_status ggml_tallocr_alloc(struct ggml_tallocr * talloc, struct ggml_tensor * tensor); + +// Graph allocator +/* + Example usage: + ggml_gallocr_t galloc = ggml_gallocr_new(ggml_backend_cpu_buffer_type()); + + // optional: create a worst-case graph and reserve the buffers to avoid reallocations + ggml_gallocr_reserve(galloc, build_graph(max_batch)); + + // allocate the graph + struct ggml_cgraph * graph = build_graph(batch); + ggml_gallocr_alloc_graph(galloc, graph); + + printf("compute buffer size: %zu bytes\n", ggml_gallocr_get_buffer_size(galloc, 0)); + + // evaluate the graph + ggml_backend_graph_compute(backend, graph); +*/ + +// special tensor flags for use with the graph allocator: +// ggml_set_input(): all input tensors are allocated at the beginning of the graph in non-overlapping addresses +// ggml_set_output(): output tensors are never freed and never overwritten + +typedef struct ggml_gallocr * ggml_gallocr_t; + +GGML_API ggml_gallocr_t ggml_gallocr_new(ggml_backend_buffer_type_t buft); +GGML_API ggml_gallocr_t ggml_gallocr_new_n(ggml_backend_buffer_type_t * bufts, int n_bufs); +GGML_API void ggml_gallocr_free(ggml_gallocr_t galloc); + +// pre-allocate buffers from a measure graph - does not allocate or modify the graph +// call with a worst-case graph to avoid buffer reallocations +// not strictly required for single buffer usage: ggml_gallocr_alloc_graph will reallocate the buffers automatically if needed +// returns false if the buffer allocation failed +// ggml_gallocr_resrve_n_size writes the buffer sizes per galloc buffer that would be allocated by ggml_gallocr_reserve_n to sizes +GGML_API bool ggml_gallocr_reserve(ggml_gallocr_t galloc, struct ggml_cgraph * graph); +GGML_API void ggml_gallocr_reserve_n_size( + ggml_gallocr_t galloc, + struct ggml_cgraph * graph, + const int * node_buffer_ids, + const int * leaf_buffer_ids, + size_t * sizes); +GGML_API bool ggml_gallocr_reserve_n( + ggml_gallocr_t galloc, + struct ggml_cgraph * graph, + const int * node_buffer_ids, + const int * leaf_buffer_ids); + +// automatic reallocation if the topology changes when using a single buffer +// returns false if using multiple buffers and a re-allocation is needed (call ggml_gallocr_reserve_n first to set the node buffers) +GGML_API bool ggml_gallocr_alloc_graph(ggml_gallocr_t galloc, struct ggml_cgraph * graph); + +GGML_API size_t ggml_gallocr_get_buffer_size(ggml_gallocr_t galloc, int buffer_id); + +// Utils +// Create a buffer and allocate all the tensors in a ggml_context + +// returns the size of the buffer that would be allocated by ggml_backend_alloc_ctx_tensors_from_buft. returns 0 on failure +GGML_API size_t ggml_backend_alloc_ctx_tensors_from_buft_size(struct ggml_context * ctx, ggml_backend_buffer_type_t buft); +// returns NULL on failure or if all tensors in ctx are already allocated or zero-sized +GGML_API struct ggml_backend_buffer * ggml_backend_alloc_ctx_tensors_from_buft(struct ggml_context * ctx, ggml_backend_buffer_type_t buft); +GGML_API struct ggml_backend_buffer * ggml_backend_alloc_ctx_tensors(struct ggml_context * ctx, ggml_backend_t backend); + +#ifdef __cplusplus +} +#endif diff --git a/python/freetoken/kernel/csrc/ggml/include/ggml-backend.h b/python/freetoken/kernel/csrc/ggml/include/ggml-backend.h new file mode 100644 index 0000000000..7a65d8121b --- /dev/null +++ b/python/freetoken/kernel/csrc/ggml/include/ggml-backend.h @@ -0,0 +1,463 @@ +#pragma once + +#include "ggml.h" +#include "ggml-alloc.h" + +#ifdef GGML_BACKEND_SHARED +# if defined(_WIN32) && !defined(__MINGW32__) +# ifdef GGML_BACKEND_BUILD +# define GGML_BACKEND_API __declspec(dllexport) extern +# else +# define GGML_BACKEND_API __declspec(dllimport) extern +# endif +# else +# define GGML_BACKEND_API __attribute__ ((visibility ("default"))) extern +# endif +#else +# define GGML_BACKEND_API extern +#endif + +#ifdef __cplusplus +extern "C" { +#endif + + typedef struct ggml_backend_buffer_type * ggml_backend_buffer_type_t; + typedef struct ggml_backend_buffer * ggml_backend_buffer_t; + typedef struct ggml_backend_event * ggml_backend_event_t; + typedef struct ggml_backend * ggml_backend_t; + typedef void * ggml_backend_graph_plan_t; + typedef struct ggml_backend_reg * ggml_backend_reg_t; + typedef struct ggml_backend_device * ggml_backend_dev_t; + + + // + // Backend buffer type + // + + GGML_API const char * ggml_backend_buft_name (ggml_backend_buffer_type_t buft); + GGML_API ggml_backend_buffer_t ggml_backend_buft_alloc_buffer (ggml_backend_buffer_type_t buft, size_t size); + GGML_API ggml_backend_buffer_t ggml_backend_buft_alloc_buffer_n (ggml_backend_buffer_type_t buft, struct ggml_tensor ** tensors, int n_tensors); + GGML_API size_t ggml_backend_buft_get_alignment (ggml_backend_buffer_type_t buft); + GGML_API size_t ggml_backend_buft_get_max_size (ggml_backend_buffer_type_t buft); + GGML_API size_t ggml_backend_buft_get_alloc_size (ggml_backend_buffer_type_t buft, const struct ggml_tensor * tensor); + GGML_API size_t ggml_backend_buft_get_alloc_size_n(ggml_backend_buffer_type_t buft, struct ggml_tensor ** tensors, int n_tensors); + GGML_API bool ggml_backend_buft_is_host (ggml_backend_buffer_type_t buft); + GGML_API ggml_backend_dev_t ggml_backend_buft_get_device (ggml_backend_buffer_type_t buft); + + // + // Backend buffer + // + + enum ggml_backend_buffer_usage { + GGML_BACKEND_BUFFER_USAGE_ANY = 0, + GGML_BACKEND_BUFFER_USAGE_WEIGHTS = 1, + GGML_BACKEND_BUFFER_USAGE_COMPUTE = 2, + }; + + GGML_API const char * ggml_backend_buffer_name (ggml_backend_buffer_t buffer); + GGML_API void ggml_backend_buffer_free (ggml_backend_buffer_t buffer); + GGML_API void * ggml_backend_buffer_get_base (ggml_backend_buffer_t buffer); + GGML_API size_t ggml_backend_buffer_get_size (ggml_backend_buffer_t buffer); + GGML_API enum ggml_status ggml_backend_buffer_init_tensor (ggml_backend_buffer_t buffer, struct ggml_tensor * tensor); + GGML_API size_t ggml_backend_buffer_get_alignment (ggml_backend_buffer_t buffer); + GGML_API size_t ggml_backend_buffer_get_max_size (ggml_backend_buffer_t buffer); + GGML_API size_t ggml_backend_buffer_get_alloc_size(ggml_backend_buffer_t buffer, const struct ggml_tensor * tensor); + GGML_API void ggml_backend_buffer_clear (ggml_backend_buffer_t buffer, uint8_t value); + GGML_API bool ggml_backend_buffer_is_host (ggml_backend_buffer_t buffer); + GGML_API void ggml_backend_buffer_set_usage (ggml_backend_buffer_t buffer, enum ggml_backend_buffer_usage usage); + GGML_API enum ggml_backend_buffer_usage ggml_backend_buffer_get_usage (ggml_backend_buffer_t buffer); + GGML_API ggml_backend_buffer_type_t ggml_backend_buffer_get_type (ggml_backend_buffer_t buffer); + GGML_API void ggml_backend_buffer_reset (ggml_backend_buffer_t buffer); + + // tensor copy between different backends + GGML_API void ggml_backend_tensor_copy(const struct ggml_tensor * src, struct ggml_tensor * dst); + + // + // Backend (stream) + // + + GGML_API ggml_guid_t ggml_backend_guid(ggml_backend_t backend); + GGML_API const char * ggml_backend_name(ggml_backend_t backend); + GGML_API void ggml_backend_free(ggml_backend_t backend); + + GGML_API ggml_backend_buffer_type_t ggml_backend_get_default_buffer_type(ggml_backend_t backend); + GGML_API ggml_backend_buffer_t ggml_backend_alloc_buffer(ggml_backend_t backend, size_t size); + GGML_API size_t ggml_backend_get_alignment(ggml_backend_t backend); + GGML_API size_t ggml_backend_get_max_size(ggml_backend_t backend); + + GGML_API void ggml_backend_tensor_set_async (ggml_backend_t backend, struct ggml_tensor * tensor, const void * data, size_t offset, size_t size); + GGML_API void ggml_backend_tensor_get_async (ggml_backend_t backend, const struct ggml_tensor * tensor, void * data, size_t offset, size_t size); + GGML_API void ggml_backend_tensor_set_2d_async(ggml_backend_t backend, struct ggml_tensor * tensor, const void * data, size_t offset, size_t size, size_t n_copies, size_t stride_tensor, size_t stride_data); + GGML_API void ggml_backend_tensor_get_2d_async(ggml_backend_t backend, const struct ggml_tensor * tensor, void * data, size_t offset, size_t size, size_t n_copies, size_t stride_tensor, size_t stride_data); + + // "offset" refers to the offset in tensor->data for setting/getting data + GGML_API void ggml_backend_tensor_set ( struct ggml_tensor * tensor, const void * data, size_t offset, size_t size); + GGML_API void ggml_backend_tensor_get (const struct ggml_tensor * tensor, void * data, size_t offset, size_t size); + GGML_API void ggml_backend_tensor_set_2d( struct ggml_tensor * tensor, const void * data, size_t offset, size_t size, size_t n_copies, size_t stride_tensor, size_t stride_data); + GGML_API void ggml_backend_tensor_get_2d(const struct ggml_tensor * tensor, void * data, size_t offset, size_t size, size_t n_copies, size_t stride_tensor, size_t stride_data); + GGML_API void ggml_backend_tensor_memset( struct ggml_tensor * tensor, uint8_t value, size_t offset, size_t size); + + GGML_API void ggml_backend_synchronize(ggml_backend_t backend); + + GGML_API ggml_backend_graph_plan_t ggml_backend_graph_plan_create(ggml_backend_t backend, struct ggml_cgraph * cgraph); + GGML_API void ggml_backend_graph_plan_free (ggml_backend_t backend, ggml_backend_graph_plan_t plan); + + GGML_API enum ggml_status ggml_backend_graph_plan_compute (ggml_backend_t backend, ggml_backend_graph_plan_t plan); + GGML_API enum ggml_status ggml_backend_graph_compute (ggml_backend_t backend, struct ggml_cgraph * cgraph); + GGML_API enum ggml_status ggml_backend_graph_compute_async(ggml_backend_t backend, struct ggml_cgraph * cgraph); + + // NOTE: will be removed, use device version instead + GGML_API bool ggml_backend_supports_op(ggml_backend_t backend, const struct ggml_tensor * op); + GGML_API bool ggml_backend_supports_buft(ggml_backend_t backend, ggml_backend_buffer_type_t buft); + GGML_API bool ggml_backend_offload_op(ggml_backend_t backend, const struct ggml_tensor * op); + + // asynchronous copy + // the copy is performed after all the currently queued operations in backend_src + // backend_dst will wait for the copy to complete before performing other operations + // automatic fallback to sync copy if async is not supported + GGML_API void ggml_backend_tensor_copy_async(ggml_backend_t backend_src, ggml_backend_t backend_dst, const struct ggml_tensor * src, struct ggml_tensor * dst); + + GGML_API ggml_backend_dev_t ggml_backend_get_device(ggml_backend_t backend); + + // + // Events + // + + GGML_API ggml_backend_event_t ggml_backend_event_new(ggml_backend_dev_t device); + GGML_API void ggml_backend_event_free(ggml_backend_event_t event); + GGML_API void ggml_backend_event_record(ggml_backend_event_t event, ggml_backend_t backend); + GGML_API void ggml_backend_event_synchronize(ggml_backend_event_t event); + GGML_API void ggml_backend_event_wait(ggml_backend_t backend, ggml_backend_event_t event); + + // + // Backend device + // + + enum ggml_backend_dev_type { + // CPU device using system memory + GGML_BACKEND_DEVICE_TYPE_CPU, + // GPU device using dedicated memory + GGML_BACKEND_DEVICE_TYPE_GPU, + // integrated GPU device using host memory + GGML_BACKEND_DEVICE_TYPE_IGPU, + // accelerator devices intended to be used together with the CPU backend (e.g. BLAS or AMX) + GGML_BACKEND_DEVICE_TYPE_ACCEL, + // "meta" device wrapping multiple other devices for tensor parallelism + GGML_BACKEND_DEVICE_TYPE_META, + }; + + // functionality supported by the device + struct ggml_backend_dev_caps { + // asynchronous operations + bool async; + // pinned host buffer + bool host_buffer; + // creating buffers from host ptr + bool buffer_from_host_ptr; + // event synchronization + bool events; + // mmap is supported for loading + bool mmap_support; + }; + + // all the device properties + struct ggml_backend_dev_props { + // device name + const char * name; + // device description + const char * description; + // device free memory in bytes + size_t memory_free; + // device total memory in bytes + size_t memory_total; + // device type + enum ggml_backend_dev_type type; + // device id + // for PCI devices, this should be the lower-case PCI bus id formatted as "domain:bus:device.function" (e.g. "0000:c1:00.0") + // if the id is unknown, this should be NULL + const char * device_id; + // device capabilities + struct ggml_backend_dev_caps caps; + }; + + GGML_API const char * ggml_backend_dev_name(ggml_backend_dev_t device); + GGML_API const char * ggml_backend_dev_description(ggml_backend_dev_t device); + GGML_API void ggml_backend_dev_memory(ggml_backend_dev_t device, size_t * free, size_t * total); + GGML_API enum ggml_backend_dev_type ggml_backend_dev_type(ggml_backend_dev_t device); + GGML_API void ggml_backend_dev_get_props(ggml_backend_dev_t device, struct ggml_backend_dev_props * props); + GGML_API ggml_backend_reg_t ggml_backend_dev_backend_reg(ggml_backend_dev_t device); + GGML_API ggml_backend_t ggml_backend_dev_init(ggml_backend_dev_t device, const char * params); + GGML_API ggml_backend_buffer_type_t ggml_backend_dev_buffer_type(ggml_backend_dev_t device); + GGML_API ggml_backend_buffer_type_t ggml_backend_dev_host_buffer_type(ggml_backend_dev_t device); + GGML_API ggml_backend_buffer_t ggml_backend_dev_buffer_from_host_ptr(ggml_backend_dev_t device, void * ptr, size_t size, size_t max_tensor_size); + + GGML_API bool ggml_backend_dev_supports_op(ggml_backend_dev_t device, const struct ggml_tensor * op); + GGML_API bool ggml_backend_dev_supports_buft(ggml_backend_dev_t device, ggml_backend_buffer_type_t buft); + GGML_API bool ggml_backend_dev_offload_op(ggml_backend_dev_t device, const struct ggml_tensor * op); + + // + // Backend (reg) + // + + GGML_API const char * ggml_backend_reg_name(ggml_backend_reg_t reg); + GGML_API size_t ggml_backend_reg_dev_count(ggml_backend_reg_t reg); + GGML_API ggml_backend_dev_t ggml_backend_reg_dev_get(ggml_backend_reg_t reg, size_t index); + GGML_API void * ggml_backend_reg_get_proc_address(ggml_backend_reg_t reg, const char * name); + + // Common functions that may be obtained using ggml_backend_reg_get_proc_address + + // Context management and operations for faster communication between backends, used for tensor parallelism (meta backend) + typedef void * (*ggml_backend_comm_init_t)(ggml_backend_t * backends, size_t n_backends); + typedef void (*ggml_backend_comm_free_t)(void * comm_ctx); + typedef bool (*ggml_backend_comm_allreduce_tensor_t)(void * comm_ctx, struct ggml_tensor ** tensors); + + // Split buffer type for tensor parallelism (old) + typedef ggml_backend_buffer_type_t (*ggml_backend_split_buffer_type_t)(int main_device, const float * tensor_split); + // Set the number of threads for the backend + typedef void (*ggml_backend_set_n_threads_t)(ggml_backend_t backend, int n_threads); + // Get additional buffer types provided by the device (returns a NULL-terminated array) + typedef ggml_backend_buffer_type_t * (*ggml_backend_dev_get_extra_bufts_t)(ggml_backend_dev_t device); + // Set the abort callback for the backend + typedef void (*ggml_backend_set_abort_callback_t)(ggml_backend_t backend, ggml_abort_callback abort_callback, void * abort_callback_data); + // Get a list of feature flags supported by the backend (returns a NULL-terminated array) + struct ggml_backend_feature { + const char * name; + const char * value; + }; + typedef struct ggml_backend_feature * (*ggml_backend_get_features_t)(ggml_backend_reg_t reg); + + // + // Backend registry + // + + GGML_API void ggml_backend_register(ggml_backend_reg_t reg); + + GGML_API void ggml_backend_device_register(ggml_backend_dev_t device); + + // Backend (reg) enumeration + GGML_API size_t ggml_backend_reg_count(void); + GGML_API ggml_backend_reg_t ggml_backend_reg_get(size_t index); + GGML_API ggml_backend_reg_t ggml_backend_reg_by_name(const char * name); + + // Device enumeration + GGML_API size_t ggml_backend_dev_count(void); + GGML_API ggml_backend_dev_t ggml_backend_dev_get(size_t index); + GGML_API ggml_backend_dev_t ggml_backend_dev_by_name(const char * name); + GGML_API ggml_backend_dev_t ggml_backend_dev_by_type(enum ggml_backend_dev_type type); + + // Direct backend (stream) initialization + // = ggml_backend_dev_init(ggml_backend_dev_by_name(name), params) + GGML_API ggml_backend_t ggml_backend_init_by_name(const char * name, const char * params); + // = ggml_backend_dev_init(ggml_backend_dev_by_type(type), params) + GGML_API ggml_backend_t ggml_backend_init_by_type(enum ggml_backend_dev_type type, const char * params); + // = ggml_backend_dev_init(ggml_backend_dev_by_type(GPU) OR ggml_backend_dev_by_type(CPU), NULL) + GGML_API ggml_backend_t ggml_backend_init_best(void); + + // Load a backend from a dynamic library and register it + GGML_API ggml_backend_reg_t ggml_backend_load(const char * path); + // Unload a backend if loaded dynamically and unregister it + GGML_API void ggml_backend_unload(ggml_backend_reg_t reg); + // Load all known backends from dynamic libraries + GGML_API void ggml_backend_load_all(void); + GGML_API void ggml_backend_load_all_from_path(const char * dir_path); + + // + // Backend scheduler + // + + // The backend scheduler allows for multiple backend devices to be used together + // Handles compute buffer allocation, assignment of tensors to backends, and copying of tensors between backends + // The backends are selected based on: + // - the backend that supports the operation + // - the location of the pre-allocated tensors (e.g. the weights) + /* + Example usage: + + // operations that use tensors allocated in a buffer with USAGE_WEIGHTS will be assigned + // preferably to run on the same backend as the buffer + ggml_backend_buffer_set_usage(buf_weights, GGML_BACKEND_BUFFER_USAGE_WEIGHTS); + + sched = ggml_backend_sched_new({backend_gpu, backend_gpu2, backend_cpu}, NULL, num_backends, GGML_DEFAULT_GRAPH_SIZE, false, true); + + // initialize buffers from a max size graph (optional) + reserve_graph = build_graph(sched, max_batch_size); + + // manually assign nodes to a backend (optional, should not be needed in most cases) + struct ggml_tensor * node = ggml_mul_mat(ctx, ...); + ggml_backend_sched_set_tensor_backend(sched, node, backend_gpu); + + ggml_backend_sched_reserve(sched, reserve_graph); + + // compute + graph = build_graph(sched); // the graph and its tensors are single-use in terms of allocation, multi-use in terms of computation + for (int i = 0; i < 10; ++i) { + ggml_backend_sched_graph_compute(sched, graph); // on the first iteration the graph is allocated automatically + } + + // if there are graph inputs: + graph = build_graph(sched); // get a new graph that is not allocated (the metadata for the old graph is freed once ggml_free is called) + ggml_backend_sched_reset(sched); // clear the allocation of the previous graph + ggml_backend_sched_alloc_graph(sched, graph); // explicitly allocate the new graph but do not execute it + ggml_backend_tensor_set(input_tensor, ...); // copy data to the newly allocated graph tensors + ggml_backend_sched_graph_compute(sched, graph); // execute the graph + + // as an alternative to the above it is also possible to assign the inputs to a dedicated context and + // allocate them statically via ggml_backend_alloc_ctx_tensors + } + */ + + typedef struct ggml_backend_sched * ggml_backend_sched_t; + + // Evaluation callback for each node in the graph (set with ggml_backend_sched_set_eval_callback) + // when ask == true, the scheduler wants to know if the user wants to observe this node + // this allows the scheduler to batch nodes together in order to evaluate them in a single call + // + // when ask == false, the scheduler is passing the node tensor to the user for observation + // if the user returns false, the scheduler will cancel the graph compute + // + typedef bool (*ggml_backend_sched_eval_callback)(struct ggml_tensor * t, bool ask, void * user_data); + + // Callback while copying input weights of a split + // if the user returns false the scheduler simply copies the entire weight + // the callback is called only for input weights in host buffers + // the callback is called after all non-weight inputs of the split have been copied + // `src` is the tensor in the previous split + // `dst` is the copy of `src` in the split + // `graph` is the compute graph of the split + typedef bool (*ggml_backend_sched_copy_callback)(ggml_backend_t backend, const struct ggml_tensor * src, struct ggml_tensor * dst, struct ggml_cgraph * graph, void * user_data); + + // Callback when the buffers reserved for a graph are not sufficient and have to be reserved again + // it happens on the graph allocation path and costs a synchronization of the backends + // `sched` is the scheduler + // `graph` is the compute graph being allocated + // `backend_ids_changed` is true if the backend assignment of the graph nodes changed since the last allocation + // `graph_size_changed` is true if the graph size changed since the last allocation + typedef void (*ggml_backend_sched_realloc_callback)(ggml_backend_sched_t sched, struct ggml_cgraph * graph, bool backend_ids_changed, bool graph_size_changed, void * user_data); + + // Initialize a backend scheduler, backends with low index are given priority over backends with high index + GGML_API ggml_backend_sched_t ggml_backend_sched_new(ggml_backend_t * backends, ggml_backend_buffer_type_t * bufts, int n_backends, size_t graph_size, bool parallel, bool op_offload); + GGML_API void ggml_backend_sched_free(ggml_backend_sched_t sched); + + // Initialize backend buffers from a measure graph + GGML_API void ggml_backend_sched_reserve_size(ggml_backend_sched_t sched, struct ggml_cgraph * measure_graph, size_t * sizes); + GGML_API bool ggml_backend_sched_reserve(ggml_backend_sched_t sched, struct ggml_cgraph * measure_graph); // returns success + + GGML_API int ggml_backend_sched_get_n_backends(ggml_backend_sched_t sched); + GGML_API ggml_backend_t ggml_backend_sched_get_backend(ggml_backend_sched_t sched, int i); + + // Get the number of splits of the last graph + GGML_API int ggml_backend_sched_get_n_splits(ggml_backend_sched_t sched); + GGML_API int ggml_backend_sched_get_n_copies(ggml_backend_sched_t sched); + + GGML_API ggml_backend_buffer_type_t ggml_backend_sched_get_buffer_type(ggml_backend_sched_t sched, ggml_backend_t backend); + GGML_API size_t ggml_backend_sched_get_buffer_size(ggml_backend_sched_t sched, ggml_backend_t backend); + + GGML_API void ggml_backend_sched_set_tensor_backend(ggml_backend_sched_t sched, struct ggml_tensor * node, ggml_backend_t backend); + GGML_API ggml_backend_t ggml_backend_sched_get_tensor_backend(ggml_backend_sched_t sched, struct ggml_tensor * node); + + // Split graph without allocating it + GGML_API void ggml_backend_sched_split_graph(ggml_backend_sched_t sched, struct ggml_cgraph * graph); + + // Allocate and compute graph on the backend scheduler + GGML_API bool ggml_backend_sched_alloc_graph(ggml_backend_sched_t sched, struct ggml_cgraph * graph); // returns success + GGML_API enum ggml_status ggml_backend_sched_graph_compute(ggml_backend_sched_t sched, struct ggml_cgraph * graph); + GGML_API enum ggml_status ggml_backend_sched_graph_compute_async(ggml_backend_sched_t sched, struct ggml_cgraph * graph); + GGML_API void ggml_backend_sched_synchronize(ggml_backend_sched_t sched); + + // Reset all assignments and allocators - must be called before changing the node backends or allocating a new graph. + // This in effect deallocates all tensors that were previously allocated and leaves them with dangling pointers. + // The correct way to use this API is to discard the deallocated tensors and create new ones. + GGML_API void ggml_backend_sched_reset(ggml_backend_sched_t sched); + + // Set a callback to be called for each resulting node during graph compute + GGML_API void ggml_backend_sched_set_eval_callback(ggml_backend_sched_t sched, ggml_backend_sched_eval_callback callback, void * user_data); + + // Set a callback to be called when the inputs weights of a split are being copied + GGML_API void ggml_backend_sched_set_copy_callback(ggml_backend_sched_t sched, ggml_backend_sched_copy_callback callback, void * user_data); + + // Set a callback to be called when the compute buffers of a graph are being reallocated + GGML_API void ggml_backend_sched_set_realloc_callback(ggml_backend_sched_t sched, ggml_backend_sched_realloc_callback callback, void * user_data); + + // + // Meta backend + // + +#define GGML_BACKEND_META_MAX_DEVICES 16 + + enum ggml_backend_meta_split_axis { + // tensor split by tensor dimensions: + GGML_BACKEND_SPLIT_AXIS_0 = 0, + GGML_BACKEND_SPLIT_AXIS_1 = 1, + GGML_BACKEND_SPLIT_AXIS_2 = 2, + GGML_BACKEND_SPLIT_AXIS_3 = 3, + + GGML_BACKEND_SPLIT_AXIS_MIRRORED = 10, // all values on all backends + GGML_BACKEND_SPLIT_AXIS_PARTIAL = 11, // each backend has a partial sum + + // for internal bookkeeping only: + GGML_BACKEND_SPLIT_AXIS_NONE = 98, + GGML_BACKEND_SPLIT_AXIS_UNKNOWN = 99, + }; + GGML_API const char * ggml_backend_meta_split_axis_name(enum ggml_backend_meta_split_axis split_axis); + + struct ggml_backend_meta_split_state { + enum ggml_backend_meta_split_axis axis; + + // for tensors with axis >= 0 && axis < GGML_MAX_DIMS: + // - each device has a slice of the tensor along the split axis + // - most tensors have n_segments == 1 and a contiguous slice of the tensor data + // - some tensors have an inhomogenenous data layout along the split axis, + // those tensors are divided into segments which are each individually split across devices + // (this usually happens when multiple tensors are fused into a single one) + // - ne has one entry per segment and device and that segment repeats nr times, + // in total when accounting for repetitions the segments add up to ggml_tensor::ne for that axis, + // the outer/inner loops are over segments/devices like [seg0_dev0_r0, seg0_dev1_r0, seg0_dev0_r1, seg0_dev1_r1, seg1_dev0_r0, seg1_dev1_r0], + // - for example, a transformer may have a fused QKV matrix rather than 3 matrices, those would be 3 separate segments + // that each need to be split individually across devices so that each device gets a slice of Q, K, and V, + // the Q matrix can be larger than the K and V matrices so this can either be expressed as 3 segments or as 2 segments + // where the segment for K/V repeats twice + int64_t ne[16*GGML_BACKEND_META_MAX_DEVICES]; + uint32_t nr[16]; + uint32_t n_segments; + }; + + // function to assign split states for statically allocated tensors, compute tensor split states will be assigned to be compatible: + typedef struct ggml_backend_meta_split_state(*ggml_backend_meta_get_split_state_t)(const struct ggml_tensor * tensor, void * userdata); + + // create a new meta device from "simple" devices, meta buffer type/buffer/backend is then derived from this: + // TODO: this looks a bit strange - a backend API creates a device. I think we should try + // express this as a backend registry functionality instead + GGML_API ggml_backend_dev_t ggml_backend_meta_device( + ggml_backend_dev_t * devs, size_t n_devs, ggml_backend_meta_get_split_state_t get_split_state, void * get_split_state_ud); + + // + // Utils + // + + struct ggml_backend_graph_copy { + ggml_backend_buffer_t buffer; + struct ggml_context * ctx_allocated; + struct ggml_context * ctx_unallocated; + struct ggml_cgraph * graph; + }; + + // Copy a graph to a different backend + GGML_API struct ggml_backend_graph_copy ggml_backend_graph_copy(ggml_backend_t backend, struct ggml_cgraph * graph); + GGML_API void ggml_backend_graph_copy_free(struct ggml_backend_graph_copy copy); + + typedef bool (*ggml_backend_eval_callback)(int node_index, struct ggml_tensor * t1, struct ggml_tensor * t2, void * user_data); + + // Compare the output of two backends + GGML_API bool ggml_backend_compare_graph_backend(ggml_backend_t backend1, ggml_backend_t backend2, struct ggml_cgraph * graph, ggml_backend_eval_callback callback, void * user_data, struct ggml_tensor const * const * test_nodes, size_t num_test_nodes); + + // Tensor initialization + GGML_API enum ggml_status ggml_backend_tensor_alloc(ggml_backend_buffer_t buffer, struct ggml_tensor * tensor, void * addr); + GGML_API enum ggml_status ggml_backend_view_init(struct ggml_tensor * tensor); + + // CPU buffer types are always available + GGML_API ggml_backend_buffer_t ggml_backend_cpu_buffer_from_ptr(void * ptr, size_t size); + GGML_API ggml_backend_buffer_type_t ggml_backend_cpu_buffer_type(void); + +#ifdef __cplusplus +} +#endif diff --git a/python/freetoken/kernel/csrc/ggml/include/ggml-cpu.h b/python/freetoken/kernel/csrc/ggml/include/ggml-cpu.h new file mode 100644 index 0000000000..dc6453c6ea --- /dev/null +++ b/python/freetoken/kernel/csrc/ggml/include/ggml-cpu.h @@ -0,0 +1,152 @@ +#pragma once + +#include "ggml.h" +#include "ggml-backend.h" + +#ifdef __cplusplus +extern "C" { +#endif + + // the compute plan that needs to be prepared for ggml_graph_compute() + // since https://github.com/ggml-org/ggml/issues/287 + struct ggml_cplan { + size_t work_size; // size of work buffer, calculated by `ggml_graph_plan()` + uint8_t * work_data; // work buffer, to be allocated by caller before calling to `ggml_graph_compute()` + + int n_threads; + struct ggml_threadpool * threadpool; + + // abort ggml_graph_compute when true + ggml_abort_callback abort_callback; + void * abort_callback_data; + + // use only reference implementations + bool use_ref; + }; + + // numa strategies + enum ggml_numa_strategy { + GGML_NUMA_STRATEGY_DISABLED = 0, + GGML_NUMA_STRATEGY_DISTRIBUTE = 1, + GGML_NUMA_STRATEGY_ISOLATE = 2, + GGML_NUMA_STRATEGY_NUMACTL = 3, + GGML_NUMA_STRATEGY_MIRROR = 4, + GGML_NUMA_STRATEGY_COUNT + }; + + GGML_BACKEND_API void ggml_numa_init(enum ggml_numa_strategy numa); // call once for better performance on NUMA systems + GGML_BACKEND_API bool ggml_is_numa(void); // true if init detected that system has >1 NUMA node + + GGML_BACKEND_API struct ggml_tensor * ggml_new_i32(struct ggml_context * ctx, int32_t value); + GGML_BACKEND_API struct ggml_tensor * ggml_new_f32(struct ggml_context * ctx, float value); + + GGML_BACKEND_API struct ggml_tensor * ggml_set_i32 (struct ggml_tensor * tensor, int32_t value); + GGML_BACKEND_API struct ggml_tensor * ggml_set_f32 (struct ggml_tensor * tensor, float value); + + GGML_BACKEND_API int32_t ggml_get_i32_1d(const struct ggml_tensor * tensor, int i); + GGML_BACKEND_API void ggml_set_i32_1d(const struct ggml_tensor * tensor, int i, int32_t value); + + GGML_BACKEND_API int32_t ggml_get_i32_nd(const struct ggml_tensor * tensor, int i0, int i1, int i2, int i3); + GGML_BACKEND_API void ggml_set_i32_nd(const struct ggml_tensor * tensor, int i0, int i1, int i2, int i3, int32_t value); + + GGML_BACKEND_API float ggml_get_f32_1d(const struct ggml_tensor * tensor, int i); + GGML_BACKEND_API void ggml_set_f32_1d(const struct ggml_tensor * tensor, int i, float value); + + GGML_BACKEND_API float ggml_get_f32_nd(const struct ggml_tensor * tensor, int i0, int i1, int i2, int i3); + GGML_BACKEND_API void ggml_set_f32_nd(const struct ggml_tensor * tensor, int i0, int i1, int i2, int i3, float value); + + GGML_BACKEND_API struct ggml_threadpool * ggml_threadpool_new (struct ggml_threadpool_params * params); + GGML_BACKEND_API void ggml_threadpool_free (struct ggml_threadpool * threadpool); + GGML_BACKEND_API int ggml_threadpool_get_n_threads (struct ggml_threadpool * threadpool); + GGML_BACKEND_API void ggml_threadpool_pause (struct ggml_threadpool * threadpool); + GGML_BACKEND_API void ggml_threadpool_resume (struct ggml_threadpool * threadpool); + + // ggml_graph_plan() has to be called before ggml_graph_compute() + // when plan.work_size > 0, caller must allocate memory for plan.work_data + GGML_BACKEND_API struct ggml_cplan ggml_graph_plan( + const struct ggml_cgraph * cgraph, + int n_threads, /* = GGML_DEFAULT_N_THREADS */ + struct ggml_threadpool * threadpool /* = NULL */ ); + GGML_BACKEND_API enum ggml_status ggml_graph_compute(struct ggml_cgraph * cgraph, struct ggml_cplan * cplan); + + // same as ggml_graph_compute() but the work data is allocated as a part of the context + // note: the drawback of this API is that you must have ensured that the context has enough memory for the work data + GGML_BACKEND_API enum ggml_status ggml_graph_compute_with_ctx(struct ggml_context * ctx, struct ggml_cgraph * cgraph, int n_threads); + + // + // system info + // + + // x86 + GGML_BACKEND_API int ggml_cpu_has_sse3 (void); + GGML_BACKEND_API int ggml_cpu_has_ssse3 (void); + GGML_BACKEND_API int ggml_cpu_has_avx (void); + GGML_BACKEND_API int ggml_cpu_has_avx_vnni (void); + GGML_BACKEND_API int ggml_cpu_has_avx2 (void); + GGML_BACKEND_API int ggml_cpu_has_bmi2 (void); + GGML_BACKEND_API int ggml_cpu_has_f16c (void); + GGML_BACKEND_API int ggml_cpu_has_fma (void); + GGML_BACKEND_API int ggml_cpu_has_avx512 (void); + GGML_BACKEND_API int ggml_cpu_has_avx512_vbmi(void); + GGML_BACKEND_API int ggml_cpu_has_avx512_vnni(void); + GGML_BACKEND_API int ggml_cpu_has_avx512_bf16(void); + GGML_BACKEND_API int ggml_cpu_has_amx_int8 (void); + // ARM + GGML_BACKEND_API int ggml_cpu_has_neon (void); + GGML_BACKEND_API int ggml_cpu_has_arm_fma (void); + GGML_BACKEND_API int ggml_cpu_has_fp16_va (void); + GGML_BACKEND_API int ggml_cpu_has_dotprod (void); + GGML_BACKEND_API int ggml_cpu_has_matmul_int8(void); + GGML_BACKEND_API int ggml_cpu_has_sve (void); + GGML_BACKEND_API int ggml_cpu_get_sve_cnt (void); // sve vector length in bytes + GGML_BACKEND_API int ggml_cpu_has_sme (void); + GGML_BACKEND_API int ggml_cpu_has_sme2 (void); + // other + GGML_BACKEND_API int ggml_cpu_has_riscv_v (void); + GGML_BACKEND_API int ggml_cpu_get_rvv_vlen (void); // risc-v vector length in bytes + GGML_BACKEND_API int ggml_cpu_has_vsx (void); + GGML_BACKEND_API int ggml_cpu_has_vxe (void); + GGML_BACKEND_API int ggml_cpu_has_wasm_simd (void); + GGML_BACKEND_API int ggml_cpu_has_llamafile (void); + + // Internal types and functions exposed for tests and benchmarks + + typedef void (*ggml_vec_dot_t) (int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT x, size_t bx, + const void * GGML_RESTRICT y, size_t by, int nrc); + + struct ggml_type_traits_cpu { + ggml_from_float_t from_float; + ggml_vec_dot_t vec_dot; + enum ggml_type vec_dot_type; + int64_t nrows; // number of rows to process simultaneously + }; + + GGML_BACKEND_API const struct ggml_type_traits_cpu * ggml_get_type_traits_cpu(enum ggml_type type); + + GGML_BACKEND_API void ggml_cpu_init(void); + + // + // CPU backend + // + + GGML_BACKEND_API ggml_backend_t ggml_backend_cpu_init(void); + + GGML_BACKEND_API bool ggml_backend_is_cpu (ggml_backend_t backend); + GGML_BACKEND_API void ggml_backend_cpu_set_n_threads (ggml_backend_t backend_cpu, int n_threads); + GGML_BACKEND_API void ggml_backend_cpu_set_threadpool (ggml_backend_t backend_cpu, ggml_threadpool_t threadpool); + GGML_BACKEND_API void ggml_backend_cpu_set_abort_callback(ggml_backend_t backend_cpu, ggml_abort_callback abort_callback, void * abort_callback_data); + + GGML_BACKEND_API void ggml_backend_cpu_set_use_ref(ggml_backend_t backend_cpu, bool use_ref); + + GGML_BACKEND_API ggml_backend_reg_t ggml_backend_cpu_reg(void); + + GGML_BACKEND_API void ggml_cpu_fp32_to_fp32(const float *, float *, int64_t); + GGML_BACKEND_API void ggml_cpu_fp32_to_i32 (const float *, int32_t *, int64_t); + GGML_BACKEND_API void ggml_cpu_fp32_to_fp16(const float *, ggml_fp16_t *, int64_t); + GGML_BACKEND_API void ggml_cpu_fp16_to_fp32(const ggml_fp16_t *, float *, int64_t); + GGML_BACKEND_API void ggml_cpu_fp32_to_bf16(const float *, ggml_bf16_t *, int64_t); + GGML_BACKEND_API void ggml_cpu_bf16_to_fp32(const ggml_bf16_t *, float *, int64_t); + +#ifdef __cplusplus +} +#endif diff --git a/python/freetoken/kernel/csrc/ggml/include/ggml-cuda.h b/python/freetoken/kernel/csrc/ggml/include/ggml-cuda.h new file mode 100644 index 0000000000..1cd81eeaeb --- /dev/null +++ b/python/freetoken/kernel/csrc/ggml/include/ggml-cuda.h @@ -0,0 +1,47 @@ +#pragma once + +#include "ggml.h" +#include "ggml-backend.h" + +#ifdef __cplusplus +extern "C" { +#endif + +#ifdef GGML_USE_HIP +#define GGML_CUDA_NAME "ROCm" +#define GGML_CUBLAS_NAME "hipBLAS" +#elif defined(GGML_USE_MUSA) +#define GGML_CUDA_NAME "MUSA" +#define GGML_CUBLAS_NAME "muBLAS" +#else +#define GGML_CUDA_NAME "CUDA" +#define GGML_CUBLAS_NAME "cuBLAS" +#endif +#define GGML_CUDA_MAX_DEVICES 16 + +// backend API +GGML_BACKEND_API ggml_backend_t ggml_backend_cuda_init(int device); + +GGML_BACKEND_API bool ggml_backend_is_cuda(ggml_backend_t backend); + +// device buffer +GGML_BACKEND_API ggml_backend_buffer_type_t ggml_backend_cuda_buffer_type(int device); + +// conduct allreduce operation between devices +GGML_BACKEND_API bool ggml_backend_cuda_allreduce_tensor(ggml_backend_t * backends, struct ggml_tensor ** tensors, size_t n_backends); + +// pinned host buffer for use with the CPU backend for faster copies between CPU and GPU +GGML_BACKEND_API ggml_backend_buffer_type_t ggml_backend_cuda_host_buffer_type(void); + +GGML_BACKEND_API int ggml_backend_cuda_get_device_count(void); +GGML_BACKEND_API void ggml_backend_cuda_get_device_description(int device, char * description, size_t description_size); +GGML_BACKEND_API void ggml_backend_cuda_get_device_memory(int device, size_t * free, size_t * total); + +GGML_BACKEND_API bool ggml_backend_cuda_register_host_buffer(void * buffer, size_t size); +GGML_BACKEND_API void ggml_backend_cuda_unregister_host_buffer(void * buffer); + +GGML_BACKEND_API ggml_backend_reg_t ggml_backend_cuda_reg(void); + +#ifdef __cplusplus +} +#endif diff --git a/python/freetoken/kernel/csrc/ggml/include/ggml.h b/python/freetoken/kernel/csrc/ggml/include/ggml.h new file mode 100644 index 0000000000..23ea6437e2 --- /dev/null +++ b/python/freetoken/kernel/csrc/ggml/include/ggml.h @@ -0,0 +1,3027 @@ +#pragma once + +// +// GGML Tensor Library +// +// This documentation is still a work in progress. +// If you wish some specific topics to be covered, feel free to drop a comment: +// +// https://github.com/ggml-org/whisper.cpp/issues/40 +// +// ## Overview +// +// This library implements: +// +// - a set of tensor operations +// - automatic differentiation +// - basic optimization algorithms +// +// The aim of this library is to provide a minimalistic approach for various machine learning tasks. This includes, +// but is not limited to, the following: +// +// - linear regression +// - support vector machines +// - neural networks +// +// The library allows the user to define a certain function using the available tensor operations. This function +// definition is represented internally via a computation graph. Each tensor operation in the function definition +// corresponds to a node in the graph. Having the computation graph defined, the user can choose to compute the +// function's value and/or its gradient with respect to the input variables. Optionally, the function can be optimized +// using one of the available optimization algorithms. +// +// For example, here we define the function: f(x) = a*x^2 + b +// +// { +// struct ggml_init_params params = { +// .mem_size = 16*1024*1024, +// .mem_buffer = NULL, +// }; +// +// // memory allocation happens here +// struct ggml_context * ctx = ggml_init(params); +// +// struct ggml_tensor * x = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 1); +// +// ggml_set_param(ctx, x); // x is an input variable +// +// struct ggml_tensor * a = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 1); +// struct ggml_tensor * b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 1); +// struct ggml_tensor * x2 = ggml_mul(ctx, x, x); +// struct ggml_tensor * f = ggml_add(ctx, ggml_mul(ctx, a, x2), b); +// +// ... +// } +// +// Notice that the function definition above does not involve any actual computation. The computation is performed only +// when the user explicitly requests it. For example, to compute the function's value at x = 2.0: +// +// { +// ... +// +// struct ggml_cgraph * gf = ggml_new_graph(ctx); +// ggml_build_forward_expand(gf, f); +// +// // set the input variable and parameter values +// ggml_set_f32(x, 2.0f); +// ggml_set_f32(a, 3.0f); +// ggml_set_f32(b, 4.0f); +// +// ggml_graph_compute_with_ctx(ctx, &gf, n_threads); +// +// printf("f = %f\n", ggml_get_f32_1d(f, 0)); +// +// ... +// } +// +// The actual computation is performed in the ggml_graph_compute() function. +// +// The ggml_new_tensor_...() functions create new tensors. They are allocated in the memory buffer provided to the +// ggml_init() function. You have to be careful not to exceed the memory buffer size. Therefore, you have to know +// in advance how much memory you need for your computation. Alternatively, you can allocate a large enough memory +// and after defining the computation graph, call the ggml_used_mem() function to find out how much memory was +// actually needed. +// +// The ggml_set_param() function marks a tensor as an input variable. This is used by the automatic +// differentiation and optimization algorithms. +// +// The described approach allows to define the function graph once and then compute its forward or backward graphs +// multiple times. All computations will use the same memory buffer allocated in the ggml_init() function. This way +// the user can avoid the memory allocation overhead at runtime. +// +// The library supports multi-dimensional tensors - up to 4 dimensions. The FP16 and FP32 data types are first class +// citizens, but in theory the library can be extended to support FP8 and integer data types. +// +// Each tensor operation produces a new tensor. Initially the library was envisioned to support only the use of unary +// and binary operations. Most of the available operations fall into one of these two categories. With time, it became +// clear that the library needs to support more complex operations. The way to support these operations is not clear +// yet, but a few examples are demonstrated in the following operations: +// +// - ggml_permute() +// - ggml_conv_1d_1s() +// - ggml_conv_1d_2s() +// +// For each tensor operator, the library implements a forward and backward computation function. The forward function +// computes the output tensor value given the input tensor values. The backward function computes the adjoint of the +// input tensors given the adjoint of the output tensor. For a detailed explanation of what this means, take a +// calculus class, or watch the following video: +// +// What is Automatic Differentiation? +// https://www.youtube.com/watch?v=wG_nF1awSSY +// +// +// ## Tensor data (struct ggml_tensor) +// +// The tensors are stored in memory via the ggml_tensor struct. The structure provides information about the size of +// the tensor, the data type, and the memory buffer where the tensor data is stored. Additionally, it contains +// pointers to the "source" tensors - i.e. the tensors that were used to compute the current tensor. For example: +// +// { +// struct ggml_tensor * c = ggml_add(ctx, a, b); +// +// assert(c->src[0] == a); +// assert(c->src[1] == b); +// } +// +// The multi-dimensional tensors are stored in row-major order. The ggml_tensor struct contains fields for the +// number of elements in each dimension ("ne") as well as the number of bytes ("nb", a.k.a. stride). This allows +// to store tensors that are not contiguous in memory, which is useful for operations such as transposition and +// permutation. All tensor operations have to take the stride into account and not assume that the tensor is +// contiguous in memory. +// +// The data of the tensor is accessed via the "data" pointer. For example: +// +// { +// const int nx = 2; +// const int ny = 3; +// +// struct ggml_tensor * a = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, nx, ny); +// +// for (int y = 0; y < ny; y++) { +// for (int x = 0; x < nx; x++) { +// *(float *) ((char *) a->data + y*a->nb[1] + x*a->nb[0]) = x + y; +// } +// } +// +// ... +// } +// +// Alternatively, there are helper functions, such as ggml_get_f32_1d() and ggml_set_f32_1d() that can be used. +// +// ## The matrix multiplication operator (ggml_mul_mat) +// +// TODO +// +// +// ## Multi-threading +// +// TODO +// +// +// ## Overview of ggml.c +// +// TODO +// +// +// ## SIMD optimizations +// +// TODO +// +// +// ## Debugging ggml +// +// TODO +// +// + +#ifdef GGML_SHARED +# if defined(_WIN32) && !defined(__MINGW32__) +# ifdef GGML_BUILD +# define GGML_API __declspec(dllexport) extern +# else +# define GGML_API __declspec(dllimport) extern +# endif +# else +# define GGML_API __attribute__ ((visibility ("default"))) extern +# endif +#else +# define GGML_API extern +#endif + +// TODO: support for clang +#ifdef __GNUC__ +# define GGML_DEPRECATED(func, hint) func __attribute__((deprecated(hint))) +#elif defined(_MSC_VER) +# define GGML_DEPRECATED(func, hint) __declspec(deprecated(hint)) func +#else +# define GGML_DEPRECATED(func, hint) func +#endif + +#ifndef __GNUC__ +# define GGML_ATTRIBUTE_FORMAT(...) +#elif defined(__MINGW32__) && !defined(__clang__) +# define GGML_ATTRIBUTE_FORMAT(...) __attribute__((format(gnu_printf, __VA_ARGS__))) +#else +# define GGML_ATTRIBUTE_FORMAT(...) __attribute__((format(printf, __VA_ARGS__))) +#endif + +#if defined(_WIN32) && !defined(_WIN32_WINNT) +# define _WIN32_WINNT 0x0A00 +#endif + +#include +#include +#include +#include + +#define GGML_FILE_MAGIC 0x67676d6c // "ggml" +#define GGML_FILE_VERSION 2 + +#define GGML_QNT_VERSION 2 // bump this on quantization format changes +#define GGML_QNT_VERSION_FACTOR 1000 // do not change this + +#define GGML_MAX_DIMS 4 +#define GGML_MAX_PARAMS 2048 +#define GGML_MAX_SRC 10 +#define GGML_MAX_N_THREADS 512 +#define GGML_MAX_OP_PARAMS 64 + +#ifndef GGML_MAX_NAME +# define GGML_MAX_NAME 64 +#endif + +#define GGML_DEFAULT_N_THREADS 4 +#define GGML_DEFAULT_GRAPH_SIZE 2048 + +#if UINTPTR_MAX == 0xFFFFFFFF + #define GGML_MEM_ALIGN 4 +#elif defined(__EMSCRIPTEN__) +// emscripten uses max_align_t == 8, so we need GGML_MEM_ALIGN == 8 for 64-bit wasm. +// (for 32-bit wasm, the first conditional is true and GGML_MEM_ALIGN stays 4.) +// ref: https://github.com/ggml-org/llama.cpp/pull/18628 + #define GGML_MEM_ALIGN 8 +#else + #define GGML_MEM_ALIGN 16 +#endif + +#define GGML_EXIT_SUCCESS 0 +#define GGML_EXIT_ABORTED 1 + +// TODO: convert to enum https://github.com/ggml-org/llama.cpp/pull/16187#discussion_r2388538726 +#define GGML_ROPE_TYPE_NORMAL 0 +#define GGML_ROPE_TYPE_NEOX 2 +#define GGML_ROPE_TYPE_MROPE 8 +#define GGML_ROPE_TYPE_VISION 24 +#define GGML_ROPE_TYPE_IMROPE 40 // binary: 101000 + +#define GGML_MROPE_SECTIONS 4 + +#define GGML_UNUSED(x) (void)(x) +#ifdef __CUDACC__ +template +__host__ __device__ constexpr inline void ggml_unused_vars_impl(Args&&...) noexcept {} +#define GGML_UNUSED_VARS(...) ggml_unused_vars_impl(__VA_ARGS__) +#else +#define GGML_UNUSED_VARS(...) do { (void)sizeof((__VA_ARGS__, 0)); } while(0) +#endif // __CUDACC__ + +#define GGML_PAD(x, n) (((x) + (n) - 1) & ~((n) - 1)) + +#ifndef NDEBUG +# define GGML_UNREACHABLE() do { fprintf(stderr, "statement should be unreachable\n"); abort(); } while(0) +#elif defined(__GNUC__) +# define GGML_UNREACHABLE() __builtin_unreachable() +#elif defined(_MSC_VER) +# define GGML_UNREACHABLE() __assume(0) +#else +# define GGML_UNREACHABLE() ((void) 0) +#endif + +#ifdef __cplusplus +# define GGML_NORETURN [[noreturn]] +#elif defined(_MSC_VER) +# define GGML_NORETURN __declspec(noreturn) +#else +# define GGML_NORETURN _Noreturn +#endif + +#define GGML_ABORT(...) ggml_abort(__FILE__, __LINE__, __VA_ARGS__) +#define GGML_ASSERT(x) if (!(x)) GGML_ABORT("GGML_ASSERT(%s) failed", #x) + +// used to copy the number of elements and stride in bytes of tensors into local variables. +// main purpose is to reduce code duplication and improve readability. +// +// example: +// +// GGML_TENSOR_LOCALS(int64_t, ne1, src1, ne); +// GGML_TENSOR_LOCALS(size_t, nb1, src1, nb); +// +#define GGML_TENSOR_LOCALS_1(type, prefix, pointer, array) \ + const type prefix##0 = (pointer) ? (pointer)->array[0] : 0; \ + GGML_UNUSED(prefix##0); +#define GGML_TENSOR_LOCALS_2(type, prefix, pointer, array) \ + GGML_TENSOR_LOCALS_1 (type, prefix, pointer, array) \ + const type prefix##1 = (pointer) ? (pointer)->array[1] : 0; \ + GGML_UNUSED(prefix##1); +#define GGML_TENSOR_LOCALS_3(type, prefix, pointer, array) \ + GGML_TENSOR_LOCALS_2 (type, prefix, pointer, array) \ + const type prefix##2 = (pointer) ? (pointer)->array[2] : 0; \ + GGML_UNUSED(prefix##2); +#define GGML_TENSOR_LOCALS(type, prefix, pointer, array) \ + GGML_TENSOR_LOCALS_3 (type, prefix, pointer, array) \ + const type prefix##3 = (pointer) ? (pointer)->array[3] : 0; \ + GGML_UNUSED(prefix##3); + +#define GGML_TENSOR_UNARY_OP_LOCALS \ + GGML_TENSOR_LOCALS(int64_t, ne0, src0, ne) \ + GGML_TENSOR_LOCALS(size_t, nb0, src0, nb) \ + GGML_TENSOR_LOCALS(int64_t, ne, dst, ne) \ + GGML_TENSOR_LOCALS(size_t, nb, dst, nb) + +#define GGML_TENSOR_BINARY_OP_LOCALS \ + GGML_TENSOR_LOCALS(int64_t, ne0, src0, ne) \ + GGML_TENSOR_LOCALS(size_t, nb0, src0, nb) \ + GGML_TENSOR_LOCALS(int64_t, ne1, src1, ne) \ + GGML_TENSOR_LOCALS(size_t, nb1, src1, nb) \ + GGML_TENSOR_LOCALS(int64_t, ne, dst, ne) \ + GGML_TENSOR_LOCALS(size_t, nb, dst, nb) + +#define GGML_TENSOR_TERNARY_OP_LOCALS \ + GGML_TENSOR_LOCALS(int64_t, ne0, src0, ne) \ + GGML_TENSOR_LOCALS(size_t, nb0, src0, nb) \ + GGML_TENSOR_LOCALS(int64_t, ne1, src1, ne) \ + GGML_TENSOR_LOCALS(size_t, nb1, src1, nb) \ + GGML_TENSOR_LOCALS(int64_t, ne2, src2, ne) \ + GGML_TENSOR_LOCALS(size_t, nb2, src2, nb) \ + GGML_TENSOR_LOCALS(int64_t, ne, dst, ne) \ + GGML_TENSOR_LOCALS(size_t, nb, dst, nb) + +#define GGML_TENSOR_BINARY_OP_LOCALS01 \ + GGML_TENSOR_LOCALS(int64_t, ne0, src0, ne) \ + GGML_TENSOR_LOCALS(size_t, nb0, src0, nb) \ + GGML_TENSOR_LOCALS(int64_t, ne1, src1, ne) \ + GGML_TENSOR_LOCALS(size_t, nb1, src1, nb) + +#ifdef __cplusplus +extern "C" { +#endif + + // Function type used in fatal error callbacks + typedef void (*ggml_abort_callback_t)(const char * error_message); + + // Set the abort callback (passing null will restore original abort functionality: printing a message to stdout) + // Returns the old callback for chaining + GGML_API ggml_abort_callback_t ggml_set_abort_callback(ggml_abort_callback_t callback); + + GGML_NORETURN GGML_ATTRIBUTE_FORMAT(3, 4) + GGML_API void ggml_abort(const char * file, int line, const char * fmt, ...); + + enum ggml_status { + GGML_STATUS_ALLOC_FAILED = -2, + GGML_STATUS_FAILED = -1, + GGML_STATUS_SUCCESS = 0, + GGML_STATUS_ABORTED = 1, + }; + + // get ggml_status name string + GGML_API const char * ggml_status_to_string(enum ggml_status status); + + // ieee 754-2008 half-precision float16 + // todo: make this not an integral type + typedef uint16_t ggml_fp16_t; + GGML_API float ggml_fp16_to_fp32(ggml_fp16_t); + GGML_API ggml_fp16_t ggml_fp32_to_fp16(float); + GGML_API void ggml_fp16_to_fp32_row(const ggml_fp16_t *, float *, int64_t); + GGML_API void ggml_fp32_to_fp16_row(const float *, ggml_fp16_t *, int64_t); + + // google brain half-precision bfloat16 + typedef struct { uint16_t bits; } ggml_bf16_t; + GGML_API ggml_bf16_t ggml_fp32_to_bf16(float); + GGML_API float ggml_bf16_to_fp32(ggml_bf16_t); // consider just doing << 16 + GGML_API void ggml_bf16_to_fp32_row(const ggml_bf16_t *, float *, int64_t); + GGML_API void ggml_fp32_to_bf16_row_ref(const float *, ggml_bf16_t *, int64_t); + GGML_API void ggml_fp32_to_bf16_row(const float *, ggml_bf16_t *, int64_t); + + struct ggml_object; + struct ggml_context; + struct ggml_cgraph; + + // NOTE: always add types at the end of the enum to keep backward compatibility + enum ggml_type { + GGML_TYPE_F32 = 0, + GGML_TYPE_F16 = 1, + GGML_TYPE_Q4_0 = 2, + GGML_TYPE_Q4_1 = 3, + // GGML_TYPE_Q4_2 = 4, support has been removed + // GGML_TYPE_Q4_3 = 5, support has been removed + GGML_TYPE_Q5_0 = 6, + GGML_TYPE_Q5_1 = 7, + GGML_TYPE_Q8_0 = 8, + GGML_TYPE_Q8_1 = 9, + GGML_TYPE_Q2_K = 10, + GGML_TYPE_Q3_K = 11, + GGML_TYPE_Q4_K = 12, + GGML_TYPE_Q5_K = 13, + GGML_TYPE_Q6_K = 14, + GGML_TYPE_Q8_K = 15, + GGML_TYPE_IQ2_XXS = 16, + GGML_TYPE_IQ2_XS = 17, + GGML_TYPE_IQ3_XXS = 18, + GGML_TYPE_IQ1_S = 19, + GGML_TYPE_IQ4_NL = 20, + GGML_TYPE_IQ3_S = 21, + GGML_TYPE_IQ2_S = 22, + GGML_TYPE_IQ4_XS = 23, + GGML_TYPE_I8 = 24, + GGML_TYPE_I16 = 25, + GGML_TYPE_I32 = 26, + GGML_TYPE_I64 = 27, + GGML_TYPE_F64 = 28, + GGML_TYPE_IQ1_M = 29, + GGML_TYPE_BF16 = 30, + // GGML_TYPE_Q4_0_4_4 = 31, support has been removed from gguf files + // GGML_TYPE_Q4_0_4_8 = 32, + // GGML_TYPE_Q4_0_8_8 = 33, + GGML_TYPE_TQ1_0 = 34, + GGML_TYPE_TQ2_0 = 35, + // GGML_TYPE_IQ4_NL_4_4 = 36, + // GGML_TYPE_IQ4_NL_4_8 = 37, + // GGML_TYPE_IQ4_NL_8_8 = 38, + GGML_TYPE_MXFP4 = 39, // MXFP4 (1 block) + GGML_TYPE_NVFP4 = 40, // NVFP4 (4 blocks, E4M3 scale) + GGML_TYPE_Q1_0 = 41, + GGML_TYPE_Q2_0 = 42, + GGML_TYPE_COUNT = 43, + }; + + // [TAG_GGML_PREC] + // this enum is used to declare the allowed numerical precision/data-types types that can be used during the compute of an op + // the declared types can be: + // - result accumulation type + // - source tensor data representation type + // - etc. + // the precision parameters are stored as ggml_tensor.op_params to the respective ops + enum ggml_prec { + GGML_PREC_UNDEFINED = 0, + GGML_PREC_DEFAULT = 0, // note: deprecated, use GGML_PREC_UNDEFINED + GGML_PREC_F32 = 10, + GGML_PREC_BF16 = 15, + GGML_PREC_F16 = 20, + GGML_PREC_Q8 = 30, + GGML_PREC_Q4 = 40, + }; + + // op hint + enum ggml_op_hint { + GGML_HINT_NONE = 0, + GGML_HINT_SRC0_IS_HADAMARD = 1, + }; + + // model file types + enum ggml_ftype { + GGML_FTYPE_UNKNOWN = -1, + GGML_FTYPE_ALL_F32 = 0, + GGML_FTYPE_MOSTLY_F16 = 1, // except 1d tensors + GGML_FTYPE_MOSTLY_Q4_0 = 2, // except 1d tensors + GGML_FTYPE_MOSTLY_Q4_1 = 3, // except 1d tensors + GGML_FTYPE_MOSTLY_Q4_1_SOME_F16 = 4, // tok_embeddings.weight and output.weight are F16 + GGML_FTYPE_MOSTLY_Q8_0 = 7, // except 1d tensors + GGML_FTYPE_MOSTLY_Q5_0 = 8, // except 1d tensors + GGML_FTYPE_MOSTLY_Q5_1 = 9, // except 1d tensors + GGML_FTYPE_MOSTLY_Q2_K = 10, // except 1d tensors + GGML_FTYPE_MOSTLY_Q3_K = 11, // except 1d tensors + GGML_FTYPE_MOSTLY_Q4_K = 12, // except 1d tensors + GGML_FTYPE_MOSTLY_Q5_K = 13, // except 1d tensors + GGML_FTYPE_MOSTLY_Q6_K = 14, // except 1d tensors + GGML_FTYPE_MOSTLY_IQ2_XXS = 15, // except 1d tensors + GGML_FTYPE_MOSTLY_IQ2_XS = 16, // except 1d tensors + GGML_FTYPE_MOSTLY_IQ3_XXS = 17, // except 1d tensors + GGML_FTYPE_MOSTLY_IQ1_S = 18, // except 1d tensors + GGML_FTYPE_MOSTLY_IQ4_NL = 19, // except 1d tensors + GGML_FTYPE_MOSTLY_IQ3_S = 20, // except 1d tensors + GGML_FTYPE_MOSTLY_IQ2_S = 21, // except 1d tensors + GGML_FTYPE_MOSTLY_IQ4_XS = 22, // except 1d tensors + GGML_FTYPE_MOSTLY_IQ1_M = 23, // except 1d tensors + GGML_FTYPE_MOSTLY_BF16 = 24, // except 1d tensors + GGML_FTYPE_MOSTLY_MXFP4 = 25, // except 1d tensors + GGML_FTYPE_MOSTLY_NVFP4 = 26, // except 1d tensors + GGML_FTYPE_MOSTLY_Q1_0 = 27, // except 1d tensors + GGML_FTYPE_MOSTLY_Q2_0 = 28, // except 1d tensors + }; + + // available tensor operations: + enum ggml_op { + GGML_OP_NONE = 0, + + GGML_OP_DUP, + GGML_OP_ADD, + GGML_OP_ADD_ID, + GGML_OP_ADD1, + GGML_OP_ACC, + GGML_OP_SUB, + GGML_OP_MUL, + GGML_OP_DIV, + GGML_OP_SQR, + GGML_OP_SQRT, + GGML_OP_LOG, + GGML_OP_SIN, + GGML_OP_COS, + GGML_OP_SUM, + GGML_OP_SUM_ROWS, + GGML_OP_CUMSUM, + GGML_OP_MEAN, + GGML_OP_ARGMAX, + GGML_OP_COUNT_EQUAL, + GGML_OP_REPEAT, + GGML_OP_REPEAT_BACK, + GGML_OP_CONCAT, + GGML_OP_SILU_BACK, + GGML_OP_NORM, // normalize + GGML_OP_RMS_NORM, + GGML_OP_RMS_NORM_BACK, + GGML_OP_GROUP_NORM, + GGML_OP_L2_NORM, + + GGML_OP_MUL_MAT, + GGML_OP_MUL_MAT_ID, + GGML_OP_OUT_PROD, + + GGML_OP_SCALE, + GGML_OP_SET, + GGML_OP_CPY, + GGML_OP_CONT, + GGML_OP_RESHAPE, + GGML_OP_VIEW, + GGML_OP_PERMUTE, + GGML_OP_TRANSPOSE, + GGML_OP_GET_ROWS, + GGML_OP_GET_ROWS_BACK, + GGML_OP_SET_ROWS, + GGML_OP_DIAG, + GGML_OP_DIAG_MASK_INF, + GGML_OP_DIAG_MASK_ZERO, + GGML_OP_SOFT_MAX, + GGML_OP_SOFT_MAX_BACK, + GGML_OP_ROPE, + GGML_OP_ROPE_BACK, + GGML_OP_CLAMP, + GGML_OP_CONV_TRANSPOSE_1D, + GGML_OP_IM2COL, + GGML_OP_IM2COL_BACK, + GGML_OP_IM2COL_3D, + GGML_OP_COL2IM_1D, + GGML_OP_CONV_2D, + GGML_OP_CONV_3D, + GGML_OP_CONV_2D_DW, + GGML_OP_CONV_TRANSPOSE_2D, + GGML_OP_POOL_1D, + GGML_OP_POOL_2D, + GGML_OP_POOL_2D_BACK, + GGML_OP_UPSCALE, + GGML_OP_PAD, + GGML_OP_PAD_REFLECT_1D, + GGML_OP_ROLL, + GGML_OP_ARANGE, + GGML_OP_TIMESTEP_EMBEDDING, + GGML_OP_ARGSORT, + GGML_OP_TOP_K, + GGML_OP_LEAKY_RELU, + GGML_OP_TRI, + GGML_OP_FILL, + + GGML_OP_FLASH_ATTN_EXT, + GGML_OP_FLASH_ATTN_BACK, + GGML_OP_SSM_CONV, + GGML_OP_SSM_SCAN, + GGML_OP_WIN_PART, + GGML_OP_WIN_UNPART, + GGML_OP_GET_REL_POS, + GGML_OP_ADD_REL_POS, + GGML_OP_RWKV_WKV6, + GGML_OP_GATED_LINEAR_ATTN, + GGML_OP_RWKV_WKV7, + GGML_OP_SOLVE_TRI, + GGML_OP_GATED_DELTA_NET, + GGML_OP_LIGHTNING_INDEXER, + GGML_OP_DSV4_HC_COMB, + GGML_OP_DSV4_HC_PRE, + GGML_OP_DSV4_HC_POST, + + GGML_OP_UNARY, + + GGML_OP_MAP_CUSTOM1, + GGML_OP_MAP_CUSTOM2, + GGML_OP_MAP_CUSTOM3, + + GGML_OP_CUSTOM, + + GGML_OP_CROSS_ENTROPY_LOSS, + GGML_OP_CROSS_ENTROPY_LOSS_BACK, + GGML_OP_OPT_STEP_ADAMW, + GGML_OP_OPT_STEP_SGD, + + GGML_OP_GLU, + + GGML_OP_COUNT, + }; + + enum ggml_unary_op { + GGML_UNARY_OP_ABS, + GGML_UNARY_OP_SGN, + GGML_UNARY_OP_NEG, + GGML_UNARY_OP_STEP, + GGML_UNARY_OP_TANH, + GGML_UNARY_OP_ELU, + GGML_UNARY_OP_RELU, + GGML_UNARY_OP_SIGMOID, + GGML_UNARY_OP_GELU, + GGML_UNARY_OP_GELU_QUICK, + GGML_UNARY_OP_SILU, + GGML_UNARY_OP_HARDSWISH, + GGML_UNARY_OP_HARDSIGMOID, + GGML_UNARY_OP_EXP, + GGML_UNARY_OP_EXPM1, + GGML_UNARY_OP_SOFTPLUS, + GGML_UNARY_OP_GELU_ERF, + GGML_UNARY_OP_XIELU, + GGML_UNARY_OP_FLOOR, + GGML_UNARY_OP_CEIL, + GGML_UNARY_OP_ROUND, + GGML_UNARY_OP_TRUNC, + + GGML_UNARY_OP_COUNT, + }; + + enum ggml_glu_op { + GGML_GLU_OP_REGLU, + GGML_GLU_OP_GEGLU, + GGML_GLU_OP_SWIGLU, + GGML_GLU_OP_SWIGLU_OAI, + GGML_GLU_OP_GEGLU_ERF, + GGML_GLU_OP_GEGLU_QUICK, + GGML_GLU_OP_SWIGLU_CLAMP, + + GGML_GLU_OP_COUNT, + }; + + enum ggml_object_type { + GGML_OBJECT_TYPE_TENSOR, + GGML_OBJECT_TYPE_GRAPH, + GGML_OBJECT_TYPE_WORK_BUFFER + }; + + enum ggml_log_level { + GGML_LOG_LEVEL_NONE = 0, + GGML_LOG_LEVEL_DEBUG = 1, + GGML_LOG_LEVEL_INFO = 2, + GGML_LOG_LEVEL_WARN = 3, + GGML_LOG_LEVEL_ERROR = 4, + GGML_LOG_LEVEL_CONT = 5, // continue previous log + }; + + // this tensor... + enum ggml_tensor_flag { + GGML_TENSOR_FLAG_INPUT = 1, // ...is an input for the GGML compute graph + GGML_TENSOR_FLAG_OUTPUT = 2, // ...is an output for the GGML compute graph + GGML_TENSOR_FLAG_PARAM = 4, // ...contains trainable parameters + GGML_TENSOR_FLAG_LOSS = 8, // ...defines loss for numerical optimization (multiple loss tensors add up) + GGML_TENSOR_FLAG_COMPUTE = 16, // ...must be computed + }; + + enum ggml_tri_type { + GGML_TRI_TYPE_UPPER_DIAG = 0, + GGML_TRI_TYPE_UPPER = 1, + GGML_TRI_TYPE_LOWER_DIAG = 2, + GGML_TRI_TYPE_LOWER = 3 + }; + + struct ggml_init_params { + // memory pool + size_t mem_size; // bytes + void * mem_buffer; // if NULL, memory will be allocated internally + bool no_alloc; // don't allocate memory for the tensor data + }; + + // n-dimensional tensor + struct ggml_tensor { + enum ggml_type type; + + struct ggml_backend_buffer * buffer; + + int64_t ne[GGML_MAX_DIMS]; // number of elements + size_t nb[GGML_MAX_DIMS]; // stride in bytes: + // nb[0] = ggml_type_size(type) + // nb[1] = nb[0] * (ne[0] / ggml_blck_size(type)) + padding + // nb[i] = nb[i-1] * ne[i-1] + + // compute data + enum ggml_op op; + + // op params - allocated as int32_t for alignment + int32_t op_params[GGML_MAX_OP_PARAMS / sizeof(int32_t)]; + + int32_t flags; + + struct ggml_tensor * src[GGML_MAX_SRC]; + + // source tensor and offset for views + struct ggml_tensor * view_src; + size_t view_offs; + + void * data; + + char name[GGML_MAX_NAME]; + + void * extra; // extra things e.g. for ggml-cuda.cu + + char padding[8]; + }; + + static const size_t GGML_TENSOR_SIZE = sizeof(struct ggml_tensor); + + // Abort callback + // If not NULL, called before ggml computation + // If it returns true, the computation is aborted + typedef bool (*ggml_abort_callback)(void * data); + + + // + // GUID + // + + // GUID types + typedef uint8_t ggml_guid[16]; + typedef ggml_guid * ggml_guid_t; + + GGML_API bool ggml_guid_matches(ggml_guid_t guid_a, ggml_guid_t guid_b); + + // misc + + GGML_API const char * ggml_version(void); + GGML_API const char * ggml_commit(void); + + GGML_API void ggml_time_init(void); // call this once at the beginning of the program + GGML_API int64_t ggml_time_ms(void); + GGML_API int64_t ggml_time_us(void); + GGML_API int64_t ggml_cycles(void); + GGML_API int64_t ggml_cycles_per_ms(void); + + // accepts a UTF-8 path, even on Windows + GGML_API FILE * ggml_fopen(const char * fname, const char * mode); + + GGML_API void ggml_print_object (const struct ggml_object * obj); + GGML_API void ggml_print_objects(const struct ggml_context * ctx); + + GGML_API int64_t ggml_nelements (const struct ggml_tensor * tensor); + GGML_API int64_t ggml_nrows (const struct ggml_tensor * tensor); + GGML_API size_t ggml_nbytes (const struct ggml_tensor * tensor); + GGML_API size_t ggml_nbytes_pad(const struct ggml_tensor * tensor); // same as ggml_nbytes() but padded to GGML_MEM_ALIGN + + GGML_API int64_t ggml_blck_size(enum ggml_type type); + GGML_API size_t ggml_type_size(enum ggml_type type); // size in bytes for all elements in a block + GGML_API size_t ggml_row_size (enum ggml_type type, int64_t ne); // size in bytes for all elements in a row + + GGML_DEPRECATED( + GGML_API double ggml_type_sizef(enum ggml_type type), // ggml_type_size()/ggml_blck_size() as float + "use ggml_row_size() instead"); + + GGML_API const char * ggml_type_name(enum ggml_type type); + GGML_API const char * ggml_op_name (enum ggml_op op); + GGML_API const char * ggml_op_symbol(enum ggml_op op); + + GGML_API const char * ggml_unary_op_name(enum ggml_unary_op op); + GGML_API const char * ggml_glu_op_name(enum ggml_glu_op op); + GGML_API const char * ggml_op_desc(const struct ggml_tensor * t); // unary or op name + + GGML_API size_t ggml_element_size(const struct ggml_tensor * tensor); + + GGML_API bool ggml_is_quantized(enum ggml_type type); + + // TODO: temporary until model loading of ggml examples is refactored + GGML_API enum ggml_type ggml_ftype_to_ggml_type(enum ggml_ftype ftype); + + GGML_API bool ggml_is_transposed(const struct ggml_tensor * tensor); + GGML_API bool ggml_is_permuted (const struct ggml_tensor * tensor); + GGML_API bool ggml_is_empty (const struct ggml_tensor * tensor); + GGML_API bool ggml_is_view (const struct ggml_tensor * tensor); + GGML_API bool ggml_is_scalar (const struct ggml_tensor * tensor); + GGML_API bool ggml_is_vector (const struct ggml_tensor * tensor); + GGML_API bool ggml_is_matrix (const struct ggml_tensor * tensor); + GGML_API bool ggml_is_3d (const struct ggml_tensor * tensor); + GGML_API int ggml_n_dims (const struct ggml_tensor * tensor); // returns 1 for scalars + + // returns whether the tensor elements can be iterated over with a flattened index (no gaps, no permutation) + GGML_API bool ggml_is_contiguous (const struct ggml_tensor * tensor); + GGML_API bool ggml_is_contiguous_0(const struct ggml_tensor * tensor); // same as ggml_is_contiguous() + GGML_API bool ggml_is_contiguous_1(const struct ggml_tensor * tensor); // contiguous for dims >= 1 + GGML_API bool ggml_is_contiguous_2(const struct ggml_tensor * tensor); // contiguous for dims >= 2 + + GGML_API bool ggml_is_contiguous_to_1(const struct ggml_tensor * tensor); // contiguous for dims < 1 + GGML_API bool ggml_is_contiguous_to_2(const struct ggml_tensor * tensor); // contiguous for dims < 2 + GGML_API bool ggml_is_contiguous_to_3(const struct ggml_tensor * tensor); // contiguous for dims < 3 + + // returns whether the tensor elements are allocated as one contiguous block of memory (no gaps, but permutation ok) + GGML_API bool ggml_is_contiguously_allocated(const struct ggml_tensor * tensor); + + // true for tensor that is stored in memory as CxWxHxN and has been permuted to WxHxCxN + GGML_API bool ggml_is_contiguous_channels(const struct ggml_tensor * tensor); + + // true if the elements in dimension 0 are contiguous, or there is just 1 block of elements + GGML_API bool ggml_is_contiguous_rows(const struct ggml_tensor * tensor); + + GGML_API bool ggml_are_same_shape (const struct ggml_tensor * t0, const struct ggml_tensor * t1); + GGML_API bool ggml_are_same_stride(const struct ggml_tensor * t0, const struct ggml_tensor * t1); + + GGML_API bool ggml_can_repeat(const struct ggml_tensor * t0, const struct ggml_tensor * t1); + + // use this to compute the memory overhead of a tensor + GGML_API size_t ggml_tensor_overhead(void); + + GGML_API bool ggml_validate_row_data(enum ggml_type type, const void * data, size_t nbytes); + + // main + + GGML_API struct ggml_context * ggml_init (struct ggml_init_params params); + GGML_API void ggml_reset(struct ggml_context * ctx); + GGML_API void ggml_free (struct ggml_context * ctx); + + GGML_API size_t ggml_used_mem(const struct ggml_context * ctx); + + GGML_API bool ggml_get_no_alloc(struct ggml_context * ctx); + GGML_API void ggml_set_no_alloc(struct ggml_context * ctx, bool no_alloc); + + GGML_API void * ggml_get_mem_buffer (const struct ggml_context * ctx); + GGML_API size_t ggml_get_mem_size (const struct ggml_context * ctx); + GGML_API size_t ggml_get_max_tensor_size(const struct ggml_context * ctx); + + GGML_API struct ggml_tensor * ggml_new_tensor( + struct ggml_context * ctx, + enum ggml_type type, + int n_dims, + const int64_t *ne); + + GGML_API struct ggml_tensor * ggml_new_tensor_1d( + struct ggml_context * ctx, + enum ggml_type type, + int64_t ne0); + + GGML_API struct ggml_tensor * ggml_new_tensor_2d( + struct ggml_context * ctx, + enum ggml_type type, + int64_t ne0, + int64_t ne1); + + GGML_API struct ggml_tensor * ggml_new_tensor_3d( + struct ggml_context * ctx, + enum ggml_type type, + int64_t ne0, + int64_t ne1, + int64_t ne2); + + GGML_API struct ggml_tensor * ggml_new_tensor_4d( + struct ggml_context * ctx, + enum ggml_type type, + int64_t ne0, + int64_t ne1, + int64_t ne2, + int64_t ne3); + + GGML_API void * ggml_new_buffer(struct ggml_context * ctx, size_t nbytes); + + GGML_API struct ggml_tensor * ggml_dup_tensor (struct ggml_context * ctx, const struct ggml_tensor * src); + GGML_API struct ggml_tensor * ggml_view_tensor(struct ggml_context * ctx, struct ggml_tensor * src); + + // Context tensor enumeration and lookup + GGML_API struct ggml_tensor * ggml_get_first_tensor(const struct ggml_context * ctx); + GGML_API struct ggml_tensor * ggml_get_next_tensor (const struct ggml_context * ctx, struct ggml_tensor * tensor); + GGML_API struct ggml_tensor * ggml_get_tensor(struct ggml_context * ctx, const char * name); + + // Converts a flat index into coordinates + GGML_API void ggml_unravel_index(const struct ggml_tensor * tensor, int64_t i, int64_t * i0, int64_t * i1, int64_t * i2, int64_t * i3); + + GGML_API enum ggml_unary_op ggml_get_unary_op(const struct ggml_tensor * tensor); + GGML_API enum ggml_glu_op ggml_get_glu_op(const struct ggml_tensor * tensor); + + GGML_API void * ggml_get_data (const struct ggml_tensor * tensor); + GGML_API float * ggml_get_data_f32(const struct ggml_tensor * tensor); + + GGML_API const char * ggml_get_name (const struct ggml_tensor * tensor); + GGML_API struct ggml_tensor * ggml_set_name ( struct ggml_tensor * tensor, const char * name); + GGML_ATTRIBUTE_FORMAT(2, 3) + GGML_API struct ggml_tensor * ggml_format_name( struct ggml_tensor * tensor, const char * fmt, ...); + + // Tensor flags + GGML_API void ggml_set_input(struct ggml_tensor * tensor); + GGML_API void ggml_set_output(struct ggml_tensor * tensor); + GGML_API void ggml_set_param(struct ggml_tensor * tensor); + GGML_API void ggml_set_loss(struct ggml_tensor * tensor); + + // + // operations on tensors with backpropagation + // + + GGML_API struct ggml_tensor * ggml_dup( + struct ggml_context * ctx, + struct ggml_tensor * a); + + // in-place, returns view(a) + GGML_API struct ggml_tensor * ggml_dup_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_add( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b); + + GGML_API struct ggml_tensor * ggml_add_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b); + + GGML_API struct ggml_tensor * ggml_add_cast( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + enum ggml_type type); + + // dst[i0, i1, i2] = a[i0, i1, i2] + b[i0, ids[i1, i2]] + GGML_API struct ggml_tensor * ggml_add_id( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + struct ggml_tensor * ids); + + GGML_DEPRECATED(GGML_API struct ggml_tensor * ggml_add1( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b), + "use ggml_add instead"); + + GGML_DEPRECATED(GGML_API struct ggml_tensor * ggml_add1_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b), + "use ggml_add_inplace instead"); + + // dst = a + // view(dst, nb1, nb2, nb3, offset) += b + // return dst + GGML_API struct ggml_tensor * ggml_acc( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + size_t nb1, + size_t nb2, + size_t nb3, + size_t offset); + + GGML_API struct ggml_tensor * ggml_acc_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + size_t nb1, + size_t nb2, + size_t nb3, + size_t offset); + + GGML_API struct ggml_tensor * ggml_sub( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b); + + GGML_API struct ggml_tensor * ggml_sub_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b); + + GGML_API struct ggml_tensor * ggml_mul( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b); + + GGML_API struct ggml_tensor * ggml_mul_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b); + + GGML_API struct ggml_tensor * ggml_div( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b); + + GGML_API struct ggml_tensor * ggml_div_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b); + + GGML_API struct ggml_tensor * ggml_sqr( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_sqr_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_sqrt( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_sqrt_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_log( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_log_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_expm1( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_expm1_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_softplus( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_softplus_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_sin( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_sin_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_cos( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_cos_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a); + + // return scalar + GGML_API struct ggml_tensor * ggml_sum( + struct ggml_context * ctx, + struct ggml_tensor * a); + + // sums along rows, with input shape [a,b,c,d] return shape [1,b,c,d] + GGML_API struct ggml_tensor * ggml_sum_rows( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_cumsum( + struct ggml_context * ctx, + struct ggml_tensor * a); + + // mean along rows + GGML_API struct ggml_tensor * ggml_mean( + struct ggml_context * ctx, + struct ggml_tensor * a); + + // argmax along rows + GGML_API struct ggml_tensor * ggml_argmax( + struct ggml_context * ctx, + struct ggml_tensor * a); + + // count number of equal elements in a and b + GGML_API struct ggml_tensor * ggml_count_equal( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b); + + // if a is the same shape as b, and a is not parameter, return a + // otherwise, return a new tensor: repeat(a) to fit in b + GGML_API struct ggml_tensor * ggml_repeat( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b); + + // repeat a to the specified shape + GGML_API struct ggml_tensor * ggml_repeat_4d( + struct ggml_context * ctx, + struct ggml_tensor * a, + int64_t ne0, + int64_t ne1, + int64_t ne2, + int64_t ne3); + + // sums repetitions in a into shape of b + GGML_API struct ggml_tensor * ggml_repeat_back( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b); // sum up values that are adjacent in dims > 0 instead of repeated with same stride + + // concat a and b along dim + // used in stable-diffusion + GGML_API struct ggml_tensor * ggml_concat( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + int dim); + + GGML_API struct ggml_tensor * ggml_abs( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_abs_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_sgn( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_sgn_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_neg( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_neg_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_step( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_step_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_tanh( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_tanh_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_elu( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_elu_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_relu( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_leaky_relu( + struct ggml_context * ctx, + struct ggml_tensor * a, float negative_slope, bool inplace); + + GGML_API struct ggml_tensor * ggml_relu_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_sigmoid( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_sigmoid_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_gelu( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_gelu_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a); + + // GELU using erf (error function) when possible + // some backends may fallback to approximation based on Abramowitz and Stegun formula + GGML_API struct ggml_tensor * ggml_gelu_erf( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_gelu_erf_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_gelu_quick( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_gelu_quick_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_silu( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_silu_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a); + + // a - dy + // b - x + GGML_API struct ggml_tensor * ggml_silu_back( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b); + + // hardswish(x) = x * relu6(x + 3) / 6 + GGML_API struct ggml_tensor * ggml_hardswish( + struct ggml_context * ctx, + struct ggml_tensor * a); + + // hardsigmoid(x) = relu6(x + 3) / 6 + GGML_API struct ggml_tensor * ggml_hardsigmoid( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_exp( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_exp_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_floor( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_floor_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_ceil( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_ceil_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_round( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_round_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a); + + /** + * Truncates the fractional part of each element in the tensor (towards zero). + * For example: trunc(3.7) = 3.0, trunc(-2.9) = -2.0 + * Similar to std::trunc in C/C++. + */ + + GGML_API struct ggml_tensor * ggml_trunc( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_trunc_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a); + + + + // xIELU activation function + // x = x * (c_a(alpha_n) + c_b(alpha_p, beta) * sigmoid(beta * x)) + eps * (x > 0) + // where c_a = softplus and c_b(a, b) = softplus(a) + b are constraining functions + // that constrain the positive and negative source alpha values respectively + GGML_API struct ggml_tensor * ggml_xielu( + struct ggml_context * ctx, + struct ggml_tensor * a, + float alpha_n, + float alpha_p, + float beta, + float eps); + + // gated linear unit ops + // A: n columns, r rows, + // result is n / 2 columns, r rows, + // expects gate in second half of row, unless swapped is true + GGML_API struct ggml_tensor * ggml_glu( + struct ggml_context * ctx, + struct ggml_tensor * a, + enum ggml_glu_op op, + bool swapped); + + GGML_API struct ggml_tensor * ggml_reglu( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_reglu_swapped( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_geglu( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_geglu_swapped( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_swiglu( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_swiglu_swapped( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_geglu_erf( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_geglu_erf_swapped( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_geglu_quick( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_geglu_quick_swapped( + struct ggml_context * ctx, + struct ggml_tensor * a); + + // A: n columns, r rows, + // B: n columns, r rows, + GGML_API struct ggml_tensor * ggml_glu_split( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + enum ggml_glu_op op); + + GGML_API struct ggml_tensor * ggml_reglu_split( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b); + + GGML_API struct ggml_tensor * ggml_geglu_split( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b); + + GGML_API struct ggml_tensor * ggml_swiglu_split( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b); + + GGML_API struct ggml_tensor * ggml_geglu_erf_split( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b); + + GGML_API struct ggml_tensor * ggml_geglu_quick_split( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b); + + GGML_API struct ggml_tensor * ggml_swiglu_oai( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + float alpha, + float limit); + + GGML_API struct ggml_tensor * ggml_swiglu_clamp( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + float limit); + + // normalize along rows + GGML_API struct ggml_tensor * ggml_norm( + struct ggml_context * ctx, + struct ggml_tensor * a, + float eps); + + GGML_API struct ggml_tensor * ggml_norm_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + float eps); + + GGML_API struct ggml_tensor * ggml_rms_norm( + struct ggml_context * ctx, + struct ggml_tensor * a, + float eps); + + GGML_API struct ggml_tensor * ggml_rms_norm_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + float eps); + + // group normalize along ne0*ne1*n_groups + // used in stable-diffusion + GGML_API struct ggml_tensor * ggml_group_norm( + struct ggml_context * ctx, + struct ggml_tensor * a, + int n_groups, + float eps); + + GGML_API struct ggml_tensor * ggml_group_norm_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + int n_groups, + float eps); + + // l2 normalize along rows + // used in rwkv v7 + GGML_API struct ggml_tensor * ggml_l2_norm( + struct ggml_context * ctx, + struct ggml_tensor * a, + float eps); + + GGML_API struct ggml_tensor * ggml_l2_norm_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + float eps); + + // a - x + // b - dy + GGML_API struct ggml_tensor * ggml_rms_norm_back( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + float eps); + + // [TAG_GGML_PREC] + // set the minimum required accumulator type for the implementation to use during the compute + // for example: + // - GGML_PREC_F32 - requires accumulation of the results in F32 + // - GGML_PREC_BF16 - can accumulate the results in BF16, F32 + // - GGML_PREC_F16 - can accumulate the results in F16, F32 + // - GGML_PREC_Q8 - not allowed + // - GGML_PREC_Q4 - not allowed + // + // return false on faliure + GGML_API bool ggml_prec_set_acc( + struct ggml_tensor * a, + enum ggml_prec prec); + + // [TAG_GGML_PREC] + // set the smallest rank that the implementation can use to internally convert the src[idx] data to + // ranks in decreasing order: + // - GGML_PREC_F32 - GGML_TYPE_F32 + // - GGML_PREC_BF16 - GGML_TYPE_BF16 + // - GGML_PREC_F16 - GGML_TYPE_F16, + // - GGML_PREC_Q8 - GGML_TYPE_Q8_0, GGML_TYPE_Q8_1, GGML_TYPE_Q8_K, etc. + // - GGML_PREC_Q4 - GGML_TYPE_Q4_0, GGML_TYPE_Q4_1, GGML_TYPE_Q4_K, GGML_TYPE_NVFP4, GGML_TYPE_MXFP4, etc. + // + // for example: + // - ggml_prec_set_src(a, GGML_PREC_Q8, 1): + // - allows the implementation to quantize F32, BF16, F16 data of src[1] down to GGML_TYPE_Q8_0 + // - cannot quantize it down to GGML_TYPE_Q4_0 or GGML_TYPE_NVFP4 + // - ggml_prec_set_src(a, GGML_PREC_Q4, 1): + // - allows the implementation to quantize F32, BF16, F16 data of src[1] down to 4-bit datatypes such as GGML_TYPE_Q4_K, GGML_TYPE_NVFP4 etc. + // + // return false on faliure + GGML_API bool ggml_prec_set_src( + struct ggml_tensor * a, + enum ggml_prec prec, + int idx); + + // A: k columns, n rows => [ne03, ne02, n, k] + // B: k columns, m rows (i.e. we transpose it internally) => [ne03 * x, ne02 * y, m, k] + // result is n columns, m rows => [ne03 * x, ne02 * y, m, n] + GGML_API struct ggml_tensor * ggml_mul_mat( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b); + + // change the precision of a matrix multiplication + // set to GGML_PREC_F32 for higher precision (useful for phi-2) + GGML_DEPRECATED(GGML_API void ggml_mul_mat_set_prec( + struct ggml_tensor * a, + enum ggml_prec prec), + "use ggml_prec_set_acc() instead"); + + // change the hint of a matrix multiplication + GGML_API void ggml_mul_mat_set_hint( + struct ggml_tensor * a, + enum ggml_op_hint hint); + + // indirect matrix multiplication + // + // [TAG_MUL_MAT_ID_DUPLICATE_ID] + // note: this op has undefined behaviour if a given `a` is referenced more than once for a given `b` + // (i.e. duplicate experts per token are not guaranteed to be supported in all backends) + // TODO: add backend support: https://github.com/ggml-org/llama.cpp/pull/29998#issuecomment-6153742533 + GGML_API struct ggml_tensor * ggml_mul_mat_id( + struct ggml_context * ctx, + struct ggml_tensor * as, + struct ggml_tensor * b, + struct ggml_tensor * ids); + + // A: m columns, n rows, + // B: p columns, n rows, + // result is m columns, p rows + GGML_API struct ggml_tensor * ggml_out_prod( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b); + + // + // operations on tensors without backpropagation + // + + GGML_API struct ggml_tensor * ggml_scale( + struct ggml_context * ctx, + struct ggml_tensor * a, + float s); + + // in-place, returns view(a) + GGML_API struct ggml_tensor * ggml_scale_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + float s); + + // x = s * a + b + GGML_API struct ggml_tensor * ggml_scale_bias( + struct ggml_context * ctx, + struct ggml_tensor * a, + float s, + float b); + + GGML_API struct ggml_tensor * ggml_scale_bias_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + float s, + float b); + + // b -> view(a,offset,nb1,nb2,3), return modified a + GGML_API struct ggml_tensor * ggml_set( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + size_t nb1, + size_t nb2, + size_t nb3, + size_t offset); // in bytes + + // b -> view(a,offset,nb1,nb2,3), return view(a) + GGML_API struct ggml_tensor * ggml_set_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + size_t nb1, + size_t nb2, + size_t nb3, + size_t offset); // in bytes + + GGML_API struct ggml_tensor * ggml_set_1d( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + size_t offset); // in bytes + + GGML_API struct ggml_tensor * ggml_set_1d_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + size_t offset); // in bytes + + // b -> view(a,offset,nb1,nb2,3), return modified a + GGML_API struct ggml_tensor * ggml_set_2d( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + size_t nb1, + size_t offset); // in bytes + + // b -> view(a,offset,nb1,nb2,3), return view(a) + GGML_API struct ggml_tensor * ggml_set_2d_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + size_t nb1, + size_t offset); // in bytes + + // a -> b, return view(b) + GGML_API struct ggml_tensor * ggml_cpy( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b); + + // note: casting from f32 to i32 will discard the fractional part + GGML_API struct ggml_tensor * ggml_cast( + struct ggml_context * ctx, + struct ggml_tensor * a, + enum ggml_type type); + + // make contiguous + GGML_API struct ggml_tensor * ggml_cont( + struct ggml_context * ctx, + struct ggml_tensor * a); + + // make contiguous, with new shape + GGML_API struct ggml_tensor * ggml_cont_1d( + struct ggml_context * ctx, + struct ggml_tensor * a, + int64_t ne0); + + GGML_API struct ggml_tensor * ggml_cont_2d( + struct ggml_context * ctx, + struct ggml_tensor * a, + int64_t ne0, + int64_t ne1); + + GGML_API struct ggml_tensor * ggml_cont_3d( + struct ggml_context * ctx, + struct ggml_tensor * a, + int64_t ne0, + int64_t ne1, + int64_t ne2); + + GGML_API struct ggml_tensor * ggml_cont_4d( + struct ggml_context * ctx, + struct ggml_tensor * a, + int64_t ne0, + int64_t ne1, + int64_t ne2, + int64_t ne3); + + // return view(a), b specifies the new shape + // TODO: when we start computing gradient, make a copy instead of view + GGML_API struct ggml_tensor * ggml_reshape( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b); + + // return view(a) + // TODO: when we start computing gradient, make a copy instead of view + GGML_API struct ggml_tensor * ggml_reshape_1d( + struct ggml_context * ctx, + struct ggml_tensor * a, + int64_t ne0); + + GGML_API struct ggml_tensor * ggml_reshape_2d( + struct ggml_context * ctx, + struct ggml_tensor * a, + int64_t ne0, + int64_t ne1); + + // return view(a) + // TODO: when we start computing gradient, make a copy instead of view + GGML_API struct ggml_tensor * ggml_reshape_3d( + struct ggml_context * ctx, + struct ggml_tensor * a, + int64_t ne0, + int64_t ne1, + int64_t ne2); + + GGML_API struct ggml_tensor * ggml_reshape_4d( + struct ggml_context * ctx, + struct ggml_tensor * a, + int64_t ne0, + int64_t ne1, + int64_t ne2, + int64_t ne3); + + // offset in bytes + GGML_API struct ggml_tensor * ggml_view_1d( + struct ggml_context * ctx, + struct ggml_tensor * a, + int64_t ne0, + size_t offset); + + GGML_API struct ggml_tensor * ggml_view_2d( + struct ggml_context * ctx, + struct ggml_tensor * a, + int64_t ne0, + int64_t ne1, + size_t nb1, // row stride in bytes + size_t offset); + + GGML_API struct ggml_tensor * ggml_view_3d( + struct ggml_context * ctx, + struct ggml_tensor * a, + int64_t ne0, + int64_t ne1, + int64_t ne2, + size_t nb1, // row stride in bytes + size_t nb2, // slice stride in bytes + size_t offset); + + GGML_API struct ggml_tensor * ggml_view_4d( + struct ggml_context * ctx, + struct ggml_tensor * a, + int64_t ne0, + int64_t ne1, + int64_t ne2, + int64_t ne3, + size_t nb1, // row stride in bytes + size_t nb2, // slice stride in bytes + size_t nb3, + size_t offset); + + GGML_API struct ggml_tensor * ggml_permute( + struct ggml_context * ctx, + struct ggml_tensor * a, + int axis0, + int axis1, + int axis2, + int axis3); + + // alias for ggml_permute(ctx, a, 1, 0, 2, 3) + GGML_API struct ggml_tensor * ggml_transpose( + struct ggml_context * ctx, + struct ggml_tensor * a); + + // supports 4D a: + // a [n_embd, ne1, ne2, ne3] + // b I32 [n_rows, ne2, ne3, 1] + // + // return [n_embd, n_rows, ne2, ne3] + GGML_API struct ggml_tensor * ggml_get_rows( + struct ggml_context * ctx, + struct ggml_tensor * a, // data + struct ggml_tensor * b); // row indices + + GGML_API struct ggml_tensor * ggml_get_rows_back( + struct ggml_context * ctx, + struct ggml_tensor * a, // gradients of ggml_get_rows result + struct ggml_tensor * b, // row indices + struct ggml_tensor * c); // data for ggml_get_rows, only used for its shape + + // a TD [n_embd, ne1, ne2, ne3] + // b TS [n_embd, n_rows, ne02, ne03] | ne02 == ne2, ne03 == ne3 + // c I64 [n_rows, ne11, ne12, 1] | c[i] in [0, ne1) + // + // undefined behavior if destination rows overlap + // + // broadcast: + // ne2 % ne11 == 0 + // ne3 % ne12 == 0 + // + // return view(a) + GGML_API struct ggml_tensor * ggml_set_rows( + struct ggml_context * ctx, + struct ggml_tensor * a, // destination + struct ggml_tensor * b, // source + struct ggml_tensor * c); // row indices + + GGML_API struct ggml_tensor * ggml_diag( + struct ggml_context * ctx, + struct ggml_tensor * a); + + // set elements above the diagonal to -INF + GGML_API struct ggml_tensor * ggml_diag_mask_inf( + struct ggml_context * ctx, + struct ggml_tensor * a, + int n_past); + + // in-place, returns view(a) + GGML_API struct ggml_tensor * ggml_diag_mask_inf_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + int n_past); + + // set elements above the diagonal to 0 + GGML_API struct ggml_tensor * ggml_diag_mask_zero( + struct ggml_context * ctx, + struct ggml_tensor * a, + int n_past); + + // in-place, returns view(a) + GGML_API struct ggml_tensor * ggml_diag_mask_zero_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + int n_past); + + GGML_API struct ggml_tensor * ggml_clamp( + struct ggml_context * ctx, + struct ggml_tensor * a, + float min, + float max); + + // in-place, returns view(a) + GGML_API struct ggml_tensor * ggml_clamp_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + float min, + float max); + + GGML_API struct ggml_tensor * ggml_soft_max( + struct ggml_context * ctx, + struct ggml_tensor * a); + + // in-place, returns view(a) + GGML_API struct ggml_tensor * ggml_soft_max_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a); + + // a [ne0, ne01, ne02, ne03] + // mask [ne0, ne11, ne12, ne13] | ne11 >= ne01, F16 or F32, optional + // + // broadcast: + // ne02 % ne12 == 0 + // ne03 % ne13 == 0 + // + // fused soft_max(a*scale + mask*(ALiBi slope)) + // max_bias = 0.0f for no ALiBi + GGML_API struct ggml_tensor * ggml_soft_max_ext( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * mask, + float scale, + float max_bias); + + GGML_API struct ggml_tensor * ggml_soft_max_ext_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * mask, + float scale, + float max_bias); + + GGML_API void ggml_soft_max_add_sinks( + struct ggml_tensor * a, + struct ggml_tensor * sinks); + + GGML_API struct ggml_tensor * ggml_soft_max_ext_back( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + float scale, + float max_bias); + + // in-place, returns view(a) + GGML_API struct ggml_tensor * ggml_soft_max_ext_back_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + float scale, + float max_bias); + + // rotary position embedding + // if (mode & 1) - skip n_past elements (NOT SUPPORTED) + // if (mode & GGML_ROPE_TYPE_NEOX) - GPT-NeoX style + // + // b is an int32 vector with size a->ne[2], it contains the positions + GGML_API struct ggml_tensor * ggml_rope( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + int n_dims, + int mode); + + // in-place, returns view(a) + GGML_API struct ggml_tensor * ggml_rope_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + int n_dims, + int mode); + + // RoPE operations with extended options + // a is the input tensor to apply RoPE to, shape [n_embd, n_head, n_token] + // b is an int32 vector with size n_token + // c is freq factors (e.g. phi3-128k), (optional) + // mode can be GGML_ROPE_TYPE_NORMAL or NEOX; for MROPE and VISION mode, use ggml_rope_multi + // + // pseudo-code for computing theta: + // for i in [0, n_dims/2): + // theta[i] = b[i] * powf(freq_base, -2.0 * i / n_dims); + // theta[i] = theta[i] / c[i]; # if c is provided, divide theta by c + // theta[i] = rope_yarn(theta[i], ...); # note: theta = theta * freq_scale is applied here + // + // other params are used by YaRN RoPE scaling, these default values will disable YaRN: + // freq_scale = 1.0f + // ext_factor = 0.0f + // attn_factor = 1.0f + // beta_fast = 0.0f + // beta_slow = 0.0f + // + // example: + // (marking: c = cos, s = sin, 0 = unrotated) + // given a single head with size = 8 --> [00000000] + // GGML_ROPE_TYPE_NORMAL n_dims = 4 --> [cscs0000] + // GGML_ROPE_TYPE_NORMAL n_dims = 8 --> [cscscscs] + // GGML_ROPE_TYPE_NEOX n_dims = 4 --> [ccss0000] + // GGML_ROPE_TYPE_NEOX n_dims = 8 --> [ccccssss] + GGML_API struct ggml_tensor * ggml_rope_ext( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + struct ggml_tensor * c, + int n_dims, + int mode, + int n_ctx_orig, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow); + + // multi-dimensional RoPE, for Qwen-VL and similar vision models + // mode can be either VISION, MROPE, IMROPE, cannot be combined with NORMAL or NEOX + // sections specify how many dimensions to rotate in each section: + // section length is equivalent to number of cos/sin pairs, NOT the number of dims + // (i.e. sum of 4 sections are expected to be n_dims/2) + // last sections can be 0, means ignored + // all other options are identical to ggml_rope_ext + // + // important note: + // - NEOX ordering is automatically applied and cannot be disabled for MROPE and VISION + // if you need normal ordering, there are 2 methods: + // (1) split the tensor manually using ggml_view + // (2) permute the weight upon conversion + // - for VISION, n_dims must be head_size/2 + // + // example M-RoPE: + // given sections = [t=4, y=2, x=2, 0] + // given a single head with size = 18 --> [000000000000000000] + // GGML_ROPE_TYPE_MROPE n_dims = 16 --> [ttttyyxxttttyyxx00] (cos/sin are applied in NEOX ordering) + // GGML_ROPE_TYPE_IMROPE n_dims = 16 --> [ttyxttyxttyxttyx00] (interleaved M-RoPE, still NEOX ordering) + // note: the theta for each dim is computed the same way as ggml_rope_ext, no matter the section + // in other words, idx used for theta: [0123456789... until n_dims/2], not reset for each section + // + // example vision RoPE: + // given sections = [y=4, x=4, 0, 0] (last 2 sections are ignored) + // given a single head with size = 8 --> [00000000] + // GGML_ROPE_TYPE_VISION n_dims = 4 --> [yyyyxxxx] + // other values of n_dims are untested and is undefined behavior + // note: unlike MROPE, the theta for each dim is computed differently for each section + // in other words, idx used for theta: [0123] for y section, then [0123] for x section + GGML_API struct ggml_tensor * ggml_rope_multi( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + struct ggml_tensor * c, + int n_dims, + int sections[GGML_MROPE_SECTIONS], + int mode, + int n_ctx_orig, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow); + + // in-place, returns view(a) + GGML_API struct ggml_tensor * ggml_rope_ext_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + struct ggml_tensor * c, + int n_dims, + int mode, + int n_ctx_orig, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow); + + GGML_API struct ggml_tensor * ggml_rope_multi_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + struct ggml_tensor * c, + int n_dims, + int sections[GGML_MROPE_SECTIONS], + int mode, + int n_ctx_orig, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow); + + GGML_DEPRECATED(GGML_API struct ggml_tensor * ggml_rope_custom( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + int n_dims, + int mode, + int n_ctx_orig, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow), + "use ggml_rope_ext instead"); + + GGML_DEPRECATED(GGML_API struct ggml_tensor * ggml_rope_custom_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + int n_dims, + int mode, + int n_ctx_orig, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow), + "use ggml_rope_ext_inplace instead"); + + // compute correction dims for YaRN RoPE scaling + GGML_API void ggml_rope_yarn_corr_dims( + int n_dims, int n_ctx_orig, float freq_base, float beta_fast, float beta_slow, float dims[2]); + + // rotary position embedding backward, i.e compute dx from dy + // a - dy + GGML_API struct ggml_tensor * ggml_rope_ext_back( + struct ggml_context * ctx, + struct ggml_tensor * a, // gradients of ggml_rope result + struct ggml_tensor * b, // positions + struct ggml_tensor * c, // freq factors + int n_dims, + int mode, + int n_ctx_orig, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow); + + GGML_API struct ggml_tensor * ggml_rope_multi_back( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + struct ggml_tensor * c, + int n_dims, + int sections[4], + int mode, + int n_ctx_orig, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow); + + // set the offset dims for RoPE + // a must be GGML_OP_ROPE or GGML_OP_ROPE_BACK + // vision RoPE is not supported + // example: (marking: x = rotated, 0 = unrotated) + // n_embd = 10, n_dims = 4, offset = 2 --> [00xxxx0000] + GGML_API struct ggml_tensor * ggml_rope_set_offset( + struct ggml_tensor * a, + int n_offs); + + // im2col + // converts data into a format that effectively results in a convolution when combined with matrix multiplication + GGML_API struct ggml_tensor * ggml_im2col( + struct ggml_context * ctx, + struct ggml_tensor * a, // convolution kernel + struct ggml_tensor * b, // data + int s0, // stride dimension 0 + int s1, // stride dimension 1 + int p0, // padding dimension 0 + int p1, // padding dimension 1 + int d0, // dilation dimension 0 + int d1, // dilation dimension 1 + bool is_2D, + enum ggml_type dst_type); + + GGML_API struct ggml_tensor * ggml_im2col_back( + struct ggml_context * ctx, + struct ggml_tensor * a, // convolution kernel + struct ggml_tensor * b, // gradient of im2col output + int64_t * ne, // shape of im2col input + int s0, // stride dimension 0 + int s1, // stride dimension 1 + int p0, // padding dimension 0 + int p1, // padding dimension 1 + int d0, // dilation dimension 0 + int d1, // dilation dimension 1 + bool is_2D); + + // col2im_1d: scatter-add GEMM columns back to 1D signal + // a: [K*OC, T_in] (columns from matmul, K = a->ne[0]/OC) + // result: [T_out, OC] where T_out = (T_in - 1)*s0 + K - 2*p0 + GGML_API struct ggml_tensor * ggml_col2im_1d( + struct ggml_context * ctx, + struct ggml_tensor * a, // columns [K*OC, T_in] + int s0, // stride + int oc, // output channels + int p0); // padding to crop from both sides + + GGML_API struct ggml_tensor * ggml_conv_1d( + struct ggml_context * ctx, + struct ggml_tensor * a, // convolution kernel + struct ggml_tensor * b, // data + int s0, // stride + int p0, // padding + int d0); // dilation + + // conv_1d with padding = half + // alias for ggml_conv_1d(a, b, s, a->ne[0]/2, d) + GGML_API struct ggml_tensor* ggml_conv_1d_ph( + struct ggml_context * ctx, + struct ggml_tensor * a, // convolution kernel + struct ggml_tensor * b, // data + int s, // stride + int d); // dilation + + // depthwise + // TODO: this is very likely wrong for some cases! - needs more testing + GGML_API struct ggml_tensor * ggml_conv_1d_dw( + struct ggml_context * ctx, + struct ggml_tensor * a, // convolution kernel + struct ggml_tensor * b, // data + int s0, // stride + int p0, // padding + int d0); // dilation + + GGML_API struct ggml_tensor * ggml_conv_1d_dw_ph( + struct ggml_context * ctx, + struct ggml_tensor * a, // convolution kernel + struct ggml_tensor * b, // data + int s0, // stride + int d0); // dilation + + GGML_API struct ggml_tensor * ggml_conv_transpose_1d( + struct ggml_context * ctx, + struct ggml_tensor * a, // convolution kernel + struct ggml_tensor * b, // data + int s0, // stride + int p0, // padding + int d0); // dilation + + GGML_API struct ggml_tensor * ggml_conv_2d( + struct ggml_context * ctx, + struct ggml_tensor * a, // convolution kernel + struct ggml_tensor * b, // data + int s0, // stride dimension 0 + int s1, // stride dimension 1 + int p0, // padding dimension 0 + int p1, // padding dimension 1 + int d0, // dilation dimension 0 + int d1); // dilation dimension 1 + + GGML_API struct ggml_tensor * ggml_im2col_3d( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + int64_t IC, + int s0, // stride width + int s1, // stride height + int s2, // stride depth + int p0, // padding width + int p1, // padding height + int p2, // padding depth + int d0, // dilation width + int d1, // dilation height + int d2, // dilation depth + enum ggml_type dst_type); + + // a: [OC*IC, KD, KH, KW] + // b: [N*IC, ID, IH, IW] + // result: [N*OC, OD, OH, OW] + GGML_API struct ggml_tensor * ggml_conv_3d( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + int64_t IC, + int s0, // stride width + int s1, // stride height + int s2, // stride depth + int p0, // padding width + int p1, // padding height + int p2, // padding depth + int d0, // dilation width + int d1, // dilation height + int d2 // dilation depth + ); + + // kernel size is a->ne[0] x a->ne[1] + // stride is equal to kernel size + // padding is zero + // example: + // a: 16 16 3 768 + // b: 1024 1024 3 1 + // res: 64 64 768 1 + // used in sam + GGML_API struct ggml_tensor * ggml_conv_2d_sk_p0( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b); + + // kernel size is a->ne[0] x a->ne[1] + // stride is 1 + // padding is half + // example: + // a: 3 3 256 256 + // b: 64 64 256 1 + // res: 64 64 256 1 + // used in sam + GGML_API struct ggml_tensor * ggml_conv_2d_s1_ph( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b); + + // depthwise (via im2col and mul_mat) + GGML_API struct ggml_tensor * ggml_conv_2d_dw( + struct ggml_context * ctx, + struct ggml_tensor * a, // convolution kernel + struct ggml_tensor * b, // data + int s0, // stride dimension 0 + int s1, // stride dimension 1 + int p0, // padding dimension 0 + int p1, // padding dimension 1 + int d0, // dilation dimension 0 + int d1); // dilation dimension 1 + + // Depthwise 2D convolution + // may be faster than ggml_conv_2d_dw, but not available in all backends + // a: KW KH 1 C convolution kernel + // b: W H C N input data + // res: W_out H_out C N + GGML_API struct ggml_tensor * ggml_conv_2d_dw_direct( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + int stride0, + int stride1, + int pad0, + int pad1, + int dilation0, + int dilation1); + + GGML_API struct ggml_tensor * ggml_conv_transpose_2d_p0( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + int stride); + + GGML_API struct ggml_tensor * ggml_conv_2d_direct( + struct ggml_context * ctx, + struct ggml_tensor * a, // convolution kernel [KW, KH, IC, OC] + struct ggml_tensor * b, // input data [W, H, C, N] + int s0, // stride dimension 0 + int s1, // stride dimension 1 + int p0, // padding dimension 0 + int p1, // padding dimension 1 + int d0, // dilation dimension 0 + int d1); // dilation dimension 1 + + GGML_API struct ggml_tensor * ggml_conv_3d_direct( + struct ggml_context * ctx, + struct ggml_tensor * a, // kernel [KW, KH, KD, IC * OC] + struct ggml_tensor * b, // input [W, H, D, C * N] + int s0, // stride + int s1, + int s2, + int p0, // padding + int p1, + int p2, + int d0, // dilation + int d1, + int d2, + int n_channels, + int n_batch, + int n_channels_out); + + enum ggml_op_pool { + GGML_OP_POOL_MAX, + GGML_OP_POOL_AVG, + GGML_OP_POOL_COUNT, + }; + + GGML_API struct ggml_tensor * ggml_pool_1d( + struct ggml_context * ctx, + struct ggml_tensor * a, + enum ggml_op_pool op, + int k0, // kernel size + int s0, // stride + int p0); // padding + + // the result will have 2*p0 padding for the first dimension + // and 2*p1 padding for the second dimension + GGML_API struct ggml_tensor * ggml_pool_2d( + struct ggml_context * ctx, + struct ggml_tensor * a, + enum ggml_op_pool op, + int k0, + int k1, + int s0, + int s1, + float p0, + float p1); + + GGML_API struct ggml_tensor * ggml_pool_2d_back( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * af, // "a"/input used in forward pass + enum ggml_op_pool op, + int k0, + int k1, + int s0, + int s1, + float p0, + float p1); + + enum ggml_scale_mode { + GGML_SCALE_MODE_NEAREST = 0, + GGML_SCALE_MODE_BILINEAR = 1, + GGML_SCALE_MODE_BICUBIC = 2, + + GGML_SCALE_MODE_COUNT + }; + + enum ggml_scale_flag { + GGML_SCALE_FLAG_ALIGN_CORNERS = (1 << 8), + GGML_SCALE_FLAG_ANTIALIAS = (1 << 9), + }; + + // interpolate + // multiplies ne0 and ne1 by scale factor + GGML_API struct ggml_tensor * ggml_upscale( + struct ggml_context * ctx, + struct ggml_tensor * a, + int scale_factor, + enum ggml_scale_mode mode); + + // interpolate + // interpolate scale to specified dimensions + GGML_DEPRECATED(GGML_API struct ggml_tensor * ggml_upscale_ext( + struct ggml_context * ctx, + struct ggml_tensor * a, + int ne0, + int ne1, + int ne2, + int ne3, + enum ggml_scale_mode mode), + "use ggml_interpolate instead"); + + // Up- or downsamples the input to the specified size. + // 2D scale modes (eg. bilinear) are applied to the first two dimensions. + GGML_API struct ggml_tensor * ggml_interpolate( + struct ggml_context * ctx, + struct ggml_tensor * a, + int64_t ne0, + int64_t ne1, + int64_t ne2, + int64_t ne3, + uint32_t mode); // ggml_scale_mode [ | ggml_scale_flag...] + + // pad each dimension with zeros: [x, ..., x] -> [x, ..., x, 0, ..., 0] + GGML_API struct ggml_tensor * ggml_pad( + struct ggml_context * ctx, + struct ggml_tensor * a, + int p0, + int p1, + int p2, + int p3); + + // pad each dimension with values on the other side of the torus (looping around) + GGML_API struct ggml_tensor * ggml_pad_circular( + struct ggml_context * ctx, + struct ggml_tensor * a, + int p0, + int p1, + int p2, + int p3); + + GGML_API struct ggml_tensor * ggml_pad_ext( + struct ggml_context * ctx, + struct ggml_tensor * a, + int lp0, + int rp0, + int lp1, + int rp1, + int lp2, + int rp2, + int lp3, + int rp3 + ); + + // pad each dimension with values on the other side of the torus (looping around) + GGML_API struct ggml_tensor * ggml_pad_ext_circular( + struct ggml_context * ctx, + struct ggml_tensor * a, + int lp0, + int rp0, + int lp1, + int rp1, + int lp2, + int rp2, + int lp3, + int rp3); + + // pad each dimension with reflection: [a, b, c, d] -> [b, a, b, c, d, c] + GGML_API struct ggml_tensor * ggml_pad_reflect_1d( + struct ggml_context * ctx, + struct ggml_tensor * a, + int p0, + int p1); + + // Move tensor elements by an offset given for each dimension. Elements that + // are shifted beyond the last position are wrapped around to the beginning. + GGML_API struct ggml_tensor * ggml_roll( + struct ggml_context * ctx, + struct ggml_tensor * a, + int shift0, + int shift1, + int shift2, + int shift3); + + // Convert matrix into a triangular one (upper, strict upper, lower or strict lower) by writing + // zeroes everywhere outside the masked area + GGML_API struct ggml_tensor * ggml_tri( + struct ggml_context * ctx, + struct ggml_tensor * a, + enum ggml_tri_type type); + + // Fill tensor a with constant c + GGML_API struct ggml_tensor * ggml_fill( + struct ggml_context * ctx, + struct ggml_tensor * a, + float c); + + GGML_API struct ggml_tensor * ggml_fill_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + float c); + + // Ref: https://github.com/CompVis/stable-diffusion/blob/main/ldm/modules/diffusionmodules/util.py#L151 + // timesteps: [N,] + // return: [N, dim] + GGML_API struct ggml_tensor * ggml_timestep_embedding( + struct ggml_context * ctx, + struct ggml_tensor * timesteps, + int dim, + int max_period); + + // sort rows + enum ggml_sort_order { + GGML_SORT_ORDER_ASC, + GGML_SORT_ORDER_DESC, + }; + + GGML_API struct ggml_tensor * ggml_argsort( + struct ggml_context * ctx, + struct ggml_tensor * a, + enum ggml_sort_order order); + + // similar to ggml_top_k but implemented as `argsort` + `view` + GGML_API struct ggml_tensor * ggml_argsort_top_k( + struct ggml_context * ctx, + struct ggml_tensor * a, + int k); + + // top k elements per row + // note: the resulting top k indices are in no particular order + GGML_API struct ggml_tensor * ggml_top_k( + struct ggml_context * ctx, + struct ggml_tensor * a, + int k); + + GGML_API struct ggml_tensor * ggml_arange( + struct ggml_context * ctx, + float start, + float stop, + float step); + + // q: [n_embd_k, n_batch, n_head, ne3 ] + // k: [n_embd_k, n_kv, n_head_kv, ne3 ] + // v: [n_embd_v, n_kv, n_head_kv, ne3 ] !! not transposed !! + // mask: [n_kv, n_batch, ne32, ne33] + // res: [n_embd_v, n_head, n_batch, ne3 ] !! permuted !! + // + // broadcast: + // n_head % n_head_kv == 0 + // n_head % ne32 == 0 + // ne3 % ne33 == 0 + // + GGML_API struct ggml_tensor * ggml_flash_attn_ext( + struct ggml_context * ctx, + struct ggml_tensor * q, + struct ggml_tensor * k, + struct ggml_tensor * v, + struct ggml_tensor * mask, + float scale, + float max_bias, + float logit_softcap); + + GGML_DEPRECATED(GGML_API void ggml_flash_attn_ext_set_prec( + struct ggml_tensor * a, + enum ggml_prec prec), + "use ggml_prec_set_acc() instead"); + + GGML_API enum ggml_prec ggml_flash_attn_ext_get_prec( + const struct ggml_tensor * a); + + // Use finite mask entries as a sparse K/V set. Set 0 to disable. + // n_kv_max must bound the number of finite entries in every mask row. + GGML_API void ggml_flash_attn_ext_set_n_kv_max( + struct ggml_tensor * a, + int32_t n_kv_max); + + GGML_API void ggml_flash_attn_ext_add_sinks( + struct ggml_tensor * a, + struct ggml_tensor * sinks); + + // TODO: needs to be adapted to ggml_flash_attn_ext + GGML_API struct ggml_tensor * ggml_flash_attn_back( + struct ggml_context * ctx, + struct ggml_tensor * q, + struct ggml_tensor * k, + struct ggml_tensor * v, + struct ggml_tensor * d, + bool masked); + + GGML_API struct ggml_tensor * ggml_ssm_conv( + struct ggml_context * ctx, + struct ggml_tensor * sx, + struct ggml_tensor * c); + + GGML_API struct ggml_tensor * ggml_ssm_scan( + struct ggml_context * ctx, + struct ggml_tensor * s, + struct ggml_tensor * x, + struct ggml_tensor * dt, + struct ggml_tensor * A, + struct ggml_tensor * B, + struct ggml_tensor * C, + struct ggml_tensor * ids, + int64_t K); + + // partition into non-overlapping windows with padding if needed + // example: + // a: 768 64 64 1 + // w: 14 + // res: 768 14 14 25 + // used in sam + GGML_API struct ggml_tensor * ggml_win_part( + struct ggml_context * ctx, + struct ggml_tensor * a, + int w); + + // reverse of ggml_win_part + // used in sam + GGML_API struct ggml_tensor * ggml_win_unpart( + struct ggml_context * ctx, + struct ggml_tensor * a, + int w0, + int h0, + int w); + + GGML_API struct ggml_tensor * ggml_unary( + struct ggml_context * ctx, + struct ggml_tensor * a, + enum ggml_unary_op op); + + GGML_API struct ggml_tensor * ggml_unary_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + enum ggml_unary_op op); + + // used in sam + GGML_API struct ggml_tensor * ggml_get_rel_pos( + struct ggml_context * ctx, + struct ggml_tensor * a, + int qh, + int kh); + + // used in sam + GGML_API struct ggml_tensor * ggml_add_rel_pos( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * pw, + struct ggml_tensor * ph); + + GGML_API struct ggml_tensor * ggml_add_rel_pos_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * pw, + struct ggml_tensor * ph); + + GGML_API struct ggml_tensor * ggml_rwkv_wkv6( + struct ggml_context * ctx, + struct ggml_tensor * k, + struct ggml_tensor * v, + struct ggml_tensor * r, + struct ggml_tensor * tf, + struct ggml_tensor * td, + struct ggml_tensor * state); + + GGML_API struct ggml_tensor * ggml_gated_linear_attn( + struct ggml_context * ctx, + struct ggml_tensor * k, + struct ggml_tensor * v, + struct ggml_tensor * q, + struct ggml_tensor * g, + struct ggml_tensor * state, + float scale); + + GGML_API struct ggml_tensor * ggml_rwkv_wkv7( + struct ggml_context * ctx, + struct ggml_tensor * r, + struct ggml_tensor * w, + struct ggml_tensor * k, + struct ggml_tensor * v, + struct ggml_tensor * a, + struct ggml_tensor * b, + struct ggml_tensor * state); + + /* Solves a specific equation of the form Ax=B, where A is a triangular matrix + * without zeroes on the diagonal (i.e. invertible). + * B can have any number of columns, but must have the same number of rows as A + * If A is [n, n] and B is [n, m], then the result will be [n, m] as well + * Has O(n^3) complexity (unlike most matrix ops out there), so use on cases + * where n > 100 sparingly, pre-chunk if necessary. + * + * If left = false, solves xA=B instead + * If lower = false, assumes upper triangular instead + * If uni = true, assumes diagonal of A to be all ones (will override actual values) + * + * TODO: currently only lower, right, non-unitriangular variant is implemented + */ + GGML_API struct ggml_tensor * ggml_solve_tri( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + bool left, + bool lower, + bool uni); + + // TODO: add ggml_gated_delta_net_set_bcast() to be able to configure Q, K broadcast type: tiled vs interleaved [TAG_GGML_GDN_BCAST] + // ref: https://github.com/ggml-org/llama.cpp/pull/19468#discussion_r2786394306 + // + // tensor shapes (S_k == S_v, H_v % H_k == 0): + // q, k : [S_k, H_k, n_tokens, n_seqs] + // v : [S_v, H_v, n_tokens, n_seqs] + // g : [1, H_v, n_tokens, n_seqs] (scalar gate) or [S_v, H_v, n_tokens, n_seqs] (KDA) + // beta : [1, H_v, n_tokens, n_seqs] + // state : [S_v, S_v, H_v, n_seqs] -- initial recurrent state s0 + // + // the output packs the attention scores [S_v, H_v, n_tokens, n_seqs] followed by K state + // snapshots, most-recent first (slot 0 = final state, slot s = state s tokens back). K == 1 + // keeps only the final state; when n_tokens < K only slots 0..n_tokens-1 are written. + GGML_API struct ggml_tensor * ggml_gated_delta_net( + struct ggml_context * ctx, + struct ggml_tensor * q, + struct ggml_tensor * k, + struct ggml_tensor * v, + struct ggml_tensor * g, + struct ggml_tensor * beta, + struct ggml_tensor * state, + int64_t K); + + // DSA lightning indexer + // + // q: [n_embd_idx, n_head_idx, n_batch, ne3 ] + // k: [n_embd_idx, 1, n_kv, ne3 ] + // weights: [n_head_idx, n_batch, 1, ne3 ] !! prescaled !! + // mask: [n_kv, n_batch, 1, ne33] !! f16 !! + // res: [n_kv, n_batch, 1, ne3 ] + // + // broadcast: + // ne3 % ne33 == 0 + // + GGML_API struct ggml_tensor * ggml_lightning_indexer( + struct ggml_context * ctx, + struct ggml_tensor * q, + struct ggml_tensor * k, + struct ggml_tensor * weights, + struct ggml_tensor * mask); + + // DeepSeek V4 hyper-connections (ref. https://arxiv.org/pdf/2512.24880) + // In short these operations are replacements for the original residual connection (x = transformer(x) + x) + // using a richer representation through streams. + // + // hc_comb: mixes [(2 + hc)*hc, n_tokens], scale [3], base [(2 + hc)*hc] + // -> [dst_hc, src_hc, n_tokens] + // logits[dst, src, t] = mixes[2*hc + dst + hc*src, t]*scale[2] + // + base[2*hc + dst + hc*src] + // Softmax over dst, add eps, normalize over src, then repeat normalization + // over dst followed by src for iterations 1 through n_iter - 1. + GGML_API struct ggml_tensor * ggml_dsv4_hc_comb( + struct ggml_context * ctx, + struct ggml_tensor * mixes, + struct ggml_tensor * scale, + struct ggml_tensor * base, + float eps, + int32_t n_iter); + + // hc_pre: x [n_embd, hc, n_tokens], weights [hc, n_tokens] -> [n_embd, n_tokens] + // result[i, t] = sum_h x[i, h, t]*weights[h, t] + // + GGML_API struct ggml_tensor * ggml_dsv4_hc_pre( + struct ggml_context * ctx, + struct ggml_tensor * x, + struct ggml_tensor * weights); + + // hc_pre with a per-element gate (Qwen3.8-Flash-Next): gate [n_embd, hc, n_tokens] + // result[i, t] = scale*sum_h x[i, h, t]*sigmoid(gate[i, h, t]) + // + GGML_API struct ggml_tensor * ggml_dsv4_hc_pre_gated( + struct ggml_context * ctx, + struct ggml_tensor * x, + struct ggml_tensor * gate, + float scale); + + // hc_post: x [n_embd, n_tokens], residual [n_embd, hc, n_tokens], + // post [hc, n_tokens], comb [dst_hc, src_hc, n_tokens] + // -> [n_embd, hc, n_tokens] + // result[i, dst, t] = x[i, t]*post[dst, t] + // + sum_src residual[i, src, t]*comb[dst, src, t] + // comb == NULL uses the identity: result[i, dst, t] = x[i, t]*post[dst, t] + residual[i, dst, t] + // + GGML_API struct ggml_tensor * ggml_dsv4_hc_post( + struct ggml_context * ctx, + struct ggml_tensor * x, + struct ggml_tensor * residual, + struct ggml_tensor * post, + struct ggml_tensor * comb); + + // custom operators + + typedef void (*ggml_custom1_op_t)(struct ggml_tensor * dst , const struct ggml_tensor * a, int ith, int nth, void * userdata); + typedef void (*ggml_custom2_op_t)(struct ggml_tensor * dst , const struct ggml_tensor * a, const struct ggml_tensor * b, int ith, int nth, void * userdata); + typedef void (*ggml_custom3_op_t)(struct ggml_tensor * dst , const struct ggml_tensor * a, const struct ggml_tensor * b, const struct ggml_tensor * c, int ith, int nth, void * userdata); + +#define GGML_N_TASKS_MAX (-1) + // n_tasks == GGML_N_TASKS_MAX means to use max number of tasks + + GGML_API struct ggml_tensor * ggml_map_custom1( + struct ggml_context * ctx, + struct ggml_tensor * a, + ggml_custom1_op_t fun, + int n_tasks, + void * userdata); + + GGML_API struct ggml_tensor * ggml_map_custom1_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + ggml_custom1_op_t fun, + int n_tasks, + void * userdata); + + GGML_API struct ggml_tensor * ggml_map_custom2( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + ggml_custom2_op_t fun, + int n_tasks, + void * userdata); + + GGML_API struct ggml_tensor * ggml_map_custom2_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + ggml_custom2_op_t fun, + int n_tasks, + void * userdata); + + GGML_API struct ggml_tensor * ggml_map_custom3( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + struct ggml_tensor * c, + ggml_custom3_op_t fun, + int n_tasks, + void * userdata); + + GGML_API struct ggml_tensor * ggml_map_custom3_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + struct ggml_tensor * c, + ggml_custom3_op_t fun, + int n_tasks, + void * userdata); + + typedef void (*ggml_custom_op_t)(struct ggml_tensor * dst , int ith, int nth, void * userdata); + + GGML_API struct ggml_tensor * ggml_custom_4d( + struct ggml_context * ctx, + enum ggml_type type, + int64_t ne0, + int64_t ne1, + int64_t ne2, + int64_t ne3, + struct ggml_tensor ** args, + int n_args, + ggml_custom_op_t fun, + int n_tasks, + void * userdata); + + GGML_API struct ggml_tensor * ggml_custom_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor ** args, + int n_args, + ggml_custom_op_t fun, + int n_tasks, + void * userdata); + + // loss function + + GGML_API struct ggml_tensor * ggml_cross_entropy_loss( + struct ggml_context * ctx, + struct ggml_tensor * a, // logits + struct ggml_tensor * b); // labels + + GGML_API struct ggml_tensor * ggml_cross_entropy_loss_back( + struct ggml_context * ctx, + struct ggml_tensor * a, // logits + struct ggml_tensor * b, // labels + struct ggml_tensor * c); // gradients of cross_entropy_loss result + + // AdamW optimizer step + // Paper: https://arxiv.org/pdf/1711.05101v3.pdf + // PyTorch: https://pytorch.org/docs/stable/generated/torch.optim.AdamW.html + GGML_API struct ggml_tensor * ggml_opt_step_adamw( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * grad, + struct ggml_tensor * m, + struct ggml_tensor * v, + struct ggml_tensor * adamw_params); // parameters such as the learning rate + + // stochastic gradient descent step (with weight decay) + GGML_API struct ggml_tensor * ggml_opt_step_sgd( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * grad, + struct ggml_tensor * sgd_params); // alpha, weight decay + + // build forward multiple tensors and select one of them for computing + // this is useful for creating graphs that have constant topology but compute different things based on the input + // ref: https://github.com/ggml-org/llama.cpp/pull/18550 + // + // nodes: + // | - build forward into the graph but do not compute + // c - build forward into the graph and compute + // + // | | ... c ... | + // | | ... c ... | + // | | ... c ... | + // [0 1 ... idx ... n-1] <-- ggml_build_forward_select(..., n, idx) + // c + // c + // + // example: + // struct ggml_tensor * curs[3]; + // + // curs[0] = compute0(...); + // curs[1] = compute1(...); + // curs[2] = compute2(...); + // + // int idx = select_branch(some_input); + // + // struct ggml_tensor * out = ggml_build_forward_select(cgraph, curs, 3, idx); + // + GGML_API struct ggml_tensor * ggml_build_forward_select( + struct ggml_cgraph * cgraph, + struct ggml_tensor ** tensors, + int n_tensors, + int idx); + + GGML_API void ggml_build_forward_expand( + struct ggml_cgraph * cgraph, + struct ggml_tensor * tensor); + + // add the tensor and its parents to the graph without marking them for compute + // the flag is set later, when the tensor is reached from a node that computes + GGML_API void ggml_build_forward_order( + struct ggml_cgraph * cgraph, + struct ggml_tensor * tensor); + + GGML_API void ggml_build_backward_expand( + struct ggml_context * ctx, // context for gradient computation + struct ggml_cgraph * cgraph, + struct ggml_tensor ** grad_accs); + + // graph allocation in a context + GGML_API struct ggml_cgraph * ggml_new_graph (struct ggml_context * ctx); // size = GGML_DEFAULT_GRAPH_SIZE, grads = false + GGML_API struct ggml_cgraph * ggml_new_graph_custom(struct ggml_context * ctx, size_t size, bool grads); + GGML_API struct ggml_cgraph * ggml_graph_dup (struct ggml_context * ctx, struct ggml_cgraph * cgraph, bool force_grads); + GGML_API void ggml_graph_cpy (struct ggml_cgraph * src, struct ggml_cgraph * dst); + GGML_API void ggml_graph_reset (struct ggml_cgraph * cgraph); // set regular grads + optimizer momenta to 0, set loss grad to 1 + GGML_API void ggml_graph_clear (struct ggml_cgraph * cgraph); + + GGML_API int ggml_graph_size (struct ggml_cgraph * cgraph); + GGML_API struct ggml_tensor * ggml_graph_node (struct ggml_cgraph * cgraph, int i); // if i < 0, returns nodes[n_nodes + i] + GGML_API struct ggml_tensor ** ggml_graph_nodes (struct ggml_cgraph * cgraph); + GGML_API int ggml_graph_n_nodes(struct ggml_cgraph * cgraph); + + GGML_API void ggml_graph_add_node(struct ggml_cgraph * cgraph, struct ggml_tensor * tensor); + + GGML_API size_t ggml_graph_overhead(void); + GGML_API size_t ggml_graph_overhead_custom(size_t size, bool grads); + + GGML_API struct ggml_tensor * ggml_graph_get_tensor (const struct ggml_cgraph * cgraph, const char * name); + GGML_API struct ggml_tensor * ggml_graph_get_grad (const struct ggml_cgraph * cgraph, const struct ggml_tensor * node); + GGML_API struct ggml_tensor * ggml_graph_get_grad_acc(const struct ggml_cgraph * cgraph, const struct ggml_tensor * node); + + // print info and performance information for the graph + GGML_API void ggml_graph_print(const struct ggml_cgraph * cgraph); + + // dump the graph into a file using the dot format + GGML_API void ggml_graph_dump_dot(const struct ggml_cgraph * gb, const struct ggml_cgraph * cgraph, const char * filename); + + // TODO these functions were sandwiched in the old optimization interface, is there a better place for them? + typedef void (*ggml_log_callback)(enum ggml_log_level level, const char * text, void * user_data); + + // Set callback for all future logging events. + // If this is not called, or NULL is supplied, everything is output on stderr. + GGML_API void ggml_log_get(ggml_log_callback * log_callback, void ** user_data); + GGML_API void ggml_log_set(ggml_log_callback log_callback, void * user_data); + + GGML_API struct ggml_tensor * ggml_set_zero(struct ggml_tensor * tensor); + + // + // quantization + // + + // - ggml_quantize_init can be called multiple times with the same type + // it will only initialize the quantization tables for the first call or after ggml_quantize_free + // automatically called by ggml_quantize_chunk for convenience + // + // - ggml_quantize_free will free any memory allocated by ggml_quantize_init + // call this at the end of the program to avoid memory leaks + // + // note: these are thread-safe + // + GGML_API void ggml_quantize_init(enum ggml_type type); + GGML_API void ggml_quantize_free(void); + + // some quantization type cannot be used without an importance matrix + GGML_API bool ggml_quantize_requires_imatrix(enum ggml_type type); + + // calls ggml_quantize_init internally (i.e. can allocate memory) + GGML_API size_t ggml_quantize_chunk( + enum ggml_type type, + const float * src, + void * dst, + int64_t start, + int64_t nrows, + int64_t n_per_row, + const float * imatrix); + +#ifdef __cplusplus + // restrict not standard in C++ +# if defined(__GNUC__) +# define GGML_RESTRICT __restrict__ +# elif defined(__clang__) +# define GGML_RESTRICT __restrict +# elif defined(_MSC_VER) +# define GGML_RESTRICT __restrict +# else +# define GGML_RESTRICT +# endif +#else +# if defined (_MSC_VER) && (__STDC_VERSION__ < 201112L) +# define GGML_RESTRICT __restrict +# else +# define GGML_RESTRICT restrict +# endif +#endif + typedef void (*ggml_to_float_t) (const void * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); + typedef void (*ggml_from_float_t)(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); + + struct ggml_type_traits { + const char * type_name; + int64_t blck_size; + int64_t blck_size_interleave; // interleave elements in blocks + size_t type_size; + bool is_quantized; + ggml_to_float_t to_float; + ggml_from_float_t from_float_ref; + }; + + GGML_API const struct ggml_type_traits * ggml_get_type_traits(enum ggml_type type); + + // ggml threadpool + // TODO: currently, only a few functions are in the base ggml API, while the rest are in the CPU backend + // the goal should be to create an API that other backends can use move everything to the ggml base + + // scheduling priorities + enum ggml_sched_priority { + GGML_SCHED_PRIO_LOW = -1, + GGML_SCHED_PRIO_NORMAL, + GGML_SCHED_PRIO_MEDIUM, + GGML_SCHED_PRIO_HIGH, + GGML_SCHED_PRIO_REALTIME + }; + + // threadpool params + // Use ggml_threadpool_params_default() or ggml_threadpool_params_init() to populate the defaults + struct ggml_threadpool_params { + bool cpumask[GGML_MAX_N_THREADS]; // mask of cpu cores (all-zeros means use default affinity settings) + int n_threads; // number of threads + enum ggml_sched_priority prio; // thread priority + uint32_t poll; // polling level (0 - no polling, 100 - aggressive polling) + bool strict_cpu; // strict cpu placement + bool paused; // start in paused state + }; + + struct ggml_threadpool; // forward declaration, see ggml.c + + typedef struct ggml_threadpool * ggml_threadpool_t; + + GGML_API struct ggml_threadpool_params ggml_threadpool_params_default(int n_threads); + GGML_API void ggml_threadpool_params_init (struct ggml_threadpool_params * p, int n_threads); + GGML_API bool ggml_threadpool_params_match (const struct ggml_threadpool_params * p0, const struct ggml_threadpool_params * p1); + +#ifdef __cplusplus +} +#endif diff --git a/python/freetoken/kernel/csrc/ggml/include/gguf.h b/python/freetoken/kernel/csrc/ggml/include/gguf.h new file mode 100644 index 0000000000..b3a1e1230a --- /dev/null +++ b/python/freetoken/kernel/csrc/ggml/include/gguf.h @@ -0,0 +1,211 @@ +// This file contains functionality related to "GGUF" files, the binary file format used by ggml. +// GGUF files have the following structure: +// +// 1. File magic "GGUF" (4 bytes). +// 2. File version (uint32_t). +// 3. Number of ggml tensors in file (int64_t). +// 4. Number of key-value-pairs in file (int64_t). +// 5. For each KV pair: +// 1. The key (string). +// 2. The value type (gguf_type). +// 3a. If the value type is GGUF_TYPE_ARRAY: +// 1. The type of the array (gguf_type). +// 2. The number of elements in the array (uint64_t). +// 3. The binary representation of each element in the array. +// 3b. Otherwise: +// 1. The binary representation of the value. +// 6. For each ggml tensor: +// 1. The tensor name (string). +// 2. The number of dimensions of the tensor (uint32_t). +// 3. For each dimension: +// 1. The size of the tensor in the dimension (int64_t). +// 4. The tensor data type (ggml_type). +// 5. The tensor data offset in the tensor data binary blob (uint64_t). +// 7. The tensor data binary blob (optional, aligned). +// +// Strings are serialized as the string length (uint64_t) followed by the C string without the null terminator. +// All enums are stored as int32_t. +// All bool values are stored as int8_t. +// If the special key "general.alignment" (uint32_t) is defined it is used for alignment, +// otherwise GGUF_DEFAULT_ALIGNMENT is used. +// +// Module maintainer: Johannes Gäßler (@JohannesGaessler, johannesg@5d6.de) + +#pragma once + +#include "ggml.h" + +#include +#include + +#define GGUF_MAGIC "GGUF" +#define GGUF_VERSION 3 + +#define GGUF_KEY_GENERAL_ALIGNMENT "general.alignment" + +#define GGUF_DEFAULT_ALIGNMENT 32 + +#ifdef __cplusplus +extern "C" { +#endif + + // types that can be stored as GGUF KV data + enum gguf_type { + GGUF_TYPE_UINT8 = 0, + GGUF_TYPE_INT8 = 1, + GGUF_TYPE_UINT16 = 2, + GGUF_TYPE_INT16 = 3, + GGUF_TYPE_UINT32 = 4, + GGUF_TYPE_INT32 = 5, + GGUF_TYPE_FLOAT32 = 6, + GGUF_TYPE_BOOL = 7, + GGUF_TYPE_STRING = 8, + GGUF_TYPE_ARRAY = 9, + GGUF_TYPE_UINT64 = 10, + GGUF_TYPE_INT64 = 11, + GGUF_TYPE_FLOAT64 = 12, + GGUF_TYPE_COUNT, // marks the end of the enum + }; + + struct gguf_context; + + struct gguf_init_params { + bool no_alloc; + + // if not NULL, create a ggml_context and allocate the tensor data in it + struct ggml_context ** ctx; + }; + + // callback to simulate or wrap a FILE pointer - read up to `len` bytes at `offset` into `output` and return the number of bytes read + typedef size_t (*gguf_reader_callback_t)(void * userdata, void * output, uint64_t offset, size_t len); + + GGML_API struct gguf_context * gguf_init_empty(void); + GGML_API struct gguf_context * gguf_init_from_file_ptr(FILE * file, struct gguf_init_params params); + GGML_API struct gguf_context * gguf_init_from_file(const char * fname, struct gguf_init_params params); + GGML_API struct gguf_context * gguf_init_from_buffer(const void * data, size_t size, struct gguf_init_params params); + + // max_chunk_read is the maximum number of bytes that the GGUF code will read at once from the callback, a value of 0 means no limit + GGML_API struct gguf_context * gguf_init_from_callback(gguf_reader_callback_t callback, void * userdata, size_t max_chunk_read, uint64_t max_expected_size, struct gguf_init_params params); + + GGML_API void gguf_free(struct gguf_context * ctx); + + GGML_API const char * gguf_type_name(enum gguf_type type); + + GGML_API uint32_t gguf_get_version (const struct gguf_context * ctx); + GGML_API size_t gguf_get_alignment (const struct gguf_context * ctx); + GGML_API size_t gguf_get_data_offset(const struct gguf_context * ctx); // padded to gguf_get_alignment if and only if the gguf_context contains at least one tensor + + GGML_API int64_t gguf_get_n_kv(const struct gguf_context * ctx); + GGML_API int64_t gguf_find_key(const struct gguf_context * ctx, const char * key); // returns -1 if key is not found + GGML_API const char * gguf_get_key (const struct gguf_context * ctx, int64_t key_id); + + GGML_API enum gguf_type gguf_get_kv_type (const struct gguf_context * ctx, int64_t key_id); + GGML_API enum gguf_type gguf_get_arr_type(const struct gguf_context * ctx, int64_t key_id); + + // will abort if the wrong type is used for the key + GGML_API uint8_t gguf_get_val_u8 (const struct gguf_context * ctx, int64_t key_id); + GGML_API int8_t gguf_get_val_i8 (const struct gguf_context * ctx, int64_t key_id); + GGML_API uint16_t gguf_get_val_u16 (const struct gguf_context * ctx, int64_t key_id); + GGML_API int16_t gguf_get_val_i16 (const struct gguf_context * ctx, int64_t key_id); + GGML_API uint32_t gguf_get_val_u32 (const struct gguf_context * ctx, int64_t key_id); + GGML_API int32_t gguf_get_val_i32 (const struct gguf_context * ctx, int64_t key_id); + GGML_API float gguf_get_val_f32 (const struct gguf_context * ctx, int64_t key_id); + GGML_API uint64_t gguf_get_val_u64 (const struct gguf_context * ctx, int64_t key_id); + GGML_API int64_t gguf_get_val_i64 (const struct gguf_context * ctx, int64_t key_id); + GGML_API double gguf_get_val_f64 (const struct gguf_context * ctx, int64_t key_id); + GGML_API bool gguf_get_val_bool(const struct gguf_context * ctx, int64_t key_id); + GGML_API const char * gguf_get_val_str (const struct gguf_context * ctx, int64_t key_id); + GGML_API const void * gguf_get_val_data(const struct gguf_context * ctx, int64_t key_id); + GGML_API size_t gguf_get_arr_n (const struct gguf_context * ctx, int64_t key_id); + + // get raw pointer to the first element of the array with the given key_id + // for bool arrays, note that they are always stored as int8 on all platforms (usually this makes no difference) + GGML_API const void * gguf_get_arr_data(const struct gguf_context * ctx, int64_t key_id); + + // get ith C string from array with given key_id + GGML_API const char * gguf_get_arr_str (const struct gguf_context * ctx, int64_t key_id, size_t i); + + GGML_API int64_t gguf_get_n_tensors (const struct gguf_context * ctx); + GGML_API int64_t gguf_find_tensor (const struct gguf_context * ctx, const char * name); // returns -1 if the tensor is not found + GGML_API size_t gguf_get_tensor_offset(const struct gguf_context * ctx, int64_t tensor_id); + GGML_API const char * gguf_get_tensor_name (const struct gguf_context * ctx, int64_t tensor_id); + GGML_API const int64_t * gguf_get_tensor_ne (const struct gguf_context * ctx, int64_t tensor_id); // returns ne, an array of GGML_MAX_DIMS elements; ne[dim] is 1 for dim >= n_dims + GGML_API enum ggml_type gguf_get_tensor_type (const struct gguf_context * ctx, int64_t tensor_id); + GGML_API size_t gguf_get_tensor_size (const struct gguf_context * ctx, int64_t tensor_id); + + // removes key if it exists, returns id that the key had prior to removal (-1 if it didn't exist) + GGML_API int64_t gguf_remove_key(struct gguf_context * ctx, const char * key); + + // overrides an existing KV pair or adds a new one, the new KV pair is always at the back + GGML_API void gguf_set_val_u8 (struct gguf_context * ctx, const char * key, uint8_t val); + GGML_API void gguf_set_val_i8 (struct gguf_context * ctx, const char * key, int8_t val); + GGML_API void gguf_set_val_u16 (struct gguf_context * ctx, const char * key, uint16_t val); + GGML_API void gguf_set_val_i16 (struct gguf_context * ctx, const char * key, int16_t val); + GGML_API void gguf_set_val_u32 (struct gguf_context * ctx, const char * key, uint32_t val); + GGML_API void gguf_set_val_i32 (struct gguf_context * ctx, const char * key, int32_t val); + GGML_API void gguf_set_val_f32 (struct gguf_context * ctx, const char * key, float val); + GGML_API void gguf_set_val_u64 (struct gguf_context * ctx, const char * key, uint64_t val); + GGML_API void gguf_set_val_i64 (struct gguf_context * ctx, const char * key, int64_t val); + GGML_API void gguf_set_val_f64 (struct gguf_context * ctx, const char * key, double val); + GGML_API void gguf_set_val_bool(struct gguf_context * ctx, const char * key, bool val); + GGML_API void gguf_set_val_str (struct gguf_context * ctx, const char * key, const char * val); + + // creates a new array with n elements of the given type and copies the corresponding number of bytes from data + GGML_API void gguf_set_arr_data(struct gguf_context * ctx, const char * key, enum gguf_type type, const void * data, size_t n); + + // creates a new array with n strings and copies the corresponding strings from data + GGML_API void gguf_set_arr_str (struct gguf_context * ctx, const char * key, const char ** data, size_t n); + + // set or add KV pairs from another context + GGML_API void gguf_set_kv(struct gguf_context * ctx, const struct gguf_context * src); + + // add tensor to GGUF context, tensor name must be unique + GGML_API void gguf_add_tensor(struct gguf_context * ctx, const struct ggml_tensor * tensor); + + // after changing a tensor's type, the offsets of all tensors with higher indices are immediately recalculated + // in such a way that the tensor data remains as one contiguous block (except for padding) + GGML_API void gguf_set_tensor_type(struct gguf_context * ctx, const char * name, enum ggml_type type); + + // assumes that at least gguf_get_tensor_size bytes can be read from data + GGML_API void gguf_set_tensor_data(struct gguf_context * ctx, const char * name, const void * data); + + // writing gguf files can be done in 3 ways: + // + // - write the entire gguf_context to a binary file in a single pass: + // + // gguf_write_to_file(ctx, fname, /*only_meta =*/ false); + // + // - write only the meta data to a file, then re-open the file and append the tensor data: + // + // gguf_write_to_file(ctx, fname, /*only_meta =*/ true); + // FILE * f = fopen(fname, "ab"); + // fwrite(f, ...); // write tensor data + // fclose(f); + // + // - first prepare a file with a placeholder for the meta data, write the tensor data, then write the meta data: + // + // FILE * f = fopen(fname, "wb"); + // const size_t size_meta = gguf_get_meta_size(ctx); + // fseek(f, size_meta, SEEK_SET); + // fwrite(f, ...); // write tensor data + // void * data = malloc(size_meta); + // gguf_get_meta_data(ctx, data); + // rewind(f); + // fwrite(data, 1, data, f); + // free(data); + // fclose(f); + // + + // write the entire context to a binary file + GGML_API bool gguf_write_to_file_ptr(const struct gguf_context * ctx, FILE * file, bool only_meta); + GGML_API bool gguf_write_to_file(const struct gguf_context * ctx, const char * fname, bool only_meta); + + // get the size in bytes of the meta data (header, kv pairs, tensor info) including padding + GGML_API size_t gguf_get_meta_size(const struct gguf_context * ctx); + + // writes the meta data to pointer "data" + GGML_API void gguf_get_meta_data(const struct gguf_context * ctx, void * data); + +#ifdef __cplusplus +} +#endif diff --git a/python/freetoken/kernel/csrc/ggml/src/ggml-backend-impl.h b/python/freetoken/kernel/csrc/ggml/src/ggml-backend-impl.h new file mode 100644 index 0000000000..dc7367dec4 --- /dev/null +++ b/python/freetoken/kernel/csrc/ggml/src/ggml-backend-impl.h @@ -0,0 +1,292 @@ +#pragma once + +// ggml-backend internal header + +#include "ggml-backend.h" + +#ifdef __cplusplus +extern "C" { +#endif + + #define GGML_BACKEND_API_VERSION 3 + + // + // Backend buffer type + // + + struct ggml_backend_buffer_type_i { + const char * (*get_name) (ggml_backend_buffer_type_t buft); + // allocate a buffer of this type + ggml_backend_buffer_t (*alloc_buffer) (ggml_backend_buffer_type_t buft, size_t size); + // (optional) allocate tensors from a list into a buffer of this type (defaults to alloc_buffer + linear allocator) + ggml_backend_buffer_t (*alloc_buffer_n) (ggml_backend_buffer_type_t buft, struct ggml_tensor ** tensors, int n_tensors); + // tensor alignment + size_t (*get_alignment) (ggml_backend_buffer_type_t buft); + // (optional) max buffer size that can be allocated (defaults to SIZE_MAX) + size_t (*get_max_size) (ggml_backend_buffer_type_t buft); + // (optional) data size needed to allocate the tensor, including padding (defaults to ggml_nbytes) + size_t (*get_alloc_size) (ggml_backend_buffer_type_t buft, const struct ggml_tensor * tensor); + // (optional) total data size needed to allocate the given tensors, including padding and splitting (defaults to per-tensor get_alloc_size) + size_t (*get_alloc_size_n)(ggml_backend_buffer_type_t buft, struct ggml_tensor ** tensors, int n_tensors); + // (optional) check if tensor data is in host memory and uses standard ggml tensor layout (defaults to false) + bool (*is_host) (ggml_backend_buffer_type_t buft); + }; + + struct ggml_backend_buffer_type { + struct ggml_backend_buffer_type_i iface; + ggml_backend_dev_t device; + void * context; + }; + + // [TAG_ALLOC_SIZE_EXPAND] + // returns true for ops that may require additional memory for fleeting data on some backends, + // i.e. the backend buffer type's get_alloc_size may return more than ggml_nbytes for the output tensor + GGML_API bool ggml_op_alloc_size_may_expand(enum ggml_op op); + + // + // Backend buffer + // + + struct ggml_backend_buffer_i { + // (optional) free the buffer + void (*free_buffer) (ggml_backend_buffer_t buffer); + // base address of the buffer + void * (*get_base) (ggml_backend_buffer_t buffer); + // (optional) initialize a tensor in the buffer (eg. add tensor extras) + enum ggml_status (*init_tensor)(ggml_backend_buffer_t buffer, struct ggml_tensor * tensor); + // tensor data access + void (*memset_tensor)(ggml_backend_buffer_t buffer, struct ggml_tensor * tensor, uint8_t value, size_t offset, size_t size); + void (*set_tensor) (ggml_backend_buffer_t buffer, struct ggml_tensor * tensor, const void * data, size_t offset, size_t size); + void (*get_tensor) (ggml_backend_buffer_t buffer, const struct ggml_tensor * tensor, void * data, size_t offset, size_t size); + // (optional) 2d data copies + void (*set_tensor_2d)(ggml_backend_buffer_t buffer, struct ggml_tensor * tensor, const void * data, size_t offset, size_t size, size_t n_copies, size_t stride_tensor, size_t stride_data); + void (*get_tensor_2d)(ggml_backend_buffer_t buffer, const struct ggml_tensor * tensor, void * data, size_t offset, size_t size, size_t n_copies, size_t stride_tensor, size_t stride_data); + + // (optional) tensor copy: dst is in the buffer, src may be in any buffer, including buffers from a different backend (return false if not supported) + bool (*cpy_tensor) (ggml_backend_buffer_t buffer, const struct ggml_tensor * src, struct ggml_tensor * dst); + // clear the entire buffer + void (*clear) (ggml_backend_buffer_t buffer, uint8_t value); + // (optional) reset any internal state due to tensor initialization, such as tensor extras + void (*reset) (ggml_backend_buffer_t buffer); + }; + + struct ggml_backend_buffer { + struct ggml_backend_buffer_i iface; + ggml_backend_buffer_type_t buft; + void * context; + size_t size; + enum ggml_backend_buffer_usage usage; + }; + + GGML_API ggml_backend_buffer_t ggml_backend_buffer_init( + ggml_backend_buffer_type_t buft, + struct ggml_backend_buffer_i iface, + void * context, + size_t size); + + // do not use directly, use ggml_backend_tensor_copy instead + GGML_API bool ggml_backend_buffer_copy_tensor(const struct ggml_tensor * src, struct ggml_tensor * dst); + + // multi-buffer + // buffer that contains a collection of buffers + GGML_API ggml_backend_buffer_t ggml_backend_multi_buffer_alloc_buffer(ggml_backend_buffer_t * buffers, size_t n_buffers); + GGML_API bool ggml_backend_buffer_is_multi_buffer(ggml_backend_buffer_t buffer); + GGML_API void ggml_backend_multi_buffer_set_usage(ggml_backend_buffer_t buffer, enum ggml_backend_buffer_usage usage); + GGML_API void ggml_backend_meta_buffer_set_usage (ggml_backend_buffer_t buffer, enum ggml_backend_buffer_usage usage); + + // + // Backend (meta) + // + + GGML_API bool ggml_backend_is_meta (ggml_backend_t backend); + GGML_API bool ggml_backend_buffer_is_meta(ggml_backend_buffer_t buf); + GGML_API bool ggml_backend_buft_is_meta (ggml_backend_buffer_type_t buft); + + GGML_API size_t ggml_backend_meta_n_backends (ggml_backend_t meta_backend); + GGML_API ggml_backend_t ggml_backend_meta_simple_backend(ggml_backend_t meta_backend, size_t index); + + // + // Backend (stream) + // + + // passed to graph_optimize so the backend can add allocation dependencies: + // if the backend executes parts of the graph out of order (e.g. on concurrent streams), + // it must keep the affected tensors allocated until a node where execution is known to have joined + struct ggml_backend_graph_optimize_params { + // keep `tensor` allocated at least until `until` (a node of the same graph) has been computed + // can be called multiple times for the same tensor: the longest lifetime applies + void (*add_alloc_dep)(void * user_data, struct ggml_tensor * tensor, struct ggml_tensor * until); + void * user_data; + }; + + struct ggml_backend_i { + const char * (*get_name)(ggml_backend_t backend); + + void (*free)(ggml_backend_t backend); + + // (optional) asynchronous tensor data access + void (*set_tensor_async) (ggml_backend_t backend, struct ggml_tensor * tensor, const void * data, size_t offset, size_t size); + void (*get_tensor_async) (ggml_backend_t backend, const struct ggml_tensor * tensor, void * data, size_t offset, size_t size); + void (*set_tensor_2d_async)(ggml_backend_t backend, struct ggml_tensor * tensor, const void * data, size_t offset, size_t size, size_t n_copies, size_t stride_tensor, size_t stride_data); + void (*get_tensor_2d_async)(ggml_backend_t backend, const struct ggml_tensor * tensor, void * data, size_t offset, size_t size, size_t n_copies, size_t stride_tensor, size_t stride_data); + bool (*cpy_tensor_async)(ggml_backend_t backend_src, ggml_backend_t backend_dst, const struct ggml_tensor * src, struct ggml_tensor * dst); + + // (optional) complete all pending operations (required if the backend supports async operations) + void (*synchronize)(ggml_backend_t backend); + + // (optional) graph plans (not used currently) + // compute graph with a plan + ggml_backend_graph_plan_t (*graph_plan_create) (ggml_backend_t backend, const struct ggml_cgraph * cgraph); + void (*graph_plan_free) (ggml_backend_t backend, ggml_backend_graph_plan_t plan); + // update the plan with a new graph - this should be faster than creating a new plan when the graph has the same topology + void (*graph_plan_update) (ggml_backend_t backend, ggml_backend_graph_plan_t plan, const struct ggml_cgraph * cgraph); + // compute the graph with the plan + enum ggml_status (*graph_plan_compute)(ggml_backend_t backend, ggml_backend_graph_plan_t plan); + + // compute graph (always async if supported by the backend) + enum ggml_status (*graph_compute) (ggml_backend_t backend, struct ggml_cgraph * cgraph); + + // (optional) event synchronization + // record an event on this stream + void (*event_record)(ggml_backend_t backend, ggml_backend_event_t event); + // wait for an event on on a different stream + void (*event_wait) (ggml_backend_t backend, ggml_backend_event_t event); + + // (optional) sort/optimize the nodes in the graph + void (*graph_optimize) (ggml_backend_t backend, struct ggml_cgraph * cgraph, struct ggml_backend_graph_optimize_params * params); + }; + + struct ggml_backend { + ggml_guid_t guid; + struct ggml_backend_i iface; + ggml_backend_dev_t device; + void * context; + }; + + struct ggml_backend_event { + struct ggml_backend_device * device; + void * context; + }; + + // + // Backend device + // + + // Note: if additional properties are needed, we should add a struct with all of them + // the current functions to obtain the properties can remain, since they are more convenient for often used properties + struct ggml_backend_device_i { + // device name: short identifier for this device, such as "CPU" or "CUDA0" + const char * (*get_name)(ggml_backend_dev_t dev); + + // device description: short informative description of the device, could be the model name + const char * (*get_description)(ggml_backend_dev_t dev); + + // device memory in bytes: 0 bytes to indicate no memory to report + void (*get_memory)(ggml_backend_dev_t dev, size_t * free, size_t * total); + + // device type + enum ggml_backend_dev_type (*get_type)(ggml_backend_dev_t dev); + + // device properties + void (*get_props)(ggml_backend_dev_t dev, struct ggml_backend_dev_props * props); + + // backend (stream) initialization + ggml_backend_t (*init_backend)(ggml_backend_dev_t dev, const char * params); + + // preferred buffer type + ggml_backend_buffer_type_t (*get_buffer_type)(ggml_backend_dev_t dev); + + // (optional) host buffer type (in system memory, typically this is a pinned memory buffer for faster transfers between host and device) + ggml_backend_buffer_type_t (*get_host_buffer_type)(ggml_backend_dev_t dev); + + // (optional) buffer from pointer: create a buffer from a host pointer (useful for memory mapped models and importing data from other libraries) + ggml_backend_buffer_t (*buffer_from_host_ptr)(ggml_backend_dev_t dev, void * ptr, size_t size, size_t max_tensor_size); + + // check if the backend can compute an operation + bool (*supports_op)(ggml_backend_dev_t dev, const struct ggml_tensor * op); + + // check if the backend can use tensors allocated in a buffer type + bool (*supports_buft)(ggml_backend_dev_t dev, ggml_backend_buffer_type_t buft); + + // (optional) check if the backend wants to run an operation, even if the weights are allocated in an incompatible buffer + // these should be expensive operations that may benefit from running on this backend instead of the CPU backend + bool (*offload_op)(ggml_backend_dev_t dev, const struct ggml_tensor * op); + + // (optional) event synchronization + ggml_backend_event_t (*event_new) (ggml_backend_dev_t dev); + void (*event_free) (ggml_backend_dev_t dev, ggml_backend_event_t event); + void (*event_synchronize) (ggml_backend_dev_t dev, ggml_backend_event_t event); + }; + + struct ggml_backend_device { + struct ggml_backend_device_i iface; + ggml_backend_reg_t reg; + void * context; + }; + + // + // Backend (reg) + // + + struct ggml_backend_reg_i { + const char * (*get_name)(ggml_backend_reg_t reg); + + // enumerate available devices + size_t (*get_device_count)(ggml_backend_reg_t reg); + ggml_backend_dev_t (*get_device)(ggml_backend_reg_t reg, size_t index); + + // (optional) get a pointer to a function in the backend + // backends can add custom functions that are not part of the standard ggml-backend interface + void * (*get_proc_address)(ggml_backend_reg_t reg, const char * name); + }; + + struct ggml_backend_reg { + int api_version; // initialize to GGML_BACKEND_API_VERSION + struct ggml_backend_reg_i iface; + void * context; + }; + + // Add backend dynamic loading support to the backend + + // Initialize the backend + typedef ggml_backend_reg_t (*ggml_backend_init_t)(void); + // Optional: obtain a score for the backend based on the system configuration + // Higher scores are preferred, 0 means the backend is not supported in the current system + typedef int (*ggml_backend_score_t)(void); + +#ifdef GGML_BACKEND_DL +# ifdef __cplusplus +# define GGML_BACKEND_DL_IMPL(reg_fn) \ + extern "C" { \ + GGML_BACKEND_API ggml_backend_reg_t ggml_backend_init(void); \ + } \ + ggml_backend_reg_t ggml_backend_init(void) { \ + return reg_fn(); \ + } +# define GGML_BACKEND_DL_SCORE_IMPL(score_fn) \ + extern "C" { \ + GGML_BACKEND_API int ggml_backend_score(void); \ + } \ + int ggml_backend_score(void) { \ + return score_fn(); \ + } +# else +# define GGML_BACKEND_DL_IMPL(reg_fn) \ + GGML_BACKEND_API ggml_backend_reg_t ggml_backend_init(void); \ + ggml_backend_reg_t ggml_backend_init(void) { \ + return reg_fn(); \ + } +# define GGML_BACKEND_DL_SCORE_IMPL(score_fn) \ + GGML_BACKEND_API int ggml_backend_score(void); \ + int ggml_backend_score(void) { \ + return score_fn(); \ + } +# endif +#else +# define GGML_BACKEND_DL_IMPL(reg_fn) +# define GGML_BACKEND_DL_SCORE_IMPL(score_fn) +#endif + +#ifdef __cplusplus +} +#endif diff --git a/python/freetoken/kernel/csrc/ggml/src/ggml-common.h b/python/freetoken/kernel/csrc/ggml/src/ggml-common.h new file mode 100644 index 0000000000..1dbbe326d0 --- /dev/null +++ b/python/freetoken/kernel/csrc/ggml/src/ggml-common.h @@ -0,0 +1,1911 @@ +#ifndef GGML_COMMON_DECL + +#if defined(GGML_COMMON_DECL_C) +#include + +typedef uint16_t ggml_half; +typedef uint32_t ggml_half2; + +#define GGML_COMMON_AGGR_U +#define GGML_COMMON_AGGR_S + +#define GGML_COMMON_DECL +#elif defined(GGML_COMMON_DECL_CPP) +#include + +typedef uint16_t ggml_half; +typedef uint32_t ggml_half2; + +// std-c++ allow anonymous unions but some compiler warn on it +#define GGML_COMMON_AGGR_U data +// std-c++ do not allow it. +#define GGML_COMMON_AGGR_S data + +#define GGML_COMMON_DECL +#elif defined(GGML_COMMON_DECL_METAL) +#include + +typedef half ggml_half; +typedef half2 ggml_half2; + +#define GGML_COMMON_AGGR_U +#define GGML_COMMON_AGGR_S + +#define GGML_COMMON_DECL +#elif defined(GGML_COMMON_DECL_CUDA) +#if defined(GGML_COMMON_DECL_MUSA) +#include +#else +#include +#endif +#include + +typedef half ggml_half; +typedef half2 ggml_half2; + +#define GGML_COMMON_AGGR_U +#define GGML_COMMON_AGGR_S data + +#define GGML_COMMON_DECL +#elif defined(GGML_COMMON_DECL_HIP) +#include +#include + +typedef half ggml_half; +typedef half2 ggml_half2; + +#define GGML_COMMON_AGGR_U +#define GGML_COMMON_AGGR_S data + +#define GGML_COMMON_DECL +#elif defined(GGML_COMMON_DECL_SYCL) +#include +#include + +typedef sycl::half ggml_half; +typedef sycl::half2 ggml_half2; + +#define GGML_COMMON_AGGR_U +#define GGML_COMMON_AGGR_S data + +#define GGML_COMMON_DECL +#endif + +#if defined(GGML_COMMON_DECL) + +#ifndef __cplusplus +#ifndef static_assert +#if defined(__STDC_VERSION__) && (__STDC_VERSION__ >= 201100L) +#define static_assert(cond, msg) _Static_assert(cond, msg) +#else +#define static_assert(cond, msg) struct global_scope_noop_trick +#endif +#endif +#endif // __cplusplus + +// QK = number of values after dequantization +// QK_K = super-block size + +#define QK_K 256 +#define K_SCALE_SIZE 12 + +#if defined(GGML_COMMON_DECL_CUDA) || defined(GGML_COMMON_DECL_HIP) || defined(GGML_COMMON_DECL_SYCL) +// QR = QK / number of values before dequantization +// QI = number of 32 bit integers before dequantization + +#define QI1_0 (QK1_0 / 32) +#define QR1_0 1 + +#define QI2_0 (QK2_0 / 32) +#define QR2_0 1 + + +#define QI4_0 (QK4_0 / (4 * QR4_0)) +#define QR4_0 2 + +#define QI4_1 (QK4_1 / (4 * QR4_1)) +#define QR4_1 2 + +#define QI_MXFP4 (QK_MXFP4 / (4 * QR_MXFP4)) +#define QR_MXFP4 2 + +#define QI_NVFP4 (QK_NVFP4 / (4 * QR_NVFP4)) +#define QR_NVFP4 2 + +#define QI5_0 (QK5_0 / (4 * QR5_0)) +#define QR5_0 2 + +#define QI5_1 (QK5_1 / (4 * QR5_1)) +#define QR5_1 2 + +#define QI8_0 (QK8_0 / (4 * QR8_0)) +#define QR8_0 1 + +#define QI8_1 (QK8_1 / (4 * QR8_1)) +#define QR8_1 1 + +#define QI2_K (QK_K / (4*QR2_K)) +#define QR2_K 4 + +#define QI3_K (QK_K / (4*QR3_K)) +#define QR3_K 4 + +#define QI4_K (QK_K / (4*QR4_K)) +#define QR4_K 2 + +#define QI5_K (QK_K / (4*QR5_K)) +#define QR5_K 2 + +#define QI6_K (QK_K / (4*QR6_K)) +#define QR6_K 2 + +#define QI2_XXS (QK_K / (4*QR2_XXS)) +#define QR2_XXS 4 + +#define QI2_XS (QK_K / (4*QR2_XS)) +#define QR2_XS 4 + +#define QI2_S (QK_K / (4*QR2_S)) +#define QR2_S 4 + +#define QI3_XXS (QK_K / (4*QR3_XXS)) +#define QR3_XXS 4 + +#define QI3_XS (QK_K / (4*QR3_XS)) +#define QR3_XS 4 + +#define QI1_S (QK_K / (4*QR1_S)) +#define QR1_S 8 + +#define QI1_M (QK_K / (4*QR1_M)) +#define QR1_M 8 + +#define QI4_NL (QK4_NL / (4*QR4_NL)) +#define QR4_NL 2 + +#define QI4_XS (QK_K / (4*QR4_XS)) +#define QR4_XS 2 + +#define QI3_S (QK_K / (4*QR3_S)) +#define QR3_S 4 + +#endif // GGML_COMMON_DECL_CUDA || GGML_COMMON_DECL_HIP + +#ifdef _MSC_VER +#define GGML_EXTENSION +#else // _MSC_VER +#define GGML_EXTENSION __extension__ +#endif // _MSC_VER + +#define QK1_0 128 +typedef struct { + ggml_half d; // delta + uint8_t qs[QK1_0 / 8]; // bits / quants +} block_q1_0; +static_assert(sizeof(block_q1_0) == sizeof(ggml_half) + QK1_0 / 8, "wrong q1_0 block size/padding"); + +#define QK2_0 64 +typedef struct { + ggml_half d; // delta (scale) + uint8_t qs[QK2_0 / 4]; // 2 bits per element +} block_q2_0; +static_assert(sizeof(block_q2_0) == sizeof(ggml_half) + QK2_0 / 4, "wrong q2_0 block size/padding"); + +#define QK4_0 32 +typedef struct { + ggml_half d; // delta + uint8_t qs[QK4_0 / 2]; // nibbles / quants +} block_q4_0; +static_assert(sizeof(block_q4_0) == sizeof(ggml_half) + QK4_0 / 2, "wrong q4_0 block size/padding"); + +#define QK4_1 32 +typedef struct { + GGML_EXTENSION union { + struct { + ggml_half d; // delta + ggml_half m; // min + } GGML_COMMON_AGGR_S; + ggml_half2 dm; + } GGML_COMMON_AGGR_U; + uint8_t qs[QK4_1 / 2]; // nibbles / quants +} block_q4_1; +static_assert(sizeof(block_q4_1) == 2 * sizeof(ggml_half) + QK4_1 / 2, "wrong q4_1 block size/padding"); + +#define QK_MXFP4 32 +typedef struct { + uint8_t e; // E8M0 + uint8_t qs[QK_MXFP4/2]; +} block_mxfp4; +static_assert(sizeof(block_mxfp4) == sizeof(uint8_t) + QK_MXFP4/2, "wrong mxfp4 block size/padding"); + +#define QK_NVFP4 64 +#define QK_NVFP4_SUB 16 // sub-block size for per-group scales +typedef struct { + uint8_t d[QK_NVFP4/QK_NVFP4_SUB]; // UE4M3 scales (4 bytes, one per 16-element sub-block) + uint8_t qs[QK_NVFP4/2]; // packed 4-bit E2M1 values (32 bytes) +} block_nvfp4; +static_assert(sizeof(block_nvfp4) == sizeof(uint8_t)*(QK_NVFP4/QK_NVFP4_SUB) + QK_NVFP4/2, "wrong nvfp4 block size/padding"); + +#define QK5_0 32 +typedef struct { + ggml_half d; // delta + uint8_t qh[4]; // 5-th bit of quants + uint8_t qs[QK5_0 / 2]; // nibbles / quants +} block_q5_0; +static_assert(sizeof(block_q5_0) == sizeof(ggml_half) + sizeof(uint32_t) + QK5_0 / 2, "wrong q5_0 block size/padding"); + +#define QK5_1 32 +typedef struct { + GGML_EXTENSION union { + struct { + ggml_half d; // delta + ggml_half m; // min + } GGML_COMMON_AGGR_S; + ggml_half2 dm; + } GGML_COMMON_AGGR_U; + uint8_t qh[4]; // 5-th bit of quants + uint8_t qs[QK5_1 / 2]; // nibbles / quants +} block_q5_1; +static_assert(sizeof(block_q5_1) == 2 * sizeof(ggml_half) + sizeof(uint32_t) + QK5_1 / 2, "wrong q5_1 block size/padding"); + +#define QK8_0 32 +typedef struct { + ggml_half d; // delta + int8_t qs[QK8_0]; // quants +} block_q8_0; +static_assert(sizeof(block_q8_0) == sizeof(ggml_half) + QK8_0, "wrong q8_0 block size/padding"); + +#define QK8_1 32 +typedef struct { + GGML_EXTENSION union { + struct { + ggml_half d; // delta + ggml_half s; // d * sum(qs[i]) + } GGML_COMMON_AGGR_S; + ggml_half2 ds; + } GGML_COMMON_AGGR_U; + int8_t qs[QK8_1]; // quants +} block_q8_1; +static_assert(sizeof(block_q8_1) == 2*sizeof(ggml_half) + QK8_1, "wrong q8_1 block size/padding"); + +// +// Ternary quantization +// + +// 1.6875 bpw +typedef struct { + uint8_t qs[(QK_K - 4 * QK_K / 64) / 5]; // 5 elements per byte (3^5 = 243 < 256) + uint8_t qh[QK_K/64]; // 4 elements per byte + ggml_half d; +} block_tq1_0; +static_assert(sizeof(block_tq1_0) == sizeof(ggml_half) + QK_K / 64 + (QK_K - 4 * QK_K / 64) / 5, "wrong tq1_0 block size/padding"); + +// 2.0625 bpw +typedef struct { + uint8_t qs[QK_K/4]; // 2 bits per element + ggml_half d; +} block_tq2_0; +static_assert(sizeof(block_tq2_0) == sizeof(ggml_half) + QK_K / 4, "wrong tq2_0 block size/padding"); + +// +// Super-block quantization structures +// + +// 2-bit quantization +// weight is represented as x = a * q + b +// 16 blocks of 16 elements each +// Effectively 2.625 bits per weight +typedef struct { + uint8_t scales[QK_K/16]; // scales and mins, quantized with 4 bits + uint8_t qs[QK_K/4]; // quants + GGML_EXTENSION union { + struct { + ggml_half d; // super-block scale for quantized scales + ggml_half dmin; // super-block scale for quantized mins + } GGML_COMMON_AGGR_S; + ggml_half2 dm; + } GGML_COMMON_AGGR_U; +} block_q2_K; +static_assert(sizeof(block_q2_K) == 2*sizeof(ggml_half) + QK_K/16 + QK_K/4, "wrong q2_K block size/padding"); + +// 3-bit quantization +// weight is represented as x = a * q +// 16 blocks of 16 elements each +// Effectively 3.4375 bits per weight +typedef struct { + uint8_t hmask[QK_K/8]; // quants - high bit + uint8_t qs[QK_K/4]; // quants - low 2 bits + uint8_t scales[12]; // scales, quantized with 6 bits + ggml_half d; // super-block scale +} block_q3_K; +static_assert(sizeof(block_q3_K) == sizeof(ggml_half) + QK_K / 4 + QK_K / 8 + 12, "wrong q3_K block size/padding"); + +// 4-bit quantization +// 8 blocks of 32 elements each +// weight is represented as x = a * q + b +// Effectively 4.5 bits per weight +typedef struct { + GGML_EXTENSION union { + struct { + ggml_half d; // super-block scale for quantized scales + ggml_half dmin; // super-block scale for quantized mins + } GGML_COMMON_AGGR_S; + ggml_half2 dm; + } GGML_COMMON_AGGR_U; + uint8_t scales[K_SCALE_SIZE]; // scales and mins, quantized with 6 bits + uint8_t qs[QK_K/2]; // 4--bit quants +} block_q4_K; +static_assert(sizeof(block_q4_K) == 2*sizeof(ggml_half) + K_SCALE_SIZE + QK_K/2, "wrong q4_K block size/padding"); + +// 5-bit quantization +// 8 blocks of 32 elements each +// weight is represented as x = a * q + b +// Effectively 5.5 bits per weight +typedef struct { + GGML_EXTENSION union { + struct { + ggml_half d; // super-block scale for quantized scales + ggml_half dmin; // super-block scale for quantized mins + } GGML_COMMON_AGGR_S; + ggml_half2 dm; + } GGML_COMMON_AGGR_U; + uint8_t scales[K_SCALE_SIZE]; // scales and mins, quantized with 6 bits + uint8_t qh[QK_K/8]; // quants, high bit + uint8_t qs[QK_K/2]; // quants, low 4 bits +} block_q5_K; +static_assert(sizeof(block_q5_K) == 2*sizeof(ggml_half) + K_SCALE_SIZE + QK_K/2 + QK_K/8, "wrong q5_K block size/padding"); + +// 6-bit quantization +// weight is represented as x = a * q +// 16 blocks of 16 elements each +// Effectively 6.5625 bits per weight +typedef struct { + uint8_t ql[QK_K/2]; // quants, lower 4 bits + uint8_t qh[QK_K/4]; // quants, upper 2 bits + int8_t scales[QK_K/16]; // scales, quantized with 8 bits + ggml_half d; // super-block scale +} block_q6_K; +static_assert(sizeof(block_q6_K) == sizeof(ggml_half) + QK_K / 16 + 3*QK_K/4, "wrong q6_K block size/padding"); + +// This is only used for intermediate quantization and dot products +typedef struct { + float d; // delta + int8_t qs[QK_K]; // quants + int16_t bsums[QK_K/16]; // sum of quants in groups of 16 +} block_q8_K; +static_assert(sizeof(block_q8_K) == sizeof(float) + QK_K + QK_K/16*sizeof(int16_t), "wrong q8_K block size/padding"); + +// (Almost) "true" 2-bit quantization. +// Due to the need to use blocks as per ggml design, it ends up using +// 2.0625 bpw because of the 16-bit scale for each block of 256. +typedef struct { + ggml_half d; + uint16_t qs[QK_K/8]; +} block_iq2_xxs; +static_assert(sizeof(block_iq2_xxs) == sizeof(ggml_half) + QK_K/8*sizeof(uint16_t), "wrong iq2_xxs block size/padding"); + +// 2.3125 bpw quants +typedef struct { + ggml_half d; + uint16_t qs[QK_K/8]; + uint8_t scales[QK_K/32]; +} block_iq2_xs; +static_assert(sizeof(block_iq2_xs) == sizeof(ggml_half) + QK_K/8*sizeof(uint16_t) + QK_K/32, "wrong iq2_xs block size/padding"); + +// 2.5625 bpw quants +typedef struct { + ggml_half d; + uint8_t qs[QK_K/4]; + uint8_t qh[QK_K/32]; + uint8_t scales[QK_K/32]; +} block_iq2_s; +static_assert(sizeof(block_iq2_s) == sizeof(ggml_half) + QK_K/4 + QK_K/16, "wrong iq2_s block size/padding"); + +// (Almost) "true" 3-bit quantization. +// Due to the need to use blocks as per ggml design, it ends up using +// 3.0625 bpw because of the 16-bit scale for each block of 256. +typedef struct { + ggml_half d; + uint8_t qs[3*QK_K/8]; +} block_iq3_xxs; +static_assert(sizeof(block_iq3_xxs) == sizeof(ggml_half) + 3*(QK_K/8), "wrong iq3_xxs block size/padding"); + +// 3.4375 bpw +#define IQ3S_N_SCALE QK_K/64 +typedef struct { + ggml_half d; + uint8_t qs[QK_K/4]; + uint8_t qh[QK_K/32]; + uint8_t signs[QK_K/8]; + uint8_t scales[IQ3S_N_SCALE]; +} block_iq3_s; +static_assert(sizeof(block_iq3_s) == sizeof(ggml_half) + 13*(QK_K/32) + IQ3S_N_SCALE, "wrong iq3_s block size/padding"); + +// 1.5625 bpw +typedef struct { + ggml_half d; + uint8_t qs[QK_K/8]; + uint16_t qh[QK_K/32]; +} block_iq1_s; +static_assert(sizeof(block_iq1_s) == sizeof(ggml_half) + QK_K/8 + QK_K/16, "wrong iq1_s block size/padding"); + +// 1.75 bpw +typedef struct { + uint8_t qs[QK_K/8]; // grid index, low 8 bits + uint8_t qh[QK_K/16]; // grid index, high 3 bits + grid shift bit (for two groups of 8) + uint8_t scales[QK_K/32]; // 3-bit block scales (4-bit if QK_K == 64) +} block_iq1_m; +static_assert(sizeof(block_iq1_m) == QK_K/8 + QK_K/16 + QK_K/32, "wrong iq1_m block size/padding"); + +// Used by IQ1_M quants +typedef union { + ggml_half f16; + uint16_t u16; +} iq1m_scale_t; + +// Non-linear quants +#define QK4_NL 32 +typedef struct { + ggml_half d; + uint8_t qs[QK4_NL/2]; +} block_iq4_nl; +static_assert(sizeof(block_iq4_nl) == sizeof(ggml_half) + QK4_NL/2, "wrong iq4_nl block size/padding"); + +typedef struct { + ggml_half d; + uint16_t scales_h; + uint8_t scales_l[QK_K/64]; + uint8_t qs[QK_K/2]; +} block_iq4_xs; +static_assert(sizeof(block_iq4_xs) == sizeof(ggml_half) + sizeof(uint16_t) + QK_K/64 + QK_K/2, "wrong iq4_xs block size/padding"); + +#endif // GGML_COMMON_DECL +#endif // GGML_COMMON_DECL + +//////////////////////////////////////////////////////////////////////////////// + +#ifndef GGML_COMMON_IMPL + +#if defined(GGML_COMMON_IMPL_C) +#include + +#define GGML_TABLE_BEGIN(type, name, size) static const type name[size] = { +#define GGML_TABLE_END() }; + +#define GGML_COMMON_IMPL +#elif defined(GGML_COMMON_IMPL_CPP) +#include + +#define GGML_TABLE_BEGIN(type, name, size) static const type name[size] = { +#define GGML_TABLE_END() }; + +#define GGML_COMMON_IMPL +#elif defined(GGML_COMMON_IMPL_METAL) +#include + +#define GGML_TABLE_BEGIN(type, name, size) static const constant type name[size] = { +#define GGML_TABLE_END() }; + +#define GGML_COMMON_IMPL +#elif defined(GGML_COMMON_IMPL_CUDA) || defined(GGML_COMMON_IMPL_HIP) || defined(GGML_COMMON_IMPL_MUSA) +#include + +#define GGML_TABLE_BEGIN(type, name, size) static const __device__ type name[size] = { +#define GGML_TABLE_END() }; + +#define GGML_COMMON_IMPL +#elif defined(GGML_COMMON_IMPL_SYCL) + +#include + +#define GGML_TABLE_BEGIN(type, name, size) static const type name[size] = { +#define GGML_TABLE_END() }; + +#define GGML_COMMON_IMPL +#endif + +#if defined(GGML_COMMON_IMPL) + +GGML_TABLE_BEGIN(uint8_t, kmask_iq2xs, 8) + 1, 2, 4, 8, 16, 32, 64, 128 +GGML_TABLE_END() + +GGML_TABLE_BEGIN(uint8_t, ksigns_iq2xs, 128) + 0, 129, 130, 3, 132, 5, 6, 135, 136, 9, 10, 139, 12, 141, 142, 15, + 144, 17, 18, 147, 20, 149, 150, 23, 24, 153, 154, 27, 156, 29, 30, 159, + 160, 33, 34, 163, 36, 165, 166, 39, 40, 169, 170, 43, 172, 45, 46, 175, + 48, 177, 178, 51, 180, 53, 54, 183, 184, 57, 58, 187, 60, 189, 190, 63, + 192, 65, 66, 195, 68, 197, 198, 71, 72, 201, 202, 75, 204, 77, 78, 207, + 80, 209, 210, 83, 212, 85, 86, 215, 216, 89, 90, 219, 92, 221, 222, 95, + 96, 225, 226, 99, 228, 101, 102, 231, 232, 105, 106, 235, 108, 237, 238, 111, + 240, 113, 114, 243, 116, 245, 246, 119, 120, 249, 250, 123, 252, 125, 126, 255, +GGML_TABLE_END() + +GGML_TABLE_BEGIN(uint64_t, ksigns64, 128) + 0x0000000000000000, 0xff000000000000ff, 0xff0000000000ff00, 0x000000000000ffff, + 0xff00000000ff0000, 0x0000000000ff00ff, 0x0000000000ffff00, 0xff00000000ffffff, + 0xff000000ff000000, 0x00000000ff0000ff, 0x00000000ff00ff00, 0xff000000ff00ffff, + 0x00000000ffff0000, 0xff000000ffff00ff, 0xff000000ffffff00, 0x00000000ffffffff, + 0xff0000ff00000000, 0x000000ff000000ff, 0x000000ff0000ff00, 0xff0000ff0000ffff, + 0x000000ff00ff0000, 0xff0000ff00ff00ff, 0xff0000ff00ffff00, 0x000000ff00ffffff, + 0x000000ffff000000, 0xff0000ffff0000ff, 0xff0000ffff00ff00, 0x000000ffff00ffff, + 0xff0000ffffff0000, 0x000000ffffff00ff, 0x000000ffffffff00, 0xff0000ffffffffff, + 0xff00ff0000000000, 0x0000ff00000000ff, 0x0000ff000000ff00, 0xff00ff000000ffff, + 0x0000ff0000ff0000, 0xff00ff0000ff00ff, 0xff00ff0000ffff00, 0x0000ff0000ffffff, + 0x0000ff00ff000000, 0xff00ff00ff0000ff, 0xff00ff00ff00ff00, 0x0000ff00ff00ffff, + 0xff00ff00ffff0000, 0x0000ff00ffff00ff, 0x0000ff00ffffff00, 0xff00ff00ffffffff, + 0x0000ffff00000000, 0xff00ffff000000ff, 0xff00ffff0000ff00, 0x0000ffff0000ffff, + 0xff00ffff00ff0000, 0x0000ffff00ff00ff, 0x0000ffff00ffff00, 0xff00ffff00ffffff, + 0xff00ffffff000000, 0x0000ffffff0000ff, 0x0000ffffff00ff00, 0xff00ffffff00ffff, + 0x0000ffffffff0000, 0xff00ffffffff00ff, 0xff00ffffffffff00, 0x0000ffffffffffff, + 0xffff000000000000, 0x00ff0000000000ff, 0x00ff00000000ff00, 0xffff00000000ffff, + 0x00ff000000ff0000, 0xffff000000ff00ff, 0xffff000000ffff00, 0x00ff000000ffffff, + 0x00ff0000ff000000, 0xffff0000ff0000ff, 0xffff0000ff00ff00, 0x00ff0000ff00ffff, + 0xffff0000ffff0000, 0x00ff0000ffff00ff, 0x00ff0000ffffff00, 0xffff0000ffffffff, + 0x00ff00ff00000000, 0xffff00ff000000ff, 0xffff00ff0000ff00, 0x00ff00ff0000ffff, + 0xffff00ff00ff0000, 0x00ff00ff00ff00ff, 0x00ff00ff00ffff00, 0xffff00ff00ffffff, + 0xffff00ffff000000, 0x00ff00ffff0000ff, 0x00ff00ffff00ff00, 0xffff00ffff00ffff, + 0x00ff00ffffff0000, 0xffff00ffffff00ff, 0xffff00ffffffff00, 0x00ff00ffffffffff, + 0x00ffff0000000000, 0xffffff00000000ff, 0xffffff000000ff00, 0x00ffff000000ffff, + 0xffffff0000ff0000, 0x00ffff0000ff00ff, 0x00ffff0000ffff00, 0xffffff0000ffffff, + 0xffffff00ff000000, 0x00ffff00ff0000ff, 0x00ffff00ff00ff00, 0xffffff00ff00ffff, + 0x00ffff00ffff0000, 0xffffff00ffff00ff, 0xffffff00ffffff00, 0x00ffff00ffffffff, + 0xffffffff00000000, 0x00ffffff000000ff, 0x00ffffff0000ff00, 0xffffffff0000ffff, + 0x00ffffff00ff0000, 0xffffffff00ff00ff, 0xffffffff00ffff00, 0x00ffffff00ffffff, + 0x00ffffffff000000, 0xffffffffff0000ff, 0xffffffffff00ff00, 0x00ffffffff00ffff, + 0xffffffffffff0000, 0x00ffffffffff00ff, 0x00ffffffffffff00, 0xffffffffffffffff, +GGML_TABLE_END() + + +GGML_TABLE_BEGIN(uint64_t, iq2xxs_grid, 256) + 0x0808080808080808, 0x080808080808082b, 0x0808080808081919, 0x0808080808082b08, + 0x0808080808082b2b, 0x0808080808190819, 0x0808080808191908, 0x08080808082b0808, + 0x08080808082b082b, 0x08080808082b2b08, 0x08080808082b2b2b, 0x0808080819080819, + 0x0808080819081908, 0x0808080819190808, 0x0808080819192b08, 0x08080808192b0819, + 0x08080808192b1908, 0x080808082b080808, 0x080808082b08082b, 0x080808082b082b2b, + 0x080808082b2b082b, 0x0808081908080819, 0x0808081908081908, 0x0808081908190808, + 0x0808081908191919, 0x0808081919080808, 0x080808192b081908, 0x080808192b192b08, + 0x0808082b08080808, 0x0808082b0808082b, 0x0808082b082b082b, 0x0808082b2b08082b, + 0x0808190808080819, 0x0808190808081908, 0x0808190808190808, 0x08081908082b0819, + 0x08081908082b1908, 0x0808190819080808, 0x080819081908082b, 0x0808190819082b08, + 0x08081908192b0808, 0x080819082b080819, 0x080819082b081908, 0x080819082b190808, + 0x080819082b2b1908, 0x0808191908080808, 0x080819190808082b, 0x0808191908082b08, + 0x08081919082b0808, 0x080819191908192b, 0x08081919192b2b19, 0x080819192b080808, + 0x080819192b190819, 0x0808192b08082b19, 0x0808192b08190808, 0x0808192b19080808, + 0x0808192b2b081908, 0x0808192b2b2b1908, 0x08082b0808080808, 0x08082b0808081919, + 0x08082b0808082b08, 0x08082b0808191908, 0x08082b08082b2b08, 0x08082b0819080819, + 0x08082b0819081908, 0x08082b0819190808, 0x08082b081919082b, 0x08082b082b082b08, + 0x08082b1908081908, 0x08082b1919080808, 0x08082b2b0808082b, 0x08082b2b08191908, + 0x0819080808080819, 0x0819080808081908, 0x0819080808190808, 0x08190808082b0819, + 0x0819080819080808, 0x08190808192b0808, 0x081908082b081908, 0x081908082b190808, + 0x081908082b191919, 0x0819081908080808, 0x0819081908082b08, 0x08190819082b0808, + 0x0819081919190808, 0x0819081919192b2b, 0x081908192b080808, 0x0819082b082b1908, + 0x0819082b19081919, 0x0819190808080808, 0x0819190808082b08, 0x08191908082b0808, + 0x08191908082b1919, 0x0819190819082b19, 0x081919082b080808, 0x0819191908192b08, + 0x08191919192b082b, 0x0819192b08080808, 0x0819192b0819192b, 0x08192b0808080819, + 0x08192b0808081908, 0x08192b0808190808, 0x08192b0819080808, 0x08192b082b080819, + 0x08192b1908080808, 0x08192b1908081919, 0x08192b192b2b0808, 0x08192b2b19190819, + 0x082b080808080808, 0x082b08080808082b, 0x082b080808082b2b, 0x082b080819081908, + 0x082b0808192b0819, 0x082b08082b080808, 0x082b08082b08082b, 0x082b0819082b2b19, + 0x082b081919082b08, 0x082b082b08080808, 0x082b082b0808082b, 0x082b190808080819, + 0x082b190808081908, 0x082b190808190808, 0x082b190819080808, 0x082b19081919192b, + 0x082b191908080808, 0x082b191919080819, 0x082b1919192b1908, 0x082b192b2b190808, + 0x082b2b0808082b08, 0x082b2b08082b0808, 0x082b2b082b191908, 0x082b2b2b19081908, + 0x1908080808080819, 0x1908080808081908, 0x1908080808190808, 0x1908080808192b08, + 0x19080808082b0819, 0x19080808082b1908, 0x1908080819080808, 0x1908080819082b08, + 0x190808081919192b, 0x19080808192b0808, 0x190808082b080819, 0x190808082b081908, + 0x190808082b190808, 0x1908081908080808, 0x19080819082b0808, 0x19080819192b0819, + 0x190808192b080808, 0x190808192b081919, 0x1908082b08080819, 0x1908082b08190808, + 0x1908082b19082b08, 0x1908082b1919192b, 0x1908082b192b2b08, 0x1908190808080808, + 0x1908190808082b08, 0x19081908082b0808, 0x190819082b080808, 0x190819082b192b19, + 0x190819190819082b, 0x19081919082b1908, 0x1908192b08080808, 0x19082b0808080819, + 0x19082b0808081908, 0x19082b0808190808, 0x19082b0819080808, 0x19082b0819081919, + 0x19082b1908080808, 0x19082b1919192b08, 0x19082b19192b0819, 0x19082b192b08082b, + 0x19082b2b19081919, 0x19082b2b2b190808, 0x1919080808080808, 0x1919080808082b08, + 0x1919080808190819, 0x1919080808192b19, 0x19190808082b0808, 0x191908082b080808, + 0x191908082b082b08, 0x1919081908081908, 0x191908191908082b, 0x191908192b2b1908, + 0x1919082b2b190819, 0x191919082b190808, 0x191919082b19082b, 0x1919191908082b2b, + 0x1919192b08080819, 0x1919192b19191908, 0x19192b0808080808, 0x19192b0808190819, + 0x19192b0808192b19, 0x19192b08192b1908, 0x19192b1919080808, 0x19192b2b08082b08, + 0x192b080808081908, 0x192b080808190808, 0x192b080819080808, 0x192b0808192b2b08, + 0x192b081908080808, 0x192b081919191919, 0x192b082b08192b08, 0x192b082b192b0808, + 0x192b190808080808, 0x192b190808081919, 0x192b191908190808, 0x192b19190819082b, + 0x192b19192b081908, 0x192b2b081908082b, 0x2b08080808080808, 0x2b0808080808082b, + 0x2b08080808082b2b, 0x2b08080819080819, 0x2b0808082b08082b, 0x2b08081908081908, + 0x2b08081908192b08, 0x2b08081919080808, 0x2b08082b08190819, 0x2b08190808080819, + 0x2b08190808081908, 0x2b08190808190808, 0x2b08190808191919, 0x2b08190819080808, + 0x2b081908192b0808, 0x2b08191908080808, 0x2b0819191908192b, 0x2b0819192b191908, + 0x2b08192b08082b19, 0x2b08192b19080808, 0x2b08192b192b0808, 0x2b082b080808082b, + 0x2b082b1908081908, 0x2b082b2b08190819, 0x2b19080808081908, 0x2b19080808190808, + 0x2b190808082b1908, 0x2b19080819080808, 0x2b1908082b2b0819, 0x2b1908190819192b, + 0x2b1908192b080808, 0x2b19082b19081919, 0x2b19190808080808, 0x2b191908082b082b, + 0x2b19190819081908, 0x2b19191919190819, 0x2b192b082b080819, 0x2b192b19082b0808, + 0x2b2b08080808082b, 0x2b2b080819190808, 0x2b2b08082b081919, 0x2b2b081908082b19, + 0x2b2b082b08080808, 0x2b2b190808192b08, 0x2b2b2b0819190808, 0x2b2b2b1908081908, +GGML_TABLE_END() + +GGML_TABLE_BEGIN(uint64_t, iq2xs_grid, 512) + 0x0808080808080808, 0x080808080808082b, 0x0808080808081919, 0x0808080808082b08, + 0x0808080808082b2b, 0x0808080808190819, 0x0808080808191908, 0x080808080819192b, + 0x0808080808192b19, 0x08080808082b0808, 0x08080808082b082b, 0x08080808082b1919, + 0x08080808082b2b08, 0x0808080819080819, 0x0808080819081908, 0x080808081908192b, + 0x0808080819082b19, 0x0808080819190808, 0x080808081919082b, 0x0808080819191919, + 0x0808080819192b08, 0x08080808192b0819, 0x08080808192b1908, 0x080808082b080808, + 0x080808082b08082b, 0x080808082b081919, 0x080808082b082b08, 0x080808082b190819, + 0x080808082b191908, 0x080808082b192b19, 0x080808082b2b0808, 0x0808081908080819, + 0x0808081908081908, 0x080808190808192b, 0x0808081908082b19, 0x0808081908190808, + 0x080808190819082b, 0x0808081908191919, 0x0808081908192b08, 0x0808081908192b2b, + 0x08080819082b0819, 0x08080819082b1908, 0x0808081919080808, 0x080808191908082b, + 0x0808081919081919, 0x0808081919082b08, 0x0808081919190819, 0x0808081919191908, + 0x08080819192b0808, 0x08080819192b2b08, 0x080808192b080819, 0x080808192b081908, + 0x080808192b190808, 0x0808082b08080808, 0x0808082b0808082b, 0x0808082b08081919, + 0x0808082b08082b08, 0x0808082b08190819, 0x0808082b08191908, 0x0808082b082b0808, + 0x0808082b19080819, 0x0808082b19081908, 0x0808082b19190808, 0x0808082b19191919, + 0x0808082b2b080808, 0x0808082b2b082b2b, 0x0808190808080819, 0x0808190808081908, + 0x080819080808192b, 0x0808190808082b19, 0x0808190808190808, 0x080819080819082b, + 0x0808190808191919, 0x0808190808192b08, 0x08081908082b0819, 0x08081908082b1908, + 0x0808190819080808, 0x080819081908082b, 0x0808190819081919, 0x0808190819082b08, + 0x0808190819190819, 0x0808190819191908, 0x080819081919192b, 0x08081908192b0808, + 0x080819082b080819, 0x080819082b081908, 0x080819082b190808, 0x0808191908080808, + 0x080819190808082b, 0x0808191908081919, 0x0808191908082b08, 0x0808191908190819, + 0x0808191908191908, 0x08081919082b0808, 0x0808191919080819, 0x0808191919081908, + 0x0808191919190808, 0x08081919192b0819, 0x080819192b080808, 0x0808192b08080819, + 0x0808192b08081908, 0x0808192b08190808, 0x0808192b082b192b, 0x0808192b19080808, + 0x0808192b1908082b, 0x0808192b2b081908, 0x08082b0808080808, 0x08082b080808082b, + 0x08082b0808081919, 0x08082b0808082b08, 0x08082b0808082b2b, 0x08082b0808190819, + 0x08082b0808191908, 0x08082b08082b0808, 0x08082b08082b1919, 0x08082b0819080819, + 0x08082b0819081908, 0x08082b0819190808, 0x08082b0819192b08, 0x08082b082b080808, + 0x08082b082b2b0808, 0x08082b082b2b2b2b, 0x08082b1908080819, 0x08082b1908081908, + 0x08082b1908190808, 0x08082b1919080808, 0x08082b192b080819, 0x08082b192b082b19, + 0x08082b2b08080808, 0x08082b2b082b0808, 0x08082b2b082b2b08, 0x08082b2b2b19192b, + 0x08082b2b2b2b0808, 0x0819080808080819, 0x0819080808081908, 0x081908080808192b, + 0x0819080808082b19, 0x0819080808190808, 0x081908080819082b, 0x0819080808191919, + 0x0819080808192b08, 0x08190808082b0819, 0x08190808082b1908, 0x0819080819080808, + 0x081908081908082b, 0x0819080819081919, 0x0819080819082b08, 0x0819080819190819, + 0x0819080819191908, 0x08190808192b0808, 0x08190808192b2b2b, 0x081908082b080819, + 0x081908082b081908, 0x081908082b190808, 0x0819081908080808, 0x081908190808082b, + 0x0819081908081919, 0x0819081908082b08, 0x0819081908190819, 0x0819081908191908, + 0x08190819082b0808, 0x0819081919080819, 0x0819081919081908, 0x0819081919190808, + 0x081908192b080808, 0x081908192b191908, 0x081908192b19192b, 0x0819082b08080819, + 0x0819082b08081908, 0x0819082b0808192b, 0x0819082b08190808, 0x0819082b19080808, + 0x0819082b192b0808, 0x0819190808080808, 0x081919080808082b, 0x0819190808081919, + 0x0819190808082b08, 0x0819190808190819, 0x0819190808191908, 0x08191908082b0808, + 0x0819190819080819, 0x0819190819081908, 0x0819190819082b19, 0x0819190819190808, + 0x08191908192b1908, 0x081919082b080808, 0x0819191908080819, 0x0819191908081908, + 0x0819191908190808, 0x0819191919080808, 0x0819192b08080808, 0x0819192b08191908, + 0x0819192b19082b19, 0x08192b0808080819, 0x08192b0808081908, 0x08192b0808190808, + 0x08192b080819082b, 0x08192b0819080808, 0x08192b0819191908, 0x08192b082b08192b, + 0x08192b1908080808, 0x08192b1908081919, 0x08192b19192b192b, 0x08192b2b19190819, + 0x08192b2b2b2b2b19, 0x082b080808080808, 0x082b08080808082b, 0x082b080808081919, + 0x082b080808082b08, 0x082b080808082b2b, 0x082b080808190819, 0x082b080808191908, + 0x082b0808082b0808, 0x082b080819080819, 0x082b080819081908, 0x082b080819190808, + 0x082b08082b080808, 0x082b08082b2b0808, 0x082b081908080819, 0x082b081908081908, + 0x082b081908190808, 0x082b081919080808, 0x082b081919082b08, 0x082b0819192b1919, + 0x082b082b08080808, 0x082b082b082b082b, 0x082b082b2b080808, 0x082b082b2b2b2b08, + 0x082b190808080819, 0x082b190808081908, 0x082b190808190808, 0x082b1908082b2b19, + 0x082b190819080808, 0x082b191908080808, 0x082b191919080819, 0x082b19191919082b, + 0x082b19192b192b19, 0x082b192b08080819, 0x082b192b08192b2b, 0x082b192b2b2b192b, + 0x082b2b0808080808, 0x082b2b0808082b08, 0x082b2b0808082b2b, 0x082b2b08082b0808, + 0x082b2b0819191919, 0x082b2b082b082b08, 0x082b2b082b2b082b, 0x082b2b19192b2b08, + 0x082b2b192b190808, 0x082b2b2b08082b08, 0x082b2b2b082b0808, 0x082b2b2b2b08082b, + 0x082b2b2b2b082b08, 0x082b2b2b2b082b2b, 0x1908080808080819, 0x1908080808081908, + 0x190808080808192b, 0x1908080808082b19, 0x1908080808190808, 0x190808080819082b, + 0x1908080808191919, 0x1908080808192b08, 0x19080808082b0819, 0x19080808082b1908, + 0x1908080819080808, 0x190808081908082b, 0x1908080819081919, 0x1908080819082b08, + 0x1908080819082b2b, 0x1908080819190819, 0x1908080819191908, 0x19080808192b0808, + 0x19080808192b1919, 0x190808082b080819, 0x190808082b081908, 0x190808082b190808, + 0x1908081908080808, 0x190808190808082b, 0x1908081908081919, 0x1908081908082b08, + 0x1908081908190819, 0x1908081908191908, 0x19080819082b0808, 0x1908081919080819, + 0x1908081919081908, 0x1908081919190808, 0x190808192b080808, 0x190808192b081919, + 0x190808192b2b082b, 0x1908082b08080819, 0x1908082b08081908, 0x1908082b08190808, + 0x1908082b0819082b, 0x1908082b082b2b19, 0x1908082b19080808, 0x1908190808080808, + 0x190819080808082b, 0x1908190808081919, 0x1908190808082b08, 0x1908190808190819, + 0x1908190808191908, 0x1908190808192b19, 0x19081908082b0808, 0x1908190819080819, + 0x1908190819081908, 0x1908190819190808, 0x190819082b080808, 0x190819082b191908, + 0x1908191908080819, 0x1908191908081908, 0x1908191908190808, 0x19081919082b1908, + 0x1908191919080808, 0x190819192b192b2b, 0x1908192b08080808, 0x1908192b08082b2b, + 0x1908192b19081908, 0x1908192b19190808, 0x19082b0808080819, 0x19082b0808081908, + 0x19082b0808190808, 0x19082b0819080808, 0x19082b0819081919, 0x19082b0819191908, + 0x19082b08192b082b, 0x19082b1908080808, 0x19082b1908190819, 0x19082b1919081908, + 0x19082b1919190808, 0x19082b19192b2b19, 0x19082b2b08081908, 0x1919080808080808, + 0x191908080808082b, 0x1919080808081919, 0x1919080808082b08, 0x1919080808190819, + 0x1919080808191908, 0x19190808082b0808, 0x19190808082b2b08, 0x1919080819080819, + 0x1919080819081908, 0x1919080819190808, 0x191908082b080808, 0x1919081908080819, + 0x1919081908081908, 0x1919081908190808, 0x1919081908191919, 0x1919081919080808, + 0x191908191908082b, 0x1919082b08080808, 0x1919082b19081908, 0x1919082b2b2b2b2b, + 0x1919190808080819, 0x1919190808081908, 0x1919190808190808, 0x19191908082b0819, + 0x1919190819080808, 0x19191908192b0808, 0x191919082b080819, 0x191919082b2b0819, + 0x1919191908080808, 0x1919191908082b08, 0x191919192b080808, 0x191919192b082b08, + 0x1919192b082b0819, 0x1919192b192b2b08, 0x1919192b2b2b0819, 0x19192b0808080808, + 0x19192b0808191908, 0x19192b0819080819, 0x19192b0819190808, 0x19192b082b192b19, + 0x19192b1908192b2b, 0x19192b1919080808, 0x19192b191908082b, 0x19192b2b2b081919, + 0x192b080808080819, 0x192b080808081908, 0x192b080808190808, 0x192b080819080808, + 0x192b080819191908, 0x192b0808192b082b, 0x192b08082b08192b, 0x192b08082b2b2b19, + 0x192b081908080808, 0x192b082b082b1908, 0x192b082b19082b2b, 0x192b082b2b19082b, + 0x192b190808080808, 0x192b19080819192b, 0x192b191908190808, 0x192b191919080808, + 0x192b191919081919, 0x192b19192b2b1908, 0x192b2b0808080819, 0x192b2b08192b2b2b, + 0x192b2b19082b1919, 0x192b2b2b0808192b, 0x192b2b2b19191908, 0x192b2b2b192b082b, + 0x2b08080808080808, 0x2b0808080808082b, 0x2b08080808081919, 0x2b08080808082b08, + 0x2b08080808190819, 0x2b08080808191908, 0x2b080808082b0808, 0x2b080808082b2b2b, + 0x2b08080819080819, 0x2b08080819081908, 0x2b08080819190808, 0x2b0808082b080808, + 0x2b0808082b08082b, 0x2b0808082b2b2b08, 0x2b0808082b2b2b2b, 0x2b08081908080819, + 0x2b08081908081908, 0x2b0808190808192b, 0x2b08081908190808, 0x2b08081919080808, + 0x2b08081919190819, 0x2b08081919192b19, 0x2b08082b08080808, 0x2b08082b082b0808, + 0x2b08082b2b080808, 0x2b08082b2b08082b, 0x2b08082b2b2b0808, 0x2b08082b2b2b2b08, + 0x2b08190808080819, 0x2b08190808081908, 0x2b08190808190808, 0x2b0819080819082b, + 0x2b08190808191919, 0x2b08190819080808, 0x2b081908192b0808, 0x2b0819082b082b19, + 0x2b08191908080808, 0x2b08191919081908, 0x2b0819192b2b1919, 0x2b08192b08192b08, + 0x2b08192b192b2b2b, 0x2b082b0808080808, 0x2b082b0808082b08, 0x2b082b08082b1919, + 0x2b082b0819192b2b, 0x2b082b082b080808, 0x2b082b082b08082b, 0x2b082b082b2b2b08, + 0x2b082b190808192b, 0x2b082b2b082b082b, 0x2b082b2b2b080808, 0x2b082b2b2b082b08, + 0x2b082b2b2b19192b, 0x2b082b2b2b2b2b08, 0x2b19080808080819, 0x2b19080808081908, + 0x2b19080808190808, 0x2b19080819080808, 0x2b1908081919192b, 0x2b1908082b081908, + 0x2b19081908080808, 0x2b190819082b082b, 0x2b190819192b1908, 0x2b19082b1919192b, + 0x2b19082b2b082b19, 0x2b19190808080808, 0x2b19190808081919, 0x2b19190819081908, + 0x2b19190819190808, 0x2b19190819192b08, 0x2b191919082b2b19, 0x2b1919192b190808, + 0x2b1919192b19082b, 0x2b19192b19080819, 0x2b192b0819190819, 0x2b192b082b2b192b, + 0x2b192b1919082b19, 0x2b192b2b08191919, 0x2b192b2b192b0808, 0x2b2b080808080808, + 0x2b2b08080808082b, 0x2b2b080808082b08, 0x2b2b080808082b2b, 0x2b2b0808082b0808, + 0x2b2b0808082b2b2b, 0x2b2b08082b2b0808, 0x2b2b081919190819, 0x2b2b081919192b19, + 0x2b2b08192b2b192b, 0x2b2b082b08080808, 0x2b2b082b0808082b, 0x2b2b082b08082b08, + 0x2b2b082b082b2b2b, 0x2b2b082b2b080808, 0x2b2b082b2b2b0808, 0x2b2b190819080808, + 0x2b2b19082b191919, 0x2b2b192b192b1919, 0x2b2b192b2b192b08, 0x2b2b2b0808082b2b, + 0x2b2b2b08082b0808, 0x2b2b2b08082b082b, 0x2b2b2b08082b2b08, 0x2b2b2b082b2b0808, + 0x2b2b2b082b2b2b08, 0x2b2b2b1908081908, 0x2b2b2b192b081908, 0x2b2b2b192b08192b, + 0x2b2b2b2b082b2b08, 0x2b2b2b2b082b2b2b, 0x2b2b2b2b2b190819, 0x2b2b2b2b2b2b2b2b, +GGML_TABLE_END() + +GGML_TABLE_BEGIN(uint64_t, iq2s_grid, 1024) + 0x0808080808080808, 0x080808080808082b, 0x0808080808081919, 0x0808080808082b08, + 0x0808080808082b2b, 0x0808080808190819, 0x0808080808191908, 0x080808080819192b, + 0x0808080808192b19, 0x08080808082b0808, 0x08080808082b082b, 0x08080808082b1919, + 0x08080808082b2b08, 0x0808080819080819, 0x0808080819081908, 0x080808081908192b, + 0x0808080819082b19, 0x0808080819190808, 0x080808081919082b, 0x0808080819191919, + 0x0808080819192b08, 0x08080808192b0819, 0x08080808192b1908, 0x08080808192b192b, + 0x08080808192b2b19, 0x080808082b080808, 0x080808082b08082b, 0x080808082b081919, + 0x080808082b082b08, 0x080808082b190819, 0x080808082b191908, 0x080808082b2b0808, + 0x080808082b2b1919, 0x080808082b2b2b2b, 0x0808081908080819, 0x0808081908081908, + 0x080808190808192b, 0x0808081908082b19, 0x0808081908190808, 0x080808190819082b, + 0x0808081908191919, 0x0808081908192b08, 0x08080819082b0819, 0x08080819082b1908, + 0x0808081919080808, 0x080808191908082b, 0x0808081919081919, 0x0808081919082b08, + 0x0808081919190819, 0x0808081919191908, 0x080808191919192b, 0x0808081919192b19, + 0x08080819192b0808, 0x08080819192b1919, 0x08080819192b2b08, 0x080808192b080819, + 0x080808192b081908, 0x080808192b190808, 0x080808192b19082b, 0x080808192b191919, + 0x080808192b2b0819, 0x080808192b2b1908, 0x0808082b08080808, 0x0808082b0808082b, + 0x0808082b08081919, 0x0808082b08082b08, 0x0808082b08190819, 0x0808082b08191908, + 0x0808082b082b0808, 0x0808082b082b2b2b, 0x0808082b19080819, 0x0808082b19081908, + 0x0808082b1908192b, 0x0808082b19082b19, 0x0808082b19190808, 0x0808082b19191919, + 0x0808082b2b080808, 0x0808082b2b081919, 0x0808082b2b082b2b, 0x0808082b2b191908, + 0x0808082b2b2b082b, 0x0808190808080819, 0x0808190808081908, 0x080819080808192b, + 0x0808190808082b19, 0x0808190808190808, 0x080819080819082b, 0x0808190808191919, + 0x0808190808192b08, 0x08081908082b0819, 0x08081908082b1908, 0x08081908082b192b, + 0x08081908082b2b19, 0x0808190819080808, 0x080819081908082b, 0x0808190819081919, + 0x0808190819082b08, 0x0808190819082b2b, 0x0808190819190819, 0x0808190819191908, + 0x080819081919192b, 0x0808190819192b19, 0x08081908192b0808, 0x08081908192b082b, + 0x08081908192b1919, 0x080819082b080819, 0x080819082b081908, 0x080819082b08192b, + 0x080819082b082b19, 0x080819082b190808, 0x080819082b191919, 0x080819082b192b08, + 0x080819082b2b0819, 0x080819082b2b1908, 0x0808191908080808, 0x080819190808082b, + 0x0808191908081919, 0x0808191908082b08, 0x0808191908082b2b, 0x0808191908190819, + 0x0808191908191908, 0x080819190819192b, 0x0808191908192b19, 0x08081919082b0808, + 0x08081919082b1919, 0x08081919082b2b08, 0x0808191919080819, 0x0808191919081908, + 0x080819191908192b, 0x0808191919082b19, 0x0808191919190808, 0x080819191919082b, + 0x0808191919191919, 0x0808191919192b08, 0x08081919192b0819, 0x08081919192b1908, + 0x080819192b080808, 0x080819192b08082b, 0x080819192b081919, 0x080819192b082b08, + 0x080819192b190819, 0x080819192b191908, 0x080819192b2b0808, 0x0808192b08080819, + 0x0808192b08081908, 0x0808192b0808192b, 0x0808192b08082b19, 0x0808192b08190808, + 0x0808192b08191919, 0x0808192b19080808, 0x0808192b19081919, 0x0808192b19082b08, + 0x0808192b19190819, 0x0808192b19191908, 0x0808192b192b0808, 0x0808192b2b080819, + 0x0808192b2b081908, 0x0808192b2b190808, 0x08082b0808080808, 0x08082b080808082b, + 0x08082b0808081919, 0x08082b0808082b08, 0x08082b0808190819, 0x08082b0808191908, + 0x08082b080819192b, 0x08082b0808192b19, 0x08082b08082b0808, 0x08082b08082b1919, + 0x08082b08082b2b2b, 0x08082b0819080819, 0x08082b0819081908, 0x08082b081908192b, + 0x08082b0819082b19, 0x08082b0819190808, 0x08082b081919082b, 0x08082b0819191919, + 0x08082b0819192b08, 0x08082b08192b0819, 0x08082b08192b1908, 0x08082b082b080808, + 0x08082b082b081919, 0x08082b082b191908, 0x08082b082b2b2b2b, 0x08082b1908080819, + 0x08082b1908081908, 0x08082b1908190808, 0x08082b190819082b, 0x08082b1908191919, + 0x08082b1908192b08, 0x08082b19082b0819, 0x08082b1919080808, 0x08082b1919081919, + 0x08082b1919082b08, 0x08082b1919190819, 0x08082b1919191908, 0x08082b19192b0808, + 0x08082b192b080819, 0x08082b192b190808, 0x08082b2b08080808, 0x08082b2b08190819, + 0x08082b2b08191908, 0x08082b2b082b082b, 0x08082b2b082b2b08, 0x08082b2b082b2b2b, + 0x08082b2b19190808, 0x08082b2b2b192b19, 0x0819080808080819, 0x0819080808081908, + 0x081908080808192b, 0x0819080808082b19, 0x0819080808190808, 0x081908080819082b, + 0x0819080808191919, 0x0819080808192b08, 0x08190808082b0819, 0x08190808082b1908, + 0x08190808082b192b, 0x0819080819080808, 0x081908081908082b, 0x0819080819081919, + 0x0819080819082b08, 0x0819080819190819, 0x0819080819191908, 0x081908081919192b, + 0x0819080819192b19, 0x08190808192b0808, 0x08190808192b082b, 0x08190808192b1919, + 0x08190808192b2b08, 0x081908082b080819, 0x081908082b081908, 0x081908082b08192b, + 0x081908082b190808, 0x081908082b191919, 0x081908082b192b08, 0x081908082b2b0819, + 0x081908082b2b1908, 0x0819081908080808, 0x081908190808082b, 0x0819081908081919, + 0x0819081908082b08, 0x0819081908082b2b, 0x0819081908190819, 0x0819081908191908, + 0x081908190819192b, 0x0819081908192b19, 0x08190819082b0808, 0x08190819082b082b, + 0x08190819082b1919, 0x08190819082b2b08, 0x0819081919080819, 0x0819081919081908, + 0x081908191908192b, 0x0819081919082b19, 0x0819081919190808, 0x081908191919082b, + 0x0819081919191919, 0x0819081919192b08, 0x08190819192b0819, 0x08190819192b1908, + 0x081908192b080808, 0x081908192b08082b, 0x081908192b081919, 0x081908192b082b08, + 0x081908192b190819, 0x081908192b191908, 0x0819082b08080819, 0x0819082b08081908, + 0x0819082b08082b19, 0x0819082b08190808, 0x0819082b08191919, 0x0819082b082b0819, + 0x0819082b082b1908, 0x0819082b19080808, 0x0819082b19081919, 0x0819082b19190819, + 0x0819082b19191908, 0x0819082b2b080819, 0x0819082b2b081908, 0x0819082b2b190808, + 0x0819190808080808, 0x081919080808082b, 0x0819190808081919, 0x0819190808082b08, + 0x0819190808190819, 0x0819190808191908, 0x081919080819192b, 0x0819190808192b19, + 0x08191908082b0808, 0x08191908082b1919, 0x08191908082b2b08, 0x0819190819080819, + 0x0819190819081908, 0x081919081908192b, 0x0819190819082b19, 0x0819190819190808, + 0x081919081919082b, 0x0819190819191919, 0x0819190819192b08, 0x08191908192b0819, + 0x08191908192b1908, 0x081919082b080808, 0x081919082b08082b, 0x081919082b081919, + 0x081919082b082b08, 0x081919082b190819, 0x081919082b191908, 0x081919082b2b0808, + 0x0819191908080819, 0x0819191908081908, 0x081919190808192b, 0x0819191908082b19, + 0x0819191908190808, 0x081919190819082b, 0x0819191908191919, 0x0819191908192b08, + 0x08191919082b0819, 0x08191919082b1908, 0x0819191919080808, 0x081919191908082b, + 0x0819191919081919, 0x0819191919082b08, 0x0819191919190819, 0x0819191919191908, + 0x08191919192b0808, 0x081919192b080819, 0x081919192b081908, 0x081919192b190808, + 0x0819192b08080808, 0x0819192b08081919, 0x0819192b08082b08, 0x0819192b08190819, + 0x0819192b08191908, 0x0819192b082b0808, 0x0819192b19080819, 0x0819192b19081908, + 0x0819192b19190808, 0x0819192b2b080808, 0x0819192b2b2b2b2b, 0x08192b0808080819, + 0x08192b0808081908, 0x08192b080808192b, 0x08192b0808082b19, 0x08192b0808190808, + 0x08192b0808191919, 0x08192b0808192b08, 0x08192b08082b0819, 0x08192b0819080808, + 0x08192b081908082b, 0x08192b0819081919, 0x08192b0819082b08, 0x08192b0819190819, + 0x08192b0819191908, 0x08192b08192b0808, 0x08192b082b080819, 0x08192b082b081908, + 0x08192b1908080808, 0x08192b190808082b, 0x08192b1908081919, 0x08192b1908082b08, + 0x08192b1908190819, 0x08192b1908191908, 0x08192b19082b0808, 0x08192b1919080819, + 0x08192b1919081908, 0x08192b1919190808, 0x08192b19192b2b19, 0x08192b192b2b082b, + 0x08192b2b08081908, 0x08192b2b08190808, 0x08192b2b19080808, 0x08192b2b1919192b, + 0x082b080808080808, 0x082b08080808082b, 0x082b080808081919, 0x082b080808082b08, + 0x082b080808190819, 0x082b080808191908, 0x082b08080819192b, 0x082b080808192b19, + 0x082b0808082b0808, 0x082b0808082b1919, 0x082b0808082b2b2b, 0x082b080819080819, + 0x082b080819081908, 0x082b080819190808, 0x082b08081919082b, 0x082b080819191919, + 0x082b0808192b1908, 0x082b08082b080808, 0x082b08082b082b2b, 0x082b08082b191908, + 0x082b08082b2b2b2b, 0x082b081908080819, 0x082b081908081908, 0x082b081908190808, + 0x082b08190819082b, 0x082b081908191919, 0x082b0819082b0819, 0x082b081919080808, + 0x082b08191908082b, 0x082b081919081919, 0x082b081919190819, 0x082b081919191908, + 0x082b0819192b0808, 0x082b08192b080819, 0x082b08192b081908, 0x082b08192b190808, + 0x082b082b08080808, 0x082b082b08082b2b, 0x082b082b082b082b, 0x082b082b082b2b08, + 0x082b082b082b2b2b, 0x082b082b19081908, 0x082b082b19190808, 0x082b082b2b082b08, + 0x082b082b2b082b2b, 0x082b082b2b2b2b08, 0x082b190808080819, 0x082b190808081908, + 0x082b19080808192b, 0x082b190808082b19, 0x082b190808190808, 0x082b190808191919, + 0x082b190808192b08, 0x082b1908082b0819, 0x082b1908082b1908, 0x082b190819080808, + 0x082b19081908082b, 0x082b190819081919, 0x082b190819082b08, 0x082b190819190819, + 0x082b190819191908, 0x082b1908192b0808, 0x082b19082b080819, 0x082b19082b081908, + 0x082b19082b190808, 0x082b191908080808, 0x082b191908081919, 0x082b191908082b08, + 0x082b191908190819, 0x082b191908191908, 0x082b1919082b0808, 0x082b191919080819, + 0x082b191919081908, 0x082b191919190808, 0x082b1919192b192b, 0x082b19192b080808, + 0x082b192b08080819, 0x082b192b08081908, 0x082b192b08190808, 0x082b192b19080808, + 0x082b192b19192b19, 0x082b2b0808080808, 0x082b2b0808081919, 0x082b2b0808190819, + 0x082b2b0808191908, 0x082b2b0819080819, 0x082b2b0819081908, 0x082b2b0819190808, + 0x082b2b082b082b2b, 0x082b2b082b2b2b2b, 0x082b2b1908080819, 0x082b2b1908081908, + 0x082b2b1908190808, 0x082b2b192b191919, 0x082b2b2b08082b2b, 0x082b2b2b082b082b, + 0x082b2b2b192b1908, 0x082b2b2b2b082b08, 0x082b2b2b2b082b2b, 0x1908080808080819, + 0x1908080808081908, 0x190808080808192b, 0x1908080808082b19, 0x1908080808190808, + 0x190808080819082b, 0x1908080808191919, 0x1908080808192b08, 0x1908080808192b2b, + 0x19080808082b0819, 0x19080808082b1908, 0x19080808082b192b, 0x1908080819080808, + 0x190808081908082b, 0x1908080819081919, 0x1908080819082b08, 0x1908080819082b2b, + 0x1908080819190819, 0x1908080819191908, 0x190808081919192b, 0x1908080819192b19, + 0x19080808192b0808, 0x19080808192b082b, 0x19080808192b1919, 0x190808082b080819, + 0x190808082b081908, 0x190808082b190808, 0x190808082b191919, 0x190808082b192b08, + 0x190808082b2b0819, 0x190808082b2b1908, 0x1908081908080808, 0x190808190808082b, + 0x1908081908081919, 0x1908081908082b08, 0x1908081908190819, 0x1908081908191908, + 0x190808190819192b, 0x1908081908192b19, 0x19080819082b0808, 0x19080819082b082b, + 0x19080819082b1919, 0x1908081919080819, 0x1908081919081908, 0x190808191908192b, + 0x1908081919082b19, 0x1908081919190808, 0x190808191919082b, 0x1908081919191919, + 0x1908081919192b08, 0x19080819192b0819, 0x19080819192b1908, 0x190808192b080808, + 0x190808192b08082b, 0x190808192b081919, 0x190808192b082b08, 0x190808192b190819, + 0x190808192b191908, 0x190808192b2b0808, 0x1908082b08080819, 0x1908082b08081908, + 0x1908082b08190808, 0x1908082b0819082b, 0x1908082b08191919, 0x1908082b08192b08, + 0x1908082b082b1908, 0x1908082b19080808, 0x1908082b19081919, 0x1908082b19082b08, + 0x1908082b19190819, 0x1908082b19191908, 0x1908082b192b0808, 0x1908082b2b080819, + 0x1908082b2b081908, 0x1908190808080808, 0x190819080808082b, 0x1908190808081919, + 0x1908190808082b08, 0x1908190808082b2b, 0x1908190808190819, 0x1908190808191908, + 0x190819080819192b, 0x1908190808192b19, 0x19081908082b0808, 0x19081908082b082b, + 0x19081908082b1919, 0x19081908082b2b08, 0x1908190819080819, 0x1908190819081908, + 0x190819081908192b, 0x1908190819082b19, 0x1908190819190808, 0x190819081919082b, + 0x1908190819191919, 0x1908190819192b08, 0x19081908192b0819, 0x19081908192b1908, + 0x190819082b080808, 0x190819082b08082b, 0x190819082b081919, 0x190819082b082b08, + 0x190819082b190819, 0x190819082b191908, 0x190819082b2b0808, 0x1908191908080819, + 0x1908191908081908, 0x190819190808192b, 0x1908191908082b19, 0x1908191908190808, + 0x190819190819082b, 0x1908191908191919, 0x1908191908192b08, 0x19081919082b0819, + 0x19081919082b1908, 0x1908191919080808, 0x190819191908082b, 0x1908191919081919, + 0x1908191919082b08, 0x1908191919190819, 0x1908191919191908, 0x19081919192b0808, + 0x19081919192b2b2b, 0x190819192b080819, 0x190819192b081908, 0x190819192b190808, + 0x1908192b08080808, 0x1908192b0808082b, 0x1908192b08081919, 0x1908192b08082b08, + 0x1908192b08190819, 0x1908192b08191908, 0x1908192b082b0808, 0x1908192b19080819, + 0x1908192b19081908, 0x1908192b19190808, 0x1908192b2b080808, 0x1908192b2b2b1919, + 0x19082b0808080819, 0x19082b0808081908, 0x19082b0808082b19, 0x19082b0808190808, + 0x19082b080819082b, 0x19082b0808191919, 0x19082b0808192b08, 0x19082b08082b0819, + 0x19082b08082b1908, 0x19082b0819080808, 0x19082b081908082b, 0x19082b0819081919, + 0x19082b0819082b08, 0x19082b0819190819, 0x19082b0819191908, 0x19082b08192b0808, + 0x19082b082b081908, 0x19082b082b190808, 0x19082b1908080808, 0x19082b190808082b, + 0x19082b1908081919, 0x19082b1908082b08, 0x19082b1908190819, 0x19082b1908191908, + 0x19082b19082b0808, 0x19082b1919080819, 0x19082b1919081908, 0x19082b1919190808, + 0x19082b192b080808, 0x19082b192b19192b, 0x19082b2b08080819, 0x19082b2b08081908, + 0x19082b2b08190808, 0x19082b2b19080808, 0x1919080808080808, 0x191908080808082b, + 0x1919080808081919, 0x1919080808082b08, 0x1919080808190819, 0x1919080808191908, + 0x191908080819192b, 0x1919080808192b19, 0x19190808082b0808, 0x19190808082b082b, + 0x19190808082b1919, 0x19190808082b2b08, 0x1919080819080819, 0x1919080819081908, + 0x191908081908192b, 0x1919080819082b19, 0x1919080819190808, 0x191908081919082b, + 0x1919080819191919, 0x1919080819192b08, 0x19190808192b0819, 0x19190808192b1908, + 0x191908082b080808, 0x191908082b08082b, 0x191908082b081919, 0x191908082b082b08, + 0x191908082b190819, 0x191908082b191908, 0x1919081908080819, 0x1919081908081908, + 0x191908190808192b, 0x1919081908082b19, 0x1919081908190808, 0x191908190819082b, + 0x1919081908191919, 0x1919081908192b08, 0x19190819082b0819, 0x19190819082b1908, + 0x1919081919080808, 0x191908191908082b, 0x1919081919081919, 0x1919081919082b08, + 0x1919081919190819, 0x1919081919191908, 0x19190819192b0808, 0x191908192b080819, + 0x191908192b081908, 0x191908192b190808, 0x1919082b08080808, 0x1919082b08081919, + 0x1919082b08082b08, 0x1919082b08190819, 0x1919082b08191908, 0x1919082b082b0808, + 0x1919082b19080819, 0x1919082b19081908, 0x1919082b19190808, 0x1919082b192b2b19, + 0x1919082b2b080808, 0x1919190808080819, 0x1919190808081908, 0x191919080808192b, + 0x1919190808082b19, 0x1919190808190808, 0x191919080819082b, 0x1919190808191919, + 0x1919190808192b08, 0x19191908082b0819, 0x19191908082b1908, 0x1919190819080808, + 0x191919081908082b, 0x1919190819081919, 0x1919190819082b08, 0x1919190819190819, + 0x1919190819191908, 0x19191908192b0808, 0x191919082b080819, 0x191919082b081908, + 0x191919082b190808, 0x1919191908080808, 0x191919190808082b, 0x1919191908081919, + 0x1919191908082b08, 0x1919191908190819, 0x1919191908191908, 0x19191919082b0808, + 0x1919191919080819, 0x1919191919081908, 0x1919191919190808, 0x191919192b080808, + 0x1919192b08080819, 0x1919192b08081908, 0x1919192b08190808, 0x1919192b082b192b, + 0x1919192b19080808, 0x19192b0808080808, 0x19192b080808082b, 0x19192b0808081919, + 0x19192b0808082b08, 0x19192b0808190819, 0x19192b0808191908, 0x19192b08082b0808, + 0x19192b0819080819, 0x19192b0819081908, 0x19192b0819190808, 0x19192b0819192b2b, + 0x19192b082b080808, 0x19192b1908080819, 0x19192b1908081908, 0x19192b1908190808, + 0x19192b1919080808, 0x19192b2b08080808, 0x19192b2b08192b19, 0x19192b2b2b081919, + 0x19192b2b2b2b2b08, 0x192b080808080819, 0x192b080808081908, 0x192b08080808192b, + 0x192b080808190808, 0x192b08080819082b, 0x192b080808191919, 0x192b080808192b08, + 0x192b0808082b0819, 0x192b0808082b1908, 0x192b080819080808, 0x192b080819081919, + 0x192b080819082b08, 0x192b080819190819, 0x192b080819191908, 0x192b0808192b0808, + 0x192b08082b081908, 0x192b08082b190808, 0x192b081908080808, 0x192b08190808082b, + 0x192b081908081919, 0x192b081908082b08, 0x192b081908190819, 0x192b081908191908, + 0x192b0819082b0808, 0x192b081919080819, 0x192b081919081908, 0x192b081919190808, + 0x192b08192b080808, 0x192b08192b192b19, 0x192b082b08081908, 0x192b082b08190808, + 0x192b082b19080808, 0x192b082b1919192b, 0x192b082b2b2b0819, 0x192b190808080808, + 0x192b190808081919, 0x192b190808082b08, 0x192b190808190819, 0x192b190808191908, + 0x192b1908082b0808, 0x192b190819080819, 0x192b190819081908, 0x192b190819190808, + 0x192b19082b080808, 0x192b191908080819, 0x192b191908081908, 0x192b191908190808, + 0x192b191919080808, 0x192b191919082b2b, 0x192b1919192b2b08, 0x192b19192b19082b, + 0x192b192b08080808, 0x192b192b2b191908, 0x192b2b0808080819, 0x192b2b0808081908, + 0x192b2b0808190808, 0x192b2b08192b1919, 0x192b2b082b192b08, 0x192b2b1908080808, + 0x192b2b19082b2b2b, 0x192b2b2b1908082b, 0x192b2b2b2b2b0819, 0x2b08080808080808, + 0x2b0808080808082b, 0x2b08080808081919, 0x2b08080808082b08, 0x2b08080808190819, + 0x2b08080808191908, 0x2b08080808192b19, 0x2b080808082b0808, 0x2b080808082b1919, + 0x2b08080819080819, 0x2b08080819081908, 0x2b08080819190808, 0x2b0808081919082b, + 0x2b08080819191919, 0x2b08080819192b08, 0x2b080808192b0819, 0x2b0808082b080808, + 0x2b0808082b081919, 0x2b0808082b190819, 0x2b0808082b191908, 0x2b08081908080819, + 0x2b08081908081908, 0x2b08081908082b19, 0x2b08081908190808, 0x2b0808190819082b, + 0x2b08081908191919, 0x2b08081908192b08, 0x2b080819082b0819, 0x2b080819082b1908, + 0x2b08081919080808, 0x2b0808191908082b, 0x2b08081919081919, 0x2b08081919082b08, + 0x2b08081919190819, 0x2b08081919191908, 0x2b0808192b080819, 0x2b0808192b081908, + 0x2b0808192b190808, 0x2b0808192b2b2b19, 0x2b08082b08080808, 0x2b08082b08081919, + 0x2b08082b08082b2b, 0x2b08082b08190819, 0x2b08082b08191908, 0x2b08082b19080819, + 0x2b08082b19081908, 0x2b08082b19190808, 0x2b08190808080819, 0x2b08190808081908, + 0x2b0819080808192b, 0x2b08190808082b19, 0x2b08190808190808, 0x2b0819080819082b, + 0x2b08190808191919, 0x2b08190808192b08, 0x2b081908082b0819, 0x2b08190819080808, + 0x2b0819081908082b, 0x2b08190819081919, 0x2b08190819082b08, 0x2b08190819190819, + 0x2b08190819191908, 0x2b081908192b0808, 0x2b0819082b080819, 0x2b0819082b081908, + 0x2b0819082b190808, 0x2b08191908080808, 0x2b0819190808082b, 0x2b08191908081919, + 0x2b08191908082b08, 0x2b08191908190819, 0x2b08191908191908, 0x2b081919082b0808, + 0x2b08191919080819, 0x2b08191919081908, 0x2b08191919190808, 0x2b0819192b080808, + 0x2b0819192b082b2b, 0x2b08192b08080819, 0x2b08192b08081908, 0x2b08192b08190808, + 0x2b08192b082b2b19, 0x2b08192b19080808, 0x2b082b0808080808, 0x2b082b0808081919, + 0x2b082b0808190819, 0x2b082b0808191908, 0x2b082b0819080819, 0x2b082b0819081908, + 0x2b082b0819190808, 0x2b082b082b2b082b, 0x2b082b1908080819, 0x2b082b1908081908, + 0x2b082b1919080808, 0x2b082b19192b1919, 0x2b082b2b082b082b, 0x2b082b2b19192b08, + 0x2b082b2b19192b2b, 0x2b082b2b2b08082b, 0x2b082b2b2b2b082b, 0x2b19080808080819, + 0x2b19080808081908, 0x2b19080808082b19, 0x2b19080808190808, 0x2b1908080819082b, + 0x2b19080808191919, 0x2b19080808192b08, 0x2b190808082b1908, 0x2b19080819080808, + 0x2b1908081908082b, 0x2b19080819081919, 0x2b19080819082b08, 0x2b19080819190819, + 0x2b19080819191908, 0x2b190808192b0808, 0x2b1908082b080819, 0x2b1908082b081908, + 0x2b1908082b190808, 0x2b19081908080808, 0x2b19081908081919, 0x2b19081908190819, + 0x2b19081908191908, 0x2b19081919080819, 0x2b19081919081908, 0x2b19081919190808, + 0x2b19081919192b2b, 0x2b19082b08080819, 0x2b19082b08081908, 0x2b19082b08190808, + 0x2b19082b19080808, 0x2b19082b2b2b192b, 0x2b19190808080808, 0x2b1919080808082b, + 0x2b19190808081919, 0x2b19190808082b08, 0x2b19190808190819, 0x2b19190808191908, + 0x2b191908082b0808, 0x2b19190819080819, 0x2b19190819081908, 0x2b19190819190808, + 0x2b1919082b080808, 0x2b1919082b19192b, 0x2b19191908080819, 0x2b19191908081908, + 0x2b19191908190808, 0x2b19191919080808, 0x2b1919192b192b08, 0x2b1919192b2b0819, + 0x2b19192b08080808, 0x2b19192b1908192b, 0x2b19192b192b1908, 0x2b192b0808080819, + 0x2b192b0808081908, 0x2b192b0808190808, 0x2b192b08082b192b, 0x2b192b0819080808, + 0x2b192b082b2b2b19, 0x2b192b1908080808, 0x2b192b1919082b19, 0x2b192b191919082b, + 0x2b192b2b2b190808, 0x2b2b080808080808, 0x2b2b080808081919, 0x2b2b080808082b2b, + 0x2b2b080808191908, 0x2b2b0808082b082b, 0x2b2b0808082b2b2b, 0x2b2b080819080819, + 0x2b2b080819081908, 0x2b2b080819190808, 0x2b2b08082b2b082b, 0x2b2b08082b2b2b2b, + 0x2b2b081919080808, 0x2b2b0819192b1919, 0x2b2b082b0808082b, 0x2b2b082b08082b2b, + 0x2b2b082b082b082b, 0x2b2b082b082b2b08, 0x2b2b082b082b2b2b, 0x2b2b082b2b08082b, + 0x2b2b082b2b082b08, 0x2b2b082b2b082b2b, 0x2b2b082b2b2b2b08, 0x2b2b190808080819, + 0x2b2b190808081908, 0x2b2b190808190808, 0x2b2b190819080808, 0x2b2b19082b082b19, + 0x2b2b19082b2b1908, 0x2b2b191908080808, 0x2b2b191908192b19, 0x2b2b192b19190819, + 0x2b2b2b0808082b2b, 0x2b2b2b08082b2b08, 0x2b2b2b082b2b082b, 0x2b2b2b1919191908, + 0x2b2b2b192b08192b, 0x2b2b2b2b08082b08, 0x2b2b2b2b08082b2b, 0x2b2b2b2b082b0808, + 0x2b2b2b2b082b082b, 0x2b2b2b2b082b2b08, 0x2b2b2b2b2b082b08, 0x2b2b2b2b2b2b2b2b, +GGML_TABLE_END() + +GGML_TABLE_BEGIN(uint32_t, iq3xxs_grid, 256) + 0x04040404, 0x04040414, 0x04040424, 0x04040c0c, 0x04040c1c, 0x04040c3e, 0x04041404, 0x04041414, + 0x04041c0c, 0x04042414, 0x04043e1c, 0x04043e2c, 0x040c040c, 0x040c041c, 0x040c0c04, 0x040c0c14, + 0x040c140c, 0x040c142c, 0x040c1c04, 0x040c1c14, 0x040c240c, 0x040c2c24, 0x040c3e04, 0x04140404, + 0x04140414, 0x04140424, 0x04140c0c, 0x04141404, 0x04141414, 0x04141c0c, 0x04141c1c, 0x04141c3e, + 0x04142c0c, 0x04142c3e, 0x04143e2c, 0x041c040c, 0x041c043e, 0x041c0c04, 0x041c0c14, 0x041c142c, + 0x041c3e04, 0x04240c1c, 0x04241c3e, 0x04242424, 0x04242c3e, 0x04243e1c, 0x04243e2c, 0x042c040c, + 0x042c043e, 0x042c1c14, 0x042c2c14, 0x04341c2c, 0x04343424, 0x043e0c04, 0x043e0c24, 0x043e0c34, + 0x043e241c, 0x043e340c, 0x0c04040c, 0x0c04041c, 0x0c040c04, 0x0c040c14, 0x0c04140c, 0x0c04141c, + 0x0c041c04, 0x0c041c14, 0x0c041c24, 0x0c04243e, 0x0c042c04, 0x0c0c0404, 0x0c0c0414, 0x0c0c0c0c, + 0x0c0c1404, 0x0c0c1414, 0x0c14040c, 0x0c14041c, 0x0c140c04, 0x0c140c14, 0x0c14140c, 0x0c141c04, + 0x0c143e14, 0x0c1c0404, 0x0c1c0414, 0x0c1c1404, 0x0c1c1c0c, 0x0c1c2434, 0x0c1c3434, 0x0c24040c, + 0x0c24042c, 0x0c242c04, 0x0c2c1404, 0x0c2c1424, 0x0c2c2434, 0x0c2c3e0c, 0x0c34042c, 0x0c3e1414, + 0x0c3e2404, 0x14040404, 0x14040414, 0x14040c0c, 0x14040c1c, 0x14041404, 0x14041414, 0x14041434, + 0x14041c0c, 0x14042414, 0x140c040c, 0x140c041c, 0x140c042c, 0x140c0c04, 0x140c0c14, 0x140c140c, + 0x140c1c04, 0x140c341c, 0x140c343e, 0x140c3e04, 0x14140404, 0x14140414, 0x14140c0c, 0x14140c3e, + 0x14141404, 0x14141414, 0x14141c3e, 0x14142404, 0x14142c2c, 0x141c040c, 0x141c0c04, 0x141c0c24, + 0x141c3e04, 0x141c3e24, 0x14241c2c, 0x14242c1c, 0x142c041c, 0x142c143e, 0x142c240c, 0x142c3e24, + 0x143e040c, 0x143e041c, 0x143e0c34, 0x143e242c, 0x1c04040c, 0x1c040c04, 0x1c040c14, 0x1c04140c, + 0x1c04141c, 0x1c042c04, 0x1c04342c, 0x1c043e14, 0x1c0c0404, 0x1c0c0414, 0x1c0c1404, 0x1c0c1c0c, + 0x1c0c2424, 0x1c0c2434, 0x1c14040c, 0x1c14041c, 0x1c140c04, 0x1c14142c, 0x1c142c14, 0x1c143e14, + 0x1c1c0c0c, 0x1c1c1c1c, 0x1c241c04, 0x1c24243e, 0x1c243e14, 0x1c2c0404, 0x1c2c0434, 0x1c2c1414, + 0x1c2c2c2c, 0x1c340c24, 0x1c341c34, 0x1c34341c, 0x1c3e1c1c, 0x1c3e3404, 0x24040424, 0x24040c3e, + 0x24041c2c, 0x24041c3e, 0x24042c1c, 0x24042c3e, 0x240c3e24, 0x24141404, 0x24141c3e, 0x24142404, + 0x24143404, 0x24143434, 0x241c043e, 0x241c242c, 0x24240424, 0x24242c0c, 0x24243424, 0x242c142c, + 0x242c241c, 0x242c3e04, 0x243e042c, 0x243e0c04, 0x243e0c14, 0x243e1c04, 0x2c040c14, 0x2c04240c, + 0x2c043e04, 0x2c0c0404, 0x2c0c0434, 0x2c0c1434, 0x2c0c2c2c, 0x2c140c24, 0x2c141c14, 0x2c143e14, + 0x2c1c0414, 0x2c1c2c1c, 0x2c240c04, 0x2c24141c, 0x2c24143e, 0x2c243e14, 0x2c2c0414, 0x2c2c1c0c, + 0x2c342c04, 0x2c3e1424, 0x2c3e2414, 0x34041424, 0x34042424, 0x34042434, 0x34043424, 0x340c140c, + 0x340c340c, 0x34140c3e, 0x34143424, 0x341c1c04, 0x341c1c34, 0x34242424, 0x342c042c, 0x342c2c14, + 0x34341c1c, 0x343e041c, 0x343e140c, 0x3e04041c, 0x3e04042c, 0x3e04043e, 0x3e040c04, 0x3e041c14, + 0x3e042c14, 0x3e0c1434, 0x3e0c2404, 0x3e140c14, 0x3e14242c, 0x3e142c14, 0x3e1c0404, 0x3e1c0c2c, + 0x3e1c1c1c, 0x3e1c3404, 0x3e24140c, 0x3e24240c, 0x3e2c0404, 0x3e2c0414, 0x3e2c1424, 0x3e341c04, +GGML_TABLE_END() + +GGML_TABLE_BEGIN(uint32_t, iq3s_grid, 512) + 0x01010101, 0x01010103, 0x01010105, 0x0101010b, 0x0101010f, 0x01010301, 0x01010303, 0x01010305, + 0x01010309, 0x0101030d, 0x01010501, 0x01010503, 0x0101050b, 0x01010707, 0x01010901, 0x01010905, + 0x0101090b, 0x0101090f, 0x01010b03, 0x01010b07, 0x01010d01, 0x01010d05, 0x01010f03, 0x01010f09, + 0x01010f0f, 0x01030101, 0x01030103, 0x01030105, 0x01030109, 0x01030301, 0x01030303, 0x0103030b, + 0x01030501, 0x01030507, 0x0103050f, 0x01030703, 0x0103070b, 0x01030909, 0x01030d03, 0x01030d0b, + 0x01030f05, 0x01050101, 0x01050103, 0x0105010b, 0x0105010f, 0x01050301, 0x01050307, 0x0105030d, + 0x01050503, 0x0105050b, 0x01050701, 0x01050709, 0x01050905, 0x0105090b, 0x0105090f, 0x01050b03, + 0x01050b07, 0x01050f01, 0x01050f07, 0x01070107, 0x01070303, 0x0107030b, 0x01070501, 0x01070505, + 0x01070703, 0x01070707, 0x0107070d, 0x01070909, 0x01070b01, 0x01070b05, 0x01070d0f, 0x01070f03, + 0x01070f0b, 0x01090101, 0x01090307, 0x0109030f, 0x01090503, 0x01090509, 0x01090705, 0x01090901, + 0x01090907, 0x01090b03, 0x01090f01, 0x010b0105, 0x010b0109, 0x010b0501, 0x010b0505, 0x010b050d, + 0x010b0707, 0x010b0903, 0x010b090b, 0x010b090f, 0x010b0d0d, 0x010b0f07, 0x010d010d, 0x010d0303, + 0x010d0307, 0x010d0703, 0x010d0b05, 0x010d0f03, 0x010f0101, 0x010f0105, 0x010f0109, 0x010f0501, + 0x010f0505, 0x010f050d, 0x010f0707, 0x010f0b01, 0x010f0b09, 0x03010101, 0x03010103, 0x03010105, + 0x03010109, 0x03010301, 0x03010303, 0x03010307, 0x0301030b, 0x0301030f, 0x03010501, 0x03010505, + 0x03010703, 0x03010709, 0x0301070d, 0x03010b09, 0x03010b0d, 0x03010d03, 0x03010f05, 0x03030101, + 0x03030103, 0x03030107, 0x0303010d, 0x03030301, 0x03030309, 0x03030503, 0x03030701, 0x03030707, + 0x03030903, 0x03030b01, 0x03030b05, 0x03030f01, 0x03030f0d, 0x03050101, 0x03050305, 0x0305030b, + 0x0305030f, 0x03050501, 0x03050509, 0x03050705, 0x03050901, 0x03050907, 0x03050b0b, 0x03050d01, + 0x03050f05, 0x03070103, 0x03070109, 0x0307010f, 0x03070301, 0x03070307, 0x03070503, 0x0307050f, + 0x03070701, 0x03070709, 0x03070903, 0x03070d05, 0x03070f01, 0x03090107, 0x0309010b, 0x03090305, + 0x03090309, 0x03090703, 0x03090707, 0x03090905, 0x0309090d, 0x03090b01, 0x03090b09, 0x030b0103, + 0x030b0301, 0x030b0307, 0x030b0503, 0x030b0701, 0x030b0705, 0x030b0b03, 0x030d0501, 0x030d0509, + 0x030d050f, 0x030d0909, 0x030d090d, 0x030f0103, 0x030f0107, 0x030f0301, 0x030f0305, 0x030f0503, + 0x030f070b, 0x030f0903, 0x030f0d05, 0x030f0f01, 0x05010101, 0x05010103, 0x05010107, 0x0501010b, + 0x0501010f, 0x05010301, 0x05010305, 0x05010309, 0x0501030d, 0x05010503, 0x05010507, 0x0501050f, + 0x05010701, 0x05010705, 0x05010903, 0x05010907, 0x0501090b, 0x05010b01, 0x05010b05, 0x05010d0f, + 0x05010f01, 0x05010f07, 0x05010f0b, 0x05030101, 0x05030105, 0x05030301, 0x05030307, 0x0503030f, + 0x05030505, 0x0503050b, 0x05030703, 0x05030709, 0x05030905, 0x05030b03, 0x05050103, 0x05050109, + 0x0505010f, 0x05050503, 0x05050507, 0x05050701, 0x0505070f, 0x05050903, 0x05050b07, 0x05050b0f, + 0x05050f03, 0x05050f09, 0x05070101, 0x05070105, 0x0507010b, 0x05070303, 0x05070505, 0x05070509, + 0x05070703, 0x05070707, 0x05070905, 0x05070b01, 0x05070d0d, 0x05090103, 0x0509010f, 0x05090501, + 0x05090507, 0x05090705, 0x0509070b, 0x05090903, 0x05090f05, 0x05090f0b, 0x050b0109, 0x050b0303, + 0x050b0505, 0x050b070f, 0x050b0901, 0x050b0b07, 0x050b0f01, 0x050d0101, 0x050d0105, 0x050d010f, + 0x050d0503, 0x050d0b0b, 0x050d0d03, 0x050f010b, 0x050f0303, 0x050f050d, 0x050f0701, 0x050f0907, + 0x050f0b01, 0x07010105, 0x07010303, 0x07010307, 0x0701030b, 0x0701030f, 0x07010505, 0x07010703, + 0x07010707, 0x0701070b, 0x07010905, 0x07010909, 0x0701090f, 0x07010b03, 0x07010d07, 0x07010f03, + 0x07030103, 0x07030107, 0x0703010b, 0x07030309, 0x07030503, 0x07030507, 0x07030901, 0x07030d01, + 0x07030f05, 0x07030f0d, 0x07050101, 0x07050305, 0x07050501, 0x07050705, 0x07050709, 0x07050b01, + 0x07070103, 0x07070301, 0x07070309, 0x07070503, 0x07070507, 0x0707050f, 0x07070701, 0x07070903, + 0x07070907, 0x0707090f, 0x07070b0b, 0x07070f07, 0x07090107, 0x07090303, 0x0709030d, 0x07090505, + 0x07090703, 0x07090b05, 0x07090d01, 0x07090d09, 0x070b0103, 0x070b0301, 0x070b0305, 0x070b050b, + 0x070b0705, 0x070b0909, 0x070b0b0d, 0x070b0f07, 0x070d030d, 0x070d0903, 0x070f0103, 0x070f0107, + 0x070f0501, 0x070f0505, 0x070f070b, 0x09010101, 0x09010109, 0x09010305, 0x09010501, 0x09010509, + 0x0901050f, 0x09010705, 0x09010903, 0x09010b01, 0x09010f01, 0x09030105, 0x0903010f, 0x09030303, + 0x09030307, 0x09030505, 0x09030701, 0x0903070b, 0x09030907, 0x09030b03, 0x09030b0b, 0x09050103, + 0x09050107, 0x09050301, 0x0905030b, 0x09050503, 0x09050707, 0x09050901, 0x09050b0f, 0x09050d05, + 0x09050f01, 0x09070109, 0x09070303, 0x09070307, 0x09070501, 0x09070505, 0x09070703, 0x0907070b, + 0x09090101, 0x09090105, 0x09090509, 0x0909070f, 0x09090901, 0x09090f03, 0x090b010b, 0x090b010f, + 0x090b0503, 0x090b0d05, 0x090d0307, 0x090d0709, 0x090d0d01, 0x090f0301, 0x090f030b, 0x090f0701, + 0x090f0907, 0x090f0b03, 0x0b010105, 0x0b010301, 0x0b010309, 0x0b010505, 0x0b010901, 0x0b010909, + 0x0b01090f, 0x0b010b05, 0x0b010d0d, 0x0b010f09, 0x0b030103, 0x0b030107, 0x0b03010b, 0x0b030305, + 0x0b030503, 0x0b030705, 0x0b030f05, 0x0b050101, 0x0b050303, 0x0b050507, 0x0b050701, 0x0b05070d, + 0x0b050b07, 0x0b070105, 0x0b07010f, 0x0b070301, 0x0b07050f, 0x0b070909, 0x0b070b03, 0x0b070d0b, + 0x0b070f07, 0x0b090103, 0x0b090109, 0x0b090501, 0x0b090705, 0x0b09090d, 0x0b0b0305, 0x0b0b050d, + 0x0b0b0b03, 0x0b0b0b07, 0x0b0d0905, 0x0b0f0105, 0x0b0f0109, 0x0b0f0505, 0x0d010303, 0x0d010307, + 0x0d01030b, 0x0d010703, 0x0d010707, 0x0d010d01, 0x0d030101, 0x0d030501, 0x0d03050f, 0x0d030d09, + 0x0d050305, 0x0d050709, 0x0d050905, 0x0d050b0b, 0x0d050d05, 0x0d050f01, 0x0d070101, 0x0d070309, + 0x0d070503, 0x0d070901, 0x0d09050b, 0x0d090907, 0x0d090d05, 0x0d0b0101, 0x0d0b0107, 0x0d0b0709, + 0x0d0b0d01, 0x0d0d010b, 0x0d0d0901, 0x0d0f0303, 0x0d0f0307, 0x0f010101, 0x0f010109, 0x0f01010f, + 0x0f010501, 0x0f010505, 0x0f01070d, 0x0f010901, 0x0f010b09, 0x0f010d05, 0x0f030105, 0x0f030303, + 0x0f030509, 0x0f030907, 0x0f03090b, 0x0f050103, 0x0f050109, 0x0f050301, 0x0f05030d, 0x0f050503, + 0x0f050701, 0x0f050b03, 0x0f070105, 0x0f070705, 0x0f07070b, 0x0f070b07, 0x0f090103, 0x0f09010b, + 0x0f090307, 0x0f090501, 0x0f090b01, 0x0f0b0505, 0x0f0b0905, 0x0f0d0105, 0x0f0d0703, 0x0f0f0101, +GGML_TABLE_END() + +// TODO: fix name to kvalues_iq4_nl +GGML_TABLE_BEGIN(int8_t, kvalues_iq4nl, 16) + -127, -104, -83, -65, -49, -35, -22, -10, 1, 13, 25, 38, 53, 69, 89, 113, +GGML_TABLE_END() + +// e2m1 values (doubled), shared by MXFP4 and NVFP4 +// ref: https://www.opencompute.org/documents/ocp-microscaling-formats-mx-v1-0-spec-final-pdf +GGML_TABLE_BEGIN(int8_t, kvalues_fp4, 16) + 0, 1, 2, 3, 4, 6, 8, 12, 0, -1, -2, -3, -4, -6, -8, -12, +GGML_TABLE_END() +#define kvalues_mxfp4 kvalues_fp4 + +#define NGRID_IQ1S 2048 +#define IQ1S_DELTA 0.125f +#define IQ1M_DELTA 0.125f +#if defined(GGML_COMMON_IMPL_C) || defined(GGML_COMMON_IMPL_CPP) +GGML_TABLE_BEGIN(uint64_t, iq1s_grid, NGRID_IQ1S) + 0xffffffffffffffff, 0xffffffffffffff01, 0xffffffffffff0000, 0xffffffffffff01ff, + 0xffffffffffff0101, 0xffffffffff00ff00, 0xffffffffff000000, 0xffffffffff01ffff, + 0xffffffffff01ff01, 0xffffffffff0101ff, 0xffffffffff010101, 0xffffffff00ff0000, + 0xffffffff0000ff00, 0xffffffff000000ff, 0xffffffff00000001, 0xffffffff00010000, + 0xffffffff01ffffff, 0xffffffff01ffff01, 0xffffffff01ff01ff, 0xffffffff01ff0101, + 0xffffffff01000000, 0xffffffff0101ffff, 0xffffffff0101ff01, 0xffffffff010101ff, + 0xffffffff01010101, 0xffffff00ffff00ff, 0xffffff00ffff0000, 0xffffff00ff00ff00, + 0xffffff00ff0000ff, 0xffffff00ff000001, 0xffffff00ff000100, 0xffffff00ff000101, + 0xffffff00ff010000, 0xffffff0000ffff00, 0xffffff0000ff0001, 0xffffff0000ff0100, + 0xffffff000000ff01, 0xffffff0000000000, 0xffffff0000000101, 0xffffff000001ff00, + 0xffffff00000100ff, 0xffffff0000010001, 0xffffff00000101ff, 0xffffff0001ff0000, + 0xffffff000100ff00, 0xffffff00010000ff, 0xffffff0001000001, 0xffffff0001010000, + 0xffffff01ffffffff, 0xffffff01ffffff01, 0xffffff01ffff01ff, 0xffffff01ffff0101, + 0xffffff01ff000000, 0xffffff01ff01ffff, 0xffffff01ff01ff01, 0xffffff01ff0101ff, + 0xffffff01ff010101, 0xffffff0100ff0000, 0xffffff010000ff00, 0xffffff0100000100, + 0xffffff01000100ff, 0xffffff0100010100, 0xffffff0101ffffff, 0xffffff0101ffff01, + 0xffffff0101ff01ff, 0xffffff0101ff0101, 0xffffff010100ff00, 0xffffff0101000000, + 0xffffff0101000100, 0xffffff010101ffff, 0xffffff010101ff01, 0xffffff01010101ff, + 0xffffff0101010101, 0xffff00ffff00ff00, 0xffff00ffff0000ff, 0xffff00ffff000001, + 0xffff00ffff010000, 0xffff00ff00ffff00, 0xffff00ff00ff0100, 0xffff00ff00000000, + 0xffff00ff00000101, 0xffff00ff000100ff, 0xffff00ff00010000, 0xffff00ff0100ff00, + 0xffff00ff01000100, 0xffff00ff01010000, 0xffff0000ffffff00, 0xffff0000ffff00ff, + 0xffff0000ffff0000, 0xffff0000ffff0001, 0xffff0000ff000000, 0xffff0000ff0001ff, + 0xffff0000ff000101, 0xffff0000ff010100, 0xffff000000ffffff, 0xffff000000ff0000, + 0xffff000000ff0101, 0xffff00000000ffff, 0xffff00000000ff00, 0xffff0000000000ff, + 0xffff000000000000, 0xffff000000000001, 0xffff000000000100, 0xffff00000001ffff, + 0xffff00000001ff01, 0xffff000000010000, 0xffff0000000101ff, 0xffff000000010101, + 0xffff000001ffff00, 0xffff00000100ff00, 0xffff000001000000, 0xffff0000010001ff, + 0xffff000001000101, 0xffff00000101ff00, 0xffff0000010100ff, 0xffff000001010000, + 0xffff000001010001, 0xffff000001010100, 0xffff0001ff0000ff, 0xffff0001ff000100, + 0xffff000100ffff00, 0xffff000100ff00ff, 0xffff00010000ffff, 0xffff00010000ff01, + 0xffff000100000000, 0xffff0001000001ff, 0xffff00010001ffff, 0xffff00010001ff00, + 0xffff000100010001, 0xffff000100010100, 0xffff000101ff0000, 0xffff00010100ff00, + 0xffff0001010000ff, 0xffff000101000100, 0xffff01ffffffffff, 0xffff01ffffffff01, + 0xffff01ffffff01ff, 0xffff01ffffff0101, 0xffff01ffff000000, 0xffff01ffff01ffff, + 0xffff01ffff01ff01, 0xffff01ffff0101ff, 0xffff01ffff010101, 0xffff01ff00ff0000, + 0xffff01ff0000ff00, 0xffff01ff00000001, 0xffff01ff00010000, 0xffff01ff01ffffff, + 0xffff01ff01ffff01, 0xffff01ff01ff01ff, 0xffff01ff01ff0101, 0xffff01ff01000000, + 0xffff01ff0101ffff, 0xffff01ff0101ff01, 0xffff01ff010101ff, 0xffff01ff01010101, + 0xffff0100ffff0000, 0xffff0100ff00ff00, 0xffff0100ff0000ff, 0xffff0100ff000100, + 0xffff0100ff0100ff, 0xffff0100ff010000, 0xffff010000ffff00, 0xffff01000000ffff, + 0xffff01000000ff00, 0xffff010000000000, 0xffff01000001ff00, 0xffff0100000100ff, + 0xffff010000010100, 0xffff01000100ff00, 0xffff0100010000ff, 0xffff010001000001, + 0xffff010001000100, 0xffff010001010000, 0xffff0101ffffffff, 0xffff0101ffffff01, + 0xffff0101ffff01ff, 0xffff0101ffff0101, 0xffff0101ff000000, 0xffff0101ff01ffff, + 0xffff0101ff01ff01, 0xffff0101ff0101ff, 0xffff0101ff010101, 0xffff010100ff0000, + 0xffff01010000ff00, 0xffff010100000100, 0xffff01010001ff00, 0xffff010100010000, + 0xffff010101ffffff, 0xffff010101ffff01, 0xffff010101ff0000, 0xffff010101ff01ff, + 0xffff010101ff0101, 0xffff010101000000, 0xffff01010101ffff, 0xffff01010101ff01, + 0xffff0101010101ff, 0xffff010101010101, 0xff00ffffff00ffff, 0xff00ffffff00ff00, + 0xff00ffffff0000ff, 0xff00ffffff000100, 0xff00ffffff0100ff, 0xff00ffffff010000, + 0xff00ffff00ffff00, 0xff00ffff00ff00ff, 0xff00ffff0000ffff, 0xff00ffff00000000, + 0xff00ffff000001ff, 0xff00ffff0001ff00, 0xff00ffff000100ff, 0xff00ffff00010000, + 0xff00ffff00010100, 0xff00ffff0100ff00, 0xff00ffff010000ff, 0xff00ffff01000001, + 0xff00ffff0101ff00, 0xff00ffff01010000, 0xff00ff00ffffff00, 0xff00ff00ffff00ff, + 0xff00ff00ffff0001, 0xff00ff00ffff0100, 0xff00ff00ff00ffff, 0xff00ff00ff00ff01, + 0xff00ff00ff000000, 0xff00ff00ff0001ff, 0xff00ff00ff01ff00, 0xff00ff00ff0100ff, + 0xff00ff00ff010100, 0xff00ff0000ff0000, 0xff00ff0000ff0101, 0xff00ff000000ffff, + 0xff00ff000000ff00, 0xff00ff000000ff01, 0xff00ff00000000ff, 0xff00ff0000000000, + 0xff00ff0000000001, 0xff00ff0000000100, 0xff00ff000001ffff, 0xff00ff0000010000, + 0xff00ff0001ff00ff, 0xff00ff000100ff01, 0xff00ff0001000000, 0xff00ff000101ff00, + 0xff00ff00010100ff, 0xff00ff01ff00ff00, 0xff00ff01ff0000ff, 0xff00ff01ff000001, + 0xff00ff01ff010000, 0xff00ff0100ffffff, 0xff00ff0100ff0001, 0xff00ff0100ff0100, + 0xff00ff010000ff01, 0xff00ff0100000000, 0xff00ff01000001ff, 0xff00ff0100000101, + 0xff00ff01000100ff, 0xff00ff0100010001, 0xff00ff0101ff0000, 0xff00ff010100ff00, + 0xff00ff01010000ff, 0xff00ff0101000001, 0xff00ff0101010000, 0xff0000ffffffff00, + 0xff0000ffffff0001, 0xff0000ffffff0100, 0xff0000ffff0000ff, 0xff0000ffff000000, + 0xff0000ffff0001ff, 0xff0000ffff000100, 0xff0000ffff01ff00, 0xff0000ffff010001, + 0xff0000ff00ffff00, 0xff0000ff00ff0000, 0xff0000ff00ff0001, 0xff0000ff00ff01ff, + 0xff0000ff00ff0101, 0xff0000ff0000ff00, 0xff0000ff000000ff, 0xff0000ff00000000, + 0xff0000ff00000001, 0xff0000ff00000100, 0xff0000ff0001ff01, 0xff0000ff00010000, + 0xff0000ff000101ff, 0xff0000ff01ff00ff, 0xff0000ff01ff0100, 0xff0000ff0100ffff, + 0xff0000ff010000ff, 0xff0000ff01000000, 0xff0000ff010001ff, 0xff0000ff01000100, + 0xff0000ff01000101, 0xff0000ff0101ff00, 0xff0000ff010100ff, 0xff0000ff01010000, + 0xff0000ff01010100, 0xff000000ffffff01, 0xff000000ffff0000, 0xff000000ffff0101, + 0xff000000ff00ff00, 0xff000000ff0000ff, 0xff000000ff000000, 0xff000000ff000001, + 0xff000000ff000100, 0xff000000ff01ffff, 0xff000000ff01ff01, 0xff000000ff010000, + 0xff000000ff0101ff, 0xff000000ff010101, 0xff00000000ffff00, 0xff00000000ff00ff, + 0xff00000000ff0000, 0xff00000000ff0001, 0xff0000000000ff00, 0xff0000000000ff01, + 0xff000000000000ff, 0xff00000000000000, 0xff00000000000001, 0xff00000000000100, + 0xff00000000000101, 0xff0000000001ff00, 0xff000000000100ff, 0xff00000000010000, + 0xff00000000010001, 0xff00000000010100, 0xff00000001ffffff, 0xff00000001ffff01, + 0xff00000001ff00ff, 0xff00000001ff0000, 0xff00000001ff01ff, 0xff00000001ff0101, + 0xff0000000100ffff, 0xff0000000100ff00, 0xff000000010000ff, 0xff00000001000000, + 0xff00000001000001, 0xff00000001000100, 0xff00000001000101, 0xff0000000101ffff, + 0xff0000000101ff01, 0xff00000001010000, 0xff000001ffffff00, 0xff000001ffff00ff, + 0xff000001ffff0000, 0xff000001ffff0001, 0xff000001ff000000, 0xff000001ff000001, + 0xff000001ff0001ff, 0xff000001ff000101, 0xff000001ff01ff00, 0xff000001ff010001, + 0xff00000100ffffff, 0xff00000100ffff01, 0xff00000100ff00ff, 0xff00000100ff0000, + 0xff00000100ff01ff, 0xff00000100ff0101, 0xff0000010000ff00, 0xff00000100000000, + 0xff00000100000001, 0xff000001000001ff, 0xff00000100000100, 0xff0000010001ff00, + 0xff000001000100ff, 0xff00000100010000, 0xff000001000101ff, 0xff00000100010100, + 0xff00000100010101, 0xff00000101ff0001, 0xff00000101ff0101, 0xff0000010100ff01, + 0xff00000101000000, 0xff000001010100ff, 0xff00000101010100, 0xff0001ffff00ff00, + 0xff0001ffff000001, 0xff0001ffff010000, 0xff0001ff00ffff00, 0xff0001ff00ff00ff, + 0xff0001ff00ff0001, 0xff0001ff00ff0100, 0xff0001ff0000ffff, 0xff0001ff00000000, + 0xff0001ff000001ff, 0xff0001ff00000101, 0xff0001ff0001ffff, 0xff0001ff0001ff00, + 0xff0001ff000100ff, 0xff0001ff00010001, 0xff0001ff00010100, 0xff0001ff01ff0000, + 0xff0001ff0100ff00, 0xff0001ff010000ff, 0xff0001ff01010000, 0xff000100ff00ffff, + 0xff000100ff00ff01, 0xff000100ff000000, 0xff000100ff000101, 0xff000100ff01ff00, + 0xff000100ff010000, 0xff00010000ffff01, 0xff00010000ff00ff, 0xff00010000ff0000, + 0xff00010000ff01ff, 0xff0001000000ff00, 0xff000100000000ff, 0xff00010000000000, + 0xff00010000000001, 0xff00010000000100, 0xff00010000000101, 0xff0001000001ffff, + 0xff00010000010000, 0xff00010000010101, 0xff00010001ff0100, 0xff0001000100ff00, + 0xff0001000100ff01, 0xff00010001000000, 0xff000100010001ff, 0xff0001000101ff00, + 0xff00010001010001, 0xff00010001010100, 0xff000101ffff0100, 0xff000101ff000001, + 0xff000101ff0100ff, 0xff000101ff010001, 0xff00010100ff00ff, 0xff00010100ff0001, + 0xff00010100ff0100, 0xff0001010000ffff, 0xff0001010000ff01, 0xff00010100000000, + 0xff000101000001ff, 0xff0001010001ff00, 0xff00010100010001, 0xff00010100010100, + 0xff00010101ff0000, 0xff0001010100ff00, 0xff00010101000001, 0xff00010101000101, + 0xff01ffffffffffff, 0xff01ffffffffff01, 0xff01ffffffff01ff, 0xff01ffffffff0101, + 0xff01ffffff000000, 0xff01ffffff01ffff, 0xff01ffffff01ff01, 0xff01ffffff010000, + 0xff01ffffff0101ff, 0xff01ffffff010101, 0xff01ffff00ff0000, 0xff01ffff0000ff00, + 0xff01ffff00000100, 0xff01ffff0001ff00, 0xff01ffff00010000, 0xff01ffff01ffffff, + 0xff01ffff01ffff01, 0xff01ffff01ff01ff, 0xff01ffff01ff0101, 0xff01ffff01000000, + 0xff01ffff0101ffff, 0xff01ffff0101ff01, 0xff01ffff01010000, 0xff01ffff010101ff, + 0xff01ffff01010101, 0xff01ff00ffff0000, 0xff01ff00ff00ff00, 0xff01ff00ff0000ff, + 0xff01ff00ff000100, 0xff01ff00ff010000, 0xff01ff0000ffff01, 0xff01ff0000ff00ff, + 0xff01ff0000ff0100, 0xff01ff0000000000, 0xff01ff00000001ff, 0xff01ff0000000101, + 0xff01ff000001ff00, 0xff01ff00000100ff, 0xff01ff0000010000, 0xff01ff0000010001, + 0xff01ff0001ff0000, 0xff01ff000100ffff, 0xff01ff0001000001, 0xff01ff0001000100, + 0xff01ff0001010000, 0xff01ff01ffffff00, 0xff01ff01ffff01ff, 0xff01ff01ffff0101, + 0xff01ff01ff00ff00, 0xff01ff01ff000000, 0xff01ff01ff01ffff, 0xff01ff01ff01ff01, + 0xff01ff01ff0101ff, 0xff01ff01ff010101, 0xff01ff0100ff0000, 0xff01ff010000ff00, + 0xff01ff0100000001, 0xff01ff0100000100, 0xff01ff0100010000, 0xff01ff0101ffff00, + 0xff01ff0101ff01ff, 0xff01ff0101ff0101, 0xff01ff010100ff00, 0xff01ff0101000000, + 0xff01ff010101ffff, 0xff01ff010101ff01, 0xff01ff01010101ff, 0xff01ff0101010101, + 0xff0100ffffff0000, 0xff0100ffff0000ff, 0xff0100ffff000001, 0xff0100ffff000100, + 0xff0100ffff010000, 0xff0100ff00ff00ff, 0xff0100ff00ff0000, 0xff0100ff00ff0001, + 0xff0100ff00ff0100, 0xff0100ff0000ff01, 0xff0100ff00000000, 0xff0100ff000001ff, + 0xff0100ff00000101, 0xff0100ff00010001, 0xff0100ff01ff0000, 0xff0100ff0100ff00, + 0xff0100ff010000ff, 0xff0100ff01000100, 0xff0100ff0101ff00, 0xff0100ff01010000, + 0xff010000ffff0100, 0xff010000ff000000, 0xff010000ff01ff00, 0xff010000ff010100, + 0xff01000000ffffff, 0xff01000000ff0000, 0xff01000000ff01ff, 0xff0100000000ff00, + 0xff010000000000ff, 0xff01000000000000, 0xff01000000000100, 0xff0100000001ff01, + 0xff01000000010000, 0xff010000000101ff, 0xff01000001ff0100, 0xff0100000100ffff, + 0xff010000010000ff, 0xff01000001000000, 0xff010000010001ff, 0xff01000001000101, + 0xff0100000101ff00, 0xff010000010100ff, 0xff01000001010001, 0xff01000001010100, + 0xff010001ffff0000, 0xff010001ff00ffff, 0xff010001ff00ff01, 0xff010001ff000100, + 0xff010001ff010000, 0xff01000100ffff00, 0xff01000100ff0100, 0xff01000100000000, + 0xff0100010001ffff, 0xff0100010001ff00, 0xff01000100010100, 0xff01000101ff00ff, + 0xff01000101ff0001, 0xff0100010100ffff, 0xff01000101000101, 0xff0101ffffffffff, + 0xff0101ffffffff01, 0xff0101ffffff01ff, 0xff0101ffffff0101, 0xff0101ffff000000, + 0xff0101ffff01ffff, 0xff0101ffff01ff01, 0xff0101ffff0101ff, 0xff0101ffff010101, + 0xff0101ff00ff0000, 0xff0101ff0000ff00, 0xff0101ff000000ff, 0xff0101ff00010000, + 0xff0101ff01ffffff, 0xff0101ff01ffff01, 0xff0101ff01ff01ff, 0xff0101ff01ff0101, + 0xff0101ff0101ffff, 0xff0101ff0101ff01, 0xff0101ff010101ff, 0xff0101ff01010101, + 0xff010100ffff0100, 0xff010100ff00ff00, 0xff010100ff0000ff, 0xff010100ff000100, + 0xff010100ff010000, 0xff01010000ff0001, 0xff01010000ff0100, 0xff0101000000ff01, + 0xff01010000000000, 0xff0101000001ff00, 0xff010100000100ff, 0xff01010000010001, + 0xff01010000010100, 0xff01010001ff0000, 0xff0101000100ffff, 0xff01010001000001, + 0xff01010001000100, 0xff010100010100ff, 0xff01010001010000, 0xff010101ffffffff, + 0xff010101ffffff01, 0xff010101ffff01ff, 0xff010101ffff0101, 0xff010101ff01ffff, + 0xff010101ff01ff01, 0xff010101ff0101ff, 0xff010101ff010101, 0xff01010100ff0000, + 0xff0101010000ff00, 0xff01010100000001, 0xff01010100000100, 0xff01010100010000, + 0xff01010101ffffff, 0xff01010101ffff01, 0xff01010101ff01ff, 0xff01010101ff0101, + 0xff01010101000000, 0xff0101010101ffff, 0xff0101010101ff01, 0xff010101010101ff, + 0xff01010101010101, 0x00ffffffffff0000, 0x00ffffffff00ff00, 0x00ffffffff000001, + 0x00ffffffff010000, 0x00ffffff00ff0100, 0x00ffffff0000ff01, 0x00ffffff00000000, + 0x00ffffff000001ff, 0x00ffffff00000101, 0x00ffffff0001ff00, 0x00ffffff000100ff, + 0x00ffffff00010001, 0x00ffffff010000ff, 0x00ffffff01000100, 0x00ffffff0101ff00, + 0x00ffffff01010001, 0x00ffff00ffffffff, 0x00ffff00ffffff00, 0x00ffff00ffff00ff, + 0x00ffff00ffff0001, 0x00ffff00ffff0100, 0x00ffff00ff00ff01, 0x00ffff00ff000000, + 0x00ffff00ff000001, 0x00ffff00ff0001ff, 0x00ffff00ff000101, 0x00ffff00ff01ff00, + 0x00ffff00ff010001, 0x00ffff00ff010100, 0x00ffff0000ff0000, 0x00ffff0000ff01ff, + 0x00ffff0000ff0101, 0x00ffff000000ff00, 0x00ffff00000000ff, 0x00ffff0000000000, + 0x00ffff0000000001, 0x00ffff0000000100, 0x00ffff0000000101, 0x00ffff0000010000, + 0x00ffff00000101ff, 0x00ffff0000010101, 0x00ffff0001ffff00, 0x00ffff0001ff00ff, + 0x00ffff0001ff0001, 0x00ffff000100ffff, 0x00ffff000100ff01, 0x00ffff0001000000, + 0x00ffff000101ffff, 0x00ffff000101ff00, 0x00ffff000101ff01, 0x00ffff01ffff0000, + 0x00ffff01ff00ff00, 0x00ffff01ff0000ff, 0x00ffff01ff000001, 0x00ffff01ff010000, + 0x00ffff0100ffff00, 0x00ffff010000ff01, 0x00ffff0100000000, 0x00ffff0100000101, + 0x00ffff01000100ff, 0x00ffff0100010100, 0x00ffff0101ff0100, 0x00ffff01010000ff, + 0x00ffff0101010000, 0x00ff00ffffffff00, 0x00ff00ffff000000, 0x00ff00ffff000100, + 0x00ff00ffff010100, 0x00ff00ff00ff0000, 0x00ff00ff00ff01ff, 0x00ff00ff00ff0101, + 0x00ff00ff0000ff00, 0x00ff00ff000000ff, 0x00ff00ff00000000, 0x00ff00ff00000001, + 0x00ff00ff0001ff00, 0x00ff00ff0001ff01, 0x00ff00ff00010000, 0x00ff00ff000101ff, + 0x00ff00ff00010101, 0x00ff00ff01ffff00, 0x00ff00ff01ff0001, 0x00ff00ff01ff0100, + 0x00ff00ff0100ffff, 0x00ff00ff0100ff01, 0x00ff00ff01000000, 0x00ff00ff0101ffff, + 0x00ff00ff0101ff00, 0x00ff00ff01010100, 0x00ff0000ffffff00, 0x00ff0000ffffff01, + 0x00ff0000ffff0000, 0x00ff0000ffff0101, 0x00ff0000ff00ff00, 0x00ff0000ff0000ff, + 0x00ff0000ff000000, 0x00ff0000ff000001, 0x00ff0000ff000100, 0x00ff0000ff01ffff, + 0x00ff0000ff010000, 0x00ff0000ff010101, 0x00ff000000ffff00, 0x00ff000000ff00ff, + 0x00ff000000ff0000, 0x00ff000000ff0001, 0x00ff000000ff0100, 0x00ff00000000ffff, + 0x00ff00000000ff00, 0x00ff0000000000ff, 0x00ff000000000000, 0x00ff000000000001, + 0x00ff0000000001ff, 0x00ff000000000100, 0x00ff00000001ff00, 0x00ff0000000100ff, + 0x00ff000000010000, 0x00ff000000010001, 0x00ff000000010100, 0x00ff000001ffff01, + 0x00ff000001ff00ff, 0x00ff000001ff0000, 0x00ff000001ff01ff, 0x00ff00000100ff00, + 0x00ff0000010000ff, 0x00ff000001000000, 0x00ff000001000001, 0x00ff000001000100, + 0x00ff000001000101, 0x00ff000001010000, 0x00ff0000010101ff, 0x00ff000001010101, + 0x00ff0001ffffff00, 0x00ff0001ffff0000, 0x00ff0001ffff0100, 0x00ff0001ff0000ff, + 0x00ff0001ff000000, 0x00ff0001ff0001ff, 0x00ff0001ff000101, 0x00ff0001ff01ff00, + 0x00ff0001ff0100ff, 0x00ff0001ff010100, 0x00ff000100ffffff, 0x00ff000100ffff01, + 0x00ff000100ff0000, 0x00ff000100ff01ff, 0x00ff00010000ffff, 0x00ff00010000ff00, + 0x00ff00010000ff01, 0x00ff000100000000, 0x00ff000100000001, 0x00ff000100000100, + 0x00ff00010001ff01, 0x00ff000100010000, 0x00ff0001000101ff, 0x00ff000101ffff00, + 0x00ff000101ff0000, 0x00ff000101ff0101, 0x00ff0001010000ff, 0x00ff000101000000, + 0x00ff00010101ff00, 0x00ff0001010100ff, 0x00ff000101010001, 0x00ff01ffffff0000, + 0x00ff01ffff00ff00, 0x00ff01ffff000000, 0x00ff01ffff000101, 0x00ff01ffff010000, + 0x00ff01ff00ffff01, 0x00ff01ff00ff0100, 0x00ff01ff0000ffff, 0x00ff01ff00000000, + 0x00ff01ff000001ff, 0x00ff01ff0001ff00, 0x00ff01ff000100ff, 0x00ff01ff00010001, + 0x00ff01ff00010100, 0x00ff01ff01ff0000, 0x00ff01ff0100ff00, 0x00ff01ff010000ff, + 0x00ff01ff01000001, 0x00ff01ff01000100, 0x00ff01ff01010000, 0x00ff0100ffffff00, + 0x00ff0100ffff0000, 0x00ff0100ffff0001, 0x00ff0100ffff0101, 0x00ff0100ff00ffff, + 0x00ff0100ff0000ff, 0x00ff0100ff000000, 0x00ff0100ff0001ff, 0x00ff0100ff01ff00, + 0x00ff0100ff0100ff, 0x00ff0100ff010001, 0x00ff010000ffffff, 0x00ff010000ff0000, + 0x00ff010000ff0101, 0x00ff01000000ff00, 0x00ff01000000ff01, 0x00ff0100000000ff, + 0x00ff010000000000, 0x00ff010000000001, 0x00ff010000000100, 0x00ff01000001ffff, + 0x00ff01000001ff01, 0x00ff010000010000, 0x00ff010000010001, 0x00ff010000010101, + 0x00ff010001ff0001, 0x00ff010001ff0100, 0x00ff01000100ff01, 0x00ff010001000000, + 0x00ff010001000001, 0x00ff0100010001ff, 0x00ff01000101ff00, 0x00ff0100010100ff, + 0x00ff010001010001, 0x00ff010001010100, 0x00ff0101ff000001, 0x00ff010100ff00ff, + 0x00ff010100ff0001, 0x00ff010100ff0100, 0x00ff010100000000, 0x00ff0101000001ff, + 0x00ff010100000101, 0x00ff0101000100ff, 0x00ff010100010100, 0x00ff0101010000ff, + 0x00ff010101010000, 0x0000ffffffffff00, 0x0000ffffffff00ff, 0x0000ffffffff0000, + 0x0000ffffffff0001, 0x0000ffffffff0100, 0x0000ffffff00ff01, 0x0000ffffff000000, + 0x0000ffffff000101, 0x0000ffffff01ff00, 0x0000ffffff0100ff, 0x0000ffffff010100, + 0x0000ffff00ffffff, 0x0000ffff00ff0000, 0x0000ffff00ff01ff, 0x0000ffff0000ff00, + 0x0000ffff000000ff, 0x0000ffff00000000, 0x0000ffff00000001, 0x0000ffff00000100, + 0x0000ffff00010000, 0x0000ffff000101ff, 0x0000ffff01ff0001, 0x0000ffff01ff0100, + 0x0000ffff01000000, 0x0000ffff010001ff, 0x0000ffff0101ffff, 0x0000ffff0101ff00, + 0x0000ffff01010001, 0x0000ffff01010100, 0x0000ff00ffff0000, 0x0000ff00ffff01ff, + 0x0000ff00ffff0100, 0x0000ff00ffff0101, 0x0000ff00ff00ff00, 0x0000ff00ff0000ff, + 0x0000ff00ff000000, 0x0000ff00ff000001, 0x0000ff00ff0001ff, 0x0000ff00ff000100, + 0x0000ff00ff01ffff, 0x0000ff00ff010000, 0x0000ff00ff010001, 0x0000ff00ff0101ff, + 0x0000ff00ff010101, 0x0000ff0000ffff00, 0x0000ff0000ff00ff, 0x0000ff0000ff0000, + 0x0000ff0000ff0001, 0x0000ff0000ff0100, 0x0000ff000000ffff, 0x0000ff000000ff00, + 0x0000ff000000ff01, 0x0000ff00000000ff, 0x0000ff0000000000, 0x0000ff0000000001, + 0x0000ff00000001ff, 0x0000ff0000000100, 0x0000ff0000000101, 0x0000ff000001ff00, + 0x0000ff00000100ff, 0x0000ff0000010000, 0x0000ff0000010001, 0x0000ff0000010100, + 0x0000ff0001ffff01, 0x0000ff0001ff0000, 0x0000ff000100ff00, 0x0000ff00010000ff, + 0x0000ff0001000000, 0x0000ff0001000001, 0x0000ff0001000100, 0x0000ff000101ffff, + 0x0000ff0001010000, 0x0000ff0001010101, 0x0000ff01ffffff00, 0x0000ff01ffff0001, + 0x0000ff01ff00ff01, 0x0000ff01ff000000, 0x0000ff01ff000101, 0x0000ff01ff01ff00, + 0x0000ff01ff0100ff, 0x0000ff0100ffff01, 0x0000ff0100ff0000, 0x0000ff0100ff0101, + 0x0000ff010000ff00, 0x0000ff01000000ff, 0x0000ff0100000000, 0x0000ff0100000001, + 0x0000ff0100000100, 0x0000ff010001ff01, 0x0000ff0100010000, 0x0000ff0101ff0000, + 0x0000ff010100ffff, 0x0000ff010100ff01, 0x0000ff0101000000, 0x0000ff0101000100, + 0x0000ff0101000101, 0x0000ff01010100ff, 0x000000ffffff00ff, 0x000000ffffff0000, + 0x000000ffff00ff00, 0x000000ffff0000ff, 0x000000ffff000000, 0x000000ffff000001, + 0x000000ffff0001ff, 0x000000ffff000100, 0x000000ffff01ff00, 0x000000ffff010000, + 0x000000ffff0101ff, 0x000000ffff010101, 0x000000ff00ffff00, 0x000000ff00ff00ff, + 0x000000ff00ff0000, 0x000000ff00ff0001, 0x000000ff00ff0100, 0x000000ff00ff0101, + 0x000000ff0000ffff, 0x000000ff0000ff00, 0x000000ff000000ff, 0x000000ff00000000, + 0x000000ff00000001, 0x000000ff000001ff, 0x000000ff00000100, 0x000000ff00000101, + 0x000000ff0001ff00, 0x000000ff0001ff01, 0x000000ff000100ff, 0x000000ff00010000, + 0x000000ff00010001, 0x000000ff00010100, 0x000000ff01ffffff, 0x000000ff01ff01ff, + 0x000000ff01ff0101, 0x000000ff0100ff00, 0x000000ff010000ff, 0x000000ff01000000, + 0x000000ff01000001, 0x000000ff01000100, 0x000000ff0101ff00, 0x000000ff010100ff, + 0x000000ff01010000, 0x000000ff01010101, 0x00000000ffffff00, 0x00000000ffffff01, + 0x00000000ffff00ff, 0x00000000ffff0000, 0x00000000ffff0001, 0x00000000ffff0100, + 0x00000000ff00ffff, 0x00000000ff00ff00, 0x00000000ff00ff01, 0x00000000ff0000ff, + 0x00000000ff000000, 0x00000000ff000001, 0x00000000ff000100, 0x00000000ff000101, + 0x00000000ff01ff00, 0x00000000ff0100ff, 0x00000000ff010000, 0x00000000ff010001, + 0x00000000ff010100, 0x0000000000ffffff, 0x0000000000ffff00, 0x0000000000ffff01, + 0x0000000000ff00ff, 0x0000000000ff0000, 0x0000000000ff0001, 0x0000000000ff01ff, + 0x0000000000ff0100, 0x000000000000ffff, 0x000000000000ff00, 0x000000000000ff01, + 0x00000000000000ff, 0x0000000000000000, 0x0000000000000001, 0x00000000000001ff, + 0x0000000000000100, 0x0000000000000101, 0x000000000001ffff, 0x000000000001ff00, + 0x00000000000100ff, 0x0000000000010000, 0x0000000000010001, 0x00000000000101ff, + 0x0000000000010100, 0x0000000000010101, 0x0000000001ffff00, 0x0000000001ff00ff, + 0x0000000001ff0000, 0x0000000001ff0100, 0x0000000001ff0101, 0x000000000100ffff, + 0x000000000100ff00, 0x00000000010000ff, 0x0000000001000000, 0x0000000001000001, + 0x00000000010001ff, 0x0000000001000100, 0x000000000101ff00, 0x00000000010100ff, + 0x0000000001010000, 0x0000000001010001, 0x0000000001010100, 0x00000001ffffffff, + 0x00000001ffffff00, 0x00000001ffffff01, 0x00000001ffff00ff, 0x00000001ffff0001, + 0x00000001ffff01ff, 0x00000001ffff0100, 0x00000001ff00ff00, 0x00000001ff0000ff, + 0x00000001ff000000, 0x00000001ff0001ff, 0x00000001ff000100, 0x00000001ff01ffff, + 0x00000001ff01ff00, 0x00000001ff01ff01, 0x00000001ff0100ff, 0x00000001ff010000, + 0x00000001ff010001, 0x00000001ff0101ff, 0x00000001ff010100, 0x0000000100ffff00, + 0x0000000100ff0000, 0x0000000100ff0001, 0x0000000100ff01ff, 0x0000000100ff0100, + 0x0000000100ff0101, 0x000000010000ffff, 0x000000010000ff00, 0x000000010000ff01, + 0x00000001000000ff, 0x0000000100000000, 0x0000000100000001, 0x00000001000001ff, + 0x0000000100000100, 0x0000000100000101, 0x000000010001ff00, 0x00000001000100ff, + 0x0000000100010000, 0x0000000100010100, 0x0000000101ffff01, 0x0000000101ff0000, + 0x0000000101ff0001, 0x0000000101ff01ff, 0x0000000101ff0100, 0x0000000101ff0101, + 0x000000010100ff00, 0x0000000101000000, 0x0000000101000101, 0x000000010101ff01, + 0x0000000101010000, 0x0000000101010001, 0x00000001010101ff, 0x0000000101010100, + 0x000001ffffff00ff, 0x000001ffffff0000, 0x000001ffffff0001, 0x000001ffffff0100, + 0x000001ffff00ffff, 0x000001ffff000000, 0x000001ffff0001ff, 0x000001ffff01ff00, + 0x000001ffff010101, 0x000001ff00ff0000, 0x000001ff00ff01ff, 0x000001ff00ff0101, + 0x000001ff0000ff00, 0x000001ff000000ff, 0x000001ff00000000, 0x000001ff00000001, + 0x000001ff000001ff, 0x000001ff00000100, 0x000001ff0001ffff, 0x000001ff0001ff01, + 0x000001ff000100ff, 0x000001ff00010000, 0x000001ff01ffff01, 0x000001ff01ff0100, + 0x000001ff0100ffff, 0x000001ff0100ff01, 0x000001ff01000000, 0x000001ff010001ff, + 0x000001ff0101ff00, 0x000001ff01010100, 0x00000100ffffff00, 0x00000100ffffff01, + 0x00000100ffff0000, 0x00000100ffff0101, 0x00000100ff00ff00, 0x00000100ff0000ff, + 0x00000100ff000000, 0x00000100ff000001, 0x00000100ff000100, 0x00000100ff010000, + 0x0000010000ffff00, 0x0000010000ff00ff, 0x0000010000ff0000, 0x0000010000ff0001, + 0x0000010000ff0100, 0x000001000000ffff, 0x000001000000ff00, 0x000001000000ff01, + 0x00000100000000ff, 0x0000010000000000, 0x0000010000000001, 0x00000100000001ff, + 0x0000010000000100, 0x0000010000000101, 0x000001000001ff00, 0x00000100000100ff, + 0x0000010000010000, 0x0000010000010001, 0x0000010000010100, 0x0000010001ffff00, + 0x0000010001ff0000, 0x0000010001ff0100, 0x000001000100ff00, 0x00000100010000ff, + 0x0000010001000000, 0x0000010001000001, 0x00000100010001ff, 0x0000010001000100, + 0x0000010001010000, 0x00000101ffff00ff, 0x00000101ffff01ff, 0x00000101ff000000, + 0x00000101ff000101, 0x00000101ff01ffff, 0x00000101ff010000, 0x00000101ff010001, + 0x00000101ff010100, 0x0000010100ff0000, 0x0000010100ff01ff, 0x0000010100ff0100, + 0x000001010000ff00, 0x0000010100000000, 0x0000010100000001, 0x00000101000001ff, + 0x0000010100000100, 0x000001010001ff01, 0x0000010100010000, 0x00000101000101ff, + 0x0000010100010101, 0x0000010101ffff00, 0x0000010101ff0101, 0x000001010100ff01, + 0x0000010101000000, 0x0000010101000001, 0x00000101010001ff, 0x0000010101000101, + 0x000001010101ff00, 0x0001ffffffff0000, 0x0001ffffff0000ff, 0x0001ffffff000001, + 0x0001ffffff000100, 0x0001ffffff010000, 0x0001ffff00ff00ff, 0x0001ffff0000ffff, + 0x0001ffff00000000, 0x0001ffff00000001, 0x0001ffff000001ff, 0x0001ffff00000101, + 0x0001ffff0001ff00, 0x0001ffff000100ff, 0x0001ffff00010001, 0x0001ffff00010100, + 0x0001ffff01ffff00, 0x0001ffff01000001, 0x0001ffff01010000, 0x0001ff00ffffff00, + 0x0001ff00ffff00ff, 0x0001ff00ffff0001, 0x0001ff00ffff0100, 0x0001ff00ff00ff01, + 0x0001ff00ff000000, 0x0001ff00ff01ff00, 0x0001ff00ff01ff01, 0x0001ff00ff010001, + 0x0001ff00ff010100, 0x0001ff0000ff0000, 0x0001ff0000ff0100, 0x0001ff000000ff00, + 0x0001ff0000000000, 0x0001ff0000000001, 0x0001ff0000000100, 0x0001ff0000010000, + 0x0001ff0000010001, 0x0001ff0000010101, 0x0001ff0001ff00ff, 0x0001ff0001ff0101, + 0x0001ff000100ff01, 0x0001ff0001000000, 0x0001ff000101ff00, 0x0001ff0001010001, + 0x0001ff0001010100, 0x0001ff01ff00ff00, 0x0001ff01ff000001, 0x0001ff01ff000100, + 0x0001ff0100ffffff, 0x0001ff0100ffff00, 0x0001ff0100ff0001, 0x0001ff0100000000, + 0x0001ff0100000001, 0x0001ff01000001ff, 0x0001ff010001ffff, 0x0001ff0101ff0000, + 0x0001ff010100ff00, 0x0001ff0101000001, 0x0001ff0101010000, 0x000100ffff00ff00, + 0x000100ffff00ff01, 0x000100ffff000000, 0x000100ffff000001, 0x000100ffff000101, + 0x000100ffff01ff00, 0x000100ffff010001, 0x000100ffff010100, 0x000100ff00ffffff, + 0x000100ff00ffff01, 0x000100ff00ff0000, 0x000100ff00ff01ff, 0x000100ff00ff0101, + 0x000100ff0000ff00, 0x000100ff000000ff, 0x000100ff00000000, 0x000100ff00000001, + 0x000100ff00000100, 0x000100ff00000101, 0x000100ff0001ffff, 0x000100ff0001ff01, + 0x000100ff00010000, 0x000100ff01ff00ff, 0x000100ff01ff0000, 0x000100ff01ff0100, + 0x000100ff0100ffff, 0x000100ff0100ff01, 0x000100ff010000ff, 0x000100ff01000000, + 0x000100ff01000001, 0x000100ff010001ff, 0x000100ff01000101, 0x000100ff0101ff00, + 0x000100ff010100ff, 0x000100ff01010100, 0x00010000ffff0000, 0x00010000ffff01ff, + 0x00010000ffff0101, 0x00010000ff00ff00, 0x00010000ff000000, 0x00010000ff000001, + 0x00010000ff000100, 0x0001000000ff00ff, 0x0001000000ff0000, 0x0001000000ff0001, + 0x0001000000ff0100, 0x000100000000ffff, 0x000100000000ff00, 0x00010000000000ff, + 0x0001000000000000, 0x0001000000000001, 0x0001000000000100, 0x000100000001ff00, + 0x00010000000100ff, 0x0001000000010000, 0x0001000000010001, 0x0001000000010100, + 0x0001000001ff0001, 0x0001000001ff0100, 0x0001000001ff0101, 0x000100000100ff00, + 0x0001000001000000, 0x0001000001000001, 0x0001000001000100, 0x0001000001000101, + 0x000100000101ff01, 0x0001000001010000, 0x0001000001010001, 0x00010000010101ff, + 0x00010001ffffff01, 0x00010001ffff0100, 0x00010001ff000000, 0x00010001ff01ffff, + 0x00010001ff010001, 0x00010001ff0101ff, 0x00010001ff010100, 0x0001000100ffffff, + 0x0001000100ff0000, 0x0001000100ff01ff, 0x0001000100ff0101, 0x000100010000ff00, + 0x00010001000000ff, 0x0001000100000000, 0x0001000100000001, 0x00010001000001ff, + 0x0001000100000101, 0x000100010001ffff, 0x0001000100010000, 0x00010001000101ff, + 0x0001000101ffffff, 0x0001000101ffff01, 0x0001000101ff0000, 0x0001000101ff0101, + 0x00010001010000ff, 0x0001000101000001, 0x00010001010001ff, 0x0001000101000100, + 0x000100010101ffff, 0x00010001010100ff, 0x0001000101010001, 0x0001000101010101, + 0x000101ffff000001, 0x000101ffff000100, 0x000101ffff010000, 0x000101ff00ffff00, + 0x000101ff0000ff01, 0x000101ff00000000, 0x000101ff00000101, 0x000101ff0001ff00, + 0x000101ff00010100, 0x000101ff01ff0000, 0x000101ff0100ff00, 0x000101ff010001ff, + 0x000101ff01010001, 0x00010100ffffff00, 0x00010100ffff00ff, 0x00010100ff00ffff, + 0x00010100ff000000, 0x00010100ff01ff00, 0x00010100ff0100ff, 0x00010100ff010001, + 0x00010100ff010100, 0x0001010000ffffff, 0x0001010000ffff00, 0x0001010000ff0000, + 0x0001010000ff0001, 0x0001010000ff01ff, 0x000101000000ff00, 0x00010100000000ff, + 0x0001010000000000, 0x0001010000000001, 0x0001010000000100, 0x000101000001ffff, + 0x0001010000010000, 0x0001010000010101, 0x0001010001ffff01, 0x0001010001ff00ff, + 0x0001010001ff0101, 0x0001010001000000, 0x000101000101ff00, 0x00010100010100ff, + 0x0001010001010000, 0x0001010001010100, 0x00010101ff00ff00, 0x00010101ff000001, + 0x00010101ff0001ff, 0x0001010100ffff00, 0x0001010100ff00ff, 0x0001010100ff0100, + 0x000101010000ffff, 0x0001010100000000, 0x00010101000001ff, 0x0001010100000101, + 0x00010101000100ff, 0x0001010100010000, 0x0001010100010100, 0x0001010101ff0001, + 0x00010101010000ff, 0x00010101010001ff, 0x0001010101000101, 0x0001010101010001, + 0x01ffffffffffffff, 0x01ffffffffffff01, 0x01ffffffffff01ff, 0x01ffffffffff0101, + 0x01ffffffff01ffff, 0x01ffffffff01ff01, 0x01ffffffff0101ff, 0x01ffffffff010101, + 0x01ffffff00ff0000, 0x01ffffff0000ffff, 0x01ffffff0000ff00, 0x01ffffff000000ff, + 0x01ffffff00000001, 0x01ffffff00000100, 0x01ffffff00010000, 0x01ffffff01ffffff, + 0x01ffffff01ffff01, 0x01ffffff01ff01ff, 0x01ffffff01ff0101, 0x01ffffff01000000, + 0x01ffffff0101ffff, 0x01ffffff0101ff01, 0x01ffffff010101ff, 0x01ffffff01010101, + 0x01ffff00ffff0000, 0x01ffff00ff00ff00, 0x01ffff00ff0000ff, 0x01ffff00ff000001, + 0x01ffff00ff000100, 0x01ffff00ff010000, 0x01ffff0000ffff00, 0x01ffff0000ff00ff, + 0x01ffff0000ff0100, 0x01ffff000000ffff, 0x01ffff000000ff01, 0x01ffff0000000000, + 0x01ffff0000000001, 0x01ffff00000001ff, 0x01ffff0000000100, 0x01ffff00000100ff, + 0x01ffff0000010001, 0x01ffff0000010100, 0x01ffff0001ff0000, 0x01ffff0001ff0100, + 0x01ffff00010000ff, 0x01ffff0001000001, 0x01ffff0001000100, 0x01ffff0001010000, + 0x01ffff01ffffffff, 0x01ffff01ffffff01, 0x01ffff01ffff01ff, 0x01ffff01ffff0101, + 0x01ffff01ff000000, 0x01ffff01ff01ffff, 0x01ffff01ff01ff01, 0x01ffff01ff0101ff, + 0x01ffff01ff010101, 0x01ffff010000ff00, 0x01ffff01000000ff, 0x01ffff0100000100, + 0x01ffff0100010000, 0x01ffff0101ffffff, 0x01ffff0101ffff01, 0x01ffff0101ff01ff, + 0x01ffff0101ff0101, 0x01ffff0101000000, 0x01ffff010101ffff, 0x01ffff010101ff01, + 0x01ffff01010101ff, 0x01ffff0101010101, 0x01ff00ffff0000ff, 0x01ff00ffff000100, + 0x01ff00ff00ffff00, 0x01ff00ff00ff00ff, 0x01ff00ff0000ff00, 0x01ff00ff00000000, + 0x01ff00ff00000101, 0x01ff00ff0001ff00, 0x01ff00ff000100ff, 0x01ff00ff00010100, + 0x01ff00ff010000ff, 0x01ff00ff01000100, 0x01ff0000ffffff00, 0x01ff0000ffff0100, + 0x01ff0000ff00ff01, 0x01ff0000ff000000, 0x01ff0000ff000101, 0x01ff0000ff010001, + 0x01ff0000ff010100, 0x01ff000000ffffff, 0x01ff000000ffff00, 0x01ff000000ff0000, + 0x01ff000000ff01ff, 0x01ff00000000ff00, 0x01ff0000000000ff, 0x01ff000000000000, + 0x01ff000000000001, 0x01ff000000000100, 0x01ff000000000101, 0x01ff000000010000, + 0x01ff000000010001, 0x01ff0000000101ff, 0x01ff000000010101, 0x01ff000001ffff00, + 0x01ff000001ff00ff, 0x01ff000001ff0001, 0x01ff000001ff0100, 0x01ff00000100ffff, + 0x01ff00000100ff01, 0x01ff000001000000, 0x01ff0000010001ff, 0x01ff000001010001, + 0x01ff0001ff00ff00, 0x01ff0001ff000001, 0x01ff0001ff000100, 0x01ff0001ff010000, + 0x01ff000100ffff00, 0x01ff000100ff00ff, 0x01ff000100ff0100, 0x01ff000100ff0101, + 0x01ff00010000ffff, 0x01ff000100000000, 0x01ff000100000100, 0x01ff000100000101, + 0x01ff00010001ff00, 0x01ff000100010001, 0x01ff000100010101, 0x01ff000101ff0000, + 0x01ff00010100ff00, 0x01ff000101000101, 0x01ff0001010100ff, 0x01ff01ffffffffff, + 0x01ff01ffffffff01, 0x01ff01ffffff01ff, 0x01ff01ffffff0101, 0x01ff01ffff000000, + 0x01ff01ffff01ffff, 0x01ff01ffff01ff01, 0x01ff01ffff0101ff, 0x01ff01ffff010101, + 0x01ff01ff00ffff00, 0x01ff01ff00ff0000, 0x01ff01ff0000ff00, 0x01ff01ff000000ff, + 0x01ff01ff00000100, 0x01ff01ff00010000, 0x01ff01ff00010100, 0x01ff01ff01ffffff, + 0x01ff01ff01ffff01, 0x01ff01ff01ff01ff, 0x01ff01ff01ff0101, 0x01ff01ff01000000, + 0x01ff01ff0101ffff, 0x01ff01ff0101ff01, 0x01ff01ff010101ff, 0x01ff01ff01010101, + 0x01ff0100ffff0000, 0x01ff0100ffff0001, 0x01ff0100ff00ff00, 0x01ff0100ff0000ff, + 0x01ff0100ff000001, 0x01ff0100ff010000, 0x01ff010000ffff00, 0x01ff010000ff00ff, + 0x01ff010000ff0001, 0x01ff010000ff0100, 0x01ff01000000ffff, 0x01ff01000000ff01, + 0x01ff010000000000, 0x01ff010000000101, 0x01ff01000001ff00, 0x01ff0100000100ff, + 0x01ff010001ff0000, 0x01ff010001000001, 0x01ff010001000100, 0x01ff010001010000, + 0x01ff0101ffffffff, 0x01ff0101ffffff01, 0x01ff0101ffff01ff, 0x01ff0101ffff0101, + 0x01ff0101ff000000, 0x01ff0101ff01ffff, 0x01ff0101ff01ff01, 0x01ff0101ff0101ff, + 0x01ff0101ff010101, 0x01ff010100ff0000, 0x01ff01010000ff00, 0x01ff0101000000ff, + 0x01ff010100000001, 0x01ff010101ffffff, 0x01ff010101ffff01, 0x01ff010101ff01ff, + 0x01ff010101ff0101, 0x01ff010101000000, 0x01ff01010101ffff, 0x01ff01010101ff01, + 0x01ff0101010101ff, 0x01ff010101010101, 0x0100ffffffff0000, 0x0100ffffff00ff00, + 0x0100ffffff000001, 0x0100ffffff0001ff, 0x0100ffffff000100, 0x0100ffffff010000, + 0x0100ffff00ffff00, 0x0100ffff00ff0001, 0x0100ffff00ff0100, 0x0100ffff00000000, + 0x0100ffff000001ff, 0x0100ffff00000101, 0x0100ffff00010100, 0x0100ffff00010101, + 0x0100ffff01ff0000, 0x0100ffff0100ff00, 0x0100ffff010000ff, 0x0100ffff01000001, + 0x0100ffff01000100, 0x0100ffff01010000, 0x0100ff00ffffff00, 0x0100ff00ffff00ff, + 0x0100ff00ffff0001, 0x0100ff00ffff0100, 0x0100ff00ff00ffff, 0x0100ff00ff000000, + 0x0100ff00ff0001ff, 0x0100ff00ff000101, 0x0100ff00ff01ff00, 0x0100ff00ff0100ff, + 0x0100ff00ff010001, 0x0100ff00ff010100, 0x0100ff0000ffffff, 0x0100ff0000ff0000, + 0x0100ff000000ffff, 0x0100ff000000ff00, 0x0100ff00000000ff, 0x0100ff0000000000, + 0x0100ff0000000001, 0x0100ff0000000100, 0x0100ff000001ff01, 0x0100ff0000010000, + 0x0100ff0001ff00ff, 0x0100ff0001ff0001, 0x0100ff000100ff01, 0x0100ff0001000000, + 0x0100ff00010001ff, 0x0100ff000101ff00, 0x0100ff00010100ff, 0x0100ff0001010001, + 0x0100ff0001010100, 0x0100ff01ffff0000, 0x0100ff01ff00ff00, 0x0100ff01ff0000ff, + 0x0100ff01ff000100, 0x0100ff01ff010000, 0x0100ff0100ff00ff, 0x0100ff0100ff0001, + 0x0100ff0100ff0100, 0x0100ff010000ffff, 0x0100ff010000ff01, 0x0100ff0100000000, + 0x0100ff01000001ff, 0x0100ff0100010001, 0x0100ff0100010100, 0x0100ff0101ff0000, + 0x0100ff01010000ff, 0x0100ff0101000001, 0x0100ff0101010100, 0x010000ffffffff00, + 0x010000ffffff00ff, 0x010000ffffff0001, 0x010000ffff00ffff, 0x010000ffff000000, + 0x010000ffff0001ff, 0x010000ffff010001, 0x010000ff00ffffff, 0x010000ff00ff0101, + 0x010000ff0000ff00, 0x010000ff000000ff, 0x010000ff00000000, 0x010000ff00000001, + 0x010000ff000001ff, 0x010000ff00000100, 0x010000ff0001ffff, 0x010000ff0001ff00, + 0x010000ff0001ff01, 0x010000ff00010000, 0x010000ff01ff00ff, 0x010000ff01ff0001, + 0x010000ff0100ff01, 0x010000ff010000ff, 0x010000ff01000000, 0x010000ff010001ff, + 0x010000ff0101ff00, 0x010000ff01010100, 0x01000000ffffffff, 0x01000000ffff0000, + 0x01000000ffff01ff, 0x01000000ffff0101, 0x01000000ff00ffff, 0x01000000ff00ff00, + 0x01000000ff0000ff, 0x01000000ff000000, 0x01000000ff000001, 0x01000000ff000100, + 0x01000000ff01ff00, 0x01000000ff010000, 0x01000000ff010100, 0x01000000ff010101, + 0x0100000000ffff00, 0x0100000000ff00ff, 0x0100000000ff0000, 0x0100000000ff0001, + 0x0100000000ff0100, 0x010000000000ffff, 0x010000000000ff00, 0x010000000000ff01, + 0x01000000000000ff, 0x0100000000000000, 0x0100000000000001, 0x01000000000001ff, + 0x0100000000000100, 0x0100000000000101, 0x010000000001ff00, 0x01000000000100ff, + 0x0100000000010000, 0x0100000000010001, 0x0100000000010100, 0x0100000001ffff00, + 0x0100000001ff0000, 0x0100000001ff01ff, 0x010000000100ff00, 0x010000000100ff01, + 0x01000000010000ff, 0x0100000001000000, 0x0100000001000001, 0x0100000001000100, + 0x0100000001000101, 0x010000000101ffff, 0x010000000101ff01, 0x0100000001010000, + 0x01000000010101ff, 0x0100000001010101, 0x01000001ffffff00, 0x01000001ffff00ff, + 0x01000001ff00ffff, 0x01000001ff000000, 0x01000001ff000100, 0x01000001ff01ffff, + 0x01000001ff010001, 0x01000001ff010100, 0x0100000100ff0000, 0x0100000100ff01ff, + 0x0100000100ff0100, 0x010000010000ff00, 0x010000010000ff01, 0x0100000100000000, + 0x0100000100000001, 0x0100000100000100, 0x0100000100010000, 0x01000001000101ff, + 0x0100000101ffff01, 0x0100000101ff00ff, 0x0100000101ff0100, 0x0100000101ff0101, + 0x010000010100ff01, 0x01000001010000ff, 0x0100000101000000, 0x01000001010100ff, + 0x0100000101010001, 0x0100000101010100, 0x010001ffffff0000, 0x010001ffff000001, + 0x010001ffff000100, 0x010001ffff010000, 0x010001ff00ffff00, 0x010001ff00ff0001, + 0x010001ff0000ffff, 0x010001ff0000ff01, 0x010001ff00000000, 0x010001ff00000001, + 0x010001ff00000101, 0x010001ff000100ff, 0x010001ff00010000, 0x010001ff01ff0000, + 0x010001ff0100ff00, 0x010001ff01000001, 0x010001ff01000100, 0x010001ff01010000, + 0x01000100ffff00ff, 0x01000100ffff0001, 0x01000100ffff0100, 0x01000100ff00ffff, + 0x01000100ff00ff01, 0x01000100ff000000, 0x01000100ff0001ff, 0x01000100ff000101, + 0x01000100ff01ffff, 0x01000100ff01ff00, 0x01000100ff0100ff, 0x01000100ff010001, + 0x0100010000ffffff, 0x0100010000ffff01, 0x0100010000ff0000, 0x0100010000ff01ff, + 0x0100010000ff0101, 0x010001000000ff00, 0x01000100000000ff, 0x0100010000000000, + 0x0100010000000001, 0x0100010000000100, 0x010001000001ff01, 0x0100010000010000, + 0x0100010000010001, 0x0100010000010101, 0x0100010001ffff00, 0x0100010001ff00ff, + 0x010001000100ffff, 0x010001000100ff01, 0x0100010001000000, 0x0100010001000101, + 0x010001000101ff00, 0x0100010001010001, 0x01000101ffff0000, 0x01000101ff000000, + 0x01000101ff010000, 0x0100010100ff00ff, 0x0100010100ff0001, 0x0100010100ff0100, + 0x010001010000ffff, 0x0100010100000000, 0x01000101000001ff, 0x010001010001ff00, + 0x0100010101ff0000, 0x010001010100ff00, 0x01000101010000ff, 0x0100010101000000, + 0x0100010101000001, 0x0101ffffffffffff, 0x0101ffffffffff01, 0x0101ffffffff01ff, + 0x0101ffffffff0101, 0x0101ffffff000000, 0x0101ffffff01ffff, 0x0101ffffff01ff01, + 0x0101ffffff0101ff, 0x0101ffffff010101, 0x0101ffff00ff0000, 0x0101ffff0000ff00, + 0x0101ffff000000ff, 0x0101ffff00000001, 0x0101ffff00000100, 0x0101ffff01ffffff, + 0x0101ffff01ffff01, 0x0101ffff01ff01ff, 0x0101ffff01ff0101, 0x0101ffff01000000, + 0x0101ffff0101ffff, 0x0101ffff0101ff01, 0x0101ffff010101ff, 0x0101ffff01010101, + 0x0101ff00ffff0000, 0x0101ff00ffff0100, 0x0101ff00ff00ff00, 0x0101ff00ff0000ff, + 0x0101ff00ff000001, 0x0101ff00ff000100, 0x0101ff00ff000101, 0x0101ff0000ff0001, + 0x0101ff0000ff0100, 0x0101ff000000ff00, 0x0101ff0000000000, 0x0101ff00000001ff, + 0x0101ff0000000101, 0x0101ff000001ff00, 0x0101ff00000100ff, 0x0101ff0001ff0000, + 0x0101ff000100ffff, 0x0101ff000100ff01, 0x0101ff0001000001, 0x0101ff0001000100, + 0x0101ff01ffffff01, 0x0101ff01ffff01ff, 0x0101ff01ffff0101, 0x0101ff01ff00ffff, + 0x0101ff01ff000100, 0x0101ff01ff01ff01, 0x0101ff01ff0101ff, 0x0101ff01ff010101, + 0x0101ff0100ff0000, 0x0101ff010000ff00, 0x0101ff0100000001, 0x0101ff0100000100, + 0x0101ff0100010000, 0x0101ff0101ffffff, 0x0101ff0101ffff01, 0x0101ff0101ff01ff, + 0x0101ff0101ff0101, 0x0101ff0101000000, 0x0101ff010101ffff, 0x0101ff010101ff01, + 0x0101ff01010101ff, 0x0101ff0101010101, 0x010100ffff000100, 0x010100ffff010000, + 0x010100ff00ffff00, 0x010100ff00ff00ff, 0x010100ff0000ffff, 0x010100ff000000ff, + 0x010100ff00000000, 0x010100ff000001ff, 0x010100ff00000101, 0x010100ff0001ff00, + 0x010100ff00010000, 0x010100ff00010001, 0x010100ff000101ff, 0x010100ff00010100, + 0x010100ff01ff0000, 0x01010000ffff0001, 0x01010000ffff0100, 0x01010000ff00ffff, + 0x01010000ff00ff01, 0x01010000ff000000, 0x01010000ff0001ff, 0x01010000ff010001, + 0x01010000ff010100, 0x0101000000ffff01, 0x0101000000ff0000, 0x010100000000ff00, + 0x01010000000000ff, 0x0101000000000000, 0x0101000000000001, 0x0101000000000100, + 0x0101000000010000, 0x0101000000010101, 0x0101000001ffff00, 0x0101000001ff00ff, + 0x0101000001ff0000, 0x0101000001ff0001, 0x0101000001ff0100, 0x010100000100ff01, + 0x0101000001000000, 0x01010000010001ff, 0x01010001ffff0000, 0x01010001ff00ff00, + 0x01010001ff000001, 0x01010001ff000101, 0x01010001ff01ff00, 0x01010001ff010000, + 0x0101000100ff00ff, 0x0101000100ff0001, 0x0101000100ff0101, 0x010100010000ff01, + 0x0101000100000000, 0x0101000100000001, 0x01010001000001ff, 0x010100010001ffff, + 0x010100010001ff01, 0x0101000101ff0001, 0x010100010100ffff, 0x0101000101000000, + 0x0101000101000001, 0x0101000101000100, 0x010100010101ff00, 0x01010001010100ff, + 0x0101000101010001, 0x010101ffffffffff, 0x010101ffffffff01, 0x010101ffffff01ff, + 0x010101ffffff0101, 0x010101ffff01ffff, 0x010101ffff01ff01, 0x010101ffff0101ff, + 0x010101ffff010101, 0x010101ff0000ff00, 0x010101ff000000ff, 0x010101ff00000001, + 0x010101ff00000100, 0x010101ff01ffffff, 0x010101ff01ffff01, 0x010101ff01ff01ff, + 0x010101ff01ff0101, 0x010101ff01000000, 0x010101ff0101ffff, 0x010101ff0101ff01, + 0x010101ff010101ff, 0x010101ff01010101, 0x01010100ffff0000, 0x01010100ff0000ff, + 0x01010100ff000100, 0x01010100ff01ff00, 0x01010100ff010000, 0x0101010000ffff00, + 0x010101000000ffff, 0x0101010000000000, 0x0101010000000101, 0x010101000001ff00, + 0x0101010000010001, 0x0101010000010100, 0x010101000100ffff, 0x0101010001000001, + 0x01010101ffffffff, 0x01010101ffffff01, 0x01010101ffff01ff, 0x01010101ffff0101, + 0x01010101ff01ffff, 0x01010101ff01ff01, 0x01010101ff0101ff, 0x01010101ff010101, + 0x010101010000ff00, 0x01010101000000ff, 0x0101010100000001, 0x0101010101ffffff, + 0x0101010101ffff01, 0x0101010101ff01ff, 0x0101010101ff0101, 0x0101010101000000, + 0x010101010101ffff, 0x010101010101ff01, 0x01010101010101ff, 0x0101010101010101, +GGML_TABLE_END() +#else +GGML_TABLE_BEGIN(uint32_t, iq1s_grid_gpu, NGRID_IQ1S) + 0x00000000, 0x00000002, 0x00000101, 0x00000200, 0x00000202, 0x00010001, 0x00010101, 0x00020000, + 0x00020002, 0x00020200, 0x00020202, 0x01000101, 0x01010001, 0x01010100, 0x01010102, 0x01020101, + 0x02000000, 0x02000002, 0x02000200, 0x02000202, 0x02010101, 0x02020000, 0x02020002, 0x02020200, + 0x02020202, 0x00000110, 0x00000111, 0x00010011, 0x00010110, 0x00010112, 0x00010211, 0x00010212, + 0x00020111, 0x01000011, 0x01000112, 0x01000211, 0x01010012, 0x01010111, 0x01010212, 0x01020011, + 0x01020110, 0x01020112, 0x01020210, 0x02000111, 0x02010011, 0x02010110, 0x02010112, 0x02020111, + 0x00000020, 0x00000022, 0x00000220, 0x00000222, 0x00010121, 0x00020020, 0x00020022, 0x00020220, + 0x00020222, 0x01000121, 0x01010021, 0x01010221, 0x01020120, 0x01020221, 0x02000020, 0x02000022, + 0x02000220, 0x02000222, 0x02010021, 0x02010121, 0x02010221, 0x02020020, 0x02020022, 0x02020220, + 0x02020222, 0x00011001, 0x00011100, 0x00011102, 0x00021101, 0x01001001, 0x01001201, 0x01011101, + 0x01011202, 0x01021100, 0x01021101, 0x02011001, 0x02011201, 0x02021101, 0x00001011, 0x00001110, + 0x00001111, 0x00001112, 0x00011111, 0x00011210, 0x00011212, 0x00021211, 0x01001010, 0x01001111, + 0x01001212, 0x01011010, 0x01011011, 0x01011110, 0x01011111, 0x01011112, 0x01011211, 0x01021010, + 0x01021012, 0x01021111, 0x01021210, 0x01021212, 0x02001011, 0x02011011, 0x02011111, 0x02011210, + 0x02011212, 0x02021011, 0x02021110, 0x02021111, 0x02021112, 0x02021211, 0x00011120, 0x00011221, + 0x01001021, 0x01001120, 0x01011020, 0x01011022, 0x01011121, 0x01011220, 0x01021020, 0x01021021, + 0x01021122, 0x01021221, 0x02001121, 0x02011021, 0x02011120, 0x02011221, 0x00002000, 0x00002002, + 0x00002200, 0x00002202, 0x00012101, 0x00022000, 0x00022002, 0x00022200, 0x00022202, 0x01002101, + 0x01012001, 0x01012102, 0x01022101, 0x02002000, 0x02002002, 0x02002200, 0x02002202, 0x02012101, + 0x02022000, 0x02022002, 0x02022200, 0x02022202, 0x00002111, 0x00012011, 0x00012110, 0x00012211, + 0x00022110, 0x00022111, 0x01002011, 0x01012010, 0x01012011, 0x01012111, 0x01022011, 0x01022110, + 0x01022211, 0x02012011, 0x02012110, 0x02012112, 0x02012211, 0x02022111, 0x00002020, 0x00002022, + 0x00002220, 0x00002222, 0x00012121, 0x00022020, 0x00022022, 0x00022220, 0x00022222, 0x01002121, + 0x01012021, 0x01012221, 0x01022021, 0x01022121, 0x02002020, 0x02002022, 0x02002121, 0x02002220, + 0x02002222, 0x02012121, 0x02022020, 0x02022022, 0x02022220, 0x02022222, 0x00110000, 0x00110001, + 0x00110100, 0x00110201, 0x00120100, 0x00120101, 0x01100001, 0x01100100, 0x01110000, 0x01110101, + 0x01110200, 0x01120001, 0x01120100, 0x01120101, 0x01120201, 0x02110001, 0x02110100, 0x02110102, + 0x02120001, 0x02120101, 0x00100011, 0x00100110, 0x00100112, 0x00100211, 0x00110010, 0x00110012, + 0x00110111, 0x00110210, 0x00120011, 0x00120110, 0x00120211, 0x01100111, 0x01100212, 0x01110010, + 0x01110011, 0x01110012, 0x01110110, 0x01110111, 0x01110112, 0x01110211, 0x01120010, 0x01120111, + 0x02100110, 0x02110012, 0x02110111, 0x02120011, 0x02120110, 0x00110021, 0x00110120, 0x00110122, + 0x00120121, 0x01100020, 0x01100122, 0x01100221, 0x01110022, 0x01110121, 0x01110220, 0x01110222, + 0x01120120, 0x01120122, 0x02100121, 0x02110021, 0x02110120, 0x02110122, 0x02120121, 0x00101001, + 0x00101102, 0x00101201, 0x00111100, 0x00111101, 0x00111200, 0x00111201, 0x00121001, 0x00121102, + 0x01101001, 0x01101101, 0x01101102, 0x01101200, 0x01101202, 0x01111001, 0x01111100, 0x01111101, + 0x01111102, 0x01111201, 0x01121002, 0x01121101, 0x01121200, 0x02101100, 0x02101201, 0x02111000, + 0x02111100, 0x02111101, 0x02111200, 0x02111201, 0x02111202, 0x02121001, 0x02121100, 0x02121101, + 0x02121201, 0x00101012, 0x00101111, 0x00101212, 0x00111011, 0x00111110, 0x00111111, 0x00111112, + 0x00111211, 0x00121010, 0x00121012, 0x00121111, 0x00121210, 0x00121212, 0x01101011, 0x01101110, + 0x01101111, 0x01101112, 0x01111011, 0x01111012, 0x01111110, 0x01111111, 0x01111112, 0x01111211, + 0x01111212, 0x01121011, 0x01121110, 0x01121111, 0x01121112, 0x01121211, 0x02101010, 0x02101012, + 0x02101110, 0x02101111, 0x02101210, 0x02101212, 0x02111010, 0x02111011, 0x02111110, 0x02111111, + 0x02111112, 0x02111211, 0x02111212, 0x02121010, 0x02121012, 0x02121111, 0x00101021, 0x00101120, + 0x00101121, 0x00101122, 0x00111121, 0x00111122, 0x00111220, 0x00111222, 0x00121021, 0x00121122, + 0x01101020, 0x01101022, 0x01101120, 0x01101121, 0x01101220, 0x01101222, 0x01111021, 0x01111121, + 0x01111122, 0x01111220, 0x01111221, 0x01121021, 0x01121120, 0x01121121, 0x01121220, 0x01121221, + 0x01121222, 0x02101122, 0x02101222, 0x02111022, 0x02111121, 0x02121120, 0x02121221, 0x00112001, + 0x00112102, 0x00122101, 0x01102001, 0x01102100, 0x01102102, 0x01102201, 0x01112000, 0x01112101, + 0x01112200, 0x01112202, 0x01122000, 0x01122001, 0x01122100, 0x01122102, 0x01122201, 0x02102101, + 0x02112001, 0x02112100, 0x02122101, 0x00112010, 0x00112012, 0x00112111, 0x00112212, 0x00122011, + 0x00122111, 0x01102012, 0x01102110, 0x01102111, 0x01102210, 0x01112011, 0x01112110, 0x01112111, + 0x01112112, 0x01112211, 0x01112212, 0x01122010, 0x01122111, 0x01122212, 0x02102211, 0x02112011, + 0x02112012, 0x02112111, 0x02112210, 0x02122011, 0x02122112, 0x02122211, 0x00102221, 0x00112122, + 0x00122120, 0x00122122, 0x01102120, 0x01102122, 0x01102221, 0x01112020, 0x01112022, 0x01112121, + 0x01112220, 0x01122021, 0x01122122, 0x01122221, 0x02102121, 0x02112021, 0x02112122, 0x02112222, + 0x00200000, 0x00200002, 0x00200200, 0x00200202, 0x00210101, 0x00220000, 0x00220002, 0x00220101, + 0x00220200, 0x00220202, 0x01200101, 0x01210001, 0x01210201, 0x01220001, 0x01220101, 0x02200000, + 0x02200002, 0x02200200, 0x02200202, 0x02210101, 0x02220000, 0x02220002, 0x02220101, 0x02220200, + 0x02220202, 0x00200111, 0x00210011, 0x00210110, 0x00210211, 0x00220111, 0x01200012, 0x01200110, + 0x01200211, 0x01210111, 0x01210210, 0x01210212, 0x01220011, 0x01220110, 0x01220111, 0x01220112, + 0x02200111, 0x02210010, 0x02210112, 0x02210211, 0x02220111, 0x00200021, 0x00200220, 0x00200222, + 0x00210021, 0x00210121, 0x00220020, 0x00220022, 0x00220220, 0x00220222, 0x01200121, 0x01210021, + 0x01210122, 0x01210221, 0x01220121, 0x02200021, 0x02200220, 0x02200222, 0x02210021, 0x02210121, + 0x02220020, 0x02220022, 0x02220220, 0x02220222, 0x00201101, 0x00211100, 0x00211102, 0x00211201, + 0x00221101, 0x01201100, 0x01201101, 0x01201102, 0x01201201, 0x01211002, 0x01211101, 0x01211200, + 0x01211202, 0x01221102, 0x02201101, 0x02211001, 0x02211100, 0x02211201, 0x02221001, 0x02221101, + 0x00201211, 0x00211111, 0x00221011, 0x00221211, 0x01201010, 0x01201111, 0x01201210, 0x01211011, + 0x01211110, 0x01211111, 0x01211211, 0x01221012, 0x01221111, 0x01221210, 0x02201211, 0x02211010, + 0x02211110, 0x02211111, 0x02211210, 0x02211212, 0x02221011, 0x02221110, 0x02221112, 0x02221211, + 0x00201121, 0x00211020, 0x00211022, 0x00211221, 0x00221121, 0x01201021, 0x01201221, 0x01211121, + 0x01221020, 0x01221021, 0x01221221, 0x02201120, 0x02201122, 0x02211020, 0x02211222, 0x00202000, + 0x00202002, 0x00202200, 0x00202202, 0x00212101, 0x00222000, 0x00222002, 0x00222200, 0x00222202, + 0x01202101, 0x01212001, 0x01212100, 0x01222101, 0x02202000, 0x02202002, 0x02202200, 0x02202202, + 0x02222000, 0x02222002, 0x02222200, 0x02222202, 0x00202211, 0x00212011, 0x00212110, 0x00212211, + 0x00222111, 0x01202112, 0x01202211, 0x01212012, 0x01212111, 0x01222011, 0x01222110, 0x01222112, + 0x01222211, 0x02202111, 0x02212010, 0x02212112, 0x02212211, 0x02222110, 0x02222111, 0x00202020, + 0x00202022, 0x00202220, 0x00202222, 0x00222020, 0x00222022, 0x00222220, 0x00222222, 0x01202121, + 0x01212021, 0x01212122, 0x01212221, 0x01222121, 0x02202020, 0x02202022, 0x02202220, 0x02202222, + 0x02212121, 0x02222020, 0x02222022, 0x02222220, 0x02222222, 0x10000101, 0x10010001, 0x10010102, + 0x10020101, 0x11000201, 0x11010002, 0x11010101, 0x11010200, 0x11010202, 0x11020001, 0x11020100, + 0x11020102, 0x12010100, 0x12010201, 0x12020001, 0x12020102, 0x10000010, 0x10000011, 0x10000110, + 0x10000112, 0x10000211, 0x10010012, 0x10010111, 0x10010112, 0x10010210, 0x10010212, 0x10020011, + 0x10020112, 0x10020211, 0x11000111, 0x11000210, 0x11000212, 0x11010011, 0x11010110, 0x11010111, + 0x11010112, 0x11010211, 0x11010212, 0x11020111, 0x11020210, 0x11020212, 0x12000011, 0x12000110, + 0x12000112, 0x12010010, 0x12010012, 0x12010111, 0x12020010, 0x12020011, 0x12020012, 0x10000121, + 0x10010021, 0x10010120, 0x10010122, 0x10020121, 0x11000021, 0x11010022, 0x11010121, 0x11010222, + 0x11020120, 0x11020221, 0x12000221, 0x12010120, 0x12020121, 0x10001001, 0x10011101, 0x10011201, + 0x10021201, 0x11001101, 0x11001200, 0x11001202, 0x11011001, 0x11011100, 0x11011101, 0x11011102, + 0x11021001, 0x11021002, 0x11021101, 0x11021200, 0x11021202, 0x12001001, 0x12001102, 0x12001201, + 0x12011000, 0x12011002, 0x12011101, 0x12021000, 0x12021001, 0x12021201, 0x10001011, 0x10001012, + 0x10001111, 0x10001212, 0x10011011, 0x10011110, 0x10011111, 0x10011112, 0x10011211, 0x10021010, + 0x10021111, 0x10021212, 0x11001011, 0x11001110, 0x11001111, 0x11001112, 0x11001211, 0x11011010, + 0x11011011, 0x11011110, 0x11011111, 0x11011112, 0x11011210, 0x11011211, 0x11021011, 0x11021110, + 0x11021111, 0x11021112, 0x11021211, 0x12001012, 0x12001110, 0x12001111, 0x12001210, 0x12011011, + 0x12011110, 0x12011111, 0x12011112, 0x12011211, 0x12011212, 0x12021111, 0x12021210, 0x12021212, + 0x10001021, 0x10001121, 0x10001221, 0x10011120, 0x10011121, 0x10011220, 0x10011222, 0x10021021, + 0x10021120, 0x10021221, 0x11001020, 0x11001022, 0x11001121, 0x11001220, 0x11011020, 0x11011021, + 0x11011022, 0x11011121, 0x11011122, 0x11011221, 0x11021022, 0x11021121, 0x11021220, 0x12001021, + 0x12001121, 0x12001222, 0x12011120, 0x12011121, 0x12021021, 0x12021120, 0x12021122, 0x10002101, + 0x10012001, 0x10012101, 0x10012202, 0x10022101, 0x11002002, 0x11002201, 0x11012000, 0x11012101, + 0x11012200, 0x11022001, 0x11022100, 0x11022102, 0x11022201, 0x12002101, 0x12012001, 0x12012100, + 0x12012102, 0x12012201, 0x12022101, 0x10002011, 0x10002111, 0x10002112, 0x10002212, 0x10012010, + 0x10012110, 0x10012111, 0x10012210, 0x10022011, 0x10022110, 0x10022112, 0x11002010, 0x11002111, + 0x11002212, 0x11012011, 0x11012012, 0x11012110, 0x11012111, 0x11012112, 0x11012211, 0x11022010, + 0x11022012, 0x11022111, 0x11022112, 0x11022212, 0x12002112, 0x12002211, 0x12012012, 0x12012111, + 0x12012112, 0x12012210, 0x12022011, 0x12022110, 0x12022112, 0x12022211, 0x10012122, 0x11002120, + 0x11002122, 0x11002221, 0x11012121, 0x11012220, 0x11012222, 0x11022120, 0x11022221, 0x12012120, + 0x12022121, 0x10100001, 0x10100100, 0x10100101, 0x10100102, 0x10100201, 0x10110002, 0x10110101, + 0x10110202, 0x10120001, 0x10120100, 0x10120201, 0x11100000, 0x11100101, 0x11100200, 0x11110001, + 0x11110100, 0x11110101, 0x11110102, 0x11110201, 0x11120101, 0x11120200, 0x12100102, 0x12100201, + 0x12110101, 0x12110200, 0x12120000, 0x12120001, 0x12120102, 0x12120201, 0x10100111, 0x10100210, + 0x10100211, 0x10100212, 0x10110011, 0x10110110, 0x10110111, 0x10110112, 0x10110210, 0x10110211, + 0x10120010, 0x10120111, 0x10120112, 0x10120210, 0x10120212, 0x11100011, 0x11100110, 0x11100111, + 0x11100112, 0x11100211, 0x11110010, 0x11110011, 0x11110012, 0x11110110, 0x11110111, 0x11110112, + 0x11110210, 0x11110211, 0x11110212, 0x11120011, 0x11120110, 0x11120111, 0x11120112, 0x11120211, + 0x12100012, 0x12100111, 0x12110011, 0x12110110, 0x12110111, 0x12110112, 0x12110211, 0x12120010, + 0x12120111, 0x12120212, 0x10100021, 0x10100122, 0x10110022, 0x10110121, 0x10110222, 0x10120021, + 0x10120120, 0x11100022, 0x11100121, 0x11100222, 0x11110021, 0x11110120, 0x11110121, 0x11110122, + 0x11110221, 0x11120022, 0x11120121, 0x12100121, 0x12110020, 0x12110022, 0x12110121, 0x12110221, + 0x12110222, 0x12120120, 0x10101100, 0x10101101, 0x10111001, 0x10111100, 0x10111101, 0x10111102, + 0x10111200, 0x10111201, 0x10121001, 0x10121101, 0x10121200, 0x10121202, 0x11101001, 0x11101100, + 0x11101101, 0x11101102, 0x11101201, 0x11101202, 0x11111000, 0x11111001, 0x11111100, 0x11111101, + 0x11111102, 0x11111200, 0x11111201, 0x11111202, 0x11121001, 0x11121002, 0x11121100, 0x11121101, + 0x11121102, 0x11121201, 0x12101000, 0x12101200, 0x12101202, 0x12111001, 0x12111100, 0x12111101, + 0x12111102, 0x12111201, 0x12121001, 0x12121100, 0x12121101, 0x12121202, 0x10101011, 0x10101012, + 0x10101110, 0x10101111, 0x10101112, 0x10101211, 0x10111010, 0x10111011, 0x10111012, 0x10111110, + 0x10111111, 0x10111112, 0x10111211, 0x10111212, 0x10121011, 0x10121110, 0x10121111, 0x10121112, + 0x10121211, 0x11101010, 0x11101011, 0x11101012, 0x11101110, 0x11101111, 0x11101112, 0x11101210, + 0x11101211, 0x11111010, 0x11111011, 0x11111012, 0x11111110, 0x11111111, 0x11111112, 0x11111210, + 0x11111211, 0x11111212, 0x11121010, 0x11121011, 0x11121110, 0x11121111, 0x11121112, 0x11121210, + 0x11121211, 0x11121212, 0x12101011, 0x12101110, 0x12101111, 0x12101211, 0x12101212, 0x12111010, + 0x12111011, 0x12111110, 0x12111111, 0x12111112, 0x12111210, 0x12111211, 0x12121011, 0x12121110, + 0x12121111, 0x12121112, 0x12121211, 0x10101020, 0x10101021, 0x10101022, 0x10101120, 0x10101122, + 0x10101220, 0x10101221, 0x10111021, 0x10111120, 0x10111121, 0x10111220, 0x10111221, 0x10121020, + 0x10121021, 0x10121022, 0x10121120, 0x10121121, 0x10121122, 0x10121220, 0x10121221, 0x11101021, + 0x11101121, 0x11101122, 0x11101220, 0x11101221, 0x11101222, 0x11111020, 0x11111021, 0x11111022, + 0x11111120, 0x11111121, 0x11111122, 0x11111220, 0x11111221, 0x11111222, 0x11121021, 0x11121120, + 0x11121121, 0x11121221, 0x12101022, 0x12101121, 0x12101122, 0x12101220, 0x12101221, 0x12101222, + 0x12111021, 0x12111121, 0x12111222, 0x12121022, 0x12121121, 0x12121122, 0x12121220, 0x12121221, + 0x10102100, 0x10102101, 0x10102102, 0x10102201, 0x10112000, 0x10112101, 0x10112200, 0x10122001, + 0x10122202, 0x11102101, 0x11102200, 0x11102202, 0x11112001, 0x11112100, 0x11112101, 0x11112102, + 0x11112200, 0x11112201, 0x11122000, 0x11122002, 0x11122100, 0x11122101, 0x12102002, 0x12102201, + 0x12112000, 0x12112002, 0x12112101, 0x12112200, 0x12122001, 0x12122201, 0x10102011, 0x10102012, + 0x10102111, 0x10102212, 0x10112011, 0x10112110, 0x10112111, 0x10112112, 0x10112211, 0x10122111, + 0x11102011, 0x11102110, 0x11102111, 0x11102112, 0x11102211, 0x11112010, 0x11112011, 0x11112012, + 0x11112110, 0x11112111, 0x11112112, 0x11112210, 0x11112211, 0x11112212, 0x11122011, 0x11122110, + 0x11122111, 0x11122112, 0x11122211, 0x12102011, 0x12102111, 0x12102211, 0x12112011, 0x12112110, + 0x12112111, 0x12112112, 0x12112210, 0x12112211, 0x12122111, 0x10102120, 0x10102220, 0x10112121, + 0x10112222, 0x10122020, 0x10122121, 0x10122122, 0x10122221, 0x11102121, 0x11102220, 0x11102221, + 0x11112021, 0x11112121, 0x11112122, 0x11112220, 0x11112221, 0x11122022, 0x11122121, 0x11122220, + 0x11122222, 0x12102021, 0x12102222, 0x12112022, 0x12112121, 0x12112122, 0x12112220, 0x12112222, + 0x12122021, 0x10200101, 0x10210100, 0x10210102, 0x10210201, 0x10220101, 0x11200100, 0x11210000, + 0x11210101, 0x11210102, 0x11210200, 0x11210202, 0x11220001, 0x11220100, 0x11220102, 0x11220201, + 0x12200001, 0x12210102, 0x12220101, 0x10200011, 0x10200110, 0x10200112, 0x10200211, 0x10210012, + 0x10210111, 0x10220011, 0x10220012, 0x10220112, 0x10220211, 0x11200111, 0x11200211, 0x11210011, + 0x11210111, 0x11210112, 0x11210211, 0x11220111, 0x11220112, 0x11220212, 0x12200110, 0x12200212, + 0x12210012, 0x12210111, 0x12220011, 0x12220112, 0x12220211, 0x10210021, 0x10210122, 0x10210221, + 0x11200020, 0x11200021, 0x11200122, 0x11210121, 0x11210122, 0x11210220, 0x11220020, 0x12200121, + 0x12210021, 0x12210122, 0x12220121, 0x10211001, 0x10211002, 0x10211101, 0x10211102, 0x10211202, + 0x10221001, 0x10221102, 0x10221201, 0x11201000, 0x11201002, 0x11201101, 0x11201200, 0x11201202, + 0x11211001, 0x11211100, 0x11211101, 0x11211102, 0x11211201, 0x11211202, 0x11221000, 0x11221002, + 0x11221101, 0x12201100, 0x12201101, 0x12201201, 0x12211000, 0x12211002, 0x12211100, 0x12211101, + 0x12211102, 0x12211200, 0x12211202, 0x12221001, 0x12221100, 0x12221201, 0x10201111, 0x10201210, + 0x10201212, 0x10211011, 0x10211111, 0x10211112, 0x10211211, 0x11201110, 0x11201111, 0x11201112, + 0x11201211, 0x11211010, 0x11211011, 0x11211110, 0x11211111, 0x11211112, 0x11211211, 0x11221011, + 0x11221110, 0x11221111, 0x11221112, 0x11221211, 0x12201112, 0x12201211, 0x12201212, 0x12211011, + 0x12211111, 0x12211112, 0x12211211, 0x12211212, 0x12221012, 0x12221111, 0x12221112, 0x12221210, + 0x10201022, 0x10201221, 0x10211121, 0x10221020, 0x10221122, 0x10221220, 0x10221221, 0x11201020, + 0x11201121, 0x11201220, 0x11201222, 0x11211021, 0x11211120, 0x11211121, 0x11211122, 0x11211220, + 0x11211222, 0x11221020, 0x11221121, 0x11221220, 0x12201020, 0x12201022, 0x12201121, 0x12201222, + 0x12211120, 0x12211122, 0x12211220, 0x12211221, 0x12221020, 0x12221120, 0x12221122, 0x12221222, + 0x10212102, 0x10212201, 0x10222101, 0x11202001, 0x11212002, 0x11212101, 0x11212202, 0x11222001, + 0x11222201, 0x12202101, 0x12212001, 0x12212200, 0x12222102, 0x10202011, 0x10202110, 0x10212010, + 0x10212111, 0x10222011, 0x10222110, 0x10222112, 0x10222211, 0x11202010, 0x11202011, 0x11202111, + 0x11202112, 0x11202210, 0x11212011, 0x11212110, 0x11212111, 0x11212112, 0x11212211, 0x11222010, + 0x11222111, 0x11222212, 0x12202012, 0x12202110, 0x12202212, 0x12212111, 0x12222011, 0x12222110, + 0x12222111, 0x12222211, 0x10212021, 0x10212122, 0x10212220, 0x11202021, 0x11202120, 0x11202221, + 0x11212020, 0x11212121, 0x11212220, 0x11212222, 0x11222120, 0x11222121, 0x11222221, 0x12202122, + 0x12212120, 0x12212220, 0x12212222, 0x12222122, 0x20000000, 0x20000002, 0x20000200, 0x20000202, + 0x20020000, 0x20020002, 0x20020200, 0x20020202, 0x21000101, 0x21010000, 0x21010001, 0x21010100, + 0x21010102, 0x21010201, 0x21020101, 0x22000000, 0x22000002, 0x22000200, 0x22000202, 0x22010101, + 0x22020000, 0x22020002, 0x22020200, 0x22020202, 0x20000111, 0x20010011, 0x20010110, 0x20010112, + 0x20010211, 0x20020111, 0x21000011, 0x21000110, 0x21000211, 0x21010010, 0x21010012, 0x21010111, + 0x21010112, 0x21010210, 0x21010211, 0x21020110, 0x21020112, 0x21020211, 0x22000111, 0x22000211, + 0x22010110, 0x22010112, 0x22010211, 0x22020111, 0x20000020, 0x20000022, 0x20000220, 0x20000222, + 0x20010121, 0x20020020, 0x20020022, 0x20020220, 0x20020222, 0x21010021, 0x21010120, 0x21010221, + 0x21020121, 0x22000020, 0x22000022, 0x22000220, 0x22000222, 0x22010121, 0x22020020, 0x22020022, + 0x22020220, 0x22020222, 0x20011100, 0x20011201, 0x21001001, 0x21001100, 0x21011001, 0x21011101, + 0x21011202, 0x21021001, 0x21021100, 0x21021201, 0x22011100, 0x22011201, 0x20001011, 0x20001211, + 0x20011012, 0x20011111, 0x20011212, 0x20021112, 0x20021211, 0x21001010, 0x21001011, 0x21001111, + 0x21001210, 0x21011011, 0x21011110, 0x21011111, 0x21011112, 0x21011211, 0x21011212, 0x21021111, + 0x21021112, 0x21021210, 0x21021212, 0x22001011, 0x22001110, 0x22001112, 0x22001211, 0x22011010, + 0x22011012, 0x22011111, 0x22011210, 0x22021112, 0x20011021, 0x20011122, 0x20011221, 0x20021121, + 0x21001021, 0x21001120, 0x21001221, 0x21001222, 0x21011020, 0x21011121, 0x21011221, 0x21011222, + 0x21021021, 0x21021122, 0x21021222, 0x22001121, 0x22011021, 0x22011222, 0x22021120, 0x20002000, + 0x20002002, 0x20002200, 0x20002202, 0x20012101, 0x20022000, 0x20022002, 0x20022200, 0x20022202, + 0x21002001, 0x21002101, 0x21012001, 0x21012100, 0x21012201, 0x21022101, 0x21022201, 0x22002000, + 0x22002002, 0x22002200, 0x22002202, 0x22012101, 0x22022000, 0x22022002, 0x22022200, 0x22022202, + 0x20002111, 0x20002112, 0x20012011, 0x20012110, 0x20012112, 0x20022111, 0x21002011, 0x21002110, + 0x21002112, 0x21002211, 0x21012010, 0x21012012, 0x21012111, 0x21012212, 0x21022011, 0x21022110, + 0x22002111, 0x22012112, 0x22012211, 0x22022111, 0x20002020, 0x20002022, 0x20002220, 0x20002222, + 0x20012121, 0x20022020, 0x20022022, 0x20022220, 0x20022222, 0x21002121, 0x21012021, 0x21012120, + 0x21012122, 0x22002020, 0x22002022, 0x22002220, 0x22002222, 0x22012121, 0x22022020, 0x22022022, + 0x22022220, 0x22022222, 0x20100101, 0x20110001, 0x20110102, 0x20110200, 0x20110201, 0x20120101, + 0x21100001, 0x21100102, 0x21100201, 0x21110101, 0x21110200, 0x21110202, 0x21120201, 0x21120202, + 0x22100101, 0x22110001, 0x22110100, 0x22110102, 0x22110201, 0x22120101, 0x20100011, 0x20100110, + 0x20100112, 0x20100211, 0x20110010, 0x20110111, 0x20110210, 0x20110212, 0x20120011, 0x20120110, + 0x20120112, 0x20120211, 0x21100010, 0x21100111, 0x21110010, 0x21110011, 0x21110110, 0x21110111, + 0x21110112, 0x21110211, 0x21120012, 0x21120111, 0x22100110, 0x22100112, 0x22110012, 0x22110111, + 0x22110210, 0x22120011, 0x22120110, 0x22120112, 0x22120211, 0x20100121, 0x20110021, 0x20110120, + 0x20110221, 0x20120121, 0x21100120, 0x21100122, 0x21100221, 0x21110020, 0x21110022, 0x21110121, + 0x21110220, 0x21120122, 0x21120221, 0x22100121, 0x22110120, 0x22110122, 0x22120221, 0x20101001, + 0x20101100, 0x20101102, 0x20111000, 0x20111101, 0x20111200, 0x20121102, 0x21101000, 0x21101202, + 0x21111001, 0x21111100, 0x21111101, 0x21111102, 0x21111200, 0x21111201, 0x21121000, 0x21121001, + 0x21121002, 0x21121101, 0x22101100, 0x22101102, 0x22111002, 0x22111100, 0x22111101, 0x22111200, + 0x22121001, 0x22121201, 0x20101010, 0x20101111, 0x20101210, 0x20101212, 0x20111010, 0x20111011, + 0x20111110, 0x20111111, 0x20111112, 0x20111211, 0x20121011, 0x20121111, 0x20121211, 0x20121212, + 0x21101011, 0x21101110, 0x21101111, 0x21101112, 0x21101211, 0x21111010, 0x21111011, 0x21111012, + 0x21111110, 0x21111111, 0x21111112, 0x21111210, 0x21111211, 0x21111212, 0x21121011, 0x21121110, + 0x21121111, 0x21121112, 0x21121211, 0x22101011, 0x22101111, 0x22101210, 0x22111011, 0x22111012, + 0x22111110, 0x22111111, 0x22111112, 0x22111211, 0x22111212, 0x22121010, 0x22121012, 0x22121111, + 0x22121210, 0x22121212, 0x20101021, 0x20101120, 0x20111020, 0x20111121, 0x20111221, 0x20121020, + 0x20121122, 0x20121221, 0x21101121, 0x21101220, 0x21101221, 0x21111021, 0x21111022, 0x21111121, + 0x21111122, 0x21111221, 0x21121121, 0x21121220, 0x22101022, 0x22101120, 0x22101221, 0x22101222, + 0x22111022, 0x22111120, 0x22111121, 0x22121120, 0x22121122, 0x22121221, 0x20102101, 0x20112102, + 0x20112201, 0x20122101, 0x21102001, 0x21102102, 0x21112000, 0x21112002, 0x21112101, 0x21112102, + 0x21112202, 0x21122100, 0x21122101, 0x22102101, 0x22112001, 0x22112102, 0x22112201, 0x22122101, + 0x20102110, 0x20102112, 0x20102211, 0x20112010, 0x20112012, 0x20112111, 0x20112210, 0x20112212, + 0x20122010, 0x20122011, 0x20122110, 0x20122112, 0x21102010, 0x21102012, 0x21102111, 0x21102210, + 0x21102212, 0x21112011, 0x21112110, 0x21112111, 0x21112112, 0x21112211, 0x21122012, 0x21122111, + 0x21122112, 0x21122212, 0x22102011, 0x22102110, 0x22112010, 0x22112012, 0x22112111, 0x22112212, + 0x22122011, 0x22122112, 0x20102121, 0x20112121, 0x20122121, 0x21102120, 0x21102122, 0x21102221, + 0x21112020, 0x21112121, 0x21112220, 0x21122021, 0x22102121, 0x22112021, 0x22112120, 0x22112121, + 0x22112122, 0x20200000, 0x20200002, 0x20200200, 0x20200202, 0x20210101, 0x20220000, 0x20220002, + 0x20220200, 0x20220202, 0x21200101, 0x21210001, 0x21210100, 0x21210102, 0x21210201, 0x22200000, + 0x22200002, 0x22200200, 0x22200202, 0x22210101, 0x22220000, 0x22220002, 0x22220200, 0x22220202, + 0x20200111, 0x20200211, 0x20210011, 0x20210110, 0x20210112, 0x20210211, 0x20210212, 0x21200112, + 0x21200211, 0x21210011, 0x21210111, 0x21210210, 0x21210212, 0x21220011, 0x21220110, 0x22200111, + 0x22210010, 0x22210012, 0x22210112, 0x22210211, 0x20200022, 0x20200220, 0x20200222, 0x20210020, + 0x20210221, 0x20220022, 0x20220220, 0x20220222, 0x21200121, 0x21210021, 0x21210122, 0x21210221, + 0x21220121, 0x22200020, 0x22200022, 0x22200220, 0x22200222, 0x22210121, 0x22220020, 0x22220022, + 0x22220220, 0x22220222, 0x20211201, 0x20221101, 0x21201001, 0x21201100, 0x21211000, 0x21211100, + 0x21211101, 0x21211200, 0x21211202, 0x21221001, 0x21221101, 0x21221102, 0x21221200, 0x21221201, + 0x22201101, 0x20201112, 0x20201211, 0x20211010, 0x20211012, 0x20211111, 0x20211210, 0x20221112, + 0x20221211, 0x21201012, 0x21201111, 0x21211011, 0x21211110, 0x21211111, 0x21211112, 0x21211211, + 0x21221111, 0x21221212, 0x22201011, 0x22201110, 0x22201111, 0x22201112, 0x22201211, 0x22211012, + 0x22211111, 0x22211210, 0x20201121, 0x20211021, 0x20211122, 0x20211222, 0x20221021, 0x20221121, + 0x21201120, 0x21201122, 0x21201222, 0x21211022, 0x21211121, 0x21211122, 0x21211220, 0x21221020, + 0x21221022, 0x22201122, 0x22211020, 0x22211121, 0x22211122, 0x22211221, 0x22221021, 0x22221120, + 0x22221122, 0x20202000, 0x20202002, 0x20202200, 0x20202202, 0x20222000, 0x20222002, 0x20222200, + 0x20222202, 0x21212001, 0x21212100, 0x21212102, 0x21212201, 0x22202000, 0x22202002, 0x22202200, + 0x22202202, 0x22212101, 0x22222000, 0x22222002, 0x22222200, 0x22222202, 0x20202111, 0x20212110, + 0x20212211, 0x20222011, 0x20222111, 0x21202011, 0x21212010, 0x21212111, 0x21212212, 0x21222011, + 0x21222112, 0x21222211, 0x22212010, 0x22212112, 0x20202020, 0x20202022, 0x20202220, 0x20202222, + 0x20222020, 0x20222022, 0x20222220, 0x20222222, 0x21212021, 0x21212120, 0x21212122, 0x22202020, + 0x22202022, 0x22202220, 0x22202222, 0x22212121, 0x22222020, 0x22222022, 0x22222220, 0x22222222, +GGML_TABLE_END() +#endif + +#endif // GGML_COMMON_IMPL +#endif // GGML_COMMON_IMPL diff --git a/python/freetoken/kernel/csrc/ggml/src/ggml-impl.h b/python/freetoken/kernel/csrc/ggml/src/ggml-impl.h new file mode 100644 index 0000000000..ae26e0c23b --- /dev/null +++ b/python/freetoken/kernel/csrc/ggml/src/ggml-impl.h @@ -0,0 +1,795 @@ +#pragma once + +// GGML internal header + +#include "ggml.h" +#include "gguf.h" + +#include +#include +#include // load `stdlib.h` before other headers to work around MinGW bug: https://sourceforge.net/p/mingw-w64/bugs/192/ +#include +#include +#include + +#ifdef __ARM_FEATURE_SVE +#include +#endif // __ARM_FEATURE_SVE + +#if defined(__ARM_NEON) && !defined(__CUDACC__) && !defined(__MUSACC__) +// if YCM cannot find , make a symbolic link to it, for example: +// +// $ ln -sfn /Library/Developer/CommandLineTools/usr/lib/clang/13.1.6/include/arm_neon.h ./src/ +// +#include +#endif + +#ifdef __cplusplus +extern "C" { +#endif + +void ggml_print_backtrace(void); + +uint64_t ggml_graph_next_uid(void); + +#ifndef MIN +# define MIN(a, b) ((a) < (b) ? (a) : (b)) +#endif + +#ifndef MAX +# define MAX(a, b) ((a) > (b) ? (a) : (b)) +#endif + +// required for mmap as gguf only guarantees 32-byte alignment +#define TENSOR_ALIGNMENT 32 + +// static_assert should be a #define, but if it's not, +// fall back to the _Static_assert C11 keyword. +// if C99 - static_assert is noop +// ref: https://stackoverflow.com/a/53923785/4039976 +#ifndef __cplusplus + #ifndef static_assert + #if defined(__STDC_VERSION__) && (__STDC_VERSION__ >= 201100L) + #define static_assert(cond, msg) _Static_assert(cond, msg) + #else + #define static_assert(cond, msg) struct global_scope_noop_trick + #endif + #endif +#endif + +static inline int ggml_up32(int n) { + return (n + 31) & ~31; +} + +//static inline int ggml_up64(int n) { +// return (n + 63) & ~63; +//} + +static inline int ggml_up(int n, int m) { + // assert m is a power of 2 + GGML_ASSERT((m & (m - 1)) == 0); + return (n + m - 1) & ~(m - 1); +} + +// TODO: move to ggml.h? (won't be able to inline) +static bool ggml_are_same_layout(const struct ggml_tensor * a, const struct ggml_tensor * b) { + if (a->type != b->type) { + return false; + } + for (int i = 0; i < GGML_MAX_DIMS; i++) { + if (a->ne[i] != b->ne[i]) { + return false; + } + if (a->nb[i] != b->nb[i]) { + return false; + } + } + return true; +} + +static bool ggml_op_is_empty(enum ggml_op op) { + switch (op) { + case GGML_OP_NONE: + case GGML_OP_RESHAPE: + case GGML_OP_TRANSPOSE: + case GGML_OP_VIEW: + case GGML_OP_PERMUTE: + return true; + default: + return false; + } +} + +static inline bool ggml_impl_is_view(const struct ggml_tensor * t) { + return t->view_src != NULL; +} + +static inline float ggml_compute_softplus_f32(float input) { + return (input > 20.0f) ? input : logf(1 + expf(input)); +} +// +// logging +// + +GGML_ATTRIBUTE_FORMAT(2, 3) +GGML_API void ggml_log_internal (enum ggml_log_level level, const char * format, ...); +GGML_API void ggml_log_callback_default(enum ggml_log_level level, const char * text, void * user_data); + +#define GGML_LOG(...) ggml_log_internal(GGML_LOG_LEVEL_NONE , __VA_ARGS__) +#define GGML_LOG_INFO(...) ggml_log_internal(GGML_LOG_LEVEL_INFO , __VA_ARGS__) +#define GGML_LOG_WARN(...) ggml_log_internal(GGML_LOG_LEVEL_WARN , __VA_ARGS__) +#define GGML_LOG_ERROR(...) ggml_log_internal(GGML_LOG_LEVEL_ERROR, __VA_ARGS__) +#define GGML_LOG_DEBUG(...) ggml_log_internal(GGML_LOG_LEVEL_DEBUG, __VA_ARGS__) +#define GGML_LOG_CONT(...) ggml_log_internal(GGML_LOG_LEVEL_CONT , __VA_ARGS__) + +#define GGML_DEBUG 0 + +#if (GGML_DEBUG >= 1) +#define GGML_PRINT_DEBUG(...) GGML_LOG_DEBUG(__VA_ARGS__) +#else +#define GGML_PRINT_DEBUG(...) +#endif + +#if (GGML_DEBUG >= 5) +#define GGML_PRINT_DEBUG_5(...) GGML_LOG_DEBUG(__VA_ARGS__) +#else +#define GGML_PRINT_DEBUG_5(...) +#endif + +#if (GGML_DEBUG >= 10) +#define GGML_PRINT_DEBUG_10(...) GGML_LOG_DEBUG(__VA_ARGS__) +#else +#define GGML_PRINT_DEBUG_10(...) +#endif + +// tensor params + +static void ggml_set_op_params(struct ggml_tensor * tensor, const void * params, size_t params_size) { + GGML_ASSERT(tensor != NULL); // silence -Warray-bounds warnings + assert(params_size <= GGML_MAX_OP_PARAMS); + memcpy(tensor->op_params, params, params_size); +} + +static int32_t ggml_get_op_params_i32(const struct ggml_tensor * tensor, uint32_t i) { + assert(i < GGML_MAX_OP_PARAMS / sizeof(int32_t)); + return ((const int32_t *)(tensor->op_params))[i]; +} + +static float ggml_get_op_params_f32(const struct ggml_tensor * tensor, uint32_t i) { + assert(i < GGML_MAX_OP_PARAMS / sizeof(float)); + return ((const float *)(tensor->op_params))[i]; +} + +// [TAG_GGML_PREC] +// - GGML_OP_MUL_MAT +// 0 - acc +// 1 - hint +// 2 - src0 precision +// 3 - src1 precision +// +// - GGML_OP_MUL_MAT_ID +// 0 - acc +// 1 - hint +// 2 - src0 precision +// 3 - src1 precision +static void ggml_set_op_params_i32(struct ggml_tensor * tensor, uint32_t i, int32_t value) { + assert(i < GGML_MAX_OP_PARAMS / sizeof(int32_t)); + ((int32_t *)(tensor->op_params))[i] = value; +} + +static void ggml_set_op_params_f32(struct ggml_tensor * tensor, uint32_t i, float value) { + assert(i < GGML_MAX_OP_PARAMS / sizeof(float)); + ((float *)(tensor->op_params))[i] = value; +} + +struct ggml_map_custom1_op_params { + ggml_custom1_op_t fun; + int n_tasks; + void * userdata; +}; + +struct ggml_map_custom2_op_params { + ggml_custom2_op_t fun; + int n_tasks; + void * userdata; +}; + +struct ggml_map_custom3_op_params { + ggml_custom3_op_t fun; + int n_tasks; + void * userdata; +}; + +struct ggml_custom_op_params { + ggml_custom_op_t fun; + int n_tasks; + void * userdata; +}; + +// bitset + +typedef uint32_t ggml_bitset_t; + +static_assert(sizeof(ggml_bitset_t) == 4, "bitset_t constants must be updated"); +#define BITSET_SHR 5 // log2(sizeof(ggml_bitset_t)*8) +#define BITSET_MASK (sizeof(ggml_bitset_t)*8 - 1) + +static size_t ggml_bitset_size(size_t n) { + return (n + BITSET_MASK) >> BITSET_SHR; +} + +static inline bool ggml_bitset_get(const ggml_bitset_t * bitset, size_t i) { + return !!(bitset[i >> BITSET_SHR] & (1u << (i & BITSET_MASK))); +} + +static inline void ggml_bitset_set(ggml_bitset_t * bitset, size_t i) { + bitset[i >> BITSET_SHR] |= (1u << (i & BITSET_MASK)); +} + +static inline void ggml_bitset_clear(ggml_bitset_t * bitset, size_t i) { + bitset[i >> BITSET_SHR] &= ~(1u << (i & BITSET_MASK)); +} + +// hash set + +#define GGML_HASHSET_FULL ((size_t)-1) +#define GGML_HASHSET_ALREADY_EXISTS ((size_t)-2) + +struct ggml_hash_set { + size_t size; + ggml_bitset_t * used; // whether or not the keys are in use i.e. set + struct ggml_tensor ** keys; // actual tensors in the set, keys[i] is only defined if ggml_bitset_get(used, i) +}; + +struct ggml_hash_set ggml_hash_set_new(size_t size); +void ggml_hash_set_free(struct ggml_hash_set * hash_set); + +// returns the minimum size for a hash set that can hold min_sz elements +size_t ggml_hash_size(size_t min_sz); + +// remove all elements from the hash set +void ggml_hash_set_reset(struct ggml_hash_set * hash_set); + +// returns true if key is in the hash set +static bool ggml_hash_contains(const struct ggml_hash_set * hash_set, struct ggml_tensor * key); + +// returns GGML_HASHSET_FULL if table is full, otherwise the current index of the key or where it should be inserted +static size_t ggml_hash_find(const struct ggml_hash_set * hash_set, const struct ggml_tensor * key); + +// returns GGML_HASHSET_ALREADY_EXISTS if key already exists, index otherwise, asserts if table is full +static size_t ggml_hash_insert(struct ggml_hash_set * hash_set, struct ggml_tensor * key); + +// return index, asserts if table is full +static size_t ggml_hash_find_or_insert(struct ggml_hash_set * hash_set, struct ggml_tensor * key); + +// hash function for ggml_tensor +static inline size_t ggml_hash(const struct ggml_tensor * p) { + // the last 4 bits are always zero due to alignment + return (size_t)(uintptr_t)p >> 4; +} + +static size_t ggml_hash_find(const struct ggml_hash_set * hash_set, const struct ggml_tensor * key) { + size_t h = ggml_hash(key) % hash_set->size; + + // linear probing + size_t i = h; + while (ggml_bitset_get(hash_set->used, i) && hash_set->keys[i] != key) { + i = (i + 1) % hash_set->size; + if (i == h) { + // visited all hash table entries -> not found + return GGML_HASHSET_FULL; + } + } + return i; +} + +static bool ggml_hash_contains(const struct ggml_hash_set * hash_set, struct ggml_tensor * key) { + size_t i = ggml_hash_find(hash_set, key); + return i != GGML_HASHSET_FULL && ggml_bitset_get(hash_set->used, i); +} + +static size_t ggml_hash_insert(struct ggml_hash_set * hash_set, struct ggml_tensor * key) { + size_t h = ggml_hash(key) % hash_set->size; + + // linear probing + size_t i = h; + do { + if (!ggml_bitset_get(hash_set->used, i)) { + ggml_bitset_set(hash_set->used, i); + hash_set->keys[i] = key; + return i; + } + if (hash_set->keys[i] == key) { + return GGML_HASHSET_ALREADY_EXISTS; + } + i = (i + 1) % hash_set->size; + } while (i != h); + + // visited all hash table entries -> not found + GGML_ABORT("fatal error"); +} + +static size_t ggml_hash_find_or_insert(struct ggml_hash_set * hash_set, struct ggml_tensor * key) { + size_t h = ggml_hash(key) % hash_set->size; + + // linear probing + size_t i = h; + do { + if (!ggml_bitset_get(hash_set->used, i)) { + ggml_bitset_set(hash_set->used, i); + hash_set->keys[i] = key; + return i; + } + if (hash_set->keys[i] == key) { + return i; + } + i = (i + 1) % hash_set->size; + } while (i != h); + + // visited all hash table entries -> not found + GGML_ABORT("fatal error"); +} + +// computation graph + +enum ggml_cgraph_eval_order { + GGML_CGRAPH_EVAL_ORDER_LEFT_TO_RIGHT = 0, + GGML_CGRAPH_EVAL_ORDER_RIGHT_TO_LEFT, + GGML_CGRAPH_EVAL_ORDER_COUNT +}; + +struct ggml_cgraph { + int size; // maximum number of nodes/leafs/grads/grad_accs + int n_nodes; // number of nodes currently in use + int n_leafs; // number of leafs currently in use + + struct ggml_tensor ** nodes; // tensors with data that can change if the graph is evaluated + struct ggml_tensor ** grads; // the outputs of these tensors are the gradients of the nodes + struct ggml_tensor ** grad_accs; // accumulators for node gradients + struct ggml_tensor ** leafs; // tensors with constant data + int32_t * use_counts;// number of uses of each tensor, indexed by hash table slot + + struct ggml_hash_set visited_hash_set; + + enum ggml_cgraph_eval_order order; + + // an optional identifier that can be utilized to recognize same graphs if two non-zero values match + // a value of 0 means it is not set and should be ignored + uint64_t uid; +}; + +// returns a slice of cgraph with nodes [i0, i1) +// the slice does not have leafs or gradients +// if you need the gradients, get them from the original graph +struct ggml_cgraph ggml_graph_view(struct ggml_cgraph * cgraph, int i0, int i1); + +// ggml-alloc.c: true if the operation can reuse memory from its sources +GGML_API bool ggml_op_can_inplace(enum ggml_op op); + + +// Memory allocation + +GGML_API void * ggml_aligned_malloc(size_t size); +GGML_API void ggml_aligned_free(void * ptr, size_t size); + +// FP16 <-> FP32 +// ref: https://github.com/Maratyszcza/FP16 + +static inline float fp32_from_bits(uint32_t w) { + union { + uint32_t as_bits; + float as_value; + } fp32; + fp32.as_bits = w; + return fp32.as_value; +} + +static inline uint32_t fp32_to_bits(float f) { + union { + float as_value; + uint32_t as_bits; + } fp32; + fp32.as_value = f; + return fp32.as_bits; +} + +static inline float ggml_compute_fp16_to_fp32(ggml_fp16_t h) { + const uint32_t w = (uint32_t) h << 16; + const uint32_t sign = w & UINT32_C(0x80000000); + const uint32_t two_w = w + w; + + const uint32_t exp_offset = UINT32_C(0xE0) << 23; +#if (defined(__STDC_VERSION__) && (__STDC_VERSION__ >= 199901L) || defined(__GNUC__) && !defined(__STRICT_ANSI__)) && (!defined(__cplusplus) || __cplusplus >= 201703L) + const float exp_scale = 0x1.0p-112f; +#else + const float exp_scale = fp32_from_bits(UINT32_C(0x7800000)); +#endif + const float normalized_value = fp32_from_bits((two_w >> 4) + exp_offset) * exp_scale; + + const uint32_t magic_mask = UINT32_C(126) << 23; + const float magic_bias = 0.5f; + const float denormalized_value = fp32_from_bits((two_w >> 17) | magic_mask) - magic_bias; + + const uint32_t denormalized_cutoff = UINT32_C(1) << 27; + const uint32_t result = sign | + (two_w < denormalized_cutoff ? fp32_to_bits(denormalized_value) : fp32_to_bits(normalized_value)); + return fp32_from_bits(result); +} + +static inline ggml_fp16_t ggml_compute_fp32_to_fp16(float f) { +#if (defined(__STDC_VERSION__) && (__STDC_VERSION__ >= 199901L) || defined(__GNUC__) && !defined(__STRICT_ANSI__)) && (!defined(__cplusplus) || __cplusplus >= 201703L) + const float scale_to_inf = 0x1.0p+112f; + const float scale_to_zero = 0x1.0p-110f; +#else + const float scale_to_inf = fp32_from_bits(UINT32_C(0x77800000)); + const float scale_to_zero = fp32_from_bits(UINT32_C(0x08800000)); +#endif + float base = (fabsf(f) * scale_to_inf) * scale_to_zero; + + const uint32_t w = fp32_to_bits(f); + const uint32_t shl1_w = w + w; + const uint32_t sign = w & UINT32_C(0x80000000); + uint32_t bias = shl1_w & UINT32_C(0xFF000000); + if (bias < UINT32_C(0x71000000)) { + bias = UINT32_C(0x71000000); + } + + base = fp32_from_bits((bias >> 1) + UINT32_C(0x07800000)) + base; + const uint32_t bits = fp32_to_bits(base); + const uint32_t exp_bits = (bits >> 13) & UINT32_C(0x00007C00); + const uint32_t mantissa_bits = bits & UINT32_C(0x00000FFF); + const uint32_t nonsign = exp_bits + mantissa_bits; + return (sign >> 16) | (shl1_w > UINT32_C(0xFF000000) ? UINT16_C(0x7E00) : nonsign); +} + +#define GGML_COMPUTE_FP16_TO_FP32(x) ggml_compute_fp16_to_fp32(x) +#define GGML_COMPUTE_FP32_TO_FP16(x) ggml_compute_fp32_to_fp16(x) + +#define GGML_FP16_TO_FP32(x) GGML_COMPUTE_FP16_TO_FP32(x) +#define GGML_FP32_TO_FP16(x) GGML_COMPUTE_FP32_TO_FP16(x) + +static inline float ggml_e8m0_to_fp32(uint8_t x) { + uint32_t bits; // Stores the raw bit representation of the float + + // Handle special case for minimum exponent (denormalized float) + if (x == 0) { + // Bit pattern for 2^(-127): + // - Sign bit: 0 (positive) + // - Exponent: 0 (denormalized number) + // - Mantissa: 0x400000 (0.5 in fractional form) + // Value = 0.5 * 2^(-126) = 2^(-127) + bits = 0x00400000; + } + // note: disabled as we don't need to handle NaNs + //// Handle special case for NaN (all bits set) + //else if (x == 0xFF) { + // // Standard quiet NaN pattern: + // // - Sign bit: 0 + // // - Exponent: all 1s (0xFF) + // // - Mantissa: 0x400000 (quiet NaN flag) + // bits = 0x7FC00000; + //} + // Normalized values (most common case) + else { + // Construct normalized float by shifting exponent into position: + // - Exponent field: 8 bits (positions 30-23) + // - Mantissa: 0 (implicit leading 1) + // Value = 2^(x - 127) + bits = (uint32_t) x << 23; + } + + float result; // Final float value + // Safely reinterpret bit pattern as float without type-punning issues + memcpy(&result, &bits, sizeof(float)); + return result; +} + +// Equal to ggml_e8m0_to_fp32/2 +// Useful with MXFP4 quantization since the E0M2 values are doubled +static inline float ggml_e8m0_to_fp32_half(uint8_t x) { + uint32_t bits; + + // For x < 2: use precomputed denormal patterns + if (x < 2) { + // 0x00200000 = 2^(-128), 0x00400000 = 2^(-127) + bits = 0x00200000 << x; + } + // For x >= 2: normalized exponent adjustment + else { + // 0.5 * 2^(x-127) = 2^(x-128) = normalized with exponent (x-1) + bits = (uint32_t)(x - 1) << 23; + } + // Note: NaNs are not handled here + + float result; + memcpy(&result, &bits, sizeof(float)); + return result; +} + +#define GGML_E8M0_TO_FP32(x) ggml_e8m0_to_fp32(x) +#define GGML_E8M0_TO_FP32_HALF(x) ggml_e8m0_to_fp32_half(x) + +// UE4M3: unsigned, 4 exp bits (bias=7), 3 mantissa bits +// Returns value * 0.5 to match kvalues_mxfp4 convention (kvalues = 2 * E2M1_float) +static inline float ggml_ue4m3_to_fp32(uint8_t x) { + if (x == 0 || x == 0x7F) { + return 0.0f; + } + int exp = (x >> 3) & 0xF; + int man = x & 0x7; + float raw; + if (exp == 0) { + raw = ldexpf((float) man, -9); + } else { + raw = ldexpf(1.0f + (float) man / 8.0f, exp - 7); + } + return raw * 0.5f; +} + +static inline uint8_t ggml_fp32_to_ue4m3(float x) { + if (!(x > 0.0f)) { + return 0; + } + if (x > 448.0f) { + x = 448.0f; + } + uint32_t bits; + memcpy(&bits, &x, 4); + int fp32_exp = ((bits >> 23) & 0xFF) - 127; + int fp32_man = (bits >> 20) & 0x7; + int ue4m3_exp = fp32_exp + 7; + if (ue4m3_exp <= 0) { + // subnormal: value = man * 2^-9, man = round(x * 2^9) + int man = (int) (x * 512.0f + 0.5f); + if (man > 7) { + man = 7; + } + if (man < 1) { + return 0; + } + return (uint8_t) man; + } + if (ue4m3_exp >= 15) { + return 0x7E; + } + int round_bit = (bits >> 19) & 1; + int ue4m3_man = fp32_man + round_bit; + if (ue4m3_man > 7) { + ue4m3_man = 0; + ue4m3_exp++; + if (ue4m3_exp >= 15) { + return 0x7E; + } + } + return (uint8_t) ((ue4m3_exp << 3) | ue4m3_man); +} + +/** + * Converts brain16 to float32. + * + * The bfloat16 floating point format has the following structure: + * + * ┌sign + * │ + * │ ┌exponent + * │ │ + * │ │ ┌mantissa + * │ │ │ + * │┌──┴───┐┌─┴───┐ + * 0b0000000000000000 brain16 + * + * Since bf16 has the same number of exponent bits as a 32bit float, + * encoding and decoding numbers becomes relatively straightforward. + * + * ┌sign + * │ + * │ ┌exponent + * │ │ + * │ │ ┌mantissa + * │ │ │ + * │┌──┴───┐┌─┴───────────────────┐ + * 0b00000000000000000000000000000000 IEEE binary32 + * + * For comparison, the standard fp16 format has fewer exponent bits. + * + * ┌sign + * │ + * │ ┌exponent + * │ │ + * │ │ ┌mantissa + * │ │ │ + * │┌─┴─┐┌─┴──────┐ + * 0b0000000000000000 IEEE binary16 + * + * @see IEEE 754-2008 + */ +static inline float ggml_compute_bf16_to_fp32(ggml_bf16_t h) { + union { + float f; + uint32_t i; + } u; + u.i = (uint32_t)h.bits << 16; + return u.f; +} + +/** + * Converts float32 to brain16. + * + * This is binary identical with Google Brain float conversion. + * Floats shall round to nearest even, and NANs shall be quiet. + * Subnormals aren't flushed to zero, except perhaps when used. + * This code should vectorize nicely if using modern compilers. + */ +static inline ggml_bf16_t ggml_compute_fp32_to_bf16(float s) { + ggml_bf16_t h; + union { + float f; + uint32_t i; + } u; + u.f = s; + if ((u.i & 0x7fffffff) > 0x7f800000) { /* nan */ + h.bits = (u.i >> 16) | 64; /* force to quiet */ + return h; + } + h.bits = (u.i + (0x7fff + ((u.i >> 16) & 1))) >> 16; + return h; +} + +#define GGML_FP32_TO_BF16(x) ggml_compute_fp32_to_bf16(x) +#define GGML_BF16_TO_FP32(x) ggml_compute_bf16_to_fp32(x) + +static inline int32_t ggml_node_get_use_count(const struct ggml_cgraph * cgraph, int node_idx) { + const struct ggml_tensor * node = cgraph->nodes[node_idx]; + + size_t hash_pos = ggml_hash_find(&cgraph->visited_hash_set, node); + if (!ggml_bitset_get(cgraph->visited_hash_set.used, hash_pos)) { + return 0; + } + return cgraph->use_counts[hash_pos]; +} + +// return true if the node's results are only used by N other nodes +// and can be fused into their calculations. +static inline bool ggml_node_has_n_uses(const struct ggml_cgraph * cgraph, int node_idx, int32_t n_uses) { + const struct ggml_tensor * node = cgraph->nodes[node_idx]; + + // check the use count against how many we're replacing + if (ggml_node_get_use_count(cgraph, node_idx) != n_uses) { + return false; + } + + // if node is a view, some other node might be using the intermediate result + // via the view source. + if (node->view_src) { + return false; + } + + // If the user requested output for the node, can't fuse + if (node->flags & GGML_TENSOR_FLAG_OUTPUT) { + return false; + } + + return true; +} + +// Returns true if nodes with indices { node_idxs } are the sequence of ggml_ops in ops[] +// and are fusable. Nodes are considered fusable according to this function if: +// - all nodes except the last have only one use and are not views/outputs (see ggml_node_has_N_uses). +// - all nodes except the last are a src of the following node. +// - all nodes are the same shape. +// TODO: Consider allowing GGML_OP_NONE nodes in between +static inline bool ggml_can_fuse_ext(const struct ggml_cgraph * cgraph, const int * node_idxs, const enum ggml_op * ops, int num_ops) { + for (int i = 0; i < num_ops; ++i) { + if (node_idxs[i] >= cgraph->n_nodes) { + return false; + } + + struct ggml_tensor * node = cgraph->nodes[node_idxs[i]]; + if (node->op != ops[i]) { + return false; + } + if ((node->flags & GGML_TENSOR_FLAG_COMPUTE) == 0) { + return false; + } + if (i < num_ops - 1 && !ggml_node_has_n_uses(cgraph, node_idxs[i], 1)) { + return false; + } + if (i > 0) { + struct ggml_tensor * prev = cgraph->nodes[node_idxs[i - 1]]; + if (node->src[0] != prev && node->src[1] != prev) { + return false; + } + if (!ggml_are_same_shape(node, prev)) { + return false; + } + } + } + return true; +} + +// same as above, for sequential indices starting at node_idx +static inline bool ggml_can_fuse(const struct ggml_cgraph * cgraph, int node_idx, const enum ggml_op * ops, int num_ops) { + assert(num_ops < 32); + + if (node_idx + num_ops > cgraph->n_nodes) { + return false; + } + + int idxs[32]; + for (int i = 0; i < num_ops; ++i) { + idxs[i] = node_idx + i; + } + + return ggml_can_fuse_ext(cgraph, idxs, ops, num_ops); +} + +GGML_API bool ggml_can_fuse_subgraph_ext(const struct ggml_cgraph * cgraph, + const int * node_idxs, + int count, + const enum ggml_op * ops, + const int * outputs, + int num_outputs); + +// Returns true if the subgraph formed by {node_idxs} can be fused +// checks whethers all nodes which are not part of outputs can be elided +// by checking if their num_uses are confined to the subgraph +static inline bool ggml_can_fuse_subgraph(const struct ggml_cgraph * cgraph, + int node_idx, + int count, + const enum ggml_op * ops, + const int * outputs, + int num_outputs) { + GGML_ASSERT(count < 32); + if (node_idx + count > cgraph->n_nodes) { + return false; + } + + int idxs[32]; + + for (int i = 0; i < count; ++i) { + idxs[i] = node_idx + i; + } + + return ggml_can_fuse_subgraph_ext(cgraph, idxs, count, ops, outputs, num_outputs); +} + +#ifdef __cplusplus +} +#endif + +#ifdef __cplusplus +#include +#include +#include + +// nicer C++ syntax for ggml_can_fuse +inline bool ggml_can_fuse(const struct ggml_cgraph * cgraph, int node_idx, std::initializer_list ops) { + return ggml_can_fuse(cgraph, node_idx, ops.begin(), (int)ops.size()); +} + +inline bool ggml_can_fuse_subgraph(const struct ggml_cgraph * cgraph, + int start_idx, + std::initializer_list ops, + std::initializer_list outputs = {}) { + return ggml_can_fuse_subgraph(cgraph, start_idx, ops.size(), ops.begin(), outputs.begin(), outputs.size()); +} + +// Return true if the edges in the graph match expectations. +inline bool ggml_check_edges(const struct ggml_cgraph * cgraph, + int start_idx, + std::initializer_list> edges) { + for (const auto & edge : edges) { + int dst_node = edge[0]; + int src_idx = edge[1]; + int src_node = edge[2]; + if (cgraph->nodes[start_idx + dst_node]->src[src_idx] != cgraph->nodes[start_idx + src_node]) { + return false; + } + } + return true; +} + +// expose GGUF internals for test code +GGML_API size_t gguf_type_size(enum gguf_type type); +GGML_API void gguf_write_to_buf(const struct gguf_context * ctx, std::vector & buf, bool only_meta); +#endif // __cplusplus diff --git a/python/freetoken/kernel/csrc/ggml/src/ggml-quants.c b/python/freetoken/kernel/csrc/ggml/src/ggml-quants.c new file mode 100644 index 0000000000..7750a72ce6 --- /dev/null +++ b/python/freetoken/kernel/csrc/ggml/src/ggml-quants.c @@ -0,0 +1,5639 @@ +#define GGML_COMMON_IMPL_C +#include "ggml-common.h" + +#include "ggml-quants.h" +#include "ggml-impl.h" +#include "ggml-cpu/ggml-cpu-impl.h" +#include "ggml-cpu.h" + +#include +#include +#include +#include +#include // for qsort +#include // for GGML_ASSERT + +#ifdef GGML_USE_OPENMP +#include +#endif + +#define GROUP_MAX_EPS 1e-15f +#define GROUP_MAX_EPS_IQ3_XXS 1e-8f +#define GROUP_MAX_EPS_IQ2_S 1e-8f +#define GROUP_MAX_EPS_IQ1_M 1e-7f +#define GROUP_MAX_EPS_IQ1_S 1e-12f + +#define UNUSED GGML_UNUSED + +static inline int best_index_int8(int n, const int8_t * val, float x) { + if (x <= val[0]) return 0; + if (x >= val[n-1]) return n-1; + int ml = 0, mu = n-1; + while (mu-ml > 1) { + int mav = (ml+mu)/2; + if (x < val[mav]) mu = mav; else ml = mav; + } + return x - val[mu-1] < val[mu] - x ? mu-1 : mu; +} + +// reference implementation for deterministic creation of model files +void quantize_row_q1_0_ref(const float * GGML_RESTRICT x, block_q1_0 * GGML_RESTRICT y, int64_t k) { + static const int qk = QK1_0; + + assert(k % qk == 0); + + const int nb = k / qk; + + for (int i = 0; i < nb; i++) { + float sum_abs = 0.0f; + for (int j = 0; j < qk; j++) { + sum_abs += fabsf(x[i*qk + j]); + } + const float d = sum_abs / qk; + + y[i].d = GGML_FP32_TO_FP16(d); + + // Clear all bits first + for (int j = 0; j < qk / 8; ++j) { + y[i].qs[j] = 0; + } + + // Just store sign of each weight directly (no normalization) + for (int j = 0; j < qk; ++j) { + const int bit_index = j; + const int byte_index = bit_index / 8; + const int bit_offset = bit_index % 8; + + if (x[i*qk + j] >= 0.0f) { + y[i].qs[byte_index] |= (1 << bit_offset); + } + } + } +} + +void quantize_row_q2_0_ref(const float * GGML_RESTRICT x, block_q2_0 * GGML_RESTRICT y, int64_t k) { + static const int qk = QK2_0; + + assert(k % qk == 0); + + const int nb = k / qk; + + for (int i = 0; i < nb; i++) { + // Compute scale as max absolute value in the block + float amax = 0.0f; + for (int j = 0; j < qk; j++) { + const float a = fabsf(x[i*qk + j]); + if (a > amax) amax = a; + } + const float d = amax; + const float id = d > 0.0f ? 1.0f / d : 0.0f; + + y[i].d = GGML_FP32_TO_FP16(d); + + // Clear quant bytes + for (int j = 0; j < qk / 4; ++j) { + y[i].qs[j] = 0; + } + + // Encode 2-bit values: round(w/d) clamped to [-1, 2], then add 1 + // 00 (-1) = -scale, 01 (0) = 0, 10 (+1) = +scale, 11 (+2) = 2*scale + for (int j = 0; j < qk; ++j) { + const float w = x[i*qk + j]; + int q = (int)roundf(w * id) + 1; + if (q < 0) q = 0; + if (q > 3) q = 3; + const int byte_index = j / 4; + const int bit_offset = (j % 4) * 2; + y[i].qs[byte_index] |= ((uint8_t)q << bit_offset); + } + } +} + +// reference implementation for deterministic creation of model files +void quantize_row_q4_0_ref(const float * GGML_RESTRICT x, block_q4_0 * GGML_RESTRICT y, int64_t k) { + static const int qk = QK4_0; + + assert(k % qk == 0); + + const int nb = k / qk; + + for (int i = 0; i < nb; i++) { + float amax = 0.0f; // absolute max + float max = 0.0f; + + for (int j = 0; j < qk; j++) { + const float v = x[i*qk + j]; + if (amax < fabsf(v)) { + amax = fabsf(v); + max = v; + } + } + + const float d = max / -8; + const float id = d ? 1.0f/d : 0.0f; + + y[i].d = GGML_FP32_TO_FP16(d); + + for (int j = 0; j < qk/2; ++j) { + const float x0 = x[i*qk + 0 + j]*id; + const float x1 = x[i*qk + qk/2 + j]*id; + + const uint8_t xi0 = MIN(15, (int8_t)(x0 + 8.5f)); + const uint8_t xi1 = MIN(15, (int8_t)(x1 + 8.5f)); + + y[i].qs[j] = xi0; + y[i].qs[j] |= xi1 << 4; + } + } +} + +void quantize_row_q4_1_ref(const float * GGML_RESTRICT x, block_q4_1 * GGML_RESTRICT y, int64_t k) { + const int qk = QK4_1; + + assert(k % qk == 0); + + const int nb = k / qk; + + for (int i = 0; i < nb; i++) { + float min = FLT_MAX; + float max = -FLT_MAX; + + for (int j = 0; j < qk; j++) { + const float v = x[i*qk + j]; + + if (v < min) min = v; + if (v > max) max = v; + } + + const float d = (max - min) / ((1 << 4) - 1); + const float id = d ? 1.0f/d : 0.0f; + + y[i].d = GGML_FP32_TO_FP16(d); + y[i].m = GGML_FP32_TO_FP16(min); + + for (int j = 0; j < qk/2; ++j) { + const float x0 = (x[i*qk + 0 + j] - min)*id; + const float x1 = (x[i*qk + qk/2 + j] - min)*id; + + const uint8_t xi0 = MIN(15, (int8_t)(x0 + 0.5f)); + const uint8_t xi1 = MIN(15, (int8_t)(x1 + 0.5f)); + + y[i].qs[j] = xi0; + y[i].qs[j] |= xi1 << 4; + } + } +} + +void quantize_row_q5_0_ref(const float * GGML_RESTRICT x, block_q5_0 * GGML_RESTRICT y, int64_t k) { + static const int qk = QK5_0; + + assert(k % qk == 0); + + const int nb = k / qk; + + for (int i = 0; i < nb; i++) { + float amax = 0.0f; // absolute max + float max = 0.0f; + + for (int j = 0; j < qk; j++) { + const float v = x[i*qk + j]; + if (amax < fabsf(v)) { + amax = fabsf(v); + max = v; + } + } + + const float d = max / -16; + const float id = d ? 1.0f/d : 0.0f; + + y[i].d = GGML_FP32_TO_FP16(d); + + uint32_t qh = 0; + + for (int j = 0; j < qk/2; ++j) { + const float x0 = x[i*qk + 0 + j]*id; + const float x1 = x[i*qk + qk/2 + j]*id; + + const uint8_t xi0 = MIN(31, (int8_t)(x0 + 16.5f)); + const uint8_t xi1 = MIN(31, (int8_t)(x1 + 16.5f)); + + y[i].qs[j] = (xi0 & 0x0F) | ((xi1 & 0x0F) << 4); + + // get the 5-th bit and store it in qh at the right position + qh |= ((xi0 & 0x10u) >> 4) << (j + 0); + qh |= ((xi1 & 0x10u) >> 4) << (j + qk/2); + } + + memcpy(&y[i].qh, &qh, sizeof(qh)); + } +} + +void quantize_row_q5_1_ref(const float * GGML_RESTRICT x, block_q5_1 * GGML_RESTRICT y, int64_t k) { + const int qk = QK5_1; + + assert(k % qk == 0); + + const int nb = k / qk; + + for (int i = 0; i < nb; i++) { + float min = FLT_MAX; + float max = -FLT_MAX; + + for (int j = 0; j < qk; j++) { + const float v = x[i*qk + j]; + + if (v < min) min = v; + if (v > max) max = v; + } + + const float d = (max - min) / ((1 << 5) - 1); + const float id = d ? 1.0f/d : 0.0f; + + y[i].d = GGML_FP32_TO_FP16(d); + y[i].m = GGML_FP32_TO_FP16(min); + + uint32_t qh = 0; + + for (int j = 0; j < qk/2; ++j) { + const float x0 = (x[i*qk + 0 + j] - min)*id; + const float x1 = (x[i*qk + qk/2 + j] - min)*id; + + const uint8_t xi0 = (uint8_t)(x0 + 0.5f); + const uint8_t xi1 = (uint8_t)(x1 + 0.5f); + + y[i].qs[j] = (xi0 & 0x0F) | ((xi1 & 0x0F) << 4); + + // get the 5-th bit and store it in qh at the right position + qh |= ((xi0 & 0x10u) >> 4) << (j + 0); + qh |= ((xi1 & 0x10u) >> 4) << (j + qk/2); + } + + memcpy(&y[i].qh, &qh, sizeof(y[i].qh)); + } +} + +// reference implementation for deterministic creation of model files +void quantize_row_q8_0_ref(const float * GGML_RESTRICT x, block_q8_0 * GGML_RESTRICT y, int64_t k) { + assert(k % QK8_0 == 0); + const int nb = k / QK8_0; + + for (int i = 0; i < nb; i++) { + float amax = 0.0f; // absolute max + + for (int j = 0; j < QK8_0; j++) { + const float v = x[i*QK8_0 + j]; + amax = MAX(amax, fabsf(v)); + } + + const float d = amax / ((1 << 7) - 1); + const float id = d ? 1.0f/d : 0.0f; + + y[i].d = GGML_FP32_TO_FP16(d); + + for (int j = 0; j < QK8_0; ++j) { + const float x0 = x[i*QK8_0 + j]*id; + + y[i].qs[j] = roundf(x0); + } + } +} + +// reference implementation for deterministic creation of model files +void quantize_row_q8_1_ref(const float * GGML_RESTRICT x, block_q8_1 * GGML_RESTRICT y, int64_t k) { + assert(QK8_1 == 32); + assert(k % QK8_1 == 0); + const int nb = k / QK8_1; + + for (int i = 0; i < nb; i++) { + float amax = 0.0f; // absolute max + + for (int j = 0; j < QK8_1; j++) { + const float v = x[i*QK8_1 + j]; + amax = MAX(amax, fabsf(v)); + } + + const float d = amax / ((1 << 7) - 1); + const float id = d ? 1.0f/d : 0.0f; + + y[i].d = GGML_FP32_TO_FP16(d); + + int sum = 0; + + for (int j = 0; j < QK8_1/2; ++j) { + const float v0 = x[i*QK8_1 + j]*id; + const float v1 = x[i*QK8_1 + QK8_1/2 + j]*id; + + y[i].qs[ j] = roundf(v0); + y[i].qs[QK8_1/2 + j] = roundf(v1); + + sum += y[i].qs[ j]; + sum += y[i].qs[QK8_1/2 + j]; + } + + y[i].s = GGML_FP32_TO_FP16(sum*d); + } +} + +static inline int best_index_mxfp4(float x, float e) { + int best_index = 0; + float best_err = fabsf(kvalues_mxfp4[0]*e - x); + for (int i = 1; i < 16; i++) { + float err = fabsf(kvalues_mxfp4[i]*e - x); + if (err < best_err) { + best_index = i; + best_err = err; + } + } + return best_index; +} + +void quantize_row_mxfp4_ref(const float * GGML_RESTRICT x, block_mxfp4 * GGML_RESTRICT y, int64_t k) { + static const int qk = QK_MXFP4; + + assert(k % qk == 0); + + const int nb = k / qk; + + for (int i = 0; i < nb; i++) { + float amax = 0.0f; // absolute max + + for (int j = 0; j < qk; j++) { + const float v = x[i*qk + j]; + + if (amax < fabsf(v)) { + amax = fabsf(v); + } + } + + const uint8_t e = amax > 0.0f ? (uint8_t) (floorf(log2f(amax)) - 2 + 127) : 0; + + const float d = GGML_E8M0_TO_FP32_HALF(e); + + y[i].e = e; + + for (int j = 0; j < qk/2; ++j) { + const uint8_t x0 = best_index_mxfp4(x[i*qk + 0 + j], d); + const uint8_t x1 = best_index_mxfp4(x[i*qk + qk/2 + j], d); + + y[i].qs[j] = x0; + y[i].qs[j] |= x1 << 4; + } + } +} + +void quantize_row_nvfp4_ref(const float * GGML_RESTRICT x, block_nvfp4 * GGML_RESTRICT y, int64_t k) { + static const int qk = QK_NVFP4; + static const int qk_sub = QK_NVFP4_SUB; + static const int n_sub = QK_NVFP4 / QK_NVFP4_SUB; + + assert(k % qk == 0); + + const int nb = k / qk; + + for (int i = 0; i < nb; i++) { + for (int s = 0; s < n_sub; s++) { + const float * xb = x + i*qk + s*qk_sub; + + float amax = 0.0f; + for (int j = 0; j < qk_sub; j++) { + if (amax < fabsf(xb[j])) { + amax = fabsf(xb[j]); + } + } + + // UE4M3 scale: amax / 6.0 maps the max E2M1 value (6.0) to amax + const uint8_t ue = ggml_fp32_to_ue4m3(amax / 6.0f); + y[i].d[s] = ue; + const float d = ggml_ue4m3_to_fp32(ue); + + for (int j = 0; j < qk_sub/2; ++j) { + const uint8_t x0 = best_index_mxfp4(xb[0 + j], d); + const uint8_t x1 = best_index_mxfp4(xb[qk_sub/2 + j], d); + + y[i].qs[s*(qk_sub/2) + j] = x0 | (x1 << 4); + } + } + } +} + +void dequantize_row_q1_0(const block_q1_0 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { + static const int qk = QK1_0; + + assert(k % qk == 0); + + const int nb = k / qk; + + for (int i = 0; i < nb; i++) { + const float d = GGML_FP16_TO_FP32(x[i].d); + const float neg_d = -d; + + for (int j = 0; j < qk; ++j) { + const int byte_index = j / 8; + const int bit_offset = j % 8; + const uint8_t bit = (x[i].qs[byte_index] >> bit_offset) & 1; + y[i*qk + j] = bit ? d : neg_d; + } + } +} + +void dequantize_row_q2_0(const block_q2_0 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { + static const int qk = QK2_0; + + assert(k % qk == 0); + + const int nb = k / qk; + + for (int i = 0; i < nb; i++) { + const float d = GGML_FP16_TO_FP32(x[i].d); + + for (int j = 0; j < qk; ++j) { + const int byte_index = j / 4; + const int bit_offset = (j % 4) * 2; + const uint8_t q = (x[i].qs[byte_index] >> bit_offset) & 0x03; + // 00=-1, 01=0, 10=+1, 11=+2 + y[i*qk + j] = ((int)q - 1) * d; + } + } +} + +void dequantize_row_q4_0(const block_q4_0 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { + static const int qk = QK4_0; + + assert(k % qk == 0); + + const int nb = k / qk; + + for (int i = 0; i < nb; i++) { + const float d = GGML_FP16_TO_FP32(x[i].d); + + for (int j = 0; j < qk/2; ++j) { + const int x0 = (x[i].qs[j] & 0x0F) - 8; + const int x1 = (x[i].qs[j] >> 4) - 8; + + y[i*qk + j + 0 ] = x0*d; + y[i*qk + j + qk/2] = x1*d; + } + } +} + +void dequantize_row_q4_1(const block_q4_1 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { + static const int qk = QK4_1; + + assert(k % qk == 0); + + const int nb = k / qk; + + for (int i = 0; i < nb; i++) { + const float d = GGML_FP16_TO_FP32(x[i].d); + const float m = GGML_FP16_TO_FP32(x[i].m); + + for (int j = 0; j < qk/2; ++j) { + const int x0 = (x[i].qs[j] & 0x0F); + const int x1 = (x[i].qs[j] >> 4); + + y[i*qk + j + 0 ] = x0*d + m; + y[i*qk + j + qk/2] = x1*d + m; + } + } +} + +void dequantize_row_q5_0(const block_q5_0 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { + static const int qk = QK5_0; + + assert(k % qk == 0); + + const int nb = k / qk; + + for (int i = 0; i < nb; i++) { + const float d = GGML_FP16_TO_FP32(x[i].d); + + uint32_t qh; + memcpy(&qh, x[i].qh, sizeof(qh)); + + for (int j = 0; j < qk/2; ++j) { + const uint8_t xh_0 = ((qh >> (j + 0)) << 4) & 0x10; + const uint8_t xh_1 = ((qh >> (j + 12)) ) & 0x10; + + const int32_t x0 = ((x[i].qs[j] & 0x0F) | xh_0) - 16; + const int32_t x1 = ((x[i].qs[j] >> 4) | xh_1) - 16; + + y[i*qk + j + 0 ] = x0*d; + y[i*qk + j + qk/2] = x1*d; + } + } +} + +void dequantize_row_q5_1(const block_q5_1 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { + static const int qk = QK5_1; + + assert(k % qk == 0); + + const int nb = k / qk; + + for (int i = 0; i < nb; i++) { + const float d = GGML_FP16_TO_FP32(x[i].d); + const float m = GGML_FP16_TO_FP32(x[i].m); + + uint32_t qh; + memcpy(&qh, x[i].qh, sizeof(qh)); + + for (int j = 0; j < qk/2; ++j) { + const uint8_t xh_0 = ((qh >> (j + 0)) << 4) & 0x10; + const uint8_t xh_1 = ((qh >> (j + 12)) ) & 0x10; + + const int x0 = (x[i].qs[j] & 0x0F) | xh_0; + const int x1 = (x[i].qs[j] >> 4) | xh_1; + + y[i*qk + j + 0 ] = x0*d + m; + y[i*qk + j + qk/2] = x1*d + m; + } + } +} + +void dequantize_row_q8_0(const block_q8_0 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { + static const int qk = QK8_0; + + assert(k % qk == 0); + + const int nb = k / qk; + + for (int i = 0; i < nb; i++) { + const float d = GGML_FP16_TO_FP32(x[i].d); + + for (int j = 0; j < qk; ++j) { + y[i*qk + j] = x[i].qs[j]*d; + } + } +} + +void dequantize_row_mxfp4(const block_mxfp4 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { + static const int qk = QK_MXFP4; + + assert(k % qk == 0); + + const int nb = k / qk; + + for (int i = 0; i < nb; i++) { + const float d = GGML_E8M0_TO_FP32_HALF(x[i].e); + + for (int j = 0; j < qk/2; ++j) { + const int8_t x0 = kvalues_mxfp4[x[i].qs[j] & 0x0F]; + const int8_t x1 = kvalues_mxfp4[x[i].qs[j] >> 4]; + + y[i*qk + j + 0 ] = x0*d; + y[i*qk + j + qk/2] = x1*d; + } + } +} + +void dequantize_row_nvfp4(const block_nvfp4 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { + static const int qk = QK_NVFP4; + static const int qk_sub = QK_NVFP4_SUB; + static const int n_sub = QK_NVFP4 / QK_NVFP4_SUB; + + assert(k % qk == 0); + + const int nb = k / qk; + + for (int i = 0; i < nb; i++) { + for (int s = 0; s < n_sub; s++) { + const float d = ggml_ue4m3_to_fp32(x[i].d[s]); + float * yb = y + i*qk + s*qk_sub; + + for (int j = 0; j < qk_sub/2; ++j) { + const int8_t v0 = kvalues_mxfp4[x[i].qs[s*(qk_sub/2) + j] & 0x0F]; + const int8_t v1 = kvalues_mxfp4[x[i].qs[s*(qk_sub/2) + j] >> 4]; + + yb[j + 0 ] = v0*d; + yb[j + qk_sub/2] = v1*d; + } + } + } +} + +// +// 2-6 bit quantization in super-blocks +// + +// +// ===================== Helper functions +// +static inline int nearest_int(float fval) { + assert(fabsf(fval) <= 4194303.f); + float val = fval + 12582912.f; + int i; memcpy(&i, &val, sizeof(int)); + return (i & 0x007fffff) - 0x00400000; +} + +static float make_qx_quants(int n, int nmax, const float * GGML_RESTRICT x, int8_t * GGML_RESTRICT L, int rmse_type, + const float * GGML_RESTRICT qw) { + float max = 0; + float amax = 0; + for (int i = 0; i < n; ++i) { + float ax = fabsf(x[i]); + if (ax > amax) { amax = ax; max = x[i]; } + } + if (amax < GROUP_MAX_EPS) { // all zero + for (int i = 0; i < n; ++i) { + L[i] = 0; + } + return 0.f; + } + float iscale = -nmax / max; + if (rmse_type == 0) { + for (int i = 0; i < n; ++i) { + int l = nearest_int(iscale * x[i]); + L[i] = nmax + MAX(-nmax, MIN(nmax-1, l)); + } + return 1/iscale; + } + bool return_early = false; + if (rmse_type < 0) { + rmse_type = -rmse_type; + return_early = true; + } + float sumlx = 0; + float suml2 = 0; +#ifdef HAVE_BUGGY_APPLE_LINKER + // use 'volatile' to prevent unroll and work around a bug in Apple ld64 1015.7 + for (volatile int i = 0; i < n; ++i) { +#else + for (int i = 0; i < n; ++i) { +#endif + int l = nearest_int(iscale * x[i]); + l = MAX(-nmax, MIN(nmax-1, l)); + L[i] = l + nmax; + float w = qw ? qw[i] : rmse_type == 1 ? x[i] * x[i] : rmse_type == 2 ? 1 : rmse_type == 3 ? fabsf(x[i]) : sqrtf(fabsf(x[i])); + sumlx += w*x[i]*l; + suml2 += w*l*l; + } + float scale = suml2 ? sumlx/suml2 : 0.0f; + if (return_early) return suml2 > 0 ? 0.5f*(scale + 1/iscale) : 1/iscale; + float best = scale * sumlx; + for (int is = -9; is <= 9; ++is) { + if (is == 0) { + continue; + } + iscale = -(nmax + 0.1f*is) / max; + sumlx = suml2 = 0; + for (int i = 0; i < n; ++i) { + int l = nearest_int(iscale * x[i]); + l = MAX(-nmax, MIN(nmax-1, l)); + float w = qw ? qw[i] : rmse_type == 1 ? x[i] * x[i] : rmse_type == 2 ? 1 : rmse_type == 3 ? fabsf(x[i]) : sqrtf(fabsf(x[i])); + sumlx += w*x[i]*l; + suml2 += w*l*l; + } + if (suml2 > 0 && sumlx*sumlx > best*suml2) { + for (int i = 0; i < n; ++i) { + int l = nearest_int(iscale * x[i]); + L[i] = nmax + MAX(-nmax, MIN(nmax-1, l)); + } + scale = sumlx/suml2; best = scale*sumlx; + } + } + return scale; +} + +static float make_q3_quants(int n, int nmax, const float * GGML_RESTRICT x, int8_t * GGML_RESTRICT L, bool do_rmse) { + float max = 0; + float amax = 0; + for (int i = 0; i < n; ++i) { + float ax = fabsf(x[i]); + if (ax > amax) { amax = ax; max = x[i]; } + } + if (amax < GROUP_MAX_EPS) { // all zero + for (int i = 0; i < n; ++i) { L[i] = 0; } + return 0.f; + } + float iscale = -nmax / max; + if (do_rmse) { + float sumlx = 0; + float suml2 = 0; + for (int i = 0; i < n; ++i) { + int l = nearest_int(iscale * x[i]); + l = MAX(-nmax, MIN(nmax-1, l)); + L[i] = l; + float w = x[i]*x[i]; + sumlx += w*x[i]*l; + suml2 += w*l*l; + } + for (int itry = 0; itry < 5; ++itry) { + int n_changed = 0; + for (int i = 0; i < n; ++i) { + float w = x[i]*x[i]; + float slx = sumlx - w*x[i]*L[i]; + if (slx > 0) { + float sl2 = suml2 - w*L[i]*L[i]; + int new_l = nearest_int(x[i] * sl2 / slx); + new_l = MAX(-nmax, MIN(nmax-1, new_l)); + if (new_l != L[i]) { + slx += w*x[i]*new_l; + sl2 += w*new_l*new_l; + if (sl2 > 0 && slx*slx*suml2 > sumlx*sumlx*sl2) { + L[i] = new_l; sumlx = slx; suml2 = sl2; + ++n_changed; + } + } + } + } + if (!n_changed) { + break; + } + } + for (int i = 0; i < n; ++i) { + L[i] += nmax; + } + return suml2 > 0.0f ? sumlx / suml2 : 0.0f; + } + for (int i = 0; i < n; ++i) { + int l = nearest_int(iscale * x[i]); + l = MAX(-nmax, MIN(nmax-1, l)); + L[i] = l + nmax; + } + return 1/iscale; +} + +static float make_qkx1_quants(int n, int nmax, const float * GGML_RESTRICT x, uint8_t * GGML_RESTRICT L, float * GGML_RESTRICT the_min, + int ntry, float alpha) { + float min = x[0]; + float max = x[0]; + for (int i = 1; i < n; ++i) { + if (x[i] < min) min = x[i]; + if (x[i] > max) max = x[i]; + } + if (max == min) { + for (int i = 0; i < n; ++i) L[i] = 0; + *the_min = 0; + return 0.f; + } + if (min > 0) min = 0; + float iscale = nmax/(max - min); + float scale = 1/iscale; + for (int itry = 0; itry < ntry; ++itry) { + float sumlx = 0; int suml2 = 0; + bool did_change = false; + for (int i = 0; i < n; ++i) { + int l = nearest_int(iscale*(x[i] - min)); + l = MAX(0, MIN(nmax, l)); + if (l != L[i]) { + L[i] = l; + did_change = true; + } + sumlx += (x[i] - min)*l; + suml2 += l*l; + } + scale = sumlx/suml2; + float sum = 0; + for (int i = 0; i < n; ++i) { + sum += x[i] - scale*L[i]; + } + min = alpha*min + (1 - alpha)*sum/n; + if (min > 0) min = 0; + iscale = 1/scale; + if (!did_change) break; + } + *the_min = -min; + return scale; +} + +static float make_qkx2_quants(int n, int nmax, const float * GGML_RESTRICT x, const float * GGML_RESTRICT weights, + uint8_t * GGML_RESTRICT L, float * GGML_RESTRICT the_min, uint8_t * GGML_RESTRICT Laux, + float rmin, float rdelta, int nstep, bool use_mad) { + float min = x[0]; + float max = x[0]; + float sum_w = weights[0]; + float sum_x = sum_w * x[0]; +#ifdef HAVE_BUGGY_APPLE_LINKER + // use 'volatile' to prevent unroll and work around a bug in Apple ld64 1015.7 + for (volatile int i = 1; i < n; ++i) { +#else + for (int i = 1; i < n; ++i) { +#endif + if (x[i] < min) min = x[i]; + if (x[i] > max) max = x[i]; + float w = weights[i]; + sum_w += w; + sum_x += w * x[i]; + } + if (min > 0) min = 0; + if (max == min) { + for (int i = 0; i < n; ++i) L[i] = 0; + *the_min = -min; + return 0.f; + } + float iscale = nmax/(max - min); + float scale = 1/iscale; + float best_error = 0; + for (int i = 0; i < n; ++i) { + int l = nearest_int(iscale*(x[i] - min)); + L[i] = MAX(0, MIN(nmax, l)); + float diff = scale * L[i] + min - x[i]; + diff = use_mad ? fabsf(diff) : diff * diff; + float w = weights[i]; + best_error += w * diff; + } + if (nstep < 1) { + *the_min = -min; + return scale; + } + for (int is = 0; is <= nstep; ++is) { + iscale = (rmin + rdelta*is + nmax)/(max - min); + float sum_l = 0, sum_l2 = 0, sum_xl = 0; + for (int i = 0; i < n; ++i) { + int l = nearest_int(iscale*(x[i] - min)); + l = MAX(0, MIN(nmax, l)); + Laux[i] = l; + float w = weights[i]; + sum_l += w*l; + sum_l2 += w*l*l; + sum_xl += w*l*x[i]; + } + float D = sum_w * sum_l2 - sum_l * sum_l; + if (D > 0) { + float this_scale = (sum_w * sum_xl - sum_x * sum_l)/D; + float this_min = (sum_l2 * sum_x - sum_l * sum_xl)/D; + if (this_min > 0) { + this_min = 0; + this_scale = sum_xl / sum_l2; + } + float cur_error = 0; + for (int i = 0; i < n; ++i) { + float diff = this_scale * Laux[i] + this_min - x[i]; + diff = use_mad ? fabsf(diff) : diff * diff; + float w = weights[i]; + cur_error += w * diff; + } + if (cur_error < best_error) { + for (int i = 0; i < n; ++i) { + L[i] = Laux[i]; + } + best_error = cur_error; + scale = this_scale; + min = this_min; + } + } + } + *the_min = -min; + return scale; +} + +static inline void get_scale_min_k4(int j, const uint8_t * GGML_RESTRICT q, uint8_t * GGML_RESTRICT d, uint8_t * GGML_RESTRICT m) { + if (j < 4) { + *d = q[j] & 63; *m = q[j + 4] & 63; + } else { + *d = (q[j+4] & 0xF) | ((q[j-4] >> 6) << 4); + *m = (q[j+4] >> 4) | ((q[j-0] >> 6) << 4); + } +} + +//========================- 2-bit (de)-quantization + +void quantize_row_q2_K_ref(const float * GGML_RESTRICT x, block_q2_K * GGML_RESTRICT y, int64_t k) { + assert(k % QK_K == 0); + const int nb = k / QK_K; + + uint8_t L[QK_K]; + uint8_t Laux[16]; + float weights[16]; + float mins[QK_K/16]; + float scales[QK_K/16]; + + const float q4scale = 15.f; + + for (int i = 0; i < nb; i++) { + float max_scale = 0; // as we are deducting the min, scales are always positive + float max_min = 0; + for (int j = 0; j < QK_K/16; ++j) { + for (int l = 0; l < 16; ++l) weights[l] = fabsf(x[16*j + l]); + scales[j] = make_qkx2_quants(16, 3, x + 16*j, weights, L + 16*j, &mins[j], Laux, -0.5f, 0.1f, 15, true); + float scale = scales[j]; + if (scale > max_scale) { + max_scale = scale; + } + float min = mins[j]; + if (min > max_min) { + max_min = min; + } + } + + if (max_scale > 0) { + float iscale = q4scale/max_scale; + for (int j = 0; j < QK_K/16; ++j) { + int l = nearest_int(iscale*scales[j]); + y[i].scales[j] = l; + } + y[i].d = GGML_FP32_TO_FP16(max_scale/q4scale); + } else { + for (int j = 0; j < QK_K/16; ++j) y[i].scales[j] = 0; + y[i].d = GGML_FP32_TO_FP16(0.f); + } + if (max_min > 0) { + float iscale = q4scale/max_min; + for (int j = 0; j < QK_K/16; ++j) { + int l = nearest_int(iscale*mins[j]); + y[i].scales[j] |= (l << 4); + } + y[i].dmin = GGML_FP32_TO_FP16(max_min/q4scale); + } else { + y[i].dmin = GGML_FP32_TO_FP16(0.f); + } + for (int j = 0; j < QK_K/16; ++j) { + const float d = GGML_FP16_TO_FP32(y[i].d) * (y[i].scales[j] & 0xF); + if (!d) continue; + const float dm = GGML_FP16_TO_FP32(y[i].dmin) * (y[i].scales[j] >> 4); + for (int ii = 0; ii < 16; ++ii) { + int l = nearest_int((x[16*j + ii] + dm)/d); + l = MAX(0, MIN(3, l)); + L[16*j + ii] = l; + } + } + + for (int j = 0; j < QK_K; j += 128) { + for (int l = 0; l < 32; ++l) { + y[i].qs[j/4 + l] = L[j + l] | (L[j + l + 32] << 2) | (L[j + l + 64] << 4) | (L[j + l + 96] << 6); + } + } + + x += QK_K; + } +} + +void dequantize_row_q2_K(const block_q2_K * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { + assert(k % QK_K == 0); + const int nb = k / QK_K; + + for (int i = 0; i < nb; i++) { + + const float d = GGML_FP16_TO_FP32(x[i].d); + const float min = GGML_FP16_TO_FP32(x[i].dmin); + + const uint8_t * q = x[i].qs; + + int is = 0; + float dl, ml; + for (int n = 0; n < QK_K; n += 128) { + int shift = 0; + for (int j = 0; j < 4; ++j) { + + uint8_t sc = x[i].scales[is++]; + dl = d * (sc & 0xF); ml = min * (sc >> 4); + for (int l = 0; l < 16; ++l) *y++ = dl * ((int8_t)((q[l] >> shift) & 3)) - ml; + + sc = x[i].scales[is++]; + dl = d * (sc & 0xF); ml = min * (sc >> 4); + for (int l = 0; l < 16; ++l) *y++ = dl * ((int8_t)((q[l+16] >> shift) & 3)) - ml; + + shift += 2; + } + q += 32; + } + } +} + +static float make_qkx3_quants(int n, int nmax, const float * GGML_RESTRICT x, const float * GGML_RESTRICT weights, + uint8_t * GGML_RESTRICT L, float * GGML_RESTRICT the_min, uint8_t * GGML_RESTRICT Laux, + float rmin, float rdelta, int nstep, bool use_mad) { + float min = x[0]; + float max = x[0]; + float sum_w = weights ? weights[0] : x[0]*x[0]; + float sum_x = sum_w * x[0]; +#ifdef HAVE_BUGGY_APPLE_LINKER + // use 'volatile' to prevent unroll and work around a bug in Apple ld64 1015.7 + for (volatile int i = 1; i < n; ++i) { +#else + for (int i = 1; i < n; ++i) { +#endif + if (x[i] < min) min = x[i]; + if (x[i] > max) max = x[i]; + float w = weights ? weights[i] : x[i]*x[i]; + sum_w += w; + sum_x += w * x[i]; + } + if (min > 0) { + min = 0; + } + if (max <= min) { + memset(L, 0, n); + *the_min = -min; + return 0.f; + } + float iscale = nmax/(max - min); + float scale = 1/iscale; + float best_mad = 0; + for (int i = 0; i < n; ++i) { + int l = nearest_int(iscale*(x[i] - min)); + L[i] = MAX(0, MIN(nmax, l)); + float diff = scale * L[i] + min - x[i]; + diff = use_mad ? fabsf(diff) : diff*diff; + float w = weights ? weights[i] : x[i]*x[i]; + best_mad += w * diff; + } + if (nstep < 1) { + *the_min = -min; + return scale; + } + for (int is = 0; is <= nstep; ++is) { + iscale = (rmin + rdelta*is + nmax)/(max - min); + float sum_l = 0, sum_l2 = 0, sum_xl = 0; + for (int i = 0; i < n; ++i) { + // min is the best fit so far and can be at or near max, so v can be inf, nan or out of range for nearest_int + const float v = iscale*(x[i] - min); + const int l = v > 0 ? nearest_int(MIN(v, nmax)) : 0; + Laux[i] = l; + float w = weights ? weights[i] : x[i]*x[i]; + sum_l += w*l; + sum_l2 += w*l*l; + sum_xl += w*l*x[i]; + } + float D = sum_w * sum_l2 - sum_l * sum_l; + if (D > 0) { + float this_scale = (sum_w * sum_xl - sum_x * sum_l)/D; + float this_min = (sum_l2 * sum_x - sum_l * sum_xl)/D; + if (this_min > 0) { + this_min = 0; + this_scale = sum_xl / sum_l2; + } + float mad = 0; + for (int i = 0; i < n; ++i) { + float diff = this_scale * Laux[i] + this_min - x[i]; + diff = use_mad ? fabsf(diff) : diff*diff; + float w = weights ? weights[i] : x[i]*x[i]; + mad += w * diff; + } + if (mad < best_mad) { + for (int i = 0; i < n; ++i) { + L[i] = Laux[i]; + } + best_mad = mad; + scale = this_scale; + min = this_min; + } + } + } + *the_min = -min; + return scale; +} + +static float make_qp_quants(int n, int nmax, const float * GGML_RESTRICT x, uint8_t * GGML_RESTRICT L, const float * quant_weights) { + float max = 0; + for (int i = 0; i < n; ++i) { + max = MAX(max, x[i]); + } + if (max < GROUP_MAX_EPS) { // all zero + for (int i = 0; i < n; ++i) { L[i] = 0; } + return 0.f; + } + float iscale = nmax / max; + for (int i = 0; i < n; ++i) { + L[i] = nearest_int(iscale * x[i]); + } + float scale = 1/iscale; + float best_mse = 0; + for (int i = 0; i < n; ++i) { + float diff = x[i] - scale*L[i]; + float w = quant_weights[i]; + best_mse += w*diff*diff; + } + for (int is = -4; is <= 4; ++is) { + if (is == 0) continue; + float iscale_is = (0.1f*is + nmax)/max; + float scale_is = 1/iscale_is; + float mse = 0; + for (int i = 0; i < n; ++i) { + int l = nearest_int(iscale_is*x[i]); + l = MIN(nmax, l); + float diff = x[i] - scale_is*l; + float w = quant_weights[i]; + mse += w*diff*diff; + } + if (mse < best_mse) { + best_mse = mse; + iscale = iscale_is; + } + } + float sumlx = 0; + float suml2 = 0; + for (int i = 0; i < n; ++i) { + int l = nearest_int(iscale * x[i]); + l = MIN(nmax, l); + L[i] = l; + float w = quant_weights[i]; + sumlx += w*x[i]*l; + suml2 += w*l*l; + } + for (int itry = 0; itry < 5; ++itry) { + int n_changed = 0; + for (int i = 0; i < n; ++i) { + float w = quant_weights[i]; + float slx = sumlx - w*x[i]*L[i]; + float sl2 = suml2 - w*L[i]*L[i]; + if (slx > 0 && sl2 > 0) { + int new_l = nearest_int(x[i] * sl2 / slx); + new_l = MIN(nmax, new_l); + if (new_l != L[i]) { + slx += w*x[i]*new_l; + sl2 += w*new_l*new_l; + if (slx*slx*suml2 > sumlx*sumlx*sl2) { + L[i] = new_l; sumlx = slx; suml2 = sl2; + ++n_changed; + } + } + } + } + if (!n_changed) { + break; + } + } + return suml2 > 0.0f ? sumlx / suml2 : 0.0f; +} + +static void quantize_row_q2_K_impl(const float * GGML_RESTRICT x, block_q2_K * GGML_RESTRICT y, int k, const float * GGML_RESTRICT quant_weights) { + GGML_ASSERT(quant_weights); + assert(k % QK_K == 0); + const int nb = k / QK_K; + const bool requantize = true; + + uint8_t L[QK_K]; + uint8_t Laux[16]; + float mins[QK_K/16]; + float scales[QK_K/16]; + float sw[QK_K/16]; + float weight[16]; + uint8_t Ls[QK_K/16], Lm[QK_K/16]; + + for (int i = 0; i < nb; i++) { + memset(sw, 0, QK_K/16*sizeof(float)); + float sumx2 = 0; + for (int j = 0; j < QK_K; ++j) sumx2 += x[j]*x[j]; + float sigma2 = sumx2/QK_K; + for (int j = 0; j < QK_K/16; ++j) { + const float * GGML_RESTRICT qw = quant_weights + QK_K * i + 16*j; + for (int l = 0; l < 16; ++l) weight[l] = qw[l] * sqrtf(sigma2 + x[16*j + l]*x[16*j + l]); + for (int l = 0; l < QK_K/16; ++l) sw[j] += weight[l]; + scales[j] = make_qkx3_quants(16, 3, x + 16*j, weight, L + 16*j, &mins[j], Laux, -0.9f, 0.05f, 36, false); + } + + float dm, mm; + dm = make_qp_quants(QK_K/16, 15, scales, Ls, sw); + mm = make_qp_quants(QK_K/16, 15, mins, Lm, sw); + + y[i].d = GGML_FP32_TO_FP16(dm); + y[i].dmin = GGML_FP32_TO_FP16(mm); + dm = GGML_FP16_TO_FP32(y[i].d); + mm = GGML_FP16_TO_FP32(y[i].dmin); + + for (int j = 0; j < QK_K/16; ++j) { + y[i].scales[j] = Ls[j] | (Lm[j] << 4); + } + + if (requantize) { + for (int j = 0; j < QK_K/16; ++j) { + const float d = dm * (y[i].scales[j] & 0xF); + if (!d) continue; + const float m = mm * (y[i].scales[j] >> 4); + for (int ii = 0; ii < 16; ++ii) { + int l = nearest_int((x[16*j + ii] + m)/d); + l = MAX(0, MIN(3, l)); + L[16*j + ii] = l; + } + } + } + + for (int j = 0; j < QK_K; j += 128) { + for (int l = 0; l < 32; ++l) { + y[i].qs[j/4 + l] = L[j + l] | (L[j + l + 32] << 2) | (L[j + l + 64] << 4) | (L[j + l + 96] << 6); + } + } + + x += QK_K; + } +} + +size_t quantize_q2_K(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) { + size_t row_size = ggml_row_size(GGML_TYPE_Q2_K, n_per_row); + if (!quant_weights) { + quantize_row_q2_K_ref(src, dst, (int64_t)nrow*n_per_row); + } + else { + char * qrow = (char *)dst; + for (int64_t row = 0; row < nrow; ++row) { + quantize_row_q2_K_impl(src, (block_q2_K*)qrow, n_per_row, quant_weights); + src += n_per_row; + qrow += row_size; + } + } + return nrow * row_size; +} + +//========================= 3-bit (de)-quantization + +void quantize_row_q3_K_ref(const float * GGML_RESTRICT x, block_q3_K * GGML_RESTRICT y, int64_t k) { + assert(k % QK_K == 0); + const int nb = k / QK_K; + + int8_t L[QK_K]; + float scales[QK_K / 16]; + + for (int i = 0; i < nb; i++) { + + float max_scale = 0; + float amax = 0; + for (int j = 0; j < QK_K/16; ++j) { + scales[j] = make_q3_quants(16, 4, x + 16*j, L + 16*j, true); + float scale = fabsf(scales[j]); + if (scale > amax) { + amax = scale; max_scale = scales[j]; + } + } + + memset(y[i].scales, 0, 12); + if (max_scale) { + float iscale = -32.f/max_scale; + for (int j = 0; j < QK_K/16; ++j) { + int8_t l = nearest_int(iscale*scales[j]); + l = MAX(-32, MIN(31, l)) + 32; + if (j < 8) { + y[i].scales[j] = l & 0xF; + } else { + y[i].scales[j-8] |= ((l & 0xF) << 4); + } + l >>= 4; + y[i].scales[j%4 + 8] |= (l << (2*(j/4))); + } + y[i].d = GGML_FP32_TO_FP16(1/iscale); + } else { + y[i].d = GGML_FP32_TO_FP16(0.f); + } + + int8_t sc; + for (int j = 0; j < QK_K/16; ++j) { + sc = j < 8 ? y[i].scales[j] & 0xF : y[i].scales[j-8] >> 4; + sc = (sc | (((y[i].scales[8 + j%4] >> (2*(j/4))) & 3) << 4)) - 32; + float d = GGML_FP16_TO_FP32(y[i].d) * sc; + if (!d) { + continue; + } + for (int ii = 0; ii < 16; ++ii) { + int l = nearest_int(x[16*j + ii]/d); + l = MAX(-4, MIN(3, l)); + L[16*j + ii] = l + 4; + } + } + + memset(y[i].hmask, 0, QK_K/8); + // We put the high-bit for the 1st 8 quants into bit 0, the next 8 into bit 1, etc. + int m = 0; + uint8_t hm = 1; + for (int j = 0; j < QK_K; ++j) { + if (L[j] > 3) { + y[i].hmask[m] |= hm; + L[j] -= 4; + } + if (++m == QK_K/8) { + m = 0; hm <<= 1; + } + } + for (int j = 0; j < QK_K; j += 128) { + for (int l = 0; l < 32; ++l) { + y[i].qs[j/4 + l] = L[j + l] | (L[j + l + 32] << 2) | (L[j + l + 64] << 4) | (L[j + l + 96] << 6); + } + } + + x += QK_K; + } +} + +void dequantize_row_q3_K(const block_q3_K * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { + assert(k % QK_K == 0); + const int nb = k / QK_K; + + const uint32_t kmask1 = 0x03030303; + const uint32_t kmask2 = 0x0f0f0f0f; + + uint32_t aux[4]; + const int8_t * scales = (const int8_t*)aux; + + for (int i = 0; i < nb; i++) { + + const float d_all = GGML_FP16_TO_FP32(x[i].d); + + const uint8_t * GGML_RESTRICT q = x[i].qs; + const uint8_t * GGML_RESTRICT hm = x[i].hmask; + uint8_t m = 1; + + memcpy(aux, x[i].scales, 12); + uint32_t tmp = aux[2]; + aux[2] = ((aux[0] >> 4) & kmask2) | (((tmp >> 4) & kmask1) << 4); + aux[3] = ((aux[1] >> 4) & kmask2) | (((tmp >> 6) & kmask1) << 4); + aux[0] = (aux[0] & kmask2) | (((tmp >> 0) & kmask1) << 4); + aux[1] = (aux[1] & kmask2) | (((tmp >> 2) & kmask1) << 4); + + int is = 0; + float dl; + for (int n = 0; n < QK_K; n += 128) { + int shift = 0; + for (int j = 0; j < 4; ++j) { + + dl = d_all * (scales[is++] - 32); + for (int l = 0; l < 16; ++l) { + *y++ = dl * ((int8_t)((q[l+ 0] >> shift) & 3) - ((hm[l+ 0] & m) ? 0 : 4)); + } + + dl = d_all * (scales[is++] - 32); + for (int l = 0; l < 16; ++l) { + *y++ = dl * ((int8_t)((q[l+16] >> shift) & 3) - ((hm[l+16] & m) ? 0 : 4)); + } + + shift += 2; + m <<= 1; + } + q += 32; + } + + } +} + +static void quantize_row_q3_K_impl(const float * GGML_RESTRICT x, block_q3_K * GGML_RESTRICT y, int64_t n_per_row, const float * GGML_RESTRICT quant_weights) { + assert(n_per_row % QK_K == 0); + const int nb = n_per_row / QK_K; + + int8_t L[QK_K]; + float scales[QK_K / 16]; + float weight[16]; + float sw[QK_K / 16]; + int8_t Ls[QK_K / 16]; + + for (int i = 0; i < nb; i++) { + + float sumx2 = 0; + for (int j = 0; j < QK_K; ++j) sumx2 += x[j]*x[j]; + float sigma2 = 2*sumx2/QK_K; + + for (int j = 0; j < QK_K/16; ++j) { + if (quant_weights) { + const float * qw = quant_weights + QK_K * i + 16*j; + for (int l = 0; l < 16; ++l) weight[l] = qw[l] * sqrtf(sigma2 + x[16*j+l]*x[16*j+l]); + } else { + for (int l = 0; l < 16; ++l) weight[l] = x[16*j+l]*x[16*j+l]; + } + float sumw = 0; + for (int l = 0; l < 16; ++l) sumw += weight[l]; + sw[j] = sumw; + + scales[j] = make_qx_quants(16, 4, x + 16*j, L + 16*j, 1, weight); + + } + + memset(y[i].scales, 0, 12); + + float d_block = make_qx_quants(QK_K/16, 32, scales, Ls, 1, sw); + for (int j = 0; j < QK_K/16; ++j) { + int l = Ls[j]; + if (j < 8) { + y[i].scales[j] = l & 0xF; + } else { + y[i].scales[j-8] |= ((l & 0xF) << 4); + } + l >>= 4; + y[i].scales[j%4 + 8] |= (l << (2*(j/4))); + } + y[i].d = GGML_FP32_TO_FP16(d_block); + + int8_t sc; + for (int j = 0; j < QK_K/16; ++j) { + sc = j < 8 ? y[i].scales[j] & 0xF : y[i].scales[j-8] >> 4; + sc = (sc | (((y[i].scales[8 + j%4] >> (2*(j/4))) & 3) << 4)) - 32; + float d = GGML_FP16_TO_FP32(y[i].d) * sc; + if (!d) { + continue; + } + for (int ii = 0; ii < 16; ++ii) { + int l = nearest_int(x[16*j + ii]/d); + l = MAX(-4, MIN(3, l)); + L[16*j + ii] = l + 4; + } + } + + memset(y[i].hmask, 0, QK_K/8); + // We put the high-bit for the 1st 8 quants into bit 0, the next 8 into bit 1, etc. + int m = 0; + uint8_t hm = 1; + for (int j = 0; j < QK_K; ++j) { + if (L[j] > 3) { + y[i].hmask[m] |= hm; + L[j] -= 4; + } + if (++m == QK_K/8) { + m = 0; hm <<= 1; + } + } + for (int j = 0; j < QK_K; j += 128) { + for (int l = 0; l < 32; ++l) { + y[i].qs[j/4 + l] = L[j + l] | (L[j + l + 32] << 2) | (L[j + l + 64] << 4) | (L[j + l + 96] << 6); + } + } + + x += QK_K; + } +} + +size_t quantize_q3_K(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) { + size_t row_size = ggml_row_size(GGML_TYPE_Q3_K, n_per_row); + if (!quant_weights) { + quantize_row_q3_K_ref(src, dst, (int64_t)nrow*n_per_row); + } + else { + char * qrow = (char *)dst; + for (int64_t row = 0; row < nrow; ++row) { + quantize_row_q3_K_impl(src, (block_q3_K*)qrow, n_per_row, quant_weights); + src += n_per_row; + qrow += row_size; + } + } + return nrow * row_size; +} + +// ====================== 4-bit (de)-quantization + +void quantize_row_q4_K_ref(const float * GGML_RESTRICT x, block_q4_K * GGML_RESTRICT y, int64_t k) { + assert(k % QK_K == 0); + const int nb = k / QK_K; + + uint8_t L[QK_K]; + uint8_t Laux[32]; + float weights[32]; + float mins[QK_K/32]; + float scales[QK_K/32]; + + for (int i = 0; i < nb; i++) { + float max_scale = 0; // as we are deducting the min, scales are always positive + float max_min = 0; + for (int j = 0; j < QK_K/32; ++j) { + //scales[j] = make_qkx1_quants(32, 15, x + 32*j, L + 32*j, &mins[j], 9, 0.5f); + float sum_x2 = 0; + for (int l = 0; l < 32; ++l) sum_x2 += x[32*j + l] * x[32*j + l]; + float av_x = sqrtf(sum_x2/32); + for (int l = 0; l < 32; ++l) weights[l] = av_x + fabsf(x[32*j + l]); + scales[j] = make_qkx2_quants(32, 15, x + 32*j, weights, L + 32*j, &mins[j], Laux, -1.f, 0.1f, 20, false); + float scale = scales[j]; + if (scale > max_scale) { + max_scale = scale; + } + float min = mins[j]; + if (min > max_min) { + max_min = min; + } + } + + float inv_scale = max_scale > 0 ? 63.f/max_scale : 0.f; + float inv_min = max_min > 0 ? 63.f/max_min : 0.f; + for (int j = 0; j < QK_K/32; ++j) { + uint8_t ls = nearest_int(inv_scale*scales[j]); + uint8_t lm = nearest_int(inv_min*mins[j]); + ls = MIN(63, ls); + lm = MIN(63, lm); + if (j < 4) { + y[i].scales[j] = ls; + y[i].scales[j+4] = lm; + } else { + y[i].scales[j+4] = (ls & 0xF) | ((lm & 0xF) << 4); + y[i].scales[j-4] |= ((ls >> 4) << 6); + y[i].scales[j-0] |= ((lm >> 4) << 6); + } + } + y[i].d = GGML_FP32_TO_FP16(max_scale/63.f); + y[i].dmin = GGML_FP32_TO_FP16(max_min/63.f); + + uint8_t sc, m; + for (int j = 0; j < QK_K/32; ++j) { + get_scale_min_k4(j, y[i].scales, &sc, &m); + const float d = GGML_FP16_TO_FP32(y[i].d) * sc; + if (!d) continue; + const float dm = GGML_FP16_TO_FP32(y[i].dmin) * m; + for (int ii = 0; ii < 32; ++ii) { + int l = nearest_int((x[32*j + ii] + dm)/d); + l = MAX(0, MIN(15, l)); + L[32*j + ii] = l; + } + } + + uint8_t * q = y[i].qs; + for (int j = 0; j < QK_K; j += 64) { + for (int l = 0; l < 32; ++l) q[l] = L[j + l] | (L[j + l + 32] << 4); + q += 32; + } + + x += QK_K; + } +} + +void dequantize_row_q4_K(const block_q4_K * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { + assert(k % QK_K == 0); + const int nb = k / QK_K; + + for (int i = 0; i < nb; i++) { + const uint8_t * q = x[i].qs; + + const float d = GGML_FP16_TO_FP32(x[i].d); + const float min = GGML_FP16_TO_FP32(x[i].dmin); + + int is = 0; + uint8_t sc, m; + for (int j = 0; j < QK_K; j += 64) { + get_scale_min_k4(is + 0, x[i].scales, &sc, &m); + const float d1 = d * sc; const float m1 = min * m; + get_scale_min_k4(is + 1, x[i].scales, &sc, &m); + const float d2 = d * sc; const float m2 = min * m; + for (int l = 0; l < 32; ++l) *y++ = d1 * (q[l] & 0xF) - m1; + for (int l = 0; l < 32; ++l) *y++ = d2 * (q[l] >> 4) - m2; + q += 32; is += 2; + } + } +} + +static void quantize_row_q4_K_impl(const float * GGML_RESTRICT x, block_q4_K * GGML_RESTRICT y, int64_t n_per_row, const float * quant_weights) { + assert(n_per_row % QK_K == 0); + const int64_t nb = n_per_row / QK_K; + + uint8_t L[QK_K]; + uint8_t Laux[32]; + uint8_t Ls[QK_K/32]; + uint8_t Lm[QK_K/32]; + float weights[32]; + float sw[QK_K/32]; + float mins[QK_K/32]; + float scales[QK_K/32]; + + for (int i = 0; i < nb; i++) { + + float sum_x2 = 0; + for (int l = 0; l < QK_K; ++l) sum_x2 += x[l] * x[l]; + float sigma2 = 2*sum_x2/QK_K; + float av_x = sqrtf(sigma2); + + for (int j = 0; j < QK_K/32; ++j) { + if (quant_weights) { + const float * qw = quant_weights + QK_K*i + 32*j; + for (int l = 0; l < 32; ++l) weights[l] = qw[l] * sqrtf(sigma2 + x[32*j + l]*x[32*j + l]); + } else { + for (int l = 0; l < 32; ++l) weights[l] = av_x + fabsf(x[32*j + l]); + } + float sumw = 0; + for (int l = 0; l < 32; ++l) sumw += weights[l]; + sw[j] = sumw; + scales[j] = make_qkx3_quants(32, 15, x + 32*j, weights, L + 32*j, &mins[j], Laux, -0.9f, 0.05f, 36, false); + } + + float d_block = make_qp_quants(QK_K/32, 63, scales, Ls, sw); + float m_block = make_qp_quants(QK_K/32, 63, mins, Lm, sw); + for (int j = 0; j < QK_K/32; ++j) { + uint8_t ls = Ls[j]; + uint8_t lm = Lm[j]; + if (j < 4) { + y[i].scales[j] = ls; + y[i].scales[j+4] = lm; + } else { + y[i].scales[j+4] = (ls & 0xF) | ((lm & 0xF) << 4); + y[i].scales[j-4] |= ((ls >> 4) << 6); + y[i].scales[j-0] |= ((lm >> 4) << 6); + } + } + y[i].d = GGML_FP32_TO_FP16(d_block); + y[i].dmin = GGML_FP32_TO_FP16(m_block); + + uint8_t sc, m; + for (int j = 0; j < QK_K/32; ++j) { + get_scale_min_k4(j, y[i].scales, &sc, &m); + const float d = GGML_FP16_TO_FP32(y[i].d) * sc; + if (!d) continue; + const float dm = GGML_FP16_TO_FP32(y[i].dmin) * m; + for (int ii = 0; ii < 32; ++ii) { + int l = nearest_int((x[32*j + ii] + dm)/d); + l = MAX(0, MIN(15, l)); + L[32*j + ii] = l; + } + } + uint8_t * q = y[i].qs; + for (int j = 0; j < QK_K; j += 64) { + for (int l = 0; l < 32; ++l) q[l] = L[j + l] | (L[j + l + 32] << 4); + q += 32; + } + + x += QK_K; + + } +} + +size_t quantize_q4_K(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) { + size_t row_size = ggml_row_size(GGML_TYPE_Q4_K, n_per_row); + if (!quant_weights) { + quantize_row_q4_K_ref(src, dst, (int64_t)nrow*n_per_row); + } + else { + char * qrow = (char *)dst; + for (int64_t row = 0; row < nrow; ++row) { + quantize_row_q4_K_impl(src, (block_q4_K*)qrow, n_per_row, quant_weights); + src += n_per_row; + qrow += row_size; + } + } + return nrow * row_size; +} + +// ====================== 5-bit (de)-quantization + +void quantize_row_q5_K_ref(const float * GGML_RESTRICT x, block_q5_K * GGML_RESTRICT y, int64_t k) { + assert(k % QK_K == 0); + const int64_t nb = k / QK_K; + + uint8_t L[QK_K]; + float mins[QK_K/32]; + float scales[QK_K/32]; + float weights[32]; + uint8_t Laux[32]; + + for (int i = 0; i < nb; i++) { + float max_scale = 0; // as we are deducting the min, scales are always positive + float max_min = 0; + for (int j = 0; j < QK_K/32; ++j) { + //scales[j] = make_qkx1_quants(32, 31, x + 32*j, L + 32*j, &mins[j], 9, 0.5f); + float sum_x2 = 0; + for (int l = 0; l < 32; ++l) sum_x2 += x[32*j + l] * x[32*j + l]; + float av_x = sqrtf(sum_x2/32); + for (int l = 0; l < 32; ++l) weights[l] = av_x + fabsf(x[32*j + l]); + scales[j] = make_qkx2_quants(32, 31, x + 32*j, weights, L + 32*j, &mins[j], Laux, -0.5f, 0.1f, 15, false); + float scale = scales[j]; + if (scale > max_scale) { + max_scale = scale; + } + float min = mins[j]; + if (min > max_min) { + max_min = min; + } + } + + float inv_scale = max_scale > 0 ? 63.f/max_scale : 0.f; + float inv_min = max_min > 0 ? 63.f/max_min : 0.f; + for (int j = 0; j < QK_K/32; ++j) { + uint8_t ls = nearest_int(inv_scale*scales[j]); + uint8_t lm = nearest_int(inv_min*mins[j]); + ls = MIN(63, ls); + lm = MIN(63, lm); + if (j < 4) { + y[i].scales[j] = ls; + y[i].scales[j+4] = lm; + } else { + y[i].scales[j+4] = (ls & 0xF) | ((lm & 0xF) << 4); + y[i].scales[j-4] |= ((ls >> 4) << 6); + y[i].scales[j-0] |= ((lm >> 4) << 6); + } + } + y[i].d = GGML_FP32_TO_FP16(max_scale/63.f); + y[i].dmin = GGML_FP32_TO_FP16(max_min/63.f); + + uint8_t sc, m; + for (int j = 0; j < QK_K/32; ++j) { + get_scale_min_k4(j, y[i].scales, &sc, &m); + const float d = GGML_FP16_TO_FP32(y[i].d) * sc; + if (!d) continue; + const float dm = GGML_FP16_TO_FP32(y[i].dmin) * m; + for (int ii = 0; ii < 32; ++ii) { + int l = nearest_int((x[32*j + ii] + dm)/d); + l = MAX(0, MIN(31, l)); + L[32*j + ii] = l; + } + } + + uint8_t * GGML_RESTRICT qh = y[i].qh; + uint8_t * GGML_RESTRICT ql = y[i].qs; + memset(qh, 0, QK_K/8); + + uint8_t m1 = 1, m2 = 2; + for (int n = 0; n < QK_K; n += 64) { + for (int j = 0; j < 32; ++j) { + int l1 = L[n + j]; + if (l1 > 15) { + l1 -= 16; qh[j] |= m1; + } + int l2 = L[n + j + 32]; + if (l2 > 15) { + l2 -= 16; qh[j] |= m2; + } + ql[j] = l1 | (l2 << 4); + } + m1 <<= 2; m2 <<= 2; + ql += 32; + } + + x += QK_K; + } +} + +void dequantize_row_q5_K(const block_q5_K * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { + assert(k % QK_K == 0); + const int64_t nb = k / QK_K; + + for (int i = 0; i < nb; i++) { + const uint8_t * ql = x[i].qs; + const uint8_t * qh = x[i].qh; + + const float d = GGML_FP16_TO_FP32(x[i].d); + const float min = GGML_FP16_TO_FP32(x[i].dmin); + + int is = 0; + uint8_t sc, m; + uint8_t u1 = 1, u2 = 2; + for (int j = 0; j < QK_K; j += 64) { + get_scale_min_k4(is + 0, x[i].scales, &sc, &m); + const float d1 = d * sc; const float m1 = min * m; + get_scale_min_k4(is + 1, x[i].scales, &sc, &m); + const float d2 = d * sc; const float m2 = min * m; + for (int l = 0; l < 32; ++l) *y++ = d1 * ((ql[l] & 0xF) + (qh[l] & u1 ? 16 : 0)) - m1; + for (int l = 0; l < 32; ++l) *y++ = d2 * ((ql[l] >> 4) + (qh[l] & u2 ? 16 : 0)) - m2; + ql += 32; is += 2; + u1 <<= 2; u2 <<= 2; + } + } +} + +static void quantize_row_q5_K_impl(const float * GGML_RESTRICT x, block_q5_K * GGML_RESTRICT y, int64_t n_per_row, const float * quant_weights) { + assert(n_per_row % QK_K == 0); + const int64_t nb = n_per_row / QK_K; + + uint8_t L[QK_K]; + uint8_t Laux[32]; + uint8_t Ls[QK_K/32]; + uint8_t Lm[QK_K/32]; + float mins[QK_K/32]; + float scales[QK_K/32]; + float sw[QK_K/32]; + float weights[32]; + + for (int i = 0; i < nb; i++) { + + float sum_x2 = 0; + for (int l = 0; l < QK_K; ++l) sum_x2 += x[l] * x[l]; + float sigma2 = 2*sum_x2/QK_K; + float av_x = sqrtf(sigma2); + + for (int j = 0; j < QK_K/32; ++j) { + if (quant_weights) { + const float * qw = quant_weights + QK_K*i + 32*j; + for (int l = 0; l < 32; ++l) weights[l] = qw[l] * sqrtf(sigma2 + x[32*j + l]*x[32*j + l]); + } else { + for (int l = 0; l < 32; ++l) weights[l] = av_x + fabsf(x[32*j + l]); + } + float sumw = 0; + for (int l = 0; l < 32; ++l) sumw += weights[l]; + sw[j] = sumw; + + scales[j] = make_qkx3_quants(32, 31, x + 32*j, weights, L + 32*j, &mins[j], Laux, -0.9f, 0.05f, 36, false); + } + + float d_block = make_qp_quants(QK_K/32, 63, scales, Ls, sw); + float m_block = make_qp_quants(QK_K/32, 63, mins, Lm, sw); + + for (int j = 0; j < QK_K/32; ++j) { + uint8_t ls = Ls[j]; + uint8_t lm = Lm[j]; + ls = MIN(63, ls); + lm = MIN(63, lm); + if (j < 4) { + y[i].scales[j] = ls; + y[i].scales[j+4] = lm; + } else { + y[i].scales[j+4] = (ls & 0xF) | ((lm & 0xF) << 4); + y[i].scales[j-4] |= ((ls >> 4) << 6); + y[i].scales[j-0] |= ((lm >> 4) << 6); + } + } + y[i].d = GGML_FP32_TO_FP16(d_block); + y[i].dmin = GGML_FP32_TO_FP16(m_block); + + uint8_t sc, m; + for (int j = 0; j < QK_K/32; ++j) { + get_scale_min_k4(j, y[i].scales, &sc, &m); + const float d = GGML_FP16_TO_FP32(y[i].d) * sc; + if (!d) continue; + const float dm = GGML_FP16_TO_FP32(y[i].dmin) * m; + for (int ii = 0; ii < 32; ++ii) { + int l = nearest_int((x[32*j + ii] + dm)/d); + l = MAX(0, MIN(31, l)); + L[32*j + ii] = l; + } + } + + uint8_t * GGML_RESTRICT qh = y[i].qh; + uint8_t * GGML_RESTRICT ql = y[i].qs; + memset(qh, 0, QK_K/8); + + uint8_t m1 = 1, m2 = 2; + for (int n = 0; n < QK_K; n += 64) { + for (int j = 0; j < 32; ++j) { + int l1 = L[n + j]; + if (l1 > 15) { + l1 -= 16; qh[j] |= m1; + } + int l2 = L[n + j + 32]; + if (l2 > 15) { + l2 -= 16; qh[j] |= m2; + } + ql[j] = l1 | (l2 << 4); + } + m1 <<= 2; m2 <<= 2; + ql += 32; + } + + x += QK_K; + + } +} + +size_t quantize_q5_K(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) { + size_t row_size = ggml_row_size(GGML_TYPE_Q5_K, n_per_row); + if (!quant_weights) { + quantize_row_q5_K_ref(src, dst, (int64_t)nrow*n_per_row); + } + else { + char * qrow = (char *)dst; + for (int64_t row = 0; row < nrow; ++row) { + quantize_row_q5_K_impl(src, (block_q5_K*)qrow, n_per_row, quant_weights); + src += n_per_row; + qrow += row_size; + } + } + return nrow * row_size; +} + +// ====================== 6-bit (de)-quantization + +void quantize_row_q6_K_ref(const float * GGML_RESTRICT x, block_q6_K * GGML_RESTRICT y, int64_t k) { + assert(k % QK_K == 0); + const int64_t nb = k / QK_K; + + int8_t L[QK_K]; + float scales[QK_K/16]; + + for (int i = 0; i < nb; i++) { + + float max_scale = 0; + float max_abs_scale = 0; + + for (int ib = 0; ib < QK_K/16; ++ib) { + + const float scale = make_qx_quants(16, 32, x + 16*ib, L + 16*ib, 1, NULL); + scales[ib] = scale; + + const float abs_scale = fabsf(scale); + if (abs_scale > max_abs_scale) { + max_abs_scale = abs_scale; + max_scale = scale; + } + + } + + if (max_abs_scale < GROUP_MAX_EPS) { + memset(&y[i], 0, sizeof(block_q6_K)); + y[i].d = GGML_FP32_TO_FP16(0.f); + x += QK_K; + continue; + } + + float iscale = -128.f/max_scale; + y[i].d = GGML_FP32_TO_FP16(1/iscale); + for (int ib = 0; ib < QK_K/16; ++ib) { + y[i].scales[ib] = MIN(127, nearest_int(iscale*scales[ib])); + } + + for (int j = 0; j < QK_K/16; ++j) { + float d = GGML_FP16_TO_FP32(y[i].d) * y[i].scales[j]; + if (!d) { + continue; + } + for (int ii = 0; ii < 16; ++ii) { + int l = nearest_int(x[16*j + ii]/d); + l = MAX(-32, MIN(31, l)); + L[16*j + ii] = l + 32; + } + } + + uint8_t * GGML_RESTRICT ql = y[i].ql; + uint8_t * GGML_RESTRICT qh = y[i].qh; + for (int j = 0; j < QK_K; j += 128) { + for (int l = 0; l < 32; ++l) { + const uint8_t q1 = L[j + l + 0] & 0xF; + const uint8_t q2 = L[j + l + 32] & 0xF; + const uint8_t q3 = L[j + l + 64] & 0xF; + const uint8_t q4 = L[j + l + 96] & 0xF; + ql[l+ 0] = q1 | (q3 << 4); + ql[l+32] = q2 | (q4 << 4); + qh[l] = (L[j + l] >> 4) | ((L[j + l + 32] >> 4) << 2) | ((L[j + l + 64] >> 4) << 4) | ((L[j + l + 96] >> 4) << 6); + } + ql += 64; + qh += 32; + } + + x += QK_K; + } +} + +void dequantize_row_q6_K(const block_q6_K * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { + assert(k % QK_K == 0); + const int64_t nb = k / QK_K; + + for (int i = 0; i < nb; i++) { + const float d = GGML_FP16_TO_FP32(x[i].d); + + const uint8_t * GGML_RESTRICT ql = x[i].ql; + const uint8_t * GGML_RESTRICT qh = x[i].qh; + const int8_t * GGML_RESTRICT sc = x[i].scales; + + for (int n = 0; n < QK_K; n += 128) { + for (int l = 0; l < 32; ++l) { + int is = l/16; + const int8_t q1 = (int8_t)((ql[l + 0] & 0xF) | (((qh[l] >> 0) & 3) << 4)) - 32; + const int8_t q2 = (int8_t)((ql[l + 32] & 0xF) | (((qh[l] >> 2) & 3) << 4)) - 32; + const int8_t q3 = (int8_t)((ql[l + 0] >> 4) | (((qh[l] >> 4) & 3) << 4)) - 32; + const int8_t q4 = (int8_t)((ql[l + 32] >> 4) | (((qh[l] >> 6) & 3) << 4)) - 32; + y[l + 0] = d * sc[is + 0] * q1; + y[l + 32] = d * sc[is + 2] * q2; + y[l + 64] = d * sc[is + 4] * q3; + y[l + 96] = d * sc[is + 6] * q4; + } + y += 128; + ql += 64; + qh += 32; + sc += 8; + } + } +} + +static void quantize_row_q6_K_impl(const float * GGML_RESTRICT x, block_q6_K * GGML_RESTRICT y, int64_t n_per_row, const float * quant_weights) { + assert(n_per_row % QK_K == 0); + const int64_t nb = n_per_row / QK_K; + + int8_t L[QK_K]; + float scales[QK_K/16]; + //float weights[16]; + + for (int i = 0; i < nb; i++) { + + //float sum_x2 = 0; + //for (int j = 0; j < QK_K; ++j) sum_x2 += x[j]*x[j]; + //float sigma2 = sum_x2/QK_K; + + float max_scale = 0; + float max_abs_scale = 0; + + for (int ib = 0; ib < QK_K/16; ++ib) { + + float scale; + if (quant_weights) { + const float * qw = quant_weights + QK_K*i + 16*ib; + //for (int j = 0; j < 16; ++j) weights[j] = qw[j] * sqrtf(sigma2 + x[16*ib + j]*x[16*ib + j]); + //scale = make_qx_quants(16, 32, x + 16*ib, L + 16*ib, 1, weights); + scale = make_qx_quants(16, 32, x + 16*ib, L + 16*ib, 1, qw); + } else { + scale = make_qx_quants(16, 32, x + 16*ib, L + 16*ib, 1, NULL); + } + scales[ib] = scale; + + const float abs_scale = fabsf(scale); + if (abs_scale > max_abs_scale) { + max_abs_scale = abs_scale; + max_scale = scale; + } + + } + + if (max_abs_scale < GROUP_MAX_EPS) { + memset(&y[i], 0, sizeof(block_q6_K)); + y[i].d = GGML_FP32_TO_FP16(0.f); + x += QK_K; + continue; + } + + float iscale = -128.f/max_scale; + y[i].d = GGML_FP32_TO_FP16(1/iscale); + for (int ib = 0; ib < QK_K/16; ++ib) { + y[i].scales[ib] = MIN(127, nearest_int(iscale*scales[ib])); + } + + for (int j = 0; j < QK_K/16; ++j) { + float d = GGML_FP16_TO_FP32(y[i].d) * y[i].scales[j]; + if (!d) { + continue; + } + for (int ii = 0; ii < 16; ++ii) { + int l = nearest_int(x[16*j + ii]/d); + l = MAX(-32, MIN(31, l)); + L[16*j + ii] = l + 32; + } + } + + uint8_t * GGML_RESTRICT ql = y[i].ql; + uint8_t * GGML_RESTRICT qh = y[i].qh; + for (int j = 0; j < QK_K; j += 128) { + for (int l = 0; l < 32; ++l) { + const uint8_t q1 = L[j + l + 0] & 0xF; + const uint8_t q2 = L[j + l + 32] & 0xF; + const uint8_t q3 = L[j + l + 64] & 0xF; + const uint8_t q4 = L[j + l + 96] & 0xF; + ql[l+ 0] = q1 | (q3 << 4); + ql[l+32] = q2 | (q4 << 4); + qh[l] = (L[j + l] >> 4) | ((L[j + l + 32] >> 4) << 2) | ((L[j + l + 64] >> 4) << 4) | ((L[j + l + 96] >> 4) << 6); + } + ql += 64; + qh += 32; + } + + x += QK_K; + + } +} + +size_t quantize_q6_K(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) { + size_t row_size = ggml_row_size(GGML_TYPE_Q6_K, n_per_row); + if (!quant_weights) { + quantize_row_q6_K_ref(src, dst, (int64_t)nrow*n_per_row); + } + else { + char * qrow = (char *)dst; + for (int64_t row = 0; row < nrow; ++row) { + quantize_row_q6_K_impl(src, (block_q6_K*)qrow, n_per_row, quant_weights); + src += n_per_row; + qrow += row_size; + } + } + return nrow * row_size; +} + +static void quantize_row_q4_0_impl(const float * GGML_RESTRICT x, block_q4_0 * GGML_RESTRICT y, int64_t n_per_row, const float * quant_weights) { + static_assert(QK4_0 == 32, "QK4_0 must be 32"); + + if (!quant_weights) { + quantize_row_q4_0_ref(x, y, n_per_row); + return; + } + + float weight[QK4_0]; + int8_t L[QK4_0]; + + float sum_x2 = 0; + for (int j = 0; j < n_per_row; ++j) sum_x2 += x[j]*x[j]; + float sigma2 = sum_x2/n_per_row; + + const int64_t nb = n_per_row/QK4_0; + for (int ib = 0; ib < nb; ++ib) { + const float * xb = x + QK4_0 * ib; + const float * qw = quant_weights + QK4_0 * ib; + for (int j = 0; j < QK4_0; ++j) weight[j] = qw[j] * sqrtf(sigma2 + xb[j]*xb[j]); + float d = make_qx_quants(QK4_0, 8, xb, L, 1, weight); + y[ib].d = GGML_FP32_TO_FP16(d); + for (int j = 0; j < 16; ++j) { + y[ib].qs[j] = L[j] | (L[j+16] << 4); + } + } +} + +size_t quantize_q1_0(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) { + if (!quant_weights) { + quantize_row_q1_0_ref(src, dst, (int64_t)nrow*n_per_row); + return nrow * ggml_row_size(GGML_TYPE_Q1_0, n_per_row); + } + size_t row_size = ggml_row_size(GGML_TYPE_Q1_0, n_per_row); + char * qrow = (char *)dst; + for (int64_t row = 0; row < nrow; ++row) { + quantize_row_q1_0_ref(src, (block_q1_0*)qrow, n_per_row); + src += n_per_row; + qrow += row_size; + } + return nrow * row_size; +} + +size_t quantize_q2_0(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) { + if (!quant_weights) { + quantize_row_q2_0_ref(src, dst, (int64_t)nrow*n_per_row); + return nrow * ggml_row_size(GGML_TYPE_Q2_0, n_per_row); + } + size_t row_size = ggml_row_size(GGML_TYPE_Q2_0, n_per_row); + char * qrow = (char *)dst; + for (int64_t row = 0; row < nrow; ++row) { + quantize_row_q2_0_ref(src, (block_q2_0*)qrow, n_per_row); + src += n_per_row; + qrow += row_size; + } + return nrow * row_size; +} + +size_t quantize_q4_0(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) { + if (!quant_weights) { + quantize_row_q4_0_ref(src, dst, (int64_t)nrow*n_per_row); + return nrow * ggml_row_size(GGML_TYPE_Q4_0, n_per_row); + } + size_t row_size = ggml_row_size(GGML_TYPE_Q4_0, n_per_row); + char * qrow = (char *)dst; + for (int64_t row = 0; row < nrow; ++row) { + quantize_row_q4_0_impl(src, (block_q4_0*)qrow, n_per_row, quant_weights); + src += n_per_row; + qrow += row_size; + } + return nrow * row_size; +} + +static void quantize_row_q4_1_impl(const float * GGML_RESTRICT x, block_q4_1 * GGML_RESTRICT y, int64_t n_per_row, const float * quant_weights) { + static_assert(QK4_1 == 32, "QK4_1 must be 32"); + + if (!quant_weights) { + quantize_row_q4_1_ref(x, y, n_per_row); + return; + } + + float weight[QK4_1]; + uint8_t L[QK4_1], Laux[QK4_1]; + + float sum_x2 = 0; + for (int j = 0; j < n_per_row; ++j) sum_x2 += x[j]*x[j]; + float sigma2 = sum_x2/n_per_row; + + const int64_t nb = n_per_row/QK4_1; + for (int ib = 0; ib < nb; ++ib) { + const float * xb = x + QK4_1 * ib; + const float * qw = quant_weights + QK4_1 * ib; + for (int j = 0; j < QK4_1; ++j) weight[j] = qw[j] * sqrtf(sigma2 + xb[j]*xb[j]); + float min; + float d = make_qkx3_quants(QK4_1, 15, xb, weight, L, &min, Laux, -0.9f, 0.05f, 36, false); + y[ib].d = GGML_FP32_TO_FP16(d); + y[ib].m = GGML_FP32_TO_FP16(-min); + for (int j = 0; j < 16; ++j) { + y[ib].qs[j] = L[j] | (L[j+16] << 4); + } + } +} + +size_t quantize_q4_1(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) { + if (!quant_weights) { + quantize_row_q4_1_ref(src, dst, (int64_t)nrow*n_per_row); + return nrow * ggml_row_size(GGML_TYPE_Q4_1, n_per_row); + } + size_t row_size = ggml_row_size(GGML_TYPE_Q4_1, n_per_row); + char * qrow = (char *)dst; + for (int64_t row = 0; row < nrow; ++row) { + quantize_row_q4_1_impl(src, (block_q4_1*)qrow, n_per_row, quant_weights); + src += n_per_row; + qrow += row_size; + } + return nrow * row_size; +} + +static void quantize_row_q5_0_impl(const float * GGML_RESTRICT x, block_q5_0 * GGML_RESTRICT y, int64_t n_per_row, const float * quant_weights) { + static_assert(QK5_0 == 32, "QK5_0 must be 32"); + + if (!quant_weights) { + quantize_row_q5_0_ref(x, y, n_per_row); + return; + } + + float weight[QK5_0]; + int8_t L[QK5_0]; + + float sum_x2 = 0; + for (int j = 0; j < n_per_row; ++j) sum_x2 += x[j]*x[j]; + float sigma2 = sum_x2/n_per_row; + + const int64_t nb = n_per_row/QK5_0; + for (int ib = 0; ib < nb; ++ib) { + const float * xb = x + QK5_0 * ib; + const float * qw = quant_weights + QK5_0 * ib; + for (int j = 0; j < QK5_0; ++j) weight[j] = qw[j] * sqrtf(sigma2 + xb[j]*xb[j]); + float d = make_qx_quants(QK5_0, 16, xb, L, 1, weight); + y[ib].d = GGML_FP32_TO_FP16(d); + + uint32_t qh = 0; + + for (int j = 0; j < 16; ++j) { + const uint8_t xi0 = L[j]; + const uint8_t xi1 = L[j+16]; + y[ib].qs[j] = (xi0 & 0x0F) | ((xi1 & 0x0F) << 4); + + // get the 5-th bit and store it in qh at the right position + qh |= ((xi0 & 0x10u) >> 4) << (j + 0); + qh |= ((xi1 & 0x10u) >> 4) << (j + QK5_0/2); + } + + memcpy(&y[ib].qh, &qh, sizeof(qh)); + } +} + +size_t quantize_q5_0(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) { + if (!quant_weights) { + quantize_row_q5_0_ref(src, dst, (int64_t)nrow*n_per_row); + return nrow * ggml_row_size(GGML_TYPE_Q5_0, n_per_row); + } + size_t row_size = ggml_row_size(GGML_TYPE_Q5_0, n_per_row); + char * qrow = (char *)dst; + for (int64_t row = 0; row < nrow; ++row) { + quantize_row_q5_0_impl(src, (block_q5_0*)qrow, n_per_row, quant_weights); + src += n_per_row; + qrow += row_size; + } + return nrow * row_size; +} + +static void quantize_row_q5_1_impl(const float * GGML_RESTRICT x, block_q5_1 * GGML_RESTRICT y, int64_t n_per_row, const float * quant_weights) { + static_assert(QK5_1 == 32, "QK5_1 must be 32"); + + if (!quant_weights) { + quantize_row_q5_1_ref(x, y, n_per_row); + return; + } + + float weight[QK5_1]; + uint8_t L[QK5_1], Laux[QK5_1]; + + float sum_x2 = 0; + for (int j = 0; j < n_per_row; ++j) sum_x2 += x[j]*x[j]; + float sigma2 = sum_x2/n_per_row; + + const int64_t nb = n_per_row/QK5_1; + for (int ib = 0; ib < nb; ++ib) { + const float * xb = x + QK5_1 * ib; + const float * qw = quant_weights + QK5_1 * ib; + for (int j = 0; j < QK5_1; ++j) weight[j] = qw[j] * sqrtf(sigma2 + xb[j]*xb[j]); + float min; + float d = make_qkx3_quants(QK5_1, 31, xb, weight, L, &min, Laux, -0.9f, 0.05f, 36, false); + y[ib].d = GGML_FP32_TO_FP16(d); + y[ib].m = GGML_FP32_TO_FP16(-min); + + uint32_t qh = 0; + for (int j = 0; j < 16; ++j) { + const uint8_t xi0 = L[j]; + const uint8_t xi1 = L[j+16]; + y[ib].qs[j] = (xi0 & 0x0F) | ((xi1 & 0x0F) << 4); + // get the 5-th bit and store it in qh at the right position + qh |= ((xi0 & 0x10u) >> 4) << (j + 0); + qh |= ((xi1 & 0x10u) >> 4) << (j + QK5_0/2); + } + memcpy(&y[ib].qh, &qh, sizeof(qh)); + } +} + +size_t quantize_q5_1(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) { + if (!quant_weights) { + quantize_row_q5_1_ref(src, dst, (int64_t)nrow*n_per_row); + return nrow * ggml_row_size(GGML_TYPE_Q5_1, n_per_row); + } + size_t row_size = ggml_row_size(GGML_TYPE_Q5_1, n_per_row); + char * qrow = (char *)dst; + for (int64_t row = 0; row < nrow; ++row) { + quantize_row_q5_1_impl(src, (block_q5_1*)qrow, n_per_row, quant_weights); + src += n_per_row; + qrow += row_size; + } + return nrow * row_size; +} + +size_t quantize_q8_0(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) { + (void)quant_weights; // not used + const size_t row_size = ggml_row_size(GGML_TYPE_Q8_0, n_per_row); + quantize_row_q8_0_ref(src, dst, (int64_t)nrow*n_per_row); + return nrow * row_size; +} + +size_t quantize_mxfp4(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) { + GGML_UNUSED(quant_weights); + quantize_row_mxfp4_ref(src, dst, (int64_t)nrow*n_per_row); + return nrow * ggml_row_size(GGML_TYPE_MXFP4, n_per_row); +} + +size_t quantize_nvfp4(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) { + GGML_UNUSED(quant_weights); + quantize_row_nvfp4_ref(src, dst, (int64_t)nrow*n_per_row); + return nrow * ggml_row_size(GGML_TYPE_NVFP4, n_per_row); +} + +// ====================== Ternary (de)-quantization (BitNet b1.58 and TriLMs) + +void quantize_row_tq1_0_ref(const float * GGML_RESTRICT x, block_tq1_0 * GGML_RESTRICT y, int64_t k) { + assert(k % QK_K == 0); + const int64_t nb = k / QK_K; + + for (int64_t i = 0; i < nb; i++) { + float amax = 0.0f; // absolute max + + for (int j = 0; j < QK_K; j++) { + const float v = x[j]; + amax = MAX(amax, fabsf(v)); + } + + const float d = amax; + const float id = d ? 1.0f/d : 0.0f; + + y[i].d = GGML_FP32_TO_FP16(d); + + // 5 elements per byte, along 32 bytes + for (size_t j = 0; j < sizeof(y->qs) - sizeof(y->qs) % 32; j += 32) { + for (size_t m = 0; m < 32; ++m) { + uint8_t q = 0; + for (size_t n = 0; n < 5; ++n) { + int xi = lroundf(x[m + n*32] * id) + 1; // -1, 0, 1 -> 0, 1, 2 + q *= 3; + q += xi; + } + // ceiling division (243 == pow(3, 5)) + q = ((uint16_t)q * 256 + (243 - 1)) / 243; + y[i].qs[j + m] = q; + } + x += 5*32; + } + // along 16 bytes + for (size_t j = sizeof(y->qs) - sizeof(y->qs) % 32; j < sizeof(y->qs); j += 16) { + for (size_t m = 0; m < 16; ++m) { + uint8_t q = 0; + for (size_t n = 0; n < 5; ++n) { + int xi = lroundf(x[m + n*16] * id) + 1; // -1, 0, 1 -> 0, 1, 2 + q *= 3; + q += xi; + } + // ceiling division (243 == pow(3, 5)) + q = ((uint16_t)q * 256 + (243 - 1)) / 243; + y[i].qs[j + m] = q; + } + x += 5*16; + } + // 4 elements per byte + for (size_t j = 0; j < sizeof(y->qh); ++j) { + uint8_t q = 0; + for (size_t m = 0; m < 4; ++m) { + // -1, 0, 1 -> 0, 1, 2 + int xi = lroundf(x[j + m*sizeof(y->qh)] * id) + 1; + q *= 3; + q += xi; + } + // shift the first value to the most significant trit + q *= 3; + // ceiling division (243 == pow(3, 5)) + q = ((uint16_t)q * 256 + (243 - 1)) / 243; + y[i].qh[j] = q; + } + x += 4*sizeof(y->qh); + } +} + +void quantize_row_tq2_0_ref(const float * GGML_RESTRICT x, block_tq2_0 * GGML_RESTRICT y, int64_t k) { + assert(k % QK_K == 0); + const int64_t nb = k / QK_K; + + for (int64_t i = 0; i < nb; i++) { + float amax = 0.0f; // absolute max + + for (int j = 0; j < QK_K; j++) { + const float v = x[j]; + amax = MAX(amax, fabsf(v)); + } + + const float d = amax; + const float id = d ? 1.0f/d : 0.0f; + + y[i].d = GGML_FP32_TO_FP16(d); + + for (size_t j = 0; j < sizeof(y->qs); j += 32) { + for (size_t m = 0; m < 32; ++m) { + uint8_t q = 0; + for (size_t n = 0; n < 4; ++n) { + // -1, 0, 1 -> 0, 1, 2 + int xi = lroundf(x[m + n*32] * id) + 1; + q += (xi & 3) << (2*n); + } + y[i].qs[j + m] = q; + } + x += 4*32; + } + } +} + +size_t quantize_tq1_0(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) { + (void)quant_weights; // not used + const size_t row_size = ggml_row_size(GGML_TYPE_TQ1_0, n_per_row); + quantize_row_tq1_0_ref(src, dst, (int64_t)nrow*n_per_row); + return nrow * row_size; +} + +size_t quantize_tq2_0(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) { + (void)quant_weights; // not used + const size_t row_size = ggml_row_size(GGML_TYPE_TQ2_0, n_per_row); + quantize_row_tq2_0_ref(src, dst, (int64_t)nrow*n_per_row); + return nrow * row_size; +} + +void dequantize_row_tq1_0(const block_tq1_0 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { + assert(k % QK_K == 0); + const int64_t nb = k / QK_K; + + const uint8_t pow3[6] = {1, 3, 9, 27, 81, 243}; + + for (int64_t i = 0; i < nb; ++i) { + + const float d = GGML_FP16_TO_FP32(x[i].d); + + for (size_t j = 0; j < sizeof(x->qs) - sizeof(x->qs) % 32; j += 32) { + for (size_t n = 0; n < 5; ++n) { + for (size_t m = 0; m < 32; ++m) { + uint8_t q = x[i].qs[j + m] * pow3[n]; + int16_t xi = ((uint16_t) q * 3) >> 8; + *y++ = (float) (xi - 1) * d; + } + } + } + for (size_t j = sizeof(x->qs) - sizeof(x->qs) % 32; j < sizeof(x->qs); j += 16) { + for (size_t n = 0; n < 5; ++n) { + for (size_t m = 0; m < 16; ++m) { + uint8_t q = x[i].qs[j + m] * pow3[n]; + int16_t xi = ((uint16_t) q * 3) >> 8; + *y++ = (float) (xi - 1) * d; + } + } + } + + for (size_t n = 0; n < 4; ++n) { + for (size_t j = 0; j < sizeof(x->qh); ++j) { + uint8_t q = x[i].qh[j] * pow3[n]; + int16_t xi = ((uint16_t) q * 3) >> 8; + *y++ = (float) (xi - 1) * d; + } + } + } +} + +void dequantize_row_tq2_0(const block_tq2_0 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { + assert(k % QK_K == 0); + const int64_t nb = k / QK_K; + + for (int64_t i = 0; i < nb; ++i) { + + const float d = GGML_FP16_TO_FP32(x[i].d); + + for (size_t j = 0; j < sizeof(x->qs); j += 32) { + for (size_t l = 0; l < 4; ++l) { + for (size_t m = 0; m < 32; ++m) { + int8_t q = (x[i].qs[j + m] >> (l*2)) & 3; + *y++ = (float) (q - 1) * d; + } + } + } + } +} + +// ====================== "True" 2-bit (de)-quantization + +void dequantize_row_iq2_xxs(const block_iq2_xxs * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { + assert(k % QK_K == 0); + const int64_t nb = k / QK_K; + + uint32_t aux32[2]; + const uint8_t * aux8 = (const uint8_t *)aux32; + + for (int i = 0; i < nb; i++) { + + const float d = GGML_FP16_TO_FP32(x[i].d); + + for (int ib32 = 0; ib32 < QK_K/32; ++ib32) { + memcpy(aux32, x[i].qs + 4*ib32, 2*sizeof(uint32_t)); + const float db = d * (0.5f + (aux32[1] >> 28)) * 0.25f; + for (int l = 0; l < 4; ++l) { + const uint8_t * grid = (const uint8_t *)(iq2xxs_grid + aux8[l]); + const uint8_t signs = ksigns_iq2xs[(aux32[1] >> 7*l) & 127]; + for (int j = 0; j < 8; ++j) { + y[j] = db * grid[j] * (signs & kmask_iq2xs[j] ? -1.f : 1.f); + } + y += 8; + } + } + } +} + +// ====================== 2.3125 bpw (de)-quantization + +void dequantize_row_iq2_xs(const block_iq2_xs * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { + assert(k % QK_K == 0); + const int64_t nb = k / QK_K; + + float db[2]; + + for (int i = 0; i < nb; i++) { + + const float d = GGML_FP16_TO_FP32(x[i].d); + + for (int ib32 = 0; ib32 < QK_K/32; ++ib32) { + db[0] = d * (0.5f + (x[i].scales[ib32] & 0xf)) * 0.25f; + db[1] = d * (0.5f + (x[i].scales[ib32] >> 4)) * 0.25f; + for (int l = 0; l < 4; ++l) { + const uint8_t * grid = (const uint8_t *)(iq2xs_grid + (x[i].qs[4*ib32 + l] & 511)); + const uint8_t signs = ksigns_iq2xs[x[i].qs[4*ib32 + l] >> 9]; + for (int j = 0; j < 8; ++j) { + y[j] = db[l/2] * grid[j] * (signs & kmask_iq2xs[j] ? -1.f : 1.f); + } + y += 8; + } + } + } +} + +// ====================== 2.5625 bpw (de)-quantization + +void dequantize_row_iq2_s(const block_iq2_s * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { + assert(k % QK_K == 0); + const int64_t nb = k / QK_K; + + float db[2]; + + for (int i = 0; i < nb; i++) { + + const float d = GGML_FP16_TO_FP32(x[i].d); + const uint8_t * qs = x[i].qs; + const uint8_t * qh = x[i].qh; + const uint8_t * signs = qs + QK_K/8; + + for (int ib32 = 0; ib32 < QK_K/32; ++ib32) { + db[0] = d * (0.5f + (x[i].scales[ib32] & 0xf)) * 0.25f; + db[1] = d * (0.5f + (x[i].scales[ib32] >> 4)) * 0.25f; + for (int l = 0; l < 4; ++l) { + const float dl = db[l/2]; + const uint8_t * grid = (const uint8_t *)(iq2s_grid + (qs[l] | (qh[ib32] << (8-2*l) & 0x300))); + for (int j = 0; j < 8; ++j) { + y[j] = dl * grid[j] * (signs[l] & kmask_iq2xs[j] ? -1.f : 1.f); + } + y += 8; + } + qs += 4; + signs += 4; + } + } +} + +// ====================== 3.0625 bpw (de)-quantization + +void dequantize_row_iq3_xxs(const block_iq3_xxs * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { + assert(k % QK_K == 0); + const int64_t nb = k / QK_K; + + uint32_t aux32; + + for (int i = 0; i < nb; i++) { + + const float d = GGML_FP16_TO_FP32(x[i].d); + const uint8_t * qs = x[i].qs; + const uint8_t * scales_and_signs = qs + QK_K/4; + + for (int ib32 = 0; ib32 < QK_K/32; ++ib32) { + memcpy(&aux32, scales_and_signs + 4*ib32, sizeof(uint32_t)); + const float db = d * (0.5f + (aux32 >> 28)) * 0.5f; + for (int l = 0; l < 4; ++l) { + const uint8_t signs = ksigns_iq2xs[(aux32 >> 7*l) & 127]; + const uint8_t * grid1 = (const uint8_t *)(iq3xxs_grid + qs[2*l+0]); + const uint8_t * grid2 = (const uint8_t *)(iq3xxs_grid + qs[2*l+1]); + for (int j = 0; j < 4; ++j) { + y[j+0] = db * grid1[j] * (signs & kmask_iq2xs[j+0] ? -1.f : 1.f); + y[j+4] = db * grid2[j] * (signs & kmask_iq2xs[j+4] ? -1.f : 1.f); + } + y += 8; + } + qs += 8; + } + } +} + +// ====================== 3.3125 bpw (de)-quantization + +void dequantize_row_iq3_s(const block_iq3_s * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { + assert(k % QK_K == 0); + const int64_t nb = k / QK_K; + + for (int i = 0; i < nb; i++) { + + const float d = GGML_FP16_TO_FP32(x[i].d); + const uint8_t * qs = x[i].qs; + const uint8_t * qh = x[i].qh; + const uint8_t * signs = x[i].signs; + + for (int ib32 = 0; ib32 < QK_K/32; ib32 += 2) { + const float db1 = d * (1 + 2*(x[i].scales[ib32/2] & 0xf)); + const float db2 = d * (1 + 2*(x[i].scales[ib32/2] >> 4)); + for (int l = 0; l < 4; ++l) { + const uint8_t * grid1 = (const uint8_t *)(iq3s_grid + (qs[2*l+0] | ((qh[0] << (8-2*l)) & 256))); + const uint8_t * grid2 = (const uint8_t *)(iq3s_grid + (qs[2*l+1] | ((qh[0] << (7-2*l)) & 256))); + for (int j = 0; j < 4; ++j) { + y[j+0] = db1 * grid1[j] * (signs[l] & kmask_iq2xs[j+0] ? -1.f : 1.f); + y[j+4] = db1 * grid2[j] * (signs[l] & kmask_iq2xs[j+4] ? -1.f : 1.f); + } + y += 8; + } + qs += 8; + signs += 4; + for (int l = 0; l < 4; ++l) { + const uint8_t * grid1 = (const uint8_t *)(iq3s_grid + (qs[2*l+0] | ((qh[1] << (8-2*l)) & 256))); + const uint8_t * grid2 = (const uint8_t *)(iq3s_grid + (qs[2*l+1] | ((qh[1] << (7-2*l)) & 256))); + for (int j = 0; j < 4; ++j) { + y[j+0] = db2 * grid1[j] * (signs[l] & kmask_iq2xs[j+0] ? -1.f : 1.f); + y[j+4] = db2 * grid2[j] * (signs[l] & kmask_iq2xs[j+4] ? -1.f : 1.f); + } + y += 8; + } + qh += 2; + qs += 8; + signs += 4; + } + } +} + +// ====================== 1.5625 bpw (de)-quantization + +void dequantize_row_iq1_s(const block_iq1_s * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { + assert(k % QK_K == 0); + const int64_t nb = k / QK_K; + + for (int i = 0; i < nb; i++) { + + const float d = GGML_FP16_TO_FP32(x[i].d); + const uint8_t * qs = x[i].qs; + const uint16_t * qh = x[i].qh; + + for (int ib = 0; ib < QK_K/32; ++ib) { + const float dl = d * (2*((qh[ib] >> 12) & 7) + 1); + const float delta = qh[ib] & 0x8000 ? -IQ1S_DELTA : IQ1S_DELTA; + for (int l = 0; l < 4; ++l) { + const int8_t * grid = (const int8_t *)(iq1s_grid + (qs[l] | (((qh[ib] >> 3*l) & 7) << 8))); + for (int j = 0; j < 8; ++j) { + y[j] = dl * (grid[j] + delta); + } + y += 8; + } + qs += 4; + } + } +} + +void dequantize_row_iq1_m(const block_iq1_m * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { + assert(k % QK_K == 0); + const int64_t nb = k / QK_K; + + float delta[4]; + uint16_t idx[4]; + + iq1m_scale_t scale; + + for (int i = 0; i < nb; i++) { + + const uint16_t * sc = (const uint16_t *)x[i].scales; + scale.u16 = (sc[0] >> 12) | ((sc[1] >> 8) & 0x00f0) | ((sc[2] >> 4) & 0x0f00) | (sc[3] & 0xf000); + const float d = GGML_FP16_TO_FP32(scale.f16); + + const uint8_t * qs = x[i].qs; + const uint8_t * qh = x[i].qh; + + for (int ib = 0; ib < QK_K/32; ++ib) { + const float dl1 = d * (2*((sc[ib/2] >> (6*(ib%2)+0)) & 0x7) + 1); + const float dl2 = d * (2*((sc[ib/2] >> (6*(ib%2)+3)) & 0x7) + 1); + + idx[0] = qs[0] | ((qh[0] << 8) & 0x700); + idx[1] = qs[1] | ((qh[0] << 4) & 0x700); + idx[2] = qs[2] | ((qh[1] << 8) & 0x700); + idx[3] = qs[3] | ((qh[1] << 4) & 0x700); + delta[0] = qh[0] & 0x08 ? -IQ1S_DELTA : IQ1S_DELTA; + delta[1] = qh[0] & 0x80 ? -IQ1S_DELTA : IQ1S_DELTA; + delta[2] = qh[1] & 0x08 ? -IQ1S_DELTA : IQ1S_DELTA; + delta[3] = qh[1] & 0x80 ? -IQ1S_DELTA : IQ1S_DELTA; + for (int l = 0; l < 2; ++l) { + const int8_t * grid = (const int8_t *)(iq1s_grid + idx[l]); + for (int j = 0; j < 8; ++j) { + y[j] = dl1 * (grid[j] + delta[l]); + } + y += 8; + } + for (int l = 2; l < 4; ++l) { + const int8_t * grid = (const int8_t *)(iq1s_grid + idx[l]); + for (int j = 0; j < 8; ++j) { + y[j] = dl2 * (grid[j] + delta[l]); + } + y += 8; + } + qs += 4; + qh += 2; + } + } +} + +void dequantize_row_iq4_nl(const block_iq4_nl * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { + assert(k % QK4_NL == 0); + const int64_t nb = k / QK4_NL; + + for (int i = 0; i < nb; i++) { + + const uint8_t * qs = x[i].qs; + + const float d = GGML_FP16_TO_FP32(x[i].d); + for (int j = 0; j < QK4_NL/2; ++j) { + y[j+ 0] = d * kvalues_iq4nl[qs[j] & 0xf]; + y[j+QK4_NL/2] = d * kvalues_iq4nl[qs[j] >> 4]; + } + y += QK4_NL; + qs += QK4_NL/2; + } +} + +void dequantize_row_iq4_xs(const block_iq4_xs * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { + assert(k % QK_K == 0); + const int64_t nb = k / QK_K; + + for (int i = 0; i < nb; i++) { + + const uint8_t * qs = x[i].qs; + + const float d = GGML_FP16_TO_FP32(x[i].d); + + for (int ib = 0; ib < QK_K/32; ++ib) { + const int ls = ((x[i].scales_l[ib/2] >> 4*(ib%2)) & 0xf) | (((x[i].scales_h >> 2*ib) & 3) << 4); + const float dl = d * (ls - 32); + for (int j = 0; j < 16; ++j) { + y[j+ 0] = dl * kvalues_iq4nl[qs[j] & 0xf]; + y[j+16] = dl * kvalues_iq4nl[qs[j] >> 4]; + } + y += 32; + qs += 16; + } + } +} + +//===================================== Q8_K ============================================== + +void quantize_row_q8_K_ref(const float * GGML_RESTRICT x, block_q8_K * GGML_RESTRICT y, int64_t k) { + assert(k % QK_K == 0); + const int64_t nb = k / QK_K; + + for (int i = 0; i < nb; i++) { + + float max = 0; + float amax = 0; + for (int j = 0; j < QK_K; ++j) { + float ax = fabsf(x[j]); + if (ax > amax) { + amax = ax; max = x[j]; + } + } + if (!amax) { + y[i].d = 0; + memset(y[i].qs, 0, QK_K); + x += QK_K; + continue; + } + //const float iscale = -128.f/max; + // We need this change for IQ2_XXS, else the AVX implementation becomes very awkward + const float iscale = -127.f/max; + for (int j = 0; j < QK_K; ++j) { + int v = nearest_int(iscale*x[j]); + y[i].qs[j] = MIN(127, v); + } + for (int j = 0; j < QK_K/16; ++j) { + int sum = 0; + for (int ii = 0; ii < 16; ++ii) { + sum += y[i].qs[j*16 + ii]; + } + y[i].bsums[j] = sum; + } + y[i].d = 1/iscale; + x += QK_K; + } +} + +void dequantize_row_q8_K(const block_q8_K * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { + assert(k % QK_K == 0); + const int64_t nb = k / QK_K; + + for (int i = 0; i < nb; i++) { + for (int j = 0; j < QK_K; ++j) { + *y++ = x[i].d * x[i].qs[j]; + } + } +} + +// ================================ IQ2 quantization ============================================= + +typedef struct { + uint64_t * grid; + int * map; + uint16_t * neighbours; +} iq2_entry_t; + +static iq2_entry_t iq2_data[4] = { + {NULL, NULL, NULL}, + {NULL, NULL, NULL}, + {NULL, NULL, NULL}, + {NULL, NULL, NULL}, +}; + +static inline int iq2_data_index(enum ggml_type type) { + GGML_ASSERT(type == GGML_TYPE_IQ2_XXS || type == GGML_TYPE_IQ2_XS || type == GGML_TYPE_IQ1_S || type == GGML_TYPE_IQ1_M || type == GGML_TYPE_IQ2_S); + return type == GGML_TYPE_IQ2_XXS ? 0 : + type == GGML_TYPE_IQ2_XS ? 1 : + type == GGML_TYPE_IQ1_S || type == GGML_TYPE_IQ1_M ? 2 : 3; +} + +static inline int iq2_grid_size(enum ggml_type type) { + GGML_ASSERT(type == GGML_TYPE_IQ2_XXS || type == GGML_TYPE_IQ2_XS || type == GGML_TYPE_IQ1_S || type == GGML_TYPE_IQ1_M || type == GGML_TYPE_IQ2_S); + return type == GGML_TYPE_IQ2_XXS ? 256 : + type == GGML_TYPE_IQ2_XS ? 512 : + type == GGML_TYPE_IQ1_S || type == GGML_TYPE_IQ1_M ? NGRID_IQ1S : 1024; +} + +static int iq2_compare_func(const void * left, const void * right) { + const int * l = (const int *)left; + const int * r = (const int *)right; + return l[0] < r[0] ? -1 : l[0] > r[0] ? 1 : l[1] < r[1] ? -1 : l[1] > r[1] ? 1 : 0; +} + +void iq2xs_init_impl(enum ggml_type type) { + const int gindex = iq2_data_index(type); + const int grid_size = iq2_grid_size(type); + if (iq2_data[gindex].grid) { + return; + } + static const uint16_t kgrid_2bit_256[256] = { + 0, 2, 5, 8, 10, 17, 20, 32, 34, 40, 42, 65, 68, 80, 88, 97, + 100, 128, 130, 138, 162, 257, 260, 272, 277, 320, 388, 408, 512, 514, 546, 642, + 1025, 1028, 1040, 1057, 1060, 1088, 1090, 1096, 1120, 1153, 1156, 1168, 1188, 1280, 1282, 1288, + 1312, 1350, 1385, 1408, 1425, 1545, 1552, 1600, 1668, 1700, 2048, 2053, 2056, 2068, 2088, 2113, + 2116, 2128, 2130, 2184, 2308, 2368, 2562, 2580, 4097, 4100, 4112, 4129, 4160, 4192, 4228, 4240, + 4245, 4352, 4360, 4384, 4432, 4442, 4480, 4644, 4677, 5120, 5128, 5152, 5157, 5193, 5248, 5400, + 5474, 5632, 5654, 6145, 6148, 6160, 6208, 6273, 6400, 6405, 6560, 6737, 8192, 8194, 8202, 8260, + 8289, 8320, 8322, 8489, 8520, 8704, 8706, 9217, 9220, 9232, 9280, 9302, 9472, 9537, 9572, 9872, + 10248, 10272, 10388, 10820, 16385, 16388, 16400, 16408, 16417, 16420, 16448, 16456, 16470, 16480, 16513, 16516, + 16528, 16640, 16672, 16737, 16768, 16773, 16897, 16912, 16968, 16982, 17000, 17408, 17416, 17440, 17536, 17561, + 17682, 17700, 17920, 18433, 18436, 18448, 18496, 18501, 18688, 18776, 18785, 18818, 19013, 19088, 20480, 20488, + 20497, 20505, 20512, 20608, 20616, 20740, 20802, 20900, 21137, 21648, 21650, 21770, 22017, 22100, 22528, 22545, + 22553, 22628, 22848, 23048, 24580, 24592, 24640, 24680, 24832, 24917, 25112, 25184, 25600, 25605, 25872, 25874, + 25988, 26690, 32768, 32770, 32778, 32833, 32898, 33028, 33048, 33088, 33297, 33793, 33796, 33808, 33813, 33856, + 33888, 34048, 34118, 34196, 34313, 34368, 34400, 34818, 35076, 35345, 36868, 36880, 36900, 36928, 37025, 37142, + 37248, 37445, 37888, 37922, 37956, 38225, 39041, 39200, 40962, 41040, 41093, 41225, 41472, 42008, 43088, 43268, + }; + static const uint16_t kgrid_2bit_512[512] = { + 0, 2, 5, 8, 10, 17, 20, 22, 25, 32, 34, 37, 40, 65, 68, 70, + 73, 80, 82, 85, 88, 97, 100, 128, 130, 133, 136, 145, 148, 153, 160, 257, + 260, 262, 265, 272, 274, 277, 280, 282, 289, 292, 320, 322, 325, 328, 337, 340, + 352, 360, 385, 388, 400, 512, 514, 517, 520, 529, 532, 544, 577, 580, 592, 597, + 640, 650, 1025, 1028, 1030, 1033, 1040, 1042, 1045, 1048, 1057, 1060, 1088, 1090, 1093, 1096, + 1105, 1108, 1110, 1120, 1153, 1156, 1168, 1280, 1282, 1285, 1288, 1297, 1300, 1312, 1345, 1348, + 1360, 1377, 1408, 1537, 1540, 1552, 1574, 1600, 1602, 1668, 2048, 2050, 2053, 2056, 2058, 2065, + 2068, 2080, 2085, 2113, 2116, 2128, 2136, 2176, 2208, 2218, 2305, 2308, 2320, 2368, 2433, 2441, + 2560, 2592, 2600, 2710, 2720, 4097, 4100, 4102, 4105, 4112, 4114, 4117, 4120, 4129, 4132, 4160, + 4162, 4165, 4168, 4177, 4180, 4192, 4202, 4225, 4228, 4240, 4352, 4354, 4357, 4360, 4369, 4372, + 4384, 4417, 4420, 4432, 4480, 4500, 4502, 4609, 4612, 4614, 4624, 4672, 4704, 5120, 5122, 5125, + 5128, 5137, 5140, 5152, 5185, 5188, 5193, 5200, 5220, 5248, 5377, 5380, 5392, 5440, 5632, 5652, + 5705, 6145, 6148, 6160, 6162, 6208, 6228, 6278, 6400, 6405, 6502, 6737, 6825, 8192, 8194, 8197, + 8200, 8202, 8209, 8212, 8224, 8257, 8260, 8272, 8320, 8352, 8449, 8452, 8464, 8512, 8520, 8549, + 8704, 8738, 8832, 8872, 9217, 9220, 9232, 9257, 9280, 9472, 9537, 9554, 9625, 9729, 9754, 9894, + 10240, 10248, 10250, 10272, 10325, 10376, 10402, 10600, 10640, 10760, 10784, 10882, 10888, 10890, 16385, 16388, + 16390, 16393, 16400, 16402, 16405, 16408, 16417, 16420, 16448, 16450, 16453, 16456, 16458, 16465, 16468, 16480, + 16485, 16513, 16516, 16528, 16640, 16642, 16645, 16648, 16657, 16660, 16672, 16705, 16708, 16720, 16768, 16773, + 16802, 16897, 16900, 16912, 16914, 16937, 16960, 17408, 17410, 17413, 17416, 17425, 17428, 17433, 17440, 17473, + 17476, 17488, 17536, 17556, 17665, 17668, 17680, 17700, 17728, 17818, 17920, 17930, 17988, 18000, 18433, 18436, + 18448, 18496, 18501, 18516, 18530, 18688, 18705, 18756, 18768, 18793, 18948, 20480, 20482, 20485, 20488, 20497, + 20500, 20512, 20520, 20545, 20548, 20560, 20608, 20737, 20740, 20752, 20757, 20800, 20802, 20992, 21060, 21162, + 21505, 21508, 21520, 21537, 21568, 21600, 21633, 21665, 21760, 21768, 21888, 21896, 22049, 22120, 22177, 22528, + 22548, 22593, 22608, 22681, 22810, 22848, 22850, 23173, 24577, 24580, 24592, 24640, 24660, 24674, 24710, 24745, + 24832, 25124, 25162, 25234, 25600, 25622, 25872, 25920, 25925, 26020, 26625, 26730, 26917, 27142, 27220, 27234, + 32768, 32770, 32773, 32776, 32785, 32788, 32800, 32810, 32833, 32836, 32848, 32896, 32898, 32936, 32938, 33025, + 33028, 33030, 33040, 33088, 33105, 33113, 33280, 33312, 33408, 33410, 33440, 33448, 33793, 33796, 33808, 33810, + 33813, 33856, 33888, 33929, 34048, 34116, 34213, 34328, 34410, 34816, 34824, 34853, 34906, 34944, 34946, 34984, + 35078, 35362, 35456, 35464, 35478, 35496, 36865, 36868, 36880, 36928, 36950, 36996, 37120, 37154, 37220, 37462, + 37513, 37888, 37893, 37956, 37968, 37976, 38185, 38288, 38290, 38465, 38993, 39078, 39241, 39445, 39520, 40960, + 40962, 40968, 40970, 40992, 41002, 41120, 41297, 41305, 41382, 41472, 41474, 41480, 41514, 41600, 41632, 42048, + 42133, 42597, 42648, 43018, 43040, 43042, 43048, 43168, 43176, 43268, 43396, 43398, 43560, 43562, 43665, 43690, + }; + static const uint16_t kgrid_1bit_2048[NGRID_IQ1S] = { + 0, 2, 5, 8, 10, 17, 21, 32, 34, 40, 42, 69, 81, 84, 86, 101, + 128, 130, 136, 138, 149, 160, 162, 168, 170, 260, 261, 273, 276, 278, 281, 282, + 293, 321, 326, 329, 338, 341, 346, 353, 356, 358, 360, 389, 401, 404, 406, 421, + 512, 514, 520, 522, 533, 544, 546, 552, 554, 581, 593, 601, 612, 617, 640, 642, + 648, 650, 657, 661, 665, 672, 674, 680, 682, 1041, 1044, 1046, 1061, 1089, 1097, 1109, + 1114, 1124, 1125, 1169, 1177, 1189, 1281, 1284, 1285, 1286, 1301, 1304, 1306, 1321, 1344, 1349, + 1354, 1360, 1361, 1364, 1365, 1366, 1369, 1376, 1378, 1381, 1384, 1386, 1409, 1425, 1429, 1432, + 1434, 1441, 1444, 1445, 1446, 1449, 1556, 1561, 1601, 1604, 1616, 1618, 1621, 1624, 1632, 1633, + 1638, 1641, 1669, 1681, 1684, 1689, 2048, 2050, 2056, 2058, 2069, 2080, 2082, 2088, 2090, 2117, + 2129, 2134, 2149, 2176, 2178, 2184, 2186, 2197, 2208, 2210, 2216, 2218, 2309, 2321, 2324, 2329, + 2340, 2341, 2369, 2384, 2385, 2389, 2401, 2404, 2409, 2449, 2452, 2454, 2457, 2469, 2560, 2562, + 2568, 2570, 2581, 2592, 2594, 2600, 2602, 2629, 2641, 2649, 2657, 2661, 2688, 2690, 2693, 2696, + 2698, 2709, 2720, 2722, 2728, 2730, 4112, 4113, 4116, 4121, 4132, 4133, 4161, 4164, 4176, 4181, + 4184, 4193, 4196, 4197, 4201, 4241, 4244, 4246, 4257, 4261, 4353, 4356, 4358, 4361, 4368, 4370, + 4373, 4376, 4385, 4388, 4393, 4421, 4426, 4432, 4433, 4434, 4436, 4437, 4438, 4441, 4448, 4453, + 4484, 4498, 4501, 4513, 4516, 4625, 4628, 4630, 4645, 4672, 4678, 4681, 4690, 4693, 4696, 4698, + 4708, 4710, 4741, 4753, 4756, 4758, 4773, 5121, 5126, 5129, 5140, 5141, 5144, 5145, 5153, 5158, + 5185, 5189, 5190, 5192, 5194, 5201, 5204, 5205, 5206, 5209, 5218, 5221, 5224, 5252, 5257, 5264, + 5268, 5269, 5272, 5273, 5274, 5281, 5284, 5285, 5289, 5378, 5381, 5386, 5393, 5396, 5397, 5398, + 5401, 5408, 5410, 5413, 5416, 5418, 5441, 5444, 5445, 5446, 5457, 5458, 5460, 5461, 5462, 5465, + 5466, 5473, 5476, 5477, 5478, 5481, 5504, 5506, 5508, 5509, 5512, 5514, 5520, 5521, 5524, 5525, + 5526, 5529, 5530, 5536, 5538, 5541, 5633, 5636, 5637, 5638, 5653, 5654, 5656, 5658, 5665, 5670, + 5696, 5698, 5700, 5701, 5704, 5706, 5713, 5717, 5718, 5720, 5721, 5729, 5732, 5733, 5736, 5737, + 5738, 5766, 5770, 5778, 5781, 5796, 5801, 6161, 6166, 6181, 6209, 6212, 6214, 6217, 6224, 6229, + 6232, 6234, 6240, 6241, 6244, 6246, 6249, 6277, 6289, 6292, 6309, 6416, 6418, 6421, 6426, 6433, + 6437, 6466, 6468, 6469, 6472, 6481, 6484, 6485, 6486, 6489, 6490, 6496, 6501, 6506, 6537, 6545, + 6546, 6549, 6552, 6561, 6566, 6569, 6665, 6678, 6692, 6694, 6724, 6726, 6729, 6736, 6738, 6741, + 6744, 6753, 6758, 6761, 6789, 6801, 6806, 6810, 8192, 8194, 8200, 8202, 8213, 8224, 8226, 8229, + 8232, 8234, 8261, 8273, 8281, 8289, 8293, 8320, 8322, 8328, 8330, 8341, 8352, 8354, 8357, 8360, + 8362, 8453, 8465, 8468, 8473, 8485, 8514, 8516, 8521, 8533, 8536, 8538, 8545, 8548, 8549, 8550, + 8581, 8592, 8598, 8601, 8613, 8705, 8712, 8714, 8721, 8725, 8736, 8738, 8744, 8746, 8773, 8785, + 8790, 8793, 8805, 8833, 8840, 8842, 8849, 8853, 8864, 8866, 8872, 8874, 9221, 9236, 9238, 9241, + 9253, 9284, 9285, 9286, 9289, 9298, 9301, 9304, 9306, 9318, 9349, 9361, 9364, 9369, 9377, 9381, + 9481, 9493, 9505, 9513, 9536, 9541, 9544, 9553, 9556, 9557, 9561, 9570, 9573, 9576, 9609, 9616, + 9620, 9621, 9624, 9626, 9633, 9636, 9638, 9641, 9733, 9744, 9746, 9753, 9765, 9793, 9801, 9813, + 9824, 9825, 9833, 9860, 9862, 9872, 9882, 10240, 10242, 10248, 10250, 10261, 10272, 10274, 10280, 10282, + 10309, 10321, 10324, 10341, 10368, 10370, 10376, 10378, 10400, 10402, 10408, 10410, 10505, 10513, 10516, 10521, + 10533, 10566, 10569, 10578, 10581, 10593, 10596, 10598, 10601, 10629, 10640, 10646, 10649, 10660, 10661, 10752, + 10754, 10760, 10762, 10784, 10786, 10792, 10794, 10821, 10833, 10838, 10841, 10853, 10880, 10882, 10888, 10890, + 10901, 10912, 10914, 10920, 10922, 16389, 16401, 16406, 16421, 16457, 16466, 16469, 16472, 16474, 16481, 16484, + 16486, 16532, 16537, 16545, 16550, 16640, 16641, 16644, 16646, 16649, 16658, 16661, 16662, 16664, 16666, 16673, + 16678, 16681, 16709, 16712, 16714, 16721, 16724, 16725, 16726, 16729, 16730, 16741, 16744, 16746, 16769, 16772, + 16774, 16784, 16786, 16789, 16800, 16801, 16802, 16901, 16913, 16916, 16918, 16933, 16961, 16978, 16981, 16986, + 16996, 17001, 17033, 17044, 17061, 17409, 17429, 17433, 17449, 17477, 17480, 17482, 17489, 17492, 17493, 17494, + 17505, 17506, 17509, 17512, 17514, 17537, 17542, 17545, 17552, 17554, 17557, 17568, 17569, 17577, 17665, 17666, + 17669, 17674, 17681, 17684, 17685, 17686, 17689, 17696, 17701, 17706, 17729, 17732, 17733, 17734, 17737, 17744, + 17745, 17748, 17749, 17750, 17752, 17753, 17761, 17764, 17765, 17766, 17769, 17794, 17796, 17797, 17800, 17809, + 17812, 17813, 17814, 17817, 17818, 17829, 17832, 17834, 17921, 17925, 17929, 17940, 17941, 17944, 17946, 17953, + 17956, 17961, 17984, 17986, 17989, 17992, 18000, 18001, 18002, 18005, 18006, 18009, 18018, 18021, 18024, 18049, + 18053, 18058, 18068, 18069, 18081, 18084, 18086, 18437, 18449, 18453, 18458, 18469, 18498, 18505, 18512, 18517, + 18520, 18529, 18532, 18534, 18537, 18565, 18577, 18580, 18582, 18585, 18597, 18689, 18693, 18694, 18698, 18704, + 18708, 18709, 18712, 18721, 18724, 18726, 18752, 18757, 18762, 18769, 18770, 18772, 18773, 18774, 18777, 18784, + 18786, 18789, 18790, 18794, 18822, 18825, 18834, 18837, 18838, 18840, 18849, 18852, 18854, 18857, 18966, 19012, + 19014, 19017, 19029, 19032, 19034, 19044, 19049, 19092, 19109, 20481, 20484, 20485, 20486, 20489, 20498, 20501, + 20506, 20513, 20516, 20521, 20544, 20549, 20552, 20561, 20564, 20565, 20566, 20569, 20581, 20584, 20614, 20617, + 20629, 20632, 20640, 20641, 20646, 20649, 20741, 20744, 20745, 20746, 20753, 20756, 20757, 20758, 20760, 20761, + 20768, 20773, 20774, 20776, 20778, 20801, 20804, 20805, 20806, 20809, 20816, 20817, 20818, 20820, 20821, 20822, + 20824, 20825, 20826, 20833, 20836, 20837, 20838, 20841, 20866, 20869, 20881, 20884, 20885, 20886, 20889, 20896, + 20901, 20906, 20993, 20998, 21010, 21013, 21018, 21025, 21028, 21058, 21061, 21066, 21073, 21076, 21077, 21078, + 21081, 21090, 21093, 21125, 21136, 21138, 21141, 21145, 21146, 21156, 21508, 21509, 21521, 21524, 21525, 21526, + 21528, 21529, 21537, 21541, 21544, 21546, 21569, 21572, 21573, 21574, 21577, 21578, 21584, 21585, 21588, 21589, + 21590, 21592, 21593, 21594, 21601, 21602, 21604, 21605, 21606, 21609, 21632, 21640, 21642, 21649, 21652, 21653, + 21654, 21657, 21665, 21668, 21669, 21674, 21761, 21762, 21764, 21765, 21766, 21769, 21776, 21777, 21778, 21780, + 21781, 21782, 21785, 21786, 21793, 21796, 21797, 21798, 21801, 21824, 21825, 21826, 21828, 21829, 21830, 21832, + 21833, 21840, 21841, 21842, 21844, 21845, 21846, 21848, 21849, 21850, 21856, 21857, 21860, 21861, 21862, 21864, + 21865, 21866, 21889, 21892, 21893, 21897, 21898, 21904, 21905, 21908, 21909, 21910, 21912, 21913, 21921, 21924, + 21925, 21926, 21929, 22016, 22017, 22018, 22020, 22022, 22024, 22025, 22033, 22036, 22037, 22040, 22041, 22048, + 22049, 22050, 22052, 22053, 22054, 22056, 22057, 22081, 22085, 22086, 22088, 22089, 22090, 22096, 22097, 22098, + 22100, 22101, 22102, 22104, 22105, 22106, 22113, 22116, 22117, 22121, 22146, 22149, 22150, 22152, 22153, 22154, + 22161, 22165, 22170, 22178, 22181, 22182, 22184, 22185, 22532, 22533, 22534, 22537, 22544, 22549, 22552, 22561, + 22570, 22597, 22600, 22602, 22609, 22612, 22613, 22614, 22616, 22617, 22624, 22626, 22628, 22629, 22658, 22665, + 22672, 22674, 22677, 22680, 22689, 22697, 22785, 22786, 22789, 22794, 22801, 22804, 22805, 22806, 22809, 22821, + 22849, 22852, 22853, 22854, 22857, 22864, 22865, 22866, 22868, 22869, 22870, 22872, 22873, 22874, 22881, 22884, + 22885, 22886, 22889, 22913, 22917, 22921, 22929, 22932, 22933, 22934, 22936, 22937, 22949, 23044, 23048, 23061, + 23066, 23072, 23077, 23078, 23081, 23109, 23112, 23113, 23121, 23125, 23126, 23128, 23129, 23138, 23141, 23144, + 23146, 23169, 23178, 23186, 23189, 23190, 23192, 23194, 23201, 24581, 24596, 24598, 24601, 24613, 24644, 24656, + 24661, 24662, 24664, 24666, 24673, 24676, 24678, 24681, 24705, 24726, 24741, 24833, 24836, 24838, 24841, 24850, + 24853, 24865, 24866, 24870, 24873, 24901, 24905, 24913, 24917, 24918, 24921, 24933, 24934, 24938, 24964, 24970, + 24978, 24981, 24993, 24998, 25001, 25105, 25110, 25113, 25152, 25153, 25158, 25173, 25174, 25176, 25184, 25221, + 25233, 25238, 25253, 25617, 25618, 25621, 25622, 25626, 25633, 25638, 25641, 25664, 25666, 25669, 25672, 25674, + 25681, 25684, 25685, 25686, 25689, 25690, 25696, 25698, 25701, 25732, 25733, 25737, 25744, 25746, 25748, 25749, + 25750, 25752, 25754, 25761, 25764, 25769, 25861, 25864, 25866, 25873, 25877, 25878, 25881, 25924, 25925, 25926, + 25929, 25936, 25937, 25940, 25941, 25942, 25945, 25953, 25956, 25957, 25958, 25961, 25990, 25993, 25994, 26001, + 26005, 26006, 26009, 26010, 26018, 26021, 26022, 26024, 26114, 26121, 26133, 26144, 26150, 26152, 26153, 26176, + 26181, 26184, 26186, 26193, 26196, 26197, 26198, 26200, 26202, 26208, 26213, 26216, 26240, 26242, 26245, 26250, + 26260, 26262, 26264, 26265, 26272, 26276, 26278, 26282, 26646, 26649, 26661, 26689, 26706, 26709, 26714, 26721, + 26729, 26757, 26769, 26776, 26790, 26881, 26884, 26896, 26901, 26913, 26916, 26918, 26921, 26944, 26945, 26949, + 26950, 26952, 26961, 26964, 26965, 26966, 26969, 26976, 26981, 26986, 27010, 27012, 27018, 27029, 27041, 27044, + 27045, 27049, 27153, 27158, 27160, 27201, 27204, 27209, 27216, 27221, 27224, 27226, 27236, 27237, 27241, 27270, + 27284, 27288, 27290, 27302, 32768, 32770, 32776, 32778, 32800, 32802, 32808, 32810, 32837, 32848, 32849, 32852, + 32854, 32857, 32869, 32896, 32898, 32904, 32906, 32917, 32928, 32930, 32936, 32938, 33029, 33041, 33044, 33046, + 33049, 33061, 33089, 33092, 33097, 33104, 33106, 33109, 33110, 33112, 33113, 33124, 33126, 33129, 33157, 33161, + 33172, 33174, 33177, 33189, 33280, 33282, 33288, 33290, 33301, 33312, 33314, 33320, 33322, 33361, 33364, 33369, + 33381, 33408, 33410, 33416, 33418, 33429, 33440, 33442, 33448, 33450, 33812, 33817, 33857, 33860, 33873, 33877, + 33882, 33889, 33892, 33897, 33940, 33945, 34049, 34057, 34066, 34069, 34074, 34086, 34089, 34112, 34113, 34117, + 34120, 34129, 34132, 34133, 34134, 34137, 34138, 34149, 34150, 34152, 34154, 34177, 34180, 34182, 34185, 34192, + 34194, 34197, 34200, 34214, 34321, 34326, 34329, 34341, 34369, 34372, 34377, 34378, 34384, 34389, 34393, 34394, + 34401, 34406, 34410, 34437, 34449, 34458, 34468, 34816, 34818, 34824, 34826, 34837, 34848, 34850, 34856, 34858, + 34881, 34885, 34897, 34900, 34905, 34917, 34921, 34944, 34946, 34952, 34954, 34965, 34976, 34978, 34984, 34986, + 35077, 35078, 35089, 35092, 35094, 35109, 35137, 35140, 35142, 35145, 35152, 35154, 35157, 35162, 35169, 35172, + 35205, 35222, 35225, 35237, 35328, 35330, 35336, 35338, 35349, 35360, 35362, 35368, 35370, 35397, 35409, 35412, + 35414, 35456, 35458, 35464, 35466, 35477, 35488, 35490, 35496, 35498, 36869, 36881, 36886, 36888, 36889, 36901, + 36929, 36934, 36937, 36949, 36952, 36954, 36969, 36970, 36997, 37009, 37012, 37014, 37017, 37029, 37121, 37124, + 37126, 37129, 37136, 37141, 37144, 37146, 37153, 37156, 37158, 37161, 37184, 37189, 37200, 37201, 37204, 37205, + 37206, 37209, 37218, 37221, 37252, 37254, 37266, 37269, 37272, 37281, 37284, 37286, 37289, 37381, 37393, 37396, + 37401, 37413, 37444, 37446, 37449, 37456, 37458, 37461, 37464, 37478, 37481, 37509, 37524, 37526, 37545, 37889, + 37892, 37894, 37904, 37909, 37912, 37926, 37952, 37962, 37969, 37972, 37973, 37974, 37976, 37977, 37984, 37985, + 37986, 37989, 38020, 38022, 38034, 38036, 38037, 38040, 38049, 38057, 38144, 38149, 38152, 38154, 38160, 38161, + 38164, 38165, 38166, 38169, 38177, 38181, 38185, 38186, 38209, 38212, 38213, 38214, 38217, 38224, 38225, 38226, + 38228, 38229, 38230, 38232, 38233, 38234, 38241, 38244, 38245, 38246, 38249, 38273, 38277, 38280, 38289, 38290, + 38292, 38293, 38294, 38297, 38298, 38304, 38306, 38309, 38312, 38314, 38401, 38404, 38416, 38421, 38425, 38432, + 38438, 38441, 38469, 38472, 38473, 38481, 38482, 38485, 38486, 38489, 38501, 38504, 38530, 38532, 38537, 38538, + 38546, 38548, 38549, 38564, 38566, 38569, 38917, 38934, 38937, 38949, 38977, 38982, 38992, 38994, 38997, 38998, + 39002, 39012, 39013, 39045, 39057, 39062, 39065, 39077, 39172, 39174, 39177, 39184, 39186, 39189, 39192, 39194, + 39200, 39201, 39204, 39206, 39232, 39234, 39237, 39240, 39242, 39249, 39252, 39253, 39254, 39257, 39266, 39269, + 39270, 39274, 39297, 39300, 39312, 39314, 39317, 39322, 39329, 39334, 39429, 39445, 39461, 39492, 39494, 39497, + 39504, 39509, 39512, 39521, 39557, 39569, 39572, 39573, 39574, 40960, 40962, 40968, 40970, 40981, 40992, 40994, + 41000, 41002, 41029, 41041, 41044, 41046, 41049, 41088, 41090, 41096, 41098, 41109, 41120, 41122, 41128, 41130, + 41221, 41225, 41233, 41236, 41238, 41241, 41242, 41286, 41289, 41297, 41301, 41304, 41306, 41313, 41316, 41349, + 41360, 41362, 41366, 41369, 41474, 41480, 41482, 41488, 41497, 41506, 41512, 41514, 41541, 41553, 41558, 41561, + 41573, 41600, 41602, 41608, 41610, 41621, 41632, 41634, 41640, 41642, 42009, 42021, 42049, 42052, 42064, 42068, + 42069, 42072, 42074, 42081, 42085, 42086, 42088, 42089, 42117, 42246, 42249, 42256, 42258, 42261, 42264, 42278, + 42281, 42306, 42309, 42321, 42324, 42325, 42326, 42329, 42341, 42346, 42369, 42372, 42373, 42374, 42377, 42386, + 42389, 42392, 42501, 42513, 42518, 42522, 42529, 42533, 42564, 42566, 42570, 42578, 42581, 42582, 42584, 42592, + 42594, 42630, 42640, 42645, 42646, 42649, 42657, 42660, 42662, 43008, 43010, 43016, 43018, 43040, 43042, 43048, + 43050, 43089, 43092, 43094, 43097, 43136, 43138, 43144, 43146, 43157, 43168, 43170, 43176, 43178, 43269, 43284, + 43289, 43297, 43301, 43329, 43344, 43349, 43354, 43361, 43366, 43369, 43408, 43414, 43520, 43522, 43528, 43530, + 43552, 43554, 43560, 43562, 43601, 43604, 43606, 43648, 43650, 43656, 43658, 43669, 43680, 43682, 43688, 43690, + }; + static const uint16_t kgrid_2bit_1024[1024] = { + 0, 2, 5, 8, 10, 17, 20, 22, 25, 32, 34, 37, 40, 65, 68, 70, + 73, 80, 82, 85, 88, 97, 100, 102, 105, 128, 130, 133, 136, 145, 148, 160, + 165, 170, 257, 260, 262, 265, 272, 274, 277, 280, 289, 292, 320, 322, 325, 328, + 337, 340, 342, 345, 352, 357, 360, 385, 388, 400, 402, 405, 417, 420, 512, 514, + 517, 520, 529, 532, 544, 554, 577, 580, 582, 585, 592, 597, 640, 645, 650, 660, + 674, 1025, 1028, 1030, 1033, 1040, 1042, 1045, 1048, 1057, 1060, 1062, 1065, 1088, 1090, 1093, + 1096, 1098, 1105, 1108, 1110, 1113, 1120, 1122, 1125, 1153, 1156, 1158, 1161, 1168, 1173, 1176, + 1185, 1188, 1280, 1282, 1285, 1288, 1290, 1297, 1300, 1302, 1305, 1312, 1317, 1320, 1345, 1348, + 1350, 1353, 1360, 1362, 1365, 1368, 1377, 1380, 1408, 1410, 1413, 1416, 1425, 1428, 1440, 1537, + 1540, 1542, 1545, 1552, 1557, 1600, 1605, 1608, 1617, 1620, 1632, 1665, 1668, 1680, 2048, 2050, + 2053, 2056, 2065, 2068, 2070, 2073, 2080, 2085, 2090, 2113, 2116, 2118, 2121, 2128, 2130, 2133, + 2136, 2145, 2148, 2176, 2181, 2196, 2218, 2305, 2308, 2320, 2322, 2325, 2328, 2337, 2368, 2373, + 2376, 2385, 2388, 2400, 2433, 2448, 2560, 2577, 2580, 2594, 2600, 2602, 2640, 2713, 4097, 4100, + 4102, 4105, 4112, 4114, 4117, 4120, 4129, 4132, 4134, 4160, 4162, 4165, 4168, 4177, 4180, 4182, + 4185, 4192, 4194, 4197, 4200, 4225, 4228, 4230, 4240, 4245, 4248, 4257, 4260, 4352, 4354, 4357, + 4360, 4362, 4369, 4372, 4374, 4377, 4384, 4386, 4389, 4392, 4417, 4420, 4422, 4425, 4432, 4434, + 4437, 4440, 4449, 4452, 4480, 4482, 4485, 4488, 4497, 4500, 4609, 4612, 4617, 4624, 4629, 4641, + 4644, 4672, 4677, 4689, 4692, 4737, 4740, 4752, 5120, 5122, 5125, 5128, 5137, 5140, 5142, 5145, + 5152, 5157, 5160, 5185, 5188, 5190, 5193, 5200, 5202, 5205, 5208, 5217, 5220, 5248, 5250, 5253, + 5256, 5265, 5268, 5280, 5377, 5380, 5382, 5385, 5392, 5394, 5397, 5400, 5409, 5412, 5440, 5442, + 5445, 5448, 5457, 5460, 5472, 5505, 5508, 5520, 5632, 5637, 5640, 5649, 5652, 5664, 5697, 5700, + 5712, 5760, 5802, 6145, 6148, 6150, 6153, 6160, 6165, 6168, 6177, 6208, 6210, 6213, 6216, 6225, + 6228, 6240, 6273, 6276, 6400, 6402, 6405, 6408, 6417, 6420, 6432, 6465, 6468, 6480, 6505, 6562, + 6660, 6672, 6720, 6742, 8192, 8194, 8197, 8200, 8209, 8212, 8214, 8217, 8224, 8229, 8234, 8257, + 8260, 8272, 8274, 8277, 8292, 8320, 8330, 8340, 8362, 8449, 8452, 8464, 8466, 8469, 8481, 8512, + 8514, 8517, 8529, 8532, 8544, 8577, 8580, 8592, 8704, 8714, 8738, 8744, 8746, 8772, 8784, 8840, + 8842, 8872, 9217, 9220, 9222, 9225, 9232, 9237, 9240, 9249, 9252, 9280, 9282, 9285, 9288, 9297, + 9300, 9312, 9345, 9348, 9360, 9472, 9477, 9480, 9489, 9492, 9504, 9537, 9540, 9552, 9574, 9600, + 9729, 9732, 9744, 9792, 9817, 10240, 10245, 10257, 10260, 10305, 10308, 10320, 10378, 10410, 10497, 10500, + 10512, 10645, 10762, 10786, 10852, 10888, 10890, 16385, 16388, 16390, 16393, 16400, 16402, 16405, 16408, 16410, + 16417, 16420, 16422, 16448, 16450, 16453, 16456, 16458, 16465, 16468, 16470, 16473, 16480, 16482, 16485, 16513, + 16516, 16528, 16533, 16536, 16545, 16548, 16640, 16642, 16645, 16648, 16657, 16660, 16662, 16665, 16672, 16674, + 16677, 16705, 16708, 16710, 16713, 16720, 16722, 16725, 16728, 16737, 16740, 16768, 16770, 16773, 16776, 16785, + 16788, 16800, 16897, 16900, 16912, 16914, 16917, 16920, 16932, 16960, 16965, 16968, 16977, 16980, 16992, 17025, + 17028, 17408, 17410, 17413, 17416, 17418, 17425, 17428, 17430, 17433, 17440, 17442, 17445, 17448, 17473, 17476, + 17478, 17481, 17488, 17490, 17493, 17496, 17505, 17508, 17536, 17538, 17541, 17544, 17553, 17556, 17568, 17665, + 17668, 17670, 17673, 17680, 17682, 17685, 17688, 17697, 17700, 17728, 17730, 17733, 17736, 17745, 17748, 17760, + 17770, 17793, 17796, 17808, 17920, 17922, 17925, 17928, 17937, 17940, 17952, 17985, 17988, 18000, 18048, 18085, + 18433, 18436, 18441, 18448, 18450, 18453, 18456, 18465, 18468, 18496, 18498, 18501, 18504, 18513, 18516, 18528, + 18564, 18576, 18688, 18690, 18693, 18696, 18705, 18708, 18720, 18753, 18756, 18768, 18816, 18838, 18945, 18948, + 18960, 19008, 20480, 20482, 20485, 20488, 20497, 20500, 20502, 20505, 20512, 20514, 20517, 20520, 20545, 20548, + 20550, 20553, 20560, 20562, 20565, 20568, 20577, 20580, 20608, 20610, 20613, 20616, 20625, 20628, 20737, 20740, + 20742, 20745, 20752, 20754, 20757, 20760, 20769, 20772, 20800, 20802, 20805, 20808, 20817, 20820, 20832, 20865, + 20868, 20880, 20992, 20997, 21000, 21009, 21012, 21024, 21057, 21060, 21072, 21097, 21120, 21505, 21508, 21510, + 21513, 21520, 21522, 21525, 21528, 21537, 21540, 21568, 21570, 21573, 21576, 21585, 21588, 21600, 21633, 21636, + 21648, 21760, 21762, 21765, 21768, 21777, 21780, 21792, 21825, 21828, 21840, 21888, 22017, 22020, 22032, 22054, + 22080, 22528, 22530, 22533, 22536, 22545, 22548, 22560, 22593, 22596, 22608, 22618, 22656, 22785, 22788, 22800, + 22848, 23040, 23065, 23173, 23208, 24577, 24580, 24582, 24592, 24594, 24597, 24600, 24609, 24612, 24640, 24645, + 24648, 24657, 24660, 24672, 24708, 24720, 24832, 24834, 24837, 24840, 24849, 24852, 24864, 24897, 24900, 24912, + 24960, 24985, 25092, 25104, 25152, 25174, 25249, 25600, 25605, 25608, 25617, 25620, 25632, 25665, 25668, 25680, + 25728, 25857, 25860, 25872, 25920, 25930, 25960, 26002, 26112, 26260, 26625, 26628, 26640, 26725, 26776, 26880, + 26922, 27202, 27297, 32768, 32770, 32773, 32776, 32785, 32788, 32793, 32800, 32805, 32833, 32836, 32848, 32850, + 32853, 32856, 32865, 32896, 32901, 32913, 32916, 33025, 33028, 33033, 33040, 33042, 33045, 33048, 33057, 33060, + 33088, 33090, 33093, 33096, 33105, 33108, 33153, 33156, 33168, 33193, 33280, 33285, 33290, 33297, 33300, 33345, + 33348, 33360, 33793, 33796, 33798, 33801, 33808, 33810, 33813, 33816, 33825, 33856, 33858, 33861, 33864, 33873, + 33876, 33888, 33921, 33924, 33936, 34048, 34050, 34053, 34056, 34065, 34068, 34080, 34113, 34116, 34128, 34176, + 34186, 34305, 34308, 34320, 34345, 34368, 34816, 34821, 34833, 34836, 34881, 34884, 34896, 34978, 35073, 35076, + 35136, 35173, 35362, 35416, 35418, 35458, 35490, 36865, 36868, 36873, 36880, 36882, 36885, 36888, 36900, 36928, + 36930, 36933, 36936, 36945, 36948, 36960, 36993, 36996, 37008, 37120, 37125, 37137, 37140, 37185, 37188, 37200, + 37210, 37377, 37380, 37392, 37440, 37542, 37888, 37890, 37893, 37896, 37905, 37908, 37920, 37953, 37956, 37968, + 38016, 38038, 38145, 38148, 38160, 38208, 38296, 38305, 38400, 38470, 38500, 38913, 38916, 38928, 38950, 38976, + 39081, 39168, 39241, 39250, 39568, 40960, 40965, 40970, 40980, 40994, 41002, 41025, 41028, 41040, 41122, 41130, + 41280, 41317, 41474, 41482, 41506, 41512, 41514, 41602, 41608, 41610, 41640, 41985, 41988, 42000, 42048, 42121, + 42148, 42240, 42265, 42577, 43018, 43048, 43170, 43348, 43398, 43528, 43530, 43552, 43554, 43560, 43656, 43690, + }; + + const int kmap_size = 43692; + //const int nwant = type == GGML_TYPE_IQ1_S ? 3 : 2; + const int nwant = type == GGML_TYPE_IQ1_S || type == GGML_TYPE_IQ1_M ? 3 : type == GGML_TYPE_IQ2_S ? 1 : 2; + const uint16_t * kgrid = type == GGML_TYPE_IQ2_XXS ? kgrid_2bit_256 : + type == GGML_TYPE_IQ2_XS ? kgrid_2bit_512 : + type == GGML_TYPE_IQ1_S || type == GGML_TYPE_IQ1_M ? kgrid_1bit_2048 : kgrid_2bit_1024; + uint64_t * kgrid_q2xs; + int * kmap_q2xs; + uint16_t * kneighbors_q2xs; + + //printf("================================================================= %s(grid_size = %d)\n", __func__, grid_size); + uint64_t * the_grid = (uint64_t *)malloc(grid_size*sizeof(uint64_t)); + for (int k = 0; k < grid_size; ++k) { + int8_t * pos = (int8_t *)(the_grid + k); + for (int i = 0; i < 8; ++i) { + int l = (kgrid[k] >> 2*i) & 0x3; + pos[i] = 2*l + 1; + } + } + kgrid_q2xs = the_grid; + iq2_data[gindex].grid = the_grid; + kmap_q2xs = (int *)malloc(kmap_size*sizeof(int)); + iq2_data[gindex].map = kmap_q2xs; + for (int i = 0; i < kmap_size; ++i) kmap_q2xs[i] = -1; + uint64_t aux64; + uint8_t * aux8 = (uint8_t *)&aux64; + for (int i = 0; i < grid_size; ++i) { + aux64 = kgrid_q2xs[i]; + uint16_t index = 0; + for (int k=0; k<8; ++k) { + uint16_t q = (aux8[k] - 1)/2; + index |= (q << 2*k); + } + kmap_q2xs[index] = i; + } + // The neighbour search runs in three passes: + // 1. Parallel: for each i, qsort and count its neighbours into n_per_i, + // and reduce the totals (num_neighbors, num_not_in_map). + // 2. Serial: prefix-sum n_per_i into offsets[], so each i has a + // pre-assigned slice of kneighbors_q2xs to write into. + // 3. Parallel: redo the qsort and write each i's neighbour list at + // offsets[i]. + int * n_per_i = (int *)malloc(kmap_size*sizeof(int)); + GGML_ASSERT(n_per_i); + int num_neighbors = 0, num_not_in_map = 0; +#ifdef GGML_USE_OPENMP + #pragma omp parallel reduction(+:num_neighbors,num_not_in_map) +#endif + { + int * dist2 = (int *)malloc(2*grid_size*sizeof(int)); + GGML_ASSERT(dist2); + int8_t pos[8]; + int i; +#ifdef GGML_USE_OPENMP + #pragma omp for schedule(dynamic, 64) +#endif + for (i = 0; i < kmap_size; ++i) { + if (kmap_q2xs[i] >= 0) { + n_per_i[i] = 0; + continue; + } + ++num_not_in_map; + for (int k = 0; k < 8; ++k) { + int l = (i >> 2*k) & 0x3; + pos[k] = 2*l + 1; + } + for (int j = 0; j < grid_size; ++j) { + const int8_t * pg = (const int8_t *)(kgrid_q2xs + j); + int d2 = 0; + for (int k = 0; k < 8; ++k) d2 += (pg[k] - pos[k])*(pg[k] - pos[k]); + dist2[2*j+0] = d2; + dist2[2*j+1] = j; + } + qsort(dist2, grid_size, 2*sizeof(int), iq2_compare_func); + int n = 0; int d2 = dist2[0]; + int nhave = 1; + for (int j = 0; j < grid_size; ++j) { + if (dist2[2*j] > d2) { + if (nhave == nwant) break; + d2 = dist2[2*j]; + ++nhave; + } + ++n; + } + n_per_i[i] = n; + num_neighbors += n; + } + free(dist2); + } + //printf("%s: %d neighbours in total\n", __func__, num_neighbors); + kneighbors_q2xs = (uint16_t *)malloc((num_neighbors + num_not_in_map)*sizeof(uint16_t)); + iq2_data[gindex].neighbours = kneighbors_q2xs; + + int * offsets = (int *)malloc(kmap_size*sizeof(int)); + GGML_ASSERT(offsets); + int counter = 0; + for (int i = 0; i < kmap_size; ++i) { + if (kmap_q2xs[i] >= 0) { + offsets[i] = -1; + continue; + } + offsets[i] = counter; + counter += 1 + n_per_i[i]; + } + +#ifdef GGML_USE_OPENMP + #pragma omp parallel +#endif + { + int * dist2 = (int *)malloc(2*grid_size*sizeof(int)); + GGML_ASSERT(dist2); + int8_t pos[8]; + int i; +#ifdef GGML_USE_OPENMP + #pragma omp for schedule(dynamic, 64) +#endif + for (i = 0; i < kmap_size; ++i) { + if (kmap_q2xs[i] >= 0) continue; + for (int k = 0; k < 8; ++k) { + int l = (i >> 2*k) & 0x3; + pos[k] = 2*l + 1; + } + for (int j = 0; j < grid_size; ++j) { + const int8_t * pg = (const int8_t *)(kgrid_q2xs + j); + int d2 = 0; + for (int k = 0; k < 8; ++k) d2 += (pg[k] - pos[k])*(pg[k] - pos[k]); + dist2[2*j+0] = d2; + dist2[2*j+1] = j; + } + qsort(dist2, grid_size, 2*sizeof(int), iq2_compare_func); + int local_counter = offsets[i]; + kmap_q2xs[i] = -(local_counter + 1); + int d2 = dist2[0]; + uint16_t * start = &kneighbors_q2xs[local_counter++]; + int n = 0, nhave = 1; + for (int j = 0; j < grid_size; ++j) { + if (dist2[2*j] > d2) { + if (nhave == nwant) break; + d2 = dist2[2*j]; + ++nhave; + } + kneighbors_q2xs[local_counter++] = dist2[2*j+1]; + ++n; + } + *start = n; + } + free(dist2); + } + free(offsets); + free(n_per_i); +} + +void iq2xs_free_impl(enum ggml_type type) { + GGML_ASSERT(type == GGML_TYPE_IQ2_XXS || type == GGML_TYPE_IQ2_XS || type == GGML_TYPE_IQ1_S || type == GGML_TYPE_IQ1_M || type == GGML_TYPE_IQ2_S); + const int gindex = iq2_data_index(type); + if (iq2_data[gindex].grid) { + free(iq2_data[gindex].grid); iq2_data[gindex].grid = NULL; + free(iq2_data[gindex].map); iq2_data[gindex].map = NULL; + free(iq2_data[gindex].neighbours); iq2_data[gindex].neighbours = NULL; + } +} + +static int iq2_find_best_neighbour(const uint16_t * GGML_RESTRICT neighbours, const uint64_t * GGML_RESTRICT grid, + const float * GGML_RESTRICT xval, const float * GGML_RESTRICT weight, float scale, int8_t * GGML_RESTRICT L) { + int num_neighbors = neighbours[0]; + GGML_ASSERT(num_neighbors > 0); + float best_d2 = FLT_MAX; + int grid_index = -1; + for (int j = 1; j <= num_neighbors; ++j) { + const int8_t * pg = (const int8_t *)(grid + neighbours[j]); + float d2 = 0; + for (int i = 0; i < 8; ++i) { + float q = pg[i]; + float diff = scale*q - xval[i]; + d2 += weight[i]*diff*diff; + } + if (d2 < best_d2) { + best_d2 = d2; grid_index = neighbours[j]; + } + } + GGML_ASSERT(grid_index >= 0); + const int8_t * pg = (const int8_t *)(grid + grid_index); + for (int i = 0; i < 8; ++i) L[i] = (pg[i] - 1)/2; + return grid_index; +} + +static void quantize_row_iq2_xxs_impl(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t n, const float * GGML_RESTRICT quant_weights) { + + const int gindex = iq2_data_index(GGML_TYPE_IQ2_XXS); + + const uint64_t * kgrid_q2xs = iq2_data[gindex].grid; + const int * kmap_q2xs = iq2_data[gindex].map; + const uint16_t * kneighbors_q2xs = iq2_data[gindex].neighbours; + + GGML_ASSERT(quant_weights && "missing quantization weights"); + GGML_ASSERT(kgrid_q2xs && "forgot to call ggml_quantize_init()?"); + GGML_ASSERT(kmap_q2xs && "forgot to call ggml_quantize_init()?"); + GGML_ASSERT(kneighbors_q2xs && "forgot to call ggml_quantize_init()?"); + GGML_ASSERT(n%QK_K == 0); + + const int kMaxQ = 3; + + const int64_t nbl = n/QK_K; + + block_iq2_xxs * y = vy; + + float scales[QK_K/32]; + float weight[32]; + float xval[32]; + int8_t L[32]; + int8_t Laux[32]; + float waux[32]; + uint8_t block_signs[4]; + uint32_t q2[2*(QK_K/32)]; + + for (int ibl = 0; ibl < nbl; ++ibl) { + + y[ibl].d = GGML_FP32_TO_FP16(0.f); + memset(q2, 0, QK_K/4); + + float max_scale = 0; + + const float * xbl = x + QK_K*ibl; + float sumx2 = 0; + for (int i = 0; i < QK_K; ++i) sumx2 += xbl[i]*xbl[i]; + float sigma2 = sumx2/QK_K; + + for (int ib = 0; ib < QK_K/32; ++ib) { + const float * xb = xbl + 32*ib; + const float * qw = quant_weights + QK_K*ibl + 32*ib; + for (int i = 0; i < 32; ++i) weight[i] = qw[i] * sqrtf(sigma2 + xb[i]*xb[i]); + for (int i = 0; i < 32; ++i) waux[i] = sqrtf(weight[i]); + for (int k = 0; k < 4; ++k) { + int nflip = 0; + uint8_t s = 0; + for (int i = 0; i < 8; ++i) { + if (xb[8*k + i] >= 0) xval[8*k + i] = xb[8*k + i]; + else { + xval[8*k + i] = -xb[8*k + i]; ++nflip; s |= (1 << i); + } + } + if (nflip%2) { + int imin = 0; float min = weight[8*k+imin]*xb[8*k+imin]*xb[8*k+imin]; + for (int i = 1; i < 8; ++i) { + float ax = weight[8*k+i]*xb[8*k+i]*xb[8*k+i]; + if (ax < min) { + min = ax; imin = i; + } + } + xval[8*k+imin] = -xval[8*k+imin]; + s ^= (1 << imin); + } + block_signs[k] = s & 127; + } + float max = xval[0]; + for (int i = 1; i < 32; ++i) max = MAX(max, xval[i]); + if (max < GROUP_MAX_EPS) { + scales[ib] = 0; + memset(L, 0, 32); + continue; + } + float scale = make_qp_quants(32, kMaxQ+1, xval, (uint8_t*)L, weight); + float eff_max = scale*kMaxQ; + if (eff_max <= 0) { + scales[ib] = 0; + memset(L, 0, 32); + continue; + } + float best = 0; + for (int is = -6; is <= 6; ++is) { + float id = (2*kMaxQ-1+is*0.1f)/eff_max; + float this_scale = 1/id; + for (int k = 0; k < 4; ++k) { + for (int i = 0; i < 8; ++i) { + int l = nearest_int(0.5f*(id*xval[8*k+i]-1)); + Laux[8*k+i] = MAX(0, MIN(kMaxQ-1, l)); + } + uint16_t u = 0; + for (int i = 0; i < 8; ++i) u |= (Laux[8*k+i] << 2*i); + int grid_index = kmap_q2xs[u]; + if (grid_index < 0) { + const uint16_t * neighbours = kneighbors_q2xs - kmap_q2xs[u] - 1; + grid_index = iq2_find_best_neighbour(neighbours, kgrid_q2xs, xval + 8*k, waux + 8*k, this_scale, Laux + 8*k); + } + } + float sumqx = 0, sumq2 = 0; + for (int i = 0; i < 32; ++i) { + float w = weight[i]; + float q = 2*Laux[i] + 1; + sumqx += w*xval[i]*q; + sumq2 += w*q*q; + } + if (sumq2 > 0 && sumqx*sumqx > best*sumq2) { + scale = sumqx/sumq2; best = scale*sumqx; + memcpy(L, Laux, 32); + } + } + if (scale > 0) { + float id = 1/scale; + for (int k = 0; k < 4; ++k) { + uint16_t u = 0; + for (int i = 0; i < 8; ++i) { + int l = nearest_int(0.5f*(id*xval[8*k+i]-1)); + l = MAX(0, MIN(kMaxQ-1, l)); + u |= (l << 2*i); + } + int grid_index = kmap_q2xs[u]; + if (grid_index < 0) { + const uint16_t * neighbours = kneighbors_q2xs - kmap_q2xs[u] - 1; + grid_index = iq2_find_best_neighbour(neighbours, kgrid_q2xs, xval + 8*k, waux + 8*k, scale, L + 8*k); + } + const int8_t * pg = (const int8_t *)(kgrid_q2xs + grid_index); + for (int i = 0; i < 8; ++i) L[8*k+i] = (pg[i] - 1)/2; + } + float sumqx = 0, sumq2 = 0; + for (int i = 0; i < 32; ++i) { + float w = weight[i]; + float q = 2*L[i] + 1; + sumqx += w*xval[i]*q; + sumq2 += w*q*q; + } + if (sumq2 > 0) scale = sumqx/sumq2; + } + if (scale < 0) { + // This should never happen, but just in case, flip scale so that it is positive (we use uint's to encode the scale) + // and correspondingly flip quant signs. + scale = -scale; + for (int k = 0; k < 4; ++k) block_signs[k] = (~block_signs[k]) & 127; + } + for (int k = 0; k < 4; ++k) { + uint16_t u = 0; + for (int i = 0; i < 8; ++i) u |= (L[8*k+i] << 2*i); + int grid_index = kmap_q2xs[u]; + if (grid_index < 0) { + printf("Oops: found point %u not on grid:", u); + for (int i = 0; i < 8; ++i) printf(" %d", L[8*k+i]); + printf("\n"); + GGML_ABORT("fatal error"); + } + q2[2*ib+0] |= ((uint32_t) grid_index << 8*k); + q2[2*ib+1] |= (block_signs[k] << 7*k); + } + GGML_ASSERT(scale >= 0); + scales[ib] = scale; + max_scale = MAX(max_scale, scale); + } + + if (!max_scale) { + memset(y[ibl].qs, 0, QK_K/4); + continue; + } + + float d = max_scale/31; + y[ibl].d = GGML_FP32_TO_FP16(d); + float id = 1/d; + for (int ib = 0; ib < QK_K/32; ++ib) { + int l = nearest_int(0.5f*(id*scales[ib]-1)); + l = MAX(0, MIN(15, l)); + q2[2*ib+1] |= ((uint32_t)l << 28); + } + memcpy(y[ibl].qs, q2, QK_K/4); + } +} + +static void quantize_row_iq2_xs_impl(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t n, const float * GGML_RESTRICT quant_weights) { + + const int gindex = iq2_data_index(GGML_TYPE_IQ2_XS); + + const uint64_t * kgrid_q2xs = iq2_data[gindex].grid; + const int * kmap_q2xs = iq2_data[gindex].map; + const uint16_t * kneighbors_q2xs = iq2_data[gindex].neighbours; + + GGML_ASSERT(quant_weights && "missing quantization weights"); + GGML_ASSERT(kmap_q2xs && "forgot to call ggml_quantize_init()?"); + GGML_ASSERT(kgrid_q2xs && "forgot to call ggml_quantize_init()?"); + GGML_ASSERT(kneighbors_q2xs && "forgot to call ggml_quantize_init()?"); + GGML_ASSERT(n%QK_K == 0); + + const int kMaxQ = 3; + + const int64_t nbl = n/QK_K; + + block_iq2_xs * y = vy; + + float scales[QK_K/16]; + float weight[16]; + float xval[16]; + int8_t L[16]; + int8_t Laux[16]; + float waux[16]; + bool is_on_grid[2]; + bool is_on_grid_aux[2]; + uint8_t block_signs[2]; + uint16_t q2[2*(QK_K/16)]; + + for (int ibl = 0; ibl < nbl; ++ibl) { + + y[ibl].d = GGML_FP32_TO_FP16(0.f); + memset(q2, 0, QK_K/4); + memset(y[ibl].scales, 0, QK_K/32); + + float max_scale = 0; + + const float * xbl = x + QK_K*ibl; + float sumx2 = 0; + for (int i = 0; i < QK_K; ++i) sumx2 += xbl[i]*xbl[i]; + float sigma2 = sumx2/QK_K; + + for (int ib = 0; ib < QK_K/16; ++ib) { + const float * xb = xbl + 16*ib; + const float * qw = quant_weights + QK_K*ibl + 16*ib; + for (int i = 0; i < 16; ++i) weight[i] = qw[i] * sqrtf(sigma2 + xb[i]*xb[i]); + for (int i = 0; i < 16; ++i) waux[i] = sqrtf(weight[i]); + for (int k = 0; k < 2; ++k) { + int nflip = 0; + uint8_t s = 0; + for (int i = 0; i < 8; ++i) { + if (xb[8*k + i] >= 0) xval[8*k + i] = xb[8*k + i]; + else { + xval[8*k + i] = -xb[8*k + i]; ++nflip; s |= (1 << i); + } + } + if (nflip%2) { + int imin = 0; float min = weight[8*k+imin]*xb[8*k+imin]*xb[8*k+imin]; + for (int i = 1; i < 8; ++i) { + float ax = weight[8*k+i]*xb[8*k+i]*xb[8*k+i]; + if (ax < min) { + min = ax; imin = i; + } + } + xval[8*k+imin] = -xval[8*k+imin]; + s ^= (1 << imin); + } + block_signs[k] = s & 127; + } + float max = xval[0]; + for (int i = 1; i < 16; ++i) max = MAX(max, xval[i]); + memset(L, 0, 16); + if (max < GROUP_MAX_EPS) { + scales[ib] = 0; + continue; + } + float best = 0; + float scale = max/(2*kMaxQ-1); + is_on_grid[0] = is_on_grid[1] = true; + for (int is = -9; is <= 9; ++is) { + float id = (2*kMaxQ-1+is*0.1f)/max; + float this_scale = 1/id; + for (int k = 0; k < 2; ++k) { + for (int i = 0; i < 8; ++i) { + int l = nearest_int(0.5f*(id*xval[8*k+i]-1)); + Laux[8*k+i] = MAX(0, MIN(kMaxQ-1, l)); + } + uint16_t u = 0; + for (int i = 0; i < 8; ++i) u |= (Laux[8*k+i] << 2*i); + int grid_index = kmap_q2xs[u]; + is_on_grid_aux[k] = true; + if (grid_index < 0) { + is_on_grid_aux[k] = false; + const uint16_t * neighbours = kneighbors_q2xs - kmap_q2xs[u] - 1; + grid_index = iq2_find_best_neighbour(neighbours, kgrid_q2xs, xval + 8*k, waux + 8*k, this_scale, Laux + 8*k); + } + } + float sumqx = 0, sumq2 = 0; + for (int i = 0; i < 16; ++i) { + float w = weight[i]; + float q = 2*Laux[i] + 1; + sumqx += w*xval[i]*q; + sumq2 += w*q*q; + } + if (sumq2 > 0 && sumqx*sumqx > best*sumq2) { + scale = sumqx/sumq2; best = scale*sumqx; + for (int i = 0; i < 16; ++i) L[i] = Laux[i]; + for (int k = 0; k < 2; ++k) is_on_grid[k] = is_on_grid_aux[k]; + } + } + int n_not_ongrid = 0; + for (int k = 0; k < 2; ++k) if (!is_on_grid[k]) ++n_not_ongrid; + if (n_not_ongrid > 0 && scale > 0) { + float id = 1/scale; + for (int k = 0; k < 2; ++k) { + if (is_on_grid[k]) continue; + uint16_t u = 0; + for (int i = 0; i < 8; ++i) { + int l = nearest_int(0.5f*(id*xval[8*k+i]-1)); + l = MAX(0, MIN(kMaxQ-1, l)); + u |= (l << 2*i); + L[8*k + i] = l; + } + int grid_index = kmap_q2xs[u]; + if (grid_index < 0) { + const uint16_t * neighbours = kneighbors_q2xs - kmap_q2xs[u] - 1; + grid_index = iq2_find_best_neighbour(neighbours, kgrid_q2xs, xval + 8*k, waux + 8*k, scale, L + 8*k); + } + } + float sumqx = 0, sumq2 = 0; + for (int i = 0; i < 16; ++i) { + float w = weight[i]; + float q = 2*L[i] + 1; + sumqx += w*xval[i]*q; + sumq2 += w*q*q; + } + if (sumq2 > 0) scale = sumqx/sumq2; + } + if (scale < 0) { + scale = -scale; + for (int k = 0; k < 2; ++k) block_signs[k] = (~block_signs[k]) & 127; + } + for (int k = 0; k < 2; ++k) { + uint16_t u = 0; + for (int i = 0; i < 8; ++i) u |= (L[8*k+i] << 2*i); + int grid_index = kmap_q2xs[u]; + if (grid_index < 0) { + printf("Oops: found point %u not on grid:", u); + for (int i = 0; i < 8; ++i) printf(" %d", L[8*k+i]); + printf("\n"); + GGML_ABORT("fatal error"); + } + q2[2*ib+k] = grid_index | (block_signs[k] << 9); + } + GGML_ASSERT(scale >= 0); + scales[ib] = scale; + max_scale = MAX(max_scale, scale); + } + + if (!max_scale) { + memset(y[ibl].qs, 0, QK_K/4); + continue; + } + + float d = max_scale/31; + y[ibl].d = GGML_FP32_TO_FP16(d); + float id = 1/d; + for (int ib = 0; ib < QK_K/16; ++ib) { + int l = nearest_int(0.5f*(id*scales[ib]-1)); + l = MAX(0, MIN(15, l)); + if (ib%2 == 0) y[ibl].scales[ib/2] = l; + else y[ibl].scales[ib/2] |= (l << 4); + } + memcpy(y[ibl].qs, q2, QK_K/4); + + } +} + +size_t quantize_iq2_xxs(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) { + GGML_ASSERT(n_per_row%QK_K == 0); + int64_t nblock = n_per_row/QK_K; + char * qrow = (char *)dst; + for (int64_t row = 0; row < nrow; ++row) { + quantize_row_iq2_xxs_impl(src, qrow, n_per_row, quant_weights); + src += n_per_row; + qrow += nblock*sizeof(block_iq2_xxs); + } + return nrow * nblock * sizeof(block_iq2_xxs); +} + +size_t quantize_iq2_xs(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) { + GGML_ASSERT(n_per_row%QK_K == 0); + int64_t nblock = n_per_row/QK_K; + char * qrow = (char *)dst; + for (int64_t row = 0; row < nrow; ++row) { + quantize_row_iq2_xs_impl(src, qrow, n_per_row, quant_weights); + src += n_per_row; + qrow += nblock*sizeof(block_iq2_xs); + } + return nrow * nblock * sizeof(block_iq2_xs); +} + +// +// ============================================= 3-bit using D4 lattice +// + +typedef struct { + uint32_t * grid; + int * map; + uint16_t * neighbours; +} iq3_entry_t; + +static iq3_entry_t iq3_data[2] = { + {NULL, NULL, NULL}, + {NULL, NULL, NULL}, +}; + +static inline int iq3_data_index(int grid_size) { + (void)grid_size; + GGML_ASSERT(grid_size == 256 || grid_size == 512); + return grid_size == 256 ? 0 : 1; +} + +static int iq3_compare_func(const void * left, const void * right) { + const int * l = (const int *)left; + const int * r = (const int *)right; + return l[0] < r[0] ? -1 : l[0] > r[0] ? 1 : l[1] < r[1] ? -1 : l[1] > r[1] ? 1 : 0; +} + +void iq3xs_init_impl(int grid_size) { + const int gindex = iq3_data_index(grid_size); + if (iq3_data[gindex].grid) { + return; + } + static const uint16_t kgrid_256[256] = { + 0, 2, 4, 9, 11, 15, 16, 18, 25, 34, 59, 61, 65, 67, 72, 74, + 81, 85, 88, 90, 97, 108, 120, 128, 130, 132, 137, 144, 146, 153, 155, 159, + 169, 175, 189, 193, 199, 200, 202, 213, 248, 267, 287, 292, 303, 315, 317, 321, + 327, 346, 362, 413, 436, 456, 460, 462, 483, 497, 513, 515, 520, 522, 529, 531, + 536, 538, 540, 551, 552, 576, 578, 585, 592, 594, 641, 643, 648, 650, 657, 664, + 698, 704, 706, 720, 729, 742, 758, 769, 773, 808, 848, 852, 870, 889, 901, 978, + 992, 1024, 1026, 1033, 1035, 1040, 1042, 1046, 1049, 1058, 1089, 1091, 1093, 1096, 1098, 1105, + 1112, 1139, 1143, 1144, 1152, 1154, 1161, 1167, 1168, 1170, 1183, 1184, 1197, 1217, 1224, 1228, + 1272, 1276, 1309, 1323, 1347, 1367, 1377, 1404, 1473, 1475, 1486, 1509, 1537, 1544, 1546, 1553, + 1555, 1576, 1589, 1594, 1600, 1602, 1616, 1625, 1636, 1638, 1665, 1667, 1672, 1685, 1706, 1722, + 1737, 1755, 1816, 1831, 1850, 1856, 1862, 1874, 1901, 1932, 1950, 1971, 2011, 2032, 2052, 2063, + 2077, 2079, 2091, 2095, 2172, 2192, 2207, 2208, 2224, 2230, 2247, 2277, 2308, 2345, 2356, 2389, + 2403, 2424, 2501, 2504, 2506, 2520, 2570, 2593, 2616, 2624, 2630, 2646, 2669, 2700, 2714, 2746, + 2754, 2795, 2824, 2835, 2839, 2874, 2882, 2905, 2984, 3028, 3042, 3092, 3108, 3110, 3124, 3153, + 3185, 3215, 3252, 3288, 3294, 3364, 3397, 3434, 3483, 3523, 3537, 3587, 3589, 3591, 3592, 3610, + 3626, 3670, 3680, 3722, 3749, 3754, 3776, 3789, 3803, 3824, 3857, 3873, 3904, 3906, 3924, 3992, + }; + static const uint16_t kgrid_512[512] = { + 0, 1, 2, 5, 7, 8, 9, 10, 12, 14, 16, 17, 21, 27, 32, 34, + 37, 39, 41, 43, 48, 50, 57, 60, 63, 64, 65, 66, 68, 72, 73, 77, + 80, 83, 87, 89, 93, 100, 113, 117, 122, 128, 129, 133, 135, 136, 139, 142, + 145, 149, 152, 156, 162, 165, 167, 169, 171, 184, 187, 195, 201, 205, 208, 210, + 217, 219, 222, 228, 232, 234, 247, 249, 253, 256, 267, 271, 273, 276, 282, 288, + 291, 297, 312, 322, 324, 336, 338, 342, 347, 353, 357, 359, 374, 379, 390, 393, + 395, 409, 426, 441, 448, 450, 452, 464, 466, 470, 475, 488, 492, 512, 513, 514, + 516, 520, 521, 523, 525, 527, 528, 530, 537, 540, 542, 556, 558, 561, 570, 576, + 577, 579, 582, 584, 588, 593, 600, 603, 609, 616, 618, 632, 638, 640, 650, 653, + 655, 656, 660, 666, 672, 675, 685, 688, 698, 705, 708, 711, 712, 715, 721, 727, + 728, 732, 737, 754, 760, 771, 773, 778, 780, 793, 795, 802, 806, 808, 812, 833, + 840, 843, 849, 856, 858, 873, 912, 916, 919, 932, 934, 961, 963, 968, 970, 977, + 989, 993, 1010, 1016, 1024, 1025, 1027, 1029, 1031, 1032, 1034, 1036, 1038, 1041, 1043, 1047, + 1048, 1050, 1057, 1059, 1061, 1064, 1066, 1079, 1080, 1083, 1085, 1088, 1090, 1096, 1099, 1103, + 1106, 1109, 1113, 1116, 1122, 1129, 1153, 1156, 1159, 1169, 1171, 1176, 1183, 1185, 1195, 1199, + 1209, 1212, 1216, 1218, 1221, 1225, 1234, 1236, 1241, 1243, 1250, 1256, 1270, 1281, 1287, 1296, + 1299, 1306, 1309, 1313, 1338, 1341, 1348, 1353, 1362, 1375, 1376, 1387, 1400, 1408, 1410, 1415, + 1425, 1453, 1457, 1477, 1481, 1494, 1496, 1507, 1512, 1538, 1545, 1547, 1549, 1551, 1554, 1561, + 1563, 1565, 1570, 1572, 1575, 1577, 1587, 1593, 1601, 1603, 1605, 1612, 1617, 1619, 1632, 1648, + 1658, 1662, 1664, 1674, 1680, 1690, 1692, 1704, 1729, 1736, 1740, 1745, 1747, 1751, 1752, 1761, + 1763, 1767, 1773, 1787, 1795, 1801, 1806, 1810, 1817, 1834, 1840, 1844, 1857, 1864, 1866, 1877, + 1882, 1892, 1902, 1915, 1934, 1953, 1985, 1987, 2000, 2002, 2013, 2048, 2052, 2058, 2064, 2068, + 2071, 2074, 2081, 2088, 2104, 2114, 2119, 2121, 2123, 2130, 2136, 2141, 2147, 2153, 2157, 2177, + 2179, 2184, 2189, 2193, 2203, 2208, 2223, 2226, 2232, 2244, 2249, 2251, 2256, 2258, 2265, 2269, + 2304, 2306, 2324, 2335, 2336, 2361, 2373, 2375, 2385, 2418, 2443, 2460, 2480, 2504, 2509, 2520, + 2531, 2537, 2562, 2568, 2572, 2578, 2592, 2596, 2599, 2602, 2614, 2620, 2625, 2627, 2629, 2634, + 2641, 2650, 2682, 2688, 2697, 2707, 2712, 2718, 2731, 2754, 2759, 2760, 2775, 2788, 2793, 2805, + 2811, 2817, 2820, 2832, 2842, 2854, 2890, 2902, 2921, 2923, 2978, 3010, 3012, 3026, 3081, 3083, + 3085, 3097, 3099, 3120, 3136, 3152, 3159, 3188, 3210, 3228, 3234, 3245, 3250, 3256, 3264, 3276, + 3281, 3296, 3349, 3363, 3378, 3392, 3395, 3420, 3440, 3461, 3488, 3529, 3531, 3584, 3588, 3591, + 3600, 3602, 3614, 3616, 3628, 3634, 3650, 3657, 3668, 3683, 3685, 3713, 3716, 3720, 3726, 3729, + 3736, 3753, 3778, 3802, 3805, 3819, 3841, 3845, 3851, 3856, 3880, 3922, 3938, 3970, 3993, 4032, + }; + + const int kmap_size = 4096; + const int nwant = grid_size == 256 ? 2 : 3; + const uint16_t * kgrid = grid_size == 256 ? kgrid_256 : kgrid_512; + uint32_t * kgrid_q3xs; + int * kmap_q3xs; + uint16_t * kneighbors_q3xs; + + //printf("================================================================= %s(grid_size = %d)\n", __func__, grid_size); + uint32_t * the_grid = (uint32_t *)malloc(grid_size*sizeof(uint32_t)); + for (int k = 0; k < grid_size; ++k) { + int8_t * pos = (int8_t *)(the_grid + k); + for (int i = 0; i < 4; ++i) { + int l = (kgrid[k] >> 3*i) & 0x7; + pos[i] = 2*l + 1; + } + } + kgrid_q3xs = the_grid; + iq3_data[gindex].grid = the_grid; + kmap_q3xs = (int *)malloc(kmap_size*sizeof(int)); + iq3_data[gindex].map = kmap_q3xs; + for (int i = 0; i < kmap_size; ++i) kmap_q3xs[i] = -1; + uint32_t aux32; + uint8_t * aux8 = (uint8_t *)&aux32; + for (int i = 0; i < grid_size; ++i) { + aux32 = kgrid_q3xs[i]; + uint16_t index = 0; + for (int k=0; k<4; ++k) { + uint16_t q = (aux8[k] - 1)/2; + index |= (q << 3*k); + } + kmap_q3xs[index] = i; + } + // See explanation of parallelism in iq2xs_init_impl + int * n_per_i = (int *)malloc(kmap_size*sizeof(int)); + GGML_ASSERT(n_per_i); + int num_neighbors = 0, num_not_in_map = 0; +#ifdef GGML_USE_OPENMP + #pragma omp parallel reduction(+:num_neighbors,num_not_in_map) +#endif + { + int * dist2 = (int *)malloc(2*grid_size*sizeof(int)); + GGML_ASSERT(dist2); + int8_t pos[4]; + int i; +#ifdef GGML_USE_OPENMP + #pragma omp for schedule(dynamic, 64) +#endif + for (i = 0; i < kmap_size; ++i) { + if (kmap_q3xs[i] >= 0) { + n_per_i[i] = 0; + continue; + } + ++num_not_in_map; + for (int k = 0; k < 4; ++k) { + int l = (i >> 3*k) & 0x7; + pos[k] = 2*l + 1; + } + for (int j = 0; j < grid_size; ++j) { + const int8_t * pg = (const int8_t *)(kgrid_q3xs + j); + int d2 = 0; + for (int k = 0; k < 4; ++k) d2 += (pg[k] - pos[k])*(pg[k] - pos[k]); + dist2[2*j+0] = d2; + dist2[2*j+1] = j; + } + qsort(dist2, grid_size, 2*sizeof(int), iq3_compare_func); + int n = 0; int d2 = dist2[0]; + int nhave = 1; + for (int j = 0; j < grid_size; ++j) { + if (dist2[2*j] > d2) { + if (nhave == nwant) break; + d2 = dist2[2*j]; + ++nhave; + } + ++n; + } + n_per_i[i] = n; + num_neighbors += n; + } + free(dist2); + } + //printf("%s: %d neighbours in total\n", __func__, num_neighbors); + kneighbors_q3xs = (uint16_t *)malloc((num_neighbors + num_not_in_map)*sizeof(uint16_t)); + iq3_data[gindex].neighbours = kneighbors_q3xs; + + int * offsets = (int *)malloc(kmap_size*sizeof(int)); + GGML_ASSERT(offsets); + int counter = 0; + for (int i = 0; i < kmap_size; ++i) { + if (kmap_q3xs[i] >= 0) { + offsets[i] = -1; + continue; + } + offsets[i] = counter; + counter += 1 + n_per_i[i]; + } + +#ifdef GGML_USE_OPENMP + #pragma omp parallel +#endif + { + int * dist2 = (int *)malloc(2*grid_size*sizeof(int)); + GGML_ASSERT(dist2); + int8_t pos[4]; + int i; +#ifdef GGML_USE_OPENMP + #pragma omp for schedule(dynamic, 64) +#endif + for (i = 0; i < kmap_size; ++i) { + if (kmap_q3xs[i] >= 0) continue; + for (int k = 0; k < 4; ++k) { + int l = (i >> 3*k) & 0x7; + pos[k] = 2*l + 1; + } + for (int j = 0; j < grid_size; ++j) { + const int8_t * pg = (const int8_t *)(kgrid_q3xs + j); + int d2 = 0; + for (int k = 0; k < 4; ++k) d2 += (pg[k] - pos[k])*(pg[k] - pos[k]); + dist2[2*j+0] = d2; + dist2[2*j+1] = j; + } + qsort(dist2, grid_size, 2*sizeof(int), iq3_compare_func); + int local_counter = offsets[i]; + kmap_q3xs[i] = -(local_counter + 1); + int d2 = dist2[0]; + uint16_t * start = &kneighbors_q3xs[local_counter++]; + int n = 0, nhave = 1; + for (int j = 0; j < grid_size; ++j) { + if (dist2[2*j] > d2) { + if (nhave == nwant) break; + d2 = dist2[2*j]; + ++nhave; + } + kneighbors_q3xs[local_counter++] = dist2[2*j+1]; + ++n; + } + *start = n; + } + free(dist2); + } + free(offsets); + free(n_per_i); +} + +void iq3xs_free_impl(int grid_size) { + GGML_ASSERT(grid_size == 256 || grid_size == 512); + const int gindex = iq3_data_index(grid_size); + if (iq3_data[gindex].grid) { + free(iq3_data[gindex].grid); iq3_data[gindex].grid = NULL; + free(iq3_data[gindex].map); iq3_data[gindex].map = NULL; + free(iq3_data[gindex].neighbours); iq3_data[gindex].neighbours = NULL; + } +} + +static int iq3_find_best_neighbour(const uint16_t * GGML_RESTRICT neighbours, const uint32_t * GGML_RESTRICT grid, + const float * GGML_RESTRICT xval, const float * GGML_RESTRICT weight, float scale, int8_t * GGML_RESTRICT L) { + int num_neighbors = neighbours[0]; + GGML_ASSERT(num_neighbors > 0); + float best_d2 = FLT_MAX; + int grid_index = -1; + for (int j = 1; j <= num_neighbors; ++j) { + const int8_t * pg = (const int8_t *)(grid + neighbours[j]); + float d2 = 0; + for (int i = 0; i < 4; ++i) { + float q = pg[i]; + float diff = scale*q - xval[i]; + d2 += weight[i]*diff*diff; + } + if (d2 < best_d2) { + best_d2 = d2; grid_index = neighbours[j]; + } + } + GGML_ASSERT(grid_index >= 0); + const int8_t * pg = (const int8_t *)(grid + grid_index); + for (int i = 0; i < 4; ++i) L[i] = (pg[i] - 1)/2; + return grid_index; +} + +static void quantize_row_iq3_xxs_impl(int grid_size, const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t n, + const float * GGML_RESTRICT quant_weights) { + + const int gindex = iq3_data_index(grid_size); + + const uint32_t * kgrid_q3xs = iq3_data[gindex].grid; + const int * kmap_q3xs = iq3_data[gindex].map; + const uint16_t * kneighbors_q3xs = iq3_data[gindex].neighbours; + + //GGML_ASSERT(quant_weights && "missing quantization weights"); + GGML_ASSERT(kgrid_q3xs && "forgot to call ggml_quantize_init()?"); + GGML_ASSERT(kmap_q3xs && "forgot to call ggml_quantize_init()?"); + GGML_ASSERT(kneighbors_q3xs && "forgot to call ggml_quantize_init()?"); + GGML_ASSERT(n%QK_K == 0); + + const int kMaxQ = 8; + + const int64_t nbl = n/QK_K; + + ggml_fp16_t * dh; + uint8_t * qs; + int block_size; + if (grid_size == 256) { + block_iq3_xxs * y = vy; + dh = &y->d; + qs = y->qs; + block_size = sizeof(block_iq3_xxs); + } else { + block_iq3_s * y = vy; + dh = &y->d; + qs = y->qs; + block_size = sizeof(block_iq3_s); + } + int quant_size = block_size - sizeof(ggml_fp16_t); + + float scales[QK_K/32]; + float weight[32]; + float xval[32]; + int8_t L[32]; + int8_t Laux[32]; + float waux[32]; + bool is_on_grid[8]; + bool is_on_grid_aux[8]; + uint8_t block_signs[8]; + uint8_t q3[3*(QK_K/8)+QK_K/32]; + uint32_t * scales_and_signs = (uint32_t *)(q3 + QK_K/4); + uint8_t * qh = q3 + 3*(QK_K/8); + + for (int ibl = 0; ibl < nbl; ++ibl) { + + dh[0] = GGML_FP32_TO_FP16(0.f); + memset(q3, 0, 3*QK_K/8+QK_K/32); + + float max_scale = 0; + + const float * xbl = x + QK_K*ibl; + float sumx2 = 0; + for (int i = 0; i < QK_K; ++i) sumx2 += xbl[i]*xbl[i]; + float sigma2 = 2*sumx2/QK_K; + + for (int ib = 0; ib < QK_K/32; ++ib) { + const float * xb = xbl + 32*ib; + if (quant_weights) { + const float * qw = quant_weights + QK_K*ibl + 32*ib; + for (int i = 0; i < 32; ++i) weight[i] = qw[i] * sqrtf(sigma2 + xb[i]*xb[i]); + } else { + for (int i = 0; i < 32; ++i) weight[i] = xb[i]*xb[i]; + } + for (int i = 0; i < 32; ++i) waux[i] = sqrtf(weight[i]); + for (int k = 0; k < 4; ++k) { + int nflip = 0; + uint8_t s = 0; + for (int i = 0; i < 8; ++i) { + if (xb[8*k + i] >= 0) xval[8*k + i] = xb[8*k + i]; + else { + xval[8*k + i] = -xb[8*k + i]; ++nflip; s |= (1 << i); + } + } + if (nflip%2) { + int imin = 0; float min = weight[8*k+imin]*xb[8*k+imin]*xb[8*k+imin]; + for (int i = 1; i < 8; ++i) { + float ax = weight[8*k+i]*xb[8*k+i]*xb[8*k+i]; + if (ax < min) { + min = ax; imin = i; + } + } + xval[8*k+imin] = -xval[8*k+imin]; + s ^= (1 << imin); + } + block_signs[k] = s & 127; + } + float max = xval[0]; + for (int i = 1; i < 32; ++i) max = MAX(max, xval[i]); + memset(L, 0, 32); + if (max < GROUP_MAX_EPS_IQ3_XXS) { + scales[ib] = 0; + continue; + } + float best = 0; + float scale = max/(2*kMaxQ-1); + for (int k = 0; k < 8; ++k) is_on_grid[k] = true; + for (int is = -15; is <= 15; ++is) { + float id = (2*kMaxQ-1+is*0.2f)/max; + float this_scale = 1/id; + for (int k = 0; k < 8; ++k) { + for (int i = 0; i < 4; ++i) { + int l = nearest_int(0.5f*(id*xval[4*k+i]-1)); + Laux[4*k+i] = MAX(0, MIN(kMaxQ-1, l)); + } + uint16_t u = 0; + for (int i = 0; i < 4; ++i) u |= (Laux[4*k+i] << 3*i); + int grid_index = kmap_q3xs[u]; + is_on_grid_aux[k] = true; + if (grid_index < 0) { + is_on_grid_aux[k] = false; + const uint16_t * neighbours = kneighbors_q3xs - kmap_q3xs[u] - 1; + grid_index = iq3_find_best_neighbour(neighbours, kgrid_q3xs, xval + 4*k, waux + 4*k, this_scale, Laux + 4*k); + } + } + float sumqx = 0, sumq2 = 0; + for (int i = 0; i < 32; ++i) { + float w = weight[i]; + float q = 2*Laux[i] + 1; + sumqx += w*xval[i]*q; + sumq2 += w*q*q; + } + if (sumq2 > 0 && sumqx*sumqx > best*sumq2) { + scale = sumqx/sumq2; best = scale*sumqx; + for (int i = 0; i < 32; ++i) L[i] = Laux[i]; + for (int k = 0; k < 8; ++k) is_on_grid[k] = is_on_grid_aux[k]; + } + } + int n_not_ongrid = 0; + for (int k = 0; k < 8; ++k) if (!is_on_grid[k]) ++n_not_ongrid; + if (n_not_ongrid > 0 && scale > 0) { + float id = 1/scale; + for (int k = 0; k < 8; ++k) { + if (is_on_grid[k]) continue; + uint16_t u = 0; + for (int i = 0; i < 4; ++i) { + int l = nearest_int(0.5f*(id*xval[4*k+i]-1)); + l = MAX(0, MIN(kMaxQ-1, l)); + u |= (l << 3*i); + } + int grid_index = kmap_q3xs[u]; + if (grid_index < 0) { + const uint16_t * neighbours = kneighbors_q3xs - kmap_q3xs[u] - 1; + grid_index = iq3_find_best_neighbour(neighbours, kgrid_q3xs, xval + 4*k, waux + 4*k, scale, L + 4*k); + } + const int8_t * pg = (const int8_t *)(kgrid_q3xs + grid_index); + for (int i = 0; i < 4; ++i) L[4*k+i] = (pg[i] - 1)/2; + } + float sumqx = 0, sumq2 = 0; + for (int i = 0; i < 32; ++i) { + float w = weight[i]; + float q = 2*L[i] + 1; + sumqx += w*xval[i]*q; + sumq2 += w*q*q; + } + if (sumq2 > 0) scale = sumqx/sumq2; + } + if (scale < 0) { + // This should never happen, but just in case, flip scale so that it is positive (we use uint's to encode the scale) + // and correspondingly flip quant signs. + scale = -scale; + for (int k = 0; k < 4; ++k) block_signs[k] = (~block_signs[k]) & 127; + } + for (int k = 0; k < 8; ++k) { + uint16_t u = 0; + for (int i = 0; i < 4; ++i) u |= (L[4*k+i] << 3*i); + int grid_index = kmap_q3xs[u]; + if (grid_index < 0) { + printf("Oops: found point %u not on grid:", u); + for (int i = 0; i < 4; ++i) printf(" %d", L[4*k+i]); + printf("\n"); + GGML_ABORT("fatal error"); + } + if (grid_size == 256) { + q3[8*ib+k] = grid_index; + } else { + q3[8*ib+k] = grid_index & 255; + qh[ib] |= ((grid_index >> 8) << k); + } + + } + scales_and_signs[ib] = block_signs[0] | (block_signs[1] << 7) | (block_signs[2] << 14) | (block_signs[3] << 21); + GGML_ASSERT(scale >= 0); + scales[ib] = scale; + max_scale = MAX(max_scale, scale); + } + + if (!max_scale) { + memset(qs, 0, quant_size); + dh += block_size/sizeof(ggml_fp16_t); + qs += block_size; + continue; + } + + float d = max_scale/31; + dh[0] = GGML_FP32_TO_FP16(d * 1.0125f); // small improvement via this fudge factor + float id = 1/d; + for (int ib = 0; ib < QK_K/32; ++ib) { + int l = nearest_int(0.5f*(id*scales[ib]-1)); + l = MAX(0, MIN(15, l)); + scales_and_signs[ib] |= ((uint32_t)l << 28); + } + memcpy(qs, q3, quant_size); + + dh += block_size/sizeof(ggml_fp16_t); + qs += block_size; + + } +} + +size_t quantize_iq3_xxs(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) { + GGML_ASSERT(n_per_row%QK_K == 0); + int64_t nblock = n_per_row/QK_K; + char * qrow = (char *)dst; + for (int64_t row = 0; row < nrow; ++row) { + quantize_row_iq3_xxs_impl(256, src, qrow, n_per_row, quant_weights); + src += n_per_row; + qrow += nblock*sizeof(block_iq3_xxs); + } + return nrow * nblock * sizeof(block_iq3_xxs); +} + +void quantize_row_iq3_xxs_ref(const float * GGML_RESTRICT x, block_iq3_xxs * GGML_RESTRICT y, int64_t k) { + assert(k % QK_K == 0); + quantize_row_iq3_xxs_impl(256, x, y, k, NULL); +} + +static void quantize_row_iq3_s_impl(int block_size, const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int n, + const float * GGML_RESTRICT quant_weights, + float * scales, + float * weight, + float * xval, + int8_t * L, + int8_t * Laux, + float * waux, + bool * is_on_grid, + bool * is_on_grid_aux, + uint8_t * block_signs) { + + const int gindex = iq3_data_index(512); + + const uint32_t * kgrid_q3xs = iq3_data[gindex].grid; + const int * kmap_q3xs = iq3_data[gindex].map; + const uint16_t * kneighbors_q3xs = iq3_data[gindex].neighbours; + + //GGML_ASSERT(quant_weights && "missing quantization weights"); + GGML_ASSERT(kgrid_q3xs && "forgot to call ggml_quantize_init()?"); + GGML_ASSERT(kmap_q3xs && "forgot to call ggml_quantize_init()?"); + GGML_ASSERT(kneighbors_q3xs && "forgot to call ggml_quantize_init()?"); + GGML_ASSERT(n%QK_K == 0); + + const int kMaxQ = 8; + + const int64_t nbl = n/QK_K; + + block_iq3_s * y = vy; + + const int bs4 = block_size/4; + const int bs8 = block_size/8; + + for (int ibl = 0; ibl < nbl; ++ibl) { + + memset(&y[ibl], 0, sizeof(block_iq3_s)); + y[ibl].d = GGML_FP32_TO_FP16(0.f); + + uint8_t * qs = y[ibl].qs; + uint8_t * qh = y[ibl].qh; + uint8_t * signs = y[ibl].signs; + + float max_scale = 0; + + const float * xbl = x + QK_K*ibl; + float sumx2 = 0; + for (int i = 0; i < QK_K; ++i) sumx2 += xbl[i]*xbl[i]; + float sigma2 = 2*sumx2/QK_K; + + for (int ib = 0; ib < QK_K/block_size; ++ib) { + const float * xb = xbl + block_size*ib; + if (quant_weights) { + const float * qw = quant_weights + QK_K*ibl + block_size*ib; + for (int i = 0; i < block_size; ++i) weight[i] = qw[i] * sqrtf(sigma2 + xb[i]*xb[i]); + } else { + for (int i = 0; i < block_size; ++i) weight[i] = xb[i]*xb[i]; + } + for (int i = 0; i < block_size; ++i) waux[i] = sqrtf(weight[i]); + for (int k = 0; k < bs8; ++k) { + uint8_t s = 0; + for (int i = 0; i < 8; ++i) { + if (xb[8*k + i] >= 0) xval[8*k + i] = xb[8*k + i]; + else { + xval[8*k + i] = -xb[8*k + i]; s |= (1 << i); + } + } + block_signs[k] = s; + } + float max = xval[0]; + for (int i = 1; i < block_size; ++i) max = MAX(max, xval[i]); + memset(L, 0, block_size); + if (!max) { + scales[ib] = 0; + continue; + } + float best = 0; + float scale = max/(2*kMaxQ-1); + for (int k = 0; k < bs4; ++k) is_on_grid[k] = false; + for (int is = -9; is <= 9; ++is) { + float id = (2*kMaxQ-1+is*0.2f)/max; + float this_scale = 1/id; + for (int k = 0; k < bs4; ++k) { + for (int i = 0; i < 4; ++i) { + int l = nearest_int(0.5f*(id*xval[4*k+i]-1)); + Laux[4*k+i] = MAX(0, MIN(kMaxQ-1, l)); + } + uint16_t u = 0; + for (int i = 0; i < 4; ++i) u |= (Laux[4*k+i] << 3*i); + int grid_index = kmap_q3xs[u]; + is_on_grid_aux[k] = true; + if (grid_index < 0) { + is_on_grid_aux[k] = false; + const uint16_t * neighbours = kneighbors_q3xs - kmap_q3xs[u] - 1; + grid_index = iq3_find_best_neighbour(neighbours, kgrid_q3xs, xval + 4*k, waux + 4*k, this_scale, Laux + 4*k); + } + } + float sumqx = 0, sumq2 = 0; + for (int i = 0; i < block_size; ++i) { + float w = weight[i]; + float q = 2*Laux[i] + 1; + sumqx += w*xval[i]*q; + sumq2 += w*q*q; + } + if (sumq2 > 0 && sumqx*sumqx > best*sumq2) { + scale = sumqx/sumq2; best = scale*sumqx; + for (int i = 0; i < block_size; ++i) L[i] = Laux[i]; + for (int k = 0; k < bs4; ++k) is_on_grid[k] = is_on_grid_aux[k]; + } + } + int n_not_ongrid = 0; + for (int k = 0; k < bs4; ++k) if (!is_on_grid[k]) ++n_not_ongrid; + if (n_not_ongrid > 0 && scale > 0) { + float id = 1/scale; + for (int k = 0; k < bs4; ++k) { + //if (is_on_grid[k]) continue; + uint16_t u = 0; + for (int i = 0; i < 4; ++i) { + int l = nearest_int(0.5f*(id*xval[4*k+i]-1)); + l = MAX(0, MIN(kMaxQ-1, l)); + u |= (l << 3*i); + } + int grid_index = kmap_q3xs[u]; + if (grid_index < 0) { + const uint16_t * neighbours = kneighbors_q3xs - kmap_q3xs[u] - 1; + grid_index = iq3_find_best_neighbour(neighbours, kgrid_q3xs, xval + 4*k, waux + 4*k, scale, L + 4*k); + } + const int8_t * pg = (const int8_t *)(kgrid_q3xs + grid_index); + for (int i = 0; i < 4; ++i) L[4*k+i] = (pg[i] - 1)/2; + } + float sumqx = 0, sumq2 = 0; + for (int i = 0; i < block_size; ++i) { + float w = weight[i]; + float q = 2*L[i] + 1; + sumqx += w*xval[i]*q; + sumq2 += w*q*q; + } + if (sumq2 > 0) scale = sumqx/sumq2; + } + if (scale < 0) { + // This should never happen, but just in case, flip scale so that it is positive (we use uint's to encode the scale) + // and correspondingly flip quant signs. + scale = -scale; + for (int k = 0; k < bs8; ++k) block_signs[k] = ~block_signs[k]; + } + for (int k = 0; k < bs4; ++k) { + uint16_t u = 0; + for (int i = 0; i < 4; ++i) u |= (L[4*k+i] << 3*i); + int grid_index = kmap_q3xs[u]; + if (grid_index < 0) { + printf("Oops: found point %u not on grid:", u); + for (int i = 0; i < 4; ++i) printf(" %d", L[4*k+i]); + printf("\n"); + GGML_ABORT("fatal error"); + } + qs[k] = grid_index & 255; + qh[(ib*bs4+k)/8] |= ((grid_index >> 8) << ((ib*bs4+k)%8)); + } + qs += bs4; + for (int k = 0; k < bs8; ++k) signs[k] = block_signs[k]; + signs += bs8; + GGML_ASSERT(scale >= 0); + scales[ib] = scale; + max_scale = MAX(max_scale, scale); + } + + if (!max_scale) { + continue; + } + + float d = max_scale/31; + y[ibl].d = GGML_FP32_TO_FP16(d * 1.033f); + float id = 1/d; + for (int ib = 0; ib < QK_K/block_size; ib += 2) { + int l1 = nearest_int(0.5f*(id*scales[ib+0]-1)); + l1 = MAX(0, MIN(15, l1)); + int l2 = nearest_int(0.5f*(id*scales[ib+1]-1)); + l2 = MAX(0, MIN(15, l2)); + y[ibl].scales[ib/2] = l1 | (l2 << 4); + } + + } +} + +#define IQ3S_BLOCK_SIZE 32 +size_t quantize_iq3_s(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) { + GGML_ASSERT(n_per_row%QK_K == 0); + int64_t nblock = n_per_row/QK_K; + float scales[QK_K/IQ3S_BLOCK_SIZE]; + float weight[IQ3S_BLOCK_SIZE]; + float xval[IQ3S_BLOCK_SIZE]; + int8_t L[IQ3S_BLOCK_SIZE]; + int8_t Laux[IQ3S_BLOCK_SIZE]; + float waux[IQ3S_BLOCK_SIZE]; + bool is_on_grid[IQ3S_BLOCK_SIZE/4]; + bool is_on_grid_aux[IQ3S_BLOCK_SIZE/4]; + uint8_t block_signs[IQ3S_BLOCK_SIZE/8]; + char * qrow = (char *)dst; + for (int64_t row = 0; row < nrow; ++row) { + quantize_row_iq3_s_impl(IQ3S_BLOCK_SIZE, src, qrow, n_per_row, quant_weights, + scales, weight, xval, L, Laux, waux, is_on_grid, is_on_grid_aux, block_signs); + src += n_per_row; + qrow += nblock*sizeof(block_iq3_s); + } + return nrow * nblock * sizeof(block_iq3_s); +} + +void quantize_row_iq3_s_ref(const float * GGML_RESTRICT x, block_iq3_s * GGML_RESTRICT y, int64_t k) { + assert(k % QK_K == 0); + quantize_iq3_s(x, y, 1, k, NULL); +} + + +// =================================== 1.5 bpw =================================================== + +static int iq1_find_best_neighbour(const uint16_t * GGML_RESTRICT neighbours, const uint64_t * GGML_RESTRICT grid, + const float * GGML_RESTRICT xval, const float * GGML_RESTRICT weight, float * scale, int8_t * GGML_RESTRICT L, int ngrid) { + int num_neighbors = neighbours[0]; + GGML_ASSERT(num_neighbors > 0); + float best_score = -FLT_MAX; + int grid_index = -1; + for (int j = 1; j <= num_neighbors; ++j) { + const int8_t * pg = (const int8_t *)(grid + neighbours[j]); + float sumqx = 0, sumq2 = 0; + for (int i = 0; i < 8; ++i) { + float q = (pg[i] - 3)/2; + float w = weight[i]; + sumqx += w*q*xval[i]; + sumq2 += w*q*q; + } + if (sumqx > 0 && sumq2 > 0 && sumqx*sumqx > best_score*sumq2) { + *scale = sumqx/sumq2; best_score = *scale * sumqx; + grid_index = neighbours[j]; + } + } + if (grid_index < 0) { + for (int i = 0; i < ngrid; ++i) { + const int8_t * grid_i = (const int8_t *)(grid + i); + float sumqx = 0, sumq2 = 0; + for (int j = 0; j < 8; ++j) { + float w = weight[j]; + float q = (grid_i[j] - 3)/2; + sumqx += w*q*xval[j]; + sumq2 += w*q*q; + } + if (sumqx > 0 && sumq2 > 0 && sumqx*sumqx > best_score*sumq2) { + *scale = sumqx/sumq2; best_score = *scale*sumqx; + grid_index = i; + } + } + } + if (grid_index < 0) { + printf("Oops, did not find grid point\n"); + printf("Have %d neighbours\n", num_neighbors); + for (int j = 1; j <= num_neighbors; ++j) { + const int8_t * pg = (const int8_t *)(grid + neighbours[j]); + float sumqx = 0, sumq2 = 0; + for (int i = 0; i < 8; ++i) { + float q = (pg[i] - 3)/2; + float w = weight[i]; + sumqx += w*q*xval[i]; + sumq2 += w*q*q; + } + printf(" neighbour %d: sumqx = %g sumq2 = %g\n", j, (double)sumqx, (double)sumq2); + } + } + GGML_ASSERT(grid_index >= 0); + //!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!! + *scale *= 1.05f; // This is a fudge factor. Don't ask me why it improves the result. + //!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!! + const int8_t * pg = (const int8_t *)(grid + grid_index); + for (int i = 0; i < 8; ++i) L[i] = (pg[i] - 1)/2; + return grid_index; +} + +static int iq1_find_best_neighbour2(const uint16_t * GGML_RESTRICT neighbours, const uint64_t * GGML_RESTRICT grid, + const float * GGML_RESTRICT xval, const float * GGML_RESTRICT weight, float scale, const float * GGML_RESTRICT xg, int8_t * GGML_RESTRICT L, int ngrid) { + int num_neighbors = neighbours[0]; + GGML_ASSERT(num_neighbors > 0); + float best_score = FLT_MAX; + int grid_index = -1; + for (int j = 1; j <= num_neighbors; ++j) { + const int8_t * pg = (const int8_t *)(grid + neighbours[j]); + float d2 = 0; + for (int i = 0; i < 8; ++i) { + float q = xg[(pg[i] - 1)/2]; + float w = weight[i]; + float diff = scale*q - xval[i]; + d2 += w*diff*diff; + } + if (d2 < best_score) { + best_score = d2; + grid_index = neighbours[j]; + } + } + if (grid_index < 0) { + for (int i = 0; i < ngrid; ++i) { + const int8_t * grid_i = (const int8_t *)(grid + i); + float d2 = 0; + for (int j = 0; j < 8; ++j) { + float w = weight[j]; + float q = xg[(grid_i[j] - 1)/2]; + float diff = scale*q - xval[i]; + d2 += w*diff*diff; + } + if (d2 < best_score) { + best_score = d2; + grid_index = i; + } + } + } + if (grid_index < 0) { + printf("Oops, did not find grid point\n"); + printf("Have %d neighbours\n", num_neighbors); + for (int j = 1; j <= num_neighbors; ++j) { + const int8_t * pg = (const int8_t *)(grid + neighbours[j]); + float sumqx = 0, sumq2 = 0; + for (int i = 0; i < 8; ++i) { + float q = xg[(pg[i] - 1)/2]; + float w = weight[i]; + sumqx += w*q*xval[i]; + sumq2 += w*q*q; + } + printf(" neighbour %d: sumqx = %g sumq2 = %g\n", j, (double)sumqx, (double)sumq2); + } + } + GGML_ASSERT(grid_index >= 0); + const int8_t * pg = (const int8_t *)(grid + grid_index); + for (int i = 0; i < 8; ++i) L[i] = (pg[i] - 1)/2; + return grid_index; +} + +static int iq1_sort_helper(const void * left, const void * right) { + const float * l = left; + const float * r = right; + return *l < *r ? -1 : *l > *r ? 1 : 0; +} + +#define IQ1S_BLOCK_SIZE 32 +#define IQ1M_BLOCK_SIZE 16 +static void quantize_row_iq1_s_impl(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t n, const float * GGML_RESTRICT quant_weights, + float * scales, + float * weight, + float * sumx, + float * sumw, + float * pairs, + int8_t * L, + uint16_t * index, + int8_t * shifts) { + + const int gindex = iq2_data_index(GGML_TYPE_IQ1_S); + + const uint64_t * kgrid_q2xs = iq2_data[gindex].grid; + const int * kmap_q2xs = iq2_data[gindex].map; + const uint16_t * kneighbors_q2xs = iq2_data[gindex].neighbours; + + GGML_ASSERT(quant_weights && "missing quantization weights"); + GGML_ASSERT(kgrid_q2xs && "forgot to call ggml_quantize_init()?"); + GGML_ASSERT(kmap_q2xs && "forgot to call ggml_quantize_init()?"); + GGML_ASSERT(kneighbors_q2xs && "forgot to call ggml_quantize_init()?"); + GGML_ASSERT(n%QK_K == 0); + + block_iq1_s * y = vy; + + const int64_t nbl = n/QK_K; + + const int block_size = IQ1S_BLOCK_SIZE; + + const float x_p[3] = {-1 + IQ1S_DELTA, IQ1S_DELTA, 1 + IQ1S_DELTA}; + const float x_m[3] = {-1 - IQ1S_DELTA, -IQ1S_DELTA, 1 - IQ1S_DELTA}; + + + int * idx = (int *)(pairs + 1); + + for (int ibl = 0; ibl < nbl; ++ibl) { + + y[ibl].d = GGML_FP32_TO_FP16(0.f); + memset(y[ibl].qs, 0, QK_K/8); + memset(y[ibl].qh, 0, QK_K/16); + + float max_scale = 0; + + const float * xbl = x + QK_K*ibl; + float sumx2 = 0; + for (int i = 0; i < QK_K; ++i) sumx2 += xbl[i]*xbl[i]; + float sigma2 = 2*sumx2/QK_K; + + for (int ib = 0; ib < QK_K/block_size; ++ib) { + const float * xb = xbl + block_size*ib; + const float * qw = quant_weights + QK_K*ibl + block_size*ib; + for (int i = 0; i < block_size; ++i) weight[i] = qw[i] * sqrtf(sigma2 + xb[i]*xb[i]); + float max = fabsf(xb[0]); + for (int i = 1; i < block_size; ++i) max = MAX(max, fabsf(xb[i])); + if (max < GROUP_MAX_EPS_IQ1_S) { + scales[ib] = 0; + shifts[ib] = 1; + memset(L, 1, block_size); + continue; + } + // Here we solve exactly the sum of squared difference (SSD) weighted minimization problem. + // With just 3 allowed quant values (-1, 0, 1), we can search exhaustively for the two + // boundaries that split the weights xb[i] into 3 groups. To do so, we sort the weights + // in ascending order, compute Si = sum[weight[j] xb[j], j = 0...i] and + // Wi = sum[weight[j], j = 0...i], and use these to quckly get get the optimum scale + // for each possible and score for each split. + for (int j = 0; j < block_size; ++j) { + pairs[2*j] = xb[j]; + idx[2*j] = j; + } + qsort(pairs, block_size, 2*sizeof(float), iq1_sort_helper); + { + sumx[0] = sumw[0] = 0; + for (int j = 0; j < block_size; ++j) { + int i = idx[2*j]; + sumx[j+1] = sumx[j] + weight[i]*xb[i]; + sumw[j+1] = sumw[j] + weight[i]; + } + } + float best_score = -FLT_MAX, scale = max; + int besti1 = -1, besti2 = -1, best_shift = 0; + for (int i1 = 0; i1 <= block_size; ++i1) { + for (int i2 = i1; i2 <= block_size; ++i2) { + float sumqx = (sumx[i1] - sumx[0])*x_p[0] + (sumx[i2] - sumx[i1])*x_p[1] + (sumx[block_size] - sumx[i2])*x_p[2]; + float sumq2 = (sumw[i1] - sumw[0])*x_p[0]*x_p[0] + (sumw[i2] - sumw[i1])*x_p[1]*x_p[1] + (sumw[block_size] - sumw[i2])*x_p[2]*x_p[2]; + if (sumq2 > 0 && sumqx*sumqx > best_score*sumq2) { + scale = sumqx/sumq2; best_score = scale*sumqx; + besti1 = i1; besti2 = i2; best_shift = 1; + } + sumqx = (sumx[i1] - sumx[0])*x_m[0] + (sumx[i2] - sumx[i1])*x_m[1] + (sumx[block_size] - sumx[i2])*x_m[2]; + sumq2 = (sumw[i1] - sumw[0])*x_m[0]*x_m[0] + (sumw[i2] - sumw[i1])*x_m[1]*x_m[1] + (sumw[block_size] - sumw[i2])*x_m[2]*x_m[2]; + if (sumq2 > 0 && sumqx*sumqx > best_score*sumq2) { + scale = sumqx/sumq2; best_score = scale*sumqx; + besti1 = i1; besti2 = i2; best_shift = -1; + } + } + } + if (besti1 < 0 || besti2 < 0 || best_shift == 0) { + scales[ib] = 0; + shifts[ib] = 1; + memset(L, 1, block_size); + continue; + } + for (int j = 0; j < besti1; ++j) L[idx[2*j]] = 0; + for (int j = besti1; j < besti2; ++j) L[idx[2*j]] = 1; + for (int j = besti2; j < block_size; ++j) L[idx[2*j]] = 2; + if (scale < 0) { + for (int j = 0; j < block_size; ++j) L[j] = 2 - L[j]; + scale = -scale; best_shift = -best_shift; + } + bool all_on_grid = true; + const float * xx = best_shift == 1 ? x_p : x_m; + for (int k = 0; k < block_size/8; ++k) { + uint16_t u = 0; + for (int j = 0; j < 8; ++j) u |= (L[8*k+j] << 2*j); + int grid_index = kmap_q2xs[u]; + if (grid_index < 0) { + all_on_grid = false; + const uint16_t * neighbours = kneighbors_q2xs - kmap_q2xs[u] - 1; + grid_index = iq1_find_best_neighbour2(neighbours, kgrid_q2xs, xb + 8*k, weight + 8*k, scale, xx, L + 8*k, NGRID_IQ1S); + GGML_ASSERT(grid_index >= 0); + } + index[k] = grid_index; + } + if (!all_on_grid) { + float sumqx = 0, sumq2 = 0; + for (int k = 0; k < block_size/8; ++k) { + const int8_t * pg = (const int8_t *)(kgrid_q2xs + index[k]); + for (int j = 0; j < 8; ++j) { + float w = weight[8*k + j]; + float q = xx[(pg[j] - 1)/2]; + sumqx += w*q*xb[8*k+j]; + sumq2 += w*q*q; + } + } + if (sumqx > 0 && sumq2 > 0) scale = sumqx/sumq2; + } + uint16_t h = 0; + for (int k = 0; k < block_size/8; ++k) { + y[ibl].qs[(block_size/8)*ib + k] = index[k] & 255; + h |= (index[k] >> 8) << 3*k; + } + y[ibl].qh[ib] = h; + GGML_ASSERT(scale >= 0); + scales[ib] = scale; + shifts[ib] = best_shift; + max_scale = MAX(max_scale, scale); + } + + if (!max_scale) { + continue; + } + + float d = max_scale/15; + y[ibl].d = GGML_FP32_TO_FP16(d*1.125f); // 1.125f is another fudge factor. Don't ask me why it is needed. + float id = 1/d; + for (int ib = 0; ib < QK_K/block_size; ++ib) { + int l = nearest_int(0.5f*(id*scales[ib]-1)); + l = MAX(0, MIN(7, l)); + if (shifts[ib] == -1) l |= 8; + y[ibl].qh[ib] |= (l << 12); + } + } +} + +size_t quantize_iq1_s(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) { + GGML_ASSERT(n_per_row%QK_K == 0); + float scales[QK_K/IQ1S_BLOCK_SIZE]; + float weight[IQ1S_BLOCK_SIZE]; + int8_t L[IQ1S_BLOCK_SIZE]; + float sumx[IQ1S_BLOCK_SIZE+1]; + float sumw[IQ1S_BLOCK_SIZE+1]; + float pairs[2*IQ1S_BLOCK_SIZE]; + uint16_t index[IQ1S_BLOCK_SIZE/8]; + int8_t shifts[QK_K/IQ1S_BLOCK_SIZE]; + int64_t nblock = n_per_row/QK_K; + char * qrow = (char *)dst; + for (int64_t row = 0; row < nrow; ++row) { + quantize_row_iq1_s_impl(src, qrow, n_per_row, quant_weights, scales, weight, sumx, sumw, pairs, L, index, shifts); + src += n_per_row; + qrow += nblock*sizeof(block_iq1_s); + } + return nrow * nblock * sizeof(block_iq1_s); +} + +static void quantize_row_iq1_m_impl(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t n, const float * GGML_RESTRICT quant_weights, + float * scales, + float * weight, + float * pairs, + int8_t * L, + uint16_t * index, + int8_t * shifts) { + + const int gindex = iq2_data_index(GGML_TYPE_IQ1_M); + + const uint64_t * kgrid_q2xs = iq2_data[gindex].grid; + const int * kmap_q2xs = iq2_data[gindex].map; + const uint16_t * kneighbors_q2xs = iq2_data[gindex].neighbours; + + //GGML_ASSERT(quant_weights && "missing quantization weights"); + GGML_ASSERT(kgrid_q2xs && "forgot to call ggml_quantize_init()?"); + GGML_ASSERT(kmap_q2xs && "forgot to call ggml_quantize_init()?"); + GGML_ASSERT(kneighbors_q2xs && "forgot to call ggml_quantize_init()?"); + GGML_ASSERT(n%QK_K == 0); + + block_iq1_m * y = vy; + + const int64_t nbl = n/QK_K; + + const int block_size = IQ1M_BLOCK_SIZE; + + const float x_p[3] = {-1 + IQ1M_DELTA, IQ1M_DELTA, 1 + IQ1M_DELTA}; + const float x_m[3] = {-1 - IQ1M_DELTA, -IQ1M_DELTA, 1 - IQ1M_DELTA}; + const uint8_t masks[4] = {0x00, 0x80, 0x08, 0x88}; + + int * idx = (int *)(pairs + 1); + + float sumqx[4], sumq2[4]; + + iq1m_scale_t s; + const float * xx; + + for (int ibl = 0; ibl < nbl; ++ibl) { + memset(y[ibl].qs, 0, QK_K/8); + memset(y[ibl].qh, 0, QK_K/16); + memset(y[ibl].scales, 0, QK_K/32); + + float max_scale = 0; + + const float * xbl = x + QK_K*ibl; + float sumx2 = 0; + for (int i = 0; i < QK_K; ++i) sumx2 += xbl[i]*xbl[i]; + float sigma2 = 2*sumx2/QK_K; + + for (int ib = 0; ib < QK_K/block_size; ++ib) { + const float * xb = xbl + block_size*ib; + if (quant_weights) { + const float * qw = quant_weights + QK_K*ibl + block_size*ib; + for (int i = 0; i < block_size; ++i) weight[i] = qw[i] * sqrtf(sigma2 + xb[i]*xb[i]); + } else { + for (int i = 0; i < block_size; ++i) weight[i] = xb[i]*xb[i]; + } + float max = fabsf(xb[0]); + for (int i = 1; i < block_size; ++i) max = MAX(max, fabsf(xb[i])); + if (max < GROUP_MAX_EPS_IQ1_M) { + scales[ib] = 0; + shifts[ib] = 0; + memset(L, 1, block_size); + continue; + } + // Here we solve exactly the sum of squared difference (SSD) weighted minimization problem. + // With just 3 allowed quant values (-1, 0, 1), we can search exhaustively for the two + // boundaries that split the weights xb[i] into 3 groups. To do so, we sort the weights + // in ascending order, compute Si = sum[weight[j] xb[j], j = 0...i] and + // Wi = sum[weight[j], j = 0...i], and use these to quckly get get the optimum scale + // for each possible and score for each split. + for (int j = 0; j < block_size; ++j) { + pairs[2*j] = xb[j]; + idx[2*j] = j; + } + qsort(pairs, block_size, 2*sizeof(float), iq1_sort_helper); + float best_score = -FLT_MAX, scale = max; + int besti1 = -1, besti2 = -1, best_k = -1; + // 0: +, + + // 1: +, - + // 2: -, + + // 3: -, - + // prefix sums are kept per half of the block because each half can use a different sign (x_p or x_m) + // since v[0]-v[1] = v[1]-v[2] = -1 for both x_p and x_m, the 3-group sum for a split collapses to T*v[2] - px[i1] - px[i2] + { + float px[2][IQ1M_BLOCK_SIZE+1]; + float pw[2][IQ1M_BLOCK_SIZE+1]; + px[0][0] = px[1][0] = 0; + pw[0][0] = pw[1][0] = 0; + for (int j = 0; j < block_size; ++j) { + const int i = idx[2*j]; + const int h = i < block_size/2 ? 0 : 1; + px[h][j+1] = px[h][j] + weight[i]*xb[i]; + px[1-h][j+1] = px[1-h][j]; + pw[h][j+1] = pw[h][j] + weight[i]; + pw[1-h][j+1] = pw[1-h][j]; + } + const float txs[2] = {px[0][block_size], px[1][block_size]}; // total weight*x per half + const float tws[2] = {pw[0][block_size], pw[1][block_size]}; // total weight per half + const float p2 = x_p[2], m2 = x_m[2]; + const float cp1 = x_p[0]*x_p[0] - x_p[1]*x_p[1]; + const float cp2 = x_p[1]*x_p[1] - x_p[2]*x_p[2]; + const float cm1 = x_m[0]*x_m[0] - x_m[1]*x_m[1]; + const float cm2 = x_m[1]*x_m[1] - x_m[2]*x_m[2]; + for (int i1 = 0; i1 <= block_size; ++i1) { + for (int i2 = i1; i2 <= block_size; ++i2) { + float qx_p[2], qx_m[2], q2_p[2], q2_m[2]; + for (int h = 0; h < 2; ++h) { + const float sx = px[h][i1] + px[h][i2]; + qx_p[h] = txs[h]*p2 - sx; + qx_m[h] = txs[h]*m2 - sx; + q2_p[h] = tws[h]*p2*p2 + pw[h][i1]*cp1 + pw[h][i2]*cp2; + q2_m[h] = tws[h]*m2*m2 + pw[h][i1]*cm1 + pw[h][i2]*cm2; + } + sumqx[0] = qx_p[0] + qx_p[1]; + sumqx[1] = qx_p[0] + qx_m[1]; + sumqx[2] = qx_m[0] + qx_p[1]; + sumqx[3] = qx_m[0] + qx_m[1]; + sumq2[0] = q2_p[0] + q2_p[1]; + sumq2[1] = q2_p[0] + q2_m[1]; + sumq2[2] = q2_m[0] + q2_p[1]; + sumq2[3] = q2_m[0] + q2_m[1]; + for (int k = 0; k < 4; ++k) { + if (sumq2[k] > 0 && sumqx[k]*sumqx[k] > best_score*sumq2[k]) { + scale = sumqx[k]/sumq2[k]; best_score = scale*sumqx[k]; + besti1 = i1; besti2 = i2; best_k = k; + } + } + } + } + } + if (besti1 < 0 || besti2 < 0 || best_k < 0) { + scales[ib] = 0; + shifts[ib] = 0; + memset(L, 1, block_size); + continue; + } + for (int j = 0; j < besti1; ++j) L[idx[2*j]] = 0; + for (int j = besti1; j < besti2; ++j) L[idx[2*j]] = 1; + for (int j = besti2; j < block_size; ++j) L[idx[2*j]] = 2; + if (scale < 0) { + for (int j = 0; j < block_size; ++j) L[j] = 2 - L[j]; + scale = -scale; + best_k = best_k == 0 ? 3 : best_k == 1 ? 2 : best_k == 2 ? 1 : 0; + } + bool all_on_grid = true; + for (int k = 0; k < block_size/8; ++k) { + if (k == 0) xx = best_k < 2 ? x_p : x_m; + else xx = best_k%2 == 0 ? x_p : x_m; + uint16_t u = 0; + for (int j = 0; j < 8; ++j) u |= (L[8*k+j] << 2*j); + int grid_index = kmap_q2xs[u]; + if (grid_index < 0) { + all_on_grid = false; + const uint16_t * neighbours = kneighbors_q2xs - kmap_q2xs[u] - 1; + grid_index = iq1_find_best_neighbour2(neighbours, kgrid_q2xs, xb + 8*k, weight + 8*k, scale, xx, L + 8*k, NGRID_IQ1S); + GGML_ASSERT(grid_index >= 0); + } + index[k] = grid_index; + } + if (!all_on_grid) { + float sumqx_f = 0, sumq2_f = 0; + for (int k = 0; k < block_size/8; ++k) { + if (k == 0) xx = best_k < 2 ? x_p : x_m; + else xx = best_k%2 == 0 ? x_p : x_m; + const int8_t * pg = (const int8_t *)(kgrid_q2xs + index[k]); + for (int j = 0; j < 8; ++j) { + float w = weight[8*k + j]; + float q = xx[(pg[j] - 1)/2]; + sumqx_f += w*q*xb[8*k+j]; + sumq2_f += w*q*q; + } + } + if (sumqx_f > 0 && sumq2_f > 0) scale = sumqx_f/sumq2_f; + } + y[ibl].qs[2*ib + 0] = index[0] & 255; + y[ibl].qs[2*ib + 1] = index[1] & 255; + y[ibl].qh[ib] = (index[0] >> 8) | ((index[1] >> 8) << 4); + GGML_ASSERT(scale >= 0); + scales[ib] = scale; + shifts[ib] = best_k; + max_scale = MAX(max_scale, scale); + } + + if (!max_scale) { + continue; + } + + uint16_t * sc = (uint16_t *)y[ibl].scales; + float d = max_scale/15; + float id = 1/d; + float sumqx_f = 0, sumq2_f = 0; + for (int ib = 0; ib < QK_K/block_size; ++ib) { + int l = nearest_int(0.5f*(id*scales[ib+0]-1)); + l = MAX(0, MIN(7, l)); + sc[ib/4] |= (l << 3*(ib%4)); + y[ibl].qh[ib] |= masks[shifts[ib]]; + const float * xb = xbl + block_size*ib; + if (quant_weights) { + const float * qw = quant_weights + QK_K*ibl + block_size*ib; + for (int i = 0; i < block_size; ++i) weight[i] = qw[i] * sqrtf(sigma2 + xb[i]*xb[i]); + } else { + for (int i = 0; i < block_size; ++i) weight[i] = xb[i]*xb[i]; + } + for (int k = 0; k < block_size/8; ++k) { + if (k == 0) xx = shifts[ib] < 2 ? x_p : x_m; + else xx = shifts[ib]%2 == 0 ? x_p : x_m; + const int8_t * pg = (const int8_t *)(kgrid_q2xs + y[ibl].qs[2*ib+k] + ((y[ibl].qh[ib] << (8 - 4*k)) & 0x700)); + for (int j = 0; j < 8; ++j) { + float w = weight[8*k + j]; + float q = xx[(pg[j] - 1)/2]*(2*l+1); + sumqx_f += w*q*xb[8*k+j]; + sumq2_f += w*q*q; + } + } + } + if (sumq2_f > 0) d = sumqx_f/sumq2_f; + s.f16 = GGML_FP32_TO_FP16(d*1.1125f); // 1.1125f is another fudge factor. Don't ask me why it is needed. + sc[0] |= ((s.u16 & 0x000f) << 12); + sc[1] |= ((s.u16 & 0x00f0) << 8); + sc[2] |= ((s.u16 & 0x0f00) << 4); + sc[3] |= ((s.u16 & 0xf000) << 0); + } +} + +size_t quantize_iq1_m(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) { + GGML_ASSERT(n_per_row%QK_K == 0); + float scales[QK_K/IQ1M_BLOCK_SIZE]; + float weight[IQ1M_BLOCK_SIZE]; + int8_t L[IQ1M_BLOCK_SIZE]; + float pairs[2*IQ1M_BLOCK_SIZE]; + uint16_t index[IQ1M_BLOCK_SIZE/8]; + int8_t shifts[QK_K/IQ1M_BLOCK_SIZE]; + int64_t nblock = n_per_row/QK_K; + char * qrow = (char *)dst; + for (int64_t row = 0; row < nrow; ++row) { + quantize_row_iq1_m_impl(src, qrow, n_per_row, quant_weights, scales, weight, pairs, L, index, shifts); + src += n_per_row; + qrow += nblock*sizeof(block_iq1_m); + } + return nrow * nblock * sizeof(block_iq1_m); +} + +// ============================ 4-bit non-linear quants + +static void quantize_row_iq4_nl_impl(const int super_block_size, const int block_size, const float * GGML_RESTRICT x, + ggml_fp16_t * dh, uint8_t * q4, uint16_t * scales_h, uint8_t * scales_l, + float * scales, float * weight, uint8_t * L, + const int8_t * values, + const float * quant_weights, + const int ntry) { + + float sigma2 = 0; + for (int j = 0; j < super_block_size; ++j) sigma2 += x[j]*x[j]; + sigma2 *= 2.f/super_block_size; + + memset(q4, 0, super_block_size/2); + dh[0] = GGML_FP32_TO_FP16(0.f); + + float max_scale = 0, amax_scale = 0; + for (int ib = 0; ib < super_block_size/block_size; ++ib) { + const float * xb = x + ib*block_size; + uint8_t * Lb = L + ib*block_size; + if (quant_weights) { + const float * qw = quant_weights + ib*block_size; + for (int j = 0; j < block_size; ++j) weight[j] = qw[j] * sqrtf(sigma2 + xb[j]*xb[j]); + } else { + for (int j = 0; j < block_size; ++j) weight[j] = xb[j]*xb[j]; + } + float amax = 0, max = 0; + for (int j = 0; j < block_size; ++j) { + float ax = fabsf(xb[j]); + if (ax > amax) { + amax = ax; max = xb[j]; + } + } + if (amax < GROUP_MAX_EPS) { + scales[ib] = 0; + continue; + } + float d = ntry > 0 ? -max/values[0] : max/values[0]; + float id = 1/d; + float sumqx = 0, sumq2 = 0; + for (int j = 0; j < block_size; ++j) { + float al = id*xb[j]; + int l = best_index_int8(16, values, al); + Lb[j] = l; + float q = values[l]; + float w = weight[j]; + sumqx += w*q*xb[j]; + sumq2 += w*q*q; + } + d = sumq2 > 0 ? sumqx/sumq2 : 0.f; + float best = d*sumqx; + for (int itry = -ntry; itry <= ntry; ++itry) { + id = (itry + values[0])/max; + sumqx = sumq2 = 0; + for (int j = 0; j < block_size; ++j) { + float al = id*xb[j]; + int l = best_index_int8(16, values, al); + float q = values[l]; + float w = weight[j]; + sumqx += w*q*xb[j]; + sumq2 += w*q*q; + } + if (sumq2 > 0 && sumqx*sumqx > best*sumq2) { + d = sumqx/sumq2; best = d * sumqx; + } + } + scales[ib] = d; + float abs_d = fabsf(d); + if (abs_d > amax_scale) { + amax_scale = abs_d; max_scale = d; + } + } + + if (super_block_size/block_size > 1) { + int nb = super_block_size/block_size; + memset(scales_h, 0, ((nb+7)/8)*sizeof(uint16_t)); + float d = -max_scale/32; + dh[0] = GGML_FP32_TO_FP16(d); + float id = d ? 1/d : 0.f; + for (int ib = 0; ib < super_block_size/block_size; ++ib) { + int l = nearest_int(id*scales[ib]); + l = MAX(-32, MIN(31, l)); + float dl = d * l; + float idl = dl ? 1/dl : 0.f; + uint8_t * Lb = L + ib*block_size; + const float * xb = x + ib*block_size; + for (int j = 0; j < block_size; ++j) { + Lb[j] = best_index_int8(16, values, idl*xb[j]); + } + l += 32; + uint8_t l_l = l & 0xf; + uint8_t l_h = l >> 4; + if (ib%2 == 0) scales_l[ib/2] = l_l; + else scales_l[ib/2] |= (l_l << 4); + scales_h[ib/8] |= (l_h << 2*(ib%8)); + } + } else { + dh[0] = GGML_FP32_TO_FP16(scales[0]); + if (ntry > 0) { + float id = scales[0] ? 1/scales[0] : 0; + for (int j = 0; j < super_block_size; ++j) { + L[j] = best_index_int8(16, values, id*x[j]); + } + } + } + + for (int i = 0; i < super_block_size/32; ++i) { + for (int j = 0; j < 16; ++j) { + q4[16*i + j] = L[32*i + j] | (L[32*i + 16 + j] << 4); + } + } +} + +size_t quantize_iq4_nl(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) { + GGML_ASSERT(n_per_row%QK4_NL == 0); + int64_t nblock = n_per_row/QK4_NL; + char * qrow = (char *)dst; + uint8_t L[QK4_NL]; + float weight[QK4_NL]; + uint16_t unused_h; + uint8_t * unused_l = NULL; + float scale; + for (int64_t row = 0; row < nrow; ++row) { + block_iq4_nl * iq4 = (block_iq4_nl *)qrow; + for (int ibl = 0; ibl < nblock; ++ibl) { + const float * qw = quant_weights ? quant_weights + QK4_NL*ibl : NULL; + quantize_row_iq4_nl_impl(QK4_NL, 32, src + QK4_NL*ibl, &iq4[ibl].d, iq4[ibl].qs, &unused_h, unused_l, + &scale, weight, L, kvalues_iq4nl, qw, 7); + } + src += n_per_row; + qrow += nblock*sizeof(block_iq4_nl); + } + return nrow * nblock * sizeof(block_iq4_nl); +} + +//void quantize_row_iq4_nl_ref(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k) { +void quantize_row_iq4_nl_ref(const float * GGML_RESTRICT x, block_iq4_nl * GGML_RESTRICT y, int64_t k) { + GGML_ASSERT(k%QK4_NL == 0); + int64_t nblock = k/QK4_NL; + uint8_t L[QK4_NL]; + float weight[QK4_NL]; + uint16_t unused_h; + uint8_t * unused_l = NULL; + float scale; + block_iq4_nl * iq4 = y; + for (int ibl = 0; ibl < nblock; ++ibl) { + quantize_row_iq4_nl_impl(QK4_NL, 32, x + QK4_NL*ibl, &iq4[ibl].d, iq4[ibl].qs, &unused_h, unused_l, + &scale, weight, L, kvalues_iq4nl, NULL, -1); + } +} + +size_t quantize_iq4_xs(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) { + GGML_ASSERT(n_per_row%QK_K == 0); + int64_t nblock = n_per_row/QK_K; + char * qrow = (char *)dst; + uint8_t L[QK_K]; + float weight[32]; + float scales[QK_K/32]; + for (int64_t row = 0; row < nrow; ++row) { + block_iq4_xs * iq4 = (block_iq4_xs *)qrow; + for (int ibl = 0; ibl < nblock; ++ibl) { + const float * qw = quant_weights ? quant_weights + QK_K*ibl : NULL; + quantize_row_iq4_nl_impl(QK_K, 32, src + QK_K*ibl, &iq4[ibl].d, iq4[ibl].qs, &iq4[ibl].scales_h, iq4[ibl].scales_l, + scales, weight, L, kvalues_iq4nl, qw, 7); + } + src += n_per_row; + qrow += nblock*sizeof(block_iq4_xs); + } + return nrow * nblock * sizeof(block_iq4_xs); +} + +void quantize_row_iq4_xs_ref(const float * GGML_RESTRICT x, block_iq4_xs * GGML_RESTRICT y, int64_t k) { + assert(k % QK_K == 0); + quantize_iq4_xs(x, y, 1, k, NULL); +} + +// =============================== 2.5625 bpw + +static void quantize_row_iq2_s_impl(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t n, const float * GGML_RESTRICT quant_weights) { + + const int gindex = iq2_data_index(GGML_TYPE_IQ2_S); + + const uint64_t * kgrid_q2xs = iq2_data[gindex].grid; + const int * kmap_q2xs = iq2_data[gindex].map; + const uint16_t * kneighbors_q2xs = iq2_data[gindex].neighbours; + + GGML_ASSERT(kmap_q2xs && "forgot to call ggml_quantize_init()?"); + GGML_ASSERT(kgrid_q2xs && "forgot to call ggml_quantize_init()?"); + GGML_ASSERT(kneighbors_q2xs && "forgot to call ggml_quantize_init()?"); + GGML_ASSERT(n%QK_K == 0); + + const int kMaxQ = 3; + + const int64_t nbl = n/QK_K; + + block_iq2_s * y = vy; + + float scales[QK_K/16]; + float weight[16]; + float xval[16]; + int8_t L[16]; + int8_t Laux[16]; + float waux[16]; + bool is_on_grid[2]; + bool is_on_grid_aux[2]; + uint8_t block_signs[2]; + + for (int ibl = 0; ibl < nbl; ++ibl) { + + memset(&y[ibl], 0, sizeof(block_iq2_s)); + y[ibl].d = GGML_FP32_TO_FP16(0.f); + + float max_scale = 0; + + const float * xbl = x + QK_K*ibl; + float sumx2 = 0; + for (int i = 0; i < QK_K; ++i) sumx2 += xbl[i]*xbl[i]; + float sigma2 = 2*sumx2/QK_K; + + for (int ib = 0; ib < QK_K/16; ++ib) { + const float * xb = xbl + 16*ib; + if (quant_weights) { + const float * qw = quant_weights + QK_K*ibl + 16*ib; + for (int i = 0; i < 16; ++i) weight[i] = qw[i] * sqrtf(sigma2 + xb[i]*xb[i]); + } else { + for (int i = 0; i < 16; ++i) weight[i] = 0.25f*sigma2 + xb[i]*xb[i]; + } + for (int i = 0; i < 16; ++i) waux[i] = sqrtf(weight[i]); + for (int k = 0; k < 2; ++k) { + uint8_t s = 0; + for (int i = 0; i < 8; ++i) { + if (xb[8*k + i] >= 0) xval[8*k + i] = xb[8*k + i]; + else { + xval[8*k + i] = -xb[8*k + i]; s |= (1 << i); + } + } + block_signs[k] = s; + } + float max = xval[0]; + for (int i = 1; i < 16; ++i) max = MAX(max, xval[i]); + memset(L, 0, 16); + if (max < GROUP_MAX_EPS_IQ2_S) { + scales[ib] = 0; + continue; + } + float best = 0; + float scale = max/(2*kMaxQ-1); + is_on_grid[0] = is_on_grid[1] = true; + for (int is = -9; is <= 9; ++is) { + float id = (2*kMaxQ-1+is*0.1f)/max; + float this_scale = 1/id; + for (int k = 0; k < 2; ++k) { + for (int i = 0; i < 8; ++i) { + int l = nearest_int(0.5f*(id*xval[8*k+i]-1)); + Laux[8*k+i] = MAX(0, MIN(kMaxQ-1, l)); + } + uint16_t u = 0; + for (int i = 0; i < 8; ++i) u |= (Laux[8*k+i] << 2*i); + int grid_index = kmap_q2xs[u]; + is_on_grid_aux[k] = true; + if (grid_index < 0) { + is_on_grid_aux[k] = false; + const uint16_t * neighbours = kneighbors_q2xs - kmap_q2xs[u] - 1; + grid_index = iq2_find_best_neighbour(neighbours, kgrid_q2xs, xval + 8*k, waux + 8*k, this_scale, Laux + 8*k); + } + } + float sumqx = 0, sumq2 = 0; + for (int i = 0; i < 16; ++i) { + float w = weight[i]; + float q = 2*Laux[i] + 1; + sumqx += w*xval[i]*q; + sumq2 += w*q*q; + } + if (sumq2 > 0 && sumqx*sumqx > best*sumq2) { + scale = sumqx/sumq2; best = scale*sumqx; + for (int i = 0; i < 16; ++i) L[i] = Laux[i]; + for (int k = 0; k < 2; ++k) is_on_grid[k] = is_on_grid_aux[k]; + } + } + int n_not_ongrid = 0; + for (int k = 0; k < 2; ++k) if (!is_on_grid[k]) ++n_not_ongrid; + if (n_not_ongrid > 0 && scale > 0) { + float id = 1/scale; + for (int k = 0; k < 2; ++k) { + if (is_on_grid[k]) continue; + uint16_t u = 0; + for (int i = 0; i < 8; ++i) { + int l = nearest_int(0.5f*(id*xval[8*k+i]-1)); + l = MAX(0, MIN(kMaxQ-1, l)); + u |= (l << 2*i); + L[8*k + i] = l; + } + int grid_index = kmap_q2xs[u]; + if (grid_index < 0) { + const uint16_t * neighbours = kneighbors_q2xs - kmap_q2xs[u] - 1; + grid_index = iq2_find_best_neighbour(neighbours, kgrid_q2xs, xval + 8*k, waux + 8*k, scale, L + 8*k); + } + } + float sumqx = 0, sumq2 = 0; + for (int i = 0; i < 16; ++i) { + float w = weight[i]; + float q = 2*L[i] + 1; + sumqx += w*xval[i]*q; + sumq2 += w*q*q; + } + if (sumq2 > 0) scale = sumqx/sumq2; + } + if (scale < 0) { + scale = -scale; + for (int k = 0; k < 2; ++k) block_signs[k] = ~block_signs[k]; + } + for (int k = 0; k < 2; ++k) { + uint16_t u = 0; + for (int i = 0; i < 8; ++i) u |= (L[8*k+i] << 2*i); + int grid_index = kmap_q2xs[u]; + if (grid_index < 0) { + printf("Oops: found point %u not on grid:", u); + for (int i = 0; i < 8; ++i) printf(" %d", L[8*k+i]); + printf("\n"); + GGML_ABORT("fatal error"); + } + const int i8 = 2*ib + k; + y[ibl].qs[i8] = grid_index & 255; + y[ibl].qh[i8/4] |= ((grid_index >> 8) << 2*(i8%4)); + y[ibl].qs[QK_K/8 + i8] = block_signs[k]; + } + GGML_ASSERT(scale >= 0); + scales[ib] = scale; + max_scale = MAX(max_scale, scale); + } + + if (!max_scale) { + continue; + } + + float d = max_scale/31; + y[ibl].d = GGML_FP32_TO_FP16(d * 0.9875f); + float id = 1/d; + for (int ib = 0; ib < QK_K/16; ++ib) { + int l = nearest_int(0.5f*(id*scales[ib]-1)); + l = MAX(0, MIN(15, l)); + if (ib%2 == 0) y[ibl].scales[ib/2] = l; + else y[ibl].scales[ib/2] |= (l << 4); + } + } +} + +size_t quantize_iq2_s(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) { + GGML_ASSERT(n_per_row%QK_K == 0); + int64_t nblock = n_per_row/QK_K; + char * qrow = (char *)dst; + for (int64_t row = 0; row < nrow; ++row) { + quantize_row_iq2_s_impl(src, qrow, n_per_row, quant_weights); + src += n_per_row; + qrow += nblock*sizeof(block_iq2_s); + } + return nrow * nblock * sizeof(block_iq2_s); +} + +void quantize_row_iq2_s_ref(const float * GGML_RESTRICT x, block_iq2_s * GGML_RESTRICT y, int64_t k) { + assert(k % QK_K == 0); + quantize_iq2_s(x, y, 1, k, NULL); +} + +// =============================== data validation + +static bool validate_float(float f, size_t i) { + if (isinf(f)) { + fprintf(stderr, "ggml_validate_row_data: found inf value at block %zu\n", i); + return false; + } + + if (isnan(f)) { + fprintf(stderr, "ggml_validate_row_data: found nan value at block %zu\n", i); + return false; + } + + return true; +} + +static bool isinf_fp16(ggml_fp16_t f) { + return (f & 0x7c00) == 0x7c00 && (f & 0x03ff) == 0; +} + +static bool isnan_fp16(ggml_fp16_t f) { + return (f & 0x7c00) == 0x7c00 && (f & 0x03ff) != 0; +} + +static bool validate_fp16(ggml_fp16_t f, size_t i) { + if (isinf_fp16(f)) { + fprintf(stderr, "ggml_validate_row_data: found inf value at block %zu\n", i); + return false; + } + + if (isnan_fp16(f)) { + fprintf(stderr, "ggml_validate_row_data: found nan value at block %zu\n", i); + return false; + } + + return true; +} + +static bool validate_e_e8m0(uint8_t e, size_t i) { + if (e == 0xff) { + fprintf(stderr, "ggml_validate_row_data: found invalid e value %d at block %zu\n", e, i); + return false; + } + + return true; +} + +#define VALIDATE_ROW_DATA_D_F16_IMPL(type, data, nb) \ + const type * q = (const type *) (data); \ + for (size_t i = 0; i < (nb); ++i) { \ + if (!validate_fp16(q[i].d, i)) { \ + return false; \ + } \ + } + +#define VALIDATE_ROW_DATA_DM_F16_IMPL(type, data, nb, d, m) \ + const type * q = (const type *) (data); \ + for (size_t i = 0; i < (nb); ++i) { \ + if (!validate_fp16(q[i].d, i) || !validate_fp16(q[i].m, i)) { \ + return false; \ + } \ + } + +#define VALIDATE_ROW_DATA_E_E8M0_IMPL(type, data, nb) \ + const type * q = (const type *) (data); \ + for (size_t i = 0; i < (nb); ++i) { \ + if (!validate_e_e8m0(q[i].e, i)) { \ + return false; \ + } \ + } + +#define VALIDATE_ROW_DATA_DVEC_F16_IMPL(type, data, nb, nr) \ + const type * q = (const type *) (data); \ + for (size_t i = 0; i < (nb); ++i) { \ + for (size_t j = 0; j < (nr); ++j) { \ + if (!validate_fp16(q[i].d[j], i)) { \ + return false; \ + } \ + } \ + } + +bool ggml_validate_row_data(enum ggml_type type, const void * data, size_t nbytes) { + if (type < 0 || type >= GGML_TYPE_COUNT) { + fprintf(stderr, "%s: invalid type %d\n", __func__, type); + return false; + } + + if (nbytes % ggml_type_size(type) != 0) { + fprintf(stderr, "%s: invalid size %zu for type %s (type size = %zu)\n", __func__, nbytes, ggml_type_name(type), ggml_type_size(type)); + return false; + } + + const size_t nb = nbytes/ggml_type_size(type); + + switch (type) { + case GGML_TYPE_BF16: + { + int nans = 0; + int infs = 0; + const unsigned short * f = (const unsigned short *) data; + for (size_t i = 0; i < nb; ++i) { + nans += (f[i] & 0x7fff) > 0x7f80; + infs += (f[i] & 0x7fff) == 0x7f80; + } + if (nans) { + fprintf(stderr, "%s: found %d NaNs in row of %zu BF16 values\n", __func__, nans, nb); + return false; + } + if (infs) { + fprintf(stderr, "%s: found %d infinities in row of %zu BF16 values\n", __func__, infs, nb); + return false; + } + } break; + case GGML_TYPE_F16: + { + const ggml_fp16_t * f = (const ggml_fp16_t *) data; + size_t i = 0; +#if defined(__AVX2__) + for (; i + 15 < nb; i += 16) { + __m256i v = _mm256_loadu_si256((const __m256i *)(f + i)); + __m256i vexp = _mm256_and_si256(v, _mm256_set1_epi16(0x7c00)); + __m256i cmp = _mm256_cmpeq_epi16(vexp, _mm256_set1_epi16(0x7c00)); + int mask = _mm256_movemask_epi8(cmp); + if (mask) { + for (size_t j = 0; j < 16; ++j) { + if (!validate_fp16(f[i + j], i + j)) { + return false; + } + } + GGML_UNREACHABLE(); + } + } +#elif defined(__ARM_NEON) + for (; i + 7 < nb; i += 8) { + uint16x8_t v = vld1q_u16(f + i); + uint16x8_t vexp = vandq_u16(v, vdupq_n_u16(0x7c00)); + uint16x8_t cmp = vceqq_u16(vexp, vdupq_n_u16(0x7c00)); + uint64_t mask = vget_lane_u64(vreinterpret_u64_u8(vshrn_n_u16(cmp, 4)), 0); + if (mask) { + for (size_t j = 0; j < 8; ++j) { + if (!validate_fp16(f[i + j], i + j)) { + return false; + } + } + GGML_UNREACHABLE(); + } + } +#endif + for (; i < nb; ++i) { + if (!validate_fp16(f[i], i)) { + return false; + } + } + } break; + case GGML_TYPE_F32: + { + const float * f = (const float *) data; + size_t i = 0; +#if defined(__AVX2__) + for (; i + 7 < nb; i += 8) { + __m256i v = _mm256_loadu_si256((const __m256i *)(f + i)); + __m256i vexp = _mm256_and_si256(v, _mm256_set1_epi32(0x7f800000)); + __m256i cmp = _mm256_cmpeq_epi32(vexp, _mm256_set1_epi32(0x7f800000)); + int mask = _mm256_movemask_epi8(cmp); + if (mask) { + for (size_t j = 0; j < 8; ++j) { + if (!validate_float(f[i + j], i + j)) { + return false; + } + } + GGML_UNREACHABLE(); + } + } +#elif defined(__ARM_NEON) + for (; i + 3 < nb; i += 4) { + uint32x4_t v = vld1q_u32((const uint32_t *)f + i); + uint32x4_t vexp = vandq_u32(v, vdupq_n_u32(0x7f800000)); + uint32x4_t cmp = vceqq_u32(vexp, vdupq_n_u32(0x7f800000)); + uint64_t mask = vget_lane_u64(vreinterpret_u64_u16(vshrn_n_u32(cmp, 8)), 0); + if (mask) { + for (size_t j = 0; j < 4; ++j) { + if (!validate_float(f[i + j], i + j)) { + return false; + } + } + GGML_UNREACHABLE(); + } + } +#endif + for (; i < nb; ++i) { + if (!validate_float(f[i], i)) { + return false; + } + } + } break; + case GGML_TYPE_F64: + { + const double * f = (const double *) data; + for (size_t i = 0; i < nb; ++i) { + if (!validate_float(f[i], i)) { + return false; + } + } + } break; + case GGML_TYPE_Q1_0: + { + VALIDATE_ROW_DATA_D_F16_IMPL(block_q1_0, data, nb); + } break; + case GGML_TYPE_Q2_0: + { + VALIDATE_ROW_DATA_D_F16_IMPL(block_q2_0, data, nb); + } break; + case GGML_TYPE_Q4_0: + { + VALIDATE_ROW_DATA_D_F16_IMPL(block_q4_0, data, nb); + } break; + case GGML_TYPE_Q4_1: + { + VALIDATE_ROW_DATA_DM_F16_IMPL(block_q4_1, data, nb, d, m); + } break; + case GGML_TYPE_Q5_0: + { + VALIDATE_ROW_DATA_D_F16_IMPL(block_q5_0, data, nb); + } break; + case GGML_TYPE_Q5_1: + { + VALIDATE_ROW_DATA_DM_F16_IMPL(block_q5_1, data, nb, d, m); + } break; + case GGML_TYPE_Q8_0: + { + VALIDATE_ROW_DATA_D_F16_IMPL(block_q8_0, data, nb); + } break; + case GGML_TYPE_MXFP4: + { + VALIDATE_ROW_DATA_E_E8M0_IMPL(block_mxfp4, data, nb); + } break; + case GGML_TYPE_NVFP4: + { + // UE4M3 scales are uint8_t — all byte values are valid + GGML_UNUSED(data); + GGML_UNUSED(nb); + } break; + case GGML_TYPE_Q2_K: + { + VALIDATE_ROW_DATA_DM_F16_IMPL(block_q2_K, data, nb, d, dmin); + } break; + case GGML_TYPE_Q3_K: + { + VALIDATE_ROW_DATA_D_F16_IMPL(block_q3_K, data, nb); + } break; + case GGML_TYPE_Q4_K: + { + VALIDATE_ROW_DATA_DM_F16_IMPL(block_q4_K, data, nb, d, dmin); + } break; + case GGML_TYPE_Q5_K: + { + VALIDATE_ROW_DATA_DM_F16_IMPL(block_q5_K, data, nb, d, dmin); + } break; + case GGML_TYPE_Q6_K: + { + VALIDATE_ROW_DATA_D_F16_IMPL(block_q6_K, data, nb); + } break; + case GGML_TYPE_Q8_K: + { + const block_q8_K * q = (const block_q8_K *) data; + for (size_t i = 0; i < nb; ++i) { + if (!validate_float(q[i].d, i)) { + return false; + } + } + } break; + case GGML_TYPE_TQ1_0: + { + VALIDATE_ROW_DATA_D_F16_IMPL(block_tq1_0, data, nb); + } break; + case GGML_TYPE_TQ2_0: + { + VALIDATE_ROW_DATA_D_F16_IMPL(block_tq2_0, data, nb); + } break; + case GGML_TYPE_IQ1_S: + { + VALIDATE_ROW_DATA_D_F16_IMPL(block_iq1_s, data, nb); + } break; + case GGML_TYPE_IQ1_M: + { + const block_iq1_m * q = (const block_iq1_m *) data; + for (size_t i = 0; i < nb; ++i) { + iq1m_scale_t scale; + const uint16_t * sc = (const uint16_t *)q[i].scales; + scale.u16 = (sc[0] >> 12) | ((sc[1] >> 8) & 0x00f0) | ((sc[2] >> 4) & 0x0f00) | (sc[3] & 0xf000); + if (!validate_fp16(scale.f16, i)) { + return false; + } + } + } break; + case GGML_TYPE_IQ2_XXS: + { + VALIDATE_ROW_DATA_D_F16_IMPL(block_iq2_xxs, data, nb); + } break; + case GGML_TYPE_IQ2_XS: + { + VALIDATE_ROW_DATA_D_F16_IMPL(block_iq2_xs, data, nb); + } break; + case GGML_TYPE_IQ2_S: + { + VALIDATE_ROW_DATA_D_F16_IMPL(block_iq2_s, data, nb); + } break; + case GGML_TYPE_IQ3_XXS: + { + VALIDATE_ROW_DATA_D_F16_IMPL(block_iq3_xxs, data, nb); + } break; + + case GGML_TYPE_IQ3_S: + { + VALIDATE_ROW_DATA_D_F16_IMPL(block_iq3_s, data, nb); + } break; + case GGML_TYPE_IQ4_XS: + { + VALIDATE_ROW_DATA_D_F16_IMPL(block_iq4_xs, data, nb); + } break; + case GGML_TYPE_IQ4_NL: + { + VALIDATE_ROW_DATA_D_F16_IMPL(block_iq4_nl, data, nb); + } break; + + case GGML_TYPE_I8: + case GGML_TYPE_I16: + case GGML_TYPE_I32: + case GGML_TYPE_I64: + // nothing to validate + break; + default: + { + fprintf(stderr, "%s: invalid type %d\n", __func__, type); + return false; + } + } + + return true; +} diff --git a/python/freetoken/kernel/csrc/ggml/src/ggml-quants.h b/python/freetoken/kernel/csrc/ggml/src/ggml-quants.h new file mode 100644 index 0000000000..75188f1af1 --- /dev/null +++ b/python/freetoken/kernel/csrc/ggml/src/ggml-quants.h @@ -0,0 +1,115 @@ +#pragma once + +#define GGML_COMMON_DECL_C +#include "ggml-common.h" + +#include "ggml.h" + +// GGML internal header + +#ifdef __cplusplus +extern "C" { +#endif + +// NOTE: these functions are defined as GGML_API because they used by the CPU backend + +// Quantization +GGML_API void quantize_row_q1_0_ref(const float * GGML_RESTRICT x, block_q1_0 * GGML_RESTRICT y, int64_t k); +GGML_API void quantize_row_q2_0_ref(const float * GGML_RESTRICT x, block_q2_0 * GGML_RESTRICT y, int64_t k); +GGML_API void quantize_row_q4_0_ref(const float * GGML_RESTRICT x, block_q4_0 * GGML_RESTRICT y, int64_t k); +GGML_API void quantize_row_q4_1_ref(const float * GGML_RESTRICT x, block_q4_1 * GGML_RESTRICT y, int64_t k); +GGML_API void quantize_row_q5_0_ref(const float * GGML_RESTRICT x, block_q5_0 * GGML_RESTRICT y, int64_t k); +GGML_API void quantize_row_q5_1_ref(const float * GGML_RESTRICT x, block_q5_1 * GGML_RESTRICT y, int64_t k); +GGML_API void quantize_row_q8_0_ref(const float * GGML_RESTRICT x, block_q8_0 * GGML_RESTRICT y, int64_t k); +GGML_API void quantize_row_q8_1_ref(const float * GGML_RESTRICT x, block_q8_1 * GGML_RESTRICT y, int64_t k); + +GGML_API void quantize_row_mxfp4_ref(const float * GGML_RESTRICT x, block_mxfp4 * GGML_RESTRICT y, int64_t k); +GGML_API void quantize_row_nvfp4_ref(const float * GGML_RESTRICT x, block_nvfp4 * GGML_RESTRICT y, int64_t k); + +GGML_API void quantize_row_q2_K_ref(const float * GGML_RESTRICT x, block_q2_K * GGML_RESTRICT y, int64_t k); +GGML_API void quantize_row_q3_K_ref(const float * GGML_RESTRICT x, block_q3_K * GGML_RESTRICT y, int64_t k); +GGML_API void quantize_row_q4_K_ref(const float * GGML_RESTRICT x, block_q4_K * GGML_RESTRICT y, int64_t k); +GGML_API void quantize_row_q5_K_ref(const float * GGML_RESTRICT x, block_q5_K * GGML_RESTRICT y, int64_t k); +GGML_API void quantize_row_q6_K_ref(const float * GGML_RESTRICT x, block_q6_K * GGML_RESTRICT y, int64_t k); +GGML_API void quantize_row_q8_K_ref(const float * GGML_RESTRICT x, block_q8_K * GGML_RESTRICT y, int64_t k); + +GGML_API void quantize_row_tq1_0_ref(const float * GGML_RESTRICT x, block_tq1_0 * GGML_RESTRICT y, int64_t k); +GGML_API void quantize_row_tq2_0_ref(const float * GGML_RESTRICT x, block_tq2_0 * GGML_RESTRICT y, int64_t k); + +GGML_API void quantize_row_iq3_xxs_ref(const float * GGML_RESTRICT x, block_iq3_xxs * GGML_RESTRICT y, int64_t k); +GGML_API void quantize_row_iq4_nl_ref (const float * GGML_RESTRICT x, block_iq4_nl * GGML_RESTRICT y, int64_t k); +GGML_API void quantize_row_iq4_xs_ref (const float * GGML_RESTRICT x, block_iq4_xs * GGML_RESTRICT y, int64_t k); +GGML_API void quantize_row_iq3_s_ref (const float * GGML_RESTRICT x, block_iq3_s * GGML_RESTRICT y, int64_t k); +GGML_API void quantize_row_iq2_s_ref (const float * GGML_RESTRICT x, block_iq2_s * GGML_RESTRICT y, int64_t k); + +// Dequantization +GGML_API void dequantize_row_q1_0(const block_q1_0 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); +GGML_API void dequantize_row_q2_0(const block_q2_0 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); +GGML_API void dequantize_row_q4_0(const block_q4_0 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); +GGML_API void dequantize_row_q4_1(const block_q4_1 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); +GGML_API void dequantize_row_q5_0(const block_q5_0 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); +GGML_API void dequantize_row_q5_1(const block_q5_1 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); +GGML_API void dequantize_row_q8_0(const block_q8_0 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); +//GGML_API void dequantize_row_q8_1(const block_q8_1 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); + +GGML_API void dequantize_row_mxfp4(const block_mxfp4 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); +GGML_API void dequantize_row_nvfp4(const block_nvfp4 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); + +GGML_API void dequantize_row_q2_K(const block_q2_K * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); +GGML_API void dequantize_row_q3_K(const block_q3_K * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); +GGML_API void dequantize_row_q4_K(const block_q4_K * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); +GGML_API void dequantize_row_q5_K(const block_q5_K * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); +GGML_API void dequantize_row_q6_K(const block_q6_K * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); +GGML_API void dequantize_row_q8_K(const block_q8_K * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); + +GGML_API void dequantize_row_tq1_0(const block_tq1_0 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); +GGML_API void dequantize_row_tq2_0(const block_tq2_0 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); + +GGML_API void dequantize_row_iq2_xxs(const block_iq2_xxs * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); +GGML_API void dequantize_row_iq2_xs (const block_iq2_xs * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); +GGML_API void dequantize_row_iq2_s (const block_iq2_s * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); +GGML_API void dequantize_row_iq3_xxs(const block_iq3_xxs * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); +GGML_API void dequantize_row_iq1_s (const block_iq1_s * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); +GGML_API void dequantize_row_iq1_m (const block_iq1_m * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); +GGML_API void dequantize_row_iq4_nl (const block_iq4_nl * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); +GGML_API void dequantize_row_iq4_xs (const block_iq4_xs * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); +GGML_API void dequantize_row_iq3_s (const block_iq3_s * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); + +// Quantization utilizing an importance matrix (a.k.a. "Activation aWare Quantization") +GGML_API size_t quantize_iq2_xxs(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix); +GGML_API size_t quantize_iq2_xs (const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix); +GGML_API size_t quantize_iq2_s (const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix); +GGML_API size_t quantize_iq3_xxs(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix); +GGML_API size_t quantize_iq1_s (const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix); +GGML_API size_t quantize_iq1_m (const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix); +GGML_API size_t quantize_iq4_nl (const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix); +GGML_API size_t quantize_iq4_xs (const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix); +GGML_API size_t quantize_iq3_s (const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix); + +GGML_API size_t quantize_tq1_0(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix); +GGML_API size_t quantize_tq2_0(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix); + +GGML_API size_t quantize_q2_K(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix); +GGML_API size_t quantize_q3_K(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix); +GGML_API size_t quantize_q4_K(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix); +GGML_API size_t quantize_q5_K(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix); +GGML_API size_t quantize_q6_K(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix); +GGML_API size_t quantize_q1_0(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix); +GGML_API size_t quantize_q2_0(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix); +GGML_API size_t quantize_q4_0(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix); +GGML_API size_t quantize_q4_1(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix); +GGML_API size_t quantize_q5_0(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix); +GGML_API size_t quantize_q5_1(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix); +GGML_API size_t quantize_q8_0(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix); + +GGML_API size_t quantize_mxfp4(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix); +GGML_API size_t quantize_nvfp4(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix); + +GGML_API void iq2xs_init_impl(enum ggml_type type); +GGML_API void iq2xs_free_impl(enum ggml_type type); +GGML_API void iq3xs_init_impl(int grid_size); +GGML_API void iq3xs_free_impl(int grid_size); + +#ifdef __cplusplus +} +#endif diff --git a/python/freetoken/kernel/csrc/gguf/dequantize.cuh b/python/freetoken/kernel/csrc/gguf/dequantize.cuh deleted file mode 100644 index 57cb1a27cf..0000000000 --- a/python/freetoken/kernel/csrc/gguf/dequantize.cuh +++ /dev/null @@ -1,583 +0,0 @@ -// copied from -// https://github.com/vllm-project/vllm/blob/4492e3a55428e161ca8db381edc28263e5da4c8d/csrc/quantization/gguf/dequantize.cuh -// copied and adapted from https://github.com/ggerganov/llama.cpp/blob/b2899/ggml-cuda/convert.cu -// Dequant functions -static __device__ __forceinline__ void dequantize_q4_0(const void* vx, const int ib, const int iqs, dfloat2& v) { - const block_q4_0* x = (const block_q4_0*)vx; - - const dfloat d = x[ib].d; - - const int vui = x[ib].qs[iqs]; - - v.x = __int2half_rn(vui & 0xF); - v.y = __int2half_rn(vui >> 4); - - v = __hsub2(v, __floats2half2_rn(8.0f, 8.0f)); - v = __hmul2(v, {d, d}); -} - -static __device__ __forceinline__ void dequantize_q4_1(const void* vx, const int ib, const int iqs, dfloat2& v) { - const block_q4_1* x = (const block_q4_1*)vx; - - const dfloat d = __low2half(x[ib].dm); - const dfloat m = __high2half(x[ib].dm); - - const int vui = x[ib].qs[iqs]; - - v.x = __int2half_rn(vui & 0xF); - v.y = __int2half_rn(vui >> 4); - - v = __hmul2(v, {d, d}); - v = __hadd2(v, {m, m}); -} - -static __device__ __forceinline__ void dequantize_q5_0(const void* vx, const int ib, const int iqs, dfloat2& v) { - const block_q5_0* x = (const block_q5_0*)vx; - - const dfloat d = x[ib].d; - - uint32_t qh; - memcpy(&qh, x[ib].qh, sizeof(qh)); - - const int xh_0 = ((qh >> (iqs + 0)) << 4) & 0x10; - const int xh_1 = ((qh >> (iqs + 12))) & 0x10; - - v.x = __int2half_rn((x[ib].qs[iqs] & 0xf) | xh_0); - v.y = __int2half_rn((x[ib].qs[iqs] >> 4) | xh_1); - - v = __hsub2(v, __floats2half2_rn(16.0f, 16.0f)); - v = __hmul2(v, {d, d}); -} - -static __device__ __forceinline__ void dequantize_q5_1(const void* vx, const int ib, const int iqs, dfloat2& v) { - const block_q5_1* x = (const block_q5_1*)vx; - - const dfloat d = __low2half(x[ib].dm); - const dfloat m = __high2half(x[ib].dm); - - uint32_t qh; - memcpy(&qh, x[ib].qh, sizeof(qh)); - - const int xh_0 = ((qh >> (iqs + 0)) << 4) & 0x10; - const int xh_1 = ((qh >> (iqs + 12))) & 0x10; - - v.x = __int2half_rn((x[ib].qs[iqs] & 0xf) | xh_0); - v.y = __int2half_rn((x[ib].qs[iqs] >> 4) | xh_1); - - v = __hmul2(v, {d, d}); - v = __hadd2(v, {m, m}); -} - -static __device__ __forceinline__ void dequantize_q8_0(const void* vx, const int ib, const int iqs, dfloat2& v) { - const block_q8_0* x = (const block_q8_0*)vx; - - const dfloat d = x[ib].d; - - v.x = __int2half_rn(x[ib].qs[iqs + 0]); - v.y = __int2half_rn(x[ib].qs[iqs + 1]); - - v = __hmul2(v, {d, d}); -} - -template -static __global__ void dequantize_block(const void* __restrict__ vx, dst_t* __restrict__ y, const int k) { - const int i = 2 * (blockDim.x * blockIdx.x + threadIdx.x); - - if (i >= k) { - return; - } - - const int ib = i / qk; // block index - const int iqs = (i % qk) / qr; // quant index - const int iybs = i - i % qk; // y block start index - const int y_offset = qr == 1 ? 1 : qk / 2; - - // dequantize - dfloat2 v; - dequantize_kernel(vx, ib, iqs, v); - - y[iybs + iqs + 0] = convert_from_half(v.x); - y[iybs + iqs + y_offset] = convert_from_half(v.y); -} - -template -static __global__ void dequantize_block_q2_K(const void* __restrict__ vx, dst_t* __restrict__ yy) { - const auto i = blockIdx.x; - const block_q2_K* x = (const block_q2_K*)vx; - - const auto tid = threadIdx.x; - const int n = tid / 32; - const int l = tid - 32 * n; - const int is = 8 * n + l / 16; - - const uint8_t q = x[i].qs[32 * n + l]; - dst_t* y = yy + i * QK_K + 128 * n; - - half dall = __low2half(x[i].dm); - half dmin = __high2half(x[i].dm); - y[l + 0] = convert_from_half(__hsub( - __hmul(dall, __int2half_rn((x[i].scales[is + 0] & 0xF) * ((q >> 0) & 3))), - __hmul(dmin, __int2half_rn(x[i].scales[is + 0] >> 4)))); - y[l + 32] = convert_from_half(__hsub( - __hmul(dall, __int2half_rn((x[i].scales[is + 2] & 0xF) * ((q >> 2) & 3))), - __hmul(dmin, __int2half_rn(x[i].scales[is + 2] >> 4)))); - y[l + 64] = convert_from_half(__hsub( - __hmul(dall, __int2half_rn((x[i].scales[is + 4] & 0xF) * ((q >> 4) & 3))), - __hmul(dmin, __int2half_rn(x[i].scales[is + 4] >> 4)))); - y[l + 96] = convert_from_half(__hsub( - __hmul(dall, __int2half_rn((x[i].scales[is + 6] & 0xF) * ((q >> 6) & 3))), - __hmul(dmin, __int2half_rn(x[i].scales[is + 6] >> 4)))); -} - -template -static __global__ void dequantize_block_q3_K(const void* __restrict__ vx, dst_t* __restrict__ yy) { - const auto i = blockIdx.x; - const block_q3_K* x = (const block_q3_K*)vx; - - const auto r = threadIdx.x / 4; - const int tid = r / 2; - const int is0 = r % 2; - const int l0 = 16 * is0 + 4 * (threadIdx.x % 4); - const int n = tid / 4; - const int j = tid - 4 * n; - - uint8_t m = 1 << (4 * n + j); - int is = 8 * n + 2 * j + is0; - int shift = 2 * j; - - int8_t us = is < 4 ? (x[i].scales[is - 0] & 0xF) | (((x[i].scales[is + 8] >> 0) & 3) << 4) - : is < 8 ? (x[i].scales[is - 0] & 0xF) | (((x[i].scales[is + 4] >> 2) & 3) << 4) - : is < 12 ? (x[i].scales[is - 8] >> 4) | (((x[i].scales[is + 0] >> 4) & 3) << 4) - : (x[i].scales[is - 8] >> 4) | (((x[i].scales[is - 4] >> 6) & 3) << 4); - half d_all = x[i].d; - half dl = __hmul(d_all, __int2half_rn(us - 32)); - - dst_t* y = yy + i * QK_K + 128 * n + 32 * j; - const uint8_t* q = x[i].qs + 32 * n; - const uint8_t* hm = x[i].hmask; - - for (int l = l0; l < l0 + 4; ++l) { - y[l] = convert_from_half(__hmul(dl, __int2half_rn((int8_t)((q[l] >> shift) & 3) - ((hm[l] & m) ? 0 : 4)))); - } -} - -static inline __device__ void get_scale_min_k4(int j, const uint8_t* q, uint8_t& d, uint8_t& m) { - if (j < 4) { - d = q[j] & 63; - m = q[j + 4] & 63; - } else { - d = (q[j + 4] & 0xF) | ((q[j - 4] >> 6) << 4); - m = (q[j + 4] >> 4) | ((q[j - 0] >> 6) << 4); - } -} - -template -static __global__ void dequantize_block_q4_K(const void* __restrict__ vx, dst_t* __restrict__ yy) { - const block_q4_K* x = (const block_q4_K*)vx; - - const auto i = blockIdx.x; - - // assume 32 threads - const auto tid = threadIdx.x; - const int il = tid / 8; - const int ir = tid % 8; - const int is = 2 * il; - const int n = 4; - - dst_t* y = yy + i * QK_K + 64 * il + n * ir; - - const half dall = __low2half(x[i].dm); - const half dmin = __high2half(x[i].dm); - - const uint8_t* q = x[i].qs + 32 * il + n * ir; - - uint8_t sc, m; - get_scale_min_k4(is + 0, x[i].scales, sc, m); - const half d1 = __hmul(dall, __int2half_rn(sc)); - const half m1 = __hmul(dmin, __int2half_rn(m)); - get_scale_min_k4(is + 1, x[i].scales, sc, m); - const half d2 = __hmul(dall, __int2half_rn(sc)); - const half m2 = __hmul(dmin, __int2half_rn(m)); - for (int l = 0; l < n; ++l) { - y[l + 0] = convert_from_half(__hsub(__hmul(d1, __int2half_rn(q[l] & 0xF)), m1)); - y[l + 32] = convert_from_half(__hsub(__hmul(d2, __int2half_rn(q[l] >> 4)), m2)); - } -} - -template -static __global__ void dequantize_block_q5_K(const void* __restrict__ vx, dst_t* __restrict__ yy) { - const block_q5_K* x = (const block_q5_K*)vx; - - const auto i = blockIdx.x; - - // assume 64 threads - this is very slightly better than the one below - const auto tid = threadIdx.x; - const int il = tid / 16; // il is in 0...3 - const int ir = tid % 16; // ir is in 0...15 - const int is = 2 * il; // is is in 0...6 - - dst_t* y = yy + i * QK_K + 64 * il + 2 * ir; - - const half dall = __low2half(x[i].dm); - const half dmin = __high2half(x[i].dm); - - const uint8_t* ql = x[i].qs + 32 * il + 2 * ir; - const uint8_t* qh = x[i].qh + 2 * ir; - - uint8_t sc, m; - get_scale_min_k4(is + 0, x[i].scales, sc, m); - const half d1 = __hmul(dall, __int2half_rn(sc)); - const half m1 = __hmul(dmin, __int2half_rn(m)); - get_scale_min_k4(is + 1, x[i].scales, sc, m); - const half d2 = __hmul(dall, __int2half_rn(sc)); - const half m2 = __hmul(dmin, __int2half_rn(m)); - - uint8_t hm = 1 << (2 * il); - y[0] = convert_from_half(__hsub(__hmul(d1, __int2half_rn((ql[0] & 0xF) + (qh[0] & hm ? 16 : 0))), m1)); - y[1] = convert_from_half(__hsub(__hmul(d1, __int2half_rn((ql[1] & 0xF) + (qh[1] & hm ? 16 : 0))), m1)); - hm <<= 1; - y[32] = convert_from_half(__hsub(__hmul(d2, __int2half_rn((ql[0] >> 4) + (qh[0] & hm ? 16 : 0))), m2)); - y[33] = convert_from_half(__hsub(__hmul(d2, __int2half_rn((ql[1] >> 4) + (qh[1] & hm ? 16 : 0))), m2)); -} - -template -static __global__ void dequantize_block_q6_K(const void* __restrict__ vx, dst_t* __restrict__ yy) { - const block_q6_K* x = (const block_q6_K*)vx; - - const auto i = blockIdx.x; - - // assume 64 threads - this is very slightly better than the one below - const auto tid = threadIdx.x; - const int ip = tid / 32; // ip is 0 or 1 - const int il = tid - 32 * ip; // 0...32 - const int is = 8 * ip + il / 16; - - dst_t* y = yy + i * QK_K + 128 * ip + il; - - const half d = x[i].d; - - const uint8_t* ql = x[i].ql + 64 * ip + il; - const uint8_t qh = x[i].qh[32 * ip + il]; - const int8_t* sc = x[i].scales + is; - - y[0] = convert_from_half( - __hmul(d, __int2half_rn(sc[0] * ((int8_t)((ql[0] & 0xF) | (((qh >> 0) & 3) << 4)) - 32)))); - y[32] = convert_from_half( - __hmul(d, __int2half_rn(sc[2] * ((int8_t)((ql[32] & 0xF) | (((qh >> 2) & 3) << 4)) - 32)))); - y[64] = convert_from_half( - __hmul(d, __int2half_rn(sc[4] * ((int8_t)((ql[0] >> 4) | (((qh >> 4) & 3) << 4)) - 32)))); - y[96] = convert_from_half( - __hmul(d, __int2half_rn(sc[6] * ((int8_t)((ql[32] >> 4) | (((qh >> 6) & 3) << 4)) - 32)))); -} - -template -static __global__ void dequantize_block_iq2_xxs(const void* __restrict__ vx, dst_t* __restrict__ yy) { - const auto i = blockIdx.x; - const block_iq2_xxs* x = (const block_iq2_xxs*)vx; - - const auto tid = threadIdx.x; - const int il = tid / 8; // 0...3 - const int ib = tid % 8; // 0...7 - dst_t* y = yy + i * QK_K + 32 * ib + 8 * il; - const uint16_t* q2 = x[i].qs + 4 * ib; - const uint8_t* aux8 = (const uint8_t*)q2; - const uint8_t* grid = (const uint8_t*)(iq2xxs_grid + aux8[il]); - const uint32_t aux32 = q2[2] | (q2[3] << 16); - const float d = __half2float(x[i].d) * (0.5f + (aux32 >> 28)) * 0.25f; - const uint8_t signs = ksigns_iq2xs[(aux32 >> 7 * il) & 127]; - for (int j = 0; j < 8; ++j) - y[j] = d * grid[j] * (signs & kmask_iq2xs[j] ? -1.f : 1.f); -} - -template -static __global__ void dequantize_block_iq2_xs(const void* __restrict__ vx, dst_t* __restrict__ yy) { - const auto i = blockIdx.x; - const block_iq2_xs* x = (const block_iq2_xs*)vx; - - const auto tid = threadIdx.x; - const int il = tid / 8; // 0...3 - const int ib = tid % 8; // 0...7 - dst_t* y = yy + i * QK_K + 32 * ib + 8 * il; - const uint16_t* q2 = x[i].qs + 4 * ib; - const uint8_t* grid = (const uint8_t*)(iq2xs_grid + (q2[il] & 511)); - const float d = __half2float(x[i].d) * (0.5f + ((x[i].scales[ib] >> 4 * (il / 2)) & 0xf)) * 0.25f; - const uint8_t signs = ksigns_iq2xs[q2[il] >> 9]; - for (int j = 0; j < 8; ++j) - y[j] = d * grid[j] * (signs & kmask_iq2xs[j] ? -1.f : 1.f); -} - -template -static __global__ void dequantize_block_iq2_s(const void* __restrict__ vx, dst_t* __restrict__ yy) { - const auto i = blockIdx.x; - const block_iq2_s* x = (const block_iq2_s*)vx; - - const auto tid = threadIdx.x; - const int il = tid / 8; // 0...3 - const int ib = tid % 8; // 0...7 - dst_t* y = yy + i * QK_K + 32 * ib + 8 * il; - const uint8_t* grid = (const uint8_t*)(iq2s_grid + (x[i].qs[4 * ib + il] | ((x[i].qh[ib] << (8 - 2 * il)) & 0x300))); - const float d = __half2float(x[i].d) * (0.5f + ((x[i].scales[ib] >> 4 * (il / 2)) & 0xf)) * 0.25f; - const uint8_t signs = x[i].qs[QK_K / 8 + 4 * ib + il]; - for (int j = 0; j < 8; ++j) - y[j] = d * grid[j] * (signs & kmask_iq2xs[j] ? -1.f : 1.f); -} - -template -static __global__ void dequantize_block_iq3_xxs(const void* __restrict__ vx, dst_t* __restrict__ yy) { - const auto i = blockIdx.x; - const block_iq3_xxs* x = (const block_iq3_xxs*)vx; - - const auto tid = threadIdx.x; - const int il = tid / 8; // 0...3 - const int ib = tid % 8; // 0...7 - dst_t* y = yy + i * QK_K + 32 * ib + 8 * il; - const uint8_t* q3 = x[i].qs + 8 * ib; - const uint16_t* gas = (const uint16_t*)(x[i].qs + QK_K / 4) + 2 * ib; - const uint8_t* grid1 = (const uint8_t*)(iq3xxs_grid + q3[2 * il + 0]); - const uint8_t* grid2 = (const uint8_t*)(iq3xxs_grid + q3[2 * il + 1]); - const uint32_t aux32 = gas[0] | (gas[1] << 16); - const float d = __half2float(x[i].d) * (0.5f + (aux32 >> 28)) * 0.5f; - const uint8_t signs = ksigns_iq2xs[(aux32 >> 7 * il) & 127]; - for (int j = 0; j < 4; ++j) { - y[j + 0] = d * grid1[j] * (signs & kmask_iq2xs[j + 0] ? -1.f : 1.f); - y[j + 4] = d * grid2[j] * (signs & kmask_iq2xs[j + 4] ? -1.f : 1.f); - } -} - -template -static __global__ void dequantize_block_iq3_s(const void* __restrict__ vx, dst_t* __restrict__ yy) { - const auto i = blockIdx.x; - const block_iq3_s* x = (const block_iq3_s*)vx; - - const auto tid = threadIdx.x; - const int il = tid / 8; // 0...3 - const int ib = tid % 8; // 0...7 - dst_t* y = yy + i * QK_K + 32 * ib + 8 * il; - const uint8_t* qs = x[i].qs + 8 * ib; - const uint8_t* grid1 = (const uint8_t*)(iq3xs_grid + (qs[2 * il + 0] | ((x[i].qh[ib] << (8 - 2 * il)) & 256))); - const uint8_t* grid2 = (const uint8_t*)(iq3xs_grid + (qs[2 * il + 1] | ((x[i].qh[ib] << (7 - 2 * il)) & 256))); - const float d = __half2float(x[i].d) * (0.5f + ((x[i].scales[ib / 2] >> 4 * (ib % 2)) & 0xf)) * 0.5f; - const uint8_t signs = x[i].signs[4 * ib + il]; - for (int j = 0; j < 4; ++j) { - y[j + 0] = d * grid1[j] * (signs & kmask_iq2xs[j + 0] ? -1.f : 1.f); - y[j + 4] = d * grid2[j] * (signs & kmask_iq2xs[j + 4] ? -1.f : 1.f); - } -} - -template -static __global__ void dequantize_block_iq1_s(const void* __restrict__ vx, dst_t* __restrict__ yy) { - const int64_t i = blockIdx.x; - const block_iq1_s* x = (const block_iq1_s*)vx; - - const int64_t tid = threadIdx.x; - const int64_t il = tid / 8; // 0...3 - const int64_t ib = tid % 8; // 0...7 - dst_t* y = yy + i * QK_K + 32 * ib + 8 * il; - const float delta = x[i].qh[ib] & 0x8000 ? -1 - IQ1S_DELTA : -1 + IQ1S_DELTA; - const float d = __half2float(x[i].d) * (2 * ((x[i].qh[ib] >> 12) & 7) + 1); - uint32_t grid32[2]; - const int8_t* q = (const int8_t*)grid32; - grid32[0] = iq1s_grid_gpu[x[i].qs[4 * ib + il] | (((x[i].qh[ib] >> 3 * il) & 7) << 8)]; - grid32[1] = (grid32[0] >> 4) & 0x0f0f0f0f; - grid32[0] &= 0x0f0f0f0f; - for (int j = 0; j < 8; ++j) { - y[j] = d * (q[j] + delta); - } -} - -template -static __global__ void dequantize_block_iq1_m(const void* __restrict__ vx, dst_t* __restrict__ yy) { - const int64_t i = blockIdx.x; - const block_iq1_m* x = (const block_iq1_m*)vx; - - const int64_t tid = threadIdx.x; - const int64_t il = tid / 8; // 0...3 - const int64_t ib = tid % 8; // 0...7 - dst_t* y = yy + i * QK_K + 32 * ib + 8 * il; - const uint16_t* sc = (const uint16_t*)x[i].scales; - iq1m_scale_t scale; - scale.u16 = (sc[0] >> 12) | ((sc[1] >> 8) & 0x00f0) | ((sc[2] >> 4) & 0x0f00) | (sc[3] & 0xf000); - const int64_t ib16 = 2 * ib + il / 2; // sc[ib16/4] >> 3*(ib16%4) -> sc[ib/2] >> 3*((2*ib+il/2)%4); - const float d = __half2float(scale.f16) * (2 * ((sc[ib16 / 4] >> 3 * (ib16 % 4)) & 0x7) + 1); - const float delta = x[i].qh[2 * ib + il / 2] & (0x08 << 4 * (il % 2)) ? -1 - IQ1M_DELTA : -1 + IQ1M_DELTA; - uint32_t grid32[2]; - const int8_t* q = (const int8_t*)grid32; - grid32[0] = iq1s_grid_gpu[x[i].qs[4 * ib + il] | (((x[i].qh[2 * ib + il / 2] >> 4 * (il % 2)) & 7) << 8)]; - grid32[1] = (grid32[0] >> 4) & 0x0f0f0f0f; - grid32[0] &= 0x0f0f0f0f; - for (int j = 0; j < 8; ++j) { - y[j] = d * (q[j] + delta); - } -} - -template -static __global__ void dequantize_block_iq4_nl(const void* __restrict__ vx, dst_t* __restrict__ yy) { - const auto i = blockIdx.x; - const block_iq4_nl* x = (const block_iq4_nl*)vx + i * (QK_K / QK4_NL); - - const auto tid = threadIdx.x; - const int il = tid / 8; // 0...3 - const int ib = tid % 8; // 0...7 - dst_t* y = yy + i * QK_K + 32 * ib + 4 * il; - const uint8_t* q4 = x[ib].qs + 4 * il; - const float d = __half2float(x[ib].d); - for (int j = 0; j < 4; ++j) { - y[j + 0] = d * kvalues_iq4nl[q4[j] & 0xf]; - y[j + 16] = d * kvalues_iq4nl[q4[j] >> 4]; - } -} - -template -static __global__ void dequantize_block_iq4_xs(const void* __restrict__ vx, dst_t* __restrict__ yy) { - const auto i = blockIdx.x; - const block_iq4_xs* x = (const block_iq4_xs*)vx; - - const auto tid = threadIdx.x; - const int il = tid / 8; // 0...3 - const int ib = tid % 8; // 0...7 - dst_t* y = yy + i * QK_K + 32 * ib + 4 * il; - const uint8_t* q4 = x[i].qs + 16 * ib + 4 * il; - const float d = __half2float(x[i].d) * - ((((x[i].scales_l[ib / 2] >> 4 * (ib % 2)) & 0xf) | (((x[i].scales_h >> 2 * ib) & 3) << 4)) - 32); - for (int j = 0; j < 4; ++j) { - y[j + 0] = d * kvalues_iq4nl[q4[j] & 0xf]; - y[j + 16] = d * kvalues_iq4nl[q4[j] >> 4]; - } -} - -template -static void -dequantize_block_cuda(const void* __restrict__ vx, dst_t* __restrict__ y, const int k, cudaStream_t stream) { - const int num_blocks = (k + 2 * CUDA_DEQUANTIZE_BLOCK_SIZE - 1) / (2 * CUDA_DEQUANTIZE_BLOCK_SIZE); - dequantize_block<<>>(vx, y, k); -} - -template -static void dequantize_row_q2_K_cuda(const void* vx, dst_t* y, const int k, cudaStream_t stream) { - const int nb = k / QK_K; - dequantize_block_q2_K<<>>(vx, y); -} - -template -static void dequantize_row_q3_K_cuda(const void* vx, dst_t* y, const int k, cudaStream_t stream) { - const int nb = k / QK_K; - dequantize_block_q3_K<<>>(vx, y); -} - -template -static void dequantize_row_q4_K_cuda(const void* vx, dst_t* y, const int k, cudaStream_t stream) { - const int nb = k / QK_K; - dequantize_block_q4_K<<>>(vx, y); -} - -template -static void dequantize_row_q5_K_cuda(const void* vx, dst_t* y, const int k, cudaStream_t stream) { - const int nb = k / QK_K; - dequantize_block_q5_K<<>>(vx, y); -} - -template -static void dequantize_row_q6_K_cuda(const void* vx, dst_t* y, const int k, cudaStream_t stream) { - const int nb = k / QK_K; - dequantize_block_q6_K<<>>(vx, y); -} - -template -static void dequantize_row_iq2_xxs_cuda(const void* vx, dst_t* y, const int k, cudaStream_t stream) { - const int nb = k / QK_K; - dequantize_block_iq2_xxs<<>>(vx, y); -} - -template -static void dequantize_row_iq2_xs_cuda(const void* vx, dst_t* y, const int k, cudaStream_t stream) { - const int nb = k / QK_K; - dequantize_block_iq2_xs<<>>(vx, y); -} - -template -static void dequantize_row_iq2_s_cuda(const void* vx, dst_t* y, const int k, cudaStream_t stream) { - const int nb = k / QK_K; - dequantize_block_iq2_s<<>>(vx, y); -} - -template -static void dequantize_row_iq3_xxs_cuda(const void* vx, dst_t* y, const int k, cudaStream_t stream) { - const int nb = k / QK_K; - dequantize_block_iq3_xxs<<>>(vx, y); -} - -template -static void dequantize_row_iq3_s_cuda(const void* vx, dst_t* y, const int k, cudaStream_t stream) { - const int nb = k / QK_K; - dequantize_block_iq3_s<<>>(vx, y); -} - -template -static void dequantize_row_iq1_s_cuda(const void* vx, dst_t* y, const int k, cudaStream_t stream) { - const int nb = k / QK_K; - dequantize_block_iq1_s<<>>(vx, y); -} - -template -static void dequantize_row_iq1_m_cuda(const void* vx, dst_t* y, const int k, cudaStream_t stream) { - const int nb = k / QK_K; - dequantize_block_iq1_m<<>>(vx, y); -} - -template -static void dequantize_row_iq4_nl_cuda(const void* vx, dst_t* y, const int k, cudaStream_t stream) { - const int nb = (k + QK_K - 1) / QK_K; - dequantize_block_iq4_nl<<>>(vx, y); -} - -template -static void dequantize_row_iq4_xs_cuda(const void* vx, dst_t* y, const int k, cudaStream_t stream) { - const int nb = (k + QK_K - 1) / QK_K; - dequantize_block_iq4_xs<<>>(vx, y); -} - -template -static to_cuda_ggml_t ggml_get_to_cuda(int64_t type) { - switch (type) { - case 2: - return dequantize_block_cuda; - case 3: - return dequantize_block_cuda; - case 6: - return dequantize_block_cuda; - case 7: - return dequantize_block_cuda; - case 8: - return dequantize_block_cuda; - case 10: - return dequantize_row_q2_K_cuda; - case 11: - return dequantize_row_q3_K_cuda; - case 12: - return dequantize_row_q4_K_cuda; - case 13: - return dequantize_row_q5_K_cuda; - case 14: - return dequantize_row_q6_K_cuda; - case 16: - return dequantize_row_iq2_xxs_cuda; - case 17: - return dequantize_row_iq2_xs_cuda; - case 18: - return dequantize_row_iq3_xxs_cuda; - case 19: - return dequantize_row_iq1_s_cuda; - case 20: - return dequantize_row_iq4_nl_cuda; - case 21: - return dequantize_row_iq3_s_cuda; - case 22: - return dequantize_row_iq2_s_cuda; - case 23: - return dequantize_row_iq4_xs_cuda; - case 29: - return dequantize_row_iq1_m_cuda; - default: - return nullptr; - } -} diff --git a/python/freetoken/kernel/csrc/gguf/dispatch.h b/python/freetoken/kernel/csrc/gguf/dispatch.h deleted file mode 100644 index f42a216332..0000000000 --- a/python/freetoken/kernel/csrc/gguf/dispatch.h +++ /dev/null @@ -1,28 +0,0 @@ -// Minimal AT_DISPATCH helper for the vendored GGUF kernels (borrowed from -// sgl-kernel csrc/quantization/gguf, which are ports of llama.cpp). The donor -// pulls these macros from its large include/utils.h; we only need the float -// dispatch, so vendor just that to keep the JIT compile self-contained. -#pragma once - -#include - -#ifndef WARP_SIZE -#define WARP_SIZE 32 -#endif - -// Warp-shuffle wrappers the donor pulls from sgl-kernel's utils.h (CUDA variants). -#ifndef SGLANG_SHFL_XOR_SYNC -#define SGLANG_SHFL_XOR_SYNC(mask, var, lane_mask) __shfl_xor_sync((mask), (var), (lane_mask)) -#endif -#ifndef SGLANG_SHFL_XOR_SYNC_WIDTH -#define SGLANG_SHFL_XOR_SYNC_WIDTH(mask, var, lane_mask, width) \ - __shfl_xor_sync((mask), (var), (lane_mask), (width)) -#endif - -#define DISPATCH_CASE_FLOAT_TYPES(...) \ - AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \ - AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__) \ - AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__) - -#define DISPATCH_FLOAT_TYPES(TYPE, NAME, ...) \ - AT_DISPATCH_SWITCH(TYPE, NAME, DISPATCH_CASE_FLOAT_TYPES(__VA_ARGS__)) diff --git a/python/freetoken/kernel/csrc/gguf/ggml-common.h b/python/freetoken/kernel/csrc/gguf/ggml-common.h deleted file mode 100644 index 88c21a4ab3..0000000000 --- a/python/freetoken/kernel/csrc/gguf/ggml-common.h +++ /dev/null @@ -1,1029 +0,0 @@ -// adapted from -// https://github.com/vllm-project/vllm/blob/4492e3a55428e161ca8db381edc28263e5da4c8d/csrc/quantization/gguf/ggml-common.h -// copied from https://github.com/ggerganov/llama.cpp/blob/b2899/ggml-common.h -#define QK_K 256 -#define K_QUANTS_PER_ITERATION 2 -#define WARP_SIZE_GGUF 32 -#define K_SCALE_SIZE 12 -#define CUDA_DEQUANTIZE_BLOCK_SIZE 256 -#define CUDA_QUANTIZE_BLOCK_SIZE 256 -#define GGML_CUDA_DMMV_X 32 -#define GGML_CUDA_MMV_Y 1 - -// Data Structures -// QK = number of values after dequantization -// QR = QK / number of values before dequantization -// QI = number of 32 bit integers before dequantization - -#define QK4_0 32 -#define QR4_0 2 -#define QI4_0 (QK4_0 / (4 * QR4_0)) -typedef struct { - half d; // delta - uint8_t qs[QK4_0 / 2]; // nibbles / quants -} block_q4_0; - -#define QK4_1 32 -#define QR4_1 2 -#define QI4_1 (QK4_1 / (4 * QR4_1)) -typedef struct { - half2 dm; // dm.x = delta, dm.y = min - uint8_t qs[QK4_1 / 2]; // nibbles / quants -} block_q4_1; - -#define QK5_0 32 -#define QR5_0 2 -#define QI5_0 (QK5_0 / (4 * QR5_0)) -typedef struct { - half d; // delta - uint8_t qh[4]; // 5-th bit of quants - uint8_t qs[QK5_0 / 2]; // nibbles / quants -} block_q5_0; - -#define QK5_1 32 -#define QR5_1 2 -#define QI5_1 (QK5_1 / (4 * QR5_1)) -typedef struct { - half2 dm; // dm.x = delta, dm.y = min - uint8_t qh[4]; // 5-th bit of quants - uint8_t qs[QK5_1 / 2]; // nibbles / quants -} block_q5_1; - -#define QK8_0 32 -#define QR8_0 1 -#define QI8_0 (QK8_0 / (4 * QR8_0)) -typedef struct { - half d; // delta - int8_t qs[QK8_0]; // quants -} block_q8_0; - -#define QK8_1 32 -#define QR8_1 1 -#define QI8_1 (QK8_1 / (4 * QR8_1)) -typedef struct { - half2 ds; // ds.x = delta, ds.y = sum - int8_t qs[QK8_0]; // quants -} block_q8_1; - -#define QR2_K 4 -#define QI2_K (QK_K / (4 * QR2_K)) -typedef struct { - uint8_t scales[QK_K / 16]; // scales and mins, quantized with 4 bits - uint8_t qs[QK_K / 4]; // quants - half2 dm; // super-block scale for quantized scales/mins -} block_q2_K; - -#define QR3_K 4 -#define QI3_K (QK_K / (4 * QR3_K)) -typedef struct { - uint8_t hmask[QK_K / 8]; // quants - high bit - uint8_t qs[QK_K / 4]; // quants - low 2 bits - uint8_t scales[K_SCALE_SIZE]; // scales, quantized with 6 bits - half d; // super-block scale -} block_q3_K; - -#define QR4_K 2 -#define QI4_K (QK_K / (4 * QR4_K)) -typedef struct { - half2 dm; // super-block scale for quantized scales/mins - uint8_t scales[3 * QK_K / 64]; // scales, quantized with 6 bits - uint8_t qs[QK_K / 2]; // 4--bit quants -} block_q4_K; - -#define QR5_K 2 -#define QI5_K (QK_K / (4 * QR5_K)) -typedef struct { - half2 dm; // super-block scale for quantized scales/mins - uint8_t scales[K_SCALE_SIZE]; // scales and mins, quantized with 6 bits - uint8_t qh[QK_K / 8]; // quants, high bit - uint8_t qs[QK_K / 2]; // quants, low 4 bits -} block_q5_K; - -#define QR6_K 2 -#define QI6_K (QK_K / (4 * QR6_K)) -typedef struct { - uint8_t ql[QK_K / 2]; // quants, lower 4 bits - uint8_t qh[QK_K / 4]; // quants, upper 2 bits - int8_t scales[QK_K / 16]; // scales - half d; // delta -} block_q6_K; - -#define QR2_XXS 8 -#define QI2_XXS (QK_K / (4 * QR2_XXS)) -typedef struct { - half d; - uint16_t qs[QK_K / 8]; -} block_iq2_xxs; - -#define QR2_XS 8 -#define QI2_XS (QK_K / (4 * QR2_XS)) -typedef struct { - half d; - uint16_t qs[QK_K / 8]; - uint8_t scales[QK_K / 32]; -} block_iq2_xs; - -#define QR2_S 8 -#define QI2_S (QK_K / (4 * QR2_S)) -typedef struct { - half d; - uint8_t qs[QK_K / 4]; - uint8_t qh[QK_K / 32]; - uint8_t scales[QK_K / 32]; -} block_iq2_s; - -#define QR3_XXS 8 -#define QI3_XXS (QK_K / (4 * QR3_XXS)) -typedef struct { - half d; - uint8_t qs[3 * (QK_K / 8)]; -} block_iq3_xxs; - -#define QR3_XS 8 -#define QI3_XS (QK_K / (4 * QR3_XS)) -#define IQ3S_N_SCALE QK_K / 64 -typedef struct { - half d; - uint8_t qs[QK_K / 4]; - uint8_t qh[QK_K / 32]; - uint8_t signs[QK_K / 8]; - uint8_t scales[IQ3S_N_SCALE]; -} block_iq3_s; - -// 1.5625 bpw -#define QR1_S 8 -#define QI1_S (QK_K / (4 * QR1_S)) -typedef struct { - half d; - uint8_t qs[QK_K / 8]; - uint16_t qh[QK_K / 32]; -} block_iq1_s; - -// 1.75 bpw -#define QR1_M 8 -#define QI1_M (QK_K / (4 * QR1_M)) -typedef struct { - uint8_t qs[QK_K / 8]; // grid index, low 8 bits - uint8_t qh[QK_K / 16]; // grid index, high 3 bits + grid shift bit (for two groups of 8) - uint8_t scales[QK_K / 32]; // 3-bit block scales (4-bit if QK_K == 64) -} block_iq1_m; - -// Used by IQ1_M quants -typedef union { - half f16; - uint16_t u16; -} iq1m_scale_t; - -#define QK4_NL 32 -#define QR4_NL 2 -#define QI4_NL (QK4_NL / (4 * QR4_NL)) -typedef struct { - half d; - uint8_t qs[QK4_NL / 2]; -} block_iq4_nl; - -#define QR4_XS 8 -#define QI4_XS (QK_K / (4 * QR4_XS)) -typedef struct { - half d; - uint16_t scales_h; - uint8_t scales_l[QK_K / 64]; - uint8_t qs[QK_K / 2]; -} block_iq4_xs; - -static const __device__ uint64_t iq2xxs_grid[256] = { - 0x0808080808080808, 0x080808080808082b, 0x0808080808081919, 0x0808080808082b08, 0x0808080808082b2b, - 0x0808080808190819, 0x0808080808191908, 0x08080808082b0808, 0x08080808082b082b, 0x08080808082b2b08, - 0x08080808082b2b2b, 0x0808080819080819, 0x0808080819081908, 0x0808080819190808, 0x0808080819192b08, - 0x08080808192b0819, 0x08080808192b1908, 0x080808082b080808, 0x080808082b08082b, 0x080808082b082b2b, - 0x080808082b2b082b, 0x0808081908080819, 0x0808081908081908, 0x0808081908190808, 0x0808081908191919, - 0x0808081919080808, 0x080808192b081908, 0x080808192b192b08, 0x0808082b08080808, 0x0808082b0808082b, - 0x0808082b082b082b, 0x0808082b2b08082b, 0x0808190808080819, 0x0808190808081908, 0x0808190808190808, - 0x08081908082b0819, 0x08081908082b1908, 0x0808190819080808, 0x080819081908082b, 0x0808190819082b08, - 0x08081908192b0808, 0x080819082b080819, 0x080819082b081908, 0x080819082b190808, 0x080819082b2b1908, - 0x0808191908080808, 0x080819190808082b, 0x0808191908082b08, 0x08081919082b0808, 0x080819191908192b, - 0x08081919192b2b19, 0x080819192b080808, 0x080819192b190819, 0x0808192b08082b19, 0x0808192b08190808, - 0x0808192b19080808, 0x0808192b2b081908, 0x0808192b2b2b1908, 0x08082b0808080808, 0x08082b0808081919, - 0x08082b0808082b08, 0x08082b0808191908, 0x08082b08082b2b08, 0x08082b0819080819, 0x08082b0819081908, - 0x08082b0819190808, 0x08082b081919082b, 0x08082b082b082b08, 0x08082b1908081908, 0x08082b1919080808, - 0x08082b2b0808082b, 0x08082b2b08191908, 0x0819080808080819, 0x0819080808081908, 0x0819080808190808, - 0x08190808082b0819, 0x0819080819080808, 0x08190808192b0808, 0x081908082b081908, 0x081908082b190808, - 0x081908082b191919, 0x0819081908080808, 0x0819081908082b08, 0x08190819082b0808, 0x0819081919190808, - 0x0819081919192b2b, 0x081908192b080808, 0x0819082b082b1908, 0x0819082b19081919, 0x0819190808080808, - 0x0819190808082b08, 0x08191908082b0808, 0x08191908082b1919, 0x0819190819082b19, 0x081919082b080808, - 0x0819191908192b08, 0x08191919192b082b, 0x0819192b08080808, 0x0819192b0819192b, 0x08192b0808080819, - 0x08192b0808081908, 0x08192b0808190808, 0x08192b0819080808, 0x08192b082b080819, 0x08192b1908080808, - 0x08192b1908081919, 0x08192b192b2b0808, 0x08192b2b19190819, 0x082b080808080808, 0x082b08080808082b, - 0x082b080808082b2b, 0x082b080819081908, 0x082b0808192b0819, 0x082b08082b080808, 0x082b08082b08082b, - 0x082b0819082b2b19, 0x082b081919082b08, 0x082b082b08080808, 0x082b082b0808082b, 0x082b190808080819, - 0x082b190808081908, 0x082b190808190808, 0x082b190819080808, 0x082b19081919192b, 0x082b191908080808, - 0x082b191919080819, 0x082b1919192b1908, 0x082b192b2b190808, 0x082b2b0808082b08, 0x082b2b08082b0808, - 0x082b2b082b191908, 0x082b2b2b19081908, 0x1908080808080819, 0x1908080808081908, 0x1908080808190808, - 0x1908080808192b08, 0x19080808082b0819, 0x19080808082b1908, 0x1908080819080808, 0x1908080819082b08, - 0x190808081919192b, 0x19080808192b0808, 0x190808082b080819, 0x190808082b081908, 0x190808082b190808, - 0x1908081908080808, 0x19080819082b0808, 0x19080819192b0819, 0x190808192b080808, 0x190808192b081919, - 0x1908082b08080819, 0x1908082b08190808, 0x1908082b19082b08, 0x1908082b1919192b, 0x1908082b192b2b08, - 0x1908190808080808, 0x1908190808082b08, 0x19081908082b0808, 0x190819082b080808, 0x190819082b192b19, - 0x190819190819082b, 0x19081919082b1908, 0x1908192b08080808, 0x19082b0808080819, 0x19082b0808081908, - 0x19082b0808190808, 0x19082b0819080808, 0x19082b0819081919, 0x19082b1908080808, 0x19082b1919192b08, - 0x19082b19192b0819, 0x19082b192b08082b, 0x19082b2b19081919, 0x19082b2b2b190808, 0x1919080808080808, - 0x1919080808082b08, 0x1919080808190819, 0x1919080808192b19, 0x19190808082b0808, 0x191908082b080808, - 0x191908082b082b08, 0x1919081908081908, 0x191908191908082b, 0x191908192b2b1908, 0x1919082b2b190819, - 0x191919082b190808, 0x191919082b19082b, 0x1919191908082b2b, 0x1919192b08080819, 0x1919192b19191908, - 0x19192b0808080808, 0x19192b0808190819, 0x19192b0808192b19, 0x19192b08192b1908, 0x19192b1919080808, - 0x19192b2b08082b08, 0x192b080808081908, 0x192b080808190808, 0x192b080819080808, 0x192b0808192b2b08, - 0x192b081908080808, 0x192b081919191919, 0x192b082b08192b08, 0x192b082b192b0808, 0x192b190808080808, - 0x192b190808081919, 0x192b191908190808, 0x192b19190819082b, 0x192b19192b081908, 0x192b2b081908082b, - 0x2b08080808080808, 0x2b0808080808082b, 0x2b08080808082b2b, 0x2b08080819080819, 0x2b0808082b08082b, - 0x2b08081908081908, 0x2b08081908192b08, 0x2b08081919080808, 0x2b08082b08190819, 0x2b08190808080819, - 0x2b08190808081908, 0x2b08190808190808, 0x2b08190808191919, 0x2b08190819080808, 0x2b081908192b0808, - 0x2b08191908080808, 0x2b0819191908192b, 0x2b0819192b191908, 0x2b08192b08082b19, 0x2b08192b19080808, - 0x2b08192b192b0808, 0x2b082b080808082b, 0x2b082b1908081908, 0x2b082b2b08190819, 0x2b19080808081908, - 0x2b19080808190808, 0x2b190808082b1908, 0x2b19080819080808, 0x2b1908082b2b0819, 0x2b1908190819192b, - 0x2b1908192b080808, 0x2b19082b19081919, 0x2b19190808080808, 0x2b191908082b082b, 0x2b19190819081908, - 0x2b19191919190819, 0x2b192b082b080819, 0x2b192b19082b0808, 0x2b2b08080808082b, 0x2b2b080819190808, - 0x2b2b08082b081919, 0x2b2b081908082b19, 0x2b2b082b08080808, 0x2b2b190808192b08, 0x2b2b2b0819190808, - 0x2b2b2b1908081908, -}; - -static const __device__ uint64_t iq2xs_grid[512] = { - 0x0808080808080808, 0x080808080808082b, 0x0808080808081919, 0x0808080808082b08, 0x0808080808082b2b, - 0x0808080808190819, 0x0808080808191908, 0x080808080819192b, 0x0808080808192b19, 0x08080808082b0808, - 0x08080808082b082b, 0x08080808082b1919, 0x08080808082b2b08, 0x0808080819080819, 0x0808080819081908, - 0x080808081908192b, 0x0808080819082b19, 0x0808080819190808, 0x080808081919082b, 0x0808080819191919, - 0x0808080819192b08, 0x08080808192b0819, 0x08080808192b1908, 0x080808082b080808, 0x080808082b08082b, - 0x080808082b081919, 0x080808082b082b08, 0x080808082b190819, 0x080808082b191908, 0x080808082b192b19, - 0x080808082b2b0808, 0x0808081908080819, 0x0808081908081908, 0x080808190808192b, 0x0808081908082b19, - 0x0808081908190808, 0x080808190819082b, 0x0808081908191919, 0x0808081908192b08, 0x0808081908192b2b, - 0x08080819082b0819, 0x08080819082b1908, 0x0808081919080808, 0x080808191908082b, 0x0808081919081919, - 0x0808081919082b08, 0x0808081919190819, 0x0808081919191908, 0x08080819192b0808, 0x08080819192b2b08, - 0x080808192b080819, 0x080808192b081908, 0x080808192b190808, 0x0808082b08080808, 0x0808082b0808082b, - 0x0808082b08081919, 0x0808082b08082b08, 0x0808082b08190819, 0x0808082b08191908, 0x0808082b082b0808, - 0x0808082b19080819, 0x0808082b19081908, 0x0808082b19190808, 0x0808082b19191919, 0x0808082b2b080808, - 0x0808082b2b082b2b, 0x0808190808080819, 0x0808190808081908, 0x080819080808192b, 0x0808190808082b19, - 0x0808190808190808, 0x080819080819082b, 0x0808190808191919, 0x0808190808192b08, 0x08081908082b0819, - 0x08081908082b1908, 0x0808190819080808, 0x080819081908082b, 0x0808190819081919, 0x0808190819082b08, - 0x0808190819190819, 0x0808190819191908, 0x080819081919192b, 0x08081908192b0808, 0x080819082b080819, - 0x080819082b081908, 0x080819082b190808, 0x0808191908080808, 0x080819190808082b, 0x0808191908081919, - 0x0808191908082b08, 0x0808191908190819, 0x0808191908191908, 0x08081919082b0808, 0x0808191919080819, - 0x0808191919081908, 0x0808191919190808, 0x08081919192b0819, 0x080819192b080808, 0x0808192b08080819, - 0x0808192b08081908, 0x0808192b08190808, 0x0808192b082b192b, 0x0808192b19080808, 0x0808192b1908082b, - 0x0808192b2b081908, 0x08082b0808080808, 0x08082b080808082b, 0x08082b0808081919, 0x08082b0808082b08, - 0x08082b0808082b2b, 0x08082b0808190819, 0x08082b0808191908, 0x08082b08082b0808, 0x08082b08082b1919, - 0x08082b0819080819, 0x08082b0819081908, 0x08082b0819190808, 0x08082b0819192b08, 0x08082b082b080808, - 0x08082b082b2b0808, 0x08082b082b2b2b2b, 0x08082b1908080819, 0x08082b1908081908, 0x08082b1908190808, - 0x08082b1919080808, 0x08082b192b080819, 0x08082b192b082b19, 0x08082b2b08080808, 0x08082b2b082b0808, - 0x08082b2b082b2b08, 0x08082b2b2b19192b, 0x08082b2b2b2b0808, 0x0819080808080819, 0x0819080808081908, - 0x081908080808192b, 0x0819080808082b19, 0x0819080808190808, 0x081908080819082b, 0x0819080808191919, - 0x0819080808192b08, 0x08190808082b0819, 0x08190808082b1908, 0x0819080819080808, 0x081908081908082b, - 0x0819080819081919, 0x0819080819082b08, 0x0819080819190819, 0x0819080819191908, 0x08190808192b0808, - 0x08190808192b2b2b, 0x081908082b080819, 0x081908082b081908, 0x081908082b190808, 0x0819081908080808, - 0x081908190808082b, 0x0819081908081919, 0x0819081908082b08, 0x0819081908190819, 0x0819081908191908, - 0x08190819082b0808, 0x0819081919080819, 0x0819081919081908, 0x0819081919190808, 0x081908192b080808, - 0x081908192b191908, 0x081908192b19192b, 0x0819082b08080819, 0x0819082b08081908, 0x0819082b0808192b, - 0x0819082b08190808, 0x0819082b19080808, 0x0819082b192b0808, 0x0819190808080808, 0x081919080808082b, - 0x0819190808081919, 0x0819190808082b08, 0x0819190808190819, 0x0819190808191908, 0x08191908082b0808, - 0x0819190819080819, 0x0819190819081908, 0x0819190819082b19, 0x0819190819190808, 0x08191908192b1908, - 0x081919082b080808, 0x0819191908080819, 0x0819191908081908, 0x0819191908190808, 0x0819191919080808, - 0x0819192b08080808, 0x0819192b08191908, 0x0819192b19082b19, 0x08192b0808080819, 0x08192b0808081908, - 0x08192b0808190808, 0x08192b080819082b, 0x08192b0819080808, 0x08192b0819191908, 0x08192b082b08192b, - 0x08192b1908080808, 0x08192b1908081919, 0x08192b19192b192b, 0x08192b2b19190819, 0x08192b2b2b2b2b19, - 0x082b080808080808, 0x082b08080808082b, 0x082b080808081919, 0x082b080808082b08, 0x082b080808082b2b, - 0x082b080808190819, 0x082b080808191908, 0x082b0808082b0808, 0x082b080819080819, 0x082b080819081908, - 0x082b080819190808, 0x082b08082b080808, 0x082b08082b2b0808, 0x082b081908080819, 0x082b081908081908, - 0x082b081908190808, 0x082b081919080808, 0x082b081919082b08, 0x082b0819192b1919, 0x082b082b08080808, - 0x082b082b082b082b, 0x082b082b2b080808, 0x082b082b2b2b2b08, 0x082b190808080819, 0x082b190808081908, - 0x082b190808190808, 0x082b1908082b2b19, 0x082b190819080808, 0x082b191908080808, 0x082b191919080819, - 0x082b19191919082b, 0x082b19192b192b19, 0x082b192b08080819, 0x082b192b08192b2b, 0x082b192b2b2b192b, - 0x082b2b0808080808, 0x082b2b0808082b08, 0x082b2b0808082b2b, 0x082b2b08082b0808, 0x082b2b0819191919, - 0x082b2b082b082b08, 0x082b2b082b2b082b, 0x082b2b19192b2b08, 0x082b2b192b190808, 0x082b2b2b08082b08, - 0x082b2b2b082b0808, 0x082b2b2b2b08082b, 0x082b2b2b2b082b08, 0x082b2b2b2b082b2b, 0x1908080808080819, - 0x1908080808081908, 0x190808080808192b, 0x1908080808082b19, 0x1908080808190808, 0x190808080819082b, - 0x1908080808191919, 0x1908080808192b08, 0x19080808082b0819, 0x19080808082b1908, 0x1908080819080808, - 0x190808081908082b, 0x1908080819081919, 0x1908080819082b08, 0x1908080819082b2b, 0x1908080819190819, - 0x1908080819191908, 0x19080808192b0808, 0x19080808192b1919, 0x190808082b080819, 0x190808082b081908, - 0x190808082b190808, 0x1908081908080808, 0x190808190808082b, 0x1908081908081919, 0x1908081908082b08, - 0x1908081908190819, 0x1908081908191908, 0x19080819082b0808, 0x1908081919080819, 0x1908081919081908, - 0x1908081919190808, 0x190808192b080808, 0x190808192b081919, 0x190808192b2b082b, 0x1908082b08080819, - 0x1908082b08081908, 0x1908082b08190808, 0x1908082b0819082b, 0x1908082b082b2b19, 0x1908082b19080808, - 0x1908190808080808, 0x190819080808082b, 0x1908190808081919, 0x1908190808082b08, 0x1908190808190819, - 0x1908190808191908, 0x1908190808192b19, 0x19081908082b0808, 0x1908190819080819, 0x1908190819081908, - 0x1908190819190808, 0x190819082b080808, 0x190819082b191908, 0x1908191908080819, 0x1908191908081908, - 0x1908191908190808, 0x19081919082b1908, 0x1908191919080808, 0x190819192b192b2b, 0x1908192b08080808, - 0x1908192b08082b2b, 0x1908192b19081908, 0x1908192b19190808, 0x19082b0808080819, 0x19082b0808081908, - 0x19082b0808190808, 0x19082b0819080808, 0x19082b0819081919, 0x19082b0819191908, 0x19082b08192b082b, - 0x19082b1908080808, 0x19082b1908190819, 0x19082b1919081908, 0x19082b1919190808, 0x19082b19192b2b19, - 0x19082b2b08081908, 0x1919080808080808, 0x191908080808082b, 0x1919080808081919, 0x1919080808082b08, - 0x1919080808190819, 0x1919080808191908, 0x19190808082b0808, 0x19190808082b2b08, 0x1919080819080819, - 0x1919080819081908, 0x1919080819190808, 0x191908082b080808, 0x1919081908080819, 0x1919081908081908, - 0x1919081908190808, 0x1919081908191919, 0x1919081919080808, 0x191908191908082b, 0x1919082b08080808, - 0x1919082b19081908, 0x1919082b2b2b2b2b, 0x1919190808080819, 0x1919190808081908, 0x1919190808190808, - 0x19191908082b0819, 0x1919190819080808, 0x19191908192b0808, 0x191919082b080819, 0x191919082b2b0819, - 0x1919191908080808, 0x1919191908082b08, 0x191919192b080808, 0x191919192b082b08, 0x1919192b082b0819, - 0x1919192b192b2b08, 0x1919192b2b2b0819, 0x19192b0808080808, 0x19192b0808191908, 0x19192b0819080819, - 0x19192b0819190808, 0x19192b082b192b19, 0x19192b1908192b2b, 0x19192b1919080808, 0x19192b191908082b, - 0x19192b2b2b081919, 0x192b080808080819, 0x192b080808081908, 0x192b080808190808, 0x192b080819080808, - 0x192b080819191908, 0x192b0808192b082b, 0x192b08082b08192b, 0x192b08082b2b2b19, 0x192b081908080808, - 0x192b082b082b1908, 0x192b082b19082b2b, 0x192b082b2b19082b, 0x192b190808080808, 0x192b19080819192b, - 0x192b191908190808, 0x192b191919080808, 0x192b191919081919, 0x192b19192b2b1908, 0x192b2b0808080819, - 0x192b2b08192b2b2b, 0x192b2b19082b1919, 0x192b2b2b0808192b, 0x192b2b2b19191908, 0x192b2b2b192b082b, - 0x2b08080808080808, 0x2b0808080808082b, 0x2b08080808081919, 0x2b08080808082b08, 0x2b08080808190819, - 0x2b08080808191908, 0x2b080808082b0808, 0x2b080808082b2b2b, 0x2b08080819080819, 0x2b08080819081908, - 0x2b08080819190808, 0x2b0808082b080808, 0x2b0808082b08082b, 0x2b0808082b2b2b08, 0x2b0808082b2b2b2b, - 0x2b08081908080819, 0x2b08081908081908, 0x2b0808190808192b, 0x2b08081908190808, 0x2b08081919080808, - 0x2b08081919190819, 0x2b08081919192b19, 0x2b08082b08080808, 0x2b08082b082b0808, 0x2b08082b2b080808, - 0x2b08082b2b08082b, 0x2b08082b2b2b0808, 0x2b08082b2b2b2b08, 0x2b08190808080819, 0x2b08190808081908, - 0x2b08190808190808, 0x2b0819080819082b, 0x2b08190808191919, 0x2b08190819080808, 0x2b081908192b0808, - 0x2b0819082b082b19, 0x2b08191908080808, 0x2b08191919081908, 0x2b0819192b2b1919, 0x2b08192b08192b08, - 0x2b08192b192b2b2b, 0x2b082b0808080808, 0x2b082b0808082b08, 0x2b082b08082b1919, 0x2b082b0819192b2b, - 0x2b082b082b080808, 0x2b082b082b08082b, 0x2b082b082b2b2b08, 0x2b082b190808192b, 0x2b082b2b082b082b, - 0x2b082b2b2b080808, 0x2b082b2b2b082b08, 0x2b082b2b2b19192b, 0x2b082b2b2b2b2b08, 0x2b19080808080819, - 0x2b19080808081908, 0x2b19080808190808, 0x2b19080819080808, 0x2b1908081919192b, 0x2b1908082b081908, - 0x2b19081908080808, 0x2b190819082b082b, 0x2b190819192b1908, 0x2b19082b1919192b, 0x2b19082b2b082b19, - 0x2b19190808080808, 0x2b19190808081919, 0x2b19190819081908, 0x2b19190819190808, 0x2b19190819192b08, - 0x2b191919082b2b19, 0x2b1919192b190808, 0x2b1919192b19082b, 0x2b19192b19080819, 0x2b192b0819190819, - 0x2b192b082b2b192b, 0x2b192b1919082b19, 0x2b192b2b08191919, 0x2b192b2b192b0808, 0x2b2b080808080808, - 0x2b2b08080808082b, 0x2b2b080808082b08, 0x2b2b080808082b2b, 0x2b2b0808082b0808, 0x2b2b0808082b2b2b, - 0x2b2b08082b2b0808, 0x2b2b081919190819, 0x2b2b081919192b19, 0x2b2b08192b2b192b, 0x2b2b082b08080808, - 0x2b2b082b0808082b, 0x2b2b082b08082b08, 0x2b2b082b082b2b2b, 0x2b2b082b2b080808, 0x2b2b082b2b2b0808, - 0x2b2b190819080808, 0x2b2b19082b191919, 0x2b2b192b192b1919, 0x2b2b192b2b192b08, 0x2b2b2b0808082b2b, - 0x2b2b2b08082b0808, 0x2b2b2b08082b082b, 0x2b2b2b08082b2b08, 0x2b2b2b082b2b0808, 0x2b2b2b082b2b2b08, - 0x2b2b2b1908081908, 0x2b2b2b192b081908, 0x2b2b2b192b08192b, 0x2b2b2b2b082b2b08, 0x2b2b2b2b082b2b2b, - 0x2b2b2b2b2b190819, 0x2b2b2b2b2b2b2b2b, -}; - -static const __device__ uint64_t iq2s_grid[1024] = { - 0x0808080808080808, 0x080808080808082b, 0x0808080808081919, 0x0808080808082b08, 0x0808080808082b2b, - 0x0808080808190819, 0x0808080808191908, 0x080808080819192b, 0x0808080808192b19, 0x08080808082b0808, - 0x08080808082b082b, 0x08080808082b1919, 0x08080808082b2b08, 0x0808080819080819, 0x0808080819081908, - 0x080808081908192b, 0x0808080819082b19, 0x0808080819190808, 0x080808081919082b, 0x0808080819191919, - 0x0808080819192b08, 0x08080808192b0819, 0x08080808192b1908, 0x08080808192b192b, 0x08080808192b2b19, - 0x080808082b080808, 0x080808082b08082b, 0x080808082b081919, 0x080808082b082b08, 0x080808082b190819, - 0x080808082b191908, 0x080808082b2b0808, 0x080808082b2b1919, 0x080808082b2b2b2b, 0x0808081908080819, - 0x0808081908081908, 0x080808190808192b, 0x0808081908082b19, 0x0808081908190808, 0x080808190819082b, - 0x0808081908191919, 0x0808081908192b08, 0x08080819082b0819, 0x08080819082b1908, 0x0808081919080808, - 0x080808191908082b, 0x0808081919081919, 0x0808081919082b08, 0x0808081919190819, 0x0808081919191908, - 0x080808191919192b, 0x0808081919192b19, 0x08080819192b0808, 0x08080819192b1919, 0x08080819192b2b08, - 0x080808192b080819, 0x080808192b081908, 0x080808192b190808, 0x080808192b19082b, 0x080808192b191919, - 0x080808192b2b0819, 0x080808192b2b1908, 0x0808082b08080808, 0x0808082b0808082b, 0x0808082b08081919, - 0x0808082b08082b08, 0x0808082b08190819, 0x0808082b08191908, 0x0808082b082b0808, 0x0808082b082b2b2b, - 0x0808082b19080819, 0x0808082b19081908, 0x0808082b1908192b, 0x0808082b19082b19, 0x0808082b19190808, - 0x0808082b19191919, 0x0808082b2b080808, 0x0808082b2b081919, 0x0808082b2b082b2b, 0x0808082b2b191908, - 0x0808082b2b2b082b, 0x0808190808080819, 0x0808190808081908, 0x080819080808192b, 0x0808190808082b19, - 0x0808190808190808, 0x080819080819082b, 0x0808190808191919, 0x0808190808192b08, 0x08081908082b0819, - 0x08081908082b1908, 0x08081908082b192b, 0x08081908082b2b19, 0x0808190819080808, 0x080819081908082b, - 0x0808190819081919, 0x0808190819082b08, 0x0808190819082b2b, 0x0808190819190819, 0x0808190819191908, - 0x080819081919192b, 0x0808190819192b19, 0x08081908192b0808, 0x08081908192b082b, 0x08081908192b1919, - 0x080819082b080819, 0x080819082b081908, 0x080819082b08192b, 0x080819082b082b19, 0x080819082b190808, - 0x080819082b191919, 0x080819082b192b08, 0x080819082b2b0819, 0x080819082b2b1908, 0x0808191908080808, - 0x080819190808082b, 0x0808191908081919, 0x0808191908082b08, 0x0808191908082b2b, 0x0808191908190819, - 0x0808191908191908, 0x080819190819192b, 0x0808191908192b19, 0x08081919082b0808, 0x08081919082b1919, - 0x08081919082b2b08, 0x0808191919080819, 0x0808191919081908, 0x080819191908192b, 0x0808191919082b19, - 0x0808191919190808, 0x080819191919082b, 0x0808191919191919, 0x0808191919192b08, 0x08081919192b0819, - 0x08081919192b1908, 0x080819192b080808, 0x080819192b08082b, 0x080819192b081919, 0x080819192b082b08, - 0x080819192b190819, 0x080819192b191908, 0x080819192b2b0808, 0x0808192b08080819, 0x0808192b08081908, - 0x0808192b0808192b, 0x0808192b08082b19, 0x0808192b08190808, 0x0808192b08191919, 0x0808192b19080808, - 0x0808192b19081919, 0x0808192b19082b08, 0x0808192b19190819, 0x0808192b19191908, 0x0808192b192b0808, - 0x0808192b2b080819, 0x0808192b2b081908, 0x0808192b2b190808, 0x08082b0808080808, 0x08082b080808082b, - 0x08082b0808081919, 0x08082b0808082b08, 0x08082b0808190819, 0x08082b0808191908, 0x08082b080819192b, - 0x08082b0808192b19, 0x08082b08082b0808, 0x08082b08082b1919, 0x08082b08082b2b2b, 0x08082b0819080819, - 0x08082b0819081908, 0x08082b081908192b, 0x08082b0819082b19, 0x08082b0819190808, 0x08082b081919082b, - 0x08082b0819191919, 0x08082b0819192b08, 0x08082b08192b0819, 0x08082b08192b1908, 0x08082b082b080808, - 0x08082b082b081919, 0x08082b082b191908, 0x08082b082b2b2b2b, 0x08082b1908080819, 0x08082b1908081908, - 0x08082b1908190808, 0x08082b190819082b, 0x08082b1908191919, 0x08082b1908192b08, 0x08082b19082b0819, - 0x08082b1919080808, 0x08082b1919081919, 0x08082b1919082b08, 0x08082b1919190819, 0x08082b1919191908, - 0x08082b19192b0808, 0x08082b192b080819, 0x08082b192b190808, 0x08082b2b08080808, 0x08082b2b08190819, - 0x08082b2b08191908, 0x08082b2b082b082b, 0x08082b2b082b2b08, 0x08082b2b082b2b2b, 0x08082b2b19190808, - 0x08082b2b2b192b19, 0x0819080808080819, 0x0819080808081908, 0x081908080808192b, 0x0819080808082b19, - 0x0819080808190808, 0x081908080819082b, 0x0819080808191919, 0x0819080808192b08, 0x08190808082b0819, - 0x08190808082b1908, 0x08190808082b192b, 0x0819080819080808, 0x081908081908082b, 0x0819080819081919, - 0x0819080819082b08, 0x0819080819190819, 0x0819080819191908, 0x081908081919192b, 0x0819080819192b19, - 0x08190808192b0808, 0x08190808192b082b, 0x08190808192b1919, 0x08190808192b2b08, 0x081908082b080819, - 0x081908082b081908, 0x081908082b08192b, 0x081908082b190808, 0x081908082b191919, 0x081908082b192b08, - 0x081908082b2b0819, 0x081908082b2b1908, 0x0819081908080808, 0x081908190808082b, 0x0819081908081919, - 0x0819081908082b08, 0x0819081908082b2b, 0x0819081908190819, 0x0819081908191908, 0x081908190819192b, - 0x0819081908192b19, 0x08190819082b0808, 0x08190819082b082b, 0x08190819082b1919, 0x08190819082b2b08, - 0x0819081919080819, 0x0819081919081908, 0x081908191908192b, 0x0819081919082b19, 0x0819081919190808, - 0x081908191919082b, 0x0819081919191919, 0x0819081919192b08, 0x08190819192b0819, 0x08190819192b1908, - 0x081908192b080808, 0x081908192b08082b, 0x081908192b081919, 0x081908192b082b08, 0x081908192b190819, - 0x081908192b191908, 0x0819082b08080819, 0x0819082b08081908, 0x0819082b08082b19, 0x0819082b08190808, - 0x0819082b08191919, 0x0819082b082b0819, 0x0819082b082b1908, 0x0819082b19080808, 0x0819082b19081919, - 0x0819082b19190819, 0x0819082b19191908, 0x0819082b2b080819, 0x0819082b2b081908, 0x0819082b2b190808, - 0x0819190808080808, 0x081919080808082b, 0x0819190808081919, 0x0819190808082b08, 0x0819190808190819, - 0x0819190808191908, 0x081919080819192b, 0x0819190808192b19, 0x08191908082b0808, 0x08191908082b1919, - 0x08191908082b2b08, 0x0819190819080819, 0x0819190819081908, 0x081919081908192b, 0x0819190819082b19, - 0x0819190819190808, 0x081919081919082b, 0x0819190819191919, 0x0819190819192b08, 0x08191908192b0819, - 0x08191908192b1908, 0x081919082b080808, 0x081919082b08082b, 0x081919082b081919, 0x081919082b082b08, - 0x081919082b190819, 0x081919082b191908, 0x081919082b2b0808, 0x0819191908080819, 0x0819191908081908, - 0x081919190808192b, 0x0819191908082b19, 0x0819191908190808, 0x081919190819082b, 0x0819191908191919, - 0x0819191908192b08, 0x08191919082b0819, 0x08191919082b1908, 0x0819191919080808, 0x081919191908082b, - 0x0819191919081919, 0x0819191919082b08, 0x0819191919190819, 0x0819191919191908, 0x08191919192b0808, - 0x081919192b080819, 0x081919192b081908, 0x081919192b190808, 0x0819192b08080808, 0x0819192b08081919, - 0x0819192b08082b08, 0x0819192b08190819, 0x0819192b08191908, 0x0819192b082b0808, 0x0819192b19080819, - 0x0819192b19081908, 0x0819192b19190808, 0x0819192b2b080808, 0x0819192b2b2b2b2b, 0x08192b0808080819, - 0x08192b0808081908, 0x08192b080808192b, 0x08192b0808082b19, 0x08192b0808190808, 0x08192b0808191919, - 0x08192b0808192b08, 0x08192b08082b0819, 0x08192b0819080808, 0x08192b081908082b, 0x08192b0819081919, - 0x08192b0819082b08, 0x08192b0819190819, 0x08192b0819191908, 0x08192b08192b0808, 0x08192b082b080819, - 0x08192b082b081908, 0x08192b1908080808, 0x08192b190808082b, 0x08192b1908081919, 0x08192b1908082b08, - 0x08192b1908190819, 0x08192b1908191908, 0x08192b19082b0808, 0x08192b1919080819, 0x08192b1919081908, - 0x08192b1919190808, 0x08192b19192b2b19, 0x08192b192b2b082b, 0x08192b2b08081908, 0x08192b2b08190808, - 0x08192b2b19080808, 0x08192b2b1919192b, 0x082b080808080808, 0x082b08080808082b, 0x082b080808081919, - 0x082b080808082b08, 0x082b080808190819, 0x082b080808191908, 0x082b08080819192b, 0x082b080808192b19, - 0x082b0808082b0808, 0x082b0808082b1919, 0x082b0808082b2b2b, 0x082b080819080819, 0x082b080819081908, - 0x082b080819190808, 0x082b08081919082b, 0x082b080819191919, 0x082b0808192b1908, 0x082b08082b080808, - 0x082b08082b082b2b, 0x082b08082b191908, 0x082b08082b2b2b2b, 0x082b081908080819, 0x082b081908081908, - 0x082b081908190808, 0x082b08190819082b, 0x082b081908191919, 0x082b0819082b0819, 0x082b081919080808, - 0x082b08191908082b, 0x082b081919081919, 0x082b081919190819, 0x082b081919191908, 0x082b0819192b0808, - 0x082b08192b080819, 0x082b08192b081908, 0x082b08192b190808, 0x082b082b08080808, 0x082b082b08082b2b, - 0x082b082b082b082b, 0x082b082b082b2b08, 0x082b082b082b2b2b, 0x082b082b19081908, 0x082b082b19190808, - 0x082b082b2b082b08, 0x082b082b2b082b2b, 0x082b082b2b2b2b08, 0x082b190808080819, 0x082b190808081908, - 0x082b19080808192b, 0x082b190808082b19, 0x082b190808190808, 0x082b190808191919, 0x082b190808192b08, - 0x082b1908082b0819, 0x082b1908082b1908, 0x082b190819080808, 0x082b19081908082b, 0x082b190819081919, - 0x082b190819082b08, 0x082b190819190819, 0x082b190819191908, 0x082b1908192b0808, 0x082b19082b080819, - 0x082b19082b081908, 0x082b19082b190808, 0x082b191908080808, 0x082b191908081919, 0x082b191908082b08, - 0x082b191908190819, 0x082b191908191908, 0x082b1919082b0808, 0x082b191919080819, 0x082b191919081908, - 0x082b191919190808, 0x082b1919192b192b, 0x082b19192b080808, 0x082b192b08080819, 0x082b192b08081908, - 0x082b192b08190808, 0x082b192b19080808, 0x082b192b19192b19, 0x082b2b0808080808, 0x082b2b0808081919, - 0x082b2b0808190819, 0x082b2b0808191908, 0x082b2b0819080819, 0x082b2b0819081908, 0x082b2b0819190808, - 0x082b2b082b082b2b, 0x082b2b082b2b2b2b, 0x082b2b1908080819, 0x082b2b1908081908, 0x082b2b1908190808, - 0x082b2b192b191919, 0x082b2b2b08082b2b, 0x082b2b2b082b082b, 0x082b2b2b192b1908, 0x082b2b2b2b082b08, - 0x082b2b2b2b082b2b, 0x1908080808080819, 0x1908080808081908, 0x190808080808192b, 0x1908080808082b19, - 0x1908080808190808, 0x190808080819082b, 0x1908080808191919, 0x1908080808192b08, 0x1908080808192b2b, - 0x19080808082b0819, 0x19080808082b1908, 0x19080808082b192b, 0x1908080819080808, 0x190808081908082b, - 0x1908080819081919, 0x1908080819082b08, 0x1908080819082b2b, 0x1908080819190819, 0x1908080819191908, - 0x190808081919192b, 0x1908080819192b19, 0x19080808192b0808, 0x19080808192b082b, 0x19080808192b1919, - 0x190808082b080819, 0x190808082b081908, 0x190808082b190808, 0x190808082b191919, 0x190808082b192b08, - 0x190808082b2b0819, 0x190808082b2b1908, 0x1908081908080808, 0x190808190808082b, 0x1908081908081919, - 0x1908081908082b08, 0x1908081908190819, 0x1908081908191908, 0x190808190819192b, 0x1908081908192b19, - 0x19080819082b0808, 0x19080819082b082b, 0x19080819082b1919, 0x1908081919080819, 0x1908081919081908, - 0x190808191908192b, 0x1908081919082b19, 0x1908081919190808, 0x190808191919082b, 0x1908081919191919, - 0x1908081919192b08, 0x19080819192b0819, 0x19080819192b1908, 0x190808192b080808, 0x190808192b08082b, - 0x190808192b081919, 0x190808192b082b08, 0x190808192b190819, 0x190808192b191908, 0x190808192b2b0808, - 0x1908082b08080819, 0x1908082b08081908, 0x1908082b08190808, 0x1908082b0819082b, 0x1908082b08191919, - 0x1908082b08192b08, 0x1908082b082b1908, 0x1908082b19080808, 0x1908082b19081919, 0x1908082b19082b08, - 0x1908082b19190819, 0x1908082b19191908, 0x1908082b192b0808, 0x1908082b2b080819, 0x1908082b2b081908, - 0x1908190808080808, 0x190819080808082b, 0x1908190808081919, 0x1908190808082b08, 0x1908190808082b2b, - 0x1908190808190819, 0x1908190808191908, 0x190819080819192b, 0x1908190808192b19, 0x19081908082b0808, - 0x19081908082b082b, 0x19081908082b1919, 0x19081908082b2b08, 0x1908190819080819, 0x1908190819081908, - 0x190819081908192b, 0x1908190819082b19, 0x1908190819190808, 0x190819081919082b, 0x1908190819191919, - 0x1908190819192b08, 0x19081908192b0819, 0x19081908192b1908, 0x190819082b080808, 0x190819082b08082b, - 0x190819082b081919, 0x190819082b082b08, 0x190819082b190819, 0x190819082b191908, 0x190819082b2b0808, - 0x1908191908080819, 0x1908191908081908, 0x190819190808192b, 0x1908191908082b19, 0x1908191908190808, - 0x190819190819082b, 0x1908191908191919, 0x1908191908192b08, 0x19081919082b0819, 0x19081919082b1908, - 0x1908191919080808, 0x190819191908082b, 0x1908191919081919, 0x1908191919082b08, 0x1908191919190819, - 0x1908191919191908, 0x19081919192b0808, 0x19081919192b2b2b, 0x190819192b080819, 0x190819192b081908, - 0x190819192b190808, 0x1908192b08080808, 0x1908192b0808082b, 0x1908192b08081919, 0x1908192b08082b08, - 0x1908192b08190819, 0x1908192b08191908, 0x1908192b082b0808, 0x1908192b19080819, 0x1908192b19081908, - 0x1908192b19190808, 0x1908192b2b080808, 0x1908192b2b2b1919, 0x19082b0808080819, 0x19082b0808081908, - 0x19082b0808082b19, 0x19082b0808190808, 0x19082b080819082b, 0x19082b0808191919, 0x19082b0808192b08, - 0x19082b08082b0819, 0x19082b08082b1908, 0x19082b0819080808, 0x19082b081908082b, 0x19082b0819081919, - 0x19082b0819082b08, 0x19082b0819190819, 0x19082b0819191908, 0x19082b08192b0808, 0x19082b082b081908, - 0x19082b082b190808, 0x19082b1908080808, 0x19082b190808082b, 0x19082b1908081919, 0x19082b1908082b08, - 0x19082b1908190819, 0x19082b1908191908, 0x19082b19082b0808, 0x19082b1919080819, 0x19082b1919081908, - 0x19082b1919190808, 0x19082b192b080808, 0x19082b192b19192b, 0x19082b2b08080819, 0x19082b2b08081908, - 0x19082b2b08190808, 0x19082b2b19080808, 0x1919080808080808, 0x191908080808082b, 0x1919080808081919, - 0x1919080808082b08, 0x1919080808190819, 0x1919080808191908, 0x191908080819192b, 0x1919080808192b19, - 0x19190808082b0808, 0x19190808082b082b, 0x19190808082b1919, 0x19190808082b2b08, 0x1919080819080819, - 0x1919080819081908, 0x191908081908192b, 0x1919080819082b19, 0x1919080819190808, 0x191908081919082b, - 0x1919080819191919, 0x1919080819192b08, 0x19190808192b0819, 0x19190808192b1908, 0x191908082b080808, - 0x191908082b08082b, 0x191908082b081919, 0x191908082b082b08, 0x191908082b190819, 0x191908082b191908, - 0x1919081908080819, 0x1919081908081908, 0x191908190808192b, 0x1919081908082b19, 0x1919081908190808, - 0x191908190819082b, 0x1919081908191919, 0x1919081908192b08, 0x19190819082b0819, 0x19190819082b1908, - 0x1919081919080808, 0x191908191908082b, 0x1919081919081919, 0x1919081919082b08, 0x1919081919190819, - 0x1919081919191908, 0x19190819192b0808, 0x191908192b080819, 0x191908192b081908, 0x191908192b190808, - 0x1919082b08080808, 0x1919082b08081919, 0x1919082b08082b08, 0x1919082b08190819, 0x1919082b08191908, - 0x1919082b082b0808, 0x1919082b19080819, 0x1919082b19081908, 0x1919082b19190808, 0x1919082b192b2b19, - 0x1919082b2b080808, 0x1919190808080819, 0x1919190808081908, 0x191919080808192b, 0x1919190808082b19, - 0x1919190808190808, 0x191919080819082b, 0x1919190808191919, 0x1919190808192b08, 0x19191908082b0819, - 0x19191908082b1908, 0x1919190819080808, 0x191919081908082b, 0x1919190819081919, 0x1919190819082b08, - 0x1919190819190819, 0x1919190819191908, 0x19191908192b0808, 0x191919082b080819, 0x191919082b081908, - 0x191919082b190808, 0x1919191908080808, 0x191919190808082b, 0x1919191908081919, 0x1919191908082b08, - 0x1919191908190819, 0x1919191908191908, 0x19191919082b0808, 0x1919191919080819, 0x1919191919081908, - 0x1919191919190808, 0x191919192b080808, 0x1919192b08080819, 0x1919192b08081908, 0x1919192b08190808, - 0x1919192b082b192b, 0x1919192b19080808, 0x19192b0808080808, 0x19192b080808082b, 0x19192b0808081919, - 0x19192b0808082b08, 0x19192b0808190819, 0x19192b0808191908, 0x19192b08082b0808, 0x19192b0819080819, - 0x19192b0819081908, 0x19192b0819190808, 0x19192b0819192b2b, 0x19192b082b080808, 0x19192b1908080819, - 0x19192b1908081908, 0x19192b1908190808, 0x19192b1919080808, 0x19192b2b08080808, 0x19192b2b08192b19, - 0x19192b2b2b081919, 0x19192b2b2b2b2b08, 0x192b080808080819, 0x192b080808081908, 0x192b08080808192b, - 0x192b080808190808, 0x192b08080819082b, 0x192b080808191919, 0x192b080808192b08, 0x192b0808082b0819, - 0x192b0808082b1908, 0x192b080819080808, 0x192b080819081919, 0x192b080819082b08, 0x192b080819190819, - 0x192b080819191908, 0x192b0808192b0808, 0x192b08082b081908, 0x192b08082b190808, 0x192b081908080808, - 0x192b08190808082b, 0x192b081908081919, 0x192b081908082b08, 0x192b081908190819, 0x192b081908191908, - 0x192b0819082b0808, 0x192b081919080819, 0x192b081919081908, 0x192b081919190808, 0x192b08192b080808, - 0x192b08192b192b19, 0x192b082b08081908, 0x192b082b08190808, 0x192b082b19080808, 0x192b082b1919192b, - 0x192b082b2b2b0819, 0x192b190808080808, 0x192b190808081919, 0x192b190808082b08, 0x192b190808190819, - 0x192b190808191908, 0x192b1908082b0808, 0x192b190819080819, 0x192b190819081908, 0x192b190819190808, - 0x192b19082b080808, 0x192b191908080819, 0x192b191908081908, 0x192b191908190808, 0x192b191919080808, - 0x192b191919082b2b, 0x192b1919192b2b08, 0x192b19192b19082b, 0x192b192b08080808, 0x192b192b2b191908, - 0x192b2b0808080819, 0x192b2b0808081908, 0x192b2b0808190808, 0x192b2b08192b1919, 0x192b2b082b192b08, - 0x192b2b1908080808, 0x192b2b19082b2b2b, 0x192b2b2b1908082b, 0x192b2b2b2b2b0819, 0x2b08080808080808, - 0x2b0808080808082b, 0x2b08080808081919, 0x2b08080808082b08, 0x2b08080808190819, 0x2b08080808191908, - 0x2b08080808192b19, 0x2b080808082b0808, 0x2b080808082b1919, 0x2b08080819080819, 0x2b08080819081908, - 0x2b08080819190808, 0x2b0808081919082b, 0x2b08080819191919, 0x2b08080819192b08, 0x2b080808192b0819, - 0x2b0808082b080808, 0x2b0808082b081919, 0x2b0808082b190819, 0x2b0808082b191908, 0x2b08081908080819, - 0x2b08081908081908, 0x2b08081908082b19, 0x2b08081908190808, 0x2b0808190819082b, 0x2b08081908191919, - 0x2b08081908192b08, 0x2b080819082b0819, 0x2b080819082b1908, 0x2b08081919080808, 0x2b0808191908082b, - 0x2b08081919081919, 0x2b08081919082b08, 0x2b08081919190819, 0x2b08081919191908, 0x2b0808192b080819, - 0x2b0808192b081908, 0x2b0808192b190808, 0x2b0808192b2b2b19, 0x2b08082b08080808, 0x2b08082b08081919, - 0x2b08082b08082b2b, 0x2b08082b08190819, 0x2b08082b08191908, 0x2b08082b19080819, 0x2b08082b19081908, - 0x2b08082b19190808, 0x2b08190808080819, 0x2b08190808081908, 0x2b0819080808192b, 0x2b08190808082b19, - 0x2b08190808190808, 0x2b0819080819082b, 0x2b08190808191919, 0x2b08190808192b08, 0x2b081908082b0819, - 0x2b08190819080808, 0x2b0819081908082b, 0x2b08190819081919, 0x2b08190819082b08, 0x2b08190819190819, - 0x2b08190819191908, 0x2b081908192b0808, 0x2b0819082b080819, 0x2b0819082b081908, 0x2b0819082b190808, - 0x2b08191908080808, 0x2b0819190808082b, 0x2b08191908081919, 0x2b08191908082b08, 0x2b08191908190819, - 0x2b08191908191908, 0x2b081919082b0808, 0x2b08191919080819, 0x2b08191919081908, 0x2b08191919190808, - 0x2b0819192b080808, 0x2b0819192b082b2b, 0x2b08192b08080819, 0x2b08192b08081908, 0x2b08192b08190808, - 0x2b08192b082b2b19, 0x2b08192b19080808, 0x2b082b0808080808, 0x2b082b0808081919, 0x2b082b0808190819, - 0x2b082b0808191908, 0x2b082b0819080819, 0x2b082b0819081908, 0x2b082b0819190808, 0x2b082b082b2b082b, - 0x2b082b1908080819, 0x2b082b1908081908, 0x2b082b1919080808, 0x2b082b19192b1919, 0x2b082b2b082b082b, - 0x2b082b2b19192b08, 0x2b082b2b19192b2b, 0x2b082b2b2b08082b, 0x2b082b2b2b2b082b, 0x2b19080808080819, - 0x2b19080808081908, 0x2b19080808082b19, 0x2b19080808190808, 0x2b1908080819082b, 0x2b19080808191919, - 0x2b19080808192b08, 0x2b190808082b1908, 0x2b19080819080808, 0x2b1908081908082b, 0x2b19080819081919, - 0x2b19080819082b08, 0x2b19080819190819, 0x2b19080819191908, 0x2b190808192b0808, 0x2b1908082b080819, - 0x2b1908082b081908, 0x2b1908082b190808, 0x2b19081908080808, 0x2b19081908081919, 0x2b19081908190819, - 0x2b19081908191908, 0x2b19081919080819, 0x2b19081919081908, 0x2b19081919190808, 0x2b19081919192b2b, - 0x2b19082b08080819, 0x2b19082b08081908, 0x2b19082b08190808, 0x2b19082b19080808, 0x2b19082b2b2b192b, - 0x2b19190808080808, 0x2b1919080808082b, 0x2b19190808081919, 0x2b19190808082b08, 0x2b19190808190819, - 0x2b19190808191908, 0x2b191908082b0808, 0x2b19190819080819, 0x2b19190819081908, 0x2b19190819190808, - 0x2b1919082b080808, 0x2b1919082b19192b, 0x2b19191908080819, 0x2b19191908081908, 0x2b19191908190808, - 0x2b19191919080808, 0x2b1919192b192b08, 0x2b1919192b2b0819, 0x2b19192b08080808, 0x2b19192b1908192b, - 0x2b19192b192b1908, 0x2b192b0808080819, 0x2b192b0808081908, 0x2b192b0808190808, 0x2b192b08082b192b, - 0x2b192b0819080808, 0x2b192b082b2b2b19, 0x2b192b1908080808, 0x2b192b1919082b19, 0x2b192b191919082b, - 0x2b192b2b2b190808, 0x2b2b080808080808, 0x2b2b080808081919, 0x2b2b080808082b2b, 0x2b2b080808191908, - 0x2b2b0808082b082b, 0x2b2b0808082b2b2b, 0x2b2b080819080819, 0x2b2b080819081908, 0x2b2b080819190808, - 0x2b2b08082b2b082b, 0x2b2b08082b2b2b2b, 0x2b2b081919080808, 0x2b2b0819192b1919, 0x2b2b082b0808082b, - 0x2b2b082b08082b2b, 0x2b2b082b082b082b, 0x2b2b082b082b2b08, 0x2b2b082b082b2b2b, 0x2b2b082b2b08082b, - 0x2b2b082b2b082b08, 0x2b2b082b2b082b2b, 0x2b2b082b2b2b2b08, 0x2b2b190808080819, 0x2b2b190808081908, - 0x2b2b190808190808, 0x2b2b190819080808, 0x2b2b19082b082b19, 0x2b2b19082b2b1908, 0x2b2b191908080808, - 0x2b2b191908192b19, 0x2b2b192b19190819, 0x2b2b2b0808082b2b, 0x2b2b2b08082b2b08, 0x2b2b2b082b2b082b, - 0x2b2b2b1919191908, 0x2b2b2b192b08192b, 0x2b2b2b2b08082b08, 0x2b2b2b2b08082b2b, 0x2b2b2b2b082b0808, - 0x2b2b2b2b082b082b, 0x2b2b2b2b082b2b08, 0x2b2b2b2b2b082b08, 0x2b2b2b2b2b2b2b2b, -}; - -static const __device__ uint32_t iq3xxs_grid[256] = { - 0x04040404, 0x04040414, 0x04040424, 0x04040c0c, 0x04040c1c, 0x04040c3e, 0x04041404, 0x04041414, 0x04041c0c, - 0x04042414, 0x04043e1c, 0x04043e2c, 0x040c040c, 0x040c041c, 0x040c0c04, 0x040c0c14, 0x040c140c, 0x040c142c, - 0x040c1c04, 0x040c1c14, 0x040c240c, 0x040c2c24, 0x040c3e04, 0x04140404, 0x04140414, 0x04140424, 0x04140c0c, - 0x04141404, 0x04141414, 0x04141c0c, 0x04141c1c, 0x04141c3e, 0x04142c0c, 0x04142c3e, 0x04143e2c, 0x041c040c, - 0x041c043e, 0x041c0c04, 0x041c0c14, 0x041c142c, 0x041c3e04, 0x04240c1c, 0x04241c3e, 0x04242424, 0x04242c3e, - 0x04243e1c, 0x04243e2c, 0x042c040c, 0x042c043e, 0x042c1c14, 0x042c2c14, 0x04341c2c, 0x04343424, 0x043e0c04, - 0x043e0c24, 0x043e0c34, 0x043e241c, 0x043e340c, 0x0c04040c, 0x0c04041c, 0x0c040c04, 0x0c040c14, 0x0c04140c, - 0x0c04141c, 0x0c041c04, 0x0c041c14, 0x0c041c24, 0x0c04243e, 0x0c042c04, 0x0c0c0404, 0x0c0c0414, 0x0c0c0c0c, - 0x0c0c1404, 0x0c0c1414, 0x0c14040c, 0x0c14041c, 0x0c140c04, 0x0c140c14, 0x0c14140c, 0x0c141c04, 0x0c143e14, - 0x0c1c0404, 0x0c1c0414, 0x0c1c1404, 0x0c1c1c0c, 0x0c1c2434, 0x0c1c3434, 0x0c24040c, 0x0c24042c, 0x0c242c04, - 0x0c2c1404, 0x0c2c1424, 0x0c2c2434, 0x0c2c3e0c, 0x0c34042c, 0x0c3e1414, 0x0c3e2404, 0x14040404, 0x14040414, - 0x14040c0c, 0x14040c1c, 0x14041404, 0x14041414, 0x14041434, 0x14041c0c, 0x14042414, 0x140c040c, 0x140c041c, - 0x140c042c, 0x140c0c04, 0x140c0c14, 0x140c140c, 0x140c1c04, 0x140c341c, 0x140c343e, 0x140c3e04, 0x14140404, - 0x14140414, 0x14140c0c, 0x14140c3e, 0x14141404, 0x14141414, 0x14141c3e, 0x14142404, 0x14142c2c, 0x141c040c, - 0x141c0c04, 0x141c0c24, 0x141c3e04, 0x141c3e24, 0x14241c2c, 0x14242c1c, 0x142c041c, 0x142c143e, 0x142c240c, - 0x142c3e24, 0x143e040c, 0x143e041c, 0x143e0c34, 0x143e242c, 0x1c04040c, 0x1c040c04, 0x1c040c14, 0x1c04140c, - 0x1c04141c, 0x1c042c04, 0x1c04342c, 0x1c043e14, 0x1c0c0404, 0x1c0c0414, 0x1c0c1404, 0x1c0c1c0c, 0x1c0c2424, - 0x1c0c2434, 0x1c14040c, 0x1c14041c, 0x1c140c04, 0x1c14142c, 0x1c142c14, 0x1c143e14, 0x1c1c0c0c, 0x1c1c1c1c, - 0x1c241c04, 0x1c24243e, 0x1c243e14, 0x1c2c0404, 0x1c2c0434, 0x1c2c1414, 0x1c2c2c2c, 0x1c340c24, 0x1c341c34, - 0x1c34341c, 0x1c3e1c1c, 0x1c3e3404, 0x24040424, 0x24040c3e, 0x24041c2c, 0x24041c3e, 0x24042c1c, 0x24042c3e, - 0x240c3e24, 0x24141404, 0x24141c3e, 0x24142404, 0x24143404, 0x24143434, 0x241c043e, 0x241c242c, 0x24240424, - 0x24242c0c, 0x24243424, 0x242c142c, 0x242c241c, 0x242c3e04, 0x243e042c, 0x243e0c04, 0x243e0c14, 0x243e1c04, - 0x2c040c14, 0x2c04240c, 0x2c043e04, 0x2c0c0404, 0x2c0c0434, 0x2c0c1434, 0x2c0c2c2c, 0x2c140c24, 0x2c141c14, - 0x2c143e14, 0x2c1c0414, 0x2c1c2c1c, 0x2c240c04, 0x2c24141c, 0x2c24143e, 0x2c243e14, 0x2c2c0414, 0x2c2c1c0c, - 0x2c342c04, 0x2c3e1424, 0x2c3e2414, 0x34041424, 0x34042424, 0x34042434, 0x34043424, 0x340c140c, 0x340c340c, - 0x34140c3e, 0x34143424, 0x341c1c04, 0x341c1c34, 0x34242424, 0x342c042c, 0x342c2c14, 0x34341c1c, 0x343e041c, - 0x343e140c, 0x3e04041c, 0x3e04042c, 0x3e04043e, 0x3e040c04, 0x3e041c14, 0x3e042c14, 0x3e0c1434, 0x3e0c2404, - 0x3e140c14, 0x3e14242c, 0x3e142c14, 0x3e1c0404, 0x3e1c0c2c, 0x3e1c1c1c, 0x3e1c3404, 0x3e24140c, 0x3e24240c, - 0x3e2c0404, 0x3e2c0414, 0x3e2c1424, 0x3e341c04, -}; - -static const __device__ uint32_t iq3xs_grid[512] = { - 0x04040404, 0x0404040c, 0x04040414, 0x0404042c, 0x0404043e, 0x04040c04, 0x04040c0c, 0x04040c14, 0x04040c24, - 0x04040c34, 0x04041404, 0x0404140c, 0x0404142c, 0x04041c1c, 0x04042404, 0x04042414, 0x0404242c, 0x0404243e, - 0x04042c0c, 0x04042c1c, 0x04043404, 0x04043414, 0x04043e0c, 0x04043e24, 0x04043e3e, 0x040c0404, 0x040c040c, - 0x040c0414, 0x040c0424, 0x040c0c04, 0x040c0c0c, 0x040c0c2c, 0x040c1404, 0x040c141c, 0x040c143e, 0x040c1c0c, - 0x040c1c2c, 0x040c2424, 0x040c340c, 0x040c342c, 0x040c3e14, 0x04140404, 0x0414040c, 0x0414042c, 0x0414043e, - 0x04140c04, 0x04140c1c, 0x04140c34, 0x0414140c, 0x0414142c, 0x04141c04, 0x04141c24, 0x04142414, 0x0414242c, - 0x0414243e, 0x04142c0c, 0x04142c1c, 0x04143e04, 0x04143e1c, 0x041c041c, 0x041c0c0c, 0x041c0c2c, 0x041c1404, - 0x041c1414, 0x041c1c0c, 0x041c1c1c, 0x041c1c34, 0x041c2424, 0x041c2c04, 0x041c2c14, 0x041c343e, 0x041c3e0c, - 0x041c3e2c, 0x04240404, 0x04240c1c, 0x04240c3e, 0x0424140c, 0x04241424, 0x04241c14, 0x04242404, 0x0424241c, - 0x04242c0c, 0x04243e04, 0x042c0414, 0x042c0424, 0x042c1404, 0x042c1414, 0x042c1434, 0x042c1c1c, 0x042c240c, - 0x042c242c, 0x042c243e, 0x042c3434, 0x042c3e1c, 0x04340434, 0x04340c0c, 0x04340c1c, 0x04341c0c, 0x04342c14, - 0x04343e0c, 0x043e0404, 0x043e0414, 0x043e0424, 0x043e1404, 0x043e1414, 0x043e1434, 0x043e1c1c, 0x043e2c04, - 0x043e2c24, 0x0c040404, 0x0c04040c, 0x0c040414, 0x0c040424, 0x0c040c04, 0x0c040c0c, 0x0c040c1c, 0x0c040c2c, - 0x0c040c3e, 0x0c041404, 0x0c041414, 0x0c041c0c, 0x0c041c24, 0x0c041c34, 0x0c042c24, 0x0c042c34, 0x0c04340c, - 0x0c043e14, 0x0c0c0404, 0x0c0c040c, 0x0c0c041c, 0x0c0c0434, 0x0c0c0c04, 0x0c0c0c24, 0x0c0c140c, 0x0c0c1c04, - 0x0c0c1c1c, 0x0c0c240c, 0x0c0c2c04, 0x0c0c2c14, 0x0c0c3e04, 0x0c0c3e34, 0x0c140404, 0x0c140c14, 0x0c140c2c, - 0x0c140c3e, 0x0c141404, 0x0c141424, 0x0c141c14, 0x0c142404, 0x0c14241c, 0x0c142c2c, 0x0c143404, 0x0c143e14, - 0x0c1c040c, 0x0c1c0424, 0x0c1c043e, 0x0c1c0c04, 0x0c1c0c1c, 0x0c1c140c, 0x0c1c143e, 0x0c1c1c04, 0x0c1c1c24, - 0x0c1c240c, 0x0c1c3414, 0x0c1c3e04, 0x0c24041c, 0x0c24042c, 0x0c240c14, 0x0c240c24, 0x0c241c0c, 0x0c241c1c, - 0x0c242414, 0x0c242434, 0x0c242c04, 0x0c242c24, 0x0c2c040c, 0x0c2c0c04, 0x0c2c0c1c, 0x0c2c140c, 0x0c2c1c04, - 0x0c2c1c14, 0x0c2c2c0c, 0x0c341404, 0x0c341424, 0x0c34143e, 0x0c342424, 0x0c342434, 0x0c3e040c, 0x0c3e041c, - 0x0c3e0c04, 0x0c3e0c14, 0x0c3e140c, 0x0c3e1c2c, 0x0c3e240c, 0x0c3e3414, 0x0c3e3e04, 0x14040404, 0x1404040c, - 0x1404041c, 0x1404042c, 0x1404043e, 0x14040c04, 0x14040c14, 0x14040c24, 0x14040c34, 0x1404140c, 0x1404141c, - 0x1404143e, 0x14041c04, 0x14041c14, 0x1404240c, 0x1404241c, 0x1404242c, 0x14042c04, 0x14042c14, 0x1404343e, - 0x14043e04, 0x14043e1c, 0x14043e2c, 0x140c0404, 0x140c0414, 0x140c0c04, 0x140c0c1c, 0x140c0c3e, 0x140c1414, - 0x140c142c, 0x140c1c0c, 0x140c1c24, 0x140c2414, 0x140c2c0c, 0x1414040c, 0x14140424, 0x1414043e, 0x1414140c, - 0x1414141c, 0x14141c04, 0x14141c3e, 0x1414240c, 0x14142c1c, 0x14142c3e, 0x14143e0c, 0x14143e24, 0x141c0404, - 0x141c0414, 0x141c042c, 0x141c0c0c, 0x141c1414, 0x141c1424, 0x141c1c0c, 0x141c1c1c, 0x141c2414, 0x141c2c04, - 0x141c3434, 0x1424040c, 0x1424043e, 0x14241404, 0x1424141c, 0x14241c14, 0x14241c2c, 0x1424240c, 0x14243e14, - 0x14243e2c, 0x142c0424, 0x142c0c0c, 0x142c1414, 0x142c1c3e, 0x142c2404, 0x142c2c1c, 0x142c3e04, 0x14340404, - 0x14340414, 0x1434043e, 0x1434140c, 0x14342c2c, 0x1434340c, 0x143e042c, 0x143e0c0c, 0x143e1434, 0x143e1c04, - 0x143e241c, 0x143e2c04, 0x1c040414, 0x1c040c0c, 0x1c040c1c, 0x1c040c2c, 0x1c040c3e, 0x1c041414, 0x1c041c0c, - 0x1c041c1c, 0x1c041c2c, 0x1c042414, 0x1c042424, 0x1c04243e, 0x1c042c0c, 0x1c04341c, 0x1c043e0c, 0x1c0c040c, - 0x1c0c041c, 0x1c0c042c, 0x1c0c0c24, 0x1c0c140c, 0x1c0c141c, 0x1c0c2404, 0x1c0c3404, 0x1c0c3e14, 0x1c0c3e34, - 0x1c140404, 0x1c140c14, 0x1c141404, 0x1c141c14, 0x1c141c24, 0x1c142c04, 0x1c1c040c, 0x1c1c0c04, 0x1c1c0c24, - 0x1c1c140c, 0x1c1c141c, 0x1c1c143e, 0x1c1c1c04, 0x1c1c240c, 0x1c1c241c, 0x1c1c243e, 0x1c1c2c2c, 0x1c1c3e1c, - 0x1c24041c, 0x1c240c0c, 0x1c240c34, 0x1c241414, 0x1c241c0c, 0x1c242c14, 0x1c243404, 0x1c243424, 0x1c2c040c, - 0x1c2c0c04, 0x1c2c0c14, 0x1c2c142c, 0x1c2c1c14, 0x1c2c2424, 0x1c2c2c34, 0x1c2c3e1c, 0x1c340c34, 0x1c34240c, - 0x1c3e040c, 0x1c3e041c, 0x1c3e1404, 0x1c3e1414, 0x1c3e1c2c, 0x24040404, 0x24040424, 0x24040c14, 0x24041404, - 0x24041424, 0x2404143e, 0x24041c14, 0x2404240c, 0x24042c04, 0x24043e04, 0x240c0414, 0x240c043e, 0x240c0c0c, - 0x240c0c1c, 0x240c1414, 0x240c1c04, 0x240c1c2c, 0x240c241c, 0x240c2c0c, 0x240c2c2c, 0x2414040c, 0x2414041c, - 0x24140c04, 0x24140c2c, 0x2414140c, 0x24141c1c, 0x24142404, 0x24142c3e, 0x24143414, 0x24143e04, 0x241c0424, - 0x241c0c0c, 0x241c0c1c, 0x241c1404, 0x241c1414, 0x241c1c0c, 0x241c1c2c, 0x24240404, 0x24240414, 0x24241424, - 0x24241c3e, 0x24242404, 0x24243e0c, 0x242c042c, 0x242c043e, 0x242c140c, 0x242c3414, 0x24340c1c, 0x24341c24, - 0x24343404, 0x243e0c04, 0x243e0c2c, 0x243e1c04, 0x243e241c, 0x243e2c0c, 0x2c040414, 0x2c040c04, 0x2c040c24, - 0x2c041414, 0x2c042404, 0x2c042424, 0x2c04243e, 0x2c042c14, 0x2c043434, 0x2c043e24, 0x2c0c040c, 0x2c0c041c, - 0x2c0c042c, 0x2c0c0c14, 0x2c0c140c, 0x2c0c1c14, 0x2c0c3e14, 0x2c140404, 0x2c140c0c, 0x2c14141c, 0x2c141c04, - 0x2c141c34, 0x2c142c1c, 0x2c1c0414, 0x2c1c043e, 0x2c1c0c04, 0x2c1c143e, 0x2c1c2424, 0x2c1c2c0c, 0x2c1c342c, - 0x2c1c3e1c, 0x2c24040c, 0x2c240424, 0x2c241404, 0x2c241c14, 0x2c242434, 0x2c2c0c14, 0x2c2c1434, 0x2c2c2c0c, - 0x2c2c2c1c, 0x2c342414, 0x2c3e0414, 0x2c3e0424, 0x2c3e1414, 0x34040c0c, 0x34040c1c, 0x34040c2c, 0x34041c0c, - 0x34041c1c, 0x34043404, 0x340c0404, 0x340c1404, 0x340c143e, 0x340c3424, 0x34140c14, 0x34141c24, 0x34142414, - 0x34142c2c, 0x34143414, 0x34143e04, 0x341c0404, 0x341c0c24, 0x341c140c, 0x341c2404, 0x3424142c, 0x3424241c, - 0x34243414, 0x342c0404, 0x342c041c, 0x342c1c24, 0x342c3404, 0x3434042c, 0x34342404, 0x343e0c0c, 0x343e0c1c, - 0x3e040404, 0x3e040424, 0x3e04043e, 0x3e041404, 0x3e041414, 0x3e041c34, 0x3e042404, 0x3e042c24, 0x3e043414, - 0x3e0c0414, 0x3e0c0c0c, 0x3e0c1424, 0x3e0c241c, 0x3e0c242c, 0x3e14040c, 0x3e140424, 0x3e140c04, 0x3e140c34, - 0x3e14140c, 0x3e141c04, 0x3e142c0c, 0x3e1c0414, 0x3e1c1c14, 0x3e1c1c2c, 0x3e1c2c1c, 0x3e24040c, 0x3e24042c, - 0x3e240c1c, 0x3e241404, 0x3e242c04, 0x3e2c1414, 0x3e2c2414, 0x3e340414, 0x3e341c0c, 0x3e3e0404, -}; - -#define IQ1S_DELTA 0.125f -#define IQ1M_DELTA 0.125f -static const __device__ uint64_t iq1s_grid_gpu[2048] = { - 0x00000000, 0x00000002, 0x00000101, 0x00000200, 0x00000202, 0x00010001, 0x00010101, 0x00020000, 0x00020002, - 0x00020200, 0x00020202, 0x01000101, 0x01010001, 0x01010100, 0x01010102, 0x01020101, 0x02000000, 0x02000002, - 0x02000200, 0x02000202, 0x02010101, 0x02020000, 0x02020002, 0x02020200, 0x02020202, 0x00000110, 0x00000111, - 0x00010011, 0x00010110, 0x00010112, 0x00010211, 0x00010212, 0x00020111, 0x01000011, 0x01000112, 0x01000211, - 0x01010012, 0x01010111, 0x01010212, 0x01020011, 0x01020110, 0x01020112, 0x01020210, 0x02000111, 0x02010011, - 0x02010110, 0x02010112, 0x02020111, 0x00000020, 0x00000022, 0x00000220, 0x00000222, 0x00010121, 0x00020020, - 0x00020022, 0x00020220, 0x00020222, 0x01000121, 0x01010021, 0x01010221, 0x01020120, 0x01020221, 0x02000020, - 0x02000022, 0x02000220, 0x02000222, 0x02010021, 0x02010121, 0x02010221, 0x02020020, 0x02020022, 0x02020220, - 0x02020222, 0x00011001, 0x00011100, 0x00011102, 0x00021101, 0x01001001, 0x01001201, 0x01011101, 0x01011202, - 0x01021100, 0x01021101, 0x02011001, 0x02011201, 0x02021101, 0x00001011, 0x00001110, 0x00001111, 0x00001112, - 0x00011111, 0x00011210, 0x00011212, 0x00021211, 0x01001010, 0x01001111, 0x01001212, 0x01011010, 0x01011011, - 0x01011110, 0x01011111, 0x01011112, 0x01011211, 0x01021010, 0x01021012, 0x01021111, 0x01021210, 0x01021212, - 0x02001011, 0x02011011, 0x02011111, 0x02011210, 0x02011212, 0x02021011, 0x02021110, 0x02021111, 0x02021112, - 0x02021211, 0x00011120, 0x00011221, 0x01001021, 0x01001120, 0x01011020, 0x01011022, 0x01011121, 0x01011220, - 0x01021020, 0x01021021, 0x01021122, 0x01021221, 0x02001121, 0x02011021, 0x02011120, 0x02011221, 0x00002000, - 0x00002002, 0x00002200, 0x00002202, 0x00012101, 0x00022000, 0x00022002, 0x00022200, 0x00022202, 0x01002101, - 0x01012001, 0x01012102, 0x01022101, 0x02002000, 0x02002002, 0x02002200, 0x02002202, 0x02012101, 0x02022000, - 0x02022002, 0x02022200, 0x02022202, 0x00002111, 0x00012011, 0x00012110, 0x00012211, 0x00022110, 0x00022111, - 0x01002011, 0x01012010, 0x01012011, 0x01012111, 0x01022011, 0x01022110, 0x01022211, 0x02012011, 0x02012110, - 0x02012112, 0x02012211, 0x02022111, 0x00002020, 0x00002022, 0x00002220, 0x00002222, 0x00012121, 0x00022020, - 0x00022022, 0x00022220, 0x00022222, 0x01002121, 0x01012021, 0x01012221, 0x01022021, 0x01022121, 0x02002020, - 0x02002022, 0x02002121, 0x02002220, 0x02002222, 0x02012121, 0x02022020, 0x02022022, 0x02022220, 0x02022222, - 0x00110000, 0x00110001, 0x00110100, 0x00110201, 0x00120100, 0x00120101, 0x01100001, 0x01100100, 0x01110000, - 0x01110101, 0x01110200, 0x01120001, 0x01120100, 0x01120101, 0x01120201, 0x02110001, 0x02110100, 0x02110102, - 0x02120001, 0x02120101, 0x00100011, 0x00100110, 0x00100112, 0x00100211, 0x00110010, 0x00110012, 0x00110111, - 0x00110210, 0x00120011, 0x00120110, 0x00120211, 0x01100111, 0x01100212, 0x01110010, 0x01110011, 0x01110012, - 0x01110110, 0x01110111, 0x01110112, 0x01110211, 0x01120010, 0x01120111, 0x02100110, 0x02110012, 0x02110111, - 0x02120011, 0x02120110, 0x00110021, 0x00110120, 0x00110122, 0x00120121, 0x01100020, 0x01100122, 0x01100221, - 0x01110022, 0x01110121, 0x01110220, 0x01110222, 0x01120120, 0x01120122, 0x02100121, 0x02110021, 0x02110120, - 0x02110122, 0x02120121, 0x00101001, 0x00101102, 0x00101201, 0x00111100, 0x00111101, 0x00111200, 0x00111201, - 0x00121001, 0x00121102, 0x01101001, 0x01101101, 0x01101102, 0x01101200, 0x01101202, 0x01111001, 0x01111100, - 0x01111101, 0x01111102, 0x01111201, 0x01121002, 0x01121101, 0x01121200, 0x02101100, 0x02101201, 0x02111000, - 0x02111100, 0x02111101, 0x02111200, 0x02111201, 0x02111202, 0x02121001, 0x02121100, 0x02121101, 0x02121201, - 0x00101012, 0x00101111, 0x00101212, 0x00111011, 0x00111110, 0x00111111, 0x00111112, 0x00111211, 0x00121010, - 0x00121012, 0x00121111, 0x00121210, 0x00121212, 0x01101011, 0x01101110, 0x01101111, 0x01101112, 0x01111011, - 0x01111012, 0x01111110, 0x01111111, 0x01111112, 0x01111211, 0x01111212, 0x01121011, 0x01121110, 0x01121111, - 0x01121112, 0x01121211, 0x02101010, 0x02101012, 0x02101110, 0x02101111, 0x02101210, 0x02101212, 0x02111010, - 0x02111011, 0x02111110, 0x02111111, 0x02111112, 0x02111211, 0x02111212, 0x02121010, 0x02121012, 0x02121111, - 0x00101021, 0x00101120, 0x00101121, 0x00101122, 0x00111121, 0x00111122, 0x00111220, 0x00111222, 0x00121021, - 0x00121122, 0x01101020, 0x01101022, 0x01101120, 0x01101121, 0x01101220, 0x01101222, 0x01111021, 0x01111121, - 0x01111122, 0x01111220, 0x01111221, 0x01121021, 0x01121120, 0x01121121, 0x01121220, 0x01121221, 0x01121222, - 0x02101122, 0x02101222, 0x02111022, 0x02111121, 0x02121120, 0x02121221, 0x00112001, 0x00112102, 0x00122101, - 0x01102001, 0x01102100, 0x01102102, 0x01102201, 0x01112000, 0x01112101, 0x01112200, 0x01112202, 0x01122000, - 0x01122001, 0x01122100, 0x01122102, 0x01122201, 0x02102101, 0x02112001, 0x02112100, 0x02122101, 0x00112010, - 0x00112012, 0x00112111, 0x00112212, 0x00122011, 0x00122111, 0x01102012, 0x01102110, 0x01102111, 0x01102210, - 0x01112011, 0x01112110, 0x01112111, 0x01112112, 0x01112211, 0x01112212, 0x01122010, 0x01122111, 0x01122212, - 0x02102211, 0x02112011, 0x02112012, 0x02112111, 0x02112210, 0x02122011, 0x02122112, 0x02122211, 0x00102221, - 0x00112122, 0x00122120, 0x00122122, 0x01102120, 0x01102122, 0x01102221, 0x01112020, 0x01112022, 0x01112121, - 0x01112220, 0x01122021, 0x01122122, 0x01122221, 0x02102121, 0x02112021, 0x02112122, 0x02112222, 0x00200000, - 0x00200002, 0x00200200, 0x00200202, 0x00210101, 0x00220000, 0x00220002, 0x00220101, 0x00220200, 0x00220202, - 0x01200101, 0x01210001, 0x01210201, 0x01220001, 0x01220101, 0x02200000, 0x02200002, 0x02200200, 0x02200202, - 0x02210101, 0x02220000, 0x02220002, 0x02220101, 0x02220200, 0x02220202, 0x00200111, 0x00210011, 0x00210110, - 0x00210211, 0x00220111, 0x01200012, 0x01200110, 0x01200211, 0x01210111, 0x01210210, 0x01210212, 0x01220011, - 0x01220110, 0x01220111, 0x01220112, 0x02200111, 0x02210010, 0x02210112, 0x02210211, 0x02220111, 0x00200021, - 0x00200220, 0x00200222, 0x00210021, 0x00210121, 0x00220020, 0x00220022, 0x00220220, 0x00220222, 0x01200121, - 0x01210021, 0x01210122, 0x01210221, 0x01220121, 0x02200021, 0x02200220, 0x02200222, 0x02210021, 0x02210121, - 0x02220020, 0x02220022, 0x02220220, 0x02220222, 0x00201101, 0x00211100, 0x00211102, 0x00211201, 0x00221101, - 0x01201100, 0x01201101, 0x01201102, 0x01201201, 0x01211002, 0x01211101, 0x01211200, 0x01211202, 0x01221102, - 0x02201101, 0x02211001, 0x02211100, 0x02211201, 0x02221001, 0x02221101, 0x00201211, 0x00211111, 0x00221011, - 0x00221211, 0x01201010, 0x01201111, 0x01201210, 0x01211011, 0x01211110, 0x01211111, 0x01211211, 0x01221012, - 0x01221111, 0x01221210, 0x02201211, 0x02211010, 0x02211110, 0x02211111, 0x02211210, 0x02211212, 0x02221011, - 0x02221110, 0x02221112, 0x02221211, 0x00201121, 0x00211020, 0x00211022, 0x00211221, 0x00221121, 0x01201021, - 0x01201221, 0x01211121, 0x01221020, 0x01221021, 0x01221221, 0x02201120, 0x02201122, 0x02211020, 0x02211222, - 0x00202000, 0x00202002, 0x00202200, 0x00202202, 0x00212101, 0x00222000, 0x00222002, 0x00222200, 0x00222202, - 0x01202101, 0x01212001, 0x01212100, 0x01222101, 0x02202000, 0x02202002, 0x02202200, 0x02202202, 0x02222000, - 0x02222002, 0x02222200, 0x02222202, 0x00202211, 0x00212011, 0x00212110, 0x00212211, 0x00222111, 0x01202112, - 0x01202211, 0x01212012, 0x01212111, 0x01222011, 0x01222110, 0x01222112, 0x01222211, 0x02202111, 0x02212010, - 0x02212112, 0x02212211, 0x02222110, 0x02222111, 0x00202020, 0x00202022, 0x00202220, 0x00202222, 0x00222020, - 0x00222022, 0x00222220, 0x00222222, 0x01202121, 0x01212021, 0x01212122, 0x01212221, 0x01222121, 0x02202020, - 0x02202022, 0x02202220, 0x02202222, 0x02212121, 0x02222020, 0x02222022, 0x02222220, 0x02222222, 0x10000101, - 0x10010001, 0x10010102, 0x10020101, 0x11000201, 0x11010002, 0x11010101, 0x11010200, 0x11010202, 0x11020001, - 0x11020100, 0x11020102, 0x12010100, 0x12010201, 0x12020001, 0x12020102, 0x10000010, 0x10000011, 0x10000110, - 0x10000112, 0x10000211, 0x10010012, 0x10010111, 0x10010112, 0x10010210, 0x10010212, 0x10020011, 0x10020112, - 0x10020211, 0x11000111, 0x11000210, 0x11000212, 0x11010011, 0x11010110, 0x11010111, 0x11010112, 0x11010211, - 0x11010212, 0x11020111, 0x11020210, 0x11020212, 0x12000011, 0x12000110, 0x12000112, 0x12010010, 0x12010012, - 0x12010111, 0x12020010, 0x12020011, 0x12020012, 0x10000121, 0x10010021, 0x10010120, 0x10010122, 0x10020121, - 0x11000021, 0x11010022, 0x11010121, 0x11010222, 0x11020120, 0x11020221, 0x12000221, 0x12010120, 0x12020121, - 0x10001001, 0x10011101, 0x10011201, 0x10021201, 0x11001101, 0x11001200, 0x11001202, 0x11011001, 0x11011100, - 0x11011101, 0x11011102, 0x11021001, 0x11021002, 0x11021101, 0x11021200, 0x11021202, 0x12001001, 0x12001102, - 0x12001201, 0x12011000, 0x12011002, 0x12011101, 0x12021000, 0x12021001, 0x12021201, 0x10001011, 0x10001012, - 0x10001111, 0x10001212, 0x10011011, 0x10011110, 0x10011111, 0x10011112, 0x10011211, 0x10021010, 0x10021111, - 0x10021212, 0x11001011, 0x11001110, 0x11001111, 0x11001112, 0x11001211, 0x11011010, 0x11011011, 0x11011110, - 0x11011111, 0x11011112, 0x11011210, 0x11011211, 0x11021011, 0x11021110, 0x11021111, 0x11021112, 0x11021211, - 0x12001012, 0x12001110, 0x12001111, 0x12001210, 0x12011011, 0x12011110, 0x12011111, 0x12011112, 0x12011211, - 0x12011212, 0x12021111, 0x12021210, 0x12021212, 0x10001021, 0x10001121, 0x10001221, 0x10011120, 0x10011121, - 0x10011220, 0x10011222, 0x10021021, 0x10021120, 0x10021221, 0x11001020, 0x11001022, 0x11001121, 0x11001220, - 0x11011020, 0x11011021, 0x11011022, 0x11011121, 0x11011122, 0x11011221, 0x11021022, 0x11021121, 0x11021220, - 0x12001021, 0x12001121, 0x12001222, 0x12011120, 0x12011121, 0x12021021, 0x12021120, 0x12021122, 0x10002101, - 0x10012001, 0x10012101, 0x10012202, 0x10022101, 0x11002002, 0x11002201, 0x11012000, 0x11012101, 0x11012200, - 0x11022001, 0x11022100, 0x11022102, 0x11022201, 0x12002101, 0x12012001, 0x12012100, 0x12012102, 0x12012201, - 0x12022101, 0x10002011, 0x10002111, 0x10002112, 0x10002212, 0x10012010, 0x10012110, 0x10012111, 0x10012210, - 0x10022011, 0x10022110, 0x10022112, 0x11002010, 0x11002111, 0x11002212, 0x11012011, 0x11012012, 0x11012110, - 0x11012111, 0x11012112, 0x11012211, 0x11022010, 0x11022012, 0x11022111, 0x11022112, 0x11022212, 0x12002112, - 0x12002211, 0x12012012, 0x12012111, 0x12012112, 0x12012210, 0x12022011, 0x12022110, 0x12022112, 0x12022211, - 0x10012122, 0x11002120, 0x11002122, 0x11002221, 0x11012121, 0x11012220, 0x11012222, 0x11022120, 0x11022221, - 0x12012120, 0x12022121, 0x10100001, 0x10100100, 0x10100101, 0x10100102, 0x10100201, 0x10110002, 0x10110101, - 0x10110202, 0x10120001, 0x10120100, 0x10120201, 0x11100000, 0x11100101, 0x11100200, 0x11110001, 0x11110100, - 0x11110101, 0x11110102, 0x11110201, 0x11120101, 0x11120200, 0x12100102, 0x12100201, 0x12110101, 0x12110200, - 0x12120000, 0x12120001, 0x12120102, 0x12120201, 0x10100111, 0x10100210, 0x10100211, 0x10100212, 0x10110011, - 0x10110110, 0x10110111, 0x10110112, 0x10110210, 0x10110211, 0x10120010, 0x10120111, 0x10120112, 0x10120210, - 0x10120212, 0x11100011, 0x11100110, 0x11100111, 0x11100112, 0x11100211, 0x11110010, 0x11110011, 0x11110012, - 0x11110110, 0x11110111, 0x11110112, 0x11110210, 0x11110211, 0x11110212, 0x11120011, 0x11120110, 0x11120111, - 0x11120112, 0x11120211, 0x12100012, 0x12100111, 0x12110011, 0x12110110, 0x12110111, 0x12110112, 0x12110211, - 0x12120010, 0x12120111, 0x12120212, 0x10100021, 0x10100122, 0x10110022, 0x10110121, 0x10110222, 0x10120021, - 0x10120120, 0x11100022, 0x11100121, 0x11100222, 0x11110021, 0x11110120, 0x11110121, 0x11110122, 0x11110221, - 0x11120022, 0x11120121, 0x12100121, 0x12110020, 0x12110022, 0x12110121, 0x12110221, 0x12110222, 0x12120120, - 0x10101100, 0x10101101, 0x10111001, 0x10111100, 0x10111101, 0x10111102, 0x10111200, 0x10111201, 0x10121001, - 0x10121101, 0x10121200, 0x10121202, 0x11101001, 0x11101100, 0x11101101, 0x11101102, 0x11101201, 0x11101202, - 0x11111000, 0x11111001, 0x11111100, 0x11111101, 0x11111102, 0x11111200, 0x11111201, 0x11111202, 0x11121001, - 0x11121002, 0x11121100, 0x11121101, 0x11121102, 0x11121201, 0x12101000, 0x12101200, 0x12101202, 0x12111001, - 0x12111100, 0x12111101, 0x12111102, 0x12111201, 0x12121001, 0x12121100, 0x12121101, 0x12121202, 0x10101011, - 0x10101012, 0x10101110, 0x10101111, 0x10101112, 0x10101211, 0x10111010, 0x10111011, 0x10111012, 0x10111110, - 0x10111111, 0x10111112, 0x10111211, 0x10111212, 0x10121011, 0x10121110, 0x10121111, 0x10121112, 0x10121211, - 0x11101010, 0x11101011, 0x11101012, 0x11101110, 0x11101111, 0x11101112, 0x11101210, 0x11101211, 0x11111010, - 0x11111011, 0x11111012, 0x11111110, 0x11111111, 0x11111112, 0x11111210, 0x11111211, 0x11111212, 0x11121010, - 0x11121011, 0x11121110, 0x11121111, 0x11121112, 0x11121210, 0x11121211, 0x11121212, 0x12101011, 0x12101110, - 0x12101111, 0x12101211, 0x12101212, 0x12111010, 0x12111011, 0x12111110, 0x12111111, 0x12111112, 0x12111210, - 0x12111211, 0x12121011, 0x12121110, 0x12121111, 0x12121112, 0x12121211, 0x10101020, 0x10101021, 0x10101022, - 0x10101120, 0x10101122, 0x10101220, 0x10101221, 0x10111021, 0x10111120, 0x10111121, 0x10111220, 0x10111221, - 0x10121020, 0x10121021, 0x10121022, 0x10121120, 0x10121121, 0x10121122, 0x10121220, 0x10121221, 0x11101021, - 0x11101121, 0x11101122, 0x11101220, 0x11101221, 0x11101222, 0x11111020, 0x11111021, 0x11111022, 0x11111120, - 0x11111121, 0x11111122, 0x11111220, 0x11111221, 0x11111222, 0x11121021, 0x11121120, 0x11121121, 0x11121221, - 0x12101022, 0x12101121, 0x12101122, 0x12101220, 0x12101221, 0x12101222, 0x12111021, 0x12111121, 0x12111222, - 0x12121022, 0x12121121, 0x12121122, 0x12121220, 0x12121221, 0x10102100, 0x10102101, 0x10102102, 0x10102201, - 0x10112000, 0x10112101, 0x10112200, 0x10122001, 0x10122202, 0x11102101, 0x11102200, 0x11102202, 0x11112001, - 0x11112100, 0x11112101, 0x11112102, 0x11112200, 0x11112201, 0x11122000, 0x11122002, 0x11122100, 0x11122101, - 0x12102002, 0x12102201, 0x12112000, 0x12112002, 0x12112101, 0x12112200, 0x12122001, 0x12122201, 0x10102011, - 0x10102012, 0x10102111, 0x10102212, 0x10112011, 0x10112110, 0x10112111, 0x10112112, 0x10112211, 0x10122111, - 0x11102011, 0x11102110, 0x11102111, 0x11102112, 0x11102211, 0x11112010, 0x11112011, 0x11112012, 0x11112110, - 0x11112111, 0x11112112, 0x11112210, 0x11112211, 0x11112212, 0x11122011, 0x11122110, 0x11122111, 0x11122112, - 0x11122211, 0x12102011, 0x12102111, 0x12102211, 0x12112011, 0x12112110, 0x12112111, 0x12112112, 0x12112210, - 0x12112211, 0x12122111, 0x10102120, 0x10102220, 0x10112121, 0x10112222, 0x10122020, 0x10122121, 0x10122122, - 0x10122221, 0x11102121, 0x11102220, 0x11102221, 0x11112021, 0x11112121, 0x11112122, 0x11112220, 0x11112221, - 0x11122022, 0x11122121, 0x11122220, 0x11122222, 0x12102021, 0x12102222, 0x12112022, 0x12112121, 0x12112122, - 0x12112220, 0x12112222, 0x12122021, 0x10200101, 0x10210100, 0x10210102, 0x10210201, 0x10220101, 0x11200100, - 0x11210000, 0x11210101, 0x11210102, 0x11210200, 0x11210202, 0x11220001, 0x11220100, 0x11220102, 0x11220201, - 0x12200001, 0x12210102, 0x12220101, 0x10200011, 0x10200110, 0x10200112, 0x10200211, 0x10210012, 0x10210111, - 0x10220011, 0x10220012, 0x10220112, 0x10220211, 0x11200111, 0x11200211, 0x11210011, 0x11210111, 0x11210112, - 0x11210211, 0x11220111, 0x11220112, 0x11220212, 0x12200110, 0x12200212, 0x12210012, 0x12210111, 0x12220011, - 0x12220112, 0x12220211, 0x10210021, 0x10210122, 0x10210221, 0x11200020, 0x11200021, 0x11200122, 0x11210121, - 0x11210122, 0x11210220, 0x11220020, 0x12200121, 0x12210021, 0x12210122, 0x12220121, 0x10211001, 0x10211002, - 0x10211101, 0x10211102, 0x10211202, 0x10221001, 0x10221102, 0x10221201, 0x11201000, 0x11201002, 0x11201101, - 0x11201200, 0x11201202, 0x11211001, 0x11211100, 0x11211101, 0x11211102, 0x11211201, 0x11211202, 0x11221000, - 0x11221002, 0x11221101, 0x12201100, 0x12201101, 0x12201201, 0x12211000, 0x12211002, 0x12211100, 0x12211101, - 0x12211102, 0x12211200, 0x12211202, 0x12221001, 0x12221100, 0x12221201, 0x10201111, 0x10201210, 0x10201212, - 0x10211011, 0x10211111, 0x10211112, 0x10211211, 0x11201110, 0x11201111, 0x11201112, 0x11201211, 0x11211010, - 0x11211011, 0x11211110, 0x11211111, 0x11211112, 0x11211211, 0x11221011, 0x11221110, 0x11221111, 0x11221112, - 0x11221211, 0x12201112, 0x12201211, 0x12201212, 0x12211011, 0x12211111, 0x12211112, 0x12211211, 0x12211212, - 0x12221012, 0x12221111, 0x12221112, 0x12221210, 0x10201022, 0x10201221, 0x10211121, 0x10221020, 0x10221122, - 0x10221220, 0x10221221, 0x11201020, 0x11201121, 0x11201220, 0x11201222, 0x11211021, 0x11211120, 0x11211121, - 0x11211122, 0x11211220, 0x11211222, 0x11221020, 0x11221121, 0x11221220, 0x12201020, 0x12201022, 0x12201121, - 0x12201222, 0x12211120, 0x12211122, 0x12211220, 0x12211221, 0x12221020, 0x12221120, 0x12221122, 0x12221222, - 0x10212102, 0x10212201, 0x10222101, 0x11202001, 0x11212002, 0x11212101, 0x11212202, 0x11222001, 0x11222201, - 0x12202101, 0x12212001, 0x12212200, 0x12222102, 0x10202011, 0x10202110, 0x10212010, 0x10212111, 0x10222011, - 0x10222110, 0x10222112, 0x10222211, 0x11202010, 0x11202011, 0x11202111, 0x11202112, 0x11202210, 0x11212011, - 0x11212110, 0x11212111, 0x11212112, 0x11212211, 0x11222010, 0x11222111, 0x11222212, 0x12202012, 0x12202110, - 0x12202212, 0x12212111, 0x12222011, 0x12222110, 0x12222111, 0x12222211, 0x10212021, 0x10212122, 0x10212220, - 0x11202021, 0x11202120, 0x11202221, 0x11212020, 0x11212121, 0x11212220, 0x11212222, 0x11222120, 0x11222121, - 0x11222221, 0x12202122, 0x12212120, 0x12212220, 0x12212222, 0x12222122, 0x20000000, 0x20000002, 0x20000200, - 0x20000202, 0x20020000, 0x20020002, 0x20020200, 0x20020202, 0x21000101, 0x21010000, 0x21010001, 0x21010100, - 0x21010102, 0x21010201, 0x21020101, 0x22000000, 0x22000002, 0x22000200, 0x22000202, 0x22010101, 0x22020000, - 0x22020002, 0x22020200, 0x22020202, 0x20000111, 0x20010011, 0x20010110, 0x20010112, 0x20010211, 0x20020111, - 0x21000011, 0x21000110, 0x21000211, 0x21010010, 0x21010012, 0x21010111, 0x21010112, 0x21010210, 0x21010211, - 0x21020110, 0x21020112, 0x21020211, 0x22000111, 0x22000211, 0x22010110, 0x22010112, 0x22010211, 0x22020111, - 0x20000020, 0x20000022, 0x20000220, 0x20000222, 0x20010121, 0x20020020, 0x20020022, 0x20020220, 0x20020222, - 0x21010021, 0x21010120, 0x21010221, 0x21020121, 0x22000020, 0x22000022, 0x22000220, 0x22000222, 0x22010121, - 0x22020020, 0x22020022, 0x22020220, 0x22020222, 0x20011100, 0x20011201, 0x21001001, 0x21001100, 0x21011001, - 0x21011101, 0x21011202, 0x21021001, 0x21021100, 0x21021201, 0x22011100, 0x22011201, 0x20001011, 0x20001211, - 0x20011012, 0x20011111, 0x20011212, 0x20021112, 0x20021211, 0x21001010, 0x21001011, 0x21001111, 0x21001210, - 0x21011011, 0x21011110, 0x21011111, 0x21011112, 0x21011211, 0x21011212, 0x21021111, 0x21021112, 0x21021210, - 0x21021212, 0x22001011, 0x22001110, 0x22001112, 0x22001211, 0x22011010, 0x22011012, 0x22011111, 0x22011210, - 0x22021112, 0x20011021, 0x20011122, 0x20011221, 0x20021121, 0x21001021, 0x21001120, 0x21001221, 0x21001222, - 0x21011020, 0x21011121, 0x21011221, 0x21011222, 0x21021021, 0x21021122, 0x21021222, 0x22001121, 0x22011021, - 0x22011222, 0x22021120, 0x20002000, 0x20002002, 0x20002200, 0x20002202, 0x20012101, 0x20022000, 0x20022002, - 0x20022200, 0x20022202, 0x21002001, 0x21002101, 0x21012001, 0x21012100, 0x21012201, 0x21022101, 0x21022201, - 0x22002000, 0x22002002, 0x22002200, 0x22002202, 0x22012101, 0x22022000, 0x22022002, 0x22022200, 0x22022202, - 0x20002111, 0x20002112, 0x20012011, 0x20012110, 0x20012112, 0x20022111, 0x21002011, 0x21002110, 0x21002112, - 0x21002211, 0x21012010, 0x21012012, 0x21012111, 0x21012212, 0x21022011, 0x21022110, 0x22002111, 0x22012112, - 0x22012211, 0x22022111, 0x20002020, 0x20002022, 0x20002220, 0x20002222, 0x20012121, 0x20022020, 0x20022022, - 0x20022220, 0x20022222, 0x21002121, 0x21012021, 0x21012120, 0x21012122, 0x22002020, 0x22002022, 0x22002220, - 0x22002222, 0x22012121, 0x22022020, 0x22022022, 0x22022220, 0x22022222, 0x20100101, 0x20110001, 0x20110102, - 0x20110200, 0x20110201, 0x20120101, 0x21100001, 0x21100102, 0x21100201, 0x21110101, 0x21110200, 0x21110202, - 0x21120201, 0x21120202, 0x22100101, 0x22110001, 0x22110100, 0x22110102, 0x22110201, 0x22120101, 0x20100011, - 0x20100110, 0x20100112, 0x20100211, 0x20110010, 0x20110111, 0x20110210, 0x20110212, 0x20120011, 0x20120110, - 0x20120112, 0x20120211, 0x21100010, 0x21100111, 0x21110010, 0x21110011, 0x21110110, 0x21110111, 0x21110112, - 0x21110211, 0x21120012, 0x21120111, 0x22100110, 0x22100112, 0x22110012, 0x22110111, 0x22110210, 0x22120011, - 0x22120110, 0x22120112, 0x22120211, 0x20100121, 0x20110021, 0x20110120, 0x20110221, 0x20120121, 0x21100120, - 0x21100122, 0x21100221, 0x21110020, 0x21110022, 0x21110121, 0x21110220, 0x21120122, 0x21120221, 0x22100121, - 0x22110120, 0x22110122, 0x22120221, 0x20101001, 0x20101100, 0x20101102, 0x20111000, 0x20111101, 0x20111200, - 0x20121102, 0x21101000, 0x21101202, 0x21111001, 0x21111100, 0x21111101, 0x21111102, 0x21111200, 0x21111201, - 0x21121000, 0x21121001, 0x21121002, 0x21121101, 0x22101100, 0x22101102, 0x22111002, 0x22111100, 0x22111101, - 0x22111200, 0x22121001, 0x22121201, 0x20101010, 0x20101111, 0x20101210, 0x20101212, 0x20111010, 0x20111011, - 0x20111110, 0x20111111, 0x20111112, 0x20111211, 0x20121011, 0x20121111, 0x20121211, 0x20121212, 0x21101011, - 0x21101110, 0x21101111, 0x21101112, 0x21101211, 0x21111010, 0x21111011, 0x21111012, 0x21111110, 0x21111111, - 0x21111112, 0x21111210, 0x21111211, 0x21111212, 0x21121011, 0x21121110, 0x21121111, 0x21121112, 0x21121211, - 0x22101011, 0x22101111, 0x22101210, 0x22111011, 0x22111012, 0x22111110, 0x22111111, 0x22111112, 0x22111211, - 0x22111212, 0x22121010, 0x22121012, 0x22121111, 0x22121210, 0x22121212, 0x20101021, 0x20101120, 0x20111020, - 0x20111121, 0x20111221, 0x20121020, 0x20121122, 0x20121221, 0x21101121, 0x21101220, 0x21101221, 0x21111021, - 0x21111022, 0x21111121, 0x21111122, 0x21111221, 0x21121121, 0x21121220, 0x22101022, 0x22101120, 0x22101221, - 0x22101222, 0x22111022, 0x22111120, 0x22111121, 0x22121120, 0x22121122, 0x22121221, 0x20102101, 0x20112102, - 0x20112201, 0x20122101, 0x21102001, 0x21102102, 0x21112000, 0x21112002, 0x21112101, 0x21112102, 0x21112202, - 0x21122100, 0x21122101, 0x22102101, 0x22112001, 0x22112102, 0x22112201, 0x22122101, 0x20102110, 0x20102112, - 0x20102211, 0x20112010, 0x20112012, 0x20112111, 0x20112210, 0x20112212, 0x20122010, 0x20122011, 0x20122110, - 0x20122112, 0x21102010, 0x21102012, 0x21102111, 0x21102210, 0x21102212, 0x21112011, 0x21112110, 0x21112111, - 0x21112112, 0x21112211, 0x21122012, 0x21122111, 0x21122112, 0x21122212, 0x22102011, 0x22102110, 0x22112010, - 0x22112012, 0x22112111, 0x22112212, 0x22122011, 0x22122112, 0x20102121, 0x20112121, 0x20122121, 0x21102120, - 0x21102122, 0x21102221, 0x21112020, 0x21112121, 0x21112220, 0x21122021, 0x22102121, 0x22112021, 0x22112120, - 0x22112121, 0x22112122, 0x20200000, 0x20200002, 0x20200200, 0x20200202, 0x20210101, 0x20220000, 0x20220002, - 0x20220200, 0x20220202, 0x21200101, 0x21210001, 0x21210100, 0x21210102, 0x21210201, 0x22200000, 0x22200002, - 0x22200200, 0x22200202, 0x22210101, 0x22220000, 0x22220002, 0x22220200, 0x22220202, 0x20200111, 0x20200211, - 0x20210011, 0x20210110, 0x20210112, 0x20210211, 0x20210212, 0x21200112, 0x21200211, 0x21210011, 0x21210111, - 0x21210210, 0x21210212, 0x21220011, 0x21220110, 0x22200111, 0x22210010, 0x22210012, 0x22210112, 0x22210211, - 0x20200022, 0x20200220, 0x20200222, 0x20210020, 0x20210221, 0x20220022, 0x20220220, 0x20220222, 0x21200121, - 0x21210021, 0x21210122, 0x21210221, 0x21220121, 0x22200020, 0x22200022, 0x22200220, 0x22200222, 0x22210121, - 0x22220020, 0x22220022, 0x22220220, 0x22220222, 0x20211201, 0x20221101, 0x21201001, 0x21201100, 0x21211000, - 0x21211100, 0x21211101, 0x21211200, 0x21211202, 0x21221001, 0x21221101, 0x21221102, 0x21221200, 0x21221201, - 0x22201101, 0x20201112, 0x20201211, 0x20211010, 0x20211012, 0x20211111, 0x20211210, 0x20221112, 0x20221211, - 0x21201012, 0x21201111, 0x21211011, 0x21211110, 0x21211111, 0x21211112, 0x21211211, 0x21221111, 0x21221212, - 0x22201011, 0x22201110, 0x22201111, 0x22201112, 0x22201211, 0x22211012, 0x22211111, 0x22211210, 0x20201121, - 0x20211021, 0x20211122, 0x20211222, 0x20221021, 0x20221121, 0x21201120, 0x21201122, 0x21201222, 0x21211022, - 0x21211121, 0x21211122, 0x21211220, 0x21221020, 0x21221022, 0x22201122, 0x22211020, 0x22211121, 0x22211122, - 0x22211221, 0x22221021, 0x22221120, 0x22221122, 0x20202000, 0x20202002, 0x20202200, 0x20202202, 0x20222000, - 0x20222002, 0x20222200, 0x20222202, 0x21212001, 0x21212100, 0x21212102, 0x21212201, 0x22202000, 0x22202002, - 0x22202200, 0x22202202, 0x22212101, 0x22222000, 0x22222002, 0x22222200, 0x22222202, 0x20202111, 0x20212110, - 0x20212211, 0x20222011, 0x20222111, 0x21202011, 0x21212010, 0x21212111, 0x21212212, 0x21222011, 0x21222112, - 0x21222211, 0x22212010, 0x22212112, 0x20202020, 0x20202022, 0x20202220, 0x20202222, 0x20222020, 0x20222022, - 0x20222220, 0x20222222, 0x21212021, 0x21212120, 0x21212122, 0x22202020, 0x22202022, 0x22202220, 0x22202222, - 0x22212121, 0x22222020, 0x22222022, 0x22222220, 0x22222222, -}; - -static const __device__ uint8_t ksigns_iq2xs[128] = { - 0, 129, 130, 3, 132, 5, 6, 135, 136, 9, 10, 139, 12, 141, 142, 15, 144, 17, 18, 147, 20, 149, - 150, 23, 24, 153, 154, 27, 156, 29, 30, 159, 160, 33, 34, 163, 36, 165, 166, 39, 40, 169, 170, 43, - 172, 45, 46, 175, 48, 177, 178, 51, 180, 53, 54, 183, 184, 57, 58, 187, 60, 189, 190, 63, 192, 65, - 66, 195, 68, 197, 198, 71, 72, 201, 202, 75, 204, 77, 78, 207, 80, 209, 210, 83, 212, 85, 86, 215, - 216, 89, 90, 219, 92, 221, 222, 95, 96, 225, 226, 99, 228, 101, 102, 231, 232, 105, 106, 235, 108, 237, - 238, 111, 240, 113, 114, 243, 116, 245, 246, 119, 120, 249, 250, 123, 252, 125, 126, 255, -}; - -static const __device__ uint64_t ksigns64[128] = { - 0x0000000000000000, 0xff000000000000ff, 0xff0000000000ff00, 0x000000000000ffff, 0xff00000000ff0000, - 0x0000000000ff00ff, 0x0000000000ffff00, 0xff00000000ffffff, 0xff000000ff000000, 0x00000000ff0000ff, - 0x00000000ff00ff00, 0xff000000ff00ffff, 0x00000000ffff0000, 0xff000000ffff00ff, 0xff000000ffffff00, - 0x00000000ffffffff, 0xff0000ff00000000, 0x000000ff000000ff, 0x000000ff0000ff00, 0xff0000ff0000ffff, - 0x000000ff00ff0000, 0xff0000ff00ff00ff, 0xff0000ff00ffff00, 0x000000ff00ffffff, 0x000000ffff000000, - 0xff0000ffff0000ff, 0xff0000ffff00ff00, 0x000000ffff00ffff, 0xff0000ffffff0000, 0x000000ffffff00ff, - 0x000000ffffffff00, 0xff0000ffffffffff, 0xff00ff0000000000, 0x0000ff00000000ff, 0x0000ff000000ff00, - 0xff00ff000000ffff, 0x0000ff0000ff0000, 0xff00ff0000ff00ff, 0xff00ff0000ffff00, 0x0000ff0000ffffff, - 0x0000ff00ff000000, 0xff00ff00ff0000ff, 0xff00ff00ff00ff00, 0x0000ff00ff00ffff, 0xff00ff00ffff0000, - 0x0000ff00ffff00ff, 0x0000ff00ffffff00, 0xff00ff00ffffffff, 0x0000ffff00000000, 0xff00ffff000000ff, - 0xff00ffff0000ff00, 0x0000ffff0000ffff, 0xff00ffff00ff0000, 0x0000ffff00ff00ff, 0x0000ffff00ffff00, - 0xff00ffff00ffffff, 0xff00ffffff000000, 0x0000ffffff0000ff, 0x0000ffffff00ff00, 0xff00ffffff00ffff, - 0x0000ffffffff0000, 0xff00ffffffff00ff, 0xff00ffffffffff00, 0x0000ffffffffffff, 0xffff000000000000, - 0x00ff0000000000ff, 0x00ff00000000ff00, 0xffff00000000ffff, 0x00ff000000ff0000, 0xffff000000ff00ff, - 0xffff000000ffff00, 0x00ff000000ffffff, 0x00ff0000ff000000, 0xffff0000ff0000ff, 0xffff0000ff00ff00, - 0x00ff0000ff00ffff, 0xffff0000ffff0000, 0x00ff0000ffff00ff, 0x00ff0000ffffff00, 0xffff0000ffffffff, - 0x00ff00ff00000000, 0xffff00ff000000ff, 0xffff00ff0000ff00, 0x00ff00ff0000ffff, 0xffff00ff00ff0000, - 0x00ff00ff00ff00ff, 0x00ff00ff00ffff00, 0xffff00ff00ffffff, 0xffff00ffff000000, 0x00ff00ffff0000ff, - 0x00ff00ffff00ff00, 0xffff00ffff00ffff, 0x00ff00ffffff0000, 0xffff00ffffff00ff, 0xffff00ffffffff00, - 0x00ff00ffffffffff, 0x00ffff0000000000, 0xffffff00000000ff, 0xffffff000000ff00, 0x00ffff000000ffff, - 0xffffff0000ff0000, 0x00ffff0000ff00ff, 0x00ffff0000ffff00, 0xffffff0000ffffff, 0xffffff00ff000000, - 0x00ffff00ff0000ff, 0x00ffff00ff00ff00, 0xffffff00ff00ffff, 0x00ffff00ffff0000, 0xffffff00ffff00ff, - 0xffffff00ffffff00, 0x00ffff00ffffffff, 0xffffffff00000000, 0x00ffffff000000ff, 0x00ffffff0000ff00, - 0xffffffff0000ffff, 0x00ffffff00ff0000, 0xffffffff00ff00ff, 0xffffffff00ffff00, 0x00ffffff00ffffff, - 0x00ffffffff000000, 0xffffffffff0000ff, 0xffffffffff00ff00, 0x00ffffffff00ffff, 0xffffffffffff0000, - 0x00ffffffffff00ff, 0x00ffffffffffff00, 0xffffffffffffffff, -}; - -static const __device__ uint8_t kmask_iq2xs[8] = {1, 2, 4, 8, 16, 32, 64, 128}; -static const __device__ int8_t kvalues_iq4nl[16] = { - -127, -104, -83, -65, -49, -35, -22, -10, 1, 13, 25, 38, 53, 69, 89, 113}; - -typedef half dfloat; // dequantize float -typedef half2 dfloat2; -typedef void (*dequantize_kernel_t)(const void* vx, const int ib, const int iqs, dfloat2& v); -template -using to_cuda_ggml_t = void (*)(const void* __restrict__ x, dst_t* __restrict__ y, int k, cudaStream_t stream); -typedef float (*vec_dot_q_cuda_t)(const void* __restrict__ vbq, const block_q8_1* __restrict__ bq8_1, const int& iqs); -typedef void (*allocate_tiles_cuda_t)(int** x_ql, half2** x_dm, int** x_qh, int** x_sc); -typedef void (*load_tiles_cuda_t)( - const void* __restrict__ vx, - int* __restrict__ x_ql, - half2* __restrict__ x_dm, - int* __restrict__ x_qh, - int* __restrict__ x_sc, - const int& i_offset, - const int& i_max, - const int& k, - const int& blocks_per_row); -typedef float (*vec_dot_q_mul_mat_cuda_t)( - const int* __restrict__ x_ql, - const half2* __restrict__ x_dm, - const int* __restrict__ x_qh, - const int* __restrict__ x_sc, - const int* __restrict__ y_qs, - const half2* __restrict__ y_ms, - const int& i, - const int& j, - const int& k); - -// Utility function - -template -static __device__ __forceinline__ dst_t convert_from_half(half val) { - return val; -} - -template <> -__device__ __forceinline__ c10::BFloat16 convert_from_half(half val) { -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 || defined(USE_MUSA) - return __float2bfloat16(__half2float(val)); -#else - return __half2float(val); -#endif // defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 || defined(USE_MUSA) -} - -template <> -__device__ __forceinline__ float convert_from_half(half val) { - return __half2float(val); -} - -#if defined(USE_ROCM) - -#ifndef __has_builtin -#define __has_builtin(x) 0 -#endif - -typedef int8_t int8x4_t __attribute__((ext_vector_type(4))); -static __device__ __forceinline__ int __vsubss4(const int a, const int b) { - const int8x4_t va = reinterpret_cast(a); - const int8x4_t vb = reinterpret_cast(b); -#if __has_builtin(__builtin_elementwise_sub_sat) - const int8x4_t c = __builtin_elementwise_sub_sat(va, vb); - return reinterpret_cast(c); -#else - int8x4_t c; - int16_t tmp; -#pragma unroll - for (int i = 0; i < 4; i++) { - tmp = va[i] - vb[i]; - if (tmp > std::numeric_limits::max()) tmp = std::numeric_limits::max(); - if (tmp < std::numeric_limits::min()) tmp = std::numeric_limits::min(); - c[i] = tmp; - } - return reinterpret_cast(c); -#endif // __has_builtin(__builtin_elementwise_sub_sat) -} - -static __device__ __forceinline__ int __dp4a(const int a, const int b, int c) { -#if __has_builtin(__builtin_amdgcn_sdot4) - c = __builtin_amdgcn_sdot4(a, b, c, false); -#else - const int8x4_t va = reinterpret_cast(a); - const int8x4_t vb = reinterpret_cast(b); - c += va[0] * vb[0] + va[1] * vb[1] + va[2] * vb[2] + va[3] * vb[3]; -#endif - return c; -} - -static __device__ __forceinline__ uint32_t __vcmpeq4(const uint32_t a, const uint32_t b) { - uint32_t neq = a ^ b; - return !(neq & 0xff000000) * 0xff000000 | !(neq & 0x00ff0000) * 0x00ff0000 | !(neq & 0x0000ff00) * 0x0000ff00 | - !(neq & 0x000000ff) * 0x000000ff; -} - -static __device__ __forceinline__ uint32_t __vsub4(const uint32_t a, const uint32_t b) { - return (static_cast(((a & 0xff000000) >> 24) - ((b & 0xff000000) >> 24)) << 24) + - (static_cast(((a & 0x00ff0000) >> 16) - ((b & 0x00ff0000) >> 16)) << 16) + - (static_cast(((a & 0x0000ff00) >> 8) - ((b & 0x0000ff00) >> 8)) << 8) + - (static_cast(((a & 0x000000ff) >> 0) - ((b & 0x000000ff) >> 0)) << 0); -} -#endif // defined(USE_ROCM) diff --git a/python/freetoken/kernel/csrc/gguf/gguf_kernel.cu b/python/freetoken/kernel/csrc/gguf/gguf_kernel.cu deleted file mode 100644 index d88960d5fd..0000000000 --- a/python/freetoken/kernel/csrc/gguf/gguf_kernel.cu +++ /dev/null @@ -1,848 +0,0 @@ -// Adatped from -// https://github.com/vllm-project/vllm/blob/755ed7b05be4743237d3339c4ff8c22bcaae04f4/csrc/quantization/gguf/gguf_kernel.cu -#include -#include -#include -#include - -// dont use clang-format here, it breaks the include order -// clang-format off -#include "dispatch.h" - -#include "ggml-common.h" -#include "vecdotq.cuh" -#include "dequantize.cuh" -#include "mmvq.cuh" -#include "mmq.cuh" -#include "moe.cuh" -#include "moe_vec.cuh" -// clang-format off - -// Q8 gemv -template -static __global__ void -quantize_q8_1(const scalar_t* __restrict__ x, void* __restrict__ vy, const int kx, const int kx_padded) { - const auto ix = blockDim.x * blockIdx.x + threadIdx.x; - if (ix >= kx_padded) { - return; - } - const auto iy = blockDim.y * blockIdx.y + threadIdx.y; - const int i_padded = iy * kx_padded + ix; - - block_q8_1* y = (block_q8_1*)vy; - - const int ib = i_padded / QK8_1; // block index - const int iqs = i_padded % QK8_1; // quant index - - const float xi = ix < kx ? static_cast(x[iy * kx + ix]) : 0.0f; - float amax = fabsf(xi); - float sum = xi; - -#pragma unroll - for (int mask = 16; mask > 0; mask >>= 1) { - amax = fmaxf(amax, SGLANG_SHFL_XOR_SYNC_WIDTH(uint32_t(-1), amax, mask, 32)); - sum += SGLANG_SHFL_XOR_SYNC_WIDTH(uint32_t(-1), sum, mask, 32); - } - - const float d = amax / 127; - const int8_t q = amax == 0.0f ? 0 : roundf(xi / d); - - y[ib].qs[iqs] = q; - - if (iqs > 0) { - return; - } - - y[ib].ds.x = __float2half(d); - y[ib].ds.y = __float2half(sum); -} - -template -static void quantize_row_q8_1_cuda(const scalar_t* x, void* vy, const int kx, const int ky, cudaStream_t stream) { - const int64_t kx_padded = (kx + 512 - 1) / 512 * 512; - const int block_num_x = (kx_padded + CUDA_QUANTIZE_BLOCK_SIZE - 1) / CUDA_QUANTIZE_BLOCK_SIZE; - constexpr int MAX_BLOCK_SIZE = 65535; - for (int off = 0; off < ky; off += MAX_BLOCK_SIZE) { - const int num_blocks_y = std::min(ky, off + MAX_BLOCK_SIZE) - off; - const dim3 num_blocks(block_num_x, num_blocks_y, 1); - const dim3 block_size(CUDA_DEQUANTIZE_BLOCK_SIZE, 1, 1); - quantize_q8_1<<>>( - &x[off * kx], (int32_t*)vy + off * (kx_padded / 32 * 9), kx, kx_padded); - } -} - -torch::Tensor ggml_dequantize( - torch::Tensor W, // quant weight - int64_t type, - int64_t m, - int64_t n, - std::optional const& dtype) { - const at::cuda::OptionalCUDAGuard device_guard(device_of(W)); - auto dtype_ = dtype.value_or(torch::kFloat16); - auto options = torch::TensorOptions().dtype(dtype_).device(W.device()); - at::Tensor DW = torch::empty({m, n}, options); - cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream(); - - DISPATCH_FLOAT_TYPES(DW.scalar_type(), "ggml_dequantize", [&] { - auto to_cuda = ggml_get_to_cuda(type); - to_cuda((void*)W.data_ptr(), (scalar_t*)DW.data_ptr(), m * n, stream); - }); - - return DW; -} - -torch::Tensor ggml_mul_mat_vec_a8( - torch::Tensor W, // quant weight - torch::Tensor X, // input - int64_t type, - int64_t row) { - int col = X.sizes()[1]; - int vecs = X.sizes()[0]; - const int padded = (col + 512 - 1) / 512 * 512; - const at::cuda::OptionalCUDAGuard device_guard(device_of(X)); - auto options = torch::TensorOptions().dtype(X.dtype()).device(W.device()); - at::Tensor Y = torch::empty({vecs, row}, options); - cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream(); - options = torch::TensorOptions().dtype(torch::kInt32).device(W.device()); - at::Tensor quant_X = torch::empty({vecs, padded / 32 * 9}, options); - DISPATCH_FLOAT_TYPES(X.scalar_type(), "ggml_mul_mat_vec_a8", [&] { - quantize_row_q8_1_cuda((scalar_t*)X.data_ptr(), (void*)quant_X.data_ptr(), col, vecs, stream); - switch (type) { - case 2: - mul_mat_vec_q4_0_q8_1_cuda( - (void*)W.data_ptr(), (void*)quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), col, row, vecs, stream); - break; - case 3: - mul_mat_vec_q4_1_q8_1_cuda( - (void*)W.data_ptr(), (void*)quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), col, row, vecs, stream); - break; - case 6: - mul_mat_vec_q5_0_q8_1_cuda( - (void*)W.data_ptr(), (void*)quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), col, row, vecs, stream); - break; - case 7: - mul_mat_vec_q5_1_q8_1_cuda( - (void*)W.data_ptr(), (void*)quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), col, row, vecs, stream); - break; - case 8: - mul_mat_vec_q8_0_q8_1_cuda( - (void*)W.data_ptr(), (void*)quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), col, row, vecs, stream); - break; - case 10: - mul_mat_vec_q2_K_q8_1_cuda( - (void*)W.data_ptr(), (void*)quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), col, row, vecs, stream); - break; - case 11: - mul_mat_vec_q3_K_q8_1_cuda( - (void*)W.data_ptr(), (void*)quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), col, row, vecs, stream); - break; - case 12: - mul_mat_vec_q4_K_q8_1_cuda( - (void*)W.data_ptr(), (void*)quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), col, row, vecs, stream); - break; - case 13: - mul_mat_vec_q5_K_q8_1_cuda( - (void*)W.data_ptr(), (void*)quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), col, row, vecs, stream); - break; - case 14: - mul_mat_vec_q6_K_q8_1_cuda( - (void*)W.data_ptr(), (void*)quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), col, row, vecs, stream); - break; - case 16: - mul_mat_vec_iq2_xxs_q8_1_cuda( - (void*)W.data_ptr(), (void*)quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), col, row, vecs, stream); - break; - case 17: - mul_mat_vec_iq2_xs_q8_1_cuda( - (void*)W.data_ptr(), (void*)quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), col, row, vecs, stream); - break; - case 18: - mul_mat_vec_iq3_xxs_q8_1_cuda( - (void*)W.data_ptr(), (void*)quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), col, row, vecs, stream); - break; - case 19: - mul_mat_vec_iq1_s_q8_1_cuda( - (void*)W.data_ptr(), (void*)quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), col, row, vecs, stream); - break; - case 20: - mul_mat_vec_iq4_nl_q8_1_cuda( - (void*)W.data_ptr(), (void*)quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), col, row, vecs, stream); - break; - case 21: - mul_mat_vec_iq3_s_q8_1_cuda( - (void*)W.data_ptr(), (void*)quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), col, row, vecs, stream); - break; - case 22: - mul_mat_vec_iq2_s_q8_1_cuda( - (void*)W.data_ptr(), (void*)quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), col, row, vecs, stream); - break; - case 23: - mul_mat_vec_iq4_xs_q8_1_cuda( - (void*)W.data_ptr(), (void*)quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), col, row, vecs, stream); - break; - case 29: - mul_mat_vec_iq1_m_q8_1_cuda( - (void*)W.data_ptr(), (void*)quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), col, row, vecs, stream); - break; - } - }); - return Y; -} - -torch::Tensor ggml_mul_mat_a8( - torch::Tensor W, // quant weight - torch::Tensor X, // input - int64_t type, - int64_t row) { - int col = X.sizes()[1]; - int padded = (col + 512 - 1) / 512 * 512; - int batch = X.sizes()[0]; - const at::cuda::OptionalCUDAGuard device_guard(device_of(X)); - auto options = torch::TensorOptions().dtype(X.dtype()).device(W.device()); - at::Tensor Y = torch::empty({batch, row}, options); - cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream(); - options = torch::TensorOptions().dtype(torch::kInt32).device(W.device()); - at::Tensor quant_X = torch::empty({batch, padded / 32 * 9}, options); - DISPATCH_FLOAT_TYPES(X.scalar_type(), "ggml_mul_mat_a8", [&] { - quantize_row_q8_1_cuda((scalar_t*)X.data_ptr(), (void*)quant_X.data_ptr(), col, batch, stream); - - switch (type) { - case 2: - ggml_mul_mat_q4_0_q8_1_cuda( - (void*)W.data_ptr(), - (void*)quant_X.data_ptr(), - (scalar_t*)Y.data_ptr(), - col, - row, - batch, - padded, - row, - stream); - break; - case 3: - ggml_mul_mat_q4_1_q8_1_cuda( - (void*)W.data_ptr(), - (void*)quant_X.data_ptr(), - (scalar_t*)Y.data_ptr(), - col, - row, - batch, - padded, - row, - stream); - break; - case 6: - ggml_mul_mat_q5_0_q8_1_cuda( - (void*)W.data_ptr(), - (void*)quant_X.data_ptr(), - (scalar_t*)Y.data_ptr(), - col, - row, - batch, - padded, - row, - stream); - break; - case 7: - ggml_mul_mat_q5_1_q8_1_cuda( - (void*)W.data_ptr(), - (void*)quant_X.data_ptr(), - (scalar_t*)Y.data_ptr(), - col, - row, - batch, - padded, - row, - stream); - break; - case 8: - ggml_mul_mat_q8_0_q8_1_cuda( - (void*)W.data_ptr(), - (void*)quant_X.data_ptr(), - (scalar_t*)Y.data_ptr(), - col, - row, - batch, - padded, - row, - stream); - break; - case 10: - ggml_mul_mat_q2_K_q8_1_cuda( - (void*)W.data_ptr(), - (void*)quant_X.data_ptr(), - (scalar_t*)Y.data_ptr(), - col, - row, - batch, - padded, - row, - stream); - break; - case 11: - ggml_mul_mat_q3_K_q8_1_cuda( - (void*)W.data_ptr(), - (void*)quant_X.data_ptr(), - (scalar_t*)Y.data_ptr(), - col, - row, - batch, - padded, - row, - stream); - break; - case 12: - ggml_mul_mat_q4_K_q8_1_cuda( - (void*)W.data_ptr(), - (void*)quant_X.data_ptr(), - (scalar_t*)Y.data_ptr(), - col, - row, - batch, - padded, - row, - stream); - break; - case 13: - ggml_mul_mat_q5_K_q8_1_cuda( - (void*)W.data_ptr(), - (void*)quant_X.data_ptr(), - (scalar_t*)Y.data_ptr(), - col, - row, - batch, - padded, - row, - stream); - break; - case 14: - ggml_mul_mat_q6_K_q8_1_cuda( - (void*)W.data_ptr(), - (void*)quant_X.data_ptr(), - (scalar_t*)Y.data_ptr(), - col, - row, - batch, - padded, - row, - stream); - break; - } - }); - return Y; -} - -torch::Tensor ggml_moe_a8( - torch::Tensor X, // input - torch::Tensor W, // expert weights - torch::Tensor sorted_token_ids, - torch::Tensor expert_ids, - torch::Tensor num_tokens_post_padded, - int64_t type, - int64_t row, - int64_t top_k, - int64_t tokens) { - int col = X.sizes()[1]; - int padded = (col + 512 - 1) / 512 * 512; - const at::cuda::OptionalCUDAGuard device_guard(device_of(X)); - auto options = torch::TensorOptions().dtype(X.dtype()).device(W.device()); - at::Tensor Y = torch::empty({tokens * top_k, row}, options); - cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream(); - options = torch::TensorOptions().dtype(torch::kInt32).device(W.device()); - at::Tensor quant_X = torch::empty({tokens, padded / 32 * 9}, options); - DISPATCH_FLOAT_TYPES(X.scalar_type(), "ggml_moe_a8", [&] { - quantize_row_q8_1_cuda((scalar_t*)X.data_ptr(), (void*)quant_X.data_ptr(), col, tokens, stream); - switch (type) { - case 2: - ggml_moe_q4_0_q8_1_cuda( - (void*)quant_X.data_ptr(), - (void*)W.data_ptr(), - (scalar_t*)Y.data_ptr(), - (int*)sorted_token_ids.data_ptr(), - (int*)expert_ids.data_ptr(), - (int*)num_tokens_post_padded.data_ptr(), - W.stride(0), - col, - row, - tokens, - padded, - row, - top_k, - sorted_token_ids.sizes()[0], - stream); - break; - case 3: - ggml_moe_q4_1_q8_1_cuda( - (void*)quant_X.data_ptr(), - (void*)W.data_ptr(), - (scalar_t*)Y.data_ptr(), - (int*)sorted_token_ids.data_ptr(), - (int*)expert_ids.data_ptr(), - (int*)num_tokens_post_padded.data_ptr(), - W.stride(0), - col, - row, - tokens, - padded, - row, - top_k, - sorted_token_ids.sizes()[0], - stream); - break; - case 6: - ggml_moe_q5_0_q8_1_cuda( - (void*)quant_X.data_ptr(), - (void*)W.data_ptr(), - (scalar_t*)Y.data_ptr(), - (int*)sorted_token_ids.data_ptr(), - (int*)expert_ids.data_ptr(), - (int*)num_tokens_post_padded.data_ptr(), - W.stride(0), - col, - row, - tokens, - padded, - row, - top_k, - sorted_token_ids.sizes()[0], - stream); - break; - case 7: - ggml_moe_q5_1_q8_1_cuda( - (void*)quant_X.data_ptr(), - (void*)W.data_ptr(), - (scalar_t*)Y.data_ptr(), - (int*)sorted_token_ids.data_ptr(), - (int*)expert_ids.data_ptr(), - (int*)num_tokens_post_padded.data_ptr(), - W.stride(0), - col, - row, - tokens, - padded, - row, - top_k, - sorted_token_ids.sizes()[0], - stream); - break; - case 8: - ggml_moe_q8_0_q8_1_cuda( - (void*)quant_X.data_ptr(), - (void*)W.data_ptr(), - (scalar_t*)Y.data_ptr(), - (int*)sorted_token_ids.data_ptr(), - (int*)expert_ids.data_ptr(), - (int*)num_tokens_post_padded.data_ptr(), - W.stride(0), - col, - row, - tokens, - padded, - row, - top_k, - sorted_token_ids.sizes()[0], - stream); - break; - case 10: - ggml_moe_q2_K_q8_1_cuda( - (void*)quant_X.data_ptr(), - (void*)W.data_ptr(), - (scalar_t*)Y.data_ptr(), - (int*)sorted_token_ids.data_ptr(), - (int*)expert_ids.data_ptr(), - (int*)num_tokens_post_padded.data_ptr(), - W.stride(0), - col, - row, - tokens, - padded, - row, - top_k, - sorted_token_ids.sizes()[0], - stream); - break; - case 11: - ggml_moe_q3_K_q8_1_cuda( - (void*)quant_X.data_ptr(), - (void*)W.data_ptr(), - (scalar_t*)Y.data_ptr(), - (int*)sorted_token_ids.data_ptr(), - (int*)expert_ids.data_ptr(), - (int*)num_tokens_post_padded.data_ptr(), - W.stride(0), - col, - row, - tokens, - padded, - row, - top_k, - sorted_token_ids.sizes()[0], - stream); - break; - case 12: - ggml_moe_q4_K_q8_1_cuda( - (void*)quant_X.data_ptr(), - (void*)W.data_ptr(), - (scalar_t*)Y.data_ptr(), - (int*)sorted_token_ids.data_ptr(), - (int*)expert_ids.data_ptr(), - (int*)num_tokens_post_padded.data_ptr(), - W.stride(0), - col, - row, - tokens, - padded, - row, - top_k, - sorted_token_ids.sizes()[0], - stream); - break; - case 13: - ggml_moe_q5_K_q8_1_cuda( - (void*)quant_X.data_ptr(), - (void*)W.data_ptr(), - (scalar_t*)Y.data_ptr(), - (int*)sorted_token_ids.data_ptr(), - (int*)expert_ids.data_ptr(), - (int*)num_tokens_post_padded.data_ptr(), - W.stride(0), - col, - row, - tokens, - padded, - row, - top_k, - sorted_token_ids.sizes()[0], - stream); - break; - case 14: - ggml_moe_q6_K_q8_1_cuda( - (void*)quant_X.data_ptr(), - (void*)W.data_ptr(), - (scalar_t*)Y.data_ptr(), - (int*)sorted_token_ids.data_ptr(), - (int*)expert_ids.data_ptr(), - (int*)num_tokens_post_padded.data_ptr(), - W.stride(0), - col, - row, - tokens, - padded, - row, - top_k, - sorted_token_ids.sizes()[0], - stream); - break; - } - }); - return Y; -} - -torch::Tensor ggml_moe_a8_vec( - torch::Tensor X, // input - torch::Tensor W, // expert weights - torch::Tensor topk_ids, - int64_t top_k, - int64_t type, - int64_t row, - int64_t tokens) { - int col = X.sizes()[1]; - const int padded = (col + 512 - 1) / 512 * 512; - const at::cuda::OptionalCUDAGuard device_guard(device_of(X)); - auto options = torch::TensorOptions().dtype(X.dtype()).device(W.device()); - at::Tensor Y = torch::zeros({tokens * top_k, row}, options); - cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream(); - options = torch::TensorOptions().dtype(torch::kInt32).device(W.device()); - at::Tensor quant_X = torch::empty({tokens, padded / 32 * 9}, options); - DISPATCH_FLOAT_TYPES(X.scalar_type(), "ggml_moe_vec_a8", [&] { - quantize_row_q8_1_cuda((scalar_t*)X.data_ptr(), (void*)quant_X.data_ptr(), col, tokens, stream); - switch (type) { - case 2: - moe_vec_q4_0_q8_1_cuda( - (void*)W.data_ptr(), - (void*)quant_X.data_ptr(), - (scalar_t*)Y.data_ptr(), - (int*)topk_ids.data_ptr(), - top_k, - tokens, - col, - row, - quant_X.stride(0), - stream); - break; - case 3: - moe_vec_q4_1_q8_1_cuda( - (void*)W.data_ptr(), - (void*)quant_X.data_ptr(), - (scalar_t*)Y.data_ptr(), - (int*)topk_ids.data_ptr(), - top_k, - tokens, - col, - row, - quant_X.stride(0), - stream); - break; - case 6: - moe_vec_q5_0_q8_1_cuda( - (void*)W.data_ptr(), - (void*)quant_X.data_ptr(), - (scalar_t*)Y.data_ptr(), - (int*)topk_ids.data_ptr(), - top_k, - tokens, - col, - row, - quant_X.stride(0), - stream); - break; - case 7: - moe_vec_q5_1_q8_1_cuda( - (void*)W.data_ptr(), - (void*)quant_X.data_ptr(), - (scalar_t*)Y.data_ptr(), - (int*)topk_ids.data_ptr(), - top_k, - tokens, - col, - row, - quant_X.stride(0), - stream); - break; - case 8: - moe_vec_q8_0_q8_1_cuda( - (void*)W.data_ptr(), - (void*)quant_X.data_ptr(), - (scalar_t*)Y.data_ptr(), - (int*)topk_ids.data_ptr(), - top_k, - tokens, - col, - row, - quant_X.stride(0), - stream); - break; - case 10: - moe_vec_q2_K_q8_1_cuda( - (void*)W.data_ptr(), - (void*)quant_X.data_ptr(), - (scalar_t*)Y.data_ptr(), - (int*)topk_ids.data_ptr(), - top_k, - tokens, - col, - row, - quant_X.stride(0), - stream); - break; - case 11: - moe_vec_q3_K_q8_1_cuda( - (void*)W.data_ptr(), - (void*)quant_X.data_ptr(), - (scalar_t*)Y.data_ptr(), - (int*)topk_ids.data_ptr(), - top_k, - tokens, - col, - row, - quant_X.stride(0), - stream); - break; - case 12: - moe_vec_q4_K_q8_1_cuda( - (void*)W.data_ptr(), - (void*)quant_X.data_ptr(), - (scalar_t*)Y.data_ptr(), - (int*)topk_ids.data_ptr(), - top_k, - tokens, - col, - row, - quant_X.stride(0), - stream); - break; - case 13: - moe_vec_q5_K_q8_1_cuda( - (void*)W.data_ptr(), - (void*)quant_X.data_ptr(), - (scalar_t*)Y.data_ptr(), - (int*)topk_ids.data_ptr(), - top_k, - tokens, - col, - row, - quant_X.stride(0), - stream); - break; - case 14: - moe_vec_q6_K_q8_1_cuda( - (void*)W.data_ptr(), - (void*)quant_X.data_ptr(), - (scalar_t*)Y.data_ptr(), - (int*)topk_ids.data_ptr(), - top_k, - tokens, - col, - row, - quant_X.stride(0), - stream); - break; - case 16: - moe_vec_iq2_xxs_q8_1_cuda( - (void*)W.data_ptr(), - (void*)quant_X.data_ptr(), - (scalar_t*)Y.data_ptr(), - (int*)topk_ids.data_ptr(), - top_k, - tokens, - col, - row, - quant_X.stride(0), - stream); - break; - case 17: - moe_vec_iq2_xs_q8_1_cuda( - (void*)W.data_ptr(), - (void*)quant_X.data_ptr(), - (scalar_t*)Y.data_ptr(), - (int*)topk_ids.data_ptr(), - top_k, - tokens, - col, - row, - quant_X.stride(0), - stream); - break; - case 18: - moe_vec_iq3_xxs_q8_1_cuda( - (void*)W.data_ptr(), - (void*)quant_X.data_ptr(), - (scalar_t*)Y.data_ptr(), - (int*)topk_ids.data_ptr(), - top_k, - tokens, - col, - row, - quant_X.stride(0), - stream); - break; - case 19: - moe_vec_iq1_s_q8_1_cuda( - (void*)W.data_ptr(), - (void*)quant_X.data_ptr(), - (scalar_t*)Y.data_ptr(), - (int*)topk_ids.data_ptr(), - top_k, - tokens, - col, - row, - quant_X.stride(0), - stream); - break; - case 20: - moe_vec_iq4_nl_q8_1_cuda( - (void*)W.data_ptr(), - (void*)quant_X.data_ptr(), - (scalar_t*)Y.data_ptr(), - (int*)topk_ids.data_ptr(), - top_k, - tokens, - col, - row, - quant_X.stride(0), - stream); - break; - case 21: - moe_vec_iq3_s_q8_1_cuda( - (void*)W.data_ptr(), - (void*)quant_X.data_ptr(), - (scalar_t*)Y.data_ptr(), - (int*)topk_ids.data_ptr(), - top_k, - tokens, - col, - row, - quant_X.stride(0), - stream); - break; - case 22: - moe_vec_iq2_s_q8_1_cuda( - (void*)W.data_ptr(), - (void*)quant_X.data_ptr(), - (scalar_t*)Y.data_ptr(), - (int*)topk_ids.data_ptr(), - top_k, - tokens, - col, - row, - quant_X.stride(0), - stream); - break; - case 23: - moe_vec_iq4_xs_q8_1_cuda( - (void*)W.data_ptr(), - (void*)quant_X.data_ptr(), - (scalar_t*)Y.data_ptr(), - (int*)topk_ids.data_ptr(), - top_k, - tokens, - col, - row, - quant_X.stride(0), - stream); - break; - case 29: - moe_vec_iq1_m_q8_1_cuda( - (void*)W.data_ptr(), - (void*)quant_X.data_ptr(), - (scalar_t*)Y.data_ptr(), - (int*)topk_ids.data_ptr(), - top_k, - tokens, - col, - row, - quant_X.stride(0), - stream); - break; - } - }); - return Y; -} - -int64_t ggml_moe_get_block_size(int64_t type) { - switch (type) { - case 2: - return MOE_X_Q4_0; - case 3: - return MOE_X_Q4_1; - case 6: - return MOE_X_Q5_0; - case 7: - return MOE_X_Q5_1; - case 8: - return MOE_X_Q8_0; - case 10: - return MOE_X_Q2_K; - case 11: - return MOE_X_Q3_K; - case 12: - return MOE_X_Q4_K; - case 13: - return MOE_X_Q5_K; - case 14: - return MOE_X_Q6_K; - } - return 0; -} - -// ---- FreeToken pybind bindings (donor registers these via TORCH_LIBRARY; we -// expose them through torch.utils.cpp_extension.load's pybind module instead) ---- -#include -PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { - m.def("ggml_dequantize", &ggml_dequantize, ""); - m.def("ggml_mul_mat_vec_a8", &ggml_mul_mat_vec_a8, ""); - m.def("ggml_mul_mat_a8", &ggml_mul_mat_a8, ""); - m.def("ggml_moe_a8", &ggml_moe_a8, ""); - m.def("ggml_moe_a8_vec", &ggml_moe_a8_vec, ""); - m.def("ggml_moe_get_block_size", &ggml_moe_get_block_size, ""); -} diff --git a/python/freetoken/kernel/csrc/gguf/mmq.cuh b/python/freetoken/kernel/csrc/gguf/mmq.cuh deleted file mode 100644 index 9838c14f01..0000000000 --- a/python/freetoken/kernel/csrc/gguf/mmq.cuh +++ /dev/null @@ -1,881 +0,0 @@ -// copied from -// https://github.com/vllm-project/vllm/blob/4492e3a55428e161ca8db381edc28263e5da4c8d/csrc/quantization/gguf/mmq.cuh -// copied from https://github.com/ggerganov/llama.cpp/blob/b2899/ggml-cuda/mmq.cu -template < - typename scalar_t, - int qk, - int qr, - int qi, - bool need_sum, - typename block_q_t, - int mmq_x, - int mmq_y, - int nwarps, - allocate_tiles_cuda_t allocate_tiles, - load_tiles_cuda_t load_tiles, - int vdr, - vec_dot_q_mul_mat_cuda_t vec_dot> -static __device__ __forceinline__ void mul_mat_q( - const void* __restrict__ vx, - const void* __restrict__ vy, - scalar_t* __restrict__ dst, - const int ncols_x, - const int nrows_x, - const int ncols_y, - const int nrows_y, - const int nrows_dst) { - const block_q_t* x = (const block_q_t*)vx; - const block_q8_1* y = (const block_q8_1*)vy; - - const int blocks_per_row_x = ncols_x / qk; - const int blocks_per_col_y = nrows_y / QK8_1; - const int blocks_per_warp = WARP_SIZE_GGUF / qi; - - const int& ncols_dst = ncols_y; - - const auto row_dst_0 = blockIdx.x * mmq_y; - const int& row_x_0 = row_dst_0; - - const auto col_dst_0 = blockIdx.y * mmq_x; - const int& col_y_0 = col_dst_0; - - int* tile_x_ql = nullptr; - half2* tile_x_dm = nullptr; - int* tile_x_qh = nullptr; - int* tile_x_sc = nullptr; - - allocate_tiles(&tile_x_ql, &tile_x_dm, &tile_x_qh, &tile_x_sc); - - __shared__ int tile_y_qs[mmq_x * WARP_SIZE_GGUF]; - __shared__ half2 tile_y_ds[mmq_x * WARP_SIZE_GGUF / QI8_1]; - - float sum[mmq_y / WARP_SIZE_GGUF][mmq_x / nwarps] = {{0.0f}}; - - for (int ib0 = 0; ib0 < blocks_per_row_x; ib0 += blocks_per_warp) { - load_tiles( - x + row_x_0 * blocks_per_row_x + ib0, - tile_x_ql, - tile_x_dm, - tile_x_qh, - tile_x_sc, - threadIdx.y, - nrows_x - row_x_0 - 1, - threadIdx.x, - blocks_per_row_x); - -#pragma unroll - for (int ir = 0; ir < qr && ib0 + ir * blocks_per_warp / qr < blocks_per_row_x; ++ir) { - const auto kqs = ir * WARP_SIZE_GGUF + threadIdx.x; - const int kbxd = kqs / QI8_1; - -#pragma unroll - for (int i = 0; i < mmq_x; i += nwarps) { - const int col_y_eff = min(col_y_0 + threadIdx.y + i, ncols_y - 1); // to prevent out-of-bounds memory accesses - const block_q8_1* by0 = &y[col_y_eff * blocks_per_col_y + ib0 * (qk / QK8_1) + kbxd]; - const int index_y = (threadIdx.y + i) * WARP_SIZE_GGUF + kqs % WARP_SIZE_GGUF; - tile_y_qs[index_y] = get_int_from_int8_aligned(by0->qs, threadIdx.x % QI8_1); - } - -#pragma unroll - for (int ids0 = 0; ids0 < mmq_x; ids0 += nwarps * QI8_1) { - const int ids = (ids0 + threadIdx.y * QI8_1 + threadIdx.x / (WARP_SIZE_GGUF / QI8_1)) % mmq_x; - const auto kby = threadIdx.x % (WARP_SIZE_GGUF / QI8_1); - const int col_y_eff = min(col_y_0 + ids, ncols_y - 1); - - // if the sum is not needed it's faster to transform the scale to f32 ahead of time - const half2* dsi_src = - &y[col_y_eff * blocks_per_col_y + ib0 * (qk / QK8_1) + ir * (WARP_SIZE_GGUF / QI8_1) + kby].ds; - half2* dsi_dst = &tile_y_ds[ids * (WARP_SIZE_GGUF / QI8_1) + kby]; - if (need_sum) { - *dsi_dst = *dsi_src; - } else { - float* dfi_dst = (float*)dsi_dst; - *dfi_dst = __low2float(*dsi_src); - } - } - - __syncthreads(); - - // #pragma unroll // unrolling this loop causes too much register pressure - for (int k = ir * WARP_SIZE_GGUF / qr; k < (ir + 1) * WARP_SIZE_GGUF / qr; k += vdr) { -#pragma unroll - for (int j = 0; j < mmq_x; j += nwarps) { -#pragma unroll - for (int i = 0; i < mmq_y; i += WARP_SIZE_GGUF) { - sum[i / WARP_SIZE_GGUF][j / nwarps] += vec_dot( - tile_x_ql, tile_x_dm, tile_x_qh, tile_x_sc, tile_y_qs, tile_y_ds, threadIdx.x + i, threadIdx.y + j, k); - } - } - } - __syncthreads(); - } - } - -#pragma unroll - for (int j = 0; j < mmq_x; j += nwarps) { - const auto col_dst = col_dst_0 + j + threadIdx.y; - if (col_dst >= ncols_dst) { - return; - } - -#pragma unroll - for (int i = 0; i < mmq_y; i += WARP_SIZE_GGUF) { - const auto row_dst = row_dst_0 + threadIdx.x + i; - if (row_dst >= nrows_dst) { - continue; - } - dst[col_dst * nrows_dst + row_dst] = sum[i / WARP_SIZE_GGUF][j / nwarps]; - } - } -} - -#if defined(USE_ROCM) -#define MMQ_X_Q4_0 64 -#define MMQ_Y_Q4_0 128 -#define NWARPS_Q4_0 8 -#else -#define MMQ_X_Q4_0 4 -#define MMQ_Y_Q4_0 32 -#define NWARPS_Q4_0 4 -#endif - -template -static __global__ void -#if defined(USE_ROCM) -__launch_bounds__(WARP_SIZE_GGUF* NWARPS_Q4_0, 2) -#endif - mul_mat_q4_0( - const void* __restrict__ vx, - const void* __restrict__ vy, - scalar_t* __restrict__ dst, - const int ncols_x, - const int nrows_x, - const int ncols_y, - const int nrows_y, - const int nrows_dst) { - const int mmq_x = MMQ_X_Q4_0; - const int mmq_y = MMQ_Y_Q4_0; - const int nwarps = NWARPS_Q4_0; - - mul_mat_q< - scalar_t, - QK4_0, - QR4_0, - QI4_0, - true, - block_q4_0, - mmq_x, - mmq_y, - nwarps, - allocate_tiles_q4_0, - load_tiles_q4_0, - VDR_Q4_0_Q8_1_MMQ, - vec_dot_q4_0_q8_1_mul_mat>(vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); -} - -template -static void ggml_mul_mat_q4_0_q8_1_cuda( - const void* vx, - const void* vy, - scalar_t* dst, - const int ncols_x, - const int nrows_x, - const int ncols_y, - const int nrows_y, - const int nrows_dst, - cudaStream_t stream) { - int mmq_x = MMQ_X_Q4_0; - int mmq_y = MMQ_Y_Q4_0; - int nwarps = NWARPS_Q4_0; - - const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y; - const int block_num_y = (ncols_y + mmq_x - 1) / mmq_x; - const dim3 block_nums(block_num_x, block_num_y, 1); - const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1); - - if (nrows_x % mmq_y == 0) { - const bool need_check = false; - mul_mat_q4_0 - <<>>(vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); - } else { - const bool need_check = true; - mul_mat_q4_0 - <<>>(vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); - } -} - -#if defined(USE_ROCM) -#define MMQ_X_Q4_1 64 -#define MMQ_Y_Q4_1 128 -#define NWARPS_Q4_1 8 -#else -#define MMQ_X_Q4_1 4 -#define MMQ_Y_Q4_1 32 -#define NWARPS_Q4_1 4 -#endif - -template -static __global__ void -#if defined(USE_ROCM) -__launch_bounds__(WARP_SIZE_GGUF* NWARPS_Q4_1, 2) -#endif - mul_mat_q4_1( - const void* __restrict__ vx, - const void* __restrict__ vy, - scalar_t* __restrict__ dst, - const int ncols_x, - const int nrows_x, - const int ncols_y, - const int nrows_y, - const int nrows_dst) { - const int mmq_x = MMQ_X_Q4_1; - const int mmq_y = MMQ_Y_Q4_1; - const int nwarps = NWARPS_Q4_1; - - mul_mat_q< - scalar_t, - QK4_1, - QR4_1, - QI4_1, - true, - block_q4_1, - mmq_x, - mmq_y, - nwarps, - allocate_tiles_q4_1, - load_tiles_q4_1, - VDR_Q4_1_Q8_1_MMQ, - vec_dot_q4_1_q8_1_mul_mat>(vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); -} - -template -static void ggml_mul_mat_q4_1_q8_1_cuda( - const void* vx, - const void* vy, - scalar_t* dst, - const int ncols_x, - const int nrows_x, - const int ncols_y, - const int nrows_y, - const int nrows_dst, - cudaStream_t stream) { - int mmq_x = MMQ_X_Q4_1; - int mmq_y = MMQ_Y_Q4_1; - int nwarps = NWARPS_Q4_1; - - const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y; - const int block_num_y = (ncols_y + mmq_x - 1) / mmq_x; - const dim3 block_nums(block_num_x, block_num_y, 1); - const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1); - - if (nrows_x % mmq_y == 0) { - const bool need_check = false; - mul_mat_q4_1 - <<>>(vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); - } else { - const bool need_check = true; - mul_mat_q4_1 - <<>>(vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); - } -} - -#if defined(USE_ROCM) -#define MMQ_X_Q5_0 64 -#define MMQ_Y_Q5_0 128 -#define NWARPS_Q5_0 8 -#else -#define MMQ_X_Q5_0 4 -#define MMQ_Y_Q5_0 32 -#define NWARPS_Q5_0 4 -#endif - -template -static __global__ void -#if defined(USE_ROCM) -__launch_bounds__(WARP_SIZE_GGUF* NWARPS_Q5_0, 2) -#endif - mul_mat_q5_0( - const void* __restrict__ vx, - const void* __restrict__ vy, - scalar_t* __restrict__ dst, - const int ncols_x, - const int nrows_x, - const int ncols_y, - const int nrows_y, - const int nrows_dst) { - const int mmq_x = MMQ_X_Q5_0; - const int mmq_y = MMQ_Y_Q5_0; - const int nwarps = NWARPS_Q5_0; - - mul_mat_q< - scalar_t, - QK5_0, - QR5_0, - QI5_0, - false, - block_q5_0, - mmq_x, - mmq_y, - nwarps, - allocate_tiles_q5_0, - load_tiles_q5_0, - VDR_Q5_0_Q8_1_MMQ, - vec_dot_q5_0_q8_1_mul_mat>(vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); -} - -template -static void ggml_mul_mat_q5_0_q8_1_cuda( - const void* vx, - const void* vy, - scalar_t* dst, - const int ncols_x, - const int nrows_x, - const int ncols_y, - const int nrows_y, - const int nrows_dst, - cudaStream_t stream) { - const int mmq_x = MMQ_X_Q5_0; - const int mmq_y = MMQ_Y_Q5_0; - const int nwarps = NWARPS_Q5_0; - - const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y; - const int block_num_y = (ncols_y + mmq_x - 1) / mmq_x; - const dim3 block_nums(block_num_x, block_num_y, 1); - const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1); - - if (nrows_x % mmq_y == 0) { - const bool need_check = false; - mul_mat_q5_0 - <<>>(vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); - } else { - const bool need_check = true; - mul_mat_q5_0 - <<>>(vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); - } -} - -#if defined(USE_ROCM) -#define MMQ_X_Q5_1 64 -#define MMQ_Y_Q5_1 128 -#define NWARPS_Q5_1 8 -#else -#define MMQ_X_Q5_1 4 -#define MMQ_Y_Q5_1 32 -#define NWARPS_Q5_1 4 -#endif - -template -static __global__ void -#if defined(USE_ROCM) -__launch_bounds__(WARP_SIZE_GGUF* NWARPS_Q5_1, 2) -#endif - mul_mat_q5_1( - const void* __restrict__ vx, - const void* __restrict__ vy, - scalar_t* __restrict__ dst, - const int ncols_x, - const int nrows_x, - const int ncols_y, - const int nrows_y, - const int nrows_dst) { - const int mmq_x = MMQ_X_Q5_1; - const int mmq_y = MMQ_Y_Q5_1; - const int nwarps = NWARPS_Q5_1; - - mul_mat_q< - scalar_t, - QK5_1, - QR5_1, - QI5_1, - true, - block_q5_1, - mmq_x, - mmq_y, - nwarps, - allocate_tiles_q5_1, - load_tiles_q5_1, - VDR_Q5_1_Q8_1_MMQ, - vec_dot_q5_1_q8_1_mul_mat>(vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); -} - -template -static void ggml_mul_mat_q5_1_q8_1_cuda( - const void* vx, - const void* vy, - scalar_t* dst, - const int ncols_x, - const int nrows_x, - const int ncols_y, - const int nrows_y, - const int nrows_dst, - cudaStream_t stream) { - const int mmq_x = MMQ_X_Q5_1; - const int mmq_y = MMQ_Y_Q5_1; - const int nwarps = NWARPS_Q5_1; - - const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y; - const int block_num_y = (ncols_y + mmq_x - 1) / mmq_x; - const dim3 block_nums(block_num_x, block_num_y, 1); - const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1); - - if (nrows_x % mmq_y == 0) { - const bool need_check = false; - mul_mat_q5_1 - <<>>(vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); - } else { - const bool need_check = true; - mul_mat_q5_1 - <<>>(vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); - } -} - -#if defined(USE_ROCM) -#define MMQ_X_Q8_0 64 -#define MMQ_Y_Q8_0 128 -#define NWARPS_Q8_0 8 -#else -#define MMQ_X_Q8_0 4 -#define MMQ_Y_Q8_0 32 -#define NWARPS_Q8_0 4 -#endif - -template -static __global__ void -#if defined(USE_ROCM) -__launch_bounds__(WARP_SIZE_GGUF* NWARPS_Q8_0, 2) -#endif - mul_mat_q8_0( - const void* __restrict__ vx, - const void* __restrict__ vy, - scalar_t* __restrict__ dst, - const int ncols_x, - const int nrows_x, - const int ncols_y, - const int nrows_y, - const int nrows_dst) { - const int mmq_x = MMQ_X_Q8_0; - const int mmq_y = MMQ_Y_Q8_0; - const int nwarps = NWARPS_Q8_0; - - mul_mat_q< - scalar_t, - QK8_0, - QR8_0, - QI8_0, - false, - block_q8_0, - mmq_x, - mmq_y, - nwarps, - allocate_tiles_q8_0, - load_tiles_q8_0, - VDR_Q8_0_Q8_1_MMQ, - vec_dot_q8_0_q8_1_mul_mat>(vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); -} - -template -static void ggml_mul_mat_q8_0_q8_1_cuda( - const void* vx, - const void* vy, - scalar_t* dst, - const int ncols_x, - const int nrows_x, - const int ncols_y, - const int nrows_y, - const int nrows_dst, - cudaStream_t stream) { - const int mmq_x = MMQ_X_Q8_0; - const int mmq_y = MMQ_Y_Q8_0; - const int nwarps = NWARPS_Q8_0; - - const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y; - const int block_num_y = (ncols_y + mmq_x - 1) / mmq_x; - const dim3 block_nums(block_num_x, block_num_y, 1); - const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1); - - if (nrows_x % mmq_y == 0) { - const bool need_check = false; - mul_mat_q8_0 - <<>>(vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); - } else { - const bool need_check = true; - mul_mat_q8_0 - <<>>(vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); - } -} - -#if defined(USE_ROCM) -#define MMQ_X_Q2_K 64 -#define MMQ_Y_Q2_K 128 -#define NWARPS_Q2_K 8 -#else -#define MMQ_X_Q2_K 4 -#define MMQ_Y_Q2_K 32 -#define NWARPS_Q2_K 4 -#endif - -template -static __global__ void -#if defined(USE_ROCM) -__launch_bounds__(WARP_SIZE_GGUF* NWARPS_Q2_K, 2) -#endif - mul_mat_q2_K( - const void* __restrict__ vx, - const void* __restrict__ vy, - scalar_t* __restrict__ dst, - const int ncols_x, - const int nrows_x, - const int ncols_y, - const int nrows_y, - const int nrows_dst) { - const int mmq_x = MMQ_X_Q2_K; - const int mmq_y = MMQ_Y_Q2_K; - const int nwarps = NWARPS_Q2_K; - - mul_mat_q< - scalar_t, - QK_K, - QR2_K, - QI2_K, - false, - block_q2_K, - mmq_x, - mmq_y, - nwarps, - allocate_tiles_q2_K, - load_tiles_q2_K, - VDR_Q2_K_Q8_1_MMQ, - vec_dot_q2_K_q8_1_mul_mat>(vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); -} - -template -static void ggml_mul_mat_q2_K_q8_1_cuda( - const void* vx, - const void* vy, - scalar_t* dst, - const int ncols_x, - const int nrows_x, - const int ncols_y, - const int nrows_y, - const int nrows_dst, - cudaStream_t stream) { - const int mmq_x = MMQ_X_Q2_K; - const int mmq_y = MMQ_Y_Q2_K; - const int nwarps = NWARPS_Q2_K; - - const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y; - const int block_num_y = (ncols_y + mmq_x - 1) / mmq_x; - const dim3 block_nums(block_num_x, block_num_y, 1); - const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1); - - if (nrows_x % mmq_y == 0) { - const bool need_check = false; - mul_mat_q2_K - <<>>(vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); - } else { - const bool need_check = true; - mul_mat_q2_K - <<>>(vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); - } -} - -#if defined(USE_ROCM) -#define MMQ_X_Q3_K 64 -#define MMQ_Y_Q3_K 128 -#define NWARPS_Q3_K 8 -#else -#define MMQ_X_Q3_K 4 -#define MMQ_Y_Q3_K 32 -#define NWARPS_Q3_K 4 -#endif - -template -static __global__ void -#if defined(USE_ROCM) -__launch_bounds__(WARP_SIZE_GGUF* NWARPS_Q3_K, 2) -#endif - mul_mat_q3_K( - const void* __restrict__ vx, - const void* __restrict__ vy, - scalar_t* __restrict__ dst, - const int ncols_x, - const int nrows_x, - const int ncols_y, - const int nrows_y, - const int nrows_dst) { - - const int mmq_x = MMQ_X_Q3_K; - const int mmq_y = MMQ_Y_Q3_K; - const int nwarps = NWARPS_Q3_K; - - mul_mat_q< - scalar_t, - QK_K, - QR3_K, - QI3_K, - false, - block_q3_K, - mmq_x, - mmq_y, - nwarps, - allocate_tiles_q3_K, - load_tiles_q3_K, - VDR_Q3_K_Q8_1_MMQ, - vec_dot_q3_K_q8_1_mul_mat>(vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); -} - -template -static void ggml_mul_mat_q3_K_q8_1_cuda( - const void* vx, - const void* vy, - scalar_t* dst, - const int ncols_x, - const int nrows_x, - const int ncols_y, - const int nrows_y, - const int nrows_dst, - cudaStream_t stream) { - const int mmq_x = MMQ_X_Q3_K; - const int mmq_y = MMQ_Y_Q3_K; - const int nwarps = NWARPS_Q3_K; - - const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y; - const int block_num_y = (ncols_y + mmq_x - 1) / mmq_x; - const dim3 block_nums(block_num_x, block_num_y, 1); - const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1); - - if (nrows_x % mmq_y == 0) { - const bool need_check = false; - mul_mat_q3_K - <<>>(vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); - } else { - const bool need_check = true; - mul_mat_q3_K - <<>>(vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); - } -} - -#if defined(USE_ROCM) -#define MMQ_X_Q4_K 64 -#define MMQ_Y_Q4_K 128 -#define NWARPS_Q4_K 8 -#else -#define MMQ_X_Q4_K 4 -#define MMQ_Y_Q4_K 32 -#define NWARPS_Q4_K 4 -#endif - -template -static __global__ void -#if defined(USE_ROCM) -__launch_bounds__(WARP_SIZE_GGUF* NWARPS_Q4_K, 2) -#endif - mul_mat_q4_K( - const void* __restrict__ vx, - const void* __restrict__ vy, - scalar_t* __restrict__ dst, - const int ncols_x, - const int nrows_x, - const int ncols_y, - const int nrows_y, - const int nrows_dst) { - const int mmq_x = MMQ_X_Q4_K; - const int mmq_y = MMQ_Y_Q4_K; - const int nwarps = NWARPS_Q4_K; - - mul_mat_q< - scalar_t, - QK_K, - QR4_K, - QI4_K, - true, - block_q4_K, - mmq_x, - mmq_y, - nwarps, - allocate_tiles_q4_K, - load_tiles_q4_K, - VDR_Q4_K_Q8_1_MMQ, - vec_dot_q4_K_q8_1_mul_mat>(vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); -} - -template -static void ggml_mul_mat_q4_K_q8_1_cuda( - const void* vx, - const void* vy, - scalar_t* dst, - const int ncols_x, - const int nrows_x, - const int ncols_y, - const int nrows_y, - const int nrows_dst, - cudaStream_t stream) { - const int mmq_x = MMQ_X_Q4_K; - const int mmq_y = MMQ_Y_Q4_K; - const int nwarps = NWARPS_Q4_K; - - const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y; - const int block_num_y = (ncols_y + mmq_x - 1) / mmq_x; - const dim3 block_nums(block_num_x, block_num_y, 1); - const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1); - - if (nrows_x % mmq_y == 0) { - const bool need_check = false; - mul_mat_q4_K - <<>>(vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); - } else { - const bool need_check = true; - mul_mat_q4_K - <<>>(vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); - } -} - -#if defined(USE_ROCM) -#define MMQ_X_Q5_K 64 -#define MMQ_Y_Q5_K 128 -#define NWARPS_Q5_K 8 -#else -#define MMQ_X_Q5_K 4 -#define MMQ_Y_Q5_K 32 -#define NWARPS_Q5_K 4 -#endif - -template -static __global__ void -#if defined(USE_ROCM) -__launch_bounds__(WARP_SIZE_GGUF* NWARPS_Q5_K, 2) -#endif - mul_mat_q5_K( - const void* __restrict__ vx, - const void* __restrict__ vy, - scalar_t* __restrict__ dst, - const int ncols_x, - const int nrows_x, - const int ncols_y, - const int nrows_y, - const int nrows_dst) { - const int mmq_x = MMQ_X_Q5_K; - const int mmq_y = MMQ_Y_Q5_K; - const int nwarps = NWARPS_Q5_K; - - mul_mat_q< - scalar_t, - QK_K, - QR5_K, - QI5_K, - true, - block_q5_K, - mmq_x, - mmq_y, - nwarps, - allocate_tiles_q5_K, - load_tiles_q5_K, - VDR_Q5_K_Q8_1_MMQ, - vec_dot_q5_K_q8_1_mul_mat>(vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); -} - -template -static void ggml_mul_mat_q5_K_q8_1_cuda( - const void* vx, - const void* vy, - scalar_t* dst, - const int ncols_x, - const int nrows_x, - const int ncols_y, - const int nrows_y, - const int nrows_dst, - cudaStream_t stream) { - const int mmq_x = MMQ_X_Q5_K; - const int mmq_y = MMQ_Y_Q5_K; - const int nwarps = NWARPS_Q5_K; - - const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y; - const int block_num_y = (ncols_y + mmq_x - 1) / mmq_x; - const dim3 block_nums(block_num_x, block_num_y, 1); - const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1); - - if (nrows_x % mmq_y == 0) { - const bool need_check = false; - mul_mat_q5_K - <<>>(vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); - } else { - const bool need_check = true; - mul_mat_q5_K - <<>>(vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); - } -} - -#if defined(USE_ROCM) -#define MMQ_X_Q6_K 64 -#define MMQ_Y_Q6_K 128 -#define NWARPS_Q6_K 8 -#else -#define MMQ_X_Q6_K 4 -#define MMQ_Y_Q6_K 32 -#define NWARPS_Q6_K 4 -#endif - -template -static __global__ void -#if defined(USE_ROCM) -__launch_bounds__(WARP_SIZE_GGUF* NWARPS_Q6_K, 2) -#endif - mul_mat_q6_K( - const void* __restrict__ vx, - const void* __restrict__ vy, - scalar_t* __restrict__ dst, - const int ncols_x, - const int nrows_x, - const int ncols_y, - const int nrows_y, - const int nrows_dst) { - const int mmq_x = MMQ_X_Q6_K; - const int mmq_y = MMQ_Y_Q6_K; - const int nwarps = NWARPS_Q6_K; - - mul_mat_q< - scalar_t, - QK_K, - QR6_K, - QI6_K, - false, - block_q6_K, - mmq_x, - mmq_y, - nwarps, - allocate_tiles_q6_K, - load_tiles_q6_K, - VDR_Q6_K_Q8_1_MMQ, - vec_dot_q6_K_q8_1_mul_mat>(vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); -} - -template -static void ggml_mul_mat_q6_K_q8_1_cuda( - const void* vx, - const void* vy, - scalar_t* dst, - const int ncols_x, - const int nrows_x, - const int ncols_y, - const int nrows_y, - const int nrows_dst, - cudaStream_t stream) { - const int mmq_x = MMQ_X_Q6_K; - const int mmq_y = MMQ_Y_Q6_K; - const int nwarps = NWARPS_Q6_K; - - const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y; - const int block_num_y = (ncols_y + mmq_x - 1) / mmq_x; - const dim3 block_nums(block_num_x, block_num_y, 1); - const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1); - - if (nrows_x % mmq_y == 0) { - const bool need_check = false; - mul_mat_q6_K - <<>>(vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); - } else { - const bool need_check = true; - mul_mat_q6_K - <<>>(vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); - } -} diff --git a/python/freetoken/kernel/csrc/gguf/mmvq.cuh b/python/freetoken/kernel/csrc/gguf/mmvq.cuh deleted file mode 100644 index 7331731ace..0000000000 --- a/python/freetoken/kernel/csrc/gguf/mmvq.cuh +++ /dev/null @@ -1,352 +0,0 @@ -// copied from -// https://github.com/vllm-project/vllm/blob/4492e3a55428e161ca8db381edc28263e5da4c8d/csrc/quantization/gguf/mmvq.cuh -// copied and adapted from https://github.com/ggerganov/llama.cpp/blob/b2899/ggml-cuda/mmvq.cu -template -static __global__ void mul_mat_vec_q( - const void* __restrict__ vx, - const void* __restrict__ vy, - scalar_t* __restrict__ dst, - const int ncols, - const int nrows, - const int nvecs) { - const auto row = blockIdx.x * blockDim.y + threadIdx.y; - const auto vec = blockIdx.y; - - if (row >= nrows || vec >= nvecs) { - return; - } - - const int blocks_per_row = ncols / qk; - const int blocks_per_warp = vdr * WARP_SIZE / qi; - const int nrows_y = (ncols + 512 - 1) / 512 * 512; - - // partial sum for each thread - float tmp = 0.0f; - - const block_q_t* x = (const block_q_t*)vx; - const block_q8_1* y = (const block_q8_1*)vy; - - for (auto i = threadIdx.x / (qi / vdr); i < blocks_per_row; i += blocks_per_warp) { - const int ibx = row * blocks_per_row + i; // x block index - - const int iby = vec * (nrows_y / QK8_1) + i * (qk / QK8_1); // y block index that aligns with ibx - - const int iqs = vdr * (threadIdx.x % (qi / vdr)); // x block quant index when casting the quants to int - - tmp += vec_dot_q_cuda(&x[ibx], &y[iby], iqs); - } - - // sum up partial sums and write back result -#pragma unroll - for (int mask = WARP_SIZE / 2; mask > 0; mask >>= 1) { - tmp += SGLANG_SHFL_XOR_SYNC(uint32_t(-1), tmp, mask); - } - - if (threadIdx.x == 0) { - dst[vec * nrows + row] = tmp; - } -} - -template -static void mul_mat_vec_q4_0_q8_1_cuda( - const void* vx, - const void* vy, - scalar_t* dst, - const int ncols, - const int nrows, - const int nvecs, - cudaStream_t stream) { - const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, nvecs, 1); - const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - mul_mat_vec_q - <<>>(vx, vy, dst, ncols, nrows, nvecs); -} - -template -static void mul_mat_vec_q4_1_q8_1_cuda( - const void* vx, - const void* vy, - scalar_t* dst, - const int ncols, - const int nrows, - const int nvecs, - cudaStream_t stream) { - const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, nvecs, 1); - const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - mul_mat_vec_q - <<>>(vx, vy, dst, ncols, nrows, nvecs); -} - -template -static void mul_mat_vec_q5_0_q8_1_cuda( - const void* vx, - const void* vy, - scalar_t* dst, - const int ncols, - const int nrows, - const int nvecs, - cudaStream_t stream) { - const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, nvecs, 1); - const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - mul_mat_vec_q - <<>>(vx, vy, dst, ncols, nrows, nvecs); -} - -template -static void mul_mat_vec_q5_1_q8_1_cuda( - const void* vx, - const void* vy, - scalar_t* dst, - const int ncols, - const int nrows, - const int nvecs, - cudaStream_t stream) { - const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, nvecs, 1); - const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - mul_mat_vec_q - <<>>(vx, vy, dst, ncols, nrows, nvecs); -} - -template -static void mul_mat_vec_q8_0_q8_1_cuda( - const void* vx, - const void* vy, - scalar_t* dst, - const int ncols, - const int nrows, - const int nvecs, - cudaStream_t stream) { - const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, nvecs, 1); - const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - mul_mat_vec_q - <<>>(vx, vy, dst, ncols, nrows, nvecs); -} - -template -static void mul_mat_vec_q2_K_q8_1_cuda( - const void* vx, - const void* vy, - scalar_t* dst, - const int ncols, - const int nrows, - const int nvecs, - cudaStream_t stream) { - const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, nvecs, 1); - const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - mul_mat_vec_q - <<>>(vx, vy, dst, ncols, nrows, nvecs); -} - -template -static void mul_mat_vec_q3_K_q8_1_cuda( - const void* vx, - const void* vy, - scalar_t* dst, - const int ncols, - const int nrows, - const int nvecs, - cudaStream_t stream) { - const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, nvecs, 1); - const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - mul_mat_vec_q - <<>>(vx, vy, dst, ncols, nrows, nvecs); -} - -template -static void mul_mat_vec_q4_K_q8_1_cuda( - const void* vx, - const void* vy, - scalar_t* dst, - const int ncols, - const int nrows, - const int nvecs, - cudaStream_t stream) { - const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, nvecs, 1); - const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - mul_mat_vec_q - <<>>(vx, vy, dst, ncols, nrows, nvecs); -} - -template -static void mul_mat_vec_q5_K_q8_1_cuda( - const void* vx, - const void* vy, - scalar_t* dst, - const int ncols, - const int nrows, - const int nvecs, - cudaStream_t stream) { - const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, nvecs, 1); - const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - mul_mat_vec_q - <<>>(vx, vy, dst, ncols, nrows, nvecs); -} - -template -static void mul_mat_vec_q6_K_q8_1_cuda( - const void* vx, - const void* vy, - scalar_t* dst, - const int ncols, - const int nrows, - const int nvecs, - cudaStream_t stream) { - const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, nvecs, 1); - const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - mul_mat_vec_q - <<>>(vx, vy, dst, ncols, nrows, nvecs); -} - -template -static void mul_mat_vec_iq2_xxs_q8_1_cuda( - const void* vx, - const void* vy, - scalar_t* dst, - const int ncols, - const int nrows, - const int nvecs, - cudaStream_t stream) { - const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, nvecs, 1); - const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - mul_mat_vec_q - <<>>(vx, vy, dst, ncols, nrows, nvecs); -} - -template -static void mul_mat_vec_iq2_xs_q8_1_cuda( - const void* vx, - const void* vy, - scalar_t* dst, - const int ncols, - const int nrows, - const int nvecs, - cudaStream_t stream) { - const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, nvecs, 1); - const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - mul_mat_vec_q - <<>>(vx, vy, dst, ncols, nrows, nvecs); -} - -template -static void mul_mat_vec_iq2_s_q8_1_cuda( - const void* vx, - const void* vy, - scalar_t* dst, - const int ncols, - const int nrows, - const int nvecs, - cudaStream_t stream) { - const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, nvecs, 1); - const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - mul_mat_vec_q - <<>>(vx, vy, dst, ncols, nrows, nvecs); -} - -template -static void mul_mat_vec_iq3_xxs_q8_1_cuda( - const void* vx, - const void* vy, - scalar_t* dst, - const int ncols, - const int nrows, - const int nvecs, - cudaStream_t stream) { - const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, nvecs, 1); - const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - mul_mat_vec_q - <<>>(vx, vy, dst, ncols, nrows, nvecs); -} - -template -static void mul_mat_vec_iq1_s_q8_1_cuda( - const void* vx, - const void* vy, - scalar_t* dst, - const int ncols, - const int nrows, - const int nvecs, - cudaStream_t stream) { - const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, nvecs, 1); - const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - mul_mat_vec_q - <<>>(vx, vy, dst, ncols, nrows, nvecs); -} - -template -static void mul_mat_vec_iq1_m_q8_1_cuda( - const void* vx, - const void* vy, - scalar_t* dst, - const int ncols, - const int nrows, - const int nvecs, - cudaStream_t stream) { - const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, nvecs, 1); - const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - mul_mat_vec_q - <<>>(vx, vy, dst, ncols, nrows, nvecs); -} - -template -static void mul_mat_vec_iq4_nl_q8_1_cuda( - const void* vx, - const void* vy, - scalar_t* dst, - const int ncols, - const int nrows, - const int nvecs, - cudaStream_t stream) { - const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, nvecs, 1); - const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - mul_mat_vec_q - <<>>(vx, vy, dst, ncols, nrows, nvecs); -} - -template -static void mul_mat_vec_iq4_xs_q8_1_cuda( - const void* vx, - const void* vy, - scalar_t* dst, - const int ncols, - const int nrows, - const int nvecs, - cudaStream_t stream) { - const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, nvecs, 1); - const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - mul_mat_vec_q - <<>>(vx, vy, dst, ncols, nrows, nvecs); -} - -template -static void mul_mat_vec_iq3_s_q8_1_cuda( - const void* vx, - const void* vy, - scalar_t* dst, - const int ncols, - const int nrows, - const int nvecs, - cudaStream_t stream) { - const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, nvecs, 1); - const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - mul_mat_vec_q - <<>>(vx, vy, dst, ncols, nrows, nvecs); -} diff --git a/python/freetoken/kernel/csrc/gguf/moe.cuh b/python/freetoken/kernel/csrc/gguf/moe.cuh deleted file mode 100644 index 91e434fa60..0000000000 --- a/python/freetoken/kernel/csrc/gguf/moe.cuh +++ /dev/null @@ -1,1379 +0,0 @@ -// copied from -// https://github.com/vllm-project/vllm/blob/4492e3a55428e161ca8db381edc28263e5da4c8d/csrc/quantization/gguf/moe.cuh -#include - -/* Adapted from ./csrc/quantization/gguf/mmq.cuh - */ -template < - typename scalar_t, - int qk, - int qr, - int qi, - bool need_sum, - typename block_q_t, - int mmq_x, - int mmq_y, - int nwarps, - allocate_tiles_cuda_t allocate_tiles, - load_tiles_cuda_t load_tiles, - int vdr, - vec_dot_q_mul_mat_cuda_t vec_dot> -static __device__ __forceinline__ void moe_q( - const void* __restrict__ vx, - const void* __restrict__ vy, - scalar_t* __restrict__ dst, - const int* __restrict__ sorted_token_ids, - const int* __restrict__ expert_ids, - const int* __restrict__ num_tokens_post_padded, - const int exp_stride, - const int ncols_x, - const int nrows_x, - const int ncols_y, - const int nrows_y, - const int nrows_dst, - const int top_k) { - const int blocks_per_row_x = ncols_x / qk; - const int blocks_per_col_y = nrows_y / QK8_1; - const int blocks_per_warp = WARP_SIZE_GGUF / qi; - - const int ncols_dst = ncols_y * top_k; - - const auto row_dst_0 = blockIdx.x * mmq_y; - const int& row_x_0 = row_dst_0; - - const auto col_dst_0 = blockIdx.y * mmq_x; - - int token_offs[mmq_x / nwarps]; - for (int i = 0; i < mmq_x; i += nwarps) { - token_offs[i / nwarps] = sorted_token_ids[col_dst_0 + threadIdx.y + i]; - } - - const int exp_idx = expert_ids[blockIdx.y]; - if (exp_idx > 255 || exp_idx < 0) return; - if (blockIdx.y * mmq_x > num_tokens_post_padded[0]) return; - - const block_q_t* x = (const block_q_t*)((char*)vx + exp_idx * exp_stride); - const block_q8_1* y = (const block_q8_1*)(vy); - - int* tile_x_ql = nullptr; - half2* tile_x_dm = nullptr; - int* tile_x_qh = nullptr; - int* tile_x_sc = nullptr; - - allocate_tiles(&tile_x_ql, &tile_x_dm, &tile_x_qh, &tile_x_sc); - - __shared__ int tile_y_qs[mmq_x * WARP_SIZE_GGUF]; - __shared__ half2 tile_y_ds[mmq_x * WARP_SIZE_GGUF / QI8_1]; - - float sum[mmq_y / WARP_SIZE_GGUF][mmq_x / nwarps] = {{0.0f}}; - - for (int ib0 = 0; ib0 < blocks_per_row_x; ib0 += blocks_per_warp) { - load_tiles( - x + row_x_0 * blocks_per_row_x + ib0, - tile_x_ql, - tile_x_dm, - tile_x_qh, - tile_x_sc, - threadIdx.y, - nrows_x - row_x_0 - 1, - threadIdx.x, - blocks_per_row_x); - - const int n_per_r = ((qk * blocks_per_warp) / qr); -#pragma unroll - for (int ir = 0; ir < qr && ib0 * qk + ir * n_per_r < ncols_x; ++ir) { - const auto kqs = ir * WARP_SIZE_GGUF + threadIdx.x; - const int kbxd = kqs / QI8_1; - -#pragma unroll - for (int i = 0; i < mmq_x; i += nwarps) { - const int col_y_eff = token_offs[i / nwarps] / top_k; - const int block_x = ib0 * (qk / QK8_1) + kbxd; - if (col_y_eff < ncols_y && block_x < blocks_per_col_y) { - const block_q8_1* by0 = &y[col_y_eff * blocks_per_col_y + block_x]; - const int index_y = (threadIdx.y + i) * WARP_SIZE_GGUF + kqs % WARP_SIZE_GGUF; - tile_y_qs[index_y] = get_int_from_int8_aligned(by0->qs, threadIdx.x % QI8_1); - } - } - - if (threadIdx.x < n_per_r / QK8_1) { - const auto kby = threadIdx.x % (WARP_SIZE_GGUF / QI8_1); - const int col_y_eff = token_offs[threadIdx.y] / top_k; - const int block_x = ib0 * (qk / QK8_1) + ir * (WARP_SIZE_GGUF / QI8_1) + kby; - - if (col_y_eff < ncols_y && block_x < blocks_per_col_y) { - const half2* dsi_src = &y[col_y_eff * blocks_per_col_y + block_x].ds; - half2* dsi_dst = &tile_y_ds[threadIdx.y * (WARP_SIZE_GGUF / QI8_1) + kby]; - - if (need_sum) { - *dsi_dst = *dsi_src; - } else { - float* dfi_dst = (float*)dsi_dst; - *dfi_dst = __low2float(*dsi_src); - } - } - } - __syncthreads(); - - // #pragma unroll // unrolling this loop causes too much register pressure - for (int k = ir * WARP_SIZE_GGUF / qr; k < (ir + 1) * WARP_SIZE_GGUF / qr; k += vdr) { -#pragma unroll - for (int j = 0; j < mmq_x; j += nwarps) { -#pragma unroll - for (int i = 0; i < mmq_y; i += WARP_SIZE_GGUF) { - sum[i / WARP_SIZE_GGUF][j / nwarps] += vec_dot( - tile_x_ql, tile_x_dm, tile_x_qh, tile_x_sc, tile_y_qs, tile_y_ds, threadIdx.x + i, threadIdx.y + j, k); - } - } - } - __syncthreads(); - } - } - -#pragma unroll - for (int j = 0; j < mmq_x; j += nwarps) { - const int col_dst = token_offs[j / nwarps]; - if (col_dst >= ncols_dst) { - return; - } - -#pragma unroll - for (int i = 0; i < mmq_y; i += WARP_SIZE_GGUF) { - const auto row_dst = row_dst_0 + threadIdx.x + i; - if (row_dst >= nrows_dst) { - continue; - } - dst[col_dst * nrows_dst + row_dst] = sum[i / WARP_SIZE_GGUF][j / nwarps]; - } - } -} - -#if defined(USE_ROCM) -#define MOE_X_Q4_0 8 -#define MOE_Y_Q4_0 128 -#define NWARPS_Q4_0 8 -#else -#define MOE_X_Q4_0 4 -#define MOE_Y_Q4_0 32 -#define NWARPS_Q4_0 4 -#endif - -template -static __global__ void -#if defined(USE_ROCM) -__launch_bounds__(WARP_SIZE_GGUF* NWARPS_Q4_0, 2) -#endif - moe_q4_0( - const void* __restrict__ vx, - const void* __restrict__ vy, - scalar_t* __restrict__ dst, - const int* sorted_token_ids, - const int* expert_ids, - const int* num_tokens_post_padded, - const int exp_stride, - const int ncols_x, - const int nrows_x, - const int ncols_y, - const int nrows_y, - const int nrows_dst, - const int top_k) { - const int mmq_x = MOE_X_Q4_0; - const int mmq_y = MOE_Y_Q4_0; - const int nwarps = NWARPS_Q4_0; - - moe_q< - scalar_t, - QK4_0, - QR4_0, - QI4_0, - true, - block_q4_0, - mmq_x, - mmq_y, - nwarps, - allocate_tiles_q4_0, - load_tiles_q4_0, - VDR_Q4_0_Q8_1_MMQ, - vec_dot_q4_0_q8_1_mul_mat>( - vx, - vy, - dst, - sorted_token_ids, - expert_ids, - num_tokens_post_padded, - exp_stride, - ncols_x, - nrows_x, - ncols_y, - nrows_y, - nrows_dst, - top_k); -} - -template -static void ggml_moe_q4_0_q8_1_cuda( - const void* inp, - const void* w, - scalar_t* dst, - const int* sorted_token_ids, - const int* expert_ids, - const int* num_tokens_post_padded, - const int exp_stride, - const int ncols_x, - const int nrows_x, - const int ncols_y, - const int nrows_y, - const int nrows_dst, - const int top_k, - const int tokens_post_padded, - cudaStream_t stream) { - int mmq_x = MOE_X_Q4_0; - int mmq_y = MOE_Y_Q4_0; - int nwarps = NWARPS_Q4_0; - - const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y; - const int block_num_y = (tokens_post_padded) / mmq_x; - const dim3 block_nums(block_num_x, block_num_y, 1); - const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1); - - if (nrows_x % mmq_y == 0) { - constexpr bool need_check = false; - moe_q4_0<<>>( - w, - inp, - dst, - sorted_token_ids, - expert_ids, - num_tokens_post_padded, - exp_stride, - ncols_x, - nrows_x, - ncols_y, - nrows_y, - nrows_dst, - top_k); - } else { - constexpr bool need_check = true; - moe_q4_0<<>>( - w, - inp, - dst, - sorted_token_ids, - expert_ids, - num_tokens_post_padded, - exp_stride, - ncols_x, - nrows_x, - ncols_y, - nrows_y, - nrows_dst, - top_k); - } -} - -#if defined(USE_ROCM) -#define MOE_X_Q4_1 8 -#define MOE_Y_Q4_1 128 -#define NWARPS_Q4_1 8 -#else -#define MOE_X_Q4_1 4 -#define MOE_Y_Q4_1 32 -#define NWARPS_Q4_1 4 -#endif - -template -static __global__ void -#if defined(USE_ROCM) -__launch_bounds__(WARP_SIZE_GGUF* NWARPS_Q4_1, 2) -#endif - moe_q4_1( - const void* __restrict__ vx, - const void* __restrict__ vy, - scalar_t* __restrict__ dst, - const int* sorted_token_ids, - const int* expert_ids, - const int* num_tokens_post_padded, - const int exp_stride, - const int ncols_x, - const int nrows_x, - const int ncols_y, - const int nrows_y, - const int nrows_dst, - const int top_k) { - const int mmq_x = MOE_X_Q4_1; - const int mmq_y = MOE_Y_Q4_1; - const int nwarps = NWARPS_Q4_1; - - moe_q< - scalar_t, - QK4_1, - QR4_1, - QI4_1, - true, - block_q4_1, - mmq_x, - mmq_y, - nwarps, - allocate_tiles_q4_1, - load_tiles_q4_1, - VDR_Q4_1_Q8_1_MMQ, - vec_dot_q4_1_q8_1_mul_mat>( - vx, - vy, - dst, - sorted_token_ids, - expert_ids, - num_tokens_post_padded, - exp_stride, - ncols_x, - nrows_x, - ncols_y, - nrows_y, - nrows_dst, - top_k); -} - -template -static void ggml_moe_q4_1_q8_1_cuda( - const void* inp, - const void* w, - scalar_t* dst, - const int* sorted_token_ids, - const int* expert_ids, - const int* num_tokens_post_padded, - const int exp_stride, - const int ncols_x, - const int nrows_x, - const int ncols_y, - const int nrows_y, - const int nrows_dst, - const int top_k, - const int tokens_post_padded, - cudaStream_t stream) { - int mmq_x = MOE_X_Q4_1; - int mmq_y = MOE_Y_Q4_1; - int nwarps = NWARPS_Q4_1; - - const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y; - const int block_num_y = (tokens_post_padded) / mmq_x; - const dim3 block_nums(block_num_x, block_num_y, 1); - const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1); - - if (nrows_x % mmq_y == 0) { - constexpr bool need_check = false; - moe_q4_1<<>>( - w, - inp, - dst, - sorted_token_ids, - expert_ids, - num_tokens_post_padded, - exp_stride, - ncols_x, - nrows_x, - ncols_y, - nrows_y, - nrows_dst, - top_k); - } else { - constexpr bool need_check = true; - moe_q4_1<<>>( - w, - inp, - dst, - sorted_token_ids, - expert_ids, - num_tokens_post_padded, - exp_stride, - ncols_x, - nrows_x, - ncols_y, - nrows_y, - nrows_dst, - top_k); - } -} - -#if defined(USE_ROCM) -#define MOE_X_Q5_0 8 -#define MOE_Y_Q5_0 128 -#define NWARPS_Q5_0 8 -#else -#define MOE_X_Q5_0 4 -#define MOE_Y_Q5_0 32 -#define NWARPS_Q5_0 4 -#endif - -template -static __global__ void -#if defined(USE_ROCM) -__launch_bounds__(WARP_SIZE_GGUF* NWARPS_Q5_0, 2) -#endif - moe_q5_0( - const void* __restrict__ vx, - const void* __restrict__ vy, - scalar_t* __restrict__ dst, - const int* sorted_token_ids, - const int* expert_ids, - const int* num_tokens_post_padded, - const int exp_stride, - const int ncols_x, - const int nrows_x, - const int ncols_y, - const int nrows_y, - const int nrows_dst, - const int top_k) { - const int mmq_x = MOE_X_Q5_0; - const int mmq_y = MOE_Y_Q5_0; - const int nwarps = NWARPS_Q5_0; - - moe_q< - scalar_t, - QK5_0, - QR5_0, - QI5_0, - false, - block_q5_0, - mmq_x, - mmq_y, - nwarps, - allocate_tiles_q5_0, - load_tiles_q5_0, - VDR_Q5_0_Q8_1_MMQ, - vec_dot_q5_0_q8_1_mul_mat>( - vx, - vy, - dst, - sorted_token_ids, - expert_ids, - num_tokens_post_padded, - exp_stride, - ncols_x, - nrows_x, - ncols_y, - nrows_y, - nrows_dst, - top_k); -} - -template -static void ggml_moe_q5_0_q8_1_cuda( - const void* inp, - const void* w, - scalar_t* dst, - const int* sorted_token_ids, - const int* expert_ids, - const int* num_tokens_post_padded, - const int exp_stride, - const int ncols_x, - const int nrows_x, - const int ncols_y, - const int nrows_y, - const int nrows_dst, - const int top_k, - const int tokens_post_padded, - cudaStream_t stream) { - const int mmq_x = MOE_X_Q5_0; - const int mmq_y = MOE_Y_Q5_0; - const int nwarps = NWARPS_Q5_0; - - const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y; - const int block_num_y = (tokens_post_padded) / mmq_x; - const dim3 block_nums(block_num_x, block_num_y, 1); - const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1); - - if (nrows_x % mmq_y == 0) { - constexpr bool need_check = false; - moe_q5_0<<>>( - w, - inp, - dst, - sorted_token_ids, - expert_ids, - num_tokens_post_padded, - exp_stride, - ncols_x, - nrows_x, - ncols_y, - nrows_y, - nrows_dst, - top_k); - } else { - constexpr bool need_check = true; - moe_q5_0<<>>( - w, - inp, - dst, - sorted_token_ids, - expert_ids, - num_tokens_post_padded, - exp_stride, - ncols_x, - nrows_x, - ncols_y, - nrows_y, - nrows_dst, - top_k); - } -} - -#if defined(USE_ROCM) -#define MOE_X_Q5_1 8 -#define MOE_Y_Q5_1 128 -#define NWARPS_Q5_1 8 -#else -#define MOE_X_Q5_1 4 -#define MOE_Y_Q5_1 32 -#define NWARPS_Q5_1 4 -#endif - -template -static __global__ void -#if defined(USE_ROCM) -__launch_bounds__(WARP_SIZE_GGUF* NWARPS_Q5_1, 2) -#endif - moe_q5_1( - const void* __restrict__ vx, - const void* __restrict__ vy, - scalar_t* __restrict__ dst, - const int* sorted_token_ids, - const int* expert_ids, - const int* num_tokens_post_padded, - const int exp_stride, - const int ncols_x, - const int nrows_x, - const int ncols_y, - const int nrows_y, - const int nrows_dst, - const int top_k) { - const int mmq_x = MOE_X_Q5_1; - const int mmq_y = MOE_Y_Q5_1; - const int nwarps = NWARPS_Q5_1; - - moe_q< - scalar_t, - QK5_1, - QR5_1, - QI5_1, - true, - block_q5_1, - mmq_x, - mmq_y, - nwarps, - allocate_tiles_q5_1, - load_tiles_q5_1, - VDR_Q5_1_Q8_1_MMQ, - vec_dot_q5_1_q8_1_mul_mat>( - vx, - vy, - dst, - sorted_token_ids, - expert_ids, - num_tokens_post_padded, - exp_stride, - ncols_x, - nrows_x, - ncols_y, - nrows_y, - nrows_dst, - top_k); -} - -template -static void ggml_moe_q5_1_q8_1_cuda( - const void* inp, - const void* w, - scalar_t* dst, - const int* sorted_token_ids, - const int* expert_ids, - const int* num_tokens_post_padded, - const int exp_stride, - const int ncols_x, - const int nrows_x, - const int ncols_y, - const int nrows_y, - const int nrows_dst, - const int top_k, - const int tokens_post_padded, - cudaStream_t stream) { - const int mmq_x = MOE_X_Q5_1; - const int mmq_y = MOE_Y_Q5_1; - const int nwarps = NWARPS_Q5_1; - - const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y; - const int block_num_y = (tokens_post_padded) / mmq_x; - const dim3 block_nums(block_num_x, block_num_y, 1); - const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1); - - if (nrows_x % mmq_y == 0) { - constexpr bool need_check = false; - moe_q5_1<<>>( - w, - inp, - dst, - sorted_token_ids, - expert_ids, - num_tokens_post_padded, - exp_stride, - ncols_x, - nrows_x, - ncols_y, - nrows_y, - nrows_dst, - top_k); - } else { - constexpr bool need_check = true; - moe_q5_1<<>>( - w, - inp, - dst, - sorted_token_ids, - expert_ids, - num_tokens_post_padded, - exp_stride, - ncols_x, - nrows_x, - ncols_y, - nrows_y, - nrows_dst, - top_k); - } -} - -#if defined(USE_ROCM) -#define MOE_X_Q8_0 8 -#define MOE_Y_Q8_0 128 -#define NWARPS_Q8_0 8 -#else -#define MOE_X_Q8_0 4 -#define MOE_Y_Q8_0 32 -#define NWARPS_Q8_0 4 -#endif - -template -static __global__ void -#if defined(USE_ROCM) -__launch_bounds__(WARP_SIZE_GGUF* NWARPS_Q8_0, 2) -#endif - moe_q8_0( - const void* __restrict__ vx, - const void* __restrict__ vy, - scalar_t* __restrict__ dst, - const int* sorted_token_ids, - const int* expert_ids, - const int* num_tokens_post_padded, - const int exp_stride, - const int ncols_x, - const int nrows_x, - const int ncols_y, - const int nrows_y, - const int nrows_dst, - const int top_k) { - const int mmq_x = MOE_X_Q8_0; - const int mmq_y = MOE_Y_Q8_0; - const int nwarps = NWARPS_Q8_0; - - moe_q< - scalar_t, - QK8_0, - QR8_0, - QI8_0, - false, - block_q8_0, - mmq_x, - mmq_y, - nwarps, - allocate_tiles_q8_0, - load_tiles_q8_0, - VDR_Q8_0_Q8_1_MMQ, - vec_dot_q8_0_q8_1_mul_mat>( - vx, - vy, - dst, - sorted_token_ids, - expert_ids, - num_tokens_post_padded, - exp_stride, - ncols_x, - nrows_x, - ncols_y, - nrows_y, - nrows_dst, - top_k); -} - -template -static void ggml_moe_q8_0_q8_1_cuda( - const void* inp, - const void* w, - scalar_t* dst, - const int* sorted_token_ids, - const int* expert_ids, - const int* num_tokens_post_padded, - const int exp_stride, - const int ncols_x, - const int nrows_x, - const int ncols_y, - const int nrows_y, - const int nrows_dst, - const int top_k, - const int tokens_post_padded, - cudaStream_t stream) { - const int mmq_x = MOE_X_Q8_0; - const int mmq_y = MOE_Y_Q8_0; - const int nwarps = NWARPS_Q8_0; - - const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y; - const int block_num_y = (tokens_post_padded) / mmq_x; - const dim3 block_nums(block_num_x, block_num_y, 1); - const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1); - - if (nrows_x % mmq_y == 0) { - constexpr bool need_check = false; - moe_q8_0<<>>( - w, - inp, - dst, - sorted_token_ids, - expert_ids, - num_tokens_post_padded, - exp_stride, - ncols_x, - nrows_x, - ncols_y, - nrows_y, - nrows_dst, - top_k); - } else { - constexpr bool need_check = true; - moe_q8_0<<>>( - w, - inp, - dst, - sorted_token_ids, - expert_ids, - num_tokens_post_padded, - exp_stride, - ncols_x, - nrows_x, - ncols_y, - nrows_y, - nrows_dst, - top_k); - } -} - -#if defined(USE_ROCM) -#define MOE_X_Q2_K 8 -#define MOE_Y_Q2_K 128 -#define NWARPS_Q2_K 8 -#else -#define MOE_X_Q2_K 4 -#define MOE_Y_Q2_K 32 -#define NWARPS_Q2_K 4 -#endif - -template -static __global__ void -#if defined(USE_ROCM) -__launch_bounds__(WARP_SIZE_GGUF* NWARPS_Q2_K, 2) -#endif - moe_q2_K( - const void* __restrict__ vx, - const void* __restrict__ vy, - scalar_t* __restrict__ dst, - const int* sorted_token_ids, - const int* expert_ids, - const int* num_tokens_post_padded, - const int exp_stride, - const int ncols_x, - const int nrows_x, - const int ncols_y, - const int nrows_y, - const int nrows_dst, - const int top_k) { - const int mmq_x = MOE_X_Q2_K; - const int mmq_y = MOE_Y_Q2_K; - const int nwarps = NWARPS_Q2_K; - - moe_q< - scalar_t, - QK_K, - QR2_K, - QI2_K, - false, - block_q2_K, - mmq_x, - mmq_y, - nwarps, - allocate_tiles_q2_K, - load_tiles_q2_K, - VDR_Q2_K_Q8_1_MMQ, - vec_dot_q2_K_q8_1_mul_mat>( - vx, - vy, - dst, - sorted_token_ids, - expert_ids, - num_tokens_post_padded, - exp_stride, - ncols_x, - nrows_x, - ncols_y, - nrows_y, - nrows_dst, - top_k); -} - -template -static void ggml_moe_q2_K_q8_1_cuda( - const void* inp, - const void* w, - scalar_t* dst, - const int* sorted_token_ids, - const int* expert_ids, - const int* num_tokens_post_padded, - const int exp_stride, - const int ncols_x, - const int nrows_x, - const int ncols_y, - const int nrows_y, - const int nrows_dst, - const int top_k, - const int tokens_post_padded, - cudaStream_t stream) { - const int mmq_x = MOE_X_Q2_K; - const int mmq_y = MOE_Y_Q2_K; - const int nwarps = NWARPS_Q2_K; - - const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y; - const int block_num_y = (tokens_post_padded) / mmq_x; - const dim3 block_nums(block_num_x, block_num_y, 1); - const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1); - - if (nrows_x % mmq_y == 0) { - constexpr bool need_check = false; - moe_q2_K<<>>( - w, - inp, - dst, - sorted_token_ids, - expert_ids, - num_tokens_post_padded, - exp_stride, - ncols_x, - nrows_x, - ncols_y, - nrows_y, - nrows_dst, - top_k); - } else { - constexpr bool need_check = true; - moe_q2_K<<>>( - w, - inp, - dst, - sorted_token_ids, - expert_ids, - num_tokens_post_padded, - exp_stride, - ncols_x, - nrows_x, - ncols_y, - nrows_y, - nrows_dst, - top_k); - } -} - -#if defined(USE_ROCM) -#define MOE_X_Q3_K 8 -#define MOE_Y_Q3_K 128 -#define NWARPS_Q3_K 8 -#else -#define MOE_X_Q3_K 4 -#define MOE_Y_Q3_K 32 -#define NWARPS_Q3_K 4 -#endif - -template -static __global__ void -#if defined(USE_ROCM) -__launch_bounds__(WARP_SIZE_GGUF* NWARPS_Q3_K, 2) -#endif - moe_q3_K( - const void* __restrict__ vx, - const void* __restrict__ vy, - scalar_t* __restrict__ dst, - const int* sorted_token_ids, - const int* expert_ids, - const int* num_tokens_post_padded, - const int exp_stride, - const int ncols_x, - const int nrows_x, - const int ncols_y, - const int nrows_y, - const int nrows_dst, - const int top_k) { - - const int mmq_x = MOE_X_Q3_K; - const int mmq_y = MOE_Y_Q3_K; - const int nwarps = NWARPS_Q3_K; - - moe_q< - scalar_t, - QK_K, - QR3_K, - QI3_K, - false, - block_q3_K, - mmq_x, - mmq_y, - nwarps, - allocate_tiles_q3_K, - load_tiles_q3_K, - VDR_Q3_K_Q8_1_MMQ, - vec_dot_q3_K_q8_1_mul_mat>( - vx, - vy, - dst, - sorted_token_ids, - expert_ids, - num_tokens_post_padded, - exp_stride, - ncols_x, - nrows_x, - ncols_y, - nrows_y, - nrows_dst, - top_k); -} -template -static void ggml_moe_q3_K_q8_1_cuda( - const void* inp, - const void* w, - scalar_t* dst, - const int* sorted_token_ids, - const int* expert_ids, - const int* num_tokens_post_padded, - const int exp_stride, - const int ncols_x, - const int nrows_x, - const int ncols_y, - const int nrows_y, - const int nrows_dst, - const int top_k, - const int tokens_post_padded, - cudaStream_t stream) { - const int mmq_x = MOE_X_Q3_K; - const int mmq_y = MOE_Y_Q3_K; - const int nwarps = NWARPS_Q3_K; - - const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y; - const int block_num_y = (tokens_post_padded) / mmq_x; - const dim3 block_nums(block_num_x, block_num_y, 1); - const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1); - - if (nrows_x % mmq_y == 0) { - constexpr bool need_check = false; - moe_q3_K<<>>( - w, - inp, - dst, - sorted_token_ids, - expert_ids, - num_tokens_post_padded, - exp_stride, - ncols_x, - nrows_x, - ncols_y, - nrows_y, - nrows_dst, - top_k); - } else { - constexpr bool need_check = true; - moe_q3_K<<>>( - w, - inp, - dst, - sorted_token_ids, - expert_ids, - num_tokens_post_padded, - exp_stride, - ncols_x, - nrows_x, - ncols_y, - nrows_y, - nrows_dst, - top_k); - } -} - -#if defined(USE_ROCM) -#define MOE_X_Q4_K 8 -#define MOE_Y_Q4_K 128 -#define NWARPS_Q4_K 8 -#else -#define MOE_X_Q4_K 4 -#define MOE_Y_Q4_K 32 -#define NWARPS_Q4_K 4 -#endif - -template -static __global__ void -#if defined(USE_ROCM) -__launch_bounds__(WARP_SIZE_GGUF* NWARPS_Q4_K, 2) -#endif - moe_q4_K( - const void* __restrict__ vx, - const void* __restrict__ vy, - scalar_t* __restrict__ dst, - const int* sorted_token_ids, - const int* expert_ids, - const int* num_tokens_post_padded, - const int exp_stride, - const int ncols_x, - const int nrows_x, - const int ncols_y, - const int nrows_y, - const int nrows_dst, - const int top_k) { - const int mmq_x = MOE_X_Q4_K; - const int mmq_y = MOE_Y_Q4_K; - const int nwarps = NWARPS_Q4_K; - - moe_q< - scalar_t, - QK_K, - QR4_K, - QI4_K, - true, - block_q4_K, - mmq_x, - mmq_y, - nwarps, - allocate_tiles_q4_K, - load_tiles_q4_K, - VDR_Q4_K_Q8_1_MMQ, - vec_dot_q4_K_q8_1_mul_mat>( - vx, - vy, - dst, - sorted_token_ids, - expert_ids, - num_tokens_post_padded, - exp_stride, - ncols_x, - nrows_x, - ncols_y, - nrows_y, - nrows_dst, - top_k); -} - -template -static void ggml_moe_q4_K_q8_1_cuda( - const void* inp, - const void* w, - scalar_t* dst, - const int* sorted_token_ids, - const int* expert_ids, - const int* num_tokens_post_padded, - const int exp_stride, - const int ncols_x, - const int nrows_x, - const int ncols_y, - const int nrows_y, - const int nrows_dst, - const int top_k, - const int tokens_post_padded, - cudaStream_t stream) { - const int mmq_x = MOE_X_Q4_K; - const int mmq_y = MOE_Y_Q4_K; - const int nwarps = NWARPS_Q4_K; - - const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y; - const int block_num_y = (tokens_post_padded) / mmq_x; - const dim3 block_nums(block_num_x, block_num_y, 1); - const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1); - - if (nrows_x % mmq_y == 0) { - constexpr bool need_check = false; - moe_q4_K<<>>( - w, - inp, - dst, - sorted_token_ids, - expert_ids, - num_tokens_post_padded, - exp_stride, - ncols_x, - nrows_x, - ncols_y, - nrows_y, - nrows_dst, - top_k); - } else { - constexpr bool need_check = true; - moe_q4_K<<>>( - w, - inp, - dst, - sorted_token_ids, - expert_ids, - num_tokens_post_padded, - exp_stride, - ncols_x, - nrows_x, - ncols_y, - nrows_y, - nrows_dst, - top_k); - } -} - -#if defined(USE_ROCM) -#define MOE_X_Q5_K 8 -#define MOE_Y_Q5_K 128 -#define NWARPS_Q5_K 8 -#else -#define MOE_X_Q5_K 4 -#define MOE_Y_Q5_K 32 -#define NWARPS_Q5_K 4 -#endif - -template -static __global__ void -#if defined(USE_ROCM) -__launch_bounds__(WARP_SIZE_GGUF* NWARPS_Q5_K, 2) -#endif - moe_q5_K( - const void* __restrict__ vx, - const void* __restrict__ vy, - scalar_t* __restrict__ dst, - const int* sorted_token_ids, - const int* expert_ids, - const int* num_tokens_post_padded, - const int exp_stride, - const int ncols_x, - const int nrows_x, - const int ncols_y, - const int nrows_y, - const int nrows_dst, - const int top_k) { - const int mmq_x = MOE_X_Q5_K; - const int mmq_y = MOE_Y_Q5_K; - const int nwarps = NWARPS_Q5_K; - - moe_q< - scalar_t, - QK_K, - QR5_K, - QI5_K, - true, - block_q5_K, - mmq_x, - mmq_y, - nwarps, - allocate_tiles_q5_K, - load_tiles_q5_K, - VDR_Q5_K_Q8_1_MMQ, - vec_dot_q5_K_q8_1_mul_mat>( - vx, - vy, - dst, - sorted_token_ids, - expert_ids, - num_tokens_post_padded, - exp_stride, - ncols_x, - nrows_x, - ncols_y, - nrows_y, - nrows_dst, - top_k); -} - -template -static void ggml_moe_q5_K_q8_1_cuda( - const void* inp, - const void* w, - scalar_t* dst, - const int* sorted_token_ids, - const int* expert_ids, - const int* num_tokens_post_padded, - const int exp_stride, - const int ncols_x, - const int nrows_x, - const int ncols_y, - const int nrows_y, - const int nrows_dst, - const int top_k, - const int tokens_post_padded, - cudaStream_t stream) { - const int mmq_x = MOE_X_Q5_K; - const int mmq_y = MOE_Y_Q5_K; - const int nwarps = NWARPS_Q5_K; - - const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y; - const int block_num_y = (tokens_post_padded) / mmq_x; - const dim3 block_nums(block_num_x, block_num_y, 1); - const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1); - - if (nrows_x % mmq_y == 0) { - constexpr bool need_check = false; - moe_q5_K<<>>( - w, - inp, - dst, - sorted_token_ids, - expert_ids, - num_tokens_post_padded, - exp_stride, - ncols_x, - nrows_x, - ncols_y, - nrows_y, - nrows_dst, - top_k); - } else { - constexpr bool need_check = true; - moe_q5_K<<>>( - w, - inp, - dst, - sorted_token_ids, - expert_ids, - num_tokens_post_padded, - exp_stride, - ncols_x, - nrows_x, - ncols_y, - nrows_y, - nrows_dst, - top_k); - } -} - -#if defined(USE_ROCM) -#define MOE_X_Q6_K 8 -#define MOE_Y_Q6_K 128 -#define NWARPS_Q6_K 8 -#else -#define MOE_X_Q6_K 4 -#define MOE_Y_Q6_K 32 -#define NWARPS_Q6_K 4 -#endif - -template -static __global__ void -#if defined(USE_ROCM) -__launch_bounds__(WARP_SIZE_GGUF* NWARPS_Q6_K, 2) -#endif - moe_q6_K( - const void* __restrict__ vx, - const void* __restrict__ vy, - scalar_t* __restrict__ dst, - const int* sorted_token_ids, - const int* expert_ids, - const int* num_tokens_post_padded, - const int exp_stride, - const int ncols_x, - const int nrows_x, - const int ncols_y, - const int nrows_y, - const int nrows_dst, - const int top_k) { - const int mmq_x = MOE_X_Q6_K; - const int mmq_y = MOE_Y_Q6_K; - const int nwarps = NWARPS_Q6_K; - - moe_q< - scalar_t, - QK_K, - QR6_K, - QI6_K, - false, - block_q6_K, - mmq_x, - mmq_y, - nwarps, - allocate_tiles_q6_K, - load_tiles_q6_K, - VDR_Q6_K_Q8_1_MMQ, - vec_dot_q6_K_q8_1_mul_mat>( - vx, - vy, - dst, - sorted_token_ids, - expert_ids, - num_tokens_post_padded, - exp_stride, - ncols_x, - nrows_x, - ncols_y, - nrows_y, - nrows_dst, - top_k); -} - -template -static void ggml_moe_q6_K_q8_1_cuda( - const void* inp, - const void* w, - scalar_t* dst, - const int* sorted_token_ids, - const int* expert_ids, - const int* num_tokens_post_padded, - const int exp_stride, - const int ncols_x, - const int nrows_x, - const int ncols_y, - const int nrows_y, - const int nrows_dst, - const int top_k, - const int tokens_post_padded, - cudaStream_t stream) { - const int mmq_x = MOE_X_Q6_K; - const int mmq_y = MOE_Y_Q6_K; - const int nwarps = NWARPS_Q6_K; - - const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y; - const int block_num_y = (tokens_post_padded) / mmq_x; - const dim3 block_nums(block_num_x, block_num_y, 1); - const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1); - - if (nrows_x % mmq_y == 0) { - constexpr bool need_check = false; - moe_q6_K<<>>( - w, - inp, - dst, - sorted_token_ids, - expert_ids, - num_tokens_post_padded, - exp_stride, - ncols_x, - nrows_x, - ncols_y, - nrows_y, - nrows_dst, - top_k); - } else { - constexpr bool need_check = true; - moe_q6_K<<>>( - w, - inp, - dst, - sorted_token_ids, - expert_ids, - num_tokens_post_padded, - exp_stride, - ncols_x, - nrows_x, - ncols_y, - nrows_y, - nrows_dst, - top_k); - } -} diff --git a/python/freetoken/kernel/csrc/gguf/moe_vec.cuh b/python/freetoken/kernel/csrc/gguf/moe_vec.cuh deleted file mode 100644 index 8cef9e080a..0000000000 --- a/python/freetoken/kernel/csrc/gguf/moe_vec.cuh +++ /dev/null @@ -1,413 +0,0 @@ -// copied from -// https://github.com/vllm-project/vllm/blob/4492e3a55428e161ca8db381edc28263e5da4c8d/csrc/quantization/gguf/moe_vec.cuh -// copied and adapted from -// https://github.com/ggerganov/llama.cpp/blob/b2899/ggml-cuda/mmvq.cu -template -static __global__ void moe_vec_q( - const void* __restrict__ vx, - const void* __restrict__ vy, - scalar_t* __restrict__ dst, - const int* topk_ids, - const int topk, - const int ncols, - const int nrows, - const int token_stride) { - const auto row = blockIdx.x * blockDim.y + threadIdx.y; - - const auto token = blockIdx.z / topk; - const auto expert = (topk_ids)[blockIdx.z]; - - if (row >= nrows) { - return; - } - - const int blocks_per_row = ncols / qk; - const int blocks_per_warp = vdr * WARP_SIZE / qi; - - // partial sum for each thread - float tmp = 0.0f; - - const block_q_t* x = ((const block_q_t*)vx) + expert * nrows * blocks_per_row; - const block_q8_1* y = (const block_q8_1*)(((const int*)vy) + token * token_stride); - - for (auto i = threadIdx.x / (qi / vdr); i < blocks_per_row; i += blocks_per_warp) { - const int ibx = row * blocks_per_row + i; // x block index - - const int iby = i * (qk / QK8_1); // y block index that aligns with ibx - - const int iqs = vdr * (threadIdx.x % (qi / vdr)); // x block quant index when casting the quants to int - - tmp += vec_dot_q_cuda(&x[ibx], &y[iby], iqs); - } - - // sum up partial sums and write back result -#pragma unroll - for (int mask = WARP_SIZE / 2; mask > 0; mask >>= 1) { - tmp += SGLANG_SHFL_XOR_SYNC(uint32_t(-1), tmp, mask); - } - - if (threadIdx.x == 0) { - dst[blockIdx.z * nrows + row] = tmp; - } -} - -template -static void moe_vec_q4_0_q8_1_cuda( - const void* vx, - const void* vy, - scalar_t* dst, - const int* topk_ids, - const int top_k, - const int tokens, - const int ncols, - const int nrows, - const int token_stride, - cudaStream_t stream) { - const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, 1, tokens * top_k); - const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - moe_vec_q - <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); -} - -template -static void moe_vec_q4_1_q8_1_cuda( - const void* vx, - const void* vy, - scalar_t* dst, - const int* topk_ids, - const int top_k, - const int tokens, - const int ncols, - const int nrows, - const int token_stride, - cudaStream_t stream) { - const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, 1, tokens * top_k); - const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - moe_vec_q - <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); -} - -template -static void moe_vec_q5_0_q8_1_cuda( - const void* vx, - const void* vy, - scalar_t* dst, - const int* topk_ids, - const int top_k, - const int tokens, - const int ncols, - const int nrows, - const int token_stride, - cudaStream_t stream) { - const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, 1, tokens * top_k); - const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - moe_vec_q - <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); -} - -template -static void moe_vec_q5_1_q8_1_cuda( - const void* vx, - const void* vy, - scalar_t* dst, - const int* topk_ids, - const int top_k, - const int tokens, - const int ncols, - const int nrows, - const int token_stride, - cudaStream_t stream) { - const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, 1, tokens * top_k); - const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - moe_vec_q - <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); -} - -template -static void moe_vec_q8_0_q8_1_cuda( - const void* vx, - const void* vy, - scalar_t* dst, - const int* topk_ids, - const int top_k, - const int tokens, - const int ncols, - const int nrows, - const int token_stride, - cudaStream_t stream) { - const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, 1, tokens * top_k); - const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - moe_vec_q - <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); -} - -template -static void moe_vec_q2_K_q8_1_cuda( - const void* vx, - const void* vy, - scalar_t* dst, - const int* topk_ids, - const int top_k, - const int tokens, - const int ncols, - const int nrows, - const int token_stride, - cudaStream_t stream) { - const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, 1, tokens * top_k); - const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - moe_vec_q - <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); -} - -template -static void moe_vec_q3_K_q8_1_cuda( - const void* vx, - const void* vy, - scalar_t* dst, - const int* topk_ids, - const int top_k, - const int tokens, - const int ncols, - const int nrows, - const int token_stride, - cudaStream_t stream) { - const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, 1, tokens * top_k); - const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - moe_vec_q - <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); -} - -template -static void moe_vec_q4_K_q8_1_cuda( - const void* vx, - const void* vy, - scalar_t* dst, - const int* topk_ids, - const int top_k, - const int tokens, - const int ncols, - const int nrows, - const int token_stride, - cudaStream_t stream) { - const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, 1, tokens * top_k); - const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - moe_vec_q - <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); -} - -template -static void moe_vec_q5_K_q8_1_cuda( - const void* vx, - const void* vy, - scalar_t* dst, - const int* topk_ids, - const int top_k, - const int tokens, - const int ncols, - const int nrows, - const int token_stride, - cudaStream_t stream) { - const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, 1, tokens * top_k); - const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - moe_vec_q - <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); -} - -template -static void moe_vec_q6_K_q8_1_cuda( - const void* vx, - const void* vy, - scalar_t* dst, - const int* topk_ids, - const int top_k, - const int tokens, - const int ncols, - const int nrows, - const int token_stride, - cudaStream_t stream) { - const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, 1, tokens * top_k); - const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - moe_vec_q - <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); -} - -template -static void moe_vec_iq2_xxs_q8_1_cuda( - const void* vx, - const void* vy, - scalar_t* dst, - const int* topk_ids, - const int top_k, - const int tokens, - const int ncols, - const int nrows, - const int token_stride, - cudaStream_t stream) { - const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, 1, tokens * top_k); - const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - moe_vec_q - <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); -} - -template -static void moe_vec_iq2_xs_q8_1_cuda( - const void* vx, - const void* vy, - scalar_t* dst, - const int* topk_ids, - const int top_k, - const int tokens, - const int ncols, - const int nrows, - const int token_stride, - cudaStream_t stream) { - const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, 1, tokens * top_k); - const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - moe_vec_q - <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); -} - -template -static void moe_vec_iq2_s_q8_1_cuda( - const void* vx, - const void* vy, - scalar_t* dst, - const int* topk_ids, - const int top_k, - const int tokens, - const int ncols, - const int nrows, - const int token_stride, - cudaStream_t stream) { - const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, 1, tokens * top_k); - const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - moe_vec_q - <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); -} - -template -static void moe_vec_iq3_xxs_q8_1_cuda( - const void* vx, - const void* vy, - scalar_t* dst, - const int* topk_ids, - const int top_k, - const int tokens, - const int ncols, - const int nrows, - const int token_stride, - cudaStream_t stream) { - const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, 1, tokens * top_k); - const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - moe_vec_q - <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); -} - -template -static void moe_vec_iq1_s_q8_1_cuda( - const void* vx, - const void* vy, - scalar_t* dst, - const int* topk_ids, - const int top_k, - const int tokens, - const int ncols, - const int nrows, - const int token_stride, - cudaStream_t stream) { - const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, 1, tokens * top_k); - const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - moe_vec_q - <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); -} - -template -static void moe_vec_iq1_m_q8_1_cuda( - const void* vx, - const void* vy, - scalar_t* dst, - const int* topk_ids, - const int top_k, - const int tokens, - const int ncols, - const int nrows, - const int token_stride, - cudaStream_t stream) { - const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, 1, tokens * top_k); - const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - moe_vec_q - <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); -} - -template -static void moe_vec_iq4_nl_q8_1_cuda( - const void* vx, - const void* vy, - scalar_t* dst, - const int* topk_ids, - const int top_k, - const int tokens, - const int ncols, - const int nrows, - const int token_stride, - cudaStream_t stream) { - const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, 1, tokens * top_k); - const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - moe_vec_q - <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); -} - -template -static void moe_vec_iq4_xs_q8_1_cuda( - const void* vx, - const void* vy, - scalar_t* dst, - const int* topk_ids, - const int top_k, - const int tokens, - const int ncols, - const int nrows, - const int token_stride, - cudaStream_t stream) { - const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, 1, tokens * top_k); - const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - moe_vec_q - <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); -} - -template -static void moe_vec_iq3_s_q8_1_cuda( - const void* vx, - const void* vy, - scalar_t* dst, - const int* topk_ids, - const int top_k, - const int tokens, - const int ncols, - const int nrows, - const int token_stride, - cudaStream_t stream) { - const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, 1, tokens * top_k); - const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - moe_vec_q - <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); -} diff --git a/python/freetoken/kernel/csrc/gguf/vecdotq.cuh b/python/freetoken/kernel/csrc/gguf/vecdotq.cuh deleted file mode 100644 index 08b4cd269c..0000000000 --- a/python/freetoken/kernel/csrc/gguf/vecdotq.cuh +++ /dev/null @@ -1,2037 +0,0 @@ -// copied from -// https://github.com/vllm-project/vllm/blob/4492e3a55428e161ca8db381edc28263e5da4c8d/csrc/quantization/gguf/vecdotq.cuh -// copied and adapted from https://github.com/ggerganov/llama.cpp/blob/b2899/ggml-cuda/vecdotq.cuh -// and https://github.com/ggerganov/llama.cpp/blob/b2899/ggml-cuda/mmq.cu -static __device__ __forceinline__ int get_int_b2(const void* x, const int& i32) { - const uint16_t* x16 = (const uint16_t*)x; // assume at least 2 byte alignment - - int x32 = x16[2 * i32 + 0] << 0; - x32 |= x16[2 * i32 + 1] << 16; - - return x32; -} - -static __device__ __forceinline__ int get_int_b4(const void* x, const int& i32) { - return ((const int*)x)[i32]; // assume at least 4 byte alignment -} - -static __device__ __forceinline__ int get_int_from_int8(const int8_t* x8, const int& i32) { - const uint16_t* x16 = (const uint16_t*)(x8 + sizeof(int) * i32); // assume at least 2 byte alignment - int x32 = 0; - x32 |= x16[0] << 0; - x32 |= x16[1] << 16; - return x32; -} - -static __device__ __forceinline__ int get_int_from_uint8(const uint8_t* x8, const int& i32) { - const uint16_t* x16 = (const uint16_t*)(x8 + sizeof(int) * i32); // assume at least 2 byte alignment - int x32 = 0; - x32 |= x16[0] << 0; - x32 |= x16[1] << 16; - return x32; -} - -static __device__ __forceinline__ int get_int_from_int8_aligned(const int8_t* x8, const int& i32) { - return *((const int*)(x8 + sizeof(int) * i32)); // assume at least 4 byte alignment -} - -static __device__ __forceinline__ int get_int_from_uint8_aligned(const uint8_t* x8, const int& i32) { - return *((const int*)(x8 + sizeof(int) * i32)); // assume at least 4 byte alignment -} - -// VDR = vec dot ratio, how many contiguous integers each thread processes when the vec dot kernel is called -// MMVQ = mul_mat_vec_q, MMQ = mul_mat_q - -#define VDR_Q4_0_Q8_1_MMVQ 2 -#define VDR_Q4_0_Q8_1_MMQ 4 - -template -static __device__ __forceinline__ float -vec_dot_q4_0_q8_1_impl(const int* v, const int* u, const float& d4, const half2& ds8) { -#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 610 || defined USE_ROCM || defined USE_MUSA - int sumi = 0; - -#pragma unroll - for (int i = 0; i < vdr; ++i) { - const int vi0 = (v[i] >> 0) & 0x0F0F0F0F; - const int vi1 = (v[i] >> 4) & 0x0F0F0F0F; - - // SIMD dot product of quantized values - sumi = __dp4a(vi0, u[2 * i + 0], sumi); - sumi = __dp4a(vi1, u[2 * i + 1], sumi); - } - - const float2 ds8f = __half22float2(ds8); - - // second part effectively subtracts 8 from each quant value - return d4 * (sumi * ds8f.x - (8 * vdr / QI4_0) * ds8f.y); -#endif -} - -#define VDR_Q4_1_Q8_1_MMVQ 2 -#define VDR_Q4_1_Q8_1_MMQ 4 - -template -static __device__ __forceinline__ float -vec_dot_q4_1_q8_1_impl(const int* v, const int* u, const half2& dm4, const half2& ds8) { -#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 610 || defined USE_ROCM || defined USE_MUSA - int sumi = 0; - -#pragma unroll - for (int i = 0; i < vdr; ++i) { - const int vi0 = (v[i] >> 0) & 0x0F0F0F0F; - const int vi1 = (v[i] >> 4) & 0x0F0F0F0F; - - // SIMD dot product of quantized values - sumi = __dp4a(vi0, u[2 * i + 0], sumi); - sumi = __dp4a(vi1, u[2 * i + 1], sumi); - } - - const float2 tmp = __half22float2(__hmul2(dm4, ds8)); - const float d4d8 = tmp.x; - const float m4s8 = tmp.y; - - // scale second part of sum by QI8_1/(vdr * QR4_1) to compensate for multiple threads adding it - return sumi * d4d8 + m4s8 / (QI8_1 / (vdr * QR4_1)); -#endif -} - -#define VDR_Q5_0_Q8_1_MMVQ 2 -#define VDR_Q5_0_Q8_1_MMQ 4 - -template -static __device__ __forceinline__ float -vec_dot_q5_0_q8_1_impl(const int* vl, const int* vh, const int* u, const float& d5, const half2& ds8) { -#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 610 || defined USE_ROCM || defined USE_MUSA - int sumi = 0; - -#pragma unroll - for (int i = 0; i < vdr; ++i) { - int vi0 = (vl[i] >> 0) & 0x0F0F0F0F; // lower 4 qs bits, still need qh as 5th bits - vi0 |= (vh[i] << 4) & 0x00000010; // 0 -> 4 - vi0 |= (vh[i] << 11) & 0x00001000; // 1 -> 12 - vi0 |= (vh[i] << 18) & 0x00100000; // 2 -> 20 - vi0 |= (vh[i] << 25) & 0x10000000; // 3 -> 28 - sumi = __dp4a(vi0, u[2 * i + 0], sumi); // SIMD dot product of quantized values - - int vi1 = (vl[i] >> 4) & 0x0F0F0F0F; // upper 4 qs bits, still need qh as 5th bits - vi1 |= (vh[i] >> 12) & 0x00000010; // 16 -> 4 - vi1 |= (vh[i] >> 5) & 0x00001000; // 17 -> 12 - vi1 |= (vh[i] << 2) & 0x00100000; // 18 -> 20 - vi1 |= (vh[i] << 9) & 0x10000000; // 19 -> 28 - sumi = __dp4a(vi1, u[2 * i + 1], sumi); // SIMD dot product of quantized values - } - - const float2 ds8f = __half22float2(ds8); - - // second part effectively subtracts 16 from each quant value - return d5 * (sumi * ds8f.x - (16 * vdr / QI5_0) * ds8f.y); -#endif -} - -#define VDR_Q5_1_Q8_1_MMVQ 2 -#define VDR_Q5_1_Q8_1_MMQ 4 - -template -static __device__ __forceinline__ float -vec_dot_q5_1_q8_1_impl(const int* vl, const int* vh, const int* u, const half2& dm5, const half2& ds8) { -#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 610 || defined USE_ROCM || defined USE_MUSA - int sumi = 0; - -#pragma unroll - for (int i = 0; i < vdr; ++i) { - int vi0 = (vl[i] >> 0) & 0x0F0F0F0F; // lower 4 qs bits, still need qh as 5th bits - vi0 |= (vh[i] << 4) & 0x00000010; // 0 -> 4 - vi0 |= (vh[i] << 11) & 0x00001000; // 1 -> 12 - vi0 |= (vh[i] << 18) & 0x00100000; // 2 -> 20 - vi0 |= (vh[i] << 25) & 0x10000000; // 3 -> 28 - sumi = __dp4a(vi0, u[2 * i + 0], sumi); // SIMD dot product of quantized values - - int vi1 = (vl[i] >> 4) & 0x0F0F0F0F; // upper 4 qs bits, still need qh as 5th bits - vi1 |= (vh[i] >> 12) & 0x00000010; // 16 -> 4 - vi1 |= (vh[i] >> 5) & 0x00001000; // 17 -> 12 - vi1 |= (vh[i] << 2) & 0x00100000; // 18 -> 20 - vi1 |= (vh[i] << 9) & 0x10000000; // 19 -> 28 - sumi = __dp4a(vi1, u[2 * i + 1], sumi); // SIMD dot product of quantized values - } - - const float2 tmp = __half22float2(__hmul2(dm5, ds8)); - const float d5d8 = tmp.x; - const float m5s8 = tmp.y; - - // scale second part of sum by QI5_1 / vdr to compensate for multiple threads adding it - return sumi * d5d8 + m5s8 / (QI5_1 / vdr); -#endif -} - -#define VDR_Q8_0_Q8_1_MMVQ 2 -#define VDR_Q8_0_Q8_1_MMQ 8 - -template -static __device__ __forceinline__ float -vec_dot_q8_0_q8_1_impl(const int* v, const int* u, const float& d8_0, const float& d8_1) { -#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 610 || defined USE_ROCM || defined USE_MUSA - int sumi = 0; - -#pragma unroll - for (int i = 0; i < vdr; ++i) { - // SIMD dot product of quantized values - sumi = __dp4a(v[i], u[i], sumi); - } - return d8_0 * d8_1 * sumi; -#endif -} - -template -static __device__ __forceinline__ float -vec_dot_q8_1_q8_1_impl(const int* v, const int* u, const half2& dm8, const half2& ds8) { -#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 610 || defined USE_ROCM || defined USE_MUSA - - int sumi = 0; - -#pragma unroll - for (int i = 0; i < vdr; ++i) { - // SIMD dot product of quantized values - sumi = __dp4a(v[i], u[i], sumi); - } - - const float2 tmp = __half22float2(__hmul2(dm8, ds8)); - const float d8d8 = tmp.x; - const float m8s8 = tmp.y; - - // scale second part of sum by QI8_1/ vdr to compensate for multiple threads adding it - return sumi * d8d8 + m8s8 / (QI8_1 / vdr); -#endif -} - -#define VDR_Q2_K_Q8_1_MMVQ 1 -#define VDR_Q2_K_Q8_1_MMQ 2 - -// contiguous v/x values -static __device__ __forceinline__ float vec_dot_q2_K_q8_1_impl_mmvq( - const int& v, - const int* __restrict__ u, - const uint8_t* __restrict__ scales, - const half2& dm2, - const float* __restrict__ d8) { -#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 610 || defined USE_ROCM || defined USE_MUSA - float sumf_d = 0.0f; - float sumf_m = 0.0f; - -#pragma unroll - for (int i = 0; i < QR2_K; ++i) { - const int sc = scales[2 * i]; - - const int vi = (v >> (2 * i)) & 0x03030303; - - sumf_d += d8[i] * (__dp4a(vi, u[i], 0) * (sc & 0xF)); // SIMD dot product - - // fill int with 4x m - int m = sc >> 4; - m |= m << 8; - m |= m << 16; - sumf_m += d8[i] * __dp4a(m, u[i], 0); // multiply constant q2_K part with sum of q8_1 values - } - - const float2 dm2f = __half22float2(dm2); - - return dm2f.x * sumf_d - dm2f.y * sumf_m; -#endif -} - -static __device__ __forceinline__ float vec_dot_q2_K_q8_1_impl_mmq( - const int* __restrict__ v, - const int* __restrict__ u, - const uint8_t* __restrict__ scales, - const half2& dm2, - const float& d8) { -#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 610 || defined USE_ROCM || defined USE_MUSA - int sumi_d = 0; - int sumi_m = 0; - -#pragma unroll - for (int i0 = 0; i0 < QI8_1; i0 += QI8_1 / 2) { - int sumi_d_sc = 0; - - const int sc = scales[i0 / (QI8_1 / 2)]; - - // fill int with 4x m - int m = sc >> 4; - m |= m << 8; - m |= m << 16; - -#pragma unroll - for (int i = i0; i < i0 + QI8_1 / 2; ++i) { - sumi_d_sc = __dp4a(v[i], u[i], sumi_d_sc); // SIMD dot product - sumi_m = __dp4a(m, u[i], sumi_m); // multiply sum of q8_1 values with m - } - - sumi_d += sumi_d_sc * (sc & 0xF); - } - - const float2 dm2f = __half22float2(dm2); - - return d8 * (dm2f.x * sumi_d - dm2f.y * sumi_m); -#endif -} - -#define VDR_Q3_K_Q8_1_MMVQ 1 -#define VDR_Q3_K_Q8_1_MMQ 2 - -// contiguous v/x values -static __device__ __forceinline__ float vec_dot_q3_K_q8_1_impl_mmvq( - const int& vl, - const int& vh, - const int* __restrict__ u, - const uint8_t* __restrict__ scales, - const int& scale_offset, - const float& d3, - const float* __restrict__ d8) { -#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 610 || defined USE_ROCM || defined USE_MUSA - - float sumf = 0.0f; - -#pragma unroll - for (int i = 0; i < QR3_K; ++i) { - const int isc = scale_offset + 2 * i; - - const int isc_low = isc % (QK_K / 32); - const int sc_shift_low = 4 * (isc / (QK_K / 32)); - const int sc_low = (scales[isc_low] >> sc_shift_low) & 0xF; - - const int isc_high = isc % (QK_K / 64); - const int sc_shift_high = 2 * (isc / (QK_K / 64)); - const int sc_high = ((scales[(QK_K / 32) + isc_high] >> sc_shift_high) & 3) << 4; - - const int sc = (sc_low | sc_high) - 32; - - const int vil = (vl >> (2 * i)) & 0x03030303; - - const int vih = ((vh >> i) << 2) & 0x04040404; - - const int vi = __vsubss4(vil, vih); - - sumf += d8[i] * (__dp4a(vi, u[i], 0) * sc); // SIMD dot product - } - - return d3 * sumf; -#endif -} - -static __device__ __forceinline__ float vec_dot_q3_K_q8_1_impl_mmq( - const int* __restrict__ v, - const int* __restrict__ u, - const int8_t* __restrict__ scales, - const float& d3, - const float& d8) { -#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 610 || defined USE_ROCM || defined USE_MUSA - int sumi = 0; - -#pragma unroll - for (int i0 = 0; i0 < QR3_K * VDR_Q3_K_Q8_1_MMQ; i0 += QI8_1 / 2) { - int sumi_sc = 0; - - for (int i = i0; i < i0 + QI8_1 / 2; ++i) { - sumi_sc = __dp4a(v[i], u[i], sumi_sc); // SIMD dot product - } - - sumi += sumi_sc * scales[i0 / (QI8_1 / 2)]; - } - - return d3 * d8 * sumi; -#endif -} - -#define VDR_Q4_K_Q8_1_MMVQ 2 -#define VDR_Q4_K_Q8_1_MMQ 8 - -// contiguous v/x values -static __device__ __forceinline__ float vec_dot_q4_K_q8_1_impl_vmmq( - const int* __restrict__ v, - const int* __restrict__ u, - const uint8_t* __restrict__ sc, - const uint8_t* __restrict__ m, - const half2& dm4, - const float* __restrict__ d8) { -#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 610 || defined USE_ROCM || defined USE_MUSA - - float sumf_d = 0.0f; - float sumf_m = 0.0f; - -#pragma unroll - for (int i = 0; i < QR4_K; ++i) { - const int v0i = (v[0] >> (4 * i)) & 0x0F0F0F0F; - const int v1i = (v[1] >> (4 * i)) & 0x0F0F0F0F; - - const int dot1 = __dp4a(v1i, u[2 * i + 1], __dp4a(v0i, u[2 * i + 0], 0)); // SIMD dot product - const int dot2 = __dp4a(0x01010101, u[2 * i + 1], __dp4a(0x01010101, u[2 * i + 0], 0)); // sum of u - - sumf_d += d8[i] * (dot1 * sc[i]); - sumf_m += d8[i] * (dot2 * m[i]); // multiply constant part of q4_K with sum of q8_1 values - } - - const float2 dm4f = __half22float2(dm4); - return dm4f.x * sumf_d - dm4f.y * sumf_m; -#endif -} - -static __device__ __forceinline__ float vec_dot_q4_K_q8_1_impl_mmq( - const int* __restrict__ v, - const int* __restrict__ u, - const uint8_t* __restrict__ sc, - const uint8_t* __restrict__ m, - const half2& dm4, - const half2* __restrict__ ds8) { -#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 610 || defined USE_ROCM || defined USE_MUSA - float sumf_d = 0.0f; - float sumf_m = 0.0f; - -#pragma unroll - for (int i = 0; i < QR4_K * VDR_Q4_K_Q8_1_MMQ / QI8_1; ++i) { - int sumi_d = 0; - -#pragma unroll - for (int j = 0; j < QI8_1; ++j) { - sumi_d = __dp4a((v[j] >> (4 * i)) & 0x0F0F0F0F, u[i * QI8_1 + j], sumi_d); // SIMD dot product - } - - const float2 ds8f = __half22float2(ds8[i]); - - sumf_d += ds8f.x * (sc[i] * sumi_d); - sumf_m += ds8f.y * m[i]; // sum of q8_1 block * q4_K min val - } - - const float2 dm4f = __half22float2(dm4); - - return dm4f.x * sumf_d - dm4f.y * sumf_m; -#endif -} - -#define VDR_Q5_K_Q8_1_MMVQ 2 -#define VDR_Q5_K_Q8_1_MMQ 8 - -static __device__ __forceinline__ float vec_dot_q5_K_q8_1_impl_vmmq( - const int* __restrict__ vl, - const int* __restrict__ vh, - const int* __restrict__ u, - const uint8_t* __restrict__ sc, - const uint8_t* __restrict__ m, - const half2& dm5, - const float* __restrict__ d8) { -#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 610 || defined USE_ROCM || defined USE_MUSA - - float sumf_d = 0.0f; - float sumf_m = 0.0f; - -#pragma unroll - for (int i = 0; i < QR5_K; ++i) { - const int vl0i = (vl[0] >> (4 * i)) & 0x0F0F0F0F; - const int vl1i = (vl[1] >> (4 * i)) & 0x0F0F0F0F; - - const int vh0i = ((vh[0] >> i) << 4) & 0x10101010; - const int vh1i = ((vh[1] >> i) << 4) & 0x10101010; - - const int v0i = vl0i | vh0i; - const int v1i = vl1i | vh1i; - - const int dot1 = __dp4a(v0i, u[2 * i + 0], __dp4a(v1i, u[2 * i + 1], 0)); // SIMD dot product - const int dot2 = __dp4a(0x01010101, u[2 * i + 0], __dp4a(0x01010101, u[2 * i + 1], 0)); // sum of u - - sumf_d += d8[i] * (dot1 * sc[i]); - sumf_m += d8[i] * (dot2 * m[i]); - } - - const float2 dm5f = __half22float2(dm5); - return dm5f.x * sumf_d - dm5f.y * sumf_m; -#endif -} - -static __device__ __forceinline__ float vec_dot_q5_K_q8_1_impl_mmq( - const int* __restrict__ v, - const int* __restrict__ u, - const uint8_t* __restrict__ sc, - const uint8_t* __restrict__ m, - const half2& dm4, - const half2* __restrict__ ds8) { -#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 610 || defined USE_ROCM || defined USE_MUSA - float sumf_d = 0.0f; - float sumf_m = 0.0f; - -#pragma unroll - for (int i = 0; i < QR5_K * VDR_Q5_K_Q8_1_MMQ / QI8_1; ++i) { - int sumi_d = 0; - -#pragma unroll - for (int j = 0; j < QI8_1; ++j) { - sumi_d = __dp4a(v[i * QI8_1 + j], u[i * QI8_1 + j], sumi_d); // SIMD dot product - } - - const float2 ds8f = __half22float2(ds8[i]); - - sumf_d += ds8f.x * (sc[i] * sumi_d); - sumf_m += ds8f.y * m[i]; // sum of q8_1 block * q4_K min val - } - - const float2 dm4f = __half22float2(dm4); - - return dm4f.x * sumf_d - dm4f.y * sumf_m; -#endif -} - -#define VDR_Q6_K_Q8_1_MMVQ 1 -#define VDR_Q6_K_Q8_1_MMQ 8 - -// contiguous v/x values -static __device__ __forceinline__ float vec_dot_q6_K_q8_1_impl_mmvq( - const int& vl, - const int& vh, - const int* __restrict__ u, - const int8_t* __restrict__ scales, - const float& d, - const float* __restrict__ d8) { -#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 610 || defined USE_ROCM || defined USE_MUSA - float sumf = 0.0f; - -#pragma unroll - for (int i = 0; i < QR6_K; ++i) { - const int sc = scales[4 * i]; - const int vil = (vl >> (4 * i)) & 0x0F0F0F0F; - const int vih = ((vh >> (4 * i)) << 4) & 0x30303030; - const int vi = __vsubss4((vil | vih), 0x20202020); // vi = (vil | vih) - 32 - - sumf += d8[i] * (__dp4a(vi, u[i], 0) * sc); // SIMD dot product - } - - return d * sumf; -#endif -} - -static __device__ __forceinline__ float vec_dot_q6_K_q8_1_impl_mmq( - const int* __restrict__ v, - const int* __restrict__ u, - const int8_t* __restrict__ sc, - const float& d6, - const float* __restrict__ d8) { -#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 610 || defined USE_ROCM || defined USE_MUSA - float sumf_d = 0.0f; - -#pragma unroll - for (int i0 = 0; i0 < VDR_Q6_K_Q8_1_MMQ; i0 += 4) { - int2 sumi_d = {0, 0}; // 2 q6_K scales per q8_1 scale - -#pragma unroll - for (int i = i0; i < i0 + 2; ++i) { - sumi_d.x = __dp4a(v[2 * i + 0], u[2 * i + 0], sumi_d.x); // SIMD dot product - sumi_d.x = __dp4a(v[2 * i + 1], u[2 * i + 1], sumi_d.x); // SIMD dot product - - sumi_d.y = __dp4a(v[2 * i + 4], u[2 * i + 4], sumi_d.y); // SIMD dot product - sumi_d.y = __dp4a(v[2 * i + 5], u[2 * i + 5], sumi_d.y); // SIMD dot product - } - - sumf_d += d8[i0 / 4] * (sc[i0 / 2 + 0] * sumi_d.x + sc[i0 / 2 + 1] * sumi_d.y); - } - - return d6 * sumf_d; -#endif -} - -static __device__ __forceinline__ float -vec_dot_q4_0_q8_1(const void* __restrict__ vbq, const block_q8_1* __restrict__ bq8_1, const int& iqs) { - const block_q4_0* bq4_0 = (const block_q4_0*)vbq; - - int v[VDR_Q4_0_Q8_1_MMVQ]; - int u[2 * VDR_Q4_0_Q8_1_MMVQ]; - -#pragma unroll - for (int i = 0; i < VDR_Q4_0_Q8_1_MMVQ; ++i) { - v[i] = get_int_from_uint8(bq4_0->qs, iqs + i); - u[2 * i + 0] = get_int_from_int8_aligned(bq8_1->qs, iqs + i); - u[2 * i + 1] = get_int_from_int8_aligned(bq8_1->qs, iqs + i + QI4_0); - } - - return vec_dot_q4_0_q8_1_impl(v, u, __half2float(bq4_0->d), bq8_1->ds); -} - -template -static __device__ __forceinline__ void allocate_tiles_q4_0(int** x_ql, half2** x_dm, int** x_qh, int** x_sc) { - __shared__ int tile_x_qs[mmq_y * (WARP_SIZE_GGUF) + mmq_y]; - __shared__ float tile_x_d[mmq_y * (WARP_SIZE_GGUF / QI4_0) + mmq_y / QI4_0]; - *x_ql = tile_x_qs; - *x_dm = (half2*)tile_x_d; -} - -template -static __device__ __forceinline__ void load_tiles_q4_0( - const void* __restrict__ vx, - int* __restrict__ x_ql, - half2* __restrict__ x_dm, - int* __restrict__ x_qh, - int* __restrict__ x_sc, - const int& i_offset, - const int& i_max, - const int& k, - const int& blocks_per_row) { - const int kbx = k / QI4_0; - const int kqsx = k % QI4_0; - - const block_q4_0* bx0 = (const block_q4_0*)vx; - float* x_dmf = (float*)x_dm; - -#pragma unroll - for (int i0 = 0; i0 < mmq_y; i0 += nwarps) { - int i = i0 + i_offset; - if (need_check) { - i = min(i, i_max); - } - const block_q4_0* bxi = bx0 + i * blocks_per_row + kbx; - x_ql[i * (WARP_SIZE_GGUF + 1) + k] = get_int_from_uint8(bxi->qs, kqsx); - // x_dmf[i * (WARP_SIZE_GGUF/QI4_0) + i / QI4_0 + kbx] = bxi->d; - } - - const int blocks_per_tile_x_row = WARP_SIZE_GGUF / QI4_0; - const int kbxd = k % blocks_per_tile_x_row; - -#pragma unroll - for (int i0 = 0; i0 < mmq_y; i0 += nwarps * QI4_0) { - int i = i0 + i_offset * QI4_0 + k / blocks_per_tile_x_row; - if (need_check) { - i = min(i, i_max); - } - const block_q4_0* bxi = bx0 + i * blocks_per_row + kbxd; - x_dmf[i * (WARP_SIZE_GGUF / QI4_0) + i / QI4_0 + kbxd] = __half2float(bxi->d); - } -} - -static __device__ __forceinline__ float vec_dot_q4_0_q8_1_mul_mat( - const int* __restrict__ x_ql, - const half2* __restrict__ x_dm, - const int* __restrict__ x_qh, - const int* __restrict__ x_sc, - const int* __restrict__ y_qs, - const half2* __restrict__ y_ds, - const int& i, - const int& j, - const int& k) { - (void)x_qh; - (void)x_sc; - - const int kyqs = k % (QI8_1 / 2) + QI8_1 * (k / (QI8_1 / 2)); - const float* x_dmf = (const float*)x_dm; - - int u[2 * VDR_Q4_0_Q8_1_MMQ]; - -#pragma unroll - for (int l = 0; l < VDR_Q4_0_Q8_1_MMQ; ++l) { - u[2 * l + 0] = y_qs[j * WARP_SIZE_GGUF + (kyqs + l) % WARP_SIZE_GGUF]; - u[2 * l + 1] = y_qs[j * WARP_SIZE_GGUF + (kyqs + l + QI4_0) % WARP_SIZE_GGUF]; - } - - return vec_dot_q4_0_q8_1_impl( - &x_ql[i * (WARP_SIZE_GGUF + 1) + k], - u, - x_dmf[i * (WARP_SIZE_GGUF / QI4_0) + i / QI4_0 + k / QI4_0], - y_ds[j * (WARP_SIZE_GGUF / QI8_1) + (2 * k / QI8_1) % (WARP_SIZE_GGUF / QI8_1)]); -} - -static __device__ __forceinline__ float -vec_dot_q4_1_q8_1(const void* __restrict__ vbq, const block_q8_1* __restrict__ bq8_1, const int& iqs) { - const block_q4_1* bq4_1 = (const block_q4_1*)vbq; - - int v[VDR_Q4_1_Q8_1_MMVQ]; - int u[2 * VDR_Q4_1_Q8_1_MMVQ]; - -#pragma unroll - for (int i = 0; i < VDR_Q4_1_Q8_1_MMVQ; ++i) { - v[i] = get_int_from_uint8_aligned(bq4_1->qs, iqs + i); - u[2 * i + 0] = get_int_from_int8_aligned(bq8_1->qs, iqs + i); - u[2 * i + 1] = get_int_from_int8_aligned(bq8_1->qs, iqs + i + QI4_1); - } - - return vec_dot_q4_1_q8_1_impl(v, u, bq4_1->dm, bq8_1->ds); -} - -template -static __device__ __forceinline__ void allocate_tiles_q4_1(int** x_ql, half2** x_dm, int** x_qh, int** x_sc) { - __shared__ int tile_x_qs[mmq_y * (WARP_SIZE_GGUF) + +mmq_y]; - __shared__ half2 tile_x_dm[mmq_y * (WARP_SIZE_GGUF / QI4_1) + mmq_y / QI4_1]; - *x_ql = tile_x_qs; - *x_dm = tile_x_dm; -} - -template -static __device__ __forceinline__ void load_tiles_q4_1( - const void* __restrict__ vx, - int* __restrict__ x_ql, - half2* __restrict__ x_dm, - int* __restrict__ x_qh, - int* __restrict__ x_sc, - const int& i_offset, - const int& i_max, - const int& k, - const int& blocks_per_row) { - const int kbx = k / QI4_1; - const int kqsx = k % QI4_1; - - const block_q4_1* bx0 = (const block_q4_1*)vx; - -#pragma unroll - for (int i0 = 0; i0 < mmq_y; i0 += nwarps) { - int i = i0 + i_offset; - if (need_check) { - i = min(i, i_max); - } - const block_q4_1* bxi = bx0 + i * blocks_per_row + kbx; - x_ql[i * (WARP_SIZE_GGUF + 1) + k] = get_int_from_uint8_aligned(bxi->qs, kqsx); - } - - const int blocks_per_tile_x_row = WARP_SIZE_GGUF / QI4_1; - const int kbxd = k % blocks_per_tile_x_row; - -#pragma unroll - for (int i0 = 0; i0 < mmq_y; i0 += nwarps * QI4_1) { - int i = i0 + i_offset * QI4_1 + k / blocks_per_tile_x_row; - if (need_check) { - i = min(i, i_max); - } - const block_q4_1* bxi = bx0 + i * blocks_per_row + kbxd; - x_dm[i * (WARP_SIZE_GGUF / QI4_1) + i / QI4_1 + kbxd] = bxi->dm; - } -} - -static __device__ __forceinline__ float vec_dot_q4_1_q8_1_mul_mat( - const int* __restrict__ x_ql, - const half2* __restrict__ x_dm, - const int* __restrict__ x_qh, - const int* __restrict__ x_sc, - const int* __restrict__ y_qs, - const half2* __restrict__ y_ds, - const int& i, - const int& j, - const int& k) { - const int kyqs = k % (QI8_1 / 2) + QI8_1 * (k / (QI8_1 / 2)); - - int u[2 * VDR_Q4_1_Q8_1_MMQ]; - -#pragma unroll - for (int l = 0; l < VDR_Q4_1_Q8_1_MMQ; ++l) { - u[2 * l + 0] = y_qs[j * WARP_SIZE_GGUF + (kyqs + l) % WARP_SIZE_GGUF]; - u[2 * l + 1] = y_qs[j * WARP_SIZE_GGUF + (kyqs + l + QI4_1) % WARP_SIZE_GGUF]; - } - - return vec_dot_q4_1_q8_1_impl( - &x_ql[i * (WARP_SIZE_GGUF + 1) + k], - u, - x_dm[i * (WARP_SIZE_GGUF / QI4_1) + i / QI4_1 + k / QI4_1], - y_ds[j * (WARP_SIZE_GGUF / QI8_1) + (2 * k / QI8_1) % (WARP_SIZE_GGUF / QI8_1)]); -} - -static __device__ __forceinline__ float -vec_dot_q5_0_q8_1(const void* __restrict__ vbq, const block_q8_1* __restrict__ bq8_1, const int& iqs) { - const block_q5_0* bq5_0 = (const block_q5_0*)vbq; - - int vl[VDR_Q5_0_Q8_1_MMVQ]; - int vh[VDR_Q5_0_Q8_1_MMVQ]; - int u[2 * VDR_Q5_0_Q8_1_MMVQ]; - -#pragma unroll - for (int i = 0; i < VDR_Q5_0_Q8_1_MMVQ; ++i) { - vl[i] = get_int_from_uint8(bq5_0->qs, iqs + i); - vh[i] = get_int_from_uint8(bq5_0->qh, 0) >> (4 * (iqs + i)); - u[2 * i + 0] = get_int_from_int8_aligned(bq8_1->qs, iqs + i); - u[2 * i + 1] = get_int_from_int8_aligned(bq8_1->qs, iqs + i + QI5_0); - } - - return vec_dot_q5_0_q8_1_impl(vl, vh, u, __half2float(bq5_0->d), bq8_1->ds); -} - -template -static __device__ __forceinline__ void allocate_tiles_q5_0(int** x_ql, half2** x_dm, int** x_qh, int** x_sc) { - __shared__ int tile_x_ql[mmq_y * (2 * WARP_SIZE_GGUF) + mmq_y]; - __shared__ float tile_x_d[mmq_y * (WARP_SIZE_GGUF / QI5_0) + mmq_y / QI5_0]; - - *x_ql = tile_x_ql; - *x_dm = (half2*)tile_x_d; -} - -template -static __device__ __forceinline__ void load_tiles_q5_0( - const void* __restrict__ vx, - int* __restrict__ x_ql, - half2* __restrict__ x_dm, - int* __restrict__ x_qh, - int* __restrict__ x_sc, - const int& i_offset, - const int& i_max, - const int& k, - const int& blocks_per_row) { - const int kbx = k / QI5_0; - const int kqsx = k % QI5_0; - - const block_q5_0* bx0 = (const block_q5_0*)vx; - -#pragma unroll - for (int i0 = 0; i0 < mmq_y; i0 += nwarps) { - int i = i0 + i_offset; - - if (need_check) { - i = min(i, i_max); - } - const block_q5_0* bxi = bx0 + i * blocks_per_row + kbx; - const int ql = get_int_from_uint8(bxi->qs, kqsx); - const int qh = get_int_from_uint8(bxi->qh, 0) >> (4 * (k % QI5_0)); - - int qs0 = (ql >> 0) & 0x0F0F0F0F; - qs0 |= (qh << 4) & 0x00000010; // 0 -> 4 - qs0 |= (qh << 11) & 0x00001000; // 1 -> 12 - qs0 |= (qh << 18) & 0x00100000; // 2 -> 20 - qs0 |= (qh << 25) & 0x10000000; // 3 -> 28 - qs0 = __vsubss4(qs0, 0x10101010); // subtract 16 - - x_ql[i * (2 * WARP_SIZE_GGUF + 1) + 2 * k + 0] = qs0; - - int qs1 = (ql >> 4) & 0x0F0F0F0F; - qs1 |= (qh >> 12) & 0x00000010; // 16 -> 4 - qs1 |= (qh >> 5) & 0x00001000; // 17 -> 12 - qs1 |= (qh << 2) & 0x00100000; // 18 -> 20 - qs1 |= (qh << 9) & 0x10000000; // 19 -> 28 - qs1 = __vsubss4(qs1, 0x10101010); // subtract 16 - - x_ql[i * (2 * WARP_SIZE_GGUF + 1) + 2 * k + 1] = qs1; - } - - const int blocks_per_tile_x_row = WARP_SIZE_GGUF / QI5_0; - const int kbxd = k % blocks_per_tile_x_row; - float* x_dmf = (float*)x_dm; - -#pragma unroll - for (int i0 = 0; i0 < mmq_y; i0 += nwarps * QI5_0) { - int i = i0 + i_offset * QI5_0 + k / blocks_per_tile_x_row; - - if (need_check) { - i = min(i, i_max); - } - - const block_q5_0* bxi = bx0 + i * blocks_per_row + kbxd; - x_dmf[i * (WARP_SIZE_GGUF / QI5_0) + i / QI5_0 + kbxd] = __half2float(bxi->d); - } -} - -static __device__ __forceinline__ float vec_dot_q5_0_q8_1_mul_mat( - const int* __restrict__ x_ql, - const half2* __restrict__ x_dm, - const int* __restrict__ x_qh, - const int* __restrict__ x_sc, - const int* __restrict__ y_qs, - const half2* __restrict__ y_ds, - const int& i, - const int& j, - const int& k) { - const int kyqs = k % (QI8_1 / 2) + QI8_1 * (k / (QI8_1 / 2)); - const int index_bx = i * (WARP_SIZE_GGUF / QI5_0) + i / QI5_0 + k / QI5_0; - const float* x_dmf = (const float*)x_dm; - const float* y_df = (const float*)y_ds; - - int u[2 * VDR_Q5_0_Q8_1_MMQ]; - -#pragma unroll - for (int l = 0; l < VDR_Q5_0_Q8_1_MMQ; ++l) { - u[2 * l + 0] = y_qs[j * WARP_SIZE_GGUF + (kyqs + l) % WARP_SIZE_GGUF]; - u[2 * l + 1] = y_qs[j * WARP_SIZE_GGUF + (kyqs + l + QI5_0) % WARP_SIZE_GGUF]; - } - - return vec_dot_q8_0_q8_1_impl( - &x_ql[i * (2 * WARP_SIZE_GGUF + 1) + 2 * k], - u, - x_dmf[index_bx], - y_df[j * (WARP_SIZE_GGUF / QI8_1) + (2 * k / QI8_1) % (WARP_SIZE_GGUF / QI8_1)]); -} - -static __device__ __forceinline__ float -vec_dot_q5_1_q8_1(const void* __restrict__ vbq, const block_q8_1* __restrict__ bq8_1, const int& iqs) { - const block_q5_1* bq5_1 = (const block_q5_1*)vbq; - - int vl[VDR_Q5_1_Q8_1_MMVQ]; - int vh[VDR_Q5_1_Q8_1_MMVQ]; - int u[2 * VDR_Q5_1_Q8_1_MMVQ]; - -#pragma unroll - for (int i = 0; i < VDR_Q5_1_Q8_1_MMVQ; ++i) { - vl[i] = get_int_from_uint8_aligned(bq5_1->qs, iqs + i); - vh[i] = get_int_from_uint8_aligned(bq5_1->qh, 0) >> (4 * (iqs + i)); - u[2 * i + 0] = get_int_from_int8_aligned(bq8_1->qs, iqs + i); - u[2 * i + 1] = get_int_from_int8_aligned(bq8_1->qs, iqs + i + QI5_1); - } - - return vec_dot_q5_1_q8_1_impl(vl, vh, u, bq5_1->dm, bq8_1->ds); -} - -template -static __device__ __forceinline__ void allocate_tiles_q5_1(int** x_ql, half2** x_dm, int** x_qh, int** x_sc) { - __shared__ int tile_x_ql[mmq_y * (2 * WARP_SIZE_GGUF) + mmq_y]; - __shared__ half2 tile_x_dm[mmq_y * (WARP_SIZE_GGUF / QI5_1) + mmq_y / QI5_1]; - - *x_ql = tile_x_ql; - *x_dm = tile_x_dm; -} - -template -static __device__ __forceinline__ void load_tiles_q5_1( - const void* __restrict__ vx, - int* __restrict__ x_ql, - half2* __restrict__ x_dm, - int* __restrict__ x_qh, - int* __restrict__ x_sc, - const int& i_offset, - const int& i_max, - const int& k, - const int& blocks_per_row) { - const int kbx = k / QI5_1; - const int kqsx = k % QI5_1; - - const block_q5_1* bx0 = (const block_q5_1*)vx; - -#pragma unroll - for (int i0 = 0; i0 < mmq_y; i0 += nwarps) { - int i = i0 + i_offset; - - if (need_check) { - i = min(i, i_max); - } - - const block_q5_1* bxi = bx0 + i * blocks_per_row + kbx; - - const int ql = get_int_from_uint8_aligned(bxi->qs, kqsx); - const int qh = get_int_from_uint8_aligned(bxi->qh, 0) >> (4 * (k % QI5_1)); - - int qs0 = (ql >> 0) & 0x0F0F0F0F; - qs0 |= (qh << 4) & 0x00000010; // 0 -> 4 - qs0 |= (qh << 11) & 0x00001000; // 1 -> 12 - qs0 |= (qh << 18) & 0x00100000; // 2 -> 20 - qs0 |= (qh << 25) & 0x10000000; // 3 -> 28 - - x_ql[i * (2 * WARP_SIZE_GGUF + 1) + 2 * k + 0] = qs0; - - int qs1 = (ql >> 4) & 0x0F0F0F0F; - qs1 |= (qh >> 12) & 0x00000010; // 16 -> 4 - qs1 |= (qh >> 5) & 0x00001000; // 17 -> 12 - qs1 |= (qh << 2) & 0x00100000; // 18 -> 20 - qs1 |= (qh << 9) & 0x10000000; // 19 -> 28 - - x_ql[i * (2 * WARP_SIZE_GGUF + 1) + 2 * k + 1] = qs1; - } - - const int blocks_per_tile_x_row = WARP_SIZE_GGUF / QI5_1; - const int kbxd = k % blocks_per_tile_x_row; - -#pragma unroll - for (int i0 = 0; i0 < mmq_y; i0 += nwarps * QI5_1) { - int i = i0 + i_offset * QI5_1 + k / blocks_per_tile_x_row; - - if (need_check) { - i = min(i, i_max); - } - - const block_q5_1* bxi = bx0 + i * blocks_per_row + kbxd; - - x_dm[i * (WARP_SIZE_GGUF / QI5_1) + i / QI5_1 + kbxd] = bxi->dm; - } -} - -static __device__ __forceinline__ float vec_dot_q5_1_q8_1_mul_mat( - const int* __restrict__ x_ql, - const half2* __restrict__ x_dm, - const int* __restrict__ x_qh, - const int* __restrict__ x_sc, - const int* __restrict__ y_qs, - const half2* __restrict__ y_ds, - const int& i, - const int& j, - const int& k) { - const int kyqs = k % (QI8_1 / 2) + QI8_1 * (k / (QI8_1 / 2)); - const int index_bx = i * (WARP_SIZE_GGUF / QI5_1) + +i / QI5_1 + k / QI5_1; - - int u[2 * VDR_Q5_1_Q8_1_MMQ]; - -#pragma unroll - for (int l = 0; l < VDR_Q5_1_Q8_1_MMQ; ++l) { - u[2 * l + 0] = y_qs[j * WARP_SIZE_GGUF + (kyqs + l) % WARP_SIZE_GGUF]; - u[2 * l + 1] = y_qs[j * WARP_SIZE_GGUF + (kyqs + l + QI5_1) % WARP_SIZE_GGUF]; - } - - return vec_dot_q8_1_q8_1_impl( - &x_ql[i * (2 * WARP_SIZE_GGUF + 1) + 2 * k], - u, - x_dm[index_bx], - y_ds[j * (WARP_SIZE_GGUF / QI8_1) + (2 * k / QI8_1) % (WARP_SIZE_GGUF / QI8_1)]); -} - -static __device__ __forceinline__ float -vec_dot_q8_0_q8_1(const void* __restrict__ vbq, const block_q8_1* __restrict__ bq8_1, const int& iqs) { - const block_q8_0* bq8_0 = (const block_q8_0*)vbq; - - int v[VDR_Q8_0_Q8_1_MMVQ]; - int u[VDR_Q8_0_Q8_1_MMVQ]; - -#pragma unroll - for (int i = 0; i < VDR_Q8_0_Q8_1_MMVQ; ++i) { - v[i] = get_int_from_int8(bq8_0->qs, iqs + i); - u[i] = get_int_from_int8_aligned(bq8_1->qs, iqs + i); - } - - return vec_dot_q8_0_q8_1_impl(v, u, __half2float(bq8_0->d), __low2float(bq8_1->ds)); -} - -template -static __device__ __forceinline__ void allocate_tiles_q8_0(int** x_ql, half2** x_dm, int** x_qh, int** x_sc) { - __shared__ int tile_x_qs[mmq_y * (WARP_SIZE_GGUF) + mmq_y]; - __shared__ float tile_x_d[mmq_y * (WARP_SIZE_GGUF / QI8_0) + mmq_y / QI8_0]; - - *x_ql = tile_x_qs; - *x_dm = (half2*)tile_x_d; -} - -template -static __device__ __forceinline__ void load_tiles_q8_0( - const void* __restrict__ vx, - int* __restrict__ x_ql, - half2* __restrict__ x_dm, - int* __restrict__ x_qh, - int* __restrict__ x_sc, - const int& i_offset, - const int& i_max, - const int& k, - const int& blocks_per_row) { - const int kbx = k / QI8_0; - const int kqsx = k % QI8_0; - float* x_dmf = (float*)x_dm; - - const block_q8_0* bx0 = (const block_q8_0*)vx; - -#pragma unroll - for (int i0 = 0; i0 < mmq_y; i0 += nwarps) { - int i = i0 + i_offset; - - if (need_check) { - i = min(i, i_max); - } - const block_q8_0* bxi = bx0 + i * blocks_per_row + kbx; - x_ql[i * (WARP_SIZE_GGUF + 1) + k] = get_int_from_int8(bxi->qs, kqsx); - } - - const int blocks_per_tile_x_row = WARP_SIZE_GGUF / QI8_0; - const int kbxd = k % blocks_per_tile_x_row; - -#pragma unroll - for (int i0 = 0; i0 < mmq_y; i0 += nwarps * QI8_0) { - int i = i0 + i_offset * QI8_0 + k / blocks_per_tile_x_row; - - if (need_check) { - i = min(i, i_max); - } - const block_q8_0* bxi = bx0 + i * blocks_per_row + kbxd; - x_dmf[i * (WARP_SIZE_GGUF / QI8_0) + i / QI8_0 + kbxd] = __half2float(bxi->d); - } -} - -static __device__ __forceinline__ float vec_dot_q8_0_q8_1_mul_mat( - const int* __restrict__ x_ql, - const half2* __restrict__ x_dm, - const int* __restrict__ x_qh, - const int* __restrict__ x_sc, - const int* __restrict__ y_qs, - const half2* __restrict__ y_ds, - const int& i, - const int& j, - const int& k) { - const float* x_dmf = (const float*)x_dm; - const float* y_df = (const float*)y_ds; - - return vec_dot_q8_0_q8_1_impl( - &x_ql[i * (WARP_SIZE_GGUF + 1) + k], - &y_qs[j * WARP_SIZE_GGUF + k], - x_dmf[i * (WARP_SIZE_GGUF / QI8_0) + i / QI8_0 + k / QI8_0], - y_df[j * (WARP_SIZE_GGUF / QI8_1) + k / QI8_1]); -} - -static __device__ __forceinline__ float -vec_dot_q2_K_q8_1(const void* __restrict__ vbq, const block_q8_1* __restrict__ bq8_1, const int& iqs) { - const block_q2_K* bq2_K = (const block_q2_K*)vbq; - - const int bq8_offset = QR2_K * (iqs / QI8_1); - const int scale_offset = iqs - iqs % QI8_1 + (iqs % QI8_1) / (QI8_1 / 2); - - const uint8_t* scales = bq2_K->scales + scale_offset; - - const int v = get_int_from_uint8_aligned(bq2_K->qs, iqs); - int u[QR2_K]; - float d8[QR2_K]; - -#pragma unroll - for (int i = 0; i < QR2_K; ++i) { - u[i] = get_int_from_int8_aligned(bq8_1[bq8_offset + i].qs, iqs % QI8_1); - d8[i] = __low2float(bq8_1[bq8_offset + i].ds); - } - - return vec_dot_q2_K_q8_1_impl_mmvq(v, u, scales, bq2_K->dm, d8); -} - -template -static __device__ __forceinline__ void allocate_tiles_q2_K(int** x_ql, half2** x_dm, int** x_qh, int** x_sc) { - __shared__ int tile_x_ql[mmq_y * (WARP_SIZE_GGUF) + mmq_y]; - __shared__ half2 tile_x_dm[mmq_y * (WARP_SIZE_GGUF / QI2_K) + mmq_y / QI2_K]; - __shared__ int tile_x_sc[mmq_y * (WARP_SIZE_GGUF / 4) + mmq_y / 4]; - - *x_ql = tile_x_ql; - *x_dm = tile_x_dm; - *x_sc = tile_x_sc; -} - -template -static __device__ __forceinline__ void load_tiles_q2_K( - const void* __restrict__ vx, - int* __restrict__ x_ql, - half2* __restrict__ x_dm, - int* __restrict__ x_qh, - int* __restrict__ x_sc, - const int& i_offset, - const int& i_max, - const int& k, - const int& blocks_per_row) { - const int kbx = k / QI2_K; - const int kqsx = k % QI2_K; - - const block_q2_K* bx0 = (const block_q2_K*)vx; - -#pragma unroll - for (int i0 = 0; i0 < mmq_y; i0 += nwarps) { - int i = i0 + i_offset; - - if (need_check) { - i = min(i, i_max); - } - const block_q2_K* bxi = bx0 + i * blocks_per_row + kbx; - x_ql[i * (WARP_SIZE_GGUF + 1) + k] = get_int_from_uint8_aligned(bxi->qs, kqsx); - } - - const int blocks_per_tile_x_row = WARP_SIZE_GGUF / QI2_K; - const int kbxd = k % blocks_per_tile_x_row; - -#pragma unroll - for (int i0 = 0; i0 < mmq_y; i0 += nwarps * QI2_K) { - int i = (i0 + i_offset * QI2_K + k / blocks_per_tile_x_row) % mmq_y; - - if (need_check) { - i = min(i, i_max); - } - const block_q2_K* bxi = bx0 + i * blocks_per_row + kbxd; - x_dm[i * (WARP_SIZE_GGUF / QI2_K) + i / QI2_K + kbxd] = bxi->dm; - } - -#pragma unroll - for (int i0 = 0; i0 < mmq_y; i0 += nwarps * 4) { - int i = i0 + i_offset * 4 + k / (WARP_SIZE_GGUF / 4); - - if (need_check) { - i = min(i, i_max); - } - const block_q2_K* bxi = bx0 + i * blocks_per_row + (k % (WARP_SIZE_GGUF / 4)) / (QI2_K / 4); - x_sc[i * (WARP_SIZE_GGUF / 4) + i / 4 + k % (WARP_SIZE_GGUF / 4)] = - get_int_from_uint8_aligned(bxi->scales, k % (QI2_K / 4)); - } -} - -static __device__ __forceinline__ float vec_dot_q2_K_q8_1_mul_mat( - const int* __restrict__ x_ql, - const half2* __restrict__ x_dm, - const int* __restrict__ x_qh, - const int* __restrict__ x_sc, - const int* __restrict__ y_qs, - const half2* __restrict__ y_ds, - const int& i, - const int& j, - const int& k) { - const int kbx = k / QI2_K; - const int ky = (k % QI2_K) * QR2_K; - const float* y_df = (const float*)y_ds; - - int v[QR2_K * VDR_Q2_K_Q8_1_MMQ]; - - const int kqsx = i * (WARP_SIZE_GGUF + 1) + kbx * QI2_K + (QI2_K / 2) * (ky / (2 * QI2_K)) + ky % (QI2_K / 2); - const int shift = 2 * ((ky % (2 * QI2_K)) / (QI2_K / 2)); - -#pragma unroll - for (int l = 0; l < QR2_K * VDR_Q2_K_Q8_1_MMQ; ++l) { - v[l] = (x_ql[kqsx + l] >> shift) & 0x03030303; - } - - const uint8_t* scales = ((const uint8_t*)&x_sc[i * (WARP_SIZE_GGUF / 4) + i / 4 + kbx * 4]) + ky / 4; - - const int index_y = j * WARP_SIZE_GGUF + (QR2_K * k) % WARP_SIZE_GGUF; - return vec_dot_q2_K_q8_1_impl_mmq( - v, &y_qs[index_y], scales, x_dm[i * (WARP_SIZE_GGUF / QI2_K) + i / QI2_K + kbx], y_df[index_y / QI8_1]); -} - -static __device__ __forceinline__ float -vec_dot_q3_K_q8_1(const void* __restrict__ vbq, const block_q8_1* __restrict__ bq8_1, const int& iqs) { - const block_q3_K* bq3_K = (const block_q3_K*)vbq; - - const int bq8_offset = QR3_K * (iqs / (QI3_K / 2)); - const int scale_offset = iqs - iqs % QI8_1 + (iqs % QI8_1) / (QI8_1 / 2); - - const float d = __half2float(bq3_K->d); - - const int vl = get_int_from_uint8(bq3_K->qs, iqs); - - // invert the mask with ~ so that a 0/1 results in 4/0 being subtracted - const int vh = ~get_int_from_uint8(bq3_K->hmask, iqs % (QI3_K / 2)) >> bq8_offset; - - int u[QR3_K]; - float d8[QR3_K]; - -#pragma unroll - for (int i = 0; i < QR3_K; ++i) { - u[i] = get_int_from_int8_aligned(bq8_1[bq8_offset + i].qs, iqs % QI8_1); - d8[i] = __low2float(bq8_1[bq8_offset + i].ds); - } - - return vec_dot_q3_K_q8_1_impl_mmvq(vl, vh, u, bq3_K->scales, scale_offset, d, d8); -} - -template -static __device__ __forceinline__ void allocate_tiles_q3_K(int** x_ql, half2** x_dm, int** x_qh, int** x_sc) { - __shared__ int tile_x_ql[mmq_y * (WARP_SIZE_GGUF) + mmq_y]; - __shared__ half2 tile_x_dm[mmq_y * (WARP_SIZE_GGUF / QI3_K) + mmq_y / QI3_K]; - __shared__ int tile_x_qh[mmq_y * (WARP_SIZE_GGUF / 2) + mmq_y / 2]; - __shared__ int tile_x_sc[mmq_y * (WARP_SIZE_GGUF / 4) + mmq_y / 4]; - - *x_ql = tile_x_ql; - *x_dm = tile_x_dm; - *x_qh = tile_x_qh; - *x_sc = tile_x_sc; -} - -template -static __device__ __forceinline__ void load_tiles_q3_K( - const void* __restrict__ vx, - int* __restrict__ x_ql, - half2* __restrict__ x_dm, - int* __restrict__ x_qh, - int* __restrict__ x_sc, - const int& i_offset, - const int& i_max, - const int& k, - const int& blocks_per_row) { - const int kbx = k / QI3_K; - const int kqsx = k % QI3_K; - - const block_q3_K* bx0 = (const block_q3_K*)vx; - -#pragma unroll - for (int i0 = 0; i0 < mmq_y; i0 += nwarps) { - int i = i0 + i_offset; - if (need_check) { - i = min(i, i_max); - } - const block_q3_K* bxi = bx0 + i * blocks_per_row + kbx; - x_ql[i * (WARP_SIZE_GGUF + 1) + k] = get_int_from_uint8(bxi->qs, kqsx); - } - - const int blocks_per_tile_x_row = WARP_SIZE_GGUF / QI3_K; - const int kbxd = k % blocks_per_tile_x_row; - float* x_dmf = (float*)x_dm; - -#pragma unroll - for (int i0 = 0; i0 < mmq_y; i0 += nwarps * QI3_K) { - int i = (i0 + i_offset * QI3_K + k / blocks_per_tile_x_row) % mmq_y; - if (need_check) { - i = min(i, i_max); - } - const block_q3_K* bxi = bx0 + i * blocks_per_row + kbxd; - x_dmf[i * (WARP_SIZE_GGUF / QI3_K) + i / QI3_K + kbxd] = __half2float(bxi->d); - } - -#pragma unroll - for (int i0 = 0; i0 < mmq_y; i0 += nwarps * 2) { - int i = i0 + i_offset * 2 + k / (WARP_SIZE_GGUF / 2); - if (need_check) { - i = min(i, i_max); - } - const block_q3_K* bxi = bx0 + i * blocks_per_row + (k % (WARP_SIZE_GGUF / 2)) / (QI3_K / 2); - // invert the mask with ~ so that a 0/1 results in 4/0 being subtracted - x_qh[i * (WARP_SIZE_GGUF / 2) + i / 2 + k % (WARP_SIZE_GGUF / 2)] = - ~get_int_from_uint8(bxi->hmask, k % (QI3_K / 2)); - } - -#pragma unroll - for (int i0 = 0; i0 < mmq_y; i0 += nwarps * 4) { - int i = i0 + i_offset * 4 + k / (WARP_SIZE_GGUF / 4); - if (need_check) { - i = min(i, i_max); - } - const block_q3_K* bxi = bx0 + i * blocks_per_row + (k % (WARP_SIZE_GGUF / 4)) / (QI3_K / 4); - - const int ksc = k % (QI3_K / 4); - - const int ksc_low = ksc % (QI3_K / 8); - const int shift_low = 4 * (ksc / (QI3_K / 8)); - const int sc_low = (get_int_from_uint8(bxi->scales, ksc_low) >> shift_low) & 0x0F0F0F0F; - - const int ksc_high = QI3_K / 8; - const int shift_high = 2 * ksc; - const int sc_high = ((get_int_from_uint8(bxi->scales, ksc_high) >> shift_high) << 4) & 0x30303030; - - const int sc = __vsubss4(sc_low | sc_high, 0x20202020); - - x_sc[i * (WARP_SIZE_GGUF / 4) + i / 4 + k % (WARP_SIZE_GGUF / 4)] = sc; - } -} - -static __device__ __forceinline__ float vec_dot_q3_K_q8_1_mul_mat( - const int* __restrict__ x_ql, - const half2* __restrict__ x_dm, - const int* __restrict__ x_qh, - const int* __restrict__ x_sc, - const int* __restrict__ y_qs, - const half2* __restrict__ y_ds, - const int& i, - const int& j, - const int& k) { - const int kbx = k / QI3_K; - const int ky = (k % QI3_K) * QR3_K; - const float* x_dmf = (const float*)x_dm; - const float* y_df = (const float*)y_ds; - - const int8_t* scales = ((const int8_t*)(x_sc + i * (WARP_SIZE_GGUF / 4) + i / 4 + kbx * 4)) + ky / 4; - - int v[QR3_K * VDR_Q3_K_Q8_1_MMQ]; - -#pragma unroll - for (int l = 0; l < QR3_K * VDR_Q3_K_Q8_1_MMQ; ++l) { - const int kqsx = i * (WARP_SIZE_GGUF + 1) + kbx * QI3_K + (QI3_K / 2) * (ky / (2 * QI3_K)) + ky % (QI3_K / 2); - const int shift = 2 * ((ky % 32) / 8); - const int vll = (x_ql[kqsx + l] >> shift) & 0x03030303; - - const int vh = x_qh[i * (WARP_SIZE_GGUF / 2) + i / 2 + kbx * (QI3_K / 2) + (ky + l) % 8] >> ((ky + l) / 8); - const int vlh = (vh << 2) & 0x04040404; - - v[l] = __vsubss4(vll, vlh); - } - - const int index_y = j * WARP_SIZE_GGUF + (k * QR3_K) % WARP_SIZE_GGUF; - return vec_dot_q3_K_q8_1_impl_mmq( - v, &y_qs[index_y], scales, x_dmf[i * (WARP_SIZE_GGUF / QI3_K) + i / QI3_K + kbx], y_df[index_y / QI8_1]); -} - -static __device__ __forceinline__ float -vec_dot_q4_K_q8_1(const void* __restrict__ vbq, const block_q8_1* __restrict__ bq8_1, const int& iqs) { - const block_q4_K* bq4_K = (const block_q4_K*)vbq; - - int v[2]; - int u[2 * QR4_K]; - float d8[QR4_K]; - - // iqs is in 0,2..30. bq8_offset = iqs/4 -> bq8_offset = 0, 2, 4, 6 - const int bq8_offset = QR4_K * ((iqs / 2) / (QI8_1 / 2)); - - // iqs = 0....3 -> bq8_offset = 0, want q4_offset = 0, 4, 8, 12 - // iqs = 4....7 -> bq8_offset = 2, want q4_offset = 32, 36, 40, 44 - // iqs = 8...11 -> bq8_offset = 4, want q4_offset = 64, 68, 72, 76 - // iqs = 12..15 -> bq8_offset = 6, want q4_offset = 96, 100, 104, 108 - - const int* q4 = (const int*)(bq4_K->qs + 16 * bq8_offset + 4 * ((iqs / 2) % 4)); - v[0] = q4[0]; - v[1] = q4[4]; - - const uint16_t* scales = (const uint16_t*)bq4_K->scales; - uint16_t aux[2]; - const int j = bq8_offset / 2; - if (j < 2) { - aux[0] = scales[j + 0] & 0x3f3f; - aux[1] = scales[j + 2] & 0x3f3f; - } else { - aux[0] = ((scales[j + 2] >> 0) & 0x0f0f) | ((scales[j - 2] & 0xc0c0) >> 2); - aux[1] = ((scales[j + 2] >> 4) & 0x0f0f) | ((scales[j - 0] & 0xc0c0) >> 2); - } - const uint8_t* sc = (const uint8_t*)aux; - const uint8_t* m = sc + 2; - - for (int i = 0; i < QR4_K; ++i) { - const block_q8_1* bq8i = bq8_1 + bq8_offset + i; - d8[i] = __low2float(bq8i->ds); - - const int* q8 = (const int*)bq8i->qs + ((iqs / 2) % 4); - u[2 * i + 0] = q8[0]; - u[2 * i + 1] = q8[4]; - } - - return vec_dot_q4_K_q8_1_impl_vmmq(v, u, sc, m, bq4_K->dm, d8); -} - -template -static __device__ __forceinline__ void allocate_tiles_q4_K(int** x_ql, half2** x_dm, int** x_qh, int** x_sc) { - __shared__ int tile_x_ql[mmq_y * (WARP_SIZE_GGUF) + mmq_y]; - __shared__ half2 tile_x_dm[mmq_y * (WARP_SIZE_GGUF / QI4_K) + mmq_y / QI4_K]; - __shared__ int tile_x_sc[mmq_y * (WARP_SIZE_GGUF / 8) + mmq_y / 8]; - - *x_ql = tile_x_ql; - *x_dm = tile_x_dm; - *x_sc = tile_x_sc; -} - -template -static __device__ __forceinline__ void load_tiles_q4_K( - const void* __restrict__ vx, - int* __restrict__ x_ql, - half2* __restrict__ x_dm, - int* __restrict__ x_qh, - int* __restrict__ x_sc, - const int& i_offset, - const int& i_max, - const int& k, - const int& blocks_per_row) { - const int kbx = k / QI4_K; // == 0 if QK_K == 256 - const int kqsx = k % QI4_K; // == k if QK_K == 256 - - const block_q4_K* bx0 = (const block_q4_K*)vx; - -#pragma unroll - for (int i0 = 0; i0 < mmq_y; i0 += nwarps) { - int i = i0 + i_offset; - - if (need_check) { - i = min(i, i_max); - } - const block_q4_K* bxi = bx0 + i * blocks_per_row + kbx; - x_ql[i * (WARP_SIZE_GGUF + 1) + k] = get_int_from_uint8_aligned(bxi->qs, kqsx); - } - - const int blocks_per_tile_x_row = WARP_SIZE_GGUF / QI4_K; // == 1 if QK_K == 256 - const int kbxd = k % blocks_per_tile_x_row; // == 0 if QK_K == 256 - -#pragma unroll - for (int i0 = 0; i0 < mmq_y; i0 += nwarps * QI4_K) { - int i = (i0 + i_offset * QI4_K + k / blocks_per_tile_x_row) % mmq_y; - if (need_check) { - i = min(i, i_max); - } - const block_q4_K* bxi = bx0 + i * blocks_per_row + kbxd; - x_dm[i * (WARP_SIZE_GGUF / QI4_K) + i / QI4_K + kbxd] = bxi->dm; - } - -#pragma unroll - for (int i0 = 0; i0 < mmq_y; i0 += nwarps * 8) { - int i = (i0 + i_offset * 8 + k / (WARP_SIZE_GGUF / 8)) % mmq_y; - - if (need_check) { - i = min(i, i_max); - } - - const block_q4_K* bxi = bx0 + i * blocks_per_row + (k % (WARP_SIZE_GGUF / 8)) / (QI4_K / 8); - - const int* scales = (const int*)bxi->scales; - - const int ksc = k % (WARP_SIZE_GGUF / 8); - // scale arrangement after the following two lines: sc0,...,sc3, sc4,...,sc7, m0,...,m3, m4,...,m8 - int scales8 = (scales[(ksc % 2) + (ksc != 0)] >> (4 * (ksc & (ksc / 2)))) & 0x0F0F0F0F; // lower 4 bits - scales8 |= (scales[ksc / 2] >> (2 * (ksc % 2))) & 0x30303030; // upper 2 bits - - x_sc[i * (WARP_SIZE_GGUF / 8) + i / 8 + ksc] = scales8; - } -} - -static __device__ __forceinline__ float vec_dot_q4_K_q8_1_mul_mat( - const int* __restrict__ x_ql, - const half2* __restrict__ x_dm, - const int* __restrict__ x_qh, - const int* __restrict__ x_sc, - const int* __restrict__ y_qs, - const half2* __restrict__ y_ds, - const int& i, - const int& j, - const int& k) { - (void)x_qh; - - const uint8_t* sc = ((const uint8_t*)&x_sc[i * (WARP_SIZE_GGUF / 8) + i / 8 + k / 16]) + 2 * ((k % 16) / 8); - - const int index_y = j * WARP_SIZE_GGUF + (QR4_K * k) % WARP_SIZE_GGUF; - return vec_dot_q4_K_q8_1_impl_mmq( - &x_ql[i * (WARP_SIZE_GGUF + 1) + k], - &y_qs[index_y], - sc, - sc + 8, - x_dm[i * (WARP_SIZE_GGUF / QI4_K) + i / QI4_K], - &y_ds[index_y / QI8_1]); -} - -static __device__ __forceinline__ float -vec_dot_q5_K_q8_1(const void* __restrict__ vbq, const block_q8_1* __restrict__ bq8_1, const int& iqs) { - const block_q5_K* bq5_K = (const block_q5_K*)vbq; - - int vl[2]; - int vh[2]; - int u[2 * QR5_K]; - float d8[QR5_K]; - - const int bq8_offset = QR5_K * ((iqs / 2) / (QI8_1 / 2)); - const int* ql = (const int*)(bq5_K->qs + 16 * bq8_offset + 4 * ((iqs / 2) % 4)); - const int* qh = (const int*)(bq5_K->qh + 4 * ((iqs / 2) % 4)); - - vl[0] = ql[0]; - vl[1] = ql[4]; - - vh[0] = qh[0] >> bq8_offset; - vh[1] = qh[4] >> bq8_offset; - - const uint16_t* scales = (const uint16_t*)bq5_K->scales; - uint16_t aux[2]; - const int j = bq8_offset / 2; - if (j < 2) { - aux[0] = scales[j + 0] & 0x3f3f; - aux[1] = scales[j + 2] & 0x3f3f; - } else { - aux[0] = ((scales[j + 2] >> 0) & 0x0f0f) | ((scales[j - 2] & 0xc0c0) >> 2); - aux[1] = ((scales[j + 2] >> 4) & 0x0f0f) | ((scales[j - 0] & 0xc0c0) >> 2); - } - const uint8_t* sc = (const uint8_t*)aux; - const uint8_t* m = sc + 2; - -#pragma unroll - for (int i = 0; i < QR5_K; ++i) { - const block_q8_1* bq8i = bq8_1 + bq8_offset + i; - d8[i] = __low2float(bq8i->ds); - - const int* q8 = (const int*)bq8i->qs + ((iqs / 2) % 4); - u[2 * i + 0] = q8[0]; - u[2 * i + 1] = q8[4]; - } - - return vec_dot_q5_K_q8_1_impl_vmmq(vl, vh, u, sc, m, bq5_K->dm, d8); -} - -template -static __device__ __forceinline__ void allocate_tiles_q5_K(int** x_ql, half2** x_dm, int** x_qh, int** x_sc) { - __shared__ int tile_x_ql[mmq_y * (2 * WARP_SIZE_GGUF) + mmq_y]; - __shared__ half2 tile_x_dm[mmq_y * (WARP_SIZE_GGUF / QI5_K) + mmq_y / QI5_K]; - __shared__ int tile_x_sc[mmq_y * (WARP_SIZE_GGUF / 8) + mmq_y / 8]; - - *x_ql = tile_x_ql; - *x_dm = tile_x_dm; - *x_sc = tile_x_sc; -} - -template -static __device__ __forceinline__ void load_tiles_q5_K( - const void* __restrict__ vx, - int* __restrict__ x_ql, - half2* __restrict__ x_dm, - int* __restrict__ x_qh, - int* __restrict__ x_sc, - const int& i_offset, - const int& i_max, - const int& k, - const int& blocks_per_row) { - const int kbx = k / QI5_K; // == 0 if QK_K == 256 - const int kqsx = k % QI5_K; // == k if QK_K == 256 - - const block_q5_K* bx0 = (const block_q5_K*)vx; - -#pragma unroll - for (int i0 = 0; i0 < mmq_y; i0 += nwarps) { - int i = i0 + i_offset; - - if (need_check) { - i = min(i, i_max); - } - - const block_q5_K* bxi = bx0 + i * blocks_per_row + kbx; - const int ky = QR5_K * kqsx; - - const int ql = get_int_from_uint8_aligned(bxi->qs, kqsx); - const int ql0 = (ql >> 0) & 0x0F0F0F0F; - const int ql1 = (ql >> 4) & 0x0F0F0F0F; - - const int qh = get_int_from_uint8_aligned(bxi->qh, kqsx % (QI5_K / 4)); - const int qh0 = ((qh >> (2 * (kqsx / (QI5_K / 4)) + 0)) << 4) & 0x10101010; - const int qh1 = ((qh >> (2 * (kqsx / (QI5_K / 4)) + 1)) << 4) & 0x10101010; - - const int kq0 = ky - ky % (QI5_K / 2) + k % (QI5_K / 4) + 0; - const int kq1 = ky - ky % (QI5_K / 2) + k % (QI5_K / 4) + (QI5_K / 4); - - x_ql[i * (2 * WARP_SIZE_GGUF + 1) + kq0] = ql0 | qh0; - x_ql[i * (2 * WARP_SIZE_GGUF + 1) + kq1] = ql1 | qh1; - } - - const int blocks_per_tile_x_row = WARP_SIZE_GGUF / QI5_K; // == 1 if QK_K == 256 - const int kbxd = k % blocks_per_tile_x_row; // == 0 if QK_K == 256 - -#pragma unroll - for (int i0 = 0; i0 < mmq_y; i0 += nwarps * QI5_K) { - int i = (i0 + i_offset * QI5_K + k / blocks_per_tile_x_row) % mmq_y; - - if (need_check) { - i = min(i, i_max); - } - - const block_q5_K* bxi = bx0 + i * blocks_per_row + kbxd; - x_dm[i * (WARP_SIZE_GGUF / QI5_K) + i / QI5_K + kbxd] = bxi->dm; - } - -#pragma unroll - for (int i0 = 0; i0 < mmq_y; i0 += nwarps * 8) { - int i = (i0 + i_offset * 8 + k / (WARP_SIZE_GGUF / 8)) % mmq_y; - - if (need_check) { - i = min(i, i_max); - } - - const block_q5_K* bxi = bx0 + i * blocks_per_row + (k % (WARP_SIZE_GGUF / 8)) / (QI5_K / 8); - - const int* scales = (const int*)bxi->scales; - - const int ksc = k % (WARP_SIZE_GGUF / 8); - - // scale arrangement after the following two lines: sc0,...,sc3, sc4,...,sc7, m0,...,m3, m4,...,m8 - int scales8 = (scales[(ksc % 2) + (ksc != 0)] >> (4 * (ksc & (ksc / 2)))) & 0x0F0F0F0F; // lower 4 bits - scales8 |= (scales[ksc / 2] >> (2 * (ksc % 2))) & 0x30303030; // upper 2 bits - - x_sc[i * (WARP_SIZE_GGUF / 8) + i / 8 + ksc] = scales8; - } -} - -static __device__ __forceinline__ float vec_dot_q5_K_q8_1_mul_mat( - const int* __restrict__ x_ql, - const half2* __restrict__ x_dm, - const int* __restrict__ x_qh, - const int* __restrict__ x_sc, - const int* __restrict__ y_qs, - const half2* __restrict__ y_ds, - const int& i, - const int& j, - const int& k) { - const uint8_t* sc = ((const uint8_t*)&x_sc[i * (WARP_SIZE_GGUF / 8) + i / 8 + k / 16]) + 2 * ((k % 16) / 8); - - const int index_x = i * (QR5_K * WARP_SIZE_GGUF + 1) + QR5_K * k; - const int index_y = j * WARP_SIZE_GGUF + (QR5_K * k) % WARP_SIZE_GGUF; - return vec_dot_q5_K_q8_1_impl_mmq( - &x_ql[index_x], - &y_qs[index_y], - sc, - sc + 8, - x_dm[i * (WARP_SIZE_GGUF / QI5_K) + i / QI5_K], - &y_ds[index_y / QI8_1]); -} - -static __device__ __forceinline__ float -vec_dot_q6_K_q8_1(const void* __restrict__ vbq, const block_q8_1* __restrict__ bq8_1, const int& iqs) { - const block_q6_K* bq6_K = (const block_q6_K*)vbq; - - const int bq8_offset = 2 * QR6_K * (iqs / (QI6_K / 2)) + (iqs % (QI6_K / 2)) / (QI6_K / 4); - const int scale_offset = (QI6_K / 4) * (iqs / (QI6_K / 2)) + (iqs % (QI6_K / 2)) / (QI6_K / 8); - const int vh_shift = 2 * ((iqs % (QI6_K / 2)) / (QI6_K / 4)); - - const int vl = get_int_from_uint8(bq6_K->ql, iqs); - const int vh = get_int_from_uint8(bq6_K->qh, (QI6_K / 4) * (iqs / (QI6_K / 2)) + iqs % (QI6_K / 4)) >> vh_shift; - - const int8_t* scales = bq6_K->scales + scale_offset; - - int u[QR6_K]; - float d8[QR6_K]; - -#pragma unroll - for (int i = 0; i < QR6_K; ++i) { - u[i] = get_int_from_int8_aligned(bq8_1[bq8_offset + 2 * i].qs, iqs % QI8_1); - d8[i] = __low2float(bq8_1[bq8_offset + 2 * i].ds); - } - - return vec_dot_q6_K_q8_1_impl_mmvq(vl, vh, u, scales, __half2float(bq6_K->d), d8); -} - -template -static __device__ __forceinline__ void allocate_tiles_q6_K(int** x_ql, half2** x_dm, int** x_qh, int** x_sc) { - __shared__ int tile_x_ql[mmq_y * (2 * WARP_SIZE_GGUF) + mmq_y]; - __shared__ half2 tile_x_dm[mmq_y * (WARP_SIZE_GGUF / QI6_K) + mmq_y / QI6_K]; - __shared__ int tile_x_sc[mmq_y * (WARP_SIZE_GGUF / 8) + mmq_y / 8]; - - *x_ql = tile_x_ql; - *x_dm = tile_x_dm; - *x_sc = tile_x_sc; -} - -template -static __device__ __forceinline__ void load_tiles_q6_K( - const void* __restrict__ vx, - int* __restrict__ x_ql, - half2* __restrict__ x_dm, - int* __restrict__ x_qh, - int* __restrict__ x_sc, - const int& i_offset, - const int& i_max, - const int& k, - const int& blocks_per_row) { - const int kbx = k / QI6_K; // == 0 if QK_K == 256 - const int kqsx = k % QI6_K; // == k if QK_K == 256 - - const block_q6_K* bx0 = (const block_q6_K*)vx; - -#pragma unroll - for (int i0 = 0; i0 < mmq_y; i0 += nwarps) { - int i = i0 + i_offset; - - if (need_check) { - i = min(i, i_max); - } - - const block_q6_K* bxi = bx0 + i * blocks_per_row + kbx; - const int ky = QR6_K * kqsx; - - const int ql = get_int_from_uint8(bxi->ql, kqsx); - const int ql0 = (ql >> 0) & 0x0F0F0F0F; - const int ql1 = (ql >> 4) & 0x0F0F0F0F; - - const int qh = get_int_from_uint8(bxi->qh, (QI6_K / 4) * (kqsx / (QI6_K / 2)) + kqsx % (QI6_K / 4)); - const int qh0 = ((qh >> (2 * ((kqsx % (QI6_K / 2)) / (QI6_K / 4)))) << 4) & 0x30303030; - const int qh1 = (qh >> (2 * ((kqsx % (QI6_K / 2)) / (QI6_K / 4)))) & 0x30303030; - - const int kq0 = ky - ky % QI6_K + k % (QI6_K / 2) + 0; - const int kq1 = ky - ky % QI6_K + k % (QI6_K / 2) + (QI6_K / 2); - - x_ql[i * (2 * WARP_SIZE_GGUF + 1) + kq0] = __vsubss4(ql0 | qh0, 0x20202020); - x_ql[i * (2 * WARP_SIZE_GGUF + 1) + kq1] = __vsubss4(ql1 | qh1, 0x20202020); - } - - const int blocks_per_tile_x_row = WARP_SIZE_GGUF / QI6_K; // == 1 if QK_K == 256 - const int kbxd = k % blocks_per_tile_x_row; // == 0 if QK_K == 256 - float* x_dmf = (float*)x_dm; - -#pragma unroll - for (int i0 = 0; i0 < mmq_y; i0 += nwarps * QI6_K) { - int i = (i0 + i_offset * QI6_K + k / blocks_per_tile_x_row) % mmq_y; - - if (need_check) { - i = min(i, i_max); - } - - const block_q6_K* bxi = bx0 + i * blocks_per_row + kbxd; - - x_dmf[i * (WARP_SIZE_GGUF / QI6_K) + i / QI6_K + kbxd] = __half2float(bxi->d); - } - -#pragma unroll - for (int i0 = 0; i0 < mmq_y; i0 += nwarps * 8) { - int i = (i0 + i_offset * 8 + k / (WARP_SIZE_GGUF / 8)) % mmq_y; - - if (need_check) { - i = min(i, i_max); - } - - const block_q6_K* bxi = bx0 + i * blocks_per_row + (k % (WARP_SIZE_GGUF / 8)) / 4; - - x_sc[i * (WARP_SIZE_GGUF / 8) + i / 8 + k % (WARP_SIZE_GGUF / 8)] = get_int_from_int8(bxi->scales, k % (QI6_K / 8)); - } -} - -static __device__ __forceinline__ float vec_dot_q6_K_q8_1_mul_mat( - const int* __restrict__ x_ql, - const half2* __restrict__ x_dm, - const int* __restrict__ x_qh, - const int* __restrict__ x_sc, - const int* __restrict__ y_qs, - const half2* __restrict__ y_ds, - const int& i, - const int& j, - const int& k) { - const float* x_dmf = (const float*)x_dm; - const float* y_df = (const float*)y_ds; - - const int8_t* sc = ((const int8_t*)&x_sc[i * (WARP_SIZE_GGUF / 8) + i / 8 + k / 8]); - - const int index_x = i * (QR6_K * WARP_SIZE_GGUF + 1) + QR6_K * k; - const int index_y = j * WARP_SIZE_GGUF + (QR6_K * k) % WARP_SIZE_GGUF; - return vec_dot_q6_K_q8_1_impl_mmq( - &x_ql[index_x], &y_qs[index_y], sc, x_dmf[i * (WARP_SIZE_GGUF / QI6_K) + i / QI6_K], &y_df[index_y / QI8_1]); -} - -static __device__ __forceinline__ float -vec_dot_iq2_xxs_q8_1(const void* __restrict__ vbq, const block_q8_1* __restrict__ bq8_1, const int& iqs) { - const block_iq2_xxs* bq2 = (const block_iq2_xxs*)vbq; - - const int ib32 = iqs; - const uint16_t* q2 = bq2->qs + 4 * ib32; - const uint8_t* aux8 = (const uint8_t*)q2; - const int8_t* q8 = bq8_1[ib32].qs; - uint32_t aux32 = q2[2] | (q2[3] << 16); - int sumi = 0; - for (int l = 0; l < 4; ++l) { - const uint8_t* grid = (const uint8_t*)(iq2xxs_grid + aux8[l]); - const uint8_t signs = ksigns_iq2xs[aux32 & 127]; - for (int j = 0; j < 8; ++j) { - sumi += q8[j] * grid[j] * (signs & kmask_iq2xs[j] ? -1 : 1); - } - q8 += 8; - aux32 >>= 7; - } - const float d = __half2float(bq2->d) * (0.5f + aux32) * __half2float(bq8_1[ib32].ds.x) * 0.25f; - return d * sumi; -} - -static __device__ __forceinline__ float -vec_dot_iq2_xs_q8_1(const void* __restrict__ vbq, const block_q8_1* __restrict__ bq8_1, const int& iqs) { - const block_iq2_xs* bq2 = (const block_iq2_xs*)vbq; - - const int ib32 = iqs; - const uint16_t* q2 = bq2->qs + 4 * ib32; - const int8_t* q8 = bq8_1[ib32].qs; - const uint8_t ls1 = bq2->scales[ib32] & 0xf; - const uint8_t ls2 = bq2->scales[ib32] >> 4; - int sumi1 = 0; - for (int l = 0; l < 2; ++l) { - const uint8_t* grid = (const uint8_t*)(iq2xs_grid + (q2[l] & 511)); - const uint8_t signs = ksigns_iq2xs[q2[l] >> 9]; - for (int j = 0; j < 8; ++j) { - sumi1 += q8[j] * grid[j] * (signs & kmask_iq2xs[j] ? -1 : 1); - } - q8 += 8; - } - int sumi2 = 0; - for (int l = 2; l < 4; ++l) { - const uint8_t* grid = (const uint8_t*)(iq2xs_grid + (q2[l] & 511)); - const uint8_t signs = ksigns_iq2xs[q2[l] >> 9]; - for (int j = 0; j < 8; ++j) { - sumi2 += q8[j] * grid[j] * (signs & kmask_iq2xs[j] ? -1 : 1); - } - q8 += 8; - } - const float d = __half2float(bq2->d) * __half2float(bq8_1[ib32].ds.x) * 0.25f; - return d * ((0.5f + ls1) * sumi1 + (0.5f + ls2) * sumi2); -} - -static __device__ __forceinline__ float -vec_dot_iq2_s_q8_1(const void* __restrict__ vbq, const block_q8_1* __restrict__ bq8_1, const int& iqs) { -#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 610 || defined USE_ROCM || defined USE_MUSA - const block_iq2_s* bq2 = (const block_iq2_s*)vbq; - - const int ib32 = iqs; - const int8_t* q8 = bq8_1[ib32].qs; - const uint8_t* signs = bq2->qs + QK_K / 8 + 4 * ib32; - const uint8_t ls1 = bq2->scales[ib32] & 0xf; - const uint8_t ls2 = bq2->scales[ib32] >> 4; - int sumi1 = 0; - for (int l = 0; l < 2; ++l) { - const uint32_t* grid = - (const uint32_t*)(iq2s_grid + (bq2->qs[4 * ib32 + l] | ((bq2->qh[ib32] << (8 - 2 * l)) & 0x300))); - const uint32_t signs0 = __vcmpeq4(((signs[l] & 0xf) * 0x01010101) & 0x08040201, 0x08040201); - const uint32_t signs1 = __vcmpeq4(((signs[l] >> 4) * 0x01010101) & 0x08040201, 0x08040201); - const int grid_l = __vsub4(grid[0] ^ signs0, signs0); - const int grid_h = __vsub4(grid[1] ^ signs1, signs1); - sumi1 = __dp4a(grid_l, *((const int*)q8 + 0), sumi1); - sumi1 = __dp4a(grid_h, *((const int*)q8 + 1), sumi1); - q8 += 8; - } - int sumi2 = 0; - for (int l = 2; l < 4; ++l) { - const uint32_t* grid = - (const uint32_t*)(iq2s_grid + (bq2->qs[4 * ib32 + l] | ((bq2->qh[ib32] << (8 - 2 * l)) & 0x300))); - const uint32_t signs0 = __vcmpeq4(((signs[l] & 0xf) * 0x01010101) & 0x08040201, 0x08040201); - const uint32_t signs1 = __vcmpeq4(((signs[l] >> 4) * 0x01010101) & 0x08040201, 0x08040201); - const int grid_l = __vsub4(grid[0] ^ signs0, signs0); - const int grid_h = __vsub4(grid[1] ^ signs1, signs1); - sumi2 = __dp4a(grid_l, *((const int*)q8 + 0), sumi2); - sumi2 = __dp4a(grid_h, *((const int*)q8 + 1), sumi2); - q8 += 8; - } - const float d = __half2float(bq2->d) * __low2float(bq8_1[ib32].ds) * 0.25f; - return d * ((0.5f + ls1) * sumi1 + (0.5f + ls2) * sumi2); -#endif -} - -static __device__ __forceinline__ float -vec_dot_iq3_xxs_q8_1(const void* __restrict__ vbq, const block_q8_1* __restrict__ bq8_1, const int& iqs) { -#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 610 || defined USE_ROCM || defined USE_MUSA - const block_iq3_xxs* bq2 = (const block_iq3_xxs*)vbq; - - const int ib32 = iqs; - const uint8_t* q3 = bq2->qs + 8 * ib32; - const uint16_t* gas = (const uint16_t*)(bq2->qs + QK_K / 4) + 2 * ib32; - const int8_t* q8 = bq8_1[ib32].qs; - uint32_t aux32 = gas[0] | (gas[1] << 16); - int sumi = 0; - for (int l = 0; l < 4; ++l) { - const uint32_t* grid1 = iq3xxs_grid + q3[2 * l + 0]; - const uint32_t* grid2 = iq3xxs_grid + q3[2 * l + 1]; - const uint32_t* signs = (const uint32_t*)(ksigns64 + (aux32 & 127)); - const int grid_l = __vsub4(grid1[0] ^ signs[0], signs[0]); - const int grid_h = __vsub4(grid2[0] ^ signs[1], signs[1]); - sumi = __dp4a(grid_l, *((int*)q8 + 0), sumi); - sumi = __dp4a(grid_h, *((int*)q8 + 1), sumi); - q8 += 8; - aux32 >>= 7; - } - const float d = __half2float(bq2->d) * (0.5f + aux32) * __low2float(bq8_1[ib32].ds) * 0.5f; - return d * sumi; -#endif -} - -static __device__ __forceinline__ float -vec_dot_iq3_s_q8_1(const void* __restrict__ vbq, const block_q8_1* __restrict__ bq8_1, const int& iqs) { -#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 610 || defined USE_ROCM || defined USE_MUSA - const block_iq3_s* bq2 = (const block_iq3_s*)vbq; - - const int ib32 = iqs; - const uint8_t* qs = bq2->qs + 8 * ib32; - const int8_t* q8 = bq8_1[ib32].qs; - int sumi = 0; - for (int l = 0; l < 4; ++l) { - const uint32_t* grid1 = iq3xs_grid + (qs[2 * l + 0] | ((bq2->qh[ib32] << (8 - 2 * l)) & 256)); - const uint32_t* grid2 = iq3xs_grid + (qs[2 * l + 1] | ((bq2->qh[ib32] << (7 - 2 * l)) & 256)); - uint32_t signs0 = __vcmpeq4(((bq2->signs[4 * ib32 + l] & 0xf) * 0x01010101) & 0x08040201, 0x08040201); - uint32_t signs1 = __vcmpeq4(((bq2->signs[4 * ib32 + l] >> 4) * 0x01010101) & 0x08040201, 0x08040201); - const int grid_l = __vsub4(grid1[0] ^ signs0, signs0); - const int grid_h = __vsub4(grid2[0] ^ signs1, signs1); - sumi = __dp4a(grid_l, *((int*)q8 + 0), sumi); - sumi = __dp4a(grid_h, *((int*)q8 + 1), sumi); - q8 += 8; - } - const float d = __half2float(bq2->d) * (0.5f + ((bq2->scales[ib32 / 2] >> 4 * (ib32 % 2)) & 0xf)) * - __low2float(bq8_1[ib32].ds) * 0.5f; - return d * sumi; -#endif -} - -static __device__ __forceinline__ float -vec_dot_iq1_s_q8_1(const void* __restrict__ vbq, const block_q8_1* __restrict__ bq8_1, const int& iqs) { -#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 610 || defined USE_ROCM || defined USE_MUSA - const block_iq1_s* bq1 = (const block_iq1_s*)vbq; - - const int qs_packed = get_int_b2(bq1->qs, iqs); - const uint8_t* qs = (const uint8_t*)&qs_packed; - - const int qh = bq1->qh[iqs]; - - int sumi = 0; -#pragma unroll - for (int l0 = 0; l0 < 8; l0 += 2) { - const int grid = iq1s_grid_gpu[qs[l0 / 2] | (((qh >> 3 * (l0 / 2)) & 0x07) << 8)]; - - const int grid0 = (grid >> 0) & 0x0F0F0F0F; - const int grid1 = (grid >> 4) & 0x0F0F0F0F; - - const int u0 = get_int_b4(bq8_1[iqs].qs, l0 + 0); - const int u1 = get_int_b4(bq8_1[iqs].qs, l0 + 1); - - sumi = __dp4a(grid0, u0, sumi); - sumi = __dp4a(grid1, u1, sumi); - } - - const float d1q = __half2float(bq1->d) * (((qh >> 11) & 0x0E) + 1); - const float delta = -1.0f + IQ1S_DELTA - (qh & 0x8000) * (2.0f * IQ1S_DELTA / 0x8000); - const float2 ds = __half22float2(bq8_1[iqs].ds); - return d1q * (ds.x * sumi + ds.y * delta); -#endif -} - -static __device__ __forceinline__ float -vec_dot_iq1_m_q8_1(const void* __restrict__ vbq, const block_q8_1* __restrict__ bq8_1, const int& iqs) { -#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 610 || defined USE_ROCM || defined USE_MUSA - - const block_iq1_m* bq1 = (const block_iq1_m*)vbq; - - const int qs_packed = get_int_b4(bq1->qs, iqs); - const uint8_t* qs = (const uint8_t*)&qs_packed; - - int sumi[2] = {0}; - float sumf[2] = {0.0f}; -#pragma unroll - for (int l0 = 0; l0 < 8; l0 += 2) { - const int qhl = bq1->qh[2 * iqs + l0 / 4] >> (4 * ((l0 / 2) % 2)); - - const int grid = iq1s_grid_gpu[qs[l0 / 2] | ((qhl & 0x07) << 8)]; - - const int grid0 = (grid >> 0) & 0x0F0F0F0F; - const int grid1 = (grid >> 4) & 0x0F0F0F0F; - - const int u0 = get_int_b4(bq8_1[iqs].qs, l0 + 0); - const int u1 = get_int_b4(bq8_1[iqs].qs, l0 + 1); - - sumi[l0 / 4] = __dp4a(grid0, u0, sumi[l0 / 4]); - sumi[l0 / 4] = __dp4a(grid1, u1, sumi[l0 / 4]); - - const float delta = -1.0f + IQ1M_DELTA - (qhl & 0x08) * (2.0f * IQ1M_DELTA / 0x08); - int sumy = 0; - sumy = __dp4a(u0, 0x01010101, sumy); - sumy = __dp4a(u1, 0x01010101, sumy); - sumf[l0 / 4] += delta * sumy; - } - - const uint16_t* sc = (const uint16_t*)bq1->scales; - - iq1m_scale_t scale; - scale.u16 = (sc[0] >> 12) | ((sc[1] >> 8) & 0x00F0) | ((sc[2] >> 4) & 0x0F00) | (sc[3] & 0xF000); - const float d = __half2float(scale.f16) * __low2float(bq8_1[iqs].ds); - - const int tmp = sc[iqs / 2] >> (6 * (iqs % 2)); - const int sc0 = 2 * ((tmp >> 0) & 0x07) + 1; - const int sc1 = 2 * ((tmp >> 3) & 0x07) + 1; - return d * ((sumi[0] + sumf[0]) * sc0 + (sumi[1] + sumf[1]) * sc1); -#endif -} - -static __device__ __forceinline__ void -get_int_from_table_16(const uint32_t& q4, const uint8_t* values, int& val1, int& val2) { - uint32_t aux32; - const uint8_t* q8 = (const uint8_t*)&aux32; - aux32 = q4 & 0x0f0f0f0f; - uint16_t v1 = values[q8[0]] | (values[q8[1]] << 8); - uint16_t v2 = values[q8[2]] | (values[q8[3]] << 8); - val1 = v1 | (v2 << 16); - aux32 = (q4 >> 4) & 0x0f0f0f0f; - v1 = values[q8[0]] | (values[q8[1]] << 8); - v2 = values[q8[2]] | (values[q8[3]] << 8); - val2 = v1 | (v2 << 16); -} - -static __device__ __forceinline__ float -vec_dot_iq4_nl_q8_1(const void* __restrict__ vbq, const block_q8_1* __restrict__ bq8_1, const int& iqs) { -#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 610 || defined USE_ROCM || defined USE_MUSA - - const block_iq4_nl* bq = (const block_iq4_nl*)vbq; - - const uint16_t* q4 = (const uint16_t*)bq->qs + 2 * iqs; - const int32_t* q8 = (const int32_t*)bq8_1->qs + iqs; - - const uint8_t* values = (const uint8_t*)kvalues_iq4nl; - - int v1, v2; - int sumi1 = 0, sumi2 = 0; - for (int l = 0; l < VDR_Q4_0_Q8_1_MMVQ; ++l) { - const uint32_t aux = q4[2 * l] | (q4[2 * l + 1] << 16); - get_int_from_table_16(aux, values, v1, v2); - sumi1 = __dp4a(v1, q8[l + 0], sumi1); - sumi2 = __dp4a(v2, q8[l + 4], sumi2); - } - const float d = __half2float(bq->d) * __low2float(bq8_1->ds); - return d * (sumi1 + sumi2); -#endif -} - -static __device__ __forceinline__ float -vec_dot_iq4_xs_q8_1(const void* __restrict__ vbq, const block_q8_1* __restrict__ bq8_1, const int& iqs) { -#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 610 || defined USE_ROCM || defined USE_MUSA - const block_iq4_xs* bq4 = (const block_iq4_xs*)vbq; - const uint8_t* values = (const uint8_t*)kvalues_iq4nl; - - // iqs is 0...7 - const int ib32 = iqs; - const int32_t* q8 = (const int*)bq8_1[ib32].qs; - const uint32_t* q4 = (const uint32_t*)bq4->qs + 4 * ib32; - const int8_t ls = ((bq4->scales_l[ib32 / 2] >> 4 * (ib32 % 2)) & 0xf) | (((bq4->scales_h >> 2 * ib32) & 3) << 4); - const float d = __half2float(bq4->d) * (ls - 32) * __low2float(bq8_1[ib32].ds); - int v1, v2; - int sumi1 = 0, sumi2 = 0; - for (int j = 0; j < 4; ++j) { - get_int_from_table_16(q4[j], values, v1, v2); - sumi1 = __dp4a(v1, q8[j + 0], sumi1); - sumi2 = __dp4a(v2, q8[j + 4], sumi2); - } - return d * (sumi1 + sumi2); -#endif -} diff --git a/python/freetoken/kernel/csrc/include/freetoken/hip_compat.h b/python/freetoken/kernel/csrc/include/freetoken/hip_compat.h new file mode 100644 index 0000000000..8abcd1a098 --- /dev/null +++ b/python/freetoken/kernel/csrc/include/freetoken/hip_compat.h @@ -0,0 +1,174 @@ +#pragma once + +// HIP compatibility shim: maps CUDA runtime API names to HIP equivalents so +// the same C++ source compiles under both nvcc and hipcc. Include this instead +// of directly when the file needs the runtime API. +// +// On NVIDIA platforms the CUDA headers are included as-is and every macro below +// resolves to the original CUDA symbol, so there is zero overhead. +// +// Supported ROCm targets: +// gfx1100 — RX 7900 XTX / XT +// gfx1101 — RX 7900 GRE +// gfx1102 — RX 7700 / XT +// gfx1103 — RX 7600 / XT +// gfx1200 — RX 9060 family +// gfx1201 — RX 9070 family / Radeon AI PRO R9700 + +#if defined(__HIP_PLATFORM_AMD__) || defined(USE_ROCM) + +#define FREETOKEN_USE_ROCM 1 + +// --- HIP runtime headers --- +#include +#include + +// --- API name mapping (CUDA -> HIP) --- +// HIP already defines most cuda* names as macros that expand to hip* equivalents +// via hip_runtime.h, but a few are missing or differ in signature. Define them +// here so call-sites stay unchanged. + +#ifndef cudaSuccess +#define cudaSuccess hipSuccess +#endif + +#ifndef cudaError_t +#define cudaError_t hipError_t +#endif + +#ifndef cudaGetErrorString +#define cudaGetErrorString hipGetErrorString +#endif + +#ifndef cudaGetLastError +#define cudaGetLastError hipGetLastError +#endif + +#ifndef cudaMalloc +#define cudaMalloc hipMalloc +#endif + +#ifndef cudaFree +#define cudaFree hipFree +#endif + +#ifndef cudaMemset +#define cudaMemset hipMemset +#endif + +#ifndef cudaMallocHost +#define cudaMallocHost hipHostMalloc +#endif + +#ifndef cudaFreeHost +#define cudaFreeHost hipHostFree +#endif + +#ifndef cudaHostAlloc +#define cudaHostAlloc hipHostMalloc +#endif + +#ifndef cudaHostRegister +#define cudaHostRegister hipHostRegister +#endif + +#ifndef cudaHostRegisterPortable +#define cudaHostRegisterPortable hipHostRegisterPortable +#endif + +#ifndef cudaHostRegisterMapped +#define cudaHostRegisterMapped hipHostRegisterMapped +#endif + +#ifndef cudaHostAllocPortable +#define cudaHostAllocPortable hipHostMallocPortable +#endif + +#ifndef cudaHostAllocMapped +#define cudaHostAllocMapped hipHostMallocMapped +#endif + +#ifndef cudaHostGetDevicePointer +#define cudaHostGetDevicePointer hipHostGetDevicePointer +#endif + +#ifndef cudaGetDevice +#define cudaGetDevice hipGetDevice +#endif + +#ifndef cudaDriverGetVersion +#define cudaDriverGetVersion hipDriverGetVersion +#endif + +#ifndef cudaDeviceGetAttribute +#define cudaDeviceGetAttribute hipDeviceGetAttribute +#endif + +#ifndef cudaDevAttrUnifiedAddressing +#define cudaDevAttrUnifiedAddressing hipDeviceAttributeUnifiedAddressing +#endif + +#ifndef cudaDevAttrCanUseHostPointerForRegisteredMem +// HIP does not expose this attribute; assume UVA identity on ROCm (true on Linux). +// TODO(ROCm): re-enable proper UVA query if HIP adds this attribute. +#define cudaDevAttrCanUseHostPointerForRegisteredMem hipDeviceAttributeUnifiedAddressing +#endif + +#ifndef cudaFuncSetAttribute +#define cudaFuncSetAttribute hipFuncSetAttribute +#endif + +#ifndef cudaFuncAttributeMaxDynamicSharedMemorySize +#define cudaFuncAttributeMaxDynamicSharedMemorySize hipFuncAttributeMaxDynamicSharedMemorySize +#endif + +#ifndef cudaLaunchKernelEx +// ROCm 7 exposes the CUDA-compatible extended launch configuration through HIP. +#define cudaLaunchKernelEx hipLaunchKernelEx +#endif + +#ifndef cudaLaunchConfig_t +#define cudaLaunchConfig_t hipLaunchConfig_t +#endif + +#ifndef cudaLaunchAttribute +#define cudaLaunchAttribute hipLaunchAttribute +#endif + +#ifndef cudaLaunchAttributeProgrammaticStreamSerialization +// PDL (Programmatic Dependent Launch) is NVIDIA-specific. +// TODO(ROCm): PDL has no ROCm equivalent — disabled, may affect overlap scheduling latency. +#define cudaLaunchAttributeProgrammaticStreamSerialization 0 +#endif + +#ifndef cudaStream_t +#define cudaStream_t hipStream_t +#endif + +#ifndef cudaStreamSynchronize +#define cudaStreamSynchronize hipStreamSynchronize +#endif + +#ifndef cudaLaunchHostFunc +#define cudaLaunchHostFunc hipLaunchHostFunc +#endif + +#ifndef CUDART_CB +#define CUDART_CB +#endif + +#ifndef __grid_constant__ +#define __grid_constant__ +#endif + +#ifndef dim3 +// HIP already provides dim3; this is a no-op guard. +#endif + +#else // NVIDIA CUDA path + +#define FREETOKEN_USE_ROCM 0 + +#include + +#endif diff --git a/python/freetoken/kernel/csrc/include/freetoken/tensor.h b/python/freetoken/kernel/csrc/include/freetoken/tensor.h index 9b591a87dc..3f711a5f7b 100644 --- a/python/freetoken/kernel/csrc/include/freetoken/tensor.h +++ b/python/freetoken/kernel/csrc/include/freetoken/tensor.h @@ -34,7 +34,7 @@ struct DeviceRef; inline constexpr auto kAnyDeviceID = -1; inline constexpr auto kAnySize = static_cast(-1); -inline constexpr auto kNullSize = static_cast(0); +inline constexpr auto kNullSize = static_cast(-1); inline constexpr auto kNullDType = static_cast(18u); inline constexpr auto kNullDevice = static_cast(-1); @@ -147,7 +147,7 @@ struct SymbolicSize { RuntimeCheck(!this->has_value(), "Size value already set"); m_value = value; } - auto has_value() const -> bool { return m_value != 0; } + auto has_value() const -> bool { return m_value != details::kNullSize; } auto get_value() const -> std::optional { return this->has_value() ? std::optional{m_value} : std::nullopt; } @@ -385,6 +385,10 @@ struct TensorMatcher { auto with_dtype(DTypeRef &&dtype) && -> TensorMatcher && { m_init_dtype(); m_dtype.rebind(*dtype); + // no options given: keep the ones already set on the shared symbol + if constexpr (sizeof...(Ts) > 0) { + m_dtype->set_options(); + } return std::move(*this); } @@ -400,6 +404,9 @@ struct TensorMatcher { auto with_device(DeviceRef &&device) && -> TensorMatcher && { m_init_device(); m_device.rebind(*device); + if constexpr (sizeof...(Codes) > 0) { + m_device->set_options(); + } return std::move(*this); } diff --git a/python/freetoken/kernel/csrc/include/freetoken/utils.cuh b/python/freetoken/kernel/csrc/include/freetoken/utils.cuh index 8e917832c1..091642a8fd 100644 --- a/python/freetoken/kernel/csrc/include/freetoken/utils.cuh +++ b/python/freetoken/kernel/csrc/include/freetoken/utils.cuh @@ -1,5 +1,6 @@ #pragma once +#include #include #include @@ -42,16 +43,21 @@ __always_inline __device__ auto offset(const T *ptr, U... offset) -> const namespace PDL { +// griddepcontrol needs sm_90+; older archs and HIP keep normal stream ordering. template __always_inline __device__ void wait() { +#if !FREETOKEN_USE_ROCM && __CUDA_ARCH__ >= 900 if constexpr (kUsePDL) { asm volatile("griddepcontrol.wait;" ::: "memory"); } +#endif } template __always_inline __device__ void launch() { +#if !FREETOKEN_USE_ROCM && __CUDA_ARCH__ >= 900 if constexpr (kUsePDL) { asm volatile("griddepcontrol.launch_dependents;" :::); } +#endif } } // namespace PDL @@ -115,6 +121,10 @@ public: } auto with_attr(bool use_pdl) -> LaunchKernel & { +#if FREETOKEN_USE_ROCM + (void)use_pdl; + m_config.numAttrs = 0; +#else if (use_pdl) { m_attr_cache.id = ::cudaLaunchAttributeProgrammaticStreamSerialization; m_attr_cache.val.programmaticStreamSerializationAllowed = 1; @@ -123,6 +133,7 @@ public: } else { m_config.numAttrs = 0; } +#endif return *this; } @@ -138,7 +149,9 @@ private: return config; } cudaLaunchConfig_t m_config; +#if !FREETOKEN_USE_ROCM cudaLaunchAttribute m_attr_cache; +#endif }; } // namespace host diff --git a/python/freetoken/kernel/csrc/jit/fast_index_copy.cuh b/python/freetoken/kernel/csrc/jit/fast_index_copy.cuh index bb83c23ed2..e37e3d48ad 100644 --- a/python/freetoken/kernel/csrc/jit/fast_index_copy.cuh +++ b/python/freetoken/kernel/csrc/jit/fast_index_copy.cuh @@ -20,6 +20,13 @@ struct device_vec { namespace details { +#if FREETOKEN_USE_ROCM +// Use native Clang vector types so each __builtin_nontemporal_load/store +// maps to one vectorized 8B/16B global memory instruction. +using native_uint2 = uint32_t __attribute__((ext_vector_type(2))); +using native_uint4 = uint32_t __attribute__((ext_vector_type(4))); +#endif + template inline constexpr auto get_mem_package() { if constexpr (kUnit == 16) { @@ -34,40 +41,68 @@ inline constexpr auto get_mem_package() { } __always_inline __device__ auto load_nc(const uint1* __restrict__ src) -> uint1 { +#if FREETOKEN_USE_ROCM + return uint1{__builtin_nontemporal_load(&src->x)}; +#else uint32_t tmp; asm volatile("ld.global.L1::no_allocate.b32 %0,[%1];" : "=r"(tmp) : "l"(src)); return uint1{tmp}; +#endif } __always_inline __device__ auto load_nc(const uint2* __restrict__ src) -> uint2 { +#if FREETOKEN_USE_ROCM + const auto value = __builtin_nontemporal_load(reinterpret_cast(src)); + return __builtin_bit_cast(uint2, value); +#else uint32_t tmp0, tmp1; asm volatile("ld.global.L1::no_allocate.v2.b32 {%0,%1},[%2];" : "=r"(tmp0), "=r"(tmp1) : "l"(src)); return uint2{tmp0, tmp1}; +#endif } __always_inline __device__ auto load_nc(const uint4* __restrict__ src) -> uint4 { +#if FREETOKEN_USE_ROCM + const auto value = __builtin_nontemporal_load(reinterpret_cast(src)); + return __builtin_bit_cast(uint4, value); +#else uint32_t tmp0, tmp1, tmp2, tmp3; asm volatile("ld.global.L1::no_allocate.v4.b32 {%0,%1,%2,%3},[%4];" : "=r"(tmp0), "=r"(tmp1), "=r"(tmp2), "=r"(tmp3) : "l"(src)); return uint4{tmp0, tmp1, tmp2, tmp3}; +#endif } __always_inline __device__ void store_nc(uint1* __restrict__ dst, const uint1& value) { +#if FREETOKEN_USE_ROCM + __builtin_nontemporal_store(value.x, &dst->x); +#else uint32_t tmp = value.x; asm volatile("st.global.wt.b32 [%0],%1;" ::"l"(dst), "r"(tmp)); +#endif } __always_inline __device__ void store_nc(uint2* __restrict__ dst, const uint2& value) { +#if FREETOKEN_USE_ROCM + __builtin_nontemporal_store( + __builtin_bit_cast(native_uint2, value), reinterpret_cast(dst)); +#else uint32_t tmp0 = value.x; uint32_t tmp1 = value.y; asm volatile("st.global.wt.v2.b32 [%0],{%1,%2};" ::"l"(dst), "r"(tmp0), "r"(tmp1)); +#endif } __always_inline __device__ void store_nc(uint4* __restrict__ dst, const uint4& value) { +#if FREETOKEN_USE_ROCM + __builtin_nontemporal_store( + __builtin_bit_cast(native_uint4, value), reinterpret_cast(dst)); +#else uint32_t tmp0 = value.x; uint32_t tmp1 = value.y; uint32_t tmp2 = value.z; uint32_t tmp3 = value.w; asm volatile("st.global.wt.v4.b32 [%0],{%1,%2,%3,%4};" ::"l"(dst), "r"(tmp0), "r"(tmp1), "r"(tmp2), "r"(tmp3)); +#endif } __always_inline __device__ void wait_flag_clear(const int32_t* __restrict__ flag_ptr) { @@ -75,7 +110,7 @@ __always_inline __device__ void wait_flag_clear(const int32_t* __restrict__ flag auto* flag = reinterpret_cast(const_cast(flag_ptr)); uint32_t sleep_ns = 128; while (atomicAdd(flag, 0) > 0) { -#if __CUDA_ARCH__ >= 700 +#if !FREETOKEN_USE_ROCM && __CUDA_ARCH__ >= 700 __nanosleep(sleep_ns); #endif sleep_ns = sleep_ns < 2048 ? (sleep_ns << 1) : 2048; @@ -147,7 +182,7 @@ inline bool host_ptr_identity() { } inline void* device_alias(void* ptr, DLDevice dev) { - if (dev.device_type == kDLCUDA || host_ptr_identity()) { + if (dev.device_type == kDLCUDA || dev.device_type == kDLROCM || host_ptr_identity()) { return ptr; } void* mapped = nullptr; @@ -269,7 +304,7 @@ inline auto get_sync_flag_ptr( auto flag_dtype = host::SymbolicDType{}; host::TensorMatcher({1}) .with_dtype(flag_dtype) - .with_device(device) + .with_device(device) .verify(sync_flag); return static_cast(sync_flag.data_ptr()); } @@ -344,17 +379,17 @@ struct FastIndexCopyKernel { TensorMatcher({-1, D}) .with_dtype(data_dtype) - .with_device() + .with_device() .verify(src); TensorMatcher({-1, D}) .with_dtype(data_dtype) - .with_device() + .with_device() .verify(dst); TensorMatcher({L}) .with_dtype(indices_dtype) - .with_device(device) + .with_device(device) .verify(src_indices) .verify(dst_indices); @@ -363,7 +398,7 @@ struct FastIndexCopyKernel { const auto num_indices_tensor = num_indices.value(); TensorMatcher({1}) .with_dtype(num_indices_dtype) - .with_device(device) + .with_device(device) .verify(num_indices_tensor); num_indices_data_ptr = static_cast(num_indices_tensor.data_ptr()); @@ -529,14 +564,14 @@ struct MultiIndexCopyKernel { auto indices_dtype = SymbolicDType{}; auto num_indices_dtype = SymbolicDType{}; - TensorMatcher({B}).with_dtype(ptr_dtype).with_device(device) + 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) + 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) + TensorMatcher({1}).with_dtype(num_indices_dtype).with_device(device) .verify(num_indices.value()); valid_length = static_cast(num_indices.value().data_ptr()); } @@ -560,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/csrc/jit/index.cu b/python/freetoken/kernel/csrc/jit/index.cu index ca0e1db26e..4832e0a460 100644 --- a/python/freetoken/kernel/csrc/jit/index.cu +++ b/python/freetoken/kernel/csrc/jit/index.cu @@ -35,9 +35,10 @@ __global__ __launch_bounds__(kNumThreads, kMaxOccupancy) void // using namespace device; constexpr auto kSize = kElementSize; constexpr auto kSizePerWarp = kSize / kNumSplits; - constexpr auto kWarpPerBlock = static_cast(kNumThreads / 32); + constexpr auto kWarpPerBlock = + static_cast(kNumThreads / kWarpThreads); - static_assert(kNumThreads % 32 == 0); + static_assert(kNumThreads % kWarpThreads == 0); static_assert(std::has_single_bit(kNumSplits)); static_assert(kElementSize % kNumSplits == 0); @@ -66,9 +67,10 @@ __global__ __launch_bounds__(kNumThreads, kMaxOccupancy) void // using namespace device; constexpr auto kSize = kElementSize; constexpr auto kSizePerWarp = kSize / kNumSplits; - constexpr auto kWarpPerBlock = static_cast(kNumThreads / 32); + constexpr auto kWarpPerBlock = + static_cast(kNumThreads / kWarpThreads); - static_assert(kNumThreads % 32 == 0); + static_assert(kNumThreads % kWarpThreads == 0); static_assert(std::has_single_bit(kNumSplits)); static_assert(kElementSize % kNumSplits == 0); @@ -114,15 +116,15 @@ struct IndexKernel { TensorMatcher({-1, D}) // .with_dtype(weights_dtype_) - .with_device(device_) + .with_device(device_) .verify(weights); TensorMatcher({L, D}) // .with_dtype(weights_dtype_) - .with_device(device_) + .with_device(device_) .verify(output); TensorMatcher({L}) // .with_dtype(indices_dtype_) - .with_device(device_) + .with_device(device_) .verify(indices); const auto device = device_.unwrap(); @@ -132,7 +134,12 @@ struct IndexKernel { RuntimeCheck(entry_size == element_size, "IndexKernel: element_size mismatch."); - constexpr auto kWarpPerBlock = num_threads / 32; + // Empty batches are valid no-ops; do not launch a zero-sized grid. + if (num_indices == 0) { + return; + } + + constexpr auto kWarpPerBlock = num_threads / device::kWarpThreads; const auto num_warps = num_splits * num_indices; const auto num_blocks = div_ceil(num_warps, kWarpPerBlock); const auto params = IndexKernelParams{ diff --git a/python/freetoken/kernel/csrc/jit/store.cu b/python/freetoken/kernel/csrc/jit/store.cu index 8d84d76ef1..18fd45346a 100644 --- a/python/freetoken/kernel/csrc/jit/store.cu +++ b/python/freetoken/kernel/csrc/jit/store.cu @@ -72,18 +72,18 @@ struct StoreKernel { TensorMatcher({-1, D}) // .with_strides({X, 1}) - .with_device(device_) + .with_device(device_) .with_dtype(dtype_) .verify(k_cache) .verify(v_cache); TensorMatcher({L, D}) // .with_strides({Y, 1}) - .with_device(device_) + .with_device(device_) .with_dtype(dtype_) .verify(k) .verify(v); TensorMatcher({L}) // - .with_device(device_) + .with_device(device_) .with_dtype(indices_dtype_) .verify(indices); @@ -96,6 +96,11 @@ struct StoreKernel { const auto kv_cache_stride = X.unwrap() * dtype_size; const auto kv_input_stride = Y.unwrap() * dtype_size; + // Empty batches are valid no-ops; do not launch a zero-sized grid. + if (length == 0) { + return; + } + const auto params = StoreKernelParams{ .k_cache = k_cache.data_ptr(), .v_cache = v_cache.data_ptr(), @@ -107,8 +112,8 @@ struct StoreKernel { .length = length, }; - constexpr auto kWarpPerBlock = num_threads / 32; - static_assert(num_threads % 32 == 0); + constexpr auto kWarpPerBlock = num_threads / device::kWarpThreads; + static_assert(num_threads % device::kWarpThreads == 0); const auto num_blocks = div_ceil(length, kWarpPerBlock); const auto kernel = use_int32 ? store_kv_cache -#include +#include #include namespace { +// A failed runtime call latches its status in the calling thread until someone +// reads it. TORCH_CHECK reports the error but does not clear it, so the next +// unrelated CUDA call -- torch's, in practice -- reports *this* failure instead +// of its own. Drain it before raising. The failures these calls can originate +// are non-sticky, so draining leaves the context usable; a sticky error latched +// elsewhere will simply re-latch on the next call, which is correct. +#define FT_CUDA_CHECK(expr, ...) \ + do { \ + const cudaError_t ft_err_ = (expr); \ + if (ft_err_ != cudaSuccess) { \ + [[maybe_unused]] const cudaError_t ft_drained_ = cudaGetLastError(); \ + TORCH_CHECK(false, __VA_ARGS__, cudaGetErrorString(ft_err_)); \ + } \ + } while (0) + void free_pinned(void *ptr) { - if (ptr != nullptr) { - cudaFreeHost(ptr); + if (ptr == nullptr) { + return; + } + // A from_blob deleter runs during GC, with no exception to attribute a failure + // to -- so this drains and never throws. Left latched it would be the worst + // case of the bug above: a stale status with no visible origin at all. + if (cudaFreeHost(ptr) != cudaSuccess) { + [[maybe_unused]] const cudaError_t drained = cudaGetLastError(); } } @@ -34,9 +55,8 @@ torch::Tensor create_pinned_tensor_like(torch::Tensor input) { const size_t alloc_nbytes = static_cast(nbytes == 0 ? 1 : nbytes); void *data_ptr = nullptr; - const cudaError_t alloc_err = cudaMallocHost(&data_ptr, alloc_nbytes); - TORCH_CHECK(alloc_err == cudaSuccess, - "cudaMallocHost failed: ", cudaGetErrorString(alloc_err)); + FT_CUDA_CHECK(cudaMallocHost(&data_ptr, alloc_nbytes), + "cudaMallocHost failed: "); auto options = input.options().device(torch::kCPU).pinned_memory(true); @@ -58,10 +78,9 @@ torch::Tensor alloc_pinned_tensor(std::vector sizes, // Portable + mapped: the offload gather kernel reads these banks straight // from host memory (zero-copy), which requires device-mapped pinned pages. void *data_ptr = nullptr; - const cudaError_t alloc_err = cudaHostAlloc( - &data_ptr, alloc_nbytes, cudaHostAllocPortable | cudaHostAllocMapped); - TORCH_CHECK(alloc_err == cudaSuccess, - "cudaHostAlloc failed: ", cudaGetErrorString(alloc_err)); + FT_CUDA_CHECK(cudaHostAlloc(&data_ptr, alloc_nbytes, + cudaHostAllocPortable | cudaHostAllocMapped), + "cudaHostAlloc failed: "); auto options = torch::TensorOptions() .dtype(dtype) @@ -76,37 +95,42 @@ torch::Tensor alloc_pinned_tensor(std::vector sizes, // device address). Zero-copy consumers resolve bank base addresses through these. bool host_ptr_identity() { int device = 0; - const cudaError_t err = cudaGetDevice(&device); - TORCH_CHECK(err == cudaSuccess, "cudaGetDevice failed: ", cudaGetErrorString(err)); + FT_CUDA_CHECK(cudaGetDevice(&device), "cudaGetDevice failed: "); + // An unqueryable attribute stays 0 and answers "no identity", which is the safe + // answer: device_ptr() then goes through host_device_ptr(), the real translation, + // correct on every platform. Only the latch is new here -- do not raise, or an + // attribute query that used to degrade gracefully becomes fatal on the offload path. int uva = 0, reg = 0; - cudaDeviceGetAttribute(&uva, cudaDevAttrUnifiedAddressing, device); - cudaDeviceGetAttribute(®, cudaDevAttrCanUseHostPointerForRegisteredMem, device); + if (cudaDeviceGetAttribute(&uva, cudaDevAttrUnifiedAddressing, device) != cudaSuccess) { + [[maybe_unused]] const cudaError_t drained = cudaGetLastError(); + uva = 0; + } + if (cudaDeviceGetAttribute(®, cudaDevAttrCanUseHostPointerForRegisteredMem, + device) != cudaSuccess) { + [[maybe_unused]] const cudaError_t drained = cudaGetLastError(); + reg = 0; + } return uva == 1 && reg == 1; } int64_t host_device_ptr(int64_t host_ptr) { void *dev_ptr = nullptr; - const cudaError_t err = - cudaHostGetDevicePointer(&dev_ptr, reinterpret_cast(host_ptr), 0); - TORCH_CHECK(err == cudaSuccess, - "cudaHostGetDevicePointer failed (host memory must be pinned+mapped): ", - cudaGetErrorString(err)); + FT_CUDA_CHECK( + cudaHostGetDevicePointer(&dev_ptr, reinterpret_cast(host_ptr), 0), + "cudaHostGetDevicePointer failed (host memory must be pinned+mapped): "); return reinterpret_cast(dev_ptr); } void host_register(int64_t addr, int64_t nbytes) { - const cudaError_t err = - cudaHostRegister(reinterpret_cast(addr), static_cast(nbytes), - cudaHostRegisterPortable | cudaHostRegisterMapped); - TORCH_CHECK(err == cudaSuccess, - "cudaHostRegister failed: ", cudaGetErrorString(err)); + FT_CUDA_CHECK(cudaHostRegister(reinterpret_cast(addr), + static_cast(nbytes), + cudaHostRegisterPortable | cudaHostRegisterMapped), + "cudaHostRegister failed: "); } int64_t driver_cuda_version() { int version = 0; // stays 0 when no driver is installed - const cudaError_t err = cudaDriverGetVersion(&version); - TORCH_CHECK(err == cudaSuccess, - "cudaDriverGetVersion failed: ", cudaGetErrorString(err)); + FT_CUDA_CHECK(cudaDriverGetVersion(&version), "cudaDriverGetVersion failed: "); return version; } diff --git a/python/freetoken/kernel/csrc/ple_store/ple_store_ext.cpp b/python/freetoken/kernel/csrc/row_store/row_store_ext.cpp similarity index 83% rename from python/freetoken/kernel/csrc/ple_store/ple_store_ext.cpp rename to python/freetoken/kernel/csrc/row_store/row_store_ext.cpp index 2d54d692d6..e33b18a8e0 100644 --- a/python/freetoken/kernel/csrc/ple_store/ple_store_ext.cpp +++ b/python/freetoken/kernel/csrc/row_store/row_store_ext.cpp @@ -1,7 +1,12 @@ -// Disk-backed PLE row store: rows read straight from the checkpoint's fp8 shard tensors -// through an extent table (PleRowSource in ple_ssd.py). Engine-thread only, no locks. -// Duplicate rows in one fill dedup into ONE batched read round; no RAM cache and no -// per-sequence state. Hash reference: tests/models/qwen4_exp/test_ple_disk.py. +// Disk-backed row store: fixed-width rows read straight from checkpoint shard tensors +// through an extent table (row i of extent e at extent_base[e] + i * row_stride). +// Engine-thread only, no locks. Duplicate rows in one fill dedup into ONE batched read +// round; no RAM cache and no per-sequence state. +// +// RowStore -- hash-agnostic: stage_rows(row_ids) + flush(); callers supply row ids. +// PleStore -- RowStore plus the Qwen3.8-Flash-Next PLE n-gram hash: stage(tokens). +// Hash reference: tests/models/qwen4_exp/test_ple_disk.py. +// // Platform seams: TableFile (O_DIRECT+pread; Win: NO_BUFFERING), BatchReader (io_uring, // pread-pool fallback = the portable shape), cumemop_* (dlopen libcuda; Win: nvcuda). @@ -183,7 +188,9 @@ class ThreadPoolBatchReader final : public BatchReader { ThreadPoolBatchReader() { // these threads block on I/O, not compute, so the core count is only a default unsigned n = std::max(1u, std::min(kReaderThreads, std::thread::hardware_concurrency())); - if (const char *env = std::getenv("FREETOKEN_PLE_READER_THREADS")) { + const char *env = std::getenv("FREETOKEN_ROW_STORE_READER_THREADS"); + if (env == nullptr) env = std::getenv("FREETOKEN_PLE_READER_THREADS"); // the PLE-era name + if (env != nullptr) { // kBatchEntries is the submit depth, so threads past it never get a read const int v = std::atoi(env); if (v > 0) n = std::min((unsigned)v, kBatchEntries); @@ -432,29 +439,24 @@ int64_t pos_mod(int64_t v, int64_t m) { return r < 0 ? r + m : r; } -class PleStore { +class RowStore { struct Extent { const TableFile *file; int64_t base; }; public: - PleStore(std::vector paths, std::vector extent_file, + RowStore(std::vector paths, std::vector extent_file, std::vector extent_base, int64_t rows_per_extent, int64_t row_bytes, - int64_t row_stride, std::vector multipliers, std::vector head_vocab_sizes, - std::vector head_offsets, int64_t eos_token_id, bool use_io_uring) - : row_bytes_(row_bytes), - row_stride_(row_stride), - rows_per_extent_(rows_per_extent), - mult_(std::move(multipliers)), - sizes_(std::move(head_vocab_sizes)), - offsets_(std::move(head_offsets)), - eos_(eos_token_id) { - if (mult_.size() != 3 || sizes_.size() != offsets_.size() || sizes_.empty()) - throw std::runtime_error("PLE hash geometry: want 3 multipliers and equal-length head tables"); + int64_t row_stride, bool use_io_uring) + : row_bytes_(row_bytes), row_stride_(row_stride), rows_per_extent_(rows_per_extent) { + if (row_bytes_ <= 0 || row_stride_ < row_bytes_ || rows_per_extent_ <= 0) + throw std::runtime_error("row store geometry: want row_bytes > 0, row_stride >= row_bytes, rows_per_extent > 0"); if (row_bytes_ > kPage) - throw std::runtime_error("PLE row_bytes " + std::to_string(row_bytes_) + + throw std::runtime_error("row_bytes " + std::to_string(row_bytes_) + " exceeds a page; bounce slots assume one-page rows"); + if (extent_file.size() != extent_base.size() || extent_file.empty()) + throw std::runtime_error("row store geometry: extent_file and extent_base must be equal-length and non-empty"); for (const std::string &p : paths) files_.push_back(std::make_unique(p)); const int64_t extent_bytes = (rows_per_extent_ - 1) * row_stride_ + row_bytes_; @@ -472,36 +474,33 @@ class PleStore { } // reader first: a still-running read must not land in freed bounce memory - ~PleStore() { + virtual ~RowStore() { reader_.reset(); free(bounce_); } - PleStore(const PleStore &) = delete; - PleStore &operator=(const PleStore &) = delete; + RowStore(const RowStore &) = delete; + RowStore &operator=(const RowStore &) = delete; - // Row ids for the token at w[2] with context (w[0], w[1]); mirrors NGramEmbedding.row_ids incl. the eos barrier. - void hash_rows(const int64_t *w, int64_t *rows) { - const int64_t prev1 = w[1]; - const int64_t prev2 = prev1 == eos_ ? eos_ : w[0]; - const int64_t bigram = wrap_mul(w[2], mult_[0]) ^ wrap_mul(prev1, mult_[1]); - const int64_t trigram = bigram ^ wrap_mul(prev2, mult_[2]); - const size_t half = sizes_.size() / 2; - for (size_t h = 0; h < sizes_.size(); h++) - rows[h] = pos_mod(h < half ? bigram : trigram, sizes_[h]) + offsets_[h]; - } + int64_t row_bytes() const { return row_bytes_; } + int64_t total_rows() const { return (int64_t)extents_.size() * rows_per_extent_; } - // Hash and queue one run: tokens_addr holds n+2 ids, the leading two are context. No I/O until flush(). - void stage(uintptr_t tokens_addr, int64_t n, uintptr_t staging_addr) { - const int64_t *tokens = reinterpret_cast(tokens_addr); + // Queue n row ids (int64 at rows_addr); row i lands at staging_addr + i * dst_stride + // (dst_stride = 0 packs rows back to back). Out-of-range ids throw before any I/O is + // issued. No I/O until flush(). + void stage_rows(uintptr_t rows_addr, int64_t n, uintptr_t staging_addr, int64_t dst_stride) { + const int64_t *rows = reinterpret_cast(rows_addr); uint8_t *staging = reinterpret_cast(staging_addr); - const size_t heads = sizes_.size(); - std::vector rows(heads); + if (dst_stride == 0) dst_stride = row_bytes_; + if (dst_stride < row_bytes_) + throw std::runtime_error("stage_rows: dst_stride below row_bytes would overlap rows"); + const int64_t total = total_rows(); for (int64_t i = 0; i < n; i++) { - hash_rows(tokens + i, rows.data()); - for (size_t h = 0; h < heads; h++) - request_row(rows[h], staging + ((size_t)i * heads + h) * row_bytes_); + if (rows[i] < 0 || rows[i] >= total) + throw std::out_of_range("stage_rows: row id " + std::to_string(rows[i]) + " outside [0, " + + std::to_string(total) + ")"); } + for (int64_t i = 0; i < n; i++) request_row(rows[i], staging + (size_t)i * (size_t)dst_stride); } // One batched disk round for everything staged; signals even when nothing was. @@ -519,15 +518,7 @@ class PleStore { std::to_string(files_.size()) + " files"; } - private: - struct Pending { - const TableFile *file; - int64_t read_off; - int64_t read_len; - int64_t row_off; // row payload start inside the read buffer - std::vector dsts; - }; - + protected: // Queue dst on this fill's pending batch; duplicate rows fan out from one read. void request_row(int64_t row_id, uint8_t *dst) { auto pit = pending_index_.find(row_id); @@ -548,11 +539,20 @@ class PleStore { pending_.push_back(std::move(p)); } + private: + struct Pending { + const TableFile *file; + int64_t read_off; + int64_t read_len; + int64_t row_off; // row payload start inside the read buffer + std::vector dsts; + }; + // Read every pending row in reader-capacity batches and fan out the copies. void flush_pending() { if (pending_.empty()) return; struct Cleanup { - PleStore *s; + RowStore *s; ~Cleanup() { s->pending_.clear(); s->pending_index_.clear(); @@ -586,8 +586,6 @@ class PleStore { } int64_t row_bytes_, row_stride_, rows_per_extent_; - std::vector mult_, sizes_, offsets_; - int64_t eos_; std::vector> files_; std::vector extents_; std::unique_ptr reader_; @@ -597,10 +595,71 @@ class PleStore { std::unordered_map pending_index_; }; +// The Qwen3.8-Flash-Next PLE table: RowStore plus the checkpoint's n-gram hash, so a fill +// takes raw token runs and never crosses the Python boundary per row. +class PleStore final : public RowStore { + public: + PleStore(std::vector paths, std::vector extent_file, + std::vector extent_base, int64_t rows_per_extent, int64_t row_bytes, + int64_t row_stride, std::vector multipliers, std::vector head_vocab_sizes, + std::vector head_offsets, int64_t eos_token_id, bool use_io_uring) + : RowStore(std::move(paths), std::move(extent_file), std::move(extent_base), rows_per_extent, + row_bytes, row_stride, use_io_uring), + mult_(std::move(multipliers)), + sizes_(std::move(head_vocab_sizes)), + offsets_(std::move(head_offsets)), + eos_(eos_token_id) { + if (mult_.size() != 3 || sizes_.size() != offsets_.size() || sizes_.empty()) + throw std::runtime_error("PLE hash geometry: want 3 multipliers and equal-length head tables"); + } + + // Row ids for the token at w[2] with context (w[0], w[1]); mirrors NGramEmbedding.row_ids incl. the eos barrier. + void hash_rows(const int64_t *w, int64_t *rows) { + const int64_t prev1 = w[1]; + const int64_t prev2 = prev1 == eos_ ? eos_ : w[0]; + const int64_t bigram = wrap_mul(w[2], mult_[0]) ^ wrap_mul(prev1, mult_[1]); + const int64_t trigram = bigram ^ wrap_mul(prev2, mult_[2]); + const size_t half = sizes_.size() / 2; + for (size_t h = 0; h < sizes_.size(); h++) + rows[h] = pos_mod(h < half ? bigram : trigram, sizes_[h]) + offsets_[h]; + } + + // Hash and queue one run: tokens_addr holds n+2 ids, the leading two are context. No I/O until flush(). + void stage(uintptr_t tokens_addr, int64_t n, uintptr_t staging_addr) { + const int64_t *tokens = reinterpret_cast(tokens_addr); + uint8_t *staging = reinterpret_cast(staging_addr); + const size_t heads = sizes_.size(); + std::vector rows(heads); + for (int64_t i = 0; i < n; i++) { + hash_rows(tokens + i, rows.data()); + for (size_t h = 0; h < heads; h++) + request_row(rows[h], staging + ((size_t)i * heads + h) * (size_t)row_bytes()); + } + } + + private: + std::vector mult_, sizes_, offsets_; + int64_t eos_; +}; + } // namespace PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { - py::class_(m, "PleStore") + py::class_(m, "RowStore") + .def(py::init, std::vector, std::vector, + int64_t, int64_t, int64_t, bool>(), + py::arg("paths"), py::arg("extent_file"), py::arg("extent_base"), + py::arg("rows_per_extent"), py::arg("row_bytes"), py::arg("row_stride"), + py::arg("use_io_uring") = true) + .def("stage_rows", &RowStore::stage_rows, py::arg("rows_addr"), py::arg("n"), + py::arg("staging_addr"), py::arg("dst_stride") = 0, + py::call_guard()) + .def("flush", &RowStore::flush, py::arg("signal_addr") = 0, + py::call_guard()) + .def("io_backend", &RowStore::io_backend) + .def_property_readonly("row_bytes", &RowStore::row_bytes) + .def_property_readonly("total_rows", &RowStore::total_rows); + py::class_(m, "PleStore") .def(py::init, std::vector, std::vector, int64_t, int64_t, int64_t, std::vector, std::vector, std::vector, int64_t, bool>(), @@ -610,10 +669,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { py::arg("head_vocab_sizes"), py::arg("head_offsets"), py::arg("eos_token_id"), py::arg("use_io_uring") = true) .def("stage", &PleStore::stage, py::arg("tokens_addr"), py::arg("n"), - py::arg("staging_addr"), py::call_guard()) - .def("flush", &PleStore::flush, py::arg("signal_addr") = 0, - py::call_guard()) - .def("io_backend", &PleStore::io_backend); + py::arg("staging_addr"), py::call_guard()); m.def("memop_write", &memop_write, py::arg("stream"), py::arg("addr"), py::arg("value")); m.def("memop_wait_geq", &memop_wait_geq, py::arg("stream"), py::arg("addr"), py::arg("value")); m.def("memop_wait_reset", &memop_wait_reset, py::arg("stream"), py::arg("flag_addr")); 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/ggml.py b/python/freetoken/kernel/ggml.py new file mode 100644 index 0000000000..7a2765c5c0 --- /dev/null +++ b/python/freetoken/kernel/ggml.py @@ -0,0 +1,168 @@ +"""llama.cpp's ggml-cuda kernels (vendored under ``csrc/ggml``) as a tvm-ffi module: prebuilt in the kernel cache, else compiled on first use. + +Weights stay in their packed ggml blocks. ``mul_mat`` / ``mul_mat_id`` pick MMVQ or MMQ the way llama.cpp's mul_mat / mul_mat_id do. +Activations are fp32, bf16 or fp16 and are quantized to q8_1 first; results come out in the activation dtype and are accumulated in fp32. +""" + +from __future__ import annotations + +import functools +import pathlib +from typing import TYPE_CHECKING + +import torch + +from .utils import _is_rocm, load_aot + +if TYPE_CHECKING: + from tvm_ffi import Module + +_ROOT = pathlib.Path(__file__).parent / "csrc" / "ggml" + + +def load_module(build_directory: str | None = None) -> Module: + cuda_dir = _ROOT / "ggml-cuda" + sources = [ + _ROOT / "freetoken" / "ggml_shim.cu", + _ROOT / "freetoken" / "ggml_ops.cu", + *(cuda_dir / f for f in ("mmvq.cu", "mmq.cu", "mmid.cu", "quantize.cu", "convert.cu", "getrows.cu")), + *sorted((cuda_dir / "template-instances").glob("mmq-instance-*.cu")), + ] + if _is_rocm(): + flags = ["-ffast-math", "-DGGML_USE_HIP"] + else: + # a multi-arch kernel-cache build is ~6x smaller with compressed SASS + flags = ["-use_fast_math", "--extended-lambda", "--compress-mode=size"] + return load_aot( + "ggml", + cuda_files=[str(s) for s in sources], + extra_cuda_cflags=flags, + extra_include_paths=[str(_ROOT / "include"), str(_ROOT / "src"), str(cuda_dir)], + build_directory=build_directory, + ) + + +@functools.cache +def _module() -> Module: + return load_module() + + +@functools.cache +def _mmq_tail_bytes(ggml_type: int, k: int) -> int: + return _module().mmq_tail_bytes(ggml_type, k) + + +def _check_mmq_tail(w: torch.Tensor, ggml_type: int, k: int) -> None: + """MMQ tiles read up to MATRIX_ROW_PADDING values past the last row of ``w``, so its storage must hold them.""" + tail = _mmq_tail_bytes(ggml_type, k) + if tail == 0 or w.numel() == 0: + return + end = w.storage_offset() + sum((n - 1) * s for n, s in zip(w.shape, w.stride())) + 1 + if end + tail > w.untyped_storage().nbytes(): + raise ValueError(f"ggml type {ggml_type} K={k} weight needs {tail} bytes of zero padding after its last row for MMQ") + + +def mul_mat(w: torch.Tensor, ggml_type: int, x: torch.Tensor, *, out: torch.Tensor | None = None, xq: torch.Tensor | None = None) -> torch.Tensor: + """``x[M, K] @ w[N, K]^T`` in ``x``'s dtype with ``w`` packed as ``[N, row_bytes]`` uint8, or written into the ``[M, N]`` view ``out``. + + ``xq`` is ``quantize_q8_1(x)``: the vector kernel reads it instead of quantizing ``x`` again. + """ + m, k = x.shape + if out is None: + out = torch.empty(m, w.shape[0], dtype=x.dtype, device=x.device) + if not uses_mmvq(ggml_type, m): + _check_mmq_tail(w, ggml_type, k) + _module().mul_mat(w, ggml_type, x.contiguous(), out, xq) + return out + + +def quantize_q8_1(x: torch.Tensor) -> torch.Tensor: + """``x[M, K]`` as the vector kernel's q8_1 blocks, to share among GEMVs that read the same ``x``.""" + m, k = x.shape + out = torch.empty(m, _q8_1_row_bytes(k), dtype=torch.uint8, device=x.device) + _module().quantize_q8_1(x.contiguous(), out) + return out + + +@functools.cache +def _q8_1_row_bytes(k: int) -> int: + return _module().q8_1_row_bytes(k) + + +@functools.cache +def uses_mmvq(ggml_type: int, m: int) -> bool: + """Whether ``mul_mat`` of ``m`` rows on this GPU runs the vector kernel, the one that reads ``xq``.""" + return _module().uses_mmvq(ggml_type, m) + + +def mul_mat_id(w: torch.Tensor, ggml_type: int, x: torch.Tensor, ids: torch.Tensor, scale: torch.Tensor | None = None) -> torch.Tensor: + """Per-route expert GEMM over ``w[E, N, row_bytes]``: returns ``[T, k, N]`` in ``x``'s dtype; ``x`` is ``[T, K]`` (shared by the k routes) or ``[T, k, K]``. + + ``scale`` is an fp32 ``[E]`` view (strided is fine) that multiplies each NVFP4 expert's rows; it needs ``T <= mmvq_max_tokens``. + """ + t, top_k = ids.shape + out = torch.empty(t, top_k, w.shape[1], dtype=x.dtype, device=x.device) + if t > mmvq_max_tokens(ggml_type): + _check_mmq_tail(w, ggml_type, x.shape[-1]) + _module().mul_mat_id(w, ggml_type, (x.unsqueeze(1) if x.dim() == 2 else x).contiguous(), ids, out, scale) + return out + + +def mul_mat_id_glu( + w: torch.Tensor, + ggml_type: int, + x: torch.Tensor, + ids: torch.Tensor, + glu: str, + limit: float = 0.0, + gate_scale: torch.Tensor | None = None, + up_scale: torch.Tensor | None = None, +) -> torch.Tensor: + """``glu(gate, up)`` per route in one pass for small ``T``: ``w[E, 2I, row_bytes]`` holds gate rows then up rows; returns ``[T, k, I]`` in ``x``'s dtype. + + ``glu`` is one of ``swiglu``, ``geglu`` (tanh GELU), ``swiglu_oai`` (alpha 1.702, limit 7) and ``swiglu_clamp`` (alpha 1, clamped at ``limit``). + ``gate_scale`` / ``up_scale`` are fp32 ``[E]`` views that multiply each NVFP4 expert's gate / up results before the glu. + """ + t, top_k = ids.shape + out = torch.empty(t, top_k, w.shape[1] // 2, dtype=x.dtype, device=x.device) + _module().mul_mat_id_glu(w, ggml_type, x.contiguous(), ids, out, glu, limit, gate_scale, up_scale) + return out + + +def moe_weighted_sum(y: torch.Tensor, weights: torch.Tensor) -> torch.Tensor: + """``sum_j weights[t, j] * y[t, j, :]`` over ``y[T, k, N]``, accumulated in fp32 and returned in ``y``'s dtype.""" + out = torch.empty(y.shape[0], y.shape[2], dtype=y.dtype, device=y.device) + _module().moe_weighted_sum(y.contiguous(), weights.float().contiguous(), out) + return out + + +@functools.cache +def mmvq_max_tokens(ggml_type: int) -> int: + """Largest token count the small-batch MoE kernels take on this GPU.""" + return _module().mmvq_max_tokens(ggml_type) + + +def dequantize(w: torch.Tensor, ggml_type: int, k: int, dtype: torch.dtype = torch.bfloat16) -> torch.Tensor: + out = torch.empty(w.numel() // w.shape[-1], k, dtype=dtype, device=w.device) + _module().dequantize(w, ggml_type, k, out) + return out + + +def get_rows(w: torch.Tensor, ggml_type: int, k: int, ids: torch.Tensor, dtype: torch.dtype = torch.bfloat16) -> torch.Tensor: + out = torch.empty(ids.shape[0], k, dtype=dtype, device=w.device) + _module().get_rows(w, ggml_type, k, ids, out) + return out + + +__all__ = [ + "dequantize", + "get_rows", + "load_module", + "mmvq_max_tokens", + "moe_weighted_sum", + "mul_mat", + "mul_mat_id", + "mul_mat_id_glu", + "quantize_q8_1", + "uses_mmvq", +] diff --git a/python/freetoken/kernel/ggml_cpu/__init__.py b/python/freetoken/kernel/ggml_cpu/__init__.py new file mode 100644 index 0000000000..909c973e3c --- /dev/null +++ b/python/freetoken/kernel/ggml_cpu/__init__.py @@ -0,0 +1,37 @@ +"""ggml-cpu row kernels for GGUF experts on the CPU MoE executor. + +setup.py builds one module per x86 variant of llama.cpp's GGML_CPU_ALL_VARIANTS; select() picks the one this CPU runs best. +""" + +from __future__ import annotations + +import functools +import importlib +import os +import pkgutil +from typing import NamedTuple + + +class Variant(NamedTuple): + name: str + api: int # address of the variant's kernel table, for _cpu_moe.set_ggml_cpu + + +@functools.cache +def select() -> Variant | None: + """The highest-scoring variant this CPU runs (upstream's rule), or the one FREETOKEN_GGML_CPU_VARIANT names; None when the build has none.""" + names = sorted(m.name[1:] for m in pkgutil.iter_modules(__path__) if m.name.startswith("_")) + forced = os.environ.get("FREETOKEN_GGML_CPU_VARIANT") + if forced: + if forced not in names: + raise RuntimeError(f"FREETOKEN_GGML_CPU_VARIANT={forced}: this build has the ggml-cpu variants {names}") + names = [forced] + best, best_score = None, 0 + for name in names: + # importing a variant runs no code built for its ISA extensions + mod = importlib.import_module(f"{__name__}._{name}") + if (score := mod.score()) > best_score: + best, best_score = Variant(name, mod.api()), score + if forced and best is None: + raise RuntimeError(f"FREETOKEN_GGML_CPU_VARIANT={forced}: this CPU lacks an ISA extension the variant is built for") + return best diff --git a/python/freetoken/kernel/gguf.py b/python/freetoken/kernel/gguf.py deleted file mode 100644 index 04a1656098..0000000000 --- a/python/freetoken/kernel/gguf.py +++ /dev/null @@ -1,142 +0,0 @@ -"""Borrowed llama.cpp GGUF dequant/GEMM CUDA kernels, JIT-compiled on first use. - -The ``.cu``/``.cuh`` under ``csrc/gguf/`` are vendored verbatim from sgl-kernel -(``csrc/quantization/gguf/``), which are themselves ports of llama.cpp. We compile -them through ``torch.utils.cpp_extension.load`` (the same toolchain sglang/vllm use) -into a torch-op module and expose the handful of ops the GGUF path needs. This is a -separate, torch-native extension that sits alongside FreeToken's tvm-ffi kernels. - -All ops keep the weight in its native GGUF block layout (packed ``uint8`` rows) and -dequantize *inside* the kernel -- no bf16 copy of the weight is ever materialized. -""" - -from __future__ import annotations - -import functools -import os -import pathlib -import shutil - -import torch - -_CSRC = pathlib.Path(__file__).parent / "csrc" / "gguf" - - -def _host_compiler() -> str | None: - """A host compiler nvcc + libtorch headers accept. - - The system default gcc can be too new for the torch headers (gcc 16 hard-errors), - and on this toolchain even nvcc+gcc-13 trips a non-conformant ``typename - decltype`` in ``List_inl.h`` once ``torch::Tensor`` is instantiated -- but nvcc - with ``clang++`` as host compiles it cleanly. So prefer clang++, then fall back - to an older gcc. Override with ``FREETOKEN_GGUF_HOST_CXX``. - """ - override = os.environ.get("FREETOKEN_GGUF_HOST_CXX") - if override: - return override - for cxx in ("clang++", "g++-13", "g++-14", "g++-15"): - if shutil.which(cxx): - return cxx - return None - - -def _c_compiler_for(cxx: str) -> str: - base = os.path.basename(cxx) - if "clang" in base: - return shutil.which("clang") or "clang" - cc = base.replace("g++", "gcc") - return shutil.which(cc) or cc - -@functools.cache -def _module(): - from torch.utils.cpp_extension import load - - extra_cuda_cflags = ["-O3", "--expt-relaxed-constexpr"] - host_cxx = _host_compiler() - if host_cxx is not None: - # Point both nvcc's host pass (-ccbin) and torch's C++ compile (CXX) at a - # libtorch/nvcc-compatible compiler. Force (not setdefault): the system - # default (CXX unset -> g++) can be a gcc too new for the torch headers. - cxx_path = shutil.which(host_cxx) or host_cxx - extra_cuda_cflags += ["-ccbin", cxx_path] - os.environ["CXX"] = cxx_path - os.environ["CC"] = _c_compiler_for(cxx_path) - - # gguf_kernel.cu carries its own PYBIND11_MODULE (appended at the end), so a - # plain `load` of the single source compiles + binds the ggml_* ops. - return load( - name="freetoken_gguf_kernels", - sources=[str(_CSRC / "gguf_kernel.cu")], - extra_include_paths=[str(_CSRC)], - extra_cuda_cflags=extra_cuda_cflags, - verbose=True, - ) - - -# ---- thin typed wrappers (signatures mirror sgl_kernel.quantization.gguf) ---- - - -def ggml_dequantize( - weight: torch.Tensor, quant_type: int, m: int, n: int, dtype: torch.dtype | None = None -) -> torch.Tensor: - """Dequantize a packed GGUF weight ``[m, row_bytes]`` to a dense ``[m, n]`` tensor.""" - return _module().ggml_dequantize(weight, quant_type, m, n, dtype) - - -def ggml_mul_mat_vec_a8( - weight: torch.Tensor, x: torch.Tensor, quant_type: int, row: int -) -> torch.Tensor: - """MMVQ: small-batch GEMV with on-the-fly dequant. ``row`` = output features.""" - return _module().ggml_mul_mat_vec_a8(weight, x, quant_type, row) - - -def ggml_mul_mat_a8( - weight: torch.Tensor, x: torch.Tensor, quant_type: int, row: int -) -> torch.Tensor: - """MMQ: large-batch quantized matmul. ``row`` = output features.""" - return _module().ggml_mul_mat_a8(weight, x, quant_type, row) - - -def ggml_moe_a8( - x: torch.Tensor, - weight: torch.Tensor, - sorted_token_ids: torch.Tensor, - expert_ids: torch.Tensor, - num_tokens_post_padded: torch.Tensor, - quant_type: int, - row: int, - top_k: int, - tokens: int, -) -> torch.Tensor: - """MMQ grouped expert matmul over stacked experts ``weight[E, row, *]``.""" - return _module().ggml_moe_a8( - x, weight, sorted_token_ids, expert_ids, num_tokens_post_padded, - quant_type, row, top_k, tokens, - ) - - -def ggml_moe_a8_vec( - x: torch.Tensor, - weight: torch.Tensor, - topk_ids: torch.Tensor, - top_k: int, - quant_type: int, - row: int, - tokens: int, -) -> torch.Tensor: - """MMVQ grouped expert GEMV over stacked experts ``weight[E, row, *]``.""" - return _module().ggml_moe_a8_vec(x, weight, topk_ids, top_k, quant_type, row, tokens) - - -def ggml_moe_get_block_size(quant_type: int) -> int: - return _module().ggml_moe_get_block_size(quant_type) - - -__all__ = [ - "ggml_dequantize", - "ggml_mul_mat_vec_a8", - "ggml_mul_mat_a8", - "ggml_moe_a8", - "ggml_moe_a8_vec", - "ggml_moe_get_block_size", -] diff --git a/python/freetoken/kernel/gluon/__init__.py b/python/freetoken/kernel/gluon/__init__.py new file mode 100644 index 0000000000..eaf4365ed6 --- /dev/null +++ b/python/freetoken/kernel/gluon/__init__.py @@ -0,0 +1,35 @@ +"""Gluon versions of Triton kernels for specific GPUs, written against the Gluon API of Triton :data:`TRITON_VERSION`. + +Callers route to a Gluon kernel only when :func:`available` is true and the kernel supports the device and inputs; otherwise the Triton kernel runs. +""" + +from __future__ import annotations + +import functools + +from freetoken.kernel.backend import is_rocm +from freetoken.kernel.utils import _env_enabled +from freetoken.utils import init_logger + +logger = init_logger(__name__) + +DISABLE_GLUON_ENV = "FREETOKEN_DISABLE_GLUON" +# Gluon is experimental and its API changes between Triton releases +TRITON_VERSION = (3, 8) + + +@functools.cache +def available() -> bool: + """True on CUDA with Triton :data:`TRITON_VERSION`, unless FREETOKEN_DISABLE_GLUON is set; call it only for a GPU that has a Gluon kernel, as another Triton version logs a warning.""" + if _env_enabled(DISABLE_GLUON_ENV) or is_rocm(): + return False + import triton + from packaging.version import Version + + if Version(triton.__version__).release[:2] == TRITON_VERSION: + return True + logger.warning( + "Gluon kernels need Triton %s, found %s; this GPU runs the slower Triton kernels instead (%s=1 silences this)", + ".".join(map(str, TRITON_VERSION)), triton.__version__, DISABLE_GLUON_ENV, + ) + return False diff --git a/python/freetoken/kernel/gluon/nvfp4/linear/sm120_gemv.py b/python/freetoken/kernel/gluon/nvfp4/linear/sm120_gemv.py new file mode 100644 index 0000000000..cb63149a81 --- /dev/null +++ b/python/freetoken/kernel/gluon/nvfp4/linear/sm120_gemv.py @@ -0,0 +1,230 @@ +"""NVFP4 W4A16 GEMV and skinny GEMM for sm_120 on the row-major weight: [N, K/2] e2m1 codes, flat 128x4 e4m3 block scales, one fp32 global. + +:func:`gemv_fma` decodes each code word once on the CUDA cores and reuses it for every token; :func:`gemv_mma` runs bf16 mma.sync with the +weight rows on M and the tokens on N. Both keep K a constexpr with the K loop unrolled and load the activation with evict_last, so it +stays in L2 while the weight stream passes; the weight keeps the default policy, as evict_first on it slows the kernels after it in a +model step (lines it would have evicted stay in L2 for them to evict). :func:`gemv_config` picks the kernel and tile for a shape. +""" + +import functools + +import torch +import triton +from triton.experimental import gluon +from triton.experimental.gluon import language as gl +from triton.experimental.gluon.language.nvidia.ampere import mma_v2 + +from ..utils import e2m1x8_bf16x2, scale_bf16x2 + +# one code word (k = 8w + 0..7) times its f16x2 scale pair -> 8 fp32; e2m1 times e4m3 is exact in f16 +_DECODE = gl.constexpr(""" +{ +.reg .b8 b0, b1, b2, b3; +.reg .b32 h0, h1, h2, h3; +.reg .b16 l0, l1, l2, l3, l4, l5, l6, l7; +mov.b32 {b0, b1, b2, b3}, $8; +cvt.rn.f16x2.e2m1x2 h0, b0; +cvt.rn.f16x2.e2m1x2 h1, b1; +cvt.rn.f16x2.e2m1x2 h2, b2; +cvt.rn.f16x2.e2m1x2 h3, b3; +mul.rn.f16x2 h0, h0, $9; +mul.rn.f16x2 h1, h1, $9; +mul.rn.f16x2 h2, h2, $9; +mul.rn.f16x2 h3, h3, $9; +mov.b32 {l0, l1}, h0; +mov.b32 {l2, l3}, h1; +mov.b32 {l4, l5}, h2; +mov.b32 {l6, l7}, h3; +cvt.f32.f16 $0, l0; +cvt.f32.f16 $1, l1; +cvt.f32.f16 $2, l2; +cvt.f32.f16 $3, l3; +cvt.f32.f16 $4, l4; +cvt.f32.f16 $5, l5; +cvt.f32.f16 $6, l6; +cvt.f32.f16 $7, l7; +} +""") + + +@gluon.jit +def _scale_row(row, K_TILES: gl.constexpr): + # 128x4 atoms: [row tile, k tile, row % 32, (row % 128) // 32, group % 4] + return (row // 128).to(gl.int64) * (K_TILES * 512) + (row % 32) * 16 + (row % 128) // 32 * 4 + + +@gluon.jit +def _gemv_fma_kernel(x_ptr, w_ptr, s_ptr, g_ptr, out_ptr, T, stride_x, stride_o, + K: gl.constexpr, T_C: gl.constexpr, R: gl.constexpr, V: gl.constexpr): + W: gl.constexpr = gl.num_warps() + BN: gl.constexpr = R * W + BKW: gl.constexpr = 32 * V + KW: gl.constexpr = K // 8 + K_TILES: gl.constexpr = (K // 16 + 3) // 4 + # a warp owns R whole rows; each lane holds V code words and their 8 k per row and per token + L4: gl.constexpr = gl.BlockedLayout([T_C, R, V, 8], [1, 1, 32, 1], [1, W, 1, 1], [3, 2, 1, 0]) + La: gl.constexpr = gl.SliceLayout(3, L4) + Lw: gl.constexpr = gl.SliceLayout(0, La) + Lx: gl.constexpr = gl.SliceLayout(1, L4) + Lxj: gl.constexpr = gl.SliceLayout(1, La) + + # 4 consecutive CTA rows are rows r, r + 32, r + 64, r + 96: one 16-byte chunk of every 128x4 scale atom + flat = gl.program_id(0) * BN + gl.arange(0, BN, layout=gl.SliceLayout(1, Lw)) + row = (flat // 128) * 128 + (flat % 128) // 4 + (flat % 4) * 32 + wcol = gl.arange(0, BKW, layout=gl.SliceLayout(0, Lw)) + w_base = w_ptr + row.to(gl.int64)[:, None] * KW + wcol[None, :] + s_row = _scale_row(row, K_TILES) + xt = gl.arange(0, T_C, layout=gl.SliceLayout(1, gl.SliceLayout(2, Lx))) + xw = gl.arange(0, BKW, layout=gl.SliceLayout(0, gl.SliceLayout(2, Lx))) + xj = gl.arange(0, 8, layout=gl.SliceLayout(0, gl.SliceLayout(1, Lx))) + x_off = xt[:, None, None] * stride_x + (xw * 8)[None, :, None] + xj[None, None, :] + x_ok = (xt < T)[:, None, None] + + acc = gl.zeros([T_C, BN, BKW], gl.float32, La) + for ph in gl.static_range(KW // BKW): + words = gl.load(w_base + ph * BKW) + grp = (ph * BKW + wcol) // 2 + sc = gl.load(s_ptr + s_row[:, None] + (grp // 4 * 512 + grp % 4)[None, :]) + s16 = sc.to(gl.float8e4nv, bitcast=True).to(gl.float16).to(gl.uint16, bitcast=True).to(gl.uint32) + s2 = (s16 | (s16 << 16)).to(gl.int32, bitcast=True) + v0, v1, v2, v3, v4, v5, v6, v7 = gl.inline_asm_elementwise(_DECODE, "=r,=r,=r,=r,=r,=r,=r,=r,r,r", [words, s2], + dtype=(gl.float32,) * 8, is_pure=True, pack=1) + xa = gl.load(x_ptr + ph * BKW * 8 + x_off, mask=x_ok, other=0.0, eviction_policy="evict_last").to(gl.float32) + # in-thread split of the 8 k of each word: x0..x7 line up with v0..v7 + ev, od = gl.split(gl.reshape(xa, [T_C, BKW, 4, 2])) + e0, e1 = gl.split(gl.reshape(ev, [T_C, BKW, 2, 2])) + o0, o1 = gl.split(gl.reshape(od, [T_C, BKW, 2, 2])) + x0, x4 = gl.split(e0) + x2, x6 = gl.split(e1) + x1, x5 = gl.split(o0) + x3, x7 = gl.split(o1) + acc += v0[None] * gl.convert_layout(x0, Lxj)[:, None, :] + acc += v1[None] * gl.convert_layout(x1, Lxj)[:, None, :] + acc += v2[None] * gl.convert_layout(x2, Lxj)[:, None, :] + acc += v3[None] * gl.convert_layout(x3, Lxj)[:, None, :] + acc += v4[None] * gl.convert_layout(x4, Lxj)[:, None, :] + acc += v5[None] * gl.convert_layout(x5, Lxj)[:, None, :] + acc += v6[None] * gl.convert_layout(x6, Lxj)[:, None, :] + acc += v7[None] * gl.convert_layout(x7, Lxj)[:, None, :] + res = gl.sum(acc, axis=2) * gl.load(g_ptr) + Lo: gl.constexpr = gl.SliceLayout(2, La) + ot = gl.arange(0, T_C, layout=gl.SliceLayout(1, Lo)) + orow = gl.convert_layout(row, gl.SliceLayout(0, Lo)) + gl.store(out_ptr + ot[:, None] * stride_o + orow[None, :], res.to(out_ptr.dtype.element_ty), mask=(ot < T)[:, None]) + + +@gluon.jit +def _to_bf16(q): + lo = (q & 0xFFFF).to(gl.uint16).to(gl.bfloat16, bitcast=True) + hi = (q >> 16).to(gl.uint16).to(gl.bfloat16, bitcast=True) + return gl.join(lo, hi) + + +@gluon.jit +def _kperm(t, A: gl.constexpr, B: gl.constexpr): + # k_phys = 32 lid + 4 s + 2 h + e -> k_mma = 16 s + 8 h + 2 lid + e: a lane's contiguous 32 k become its fragment slots + t = gl.reshape(t, [A, B, 4, 8, 2, 2]) + t = gl.permute(t, (0, 1, 3, 4, 2, 5)) + return gl.reshape(t, [A, B, 128]) + + +@gluon.jit +def _gemv_mma_kernel(x_ptr, w_ptr, s_ptr, g_ptr, out_ptr, T, stride_x, stride_o, + K: gl.constexpr, T_C: gl.constexpr, BM: gl.constexpr, Lc: gl.constexpr, Lx: gl.constexpr, mma: gl.constexpr): + KW: gl.constexpr = gl.num_warps() + NB: gl.constexpr = K // 128 + K_TILES: gl.constexpr = (K // 16 + 3) // 4 + al: gl.constexpr = gl.DotOperandLayout(0, mma, 2) + bl: gl.constexpr = gl.DotOperandLayout(1, mma, 2) + + # batched [KW, BM, T_C] mma: warp w takes the 128-k blocks w, w + KW, ...; the warps' partials are summed at the end + row0 = gl.program_id(0) * BM + cw = gl.arange(0, KW, layout=gl.SliceLayout(1, gl.SliceLayout(2, Lc))) + cm = gl.arange(0, BM, layout=gl.SliceLayout(0, gl.SliceLayout(2, Lc))) + cj = gl.arange(0, 16, layout=gl.SliceLayout(0, gl.SliceLayout(1, Lc))) + row = row0 + cm + c_off = (row.to(gl.int64) * (K // 8))[None, :, None] + (cw * 16)[:, None, None] + cj[None, None, :] + grp = cw[:, None, None] * 8 + (cj // 2)[None, None, :] + s_off = _scale_row(row, K_TILES)[None, :, None] + (grp // 4 * 512 + grp % 4) + xw = gl.arange(0, KW, layout=gl.SliceLayout(1, gl.SliceLayout(2, Lx))) + xt = gl.arange(0, T_C, layout=gl.SliceLayout(0, gl.SliceLayout(2, Lx))) + xk = gl.arange(0, 128, layout=gl.SliceLayout(0, gl.SliceLayout(1, Lx))) + x_off = (xt * stride_x)[None, :, None] + (xw * 128)[:, None, None] + xk[None, None, :] + x_ok = (xt < T)[None, :, None] + + acc = gl.zeros([KW, BM, T_C], gl.float32, mma) + for it in gl.static_range((NB + KW - 1) // KW): + b0 = it * KW + if (it + 1) * KW <= NB: + q = gl.load(w_ptr + c_off + b0 * 16) + sc = gl.load(s_ptr + s_off + b0 * 1024) + xa = gl.load(x_ptr + x_off + b0 * 128, mask=x_ok, other=0.0, eviction_policy="evict_last") + else: + ok = (cw + b0 < NB)[:, None, None] + q = gl.load(w_ptr + c_off + b0 * 16, mask=ok, other=0) + sc = gl.load(s_ptr + s_off + b0 * 1024, mask=ok, other=0) + xa = gl.load(x_ptr + x_off + b0 * 128, mask=x_ok & (xw + b0 < NB)[:, None, None], other=0.0, eviction_policy="evict_last") + w0, w1, w2, w3 = e2m1x8_bf16x2(q.to(gl.uint32, bitcast=True), scale_bf16x2(sc)) + # join appends its axis last: this order flattens to the k pairs w0, w1, w2, w3 + p = gl.reshape(gl.join(gl.join(w0, w2), gl.join(w1, w3)), [KW, BM, 64]) + a = gl.convert_layout(_kperm(gl.reshape(_to_bf16(p), [KW, BM, 128]), KW, BM), al) + b = gl.convert_layout(gl.permute(_kperm(xa, KW, T_C), (0, 2, 1)), bl) + acc = mma_v2(a, b, acc) + # e2m1x8_bf16x2 leaves a factor 2^-8 in the decoded weight + res = gl.sum(acc, axis=0) * (gl.load(g_ptr) * 256.0) + Lo: gl.constexpr = gl.SliceLayout(0, mma) + om = gl.arange(0, BM, layout=gl.SliceLayout(1, Lo)) + ot = gl.arange(0, T_C, layout=gl.SliceLayout(0, Lo)) + gl.store(out_ptr + ot[None, :] * stride_o + (row0 + om)[:, None], res.to(out_ptr.dtype.element_ty), mask=(ot < T)[None, :]) + + +@functools.cache +def _mma_layouts(kw: int, t_c: int, bm: int): + warps = [[1 << i, 0, 0] for i in range(kw.bit_length() - 1)] + lanes = [[0, 0, 32], [0, 0, 64], [0, 1, 0], [0, 2, 0], [0, 4, 0]] + # codes [KW, BM, 16 words]: lane (gid, lid) loads words 4 lid .. 4 lid + 3 of rows gid, gid + 8, gid + 16, ... + codes = gl.DistributedLinearLayout([[0, 0, 1], [0, 0, 2], [0, 8, 0]] + [[0, 16 << i, 0] for i in range(bm.bit_length() - 5)], + [[0, 0, 4], [0, 0, 8], [0, 1, 0], [0, 2, 0], [0, 4, 0]], warps, [], [kw, bm, 16]) + # activations [KW, T_C, 128]: lane (gid, lid) loads k 32 lid .. 32 lid + 31 of tokens gid, gid + 8 + acts = gl.DistributedLinearLayout([[0, 0, 1 << i] for i in range(5)] + [[0, 8 << i, 0] for i in range(t_c.bit_length() - 4)], + lanes, warps, [], [kw, t_c, 128]) + return codes, acts, gl.NVMMADistributedLayout([2, 0], [kw, 1, 1], [1, 16, 8]) + + +def gemv_fma(x: torch.Tensor, weight: torch.Tensor, scale: torch.Tensor, g: torch.Tensor, R: int, V: int, num_warps: int) -> torch.Tensor: + """x [T, K] bf16 row-contiguous, weight [N, K/2] uint8, scale flat 128x4 e4m3, g fp32 -> [T, N]; see :func:`gemv_config` for the shape rules.""" + T, K = x.shape + N = weight.shape[0] + out = torch.empty(T, N, dtype=x.dtype, device=x.device) + _gemv_fma_kernel[(N // (R * num_warps),)](x, weight.view(torch.int32), scale.view(torch.uint8), g, out, T, x.stride(0), out.stride(0), + K=K, T_C=triton.next_power_of_2(T), R=R, V=V, num_warps=num_warps) + return out + + +def gemv_mma(x: torch.Tensor, weight: torch.Tensor, scale: torch.Tensor, g: torch.Tensor, bm: int, num_warps: int) -> torch.Tensor: + """Same contract as :func:`gemv_fma`; T <= 16.""" + T, K = x.shape + N = weight.shape[0] + t_c = max(8, triton.next_power_of_2(T)) + codes, acts, mma = _mma_layouts(num_warps, t_c, bm) + out = torch.empty(T, N, dtype=x.dtype, device=x.device) + _gemv_mma_kernel[(N // bm,)](x, weight.view(torch.int32), scale.view(torch.uint8), g, out, T, x.stride(0), out.stride(0), + K=K, T_C=t_c, BM=bm, Lc=codes, Lx=acts, mma=mma, num_warps=num_warps) + return out + + +def gemv_config(N: int, K: int, T: int, max_tokens: int = 16): + """(kernel, tile) for a [T, K] x [N, K] product, or None when neither kernel covers it (T > max_tokens <= 16, or an unsupported shape).""" + if T == 1 and K < 16384 or T == 2 and N >= 16384: + words = K // 8 + if N % 128 == 0: + # wide layers take more rows per CTA; narrow ones longer words per lane + if N >= 16384 and words % 64 == 0: + return gemv_fma, {"R": 4, "V": 2, "num_warps": 4} + for V, warps in ((4, 4), (2, 8), (1, 8)): + if words % (32 * V) == 0 and not (N >= 16384 and V == 4): + return gemv_fma, {"R": 2, "V": V, "num_warps": warps} + if T > max_tokens or K % 128: + return None + tile = {"bm": 16, "num_warps": 8} if T <= 8 or N % 32 else {"bm": 32, "num_warps": 8} + return (gemv_mma, tile) if N % tile["bm"] == 0 else None diff --git a/python/freetoken/kernel/gluon/nvfp4/moe/sm120.py b/python/freetoken/kernel/gluon/nvfp4/moe/sm120.py new file mode 100644 index 0000000000..0455b237a3 --- /dev/null +++ b/python/freetoken/kernel/gluon/nvfp4/moe/sm120.py @@ -0,0 +1,354 @@ +"""NVFP4 W4A16 MoE prefill GEMM for sm_120: a loader warp (TMA weights, cp.async scales and tokens), decode warps and mma.sync warps. + +Same contract as the Triton prefill GEMM; :mod:`freetoken.kernel.nvfp4_moe` picks it for BLOCK_SIZE_M 64 and 128 and sets its constexprs. +""" + +import torch +import triton +from triton.experimental import gluon +from triton.experimental.gluon import language as gl +from triton.experimental.gluon.language.nvidia.ampere import async_copy, mma_v2 +from triton.experimental.gluon.language.nvidia.hopper import fence_async_shared, mbarrier, tma +from triton.experimental.gluon.nvidia.hopper import TensorDescriptor + +from ..utils import e2m1x8_bf16x2, scale_bf16x2 + + +@gluon.jit +def _tile(t, num_pid_m, num_pid_n, GROUP_SIZE_M: gl.constexpr): + num_pid_in_group = GROUP_SIZE_M * num_pid_n + first_pid_m = (t // num_pid_in_group) * GROUP_SIZE_M + group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) + pid_m = first_pid_m + ((t % num_pid_in_group) % group_size_m) + pid_n = (t % num_pid_in_group) // group_size_m + return pid_m, pid_n + + +@gluon.jit +def _load_partition(a_ptr, ids_ptr, eid_ptr, w_desc, scale_ptr, bufs, bars, N, nk, num_tiles, num_pid_m, num_valid, + stride_am, ks4, n_rows, top_k: gl.constexpr, K_TAIL: gl.constexpr, GROUP_SIZE_M: gl.constexpr, + PERSISTENT: gl.constexpr): + BN: gl.constexpr = bufs[1].type.shape[1] + SPC: gl.constexpr = 256 // (bufs[1].type.shape[2] * 2) + if PERSISTENT: + # it / jt count chunks and scale blocks across tiles, so the rings keep their phase from one tile to the next + it = 0 + jt = 0 + for t in range(gl.program_id(0), num_tiles, gl.num_programs(0)): + pid_m, pid_n = _tile(t, num_pid_m, gl.cdiv(N, BN), GROUP_SIZE_M) + _load_tile(a_ptr, ids_ptr, eid_ptr, w_desc, scale_ptr, bufs, bars, pid_m, pid_n, N, nk, it, jt, num_valid, + stride_am, ks4, n_rows, top_k, K_TAIL) + it += nk + jt += gl.cdiv(nk, SPC) + else: + pid_m, pid_n = _tile(gl.program_id(0), num_pid_m, gl.cdiv(N, BN), GROUP_SIZE_M) + _load_tile(a_ptr, ids_ptr, eid_ptr, w_desc, scale_ptr, bufs, bars, pid_m, pid_n, N, nk, 0, 0, num_valid, + stride_am, ks4, n_rows, top_k, K_TAIL) + + +@gluon.jit +def _load_tile(a_ptr, ids_ptr, eid_ptr, w_desc, scale_ptr, bufs, bars, pid_m, pid_n, N, nk, it, jt, num_valid, + stride_am, ks4, n_rows, top_k: gl.constexpr, K_TAIL: gl.constexpr): + x_bufs, w_bufs, s_bufs, d_bufs = bufs + x_ready, x_empty, w_ready, w_empty, s_ready, s_empty, d_ready = bars + KS: gl.constexpr = w_bufs.type.shape[2] // 8 + XS: gl.constexpr = x_bufs.type.shape[0] + BM: gl.constexpr = x_bufs.type.shape[1] // KS + BK: gl.constexpr = KS * 16 + WS: gl.constexpr = w_bufs.type.shape[0] + KB: gl.constexpr = w_bufs.type.shape[2] + SS: gl.constexpr = s_bufs.type.shape[0] + BN: gl.constexpr = s_bufs.type.shape[1] + SPC: gl.constexpr = 256 // BK + + # 16 bytes per cp.async; adjacent lanes read one token's contiguous K + NWL: gl.constexpr = gl.num_warps() + xl: gl.constexpr = gl.BlockedLayout([1, 1, 8], [KS, 16 // KS, 2], [1, NWL, 1], [2, 0, 1]) + offs_s = gl.arange(0, KS, layout=gl.SliceLayout(1, gl.SliceLayout(2, xl))) + offs_c = gl.arange(0, 16, layout=gl.SliceLayout(0, gl.SliceLayout(1, xl))) + cols = gl.expand_dims(gl.expand_dims(offs_s * 16, 1), 2) + gl.expand_dims(gl.expand_dims(offs_c, 0), 1) + + # scale rows are K / 16 bytes: 16-byte copies when K % 256 == 0, else 4-byte copies zeroed past the row end + SV: gl.constexpr = 4 if K_TAIL else 16 + sll: gl.constexpr = gl.BlockedLayout([1, SV], [32 * SV // 16, 16 // SV], [NWL, 1], [1, 0]) + s_cols = gl.arange(0, 16, layout=gl.SliceLayout(0, sll)) + # both forms equal K // 16; the second tells the compiler the row is a multiple of 16 bytes + row_len = ks4 * 4 if K_TAIL else (ks4 // 4) * 16 + + w_row = gl.load(eid_ptr + pid_m) * N + pid_n * BN + offs_m = pid_m * BM + gl.arange(0, BM, layout=gl.SliceLayout(0, gl.SliceLayout(2, xl))) + tok = gl.load(ids_ptr + offs_m) + tok_mask = gl.expand_dims(gl.expand_dims(tok < num_valid, 0), 2) + rows = gl.expand_dims(gl.expand_dims((tok // top_k).to(gl.int64) * stride_am, 0), 2) + x_ptrs = a_ptr + rows + cols + s_rows = w_row + gl.arange(0, BN, layout=gl.SliceLayout(1, sll)) + s_ptrs = scale_ptr + s_rows.to(gl.int64)[:, None] * row_len + s_cols[None, :] + s_row_ok = (s_rows < n_rows)[:, None] + for kk in range(nk): + if kk % SPC == 0: + jl = kk // SPC + j = jt + jl + ss = j % SS + mbarrier.wait(s_empty.index(ss), (j // SS & 1) ^ 1) + async_copy.async_load(s_bufs.index(ss), s_ptrs + jl * 16, mask=s_row_ok & (jl * 16 + s_cols < row_len)[None, :]) + async_copy.mbarrier_arrive(s_ready.index(ss)) + mbarrier.arrive(s_ready.index(ss)) + i = it + kk + ws = i % WS + mbarrier.wait(w_empty.index(ws), (i // WS & 1) ^ 1) + mbarrier.expect(w_ready.index(ws), w_desc.block_type.nbytes) + tma.async_load(w_desc, [w_row, kk * KB], w_ready.index(ws), w_bufs.index(ws)) + xs = i % XS + mbarrier.wait(x_empty.index(xs), (i // XS & 1) ^ 1) + async_copy.async_load(x_bufs.index(xs).reshape([KS, BM, 16]), x_ptrs + kk * BK, mask=tok_mask) + async_copy.mbarrier_arrive(x_ready.index(xs)) + mbarrier.arrive(x_ready.index(xs)) + + +@gluon.jit +def _decode_partition(a_ptr, ids_ptr, eid_ptr, w_desc, scale_ptr, bufs, bars, N, nk, num_tiles, num_pid_m, num_valid, + stride_am, ks4, n_rows, top_k: gl.constexpr, K_TAIL: gl.constexpr, GROUP_SIZE_M: gl.constexpr, + PERSISTENT: gl.constexpr): + SPC: gl.constexpr = 256 // (bufs[1].type.shape[2] * 2) + if PERSISTENT: + it = 0 + jt = 0 + for t in range(gl.program_id(0), num_tiles, gl.num_programs(0)): + _decode_tile(bufs, bars, nk, it, jt, K_TAIL) + it += nk + jt += gl.cdiv(nk, SPC) + else: + _decode_tile(bufs, bars, nk, 0, 0, K_TAIL) + + +@gluon.jit +def _decode_tile(bufs, bars, nk, it, jt, K_TAIL: gl.constexpr): + x_bufs, w_bufs, s_bufs, d_bufs = bufs + x_ready, x_empty, w_ready, w_empty, s_ready, s_empty, d_ready = bars + WS: gl.constexpr = w_bufs.type.shape[0] + BN: gl.constexpr = w_bufs.type.shape[1] + KB: gl.constexpr = w_bufs.type.shape[2] + BK: gl.constexpr = KB * 2 + SS: gl.constexpr = s_bufs.type.shape[0] + SB: gl.constexpr = BK // 16 + DS: gl.constexpr = d_bufs.type.shape[0] + SPC: gl.constexpr = 256 // BK + NW: gl.constexpr = gl.num_warps() + + # one thread per row, so the k16 slice stores of a warp hit 32 consecutive rows + ql: gl.constexpr = gl.BlockedLayout([1, 2 * SB], [32, 1], [NW, 1], [1, 0]) + q3l: gl.constexpr = gl.BlockedLayout([1, SB, 2], [32, 1, 1], [NW, 1, 1], [2, 1, 0]) + qsl: gl.constexpr = gl.SliceLayout(2, q3l) + w32_layout: gl.constexpr = gl.NVMMASharedLayout(swizzle_byte_width=w_bufs.type.layout.swizzle_byte_width, + element_bitwidth=32, rank=2) + d32_layout: gl.constexpr = gl.NVMMASharedLayout(swizzle_byte_width=32, element_bitwidth=32, rank=2) + w32 = w_bufs._reinterpret(gl.int32, [WS, BN, KB // 4], w32_layout) + d32 = d_bufs._reinterpret(gl.int32, [DS, SB * BN, 8], d32_layout) + + # with K % 256 != 0 the last scale block covers fewer than SPC chunks + for jb in range(gl.cdiv(nk, SPC)): + j = jt + jb + ss = j % SS + mbarrier.wait(s_ready.index(ss), j // SS & 1) + for c in gl.static_range(SPC): + kk = jb * SPC + c + if not K_TAIL or kk < nk: + i = it + kk + ws = i % WS + mbarrier.wait(w_ready.index(ws), i // WS & 1) + q = w32.index(ws).load(ql) + s = s_bufs.index(ss).slice(c * SB, SB, dim=1).load(qsl) + fence_async_shared() + mbarrier.arrive(w_empty.index(ws)) + q = gl.convert_layout(gl.reshape(q, [BN, SB, 2]), q3l, assert_trivial=True).to(gl.uint32, bitcast=True) + w0, w1, w2, w3 = e2m1x8_bf16x2(q, gl.expand_dims(scale_bf16x2(s), 2)) + # join appends the new axis last, so this flattens to w0, w1, w2, w3; then k16 slice major for the MMA warps + w = gl.reshape(gl.join(gl.join(w0, w2), gl.join(w1, w3)), [BN, SB, 8]) + w = gl.permute(w, (1, 0, 2)).to(gl.int32, bitcast=True) + + ds = i % DS + # the MMA warps release x and d slots together on x_empty + mbarrier.wait(x_empty.index(ds), (i // DS & 1) ^ 1) + d32.index(ds).reshape([SB, BN, 8]).store(w) + mbarrier.arrive(d_ready.index(ds)) + fence_async_shared() + mbarrier.arrive(s_empty.index(ss)) + + +@gluon.jit +def _mma_chunks(bufs, bars, acc, it, nk, al: gl.constexpr, bl: gl.constexpr, BN: gl.constexpr, BM: gl.constexpr): + x_bufs, w_bufs, s_bufs, d_bufs = bufs + x_ready, x_empty, w_ready, w_empty, s_ready, s_empty, d_ready = bars + KS: gl.constexpr = w_bufs.type.shape[2] // 8 + XS: gl.constexpr = x_bufs.type.shape[0] + for kk in range(nk): + i = it + kk + xs = i % XS + mbarrier.wait(d_ready.index(xs), i // XS & 1) + mbarrier.wait(x_ready.index(xs), i // XS & 1) + for k16 in gl.static_range(KS): + a = d_bufs.index(xs).slice(k16 * BN, BN).load(al) + b = x_bufs.index(xs).slice(k16 * BM, BM).permute((1, 0)).load(bl) + acc = mma_v2(a, b, acc) + mbarrier.arrive(x_empty.index(xs)) + return acc + + +@gluon.jit +def _mma_partition(bufs, bars, ids_ptr, eid_ptr, global_ptr, c_ptr, topk_weights_ptr, N, nk, num_tiles, num_pid_m, + num_valid, stride_ge, stride_cm, stride_tw, mma: gl.constexpr, BN: gl.constexpr, BM: gl.constexpr, + MUL_ROUTED_WEIGHT: gl.constexpr, GROUP_SIZE_M: gl.constexpr, PERSISTENT: gl.constexpr): + al: gl.constexpr = gl.DotOperandLayout(0, mma, 2) + bl: gl.constexpr = gl.DotOperandLayout(1, mma, 2) + if PERSISTENT: + it = 0 + for t in range(gl.program_id(0), num_tiles, gl.num_programs(0)): + acc = _mma_chunks(bufs, bars, gl.zeros([BN, BM], gl.float32, mma), it, nk, al, bl, BN, BM) + it += nk + # the epilogue runs here while the loader and decode warps already fill the rings for the next tile + # stored straight from the MMA layout, as a convert_layout needs scratch shared memory on top of the rings + pid_m, pid_n = _tile(t, num_pid_m, gl.cdiv(N, BN), GROUP_SIZE_M) + slot = gl.load(eid_ptr + pid_m) + offs_n = pid_n * BN + gl.arange(0, BN, layout=gl.SliceLayout(1, mma)) + offs_t = pid_m * BM + gl.arange(0, BM, layout=gl.SliceLayout(0, mma)) + _epilogue(acc, ids_ptr, global_ptr, c_ptr, topk_weights_ptr, slot, offs_n, offs_t, N, num_valid, + stride_ge, stride_cm, stride_tw, MUL_ROUTED_WEIGHT) + else: + return _mma_chunks(bufs, bars, gl.zeros([BN, BM], gl.float32, mma), 0, nk, al, bl, BN, BM) + + +@gluon.jit +def _epilogue(acc, ids_ptr, global_ptr, c_ptr, topk_weights_ptr, slot, offs_n, offs_t, N, num_valid, + stride_ge, stride_cm, stride_tw, MUL_ROUTED_WEIGHT: gl.constexpr): + route = gl.load(ids_ptr + offs_t) + route_mask = route < num_valid + n_mask = offs_n < N + g = gl.load(global_ptr + slot.to(gl.int64) * stride_ge + offs_n, mask=n_mask, other=0.0).to(gl.float32) + # 256 undoes the 2^-8 that the decode leaves in every weight + acc = acc * (g * 256.0)[:, None] + if MUL_ROUTED_WEIGHT: + wt = gl.load(topk_weights_ptr + route * stride_tw, mask=route_mask, other=0.0) + acc = acc * wt[None, :] + c_ptrs = c_ptr + route.to(gl.int64)[None, :] * stride_cm + offs_n[:, None] + gl.store(c_ptrs, acc.to(c_ptr.dtype.element_ty), mask=route_mask[None, :] & n_mask[:, None]) + + +SCALE_SMEM_LAYOUT = gl.constexpr(gl.NVMMASharedLayout(swizzle_byte_width=0, element_bitwidth=8, rank=2)) + + +@gluon.jit +def _prefill_nvfp4_moe_kernel( + a_ptr, w_desc, scale_ptr, global_ptr, c_ptr, topk_weights_ptr, + sorted_token_ids_ptr, expert_ids_ptr, num_tokens_post_padded_ptr, + N, K, EM, num_valid_tokens, ks4, n_rows, + stride_am, stride_ge, stride_cm, stride_tw, + BLOCK_SIZE_M: gl.constexpr, + BLOCK_SIZE_N: gl.constexpr, + BLOCK_SIZE_K: gl.constexpr, + GROUP_SIZE_M: gl.constexpr, + STAGES: gl.constexpr, + W_STAGES: gl.constexpr, + PERSISTENT: gl.constexpr, + K_TAIL: gl.constexpr, + DECODE_WARPS: gl.constexpr, + MMA_WARPS_N: gl.constexpr, + LOAD_REGS: gl.constexpr, + DECODE_REGS: gl.constexpr, + MUL_ROUTED_WEIGHT: gl.constexpr, + top_k: gl.constexpr, +): + BM: gl.constexpr = BLOCK_SIZE_M + BN: gl.constexpr = BLOCK_SIZE_N + BK: gl.constexpr = BLOCK_SIZE_K + live_m = gl.cdiv(gl.load(num_tokens_post_padded_ptr), BM) + if PERSISTENT: + num_pid_m = live_m + else: + # the grid covers the most m-blocks moe_align can produce; each CTA takes the tile of its index + num_pid_m = gl.cdiv(EM, BM) + pid_m, pid_n = _tile(gl.program_id(0), num_pid_m, gl.cdiv(N, BN), GROUP_SIZE_M) + if pid_m >= live_m: + return + num_tiles = num_pid_m * gl.cdiv(N, BN) + nk = K // BK + + # x and the decoded weights are stored as k16 slices, [stage, (k // 16) * rows, 16], and share STAGES slots + x_layout: gl.constexpr = gl.NVMMASharedLayout.get_default_for([BM, 16], gl.bfloat16) + d_layout: gl.constexpr = gl.NVMMASharedLayout.get_default_for([BN, 16], gl.bfloat16) + x_bufs = gl.allocate_shared_memory(gl.bfloat16, [STAGES, (BK // 16) * BM, 16], x_layout) + w_bufs = gl.allocate_shared_memory(gl.uint8, [W_STAGES, BN, BK // 2], w_desc.layout) + s_bufs = gl.allocate_shared_memory(gl.uint8, [2, BN, 16], SCALE_SMEM_LAYOUT) + d_bufs = gl.allocate_shared_memory(gl.bfloat16, [STAGES, (BK // 16) * BN, 16], d_layout) + bufs = (x_bufs, w_bufs, s_bufs, d_bufs) + + x_ready = gl.allocate_shared_memory(gl.int64, [STAGES, 1], mbarrier.MBarrierLayout()) + x_empty = gl.allocate_shared_memory(gl.int64, [STAGES, 1], mbarrier.MBarrierLayout()) + w_ready = gl.allocate_shared_memory(gl.int64, [W_STAGES, 1], mbarrier.MBarrierLayout()) + w_empty = gl.allocate_shared_memory(gl.int64, [W_STAGES, 1], mbarrier.MBarrierLayout()) + s_ready = gl.allocate_shared_memory(gl.int64, [2, 1], mbarrier.MBarrierLayout()) + s_empty = gl.allocate_shared_memory(gl.int64, [2, 1], mbarrier.MBarrierLayout()) + d_ready = gl.allocate_shared_memory(gl.int64, [STAGES, 1], mbarrier.MBarrierLayout()) + for i in gl.static_range(STAGES): + mbarrier.init(x_ready.index(i), count=1) + mbarrier.init(x_empty.index(i), count=1) + mbarrier.init(d_ready.index(i), count=1) + for i in gl.static_range(W_STAGES): + mbarrier.init(w_ready.index(i), count=1) + mbarrier.init(w_empty.index(i), count=1) + for i in gl.static_range(2): + mbarrier.init(s_ready.index(i), count=1) + mbarrier.init(s_empty.index(i), count=1) + bars = (x_ready, x_empty, w_ready, w_empty, s_ready, s_empty, d_ready) + + NW: gl.constexpr = gl.num_warps() + mma: gl.constexpr = gl.NVMMADistributedLayout(version=[2, 0], warps_per_cta=[MMA_WARPS_N, NW // MMA_WARPS_N], + instr_shape=[16, 8]) + # tuples go inline: a tuple stored in a variable turns its constexprs into runtime values + acc = gl.warp_specialize([ + (_mma_partition, (bufs, bars, sorted_token_ids_ptr, expert_ids_ptr, global_ptr, c_ptr, topk_weights_ptr, N, nk, + num_tiles, num_pid_m, num_valid_tokens, stride_ge, stride_cm, stride_tw, mma, BN, BM, + MUL_ROUTED_WEIGHT, GROUP_SIZE_M, PERSISTENT)), + (_load_partition, (a_ptr, sorted_token_ids_ptr, expert_ids_ptr, w_desc, scale_ptr, bufs, bars, N, nk, num_tiles, + num_pid_m, num_valid_tokens, stride_am, ks4, n_rows, top_k, K_TAIL, GROUP_SIZE_M, PERSISTENT)), + (_decode_partition, (a_ptr, sorted_token_ids_ptr, expert_ids_ptr, w_desc, scale_ptr, bufs, bars, N, nk, num_tiles, + num_pid_m, num_valid_tokens, stride_am, ks4, n_rows, top_k, K_TAIL, GROUP_SIZE_M, PERSISTENT)), + ], [1, DECODE_WARPS], [LOAD_REGS, DECODE_REGS]) + + if not PERSISTENT: + # a coalesced store layout; the rings are idle by now, so the conversion's scratch fits + out_l: gl.constexpr = gl.BlockedLayout([8, 1], [BN // 8, 32 * 8 // BN], [1, NW], [0, 1]) + pid_m, pid_n = _tile(gl.program_id(0), num_pid_m, gl.cdiv(N, BN), GROUP_SIZE_M) + slot = gl.load(expert_ids_ptr + pid_m) + offs_n = pid_n * BN + gl.arange(0, BN, layout=gl.SliceLayout(1, out_l)) + offs_t = pid_m * BM + gl.arange(0, BM, layout=gl.SliceLayout(0, out_l)) + _epilogue(gl.convert_layout(acc, out_l), sorted_token_ids_ptr, global_ptr, c_ptr, topk_weights_ptr, slot, offs_n, + offs_t, N, num_valid_tokens, stride_ge, stride_cm, stride_tw, MUL_ROUTED_WEIGHT) + + +def supports(a: torch.Tensor, packed: torch.Tensor, scale: torch.Tensor, c: torch.Tensor) -> bool: + """Whether :func:`prefill_gemm` can run this GEMM on an sm_120 GPU.""" + K = packed.shape[2] * 2 + # the TMA weight loads and 16-byte cp.async need contiguous, 16-byte aligned rows + return (a.dtype == torch.bfloat16 and K % 64 == 0 + and packed.is_contiguous() and scale.is_contiguous() and c.is_contiguous() + and packed.data_ptr() % 16 == 0 and scale.data_ptr() % 16 == 0 + and a.stride(1) == 1 and a.stride(0) % 8 == 0 and a.data_ptr() % 16 == 0) + + +def prefill_gemm(a, packed, scale, glob, c, topk_weights_flat, sorted_ids, expert_ids, num_tokens_post_padded, + num_valid_tokens, top_k, mul_routed_weight, cfg): + """Launch the kernel with the tile / pipeline constexprs in ``cfg``; K must be a multiple of cfg["BLOCK_SIZE_K"].""" + S, N, KB = packed.shape + K = KB * 2 + BM, BN, BK = cfg["BLOCK_SIZE_M"], cfg["BLOCK_SIZE_N"], cfg["BLOCK_SIZE_K"] + w_layout = gl.NVMMASharedLayout.get_default_for([BN, BK // 2], gl.uint8) + w_desc = TensorDescriptor.from_tensor(packed.view(S * N, KB), [BN, BK // 2], w_layout) + max_tiles = triton.cdiv(sorted_ids.shape[0], BM) * triton.cdiv(N, BN) + if cfg["PERSISTENT"]: + max_tiles = min(max_tiles, torch.cuda.get_device_properties(a.device).multi_processor_count) + _prefill_nvfp4_moe_kernel[(max_tiles,)]( + a, w_desc, scale.view(torch.uint8), glob, c, topk_weights_flat, sorted_ids, expert_ids, num_tokens_post_padded, + N, K, sorted_ids.shape[0], num_valid_tokens, K // 64, S * N, + a.stride(0), glob.stride(0), c.stride(1), topk_weights_flat.stride(0), + K_TAIL=K % 256 != 0, MUL_ROUTED_WEIGHT=mul_routed_weight, top_k=top_k, **cfg, + ) diff --git a/python/freetoken/kernel/gluon/nvfp4/moe/sm120_w4a4.py b/python/freetoken/kernel/gluon/nvfp4/moe/sm120_w4a4.py new file mode 100644 index 0000000000..2a2c1430ce --- /dev/null +++ b/python/freetoken/kernel/gluon/nvfp4/moe/sm120_w4a4.py @@ -0,0 +1,536 @@ +"""NVFP4 W4A4 MoE prefill GEMMs for sm_120: NVFP4 weights times NVFP4 activations on mma.sync kind::mxf4nvf4. + +The activations are quantized under static global scales (the checkpoint's input_scale), so each GEMM input is e2m1 codes, per-16 e4m3 scales and one fp32 scalar. +:func:`gate_up_gemm` applies the gated activation in its epilogue and writes the down GEMM's NVFP4 input; :func:`down_gemm` scales by the routing weight and adds each route into its token's output row with bf16 atomics (the sum order, and so the result, varies from run to run). +Routing follows the W4A16 kernels (moe_align sorted ids, expert per M block, swap-AB accumulator [BN, BM]); :mod:`freetoken.kernel.nvfp4_moe` picks the constexprs. +""" + +import torch +import triton +from triton.experimental import gluon +from triton.experimental.gluon import language as gl +from triton.experimental.gluon.language.nvidia.ampere import async_copy +from triton.experimental.gluon.language.nvidia.hopper import mbarrier, tma +from triton.experimental.gluon.nvidia.hopper import TensorDescriptor + +from ..utils import mma_nvf4, scale_layout + +# gated activations of the gate_up epilogue, as gated_act_and_mul computes them +ACTIVATIONS = {"silu": 0, "gelu_tanh": 1, "swigluoai": 2, "swiglu_clamp": 3} + +SCALE_SMEM_LAYOUT = gl.constexpr(gl.SwizzledSharedLayout(1, 1, 1, [1, 0])) + + +@gluon.jit +def _tile(t, num_pid_m, num_pid_n, GROUP_SIZE_M: gl.constexpr, N_GROUP: gl.constexpr): + if N_GROUP > 0: + # N-outer strips of N_GROUP tiles: concurrent CTAs add into a narrow output strip that stays in L2 + per_strip = num_pid_m * N_GROUP + strip = t // per_strip + rem = t % per_strip + width = min(N_GROUP, num_pid_n - strip * N_GROUP) + return rem // width, strip * N_GROUP + rem % width + num_pid_in_group = GROUP_SIZE_M * num_pid_n + first_pid_m = (t // num_pid_in_group) * GROUP_SIZE_M + group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) + return first_pid_m + ((t % num_pid_in_group) % group_size_m), (t % num_pid_in_group) // group_size_m + + +@gluon.jit +def _scale_ptrs(ws_ptr, xs_ptr, ids_ptr, w_row, w_row_up, x_div, pid_m, num_valid, ks4, n_rows, + BM: gl.constexpr, BN: gl.constexpr, K_TAIL: gl.constexpr, GATED: gl.constexpr, NWL: gl.constexpr): + # scale rows are K / 16 bytes: 16-byte copies when K % 256 == 0, else 4-byte copies zeroed past the row end + SV: gl.constexpr = 4 if K_TAIL else 16 + sll: gl.constexpr = gl.BlockedLayout([1, SV], [32 * SV // 16, 16 // SV], [NWL, 1], [1, 0]) + s_cols = gl.arange(0, 16, layout=gl.SliceLayout(0, sll)) + # both forms equal K // 16; the second tells the compiler the row is a multiple of 16 bytes + row_len = ks4 * 4 if K_TAIL else (ks4 // 4) * 16 + r = gl.arange(0, BN, layout=gl.SliceLayout(1, sll)) + if GATED: + sw_rows = gl.where(r < BN // 2, w_row + r, w_row_up + r - BN // 2) + else: + sw_rows = w_row + r + sw_ptrs = ws_ptr + sw_rows.to(gl.int64)[:, None] * row_len + s_cols[None, :] + sw_ok = (sw_rows < n_rows)[:, None] + sx_tok = gl.load(ids_ptr + pid_m * BM + gl.arange(0, BM, layout=gl.SliceLayout(1, sll))) + # the x scale rows are contiguous K / 16 bytes too; row_len keeps the divisibility a stride argument loses + sx_ptrs = xs_ptr + ((sx_tok // x_div).to(gl.int64) * row_len)[:, None] + s_cols[None, :] + sx_ok = (sx_tok < num_valid)[:, None] + return sw_ptrs, sw_ok, sx_ptrs, sx_ok, s_cols, row_len + + +@gluon.jit +def _load_tile(xq_ptr, xs_ptr, ids_ptr, eid_ptr, w_desc, ws_ptr, bufs, bars, pid_m, pid_n, N, nk, it, jt, num_valid, + stride_xq, ks4, n_rows, x_div, K_TAIL: gl.constexpr, GATED: gl.constexpr): + x_bufs, w_bufs, sw_bufs, sx_bufs = bufs + ready, empty, s_ready, s_empty = bars + S: gl.constexpr = x_bufs.type.shape[0] + BM: gl.constexpr = x_bufs.type.shape[1] + KB: gl.constexpr = x_bufs.type.shape[2] + BN: gl.constexpr = w_bufs.type.shape[1] + SPC: gl.constexpr = 128 // KB + NWL: gl.constexpr = gl.num_warps() + + if GATED: + # rows [0, BN/2) of the tile are gate rows, [BN/2, BN) the matching up rows N/2 further down + w_row = gl.load(eid_ptr + pid_m) * N + pid_n * (BN // 2) + w_row_up = w_row + N // 2 + else: + w_row = gl.load(eid_ptr + pid_m) * N + pid_n * BN + w_row_up = w_row + + # x rows: 16 bytes per cp.async, KB // 16 lanes per row + xl: gl.constexpr = gl.BlockedLayout([1, 16], [32 * 16 // KB, KB // 16], [NWL, 1], [1, 0]) + tok = gl.load(ids_ptr + pid_m * BM + gl.arange(0, BM, layout=gl.SliceLayout(1, xl))) + x_ptrs = xq_ptr + ((tok // x_div).to(gl.int64) * stride_xq)[:, None] + gl.arange(0, KB, layout=gl.SliceLayout(0, xl))[None, :] + x_mask = (tok < num_valid)[:, None] + sw_ptrs, sw_ok, sx_ptrs, sx_ok, s_cols, row_len = _scale_ptrs( + ws_ptr, xs_ptr, ids_ptr, w_row, w_row_up, x_div, pid_m, num_valid, ks4, n_rows, BM, BN, K_TAIL, GATED, NWL) + + for kk in range(nk): + if kk % SPC == 0: + jl = kk // SPC + j = jt + jl + ss = j % 2 + mbarrier.wait(s_empty.index(ss), (j // 2 & 1) ^ 1) + col_ok = (jl * 16 + s_cols < row_len)[None, :] + async_copy.async_load(sw_bufs.index(ss), sw_ptrs + jl * 16, mask=sw_ok & col_ok) + async_copy.async_load(sx_bufs.index(ss), sx_ptrs + jl * 16, mask=sx_ok & col_ok) + async_copy.mbarrier_arrive(s_ready.index(ss)) + mbarrier.arrive(s_ready.index(ss)) + c = it + kk + st = c % S + mbarrier.wait(empty.index(st), (c // S & 1) ^ 1) + # the cp.async arrive goes first: it raises the pending count before the TMA expect can complete the phase + async_copy.async_load(x_bufs.index(st), x_ptrs + kk * KB, mask=x_mask) + async_copy.mbarrier_arrive(ready.index(st)) + if GATED: + mbarrier.expect(ready.index(st), 2 * w_desc.block_type.nbytes) + tma.async_load(w_desc, [w_row, kk * KB], ready.index(st), w_bufs.index(st).slice(0, BN // 2)) + tma.async_load(w_desc, [w_row_up, kk * KB], ready.index(st), w_bufs.index(st).slice(BN // 2, BN // 2)) + else: + mbarrier.expect(ready.index(st), w_desc.block_type.nbytes) + tma.async_load(w_desc, [w_row, kk * KB], ready.index(st), w_bufs.index(st)) + + +@gluon.jit +def _load_partition(xq_ptr, xs_ptr, ids_ptr, eid_ptr, w_desc, ws_ptr, bufs, bars, N, nk, num_pid_m, num_valid, + stride_xq, ks4, n_rows, x_div, K_TAIL: gl.constexpr, GROUP_SIZE_M: gl.constexpr, + GATED: gl.constexpr, N_GROUP: gl.constexpr): + BN: gl.constexpr = bufs[1].type.shape[1] + pid_m, pid_n = _tile(gl.program_id(0), num_pid_m, gl.cdiv(N, BN), GROUP_SIZE_M, N_GROUP) + _load_tile(xq_ptr, xs_ptr, ids_ptr, eid_ptr, w_desc, ws_ptr, bufs, bars, pid_m, pid_n, N, nk, 0, 0, num_valid, + stride_xq, ks4, n_rows, x_div, K_TAIL, GATED) + + +@gluon.jit +def _mma_partition(bufs, bars, nk, mma: gl.constexpr, wsl: gl.constexpr, xsl: gl.constexpr, K_TAIL: gl.constexpr): + x_bufs, w_bufs, sw_bufs, sx_bufs = bufs + ready, empty, s_ready, s_empty = bars + S: gl.constexpr = x_bufs.type.shape[0] + BM: gl.constexpr = x_bufs.type.shape[1] + KB: gl.constexpr = x_bufs.type.shape[2] + BN: gl.constexpr = w_bufs.type.shape[1] + SB: gl.constexpr = KB // 8 + SPC: gl.constexpr = 128 // KB + al: gl.constexpr = gl.DotOperandLayout(0, mma, 4) + bl: gl.constexpr = gl.DotOperandLayout(1, mma, 4) + acc = gl.zeros([BN, BM], gl.float32, mma) + # with K % 256 != 0 the last scale block covers fewer than SPC chunks + for jb in range(gl.cdiv(nk, SPC)): + ss = jb % 2 + mbarrier.wait(s_ready.index(ss), jb // 2 & 1) + for c in gl.static_range(SPC): + kk = jb * SPC + c + if not K_TAIL or kk < nk: + st = kk % S + mbarrier.wait(ready.index(st), kk // S & 1) + a = w_bufs.index(st).load(al) + b = x_bufs.index(st).permute((1, 0)).load(bl) + sw = sw_bufs.index(ss).slice(c * SB, SB, dim=1).load(wsl).to(gl.float8e4nv, bitcast=True) + sx = sx_bufs.index(ss).slice(c * SB, SB, dim=1).load(xsl).to(gl.float8e4nv, bitcast=True) + mbarrier.arrive(empty.index(st)) + acc = mma_nvf4(a, sw, b, sx, acc) + mbarrier.arrive(s_empty.index(ss)) + return acc + + +@gluon.jit +def _e2m1x2(hi, lo): + # cvt packs its first source into the high nibble + return gl.inline_asm_elementwise("{ .reg .b8 t; cvt.rn.satfinite.e2m1x2.f32 t, $1, $2; cvt.u16.u8 $0, t; }", + "=h,r,r", [hi, lo], gl.uint16, True, 1) + + +@gluon.jit +def _gated_epilogue(acc, slot, pid_m, pid_n, wg_ptr, g_in, ids_ptr, q_ptr, qs_ptr, g_out, N, num_valid, stride_ge, + ACT: gl.constexpr, act_alpha, act_limit): + """act(gate, up) of the tile's BN / 2 columns as NVFP4 under the global g_out: q [routes, N / 4], qs [routes, N / 32].""" + BN: gl.constexpr = acc.type.shape[0] + BM: gl.constexpr = acc.type.shape[1] + H2: gl.constexpr = BN // 2 + G: gl.constexpr = H2 // 16 + NW: gl.constexpr = gl.num_warps() + out_l: gl.constexpr = acc.type.layout + r = gl.arange(0, BN, layout=gl.SliceLayout(1, out_l)) + col = gl.where(r < H2, pid_n * H2 + r, N // 2 + pid_n * H2 + r - H2) + wg = gl.load(wg_ptr + slot.to(gl.int64) * stride_ge + col).to(gl.float32) * g_in + # bf16 like the W4A16 gate_up output; n-contiguous 16-byte runs, swizzled so the lanes of a warp spread over banks + h_l: gl.constexpr = gl.SwizzledSharedLayout(8, 1, 8, [0, 1]) + h_s = gl.allocate_shared_memory(gl.bfloat16, [BN, BM], h_l, (acc * wg[:, None]).to(gl.bfloat16)) + # one thread holds 16 consecutive columns of one route + TG: gl.constexpr = min(32, G) + ql2: gl.constexpr = gl.BlockedLayout([16, 1], [TG, 32 // TG], [1, NW], [0, 1]) + l4: gl.constexpr = gl.BlockedLayout([1, 8, 2, 1], [TG, 1, 1, 32 // TG], [1, 1, 1, NW], [2, 1, 0, 3]) + gate = gl.convert_layout(gl.reshape(h_s.slice(0, H2).load(ql2), [G, 8, 2, BM]), l4, assert_trivial=True).to(gl.float32) + up = gl.convert_layout(gl.reshape(h_s.slice(H2, H2).load(ql2), [G, 8, 2, BM]), l4, assert_trivial=True).to(gl.float32) + if ACT == 1: + inner = 0.7978845608028654 * (gate + 0.044715 * gate * gate * gate) + # tanh(y) = 1 - 2 / (exp(2y) + 1) + h = 0.5 * gate * (2.0 - 2.0 / (gl.exp(2.0 * inner) + 1.0)) * up + elif ACT == 2 or ACT == 3: + gate = gl.minimum(gate, act_limit) + up = gl.minimum(gl.maximum(up, -act_limit), act_limit) + if ACT == 2: + up = up + 1.0 + h = gate / (1.0 + gl.exp(-gate * act_alpha)) * up + else: + h = gate / (1.0 + gl.exp(-gate)) * up + h = h.to(gl.bfloat16).to(gl.float32) + amax = gl.max(gl.max(gl.abs(h), axis=2), axis=1) + sc = (amax / (6.0 * g_out)).to(gl.float8e4nv) + sf = sc.to(gl.float32) * g_out + inv = gl.where(sf > 0, 1.0 / sf, 0.0) + v = h * gl.expand_dims(gl.expand_dims(inv, 1), 2) + lo, hi = gl.split(gl.permute(v, (0, 1, 3, 2))) + q = _e2m1x2(hi, lo).to(gl.uint8) + ql: gl.constexpr = q.type.layout + g_i = gl.arange(0, G, layout=gl.SliceLayout(1, gl.SliceLayout(2, ql))) + b_i = gl.arange(0, 8, layout=gl.SliceLayout(0, gl.SliceLayout(2, ql))) + rq = gl.load(ids_ptr + pid_m * BM + gl.arange(0, BM, layout=gl.SliceLayout(0, gl.SliceLayout(1, ql)))) + q_cols = pid_n * (H2 // 2) + g_i[:, None] * 8 + b_i[None, :] + q_ptrs = q_ptr + (rq.to(gl.int64) * (N // 4))[None, None, :] + q_cols[:, :, None] + gl.store(q_ptrs, q, mask=(rq < num_valid)[None, None, :]) + sl: gl.constexpr = sc.type.layout + sg = gl.arange(0, G, layout=gl.SliceLayout(1, sl)) + rs = gl.load(ids_ptr + pid_m * BM + gl.arange(0, BM, layout=gl.SliceLayout(0, sl))) + s_ptrs = qs_ptr + (rs.to(gl.int64) * (N // 32))[None, :] + (pid_n * G + sg)[:, None] + gl.store(s_ptrs, sc.to(gl.uint8, bitcast=True), mask=(rs < num_valid)[None, :]) + + +@gluon.jit +def _w4a4_moe_kernel( + xq_ptr, xs_ptr, g_in_ptr, w_desc, ws_ptr, wg_ptr, out_ptr, topk_weights_ptr, + sorted_token_ids_ptr, expert_ids_ptr, num_tokens_post_padded_ptr, + q_ptr, qs_ptr, g_out_ptr, act_alpha, act_limit, + N, K, EM, num_valid_tokens, ks4, n_rows, + stride_xq, stride_ge, stride_out, stride_tw, + BLOCK_SIZE_M: gl.constexpr, BLOCK_SIZE_N: gl.constexpr, BLOCK_SIZE_K: gl.constexpr, + GROUP_SIZE_M: gl.constexpr, N_GROUP: gl.constexpr, STAGES: gl.constexpr, MMA_WARPS_N: gl.constexpr, + LOAD_WARPS: gl.constexpr, LOAD_REGS: gl.constexpr, K_TAIL: gl.constexpr, GATED: gl.constexpr, + ACT: gl.constexpr, top_k: gl.constexpr, wsl: gl.constexpr, xsl: gl.constexpr, +): + BM: gl.constexpr = BLOCK_SIZE_M + BN: gl.constexpr = BLOCK_SIZE_N + KB: gl.constexpr = BLOCK_SIZE_K // 2 + num_pid_m = gl.cdiv(EM, BM) + pid_m, pid_n = _tile(gl.program_id(0), num_pid_m, gl.cdiv(N, BN), GROUP_SIZE_M, N_GROUP) + if pid_m >= gl.cdiv(gl.load(num_tokens_post_padded_ptr), BM): + return + nk = K // BLOCK_SIZE_K + # gate_up reads one x row per token, down one row per route + x_div = top_k if GATED else 1 + + x_bufs = gl.allocate_shared_memory(gl.uint8, [STAGES, BM, KB], gl.NVMMASharedLayout.get_default_for([BM, KB], gl.uint8)) + w_bufs = gl.allocate_shared_memory(gl.uint8, [STAGES, BN, KB], w_desc.layout) + sw_bufs = gl.allocate_shared_memory(gl.uint8, [2, BN, 16], SCALE_SMEM_LAYOUT) + sx_bufs = gl.allocate_shared_memory(gl.uint8, [2, BM, 16], SCALE_SMEM_LAYOUT) + ready = gl.allocate_shared_memory(gl.int64, [STAGES, 1], mbarrier.MBarrierLayout()) + empty = gl.allocate_shared_memory(gl.int64, [STAGES, 1], mbarrier.MBarrierLayout()) + s_ready = gl.allocate_shared_memory(gl.int64, [2, 1], mbarrier.MBarrierLayout()) + s_empty = gl.allocate_shared_memory(gl.int64, [2, 1], mbarrier.MBarrierLayout()) + for i in gl.static_range(STAGES): + mbarrier.init(ready.index(i), count=1) + mbarrier.init(empty.index(i), count=1) + for i in gl.static_range(2): + mbarrier.init(s_ready.index(i), count=1) + mbarrier.init(s_empty.index(i), count=1) + + NW: gl.constexpr = gl.num_warps() + mma: gl.constexpr = gl.NVMMADistributedLayout(version=[2, 0], warps_per_cta=[MMA_WARPS_N, NW // MMA_WARPS_N], + instr_shape=[16, 8]) + # tuples go inline: a tuple stored in a variable turns its constexprs into runtime values + acc = gl.warp_specialize([ + (_mma_partition, ((x_bufs, w_bufs, sw_bufs, sx_bufs), (ready, empty, s_ready, s_empty), nk, mma, wsl, xsl, K_TAIL)), + (_load_partition, (xq_ptr, xs_ptr, sorted_token_ids_ptr, expert_ids_ptr, w_desc, ws_ptr, + (x_bufs, w_bufs, sw_bufs, sx_bufs), (ready, empty, s_ready, s_empty), N, nk, num_pid_m, + num_valid_tokens, stride_xq, ks4, n_rows, x_div, K_TAIL, GROUP_SIZE_M, GATED, N_GROUP)), + ], [LOAD_WARPS], [LOAD_REGS]) + + out_l: gl.constexpr = gl.BlockedLayout([8, 1], [BN // 8, 32 * 8 // BN], [1, NW], [0, 1]) + acc = gl.convert_layout(acc, out_l) + slot = gl.load(expert_ids_ptr + pid_m) + g_in = gl.load(g_in_ptr) + if GATED: + _gated_epilogue(acc, slot, pid_m, pid_n, wg_ptr, g_in, sorted_token_ids_ptr, q_ptr, qs_ptr, gl.load(g_out_ptr), + N, num_valid_tokens, stride_ge, ACT, act_alpha, act_limit) + else: + offs_n = pid_n * BN + gl.arange(0, BN, layout=gl.SliceLayout(1, out_l)) + n_mask = offs_n < N + route = gl.load(sorted_token_ids_ptr + pid_m * BM + gl.arange(0, BM, layout=gl.SliceLayout(0, out_l))) + route_mask = route < num_valid_tokens + wg = gl.load(wg_ptr + slot.to(gl.int64) * stride_ge + offs_n, mask=n_mask, other=0.0).to(gl.float32) * g_in + tw = gl.load(topk_weights_ptr + route * stride_tw, mask=route_mask, other=0.0) + acc = acc * wg[:, None] * tw[None, :] + out_ptrs = out_ptr + (route // top_k).to(gl.int64)[None, :] * stride_out + offs_n[:, None] + gl.atomic_add(out_ptrs, acc.to(out_ptr.dtype.element_ty), mask=route_mask[None, :] & n_mask[:, None], sem="relaxed") + + +# Persistent down GEMM: the 256 x 128 tile and 8 MMA warps of the kernel above, with the routing-weight scale and the atomics on +# separate epilogue warps. The accumulator is two 128 x 128 halves that go to the epilogue one at a time through one 32 KB buffer, +# so the next tile's mainloop overlaps the atomics of the last half. + + +@gluon.jit +def _num_tiles(ntpp_ptr, N, BM: gl.constexpr, BN: gl.constexpr): + num_pid_m = gl.cdiv(gl.load(ntpp_ptr), BM) + num_pid_n = gl.cdiv(N, BN) + return num_pid_m, num_pid_n, num_pid_m * num_pid_n + + +@gluon.jit +def _down_load_partition(xq_ptr, xs_ptr, ids_ptr, eid_ptr, ntpp_ptr, w_desc, ws_ptr, bufs, bars, N, nk, num_valid, + stride_xq, ks4, n_rows, K_TAIL: gl.constexpr, N_GROUP: gl.constexpr, GROUP_SIZE_M: gl.constexpr): + BM: gl.constexpr = bufs[0].type.shape[1] + KB: gl.constexpr = bufs[0].type.shape[2] + BN: gl.constexpr = bufs[1].type.shape[1] + SPC: gl.constexpr = 128 // KB + num_pid_m, num_pid_n, num_tiles = _num_tiles(ntpp_ptr, N, BM, BN) + # it / jt count chunks and scale blocks across tiles, so the rings keep their phase from one tile to the next + it = 0 + jt = 0 + for t in range(gl.program_id(0), num_tiles, gl.num_programs(0)): + pid_m, pid_n = _tile(t, num_pid_m, num_pid_n, GROUP_SIZE_M, N_GROUP) + _load_tile(xq_ptr, xs_ptr, ids_ptr, eid_ptr, w_desc, ws_ptr, (bufs[0], bufs[1], bufs[2], bufs[3]), + (bars[0], bars[1], bars[2], bars[3]), pid_m, pid_n, N, nk, it, jt, num_valid, stride_xq, ks4, + n_rows, 1, K_TAIL, False) + it += nk + jt += gl.cdiv(nk, SPC) + + +@gluon.jit +def _down_epi_partition(ids_ptr, ntpp_ptr, out_ptr, bufs, bars, N, num_valid, stride_out, top_k: gl.constexpr, + N_GROUP: gl.constexpr, GROUP_SIZE_M: gl.constexpr): + h_buf = bufs[4] + acc_ready, acc_empty = bars[4], bars[5] + HB: gl.constexpr = h_buf.type.shape[0] + BM: gl.constexpr = h_buf.type.shape[1] + BN: gl.constexpr = 2 * HB + CH: gl.constexpr = 64 + NWE: gl.constexpr = gl.num_warps() + # 8 consecutive n of one route per thread: one 16-byte bf16 atomic; 64 rows per pass keep the registers low + el: gl.constexpr = gl.BlockedLayout([8, 1], [CH // 8, 32 * 8 // CH], [1, NWE], [0, 1]) + offs_cl = gl.arange(0, CH, layout=gl.SliceLayout(1, el)) + offs_tl = gl.arange(0, BM, layout=gl.SliceLayout(0, el)) + num_pid_m, num_pid_n, num_tiles = _num_tiles(ntpp_ptr, N, BM, BN) + hc = 0 + for t in range(gl.program_id(0), num_tiles, gl.num_programs(0)): + pid_m, pid_n = _tile(t, num_pid_m, num_pid_n, GROUP_SIZE_M, N_GROUP) + route = gl.load(ids_ptr + pid_m * BM + offs_tl) + rmask = (route < num_valid)[None, :] + rows = (route // top_k).to(gl.int64)[None, :] * stride_out + for half in gl.static_range(2): + mbarrier.wait(acc_ready.index(0), hc & 1) + for ch in gl.static_range(HB // CH): + offs_n = pid_n * BN + half * HB + ch * CH + offs_cl + h = h_buf.slice(ch * CH, CH).load(el) + gl.atomic_add(out_ptr + rows + offs_n[:, None], h, mask=rmask & (offs_n < N)[:, None], sem="relaxed") + mbarrier.arrive(acc_empty.index(0)) + hc += 1 + + +@gluon.jit +def _down_mma_partition(bufs, bars, ids_ptr, eid_ptr, ntpp_ptr, wg_ptr, g_in_ptr, tw_ptr, N, nk, num_valid, stride_ge, + stride_tw, mma: gl.constexpr, wsl: gl.constexpr, xsl: gl.constexpr, K_TAIL: gl.constexpr, + N_GROUP: gl.constexpr, GROUP_SIZE_M: gl.constexpr): + x_bufs, w_bufs, sw_bufs, sx_bufs, h_buf = bufs + ready, empty, s_ready, s_empty, acc_ready, acc_empty = bars + S: gl.constexpr = x_bufs.type.shape[0] + BM: gl.constexpr = x_bufs.type.shape[1] + KB: gl.constexpr = x_bufs.type.shape[2] + BN: gl.constexpr = w_bufs.type.shape[1] + HB: gl.constexpr = BN // 2 + SB: gl.constexpr = KB // 8 + SPC: gl.constexpr = 128 // KB + al: gl.constexpr = gl.DotOperandLayout(0, mma, 4) + bl: gl.constexpr = gl.DotOperandLayout(1, mma, 4) + offs_nl = gl.arange(0, HB, layout=gl.SliceLayout(1, mma)) + offs_tl = gl.arange(0, BM, layout=gl.SliceLayout(0, mma)) + g_in = gl.load(g_in_ptr) + num_pid_m, num_pid_n, num_tiles = _num_tiles(ntpp_ptr, N, BM, BN) + it = 0 + jt = 0 + hc = 0 + for t in range(gl.program_id(0), num_tiles, gl.num_programs(0)): + pid_m, pid_n = _tile(t, num_pid_m, num_pid_n, GROUP_SIZE_M, N_GROUP) + acc0 = gl.zeros([HB, BM], gl.float32, mma) + acc1 = gl.zeros([HB, BM], gl.float32, mma) + for jb in range(gl.cdiv(nk, SPC)): + j = jt + jb + ss = j % 2 + mbarrier.wait(s_ready.index(ss), j // 2 & 1) + for c in gl.static_range(SPC): + kk = jb * SPC + c + if not K_TAIL or kk < nk: + i = it + kk + st = i % S + mbarrier.wait(ready.index(st), i // S & 1) + b = x_bufs.index(st).permute((1, 0)).load(bl) + sx = sx_bufs.index(ss).slice(c * SB, SB, dim=1).load(xsl).to(gl.float8e4nv, bitcast=True) + a0 = w_bufs.index(st).slice(0, HB).load(al) + a1 = w_bufs.index(st).slice(HB, HB).load(al) + sws = sw_bufs.index(ss).slice(c * SB, SB, dim=1) + sw0 = sws.slice(0, HB).load(wsl).to(gl.float8e4nv, bitcast=True) + sw1 = sws.slice(HB, HB).load(wsl).to(gl.float8e4nv, bitcast=True) + mbarrier.arrive(empty.index(st)) + acc0 = mma_nvf4(a0, sw0, b, sx, acc0) + acc1 = mma_nvf4(a1, sw1, b, sx, acc1) + mbarrier.arrive(s_empty.index(ss)) + it += nk + jt += gl.cdiv(nk, SPC) + slot = gl.load(eid_ptr + pid_m) + route = gl.load(ids_ptr + pid_m * BM + offs_tl) + rmask = route < num_valid + xs_ = gl.load(tw_ptr + route * stride_tw, mask=rmask, other=0.0) * g_in + wrow = wg_ptr + slot.to(gl.int64) * stride_ge + n0 = pid_n * BN + offs_nl + wg0 = gl.load(wrow + n0, mask=n0 < N, other=0.0).to(gl.float32) + h0 = (acc0 * wg0[:, None] * xs_[None, :]).to(gl.bfloat16) + mbarrier.wait(acc_empty.index(0), (hc & 1) ^ 1) + h_buf.store(h0) + mbarrier.arrive(acc_ready.index(0)) + hc += 1 + n1 = n0 + HB + wg1 = gl.load(wrow + n1, mask=n1 < N, other=0.0).to(gl.float32) + h1 = (acc1 * wg1[:, None] * xs_[None, :]).to(gl.bfloat16) + mbarrier.wait(acc_empty.index(0), (hc & 1) ^ 1) + h_buf.store(h1) + mbarrier.arrive(acc_ready.index(0)) + hc += 1 + + +@gluon.jit +def _w4a4_down_persistent_kernel( + xq_ptr, xs_ptr, g_in_ptr, w_desc, ws_ptr, wg_ptr, out_ptr, topk_weights_ptr, + sorted_token_ids_ptr, expert_ids_ptr, num_tokens_post_padded_ptr, + N, K, num_valid_tokens, ks4, n_rows, stride_xq, stride_ge, stride_out, stride_tw, + BLOCK_SIZE_M: gl.constexpr, BLOCK_SIZE_N: gl.constexpr, BLOCK_SIZE_K: gl.constexpr, STAGES: gl.constexpr, + K_TAIL: gl.constexpr, N_GROUP: gl.constexpr, GROUP_SIZE_M: gl.constexpr, MMA_WARPS_N: gl.constexpr, + LOAD_WARPS: gl.constexpr, LOAD_REGS: gl.constexpr, EPI_WARPS: gl.constexpr, EPI_REGS: gl.constexpr, + top_k: gl.constexpr, wsl: gl.constexpr, xsl: gl.constexpr, +): + BM: gl.constexpr = BLOCK_SIZE_M + BN: gl.constexpr = BLOCK_SIZE_N + KB: gl.constexpr = BLOCK_SIZE_K // 2 + nk = K // BLOCK_SIZE_K + x_bufs = gl.allocate_shared_memory(gl.uint8, [STAGES, BM, KB], gl.NVMMASharedLayout.get_default_for([BM, KB], gl.uint8)) + w_bufs = gl.allocate_shared_memory(gl.uint8, [STAGES, BN, KB], w_desc.layout) + sw_bufs = gl.allocate_shared_memory(gl.uint8, [2, BN, 16], SCALE_SMEM_LAYOUT) + sx_bufs = gl.allocate_shared_memory(gl.uint8, [2, BM, 16], SCALE_SMEM_LAYOUT) + # n-contiguous and swizzled, so the epilogue's 8-row reads and the MMA-layout stores spread over banks + h_buf = gl.allocate_shared_memory(gl.bfloat16, [BN // 2, BM], gl.SwizzledSharedLayout(8, 1, 8, [0, 1])) + ready = gl.allocate_shared_memory(gl.int64, [STAGES, 1], mbarrier.MBarrierLayout()) + empty = gl.allocate_shared_memory(gl.int64, [STAGES, 1], mbarrier.MBarrierLayout()) + s_ready = gl.allocate_shared_memory(gl.int64, [2, 1], mbarrier.MBarrierLayout()) + s_empty = gl.allocate_shared_memory(gl.int64, [2, 1], mbarrier.MBarrierLayout()) + acc_ready = gl.allocate_shared_memory(gl.int64, [1, 1], mbarrier.MBarrierLayout()) + acc_empty = gl.allocate_shared_memory(gl.int64, [1, 1], mbarrier.MBarrierLayout()) + for i in gl.static_range(STAGES): + mbarrier.init(ready.index(i), count=1) + mbarrier.init(empty.index(i), count=1) + for i in gl.static_range(2): + mbarrier.init(s_ready.index(i), count=1) + mbarrier.init(s_empty.index(i), count=1) + mbarrier.init(acc_ready.index(0), count=1) + mbarrier.init(acc_empty.index(0), count=1) + + NW: gl.constexpr = gl.num_warps() + mma: gl.constexpr = gl.NVMMADistributedLayout(version=[2, 0], warps_per_cta=[MMA_WARPS_N, NW // MMA_WARPS_N], + instr_shape=[16, 8]) + gl.warp_specialize([ + (_down_mma_partition, ((x_bufs, w_bufs, sw_bufs, sx_bufs, h_buf), (ready, empty, s_ready, s_empty, acc_ready, acc_empty), + sorted_token_ids_ptr, expert_ids_ptr, num_tokens_post_padded_ptr, wg_ptr, g_in_ptr, + topk_weights_ptr, N, nk, num_valid_tokens, stride_ge, stride_tw, mma, wsl, xsl, K_TAIL, + N_GROUP, GROUP_SIZE_M)), + (_down_load_partition, (xq_ptr, xs_ptr, sorted_token_ids_ptr, expert_ids_ptr, num_tokens_post_padded_ptr, w_desc, + ws_ptr, (x_bufs, w_bufs, sw_bufs, sx_bufs, h_buf), + (ready, empty, s_ready, s_empty, acc_ready, acc_empty), N, nk, num_valid_tokens, + stride_xq, ks4, n_rows, K_TAIL, N_GROUP, GROUP_SIZE_M)), + (_down_epi_partition, (sorted_token_ids_ptr, num_tokens_post_padded_ptr, out_ptr, + (x_bufs, w_bufs, sw_bufs, sx_bufs, h_buf), (ready, empty, s_ready, s_empty, acc_ready, acc_empty), + N, num_valid_tokens, stride_out, top_k, N_GROUP, GROUP_SIZE_M)), + ], [LOAD_WARPS, EPI_WARPS], [LOAD_REGS, EPI_REGS]) + + +def supports(H: int, I: int) -> bool: + """Whether the GEMMs take a layer of this hidden / intermediate size.""" + # K is a whole number of 64-wide k blocks for both GEMMs; the gated tile pairs 64 gate rows with 64 up rows at least + return H % 64 == 0 and I % 64 == 0 + + +def _weight_desc(packed, box_n, BK): + S, N, KB = packed.shape + layout = gl.NVMMASharedLayout.get_default_for([box_n, BK // 2], gl.uint8) + return TensorDescriptor.from_tensor(packed.view(S * N, KB), [box_n, BK // 2], layout) + + +def _scale_layouts(cfg, w_rows): + nw, wn = cfg["num_warps"], cfg["MMA_WARPS_N"] + kg = cfg["BLOCK_SIZE_K"] // 16 + return scale_layout(0, w_rows, kg, [wn, nw // wn]), scale_layout(1, cfg["BLOCK_SIZE_M"], kg, [wn, nw // wn]) + + +def gate_up_gemm(xq, xs, g_in, packed, scale, glob, q, qs, g_out, sorted_ids, expert_ids, num_tokens_post_padded, + num_valid_tokens, top_k, activation, act_alpha, act_limit, cfg): + """q, qs = NVFP4 under g_out of act(gate, up), where [gate; up] = (xq * xs * g_in)[route // top_k] W4A4 packed[expert]^T * glob.""" + S, N, KB = packed.shape + K = KB * 2 + BM, BN, BK = cfg["BLOCK_SIZE_M"], cfg["BLOCK_SIZE_N"], cfg["BLOCK_SIZE_K"] + assert K % BK == 0 and (N // 2) % (BN // 2) == 0 and xs.stride(0) == K // 16 + assert q.stride(0) == N // 4 and qs.stride(0) == N // 32 + wsl, xsl = _scale_layouts(cfg, BN) + grid = triton.cdiv(sorted_ids.shape[0], BM) * triton.cdiv(N, BN) + _w4a4_moe_kernel[(grid,)]( + xq, xs.view(torch.uint8), g_in, _weight_desc(packed, BN // 2, BK), scale.view(torch.uint8), glob, None, None, + sorted_ids, expert_ids, num_tokens_post_padded, q, qs.view(torch.uint8), g_out, act_alpha, act_limit, + N, K, sorted_ids.shape[0], num_valid_tokens, K // 64, S * N, xq.stride(0), glob.stride(0), 0, 0, + K_TAIL=K % 256 != 0, GATED=True, ACT=ACTIVATIONS[activation], top_k=top_k, wsl=wsl, xsl=xsl, **cfg, + ) + + +def down_gemm(xq, xs, g_in, packed, scale, glob, out, topk_weights_flat, sorted_ids, expert_ids, num_tokens_post_padded, + num_valid_tokens, top_k, cfg): + """out[route // top_k] += (xq * xs * g_in)[route] W4A4 packed[expert]^T * glob * topk_weights_flat[route]; out starts zeroed.""" + S, N, KB = packed.shape + K = KB * 2 + cfg = dict(cfg) + persistent = cfg.pop("PERSISTENT") + BM, BN, BK = cfg["BLOCK_SIZE_M"], cfg["BLOCK_SIZE_N"], cfg["BLOCK_SIZE_K"] + assert K % BK == 0 and xs.stride(0) == K // 16 + args = (xq, xs.view(torch.uint8), g_in, _weight_desc(packed, BN, BK), scale.view(torch.uint8), glob, out, + topk_weights_flat, sorted_ids, expert_ids, num_tokens_post_padded) + if persistent: + wsl, xsl = _scale_layouts(cfg, BN // 2) + grid = torch.cuda.get_device_properties(out.device).multi_processor_count + _w4a4_down_persistent_kernel[(grid,)]( + *args, N, K, num_valid_tokens, K // 64, S * N, xq.stride(0), glob.stride(0), out.stride(0), + topk_weights_flat.stride(0), K_TAIL=K % 256 != 0, top_k=top_k, wsl=wsl, xsl=xsl, **cfg, + ) + return + wsl, xsl = _scale_layouts(cfg, BN) + grid = triton.cdiv(sorted_ids.shape[0], BM) * triton.cdiv(N, BN) + _w4a4_moe_kernel[(grid,)]( + *args, None, None, None, 1.0, 1.0, + N, K, sorted_ids.shape[0], num_valid_tokens, K // 64, S * N, xq.stride(0), glob.stride(0), out.stride(0), + topk_weights_flat.stride(0), K_TAIL=K % 256 != 0, GATED=False, ACT=0, top_k=top_k, wsl=wsl, xsl=xsl, **cfg, + ) diff --git a/python/freetoken/kernel/gluon/nvfp4/utils.py b/python/freetoken/kernel/gluon/nvfp4/utils.py new file mode 100644 index 0000000000..11426ec569 --- /dev/null +++ b/python/freetoken/kernel/gluon/nvfp4/utils.py @@ -0,0 +1,82 @@ +"""NVFP4 helpers shared by the Gluon kernels.""" + +from triton.experimental import gluon +from triton.experimental.gluon import language as gl +from triton.experimental.gluon.language import _core as ttgl +from triton.experimental.gluon.language._layouts import ( + DotOperandLayout, + NVMMADistributedLayout, +) + + +@gluon.jit +def bf16x2_mul(a, b): + # sm_90+; sm_80-89 need fma.rn.bf16x2 with a -0.0 addend + return gl.inline_asm_elementwise("mul.rn.bf16x2 $0, $1, $2;", "=r,r,r", [a, b], gl.uint32, True, 1) + + +@gluon.jit +def prmt(a, b, sel: gl.constexpr): + return gl.inline_asm_elementwise(f"prmt.b32 $0, $1, $2, {sel};", "=r,r,r", [a, b], gl.uint32, True, 1) + + +@gluon.jit +def scale_bf16x2(s): + """e4m3 scale bytes -> bf16x2 words holding scale * 2^118 twice.""" + # 2^118 * the 2^-126 of e2m1x8_bf16x2 leaves 2^-8 for the caller to undo; every step is exact in bf16 + s = (s.to(gl.float8e4nv, bitcast=True).to(gl.float32) * 3.3230699894622794e35).to(gl.bfloat16) + s = s.to(gl.uint16, bitcast=True).to(gl.uint32) + return s | (s << 16) + + +@gluon.jit +def e2m1x8_bf16x2(q, s2): + """8 e2m1 codes in a uint32 -> 4 bf16x2 words in k order, each value times the bf16x2 scale s2 and 2^-126.""" + # each 16-bit half's top nibble lands in a bf16 as sign + exponent/mantissa bits, a value of e2m1 * 2^-126 + o0 = (q & 0x80008000) | ((q >> 6) & 0x01C001C0) + q1 = q << 4 + o1 = (q1 & 0x80008000) | ((q1 >> 6) & 0x01C001C0) + q2 = q << 8 + o2 = (q2 & 0x80008000) | ((q2 >> 6) & 0x01C001C0) + q3 = q << 12 + o3 = (q3 & 0x80008000) | ((q3 >> 6) & 0x01C001C0) + # o0..o3 hold the k pairs (3, 7), (2, 6), (1, 5), (0, 4) + w0 = bf16x2_mul(prmt(o3, o2, 0x5410), s2) + w1 = bf16x2_mul(prmt(o1, o0, 0x5410), s2) + w2 = bf16x2_mul(prmt(o3, o2, 0x7632), s2) + w3 = bf16x2_mul(prmt(o1, o0, 0x7632), s2) + return w0, w1, w2, w3 + + +def scale_layout(op_idx: int, rows: int, kgroups: int, mma_warps: list[int]) -> gl.DistributedLinearLayout: + """Layout of a [rows, K / 16] e4m3 scale tile for operand op_idx of an NVFP4 mma_v2; mirrors Triton's getSM120DotScaledScaleLayout.""" + def bits(n): + return n.bit_length() - 1 + + w0, w1 = mma_warps + reg = [[0, 1 << i] for i in range(bits(kgroups))] + if op_idx == 0: + lane = [[8, 0], [0, 0], [1, 0], [2, 0], [4, 0]] + warp = [[0, 0]] * bits(w1) + [[16 << i, 0] for i in range(bits(w0))] + cover = 16 * w0 + else: + lane = [[0, 0], [0, 0], [1, 0], [2, 0], [4, 0]] + warp = [[8 << i, 0] for i in range(bits(w1))] + [[0, 0]] * bits(w0) + cover = 8 * w1 + reg += [[cover << i, 0] for i in range(max(0, bits(rows) - bits(cover)))] + # bases past the tile's rows broadcast + lane = [[b[0] if b[0] < rows else 0, b[1]] for b in lane] + warp = [[b[0] if b[0] < rows else 0, b[1]] for b in warp] + return gl.DistributedLinearLayout(reg, lane, warp, [], [rows, kgroups]) + + +@ttgl.builtin +def mma_nvf4(a, a_scale, b, b_scale, acc, _semantic=None): + """acc += (a * a_scale) @ (b * b_scale) on sm_120 (mma kind::mxf4nvf4, scale_vec::4X). + + a [M, K/2] and b [K/2, N] are packed e2m1 in DotOperandLayout(k_width=4) of acc's mma_v2 layout; the e4m3 scales are [M, K/16] / [N, K/16] in scale_layout. + """ + assert isinstance(acc.type.layout, NVMMADistributedLayout) and acc.type.layout.version == [2, 0] + assert isinstance(a.type.layout, DotOperandLayout) and isinstance(b.type.layout, DotOperandLayout) + out = _semantic.dot_scaled(a, a_scale, "e2m1", b, b_scale, "e2m1", acc, False, True, True, ttgl.float32) + return ttgl.tensor(out.handle, acc.type) diff --git a/python/freetoken/kernel/moe_impl.py b/python/freetoken/kernel/moe_impl.py index ef569fa2e6..488f8c2bdf 100644 --- a/python/freetoken/kernel/moe_impl.py +++ b/python/freetoken/kernel/moe_impl.py @@ -34,7 +34,12 @@ def fused_moe_kernel_triton( even_Ks = True else: even_Ks = False - dtype = tl.bfloat16 if compute_type == torch.bfloat16 else tl.float16 + if compute_type == torch.float32: + dtype = tl.float32 + elif compute_type == torch.bfloat16: + dtype = tl.bfloat16 + else: + dtype = tl.float16 fused_moe_kernel[grid]( A, B, @@ -484,12 +489,15 @@ def mxfp4_splitk_gemv_triton( k_groups = K // 32 kgps = triton.cdiv(k_groups, num_splits) lut = get_fp4_lut(x.device) + assert N % 2 == 0 and block_n % 2 == 0, "the kernel reads byte pairs of adjacent columns" + w16 = w_blocks_t.view(torch.int16) + s16 = w_scales_t.view(torch.int16) grid = (triton.cdiv(N, block_n), routes * num_splits) mxfp4_splitk_gemv_kernel[grid]( - x, w_blocks_t, w_scales_t, bias_arg, expert_ids, partial, lut, N, K, - stride_xe, w_blocks_t.stride(0), w_blocks_t.stride(1), - w_scales_t.stride(0), w_scales_t.stride(1), bias_stride, N, + x, w16, s16, bias_arg, expert_ids, partial, lut, N, K, + stride_xe, w16.stride(0), w16.stride(1), + s16.stride(0), s16.stride(1), bias_stride, N, HAS_BIAS=has_bias, BLOCK_N=block_n, # type: ignore NUM_K_SPLITS=num_splits, K_GROUPS_PER_SPLIT=kgps, # type: ignore num_warps=num_warps, diff --git a/python/freetoken/kernel/nvfp4_moe.py b/python/freetoken/kernel/nvfp4_moe.py new file mode 100644 index 0000000000..df89410726 --- /dev/null +++ b/python/freetoken/kernel/nvfp4_moe.py @@ -0,0 +1,460 @@ +"""Kernel dispatch for the NVFP4 MoE GEMMs: which kernel runs, with which tiles and switches, and its launch. + +The kernels are templates (:mod:`freetoken.kernel.triton.nvfp4_fused_moe`, :mod:`freetoken.kernel.gluon.nvfp4.moe.sm120`, :mod:`freetoken.kernel.gluon.nvfp4.moe.sm120_w4a4`); +:mod:`freetoken.moe.fused_nvfp4` strings the GEMMs into the MoE layer. +""" + +from __future__ import annotations + +import functools +import re +from dataclasses import dataclass, replace +from typing import Any, Dict, Tuple + +import torch +import triton +import triton.language as tl +from triton.tools.tensor_descriptor import TensorDescriptor + +from freetoken.kernel.backend import is_rocm +from freetoken.kernel.triton.e4m3_compat import e4m3_kernel_view, e4m3_native +from freetoken.kernel.triton.nvfp4_fused_moe import ( + _decode_nvfp4_marlin_kernel, + _decode_nvfp4_moe_kernel, + _e2m1_lut, + _prefill_nvfp4_moe_kernel, +) +from freetoken.utils import init_logger + +logger = init_logger(__name__) + +# Decode is captured into a CUDA graph, so the config must be fixed (no triton.autotune, +# which benchmarks at run time). Tuned offline against the NVFP4 decode kernels. +# These drive the original LUT-gather decode (decode_gemm), kept only for A/B. +_DECODE_BLOCK_N = 64 +_DECODE_BLOCK_KB = 128 +_DECODE_WARPS = 4 + +# Marlin-style decode: every K-loop load must keep the word tile's register layout, so BLOCK_N * BLOCK_KW <= 128 * num_warps and BLOCK_KW <= 32 at 4 warps. +# Long K (gate_up) uses the narrow N tile for more programs at M=1; short K (down) uses the wide one. +_DECODE_MARLIN_BLOCK_N = 16 +_DECODE_MARLIN_BLOCK_KW = 16 +_DECODE_MARLIN_UNROLL = 2 +_DECODE_MARLIN_WARPS = 4 +_DECODE_MARLIN_DEEPK_BLOCK_N = 8 +_DECODE_MARLIN_DEEPK_BLOCK_KW = 32 +_DECODE_MARLIN_DEEPK_UNROLL = 2 +_DECODE_MARLIN_DEEPK_THRESHOLD = 2048 + + +def _tl_dtype(dt: torch.dtype): + if dt == torch.bfloat16: + return tl.bfloat16 + if dt == torch.float16: + return tl.float16 + return tl.float32 + + +def decode_gemm( + a: torch.Tensor, + packed: torch.Tensor, + scale: torch.Tensor, + glob: torch.Tensor, + c: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + mul_routed_weight: bool, + a_row_is_route: bool, +) -> None: + M, top_k = topk_ids.shape + N = packed.shape[1] + K = packed.shape[2] * 2 + scale = e4m3_kernel_view(scale) + total_routes = M * top_k + grid = (total_routes, triton.cdiv(N, _DECODE_BLOCK_N)) + _decode_nvfp4_moe_kernel[grid]( + a, packed, scale, glob, c, topk_weights, topk_ids, + _e2m1_lut(a.device.index), + total_routes, N, K, + a.stride(0), a.stride(1), + packed.stride(0), packed.stride(1), packed.stride(2), + scale.stride(0), scale.stride(1), scale.stride(2), + glob.stride(0), glob.stride(1), + c.stride(0), c.stride(1), c.stride(2), + topk_weights.stride(0), topk_weights.stride(1), + topk_ids.stride(0), topk_ids.stride(1), + BLOCK_SIZE_N=_DECODE_BLOCK_N, + BLOCK_SIZE_KB=_DECODE_BLOCK_KB, + TOP_K=top_k, + A_ROW_IS_ROUTE=a_row_is_route, + MUL_ROUTED_WEIGHT=mul_routed_weight, + compute_type=_tl_dtype(c.dtype), + num_warps=_DECODE_WARPS, + ) + + +def decode_gemm_marlin( + a: torch.Tensor, + packed: torch.Tensor, + scale: torch.Tensor, + glob: torch.Tensor, + c: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + mul_routed_weight: bool, + a_row_is_route: bool, +) -> None: + """Marlin-style decode GEMV: int32 wide loads + deferred reduction + (:func:`_decode_nvfp4_marlin_kernel`). ``packed`` is the uint8 ``[S, N, K//2]`` bank; + it is reinterpreted as int32 ``[S, N, K//8]`` (contiguous, K%8==0 for NVFP4).""" + M, top_k = topk_ids.shape + N = packed.shape[1] + K = packed.shape[2] * 2 + packed_i32 = packed.view(torch.int32) # [S, N, K // 8] + scale = e4m3_kernel_view(scale) + total_routes = M * top_k + deep_k = K >= _DECODE_MARLIN_DEEPK_THRESHOLD + block_n = _DECODE_MARLIN_DEEPK_BLOCK_N if deep_k else _DECODE_MARLIN_BLOCK_N + block_kw = _DECODE_MARLIN_DEEPK_BLOCK_KW if deep_k else _DECODE_MARLIN_BLOCK_KW + unroll = _DECODE_MARLIN_DEEPK_UNROLL if deep_k else _DECODE_MARLIN_UNROLL + # int64 activation loads need contiguous 16-bit rows that start 16-byte aligned + a_wide = ( + a.dtype in (torch.bfloat16, torch.float16) + and a.stride(1) == 1 and a.stride(0) % 8 == 0 and a.data_ptr() % 16 == 0 + ) + grid = (total_routes, triton.cdiv(N, block_n)) + _decode_nvfp4_marlin_kernel[grid]( + a, packed_i32, scale, glob, c, topk_weights, topk_ids, + total_routes, N, K, + a.stride(0), a.stride(1), + packed_i32.stride(0), packed_i32.stride(1), packed_i32.stride(2), + scale.stride(0), scale.stride(1), scale.stride(2), + glob.stride(0), glob.stride(1), + c.stride(0), c.stride(1), c.stride(2), + topk_weights.stride(0), topk_weights.stride(1), + topk_ids.stride(0), topk_ids.stride(1), + BLOCK_SIZE_N=block_n, + BLOCK_SIZE_KW=block_kw, + UNROLL=unroll, + TOP_K=top_k, + A_ROW_IS_ROUTE=a_row_is_route, + A_WIDE=a_wide, + MUL_ROUTED_WEIGHT=mul_routed_weight, + compute_type=_tl_dtype(c.dtype), + num_warps=_DECODE_MARLIN_WARPS, + ) + + +# rho = bf16 tensor-core FLOP / fp32 ALU FLOP per SM per clock, with fp32 accumulate (the only accumulator for bf16 MMA). +# Each M block decodes its weight tile once on the ALUs, and the MMA also runs on the padding up to BLOCK_SIZE_M. +# High rho makes decode the cost, so big blocks win; low rho makes padded MMA the cost, so small blocks win. +# Spec, FLOP/clk/SM: GeForce RTX 30/40/50 512 / 256 = 2, as GeForce runs fp32 accumulate at half rate. +# RTX A, RTX Ada Generation, RTX PRO Blackwell, A10, A40, L4, L40S 1024 / 256 = 4; A100 2048 / 128 = 16; H100 4096 / 256 = 16; H20 148 / 44 TFLOP/s = 3.4. +# Measured (cuBLAS bf16 / Triton fp32 FMA): RTX PRO 6000 4.5, RTX 5090 2.3, RTX 4090 1.9, RTX 3080 Ti 1.6. +_PRO_GPU_RE = re.compile( + r"RTX\s*A\d{3,4}" # RTX A400 .. A6000 + r"|RTX\s*\d{3,4}E?\s*(SFF\s*)?-?\s*Ada" # RTX 500 .. 6000 Ada Generation + r"|RTX\s*PRO\s*\d{3,4}" # RTX PRO 500 .. 6000 Blackwell + r"|RTX\s*\d{4}D\b" # RTX 6000D + r"|Quadro|TITAN" + r"|^(NVIDIA\s+)?(GRID\s+)?(A2|A10[GM]?|A16|A40|L2|L20|L4|L40S?)\b", + re.I, +) +# also matches names without "GeForce", e.g. "NVIDIA RTX 4090" +_GEFORCE_GPU_RE = re.compile(r"GeForce|RTX\s*[2-6]0[5-9]0(?!\d)", re.I) + + +def _tensor_rho(name: str, major: int, minor: int) -> float: + if major not in (8, 12) or (major, minor) == (8, 0): + return 3.4 if re.search(r"\bH20\b", name) else 16.0 + if _PRO_GPU_RE.search(name): + return 4.0 + if _GEFORCE_GPU_RE.search(name): + return 2.0 + logger.warning("unknown GPU %r: NVFP4 MoE prefill assumes full-rate fp32-accumulate tensor cores", name) + return 4.0 + + +# BLOCK_SIZE_M -> (BLOCK_SIZE_N, BLOCK_SIZE_KB, num_stages) +_PREFILL_TILES_H100 = {16: (64, 64, 3), 32: (64, 32, 3), 64: (64, 32, 3), 128: (64, 32, 3)} +_PREFILL_TILES_RTX_PRO = {16: (64, 64, 4), 32: (128, 64, 3), 64: (128, 64, 2), 128: (64, 32, 2)} +_PREFILL_TILES_GEFORCE = {16: (64, 64, 4), 32: (128, 64, 2), 64: (128, 32, 2)} + + +@dataclass(frozen=True) +class _PrefillTuning: + """Triton prefill GEMM tuning of one GPU class.""" + + steps: Tuple[float, float, float] # tokens per expert where BLOCK_SIZE_M becomes 32, 64, 128 + tiles: Dict[int, Tuple[int, int, int]] + pair_decode_below: float = 0.0 # tokens per expert under which FAST_DECODE stays 0 + small_tile_decode: int = 2 # FAST_DECODE of BLOCK_SIZE_M <= 32 + large_tile_decode: int = 1 # FAST_DECODE of BLOCK_SIZE_M >= 64 + scale_words: bool = False # SCALE_WORDS on BLOCK_SIZE_M 64 word-decode tiles + long_k_warps: int = 4 # num_warps of SCALE_WORDS GEMMs with K >= _LONG_K + + +_INF = float("inf") +_LONG_K = 2048 +_TUNING_H100 = _PrefillTuning((12, 24, 64), _PREFILL_TILES_H100) +_TUNING_RTX_PRO = _PrefillTuning((16, 36, 150), _PREFILL_TILES_RTX_PRO) +_TUNING_GEFORCE = _PrefillTuning((16, 150, _INF), _PREFILL_TILES_GEFORCE) +# RTX 4090 A/B: up to 32 tokens per expert the pair decode streams the weights 2-18% faster, from 51 words win 11-14%, also on 32-row tiles +_SM89_DECODE = dict(pair_decode_below=48, small_tile_decode=1) +# H100 A/B: bytes beat words at BLOCK_SIZE_M 64 / 128 by 1-17% per layer (24-31% on Gemma-4-26B-A4B's K = 704 down GEMM) +_SM90_DECODE = dict(small_tile_decode=2, large_tile_decode=2) +# RTX 4090 / 3080 Ti layer A/B: scale words + 8 warps on the long-K GEMM give 1.03-1.09x / 1.05-1.19x; with them sm_86 gains from 64-row tiles at 100 tokens per expert, sm_89 does not +_GEFORCE_SCALE_WORDS = dict(scale_words=True, long_k_warps=8) +# (rho class, compute capability) -> tuning, from sweeps on H100, RTX PRO 6000 and RTX 5090 / 4090 / 3080 Ti; (class, None) covers the rest of the class. +# The RTX tiles come from cc 8.6 / 8.9 / 12.0 sweeps (mma.sync, 99 KB shared memory); other GPUs keep the H100 tiles. +_PREFILL_TUNING = { + ("wide", (9, 0)): replace(_TUNING_H100, **_SM90_DECODE), + ("wide", None): _TUNING_H100, + ("pro", (8, 6)): _TUNING_RTX_PRO, + ("pro", (8, 9)): replace(_TUNING_RTX_PRO, **_SM89_DECODE), + ("pro", (9, 0)): replace(_TUNING_RTX_PRO, tiles=_PREFILL_TILES_H100, **_SM90_DECODE), + ("pro", (12, 0)): _TUNING_RTX_PRO, + ("pro", None): replace(_TUNING_RTX_PRO, tiles=_PREFILL_TILES_H100), + ("geforce", (8, 6)): replace(_TUNING_GEFORCE, steps=(16, 100, _INF), **_GEFORCE_SCALE_WORDS), + ("geforce", (8, 9)): replace(_TUNING_GEFORCE, **_SM89_DECODE, **_GEFORCE_SCALE_WORDS), + ("geforce", (12, 0)): _TUNING_GEFORCE, + ("geforce", None): replace(_TUNING_GEFORCE, tiles=_PREFILL_TILES_H100), +} +# steps while the Gluon kernel runs BLOCK_SIZE_M >= 64, from an RTX PRO 6000 sweep; classes without an entry keep the Triton steps +_PREFILL_BLOCK_M_STEPS_GLUON = ((3.0, (16, 36, 96)),) + + +# Gluon prefill GEMM tiles: BLOCK_SIZE_N -> (BLOCK_SIZE_K, setmaxnreg of the loader / decode warps) +_GLUON_PREFILL_TILES = {128: (64, 96, 96), 256: (32, 88, 72)} + + +def _gluon_prefill_config(N: int, block_m: int, mul_routed_weight: bool) -> Dict[str, Any]: + # BLOCK_SIZE_M 64 tiles are short (few tokens per expert, often short K): persistent CTAs overlap one tile's epilogue and ring fill with the next tile, and a third x / decoded-weight stage absorbs DRAM latency + persistent = block_m == 64 + # a 256-row tile halves the token loads per weight byte, but a partial last tile still decodes all 256 rows + # the persistent epilogue keeps one 64-bit address per element in registers, which spills with the routed-weight loads + block_n = 256 if N % 256 == 0 and not (persistent and mul_routed_weight) else 128 + block_k, load_regs, decode_regs = _GLUON_PREFILL_TILES[block_n] + return dict(BLOCK_SIZE_M=block_m, BLOCK_SIZE_N=block_n, BLOCK_SIZE_K=block_k, GROUP_SIZE_M=8, + STAGES=3 if persistent else 2, W_STAGES=4, PERSISTENT=persistent, DECODE_WARPS=4, MMA_WARPS_N=4, + LOAD_REGS=load_regs, DECODE_REGS=decode_regs, num_warps=8) + + +@functools.lru_cache(maxsize=None) +def _gluon_prefill(device_index: int): + """The Gluon module whose prefill GEMM replaces the Triton kernel on this GPU, or None.""" + from freetoken.kernel import gluon + + if torch.cuda.get_device_capability(device_index) != (12, 0) or not gluon.available(): + return None + from freetoken.kernel.gluon.nvfp4.moe import sm120 + + return sm120 + + +@functools.lru_cache(maxsize=None) +def _device_rho(device_index: int) -> float: + props = torch.cuda.get_device_properties(device_index) + return _tensor_rho(props.name, props.major, props.minor) + + +@functools.lru_cache(maxsize=None) +def _prefill_tuning(device_index: int) -> _PrefillTuning: + if is_rocm(): + # TODO: tune for AMD GPUs; they keep the H100 tuning + return _TUNING_H100 + rho = _device_rho(device_index) + rho_class = "wide" if rho >= 8.0 else "pro" if rho >= 3.0 else "geforce" + cc = torch.cuda.get_device_capability(device_index) + return _PREFILL_TUNING.get((rho_class, cc), _PREFILL_TUNING[(rho_class, None)]) + + +@functools.lru_cache(maxsize=None) +def _prefill_block_m_steps(device_index: int) -> Tuple[float, float, float]: + steps = _prefill_tuning(device_index).steps + if _gluon_prefill(device_index) is not None: + rho = _device_rho(device_index) + steps = next((s for min_rho, s in _PREFILL_BLOCK_M_STEPS_GLUON if rho >= min_rho), steps) + return steps + + +@functools.lru_cache(maxsize=None) +def _tma_device(device_index: int) -> bool: + return not is_rocm() and torch.cuda.get_device_capability(device_index) == (12, 0) + + +def _prefill_use_tma(packed: torch.Tensor, scale: torch.Tensor, K: int, device_index: int) -> bool: + # a TMA box row is >= 16 bytes and every global stride a multiple of 16 bytes, so K // 16 scale bytes need K % 256 == 0; + # the TMA scale path reads the scales as fp8, so the forced e4m3 emulation stays on pointer loads + return (_tma_device(device_index) and e4m3_native() and K % 256 == 0 + and packed.is_contiguous() and scale.is_contiguous() + and packed.data_ptr() % 16 == 0 and scale.data_ptr() % 16 == 0) + + +def _prefill_use_scale_words(scale: torch.Tensor, K: int, cfg: Dict[str, Any], device_index: int) -> bool: + # K // 16 scale bytes per row split into whole int32 words, and a K step reads at least one word + return (_prefill_tuning(device_index).scale_words and cfg["BLOCK_SIZE_M"] == 64 and cfg["BLOCK_SIZE_KB"] >= 32 + and K % 64 == 0 and scale.is_contiguous() and scale.data_ptr() % 4 == 0) + + +def _prefill_fast_decode(block_m: int, per_expert: float, device_index: int) -> int: + """FAST_DECODE of the prefill kernel: 0 keeps the fp16 pair decode (ROCm: the fast decode is PTX inline asm), 1 words, 2 bytes.""" + if is_rocm(): + return 0 + tuning = _prefill_tuning(device_index) + if per_expert < tuning.pair_decode_below: + return 0 + # the word decode halves the loop but stages the decoded tile through shared memory; small tiles keep the byte form + return tuning.small_tile_decode if block_m <= 32 else tuning.large_tile_decode + + +def prefill_config(M: int, top_k: int, num_experts: int, device_index: int) -> Dict[str, int]: + # moe_align_block_size pads each expert's token count to a multiple of BLOCK_SIZE_M, so triton.autotune cannot pick it. + to32, to64, to128 = _prefill_block_m_steps(device_index) + per_expert = M * top_k / num_experts + block_m = 128 if per_expert >= to128 else 64 if per_expert >= to64 else 32 if per_expert >= to32 else 16 + block_n, block_kb, stages = _prefill_tuning(device_index).tiles[block_m] + return dict(BLOCK_SIZE_M=block_m, BLOCK_SIZE_N=block_n, BLOCK_SIZE_KB=block_kb, + GROUP_SIZE_M=8, num_warps=4, num_stages=stages, + FAST_DECODE=_prefill_fast_decode(block_m, per_expert, device_index)) + + +def prefill_gemm( + a: torch.Tensor, + packed: torch.Tensor, + scale: torch.Tensor, + glob: torch.Tensor, + c: torch.Tensor, + topk_weights_flat: torch.Tensor, + sorted_ids: torch.Tensor, + expert_ids: torch.Tensor, + num_tokens_post_padded: torch.Tensor, + num_valid_tokens: int, + kernel_top_k: int, + mul_routed_weight: bool, + cfg: Dict[str, Any], +) -> None: + gluon = _gluon_prefill(a.device.index) + if gluon is not None and cfg["BLOCK_SIZE_M"] >= 64 and gluon.supports(a, packed, scale, c): + gluon.prefill_gemm(a, packed, scale, glob, c, topk_weights_flat, sorted_ids, expert_ids, num_tokens_post_padded, + num_valid_tokens, kernel_top_k, mul_routed_weight, + _gluon_prefill_config(packed.shape[1], cfg["BLOCK_SIZE_M"], mul_routed_weight)) + return + N = packed.shape[1] + K = packed.shape[2] * 2 + EM = sorted_ids.shape[0] + scale = e4m3_kernel_view(scale) + grid = lambda META: ( # noqa: E731 + triton.cdiv(EM, META["BLOCK_SIZE_M"]) * triton.cdiv(N, META["BLOCK_SIZE_N"]), + ) + cfg = dict(cfg) + fast = cfg.pop("FAST_DECODE", 0) + if K % 8: + fast = 0 + if fast == 1: + packed = packed.view(torch.int32) + use_tma = _prefill_use_tma(packed, scale, K, a.device.index) + # the 16-byte scale box costs registers that the BLOCK_SIZE_M >= 64 tiles do not have + tma_scales = cfg["BLOCK_SIZE_M"] <= 32 + scale_words = not use_tma and fast == 1 and _prefill_use_scale_words(scale, K, cfg, a.device.index) + if scale_words: + scale = scale.view(torch.int32) + if K >= _LONG_K: + cfg["num_warps"] = _prefill_tuning(a.device.index).long_k_warps + w_desc = s_desc = None + if use_tma: + rows = packed.shape[0] * N + box_k = cfg["BLOCK_SIZE_KB"] // 4 if fast == 1 else cfg["BLOCK_SIZE_KB"] + w_desc = TensorDescriptor.from_tensor(packed.view(rows, packed.shape[2]), [cfg["BLOCK_SIZE_N"], box_k]) + if tma_scales: + s_desc = TensorDescriptor.from_tensor(scale.view(rows, K // 16), [cfg["BLOCK_SIZE_N"], 16]) + _prefill_nvfp4_moe_kernel[grid]( + a, packed, scale, glob, c, topk_weights_flat, sorted_ids, expert_ids, + num_tokens_post_padded, + w_desc, s_desc, + N, K, EM, num_valid_tokens, + a.stride(0), a.stride(1), + packed.stride(0), packed.stride(1), packed.stride(2), + scale.stride(0), scale.stride(1), scale.stride(2), + glob.stride(0), glob.stride(1), + c.stride(1), c.stride(2), + topk_weights_flat.stride(0), + MUL_ROUTED_WEIGHT=mul_routed_weight, + top_k=kernel_top_k, + compute_type=_tl_dtype(c.dtype), + USE_TMA=use_tma, + TMA_SCALES=tma_scales, + FAST_DECODE=fast, + SCALE_WORDS=scale_words, + **cfg, + ) + + +@functools.lru_cache(maxsize=None) +def _w4a4_prefill(device_index: int): + """The Gluon module of the W4A4 prefill GEMMs on this GPU, or None.""" + from freetoken.kernel import gluon + + if torch.cuda.get_device_capability(device_index) != (12, 0) or not gluon.available(): + return None + from freetoken.kernel.gluon.nvfp4.moe import sm120_w4a4 + + return sm120_w4a4 + + +def w4a4_available(device_index: int) -> bool: + return _w4a4_prefill(device_index) is not None + + +# Below this many tokens per expert both prefills stream the weights at the same speed (RTX PRO 6000 and RTX 5090 sweeps over the +# NVFP4 MoE shapes); W4A16 keeps the exact activations there +_W4A4_MIN_TOKENS_PER_EXPERT = 4 +# (min rho, tokens per expert where BLOCK_SIZE_M becomes 128), from the same sweeps +_W4A4_BLOCK_M_STEPS = ((3.0, 128), (0.0, 64)) + + +@functools.lru_cache(maxsize=None) +def _w4a4_device(device_index: int) -> Tuple[float, int]: + props = torch.cuda.get_device_properties(device_index) + rho = _tensor_rho(props.name, props.major, props.minor) + return next(t for min_rho, t in _W4A4_BLOCK_M_STEPS if rho >= min_rho), props.L2_cache_size + + +def w4a4_prefill_config(M: int, top_k: int, num_experts: int, H: int, I: int, activation: str, dtype: torch.dtype, + device_index: int) -> Dict[str, Any] | None: + """Tiles of the W4A4 prefill GEMMs, or None where the W4A16 prefill runs instead.""" + w4a4 = _w4a4_prefill(device_index) + if w4a4 is None or dtype != torch.bfloat16 or activation not in w4a4.ACTIVATIONS or not w4a4.supports(H, I): + return None + to128, l2_bytes = _w4a4_device(device_index) + per_expert = M * top_k / num_experts + if per_expert < _W4A4_MIN_TOKENS_PER_EXPERT: + return None + block_m = 128 if per_expert >= to128 else 64 + common = dict(BLOCK_SIZE_M=block_m, GROUP_SIZE_M=8, MMA_WARPS_N=4, LOAD_WARPS=4, num_warps=8) + # the gated tile holds BN / 2 gate rows and their up rows, so I must split into whole half tiles + gate_up = dict(common, BLOCK_SIZE_N=256 if I % 128 == 0 else 128, BLOCK_SIZE_K=128 if H % 128 == 0 else 64, + N_GROUP=0, STAGES=3, LOAD_REGS=88) + if I < 1024: + # short K: the persistent kernel overlaps a tile's atomics with the next mainloop; N strips keep the [M, H] output + # in L2 once it outgrows half of it + down = dict(common, BLOCK_SIZE_N=256, BLOCK_SIZE_K=64, STAGES=4, N_GROUP=2 if M * H * 2 > l2_bytes // 2 else 0, + LOAD_REGS=56, EPI_WARPS=4, EPI_REGS=72, PERSISTENT=True) + else: + down = dict(common, BLOCK_SIZE_N=256 if H % 256 == 0 else 128, BLOCK_SIZE_K=128 if I % 128 == 0 else 64, + N_GROUP=0, STAGES=3, LOAD_REGS=88, PERSISTENT=False) + return dict(BLOCK_SIZE_M=block_m, gate_up=gate_up, down=down) + + +def w4a4_gate_up_gemm(xq, xs, g_in, packed, scale, glob, q, qs, g_out, sorted_ids, expert_ids, num_tokens_post_padded, + num_valid_tokens, top_k, activation, act_alpha, act_limit, cfg): + _w4a4_prefill(xq.device.index).gate_up_gemm( + xq, xs, g_in, packed, scale, glob, q, qs, g_out, sorted_ids, expert_ids, num_tokens_post_padded, + num_valid_tokens, top_k, activation, act_alpha, act_limit, cfg) + + +def w4a4_down_gemm(xq, xs, g_in, packed, scale, glob, out, topk_weights_flat, sorted_ids, expert_ids, + num_tokens_post_padded, num_valid_tokens, top_k, cfg): + _w4a4_prefill(xq.device.index).down_gemm( + xq, xs, g_in, packed, scale, glob, out, topk_weights_flat, sorted_ids, expert_ids, num_tokens_post_padded, + num_valid_tokens, top_k, cfg) diff --git a/python/freetoken/kernel/pynccl.py b/python/freetoken/kernel/pynccl.py index 23ea57351d..71bde7342a 100644 --- a/python/freetoken/kernel/pynccl.py +++ b/python/freetoken/kernel/pynccl.py @@ -27,6 +27,7 @@ def get_buffer(self) -> int: ... @functools.cache def _load_nccl_module() -> Module: + # TODO(ROCm): NCCL -> RCCL migration for multi-GPU tensor parallelism on AMD. return load_aot("pynccl", cuda_files=["pynccl.cu"], extra_ldflags=["-lnccl"]) diff --git a/python/freetoken/kernel/row_store.py b/python/freetoken/kernel/row_store.py new file mode 100644 index 0000000000..5b42479ebc --- /dev/null +++ b/python/freetoken/kernel/row_store.py @@ -0,0 +1,54 @@ +"""Python face of the disk row store extension (``kernel/csrc/row_store``) and its CUDA-graph sync. + +``RowStore`` reads fixed-width rows straight out of checkpoint shard files (io_uring or a pread pool, +O_DIRECT) into pinned staging by row id; ``PleStore`` adds PLE's n-gram hashing. + +Graph sync: a captured graph cannot wait for host I/O, so a captured ``lookup`` WAITs on a pinned +flag through a stream memop (``wait_reset``) that the host ``signal``s once the fill landed; where +the driver rejects memops in capture, the caller falls back to launch-gating (fill before replay). +``probe_wait_sync`` decides which, once, at startup. +""" + +from __future__ import annotations + +import torch + +from freetoken.kernel import _row_store +from freetoken.kernel.pinned import alloc_pinned_tensor + +RowStore = _row_store.RowStore +PleStore = _row_store.PleStore + + +def probe_wait_sync(mode: str, device: torch.device) -> bool: + """``mode``: ``auto`` (probe), ``wait`` (require memops) or ``gate`` (never use them).""" + if mode not in ("auto", "wait", "gate"): + raise ValueError(f"unknown row-store sync mode {mode!r}; expected auto, wait or gate") + if mode == "gate": + return False + scratch = alloc_pinned_tensor(1, dtype=torch.int64) + scratch.zero_() + stream = torch.cuda.current_stream(device) + ok = ( + _row_store.memop_write(stream.cuda_stream, scratch.data_ptr(), 7) == 0 + and _row_store.memop_wait_geq(stream.cuda_stream, scratch.data_ptr(), 7) == 0 + ) + if ok: + stream.synchronize() + ok = int(scratch[0]) == 7 + if mode == "wait" and not ok: + raise RuntimeError("wait-sync requested but stream memops are unavailable") + return ok + + +def wait_reset(stream: torch.cuda.Stream, flag: torch.Tensor) -> None: + """Enqueue WAIT(flag >= 1) then flag := 0 on ``stream`` (capturable).""" + _row_store.memop_wait_reset(stream.cuda_stream, flag.data_ptr()) + + +def signal(flag: torch.Tensor) -> None: + """Host side: release a pending WAIT (a release-store of 1).""" + _row_store.signal_flag(flag.data_ptr()) + + +__all__ = ["RowStore", "PleStore", "probe_wait_sync", "wait_reset", "signal"] diff --git a/python/freetoken/kernel/triton/activation.py b/python/freetoken/kernel/triton/activation.py index 0e4d28c8fc..9645960d11 100644 --- a/python/freetoken/kernel/triton/activation.py +++ b/python/freetoken/kernel/triton/activation.py @@ -20,8 +20,9 @@ import triton.language as tl from triton.language.extra import libdevice from triton.language.extra.cuda import gdc_wait, gdc_launch_dependents +from triton.language import target_info -from freetoken.utils.arch import is_sm90_supported +from freetoken.utils.arch import is_rocm, is_sm90_supported SILU = 0 GELU = 1 @@ -50,6 +51,8 @@ def _pdl_supported() -> bool: @triton.jit def _fast_tanh(x): + if target_info.is_hip(): + return libdevice.tanh(x) # PTX tanh.approx.f32 — single HW op, matches flashinfer math::tanh. return tl.inline_asm_elementwise( "tanh.approx.f32 $0, $1;", "=f,f", [x], @@ -59,6 +62,8 @@ def _fast_tanh(x): @triton.jit def _fast_ex2(x): + if target_info.is_hip(): + return libdevice.exp2(x) # PTX ex2.approx.f32 — matches __expf fast path used by flashinfer silu. return tl.inline_asm_elementwise( "ex2.approx.f32 $0, $1;", "=f,f", [x], @@ -141,8 +146,9 @@ def _act_and_mul( block_d = min(triton.next_power_of_2(d), 1024 if M >= 4096 else 512) num_stages = 2 if block_d == 1024 else 3 _act_and_mul_kernel[grid]( - o2, x2, d, alpha, limit, ACT=kind, ENABLE_PDL=pdl, launch_pdl=pdl, + o2, x2, d, alpha, limit, ACT=kind, ENABLE_PDL=pdl, BLOCK_D=block_d, num_warps=4, num_stages=num_stages, + **({} if is_rocm() else {"launch_pdl": pdl}), ) return out diff --git a/python/freetoken/kernel/triton/attention.py b/python/freetoken/kernel/triton/attention.py index 3f53c5fc0e..d898e7a575 100644 --- a/python/freetoken/kernel/triton/attention.py +++ b/python/freetoken/kernel/triton/attention.py @@ -6,11 +6,56 @@ 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 +# a long decode request takes splits of at least this many keys, so the partial result each one writes stays small next to the keys it reads +_MIN_SPLIT_LEN = 128 +# a short request still takes up to this many splits of _MIN_BLOCK_KV keys, so tiny requests stay parallel +_MIN_SPLITS = 8 +_MAX_KV_SPLITS = 64 +# stage 2 merges at most this many fp32 partial elements per step +_STAGE2_TILE_ELEMS = 8192 + + +def _stage2_launch(max_kv_splits: int, block_dv: int) -> tuple[int, int]: + """(splits merged per step, num_warps) for stage 2; the tile is at most the split cap rounded up to a power of two.""" + tile = min(triton.next_power_of_2(max_kv_splits), max(4, _STAGE2_TILE_ELEMS // block_dv), 64) + return tile, max(1, min(8, tile * block_dv // 2048)) + + +def _max_request_splits(keys: int) -> int: + return max(triton.cdiv(keys, _MIN_SPLIT_LEN), min(_MIN_SPLITS, triton.cdiv(keys, _MIN_BLOCK_KV)), 1) + + @functools.lru_cache(maxsize=None) def _optin_smem_bytes(device_index: int) -> int: """Per-block opt-in shared-memory budget for a CUDA device (0 if unavailable).""" @@ -18,25 +63,45 @@ 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]: +@functools.lru_cache(maxsize=None) +def _num_sms(device_index: int) -> int: + return int(torch.cuda.get_device_properties(device_index).multi_processor_count) + + +def decode_kv_split_cap(batch: int, num_q_heads: int, num_kv_heads: int, device_index: int) -> int: + """Most KV splits per request for a decode batch: about four programs per SM, within [_MIN_SPLITS, _MAX_KV_SPLITS].""" + head_blocks = triton.cdiv(num_q_heads, min(16, max(1, num_q_heads // num_kv_heads))) + return max(_MIN_SPLITS, min(_MAX_KV_SPLITS, triton.cdiv(4 * _num_sms(device_index), batch * head_blocks))) + + +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 +112,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 +129,8 @@ def _paged_attention_kernel( stride_kh, stride_vs, stride_vh, + stride_kss, + stride_vss, stride_ot, stride_oh, GROUP: tl.constexpr, @@ -68,6 +139,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 +180,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 +219,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 @@ -145,23 +260,36 @@ def _paged_attention_kernel( @triton.jit +def _decode_num_splits(effective_len, max_kv_splits, MIN_SPLIT_LEN: tl.constexpr, MIN_SPLITS: tl.constexpr, MIN_BLOCK_KV: tl.constexpr): + # same rule as _max_request_splits: long requests split by length, short ones still take up to MIN_SPLITS small splits + by_len = tl.maximum(tl.cdiv(effective_len, MIN_SPLIT_LEN), tl.minimum(MIN_SPLITS, tl.cdiv(effective_len, MIN_BLOCK_KV))) + return tl.minimum(tl.maximum(by_len, 1), max_kv_splits) + + +@triton.jit(do_not_specialize=["max_kv_splits"]) 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, q_pos_ptr, mid_o_ptr, mid_lse_ptr, - num_kv_splits_ptr, + max_kv_splits, stride_qt, stride_qh, stride_ks, stride_kh, stride_vs, stride_vh, + stride_kss, + stride_vss, stride_mid_ob, stride_mid_oh, stride_mid_os, @@ -176,9 +304,13 @@ def _decode_grouped_stage1_kernel( BLOCK_H: tl.constexpr, VALID_BLOCK_H: tl.constexpr, MIN_BLOCK_KV: tl.constexpr, + MIN_SPLIT_LEN: tl.constexpr, + MIN_SPLITS: tl.constexpr, 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) @@ -207,7 +339,7 @@ def _decode_grouped_stage1_kernel( effective_start = tl.maximum(0, q_pos - SLIDING_WINDOW + 1) effective_len = tl.maximum(0, effective_end - effective_start) - kv_splits = tl.load(num_kv_splits_ptr + batch_id) + kv_splits = _decode_num_splits(effective_len, max_kv_splits, MIN_SPLIT_LEN, MIN_SPLITS, MIN_BLOCK_KV) kv_len_per_split = ( tl.cdiv(tl.cdiv(effective_len, kv_splits), MIN_BLOCK_KV) * MIN_BLOCK_KV ) @@ -224,7 +356,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 +367,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 @@ -270,14 +437,14 @@ def _decode_grouped_stage1_kernel( tl.store(mid_lse_ptr + lse_offsets, m_i + tl.log(l_i), mask=mask_h) -@triton.jit +@triton.jit(do_not_specialize=["max_kv_splits"]) def _decode_stage2_kernel( mid_o_ptr, mid_lse_ptr, o_ptr, indptr_ptr, q_pos_ptr, - num_kv_splits_ptr, + max_kv_splits, sinks_ptr, stride_mid_ob, stride_mid_oh, @@ -287,12 +454,14 @@ def _decode_stage2_kernel( stride_lse_s, stride_ot, stride_oh, - MAX_KV_SPLITS: tl.constexpr, MIN_BLOCK_KV: tl.constexpr, + MIN_SPLIT_LEN: tl.constexpr, + MIN_SPLITS: tl.constexpr, BLOCK_DV: tl.constexpr, DV: tl.constexpr, SLIDING_WINDOW: tl.constexpr, HAS_SINKS: tl.constexpr, + SPLIT_TILE: tl.constexpr, ): batch_id = tl.program_id(0) q_head = tl.program_id(1) @@ -305,13 +474,14 @@ def _decode_stage2_kernel( effective_start = tl.maximum(0, q_pos - SLIDING_WINDOW + 1) effective_len = tl.maximum(0, effective_end - effective_start) - kv_splits = tl.load(num_kv_splits_ptr + batch_id) + kv_splits = _decode_num_splits(effective_len, max_kv_splits, MIN_SPLIT_LEN, MIN_SPLITS, MIN_BLOCK_KV) kv_len_per_split = ( tl.cdiv(tl.cdiv(effective_len, kv_splits), MIN_BLOCK_KV) * MIN_BLOCK_KV ) offs_d = tl.arange(0, BLOCK_DV) mask_d = offs_d < DV + offs_s = tl.arange(0, SPLIT_TILE) if HAS_SINKS: m_i = tl.load(sinks_ptr + q_head).to(tl.float32) l_i = 1.0 @@ -320,26 +490,27 @@ def _decode_stage2_kernel( l_i = 0.0 acc = tl.zeros((BLOCK_DV,), dtype=tl.float32) - mid_base = batch_id * stride_mid_ob + q_head * stride_mid_oh + offs_d + mid_base = batch_id * stride_mid_ob + q_head * stride_mid_oh lse_base = batch_id * stride_lse_b + q_head * stride_lse_h - for split_id in tl.range(0, MAX_KV_SPLITS, num_stages=2): - split_start = kv_len_per_split * split_id - split_end = tl.minimum(split_start + kv_len_per_split, effective_len) - - if split_end > split_start: - partial = tl.load( - mid_o_ptr + mid_base + split_id * stride_mid_os, - mask=mask_d, - other=0.0, - ) - partial_lse = tl.load(mid_lse_ptr + lse_base + split_id * stride_lse_s) - m_new = tl.maximum(partial_lse, m_i) - alpha = tl.exp(m_i - m_new) - beta = tl.exp(partial_lse - m_new) - acc = acc * alpha + partial * beta - l_i = l_i * alpha + beta - m_i = m_new + # one load round trip per SPLIT_TILE splits + for s0 in range(0, kv_splits, SPLIT_TILE): + split_id = s0 + offs_s + valid = (split_id < kv_splits) & (kv_len_per_split * split_id < effective_len) + partial_lse = tl.load(mid_lse_ptr + lse_base + split_id * stride_lse_s, mask=valid, other=-float("inf")) + partial = tl.load( + mid_o_ptr + mid_base + split_id[:, None] * stride_mid_os + offs_d[None, :], + mask=valid[:, None] & mask_d[None, :], + other=0.0, + ) + m_new = tl.maximum(m_i, tl.max(partial_lse, axis=0)) + # a tile of empty splits with no sink leaves m at -inf, where exp(-inf - -inf) is nan + m_safe = tl.where(m_new == -float("inf"), 0.0, m_new) + alpha = tl.exp(m_i - m_safe) + beta = tl.exp(partial_lse - m_safe) + acc = acc * alpha + tl.sum(partial * beta[:, None], axis=0) + l_i = l_i * alpha + tl.sum(beta, axis=0) + m_i = m_new out = tl.where(l_i == 0.0, 0.0, acc / l_i) tl.store( @@ -349,6 +520,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, @@ -358,30 +550,54 @@ def decode_paged_attention( q_positions: torch.Tensor, attn_logits: torch.Tensor, attn_lse: torch.Tensor, - num_kv_splits: torch.Tensor, - max_kv_splits: int, sm_scale: float, 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, + max_kv_splits: int | 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. + + Each request takes ``max(ceil(keys / 128), min(8, ceil(keys / 32)))`` splits, capped by ``max_kv_splits`` (default :func:`decode_kv_split_cap`) and by the split dim of ``attn_logits`` / ``attn_lse``. + The kernels compute the split count on the device, so a captured graph serves any context length. + + ``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 - assert attn_logits.shape[2] >= max_kv_splits assert attn_logits.shape[3] >= head_dim assert attn_lse.shape[0] >= batch assert attn_lse.shape[1] >= num_q_heads - assert attn_lse.shape[2] >= max_kv_splits + if sliding_window is not None and sliding_window <= 0: + # the kernels treat a window <= 0 as no window + sliding_window = None + if max_kv_splits is None: + max_kv_splits = decode_kv_split_cap(batch, num_q_heads, num_kv_heads, q.device.index) + if sliding_window: + # no request attends more than the window, so programs past its splits would only exit + max_kv_splits = min(max_kv_splits, _max_request_splits(sliding_window)) + max_kv_splits = min(max_kv_splits, attn_logits.shape[2], attn_lse.shape[2]) if sinks is not None: assert sinks.is_cuda assert sinks.dim() == 1 @@ -391,6 +607,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. @@ -398,26 +625,42 @@ def decode_paged_attention( block_h = triton.next_power_of_2(valid_block_h) block_d = triton.next_power_of_2(head_dim) block_dv = triton.next_power_of_2(head_dim) - - _decode_grouped_stage1_kernel[ - (batch, triton.cdiv(num_q_heads, valid_block_h), max_kv_splits) - ]( + grid = (batch, triton.cdiv(num_q_heads, valid_block_h), max_kv_splits) + # a grid that fits in one wave runs about one program per SM, which then needs wider tiles to keep enough loads in flight + block_n = 64 if grid[0] * grid[1] * grid[2] <= _num_sms(q.device.index) else _MIN_BLOCK_KV + num_stages = 3 + # stage 1 keeps one K tile and one V tile in shared memory + if 2 * block_n * block_d * k_cache.element_size() > 0.8 * _optin_smem_bytes(q.device.index): + block_n = _MIN_BLOCK_KV + if sliding_window: + split_len = triton.cdiv(triton.cdiv(sliding_window, max_kv_splits), _MIN_BLOCK_KV) * _MIN_BLOCK_KV + if split_len <= block_n: + # a full window's split fits in one tile, where a deeper pipeline only adds overhead + block_n, num_stages = split_len, 2 + + _decode_grouped_stage1_kernel[grid]( q, k_cache, v_cache, + k_scale_arg, + v_scale_arg, + k_block_arg, + v_block_arg, sm_scale, indptr, indices, q_positions, attn_logits, attn_lse, - num_kv_splits, + max_kv_splits, q.stride(0), q.stride(1), k_cache.stride(0), 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), @@ -428,23 +671,28 @@ def decode_paged_attention( NUM_Q_HEADS=num_q_heads, BLOCK_D=block_d, BLOCK_DV=block_dv, - BLOCK_N=32, + BLOCK_N=block_n, BLOCK_H=block_h, VALID_BLOCK_H=valid_block_h, MIN_BLOCK_KV=_MIN_BLOCK_KV, + MIN_SPLIT_LEN=_MIN_SPLIT_LEN, + MIN_SPLITS=_MIN_SPLITS, 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, + num_stages=num_stages, ) + split_tile, stage2_warps = _stage2_launch(max_kv_splits, block_dv) _decode_stage2_kernel[(batch, num_q_heads)]( attn_logits, attn_lse, o, indptr, q_positions, - num_kv_splits, + max_kv_splits, sinks_arg, attn_logits.stride(0), attn_logits.stride(1), @@ -454,14 +702,15 @@ def decode_paged_attention( attn_lse.stride(2), o.stride(0), o.stride(1), - MAX_KV_SPLITS=max_kv_splits, MIN_BLOCK_KV=_MIN_BLOCK_KV, + MIN_SPLIT_LEN=_MIN_SPLIT_LEN, + MIN_SPLITS=_MIN_SPLITS, BLOCK_DV=block_dv, DV=head_dim, SLIDING_WINDOW=sliding_window or 0, HAS_SINKS=sinks is not None, - num_warps=4, - num_stages=2, + SPLIT_TILE=split_tile, + num_warps=stage2_warps, ) return o @@ -471,6 +720,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 +738,8 @@ def _extend_attention_kernel( stride_kh, stride_vs, stride_vh, + stride_kss, + stride_vss, stride_ot, stride_oh, GROUP: tl.constexpr, @@ -495,6 +750,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 +812,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 +847,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 @@ -595,6 +888,98 @@ def _extend_attention_kernel( ) +@triton.jit +def _extend_attend_tile( + q, m_i, l_i, acc, qk_scale, k_base, v_base, k_rows, v_rows, offs_d, offs_dv, mask_d, mask_dv, mask_n, mask, + MASKED: tl.constexpr, +): + if MASKED: + k_mask = mask_n[None, :] & mask_d[:, None] + v_mask = mask_n[:, None] & mask_dv[None, :] + else: + k_mask = mask_d[:, None] + v_mask = mask_dv[None, :] + k = tl.load(k_base + k_rows[None, :] + offs_d[:, None], mask=k_mask, other=0.0) + scores = tl.dot(q.to(k.dtype), k) + if MASKED: + scores = tl.where(mask, scores, -float("inf")) + m_new = tl.maximum(m_i, tl.max(scores, axis=1) * qk_scale) + if MASKED: + m_new = tl.where(m_new == -float("inf"), -1e20, m_new) + alpha = tl.math.exp2(m_i - m_new) + p = tl.math.exp2(scores * qk_scale - m_new[:, None]) + v = tl.load(v_base + v_rows[:, None] + offs_dv[None, :], mask=v_mask, other=0.0) + acc = acc * alpha[:, None] + tl.dot(p.to(v.dtype), v) + l_i = l_i * alpha + tl.sum(p, axis=1) + return m_new, l_i, acc + + +@triton.jit +def _extend_attend_cache_tile( + q, m_i, l_i, acc, qk_scale, k_cache_ptr, v_cache_ptr, k_scale_ptr, v_scale_ptr, k_block_ptr, v_block_ptr, + slots, kv_head, stride_kcs, stride_kch, stride_vcs, stride_vch, stride_kss, stride_vss, + offs_d, offs_dv, mask_d, mask_dv, mask_n, mask, D: tl.constexpr, MASKED: tl.constexpr, KV_NVFP4: tl.constexpr, +): + """:func:`_extend_attend_tile` over an fp8 or NVFP4 cache tile: K/V are decoded to q's dtype, and an fp8 + tile's per-(token, kv_head) scales go on the dot outputs (see :func:`_kv_dequant_scale`). ``mask_n`` is + required even when ``MASKED`` is false, since the scale loads take it.""" + if KV_NVFP4: + k = load_nvfp4( + k_cache_ptr, k_block_ptr, k_scale_ptr, slots[None, :], kv_head, offs_d[:, None], + mask_n[None, :], stride_kcs, stride_kch, stride_kss, D, + ).to(q.dtype) + else: + s_k = _kv_dequant_scale(k_scale_ptr, slots, stride_kss, kv_head, mask_n) + k = _kv_load_s16( + k_cache_ptr + slots[None, :] * stride_kcs + kv_head * stride_kch + offs_d[:, None], + mask_n[None, :] & mask_d[:, None], + ).to(q.dtype) + scores = tl.dot(q.to(k.dtype), k) + # NVFP4's loader already applies both row and block scales. + if not KV_NVFP4: + scores = scores * s_k[None, :] + if MASKED: + scores = tl.where(mask, scores, -float("inf")) + m_new = tl.maximum(m_i, tl.max(scores, axis=1) * qk_scale) + if MASKED: + m_new = tl.where(m_new == -float("inf"), -1e20, m_new) + alpha = tl.math.exp2(m_i - m_new) + p = tl.math.exp2(scores * qk_scale - m_new[:, None]) + if KV_NVFP4: + v = load_nvfp4( + v_cache_ptr, v_block_ptr, v_scale_ptr, slots[:, None], kv_head, offs_dv[None, :], + mask_n[:, None], stride_vcs, stride_vch, stride_vss, D, + ).to(q.dtype) + pv = p + else: + s_v = _kv_dequant_scale(v_scale_ptr, slots, stride_vss, kv_head, mask_n) + v = _kv_load_s16( + v_cache_ptr + slots[:, None] * stride_vcs + kv_head * stride_vch + offs_dv[None, :], + mask_n[:, None] & mask_dv[None, :], + ).to(q.dtype) + # p stays unscaled for the l_i denominator below. + pv = p * s_v[None, :] + acc = acc * alpha[:, None] + tl.dot(pv.to(v.dtype), v) + l_i = l_i * alpha + tl.sum(p, axis=1) + return m_new, l_i, acc + + +@triton.jit +def _extend_tile_ranges( + first_row, diag, end, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, SLIDING_WINDOW: tl.constexpr +): + """Split keys [0, end) seen by rows first_row + [0, BLOCK_M) into skipped [0, lo), masked [lo, u0), unmasked [u0, u1) and masked [u1, end); keys below diag are causal for every row.""" + if SLIDING_WINDOW > 0: + lo = tl.minimum(tl.maximum(first_row - SLIDING_WINDOW + 1, 0) // BLOCK_N * BLOCK_N, end) + full = tl.cdiv(tl.maximum(first_row + BLOCK_M - SLIDING_WINDOW, 0), BLOCK_N) * BLOCK_N + else: + lo = 0 + full = 0 + u1 = tl.maximum(diag, lo) + u0 = tl.minimum(tl.maximum(full, lo), u1) + return lo, u0, u1 + + @triton.jit def _extend_attention_split_kernel( q_ptr, @@ -602,6 +987,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 +1009,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 +1021,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) @@ -642,7 +1035,8 @@ def _extend_attention_split_kernel( kv_start = tl.load(kv_indptr_ptr + seq_id) prefix_len = tl.load(prefix_lens_ptr + seq_id) - offs_m = block_m_id * BLOCK_M + tl.arange(0, BLOCK_M) + block_start = block_m_id * BLOCK_M + offs_m = block_start + tl.arange(0, BLOCK_M) offs_n = tl.arange(0, BLOCK_N) offs_d = tl.arange(0, BLOCK_D) offs_dv = tl.arange(0, BLOCK_DV) @@ -650,6 +1044,8 @@ def _extend_attention_split_kernel( mask_d = offs_d < D mask_dv = offs_dv < D q_abs_pos = prefix_len + offs_m + # Scores live in the log2 domain: exp2 lowers to ex2.approx.ftz, while tl.exp adds a denormal fixup per element. + qk_scale = sm_scale * 1.4426950408889634 q = tl.load( q_ptr @@ -661,7 +1057,7 @@ def _extend_attention_split_kernel( ) if HAS_SINKS: - sink = tl.load(sinks_ptr + q_head).to(tl.float32) + sink = tl.load(sinks_ptr + q_head).to(tl.float32) * 1.4426950408889634 m_i = tl.full((BLOCK_M,), sink, dtype=tl.float32) l_i = tl.full((BLOCK_M,), 1.0, dtype=tl.float32) else: @@ -669,59 +1065,68 @@ def _extend_attention_split_kernel( l_i = tl.zeros((BLOCK_M,), dtype=tl.float32) acc = tl.zeros((BLOCK_M, BLOCK_DV), dtype=tl.float32) - for start_n in tl.range(0, prefix_len, BLOCK_N): - kv_offsets = start_n + offs_n - mask_n = kv_offsets < prefix_len - key_pos = kv_offsets - final_mask = mask_m[:, None] & mask_n[None, :] - if SLIDING_WINDOW > 0: - window_mask = (key_pos[None, :] + SLIDING_WINDOW) > q_abs_pos[:, None] - final_mask = final_mask & window_mask - - skip_tile = False + kc_base = k_cache_ptr + kv_head * stride_kch + vc_base = v_cache_ptr + kv_head * stride_vch + prefix_end = tl.where(block_start < q_len, prefix_len, 0) + lo, u0, u1 = _extend_tile_ranges( + prefix_len + block_start, prefix_end // BLOCK_N * BLOCK_N, prefix_end, BLOCK_M, BLOCK_N, SLIDING_WINDOW + ) + for start_n in tl.range(u0, u1, BLOCK_N): + slots = tl.load(kv_indices_ptr + kv_start + start_n + offs_n) + if HAS_KV_SCALE: + m_i, l_i, acc = _extend_attend_cache_tile( + q, m_i, l_i, acc, qk_scale, k_cache_ptr, v_cache_ptr, k_scale_ptr, v_scale_ptr, k_block_ptr, + v_block_ptr, slots, kv_head, stride_kcs, stride_kch, stride_vcs, stride_vch, stride_kss, stride_vss, + offs_d, offs_dv, mask_d, mask_dv, offs_n < BLOCK_N, None, D, False, KV_NVFP4, + ) + else: + m_i, l_i, acc = _extend_attend_tile( + q, m_i, l_i, acc, qk_scale, kc_base, vc_base, slots * stride_kcs, slots * stride_vcs, + offs_d, offs_dv, mask_d, mask_dv, None, None, False, + ) + n_left = (u0 - lo) // BLOCK_N + for i in tl.range(0, n_left + tl.cdiv(prefix_end - u1, BLOCK_N)): + kv_offsets = tl.where(i < n_left, lo + i * BLOCK_N, u1 + (i - n_left) * BLOCK_N) + offs_n + mask_n = kv_offsets < prefix_end + final_mask = mask_n[None, :] if SLIDING_WINDOW > 0: - 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_cache_ptr - + slots[None, :] * stride_kcs - + kv_head * stride_kch - + offs_d[:, None], - mask=mask_n[None, :] & mask_d[:, None], - other=0.0, + final_mask = final_mask & ((kv_offsets[None, :] + SLIDING_WINDOW) > q_abs_pos[:, None]) + slots = tl.load(kv_indices_ptr + kv_start + kv_offsets, mask=mask_n, other=0) + if HAS_KV_SCALE: + m_i, l_i, acc = _extend_attend_cache_tile( + q, m_i, l_i, acc, qk_scale, k_cache_ptr, v_cache_ptr, k_scale_ptr, v_scale_ptr, k_block_ptr, + v_block_ptr, slots, kv_head, stride_kcs, stride_kch, stride_vcs, stride_vch, stride_kss, stride_vss, + offs_d, offs_dv, mask_d, mask_dv, mask_n, final_mask, D, True, KV_NVFP4, ) - scores = tl.dot(q.to(k.dtype), k) * sm_scale - scores = tl.where(final_mask, scores, -float("inf")) - - row_max = tl.max(scores, axis=1) - row_max_fixed = tl.where(row_max == -float("inf"), -1e20, row_max) - m_new = tl.maximum(row_max_fixed, m_i) - 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, + else: + m_i, l_i, acc = _extend_attend_tile( + q, m_i, l_i, acc, qk_scale, kc_base, vc_base, slots * stride_kcs, slots * stride_vcs, + offs_d, offs_dv, mask_d, mask_dv, mask_n, final_mask, True, ) - acc = acc * alpha[:, None] + tl.dot(p.to(v.dtype), v) - l_i = l_i * alpha + tl.sum(p, axis=1) - m_i = m_new - current_end = tl.minimum(q_len, (block_m_id + 1) * BLOCK_M) + ke_base = k_extend_ptr + kv_head * stride_keh + ve_base = v_extend_ptr + kv_head * stride_veh + current_end = tl.minimum(q_len, block_start + BLOCK_M) if HAS_BLOCKS: # rows inside a multimodal span also attend forward to the span's later keys: the tile loop must reach them block_end = tl.load(block_ends_ptr + q_start + offs_m, mask=mask_m, other=0) - prefix_len current_end = tl.minimum(q_len, tl.maximum(current_end, tl.max(block_end, axis=0))) else: block_end = tl.zeros((BLOCK_M,), dtype=tl.int32) - for start_n in tl.range(0, current_end, BLOCK_N): - local_kv_offsets = start_n + offs_n + current_end = tl.where(block_start < q_len, current_end, 0) + lo, u0, u1 = _extend_tile_ranges( + block_start, tl.minimum(block_start // BLOCK_N * BLOCK_N, current_end), current_end, BLOCK_M, BLOCK_N, + SLIDING_WINDOW, + ) + for start_n in tl.range(u0, u1, BLOCK_N): + rows = q_start + start_n + offs_n + m_i, l_i, acc = _extend_attend_tile( + q, m_i, l_i, acc, qk_scale, ke_base, ve_base, rows * stride_ket, rows * stride_vet, + offs_d, offs_dv, mask_d, mask_dv, None, None, False, + ) + n_left = (u0 - lo) // BLOCK_N + for i in tl.range(0, n_left + tl.cdiv(current_end - u1, BLOCK_N)): + local_kv_offsets = tl.where(i < n_left, lo + i * BLOCK_N, u1 + (i - n_left) * BLOCK_N) + offs_n mask_n = local_kv_offsets < current_end local_q_pos = offs_m causal_mask = local_kv_offsets[None, :] <= local_q_pos[:, None] @@ -731,41 +1136,12 @@ def _extend_attention_split_kernel( causal_mask = causal_mask & ( (local_kv_offsets[None, :] + SLIDING_WINDOW) > local_q_pos[:, None] ) - final_mask = mask_m[:, None] & mask_n[None, :] & causal_mask - - skip_tile = False - if SLIDING_WINDOW > 0: - skip_tile = tl.max(tl.max(final_mask.to(tl.int32), axis=1), axis=0) == 0 - - if not skip_tile: - k = tl.load( - k_extend_ptr - + (q_start + local_kv_offsets[None, :]) * stride_ket - + kv_head * stride_keh - + offs_d[:, None], - mask=mask_n[None, :] & mask_d[:, None], - other=0.0, - ) - scores = tl.dot(q.to(k.dtype), k) * sm_scale - scores = tl.where(final_mask, scores, -float("inf")) - - row_max = tl.max(scores, axis=1) - row_max_fixed = tl.where(row_max == -float("inf"), -1e20, row_max) - m_new = tl.maximum(row_max_fixed, m_i) - alpha = tl.exp(m_i - m_new) - p = tl.exp(scores - m_new[:, None]) - - v = tl.load( - v_extend_ptr - + (q_start + local_kv_offsets[:, None]) * stride_vet - + kv_head * stride_veh - + offs_dv[None, :], - mask=mask_n[:, None] & mask_dv[None, :], - other=0.0, - ) - acc = acc * alpha[:, None] + tl.dot(p.to(v.dtype), v) - l_i = l_i * alpha + tl.sum(p, axis=1) - m_i = m_new + final_mask = mask_n[None, :] & causal_mask + rows = q_start + local_kv_offsets + m_i, l_i, acc = _extend_attend_tile( + q, m_i, l_i, acc, qk_scale, ke_base, ve_base, rows * stride_ket, rows * stride_vet, + offs_d, offs_dv, mask_d, mask_dv, mask_n, final_mask, True, + ) out = tl.where(l_i[:, None] == 0.0, 0.0, acc / l_i[:, None]) tl.store( @@ -793,18 +1169,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 +1206,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 +1242,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 +1264,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 +1276,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 +1288,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 +1306,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 +1318,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 +1340,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 +1371,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 +1404,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 +1414,8 @@ def paged_attention( BLOCK_N=block_n, SLIDING_WINDOW=sliding_window or 0, HAS_SINKS=sinks is not None, + HAS_KV_SCALE=has_kv_scale, + KV_NVFP4=nvfp4, num_warps=8 if head_dim >= 256 else 4, num_stages=2, ) diff --git a/python/freetoken/kernel/triton/autotune_cache.py b/python/freetoken/kernel/triton/autotune_cache.py index ea95768368..2e57fadde9 100644 --- a/python/freetoken/kernel/triton/autotune_cache.py +++ b/python/freetoken/kernel/triton/autotune_cache.py @@ -34,4 +34,28 @@ {"cache_results": _CACHE_RESULTS} if _SUPPORTS_AUTOTUNE_CACHE else {} ) -__all__ = ["autotune_cache_kwargs"] + + +def bound_autotune_flush_buffer(device_index: int) -> None: + """Size triton's L2-flush buffer for autotune benchmarks to the device instead of a flat 256 MB. + + A kernel autotunes on its first launch of each key, which can be the first request, after the + cache planner has handed out nearly all VRAM; a 256 MB transient there is an OOM on a small + card. Four times the L2 still evicts everything between timed runs. + """ + import torch + + try: + from triton.backends.nvidia.driver import CudaDriver + except ImportError: + return + l2 = int(getattr(torch.cuda.get_device_properties(device_index), "L2_cache_size", 0) or 0) + nbytes = min(256 << 20, max(8 << 20, 4 * l2)) + + def get_empty_cache_for_benchmark(self): + return torch.empty(nbytes // 4, dtype=torch.int, device="cuda") + + CudaDriver.get_empty_cache_for_benchmark = get_empty_cache_for_benchmark + + +__all__ = ["autotune_cache_kwargs", "bound_autotune_flush_buffer"] diff --git a/python/freetoken/kernel/triton/dsv4/compress.py b/python/freetoken/kernel/triton/dsv4/compress.py index 77bacd42ce..3ab89ff0d3 100644 --- a/python/freetoken/kernel/triton/dsv4/compress.py +++ b/python/freetoken/kernel/triton/dsv4/compress.py @@ -31,12 +31,13 @@ def _gated_pool_kernel( kv_base = kv_ptr + b * stride_kb + offs * stride_kd sc_base = score_ptr + b * stride_sb + offs * stride_sd maxv = tl.full((BLOCK_D,), float("-inf"), tl.float32) - for r in range(R): + # Explicit unroll: ptxas does not unroll the plain loop Triton 3.8 emits, which serializes these latency-bound loads. + for r in tl.range(R, loop_unroll_factor=2): s = tl.load(sc_base + r * stride_sr, mask=mask, other=float("-inf")).to(tl.float32) maxv = tl.maximum(maxv, s) denom = tl.zeros((BLOCK_D,), tl.float32) acc = tl.zeros((BLOCK_D,), tl.float32) - for r in range(R): + for r in tl.range(R, loop_unroll_factor=2): s = tl.load(sc_base + r * stride_sr, mask=mask, other=float("-inf")).to(tl.float32) w = tl.exp(s - maxv) denom += w 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/dsv4/sparse_attn.py b/python/freetoken/kernel/triton/dsv4/sparse_attn.py index c22f891e89..639c3e5ea5 100644 --- a/python/freetoken/kernel/triton/dsv4/sparse_attn.py +++ b/python/freetoken/kernel/triton/dsv4/sparse_attn.py @@ -99,7 +99,8 @@ def _sparse_attn_paged_kernel( # multiple KV tiles in shared memory. Result is bit-identical to the two-pool form -- each # column still reads from the same pool/slot. is_win = offs_t < N_WINDOW - base = tl.where(is_win, win_ptr, cmp_ptr) # [BLOCK_T] per-column pool base pointer + # The select drops the pools' 16-byte alignment (asserted in the wrapper); without this hint Triton 3.8 gathers in 2-byte loads. + base = tl.multiple_of(tl.where(is_win, win_ptr, cmp_ptr), 16) # [BLOCK_T] per-column pool base pointer kv_ptrs = base[:, None] + idxs[:, None] * stride_wn + offs_d[None, :] * stride_wd kv = tl.load(kv_ptrs, mask=valid[:, None], other=0.0).to(tl.float32) # [BLOCK_T, D] @@ -182,7 +183,7 @@ def _sparse_attn_paged_splitk_kernel( idxs = tl.load(idx_base + offs_t * stride_it, mask=t_mask, other=-1) valid = idxs >= 0 is_win = offs_t < N_WINDOW - base = tl.where(is_win, win_ptr, cmp_ptr) + base = tl.multiple_of(tl.where(is_win, win_ptr, cmp_ptr), 16) kv_ptrs = base[:, None] + idxs[:, None] * stride_wn + offs_d[None, :] * stride_wd kv = tl.load(kv_ptrs, mask=valid[:, None], other=0.0).to(tl.float32) @@ -314,6 +315,7 @@ def sparse_attn_paged( # The kernel selects one pool base per column and gathers with a single stride pair, so both # pools must share strides. Guaranteed here by .contiguous() on their [*, D] shape. assert window_pool.stride() == cmp_pool.stride(), (window_pool.stride(), cmp_pool.stride()) + assert window_pool.data_ptr() % 16 == 0 and cmp_pool.data_ptr() % 16 == 0, (window_pool.data_ptr(), cmp_pool.data_ptr()) has_counts = cmp_counts is not None if has_counts: 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..bd3fe7ae31 --- /dev/null +++ b/python/freetoken/kernel/triton/dsv41/indexer.py @@ -0,0 +1,196 @@ +"""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 _load_fp4_e8m0_rows(pool_ptr, rows, valid, D: tl.constexpr, ROW_BYTES: tl.constexpr): + """``load_rows(..., FMT=3)`` as bf16 for ``D == 128``, with each row read as int32 words: 16 words of 8 e2m1 codes, then one word of 4 scales. + The contiguous word axis pins the gather layout; a per-channel byte gather has no contiguous axis, so its layout (and speed) follows the Triton version.""" + W: tl.constexpr = D // 8 + base = pool_ptr.to(tl.pointer_type(tl.int32)) + tl.maximum(rows, 0).to(tl.int64) * (ROW_BYTES // 4) + offs_w = tl.arange(0, W) + x = tl.load(base[:, None] + offs_w[None, :], mask=valid[:, None], other=0) + sw = tl.load(base + W, mask=valid, other=0) + s = tl.exp2(((sw[:, None] >> (8 * (offs_w // 4))[None, :]) & 0xFF).to(tl.float32) - 127.0) + x = tl.join(x, x >> 16) + x = tl.join(x, x >> 8) + x = tl.join(x, x >> 4) + # e2m1 code -> fp32 bits: magnitude to bits 22-24, sign to bit 31 (mask 0x81C00000). That is the value times 2**-126 (code 1 is subnormal), so the multiply is exact. + q = (((x << 28) >> 6) & -0x7E400000).to(tl.float32, bitcast=True) * 8.507059173023462e37 + return tl.reshape((q * s[:, :, None, None, None]).to(tl.bfloat16), (rows.shape[0], D)) + + +@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) + if FMT == 3 and D == 128: + k = _load_fp4_e8m0_rows(pool_ptr, rows, valid, D, ROW_BYTES) + else: + k = load_rows(pool_ptr, rows, valid, tl.arange(0, 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() and k_pool.data_ptr() % 4 == 0 + 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..c0fca83162 --- /dev/null +++ b/python/freetoken/kernel/triton/dsv41/pack.py @@ -0,0 +1,188 @@ +"""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)) + if e4m3_native_cx(): + s = amax * (1.0 / 6.0) + else: + # Round before the emulated magic-add; fusing it changes subnormal ties. + s = tl.math.div_rn(amax, 6.0) + s = _e4m3_round(s) + 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 1d9f744ce1..b14a7ca68b 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( @@ -59,14 +70,56 @@ def e4m3_native() -> bool: if _native is None: if FORCE_EMU: _native = False + elif torch.version.hip is not None: + # ROCm reports gfx1101 as capability (11, 0), which is not a CUDA + # compute capability and must not select the native fp8e4nv path. + _native = False else: from freetoken.gpu_select import assigned_visible_gpu # 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).""" @@ -84,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) @@ -105,6 +168,20 @@ 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 round_e4m3(x): """Round fp32 onto the e4m3 value grid (RNE), fp32 -> fp32, in a SINGLE @@ -122,3 +199,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/fp8_blockscale_moe.py b/python/freetoken/kernel/triton/fp8_blockscale_moe.py index 747605168a..429134d23f 100644 --- a/python/freetoken/kernel/triton/fp8_blockscale_moe.py +++ b/python/freetoken/kernel/triton/fp8_blockscale_moe.py @@ -35,7 +35,7 @@ def _decode_fp8_moe_kernel( stride_twm, stride_twk, stride_tidm, stride_tidk, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, # BLOCK_K == 128 (one weight scale block) TOP_K: tl.constexpr, A_ROW_IS_ROUTE: tl.constexpr, MUL_ROUTED_WEIGHT: tl.constexpr, - compute_type: tl.constexpr, + compute_type: tl.constexpr, EVEN_K: tl.constexpr, ): route_id = tl.program_id(0) n_block = tl.program_id(1) @@ -54,7 +54,10 @@ def _decode_fp8_moe_kernel( acc = tl.zeros((BLOCK_N,), dtype=tl.float32) for kb in range(tl.cdiv(K, BLOCK_K)): offs_k = kb * BLOCK_K + tl.arange(0, BLOCK_K) - k_mask = offs_k < K + if EVEN_K: + k_mask = tl.full((BLOCK_K,), True, tl.int1) + else: + k_mask = offs_k < K if e4m3_native_cx(): w = tl.load( w_slot + offs_n[:, None] * stride_wn + offs_k[None, :] * stride_wk, @@ -66,8 +69,10 @@ def _decode_fp8_moe_kernel( mask=n_mask[:, None] & k_mask[None, :], other=0, )) sc = tl.load(s_slot + sn * stride_sn + kb * stride_sk, mask=n_mask, other=0.0).to(tl.float32) - a = tl.load(a_base + offs_k * stride_ak, mask=k_mask, other=0.0).to(tl.float32) - acc += tl.sum(w * a[None, :], axis=1) * sc + # load the activation row at the weight tile shape so it gets the weight tile layout (no shared-memory layout conversion) + a_ptrs = tl.broadcast_to(a_base + offs_k[None, :] * stride_ak, (BLOCK_N, BLOCK_K)) + a = tl.load(a_ptrs, mask=k_mask[None, :], other=0.0).to(tl.float32) + acc += tl.sum(w * a, axis=1) * sc if MUL_ROUTED_WEIGHT: acc *= tl.load(topk_weights_ptr + token_id * stride_twm + route_k * stride_twk) @@ -92,7 +97,7 @@ def _decode_gemm(a, w, s, c, topk_weights, topk_ids, mul_routed_weight, a_row_is topk_weights.stride(0), topk_weights.stride(1), topk_ids.stride(0), topk_ids.stride(1), BLOCK_N=BLOCK_N, BLOCK_K=128, TOP_K=top_k, A_ROW_IS_ROUTE=a_row_is_route, MUL_ROUTED_WEIGHT=mul_routed_weight, - compute_type=_TL.get(c.dtype, tl.bfloat16), num_warps=4, + compute_type=_TL.get(c.dtype, tl.bfloat16), EVEN_K=w.shape[2] % 128 == 0, num_warps=4, ) @@ -123,8 +128,7 @@ def fused_experts_decode_fp8_blockscale( # ====================================================================================== -# Prefill: M-tiled inline-dequant grouped GEMM (moe_align_block_size), mirrors -# nvfp4_fused_moe._prefill_nvfp4_moe_kernel with fp8*bf16-block-scale dequant. +# Prefill: M-tiled inline-dequant grouped GEMM (moe_align_block_size) with fp8*bf16-block-scale dequant. # BLOCK_SIZE_N == BLOCK_SIZE_K == 128 so each tile maps to one weight scale block. # ====================================================================================== @triton.jit diff --git a/python/freetoken/kernel/triton/glm_dsa_sparse.py b/python/freetoken/kernel/triton/glm_dsa_sparse.py index 8aaeeb9ba9..bac2cd7eed 100644 --- a/python/freetoken/kernel/triton/glm_dsa_sparse.py +++ b/python/freetoken/kernel/triton/glm_dsa_sparse.py @@ -28,6 +28,9 @@ import triton import triton.language as tl +from freetoken.kernel.triton.e4m3_compat import kv_load_e4m3_tile_f32 +from freetoken.kernel.triton.kv_nvfp4 import load_nvfp4 + BLOCK_H = 16 BLOCK_T = 32 MAX_SPLITS = 32 @@ -36,11 +39,12 @@ @triton.jit def _glm_dsa_sparse_kernel( - q_ptr, pool_ptr, o_ptr, idx_ptr, cnt_ptr, + q_ptr, pool_ptr, pool_scale_ptr, pool_block_ptr, o_ptr, idx_ptr, cnt_ptr, scale, H, TOPK, stride_qb, stride_qm, stride_qh, stride_qd, stride_pn, stride_pd, + stride_ps, stride_ob, stride_om, stride_oh, stride_od, stride_ib, stride_im, stride_it, stride_nb, stride_nm, @@ -50,6 +54,8 @@ def _glm_dsa_sparse_kernel( BLOCK_T: tl.constexpr, HAS_COUNTS: tl.constexpr, HAS_ROPE: tl.constexpr, + HAS_FP8: tl.constexpr, + HAS_NVFP4: tl.constexpr, ): pid_m = tl.program_id(0) pid_b = tl.program_id(1) @@ -61,11 +67,15 @@ def _glm_dsa_sparse_kernel( q_base = q_ptr + pid_b * stride_qb + pid_m * stride_qm + offs_h[:, None] * stride_qh q_v = tl.load(q_base + offs_v[None, :] * stride_qd, mask=h_mask[:, None], other=0.0).to(tl.float32) + if HAS_NVFP4: + q_v = q_v.to(q_ptr.dtype.element_ty) if HAS_ROPE: # NoPE checkpoints (glm5_next) have D_R == 0: tl.arange needs a non-empty # span, so the whole rope half is compiled out on the constexpr. offs_r = tl.arange(0, D_R) q_r = tl.load(q_base + (D_V + offs_r[None, :]) * stride_qd, mask=h_mask[:, None], other=0.0).to(tl.float32) + if HAS_NVFP4: + q_r = q_r.to(q_ptr.dtype.element_ty) m_i = tl.full((BLOCK_H,), -float("inf"), dtype=tl.float32) l_i = tl.zeros((BLOCK_H,), dtype=tl.float32) @@ -82,11 +92,34 @@ def _glm_dsa_sparse_kernel( idxs = tl.load(idx_base + offs_t * stride_it, mask=t_mask, other=-1) valid = idxs >= 0 kv_base = pool_ptr + idxs[:, None] * stride_pn - kv_v = tl.load(kv_base + offs_v[None, :] * stride_pd, mask=valid[:, None], other=0.0).to(tl.float32) + if HAS_NVFP4: + kv_v = load_nvfp4( + pool_ptr, pool_block_ptr, pool_scale_ptr, + idxs[:, None], 0, offs_v[None, :], valid[:, None], + stride_pn, 0, stride_ps, D_V + D_R, + ).to(q_ptr.dtype.element_ty) + elif HAS_FP8: + row_scale = tl.load(pool_scale_ptr + idxs * stride_ps, mask=valid, other=0.0) + kv_v = kv_load_e4m3_tile_f32( + kv_base + offs_v[None, :] * stride_pd, valid[:, None] + ) * row_scale[:, None] + else: + kv_v = tl.load(kv_base + offs_v[None, :] * stride_pd, mask=valid[:, None], other=0.0).to(tl.float32) scores = tl.dot(q_v, tl.trans(kv_v)) if HAS_ROPE: - kv_r = tl.load(kv_base + (D_V + offs_r[None, :]) * stride_pd, mask=valid[:, None], other=0.0).to(tl.float32) + if HAS_NVFP4: + kv_r = load_nvfp4( + pool_ptr, pool_block_ptr, pool_scale_ptr, + idxs[:, None], 0, offs_r[None, :], valid[:, None], + stride_pn, 0, stride_ps, D_V + D_R, DIM_OFFSET=D_V, + ).to(q_ptr.dtype.element_ty) + elif HAS_FP8: + kv_r = kv_load_e4m3_tile_f32( + kv_base + (D_V + offs_r[None, :]) * stride_pd, valid[:, None] + ) * row_scale[:, None] + else: + kv_r = tl.load(kv_base + (D_V + offs_r[None, :]) * stride_pd, mask=valid[:, None], other=0.0).to(tl.float32) scores += tl.dot(q_r, tl.trans(kv_r)) scores = scores * scale scores = tl.where(valid[None, :], scores, -float("inf")) @@ -99,7 +132,7 @@ def _glm_dsa_sparse_kernel( acc = acc * alpha[:, None] + tl.dot(p.to(kv_v.dtype), kv_v) m_i = m_new - o = acc / l_i[:, None] + o = tl.where(l_i[:, None] > 0, acc / l_i[:, None], 0.0) o_ptrs = o_ptr + pid_b * stride_ob + pid_m * stride_om + offs_h[:, None] * stride_oh + offs_v[None, :] * stride_od tl.store(o_ptrs, o.to(o_ptr.dtype.element_ty), mask=h_mask[:, None]) @@ -200,11 +233,12 @@ def glm_dsa_decode_logits( @triton.jit def _glm_dsa_splitk_kernel( - q_ptr, pool_ptr, mid_o_ptr, mid_lse_ptr, idx_ptr, cnt_ptr, + q_ptr, pool_ptr, pool_scale_ptr, pool_block_ptr, mid_o_ptr, mid_lse_ptr, idx_ptr, cnt_ptr, scale, H, TOPK, stride_qb, stride_qm, stride_qh, stride_qd, stride_pn, stride_pd, + stride_ps, stride_mb, stride_mm, stride_mh, stride_ms, stride_md, stride_lb, stride_lm, stride_lh, stride_ls, stride_ib, stride_im, stride_it, @@ -215,6 +249,8 @@ def _glm_dsa_splitk_kernel( BLOCK_T: tl.constexpr, HAS_COUNTS: tl.constexpr, HAS_ROPE: tl.constexpr, + HAS_FP8: tl.constexpr, + HAS_NVFP4: tl.constexpr, NUM_SPLITS: tl.constexpr, ): """Stage 1 (decode flash-decoding): each program reduces one BLOCK_T-aligned slice of @@ -245,9 +281,13 @@ def _glm_dsa_splitk_kernel( if split_end > split_start: q_base = q_ptr + pid_b * stride_qb + pid_m * stride_qm + offs_h[:, None] * stride_qh q_v = tl.load(q_base + offs_v[None, :] * stride_qd, mask=h_mask[:, None], other=0.0).to(tl.float32) + if HAS_NVFP4: + q_v = q_v.to(q_ptr.dtype.element_ty) if HAS_ROPE: offs_r = tl.arange(0, D_R) q_r = tl.load(q_base + (D_V + offs_r[None, :]) * stride_qd, mask=h_mask[:, None], other=0.0).to(tl.float32) + if HAS_NVFP4: + q_r = q_r.to(q_ptr.dtype.element_ty) idx_base = idx_ptr + pid_b * stride_ib + pid_m * stride_im for start in range(split_start, split_end, BLOCK_T): @@ -256,11 +296,34 @@ def _glm_dsa_splitk_kernel( idxs = tl.load(idx_base + offs_t * stride_it, mask=t_mask, other=-1) valid = idxs >= 0 kv_base = pool_ptr + idxs[:, None] * stride_pn - kv_v = tl.load(kv_base + offs_v[None, :] * stride_pd, mask=valid[:, None], other=0.0).to(tl.float32) + if HAS_NVFP4: + kv_v = load_nvfp4( + pool_ptr, pool_block_ptr, pool_scale_ptr, + idxs[:, None], 0, offs_v[None, :], valid[:, None], + stride_pn, 0, stride_ps, D_V + D_R, + ).to(q_ptr.dtype.element_ty) + elif HAS_FP8: + row_scale = tl.load(pool_scale_ptr + idxs * stride_ps, mask=valid, other=0.0) + kv_v = kv_load_e4m3_tile_f32( + kv_base + offs_v[None, :] * stride_pd, valid[:, None] + ) * row_scale[:, None] + else: + kv_v = tl.load(kv_base + offs_v[None, :] * stride_pd, mask=valid[:, None], other=0.0).to(tl.float32) scores = tl.dot(q_v, tl.trans(kv_v)) if HAS_ROPE: - kv_r = tl.load(kv_base + (D_V + offs_r[None, :]) * stride_pd, mask=valid[:, None], other=0.0).to(tl.float32) + if HAS_NVFP4: + kv_r = load_nvfp4( + pool_ptr, pool_block_ptr, pool_scale_ptr, + idxs[:, None], 0, offs_r[None, :], valid[:, None], + stride_pn, 0, stride_ps, D_V + D_R, DIM_OFFSET=D_V, + ).to(q_ptr.dtype.element_ty) + elif HAS_FP8: + kv_r = kv_load_e4m3_tile_f32( + kv_base + (D_V + offs_r[None, :]) * stride_pd, valid[:, None] + ) * row_scale[:, None] + else: + kv_r = tl.load(kv_base + (D_V + offs_r[None, :]) * stride_pd, mask=valid[:, None], other=0.0).to(tl.float32) scores += tl.dot(q_r, tl.trans(kv_r)) scores = scores * scale scores = tl.where(valid[None, :], scores, -float("inf")) @@ -323,7 +386,7 @@ def _glm_dsa_merge_kernel( l_i = l_i * alpha + beta m_i = m_new - o = acc / l_i + o = tl.where(l_i > 0, acc / l_i, 0.0) o_ptrs = ( o_ptr + pid_b * stride_ob + pid_m * stride_om + pid_h * stride_oh + offs_v * stride_od @@ -351,7 +414,10 @@ def glm_dsa_sparse_attn( softmax_scale: float, counts: torch.Tensor | None = None, # [b, m] int32 live columns per query (device-read) d_v: int = 512, + pool_scale: torch.Tensor | None = None, # [rows] fp32, one scale per quantized latent row force_splits: int | None = None, # tests only: 0 = single-program, N = split-k N + kv_quant: str | None = None, + pool_block_scale: torch.Tensor | None = None, ) -> torch.Tensor: """Sparse MLA attention over gathered latent rows; returns ``[b, m, h, d_v]``. @@ -359,13 +425,45 @@ def glm_dsa_sparse_attn( list across all queries with STRIDE 0 -- the identity-selection dense path: every query reads the same position-ordered row list, causally bounded by its own ``counts[q] = position + 1``, with zero per-query index materialization. + + NVFP4 restores tiles to Q's compute dtype before the dot products, with FP32 + accumulators. FP32 dot operands exceed consumer GPU shared memory at width 576. """ b, m, h, d = q.shape d_r = d - d_v topk = topk_idxs.shape[-1] - assert pool.shape[-1] == d, (pool.shape, d) + if kv_quant is None: + kv_quant = "fp8" if pool_scale is not None else "none" + assert kv_quant in ("none", "fp8", "nvfp4"), kv_quant + has_fp8 = kv_quant == "fp8" + has_nvfp4 = kv_quant == "nvfp4" + stored_dim = d // 2 if has_nvfp4 else d + assert pool.ndim == 2 and pool.shape[-1] == stored_dim, (pool.shape, d) + if has_fp8 or has_nvfp4: + assert pool_scale is not None + assert pool.dtype == torch.uint8, pool.dtype + assert pool_scale.shape == (pool.shape[0],), (pool_scale.shape, pool.shape) + assert pool_scale.dtype is torch.float32, pool_scale.dtype + scale_pool = pool_scale.contiguous() + else: + assert pool_scale is None + assert pool.is_floating_point() + scale_pool = pool + if has_nvfp4: + assert d % 16 == 0 + assert pool_block_scale is not None + assert pool_block_scale.shape == (pool.shape[0], d // 16) + assert pool_block_scale.dtype == torch.uint8 + assert pool_block_scale.device == pool.device + block_pool = pool_block_scale.contiguous() + else: + assert pool_block_scale is None + block_pool = pool + assert pool.device == q.device + if has_fp8 or has_nvfp4: + assert pool_scale.device == pool.device q = q.contiguous() - pool_2d = pool.reshape(-1, d) + pool_2d = pool.reshape(-1, stored_dim) assert pool_2d.stride(-1) == 1 idx = topk_idxs.contiguous().to(torch.int32) broadcast_m = idx.shape[1] == 1 and m > 1 @@ -379,25 +477,30 @@ def glm_dsa_sparse_attn( else: cnt, stride_nb, stride_nm = idx, 0, 0 + # Packed gathers need additional layout conversions at the 512-wide latent size. + block_t = 16 if has_nvfp4 else BLOCK_T n_splits = _split_count(b, m, h, topk, q.device) if force_splits is None else force_splits if n_splits: mid_o = q.new_empty(b, m, h, n_splits, d_v, dtype=torch.float32) mid_lse = q.new_empty(b, m, h, n_splits, dtype=torch.float32) grid1 = (m * n_splits, b, triton.cdiv(h, BLOCK_H)) _glm_dsa_splitk_kernel[grid1]( - q, pool_2d, mid_o, mid_lse, idx, cnt, + q, pool_2d, scale_pool, block_pool, mid_o, mid_lse, idx, cnt, float(softmax_scale), h, topk, q.stride(0), q.stride(1), q.stride(2), q.stride(3), pool_2d.stride(0), pool_2d.stride(1), + scale_pool.stride(0) if has_fp8 or has_nvfp4 else 0, mid_o.stride(0), mid_o.stride(1), mid_o.stride(2), mid_o.stride(3), mid_o.stride(4), mid_lse.stride(0), mid_lse.stride(1), mid_lse.stride(2), mid_lse.stride(3), idx.stride(0), 0 if broadcast_m else idx.stride(1), idx.stride(2), stride_nb, stride_nm, D_V=d_v, D_R=d_r, - BLOCK_H=BLOCK_H, BLOCK_T=BLOCK_T, - HAS_COUNTS=has_counts, HAS_ROPE=d_r > 0, NUM_SPLITS=n_splits, - num_warps=4, num_stages=2, + BLOCK_H=BLOCK_H, BLOCK_T=block_t, + HAS_COUNTS=has_counts, HAS_ROPE=d_r > 0, HAS_FP8=has_fp8, HAS_NVFP4=has_nvfp4, NUM_SPLITS=n_splits, + # Consumer Blackwell caps kernel shared memory at 99 KiB. Single-stage, the float path + # needs 100 KiB; two stages (its setting before quantized KV) fit. + num_warps=4, num_stages=1 if has_fp8 or has_nvfp4 else 2, ) grid2 = (m, b, h) _glm_dsa_merge_kernel[grid2]( @@ -412,18 +515,20 @@ def glm_dsa_sparse_attn( grid = (m, b, triton.cdiv(h, BLOCK_H)) _glm_dsa_sparse_kernel[grid]( - q, pool_2d, o, idx, cnt, + q, pool_2d, scale_pool, block_pool, o, idx, cnt, float(softmax_scale), h, topk, q.stride(0), q.stride(1), q.stride(2), q.stride(3), pool_2d.stride(0), pool_2d.stride(1), + scale_pool.stride(0) if has_fp8 or has_nvfp4 else 0, o.stride(0), o.stride(1), o.stride(2), o.stride(3), idx.stride(0), 0 if broadcast_m else idx.stride(1), idx.stride(2), stride_nb, stride_nm, D_V=d_v, D_R=d_r, - BLOCK_H=BLOCK_H, BLOCK_T=BLOCK_T, - HAS_COUNTS=has_counts, HAS_ROPE=d_r > 0, - num_warps=4, num_stages=2, + BLOCK_H=BLOCK_H, BLOCK_T=block_t, + HAS_COUNTS=has_counts, HAS_ROPE=d_r > 0, HAS_FP8=has_fp8, HAS_NVFP4=has_nvfp4, + # same shared-memory split as the split-k launch above + num_warps=4, num_stages=1 if has_fp8 or has_nvfp4 else 2, ) return o diff --git a/python/freetoken/kernel/triton/hc.py b/python/freetoken/kernel/triton/hc.py index 45d71bd731..acae534c45 100644 --- a/python/freetoken/kernel/triton/hc.py +++ b/python/freetoken/kernel/triton/hc.py @@ -326,9 +326,11 @@ def _hc_combine_norm_kernel( tl.store(out_ptr + row * stride_out + offs, out, mask=mask_inner) out = out.to(tl.float32) - # Keep the two-axis reduction: flattening the padded tile is ~40% slower - # at decode sizes. - sum_sq = tl.sum(tl.sum(out * out, axis=1), axis=0) + # The standalone grouped RMSNorm reduces one padded [next_pow2(HC_DIM)] + # vector; reshape keeps the identical tree (same element + # order, same masked zeros), so the fused statistic is bit-exact vs the split. + flat = tl.reshape(out * out, [NUM_TILES_PAD * BLOCK_SIZE]) + sum_sq = tl.sum(flat) rrms = tl.rsqrt(sum_sq / HC_DIM + EPS) if launch_pdl: diff --git a/python/freetoken/kernel/triton/kv_nvfp4.py b/python/freetoken/kernel/triton/kv_nvfp4.py new file mode 100644 index 0000000000..04727deec9 --- /dev/null +++ b/python/freetoken/kernel/triton/kv_nvfp4.py @@ -0,0 +1,154 @@ +"""Row-scaled NVFP4 KV storage, with low nibble first and 16-wide blocks. + +Each (token, head) row has an FP32 scale; each block has an E4M3 scale +stored as uint8. Reconstruction is E2M1 * block_scale * row_scale. +Row-local second-level scales keep appends independent of the cached prefix. +""" + +import torch +import triton +import triton.language as tl + +from .e4m3_compat import e4m3_f32_to_u8, e4m3_u8_to_f32, round_e4m3 + + +@triton.jit +def _encode_e2m1(x): + a = tl.abs(x) + code = tl.full(x.shape, 0, tl.int32) + code = tl.where(a > 0.25, 1, code) + code = tl.where(a >= 0.75, 2, code) + code = tl.where(a > 1.25, 3, code) + code = tl.where(a >= 1.75, 4, code) + code = tl.where(a > 2.5, 5, code) + code = tl.where(a >= 3.5, 6, code) + code = tl.where(a > 5.0, 7, code) + return code | tl.where(x < 0, 8, 0) + + +@triton.jit +def _decode_e2m1(code): + # Place E2M1 bits in FP16, then compensate the exponent-bias difference (15 - 1). + bits = ((code & 8).to(tl.uint16) << 12) | ((code & 7).to(tl.uint16) << 9) + return bits.to(tl.float16, bitcast=True).to(tl.float32) * 16384.0 + + +@triton.jit +def load_nvfp4(ptr, block_ptr, row_ptr, slots, head, dims, slot_mask, + stride_slot, stride_head, stride_row, D: tl.constexpr, + DIM_OFFSET: tl.constexpr = 0): + TRANSPOSE: tl.constexpr = dims.shape[0] != 1 + WIDTH: tl.constexpr = dims.shape[0] if TRANSPOSE else dims.shape[1] + TOKENS: tl.constexpr = slots.shape[0] * slots.shape[1] + slot = slots.reshape(TOKENS).to(tl.int64) + valid = slot_mask.reshape(TOKENS) + dim = DIM_OFFSET + tl.arange(0, WIDTH) + packed = tl.load( + ptr + slot[:, None] * stride_slot + head * stride_head + (dim[None, :] // 2), + valid[:, None] & (dim[None, :] < D), other=0, + ).to(tl.int32) + codes = tl.where((dim[None, :] & 1) == 0, packed & 15, packed >> 4) + block_dim = dim // 16 + block = tl.load( + block_ptr + (slot[:, None] * stride_row + head) * (D // 16) + block_dim[None, :], + valid[:, None] & (dim[None, :] < D), other=0, + ) + row = tl.load(row_ptr + slot * stride_row + head, valid, other=0) + value = _decode_e2m1(codes) * e4m3_u8_to_f32(block) * row[:, None] + if TRANSPOSE: + return tl.trans(value) + else: + return value + + +@triton.jit +def _quantize_row(src, dst, block_ptr, row_ptr, t, h, slot, stride_src, + HEADS: tl.constexpr, D: tl.constexpr, BLOCKS: tl.constexpr): + blocks = tl.arange(0, BLOCKS) + dims = blocks[:, None] * 16 + tl.arange(0, 16)[None, :] + x = tl.load(src + t * stride_src + h * D + dims, dims < D, other=0).to(tl.float32) + amax = tl.max(tl.abs(x), 1) + row_scale = tl.maximum(tl.max(amax, 0), 1e-10) / (6.0 * 448.0) + block_scale = round_e4m3(tl.minimum(tl.div_rn(amax, 6.0 * row_scale), 448.0)) + # Quantize against the scale actually stored, including E4M3 rounding/underflow. + denom = tl.where(block_scale > 0, block_scale * row_scale, 1.0) + normalized = tl.where(block_scale[:, None] > 0, tl.div_rn(x, denom[:, None]), 0.0) + codes = _encode_e2m1(normalized).reshape(BLOCKS, 8, 2) + lo, hi = tl.split(codes) + packed = lo | (hi << 4) + byte_dims = blocks[:, None] * 8 + tl.arange(0, 8)[None, :] + tl.store(dst + (slot * HEADS + h) * (D // 2) + byte_dims, packed, byte_dims < D // 2) + tl.store(block_ptr + (slot * HEADS + h) * (D // 16) + blocks, + e4m3_f32_to_u8(block_scale), blocks < D // 16) + tl.store(row_ptr + slot * HEADS + h, row_scale) + + +@triton.jit +def _scatter_rows(src, dst, block, row, indices, stride_src, + D: tl.constexpr, BLOCKS: tl.constexpr): + t = tl.program_id(0) + slot = tl.load(indices + t).to(tl.int64) + _quantize_row(src, dst, block, row, t, 0, slot, stride_src, 1, D, BLOCKS) + + +def quantize_nvfp4_rows_to_cache(rows, out_loc, cache, scales, block_scales) -> None: + """Quantize a single MLA latent slab, with one second-level scale per token.""" + tokens, dim = rows.shape + slots = cache.shape[0] + assert dim % 16 == 0 and rows.stride(1) == 1 + assert rows.dtype in (torch.float16, torch.bfloat16, torch.float32) + assert cache.shape == (slots, dim // 2) and cache.dtype == torch.uint8 + assert scales.shape == (slots,) and scales.dtype == torch.float32 + assert block_scales.shape == (slots, dim // 16) and block_scales.dtype == torch.uint8 + assert out_loc.shape == (tokens,) and out_loc.dtype in (torch.int32, torch.int64) + for tensor in (cache, scales, block_scales, out_loc): + assert tensor.is_contiguous() and tensor.device == rows.device + assert rows.is_cuda + if tokens: + _scatter_rows[(tokens,)]( + rows, cache, block_scales, scales, out_loc, rows.stride(0), + D=dim, BLOCKS=triton.next_power_of_2(dim // 16), + num_warps=4, enable_fp_fusion=False, + ) + + +@triton.jit +def _scatter(k, v, kc, vc, kb, vb, kr, vr, indices, stride_k, stride_v, + HEADS: tl.constexpr, D: tl.constexpr, BLOCKS: tl.constexpr): + t, h = tl.program_id(0), tl.program_id(1) + slot = tl.load(indices + t).to(tl.int64) + _quantize_row(k, kc, kb, kr, t, h, slot, stride_k, HEADS, D, BLOCKS) + _quantize_row(v, vc, vb, vr, t, h, slot, stride_v, HEADS, D, BLOCKS) + + +def quantize_nvfp4_to_cache( + k: torch.Tensor, + v: torch.Tensor, + out_loc: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + k_scale: torch.Tensor, + v_scale: torch.Tensor, + k_block_scale: torch.Tensor, + v_block_scale: torch.Tensor, +) -> None: + tokens, width = k.shape + slots, heads, packed_dim = k_cache.shape + dim = packed_dim * 2 + assert dim % 16 == 0 and width == heads * dim + assert v.shape == k.shape and k.stride(1) == v.stride(1) == 1 + assert k.dtype in (torch.float16, torch.bfloat16, torch.float32) and v.dtype == k.dtype + assert out_loc.shape == (tokens,) and out_loc.dtype in (torch.int32, torch.int64) + assert out_loc.is_contiguous() + for codes, row, block in ((k_cache, k_scale, k_block_scale), (v_cache, v_scale, v_block_scale)): + assert codes.shape == (slots, heads, packed_dim) and codes.dtype == torch.uint8 + assert row.shape == (slots, heads) and row.dtype == torch.float32 + assert block.shape == (slots, heads, dim // 16) and block.dtype == torch.uint8 + assert codes.is_contiguous() and row.is_contiguous() and block.is_contiguous() + assert codes.device == row.device == block.device == k.device + assert k.is_cuda and v.device == out_loc.device == k.device + if tokens: + _scatter[(tokens, heads)](k, v, k_cache, v_cache, k_block_scale, v_block_scale, + k_scale, v_scale, out_loc, k.stride(0), v.stride(0), + HEADS=heads, D=dim, BLOCKS=triton.next_power_of_2(dim // 16), + num_warps=4, enable_fp_fusion=False) diff --git a/python/freetoken/kernel/triton/kv_quant.py b/python/freetoken/kernel/triton/kv_quant.py new file mode 100644 index 0000000000..c10ad9c4b6 --- /dev/null +++ b/python/freetoken/kernel/triton/kv_quant.py @@ -0,0 +1,249 @@ +"""FP8 (e4m3) KV-cache storage: per-token, per-head symmetric quantization. + +One KV row is one ``(token, kv_head)`` slice of ``head_dim`` elements. It is stored +as ``head_dim`` e4m3 bytes plus ONE fp32 scale shared by the whole row: + + scale = max(amax(row) / 448, eps) # 448 == e4m3 finite max + code = round_e4m3(clamp(row / scale)) # RNE onto the e4m3 grid + read = code.to(f32) * scale # in the attention kernels + +Granularity rationale: an fp32 scale per (token, head) costs ``4 / head_dim`` bytes +per element (3% at head_dim 128, 6% at 64) while tracking each key's own magnitude, +which is what keeps a quantized KV from collapsing on outlier heads. A coarser +per-tensor scale needs no storage at all but has no headroom for them; a finer +per-element "scale" is the format itself. + +Architectures below sm_89 have no fp8e4nv type in Triton (see +:mod:`freetoken.kernel.triton.e4m3_compat`), so the codes live in a plain ``uint8`` +buffer on EVERY architecture and are decoded by :func:`e4m3_u8_to_f32`. That keeps one +set of bytes and one set of numbers across GPUs, and it is why :func:`kv_codes_dtype` +is a constant rather than a question: the fp8 type never appears in a kernel +signature, so nothing here can disagree with the host that allocated the buffer. +Choosing the encode/decode per target -- by an arch probe, or by testing the pointer's +element type -- is what broke this feature twice on real hardware (see +:func:`kv_load_e4m3_tile_f32`). + +The write path replaces ``kernel.store_cache`` for a quantized pool: the plain store +kernel is a raw byte copy that requires the source and the cache to share a dtype, +and quantization is exactly the step where the two diverge. Folding the scatter into +the quantization kernel keeps that to a single launch (and a single HBM round trip) +under CUDA-graph capture, where the slot ids arrive as a device tensor. +""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + +from freetoken.kernel.triton.e4m3_compat import ( + e4m3_f32_to_u8, + round_e4m3, +) + +FP8 = torch.float8_e4m3fn +KV_SCALE_DTYPE = torch.float32 +KV_QUANT_FP8 = "fp8" + + +def kv_codes_dtype() -> torch.dtype: + """Storage dtype of one quantized KV element: e4m3 bytes in a uint8 buffer. + + A constant, deliberately. Every attempt to answer this per target -- triton's + compile-time probe, then testing the pointer's element type at the load -- ended up + picking an arm that did not match the buffer the host had just allocated (see the + module header). torch still reads these bytes as fp8 whenever real numbers are + wanted: :func:`codes_to_f32`. + """ + return torch.uint8 + + +def alloc_codes(shape: tuple[int, ...], device: torch.device) -> torch.Tensor: + """A zero-filled code buffer of :func:`kv_codes_dtype` -- bytes, on every arch. + + Zero-filling matters because the pools read slots that were never written: a stale + byte decodes to a real number (0x7F/0xFF even to NaN once reinterpreted as fp8), + while 0x00 is exactly 0.0 (same reasoning as kvcache/bsa_pool.py). + """ + return torch.zeros(shape, dtype=kv_codes_dtype(), device=device) + + +def codes_to_f32(codes: torch.Tensor) -> torch.Tensor: + """Decode a code buffer to fp32 ON THE HOST (torch's own e4m3 cast). + + Works on either storage dtype -- uint8 bytes are reinterpreted as fp8 first -- so + a test or a debugging tool reads the same numbers on sm_86 as on sm_90, and does + it through torch rather than through the software decoder it is checking. + """ + if codes.dtype is not FP8: + codes = codes.view(FP8) + return codes.to(torch.float32) + + +@triton.jit +def _kv_quant_scatter_kernel( + k_src, + v_src, + k_dst, + v_dst, + k_scale, + v_scale, + idx_ptr, + stride_xs, # K source row pitch, in elements (the qkv slice is wider than one row) + stride_vx, # V source row pitch. K and V need not share one: a .clamp() on one side + # leaves it densely packed, so reusing K's pitch reads V off its rows. + stride_kd, # K cache row pitch, in elements (== HEADS * D) + stride_vd, + stride_ks, # scale row pitch, in elements (== HEADS) + stride_vs, + D: tl.constexpr, + BLOCK_D: tl.constexpr, +): + """One program per (token, kv_head): quantize the row and write it to slot + ``idx_ptr[token]`` (the same ``out_loc`` the bf16 store scatters through).""" + t = tl.program_id(0) + h = tl.program_id(1) + pos = tl.load(idx_ptr + t).to(tl.int64) # int64: slots * row can pass 2**31 + d = tl.arange(0, BLOCK_D) + mask = d < D + + off = h * D + d + xk = tl.load(k_src + t * stride_xs + off, mask=mask, other=0.0).to(tl.float32) + xv = tl.load(v_src + t * stride_vx + off, mask=mask, other=0.0).to(tl.float32) + + # 448 == e4m3 finite max; 1e-10 is the amax floor of the activation quant in + # kernel/triton/fp8_block_linear.py (literals keep the kernel self-contained). + sk = tl.maximum(tl.max(tl.abs(xk), axis=0), 1e-10) / 448.0 + sv = tl.maximum(tl.max(tl.abs(xv), axis=0), 1e-10) / 448.0 + qk = tl.clamp(xk / sk, -448.0, 448.0) + qv = tl.clamp(xv / sv, -448.0, 448.0) + + # Straight-line, like the reader: round onto the e4m3 grid in ONE step (RNE) and + # pack the bits into the byte buffer. No fp8 type on either side -- the reason this + # feature stopped compiling twice is explained in kv_codes_dtype's docstring. + out_k = e4m3_f32_to_u8(round_e4m3(qk)) + out_v = e4m3_f32_to_u8(round_e4m3(qv)) + + tl.store(k_dst + pos * stride_kd + h * D + d, out_k, mask=mask) + tl.store(v_dst + pos * stride_vd + h * D + d, out_v, mask=mask) + tl.store(k_scale + pos * stride_ks + h, sk) + tl.store(v_scale + pos * stride_vs + h, sv) + + +def quantize_kv_to_cache( + k: torch.Tensor, + v: torch.Tensor, + out_loc: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + k_scale: torch.Tensor, + v_scale: torch.Tensor, +) -> None: + """Quantize fresh K/V rows into an fp8 KV pool. + + ``k``/``v`` : ``[T, num_kv_heads * head_dim]`` compute-dtype rows -- exactly what + the attention backends hand to ``store_kv`` (a slice of the qkv projection, + so the row pitch may be wider than the row itself). + ``out_loc`` : ``[T]`` device slot index per row (int32 or int64). + ``*_cache`` : ``[num_slots, num_kv_heads, head_dim]`` of :func:`kv_codes_dtype`. + ``*_scale`` : ``[num_slots, num_kv_heads]`` fp32, indexed by the SAME slot. + """ + tokens = k.shape[0] + if tokens == 0: + return + assert k.dim() == 2 and v.shape == k.shape, (k.shape, v.shape) + assert k.stride(1) == 1 and v.stride(1) == 1, "K/V rows must be contiguous" + heads, dim = k_cache.shape[1], k_cache.shape[2] + assert k.shape[1] == heads * dim, (tuple(k.shape), tuple(k_cache.shape)) + assert k_cache.shape == v_cache.shape, (tuple(k_cache.shape), tuple(v_cache.shape)) + assert k_scale.shape == v_scale.shape == (k_cache.shape[0], heads), ( + tuple(k_scale.shape), + tuple(k_cache.shape), + ) + assert k_cache.dtype == v_cache.dtype == kv_codes_dtype(), ( + k_cache.dtype, + kv_codes_dtype(), + ) + assert k_scale.dtype == KV_SCALE_DTYPE, k_scale.dtype + _kv_quant_scatter_kernel[(tokens, heads)]( + k, + v, + k_cache, + v_cache, + k_scale, + v_scale, + out_loc, + k.stride(0), + v.stride(0), + k_cache.stride(0), + v_cache.stride(0), + k_scale.stride(0), + v_scale.stride(0), + D=dim, + BLOCK_D=triton.next_power_of_2(dim), + num_warps=1, + ) + + +@triton.jit +def _kv_quant_rows_scatter_kernel( + src, dst, scale_ptr, idx_ptr, + stride_src, stride_dst, stride_scale, + D: tl.constexpr, + BLOCK_D: tl.constexpr, +): + """Quantize one whole latent row per program and scatter it by ``out_loc``.""" + t = tl.program_id(0) + pos = tl.load(idx_ptr + t).to(tl.int64) + d = tl.arange(0, BLOCK_D) + mask = d < D + x = tl.load(src + t * stride_src + d, mask=mask, other=0.0).to(tl.float32) + s = tl.maximum(tl.max(tl.abs(x), axis=0), 1e-10) / 448.0 + q = tl.clamp(x / s, -448.0, 448.0) + tl.store(dst + pos * stride_dst + d, e4m3_f32_to_u8(round_e4m3(q)), mask=mask) + tl.store(scale_ptr + pos * stride_scale, s) + + +def quantize_rows_to_cache( + rows: torch.Tensor, + out_loc: torch.Tensor, + cache: torch.Tensor, + scales: torch.Tensor, +) -> None: + """Quantize and scatter one scale-bearing FP8 row per token. + + MLA stores its absorbed K/V state as one latent row, rather than separate K and + V heads. Its scale granularity is therefore one FP32 value per ``(token, + layer)`` row, which is the single-head MLA case already priced by + ``kv_scale_bytes_per_token``. + """ + tokens = rows.shape[0] + if tokens == 0: + return + assert rows.dim() == 2 and rows.stride(1) == 1, tuple(rows.shape) + assert cache.dim() == 2 and cache.shape[1] == rows.shape[1], ( + tuple(cache.shape), tuple(rows.shape), + ) + assert scales.shape == (cache.shape[0],), (tuple(scales.shape), tuple(cache.shape)) + assert cache.dtype == kv_codes_dtype(), cache.dtype + assert scales.dtype == KV_SCALE_DTYPE, scales.dtype + dim = rows.shape[1] + _kv_quant_rows_scatter_kernel[(tokens,)]( + rows, cache, scales, out_loc, + rows.stride(0), cache.stride(0), scales.stride(0), + D=dim, BLOCK_D=triton.next_power_of_2(dim), + num_warps=4 if dim > 256 else 1, + ) + + +__all__ = [ + "FP8", + "KV_QUANT_FP8", + "KV_SCALE_DTYPE", + "alloc_codes", + "codes_to_f32", + "kv_codes_dtype", + "quantize_kv_to_cache", + "quantize_rows_to_cache", +] + diff --git a/python/freetoken/kernel/triton/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/mxfp4_moe.py b/python/freetoken/kernel/triton/mxfp4_moe.py index 3981d2a391..c0237f25a1 100644 --- a/python/freetoken/kernel/triton/mxfp4_moe.py +++ b/python/freetoken/kernel/triton/mxfp4_moe.py @@ -53,30 +53,36 @@ def mxfp4_splitk_gemv_kernel( expert_id = tl.load(expert_ids_ptr + pid_e) offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) mask_n = offs_n < N + # w_ptr / s_ptr are int16 views: one element holds the bytes of columns 2p (low) and 2p+1 (high) + offs_p = pid_n * (BLOCK_N // 2) + tl.arange(0, BLOCK_N // 2) + mask_p = offs_p < N // 2 k_groups = K // 32 kg0 = pid_k * K_GROUPS_PER_SPLIT kg1 = tl.minimum(kg0 + K_GROUPS_PER_SPLIT, k_groups) - acc = tl.zeros([BLOCK_N], dtype=tl.float32) + acc_even = tl.zeros([BLOCK_N // 2], dtype=tl.float32) + acc_odd = tl.zeros([BLOCK_N // 2], dtype=tl.float32) w_base = expert_id * stride_we s_base = expert_id * stride_se x_base = pid_e * stride_xe for kg in range(kg0, kg1): - s_val = tl.load(s_ptr + s_base + kg * stride_sk + offs_n, mask=mask_n, other=0) - scale_f = _e8m0_scale(s_val) + s_pair = tl.load(s_ptr + s_base + kg * stride_sk + offs_p, mask=mask_p, other=0).to(tl.int32) + scale_even = _e8m0_scale(s_pair & 0xFF) + scale_odd = _e8m0_scale((s_pair >> 8) & 0xFF) for kk in tl.static_range(16): k_packed = kg * 16 + kk - w_byte = tl.load( - w_ptr + w_base + k_packed * stride_wk + offs_n, mask=mask_n, other=0 + w_pair = tl.load( + w_ptr + w_base + k_packed * stride_wk + offs_p, mask=mask_p, other=0 ).to(tl.int32) - lo = w_byte & 0x0F - hi = (w_byte >> 4) & 0x0F x_lo = tl.load(x_ptr + x_base + kg * 32 + kk * 2).to(tl.float32) x_hi = tl.load(x_ptr + x_base + kg * 32 + kk * 2 + 1).to(tl.float32) - acc += _fp4_table_lut(lo, lut_ptr) * scale_f * x_lo - acc += _fp4_table_lut(hi, lut_ptr) * scale_f * x_hi + acc_even += _fp4_table_lut(w_pair & 0x0F, lut_ptr) * scale_even * x_lo + acc_even += _fp4_table_lut((w_pair >> 4) & 0x0F, lut_ptr) * scale_even * x_hi + acc_odd += _fp4_table_lut((w_pair >> 8) & 0x0F, lut_ptr) * scale_odd * x_lo + acc_odd += _fp4_table_lut((w_pair >> 12) & 0x0F, lut_ptr) * scale_odd * x_hi + acc = tl.reshape(tl.join(acc_even, acc_odd), (BLOCK_N,)) if HAS_BIAS and pid_k == 0: acc += tl.load( diff --git a/python/freetoken/kernel/triton/norm.py b/python/freetoken/kernel/triton/norm.py index 3f95c29f0d..9071e1dfff 100644 --- a/python/freetoken/kernel/triton/norm.py +++ b/python/freetoken/kernel/triton/norm.py @@ -30,7 +30,7 @@ import triton.language as tl from triton.language.extra.cuda import gdc_launch_dependents, gdc_wait -from freetoken.utils.arch import is_sm90_supported +from freetoken.utils.arch import is_rocm, is_sm90_supported _HEUR = {"BLOCK": lambda a: triton.next_power_of_2(a["H"])} @@ -144,7 +144,8 @@ def _rmsnorm(input, weight, eps, out, gemma: bool): pdl = contig and is_sm90_supported() _rmsnorm_kernel[(A, B)]( out, input, weight, eps, H, sxa, sxb, soa, sob, - CONTIG=contig, ENABLE_PDL=pdl, launch_pdl=pdl, GEMMA=gemma, + CONTIG=contig, ENABLE_PDL=pdl, GEMMA=gemma, + **({} if is_rocm() else {"launch_pdl": pdl}), num_warps=_num_warps(A * B), num_stages=1, ) return out @@ -172,7 +173,8 @@ def _fused_add_rmsnorm(input, residual, weight, eps, gemma: bool): pdl = contig and is_sm90_supported() _fused_add_rmsnorm_kernel[(A, B)]( input, residual, weight, eps, H, sxa, sxb, sra, srb, - CONTIG=contig, ENABLE_PDL=pdl, launch_pdl=pdl, GEMMA=gemma, + CONTIG=contig, ENABLE_PDL=pdl, GEMMA=gemma, + **({} if is_rocm() else {"launch_pdl": pdl}), num_warps=_num_warps(A * B), num_stages=1, ) diff --git a/python/freetoken/kernel/triton/nvfp4_act_quant.py b/python/freetoken/kernel/triton/nvfp4_act_quant.py new file mode 100644 index 0000000000..6f56a46068 --- /dev/null +++ b/python/freetoken/kernel/triton/nvfp4_act_quant.py @@ -0,0 +1,68 @@ +"""NVFP4 activation quantization under a static global scale: x ~= e2m1 * e4m3 block scale (per 16) * g.""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _e2m1x2(hi, lo): + # sm_100+ / sm_120a; cvt packs its first source into the high nibble + return tl.inline_asm_elementwise( + "{ .reg .b8 t; cvt.rn.satfinite.e2m1x2.f32 t, $1, $2; cvt.u16.u8 $0, t; }", + "=h,r,r", [hi, lo], dtype=tl.uint16, is_pure=True, pack=1) + + +@triton.jit +def _quant_kernel(x_ptr, q_ptr, s_ptr, g_ptr, R, K, stride_x, k_tiles, + SWIZZLED: tl.constexpr, BLOCK_M: tl.constexpr, BLOCK_K: tl.constexpr): + # with a static global every 16-wide block quantizes on its own, so a row splits over programs + rows = tl.program_id(0) * BLOCK_M + tl.arange(0, BLOCK_M) + cols = tl.program_id(1) * BLOCK_K + tl.arange(0, BLOCK_K) + row_ok = rows < R + r64 = rows.to(tl.int64) + x = tl.load(x_ptr + r64[:, None] * stride_x + cols[None, :], mask=row_ok[:, None] & (cols < K)[None, :], other=0.0) + g = tl.load(g_ptr) + xb = tl.reshape(x.to(tl.float32), [BLOCK_M, BLOCK_K // 16, 16]) + s = (tl.max(tl.abs(xb), 2) / (6.0 * g)).to(tl.float8e4nv) + sf = s.to(tl.float32) * g + inv = tl.where(sf > 0, 1.0 / sf, 0.0) + lo, hi = tl.split(tl.reshape(xb * inv[:, :, None], [BLOCK_M, BLOCK_K // 2, 2])) + oq = tl.program_id(1) * (BLOCK_K // 2) + tl.arange(0, BLOCK_K // 2) + tl.store(q_ptr + r64[:, None] * (K // 2) + oq[None, :], _e2m1x2(hi, lo).to(tl.uint8), + mask=row_ok[:, None] & (oq < K // 2)[None, :]) + sc = tl.program_id(1) * (BLOCK_K // 16) + tl.arange(0, BLOCK_K // 16) + if SWIZZLED: + # cuBLAS / CUTLASS 128x4 block-scale atoms: [row tile, k tile, row % 32, (row % 128) // 32, col % 4] + s_off = ((r64 // 128)[:, None] * k_tiles + (sc // 4)[None, :]) * 512 \ + + ((rows % 32) * 16 + (rows % 128) // 32 * 4)[:, None] + (sc % 4)[None, :] + else: + s_off = r64[:, None] * (K // 16) + sc[None, :] + tl.store(s_ptr + s_off, s.to(tl.uint8, bitcast=True), mask=row_ok[:, None] & (sc < K // 16)[None, :]) + + +def quantize_nvfp4(x: torch.Tensor, g: torch.Tensor, swizzled: bool = False) -> tuple[torch.Tensor, torch.Tensor]: + """x [R, K] -> (codes [R, K / 2] uint8, low nibble first; e4m3 block scales) under the fp32 global g [1]; K % 16 == 0. + + The scales are [R, K / 16], or with ``swizzled`` the flat cuBLAS / CUTLASS 128x4 layout that ``torch.nn.functional.scaled_mm`` takes, + rows padded to 128 and K / 16 to 4 with zeros. + """ + R, K = x.shape + assert K % 16 == 0 and x.stride(1) == 1 + q = torch.empty(R, K // 2, dtype=torch.uint8, device=x.device) + k_tiles = triton.cdiv(K // 16, 4) + if swizzled: + n = triton.cdiv(R, 128) * 128 * k_tiles * 4 + exact = R % 128 == 0 and (K // 16) % 4 == 0 + s = (torch.empty if exact else torch.zeros)(n, dtype=torch.uint8, device=x.device) + else: + s = torch.empty(R, K // 16, dtype=torch.uint8, device=x.device) + if R: + block_k = min(triton.next_power_of_2(K), 1024) + block_m = max(1, 1024 // block_k) + grid = (triton.cdiv(R, block_m), triton.cdiv(K, block_k)) + _quant_kernel[grid](x, q, s, g, R, K, x.stride(0), k_tiles, SWIZZLED=swizzled, BLOCK_M=block_m, BLOCK_K=block_k, + num_warps=4) + return q, s.view(torch.float8_e4m3fn) diff --git a/python/freetoken/kernel/triton/nvfp4_dequant.py b/python/freetoken/kernel/triton/nvfp4_dequant.py index a949e7fa0a..214e6ab71e 100644 --- a/python/freetoken/kernel/triton/nvfp4_dequant.py +++ b/python/freetoken/kernel/triton/nvfp4_dequant.py @@ -26,7 +26,7 @@ def _e2m1_lut(device_index: int) -> torch.Tensor: def _dequant_nvfp4_kernel( packed_ptr, # [S, OUT, IN // 2] uint8 scale_ptr, # [S, OUT, IN // 16] fp8-e4m3 per-block scale (checkpoint native) - global_ptr, # [S, OUT] fp16 per-output-row global scale (weight_scale_2) + global_ptr, # [S, OUT] fp32 per-output-row global scale (weight_scale_2) slots_ptr, # [N] int32 -> cache slot for each output expert out_ptr, # [N, OUT, IN] compute dtype lut_ptr, # [16] float32 E2M1 values diff --git a/python/freetoken/kernel/triton/nvfp4_fused_moe.py b/python/freetoken/kernel/triton/nvfp4_fused_moe.py index a1064bcf7d..08f2f77702 100644 --- a/python/freetoken/kernel/triton/nvfp4_fused_moe.py +++ b/python/freetoken/kernel/triton/nvfp4_fused_moe.py @@ -1,14 +1,14 @@ """Inline-dequant NVFP4 fused-MoE Triton kernels. These kernels read the NVFP4 expert cache (packed FP4 codes + fp8 block scales + -fp16 per-row global scale) *directly* and dequantize inside the GEMM K-loop, so the +fp32 per-row global scale) *directly* and dequantize inside the GEMM K-loop, so the grouped MoE GEMM never materializes a BF16 copy of the experts (no separate dequant pass, no HBM round-trip of the 4x larger BF16 weights). Layout, for an expert weight ``W[N, K]``: - ``packed[slot, n, k//2]`` uint8: two FP4 codes per byte, low nibble = even k. - ``scale[slot, n, k//16]`` fp8-e4m3: per 16-wide block scale. - - ``global[slot, n]`` fp16: per-output-row scale (``weight_scale_2``). + - ``global[slot, n]`` fp32: per-output-row scale (``weight_scale_2``). - ``W[n, k] = E2M1[code] * scale[n, k//16] * global[n]``. The global scale is constant along K, so it is applied once after the K-loop. @@ -28,8 +28,11 @@ import torch import triton import triton.language as tl +from triton.language import target_info +from triton.runtime.jit import constexpr_function from freetoken.kernel.triton.e4m3_compat import e4m3_native_cx, e4m3_u8_to_f32 +from freetoken.kernel.triton.nvfp4_linear import _nvfp4_pair_f32 _E2M1_VALUES = [ 0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, @@ -49,7 +52,7 @@ def _decode_nvfp4_moe_kernel( a_ptr, # [M, K] activations (compute dtype) packed_ptr, # [S, N, K // 2] uint8 scale_ptr, # [S, N, K // 16] fp8-e4m3 - global_ptr, # [S, N] fp16 + global_ptr, # [S, N] fp32 c_ptr, # [M, TOP_K, N] output (compute dtype) topk_weights_ptr, # [M, TOP_K] fp32 topk_ids_ptr, # [M, TOP_K] int32 -> cache slot @@ -130,16 +133,40 @@ def _decode_nvfp4_moe_kernel( tl.store(c_ptrs, accumulator.to(compute_type), mask=(route_id < total_routes) & n_mask) +@triton.jit +def _bf16x4_to_f32(v): + """Four bf16 packed in one int64 -> four fp32 (a bf16 is the high half of its fp32).""" + lo = v.to(tl.int32) + hi = (v >> 32).to(tl.int32) + return ( + (lo << 16).to(tl.float32, bitcast=True), + (lo & -65536).to(tl.float32, bitcast=True), + (hi << 16).to(tl.float32, bitcast=True), + (hi & -65536).to(tl.float32, bitcast=True), + ) + + +@triton.jit +def _f16x4_to_f32(v): + lo = v.to(tl.int32) + hi = (v >> 32).to(tl.int32) + return ( + lo.to(tl.int16).to(tl.float16, bitcast=True).to(tl.float32), + (lo >> 16).to(tl.int16).to(tl.float16, bitcast=True).to(tl.float32), + hi.to(tl.int16).to(tl.float16, bitcast=True).to(tl.float32), + (hi >> 16).to(tl.int16).to(tl.float16, bitcast=True).to(tl.float32), + ) + + @triton.jit def _decode_nvfp4_marlin_kernel( a_ptr, # [M, K] activations (compute dtype) packed_ptr, # [S, N, K // 8] int32 (8 fp4 codes per word, nibble j -> k=8*w+j) scale_ptr, # [S, N, K // 16] fp8-e4m3 - global_ptr, # [S, N] fp16 + global_ptr, # [S, N] fp32 c_ptr, # [M, TOP_K, N] output (compute dtype) topk_weights_ptr, # [M, TOP_K] fp32 topk_ids_ptr, # [M, TOP_K] int32 -> cache slot - lut_ptr, # [16] fp32 total_routes, N, K, @@ -151,9 +178,11 @@ def _decode_nvfp4_marlin_kernel( stride_tw_m, stride_tw_k, stride_tid_m, stride_tid_k, BLOCK_SIZE_N: tl.constexpr, - BLOCK_SIZE_KW: tl.constexpr, # int32 words per K-iter (covers 8*KW k-values) + BLOCK_SIZE_KW: tl.constexpr, # int32 words per tile (covers 8*KW k-values) + UNROLL: tl.constexpr, # tiles per loop trip TOP_K: tl.constexpr, A_ROW_IS_ROUTE: tl.constexpr, + A_WIDE: tl.constexpr, # activation rows are contiguous bf16/fp16, 16-byte aligned MUL_ROUTED_WEIGHT: tl.constexpr, compute_type: tl.constexpr, ): @@ -165,13 +194,18 @@ def _decode_nvfp4_marlin_kernel( * Loads the FP4 codes as **int32 words** (8 codes / 4 bytes per element) so the HBM read issues wide coalesced transactions -- the single biggest lever for the gate/up GEMM, whose mem-only ceiling jumps from ~65% to ~85% of peak read BW. - * **Defers** the K reduction: the ``[BLOCK_KW, BLOCK_N]`` partial accumulates across + * **Defers** the K reduction: the ``[BLOCK_N, BLOCK_KW]`` partial accumulates across *all* K-iters and is reduced once at the end, removing the per-iter ``tl.sum`` barrier so the weight loads of successive iters pipeline. - The e2m1 codes share one fp8 block scale per 16 k-values (== 2 words), applied per - word. Mirrors the fast kernel's route/tile mapping and epilogue so it is a drop-in - decode GEMV (CUDA-graph safe: fixed shapes, no host sync).""" + Every load in the K-loop is shaped so the compiler keeps the whole loop in the word + tile's register layout (no shared-memory layout conversions): the block scales are + loaded as a half tile and interleaved in registers (one e4m3 scale per 16 k-values + == 2 words), the e2m1 codes are decoded with the fp16 bit trick of + :func:`_nvfp4_pair_f32` (2^14 folded into the scale), and with ``A_WIDE`` the 8 + activations of a word are read as two int64 and unpacked by shifts. Mirrors the fast + kernel's route/tile mapping and epilogue so it is a drop-in decode GEMV (CUDA-graph + safe: fixed shapes, no host sync).""" route_id = tl.program_id(0) n_block_id = tl.program_id(1) token_id = route_id // TOP_K @@ -185,37 +219,70 @@ def _decode_nvfp4_marlin_kernel( a_base = a_ptr + a_row * stride_am offs_kw = tl.arange(0, BLOCK_SIZE_KW) + offs_kh = tl.arange(0, BLOCK_SIZE_KW // 2) K_WORDS = K // 8 - partial = tl.zeros((BLOCK_SIZE_KW, BLOCK_SIZE_N), dtype=tl.float32) - - packed_slot = packed_ptr + slot * stride_pe - scale_slot = scale_ptr + slot * stride_se - for kw_start in range(0, tl.cdiv(K_WORDS, BLOCK_SIZE_KW)): - widx = kw_start * BLOCK_SIZE_KW + offs_kw - w_mask = widx < K_WORDS - - word = tl.load( - packed_slot + offs_n[None, :] * stride_pn + widx[:, None] * stride_pkw, - mask=w_mask[:, None] & n_mask[None, :], other=0, - ) - # 8 codes/word fall in the same or adjacent 16-wide block -> one scale per word. - s_ptrs = scale_slot + offs_n[None, :] * stride_sn + (widx[:, None] // 2) * stride_sblk - s_mask = w_mask[:, None] & n_mask[None, :] - if e4m3_native_cx(): - scale = tl.load(s_ptrs, mask=s_mask, other=0.0).to(tl.float32) - else: - scale = e4m3_u8_to_f32(tl.load(s_ptrs, mask=s_mask, other=0)) - - kbase = 8 * widx - acc_w = tl.zeros((BLOCK_SIZE_KW, BLOCK_SIZE_N), dtype=tl.float32) - for j in tl.static_range(8): - code = (word >> (4 * j)) & 0xF - b = tl.load(lut_ptr + code) - a_j = tl.load(a_base + (kbase + j) * stride_ak, mask=w_mask, other=0.0).to(tl.float32) - acc_w += a_j[:, None] * b - partial += acc_w * scale - - accumulator = tl.sum(partial, axis=0) + partial = tl.zeros((BLOCK_SIZE_N, BLOCK_SIZE_KW), dtype=tl.float32) + + p_base = packed_ptr + slot * stride_pe + offs_n[:, None] * stride_pn + s_base = scale_ptr + slot * stride_se + offs_n[:, None] * stride_sn + for kw_start in range(0, tl.cdiv(K_WORDS, BLOCK_SIZE_KW), UNROLL): + for u in tl.static_range(UNROLL): + widx = (kw_start + u) * BLOCK_SIZE_KW + offs_kw + w_mask = widx < K_WORDS + word = tl.load( + p_base + widx[None, :] * stride_pkw, + mask=n_mask[:, None] & w_mask[None, :], other=0, + ) + hidx = (kw_start + u) * (BLOCK_SIZE_KW // 2) + offs_kh + s_ptrs = s_base + hidx[None, :] * stride_sblk + s_mask = n_mask[:, None] & (hidx[None, :] < K_WORDS // 2) + if e4m3_native_cx(): + s_half = tl.load(s_ptrs, mask=s_mask, other=0.0).to(tl.float32) + else: + s_half = e4m3_u8_to_f32(tl.load(s_ptrs, mask=s_mask, other=0)) + s_half = s_half * 16384.0 + scale = tl.interleave(s_half, s_half) + + if A_WIDE: + a64 = a_base.to(tl.pointer_type(tl.int64), bitcast=True) + v = tl.load( + a64 + widx[:, None] * 2 + tl.arange(0, 2)[None, :], + mask=w_mask[:, None], other=0, + ) + v0, v1 = tl.split(v) + if a_ptr.dtype.element_ty == tl.bfloat16: + a0, a1, a2, a3 = _bf16x4_to_f32(v0) + a4, a5, a6, a7 = _bf16x4_to_f32(v1) + else: + a0, a1, a2, a3 = _f16x4_to_f32(v0) + a4, a5, a6, a7 = _f16x4_to_f32(v1) + else: + kbase = 8 * widx + a0 = tl.load(a_base + kbase * stride_ak, mask=w_mask, other=0.0).to(tl.float32) + a1 = tl.load(a_base + (kbase + 1) * stride_ak, mask=w_mask, other=0.0).to(tl.float32) + a2 = tl.load(a_base + (kbase + 2) * stride_ak, mask=w_mask, other=0.0).to(tl.float32) + a3 = tl.load(a_base + (kbase + 3) * stride_ak, mask=w_mask, other=0.0).to(tl.float32) + a4 = tl.load(a_base + (kbase + 4) * stride_ak, mask=w_mask, other=0.0).to(tl.float32) + a5 = tl.load(a_base + (kbase + 5) * stride_ak, mask=w_mask, other=0.0).to(tl.float32) + a6 = tl.load(a_base + (kbase + 6) * stride_ak, mask=w_mask, other=0.0).to(tl.float32) + a7 = tl.load(a_base + (kbase + 7) * stride_ak, mask=w_mask, other=0.0).to(tl.float32) + + acc_w = tl.zeros((BLOCK_SIZE_N, BLOCK_SIZE_KW), dtype=tl.float32) + b0, b4 = _nvfp4_pair_f32(word) + b1, b5 = _nvfp4_pair_f32(word >> 4) + b2, b6 = _nvfp4_pair_f32(word >> 8) + b3, b7 = _nvfp4_pair_f32(word >> 12) + acc_w += a0[None, :] * b0 + acc_w += a1[None, :] * b1 + acc_w += a2[None, :] * b2 + acc_w += a3[None, :] * b3 + acc_w += a4[None, :] * b4 + acc_w += a5[None, :] * b5 + acc_w += a6[None, :] * b6 + acc_w += a7[None, :] * b7 + partial += acc_w * scale + + accumulator = tl.sum(partial, axis=1) g = tl.load(global_ptr + slot * stride_ge + offs_n * stride_gn, mask=n_mask, other=0.0).to(tl.float32) accumulator = accumulator * g @@ -227,18 +294,77 @@ def _decode_nvfp4_marlin_kernel( tl.store(c_ptrs, accumulator.to(compute_type), mask=(route_id < total_routes) & n_mask) +@triton.jit +def _e2m1_f16_x2p14(code): + """e2m1 code (low 4 bits) -> fp16 value * 2^-14: magnitude bits go to fp16 [11:9], the sign to [15]; fp16 subnormals line up the e == 0 codes.""" + bits = ((code & 7) << 9) | ((code & 8) << 12) + return bits.to(tl.uint16).to(tl.float16, bitcast=True) + + +@constexpr_function +def _bf16x2_mul_native(): + return target_info.cuda_capability_geq(9, 0) + + +@triton.jit +def _bf16x2_scale(w, s2): + if _bf16x2_mul_native(): + return tl.inline_asm_elementwise("mul.rn.bf16x2 $0, $1, $2;", "=r,r,r", [w, s2], tl.uint32, True, 1) + else: + # mul.bf16x2 is sm_90+; fma with a -0.0 addend gives the same product + return tl.inline_asm_elementwise("{ .reg .b32 z; mov.b32 z, 0x80008000; fma.rn.bf16x2 $0, $1, $2, z; }", + "=r,r,r", [w, s2], tl.uint32, True, 1) + + +@triton.jit +def _bf16x2_halves(w): + lo = (w & 0xFFFF).to(tl.uint16).to(tl.bfloat16, bitcast=True) + hi = (w >> 16).to(tl.uint16).to(tl.bfloat16, bitcast=True) + return lo, hi + + +@triton.jit +def _e2m1x8_bf16(q, s2, BLOCK_N: tl.constexpr, BLOCK_W: tl.constexpr): + """[BLOCK_N, BLOCK_W] uint32 words of 8 e2m1 codes -> [BLOCK_N, 8 * BLOCK_W] bf16 in k order, times s2 * 2^-126.""" + # each 16-bit half's top nibble becomes a bf16 sign + exponent/mantissa: e2m1 * 2^-126, exact incl. e == 0 + o0 = _bf16x2_scale((q & 0x80008000) | ((q >> 6) & 0x01C001C0), s2) + q1 = q << 4 + o1 = _bf16x2_scale((q1 & 0x80008000) | ((q1 >> 6) & 0x01C001C0), s2) + q2 = q << 8 + o2 = _bf16x2_scale((q2 & 0x80008000) | ((q2 >> 6) & 0x01C001C0), s2) + q3 = q << 12 + o3 = _bf16x2_scale((q3 & 0x80008000) | ((q3 >> 6) & 0x01C001C0), s2) + # o0..o3 hold the k pairs (3, 7), (2, 6), (1, 5), (0, 4); join appends the new axis last + k0, k4 = _bf16x2_halves(o3) + k1, k5 = _bf16x2_halves(o2) + k2, k6 = _bf16x2_halves(o1) + k3, k7 = _bf16x2_halves(o0) + w = tl.join(tl.join(tl.join(k0, k4), tl.join(k2, k6)), tl.join(tl.join(k1, k5), tl.join(k3, k7))) + return tl.reshape(w, (BLOCK_N, 8 * BLOCK_W)) + + +@triton.jit +def _e2m1x2_bf16(v, s2, BLOCK_N: tl.constexpr, BLOCK_KB: tl.constexpr): + """[BLOCK_N, BLOCK_KB] bytes of 2 e2m1 codes -> [BLOCK_N, 2 * BLOCK_KB] bf16 in k order, times s2 * 2^-126.""" + # low nibble -> low bf16, high nibble -> high bf16: sign to bit 15 / 31, exponent and mantissa to bits 8..6 / 24..22 + pair = ((v & 0x8) << 12) | ((v & 0x7) << 6) | ((v & 0x80) << 24) | ((v & 0x70) << 18) + lo, hi = _bf16x2_halves(_bf16x2_scale(pair.to(tl.uint32, bitcast=True), s2)) + return tl.reshape(tl.join(lo, hi), (BLOCK_N, 2 * BLOCK_KB)) + + @triton.jit def _prefill_nvfp4_moe_kernel( a_ptr, # [M, K] activations packed_ptr, # [S, N, K // 2] uint8 - scale_ptr, # [S, N, K // 16] fp8-e4m3 - global_ptr, # [S, N] fp16 + scale_ptr, # [S, N, K // 16] fp8-e4m3 (uint8 view before sm_89) + global_ptr, # [S, N] fp32 c_ptr, # [num_valid_tokens, N] output (flat over M*top_k) topk_weights_ptr, sorted_token_ids_ptr, expert_ids_ptr, # cache slot per M-block num_tokens_post_padded_ptr, - lut_ptr, + w_desc, # TMA descriptor over packed as [S * N, K // 2], or None + s_desc, # TMA descriptor over scale as [S * N, K // 16], or None N, K, EM, @@ -249,14 +375,20 @@ def _prefill_nvfp4_moe_kernel( stride_ge, stride_gn, stride_cm, stride_cn, stride_tw, - BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_M: tl.constexpr, # tokens per block, the moe_align padding BLOCK_SIZE_N: tl.constexpr, - BLOCK_SIZE_KB: tl.constexpr, + BLOCK_SIZE_KB: tl.constexpr, # packed bytes per K step (2 values each) GROUP_SIZE_M: tl.constexpr, MUL_ROUTED_WEIGHT: tl.constexpr, top_k: tl.constexpr, compute_type: tl.constexpr, + USE_TMA: tl.constexpr = False, # needs K % 256 == 0 + TMA_SCALES: tl.constexpr = False, # with USE_TMA: scales by TMA too, else pointer loads + FAST_DECODE: tl.constexpr = 0, # 0: fp16 pair decode, 1: packed_ptr / w_desc are int32 views of 8 codes, 2: bytes + SCALE_WORDS: tl.constexpr = False, # with FAST_DECODE == 1 and pointer loads: scale_ptr is an int32 view of 4 scales ): + # out^T = W @ X^T: wgmma reads its A operand (the decoded weights) from registers, while B must be in shared memory. + # The token block is the MMA N dimension, so a 16-token block still fills the MMA. pid = tl.program_id(axis=0) num_pid_m = tl.cdiv(EM, BLOCK_SIZE_M) num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) @@ -277,51 +409,110 @@ 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) + offs_sb = tl.arange(0, BLOCK_SIZE_KB // 8) + KW: tl.constexpr = BLOCK_SIZE_KB // 4 + offs_kw = tl.arange(0, KW) + xt_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 - scale_base = scale_ptr + slot * stride_se + offs_bn[None, :] * stride_sn - - accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) + WORDS: tl.constexpr = FAST_DECODE == 1 + if WORDS: + p_ptrs = packed_ptr + slot * stride_pe + offs_bn[:, None] * stride_pn + offs_kw[None, :] * stride_pkb + else: + p_ptrs = packed_ptr + slot * stride_pe + offs_bn[:, None] * stride_pn + offs_kb[None, :] * stride_pkb + tl.static_assert(not SCALE_WORDS or (WORDS and not USE_TMA and BLOCK_SIZE_KB >= 32)) + SW: tl.constexpr = BLOCK_SIZE_KB // 32 + if SCALE_WORDS: + offs_sw = tl.arange(0, SW) + s_ptrs = scale_ptr + slot * stride_se + offs_bn[:, None] * stride_sn + offs_sw[None, :] * stride_sblk + else: + s_ptrs = scale_ptr + slot * stride_se + offs_bn[:, None] * stride_sn + offs_sb[None, :] * stride_sblk + + accumulator = tl.zeros((BLOCK_SIZE_N, BLOCK_SIZE_M), dtype=tl.float32) K_BYTES = K // 2 + # rows past N belong to the next slot; their outputs are masked on store + w_row = (slot * N + pid_n * BLOCK_SIZE_N).to(tl.int32) + SB: tl.constexpr = BLOCK_SIZE_KB // 8 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) + k_mask = kb * 2 * BLOCK_SIZE_KB + offs_k < K + if USE_TMA and WORDS: + v = w_desc.load([w_row, kb * KW]) + elif USE_TMA: + v = w_desc.load([w_row, kb * BLOCK_SIZE_KB]).to(tl.int32) + if USE_TMA and TMA_SCALES: + # a TMA box row needs >= 16 bytes, so load the 16 scales that hold this step and pick its SB of them + s16 = s_desc.load([w_row, (kb * SB) // 16 * 16]).to(tl.uint8, bitcast=True) + s3 = tl.reshape(s16, (BLOCK_SIZE_N, 16 // SB, SB)) + pick = tl.arange(0, 16 // SB) == (kb * SB) % 16 // SB + s = tl.sum(tl.where(pick[None, :, None], s3, 0), axis=1).to(tl.uint8) + s = s.to(tl.float8e4nv, bitcast=True).to(tl.float32) + elif USE_TMA: + sblk_mask = kb * SB + offs_sb < K_BYTES // 8 + s = tl.load(s_ptrs, mask=sblk_mask[None, :], other=0.0).to(tl.float32) 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 + sblk_mask = kb * SB + offs_sb < K_BYTES // 8 + if WORDS: + v = tl.load(p_ptrs, mask=(kb * KW + offs_kw < K_BYTES // 4)[None, :], other=0) + else: + byte_mask = kb * BLOCK_SIZE_KB + offs_kb < K_BYTES + v = tl.load(p_ptrs, mask=byte_mask[None, :], other=0).to(tl.int32) + if SCALE_WORDS: + # 4-byte loads can go through the cp.async pipeline, byte loads cannot; spread each word over the 8 weight words it scales + sw = tl.load(s_ptrs, mask=(kb * SW + offs_sw < K_BYTES // 32)[None, :], other=0) + sw = tl.reshape(tl.broadcast_to(sw[:, :, None], (BLOCK_SIZE_N, SW, KW // SW)), (BLOCK_SIZE_N, KW)) + s = ((sw >> ((offs_kw // 2 % 4) * 8)[None, :]) & 0xFF).to(tl.uint8) + if e4m3_native_cx(): + s = s.to(tl.float8e4nv, bitcast=True).to(tl.float32) + else: + s = e4m3_u8_to_f32(s) + elif e4m3_native_cx(): + s = tl.load(s_ptrs, mask=sblk_mask[None, :], other=0.0).to(tl.float32) + else: + s = e4m3_u8_to_f32(tl.load(s_ptrs, mask=sblk_mask[None, :], other=0)) + if FAST_DECODE: + # 2^118 here and the decode's 2^-126 leave 2^-8, undone in the epilogue; every step is exact in bf16 + s = (s * 3.3230699894622794e35).to(tl.bfloat16).to(tl.uint16, bitcast=True).to(tl.uint32) + s = s | (s << 16) + if WORDS: + if not SCALE_WORDS: + # one scale per 16 values, two words + s = tl.reshape(tl.broadcast_to(s[:, :, None], (BLOCK_SIZE_N, SB, 2)), (BLOCK_SIZE_N, KW)) + w = _e2m1x8_bf16(v.to(tl.uint32, bitcast=True), s, BLOCK_SIZE_N, KW) + else: + # one scale per 16 values, eight bytes + s = tl.reshape(tl.broadcast_to(s[:, :, None], (BLOCK_SIZE_N, SB, 8)), (BLOCK_SIZE_N, BLOCK_SIZE_KB)) + w = _e2m1x2_bf16(v, s, BLOCK_SIZE_N, BLOCK_SIZE_KB) + else: + # one scale per 16 values; the 2^14 undoes the decode's 2^-14 + s = tl.reshape(tl.broadcast_to((s * 16384.0)[:, :, None], (BLOCK_SIZE_N, BLOCK_SIZE_KB // 8, 8)), (BLOCK_SIZE_N, BLOCK_SIZE_KB)) + s = s.to(tl.bfloat16) + # exact in bf16: e2m1 has <= 2 significant bits and e4m3 <= 4 + w_lo = _e2m1_f16_x2p14(v & 0xF).to(tl.bfloat16) * s + w_hi = _e2m1_f16_x2p14(v >> 4).to(tl.bfloat16) * s + # interleave back to k = 2 * byte + {0, 1} so one dot covers the K step + w = tl.reshape(tl.join(w_lo, w_hi), (BLOCK_SIZE_N, 2 * BLOCK_SIZE_KB)) + xt = tl.load(xt_ptrs, mask=k_mask[:, None] & token_mask[None, :], other=0.0) + accumulator = tl.dot(w.to(xt.dtype), xt, accumulator) + xt_ptrs += 2 * BLOCK_SIZE_KB * stride_ak + if not USE_TMA: + p_ptrs += (KW if WORDS else BLOCK_SIZE_KB) * stride_pkb + if not (USE_TMA and TMA_SCALES): + s_ptrs += (SW if SCALE_WORDS else BLOCK_SIZE_KB // 8) * stride_sblk g = tl.load(global_ptr + slot * stride_ge + offs_bn * stride_gn).to(tl.float32) - accumulator = accumulator * g[None, :] + if FAST_DECODE: + g = g * 256.0 + accumulator = accumulator * g[:, None] if MUL_ROUTED_WEIGHT: moe_weight = tl.load(topk_weights_ptr + offs_token * stride_tw, mask=token_mask, other=0) - accumulator = accumulator * moe_weight[:, None] + accumulator = accumulator * moe_weight[None, :] accumulator = accumulator.to(compute_type) offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) - c_ptrs = c_ptr + stride_cm * offs_token[:, None] + stride_cn * offs_cn[None, :] - c_mask = token_mask[:, None] & (offs_cn[None, :] < N) + c_ptrs = c_ptr + stride_cm * offs_token[None, :] + stride_cn * offs_cn[:, None] + c_mask = token_mask[None, :] & (offs_cn[:, None] < N) tl.store(c_ptrs, accumulator, mask=c_mask) diff --git a/python/freetoken/kernel/triton/nvfp4_linear.py b/python/freetoken/kernel/triton/nvfp4_linear.py index 53abdedc71..38f148f0ff 100644 --- a/python/freetoken/kernel/triton/nvfp4_linear.py +++ b/python/freetoken/kernel/triton/nvfp4_linear.py @@ -158,7 +158,7 @@ def _nvfp4_gemv_kernel( a_ptr, # [K] activation (compute dtype), contiguous packed_ptr, # [N, K // 8] int32 (8 fp4 codes per word, nibble j -> k = 8*w + j) scale_ptr, # [N, K // 16] fp8-e4m3 per-16 block scale - gscale_ptr, # [N] fp16 per-output-row global scale (weight_scale_2) + gscale_ptr, # [N] fp32 per-output-row global scale (weight_scale_2) out_ptr, # [N] out dtype N, K, stride_pn, stride_pkw, @@ -226,7 +226,7 @@ def _nvfp4_gemv_splitk_kernel( a_ptr, # [K] activation (compute dtype), contiguous packed_ptr, # [N, K // 8] int32 scale_ptr, # [N, K // 16] fp8-e4m3 - gscale_ptr, # [N] fp16 per-output-row global scale + gscale_ptr, # [N] fp32 per-output-row global scale part_ptr, # [SPLIT_K, N] fp32 partial sums (pre-global-scale) counter_ptr, # [cdiv(N, BLOCK_N)] int32 arrival counters, zeroed (self-resetting) out_ptr, # [N] out dtype @@ -311,7 +311,7 @@ def _gemv(a: torch.Tensor, packed_i32: torch.Tensor, scale: torch.Tensor, gscale: torch.Tensor, out_dtype: torch.dtype, transposed: bool) -> torch.Tensor: """M==1 W4A16 GEMV. ``a`` [K] compute-dtype; ``packed_i32`` logical [N, K//8] int32 (row-major or a K-major transposed view); ``scale`` logical [N, K//16] fp8; ``gscale`` - [N] fp16. Picks split-K so small ``N`` (shared expert) still fills the SMs; large ``N`` + [N] fp32. Picks split-K so small ``N`` (shared expert) still fills the SMs; large ``N`` (lm_head) stays single-pass (split_k=1).""" out_tl = _TL_DTYPE[out_dtype if out_dtype in _TL_DTYPE else torch.bfloat16] N = packed_i32.shape[0] @@ -389,7 +389,7 @@ def _nvfp4_gemm_kernel( a_ptr, # [M, K] activations packed_ptr, # logical [N, K // 8] int32 (row- or K-major via strides) scale_ptr, # logical [N, K // 16] fp8-e4m3 - gscale_ptr, # [N] fp16 + gscale_ptr, # [N] fp32 c_ptr, # [M, N] output (written when SPLIT_K == 1) part_ptr, # [SPLIT_K, M, N] fp32 partials (written when SPLIT_K > 1) M, N, K, K_WORDS, tiles_per, @@ -612,7 +612,7 @@ def _gemm_inkernel(a: torch.Tensor, packed_i32: torch.Tensor, scale: torch.Tenso def _nvfp4_dequant_rows_kernel( packed_ptr, # [N, K // 8] int32 scale_ptr, # [N, K // 16] fp8-e4m3 - gscale_ptr, # [N] fp16 + gscale_ptr, # [N] fp32 out_ptr, # [N, K] compute dtype K_WORDS, stride_pn, stride_pkw, @@ -646,7 +646,7 @@ def _nvfp4_dequant_rows_kernel( def _nvfp4_dequant_cols_kernel( packed_ptr, # K-major: element (n, kw) at kw * stride_pkw + n (stride_pn == 1) scale_ptr, - gscale_ptr, # [N] fp16 + gscale_ptr, # [N] fp32 out_ptr, # [n_rows, K] compute dtype, row-major N, K_WORDS, stride_pn, stride_pkw, @@ -781,7 +781,7 @@ def nvfp4_dense_linear( weight_global: torch.Tensor, bias: torch.Tensor | None = None, ) -> torch.Tensor: """``y = x @ dequant(weight)^T`` on the checkpoint-native row-major layout. - ``weight`` [N, K//2] uint8, ``weight_scale`` [N, K//16] fp8-e4m3, ``weight_global`` [N] fp16.""" + ``weight`` [N, K//2] uint8, ``weight_scale`` [N, K//16] fp8-e4m3, ``weight_global`` [N] fp32.""" return _linear_impl( x, weight.view(torch.int32), weight_scale, weight_global, bias, x.dtype, transposed=False ) 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/__init__.py b/python/freetoken/kernel/triton/qsa/__init__.py index 91753b7922..b04d27f050 100644 --- a/python/freetoken/kernel/triton/qsa/__init__.py +++ b/python/freetoken/kernel/triton/qsa/__init__.py @@ -3,7 +3,7 @@ from .attend import qsa_sparse_paged_attention from .compress import qsa_compress_groups, qsa_index_norm_rope, qsa_store_rows from .expand import expand_qsa_block_indices -from .score import qsa_mqa_paged +from .score import qsa_mqa_paged, qsa_mqa_paged_rows, qsa_score_tile_map from .topk import qsa_block_topk, qsa_block_topk_scratch_width __all__ = [ @@ -13,6 +13,8 @@ "qsa_compress_groups", "qsa_index_norm_rope", "qsa_mqa_paged", + "qsa_mqa_paged_rows", + "qsa_score_tile_map", "qsa_sparse_paged_attention", "qsa_store_rows", ] diff --git a/python/freetoken/kernel/triton/qsa/attend.py b/python/freetoken/kernel/triton/qsa/attend.py index 541e27c680..74ebf95c48 100644 --- a/python/freetoken/kernel/triton/qsa/attend.py +++ b/python/freetoken/kernel/triton/qsa/attend.py @@ -9,12 +9,19 @@ import triton import triton.language as tl +from freetoken.kernel.triton.e4m3_compat import kv_load_e4m3_tile_f32 +from freetoken.kernel.triton.kv_nvfp4 import load_nvfp4 + @triton.jit def _qsa_sparse_paged_gqa_splitk_kernel( q_ptr, k_cache_ptr, v_cache_ptr, + k_scale_ptr, + v_scale_ptr, + k_block_scale_ptr, + v_block_scale_ptr, indices_ptr, block_table_ptr, token_to_req_ptr, @@ -29,6 +36,8 @@ def _qsa_sparse_paged_gqa_splitk_kernel( stride_v_block, stride_v_token, stride_v_head, + stride_kss, + stride_vss, stride_indices_row, stride_table_req, stride_output_row, @@ -46,6 +55,11 @@ def _qsa_sparse_paged_gqa_splitk_kernel( NUM_TILES: tl.constexpr, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, + # e4m3 KV pool (kvcache/mha_pool.py): read codes + per-token row scales instead of + # 16-bit values. The bf16 branch below stays exactly as it was, instruction for + # instruction, for the unquantized default. + HAS_KV_SCALE: tl.constexpr, + KV_NVFP4: tl.constexpr, ) -> None: # row * stride can overflow int32 for large row counts. row = tl.program_id(0).to(tl.int64) @@ -101,24 +115,94 @@ def _qsa_sparse_paged_gqa_splitk_kernel( valid &= (physical_page >= 0) & (physical_page < num_cache_blocks) # physical_page * block stride can overflow int32 for large caches. safe_page = tl.maximum(physical_page, 0).to(tl.int64) - keys = tl.load( - k_cache_ptr - + safe_page[None, :] * stride_k_block - + page_offset[None, :] * stride_k_token - + kv_head * stride_k_head - + dim_offsets[:, None], - mask=valid[None, :], - other=0.0, - ) - values = tl.load( - v_cache_ptr - + safe_page[:, None] * stride_v_block - + page_offset[:, None] * stride_v_token - + kv_head * stride_v_head - + dim_offsets[None, :], - mask=valid[:, None], - other=0.0, - ) + if KV_NVFP4: + scale_slot = safe_page * PAGE_SIZE + page_offset + keys = load_nvfp4( + k_cache_ptr, + k_block_scale_ptr, + k_scale_ptr, + scale_slot[None, :], + kv_head, + dim_offsets[:, None], + valid[None, :], + stride_k_token, + stride_k_head, + stride_kss, + HEAD_DIM, + ).to(query.dtype) + values = load_nvfp4( + v_cache_ptr, + v_block_scale_ptr, + v_scale_ptr, + scale_slot[:, None], + kv_head, + dim_offsets[None, :], + valid[:, None], + stride_v_token, + stride_v_head, + stride_vss, + HEAD_DIM, + ).to(query.dtype) + elif HAS_KV_SCALE: + # The scale row is the slot the code lives in: QSA pins page_size to this + # kernel's PAGE_SIZE (attention/__init__.py registers page_sizes=(64,)), so + # slot = page * PAGE_SIZE + offset addresses k_scale/v_scale exactly. + scale_slot = safe_page * PAGE_SIZE + page_offset # safe_page is int64 + # Invalid columns mask the codes to 0.0 and the scale to 1.0, and slots that + # were never written read back 0.0 * 0.0 -- both buffers are zero-filled. + # Either way the operand stays finite, so the -inf row mask below is what + # decides such a column's fate rather than a NaN poisoning the row. + s_k = tl.load( + k_scale_ptr + scale_slot[None, :] * stride_kss + kv_head, + mask=valid[None, :], + other=1.0, + ) + s_v = tl.load( + v_scale_ptr + scale_slot[:, None] * stride_vss + kv_head, + mask=valid[:, None], + other=1.0, + ) + keys = ( + kv_load_e4m3_tile_f32( + k_cache_ptr + + safe_page[None, :] * stride_k_block + + page_offset[None, :] * stride_k_token + + kv_head * stride_k_head + + dim_offsets[:, None], + valid[None, :], + ) + * s_k + ).to(query.dtype) + values = ( + kv_load_e4m3_tile_f32( + v_cache_ptr + + safe_page[:, None] * stride_v_block + + page_offset[:, None] * stride_v_token + + kv_head * stride_v_head + + dim_offsets[None, :], + valid[:, None], + ) + * s_v + ).to(query.dtype) + else: + keys = tl.load( + k_cache_ptr + + safe_page[None, :] * stride_k_block + + page_offset[None, :] * stride_k_token + + kv_head * stride_k_head + + dim_offsets[:, None], + mask=valid[None, :], + other=0.0, + ) + values = tl.load( + v_cache_ptr + + safe_page[:, None] * stride_v_block + + page_offset[:, None] * stride_v_token + + kv_head * stride_v_head + + dim_offsets[None, :], + mask=valid[:, None], + other=0.0, + ) scores = tl.dot(query, keys) # Scaling scores avoids re-quantizing a scaled query to BF16. scores *= softmax_scale_log2 @@ -232,8 +316,13 @@ def qsa_sparse_paged_attention( block_table: torch.Tensor, token_to_req: torch.Tensor, out: torch.Tensor | None = None, + k_scale: torch.Tensor | None = None, + v_scale: torch.Tensor | None = None, + kv_quant: str | None = None, + k_block_scale: torch.Tensor | None = None, + v_block_scale: torch.Tensor | None = None, ) -> torch.Tensor: - """Run sparse GQA directly over paged BF16 K/V caches.""" + """Run sparse GQA over bf16, FP8, or packed NVFP4 K/V caches.""" if q.ndim != 3 or k_cache.ndim != 4 or v_cache.shape != k_cache.shape: raise ValueError("QSA sparse attention received invalid Q/K/V shapes") @@ -243,11 +332,52 @@ def qsa_sparse_paged_attention( raise ValueError("QSA sparse attention metadata has invalid shapes") if logical_indices.shape[1] <= 0: raise ValueError("QSA sparse attention requires a positive selection width") - if q.shape[2] != k_cache.shape[3] or q.shape[1] % k_cache.shape[2]: + kv_quant = kv_quant if kv_quant is not None else ("fp8" if k_scale is not None else "none") + if kv_quant not in ("none", "fp8", "nvfp4"): + raise ValueError(f"unknown QSA kv_quant {kv_quant!r}") + stored_dim = q.shape[2] // 2 if kv_quant == "nvfp4" else q.shape[2] + if ( + (kv_quant == "nvfp4" and q.shape[2] % 16) + or k_cache.shape[3] != stored_dim + or q.shape[1] % k_cache.shape[2] + ): raise ValueError("QSA sparse attention requires valid grouped-query heads") head_dim = q.shape[2] assert head_dim >= 16 and (head_dim & (head_dim - 1)) == 0 - assert q.dtype == k_cache.dtype == v_cache.dtype + if (k_scale is None) != (v_scale is None) or (k_block_scale is None) != (v_block_scale is None): + raise ValueError("QSA sparse attention requires both KV scale tensors") + nvfp4 = kv_quant == "nvfp4" + if nvfp4 and (k_scale is None or k_block_scale is None): + raise ValueError("QSA NVFP4 requires row and block scale tensors") + if not nvfp4 and (k_block_scale is not None or v_block_scale is not None): + raise ValueError("QSA block scales require kv_quant='nvfp4'") + if k_scale is not None: + # The pool hands out 1-byte e4m3 codes plus one fp32 row scale per + # (slot, kv_head); the kernel rebuilds that slot as + # page * PAGE_SIZE + page_offset, which is exact only because QSA pins + # page_size to PAGE_SIZE (page_sizes=(64,) in attention/__init__.py). + if k_cache.element_size() != 1 or v_cache.element_size() != 1: + raise ValueError("QSA KV scales require 1-byte e4m3 code caches") + if k_scale.dtype is not torch.float32 or v_scale.dtype is not torch.float32: + raise ValueError("QSA KV scale tensors must be float32") + want = (k_cache.shape[0] * k_cache.shape[1], k_cache.shape[2]) + if k_scale.shape != want or v_scale.shape != want: + raise ValueError(f"QSA KV scale tensors must have shape {want}") + assert k_scale.stride(1) == v_scale.stride(1) == 1 + if nvfp4: + block_want = (*want, head_dim // 16) + if ( + k_cache.dtype is not torch.uint8 + or v_cache.dtype is not torch.uint8 + or k_block_scale.dtype is not torch.uint8 + or v_block_scale.dtype is not torch.uint8 + or k_block_scale.shape != block_want + or v_block_scale.shape != block_want + ): + raise ValueError(f"QSA NVFP4 block scales must have shape {block_want}") + assert k_block_scale.stride(2) == v_block_scale.stride(2) == 1 + else: + assert q.dtype == k_cache.dtype == v_cache.dtype assert logical_indices.dtype == block_table.dtype == torch.int32 assert token_to_req.dtype == torch.int32 assert q.stride(2) == k_cache.stride(3) == v_cache.stride(3) == 1 @@ -277,6 +407,25 @@ def qsa_sparse_paged_attention( else: block_n, target_splits, partial_warps = 64, 1, 2 + if nvfp4 and block_n == 64 and head_dim >= 256: + # The wide bf16 profile has load_nvfp4 decode a BLOCK_N x HEAD_DIM fp32 tile for K and V + # across 64 threads (256 values per thread at HEAD_DIM 256): Triton 3.8 compiles that to + # 40 registers with ~1.4 KB of spills, ~10x slower and a ~550 MiB local-memory reservation + # outside the torch allocator. Narrow tiles at 4 warps do not spill and are faster on 3.6 too. + block_n, partial_warps = 16, 4 + + # These tiles were tuned for the GB300's 228KB LDS; AMD parts have less and it varies by arch + # (e.g. 64KB/CU on RDNA gfx11/gfx12), so k+v staging can overflow. Query the device's actual + # shared-memory limit rather than assuming a fixed size, and shrink the KV tile only when it + # doesn't fit -- so a larger-LDS arch keeps the wide tile. ROCm-gated so the NVIDIA tiles stay + # byte-identical (a small NV part could report a 48KB default and would otherwise clamp). + if torch.version.hip is not None: + _props = torch.cuda.get_device_properties(q.device) + _lds_budget = getattr(_props, "shared_memory_per_block", 64 * 1024) or 64 * 1024 + _kv_bytes = k_cache.element_size() + while block_n > 16 and 2 * block_n * q.shape[2] * _kv_bytes + 2048 > _lds_budget: + block_n //= 2 + num_tiles = triton.cdiv(logical_indices.shape[1], block_n) # Avoid empty splits when the selection width is smaller than the profile. max_useful_splits = 1 << (num_tiles.bit_length() - 1) @@ -303,6 +452,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 +472,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 +491,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..24526d0e86 100644 --- a/python/freetoken/kernel/triton/qsa/score.py +++ b/python/freetoken/kernel/triton/qsa/score.py @@ -115,6 +115,227 @@ def _qsa_mqa_paged_kernel( ) +@triton.jit +def _qsa_mqa_paged_rows_kernel( + q_ptr, + k_cache_ptr, + page_table_ptr, + tile_map_ptr, + token_to_req_ptr, + query_positions_ptr, + sequence_lengths_ptr, + visible_blocks_ptr, + logits_ptr, + stride_q_row, + stride_q_head, + stride_q_dim, + stride_cache_block, + stride_cache_token, + stride_cache_dim, + stride_table_req, + stride_table_page, + stride_logits_row, + num_tiles, + num_columns, + num_pages, + num_requests, + score_divisor, + PAGE_SIZE: tl.constexpr, + PAGE_TABLE_WIDTH: tl.constexpr, + HEAD_DIM: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_D: tl.constexpr, + ROWS: tl.constexpr, + TILES_PER_PROG: tl.constexpr, + STAGES: tl.constexpr, + COMPRESS_RATIO: tl.constexpr, +) -> None: + """Prefill scoring, ROWS query rows x 4 heads per program. + + The rows of a tile all belong to one request (``tile_map`` holds each tile's first row + and row count, cut at request boundaries), so one key tile serves every row. Each head + is its own ``[BLOCK_N, ROWS]`` dot and the head sum keeps the one-row kernel's order + ``(h0 + h1) + (h2 + h3)``: every score is bitwise the one-row kernel's. + """ + tile = tl.program_id(0) + row0 = tl.load(tile_map_ptr + tile) + nrows = tl.load(tile_map_ptr + num_tiles + tile) + dims = tl.arange(0, BLOCK_D) + lanes = tl.arange(0, ROWS) + row_live = lanes < nrows + rows = row0 + lanes + request = tl.load(token_to_req_ptr + row0) + safe_request = tl.minimum(tl.maximum(request, 0), num_requests - 1) + query_position = tl.load(query_positions_ptr + rows, mask=row_live, other=0) + sequence_length = tl.load( + sequence_lengths_ptr + safe_request, + mask=(request >= 0) & (request < num_requests), + other=0, + ) + visible = tl.minimum( + (query_position + 1) // COMPRESS_RATIO, + sequence_length // COMPRESS_RATIO, + ) + visible = tl.where(row_live, visible, 0) + if tl.program_id(1) == 0: + tl.store(visible_blocks_ptr + rows, visible, mask=row_live) + visible_max = tl.max(visible, axis=0) + tile_start = tl.program_id(1) * TILES_PER_PROG + # Top-k is bounded by visible_blocks, so columns beyond it need no value. + if tile_start * BLOCK_N >= visible_max: + return + tile_end = tl.minimum(tile_start + TILES_PER_PROG, tl.cdiv(visible_max, BLOCK_N)) + tile_end = tl.minimum(tile_end, tl.cdiv(num_columns, BLOCK_N)) + + q_rows = q_ptr + rows[None, :].to(tl.int64) * stride_q_row + dims[:, None] * stride_q_dim + q_mask = row_live[None, :] & (dims[:, None] < HEAD_DIM) + query0 = tl.load(q_rows + 0 * stride_q_head, mask=q_mask, other=0.0) + query1 = tl.load(q_rows + 1 * stride_q_head, mask=q_mask, other=0.0) + query2 = tl.load(q_rows + 2 * stride_q_head, mask=q_mask, other=0.0) + query3 = tl.load(q_rows + 3 * stride_q_head, mask=q_mask, other=0.0) + column_offsets = tl.arange(0, BLOCK_N) + out_rows = logits_ptr + rows[None, :].to(tl.int64) * stride_logits_row + for col_tile in tl.range(tile_start, tile_end, num_stages=STAGES): + columns = col_tile * BLOCK_N + column_offsets + live = columns < visible_max + logical_page = tl.minimum(columns // PAGE_SIZE, PAGE_TABLE_WIDTH - 1) + page_offset = columns % PAGE_SIZE + physical_page = tl.load( + page_table_ptr + + safe_request * stride_table_req + + logical_page * stride_table_page, + mask=live, + other=-1, + ) + page_valid = live & (physical_page >= 0) & (physical_page < num_pages) + # physical_page * block stride can overflow int32 for large caches. + safe_physical_page = tl.maximum(physical_page, 0).to(tl.int64) + keys = tl.load( + k_cache_ptr + + safe_physical_page[:, None] * stride_cache_block + + page_offset[:, None] * stride_cache_token + + dims[None, :] * stride_cache_dim, + mask=page_valid[:, None] & (dims[None, :] < HEAD_DIM), + other=0.0, + eviction_policy="evict_first", + ) + s0 = tl.maximum(tl.dot(keys, query0, out_dtype=tl.float32), 0.0) + s1 = tl.maximum(tl.dot(keys, query1, out_dtype=tl.float32), 0.0) + s2 = tl.maximum(tl.dot(keys, query2, out_dtype=tl.float32), 0.0) + s3 = tl.maximum(tl.dot(keys, query3, out_dtype=tl.float32), 0.0) + score = ((s0 + s1) + (s2 + s3)) / score_divisor + tl.store( + out_rows + columns[:, None], + tl.where(page_valid[:, None], score, -float("inf")), + mask=(columns[:, None] < visible[None, :]) & (columns[:, None] < num_columns), + ) + + +def qsa_score_tile_map( + qo_indptr: torch.Tensor, start: int, end: int, rows_per_tile: int, device: torch.device +) -> torch.Tensor: + """``[2, tiles]`` int32 (first row, row count) of rows ``[start, end)``, relative to + ``start``, in tiles of at most ``rows_per_tile`` rows that never cross a request boundary + of the host-side ``qo_indptr`` (cu_seqlens).""" + + bounds = [int(x) for x in qo_indptr.tolist()] + firsts: list[int] = [] + counts: list[int] = [] + for lo, hi in zip(bounds, bounds[1:]): + lo, hi = max(lo, start), min(hi, end) + for row in range(lo, hi, rows_per_tile): + firsts.append(row - start) + counts.append(min(rows_per_tile, hi - row)) + host = torch.tensor([firsts, counts], dtype=torch.int32) + if device.type == "cuda": + host = host.pin_memory() + return host.to(device, non_blocking=True) + + +SCORE_ROWS = 16 + + +def qsa_mqa_paged_rows( + q: torch.Tensor, + k_cache: torch.Tensor, + page_table: torch.Tensor, + tile_map: torch.Tensor, + token_to_req: torch.Tensor, + query_positions: torch.Tensor, + sequence_lengths: torch.Tensor, + compress_ratio: int, + logits: torch.Tensor, + visible_blocks: torch.Tensor, + score_scale: float | None = None, + rows_per_tile: int = SCORE_ROWS, +) -> tuple[torch.Tensor, torch.Tensor]: + """Prefill form of :func:`qsa_mqa_paged`: the same scores bit for bit, with + ``rows_per_tile`` query rows of one request per program instead of one row padded to 16 + heads. ``tile_map`` comes from :func:`qsa_score_tile_map` over the same rows and must use + the same ``rows_per_tile``. Needs exactly 4 index heads.""" + + if q.ndim != 3 or q.shape[1] != 4 or q.shape[2] <= 0: + raise ValueError("QSA row-tiled scoring needs a [rows, 4, head_dim] query") + if k_cache.ndim != 4 or k_cache.shape[2] != 1 or k_cache.shape[3] != q.shape[2]: + raise ValueError("QSA cache must be [pages, page_size, 1, head_dim] matching the query") + if token_to_req.shape != (q.shape[0],) or query_positions.shape != (q.shape[0],): + 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: + raise ValueError( + f"QSA scoring is 16-bit only, got query={q.dtype} keys={k_cache.dtype}" + ) + if tile_map.ndim != 2 or tile_map.shape[0] != 2 or tile_map.dtype != torch.int32: + raise ValueError("QSA tile map must be [2, tiles] int32") + if rows_per_tile < 16 or rows_per_tile & (rows_per_tile - 1): + raise ValueError("rows_per_tile must be a power of two >= 16 (the dot's N)") + score_divisor = math.sqrt(q.shape[2]) if score_scale is None else score_scale + columns = logits.shape[1] + tiles = tile_map.shape[1] + if not q.shape[0] or not columns or not tiles: + return logits, visible_blocks + BLOCK_N = 64 + BLOCK_D = max(16, triton.next_power_of_2(q.shape[2])) + tiles_per_program = 8 + _qsa_mqa_paged_rows_kernel[(tiles, triton.cdiv(columns, BLOCK_N * tiles_per_program))]( + q, + k_cache, + page_table, + tile_map, + token_to_req, + query_positions, + sequence_lengths, + visible_blocks, + logits, + q.stride(0), + q.stride(1), + q.stride(2), + k_cache.stride(0), + k_cache.stride(1), + k_cache.stride(3), + page_table.stride(0), + page_table.stride(1), + logits.stride(0), + tiles, + columns, + k_cache.shape[0], + page_table.shape[0], + float(score_divisor), + PAGE_SIZE=k_cache.shape[1], + PAGE_TABLE_WIDTH=page_table.shape[1], + HEAD_DIM=q.shape[2], + BLOCK_N=BLOCK_N, + BLOCK_D=BLOCK_D, + ROWS=rows_per_tile, + TILES_PER_PROG=tiles_per_program, + STAGES=2, + COMPRESS_RATIO=compress_ratio, + num_warps=4, + ) + return logits, visible_blocks + + def qsa_mqa_paged( q: torch.Tensor, k_cache: torch.Tensor, @@ -139,6 +360,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/qsa/topk.py b/python/freetoken/kernel/triton/qsa/topk.py index ae57566cab..0f37a0f9d9 100644 --- a/python/freetoken/kernel/triton/qsa/topk.py +++ b/python/freetoken/kernel/triton/qsa/topk.py @@ -17,10 +17,18 @@ top-k is a subset of that union -- a global winner beats at most k-1 columns, so it beats all but at most k-1 of its own slice -- so the answer stays exact. ``_split_plan`` derives the geometry from the buffer width alone, which is fixed when a graph is captured. + +A row wide enough to leave more candidates than one program holds resident (262,144 columns, +a 1M-token window, leaves 32 x k) gets reduction levels between the two phases: each program +of a level takes ``_MAX_RESIDENT`` consecutive candidates of the level before and writes their +own top-k, until the survivors fit the merge. The same subset argument holds per level, and +candidates stay in ascending column order throughout, so ties still go to the lowest column. """ from __future__ import annotations +from dataclasses import dataclass + import torch import triton import triton.language as tl @@ -125,6 +133,44 @@ def _compact( return above, equal +@triton.jit +def _live_candidates( + visible, + CHUNK: tl.constexpr, + N_SPLITS: tl.constexpr, + TOP_K: tl.constexpr, + LEVEL: tl.constexpr, + RESIDENT: tl.constexpr, +): + """Candidate slots level ``LEVEL`` filled for a row: its splits past the row's visible tail + are one skipped suffix, so every later stage keeps costing what the live part costs.""" + live = tl.minimum(tl.cdiv(visible, CHUNK), N_SPLITS) * TOP_K + for _ in tl.static_range(LEVEL): + live = tl.cdiv(live, RESIDENT) * TOP_K + return live + + +@triton.jit +def _emit_candidates( + key, + column, + k_eff, + key_out, + col_out, + TOP_K: tl.constexpr, + BINS: tl.constexpr, + RADIX: tl.constexpr, + PASSES: tl.constexpr, +): + """Write a resident tile's own top-k as ranked (key, column) candidates; returns how many.""" + prefix, ties = _resident_prefix(key, k_eff, BINS, RADIX, PASSES) + rank, take, above, equal = _tile_ranks(key, prefix, ties, 0, 0) + write = take & (rank < TOP_K) + tl.store(key_out + rank, key.to(tl.int32, bitcast=True), mask=write) + tl.store(col_out + rank, column, mask=write) + return above + tl.minimum(equal, ties) + + @triton.jit def _qsa_block_topk_kernel( logits_ptr, @@ -225,17 +271,95 @@ def _qsa_topk_split_kernel( if k_eff > 0: offsets = tl.arange(0, CHUNK) key = _load_keys(logits_ptr + row.to(tl.int64) * stride_logits_row + base, offsets, limit) - prefix, ties = _resident_prefix(key, k_eff, BINS, RADIX, PASSES) - rank, take, above, equal = _tile_ranks(key, prefix, ties, 0, 0) - write = take & (rank < TOP_K) - tl.store(key_ptr + slot + rank, key.to(tl.int32, bitcast=True), mask=write) - tl.store(col_ptr + slot + rank, (base + offsets).to(tl.int32), mask=write) - emitted = above + tl.minimum(equal, ties) + emitted = _emit_candidates( + key, (base + offsets).to(tl.int32), k_eff, key_ptr + slot, col_ptr + slot, + TOP_K, BINS, RADIX, PASSES, + ) # A key of 0 is the dead-column sentinel, so the merge needs no separate count per slot. pad = tl.arange(0, PAD_K) tl.store(key_ptr + slot + pad, 0, mask=(pad >= emitted) & (pad < TOP_K)) +@triton.jit +def _reduce_tile( + in_key, + in_col, + out_key, + out_col, + count, + k_eff, + TOP_K: tl.constexpr, + BLOCK: tl.constexpr, + BINS: tl.constexpr, + RADIX: tl.constexpr, + PASSES: tl.constexpr, +): + """Top-k of ``count`` candidates held as one resident tile; returns how many it wrote.""" + tl.static_assert(BLOCK <= 0xFFFF, "packed cumsum keeps 16 bits per half") + offsets = tl.arange(0, BLOCK) + inside = offsets < count + key = tl.load(in_key + offsets, mask=inside, other=0).to(tl.uint32, bitcast=True) + column = tl.load(in_col + offsets, mask=inside, other=-1) + return _emit_candidates(key, column, k_eff, out_key, out_col, TOP_K, BINS, RADIX, PASSES) + + +@triton.jit +def _qsa_topk_reduce_kernel( + visible_ptr, + in_key_ptr, + in_col_ptr, + out_key_ptr, + out_col_ptr, + stride_scratch_row, + num_columns, + TOP_K: tl.constexpr, + PAD_K: tl.constexpr, + CHUNK: tl.constexpr, + N_SPLITS: tl.constexpr, + LEVEL: tl.constexpr, + RESIDENT: tl.constexpr, + BLOCK_SMALL: tl.constexpr, + BLOCK_MID: tl.constexpr, + BINS: tl.constexpr, + RADIX: tl.constexpr, + PASSES: tl.constexpr, +) -> None: + """A reduction level: one program per (row, ``RESIDENT`` candidates of level ``LEVEL - 1``), + writing their own top-k as this level's candidates.""" + row = tl.program_id(0) + split = tl.program_id(1) + visible = tl.maximum(tl.minimum(tl.load(visible_ptr + row), num_columns), 0) + k_eff = tl.minimum(visible, TOP_K) + live = _live_candidates(visible, CHUNK, N_SPLITS, TOP_K, LEVEL - 1, RESIDENT) + base = split * RESIDENT + count = tl.minimum(tl.maximum(live - base, 0), RESIDENT) + source = row.to(tl.int64) * stride_scratch_row + base + slot = row.to(tl.int64) * stride_scratch_row + split * TOP_K + in_key, in_col = in_key_ptr + source, in_col_ptr + source + out_key, out_col = out_key_ptr + slot, out_col_ptr + slot + + emitted = 0 + # As in the merge: residency costs the whole tile, so a short range takes a narrower one. + if count > 0: + if count <= BLOCK_SMALL: + emitted = _reduce_tile( + in_key, in_col, out_key, out_col, count, k_eff, + TOP_K, BLOCK_SMALL, BINS, RADIX, PASSES, + ) + elif count <= BLOCK_MID: + emitted = _reduce_tile( + in_key, in_col, out_key, out_col, count, k_eff, + TOP_K, BLOCK_MID, BINS, RADIX, PASSES, + ) + else: + emitted = _reduce_tile( + in_key, in_col, out_key, out_col, count, k_eff, + TOP_K, RESIDENT, BINS, RADIX, PASSES, + ) + pad = tl.arange(0, PAD_K) + tl.store(out_key + pad, 0, mask=(pad >= emitted) & (pad < TOP_K)) + + @triton.jit def _merge_tile( key_row, @@ -274,6 +398,8 @@ def _qsa_topk_merge_kernel( PAD_K: tl.constexpr, CHUNK: tl.constexpr, N_SPLITS: tl.constexpr, + LEVEL: tl.constexpr, + RESIDENT: tl.constexpr, BLOCK_SMALL: tl.constexpr, BLOCK_MID: tl.constexpr, BLOCK_FULL: tl.constexpr, @@ -281,13 +407,11 @@ def _qsa_topk_merge_kernel( RADIX: tl.constexpr, PASSES: tl.constexpr, ) -> None: - """Phase 2: one program per row over the candidates phase 1 left behind.""" + """Phase 2: one program per row over the candidates the last level ``LEVEL`` left behind.""" row = tl.program_id(0) limit = tl.maximum(tl.minimum(tl.load(visible_ptr + row), num_columns), 0) k_eff = tl.minimum(limit, TOP_K) - # Candidates sit chunk-major, so the splits past the visible tail are one skipped suffix - # and the merge keeps costing what the live part of the row costs. - candidates = tl.minimum(tl.cdiv(limit, CHUNK), N_SPLITS) * TOP_K + candidates = _live_candidates(limit, CHUNK, N_SPLITS, TOP_K, LEVEL, RESIDENT) key_row = key_ptr + row.to(tl.int64) * stride_scratch_row col_row = col_ptr + row.to(tl.int64) * stride_scratch_row out_row = out_ptr + row.to(tl.int64) * stride_out_row @@ -325,27 +449,37 @@ def _qsa_topk_merge_kernel( tl.store(out_row + pad, -1, mask=(pad >= emitted) & (pad < TOP_K)) -def _split_plan(columns: int, top_k: int) -> tuple[int, int] | None: - """``(chunk, n_splits)`` for the split+merge path, or None to keep the one-program path.""" +@dataclass(frozen=True) +class _SplitPlan: + """``chunk`` columns per first-level program, and the program count of every level: the + first splits the row, each later one reduces the candidates of the level before it.""" + + chunk: int + levels: tuple[int, ...] + + +def _split_plan(columns: int, top_k: int) -> _SplitPlan | None: + """The split geometry of a ``columns``-wide buffer, or None to keep the one-program path.""" if top_k <= 0 or columns <= _MIN_CHUNK: return None max_splits = _MAX_RESIDENT // triton.next_power_of_2(top_k) if max_splits < 2: return None - chunk = max(_MIN_CHUNK, triton.next_power_of_2(-(-columns // max_splits))) - if chunk > _MAX_RESIDENT: - return None + chunk = min(_MAX_RESIDENT, max(_MIN_CHUNK, triton.next_power_of_2(-(-columns // max_splits)))) n_splits = -(-columns // chunk) # Merging n_splits*top_k candidates has to be cheaper than scanning the row once. if n_splits < 2 or n_splits * top_k >= columns: return None - return chunk, n_splits + levels = [n_splits] + while levels[-1] * top_k > _MAX_RESIDENT: + levels.append(-(-levels[-1] * top_k // _MAX_RESIDENT)) + return _SplitPlan(chunk, tuple(levels)) def qsa_block_topk_scratch_width(columns: int, top_k: int) -> int: """int32 columns of scratch ``qsa_block_topk`` wants per row; 0 when it needs none.""" plan = _split_plan(columns, top_k) - return 0 if plan is None else 2 * plan[1] * top_k + return 0 if plan is None else 2 * top_k * sum(plan.levels) def qsa_block_topk( @@ -396,24 +530,31 @@ def qsa_block_topk( ) return out - chunk, n_splits = plan - half = n_splits * top_k + width = 2 * top_k * sum(plan.levels) if scratch is None: - scratch = torch.empty((rows, 2 * half), dtype=torch.int32, device=logits.device) + scratch = torch.empty((rows, width), dtype=torch.int32, device=logits.device) elif ( scratch.ndim != 2 or scratch.shape[0] < rows - or scratch.shape[1] < 2 * half + or scratch.shape[1] < width or scratch.dtype != torch.int32 or scratch.stride(1) != 1 ): raise ValueError( f"QSA block top-k needs a row-contiguous int32 scratch of at least " - f"[{rows}, {2 * half}], got {tuple(scratch.shape)} {scratch.dtype}" + f"[{rows}, {width}], got {tuple(scratch.shape)} {scratch.dtype}" ) - keys = scratch[:rows, :half] - cols = scratch[:rows, half : 2 * half] - _qsa_topk_split_kernel[(rows, n_splits)]( + # Each level's candidates: its keys, then their columns, level after level along the row. + candidates: list[tuple[torch.Tensor, torch.Tensor]] = [] + offset = 0 + for programs in plan.levels: + half = programs * top_k + candidates.append( + (scratch[:rows, offset : offset + half], scratch[:rows, offset + half : offset + 2 * half]) + ) + offset += 2 * half + keys, cols = candidates[0] + _qsa_topk_split_kernel[(rows, plan.levels[0])]( logits, visible, keys, @@ -423,14 +564,39 @@ def qsa_block_topk( columns, TOP_K=top_k, PAD_K=pad_k, - CHUNK=chunk, + CHUNK=plan.chunk, BINS=_BINS, RADIX=_RADIX, PASSES=_PASSES, num_warps=8, num_stages=1, ) - merge_block = triton.next_power_of_2(half) + for level in range(1, len(plan.levels)): + (in_keys, in_cols), (out_keys, out_cols) = candidates[level - 1], candidates[level] + _qsa_topk_reduce_kernel[(rows, plan.levels[level])]( + visible, + in_keys, + in_cols, + out_keys, + out_cols, + keys.stride(0), + columns, + TOP_K=top_k, + PAD_K=pad_k, + CHUNK=plan.chunk, + N_SPLITS=plan.levels[0], + LEVEL=level, + RESIDENT=_MAX_RESIDENT, + BLOCK_SMALL=min(1024, _MAX_RESIDENT), + BLOCK_MID=min(4096, _MAX_RESIDENT), + BINS=_BINS, + RADIX=_RADIX, + PASSES=_PASSES, + num_warps=8, + num_stages=1, + ) + keys, cols = candidates[-1] + merge_block = triton.next_power_of_2(keys.shape[1]) _qsa_topk_merge_kernel[(rows,)]( visible, keys, @@ -441,8 +607,10 @@ def qsa_block_topk( columns, TOP_K=top_k, PAD_K=pad_k, - CHUNK=chunk, - N_SPLITS=n_splits, + CHUNK=plan.chunk, + N_SPLITS=plan.levels[0], + LEVEL=len(plan.levels) - 1, + RESIDENT=_MAX_RESIDENT, BLOCK_SMALL=min(1024, merge_block), BLOCK_MID=min(4096, merge_block), BLOCK_FULL=merge_block, diff --git a/python/freetoken/kernel/triton/sampling.py b/python/freetoken/kernel/triton/sampling.py index a970714406..38bc253386 100644 --- a/python/freetoken/kernel/triton/sampling.py +++ b/python/freetoken/kernel/triton/sampling.py @@ -92,7 +92,9 @@ 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) + # While every logit so far is -inf, shift by 0: exp(-inf - -inf) would be NaN. + shift = tl.where(new_m == -float("inf"), 0.0, new_m) + d = d * tl.exp(m - shift) + tl.sum(tl.exp(x - shift), 0) m = new_m tl.store(pm_ptr + pid, m) tl.store(pl_ptr + pid, d) diff --git a/python/freetoken/kernel/utils.py b/python/freetoken/kernel/utils.py index cdd765c7d6..9af61526ae 100644 --- a/python/freetoken/kernel/utils.py +++ b/python/freetoken/kernel/utils.py @@ -5,6 +5,7 @@ import os import pathlib import re +from functools import cache from typing import TYPE_CHECKING, Iterator, List, NamedTuple, Tuple, TypeAlias, Union if TYPE_CHECKING: @@ -20,15 +21,23 @@ DEFAULT_INCLUDE = [str(KERNEL_PATH / "include")] DEFAULT_CFLAGS = ["-std=c++20", "-O3"] DEFAULT_CUDA_CFLAGS = ["-std=c++20", "-O3", "--expt-relaxed-constexpr"] +DEFAULT_HIP_CFLAGS = ["-std=c++20", "-O3"] DEFAULT_LDFLAGS = [] -ARCH_LIST_ENV = "TVM_FFI_CUDA_ARCH_LIST" +CUDA_ARCH_LIST_ENV = "TVM_FFI_CUDA_ARCH_LIST" +ROCM_ARCH_LIST_ENV = "TVM_FFI_ROCM_ARCH_LIST" + + +def _is_rocm() -> bool: + import torch + + return getattr(torch.version, "hip", None) is not None def _cuda_arch_list() -> List[str]: """Archs a CUDA build targets: the AOT build's TVM_FFI_CUDA_ARCH_LIST, else the GPU this process is bound to.""" - arch_list = os.getenv(ARCH_LIST_ENV, "").split() + arch_list = os.getenv(CUDA_ARCH_LIST_ENV, "").split() if arch_list: return arch_list import torch @@ -39,17 +48,34 @@ def _cuda_arch_list() -> List[str]: return [f"{major}.{minor}"] +def _rocm_arch_list() -> List[str]: + """The single ROCm GPU this process is bound to, or an explicit cross-compile target.""" + from freetoken.utils.arch import get_rocm_gfx_arch + + arch = get_rocm_gfx_arch() + if arch is None: + raise RuntimeError( + "Could not determine the bound ROCm GPU architecture; select a visible GPU " + "or set FREETOKEN_ROCM_ARCH for cross-compilation" + ) + return [arch] + + @contextlib.contextmanager -def _pin_tvm_ffi_arch_ctx(arch_list: List[str]) -> Iterator[None]: - """Hand tvm-ffi the arch list through its env var for one build. Left unset, tvm-ffi asks nvidia-smi and takes the first GPU listed, which under --gpu or CUDA_VISIBLE_DEVICES on a mixed box is not the bound one.""" - if os.getenv(ARCH_LIST_ENV) or not arch_list: +def _pin_tvm_ffi_arch_ctx(arch_list: List[str], env_var: str) -> Iterator[None]: + """Pin tvm-ffi to the resolved targets for one build, then restore the caller's environment.""" + if not arch_list: yield return - os.environ[ARCH_LIST_ENV] = " ".join(arch_list) + old_arch_list = os.environ.get(env_var) + os.environ[env_var] = " ".join(arch_list) try: yield finally: - os.environ.pop(ARCH_LIST_ENV, None) + if old_arch_list is None: + os.environ.pop(env_var, None) + else: + os.environ[env_var] = old_arch_list def _cuda_cflags(extra: List[str], arch_list: List[str]) -> List[str]: @@ -63,6 +89,61 @@ def _rank(a: str) -> int: cc = max(arch_list, key=_rank).rstrip("a").replace(".", "") flags = flags + [f"-gencode=arch=compute_{cc},code=compute_{cc}"] return flags + + +def _hip_cflags(extra: List[str], arch_list: List[str]) -> List[str]: + """HIP flags for a kernel build on ROCm.""" + # TODO(ROCm): Triton autotune configs need RDNA-specific tuning (wave count, LDS size). + return DEFAULT_HIP_CFLAGS + extra + [f"--offload-arch={arch}" for arch in arch_list] + + +@cache +def _rocm_link_flags() -> List[str]: + """Make ROCm's runtime library discoverable to JIT link commands. + + Traditional ROCm installs provide ``libamdhip64.so`` under ``$ROCM_HOME/lib``. + ROCm 7.14 Python SDK images only provide the versioned soname, while TVM-FFI + still links with ``-lamdhip64``. Supply a cache-local unversioned symlink via + an explicit linker search path without modifying the Python environment. + """ + candidates: list[pathlib.Path] = [] + if os.getenv("ROCM_HOME"): + candidates.append(pathlib.Path(os.environ["ROCM_HOME"])) + try: + from torch.utils.cpp_extension import ROCM_HOME + + if ROCM_HOME: + candidates.append(pathlib.Path(ROCM_HOME)) + except ImportError: + pass + spec = importlib.util.find_spec("_rocm_sdk_core") + if spec and spec.submodule_search_locations: + candidates.append(pathlib.Path(next(iter(spec.submodule_search_locations)))) + candidates.append(pathlib.Path("/opt/rocm")) + + for rocm_home in dict.fromkeys(candidates): + library_dir = rocm_home / "lib" + unversioned = library_dir / "libamdhip64.so" + link_dir = library_dir + if not unversioned.exists(): + versioned = sorted(library_dir.glob("libamdhip64.so.*")) + if not versioned: + continue + link_dir = pathlib.Path.home() / ".cache" / "freetoken" / "rocm-lib" + link_dir.mkdir(parents=True, exist_ok=True) + compat_link = link_dir / "libamdhip64.so" + if not compat_link.exists() and not compat_link.is_symlink(): + try: + compat_link.symlink_to(versioned[-1]) + except FileExistsError: + # Multiple tensor-parallel ranks may prepare the same cache. + pass + + return [f"-L{link_dir}", f"-Wl,-rpath,{library_dir}"] + + raise RuntimeError("Unable to locate libamdhip64 for ROCm JIT linking") + + CPP_TEMPLATE_TYPE: TypeAlias = Union[int, float, bool] @@ -223,12 +304,16 @@ def load_aot( if prebuilt is not None: return prebuilt + is_rocm = _is_rocm() arch_list: List[str] = [] if cuda_files: - from freetoken.kernel._toolchain import check_nvcc_matches_torch + if is_rocm: + arch_list = _rocm_arch_list() + else: + from freetoken.kernel._toolchain import check_nvcc_matches_torch - check_nvcc_matches_torch() - arch_list = _cuda_arch_list() + check_nvcc_matches_torch() + arch_list = _cuda_arch_list() from tvm_ffi.cpp import load @@ -242,14 +327,23 @@ def load_aot( cpp_files = [str((KERNEL_PATH / "src" / f).resolve()) for f in cpp_files] cuda_files = [str((KERNEL_PATH / "src" / f).resolve()) for f in cuda_files] - with _pin_tvm_ffi_arch_ctx(arch_list): + if is_rocm: + cuda_cflags = _hip_cflags(extra_cuda_cflags, arch_list) + runtime_ldflags = _rocm_link_flags() + arch_list_env = ROCM_ARCH_LIST_ENV + else: + cuda_cflags = _cuda_cflags(extra_cuda_cflags, arch_list) + runtime_ldflags = [] + arch_list_env = CUDA_ARCH_LIST_ENV + + with _pin_tvm_ffi_arch_ctx(arch_list, arch_list_env): return load( name, cpp_files=cpp_files, cuda_files=cuda_files, extra_cflags=DEFAULT_CFLAGS + extra_cflags, - extra_cuda_cflags=_cuda_cflags(extra_cuda_cflags, arch_list), - extra_ldflags=DEFAULT_LDFLAGS + extra_ldflags, + extra_cuda_cflags=cuda_cflags, + extra_ldflags=DEFAULT_LDFLAGS + runtime_ldflags + extra_ldflags, extra_include_paths=DEFAULT_INCLUDE + extra_include_paths, build_directory=build_directory, ) @@ -272,12 +366,16 @@ def load_jit( if prebuilt is not None: return prebuilt + is_rocm = _is_rocm() arch_list: List[str] = [] if cuda_files or cuda_wrappers: - from freetoken.kernel._toolchain import check_nvcc_matches_torch + if is_rocm: + arch_list = _rocm_arch_list() + else: + from freetoken.kernel._toolchain import check_nvcc_matches_torch - check_nvcc_matches_torch() - arch_list = _cuda_arch_list() + check_nvcc_matches_torch() + arch_list = _cuda_arch_list() from tvm_ffi.cpp import load_inline @@ -300,14 +398,23 @@ def load_jit( cuda_sources = [f'#include "{path}"' for path in cuda_paths] cuda_sources += [_make_wrapper(tup) for tup in cuda_wrappers] - with _pin_tvm_ffi_arch_ctx(arch_list): + if is_rocm: + cuda_cflags = _hip_cflags(extra_cuda_cflags, arch_list) + runtime_ldflags = _rocm_link_flags() + arch_list_env = ROCM_ARCH_LIST_ENV + else: + cuda_cflags = _cuda_cflags(extra_cuda_cflags, arch_list) + runtime_ldflags = [] + arch_list_env = CUDA_ARCH_LIST_ENV + + with _pin_tvm_ffi_arch_ctx(arch_list, arch_list_env): return load_inline( name, cpp_sources=cpp_sources, cuda_sources=cuda_sources, extra_cflags=DEFAULT_CFLAGS + extra_cflags, - extra_cuda_cflags=_cuda_cflags(extra_cuda_cflags, arch_list), - extra_ldflags=DEFAULT_LDFLAGS + extra_ldflags, + extra_cuda_cflags=cuda_cflags, + extra_ldflags=DEFAULT_LDFLAGS + runtime_ldflags + extra_ldflags, extra_include_paths=DEFAULT_INCLUDE + extra_include_paths, build_directory=build_directory, ) 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..7da2011bb3 100644 --- a/python/freetoken/kvcache/base.py +++ b/python/freetoken/kvcache/base.py @@ -9,6 +9,9 @@ logger = init_logger(__name__) +# One fp32 scale per (token, slab, layer, kv head) rides alongside fp8 KV codes. +FP8_KV_SCALE_BYTES = 4 + class CacheRebuildRejected(Exception): """A runtime cache rebuild was rejected BEFORE any destructive free (e.g. the @@ -16,23 +19,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 @@ -119,6 +185,24 @@ def validate_rebuild( f"needs {mem_GB(need)} > budget {mem_GB(budget)}; old cache kept, still serving" ) + @classmethod + def rebuild_footprint( + cls, config, num_pages: int, num_swa_pages: int | None + ) -> int: + """Whole-family GPU bytes at ``(num_pages, num_swa_pages)`` -- the pool's own + tiers priced together, same kv_cost plumbing validate_rebuild uses. The engine + compares old-vs-new footprints to order a live rebuild's resizes + shrinking-first (#643); ``num_swa_pages=None`` prices the config's current + window pin.""" + import inspect + + cost_params = inspect.signature(cls.kv_cost).parameters + cost_kwargs = ( + {"num_swa_pages": num_swa_pages} if "num_swa_pages" in cost_params else {} + ) + cache_per_page, fixed_cache_size, _, _ = cls.kv_cost(config, **cost_kwargs) + return num_pages * cache_per_page + fixed_cache_size + @abstractmethod def rebuild_from_config( self, config, num_pages: int, *, num_swa_pages: int | None = None @@ -154,13 +238,44 @@ def store_kv( layer_id: int, ) -> None: ... + def k_scale(self, index: int) -> torch.Tensor | None: + """fp32 ``[num_slots, local_kv_heads]`` scale of ``k_cache(index)``, indexed by the + same slot; None when the pool is not quantized.""" + return None + + def v_scale(self, index: int) -> torch.Tensor | None: + """fp32 ``[num_slots, local_kv_heads]`` scale of ``v_cache(index)``; see k_scale.""" + return None + + def k_block_scale(self, index: int) -> torch.Tensor | None: + """NVFP4 E4M3 scales for ``k_cache(index)``, or ``None`` for other layouts.""" + return None + + def v_block_scale(self, index: int) -> torch.Tensor | None: + """NVFP4 E4M3 scales for ``v_cache(index)``, or ``None`` for other layouts.""" + return None + @property @abstractmethod def device(self) -> torch.device: ... @property @abstractmethod - def dtype(self) -> torch.dtype: ... + def dtype(self) -> torch.dtype: + """The pool's COMPUTE dtype: the dtype of the K/V rows a backend hands to + ``store_kv``, and what backends size their scratch with. A quantized (e4m3) + pool still answers 16-bit here -- code bytes handed to a scratch buffer end up + as the rhs of a ``tl.dot`` that has no fp8 path (QSA's indexer died this way at + graph capture). See :attr:`store_dtype` for what the buffer holds.""" + ... + + @property + def store_dtype(self) -> torch.dtype: + """Element type of the KV buffer: e4m3/uint8 codes on a quantized pool, else + :attr:`dtype`. Only code that touches the buffer itself needs this; attention + backends that cannot apply the row scales are refused a quantized pool up + front (``BackendInfo.supports_fp8_kv``).""" + return self.dtype @property @abstractmethod diff --git a/python/freetoken/kvcache/cache_status.py b/python/freetoken/kvcache/cache_status.py index 10169d6515..8c29dccab5 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]: @@ -39,6 +40,9 @@ def _moe() -> int: banks = getattr(engine.moe_offload_cache, "bank_caches", None) if not banks: return 0 + if getattr(engine.moe_offload_cache, "_raw_banks", False): + # pools of several bank layouts share byte segments: report the mean slot size + return int(sum(t.numel() for t in banks.values()) // engine.moe_offload_cache.cache_size) # Each bank cache is (cache_size, *row_shape); one slot's bytes = row bytes summed over # the format's banks (== cache_budget.expert_bytes_per_slot on the source rows). return int(sum(t[0].numel() * t.element_size() for t in banks.values())) @@ -120,21 +124,8 @@ def _mamba() -> int: return int(_linear_pool_min_slots(config) - 1) def _swa() -> int: - # Window-pool floor in tokens (matches the pool's own page unit x count): - # DSV4 -- the working-set floor pages _dsv4_pool_sizes clamps n_win_pages to, x P. - # radix-SWA -- the concurrent working set (+1 slot-0 sentinel) _swa_paged_num_tokens - # floors to. - # 0 for models without a window pool. - 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) * P) - if not (mc.has_swa_attention and config.cache_type == "swa_radix"): - return 0 - return int(_swa_pool_floor(config) + 1) + 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)): @@ -160,19 +151,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..484dcbbb29 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: @@ -423,123 +279,32 @@ def validate_rebuild( f"{mem_GB(need)} > budget {mem_GB(budget)}; old cache kept, still serving" ) + @classmethod + def rebuild_footprint(cls, config, num_pages: int, num_swa_pages: int | None) -> int: + # The DSV4 kv_cost does not take num_swa_pages, so price the family with the same + # byte model validate_rebuild uses (num_pages is USABLE; +1 for the dummy page). + from .v4_cost_model import _dsv4_pool_sizes, dsv4_pool_bytes + + dsv4_args = config.model_config.dsv4_args + sizes = _dsv4_pool_sizes(config, num_pages + 1, num_swa_pages=num_swa_pages) + return dsv4_pool_bytes(sizes, dsv4_args, config.max_running_req + 1) + def rebuild_from_config( self, config, num_pages: int, *, num_swa_pages: int | None = None ) -> None: - 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 +351,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/launch.py b/python/freetoken/launch.py index 5aac9389fe..e3c42154ab 100644 --- a/python/freetoken/launch.py +++ b/python/freetoken/launch.py @@ -5,6 +5,7 @@ import os import shlex import shutil +import stat import subprocess import sys import time @@ -20,7 +21,7 @@ CODEX_PROFILE = "freetoken-launch" CODEX_PROVIDER_NAME = "FreeToken" CODEX_CATALOG_NAME = "freetoken-model.json" -CODEX_PROVIDER_API_KEY_ENV = "FREETOKEN_API_KEY" +API_KEY_ENV = "FREETOKEN_API_KEY" # Used when the server reports no context length. Guessing low only costs earlier compaction. FALLBACK_CONTEXT_WINDOW = 128_000 MAX_OUTPUT_TOKENS_CAP = 32_768 @@ -71,6 +72,7 @@ class LaunchContext: extra_args: list[str] dry_run: bool assume_yes: bool = False + api_key: str | None = None @dataclass(frozen=True) @@ -167,9 +169,15 @@ def resolve_server_url(server: str | None) -> ServerURL: return ServerURL(origin=origin, openai_base_url=f"{origin}/v1") -def _get_json(url: str) -> object: - request = Request(url, headers={"Accept": "application/json"}) - with urlopen(request, timeout=5) as response: +def resolve_api_key(api_key: str | None) -> str | None: + return api_key or os.environ.get(API_KEY_ENV) or None + + +def _get_json(url: str, api_key: str | None) -> object: + headers = {"Accept": "application/json"} + if api_key: + headers["Authorization"] = f"Bearer {api_key}" + with urlopen(Request(url, headers=headers), timeout=5) as response: return json.loads(response.read().decode("utf-8")) @@ -177,10 +185,10 @@ def _positive_int(value: object) -> int | None: return value if isinstance(value, int) and not isinstance(value, bool) and value > 0 else None -def _stats_model(server: ServerURL) -> dict[str, object]: +def _stats_model(server: ServerURL, api_key: str | None) -> dict[str, object]: """The /v1/stats model card; empty when the server predates it.""" try: - payload = _get_json(f"{server.origin}/v1/stats") + payload = _get_json(f"{server.origin}/v1/stats", api_key) except (HTTPError, URLError, HTTPException, OSError, TimeoutError, ValueError): return {} model = payload.get("model") if isinstance(payload, dict) else None @@ -193,11 +201,16 @@ def _input_modalities(stats_model: dict[str, object]) -> tuple[str, ...]: return ("text", *extra) -def discover_server_model(server: ServerURL) -> ServedModel: +def discover_server_model(server: ServerURL, api_key: str | None) -> ServedModel: try: - _get_json(server.openai_base_url) - payload = _get_json(f"{server.openai_base_url}/models") + _get_json(server.openai_base_url, api_key) + payload = _get_json(f"{server.openai_base_url}/models", api_key) except (HTTPError, URLError, HTTPException, OSError, TimeoutError) as exc: + if isinstance(exc, HTTPError) and exc.code == 401: + raise RuntimeError( + f"FreeToken server at {server.origin} needs a valid API key: " + f"pass --api-key or set {API_KEY_ENV}" + ) from exc raise RuntimeError(f"Cannot connect to FreeToken server at {server.origin}: {exc}") from exc if not isinstance(payload, dict): @@ -220,7 +233,7 @@ def discover_server_model(server: ServerURL) -> ServedModel: if not models: raise RuntimeError(f"FreeToken server at {server.origin} reported no models") - stats_model = _stats_model(server) + stats_model = _stats_model(server, api_key) if context_length is None: context_length = _positive_int(stats_model.get("ctx")) @@ -236,15 +249,22 @@ def _backup_path(path: Path) -> Path: return path.with_name(f"{path.name}.{time.time_ns()}.bak") -def _write_text_with_backup(path: Path, text: str) -> None: +def _write_text_with_backup(path: Path, text: str, *, private: bool = False) -> None: path.parent.mkdir(parents=True, exist_ok=True) if path.exists() and path.read_text() != text: - shutil.copy2(path, _backup_path(path)) + backup = _backup_path(path) + # copy2 sets the mode only after it copies the content, so create the backup with the source mode first + backup.touch(mode=stat.S_IMODE(path.stat().st_mode)) + shutil.copy2(path, backup) + if private: + # before the write, so the key never sits in a file others can read + path.touch(mode=0o600) + path.chmod(0o600) path.write_text(text) -def _write_json_with_backup(path: Path, value: object) -> None: - _write_text_with_backup(path, json.dumps(value, indent=2) + "\n") +def _write_json_with_backup(path: Path, value: object, *, private: bool = False) -> None: + _write_text_with_backup(path, json.dumps(value, indent=2) + "\n", private=private) def _toml_string(value: str) -> str: @@ -371,7 +391,7 @@ def _codex_profile_text(ctx: LaunchContext, catalog_path: Path) -> str: + f"name = {_toml_string(CODEX_PROVIDER_NAME)}\n" + f"base_url = {_toml_string(ctx.server.openai_base_url)}\n" + 'wire_api = "responses"\n' - + f"env_key = {_toml_string(CODEX_PROVIDER_API_KEY_ENV)}\n" + + f"env_key = {_toml_string(API_KEY_ENV)}\n" ) @@ -408,7 +428,7 @@ def prepare_codex(ctx: LaunchContext) -> CommandSpec: ] return CommandSpec( argv=argv, - env={CODEX_PROVIDER_API_KEY_ENV: "freetoken"}, + env={API_KEY_ENV: ctx.api_key or "freetoken"}, unset_env=CODEX_CLEAR_ENV, ) @@ -423,7 +443,7 @@ def prepare_claude(ctx: LaunchContext) -> CommandSpec: "CLAUDE_CODE_MAX_OUTPUT_TOKENS": str(_max_output_tokens(ctx)), "ANTHROPIC_BASE_URL": ctx.server.origin, "ANTHROPIC_API_KEY": "", - "ANTHROPIC_AUTH_TOKEN": "freetoken", + "ANTHROPIC_AUTH_TOKEN": ctx.api_key or "freetoken", "ANTHROPIC_MODEL": model_id, "ANTHROPIC_DEFAULT_OPUS_MODEL": model_id, "ANTHROPIC_DEFAULT_SONNET_MODEL": model_id, @@ -456,13 +476,16 @@ def _opencode_model_entries( def _opencode_config(ctx: LaunchContext) -> str: model_ids = _ordered_model_ids(ctx) + options = {"baseURL": ctx.server.openai_base_url} + if ctx.api_key: + options["apiKey"] = ctx.api_key config = { "$schema": "https://opencode.ai/config.json", "provider": { OPENCODE_PROVIDER: { "npm": "@ai-sdk/openai-compatible", "name": OPENCODE_PROVIDER_NAME, - "options": {"baseURL": ctx.server.openai_base_url}, + "options": options, "models": _opencode_model_entries( model_ids, _context_window(ctx), _max_output_tokens(ctx), _accepts_images(ctx) ), @@ -614,7 +637,7 @@ def _patch_openclaw_config( new_models.append(entry) provider["baseUrl"] = ctx.server.openai_base_url - provider["apiKey"] = "freetoken-local" + provider["apiKey"] = ctx.api_key or "freetoken-local" provider["api"] = "openai-completions" provider["models"] = new_models providers[OPENCLAW_PROVIDER] = provider @@ -661,7 +684,9 @@ def prepare_openclaw(ctx: LaunchContext) -> CommandSpec: if not _openclaw_has_freetoken_provider(config): if not ctx.assume_yes and not _confirm_openclaw_first_patch(ctx, config_path): raise RuntimeError("OpenClaw config patch cancelled") - _write_json_with_backup(config_path, _patch_openclaw_config(ctx, config)) + _write_json_with_backup( + config_path, _patch_openclaw_config(ctx, config), private=bool(ctx.api_key) + ) _clear_openclaw_session_overrides(ctx.model.model_id) argv = ["openclaw", *ctx.extra_args] if ctx.extra_args else ["openclaw", "chat"] @@ -697,8 +722,8 @@ def prepare_hermes(ctx: LaunchContext) -> CommandSpec: model_section["default"] = ctx.model.model_id model_section["provider"] = "custom" model_section["base_url"] = ctx.server.openai_base_url - # A non-empty dummy: FreeToken is unauthenticated, but some clients reject an empty key. - model_section["api_key"] = HERMES_API_KEY + # Dummy for a server without --api-key; some clients reject an empty key. + model_section["api_key"] = ctx.api_key or HERMES_API_KEY # Hermes' name for the window (prompt + generation). Left unset it auto-detects, and a # miss lands on a 32k default that trips the floor below. window = _context_window(ctx) @@ -712,7 +737,9 @@ def prepare_hermes(ctx: LaunchContext) -> CommandSpec: file=sys.stderr, ) - _write_text_with_backup(config_path, yaml.safe_dump(config, sort_keys=False)) + _write_text_with_backup( + config_path, yaml.safe_dump(config, sort_keys=False), private=bool(ctx.api_key) + ) argv = ["hermes", *ctx.extra_args] if ctx.extra_args else ["hermes", "chat"] return CommandSpec(argv=argv, env={}) @@ -753,8 +780,7 @@ def prepare_dsh(ctx: LaunchContext) -> CommandSpec: ``reasoning_content``, while llm-pi-ai's JSON round-trip can change argument values and key order — drift that survives render canonicalization, cuts the prefix cache, and shows the model a rewrite of its own output. The - dummy DEEPSEEK_API_KEY satisfies dsh's non-empty key requirement; FreeToken - itself is unauthenticated.""" + dummy DEEPSEEK_API_KEY satisfies dsh's non-empty key requirement.""" settings_path = _dsh_home() / DSH_LAUNCH_SETTINGS_NAME patch_path = _dsh_home() / DSH_LAUNCH_PATCH_NAME if not ctx.dry_run: @@ -774,7 +800,7 @@ def prepare_dsh(ctx: LaunchContext) -> CommandSpec: window, output = _context_window(ctx), _max_output_tokens(ctx) modalities = ["text", "image"] if _accepts_images(ctx) else ["text"] - config["llm-deepseek"] = { + section: dict[str, object] = { "baseURL": ctx.server.openai_base_url, "models": [ { @@ -787,6 +813,10 @@ def prepare_dsh(ctx: LaunchContext) -> CommandSpec: for model_id in _ordered_model_ids(ctx) ], } + if ctx.api_key: + # not DEEPSEEK_API_KEY: dsh's web search also reads it and sends it to api.deepseek.com + section["apiKeyEnv"] = API_KEY_ENV + config["llm-deepseek"] = section config["agent-default-model"] = { "provider": "deepseek-official", "model": ctx.model.model_id, @@ -804,15 +834,14 @@ def prepare_dsh(ctx: LaunchContext) -> CommandSpec: elif len(app_args) >= 2 and app_args[0] == "--profile": profile, app_args = app_args[1], app_args[2:] argv = ["dsh", "--profile", profile, "--patch", str(patch_path), *app_args] - return CommandSpec( - argv=argv, - env={ - "DEEPSEEK_BASE_URL": ctx.server.openai_base_url, - "DEEPSEEK_API_KEY": DSH_API_KEY, - "DSH_TELEMETRY_DISABLED": "1", - }, - unset_env=CLOUD_PROVIDER_API_KEY_ENV, - ) + env = { + "DEEPSEEK_BASE_URL": ctx.server.openai_base_url, + "DEEPSEEK_API_KEY": DSH_API_KEY, + "DSH_TELEMETRY_DISABLED": "1", + } + if ctx.api_key: + env[API_KEY_ENV] = ctx.api_key + return CommandSpec(argv=argv, env=env, unset_env=CLOUD_PROVIDER_API_KEY_ENV) PREPARERS = { @@ -840,6 +869,10 @@ def parse_argv(argv: list[str], prog: str = "python -m freetoken.launch") -> arg ) parser.add_argument("agent", choices=sorted(PREPARERS)) parser.add_argument("--server", help="FreeToken server origin or /v1 URL") + parser.add_argument( + "--api-key", + help=f"API key of a server started with --api-key (default: ${API_KEY_ENV})", + ) parser.add_argument( "--dry-run", action="store_true", @@ -1028,13 +1061,16 @@ def main(argv: list[str] | None = None, prog: str = "python -m freetoken.launch" return 0 server = resolve_server_url(args.server) - model = discover_server_model(server) + api_key = resolve_api_key(args.api_key) + model = discover_server_model(server, api_key) plan_ctx = LaunchContext( server=server, model=model, extra_args=args.extra_args, dry_run=True, assume_yes=args.yes, + # the plan is printed by --dry-run, so it never holds the real key + api_key="" if api_key else None, ) spec = PREPARERS[args.agent](plan_ctx) if args.dry_run: @@ -1052,6 +1088,7 @@ def main(argv: list[str] | None = None, prog: str = "python -m freetoken.launch" extra_args=args.extra_args, dry_run=False, assume_yes=args.yes, + api_key=api_key, ) spec = PREPARERS[args.agent](ctx) if binary_path != spec.argv[0]: 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/activation.py b/python/freetoken/layers/activation.py index cd7a402a65..e38a6caebc 100644 --- a/python/freetoken/layers/activation.py +++ b/python/freetoken/layers/activation.py @@ -6,10 +6,17 @@ import torch -def silu_and_mul(x: torch.Tensor, out: torch.Tensor | None = None): +def _use_flashinfer(x: torch.Tensor, out: torch.Tensor | None) -> bool: + """flashinfer's *_and_mul kernels take fp16 / bf16 in and the same dtype out; the Triton kernels take any float dtypes.""" + import torch + from freetoken.kernel.backend import is_flashinfer_installed - if is_flashinfer_installed(): + return x.dtype in (torch.float16, torch.bfloat16) and (out is None or out.dtype == x.dtype) and is_flashinfer_installed() + + +def silu_and_mul(x: torch.Tensor, out: torch.Tensor | None = None): + if _use_flashinfer(x, out): from flashinfer import silu_and_mul else: from freetoken.kernel.triton.activation import silu_and_mul @@ -18,9 +25,7 @@ def silu_and_mul(x: torch.Tensor, out: torch.Tensor | None = None): def gelu_and_mul(x: torch.Tensor, out: torch.Tensor | None = None): - from freetoken.kernel.backend import is_flashinfer_installed - - if is_flashinfer_installed(): + if _use_flashinfer(x, out): from flashinfer import gelu_and_mul else: from freetoken.kernel.triton.activation import gelu_and_mul @@ -30,9 +35,7 @@ def gelu_and_mul(x: torch.Tensor, out: torch.Tensor | None = None): def gelu_tanh_and_mul(x: torch.Tensor, out: torch.Tensor | None = None): """tanh-approximate GELU gate (`gelu_pytorch_tanh`) followed by elementwise mul.""" - from freetoken.kernel.backend import is_flashinfer_installed - - if is_flashinfer_installed(): + if _use_flashinfer(x, out): from flashinfer import gelu_tanh_and_mul else: from freetoken.kernel.triton.activation import gelu_tanh_and_mul diff --git a/python/freetoken/layers/embedding.py b/python/freetoken/layers/embedding.py index a22216548a..eaa3608826 100644 --- a/python/freetoken/layers/embedding.py +++ b/python/freetoken/layers/embedding.py @@ -3,7 +3,6 @@ from typing import Dict 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.utils import div_ceil, nvtx_annotate @@ -13,11 +12,18 @@ class VocabParallelEmbedding(BaseOP): + """The vocab table, sharded over TP: its ``quant_method`` owns the table, looks rows up, and serves a tied LM head's product.""" + + quant_layer_kind = LayerKind.EMBEDDING + def __init__( self, num_embeddings: int, embedding_dim: int, embed_scale: float | None = None, + *, + quant_config: QuantConfig | None = None, + prefix: str = "", ): super().__init__() tp_info = get_tp_info() @@ -28,7 +34,12 @@ def __init__( start_idx = self.num_embeddings_tp * tp_rank finish_idx = min(start_idx + self.num_embeddings_tp, num_embeddings) self.vocab_range = (start_idx, finish_idx - start_idx) - self.weight = torch.empty(self.num_embeddings_tp, embedding_dim) + self.embedding_dim = embedding_dim + # the LM head subclass is a linear and builds its own method + self.quant_method = None + if self.quant_layer_kind is LayerKind.EMBEDDING: + self.quant_method = quant_method_for(quant_config, self, prefix) + self.quant_method.create_weights(self) # Gemma scales embeddings by sqrt(hidden_size). The scale is materialized in # the weight dtype (bf16) to match HF, which downcasts the scalar. The GPU # scalar is built lazily (model __init__ runs on the meta device) and cached @@ -36,17 +47,34 @@ def __init__( self._embed_scale = embed_scale self._embed_scale_t: torch.Tensor | None = None self._comm = DistributedCommunicator() + self._host_ptr: int | None = None - @nvtx_annotate("Embedding") - def forward(self, x: torch.Tensor) -> torch.Tensor: - from freetoken.kernel import indexing + 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 - y = indexing( - weights=self.weight, - indices=x, - vocab_range=self.vocab_range if self.tp_size > 1 else None, - ) + 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: + 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 = self.quant_method.embed(self, x) if self.tp_size > 1: y = self._comm.all_reduce(y) if self._embed_scale is not None: @@ -83,7 +111,6 @@ def __init__( self.in_features = embedding_dim self.out_features = self.num_embeddings_tp self.output_sizes = (self.num_embeddings_tp,) - self.quant_method = None if tied_embedding is None: self.quant_method = quant_method_for(quant_config, self, prefix) self.quant_method.create_weights(self) @@ -121,6 +148,12 @@ 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.""" + if self.tied_embedding is not None: + return self.tied_embedding.quant_method.logits(self.tied_embedding, x, 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 +164,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..1d8713cc8e --- /dev/null +++ b/python/freetoken/layers/fp8_dynamic.py @@ -0,0 +1,252 @@ +"""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-row scale (``quant_per_row``, applied +after the GEMM). ``quant_per_tensor`` is the older per-tensor cast: 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 + + +@triton.jit +def _quant_rows_kernel(x_ptr, out_ptr, scale_ptr, k, BLOCK: tl.constexpr): + """One program per row: that row's max|x|, then its cast.""" + row = tl.program_id(0).to(tl.int64) + x_row = x_ptr + row * k + acc = tl.zeros([BLOCK], dtype=tl.float32) + for start in range(0, k, BLOCK): + offs = start + tl.arange(0, BLOCK) + acc = tl.maximum(acc, tl.abs(tl.load(x_row + offs, mask=offs < k, other=0.0).to(tl.float32))) + amax = tl.maximum(tl.max(acc, axis=0), 1e-12) + tl.store(scale_ptr + row, amax / 448.0) + inv = 448.0 / amax + for start in range(0, k, BLOCK): + offs = start + tl.arange(0, BLOCK) + mask = offs < k + v = tl.load(x_row + offs, mask=mask, other=0.0).to(tl.float32) * inv + v = tl.minimum(tl.maximum(v, -448.0), 448.0) + tl.store(out_ptr + row * k + offs, v.to(tl.float8e4nv), mask=mask) + + +def quant_per_row(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """``(x_fp8, scale)`` with ``x ~= x_fp8 * scale``; ``x`` contiguous ``[M, K]``, ``scale`` fp32 ``[M, 1]``.""" + m, k = x.shape + out = torch.empty_like(x, dtype=FP8) + scale = torch.empty((m, 1), dtype=torch.float32, device=x.device) + if m: + _quant_rows_kernel[(m,)](x, out, scale, k, BLOCK=1024, num_warps=4) + 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). + + The activation scale is per row, so a row never depends on the other sequences of its + batch. It is applied after the GEMM: cuBLASLt's rowwise-scaled mode is 1.4-2.9x slower + than the tensorwise one on sm_120.""" + *lead, k = x.shape + x2 = x.reshape(-1, k).contiguous() + if x2.shape[0] == 1: # one row: per-tensor is per-row, and the scale rides in the GEMM epilogue + 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]) + x8, row_scale = quant_per_row(x2) + y = torch._scaled_mm( + x8, weight.t(), scale_a=weight_scale.new_ones(()), scale_b=weight_scale, out_dtype=x.dtype + ) + y.mul_(row_scale) + 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_row", + "quant_per_tensor", +] diff --git a/python/freetoken/layers/gguf.py b/python/freetoken/layers/gguf.py deleted file mode 100644 index ac49b1a5be..0000000000 --- a/python/freetoken/layers/gguf.py +++ /dev/null @@ -1,128 +0,0 @@ -"""Native-GGUF quantized layers: weights stay in their packed block layout and are -dequantized *inside* the borrowed llama.cpp kernels (no bf16 copy ever materialized). - -Mirrors vLLM/sglang's ``GGUFLinearMethod`` / ``GGUFEmbeddingMethod`` dispatch, ported -onto FreeToken's ``BaseOP``. FreeToken keeps fused projections (qkv, gate_up) as a -single tensor: because Q4_0/K-quants pack each *output row* independently over the -input dim, the loader can concatenate the per-shard packed rows along dim 0 (they -share an input dim, hence the same ``row_bytes``), so a fused layer is still one -``[out, row_bytes]`` qweight -- no per-shard padding bookkeeping needed. - -TP is assumed to be 1 (the gemma4 GGUF path restricts to TP=1, like the HF path). -""" - -from __future__ import annotations - -import torch - -from freetoken.models.gguf.dequant import ( - BLOCK_SHAPE, - GGML_BF16, - GGML_F16, - GGML_F32, - GGML_NAME, - GGML_Q4_0, - GGML_Q6_K, - GGML_Q8_0, - row_bytes, -) - -from .base import BaseOP - -# ggml type groups for kernel dispatch (subset we build kernels for). -_UNQUANTIZED = {GGML_F32, GGML_F16, GGML_BF16} -# standard + k-quants: both an MMVQ (small-batch GEMV) and MMQ (large-batch) kernel exist. -_MMVQ = {GGML_Q4_0, GGML_Q8_0, GGML_Q6_K} -_MMQ = {GGML_Q4_0, GGML_Q8_0, GGML_Q6_K} -_DEQUANT = {GGML_Q4_0, GGML_Q8_0, GGML_Q6_K} - -# Below this token count, the MMVQ GEMV kernel wins (matches vLLM's heuristic). -_MMVQ_SAFE = 6 - - -def fused_mul_mat_gguf(x: torch.Tensor, qweight: torch.Tensor, qweight_type: int) -> torch.Tensor: - """y = x @ dequant(qweight).T, dispatched by batch size and quant type.""" - from freetoken.kernel.gguf import ( - ggml_dequantize, - ggml_mul_mat_a8, - ggml_mul_mat_vec_a8, - ) - - out_features = qweight.shape[0] - if x.shape[0] == 0: - return x.new_empty((0, out_features)) - if qweight_type in _UNQUANTIZED: - return x @ qweight.T - if x.shape[0] <= _MMVQ_SAFE and qweight_type in _MMVQ: - return ggml_mul_mat_vec_a8(qweight, x, qweight_type, out_features) - if qweight_type in _MMQ: - return ggml_mul_mat_a8(qweight, x, qweight_type, out_features) - if qweight_type in _DEQUANT: - block, type_size = BLOCK_SHAPE[qweight_type] - in_features = qweight.shape[1] // type_size * block - weight = ggml_dequantize(qweight, qweight_type, out_features, in_features, x.dtype) - return x @ weight.T - raise NotImplementedError(f"unsupported GGUF type {GGML_NAME.get(qweight_type, qweight_type)}") - - -class GGUFLinear(BaseOP): - """Linear whose weight is a native GGUF block-quantized ``[out, row_bytes]`` tensor.""" - - def __init__( - self, - in_features: int, - out_features: int, - quant_type: int, - has_bias: bool = False, - ): - self.in_features = in_features - self.out_features = out_features - self._quant_type = quant_type - self.qweight = torch.empty(out_features, row_bytes(in_features, quant_type), dtype=torch.uint8) - self.bias = torch.empty(out_features) if has_bias else None - - def forward(self, x: torch.Tensor) -> torch.Tensor: - out = fused_mul_mat_gguf(x, self.qweight, self._quant_type) - if self.bias is not None: - out = out + self.bias - return out - - -class GGUFEmbedding(BaseOP): - """Vocab embedding stored as a native GGUF block-quantized table. - - The full table is never dequantized: only the looked-up rows are gathered (in - packed form) and dequantized per lookup, matching vLLM's ``_apply_gguf_embedding``. - """ - - def __init__( - self, - num_embeddings: int, - embedding_dim: int, - quant_type: int, - embed_scale: float | None = None, - ): - self.num_embeddings = num_embeddings - self.embedding_dim = embedding_dim - self._quant_type = quant_type - self.qweight = torch.empty( - num_embeddings, row_bytes(embedding_dim, quant_type), dtype=torch.uint8 - ) - self._embed_scale = embed_scale - self._embed_scale_t: torch.Tensor | None = None - - def forward(self, x: torch.Tensor) -> torch.Tensor: - from freetoken.kernel.gguf import ggml_dequantize - - flat = x.flatten() - rows = self.qweight.index_select(0, flat) # [n, row_bytes] packed - y = ggml_dequantize(rows, self._quant_type, flat.shape[0], self.embedding_dim, torch.bfloat16) - y = y.view(*x.shape, self.embedding_dim) - if self._embed_scale is not None: - if self._embed_scale_t is None: - self._embed_scale_t = torch.tensor(self._embed_scale, dtype=y.dtype, device=y.device) - y = y * self._embed_scale_t - return y - - -__all__ = ["GGUFLinear", "GGUFEmbedding", "fused_mul_mat_gguf"] diff --git a/python/freetoken/layers/linear.py b/python/freetoken/layers/linear.py index d8a3e8619f..e28aaa07df 100644 --- a/python/freetoken/layers/linear.py +++ b/python/freetoken/layers/linear.py @@ -63,6 +63,7 @@ def __init__( output_size: int, has_bias: bool, *, + output_sizes: List[int] | None = None, quant_config: QuantConfig | None = None, prefix: str = "", ): @@ -72,6 +73,7 @@ def __init__( local_isize=input_size, local_osize=output_size, has_bias=has_bias, + output_sizes=output_sizes, quant_config=quant_config, prefix=prefix, ) @@ -83,13 +85,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..f9628b3b5c 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): @@ -81,16 +87,12 @@ def __init__( self.strategy = strategy self.decode_target = decode_target self.prefix = prefix - # offload layers without a quant config stay on the format-tag banks (GGUF q4_0) - self.quant_method = None - if quant_config is not None or allocate_experts: - self.quant_method = quant_method_for(quant_config, self, prefix) - if allocate_experts: - self.quant_method.create_weights(self) + self.quant_method = quant_method_for(quant_config, self, prefix) + if allocate_experts: + self.quant_method.create_weights(self) def finalize(self) -> None: - if self.quant_method is not None: - self.quant_method.finalize(self) + self.quant_method.finalize(self) def _maybe_all_reduce(self, hidden_states: torch.Tensor) -> torch.Tensor: if self.tp_size > 1: @@ -103,7 +105,6 @@ def _resident_gemm( topk_weights: torch.Tensor, topk_ids: torch.Tensor, ) -> torch.Tensor: - assert self.quant_method is not None return self.quant_method.apply( hidden_states, topk_weights, topk_ids, self.quant_method.resident_view(self), layer=self, is_prefill=get_global_ctx().batch.is_prefill, @@ -188,6 +189,17 @@ def __init__( prefix=prefix, ) self.offload_cache: OffloadMoeCache | None = None + # FREETOKEN_MOE_PREFETCH: (next MoE layer's router op, its layer id, its budget); the + # underscore keeps the borrowed router out of this layer's state dict + self._lookahead: tuple[BaseOP, int, int] | None = None + + def set_lookahead(self, gate: BaseOP, target_layer: int, budget: int) -> None: + """Let this layer's GPU decode predict ``target_layer``'s experts with that layer's router.""" + self._lookahead = (gate, target_layer, budget) + + @property + def skips_inactive_slots(self) -> bool: + return self.quant_method.kernel.supports_inactive_slots def forward( self, @@ -235,7 +247,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, @@ -253,8 +277,9 @@ def prefill_forward( # ------------------------------------------------------------------ # Data movement -- one decision tree for every quant format (the banks # registry makes the cache machinery bank-count agnostic). Decode loads - # on demand; prefill streams whole layers, double-buffered when overlap - # is enabled. The kernels only ever see bank views plus row indices; + # on demand, and so does a short enough prefill chunk; a longer one + # streams whole layers, double-buffered when overlap is enabled. The + # kernels only ever see bank views plus row indices; # which kernel runs is decided afterwards, in ``_expert_gemm``. # ------------------------------------------------------------------ @@ -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,18 +308,57 @@ def _decode_routed( return executor.decode(self.layer_id, hidden_states, topk_weights, topk_ids) if cache.decode_target == "hybrid": return self._decode_hybrid(cache, hidden_states, topk_weights, topk_ids) - cache.ensure_experts(self.layer_id, topk_ids) + prefetch, verify = cache.prefetch, cache.verify + if verify is not None: + verify.begin(cache, self.layer_id, topk_ids, prefetch and prefetch.plan_of(self.layer_id)) + cache.ensure_experts(self.layer_id, topk_ids, router_logits=router_logits) + target = None + if prefetch is not None: + # on: install the next layer's prediction (forked at this block's start) on this stream, the slot + # maps' only writer; before this layer's copy join, so a late prefetch copy never delays it + target = prefetch.install(self._lookahead and self._lookahead[1]) + prefetch.join_copy(self.layer_id) cache.copy_missing() - return self._expert_gemm( + if prefetch is not None: + prefetch.join_and_count(self.layer_id, topk_ids, cache.id_of_slot, cache.num_indices) + if target is not None: + prefetch.issue_copy(target) + if verify is not None: + # after issue_copy: the next layer's prefetch copy starts when it would without the checks + verify.before_gemm(cache, topk_ids, None if target is None else prefetch.plan_of(target)) + out = self._expert_gemm( cache, hidden_states, topk_weights, topk_ids, - views=cache.bank_views(), + views=cache.pool_bank_views(self.layer_id), n=None, alphas=cache.alphas_for_slots(self.layer_id), is_prefill=False, ) + if verify is not None: + verify.after_gemm(cache) + return out + + def fork_lookahead(self, hidden_states: torch.Tensor) -> None: + """FREETOKEN_MOE_PREFETCH: at the start of this layer's MoE block, predict the next layer's + experts from this block's router input on the predictor stream, beside this layer's router + and ensure. Only on the GPU decode path that ``_decode_routed`` takes and not under TP, and + never toward a CPU-decoded layer, whose ensure would not use it.""" + cache = self.offload_cache + if ( + self._lookahead is None + or cache is None + or cache.prefetch is None + or self.tp_size > 1 + or get_global_ctx().batch.is_prefill + or cache.is_cpu_layer(self.layer_id) + or cache.decode_target == "hybrid" + ): + return + gate, target, budget = self._lookahead + if not cache.is_cpu_layer(target): + cache.prefetch.fork(target, hidden_states, gate, cache.slot_for_id[target], budget) def _decode_hybrid( self, @@ -307,19 +372,19 @@ 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). + routing split runs inside the hybrid ensure kernel and the CPU submit/sync are host nodes. + Each route is computed exactly once -- CPU-assigned routes reach the GPU as slot -1 with + weight 0 (a kernel with ``supports_inactive_slots`` skips them without a read; any other reads + slot 0, whose storage is a payload or the bank's zero fill) and the CPU ids are -1 for + 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 - 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() + # the kernel reads every route's slot: give the inactive ones slot 0, which holds a + # complete payload or the bank's zero fill (never uninitialized storage), times weight 0 + # with several quants this holds only because the cache zeroes its staging area after prefill + clamp = not self.skips_inactive_slots + gpu_slots, gpu_w, cpu_ids = cache.ensure_experts_hybrid(self.layer_id, topk_ids, topk_weights, clamp=clamp) pending = executor.decode_submit(self.layer_id, hidden_states, topk_weights, cpu_ids) # Measurement knob: FREETOKEN_HYBRID_OVERLAP=0 syncs the CPU pool *before* the @@ -329,20 +394,19 @@ 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() gpu_routed = self._expert_gemm( cache, hidden_states, gpu_w, gpu_slots, - views=cache.bank_views(), + views=cache.pool_bank_views(self.layer_id), n=None, alphas=cache.alphas_for_slots(self.layer_id), is_prefill=False, ) - cpu_routed = cpu_routed_early if not _HYBRID_OVERLAP else executor.decode_sync(pending) - return gpu_routed + cpu_routed + if not _HYBRID_OVERLAP: + return gpu_routed + cpu_routed_early + return executor.decode_sync(pending, add_to=gpu_routed) def _prefill_routed( self, @@ -350,12 +414,34 @@ def _prefill_routed( topk_weights: torch.Tensor, topk_ids: torch.Tensor, ) -> torch.Tensor: - """Prefill movement: stream whole layers -- double-buffered behind the - previous layer's GEMMs when ``prefill_overlap`` is on, else a synchronous - ``materialize_layer``. In both, position == expert id, so the routing ids - pass through unmapped.""" + """Prefill movement. A chunk within ``prefill_demand_tokens`` loads its routed experts + into the slot cache as decode does and runs the same grouped GEMM over the slots. A + longer one streams whole layers -- double-buffered behind the previous layer's GEMMs + when ``prefill_overlap`` is on, else a synchronous ``materialize_layer`` -- where + position == expert id, so the routing ids pass through unmapped.""" cache = self.offload_cache assert cache is not None + if ( + 0 < hidden_states.shape[0] <= max(cache.prefill_demand_tokens, _SMALL_PREFILL_TOKENS) + and cache.decode_target in ("gpu", "hybrid") + and not cache.is_unpinned_layer(self.layer_id) + ): + # A short extension (an agent turn over a cached prefix) touches a fraction of each + # layer's experts; fetching those beats streaming all of them (~68 GB here). It runs + # the prefill GEMM over the slot cache (#676), so it matches the streamed layer bit for + # bit. An unpinned layer's copy_missing presumes position == expert id, so it streams + # below; unpinned layers only exist with the prefill overlap off, so no layer can enter + # the overlap double buffer after an earlier layer skipped begin_prefill. + return self._expert_gemm( + cache, + hidden_states, + topk_weights, + self._ensure_unique(cache, topk_ids), + views=cache.bank_views(), + n=cache.cache_size, + alphas=cache.alphas_for_slots(self.layer_id), + is_prefill=True, + ) if cache.prefill_overlap: views = self._wait_prefill_overlap(cache) out = self._expert_gemm( @@ -372,16 +458,39 @@ def _prefill_routed( return out cache.materialize_layer(self.layer_id) cache.copy_missing() - return self._expert_gemm( + out = self._expert_gemm( cache, hidden_states, topk_weights, topk_ids, - views=cache.bank_views(self.num_experts), + views=cache.stage_views(self.layer_id, 0), n=self.num_experts, alphas=cache.alphas_for_layer(self.layer_id), is_prefill=True, ) + cache.release_prefill_layer(self.layer_id) + return out + + def _ensure_unique(self, cache: OffloadMoeCache, topk_ids: torch.Tensor) -> torch.Tensor: + """Make this layer's routed experts resident and return ``topk_ids`` mapped to slots. + + The LRU kernel is sized for decode: its block is next_pow2(ids) x next_pow2(slots), so a + prefill's tokens x top_k ids overflow Triton's element limit, and even a block under it + can fail to launch on a full GPU. It gets the unique ids (at most num_experts) in + decode-width chunks; each chunk's misses are copied before the next ensure reuses the plan.""" + uniq, inverse = torch.unique(topk_ids.reshape(-1), return_inverse=True) + n = uniq.numel() + # Each ensure stamps a newer LRU step, so a later chunk evicts the earlier chunks' slots + # only once every other slot is gone; the engine's num_experts slot floor rules that out. + assert n <= cache.cache_size, f"{n} routed experts exceed the {cache.cache_size}-slot cache" + uniq = uniq.to(torch.int32) + # scored eviction ranks an earlier chunk's fresh installs coldest, so pin them explicitly + pin_since = cache.step.clone() if cache.cache_policy_id else None + for start in range(0, n, _ENSURE_CHUNK_IDS): + part = uniq[start : start + _ENSURE_CHUNK_IDS] # a contiguous view: ensure rewrites it in place + cache.ensure_experts(self.layer_id, part, update_state=False, pin_since=pin_since, counts_as_decode=False) # not a decode step + cache.copy_missing() + return uniq[inverse].view_as(topk_ids) def _wait_prefill_overlap(self, cache: OffloadMoeCache) -> tuple[torch.Tensor, ...]: """Double-buffer choreography for this layer's overlap prefill: kick off the @@ -397,7 +506,6 @@ def _wait_prefill_overlap(self, cache: OffloadMoeCache) -> tuple[torch.Tensor, . # ------------------------------------------------------------------ # Kernel dispatch: ``views`` are the bank tensors the movement step produced (in bank registration order) and ``topk_ids`` already index their rows. - # GGUF q4_0 experts still dispatch on the cache's format tag until they get a method. # ------------------------------------------------------------------ def _expert_gemm( @@ -412,27 +520,28 @@ def _expert_gemm( alphas: tuple[torch.Tensor, torch.Tensor] | None, is_prefill: bool, ) -> torch.Tensor: - if self.quant_method is not None: - from freetoken.moe.legacy_format import canonical_role # legacy_format imports this package + from freetoken.moe.legacy_format import canonical_role # legacy_format imports this package - view = ExpertView( - {canonical_role(name): t for name, t in zip(cache.bank_schema, views)}, - slots=None if n is not None else topk_ids, n=n, alphas=alphas, - ) - return self.quant_method.apply( - hidden_states, topk_weights, topk_ids, view, layer=self, is_prefill=is_prefill - ) - fmt = cache.quant_format - if fmt == "q4_0": - # Native GGUF Q4_0 experts: dequant-in-kernel grouped GEMV (MMVQ) over the - # streamed packed banks; topk_ids already index the cache slots / layer. - from freetoken.moe.fused_q4_0 import fused_experts_gguf_q4_0 - - gate_up, down = views - return fused_experts_gguf_q4_0( - hidden_states, gate_up, down, topk_weights, topk_ids, self.activation - ) - raise AssertionError(f"offload experts without a quant method only serve q4_0 banks, got {fmt!r}") + view = ExpertView( + {canonical_role(name): t for name, t in zip(cache.bank_schema, views)}, + slots=None if n is not None else topk_ids, n=n, alphas=alphas, + ) + return self.quant_method.apply( + hidden_states, topk_weights, topk_ids, view, layer=self, is_prefill=is_prefill + ) + + +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( @@ -483,10 +592,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/__init__.py b/python/freetoken/layers/quantization/__init__.py index 62019fc57d..905b80ef5d 100644 --- a/python/freetoken/layers/quantization/__init__.py +++ b/python/freetoken/layers/quantization/__init__.py @@ -1,6 +1,7 @@ from .scheme import GroupShape, QuantKind, QuantScheme, WeightDesc from .registry import LayerKind, dialects, method_class, methods_for, register_dialect, register_method from .method import KernelSelectionError, QuantMethod, finalize_quant, select_kernel +from .embedding import EmbeddingConfig, EmbeddingKernel, EmbeddingMethod, GGUFEmbeddingMethod, UnquantizedEmbeddingMethod from .linear import ( Fp8BlockLinearMethod, Fp8TensorLinearMethod, @@ -28,6 +29,7 @@ from .configs import ( CompressedTensorsConfig, Fp8BlockConfig, + GGUFConfig, ModelOptConfig, Mxfp4Config, NoQuantConfig, @@ -42,11 +44,12 @@ "QuantKind", "GroupShape", "QuantScheme", "WeightDesc", "LayerKind", "dialects", "method_class", "methods_for", "register_dialect", "register_method", "KernelSelectionError", "select_kernel", "QuantMethod", + "EmbeddingConfig", "EmbeddingKernel", "EmbeddingMethod", "UnquantizedEmbeddingMethod", "GGUFEmbeddingMethod", "LinearConfig", "LinearKernel", "LinearMethod", "UnquantizedLinearMethod", "Fp8TensorLinearMethod", "Fp8BlockLinearMethod", "Mxfp8LinearMethod", "Nvfp4LinearMethod", "BankSpec", "ExpertView", "MoEConfig", "MoEKernel", "MoEMethod", "UnquantizedMoEMethod", "Fp8BlockMoEMethod", "Nvfp4MoEMethod", "Mxfp4MoEMethod", "Mxfp8MoEMethod", "QuantBackend", "set_quant_backend", "get_quant_backend", "NameMap", - "QuantConfig", "NoQuantConfig", "ModelOptConfig", "CompressedTensorsConfig", "Fp8BlockConfig", "Mxfp4Config", + "QuantConfig", "NoQuantConfig", "ModelOptConfig", "CompressedTensorsConfig", "Fp8BlockConfig", "GGUFConfig", "Mxfp4Config", "quant_method_for", "quantization_config_of", "set_quant_config", "get_quant_config", "finalize_quant", ] diff --git a/python/freetoken/layers/quantization/configs/__init__.py b/python/freetoken/layers/quantization/configs/__init__.py index 4bd7a6a9a7..358e6fecd6 100644 --- a/python/freetoken/layers/quantization/configs/__init__.py +++ b/python/freetoken/layers/quantization/configs/__init__.py @@ -1,11 +1,12 @@ """One QuantConfig per checkpoint dialect; importing the package registers them all.""" from .base import NoQuantConfig, QuantConfig, quantization_config_of -from . import compressed_tensors, fp8, modelopt, mxfp4 +from . import compressed_tensors, fp8, gguf, modelopt, mxfp4 from .compressed_tensors import CompressedTensorsConfig from .fp8 import Fp8BlockConfig from .modelopt import ModelOptConfig from .mxfp4 import Mxfp4Config +from .gguf import GGUFConfig _NO_QUANT = NoQuantConfig() @@ -33,6 +34,6 @@ def get_quant_config() -> QuantConfig | None: __all__ = [ "QuantConfig", "quantization_config_of", "quant_method_for", "set_quant_config", "get_quant_config", - "NoQuantConfig", "ModelOptConfig", "CompressedTensorsConfig", "Fp8BlockConfig", "Mxfp4Config", - "compressed_tensors", "fp8", "modelopt", "mxfp4", + "NoQuantConfig", "ModelOptConfig", "CompressedTensorsConfig", "Fp8BlockConfig", "Mxfp4Config", "GGUFConfig", + "compressed_tensors", "fp8", "gguf", "modelopt", "mxfp4", ] diff --git a/python/freetoken/layers/quantization/configs/base.py b/python/freetoken/layers/quantization/configs/base.py index e84dcbcbba..3dd63ecc96 100644 --- a/python/freetoken/layers/quantization/configs/base.py +++ b/python/freetoken/layers/quantization/configs/base.py @@ -6,14 +6,15 @@ from dataclasses import dataclass from typing import Any, ClassVar +from ..embedding import EmbeddingConfig from ..linear import LinearConfig from ..moe import MoEConfig from ..names import Matcher, NameMap, name_set from ..quant_backend import get_quant_backend -from ..registry import LayerKind, dialects, method_class, register_dialect +from ..registry import LayerKind, dialects, method_class, methods_for, register_dialect from ..scheme import QuantKind, QuantScheme -LAYER_CONFIGS = {LayerKind.LINEAR: LinearConfig, LayerKind.MOE: MoEConfig} +LAYER_CONFIGS = {LayerKind.LINEAR: LinearConfig, LayerKind.MOE: MoEConfig, LayerKind.EMBEDDING: EmbeddingConfig} def cfg_get(cfg: Any, key: str, default: Any = None) -> Any: @@ -39,10 +40,15 @@ def quantization_config_of(hf_config: Any) -> dict[str, Any] | None: @dataclass(frozen=True) class Stored: - """One checkpoint tensor behind a role: its suffix, and whether it holds the quant-side scale whose reciprocal the layer wants.""" + """One checkpoint tensor behind a role: its suffix, other suffixes exports of the same dialect use for it, and whether it holds the quant-side scale whose reciprocal the layer wants.""" name: str reciprocal: bool = False + aliases: tuple[str, ...] = () + + @property + def names(self) -> tuple[str, ...]: + return (self.name, *self.aliases) class QuantConfig(ABC): @@ -88,10 +94,17 @@ def get_quant_method(self, layer: Any, prefix: str): scheme = self.scheme_for(prefix) layer_kind = layer.quant_layer_kind kind = scheme.kind if scheme else QuantKind.NONE + if layer_kind is LayerKind.EMBEDDING and kind not in methods_for(layer_kind): + # a checkpoint that quantizes its linears stores the vocab table dense + scheme, kind = None, 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/configs/gguf.py b/python/freetoken/layers/quantization/configs/gguf.py new file mode 100644 index 0000000000..77859cd1f7 --- /dev/null +++ b/python/freetoken/layers/quantization/configs/gguf.py @@ -0,0 +1,90 @@ +"""GGUF: every tensor carries its own ggml type, read from the file's tensor table.""" + +from __future__ import annotations + +from typing import Any, Callable, ClassVar + +from ..ggml_types import has_expert_scales, has_weight_scale, is_quantized, type_name +from ..names import NameMap +from ..quant_backend import get_quant_backend +from ..registry import LayerKind, method_class, register_dialect +from ..scheme import QuantKind, QuantScheme, gguf_scheme +from .base import QuantConfig, Stored + + +@register_dialect +class GGUFConfig(QuantConfig): + """``tensor_types`` is the file's name -> ggml type table; ``to_gguf(name, tensor_types)`` (the family's adapter) maps a checkpoint module name to the GGUF tensors that store it.""" + + dialect = "gguf" + STORAGE: ClassVar[dict[QuantKind, dict[str, str | Stored]]] = {QuantKind.GGUF: {"weight": "weight"}} + + def __init__( + self, + tensor_types: dict[str, int], + to_gguf: Callable[[str, dict[str, int]], tuple[str, ...]], + name_map: NameMap | None = None, + unquantized: tuple[str, ...] = (), + ): + super().__init__(name_map, unquantized) + self.tensor_types = tensor_types + self.to_gguf = to_gguf + self.expert_scales = has_expert_scales(tensor_types) + self.weight_scales = any(n.endswith(".scale") and has_weight_scale(tensor_types, n.removesuffix(".scale")) for n in tensor_types) + + @classmethod + def claims(cls, q: dict[str, Any]) -> bool: + return False + + def _part_tensors(self, name: str) -> tuple[str, ...]: + keys = [] + for tensor in self.to_gguf(name, self.tensor_types): + key = f"{tensor}.weight" if f"{tensor}.weight" in self.tensor_types else tensor + if key not in self.tensor_types: + raise KeyError(f"GGUF has no tensor {key!r} for {name!r}") + keys.append(key) + return tuple(keys) + + def _part_types(self, name: str) -> tuple[int, ...]: + return tuple(self.tensor_types[key] for key in self._part_tensors(name)) + + def scheme_for_name(self, name: str) -> QuantScheme | None: + types = self._part_types(name) + if not any(is_quantized(t) for t in types): + return None + if len(types) != 1: + raise NotImplementedError(f"{name} is stored as {len(types)} quantized GGUF tensors") + return gguf_scheme(types) + + def scheme_for(self, prefix: str) -> QuantScheme | None: + if prefix in self._schemes: + return self._schemes[prefix] + names = self.name_map.to_checkpoint(prefix) + tensors = [() if self.unquantized(n) else self._part_tensors(n) for n in names] + parts = [tuple(self.tensor_types[key] for key in ts) for ts in tensors] + quantized = [any(is_quantized(t) for t in p) for p in parts] + if not any(quantized): + scheme = None + elif any(len(p) != 1 for p in parts): + found = {n: [type_name(t) for t in p] for n, p in zip(names, parts)} + raise NotImplementedError(f"fused module {prefix!r} needs one GGUF tensor per part, got {found}") + else: + scaled = any(has_weight_scale(self.tensor_types, ts[0]) for ts in tensors) + scheme = gguf_scheme((p[0] for p in parts), scaled=scaled) + self._schemes[prefix] = scheme + return scheme + + def get_quant_method(self, layer: Any, prefix: str): + if layer.quant_layer_kind is not LayerKind.MOE: + return super().get_quant_method(layer, prefix) + # every expert layer loads its GGUF banks, float experts included; the GGUF MoE kernel picks how to multiply them + parts = [self._part_types(n) for n in self.name_map.to_checkpoint(prefix)] + if any(len(p) != 1 for p in parts): + raise NotImplementedError(f"experts {prefix!r} are stored as {parts} GGUF tensors per projection") + # the scale bank is per file, so every MoE layer keeps one bank set + scheme = gguf_scheme((p[0] for p in parts), scaled=self.expert_scales) + cfg = self.layer_config(layer, LayerKind.MOE, scheme) + return method_class(QuantKind.GGUF, LayerKind.MOE)(cfg, get_quant_backend().select(LayerKind.MOE, QuantKind.GGUF)) + + +__all__ = ["GGUFConfig"] diff --git a/python/freetoken/layers/quantization/configs/modelopt.py b/python/freetoken/layers/quantization/configs/modelopt.py index 6183dfc96b..3547f186d4 100644 --- a/python/freetoken/layers/quantization/configs/modelopt.py +++ b/python/freetoken/layers/quantization/configs/modelopt.py @@ -27,7 +27,8 @@ class ModelOptConfig(QuantConfig): STORAGE: ClassVar[dict[QuantKind, dict[str, str | Stored]]] = { QuantKind.FP8_TENSOR: {"weight": "weight", "weight_scale": "weight_scale", "input_scale": "input_scale"}, QuantKind.FP8_BLOCK: {"weight": "weight", "weight_scale_inv": "weight_scale_inv"}, - QuantKind.MXFP8: {"weight": "weight", "weight_scale_inv": "weight_scale_inv"}, + # same e8m0 codes under either name: MiniMax-M3 exports weight_scale_inv, Qwen3.8-Flash-Next weight_scale + QuantKind.MXFP8: {"weight": "weight", "weight_scale_inv": Stored("weight_scale_inv", aliases=("weight_scale",))}, QuantKind.NVFP4: {"weight": "weight", "weight_scale": "weight_scale", "weight_global": "weight_scale_2", "input_scale": "input_scale"}, } diff --git a/python/freetoken/layers/quantization/embedding/__init__.py b/python/freetoken/layers/quantization/embedding/__init__.py new file mode 100644 index 0000000000..ccdbece04e --- /dev/null +++ b/python/freetoken/layers/quantization/embedding/__init__.py @@ -0,0 +1,10 @@ +from .base import EmbeddingConfig, EmbeddingKernel, EmbeddingMethod +from . import gguf, unquantized +from .gguf import GGUFEmbeddingMethod +from .unquantized import UnquantizedEmbeddingMethod + +__all__ = [ + "EmbeddingConfig", "EmbeddingKernel", "EmbeddingMethod", + "UnquantizedEmbeddingMethod", "GGUFEmbeddingMethod", + "gguf", "unquantized", +] diff --git a/python/freetoken/layers/quantization/embedding/base.py b/python/freetoken/layers/quantization/embedding/base.py new file mode 100644 index 0000000000..caa873f4f9 --- /dev/null +++ b/python/freetoken/layers/quantization/embedding/base.py @@ -0,0 +1,64 @@ +"""The embedding side: per-layer config, the kernel ABC and the Method base. The method owns the vocab table: it looks rows up, and a tied LM head multiplies by the same table through it.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from dataclasses import dataclass +from typing import Any, ClassVar + +import torch + +from ..method import QuantMethod +from ..scheme import QuantScheme + + +@dataclass(frozen=True) +class EmbeddingConfig: + """The TP-local table shape, the dtype a lookup returns, and the table's scheme.""" + + num_embeddings: int + embedding_dim: int + tp_size: int + dtype: torch.dtype + scheme: QuantScheme | None = None + + @classmethod + def from_layer(cls, layer: Any, scheme: QuantScheme | None) -> "EmbeddingConfig": + return cls(layer.num_embeddings_tp, layer.embedding_dim, layer.tp_size, torch.get_default_dtype(), scheme) + + +class EmbeddingKernel(ABC): + """One backend for one table kind: the row lookup, and ``x @ table^T`` for a tied LM head.""" + + name: ClassVar[str] + + def unusable_reason(self, cfg: EmbeddingConfig) -> str | None: + return None + + def worth_it(self, cfg: EmbeddingConfig) -> bool: + return True + + def finalize(self, layer: Any) -> None: + pass + + @abstractmethod + def embed(self, layer: Any, ids: torch.Tensor) -> torch.Tensor: ... + + @abstractmethod + def logits(self, layer: Any, x: torch.Tensor, bias: torch.Tensor | None) -> torch.Tensor: ... + + +class EmbeddingMethod(QuantMethod): + """Declares the table for its kind; the lookup and the tied head's product go to the kernel.""" + + @abstractmethod + def create_weights(self, layer: Any) -> None: ... + + def finalize(self, layer: Any) -> None: + self.kernel.finalize(layer) + + def embed(self, layer: Any, ids: torch.Tensor) -> torch.Tensor: + return self.kernel.embed(layer, ids) + + def logits(self, layer: Any, x: torch.Tensor, bias: torch.Tensor | None = None) -> torch.Tensor: + return self.kernel.logits(layer, x, bias) diff --git a/python/freetoken/layers/quantization/embedding/gguf.py b/python/freetoken/layers/quantization/embedding/gguf.py new file mode 100644 index 0000000000..98f12dd9cc --- /dev/null +++ b/python/freetoken/layers/quantization/embedding/gguf.py @@ -0,0 +1,48 @@ +"""GGUF tables: packed ggml rows, decoded per looked-up row; a tied head runs the GGUF linear kernel on the same rows.""" + +from __future__ import annotations + +from types import SimpleNamespace +from typing import Any + +import torch + +from ..ggml_types import NVFP4, Segment, pad_rows, row_bytes +from ..linear.gguf import GgmlLinearKernel +from ..registry import LayerKind, register_method +from ..scheme import QuantKind, gguf_types +from .base import EmbeddingConfig, EmbeddingKernel, EmbeddingMethod + + +class GgmlEmbeddingKernel(EmbeddingKernel): + name = "ggml" + + def unusable_reason(self, cfg: EmbeddingConfig) -> str | None: + if cfg.tp_size > 1: + return "GGUF tables are not sharded over TP" + # the vendored get_rows has no NVFP4 case, and the table's weight_scale_2 would need applying too + return "NVFP4 GGUF tables are not supported" if NVFP4 in gguf_types(cfg.scheme) else None + + def embed(self, layer: Any, ids: torch.Tensor) -> torch.Tensor: + from freetoken.kernel import ggml + + cfg = layer.quant_method.cfg + (ggml_type,) = gguf_types(cfg.scheme) + rows = ggml.get_rows(layer.weight, ggml_type, cfg.embedding_dim, ids.reshape(-1).to(torch.int32), cfg.dtype) + return rows.view(*ids.shape, cfg.embedding_dim) + + def logits(self, layer: Any, x: torch.Tensor, bias: torch.Tensor | None) -> torch.Tensor: + cfg = layer.quant_method.cfg + (ggml_type,) = gguf_types(cfg.scheme) + table = SimpleNamespace(_ggml_segments=[Segment(ggml_type, 0, 1, 0, cfg.num_embeddings)], weight=layer.weight, bias=bias) + return GgmlLinearKernel().apply(table, x) + + +@register_method(QuantKind.GGUF, LayerKind.EMBEDDING) +class GGUFEmbeddingMethod(EmbeddingMethod): + candidates = (GgmlEmbeddingKernel,) + + def create_weights(self, layer: Any) -> None: + (ggml_type,) = gguf_types(self.cfg.scheme) + h = self.cfg.embedding_dim + layer.weight = torch.empty(self.cfg.num_embeddings + pad_rows(ggml_type, h), row_bytes(ggml_type, h), dtype=torch.uint8) diff --git a/python/freetoken/layers/quantization/embedding/unquantized.py b/python/freetoken/layers/quantization/embedding/unquantized.py new file mode 100644 index 0000000000..5b755baf19 --- /dev/null +++ b/python/freetoken/layers/quantization/embedding/unquantized.py @@ -0,0 +1,33 @@ +"""Dense tables: the indexing kernel looks rows up; a tied head runs the unquantized linear kernel on the same table.""" + +from __future__ import annotations + +from types import SimpleNamespace +from typing import Any + +import torch + +from ..linear.unquantized import TorchLinearKernel +from ..registry import LayerKind, register_method +from ..scheme import QuantKind +from .base import EmbeddingKernel, EmbeddingMethod + + +class IndexingEmbeddingKernel(EmbeddingKernel): + name = "indexing" + + def embed(self, layer: Any, ids: torch.Tensor) -> torch.Tensor: + from freetoken.kernel import indexing + + return indexing(weights=layer.weight, indices=ids, vocab_range=layer.vocab_range if layer.tp_size > 1 else None) + + def logits(self, layer: Any, x: torch.Tensor, bias: torch.Tensor | None) -> torch.Tensor: + return TorchLinearKernel().apply(SimpleNamespace(weight=layer.weight, bias=bias), x) + + +@register_method(QuantKind.NONE, LayerKind.EMBEDDING) +class UnquantizedEmbeddingMethod(EmbeddingMethod): + candidates = (IndexingEmbeddingKernel,) + + def create_weights(self, layer: Any) -> None: + layer.weight = torch.empty(self.cfg.num_embeddings, self.cfg.embedding_dim) diff --git a/python/freetoken/layers/quantization/ggml_types.py b/python/freetoken/layers/quantization/ggml_types.py new file mode 100644 index 0000000000..95c2779311 --- /dev/null +++ b/python/freetoken/layers/quantization/ggml_types.py @@ -0,0 +1,119 @@ +"""ggml type facts shared by the GGUF dialect, its methods and the GGUF weight readers.""" + +from __future__ import annotations + +import re +from dataclasses import dataclass + +import torch + +from freetoken.third_party import gguf + +# ggml reads up to this many values past the end of a packed matrix (llama.cpp MATRIX_ROW_PADDING) +MATRIX_ROW_PADDING = 512 + +UNQUANTIZED = frozenset( + int(t) for t in (gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16, gguf.GGMLQuantizationType.BF16) +) + + +_FLOAT_DTYPES = { + int(gguf.GGMLQuantizationType.F32): torch.float32, + int(gguf.GGMLQuantizationType.F16): torch.float16, + int(gguf.GGMLQuantizationType.BF16): torch.bfloat16, +} + + +def float_dtype(ggml_type: int) -> torch.dtype | None: + """The torch dtype of an unquantized ggml type, None for a quantized one.""" + return _FLOAT_DTYPES.get(int(ggml_type)) + + +def type_name(ggml_type: int) -> str: + return gguf.GGMLQuantizationType(ggml_type).name + + +def is_quantized(ggml_type: int) -> bool: + return int(ggml_type) not in UNQUANTIZED + + +def row_bytes(ggml_type: int, k: int) -> int: + block, size = gguf.GGML_QUANT_SIZES[gguf.GGMLQuantizationType(ggml_type)] + if k % block: + raise ValueError(f"{type_name(ggml_type)} rows need a multiple of {block} values, got {k}") + return k // block * size + + +def pad_rows(ggml_type: int, k: int) -> int: + """Zero rows to append after a packed [rows, K] matrix so ggml's tail reads stay inside it.""" + tail = mmq_tail_bytes(ggml_type, k) + return 0 if tail == 0 else max(1, -(-tail // row_bytes(ggml_type, k))) + + +def mmq_tail_bytes(ggml_type: int, k: int) -> int: + """Bytes past the last row of a packed [rows, K] matrix that MMQ reads and that must read as zeros.""" + rem = k % MATRIX_ROW_PADDING + return 0 if rem == 0 or not is_quantized(ggml_type) else row_bytes(ggml_type, MATRIX_ROW_PADDING - rem) + + +@dataclass(frozen=True) +class Segment: + """A run of consecutive fused parts that share one ggml type: one packed matrix and one GEMM.""" + + ggml_type: int + first_part: int + num_parts: int + out_start: int + rows: int + + +def segments(types: tuple[int, ...], output_sizes: tuple[int, ...]) -> list[Segment]: + if len(types) != len(output_sizes): + raise ValueError(f"{len(types)} ggml types for {len(output_sizes)} fused parts") + out: list[Segment] = [] + start = 0 + for i, (t, n) in enumerate(zip(types, output_sizes)): + if out and out[-1].ggml_type == t: + last = out[-1] + out[-1] = Segment(t, last.first_part, last.num_parts + 1, last.out_start, last.rows + n) + else: + out.append(Segment(int(t), i, 1, start, n)) + start += n + return out + + +def segment_param(index: int) -> str: + """State-dict leaf of segment ``index``: ``weight`` for the first, then ``weight_1``, ``weight_2``.""" + return "weight" if index == 0 else f"weight_{index}" + + +NVFP4 = int(gguf.GGMLQuantizationType.NVFP4) + +# per-expert fp32 weight scales llama.cpp reads next to separate gate / up / down expert tensors (llama-model.cpp ffn_*_exps_s) +_EXPERT_SCALE_RE = re.compile(r"^blk\.\d+\.ffn_(gate|up|down)_exps\.scale$") + + +def scale_name(tensor: str) -> str: + """``.scale`` of the GGUF tensor ```` or ``.weight``.""" + return f"{tensor.removesuffix('.weight')}.scale" + + +def has_weight_scale(tensor_types: dict[str, int], tensor: str) -> bool: + """Whether llama.cpp's converter stored an NVFP4 tensor's weight_scale_2 as ``.scale``; other types keep no weight scale.""" + weight = tensor if tensor.endswith(".weight") else f"{tensor}.weight" + return tensor_types.get(weight, tensor_types.get(tensor)) == NVFP4 and scale_name(tensor) in tensor_types + + +def has_expert_scales(tensor_types: dict[str, int]) -> bool: + """Whether any NVFP4 expert tensor of the file carries per-expert scales, which then take a scale bank in every MoE layer.""" + merged = [n for n in tensor_types if n.endswith("ffn_gate_up_exps.scale") and tensor_types.get(n.replace(".scale", ".weight")) == NVFP4] + if merged: + # llama.cpp scales only separate gate / up expert tensors and would drop this one + raise NotImplementedError(f"per-expert scales of merged gate_up NVFP4 experts are not supported (e.g. {merged[0]})") + return any(_EXPERT_SCALE_RE.match(n) and has_weight_scale(tensor_types, n.removesuffix(".scale")) for n in tensor_types) + + +__all__ = [ + "MATRIX_ROW_PADDING", "NVFP4", "Segment", "float_dtype", "has_expert_scales", "has_weight_scale", "is_quantized", "mmq_tail_bytes", + "pad_rows", "row_bytes", "scale_name", "segment_param", "segments", "type_name", +] diff --git a/python/freetoken/layers/quantization/linear/__init__.py b/python/freetoken/layers/quantization/linear/__init__.py index 8236abe018..781f2b458a 100644 --- a/python/freetoken/layers/quantization/linear/__init__.py +++ b/python/freetoken/layers/quantization/linear/__init__.py @@ -1,13 +1,14 @@ from .base import LinearConfig, LinearKernel, LinearMethod -from . import fp8_block, fp8_tensor, mxfp8, nvfp4, unquantized +from . import fp8_block, fp8_tensor, gguf, mxfp8, nvfp4, unquantized from .fp8_block import Fp8BlockLinearMethod from .fp8_tensor import Fp8TensorLinearMethod from .mxfp8 import Mxfp8LinearMethod from .nvfp4 import Nvfp4LinearMethod from .unquantized import UnquantizedLinearMethod +from .gguf import GGUFLinearMethod __all__ = [ "LinearConfig", "LinearKernel", "LinearMethod", - "UnquantizedLinearMethod", "Fp8TensorLinearMethod", "Fp8BlockLinearMethod", "Mxfp8LinearMethod", "Nvfp4LinearMethod", - "fp8_block", "fp8_tensor", "mxfp8", "nvfp4", "unquantized", + "UnquantizedLinearMethod", "Fp8TensorLinearMethod", "Fp8BlockLinearMethod", "Mxfp8LinearMethod", "Nvfp4LinearMethod", "GGUFLinearMethod", + "fp8_block", "fp8_tensor", "gguf", "mxfp8", "nvfp4", "unquantized", ] 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/linear/gguf.py b/python/freetoken/layers/quantization/linear/gguf.py new file mode 100644 index 0000000000..badbe92b49 --- /dev/null +++ b/python/freetoken/layers/quantization/linear/gguf.py @@ -0,0 +1,62 @@ +"""GGUF Linear: packed ggml rows multiplied by llama.cpp's MMVQ / MMQ kernels; a fused module runs one GEMM per ggml type, float parts as plain bf16 rows.""" + +from __future__ import annotations + +from typing import Any + +import torch + +from ..ggml_types import is_quantized, pad_rows, row_bytes, segment_param, segments +from ..registry import LayerKind, register_method +from ..scheme import QuantKind, gguf_types +from .base import LinearConfig, LinearKernel, LinearMethod + + + +class GgmlLinearKernel(LinearKernel): + name = "ggml" + + def unusable_reason(self, cfg: LinearConfig) -> str | None: + return None + + def apply(self, layer: Any, x: torch.Tensor) -> torch.Tensor: + from freetoken.kernel import ggml + + x2 = x.reshape(-1, x.shape[-1]) + segs = layer._ggml_segments + m = x2.shape[0] + y = torch.empty(m, sum(s.rows for s in segs), dtype=x.dtype, device=x.device) + # every segment reads the same x: quantize it once for all the vector-kernel segments + xq = ggml.quantize_q8_1(x2) if sum(is_quantized(s.ggml_type) and ggml.uses_mmvq(s.ggml_type, m) for s in segs) > 1 else None + for i, s in enumerate(segs): + w = getattr(layer, segment_param(i))[: s.rows] + out = y[:, s.out_start : s.out_start + s.rows] + if is_quantized(s.ggml_type): + ggml.mul_mat(w, s.ggml_type, x2, out=out, xq=xq) + else: + torch.mm(x2, w.t(), out=out) + # llama.cpp multiplies an NVFP4 result by the tensor's weight_scale_2 (build_lora_mm) + if getattr(layer, "ggml_scale", None) is not None: + y.mul_(layer.ggml_scale) + if layer.bias is not None: + y = y + layer.bias + return y.reshape(*x.shape[:-1], y.shape[-1]) + + +@register_method(QuantKind.GGUF, LayerKind.LINEAR) +class GGUFLinearMethod(LinearMethod): + candidates = (GgmlLinearKernel,) + + def create_weights(self, layer: Any) -> None: + g = self.cfg + segs = segments(gguf_types(g.scheme), g.output_sizes) + layer._ggml_segments = segs + k = g.in_features + for i, s in enumerate(segs): + if is_quantized(s.ggml_type): + setattr(layer, segment_param(i), torch.empty(s.rows + pad_rows(s.ggml_type, k), row_bytes(s.ggml_type, k), dtype=torch.uint8)) + else: + setattr(layer, segment_param(i), torch.empty(s.rows, k)) + if g.scheme.weight.scale is not None: + # one value per output row: each part's weight_scale_2, 1 for parts without one + layer.ggml_scale = torch.empty(sum(s.rows for s in segs), dtype=torch.float32) diff --git a/python/freetoken/layers/quantization/linear/nvfp4.py b/python/freetoken/layers/quantization/linear/nvfp4.py index 0953b82501..1002a70b3c 100644 --- a/python/freetoken/layers/quantization/linear/nvfp4.py +++ b/python/freetoken/layers/quantization/linear/nvfp4.py @@ -1,19 +1,25 @@ -"""e2m1 weight packed two per byte, e4m3 scales per 16 input elements, one fp16 global per output row.""" +"""e2m1 weight packed two per byte, e4m3 scales per 16 input elements, one fp32 global per output row.""" from __future__ import annotations +import functools from types import SimpleNamespace from typing import Any import torch from freetoken.kernel import backend +from freetoken.utils import init_logger from ..registry import LayerKind, register_method from ..scheme import NVFP4_GROUP as GROUP, QuantKind from .base import LinearConfig, LinearKernel, LinearMethod +logger = init_logger(__name__) + FP8 = torch.float8_e4m3fn +# fused parts whose weight globals are further apart than this ratio get a load-time warning +_PART_GLOBAL_WARN_RATIO = 1.05 class TritonNvfp4LinearKernel(LinearKernel): @@ -32,12 +38,188 @@ def apply(self, layer: Any, x: torch.Tensor) -> torch.Tensor: return nvfp4_dense_linear_t(x, layer.weight, layer.weight_scale, layer.weight_global, layer.bias) +@functools.cache +def _warn_part_globals() -> None: + logger.warning_rank0( + "an NVFP4 linear fuses parts whose weight global scales differ by more than %d%%; the W4A4 kernels use the largest " + "for every part, which scales the other parts' outputs by their ratio; --quant-backend linear.nvfp4=triton keeps " + "each part's global", round((_PART_GLOBAL_WARN_RATIO - 1) * 100), + ) + + +def _tensor_global(layer: Any) -> torch.Tensor: + """The one weight global the W4A4 GEMMs and the small-batch kernels take: the largest per-part global, the block scales unchanged.""" + g = layer.weight_global.float() + g_max = g.max() + if g_max > _PART_GLOBAL_WARN_RATIO * g.min(): + _warn_part_globals() + return g_max + + +def _swizzle_128x4(scale: torch.Tensor) -> torch.Tensor: + """[N, K/16] block scales -> the flat cuBLAS / CUTLASS 128x4 layout, rows padded to 128 and columns to 4 with zeros.""" + rows, cols = scale.shape + rb, cb = -(-rows // 128), -(-cols // 4) + padded = torch.zeros(rb * 128, cb * 4, dtype=scale.dtype, device=scale.device) + padded[:rows, :cols] = scale + blocks = padded.view(rb, 128, cb, 4).permute(0, 2, 1, 3) + return blocks.reshape(-1, 4, 32, 4).transpose(1, 2).reshape(-1).contiguous() + + +@functools.cache +def _sm120_gemv(device_index: int): + """The row-major W4A16 GEMV / skinny GEMM module for this GPU, or None.""" + from freetoken.kernel import gluon + + if torch.cuda.get_device_capability(device_index) != (12, 0) or not gluon.available(): + return None + from freetoken.kernel.gluon.nvfp4.linear import sm120_gemv + + return sm120_gemv + + +def _small_batch(layer: Any, x: torch.Tensor, max_tokens: int) -> torch.Tensor | None: + """W4A16 on the W4A4 kernels' own weight for decode-sized batches, where the W4A4 GEMMs leave bandwidth unused; None to run the GEMM.""" + gemv = _sm120_gemv(x.device.index) + if gemv is None or x.dtype != torch.bfloat16 or x.shape[0] == 0: + return None + route = gemv.gemv_config(layer.weight.shape[0], x.shape[1], x.shape[0], max_tokens) + if route is None: + return None + kernel, tile = route + return kernel(x if x.stride(1) == 1 else x.contiguous(), layer.weight, layer.weight_scale, layer.weight_global, **tile) + + +@functools.cache +def _flashinfer_cuda_major() -> int: + """The CUDA version flashinfer's b12x GEMM checks at call time: nvcc's, else torch's.""" + from flashinfer.jit.cpp_ext import get_cuda_version + + return get_cuda_version().major + + +def _small_batch_usable() -> bool: + return torch.cuda.is_available() and _sm120_gemv(torch.cuda.current_device()) is not None + + +class CublasNvfp4LinearKernel(LinearKernel): + """W4A4: activations quantized under the checkpoint's input_scale, then torch's NVFP4 scaled_mm (cuBLASLt) on the row-major weight.""" + + name = "cublas" + + def unusable_reason(self, cfg: LinearConfig) -> str | None: + if backend.is_rocm() or not torch.cuda.is_available(): + return "needs CUDA" + if cfg.scheme is None or not cfg.scheme.has("input_scale"): + return "the weight declares no static activation scale (W4A16 checkpoint)" + if backend.device_capability() < (10, 0): + return "NVFP4 tensor cores need sm_100+" + if not hasattr(torch.nn.functional, "scaled_mm") or not hasattr(torch, "float4_e2m1fn_x2"): + return f"torch {torch.__version__} has no NVFP4 scaled_mm" + if cfg.out_features % 16 or cfg.in_features % 32: + return f"scaled_mm needs N % 16 == 0 and K % 32 == 0, got N={cfg.out_features} K={cfg.in_features}" + return None + + def worth_it(self, cfg: LinearConfig) -> bool: + # cuBLAS alone runs decode-sized batches up to 1.4x slower than the triton GEMV; the small-batch kernels cover those on sm_120 + return _small_batch_usable() + + def finalize(self, layer: Any) -> None: + g_w = _tensor_global(layer) + layer.weight_scale = _swizzle_128x4(layer.weight_scale) + layer.weight_global = g_w.reshape(()) + layer.input_scale = layer.input_scale.float().reshape(()) + + def apply(self, layer: Any, x: torch.Tensor) -> torch.Tensor: + from torch.nn import functional as F + + from freetoken.kernel.triton.nvfp4_act_quant import quantize_nvfp4 + + *lead, K = x.shape + out = _small_batch(layer, x.reshape(-1, K), 16) + if out is not None: + out = out.reshape(*lead, layer.weight.shape[0]) + return out + layer.bias.to(out.dtype) if layer.bias is not None else out + xq, xs = quantize_nvfp4(x.reshape(-1, K), layer.input_scale.reshape(1), swizzled=True) + block, tensor = F.ScalingType.BlockWise1x16, F.ScalingType.TensorWise + swizzle = [F.SwizzleType.SWIZZLE_32_4_4, F.SwizzleType.NO_SWIZZLE] + out = F.scaled_mm( + xq.view(torch.float4_e2m1fn_x2), layer.weight.view(torch.float4_e2m1fn_x2).t(), + [xs, layer.input_scale], [block, tensor], [layer.weight_scale, layer.weight_global], [block, tensor], + swizzle_a=swizzle, swizzle_b=swizzle, output_dtype=x.dtype, + ).reshape(*lead, layer.weight.shape[0]) + return out + layer.bias.to(out.dtype) if layer.bias is not None else out + + +class B12xNvfp4LinearKernel(LinearKernel): + """W4A4 through flashinfer alone: its activation quant and its SM12x CuTe-DSL GEMM (mm_fp4 backend b12x) on the row-major weight.""" + + name = "b12x" + + def unusable_reason(self, cfg: LinearConfig) -> str | None: + if backend.is_rocm() or not torch.cuda.is_available(): + return "needs CUDA" + if cfg.scheme is None or not cfg.scheme.has("input_scale"): + return "the weight declares no static activation scale (W4A16 checkpoint)" + cc = backend.device_capability() + if cc[0] != 12: + return f"b12x requires sm_120 / sm_121, got sm_{cc[0]}{cc[1]}" + if not backend.is_flashinfer_installed(): + return "flashinfer is not installed" + drv = backend.driver_cuda_version() + if drv is not None and drv < 13000: + return f"b12x JIT-compiles through the driver and needs CUDA 13+, the driver supports {drv // 1000}.{drv % 1000 // 10}" + if cfg.in_features % 32 or cfg.out_features % 8: + # flashinfer checks only K; N % 8 != 0 returns wrong output + return f"b12x needs K % 32 == 0 and N % 8 == 0, got K={cfg.in_features} N={cfg.out_features}" + try: + from flashinfer import fp4_quantize, mm_fp4, nvfp4_block_scale_interleave # noqa: F401 + + cuda_major = _flashinfer_cuda_major() + except Exception as exc: + return f"flashinfer FP4 GEMM symbols are unusable ({exc!r})" + if cuda_major < 13: + return f"flashinfer's b12x GEMM needs a CUDA 13 toolkit (nvcc, else torch), got CUDA {cuda_major}" + return None + + def worth_it(self, cfg: LinearConfig) -> bool: + return _small_batch_usable() + + def finalize(self, layer: Any) -> None: + from flashinfer import nvfp4_block_scale_interleave + + g_w = _tensor_global(layer) + layer.weight_scale = nvfp4_block_scale_interleave(layer.weight_scale.view(torch.uint8)).view(FP8) + # flashinfer takes the quant-side activation global and alpha = 1 / (that * the weight's quant-side global) + layer.input_scale_inv = (1.0 / layer.input_scale.float()).reshape(1) + layer.alpha = (layer.input_scale.float() * g_w).reshape(1) + layer.weight_global = g_w.reshape(()) + + def apply(self, layer: Any, x: torch.Tensor) -> torch.Tensor: + from flashinfer import fp4_quantize, mm_fp4 + + *lead, K = x.shape + x2 = x.reshape(-1, K) + if x2.shape[0] == 0: + return x.new_empty(*lead, layer.weight.shape[0]) + # b12x's W4A4 GEMM beats the W4A16 kernels from 9 tokens on GeForce, where bf16 mma accumulates in fp32 at half rate + out = _small_batch(layer, x2, 8) + if out is not None: + out = out.reshape(*lead, layer.weight.shape[0]) + return out + layer.bias.to(out.dtype) if layer.bias is not None else out + xq, xs = fp4_quantize(x2.contiguous(), layer.input_scale_inv, is_sf_swizzled_layout=True) + out = mm_fp4(xq, layer.weight.T, xs, layer.weight_scale, layer.alpha, x.dtype, backend="b12x").reshape(*lead, layer.weight.shape[0]) + return out + layer.bias.to(out.dtype) if layer.bias is not None else out + + class MarlinNvfp4LinearKernel(LinearKernel): """vLLM's fused FP4 Marlin GEMM (W4A16); needs one global scale for the whole weight.""" name = "marlin" def unusable_reason(self, cfg: LinearConfig) -> str | None: + if backend.is_rocm(): + return "Marlin is CUDA-only and unavailable on ROCm; use triton" if not backend.is_vllm_installed(): return "vLLM is not installed" if len(cfg.output_sizes) > 1: @@ -74,7 +256,7 @@ def finalize(self, layer: Any) -> None: dev = layer.weight.device layer.weight = torch.empty(0, dtype=torch.uint8, device=dev) layer.weight_scale = torch.empty(0, dtype=FP8, device=dev) - layer.weight_global = torch.empty(0, dtype=torch.float16, device=dev) + layer.weight_global = torch.empty(0, dtype=torch.float32, device=dev) def apply(self, layer: Any, x: torch.Tensor) -> torch.Tensor: from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import apply_fp4_marlin_linear @@ -111,7 +293,7 @@ def apply(self, layer: Any, x: torch.Tensor) -> torch.Tensor: @register_method(QuantKind.NVFP4, LayerKind.LINEAR) class Nvfp4LinearMethod(LinearMethod): - candidates = (TritonNvfp4LinearKernel, MarlinNvfp4LinearKernel, EmulationNvfp4LinearKernel) + candidates = (B12xNvfp4LinearKernel, CublasNvfp4LinearKernel, TritonNvfp4LinearKernel, MarlinNvfp4LinearKernel, EmulationNvfp4LinearKernel) def create_weights(self, layer: Any) -> None: g = self.cfg @@ -119,6 +301,6 @@ def create_weights(self, layer: Any) -> None: raise ValueError(f"nvfp4 needs in_features divisible by {GROUP}, got {g.in_features}") layer.weight = torch.empty(g.out_features, g.in_features // 2, dtype=torch.uint8) layer.weight_scale = torch.empty(g.out_features, g.in_features // GROUP, dtype=FP8) - layer.weight_global = torch.empty(g.out_features, dtype=torch.float16) + layer.weight_global = torch.empty(g.out_features, dtype=torch.float32) # no W4A16 kernel reads it; declared so a W4A4 checkpoint loads complete and a W4A4 kernel finds it in place layer.input_scale = torch.empty((), dtype=torch.float32) if self.scheme.has("input_scale") else None diff --git a/python/freetoken/layers/quantization/moe/__init__.py b/python/freetoken/layers/quantization/moe/__init__.py index a5f2dc6645..1230f36f08 100644 --- a/python/freetoken/layers/quantization/moe/__init__.py +++ b/python/freetoken/layers/quantization/moe/__init__.py @@ -1,13 +1,14 @@ from .base import BankSpec, ExpertView, MoEConfig, MoEKernel, MoEMethod -from . import fp8_block, mxfp4, mxfp8, nvfp4, unquantized +from . import fp8_block, gguf, mxfp4, mxfp8, nvfp4, unquantized from .fp8_block import Fp8BlockMoEMethod from .mxfp4 import Mxfp4MoEMethod from .mxfp8 import Mxfp8MoEMethod from .nvfp4 import Nvfp4MoEMethod from .unquantized import UnquantizedMoEMethod +from .gguf import GGUFMoEMethod __all__ = [ "BankSpec", "ExpertView", "MoEConfig", "MoEKernel", "MoEMethod", - "UnquantizedMoEMethod", "Fp8BlockMoEMethod", "Nvfp4MoEMethod", "Mxfp4MoEMethod", "Mxfp8MoEMethod", - "fp8_block", "mxfp4", "mxfp8", "nvfp4", "unquantized", + "UnquantizedMoEMethod", "Fp8BlockMoEMethod", "Nvfp4MoEMethod", "Mxfp4MoEMethod", "Mxfp8MoEMethod", "GGUFMoEMethod", + "fp8_block", "gguf", "mxfp4", "mxfp8", "nvfp4", "unquantized", ] diff --git a/python/freetoken/layers/quantization/moe/base.py b/python/freetoken/layers/quantization/moe/base.py index 89f13ed09d..b17a91359b 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 @@ -92,8 +96,8 @@ def fused_piece(pieces: dict[str, torch.Tensor], role: str) -> torch.Tensor: def global_rows(piece: torch.Tensor, rows: int) -> torch.Tensor: - """Per-expert global scale as one fp16 value per output row (a scalar per expert is broadcast).""" - flat = piece.reshape(piece.shape[0], -1).to(torch.float16) + """Per-expert global scale as one fp32 value per output row (a scalar per expert is broadcast).""" + flat = piece.reshape(piece.shape[0], -1).to(torch.float32) if flat.shape[1] == rows: return flat assert flat.shape[1] == 1, (flat.shape, rows) @@ -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 @@ -143,13 +151,15 @@ def unusable_reason(self, cfg: MoEConfig) -> str | None: def worth_it(self, cfg: MoEConfig) -> bool: return True + def cpu_options(self, cfg: MoEConfig) -> dict[str, int]: + """Integer options of ``cpu_format`` for these experts.""" + return {} + 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 +200,19 @@ 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 cpu_spec(self) -> tuple[str | None, dict[str, int]]: + """The CPU executor format and options of this layer's experts.""" + return self.cpu_format, self.kernel.cpu_options(self.cfg) + + 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..6cda6b1ecd 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 @@ -19,11 +19,14 @@ def _pad_scale(rows: int, cols: int) -> int: class TritonFp8BlockMoEKernel(MoEKernel): name = "triton" + cpu_format = "fp8_block" 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=True, 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 +67,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/gguf.py b/python/freetoken/layers/quantization/moe/gguf.py new file mode 100644 index 0000000000..01e754d644 --- /dev/null +++ b/python/freetoken/layers/quantization/moe/gguf.py @@ -0,0 +1,168 @@ +"""GGUF experts: packed ggml rows in two banks (gate|up and down), multiplied by llama.cpp's mul_mat_id kernels; float experts run on the bf16 fused MoE.""" + +from __future__ import annotations + +import math + +import torch + +from ..ggml_types import NVFP4, float_dtype, is_quantized, row_bytes, type_name +from ..registry import LayerKind, register_method +from ..scheme import QuantKind, gguf_types +from .base import BankSpec, ExpertView, MoEConfig, MoEKernel, MoEMethod, gated_epilogue_reason, is_resident, limit_or_inf + + + +def fused_glu(layer) -> tuple[str, float] | None: + """The MMVQ epilogue (``ggml.mul_mat_id_glu``) that computes this layer's gated activation, or None to run it as a separate kernel.""" + act, alpha, limit = layer.activation, float(layer.alpha), limit_or_inf(layer) + if act == "silu" and alpha == 1.0 and limit == math.inf: + return "swiglu", 0.0 + if act == "gelu_tanh": + return "geglu", 0.0 + if act == "swiglu_clamp" and alpha == 1.0 and limit != math.inf: + return "swiglu_clamp", limit + # ggml's fused swiglu_oai keeps its default alpha and limit + if act == "swigluoai" and math.isclose(alpha, 1.702, rel_tol=1e-6) and math.isclose(limit, 7.0, rel_tol=1e-6): + return "swiglu_oai", 0.0 + return None + + +def expert_types(cfg: MoEConfig) -> tuple[int, int]: + gate, up, down = gguf_types(cfg.scheme) + if gate != up: + raise NotImplementedError(f"GGUF experts with gate {type_name(gate)} and up {type_name(up)} need one gate_up bank type") + return gate, down + + +def _bank_spec(ggml_type: int, rows: int, k: int) -> BankSpec: + dtype = float_dtype(ggml_type) + return BankSpec((rows, k), dtype) if dtype is not None else BankSpec((rows, row_bytes(ggml_type, k)), torch.uint8) + + +_DECODE_CONFIG = {"BLOCK_SIZE_N": 32, "BLOCK_SIZE_K": 64, "num_warps": 8} + + +def float_mul_mat_id(w: torch.Tensor, x: torch.Tensor, ids: torch.Tensor, topk_weights: torch.Tensor, config: dict | None) -> torch.Tensor: + """``ggml.mul_mat_id`` for float experts ``w[E, N, K]``: same ``x`` and ``ids``, same ``[T, k, N]`` output in ``x``'s dtype, no router weights. + + ``config`` runs the sorted-route kernel (prefill, resident layers); None runs the per-route decode kernel.""" + from freetoken.kernel import fused_moe_decode_kernel_triton, fused_moe_kernel_triton + from freetoken.moe.fused import moe_align_block_size + + t, top_k = ids.shape + per_route = x.dim() == 3 + # f16 experts take f16 activations, as llama.cpp's f16 path does + a = x.reshape(-1, x.shape[-1]).to(w.dtype).contiguous() + out = torch.empty(t, top_k, w.shape[1], dtype=x.dtype, device=x.device) + if config is None: + fused_moe_decode_kernel_triton(a, w, out, topk_weights, ids, False, per_route, _DECODE_CONFIG, compute_type=x.dtype) + else: + sorted_ids, expert_ids, num_post_padded = moe_align_block_size(ids, config["BLOCK_SIZE_M"], w.shape[0]) + fused_moe_kernel_triton( + a, w, out, topk_weights, ids, sorted_ids, expert_ids, num_post_padded, False, 1 if per_route else top_k, config, compute_type=x.dtype + ) + return out + + + +class GgmlMoEKernel(MoEKernel): + name = "ggml" + cpu_format = "gguf" + + def unusable_reason(self, cfg: MoEConfig) -> str | None: + if not is_quantized(expert_types(cfg)[0]): + reason = gated_epilogue_reason(cfg) + if reason: + return f"float GGUF experts run on the bf16 fused MoE: {reason}" + return self._common_reject(cfg, tp_ok=False, cpu_ok=True, plain_silu_only=False) + + def layout(self, cfg: MoEConfig) -> dict[str, BankSpec]: + gate_up, down = expert_types(cfg) + i, h = cfg.local_intermediate, cfg.hidden + banks = {"gate_up": _bank_spec(gate_up, 2 * i, h), "down": _bank_spec(down, h, i)} + if cfg.scheme.weight.scale is not None: + # each expert's NVFP4 weight_scale_2 as (gate, up, down, 1): a 16-byte row, which the slot cache's copies need + banks["global"] = BankSpec((4,), torch.float32) + return banks + + def pack(self, pieces, cfg: MoEConfig, out): + i = cfg.local_intermediate + out["gate_up"][:, :i].copy_(pieces["gate"]) + out["gate_up"][:, i:].copy_(pieces["up"]) + out["down"].copy_(pieces["down"]) + return {} + + def apply(self, layer, x, topk_weights, topk_ids, view: ExpertView, *, is_prefill: bool): + from freetoken.kernel import ggml + from freetoken.layers.activation import gated_act_and_mul + + cfg = layer.quant_method.cfg + gate_up_type, down_type = expert_types(cfg) + if not is_quantized(gate_up_type) and not is_quantized(down_type): + return self._apply_float(layer, x, topk_weights, topk_ids, view, gate_up_type, is_prefill=is_prefill) + assert not layer.apply_router_weight_on_input, "GGUF experts apply the router weights on the output" + ids = topk_ids.to(torch.int32) + gate_up = view.tensors["gate_up"] + # a float half runs on the bf16 fused MoE's grouped GEMM, with its config for this layer's shape + config = None + if not (is_quantized(gate_up_type) and is_quantized(down_type)) and (is_prefill or is_resident(layer)): + from freetoken.moe.fused import try_get_optimal_moe_config + + e, inter_size, hidden = gate_up.shape[0], cfg.local_intermediate, cfg.hidden + config = try_get_optimal_moe_config((e, 2 * inter_size, hidden), (e, hidden, inter_size), ids.shape[1], x.shape[0]) + + def mul_mat_id(w: torch.Tensor, ggml_type: int, a: torch.Tensor) -> torch.Tensor: + if is_quantized(ggml_type): + return ggml.mul_mat_id(w, ggml_type, a, ids) + return float_mul_mat_id(w, a, ids, topk_weights, config) + + # llama.cpp's build_moe_ffn: each NVFP4 GEMM result times the expert's weight_scale_2, gate and up before the glu, down before the routing weight + scales = view.tensors.get("global") + gate_up_scales = scales if scales is not None and gate_up_type == NVFP4 else None + down_scales = scales if scales is not None and down_type == NVFP4 else None + glu = fused_glu(layer) + small = x.shape[0] <= min(ggml.mmvq_max_tokens(t) for t in (gate_up_type, down_type) if is_quantized(t)) + if glu is not None and small and is_quantized(gate_up_type): + gu = (None, None) if gate_up_scales is None else (gate_up_scales[:, 0], gate_up_scales[:, 1]) + inter = ggml.mul_mat_id_glu(gate_up, gate_up_type, x, ids, *glu, *gu) + else: + h = mul_mat_id(gate_up, gate_up_type, x) + t, k, two_i = h.shape + if gate_up_scales is not None: + h.view(t, k, 2, two_i // 2).mul_(gate_up_scales[ids, :2].unsqueeze(-1)) + inter = torch.empty(t, k, two_i // 2, dtype=h.dtype, device=h.device) + gated_act_and_mul(layer.activation, h.view(-1, two_i), inter.view(-1, two_i // 2), alpha=float(layer.alpha), limit=limit_or_inf(layer)) + del h # free the gate_up output before the down GEMM allocates its own + if down_scales is not None and x.shape[0] <= ggml.mmvq_max_tokens(down_type): + y = ggml.mul_mat_id(view.tensors["down"], down_type, inter, ids, scale=down_scales[:, 2]) + else: + y = mul_mat_id(view.tensors["down"], down_type, inter) + if down_scales is not None: + topk_weights = topk_weights * down_scales[ids, 2] + return ggml.moe_weighted_sum(y, topk_weights) + + @staticmethod + def _apply_float(layer, x, topk_weights, topk_ids, view: ExpertView, ggml_type: int, *, is_prefill: bool): + """Float experts: their banks in their own dtype, through the bf16 fused MoE (the unquantized kernel's path).""" + from freetoken.moe.fused import fused_experts_decode_impl, fused_experts_impl + + gate_up, down = view.tensors["gate_up"], view.tensors["down"] + if gate_up.dtype != x.dtype: + raise NotImplementedError(f"GGUF {type_name(ggml_type)} experts need {gate_up.dtype} activations, got {x.dtype}") + impl = fused_experts_impl if is_prefill or is_resident(layer) else fused_experts_decode_impl + return impl(x, gate_up, down, topk_weights, topk_ids, layer.activation, layer.apply_router_weight_on_input, float(layer.alpha), limit_or_inf(layer)) + + +@register_method(QuantKind.GGUF, LayerKind.MOE) +class GGUFMoEMethod(MoEMethod): + candidates = (GgmlMoEKernel,) + + def cpu_spec(self) -> tuple[str | None, dict[str, int]]: + gate_up, down = expert_types(self.cfg) + # BF16 experts keep the executor's bf16 kernels, which use AVX512-BF16 where the CPU has it; ggml-cpu is built for AVX2 only + if float_dtype(gate_up) is torch.bfloat16 and float_dtype(down) is torch.bfloat16: + return "bf16", {} + # an NVFP4 layer of a file with weight scales reads the layout's global bank too + scaled = self.cfg.scheme.weight.scale is not None and NVFP4 in (gate_up, down) + return "gguf_scaled" if scaled else "gguf", {"gate_up_type": gate_up, "down_type": down} diff --git a/python/freetoken/layers/quantization/moe/mxfp4.py b/python/freetoken/layers/quantization/moe/mxfp4.py index aa020ee33a..a17b5bbe09 100644 --- a/python/freetoken/layers/quantization/moe/mxfp4.py +++ b/python/freetoken/layers/quantization/moe/mxfp4.py @@ -12,17 +12,24 @@ 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 cpu_options(self, cfg: MoEConfig) -> dict[str, int]: + return {"act_block": int(cfg.act_block)} def layout(self, cfg: MoEConfig) -> dict[str, BankSpec]: i, h = cfg.intermediate, cfg.hidden @@ -44,13 +51,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 +114,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..924401ef7f 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,27 +41,33 @@ 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), - "gate_up_global": BankSpec((2 * i,), torch.float16), + "gate_up_global": BankSpec((2 * i,), torch.float32), "down": BankSpec((h, i // 2), torch.uint8), "down_scale": BankSpec((h, i // GROUP), FP8), - "down_global": BankSpec((h,), torch.float16), + "down_global": BankSpec((h,), torch.float32), } 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)) @@ -78,6 +84,56 @@ def apply(self, layer, x, topk_weights, topk_ids, view: ExpertView, *, is_prefil return fused_experts_decode_nvfp4_marlin(x, *banks, topk_weights, topk_ids, layer.activation, layer.apply_router_weight_on_input, alpha, limit) +class W4a4Nvfp4MoEKernel(TritonNvfp4MoEKernel): + """The Triton kernels' banks and decode, with the prefill GEMMs on NVFP4 activations (W4A4) where the GPU and the batch favour it. + + Needs a checkpoint whose experts declare static activation scales; the GPU-resident alphas carry them per expert (gate_up_alpha: the larger of gate's and up's input_scale). + """ + + name = "w4a4" + + def unusable_reason(self, cfg: MoEConfig) -> str | None: + reason = super().unusable_reason(cfg) + if reason: + return reason + if cfg.scheme is None or not cfg.scheme.has("input_scale"): + return "the experts declare no static activation scale (W4A16 checkpoint)" + if cfg.apply_router_weight_on_input: + return "the routing weight scales the input, which is quantized once per token" + return None + + def worth_it(self, cfg: MoEConfig) -> bool: + from freetoken.kernel import nvfp4_moe + + return not backend.is_rocm() and torch.cuda.is_available() and nvfp4_moe.w4a4_available(torch.cuda.current_device()) + + def layout(self, cfg: MoEConfig) -> dict[str, BankSpec]: + return { + **super().layout(cfg), + "gate_up_alpha": BankSpec((), torch.float32, resident=True), + "down_alpha": BankSpec((), torch.float32, resident=True), + } + + def pack(self, pieces, cfg: MoEConfig, out): + super().pack(pieces, cfg, out) + if not all(f"{p}_input" in pieces for p in ("gate", "up", "down")): + raise ValueError("w4a4 nvfp4 MoE kernel: the checkpoint has no expert input_scale tensors") + gate, up, down = (pieces[f"{p}_input"].reshape(len(pieces[f"{p}_input"]), -1)[:, 0].float() for p in ("gate", "up", "down")) + return {"gate_up_alpha": torch.maximum(gate, up), "down_alpha": down} + + def apply(self, layer, x, topk_weights, topk_ids, view: ExpertView, *, is_prefill: bool): + if not is_prefill: + return super().apply(layer, x, topk_weights, topk_ids, view, is_prefill=False) + from freetoken.moe.fused_nvfp4 import fused_experts_nvfp4_w4a4 + + t = view.tensors + assert view.alphas is not None + return fused_experts_nvfp4_w4a4( + x, t["gate_up"], t["gate_up_scale"], t["gate_up_global"], t["down"], t["down_scale"], t["down_global"], + view.alphas[0], view.alphas[1], topk_weights, topk_ids, view.n, layer.activation, float(layer.alpha), limit_or_inf(layer), + ) + + # --------------------------------------------------------------------------- # Marlin: vLLM's fused Marlin MoE (W4A16, dequant in kernel) over Marlin-tiled banks. # The Marlin kernels come from vLLM (not sgl-kernel) deliberately: vLLM ships the @@ -133,7 +189,7 @@ def _marlin_pack_proj( """Pack one expert projection (on GPU) into the Marlin layout. ``packed`` ``[N, K//2]`` uint8 e2m1 pairs, ``scale`` ``[N, K//16]`` e4m3, - ``row_global`` ``[N]`` fp16 per-row global (rows of a merged gate_up may carry + ``row_global`` ``[N]`` fp32 per-row global (rows of a merged gate_up may carry w1's and w3's different globals). Returns ``(qweight [K//16, 2N] int32, scales [K//16, N] e4m3-coded, global scalar bf16)``. @@ -228,9 +284,11 @@ class MarlinNvfp4MoEKernel(MoEKernel): max_slots = MARLIN_MAX_SLOTS def unusable_reason(self, cfg: MoEConfig) -> str | None: + if backend.is_rocm(): + return "Marlin is CUDA-only and unavailable on ROCm; use triton" 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 +299,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 +310,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) @@ -491,19 +549,21 @@ class B12xNvfp4MoEKernel(MoEKernel): name = "b12x" def unusable_reason(self, cfg: MoEConfig) -> str | None: + if backend.is_rocm(): + return "b12x is CUDA-only and unavailable on ROCm; use triton" cc = backend.device_capability() if cc < (12, 0): 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 +572,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 +585,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() @@ -566,10 +626,4 @@ 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") + candidates = (W4a4Nvfp4MoEKernel, TritonNvfp4MoEKernel, MarlinNvfp4MoEKernel, B12xNvfp4MoEKernel) 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/registry.py b/python/freetoken/layers/quantization/registry.py index 2d13cf50cc..c561e55edc 100644 --- a/python/freetoken/layers/quantization/registry.py +++ b/python/freetoken/layers/quantization/registry.py @@ -11,6 +11,7 @@ class LayerKind(Enum): LINEAR = "linear" MOE = "moe" + EMBEDDING = "embedding" def __str__(self) -> str: return self.value diff --git a/python/freetoken/layers/quantization/scheme.py b/python/freetoken/layers/quantization/scheme.py index 668eb6fb37..a3aff0bc0f 100644 --- a/python/freetoken/layers/quantization/scheme.py +++ b/python/freetoken/layers/quantization/scheme.py @@ -20,12 +20,14 @@ class QuantKind(Enum): MXFP8 = "mxfp8" NVFP4 = "nvfp4" MXFP4 = "mxfp4" + GGUF = "gguf" 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 +82,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: @@ -95,3 +106,15 @@ def nvfp4_scheme(*, input_scale: bool) -> QuantScheme: def mxfp4_scheme() -> QuantScheme: return QuantScheme(QuantKind.MXFP4, WeightDesc("e2m1", (1, MX_GROUP), "e8m0"), {"weight", "weight_scale"}) + + +def gguf_scheme(types: Iterable[int], scaled: bool = False) -> QuantScheme: + """Packed ggml blocks, one ggml type per fused part (``types`` follows the parts' order). + + ``scaled``: the NVFP4 tensors carry the fp32 ``.scale`` (weight_scale_2) that llama.cpp's converter writes.""" + return QuantScheme(QuantKind.GGUF, WeightDesc("ggml:" + ",".join(str(int(t)) for t in types), None, "f32" if scaled else None), {"weight"}) + + +def gguf_types(scheme: QuantScheme) -> tuple[int, ...]: + assert scheme.kind is QuantKind.GGUF, scheme + return tuple(int(t) for t in scheme.weight.elem.removeprefix("ggml:").split(",")) diff --git a/python/freetoken/layers/rotary.py b/python/freetoken/layers/rotary.py index 0d04645d99..d0178f3bc1 100644 --- a/python/freetoken/layers/rotary.py +++ b/python/freetoken/layers/rotary.py @@ -2,7 +2,7 @@ import functools import math -from typing import Any, Callable, Dict, Tuple +from typing import Any, Mapping, Tuple import torch @@ -14,50 +14,21 @@ def __init__( self, head_size: int, rotary_dim: int, - max_position_embeddings: int, - base: float, - post_process: None | Callable[[torch.Tensor], torch.Tensor] = None, - proportional: bool = False, - attention_factor: float = 1.0, + cos_sin_cache: torch.Tensor, is_neox: bool = True, ) -> None: super().__init__() + assert 0 < rotary_dim <= head_size + assert rotary_dim % 2 == 0 + assert head_size in [64, 128, 256, 512] self.head_size = head_size self.rotary_dim = rotary_dim # NeoX (half-rotation, HF default) vs GPT-J interleaved (adjacent pairs, # ``rope_interleave`` models: GLM MLA lineage). Both underlying kernels # accept the flag; the cos/sin cache layout is identical. self.is_neox = is_neox - if proportional: - assert 0 < rotary_dim <= head_size - assert rotary_dim % 2 == 0 - inv_freq = 1.0 / ( - base ** (torch.arange(0, head_size, 2, dtype=torch.float) / head_size) - ) - if rotary_dim < head_size: - inv_freq[rotary_dim // 2 :] = 0.0 - else: - # Standard (NeoX) rope. Supports partial rotary (rotary_dim < head_size): - # rope is applied to the first ``rotary_dim`` dims of each head, the rest pass - # through. Frequencies are spaced over ``rotary_dim`` (matches HF default - # partial rope, e.g. Qwen3.5 partial_rotary_factor, MiniMax-M2's - # ``apply_rotary_pos_emb``). Full rope is rotary_dim == head_size and is - # unaffected. ``head_size`` is passed to flashinfer separately so it rotates - # only the first ``rotary_dim`` dims. - assert 0 < rotary_dim <= head_size - assert rotary_dim % 2 == 0 - inv_freq = 1.0 / ( - base ** (torch.arange(0, rotary_dim, 2, dtype=torch.float) / rotary_dim) - ) - if post_process is not None: - inv_freq = post_process(inv_freq) - t = torch.arange(max_position_embeddings, dtype=torch.float) - freqs = torch.einsum("i,j -> ij", t, inv_freq) - cos = freqs.cos() * attention_factor - sin = freqs.sin() * attention_factor - # buffer, so don't load/save - self._cos_sin_cache = torch.cat((cos, sin), dim=-1) - assert self.head_size in [64, 128, 256, 512] + # shared with every rope of the same frequencies (rope_cos_sin_table), so never written + self._cos_sin_cache = cos_sin_cache from freetoken.kernel.backend import is_flashinfer_installed @@ -160,6 +131,7 @@ def __init__( ) -> None: super().__init__(*args, **kwargs) assert self.is_neox, "mrope is defined on the NeoX half-rotation layout" + assert self._cos_sin_cache.shape[1] == self.rotary_dim, "mrope reads a rotary_dim-wide table, not a proportional rope's" half = self.rotary_dim // 2 assert sum(mrope_section) == half, (mrope_section, half) self._section_table = build_section_table(tuple(mrope_section), layout) @@ -193,56 +165,50 @@ def forward( return query, key -def _get_rope( - head_dim: int, - rotary_dim: int, - max_position: int, - base: float, - rope_scaling: Dict[str, Any] | None = None, - is_neox: bool = True, -) -> RotaryEmbedding: - if rope_scaling is None: - return RotaryEmbedding(head_dim, rotary_dim, max_position, base, is_neox=is_neox) - # need to test some cases: - match rope_scaling["rope_type"]: +def rope_table_positions(max_position: int, rope_scaling: Mapping[str, Any] | None) -> int: + """Rows the cos/sin table needs: YaRN serves ``original_max_position_embeddings * factor``, + above ``max_position_embeddings`` when only the rope parameters are overridden.""" + if rope_scaling is None or rope_scaling.get("rope_type") != "yarn": + return max_position + extended = int(rope_scaling["original_max_position_embeddings"] * rope_scaling["factor"]) + return max(max_position, extended) + + +def _rope_frequencies( + head_size: int, rotary_dim: int, base: float, rope_scaling: Mapping[str, Any] | None +) -> Tuple[torch.Tensor, float]: + """``inv_freq`` over the table's half width, and the factor cos and sin are both scaled by.""" + rope_type = "default" if rope_scaling is None else rope_scaling["rope_type"] + if rope_type == "proportional": + # Frequencies spaced over the whole head; the pairs past rotary_dim do not rotate. + inv_freq = 1.0 / (base ** (torch.arange(0, head_size, 2, dtype=torch.float) / head_size)) + inv_freq[rotary_dim // 2 :] = 0.0 + return inv_freq, 1.0 + # Standard (NeoX) rope. Supports partial rotary (rotary_dim < head_size): rope is applied + # to the first ``rotary_dim`` dims of each head, the rest pass through, and frequencies + # are spaced over ``rotary_dim`` (HF default partial rope, e.g. Qwen3.5 + # partial_rotary_factor, MiniMax-M2's ``apply_rotary_pos_emb``). + inv_freq = 1.0 / (base ** (torch.arange(0, rotary_dim, 2, dtype=torch.float) / rotary_dim)) + match rope_type: case "default": - return RotaryEmbedding(head_dim, rotary_dim, max_position, base, is_neox=is_neox) - - case "proportional": - return RotaryEmbedding( - head_dim, - rotary_dim, - max_position, - base, - proportional=True, - is_neox=is_neox, - ) + return inv_freq, 1.0 case "llama3": scaling_factor: float = rope_scaling["factor"] low_freq_factor: float = rope_scaling["low_freq_factor"] high_freq_factor: float = rope_scaling["high_freq_factor"] original_max_position: int = rope_scaling["original_max_position_embeddings"] - - def post_process(inv_freq: torch.Tensor) -> torch.Tensor: - # no smooth if low_freq_factor == high_freq_factor - wave_len = 2 * math.pi / inv_freq - if low_freq_factor == high_freq_factor: - return torch.where( - wave_len < original_max_position / high_freq_factor, - inv_freq, - inv_freq / scaling_factor, - ) - - delta = high_freq_factor - low_freq_factor - smooth = (original_max_position / wave_len - low_freq_factor) / delta - smooth = torch.clamp(smooth, 0, 1) - factor = (1 - smooth) / scaling_factor + smooth - return factor * inv_freq - - return RotaryEmbedding( - head_dim, rotary_dim, max_position, base, post_process, is_neox=is_neox - ) + wave_len = 2 * math.pi / inv_freq + # no smooth if low_freq_factor == high_freq_factor + if low_freq_factor == high_freq_factor: + return torch.where( + wave_len < original_max_position / high_freq_factor, + inv_freq, + inv_freq / scaling_factor, + ), 1.0 + delta = high_freq_factor - low_freq_factor + smooth = torch.clamp((original_max_position / wave_len - low_freq_factor) / delta, 0, 1) + return ((1 - smooth) / scaling_factor + smooth) * inv_freq, 1.0 case "yarn": factor: float = rope_scaling["factor"] @@ -288,27 +254,54 @@ def _find_correction_dim(num_rotations: float) -> float: high = min(high, rotary_dim - 1) if low == high: # HF nudges instead of flooring the gap at 1 ("truncate": false) high += 0.001 - - def post_process(inv_freq: torch.Tensor) -> torch.Tensor: - ramp = torch.clamp( - (torch.arange(rotary_dim // 2, dtype=torch.float32) - low) / (high - low), - 0, 1, - ) - return (inv_freq / factor) * ramp + inv_freq * (1 - ramp) - - return RotaryEmbedding( - head_dim, - rotary_dim, - max_position, - base, - post_process, - attention_factor=float(attention_factor), - is_neox=is_neox, + ramp = torch.clamp( + (torch.arange(rotary_dim // 2, dtype=torch.float32) - low) / (high - low), 0, 1 ) + return (inv_freq / factor) * ramp + inv_freq * (1 - ramp), float(attention_factor) raise ValueError(f"Unsupported {rope_scaling = }") +@functools.cache +def _cos_sin_table( + head_size: int, + rotary_dim: int, + positions: int, + base: float, + rope_scaling: Tuple[Tuple[str, Any], ...] | None, + device: torch.device, +) -> torch.Tensor: + inv_freq, attention_factor = _rope_frequencies( + head_size, rotary_dim, base, dict(rope_scaling) if rope_scaling is not None else None + ) + freqs = torch.arange(positions, dtype=torch.float)[:, None] * inv_freq[None, :] + return torch.cat((freqs.cos() * attention_factor, freqs.sin() * attention_factor), dim=-1) + + +def rope_cos_sin_table( + head_size: int, + rotary_dim: int, + max_position: int, + base: float, + rope_scaling: Tuple[Tuple[str, Any], ...] | None = None, +) -> torch.Tensor: + """The ``[positions, width]`` cos|sin table, one per distinct rope: only proportional rope + depends on the head size, so an indexer whose heads differ in width reads the attention's table.""" + rope_map = dict(rope_scaling) if rope_scaling is not None else None + proportional = rope_map is not None and rope_map["rope_type"] == "proportional" + device = torch.get_default_device() + if device.type == "cuda" and device.index is None: + device = torch.device("cuda", torch.cuda.current_device()) + return _cos_sin_table( + head_size if proportional else rotary_dim, + rotary_dim, + rope_table_positions(max_position, rope_map), + base, + rope_scaling, + device, + ) + + _ROPE_DEVICE: torch.device | None = None @@ -328,17 +321,15 @@ def get_rope( mrope_section: Tuple[int, ...] | None = None, mrope_layout: str = "interleaved", ) -> RotaryEmbedding: - rope_map = dict(rope_scaling) if rope_scaling is not None else None - def build() -> RotaryEmbedding: + table = rope_cos_sin_table(head_dim, rotary_dim, max_position, base, rope_scaling) if mrope_section is not None: - assert rope_map is None or rope_map.get("rope_type", "default") == "default" return MRotaryEmbedding( - head_dim, rotary_dim, max_position, base, + head_dim, rotary_dim, table, is_neox=is_neox, mrope_section=tuple(mrope_section), layout=mrope_layout, ) - return _get_rope(head_dim, rotary_dim, max_position, base, rope_map, is_neox) + return RotaryEmbedding(head_dim, rotary_dim, table, is_neox=is_neox) t = torch.tensor([]) if t.device == torch.device("meta"): @@ -358,5 +349,7 @@ def build() -> RotaryEmbedding: "RotaryEmbedding", "build_section_table", "get_rope", + "rope_cos_sin_table", + "rope_table_positions", "set_rope_device", ] diff --git a/python/freetoken/memory.py b/python/freetoken/memory.py new file mode 100644 index 0000000000..32832aa9a6 --- /dev/null +++ b/python/freetoken/memory.py @@ -0,0 +1,294 @@ +"""Conservative host-memory admission for bare-metal and cgrouped processes. + +The kernel's host-wide ``MemAvailable`` figure is not an allocation budget for a +container. This module combines it with the process's tightest finite cgroup +budget and deliberately keeps the probing code independent of torch and the +rest of the runtime so startup checks and diagnostic tools can share it. +""" + +from __future__ import annotations + +import re +from pathlib import Path + + +# Kernel pseudo-files are small, but bound every read and hierarchy walk. A +# container may expose a synthetic or partially mounted cgroup tree; admission +# must not turn an unexpected file into an unbounded read or traversal. +_CGROUP_VALUE_MAX_BYTES = 128 +_CGROUP_MEMBERSHIP_MAX_BYTES = 64 * 1024 +_CGROUP_MAX_ANCESTORS = 64 +_MEMINFO_MAX_BYTES = 256 * 1024 +_UINT64_MAX = (1 << 64) - 1 + +# Cgroup v1 represents "unlimited" with a page-aligned value just below the +# signed 64-bit maximum (commonly 9223372036854771712). No supported host can +# offer an exbibyte of usable RAM, so this threshold safely covers the kernel's +# sentinel variants without mistaking a practical finite limit for unlimited. +_CGROUP_V1_UNLIMITED_MIN = 1 << 60 + + +def _read_bounded_ascii(path: Path, max_bytes: int) -> tuple[bool, str | None]: + """Return ``(present, exact_text)`` for a small kernel pseudo-file. + + A missing path means that source/controller is unavailable. Other I/O + errors are authoritative probe failures and raise so the public resolver can + fail closed instead of silently treating the process as unconstrained. + ``None`` text means the present file was oversized or not ASCII. + """ + path = Path(path) + try: + with path.open("rb") as stream: + raw = stream.read(max_bytes + 1) + except (FileNotFoundError, NotADirectoryError): + return False, None + except OSError as exc: + raise ValueError(f"cannot read memory control file {path}: {exc}") from exc + if len(raw) > max_bytes: + return True, None + try: + return True, raw.decode("ascii") + except UnicodeDecodeError: + return True, None + + +def _parse_u64_counter(text: str | None) -> int | None: + """Parse one non-negative decimal counter with at most one final newline.""" + match = re.fullmatch(r"(\d+)\n?", text or "") + if match is None: + return None + value = int(match.group(1)) + return value if value <= _UINT64_MAX else None + + +def _host_memory_available(meminfo_path: Path) -> int | None: + """Read host-wide ``MemAvailable`` bytes, or ``None`` when /proc is absent.""" + present, text = _read_bounded_ascii(meminfo_path, _MEMINFO_MAX_BYTES) + if not present: + return None + if text is None: + raise ValueError(f"malformed or oversized memory information: {meminfo_path}") + match = re.search( + r"^MemAvailable:[ \t]+(\d+)[ \t]+kB[ \t]*(?:\n|$)", + text, + flags=re.MULTILINE, + ) + if match is None: + raise ValueError(f"MemAvailable is unavailable or malformed: {meminfo_path}") + kib = int(match.group(1)) + if kib > _UINT64_MAX // 1024: + raise ValueError(f"MemAvailable overflows a byte counter: {meminfo_path}") + return kib * 1024 + + +def _cgroup_pair_remaining( + directory: Path, + limit_name: str, + current_name: str, + *, + v1: bool = False, +) -> tuple[bool, int | None]: + """Return ``(limit_present, finite_remaining_or_none)`` for one cgroup.""" + limit_path = Path(directory) / limit_name + current_path = Path(directory) / current_name + limit_present, limit_text = _read_bounded_ascii(limit_path, _CGROUP_VALUE_MAX_BYTES) + if not limit_present: + return False, None + if not v1 and limit_text in ("max", "max\n"): + return True, None + + limit = _parse_u64_counter(limit_text) + if limit is None: + raise ValueError(f"malformed cgroup memory limit: {limit_path}") + if v1 and limit >= _CGROUP_V1_UNLIMITED_MIN: + return True, None + + current_present, current_text = _read_bounded_ascii( + current_path, _CGROUP_VALUE_MAX_BYTES + ) + if not current_present: + raise ValueError(f"cgroup memory usage is unavailable: {current_path}") + current = _parse_u64_counter(current_text) + if current is None: + raise ValueError(f"malformed cgroup memory usage: {current_path}") + return True, max(0, limit - current) + + +def _read_cgroup_memberships(proc_cgroup_path: Path) -> tuple[str | None, str | None]: + """Return the v2 and v1-memory paths from bounded ``/proc/self/cgroup``. + + Hybrid hosts list leftover v1 hierarchies (``1:net_cls:/``) above the unified + ``0::/...`` line; non-memory v1 controllers are simply skipped. + """ + present, text = _read_bounded_ascii(proc_cgroup_path, _CGROUP_MEMBERSHIP_MAX_BYTES) + if not present: + return None, None + if text is None: + raise ValueError( + f"malformed or oversized cgroup membership: {proc_cgroup_path}" + ) + + v2_path = None + v1_path = None + for line in text.splitlines(): + fields = line.split(":", 2) + if len(fields) != 3: + continue + hierarchy, controllers, member_path = fields + if not member_path.startswith("/"): + continue + # Real cgroup membership paths are absolute and contain no dot + # components. Reject lexical escapes in synthetic/malformed proc data. + parts = Path(member_path).parts[1:] + if any(part in ("", ".", "..") for part in parts): + continue + if hierarchy == "0" and not controllers: + v2_path = member_path + elif "memory" in controllers.split(","): + v1_path = member_path + return v2_path, v1_path + + +def _control_file_exists(path: Path) -> bool: + """Existence probe that fails closed on errors other than a missing path.""" + path = Path(path) + try: + path.stat() + return True + except (FileNotFoundError, NotADirectoryError): + return False + except OSError as exc: + raise ValueError(f"cannot inspect cgroup control file {path}: {exc}") from exc + + +def _cgroup_control_directory( + root: Path, member_path: str | None, limit_name: str +) -> Path | None: + """Map a proc membership path into a mounted, possibly namespaced tree.""" + root = Path(root) + if member_path: + parts = Path(member_path).parts[1:] + if len(parts) >= _CGROUP_MAX_ANCESTORS: + raise ValueError( + f"cgroup membership nesting exceeds {_CGROUP_MAX_ANCESTORS - 1} levels" + ) + candidate = root.joinpath(*parts) + # A leaf can omit controller files when only an ancestor has the memory + # controller enabled. Find the nearest visible ancestor without a + # filesystem scan or a walk outside the mounted hierarchy. + for _ in range(_CGROUP_MAX_ANCESTORS): + if _control_file_exists(candidate / limit_name): + return candidate + if candidate == root: + break + parent = candidate.parent + if parent == candidate or (parent != root and root not in parent.parents): + break + candidate = parent + + # Container runtimes may mount the process's own subgroup as the hierarchy + # root while /proc still reports a host-side membership path. + if _control_file_exists(root / limit_name): + return root + return None + + +def _cgroup_hierarchy_remaining( + root: Path, + member_path: str | None, + limit_name: str, + current_name: str, + *, + v1: bool = False, +) -> tuple[bool, int | None]: + """Return ``(hierarchy_seen, tightest finite ancestor headroom)``.""" + root = Path(root) + current = _cgroup_control_directory(root, member_path, limit_name) + if current is None: + return False, None + + seen = False + remaining = None + for _ in range(_CGROUP_MAX_ANCESTORS): + present, candidate = _cgroup_pair_remaining( + current, limit_name, current_name, v1=v1 + ) + seen = seen or present + if candidate is not None: + remaining = candidate if remaining is None else min(remaining, candidate) + if current == root: + break + parent = current.parent + if parent == current or (parent != root and root not in parent.parents): + break + current = parent + return seen, remaining + + +def _cgroup_memory_remaining( + cgroup_root: Path = Path("/sys/fs/cgroup"), + proc_cgroup_path: Path = Path("/proc/self/cgroup"), +) -> int | None: + """Return this process's tightest finite cgroup memory headroom. + + Cgroup v2 is authoritative whenever its memory controller is visible. + ``max`` is a known-unlimited value. Only when v2 is absent do we try the + conventional, bounded v1 controller roots; no filesystem scan is used. + Malformed, overflowing, or unreadable visible controls raise ``ValueError``. + """ + root = Path(cgroup_root) + v2_path, v1_path = _read_cgroup_memberships(proc_cgroup_path) + v2_seen, remaining = _cgroup_hierarchy_remaining( + root, v2_path, "memory.max", "memory.current" + ) + if v2_seen: + return remaining + + for v1_root in (root / "memory", root): + v1_seen, remaining = _cgroup_hierarchy_remaining( + v1_root, + v1_path, + "memory.limit_in_bytes", + "memory.usage_in_bytes", + v1=True, + ) + if v1_seen: + return remaining + return None + + +def effective_memory_available( + *, + meminfo_path: Path = Path("/proc/meminfo"), + cgroup_root: Path = Path("/sys/fs/cgroup"), + proc_cgroup_path: Path = Path("/proc/self/cgroup"), +) -> int | None: + """Return the effective available host-memory budget in bytes. + + The result is ``min(MemAvailable, tightest finite cgroup headroom)`` when + both measurements exist, either known bound when only one exists, and + ``None`` only when neither source is available -- no procfs, a non-Linux + host, or a cgroup tree without a visible memory controller. ``None`` means + *unknown* and is never conflated with ``0``; callers keep their historical + best-effort behaviour on it. A known zero remains zero. + + A *present* but malformed, overflowing, or permission-denied control file + is an authoritative signal that the budget cannot be trusted, so it fails + closed to ``0`` rather than being mistaken for an unlimited controller. + """ + try: + host_available = _host_memory_available(Path(meminfo_path)) + cgroup_remaining = _cgroup_memory_remaining( + Path(cgroup_root), Path(proc_cgroup_path) + ) + except ValueError: + return 0 + + if host_available is None: + return cgroup_remaining + if cgroup_remaining is None: + return host_available + return min(host_available, cgroup_remaining) + + +__all__ = ["effective_memory_available"] diff --git a/python/freetoken/message/frontend.py b/python/freetoken/message/frontend.py index 24725567b9..35620c7e8e 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,16 @@ 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 entries released with this reply (see + # tokenizer.detokenize.build_logprobs_entry), normally this token's own; None when + # there are none. Each entry's "text_offset" is where its token starts in the generated + # text. With stop strings an entry waits for its token's text, and entries + # of tokens a matched stop string trimmed arrive marked "stop_trimmed"; the token it + # ends inside carries "stop_visible", the length of its text that was released. + logprobs: list[dict] | None = None @dataclass diff --git a/python/freetoken/message/tokenizer.py b/python/freetoken/message/tokenizer.py index 9442ee059e..fcd599681f 100644 --- a/python/freetoken/message/tokenizer.py +++ b/python/freetoken/message/tokenizer.py @@ -36,6 +36,17 @@ 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 + # The prompt left a reasoning block open (SamplingParams.starts_in_reasoning). + starts_in_reasoning: 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 +59,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 +89,11 @@ 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 + # Prompt ids served as they are (``text`` is then ignored for tokenization): a client + # token-id prompt (OpenAI / vLLM completions; a list, range-checked by the worker), or + # the int32 tensor the frontend already rendered and counted (no second tokenization). + input_ids: Any = None @dataclass diff --git a/python/freetoken/mm/media.py b/python/freetoken/mm/media.py index 5c3c4ef789..90ef790de3 100644 --- a/python/freetoken/mm/media.py +++ b/python/freetoken/mm/media.py @@ -1,6 +1,8 @@ """Client-supplied image handling: ref collection, the accept gate, and byte fetching. Raises plain ValueError on bad input; the server layer maps it to its wire error type. +The message names the class of failure only: no local path, allowlist root, errno or +library text reaches the client (that detail goes to the server log). """ from __future__ import annotations @@ -8,10 +10,15 @@ import base64 from typing import TYPE_CHECKING, Any +from freetoken.utils import init_logger + if TYPE_CHECKING: from freetoken.server.args import ServerArgs +logger = init_logger(__name__) + _MAX_IMAGE_BYTES = 32 << 20 +_MAX_REDIRECTS = 5 def collect_image_refs(messages: list[dict[str, Any]]) -> list[dict[str, Any]]: @@ -56,6 +63,57 @@ def _check_media_domain(url: str, config: ServerArgs) -> None: ) +def _check_redirect_target(url: str, config: ServerArgs) -> None: + """A redirect hop is the remote server's choice, not the client's URL: it must pass the + allowlist like the original, and may never land on a loopback / private / link-local + address (an allowed public host 302-ing into the operator's network is the SSRF shape).""" + import ipaddress + from urllib.parse import urlparse + + if not url.startswith(("http://", "https://")): + raise ValueError("the URL redirects to a non-http(s) location") + try: + _check_media_domain(url, config) + except ValueError: + raise ValueError("the URL redirects outside the allowed domains") from None + host = (urlparse(url).hostname or "").rstrip(".").lower() + if host == "localhost" or host.endswith(".localhost"): + raise ValueError("the URL redirects to a private or loopback address") + try: + ip = ipaddress.ip_address(host) + except ValueError: + return # a name; the allowlist above is the gate for names (no DNS resolution here) + if isinstance(ip, ipaddress.IPv6Address) and ip.ipv4_mapped is not None: + ip = ip.ipv4_mapped + if not ip.is_global or ip.is_multicast: # covers loopback, private, link-local, CGNAT, unspecified + raise ValueError("the URL redirects to a private or loopback address") + + +async def _fetch_remote_media(url: str, config: ServerArgs) -> bytes: + """GET an http(s) ref, following redirects one hop at a time so every target is re-checked.""" + import httpx + + _check_media_domain(url, config) + async with httpx.AsyncClient(timeout=30.0, follow_redirects=False) as client: + for _ in range(_MAX_REDIRECTS + 1): + async with client.stream("GET", url) as resp: + if resp.next_request is not None: # a 3xx with a Location: httpx built the hop + url = str(resp.next_request.url) + _check_redirect_target(url, config) + continue + resp.raise_for_status() + declared = resp.headers.get("content-length", "") + if declared.isdigit() and int(declared) > _MAX_IMAGE_BYTES: + raise ValueError(f"image exceeds {_MAX_IMAGE_BYTES} bytes") + buf = bytearray() + async for chunk in resp.aiter_bytes(): + buf += chunk + if len(buf) > _MAX_IMAGE_BYTES: + raise ValueError(f"image exceeds {_MAX_IMAGE_BYTES} bytes") + return bytes(buf) + raise ValueError(f"the URL redirects more than {_MAX_REDIRECTS} times") + + def _load_local_media(url: str, config: ServerArgs) -> bytes: """Read a file:// ref; the resolved path must be a strict subpath of --allowed-local-media-path.""" from pathlib import Path @@ -66,45 +124,60 @@ def _load_local_media(url: str, config: ServerArgs) -> bytes: if not root: raise ValueError("cannot load local files without --allowed-local-media-path") spec = urlparse(url) - filepath = Path(url2pathname((spec.netloc or "") + (spec.path or ""))) + # RFC 8089: empty authority and localhost identify the local machine. + if spec.netloc and spec.netloc.lower().rstrip(".") != "localhost": + raise ValueError("file:// URLs must have no host or use localhost") + filepath = Path(url2pathname(spec.path or "")) # resolve() follows symlinks, so a link inside the root escaping it is rejected too resolved = filepath.resolve() if Path(root).resolve() not in resolved.parents: - raise ValueError( - f"the file path {filepath} must be a subpath of " - f"--allowed-local-media-path {root}" - ) + raise ValueError("the file path must be a subpath of --allowed-local-media-path") return resolved.read_bytes() +async def _fetch_one(ref: dict[str, Any], config: ServerArgs) -> bytes: + """One ref to bytes. Our own ValueErrors are already client-safe and pass through; anything + a library raises is reduced to its failure class and logged with the detail.""" + data = ref.get("data") or "" + try: + if ref.get("kind") == "b64": + return base64.b64decode(data) + if data.startswith("data:"): + return base64.b64decode(data.split(",", 1)[1]) + except Exception as exc: # noqa: BLE001 -- binascii.Error / IndexError, input-driven + raise ValueError("image data is not valid base64") from exc + if data.startswith(("http://", "https://")): + import httpx + + try: + return await _fetch_remote_media(data, config) + except ValueError: + raise + except httpx.HTTPStatusError as exc: + raise ValueError(f"the image URL returned HTTP {exc.response.status_code}") from exc + except httpx.TimeoutException as exc: + raise ValueError("timed out fetching the image URL") from exc + except Exception as exc: # noqa: BLE001 -- transport / protocol / bad URL + logger.warning("image URL fetch failed: %r", exc) + raise ValueError("the image URL is unreachable") from exc + if data.startswith("file://"): + try: + return _load_local_media(data, config) + except ValueError: + raise + except OSError as exc: + logger.warning("local image read failed: %r", exc) + raise ValueError("the local file is missing or unreadable") from exc + raise ValueError("unsupported image source (expect http(s)/file url or base64)") + + async def fetch_image_bytes(refs: list[dict[str, Any]], config: ServerArgs) -> list[bytes]: """Resolve collected image refs (URLs / base64) to raw bytes, in order.""" out: list[bytes] = [] for ref in refs: - data = ref.get("data") or "" try: - if ref.get("kind") == "b64": - out.append(base64.b64decode(data)) - elif data.startswith("data:"): - out.append(base64.b64decode(data.split(",", 1)[1])) - elif data.startswith(("http://", "https://")): - import httpx - - _check_media_domain(data, config) - # redirect targets are not re-checked against the allowlist - async with httpx.AsyncClient(timeout=30.0, follow_redirects=True) as client: - resp = await client.get(data) - resp.raise_for_status() - if len(resp.content) > _MAX_IMAGE_BYTES: - raise ValueError(f"image exceeds {_MAX_IMAGE_BYTES} bytes") - out.append(resp.content) - elif data.startswith("file://"): - out.append(_load_local_media(data, config)) - else: - raise ValueError( - "unsupported image source (expect http(s)/file url or base64)" - ) - except Exception as exc: # noqa: BLE001 -- input-driven, client-classifiable + out.append(await _fetch_one(ref, config)) + except ValueError as exc: raise ValueError(f"could not load image: {exc}") from exc return out diff --git a/python/freetoken/mm/processor.py b/python/freetoken/mm/processor.py index 48211a5b80..84463bc2eb 100644 --- a/python/freetoken/mm/processor.py +++ b/python/freetoken/mm/processor.py @@ -200,7 +200,8 @@ def get_mm_processor(model_path: str, mm: MultimodalConfig | None = None) -> MMP served = [e for e in spec.encoders if getattr(config, e.config_key, None) is not None and e.kind not in mm.disabled_encoders] if spec.mm_processor is None or not served: return None - check_mm_pad_shift(config.text_config.vocab_size) + # families with a flat checkpoint config (deepseek_v4) carry vocab_size at the top level + check_mm_pad_shift(getattr(config, "text_config", config).vocab_size) module, _, cls = spec.mm_processor.partition(":") return getattr(importlib.import_module(module), cls)(config, model_path, mm) diff --git a/python/freetoken/mm/processors/deepseek_v4.py b/python/freetoken/mm/processors/deepseek_v4.py new file mode 100644 index 0000000000..5cea991994 --- /dev/null +++ b/python/freetoken/mm/processors/deepseek_v4.py @@ -0,0 +1,129 @@ +"""DeepSeek-V4 image processor: each ``<｜deepseek_image｜>`` expands into the image's block of tokens. + +Where the reference streams ``vocab_size + type`` pseudo ids, the whole span carries one +content pad id and the block's embeddings -- sentinels included -- come from the encoder. +""" + +from __future__ import annotations + +import io +import struct +from typing import Any + +import torch + +from freetoken.message import MMItem +from freetoken.mm import mm_pad_value +from freetoken.mm.config import MultimodalConfig +from freetoken.mm.processor import MMResult, MMProcessor, PromptReplacement, _find_all, content_hash +from freetoken.models.deepseek_v4.config import IMAGE_TOKEN_ID, DSV4VisionConfig, parse_vision_config +from freetoken.models.deepseek_v4.image_processor import load_image +from freetoken.models.deepseek_v4.vision import COMPRESS_PAD_TO, build_image_block + + +class DSV4MMProcessor(MMProcessor): + def __init__(self, hf_config: Any, model_path: str, mm: MultimodalConfig) -> None: + super().__init__(model_path, mm) + vc = parse_vision_config(hf_config) + if vc is None: + raise ValueError(f"{model_path} carries no vision tower (vision_n_layers <= 0)") + self.vc: DSV4VisionConfig = vc + self.image_token_id = IMAGE_TOKEN_ID + self.placeholder = [self.image_token_id] + # --image-max-tokens lowers the checkpoint's per-image block cap + self.max_n_token = min(vc.vision_max_n_token, mm.image_max_tokens) if mm.image_max_tokens else vc.vision_max_n_token + + def process(self, images: list[Any]) -> list[MMItem]: + items: list[MMItem] = [] + for image in images: + patches, n_vit_h, n_vit_w, n_llm_h, n_llm_w = load_image(image, self.vc, self.max_n_token) + # layout, hash and token count are only known at the insertion offset + items.append( + MMItem( + modality="image", + hash=0, + pad_value=0, + offsets=[], + feature=patches, + model_specific_data={ + "n_vit_h": n_vit_h, + "n_vit_w": n_vit_w, + "n_llm_h": n_llm_h, + "n_llm_w": n_llm_w, + }, + ) + ) + return items + + def prompt_replacement(self, item: MMItem, start: int) -> PromptReplacement: + """The DSV4-only ``start`` is where the block lands in input_ids; only our apply below calls this.""" + n_llm_h, n_llm_w = item.n_llm_h, item.n_llm_w + types, _perm = build_image_block(n_llm_h, n_llm_w, start) + # the lead pads align the block, so the same image at another offset is another + # sequence: the layout belongs in the hash or the two would share a radix key + compress_pad = COMPRESS_PAD_TO - 1 - start % COMPRESS_PAD_TO + item.hash = content_hash(item.feature, struct.pack("<3i", n_llm_h, n_llm_w, compress_pad)) + item.pad_value = mm_pad_value(item.hash) + item.model_specific_data["start"] = start + return PromptReplacement([self.image_token_id] * types.numel()) + + def apply(self, input_ids: torch.Tensor, images: list[bytes]) -> MMResult: + """MMProcessor.apply, save that each replacement is built knowing the offset it lands at.""" + from PIL import Image + + ids = input_ids.tolist() + target = self.placeholder + slots = _find_all(ids, target) + if len(slots) != len(images): + raise ValueError( + f"prompt renders {len(slots)} image placeholders but the request " + f"carries {len(images)} images" + ) + + pils = [Image.open(io.BytesIO(raw)).convert("RGB") for raw in images] + items = self.process(pils) + out: list[int] = [] + cursor = 0 + for slot, item in zip(slots, items): + out.extend(ids[cursor:slot]) + base = len(out) + repl = self.prompt_replacement(item, base) + full = list(repl.full) + spans = repl.embed_spans() + # embedding slots carry the content pad id so radix keys and the model's scatter mask see the image + for lo, hi in spans: + full[lo:hi] = [item.pad_value] * (hi - lo) + out.extend(full) + item.offsets = [[base + lo, base + hi] for lo, hi in spans] + item.validate() + cursor = slot + len(target) + out.extend(ids[cursor:]) + new_ids = torch.tensor(out, dtype=input_ids.dtype) + + positions, delta = self.positions(len(out), items) or (None, 0) + return MMResult(new_ids, items, positions, delta) + + def dummy_items(self, dtype: torch.dtype, device: torch.device) -> list[MMItem]: + # The smallest block the tower can run: one 2x2 merge grid out of a 6x6 patch grid. + r, p = self.vc.vision_downsample_ratio, self.vc.vision_patch_size + n_vit, n_llm = 2 * r, 2 + types, _perm = build_image_block(n_llm, n_llm, 0) + return [ + MMItem( + modality="image", + hash=0, + pad_value=0, + offsets=[[0, types.numel()]], + feature=torch.zeros(n_vit * n_vit, 3, p, p, dtype=dtype, device=device), + model_specific_data={ + "n_vit_h": n_vit, + "n_vit_w": n_vit, + "n_llm_h": n_llm, + "n_llm_w": n_llm, + "start": 0, + }, + ) + ] + + +__all__ = ["DSV4MMProcessor"] 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 d7b504a769..0d4faae420 100644 --- a/python/freetoken/models/config.py +++ b/python/freetoken/models/config.py @@ -1,4 +1,6 @@ from __future__ import annotations + +import os from dataclasses import dataclass from typing import Any, ClassVar, Dict, List, Literal, Tuple, TypeAlias @@ -7,6 +9,23 @@ # State-dict key prefixes of the vision stack; load_weight drops them when the engine serves text-only. VISION_KEY_PREFIXES = ("vision_tower.", "embed_vision.", "vision_embedder.", "visual.") +_ENV_TRUE = {"1", "true", "yes", "on"} + + +def fp8_dense_enabled() -> bool: + """Load-time FP8 for a model's bf16 attention / GDN projections (opt-in, default OFF): + per-tensor e4m3 weights, per-tensor dynamic activation scale, cuBLASLt W8A8 GEMMs + (``torch._scaled_mm``, sm_89+). ``FREETOKEN_FP8_DENSE=1``.""" + return os.getenv("FREETOKEN_FP8_DENSE", "0").strip().lower() in _ENV_TRUE + + +def fp8_lmhead_enabled() -> bool: + """Load-time FP8 for the lm_head as well (opt-in, default OFF, needs FREETOKEN_FP8_DENSE=1 + for the rest). Separate from :func:`fp8_dense_enabled` because this one moves the logits: + a per-tensor e4m3 vocab matrix changes every sampled token's score, so it carries its own + quality gate. ``FREETOKEN_FP8_LMHEAD=1``.""" + return os.getenv("FREETOKEN_FP8_LMHEAD", "0").strip().lower() in _ENV_TRUE + def detect_expert_quant(hf_config: Any) -> str: """Routed-expert quantization from a checkpoint's ``quantization_config``: ``"nvfp4"`` for @@ -92,6 +111,13 @@ class RotaryConfig: mrope_section: list | None = None mrope_layout: str = "contiguous" # see freetoken.layers.rotary.build_section_table + @property + def table_positions(self) -> int: + """Positions the cos/sin table covers: the longest sequence this rope can serve.""" + from freetoken.layers.rotary import rope_table_positions + + return rope_table_positions(self.max_position, self.scaling) + def mrope_layout_from_rope_params(rope_params: Any) -> str: """rope_parameters flags -> layout name: mrope_interleaved_glm, else mrope_interleaved, else contiguous.""" @@ -198,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" @@ -206,6 +233,11 @@ 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" + # image token spans attend to each other in both directions, so prefill chunks end + # before a block instead of splitting it (the scheduler's keep_images_whole gate) + bidirectional_mm_blocks: bool = False AttentionGroupConfig: TypeAlias = ( @@ -320,6 +352,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. @@ -441,7 +476,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, ...]: @@ -495,7 +530,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..5bbb9b8e72 100644 --- a/python/freetoken/models/deepseek_v4/__init__.py +++ b/python/freetoken/models/deepseek_v4/__init__.py @@ -10,19 +10,26 @@ 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 -from .config import parse_config +from .config import DSV4VisionConfig, parse_config, parse_vision_config from .model import DeepseekV4ForCausalLM -from .weight import iter_expert_pieces, iter_weights +from .vision import DSV4Aligner, DSV4Vision, DSV4VisionTower +from .weight import iter_expert_pieces, iter_vision_weights, iter_weights __all__ = [ "DeepseekV4Args", + "DSV4Aligner", + "DSV4Vision", + "DSV4VisionConfig", + "DSV4VisionTower", "load_args", "parse_config", + "parse_vision_config", "DeepseekV4ForCausalLM", "iter_weights", "iter_expert_pieces", + "iter_vision_weights", ] diff --git a/python/freetoken/models/deepseek_v4/args.py b/python/freetoken/models/deepseek_v4/args.py index b75d2800df..d2646a2b3b 100644 --- a/python/freetoken/models/deepseek_v4/args.py +++ b/python/freetoken/models/deepseek_v4/args.py @@ -69,6 +69,18 @@ class DeepseekV4Args: hc_sinkhorn_iters: int = 20 hc_eps: float = 1e-6 + # ----- vision tower (DeepSeek-V4-Flash-Vision-Exp; 0 on the text-only releases) ----- + 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 = 384 + vision_min_pixels: int = 147456 + vision_max_wh_ratio: int = 8 + def __post_init__(self) -> None: # JSON lists -> tuple so the dataclass stays hashable / immutable-ish. if isinstance(self.compress_ratios, list): diff --git a/python/freetoken/models/deepseek_v4/attention.py b/python/freetoken/models/deepseek_v4/attention.py index cb0dd3150b..91bc8ac095 100644 --- a/python/freetoken/models/deepseek_v4/attention.py +++ b/python/freetoken/models/deepseek_v4/attention.py @@ -38,6 +38,7 @@ def __init__(self, layer_id: int, args: DeepseekV4Args, *, quant_config=None, pr self.window_size = args.window_size self.compress_ratio = args.compress_ratios[layer_id] self.eps = args.norm_eps + self.max_image_tokens = args.vision_max_n_token self.attn_sink = torch.empty(self.n_heads, dtype=torch.float32) # the latent projections are replicated; wq_b shards over heads, wo_b over the output groups @@ -104,7 +105,53 @@ def _wo(self, o: torch.Tensor, bsz: int, seqlen: int) -> torch.Tensor: o = torch.einsum("bsgd,grd->bsgr", o, wo_a).flatten(2) return self.wo_b.forward(o) - def _prefill_segment(self, x_seg, qr_seg, kv_seg, ti: int, start_pos: int, n: int): + @staticmethod + def visible_window_cols( + spans, start_pos: int, n: int, win: int, max_img: int, device: torch.device, + addr_lo: int = 0, width: int | None = None, + ) -> torch.Tensor: + """Window candidate columns (``-1`` padded) for a segment containing image spans. + + Port of the reference ``get_image_visible`` + ``get_window_topk_idxs_visible`` + pair: span tokens see their whole span bidirectionally (capped at ``max_img``) + instead of the 128-window, while text tokens -- and span tokens near the span + start -- keep the ordinary window into history. Returns ``[n, width]`` in + request-absolute coordinates with live entries in ``[addr_lo, span_end)``; the + caller translates to slot-lut indices (cold: as-is; extend: minus ``w_lo``). + """ + if width is None: + width = win + max_img + left, right = Attention._span_left_right(spans, start_pos, n, max_img, device) + idx = start_pos + torch.arange(n, device=device) + left_add = (left - (win - 1)).clamp(min=0) + starts = (idx - (win - 1) - left_add).clamp(min=addr_lo) + cols = starts.unsqueeze(1) + torch.arange(width, device=device) + return torch.where(cols > (idx + right).unsqueeze(1), -1, cols) + + @staticmethod + def _span_left_right( + spans, start_pos: int, n: int, max_img: int, device: torch.device + ) -> tuple[torch.Tensor, torch.Tensor]: + """Per-token visible counts to the left/right within each image span. + + Metadata twin of the reference ``get_image_visible`` (which derives spans from + virtual token ids): ``left[q] = q - span_start`` clamped to ``max_img - 1``, + ``right[q] = span_end - q`` clamped to ``max_img``, zero outside spans. ``spans`` + are request-absolute positions inside ``[start_pos, start_pos + n)``. + """ + left = torch.zeros(n, dtype=torch.int64, device=device) + right = torch.zeros(n, dtype=torch.int64, device=device) + for s, e in spans: + lo, hi = s - start_pos, e - start_pos + idx = torch.arange(lo, hi, device=device) + left[lo:hi] = (idx - lo).clamp(max=max_img - 1) + # the END sentinel sits at e-1: visibility reaches the END TOKEN, not past it + right[lo:hi] = ((hi - 1) - idx).clamp(max=max_img) + return left, right + + def _prefill_segment( + self, x_seg, qr_seg, kv_seg, ti: int, start_pos: int, n: int, spans=None + ): """One request's prefill work inside a (possibly ragged) batch: persist its window KV, advance its compressor/indexer carry, and return its per-query candidate lists. @@ -113,6 +160,16 @@ def _prefill_segment(self, x_seg, qr_seg, kv_seg, ti: int, start_pos: int, n: in ``(win_global [1, n, w], cmp_global [1, n, c] | None)`` at natural widths (cold: w = min(n, win); extend: w = win); the caller pads BOTH halves to the batch-uniform widths and concatenates. + + ``spans`` (VL requests) are the request's image-block ranges ``[(start, end)]`` inside + ``[start_pos, start_pos + n)``. Vision checkpoints set ``bidirectional_mm_blocks`` on the + attention group, which makes prefill end a chunk before a block instead of splitting it; + only a block longer than one chunk budget is still cut. Inside a span, every token + sees the whole span (bidirectional, capped at ``vision_max_n_token``) instead of only its + 128-window, matching the reference's ``get_window_topk_idxs_visible``. All span-visible + candidates are intra-segment, so the live full-page map covers them (the 128-slot ring + never needs to serve >128-back reads). ``spans=None`` keeps the historical path + bit-identical. """ win, ratio, device = self.window_size, self.compress_ratio, x_seg.device end = start_pos + n @@ -120,23 +177,38 @@ def _prefill_segment(self, x_seg, qr_seg, kv_seg, ti: int, start_pos: int, n: in self.attn.store_window(kv_seg, self.layer_id, slots) if start_pos == 0: - win_cols = get_window_topk_idxs(win, 1, n, 0).to(device) - # natural width min(n, win); the caller pads to the batch-uniform width + if spans: + max_img = self.max_image_tokens + cols = Attention.visible_window_cols( + spans, 0, n, win, max_img, device, addr_lo=0, + width=min(n, win + max_img), + ) + win_cols = cols.unsqueeze(0) + else: + win_cols = get_window_topk_idxs(win, 1, n, 0).to(device) + # natural width min(n, win) -- win+max_img for span segments; the caller pads win_global = self.attn.win_cols_to_global(win_cols, slots) else: # Candidates span [w_lo, end): each new query's 128-sliding window over the retained # prefix plus the new tokens, causal-masked to -1 past its own position. w_lo = max(0, start_pos - win + 1) ws_pool = self.attn.window_slots_of(ti, w_lo, end) - abs_p = start_pos + torch.arange(n, device=device).unsqueeze(1) - cand = (abs_p - win + 1).clamp(min=w_lo) + torch.arange(win, device=device) - win_cols = torch.where(cand > abs_p, -1, cand - w_lo).unsqueeze(0) + if spans: + max_img = self.max_image_tokens + cols = Attention.visible_window_cols( + spans, start_pos, n, win, max_img, device, addr_lo=w_lo, + ) + win_cols = torch.where(cols < 0, cols, cols - w_lo).unsqueeze(0) + else: + abs_p = start_pos + torch.arange(n, device=device).unsqueeze(1) + cand = (abs_p - win + 1).clamp(min=w_lo) + torch.arange(win, device=device) + win_cols = torch.where(cand > abs_p, -1, cand - w_lo).unsqueeze(0) win_global = self.attn.win_cols_to_global(win_cols, ws_pool) if not ratio: return win_global, None - # Only the compressor/indexer read the matched tail page's slot; resolving it costs a - # host sync (.item()), so do it after the ratio-0 early-out. + # Only the compressor/indexer read the matched tail page's slot; resolving it costs + # a host sync (.item()), so do it after the ratio-0 early-out. tail_ws = ( int(self.attn.window_slots_of(ti, start_pos - 1, start_pos).item()) if start_pos > 0 else None @@ -157,7 +229,7 @@ def _prefill_segment(self, x_seg, qr_seg, kv_seg, ti: int, start_pos: int, n: in self.compressor.forward(x_seg, start_pos, slots, tail_window_slot=tail_ws, ti=ti) return win_global, self.attn.blocks_to_global(blocks, ratio, ti=ti) - def forward_ragged(self, x, segments, flat_positions): + def forward_ragged(self, x, segments, flat_positions, spans_by_ti=None): """Ragged batched prefill (cu_seqlens). ``x`` is [1, T, dim] -- the requests' NEW token streams concatenated (NO padding); ``segments`` is [(offset, n, table_idx, start_pos)] tiling [0, T); ``flat_positions`` [T] is each token's ABSOLUTE position, so a radix-hit @@ -197,9 +269,18 @@ def forward_ragged(self, x, segments, flat_positions): win_parts: list[torch.Tensor] = [] cmp_parts: list[torch.Tensor | None] = [] max_c = 0 + any_spans = False for off, n, ti, start_pos in segments: + spans = None + if spans_by_ti is not None: + req_spans = spans_by_ti.get(ti) + if req_spans: + end = start_pos + n + spans = [(s, e) for (s, e) in req_spans if s >= start_pos and e <= end] + any_spans = any_spans or bool(spans) win_global, cmp_global = self._prefill_segment( - x[:, off:off + n], qr[:, off:off + n], kv[0, off:off + n], ti, start_pos, n + x[:, off:off + n], qr[:, off:off + n], kv[0, off:off + n], ti, start_pos, n, + spans=spans, ) win_parts.append(win_global) cmp_parts.append(cmp_global) @@ -213,6 +294,11 @@ def forward_ragged(self, x, segments, flat_positions): # is enough to flip a greedy near-tie. bs>1 pads every part to win (the historical # batched behavior). n_window = win if len(segments) > 1 else win_parts[0].shape[-1] + if any_spans: + # Span segments carry win+max_img-wide candidates; pad every part to the widest + # so the single flat launch stays rectangular. Text-only batches (any_spans + # False) keep the historical widths bit-identically. + n_window = max(n_window, *(w.shape[-1] for w in win_parts)) flat = [] for i, (off, n, ti, _start) in enumerate(segments): wg = win_parts[i] diff --git a/python/freetoken/models/deepseek_v4/config.py b/python/freetoken/models/deepseek_v4/config.py index 1005b8a3b3..9671c68290 100644 --- a/python/freetoken/models/deepseek_v4/config.py +++ b/python/freetoken/models/deepseek_v4/config.py @@ -13,12 +13,38 @@ from __future__ import annotations +from dataclasses import dataclass from typing import Any from freetoken.models.config import DSV4AttentionGroupConfig, ModelConfig, RotaryConfig from .args import load_args +# ``<｜deepseek_image｜>``: the only vocab-level image token. The reference expands each +# occurrence into a block of ``vocab_size + type`` pseudo ids; here the block keeps this +# one id and the image embeddings are scattered over the span it occupies. +IMAGE_TOKEN_ID = 129264 + + +@dataclass(frozen=True) +class DSV4VisionConfig: + """The tower's dims, split out of the checkpoint's flat inference config so the engine's + section nulling and the loader treat DSV4 like every other family with a vision tower. + Field names follow the reference (``inference/model.py`` ModelArgs) one to one.""" + + vision_n_layers: int + vision_dim: int + vision_n_heads: int + vision_inter_dim: int + vision_patch_size: int + vision_rope_theta: float + vision_downsample_ratio: int + vision_max_n_token: int + vision_min_pixels: int + vision_max_wh_ratio: int + # the language model's hidden size the aligner projects into (ModelArgs.dim) + text_dim: int + def parse_config(hf_config: Any) -> ModelConfig: model_path = getattr(hf_config, "_name_or_path", None) or getattr( @@ -46,6 +72,8 @@ def parse_config(hf_config: Any) -> ModelConfig: if max_position <= 0: max_position = int(args.original_seq_len * args.rope_factor) + vision_config = parse_vision_config(hf_config) + return ModelConfig( num_layers=args.n_layers, num_qo_heads=args.n_heads, @@ -88,9 +116,38 @@ def parse_config(hf_config: Any) -> ModelConfig: num_kv_heads=1, # MLA-style shared latent (K == V) head_dim=args.head_dim, sliding_window=args.window_size, + bidirectional_mm_blocks=vision_config is not None, ), ), + vision_config=vision_config, + image_token_id=IMAGE_TOKEN_ID if vision_config is not None else None, + ) + + +def parse_vision_config(hf_config: Any) -> DSV4VisionConfig | None: + """The tower's dims from the checkpoint config, or None when it carries no vision stack. + + The vision dims sit at the top level of the checkpoint config, next to the language + model's, and the engine nulls ``vision_n_layers`` on the copy it hands the parser when + this process serves text-only (--text-model-only / --mm-disable). That null is what + decides, so the built tower always matches the weights the loader is told to read. + """ + n_layers = int(getattr(hf_config, "vision_n_layers", 0) or 0) + if n_layers <= 0: + return None + return DSV4VisionConfig( + vision_n_layers=n_layers, + vision_dim=int(hf_config.vision_dim), + vision_n_heads=int(hf_config.vision_n_heads), + vision_inter_dim=int(hf_config.vision_inter_dim), + vision_patch_size=int(hf_config.vision_patch_size), + vision_rope_theta=float(hf_config.vision_rope_theta), + vision_downsample_ratio=int(hf_config.vision_downsample_ratio), + vision_max_n_token=int(hf_config.vision_max_n_token), + vision_min_pixels=int(hf_config.vision_min_pixels), + vision_max_wh_ratio=int(hf_config.vision_max_wh_ratio), + text_dim=int(hf_config.hidden_size), ) -__all__ = ["parse_config"] +__all__ = ["IMAGE_TOKEN_ID", "DSV4VisionConfig", "parse_config", "parse_vision_config"] diff --git a/python/freetoken/models/deepseek_v4/image_processor.py b/python/freetoken/models/deepseek_v4/image_processor.py new file mode 100644 index 0000000000..16742b1e3e --- /dev/null +++ b/python/freetoken/models/deepseek_v4/image_processor.py @@ -0,0 +1,50 @@ +"""DeepSeek-V4 image preprocessing: one image -> ViT patches plus its grid metadata. + +Port of the reference ``inference/image_processor.py:load_image``. The PIL work stays on the +CPU and must be bit-identical across runs: an image's radix key rides on these patches. +""" + +from __future__ import annotations + +import math + +import numpy as np +import torch +from PIL import Image, ImageOps + +from .config import DSV4VisionConfig +from .vision import safe_resize + + +def load_image(image: Image.Image, vc: DSV4VisionConfig, max_n_token: int | None = None) -> tuple[torch.Tensor, int, int, int, int]: + """One image -> ``(patches bf16 [n_vit_h*n_vit_w, 3, p, p], n_vit_h, n_vit_w, n_llm_h, n_llm_w)``. + + ``max_n_token`` lowers the checkpoint's per-image block cap (``--image-max-tokens``). + """ + image = image.convert("RGB") + p = vc.vision_patch_size + width, height = image.size + if vc.vision_max_wh_ratio is not None and width > height * vc.vision_max_wh_ratio: + width = height * vc.vision_max_wh_ratio + if 0 < width * height < vc.vision_min_pixels: + ratio = (vc.vision_min_pixels / (width * height)) ** 0.5 + width = int(width * ratio) + height = int(height * ratio) + best_width = math.ceil(width / p) * p + best_height = math.ceil(height / p) * p + n_llm_h, n_llm_w, best_height, best_width = safe_resize( + height, width, best_height, best_width, p, vc.vision_downsample_ratio, + vc.vision_max_n_token if max_n_token is None else min(vc.vision_max_n_token, max_n_token), + ) + n_vit_h, n_vit_w = best_height // p, best_width // p + if vc.vision_max_wh_ratio is not None and image.width >= vc.vision_max_wh_ratio * image.height: + image = image.resize((best_width, best_height)) + else: + image = ImageOps.pad(image, (best_width, best_height), 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) + patches = x.reshape(3, n_vit_h, p, n_vit_w, p).permute(1, 3, 0, 2, 4).reshape(n_vit_h * n_vit_w, 3, p, p) + return patches, n_vit_h, n_vit_w, n_llm_h, n_llm_w + + +__all__ = ["load_image"] diff --git a/python/freetoken/models/deepseek_v4/model.py b/python/freetoken/models/deepseek_v4/model.py index a2d4387ade..0c3441b951 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 @@ -29,11 +29,12 @@ from freetoken.kernel.triton.dsv4.hc import hc_post_combine, hc_pre_combine from freetoken.kernel.triton.dsv4.sinkhorn import hc_split_sinkhorn from freetoken.layers import BaseOP, OPList, ParallelLMHead, RMSNorm, VocabParallelEmbedding -from freetoken.models.blocks import BaseLLMModel +from freetoken.models.blocks import BaseLLMModel, embed_input_ids from .args import DeepseekV4Args from .attention import Attention from .moe import MoE +from .vision import COMPRESS_PAD_TO, DSV4Vision # Re-exports: keep every class/helper previously defined here importable from .model # (external import stability; the moved definitions live in their own modules). @@ -85,7 +86,7 @@ def hc_post(self, x, residual, post, comb): ) return y.view(shape) - def prefill_batched(self, x, input_ids, segments, flat_positions): + def prefill_batched(self, x, input_ids, segments, flat_positions, spans_by_ti=None): # Ragged batched prefill (cu_seqlens, no padding; bs >= 1, cold and radix-hit segments # mixed freely). ``x`` is [1, T, hc_mult, dim] -- the requests' token streams # concatenated. Per-token ops (HC / norm / MoE) run batched over ALL T tokens (the @@ -93,10 +94,12 @@ def prefill_batched(self, x, input_ids, segments, flat_positions): # T queries (Attention.forward_ragged), with the stateful compressor/indexer looped per # request. ``segments`` = [(offset, extend_len, table_idx, start_pos)] off the attention # metadata; ``flat_positions`` [T] = per-token ABSOLUTE position (batch.positions). + # ``spans_by_ti`` maps a page-table row to its image spans: the sparse attention needs + # the pairs, the block-tiled backends read the flat ends (batch.mm_block_ends) instead. residual = x x, post, comb = self.hc_pre(x, self.hc_attn_fn, self.hc_attn_scale, self.hc_attn_base) x = self.attn_norm.forward(x) - x = self.attn.forward_ragged(x, segments, flat_positions) + x = self.attn.forward_ragged(x, segments, flat_positions, spans_by_ti=spans_by_ti) x = self.hc_post(x, residual, post, comb) residual = x @@ -152,6 +155,7 @@ def hc_head(self, x): def prefill_batched( self, input_ids: torch.Tensor, segments, flat_positions: torch.Tensor, + spans_by_ti: dict | None = None, ) -> torch.Tensor: # Ragged batched prefill (bs >= 1). ``input_ids`` is [1, T] -- the requests' NEW tokens # concatenated (cu_seqlens, no padding); each request starts at its own cached_len @@ -163,10 +167,13 @@ def prefill_batched( # metadata; ``flat_positions`` [T] is the scheduler-staged batch.positions (per-token # ABSOLUTE position); the head picks each request's final token off the attention # metadata -> its next-token logits row. - h = self.embed.forward(input_ids.view(-1)).view(1, -1, self.args.dim) + batch = get_global_ctx().batch + # on a chunk that carries image rows, batch.mm_embeds' leading columns replace the + # embedding of every row in batch.mm_rows -- the whole image span, sentinels included + h = embed_input_ids(self.embed, input_ids.view(-1), batch).view(1, -1, self.args.dim) h = h.unsqueeze(2).repeat(1, 1, self.hc_mult, 1) for layer in self.layers.op_list: - h = layer.prefill_batched(h, input_ids, segments, flat_positions) + h = layer.prefill_batched(h, input_ids, segments, flat_positions, spans_by_ti) h = self.hc_head(h) h = self.norm.forward(h) return self.head.forward(h[0]) # [B, vocab] @@ -200,6 +207,26 @@ def decode( return self.head.forward(h[:, -1]) +def _image_spans_by_ti(reqs) -> dict | None: + """The batch's image spans as ``(IMAGE_START position, block end)`` pairs per page-table row. + + Only this model reads the pairs (built here, not staged on the batch): its sparse + attention builds candidate lists per span. The reference ``get_image_visible`` counts + visibility within [IMAGE_START, IMAGE_END], so the span starts at the START sentinel + (after ``build_image_block``'s alignment lead pads), not at the block's first token. + """ + spans: dict[int, list[tuple[int, int]]] = {} + for req in reqs: + req_spans = [ + (lo + COMPRESS_PAD_TO - 1 - lo % COMPRESS_PAD_TO, hi) + for item in req.mm_items or () + for lo, hi in item.offsets + ] + if req_spans: + spans[req.table_idx] = req_spans + return spans or None + + class DeepseekV4ForCausalLM(BaseLLMModel): """Engine adapter: a registered :class:`BaseLLMModel` wrapping the DSV4 transformer. @@ -211,6 +238,12 @@ def __init__(self, config): self._args: DeepseekV4Args = config.dsv4_args self.model = Transformer(self._args, config.quant, strategy=config.moe_strategy, decode_target=config.decode_target, prefix="model") self._bound = False + # Vision (DeepSeek-V4-Flash-Vision-Exp). config.is_multimodal is the engine's + # resolved gate, so a text-only process never builds weights the loader skips. + # The whole stack sits under one ``visual.`` mount, so the shared vision-key + # filter (VISION_KEY_PREFIXES) covers every image-only parameter. + if config.is_multimodal: + self.visual = DSV4Vision(config.vision_config, config.hidden_size) def _ensure_bound(self) -> None: if self._bound: @@ -226,6 +259,12 @@ def mark_for_rebind(self) -> None: block count, freqs) depends on the new pool's geometry, so re-derive it via _ensure_bound.""" self._bound = False + def place_encoder_weights(self, mode: str) -> None: + self.visual.place_weights(mode) + + def encode(self, item): + return self.visual.forward(item) + def forward(self) -> torch.Tensor: self._ensure_bound() batch = get_global_ctx().batch @@ -240,6 +279,7 @@ def forward(self) -> torch.Tensor: # cross requests. return self.model.prefill_batched( input_ids.view(1, -1), md.segments, batch.positions.long(), + spans_by_ti=_image_spans_by_ti(batch.reqs), ) # DECODE (bs>=1): per-row position (GPU int tensor -> no host syncs / graph safe). The # compressed staging cap is the max position any row reaches (eager); a static max_seq-1 diff --git a/python/freetoken/models/deepseek_v4/moe.py b/python/freetoken/models/deepseek_v4/moe.py index 6b6245674d..65de7f85fe 100644 --- a/python/freetoken/models/deepseek_v4/moe.py +++ b/python/freetoken/models/deepseek_v4/moe.py @@ -9,24 +9,34 @@ from freetoken.kernel.triton.dsv4.bf16_linear import bf16_linear_fp32 from freetoken.kernel.triton.dsv4.swiglu import fused_swiglu from freetoken.layers import BaseOP, LinearColParallelMerged, LinearRowParallel, OffloadMoELayer +from freetoken.mm import MM_PAD_SHIFT_VALUE from .args import DeepseekV4Args class Gate(BaseOP): - """MoE router: sqrtsoftplus scoring + hash routing (first ``n_hash_layers``).""" + """MoE router: sqrtsoftplus scoring + hash routing (first ``n_hash_layers``). + + A vision checkpoint gives every layer a ``bias_vl`` too: the expert-selection bias for + image rows, which never touches the routing weights. Hash layers keep routing text rows + through ``tid2eid`` and pick image rows by ``(scores + bias_vl).topk``, mirroring the + reference ``Gate.forward``. + """ def __init__(self, layer_id: int, args: DeepseekV4Args): self.topk = args.n_activated_experts self.score_func = args.score_func self.route_scale = args.route_scale self.hash = layer_id < args.n_hash_layers + vl = args.vision_n_layers > 0 self.weight = torch.empty(args.n_routed_experts, args.dim, dtype=torch.bfloat16) if self.hash: self.tid2eid = torch.empty(args.vocab_size, args.n_activated_experts, dtype=torch.int64) - self.bias = None + # present on hash layers of a vision checkpoint, loaded but never applied + self.bias = torch.empty(args.n_routed_experts, dtype=torch.float32) if vl else None else: self.bias = torch.empty(args.n_routed_experts, dtype=torch.float32) + self.bias_vl = torch.empty(args.n_routed_experts, dtype=torch.float32) if vl else None def forward(self, x: torch.Tensor, input_ids: torch.Tensor): scores = bf16_linear_fp32(x, self.weight) @@ -37,11 +47,20 @@ def forward(self, x: torch.Tensor, input_ids: torch.Tensor): else: scores = F.softplus(scores).sqrt() original_scores = scores - if self.bias is not None: - scores = scores + self.bias + # image rows are the content pad ids of an expanded image span + image_mask = input_ids >= MM_PAD_SHIFT_VALUE if self.bias_vl is not None else None if self.hash: - indices = self.tid2eid[input_ids] + if image_mask is None: + indices = self.tid2eid[input_ids] + else: + indices = self.tid2eid[torch.where(image_mask, 0, input_ids)] + vl_indices = (scores + self.bias_vl).topk(self.topk, dim=-1)[1] + indices = torch.where(image_mask.unsqueeze(-1), vl_indices.to(indices.dtype), indices) else: + if image_mask is None: + scores = scores + self.bias + else: + scores = scores + torch.where(image_mask.unsqueeze(-1), self.bias_vl, self.bias) indices = scores.topk(self.topk, dim=-1)[1] weights = original_scores.gather(1, indices) if self.score_func != "softmax": @@ -101,12 +120,13 @@ def _prefill_routed( cache = self.offload_cache assert cache is not None # unpinned (LOCKED) layers must take the base materialize path: their copy_missing is the whole-layer pageable branch with position == expert id, which ensure_experts's LRU slot remap would contradict (the GEMM would gather other experts' weights) + # one threshold for all layers (the smallest pool): with overlap prefill, the layers of a chunk must all stream or all take the slot path if ( - hidden_states.shape[0] * self.top_k >= self.num_experts + hidden_states.shape[0] * self.top_k >= min(self.num_experts, *cache.pool_sizes) or cache.is_unpinned_layer(self.layer_id) ): return super()._prefill_routed(hidden_states, topk_weights, topk_ids) - cache.ensure_experts(self.layer_id, topk_ids) # in-place expert-id -> slot + cache.ensure_experts(self.layer_id, topk_ids, update_state=False) # in-place expert-id -> slot; not a decode step cache.copy_missing() if cache.collect_stats: cache.record_decode_stats(self.layer_id) @@ -115,7 +135,7 @@ def _prefill_routed( hidden_states, topk_weights, topk_ids, - views=cache.bank_views(), + views=cache.pool_bank_views(self.layer_id), n=None, alphas=cache.alphas_for_slots(self.layer_id), is_prefill=True, diff --git a/python/freetoken/models/deepseek_v4/vision.py b/python/freetoken/models/deepseek_v4/vision.py new file mode 100644 index 0000000000..44549460ef --- /dev/null +++ b/python/freetoken/models/deepseek_v4/vision.py @@ -0,0 +1,336 @@ +"""DeepSeek-V4 vision tower, aligner and image-block layout. + +Port of the reference ``inference/vision.py`` / ``inference/image_processor.py`` from +deepseek-ai/DeepSeek-V4-Flash-Vision-Exp. The tower, the aligner and the four sentinel +vectors hang under one :class:`DSV4Vision` container the wrapper mounts as ``visual``, +so their state-dict keys sit behind the ``visual.`` prefix every VLM family filters on +(the checkpoint's own unprefixed names, ``vision.patch_embed.proj.weight``, +``aligner.w1.weight``, ``image_start``, ..., are renamed on load); every tensor is bf16 +exactly as stored. + +Image tokens occupy plain sequential positions: the 2-D structure lives inside the +tower (2-D RoPE over the patch grid) and in the N-layout of the token block the +:func:`build_image_block` / :func:`assemble_block` pair renders. All block embeddings, +sentinels included, come from the vision side, so the text embedding row of the +placeholder id is never read -- the span is overwritten before the first layer. +""" + +from __future__ import annotations + +import math +from functools import lru_cache + +import torch +import torch.nn.functional as F + +from freetoken.layers import BaseOP, OPList +from freetoken.models.weight_stream import BlockWeightStreamer + +from .config import DSV4VisionConfig + +# Sentinel token types inside an image block (reference inference/image_processor.py). +IMAGE_START, IMAGE_PAD, IMAGE, IMAGE_NEW_LINE, IMAGE_END = range(5) +# An image block is padded so its first row token lands on a compressor stride boundary. +COMPRESS_PAD_TO = 4 + + +@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 _Linear(BaseOP): + """Linear in checkpoint-native layout: ``weight`` [out, in] plus an optional ``bias``.""" + + def __init__(self, in_features: int, out_features: int, bias: bool = True): + self.weight = torch.empty(out_features, in_features, dtype=torch.bfloat16) + self.bias = torch.empty(out_features, dtype=torch.bfloat16) if bias else None + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return F.linear(x, self.weight, self.bias) + + +class RMSNorm(BaseOP): + def __init__(self, dim: int, eps: float = 1e-6): + self.eps = eps + self.weight = torch.empty(dim, dtype=torch.bfloat16) + + 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(BaseOP): + def __init__(self, patch_size: int, dim: int): + self.proj = _Linear(3 * patch_size**2, dim) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.proj.forward(x.flatten(1)) + + +class VisionAttention(BaseOP): + def __init__(self, dim: int, n_heads: int): + self.n_heads = n_heads + self.head_dim = dim // n_heads + self.wqkv = _Linear(dim, 3 * dim) + self.wo = _Linear(dim, 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.forward(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.forward(o.transpose(0, 1).reshape(n, -1)) + + +class VisionMLP(BaseOP): + def __init__(self, dim: int, inter_dim: int): + self.w1 = _Linear(dim, 2 * inter_dim, bias=False) + self.w2 = _Linear(inter_dim, dim, bias=False) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + gate, up = self.w1.forward(x).chunk(2, dim=-1) + return self.w2.forward(F.silu(gate) * up) + + +class VisionBlock(BaseOP): + def __init__(self, dim: int, n_heads: int, inter_dim: int, eps: float): + self.norm1 = RMSNorm(dim, eps) + self.attn = VisionAttention(dim, n_heads) + self.norm2 = RMSNorm(dim, eps) + self.mlp = VisionMLP(dim, inter_dim) + + def forward(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: + x = x + self.attn.forward(self.norm1.forward(x), cos, sin) + return x + self.mlp.forward(self.norm2.forward(x)) + + +class DSV4VisionTower(BaseOP): + """Native-resolution ViT: bidirectional attention inside one image, 2-D RoPE over its patch grid.""" + + def __init__(self, vc: DSV4VisionConfig): + self.rope_dim = vc.vision_dim // vc.vision_n_heads // 2 + self.rope_theta = vc.vision_rope_theta + self.patch_embed = PatchEmbed(vc.vision_patch_size, vc.vision_dim) + self.blocks = OPList( + [ + VisionBlock(vc.vision_dim, vc.vision_n_heads, vc.vision_inter_dim, eps=1e-6) + for _ in range(vc.vision_n_layers) + ] + ) + self.norm = RMSNorm(vc.vision_dim, 1e-6) + self._streamer: BlockWeightStreamer | None = None + + def place_weights(self, mode: str) -> None: + """gpu: every tensor resident; host: the block stack in pinned banks, two blocks streamed at a time (patch embed and final norm stay resident).""" + 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 _blocks(self): + if self._streamer is None: + return enumerate(self.blocks.op_list) + return self._streamer.blocks(self.blocks.op_list) + + def forward(self, patches: torch.Tensor, n_h: int, n_w: int) -> torch.Tensor: + patches = patches.to(self.patch_embed.proj.weight.dtype) + x = self.patch_embed.forward(patches) + cos, sin = get_vision_cos_sin(n_h, n_w, self.rope_dim, self.rope_theta) + cos, sin = cos.to(x.device), sin.to(x.device) + for _, block in self._blocks(): + x = block.forward(x, cos, sin) + return self.norm.forward(x) + + +class DSV4Aligner(BaseOP): + """3x3/stride-3 merge of every 3x3 patch grid into one token, then a 2-layer erf-GELU MLP.""" + + def __init__(self, vc: DSV4VisionConfig): + self.downsample_ratio = vc.vision_downsample_ratio + in_dim = vc.vision_dim * self.downsample_ratio**2 + self.w1 = _Linear(in_dim, vc.text_dim) + self.w2 = _Linear(vc.text_dim, vc.text_dim) + + def forward(self, x: torch.Tensor, n_h: int, n_w: int) -> torch.Tensor: + dtype = self.w1.weight.dtype + x = x.to(dtype) + 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.forward(F.gelu(self.w1.forward(x))) + + +class DSV4Vision(BaseOP): + """The wrapper's whole vision stack under one ``visual.`` mount, like the V4.1 sibling.""" + + def __init__(self, vc: DSV4VisionConfig, text_dim: int): + self.vision = DSV4VisionTower(vc) + self.aligner = DSV4Aligner(vc) + for name in ("image_start", "image_end", "image_newline", "image_pad"): + setattr(self, name, torch.empty(text_dim, dtype=torch.bfloat16)) + + def place_weights(self, mode: str) -> None: + self.vision.place_weights(mode) + + @torch.inference_mode() + def forward(self, item): + """One image -> the embeddings of its whole token block, sentinels included. + + The block's layout is a function of where it lands in the prompt, so the item + carries the grid and the offset; ``build_image_block`` re-derives the same + ``(types, perm)`` the processor tokenized with. + """ + device = self.vision.patch_embed.proj.weight.device + patches = item.feature.to(device, non_blocking=True) + aligned = self.aligner.forward( + self.vision.forward(patches, item.n_vit_h, item.n_vit_w), item.n_vit_h, item.n_vit_w + ) + types, perm = build_image_block(item.n_llm_h, item.n_llm_w, item.start) + sentinels = torch.stack( + [self.image_start, self.image_pad, self.image_pad, self.image_newline, self.image_end] + ) + return assemble_block(aligned, types, perm, sentinels) + + +def grid_tokens(best_height: int, best_width: int, patch_size: int, downsample_ratio: int): + """LLM-token footprint of a resized image: the N-layout row/align padding included.""" + n_llm_h = math.ceil((best_height // patch_size) / downsample_ratio) + n_llm_w = math.ceil((best_width // patch_size) / downsample_ratio) + num_tokens = n_llm_h * (n_llm_w + 1) + 2 + if n_llm_h % 2 == 1: + num_tokens += n_llm_w + 1 + num_tokens += (n_llm_h + 1) // 2 * (n_llm_w + 1) % 2 * 2 + return n_llm_h, n_llm_w, num_tokens + + +def solve_resize_ratio(height, width, patch_size, downsample_ratio, max_n_token): + r = height / width + max_w_float = math.sqrt((max_n_token - 2) / r + 0.25) - 0.5 + max_h_float = max_w_float * r + if max_w_float < 1.0: + max_w = 1 + max_h = (max_n_token - 2) // (max_w + 1) + if max_h % 2 == 1: + max_h -= 1 + best_width = max_w * patch_size * downsample_ratio + best_height = max_h * patch_size * downsample_ratio + elif max_h_float < 2.0: + max_h = 2 + max_w = ((max_n_token - 2) // max_h) - 1 + assert max_w > 1 + best_width = max_w * patch_size * downsample_ratio + best_height = max_h * patch_size * downsample_ratio + else: + max_w = math.floor(max_w_float) + max_h = math.floor(max_h_float) + if max_h % 2 == 1: + max_h -= 1 + beta = min(max_w * patch_size * downsample_ratio / width, max_h * patch_size * downsample_ratio / height) + best_width = math.floor(width * beta / patch_size) * patch_size + best_height = math.floor(height * beta / patch_size) * patch_size + n_llm_h, n_llm_w, num_tokens = grid_tokens(best_height, best_width, patch_size, downsample_ratio) + return n_llm_h, n_llm_w, best_height, best_width, num_tokens + + +def safe_resize(height, width, best_height, best_width, patch_size, downsample_ratio, max_n_token): + """Largest grid whose block still fits ``max_n_token`` once the alignment pads are added.""" + max_n_token -= COMPRESS_PAD_TO - 1 + n_llm_h, n_llm_w, num_tokens = grid_tokens(best_height, best_width, patch_size, downsample_ratio) + budget = max_n_token + while num_tokens > max_n_token: + n_llm_h, n_llm_w, best_height, best_width, num_tokens = solve_resize_ratio( + height, width, patch_size, downsample_ratio, budget + ) + budget -= 1 + return n_llm_h, n_llm_w, best_height, best_width + + +def build_image_block(n_llm_h: int, n_llm_w: int, start_pos: int): + """The block's token types in their final order, plus the aligner-row order of its IMAGE slots. + + ``types`` [n_out] places every token of the block (lead pads, START, the row-interleaved + IMAGE tokens, trailing pads, END). The reference streams them as ``vocab_size + types``; + here the span carries the image's content pad id and the assembled embeddings are + scattered over it. ``perm`` [G] maps aligner outputs (row-major grid) into the IMAGE slots. + + ``start_pos`` is the block's absolute position in the prompt: the lead-pad count aligns + the first row token to a compressor stride boundary, so the layout -- and the block's + length -- depend on where the image lands. + """ + compress_pad = COMPRESS_PAD_TO - 1 - start_pos % COMPRESS_PAD_TO + pad_h = n_llm_h % 2 + rows = n_llm_h + pad_h + row_len = n_llm_w + 1 + pad_last = rows // 2 * row_len % 2 * 2 + types = torch.tensor( + ([IMAGE] * n_llm_w + [IMAGE_NEW_LINE]) * n_llm_h + [IMAGE_PAD] * (row_len * pad_h), + dtype=torch.int64, + ) + order = torch.arange(rows * row_len).view(rows // 2, 2, row_len).transpose(1, 2).reshape(-1) + image_idx = torch.full((rows * row_len,), -1, dtype=torch.int64) + image_idx.view(rows, row_len)[:n_llm_h, :n_llm_w] = torch.arange(n_llm_h * n_llm_w).view(n_llm_h, n_llm_w) + perm = image_idx[order] + perm = perm[perm >= 0] + types = torch.cat( + [ + torch.full((compress_pad,), IMAGE_PAD, dtype=torch.int64), + torch.tensor([IMAGE_START]), + types[order], + torch.full((pad_last,), IMAGE_PAD, dtype=torch.int64), + torch.tensor([IMAGE_END]), + ] + ) + return types, perm + + +def assemble_block( + aligned: torch.Tensor, + types: torch.Tensor, + perm: torch.Tensor, + sentinels: torch.Tensor, +) -> torch.Tensor: + """Block embeddings [n_out, dim] from the aligner rows and the learned sentinels. + + ``sentinels`` [5, dim] is ``(start, pad, pad, newline, end)``, indexable by the type + enum exactly like the reference ``merge_image_embeddings``. + """ + block = sentinels[types.to(sentinels.device)] + block[types == IMAGE] = aligned[perm.to(aligned.device)] + return block + + +__all__ = [ + "COMPRESS_PAD_TO", + "DSV4Aligner", + "DSV4Vision", + "DSV4VisionTower", + "IMAGE", + "IMAGE_END", + "IMAGE_NEW_LINE", + "IMAGE_PAD", + "IMAGE_START", + "assemble_block", + "build_image_block", + "grid_tokens", + "safe_resize", + "solve_resize_ratio", +] diff --git a/python/freetoken/models/deepseek_v4/weight.py b/python/freetoken/models/deepseek_v4/weight.py index 4728f7f9c3..8bc4c1f098 100644 --- a/python/freetoken/models/deepseek_v4/weight.py +++ b/python/freetoken/models/deepseek_v4/weight.py @@ -78,12 +78,16 @@ def iter_weights( *, include_moe_experts: bool = True, include_non_moe: bool = True, + include_vision: bool = True, ): """Stream resident (non-expert) weights as ``(name, tensor)`` keyed to engine params. Routed MXFP4 experts come from the offload cache, so ``include_moe_experts`` must be False (DeepSeek-V4 only runs ``--moe-strategy offload``). Tensors yielded in checkpoint dtype (fp8 + e8m0 preserved); ``wo_a`` dequantized to bf16 to match the reference einsum. + ``include_vision`` (the engine's active-encoder set) drops the ``visual.*`` tower + tensors; the router biases stay in either way, since the gate builds + them on any vision checkpoint. """ if include_moe_experts: raise ValueError( @@ -112,6 +116,9 @@ def linear(src: str, dst: str): for nm in ("hc_head_fn", "hc_head_base", "hc_head_scale"): yield f"model.{nm}", get(nm) + if include_vision: + yield from _iter_vision(reader, args) + for L in range(args.n_layers): a = f"layers.{L}.attn" m = f"model.{a}" @@ -147,8 +154,13 @@ def linear(src: str, dst: str): yield f"model.{g}.weight", get(f"{g}.weight") if L < args.n_hash_layers: yield f"model.{g}.tid2eid", get(f"{g}.tid2eid") + if args.vision_n_layers > 0: + # a vision checkpoint gives its hash layers a text bias too + yield f"model.{g}.bias", get(f"{g}.bias") else: yield f"model.{g}.bias", get(f"{g}.bias") + if args.vision_n_layers > 0: + yield f"model.{g}.bias_vl", get(f"{g}.bias_vl") for proj in ("w1", "w2", "w3"): src = f"layers.{L}.ffn.shared_experts.{proj}" yield from linear(src, f"model.{src}") @@ -162,6 +174,58 @@ def linear(src: str, dst: str): reader.close() +# -------------------------------------------------------------------------------------- +# Vision tower. +# -------------------------------------------------------------------------------------- +# The wrapper mounts the tower, the aligner and the sentinels under one ``visual`` +# container, so every engine-side name carries the ``visual.`` prefix the shared +# VISION_KEY_PREFIXES filter matches; the checkpoint's unprefixed names +# (``vision.*``, ``aligner.*``, ``image_*``) are renamed here on load. +_VISION_BLOCK_SUFFIXES = ( + "norm1.weight", + "attn.wqkv.weight", + "attn.wqkv.bias", + "attn.wo.weight", + "attn.wo.bias", + "norm2.weight", + "mlp.w1.weight", + "mlp.w2.weight", +) +_VISION_TAIL = ( + "vision.patch_embed.proj.weight", + "vision.patch_embed.proj.bias", + "vision.norm.weight", + "aligner.w1.weight", + "aligner.w1.bias", + "aligner.w2.weight", + "aligner.w2.bias", + "image_start", + "image_end", + "image_newline", + "image_pad", +) + + +def _iter_vision(reader, args: DeepseekV4Args): + if args.vision_n_layers <= 0: + return + for L in range(args.vision_n_layers): + for suffix in _VISION_BLOCK_SUFFIXES: + yield f"visual.vision.blocks.{L}.{suffix}", reader.get(f"vision.blocks.{L}.{suffix}") + for name in _VISION_TAIL: + yield f"visual.{name}", reader.get(name) + + +def iter_vision_weights(model_path: str, device): + """The vision tower alone, for the engine's encoder-only load and the FTW converter.""" + args = load_args(model_path, max_batch_size=1) + reader = _ShardReader(model_path, _weight_map(model_path), device) + try: + yield from _iter_vision(reader, args) + finally: + reader.close() + + # -------------------------------------------------------------------------------------- # Routed MXFP4 expert pieces. # -------------------------------------------------------------------------------------- @@ -173,7 +237,7 @@ def linear(src: str, dst: str): _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): +def iter_expert_pieces(model_path: str, config, kind: QuantKind, *, layers=None, parallel: bool | None = False, workers: int = 8, chunk: int = 8 << 20): """Routed experts, one piece per expert: ``{gate, up, down}`` e2m1 pairs and their e8m0 ``_scale`` companions (``w1`` / ``w3`` / ``w2``). The MTP layer's experts are skipped.""" if kind is not QuantKind.MXFP4: @@ -185,10 +249,11 @@ def iter_expert_pieces(model_path: str, config, kind: QuantKind, *, parallel: bo args = load_args(model_path, max_batch_size=1) L, E = args.n_layers, args.n_routed_experts + keep = None if layers is None else set(layers) def locate(raw_name: str): m = _EXPERT_RE.match(raw_name) - if m is None or int(m["layer"]) >= L: + if m is None or int(m["layer"]) >= L or (keep is not None and int(m["layer"]) not in keep): return None return int(m["layer"]), int(m["expert"]), _PROJ_ROLE[m["proj"]] + _KIND_SUFFIX[m["kind"]] @@ -199,7 +264,7 @@ def locate(raw_name: str): def _serial(): reader = _ShardReader(model_path, _weight_map(model_path), torch.device("cpu")) try: - for li in tqdm(range(L), desc="Loading DSV4 experts (serial)", disable=not get_tp_info().is_primary()): + for li in tqdm(range(L) if keep is None else sorted(keep), desc="Loading DSV4 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"): @@ -212,4 +277,4 @@ def _serial(): return per_expert_pieces(_serial(), locate, tensors_per_expert=6) -__all__ = ["iter_weights", "iter_expert_pieces"] +__all__ = ["iter_weights", "iter_expert_pieces", "iter_vision_weights"] 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..962bd19589 --- /dev/null +++ b/python/freetoken/models/deepseek_v41/weight.py @@ -0,0 +1,243 @@ +"""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, *, layers=None, 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 + keep = None if layers is None else set(layers) + + def locate(raw_name: str): + m = _EXPERT_RE.match(raw_name) + if m is None or int(m["layer"]) >= L or (keep is not None and int(m["layer"]) not in keep): + 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) if keep is None else sorted(keep), 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/gemma4/__init__.py b/python/freetoken/models/gemma4/__init__.py index 9f41f080fa..4fcf15a373 100644 --- a/python/freetoken/models/gemma4/__init__.py +++ b/python/freetoken/models/gemma4/__init__.py @@ -1,11 +1,6 @@ from .attention import Gemma4Attention from .config import parse_config -from .gguf import ( - dummy_q4_0_expert_sources, - iter_gguf_weights, - load_q4_0_expert_sources, - parse_gguf_config, -) +from .gguf import iter_gguf_weights, parse_gguf_config, to_gguf from .model import Gemma4ForCausalLM, Gemma4ForConditionalGeneration, Gemma4UnifiedForConditionalGeneration from .moe import Gemma4Router from .vision import Gemma4MultimodalEmbedder, Gemma4UnifiedVisionEmbedder, Gemma4VisionModel @@ -32,6 +27,5 @@ "iter_weights_parallel", "parse_gguf_config", "iter_gguf_weights", - "load_q4_0_expert_sources", - "dummy_q4_0_expert_sources", + "to_gguf", ] diff --git a/python/freetoken/models/gemma4/gguf.py b/python/freetoken/models/gemma4/gguf.py index 437822b519..3d4de920b8 100644 --- a/python/freetoken/models/gemma4/gguf.py +++ b/python/freetoken/models/gemma4/gguf.py @@ -6,12 +6,13 @@ ``gemma4.config.parse_config`` -- only the source is GGUF KV metadata instead of a HF config object. transformers' own GGUF->config conversion is rejected by the gemma4 strict dataclass (per-layer ``num_key_value_heads`` array), so we read the -metadata directly. ``expert_quant`` is set to ``"q4_0"`` to route the routed experts -through the native-Q4_0 offload-cache path. +metadata directly. Every tensor keeps its packed ggml type (``QuantKind.GGUF``); the +routed experts load as GGUF expert banks (``freetoken.models.gguf.experts``). """ from __future__ import annotations +import re from typing import TYPE_CHECKING, Iterator import torch @@ -22,7 +23,8 @@ RotaryConfig, SWAAttentionGroupConfig, ) -from freetoken.models.gguf.dequant import GGML_Q4_0, GGML_Q6_K, dequantize, row_bytes +from freetoken.models.gguf.reader import gguf_tensor_index +from freetoken.models.gguf.tensors import check_scales, float_tensor, fused, matrix, weight_scale if TYPE_CHECKING: from freetoken.models.gguf.config import GgufConfigShim @@ -36,12 +38,9 @@ def _full_rotary_dim(shim: "GgufConfigShim", full_head_dim: int) -> int: which drives the unrotated tail's frequencies to zero (ggml ops.cpp ``theta / ff``). Count the untouched entries to get the rotated width back. """ - from freetoken.models.gguf.reader import iter_gguf_tensors - - for t in iter_gguf_tensors(shim.model_path): - if t.name == "rope_freqs.weight": # F32, so the packed bytes ARE the values - freqs = t.packed().reshape(-1).view(torch.float32) - return int((freqs == 1.0).sum().item()) * 2 + info = gguf_tensor_index(shim.model_path).get("rope_freqs.weight") + if info is not None: + return int((float_tensor(info) == 1.0).sum().item()) * 2 # A metadata-only GGUF (an FTW dir's source_metadata.gguf) has no tensor table. Every # released gemma4 sets partial_rotary_factor 0.25 on the full layers, and the converter # asserts rope_type == "proportional" (llama.cpp conversion/gemma.py). @@ -49,6 +48,7 @@ def _full_rotary_dim(shim: "GgufConfigShim", full_head_dim: int) -> int: def parse_gguf_config(shim: "GgufConfigShim") -> ModelConfig: + check_scales(shim.model_path, known=("ffn_gate_inp.scale", "ffn_down_exps.scale")) m = shim.metadata def g(key: str): @@ -57,8 +57,11 @@ def g(key: str): raise KeyError(f"missing GGUF metadata key gemma4.{key}") return val + if int(m.get("gemma4.embedding_length_per_layer_input", 0)) or int(m.get("gemma4.attention.shared_kv_layers", 0)): + raise NotImplementedError("gemma4 E2B / E4B GGUF (per-layer input embeddings, shared KV layers) is not supported") num_layers = int(g("block_count")) hidden = int(g("embedding_length")) + num_experts = int(m.get("gemma4.expert_count", 0)) num_qo_heads = int(g("attention.head_count")) kv_per_layer = g("attention.head_count_kv") # per-layer list # True -> sliding-window (SWA) layer, False -> full attention. @@ -105,15 +108,14 @@ def g(key: str): rms_norm_eps=float(g("attention.layer_norm_rms_epsilon")), tie_word_embeddings=bool(shim.tie_word_embeddings), rotary_config=full_rotary, - num_experts=int(g("expert_count")), - num_experts_per_tok=int(g("expert_used_count")), - moe_intermediate_size=int(g("expert_feed_forward_length")), + num_experts=num_experts, + num_experts_per_tok=int(g("expert_used_count")) if num_experts else 0, + moe_intermediate_size=int(g("expert_feed_forward_length")) if num_experts else 0, norm_topk_prob=True, model_type="gemma4", architectures=list(shim.architectures), - moe_enabled=True, - expert_quant="q4_0", - moe_weight_format="q4_0", + moe_enabled=num_experts > 0, + expert_quant="gguf" if num_experts else "none", use_qk_norm=True, attn_sm_scale=1.0, final_logit_softcapping=float(g("final_logit_softcapping")), @@ -144,31 +146,60 @@ def g(key: str): # Weight loading: GGUF tensor names -> FreeToken gemma4 module params. # -------------------------------------------------------------------------------------- -# Per-layer 1:1 norm/scale tensors (gguf suffix -> freetoken module-relative name). All -# are tiny F32 tensors dequantized to bf16. +# Per-layer 1:1 norm/scale tensors (gguf suffix -> freetoken module-relative name), all F32. _LAYER_SCALAR_MAP = { "attn_norm.weight": "input_layernorm.weight", "attn_q_norm.weight": "self_attn.q_norm.weight", "attn_k_norm.weight": "self_attn.k_norm.weight", "post_attention_norm.weight": "post_attention_layernorm.weight", "ffn_norm.weight": "pre_feedforward_layernorm.weight", - "pre_ffw_norm_2.weight": "feed_forward.pre_feedforward_layernorm_2.weight", "post_ffw_norm.weight": "feed_forward.post_feedforward_layernorm.weight", + "layer_output_scale.weight": "feed_forward.layer_scalar", +} +# the MoE layers' second feed-forward branch and router +_MOE_LAYER_SCALAR_MAP = { + "pre_ffw_norm_2.weight": "feed_forward.pre_feedforward_layernorm_2.weight", "post_ffw_norm_1.weight": "feed_forward.post_feedforward_layernorm_1.weight", "post_ffw_norm_2.weight": "feed_forward.post_feedforward_layernorm_2.weight", - "layer_output_scale.weight": "feed_forward.layer_scalar", "ffn_gate_inp.weight": "feed_forward.router.proj.weight", "ffn_gate_inp.scale": "feed_forward.router.scale", "ffn_down_exps.scale": "feed_forward.router.per_expert_scale", } -# Expert weight tensors (handled by the offload bank loader, not iter_gguf_weights). -_EXPERT_SUFFIXES = ("ffn_gate_up_exps.weight", "ffn_down_exps.weight") - -def _to_bf16(t) -> torch.Tensor: - """Dequantize a GgufTensor (F32/F16/Q*) to a dense bf16 tensor of its torch shape.""" - flat = dequantize(t.packed().reshape(-1), t.ggml_type, torch.bfloat16) - return flat.reshape(t.shape) +# checkpoint module leaf (gemma4 HF naming) -> GGUF tensor +_LAYER_LEAVES = { + "self_attn.q_proj": "attn_q", + "self_attn.k_proj": "attn_k", + "self_attn.o_proj": "attn_output", + "mlp.gate_proj": "ffn_gate", + "mlp.up_proj": "ffn_up", + "mlp.down_proj": "ffn_down", + "experts.0.gate_proj": "ffn_gate_exps", + "experts.0.up_proj": "ffn_up_exps", + "experts.0.down_proj": "ffn_down_exps", +} +_LAYER_RE = re.compile(r"^model\.layers\.(\d+)\.(.+)$") + + +def to_gguf(name: str, tensors: dict[str, int]) -> tuple[str, ...]: + """Checkpoint module name -> the GGUF tensors (without ``.weight``) that store it; empty for modules the GGUF has no matrix for.""" + if name == "model.embed_tokens": + return ("token_embd",) + if name == "lm_head": + return ("output",) if "output.weight" in tensors else () + m = _LAYER_RE.match(name) + if m is None: + return () + blk, leaf = f"blk.{m.group(1)}", m.group(2) + if leaf in ("self_attn.q_proj", "self_attn.k_proj", "self_attn.v_proj") and f"{blk}.attn_qkv.weight" in tensors: + return (f"{blk}.attn_qkv",) + if leaf == "self_attn.v_proj": + # the full-attention layers store no attn_v: k is reused as v + return (f"{blk}.attn_v",) if f"{blk}.attn_v.weight" in tensors else (f"{blk}.attn_k",) + if leaf in ("experts.0.gate_proj", "experts.0.up_proj") and f"{blk}.ffn_gate_up_exps.weight" in tensors: + return (f"{blk}.ffn_gate_up_exps",) + tensor = _LAYER_LEAVES.get(leaf) + return (f"{blk}.{tensor}",) if tensor is not None else () def _require_tp1(what: str) -> None: @@ -191,284 +222,51 @@ def iter_gguf_weights( include_moe_experts: bool, include_non_moe: bool, ) -> Iterator[tuple[str, torch.Tensor]]: - """Yield (param_name, tensor) for every non-expert gemma4 param. - - Quantized projections (attention q/k/v/o, shared-MLP gate/up/down) and the - embedding stay in their native packed block layout and are yielded as ``.qweight`` - (uint8); norms / router / per-layer scalars dequantize to bf16. q/k/v and gate/up - are fused by concatenating packed rows along the output dim (same input dim -> - same ``row_bytes``). Routed experts are served from the offload cache, so they are - skipped here (asserts the offload contract like the other MoE models). - """ - from freetoken.models.gguf.reader import iter_gguf_tensors + """The dense state dict from the GGUF; the routed experts load as GGUF expert banks (``freetoken.models.gguf.experts``).""" + if not include_non_moe: + return from freetoken.utils import cached_load_hf_config - assert not include_moe_experts, ( - "gemma4 GGUF stores experts as Q4_0 and only supports the offload backend; " - "experts are loaded into the offload cache via load_q4_0_expert_sources()." - ) - assert include_non_moe _require_tp1("weight loading") - - # Full-attention layers ship no attn_v (k reused as v); SWA layers do. Knowing which - # is which lets us emit the fused qkv as soon as its parts are present. config = parse_gguf_config(cached_load_hf_config(model_path)) - k_eq_v_layers = { - lid - for lid in range(config.num_layers) - if isinstance(config.attention_group_for_layer(lid), FullAttentionGroupConfig) - and config.attention_group_for_layer(lid).k_eq_v - } - - # Per-layer fusion buffers: layer -> {slot: packed[out, row_bytes]}. - qkv_buf: dict[int, dict[str, torch.Tensor]] = {} - gate_up_buf: dict[int, dict[str, torch.Tensor]] = {} - - def layer_of(name: str) -> int: - return int(name.split(".")[1]) - - for t in iter_gguf_tensors(model_path): - name = t.name - if name == "token_embd.weight": - yield "model.embed_tokens.qweight", t.packed() # Q6_K packed table - continue - if name == "output_norm.weight": - yield "model.norm.weight", _to_bf16(t) - continue - if name == "rope_freqs.weight": - continue # rope frequencies recomputed in-engine - if not name.startswith("blk."): - continue - if any(name.endswith(sfx) for sfx in _EXPERT_SUFFIXES): - continue # routed experts -> offload banks - - layer = layer_of(name) - suffix = name.split(".", 2)[2] # after "blk.N." - base = f"model.layers.{layer}" - - if suffix in _LAYER_SCALAR_MAP: - rel = _LAYER_SCALAR_MAP[suffix] - tensor = _to_bf16(t) - if rel.endswith("layer_scalar"): - tensor = tensor.reshape(1) - yield f"{base}.{rel}", tensor - continue - - # Quantized projections: keep packed; fuse q/k/v and gate/up. - if suffix == "attn_q.weight": - qkv_buf.setdefault(layer, {})["q"] = t.packed() - elif suffix == "attn_k.weight": - qkv_buf.setdefault(layer, {})["k"] = t.packed() - elif suffix == "attn_v.weight": - qkv_buf.setdefault(layer, {})["v"] = t.packed() - elif suffix == "ffn_gate.weight": - gate_up_buf.setdefault(layer, {})["gate"] = t.packed() - elif suffix == "ffn_up.weight": - gate_up_buf.setdefault(layer, {})["up"] = t.packed() - elif suffix == "attn_output.weight": - yield f"{base}.self_attn.o_proj.qweight", t.packed() - elif suffix == "ffn_down.weight": - yield f"{base}.feed_forward.shared_mlp.down_proj.qweight", t.packed() + index = gguf_tensor_index(model_path) + hidden = config.hidden_size + + def t(name: str): + return index[name if name in index else f"{name}.weight"] + + embed = t("token_embd") + yield from fused("model.embed_tokens", [embed], [matrix(embed)], hidden) + if "output.weight" in index: + out = t("output") + yield from fused("lm_head", [out], [matrix(out)], hidden) + yield "model.norm.weight", float_tensor(t("output_norm")) + + scalar_map = {**_LAYER_SCALAR_MAP, **_MOE_LAYER_SCALAR_MAP} if config.moe_enabled else _LAYER_SCALAR_MAP + for layer in range(config.num_layers): + blk, base = f"blk.{layer}", f"model.layers.{layer}" + for suffix, rel in scalar_map.items(): + tensor = float_tensor(index[f"{blk}.{suffix}"]) + if suffix == "ffn_down_exps.scale" and weight_scale(index[f"{blk}.ffn_down_exps.weight"]) is not None: + # the converter folds per_expert_scale into the NVFP4 down weight_scale_2, which the expert kernel applies + tensor = torch.ones_like(tensor) + yield f"{base}.{rel}", tensor.reshape(1) if rel.endswith("layer_scalar") else tensor + if f"{blk}.attn_qkv.weight" in index: + # convert_hf_to_gguf.py --fuse-qkv stores q, k and v as one attn_qkv matrix + qkv = [t(f"{blk}.attn_qkv")] else: - raise ValueError(f"unmapped gemma4 GGUF tensor: {name}") - - # Emit fused qkv once all parts are present: k_eq_v full layers reuse k as v - # (no attn_v tensor), SWA layers wait for their separate attn_v. - slots = qkv_buf.get(layer) - if slots is not None and "q" in slots and "k" in slots: - if layer in k_eq_v_layers: - v = slots["k"] - elif "v" in slots: - v = slots["v"] - else: - v = None - if v is not None: - yield f"{base}.self_attn.qkv_proj.qweight", torch.cat( - [slots["q"], slots["k"], v], dim=0 - ) - del qkv_buf[layer] - gu = gate_up_buf.get(layer) - if gu is not None and "gate" in gu and "up" in gu: - yield f"{base}.feed_forward.shared_mlp.gate_up_proj.qweight", torch.cat( - [gu["gate"], gu["up"]], dim=0 - ) - del gate_up_buf[layer] - - assert not qkv_buf, f"incomplete qkv groups: {sorted(qkv_buf)}" - assert not gate_up_buf, f"incomplete gate_up groups: {sorted(gate_up_buf)}" - - -# -------------------------------------------------------------------------------------- -# Model layer swap: dense bf16 Linear/Embedding -> native GGUF-quant ops. -# -------------------------------------------------------------------------------------- - - -def is_gguf_model(config: ModelConfig) -> bool: - """True when the model was parsed from a GGUF checkpoint (native-quant path).""" - return getattr(config, "moe_weight_format", None) == "q4_0" - - -class GGUFTiedLMHead: - """Tied LM head over a native Q6_K embedding table (logits via ggml matmul). - - Holds only a reference to the GGUF embedding (no params of its own -> empty - state_dict), mirroring ``ParallelLMHead`` with ``tie_word_embeddings``. TP=1 only. - """ - - def __init__(self, embedding, quant_type: int): - self._embedding = embedding - self._quant_type = quant_type - - def state_dict(self, *, prefix: str = "", result=None): - return result if result is not None else {} - - def load_state_dict(self, state_dict, *, prefix: str = "", _internal: bool = False): - state_dict.pop(f"{prefix}.weight", None) - state_dict.pop(f"{prefix}.bias", None) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - from freetoken.core import get_global_ctx - from freetoken.layers.gguf import fused_mul_mat_gguf - - batch = get_global_ctx().batch - if batch.is_prefill: - indices = batch.attn_metadata.get_last_indices(batch.size) - x = x[indices].contiguous() - return fused_mul_mat_gguf(x, self._embedding.qweight, self._quant_type) - - -def convert_gemma4_to_gguf(model, config: ModelConfig) -> None: - """In place: replace gemma4's dense projections + embedding with native GGUF ops. - - Quantized in the checkpoint -> swapped: attention qkv/o, shared-MLP gate_up/down - (all Q4_0) and the token embedding (Q6_K, also the tied LM head). Left as dense - bf16 (F32 in the GGUF): the router gate, all RMSNorms, the per-layer scalars, and - the routed experts (served from the offload cache). - """ - from freetoken.layers.gguf import GGUFEmbedding, GGUFLinear - - def swap_linear(owner, attr, quant_type=GGML_Q4_0): - lin = getattr(owner, attr) - out_features, in_features = lin.weight.shape - setattr( - owner, - attr, - GGUFLinear(in_features, out_features, quant_type, has_bias=lin.bias is not None), - ) - - inner = model.model - embed = GGUFEmbedding( - num_embeddings=config.vocab_size, - embedding_dim=config.hidden_size, - quant_type=GGML_Q6_K, - embed_scale=config.embedding_scale, - ) - inner.embed_tokens = embed - - for layer in inner.layers.op_list: - swap_linear(layer.self_attn, "qkv_proj") - swap_linear(layer.self_attn, "o_proj") - swap_linear(layer.feed_forward.shared_mlp, "gate_up_proj") - swap_linear(layer.feed_forward.shared_mlp, "down_proj") - - if config.tie_word_embeddings: - model.lm_head = GGUFTiedLMHead(embed, GGML_Q6_K) - - -# -------------------------------------------------------------------------------------- -# Routed-expert host banks (native Q4_0) for the offload cache. -# -------------------------------------------------------------------------------------- - -def _q4_0_expert_specs(config: ModelConfig) -> dict[str, tuple[tuple[int, ...], torch.dtype]]: - E = config.num_experts - H, I = config.hidden_size, config.moe_intermediate_size - return { - "gate_up": ((E, 2 * I, row_bytes(H, GGML_Q4_0)), torch.uint8), - "down": ((E, H, row_bytes(I, GGML_Q4_0)), torch.uint8), - } - - -def load_q4_0_expert_sources( - model_path: str, config: ModelConfig, *, layer_sink=None -) -> dict[str, list[torch.Tensor]]: - """Per-layer host banks of the routed experts' native Q4_0 block bytes. - - ``gate_up`` is one ``[E, 2I, H//32*18]`` tensor per layer and ``down`` one - ``[E, H, I//32*18]`` per layer (independent :class:`HostBank` allocations) -- each - expert's packed rows verbatim from the GGUF (no dequant), whole layers arriving in - one shot (gate_up + down = 2 writes/layer), so the offload cache streams whole - experts to the ggml MoE kernels. - - ``layer_sink=None`` (serving): pin each layer's two banks as they complete via an - internally-owned :class:`PinPipeline` (or, on a CUDA-less host, allocate the mmap - banks but never pin -- the CPU executor reads them pageable). ``layer_sink`` given - (converter): the completion tracker fires into it instead -- nothing here is pinned, - and the sink may release banks it has written out, so the returned tensors are only - valid until then (the caller owns that tradeoff). - """ - from freetoken.models.gguf.reader import iter_gguf_tensors - from freetoken.moe.host_banks import LayerCompletionTracker, PinPipeline, alloc_layer_banks - - _require_tp1("expert banks") - L, E = config.num_layers, config.num_experts - H, I = config.hidden_size, config.moe_intermediate_size - h_bytes, i_bytes = row_bytes(H, GGML_Q4_0), row_bytes(I, GGML_Q4_0) - hb = alloc_layer_banks(_q4_0_expert_specs(config), L) # lazy anon mmaps (unpinned) - banks = {name: [b.tensor for b in hb[name]] for name in hb} - seen_gu, seen_dn = set(), set() - - def _load(sink) -> None: - tracker = LayerCompletionTracker(2, hb, sink) if sink is not None else None # gate_up + down - for t in iter_gguf_tensors(model_path): - if not t.name.startswith("blk."): - continue - layer = int(t.name.split(".")[1]) - if t.name.endswith("ffn_gate_up_exps.weight"): - banks["gate_up"][layer].copy_(t.packed().reshape(E, 2 * I, h_bytes)) - seen_gu.add(layer) - elif t.name.endswith("ffn_down_exps.weight"): - banks["down"][layer].copy_(t.packed().reshape(E, H, i_bytes)) - seen_dn.add(layer) - else: - continue - if tracker is not None: - tracker.note(layer) - - if layer_sink is not None: - _load(layer_sink) - elif torch.cuda.is_available(): - with PinPipeline() as pins: - _load(pins) - else: - _load(None) # CUDA-less: mmap banks stay pageable, never pinned - - want = set(range(L)) - assert seen_gu == want and seen_dn == want, ( - f"missing Q4_0 expert layers: gate_up {sorted(want - seen_gu)}, " - f"down {sorted(want - seen_dn)}" - ) - return banks - - -def dummy_q4_0_expert_sources(config: ModelConfig) -> dict[str, list[torch.Tensor]]: - """Random Q4_0 expert banks shaped like ``load_q4_0_expert_sources`` output.""" - from freetoken.moe.host_banks import alloc_layer_banks, pin_banks - - L = config.num_layers - hb = alloc_layer_banks(_q4_0_expert_specs(config), L) - banks = {name: [b.tensor for b in hb[name]] for name in hb} - for t in banks["gate_up"] + banks["down"]: - t.random_(0, 256) - if torch.cuda.is_available(): - pin_banks(hb) # match the other dummies: pin-after-fill (no-op mmap fill on CPU-only) - return banks + q, k = t(f"{blk}.attn_q"), t(f"{blk}.attn_k") + qkv = [q, k, index.get(f"{blk}.attn_v.weight", k)] + yield from fused(f"{base}.self_attn.qkv_proj", qkv, [matrix(p) for p in qkv], hidden) + o = t(f"{blk}.attn_output") + yield from fused(f"{base}.self_attn.o_proj", [o], [matrix(o)], int(o.ne[0])) + gate, up, down = (t(f"{blk}.ffn_{p}") for p in ("gate", "up", "down")) + yield from fused(f"{base}.feed_forward.shared_mlp.gate_up_proj", [gate, up], [matrix(gate), matrix(up)], hidden) + yield from fused(f"{base}.feed_forward.shared_mlp.down_proj", [down], [matrix(down)], int(down.ne[0])) __all__ = [ "parse_gguf_config", "iter_gguf_weights", - "convert_gemma4_to_gguf", - "is_gguf_model", - "load_q4_0_expert_sources", - "dummy_q4_0_expert_sources", + "to_gguf", ] diff --git a/python/freetoken/models/gemma4/model.py b/python/freetoken/models/gemma4/model.py index 5c91b90571..1b9d0da7d9 100644 --- a/python/freetoken/models/gemma4/model.py +++ b/python/freetoken/models/gemma4/model.py @@ -61,6 +61,8 @@ def __init__(self, config: ModelConfig, *, prefix: str = "model"): num_embeddings=config.vocab_size, embedding_dim=config.hidden_size, embed_scale=config.embedding_scale, + quant_config=config.quant, + prefix=f"{prefix}.embed_tokens", ) self.layers = OPList( [ @@ -91,13 +93,6 @@ def __init__(self, config: ModelConfig): self._final_logit_softcapping = config.final_logit_softcapping super().__init__() - # GGUF checkpoints carry native block-quantized weights: swap the dense - # projections + embedding for GGUF-quant ops (experts stay on the offload cache). - from .gguf import convert_gemma4_to_gguf, is_gguf_model - - if is_gguf_model(config): - convert_gemma4_to_gguf(self, config) - def forward(self) -> torch.Tensor: output = self.model.forward(get_global_ctx().batch.input_ids) logits = self.lm_head.forward(output) diff --git a/python/freetoken/models/gemma4/weight.py b/python/freetoken/models/gemma4/weight.py index ba990c007b..7797c7842e 100644 --- a/python/freetoken/models/gemma4/weight.py +++ b/python/freetoken/models/gemma4/weight.py @@ -66,15 +66,15 @@ def _nvfp4_dense_parts(reader: ShardReader, raw_base: str): - """Load an NVFP4 dense weight as the NVFP4 linear method's buffers: weight uint8 [O, IN//2], weight_scale fp8-e4m3 block [O, IN//16], weight_global fp16 [O] (the per-tensor weight_scale_2 per output row), input_scale fp32 scalar or None when the export has none.""" + """Load an NVFP4 dense weight as the NVFP4 linear method's buffers: weight uint8 [O, IN//2], weight_scale fp8-e4m3 block [O, IN//16], weight_global fp32 [O] (the per-tensor weight_scale_2 per output row), input_scale fp32 scalar or None when the export has none.""" w = reader.get_tensor(raw_base + ".weight") s = reader.get_tensor(raw_base + ".weight_scale") - g = reader.get_tensor(raw_base + ".weight_scale_2").reshape(1).to(torch.float16) + g = reader.get_tensor(raw_base + ".weight_scale_2").reshape(1).to(torch.float32) g = g.expand(w.shape[0]).contiguous() assert ( w.dtype is torch.uint8 and s.dtype is torch.float8_e4m3fn - and g.dtype is torch.float16 + and g.dtype is torch.float32 ), f"unexpected NVFP4 dense dtypes at {raw_base}: {w.dtype}/{s.dtype}/{g.dtype}" a = reader.get_tensor(raw_base + ".input_scale").reshape(()).to(torch.float32) if reader.has(raw_base + ".input_scale") else None return w, s, g, a @@ -196,7 +196,7 @@ def merge_info(key: str) -> tuple[str, MergeRule] | None: continue # Per-expert NVFP4 tensors go to the offload cache (expert pieces), - # not this dense pass; fused bf16/q4_0 experts lack ".experts.." so are unaffected. + # not this dense pass; fused bf16 experts lack ".experts.." so are unaffected. if _NVFP4_EXPERT_RE.search(raw_name): continue diff --git a/python/freetoken/models/gguf/__init__.py b/python/freetoken/models/gguf/__init__.py index 75e40e4faa..20a97b4610 100644 --- a/python/freetoken/models/gguf/__init__.py +++ b/python/freetoken/models/gguf/__init__.py @@ -1,26 +1,19 @@ -from .dequant import GGML_NAME, dequantize from .reader import ( FTW_METADATA_GGUF, - GgufTensor, gguf_architecture, gguf_config_source, gguf_tensor_names, is_gguf_path, - iter_gguf_tensors, load_gguf_metadata, write_metadata_gguf, ) __all__ = [ - "GGML_NAME", - "dequantize", "FTW_METADATA_GGUF", - "GgufTensor", "gguf_architecture", "gguf_config_source", "gguf_tensor_names", "is_gguf_path", - "iter_gguf_tensors", "load_gguf_metadata", "write_metadata_gguf", ] diff --git a/python/freetoken/models/gguf/config.py b/python/freetoken/models/gguf/config.py index 63b1a18b97..5767d4eb1c 100644 --- a/python/freetoken/models/gguf/config.py +++ b/python/freetoken/models/gguf/config.py @@ -18,6 +18,20 @@ # reuses the model classes but a GGUF parse_config / iter_weights). GGUF_ARCH_TO_REGISTRY: dict[str, str] = { "gemma4": "Gemma4GGUFForCausalLM", + "glm4moe": "Glm4MoeGGUFForCausalLM", + # glm5-next: the llama.cpp spelling; glm5next: unsloth's first release + "glm5-next": "Glm5NextGGUFForCausalLM", + "glm5next": "Glm5NextGGUFForCausalLM", + "llama": "LlamaGGUFForCausalLM", + "minimax-m2": "MiniMaxM2GGUFForCausalLM", + "minimax-m3": "MiniMaxM3GGUFForCausalLM", + "muse-glimmer": "MuseGlimmerGGUFForCausalLM", + "qwen2": "Qwen2GGUFForCausalLM", + "qwen3": "Qwen3GGUFForCausalLM", + "qwen3moe": "Qwen3MoeGGUFForCausalLM", + "qwen35": "Qwen3_5GGUFForCausalLM", + "qwen35moe": "Qwen3_5MoeGGUFForCausalLM", + "qwen4exp": "Qwen4ExpGGUFForCausalLM", } diff --git a/python/freetoken/models/gguf/dequant.py b/python/freetoken/models/gguf/dequant.py deleted file mode 100644 index 77c3ea0102..0000000000 --- a/python/freetoken/models/gguf/dequant.py +++ /dev/null @@ -1,153 +0,0 @@ -"""GGML block-quant dequantization in pure torch (the formats this repo's GGUF -checkpoints use: Q4_0, Q6_K, plus trivial F32/F16/BF16). - -This is the *reference / CPU* path, NOT the engine's hot path: GGUF weights stay -packed and are dequantized inside the borrowed ggml CUDA kernels (see -``freetoken.kernel.gguf``). These routines are used only to (a) materialize the few -dense F32/F16 tensors at load (norms, scales, router) via :func:`dequantize`, and -(b) cross-check the CUDA kernels in tests. The ``BLOCK_SHAPE`` table and -:func:`row_bytes` are the type metadata the packed (kernel) path also relies on. - -Each ``dequant_*`` takes the raw little-endian bytes as a ``uint8`` tensor whose -final axis spans whole blocks, and returns the values in *storage order* (ggml's -fastest axis first); the caller reshapes to the torch shape (``dims[::-1]``). The -math mirrors ``ggml-quants.c``. -""" - -from __future__ import annotations - -import torch - -# ggml_type enum values (subset present in these checkpoints). -GGML_F32 = 0 -GGML_F16 = 1 -GGML_Q4_0 = 2 -GGML_Q8_0 = 8 -GGML_Q6_K = 14 -GGML_BF16 = 30 - -# (block numel, bytes per block) per ggml type. -BLOCK_SHAPE: dict[int, tuple[int, int]] = { - GGML_F32: (1, 4), - GGML_F16: (1, 2), - GGML_BF16: (1, 2), - GGML_Q4_0: (32, 18), - GGML_Q8_0: (32, 34), - GGML_Q6_K: (256, 210), -} - -GGML_NAME = { - GGML_F32: "F32", - GGML_F16: "F16", - GGML_BF16: "BF16", - GGML_Q4_0: "Q4_0", - GGML_Q8_0: "Q8_0", - GGML_Q6_K: "Q6_K", -} - - -def row_bytes(numel: int, ggml_type: int) -> int: - """Packed byte length of one row of ``numel`` elements in ``ggml_type`` blocks. - - Single source of truth for the ``numel // block * type_size`` math shared by the - packed-weight ops (``GGUFLinear``/``GGUFEmbedding``) and the expert bank loaders. - """ - block, type_size = BLOCK_SHAPE[ggml_type] - assert numel % block == 0, ( - f"{numel} not a multiple of block {block} for {GGML_NAME.get(ggml_type, ggml_type)}" - ) - return numel // block * type_size - - -def _f16_scales(raw: torch.Tensor, lo: int, hi: int) -> torch.Tensor: - """Reinterpret bytes ``[lo:hi]`` (2 per block) of each block row as fp16 -> fp32 [N,1].""" - return raw[:, lo:hi].contiguous().view(torch.float16).to(torch.float32) - - -def dequant_q4_0(raw: torch.Tensor, out_dtype: torch.dtype) -> torch.Tensor: - """Q4_0: per 32-elem block = fp16 scale ``d`` + 16 packed nibbles; ``w = d*(q-8)``. - - Byte ``j`` of the 16 holds element ``j`` in its low nibble and ``j+16`` in its high - nibble, so storage order within the block is ``[lo0..lo15, hi0..hi15]``. - """ - raw = raw.reshape(-1, 18) - d = _f16_scales(raw, 0, 2) # [N,1] - qs = raw[:, 2:18] # [N,16] uint8 - lo = (qs & 0x0F).to(torch.float32) - hi = (qs >> 4).to(torch.float32) - q = torch.cat([lo, hi], dim=1) # [N,32] - return ((q - 8.0) * d).reshape(-1).to(out_dtype) - - -def dequant_q6_k(raw: torch.Tensor, out_dtype: torch.dtype) -> torch.Tensor: - """Q6_K: 256-elem super-block = 128B low nibbles + 64B high 2-bits + 16 int8 - sub-scales + fp16 ``d``. Direct vectorization of ggml's two-half loop.""" - raw = raw.reshape(-1, 210) - n = raw.shape[0] - ql = raw[:, 0:128] # [n,128] - qh = raw[:, 128:192] # [n,64] - sc = raw[:, 192:208].view(torch.int8).to(torch.float32) # [n,16] - d = _f16_scales(raw, 208, 210) # [n,1] - - y = torch.empty((n, 256), dtype=torch.float32, device=raw.device) - # l in 0..15 -> is=0; l in 16..31 -> is=1 (per ggml: is = l/16). - is_idx = (torch.arange(32, device=raw.device) // 16) # [32] in {0,1} - for h in range(2): # two 128-elem halves of the super-block - qlh = ql[:, h * 64:(h + 1) * 64] # [n,64] - qhh = qh[:, h * 32:(h + 1) * 32] # [n,32] - sch = sc[:, h * 8:(h + 1) * 8] # [n,8] - a = qlh[:, 0:32].to(torch.int32) # ql[l] - b = qlh[:, 32:64].to(torch.int32) # ql[l+32] - hb = qhh.to(torch.int32) # qh[l] - q1 = ((a & 0x0F) | (((hb >> 0) & 3) << 4)) - 32 - q2 = ((b & 0x0F) | (((hb >> 2) & 3) << 4)) - 32 - q3 = ((a >> 4) | (((hb >> 4) & 3) << 4)) - 32 - q4 = ((b >> 4) | (((hb >> 6) & 3) << 4)) - 32 - s1 = sch.index_select(1, is_idx + 0).to(torch.float32) - s2 = sch.index_select(1, is_idx + 2).to(torch.float32) - s3 = sch.index_select(1, is_idx + 4).to(torch.float32) - s4 = sch.index_select(1, is_idx + 6).to(torch.float32) - base = h * 128 - y[:, base + 0:base + 32] = d * s1 * q1.to(torch.float32) - y[:, base + 32:base + 64] = d * s2 * q2.to(torch.float32) - y[:, base + 64:base + 96] = d * s3 * q3.to(torch.float32) - y[:, base + 96:base + 128] = d * s4 * q4.to(torch.float32) - return y.reshape(-1).to(out_dtype) - - -_DEQUANT = { - GGML_Q4_0: dequant_q4_0, - GGML_Q6_K: dequant_q6_k, -} - - -def dequantize(raw: torch.Tensor, ggml_type: int, out_dtype: torch.dtype) -> torch.Tensor: - """Dequantize ``raw`` (uint8) of any supported ggml type to flat ``out_dtype``.""" - if ggml_type == GGML_F32: - return raw.view(torch.float32).to(out_dtype) - if ggml_type == GGML_F16: - return raw.view(torch.float16).to(out_dtype) - if ggml_type == GGML_BF16: - return raw.view(torch.bfloat16).to(out_dtype) - fn = _DEQUANT.get(ggml_type) - if fn is None: - raise NotImplementedError( - f"dequant for ggml type {GGML_NAME.get(ggml_type, ggml_type)} not implemented" - ) - return fn(raw, out_dtype) - - -__all__ = [ - "GGML_F32", - "GGML_F16", - "GGML_BF16", - "GGML_Q4_0", - "GGML_Q8_0", - "GGML_Q6_K", - "GGML_NAME", - "BLOCK_SHAPE", - "row_bytes", - "dequant_q4_0", - "dequant_q6_k", - "dequantize", -] diff --git a/python/freetoken/models/gguf/experts.py b/python/freetoken/models/gguf/experts.py new file mode 100644 index 0000000000..0efb5982fb --- /dev/null +++ b/python/freetoken/models/gguf/experts.py @@ -0,0 +1,131 @@ +"""Routed-expert banks straight from a GGUF's ``blk.N.ffn_{gate,up,down}_exps`` tensors, in their packed ggml rows: pinned host banks for the offload cache, or device banks for resident experts. + +Each layer keeps its own ggml types, so each layer gets its own bank shapes: ``gate_up`` [E, 2I, row_bytes] (expert e's gate rows, then its up rows) and ``down`` [E, H, row_bytes]; float experts keep their dtype ([E, 2I, H] and [E, H, I]). +A GGUF that stores ``ffn_gate_up_exps`` (gemma4) already holds the gate_up rows in that order. +When NVFP4 experts carry per-expert ``.scale`` tensors, every layer also gets ``global`` [E, 4] fp32: the (gate, up, down, 1) weight_scale_2, 1 where a tensor has none.""" + +from __future__ import annotations + +import contextlib +import os + +import torch + +from freetoken.layers.quantization import QuantKind +from freetoken.layers.quantization.ggml_types import float_dtype, has_expert_scales, mmq_tail_bytes, pad_rows, type_name +from freetoken.utils import init_logger + +from .reader import GgufTensorInfo, gguf_tensor_index +from .tensors import weight_scale + +logger = init_logger(__name__) + + +def _expert_rows(info: GgufTensorInfo) -> torch.Tensor: + """[E, rows, row_bytes] packed ggml rows, or [E, rows, K] for a float tensor.""" + k, rows, experts = info.ne + raw = torch.from_numpy(info.data()).view(experts, rows, info.nbytes // (experts * rows)) + dtype = float_dtype(info.ggml_type) + return raw if dtype is None else raw.view(dtype) + + +def _drop_page_cache(info: GgufTensorInfo) -> None: + try: + fd = os.open(info.path, os.O_RDONLY) + try: + os.posix_fadvise(fd, info.offset, info.nbytes, os.POSIX_FADV_DONTNEED) + finally: + os.close(fd) + except OSError: + pass + + +def _device_bank(shape: tuple[int, ...], ggml_type: int, k: int, dtype: torch.dtype, device: torch.device) -> torch.Tensor: + """An [E, rows, row] device bank; a zero row follows the last one when ggml's MMQ tiles read past it.""" + experts, rows, row = shape + flat = torch.empty(experts * rows + pad_rows(ggml_type, k), row, dtype=dtype, device=device) + flat[experts * rows :].zero_() + return flat[: experts * rows].view(shape) + + +def _expert_scales(gate_up: list[GgufTensorInfo], down: GgufTensorInfo, experts: int) -> torch.Tensor: + """[E, 4] fp32 (gate, up, down, 1) weight_scale_2 of one layer's experts.""" + gate, up = gate_up if len(gate_up) == 2 else (gate_up[0], gate_up[0]) + out = torch.ones(experts, 4, dtype=torch.float32) + for col, info in enumerate((gate, up, down)): + s = weight_scale(info) + if s is None: + continue + if s.numel() != experts: + raise NotImplementedError(f"{info.name}: {s.numel()} scales for {experts} experts") + out[:, col] = s + return out + + +def load_gguf_expert_banks(model_path: str, num_layers: int, *, first_layer: int = 0, dummy: bool = False, device: torch.device | None = None): + """``ExpertBanks`` with one bank per (layer, role): pinned host banks for the offload cache, filled from the file and pinned while later layers load, or resident banks on ``device``.""" + from freetoken.moe.expert_banks import ExpertBanks + from freetoken.moe.host_banks import HostBank, PinPipeline + + index = gguf_tensor_index(model_path) + scaled = has_expert_scales({name: info.ggml_type for name, info in index.items()}) + roles = ("gate_up", "down", "global") if scaled else ("gate_up", "down") + sources: dict[str, list[torch.Tensor]] = {role: [] for role in roles} + tails: dict[str, list[int]] = {role: [] for role in roles} + tags = [] + with PinPipeline() if device is None else contextlib.nullcontext() as pins: + for layer in range(num_layers): + blk = f"blk.{first_layer + layer}" + down = index[f"{blk}.ffn_down_exps.weight"] + if f"{blk}.ffn_gate_up_exps.weight" in index: + gate = index[f"{blk}.ffn_gate_up_exps.weight"] + parts = [gate] + else: + gate, up = index[f"{blk}.ffn_gate_exps.weight"], index[f"{blk}.ffn_up_exps.weight"] + if gate.ggml_type != up.ggml_type: + raise NotImplementedError(f"{blk}: gate and up experts have different ggml types") + parts = [gate, up] + tags.append(("ggml", type_name(gate.ggml_type), type_name(down.ggml_type))) + tails["gate_up"].append(mmq_tail_bytes(gate.ggml_type, gate.ne[0])) + tails["down"].append(mmq_tail_bytes(down.ggml_type, down.ne[0])) + if scaled: + tails["global"].append(0) + rows = [_expert_rows(info) for info in parts] + d = _expert_rows(down) + gate_up_shape = (rows[0].shape[0], sum(r.shape[1] for r in rows), rows[0].shape[2]) + experts = rows[0].shape[0] + if device is None: + host = {"gate_up": HostBank(gate_up_shape, rows[0].dtype), "down": HostBank(tuple(d.shape), d.dtype)} + if scaled: + host["global"] = HostBank((experts, 4), torch.float32) + banks = {role: bank.tensor for role, bank in host.items()} + else: + banks = { + "gate_up": _device_bank(gate_up_shape, gate.ggml_type, gate.ne[0], rows[0].dtype, device), + "down": _device_bank(tuple(d.shape), down.ggml_type, down.ne[0], d.dtype, device), + } + if scaled: + banks["global"] = torch.empty(experts, 4, dtype=torch.float32, device=device) + if scaled: + banks["global"].copy_(_expert_scales(parts, down, experts)) + if dummy: + banks["gate_up"].zero_() + banks["down"].zero_() + else: + start = 0 + for r in rows: + banks["gate_up"][:, start : start + r.shape[1]].copy_(r) + start += r.shape[1] + banks["down"].copy_(d) + for info in (*parts, down): + _drop_page_cache(info) + if device is None: + pins(layer, host) + for role, bank in banks.items(): + sources[role].append(bank) + total = sum(t.numel() * t.element_size() for ts in sources.values() for t in ts) + logger.info_rank0(f"GGUF expert banks: {num_layers} layers, {total / 2**30:.1f} GiB {'on ' + str(device) if device is not None else 'in host RAM'}") + return ExpertBanks("gguf_scaled" if scaled else "gguf", sources, kind=QuantKind.GGUF, kernel="ggml", layout=None, layer_tags=tags, row_tail_bytes=tails) + + +__all__ = ["load_gguf_expert_banks"] diff --git a/python/freetoken/models/gguf/reader.py b/python/freetoken/models/gguf/reader.py index b950d92981..ff7350941e 100644 --- a/python/freetoken/models/gguf/reader.py +++ b/python/freetoken/models/gguf/reader.py @@ -2,22 +2,20 @@ Thin layer over ``gguf.GGUFReader`` (gguf-py). Metadata is read into a plain dict keyed by the GGUF field name (``general.architecture``, ``gemma4.block_count`` ...); -tensors are exposed as ``GgufTensor`` records carrying the *torch* shape (ggml dims -reversed), the ggml quant type, and a zero-copy ``uint8`` view of the packed block -bytes laid out as ``[rows, row_bytes]`` (rows = product of all but the fastest ggml -dim; row_bytes spans whole quant blocks of the fastest dim). +tensors are indexed as ``GgufTensorInfo`` records (ggml dims, ggml type, and where the +packed bytes sit in the memory-mapped file). """ from __future__ import annotations import functools import os +import re import struct from dataclasses import dataclass -from typing import Any, Iterator +from typing import Any import numpy as np -import torch def is_gguf_path(model_path: str) -> bool: @@ -65,7 +63,7 @@ def write_metadata_gguf(source_gguf: str, dest_path: str) -> None: Validates by re-parsing: the copy must list zero tensors and expose the identical KV key set (the KV *bytes* are copied verbatim, so identical keys imply identical values). """ - import gguf + from freetoken.third_party import gguf reader = gguf.GGUFReader(source_gguf) assert reader.tensors, f"{source_gguf}: no tensors to bound the KV section" @@ -98,20 +96,6 @@ def write_metadata_gguf(source_gguf: str, dest_path: str) -> None: ) -@dataclass(frozen=True) -class GgufTensor: - name: str - shape: tuple[int, ...] # torch order (ggml dims reversed) - ggml_type: int - rows: int # product of shape[:-1] over the *ggml* layout = blocks-major rows - row_bytes: int # packed bytes per row (whole quant blocks of the fastest dim) - _raw: np.ndarray # uint8 view, shape [rows, row_bytes] - - def packed(self) -> torch.Tensor: - """Zero-copy ``[rows, row_bytes]`` uint8 tensor of the native block bytes.""" - return torch.from_numpy(self._raw) - - def _field_value(reader, name: str) -> Any: field = reader.fields.get(name) if field is None: @@ -121,7 +105,7 @@ def _field_value(reader, name: str) -> Any: @functools.cache def _reader(model_path: str): - import gguf + from freetoken.third_party import gguf return gguf.GGUFReader(model_path) @@ -133,6 +117,57 @@ def load_gguf_metadata(model_path: str) -> dict[str, Any]: return {name: field.contents() for name, field in reader.fields.items()} +_SPLIT_RE = re.compile(r"-(\d{5})-of-(\d{5})\.gguf$") + + +@functools.cache +def gguf_split_paths(model_path: str) -> tuple[str, ...]: + """Every file of a split GGUF (``name-00001-of-00002.gguf``), in order; a single file is its own list.""" + count = int(_field_value(_reader(model_path), "split.count") or 1) + if count == 1: + return (model_path,) + match = _SPLIT_RE.search(model_path) + if match is None: + raise ValueError(f"{model_path} is one of {count} GGUF splits but is not named like name-00001-of-{count:05d}.gguf") + stem = model_path[: match.start()] + paths = tuple(f"{stem}-{i:05d}-of-{count:05d}.gguf" for i in range(1, count + 1)) + missing = [p for p in paths if not os.path.isfile(p)] + if missing: + raise FileNotFoundError(f"GGUF split files missing: {missing}") + return paths + + +@dataclass(frozen=True) +class GgufTensorInfo: + """Where one tensor's packed bytes sit: ``ne`` is in ggml order (fastest dim first).""" + + name: str + path: str + offset: int + nbytes: int + ggml_type: int + ne: tuple[int, ...] + # the model path the index was built from, which resolves tensors in any split + root: str = "" + + def data(self) -> np.ndarray: + """Zero-copy uint8 view of the bytes (memory-mapped).""" + reader = _reader(self.path) + return reader.data[self.offset : self.offset + self.nbytes] + + +@functools.cache +def gguf_tensor_index(model_path: str) -> dict[str, GgufTensorInfo]: + """Every tensor of every split, by name.""" + index: dict[str, GgufTensorInfo] = {} + for path in gguf_split_paths(model_path): + for t in _reader(path).tensors: + index[t.name] = GgufTensorInfo( + t.name, path, int(t.data_offset), int(t.n_bytes), int(t.tensor_type), tuple(int(x) for x in t.shape), model_path + ) + return index + + def gguf_architecture(model_path: str) -> str: arch = _field_value(_reader(model_path), "general.architecture") if arch is None: @@ -140,39 +175,8 @@ def gguf_architecture(model_path: str) -> str: return str(arch) -def iter_gguf_tensors(model_path: str) -> Iterator[GgufTensor]: - """Yield every tensor with its torch shape, ggml type, and packed block bytes.""" - import gguf - - reader = _reader(model_path) - for t in reader.tensors: - ne = [int(s) for s in t.shape] # ggml order, fastest dim first - torch_shape = tuple(reversed(ne)) - block, type_size = gguf.GGML_QUANT_SIZES[t.tensor_type] - n_fast = ne[0] - if n_fast % block != 0: - raise ValueError( - f"{t.name}: fastest dim {n_fast} not a multiple of block {block} " - f"for {t.tensor_type.name}" - ) - row_bytes = n_fast // block * type_size - rows = int(np.prod(ne[1:])) if len(ne) > 1 else 1 - # gguf-py returns quantized tensors as raw uint8 but F32/F16 as typed arrays; - # normalize everything to a flat byte view before shaping into [rows, row_bytes]. - flat = np.ascontiguousarray(t.data).reshape(-1).view(np.uint8) - raw = flat.reshape(rows, row_bytes) - yield GgufTensor( - name=t.name, - shape=torch_shape, - ggml_type=int(t.tensor_type), - rows=rows, - row_bytes=row_bytes, - _raw=raw, - ) - - def gguf_tensor_names(model_path: str) -> set[str]: - return {t.name for t in _reader(model_path).tensors} + return {t.name for path in gguf_split_paths(model_path) for t in _reader(path).tensors} __all__ = [ @@ -181,9 +185,10 @@ def gguf_tensor_names(model_path: str) -> set[str]: "OUTPUT_WEIGHT_PRESENT_KV", "gguf_config_source", "write_metadata_gguf", - "GgufTensor", + "GgufTensorInfo", + "gguf_split_paths", + "gguf_tensor_index", "load_gguf_metadata", "gguf_architecture", - "iter_gguf_tensors", "gguf_tensor_names", ] diff --git a/python/freetoken/models/gguf/tensors.py b/python/freetoken/models/gguf/tensors.py new file mode 100644 index 0000000000..f2cb0e0038 --- /dev/null +++ b/python/freetoken/models/gguf/tensors.py @@ -0,0 +1,100 @@ +"""GGUF tensors as state-dict entries: float tensors in torch order, quantized ones as packed ggml rows grouped into the GGUF Linear's per-type segments.""" + +from __future__ import annotations + +from typing import Iterator + +import numpy as np +import torch + +from freetoken.layers.quantization.ggml_types import NVFP4, is_quantized, pad_rows, scale_name, segment_param, segments + +from .reader import GgufTensorInfo, gguf_tensor_index + +_FLOAT_DTYPES = {0: torch.float32, 1: torch.float16, 30: torch.bfloat16} + + +def float_tensor(info: GgufTensorInfo) -> torch.Tensor: + """A float tensor in torch order (ggml dims reversed).""" + return torch.from_numpy(np.array(info.data())).view(_FLOAT_DTYPES[info.ggml_type]).view(*reversed(info.ne)) + + +def packed_rows(info: GgufTensorInfo) -> torch.Tensor: + """Packed ggml rows ``[rows, row_bytes]`` of a quantized tensor.""" + rows = int(np.prod(info.ne[1:])) + return torch.from_numpy(info.data()).view(rows, info.nbytes // rows) + + +def matrix(info: GgufTensorInfo) -> torch.Tensor: + return packed_rows(info) if is_quantized(info.ggml_type) else float_tensor(info) + + +def check_scales(model_path: str, *, nvfp4: bool = True, known: tuple[str, ...] = ()) -> None: + """Refuse ``.scale`` / ``.input_scale`` tensors that nothing reads. + + With ``nvfp4``, an NVFP4 tensor's ``.scale`` (its weight_scale_2) is read and its ``.input_scale`` goes unused, as in llama.cpp. + ``known`` lists tensor-name suffixes the family's loader reads itself.""" + types = {name: info.ggml_type for name, info in gguf_tensor_index(model_path).items()} + unread = [] + for name in types: + weight = name.removesuffix(".input_scale") if name.endswith(".input_scale") else name.removesuffix(".scale") + if weight == name or name.endswith(known): + continue + if not (nvfp4 and types.get(f"{weight}.weight", types.get(weight)) == NVFP4): + unread.append(name) + if unread: + what = "scale tensors outside NVFP4 weights" if nvfp4 else "tensors with per-tensor scales" + raise NotImplementedError(f"{model_path}: GGUF {what} are not supported (e.g. {sorted(unread)[0]})") + + +def weight_scale(info: GgufTensorInfo) -> torch.Tensor | None: + """The fp32 weight_scale_2 llama.cpp's converter writes next to an NVFP4 tensor as ``.scale`` (one value, or one per expert); None if there is none.""" + if info.ggml_type != NVFP4: + return None + scale = gguf_tensor_index(info.root or info.path).get(scale_name(info.name)) + return None if scale is None else float_tensor(scale).float().reshape(-1) + + +def _row_scales(prefix: str, infos: list[GgufTensorInfo], parts: list[torch.Tensor], pad: int) -> torch.Tensor | None: + scales = [weight_scale(i) for i in infos] + if all(s is None for s in scales): + return None + rows = [] + for info, part, s in zip(infos, parts, scales): + # llama.cpp reads a dense tensor's scale as one value + if s is not None and s.numel() != 1: + raise NotImplementedError(f"{prefix}: {info.name} has {s.numel()} scales, expected one") + rows.append(torch.ones(part.shape[0]) if s is None else s.expand(part.shape[0])) + return torch.cat([*rows, torch.ones(pad)]) + + +def fused(prefix: str, infos: list[GgufTensorInfo], parts: list[torch.Tensor], k: int, pad_to: int = 0) -> Iterator[tuple[str, torch.Tensor]]: + """One fused module: float parts concatenated into ``weight``; with any quantized part, one matrix per ggml-type run (float runs as bf16 rows, quantized runs packed). + + NVFP4 parts with a ``.scale`` add ``ggml_scale``: each part's scale on its output rows, fp32.""" + types = [i.ggml_type for i in infos] + pad = (-sum(p.shape[0] for p in parts)) % pad_to if pad_to else 0 + if not any(is_quantized(t) for t in types): + rows = [p.to(torch.bfloat16) for p in parts] + if pad: + rows.append(torch.zeros(pad, *rows[0].shape[1:], dtype=torch.bfloat16)) + yield f"{prefix}.weight", torch.cat(rows, dim=0) + return + segs = segments(tuple(types), tuple(p.shape[0] for p in parts)) + for i, seg in enumerate(segs): + rows = parts[seg.first_part : seg.first_part + seg.num_parts] + zeros = pad if i == len(segs) - 1 else 0 + if is_quantized(seg.ggml_type): + zeros += pad_rows(seg.ggml_type, k) + dtype = torch.uint8 + else: + rows = [r.to(torch.bfloat16) for r in rows] + dtype = torch.bfloat16 + tail = [torch.zeros(zeros, rows[0].shape[1], dtype=dtype)] if zeros else [] + yield f"{prefix}.{segment_param(i)}", torch.cat([*rows, *tail], dim=0) + scales = _row_scales(prefix, infos, parts, pad) + if scales is not None: + yield f"{prefix}.ggml_scale", scales + + +__all__ = ["check_scales", "float_tensor", "fused", "matrix", "packed_rows", "weight_scale"] diff --git a/python/freetoken/models/gguf/tokenizer.py b/python/freetoken/models/gguf/tokenizer.py index 6d5481c177..26e447aad2 100644 --- a/python/freetoken/models/gguf/tokenizer.py +++ b/python/freetoken/models/gguf/tokenizer.py @@ -1,74 +1,119 @@ """Build a HF fast tokenizer from a GGUF file's embedded tokenizer metadata. -transformers' ``AutoTokenizer.from_pretrained(gguf_file=...)`` first builds the HF -config, which the gemma4 strict dataclass rejects (per-layer ``num_key_value_heads`` -array). So we call the GGUF->fast tokenizer converter directly on the -``tokenizer.ggml.*`` metadata, bypassing config entirely. +transformers' ``AutoTokenizer.from_pretrained(gguf_file=...)`` first builds the HF model config from the GGUF, which needs the PyPI ``gguf`` package and an architecture transformers knows (``qwen4exp`` is not one). +So we run only its tokenizer half: ``get_gguf_tokenizer`` and ``convert_gguf_tokenizer``. """ from __future__ import annotations -from typing import Any +import json -from .reader import gguf_architecture, load_gguf_metadata +from .reader import load_gguf_metadata -# GGUF architecture -> transformers GGUF tokenizer-converter key. -_TOKENIZER_ARCH = {"gemma4": "gemma4_text"} +# end-of-generation token names from llama.cpp (llama_vocab::impl::load in src/llama-vocab.cpp): a GGUF stores one eos id and no stop list +_EOG_NAMES = frozenset(( + "<|eot_id|>", "<|im_end|>", "<|end|>", "<|return|>", "<|call|>", "<|flush|>", "<|calls|>", "", + "<|endoftext|>", "", "<|eom_id|>", "", "_", "[EOT]", "[EOS]", "<|end_of_text|>", "", + "", "", "<|tool_response>", "<\uff5cend\u2581of\u2581sentence\uff5c>", "[e~[", "<|ifm|im_end|>", +)) + +# BPE pre-tokenizers that llama.cpp adds a bos for when the file has no add_bos_token key (llama_vocab::impl::load) +_BOS_PRE = frozenset(( + "llama3", "llama-v3", "llama-bpe", "falcon3", "falcon-h1", "pixtral", "midm-2.0", "lfm2", "jina-v5-nano", "tekken", "chameleon", +)) + +# Llama 3 and GLM pre-tokenizers, whose HF tokenizer and llama.cpp take a word that is one vocabulary token as is instead of merging it +_IGNORE_MERGES_PRE = frozenset(("llama3", "llama-v3", "llama-bpe", "glm4", "glm5")) + + +def _add_bos(meta: dict) -> bool: + if "tokenizer.ggml.add_bos_token" in meta: + return bool(meta["tokenizer.ggml.add_bos_token"]) + model = meta.get("tokenizer.ggml.model") + return model == "llama" or (model == "gpt2" and meta.get("tokenizer.ggml.pre") in _BOS_PRE) + + +def _fix_added_tokens(fast, meta: dict): + """Register the control and user-defined tokens that transformers skipped at their file ids, and drop added tokens past the file's vocabulary.""" + from tokenizers import Tokenizer + from transformers.convert_slow_tokenizer import bytes_to_unicode + + tokens = meta["tokenizer.ggml.tokens"] + spelling = bytes_to_unicode() + alphabet = set(spelling.values()) + # transformers' byte-level vocabulary spells a token with characters outside the byte alphabet in bytes and then skips it (DeepSeek's full-width-bar tokens, MiniMax's tokens with spaces) + skipped = [ + (i, text, kind == 3) + for i, (text, kind) in enumerate(zip(tokens, meta.get("tokenizer.ggml.token_type") or ())) + if kind in (3, 4) and meta.get("tokenizer.ggml.model") == "gpt2" and not set(text) <= alphabet and fast.token_to_id(text) != i + ] + # its GPT2Tokenizer also adds a default <|endoftext|> when the vocabulary lacks one + if not skipped and fast.get_vocab_size(with_added_tokens=True) <= len(tokens): + return fast + # a piece the merges use or build stays a plain BPE token, since renaming it breaks the merges (unsloth's MiniMax-M3 marks U+FFFD as control) + pieces = {p for m in meta.get("tokenizer.ggml.merges") or () for p in (*m.split(" "), m.replace(" ", ""))} + state = json.loads(fast.to_str()) + vocab = state["model"]["vocab"] + added = [t for t in state["added_tokens"] if t["id"] < len(tokens)] + for i, text, special in skipped: + spelled = "".join(spelling[b] for b in text.encode("utf-8")) + if spelled in pieces or vocab.get(spelled) != i: + continue + # renamed so the added token keeps the file id instead of taking a new one + vocab[text] = vocab.pop(spelled) + added.append({"id": i, "content": text, "single_word": False, "lstrip": False, "rstrip": False, "normalized": False, "special": special}) + state["added_tokens"] = sorted(added, key=lambda t: t["id"]) + return Tokenizer.from_str(json.dumps(state)) + + +def _known_specials(special: dict, fast) -> dict: + """Drop a bos/eos/unk/pad the vocabulary lacks, which PreTrainedTokenizerFast would add past it (unsloth's MiniMax-M3 unk is U+FFFD stored decoded).""" + vocab = fast.get_vocab(with_added_tokens=True) + return {k: v for k, v in special.items() if k == "chat_template" or v is None or v in vocab} def load_gguf_tokenizer(model_path: str): from transformers import PreTrainedTokenizerFast - from transformers.integrations.ggml import convert_gguf_tokenizer + from transformers.integrations.gguf import convert_gguf_tokenizer, get_gguf_tokenizer + arch, tok_dict, special = get_gguf_tokenizer(model_path) + fast, _ = convert_gguf_tokenizer(arch, tok_dict) meta = load_gguf_metadata(model_path) - arch = gguf_architecture(model_path) - conv_arch = _TOKENIZER_ARCH.get(arch, arch) - tok_dict: dict[str, Any] = { - k[len("tokenizer.ggml.") :]: v - for k, v in meta.items() - if k.startswith("tokenizer.ggml.") - } - fast, _extra = convert_gguf_tokenizer(conv_arch, tok_dict) - - tokens = tok_dict["tokens"] - - def tok_for(id_key: str, default: str) -> str: - tid = meta.get(f"tokenizer.ggml.{id_key}") - return tokens[int(tid)] if tid is not None and int(tid) < len(tokens) else default - - # gemma4 chat turns end with ; prefer it as eos so chat generation halts - # (the formal is also a stop id, see gguf_eos_token_ids). - turn_end = "" if "" in tokens else None - tokenizer = PreTrainedTokenizerFast( + fast = _fix_added_tokens(fast, meta) + pre = meta.get("tokenizer.ggml.pre") if meta.get("tokenizer.ggml.model") == "gpt2" else None + if pre in ("glm4", "glm5"): + # llama.cpp spells the same GLM pre-tokenizer "glm4" and "glm5" (llama-vocab.cpp maps both + # to CHATGLM4 with ignore_merges); transformers maps no split for either + from tokenizers import Regex, pre_tokenizers + from transformers.integrations.gguf.gguf_tokenizer_mapping import GGUF_PRE_TOKENIZER_SPLITS + + fast.pre_tokenizer = pre_tokenizers.Sequence([ + pre_tokenizers.Split(Regex(GGUF_PRE_TOKENIZER_SPLITS["llama3"]), behavior="isolated"), + pre_tokenizers.ByteLevel(add_prefix_space=False, use_regex=False), + ]) + if pre in _IGNORE_MERGES_PRE: + fast.model.ignore_merges = True + # special: bos/eos/unk/pad as the file names them (None when absent) and the chat template + return PreTrainedTokenizerFast( tokenizer_object=fast, - bos_token=tok_for("bos_token_id", ""), - eos_token=turn_end or tok_for("eos_token_id", ""), - unk_token=tok_for("unknown_token_id", ""), - pad_token=tok_for("padding_token_id", ""), + **_known_specials(special, fast), + add_bos_token=_add_bos(meta), + add_eos_token=bool(meta.get("tokenizer.ggml.add_eos_token")), ) - chat_template = meta.get("tokenizer.chat_template") - if chat_template: - tokenizer.chat_template = chat_template - return tokenizer -def gguf_eos_token_ids(model_path: str, tokenizer) -> set[int]: - """Stop ids for GGUF generation: the formal plus the chat turn end .""" +def gguf_eos_token_ids(model_path: str) -> set[int]: + """Stop ids for GGUF generation, picked the way llama.cpp picks its end-of-generation tokens.""" meta = load_gguf_metadata(model_path) tokens = meta["tokenizer.ggml.tokens"] - ids: set[int] = set() - if tokenizer.eos_token_id is not None: - ids.add(int(tokenizer.eos_token_id)) - eid = meta.get("tokenizer.ggml.eos_token_id") - if eid is not None: - ids.add(int(eid)) - # Look the stop tokens up in the vocab directly (convert_tokens_to_ids would map an - # absent name to , wrongly adding it as a stop id). - for name in ("", ""): - try: - ids.add(tokens.index(name)) - except ValueError: - pass + ids = {i for i, t in enumerate(tokens) if t in _EOG_NAMES} + ids |= {int(meta[f"tokenizer.ggml.{k}_token_id"]) for k in ("eos", "eot", "eom") if f"tokenizer.ggml.{k}_token_id" in meta} + names = {tokens[i]: i for i in ids if i < len(tokens)} + # llama.cpp's exceptions: <|end|> only ends a message when <|call|> and <|return|> end the turn (gpt-oss), and gemma4's is a normal token + if "<|end|>" in names and names.keys() & {"<|call|>", "<|calls|>"} and names.keys() & {"<|return|>", "<|flush|>"}: + ids.discard(names["<|end|>"]) + if "" in names and names.keys() & {"<|tool_response>", "<|plamo:eos|>"}: + ids.discard(names[""]) return ids diff --git a/python/freetoken/models/glm4_moe/__init__.py b/python/freetoken/models/glm4_moe/__init__.py index b8748b1d44..7f093c6762 100644 --- a/python/freetoken/models/glm4_moe/__init__.py +++ b/python/freetoken/models/glm4_moe/__init__.py @@ -1,4 +1,5 @@ from .config import parse_config +from .gguf import iter_gguf_weights, parse_gguf_config, to_gguf from .model import Glm4MoeForCausalLM from .weight import iter_weights, nvfp4_expert_spec @@ -7,4 +8,7 @@ "Glm4MoeForCausalLM", "parse_config", "iter_weights", + "iter_gguf_weights", + "parse_gguf_config", + "to_gguf", ] diff --git a/python/freetoken/models/glm4_moe/attention.py b/python/freetoken/models/glm4_moe/attention.py index 7a4b335ef1..65b70bf155 100644 --- a/python/freetoken/models/glm4_moe/attention.py +++ b/python/freetoken/models/glm4_moe/attention.py @@ -3,7 +3,7 @@ from typing import TYPE_CHECKING from freetoken.core import get_global_ctx -from freetoken.layers import BaseOP, RMSNorm +from freetoken.layers import BaseOP, LinearReplicated, RMSNorm from freetoken.layers.rotary import get_rope from freetoken.utils import nvtx_annotate @@ -27,9 +27,10 @@ class Glm4MoeAttention(BaseOP): every ``self_attn`` out of quantization (qkvo stay bf16), so we reproduce that exactly by storing the bf16 weights compressed (~10.7 bits/weight) and decoding them bit-for-bit in the forward. This fits a 32 GB VRAM target without the accuracy loss fp8 attention would add. + From a GGUF file the projections are the file's packed ggml linears instead. """ - def __init__(self, config: ModelConfig, layer_id: int): + def __init__(self, config: ModelConfig, layer_id: int, *, prefix: str = ""): head_dim = config.head_dim self.layer_id = layer_id self.num_qo_heads = config.num_qo_heads @@ -38,10 +39,17 @@ def __init__(self, config: ModelConfig, layer_id: int): self.kv_attn_dim = self.num_kv_heads * head_dim self.head_dim = head_dim - self.q_proj = LinearDF11(config.hidden_size, self.qo_attn_dim, has_bias=config.has_attn_bias) - self.k_proj = LinearDF11(config.hidden_size, self.kv_attn_dim, has_bias=config.has_attn_bias) - self.v_proj = LinearDF11(config.hidden_size, self.kv_attn_dim, has_bias=config.has_attn_bias) - self.o_proj = LinearDF11(self.qo_attn_dim, config.hidden_size, has_bias=False) + gguf = config.quant is not None and config.quant.dialect == "gguf" + + def linear(in_features: int, out_features: int, has_bias: bool, name: str) -> BaseOP: + if gguf: + return LinearReplicated(in_features, out_features, has_bias, quant_config=config.quant, prefix=f"{prefix}.{name}") + return LinearDF11(in_features, out_features, has_bias=has_bias) + + self.q_proj = linear(config.hidden_size, self.qo_attn_dim, config.has_attn_bias, "q_proj") + self.k_proj = linear(config.hidden_size, self.kv_attn_dim, config.has_attn_bias, "k_proj") + self.v_proj = linear(config.hidden_size, self.kv_attn_dim, config.has_attn_bias, "v_proj") + self.o_proj = linear(self.qo_attn_dim, config.hidden_size, False, "o_proj") if config.use_qk_norm: self.q_norm = RMSNorm(head_dim, eps=config.rms_norm_eps) diff --git a/python/freetoken/models/glm4_moe/gguf.py b/python/freetoken/models/glm4_moe/gguf.py new file mode 100644 index 0000000000..a477706b86 --- /dev/null +++ b/python/freetoken/models/glm4_moe/gguf.py @@ -0,0 +1,179 @@ +"""GLM-4.5 / 4.6 / 4.7 from a llama.cpp ``glm4moe`` GGUF: config from the GGUF metadata, tensors by name; the routed experts load as GGUF expert banks (``freetoken.models.gguf.experts``). + +The llama.cpp converter (conversion/glm.py) only renames the tensors and stacks the experts: both sides run NeoX rope on the first ``rope.dimension_count`` dims of a head, so q/k carry no rope permute. +Files that keep the MTP layer append it as the last block, which is not loaded. +""" + +from __future__ import annotations + +import re +from dataclasses import replace +from types import SimpleNamespace +from typing import TYPE_CHECKING, Iterator + +import torch + +from freetoken.models.gguf.reader import GgufTensorInfo, gguf_tensor_index +from freetoken.models.gguf.tensors import float_tensor, fused, matrix + +from .config import parse_config + +if TYPE_CHECKING: + from freetoken.models.config import ModelConfig + from freetoken.models.gguf.config import GgufConfigShim + +_ARCH = "glm4moe" +_LAYER_LEAVES = { + "self_attn.q_proj": "attn_q", + "self_attn.k_proj": "attn_k", + "self_attn.v_proj": "attn_v", + "self_attn.o_proj": "attn_output", + "mlp.gate_proj": "ffn_gate", + "mlp.up_proj": "ffn_up", + "mlp.down_proj": "ffn_down", + "mlp.gate": "ffn_gate_inp", + "mlp.shared_experts.gate_proj": "ffn_gate_shexp", + "mlp.shared_experts.up_proj": "ffn_up_shexp", + "mlp.shared_experts.down_proj": "ffn_down_shexp", + "mlp.experts.0.gate_proj": "ffn_gate_exps", + "mlp.experts.0.up_proj": "ffn_up_exps", + "mlp.experts.0.down_proj": "ffn_down_exps", +} +_TOP_LEAVES = {"model.embed_tokens": "token_embd", "lm_head": "output"} +_LAYER_RE = re.compile(r"^model\.layers\.(\d+)\.(.+)$") +# llama.cpp's LLAMA_EXPERT_GATING_FUNC_TYPE_*: 0 (unset) means sigmoid for this arch +_SIGMOID_GATING = (0, 2) + + +def to_gguf(name: str, tensors: dict[str, int]) -> tuple[str, ...]: + """Checkpoint module name -> the GGUF tensors (without ``.weight``) that store it; empty for modules the GGUF has no matrix for.""" + if name in _TOP_LEAVES: + return (_TOP_LEAVES[name],) + m = _LAYER_RE.match(name) + if m is None: + return () + tensor = _LAYER_LEAVES.get(m.group(2)) + return (f"blk.{m.group(1)}.{tensor}",) if tensor is not None else () + + +def _check_supported(shim: GgufConfigShim) -> None: + md = shim.metadata + if any(k.startswith("prism.hadamard.") for k in md): + raise NotImplementedError(f"{shim.model_path}: Hadamard-rotated (prism) GGUF weights are not supported") + names = sorted(gguf_tensor_index(shim.model_path)) + if not names: + raise NotImplementedError(f"{shim.model_path}: a metadata-only glm4moe GGUF has no tensor table to tell whether the attention has q/k norms and biases") + if "blk.0.attn_q.weight" not in names: + raise NotImplementedError(f"{shim.model_path}: no decoder layer tensors (blk.0.attn_q), e.g. an MTP-only draft GGUF") + scaled = [n for n in names if n.endswith((".scale", ".input_scale"))] + if scaled: + raise NotImplementedError(f"{shim.model_path}: GGUF tensors with per-tensor scales are not supported (e.g. {scaled[0]})") + fused_qkv = [n for n in names if ".attn_qkv." in n] + if fused_qkv: + raise NotImplementedError(f"{shim.model_path}: a fused attn_qkv glm4moe GGUF is not supported (e.g. {fused_qkv[0]})") + gating = int(md.get(f"{_ARCH}.expert_gating_func", 0)) + if gating not in _SIGMOID_GATING: + raise NotImplementedError(f"{shim.model_path}: expert gating function {gating} is not supported, GLM routes with sigmoid") + scaling = md.get(f"{_ARCH}.rope.scaling.type", "none") + if scaling != "none": + raise NotImplementedError(f"{shim.model_path}: rope scaling {scaling!r} is not supported") + + +def _hf_config(shim: GgufConfigShim) -> SimpleNamespace: + """The HF config fields ``parse_config`` reads, from the GGUF metadata.""" + md = shim.metadata + + def kv(key: str): + return md[f"{_ARCH}.{key}"] + + _check_supported(shim) + index = gguf_tensor_index(shim.model_path) + head_dim = int(kv("attention.key_length")) + return SimpleNamespace( + model_type="glm4_moe", + architectures=list(shim.architectures), + hidden_act="silu", + hidden_size=int(kv("embedding_length")), + # block_count also counts the MTP layers a GGUF may carry after the decoder + num_hidden_layers=int(kv("block_count")) - int(md.get(f"{_ARCH}.nextn_predict_layers", 0)), + num_attention_heads=int(kv("attention.head_count")), + num_key_value_heads=int(kv("attention.head_count_kv")), + head_dim=head_dim, + rope_theta=float(kv("rope.freq_base")), + partial_rotary_factor=int(kv("rope.dimension_count")) / head_dim, + max_position_embeddings=int(kv("context_length")), + rms_norm_eps=float(kv("attention.layer_norm_rms_epsilon")), + vocab_size=int(shim.vocab_size), + tie_word_embeddings=bool(shim.tie_word_embeddings), + intermediate_size=int(kv("feed_forward_length")), + n_routed_experts=int(kv("expert_count")), + num_experts_per_tok=int(kv("expert_used_count")), + moe_intermediate_size=int(kv("expert_feed_forward_length")), + n_shared_experts=int(kv("expert_shared_count")), + first_k_dense_replace=int(md.get(f"{_ARCH}.leading_dense_block_count", 0)), + # llama.cpp neither renormalizes nor scales the router weights when the keys are absent + norm_topk_prob=bool(md.get(f"{_ARCH}.expert_weights_norm", False)), + routed_scaling_factor=float(md.get(f"{_ARCH}.expert_weights_scale", 0.0)) or 1.0, + n_group=max(1, int(md.get(f"{_ARCH}.expert_group_count", 1))), + topk_group=max(1, int(md.get(f"{_ARCH}.expert_group_used_count", 1))), + # GLM-4.5-Air has no q/k norms; the converter writes no flag for them or for the qkv biases + use_qk_norm="blk.0.attn_q_norm.weight" in index, + attention_bias="blk.0.attn_q.bias" in index, + ) + + +def parse_gguf_config(shim: GgufConfigShim) -> ModelConfig: + return replace(parse_config(_hf_config(shim)), expert_quant="gguf") + + +def iter_gguf_weights( + model_path: str, + device: torch.device, + *, + include_moe_experts: bool, + include_non_moe: bool, +) -> Iterator[tuple[str, torch.Tensor]]: + """The dense state dict from the GGUF; the routed experts load as GGUF expert banks (``freetoken.models.gguf.experts``).""" + if not include_non_moe: + return + from freetoken.utils import cached_load_hf_config + + hf = _hf_config(cached_load_hf_config(model_path)) + index = gguf_tensor_index(model_path) + + def t(name: str) -> GgufTensorInfo: + return index[name if name in index else f"{name}.weight"] + + def linear(module: str, tensor: str) -> Iterator[tuple[str, torch.Tensor]]: + info = t(tensor) + yield from fused(module, [info], [matrix(info)], int(info.ne[0])) + if f"{tensor}.bias" in index: + yield f"{module}.bias", float_tensor(index[f"{tensor}.bias"]) + + yield from linear("model.embed_tokens", "token_embd") + if not hf.tie_word_embeddings: + yield from linear("lm_head", "output") + yield "model.norm.weight", float_tensor(t("output_norm")) + + for layer in range(hf.num_hidden_layers): + blk, pre = f"blk.{layer}", f"model.layers.{layer}" + yield f"{pre}.input_layernorm.weight", float_tensor(t(f"{blk}.attn_norm")) + yield f"{pre}.post_attention_layernorm.weight", float_tensor(t(f"{blk}.post_attention_norm")) + sa = f"{pre}.self_attn" + for proj, tensor in (("q_proj", "attn_q"), ("k_proj", "attn_k"), ("v_proj", "attn_v"), ("o_proj", "attn_output")): + yield from linear(f"{sa}.{proj}", f"{blk}.{tensor}") + if hf.use_qk_norm: + yield f"{sa}.q_norm.weight", float_tensor(t(f"{blk}.attn_q_norm")) + yield f"{sa}.k_norm.weight", float_tensor(t(f"{blk}.attn_k_norm")) + mlp = f"{pre}.mlp" + if layer < hf.first_k_dense_replace: + for proj in ("gate", "up", "down"): + yield from linear(f"{mlp}.{proj}_proj", f"{blk}.ffn_{proj}") + continue + yield f"{mlp}.gate.weight", float_tensor(t(f"{blk}.ffn_gate_inp")) + yield f"{mlp}.e_score_correction_bias", float_tensor(t(f"{blk}.exp_probs_b.bias")) + for proj in ("gate", "up", "down"): + yield from linear(f"{mlp}.shared_experts.{proj}_proj", f"{blk}.ffn_{proj}_shexp") + + +__all__ = ["iter_gguf_weights", "parse_gguf_config", "to_gguf"] diff --git a/python/freetoken/models/glm4_moe/model.py b/python/freetoken/models/glm4_moe/model.py index b5c647e146..8824dfb54b 100644 --- a/python/freetoken/models/glm4_moe/model.py +++ b/python/freetoken/models/glm4_moe/model.py @@ -4,7 +4,7 @@ import torch from freetoken.core import get_global_ctx -from freetoken.layers import BaseOP, OPList, ParallelLMHead, RMSNormFused +from freetoken.layers import BaseOP, OPList, ParallelLMHead, RMSNormFused, VocabParallelEmbedding from freetoken.models.blocks import BaseLLMModel from freetoken.utils import nvtx_annotate @@ -19,7 +19,7 @@ class Glm4MoeDecoderLayer(BaseOP): def __init__(self, config: ModelConfig, layer_id: int, *, prefix: str = ""): - self.self_attn = Glm4MoeAttention(config, layer_id) + self.self_attn = Glm4MoeAttention(config, layer_id, prefix=f"{prefix}.self_attn") if layer_id >= config.first_k_dense_replace: self.mlp: BaseOP = Glm4MoeSparseBlock(config, layer_id, prefix=f"{prefix}.mlp") else: @@ -45,10 +45,15 @@ def forward( class Glm4MoeModel(BaseOP): def __init__(self, config: ModelConfig, *, prefix: str = "model"): - self.embed_tokens = EmbeddingDF11( - num_embeddings=config.vocab_size, - embedding_dim=config.hidden_size, - ) + if config.quant is not None and config.quant.dialect == "gguf": + self.embed_tokens: BaseOP = VocabParallelEmbedding( + config.vocab_size, config.hidden_size, quant_config=config.quant, prefix=f"{prefix}.embed_tokens" + ) + else: + self.embed_tokens = EmbeddingDF11( + num_embeddings=config.vocab_size, + embedding_dim=config.hidden_size, + ) self.layers = OPList( [ Glm4MoeDecoderLayer(config, layer_id, prefix=f"{prefix}.layers.{layer_id}") diff --git a/python/freetoken/models/glm4_moe/weight.py b/python/freetoken/models/glm4_moe/weight.py index aa6538dcf9..0ca79efbd8 100644 --- a/python/freetoken/models/glm4_moe/weight.py +++ b/python/freetoken/models/glm4_moe/weight.py @@ -78,12 +78,12 @@ def _iter_nvfp4_resident( """Yield native NVFP4 buffers for an always-resident Linear. ``weight`` (packed fp4) + ``weight_scale`` (fp8 block scale) verbatim; per-tensor - ``weight_scale_2`` broadcast to per-row fp16 ``weight_global`` as the NVFP4 linear + ``weight_scale_2`` broadcast to per-row fp32 ``weight_global`` as the NVFP4 linear method / dequant_nvfp4 expect. lossless vs checkpoint; same dequant math as routed experts. """ packed = reader.get(f"{src_prefix}.weight") # [OUT, IN//2] uint8 scale = reader.get(f"{src_prefix}.weight_scale") # [OUT, IN//16] fp8-e4m3 - g = reader.get(f"{src_prefix}.weight_scale_2").reshape(()).to(torch.float16) # scalar + g = reader.get(f"{src_prefix}.weight_scale_2").reshape(()).to(torch.float32) # scalar yield f"{dst_prefix}.weight", packed yield f"{dst_prefix}.weight_scale", scale yield f"{dst_prefix}.weight_global", g.expand(packed.shape[0]).contiguous() diff --git a/python/freetoken/models/glm5_next/__init__.py b/python/freetoken/models/glm5_next/__init__.py index 7eee33dfee..0c548da7f6 100644 --- a/python/freetoken/models/glm5_next/__init__.py +++ b/python/freetoken/models/glm5_next/__init__.py @@ -1,4 +1,5 @@ from .config import VisionConfig, parse_config, parse_vision_config +from .gguf import iter_gguf_weights, parse_gguf_config, to_gguf from .model import Glm5NextForCausalLM, Glm5NextForConditionalGeneration from .vision import Glm5NextVisionModel from .weight import iter_expert_pieces, iter_vision_weights, iter_weights, nvfp4_expert_spec @@ -10,8 +11,11 @@ "Glm5NextVisionModel", "VisionConfig", "parse_config", + "parse_gguf_config", "parse_vision_config", "iter_vision_weights", + "iter_gguf_weights", "iter_weights", + "to_gguf", "iter_expert_pieces", ] diff --git a/python/freetoken/models/glm5_next/args.py b/python/freetoken/models/glm5_next/args.py index c25fef4856..c771467b89 100644 --- a/python/freetoken/models/glm5_next/args.py +++ b/python/freetoken/models/glm5_next/args.py @@ -24,16 +24,16 @@ from dataclasses import dataclass from typing import Any, Tuple -# Layer-type strings used by the checkpoint's ``layer_types`` field. +# Layer-type strings after transformers' legacy remap, which renames the checkpoint's deepseek_sparse_attention. KDA_LAYER = "linear_attention" -DSA_LAYER = "deepseek_sparse_attention" +DSA_LAYER = "indexed_attention" @dataclass(frozen=True) class Glm5NextArgs: hidden_size: int num_heads: int - # ---- MLA (the 11 "deepseek_sparse_attention" layers) ---- + # ---- MLA (the 11 "indexed_attention" layers) ---- q_lora_rank: int kv_lora_rank: int qk_nope_head_dim: int @@ -119,8 +119,10 @@ def load_args(hf_config: Any) -> Glm5NextArgs: are accepted as fallbacks.""" text = _get(hf_config, "text_config", hf_config) + from transformers.configuration_utils import remap_legacy_layer_types + num_layers = int(text.num_hidden_layers) - layer_types = tuple(_get(text, "layer_types", ()) or ()) + layer_types = tuple(remap_legacy_layer_types(list(_get(text, "layer_types", ()) or ()))) if not layer_types: raise ValueError("glm5_next config is missing layer_types") if len(layer_types) != num_layers: diff --git a/python/freetoken/models/glm5_next/gguf.py b/python/freetoken/models/glm5_next/gguf.py new file mode 100644 index 0000000000..73eb0a82ce --- /dev/null +++ b/python/freetoken/models/glm5_next/gguf.py @@ -0,0 +1,367 @@ +"""GLM-5.3-Flash from a llama.cpp ``glm5-next`` GGUF (unsloth ``GLM-5.3-Flash-GGUF``): config from the GGUF metadata, tensor names, and the converter's transforms undone. + +The GGUF geometry is identical to the HF checkpoint (34 KDA linear-attention layers, +11 NoPE-MLA/DSA layers with a kpool-compressed indexer on a 3+1 hybrid block layout, +mHC x4 residual streams, 288-expert sigmoid/noaux_tc MoE with a shared expert and +clamped SwiGLU), so ``parse_gguf_config`` derives a text-config-shaped namespace from +the GGUF KV metadata and reuses ``glm5_next.config.parse_config`` wholesale. + +The llama.cpp converter (unsloth's glm5-next writer; tensor names identical across +the ``glm5-next`` and the older ``glm5next`` arch spellings) stores KDA's ``A_log`` +pre-transformed as ``-exp(A_log)`` -- inverted at load -- and splits the MLA +``kv_b_proj`` into per-head 3-D ``attn_k_b`` / ``attn_v_b`` tensors (k stored +transposed, [heads, kv_lora, nope] next to [heads, v, kv_lora]); the halves may be +quantized (dynamic quants store them Q8_0), so they dequantize and fuse into the +dense layout the module's MLA absorption reads. The trailing MTP block (``blk.N.nextn.*``, ``nextn_predict_layers == 1``) is +dropped exactly like the HF loader drops the never-read MTP layer; routed experts +serve from the GGUF expert banks (``freetoken.models.gguf.experts``). +""" + +from __future__ import annotations + +import re +from dataclasses import replace +from types import SimpleNamespace +from typing import TYPE_CHECKING, Iterator + +import numpy as np +import torch + +from freetoken.layers.quantization.ggml_types import is_quantized +from freetoken.models.gguf.reader import GgufTensorInfo, gguf_tensor_index, gguf_tensor_names +from freetoken.models.gguf.tensors import check_scales, fused, matrix + +from .config import parse_config + +if TYPE_CHECKING: + from freetoken.models.config import ModelConfig + from freetoken.models.gguf.config import GgufConfigShim + +# HF config fields the GGUF metadata does not carry (GLM-5.3-Flash's config.json values). +_FAMILY_DEFAULTS = dict( + hidden_act="silu", + attention_bias=False, + # NoPE means no rope on the main attention; the indexer's interleaved rope is a + # served-model fact the GGUF KV predates, and getting it wrong scores the indexer + # keys incorrectly (DSA retrieval degrades to garbage). + indexer_rope_interleave=True, +) + +# checkpoint module leaf (glm5_next HF naming, language-model rooted) -> GGUF tensors +_LAYER_LEAVES: dict[str, tuple[str, ...]] = { + # KDA input projections (the six leaves of the fused in_proj) + "self_attn.q_proj": ("attn_q",), + "self_attn.k_proj": ("attn_k",), + "self_attn.v_proj": ("attn_v",), + "self_attn.b_proj": ("ssm_beta",), + "self_attn.f_a_proj": ("ssm_f_a",), + "self_attn.g_a_proj": ("ssm_g_a",), + "self_attn.f_b_proj": ("ssm_f_b",), + "self_attn.g_b_proj": ("ssm_g_b",), + # both layer kinds name their output projection the same + "self_attn.o_proj": ("attn_output",), + # DSA (MLA) projections + "self_attn.q_a_proj": ("attn_q_a",), + "self_attn.q_b_proj": ("attn_q_b",), + "self_attn.kv_a_proj_with_mqa": ("attn_kv_a_mqa",), + # two tensors, dequantized and fused at load: the module keeps its dense bf16 kv_b + # (exempted in the registry's unquantized_modules, so quantized halves are fine too) + "self_attn.kv_b_proj": ("attn_k_b", "attn_v_b"), + # kpool indexer + "self_attn.indexer.wq_b": ("indexer.attn_q_b",), + "self_attn.indexer.wk": ("indexer.attn_k",), + "self_attn.indexer.weights_proj": ("indexer.proj",), + # dense MLP of the leading blocks + "mlp.gate_proj": ("ffn_gate",), + "mlp.up_proj": ("ffn_up",), + "mlp.down_proj": ("ffn_down",), + # shared expert of the MoE blocks + "mlp.shared_experts.gate_proj": ("ffn_gate_shexp",), + "mlp.shared_experts.up_proj": ("ffn_up_shexp",), + "mlp.shared_experts.down_proj": ("ffn_down_shexp",), + # routed experts (GGUF expert banks) + "mlp.experts.0.gate_proj": ("ffn_gate_exps",), + "mlp.experts.0.up_proj": ("ffn_up_exps",), + "mlp.experts.0.down_proj": ("ffn_down_exps",), +} +_TOP_LEAVES: dict[str, tuple[str, ...]] = { + # checkpoint-space names: the language-model root is already applied when to_gguf runs + "model.language_model.embed_tokens": ("token_embd",), + "lm_head": ("output",), +} +_LAYER_RE = re.compile(r"^model\.language_model\.layers\.(\d+)\.(.+)$") + + +def to_gguf(name: str, tensors: dict[str, int]) -> tuple[str, ...]: + """Checkpoint module name -> the GGUF tensors (without ``.weight``) that store it; empty for modules the GGUF has no matrix for.""" + if name in _TOP_LEAVES: + return _TOP_LEAVES[name] + m = _LAYER_RE.match(name) + if m is None: + return () + return tuple(f"blk.{m.group(1)}.{leaf}" for leaf in _LAYER_LEAVES.get(m.group(2), ())) + + +def _layer_types(shim: "GgufConfigShim", num_layers: int) -> tuple[str, ...]: + """Per-decoder-layer kinds: DSA where the MLA projections sit, KDA elsewhere, from the + per-block ``attention.head_count_kv`` list (0 on the KDA blocks, 1 on the MLA ones). + + Every public glm5-next GGUF carries the list (the generic scalar spelling cannot express + the hybrid; llama.cpp broadcasts it and then fails tensor creation), so anything else is + refused rather than guessed at.""" + from .args import DSA_LAYER, KDA_LAYER + + per_block = shim.metadata.get(f"{shim.model_type}.attention.head_count_kv") + if not isinstance(per_block, list) or len(per_block) < num_layers: + raise ValueError( + f"{shim.model_path}: glm5-next GGUF lacks the per-block attention.head_count_kv list " + "the KDA/MLA hybrid layout needs (reconvert with a current llama.cpp)" + ) + return tuple(DSA_LAYER if int(kv) else KDA_LAYER for kv in per_block[:num_layers]) + + +def parse_gguf_config(shim: "GgufConfigShim") -> ModelConfig: + md = shim.metadata + prefix = shim.model_type + + check_scales(shim.model_path) + + def kv(key: str, default=None): + return md.get(f"{prefix}.{key}", default) + + def need(key: str): + val = kv(key) + if val is None: + raise KeyError(f"missing GGUF metadata key {prefix}.{key}") + return val + + # block_count also counts the MTP block a GGUF carries after the decoder + n_mtp = int(kv("nextn_predict_layers", 0) or 0) + num_layers = int(need("block_count")) - n_mtp + layer_types = _layer_types(shim, num_layers) + first_dense = int(need("leading_dense_block_count")) + + heads = int(need("attention.head_count")) + qk_nope, v_head = int(need("attention.key_length_mla")), int(need("attention.value_length_mla")) + # NoPE: the main attention carries no rotary dims (rope survives only in the indexer) + if int(kv("rope.dimension_count", 0) or 0) != 0: + raise NotImplementedError("glm5-next GGUF: rotary dims on the main attention are not served") + + # KDA head count: one q/k/v head set of kda.head_dim each; the conv channels or the + # beta rows count them, falling back to the shared head_count (the two coincide on + # every release) when the tensor table is gone (metadata-only GGUF). + kda_head_dim = int(need("kda.head_dim")) + index = gguf_tensor_index(shim.model_path) + conv_q = next((i for i in index.values() if i.name.endswith("ssm_conv1d_q.weight")), None) + kda_heads = conv_q.ne[2] // kda_head_dim if conv_q is not None else heads + + kpool = int(kv("attention.indexer.kpool", 1) or 1) + # the compressor tensors exist exactly when the indexer pools; a metadata-only GGUF + # keeps the fact only through kpool > 1 + names = gguf_tensor_names(shim.model_path) + kpool_compress = any(n.endswith("indexer_compressor_gate.weight") for n in names) if names else kpool > 1 + indexer_types_kv = kv("attention.indexer.types") + if indexer_types_kv is not None and not all(indexer_types_kv): + # llama.cpp reads False as "no full indexer on this layer", which the model cannot express + raise NotImplementedError(f"{shim.model_path}: glm5-next GGUF with layers that lack a full indexer is not served") + indexer_types = ("full",) * num_layers + + # the clamp is per-layer in the GGUF (and the shared expert carries its own array); the model + # takes one value, so anything non-uniform is refused instead of silently taking the first + clamp = need("swiglu_clamp_exp") + clamp_vals = list(clamp) if isinstance(clamp, list) else [clamp] + clamp_shexp = kv("swiglu_clamp_shexp") + shexp_vals = list(clamp_shexp) if isinstance(clamp_shexp, list) else [clamp_shexp] if clamp_shexp is not None else [] + if len(set(map(float, clamp_vals + shexp_vals))) != 1: + raise NotImplementedError(f"{shim.model_path}: non-uniform swiglu_clamp_exp / swiglu_clamp_shexp is not served") + if int(kv("expert_gating_func", 2) or 2) != 2: + raise NotImplementedError(f"{shim.model_path}: expert_gating_func other than sigmoid (2) is not served") + text = SimpleNamespace( + **_FAMILY_DEFAULTS, + num_hidden_layers=num_layers, + layer_types=tuple(layer_types), + mlp_layer_types=("dense",) * first_dense + ("sparse",) * (num_layers - first_dense), + first_k_dense_replace=first_dense, + hidden_size=int(need("embedding_length")), + num_attention_heads=heads, + vocab_size=int(shim.vocab_size), + intermediate_size=int(need("feed_forward_length")), + rms_norm_eps=float(need("attention.layer_norm_rms_epsilon")), + max_position_embeddings=int(need("context_length")), + q_lora_rank=int(need("attention.q_lora_rank")), + kv_lora_rank=int(need("attention.kv_lora_rank")), + qk_nope_head_dim=qk_nope, + qk_rope_head_dim=0, + v_head_dim=v_head, + mla_use_nope=True, + index_n_heads=int(need("attention.indexer.head_count")), + index_head_dim=int(need("attention.indexer.key_length")), + index_topk=int(need("attention.indexer.top_k")), + indexer_types=indexer_types, + index_kpool=kpool, + index_kpool_compress=kpool_compress, + # old converters predate the key; both released checkpoints set it true + index_kpool_always_select_tail=bool(kv("attention.indexer.kpool_select_tail", True)), + linear_attn_config={ + "num_heads": kda_heads, + "head_dim": kda_head_dim, + "short_conv_kernel_size": int(need("ssm.conv_kernel")), + "gate_lower_bound": float(need("kda.gate_lower_bound")), + }, + mhc=True, + hc_mult=int(need("hyper_connection.count")), + hc_sinkhorn_iters=int(need("hyper_connection.sinkhorn_iterations")), + hc_eps=float(need("hyper_connection.epsilon")), + n_routed_experts=int(need("expert_count")), + n_shared_experts=int(kv("expert_shared_count", 1) or 1), + num_experts_per_tok=int(need("expert_used_count")), + moe_intermediate_size=int(need("expert_feed_forward_length")), + norm_topk_prob=bool(need("expert_weights_norm")), + routed_scaling_factor=float(need("expert_weights_scale")), + n_group=int(kv("expert_group_count", 1) or 1), + topk_group=int(kv("expert_group_used_count", 1) or 1), + tie_word_embeddings=bool(shim.tie_word_embeddings), + swiglu_limit=float(clamp_vals[0]), + ) + hf_like = SimpleNamespace( + text_config=text, + model_type="glm5_next", + architectures=list(shim.architectures), + vision_config=None, + image_token_id=None, + ) + config = parse_config(hf_like) + return replace(config, expert_quant="gguf") + + +# ---------------------------------------------------------------------------------------------- +# Weights +# ---------------------------------------------------------------------------------------------- + +def _float(info: GgufTensorInfo) -> torch.Tensor: + """A float tensor in torch order, dequantizing packed ggml rows: dynamic quants quantize + tensors the module consumes as floats (the kv_b halves, the mHC mixers, the compressor gate). + Trailing unit dims in ``ne`` are dropped -- some writers ship the conv kernels 4-D + (llama.cpp's glm5-next loader tries that shape first) and a trailing 1 carries no elements.""" + ne = info.ne + while len(ne) > 1 and ne[-1] == 1: + ne = ne[:-1] + if not is_quantized(info.ggml_type): + from freetoken.models.gguf.tensors import float_tensor + + return float_tensor(info).reshape(tuple(reversed(ne))) + from freetoken.third_party.gguf import quants + from freetoken.third_party.gguf.constants import GGMLQuantizationType + + flat = quants.dequantize(np.array(info.data()), GGMLQuantizationType(info.ggml_type)) + return torch.from_numpy(flat).view(*reversed(ne)) + + +def _kv_b(kb: GgufTensorInfo, vb: GgufTensorInfo, num_heads: int, kv_lora: int) -> torch.Tensor: + """Per-head k [kv_lora, nope] (stored transposed) + v [v, kv_lora] -> the module's fused dense [heads * (nope + v), kv_lora].""" + k = _float(kb).transpose(1, 2) # [heads, nope, kv_lora] + fused_w = torch.cat([k, _float(vb)], dim=1) + assert fused_w.shape[0] == num_heads, (fused_w.shape, num_heads) + return fused_w.reshape(num_heads * fused_w.shape[1], kv_lora).to(torch.bfloat16).contiguous() + + +def iter_gguf_weights( + model_path: str, + device: torch.device, + *, + include_moe_experts: bool, + include_non_moe: bool, + include_vision: bool = True, +) -> Iterator[tuple[str, torch.Tensor]]: + """The dense state dict from the GGUF; the routed experts load as offload banks (``freetoken.models.gguf.experts``).""" + if not include_non_moe: + return + from freetoken.utils import cached_load_hf_config + + config = parse_gguf_config(cached_load_hf_config(model_path)) + args = config.glm5_args + index = gguf_tensor_index(model_path) + hidden = config.hidden_size + + def t(name: str) -> GgufTensorInfo: + return index[name if name in index else f"{name}.weight"] + + def bf16(name: str) -> torch.Tensor: + return _float(t(name)).to(torch.bfloat16) + + yield from fused("model.embed_tokens", [t("token_embd")], [matrix(t("token_embd"))], hidden) + yield from fused("lm_head", [t("output")], [matrix(t("output"))], hidden) + yield "model.norm.weight", bf16("output_norm") + + for layer in range(config.num_layers): + blk, pre = f"blk.{layer}", f"model.layers.{layer}" + yield f"{pre}.input_layernorm.weight", bf16(f"{blk}.attn_norm") + yield f"{pre}.post_attention_layernorm.weight", bf16(f"{blk}.ffn_norm") + # mHC mixing tensors, fp32 on every layer + for sub in ("attn", "ffn"): + for part in ("fn", "base", "scale"): + yield f"{pre}.hc_{sub}_{part}", _float(t(f"{blk}.hc_{sub}_{part}")) + + if args.is_kda_layer(layer): + sa = f"{pre}.self_attn" + p, h, d = args.linear_num_heads * args.linear_head_dim, args.linear_num_heads, args.linear_head_dim + group = [t(f"{blk}.{n}") for n in ("attn_q", "attn_k", "attn_v", "ssm_beta", "ssm_f_a", "ssm_g_a")] + yield from fused(f"{sa}.in_proj", group, [matrix(i) for i in group], hidden) + # the GGUF conv kernels already carry the [channels, 1, kernel] singleton dim + conv = torch.cat([_float(t(f"{blk}.ssm_conv1d_{q}")) for q in ("q", "k", "v")], dim=0) + yield f"{sa}.conv1d.weight", conv.to(torch.bfloat16) + for proj, src, k in (("f_b_proj", "ssm_f_b", d), ("g_b_proj", "ssm_g_b", d), ("o_proj", "attn_output", p)): + info = t(f"{blk}.{src}") + yield from fused(f"{sa}.{proj}", [info], [matrix(info)], k) + # the converter stores A_log pre-transformed as -exp(A_log); the kernels want the raw log + yield f"{sa}.A_log", torch.log(-_float(t(f"{blk}.ssm_a"))) + yield f"{sa}.dt_bias", _float(t(f"{blk}.ssm_dt.bias")) + yield f"{sa}.o_norm.weight", bf16(f"{blk}.ssm_norm") + else: + sa = f"{pre}.self_attn" + for proj, src, k in ( + ("q_a_proj", "attn_q_a", hidden), + ("q_b_proj", "attn_q_b", args.q_lora_rank), + ("kv_a_proj_with_mqa", "attn_kv_a_mqa", hidden), + ("o_proj", "attn_output", args.num_heads * args.v_head_dim), + ): + info = t(f"{blk}.{src}") + yield from fused(f"{sa}.{proj}", [info], [matrix(info)], k) + yield f"{sa}.kv_b_proj.weight", _kv_b( + t(f"{blk}.attn_k_b"), t(f"{blk}.attn_v_b"), args.num_heads, args.kv_lora_rank + ) + yield f"{sa}.q_a_layernorm.weight", bf16(f"{blk}.attn_q_a_norm") + yield f"{sa}.kv_a_layernorm.weight", bf16(f"{blk}.attn_kv_a_norm") + # wq_b serves its checkpoint quant (packed); wk / weights_proj are dense modules + # (glm_moe_dsa's indexer reasoning), dequantized at load when the file quantized them + yield from fused(f"{sa}.indexer.wq_b", [t(f"{blk}.indexer.attn_q_b")], [matrix(t(f"{blk}.indexer.attn_q_b"))], args.q_lora_rank) + yield f"{sa}.indexer.wk.weight", _float(t(f"{blk}.indexer.attn_k")).to(torch.bfloat16) + yield f"{sa}.indexer.weights_proj.weight", _float(t(f"{blk}.indexer.proj")).to(torch.bfloat16) + yield f"{sa}.indexer.k_norm.weight", bf16(f"{blk}.indexer.k_norm") + yield f"{sa}.indexer.k_norm.bias", bf16(f"{blk}.indexer.k_norm.bias") + yield f"{sa}.indexer.index_kpool_compress_gate", bf16(f"{blk}.indexer_compressor_gate") + yield f"{sa}.indexer.index_kpool_compress_ape", _float(t(f"{blk}.indexer_compressor_ape")) + + mlp = f"{pre}.mlp" + if layer < config.first_k_dense_replace: + for proj, src, k in ( + ("gate_proj", "ffn_gate", hidden), + ("up_proj", "ffn_up", hidden), + ("down_proj", "ffn_down", config.intermediate_size), + ): + info = t(f"{blk}.{src}") + yield from fused(f"{mlp}.{proj}", [info], [matrix(info)], k) + else: + yield f"{mlp}.gate.weight", bf16(f"{blk}.ffn_gate_inp") + yield f"{mlp}.e_score_correction_bias", _float(t(f"{blk}.exp_probs_b.bias")) + shared_i = config.moe_intermediate_size * max(1, config.n_shared_experts) + for proj, src, k in ( + ("gate_proj", "ffn_gate_shexp", hidden), + ("up_proj", "ffn_up_shexp", hidden), + ("down_proj", "ffn_down_shexp", shared_i), + ): + info = t(f"{blk}.{src}") + yield from fused(f"{mlp}.shared_experts.{proj}", [info], [matrix(info)], k) + + +__all__ = ["iter_gguf_weights", "parse_gguf_config", "to_gguf"] 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/glm5_next/model.py b/python/freetoken/models/glm5_next/model.py index d2bfa684c1..d92977dea2 100644 --- a/python/freetoken/models/glm5_next/model.py +++ b/python/freetoken/models/glm5_next/model.py @@ -142,6 +142,8 @@ def __init__(self, config: ModelConfig, *, prefix: str = "model"): self.embed_tokens = VocabParallelEmbedding( num_embeddings=config.vocab_size, embedding_dim=config.hidden_size, + quant_config=config.quant, + prefix=f"{prefix}.embed_tokens", ) self.layers = OPList( [Glm5NextDecoderLayer(config, i, prefix=f"{prefix}.layers.{i}") for i in range(config.num_layers)] diff --git a/python/freetoken/models/glm5_next/weight.py b/python/freetoken/models/glm5_next/weight.py index 6f967dd8c7..09f9169419 100644 --- a/python/freetoken/models/glm5_next/weight.py +++ b/python/freetoken/models/glm5_next/weight.py @@ -250,7 +250,7 @@ def iter_vision_weights(model_path: str, device: torch.device) -> Iterator[tuple reader.close() -def iter_expert_pieces(model_path, config, kind: QuantKind, *, parallel: bool | None = False, workers: int = 8, chunk: int = 8 << 20): +def iter_expert_pieces(model_path, config, kind: QuantKind, *, layers=None, parallel: bool | None = False, workers: int = 8, chunk: int = 8 << 20): """Block-fp8 routed experts, one piece per expert: ``{gate, up, down}`` fp8 codes and their ``_scale`` companions; other kinds use the generic readers.""" if kind is not QuantKind.FP8_BLOCK: return None @@ -260,13 +260,14 @@ def iter_expert_pieces(model_path, config, kind: QuantKind, *, parallel: bool | from freetoken.moe.expert_pieces import per_expert_pieces suffix = {"weight": "", "weight_scale_inv": "_scale"} + keep = None if layers is None else set(layers) def locate(raw_name: str): m = _FP8_EXPERT_RE.match(raw_name) if m is None: return None bank = _layer_to_bank(int(m["layer"]), config) - if bank is None: + if bank is None or (keep is not None and bank not in keep): return None return bank, int(m["expert"]), m["proj"] + suffix[m["kind"]] @@ -284,8 +285,11 @@ def _serial(): weight_map = json.load(f)["weight_map"] reader = _ShardReader(folder, weight_map, torch.device("cpu")) try: - layers = range(config.first_k_dense_replace, config.num_layers) - for layer in tqdm(layers, desc="Loading GLM-5.3 fp8 experts (serial)", disable=not get_tp_info().is_primary()): + ckpt_layers = [ + layer for layer in range(config.first_k_dense_replace, config.num_layers) + if keep is None or _layer_to_bank(layer, config) in keep + ] + for layer in tqdm(ckpt_layers, desc="Loading GLM-5.3 fp8 experts (serial)", disable=not get_tp_info().is_primary()): for e in range(config.num_experts): base = f"{_CKPT}.layers.{layer}.mlp.experts.{e}" for proj in ("gate", "up", "down"): diff --git a/python/freetoken/models/gpt_oss/weight.py b/python/freetoken/models/gpt_oss/weight.py index 686f02466f..cb5830c3db 100644 --- a/python/freetoken/models/gpt_oss/weight.py +++ b/python/freetoken/models/gpt_oss/weight.py @@ -288,7 +288,7 @@ def _source_slices(config, tp_info) -> dict[str, tuple[str, tuple[slice, ...]]]: } -def iter_expert_pieces(model_path: str, config, kind: QuantKind, *, parallel: bool = False, workers: int = 8, chunk: int = 8 << 20): +def iter_expert_pieces(model_path: str, config, kind: QuantKind, *, layers=None, parallel: bool = False, workers: int = 8, chunk: int = 8 << 20): """gpt-oss experts as HF ships them, one piece per layer: stacked ``[E, ...]`` mxfp4 ``_blocks`` / ``_scales`` / ``_bias`` tensors sliced to this rank's intermediate range.""" if kind is not QuantKind.MXFP4: @@ -298,10 +298,11 @@ def iter_expert_pieces(model_path: str, config, kind: QuantKind, *, parallel: bo tp_info = get_tp_info() slices = _source_slices(config, tp_info) num_layers, num_experts = config.num_layers, config.num_experts + keep = None if layers is None else set(layers) def _is_expert(name: str) -> bool: info = _expert_layer_and_name(name) - return info is not None and 0 <= info[0] < num_layers + return info is not None and 0 <= info[0] < num_layers and (keep is None or info[0] in keep) def _tensors(): if parallel: @@ -320,6 +321,8 @@ def _tensors(): layer_id, source = info if not 0 <= layer_id < num_layers: raise ValueError(f"Unexpected GPT-OSS expert layer in checkpoint: {name}") + if keep is not None and layer_id not in keep: + continue if source not in slices: raise ValueError(f"Unexpected GPT-OSS expert source: {name}") yield name, _read_safetensor_slice(f, name, slices[source][1]) diff --git a/python/freetoken/models/llama/__init__.py b/python/freetoken/models/llama/__init__.py index 9ffad0e95f..31ffa447cd 100644 --- a/python/freetoken/models/llama/__init__.py +++ b/python/freetoken/models/llama/__init__.py @@ -1,5 +1,6 @@ from .config import parse_config +from .gguf import iter_gguf_weights, parse_gguf_config, to_gguf from .model import LlamaForCausalLM from .weight import iter_weights -__all__ = ["LlamaForCausalLM", "parse_config", "iter_weights"] +__all__ = ["LlamaForCausalLM", "parse_config", "iter_weights", "iter_gguf_weights", "parse_gguf_config", "to_gguf"] diff --git a/python/freetoken/models/llama/gguf.py b/python/freetoken/models/llama/gguf.py new file mode 100644 index 0000000000..833a5d66b5 --- /dev/null +++ b/python/freetoken/models/llama/gguf.py @@ -0,0 +1,189 @@ +"""Llama and Mistral from a llama.cpp ``llama`` GGUF: config from the GGUF metadata, tensors by name. + +llama.cpp runs interleaved rope on this arch, so its converter (conversion/llama.py) permutes the rows of each q and k head; the loader undoes that, and the model keeps HF's NeoX rope. +The RMSNorm weights are the HF ones. +""" + +from __future__ import annotations + +import re +from types import SimpleNamespace +from typing import TYPE_CHECKING, Iterator + +import torch + +from freetoken.models.gguf.reader import GgufTensorInfo, gguf_tensor_index +from freetoken.models.gguf.tensors import check_scales, float_tensor, fused, matrix + +from .config import parse_config + +if TYPE_CHECKING: + from freetoken.models.config import ModelConfig + from freetoken.models.gguf.config import GgufConfigShim + +_ARCH = "llama" +_LAYER_LEAVES = { + "self_attn.q_proj": "attn_q", + "self_attn.k_proj": "attn_k", + "self_attn.v_proj": "attn_v", + "self_attn.o_proj": "attn_output", + "mlp.gate_proj": "ffn_gate", + "mlp.up_proj": "ffn_up", + "mlp.down_proj": "ffn_down", +} +_QKV_LEAVES = ("self_attn.q_proj", "self_attn.k_proj", "self_attn.v_proj") +_TOP_LEAVES = {"model.embed_tokens": "token_embd", "lm_head": "output"} +_LAYER_RE = re.compile(r"^model\.layers\.(\d+)\.(.+)$") + + +def to_gguf(name: str, tensors: dict[str, int]) -> tuple[str, ...]: + """Checkpoint module name -> the GGUF tensors (without ``.weight``) that store it; empty for modules the GGUF has no matrix for.""" + if name in _TOP_LEAVES: + return (_TOP_LEAVES[name],) + m = _LAYER_RE.match(name) + if m is None: + return () + blk, leaf = f"blk.{m.group(1)}", m.group(2) + if leaf in _QKV_LEAVES and f"{blk}.attn_qkv.weight" in tensors: + return (f"{blk}.attn_qkv",) + tensor = _LAYER_LEAVES.get(leaf) + return (f"{blk}.{tensor}",) if tensor is not None else () + + +def _check_supported(shim: GgufConfigShim) -> None: + if any(k.startswith("prism.hadamard.") for k in shim.metadata): + raise NotImplementedError(f"{shim.model_path}: Hadamard-rotated (prism) GGUF weights are not supported") + check_scales(shim.model_path) + names = sorted(gguf_tensor_index(shim.model_path)) + # the llama arch also carries Mixtral experts and optional biases, which the llama model does not have + extra = [n for n in names if n.endswith(".bias") or ".ffn_gate_inp." in n] + if extra: + raise NotImplementedError(f"{shim.model_path}: llama GGUF tensors with biases or MoE experts are not supported (e.g. {extra[0]})") + + +def _rope_scaling(shim: GgufConfigShim, base: float, head_dim: int) -> dict | None: + """The HF rope scaling of the file: none, or Llama 3's, which the converter stores only as the ``rope_freqs`` divisors.""" + md = shim.metadata + scaling = md.get(f"{_ARCH}.rope.scaling.type", "none") + # llama.cpp scales linearly by rope.scaling.factor (or the older rope.scale_linear) when the type key is absent + factor = float(md.get(f"{_ARCH}.rope.scaling.factor", md.get(f"{_ARCH}.rope.scale_linear", 0.0))) + attn_factor = float(md.get(f"{_ARCH}.rope.scaling.attn_factor", 1.0)) + if scaling not in ("none", "linear") or factor not in (0.0, 1.0) or attn_factor != 1.0: + raise NotImplementedError(f"{shim.model_path}: rope scaling {scaling!r} (factor {factor}, attn_factor {attn_factor}) is not supported") + index = gguf_tensor_index(shim.model_path) + if not index: + raise NotImplementedError(f"{shim.model_path}: a metadata-only llama GGUF has no rope_freqs tensor to tell Llama 3 rope from plain rope") + info = index.get("rope_freqs.weight") + if info is None: + return None + from freetoken.layers.rotary import _rope_frequencies + + divisors = float_tensor(info).float() + # Llama 3.1-3.3 differ only in factor + llama3 = { + "rope_type": "llama3", + "factor": float(divisors.max()), + "low_freq_factor": 1.0, + "high_freq_factor": 4.0, + "original_max_position_embeddings": 8192, + } + plain, _ = _rope_frequencies(head_dim, head_dim, base, None) + scaled, _ = _rope_frequencies(head_dim, head_dim, base, llama3) + if not torch.allclose(plain / scaled, divisors, rtol=1e-5): + raise NotImplementedError(f"{shim.model_path}: rope_freqs does not match Llama 3 rope scaling") + return llama3 + + +def _hf_config(shim: GgufConfigShim) -> SimpleNamespace: + """The HF config fields ``parse_config`` reads, from the GGUF metadata.""" + md = shim.metadata + + def kv(key: str): + return md[f"{_ARCH}.{key}"] + + _check_supported(shim) + hidden, heads = int(kv("embedding_length")), int(kv("attention.head_count")) + # Mistral-7B files carry no key_length + head_dim = int(md.get(f"{_ARCH}.attention.key_length", hidden // heads)) + q = gguf_tensor_index(shim.model_path).get("blk.0.attn_q.weight") + # early Mistral-Nemo conversions lack key_length, so hidden / heads gives the wrong head_dim; llama.cpp rejects them too + if q is not None and q.ne[1] != heads * head_dim: + raise NotImplementedError(f"{shim.model_path}: attn_q has {q.ne[1]} rows, not head_count x head_dim = {heads * head_dim}") + if int(md.get(f"{_ARCH}.rope.dimension_count", head_dim)) != head_dim: + raise NotImplementedError(f"{shim.model_path}: partial rope is not supported") + base = float(md.get(f"{_ARCH}.rope.freq_base", 10000.0)) + return SimpleNamespace( + model_type="llama", + architectures=list(shim.architectures), + hidden_act="silu", + hidden_size=hidden, + num_hidden_layers=int(kv("block_count")), + num_attention_heads=heads, + num_key_value_heads=int(md.get(f"{_ARCH}.attention.head_count_kv", heads)), + head_dim=head_dim, + rope_theta=base, + rope_scaling=_rope_scaling(shim, base, head_dim), + max_position_embeddings=int(kv("context_length")), + rms_norm_eps=float(kv("attention.layer_norm_rms_epsilon")), + vocab_size=int(shim.vocab_size), + tie_word_embeddings=bool(shim.tie_word_embeddings), + intermediate_size=int(kv("feed_forward_length")), + ) + + +def parse_gguf_config(shim: GgufConfigShim) -> ModelConfig: + return parse_config(_hf_config(shim)) + + +def _unpermute(rows: torch.Tensor, heads: int) -> torch.Tensor: + """Undo the converter's rope permute of one q or k matrix (``LlamaModel.permute``); it moves whole rows, so packed ggml rows work too.""" + return rows.reshape(heads, rows.shape[0] // heads // 2, 2, *rows.shape[1:]).swapaxes(1, 2).reshape(rows.shape) + + +def iter_gguf_weights( + model_path: str, + device: torch.device, + *, + include_moe_experts: bool, + include_non_moe: bool, +) -> Iterator[tuple[str, torch.Tensor]]: + """The state dict from the GGUF.""" + if not include_non_moe: + return + from freetoken.utils import cached_load_hf_config + + hf = _hf_config(cached_load_hf_config(model_path)) + index = gguf_tensor_index(model_path) + hidden, heads, kv_heads = hf.hidden_size, hf.num_attention_heads, hf.num_key_value_heads + q_rows, kv_rows = heads * hf.head_dim, kv_heads * hf.head_dim + + def t(name: str) -> GgufTensorInfo: + return index[name if name in index else f"{name}.weight"] + + yield from fused("model.embed_tokens", [t("token_embd")], [matrix(t("token_embd"))], hidden) + if not hf.tie_word_embeddings: + yield from fused("lm_head", [t("output")], [matrix(t("output"))], hidden) + yield "model.norm.weight", float_tensor(t("output_norm")) + + for layer in range(hf.num_hidden_layers): + blk, pre = f"blk.{layer}", f"model.layers.{layer}" + yield f"{pre}.input_layernorm.weight", float_tensor(t(f"{blk}.attn_norm")) + yield f"{pre}.post_attention_layernorm.weight", float_tensor(t(f"{blk}.ffn_norm")) + sa = f"{pre}.self_attn" + # convert_hf_to_gguf.py --fuse-qkv stores q, k and v as one attn_qkv matrix + if f"{blk}.attn_qkv.weight" in index: + qkv = [t(f"{blk}.attn_qkv")] * 3 + q, k, v = matrix(qkv[0]).split([q_rows, kv_rows, kv_rows]) + else: + qkv = [t(f"{blk}.attn_{p}") for p in ("q", "k", "v")] + q, k, v = (matrix(p) for p in qkv) + yield from fused(f"{sa}.qkv_proj", qkv, [_unpermute(q, heads), _unpermute(k, kv_heads), v], hidden) + out = t(f"{blk}.attn_output") + yield from fused(f"{sa}.o_proj", [out], [matrix(out)], int(out.ne[0])) + mlp = f"{pre}.mlp" + gate, up, down = (t(f"{blk}.ffn_{p}") for p in ("gate", "up", "down")) + yield from fused(f"{mlp}.gate_up_proj", [gate, up], [matrix(gate), matrix(up)], hidden) + yield from fused(f"{mlp}.down_proj", [down], [matrix(down)], int(down.ne[0])) + + +__all__ = ["iter_gguf_weights", "parse_gguf_config", "to_gguf"] diff --git a/python/freetoken/models/llama/model.py b/python/freetoken/models/llama/model.py index eb7ff28fac..107d9e0a2c 100644 --- a/python/freetoken/models/llama/model.py +++ b/python/freetoken/models/llama/model.py @@ -48,6 +48,8 @@ def __init__(self, config: ModelConfig, *, prefix: str = "model"): self.embed_tokens = VocabParallelEmbedding( num_embeddings=config.vocab_size, embedding_dim=config.hidden_size, + quant_config=config.quant, + prefix=f"{prefix}.embed_tokens", ) self.layers = OPList( [ diff --git a/python/freetoken/models/loader.py b/python/freetoken/models/loader.py index 0f310ccbb2..4a9ac94c73 100644 --- a/python/freetoken/models/loader.py +++ b/python/freetoken/models/loader.py @@ -208,11 +208,11 @@ def close(self) -> None: def nvfp4_parts_ct(f, raw_base: str) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor | None]: - """compressed-tensors NVFP4 -> ``(packed uint8 [O, IN//2], block scale fp8 [O, IN//16], per-output-row global fp16 [O], dequant-side input_scale fp32 scalar or None)`` for the NVFP4 linear method.""" + """compressed-tensors NVFP4 -> ``(packed uint8 [O, IN//2], block scale fp8 [O, IN//16], per-output-row global fp32 [O], dequant-side input_scale fp32 scalar or None)`` for the NVFP4 linear method.""" w = f.get_tensor(raw_base + ".weight_packed") s = f.get_tensor(raw_base + ".weight_scale") wg = f.get_tensor(raw_base + ".weight_global_scale").reshape(1).to(torch.float32) - g = (1.0 / wg).to(torch.float16).expand(w.shape[0]).contiguous() + g = (1.0 / wg).expand(w.shape[0]).contiguous() a = None if f.has(raw_base + ".input_global_scale"): a = (1.0 / f.get_tensor(raw_base + ".input_global_scale").reshape(()).to(torch.float32)) diff --git a/python/freetoken/models/minimax_m2/__init__.py b/python/freetoken/models/minimax_m2/__init__.py index ad276f8457..edb9896d24 100644 --- a/python/freetoken/models/minimax_m2/__init__.py +++ b/python/freetoken/models/minimax_m2/__init__.py @@ -1,4 +1,5 @@ from .config import parse_config +from .gguf import iter_gguf_weights, parse_gguf_config, to_gguf from .model import MiniMaxM2ForCausalLM from .weight import iter_weights, nvfp4_expert_spec @@ -7,4 +8,7 @@ "MiniMaxM2ForCausalLM", "parse_config", "iter_weights", + "iter_gguf_weights", + "parse_gguf_config", + "to_gguf", ] diff --git a/python/freetoken/models/minimax_m2/gguf.py b/python/freetoken/models/minimax_m2/gguf.py new file mode 100644 index 0000000000..99047b404c --- /dev/null +++ b/python/freetoken/models/minimax_m2/gguf.py @@ -0,0 +1,161 @@ +"""MiniMax-M2 from a llama.cpp ``minimax-m2`` GGUF: config from the GGUF metadata, tensors by name; the routed experts load as GGUF expert banks (``freetoken.models.gguf.experts``). + +The llama.cpp converter (conversion/minimax.py) only stacks the w1 / w2 / w3 experts: both sides run NeoX rope, so q/k carry no rope permute, and the RMSNorm weights are the HF ones. +""" + +from __future__ import annotations + +import re +from dataclasses import replace +from types import SimpleNamespace +from typing import TYPE_CHECKING, Iterator + +import torch + +from freetoken.models.gguf.reader import GgufTensorInfo, gguf_tensor_index +from freetoken.models.gguf.tensors import float_tensor, fused, matrix + +from .config import parse_config + +if TYPE_CHECKING: + from freetoken.models.config import ModelConfig + from freetoken.models.gguf.config import GgufConfigShim + +_ARCH = "minimax-m2" +_LAYER_LEAVES = { + "self_attn.q_proj": "attn_q", + "self_attn.k_proj": "attn_k", + "self_attn.v_proj": "attn_v", + "self_attn.o_proj": "attn_output", + "block_sparse_moe.gate": "ffn_gate_inp", + "block_sparse_moe.experts.0.w1": "ffn_gate_exps", + "block_sparse_moe.experts.0.w3": "ffn_up_exps", + "block_sparse_moe.experts.0.w2": "ffn_down_exps", +} +_QKV_LEAVES = ("self_attn.q_proj", "self_attn.k_proj", "self_attn.v_proj") +_TOP_LEAVES = {"model.embed_tokens": "token_embd", "lm_head": "output"} +_LAYER_RE = re.compile(r"^model\.layers\.(\d+)\.(.+)$") +_SIGMOID_GATING = 2 + + +def to_gguf(name: str, tensors: dict[str, int]) -> tuple[str, ...]: + """Checkpoint module name -> the GGUF tensors (without ``.weight``) that store it; empty for modules the GGUF has no matrix for.""" + if name in _TOP_LEAVES: + return (_TOP_LEAVES[name],) + m = _LAYER_RE.match(name) + if m is None: + return () + blk, leaf = f"blk.{m.group(1)}", m.group(2) + if leaf in _QKV_LEAVES and f"{blk}.attn_qkv.weight" in tensors: + return (f"{blk}.attn_qkv",) + tensor = _LAYER_LEAVES.get(leaf) + return (f"{blk}.{tensor}",) if tensor is not None else () + + +def _check_supported(shim: GgufConfigShim) -> None: + if any(k.startswith("prism.hadamard.") for k in shim.metadata): + raise NotImplementedError(f"{shim.model_path}: Hadamard-rotated (prism) GGUF weights are not supported") + index = gguf_tensor_index(shim.model_path) + scaled = sorted(n for n in index if n.endswith((".scale", ".input_scale"))) + if scaled: + raise NotImplementedError(f"{shim.model_path}: GGUF tensors with per-tensor scales are not supported (e.g. {scaled[0]})") + # a metadata-only GGUF (an FTW dir) has no tensor table to check + if index and "blk.0.exp_probs_b.bias" not in index: + raise NotImplementedError(f"{shim.model_path}: converted by an older llama.cpp that names the router bias exp_probs_b, which current llama.cpp refuses too") + + +def _hf_config(shim: GgufConfigShim) -> SimpleNamespace: + """The HF config fields ``parse_config`` reads, from the GGUF metadata.""" + md = shim.metadata + + def kv(key: str): + return md[f"{_ARCH}.{key}"] + + _check_supported(shim) + gating = int(md.get(f"{_ARCH}.expert_gating_func", 0)) + if gating != _SIGMOID_GATING: + raise NotImplementedError(f"{shim.model_path}: expert gating function {gating} is not supported; MiniMax-M2 routes with a sigmoid (2)") + scaling = md.get(f"{_ARCH}.rope.scaling.type", "none") + if scaling != "none": + raise NotImplementedError(f"{shim.model_path}: rope scaling {scaling!r} is not supported") + # a list holds one size per layer + expert_ff = kv("expert_feed_forward_length") + if isinstance(expert_ff, list): + if len(set(expert_ff)) != 1: + raise NotImplementedError(f"{shim.model_path}: per-layer expert sizes {sorted(set(expert_ff))} are not supported") + expert_ff = expert_ff[0] + return SimpleNamespace( + model_type="minimax_m2", + architectures=list(shim.architectures), + hidden_act="silu", + hidden_size=int(kv("embedding_length")), + num_hidden_layers=int(kv("block_count")), + num_attention_heads=int(kv("attention.head_count")), + num_key_value_heads=int(kv("attention.head_count_kv")), + head_dim=int(kv("attention.key_length")), + rotary_dim=int(kv("rope.dimension_count")), + rope_theta=float(kv("rope.freq_base")), + max_position_embeddings=int(kv("context_length")), + rms_norm_eps=float(kv("attention.layer_norm_rms_epsilon")), + vocab_size=int(shim.vocab_size), + tie_word_embeddings=bool(shim.tie_word_embeddings), + intermediate_size=int(expert_ff), + num_local_experts=int(kv("expert_count")), + num_experts_per_tok=int(kv("expert_used_count")), + # llama.cpp always builds the per-layer q/k norms of this arch + use_qk_norm=True, + ) + + +def parse_gguf_config(shim: GgufConfigShim) -> ModelConfig: + return replace(parse_config(_hf_config(shim)), expert_quant="gguf") + + +def iter_gguf_weights( + model_path: str, + device: torch.device, + *, + include_moe_experts: bool, + include_non_moe: bool, +) -> Iterator[tuple[str, torch.Tensor]]: + """The dense state dict from the GGUF; the routed experts load as GGUF expert banks (``freetoken.models.gguf.experts``).""" + if not include_non_moe: + return + from freetoken.utils import cached_load_hf_config + + hf = _hf_config(cached_load_hf_config(model_path)) + index = gguf_tensor_index(model_path) + hidden = hf.hidden_size + q_rows, kv_rows = hf.num_attention_heads * hf.head_dim, hf.num_key_value_heads * hf.head_dim + + def t(name: str) -> GgufTensorInfo: + return index[name if name in index else f"{name}.weight"] + + yield from fused("model.embed_tokens", [t("token_embd")], [matrix(t("token_embd"))], hidden) + if not hf.tie_word_embeddings: + yield from fused("lm_head", [t("output")], [matrix(t("output"))], hidden) + yield "model.norm.weight", float_tensor(t("output_norm")) + + for layer in range(hf.num_hidden_layers): + blk, pre = f"blk.{layer}", f"model.layers.{layer}" + yield f"{pre}.input_layernorm.weight", float_tensor(t(f"{blk}.attn_norm")) + yield f"{pre}.post_attention_layernorm.weight", float_tensor(t(f"{blk}.ffn_norm")) + sa = f"{pre}.self_attn" + # convert_hf_to_gguf.py --fuse-qkv stores q, k and v as one attn_qkv matrix + if f"{blk}.attn_qkv.weight" in index: + qkv = [t(f"{blk}.attn_qkv")] * 3 + parts = list(matrix(qkv[0]).split([q_rows, kv_rows, kv_rows])) + else: + qkv = [t(f"{blk}.attn_{p}") for p in ("q", "k", "v")] + parts = [matrix(p) for p in qkv] + yield from fused(f"{sa}.qkv_proj", qkv, parts, hidden) + yield f"{sa}.q_norm.weight", float_tensor(t(f"{blk}.attn_q_norm")) + yield f"{sa}.k_norm.weight", float_tensor(t(f"{blk}.attn_k_norm")) + out = t(f"{blk}.attn_output") + yield from fused(f"{sa}.o_proj", [out], [matrix(out)], int(out.ne[0])) + moe = f"{pre}.block_sparse_moe" + yield f"{moe}.gate.weight", float_tensor(t(f"{blk}.ffn_gate_inp")) + yield f"{moe}.e_score_correction_bias", float_tensor(t(f"{blk}.exp_probs_b.bias")) + + +__all__ = ["iter_gguf_weights", "parse_gguf_config", "to_gguf"] diff --git a/python/freetoken/models/minimax_m2/model.py b/python/freetoken/models/minimax_m2/model.py index 457ab56500..2f29a87353 100644 --- a/python/freetoken/models/minimax_m2/model.py +++ b/python/freetoken/models/minimax_m2/model.py @@ -41,6 +41,8 @@ def __init__(self, config: ModelConfig, *, prefix: str = "model"): self.embed_tokens = VocabParallelEmbedding( num_embeddings=config.vocab_size, embedding_dim=config.hidden_size, + quant_config=config.quant, + prefix=f"{prefix}.embed_tokens", ) self.layers = OPList( [ diff --git a/python/freetoken/models/minimax_m3/__init__.py b/python/freetoken/models/minimax_m3/__init__.py index 4abf9c8793..8d0aea79bf 100644 --- a/python/freetoken/models/minimax_m3/__init__.py +++ b/python/freetoken/models/minimax_m3/__init__.py @@ -1,4 +1,5 @@ from .config import VisionConfig, parse_config, parse_vision_config +from .gguf import iter_gguf_weights, parse_gguf_config, to_gguf from .model import MiniMaxM3ForCausalLM, MiniMaxM3ForConditionalGeneration from .vision import MiniMaxM3VisionModel from .weight import ( @@ -17,4 +18,7 @@ "parse_vision_config", "iter_vision_weights", "iter_weights", + "iter_gguf_weights", + "parse_gguf_config", + "to_gguf", ] diff --git a/python/freetoken/models/minimax_m3/attention.py b/python/freetoken/models/minimax_m3/attention.py index 0d55c07294..4fcf57e103 100644 --- a/python/freetoken/models/minimax_m3/attention.py +++ b/python/freetoken/models/minimax_m3/attention.py @@ -45,6 +45,7 @@ def __init__(self, config: ModelConfig, layer_id: int, *, prefix: str = ""): self.qkv_proj = LinearReplicated( args.hidden_size, self.qo_attn_dim + 2 * self.kv_attn_dim, has_bias=False, + output_sizes=[self.qo_attn_dim, self.kv_attn_dim, self.kv_attn_dim], quant_config=config.quant, prefix=f"{prefix}.qkv_proj", ) self.o_proj = LinearReplicated( diff --git a/python/freetoken/models/minimax_m3/gguf.py b/python/freetoken/models/minimax_m3/gguf.py new file mode 100644 index 0000000000..92742c3aae --- /dev/null +++ b/python/freetoken/models/minimax_m3/gguf.py @@ -0,0 +1,198 @@ +"""MiniMax-M3 text from a llama.cpp ``minimax-m3`` GGUF: config from the GGUF metadata, tensors by name; the routed experts load as GGUF expert banks (``freetoken.models.gguf.experts``). + +The llama.cpp converter (conversion/minimax.py) stacks the w1 / w2 / w3 experts, stores the indexer projections as F32, and adds 1 to every RMSNorm weight so llama.cpp can run plain RMSNorm; the loader subtracts the 1 for this family's (1 + w) norms. +Both sides run NeoX rope, so q/k carry no rope permute. +""" + +from __future__ import annotations + +import re +from dataclasses import replace +from types import SimpleNamespace +from typing import TYPE_CHECKING, Iterator + +import torch + +from freetoken.models.gguf.reader import GgufTensorInfo, gguf_tensor_index +from freetoken.models.gguf.tensors import float_tensor, fused, matrix + +from .config import parse_config + +if TYPE_CHECKING: + from freetoken.models.config import ModelConfig + from freetoken.models.gguf.config import GgufConfigShim + +_ARCH = "minimax-m3" +_LAYER_LEAVES = { + "self_attn.q_proj": "attn_q", + "self_attn.k_proj": "attn_k", + "self_attn.v_proj": "attn_v", + "self_attn.o_proj": "attn_output", + "self_attn.index_q_proj": "indexer.q_proj", + "self_attn.index_k_proj": "indexer.k_proj", + "mlp.gate_proj": "ffn_gate", + "mlp.up_proj": "ffn_up", + "mlp.down_proj": "ffn_down", + "block_sparse_moe.gate": "ffn_gate_inp", + "block_sparse_moe.shared_experts.gate_proj": "ffn_gate_shexp", + "block_sparse_moe.shared_experts.up_proj": "ffn_up_shexp", + "block_sparse_moe.shared_experts.down_proj": "ffn_down_shexp", + "block_sparse_moe.experts.0.w1": "ffn_gate_exps", + "block_sparse_moe.experts.0.w3": "ffn_up_exps", + "block_sparse_moe.experts.0.w2": "ffn_down_exps", +} +_TOP_LEAVES = {"model.embed_tokens": "token_embd", "lm_head": "output"} +_LAYER_RE = re.compile(r"^model\.layers\.(\d+)\.(.+)$") +_SIGMOID_GATING = 2 +# llama.cpp hard-codes these for the arch (src/llama-graph.cpp LLM_FFN_SWIGLU_OAI_MOE) +_SWIGLU_ALPHA, _SWIGLU_LIMIT = 1.702, 7.0 + + +def to_gguf(name: str, tensors: dict[str, int]) -> tuple[str, ...]: + """Checkpoint module name -> the GGUF tensors (without ``.weight``) that store it; empty for modules the GGUF has no matrix for.""" + if name in _TOP_LEAVES: + return (_TOP_LEAVES[name],) + m = _LAYER_RE.match(name) + if m is None: + return () + tensor = _LAYER_LEAVES.get(m.group(2)) + return (f"blk.{m.group(1)}.{tensor}",) if tensor is not None else () + + +def _check_supported(shim: GgufConfigShim) -> None: + if any(k.startswith("prism.hadamard.") for k in shim.metadata): + raise NotImplementedError(f"{shim.model_path}: Hadamard-rotated (prism) GGUF weights are not supported") + scaled = sorted(n for n in gguf_tensor_index(shim.model_path) if n.endswith((".scale", ".input_scale"))) + if scaled: + raise NotImplementedError(f"{shim.model_path}: GGUF tensors with per-tensor scales are not supported (e.g. {scaled[0]})") + + +def _hf_config(shim: GgufConfigShim) -> SimpleNamespace: + """The HF text config fields ``parse_config`` reads, from the GGUF metadata.""" + md = shim.metadata + + def kv(key: str): + return md[f"{_ARCH}.{key}"] + + _check_supported(shim) + if f"{_ARCH}.attention.indexer.head_count" not in md: + raise NotImplementedError(f"{shim.model_path}: converted by an older llama.cpp without the sparse-attention indexer, which current llama.cpp refuses too") + gating = int(kv("expert_gating_func")) + if gating != _SIGMOID_GATING: + raise NotImplementedError(f"{shim.model_path}: expert gating function {gating} is not supported; MiniMax-M3 routes with a sigmoid (2)") + if not md.get(f"{_ARCH}.expert_weights_norm", False): + raise NotImplementedError(f"{shim.model_path}: routing without renormalized expert weights is not supported") + scaling = md.get(f"{_ARCH}.rope.scaling.type", "none") + if scaling != "none": + raise NotImplementedError(f"{shim.model_path}: rope scaling {scaling!r} is not supported") + expert_ff = kv("expert_feed_forward_length") + if isinstance(expert_ff, list): + if len(set(expert_ff)) != 1: + raise NotImplementedError(f"{shim.model_path}: per-layer expert sizes {sorted(set(expert_ff))} are not supported") + expert_ff = expert_ff[0] + num_layers = int(kv("block_count")) + dense = int(md.get(f"{_ARCH}.leading_dense_block_count", 0)) + # llama.cpp runs sparse attention and the MoE FFN on every layer after the leading dense ones + sparse = [0] * dense + [1] * (num_layers - dense) + return SimpleNamespace( + model_type="minimax_m3", + architectures=list(shim.architectures), + hidden_act="swigluoai", + swiglu_alpha=_SWIGLU_ALPHA, + swiglu_limit=_SWIGLU_LIMIT, + hidden_size=int(kv("embedding_length")), + num_hidden_layers=num_layers, + num_attention_heads=int(kv("attention.head_count")), + num_key_value_heads=int(kv("attention.head_count_kv")), + head_dim=int(kv("attention.key_length")), + rotary_dim=int(kv("rope.dimension_count")), + rope_theta=float(kv("rope.freq_base")), + max_position_embeddings=int(kv("context_length")), + rms_norm_eps=float(kv("attention.layer_norm_rms_epsilon")), + vocab_size=int(shim.vocab_size), + tie_word_embeddings=bool(shim.tie_word_embeddings), + dense_intermediate_size=int(kv("feed_forward_length")), + intermediate_size=int(expert_ff), + shared_intermediate_size=int(expert_ff), + n_shared_experts=int(kv("expert_shared_count")), + num_local_experts=int(kv("expert_count")), + num_experts_per_tok=int(kv("expert_used_count")), + # llama.cpp skips the scale when the key is absent or 0 + routed_scaling_factor=float(md.get(f"{_ARCH}.expert_weights_scale", 0.0)) or 1.0, + use_routing_bias=True, + moe_layer_freq=sparse, + sparse_attention_config={ + "use_sparse_attention": True, + "sparse_attention_freq": sparse, + "sparse_num_index_heads": int(kv("attention.indexer.head_count")), + "sparse_index_dim": int(kv("attention.indexer.key_length")), + "sparse_topk_blocks": int(kv("attention.indexer.top_k")), + "sparse_block_size": int(kv("attention.indexer.block_size")), + "sparse_local_block": int(kv("attention.indexer.local_blocks")), + "sparse_init_block": 0, + "sparse_score_type": "max", + }, + ) + + +def parse_gguf_config(shim: GgufConfigShim) -> ModelConfig: + return replace(parse_config(_hf_config(shim)), expert_quant="gguf") + + +def iter_gguf_weights( + model_path: str, + device: torch.device, + *, + include_moe_experts: bool, + include_non_moe: bool, +) -> Iterator[tuple[str, torch.Tensor]]: + """The dense text state dict from the GGUF; the routed experts load as GGUF expert banks (``freetoken.models.gguf.experts``).""" + if not include_non_moe: + return + from freetoken.utils import cached_load_hf_config + + config = parse_gguf_config(cached_load_hf_config(model_path)) + args = config.m3_args + index = gguf_tensor_index(model_path) + hidden = config.hidden_size + + def t(name: str) -> GgufTensorInfo: + return index[name if name in index else f"{name}.weight"] + + def norm(name: str) -> torch.Tensor: + return float_tensor(t(name)).float() - 1.0 + + yield from fused("model.embed_tokens", [t("token_embd")], [matrix(t("token_embd"))], hidden) + if not config.tie_word_embeddings: + yield from fused("lm_head", [t("output")], [matrix(t("output"))], hidden) + yield "model.norm.weight", norm("output_norm") + + for layer in range(config.num_layers): + blk, pre = f"blk.{layer}", f"model.layers.{layer}" + yield f"{pre}.input_layernorm.weight", norm(f"{blk}.attn_norm") + yield f"{pre}.post_attention_layernorm.weight", norm(f"{blk}.ffn_norm") + sa = f"{pre}.self_attn" + qkv = [t(f"{blk}.attn_{p}") for p in ("q", "k", "v")] + yield from fused(f"{sa}.qkv_proj", qkv, [matrix(p) for p in qkv], hidden) + out = t(f"{blk}.attn_output") + yield from fused(f"{sa}.o_proj", [out], [matrix(out)], int(out.ne[0])) + yield f"{sa}.q_norm.weight", norm(f"{blk}.attn_q_norm") + yield f"{sa}.k_norm.weight", norm(f"{blk}.attn_k_norm") + if args.is_sparse_layer(layer): + iq, ik = t(f"{blk}.indexer.q_proj"), t(f"{blk}.indexer.k_proj") + yield from fused(f"{sa}.index_qk_proj", [iq, ik], [matrix(iq), matrix(ik)], hidden) + yield f"{sa}.index_q_norm.weight", norm(f"{blk}.indexer.q_norm") + yield f"{sa}.index_k_norm.weight", norm(f"{blk}.indexer.k_norm") + if layer in args.moe_layer_ids: + moe = f"{pre}.block_sparse_moe" + yield f"{moe}.gate.weight", float_tensor(t(f"{blk}.ffn_gate_inp")).float() + yield f"{moe}.e_score_correction_bias", float_tensor(t(f"{blk}.exp_probs_b.bias")).float() + mlp, parts = f"{moe}.shared_experts", [t(f"{blk}.ffn_{p}_shexp") for p in ("gate", "up", "down")] + else: + mlp, parts = f"{pre}.mlp", [t(f"{blk}.ffn_{p}") for p in ("gate", "up", "down")] + gate, up, down = parts + yield from fused(f"{mlp}.gate_up_proj", [gate, up], [matrix(gate), matrix(up)], hidden) + yield from fused(f"{mlp}.down_proj", [down], [matrix(down)], int(down.ne[0])) + + +__all__ = ["iter_gguf_weights", "parse_gguf_config", "to_gguf"] diff --git a/python/freetoken/models/minimax_m3/mlp.py b/python/freetoken/models/minimax_m3/mlp.py index d5332eed0c..62b02e8040 100644 --- a/python/freetoken/models/minimax_m3/mlp.py +++ b/python/freetoken/models/minimax_m3/mlp.py @@ -36,6 +36,7 @@ def __init__( ): self.gate_up_proj = LinearReplicated( hidden_size, 2 * intermediate_size, has_bias=False, + output_sizes=[intermediate_size, intermediate_size], quant_config=quant_config, prefix=f"{prefix}.gate_up_proj", ) self.down_proj = LinearReplicated( diff --git a/python/freetoken/models/minimax_m3/model.py b/python/freetoken/models/minimax_m3/model.py index d78fc10afd..b6f271b35c 100644 --- a/python/freetoken/models/minimax_m3/model.py +++ b/python/freetoken/models/minimax_m3/model.py @@ -71,6 +71,8 @@ def __init__(self, config: ModelConfig, *, prefix: str = "model"): self.embed_tokens = VocabParallelEmbedding( num_embeddings=config.vocab_size, embedding_dim=config.hidden_size, + quant_config=config.quant, + prefix=f"{prefix}.embed_tokens", ) self.layers = OPList( [ diff --git a/python/freetoken/models/muse_glimmer/__init__.py b/python/freetoken/models/muse_glimmer/__init__.py index 4ee14808c4..ff8d3d1710 100644 --- a/python/freetoken/models/muse_glimmer/__init__.py +++ b/python/freetoken/models/muse_glimmer/__init__.py @@ -1,5 +1,6 @@ from .attention import MuseGlimmerAttention from .config import VisionConfig, parse_config, parse_vision_config +from .gguf import iter_gguf_weights, parse_gguf_config, to_gguf from .model import MuseGlimmerForCausalLM, MuseGlimmerForConditionalGeneration from .vision import MuseGlimmerVisionModel from .weight import iter_vision_weights, iter_weights @@ -14,4 +15,7 @@ "parse_vision_config", "iter_vision_weights", "iter_weights", + "iter_gguf_weights", + "parse_gguf_config", + "to_gguf", ] diff --git a/python/freetoken/models/muse_glimmer/attention.py b/python/freetoken/models/muse_glimmer/attention.py index 7c814d4c1b..a4af6e5210 100644 --- a/python/freetoken/models/muse_glimmer/attention.py +++ b/python/freetoken/models/muse_glimmer/attention.py @@ -18,16 +18,16 @@ class MuseGlimmerAttention(BaseOP): """Gated GQA for one sliding-window or full-attention (NoPE) layer. - q, k, v, gate = split(qkvg_proj(x)) + q, k, v = split(qkv_proj(x)) q = qk_norm(q); k = qk_norm(k) # weightless per-head RMSNorm q, k = rope(q, k) # sliding layers only; full layers are NoPE attn = paged_attention(q, k, v) # sm_scale carries qk_scale_factor/sqrt(d) - out = o_proj(attn * sigmoid(gate)) + out = o_proj(attn * sigmoid(gate_proj(x))) The reference multiplies q by ``qk_scale_factor`` before rope; rope is a rotation, so the factor commutes and is folded into ``AttentionSpec.sm_scale`` instead (parse_config - sets ``attn_sm_scale = qk_scale_factor / sqrt(head_dim)``). The gate is computed from - the same normed layer input as q/k/v, so it rides the fused projection. + sets ``attn_sm_scale = qk_scale_factor / sqrt(head_dim)``). The gate stays its own + linear: the NVFP4 release quantizes it with a different global scale than q/k/v. """ def __init__(self, config: ModelConfig, layer_id: int, *, prefix: str = ""): @@ -41,10 +41,14 @@ def __init__(self, config: ModelConfig, layer_id: int, *, prefix: str = ""): self.qo_attn_dim = self.num_q * head_dim self.kv_attn_dim = self.num_kv * head_dim - self._qkvg_split = [self.qo_attn_dim, self.kv_attn_dim, self.kv_attn_dim, self.qo_attn_dim] - self.qkvg_proj = LinearColParallelMerged( - config.hidden_size, self._qkvg_split, has_bias=False, - quant_config=config.quant, prefix=f"{prefix}.qkvg_proj", + self._qkv_split = [self.qo_attn_dim, 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.gate_proj = LinearColParallelMerged( + config.hidden_size, [self.qo_attn_dim], has_bias=False, + quant_config=config.quant, prefix=f"{prefix}.gate_proj", ) self.o_proj = LinearReplicated( self.qo_attn_dim, config.hidden_size, has_bias=False, @@ -76,12 +80,13 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: ctx = get_global_ctx() T = x.shape[0] - qkvg = self.qkvg_proj.forward(x) - q, k, v, gate = torch.split(qkvg, self._qkvg_split, dim=-1) - del qkvg + gate = self.gate_proj.forward(x) + qkv = self.qkv_proj.forward(x) + q, k, v = torch.split(qkv, self._qkv_split, dim=-1) + del qkv q = self.qk_norm.forward(q.contiguous().view(T, self.num_q, self.head_dim)) k = self.qk_norm.forward(k.contiguous().view(T, self.num_kv, self.head_dim)) - v = v.contiguous() # the split view has the qkvg row stride; the KV store needs contiguous + v = v.contiguous() # the split view has the qkv row stride; the KV store needs contiguous q = q.reshape(T, self.qo_attn_dim) k = k.reshape(T, self.kv_attn_dim) diff --git a/python/freetoken/models/muse_glimmer/gguf.py b/python/freetoken/models/muse_glimmer/gguf.py new file mode 100644 index 0000000000..cad5e633d8 --- /dev/null +++ b/python/freetoken/models/muse_glimmer/gguf.py @@ -0,0 +1,185 @@ +"""Muse Glimmer text from a llama.cpp ``muse-glimmer`` GGUF: config from the GGUF metadata, tensors by name. + +The converter (conversion/muse_glimmer.py) adds 1 to the four norms of each block, permutes the rows of each q and k head for llama.cpp's interleaved rope, and writes constant q/k norms that carry ``qk_scale_factor``. +The loader subtracts the 1, undoes the permute, and reads ``qk_scale_factor`` back from the q norm. +""" + +from __future__ import annotations + +import re +from types import SimpleNamespace +from typing import TYPE_CHECKING, Iterator + +import torch + +from freetoken.models.gguf.reader import GgufTensorInfo, gguf_tensor_index +from freetoken.models.gguf.tensors import float_tensor, fused, matrix + +from .config import parse_config + +if TYPE_CHECKING: + from freetoken.models.config import ModelConfig + from freetoken.models.gguf.config import GgufConfigShim + +_ARCH = "muse-glimmer" +_LAYER_LEAVES = { + "self_attn.q_proj": "attn_q", + "self_attn.k_proj": "attn_k", + "self_attn.v_proj": "attn_v", + "self_attn.gate_proj": "attn_gate", + "self_attn.o_proj": "attn_output", + "mlp.gate_proj": "ffn_gate", + "mlp.up_proj": "ffn_up", + "mlp.down_proj": "ffn_down", +} +_QKV_LEAVES = ("self_attn.q_proj", "self_attn.k_proj", "self_attn.v_proj") +_TOP_LEAVES = {"model.embed_tokens": "token_embd", "lm_head": "output"} +_LAYER_RE = re.compile(r"^model\.layers\.(\d+)\.(.+)$") +# checkpoint norm -> GGUF norm; the GGUF stores each as weight + 1 +_SHIFTED_NORMS = { + "input_layernorm": "attn_norm", + "post_attention_layernorm": "post_attention_norm", + "pre_feedforward_layernorm": "ffn_norm", + "post_feedforward_layernorm": "post_ffw_norm", +} + + +def to_gguf(name: str, tensors: dict[str, int]) -> tuple[str, ...]: + """Checkpoint module name -> the GGUF tensors (without ``.weight``) that store it; empty for modules the GGUF has no matrix for.""" + if name in _TOP_LEAVES: + return (_TOP_LEAVES[name],) + m = _LAYER_RE.match(name) + if m is None: + return () + blk, leaf = f"blk.{m.group(1)}", m.group(2) + if leaf in _QKV_LEAVES and f"{blk}.attn_qkv.weight" in tensors: + return (f"{blk}.attn_qkv",) + tensor = _LAYER_LEAVES.get(leaf) + return (f"{blk}.{tensor}",) if tensor is not None else () + + +def _check_supported(shim: GgufConfigShim) -> None: + index = gguf_tensor_index(shim.model_path) + if not index: + raise NotImplementedError(f"{shim.model_path}: a metadata-only muse-glimmer GGUF has no q norm to read qk_scale_factor from") + if any(k.startswith("prism.hadamard.") for k in shim.metadata): + raise NotImplementedError(f"{shim.model_path}: Hadamard-rotated (prism) GGUF weights are not supported") + scaled = sorted(n for n in index if n.endswith((".scale", ".input_scale"))) + if scaled: + raise NotImplementedError(f"{shim.model_path}: GGUF tensors with per-tensor scales are not supported (e.g. {scaled[0]})") + + +def _qk_scale_factor(model_path: str, num_layers: int) -> float: + """The constant the converter wrote into every q norm; the k norms are all ones.""" + index = gguf_tensor_index(model_path) + scale = float(float_tensor(index["blk.0.attn_q_norm.weight"])[0]) + for layer in range(num_layers): + q, k = (float_tensor(index[f"blk.{layer}.attn_{p}_norm.weight"]).float() for p in ("q", "k")) + if not (torch.all(q == scale) and torch.all(k == 1.0)): + raise NotImplementedError(f"{model_path}: the q/k norms of layer {layer} are not the converter's constant qk_scale_factor / ones") + return scale + + +def _hf_config(shim: GgufConfigShim) -> SimpleNamespace: + """The HF text config fields ``parse_config`` reads, from the GGUF metadata.""" + md = shim.metadata + + def kv(key: str): + return md[f"{_ARCH}.{key}"] + + _check_supported(shim) + num_layers = int(kv("block_count")) + # a bool per layer, or llama.cpp's period: every n-th layer is full attention, and 0 means none is + pattern = md.get(f"{_ARCH}.attention.sliding_window_pattern", 4) + if isinstance(pattern, list): + swa = [bool(p) for p in pattern] + if len(swa) != num_layers: + raise ValueError(f"{shim.model_path}: sliding_window_pattern has {len(swa)} entries for {num_layers} layers") + else: + swa = [int(pattern) == 0 or i % int(pattern) < int(pattern) - 1 for i in range(num_layers)] + base = float(kv("rope.freq_base")) + swa_base = float(md.get(f"{_ARCH}.rope.freq_base_swa", base)) + return SimpleNamespace( + model_type="muse_glimmer", + architectures=list(shim.architectures), + hidden_activation="silu", + hidden_size=int(kv("embedding_length")), + num_hidden_layers=num_layers, + num_attention_heads=int(kv("attention.head_count")), + num_key_value_heads=int(kv("attention.head_count_kv")), + head_dim=int(kv("attention.key_length")), + intermediate_size=int(kv("feed_forward_length")), + max_position_embeddings=int(kv("context_length")), + rms_norm_eps=float(kv("attention.layer_norm_rms_epsilon")), + # llama.cpp hard-codes the post-norm eps and NoPE on the full layers (src/models/muse-glimmer.cpp) + post_norm_eps=1e-8, + layer_types=["sliding_attention" if s else "full_attention" for s in swa], + layer_rope_theta=[swa_base if s else 0.0 for s in swa], + sliding_window=int(kv("attention.sliding_window")), + qk_scale_factor=_qk_scale_factor(shim.model_path, num_layers), + # llama.cpp's default when the key is absent; 0 turns the cap off + final_logit_softcapping=float(md.get(f"{_ARCH}.final_logit_softcapping", 30.0)) or None, + output_multiplier=float(kv("logit_scale")), + vocab_size=int(shim.vocab_size), + tie_word_embeddings=bool(shim.tie_word_embeddings), + ) + + +def parse_gguf_config(shim: GgufConfigShim) -> ModelConfig: + return parse_config(_hf_config(shim)) + + +def _unpermute(rows: torch.Tensor, heads: int) -> torch.Tensor: + """Undo the converter's rope permute of one q or k matrix (``_unpermute_for_rope``); it moves whole rows, so packed ggml rows work too.""" + return rows.reshape(heads, rows.shape[0] // heads // 2, 2, *rows.shape[1:]).swapaxes(1, 2).reshape(rows.shape) + + +def iter_gguf_weights( + model_path: str, + device: torch.device, + *, + include_moe_experts: bool, + include_non_moe: bool, +) -> Iterator[tuple[str, torch.Tensor]]: + """The text state dict from the GGUF.""" + if not include_non_moe: + return + from freetoken.utils import cached_load_hf_config + + hf = _hf_config(cached_load_hf_config(model_path)) + index = gguf_tensor_index(model_path) + hidden, heads, kv_heads = hf.hidden_size, hf.num_attention_heads, hf.num_key_value_heads + q_rows, kv_rows = heads * hf.head_dim, kv_heads * hf.head_dim + + def t(name: str) -> GgufTensorInfo: + return index[name if name in index else f"{name}.weight"] + + yield from fused("model.embed_tokens", [t("token_embd")], [matrix(t("token_embd"))], hidden) + if not hf.tie_word_embeddings: + yield from fused("lm_head", [t("output")], [matrix(t("output"))], hidden) + yield "model.norm.weight", float_tensor(t("output_norm")) + + for layer in range(hf.num_hidden_layers): + blk, pre = f"blk.{layer}", f"model.layers.{layer}" + for norm, tensor in _SHIFTED_NORMS.items(): + yield f"{pre}.{norm}.weight", float_tensor(t(f"{blk}.{tensor}")).float() - 1.0 + sa = f"{pre}.self_attn" + # convert_hf_to_gguf.py --fuse-qkv stores q, k and v as one attn_qkv matrix + if f"{blk}.attn_qkv.weight" in index: + qkv = [t(f"{blk}.attn_qkv")] * 3 + q, k, v = matrix(qkv[0]).split([q_rows, kv_rows, kv_rows]) + else: + qkv = [t(f"{blk}.attn_{p}") for p in ("q", "k", "v")] + q, k, v = (matrix(p) for p in qkv) + yield from fused(f"{sa}.qkv_proj", qkv, [_unpermute(q, heads), _unpermute(k, kv_heads), v], hidden) + gate = t(f"{blk}.attn_gate") + yield from fused(f"{sa}.gate_proj", [gate], [matrix(gate)], hidden) + out = t(f"{blk}.attn_output") + yield from fused(f"{sa}.o_proj", [out], [matrix(out)], int(out.ne[0])) + mlp = f"{pre}.mlp" + gate, up, down = (t(f"{blk}.ffn_{p}") for p in ("gate", "up", "down")) + yield from fused(f"{mlp}.gate_up_proj", [gate, up], [matrix(gate), matrix(up)], hidden) + yield from fused(f"{mlp}.down_proj", [down], [matrix(down)], int(down.ne[0])) + + +__all__ = ["iter_gguf_weights", "parse_gguf_config", "to_gguf"] diff --git a/python/freetoken/models/muse_glimmer/model.py b/python/freetoken/models/muse_glimmer/model.py index 50bd450670..d51c211743 100644 --- a/python/freetoken/models/muse_glimmer/model.py +++ b/python/freetoken/models/muse_glimmer/model.py @@ -105,6 +105,8 @@ def __init__(self, config: ModelConfig, *, prefix: str = "model"): self.embed_tokens = VocabParallelEmbedding( num_embeddings=config.vocab_size, embedding_dim=config.hidden_size, + quant_config=config.quant, + prefix=f"{prefix}.embed_tokens", ) # NormedEmbedding: a weightless RMSNorm on top of the embeddings. Kept separate # from the embedding matrix (the reference cannot fold it either -- the DFlash diff --git a/python/freetoken/models/muse_glimmer/weight.py b/python/freetoken/models/muse_glimmer/weight.py index 8b9bb21317..11a1817222 100644 --- a/python/freetoken/models/muse_glimmer/weight.py +++ b/python/freetoken/models/muse_glimmer/weight.py @@ -26,13 +26,10 @@ _VISION_QKV_PARTS = (".attn.q_proj", ".attn.k_proj", ".attn.v_proj") # Fused projections, concatenated on the output dim in this exact order to match the -# model's merged-linear splits. The attention gate rides the q/k/v fusion (it is computed -# from the same layer input); ``.self_attn.gate_proj`` and the SwiGLU ``.mlp.gate_proj`` -# are disambiguated by the full suffix. +# model's merged-linear splits. The full suffix keeps the standalone attention gate +# ``.self_attn.gate_proj`` out of the SwiGLU ``.mlp.gate_proj`` fusion. _FUSIONS: dict[str, tuple[str, ...]] = { - ".self_attn.qkvg_proj": ( - ".self_attn.q_proj", ".self_attn.k_proj", ".self_attn.v_proj", ".self_attn.gate_proj", - ), + ".self_attn.qkv_proj": (".self_attn.q_proj", ".self_attn.k_proj", ".self_attn.v_proj"), ".mlp.gate_up_proj": (".mlp.gate_proj", ".mlp.up_proj"), } @@ -134,8 +131,8 @@ def _iter_weights_compressed_tensors( Every text Linear (q/k/v/o, the attention gate, the MLP) is kept native NVFP4 (W4A16): ``.weight`` (uint8 packed) + ``.weight_scale`` (fp8 block) + ``.weight_global`` - (fp16 per-row, the reciprocal of the stored quant-side global). q/k/v/gate fuse into - ``qkvg_proj`` and gate/up into ``gate_up_proj`` on the output dim, each part keeping + (fp32 per-row, the reciprocal of the stored quant-side global). q/k/v fuse into + ``qkv_proj`` and gate/up into ``gate_up_proj`` on the output dim, each part keeping its own scales, so the fused FP4 weights are exact. Embeddings, norms, lm_head and the vision tower are bf16 (the checkpoint's ignore list). Scale lookups go through the shard-map reader: they can land in a different shard than their weight_packed.""" diff --git a/python/freetoken/models/nvfp4_banks.py b/python/freetoken/models/nvfp4_banks.py index 2dfdaaa35f..6d16c5ada2 100644 --- a/python/freetoken/models/nvfp4_banks.py +++ b/python/freetoken/models/nvfp4_banks.py @@ -4,7 +4,7 @@ import os import re from dataclasses import dataclass -from typing import Callable +from typing import Callable, Collection import safetensors import torch @@ -36,7 +36,29 @@ def _canon_kind(spec: "Nvfp4ExpertSourceSpec", kind: str) -> str: def _ingest_global(spec: "Nvfp4ExpertSourceSpec", tensor: torch.Tensor) -> torch.Tensor: if spec.global_reciprocal: tensor = 1.0 / tensor.float() - return tensor.to(torch.float16) + return tensor.to(torch.float32) + + +def _activation_scale(): + """The checkpoint tensor of the NVFP4 static activation scale under the installed dialect, or None.""" + from freetoken.layers.quantization import QuantKind, get_quant_config + + try: + quant = get_quant_config() + except RuntimeError: + return None + if quant is None or QuantKind.NVFP4 not in quant.STORAGE: + return None + return quant.stored_tensors(QuantKind.NVFP4).get("input_scale") + + +def _ingest(spec: "Nvfp4ExpertSourceSpec", act, role: str, tensor: torch.Tensor) -> torch.Tensor: + if role.endswith("_global"): + return _ingest_global(spec, tensor) + if role.endswith("_input"): + tensor = tensor.float() + return 1.0 / tensor if act.reciprocal else tensor + return tensor def _num_moe_layers(config) -> int: @@ -59,6 +81,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`` and ``_input`` are + per-tensor scalars and have 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", "_input")): + 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] @@ -68,6 +120,7 @@ def iter_nvfp4_expert_pieces( config, spec: Nvfp4ExpertSourceSpec, *, + layers: Collection[int] | None = None, parallel: bool = False, workers: int = 8, chunk: int = 8 << 20, @@ -76,10 +129,12 @@ def iter_nvfp4_expert_pieces( ): """One piece per routed expert: ``gate`` / ``up`` / ``down`` codes plus their ``_scale`` (fp8 block scales) and ``_global`` (the per-tensor scale, reciprocal for quant-side dialects, - fp16) companions, straight from the safetensors shards. + fp32) companions, straight from the safetensors shards. When every expert also stores the + dialect's static activation scale, the pieces carry it as fp32 ``_input`` (dequant side). Serial reads walk the shards in order; ``parallel`` uses the chunked O_DIRECT reader. Either way tensors of one expert may span shards, so they are grouped by (layer, expert) as they land. + ``layers`` (bank layer ids, default all) reads only those layers' experts. """ from freetoken.models.loader import drop_page_cache as _drop from freetoken.models.loader import safetensors_weight_map @@ -89,13 +144,15 @@ def iter_nvfp4_expert_pieces( folder = download_hf_weight(model_path) weight_map = safetensors_weight_map(folder) + keep = None if layers is None else set(layers) + act = _activation_scale() wanted: dict[str, tuple[int, int, str]] = {} for name in weight_map: match = spec.key_pattern.match(name) if match is None: continue bank_layer = _bank_layer(spec, int(match.group("layer")), config) - if bank_layer is None: + if bank_layer is None or (keep is not None and bank_layer not in keep): continue proj = match.group("proj") if proj not in spec.proj_to_role: @@ -104,9 +161,20 @@ def iter_nvfp4_expert_pieces( if kind not in ("weight", "weight_scale", "weight_scale_2"): raise ValueError(f"{spec.desc}: unknown NVFP4 expert tensor kind {kind!r}") wanted[name] = (bank_layer, int(match.group("expert")), spec.proj_to_role[proj] + _kind_suffix(kind)) - expected = _num_moe_layers(config) * config.num_experts * 9 - if len(wanted) != expected: - raise ValueError(f"{spec.desc}: found {len(wanted)} expert tensors, expected {expected}") + if kind == "weight" and act is not None: + act_name = name[: match.start("kind")] + act.name + if act_name in weight_map: + wanted[act_name] = (bank_layer, int(match.group("expert")), spec.proj_to_role[proj] + "_input") + experts = (_num_moe_layers(config) if keep is None else len(keep)) * config.num_experts + acts = [name for name, (_, _, role) in wanted.items() if role.endswith("_input")] + # a partial set serves no W4A4 kernel; the pieces then carry none + if len(acts) != experts * 3: + for name in acts: + del wanted[name] + acts = [] + per_expert = 12 if acts else 9 + if len(wanted) != experts * per_expert: + raise ValueError(f"{spec.desc}: found {len(wanted)} expert tensors, expected {experts * per_expert}") def _serial(): by_shard: dict[str, list[str]] = collections.defaultdict(list) @@ -118,21 +186,28 @@ def _serial(): drop(path) with safetensors.safe_open(path, framework="pt", device="cpu") as f: for name in by_shard[shard]: - tensor = f.get_tensor(name) - if wanted[name][2].endswith("_global"): - tensor = _ingest_global(spec, tensor) - yield name, tensor + yield name, _ingest(spec, act, wanted[name][2], f.get_tensor(name)) drop(path) def _parallel(): from freetoken.models.weight import iter_expert_tensors_parallel for name, tensor in iter_expert_tensors_parallel(folder, lambda n: n in wanted, workers=workers, chunk=chunk): - if wanted[name][2].endswith("_global"): - tensor = _ingest_global(spec, tensor) - yield name, tensor + yield name, _ingest(spec, act, wanted[name][2], tensor) + + 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) - return per_expert_pieces(_parallel() if parallel else _serial(), wanted.get, tensors_per_expert=9) + stream = _sharded() + return per_expert_pieces(stream, wanted.get, tensors_per_expert=per_expert) __all__ = ["Nvfp4ExpertSourceSpec", "iter_nvfp4_expert_pieces"] diff --git a/python/freetoken/models/qwen2/__init__.py b/python/freetoken/models/qwen2/__init__.py index 43ca84f82a..3917c188fd 100644 --- a/python/freetoken/models/qwen2/__init__.py +++ b/python/freetoken/models/qwen2/__init__.py @@ -1,5 +1,6 @@ from .config import parse_config +from .gguf import iter_gguf_weights, parse_gguf_config, to_gguf from .model import Qwen2ForCausalLM from .weight import iter_weights -__all__ = ["Qwen2ForCausalLM", "parse_config", "iter_weights"] +__all__ = ["Qwen2ForCausalLM", "parse_config", "iter_weights", "iter_gguf_weights", "parse_gguf_config", "to_gguf"] diff --git a/python/freetoken/models/qwen2/gguf.py b/python/freetoken/models/qwen2/gguf.py new file mode 100644 index 0000000000..501c113a3a --- /dev/null +++ b/python/freetoken/models/qwen2/gguf.py @@ -0,0 +1,144 @@ +"""Qwen2 / Qwen2.5 from a llama.cpp ``qwen2`` GGUF: config from the GGUF metadata, tensors by name. + +The llama.cpp converter (conversion/qwen.py) stores these tensors unchanged: both sides run NeoX rope, so q/k carry no rope permute, and the RMSNorm weights are the HF ones. +""" + +from __future__ import annotations + +import re +from types import SimpleNamespace +from typing import TYPE_CHECKING, Iterator + +import torch + +from freetoken.models.gguf.reader import GgufTensorInfo, gguf_tensor_index +from freetoken.models.gguf.tensors import float_tensor, fused, matrix + +from .config import parse_config + +if TYPE_CHECKING: + from freetoken.models.config import ModelConfig + from freetoken.models.gguf.config import GgufConfigShim + +_ARCH = "qwen2" +_LAYER_LEAVES = { + "self_attn.q_proj": "attn_q", + "self_attn.k_proj": "attn_k", + "self_attn.v_proj": "attn_v", + "self_attn.o_proj": "attn_output", + "mlp.gate_proj": "ffn_gate", + "mlp.up_proj": "ffn_up", + "mlp.down_proj": "ffn_down", +} +_QKV_LEAVES = ("self_attn.q_proj", "self_attn.k_proj", "self_attn.v_proj") +_TOP_LEAVES = {"model.embed_tokens": "token_embd", "lm_head": "output"} +_LAYER_RE = re.compile(r"^model\.layers\.(\d+)\.(.+)$") + + +def to_gguf(name: str, tensors: dict[str, int]) -> tuple[str, ...]: + """Checkpoint module name -> the GGUF tensors (without ``.weight``) that store it; empty for modules the GGUF has no matrix for.""" + if name in _TOP_LEAVES: + return (_TOP_LEAVES[name],) + m = _LAYER_RE.match(name) + if m is None: + return () + blk, leaf = f"blk.{m.group(1)}", m.group(2) + if leaf in _QKV_LEAVES and f"{blk}.attn_qkv.weight" in tensors: + return (f"{blk}.attn_qkv",) + tensor = _LAYER_LEAVES.get(leaf) + return (f"{blk}.{tensor}",) if tensor is not None else () + + +def _check_supported(shim: GgufConfigShim) -> None: + if any(k.startswith("prism.hadamard.") for k in shim.metadata): + raise NotImplementedError(f"{shim.model_path}: Hadamard-rotated (prism) GGUF weights are not supported") + scaled = sorted(n for n in gguf_tensor_index(shim.model_path) if n.endswith((".scale", ".input_scale"))) + if scaled: + raise NotImplementedError(f"{shim.model_path}: GGUF tensors with per-tensor scales are not supported (e.g. {scaled[0]})") + + +def _hf_config(shim: GgufConfigShim) -> SimpleNamespace: + """The HF config fields ``parse_config`` reads, from the GGUF metadata.""" + md = shim.metadata + + def kv(key: str): + return md[f"{_ARCH}.{key}"] + + _check_supported(shim) + rope = {"rope_theta": float(kv("rope.freq_base")), "rope_type": "default"} + scaling = md.get(f"{_ARCH}.rope.scaling.type", "none") + if scaling == "yarn": + rope.update( + rope_type="yarn", + factor=float(kv("rope.scaling.factor")), + original_max_position_embeddings=int(kv("rope.scaling.original_context_length")), + ) + elif scaling != "none": + raise NotImplementedError(f"{shim.model_path}: rope scaling {scaling!r} is not supported") + # Qwen2.5 files carry no key_length; parse_config then derives head_dim from the hidden size + key_length = md.get(f"{_ARCH}.attention.key_length") + return SimpleNamespace( + model_type="qwen2", + architectures=list(shim.architectures), + hidden_act="silu", + hidden_size=int(kv("embedding_length")), + num_hidden_layers=int(kv("block_count")), + num_attention_heads=int(kv("attention.head_count")), + num_key_value_heads=int(kv("attention.head_count_kv")), + head_dim=int(key_length) if key_length is not None else None, + rope_scaling=rope, + max_position_embeddings=int(kv("context_length")), + rms_norm_eps=float(kv("attention.layer_norm_rms_epsilon")), + vocab_size=int(shim.vocab_size), + tie_word_embeddings=bool(shim.tie_word_embeddings), + intermediate_size=int(kv("feed_forward_length")), + ) + + +def parse_gguf_config(shim: GgufConfigShim) -> ModelConfig: + return parse_config(_hf_config(shim)) + + +def iter_gguf_weights( + model_path: str, + device: torch.device, + *, + include_moe_experts: bool, + include_non_moe: bool, +) -> Iterator[tuple[str, torch.Tensor]]: + """The state dict from the GGUF.""" + if not include_non_moe: + return + from freetoken.utils import cached_load_hf_config + + hf = _hf_config(cached_load_hf_config(model_path)) + index = gguf_tensor_index(model_path) + hidden = hf.hidden_size + + def t(name: str) -> GgufTensorInfo: + return index[name if name in index else f"{name}.weight"] + + yield from fused("model.embed_tokens", [t("token_embd")], [matrix(t("token_embd"))], hidden) + if not hf.tie_word_embeddings: + yield from fused("lm_head", [t("output")], [matrix(t("output"))], hidden) + yield "model.norm.weight", float_tensor(t("output_norm")) + + for layer in range(hf.num_hidden_layers): + blk, pre = f"blk.{layer}", f"model.layers.{layer}" + yield f"{pre}.input_layernorm.weight", float_tensor(t(f"{blk}.attn_norm")) + yield f"{pre}.post_attention_layernorm.weight", float_tensor(t(f"{blk}.ffn_norm")) + sa = f"{pre}.self_attn" + # convert_hf_to_gguf.py --fuse-qkv stores q, k and v as one attn_qkv matrix + qkv = [t(f"{blk}.attn_qkv")] if f"{blk}.attn_qkv.weight" in index else [t(f"{blk}.attn_{p}") for p in ("q", "k", "v")] + yield from fused(f"{sa}.qkv_proj", qkv, [matrix(p) for p in qkv], hidden) + bias = index.get(f"{blk}.attn_qkv.bias") + yield f"{sa}.qkv_proj.bias", float_tensor(bias) if bias is not None else torch.cat([float_tensor(t(f"{blk}.attn_{p}.bias")) for p in ("q", "k", "v")]) + out = t(f"{blk}.attn_output") + yield from fused(f"{sa}.o_proj", [out], [matrix(out)], int(out.ne[0])) + mlp = f"{pre}.mlp" + gate, up, down = (t(f"{blk}.ffn_{p}") for p in ("gate", "up", "down")) + yield from fused(f"{mlp}.gate_up_proj", [gate, up], [matrix(gate), matrix(up)], hidden) + yield from fused(f"{mlp}.down_proj", [down], [matrix(down)], int(down.ne[0])) + + +__all__ = ["iter_gguf_weights", "parse_gguf_config", "to_gguf"] diff --git a/python/freetoken/models/qwen2/model.py b/python/freetoken/models/qwen2/model.py index ed56fde25a..94bc2beb6b 100644 --- a/python/freetoken/models/qwen2/model.py +++ b/python/freetoken/models/qwen2/model.py @@ -46,6 +46,8 @@ def __init__(self, config: ModelConfig, *, prefix: str = "model"): self.embed_tokens = VocabParallelEmbedding( num_embeddings=config.vocab_size, embedding_dim=config.hidden_size, + quant_config=config.quant, + prefix=f"{prefix}.embed_tokens", ) self.layers = OPList( [ diff --git a/python/freetoken/models/qwen3/__init__.py b/python/freetoken/models/qwen3/__init__.py index 71bda9781f..fef7c8fb99 100644 --- a/python/freetoken/models/qwen3/__init__.py +++ b/python/freetoken/models/qwen3/__init__.py @@ -1,5 +1,6 @@ from .config import parse_config +from .gguf import iter_gguf_weights, parse_gguf_config, to_gguf from .model import Qwen3ForCausalLM from .weight import iter_weights -__all__ = ["Qwen3ForCausalLM", "parse_config", "iter_weights"] +__all__ = ["Qwen3ForCausalLM", "parse_config", "iter_weights", "iter_gguf_weights", "parse_gguf_config", "to_gguf"] diff --git a/python/freetoken/models/qwen3/gguf.py b/python/freetoken/models/qwen3/gguf.py new file mode 100644 index 0000000000..8e0388fe3c --- /dev/null +++ b/python/freetoken/models/qwen3/gguf.py @@ -0,0 +1,141 @@ +"""Qwen3 dense from a llama.cpp ``qwen3`` GGUF: config from the GGUF metadata, tensors by name. + +The llama.cpp converter (conversion/qwen.py) stores these tensors unchanged: both sides run NeoX rope, so q/k carry no rope permute, and the RMSNorm weights are the HF ones. +""" + +from __future__ import annotations + +import re +from types import SimpleNamespace +from typing import TYPE_CHECKING, Iterator + +import torch + +from freetoken.models.gguf.reader import GgufTensorInfo, gguf_tensor_index +from freetoken.models.gguf.tensors import check_scales, float_tensor, fused, matrix + +from .config import parse_config + +if TYPE_CHECKING: + from freetoken.models.config import ModelConfig + from freetoken.models.gguf.config import GgufConfigShim + +_ARCH = "qwen3" +_LAYER_LEAVES = { + "self_attn.q_proj": "attn_q", + "self_attn.k_proj": "attn_k", + "self_attn.v_proj": "attn_v", + "self_attn.o_proj": "attn_output", + "mlp.gate_proj": "ffn_gate", + "mlp.up_proj": "ffn_up", + "mlp.down_proj": "ffn_down", +} +_QKV_LEAVES = ("self_attn.q_proj", "self_attn.k_proj", "self_attn.v_proj") +_TOP_LEAVES = {"model.embed_tokens": "token_embd", "lm_head": "output"} +_LAYER_RE = re.compile(r"^model\.layers\.(\d+)\.(.+)$") + + +def to_gguf(name: str, tensors: dict[str, int]) -> tuple[str, ...]: + """Checkpoint module name -> the GGUF tensors (without ``.weight``) that store it; empty for modules the GGUF has no matrix for.""" + if name in _TOP_LEAVES: + return (_TOP_LEAVES[name],) + m = _LAYER_RE.match(name) + if m is None: + return () + blk, leaf = f"blk.{m.group(1)}", m.group(2) + if leaf in _QKV_LEAVES and f"{blk}.attn_qkv.weight" in tensors: + return (f"{blk}.attn_qkv",) + tensor = _LAYER_LEAVES.get(leaf) + return (f"{blk}.{tensor}",) if tensor is not None else () + + +def _check_supported(shim: GgufConfigShim) -> None: + if any(k.startswith("prism.hadamard.") for k in shim.metadata): + raise NotImplementedError(f"{shim.model_path}: Hadamard-rotated (prism) GGUF weights are not supported") + check_scales(shim.model_path) + + +def _hf_config(shim: GgufConfigShim) -> SimpleNamespace: + """The HF config fields ``parse_config`` reads, from the GGUF metadata.""" + md = shim.metadata + + def kv(key: str): + return md[f"{_ARCH}.{key}"] + + _check_supported(shim) + rope = {"rope_theta": float(kv("rope.freq_base")), "rope_type": "default"} + scaling = md.get(f"{_ARCH}.rope.scaling.type", "none") + if scaling == "yarn": + rope.update( + rope_type="yarn", + factor=float(kv("rope.scaling.factor")), + original_max_position_embeddings=int(kv("rope.scaling.original_context_length")), + ) + elif scaling != "none": + raise NotImplementedError(f"{shim.model_path}: rope scaling {scaling!r} is not supported") + key_length = md.get(f"{_ARCH}.attention.key_length") + return SimpleNamespace( + model_type="qwen3", + architectures=list(shim.architectures), + hidden_act="silu", + hidden_size=int(kv("embedding_length")), + num_hidden_layers=int(kv("block_count")), + num_attention_heads=int(kv("attention.head_count")), + num_key_value_heads=int(kv("attention.head_count_kv")), + head_dim=int(key_length) if key_length is not None else None, + rope_scaling=rope, + max_position_embeddings=int(kv("context_length")), + rms_norm_eps=float(kv("attention.layer_norm_rms_epsilon")), + vocab_size=int(shim.vocab_size), + tie_word_embeddings=bool(shim.tie_word_embeddings), + intermediate_size=int(kv("feed_forward_length")), + ) + + +def parse_gguf_config(shim: GgufConfigShim) -> ModelConfig: + return parse_config(_hf_config(shim)) + + +def iter_gguf_weights( + model_path: str, + device: torch.device, + *, + include_moe_experts: bool, + include_non_moe: bool, +) -> Iterator[tuple[str, torch.Tensor]]: + """The state dict from the GGUF.""" + if not include_non_moe: + return + from freetoken.utils import cached_load_hf_config + + hf = _hf_config(cached_load_hf_config(model_path)) + index = gguf_tensor_index(model_path) + hidden = hf.hidden_size + + def t(name: str) -> GgufTensorInfo: + return index[name if name in index else f"{name}.weight"] + + yield from fused("model.embed_tokens", [t("token_embd")], [matrix(t("token_embd"))], hidden) + if not hf.tie_word_embeddings: + yield from fused("lm_head", [t("output")], [matrix(t("output"))], hidden) + yield "model.norm.weight", float_tensor(t("output_norm")) + + for layer in range(hf.num_hidden_layers): + blk, pre = f"blk.{layer}", f"model.layers.{layer}" + yield f"{pre}.input_layernorm.weight", float_tensor(t(f"{blk}.attn_norm")) + yield f"{pre}.post_attention_layernorm.weight", float_tensor(t(f"{blk}.ffn_norm")) + sa = f"{pre}.self_attn" + # convert_hf_to_gguf.py --fuse-qkv stores q, k and v as one attn_qkv matrix + qkv = [t(f"{blk}.attn_qkv")] if f"{blk}.attn_qkv.weight" in index else [t(f"{blk}.attn_{p}") for p in ("q", "k", "v")] + yield from fused(f"{sa}.qkv_proj", qkv, [matrix(p) for p in qkv], hidden) + yield f"{sa}.q_norm.weight", float_tensor(t(f"{blk}.attn_q_norm")) + yield f"{sa}.k_norm.weight", float_tensor(t(f"{blk}.attn_k_norm")) + out = t(f"{blk}.attn_output") + yield from fused(f"{sa}.o_proj", [out], [matrix(out)], int(out.ne[0])) + mlp = f"{pre}.mlp" + gate, up, down = (t(f"{blk}.ffn_{p}") for p in ("gate", "up", "down")) + yield from fused(f"{mlp}.gate_up_proj", [gate, up], [matrix(gate), matrix(up)], hidden) + yield from fused(f"{mlp}.down_proj", [down], [matrix(down)], int(down.ne[0])) + + +__all__ = ["iter_gguf_weights", "parse_gguf_config", "to_gguf"] diff --git a/python/freetoken/models/qwen3/model.py b/python/freetoken/models/qwen3/model.py index 6ac7d2fa8f..ae0d4394f5 100644 --- a/python/freetoken/models/qwen3/model.py +++ b/python/freetoken/models/qwen3/model.py @@ -46,6 +46,8 @@ def __init__(self, config: ModelConfig, *, prefix: str = "model"): self.embed_tokens = VocabParallelEmbedding( num_embeddings=config.vocab_size, embedding_dim=config.hidden_size, + quant_config=config.quant, + prefix=f"{prefix}.embed_tokens", ) self.layers = OPList( [ diff --git a/python/freetoken/models/qwen3_5_moe/__init__.py b/python/freetoken/models/qwen3_5_moe/__init__.py index 018b1c1e97..abe8fd1c2f 100644 --- a/python/freetoken/models/qwen3_5_moe/__init__.py +++ b/python/freetoken/models/qwen3_5_moe/__init__.py @@ -1,4 +1,5 @@ from .config import parse_config +from .gguf import iter_gguf_weights, parse_gguf_config, to_gguf from .model import ( Qwen3_5ForCausalLM, Qwen3_5ForConditionalGeneration, @@ -18,4 +19,7 @@ "iter_weights_parallel", "iter_expert_pieces", "nvfp4_expert_spec", + "iter_gguf_weights", + "parse_gguf_config", + "to_gguf", ] diff --git a/python/freetoken/models/qwen3_5_moe/attention.py b/python/freetoken/models/qwen3_5_moe/attention.py index 7a9e6a333a..0a736a1042 100644 --- a/python/freetoken/models/qwen3_5_moe/attention.py +++ b/python/freetoken/models/qwen3_5_moe/attention.py @@ -4,7 +4,7 @@ import torch from freetoken.core import get_global_ctx -from freetoken.layers import BaseOP, GemmaRMSNorm, LinearColParallelMerged, LinearReplicated +from freetoken.layers import BaseOP, GemmaPlusOneRMSNorm, LinearColParallelMerged, LinearReplicated from freetoken.layers.rotary import get_rope from freetoken.utils import nvtx_annotate @@ -41,10 +41,8 @@ def __init__(self, config: ModelConfig, layer_id: int, *, prefix: str = ""): config.hidden_size, self._qkv_split, has_bias=False, quant_config=config.quant, prefix=f"{prefix}.qkv_proj", ) - # Qwen3.5 uses Gemma-style (1+weight) RMSNorm; the weight loader bakes the +1 - # into the stored weight (GemmaRMSNorm scales by the raw weight). - self.q_norm = GemmaRMSNorm(head_dim, eps=config.rms_norm_eps) - self.k_norm = GemmaRMSNorm(head_dim, eps=config.rms_norm_eps) + self.q_norm = GemmaPlusOneRMSNorm(head_dim, eps=config.rms_norm_eps) + self.k_norm = GemmaPlusOneRMSNorm(head_dim, eps=config.rms_norm_eps) self.rotary = get_rope( head_dim=head_dim, rotary_dim=config.rotary_config.rotary_dim, diff --git a/python/freetoken/models/qwen3_5_moe/gdn.py b/python/freetoken/models/qwen3_5_moe/gdn.py index d5fe2aa8fe..ae584b7d1b 100644 --- a/python/freetoken/models/qwen3_5_moe/gdn.py +++ b/python/freetoken/models/qwen3_5_moe/gdn.py @@ -30,9 +30,10 @@ class Qwen3_5GatedDeltaNet(BaseOP): 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, *, quant_config: QuantConfig | None = None, - prefix: str = "", + prefix: str = "", out_heads_tiled: bool = False, ): self.layer_id = layer_id + self._out_heads_tiled = out_heads_tiled # The fla chunk/decode kernels read+write the recurrent state and the per-chunk h as # [V, K] while the LinearStatePool declares it [K, V]; these coincide (and the # hybrid-radix snapshot scatter h[h_row]->slot is a plain copy) only when the two head @@ -90,14 +91,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 +170,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) @@ -195,6 +200,10 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: core_out = core_out.reshape(-1, self.head_v_dim) z = z.reshape(-1, self.head_v_dim) out = self.norm.forward(core_out, z).reshape(total, -1) + if self._out_heads_tiled: + # grouped value heads (k-head major) -> tiled (k-head minor), the column order of a GGUF ssm_out + r = self.num_v_heads // self.num_k_heads + out = out.view(total, self.num_k_heads, r, self.head_v_dim).transpose(1, 2).reshape(total, -1) return self.out_proj.forward(out) diff --git a/python/freetoken/models/qwen3_5_moe/gguf.py b/python/freetoken/models/qwen3_5_moe/gguf.py new file mode 100644 index 0000000000..caf1268008 --- /dev/null +++ b/python/freetoken/models/qwen3_5_moe/gguf.py @@ -0,0 +1,221 @@ +"""Qwen3.5 / 3.6 / 3.8 dense and MoE from a llama.cpp ``qwen35`` / ``qwen35moe`` GGUF: config from the GGUF metadata, tensor names, and the converter's transforms undone. + +The llama.cpp converter (conversion/qwen.py) stores ``A_log`` as ``-exp(A_log)``, the squeezed conv kernels, and the GDN value heads in tiled order (k-head minor). Its zero-centred norms are stored as ``w + 1``; the loader subtracts the 1 in fp32 and the norms add it back in fp32 at runtime. Row reorders are undone on the packed bytes; ``ssm_out`` keeps its tiled columns and the GDN layer permutes its input instead. +""" + +from __future__ import annotations + +import re +from dataclasses import replace +from types import SimpleNamespace +from typing import TYPE_CHECKING, Iterator + +import torch + +from freetoken.layers.quantization.ggml_types import is_quantized +from freetoken.models.gguf.reader import GgufTensorInfo, gguf_tensor_index +from freetoken.models.gguf.tensors import check_scales, float_tensor, fused, matrix + +from .config import parse_config + +if TYPE_CHECKING: + from freetoken.models.config import ModelConfig + from freetoken.models.gguf.config import GgufConfigShim + +_LAYER_LEAVES: dict[str, tuple[str, ...]] = { + "linear_attn.in_proj_qkv": ("attn_qkv",), + "linear_attn.in_proj_z": ("attn_gate",), + "linear_attn.in_proj_b": ("ssm_beta",), + "linear_attn.in_proj_a": ("ssm_alpha",), + "linear_attn.out_proj": ("ssm_out",), + "self_attn.q_proj": ("attn_q",), + "self_attn.k_proj": ("attn_k",), + "self_attn.v_proj": ("attn_v",), + "self_attn.o_proj": ("attn_output",), + "mlp.gate_proj": ("ffn_gate",), + "mlp.up_proj": ("ffn_up",), + "mlp.down_proj": ("ffn_down",), + "mlp.gate": ("ffn_gate_inp",), + "mlp.shared_expert_gate": ("ffn_gate_inp_shexp",), + "mlp.shared_expert.gate_proj": ("ffn_gate_shexp",), + "mlp.shared_expert.up_proj": ("ffn_up_shexp",), + "mlp.shared_expert.down_proj": ("ffn_down_shexp",), + "mlp.experts.0.gate_proj": ("ffn_gate_exps",), + "mlp.experts.0.up_proj": ("ffn_up_exps",), + "mlp.experts.0.down_proj": ("ffn_down_exps",), +} +_TOP_LEAVES: dict[str, tuple[str, ...]] = { + "model.language_model.embed_tokens": ("token_embd",), + "lm_head": ("output",), +} +_LAYER_RE = re.compile(r"^model\.language_model\.layers\.(\d+)\.(.+)$") + + +def to_gguf(name: str, tensors: dict[str, int]) -> tuple[str, ...]: + """Checkpoint module name -> the GGUF tensors (without ``.weight``) that store it; empty for modules the GGUF has no matrix for.""" + if name in _TOP_LEAVES: + return _TOP_LEAVES[name] + m = _LAYER_RE.match(name) + if m is None: + return () + blk, leaf = f"blk.{m.group(1)}", m.group(2) + if leaf in ("mlp.experts.0.gate_proj", "mlp.experts.0.up_proj") and f"{blk}.ffn_gate_up_exps.weight" in tensors: + return (f"{blk}.ffn_gate_up_exps",) + return tuple(f"{blk}.{t}" for t in _LAYER_LEAVES.get(leaf, ())) + + +def _check_supported(shim: GgufConfigShim) -> None: + if any(k.startswith("prism.hadamard.") for k in shim.metadata): + raise NotImplementedError(f"{shim.model_path}: Hadamard-rotated (prism) GGUF weights are not supported") + check_scales(shim.model_path) + + +def parse_gguf_config(shim: GgufConfigShim) -> ModelConfig: + md = shim.metadata + arch = shim.model_type + + def kv(key: str): + return md[f"{arch}.{key}"] + + _check_supported(shim) + head_dim = int(kv("attention.key_length")) + n_v = int(kv("ssm.time_step_rank")) + # block_count also counts the MTP layers a GGUF may carry after the decoder + num_layers = int(kv("block_count")) - int(md.get(f"{arch}.nextn_predict_layers", 0)) + recurrent = md.get(f"{arch}.attention.recurrent_layers") + if recurrent is not None: + layer_types = ["linear_attention" if r else "full_attention" for r in list(recurrent)[:num_layers]] + else: + interval = int(kv("full_attention_interval")) + layer_types = ["full_attention" if (i + 1) % interval == 0 else "linear_attention" for i in range(num_layers)] + moe = arch == "qwen35moe" + text = SimpleNamespace( + hidden_act="silu", + hidden_size=int(kv("embedding_length")), + num_hidden_layers=num_layers, + layer_types=layer_types, + num_attention_heads=int(kv("attention.head_count")), + num_key_value_heads=int(kv("attention.head_count_kv")), + head_dim=head_dim, + rope_parameters={ + "rope_type": "default", + "rope_theta": float(kv("rope.freq_base")), + "partial_rotary_factor": int(kv("rope.dimension_count")) / head_dim, + "mrope_section": [int(s) for s in kv("rope.dimension_sections")][:3], + "mrope_interleaved": True, + }, + max_position_embeddings=int(kv("context_length")), + rms_norm_eps=float(kv("attention.layer_norm_rms_epsilon")), + vocab_size=int(shim.vocab_size), + tie_word_embeddings=bool(shim.tie_word_embeddings), + intermediate_size=0 if moe else int(kv("feed_forward_length")), + num_experts=int(kv("expert_count")) if moe else 0, + num_experts_per_tok=int(kv("expert_used_count")) if moe else 0, + moe_intermediate_size=int(kv("expert_feed_forward_length")) if moe else 0, + shared_expert_intermediate_size=int(kv("expert_shared_feed_forward_length")) if moe else 0, + linear_num_key_heads=int(kv("ssm.group_count")), + linear_num_value_heads=n_v, + linear_key_head_dim=int(kv("ssm.state_size")), + linear_value_head_dim=int(kv("ssm.inner_size")) // n_v, + linear_conv_kernel_dim=int(kv("ssm.conv_kernel")), + ) + hf_like = SimpleNamespace( + text_config=text, + model_type="qwen3_5_moe" if moe else "qwen3_5", + architectures=list(shim.architectures), + vision_config=None, + ) + config = parse_config(hf_like) + return replace(config, expert_quant="gguf") if moe else config + + +# ---------------------------------------------------------------------------------------------- +# Weights +# ---------------------------------------------------------------------------------------------- + +def _untile(t: torch.Tensor, num_k_heads: int, head_dim: int, dim: int = 0) -> torch.Tensor: + """Tiled value heads (k-head minor) back to grouped (k-head major) along ``dim``.""" + n = t.shape[dim] + r = n // (num_k_heads * head_dim) + idx = torch.arange(n).view(r, num_k_heads, head_dim).transpose(0, 1).reshape(n) + return t.index_select(dim, idx) + + +def iter_gguf_weights( + model_path: str, + device: torch.device, + *, + include_moe_experts: bool, + include_non_moe: bool, + include_vision: bool = True, +) -> Iterator[tuple[str, torch.Tensor]]: + """The dense state dict from the GGUF; the routed experts load as offload banks (``freetoken.models.gguf.experts``).""" + if not include_non_moe: + return + from freetoken.utils import cached_load_hf_config + + config = parse_gguf_config(cached_load_hf_config(model_path)) + index = gguf_tensor_index(model_path) + g = config.linear_attention_group() + hidden = config.hidden_size + + def t(name: str) -> GgufTensorInfo: + return index[name if name in index else f"{name}.weight"] + + def norm(name: str) -> torch.Tensor: + return float_tensor(t(name)).float() - 1.0 + + yield from fused("model.embed_tokens", [t("token_embd")], [matrix(t("token_embd"))], hidden) + if not config.tie_word_embeddings: + yield from fused("lm_head", [t("output")], [matrix(t("output"))], hidden) + yield "model.norm.weight", norm("output_norm") + + for layer in range(config.num_layers): + blk, pre = f"blk.{layer}", f"model.layers.{layer}" + yield f"{pre}.input_layernorm.weight", norm(f"{blk}.attn_norm") + yield f"{pre}.post_attention_layernorm.weight", norm(f"{blk}.post_attention_norm") + if config.is_linear_layer(layer): + la = f"{pre}.linear_attn" + nk, dk = g.num_key_heads, g.key_head_dim + qkv, gate = t(f"{blk}.attn_qkv"), t(f"{blk}.attn_gate") + qkv_rows = matrix(qkv) + v_start = 2 * nk * dk + qkv_rows = torch.cat([qkv_rows[:v_start], _untile(qkv_rows[v_start:], nk, g.value_head_dim)]) + z_rows = _untile(matrix(gate), nk, g.value_head_dim) + beta_info, alpha_info = t(f"{blk}.ssm_beta"), t(f"{blk}.ssm_alpha") + beta, alpha = (_untile(matrix(i), nk, 1) for i in (beta_info, alpha_info)) + if is_quantized(qkv.ggml_type) or is_quantized(gate.ggml_type): + yield from fused(f"{la}.in_proj_qkvz", [qkv, gate], [qkv_rows, z_rows], hidden) + yield from fused(f"{la}.in_proj_ba", [beta_info, alpha_info], [beta, alpha], hidden) + else: + yield from fused(f"{la}.in_proj", [qkv, gate, beta_info, alpha_info], [qkv_rows, z_rows, beta, alpha], hidden) + conv = float_tensor(t(f"{blk}.ssm_conv1d")).float() + conv = torch.cat([conv[:v_start], _untile(conv[v_start:], nk, g.value_head_dim)]) + yield f"{la}.conv1d.weight", conv.unsqueeze(1) + yield f"{la}.dt_bias", _untile(float_tensor(t(f"{blk}.ssm_dt.bias")).float(), nk, 1) + yield f"{la}.A_log", _untile(torch.log(-float_tensor(t(f"{blk}.ssm_a")).float()), nk, 1) + yield f"{la}.norm.weight", float_tensor(t(f"{blk}.ssm_norm")) + yield from fused(f"{la}.out_proj", [t(f"{blk}.ssm_out")], [matrix(t(f"{blk}.ssm_out"))], g.num_value_heads * g.value_head_dim) + else: + sa = f"{pre}.self_attn" + parts = [t(f"{blk}.attn_{p}") for p in ("q", "k", "v")] + yield from fused(f"{sa}.qkv_proj", parts, [matrix(p) for p in parts], hidden) + out = t(f"{blk}.attn_output") + yield from fused(f"{sa}.o_proj", [out], [matrix(out)], int(out.ne[0])) + yield f"{sa}.q_norm.weight", norm(f"{blk}.attn_q_norm") + yield f"{sa}.k_norm.weight", norm(f"{blk}.attn_k_norm") + + mlp = f"{pre}.mlp" + if config.moe_enabled: + yield f"{mlp}.gate.weight", float_tensor(t(f"{blk}.ffn_gate_inp")) + sg, su, sd = (t(f"{blk}.ffn_{p}_shexp") for p in ("gate", "up", "down")) + yield from fused(f"{mlp}.shared_expert.gate_up_proj", [sg, su], [matrix(sg), matrix(su)], hidden) + yield from fused(f"{mlp}.shared_expert.down_proj", [sd], [matrix(sd)], config.shared_expert_intermediate_size) + yield f"{mlp}.shared_expert_gate.weight", float_tensor(t(f"{blk}.ffn_gate_inp_shexp")).reshape(1, -1) + else: + fg, fu, fd = (t(f"{blk}.ffn_{p}") for p in ("gate", "up", "down")) + yield from fused(f"{mlp}.gate_up_proj", [fg, fu], [matrix(fg), matrix(fu)], hidden) + yield from fused(f"{mlp}.down_proj", [fd], [matrix(fd)], config.intermediate_size) + + +__all__ = ["iter_gguf_weights", "parse_gguf_config", "to_gguf"] diff --git a/python/freetoken/models/qwen3_5_moe/model.py b/python/freetoken/models/qwen3_5_moe/model.py index 8cdd41ae35..88c6ae1345 100644 --- a/python/freetoken/models/qwen3_5_moe/model.py +++ b/python/freetoken/models/qwen3_5_moe/model.py @@ -6,7 +6,7 @@ from freetoken.core import get_global_ctx from freetoken.layers import ( BaseOP, - GemmaRMSNorm, + GemmaPlusOneRMSNormFused, OPList, ParallelLMHead, VocabParallelEmbedding, @@ -46,6 +46,8 @@ def __init__(self, config: ModelConfig, layer_id: int, *, prefix: str = ""): layer_id=layer_id, quant_config=config.quant, prefix=f"{prefix}.linear_attn", + # llama.cpp stores every qwen35 ssm_out with tiled value-head columns + out_heads_tiled=config.quant is not None and config.quant.dialect == "gguf", ) else: self.self_attn = Qwen3_5Attention(config, layer_id, prefix=f"{prefix}.self_attn") @@ -56,20 +58,15 @@ def __init__(self, config: ModelConfig, layer_id: int, *, prefix: str = ""): if config.moe_enabled else Qwen3_5DenseMLP(config, prefix=f"{prefix}.mlp") ) - self.input_layernorm = GemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps) - self.post_attention_layernorm = GemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.input_layernorm = GemmaPlusOneRMSNormFused(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = GemmaPlusOneRMSNormFused(config.hidden_size, eps=config.rms_norm_eps) @nvtx_annotate("Layer_{}", layer_id_field="_layer_id") def forward(self, hidden: torch.Tensor, residual: torch.Tensor | None): - # Residual-stream form: fuse each residual-add into the next RMSNorm - # (GemmaRMSNorm.forward_add_residual) so add + norm are one kernel per sublayer. - if residual is None: - residual = hidden - hidden = self.input_layernorm.forward(hidden) - else: - hidden, residual = self.input_layernorm.forward_add_residual(hidden, residual) + # Residual-stream form: fuse each residual-add into the next RMSNorm so add + norm are one kernel per sublayer. + hidden, residual = self.input_layernorm.forward(hidden, residual) hidden = self.linear_attn.forward(hidden) if self._is_linear else self.self_attn.forward(hidden) - hidden, residual = self.post_attention_layernorm.forward_add_residual(hidden, residual) + hidden, residual = self.post_attention_layernorm.forward(hidden, residual) hidden = self.mlp.forward(hidden) return hidden, residual @@ -79,6 +76,8 @@ def __init__(self, config: ModelConfig, *, prefix: str = "model"): self.embed_tokens = VocabParallelEmbedding( num_embeddings=config.vocab_size, embedding_dim=config.hidden_size, + quant_config=config.quant, + prefix=f"{prefix}.embed_tokens", ) self.layers = OPList( [ @@ -86,14 +85,14 @@ def __init__(self, config: ModelConfig, *, prefix: str = "model"): for layer_id in range(config.num_layers) ] ) - self.norm = GemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.norm = GemmaPlusOneRMSNormFused(config.hidden_size, eps=config.rms_norm_eps) def forward(self, input_ids: torch.Tensor) -> torch.Tensor: x = embed_input_ids(self.embed_tokens, input_ids, get_global_ctx().batch) residual: torch.Tensor | None = None for layer in self.layers.op_list: x, residual = layer.forward(x, residual) - x, _ = self.norm.forward_add_residual(x, residual) + x, _ = self.norm.forward(x, residual) return x diff --git a/python/freetoken/models/qwen3_5_moe/weight.py b/python/freetoken/models/qwen3_5_moe/weight.py index 1bac8227b0..afd32eb2c0 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 @@ -35,13 +35,6 @@ # role -> the expert bank reader's canonical (ModelOpt) tensor kind _BANK_KINDS = {"weight": "weight", "weight_scale": "weight_scale", "weight_global": "weight_scale_2"} -# Gemma-style (1+weight) RMSNorm weights; the GDN gated norm (linear_attn.norm) is a plain weight*x norm -_GEMMA_NORM_SUFFIXES = ( - ".input_layernorm.weight", - ".post_attention_layernorm.weight", - ".self_attn.q_norm.weight", - ".self_attn.k_norm.weight", -) # leaves the model builds as Linear layers: only their tensors are read under the QuantConfig, the rest passes through as stored _LINEAR_LEAVES = frozenset({ "q_proj", "k_proj", "v_proj", "o_proj", "in_proj_qkv", "in_proj_z", "in_proj_b", "in_proj_a", "out_proj", @@ -63,10 +56,6 @@ def _rename(raw_name: str) -> str | None: return rename_vl_prefix(raw_name) -def _is_gemma_norm(name: str) -> bool: - return name == "model.norm.weight" or name.endswith(_GEMMA_NORM_SUFFIXES) - - def _per_row_scale(scale: torch.Tensor, rows: int) -> torch.Tensor: """Per-tensor scalar or per-channel ``[rows, 1]`` fp8 scale -> fp32 ``[rows]``; any other count is refused rather than broadcast onto the wrong rows.""" flat = scale.reshape(-1).to(torch.float32) @@ -149,7 +138,7 @@ def add(self, name: str, tensor: torch.Tensor) -> list[tuple[str, torch.Tensor]] if module.rpartition(".")[2] not in _LINEAR_LEAVES: return None stored = self.stored(module) - roles = {"weight": "weight"} if stored is None else {e.name: r for r, e in self.quant.storage(stored).items()} + roles = {"weight": "weight"} if stored is None else {n: r for r, e in self.quant.storage(stored).items() for n in e.names} role = roles.get(suffix) if role is None: if suffix in _DROPPED_SUFFIXES: @@ -227,7 +216,7 @@ def _check(self, target: str, scheme: QuantScheme, part: dict[str, torch.Tensor] g = part["weight_global"].reshape(-1).to(torch.float32) if g.numel() != 1: raise ValueError(f"{target}: weight_global has {g.numel()} elements, expected one per-tensor scale") - out["weight_global"] = g.to(torch.float16).expand(rows).contiguous() + out["weight_global"] = g.expand(rows).contiguous() if "input_scale" in part: out["input_scale"] = part["input_scale"].reshape(()).to(torch.float32) return out @@ -241,9 +230,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 +242,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): @@ -276,8 +263,6 @@ def _iter_shards(model_path: str, device: torch.device, reader: _DenseReader | N emitted = reader.add(name, tensor) if emitted is not None: yield from emitted - elif _is_gemma_norm(name): - yield name, tensor + 1.0 # (1 + weight) baked into the stored norm weight else: yield name, tensor if reader is not None and reader.pending: @@ -327,7 +312,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 +323,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 ( @@ -375,7 +331,7 @@ def _moe_dims(model_config): ) -def iter_expert_pieces(model_path, config, kind: QuantKind, *, parallel: bool | None = False, workers: int = 8, chunk: int = 8 << 20): +def iter_expert_pieces(model_path, config, kind: QuantKind, *, layers=None, parallel: bool | None = False, workers: int = 8, chunk: int = 8 << 20): """Block-fp8 routed experts, one piece per expert: ``{gate, up, down}`` fp8 codes and their ``_scale`` (block scale) companions, named as the checkpoint's dialect stores them. Other expert kinds use the generic readers.""" if kind is not QuantKind.FP8_BLOCK: @@ -386,6 +342,7 @@ def iter_expert_pieces(model_path, config, kind: QuantKind, *, parallel: bool | from freetoken.moe.expert_pieces import per_expert_pieces L, E, H, I, dense = _moe_dims(config) + keep = None if layers is None else set(layers) scale = get_quant_config().stored_tensors(QuantKind.FP8_BLOCK)["weight_scale_inv"].name key_re = re.compile(_FP8_EXPERT_KEY_RE.format(scale=re.escape(scale))) suffix = {"weight": "", scale: "_scale"} @@ -397,20 +354,22 @@ def locate(raw_name: str): li = int(m["layer"]) - dense if not 0 <= li < L: raise ValueError(f"unexpected routed-expert layer in {raw_name}") + if keep is not None and li not in keep: + return None return li, int(m["expert"]), m["proj"] + suffix[m["kind"]] if parallel is None: parallel = experts_scattered(model_path) if parallel: tensors = iter_expert_tensors_parallel( - model_path, lambda n: key_re.match(n) is not None, workers=workers, chunk=chunk + 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, torch.device("cpu")) try: - for li in tqdm(range(L), desc="Loading fp8 experts (serial)", disable=not get_tp_info().is_primary()): + for li in tqdm(range(L) if keep is None else sorted(keep), desc="Loading fp8 experts (serial)", disable=not get_tp_info().is_primary()): for e in range(E): base = f"model.language_model.layers.{dense + li}.mlp.experts.{e}" for proj in ("gate", "up", "down"): diff --git a/python/freetoken/models/qwen3_moe/__init__.py b/python/freetoken/models/qwen3_moe/__init__.py index 59c74fb0cb..cb96a1921b 100644 --- a/python/freetoken/models/qwen3_moe/__init__.py +++ b/python/freetoken/models/qwen3_moe/__init__.py @@ -1,4 +1,5 @@ from .config import parse_config +from .gguf import iter_gguf_weights, parse_gguf_config, to_gguf from .model import Qwen3MoeForCausalLM from .weight import iter_weights, iter_weights_parallel @@ -7,4 +8,7 @@ "parse_config", "iter_weights", "iter_weights_parallel", + "iter_gguf_weights", + "parse_gguf_config", + "to_gguf", ] diff --git a/python/freetoken/models/qwen3_moe/gguf.py b/python/freetoken/models/qwen3_moe/gguf.py new file mode 100644 index 0000000000..3a8ef98a7a --- /dev/null +++ b/python/freetoken/models/qwen3_moe/gguf.py @@ -0,0 +1,147 @@ +"""Qwen3 MoE from a llama.cpp ``qwen3moe`` GGUF: config from the GGUF metadata, tensors by name; the routed experts load as GGUF expert banks (``freetoken.models.gguf.experts``). + +The llama.cpp converter (conversion/qwen.py) stores these tensors unchanged: both sides run NeoX rope, so q/k carry no rope permute, and the RMSNorm weights are the HF ones. +""" + +from __future__ import annotations + +import re +from dataclasses import replace +from types import SimpleNamespace +from typing import TYPE_CHECKING, Iterator + +import torch + +from freetoken.models.gguf.reader import GgufTensorInfo, gguf_tensor_index +from freetoken.models.gguf.tensors import check_scales, float_tensor, fused, matrix + +from .config import parse_config + +if TYPE_CHECKING: + from freetoken.models.config import ModelConfig + from freetoken.models.gguf.config import GgufConfigShim + +_ARCH = "qwen3moe" +_LAYER_LEAVES = { + "self_attn.q_proj": "attn_q", + "self_attn.k_proj": "attn_k", + "self_attn.v_proj": "attn_v", + "self_attn.o_proj": "attn_output", + "mlp.gate": "ffn_gate_inp", + "mlp.experts.0.gate_proj": "ffn_gate_exps", + "mlp.experts.0.up_proj": "ffn_up_exps", + "mlp.experts.0.down_proj": "ffn_down_exps", +} +_QKV_LEAVES = ("self_attn.q_proj", "self_attn.k_proj", "self_attn.v_proj") +_TOP_LEAVES = {"model.embed_tokens": "token_embd", "lm_head": "output"} +_LAYER_RE = re.compile(r"^model\.layers\.(\d+)\.(.+)$") + + +def to_gguf(name: str, tensors: dict[str, int]) -> tuple[str, ...]: + """Checkpoint module name -> the GGUF tensors (without ``.weight``) that store it; empty for modules the GGUF has no matrix for.""" + if name in _TOP_LEAVES: + return (_TOP_LEAVES[name],) + m = _LAYER_RE.match(name) + if m is None: + return () + blk, leaf = f"blk.{m.group(1)}", m.group(2) + if leaf in _QKV_LEAVES and f"{blk}.attn_qkv.weight" in tensors: + return (f"{blk}.attn_qkv",) + if leaf in ("mlp.experts.0.gate_proj", "mlp.experts.0.up_proj") and f"{blk}.ffn_gate_up_exps.weight" in tensors: + return (f"{blk}.ffn_gate_up_exps",) + tensor = _LAYER_LEAVES.get(leaf) + return (f"{blk}.{tensor}",) if tensor is not None else () + + +def _check_supported(shim: GgufConfigShim) -> None: + if any(k.startswith("prism.hadamard.") for k in shim.metadata): + raise NotImplementedError(f"{shim.model_path}: Hadamard-rotated (prism) GGUF weights are not supported") + check_scales(shim.model_path) + + +def _hf_config(shim: GgufConfigShim) -> SimpleNamespace: + """The HF config fields ``parse_config`` reads, from the GGUF metadata.""" + md = shim.metadata + + def kv(key: str): + return md[f"{_ARCH}.{key}"] + + _check_supported(shim) + rope = {"rope_theta": float(kv("rope.freq_base")), "rope_type": "default"} + scaling = md.get(f"{_ARCH}.rope.scaling.type", "none") + if scaling == "yarn": + rope.update( + rope_type="yarn", + factor=float(kv("rope.scaling.factor")), + original_max_position_embeddings=int(kv("rope.scaling.original_context_length")), + ) + elif scaling != "none": + raise NotImplementedError(f"{shim.model_path}: rope scaling {scaling!r} is not supported") + key_length = md.get(f"{_ARCH}.attention.key_length") + return SimpleNamespace( + model_type="qwen3_moe", + architectures=list(shim.architectures), + hidden_act="silu", + hidden_size=int(kv("embedding_length")), + num_hidden_layers=int(kv("block_count")), + num_attention_heads=int(kv("attention.head_count")), + num_key_value_heads=int(kv("attention.head_count_kv")), + head_dim=int(key_length) if key_length is not None else None, + rope_scaling=rope, + max_position_embeddings=int(kv("context_length")), + rms_norm_eps=float(kv("attention.layer_norm_rms_epsilon")), + vocab_size=int(shim.vocab_size), + tie_word_embeddings=bool(shim.tie_word_embeddings), + intermediate_size=int(kv("feed_forward_length")), + num_experts=int(kv("expert_count")), + num_experts_per_tok=int(kv("expert_used_count")), + moe_intermediate_size=int(kv("expert_feed_forward_length")), + # llama.cpp always renormalizes the top-k weights of a qwen3moe router + norm_topk_prob=True, + ) + + +def parse_gguf_config(shim: GgufConfigShim) -> ModelConfig: + return replace(parse_config(_hf_config(shim)), expert_quant="gguf") + + +def iter_gguf_weights( + model_path: str, + device: torch.device, + *, + include_moe_experts: bool, + include_non_moe: bool, +) -> Iterator[tuple[str, torch.Tensor]]: + """The dense state dict from the GGUF; the routed experts load as GGUF expert banks (``freetoken.models.gguf.experts``).""" + if not include_non_moe: + return + from freetoken.utils import cached_load_hf_config + + hf = _hf_config(cached_load_hf_config(model_path)) + index = gguf_tensor_index(model_path) + hidden = hf.hidden_size + + def t(name: str) -> GgufTensorInfo: + return index[name if name in index else f"{name}.weight"] + + yield from fused("model.embed_tokens", [t("token_embd")], [matrix(t("token_embd"))], hidden) + if not hf.tie_word_embeddings: + yield from fused("lm_head", [t("output")], [matrix(t("output"))], hidden) + yield "model.norm.weight", float_tensor(t("output_norm")) + + for layer in range(hf.num_hidden_layers): + blk, pre = f"blk.{layer}", f"model.layers.{layer}" + yield f"{pre}.input_layernorm.weight", float_tensor(t(f"{blk}.attn_norm")) + yield f"{pre}.post_attention_layernorm.weight", float_tensor(t(f"{blk}.ffn_norm")) + sa = f"{pre}.self_attn" + # convert_hf_to_gguf.py --fuse-qkv stores q, k and v as one attn_qkv matrix + qkv = [t(f"{blk}.attn_qkv")] if f"{blk}.attn_qkv.weight" in index else [t(f"{blk}.attn_{p}") for p in ("q", "k", "v")] + yield from fused(f"{sa}.qkv_proj", qkv, [matrix(p) for p in qkv], hidden) + yield f"{sa}.q_norm.weight", float_tensor(t(f"{blk}.attn_q_norm")) + yield f"{sa}.k_norm.weight", float_tensor(t(f"{blk}.attn_k_norm")) + out = t(f"{blk}.attn_output") + yield from fused(f"{sa}.o_proj", [out], [matrix(out)], int(out.ne[0])) + yield f"{pre}.mlp.gate.weight", float_tensor(t(f"{blk}.ffn_gate_inp")) + + +__all__ = ["iter_gguf_weights", "parse_gguf_config", "to_gguf"] diff --git a/python/freetoken/models/qwen3_moe/model.py b/python/freetoken/models/qwen3_moe/model.py index 5d89656bda..7d81dd2673 100644 --- a/python/freetoken/models/qwen3_moe/model.py +++ b/python/freetoken/models/qwen3_moe/model.py @@ -47,6 +47,8 @@ def __init__(self, config: ModelConfig, *, prefix: str = "model"): self.embed_tokens = VocabParallelEmbedding( num_embeddings=config.vocab_size, embedding_dim=config.hidden_size, + quant_config=config.quant, + prefix=f"{prefix}.embed_tokens", ) self.layers = OPList( [ diff --git a/python/freetoken/models/qwen4_exp/__init__.py b/python/freetoken/models/qwen4_exp/__init__.py index 0fe4a4ffa8..a623985636 100644 --- a/python/freetoken/models/qwen4_exp/__init__.py +++ b/python/freetoken/models/qwen4_exp/__init__.py @@ -10,8 +10,10 @@ """ from .config import parse_config +from .gguf import iter_gguf_weights, parse_gguf_config, to_gguf from .model import Qwen4ExpForCausalLM, Qwen4ExpForConditionalGeneration from .weight import ( + check_ftw_weights, ftw_side_files, nvfp4_expert_spec, iter_vision_weights, @@ -24,6 +26,7 @@ from freetoken.models.qwen3_5_moe.weight import iter_expert_pieces __all__ = [ + "check_ftw_weights", "ftw_side_files", "nvfp4_expert_spec", "Qwen4ExpForCausalLM", @@ -33,4 +36,7 @@ "load_ple_table", "parse_config", "iter_expert_pieces", + "iter_gguf_weights", + "parse_gguf_config", + "to_gguf", ] diff --git a/python/freetoken/models/qwen4_exp/attention.py b/python/freetoken/models/qwen4_exp/attention.py index 8f6b938870..80fcaade1f 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 @@ -85,7 +93,7 @@ def __init__(self, config: ModelConfig, layer_id: int, *, prefix: str = "") -> N self.eps = config.rms_norm_eps self._split = [self.num_heads * self.head_dim, self.num_kv_heads * self.head_dim] self.index_qk_proj = LinearReplicated( - args.hidden_size, sum(self._split), has_bias=False, + args.hidden_size, sum(self._split), has_bias=False, output_sizes=self._split, quant_config=config.quant, prefix=f"{prefix}.index_qk_proj", ) self.q_layernorm = GemmaPlusOneRMSNorm(self.head_dim, eps=self.eps) @@ -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..cea9120730 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, @@ -43,6 +45,8 @@ class Qwen4ExpArgs: index_budget: int index_ratio: int image_token_id: int | None = None + # GGUF stores ssm_out with its value-head columns in llama.cpp's tiled order + gdn_out_heads_tiled: bool = False @property def index_topk_blocks(self) -> int: @@ -99,9 +103,9 @@ def ple_slot_states(args: Qwen4ExpArgs) -> Tuple[SlotStateSpec, ...]: def _layer_types(text: Any) -> list[str]: layer_types = getattr(text, "layer_types", None) if layer_types is not None: - # HF Qwen4ExpTextConfig rewrites full_attention to qwen_sparse_attention in __post_init__. + # HF Qwen4ExpTextConfig rewrites full_attention to indexed_attention in __post_init__. return [ - "full_attention" if t == "qwen_sparse_attention" else t for t in layer_types + "full_attention" if t == "indexed_attention" else t for t in layer_types ] # Fall back to full_attention_interval: every Nth layer (1-indexed) is full. interval = int(getattr(text, "full_attention_interval", 4)) @@ -112,6 +116,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 +162,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 +283,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..dbb44c28c3 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 @@ -32,9 +35,11 @@ class Qwen4ExpGatedDeltaNet(BaseOP): 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 = "", + *, quant_config: QuantConfig | None = None, prefix: str = "", out_heads_tiled: bool = False, + attn_quant: str = "none", ): self.layer_id = layer_id + self._out_heads_tiled = out_heads_tiled # The fla chunk/decode kernels read+write the recurrent state and the per-chunk h as # [V, K] while the LinearStatePool declares it [K, V]; these coincide (and the # hybrid-radix snapshot scatter h[h_row]->slot is a plain copy) only when the two head @@ -50,13 +55,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 +102,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 +134,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 +185,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 +204,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 +215,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: @@ -197,6 +245,10 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: core_out = core_out.reshape(-1, self.head_v_dim) z = z.reshape(-1, self.head_v_dim) out = self.norm.forward(core_out, z).reshape(total, -1) + if self._out_heads_tiled: + # grouped value heads (k-head major) -> tiled (k-head minor), the column order of a GGUF ssm_out + r = nv // nk + out = out.view(total, nk, r, self.head_v_dim).transpose(1, 2).reshape(total, -1) return self.out_proj.forward(out) diff --git a/python/freetoken/models/qwen4_exp/gguf.py b/python/freetoken/models/qwen4_exp/gguf.py new file mode 100644 index 0000000000..f29f65337e --- /dev/null +++ b/python/freetoken/models/qwen4_exp/gguf.py @@ -0,0 +1,271 @@ +"""Qwen3.8-Flash-Next from a llama.cpp ``qwen4exp`` GGUF: config from the GGUF metadata, tensor names, and the converter's transforms undone. + +The llama.cpp converter (conversion/qwen.py, conversion/qwen4exp.py) stores every zero-centred norm as ``w + 1``, ``A_log`` as ``-exp(A_log)``, the squeezed conv kernels, the indexer projection split into q and k, and the GDN value heads in tiled order (k-head minor). Row reorders are undone on the packed bytes; ``ssm_out`` keeps its tiled columns and the GDN layer permutes its input instead (``Qwen4ExpArgs.gdn_out_heads_tiled``). +""" + +from __future__ import annotations + +import re +from dataclasses import replace +from types import SimpleNamespace +from typing import TYPE_CHECKING, Iterator + +import torch + +from freetoken.layers.quantization.ggml_types import is_quantized +from freetoken.models.gguf.reader import GgufTensorInfo, gguf_tensor_index +from freetoken.models.gguf.tensors import check_scales, float_tensor, fused, matrix + +from .config import parse_config + +if TYPE_CHECKING: + from freetoken.models.config import ModelConfig + from freetoken.models.gguf.config import GgufConfigShim + +_ARCH = "qwen4exp" + +# HF config fields the GGUF metadata does not carry (Qwen3.8-Flash-Next's own config.json values) +_FAMILY_DEFAULTS = dict( + hidden_act="silu", + output_gate_type="sigmoid", + indexer_kv_heads=1, + ngram_vocab_size_base=20000000, + make_ngram_vocab_size_divisible_by=128, + split_ngram_parts=128, + norm_topk_prob=True, +) + +_LAYER_LEAVES: dict[str, tuple[str, ...]] = { + "linear_attn.in_proj_qkv": ("attn_qkv",), + "linear_attn.in_proj_z": ("attn_gate",), + "linear_attn.in_proj_b": ("ssm_beta",), + "linear_attn.in_proj_a": ("ssm_alpha",), + "linear_attn.out_proj": ("ssm_out",), + "self_attn.q_proj": ("attn_q",), + "self_attn.k_proj": ("attn_k",), + "self_attn.v_proj": ("attn_v",), + "self_attn.o_proj": ("attn_output",), + "self_attn.indexer.index_q_proj": ("indexer.q_proj",), + "self_attn.indexer.index_k_proj": ("indexer.k_proj",), + "mlp.gate": ("ffn_gate_inp",), + "mlp.shared_expert_gate": ("ffn_gate_inp_shexp",), + "mlp.shared_expert.gate_proj": ("ffn_gate_shexp",), + "mlp.shared_expert.up_proj": ("ffn_up_shexp",), + "mlp.shared_expert.down_proj": ("ffn_down_shexp",), + "mlp.experts.0.gate_proj": ("ffn_gate_exps",), + "mlp.experts.0.up_proj": ("ffn_up_exps",), + "mlp.experts.0.down_proj": ("ffn_down_exps",), + "attn_hyper_connection.input_mix_weight_down": ("hc_attn_down",), + "attn_hyper_connection.block_inject_weight": ("hc_attn_inject",), + "attn_hyper_connection.input_mix_weight_up": ("hc_attn_up",), + "mlp_hyper_connection.input_mix_weight_down": ("hc_ffn_down",), + "mlp_hyper_connection.block_inject_weight": ("hc_ffn_inject",), + "mlp_hyper_connection.input_mix_weight_up": ("hc_ffn_up",), + "ple.key_proj": ("ple_key",), + "ple.value_proj": ("ple_value",), +} +_TOP_LEAVES: dict[str, tuple[str, ...]] = { + "model.language_model.embed_tokens": ("token_embd",), + "lm_head": ("output",), + "model.language_model.hyper_connection_mixer.input_mix_weight_down": ("output_hc_down",), + "model.language_model.hyper_connection_mixer.input_mix_weight_up": ("output_hc_up",), +} +_LAYER_RE = re.compile(r"^model\.language_model\.layers\.(\d+)\.(.+)$") + + +def to_gguf(name: str, tensors: dict[str, int]) -> tuple[str, ...]: + """Checkpoint module name -> the GGUF tensors (without ``.weight``) that store it; empty for modules the GGUF has no matrix for.""" + if name in _TOP_LEAVES: + return _TOP_LEAVES[name] + m = _LAYER_RE.match(name) + if m is None: + return () + return tuple(f"blk.{m.group(1)}.{leaf}" for leaf in _LAYER_LEAVES.get(m.group(2), ())) + + +def parse_gguf_config(shim: GgufConfigShim) -> ModelConfig: + # NVFP4 weight scales are unvalidated on this arch, so refuse them instead of risking wrong output + check_scales(shim.model_path, nvfp4=False) + # a metadata-only GGUF (an FTW dir's source_metadata.gguf) has no tensor table to check + index = gguf_tensor_index(shim.model_path) + if index and "blk.0.hc_attn_down.weight" not in index: + raise NotImplementedError(f"{shim.model_path}: no decoder layer tensors (blk.0.hc_attn_down), e.g. an MTP-only draft GGUF") + md = shim.metadata + + def kv(key: str): + return md[f"{_ARCH}.{key}"] + + head_dim = int(kv("attention.key_length")) + sections = [int(s) for s in kv("rope.dimension_sections")][:3] + ratios = [int(r) for r in kv("attention.compress_ratios")] + n_v = int(kv("ssm.time_step_rank")) + text = SimpleNamespace( + **_FAMILY_DEFAULTS, + hidden_size=int(kv("embedding_length")), + # block_count also counts the MTP layers a GGUF may carry after the decoder + num_hidden_layers=int(kv("block_count")) - int(md.get(f"{_ARCH}.nextn_predict_layers", 0)), + num_attention_heads=int(kv("attention.head_count")), + num_key_value_heads=int(kv("attention.head_count_kv")), + head_dim=head_dim, + rope_parameters={ + "rope_type": "default", + "rope_theta": float(kv("rope.freq_base")), + "partial_rotary_factor": int(kv("rope.dimension_count")) / head_dim, + "mrope_section": sections, + "mrope_interleaved": True, + }, + max_position_embeddings=int(kv("context_length")), + full_attention_interval=int(kv("full_attention_interval")), + rms_norm_eps=float(kv("attention.layer_norm_rms_epsilon")), + vocab_size=int(shim.vocab_size), + tie_word_embeddings=bool(shim.tie_word_embeddings), + num_experts=int(kv("expert_count")), + num_experts_per_tok=int(kv("expert_used_count")), + moe_intermediate_size=int(kv("expert_feed_forward_length")), + shared_expert_intermediate_size=int(kv("expert_shared_feed_forward_length")), + linear_num_key_heads=int(kv("ssm.group_count")), + linear_num_value_heads=n_v, + linear_key_head_dim=int(kv("ssm.state_size")), + linear_value_head_dim=int(kv("ssm.inner_size")) // n_v, + linear_conv_kernel_dim=int(kv("ssm.conv_kernel")), + hc_count=int(kv("hyper_connection.count")), + hc_lowrank=int(kv("hyper_connection.low_rank")), + indexer_n_heads=int(kv("attention.indexer.head_count")), + indexer_head_dim=int(kv("attention.indexer.key_length")), + indexer_budget=int(kv("attention.indexer.top_k")), + indexer_compress_ratio=max(ratios), + ple_layer_ids=[int(i) + 1 for i in kv("ple.layers")], + # the GGUF key is one n-gram head's row; HF counts all (ngram_size - 1) * heads_per_ngram heads + ple_embed_dim=int(kv("embedding_length_per_layer_input")) * (int(kv("ple.ngram_size")) - 1) * int(kv("ple.heads_per_ngram")), + ple_conv_kernel_size=int(kv("ple.conv_kernel")), + ngram_size=int(kv("ple.ngram_size")), + heads_per_ngram=int(kv("ple.heads_per_ngram")), + eos_token_id=int(kv("ple.eos_token_id")), + ) + image_token_id = md.get(f"{_ARCH}.ple.image_token_id") + hf_like = SimpleNamespace( + text_config=text, + model_type="qwen4_exp", + architectures=list(shim.architectures), + image_token_id=None if image_token_id is None else int(image_token_id), + vision_config=None, + ) + config = parse_config(hf_like) + return replace(config, expert_quant="gguf", qwen4_args=replace(config.qwen4_args, gdn_out_heads_tiled=True)) + + +# ---------------------------------------------------------------------------------------------- +# Weights +# ---------------------------------------------------------------------------------------------- + +def _untile(t: torch.Tensor, num_k_heads: int, head_dim: int, dim: int = 0) -> torch.Tensor: + """Tiled value heads (k-head minor) back to grouped (k-head major) along ``dim``.""" + n = t.shape[dim] + r = n // (num_k_heads * head_dim) + idx = torch.arange(n).view(r, num_k_heads, head_dim).transpose(0, 1).reshape(n) + return t.index_select(dim, idx) + + +def iter_gguf_weights( + model_path: str, + device: torch.device, + *, + include_moe_experts: bool, + include_non_moe: bool, + include_vision: bool = True, +) -> Iterator[tuple[str, torch.Tensor]]: + """The dense state dict from the GGUF; the routed experts load as offload banks (``freetoken.models.gguf.experts``).""" + if not include_non_moe: + return + from freetoken.utils import cached_load_hf_config + + config = parse_gguf_config(cached_load_hf_config(model_path)) + index = gguf_tensor_index(model_path) + md = cached_load_hf_config(model_path).metadata + g = config.linear_attention_group() + hidden = config.hidden_size + + def t(name: str) -> GgufTensorInfo: + return index[name if name in index else f"{name}.weight"] + + def norm(name: str) -> torch.Tensor: + return float_tensor(t(name)).float() - 1.0 + + def hc(prefix: str, src: str) -> Iterator[tuple[str, torch.Tensor]]: + down, inject, up = (t(f"{src}_{p}") for p in ("down", "inject", "up")) + yield from fused(f"{prefix}.input_mix_weight_down_block_inject", [down, inject], [matrix(down), matrix(inject)], + hidden * config.qwen4_args.hc_count, pad_to=16) + yield from fused(f"{prefix}.input_mix_weight_up", [up], [matrix(up)], config.qwen4_args.hc_lowrank) + yield f"{prefix}.hc_norm.weight", norm(f"{src}_norm") + + yield from fused("model.embed_tokens", [t("token_embd")], [matrix(t("token_embd"))], hidden) + yield from fused("lm_head", [t("output")], [matrix(t("output"))], hidden) + for proj, k in (("down", hidden * config.qwen4_args.hc_count), ("up", config.qwen4_args.hc_lowrank)): + info = t(f"output_hc_{proj}") + yield from fused(f"model.hyper_connection_mixer.input_mix_weight_{proj}", [info], [matrix(info)], k) + yield "model.hyper_connection_mixer.hc_norm.weight", norm("output_hc_norm") + + for layer in range(config.num_layers): + blk, pre = f"blk.{layer}", f"model.layers.{layer}" + if config.is_linear_layer(layer): + la = f"{pre}.linear_attn" + nk, dk = g.num_key_heads, g.key_head_dim + qkv, gate = t(f"{blk}.attn_qkv"), t(f"{blk}.attn_gate") + qkv_rows = matrix(qkv) + v_start = 2 * nk * dk + qkv_rows = torch.cat([qkv_rows[:v_start], _untile(qkv_rows[v_start:], nk, g.value_head_dim)]) + z_rows = _untile(matrix(gate), nk, g.value_head_dim) + beta_info, alpha_info = t(f"{blk}.ssm_beta"), t(f"{blk}.ssm_alpha") + beta, alpha = (_untile(matrix(i), nk, 1) for i in (beta_info, alpha_info)) + if is_quantized(qkv.ggml_type) or is_quantized(gate.ggml_type): + yield from fused(f"{la}.in_proj_qkvz", [qkv, gate], [qkv_rows, z_rows], hidden) + yield from fused(f"{la}.in_proj_ba", [beta_info, alpha_info], [beta, alpha], hidden) + else: + yield from fused(f"{la}.in_proj", [qkv, gate, beta_info, alpha_info], [qkv_rows, z_rows, beta, alpha], hidden) + conv = float_tensor(t(f"{blk}.ssm_conv1d")).float() + conv = torch.cat([conv[:v_start], _untile(conv[v_start:], nk, g.value_head_dim)]) + yield f"{la}.conv1d.weight", conv.unsqueeze(1) + yield f"{la}.dt_bias", _untile(float_tensor(t(f"{blk}.ssm_dt.bias")).float(), nk, 1) + yield f"{la}.A_log", _untile(torch.log(-float_tensor(t(f"{blk}.ssm_a")).float()), nk, 1) + yield f"{la}.norm.weight", float_tensor(t(f"{blk}.ssm_norm")) + yield from fused(f"{la}.out_proj", [t(f"{blk}.ssm_out")], [matrix(t(f"{blk}.ssm_out"))], g.num_value_heads * g.value_head_dim) + else: + sa = f"{pre}.self_attn" + parts = [t(f"{blk}.attn_{p}") for p in ("q", "k", "v")] + yield from fused(f"{sa}.qkv_proj", parts, [matrix(p) for p in parts], hidden) + out = t(f"{blk}.attn_output") + yield from fused(f"{sa}.o_proj", [out], [matrix(out)], int(out.ne[0])) + yield f"{sa}.q_norm.weight", norm(f"{blk}.attn_q_norm") + yield f"{sa}.k_norm.weight", norm(f"{blk}.attn_k_norm") + iq, ik = t(f"{blk}.indexer.q_proj"), t(f"{blk}.indexer.k_proj") + yield from fused(f"{sa}.indexer.index_qk_proj", [iq, ik], [matrix(iq), matrix(ik)], hidden) + yield f"{sa}.indexer.q_layernorm.weight", norm(f"{blk}.indexer.q_norm") + yield f"{sa}.indexer.k_layernorm.weight", norm(f"{blk}.indexer.k_norm") + + mlp = f"{pre}.mlp" + yield f"{mlp}.gate.weight", float_tensor(t(f"{blk}.ffn_gate_inp")) + sg, su, sd = (t(f"{blk}.ffn_{p}_shexp") for p in ("gate", "up", "down")) + yield from fused(f"{mlp}.shared_expert.gate_up_proj", [sg, su], [matrix(sg), matrix(su)], hidden) + yield from fused(f"{mlp}.shared_expert.down_proj", [sd], [matrix(sd)], config.shared_expert_intermediate_size) + yield f"{mlp}.shared_expert_gate.weight", float_tensor(t(f"{blk}.ffn_gate_inp_shexp")).reshape(1, -1) + yield from hc(f"{pre}.attn_hyper_connection", f"{blk}.hc_attn") + yield from hc(f"{pre}.mlp_hyper_connection", f"{blk}.hc_ffn") + + if layer in config.qwen4_args.ple_layer_ids: + ple = f"{pre}.ple" + for proj in ("key", "value"): + info = t(f"{blk}.ple_{proj}") + yield from fused(f"{ple}.{proj}_proj", [info], [matrix(info)], config.qwen4_args.ple_embed_dim) + for n in ("key", "query", "conv"): + yield f"{ple}.norm_{n}.weight", norm(f"{blk}.ple_norm_{n}") + yield f"{ple}.conv1d.weight", float_tensor(t(f"{blk}.ple_conv1d")).float().unsqueeze(1) + yield f"{ple}.ple_embedding.layer_multipliers", torch.tensor([int(x) for x in md[f"{_ARCH}.ple.layer_multipliers"]], dtype=torch.int64) + yield f"{ple}.ple_embedding.ngram_heads_vocab_sizes", torch.tensor([int(x) for x in md[f"{_ARCH}.ple.head_vocab_sizes"]], dtype=torch.int64) + yield f"{ple}.ple_embedding.ngram_heads_offsets", torch.tensor([int(x) for x in md[f"{_ARCH}.ple.head_offsets"]], dtype=torch.int64) + + +def ple_table_info(model_path: str) -> GgufTensorInfo: + return gguf_tensor_index(model_path)["per_layer_token_embd.weight"] + + +__all__ = ["iter_gguf_weights", "parse_gguf_config", "ple_table_info", "to_gguf"] diff --git a/python/freetoken/models/qwen4_exp/hc.py b/python/freetoken/models/qwen4_exp/hc.py index f13808dccd..2071916db0 100644 --- a/python/freetoken/models/qwen4_exp/hc.py +++ b/python/freetoken/models/qwen4_exp/hc.py @@ -17,6 +17,7 @@ from freetoken.kernel.triton.hc import ( grouped_gemma_rmsnorm, hc_combine, + hc_combine_norm, hc_gate_mix, hc_silu, ) @@ -98,6 +99,7 @@ def __init__(self, config: ModelConfig, use_combine: bool = True, *, prefix: str self.pad_size = (-(self.lowrank + self.hc_count)) % 16 self.input_mix_weight_down_block_inject = LinearReplicated( width, self.lowrank + self.hc_count + self.pad_size, has_bias=False, + output_sizes=[self.lowrank, self.hc_count + self.pad_size], quant_config=config.quant, prefix=f"{prefix}.input_mix_weight_down_block_inject", ) else: @@ -119,14 +121,16 @@ def _down(self, rn: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor | None]: # both slices keep unit inner stride, so the kernels read them without a copy return down[:, : self.lowrank], down[:, self.lowrank : self.lowrank + self.hc_count] - def _mix_kernel(self, R: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor | None]: - rn = grouped_gemma_rmsnorm(R, self.hc_norm.weight, self.hc_norm.eps, self.hc_count) + def _mix_kernel(self, R: torch.Tensor, rn: torch.Tensor | None = None) -> Tuple[torch.Tensor, torch.Tensor | None]: + if rn is None: + rn = grouped_gemma_rmsnorm(R, self.hc_norm.weight, self.hc_norm.eps, self.hc_count) lora, s = self._down(rn) gate = self.input_mix_weight_up.forward(hc_silu(lora, self.hc_count)) return hc_gate_mix(rn, gate, self.hc_count), s - def _mix_torch(self, R: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor | None]: - rn = grouped_plus_one_rms_norm(R, self.hc_norm.weight, self.hc_norm.eps, self.hc_count) + def _mix_torch(self, R: torch.Tensor, rn: torch.Tensor | None = None) -> Tuple[torch.Tensor, torch.Tensor | None]: + if rn is None: + rn = grouped_plus_one_rms_norm(R, self.hc_norm.weight, self.hc_norm.eps, self.hc_count) lora, s = self._down(rn) lora = F.silu(lora.float() / self.hc_count) gate = self.input_mix_weight_up.forward(lora.to(R.dtype)) @@ -140,9 +144,12 @@ def _combine_torch(self, R: torch.Tensor, y: torch.Tensor, s: torch.Tensor) -> t out = out + y.float().unsqueeze(-2) * inject.unsqueeze(-1) return out.flatten(-2).to(R.dtype) - def mix(self, R: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor | None]: - """Return the block input ``x [T, hidden]`` and the inject logits ``s [T, hc_count]`` (None if no combine).""" - return self._mix_kernel(R) if R.is_cuda else self._mix_torch(R) + def mix(self, R: torch.Tensor, rn: torch.Tensor | None = None) -> Tuple[torch.Tensor, torch.Tensor | None]: + """Return the block input ``x [T, hidden]`` and the inject logits ``s [T, hc_count]`` (None if no combine). + + ``rn`` is the pre-normalized R from a preceding :meth:`combine_norm` (same norm weight as + this block's ``hc_norm``), letting the caller fuse the two kernels; None recomputes it.""" + return self._mix_kernel(R, rn) if R.is_cuda else self._mix_torch(R, rn) def combine(self, R: torch.Tensor, y: torch.Tensor, s: torch.Tensor) -> torch.Tensor: """Inject the block output ``y [T, hidden]`` back into every stream of ``R``.""" @@ -150,5 +157,15 @@ def combine(self, R: torch.Tensor, y: torch.Tensor, s: torch.Tensor) -> torch.Te return hc_combine(R, y, s, self.hc_count) return self._combine_torch(R, y, s) + def combine_norm( + self, R: torch.Tensor, y: torch.Tensor, s: torch.Tensor, next_block: "GatedResidual" + ) -> Tuple[torch.Tensor, torch.Tensor | None]: + """combine() fused with ``next_block``'s hc_norm; returns (R', Rn'). CPU keeps the split path (Rn' = None).""" + if not R.is_cuda: + return self.combine(R, y, s), None + return hc_combine_norm( + R, y, s, next_block.hc_norm.weight, next_block.hc_norm.eps, self.hc_count + ) + __all__ = ["GatedResidual", "GroupedPlusOneRMSNorm", "grouped_plus_one_rms_norm"] diff --git a/python/freetoken/models/qwen4_exp/model.py b/python/freetoken/models/qwen4_exp/model.py index 258badac42..63fa971745 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,8 @@ def build_linear_mixer(config: ModelConfig, layer_id: int, prefix: str) -> BaseO output_gate=g.output_gate, quant_config=config.quant, prefix=prefix, + out_heads_tiled=config.qwen4_args.gdn_out_heads_tiled, + attn_quant=config.attn_quant, ) @@ -73,17 +76,19 @@ def __init__(self, config: ModelConfig, layer_id: int, *, prefix: str = "") -> N ) @nvtx_annotate("Layer_{}", layer_id_field="_layer_id") - def forward(self, hidden: torch.Tensor, batch: Batch) -> torch.Tensor: + def forward(self, hidden: torch.Tensor, batch: Batch, rn: torch.Tensor | None = None): + """Returns (R, mlp_out, inject) un-combined: the caller fuses the mlp combine with the + next consumer's norm (next layer's attn mix, or the final mixer).""" if self.ple is not None: hidden = hidden + self.ple.forward(hidden, batch) - block_input, inject = self.attn_hyper_connection.mix(hidden) + block_input, inject = self.attn_hyper_connection.mix(hidden, rn=rn) if self._is_linear: block_output = self.linear_attn.forward(block_input) else: block_output = self.self_attn.forward(block_input, batch) - hidden = self.attn_hyper_connection.combine(hidden, block_output, inject) - block_input, inject = self.mlp_hyper_connection.mix(hidden) - return self.mlp_hyper_connection.combine(hidden, self.mlp.forward(block_input), inject) + hidden, rn_mlp = self.attn_hyper_connection.combine_norm(hidden, block_output, inject, self.mlp_hyper_connection) + block_input, inject = self.mlp_hyper_connection.mix(hidden, rn=rn_mlp) + return hidden, self.mlp.forward(block_input), inject class Qwen4ExpModel(BaseOP): @@ -92,6 +97,8 @@ def __init__(self, config: ModelConfig, *, prefix: str = "model") -> None: self.embed_tokens = VocabParallelEmbedding( num_embeddings=config.vocab_size, embedding_dim=config.hidden_size, + quant_config=config.quant, + prefix=f"{prefix}.embed_tokens", ) self.layers = OPList( [ @@ -99,6 +106,7 @@ def __init__(self, config: ModelConfig, *, prefix: str = "model") -> None: for layer_id in range(config.num_layers) ] ) + wire_router_lookahead([layer.mlp for layer in self.layers.op_list], config) self.hyper_connection_mixer = GatedResidual(config, use_combine=False, prefix=f"{prefix}.hyper_connection_mixer") # plain tuple (not an OP child), so it never shows up in the state dict self._ple = tuple(layer.ple for layer in self.layers.op_list if layer.ple is not None) @@ -118,27 +126,48 @@ def forward(self, input_ids: torch.Tensor, batch: Batch) -> torch.Tensor: meta = build_ple_metadata(batch, self._ple[0].args, input_ids.device) for ple in self._ple: # gather the pinned-host PLE rows while the early layers run ple.start_prefetch(batch, meta) - for layer in self.layers.op_list: - hidden = layer.forward(hidden, batch) + # The last block of each layer combines through the layer boundary, fusing the + # combine with the next consumer's hc_norm -- except before a PLE layer, whose + # `hidden + ple(hidden)` happens after the boundary and changes what the norm sees. + layers = self.layers.op_list + rn = None + for i, layer in enumerate(layers): + hidden, mlp_out, inject = layer.forward(hidden, batch, rn=rn) + nxt = layers[i + 1] if i + 1 < len(layers) else None + if nxt is not None and nxt.ple is not None: + hidden = layer.mlp_hyper_connection.combine(hidden, mlp_out, inject) + rn = None + else: + consumer = nxt.attn_hyper_connection if nxt is not None else self.hyper_connection_mixer + hidden, rn = layer.mlp_hyper_connection.combine_norm(hidden, mlp_out, inject, consumer) if meta is not None: # single writer: the layers only read the context, so a second PLE layer's # prefetch sees the un-rolled window commit_ngram_context(meta, getattr(batch, "fla_metadata", None)) - return self.hyper_connection_mixer.mix(hidden)[0] + return self.hyper_connection_mixer.mix(hidden, rn=rn)[0] class Qwen4ExpForCausalLM(BaseLLMModel): def __init__(self, config: ModelConfig) -> None: self._config = config self.model = Qwen4ExpModel(config) - self.lm_head = ParallelLMHead( - num_embeddings=config.vocab_size, - embedding_dim=config.hidden_size, - tie_word_embeddings=config.tie_word_embeddings, - tied_embedding=self.model.embed_tokens if config.tie_word_embeddings else None, - quant_config=config.quant, - prefix="lm_head", - ) + if use_fp8_lmhead(config): + # Load-time per-tensor FP8 head, for a checkpoint that ships lm_head unquantized. + # A checkpoint that declares its own scheme goes through QuantConfig instead. + from freetoken.layers.fp8_dynamic import Fp8ParallelLMHead + + self.lm_head = Fp8ParallelLMHead( + num_embeddings=config.vocab_size, embedding_dim=config.hidden_size + ) + else: + self.lm_head = ParallelLMHead( + num_embeddings=config.vocab_size, + embedding_dim=config.hidden_size, + tie_word_embeddings=config.tie_word_embeddings, + tied_embedding=self.model.embed_tokens if config.tie_word_embeddings else None, + quant_config=config.quant, + prefix="lm_head", + ) super().__init__() def load_host_tables(self, engine_config) -> int: @@ -168,11 +197,8 @@ def load_host_tables(self, engine_config) -> int: return 0 if engine_config.ple_backend == "disk": - from freetoken.utils import download_hf_weight - from .ple_disk import DiskRowTable, resolve_row_source - folder = download_hf_weight(engine_config.model_path) # one WAIT node per captured graph: the flag protocol supports a single consume assert len(ple_layers) == 1, "disk PLE backend expects exactly one PLE layer" emb, args = ple_layers[0].ple_embedding, ple_layers[0].args @@ -186,7 +212,7 @@ def load_host_tables(self, engine_config) -> int: "image_token_id": args.image_token_id, } disk_table = DiskRowTable( - resolve_row_source(folder), + resolve_row_source(engine_config.model_path), constants, max_graph_rows=max(256, engine_config.cuda_graph_max_bs or 0), max_extend_tokens=engine_config.max_extend_tokens, 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 cea6269357..fa8740897d 100644 --- a/python/freetoken/models/qwen4_exp/ple_disk.py +++ b/python/freetoken/models/qwen4_exp/ple_disk.py @@ -1,4 +1,4 @@ -"""Disk-backed PLE table (--ple-backend disk): the C++ store hashes n-gram windows and batch-reads rows from the checkpoint's fp8 shard tensors into pinned staging; the captured ``lookup`` is a fixed-shape H2D copy + dequant. +"""Disk-backed PLE table (--ple-backend disk): the C++ store hashes n-gram windows and batch-reads rows from the checkpoint's fp8 or bf16 shard tensors into pinned staging; the captured ``lookup`` is a fixed-shape H2D copy + dequant. Hash windows are pure functions of ``req.input_ids`` + ``device_len`` (prefix hits, restores and COW forks need no bookkeeping); the decode input token lives device-side under overlap scheduling and is read back here. """ @@ -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 @@ -22,7 +26,7 @@ from .weight import ( _PLE_SCALE_SUFFIX, _PLE_SHARD_RE, - _PLE_ST_DTYPE, + _PLE_ST_DTYPES, _ple_table_files, _safetensors_header, ) @@ -50,6 +54,10 @@ class PleRowSource: row_bytes: int row_stride: int scale: float + # packed ggml rows (a GGUF table) of ``row_values`` values each; None for float rows of ``row_dtype`` + ggml_type: int | None = None + row_values: int | None = None + row_dtype: torch.dtype = torch.float8_e4m3fn @property def total_rows(self) -> int: @@ -59,6 +67,7 @@ def total_rows(self) -> int: def source_from_safetensors(folder: str) -> PleRowSource: """Map the checkpoint's ``ngram_embedding.shard_`` tensors in place: one extent per shard, no copy.""" rows = cols = 0 + dtype: torch.dtype | None = None scale: torch.Tensor | None = None paths: list[str] = [] path_idx: dict[str, int] = {} @@ -75,8 +84,10 @@ def source_from_safetensors(folder: str) -> PleRowSource: match = _PLE_SHARD_RE.search(key) if match is None: continue - if meta["dtype"] != _PLE_ST_DTYPE: - raise ValueError(f"PLE shard {key} has dtype {meta['dtype']}, expected {_PLE_ST_DTYPE}") + shard_dtype = _PLE_ST_DTYPES.get(meta["dtype"]) + if shard_dtype is None or dtype not in (None, shard_dtype): + raise ValueError(f"PLE shard {key} has dtype {meta['dtype']}, expected one of {sorted(_PLE_ST_DTYPES)} for every shard") + dtype = shard_dtype if rows and tuple(meta["shape"]) != (rows, cols): raise ValueError(f"PLE shard {key} is {meta['shape']}, expected {[rows, cols]}") rows, cols = meta["shape"] @@ -90,19 +101,94 @@ def source_from_safetensors(folder: str) -> PleRowSource: if sorted(shards) != list(range(len(shards))) or not shards: raise ValueError(f"PLE shard indices are not contiguous 0..N-1: {sorted(shards)[:8]}") if scale is None: - raise ValueError("PLE table has no weight_scale") + if dtype is torch.float8_e4m3fn: + raise ValueError("PLE table has no weight_scale") + scale = torch.tensor(1.0) order = [shards[i] for i in range(len(shards))] - return PleRowSource(paths, [f for f, _ in order], [b for _, b in order], rows, cols, cols, float(scale)) + row_bytes = cols * dtype.itemsize + return PleRowSource(paths, [f for f, _ in order], [b for _, b in order], rows, row_bytes, row_bytes, float(scale), + row_values=cols, row_dtype=dtype) -def resolve_row_source(folder: str) -> PleRowSource: +def source_from_gguf(model_path: str) -> PleRowSource: + """The GGUF ``per_layer_token_embd`` table in place: one extent of packed ggml rows.""" + from freetoken.layers.quantization.ggml_types import row_bytes + + from .gguf import ple_table_info + + info = ple_table_info(model_path) + values, rows = info.ne + nbytes = row_bytes(info.ggml_type, values) + return PleRowSource([info.path], [0], [info.offset], rows, nbytes, nbytes, 1.0, info.ggml_type, values) + + +def resolve_row_source(model_path: str) -> PleRowSource: """Pick the row source for a checkpoint; the seam where a repacked format would plug in.""" - return source_from_safetensors(folder) + from freetoken.models.gguf.reader import gguf_config_source + + if (gguf_src := gguf_config_source(model_path)) is not None: + return source_from_gguf(gguf_src) + from freetoken.utils import download_hf_weight + + return source_from_safetensors(download_hf_weight(model_path)) + + +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, @@ -112,10 +198,15 @@ def __init__( max_extend_tokens: int = 8192, dtype: torch.dtype = torch.bfloat16, ) -> None: - from freetoken.kernel import _ple_store + from freetoken.kernel.row_store import PleStore + from freetoken.layers.quantization.ggml_types import float_dtype self.num_rows = source.total_rows - self.head_dim = source.row_bytes # fp8: one byte per element + self.head_dim = source.row_values or source.row_bytes + self._ggml_type = source.ggml_type + self._row_bytes = source.row_bytes + # float rows (fp8, or F32 / F16 / BF16 in a GGUF) are only cast; quantized ggml rows go through ggml + self._row_dtype = source.row_dtype if source.ggml_type is None else float_dtype(source.ggml_type) self.dtype = dtype self.heads = int(hash_constants["num_ngram_heads"]) self.scale = source.scale @@ -128,7 +219,7 @@ def __init__( raise ValueError( f"PLE row source holds {source.total_rows} rows but the hash addresses {need}; incomplete checkpoint?" ) - self._store = _ple_store.PleStore( + self._store = PleStore( paths=list(source.paths), extent_file=list(source.extent_file), extent_base=list(source.extent_base), @@ -142,7 +233,7 @@ def __init__( use_io_uring=os.getenv(_IO_URING_ENV, "1") != "0", ) self._device = torch.device("cuda", torch.cuda.current_device()) - self._token_bytes = self.heads * self.head_dim + self._token_bytes = self.heads * self._row_bytes # allocated up front: pinned alloc inside stream capture is illegal; one replay consumes it at a time self._graph_pinned = alloc_pinned_tensor(max_graph_rows * self._token_bytes, dtype=torch.uint8) self._graph_pinned.zero_() # padded decode lanes read whatever sits here @@ -151,43 +242,80 @@ def __init__( max_graph_rows * self._token_bytes, dtype=torch.uint8, device=self._device ) eager_bytes = max_extend_tokens * self._token_bytes - self._eager_pinned = alloc_pinned_tensor(eager_bytes, dtype=torch.uint8) - self._eager_pinned.zero_() # the warmup prefill stages nothing and reads whatever sits here + # 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 + # that last read it has run (``_eager_read[i]``, recorded right after that copy) + self._eager_pinned = [alloc_pinned_tensor(eager_bytes, dtype=torch.uint8) for _ in range(2)] + for buf in self._eager_pinned: + buf.zero_() # the warmup prefill stages nothing and reads whatever sits here + self._eager_read = [torch.cuda.Event(), torch.cuda.Event()] + self._eager_slot = 0 self._eager_dev = torch.empty(eager_bytes, dtype=torch.uint8, device=self._device) # probe picks flag-sync (graph WAITs at the consume, host fills then signals) or launch-gating - self._wait_sync = self._probe_wait_sync(os.getenv(_SYNC_ENV, "auto")) - # one flag for all graphs: the readback event orders a fill after the previous graph, so signals never overlap + 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: 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}") - def _probe_wait_sync(self, mode: str) -> bool: - from freetoken.kernel import _ple_store - - if mode == "gate": - return False - scratch = alloc_pinned_tensor(1, dtype=torch.int64) - scratch.zero_() - stream = torch.cuda.current_stream(self._device) - ok = ( - _ple_store.memop_write(stream.cuda_stream, scratch.data_ptr(), 7) == 0 - and _ple_store.memop_wait_geq(stream.cuda_stream, scratch.data_ptr(), 7) == 0 - ) - if ok: - stream.synchronize() - ok = int(scratch[0]) == 7 - if mode == "wait" and not ok: - raise RuntimeError("FREETOKEN_PLE_SYNC=wait but stream memops are unavailable") - return ok - # ---------------- 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.""" - pinned = self._graph_pinned if graph else self._eager_pinned + if graph: + # 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 + self._eager_read[self._eager_slot].synchronize() + pinned = self._eager_pinned[self._eager_slot] offset = 0 for run in runs: self._store.stage(run.data_ptr(), run.numel() - 2, pinned.data_ptr() + offset * self._token_bytes) @@ -204,36 +332,49 @@ 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 import _ple_store + from freetoken.kernel.row_store import signal # unblock the stream before surfacing; the step's output is discarded - _ple_store.signal_flag(self._flag.data_ptr()) + 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(( @@ -242,35 +383,44 @@ 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() + """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) ---------------- def lookup(self, row_ids: torch.Tensor, out: torch.Tensor | None = None) -> torch.Tensor: rows = row_ids.shape[0] + stream = torch.cuda.current_stream(self._device) capturing = torch.cuda.is_current_stream_capturing() if capturing and self._wait_sync: - from freetoken.kernel import _ple_store + from freetoken.kernel.row_store import wait_reset - _ple_store.memop_wait_reset( - torch.cuda.current_stream(self._device).cuda_stream, self._flag.data_ptr() - ) + wait_reset(stream, self._flag) pinned, dev = ( - (self._graph_pinned, self._graph_dev) if capturing else (self._eager_pinned, self._eager_dev) + (self._graph_pinned, self._graph_dev) if capturing else (self._eager_pinned[self._eager_slot], self._eager_dev) ) nbytes = rows * self._token_bytes dev[:nbytes].copy_(pinned[:nbytes], non_blocking=True) - values = dev[:nbytes].view(torch.float8_e4m3fn).to(self.dtype) + if not capturing: + self._eager_read[self._eager_slot].record(stream) + if self._row_dtype is not None: + # copy=True: rows already in self.dtype must not alias the staging buffer the next fill reuses + values = dev[:nbytes].view(self._row_dtype).to(self.dtype, copy=True) + else: + from freetoken.kernel import ggml + + values = ggml.dequantize(dev[:nbytes].view(-1, self._row_bytes), self._ggml_type, self.head_dim, self.dtype) if self.scale != 1.0: values = values * self.scale values = values.view(*row_ids.shape[:-1], -1) diff --git a/python/freetoken/models/qwen4_exp/weight.py b/python/freetoken/models/qwen4_exp/weight.py index 691651d4cc..9b4247ce13 100644 --- a/python/freetoken/models/qwen4_exp/weight.py +++ b/python/freetoken/models/qwen4_exp/weight.py @@ -1,10 +1,10 @@ -"""Qwen3.8-Flash-Next checkpoint reader (the NVFP4 and the official block-fp8 releases). +"""Qwen3.8-Flash-Next checkpoint reader (bf16, block-fp8, ModelOpt and llm-compressor exports). Three separate paths, because the checkpoint's three weight classes live in different places: -* :func:`iter_weights` -- every dense (non-expert) tensor, with the ``model.language_model.`` prefix stripped and fused where the model expects one buffer. See ``_DenseFuser``. -* :func:`load_ple_table` -- the 47.7 GiB FP8 n-gram table, 128 checkpoint shards concatenated into one pinned :class:`HostBank`. -* :func:`nvfp4_expert_spec` -- how the routed NVFP4 experts are named, for the offload cache's expert reader. +* :func:`iter_weights` -- every dense (non-expert) tensor, with the ``model.language_model.`` prefix stripped and fused where the model expects one buffer. See ``_DenseReader``. +* :func:`load_ple_table` -- the n-gram table (47.7 GiB fp8 or 95.4 GiB bf16), 128 checkpoint shards concatenated into one pinned :class:`HostBank`. +* :func:`nvfp4_expert_spec` -- how the routed NVFP4 experts are named under the checkpoint's dialect, for the offload cache's expert reader. Dropped: ``mtp.*`` (speculative head, including its stacked ``mtp.layers.0.mlp.experts.*``); ``model.visual.*`` is kept only when the model built the tower. """ @@ -24,33 +24,41 @@ from freetoken.models.qwen3_vl.weight import rename_vl_prefix from freetoken.models.config import VISION_KEY_PREFIXES -from freetoken.models.loader import drop_page_cache, iter_weight_files +from freetoken.models.loader import drop_page_cache, iter_weight_files, shard_tensor from freetoken.models.nvfp4_banks import ( Nvfp4ExpertSourceSpec, ) -from freetoken.layers.quantization import get_quant_config +from freetoken.layers.quantization import QuantConfig, QuantKind, QuantScheme, get_quant_config from freetoken.models.register import get_model_spec from freetoken.moe.host_banks import HostBank, read_range_into -from freetoken.utils import cached_load_hf_config, download_hf_weight +from freetoken.utils import cached_load_hf_config, div_even, download_hf_weight, init_logger from freetoken.utils.progress import byte_bar from tqdm import tqdm -# Routed NVFP4 experts (nvidia modelopt layout): per-expert, un-fused. Matched against the RAW -# weight_map key in nvfp4_banks. The ``model.language_model.`` anchor excludes the MTP head's -# stacked ``mtp.layers.N.mlp.experts.*`` tensors. -_EXPERT_KEY_RE = re.compile( +logger = init_logger(__name__) + +# Routed NVFP4 experts: per-expert, un-fused. Matched against the RAW weight_map key in nvfp4_banks. +# The ``model.language_model.`` anchor excludes the MTP head's stacked ``mtp.layers.N.mlp.experts.*`` tensors. +_EXPERT_KEY_RE = ( r"^model\.language_model\.layers\.(?P\d+)\.mlp\.experts\.(?P\d+)\." - r"(?Pgate_proj|up_proj|down_proj)\.(?Pweight|weight_scale|weight_scale_2)$" + r"(?Pgate_proj|up_proj|down_proj)\.(?P{kinds})$" ) _EXPERT_RE = re.compile(r"\.mlp\.experts\.\d+\.") -_NVFP4_SOURCE_SPEC = Nvfp4ExpertSourceSpec( - key_pattern=_EXPERT_KEY_RE, - proj_to_role={"gate_proj": "gate", "up_proj": "up", "down_proj": "down"}, - layer_to_bank=lambda layer, config: layer, # every layer is MoE - desc="Qwen3.8-Flash-Next NVFP4 experts", -) -# Per-tensor modelopt quant scales; consumed with their ``.weight`` (experts) or unused. -_SCALE_SUFFIXES = (".weight_scale", ".weight_scale_2", ".input_scale") +# role -> the expert bank reader's canonical (ModelOpt) tensor kind +_BANK_KINDS = {"weight": "weight", "weight_scale": "weight_scale", "weight_global": "weight_scale_2"} + +# leaves the model builds as quant-aware Linear layers: their tensors are read under the QuantConfig +_LINEAR_LEAVES = frozenset({ + "q_proj", "k_proj", "v_proj", "o_proj", "index_qk_proj", + "in_proj_qkv", "in_proj_z", "in_proj_b", "in_proj_a", "out_proj", + "gate_proj", "up_proj", "down_proj", + "input_mix_weight_down", "block_inject_weight", "input_mix_weight_up", + "key_proj", "value_proj", "lm_head", +}) +# activation scales of modules whose scheme carries no input_scale role +_DROPPED_SUFFIXES = frozenset({"input_scale", "input_global_scale"}) +_ELEM_DTYPES = {"e4m3": torch.float8_e4m3fn, "e2m1": torch.uint8} +_QUANT_DTYPES = (torch.float8_e4m3fn, torch.float8_e5m2, torch.uint8, torch.int8) # The n-gram table itself: too big for the dense state dict, loaded by load_ple_table. _PLE_TABLE_INFIX = ".ple.ple_embedding.ngram_embedding." @@ -79,9 +87,6 @@ # The top-level hyper_connection_mixer has no injection and never fuses. _PAD_TO = {"input_mix_weight_down_block_inject": 16} _HC_WITH_INJECT = (".attn_hyper_connection", ".mlp_hyper_connection") -_KIND_SUFFIXES = (".weight_scale_inv", ".weight") -_FP8_DTYPES = (torch.float8_e4m3fn, torch.float8_e5m2) -_ELEM_DTYPES = {"e4m3": torch.float8_e4m3fn} def _rename(raw_name: str) -> str | None: @@ -92,100 +97,258 @@ def _rename(raw_name: str) -> str | None: return None # n-gram table + its scale: load_ple_table if _EXPERT_RE.search(raw_name): return None # routed experts: offload source banks - if raw_name.endswith(_SCALE_SUFFIXES): + # static KV-cache scales of the quantizers; the KV cache runs in the engine's dtype + if raw_name.endswith((".k_scale", ".v_scale", ".q_scale", ".prob_scale")): return None return rename_vl_prefix(raw_name) -def _split_kind(name: str) -> tuple[str, str]: - """``name`` -> ``(module, kind)``; kind is "" for tensors that are neither a weight nor a block scale.""" - for suffix in _KIND_SUFFIXES: - if name.endswith(suffix): - return name[: -len(suffix)], suffix - return name, "" +def _per_row_scale(scale: torch.Tensor, rows: int) -> torch.Tensor: + """Per-tensor scalar or per-channel ``[rows, 1]`` scale -> fp32 ``[rows]``; any other count is refused rather than broadcast onto the wrong rows.""" + flat = scale.reshape(-1).to(torch.float32) + if flat.numel() == 1: + return flat.expand(rows).contiguous() + if flat.numel() != rows: + raise ValueError(f"scale has {flat.numel()} elements for a weight with {rows} output rows (shape {tuple(scale.shape)}); expected 1 or {rows}") + return flat.contiguous() -class _DenseFuser: - """Concatenates checkpoint projection parts into the model's merged buffers, per kind (weight / block scale). +class _DenseReader: + """Routes each Linear tensor to the buffer its module's scheme declares; packed projections are concatenated per role once every part is in. - The part table is the family's packed_modules_mapping. The QuantConfig picks the GDN in_proj layout and validates each part against the scheme the model built its buffer from. + The part table is the family's packed_modules_mapping. The QuantConfig picks the GDN in_proj layout and, through its dialect's storage names, which checkpoint tensor fills which role. """ - def __init__(self, quant, packed: tuple[tuple[str, tuple[str, ...]], ...]) -> None: + def __init__(self, quant: QuantConfig | None, packed: tuple[tuple[str, tuple[str, ...]], ...]) -> None: self.quant = quant self.groups = {fused: parts for fused, parts in packed if fused != "experts"} # experts: bank reader self.by_part: dict[str, list[tuple[str, int]]] = {} for fused, parts in self.groups.items(): for idx, part in enumerate(parts): self.by_part.setdefault(part, []).append((fused, idx)) - self.buf: dict[tuple[str, str], dict[int, torch.Tensor]] = {} + # target module -> (part count, {part: {role: tensor}}, {part: the roles its module stores}, scheme) + self.pending: dict[str, tuple[int, dict[int, dict[str, torch.Tensor]], dict[int, set[str]], QuantScheme | None]] = {} - def scheme(self, module: str): + def scheme(self, module: str) -> QuantScheme | None: return None if self.quant is None else self.quant.scheme_for(module) - def _target(self, parent: str, leaf: str) -> tuple[str, int] | None: + def target(self, module: str) -> tuple[str, int, int]: + """``(fused module, part index, part count)``; a standalone linear is its own single-part target.""" + parent, _, leaf = module.rpartition(".") candidates = self.by_part.get(leaf) if not candidates: - return None + return module, 0, 1 if len(candidates) > 1: - # GDN: quantized checkpoints split qkv|z from the bf16 b|a; same test as gdn.py + # GDN: quantized checkpoints split qkv|z from b|a; same test as gdn.py split = self.scheme(f"{parent}.in_proj_qkvz") is not None keep = {"in_proj_qkvz", "in_proj_ba"} if split else {"in_proj"} candidates = [c for c in candidates if c[0] in keep] - if not candidates: - raise ValueError(f"{parent}.{leaf}: no merged projection for the {'split' if split else 'fused'} GDN layout") fused, idx = candidates[0] if fused in _PAD_TO and not parent.endswith(_HC_WITH_INJECT): - return None - return f"{parent}.{fused}", idx + return module, 0, 1 + return f"{parent}.{fused}", idx, len(self.groups[fused]) - def check(self, module: str, name: str, tensor: torch.Tensor) -> None: - """``tensor`` (checkpoint key ``name``) must match the scheme the model built ``module`` from.""" + def add(self, name: str, tensor: torch.Tensor) -> list[tuple[str, torch.Tensor]] | None: + """Take one tensor; the emitted ``[(name, tensor)]`` once its target module is complete, ``[]`` before, None if ``name`` is not a Linear's tensor.""" + module, _, suffix = name.rpartition(".") + if module.rpartition(".")[2] not in _LINEAR_LEAVES: + return None scheme = self.scheme(module) - if name.endswith(".weight_scale_inv"): - if scheme is None or not scheme.has("weight_scale_inv"): - raise ValueError(f"{name}: {module} has no block scale in the checkpoint's quant config ({scheme})") - return - is_fp8 = tensor.dtype in _FP8_DTYPES + roles = {"weight": "weight"} if scheme is None else {n: r for r, e in self.quant.storage(scheme).items() for n in e.names} + role = roles.get(suffix) + if role is None: + if suffix in _DROPPED_SUFFIXES: + return [] + raise ValueError(f"{name}: the checkpoint's quant config declares {module} {scheme or 'unquantized'}, stored as {sorted(roles)}") + if role == "weight": + self._check_weight(name, module, scheme, tensor) + target, idx, count = self.target(module) + _, parts, expected, _ = self.pending.setdefault(target, (count, {}, {}, scheme)) + part = parts.setdefault(idx, {}) + if role in part: + raise ValueError(f"{name}: {module} already has its {role} from another checkpoint tensor") + part[role] = tensor + expected[idx] = set(roles.values()) + if scheme is not None and set(part) == expected[idx]: + parts[idx] = self._finish(module, scheme, part) + if len(parts) < count or any(set(parts[i]) != expected[i] for i in parts): + return [] + del self.pending[target] + return self._emit(target, [parts[i] for i in range(count)], scheme) + + def missing(self) -> list[str]: + """One line per incomplete module: the roles its parts still lack.""" + lines = [] + for target, (count, parts, expected, _scheme) in sorted(self.pending.items()): + lacking = sorted(set().union(*(expected[i] - set(parts[i]) for i in parts))) + if len(parts) < count: + lacking.append(f"{count - len(parts)} of {count} fused parts") + lines.append(f"{target}: missing {lacking}") + return lines + + @staticmethod + def _check_weight(name: str, module: str, scheme: QuantScheme | None, tensor: torch.Tensor) -> None: if scheme is None: - if is_fp8: + if tensor.dtype in _QUANT_DTYPES: raise ValueError(f"{name} is {tensor.dtype} but the checkpoint's quant config declares {module} unquantized") return - expected = _ELEM_DTYPES.get(scheme.weight.elem) - if expected is not None and tensor.dtype is not expected: + elem = _ELEM_DTYPES.get(scheme.weight.elem) + if elem is None: + raise ValueError(f"{name}: qwen4_exp has no dense reader for {scheme}") + if tensor.dtype is not elem: raise ValueError(f"{name} is {tensor.dtype} but the checkpoint's quant config declares {module} {scheme}") - rows, cols = (scheme.weight.group or (1, 1)) - if rows > 1 and tensor.shape[0] % rows or cols > 1 and tensor.shape[1] % cols: - raise ValueError(f"{name}: {tuple(tensor.shape)} is not a multiple of the {rows}x{cols} scale block of {module}") - - def check_unfused(self, name: str, tensor: torch.Tensor) -> None: - module, kind = _split_kind(name) - if kind == ".weight_scale_inv" or (kind == ".weight" and tensor.dtype in _FP8_DTYPES): - self.check(module, name, tensor) - - def fuse(self, name: str, tensor: torch.Tensor) -> list[tuple[str, torch.Tensor]] | None: - """Buffer a part; return the merged ``[(name, tensor)]`` once its kind is complete, ``[]`` while incomplete, ``None`` if ``name`` is not a part.""" - module, kind = _split_kind(name) - if not kind: - return None - parent, _, leaf = module.rpartition(".") - hit = self._target(parent, leaf) - if hit is None: - return None - fused, idx = hit - self.check(fused, name, tensor) - slots = self.buf.setdefault((fused, kind), {}) - slots[idx] = tensor - parts = self.groups[fused.rpartition(".")[2]] - if len(slots) < len(parts): - return [] - del self.buf[(fused, kind)] - rows = [slots[i] for i in range(len(parts))] - pad_to = _PAD_TO.get(fused.rpartition(".")[2], 0) if kind == ".weight" else 0 - pad = (-sum(t.shape[0] for t in rows)) % pad_to if pad_to else 0 - if pad: - rows.append(torch.zeros(pad, *rows[0].shape[1:], dtype=rows[0].dtype, device=rows[0].device)) - return [(fused + kind, torch.cat(rows, dim=0))] + rows, cols = tensor.shape[0], tensor.shape[1] * (2 if scheme.weight.elem == "e2m1" else 1) + block_rows, block_cols = scheme.weight.group or (1, 1) + if block_rows > 1 and rows % block_rows or block_cols > 1 and cols % block_cols: + raise ValueError(f"{name}: {rows}x{cols} is not a multiple of the {block_rows}x{block_cols} scale block of {module}") + + def _finish(self, module: str, scheme: QuantScheme, part: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: + """Validate one complete part against ``scheme`` and put its scales in the layer's form.""" + storage = self.quant.storage(scheme) + part = {role: 1.0 / t.to(torch.float32) if storage[role].reciprocal else t for role, t in part.items()} + weight = part["weight"] + rows, cols = weight.shape[0], weight.shape[1] * (2 if scheme.weight.elem == "e2m1" else 1) + block_rows, block_cols = scheme.weight.group or (1, 1) + scale_role = "weight_scale_inv" if "weight_scale_inv" in part else "weight_scale" + out = dict(part) + if block_cols < 0: + try: + out[scale_role] = _per_row_scale(part[scale_role], rows) + except ValueError as exc: + raise ValueError(f"{module}.{scale_role}: {exc}") from None + else: + want = (rows // block_rows, cols // block_cols) + if tuple(part[scale_role].shape) != want: + raise ValueError(f"{module}: {scale_role} is {tuple(part[scale_role].shape)}, expected {want} for {scheme}") + if scheme.weight.scale == "e4m3" and part[scale_role].dtype is not torch.float8_e4m3fn: + raise ValueError(f"{module}: {scale_role} is {part[scale_role].dtype} but {scheme} stores e4m3 scales") + if "weight_global" in part: + g = part["weight_global"].reshape(-1).to(torch.float32) + if g.numel() != 1: + raise ValueError(f"{module}: weight_global has {g.numel()} elements, expected one per-tensor scale") + out["weight_global"] = g.expand(rows).contiguous() + if "input_scale" in part: + out["input_scale"] = part["input_scale"].reshape(()).to(torch.float32) + return out + + @staticmethod + def _emit(target: str, parts: list[dict[str, torch.Tensor]], scheme: QuantScheme | None) -> list[tuple[str, torch.Tensor]]: + pad_to = _PAD_TO.get(target.rpartition(".")[2], 0) + if pad_to and scheme is not None: + raise ValueError(f"{target}: qwen4_exp pads this fusion to {pad_to} rows and reads it bf16 only, the checkpoint's quant config declares {scheme}") + out = [] + for role in parts[0]: + tensors = [part[role] for part in parts] + if role == "input_scale": + # fused parts read the same activation, so one range covers them: max is exact then and safe if they drift + out.append((f"{target}.{role}", torch.stack(tensors).max())) + continue + pad = (-sum(t.shape[0] for t in tensors)) % pad_to if pad_to else 0 + if pad: + tensors.append(torch.zeros(pad, *tensors[0].shape[1:], dtype=tensors[0].dtype, device=tensors[0].device)) + out.append((f"{target}.{role}", tensors[0] if len(tensors) == 1 else torch.cat(tensors, dim=0))) + return out + + +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( @@ -199,22 +362,48 @@ def iter_weights( """Yield the dense (non-expert) weights, prefix-stripped and fused to the model's buffers. Keys keep the checkpoint's module names below the stripped prefix, so the emitted set is the model's state dict minus the routed experts. - A dense projection is bf16 or 128x128 block-fp8 (``.weight`` e4m3 + ``.weight_scale_inv``) as the checkpoint's QuantConfig says: the official releases skip everything but the routed experts, the community NVFP4-FP8 requants quantize the attention / GDN projections. - 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``. + Each Linear is read under the scheme the checkpoint's QuantConfig gives it, the same answer the model built its buffers from, so bf16, per-channel / per-tensor fp8, block-fp8, MXFP8 and NVFP4 projections in any mix land as the model's state dict. + Fusions, per role: attention q|k|v -> ``qkv_proj``; GDN ``in_proj_{qkv,z,b,a}`` -> ``in_proj``, or ``in_proj_qkvz`` + ``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) + reader = _DenseReader(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(): @@ -224,14 +413,22 @@ def iter_weights( if not include_vision and name.startswith(VISION_KEY_PREFIXES): continue tensor = f.get_tensor(raw_name) - fused = fuser.fuse(name, tensor) - if fused is None: - fuser.check_unfused(name, tensor) - yield name, tensor + emitted = None if name.startswith(VISION_KEY_PREFIXES) else reader.add(name, tensor) + if emitted is None: + yield from emit(name, tensor) else: - yield from fused + for fused_name, fused_tensor in emitted: + 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 reader.pending: + lines = reader.missing() + shown = "\n ".join(lines[:8]) + (f"\n ... {len(lines) - 8} more" if len(lines) > 8 else "") + raise ValueError(f"checkpoint is missing tensors the quant config declares for {len(lines)} modules:\n {shown}") + 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]]: @@ -251,18 +448,18 @@ def iter_vision_weights(model_path: str, device: torch.device) -> Iterator[tuple @dataclass(frozen=True) class PleTable: - """The filled n-gram table: one pinned host bank plus the checkpoint's per-tensor FP8 scale.""" + """The filled n-gram table: one pinned host bank plus the checkpoint's per-tensor scale (1 for a bf16 table).""" bank: HostBank weight_scale: torch.Tensor # scalar, checkpoint dtype (bf16) @property def tensor(self) -> torch.Tensor: - """``[total_rows, ngram_head_dim]`` float8_e4m3fn view of the bank.""" + """``[total_rows, ngram_head_dim]`` view of the bank in the stored dtype.""" return self.bank.tensor -_PLE_ST_DTYPE = "F8_E4M3" +_PLE_ST_DTYPES = {"F8_E4M3": torch.float8_e4m3fn, "BF16": torch.bfloat16} def _safetensors_header(path: str) -> tuple[dict, int]: @@ -282,6 +479,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. @@ -322,12 +545,14 @@ def load_ple_table(model_path: str, qwen4_args, *, pin: bool = True, The checkpoint splits the table into ``split_ngram_parts`` equal row blocks named by shard index and scattered over the ``model-plefp8-*`` shards in header (lexicographic) order, so the bank is filled shard by shard at ``shard_index * rows_per_shard``. Each read is O_DIRECT: the - table is ~47.7 GiB and must not also sit in the page cache while the bank holds the same bytes. + table is ~47.7 GiB (fp8) and must not also sit in the page cache while the bank holds the same bytes. + An fp8 table needs the checkpoint's scalar ``weight_scale``; a bf16 table without one is unscaled. """ folder = download_hf_weight(model_path) parts: dict[int, tuple[str, int, int]] = {} # shard index -> (path, file offset, bytes) scale: torch.Tensor | None = None rows = cols = 0 + dtype: torch.dtype | None = None for path in _ple_table_files(folder): header, base = _safetensors_header(path) for key, meta in header.items(): @@ -340,8 +565,10 @@ def load_ple_table(model_path: str, qwen4_args, *, pin: bool = True, match = _PLE_SHARD_RE.search(key) if match is None: continue - if meta["dtype"] != _PLE_ST_DTYPE: - raise ValueError(f"PLE table shard {key} has unsupported dtype {meta['dtype']}") + shard_dtype = _PLE_ST_DTYPES.get(meta["dtype"]) + if shard_dtype is None or dtype not in (None, shard_dtype): + raise ValueError(f"PLE table shard {key} has dtype {meta['dtype']}, expected one of {sorted(_PLE_ST_DTYPES)} for every shard") + dtype = shard_dtype shape = meta["shape"] if rows and tuple(shape) != (rows, cols): raise ValueError(f"PLE table shard {key} is {shape}, expected {[rows, cols]}") @@ -357,10 +584,12 @@ def load_ple_table(model_path: str, qwen4_args, *, pin: bool = True, if cols != qwen4_args.ngram_head_dim: raise ValueError(f"PLE table row is {cols} wide, config says {qwen4_args.ngram_head_dim}") if scale is None: - raise ValueError("PLE table has no weight_scale") + if dtype is torch.float8_e4m3fn: + raise ValueError("PLE table has no weight_scale") + scale = torch.tensor(1.0, dtype=torch.bfloat16) - bank = HostBank((expected * rows, cols), torch.float8_e4m3fn) - shard_bytes = rows * cols + bank = HostBank((expected * rows, cols), dtype) + shard_bytes = rows * cols * dtype.itemsize bar = byte_bar(expected * shard_bytes, "Loading PLE table") try: buf = bank.memoryview() @@ -382,8 +611,19 @@ def load_ple_table(model_path: str, qwen4_args, *, pin: bool = True, # ====================================================================================== -def nvfp4_expert_spec(model_path: str, config): - return _NVFP4_SOURCE_SPEC +def nvfp4_expert_spec(model_path: str, config) -> Nvfp4ExpertSourceSpec: + """The per-expert NVFP4 layout under the checkpoint's dialect names (ModelOpt or llm-compressor).""" + quant = get_quant_config() + stored = quant.stored_tensors(QuantKind.NVFP4) + kind_map = {stored[role].name: kind for role, kind in _BANK_KINDS.items()} + return Nvfp4ExpertSourceSpec( + key_pattern=re.compile(_EXPERT_KEY_RE.format(kinds="|".join(map(re.escape, kind_map)))), + proj_to_role={"gate_proj": "gate", "up_proj": "up", "down_proj": "down"}, + layer_to_bank=lambda layer, config: layer, # every layer is MoE + desc=f"Qwen3.8-Flash-Next NVFP4 experts ({quant.dialect})", + kind_map=kind_map, + global_reciprocal=stored["weight_global"].reciprocal, + ) __all__ = [ diff --git a/python/freetoken/models/register.py b/python/freetoken/models/register.py index 2a4ced10af..4d268ffb50 100644 --- a/python/freetoken/models/register.py +++ b/python/freetoken/models/register.py @@ -48,6 +48,8 @@ class ModelSpec: # per-expert checkpoints: probe the first expert's projections for the experts container _EXPERTS_PACKED = (("experts", ("experts.0.gate_proj", "experts.0.up_proj", "experts.0.down_proj")),) _EXPERTS_W123_PACKED = (("experts", ("experts.0.w1", "experts.0.w2", "experts.0.w3")),) +# the GGUF scheme reads expert types in (gate, up, down) order +_EXPERTS_W132_PACKED = (("experts", ("experts.0.w1", "experts.0.w3", "experts.0.w2")),) # pre-stacked expert checkpoints (Qwen3-VL-MoE): the container's own two tensors are the probe _STACKED_EXPERTS_PACKED = (("experts", ("experts.gate_up_proj", "experts.down_proj")),) _QWEN3_5_PACKED = _DENSE_PACKED + ( @@ -61,6 +63,8 @@ class ModelSpec: _QWEN4_EXP_PACKED = _QWEN3_5_PACKED + ( ("input_mix_weight_down_block_inject", ("input_mix_weight_down", "block_inject_weight")), ) +# the llama.cpp converter splits the indexer projection into q and k +_QWEN4_EXP_GGUF_PACKED = _QWEN4_EXP_PACKED + (("index_qk_proj", ("index_q_proj", "index_k_proj")),) # GLM-5.3-Flash's KDA fuses its six input projections. _GLM5_NEXT_PACKED = _EXPERTS_PACKED + ( ("in_proj", ("q_proj", "k_proj", "v_proj", "b_proj", "f_a_proj", "g_a_proj")), @@ -82,8 +86,13 @@ class ModelSpec: _MINIMAX_M3_PACKED = _DENSE_PACKED + ( ("index_qk_proj", ("index_q_proj", "index_k_proj")), ) + _EXPERTS_W123_PACKED +_MINIMAX_M3_GGUF_PACKED = _DENSE_PACKED + (("index_qk_proj", ("index_q_proj", "index_k_proj")),) + _EXPERTS_W132_PACKED _MINIMAX_M3_PROCESSOR = "freetoken.mm.processors.minimax_m3:MiniMaxM3MMProcessor" _MINIMAX_M3_ENCODERS = (EncoderSpec("vision", "vision_config", ("image",)),) +# DeepSeek-V4's vision dims sit at the top level of the checkpoint config, next to the +# language model's rather than under a section of their own. +_DSV4_PROCESSOR = "freetoken.mm.processors.deepseek_v4:DSV4MMProcessor" +_DSV4_ENCODERS = (EncoderSpec("vision", "vision_n_layers", ("image",)),) _MODEL_REGISTRY: dict[str, ModelSpec] = { "LlamaForCausalLM": ModelSpec( @@ -91,6 +100,14 @@ class ModelSpec: "LlamaForCausalLM", packed_modules_mapping=_DENSE_PACKED, ), + # llama.cpp llama GGUF (Llama, Mistral): the llama model classes with every tensor in its packed ggml type + "LlamaGGUFForCausalLM": ModelSpec( + "freetoken.models.llama", + "LlamaForCausalLM", + parse_config="parse_gguf_config", + iter_weights="iter_gguf_weights", + packed_modules_mapping=_DENSE_PACKED, + ), "Qwen2ForCausalLM": ModelSpec( "freetoken.models.qwen2", "Qwen2ForCausalLM", @@ -106,6 +123,28 @@ class ModelSpec: "Qwen3MoeForCausalLM", packed_modules_mapping=_DENSE_PACKED + _EXPERTS_PACKED, ), + # llama.cpp qwen2 / qwen3 / qwen3moe GGUF: the same model classes with every tensor in its packed ggml type + "Qwen2GGUFForCausalLM": ModelSpec( + "freetoken.models.qwen2", + "Qwen2ForCausalLM", + parse_config="parse_gguf_config", + iter_weights="iter_gguf_weights", + packed_modules_mapping=_DENSE_PACKED, + ), + "Qwen3GGUFForCausalLM": ModelSpec( + "freetoken.models.qwen3", + "Qwen3ForCausalLM", + parse_config="parse_gguf_config", + iter_weights="iter_gguf_weights", + packed_modules_mapping=_DENSE_PACKED, + ), + "Qwen3MoeGGUFForCausalLM": ModelSpec( + "freetoken.models.qwen3_moe", + "Qwen3MoeForCausalLM", + parse_config="parse_gguf_config", + iter_weights="iter_gguf_weights", + packed_modules_mapping=_DENSE_PACKED + _EXPERTS_PACKED, + ), # Qwen3-VL: the Qwen3 text tower under model.language_model. plus the shared Qwen VL vision tower with DeepStack; the MoE variant ships its experts pre-stacked. "Qwen3VLForConditionalGeneration": ModelSpec( "freetoken.models.qwen3_vl", @@ -128,6 +167,14 @@ class ModelSpec: "MiniMaxM2ForCausalLM", packed_modules_mapping=_DENSE_PACKED + _EXPERTS_W123_PACKED, ), + # llama.cpp minimax-m2 GGUF: the same model classes with every tensor in its packed ggml type + "MiniMaxM2GGUFForCausalLM": ModelSpec( + "freetoken.models.minimax_m2", + "MiniMaxM2ForCausalLM", + parse_config="parse_gguf_config", + iter_weights="iter_gguf_weights", + packed_modules_mapping=_DENSE_PACKED + _EXPERTS_W132_PACKED, + ), # MiniMax-M3 (model_type minimax_m3_vl): multimodal wrapper config (text tower in # text_config, weights under language_model.); the CLIP-style tower under vision_tower. # serves image input. GQA + block-sparse attention (lightning indexer, top-k 128-token @@ -147,14 +194,41 @@ class ModelSpec: "MiniMaxM3ForCausalLM", packed_modules_mapping=_MINIMAX_M3_PACKED, ), + # llama.cpp minimax-m3 GGUF (text only): the same model classes with every tensor in its packed ggml type + "MiniMaxM3GGUFForCausalLM": ModelSpec( + "freetoken.models.minimax_m3", + "MiniMaxM3ForCausalLM", + parse_config="parse_gguf_config", + iter_weights="iter_gguf_weights", + packed_modules_mapping=_MINIMAX_M3_GGUF_PACKED, + ), "DeepseekV4ForCausalLM": ModelSpec( "freetoken.models.deepseek_v4", "DeepseekV4ForCausalLM", # the checkpoint has no ``model.`` root checkpoint_roots=(("model.layers", "layers"), ("model.head", "head")), packed_modules_mapping=_EXPERTS_W123_PACKED, - # 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"), + # the head, the KV compressors, the indexer's scorer and the vision tower + aligner ship bf16; + # the fp8 config has no modules_to_not_convert (patterns match checkpoint names, like V4.1) + unquantized_modules=( + "head", "*.compressor.wkv", "*.compressor.wgate", "*.indexer.weights_proj", "vision.*", "aligner.*", + ), + mm_processor=_DSV4_PROCESSOR, + encoders=_DSV4_ENCODERS, + ), + # 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", @@ -178,6 +252,25 @@ class ModelSpec: packed_modules_mapping=_QWEN3_5_PACKED, unquantized_modules=_QWEN3_5_UNQUANTIZED, ), + # llama.cpp qwen35 / qwen35moe GGUF (text only): the same model classes with every tensor in its packed ggml type + "Qwen3_5GGUFForCausalLM": ModelSpec( + "freetoken.models.qwen3_5_moe", + "Qwen3_5ForCausalLM", + parse_config="parse_gguf_config", + iter_weights="iter_gguf_weights", + checkpoint_roots=_LANGUAGE_MODEL_ROOT, + packed_modules_mapping=_QWEN3_5_PACKED, + unquantized_modules=_QWEN3_5_UNQUANTIZED, + ), + "Qwen3_5MoeGGUFForCausalLM": ModelSpec( + "freetoken.models.qwen3_5_moe", + "Qwen3_5MoeForCausalLM", + parse_config="parse_gguf_config", + iter_weights="iter_gguf_weights", + checkpoint_roots=_LANGUAGE_MODEL_ROOT, + packed_modules_mapping=_QWEN3_5_PACKED, + unquantized_modules=_QWEN3_5_UNQUANTIZED, + ), # Qwen3.8-Flash-Next (model_type qwen4_exp): multimodal wrapper config (text tower in # text_config, weights under model.language_model.). 36 GDN + 12 QSA # compressed-sparse attention layers on 4 hyper-connection residual streams, a PLE @@ -190,6 +283,15 @@ class ModelSpec: encoders=_QWEN_VL_ENCODERS, packed_modules_mapping=_QWEN4_EXP_PACKED, ), + # llama.cpp qwen4exp GGUF (text only): the same model classes with every tensor in its packed ggml type + "Qwen4ExpGGUFForCausalLM": ModelSpec( + "freetoken.models.qwen4_exp", + "Qwen4ExpForCausalLM", + parse_config="parse_gguf_config", + iter_weights="iter_gguf_weights", + checkpoint_roots=_LANGUAGE_MODEL_ROOT, + packed_modules_mapping=_QWEN4_EXP_GGUF_PACKED, + ), # Dense Qwen3.x (no "Moe" in the arch name, num_experts==0, e.g. Qwen3.6-27B). Shares the # qwen3_5_moe package: the decoder routes its MLP through the dense Qwen3_5DenseMLP and the # loader handles the compressed-tensors NVFP4 layout. @@ -212,13 +314,18 @@ class ModelSpec: "freetoken.models.muse_glimmer", "MuseGlimmerForConditionalGeneration", checkpoint_roots=_LANGUAGE_MODEL_ROOT, - packed_modules_mapping=( - ("qkvg_proj", ("q_proj", "k_proj", "v_proj", "gate_proj")), - ("gate_up_proj", ("gate_proj", "up_proj")), - ), + packed_modules_mapping=_DENSE_PACKED, mm_processor=_MUSE_GLIMMER_PROCESSOR, encoders=_MUSE_GLIMMER_ENCODERS, ), + # llama.cpp muse-glimmer GGUF (text only): the same model classes with every tensor in its packed ggml type + "MuseGlimmerGGUFForCausalLM": ModelSpec( + "freetoken.models.muse_glimmer", + "MuseGlimmerForCausalLM", + parse_config="parse_gguf_config", + iter_weights="iter_gguf_weights", + packed_modules_mapping=_DENSE_PACKED, + ), "MistralForCausalLM": ModelSpec( "freetoken.models.mistral", "MistralForCausalLM", @@ -259,12 +366,14 @@ class ModelSpec: checkpoint_segments=_GEMMA4_SEGMENTS, packed_modules_mapping=_GEMMA4_PACKED, ), - # GGUF (native Q4_0/Q6_K) gemma4: same model classes, GGUF config + weight loaders. + # llama.cpp gemma4 GGUF (text only): the same model classes with every tensor in its packed ggml type "Gemma4GGUFForCausalLM": ModelSpec( "freetoken.models.gemma4", "Gemma4ForCausalLM", parse_config="parse_gguf_config", iter_weights="iter_gguf_weights", + checkpoint_segments=_GEMMA4_SEGMENTS, + packed_modules_mapping=_GEMMA4_PACKED, ), "GptOssForCausalLM": ModelSpec( "freetoken.models.gpt_oss", @@ -276,6 +385,14 @@ class ModelSpec: "Glm4MoeForCausalLM", packed_modules_mapping=_EXPERTS_PACKED, ), + # llama.cpp glm4moe GGUF (GLM-4.5, 4.5-Air, 4.6, 4.7): the same model classes with every tensor in its packed ggml type + "Glm4MoeGGUFForCausalLM": ModelSpec( + "freetoken.models.glm4_moe", + "Glm4MoeForCausalLM", + parse_config="parse_gguf_config", + iter_weights="iter_gguf_weights", + packed_modules_mapping=_EXPERTS_PACKED, + ), # GLM-5.2 (model_type glm_moe_dsa): DeepSeek-V3.2-class MLA + DSA sparse attention # with GLM-4-style sigmoid/noaux_tc MoE routing; NVFP4 routed experts served from # the offload cache. @@ -303,6 +420,19 @@ class ModelSpec: "Glm5NextForCausalLM", packed_modules_mapping=_GLM5_NEXT_PACKED, ), + # llama.cpp glm5-next GGUF (unsloth GLM-5.3-Flash-GGUF; text only): the same model + # classes with every tensor in its packed ggml type + "Glm5NextGGUFForCausalLM": ModelSpec( + "freetoken.models.glm5_next", + "Glm5NextForCausalLM", + parse_config="parse_gguf_config", + iter_weights="iter_gguf_weights", + checkpoint_roots=_LANGUAGE_MODEL_ROOT, + packed_modules_mapping=_GLM5_NEXT_PACKED, + # the MLA kv_b ships as two per-head tensors (attn_k_b transposed + attn_v_b), quantized in + # dynamic quants (unsloth UD); the reader dequantizes and fuses them into the dense module + unquantized_modules=("*.self_attn.kv_b_proj",), + ), } @@ -319,16 +449,19 @@ def _load_attr(module_path: str, attr_name: str) -> Any: def checkpoint_quant_config(model_path: str, hf_config: Any, spec: ModelSpec): - """The checkpoint's QuantConfig under the family's naming, or None for GGUF, whose native-quant ops the shared parser does not model yet.""" - from freetoken.layers.quantization import NameMap, QuantConfig + """The checkpoint's QuantConfig under the family's naming. + A GGUF family gets the GGUF dialect over the file's tensor table, named through the ``to_gguf`` its module exports.""" + from freetoken.layers.quantization import GGUFConfig, NameMap, QuantConfig + + name_map = NameMap(roots=spec.checkpoint_roots, segments=spec.checkpoint_segments, packed=spec.packed_modules_mapping) if spec.parse_config == "parse_gguf_config": - return None - return QuantConfig.from_hf( - hf_config, - name_map=NameMap(roots=spec.checkpoint_roots, segments=spec.checkpoint_segments, packed=spec.packed_modules_mapping), - unquantized=spec.unquantized_modules, - ) + to_gguf = _load_attr(spec.module, "to_gguf") + from freetoken.models.gguf.reader import gguf_tensor_index + + types = {name: info.ggml_type for name, info in gguf_tensor_index(hf_config.model_path).items()} + return GGUFConfig(types, to_gguf, name_map=name_map, unquantized=spec.unquantized_modules) + return QuantConfig.from_hf(hf_config, name_map=name_map, unquantized=spec.unquantized_modules) def get_model_class(model_architecture: str, model_config: ModelConfig): diff --git a/python/freetoken/models/weight.py b/python/freetoken/models/weight.py index e3382b1b59..19268ac619 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) @@ -262,28 +274,9 @@ def ftw_lacks_vision(model_path: str) -> bool: return not any(name.startswith(VISION_KEY_PREFIXES) for name in ftw_tensor_names(model_path, "weight")) -def load_q4_0_moe_expert_sources( - model_path: str, - model_config, - *, - dummy: bool = False, - layer_sink=None, -) -> dict: - """Load (or fabricate, with ``dummy=True``) packed GGUF Q4_0 expert source banks. - ``layer_sink`` (converter) streams each completed layer's banks; ignored for dummy.""" - _config, spec = _spec_for_model_path(model_path) - if dummy: - builder = _model_override(spec, "dummy_q4_0_expert_sources") - assert builder is not None, "model defines no dummy_q4_0_expert_sources" - return builder(model_config) - loader = _load_attr(spec.module, "load_q4_0_expert_sources") - return loader(model_path, model_config, layer_sink=layer_sink) - - __all__ = [ "load_weight", "load_vision_weight", "ftw_lacks_vision", - "load_q4_0_moe_expert_sources", "iter_expert_tensors_parallel", ] diff --git a/python/freetoken/moe/bench_profile.py b/python/freetoken/moe/bench_profile.py index d7b720738d..c939a947f2 100644 --- a/python/freetoken/moe/bench_profile.py +++ b/python/freetoken/moe/bench_profile.py @@ -1,12 +1,8 @@ """Torch-free reader for the ``ft bench bw`` hardware profile (``benchbw/.json``). -The engine consults this at MoE-backend *auto* resolution (``engine.py``) to make the -offload-vs-hybrid choice hardware-adaptive without importing the (torch-heavy) benchmark -itself. ``benchbw.py`` writes the profile; this module only reads it. - -The join key is the expert *format*: the CPU-MoE-vs-PCIe-gather bandwidth ratio the choice -rides on is dominated by ``(format, hardware)``, not by the exact model, so a profile benched -on one workload transfers to any model with the same expert format on the same GPU. +The profile holds one row per MoE layer kind (``layer_kind_key``): the CPU MoE and PCIe gather +bandwidths benched for that kind on this GPU and host. The engine reads it to set each hybrid +layer's fetch split, without importing the (torch-heavy) benchmark. ``benchbw.py`` writes it. """ from __future__ import annotations @@ -18,17 +14,8 @@ logger = init_logger(__name__) -# Engine ``expert_quant`` (models/config.py) -> benchbw format key (offload_cache._BANK_SCHEMAS -# / benchbw._offload_bank_specs). Only the offload-family formats with a CPU MoE weight path can -# ever resolve to hybrid; anything not listed falls through unmapped and finds no profile entry -# (-> None -> offload), which is the safe default. -_QUANT_TO_BENCH_FORMAT = { - "nvfp4": "nvfp4", - "ds_fp4": "ds_fp4", - "mxfp4": "mxfp4_triton", - "bf16": "bf16", - "fp8_block": "fp8_block", -} +# Bump when the row key or the measurements change: rows of another version are not comparable. +PROFILE_VERSION = 6 def _cache_dir() -> str: @@ -101,6 +88,12 @@ def _usable_profile( return None if not isinstance(prof, dict): return None + if prof.get("version") != PROFILE_VERSION: + logger.warning( + f"benchbw profile {src} has version {prof.get('version')}, this engine reads " + f"{PROFILE_VERSION}; ignoring it, rerun `ft bench bw`" + ) + return None prof_gpu = (prof.get("gpu") or {}).get("name") if gpu_name and prof_gpu and prof_gpu != gpu_name: logger.warning( @@ -111,81 +104,34 @@ def _usable_profile( return prof -def load_backend_recommendation( - quant_format: str, - gpu_name: str | None = None, - path: str | None = None, - gpu_uuid: str | None = None, -) -> str | None: - """Bench-recommended offload-family backend for ``quant_format`` on this GPU, or ``None``. +def layer_kind_key(fmt: str, options: dict[str, int], hidden: int, inter: int, top_k: int, cpu_threads: int) -> str: + """Row key of one MoE layer kind, e.g. ``gguf[down_type=8,gate_up_type=8]:2816x704:top8:t12``. - Returns ``"hybrid"`` only when *every* benched workload sharing this expert format - recommended hybrid (CPU MoE BW > threshold x PCIe gather BW); a mixed verdict (a - near-threshold format) resolves conservatively to ``"offload"``. ``None`` means "no usable - profile" (see ``_usable_profile``) or no entry for this format. The caller keeps its own - default (offload) on ``None``. - """ - fmt = _QUANT_TO_BENCH_FORMAT.get(quant_format, quant_format) - prof = _usable_profile(gpu_name, path, gpu_uuid) - if prof is None: - return None + ``cpu_threads`` is the CPU executor's worker count: the per-step pool overhead makes the CPU bandwidth depend on it and on ``top_k``.""" + opts = ",".join(f"{k}={v}" for k, v in sorted(options.items())) + fmt = f"{fmt}[{opts}]" if opts else fmt + return f"{fmt}:{hidden}x{inter}:top{top_k}:t{cpu_threads}" - # Preferred: the per-dtype tuning verdicts (`ft bench bw --dtype`), a direct format->backend - # map -- the axis the backend pick is meant to key on. - dtypes = prof.get("dtypes") - if isinstance(dtypes, dict) and dtypes.get(fmt) in ("hybrid", "offload"): - return dtypes[fmt] - # Fallback: a per-model profile (`ft bench bw --model`). Aggregate the workloads sharing this - # format -- unanimous hybrid -> hybrid; any offload (a near-threshold split) -> offload. - workloads = prof.get("workloads") - if not isinstance(workloads, dict): - return None - picks = [ - entry["recommended"] - for wl in workloads.values() - if isinstance(wl, dict) - for entry in [(wl.get("kernels") or {}).get(fmt)] - if isinstance(entry, dict) and entry.get("recommended") - ] - if not picks: +def fetch_fraction(row: dict | None) -> float | None: + """The hybrid fetch split of a row: of a decode step's expert misses, the fraction to fetch over PCIe so the fetch and the CPU compute finish together. + + From the overlapped pair (both sides measured while running concurrently): pcie_ov / (pcie_ov + cpu_ov); without it, pcie / cpu. Clamped to 1.""" + if not isinstance(row, dict): return None - return "hybrid" if all(p == "hybrid" for p in picks) else "offload" - - -def load_hybrid_fetch_fraction( - quant_format: str, - gpu_name: str | None = None, - path: str | None = None, - gpu_uuid: str | None = None, -) -> float | None: - """Benched hybrid fetch fraction for ``quant_format``, or ``None``. - - The hybrid backend's bandwidth-matched fetch split: of a decode step's expert misses, - fetch this fraction over PCIe and compute the rest on the CPU, so both finish together. - Preferred source is the *overlapped* pair (CPU MoE and PCIe gather measured while - running concurrently -- the real contention regime): fetched/misses = pcie_ov / - (pcie_ov + cpu_ov). Older profiles without it fall back to the standalone bandwidths - under a full-DRAM-contention assumption (cpu keeps cpu - pcie under DMA), which - reduces to pcie/cpu. Per-dtype entry first, then any per-model entry with this format. - ``None`` = no usable profile; clamped to [0, 1]. - """ - fmt = _QUANT_TO_BENCH_FORMAT.get(quant_format, quant_format) + cpu_ov, pcie_ov = row.get("cpu_moe_overlap_gbs"), row.get("pcie_gather_overlap_gbs") + if cpu_ov and pcie_ov: + return min(1.0, pcie_ov / (pcie_ov + cpu_ov)) + cpu, pcie = row.get("cpu_moe_gbs"), row.get("pcie_gather_gbs") + if cpu and pcie: + return min(1.0, pcie / cpu) + return None + + +def load_layer_kinds(gpu_name: str | None = None, path: str | None = None, gpu_uuid: str | None = None) -> dict | None: + """The rows of the usable profile for this GPU (key -> row), or None without one.""" prof = _usable_profile(gpu_name, path, gpu_uuid) if prof is None: return None - entries = [(prof.get("dtype_kernels") or {}).get(fmt)] + [ - (wl.get("kernels") or {}).get(fmt) - for wl in (prof.get("workloads") or {}).values() - if isinstance(wl, dict) - ] - for entry in entries: - if not isinstance(entry, dict): - continue - cpu_ov, pcie_ov = entry.get("cpu_moe_overlap_gbs"), entry.get("pcie_gather_overlap_gbs") - if cpu_ov and pcie_ov: - return min(1.0, pcie_ov / (pcie_ov + cpu_ov)) - cpu, pcie = entry.get("cpu_moe_gbs"), entry.get("pcie_gather_gbs") - if cpu and pcie: - return min(1.0, pcie / cpu) - return None + rows = prof.get("layer_kinds") + return rows if isinstance(rows, dict) else {} diff --git a/python/freetoken/moe/benchbw.py b/python/freetoken/moe/benchbw.py index 83b3d19665..47c5ca896d 100644 --- a/python/freetoken/moe/benchbw.py +++ b/python/freetoken/moe/benchbw.py @@ -11,23 +11,24 @@ PCIe on a decode miss). These are what the hybrid-vs-offload choice rides on. When both work, the same pair is also measured *overlapped* (running concurrently, contending for host DRAM) -- those two achieved bandwidths set the hybrid backend's - per-step fetch split (``load_hybrid_fetch_fraction``). + per-step fetch split (``bench_profile.fetch_fraction``). The rule, per (workload, format): recommend ``hybrid`` when the CPU MoE kernel bandwidth exceeds ``threshold`` x the PCIe gather bandwidth (default 2x), otherwise ``offload``. Formats with no CPU MoE weight path (e.g. block-fp8) can't be computed on the CPU, so they always resolve to ``offload``. -The hybrid-vs-offload choice is dtype-dominated, so the default is a **per-dtype tuning bench** -(``--dtype``): one bench per expert format against a canonical geometry -- the minimal set the -runtime backend pick matches on. ``--model`` additionally benches specific model geometries for -per-model detail. Results are written to a JSON file per GPU (default ``$XDG_CACHE_HOME/ -freetoken/benchbw/.json``) so the choice is reproducible and a multi-GPU box keeps -one profile per card. +Each bench writes one row per MoE layer kind: format, options, expert shape, top_k and CPU +worker threads (``bench_profile.layer_kind_key``). ``--model `` benches every kind +a real checkpoint has, which is what the hybrid fetch split matches per layer; ``--dtype`` and +the ``--model`` presets bench canonical geometries. The rows live in a JSON file per GPU +(default ``$XDG_CACHE_HOME/freetoken/benchbw/.json``); a run replaces the rows it +benched and keeps the others. ft bench bw # per-dtype tuning bench (default: all formats) ft bench bw --dtype nvfp4,bf16 # only these formats ft bench bw --model qwen3.6-moe # per-model detail instead + ft bench bw --model /path/to/model.gguf # every MoE layer kind of a checkpoint ft bench bw --dtype all --model all # both ft bench bw --gpu GPU-2f3a... # bench a specific GPU (UUID or nvidia-smi index) """ @@ -38,9 +39,11 @@ import contextlib import json import logging +import math import os import socket import statistics +import subprocess import threading import time from dataclasses import dataclass @@ -56,14 +59,16 @@ single_gpu_arg, ) from freetoken.kernel.pinned import alloc_pinned_tensor +from freetoken.memory import effective_memory_available +from freetoken.moe.bench_profile import PROFILE_VERSION from freetoken.moe.cpu_executor import physical_core_cpus, resolve_threads_and_affinity from freetoken.utils import init_logger logger = init_logger(__name__) # Formats the CPU MoE C++ kernel can compute AND this bench can build banks for; anything -# else is offload-only here. (The kernel also does q4_0, but this bench has no q4_0 banks.) -_CPU_MOE_FORMATS = frozenset({"bf16", "nvfp4", "mxfp4_triton", "ds_fp4"}) +# else is offload-only here. +_CPU_MOE_FORMATS = frozenset({"bf16", "nvfp4", "mxfp4_triton", "ds_fp4", "fp8_block", "gguf", "gguf_scaled"}) # Formats this bench can build synthetic (correctly-sized) banks for. _BUILDABLE_FORMATS = frozenset({"bf16", "nvfp4", "fp8_block", "mxfp4_triton", "ds_fp4"}) # Friendlier CLI/display aliases for the internal quant_format strings. @@ -108,6 +113,14 @@ 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 + options: tuple[tuple[str, int], ...] = () # CPU executor options of the format, e.g. the GGUF types + + def format_options(self, fmt: str) -> dict[str, int]: + opts = dict(self.options) + if self.act_block is not None and fmt == "ds_fp4": + opts.setdefault("act_block", self.act_block) + return opts # Preset workloads. Dims from the model configs / benchmarks/bench_offload_cache_copy.py. @@ -120,7 +133,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",), @@ -129,20 +143,18 @@ class Workload: "minimax-m2.5": Workload("minimax-m2.5", 3072, 1536, 256, 8, ("nvfp4",)), } -# Per-dtype canonical geometries for the TUNING bench (`--dtype`). The hybrid-vs-offload choice is -# dtype-dominated -- the CPU-MoE-vs-PCIe-gather bandwidth ratio rides on the expert *format* and -# the hardware, not the exact model -- so measuring one representative geometry per format is -# enough for the runtime backend pick (the engine also matches by dtype). This is the minimal set -# the Desktop "detect" flow runs; `--model` still benches specific model geometries for per-model -# detail. Geometries mirror the model families that ship each format (mxfp4 keeps gpt-oss's large -# experts -- the reason that format sits near the 2x threshold). +# Per-dtype canonical geometries for the TUNING bench (`--dtype`): one representative geometry per +# expert format, the set the Desktop "detect" flow runs. A model's layers match these rows only when +# their kind is the same; `--model ` benches a model's own kinds. Geometries mirror the +# model families that ship each format (mxfp4 keeps gpt-oss's large experts -- the reason that +# format sits near the 2x threshold). DTYPE_WORKLOADS: dict[str, Workload] = { "bf16": Workload("dtype:bf16", 2048, 768, 128, 8, ("bf16",)), "nvfp4": Workload("dtype:nvfp4", 3072, 1536, 128, 8, ("nvfp4",)), "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), } @@ -153,29 +165,10 @@ def default_out_path(gpu_uuid: str | None = None) -> str: return default_profile_path(gpu_uuid) -def _cgroup_mem_headroom() -> int | None: - """Free bytes under this process's cgroup v2 memory limit, or None if unlimited.""" - try: - with open("/sys/fs/cgroup/memory.max") as f: - raw = f.read().strip() - if raw == "max": - return None - limit = int(raw) - with open("/sys/fs/cgroup/memory.current") as f: - used = int(f.read().strip()) - return max(0, limit - used) - except (OSError, ValueError): - return None - - def _available_ram_bytes() -> int: - """Free host RAM, clamped to the cgroup memory limit so a container isn't over-estimated.""" - try: - host = os.sysconf("SC_AVPHYS_PAGES") * os.sysconf("SC_PAGE_SIZE") - except (ValueError, OSError, AttributeError): - host = 8 << 30 - cg = _cgroup_mem_headroom() - return min(host, cg) if cg is not None else host + """Effective free RAM, retaining the legacy fallback only when probing is absent.""" + available = effective_memory_available() + return (8 << 30) if available is None else available # ============================== ceilings (hardware) ============================== @@ -282,9 +275,25 @@ def _timed(copy) -> float: # ============================== bank geometry ============================== -def _offload_bank_specs(fmt: str, H: int, I: int) -> dict[str, tuple[int, torch.dtype]]: +def _gguf_bank_shapes(fmt: str, opts: dict, H: int, I: int) -> dict[str, tuple[tuple[int, ...], torch.dtype]]: + """Per-expert (row shape, dtype) of the GGUF banks: ggml rows as bytes, float rows in their dtype.""" + from freetoken.layers.quantization.ggml_types import float_dtype, row_bytes + + def rows(ggml_type: int, n: int, k: int): + dtype = float_dtype(ggml_type) + return ((n, k), dtype) if dtype is not None else ((n, row_bytes(ggml_type, k)), torch.uint8) + + shapes = {"gate_up": rows(opts["gate_up_type"], 2 * I, H), "down": rows(opts["down_type"], H, I)} + if fmt == "gguf_scaled": + shapes["global"] = ((4,), torch.float32) + return shapes + + +def _offload_bank_specs(fmt: str, H: int, I: int, opts: dict | None = None) -> dict[str, tuple[int, torch.dtype]]: """Per-bank (elems_per_expert, dtype) for the flat offload host banks (the gather).""" - u8, bf16, f16, f8 = torch.uint8, torch.bfloat16, torch.float16, torch.float8_e4m3fn + u8, bf16, f32, f8 = torch.uint8, torch.bfloat16, torch.float32, torch.float8_e4m3fn + if fmt in ("gguf", "gguf_scaled"): + return {name: (math.prod(shape), dtype) for name, (shape, dtype) in _gguf_bank_shapes(fmt, opts, H, I).items()} if fmt == "bf16": return {"gate_up": (2 * I * H, bf16), "down": (H * I, bf16)} if fmt == "fp8_block": @@ -295,8 +304,8 @@ def _offload_bank_specs(fmt: str, H: int, I: int) -> dict[str, tuple[int, torch. if fmt == "nvfp4": return { "gate_up_packed": (2 * I * (H // 2), u8), "gate_up_scale": (2 * I * (H // 16), u8), - "gate_up_global": (2 * I, f16), "down_packed": (H * (I // 2), u8), - "down_scale": (H * (I // 16), u8), "down_global": (H, f16), + "gate_up_global": (2 * I, f32), "down_packed": (H * (I // 2), u8), + "down_scale": (H * (I // 16), u8), "down_global": (H, f32), } if fmt == "mxfp4_triton": return { @@ -312,8 +321,8 @@ def _offload_bank_specs(fmt: str, H: int, I: int) -> dict[str, tuple[int, torch. raise NotImplementedError(fmt) -def _expert_bytes(fmt: str, H: int, I: int) -> int: - return sum(elems * dt.itemsize for elems, dt in _offload_bank_specs(fmt, H, I).values()) +def _expert_bytes(fmt: str, H: int, I: int, opts: dict | None = None) -> int: + return sum(elems * dt.itemsize for elems, dt in _offload_bank_specs(fmt, H, I, opts).values()) def _synth_experts(E: int, expert_bytes: int) -> int: @@ -324,7 +333,7 @@ def _synth_experts(E: int, expert_bytes: int) -> int: return min(E, max(1, _SYNTH_BANK_BUDGET // max(1, expert_bytes))) -def _cpu_moe_bank_sources(fmt: str, H: int, I: int, E: int) -> dict: +def _cpu_moe_bank_sources(fmt: str, H: int, I: int, E: int, opts: dict | None = None) -> dict: """3D pinned banks in the exact layout ``CpuMoeExecutor`` expects for ``fmt``. Scale banks that decode through an e8m0 exponent (mxfp4/ds_fp4) are set to a unit @@ -335,22 +344,54 @@ def _cpu_moe_bank_sources(fmt: str, H: int, I: int, E: int) -> dict: def pin(*shape, dtype): return alloc_pinned_tensor(*shape, dtype=dtype) + if fmt in ("gguf", "gguf_scaled"): + from freetoken.third_party.gguf import GGML_QUANT_SIZES, GGMLQuantizationType + + b = {} + for name, (shape, dtype) in _gguf_bank_shapes(fmt, opts, H, I).items(): + b[name] = pin(E, *shape, dtype=dtype) + if name == "global": + b[name].fill_(1.0) + elif dtype == torch.uint8: + b[name].zero_() # zero ggml blocks: every fp16 block scale is 0, never denormal + if opts["gate_up_type" if name == "gate_up" else "down_type"] == GGMLQuantizationType.MXFP4: + # each block starts with an e8m0 scale byte; 0 decodes to 2^-127, a denormal that slows the CPU kernel + b[name].view(E, -1, GGML_QUANT_SIZES[GGMLQuantizationType.MXFP4][1])[..., 0] = 127 + else: + b[name].fill_(0.02) + return b if fmt == "bf16": gate_up, down = pin(E, 2 * I, H, dtype=torch.bfloat16), pin(E, H, I, dtype=torch.bfloat16) gate_up.fill_(0.02) # uninitialized bf16 can be denormal -> x86 FP slowdown down.fill_(0.02) return {"gate_up": gate_up, "down": down} + if fmt == "fp8_block": + def nb(n): + return (n + 127) // 128 + b = { + "gate_up": pin(E, 2 * I, H, dtype=torch.float8_e4m3fn), + "gate_up_scale": pin(E, nb(2 * I), nb(H), dtype=torch.bfloat16), + "down": pin(E, H, I, dtype=torch.float8_e4m3fn), + "down_scale": pin(E, nb(H), nb(I), dtype=torch.bfloat16), + } + # e4m3 0x00 is a clean zero and the scales are unit-ish, so nothing lands in + # the fp32 denormal range (which would slow the GEMV and skew the timing). + b["gate_up"].view(torch.uint8).fill_(0x38) # 1.0 + b["down"].view(torch.uint8).fill_(0x38) + b["gate_up_scale"].fill_(1.0) + b["down_scale"].fill_(1.0) + return b if fmt == "nvfp4": b = { "gate_up_packed": pin(E, 2 * I, H // 2, dtype=torch.uint8), "gate_up_scale": pin(E, 2 * I, H // 16, dtype=torch.uint8), - "gate_up_global": pin(E, 2 * I, dtype=torch.float16), + "gate_up_global": pin(E, 2 * I, dtype=torch.float32), "down_packed": pin(E, H, I // 2, dtype=torch.uint8), "down_scale": pin(E, H, I // 16, dtype=torch.uint8), - "down_global": pin(E, H, dtype=torch.float16), + "down_global": pin(E, H, dtype=torch.float32), } b["gate_up_global"].fill_(1.0) - b["down_global"].fill_(1.0) # e4m3 scales decode to normal float32; globals are fp16 + b["down_global"].fill_(1.0) # e4m3 scales decode to normal float32 return b if fmt == "mxfp4_triton": # transposed split-K layout b = { @@ -389,19 +430,16 @@ def _build_gather_rig(fmt: str, wl: Workload, device: torch.device): bench and the overlapped CPU+PCIe bench.""" from freetoken.moe.offload_cache import OffloadMoeCache - H, I = wl.hidden, wl.inter - E = _synth_experts(wl.experts, _expert_bytes(fmt, H, I)) - specs = _offload_bank_specs(fmt, H, I) + H, I, opts = wl.hidden, wl.inter, wl.format_options(fmt) + E = _synth_experts(wl.experts, _expert_bytes(fmt, H, I, opts)) + specs = _offload_bank_specs(fmt, H, I, opts) cache = OffloadMoeCache(num_layers=1, num_experts=E, cache_size=E, device=device, quant_format=fmt) total_bytes = 0 + sources = {} for name, (elems, dtype) in specs.items(): - src = alloc_pinned_tensor(E, elems, dtype=dtype) # cudaHostAlloc (mapped) like the loaders - dst = torch.empty(E, elems, dtype=dtype, device=device) - cache.bank_sources[name] = [src] - cache.bank_caches[name] = dst - cache.banks.append(([src], dst)) + sources[name] = [alloc_pinned_tensor(E, elems, dtype=dtype)] # cudaHostAlloc (mapped) like the loaders total_bytes += E * elems * dtype.itemsize - cache._build_copy_plan() # enable the fused path (production default) when aligned + cache.set_bank_sources(sources) # builds the fused copy plan (production default) when aligned cache.evict_slots[:E].copy_(torch.arange(E, dtype=torch.int32, device=device)) cache.src_indices[:E].copy_(torch.randperm(E, device=device).to(torch.int32)) @@ -421,7 +459,7 @@ def measure_pcie_gather_bw(fmt: str, wl: Workload, device: torch.device, iters: (all synth experts, scattered sources). Reusing the cache gives the kernel-correct int64 ``num_indices``. Timed with CUDA events. """ - eb = _expert_bytes(fmt, wl.hidden, wl.inter) + eb = _expert_bytes(fmt, wl.hidden, wl.inter, wl.format_options(fmt)) cache, total_bytes = _build_gather_rig(fmt, wl, device) E = cache.num_experts @@ -460,7 +498,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, layer_formats=[(fmt, wl.format_options(fmt))] * cache.num_layers, ) @@ -471,17 +509,16 @@ def _cpu_moe_step_fns(ex, wl: Workload, E: int): disjoint experts (working set >> LLC) -- the DRAM-bound regime hybrid decode actually runs in; ``run()`` executes one step (releases the GIL inside the extension).""" bs, top_k = 1, wl.top_k - io = ex._io_for(bs) + io = ex._io io["x"].copy_(torch.randn(bs, wl.hidden, dtype=torch.bfloat16) * 0.1) io["w"].copy_(torch.rand(bs, top_k, dtype=torch.float32)) - task = ex._task_for(0, bs) base = torch.arange(top_k, dtype=torch.int32) def set_ids(i: int) -> None: io["ids"].copy_((((i * top_k) + base) % E).view(bs, top_k)) def run() -> None: - ex._ext.run_task(task) + ex._ext.run(0, bs) return set_ids, run @@ -513,25 +550,25 @@ def measure_cpu_moe_bw(fmt: str, wl: Workload, iters: int = 64, num_threads: int the SAME banks (built once) to populate the diagnostic ``isa_sweep`` breakdown; tiers above the CPU's support clamp down and dedupe by the resulting isa name. """ - H, I = wl.hidden, wl.inter - eb = _expert_bytes(fmt, H, I) + H, I, opts = wl.hidden, wl.inter, wl.format_options(fmt) + eb = _expert_bytes(fmt, H, I, opts) E = _synth_experts(wl.experts, eb) - banks = _cpu_moe_bank_sources(fmt, H, I, E) + banks = _cpu_moe_bank_sources(fmt, H, I, E, opts) - def run(isa: str | None) -> tuple[str, float]: + def run(isa: str | None) -> tuple[str, float, list[int]]: with _forced_isa(isa): ex = _build_cpu_moe_executor(fmt, wl, banks, num_threads, E) bw = round(_time_cpu_moe(ex, wl, eb, E, iters), 2) - label = ex.isa + label, cores = ex.isa, list(ex.core_ids) del ex - return label, bw + return label, bw, cores - auto_isa, auto_bw = run(None) # recommendation basis: the tier production would pick + auto_isa, auto_bw, cores = run(None) # recommendation basis: the tier production would pick sweep = {auto_isa: auto_bw} for isa in isas or []: - label, bw = run(isa) + label, bw, _ = run(isa) sweep[label] = max(sweep.get(label, 0.0), bw) # dedupe clamped tiers, keep best - return {"bw_gbs": auto_bw, "isa": auto_isa, "isa_sweep": sweep, + return {"bw_gbs": auto_bw, "isa": auto_isa, "isa_sweep": sweep, "cores": cores, "expert_bytes": eb, "synth_experts": E} @@ -547,15 +584,15 @@ def measure_overlap_bw(fmt: str, wl: Workload, device: torch.device, backend its bandwidth-matched fetch split: fetched : cpu-computed = pcie_ov : cpu_ov. Both sides hammer their own pinned banks flat out for ``seconds`` from a shared - barrier: a worker thread loops bs=1 CPU decode steps (``run_task`` releases the GIL) + barrier: a worker thread loops bs=1 CPU decode steps (``run`` releases the GIL) while the main thread loops full-layer ``copy_missing`` gathers (synchronized per copy so the DMA is really in flight, not just enqueued). Each side reports bytes / its own elapsed; the two windows differ by at most one CPU step + one gather. """ - H, I = wl.hidden, wl.inter - eb = _expert_bytes(fmt, H, I) + H, I, opts = wl.hidden, wl.inter, wl.format_options(fmt) + eb = _expert_bytes(fmt, H, I, opts) E = _synth_experts(wl.experts, eb) - ex = _build_cpu_moe_executor(fmt, wl, _cpu_moe_bank_sources(fmt, H, I, E), num_threads, E) + ex = _build_cpu_moe_executor(fmt, wl, _cpu_moe_bank_sources(fmt, H, I, E, opts), num_threads, E) set_ids, run_step = _cpu_moe_step_fns(ex, wl, E) cache, gather_bytes = _build_gather_rig(fmt, wl, device) @@ -616,17 +653,15 @@ def _bench_format(fmt: str, wl: Workload, device: torch.device, threshold: float one format can't kill the run; an unexpected exception is left to propagate. """ entry: dict = { - "expert_bytes": None, "synth_experts": None, "cpu_moe_gbs": None, "cpu_moe_isa": None, - "isa_sweep": None, "pcie_gather_gbs": None, + "cpu_moe_gbs": None, "cpu_moe_isa": None, + "cpu_cores": None, "isa_sweep": None, "pcie_gather_gbs": None, "cpu_moe_overlap_gbs": None, "pcie_gather_overlap_gbs": None, - "ratio": None, "recommended": None, + "recommended": None, "note": None, } try: g = measure_pcie_gather_bw(fmt, wl, device, pcie_iters) entry["pcie_gather_gbs"] = round(g["bw_gbs"], 2) - entry["expert_bytes"] = g["expert_bytes"] - entry["synth_experts"] = g["synth_experts"] except (ImportError, RuntimeError) as e: logger.warning(f"benchbw: PCIe gather bench failed for {wl.name}/{fmt}: {e}") _note(entry, f"pcie gather unavailable ({e})") @@ -640,28 +675,29 @@ def _bench_format(fmt: str, wl: Workload, device: torch.device, threshold: float c = measure_cpu_moe_bw(fmt, wl, cpu_iters, cpu_threads, isas=isas) entry["cpu_moe_gbs"] = c["bw_gbs"] # auto/best tier (what production runs) entry["cpu_moe_isa"] = c["isa"] + entry["cpu_cores"] = c["cores"] if isas: # a sweep was requested -> keep the per-tier breakdown entry["isa_sweep"] = c["isa_sweep"] # The tier label comes from the bf16 dot chooser, so it can over-report tiers # the actual per-format kernel doesn't distinguish. Disclose the real story. - if fmt == "nvfp4": # WF_NVFP4 is the only use_vnni format + if fmt == "nvfp4": # nvfp4 is the only use_vnni format _note(entry, "isa tier labels are nominal: this W4A8 kernel rides AVX-VNNI " "regardless of the forced tier") - elif fmt in ("mxfp4_triton", "ds_fp4"): + elif fmt in ("mxfp4_triton", "ds_fp4", "fp8_block"): _note(entry, "isa tier labels are nominal: avx512bf16 == avx512f here " "(3 real kernels: scalar/avx2/avx512, no bf16/VNNI variant)") - entry["expert_bytes"] = entry["expert_bytes"] or c["expert_bytes"] - entry["synth_experts"] = entry["synth_experts"] or c["synth_experts"] + elif fmt in ("gguf", "gguf_scaled"): + _note(entry, "isa tier labels are nominal: GGUF rows run on the ggml-cpu variant " + "picked at load, regardless of the forced tier") except (ImportError, RuntimeError) as e: logger.warning(f"benchbw: CPU MoE bench failed for {wl.name}/{fmt}: {e}") _note(entry, f"cpu moe unavailable ({e})") cpu_g, pcie_g = entry["cpu_moe_gbs"], entry["pcie_gather_gbs"] if cpu_g is not None and pcie_g: # pcie_g truthy also rules out a div-by-zero - entry["ratio"] = round(cpu_g / pcie_g, 3) entry["recommended"] = recommend(cpu_g, pcie_g, threshold) # Both sides work standalone -> also measure them contended (concurrently). This - # pair sets the hybrid backend's fetch split (load_hybrid_fetch_fraction); the + # pair sets the hybrid backend's fetch split (bench_profile.fetch_fraction); the # hybrid-vs-offload verdict above stays on the standalone numbers. try: o = measure_overlap_bw(fmt, wl, device, cpu_threads) @@ -679,6 +715,89 @@ def _bench_format(fmt: str, wl: Workload, device: torch.device, threshold: float return entry +def _kind(fmt: str, wl: Workload) -> tuple: + """A layer kind without the CPU thread count, which the row key adds once the executor has sized its pool.""" + return (fmt, tuple(sorted(wl.format_options(fmt).items())), wl.hidden, wl.inter, wl.top_k) + + +def checkpoint_layer_kinds(model_path: str, device: torch.device) -> dict[tuple, Workload]: + """Each distinct MoE layer kind of a checkpoint as hybrid decode would serve it: kind -> workload.""" + from freetoken.distributed import DistributedInfo, set_tp_info, try_get_tp_info + from freetoken.engine.config import EngineConfig + from freetoken.engine.engine import meta_offload_layer_methods + from freetoken.layers import set_rope_device + + if try_get_tp_info() is None: + set_tp_info(rank=0, size=1) + tp = try_get_tp_info() + cfg = EngineConfig(model_path=model_path, tp_info=DistributedInfo(tp.rank, tp.size), dtype=torch.bfloat16, moe_strategy="hybrid") + if not getattr(cfg.model_config, "is_moe", False): + raise RuntimeError(f"{model_path} has no MoE layers") + set_rope_device(device) + methods = meta_offload_layer_methods(cfg, "hybrid") + if methods is None: + raise RuntimeError(f"{model_path}: its MoE layers bind no expert method") + kinds: dict[tuple, Workload] = {} + for m in methods: + fmt, opts = m.cpu_spec() + c = m.cfg + wl = Workload( + f"{os.path.basename(model_path.rstrip('/'))}:{fmt}", c.hidden, c.local_intermediate, c.num_experts, c.top_k, (fmt,), + activation=c.activation, swiglu_alpha=float(c.alpha), swiglu_limit=c.limit, options=tuple(sorted(opts.items())), + ) + kinds.setdefault(_kind(fmt, wl), wl) + return kinds + + +def _freetoken_build() -> dict: + """FreeToken version and, in a git checkout, its commit: rows benched on older kernels may be stale.""" + from freetoken.version import __version__ + + try: + commit = subprocess.run( + ["git", "-C", os.path.dirname(os.path.abspath(__file__)), "rev-parse", "--short=12", "HEAD"], + capture_output=True, text=True, timeout=5, check=True, + ).stdout.strip() or None + except (OSError, subprocess.SubprocessError): + commit = None + return {"version": __version__, "commit": commit} + + +def _cpu_info() -> dict: + model = numa = None + try: + with open("/proc/cpuinfo") as f: + model = next((line.split(":", 1)[1].strip() for line in f if line.startswith("model name")), None) + except OSError: + pass + try: + numa = sum(1 for n in os.listdir("/sys/devices/system/node") if n.startswith("node") and n[4:].isdigit()) or None + except OSError: + pass + return {"model": model, "physical_cores": len(physical_core_cpus()), "numa_nodes": numa} + + +def _merge_previous(result: dict, path: str) -> None: + """Keep the rows of the profile already at ``path`` (same GPU and version) that this run did not re-bench; a re-benched row keeps its other sources.""" + try: + with open(path) as f: + prev = json.load(f) + except (OSError, ValueError): + return + if not isinstance(prev, dict) or prev.get("version") != result["version"] or (prev.get("gpu") or {}).get("uuid") != result["gpu"]["uuid"]: + return + prev_rows = prev.get("layer_kinds") + if not isinstance(prev_rows, dict): + return + rows = dict(prev_rows) + for key, row in result["layer_kinds"].items(): + old = prev_rows.get(key) + if isinstance(old, dict) and isinstance(old.get("sources"), list): + row["sources"] = sorted(set(old["sources"]) | set(row["sources"])) + rows[key] = row + result["layer_kinds"] = rows + + def _atomic_write_json(path: str, obj: dict) -> None: os.makedirs(os.path.dirname(os.path.abspath(path)) or ".", exist_ok=True) tmp = f"{path}.tmp.{os.getpid()}" @@ -702,7 +821,11 @@ def run_benchbw( pcie_iters: int = 30, kernel_cpu_iters: int = 64, kernel_pcie_iters: int = 20, + checkpoints: tuple[str, ...] = (), ) -> dict: + from freetoken.moe.bench_profile import layer_kind_key + from freetoken.moe.cpu_executor import _format_label + if not torch.cuda.is_available(): raise RuntimeError( "benchbw needs a CUDA device to measure PCIe bandwidth (both offload and " @@ -720,12 +843,29 @@ def run_benchbw( torch.cuda.set_device(device) gpu = gpu_identity(device_index) + # kind -> (workload, sources); checkpoints are read before any timing so a bad one fails fast + units: dict[tuple, tuple[Workload, list[str]]] = {} + + def add(fmt: str, wl: Workload, source: str) -> None: + sources = units.setdefault(_kind(fmt, wl), (wl, []))[1] + if source not in sources: + sources.append(source) + + for fmt in dtypes or (): + add(fmt, DTYPE_WORKLOADS[fmt], f"dtype:{_FORMAT_DISPLAY.get(fmt, fmt)}") + for mname in models: + for fmt in formats or WORKLOADS[mname].formats: + add(fmt, WORKLOADS[mname], f"preset:{mname}") + for path in checkpoints: + logger.info(f"benchbw: reading the MoE layer kinds of {path} ...") + for kind, wl in checkpoint_layer_kinds(path, device).items(): + add(kind[0], wl, path) + # Machine-parseable progress on stdout (opt-in), so the daemon/Desktop can stream feedback # while the bench runs -- mirrors ft checkpoint's FTCONVERT lines. `done`/`total` count the - # measured units (ceilings + one per benched format); `label` is what's measuring now. + # measured units (ceilings + one per layer kind); `label` is what's measuring now. prog_on = os.environ.get("FREETOKEN_BENCH_PROGRESS") == "1" - n_model_units = sum(len(formats or WORKLOADS[m].formats) for m in models) - total = 1 + len(dtypes or ()) + n_model_units + total = 1 + len(units) done = 0 def _prog(label: str) -> None: @@ -743,64 +883,52 @@ def _prog(label: str) -> None: exec_log = logging.getLogger("freetoken.moe.cpu_executor") prev_level = exec_log.level exec_log.setLevel(logging.WARNING) - dtypes_out: dict = {} - workloads_out: dict = {} + now = datetime.now(timezone.utc).astimezone().isoformat(timespec="seconds") + build = _freetoken_build() + rows: dict = {} try: - # dtype tuning: one bench per expert format (canonical geometry). This is what the - # runtime backend pick keys off -- matched by dtype in bench_profile.py. - for fmt in dtypes or (): - wl = DTYPE_WORKLOADS[fmt] - _prog(f"dtype:{_FORMAT_DISPLAY.get(fmt, fmt)}") - logger.info(f"benchbw: dtype {_FORMAT_DISPLAY.get(fmt, fmt)} real kernels ...") - dtypes_out[fmt] = _bench_format( - fmt, wl, device, threshold, cpu_threads, kernel_cpu_iters, kernel_pcie_iters, isas - ) + for (fmt, opts, hidden, inter, top_k), (wl, sources) in units.items(): + label = f"{_format_label(fmt, opts)} {hidden}x{inter} top{top_k}" + _prog(label) + logger.info(f"benchbw: {label} real kernels ...") + entry = _bench_format(fmt, wl, device, threshold, cpu_threads, kernel_cpu_iters, kernel_pcie_iters, isas) done += 1 - # optional per-model detail (--model): specific geometries, unchanged. - for mname in models: - wl = WORKLOADS[mname] - fmts = formats or wl.formats - kernels = {} - for fmt in fmts: - _prog(f"{mname}:{_FORMAT_DISPLAY.get(fmt, fmt)}") - logger.info(f"benchbw: {mname}/{fmt} real kernels ...") - kernels[fmt] = _bench_format( - fmt, wl, device, threshold, cpu_threads, kernel_cpu_iters, kernel_pcie_iters, isas - ) - done += 1 - workloads_out[mname] = { - "model": {"name": wl.name, "hidden": wl.hidden, "inter": wl.inter, - "experts": wl.experts, "top_k": wl.top_k}, - "kernels": kernels, - "recommended_moe_backend": {f: kernels[f]["recommended"] for f in fmts}, + if not entry["cpu_cores"]: + logger.warning(f"benchbw: no CPU MoE result for {label}, so no row: {entry['note']}") + continue + threads = len(entry["cpu_cores"]) + rows[layer_kind_key(fmt, dict(opts), hidden, inter, top_k, threads)] = { + "format": fmt, "options": dict(opts), "hidden": hidden, "inter": inter, "top_k": top_k, + "cpu_threads": threads, "activation": wl.activation, "sources": sources, + "threshold": threshold, "measured_at": now, "freetoken": build, **entry, } finally: exec_log.setLevel(prev_level) _prog("done") result = { - "version": 4, - "timestamp": datetime.now(timezone.utc).astimezone().isoformat(timespec="seconds"), + "version": PROFILE_VERSION, + "timestamp": now, "epoch": int(time.time()), "host": socket.gethostname(), # index is the CUDA ordinal the bench ran on; uuid keys the profile file "gpu": {"index": device_index, "name": gpu["name"], "uuid": gpu["uuid"]}, - "cpu": {"physical_cores": len(physical_core_cpus()), "threads_used": cpu["threads"]}, - "threshold": threshold, + "cpu": _cpu_info(), "ceilings": { "cpu_stream_read_gbs": round(cpu["bw_gbs"], 2), + "cpu_stream_threads": cpu["threads"], "pcie_linear_h2d_gbs": round(pcie["h2d_gbs"], 2), "pcie_linear_d2h_gbs": round(pcie["d2h_gbs"], 2), }, - # Per-dtype verdicts (the tuning axis) + optional per-model detail. - "dtypes": {f: e["recommended"] for f, e in dtypes_out.items()}, - "dtype_kernels": dtypes_out, - "workloads": workloads_out, + "layer_kinds": rows, } out_path = os.path.expanduser(out_path or default_out_path(gpu["uuid"])) + benched = list(rows) + _merge_previous(result, out_path) _atomic_write_json(out_path, result) result["out_path"] = out_path + result["benched"] = benched if prog_on: print(f"FTBENCH_OUT {out_path}", flush=True) return result @@ -818,8 +946,9 @@ def _print_kernels(kernels: dict, iw: int) -> None: f"{'CPU/PCIe':>9} backend") for fmt, e in kernels.items(): disp = _FORMAT_DISPLAY.get(fmt, fmt) - eb = f"{e['expert_bytes'] / 2**20:.2f} MB" if e["expert_bytes"] else "n/a" - ratio = f"{e['ratio']:.2f}x" if e["ratio"] is not None else "—" + eb = f"{_expert_bytes(fmt, e['hidden'], e['inter'], e['options']) / 2**20:.2f} MB" + cpu_g, pcie_g = e["cpu_moe_gbs"], e["pcie_gather_gbs"] + ratio = f"{cpu_g / pcie_g:.2f}x" if cpu_g is not None and pcie_g else "—" print(f" {disp:<8} {eb:>9} {_gbs(e['cpu_moe_gbs'])} {_gbs(e['pcie_gather_gbs'])} " f"{ratio:>9} {e['recommended']}") c_ov, p_ov = e.get("cpu_moe_overlap_gbs"), e.get("pcie_gather_overlap_gbs") @@ -836,27 +965,24 @@ def _print_kernels(kernels: dict, iw: int) -> None: def _print_report(r: dict) -> None: - c = r["ceilings"] + from freetoken.moe.cpu_executor import _format_label + + c, cpu, rows = r["ceilings"], r["cpu"], r["layer_kinds"] print() print(f" host {r['host']} gpu cuda:{r['gpu']['index']} ({r['gpu']['name']}) " - f"cpu {r['cpu']['physical_cores']}c/{r['cpu']['threads_used']}t") + f"cpu {cpu['model']} ({cpu['physical_cores']}c, {cpu['numa_nodes']} NUMA)") print(f" ceilings: CPU STREAM read {c['cpu_stream_read_gbs']:.1f} | " - f"PCIe linear H2D {c['pcie_linear_h2d_gbs']:.1f} D2H {c['pcie_linear_d2h_gbs']:.1f} GB/s" - f" (threshold {r['threshold']}x)") - # one global ISA-name width so every sweep's GB/s lands in the same column; the - # redundant "()" tag is dropped -- the format is already in the row header. - all_kernels = list((r.get("dtype_kernels") or {}).values()) + [ - e for wl in r["workloads"].values() for e in wl["kernels"].values() - ] - iw = max((len(k.split("(")[0]) for e in all_kernels for k in (e.get("isa_sweep") or {})), default=0) - if r.get("dtype_kernels"): - print("\n per-dtype (tuning — what the runtime backend pick matches on)") - _print_kernels(r["dtype_kernels"], iw) - for name, w in r["workloads"].items(): - m = w["model"] - print(f"\n {name} H={m['hidden']} I={m['inter']} E={m['experts']} top_k={m['top_k']}") - _print_kernels(w["kernels"], iw) - print(f"\n saved: {r['out_path']}\n") + f"PCIe linear H2D {c['pcie_linear_h2d_gbs']:.1f} D2H {c['pcie_linear_d2h_gbs']:.1f} GB/s") + benched = [rows[key] for key in r["benched"]] + # one global ISA-name width so every sweep's GB/s lands in the same column + iw = max((len(k.split("(")[0]) for e in benched for k in (e.get("isa_sweep") or {})), default=0) + for e in benched: + label = _format_label(e["format"], tuple(sorted(e["options"].items()))) + sources = ", ".join(os.path.basename(s.rstrip("/")) for s in e["sources"]) + print(f"\n {label} H={e['hidden']} I={e['inter']} top_k={e['top_k']} threads={e['cpu_threads']} " + f" ({sources}; threshold {e['threshold']}x)") + _print_kernels({e["format"]: e}, iw) + print(f"\n saved: {r['out_path']} ({len(benched)} rows benched, {len(rows)} in the profile)\n") # ================================ CLI plumbing ================================ @@ -884,13 +1010,14 @@ def _positive_float(s: str) -> float: def _model_list(s: str) -> tuple[str, ...]: + """Preset names and checkpoint paths; a path is any item that exists on disk.""" if s.strip() == "all": return tuple(WORKLOADS) items = [x.strip() for x in s.split(",") if x.strip()] - bad = [x for x in items if x not in WORKLOADS] + bad = [x for x in items if x not in WORKLOADS and not os.path.exists(x)] if bad or not items: raise argparse.ArgumentTypeError( - f"'all' or a comma-separated subset of {sorted(WORKLOADS)}, got {s!r}" + f"'all', or a comma-separated list of checkpoint paths and presets from {sorted(WORKLOADS)}, got {s!r}" ) return tuple(dict.fromkeys(items)) @@ -936,11 +1063,12 @@ def main(argv: list[str] | None = None, prog: str = "ft bench bw") -> int: formatter_class=argparse.RawDescriptionHelpFormatter) p.add_argument("--dtype", type=_dtype_list, default=None, help=f"tuning bench: 'all' (default when no --model), or comma-separated from " - f"{sorted(DTYPE_WORKLOADS)} -- one bench per expert format (what the runtime " - f"backend pick matches on)") + f"{sorted(DTYPE_WORKLOADS)} -- one bench per expert format at a canonical geometry") p.add_argument("--model", type=_model_list, default=None, help=f"per-model detail instead of/in addition to --dtype: 'all', or comma-" - f"separated from {sorted(WORKLOADS)}") + f"separated presets from {sorted(WORKLOADS)} and checkpoint paths (a " + f"checkpoint benches each of its MoE layer kinds, which the hybrid fetch " + f"split matches per layer)") p.add_argument("--formats", type=_format_list, default=None, help="override a --model's native expert formats (comma-separated)") p.add_argument("--isa", type=_isa_list, default=None, dest="isas", @@ -971,11 +1099,11 @@ def main(argv: list[str] | None = None, prog: str = "ft bench bw") -> int: except (ValueError, RuntimeError) as e: p.error(str(e)) - models = ns.model or () + models = tuple(m for m in ns.model or () if m in WORKLOADS) + checkpoints = tuple(m for m in ns.model or () if m not in WORKLOADS) dtypes = ns.dtype - # Default (no selection): the full per-dtype tuning bench -- the minimal set the runtime - # backend pick needs. --model alone skips it; --dtype + --model runs both. - if not models and not dtypes: + # Default (no selection): the full per-dtype tuning bench. --model alone skips it; --dtype + --model runs both. + if not ns.model and not dtypes: dtypes = tuple(DTYPE_WORKLOADS) try: @@ -984,6 +1112,7 @@ def main(argv: list[str] | None = None, prog: str = "ft bench bw") -> int: dtypes=dtypes, formats=ns.formats, isas=ns.isas, cpu_threads=ns.cpu_threads, cpu_iters=ns.cpu_iters, pcie_bytes=ns.pcie_mib << 20, pcie_iters=ns.pcie_iters, kernel_cpu_iters=ns.kernel_cpu_iters, kernel_pcie_iters=ns.kernel_pcie_iters, + checkpoints=checkpoints, ) except (RuntimeError, OSError) as e: print(f"error: {e}") diff --git a/python/freetoken/moe/cpu_executor.py b/python/freetoken/moe/cpu_executor.py index e2e48d1461..1f275749d7 100644 --- a/python/freetoken/moe/cpu_executor.py +++ b/python/freetoken/moe/cpu_executor.py @@ -1,18 +1,15 @@ """Python wrapper around the ``_cpu_moe`` C++ executor (--moe-strategy cpu). -Owns the persistent CPU worker pool, the per-batch-size pinned IO buffers, and -the per-(layer, batch-size) host-func task descriptors. ``decode`` issues the -whole CUDA-graph-capturable sequence on the current stream: +Owns the persistent CPU worker pool and one pinned IO block of ``max_tokens`` rows that +every layer and batch size shares (one task is in flight at a time). ``decode`` issues +the whole CUDA-graph-capturable sequence on the current stream: - D2H (hidden, topk_ids, topk_weights -> pinned) - -> submit host node (cudaLaunchHostFunc: enqueue MoE task to the pool) - -> sync host node (cudaLaunchHostFunc: spin until the pool drains) + D2H (hidden, topk_ids, topk_weights -> pinned IO rows [:bs]) + -> submit (doorbell naming (layer, bs), or a cudaLaunchHostFunc host node) + -> sync (wait for the CPU pool to finish) -> H2D (pinned expert output -> GPU) -Buffers and tasks are allocated lazily per batch size. GraphRunner runs an eager -``model.forward()`` at each batch size immediately before capturing it, so the -first (eager) call materializes the stable pinned buffers + task pointers that -the subsequent capture embeds in its host/memcpy nodes. +With the doorbell on CUDA, a publish kernel does the D2H and the doorbell, and skips the step when no route goes to the CPU; a combine kernel does the H2D (csrc/cpu_moe/gpu_sync_io.h). """ from __future__ import annotations @@ -35,26 +32,27 @@ # on a 75-layer model). Instead the GPU raises a mapped-pinned "ready" flag at submit; a # persistent CPU coordinator (in _cpu_moe) polls it, runs the layer, and sets a "done" # flag the GPU waits on at sync -- no host-func round-trip. Both GPU-side operations are -# STREAM MEMORY OPERATIONS (cuStreamWriteValue64 / cuStreamWaitValue64, resolved from the -# driver at runtime): they execute on the GPU front end with no SM-resident kernel, so +# STREAM MEMORY OPERATIONS (cuStreamWriteValue64 / cuStreamWaitValue64 on CUDA, see +# gpu_sync.h): they execute on the GPU front end with no SM-resident kernel, so # GPU utilization stays truthful during the CPU compute window. (The first cut used a # spin-wait kernel; that pinned reported utilization at 99% and laptop CPU/GPU dynamic # power schedulers responded by clamping the CPU frequency -- a net decode regression on -# power-coupled edge devices.) Each (layer, decode batch size) pair gets its own flag -# slot, so every captured decode graph rides the handshake. Where memops are unavailable -# (Windows WDDM, vGPU, old drivers -- functionally probed at startup) or the slot -# capacity is exceeded, decode keeps the host-func path (functional, slower). A Python -# watchdog thread turns a wedged coordinator into a loud RuntimeError (via err[] + -# raise_if_unhealthy) instead of an indefinite stream stall. +# power-coupled edge devices.) The doorbell word names the (layer, decode batch size) +# task, so one ready/done pair serves every layer and every captured graph (gpu_sync.h). +# Where memops are unavailable (some driver models, vGPU, old drivers -- functionally probed at +# startup), decode keeps the host-func path (functional, slower). ROCm 7.14 cannot replay +# that path from a graph, so capture fails loudly there instead. A Python watchdog thread +# turns a wedged coordinator into a loud RuntimeError (via raise_if_unhealthy) instead of +# an indefinite stream stall. # Caveat: the coordinator busy-polls one core while decode traffic flows (idle backoff # otherwise); FREETOKEN_CPU_MOE_FLAG_SYNC=0 opts out entirely. _FLAG_SYNC = os.getenv("FREETOKEN_CPU_MOE_FLAG_SYNC", "1") != "0" -# Flag slots per MoE layer: covers this many distinct decode batch sizes (captured graph -# sizes plus any eager padded sizes); more than that is unheard of, and the overflow -# just keeps the host-func path for the extra combos. -_FLAG_SLOTS_PER_LAYER = 16 +# Pool threads spin up to 50 ms before they park. FREETOKEN_CPU_MOE_SPIN=0 opts out. +_SPIN_WAIT = os.getenv("FREETOKEN_CPU_MOE_SPIN", "1") != "0" +# ds_fp4 runs its input FP8 round-trip on the GPU when it can; tests turn this off to compare against the CPU round-trip. +_DSFP4_GPU_PREQUANT = True -# Activation ids must match ActKind in csrc/cpu_moe/cpu_moe_ext.cpp. Id 3 is the +# Activation ids must match ActKind in csrc/cpu_moe/common.h. Id 3 is the # clamped (up + 1) swiglu: "swigluoai" runs it in the generic GEMV epilogue, # "gpt_oss_swiglu" is the same math fused inside the mxfp4 kernel. _ACT_IDS = { @@ -68,9 +66,6 @@ "swiglu_clamp": 4, } -# Weight-format ids must match WFmt in csrc/cpu_moe/cpu_moe_ext.cpp. -_WFMT_IDS = {"bf16": 0, "nvfp4": 1, "mxfp4_triton": 2, "ds_fp4": 3, "q4_0": 4} - def compiled_extension_supports(activation: str) -> bool: """Whether the compiled ``_cpu_moe`` extension can serve ``activation`` @@ -116,6 +111,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. @@ -141,9 +147,19 @@ def resolve_threads_and_affinity(requested: int) -> tuple[int, list[int]]: return len(reps), list(reps) +def _format_label(name: str, opts: tuple) -> str: + """One format instance as the startup log names it, e.g. ``gguf(gate_up=IQ2_XXS, down=IQ4_NL)``.""" + if name in ("gguf", "gguf_scaled"): + from freetoken.layers.quantization.ggml_types import type_name + + gate_up, down = dict(opts)["gate_up_type"], dict(opts)["down_type"] + return f"{name}(gate_up={type_name(gate_up)}, down={type_name(down)})" + return name + (f"({', '.join(f'{k}={v}' for k, v in opts)})" if opts else "") + + class CpuMoeExecutor: """Decode-time CPU expert compute over an ``OffloadMoeCache``'s host banks - (bf16, nvfp4, mxfp4_triton, ds_fp4 or q4_0 — see ``_WFMT_IDS`` / ``_resolve_banks``).""" + (the formats ``_cpu_moe.weight_formats()`` lists — see ``_resolve_banks``).""" def __init__( self, @@ -157,25 +173,43 @@ def __init__( device: torch.device, swiglu_alpha: float = 1.702, swiglu_limit: float | None = None, - fmt: str | None = None, + layer_formats: list[tuple[str, dict[str, int]]] | None = None, ) -> None: from freetoken.kernel import _cpu_moe from freetoken.moe.legacy_format import canonical_role - fmt = fmt or cache.quant_format - if fmt not in _WFMT_IDS: - raise NotImplementedError( - f"--moe-strategy cpu/hybrid computes experts on the CPU and supports " - f"{sorted(_WFMT_IDS)} formats, but this checkpoint's experts are " - f"{fmt!r}; use --moe-strategy offload (GPU-side dequant) instead." - ) + # layer_formats[l] = (format, options) of layer l; None runs every layer in cache.quant_format + if layer_formats is None: + layer_formats = [(cache.quant_format, {})] * int(cache.num_layers) + assert len(layer_formats) == cache.num_layers, (len(layer_formats), cache.num_layers) + keys = [(name, tuple(sorted(opts.items()))) for name, opts in layer_formats] + specs = list(dict.fromkeys(keys)) + names = sorted({name for name, _ in specs}) + for f in names: + if f not in _cpu_moe.weight_formats(): + raise NotImplementedError( + f"--moe-strategy cpu/hybrid computes experts on the CPU and supports " + f"{sorted(_cpu_moe.weight_formats())} formats, but this checkpoint's experts are " + f"{f!r}; use --moe-strategy offload (GPU-side dequant) instead." + ) + if f in ("gguf", "gguf_scaled"): + from freetoken.kernel import ggml_cpu + + if (variant := ggml_cpu.select()) is None: + raise NotImplementedError( + "--moe-strategy cpu/hybrid cannot compute GGUF experts here: this build has no ggml-cpu kernels (x86-64 only); " + "use --moe-strategy offload instead." + ) + _cpu_moe.set_ggml_cpu(variant.api) if activation not in _ACT_IDS: raise NotImplementedError(f"CPU MoE backend: unsupported activation {activation!r}") + if apply_router_weight_on_input and (output_only := {"mxfp4_triton", "ds_fp4"} & set(names)): + raise NotImplementedError(f"CPU MoE backend: {sorted(output_only)} experts apply the router weight on the output only") # ABI probe: a stale prebuilt _cpu_moe.so accepts newer act ids without # error and silently computes the wrong activation in the generic # epilogue -- fail loudly with the rebuild instruction instead. (mxfp4 # handles its act inside the kernel and predates the marker.) - if _ACT_IDS[activation] >= 3 and fmt != "mxfp4_triton": + if _ACT_IDS[activation] >= 3 and names != ["mxfp4_triton"]: supported = getattr(_cpu_moe, "max_generic_act_id", lambda: 2)() if _ACT_IDS[activation] > supported: raise RuntimeError( @@ -188,35 +222,60 @@ def __init__( self.num_layers = int(cache.num_layers) self.num_experts = int(cache.num_experts) self.top_k = int(top_k) - self.quant_format = fmt + self.quant_format = ",".join(names) self.device = device self.max_tokens = int(max_tokens) self.apply_router_weight_on_input = bool(apply_router_weight_on_input) # The per-layer tensors and their pointer tables must outlive the executor # (C++ holds raw addresses into both). self._banks: list[torch.Tensor] = [] - ptrs, (self.H, self.I) = self._resolve_banks( - {canonical_role(name): per_layer for name, per_layer in cache.bank_sources.items()}, fmt - ) + by_role = {canonical_role(name): per_layer for name, per_layer in cache.bank_sources.items()} + # every format's tables address every layer, but a format reads only its own layers' rows + formats, shapes = [], set() + for name, opts in specs: + banks, options, shape = self._resolve_banks(by_role, name, keys.index((name, opts)), dict(opts)) + formats.append((name, banks, {**dict(opts), **options})) + shapes.add(shape) + if len(shapes) != 1: + raise NotImplementedError(f"CPU MoE formats disagree on the expert shape (hidden, intermediate): {sorted(shapes)}") + (self.H, self.I), = shapes + layer_counts = {spec: keys.count(spec) for spec in specs} + + # ds_fp4: the reference FP8 activation round-trip is a scalar per-element chain + # that the C++ side runs single-threaded on the CUDA host-callback thread -- + # straight on the decode critical path (~0.3ms/layer at H=4096, every worker and + # the GPU waiting on it). When a GPU is present we run the numerically identical + # round-trip as a captured GPU kernel BEFORE the D2H (see decode_submit) and tell + # the C++ side to skip its own. Measured on DeepSeek-V4-Flash bs=1 decode: + # 12.85 -> 15.65 tok/s, output bit-identical (tests/moe/test_dsfp4_prequant.py). + # the W4A8 activation round-trip block follows the checkpoint's fp8 block (128 on V4, 32 on V4.1) + blocks = {options.get("act_block") for name, _, options in formats if name == "ds_fp4"} + assert None not in blocks, "ds_fp4 experts need the act_block option, the checkpoint's fp8 activation quant block" + # the GPU round-trip rewrites the whole input, so only when every layer is ds_fp4 with one block + self._gpu_prequant = names == ["ds_fp4"] and len(blocks) == 1 and device.type == "cuda" and _DSFP4_GPU_PREQUANT + self._act_block = next(iter(blocks)) if self._gpu_prequant else None + for name, _, options in formats: + if name == "ds_fp4": + options["input_prequant"] = int(self._gpu_prequant) + if self._gpu_prequant: + logger.info_rank0( + "ds_fp4: input FP8 round-trip moved to the GPU " + "(bit-identical grid; the CPU-side scalar round-trip is skipped)" + ) # Decide the flag handshake up front (env + device + a functional stream-memop # probe): its coordinator needs a core of its own, which the auto thread sizing # below reserves (a coordinator time-slicing against the GEMV workers measurably # destabilizes throughput on fully-subscribed boxes). self._flag_sync = _FLAG_SYNC and device.type == "cuda" - self._cpu_moe = _cpu_moe # module ref for the decode-path memop calls - if self._flag_sync: - probe_scratch = alloc_pinned_tensor(1, dtype=torch.int64) - probe_scratch.zero_() - if not _cpu_moe.memops_probe( - torch.cuda.current_stream().cuda_stream, probe_scratch.data_ptr() - ): - logger.info_rank0( - "cpu-moe flag handshake unavailable: CUDA stream memory operations " - "are not supported here (Windows WDDM / vGPU / old driver); using " - "the cudaLaunchHostFunc sync" - ) - self._flag_sync = False + if self._flag_sync and not _cpu_moe.memops_probe(torch.cuda.current_stream().cuda_stream): + logger.info_rank0( + "cpu-moe flag handshake unavailable: stream memory operations failed the " + "startup probe; using the cudaLaunchHostFunc sync" + ) + self._flag_sync = False + # ROCm 7.14 returns wrong results when it replays host-function nodes, so there only the doorbell is graph-safe. + self._graph_capture_safe = self._flag_sync or _cpu_moe.host_func_capture_safe() nthreads, core_ids = resolve_threads_and_affinity(num_threads) coord_core = -1 @@ -227,6 +286,17 @@ def __init__( nthreads -= 1 core_ids = core_ids[:-1] self._coord_core = coord_core + alloc_io = ( + (lambda *shape, dtype: torch.empty(*shape, dtype=dtype)) + if device.type == "cpu" + else alloc_pinned_tensor + ) + self._io = { + "x": alloc_io(self.max_tokens, self.H, dtype=torch.bfloat16), + "ids": alloc_io(self.max_tokens, self.top_k, dtype=torch.int32), + "w": alloc_io(self.max_tokens, self.top_k, dtype=torch.float32), + "y": alloc_io(self.max_tokens, self.H, dtype=torch.bfloat16), + } self._ext = _cpu_moe.CpuMoeExecutor( num_threads=nthreads, num_layers=self.num_layers, @@ -237,15 +307,21 @@ def __init__( max_tokens=self.max_tokens, activation_id=_ACT_IDS[activation], apply_router_weight_on_input=1 if apply_router_weight_on_input else 0, - weight_format=_WFMT_IDS[fmt], + formats=formats, + layer_formats=[specs.index(k) for k in keys], swiglu_alpha=float(swiglu_alpha), swiglu_limit=float(swiglu_limit) if swiglu_limit is not None else float("inf"), core_ids=core_ids, - **ptrs, + x_ptr=self._io["x"].data_ptr(), + ids_ptr=self._io["ids"].data_ptr(), + w_ptr=self._io["w"].data_ptr(), + y_ptr=self._io["y"].data_ptr(), ) self.num_threads = nthreads self.core_ids = core_ids 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)) @@ -256,34 +332,24 @@ def __init__( ) torch.set_num_threads(clamp) - self._io: dict[int, dict[str, torch.Tensor]] = {} - self._tasks: dict[tuple[int, int], int] = {} - - # Flag-based handshake: mapped-pinned ready/done/err int64 arrays (one slot per - # (MoE layer, decode batch size) pair, allocated as tasks are created) + a - # persistent CPU coordinator that polls ready[], runs the slot's task on the - # pool, and sets done[]. Binary per-step protocol (GPU memops: done=0, ready=1; - # coordinator: consume ready, run, done=1; GPU waits done>=1 -- the WAIT - # immediate is constant, so CUDA-graph replays are safe). err[] is raised by the - # WATCHDOG thread when a ready flag stays unanswered (dead coordinator): it - # poisons done to unblock the stream and raise_if_unhealthy() turns the step - # into a loud error instead of silent stale output. Buffers are kept alive on - # self so the coordinator's pinned pointers stay valid for the executor's - # lifetime (flag_sync itself was decided above, before thread sizing). - self._ready = self._done = self._err = None - self._flag_slots: dict[tuple[int, int], int] = {} # (layer_id, bs) -> slot - self._flag_capacity = self.num_layers * _FLAG_SLOTS_PER_LAYER + # Flag-based handshake: the extension owns the ready/done words (gpu_sync.h) and a + # persistent CPU coordinator that polls the doorbell, runs the task it names on the + # pool, and sets done for the GPU. The WATCHDOG thread marks the executor unhealthy + # when a doorbell stays unanswered (dead coordinator): it poisons done to unblock the + # stream and raise_if_unhealthy() turns the step into a loud error instead of silent + # stale output (flag_sync itself was decided above, before thread sizing). + self._unhealthy = False + self._kernel_io = False if self._flag_sync: - self._ready = alloc_pinned_tensor(self._flag_capacity, dtype=torch.int64) - self._done = alloc_pinned_tensor(self._flag_capacity, dtype=torch.int64) - self._err = alloc_pinned_tensor(self._flag_capacity, dtype=torch.int64) - self._ready.zero_() - self._done.zero_() - self._err.zero_() - self._ext.start_flag_coordinator( - self._ready.data_ptr(), self._done.data_ptr(), self._flag_capacity, - self._coord_core, - ) + self._ext.start_flag_coordinator(self._coord_core) + # ROCm keeps the copies until the kernel IO has been checked on AMD GPUs. + if torch.version.hip is None: + reason = _cpu_moe.kernel_io_probe(torch.cuda.current_stream().cuda_stream) + if reason: + logger.warning_rank0(f"cpu-moe kernel IO unavailable ({reason}); using D2H/H2D copies around the doorbell") + else: + self._ext.enable_kernel_io() + self._kernel_io = True self._watchdog_stop = False # The thread target holds a WEAKREF and re-derefs it each tick: a bound # method would strong-reference the executor forever (the loop never ends @@ -300,26 +366,16 @@ def __init__( ) self._watchdog.start() - # ds_fp4: the reference FP8 activation round-trip is a scalar per-element chain - # that the C++ side runs single-threaded on the CUDA host-callback thread -- - # straight on the decode critical path (~0.3ms/layer at H=4096, every worker and - # the GPU waiting on it). When a GPU is present we run the numerically identical - # round-trip as a captured GPU kernel BEFORE the D2H (see decode_submit) and tell - # the C++ side to skip its own. Measured on DeepSeek-V4-Flash bs=1 decode: - # 12.85 -> 15.65 tok/s, output bit-identical (tests/moe/test_dsfp4_prequant.py). - self._gpu_prequant = fmt == "ds_fp4" and device.type == "cuda" - if self._gpu_prequant: - self._ext.set_input_prequant(True) - logger.info_rank0( - "ds_fp4: input FP8 round-trip moved to the GPU " - "(bit-identical grid; the CPU-side scalar round-trip is skipped)" - ) + fmt_label = "; ".join(f"{_format_label(name, opts)}: {n} layers" for (name, opts), n in layer_counts.items()) logger.info_rank0( f"CPU MoE executor ready: threads={nthreads} (pinned to cores " - f"{core_ids[0]}..{core_ids[-1]}) isa={self.isa} fmt={fmt} " + f"{core_ids[0]}..{core_ids[-1]}) isa={self.isa} " + f"fmt={fmt_label} " 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} " + f"handshake={'kernel' if self._kernel_io else 'memop' if self._flag_sync else 'host-func'}" ) def _make_table(self, layers: list[torch.Tensor]) -> torch.Tensor: @@ -337,193 +393,138 @@ def _make_table(self, layers: list[torch.Tensor]) -> torch.Tensor: self._banks.extend(layers) return table - def _resolve_banks(self, banks: dict, fmt: str) -> tuple[dict, tuple[int, int]]: - """Return (pointer kwargs for the C++ ctor, (H, I)) for the given format. + def _tables(self, **roles: list[torch.Tensor]) -> dict[str, int]: + """Per-layer table addresses keyed by canonical bank role.""" + return {role: self._make_table(layers).data_ptr() for role, layers in roles.items()} + + def _resolve_fp8_block_banks(self, banks: dict, rep: int = 0) -> tuple[dict, dict, tuple[int, int]]: + """Block-FP8 schema: row-major e4m3 + one bf16 scale per 128x128 block; the scale banks pad their last dim, so C++ gets the real row stride.""" + gu, gus = banks["gate_up"], banks["gate_up_scale"] + dn, dns = banks["down"], banks["down_scale"] + assert gu[rep].element_size() == 1 and dn[rep].element_size() == 1, (gu[rep].dtype, dn[rep].dtype) + assert gus[rep].dtype == torch.bfloat16 and dns[rep].dtype == torch.bfloat16, (gus[rep].dtype, dns[rep].dtype) + I = int(gu[rep].shape[1] // 2) + H = int(gu[rep].shape[2]) + assert gu[rep].shape[1] == 2 * I + assert tuple(dn[rep].shape[1:]) == (H, I), (dn[rep].shape, H, I) + assert gus[rep].shape[1] == (2 * I + 127) // 128 and gus[rep].shape[2] >= (H + 127) // 128, (gus[rep].shape, H, I) + assert dns[rep].shape[1] == (H + 127) // 128 and dns[rep].shape[2] >= (I + 127) // 128, (dns[rep].shape, H, I) + banks = self._tables(gate_up=gu, down=dn, gate_up_scale=gus, down_scale=dns) + options = dict(gu_scale_stride=int(gus[rep].stride(1)), dn_scale_stride=int(dns[rep].stride(1))) + return banks, options, (H, I) + + def _resolve_banks(self, banks: dict, fmt: str, rep: int = 0, opts: dict | None = None) -> tuple[dict, dict, tuple[int, int]]: + """Return (banks, options, (H, I)) for the C++ ctor: banks maps each canonical role + the format reads to its per-layer table's address, options holds its integer settings. ``banks[name]`` is a list of ``num_layers`` ``[num_experts, ...]`` tensors (the per-layer host bank contract); shapes are read from the first layer so - per-partition (TP) sizes are exact. Unused pointers are 0. Every pointer kwarg - is actually a per-layer table's address (see ``_make_table``), not a single - bank's -- the C++ ctor resolves ``tbl[layer_id]`` per task. + per-partition (TP) sizes are exact. Every bank address is a per-layer table + (see ``_make_table``), not a single bank's -- the C++ side resolves + ``tbl[layer_id]`` per task. """ + if fmt in ("gguf", "gguf_scaled"): + return self._resolve_gguf_banks(banks, rep, int(opts["gate_up_type"]), int(opts["down_type"]), scaled=fmt == "gguf_scaled") if fmt == "bf16": gate_up = banks["gate_up"] down = banks["down"] - if gate_up[0].dtype != torch.bfloat16 or down[0].dtype != torch.bfloat16: + if gate_up[rep].dtype != torch.bfloat16 or down[rep].dtype != torch.bfloat16: raise NotImplementedError( - f"bf16 CPU MoE requires bf16 banks, got {gate_up[0].dtype}/{down[0].dtype}" + f"bf16 CPU MoE requires bf16 banks, got {gate_up[rep].dtype}/{down[rep].dtype}" ) - H = int(gate_up[0].shape[2]) - I = int(gate_up[0].shape[1] // 2) - assert gate_up[0].shape[1] == 2 * I - assert tuple(down[0].shape[1:]) == (H, I), (down[0].shape, H, I) - ptrs = dict( - gate_up_ptr=self._make_table(gate_up).data_ptr(), - down_ptr=self._make_table(down).data_ptr(), - gate_up_scale_ptr=0, - gate_up_global_ptr=0, - down_scale_ptr=0, - down_global_ptr=0, - gate_up_bias_ptr=0, - down_bias_ptr=0, - ) - return ptrs, (H, I) - - if fmt == "q4_0": - return self._resolve_q4_0_banks(banks) + H = int(gate_up[rep].shape[2]) + I = int(gate_up[rep].shape[1] // 2) + assert gate_up[rep].shape[1] == 2 * I + assert tuple(down[rep].shape[1:]) == (H, I), (down[rep].shape, H, I) + banks = self._tables(gate_up=gate_up, down=down) + return banks, {}, (H, I) if fmt == "mxfp4_triton": - return self._resolve_mxfp4_banks(banks) + return self._resolve_mxfp4_banks(banks, rep) if fmt == "ds_fp4": - return self._resolve_dsfp4_banks(banks) + return self._resolve_dsfp4_banks(banks, rep) + + if fmt == "fp8_block": + return self._resolve_fp8_block_banks(banks, rep) - # nvfp4: packed e2m1 (2/byte) + fp8-e4m3 per-16 block scales + fp16 row globals. + # nvfp4: packed e2m1 (2/byte) + fp8-e4m3 per-16 block scales + fp32 row globals. gup, gus, gug = banks["gate_up"], banks["gate_up_scale"], banks["gate_up_global"] dnp, dns, dng = banks["down"], banks["down_scale"], banks["down_global"] - assert gup[0].dtype == torch.uint8 and dnp[0].dtype == torch.uint8, (gup[0].dtype, dnp[0].dtype) - assert gus[0].element_size() == 1 and dns[0].element_size() == 1, "block scales must be 1 byte" - assert gug[0].dtype == torch.float16 and dng[0].dtype == torch.float16, (gug[0].dtype, dng[0].dtype) - I = int(gup[0].shape[1] // 2) - H = int(gup[0].shape[2] * 2) - assert gup[0].shape[1] == 2 * I + assert gup[rep].dtype == torch.uint8 and dnp[rep].dtype == torch.uint8, (gup[rep].dtype, dnp[rep].dtype) + assert gus[rep].element_size() == 1 and dns[rep].element_size() == 1, "block scales must be 1 byte" + assert gug[rep].dtype == torch.float32 and dng[rep].dtype == torch.float32, (gug[rep].dtype, dng[rep].dtype) + I = int(gup[rep].shape[1] // 2) + H = int(gup[rep].shape[2] * 2) + assert gup[rep].shape[1] == 2 * I assert H % 16 == 0 and I % 16 == 0, (H, I) - assert tuple(dnp[0].shape[1:]) == (H, I // 2), (dnp[0].shape, H, I) - assert tuple(gus[0].shape[1:]) == (2 * I, H // 16), (gus[0].shape, I, H) - assert tuple(dns[0].shape[1:]) == (H, I // 16), (dns[0].shape, H, I) - assert tuple(gug[0].shape[1:]) == (2 * I,) and tuple(dng[0].shape[1:]) == (H,) - ptrs = dict( - gate_up_ptr=self._make_table(gup).data_ptr(), - down_ptr=self._make_table(dnp).data_ptr(), - gate_up_scale_ptr=self._make_table(gus).data_ptr(), - gate_up_global_ptr=self._make_table(gug).data_ptr(), - down_scale_ptr=self._make_table(dns).data_ptr(), - down_global_ptr=self._make_table(dng).data_ptr(), - gate_up_bias_ptr=0, - down_bias_ptr=0, - ) - return ptrs, (H, I) + assert tuple(dnp[rep].shape[1:]) == (H, I // 2), (dnp[rep].shape, H, I) + assert tuple(gus[rep].shape[1:]) == (2 * I, H // 16), (gus[rep].shape, I, H) + assert tuple(dns[rep].shape[1:]) == (H, I // 16), (dns[rep].shape, H, I) + assert tuple(gug[rep].shape[1:]) == (2 * I,) and tuple(dng[rep].shape[1:]) == (H,) + banks = self._tables(gate_up=gup, down=dnp, gate_up_scale=gus, gate_up_global=gug, down_scale=dns, down_global=dng) + return banks, {}, (H, I) - def _resolve_q4_0_banks(self, banks: dict) -> tuple[dict, tuple[int, int]]: - """Native GGUF Q4_0 schema (gemma4 GGUF): per-32 blocks (fp16 scale + 16 nibble - bytes), row-major over K -- the *same* packed banks the GPU offload path streams. - gate_up is [S, 2I, H//32*18], down is [S, H, I//32*18]; the C++ W4A16 GEMV reads a - row in place (18 bytes / 32 K) and dequantizes weights inside the K-loop.""" - gate_up, down = banks["gate_up"], banks["down"] - assert gate_up[0].dtype == torch.uint8 and down[0].dtype == torch.uint8, ( - gate_up[0].dtype, down[0].dtype, - ) - I = int(gate_up[0].shape[1] // 2) - H = int(down[0].shape[1]) - assert gate_up[0].shape[1] == 2 * I - assert H % 32 == 0 and I % 32 == 0, (H, I) - assert int(gate_up[0].shape[2]) == (H // 32) * 18, (gate_up[0].shape, H) - assert int(down[0].shape[2]) == (I // 32) * 18, (down[0].shape, I) - ptrs = dict( - gate_up_ptr=self._make_table(gate_up).data_ptr(), - down_ptr=self._make_table(down).data_ptr(), - gate_up_scale_ptr=0, - gate_up_global_ptr=0, - down_scale_ptr=0, - down_global_ptr=0, - gate_up_bias_ptr=0, - down_bias_ptr=0, - ) - return ptrs, (H, I) + def _resolve_gguf_banks(self, banks: dict, rep: int, gate_up_type: int, down_type: int, scaled: bool = False) -> tuple[dict, dict, tuple[int, int]]: + """GGUF rows read in place: gate_up [E, 2I, row_bytes(H)] (gate rows, then up rows) and down [E, H, row_bytes(I)], one ggml type each; float rows keep their dtype. + + ``scaled`` adds global [E, 4] fp32, each expert's (gate, up, down, 1) NVFP4 weight_scale_2.""" + from freetoken.layers.quantization.ggml_types import float_dtype, row_bytes, type_name - def _resolve_mxfp4_banks(self, banks: dict) -> tuple[dict, tuple[int, int]]: + gate_up, down = banks["gate_up"], banks["down"] + H, I = int(down[rep].shape[1]), int(gate_up[rep].shape[1] // 2) + for role, t, ggml_type, k in (("gate_up", gate_up[rep], gate_up_type, H), ("down", down[rep], down_type, I)): + dtype = float_dtype(ggml_type) if float_dtype(ggml_type) is not None else torch.uint8 + if t.dtype != dtype or t.shape[-1] * t.element_size() != row_bytes(ggml_type, k): + raise ValueError(f"MoE layer {rep} {role} bank {tuple(t.shape)} {t.dtype} is not {type_name(ggml_type)} rows of {k} values") + types = {"gate_up_type": gate_up_type, "down_type": down_type} + if not scaled: + return self._tables(gate_up=gate_up, down=down), types, (H, I) + scales = banks["global"] + if scales[rep].dtype != torch.float32 or tuple(scales[rep].shape) != (gate_up[rep].shape[0], 4) or not scales[rep].is_contiguous(): + raise ValueError(f"MoE layer {rep} global bank {tuple(scales[rep].shape)} {scales[rep].dtype} is not contiguous [experts, 4] fp32") + return self._tables(gate_up=gate_up, down=down, **{"global": scales}), types, (H, I) + + def _resolve_mxfp4_banks(self, banks: dict, rep: int = 0) -> tuple[dict, dict, tuple[int, int]]: """gpt-oss mxfp4 ``mxfp4_triton`` schema: transposed split-K blocks/scales (N innermost) + per-output-row biases. The C++ kernel streams K and accumulates a contiguous N-block, so the GPU-tiled layout is read in place (no repack, no extra host memory). Block scales are e8m0 (1 byte / 32 K).""" gub, gus, gob = banks["gate_up"], banks["gate_up_scale"], banks["gate_up_bias"] dnb, dns, dob = banks["down"], banks["down_scale"], banks["down_bias"] - assert gub[0].dtype == torch.uint8 and dnb[0].dtype == torch.uint8, (gub[0].dtype, dnb[0].dtype) - assert gus[0].dtype == torch.uint8 and dns[0].dtype == torch.uint8, (gus[0].dtype, dns[0].dtype) - assert gob[0].dtype == torch.bfloat16 and dob[0].dtype == torch.bfloat16, (gob[0].dtype, dob[0].dtype) + assert gub[rep].dtype == torch.uint8 and dnb[rep].dtype == torch.uint8, (gub[rep].dtype, dnb[rep].dtype) + assert gus[rep].dtype == torch.uint8 and dns[rep].dtype == torch.uint8, (gus[rep].dtype, dns[rep].dtype) + assert gob[rep].dtype == torch.bfloat16 and dob[rep].dtype == torch.bfloat16, (gob[rep].dtype, dob[rep].dtype) # gate_up_blocks [E, H//2, 2I]; down_blocks [E, I//2, H] - H = int(gub[0].shape[1] * 2) - I = int(gub[0].shape[2] // 2) - assert gub[0].shape[2] == 2 * I + H = int(gub[rep].shape[1] * 2) + I = int(gub[rep].shape[2] // 2) + assert gub[rep].shape[2] == 2 * I assert H % 32 == 0 and I % 32 == 0, (H, I) - assert tuple(dnb[0].shape[1:]) == (I // 2, H), (dnb[0].shape, H, I) - assert tuple(gus[0].shape[1:]) == (H // 32, 2 * I), (gus[0].shape, H, I) - assert tuple(dns[0].shape[1:]) == (I // 32, H), (dns[0].shape, H, I) - assert tuple(gob[0].shape[1:]) == (2 * I,) and tuple(dob[0].shape[1:]) == (H,) - ptrs = dict( - gate_up_ptr=self._make_table(gub).data_ptr(), - down_ptr=self._make_table(dnb).data_ptr(), - gate_up_scale_ptr=self._make_table(gus).data_ptr(), - gate_up_global_ptr=0, - down_scale_ptr=self._make_table(dns).data_ptr(), - down_global_ptr=0, - gate_up_bias_ptr=self._make_table(gob).data_ptr(), - down_bias_ptr=self._make_table(dob).data_ptr(), - ) - return ptrs, (H, I) - - def _resolve_dsfp4_banks(self, banks: dict) -> tuple[dict, tuple[int, int]]: + assert tuple(dnb[rep].shape[1:]) == (I // 2, H), (dnb[rep].shape, H, I) + assert tuple(gus[rep].shape[1:]) == (H // 32, 2 * I), (gus[rep].shape, H, I) + assert tuple(dns[rep].shape[1:]) == (I // 32, H), (dns[rep].shape, H, I) + assert tuple(gob[rep].shape[1:]) == (2 * I,) and tuple(dob[rep].shape[1:]) == (H,) + banks = self._tables(gate_up=gub, down=dnb, gate_up_scale=gus, down_scale=dns, gate_up_bias=gob, down_bias=dob) + return banks, {}, (H, I) + + def _resolve_dsfp4_banks(self, banks: dict, rep: int = 0) -> tuple[dict, dict, tuple[int, int]]: """DeepSeek-V4 ``ds_fp4`` schema: row-major e2m1 (2/byte) + e8m0 per-32 block scales, no global, no bias. Layout matches nvfp4 (K contiguous per output row), so the C++ GEMV reads it in place. The kernel additionally FP8-round-trips the - 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) - assert gus[0].element_size() == 1 and dns[0].element_size() == 1, "block scales must be 1 byte" - 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 tuple(dnp[0].shape[1:]) == (H, I // 2), (dnp[0].shape, H, I) - assert tuple(gus[0].shape[1:]) == (2 * I, H // 32), (gus[0].shape, I, H) - assert tuple(dns[0].shape[1:]) == (H, I // 32), (dns[0].shape, H, I) - ptrs = dict( - gate_up_ptr=self._make_table(gup).data_ptr(), - down_ptr=self._make_table(dnp).data_ptr(), - gate_up_scale_ptr=self._make_table(gus).data_ptr(), - gate_up_global_ptr=0, - down_scale_ptr=self._make_table(dns).data_ptr(), - down_global_ptr=0, - gate_up_bias_ptr=0, - down_bias_ptr=0, - ) - return ptrs, (H, I) - - def _io_for(self, bs: int) -> dict[str, torch.Tensor]: - io = self._io.get(bs) - if io is None: - io = { - "x": alloc_pinned_tensor(bs, self.H, dtype=torch.bfloat16), - "ids": alloc_pinned_tensor(bs, self.top_k, dtype=torch.int32), - "w": alloc_pinned_tensor(bs, self.top_k, dtype=torch.float32), - "y": alloc_pinned_tensor(bs, self.H, dtype=torch.bfloat16), - } - self._io[bs] = io - return io - - def _task_for(self, layer_id: int, bs: int) -> int: - key = (layer_id, bs) - task = self._tasks.get(key) - if task is None: - io = self._io_for(bs) - task = self._ext.create_task( - layer_id, - bs, - io["x"].data_ptr(), - io["ids"].data_ptr(), - io["w"].data_ptr(), - io["y"].data_ptr(), - ) - self._tasks[key] = task - # Allocate this (layer, bs) combo a flag slot and register its task with the - # coordinator. Combos past the slot capacity keep the host-func path. - if self._flag_sync and key not in self._flag_slots: - slot = len(self._flag_slots) - if slot < self._flag_capacity: - self._flag_slots[key] = slot - self._ext.register_flag_task(slot, task) - return task + assert gup[rep].dtype == torch.uint8 and dnp[rep].dtype == torch.uint8, (gup[rep].dtype, dnp[rep].dtype) + assert gus[rep].element_size() == 1 and dns[rep].element_size() == 1, "block scales must be 1 byte" + I = int(gup[rep].shape[1] // 2) + H = int(gup[rep].shape[2] * 2) + assert gup[rep].shape[1] == 2 * I + assert H % 32 == 0 and I % 32 == 0, (H, I) # e8m0 scales per 32 along K + assert tuple(dnp[rep].shape[1:]) == (H, I // 2), (dnp[rep].shape, H, I) + assert tuple(gus[rep].shape[1:]) == (2 * I, H // 32), (gus[rep].shape, I, H) + assert tuple(dns[rep].shape[1:]) == (H, I // 32), (dns[rep].shape, H, I) + banks = self._tables(gate_up=gup, down=dnp, gate_up_scale=gus, down_scale=dns) + return banks, {}, (H, I) def decode( self, @@ -547,7 +548,7 @@ def decode_submit( topk_weights: torch.Tensor, topk_ids: torch.Tensor, ) -> tuple: - """Issue the D2H copies + the CPU-pool submit host node, then return without + """Issue the D2H copies + the CPU-pool submit host node (or the publish kernel), then return without waiting. Lets a caller (the hybrid backend) enqueue GPU work between this and :meth:`decode_sync` so the CPU compute overlaps the GPU GEMM / PCIe fetch. @@ -556,7 +557,14 @@ def decode_submit( :meth:`decode_sync`. The output tensor is allocated here so it stays live (and distinct from the interleaved GPU work) across the overlap window.""" bs = hidden_states.shape[0] - io = self._io_for(bs) + if bs > self.max_tokens: + raise ValueError(f"CPU MoE decode batch {bs} exceeds the executor's max_tokens {self.max_tokens}") + if not self._graph_capture_safe and torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "CPU MoE decode cannot be captured into a CUDA graph here: the stream-memop " + "doorbell is off or failed its startup probe, and the host-function sync is not " + "replay-safe on this runtime; run with --cuda-graph-max-bs 0" + ) if self._gpu_prequant: # DSV4: apply the reference FP8 round-trip on the GPU (the same kernel the @@ -564,93 +572,88 @@ 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) + + stream = torch.cuda.current_stream().cuda_stream + if self._kernel_io and hidden_states.dtype == torch.bfloat16: + x = hidden_states.contiguous() + ids = topk_ids.to(torch.int32).contiguous() + w = topk_weights.to(torch.float32).contiguous() + assert x.shape == (bs, self.H) and ids.shape == w.shape == (bs, self.top_k), (x.shape, ids.shape, w.shape) + self._ext.flag_publish(stream, x.data_ptr(), ids.data_ptr(), w.data_ptr(), layer_id, bs) + return (bs, torch.empty_like(x), True) # D2H: ship this step's activations + routing to pinned host memory. - io["x"].copy_(hidden_states, non_blocking=True) - io["ids"].copy_(topk_ids.to(torch.int32), non_blocking=True) - io["w"].copy_(topk_weights.to(torch.float32), non_blocking=True) + io = self._io + io["x"][:bs].copy_(hidden_states, non_blocking=True) + io["ids"][:bs].copy_(topk_ids.to(torch.int32), non_blocking=True) + io["w"][:bs].copy_(topk_weights.to(torch.float32), non_blocking=True) - task = self._task_for(layer_id, bs) out = torch.empty_like(hidden_states) - slot = self._flag_slots.get((layer_id, bs)) if self._flag_sync else None - if slot is not None: - # Front-end memops: done[slot]=0 then ready[slot]=1 (the coordinator's - # doorbell). No kernel launched; no host-func round trip. - self._cpu_moe.memop_submit( - torch.cuda.current_stream().cuda_stream, - self._done.data_ptr(), self._ready.data_ptr(), slot, - ) + if self._flag_sync: + # Front-end memops: reset done, then ring the doorbell with (layer, bs). + # No kernel launched; no host-func round trip. + self._ext.flag_signal(stream, layer_id, bs) else: - stream = torch.cuda.current_stream().cuda_stream - self._ext.submit_with_cuda_stream(stream, task) - return (bs, task, out, slot) + self._ext.submit_with_cuda_stream(stream, layer_id, bs) + return (bs, out, False) - def decode_sync(self, pending: tuple) -> torch.Tensor: + def decode_sync(self, pending: tuple, add_to: torch.Tensor | None = None) -> torch.Tensor: """Issue the CPU-pool sync + the H2D result copy for a prior :meth:`decode_submit`, and return the GPU output tensor. With flag-sync the wait is a front-end stream - memop on done[slot] (set by the CPU coordinator); otherwise a cudaLaunchHostFunc.""" - bs, task, out, slot = pending - if slot is not None: - # Front-end WAIT(done[slot] >= 1): blocks this stream's later nodes without + memop on done (set by the CPU coordinator); otherwise a cudaLaunchHostFunc. + With ``add_to``, return ``add_to`` + the CPU output; the kernel path adds in place into ``add_to``.""" + bs, out, kernel_io = pending + stream = torch.cuda.current_stream().cuda_stream + if self._flag_sync: + # Front-end WAIT(done >= 1): blocks this stream's later nodes without # occupying an SM, so GPU utilization stays truthful during the CPU window. - self._cpu_moe.memop_sync( - torch.cuda.current_stream().cuda_stream, self._done.data_ptr(), slot, - ) + self._ext.flag_wait(stream) else: - stream = torch.cuda.current_stream().cuda_stream - self._ext.sync_with_cuda_stream(stream, task) - io = self._io[bs] - out.copy_(io["y"], non_blocking=True) - return out - - def _watchdog_tick(self, suspects: dict) -> None: + self._ext.sync_with_cuda_stream(stream) + if kernel_io: + if add_to is not None and add_to.dtype == torch.bfloat16 and add_to.is_contiguous() and add_to.shape == out.shape: + self._ext.flag_combine(stream, add_to.data_ptr(), bs, True) + return add_to + self._ext.flag_combine(stream, out.data_ptr(), bs, False) + return out if add_to is None else add_to + out + out.copy_(self._io["y"][:bs], non_blocking=True) + return out if add_to is None else add_to + out + + def _watchdog_tick(self, suspect: dict) -> None: """One watchdog sampling round (called every 2 s by ``_watchdog_main``). - A slot is only declared dead when THREE things hold across >=10 s: its doorbell - is still pending (ready==1 && done==0), it was already pending when first - suspected, and the coordinator has served NOTHING on it since (flag_served_count - unchanged). The served-count criterion kills the false-positive window: two - point samples can land on the same slot's (independent, us-scale) pending - windows under heavy external load, but a coordinator that made progress in - between is alive by definition. ``suspects`` maps slot -> (first_seen, - served_at_first_sight) and persists across ticks.""" - stuck = (self._ready == 1) & (self._done == 0) - if not bool(stuck.any()): - suspects.clear() + The doorbell is only declared dead when THREE things hold across >=10 s: it is + still pending (ready set, done==0), it was already pending when first suspected, and + the coordinator has served NOTHING since (flag_served unchanged). The served-count + criterion kills the false-positive window: two point samples can land on + (independent, us-scale) pending windows under heavy external load, but a + coordinator that made progress in between is alive by definition. ``suspect`` + holds (first_seen, served_at_first_sight) and persists across ticks.""" + if not self._ext.flag_pending(): + suspect.clear() return now = time.monotonic() - pending = set(stuck.nonzero().flatten().tolist()) - for slot in list(suspects): - if slot not in pending: - del suspects[slot] - dead = [] - for slot in pending: - served = self._ext.flag_served_count(slot) - first_seen, served_then = suspects.get(slot, (None, None)) - if first_seen is None or served != served_then: - suspects[slot] = (now, served) # new suspect, or alive-but-loaded: rearm - continue - if now - first_seen >= 10.0: - dead.append(slot) - if not dead: + served = self._ext.flag_served() + if not suspect or suspect["served"] != served: + suspect.update(first_seen=now, served=served) # new suspect, or alive-but-loaded: rearm + return + if now - suspect["first_seen"] < 10.0: return logger.error( - f"cpu-moe flag watchdog: slots {dead} unanswered for >10s with no coordinator " - "progress (wedged/dead); poisoning done[] and failing the next step" + "cpu-moe flag watchdog: doorbell unanswered for >10s with no coordinator " + "progress (wedged/dead); poisoning done and failing the next step" ) - for i in dead: - self._err[i] = 1 - for i in dead: - self._done[i] = 1 # after err: unblock the stream into a checked failure - suspects.pop(i, None) + self._unhealthy = True + self._ext.poison_flag() # after the mark: unblock the stream into a checked failure + suspect.clear() def raise_if_unhealthy(self) -> None: """Raise if the flag watchdog fired (a doorbell stayed unanswered because the coordinator never responded). Called by the engine once per forward -- a single - pinned read -- so a dead coordinator surfaces as a loud error on the next step + flag read -- so a dead coordinator surfaces as a loud error on the next step instead of silently shipping stale expert outputs.""" - if self._err is not None and bool((self._err != 0).any()): + if self._unhealthy: raise RuntimeError( "CPU MoE flag-handshake watchdog fired: a decode step's doorbell was " "never answered by the coordinator thread (its outputs cannot be " @@ -663,13 +666,13 @@ def raise_if_unhealthy(self) -> None: def _watchdog_main(executor_ref) -> None: """Watchdog daemon body: weakref-deref per tick so the thread never keeps a dead executor alive (see the start site in ``CpuMoeExecutor.__init__``).""" - suspects: dict = {} + suspect: dict = {} while True: time.sleep(2.0) executor = executor_ref() - if executor is None or executor._watchdog_stop or executor._ready is None: + if executor is None or executor._watchdog_stop: return try: - executor._watchdog_tick(suspects) + executor._watchdog_tick(suspect) finally: del executor # drop the strong ref before the next sleep diff --git a/python/freetoken/moe/expert_banks.py b/python/freetoken/moe/expert_banks.py index e3a768b6d0..213b9fd2b5 100644 --- a/python/freetoken/moe/expert_banks.py +++ b/python/freetoken/moe/expert_banks.py @@ -2,8 +2,8 @@ The expert kernel (``QuantMethod.kernel``) owns the bank layout and the pack step; the checkpoint side delivers pieces (``moe.expert_pieces``) and this module fills the pinned host -banks from them (``build_expert_banks``). The GGUF q4_0 experts still -use their own providers until they get a method. +banks from them (``build_expert_banks``). GGUF experts keep their packed ggml rows +(``models.gguf.experts``). """ from __future__ import annotations @@ -11,15 +11,17 @@ import glob import math import os +from contextlib import contextmanager, nullcontext from dataclasses import dataclass, field import torch from freetoken.layers.quantization import QuantKind +from freetoken.memory import effective_memory_available from freetoken.utils import init_logger -from .host_banks import alloc_layer_banks -from .offload_cache import _BANK_BYTES_PER_EXPERT, _BANK_SCHEMAS +from .host_banks import HostBank +from .offload_cache import _BANK_BYTES_PER_EXPERT, placeholder_bank logger = init_logger(__name__) @@ -27,6 +29,19 @@ _PARALLEL_READER_SUPPORTED = hasattr(os, "O_DIRECT") and hasattr(os, "preadv") +@contextmanager +def _single_threaded_host_copies(): + """Keep small NVFP4 bank placements off PyTorch's intra-op worker pool.""" + previous = torch.get_num_threads() + if previous > 1: + torch.set_num_threads(1) + try: + yield + finally: + if previous > 1: + torch.set_num_threads(previous) + + @dataclass(frozen=True) class ExpertBanks: """Loaded expert banks, normalized for ``OffloadMoeCache`` wiring.""" @@ -48,6 +63,12 @@ class ExpertBanks: kind: QuantKind | None = None kernel: str | None = None layout: dict | None = None + # per-layer quant tag (e.g. the gate_up and down GGUF types), for layers whose banks match in dtype and row shape but not in quant; None = one quant + layer_tags: list | None = None + # per bank, per layer: bytes past an expert matrix's last row that the kernel reads and that must read as zeros (ggml MMQ); None = no such reads + row_tail_bytes: dict[str, list[int]] | None = None + # per-layer quant_format of experts in several formats (``quant_format``/``kind``/``kernel`` then name the most common one); None = every layer is ``quant_format`` + layer_formats: list[str] | None = None def _dummy_fill(role: str, tensor: torch.Tensor) -> None: @@ -75,6 +96,8 @@ def build_expert_banks( device: torch.device, layer_sink=None, dummy: bool = False, + resident: bool = False, + layers: list[int] | None = None, ) -> ExpertBanks: """Fill host banks in the kernel's layout from a stream of expert pieces. @@ -83,6 +106,9 @@ 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). + ``layers`` fills only those bank layers (the ones ``method`` serves, for experts of several + formats): pieces of other layers are skipped and the other layers' sources are None. """ from freetoken.moe.host_banks import LayerCompletionTracker, PinPipeline, pin_banks from freetoken.moe.legacy_format import legacy_format_for @@ -90,24 +116,34 @@ def build_expert_banks( kernel = method.kernel layout = method.layout() E = method.cfg.num_experts + ids = list(range(num_layers)) if layers is None else sorted(layers) + filled = set(ids) 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: {l: torch.empty(shape, dtype=dtype, device=device) for l in ids} for role, (shape, dtype) in specs.items()} + else: + hb = {role: {l: HostBank(shape, dtype) for l in ids} for role, (shape, dtype) in specs.items()} + banks = {role: {l: b.tensor for l, b in per_layer.items()} for role, per_layer in hb.items()} alphas = { role: torch.empty(num_layers * E, dtype=spec.dtype, device=device) for role, spec in layout.items() if spec.resident } + def _sources() -> dict[str, list]: + return {role: [per_layer.get(l) for l in range(num_layers)] for role, per_layer in banks.items()} + if dummy: for role, per_layer in banks.items(): - for tensor in per_layer: + for tensor in per_layer.values(): _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, + legacy_format_for(method.kind, kernel.name), _sources(), gate_up_alpha=alphas.get("gate_up_alpha"), down_alpha=alphas.get("down_alpha"), kind=method.kind, kernel=kernel.name, layout=layout, ) @@ -119,6 +155,8 @@ def _fill(sink) -> None: for layer_id, e0, e1, piece in pieces: if not (0 <= layer_id < num_layers and 0 <= e0 < e1 <= E): raise ValueError(f"expert piece out of range: layer {layer_id}, experts {e0}:{e1} of {num_layers} x {E}") + if layer_id not in filled: + continue # refuse before writing: a duplicate row would also complete the layer early and hand the sink a half-filled bank if written[layer_id, e0:e1].any(): raise ValueError(f"expert rows written more than once: layer {layer_id}, experts {e0}:{e1}") @@ -130,94 +168,120 @@ def _fill(sink) -> None: if tracker is not None: for _ in range(e1 - e0): tracker.note(layer_id) - missing = (written == 0).nonzero().tolist() + missing = (written[ids] == 0).nonzero().tolist() if missing: + missing = [[ids[l], e] for l, e in missing] raise ValueError(f"expert banks were not filled: {len(missing)} (layer, expert) rows missing (first {missing[:4]})") - if layer_sink is not None: - _fill(layer_sink) - elif torch.cuda.is_available(): - with PinPipeline() as pins: - _fill(pins) - else: - _fill(None) + # Per-expert NVFP4 pieces produce thousands of small CPU copies. On high-core + # hosts the intra-op pool costs more than the copies and contends with the + # checkpoint reader. Restore its setting even if packing or pinning fails. + copy_threads = ( + _single_threaded_host_copies() + if method.kind is QuantKind.NVFP4 and not resident + else nullcontext() + ) + with copy_threads: + if layer_sink is not None: + _fill(layer_sink) + elif torch.cuda.is_available() and not resident: + with PinPipeline() as pins: + _fill(pins) + else: + _fill(None) return ExpertBanks( - legacy_format_for(method.kind, kernel.name), banks, + legacy_format_for(method.kind, kernel.name), _sources(), gate_up_alpha=alphas.get("gate_up_alpha"), down_alpha=alphas.get("down_alpha"), streamed=layer_sink is not None, kind=method.kind, kernel=kernel.name, layout=layout, ) -_PARALLEL_CHUNK = 8 << 20 # default O_DIRECT chunk for the parallel reader +def method_groups(layer_methods) -> list[tuple]: + """``[(method, [bank layers])]`` of the layers' expert methods, one entry per (kind, kernel), most layers first.""" + groups: dict = {} + for layer, method in enumerate(layer_methods): + groups.setdefault((method.kind, method.kernel.name), (method, []))[1].append(layer) + return sorted(groups.values(), key=lambda g: (-len(g[1]), g[1][0])) -def _q4_0_banks(model_path, model_config, device, dtype, dummy, parallel=False, workers=8, chunk=_PARALLEL_CHUNK, decode_target="gpu", layer_sink=None) -> ExpertBanks: - if parallel: - raise NotImplementedError( - "parallel reader not implemented for q4_0: GGUF is a single packed file " - "(not safetensors), so the common reader doesn't apply -- it needs a GGUF-native " - "parallel reader (parse the tensor table, chunked O_DIRECT over the one file)" - ) - from freetoken.models.weight import load_q4_0_moe_expert_sources - - # Native GGUF Q4_0 routed experts: packed block bytes streamed to the GPU and - # dequantized inside the borrowed ggml MoE kernels (no bf16 expert copy). Banks are - # per-layer HostBanks (pin-after-fill), so conversion streams each completed layer's - # gate_up + down straight through the sink (dummy fabricates in one shot -> not streamed). - sink = None if dummy else layer_sink - sources = load_q4_0_moe_expert_sources(model_path, model_config, dummy=dummy, layer_sink=sink) +def is_mixed(layer_methods) -> bool: + return layer_methods is not None and len(method_groups(layer_methods)) > 1 + + +def merge_expert_banks(parts: list[ExpertBanks], layer_methods, num_experts: int) -> ExpertBanks: + """One ``ExpertBanks`` over the per-format ``parts`` (each filled only on its layers): the banks are the + union of the roles, a layer without a role gets a zero-byte placeholder (see ``OffloadMoeCache.layer_layouts``).""" + from freetoken.moe.legacy_format import legacy_format_for + + num_layers = len(layer_methods) + sources: dict[str, list] = {} + for part in parts: + if part.gate_up_alpha is not None or part.down_alpha is not None: + raise ValueError(f"experts of several formats cannot use the {part.kernel} kernel, which keeps GPU-resident alphas") + for role, per_layer in part.sources.items(): + merged = sources.setdefault(role, [None] * num_layers) + for layer, tensor in enumerate(per_layer): + if tensor is not None: + assert merged[layer] is None, (role, layer) + merged[layer] = tensor + for per_layer in sources.values(): + for layer, tensor in enumerate(per_layer): + if tensor is None: + per_layer[layer] = placeholder_bank(num_experts) + formats = [legacy_format_for(m.kind, m.kernel.name) for m in layer_methods] + primary = parts[0] return ExpertBanks( - "q4_0", {name: sources[name] for name in _BANK_SCHEMAS["q4_0"]}, streamed=sink is not None + primary.quant_format, sources, streamed=all(p.streamed for p in parts), + kind=primary.kind, kernel=primary.kernel, layer_tags=formats, layer_formats=formats, ) -# expert formats that still load through their own provider (GGUF) -_PROVIDERS = { - "q4_0": _q4_0_banks, -} - +def _mixed_expert_banks(model_path, model_config, layer_methods, device, dummy, parallel, workers, chunk, layer_sink=None) -> ExpertBanks: + from freetoken.moe.expert_pieces import iter_expert_pieces -def _legacy_expert_banks(model_path, model_config, device, dtype, dummy, parallel, workers, chunk, decode_target="gpu", layer_sink=None) -> ExpertBanks: - expert_quant = model_config.expert_quant - if expert_quant not in _PROVIDERS: - raise ValueError( - f"{expert_quant!r} experts load through their MoE quant method; " - f"only {sorted(_PROVIDERS)} still have a format provider" + num_layers = model_config.num_moe_layers + assert len(layer_methods) == num_layers, (len(layer_methods), num_layers) + groups = method_groups(layer_methods) + # resolve every format's reader before any bank exists, so a missing parallel reader falls back to serial with nothing allocated + readers = [ + None if dummy else iter_expert_pieces( + model_path, model_config, method.kind, layers=layers, parallel=parallel, workers=workers, chunk=chunk ) - return _PROVIDERS[expert_quant]( - model_path, model_config, device, dtype, dummy, - parallel=parallel, workers=workers, chunk=chunk, decode_target=decode_target, - layer_sink=layer_sink, - ) + for method, layers in groups + ] + parts = [] + for (method, layers), pieces in zip(groups, readers): + parts.append(build_expert_banks( + method, num_layers, pieces, device=device, layer_sink=layer_sink, dummy=dummy, layers=layers, + )) + logger.info_rank0(f"expert banks: {method.kind} / {method.kernel.name} on {len(layers)} of {num_layers} MoE layers") + return merge_expert_banks(parts, layer_methods, model_config.num_experts) + +_PARALLEL_CHUNK = 8 << 20 # default O_DIRECT chunk for the parallel reader -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: """Best-effort: can free host RAM hold the expert banks plus the parallel reader's one extra (non-reclaimable) whole-shard buffer? Unknown (non-local path / no /proc) -> True, i.e. keep the fast path. Banks ~= checkpoint size (experts dominate); transient ~= the - largest shard. Uses MemAvailable (counts reclaimable cache) -- the OOM-relevant figure.""" - avail = None - try: - with open("/proc/meminfo") as f: - for line in f: - if line.startswith("MemAvailable:"): - avail = int(line.split()[1]) * 1024 - break - except OSError: - pass + largest shard. Uses the effective available memory (host ``MemAvailable``, which counts + reclaimable cache, clamped by the process's tightest finite cgroup budget) -- the + OOM-relevant figure on bare metal and inside a ``--memory``-limited container alike.""" + avail = effective_memory_available() if avail is None: return True try: # resolve a hub id to its local cache dir (no-op for a local path) so glob sees the shards @@ -245,19 +309,24 @@ def ftw_bank_bytes(model_path: str) -> int | None: return sum(t["nbytes"] for t in tensors if t.get("kind") == "experts_bank") -def bank_bytes_estimate(model_config, method=None) -> int | None: +def _method_expert_bytes(method) -> int: + return method.cfg.num_experts * sum( + math.prod(spec.shape) * torch.empty((), dtype=spec.dtype).element_size() + for spec in method.layout().values() if not spec.resident + ) + + +def bank_bytes_estimate(model_config, method=None, layer_methods=None) -> int | None: """Estimated total expert-bank bytes of a raw checkpoint before loading it. - With a bound expert ``method`` the kernel's layout gives the exact host bytes; otherwise the - format-tag table sizes the GGUF format. ``None`` for unknown formats or missing dims - (callers then skip the pre-load sizing).""" + With a bound expert ``method`` the kernel's layout gives the exact host bytes (``layer_methods``, + one per MoE layer, for experts of several formats); otherwise the format-tag table sizes the + format. ``None`` for unknown formats or missing dims (callers then skip the pre-load sizing).""" layers = getattr(model_config, "num_moe_layers", None) + if layer_methods is not None: + return sum(_method_expert_bytes(m) for m in layer_methods) if method is not None and layers: - per_expert = sum( - math.prod(spec.shape) * torch.empty((), dtype=spec.dtype).element_size() - for spec in method.layout().values() if not spec.resident - ) - return layers * method.cfg.num_experts * per_expert + return layers * _method_expert_bytes(method) expert_quant = getattr(model_config, "expert_quant", "none") fmt = expert_quant if expert_quant != "none" else ( getattr(model_config, "moe_weight_format", None) or "bf16" @@ -278,14 +347,14 @@ def load_expert_banks( *, method=None, device: torch.device, - dtype: torch.dtype, dummy: bool = False, parallel: bool | None = None, workers: int = 8, chunk: int = _PARALLEL_CHUNK, - decode_target: str = "gpu", layer_sink=None, layer_residency: list[str] | None = None, + resident: bool = False, + layer_methods: list | None = None, ) -> ExpertBanks: """Load (or fabricate, with ``dummy=True``) the expert banks. Two paths, both returning the same normalized ``ExpertBanks`` and both pinning after fill: @@ -305,19 +374,32 @@ def load_expert_banks( streamable formats, so callers must check ``streamed`` rather than assume it fired. ``method`` (the bound expert quant method of the model's offload layers) selects - the generic path: the family's pieces packed by the method's kernel. Without it only the - GGUF q4_0 format loads, through its own provider. + the generic path: the family's pieces packed by the method's kernel. Without it only an + FTW checkpoint loads. ``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. + + ``layer_methods`` (one bound method per MoE layer) serves experts of several formats: each + format fills its own layers and the banks merge over the union of roles (``merge_expert_banks``). """ - 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 + mixed = is_mixed(layer_methods) + if mixed and (resident or layer_residency is not None): + raise ValueError("experts of several formats are served from the offload slot cache only (no resident experts, no --moe-cpu-layers)") + 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 @@ -353,9 +435,16 @@ def load_expert_banks( from freetoken.moe.host_banks import requested_residency 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 _legacy_expert_banks(model_path, model_config, device, dtype, dummy, par, workers, chunk, decode_target, layer_sink) + if mixed: + return _mixed_expert_banks(model_path, model_config, layer_methods, device, dummy, par, workers, chunk, layer_sink) + if method is not None and method.kind is QuantKind.GGUF: + from freetoken.models.gguf.experts import load_gguf_expert_banks + + # each layer keeps its own ggml types, so the banks come from the file, not from one method's layout + return load_gguf_expert_banks(model_path, model_config.num_moe_layers, first_layer=model_config.first_k_dense_replace, dummy=dummy, device=device if resident else None) + if method is None: + raise ValueError("the experts of a raw checkpoint are packed by the offload layers' MoE quant method; pass method=") + return _method_expert_banks(model_path, model_config, method, device, dummy, par, workers, chunk, layer_sink, resident) with requested_residency(layer_residency) as residency_plan: try: @@ -393,3 +482,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/expert_pieces.py b/python/freetoken/moe/expert_pieces.py index 79f655a951..74c4a41b0f 100644 --- a/python/freetoken/moe/expert_pieces.py +++ b/python/freetoken/moe/expert_pieces.py @@ -8,7 +8,7 @@ from __future__ import annotations -from typing import Callable, Iterable, Iterator +from typing import Callable, Collection, Iterable, Iterator import torch from freetoken.layers.quantization import QuantKind @@ -38,24 +38,34 @@ def _model_hook(spec, name: str): def iter_expert_pieces( - model_path: str, config, kind: QuantKind, *, parallel: bool = False, workers: int = 8, chunk: int = 8 << 20 + model_path: str, + config, + kind: QuantKind, + *, + layers: Collection[int] | None = None, + parallel: bool = False, + workers: int = 8, + chunk: int = 8 << 20, ) -> Iterator[Piece]: """The pieces of ``model_path``'s routed experts, stored as ``kind``. - A family may ship its own ``iter_expert_pieces(model_path, config, kind, *, parallel, workers, - chunk)`` for the storage forms only it uses (it returns None for the rest); otherwise bf16 + A family may ship its own ``iter_expert_pieces(model_path, config, kind, *, layers, parallel, + workers, chunk)`` for the storage forms only it uses (it returns None for the rest); otherwise bf16 experts come from the family's stacked ``iter_weights`` and NVFP4 experts from its ``nvfp4_expert_spec``. The reader is resolved here, before any bank is allocated, so a missing parallel reader raises ``NotImplementedError`` while a serial fallback is still cheap. + + ``layers`` (bank layer ids, default all) limits the pieces to those layers, for experts of several + formats: a reader skips the other layers' tensors and checks only these layers for completeness. """ spec = get_model_spec(config.architectures[0]) hook = _model_hook(spec, "iter_expert_pieces") if hook is not None: - pieces = hook(model_path, config, kind, parallel=parallel, workers=workers, chunk=chunk) + pieces = hook(model_path, config, kind, layers=layers, parallel=parallel, workers=workers, chunk=chunk) if pieces is not None: return pieces if kind is QuantKind.NONE: - return _bf16_pieces(model_path, config, spec, parallel=parallel, workers=workers, chunk=chunk) + return _bf16_pieces(model_path, config, spec, layers=layers, parallel=parallel, workers=workers, chunk=chunk) if kind is QuantKind.NVFP4: spec_hook = _model_hook(spec, "nvfp4_expert_spec") if spec_hook is None: @@ -63,7 +73,7 @@ def iter_expert_pieces( from freetoken.models.nvfp4_banks import iter_nvfp4_expert_pieces return iter_nvfp4_expert_pieces( - model_path, config, spec_hook(model_path, config), parallel=parallel, workers=workers, chunk=chunk + model_path, config, spec_hook(model_path, config), layers=layers, parallel=parallel, workers=workers, chunk=chunk ) raise NotImplementedError(f"{spec.module} provides no expert reader for {kind!r} experts") @@ -81,18 +91,31 @@ def packed_expert_source_info(key: str) -> tuple[int, str] | None: return None -def stacked_expert_pieces(tensors: Iterable[tuple[str, torch.Tensor]], config) -> Iterator[Piece]: - """Whole-layer ``experts.gate_up_proj`` / ``experts.down_proj`` tensors -> one piece per layer.""" +def _source_layer(key: str) -> int | None: + """Checkpoint layer of a ``model.layers.N....`` key.""" + parts = key.split(".") + return int(parts[2]) if len(parts) > 2 and parts[:2] == ["model", "layers"] and parts[2].isdigit() else None + + +def stacked_expert_pieces(tensors: Iterable[tuple[str, torch.Tensor]], config, *, layers: Collection[int] | None = None) -> Iterator[Piece]: + """Whole-layer ``experts.gate_up_proj`` / ``experts.down_proj`` tensors -> one piece per layer, for ``layers`` (bank layer ids, default all).""" + keep = None if layers is None else set(layers) pending: dict[int, dict[str, torch.Tensor]] = {} num_experts = config.num_experts for name, tensor in tensors: info = packed_expert_source_info(name) if info is None: + # another format's layer stores its experts in another form + layer = _source_layer(name) + if keep is not None and layer is not None and bank_layer_of(config, layer) not in keep: + continue raise ValueError(f"Unexpected expert weight key: {name}") layer, packed_name = info bank_layer = bank_layer_of(config, layer) if bank_layer is None: raise ValueError(f"Unexpected MoE expert layer {layer}; expected a routed-expert layer") + if keep is not None and bank_layer not in keep: + continue if tensor.size(0) != num_experts: raise ValueError(f"Unexpected {packed_name} expert count {tensor.size(0)}; expected {num_experts}") piece = pending.setdefault(bank_layer, {}) @@ -104,7 +127,7 @@ def stacked_expert_pieces(tensors: Iterable[tuple[str, torch.Tensor]], config) - raise ValueError(f"Missing MoE expert source layers: {sorted(pending)}") -def _bf16_pieces(model_path: str, config, spec, *, parallel: bool, workers: int, chunk: int) -> Iterator[Piece]: +def _bf16_pieces(model_path: str, config, spec, *, layers: Collection[int] | None, parallel: bool, workers: int, chunk: int) -> Iterator[Piece]: device = torch.device("cpu") if parallel: iter_weights = _model_hook(spec, "iter_weights_parallel") @@ -116,7 +139,7 @@ def _bf16_pieces(model_path: str, config, spec, *, parallel: bool, workers: int, else: iter_weights = _load_attr(spec.module, spec.iter_weights) tensors = iter_weights(model_path, device, include_moe_experts=True, include_non_moe=False) - return stacked_expert_pieces(tensors, config) + return stacked_expert_pieces(tensors, config, layers=layers) def per_expert_pieces( @@ -128,7 +151,7 @@ def per_expert_pieces( """Group per-expert tensors into one piece per expert, in whatever order they arrive. ``locate(name)`` -> ``(bank_layer, expert, role)`` or None to skip; a role tensor is one - expert's ``[...]`` row and becomes the piece's ``[1, ...]`` entry. Globals may be scalars. + expert's ``[...]`` row and becomes the piece's ``[1, ...]`` entry. Globals and activation scales may be scalars. """ pending: dict[tuple[int, int], dict[str, torch.Tensor]] = {} for name, tensor in tensors: @@ -137,7 +160,7 @@ def per_expert_pieces( continue bank_layer, expert, role = where piece = pending.setdefault((bank_layer, expert), {}) - piece[role] = tensor.reshape(1, -1) if role.endswith("_global") else tensor.unsqueeze(0) + piece[role] = tensor.reshape(1, -1) if role.endswith(("_global", "_input")) else tensor.unsqueeze(0) if len(piece) == tensors_per_expert: del pending[(bank_layer, expert)] yield bank_layer, expert, expert + 1, piece diff --git a/python/freetoken/moe/fused.py b/python/freetoken/moe/fused.py index a90aa4c792..2f8af47c89 100644 --- a/python/freetoken/moe/fused.py +++ b/python/freetoken/moe/fused.py @@ -44,6 +44,9 @@ def fused_topk( return fused_topk_softmax(gating_output, topk, renormalize, num_token_non_padded) +_SGL_ALIGN_MAX_EXPERTS = 1024 + + def moe_align_block_size( topk_ids: torch.Tensor, block_size: int, num_experts: int ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: @@ -84,9 +87,23 @@ def moe_align_block_size( - The padding ensures that the total number of tokens is now divisible by block_size for proper block matrix operations. """ - from freetoken.kernel.backend import is_sgl_kernel_installed + from freetoken.kernel.backend import is_rocm, is_sgl_kernel_installed + + if is_rocm(): + import triton + from packaging.version import Version + + # FIXME: remove once Triton 3.8 is unsupported. Its AMD backend folds compares on tl.histogram results to false. + # https://github.com/triton-lang/triton/pull/11246 + if Version(triton.__version__).release[:2] == (3, 8): + from freetoken.kernel import moe_align_block_size_triton + + return moe_align_block_size_triton(topk_ids, block_size, num_experts) - if not is_sgl_kernel_installed(): + # sgl_kernel's align kernel takes num_experts + 1 counters in one block and reads out of + # bounds past 1024 of them (0.4.5 and 0.4.9). A prefill over an offload slot cache passes the + # slot count here (#676), which is thousands; the Triton kernel has no such limit. + if not is_sgl_kernel_installed() or num_experts + 1 > _SGL_ALIGN_MAX_EXPERTS: from freetoken.kernel.triton.moe_align import ( moe_align_block_size as triton_moe_align_block_size, ) diff --git a/python/freetoken/moe/fused_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..6bedd57b12 100644 --- a/python/freetoken/moe/fused_nvfp4.py +++ b/python/freetoken/moe/fused_nvfp4.py @@ -7,134 +7,13 @@ from __future__ import annotations -from typing import Any, Dict - import torch -import triton -import triton.language as tl -from freetoken.kernel import moe_sum_reduce_triton -from freetoken.kernel.triton.e4m3_compat import e4m3_kernel_view -from freetoken.kernel.triton.nvfp4_fused_moe import ( - _decode_nvfp4_marlin_kernel, - _decode_nvfp4_moe_kernel, - _e2m1_lut, - _prefill_nvfp4_moe_kernel, -) +from freetoken.kernel import moe_sum_reduce_triton, nvfp4_moe +from freetoken.kernel.triton.nvfp4_act_quant import quantize_nvfp4 from freetoken.layers import gated_act_and_mul from freetoken.moe.fused import moe_align_block_size -# Decode is captured into a CUDA graph, so the config must be fixed (no triton.autotune, -# which benchmarks at run time). Tuned offline against the NVFP4 decode kernels. -# These drive the original LUT-gather decode (_decode_gemm), kept only for A/B. -_DECODE_BLOCK_N = 64 -_DECODE_BLOCK_KB = 128 -_DECODE_WARPS = 4 - -# Marlin-style decode config (int32 wide loads + deferred reduction). Offline sweep over -# the qwen35/qwen3moe (I=512/768) decode shapes picked BLOCK_N=16, BLOCK_KW=16 (== 128 -# k-values/iter), 4 warps -- the wide load lifts the gate/up GEMM ~43%->~51% of peak BW. -_DECODE_MARLIN_BLOCK_N = 16 -_DECODE_MARLIN_BLOCK_KW = 16 -_DECODE_MARLIN_WARPS = 4 -# Deep-K variant: at K > 2048 (qwen4_exp gate_up, K=2560) a narrower N tile with the whole -# K strip in one program iteration measures ~13% faster (18.6 vs 21.0us); short-K shapes -# regress under it, so the split is by K, not by gemm position. -_DECODE_MARLIN_DEEPK_BLOCK_N = 8 -_DECODE_MARLIN_DEEPK_BLOCK_KW = 128 -_DECODE_MARLIN_DEEPK_THRESHOLD = 2048 - - -def _tl_dtype(dt: torch.dtype): - if dt == torch.bfloat16: - return tl.bfloat16 - if dt == torch.float16: - return tl.float16 - return tl.float32 - - -def _decode_gemm( - a: torch.Tensor, - packed: torch.Tensor, - scale: torch.Tensor, - glob: torch.Tensor, - c: torch.Tensor, - topk_weights: torch.Tensor, - topk_ids: torch.Tensor, - mul_routed_weight: bool, - a_row_is_route: bool, -) -> None: - M, top_k = topk_ids.shape - N = packed.shape[1] - K = packed.shape[2] * 2 - scale = e4m3_kernel_view(scale) - total_routes = M * top_k - grid = (total_routes, triton.cdiv(N, _DECODE_BLOCK_N)) - _decode_nvfp4_moe_kernel[grid]( - a, packed, scale, glob, c, topk_weights, topk_ids, - _e2m1_lut(a.device.index), - total_routes, N, K, - a.stride(0), a.stride(1), - packed.stride(0), packed.stride(1), packed.stride(2), - scale.stride(0), scale.stride(1), scale.stride(2), - glob.stride(0), glob.stride(1), - c.stride(0), c.stride(1), c.stride(2), - topk_weights.stride(0), topk_weights.stride(1), - topk_ids.stride(0), topk_ids.stride(1), - BLOCK_SIZE_N=_DECODE_BLOCK_N, - BLOCK_SIZE_KB=_DECODE_BLOCK_KB, - TOP_K=top_k, - A_ROW_IS_ROUTE=a_row_is_route, - MUL_ROUTED_WEIGHT=mul_routed_weight, - compute_type=_tl_dtype(c.dtype), - num_warps=_DECODE_WARPS, - ) - - -def _decode_gemm_marlin( - a: torch.Tensor, - packed: torch.Tensor, - scale: torch.Tensor, - glob: torch.Tensor, - c: torch.Tensor, - topk_weights: torch.Tensor, - topk_ids: torch.Tensor, - mul_routed_weight: bool, - a_row_is_route: bool, -) -> None: - """Marlin-style decode GEMV: int32 wide loads + deferred reduction - (:func:`_decode_nvfp4_marlin_kernel`). ``packed`` is the uint8 ``[S, N, K//2]`` bank; - it is reinterpreted as int32 ``[S, N, K//8]`` (contiguous, K%8==0 for NVFP4).""" - M, top_k = topk_ids.shape - N = packed.shape[1] - K = packed.shape[2] * 2 - packed_i32 = packed.view(torch.int32) # [S, N, K // 8] - scale = e4m3_kernel_view(scale) - total_routes = M * top_k - deep_k = K > _DECODE_MARLIN_DEEPK_THRESHOLD - block_n = _DECODE_MARLIN_DEEPK_BLOCK_N if deep_k else _DECODE_MARLIN_BLOCK_N - block_kw = _DECODE_MARLIN_DEEPK_BLOCK_KW if deep_k else _DECODE_MARLIN_BLOCK_KW - grid = (total_routes, triton.cdiv(N, block_n)) - _decode_nvfp4_marlin_kernel[grid]( - a, packed_i32, scale, glob, c, topk_weights, topk_ids, - _e2m1_lut(a.device.index), - total_routes, N, K, - a.stride(0), a.stride(1), - packed_i32.stride(0), packed_i32.stride(1), packed_i32.stride(2), - scale.stride(0), scale.stride(1), scale.stride(2), - glob.stride(0), glob.stride(1), - c.stride(0), c.stride(1), c.stride(2), - topk_weights.stride(0), topk_weights.stride(1), - topk_ids.stride(0), topk_ids.stride(1), - BLOCK_SIZE_N=block_n, - BLOCK_SIZE_KW=block_kw, - TOP_K=top_k, - A_ROW_IS_ROUTE=a_row_is_route, - MUL_ROUTED_WEIGHT=mul_routed_weight, - compute_type=_tl_dtype(c.dtype), - num_warps=_DECODE_MARLIN_WARPS, - ) - def _fused_experts_decode_nvfp4( gemm_fn, @@ -153,8 +32,8 @@ def _fused_experts_decode_nvfp4( act_limit: float = 7.0, ) -> torch.Tensor: """Shared decode body (gemm1 -> act -> gemm2 -> sum-reduce); ``gemm_fn`` is either - the marlin-style int32 GEMV (:func:`_decode_gemm_marlin`) or the original LUT-gather - GEMV (:func:`_decode_gemm`), both with the same calling convention.""" + the marlin-style int32 GEMV (:func:`freetoken.kernel.nvfp4_moe.decode_gemm_marlin`) or the original LUT-gather + GEMV (:func:`freetoken.kernel.nvfp4_moe.decode_gemm`), both with the same calling convention.""" M, H = hidden_states.shape top_k = topk_ids.shape[1] two_i = gate_up_packed.shape[1] @@ -195,7 +74,7 @@ def fused_experts_decode_nvfp4_marlin( ) -> torch.Tensor: """Decode inline-NVFP4 MoE using the Marlin-style int32 wide-load GEMV.""" return _fused_experts_decode_nvfp4( - _decode_gemm_marlin, + nvfp4_moe.decode_gemm_marlin, hidden_states, gate_up_packed, gate_up_scale, gate_up_global, down_packed, down_scale, down_global, topk_weights, topk_ids, activation, apply_router_weight_on_input, @@ -221,7 +100,7 @@ def fused_experts_decode_nvfp4_serial( """Original LUT-gather decode (one program per route, full K reduction). Retained for A/B benchmarking against the marlin decode path; not on the production decode path.""" return _fused_experts_decode_nvfp4( - _decode_gemm, + nvfp4_moe.decode_gemm, hidden_states, gate_up_packed, gate_up_scale, gate_up_global, down_packed, down_scale, down_global, topk_weights, topk_ids, activation, apply_router_weight_on_input, @@ -229,57 +108,6 @@ def fused_experts_decode_nvfp4_serial( ) -def _prefill_config(M: int) -> Dict[str, int]: - # ``BLOCK_SIZE_M`` is coupled to host-side ``moe_align_block_size`` (token padding), - # so it cannot be picked by triton.autotune; these were chosen by an offline sweep - # over (BLOCK_M, BLOCK_N, BLOCK_KB, num_warps, num_stages) for the MiniMax-M2 shapes. - if M <= 64: - return dict(BLOCK_SIZE_M=16, BLOCK_SIZE_N=64, BLOCK_SIZE_KB=32, - GROUP_SIZE_M=1, num_warps=8, num_stages=4) - return dict(BLOCK_SIZE_M=32, BLOCK_SIZE_N=64, BLOCK_SIZE_KB=32, - GROUP_SIZE_M=8, num_warps=8, num_stages=4) - - -def _prefill_gemm( - a: torch.Tensor, - packed: torch.Tensor, - scale: torch.Tensor, - glob: torch.Tensor, - c: torch.Tensor, - topk_weights_flat: torch.Tensor, - sorted_ids: torch.Tensor, - expert_ids: torch.Tensor, - num_tokens_post_padded: torch.Tensor, - num_valid_tokens: int, - kernel_top_k: int, - mul_routed_weight: bool, - cfg: Dict[str, Any], -) -> None: - N = packed.shape[1] - K = packed.shape[2] * 2 - EM = sorted_ids.shape[0] - scale = e4m3_kernel_view(scale) - grid = lambda META: ( # noqa: E731 - triton.cdiv(EM, META["BLOCK_SIZE_M"]) * triton.cdiv(N, META["BLOCK_SIZE_N"]), - ) - _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), - scale.stride(0), scale.stride(1), scale.stride(2), - glob.stride(0), glob.stride(1), - c.stride(1), c.stride(2), - topk_weights_flat.stride(0), - MUL_ROUTED_WEIGHT=mul_routed_weight, - top_k=kernel_top_k, - compute_type=_tl_dtype(c.dtype), - **cfg, - ) - - def fused_experts_nvfp4( hidden_states: torch.Tensor, gate_up_packed: torch.Tensor, @@ -304,14 +132,14 @@ def fused_experts_nvfp4( two_i = gate_up_packed.shape[1] inter = two_i // 2 dev, dt = hidden_states.device, hidden_states.dtype - cfg = _prefill_config(M) + cfg = nvfp4_moe.prefill_config(M, top_k, num_experts, hidden_states.device.index) sorted_ids, expert_ids, ntpp = moe_align_block_size(topk_ids, cfg["BLOCK_SIZE_M"], num_experts) tw = topk_weights.reshape(-1).contiguous() num_valid = topk_ids.numel() ic1 = torch.empty((M, top_k, two_i), device=dev, dtype=dt) - _prefill_gemm( + nvfp4_moe.prefill_gemm( hidden_states, gate_up_packed, gate_up_scale, gate_up_global, ic1, tw, sorted_ids, expert_ids, ntpp, num_valid, top_k, apply_router_weight_on_input, cfg, @@ -319,7 +147,7 @@ def fused_experts_nvfp4( ic2 = torch.empty((M * top_k, inter), device=dev, dtype=dt) gated_act_and_mul(activation, ic1.view(-1, two_i), ic2, alpha=act_alpha, limit=act_limit) ic3 = torch.empty((M, top_k, H), device=dev, dtype=dt) - _prefill_gemm( + nvfp4_moe.prefill_gemm( ic2, down_packed, down_scale, down_global, ic3, tw, sorted_ids, expert_ids, ntpp, num_valid, 1, not apply_router_weight_on_input, cfg, @@ -329,8 +157,60 @@ def fused_experts_nvfp4( return out +def fused_experts_nvfp4_w4a4( + hidden_states: torch.Tensor, + gate_up_packed: torch.Tensor, + gate_up_scale: torch.Tensor, + gate_up_global: torch.Tensor, + down_packed: torch.Tensor, + down_scale: torch.Tensor, + down_global: torch.Tensor, + gate_up_input_scale: torch.Tensor, + down_input_scale: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + num_experts: int, + activation: str = "silu", + act_alpha: float = 1.702, + act_limit: float = 7.0, +) -> torch.Tensor: + """Prefill NVFP4 MoE with NVFP4 activations (W4A4) where the GPU and the batch favour it, else :func:`fused_experts_nvfp4`. + + ``*_input_scale`` are the per-expert static activation scales of the ``num_experts`` banks; each GEMM input is quantized under their max, as vLLM does. + """ + M, H = hidden_states.shape + top_k = topk_ids.shape[1] + inter = gate_up_packed.shape[1] // 2 + cfg = nvfp4_moe.w4a4_prefill_config(M, top_k, num_experts, H, inter, activation, hidden_states.dtype, hidden_states.device.index) + if cfg is None: + return fused_experts_nvfp4( + hidden_states, gate_up_packed, gate_up_scale, gate_up_global, down_packed, down_scale, down_global, + topk_weights, topk_ids, num_experts, activation, False, act_alpha, act_limit, + ) + dev = hidden_states.device + g = torch.stack((gate_up_input_scale[:num_experts], down_input_scale[:num_experts])).float().amax(dim=1) + sorted_ids, expert_ids, ntpp = moe_align_block_size(topk_ids, cfg["BLOCK_SIZE_M"], num_experts) + tw = topk_weights.reshape(-1).contiguous() + num_valid = topk_ids.numel() + + xq, xs = quantize_nvfp4(hidden_states, g[0:1]) + q = torch.empty((num_valid, inter // 2), device=dev, dtype=torch.uint8) + qs = torch.empty((num_valid, inter // 16), device=dev, dtype=torch.float8_e4m3fn) + nvfp4_moe.w4a4_gate_up_gemm( + xq, xs, g[0:1], gate_up_packed, gate_up_scale, gate_up_global, q, qs, g[1:2], sorted_ids, expert_ids, ntpp, + num_valid, top_k, activation, act_alpha, act_limit, cfg["gate_up"], + ) + out = torch.zeros_like(hidden_states) + nvfp4_moe.w4a4_down_gemm( + q, qs, g[1:2], down_packed, down_scale, down_global, out, tw, sorted_ids, expert_ids, ntpp, num_valid, top_k, + cfg["down"], + ) + return out + + __all__ = [ "fused_experts_decode_nvfp4_marlin", "fused_experts_decode_nvfp4_serial", "fused_experts_nvfp4", + "fused_experts_nvfp4_w4a4", ] diff --git a/python/freetoken/moe/fused_q4_0.py b/python/freetoken/moe/fused_q4_0.py deleted file mode 100644 index cdab82bf99..0000000000 --- a/python/freetoken/moe/fused_q4_0.py +++ /dev/null @@ -1,53 +0,0 @@ -"""Grouped expert GEMM over native GGUF Q4_0 banks (borrowed ggml MoE kernels). - -Ports vLLM/sglang's ``_fused_moe_gguf`` MMVQ path onto FreeToken's offload-cache -interface: the experts are streamed to the GPU as packed Q4_0 block bytes and -dequantized *inside* ``ggml_moe_a8_vec`` -- no bf16 expert copy is materialized. We -use the MMVQ (vector) kernel for both prefill and decode: it consumes ``topk_ids`` -directly (no ``moe_align_block_size`` needed) and on small batches it is the right -choice anyway. ``topk_ids`` already index the streamed cache slots (decode) or the -materialized layer positions (prefill). -""" - -from __future__ import annotations - -import torch - -from freetoken.layers.activation import gelu_and_mul, gelu_tanh_and_mul, silu_and_mul -from freetoken.models.gguf.dequant import GGML_Q4_0 - -_ACT = {"silu": silu_and_mul, "gelu": gelu_and_mul, "gelu_tanh": gelu_tanh_and_mul} - - -def fused_experts_gguf_q4_0( - hidden_states: torch.Tensor, - gate_up_q: torch.Tensor, # [num_slots, 2I, H//32*18] uint8 - down_q: torch.Tensor, # [num_slots, H, I//32*18] uint8 - topk_weights: torch.Tensor, - topk_ids: torch.Tensor, - activation: str, -) -> torch.Tensor: - from freetoken.kernel.gguf import ggml_moe_a8_vec - - act_fn = _ACT.get(activation) - if act_fn is None: - raise ValueError(f"unsupported MoE activation {activation!r}") - - num_tokens = hidden_states.shape[0] - n2 = gate_up_q.shape[1] # 2 * intermediate - h = down_q.shape[1] # hidden - top_k = topk_ids.shape[1] - qt = int(GGML_Q4_0) - - # gate_up: [num_tokens*top_k, 2I] -> activation -> [num_tokens*top_k, I] - gate_up = ggml_moe_a8_vec(hidden_states, gate_up_q, topk_ids, top_k, qt, n2, num_tokens) - inter = act_fn(gate_up) - # down: each of the num_tokens*top_k intermediate rows uses its own expert id. - out = ggml_moe_a8_vec(inter, down_q, topk_ids, 1, qt, h, num_tokens * top_k) - out = out.reshape(num_tokens, top_k, h) * topk_weights.reshape(num_tokens, top_k, 1).to( - out.dtype - ) - return out.sum(dim=1) - - -__all__ = ["fused_experts_gguf_q4_0"] diff --git a/python/freetoken/moe/host_banks.py b/python/freetoken/moe/host_banks.py index 91e2648773..95562a8855 100644 --- a/python/freetoken/moe/host_banks.py +++ b/python/freetoken/moe/host_banks.py @@ -204,18 +204,6 @@ def alloc_banks(specs: dict[str, tuple[tuple[int, ...], torch.dtype]]) -> dict[s return {name: HostBank(shape, dtype) for name, (shape, dtype) in specs.items()} -def alloc_layer_banks( - specs: dict[str, tuple[tuple[int, ...], torch.dtype]], num_layers: int -) -> dict[str, list[HostBank]]: - """Allocate per-layer host banks: ``{name: ([num_experts, ...] row shape, dtype)}`` - -> one independently allocated (page-aligned, independently pin/lock-able) - ``HostBank`` per layer per name.""" - return { - name: [HostBank(shape, dtype) for _ in range(num_layers)] - for name, (shape, dtype) in specs.items() - } - - class _ResidencyPlan: """Per-layer ``HostResidency`` labels, ambiently visible to the bank settle points. @@ -265,13 +253,13 @@ def _settle(bank: HostBank, residency: str) -> None: bank.lock() -def pin_banks(banks: dict[str, HostBank | list[HostBank]]) -> None: +def pin_banks(banks: dict[str, HostBank | list[HostBank] | dict[int, HostBank]]) -> None: """Settle every bank after it has been filled -- pin-after-fill by default. - List-valued entries are per-layer and honor the ambient :func:`requested_residency` plan; scalar banks always pin.""" + List-valued entries (indexed by layer id) and dict-valued ones (layer id -> bank) are per-layer and honor the ambient :func:`requested_residency` plan; scalar banks always pin.""" plan = _requested_residency for bank in banks.values(): - if isinstance(bank, list): - for layer_id, layer_bank in enumerate(bank): + if isinstance(bank, (list, dict)): + for layer_id, layer_bank in bank.items() if isinstance(bank, dict) else enumerate(bank): residency = ( HostResidency.PINNED.value if plan is None else plan.residency_for(layer_id) @@ -479,7 +467,6 @@ def rd(i: int) -> None: "LayerCompletionTracker", "PinPipeline", "alloc_banks", - "alloc_layer_banks", "born_pinned_default", "pin_banks", "read_file_into", diff --git a/python/freetoken/moe/legacy_format.py b/python/freetoken/moe/legacy_format.py index 973a3dba52..8f347e5c30 100644 --- a/python/freetoken/moe/legacy_format.py +++ b/python/freetoken/moe/legacy_format.py @@ -30,12 +30,16 @@ def legacy_bank_names(quant_format: str) -> dict[str, str]: (QuantKind.NONE, "fused"): "bf16", (QuantKind.FP8_BLOCK, "triton"): "fp8_block", (QuantKind.NVFP4, "triton"): "nvfp4", + # same banks as triton, plus the resident input scales; a "nvfp4" file reads back as triton + (QuantKind.NVFP4, "w4a4"): "nvfp4", (QuantKind.NVFP4, "marlin"): "nvfp4_marlin", (QuantKind.NVFP4, "b12x"): "nvfp4_b12x", (QuantKind.MXFP4, "triton_gptoss"): "mxfp4_triton", (QuantKind.MXFP4, "triton"): "ds_fp4", } -_KIND_KERNEL = {fmt: kk for kk, fmt in LEGACY_FORMAT.items()} +_KIND_KERNEL: dict[str, tuple[QuantKind, str]] = {} +for _kk, _fmt in LEGACY_FORMAT.items(): + _KIND_KERNEL.setdefault(_fmt, _kk) def legacy_format_for(kind: QuantKind, kernel: str) -> str: diff --git a/python/freetoken/moe/offload_cache.py b/python/freetoken/moe/offload_cache.py index 7abb9c8003..99d799aabb 100644 --- a/python/freetoken/moe/offload_cache.py +++ b/python/freetoken/moe/offload_cache.py @@ -2,8 +2,10 @@ import math import os +import weakref +from contextlib import contextmanager from dataclasses import dataclass -from typing import Iterator +from typing import TYPE_CHECKING, Iterator import torch from flashlib.kernels.slot_cache import N_STATS, Stat @@ -13,6 +15,9 @@ # (kept for A/B profiling). Falls back to per-bank automatically if a bank's row bytes or # base address are not 16-byte aligned. _FUSED_COPY = os.getenv("FREETOKEN_FUSED_COPY", "1").strip().lower() not in {"0", "false", "no", "off"} +# Opt-in: run the fused copy_missing on the slim kernel (a few small CTAs, leaves most SMs to +# concurrent compute) instead of the 8-CTAs-per-bank one. Off by default. +_SLIM_COPY = os.getenv("FREETOKEN_MOE_SLIM_COPY", "0").strip().lower() in {"1", "true", "yes", "on"} # cudaMemcpyBatchAsync silently degrades to a SYNCHRONOUS copy when a batch mixes # large entries with sub-~256KB entries on registered host memory (H100 + CUDA 13.0, @@ -25,8 +30,13 @@ # 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 +if TYPE_CHECKING: + from freetoken.moe.offload_kernels import HybridRoutes + logger = init_logger(__name__) # quant_format -> bank names, in registration order: the single place a format's bank @@ -42,11 +52,12 @@ # Half the host/cache footprint of bf16; the grouped GEMM (kernel/triton/fp8_blockscale_moe) # reads the routed fp8 rows directly and dequantizes in the K-loop (no bf16 materialization). "fp8_block": ("gate_up", "gate_up_scale", "down", "down_scale"), - # native GGUF Q4_0 experts: packed block bytes per output row, dequantized inside - # the borrowed ggml MoE kernels. gate_up [L*E, 2I, H//32*18], down [L*E, H, I//32*18]. - "q4_0": ("gate_up", "down"), + # GGUF experts in their own ggml type per layer: gate_up [L*E, 2I, row_bytes] (gate rows, then up rows), down [L*E, H, row_bytes]; float experts keep their dtype + "gguf": ("gate_up", "down"), + # GGUF experts whose NVFP4 tensors carry per-expert scales: global [L*E, 4] fp32 (gate, up, down, 1) weight_scale_2 + "gguf_scaled": ("gate_up", "down", "global"), # native ModelOpt rows for the Triton inline-dequant kernels: packed e2m1 codes + - # fp8-e4m3 per-16 block scales + per-output-row fp16 globals (w1/w3 carry distinct + # fp8-e4m3 per-16 block scales + per-output-row fp32 globals (w1/w3 carry distinct # globals, and folding them into the e4m3 block scales would underflow) "nvfp4": ( "gate_up_packed", @@ -73,7 +84,7 @@ "down_bias", ), # DeepSeek-V4 FP4: packed e2m1 codes + e8m0 per-32 block scales, no global scale - # (4 banks). Read by DeepSeek-V4's own DS-FP4 grouped GEMV kernels via bank_views(). + # (4 banks). Read by DeepSeek-V4's own DS-FP4 grouped GEMV kernels via pool_bank_views(). "ds_fp4": ("gate_up_packed", "gate_up_scale", "down_packed", "down_scale"), } @@ -89,8 +100,7 @@ (2 * I // 128) * fp8_block_scale_pad(2 * I // 128, H // 128) + (H // 128) * fp8_block_scale_pad(H // 128, I // 128) ) * 2, - "q4_0": lambda H, I: 2 * I * (H // 32) * 18 + H * (I // 32) * 18, - "nvfp4": lambda H, I: 2 * I * (H // 2 + H // 16 + 2) + H * (I // 2 + I // 16 + 2), + "nvfp4": lambda H, I: 2 * I * (H // 2 + H // 16 + 4) + H * (I // 2 + I // 16 + 4), "mxfp4": lambda H, I: 2 * I * (H // 2 + H // 32 + 2) + H * (I // 2 + I // 32 + 2), "ds_fp4": lambda H, I: 2 * I * (H // 2 + H // 32) + H * (I // 2 + I // 32), } @@ -100,6 +110,19 @@ MARLIN_MAX_CACHE_SIZE = 992 +def placeholder_bank(num_experts: int) -> torch.Tensor: + """The source of a bank whose role a layer's expert kernel does not read (experts of several formats): zero bytes per row.""" + return torch.empty((num_experts, 0), dtype=torch.uint8) + + +def bank_layouts(sources: dict[str, list[torch.Tensor]], layer_tags: list | None = None) -> list[tuple]: + """Each layer's layout key: its quant tag and every bank's (dtype, row shape). A slot pool holds layers of one key.""" + num_layers = len(next(iter(sources.values()))) + tags = layer_tags if layer_tags is not None else [None] * num_layers + assert len(tags) == num_layers, (len(tags), num_layers) + return [(tags[l], tuple((t[l].dtype, tuple(t[l].shape[1:])) for t in sources.values())) for l in range(num_layers)] + + @dataclass class OffloadMoeCache: num_layers: int @@ -107,6 +130,9 @@ class OffloadMoeCache: cache_size: int device: torch.device cache_policy: str = "lru" + # A prefill chunk of at most this many tokens admits the experts it routes to into the slot + # cache, as decode does, instead of streaming every expert of each layer; 0 streams them all. + prefill_demand_tokens: int = 0 prefill_overlap: bool = False # Prefill hit/miss split: experts already resident in the slot cache (slots # >= 2 * num_experts) are gathered device-side into the double buffer instead @@ -135,22 +161,61 @@ class OffloadMoeCache: # of that step's misses are computed on the CPU. 0 -> never fetch (CPU does every # miss, the GPU cache stays cold); large -> behaves like pure offload. hybrid_max_fetch: int = 1 - # hybrid only: when > 0, replaces the fixed cap with a per-step fraction -- fetch + # hybrid only, one per MoE layer: when > 0, replaces the fixed cap with a per-step fraction -- fetch # ~fraction * misses experts over PCIe (rounded to whichever integer balances the # overlap best), the CPU computes the rest. The engine sets it to the benched # pcie_bw / cpu_bw ratio so the PCIe fetch and the CPU overflow GEMV take equal # time (perfect overlap): fetched : cpu = pcie : cpu - pcie. - hybrid_fetch_fraction: float = 0.0 + hybrid_fetch_fraction: list[float] | None = None # 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 + # Per-layer kernel layouts for experts of several formats (replaces ``layout``). The banks are the union of their roles; + # a layer whose kernel reads no such role carries a zero-byte [num_experts, 0] uint8 placeholder (see placeholder_bank). + layer_layouts: list[dict] | None = None + # Each layer belongs to one pool: a contiguous slot range with its own LRU scope; pool p owns slots [sum(pool_sizes[:p]), +pool_sizes[p]). Slot ids in slot_for_id are local to the pool. None = one pool for all layers. + pool_of_layer: list[int] | None = None + pool_sizes: list[int] | None = None + # GPU decode only: the caller's work that is independent of the routed experts (qwen4_exp: the + # shared expert) runs on decode_copy_stream beside the miss copy. None = FREETOKEN_MOE_COPY_OVERLAP. + decode_copy_overlap: bool | None = None + # GPU decode only: cross-layer expert prefetch (moe/prefetch.py). None = FREETOKEN_MOE_PREFETCH. + prefetch_mode: str | None = None + # GPU decode only: debug slot checks (moe/verify.py). None = FREETOKEN_MOE_PREFETCH_VERIFY. + verify_mode: str | None = None + # GPU decode only: debug slot writer history + audit every N decode steps (moe/slot_audit.py). + # None = FREETOKEN_MOE_SLOT_AUDIT, 0 = off. + slot_audit: int | None = None def __post_init__(self) -> None: - policy_ids = {"lru": 0} - assert self.cache_policy in policy_ids + from freetoken.moe.scored_ensure import POLICY_IDS + + assert self.cache_policy in POLICY_IDS, self.cache_policy assert self.decode_target in ("gpu", "cpu", "hybrid"), self.decode_target - if self.layout is None: + if self.cache_policy != "lru" and self.decode_target != "gpu": + # the hybrid and CPU decode paths run their own kernels and would silently stay LRU + raise ValueError( + f"--moe-cache-policy {self.cache_policy} needs GPU decode, not {self.decode_target!r}" + ) + if self.prefill_demand_tokens and self.decode_target != "gpu": + raise ValueError( + f"--moe-prefill-demand-tokens needs GPU decode; this cache decodes on {self.decode_target!r}" + ) + if self.layer_layouts is not None: + assert self.layout is None, "pass layout or layer_layouts, not both" + assert len(self.layer_layouts) == self.num_layers, (len(self.layer_layouts), self.num_layers) + if any(spec.resident for layout in self.layer_layouts for spec in layout.values()): + raise ValueError("experts of several formats cannot use a kernel with GPU-resident alphas (marlin, b12x)") + elif self.layout is None: assert self.quant_format in _BANK_SCHEMAS, f"unknown quant_format {self.quant_format!r}" + if self.pool_sizes is None: + self.pool_of_layer = [0] * self.num_layers + self.pool_sizes = [self.cache_size] + self.pool_of_layer = list(self.pool_of_layer) + self.pool_sizes = list(self.pool_sizes) + assert len(self.pool_of_layer) == self.num_layers, (len(self.pool_of_layer), self.num_layers) + assert sorted(set(self.pool_of_layer)) == list(range(len(self.pool_sizes))), self.pool_of_layer + assert sum(self.pool_sizes) == self.cache_size and min(self.pool_sizes) > 0, (self.pool_sizes, self.cache_size) # Attached by the engine for decode_target == "cpu" (CpuMoeExecutor); None # for the GPU decode path. self.cpu_executor = None @@ -158,14 +223,25 @@ def __post_init__(self) -> None: # offload/PCIe path. Set by the engine after construction (empty = all-GPU, # all layers = the plain --moe-strategy cpu case). self.cpu_layer_ids: frozenset = frozenset() + # Bank row layouts, set by set_bank_sources: layouts[g][k] is (dtype, row shape) of bank k (layouts of two quant tags may be equal), layer_layout[l] is layer l's g. + # With one (dtype, row shape) per bank a bank cache is typed [cache_size, *row]; with several, or with pads, it is a byte segment each pool views with its own layout. + self.layouts: list[tuple[tuple[torch.dtype, tuple[int, ...]], ...]] = [] + self.layer_layout: list[int] = [0] * self.num_layers + # _pad_bytes[g][k]: zeroed bytes after the last row of every pool and staging buffer of layout g in bank k, for kernels that read past it (16-aligned) + self._pad_bytes: list[list[int]] = [] + self._raw_banks = False # num_experts floor + nvfp4_marlin slot cap, shared with the runtime-rebuild path. - self.validate_rebuild(self.cache_size) + self.validate_rebuild(self.cache_size, self.pool_sizes) assert not self.prefill_overlap or self.cache_size >= 2 * self.num_experts, ( "Prefill overlap borrows two full expert-layer buffers from the unified MoE " "cache, so cache_size must be at least 2 * num_experts " "(raise moe_cache_size or disable moe_prefill_overlap)" ) - self.cache_policy_id = policy_ids[self.cache_policy] + self.cache_policy_id = POLICY_IDS[self.cache_policy] + # rule: an unrouted expert within this many logits of the top-k-th is a near miss + self.near_miss_thr = ( + float(os.getenv("FREETOKEN_MOE_NEAR_MISS_THR", "0.25")) if self.cache_policy == "rule" else None + ) self.slot_for_id = torch.full( (self.num_layers, self.num_experts), -1, @@ -199,11 +275,16 @@ 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 # set_bank_sources). The GPU slot cache stays one unified pool per bank. - if self.layout is not None: + if self.layer_layouts is not None: + self.bank_schema = tuple(dict.fromkeys( + role for layout in self.layer_layouts for role, spec in layout.items() if not spec.resident + )) + elif self.layout is not None: self.bank_schema = tuple(role for role, spec in self.layout.items() if not spec.resident) else: self.bank_schema = _BANK_SCHEMAS[self.quant_format] @@ -223,24 +304,23 @@ def __post_init__(self) -> None: # is a one-off warm-up increment at capture time (<0.1% over a session). self.collect_stats = False # [num_layers, N_STATS] -- ensure_experts passes lru_stats[layer_id] straight to - # the kernel, which accumulates in the same launch. The stat_* tensors below stay - # for the hybrid path, whose kernel is still ours. + # the kernel, which accumulates in the same launch. The stat_* tensors below are + # the hybrid path's: views of one buffer its kernel accumulates into the same way. self.lru_stats = torch.zeros( (self.num_layers, N_STATS), dtype=torch.int64, device=self.device ) - self.stat_missing = torch.zeros((), dtype=torch.int64, device=self.device) - self.stat_active = torch.zeros((), dtype=torch.int64, device=self.device) - self.stat_calls = torch.zeros((), dtype=torch.int64, device=self.device) - # hybrid only: experts actually fetched over PCIe (<= stat_missing). The CPU - # computes stat_missing - stat_fetched of them. - self.stat_fetched = torch.zeros((), dtype=torch.int64, device=self.device) + # (missing, fetched, active, calls), then the same per layer + stat_buf = torch.zeros(4 * (self.num_layers + 1), dtype=torch.int64, device=self.device) + self.hybrid_stats = (stat_buf[:4], stat_buf[4:]) + # hybrid only: stat_fetched counts the experts actually fetched over PCIe (<= stat_missing). + # The CPU computes stat_missing - stat_fetched of them. + self.stat_missing, self.stat_fetched, self.stat_active, self.stat_calls = stat_buf[:4].unbind() # Per-layer counterparts of the scalars above (indexed by MoE-layer id). Same # device-side accumulation (graph-safe: layer_id is a static index per graph node), # so one req's per-layer miss rate is readable via decode_miss_stats_per_layer(). - self.stat_missing_layer = torch.zeros(self.num_layers, dtype=torch.int64, device=self.device) - self.stat_active_layer = torch.zeros(self.num_layers, dtype=torch.int64, device=self.device) - self.stat_fetched_layer = torch.zeros(self.num_layers, dtype=torch.int64, device=self.device) - self.stat_steps_layer = torch.zeros(self.num_layers, dtype=torch.int64, device=self.device) + per_layer = stat_buf[4:].view(self.num_layers, 4) + self.stat_missing_layer, self.stat_fetched_layer = per_layer[:, 0], per_layer[:, 1] + self.stat_active_layer, self.stat_steps_layer = per_layer[:, 2], per_layer[:, 3] # Opt-in decode routing histogram (per layer, per expert) for cache-skew # analysis. Accumulated in ``ensure_experts`` from the raw expert ids before the # kernel rewrites them to slots. Only accurate with CUDA graphs disabled (the @@ -251,24 +331,33 @@ def __post_init__(self) -> None: ) # (per-layer sources, cache) per bank, in schema order. Every piece of cache # machinery that moves bank bytes (copy_missing, the prefill double buffers, - # bank_views) iterates this list, so the slot cache is bank-count agnostic. + # the pool and staging views) iterates this list, so the slot cache is bank-count agnostic. self.banks: list[tuple[list[torch.Tensor], torch.Tensor]] = [] # Fused multi-bank copy descriptor (built by set_bank_sources/_build_copy_plan). # Source pointers are per layer (_copy_src_ptrs[layer_id] -> [num_banks] device - # tensor); dst/feat are layer-invariant. + # tensor); dst pointers are per pool and row bytes per layout. self._copy_fused_ok = False + self._copy_blocks_per_bank = 8 self._copy_dst_ptrs: torch.Tensor | None = None + self._copy_dst_ptrs_pool: list[torch.Tensor] | None = None self._copy_src_ptrs: list[torch.Tensor] | None = None - self._copy_feat_bytes: torch.Tensor | None = None + self._copy_feat_bytes: list[torch.Tensor] | None = None # The layer whose misses ensure_experts/materialize_layer staged last; consumed # by copy_missing to pick the per-layer source (part of the same pending-copy # state as evict_slots/src_indices/num_indices). # _pending_whole_layer records WHICH staged it: the pageable branch is only sound after materialize_layer self._pending_src_layer: int | None = None self._pending_whole_layer = False - # Per-bank [2, num_experts, ...] double-buffer views over the slot cache's - # first 2 * num_experts slots (set up when prefill_overlap is enabled). - self.prefill_bank_buffers: list[torch.Tensor] = [] + # The two prefill staging buffers are the first 2 * _stage_bytes[k] bytes of every bank, on top of the first pools. + # _stage_views[g][b] are buffer b's per-bank [num_experts, ...] views for layout g; _inval[b] are the global slots it overwrites. + self._stage_bytes: list[int] = [] + self._stage_views: list[list[tuple[torch.Tensor, ...] | None]] = [] + # _stage_tails[g][b]: the pad after buffer b's rows in layout g + self._stage_tails: list[list[list[torch.Tensor]]] = [] + self._inval: list[slice | torch.Tensor] = [] + # the slot audit's label for the staged copy, and the most entries its plan can have + self._pending_kind = 0 + self._pending_width: int | None = None self.prefill_copy_stream: torch.cuda.Stream | None = None self.prefill_begin_event: torch.cuda.Event | None = None self.prefill_ready_events: list[torch.cuda.Event] = [] @@ -278,7 +367,7 @@ def __post_init__(self) -> None: self._prefill_buffer_has_release_event: list[bool] = [False, False] # hit-D2D split state: pinned begin-of-chunk snapshot of slot_for_id (the # classification input; frozen for the chunk -- no decode runs inside one, - # and buffer invalidation only clears slot < 2E entries, which classify as + # and buffer invalidation only clears slots under the staging buffers, which classify as # miss regardless), the lazily resolved batch-memcpy entry point (False = # unavailable), and row counters for cache reports. self._prefill_slot_snapshot: torch.Tensor | None = None @@ -288,11 +377,300 @@ def __post_init__(self) -> None: self._batch_memcpy = None self.prefill_hit_rows = 0 self.prefill_total_rows = 0 + self._slice_pools() + # Decode copy overlap: one side stream plus a (forked, done) event pair per layer. + # Neither is cache_size-shaped, so rebuild() keeps them. + if self.decode_copy_overlap is None: + self.decode_copy_overlap = bool(ENV.MOE_COPY_OVERLAP) + self.decode_copy_stream: torch.cuda.Stream | None = None + self._decode_copy_events: list[tuple[torch.cuda.Event, torch.cuda.Event]] = [] + if self.decode_copy_overlap and self.device.type == "cuda" and self.decode_target == "gpu": + self._init_decode_copy_overlap() + self.prefetch = self._init_prefetch() + self.verify = self._init_verify() + self.audit = self._init_audit() + self._refuse_single_pool_features() + + def _refuse_single_pool_features(self) -> None: + """The scored policies, prefetch, its debug checks and the on-demand prefill GEMM address + the cache by global slot id over one typed bank, so with several slot pools (#689) they + would install and read pool-local slots in the wrong rows: refuse that combination.""" + if len(self.pool_sizes) == 1: + return + used = [ + name + for name, on in ( + (f"--moe-cache-policy {self.cache_policy}", self.cache_policy != "lru"), + ("FREETOKEN_MOE_PREFETCH", self.prefetch_mode not in (None, "off")), + ("FREETOKEN_MOE_PREFETCH_VERIFY", self.verify_mode not in (None, "off")), + ("FREETOKEN_MOE_SLOT_AUDIT", bool(self.slot_audit)), + ("--moe-prefill-demand-tokens", bool(self.prefill_demand_tokens)), + ) + if on + ] + if used: + raise ValueError( + f"{', '.join(used)} need one MoE slot pool, but this cache has {len(self.pool_sizes)} " + "(--moe-cache-scope per-layer, or experts in several quant formats)" + ) + + def _init_audit(self): + from freetoken.kernel.fast_index_copy import _skip_fast_index_copy_enabled + from freetoken.moe.slot_audit import SlotAuditor, resolve_interval + + from_env = self.slot_audit is None + self.slot_audit = resolve_interval(self.slot_audit) + if not self.slot_audit or self.device.type != "cuda": + return None + if self.decode_target != "gpu": + if not from_env: + raise ValueError(f"FREETOKEN_MOE_SLOT_AUDIT={self.slot_audit} needs GPU decode, not {self.decode_target!r}") + logger.warning(f"FREETOKEN_MOE_SLOT_AUDIT={self.slot_audit} ignored: decode target is {self.decode_target!r}") + self.slot_audit = 0 + return None + if _skip_fast_index_copy_enabled(): + logger.warning("FREETOKEN_MOE_SLOT_AUDIT ignored: FREETOKEN_SKIP_FAST_INDEX_COPY makes every copy a no-op") + self.slot_audit = 0 + return None + return SlotAuditor(self.num_layers, self.num_experts, self.device, interval=self.slot_audit) + + def _init_verify(self): + from freetoken.moe.verify import ExpertVerifier, resolve_mode + + from_env = self.verify_mode is None + self.verify_mode = resolve_mode(self.verify_mode) + if self.verify_mode == "off" or self.device.type != "cuda": + return None + if self.decode_target != "gpu": + if not from_env: + raise ValueError(f"FREETOKEN_MOE_PREFETCH_VERIFY={self.verify_mode} needs GPU decode, not {self.decode_target!r}") + logger.warning(f"FREETOKEN_MOE_PREFETCH_VERIFY={self.verify_mode} ignored: decode target is {self.decode_target!r}") + self.verify_mode = "off" + return None + return ExpertVerifier(self.num_layers, self.num_experts, self.device, mode=self.verify_mode) + + def _init_prefetch(self): + from freetoken.moe.prefetch import ExpertPrefetcher, resolve_mode + + from_env = self.prefetch_mode is None + self.prefetch_mode = resolve_mode(self.prefetch_mode) + if self.prefetch_mode == "off": + return None + if self.decode_target != "gpu": + if not from_env: + raise ValueError(f"FREETOKEN_MOE_PREFETCH={self.prefetch_mode} needs GPU decode, not {self.decode_target!r}") + # the process-wide flag skips hybrid/CPU-decoded caches, like FREETOKEN_MOE_COPY_OVERLAP + logger.warning(f"FREETOKEN_MOE_PREFETCH={self.prefetch_mode} ignored: decode target is {self.decode_target!r}") + self.prefetch_mode = "off" + return None + if self.device.type != "cuda": + return None + prefetch = ExpertPrefetcher(self.num_layers, self.num_experts, self.device, mode=self.prefetch_mode) + # a proxy, not a cycle: dropping the cache must free its slot banks without waiting for gc + prefetch.cache = weakref.proxy(self) + return prefetch + + @property + def prefetch_on(self) -> bool: + """FREETOKEN_MOE_PREFETCH=on: every ensure keys prefetched slots low priority.""" + return self.prefetch is not None and self.prefetch.mode == "on" + + def _pool_layout(self, pool: int) -> int: + return self.layer_layout[self.pool_of_layer.index(pool)] + + def _bank_rows(self, k: int, layout: int, offset: int, nrows: int) -> torch.Tensor: + """``nrows`` rows of bank ``k`` in ``layout``, starting ``offset`` bytes into the bank cache.""" + cache = self.banks[k][1] + row = self._row_bytes[layout][k] + if not self._raw_banks: + first = offset // row + return cache if first == 0 and nrows == cache.shape[0] else cache[first : first + nrows] + dtype, shape = self.layouts[layout][k] + return cache[offset : offset + nrows * row].view(dtype).view(nrows, *shape) + + def _bank_bytes(self, k: int, pool_sizes: list[int]) -> int: + """Bytes of raw bank ``k``: each pool's slots and pad, plus the largest pad, so two staging buffers fit wherever two layers of slots do.""" + layouts = [self._pool_layout(p) for p in range(len(pool_sizes))] + slots = sum(s * self._row_bytes[g][k] + self._pad_bytes[g][k] for g, s in zip(layouts, pool_sizes)) + return slots + max(pads[k] for pads in self._pad_bytes) + + def _stage_fits(self, buffers: int) -> bool: + return all(c.numel() * c.element_size() >= buffers * s for (_, c), s in zip(self.banks, self._stage_bytes)) + + def _slice_pools(self) -> None: + E = self.num_experts + bases = [0] + for size in self.pool_sizes[:-1]: + bases.append(bases[-1] + size) + self.pool_bases = bases + spans = list(zip(bases, self.pool_sizes)) + + def cut(t: torch.Tensor, b: int, s: int) -> torch.Tensor: + return t if s == t.shape[0] else t[b : b + s] + + self._pool_id_of_slot = [cut(self.id_of_slot, b, s) for b, s in spans] + self._pool_usage = [cut(self.usage, b, s) for b, s in spans] + # rows[k][p]: bytes of one slot of pool p in bank k; before the banks exist, one unit per slot with a stage of E + if self.banks: + rows = [[self._row_bytes[self._pool_layout(p)][k] for p in range(len(spans))] for k in range(len(self.banks))] + pads = [[self._pad_bytes[self._pool_layout(p)][k] for p in range(len(spans))] for k in range(len(self.banks))] + stage = self._stage_bytes + else: + rows, pads, stage = [[1] * len(spans)], [[0] * len(spans)], [E] + offsets = [] + for row, pad in zip(rows, pads): + offsets.append([0]) + for s, r, d in zip(self.pool_sizes[:-1], row, pad): + offsets[-1].append(offsets[-1][-1] + s * r + d) + self._pool_offsets = offsets + # local slots below this index overlap the prefill staging area in some bank; a pool's zero-byte rows (a role its kernel does not read) overlap nothing + self._pool_stable_from = [ + max((min(max(-(-(2 * stage[k] - offsets[k][p]) // rows[k][p]), 0), s) for k in range(len(rows)) if rows[k][p]), default=0) + for p, s in enumerate(self.pool_sizes) + ] + self._inval = [] + for b in range(2): + mask = torch.zeros(self.cache_size, dtype=torch.bool) + for k in range(len(rows)): + lo, hi = b * stage[k], (b + 1) * stage[k] + for p, (base, s) in enumerate(spans): + if not rows[k][p]: + continue + first = max((lo - offsets[k][p]) // rows[k][p], 0) + end = min(-(-(hi - offsets[k][p]) // rows[k][p]), s) + mask[base + first : base + max(first, end)] = True + idx = mask.nonzero().flatten() + if idx.numel() == 0 or int(idx[-1] - idx[0]) + 1 == idx.numel(): + self._inval.append(slice(int(idx[0]), int(idx[-1]) + 1) if idx.numel() else slice(0, 0)) + else: + self._inval.append(idx.to(self.device)) + self._stage_dirty = [False, False] + if not self.banks: + self._pool_bank_views, self._stage_views, self._stage_tails = [], [], [] + return + self._pool_bank_views = [ + tuple(self._bank_rows(k, self._pool_layout(p), offsets[k][p], s) for k in range(len(self.banks))) + for p, s in enumerate(self.pool_sizes) + ] + self._stage_views = [ + [ + tuple(self._bank_rows(k, g, b * stage[k], E) for k in range(len(self.banks))) if self._stage_fits(b + 1) else None + for b in range(2) + ] + for g in range(len(self.layouts)) + ] + self._stage_tails = [ + [[self.banks[k][1][b * stage[k] + E * row[k] :][: pad[k]] for k in range(len(self.banks)) if pad[k] and self._stage_fits(b + 1)] for b in range(2)] + for row, pad in zip(self._row_bytes, self._pad_bytes) + ] + + def stage_views(self, layer_id: int, buffer_id: int) -> tuple[torch.Tensor, ...]: + """Per-bank ``[num_experts, ...]`` views of prefill staging buffer ``buffer_id`` in ``layer_id``'s layout; row == expert id.""" + views = self._stage_views[self.layer_layout[layer_id]][buffer_id] + assert views is not None, f"the slot cache is too small for staging buffer {buffer_id}" + return views + + def pool_state(self, layer_id: int) -> tuple[torch.Tensor, torch.Tensor]: + """``(id_of_slot, usage)`` views of the pool that serves ``layer_id``.""" + p = self.pool_of_layer[layer_id] + return self._pool_id_of_slot[p], self._pool_usage[p] + + def pool_slots(self, layer_id: int) -> int: + return self.pool_sizes[self.pool_of_layer[layer_id]] + + def bank_views(self, n: int | None = None) -> tuple[torch.Tensor, ...]: + """Per-bank views of the whole ``[S]`` slot cache, or of its first ``n`` slots, for a cache + with one pool and one bank layout, where a slot id is its bank row (the on-demand prefill + GEMM over the slot cache and the fork's slot checks).""" + assert self.banks, "set_bank_sources must register the banks first" + assert len(self.pool_sizes) == 1 and not self._raw_banks, "bank_views needs one slot pool and one bank layout" + if n is None: + return tuple(cache for _, cache in self.banks) + return tuple(cache[:n] for _, cache in self.banks) + + def pool_bank_views(self, layer_id: int) -> tuple[torch.Tensor, ...]: + """Per-bank slot views of the pool that serves ``layer_id``, indexed by its local slot ids.""" + assert self.banks, "set_bank_sources must register the banks first" + return self._pool_bank_views[self.pool_of_layer[layer_id]] + def _init_decode_copy_overlap(self) -> None: + # a dedicated stream, not one of torch's 32 pooled ones: a pooled stream can alias the engine + # or capture stream, and the side work would then run serially (correct, but no overlap) + from cuda.bindings import runtime as cudart + + with torch.cuda.device(self.device): + err, handle = cudart.cudaStreamCreateWithFlags(cudart.cudaStreamNonBlocking) + if err != cudart.cudaError_t.cudaSuccess: + raise RuntimeError(f"cudaStreamCreateWithFlags failed: {err}") + self.decode_copy_stream = torch.cuda.ExternalStream(int(handle), device=self.device) + self._decode_copy_events = [ + (torch.cuda.Event(), torch.cuda.Event()) for _ in range(self.num_layers) + ] + # torch creates the CUDA event on its first record: do that now, never inside a capture + for pair in self._decode_copy_events: + for event in pair: + event.record(self.decode_copy_stream) + logger.info("MoE decode copy overlap on: expert misses copy on a side stream") + + def _init_evict_state(self) -> None: + """Scored policies' per-(layer, expert) state (moe/scored_ensure.py); None under LRU. + + It is not slot-shaped: it survives eviction and prefill invalidation, and a rebuild + only resets it. ``evict_tok`` is the decode-step clock the first ensure of each step bumps. + """ + from freetoken.moe import scored_ensure as se + + self.evict_tok = self.evict_last_tok = self.evict_lc = self.evict_ct = None + self.evict_g_table = None + n, dev = self.num_layers * self.num_experts, self.device + if self.cache_policy_id >= 1: + self.evict_tok = torch.zeros((), dtype=torch.int64, device=dev) + self.evict_last_tok = torch.full((n,), se.LAST_TOK_NEVER, dtype=torch.int64, device=dev) + if self.cache_policy_id >= 2: + self.evict_lc = torch.full((n,), se.LC_NEVER, dtype=torch.int32, device=dev) + self.evict_ct = torch.zeros((n,), dtype=torch.int64, device=dev) + self.evict_g_table = se.softplus2_table(dev) + self._alloc_evict_mirrors() + + def _alloc_evict_mirrors(self) -> None: + """Slot-shaped copies of each slot owner's state, read coalesced by the victim scan. + + A mirror is valid while ``evict_slot_owner`` equals ``id_of_slot``; the kernel keeps it + so, and anything else that installs an id into a slot must set the owner to -1. + """ + self.evict_slot_owner = self.evict_slot_last_tok = self.evict_slot_lc = self.evict_slot_ct = None + if self.evict_tok is None: + return + size, dev = self.cache_size, self.device + self.evict_slot_owner = torch.full((size,), -1, dtype=torch.int32, device=dev) + self.evict_slot_last_tok = torch.zeros((size,), dtype=torch.int64, device=dev) + if self.evict_lc is not None: + self.evict_slot_lc = torch.zeros((size,), dtype=torch.int32, device=dev) + self.evict_slot_ct = torch.zeros((size,), dtype=torch.int64, device=dev) + + def reset_evict_state(self) -> None: + if self.evict_tok is None: + return + from freetoken.moe import scored_ensure as se + + self.evict_tok.zero_() + self.evict_last_tok.fill_(se.LAST_TOK_NEVER) + self.evict_slot_owner.fill_(-1) + if self.evict_lc is not None: + self.evict_lc.fill_(se.LC_NEVER) + self.evict_ct.zero_() + + def evict_clock_layer(self) -> int: + """The first layer that ensures on the GPU each decode step; its call bumps ``evict_tok``.""" + if not self.cpu_layer_ids: + return 0 + return min(set(range(self.num_layers)) - self.cpu_layer_ids, default=0) def set_bank_sources( self, sources: dict[str, list[torch.Tensor]], layer_residency: list[str] | None = None, + layer_tags: list | None = None, + row_tail_bytes: dict[str, list[int]] | None = None, ) -> None: """Attach the host (CPU pinned) expert source banks and allocate a GPU slot cache per bank, following the format's bank schema. @@ -306,6 +684,10 @@ def set_bank_sources( ``layer_residency`` labels each layer with a ``HostResidency`` value (default: all pinned). Non-pinned (LOCKED/PAGEABLE) layers have no device address: they must already be routed to the CPU executor (``cpu_layer_ids``, set BEFORE this call), the copy plan skips their rows, and their only movement is ``copy_missing``'s whole-layer pageable prefill branch -- which is why prefill overlap is incompatible with them. + + ``layer_tags`` names each layer's quant (see :func:`bank_layouts`); layers of different tags never share a pool. + + ``row_tail_bytes`` gives, per bank and layer, how many bytes past the last row the kernel reads; every pool and staging buffer is then followed by that many zero bytes. """ from freetoken.moe.legacy_format import canonical_role from freetoken.moe.host_banks import HostResidency @@ -339,43 +721,79 @@ def set_bank_sources( for name in self.bank_schema: per_layer = sources[name] assert len(per_layer) == self.num_layers, (name, len(per_layer)) - head = per_layer[0] - if self.layout is not None: - spec = self.layout[name] - if tuple(head.shape[1:]) != tuple(spec.shape) or head.dtype != spec.dtype: - raise ValueError( - f"bank {name!r} rows are {tuple(head.shape[1:])} {head.dtype} but the expert kernel's layout " - f"wants {tuple(spec.shape)} {spec.dtype}; the banks were packed for another kernel" - ) for layer_id, source in enumerate(per_layer): assert source.is_contiguous(), f"bank {name!r} layer {layer_id} must be contiguous" assert source.size(0) == self.num_experts, (name, layer_id, source.shape) - assert source.shape == head.shape and source.dtype == head.dtype, ( - name, layer_id, source.shape, source.dtype, - ) + spec = self.layout[name] if self.layout is not None else None + if self.layer_layouts is not None: + spec = self.layer_layouts[layer_id].get(name) + if spec is None and source.numel() != 0: + raise ValueError(f"layer {layer_id}'s expert kernel reads no {name!r} bank, so it takes a zero-byte placeholder") + if spec is not None: + if tuple(source.shape[1:]) != tuple(spec.shape) or source.dtype != spec.dtype: + raise ValueError( + f"bank {name!r} rows are {tuple(source.shape[1:])} {source.dtype} but the expert kernel's layout " + f"wants {tuple(spec.shape)} {spec.dtype}; the banks were packed for another kernel" + ) self.bank_sources[name] = list(per_layer) - self.bank_caches[name] = torch.empty( - (self.cache_size, *head.shape[1:]), - dtype=head.dtype, - device=self.device, - ) - self.banks = [(self.bank_sources[n], self.bank_caches[n]) for n in self.bank_schema] + keys = bank_layouts(sources, layer_tags) + unique = list(dict.fromkeys(keys)) + self.layouts = [banks for _, banks in unique] + self.layer_layout = [unique.index(key) for key in keys] + for p in range(len(self.pool_sizes)): + mixed = {self.layer_layout[l] for l, q in enumerate(self.pool_of_layer) if q == p} + assert len(mixed) == 1, f"slot pool {p} holds layers of {len(mixed)} layouts; a pool needs one" + self._pad_bytes = [[0] * len(self.bank_schema) for _ in self.layouts] + if row_tail_bytes is not None: + tails = {canonical_role(name): per_layer for name, per_layer in row_tail_bytes.items()} + for k, name in enumerate(self.bank_schema): + for layer_id, tail in enumerate(tails[canonical_role(name)]): + g = self.layer_layout[layer_id] + self._pad_bytes[g][k] = max(self._pad_bytes[g][k], -(-tail // 16) * 16) + self._raw_banks = len(set(self.layouts)) > 1 or any(map(any, self._pad_bytes)) + self._row_bytes = [[math.prod(shape) * dtype.itemsize for dtype, shape in layout] for layout in self.layouts] + if self._raw_banks: + # every pool view and staging buffer starts on a row boundary of its own layout + assert all(r % 16 == 0 for rows in self._row_bytes for r in rows), self._row_bytes + self._stage_bytes = [ + self.num_experts * max(rows[k] for rows in self._row_bytes) + max(pads[k] for pads in self._pad_bytes) + for k in range(len(self.bank_schema)) + ] + self._alloc_banks() + if not self._stage_fits(1): + raise ValueError(f"MoE cache of {self.cache_size} slots cannot stage one whole layer of experts") + if self.prefill_overlap and not self._stage_fits(2): + logger.warning("Disabling MoE prefill overlap: the slot cache cannot hold two whole layers of experts") + self.prefill_overlap = False + self._slice_pools() self._build_copy_plan() + if self.verify is not None: + self.verify.bind(self) + if self.audit is not None: + self.audit.bind(self) if self.prefill_overlap: self._init_prefill_overlap_buffers() + def _alloc_banks(self) -> None: + for k, name in enumerate(self.bank_schema): + if self._raw_banks: + self.bank_caches[name] = torch.zeros(self._bank_bytes(k, self.pool_sizes), dtype=torch.uint8, device=self.device) + else: + dtype, shape = self.layouts[0][k] + self.bank_caches[name] = torch.zeros((self.cache_size, *shape), dtype=dtype, device=self.device) + self.banks = [(self.bank_sources[n], self.bank_caches[n]) for n in self.bank_schema] + def _build_copy_plan(self) -> None: self._build_fused_copy_plan() if self._copy_fused_ok or self.device.type != "cuda" or not self.banks: return - for name in self.bank_schema: - cache = self.bank_caches[name] - feat = math.prod(cache.shape[1:]) * cache.element_size() - if feat % 128: - raise RuntimeError( - f"MoE bank {name!r} rows are {feat} bytes (not a multiple of 128): " - f"only the fused multi-bank copy can move them, but it is disabled" - ) + for k, name in enumerate(self.bank_schema): + for rows in self._row_bytes: + if rows[k] % 128: + raise RuntimeError( + f"MoE bank {name!r} rows are {rows[k]} bytes (not a multiple of 128): " + f"only the fused multi-bank copy can move them, but it is disabled" + ) def _build_fused_copy_plan(self) -> None: """Precompute the fused multi-bank copy descriptor (base addrs + per-row bytes). @@ -387,30 +805,34 @@ def _build_fused_copy_plan(self) -> None: """ self._copy_fused_ok = False self._copy_dst_ptrs = None + self._copy_dst_ptrs_pool = None self._copy_src_ptrs = None self._copy_feat_bytes = None self._copy_dst_ptrs_host: list[int] = [] self._copy_src_ptrs_host: list[list[int]] = [] - self._copy_feat_bytes_host: list[int] = [] - self._gather_bank_ids: list[int] = [] - self._gather_dst_ptrs: torch.Tensor | None = None - self._gather_feat_bytes: torch.Tensor | None = None + self._gather_bank_ids: list[list[int]] = [] + self._gather_feat_bytes: list[torch.Tensor | None] = [] + self._gather_stage_ptrs: list[list[torch.Tensor] | None] = [] + self._gather_pool_ptrs: list[torch.Tensor | None] = [] if not _FUSED_COPY or self.device.type != "cuda" or not self.banks: return from freetoken.kernel.pinned import device_ptr - dst_ptrs, feats = [], [] + aligned = [self._stage_bytes] + self._row_bytes + self._pool_offsets + if any(n % 16 for ns in aligned for n in ns) or any(c.data_ptr() % 16 for _, c in self.banks): + return # leave fused disabled; copy_missing uses the per-bank path + dst_ptrs = [cache.data_ptr() for _, cache in self.banks] layer_src_ptrs = [[] for _ in range(self.num_layers)] - for per_layer, cache in self.banks: - feat = math.prod(per_layer[0].shape[1:]) * per_layer[0].element_size() - if feat % 16 != 0 or cache.data_ptr() % 16 != 0: - return # leave fused disabled; copy_missing uses the per-bank path + for per_layer, _ in self.banks: for layer_id, source in enumerate(per_layer): if layer_id in self._unpinned_layers: # unregistered layer: no device alias exists, and the row is never consumed (CPU decode; pageable prefill) # a 0 placeholder keeps the descriptor shape layer_src_ptrs[layer_id].append(0) continue + if source.numel() == 0: + layer_src_ptrs[layer_id].append(0) # placeholder bank: zero bytes per row, never dereferenced + continue # The kernel dereferences these on the GPU, so store each host bank's # device alias (== data_ptr() under UVA identity; differs on # Windows/WDDM). @@ -418,51 +840,65 @@ def _build_fused_copy_plan(self) -> None: if src_dev % 16 != 0: return layer_src_ptrs[layer_id].append(src_dev) - dst_ptrs.append(cache.data_ptr()) - feats.append(feat) - self._copy_dst_ptrs = torch.tensor(dst_ptrs, dtype=torch.int64, device=self.device) - self._copy_src_ptrs = [ - torch.tensor(ptrs, dtype=torch.int64, device=self.device) - for ptrs in layer_src_ptrs + + def ptrs(values: list[int]) -> torch.Tensor: + return torch.tensor(values, dtype=torch.int64, device=self.device) + + self._copy_dst_ptrs = ptrs(dst_ptrs) + self._copy_dst_ptrs_pool = [ + self._copy_dst_ptrs if all(o[p] == 0 for o in self._pool_offsets) + else ptrs([d + o[p] for d, o in zip(dst_ptrs, self._pool_offsets)]) + for p in range(len(self.pool_sizes)) ] - self._copy_feat_bytes = torch.tensor(feats, dtype=torch.int64, device=self.device) + self._copy_src_ptrs = [ptrs(values) for values in layer_src_ptrs] + self._copy_feat_bytes = [ptrs(rows) for rows in self._row_bytes] self._copy_dst_ptrs_host = dst_ptrs self._copy_src_ptrs_host = layer_src_ptrs - self._copy_feat_bytes_host = feats # hit-D2D gather serves only the big banks; small banks are whole-layer # H2D entries (see _SMALL_BANK_FEAT_BYTES), so their rows never need D2D. - self._gather_bank_ids = [i for i, f in enumerate(feats) if f >= _SMALL_BANK_FEAT_BYTES] - if len(self._gather_bank_ids) == len(feats): - self._gather_dst_ptrs = self._copy_dst_ptrs - self._gather_feat_bytes = self._copy_feat_bytes - elif self._gather_bank_ids: - self._gather_dst_ptrs = self._copy_dst_ptrs[self._gather_bank_ids].contiguous() - self._gather_feat_bytes = self._copy_feat_bytes[self._gather_bank_ids].contiguous() + # The gather copies cache rows of the layer's pool into its staging buffer, row == expert id. + self._gather_bank_ids = [[k for k, r in enumerate(rows) if r >= _SMALL_BANK_FEAT_BYTES] for rows in self._row_bytes] + for g, ids in enumerate(self._gather_bank_ids): + self._gather_feat_bytes.append(ptrs([self._row_bytes[g][k] for k in ids]) if ids else None) + self._gather_stage_ptrs.append( + [ptrs([dst_ptrs[k] + b * self._stage_bytes[k] for k in ids]) for b in range(2)] if ids else None + ) + for p in range(len(self.pool_sizes)): + ids = self._gather_bank_ids[self._pool_layout(p)] + self._gather_pool_ptrs.append(ptrs([dst_ptrs[k] + self._pool_offsets[k][p] for k in ids]) if ids else None) + self._copy_blocks_per_bank = default_blocks_per_bank(self.device) self._copy_fused_ok = True - def validate_rebuild(self, cache_size: int) -> None: + def validate_rebuild(self, cache_size: int, pool_sizes: list[int] | None = None) -> None: """Pure geometry validation of a rebuild target (no GPU side effects). Raises ``ValueError`` if ``cache_size`` is below the ``num_experts`` floor or - above the marlin slot cap. Called by :meth:`rebuild` and by the engine's + a pool is above the expert kernel's slot cap, or the bank byte segments cannot stage one layer. Called by :meth:`rebuild` and by the engine's pre-teardown check, so an invalid target rejects with the old cache intact (no destructive free first). """ if cache_size < self.num_experts: raise ValueError(f"cache_size {cache_size} < num_experts {self.num_experts}") - if self.max_slots is not None and cache_size > self.max_slots: + # decode kernels see one pool's slots, so the cap applies per pool + largest = cache_size if pool_sizes is None else max(pool_sizes) + what = f"moe_cache_size={cache_size}" if largest == cache_size else f"a MoE slot pool of {largest} slots" + if self.max_slots is not None and largest > self.max_slots: raise ValueError( - f"moe_cache_size={cache_size} exceeds the expert kernel's slot limit of {self.max_slots}; " + f"{what} exceeds the expert kernel's slot limit of {self.max_slots}; " f"pass --moe-cache-size {self.max_slots} or less, or let the default kernel serve the experts" ) - if self.layout is None and self.quant_format == "nvfp4_marlin" and cache_size > MARLIN_MAX_CACHE_SIZE: + if self.layout is None and self.quant_format == "nvfp4_marlin" and largest > MARLIN_MAX_CACHE_SIZE: raise ValueError( - f"moe_cache_size={cache_size} exceeds the marlin backend's slot limit of " + f"{what} exceeds the marlin backend's slot limit of " f"{MARLIN_MAX_CACHE_SIZE} (vLLM moe_align_block_size caps padded experts at " "1024); reduce moe_cache_size or force --quant-backend moe.nvfp4=triton" ) + if self._raw_banks and pool_sizes is not None: + for k, stage in enumerate(self._stage_bytes): + if self._bank_bytes(k, pool_sizes) < stage: + raise ValueError(f"moe_cache_size={cache_size} cannot stage one whole layer of experts") - def rebuild(self, cache_size: int) -> None: + def rebuild(self, cache_size: int, pool_sizes: list[int] | None = None) -> None: """Resize the GPU slot cache + bookkeeping to ``cache_size`` IN PLACE. Keeps the CPU/pinned ``bank_sources`` and the GPU-resident alphas; never @@ -472,9 +908,10 @@ def rebuild(self, cache_size: int) -> None: ``ctx.moe_offload_cache`` stay valid. """ assert self.bank_sources, "set_bank_sources must run before rebuild" - self.validate_rebuild(cache_size) + sizes = [cache_size] if pool_sizes is None else list(pool_sizes) + assert len(sizes) == len(self.pool_sizes) and sum(sizes) == cache_size and min(sizes) > 0, (sizes, cache_size) + self.validate_rebuild(cache_size, sizes) # 1. Tear down prefill-overlap (its buffer views alias the old bank_caches). - self.prefill_bank_buffers = [] self.prefill_copy_stream = None self.prefill_begin_event = None self.prefill_ready_events = [] @@ -485,22 +922,24 @@ def rebuild(self, cache_size: int) -> None: # 2. Drop old GPU tensors (free-before-alloc). self.banks = [] self.bank_caches = {} + self._pool_bank_views = [] + self._stage_views = [] + self._stage_tails = [] self.cache_size = cache_size + self.pool_sizes = sizes if self.device.type == "cuda": torch.cuda.synchronize(self.device) torch.cuda.empty_cache() # 3. Reallocate the slot cache from the retained host sources. - for name in self.bank_schema: - head = self.bank_sources[name][0] - self.bank_caches[name] = torch.empty( - (cache_size, *head.shape[1:]), dtype=head.dtype, device=self.device - ) - self.banks = [(self.bank_sources[n], self.bank_caches[n]) for n in self.bank_schema] - self._build_copy_plan() # slot caches were reallocated -> refresh fused-copy addrs + self._alloc_banks() # 4. Reallocate cache_size-shaped bookkeeping; reset the slot map (cold start). self.slot_for_id.fill_(-1) self.id_of_slot = torch.full((cache_size,), -1, dtype=torch.int32, device=self.device) self.usage = torch.zeros((cache_size,), dtype=torch.int64, device=self.device) + self._slice_pools() + self._build_copy_plan() # slot caches were reallocated -> refresh fused-copy addrs + if self.verify is not None: + self.verify.bind(self) plan_slots = max(self.num_experts, cache_size) self.evict_slots = torch.empty((plan_slots,), dtype=torch.int32, device=self.device) self.src_indices = torch.empty((plan_slots,), dtype=torch.int32, device=self.device) @@ -509,6 +948,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,14 +961,18 @@ def rebuild(self, cache_size: int) -> None: self.stat_fetched_layer.zero_() self.stat_steps_layer.zero_() self.decode_freq.zero_() + if self.prefetch is not None: + self.prefetch.reset() + if self.audit is not None: + self.audit.bind(self) # new slot arrays: a new history self.prefill_hit_rows = 0 self.prefill_total_rows = 0 self._hit_d2d_fallback_logged = False # geometry changed; re-log if still unusable # 5. Re-evaluate prefill overlap against the new size. - if self.prefill_overlap and cache_size < 2 * self.num_experts: + if self.prefill_overlap and not self._stage_fits(2): logger.warning( f"Disabling MoE prefill overlap on rebuild: cache_size {cache_size} " - f"< 2*num_experts {2 * self.num_experts}." + f"cannot hold two whole layers of experts (2*num_experts {2 * self.num_experts} slots of one layout)." ) self.prefill_overlap = False if self.prefill_overlap: @@ -580,8 +1025,9 @@ def alphas_for_slots(self, layer_id: int) -> tuple[torch.Tensor, torch.Tensor] | ``layer_id`` -- are ever read by the grouped GEMM.""" if self.gate_up_alpha is None: return None + id_of_slot, _ = self.pool_state(layer_id) idx = layer_id * self.num_experts + ( - self.id_of_slot.clamp(min=0).long() % self.num_experts + id_of_slot.clamp(min=0).long() % self.num_experts ) return self.gate_up_alpha[idx], self.down_alpha[idx] @@ -595,25 +1041,12 @@ def alphas_for_layer(self, layer_id: int) -> tuple[torch.Tensor, torch.Tensor] | hi = lo + self.num_experts return self.gate_up_alpha[lo:hi], self.down_alpha[lo:hi] - def bank_views(self, n: int | None = None) -> tuple[torch.Tensor, ...]: - """Per-bank cache views in registration order: the full ``[S]`` slot cache - (decode), or its first ``n`` slots (materialized layer).""" - assert self.banks, "set_bank_sources must register the banks first" - if n is None: - return tuple(cache for _, cache in self.banks) - return tuple(cache[:n] for _, cache in self.banks) - def _init_prefill_overlap_buffers(self) -> None: assert self.banks, "set_bank_sources must register the banks first" self._prefill_buffer_layer = [None, None] self._prefill_buffer_released = [True, True] self._prefill_buffer_has_release_event = [False, False] - # The double buffers borrow the slot cache's first 2 * num_experts slots - # (one full expert layer per buffer), one view per registered bank. - self.prefill_bank_buffers = [ - cache[: 2 * self.num_experts].view(2, self.num_experts, *cache.shape[1:]) - for _, cache in self.banks - ] + # the double buffers are the staging views (_slice_pools), on top of the first pools' slots if self.device.type == "cuda": self.prefill_copy_stream = torch.cuda.Stream(device=self.device) self.prefill_ready_events = [torch.cuda.Event() for _ in range(2)] @@ -632,15 +1065,36 @@ def _init_prefill_overlap_buffers(self) -> None: ) self._prefill_hit_num = torch.zeros((1,), dtype=torch.int64, device=self.device) - def _invalidate_prefill_buffer(self, buffer_id: int) -> None: - slot_start = buffer_id * self.num_experts - slot_end = slot_start + self.num_experts - old_ids = self.id_of_slot[slot_start:slot_end] - self.slot_for_id.view(-1)[old_ids[old_ids >= 0].long()] = -1 - old_ids.fill_(-1) - # usage=0 makes these slots the oldest, so the argmin(usage) victim selection in - # ensure_experts evicts them first. - self.usage[slot_start:slot_end].zero_() + def _invalidate_prefill_buffer(self, buffer_id: int, layer_id: int) -> None: + """Take buffer ``buffer_id``'s slots for staging ``layer_id``.""" + slots = self._inval[buffer_id] + if isinstance(slots, slice): + if self.audit is not None: + from freetoken.moe.slot_audit import INVALIDATE + + self.audit.record_clear(self, INVALIDATE, layer_id, slots.start, slots.stop - slots.start) + # One fixed-shape launch. The boolean-mask index + # (slot_for_id[old_ids[old_ids >= 0]] = -1) has a data-dependent shape, so every + # call hid a device-to-host sync -- and with two buffer reuses per chunk x 48 + # layers the host waited, per layer, for ALL enqueued GPU work (the cached-context + # attention included), serializing the prefill-overlap pipeline. + from freetoken.kernel.triton.moe import invalidate_prefill_slots + + if slots.stop > slots.start: + invalidate_prefill_slots( + self.id_of_slot, self.slot_for_id, self.usage, slots.start, slots.stop - slots.start + ) + else: + old_ids = self.id_of_slot[slots] + self.slot_for_id.view(-1)[old_ids[old_ids >= 0].long()] = -1 + self.id_of_slot[slots] = -1 + # usage=0 makes these slots the oldest, so the argmin(usage) victim selection in + # ensure_experts evicts them first. + self.usage[slots] = 0 + self._stage_dirty[buffer_id] = True + # the pad may hold rows of a wider layout or of a pool + for tail in self._stage_tails[self.layer_layout[layer_id]][buffer_id]: + tail.zero_() def begin_prefill(self) -> None: if not self.prefill_overlap: @@ -671,7 +1125,7 @@ def prefetch_prefill_layer(self, layer_id: int) -> None: if layer_id < 0: raise ValueError(f"Invalid prefill layer id: {layer_id}") - assert self.banks and self.prefill_bank_buffers + assert self.banks and self._stage_views buffer_id = layer_id % 2 if self._prefill_buffer_layer[buffer_id] == layer_id: @@ -682,9 +1136,13 @@ def prefetch_prefill_layer(self, layer_id: int) -> None: ) def copy() -> None: - self._invalidate_prefill_buffer(buffer_id) - for (per_layer, _), buffer in zip(self.banks, self.prefill_bank_buffers): - buffer[buffer_id].copy_(per_layer[layer_id], non_blocking=True) + self._invalidate_prefill_buffer(buffer_id, layer_id) + for (per_layer, _), view in zip(self.banks, self.stage_views(layer_id, buffer_id)): + view.copy_(per_layer[layer_id], non_blocking=True) + if self.audit is not None: + from freetoken.moe.slot_audit import PREFILL_BUFFER + + self.audit.record_range(self, PREFILL_BUFFER, layer_id, buffer_id * self.num_experts, self.num_experts, meta=False) if self._prefill_hit_d2d_active: self._prefetch_split(layer_id, buffer_id) @@ -715,10 +1173,10 @@ def _hit_d2d_usable(self) -> bool: reason = "FREETOKEN_SKIP_FAST_INDEX_COPY is set (the hit gather would be a no-op)" elif not self._copy_fused_ok: reason = "the fused copy plan is unavailable (bank alignment or FREETOKEN_FUSED_COPY=0)" - elif self.cache_size <= 2 * self.num_experts: + elif all(stable >= s for stable, s in zip(self._pool_stable_from, self.pool_sizes)): reason = ( f"cache_size {self.cache_size} leaves no hit region " - f"(needs > {2 * self.num_experts} slots)" + f"(every slot lies under the prefill staging buffers of {2 * self.num_experts} experts)" ) elif not self._resolve_batch_memcpy(): reason = "cudaMemcpyBatchAsync is unavailable" # resolve logged the specifics @@ -756,7 +1214,7 @@ def _prefetch_split(self, layer_id: int, buffer_id: int) -> None: run list comes from the begin-of-chunk snapshot because the batch API takes HOST pointer arrays. Live-vs-snapshot cannot disagree: the only chunk-internal writer (buffer invalidation) rewrites slots already below - the 2E threshold, and slots < 2E (including -1) are misses on both sides + the pool's staging threshold, and those slots (and -1) are misses on both sides -- the buffers own those slots, so their bytes are volatile within the chunk. Hit and miss row sets are disjoint, so the streams need no ordering against each other. @@ -767,43 +1225,53 @@ def _prefetch_split(self, layer_id: int, buffer_id: int) -> None: from freetoken.moe.offload_kernels import prefill_hit_compact E = self.num_experts + layout = self.layer_layout[layer_id] + pool = self.pool_of_layer[layer_id] snap = self._prefill_snapshot_np[layer_id] - hit_mask = snap >= 2 * E + hit_mask = snap >= self._pool_stable_from[pool] self.prefill_hit_rows += int(hit_mask.sum()) self.prefill_total_rows += E - if self._gather_dst_ptrs is not None: - prefill_hit_compact(self, layer_id, buffer_id) + if self._gather_feat_bytes[layout] is not None: + prefill_hit_compact(self, layer_id) # blocks_per_bank=64 vs the PCIe-tuned default of 8: HBM D2D needs the # wider grid (~22 GB/s per 1024-thread block on H100). fast_index_copy_multi_jit( - self._gather_dst_ptrs, - self._gather_dst_ptrs, - self._gather_feat_bytes, + self._gather_stage_ptrs[layout][buffer_id], + self._gather_pool_ptrs[pool], + self._gather_feat_bytes[layout], self._prefill_hit_dst, self._prefill_hit_src, self._prefill_hit_num, blocks_per_bank=64, ) + if self.audit is not None: + from freetoken.moe.slot_audit import PREFILL_HIT_D2D + + self.audit.record_plan(self, PREFILL_HIT_D2D, layer_id, self._prefill_hit_dst, None, self._prefill_hit_num, + meta=False, aux=self._prefill_hit_src, row_base=buffer_id * E) miss = np.nonzero(~hit_mask)[0] with torch.cuda.stream(self.prefill_copy_stream): if self._prefill_buffer_has_release_event[buffer_id]: self.prefill_copy_stream.wait_event(self.prefill_release_events[buffer_id]) - self._invalidate_prefill_buffer(buffer_id) + self._invalidate_prefill_buffer(buffer_id, layer_id) if miss.size: run_starts = np.concatenate(([0], np.nonzero(np.diff(miss) != 1)[0] + 1)) starts = miss[run_starts] lengths = np.diff(np.concatenate((run_starts, [miss.size]))) dst, src, nbytes = [], [], [] - for b, feat in enumerate(self._copy_feat_bytes_host): + for b, feat in enumerate(self._row_bytes[layout]): + stage = self._copy_dst_ptrs_host[b] + buffer_id * self._stage_bytes[b] + if feat == 0: + continue if feat < _SMALL_BANK_FEAT_BYTES: # Whole layer as one entry, EVEN with zero misses: it keeps every # batch entry above the driver's async floor and covers the hit # rows the gather skips for these banks. - dst.append(self._copy_dst_ptrs_host[b] + buffer_id * E * feat) + dst.append(stage) src.append(self._copy_src_ptrs_host[layer_id][b]) nbytes.append(E * feat) elif miss.size: - dst.extend(self._copy_dst_ptrs_host[b] + (buffer_id * E + starts) * feat) + dst.extend(stage + starts * feat) src.extend(self._copy_src_ptrs_host[layer_id][b] + starts * feat) nbytes.extend(lengths * feat) if dst: @@ -813,6 +1281,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, ...]: @@ -821,15 +1294,23 @@ def wait_prefill_layer(self, layer_id: int) -> tuple[torch.Tensor, ...]: marlin/b12x ``(gate_up_packed, gate_up_scale, down_packed, down_scale)``; nvfp4 native adds the two global banks after each scale bank.""" assert self.prefill_overlap - assert self.prefill_bank_buffers + assert self._stage_views self.prefetch_prefill_layer(layer_id) buffer_id = layer_id % 2 assert self._prefill_buffer_layer[buffer_id] == layer_id if self.prefill_ready_events: torch.cuda.current_stream(self.device).wait_event(self.prefill_ready_events[buffer_id]) - return tuple(buffer[buffer_id] for buffer in self.prefill_bank_buffers) + return self.stage_views(layer_id, buffer_id) def release_prefill_layer(self, layer_id: int) -> None: + if layer_id == self.num_layers - 1: + # the staged slots were invalidated and nothing registers them within the forward, so their bytes are free to clear + for b in range(2): + if self._stage_dirty[b]: + # staged rows left there may be shifted or of another quant, and kernels still read them: MMQ past a slot's last row, weight-0 routes on slot 0 + for (_, cache), stage in zip(self.banks, self._stage_bytes): + cache.view(-1).view(torch.uint8)[b * stage : (b + 1) * stage].zero_() + self._stage_dirty[b] = False if not self.prefill_overlap: return buffer_id = layer_id % 2 @@ -840,19 +1321,80 @@ 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 loads_prefill_on_demand(self, num_tokens: int) -> bool: + return 0 < num_tokens <= self.prefill_demand_tokens - if self.collect_decode_freq: + def ensure_experts( + self, + layer_id: int, + expert_ids: torch.Tensor, + *, + update_state: bool = True, + router_logits: torch.Tensor | None = None, + pin_since: torch.Tensor | None = None, + counts_as_decode: bool = True, + ) -> None: + """``update_state=False`` (small-prefill ensures) keeps a scored policy's state and + decode-step clock still; ``router_logits`` (``[rows, num_experts]``) feed the ``rule`` + policy's near misses; ``pin_since`` (a ``step`` value) keeps every slot touched after it + (a small prefill's earlier chunks). All are ignored under LRU, whose recency order + already keeps those slots.""" + from freetoken.moe.offload_kernels import ensure_experts, ensure_experts_scored + + if self.collect_decode_freq and counts_as_decode and update_state: # ``expert_ids`` still holds raw expert ids here (the kernel rewrites them to # slot ids in place), so snapshot the routing histogram before that happens. ids = expert_ids.reshape(-1).long() self.decode_freq[layer_id].scatter_add_(0, ids, torch.ones_like(ids)) self._pending_src_layer = layer_id self._pending_whole_layer = False - ensure_experts(self, layer_id, expert_ids) + audit = self.audit + if audit is not None and update_state and layer_id == self.evict_clock_layer(): + audit.begin_step() + lowpri = self.prefetch_on + width = expert_ids.numel() + if self.cache_policy_id == 0 and not lowpri: + ensure_experts(self, layer_id, expert_ids, record_stats=counts_as_decode and update_state) + else: + # prefetch on: LRU runs the vendored kernel (flashlib's cannot key low-priority slots) + ensure_experts_scored( + self, + layer_id, + expert_ids, + bump_tok=update_state and layer_id == self.evict_clock_layer(), + update_state=update_state, + router_logits=router_logits, + pin_since=pin_since, + lowpri=lowpri, + pf_count=self.prefetch.count_args(layer_id) if lowpri and update_state else None, + ) + if audit is not None: + from freetoken.moe import slot_audit as sa + + kind = sa.DEMAND_INSTALL if update_state else sa.PREFILL_INSTALL + self._pending_kind = sa.DEMAND_COPY if update_state else sa.PREFILL_COPY + self._pending_width = width + audit.record_plan(self, kind, layer_id, self.evict_slots, self.src_indices, self.num_indices, meta=True, width=width) + + def prefetch_ensure( + self, layer_id: int, query: torch.Tensor, dst_slots: torch.Tensor, src_rows: torch.Tensor, + num: torch.Tensor, stats_row: torch.Tensor, ready: torch.Tensor, budget: int | None = None, + ) -> None: + """Install the first ``budget`` of ``query``'s non-resident experts of ``layer_id`` as + low-priority slots and plan their copy into ``dst_slots``/``src_rows``/``num`` (never the + demand plan, which copy(L-1) still reads). Runs on the compute stream like every other slot + map writer.""" + from freetoken.moe.offload_kernels import prefetch_ensure_experts - def ensure_experts_hybrid(self, layer_id: int, expert_ids: torch.Tensor) -> None: + 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, weights: torch.Tensor, clamp: bool = False + ) -> HybridRoutes: """Capped-fetch LRU for the hybrid backend. Like :meth:`ensure_experts` but assigns slots to (and schedules copies for) at @@ -861,7 +1403,8 @@ def ensure_experts_hybrid(self, layer_id: int, expert_ids: torch.Tensor) -> None left non-resident and ``expert_ids`` is rewritten to their cache slot (hit or freshly fetched) or ``-1`` (overflow -> compute on the CPU). ``num_indices`` holds the capped fetch count (for ``copy_missing``); ``num_missing_full`` the pre-cap - miss count (for stats). All device-side / fixed-shape, so it is CUDA-graph safe.""" + miss count (for stats). All device-side / fixed-shape, so it is CUDA-graph safe. + Returns the routes split between GPU and CPU (:class:`freetoken.moe.offload_kernels.HybridRoutes`).""" from freetoken.moe.offload_kernels import ensure_experts_hybrid if self.collect_decode_freq: @@ -869,24 +1412,44 @@ def ensure_experts_hybrid(self, layer_id: int, expert_ids: torch.Tensor) -> None self.decode_freq[layer_id].scatter_add_(0, ids, torch.ones_like(ids)) self._pending_src_layer = layer_id self._pending_whole_layer = False - ensure_experts_hybrid( - self, layer_id, expert_ids, self.hybrid_max_fetch, self.hybrid_fetch_fraction - ) + fraction = self.hybrid_fetch_fraction[layer_id] if self.hybrid_fetch_fraction is not None else 0.0 + return ensure_experts_hybrid(self, layer_id, expert_ids, weights, self.hybrid_max_fetch, fraction, clamp) def materialize_layer(self, layer_id: int) -> None: from freetoken.moe.offload_kernels import materialize_layer 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 @@ -908,23 +1471,6 @@ def record_decode_stats(self, layer_id: int) -> None: was eight torch ops per layer per step, all captured into the decode graph. """ - def record_decode_stats_hybrid(self, layer_id: int) -> None: - """Hybrid stats: full miss count (pre-cap), the PCIe-fetched count (capped), and - the active count. The CPU computes (missing - fetched) experts. Device-side; - accumulates both the scalar totals and the per-layer breakdown.""" - assert 0 <= layer_id < self.num_layers, f"layer_id {layer_id} out of range [0, {self.num_layers})" - missing = self.num_missing_full.sum() - fetched = self.num_indices.sum() - active = self.active_mask.sum() - self.stat_missing += missing - self.stat_fetched += fetched - self.stat_active += active - self.stat_calls += 1 - self.stat_missing_layer[layer_id] += missing - self.stat_fetched_layer[layer_id] += fetched - self.stat_active_layer[layer_id] += active - self.stat_steps_layer[layer_id] += 1 - def decode_miss_stats(self) -> dict: if self.decode_target == "hybrid": active = int(self.stat_active.item()) @@ -968,6 +1514,8 @@ def decode_miss_stats_per_layer(self) -> dict: s, m, a, f = steps[L], missing[L], active[L], fetched[L] per_layer.append({ "layer": L, + "pool": self.pool_of_layer[L], + "pool_slots": self.pool_slots(L), "steps": s, "active_per_step": (a / s) if s else 0.0, "missing_per_step": (m / s) if s else 0.0, @@ -979,10 +1527,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 +1540,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 +1552,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: @@ -1019,55 +1570,86 @@ def copy_missing(self) -> None: f"pageable materialize (position == expert id); ensure_experts's " f"LRU slot remap cannot be honored without a device alias" ) - # the only copy a non-pinned layer ever needs is the non-overlap prefill materialize, which schedules the whole layer into slots [0, num_experts) with position == expert id -- a plain synchronous pageable H2D copy + # the only copy a non-pinned layer ever needs is the non-overlap prefill materialize, which schedules the whole layer into staging buffer 0 with position == expert id -- a plain synchronous pageable H2D copy # never CUDA-graph captured: prefill is not captured, and decode never reaches this branch (it routes to the CPU executor) - for per_layer, cache in self.banks: - cache[: self.num_experts].copy_(per_layer[layer_id]) + for (per_layer, _), view in zip(self.banks, self.stage_views(layer_id, 0)): + view.copy_(per_layer[layer_id]) + if self.audit is not None: + from freetoken.moe.slot_audit import PAGEABLE_COPY + + self.audit.record_range(self, PAGEABLE_COPY, layer_id, 0, self.num_experts, meta=False) return + # materialize stages rows of staging buffer 0; ensure stages slots local to the layer's pool + self.copy_rows(layer_id, self.evict_slots, self.src_indices, self.num_indices, slim=_SLIM_COPY, + kind=self._pending_kind, width=self._pending_width, whole=self._pending_whole_layer) + + def copy_rows( + self, layer_id: int, dst_slots: torch.Tensor, src_rows: torch.Tensor, num: torch.Tensor, *, slim: bool, + kind: int = 0, width: int | None = None, whole: bool = False, + ) -> None: + """Copy host rows ``src_rows[:num]`` of ``layer_id``'s banks into slots ``dst_slots[:num]`` + on the current stream (the demand plan, or a prefetch plan on the prefetch copy stream); + ``kind`` labels the copy for the slot audit and ``width`` bounds ``num`` for it. ``dst_slots`` + are local to the layer's pool, or rows of staging buffer 0 when ``whole`` (materialize).""" + self._copy_rows(layer_id, dst_slots, src_rows, num, slim=slim, whole=whole) + if self.audit is not None: + from freetoken.moe.slot_audit import OTHER_COPY + + self.audit.record_plan(self, kind or OTHER_COPY, layer_id, dst_slots, src_rows, num, meta=False, width=width) + + def _copy_rows( + self, layer_id: int, dst_slots: torch.Tensor, src_rows: torch.Tensor, num: torch.Tensor, *, slim: bool, + whole: bool = False, + ) -> None: if self._copy_fused_ok: - from freetoken.kernel.fast_index_copy import fast_index_copy_multi_jit + from freetoken.kernel.fast_index_copy import fast_index_copy_multi_jit, fast_index_copy_multi_slim_jit # One launch copies the missing rows for every bank (instead of one launch per - # bank). evict_slots/src_indices/num_indices are shared across banks; - # src_indices holds layer-local expert rows, resolved against this layer's - # source pointers (layer_id is a static int per captured graph node). - fast_index_copy_multi_jit( - self._copy_dst_ptrs, - self._copy_src_ptrs[layer_id], - self._copy_feat_bytes, - self.evict_slots, - self.src_indices, - self.num_indices, - ) + # bank). The slot/row plan is shared across banks; src_rows holds layer-local + # expert rows, resolved against this layer's source pointers (layer_id is a + # static int per captured graph node). + dst_ptrs = self._copy_dst_ptrs if whole else self._copy_dst_ptrs_pool[self.pool_of_layer[layer_id]] + feat_bytes = self._copy_feat_bytes[self.layer_layout[layer_id]] + if slim: + fast_index_copy_multi_slim_jit(dst_ptrs, self._copy_src_ptrs[layer_id], feat_bytes, dst_slots, src_rows, num) + else: + fast_index_copy_multi_jit( + dst_ptrs, self._copy_src_ptrs[layer_id], feat_bytes, dst_slots, src_rows, num, + blocks_per_bank=self._copy_blocks_per_bank, + ) return from freetoken.kernel import fast_index_copy_jit - for per_layer, cache in self.banks: - fast_index_copy_jit( - cache, - self.evict_slots, - per_layer[layer_id], - self.src_indices, - self.num_indices, - ) + views = self.stage_views(layer_id, 0) if whole else self.pool_bank_views(layer_id) + for (per_layer, _), cache in zip(self.banks, views): + if cache.numel() == 0: + continue + 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..c94b1ca922 100644 --- a/python/freetoken/moe/offload_kernels.py +++ b/python/freetoken/moe/offload_kernels.py @@ -1,6 +1,7 @@ from __future__ import annotations import os +from typing import NamedTuple import torch import triton @@ -16,7 +17,7 @@ ) -def ensure_experts(cache, layer_id: int, expert_ids: torch.Tensor) -> None: +def ensure_experts(cache, layer_id: int, expert_ids: torch.Tensor, *, record_stats: bool = True) -> None: """Make this layer's routed experts resident; rewrite ``expert_ids`` to slot ids. Delegates to flashlib's slot cache. ``id_base`` maps this layer's expert ids into the @@ -25,7 +26,49 @@ def ensure_experts(cache, layer_id: int, expert_ids: torch.Tensor) -> None: tensor. ``out_indices`` aliases the input, preserving the in-place rewrite every downstream GEMM depends on. """ + id_of_slot, usage = cache.pool_state(layer_id) lru_ensure( + expert_ids, + cache.slot_for_id.view(-1), + id_of_slot, + usage, + cache.step, + expert_ids, + cache.src_indices, + cache.evict_slots, + cache.num_indices, + stats=cache.lru_stats[layer_id] if cache.collect_stats and record_stats else None, + id_base=layer_id * cache.num_experts, + ) + + +class HybridRoutes(NamedTuple): + """One layer's routes split between the GPU and the CPU, each [bs, top_k].""" + + gpu_slots: torch.Tensor # cache slot of each route; a CPU route is -1, or 0 with clamp + gpu_w: torch.Tensor # router weights, 0 for a CPU route + cpu_ids: torch.Tensor # expert ids, -1 for a GPU route + + +def ensure_experts_scored( + cache, + layer_id: int, + expert_ids: torch.Tensor, + *, + bump_tok: bool, + update_state: bool, + router_logits: torch.Tensor | None = None, + pin_since: torch.Tensor | None = None, + lowpri: bool = False, + pf_count: tuple | None = None, +) -> None: + """``ensure_experts`` for the scored ``--moe-cache-policy`` kinds (kd, kdfb, rule), and for + every policy under FREETOKEN_MOE_PREFETCH=on (``lowpri``; ``pf_count`` is the prefetcher's + ``(stats_row, ready, num, rows)`` when this layer was prefetched this step).""" + from freetoken.moe.scored_ensure import scored_ensure + + stats_row, ready, num, rows = pf_count if pf_count is not None else (None, None, None, 0) + scored_ensure( expert_ids, cache.slot_for_id.view(-1), cache.id_of_slot, @@ -35,14 +78,62 @@ def ensure_experts(cache, layer_id: int, expert_ids: torch.Tensor) -> None: 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 ensure_experts_hybrid( - cache, layer_id: int, expert_ids: torch.Tensor, max_fetch: int, fetch_fraction: float = 0.0 +def prefetch_ensure_experts( + cache, layer_id: int, query: torch.Tensor, dst_slots: torch.Tensor, src_rows: torch.Tensor, + num: torch.Tensor, stats_row: torch.Tensor, ready: torch.Tensor, budget: int | None = None, ) -> None: + """FREETOKEN_MOE_PREFETCH=on: install the first ``budget`` of ``query``'s non-resident experts + of ``layer_id`` as low-priority slots and plan their copy (scored_ensure.prefetch_ensure).""" + from freetoken.moe.scored_ensure import prefetch_ensure + + prefetch_ensure( + query, cache.slot_for_id.view(-1), cache.id_of_slot, cache.usage, cache.step, + dst_slots, src_rows, num, stats_row, ready, budget=budget, **_evict_state(cache, layer_id), + ) + + +def ensure_experts_hybrid( + cache, + layer_id: int, + expert_ids: torch.Tensor, + weights: torch.Tensor, + max_fetch: int, + fetch_fraction: float = 0.0, + clamp: bool = False, +) -> HybridRoutes: """Capped-fetch variant of ``ensure_experts`` (hybrid backend). Identical LRU bookkeeping, but only the first ``max_fetch`` of this step's missing @@ -52,21 +143,26 @@ def ensure_experts_hybrid( split (fraction = pcie_bw / cpu_bw): fetch ~fraction of the step's misses, rounded to the integer that makes the PCIe fetch and the CPU overflow compute finish closest to together. ``num_indices`` = capped fetch count (copy_missing); ``num_missing_full`` = - pre-cap miss count (stats).""" - # Q16 fixed point so the GPU kernel and the CPU reference cap identically (no float). + pre-cap miss count (stats). + + The same launch splits the routes: ``expert_ids`` becomes the GPU slots in place, and + ``clamp`` gives CPU routes slot 0 instead of -1 for GEMM kernels that read every slot.""" + assert expert_ids.is_cuda, "the hybrid ensure kernel runs on the GPU" + # Q16 fixed point: the kernel caps in integers. frac_q16 = min(1 << 16, max(0, round(fetch_fraction * (1 << 16)))) - if not expert_ids.is_cuda: - return _ensure_experts_hybrid_cpu(cache, layer_id, expert_ids, max_fetch, frac_q16) - _ensure_experts_hybrid_gpu(cache, layer_id, expert_ids, max_fetch, frac_q16) + weights = weights.contiguous() + routes = HybridRoutes(expert_ids, torch.empty_like(weights), torch.empty_like(expert_ids)) + _ensure_experts_hybrid_gpu(cache, layer_id, expert_ids, weights, routes, clamp, max_fetch, frac_q16) + return routes -def prefill_hit_compact(cache, layer_id: int, buffer_id: int) -> None: +def prefill_hit_compact(cache, layer_id: int) -> None: """Compact this layer's cache-resident experts into gather indices, device-side. - hit = slot_for_id[layer_id][e] >= 2 * num_experts (the double buffer owns the + hit = slot_for_id[layer_id][e] >= the pool's staging threshold (the double buffer owns the slots below, so those bytes are volatile within a prefill chunk and classify as miss). Writes fixed-shape ``_prefill_hit_dst``/``_prefill_hit_src`` (buffer - row / cache slot) and the count into ``_prefill_hit_num``; one launch on the + row / pool-local slot) and the count into ``_prefill_hit_num``; one launch on the current stream, no host sync. Safe against the concurrent buffer invalidation on the copy stream: that only rewrites entries already below the threshold.""" num_experts = cache.num_experts @@ -75,15 +171,23 @@ def prefill_hit_compact(cache, layer_id: int, buffer_id: int) -> None: cache._prefill_hit_dst, cache._prefill_hit_src, cache._prefill_hit_num, - buffer_id * num_experts, - 2 * num_experts, + cache._pool_stable_from[cache.pool_of_layer[layer_id]], num_experts, BLOCK=triton.next_power_of_2(num_experts), ) def materialize_layer(cache, layer_id: int) -> None: - _materialize_layer_gpu(cache, layer_id) + if len(cache.pool_sizes) == 1: + _materialize_layer_gpu(cache, layer_id) + return + # staging buffer 0 may overlap several pools, so it only holds the layer for this forward and registers nothing + E = cache.num_experts + cache._invalidate_prefill_buffer(0, layer_id) + rows = torch.arange(E, dtype=torch.int32, device=cache.device) + cache.evict_slots[:E].copy_(rows) + cache.src_indices[:E].copy_(rows) + cache.num_indices.fill_(E) def reset_cache(cache) -> None: @@ -95,16 +199,25 @@ def reset_cache(cache) -> None: def _ensure_experts_hybrid_gpu( - cache, layer_id: int, expert_ids: torch.Tensor, max_fetch: int, frac_q16: int + cache, + layer_id: int, + expert_ids: torch.Tensor, + weights: torch.Tensor, + routes: HybridRoutes, + clamp: bool, + max_fetch: int, + frac_q16: int, ) -> None: + id_of_slot, usage = cache.pool_state(layer_id) + num_slots = id_of_slot.numel() block_e = triton.next_power_of_2(cache.num_experts) - block_c = triton.next_power_of_2(cache.cache_size) + block_c = triton.next_power_of_2(num_slots) num_warps = 8 if block_c >= 2048 else 4 _ensure_experts_hybrid_kernel[(1,)]( expert_ids, cache.slot_for_id, - cache.id_of_slot, - cache.usage, + id_of_slot, + usage, cache.step, cache.active_mask, cache.evict_slots, @@ -112,82 +225,25 @@ def _ensure_experts_hybrid_gpu( cache.num_indices, cache.num_missing_full, cache.expert_recency, + weights, + routes.cpu_ids, + routes.gpu_w, + *(cache.hybrid_stats if cache.collect_stats else (expert_ids,) * 2), layer_id, expert_ids.numel(), int(max_fetch), int(frac_q16), cache.num_experts, - cache.cache_size, + num_slots, BLOCK_E=block_e, BLOCK_C=block_c, BY_RECENCY=_HYBRID_FETCH_BY_RECENCY, + CLAMP=clamp, + STATS=cache.collect_stats, num_warps=num_warps, ) -def _ensure_experts_hybrid_cpu( - cache, layer_id: int, expert_ids: torch.Tensor, max_fetch: int, frac_q16: int -) -> None: - """CPU reference mirror of the hybrid kernel (eviction/fetch decisions bit-identical to - the GPU path; see tests/test_offload_lru_kernels.py). Fetches at most ``max_fetch`` (or - the bandwidth-matched ``~frac_q16/2^16 * misses`` when ``frac_q16`` > 0) of the missing - experts; overflow misses are rewritten to -1. With ``BY_RECENCY`` the fetch set is the - most-recently-active misses (ties -> lower id); else the lowest ids.""" - seen = [] - for expert in expert_ids.view(-1).tolist(): - if expert not in seen: - seen.append(expert) - - cache.active_mask.zero_() - step = int(cache.step.item()) + 1 - cache.step.fill_(step) - for expert in seen: - cache.active_mask[expert] = 1 - - for expert in seen: - slot = int(cache.slot_for_id[layer_id, expert].item()) - if slot != -1: - cache.usage[slot] = step - - missing = [e for e in seen if int(cache.slot_for_id[layer_id, e].item()) == -1] - if _HYBRID_FETCH_BY_RECENCY: - rec = cache.expert_recency[layer_id].tolist() - missing.sort(key=lambda e: (-rec[e], e)) - else: - missing.sort() - if frac_q16 > 0: - m, q = len(missing), 1 << 16 - lo = (m * frac_q16) >> 16 - cost = lambda f: max(f * (q - frac_q16), (m - f) * frac_q16) # noqa: E731 - max_fetch = lo if cost(lo) <= cost(lo + 1) else lo + 1 - num_fetch = min(len(missing), int(max_fetch)) - cache.num_missing_full.fill_(len(missing)) - cache.num_indices.fill_(num_fetch) - - usage = cache.usage.tolist() - for idx in range(num_fetch): - expert = missing[idx] - victim = min(range(cache.cache_size), key=lambda s: (usage[s], s)) - old_id = int(cache.id_of_slot[victim].item()) - if old_id >= 0: - cache.slot_for_id.view(-1)[old_id] = -1 - cache.id_of_slot[victim] = layer_id * cache.num_experts + expert - cache.slot_for_id[layer_id, expert] = victim - cache.usage[victim] = step - usage[victim] = step - cache.evict_slots[idx] = victim - cache.src_indices[idx] = expert # layer-local row - - if _HYBRID_FETCH_BY_RECENCY: - for expert in seen: - cache.expert_recency[layer_id, expert] = step - - # Overflow misses keep slot_for_id == -1, so the rewrite below yields -1 for them. - flat = expert_ids.view(-1) - for i in range(flat.numel()): - flat[i] = int(cache.slot_for_id[layer_id, int(flat[i].item())].item()) - - def _materialize_layer_gpu(cache, layer_id: int) -> None: block = triton.next_power_of_2(max(cache.num_experts, cache.cache_size)) _materialize_layer_kernel[(1,)]( @@ -287,7 +343,7 @@ def _materialize_layer_kernel( -@triton.jit(do_not_specialize=["layer_id", "num_active", "max_fetch", "fetch_frac_q16"]) +@triton.jit(do_not_specialize=["layer_id", "num_active", "max_fetch", "fetch_frac_q16", "cache_size", "id_of_slot_ptr", "usage_ptr"]) def _ensure_experts_hybrid_kernel( expert_ids_ptr, slot_for_id_ptr, @@ -300,15 +356,22 @@ def _ensure_experts_hybrid_kernel( num_indices_ptr, num_missing_full_ptr, expert_recency_ptr, + weights_ptr, + cpu_ids_ptr, + gpu_w_ptr, + stats_ptr, + stats_layer_ptr, layer_id, num_active, max_fetch, fetch_frac_q16, num_experts: tl.constexpr, - cache_size: tl.constexpr, + cache_size, BLOCK_E: tl.constexpr, BLOCK_C: tl.constexpr, BY_RECENCY: tl.constexpr, + CLAMP: tl.constexpr, + STATS: tl.constexpr, ): """Capped-fetch timestamp-LRU (hybrid backend). @@ -352,10 +415,11 @@ def _ensure_experts_hybrid_kernel( (lo + 1) * ((1 << 16) - fetch_frac_q16), (num_missing - lo - 1) * fetch_frac_q16 ) max_fetch = tl.where(cost_lo <= cost_hi, lo, lo + 1) - num_fetch = tl.minimum(num_missing, max_fetch) + is_hit = is_active & (slot >= 0) + # a pool smaller than num_experts may hold fewer free slots than misses + num_fetch = tl.minimum(tl.minimum(num_missing, max_fetch), cache_size - tl.sum(is_hit.to(tl.int32))) tl.store(num_missing_full_ptr, num_missing.to(tl.int64)) tl.store(num_indices_ptr, num_fetch.to(tl.int64)) - is_hit = is_active & (slot >= 0) tl.store(usage_ptr + slot, step, mask=is_hit) # Fetch-selection priority: encode (recency desc, id asc) into one strictly-ordered @@ -401,8 +465,26 @@ def _ensure_experts_hybrid_kernel( for i in tl.range(num_active): e = tl.load(expert_ids_ptr + i) s = tl.load(slot_for_id_ptr + base + e) + on_gpu = s >= 0 + tl.store(cpu_ids_ptr + i, tl.where(on_gpu, -1, e)) + tl.store(gpu_w_ptr + i, tl.where(on_gpu, tl.load(weights_ptr + i), 0.0)) + if CLAMP: + s = tl.maximum(s, 0) tl.store(expert_ids_ptr + i, s) + # stats_ptr: (missing, fetched, active, calls); stats_layer_ptr: the same per layer, [num_layers, 4] + if STATS: + n_active = tl.sum((is_active & e_mask).to(tl.int64)) + tl.store(stats_ptr + 0, tl.load(stats_ptr + 0) + num_missing) + tl.store(stats_ptr + 1, tl.load(stats_ptr + 1) + num_fetch) + tl.store(stats_ptr + 2, tl.load(stats_ptr + 2) + n_active) + tl.store(stats_ptr + 3, tl.load(stats_ptr + 3) + 1) + row = stats_layer_ptr + layer_id * 4 + tl.store(row + 0, tl.load(row + 0) + num_missing) + tl.store(row + 1, tl.load(row + 1) + num_fetch) + tl.store(row + 2, tl.load(row + 2) + n_active) + tl.store(row + 3, tl.load(row + 3) + 1) + # Bump every active expert's recency to this step (LRU on the expert): an overflow miss # computed on the CPU now ranks high if it recurs, so it gets fetched next time. if BY_RECENCY: @@ -410,14 +492,13 @@ def _ensure_experts_hybrid_kernel( tl.store(expert_recency_ptr + base + off_e, step_vec, mask=is_active & e_mask) -@triton.jit(do_not_specialize=["buffer_base"]) +@triton.jit(do_not_specialize=["threshold"]) def _prefill_hit_compact_kernel( slot_ptr, # [num_experts] int32: this layer's slot_for_id row - dst_ptr, # [num_experts] int32 out: buffer rows, compacted - src_ptr, # [num_experts] int32 out: cache slots, compacted + dst_ptr, # [num_experts] int32 out: buffer rows (expert ids), compacted + src_ptr, # [num_experts] int32 out: pool-local slots, compacted num_ptr, # [1] int64 out: hit count - buffer_base, # buffer_id * num_experts - threshold, # 2 * num_experts + threshold, # first pool-local slot outside the staging buffers num_experts, BLOCK: tl.constexpr, ): @@ -426,6 +507,6 @@ def _prefill_hit_compact_kernel( slots = tl.load(slot_ptr + offs, mask=lane, other=-1) is_hit = lane & (slots >= threshold) pos = tl.cumsum(is_hit.to(tl.int32)) - 1 - tl.store(dst_ptr + pos, (buffer_base + offs).to(tl.int32), mask=is_hit) + tl.store(dst_ptr + pos, offs.to(tl.int32), mask=is_hit) tl.store(src_ptr + pos, slots, mask=is_hit) tl.store(num_ptr, tl.sum(is_hit.to(tl.int64))) diff --git a/python/freetoken/moe/pool_policy.py b/python/freetoken/moe/pool_policy.py new file mode 100644 index 0000000000..990c18ade2 --- /dev/null +++ b/python/freetoken/moe/pool_policy.py @@ -0,0 +1,149 @@ +"""Slot pool policy for the offload MoE cache: which layers share a pool, and how many slots each pool gets. + +Pure functions. The engine supplies the constraints and hands the result to ``OffloadMoeCache``: + + groups = group_layers(scope, layouts, slot_bytes) # --moe-cache-scope, once per load + pools = make_pools(groups, num_experts, min_num_slots, slot_limit) # the minimum follows the decode batch + sizes = split_slots(total, pools, counts) # --moe-cache-split, counts from resolve_profile + OffloadMoeCache(pool_of_layer=pool_of_layer(pools, num_layers), pool_sizes=sizes, ...) +""" + +from __future__ import annotations + +import json +from dataclasses import dataclass +from pathlib import Path + +import torch + +from freetoken.utils import init_logger + +logger = init_logger(__name__) + +SCOPES = ("unified", "per-layer") + + +@dataclass(frozen=True) +class PoolSpec: + layers: tuple[int, ...] + min_num_slots: int + max_num_slots: int + + +def group_layers(scope: str, layouts: list, slot_bytes: list[int]) -> list[list[int]]: + """The layers of each pool, in placement order. + + ``unified``: one pool per layout key (see ``bank_layouts``); ``per-layer``: one pool per layer. + Pools with wider slots come first: the prefill staging buffers sit at the start of every bank, so they then cover fewer slots. + """ + if scope == "unified": + groups: dict = {} + for layer, layout in enumerate(layouts): + groups.setdefault(layout, []).append(layer) + pools = list(groups.values()) + elif scope == "per-layer": + pools = [[layer] for layer in range(len(layouts))] + else: + raise ValueError(f"unknown --moe-cache-scope {scope!r}; want one of {SCOPES}") + return sorted(pools, key=lambda layers: (-slot_bytes[layers[0]], layers[0])) + + +def make_pools(groups: list[list[int]], num_experts: int, min_num_slots: int, slot_limit: int | None) -> list[PoolSpec]: + """Each group's slot range: at least ``min_num_slots``, at most one slot per expert of its layers and the kernel's ``slot_limit``.""" + pools = [] + for layers in groups: + max_num_slots = len(layers) * num_experts + if slot_limit is not None: + max_num_slots = min(max_num_slots, slot_limit) + pools.append(PoolSpec(tuple(layers), min(min_num_slots, max_num_slots), max_num_slots)) + return pools + + +def pool_of_layer(pools: list[PoolSpec], num_layers: int) -> list[int]: + """Each layer's pool index, the form ``OffloadMoeCache`` takes.""" + owner = [-1] * num_layers + for p, pool in enumerate(pools): + for layer in pool.layers: + owner[layer] = p + assert min(owner) >= 0, owner + return owner + + +def split_slots(total: int, pools: list[PoolSpec], counts: torch.Tensor | None = None) -> list[int]: + """Slots per pool, each in [min_num_slots, max_num_slots], summing to ``total``. + + With ``counts`` (per-layer expert use from ``resolve_profile``) the slots follow expert use, else the pools' layer counts. + """ + if len(pools) == 1: + return [total] + _check_budget(total, pools) + return _split_equal(total, pools) if counts is None else _split_hot(total, pools, counts) + + +def _check_budget(total: int, pools: list[PoolSpec]) -> None: + need = sum(p.min_num_slots for p in pools) + room = sum(p.max_num_slots for p in pools) + if total < need: + raise ValueError(f"MoE cache of {total} slots is below what the pools need ({need} slots)") + if total > room: + raise ValueError(f"MoE cache of {total} slots exceeds what the pools can hold ({room} slots)") + + +def _split_equal(total: int, pools: list[PoolSpec]) -> list[int]: + sizes = [p.min_num_slots for p in pools] + weights = [len(p.layers) for p in pools] + left = total - sum(sizes) + while left > 0: + open_ = [i for i, p in enumerate(pools) if sizes[i] < p.max_num_slots] + # give the next slot to the pool with the fewest slots per layer + i = min(open_, key=lambda j: (sizes[j] / weights[j], j)) + sizes[i] += 1 + left -= 1 + return sizes + + +def _split_hot(total: int, pools: list[PoolSpec], counts: torch.Tensor) -> list[int]: + sizes = [p.min_num_slots for p in pools] + # a pool's minimum covers its hottest experts; the rest, up to its maximum, competes across pools + cand = [] + for p, pool in enumerate(pools): + hot = counts[list(pool.layers)].flatten().sort(descending=True).values + cand.extend((float(c), p) for c in hot[pool.min_num_slots : pool.max_num_slots]) + cand.sort(key=lambda c: -c[0]) + for _, p in cand[: total - sum(sizes)]: + sizes[p] += 1 + return sizes + + +def resolve_profile(spec: str, model_path: str, num_layers: int, num_experts: int) -> torch.Tensor | None: + """Expert use for ``--moe-cache-split``: None for ``equal``, else the profile at ``profile:``.""" + if spec == "equal": + return None + path = Path(spec.removeprefix("profile:")) + logger.info_rank0(f"MoE pool profile: {path}") + return load_profile(path, num_layers, num_experts, model_path) + + +def load_profile(path: str | Path, num_layers: int, num_experts: int, model_path: str = "") -> torch.Tensor: + """Per-layer expert use ``[num_layers, num_experts]`` from a profile, each row scaled to sum to 1.""" + data = json.loads(Path(path).read_text()) + model = str(data.get("model", "")) + if model and model_path and not any(n.lower().startswith(model.lower()) for n in _model_names(model_path)): + logger.warning_rank0(f"MoE pool profile {path} was taken on {model!r}, not on {model_path}") + counts = torch.tensor(data["expert_counts"], dtype=torch.float64) + if tuple(counts.shape) != (num_layers, num_experts): + raise ValueError( + f"profile {path} has {tuple(counts.shape)} expert counts; this model needs ({num_layers}, {num_experts})" + ) + # every layer decodes the same tokens, so equal row sums keep the layers comparable + return counts / counts.sum(dim=1, keepdim=True).clamp(min=1e-12) + + +def _model_names(model_path: str) -> list[str]: + # every path component, so an HF cache snapshot (models--Org--Name/snapshots/) still finds its name + names = [] + for part in reversed(Path(model_path).parts): + name = part.split("--")[-1].removesuffix(".gguf") + if name: + names.append(name) + return names diff --git a/python/freetoken/moe/prefetch.py b/python/freetoken/moe/prefetch.py new file mode 100644 index 0000000000..01129d3997 --- /dev/null +++ b/python/freetoken/moe/prefetch.py @@ -0,0 +1,438 @@ +"""Cross-layer expert prefetch for GPU decode (``FREETOKEN_MOE_PREFETCH``). + +A training-free router lookahead predicts layer L's experts from the tensor layer L-1's router +consumed (``W_gate[L] . x_{L-1}``, top-K) and keeps the first BUDGET candidates that are not +resident. At bs 2 both rows' top-K lists merge rank by rank (row 0's first, then row 1's +first, ...), duplicates dropped, under one budget. + +``measure``: after layer L's ensure, count how many of them layer L actually routed. Nothing is +copied and the cache is never touched. The predictor reads ``slot_for_id[L]`` without ordering +against the ensures; a late read only skews the counters, never the model. + +The prediction forks at the start of layer L-1's MoE block, where x_{L-1} first exists, so it runs +beside layer L-1's own router and ensure: + + compute: -fork(L)-> router_{L-1} ensure(L-1) copy(L-1) GEMM(L-1) ... ensure(L) copy(L) -join-> count(L) GEMM(L) + predict: \\-> router_L(x_{L-1}) -> top-K -> select(L) -rec sel(L) -/ + +``on``: the predictor only merges the candidates (it reads no cache state); the compute stream +installs the first BUDGET of them that are not resident as low-priority slots (prefetch_ensure: +held, usage 0, keyed between the empty slots and every resident) and a third stream copies them +from layer L's host banks with the slim copy kernel: + + compute: -fork(L)-> ... ensure(L-1) [wait sel(L)] prefetch_ensure(L) [wait pfcopy(L-1)] copy(L-1) -rec dcopy(L)-> GEMM(L-1) ... ensure(L) [wait pfcopy(L)] copy(L) GEMM(L) + predict: \\-> router_L -> top-K -> merge(L) -rec sel(L) + copy: [wait dcopy(L)] slim copy(L) -> ready(L) -rec pfcopy(L) + +Every writer of the slot maps and their mirrors runs on the compute stream, so the maps have one +writer in stream order and no cross-stream edge guards them, and the installs follow layer L's +residency right after ensure(L-1) however late the predictor runs. dcopy(L) follows +prefetch_ensure(L) and the demand copy(L-1) on the compute stream, so the copy stream reads a +finished plan and owns the host link second. prefetch_ensure(L) keeps usage == step, i.e. layer +L-1's routed slots, pinned while GEMM(L-1) may read them; the compute stream joins the prefetch copy +before its own copy(L), so ensure(L) may reuse a prefetched slot as a demand victim and GEMM(L) only +ever reads landed bytes. A demand hit on a prefetched slot sets its usage (a normal resident) and +counts as useful. Layer 0 is never predicted. +""" +from __future__ import annotations + +import torch +import triton +import triton.language as tl + +from freetoken.env import ENV +from freetoken.moe.slot_audit import PREFETCH_COPY +from freetoken.utils import init_logger + +logger = init_logger(__name__) + +MODES = ("off", "measure", "on") +# non-resident candidates per target layer: GDN layers leave a shorter copy window than QSA ones +GDN_BUDGET = 3 +ATTN_BUDGET = 4 +# counters[layer, col]; on mode adds LATE and COPIED (calls whose plan had a copy) in a wider stats row +ISSUED, USEFUL, RESIDENT_HITS, CALLS, ROWS, MISSES, LATE, COPIED = range(8) +NUM_COLS = 6 +NUM_STAT_COLS = 8 +# the select kernel's dedup tile is (rows * k)^2; larger (eager-only) decode batches skip the prediction +MAX_ROWS = 8 + + +def resolve_mode(mode: str | None = None) -> str: + mode = (ENV.MOE_PREFETCH.value if mode is None else mode).strip().lower() + if mode not in MODES: + raise ValueError(f"FREETOKEN_MOE_PREFETCH={mode!r}: expected one of {', '.join(MODES)}") + return mode + + +def default_budget(before_linear: bool) -> int: + """Non-resident candidates kept for a target layer, by the target's attention kind.""" + return GDN_BUDGET if before_linear else ATTN_BUDGET + + +def dedicated_stream(device: torch.device, *, highest_priority: bool = False) -> torch.cuda.ExternalStream: + # not one of torch's 32 pooled streams: a pooled one can alias the engine or capture stream + from cuda.bindings import runtime as cudart + + with torch.cuda.device(device): + if highest_priority: + err, _, greatest = cudart.cudaDeviceGetStreamPriorityRange() + if err == cudart.cudaError_t.cudaSuccess: + err, handle = cudart.cudaStreamCreateWithPriority(cudart.cudaStreamNonBlocking, greatest) + else: + err, handle = cudart.cudaStreamCreateWithFlags(cudart.cudaStreamNonBlocking) + if err != cudart.cudaError_t.cudaSuccess: + raise RuntimeError(f"creating a dedicated CUDA stream failed: {err}") + return torch.cuda.ExternalStream(int(handle), device=device) + + +@triton.jit(do_not_specialize=["stride", "budget"]) +def _lookahead_select_kernel( + logits_ptr, stride, resident_ptr, sel_ptr, res_ptr, budget, + ROWS: tl.constexpr, ROWS_P: tl.constexpr, E: tl.constexpr, E_P: tl.constexpr, + K: tl.constexpr, N_P: tl.constexpr, W: tl.constexpr, W_P: tl.constexpr, +): + b = tl.arange(0, ROWS_P) + e = tl.arange(0, E_P) + live = (b < ROWS)[:, None] & (e < E)[None, :] + v = tl.load(logits_ptr + b[:, None] * stride + e[None, :], mask=live, other=float("-inf")).to(tl.float32) + # candidates in merged order: position r * ROWS + b holds row b's rank-r id + j = tl.arange(0, N_P) + cand = tl.full([N_P], -1, tl.int32) + for r in tl.static_range(K): + top = tl.argmax(v, axis=1).to(tl.int32) # ties go to the lowest id + at = (j[:, None] == r * ROWS + b[None, :]) & (b < ROWS)[None, :] + cand = tl.where(tl.sum(at.to(tl.int32), axis=1) > 0, tl.sum(tl.where(at, top[None, :], 0), axis=1), cand) + v = tl.where(e[None, :] == top[:, None], float("-inf"), v) + valid = j < ROWS * K + earlier = (cand[:, None] == cand[None, :]) & (j[None, :] < j[:, None]) & valid[None, :] + first = valid & (tl.sum(earlier.to(tl.int32), axis=1) == 0) + held = tl.load(resident_ptr + cand, mask=first, other=-1) >= 0 + pick = first & (~held) + rank = tl.cumsum(pick.to(tl.int32), axis=0) - 1 + kept = first & held & (tl.cumsum(first.to(tl.int32), axis=0) <= budget) + krank = tl.cumsum(kept.to(tl.int32), axis=0) - 1 + p = tl.arange(0, W_P) + m_sel = pick[None, :] & (rank[None, :] == p[:, None]) & (p < budget)[:, None] + m_res = kept[None, :] & (krank[None, :] == p[:, None]) + # at most one candidate matches each output position + sel = tl.sum(tl.where(m_sel, cand[None, :] + 1, 0), axis=1) - 1 + res = tl.sum(tl.where(m_res, cand[None, :] + 1, 0), axis=1) - 1 + tl.store(sel_ptr + p, sel, mask=p < W) + tl.store(res_ptr + p, res, mask=p < W) + + +@triton.jit(do_not_specialize=["id_base"]) +def _prefetch_count_kernel( + slots_ptr, id_of_slot_ptr, misses_ptr, sel_ptr, res_ptr, counters_ptr, id_base, + N: tl.constexpr, N_P: tl.constexpr, W: tl.constexpr, W_P: tl.constexpr, ROWS: tl.constexpr, +): + i = tl.arange(0, N_P) + slot = tl.load(slots_ptr + i, mask=i < N, other=-1) + # ensure rewrote the routed ids to slots, all pinned to this layer's ids until the next ensure + routed = tl.where(slot >= 0, tl.load(id_of_slot_ptr + slot, mask=slot >= 0, other=0) - id_base, -1) + p = tl.arange(0, W_P) + sel = tl.load(sel_ptr + p, mask=p < W, other=-1) + res = tl.load(res_ptr + p, mask=p < W, other=-1) + sel_hit = tl.sum(((sel[:, None] == routed[None, :]) & (i < N)[None, :]).to(tl.int32), axis=1) > 0 + res_hit = tl.sum(((res[:, None] == routed[None, :]) & (i < N)[None, :]).to(tl.int32), axis=1) > 0 + issued = tl.sum((sel >= 0).to(tl.int32)) + useful = tl.sum((sel_hit & (sel >= 0)).to(tl.int32)) + res_hits = tl.sum((res_hit & (res >= 0)).to(tl.int32)) + misses = tl.load(misses_ptr) + c = tl.arange(0, 8) + add = tl.where(c == 0, issued, tl.where(c == 1, useful, tl.where(c == 2, res_hits, tl.where(c == 3, 1, ROWS)))) + add = tl.where(c == 5, misses, add.to(tl.int64)) + # one writer per layer row, stream-ordered: no atomics needed + old = tl.load(counters_ptr + c, mask=c < 6, other=0) + tl.store(counters_ptr + c, old + add, mask=c < 6) + + +def lookahead_select( + logits: torch.Tensor, resident: torch.Tensor, sel: torch.Tensor, res: torch.Tensor, *, k: int, budget: int +) -> None: + """Top-``k`` of each ``logits`` row, merged rank by rank; ``sel`` gets the first ``budget`` + distinct ids with ``resident[id] < 0`` and ``res`` the resident ids among the first ``budget`` + distinct ones, both in merge order and padded with -1. One CTA, no host sync.""" + rows, num_experts = logits.shape + width = sel.numel() + assert logits.stride(1) == 1 and resident.numel() == num_experts and resident.dtype == torch.int32 + assert sel.dtype == res.dtype == torch.int32 and res.numel() == width and 1 <= k <= num_experts + _lookahead_select_kernel[(1,)]( + logits, logits.stride(0), resident, sel, res, min(budget, width), + ROWS=rows, ROWS_P=triton.next_power_of_2(rows), + E=num_experts, E_P=triton.next_power_of_2(num_experts), + K=k, N_P=triton.next_power_of_2(rows * k), + W=width, W_P=triton.next_power_of_2(width), + num_warps=4, + ) + + +def prefetch_count( + slots: torch.Tensor, id_of_slot: torch.Tensor, misses: torch.Tensor, + sel: torch.Tensor, res: torch.Tensor, counters: torch.Tensor, *, id_base: int, rows: int, +) -> None: + """Add one call's issued / useful / resident-hit / call / row / miss counts to ``counters``.""" + assert slots.dtype == torch.int32 and slots.is_contiguous() and counters.dtype == torch.int64 + n, width = slots.numel(), sel.numel() + _prefetch_count_kernel[(1,)]( + slots, id_of_slot, misses, sel, res, counters, id_base, + N=n, N_P=triton.next_power_of_2(n), W=width, W_P=triton.next_power_of_2(width), ROWS=rows, + num_warps=1, + ) + + +@triton.jit +def _mark_ready_kernel(ready_ptr): + tl.store(ready_ptr, 1) + + +class ExpertPrefetcher: + """Predictor stream, per-layer buffers and device counters for the router lookahead. + + ``fork`` runs layer L's prediction on the dedicated stream at the start of layer L-1's MoE + block. ``measure``: ``join_and_count`` waits for it after layer L's copy is enqueued and counts + on the compute stream. ``on``: ``install`` (right after layer L-1's ensure) waits for it and + runs prefetch_ensure(L) on the compute stream; ``issue_copy`` puts its copy on a second + dedicated stream after layer L-1's demand copy; the compute stream joins the copy before its + copy(L) (``join_copy``) and ensure(L) counts the call (``count_args``). The events are created + up front, the shapes are fixed and nothing syncs the host, so every side captures into the + decode graph. Nothing here is cache_size-shaped: a rebuild keeps it and only resets the counters. + """ + + def __init__( + self, num_layers: int, num_experts: int, device: torch.device, *, + mode: str = "measure", k: int | None = None, budget: int | None = None, + ) -> None: + self.mode = mode + self.num_layers = num_layers + self.num_experts = num_experts + self.device = device + k = ENV.MOE_PREFETCH_K.value if k is None else int(k) + if k < 1: + raise ValueError(f"FREETOKEN_MOE_PREFETCH_K={k} must be >= 1") + self.k = min(k, num_experts) + # 0 = the per-layer default the model wires in + self.budget_override = ENV.MOE_PREFETCH_BUDGET.value if budget is None else int(budget) + if self.budget_override < 0: + raise ValueError(f"FREETOKEN_MOE_PREFETCH_BUDGET={self.budget_override} must be >= 0") + self.stream = dedicated_stream(device) + # (forked, selected): the predictor starts after forked and records selected after the select + self._events = [(torch.cuda.Event(), torch.cuda.Event()) for _ in range(num_layers)] + # torch creates the CUDA event on its first record: do that now, never inside a capture + for pair in self._events: + for event in pair: + event.record(self.stream) + # on mode's select sees nothing resident: prefetch_ensure applies residency and the budget + self._none_resident = torch.full((num_experts,), -1, dtype=torch.int32, device=device) + self._budget = [0] * num_layers + width = 2 * self.k # room for the whole bs-2 union + self.sel = torch.full((num_layers, width), -1, dtype=torch.int32, device=device) + self.res = torch.full((num_layers, width), -1, dtype=torch.int32, device=device) + self.stats = torch.zeros((num_layers, NUM_STAT_COLS), dtype=torch.int64, device=device) + # measure mode's six columns, a view: its count kernel and readers keep the stage-A layout + self.counters = self.stats[:, :NUM_COLS] + self.totals = torch.zeros((num_layers, NUM_STAT_COLS), dtype=torch.int64) + # the router input of each forked prediction, alive until the compute stream has waited for + # the select so the allocator cannot hand its block to later work the predictor still reads + self._inflight: list[torch.Tensor | None] = [None] * num_layers + # on mode: whether install ran prefetch_ensure for the forked prediction (its copy is then due) + self._installed = [False] * num_layers + # on mode: the cache whose slot maps prefetch_ensure writes and whose banks the copy reads + self.cache = None + self.copy_stream = None + if mode == "on": + # the slim copy is an SM kernel: at default priority a late copy queues for SMs behind the next + # compute grid while the compute stream waits on it (bench --on: a 50 us slip cost 14% per step) + self.copy_stream = dedicated_stream(device, highest_priority=True) + # (dcopy: after prefetch_ensure(L) and layer L-1's demand copy on the compute stream, + # pfcopy: after layer L's prefetch copy) + self._copy_events = [(torch.cuda.Event(), torch.cuda.Event()) for _ in range(num_layers)] + for pair in self._copy_events: + for event in pair: + event.record(self.copy_stream) + # one plan per target layer: prefetch_ensure(L+1) must not overwrite the plan copy(L) reads + self.pf_slots = torch.full((num_layers, width), -1, dtype=torch.int32, device=device) + self.pf_src = torch.full((num_layers, width), -1, dtype=torch.int32, device=device) + self.pf_num = torch.zeros((num_layers, 1), dtype=torch.int64, device=device) + self.pf_ready = torch.zeros((num_layers, 1), dtype=torch.int32, device=device) + # test hooks, read at enqueue (capture) time: a [num_layers, width] device buffer replacing + # sel (on mode still installs at most the budget), and sleeps put on the predictor / copy + # stream ahead of their work + self.sel_override: torch.Tensor | None = None + self.delay_predict_cycles = 0 + self.delay_copy_cycles = 0 + logger.info( + f"MoE prefetch {mode}: router lookahead top-{self.k}, budget " + + (str(self.budget_override) if self.budget_override else f"{GDN_BUDGET} before GDN / {ATTN_BUDGET} before attention layers") + + (", slim copies on a dedicated high-priority stream" if mode == "on" else "") + ) + + def budget_for(self, wired: int) -> int: + return min(self.budget_override or wired, self.sel.shape[1]) + + def fork(self, target: int, x: torch.Tensor, gate, resident: torch.Tensor, budget: int) -> bool: + """Predict ``target``'s experts from ``x`` (the previous layer's router input) on the + predictor stream, ordered after the work already on the current stream. Measure mode skips + the ids ``resident`` holds (a racy read of the cache); on mode reads no cache state and + leaves residency and the budget to ``install``. Returns whether it forked.""" + if x.shape[0] > MAX_ROWS: + return False + budget = self.budget_for(budget) + if self.mode == "on": + self._budget[target] = budget + resident, budget = self._none_resident, self.sel.shape[1] + forked, selected = self._events[target] + forked.record(torch.cuda.current_stream(self.device)) + self.stream.wait_event(forked) + with torch.cuda.stream(self.stream): + if self.delay_predict_cycles: + torch.cuda._sleep(self.delay_predict_cycles) + logits = gate.forward(x) + lookahead_select(logits, resident, self.sel[target], self.res[target], k=self.k, budget=budget) + if self.sel_override is not None: + self.sel[target].copy_(self.sel_override[target]) + selected.record(self.stream) + self._inflight[target] = x + self._installed[target] = False + return True + + def install(self, target: int | None) -> int | None: + """On mode, right after the previous layer's ensure: wait for ``target``'s prediction and + install its first budget non-resident candidates with prefetch_ensure on the current + (compute) stream. Returns ``target`` when its copy is due (``issue_copy``), else None.""" + if target is None or self.mode != "on" or self._inflight[target] is None: + return None + torch.cuda.current_stream(self.device).wait_event(self._events[target][1]) + self.cache.prefetch_ensure( + target, self.sel[target], self.pf_slots[target], self.pf_src[target], + self.pf_num[target], self.stats[target], self.pf_ready[target], budget=self._budget[target], + ) + self._installed[target] = True + return target + + def _pending(self, layer_id: int) -> bool: + """On mode: ``layer_id``'s prediction was installed and its copy is not joined yet.""" + return self.mode == "on" and self._installed[layer_id] and self._inflight[layer_id] is not None + + def count_args(self, layer_id: int) -> tuple | None: + """On mode: what ``layer_id``'s demand ensure counts into, or None when it was not prefetched.""" + if not self._pending(layer_id): + return None + return self.stats[layer_id], self.pf_ready[layer_id], self.pf_num[layer_id], self._inflight[layer_id].shape[0] + + def plan_of(self, layer_id: int) -> tuple[torch.Tensor, torch.Tensor] | None: + """On mode: ``layer_id``'s prefetch copy plan (slots, count) while its prefetch is in flight.""" + if not self._pending(layer_id): + return None + return self.pf_slots[layer_id], self.pf_num[layer_id] + + def join_copy(self, layer_id: int) -> None: + """On mode: the compute stream waits for ``layer_id``'s prefetch copy before its own copy, so + a prefetched slot ensure(L) picked as a demand victim is never written by two copies.""" + if self.mode != "on" or self._inflight[layer_id] is None: + return + # a prediction that was never installed still has to rejoin (a capture needs every fork joined) + done = self._copy_events[layer_id][1] if self._pending(layer_id) else self._events[layer_id][1] + torch.cuda.current_stream(self.device).wait_event(done) + self._inflight[layer_id] = None + self._installed[layer_id] = False + + def issue_copy(self, target: int) -> None: + """On mode: copy ``target``'s prefetch plan on the copy stream after the work already on the + current stream: its prefetch_ensure and the demand copy (which owns the host link first).""" + if not self._pending(target): + return + dcopy, pfcopy = self._copy_events[target] + dcopy.record(torch.cuda.current_stream(self.device)) + self.copy_stream.wait_event(dcopy) + with torch.cuda.stream(self.copy_stream): + if self.delay_copy_cycles: + torch.cuda._sleep(self.delay_copy_cycles) + self.cache.copy_rows(target, self.pf_slots[target], self.pf_src[target], self.pf_num[target], slim=True, + kind=PREFETCH_COPY) + _mark_ready_kernel[(1,)](self.pf_ready[target]) + pfcopy.record(self.copy_stream) + + def join_and_count(self, layer_id: int, slots: torch.Tensor, id_of_slot: torch.Tensor, misses: torch.Tensor) -> None: + """Measure mode: join ``layer_id``'s prediction (if one was forked this step) and count it + against the routed slots ``slots`` of this layer's ensure.""" + x = self._inflight[layer_id] if self.mode == "measure" else None + if x is None: + return + torch.cuda.current_stream(self.device).wait_event(self._events[layer_id][1]) + prefetch_count( + slots.view(-1), id_of_slot, misses, self.sel[layer_id], self.res[layer_id], self.counters[layer_id], + id_base=layer_id * self.num_experts, rows=x.shape[0], + ) + self._inflight[layer_id] = None + + def reset_counters(self) -> None: + self.stats.zero_() + + def reset(self) -> None: + """Counters and in-flight state; the streams, events and buffers stay.""" + self.stats.zero_() + self.totals.zero_() + self._inflight = [None] * self.num_layers + self._installed = [False] * self.num_layers + + def take_window(self) -> torch.Tensor: + """This window's ``[num_layers, NUM_STAT_COLS]`` stats on the host (one sync), added to + ``totals``; the device ones restart.""" + window = self.stats.cpu() + self.stats.zero_() + self.totals += window + return window + + +def summarize(counters: torch.Tensor, mode: str = "measure") -> dict: + """Aggregate ``[num_layers, NUM_COLS or NUM_STAT_COLS]`` counters into per-layer-call and + per-token rates. In on mode the misses are what was left after the prefetch, so coverage is + useful / (useful + misses), and late is the share of copied calls whose copy had not landed + when the compute stream reached it.""" + total = counters.sum(0).tolist() + calls, issued, useful = total[CALLS], total[ISSUED], total[USEFUL] + tokens = int(counters[:, ROWS].max()) if counters.numel() else 0 + wanted = useful + total[MISSES] if mode == "on" else total[MISSES] + stats = { + "layer_calls": calls, + "tokens": tokens, + "issued_per_layer": issued / calls if calls else 0.0, + "useful_per_layer": useful / calls if calls else 0.0, + "resident_hits_per_layer": total[RESIDENT_HITS] / calls if calls else 0.0, + "misses_per_layer": total[MISSES] / calls if calls else 0.0, + "precision": useful / issued if issued else 0.0, + "coverage": useful / wanted if wanted else 0.0, + "useful_per_token": useful / tokens if tokens else 0.0, + } + if mode == "on": + copied = total[COPIED] if len(total) > COPIED else 0 + stats["copied_calls"] = copied + stats["late_fraction"] = total[LATE] / copied if copied else 0.0 + return stats + + +def format_summary(stats: dict) -> str: + if "late_fraction" in stats: + return ( + f"issued/layer={stats['issued_per_layer']:.2f}, useful/layer={stats['useful_per_layer']:.2f}, " + f"precision={stats['precision']:.3f}, useful/token={stats['useful_per_token']:.1f}, " + f"misses/layer={stats['misses_per_layer']:.2f}, coverage={stats['coverage']:.3f}, " + f"late={stats['late_fraction']:.3f} of {stats['copied_calls']} copies" + ) + return ( + f"issued/layer={stats['issued_per_layer']:.2f}, useful/layer={stats['useful_per_layer']:.2f}, " + f"precision={stats['precision']:.3f}, useful/token={stats['useful_per_token']:.1f}, " + f"misses/layer={stats['misses_per_layer']:.2f}, coverage={stats['coverage']:.3f}, " + f"resident_hits/layer={stats['resident_hits_per_layer']:.2f}" + ) + + +def format_per_layer(counters: torch.Tensor) -> str: + parts = [] + for layer, row in enumerate(counters.tolist()): + if row[CALLS]: + c = row[CALLS] + parts.append(f"L{layer}={row[USEFUL] / c:.2f}/{row[ISSUED] / c:.2f}/{row[MISSES] / c:.2f}") + return ", ".join(parts) diff --git a/python/freetoken/moe/scored_ensure.py b/python/freetoken/moe/scored_ensure.py new file mode 100644 index 0000000000..ba570b5169 --- /dev/null +++ b/python/freetoken/moe/scored_ensure.py @@ -0,0 +1,533 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright 2026 The FlashLib Authors +# Adapted from flashlib 0.3.0 (kernels/slot_cache/triton/lru_ensure.py: _phase1, _stats, +# _lru_ensure_kernel, _num_warps_for), https://github.com/FlashML-org/flashlib +"""Scored-victim slot-cache admission for the MoE expert cache (``--moe-cache-policy``). + +flashlib's sequential-argmin ``lru_ensure`` with the victim key made pluggable. POLICY 0 is +LRU and matches flashlib bit for bit; the others evict the lowest packed score key: + +- 1 ``kd``: ``-(k + beta * d / L)``. ``k`` is the tokens the expert has been passed over + since its last use, ``d`` the layer-steps until its layer comes up again. +- 2 ``kdfb``: kd plus ``w * log2(decayed use count)`` with a ``halflife``-token decay. +- 3 ``rule``: kdfb plus the near-miss refresh: this layer's experts whose router logit is + within a margin of the top-k-th count as used one token ago (``last_tok``), unpinned. + +The per-(layer, expert) state survives eviction; slot-indexed mirrors of it (tagged with the +owning id) let the victim scan read it coalesced instead of gathering by id. Scores and +counts are int32 Q16 fixed point and the count update ``log2(2^x + 1)`` reads a host-built +table, so a CPU reference can reproduce every victim exactly (``tl.exp2``/``tl.log2`` are +approximate). +""" +from __future__ import annotations + +import math + +import torch +import triton +import triton.language as tl + +POLICY_IDS = {"lru": 0, "kd": 1, "kdfb": 2, "rule": 3} + +# Replay-tuned on sep5_s01 at 1650 slots; the plateau is flat for beta 1-8, w 2-4, halflife 32-128. +BETA = 4.0 +W = 3.0 +HALFLIFE = 64 + +Q = 16 # fraction bits of scores and counts +LAST_TOK_NEVER = -1_000_000 +LC_NEVER = -(1 << 30) +LC_FLOOR = -(30 << Q) +K_MAX = 1 << 14 # tokens; with the caps on beta and w below, a score fits int32 +MAX_BETA, MAX_W = 64.0, 16.0 +DECAY_SHIFT = 8 +SCORE_BIAS = 1 << 31 +MAX_SLOT_BITS = 30 +# g(u) = log2(1 + 2^-u) sampled every 2^-G_STEP_BITS for u in [0, G_RANGE]; zero beyond +G_STEP_BITS = 6 +G_RANGE = 24 +G_LAST = G_RANGE << G_STEP_BITS +# jit code may only read globals that are constexpr +_Q = tl.constexpr(Q) +_LC_NEVER = tl.constexpr(LC_NEVER) +_LC_FLOOR = tl.constexpr(LC_FLOOR) +_K_MAX = tl.constexpr(K_MAX) +_DECAY_SHIFT = tl.constexpr(DECAY_SHIFT) +_SCORE_BIAS = tl.constexpr(SCORE_BIAS) +_G_SHIFT = tl.constexpr(Q - G_STEP_BITS) +_G_LAST = tl.constexpr(G_LAST) + + +def softplus2_table(device=None) -> torch.Tensor: + """Q16 ``log2(1 + 2^-u)`` at ``u = i / 64``, two entries past ``G_LAST`` for the interpolation.""" + step = 1 << G_STEP_BITS + vals = [round(math.log2(1.0 + 2.0 ** (-i / step)) * (1 << Q)) for i in range(G_LAST + 2)] + return torch.tensor(vals, dtype=torch.int32, device=device) + + +def score_params(beta: float, w: float, halflife: int, num_layers: int) -> tuple[int, int, int, int]: + """``(BETA_STEP, W_Q4, DECAY_MUL, DT_MAX)`` shared by the kernel and the CPU reference. + + ``beta * d / L`` is ``d * BETA_STEP`` in Q16; the decay over ``dt`` layer-steps is + ``(dt * DECAY_MUL) >> DECAY_SHIFT``, with ``dt`` capped at 64 half-lives (past the -30 floor). + """ + assert 0 <= beta <= MAX_BETA and 0 <= w <= MAX_W and halflife >= 1 + span = num_layers * int(halflife) + decay_mul = round((1 << (Q + DECAY_SHIFT)) / span) + dt_max = 64 * span + assert dt_max * decay_mul < (1 << 31) + return round(beta * (1 << Q) / num_layers), round(w * 16), decay_mul, dt_max + + +@triton.jit +def _phase1(query_ptr, slot_of_id_ptr, lru_usage_ptr, num_copy_ptr, step, K, + BLOCK_K: tl.constexpr, id_base): + """Vendored unchanged: dedup the query, split hit/miss, rank the misses, bump the hits.""" + k = tl.arange(0, BLOCK_K) + kmask = k < K + q = tl.load(query_ptr + k, mask=kmask, other=-1) + id_base + s = tl.load(slot_of_id_ptr + q, mask=kmask, other=-1) + hit = kmask & (s >= 0) + miss = kmask & (s == -1) + same = ( + (q[:, None] == q[None, :]) + & (k[:, None] > k[None, :]) + & kmask[:, None] + & kmask[None, :] + ) + first = kmask & (tl.sum(same.to(tl.int32), axis=1) == 0) + first_miss = miss & first + smaller = (q[None, :] < q[:, None]) & first_miss[None, :] + rank = tl.sum(smaller.to(tl.int32), axis=1) + num_missing = tl.sum(first_miss.to(tl.int32)) + tl.store(num_copy_ptr, num_missing.to(tl.int64)) + # Duplicated hits write the same value to the same slot -- idempotent. + tl.store(lru_usage_ptr + s, step, mask=hit) + return q, kmask, miss, first, first_miss, rank, num_missing, tl.where(hit, s, -1) + + +@triton.jit +def _prefetch_phase1(query_ptr, slot_of_id_ptr, K, BLOCK_K: tl.constexpr, id_base): + """_phase1 for a prefetch: -1 entries are padding, hits are left alone (a prefetch is not an + access), and the misses rank in query order so the select's best candidates install first.""" + k = tl.arange(0, BLOCK_K) + raw = tl.load(query_ptr + k, mask=k < K, other=-1) + kmask = raw >= 0 + q = tl.where(kmask, raw + id_base, 0) + s = tl.load(slot_of_id_ptr + q, mask=kmask, other=0) + miss = kmask & (s == -1) + same = (q[:, None] == q[None, :]) & (k[:, None] > k[None, :]) & kmask[:, None] & kmask[None, :] + first_miss = miss & (tl.sum(same.to(tl.int32), axis=1) == 0) + rank = tl.cumsum(first_miss.to(tl.int32), axis=0) - 1 + return q, first_miss, rank, tl.sum(first_miss.to(tl.int32)) + + +@triton.jit +def _stats(stats_ptr, first, num_missing): + # One vectorized atomic over 3 lanes, not three scalar ones: a scalar atomic serializes the CTA. + si = tl.arange(0, 4) + v = tl.where(si == 0, tl.sum(first.to(tl.int32)), tl.where(si == 1, num_missing, 1)) + # relaxed: only later kernels read the counters; flashlib's acq_rel waits on this call's stores (~0.3 us) + tl.atomic_add(stats_ptr + si, v.to(tl.int64), mask=si < 3, sem="relaxed") + + +@triton.jit +def _decayed(lc, ct, now, DECAY_MUL: tl.constexpr, DT_MAX: tl.constexpr): + """int32 Q16 log2 count decayed from ``ct`` to ``now`` (int64 layer-steps).""" + dt = tl.minimum(tl.maximum(now - ct, 0), DT_MAX).to(tl.int32) + return lc - ((dt * DECAY_MUL) >> _DECAY_SHIFT) + + +@triton.jit +def _softplus2(x, g_ptr, mask): + """Q16 ``log2(2^x + 1)`` = ``max(x, 0) + g(|x|)``, g linearly interpolated from the table.""" + u = tl.abs(x) + i = tl.minimum(u >> _G_SHIFT, _G_LAST) + frac = u & ((1 << _G_SHIFT) - 1) + g0 = tl.load(g_ptr + i, mask=mask, other=0) + g1 = tl.load(g_ptr + i + 1, mask=mask, other=0) + g = tl.where(i < _G_LAST, g0 + (((g1 - g0) * frac) >> _G_SHIFT), 0) + return tl.maximum(x, 0) + g + + +@triton.jit +def _owner_state(mirror_ptr, state_ptr, c, oid, mirrored, stale, other): + """The slot owner's state: its coalesced mirror, or a gather by id where the mirror is stale.""" + return tl.where( + mirrored, + tl.load(mirror_ptr + c, mask=mirrored, other=other), + tl.load(state_ptr + oid, mask=stale, other=other), + ) + + +@triton.jit(do_not_specialize=["K", "num_cached", "id_base", "nm_topk", "nm_stride", "pf_rows", "pf_budget"]) +def _scored_ensure_kernel( + query_ptr, slot_of_id_ptr, id_of_slot_ptr, lru_usage_ptr, lru_step_ptr, + out_ptr, src_ptr, dst_ptr, num_copy_ptr, stats_ptr, + tok_ptr, last_tok_ptr, lc_ptr, ct_ptr, g_ptr, logits_ptr, + owner_ptr, slot_last_ptr, slot_lc_ptr, slot_ct_ptr, pin_ptr, + pf_stats_ptr, pf_ready_ptr, pf_num_ptr, + K, num_cached, id_base, nm_topk, nm_stride, nm_thr, pf_rows, pf_budget, + BLOCK_K: tl.constexpr, BLOCK_C: tl.constexpr, BLOCK_E: tl.constexpr, BLOCK_TOPK: tl.constexpr, + NM_ROWS: tl.constexpr, + USAGE_MAX: tl.constexpr, COLLECT_STATS: tl.constexpr, + POLICY: tl.constexpr, BUMP_TOK: tl.constexpr, UPDATE_STATE: tl.constexpr, PIN_SINCE: tl.constexpr, + NUM_LAYERS: tl.constexpr, NUM_EXPERTS: tl.constexpr, + BETA_STEP: tl.constexpr, W_Q4: tl.constexpr, DECAY_MUL: tl.constexpr, DT_MAX: tl.constexpr, + SLOT_BITS: tl.constexpr, + LOWPRI: tl.constexpr, PF_COUNT: tl.constexpr, PREFETCH: tl.constexpr, +): + if PREFETCH: + # not an access: the clock stays, so usage == step still pins the last demand call's slots + # (at step 0, right after a reset, every slot reads pinned: nothing installs before a demand call) + step = tl.load(lru_step_ptr) + tl.store(num_copy_ptr, tl.zeros([], tl.int64)) + tl.store(pf_ready_ptr, 0) + else: + step = tl.load(lru_step_ptr) + 1 + tl.store(lru_step_ptr, step) + if POLICY != 0: + tok = tl.load(tok_ptr) + if BUMP_TOK: + tok = tok + 1 + tl.store(tok_ptr, tok) + layer = id_base // NUM_EXPERTS + if PREFETCH: + # score from where the pass is, right after layer - 1's demand call: the target layer's + # residents are next up, not a whole pass away (else it evicts what the layer is about to route) + layer = layer - 1 + now = tok * NUM_LAYERS + layer + if POLICY == 3: + if UPDATE_STATE: + # Routed ids qualify too; they are pinned now and set to tok below, so that is harmless. + ex = tl.arange(0, BLOCK_E) + j = tl.arange(0, BLOCK_TOPK) + prev = tl.load(last_tok_ptr + id_base + ex, mask=ex < NUM_EXPERTS, other=0) + held_at = tl.load(slot_of_id_ptr + id_base + ex, mask=ex < NUM_EXPERTS, other=-1) + near = ex < 0 + for r in tl.static_range(NM_ROWS): + routed = tl.load(query_ptr + r * nm_topk + j, mask=j < nm_topk, other=-1) + row_ptr = logits_ptr + r * nm_stride + kth = tl.min(tl.load(row_ptr + routed, mask=routed >= 0, other=float("inf")).to(tl.float32), axis=0) + row = tl.load(row_ptr + ex, mask=ex < NUM_EXPERTS, other=float("-inf")).to(tl.float32) + near = near | (row >= kth - nm_thr) + refreshed = tl.maximum(prev, tok - 1) + tl.store(last_tok_ptr + id_base + ex, refreshed, mask=near) + tl.store(slot_last_ptr + held_at, refreshed, mask=near & (held_at >= 0)) + # the victim scan and the routed update below read and rewrite what this wrote + tl.debug_barrier() + if PREFETCH: + q, first_miss, rank, num_missing = _prefetch_phase1(query_ptr, slot_of_id_ptr, K, BLOCK_K, id_base) + else: + if PF_COUNT: + kk = tl.arange(0, BLOCK_K) + sk = tl.load(slot_of_id_ptr + tl.load(query_ptr + kk, mask=kk < K, other=0) + id_base, mask=kk < K, other=-1) + # a hit on a held slot with usage 0 consumes a prefetch; read it before _phase1 bumps it + was_pf = (sk >= 0) & (tl.load(lru_usage_ptr + sk, mask=sk >= 0, other=1) == 0) + tl.debug_barrier() + q, kmask, miss, first, first_miss, rank, num_missing, out = _phase1( + query_ptr, slot_of_id_ptr, lru_usage_ptr, num_copy_ptr, step, K, BLOCK_K, + id_base) + + if num_missing > 0: + # REQUIRED (flashlib): the reload below must see _phase1's hit bump, or a hit is evicted. + tl.debug_barrier() + c = tl.arange(0, BLOCK_C) + cmask = c < num_cached + u = tl.load(lru_usage_ptr + c, mask=cmask, other=USAGE_MAX) + if POLICY == 0: + umax = tl.full([BLOCK_C], USAGE_MAX, u.dtype) + if LOWPRI: + held0 = cmask & (tl.load(id_of_slot_ptr + c, mask=cmask, other=-1) >= 0) + # empty slots, then unconsumed prefetches (held, usage 0), then residents by recency + u = tl.where((u == step) | (~cmask), umax, tl.where(held0, u + 1, 0)) + else: + u = tl.where((u == step) | (~cmask), umax, u) + n_free = tl.sum((u != umax).to(tl.int32)) + else: + oid = tl.load(id_of_slot_ptr + c, mask=cmask, other=-1) + held = cmask & (oid >= 0) + owner = tl.load(owner_ptr + c, mask=cmask, other=-1) + mirrored = held & (owner == oid) + stale = held & (owner != oid) + lk = oid // NUM_EXPERTS + ahead = lk > layer + d = tl.where(ahead, lk - layer, NUM_LAYERS - layer + lk) + last = _owner_state(slot_last_ptr, last_tok_ptr, c, oid, mirrored, stale, 0) + k = tl.minimum(tl.maximum(tl.where(ahead, tok - 1, tok) - last, 0), _K_MAX).to(tl.int32) + # int32 throughout: |score| < 2^31 by the caps on k, beta and w + score = -((k << _Q) + d * BETA_STEP) + if POLICY >= 2: + lc = _owner_state(slot_lc_ptr, lc_ptr, c, oid, mirrored, stale, _LC_NEVER) + ct = _owner_state(slot_ct_ptr, ct_ptr, c, oid, mirrored, stale, 0) + lcv = _decayed(lc, ct, now, DECAY_MUL, DT_MAX) + lcv = tl.where(lc == _LC_NEVER, _LC_FLOOR, tl.maximum(lcv, _LC_FLOOR)) + score += (W_Q4 * lcv) >> 4 + # Every key is distinct and ties go to the lowest slot; an empty slot keys below any held one. + key = tl.where(held, ((score.to(tl.int64) + _SCORE_BIAS) << SLOT_BITS) | c, c.to(tl.int64)) + if LOWPRI: + # an unconsumed prefetch (held, usage 0) keys above every empty slot and below every + # resident: the caps keep score + 2^31 > 2^30 - 2^25 - 2^22, far above this band's 1 + key = tl.where(held & (u == 0), (1 << SLOT_BITS) | c.to(tl.int64), key) + evictable = cmask & (u != step) + if PIN_SINCE: + # the score ignores recency, so slots an earlier call touched since *pin_ptr stay pinned too + evictable = evictable & (u <= tl.load(pin_ptr)) + # the pinned key's slot bits are 0, so an all-pinned scan falls back to slot 0 like flashlib's argmin + pinned_key = 0x7FFFFFFFFFFFFFFF ^ ((1 << SLOT_BITS) - 1) + key = tl.where(evictable, key, pinned_key) + n_free = tl.sum(evictable.to(tl.int32)) + if PREFETCH: + # never the all-pinned fallback: that slot may be under the running GEMM; misses rank in + # query order, so the budget keeps the first non-resident candidates + n_iter = tl.minimum(tl.minimum(num_missing, n_free), pf_budget) + tl.store(num_copy_ptr, n_iter.to(tl.int64)) + tl.store(pf_stats_ptr, tl.load(pf_stats_ptr) + n_iter) + else: + n_iter = num_missing + for i in tl.range(n_iter): + if POLICY == 0: + victim = tl.argmin(u, axis=0).to(tl.int32) + else: + victim = (tl.min(key, axis=0) & ((1 << SLOT_BITS) - 1)).to(tl.int32) + # Scalar load: victims are distinct, so no earlier iteration wrote this slot. + old = tl.load(id_of_slot_ptr + victim) + if old >= 0: + tl.store(slot_of_id_ptr + old, -1) + e = tl.sum(tl.where((rank == i) & first_miss, q, 0)) + tl.store(id_of_slot_ptr + victim, e) + tl.store(slot_of_id_ptr + e, victim) + if PREFETCH: + # usage 0 on a held slot marks it low priority until a demand call touches it + tl.store(lru_usage_ptr + victim, 0) + if POLICY != 0: + tl.store(owner_ptr + victim, -1) + else: + tl.store(lru_usage_ptr + victim, step) + tl.store(dst_ptr + i, victim) + tl.store(src_ptr + i, e - id_base) # back to the caller's id space + if not PREFETCH: + out = tl.where((rank == i) & miss, victim, out) + if POLICY == 0: + u = tl.where(c == victim, umax, u) # claim in-register + else: + key = tl.where(c == victim, pinned_key, key) + + if not PREFETCH: + # Written from registers, never re-read from slot_of_id, so out_ptr may alias query_ptr. + tl.store(out_ptr + tl.arange(0, BLOCK_K), out, mask=kmask) + if PF_COUNT: + # one writer per column: pf_ensure adds ISSUED (column 0), earlier on this stream + pn = tl.load(pf_num_ptr) + late = (pn > 0) & (tl.load(pf_ready_ptr, volatile=True) == 0) + col = tl.arange(0, 8) + add = tl.where(col == 1, tl.sum((first & was_pf).to(tl.int32)).to(tl.int64), 0) + add = tl.where(col == 3, 1, add) + add = tl.where(col == 4, pf_rows, add) + add = tl.where(col == 5, num_missing.to(tl.int64), add) + add = tl.where(col == 6, late.to(tl.int64), add) + add = tl.where(col == 7, (pn > 0).to(tl.int64), add) + cols = (col == 1) | (col >= 3) + tl.store(pf_stats_ptr + col, tl.load(pf_stats_ptr + col, mask=cols, other=0) + add, mask=cols) + if POLICY != 0: + if UPDATE_STATE: + # distinct ids only, so a duplicate at bs > 1 counts once per call; out is each id's final slot + tl.store(last_tok_ptr + q, tok, mask=first) + tl.store(owner_ptr + out, q, mask=first) + tl.store(slot_last_ptr + out, tok, mask=first) + if POLICY >= 2: + lc = tl.load(lc_ptr + q, mask=first, other=_LC_NEVER) + x = _decayed(lc, tl.load(ct_ptr + q, mask=first, other=0), now, DECAY_MUL, DT_MAX) + lc_new = tl.where(lc == _LC_NEVER, 0, _softplus2(x, g_ptr, first)) + tl.store(lc_ptr + q, lc_new, mask=first) + tl.store(ct_ptr + q, now, mask=first) + tl.store(slot_lc_ptr + out, lc_new, mask=first) + tl.store(slot_ct_ptr + out, now, mask=first) + elif not PREFETCH: + # a slot filled without a state update reads the per-id state until a decode hit re-mirrors it + tl.store(owner_ptr + out, -1, mask=first_miss) + if COLLECT_STATS: + _stats(stats_ptr, first, num_missing) + + +def _num_warps_for(block_c: int) -> int: + # flashlib's rule: 8 warps from 2048 slots, more only past ~32 elements per thread. + warps = 8 if block_c >= 2048 else 4 + return max(warps, min(triton.next_power_of_2(max(block_c // 1024, 1)), 32)) + + +def scored_ensure( + query: torch.Tensor, + slot_of_id: torch.Tensor, + id_of_slot: torch.Tensor, + lru_usage: torch.Tensor, + lru_step: torch.Tensor, + out_indices: torch.Tensor, + src_indices: torch.Tensor, + dst_indices: torch.Tensor, + num_copy: torch.Tensor, + *, + policy: int, + num_layers: int, + num_experts: int, + tok: torch.Tensor | None = None, + last_tok: torch.Tensor | None = None, + lc: torch.Tensor | None = None, + ct: torch.Tensor | None = None, + g_table: torch.Tensor | None = None, + slot_owner: torch.Tensor | None = None, + slot_last_tok: torch.Tensor | None = None, + slot_lc: torch.Tensor | None = None, + slot_ct: torch.Tensor | None = None, + router_logits: torch.Tensor | None = None, + near_miss_thr: float = 0.0, + stats: torch.Tensor | None = None, + id_base: int = 0, + bump_tok: bool = False, + update_state: bool = True, + pin_since: torch.Tensor | None = None, + beta: float = BETA, + w: float = W, + halflife: int = HALFLIFE, + lowpri: bool = False, + pf_stats: torch.Tensor | None = None, + pf_ready: torch.Tensor | None = None, + pf_num: torch.Tensor | None = None, + pf_rows: int = 0, + pf_budget: int | None = None, + prefetch: bool = False, +) -> None: + """``flashlib.lru_ensure`` (sequential strategy) with the victim picked by ``policy``. + + Same contract as ``lru_ensure``; ``policy`` > 0 also reads and updates the per-(layer, + expert) state ``last_tok``/``lc``/``ct`` (flat ``layer * num_experts + expert``) against + the device token counter ``tok``, and keeps its ``[num_cached]`` mirrors ``slot_last_tok``/ + ``slot_lc``/``slot_ct`` valid where ``slot_owner`` equals ``id_of_slot`` (-1 marks a stale + mirror). ``bump_tok`` advances the counter first (once per decode step); + ``update_state=False`` leaves the state untouched (small-prefill ensures). + ``pin_since`` (a ``lru_step`` value) also pins every slot touched after it, so a small + prefill's chunked calls cannot evict each other's experts; policy > 0 only. + Policy 3 refreshes the experts whose ``router_logits`` (``[rows, num_experts]``, rows of + ``query``) are within ``near_miss_thr`` of their row's lowest routed logit. + + FREETOKEN_MOE_PREFETCH=on: ``lowpri`` keys held slots with usage 0 (unconsumed prefetches) + between the empty slots and every resident. ``pf_stats`` (a ``[8]`` int64 prefetch stats row) + also counts this demand call: hits on such slots, the call, ``pf_rows``, the misses, and whether + the prefetch plan ``pf_num`` had rows whose copy had not set ``pf_ready`` yet. + ``prefetch=True`` is :func:`prefetch_ensure`'s install mode instead, installing at most + ``pf_budget`` ids (default: all). + """ + if prefetch: + assert lowpri and pf_stats is not None and pf_ready is not None and stats is None + assert not update_state and not bump_tok and pin_since is None and router_logits is None + if pf_stats is not None: + assert lowpri and pf_stats.dtype == torch.int64 and pf_stats.numel() == 8 and pf_stats.is_contiguous() + assert pf_ready is not None and pf_ready.dtype == torch.int32 and pf_ready.numel() == 1 + assert prefetch or (pf_num is not None and pf_num.dtype == torch.int64 and pf_num.numel() == 1) + k = query.numel() + num_cached = id_of_slot.numel() + assert query.dtype == torch.int32 and query.is_contiguous() + assert slot_of_id.dtype == torch.int32 and id_of_slot.dtype == torch.int32 + assert out_indices.dtype == torch.int32 and out_indices.numel() == k + assert lru_usage.dtype == lru_step.dtype and lru_usage.numel() == num_cached + plan = min(k, num_cached) + assert src_indices.numel() >= plan and dst_indices.numel() >= plan + block_c = triton.next_power_of_2(num_cached) + slot_bits = max(block_c.bit_length() - 1, 1) + beta_step, w_q4, decay_mul, dt_max = score_params(beta, w, halflife, num_layers) + if policy != 0: + assert slot_bits <= MAX_SLOT_BITS, f"{num_cached} slots overflow the packed score key" + assert tok is not None and last_tok is not None + assert last_tok.dtype == torch.int64 and last_tok.numel() == num_layers * num_experts + assert slot_owner is not None and slot_last_tok is not None + assert slot_owner.dtype == torch.int32 and slot_owner.numel() == num_cached + assert slot_last_tok.dtype == torch.int64 and slot_last_tok.numel() == num_cached + if pin_since is not None: + assert policy != 0 and pin_since.dtype == lru_step.dtype and pin_since.numel() == 1 + if policy >= 2: + assert lc is not None and ct is not None and g_table is not None + assert lc.dtype == torch.int32 and ct.dtype == torch.int64 + assert slot_lc is not None and slot_lc.dtype == torch.int32 and slot_lc.numel() == num_cached + assert slot_ct is not None and slot_ct.dtype == torch.int64 and slot_ct.numel() == num_cached + nm_rows = nm_topk = nm_stride = 0 + if policy == 3 and router_logits is not None: + assert router_logits.ndim == 2 and router_logits.shape[1] == num_experts + assert router_logits.stride(1) == 1 and k % router_logits.shape[0] == 0 + nm_rows, nm_stride = router_logits.shape[0], router_logits.stride(0) + nm_topk = k // nm_rows + # unused pointers still need a tensor argument + dummy = lru_step + _scored_ensure_kernel[(1,)]( + query, slot_of_id, id_of_slot, lru_usage, lru_step, + out_indices, src_indices, dst_indices, num_copy, stats, + dummy if tok is None else tok, + dummy if last_tok is None else last_tok, + dummy if lc is None else lc, + dummy if ct is None else ct, + dummy if g_table is None else g_table, + dummy if nm_rows == 0 else router_logits, + dummy if slot_owner is None else slot_owner, + dummy if slot_last_tok is None else slot_last_tok, + dummy if slot_lc is None else slot_lc, + dummy if slot_ct is None else slot_ct, + dummy if pin_since is None else pin_since, + dummy if pf_stats is None else pf_stats, + dummy if pf_ready is None else pf_ready, + dummy if pf_num is None else pf_num, + k, num_cached, id_base, nm_topk, nm_stride, float(near_miss_thr), int(pf_rows), + k if pf_budget is None else int(pf_budget), + BLOCK_K=triton.next_power_of_2(k), + BLOCK_C=block_c, + BLOCK_E=triton.next_power_of_2(num_experts), + BLOCK_TOPK=triton.next_power_of_2(max(nm_topk, 1)), + NM_ROWS=nm_rows, + USAGE_MAX=torch.iinfo(lru_usage.dtype).max, + COLLECT_STATS=stats is not None, + POLICY=policy, + BUMP_TOK=bool(bump_tok) and policy != 0, + UPDATE_STATE=bool(update_state) and policy != 0, + PIN_SINCE=pin_since is not None, + NUM_LAYERS=num_layers, + NUM_EXPERTS=num_experts, + BETA_STEP=beta_step, + W_Q4=w_q4, + DECAY_MUL=decay_mul, + DT_MAX=dt_max, + SLOT_BITS=slot_bits, + LOWPRI=bool(lowpri), + PF_COUNT=pf_stats is not None and not prefetch, + PREFETCH=bool(prefetch), + num_warps=_num_warps_for(block_c), + ) + + +def prefetch_ensure( + query: torch.Tensor, + slot_of_id: torch.Tensor, + id_of_slot: torch.Tensor, + lru_usage: torch.Tensor, + lru_step: torch.Tensor, + dst_slots: torch.Tensor, + src_rows: torch.Tensor, + num_copy: torch.Tensor, + pf_stats: torch.Tensor, + pf_ready: torch.Tensor, + budget: int | None = None, + **state, +) -> None: + """Install the ids of ``query`` (-1 = padding) that are not resident as low-priority slots + (held, usage 0), at most the first ``budget`` of them in query order (default: all), and + write their copy plan to ``dst_slots``/``src_rows``/``num_copy``. + + Unlike a demand call it bumps no clock, touches no hit and no per-id state, and so keeps + every slot of the last demand call (usage == step) pinned; victims follow the demand key + (empty, then earlier unconsumed prefetches, then the coldest residents), scored as of the + previous layer's call, the position the prefetch runs at. With fewer + evictable slots than misses it installs only the first ones, in query order. Adds the + installs to ``pf_stats[0]`` and clears ``pf_ready``. + """ + scored_ensure( + query, slot_of_id, id_of_slot, lru_usage, lru_step, query, src_rows, dst_slots, num_copy, + update_state=False, lowpri=True, pf_stats=pf_stats, pf_ready=pf_ready, pf_budget=budget, prefetch=True, + **state, + ) diff --git a/python/freetoken/moe/slot_audit.py b/python/freetoken/moe/slot_audit.py new file mode 100644 index 0000000000..12d9f40ad9 --- /dev/null +++ b/python/freetoken/moe/slot_audit.py @@ -0,0 +1,688 @@ +"""Expert slot audit (``FREETOKEN_MOE_SLOT_AUDIT=N``): a debug detector for PERSISTENT slot corruption. + +Two parts, both on the device and off by default: + +- writer history: every writer of the slot maps (demand and prefetch ensures, prefill ensures, the + whole-layer materialize, prefill buffer invalidation, reset) and every writer of the slot banks + (demand, prefetch and prefill copies, the prefill double buffers) is followed on its writer's + stream by a tiny kernel keyed on the same plan (or slot range) that appends ``(seq, kind, decode + step, lru step, layer, expert row, ...)`` to a per-slot ring of the last ``DEPTH`` metadata events + and one of the last ``DEPTH`` byte events. The map writers all run on the compute stream except the + prefill buffer invalidation (on the prefill copy stream, fenced behind it); the byte writers also + run on the prefetch and prefill copy streams. ``seq`` is one global atomic counter, so the two rings + merge into one cross-stream order. Clears are recorded just before their writer, so they carry the + old owner. + Fixed shapes and no host sync: the recorders capture into the decode graph. +- audit: every N decode steps, after the forward on the compute stream, every held slot is checked: + its bytes in every bank against the host row of the (layer, expert) its ``id_of_slot`` names + (host rows gathered with the fused copy into a scratch of fewer than 64 MiB, one layer chunk at a + time), ``slot_for_id[id_of_slot[s]] == s``, and, over every id, ``id_of_slot[slot_for_id[id]] == id`` + (a stale entry routes that id to another expert's bytes). The first ``RING`` bad (slot, owner) + pairs are kept with both history rings; the engine logs them and the totals every + MOE_STATS_INTERVAL decode steps and at shutdown, with a diagnosis read from the history. +""" +from __future__ import annotations + +import math + +import torch +import triton +import triton.language as tl + +from freetoken.env import ENV +from freetoken.utils import init_logger + +logger = init_logger(__name__) + +# events kept per slot and ring (metadata / bytes) +DEPTH = 4 +# event fields +(E_SEQ, E_KIND, E_DSTEP, E_LSTEP, E_LAYER, E_ROW, E_AUX, E_OWNER, E_USAGE, E_POS) = range(10) +EVENT_FIELDS = 10 +# metadata kinds (installs name the new owner; clears the old one, recorded before the writer) +DEMAND_INSTALL, PREFETCH_INSTALL, PREFILL_INSTALL, MATERIALIZE_INSTALL = 1, 2, 3, 4 +MATERIALIZE_CLEAR, INVALIDATE, RESET = 5, 6, 7 +# byte kinds +DEMAND_COPY, PREFETCH_COPY, PREFILL_COPY, MATERIALIZE_COPY, PAGEABLE_COPY = 11, 12, 13, 14, 15 +PREFILL_BUFFER, PREFILL_SPLIT_H2D, PREFILL_HIT_D2D, OTHER_COPY = 16, 17, 18, 19 +KIND_NAMES = { + DEMAND_INSTALL: "demand install", PREFETCH_INSTALL: "prefetch install", PREFILL_INSTALL: "prefill ensure install", + MATERIALIZE_INSTALL: "materialize install", MATERIALIZE_CLEAR: "materialize clear", + INVALIDATE: "prefill buffer invalidate", RESET: "reset", + DEMAND_COPY: "demand copy", PREFETCH_COPY: "prefetch copy", PREFILL_COPY: "prefill ensure copy", + MATERIALIZE_COPY: "materialize copy", PAGEABLE_COPY: "pageable materialize copy", + PREFILL_BUFFER: "prefill buffer copy", PREFILL_SPLIT_H2D: "prefill buffer H2D (misses and small banks)", + PREFILL_HIT_D2D: "prefill buffer D2D", OTHER_COPY: "copy", +} +INSTALLS = (DEMAND_INSTALL, PREFETCH_INSTALL, PREFILL_INSTALL, MATERIALIZE_INSTALL) +# the copy each install expects next +COPY_FOR = { + DEMAND_INSTALL: DEMAND_COPY, PREFETCH_INSTALL: PREFETCH_COPY, PREFILL_INSTALL: PREFILL_COPY, + MATERIALIZE_INSTALL: MATERIALIZE_COPY, +} +# recorder row sources: the plan's src rows, slot - row_base, the range index, the old owner +ROW_PLAN, ROW_FROM_SLOT, ROW_INDEX, ROW_OLD_OWNER = range(4) +# recorder predicates: every slot, held slots, slots >= E held by the layer +PRED_ALL, PRED_HELD, PRED_LAYER_ABOVE_E = range(3) +# totals: audits, held slots seen, slots byte-checked, held slots not byte-checked, bad slots found +# (summed over audits), of them with bad bytes / a forward map miss / named by a stale slot_for_id / +# an owner out of range, slot_for_id entries past the cache, new records, repeat findings, events +# recorded, plan entries naming no slot +(T_AUDITS, T_HELD, T_BYTES_CHECKED, T_UNCHECKED, T_BAD, T_BYTES, T_FWD, T_STALE, T_OWNER_RANGE, T_MAP_RANGE, + T_NEW, T_STILL, T_EVENTS, T_BAD_PLAN) = range(14) +NUM_TOTALS = 16 +# per-slot audit flags +F_BYTES, F_FWD, F_STALE, F_OWNER_RANGE = 1, 2, 4, 8 +# record header, then the metadata ring and the byte ring as stored ([DEPTH, EVENT_FIELDS] each) +(H_AUDIT, H_DSTEP, H_LSTEP, H_SLOT, H_OWNER, H_MAPPED, H_USAGE, H_FLAGS, H_BANKS, H_FIRST_BANK, H_FIRST_BYTE, + H_BAD_WORDS, H_STALE_ID, H_STALE_N, H_META_N, H_BYTES_N) = range(16) +HEADER = 16 +RING_EVENTS = DEPTH * EVENT_FIELDS +RECORD_FIELDS = HEADER + 2 * RING_EVENTS +RING = 32 +SCRATCH_BYTES = 64 << 20 +MAX_CHUNK_ROWS = 64 +COMPARE_WORDS = 4096 +_NO_BAD = tl.constexpr(0x7FFFFFFF) +_D = tl.constexpr(DEPTH) +_F = tl.constexpr(EVENT_FIELDS) + + +def resolve_interval(interval: int | None = None) -> int: + value = ENV.MOE_SLOT_AUDIT.value if interval is None else interval + value = int(value or 0) + if value < 0: + raise ValueError(f"FREETOKEN_MOE_SLOT_AUDIT={value}: expected 0 (off) or the decode steps between audits") + return value + + +@triton.jit +def _put(vals, f, idx: tl.constexpr, v): + return tl.where(f[None, :] == idx, v.to(tl.int64)[:, None], vals) + + +@triton.jit +def _bump_step_kernel(dstep_ptr): + tl.store(dstep_ptr, tl.load(dstep_ptr) + 1) + + +@triton.jit(do_not_specialize=["kind", "layer", "start", "row_base", "width", "num_slots"]) +def _record_kernel( + hist_ptr, cur_ptr, seq_ptr, dstep_ptr, lstep_ptr, id_of_slot_ptr, usage_ptr, totals_ptr, + dst_ptr, src_ptr, aux_ptr, num_ptr, + kind, layer, start, row_base, width, num_slots, + PLAN: tl.constexpr, ROW: tl.constexpr, HAS_AUX: tl.constexpr, HAS_NUM: tl.constexpr, PRED: tl.constexpr, + E: tl.constexpr, BLOCK: tl.constexpr, +): + """Append one event per written slot to that slot's ring: plan entries ``i < min(*num, width)`` + (``PLAN``) or the slots ``start + i`` for ``i < width``.""" + i = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + n = width + if HAS_NUM: + n = tl.minimum(tl.load(num_ptr).to(tl.int32), width) + valid = i < n + if PLAN: + slot = tl.load(dst_ptr + i, mask=valid, other=-1) + else: + slot = start + i + ok = valid & (slot >= 0) & (slot < num_slots) + bad_plan = valid & ~ok + owner = tl.load(id_of_slot_ptr + slot, mask=ok, other=-3) + usage = tl.load(usage_ptr + slot, mask=ok, other=-3) + if PRED == 1: + ok = ok & (owner >= 0) + if PRED == 2: + ok = ok & (owner >= layer * E) & (owner < layer * E + E) & (slot >= E) + lay = tl.zeros([BLOCK], tl.int32) + layer + aux = tl.zeros([BLOCK], tl.int32) - 1 + if ROW == 0: + row = tl.load(src_ptr + i, mask=ok, other=-1) + elif ROW == 1: + row = slot - row_base + elif ROW == 2: + row = i + else: + row = tl.where(owner >= 0, owner % E, -1) + lay = tl.where(owner >= 0, owner // E, -1) + aux = aux + 1 + layer + if HAS_AUX: + aux = tl.load(aux_ptr + i, mask=ok, other=-1) + cnt = tl.sum(ok.to(tl.int64), axis=0) + seq = tl.atomic_add(seq_ptr, cnt) + tl.cumsum(ok.to(tl.int64), axis=0) - 1 + pos = tl.atomic_add(cur_ptr + slot, 1, mask=ok) + f = tl.arange(0, 16) + z = tl.zeros([BLOCK], tl.int64) + vals = tl.zeros([BLOCK, 16], tl.int64) - 1 + vals = _put(vals, f, 0, seq) + vals = _put(vals, f, 1, z + kind) + vals = _put(vals, f, 2, z + tl.load(dstep_ptr)) + vals = _put(vals, f, 3, z + tl.load(lstep_ptr)) + vals = _put(vals, f, 4, lay) + vals = _put(vals, f, 5, row) + vals = _put(vals, f, 6, aux) + vals = _put(vals, f, 7, owner) + vals = _put(vals, f, 8, usage) + vals = _put(vals, f, 9, i) + at = (slot.to(tl.int64) * _D + (pos % _D)) * _F + tl.store(hist_ptr + at[:, None] + f[None, :], vals, mask=ok[:, None] & (f[None, :] < _F)) + t = tl.arange(0, 2) + add = tl.where(t == 0, cnt, tl.sum(bad_plan.to(tl.int64), axis=0)) + tl.atomic_add(totals_ptr + 12 + t, add, sem="relaxed") + + +@triton.jit(do_not_specialize=["num_slots", "num_ids"]) +def _audit_meta_kernel( + id_of_slot_ptr, slot_for_id_ptr, layer_ok_ptr, count_ptr, rows_ptr, slots_ptr, + flags_ptr, stale_id_ptr, stale_n_ptr, totals_ptr, num_slots, num_ids, + E: tl.constexpr, CAP: tl.constexpr, BLOCK: tl.constexpr, +): + """Both map checks, and each byte-checkable held slot appended to its layer's list.""" + i = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + ms = i < num_slots + o = tl.load(id_of_slot_ptr + i, mask=ms, other=-1) + held = ms & (o >= 0) + in_range = held & (o < num_ids) + mapped = tl.load(slot_for_id_ptr + o, mask=in_range, other=-2) + fwd = in_range & (mapped != i) + out_of_range = held & ~in_range + fl = tl.where(fwd, 2, 0) | tl.where(out_of_range, 8, 0) + tl.atomic_or(flags_ptr + i, fl, mask=ms & (fl != 0)) + layer = tl.where(in_range, o // E, 0) + ok = in_range & (tl.load(layer_ok_ptr + layer, mask=in_range, other=0) != 0) + p = tl.atomic_add(count_ptr + layer, 1, mask=ok) + take = ok & (p < CAP) + tl.store(rows_ptr + layer * CAP + p, o - layer * E, mask=take) + tl.store(slots_ptr + layer * CAP + p, i, mask=take) + mi = i < num_ids + s = tl.load(slot_for_id_ptr + i, mask=mi, other=-1) + live = mi & (s >= 0) + s_ok = live & (s < num_slots) + owner = tl.load(id_of_slot_ptr + s, mask=s_ok, other=-1) + stale = s_ok & (owner != i) + tl.atomic_or(flags_ptr + s, 4, mask=stale) + tl.store(stale_id_ptr + s, i, mask=stale) + tl.atomic_add(stale_n_ptr + s, 1, mask=stale) + t = tl.arange(0, 16) + add = tl.where(t == 1, tl.sum(held.to(tl.int64), axis=0), 0) + add = tl.where(t == 2, tl.sum(take.to(tl.int64), axis=0), add) + add = tl.where(t == 3, tl.sum((held & ~take).to(tl.int64), axis=0), add) + add = tl.where(t == 6, tl.sum(fwd.to(tl.int64), axis=0), add) + add = tl.where(t == 7, tl.sum(stale.to(tl.int64), axis=0), add) + add = tl.where(t == 8, tl.sum(out_of_range.to(tl.int64), axis=0), add) + add = tl.where(t == 9, tl.sum((live & ~s_ok).to(tl.int64), axis=0), add) + add = tl.where((t == 0) & (tl.program_id(0) == 0), 1, add) + tl.atomic_add(totals_ptr + t, add, mask=add != 0, sem="relaxed") + + +@triton.jit +def _audit_nums_kernel(count_ptr, nums_ptr, L: tl.constexpr, C: tl.constexpr, R: tl.constexpr, BLOCK: tl.constexpr): + j = tl.arange(0, BLOCK) + m = j < L * C + cnt = tl.load(count_ptr + j // C, mask=m, other=0) + tl.store(nums_ptr + j, tl.minimum(tl.maximum(cnt - (j % C) * R, 0), R).to(tl.int64), mask=m) + + +@triton.jit +def _audit_compare_kernel(cache_ptrs, ref_ptrs, words_ptr, slots_ptr, num_ptr, bad_cnt_ptr, bad_first_ptr, + NB: tl.constexpr, BLOCK: tl.constexpr): + """Program (row, bank, chunk): one chunk of a listed slot's bank row against its gathered host row.""" + r = tl.program_id(0) + if r < tl.load(num_ptr): + b = tl.program_id(1) + words = tl.load(words_ptr + b) + off = tl.program_id(2).to(tl.int64) * BLOCK + tl.arange(0, BLOCK) + m = off < words + slot = tl.load(slots_ptr + r).to(tl.int64) + got = tl.load(tl.load(cache_ptrs + b).to(tl.pointer_type(tl.int32)) + slot * words + off, mask=m, other=0) + want = tl.load(tl.load(ref_ptrs + b).to(tl.pointer_type(tl.int32)) + r * words + off, mask=m, other=0) + bad = m & (got != want) + nbad = tl.sum(bad.to(tl.int32), axis=0) + if nbad > 0: + tl.atomic_add(bad_cnt_ptr + slot * NB + b, nbad) + tl.atomic_min(bad_first_ptr + slot * NB + b, tl.min(tl.where(bad, off, _NO_BAD), axis=0).to(tl.int32)) + + +@triton.jit(do_not_specialize=["num_slots", "num_ids"]) +def _audit_report_kernel( + id_of_slot_ptr, slot_for_id_ptr, usage_ptr, lstep_ptr, dstep_ptr, + flags_ptr, stale_id_ptr, stale_n_ptr, bad_cnt_ptr, bad_first_ptr, reported_ptr, + meta_hist_ptr, meta_cur_ptr, bytes_hist_ptr, bytes_cur_ptr, + ring_ptr, cursor_ptr, totals_ptr, count_ptr, + num_slots, num_ids, + NB: tl.constexpr, NB_P: tl.constexpr, L: tl.constexpr, L_P: tl.constexpr, RING_N: tl.constexpr, + BLOCK: tl.constexpr, +): + """Keep each new bad (slot, owner) with its history, count the findings and clear the work arrays.""" + s = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + ms = s < num_slots + b = tl.arange(0, NB_P) + cell = ms[:, None] & (b < NB)[None, :] + cnt = tl.load(bad_cnt_ptr + s[:, None] * NB + b[None, :], mask=cell, other=0) + first = tl.load(bad_first_ptr + s[:, None] * NB + b[None, :], mask=cell, other=_NO_BAD) + bank_bad = cnt > 0 + bytes_bad = tl.sum(bank_bad.to(tl.int32), axis=1) > 0 + bank_mask = tl.sum(tl.where(bank_bad, 1 << b[None, :], 0), axis=1) + first_bank = tl.min(tl.where(bank_bad, b[None, :], NB_P), axis=1) + first_word = tl.sum(tl.where(b[None, :] == first_bank[:, None], first, 0), axis=1) + fl = tl.load(flags_ptr + s, mask=ms, other=0) | tl.where(bytes_bad, 1, 0) + bad = ms & (fl != 0) + owner = tl.load(id_of_slot_ptr + s, mask=ms, other=-1) + sid = tl.load(stale_id_ptr + s, mask=ms, other=-1) + key = (owner.to(tl.int64) + 1) * (num_ids + 2) + sid + 1 + new = bad & (tl.load(reported_ptr + s, mask=ms, other=-1) != key) + tl.store(reported_ptr + s, tl.where(bad, key, -1), mask=ms) + pos = tl.atomic_add(cursor_ptr + s * 0, 1, mask=new) + keep = new & (pos < RING_N) + in_range = ms & (owner >= 0) & (owner < num_ids) + f = tl.arange(0, 16) + z = tl.zeros([BLOCK], tl.int64) + vals = tl.zeros([BLOCK, 16], tl.int64) - 1 + vals = _put(vals, f, 0, z + tl.load(totals_ptr) - 1) + vals = _put(vals, f, 1, z + tl.load(dstep_ptr)) + vals = _put(vals, f, 2, z + tl.load(lstep_ptr)) + vals = _put(vals, f, 3, s) + vals = _put(vals, f, 4, owner) + vals = _put(vals, f, 5, tl.load(slot_for_id_ptr + owner, mask=in_range, other=-2)) + vals = _put(vals, f, 6, tl.load(usage_ptr + s, mask=ms, other=-1)) + vals = _put(vals, f, 7, fl) + vals = _put(vals, f, 8, bank_mask) + vals = _put(vals, f, 9, tl.where(bytes_bad, first_bank, -1)) + vals = _put(vals, f, 10, tl.where(bytes_bad, first_word.to(tl.int64) * 4, -1)) + vals = _put(vals, f, 11, tl.sum(cnt, axis=1)) + vals = _put(vals, f, 12, sid) + vals = _put(vals, f, 13, tl.load(stale_n_ptr + s, mask=ms, other=0)) + vals = _put(vals, f, 14, tl.load(meta_cur_ptr + s, mask=ms, other=0)) + vals = _put(vals, f, 15, tl.load(bytes_cur_ptr + s, mask=ms, other=0)) + rec = ring_ptr + pos.to(tl.int64) * (16 + 2 * _D * _F) + tl.store(rec[:, None] + f[None, :], vals, mask=keep[:, None]) + j = tl.arange(0, 64) + jm = keep[:, None] & (j < _D * _F)[None, :] + src = s.to(tl.int64)[:, None] * (_D * _F) + j[None, :] + tl.store(rec[:, None] + 16 + j[None, :], tl.load(meta_hist_ptr + src, mask=jm, other=-1), mask=jm) + tl.store(rec[:, None] + 16 + _D * _F + j[None, :], tl.load(bytes_hist_ptr + src, mask=jm, other=-1), mask=jm) + t = tl.arange(0, 16) + add = tl.where(t == 4, tl.sum(bad.to(tl.int64), axis=0), 0) + add = tl.where(t == 5, tl.sum((ms & bytes_bad).to(tl.int64), axis=0), add) + add = tl.where(t == 10, tl.sum(new.to(tl.int64), axis=0), add) + add = tl.where(t == 11, tl.sum((bad & ~new).to(tl.int64), axis=0), add) + tl.atomic_add(totals_ptr + t, add, mask=add != 0, sem="relaxed") + tl.store(flags_ptr + s, tl.zeros([BLOCK], tl.int32), mask=ms) + tl.store(stale_id_ptr + s, tl.zeros([BLOCK], tl.int32) - 1, mask=ms) + tl.store(stale_n_ptr + s, tl.zeros([BLOCK], tl.int32), mask=ms) + tl.store(bad_cnt_ptr + s[:, None] * NB + b[None, :], tl.zeros([BLOCK, NB_P], tl.int32), mask=cell) + tl.store(bad_first_ptr + s[:, None] * NB + b[None, :], tl.zeros([BLOCK, NB_P], tl.int32) + _NO_BAD, mask=cell) + if tl.program_id(0) == 0: + ll = tl.arange(0, L_P) + tl.store(count_ptr + ll, tl.zeros([L_P], tl.int32), mask=ll < L) + + +class SlotAuditor: + """Per-slot writer history (recorded by ``record_*`` on the writer's stream) and the periodic + audit (``end_decode_step`` -> ``scan`` on the compute stream). ``bind`` sizes everything for a + cache's slots and banks; a rebuild rebinds and starts a new history, the totals and records stay. + Holds no reference to its cache: every call takes it.""" + + def __init__(self, num_layers: int, num_experts: int, device: torch.device, *, interval: int, + ring: int | None = None, scratch_bytes: int = SCRATCH_BYTES) -> None: + self.num_layers = num_layers + self.num_experts = num_experts + self.device = device + self.interval = interval + self.ring_n = RING if ring is None else ring + self.scratch_bytes = scratch_bytes + ring = self.ring_n + self.dstep = torch.zeros((1,), dtype=torch.int64, device=device) + self.seq = torch.zeros((1,), dtype=torch.int64, device=device) + self.totals = torch.zeros((NUM_TOTALS,), dtype=torch.int64, device=device) + self.ring = torch.full((ring, RECORD_FIELDS), -1, dtype=torch.int64, device=device) + self.cursor = torch.zeros((1,), dtype=torch.int64, device=device) + self.count = torch.zeros((num_layers,), dtype=torch.int32, device=device) + self.num_slots = 0 + self.bank_names: list[str] = [] + self._feats: list[int] = [] + self._steps = 0 + self._printed = 0 + self._last = torch.zeros((NUM_TOTALS,), dtype=torch.int64) + self._dummy = torch.zeros((1,), dtype=torch.int32, device=device) + logger.info(f"MoE slot audit: recording every slot writer, auditing every held slot each {interval} decode steps") + + # ------------------------------------------------------------------ setup + + def bind(self, cache) -> None: + """Size the history and audit arrays for ``cache``'s slots and banks (a rebuild starts a new + history) and point the byte checks at its banks.""" + caches = [bank for _, bank in cache.banks] + feats = [math.prod(bank.shape[1:]) * bank.element_size() for bank in caches] + assert all(f % 4 == 0 for f in feats), feats + if self.bank_names: + assert list(cache.bank_schema) == self.bank_names and feats == self._feats, "the bank layout changed" + else: + self.bank_names = list(cache.bank_schema) + self._feats = feats + self._alloc_scratch(cache, feats) + self.cache_ptrs = torch.tensor([bank.data_ptr() for bank in caches], dtype=torch.int64, device=self.device) + S, dev = cache.cache_size, self.device + self.num_slots = S + self.meta_hist = torch.full((S, DEPTH, EVENT_FIELDS), -1, dtype=torch.int64, device=dev) + self.bytes_hist = torch.full((S, DEPTH, EVENT_FIELDS), -1, dtype=torch.int64, device=dev) + self.meta_cur = torch.zeros((S,), dtype=torch.int32, device=dev) + self.bytes_cur = torch.zeros((S,), dtype=torch.int32, device=dev) + self.flags = torch.zeros((S,), dtype=torch.int32, device=dev) + self.stale_id = torch.full((S,), -1, dtype=torch.int32, device=dev) + self.stale_n = torch.zeros((S,), dtype=torch.int32, device=dev) + self.reported = torch.full((S,), -1, dtype=torch.int64, device=dev) + nb = len(feats) + self.bad_cnt = torch.zeros((S * nb,), dtype=torch.int32, device=dev) + self.bad_first = torch.full((S * nb,), 0x7FFFFFFF, dtype=torch.int32, device=dev) + # a correct cache holds at most num_experts slots of a layer; more are counted unchecked + per_layer = min(self.num_experts, S) + self.chunks = triton.cdiv(per_layer, self.rows) + self.cap = self.chunks * self.rows + self.list_rows = torch.zeros((self.num_layers, self.cap), dtype=torch.int32, device=dev) + self.list_slots = torch.zeros((self.num_layers, self.cap), dtype=torch.int32, device=dev) + self.nums = torch.zeros((self.num_layers * self.chunks,), dtype=torch.int64, device=dev) + # layers without a device alias (LOCKED/PAGEABLE host banks) cannot be gathered + self.layer_ok_host = [layer not in cache._unpinned_layers for layer in range(self.num_layers)] + self.layer_ok = torch.tensor(self.layer_ok_host, dtype=torch.int32, device=dev) + self.count.zero_() + + def _alloc_scratch(self, cache, feats: list[int]) -> None: + row = sum(feats) + self.rows = max(1, min(MAX_CHUNK_ROWS, (self.scratch_bytes - 1) // row)) + rows = self.rows + self.scratch = torch.empty((rows * row,), dtype=torch.uint8, device=self.device) + offsets = [rows * sum(feats[:b]) for b in range(len(feats))] + self.scratch_views = [ + self.scratch[o : o + rows * f].view(bank.dtype).view(rows, *bank.shape[1:]) + for o, f, (_, bank) in zip(offsets, feats, cache.banks) + ] + self.scratch_ptrs = torch.tensor([self.scratch.data_ptr() + o for o in offsets], dtype=torch.int64, device=self.device) + self.words = torch.tensor([f // 4 for f in feats], dtype=torch.int64, device=self.device) + self.iota = torch.arange(rows, dtype=torch.int32, device=self.device) + self._word_chunks = triton.cdiv(max(feats) // 4, COMPARE_WORDS) + logger.info(f"MoE slot audit scratch: {self.scratch.numel() / 2**20:.1f} MiB, {rows} experts a chunk") + + # ------------------------------------------------------------------ recording (graph-capturable) + + def begin_step(self) -> None: + """Advance the decode-step clock (the first GPU layer's decode ensure calls this).""" + _bump_step_kernel[(1,)](self.dstep) + + def _launch(self, cache, kind: int, layer: int, *, meta: bool, width: int, dst=None, src=None, aux=None, + num=None, start: int = 0, row: int = ROW_PLAN, row_base: int = 0, pred: int = PRED_ALL) -> None: + if width <= 0 or not self.num_slots: + return + block = min(triton.next_power_of_2(width), 128) + d = self._dummy + _record_kernel[(triton.cdiv(width, block),)]( + self.meta_hist if meta else self.bytes_hist, self.meta_cur if meta else self.bytes_cur, + self.seq, self.dstep, cache.step, cache.id_of_slot, cache.usage, self.totals, + d if dst is None else dst, d if src is None else src, d if aux is None else aux, d if num is None else num, + kind, layer, start, row_base, width, self.num_slots, + PLAN=dst is not None, ROW=row, HAS_AUX=aux is not None, HAS_NUM=num is not None, PRED=pred, + E=self.num_experts, BLOCK=block, num_warps=4, + ) + + def record_plan(self, cache, kind: int, layer: int, dst: torch.Tensor, src: torch.Tensor | None, num: torch.Tensor, + *, meta: bool, width: int | None = None, aux: torch.Tensor | None = None, row_base: int = 0) -> None: + """The writer just wrote slots ``dst[:num]`` (rows ``src[:num]``, else ``dst - row_base``).""" + width = dst.numel() if width is None else min(width, dst.numel()) + self._launch(cache, kind, layer, meta=meta, width=width, dst=dst, src=src, aux=aux, num=num, + row=ROW_PLAN if src is not None else ROW_FROM_SLOT, row_base=row_base) + + def record_range(self, cache, kind: int, layer: int, start: int, n: int, *, meta: bool) -> None: + """The writer just wrote expert ``i`` of ``layer`` into slot ``start + i`` for ``i < n``.""" + self._launch(cache, kind, layer, meta=meta, width=n, start=start, row=ROW_INDEX) + + def record_clear(self, cache, kind: int, context_layer: int, start: int, n: int, *, layer_above_e: bool = False) -> None: + """Before a writer clears held slots in ``[start, start + n)`` (with ``layer_above_e``: only + those >= E held by ``context_layer``); the event names the old owner.""" + self._launch(cache, kind, context_layer, meta=True, width=n, start=start, row=ROW_OLD_OWNER, + pred=PRED_LAYER_ABOVE_E if layer_above_e else PRED_HELD) + + # ------------------------------------------------------------------ audit + + def end_decode_step(self, cache) -> bool: + """After a decode forward, on the compute stream: audit every ``interval`` steps.""" + self._steps += 1 + if self._steps % self.interval: + return False + self.scan(cache) + return True + + def scan(self, cache) -> None: + """Audit every held slot now (enqueued on the current stream; no host sync).""" + S, num_ids = self.num_slots, self.num_layers * self.num_experts + assert S == cache.cache_size and cache.id_of_slot.numel() == S, "bind after a rebuild" + block = 1024 + _audit_meta_kernel[(triton.cdiv(max(S, num_ids), block),)]( + cache.id_of_slot, cache.slot_for_id.view(-1), self.layer_ok, self.count, self.list_rows, self.list_slots, + self.flags, self.stale_id, self.stale_n, self.totals, S, num_ids, + E=self.num_experts, CAP=self.cap, BLOCK=block, num_warps=4, + ) + n_lists = self.num_layers * self.chunks + _audit_nums_kernel[(1,)](self.count, self.nums, L=self.num_layers, C=self.chunks, R=self.rows, + BLOCK=triton.next_power_of_2(n_lists)) + nb = len(self.bank_names) + R = self.rows + for layer in range(self.num_layers): + if not self.layer_ok_host[layer]: + continue + for c in range(self.chunks): + rows = self.list_rows[layer, c * R : (c + 1) * R] + slots = self.list_slots[layer, c * R : (c + 1) * R] + num = self.nums[layer * self.chunks + c : layer * self.chunks + c + 1] + self._gather(cache, layer, rows, num) + _audit_compare_kernel[(R, nb, self._word_chunks)]( + self.cache_ptrs, self.scratch_ptrs, self.words, slots, num, self.bad_cnt, self.bad_first, + NB=nb, BLOCK=COMPARE_WORDS, num_warps=8, + ) + rblock = 256 + _audit_report_kernel[(triton.cdiv(S, rblock),)]( + cache.id_of_slot, cache.slot_for_id.view(-1), cache.usage, cache.step, self.dstep, + self.flags, self.stale_id, self.stale_n, self.bad_cnt, self.bad_first, self.reported, + self.meta_hist, self.meta_cur, self.bytes_hist, self.bytes_cur, + self.ring, self.cursor, self.totals, self.count, S, num_ids, + NB=nb, NB_P=triton.next_power_of_2(nb), L=self.num_layers, L_P=triton.next_power_of_2(self.num_layers), + RING_N=self.ring_n, BLOCK=rblock, num_warps=4, + ) + + def _gather(self, cache, layer: int, rows: torch.Tensor, num: torch.Tensor) -> None: + # the plain fused copy, not the slim one prefetch copies use: a slim-kernel bug cannot hide itself + if cache._copy_fused_ok: + from freetoken.kernel.fast_index_copy import fast_index_copy_multi_jit + + fast_index_copy_multi_jit( + self.scratch_ptrs, cache._copy_src_ptrs[layer], cache._copy_feat_bytes[cache.layer_layout[layer]], self.iota, rows, num, + ) + return + from freetoken.kernel import fast_index_copy_jit + + for (per_layer, _), view in zip(cache.banks, self.scratch_views): + fast_index_copy_jit(view, self.iota, per_layer[layer], rows, num) + + # ------------------------------------------------------------------ reporting (host syncs) + + def take_window(self) -> tuple[list[int], list[int], list[list[int]], int]: + """(totals now, totals of this window, records kept since the last call, records in all).""" + totals = self.totals.cpu() + window = (totals - self._last).tolist() + self._last = totals + cursor = int(self.cursor.item()) + upto = min(cursor, self.ring_n) + records = self.ring[self._printed : upto].cpu().tolist() if upto > self._printed else [] + self._printed = max(self._printed, upto) + return totals.tolist(), window, records, cursor + + def report_window(self, steps: int) -> list[tuple[str, str]]: + """(level, line) pairs for one MOE_STATS_INTERVAL window; nothing when no audit ran.""" + totals, window, records, cursor = self.take_window() + if not window[T_AUDITS] and not records: + return [] + bad = bool(window[T_BAD] or window[T_MAP_RANGE] or window[T_BAD_PLAN] or records) + lines = [("warning" if bad else "info", + f"MoE slot audit ({steps} decode steps, every {self.interval}): {self.format_counts(window)}; " + f"session: {totals[T_AUDITS]} audits, {totals[T_NEW]} bad (slot, owner) pairs")] + lines += [("warning", "MoE slot audit record: " + self.format_record(r)) for r in records] + if records and cursor > self.ring_n: + lines.append(("warning", f"MoE slot audit: {cursor} bad (slot, owner) pairs in all, the first {self.ring_n} are kept")) + return lines + + def report_session(self) -> list[tuple[str, str]]: + """Session totals and every kept record (shutdown).""" + totals = self.totals.cpu().tolist() + cursor = int(self.cursor.item()) + if not totals[T_AUDITS]: + return [] + bad = bool(totals[T_BAD] or totals[T_MAP_RANGE] or totals[T_BAD_PLAN]) + lines = [("warning" if bad else "info", f"MoE slot audit (session, every {self.interval} decode steps): " + f"{self.format_counts(totals)}")] + records = self.ring[: min(cursor, self.ring_n)].cpu().tolist() + lines += [("warning", "MoE slot audit record: " + self.format_record(r)) for r in records] + if cursor > self.ring_n: + lines.append(("warning", f"MoE slot audit: {cursor} bad (slot, owner) pairs in all, the first {self.ring_n} are kept")) + return lines + + @staticmethod + def format_counts(t: list[int]) -> str: + audits = t[T_AUDITS] + per = (lambda v: f"{v / audits:.0f}") if audits else str + text = ( + f"{audits} audits, {per(t[T_HELD])} held slots an audit ({per(t[T_BYTES_CHECKED])} byte-checked), " + f"bad slots={t[T_BAD]} (bytes={t[T_BYTES]}, map={t[T_FWD]}, stale slot_for_id={t[T_STALE]}, " + f"owner out of range={t[T_OWNER_RANGE]}), new={t[T_NEW]}, repeat={t[T_STILL]}, events={t[T_EVENTS]}" + ) + if t[T_UNCHECKED]: + text += f", held slots not byte-checked={t[T_UNCHECKED]}" + if t[T_MAP_RANGE]: + text += f", slot_for_id entries past the cache={t[T_MAP_RANGE]}" + if t[T_BAD_PLAN]: + text += f", plan entries naming no slot={t[T_BAD_PLAN]}" + return text + + # ------------------------------------------------------------------ record decoding + + def _name(self, flat: int) -> str: + e = self.num_experts + return f"L{flat // e}/e{flat % e}" if flat >= 0 else str(flat) + + @staticmethod + def events(rec: list[int]) -> list[dict]: + """Both rings of a record as event dicts in ``seq`` order (``ring`` 'meta' or 'bytes').""" + out = [] + for ring, base in (("meta", HEADER), ("bytes", HEADER + RING_EVENTS)): + for d in range(DEPTH): + ev = rec[base + d * EVENT_FIELDS : base + (d + 1) * EVENT_FIELDS] + if ev[E_SEQ] >= 0: + out.append({ + "ring": ring, "seq": ev[E_SEQ], "kind": ev[E_KIND], "dstep": ev[E_DSTEP], "lstep": ev[E_LSTEP], + "layer": ev[E_LAYER], "row": ev[E_ROW], "aux": ev[E_AUX], "owner": ev[E_OWNER], + "usage": ev[E_USAGE], "pos": ev[E_POS], + }) + return sorted(out, key=lambda ev: ev["seq"]) + + def _flat(self, ev: dict) -> int: + return ev["layer"] * self.num_experts + ev["row"] if ev["layer"] >= 0 and ev["row"] >= 0 else -1 + + def format_event(self, ev: dict) -> str: + e = self.num_experts + text = f"#{ev['seq']} {KIND_NAMES.get(ev['kind'], ev['kind'])} L{ev['layer']}/e{ev['row']}" + if ev["kind"] == PREFILL_HIT_D2D: + text += f" from slot {ev['aux']}" + elif ev["kind"] in (MATERIALIZE_CLEAR, INVALIDATE) and ev["aux"] >= 0: + text += f" (for layer {ev['aux']})" + seen = "slot held" if ev["ring"] == "bytes" else "then held" + if ev["kind"] in (MATERIALIZE_CLEAR, INVALIDATE, RESET): + seen = "held" + owner = f"L{ev['owner'] // e}/e{ev['owner'] % e}" if ev["owner"] >= 0 else "nothing" + return f"{text} @step {ev['dstep']}/lru {ev['lstep']} ({seen} {owner}, usage {ev['usage']})" + + def _kind(self, ev: dict) -> str: + return KIND_NAMES.get(ev["kind"], str(ev["kind"])) + + def _pending_install(self, copy: dict, meta: list[dict], byte: list[dict], before: dict) -> dict | None: + """The earlier install (before ``before``) that ``copy`` completes: an install of the id it + copied with no copy of that id between them.""" + flat = self._flat(copy) + for inst in reversed([ev for ev in meta if ev["kind"] in INSTALLS and ev["seq"] < before["seq"]]): + if self._flat(inst) == flat: + done = any(self._flat(ev) == flat and inst["seq"] < ev["seq"] < copy["seq"] for ev in byte) + return None if done else inst + return None + + def _bytes_story(self, owner: int, meta: list[dict], byte: list[dict]) -> str: + insts = [ev for ev in meta if ev["kind"] in INSTALLS and self._flat(ev) == owner] + if not insts: + return (f"no recorded install of {self._name(owner)} in the kept history " + f"(an unrecorded map writer, or it is older than the last {DEPTH} events)") + inst = insts[-1] + name = self._kind(inst) + after = [ev for ev in byte if ev["seq"] > inst["seq"]] + own = [ev for ev in after if self._flat(ev) == owner] + if own: + last = after[-1] + if self._flat(last) == owner: + return (f"{name} and {self._kind(last)} of {self._name(owner)} (#{last['seq']}), then the bytes changed " + f"with no recorded write: overwritten after the last recorded write, or that copy did not land") + text = (f"{name} and {self._kind(own[-1])} of {self._name(owner)} (#{own[-1]['seq']}), then " + f"{self._kind(last)} of {self._name(self._flat(last))} (#{last['seq']}) overwrote it") + prior = self._pending_install(last, meta, byte, inst) + if prior is not None: + text += f": the late copy of the earlier {self._kind(prior)} of {self._name(self._flat(prior))} (#{prior['seq']})" + return text + text = f"{name}, no {KIND_NAMES[COPY_FOR[inst['kind']]]}" + if not after: + return text + last = after[-1] + prior = self._pending_install(last, meta, byte, inst) + if prior is not None: + return (f"{text}; the last write, {self._kind(last)} of {self._name(self._flat(last))} (#{last['seq']}), " + f"is the copy of the earlier {self._kind(prior)} (#{prior['seq']}) and landed after this install") + where = (f"from layer {last['layer']}'s source" if last["row"] == owner % self.num_experts + else f"of {self._name(self._flat(last))}") + return f"{name} of {self._name(owner)}, then {self._kind(last)} {where} ({self._name(self._flat(last))}, #{last['seq']})" + + def diagnose(self, rec: list[int]) -> str: + """What the history says about a bad slot, in words.""" + owner, fl = rec[H_OWNER], rec[H_FLAGS] + evs = self.events(rec) + meta = [ev for ev in evs if ev["ring"] == "meta"] + byte = [ev for ev in evs if ev["ring"] == "bytes"] + parts = [] + if fl & F_OWNER_RANGE: + parts.append(f"id_of_slot={owner} names no expert") + if fl & F_BYTES and owner >= 0: + parts.append(self._bytes_story(owner, meta, byte)) + if fl & F_FWD: + parts.append(f"slot_for_id[{self._name(owner)}]={rec[H_MAPPED]}, not this slot") + if fl & F_STALE: + sid = rec[H_STALE_ID] + text = (f"stale slot_for_id: {self._name(sid)} -> this slot, which holds {self._name(owner)}" + + (f" ({rec[H_STALE_N]} such ids)" if rec[H_STALE_N] > 1 else "")) + gone = [ev for ev in meta if ev["kind"] in INSTALLS and self._flat(ev) == sid] + if gone: + later = [ev for ev in meta if ev["seq"] > gone[-1]["seq"]] + text += (f"; {self._name(sid)} was installed here by {self._kind(gone[-1])} #{gone[-1]['seq']}" + + (f" and replaced by {self._kind(later[0])} #{later[0]['seq']} without clearing " + f"its map entry" if later else "")) + parts.append(text) + # an install's recorder reads the slot right after the install on the same stream + for ev in meta: + if ev["kind"] in INSTALLS and ev["owner"] != self._flat(ev): + parts.append(f"right after the {self._kind(ev)} of {self._name(self._flat(ev))} (#{ev['seq']}) the slot " + f"named {self._name(ev['owner'])}: another map writer ran concurrently") + return "; ".join(parts) or "clean" + + def format_record(self, rec: list[int]) -> str: + fl = rec[H_FLAGS] + text = (f"audit={rec[H_AUDIT]} step={rec[H_DSTEP]} lru={rec[H_LSTEP]} slot={rec[H_SLOT]} " + f"holds {self._name(rec[H_OWNER])} (slot_for_id={rec[H_MAPPED]}, usage={rec[H_USAGE]})") + if fl & F_BYTES: + banks = [self.bank_names[b] if b < len(self.bank_names) else str(b) + for b in range(max(len(self.bank_names), 1)) if rec[H_BANKS] >> b & 1] + first = self.bank_names[rec[H_FIRST_BANK]] if 0 <= rec[H_FIRST_BANK] < len(self.bank_names) else rec[H_FIRST_BANK] + text += (f" | bytes differ in {', '.join(banks)} (first at {first}+{rec[H_FIRST_BYTE]}, " + f"{rec[H_BAD_WORDS]} bad words)") + text += f" | diagnosis: {self.diagnose(rec)}" + history = "; ".join(self.format_event(ev) for ev in self.events(rec)) + return text + f" | history ({rec[H_META_N]} map / {rec[H_BYTES_N]} byte events in all): {history or 'none'}" diff --git a/python/freetoken/moe/verify.py b/python/freetoken/moe/verify.py new file mode 100644 index 0000000000..cc71e19d9d --- /dev/null +++ b/python/freetoken/moe/verify.py @@ -0,0 +1,516 @@ +"""Decode expert-cache verifier (``FREETOKEN_MOE_PREFETCH_VERIFY``): a debug detector for slot races. + +For every GPU decode layer call it checks, on the compute stream and in its order, what the +grouped GEMM is about to read and has just read: + +- meta (before and after the GEMM): each routed entry's slot maps to the expert the router chose, + in both slot maps (``id_of_slot[s] == layer * E + id`` and ``slot_for_id[layer, id] == s``); the + router's ids are copied before ensure rewrites them in place; +- bytes before the GEMM (``full``): every routed slot's row in every bank equals the host bank row + of that expert for that layer, gathered into a scratch buffer with the fused copy; +- bytes after the GEMM (``full``): the same slots against the same scratch, so a slot overwritten + while the GEMM ran shows up even when a later copy restores it. + +Nothing here records or waits on an event: the prefetch streams keep exactly the edges they have +without it. The checks do delay the compute stream (``full`` gathers every routed expert over +PCIe), which moves where the other streams land relative to it. ``meta`` skips the byte checks and +costs a few one-CTA launches per layer. Counts accumulate per layer on the device and the first +``RING`` offending records are kept with the slot state seen at check time; the engine reports +them every MOE_STATS_INTERVAL decode steps and at shutdown. +""" +from __future__ import annotations + +import math + +import torch +import triton +import triton.language as tl + +from freetoken.env import ENV +from freetoken.utils import init_logger + +logger = init_logger(__name__) + +MODES = ("off", "meta", "full") +# counters[layer, col]: checked calls, offending records per check, routed experts left unchecked +CHECKS, META_BAD, PRE_BAD, POST_BAD, POST_META_BAD, UNCHECKED = range(6) +NUM_COLS = 6 +COL_NAMES = ("checks", "meta_bad", "pre_bad", "post_bad", "post_meta_bad", "unchecked") +META_PRE, BYTES_PRE, META_POST, BYTES_POST = range(4) +KINDS = ("meta_pre", "bytes_pre", "meta_post", "bytes_post") +# record fields; owner / mapped / usage / step are read when the check runs +(R_CALL, R_LAYER, R_ROW, R_EXPERT, R_SLOT, R_BANK, R_KIND, R_OWNER, R_MAPPED, R_USAGE, R_STEP, + R_SOURCE, R_FIRST_BYTE, R_BAD_WORDS) = range(14) +NUM_FIELDS = 14 +# R_SOURCE bits: the slot is in this call's demand copy plan, this layer's prefetch copy plan, or +# the next layer's prefetch plan (installed earlier on the compute stream) +SRC_DEMAND, SRC_PREFETCH, SRC_NEXT_PREFETCH = 1, 2, 4 +RING = 64 +# routed entries a call may have (bs * top_k); wider eager batches are counted unchecked +MAX_ENTRIES = 64 +CHUNK_WORDS = 4096 +_NF = tl.constexpr(NUM_FIELDS) +_NC = tl.constexpr(NUM_COLS) +_NO_BAD = tl.constexpr(0x7FFFFFFF) + + +def resolve_mode(mode: str | None = None) -> str: + raw = (ENV.MOE_PREFETCH_VERIFY.value if mode is None else mode).strip().lower() + if raw in ("", "0", "off", "false", "no"): + return "off" + if raw in ("1", "on", "true", "yes", "full"): + return "full" + if raw == "meta": + return "meta" + raise ValueError(f"FREETOKEN_MOE_PREFETCH_VERIFY={raw!r}: expected 0, 1 (full) or meta") + + +@triton.jit +def _put(vals, f, idx: tl.constexpr, v): + return tl.where(f[None, :] == idx, v.to(tl.int64)[:, None], vals) + + +@triton.jit +def _in_plan(s, plan_ptr, num_ptr, W: tl.constexpr, W_P: tl.constexpr): + j = tl.arange(0, W_P) + plan = tl.load(plan_ptr + j, mask=j < tl.minimum(tl.load(num_ptr), W), other=-3) + return tl.sum((s[:, None] == plan[None, :]).to(tl.int32), axis=1) > 0 + + +@triton.jit +def _source(s, dem_ptr, dem_num_ptr, pf_ptr, pf_num_ptr, nx_ptr, nx_num_ptr, + DEM_W: tl.constexpr, DEM_P: tl.constexpr, PF_W: tl.constexpr, PF_P: tl.constexpr, + HAS_PF: tl.constexpr, HAS_NX: tl.constexpr): + src = tl.where(_in_plan(s, dem_ptr, dem_num_ptr, DEM_W, DEM_P), 1, 0) + if HAS_PF: + src += tl.where(_in_plan(s, pf_ptr, pf_num_ptr, PF_W, PF_P), 2, 0) + if HAS_NX: + src += tl.where(_in_plan(s, nx_ptr, nx_num_ptr, PF_W, PF_P), 4, 0) + return src + + +@triton.jit +def _write(ring_ptr, cursor, bad, vals, RING_N: tl.constexpr): + """Append the ``bad`` rows of ``vals`` after ``cursor`` while the ring has room; the new cursor.""" + f = tl.arange(0, 16) + pos = cursor + tl.cumsum(bad.to(tl.int64), axis=0) - 1 + keep = bad & (pos < RING_N) + tl.store(ring_ptr + pos[:, None] * _NF + f[None, :], vals, mask=keep[:, None] & (f[None, :] < _NF)) + return cursor + tl.sum(bad.to(tl.int64), axis=0) + + +@triton.jit +def _meta(ids_ptr, slots_ptr, id_of_slot_ptr, slot_for_id_ptr, usage_ptr, step_ptr, + dem_ptr, dem_num_ptr, pf_ptr, pf_num_ptr, nx_ptr, nx_num_ptr, + ring_ptr, cursor, counters_ptr, col: tl.constexpr, kind: tl.constexpr, call, layer, num_slots, + N: tl.constexpr, N_P: tl.constexpr, E: tl.constexpr, + DEM_W: tl.constexpr, DEM_P: tl.constexpr, PF_W: tl.constexpr, PF_P: tl.constexpr, + HAS_PF: tl.constexpr, HAS_NX: tl.constexpr, RING_N: tl.constexpr): + """Check every routed entry's slot against both maps; returns the new ring cursor.""" + i = tl.arange(0, N_P) + m = i < N + ids = tl.load(ids_ptr + i, mask=m, other=0) + s = tl.load(slots_ptr + i, mask=m, other=-1) + ok_s = m & (s >= 0) & (s < num_slots) + ok_id = m & (ids >= 0) & (ids < E) + fid = layer * E + ids + owner = tl.load(id_of_slot_ptr + s, mask=ok_s, other=-2) + mapped = tl.load(slot_for_id_ptr + fid, mask=ok_id, other=-2) + bad = m & ~(ok_s & ok_id & (owner == fid) & (mapped == s)) + usage = tl.load(usage_ptr + s, mask=ok_s, other=-1) + src = _source(s, dem_ptr, dem_num_ptr, pf_ptr, pf_num_ptr, nx_ptr, nx_num_ptr, DEM_W, DEM_P, PF_W, PF_P, HAS_PF, HAS_NX) + f = tl.arange(0, 16) + z = tl.zeros([N_P], tl.int64) + vals = tl.zeros([N_P, 16], tl.int64) - 1 + vals = _put(vals, f, 0, z + call) + vals = _put(vals, f, 1, z + layer) + vals = _put(vals, f, 2, i) + vals = _put(vals, f, 3, ids) + vals = _put(vals, f, 4, s) + vals = _put(vals, f, 6, z + kind) + vals = _put(vals, f, 7, owner) + vals = _put(vals, f, 8, mapped) + vals = _put(vals, f, 9, usage) + vals = _put(vals, f, 10, z + tl.load(step_ptr)) + vals = _put(vals, f, 11, src) + nbad = tl.sum(bad.to(tl.int64), axis=0) + tl.store(counters_ptr + col, tl.load(counters_ptr + col) + nbad) + return _write(ring_ptr, cursor, bad, vals, RING_N) + + +@triton.jit(do_not_specialize=["layer", "num_slots"]) +def _plan_kernel( + ids_ptr, slots_ptr, id_of_slot_ptr, slot_for_id_ptr, usage_ptr, step_ptr, + dem_ptr, dem_num_ptr, pf_ptr, pf_num_ptr, nx_ptr, nx_num_ptr, + ring_ptr, cursor_ptr, counters_ptr, calls_ptr, + u_ids_ptr, u_slots_ptr, u_rows_ptr, u_num_ptr, slots_copy_ptr, + layer, num_slots, + N: tl.constexpr, N_P: tl.constexpr, E: tl.constexpr, CAP: tl.constexpr, + DEM_W: tl.constexpr, DEM_P: tl.constexpr, PF_W: tl.constexpr, PF_P: tl.constexpr, + HAS_PF: tl.constexpr, HAS_NX: tl.constexpr, BYTES: tl.constexpr, RING_N: tl.constexpr, +): + """Before the GEMM: count the call, check the maps, keep the slots for the after-GEMM checks and + list the distinct (expert, slot) pairs to byte-check (first CAP of them).""" + call = tl.load(calls_ptr + layer) + tl.store(calls_ptr + layer, call + 1) + row = counters_ptr + layer * _NC + tl.store(row + 0, tl.load(row + 0) + 1) + cursor = tl.load(cursor_ptr) + cursor = _meta(ids_ptr, slots_ptr, id_of_slot_ptr, slot_for_id_ptr, usage_ptr, step_ptr, + dem_ptr, dem_num_ptr, pf_ptr, pf_num_ptr, nx_ptr, nx_num_ptr, + ring_ptr, cursor, row, 1, 0, call, layer, num_slots, + N, N_P, E, DEM_W, DEM_P, PF_W, PF_P, HAS_PF, HAS_NX, RING_N) + tl.store(cursor_ptr, cursor) + i = tl.arange(0, N_P) + m = i < N + ids = tl.load(ids_ptr + i, mask=m, other=-1) + s = tl.load(slots_ptr + i, mask=m, other=-1) + tl.store(slots_copy_ptr + i, s, mask=m) + if BYTES: + # an unreadable slot or id is already a meta record; a duplicate pair shares its first row + ok = m & (s >= 0) & (s < num_slots) & (ids >= 0) & (ids < E) + same = (ids[:, None] == ids[None, :]) & (s[:, None] == s[None, :]) & (i[None, :] < i[:, None]) & ok[None, :] + first = ok & (tl.sum(same.to(tl.int32), axis=1) == 0) + rank = tl.cumsum(first.to(tl.int32), axis=0) - 1 + take = first & (rank < CAP) + tl.store(u_ids_ptr + rank, ids, mask=take) + tl.store(u_slots_ptr + rank, s, mask=take) + tl.store(u_rows_ptr + rank, i, mask=take) + total = tl.sum(first.to(tl.int64), axis=0) + kept = tl.minimum(total, CAP) + tl.store(u_num_ptr, kept) + tl.store(row + 5, tl.load(row + 5) + total - kept) + + +@triton.jit +def _compare_kernel(cache_ptrs, ref_ptrs, words_ptr, u_slots_ptr, u_num_ptr, bad_cnt_ptr, bad_first_ptr, + NB: tl.constexpr, BLOCK: tl.constexpr): + """Program (row, bank, chunk): compare one chunk of a routed slot's row with its host copy.""" + r = tl.program_id(0) + b = tl.program_id(1) + words = tl.load(words_ptr + b) + off = tl.program_id(2).to(tl.int64) * BLOCK + tl.arange(0, BLOCK) + m = (r < tl.load(u_num_ptr)) & (off < words) + slot = tl.load(u_slots_ptr + r).to(tl.int64) + got = tl.load(tl.load(cache_ptrs + b).to(tl.pointer_type(tl.int32)) + slot * words + off, mask=m, other=0) + want = tl.load(tl.load(ref_ptrs + b).to(tl.pointer_type(tl.int32)) + r * words + off, mask=m, other=0) + bad = m & (got != want) + nbad = tl.sum(bad.to(tl.int32), axis=0) + if nbad > 0: + tl.atomic_add(bad_cnt_ptr + r * NB + b, nbad) + tl.atomic_min(bad_first_ptr + r * NB + b, tl.min(tl.where(bad, off, _NO_BAD), axis=0).to(tl.int32)) + + +@triton.jit(do_not_specialize=["layer", "num_slots"]) +def _tally_kernel( + ids_ptr, slots_copy_ptr, id_of_slot_ptr, slot_for_id_ptr, usage_ptr, step_ptr, + dem_ptr, dem_num_ptr, pf_ptr, pf_num_ptr, nx_ptr, nx_num_ptr, + ring_ptr, cursor_ptr, counters_ptr, calls_ptr, + u_ids_ptr, u_slots_ptr, u_rows_ptr, u_num_ptr, bad_cnt_ptr, bad_first_ptr, + layer, num_slots, + N: tl.constexpr, N_P: tl.constexpr, E: tl.constexpr, CAP: tl.constexpr, CAP_P: tl.constexpr, + NB: tl.constexpr, NB_P: tl.constexpr, + DEM_W: tl.constexpr, DEM_P: tl.constexpr, PF_W: tl.constexpr, PF_P: tl.constexpr, + HAS_PF: tl.constexpr, HAS_NX: tl.constexpr, BYTES: tl.constexpr, POST: tl.constexpr, + BYTES_KIND: tl.constexpr, BYTES_COL: tl.constexpr, RING_N: tl.constexpr, +): + """Turn the compare flags into counts and records (and clear them); after the GEMM also recheck + the maps of the slots it read.""" + call = tl.load(calls_ptr + layer) - 1 + row = counters_ptr + layer * _NC + cursor = tl.load(cursor_ptr) + if POST: + cursor = _meta(ids_ptr, slots_copy_ptr, id_of_slot_ptr, slot_for_id_ptr, usage_ptr, step_ptr, + dem_ptr, dem_num_ptr, pf_ptr, pf_num_ptr, nx_ptr, nx_num_ptr, + ring_ptr, cursor, row, 4, 2, call, layer, num_slots, + N, N_P, E, DEM_W, DEM_P, PF_W, PF_P, HAS_PF, HAS_NX, RING_N) + if BYTES: + j = tl.arange(0, CAP_P * NB_P) + r = j // NB_P + b = j % NB_P + cell = (r < CAP) & (b < NB) + live = cell & (r < tl.load(u_num_ptr)) + cnt = tl.load(bad_cnt_ptr + r * NB + b, mask=live, other=0) + first = tl.load(bad_first_ptr + r * NB + b, mask=live, other=_NO_BAD) + tl.store(bad_cnt_ptr + r * NB + b, tl.zeros([CAP_P * NB_P], tl.int32), mask=cell) + tl.store(bad_first_ptr + r * NB + b, tl.zeros([CAP_P * NB_P], tl.int32) + _NO_BAD, mask=cell) + bad = live & (cnt > 0) + ids = tl.load(u_ids_ptr + r, mask=live, other=0) + s = tl.load(u_slots_ptr + r, mask=live, other=-1) + ok_s = live & (s >= 0) & (s < num_slots) + src = _source(s, dem_ptr, dem_num_ptr, pf_ptr, pf_num_ptr, nx_ptr, nx_num_ptr, DEM_W, DEM_P, PF_W, PF_P, HAS_PF, HAS_NX) + f = tl.arange(0, 16) + z = tl.zeros([CAP_P * NB_P], tl.int64) + vals = tl.zeros([CAP_P * NB_P, 16], tl.int64) - 1 + vals = _put(vals, f, 0, z + call) + vals = _put(vals, f, 1, z + layer) + vals = _put(vals, f, 2, tl.load(u_rows_ptr + r, mask=live, other=-1)) + vals = _put(vals, f, 3, ids) + vals = _put(vals, f, 4, s) + vals = _put(vals, f, 5, b) + vals = _put(vals, f, 6, z + BYTES_KIND) + vals = _put(vals, f, 7, tl.load(id_of_slot_ptr + s, mask=ok_s, other=-2)) + vals = _put(vals, f, 8, tl.load(slot_for_id_ptr + layer * E + ids, mask=live, other=-2)) + vals = _put(vals, f, 9, tl.load(usage_ptr + s, mask=ok_s, other=-1)) + vals = _put(vals, f, 10, z + tl.load(step_ptr)) + vals = _put(vals, f, 11, src) + vals = _put(vals, f, 12, first.to(tl.int64) * 4) + vals = _put(vals, f, 13, cnt) + tl.store(row + BYTES_COL, tl.load(row + BYTES_COL) + tl.sum(bad.to(tl.int64), axis=0)) + cursor = _write(ring_ptr, cursor, bad, vals, RING_N) + tl.store(cursor_ptr, cursor) + + +class ExpertVerifier: + """Device counters, the record ring and (``full``) the scratch the byte checks compare against. + + ``begin`` copies the router's ids before ensure rewrites them, ``before_gemm`` runs after the + joins and the demand copy, ``after_gemm`` right after the GEMM; all three enqueue on the current + (compute) stream with fixed shapes and no host sync, so they capture into the decode graph. + ``bind`` sizes the scratch from the banks once and refreshes the slot bank addresses in place + (a rebuild reallocates the banks); counters survive rebuilds and resets. + """ + + def __init__(self, num_layers: int, num_experts: int, device: torch.device, *, mode: str = "full", + rows: int | None = None) -> None: + from freetoken.kernel.fast_index_copy import _skip_fast_index_copy_enabled + + if mode == "full" and _skip_fast_index_copy_enabled(): + logger.warning( + "FREETOKEN_MOE_PREFETCH_VERIFY=1 with FREETOKEN_SKIP_FAST_INDEX_COPY: the host rows cannot be " + "gathered, so only the slot maps are checked (meta)" + ) + mode = "meta" + self.mode = mode + self.num_layers = num_layers + self.num_experts = num_experts + self.device = device + self.rows = ENV.MOE_PREFETCH_VERIFY_ROWS.value if rows is None else int(rows) + if self.rows < 1: + raise ValueError(f"FREETOKEN_MOE_PREFETCH_VERIFY_ROWS={self.rows} must be >= 1") + self.counters = torch.zeros((num_layers, NUM_COLS), dtype=torch.int64, device=device) + self.totals = torch.zeros((num_layers, NUM_COLS), dtype=torch.int64) + self.calls = torch.zeros((num_layers,), dtype=torch.int64, device=device) + self.ring = torch.full((RING, NUM_FIELDS), -1, dtype=torch.int64, device=device) + self.cursor = torch.zeros((1,), dtype=torch.int64, device=device) + self.ids = torch.full((MAX_ENTRIES,), -1, dtype=torch.int32, device=device) + self.slots = torch.full((MAX_ENTRIES,), -1, dtype=torch.int32, device=device) + self.u_ids = torch.zeros((self.rows,), dtype=torch.int32, device=device) + self.u_slots = torch.zeros((self.rows,), dtype=torch.int32, device=device) + self.u_rows = torch.zeros((self.rows,), dtype=torch.int32, device=device) + self.u_num = torch.zeros((1,), dtype=torch.int64, device=device) + self.iota = torch.arange(self.rows, dtype=torch.int32, device=device) + self.bank_names: list[str] = [] + self.scratch: torch.Tensor | None = None + self._printed = 0 + self._top_k = 1 + self._call = None # (layer, entries, this layer's prefetch plan or None) between begin and after_gemm + self._next = None + logger.info( + f"MoE verify {self.mode}: checking every GPU decode layer call's slot maps" + + (f" and slot bytes before and after the GEMM (up to {self.rows} experts a call)" if self.mode == "full" else "") + ) + + def bind(self, cache) -> None: + """Size the scratch for ``cache``'s banks (first call) or point the checks at its new banks.""" + caches = [bank for _, bank in cache.banks] + feats = [math.prod(bank.shape[1:]) * bank.element_size() for bank in caches] + if self.bank_names: + assert list(cache.bank_schema) == self.bank_names and feats == self._feats, "the bank layout changed" + else: + assert all(f % 4 == 0 for f in feats), feats + self.bank_names = list(cache.bank_schema) + self._feats = feats + self.cache_ptrs = torch.zeros((len(feats),), dtype=torch.int64, device=self.device) + if self.mode == "full": + self._alloc_scratch(cache, feats) + self.cache_ptrs.copy_(torch.tensor([bank.data_ptr() for bank in caches], dtype=torch.int64)) + + def _alloc_scratch(self, cache, feats: list[int]) -> None: + rows = self.rows + self.scratch = torch.empty((rows * sum(feats),), dtype=torch.uint8, device=self.device) + offsets = [rows * sum(feats[:b]) for b in range(len(feats))] + self.scratch_views = [ + self.scratch[o : o + rows * f].view(bank.dtype).view(rows, *bank.shape[1:]) + for o, f, (_, bank) in zip(offsets, feats, cache.banks) + ] + self.scratch_ptrs = torch.tensor([self.scratch.data_ptr() + o for o in offsets], dtype=torch.int64, device=self.device) + self.words = torch.tensor([f // 4 for f in feats], dtype=torch.int64, device=self.device) + self.bad_cnt = torch.zeros((rows * len(feats),), dtype=torch.int32, device=self.device) + self.bad_first = torch.full((rows * len(feats),), 0x7FFFFFFF, dtype=torch.int32, device=self.device) + self._chunks = triton.cdiv(max(feats) // 4, CHUNK_WORDS) + logger.info(f"MoE verify scratch: {self.scratch.numel() / 2**20:.1f} MiB for {rows} experts") + + def begin(self, cache, layer_id: int, topk_ids: torch.Tensor, plan: tuple | None) -> None: + """Copy the router's ids of ``layer_id`` (``plan``: this layer's prefetch copy plan, if forked).""" + assert topk_ids.dtype == torch.int32 and topk_ids.is_contiguous() + n = topk_ids.numel() + self._top_k = topk_ids.shape[-1] + if n > MAX_ENTRIES: + self.counters[layer_id, UNCHECKED] += n + self._call = None + return + self.ids[:n].copy_(topk_ids.view(-1)) + self._call = (layer_id, n, plan) + + def _plans(self, cache) -> tuple: + _, n, plan = self._call + pf, nx = plan, self._next + dummy = (self.u_slots, self.u_num) + width = (pf or nx or dummy)[0].numel() + return ( + cache.evict_slots, cache.num_indices, *(pf or dummy), *(nx or dummy), + dict(DEM_W=n, DEM_P=triton.next_power_of_2(n), PF_W=width, PF_P=triton.next_power_of_2(width), + HAS_PF=pf is not None, HAS_NX=nx is not None), + ) + + def before_gemm(self, cache, slots: torch.Tensor, next_plan: tuple | None) -> None: + """Check the maps and (``full``) the bytes of the slots ``slots`` the GEMM is about to read + (``next_plan``: the next layer's prefetch plan, if one was issued).""" + if self._call is None: + return + layer, n, _ = self._call + self._next = next_plan + *plans, plan_kw = self._plans(cache) + full = self.mode == "full" + _plan_kernel[(1,)]( + self.ids, slots, cache.id_of_slot, cache.slot_for_id.view(-1), cache.usage, cache.step, *plans, + self.ring, self.cursor, self.counters, self.calls, + self.u_ids, self.u_slots, self.u_rows, self.u_num, self.slots, + layer, cache.cache_size, + N=n, N_P=triton.next_power_of_2(n), E=self.num_experts, CAP=self.rows, + BYTES=full, RING_N=RING, num_warps=4, **plan_kw, + ) + if not full: + return + self._gather(cache, layer) + self._compare_and_tally(cache, post=False) + + def after_gemm(self, cache) -> None: + """Recheck the maps and (``full``) the bytes of the slots the GEMM just read.""" + if self._call is None: + return + if self.mode == "full": + self._compare_and_tally(cache, post=True) + else: + self._tally(cache, post=True, bytes_=False) + self._call = self._next = None + + def _gather(self, cache, layer: int) -> None: + if cache._copy_fused_ok: + from freetoken.kernel.fast_index_copy import fast_index_copy_multi_jit + + fast_index_copy_multi_jit( + self.scratch_ptrs, cache._copy_src_ptrs[layer], cache._copy_feat_bytes[cache.layer_layout[layer]], self.iota, self.u_ids, self.u_num, + ) + return + from freetoken.kernel import fast_index_copy_jit + + for (per_layer, _), view in zip(cache.banks, self.scratch_views): + fast_index_copy_jit(view, self.iota, per_layer[layer], self.u_ids, self.u_num) + + def _compare_and_tally(self, cache, *, post: bool) -> None: + nb = len(self.bank_names) + _compare_kernel[(self.rows, nb, self._chunks)]( + self.cache_ptrs, self.scratch_ptrs, self.words, self.u_slots, self.u_num, self.bad_cnt, self.bad_first, + NB=nb, BLOCK=CHUNK_WORDS, num_warps=8, + ) + self._tally(cache, post=post, bytes_=True) + + def _tally(self, cache, *, post: bool, bytes_: bool) -> None: + layer, n, _ = self._call + *plans, plan_kw = self._plans(cache) + nb = max(len(self.bank_names), 1) + _tally_kernel[(1,)]( + self.ids, self.slots, cache.id_of_slot, cache.slot_for_id.view(-1), cache.usage, cache.step, *plans, + self.ring, self.cursor, self.counters, self.calls, + self.u_ids, self.u_slots, self.u_rows, self.u_num, + self.bad_cnt if bytes_ else self.u_rows, self.bad_first if bytes_ else self.u_rows, + layer, cache.cache_size, + N=n, N_P=triton.next_power_of_2(n), E=self.num_experts, CAP=self.rows, + CAP_P=triton.next_power_of_2(self.rows), NB=nb, NB_P=triton.next_power_of_2(nb), + BYTES=bytes_, POST=post, BYTES_KIND=BYTES_POST if post else BYTES_PRE, BYTES_COL=POST_BAD if post else PRE_BAD, + RING_N=RING, num_warps=4, **plan_kw, + ) + + # ------------------------------------------------------------------ reporting (host syncs) + + def take_window(self) -> tuple[torch.Tensor, list[list[int]], int]: + """This window's ``[num_layers, NUM_COLS]`` counts (added to ``totals``; the device ones + restart), the records kept since the last call, and how many records there were in all.""" + window = self.counters.cpu() + self.counters.zero_() + self.totals += window + cursor = int(self.cursor.item()) + upto = min(cursor, RING) + records = self.ring[self._printed : upto].cpu().tolist() if upto > self._printed else [] + self._printed = max(self._printed, upto) + return window, records, cursor + + def report_window(self, steps: int) -> list[tuple[str, str]]: + """(level, line) pairs for one MOE_STATS_INTERVAL window; nothing when no call was checked.""" + window, records, cursor = self.take_window() + if not int(window[:, CHECKS].sum()) and not records: + return [] + bad = self._bad(window) or bool(records) + session = self._bad(self.totals) + lines = [("warning" if bad else "info", + f"MoE verify {self.mode} ({steps} decode steps): {self.format_counts(window)}; session bad={session}")] + lines += [("warning", "MoE verify record: " + self.format_record(r)) for r in records] + if bad and cursor > RING: + lines.append(("warning", f"MoE verify: {cursor} bad records in all, the first {RING} are kept")) + return lines + + def report_session(self) -> list[tuple[str, str]]: + """Session counts and every kept record (shutdown).""" + totals = self.totals + self.counters.cpu() + cursor = int(self.cursor.item()) + if not int(totals[:, CHECKS].sum()): + return [] + bad = self._bad(totals) + lines = [("warning" if bad else "info", f"MoE verify {self.mode} (session): {self.format_counts(totals)}")] + records = self.ring[: min(cursor, RING)].cpu().tolist() + lines += [("warning", "MoE verify record: " + self.format_record(r)) for r in records] + if cursor > RING: + lines.append(("warning", f"MoE verify: {cursor} bad records in all, the first {RING} are kept")) + return lines + + @staticmethod + def _bad(counts: torch.Tensor) -> int: + return int(counts[:, [META_BAD, PRE_BAD, POST_BAD, POST_META_BAD]].sum()) + + def format_counts(self, counts: torch.Tensor) -> str: + total = counts.sum(0).tolist() + text = f"{total[CHECKS]} layer calls checked, " + ", ".join( + f"{COL_NAMES[c]}={total[c]}" for c in (META_BAD, PRE_BAD, POST_BAD, POST_META_BAD) + ) + if total[UNCHECKED]: + text += f", unchecked experts={total[UNCHECKED]}" + per_layer = [ + f"L{layer}=" + "/".join(str(row[c]) for c in (META_BAD, PRE_BAD, POST_BAD, POST_META_BAD)) + for layer, row in enumerate(counts.tolist()) + if any(row[c] for c in (META_BAD, PRE_BAD, POST_BAD, POST_META_BAD)) + ] + if per_layer: + text += " (meta/pre/post/post_meta by layer: " + ", ".join(per_layer) + ")" + return text + + def format_record(self, rec: list[int]) -> str: + e = self.num_experts + kind = KINDS[rec[R_KIND]] if 0 <= rec[R_KIND] < len(KINDS) else str(rec[R_KIND]) + owner = rec[R_OWNER] + owner_text = f"L{owner // e}/e{owner % e}" if owner >= 0 else str(owner) + src = [name for bit, name in ((SRC_DEMAND, "demand copy"), (SRC_PREFETCH, "prefetch copy"), + (SRC_NEXT_PREFETCH, "next layer's prefetch plan")) if rec[R_SOURCE] & bit] + row = rec[R_ROW] + text = ( + f"{kind} call={rec[R_CALL]} layer={rec[R_LAYER]} row={row} (token {row // self._top_k}, rank {row % self._top_k}) " + f"expert={rec[R_EXPERT]} slot={rec[R_SLOT]}" + ) + if rec[R_BANK] >= 0: + bank = self.bank_names[rec[R_BANK]] if rec[R_BANK] < len(self.bank_names) else rec[R_BANK] + text += f" bank={bank} first_bad_byte={rec[R_FIRST_BYTE]} bad_words={rec[R_BAD_WORDS]}" + return text + ( + f" | slot now holds {owner_text}, slot_for_id={rec[R_MAPPED]}, usage={rec[R_USAGE]}, lru_step={rec[R_STEP]}, " + f"slot in: {', '.join(src) or 'no copy plan (a resident hit)'}" + ) diff --git a/python/freetoken/scheduler/cache.py b/python/freetoken/scheduler/cache.py index 8f77be9330..a9ef9dca11 100644 --- a/python/freetoken/scheduler/cache.py +++ b/python/freetoken/scheduler/cache.py @@ -29,9 +29,13 @@ def _swa_eviction_interval() -> int: _SWA_RETAIN_GAP = 16 +LOGPROB_READ_SNAPSHOT_DONATE = ("all", "boundary", "none") + + class CacheManager: def __init__(self, num_pages: int, page_size: int, page_table: torch.Tensor, type: str, - linear_state_pool=None, swa_pool=None, sliding_window_size=None): + linear_state_pool=None, swa_pool=None, sliding_window_size=None, + logprob_read_snapshot_donate: str = "all", isolate_logprob_reads: bool = False): # The `_free_slots` follows a page-aligned manner. For example, if page_size = 2, # the `_free_slots` may look like [0, 2, 4, 6, ...], and each slot represents a page. device = page_table.device @@ -44,6 +48,13 @@ def __init__(self, num_pages: int, page_size: int, page_table: torch.Tensor, typ self.swa_pool = swa_pool self.sliding_window_size = sliding_window_size self.is_hybrid = type == "hybrid_radix" + # GDN snapshots a finished logprob read (max_tokens=1 + logprobs) donates: "all" (the + # x64 prefill boundary and the live final state), "boundary" (only the former), "none". + # A burst of such reads otherwise evicts long conversations' snapshots. + if logprob_read_snapshot_donate not in LOGPROB_READ_SNAPSHOT_DONATE: + raise ValueError(f"unknown snapshot donate policy {logprob_read_snapshot_donate!r}") + self.logprob_read_snapshot_donate = logprob_read_snapshot_donate + self.isolate_logprob_reads = isolate_logprob_reads self.is_swa = type == "swa_radix" # swa_paged: this SWA model drives the global-paged swa pool -- true for BOTH the naive # (NaivePrefixCache, no reuse) and radix (SWARadixCache) paths. Gates the swa slot @@ -56,6 +67,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 +102,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 +131,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 +142,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 and len(ids) > 0: + 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) @@ -324,6 +375,11 @@ def cache_req(self, req: Req, *, finished: bool) -> None: req.cache_handle = new_handle self.lock(new_handle) + def _donate_policy(self, req: Req) -> str: + if not req.sampling_params.is_logprob_read: + return "all" + return "none" if self.isolate_logprob_reads else self.logprob_read_snapshot_donate + def _cache_req_hybrid(self, req: Req, *, finished: bool) -> None: """Hybrid (GDN) cache_req: commit KV like radix AND manage the GDN state snapshot. Prefill chunk commit: DONATE the frozen ping-pong slot (the snapshot the forward wrote @@ -344,10 +400,12 @@ def _cache_req_hybrid(self, req: Req, *, finished: bool) -> None: # is this request's dup to free. The frozen slot is consumed either way (taken by # the tree or freed here) and both ping-pong refs are dropped before # _free_req_slots so nothing double-frees. + donate = self._donate_policy(req) free_upto = old_handle.cached_len L = req.mamba_last_track_seqlen if ( - L is not None + donate != "none" + and L is not None and 0 < L <= req.cached_len and align_down(L, self.page_size) == L and req.mamba_ping_pong is not None @@ -367,7 +425,9 @@ def _cache_req_hybrid(self, req: Req, *, finished: bool) -> None: # remain as reuse points). insert_len = align_down(req.cached_len, self.page_size) keep_live = False - if insert_len == req.cached_len and insert_len > 0: + # Without a donated snapshot a node could never be resumed (match_prefix needs + # one), so a non-donating read frees its pages instead of inserting them. + if donate == "all" and insert_len == req.cached_len and insert_len > 0: prefix_len, mamba_exist = self.prefix_cache.insert( req.input_ids[:insert_len], page_indices[:insert_len], req.linear_slot_idx) self.unlock(old_handle) @@ -380,6 +440,10 @@ def _cache_req_hybrid(self, req: Req, *, finished: bool) -> None: return # Prefill chunk commit: donate the frozen snapshot at the tracked ×64 boundary. + # A non-donating read keeps its own pages and frozen slot until the finish frees them. + if self._donate_policy(req) == "none": + req.mamba_last_track_seqlen = None + return L = req.mamba_last_track_seqlen if L is None: return # no ×64 boundary crossed this chunk; req keeps its pages (committed later) @@ -548,6 +612,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 +724,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/config.py b/python/freetoken/scheduler/config.py index b2bbe0aefe..18981135e2 100644 --- a/python/freetoken/scheduler/config.py +++ b/python/freetoken/scheduler/config.py @@ -18,6 +18,10 @@ class SchedulerConfig(EngineConfig): offline_mode: bool = False decode_log_interval: int = 40 special_token_ckpt: bool = False + # GDN snapshots a finished logprob read donates (CacheManager): all | boundary | none + logprob_read_snapshot_donate: str = "all" + # logprob reads prefill alone and bypass the prefix cache, so their logprobs do not depend on other traffic + logprob_read_isolate: bool = False # networking config _unique_suffix: str = field(default_factory=_get_pid_suffix) 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..7536e7a786 100644 --- a/python/freetoken/scheduler/prefill.py +++ b/python/freetoken/scheduler/prefill.py @@ -49,6 +49,8 @@ class PrefillAdder: encoder_cache: EncoderCache | None = None # end a chunk before an image it would cut; only models whose image spans attend in both directions need it keep_images_whole: bool = False + # a logprob read (max_tokens=1 + logprobs) neither resumes from nor feeds the prefix cache + isolate_logprob_reads: bool = False # the whole budget of this pass; token_budget shrinks as requests are admitted pass_budget: int = 0 # SWA-pool tokens charged to reqs admitted so far this pass. Mirrors reserved_size: swa is @@ -71,8 +73,13 @@ 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 + if self.isolate_logprob_reads and req.sampling_params.is_logprob_read: + max_len = 0 # a resumed prefill differs numerically from a cold one + else: + max_len = max(0, req.input_len - replay) if replay else None + mr = self.cache_manager.match_req(req, max_len=max_len) handle = mr.cuda_handle cached_len = handle.cached_len # TODO: better estimate policy @@ -142,6 +149,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 +222,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 +232,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 +252,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 ) @@ -275,6 +288,8 @@ class PrefillManager: decode_manager: DecodeManager encoder_cache: EncoderCache | None = None keep_images_whole: bool = False + # a logprob read prefills in a pass of its own (see PrefillAdder for the cache side) + isolate_logprob_reads: bool = False pending_list: List[PendingReq] = field(default_factory=list) def add_one_req(self, req: UserMsg) -> None: @@ -301,6 +316,7 @@ def schedule_next_batch(self, prefill_budget: int) -> Batch | None: table_manager=self.table_manager, encoder_cache=self.encoder_cache, keep_images_whole=self.keep_images_whole, + isolate_logprob_reads=self.isolate_logprob_reads, ) reqs: List[Req] = [] chunked_list: List[PendingReq] = [] @@ -312,6 +328,10 @@ def schedule_next_batch(self, prefill_budget: int) -> Batch | None: log_cached_tokens = 0 for pending_req in self.pending_list: is_continuation = pending_req.chunked_req is not None + # break, not continue: the slicing below assumes the admitted requests are a prefix + read = self.isolate_logprob_reads and pending_req.sampling_params.is_logprob_read + if read and reqs: + break # a read never shares a pass if req := adder.try_add_one(pending_req): pending_req.chunked_req = None if isinstance(req, ChunkedReq): @@ -334,6 +354,8 @@ def schedule_next_batch(self, prefill_budget: int) -> Batch | None: log_new_tokens += req.extend_len if not is_continuation: log_cached_tokens += req.cache_handle.cached_len + if read: + break # nobody joins a read's pass else: break # We cannot add more requests if len(reqs) == 0: diff --git a/python/freetoken/scheduler/scheduler.py b/python/freetoken/scheduler/scheduler.py index ac1bf322e3..d2fb4c3c8f 100644 --- a/python/freetoken/scheduler/scheduler.py +++ b/python/freetoken/scheduler/scheduler.py @@ -1,6 +1,8 @@ from __future__ import annotations +import math import os +import time from typing import TYPE_CHECKING, List, NamedTuple, NoReturn, Set, Tuple, TypeAlias @@ -31,6 +33,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 @@ -61,6 +64,13 @@ class ForwardInput(NamedTuple): ForwardData: TypeAlias = "Tuple[ForwardInput, ForwardOutput]" +def logprob_candidates(ids: list[int], logprobs: list[float]) -> tuple[list[int], list[float]]: + """A row's top-logprob slice minus what is not a candidate: padding (-1 ids) and the + -inf tail of a row restricted to fewer allowed_token_ids than it asked for.""" + keep = [k for k, (t, lp) in enumerate(zip(ids, logprobs)) if t >= 0 and math.isfinite(lp)] + return [int(ids[k]) for k in keep], [logprobs[k] for k in keep] + + class Scheduler(SchedulerIOMixin): def __init__(self, config: SchedulerConfig): from freetoken.engine import Engine @@ -89,8 +99,18 @@ def __init__(self, config: SchedulerConfig): (g.sliding_window for g in config.model_config.kv_cache_group_specs() if g.is_swa), None, ) or getattr(self.engine.kv_cache, "sliding_window_size", None), + logprob_read_snapshot_donate=getattr(config, "logprob_read_snapshot_donate", "all"), + isolate_logprob_reads=getattr(config, "logprob_read_isolate", False), ) 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, @@ -98,6 +118,7 @@ def __init__(self, config: SchedulerConfig): self.decode_manager, encoder_cache=self.engine.encoder_cache, keep_images_whole=self._bidirectional_mm, + isolate_logprob_reads=getattr(config, "logprob_read_isolate", False), ) # some alias for easy access @@ -111,6 +132,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 +339,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,12 +374,28 @@ 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, row_top_logprobs = logprob_candidates( + outputs.top_ids_cpu[i, :requested_top].tolist(), + 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. - hit_length = not req.can_decode + # 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 @@ -370,6 +413,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, @@ -378,6 +425,16 @@ 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, + starts_in_reasoning=req.sampling_params.starts_in_reasoning, + 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, ) ) @@ -428,16 +485,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 @@ -511,10 +590,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 @@ -526,7 +612,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" ), @@ -541,6 +627,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) @@ -631,6 +722,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). @@ -736,17 +830,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, @@ -867,11 +954,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) @@ -886,6 +990,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..167db5a995 100644 --- a/python/freetoken/server/anthropic_api.py +++ b/python/freetoken/server/anthropic_api.py @@ -40,19 +40,23 @@ GenerationError, GenResult, GenSpec, + PromptAdmitted, ReasoningDelta, 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 +124,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 +167,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 +180,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 +201,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 +220,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 +314,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 +322,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 +336,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 +357,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 +406,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 +433,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 +457,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 +482,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]: @@ -463,14 +503,21 @@ def _stop_block() -> list[str]: block_index += 1 return frames - # message_start (input_tokens filled in by the terminal message_delta). - yield _event(AnthropicStreamEvent( - type="message_start", - message=AnthropicMessagesResponse( - id=f"msg_{uid}", content=[], model=model, - usage=AnthropicUsage(input_tokens=0, output_tokens=0), - ), - )) + started = False + + def _message_start(usage: AnthropicUsage) -> str: + nonlocal started + started = True + return _event(AnthropicStreamEvent( + type="message_start", + message=AnthropicMessagesResponse( + id=f"msg_{uid}", content=[], model=model, usage=usage, + ), + )) + + # message_start waits for the prompt's admission so it can carry the real input/cache + # usage; anything that must go out first (a keepalive, an error) sends it with zeros. + zero_usage = AnthropicUsage(input_tokens=0, output_tokens=0) def _open_tool(name: str | None, ordinal: int | None, stable: bool = True) -> list[str]: nonlocal block_open, tool_args_sent, tool_ordinal, tool_stable @@ -504,6 +551,15 @@ def _tool_args_delta(fragment: str) -> str: events = with_keepalive(events, KEEPALIVE_INTERVAL_S) try: async for ev in events: + if isinstance(ev, PromptAdmitted): + if not started: + yield _message_start( + _anthropic_usage(ev.prompt_tokens, 0, ev.cached_tokens, cache_report) + ) + continue + if not started: + yield _message_start(zero_usage) + if ev is KEEPALIVE: yield _event(AnthropicStreamEvent(type="ping")) @@ -576,14 +632,24 @@ 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 # `data: [DONE]` sentinel (it would be an unknown event to strict clients). + if not started: + yield _message_start(zero_usage) except GenerationError as exc: # Request-side failure (template rejection / over-length prompt). Match the # non-streaming path's classification so Claude Code treats it as a client error # rather than a server fault to retry. + if not started: + yield _message_start(zero_usage) if block_open: for f in _stop_block(): yield f @@ -591,6 +657,8 @@ def _tool_args_delta(fragment: str) -> str: type="error", error=AnthropicError(type="invalid_request_error", message=str(exc)), )) except Exception as exc: # noqa: BLE001 — surface as an Anthropic error event + if not started: + yield _message_start(zero_usage) if block_open: for f in _stop_block(): yield f 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..344b274829 100644 --- a/python/freetoken/server/api_models.py +++ b/python/freetoken/server/api_models.py @@ -74,8 +74,14 @@ class ChatCompletionRequest(BaseModel): stream: bool = False stream_options: StreamOptions | None = None stop: str | list[str] | None = None - presence_penalty: float = 0.0 - frequency_penalty: float = 0.0 + presence_penalty: float = Field(default=0.0, ge=-2.0, le=2.0, allow_inf_nan=False) + frequency_penalty: float = Field(default=0.0, ge=-2.0, le=2.0, allow_inf_nan=False) + # Sampled-token logprobs (OpenAI chat semantics): top_logprobs (0..20) requires + # logprobs=true; entries cover message-content tokens only (no prompt, reasoning or + # tool-call tokens). 400 for parser families whose hidden tokens the segmenter + # cannot tell apart from content (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 @@ -112,14 +137,29 @@ class CompletionRequest(BaseModel): stream: bool = False stream_options: StreamOptions | None = None stop: str | list[str] | None = None - presence_penalty: float = 0.0 - frequency_penalty: float = 0.0 + presence_penalty: float = Field(default=0.0, ge=-2.0, le=2.0, allow_inf_nan=False) + frequency_penalty: float = Field(default=0.0, ge=-2.0, le=2.0, allow_inf_nan=False) ignore_eos: bool = False logprobs: int | None = None echo: bool = False 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 + allowed_token_ids: list[int] | None = None + return_tokens_as_token_ids: bool = False + # 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 +168,34 @@ 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 +# prompts x n per /v1/completions request: every one is queued before any is awaited +MAX_PROMPTS = 256 + + +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..29ba88e1c7 100644 --- a/python/freetoken/server/api_server.py +++ b/python/freetoken/server/api_server.py @@ -2,6 +2,8 @@ import asyncio import contextlib +import hashlib +import hmac import json import os import signal @@ -33,7 +35,8 @@ init_logger, load_generation_sampling, ) -from pydantic import BaseModel +from pydantic import BaseModel, Field +from starlette.datastructures import Headers from .args import ServerArgs from .anthropic_api import register_anthropic_routes @@ -57,6 +60,10 @@ # shutdown is treated as expected — no ERROR log, no "failed" latch. See run_backend_supervisor. _SHUTTING_DOWN = threading.Event() BACKEND_DEATH_EXIT_GRACE_S = 10.0 +# The longest a stream stays silent before it sends an SSE comment. A long prompt prefills for +# minutes before its first token, and clients time an idle body out (Node's fetch after 300 s). +STREAM_KEEPALIVE_SECONDS = 15.0 +_STREAM_KEEPALIVE = b": keep-alive\n\n" def get_global_state() -> FrontendManager: @@ -137,6 +144,13 @@ class FrontendManager: initialized: bool = False ack_map: Dict[int, List[UserReply]] = field(default_factory=dict) event_map: Dict[int, asyncio.Event] = field(default_factory=dict) + # Abort claims live HERE, not in the maps above: a disconnect cancels the response + # task at its innermost await (inside wait_for_ack, whose finally empties both maps + # on the way out), so by the time any abort path runs, a map-based claim would no-op + # exactly when the AbortMsg matters most. Insertion-ordered, FIFO-bounded like the + # scheduler's abort tombstones; a duplicate AbortMsg for an evicted or finished uid + # is safe (the scheduler acks aborts for uids it no longer has). + aborted_uids: Dict[int, None] = field(default_factory=dict) # Stable identity for this serve process. Generated before the backend is ready so every # /health state (loading/ok/error) and /v1/stats can identify the same engine generation. instance_id: str = field(default_factory=lambda: str(uuid.uuid4())) @@ -358,24 +372,56 @@ async def stream_generate(self, uid: int): yield b"data: [DONE]\n\n" logger.debug("Finished streaming response for user %s", uid) - async def stream_with_cancellation(self, generator, request: Request, uid: int): + async def stream_with_cancellation(self, generator, request: Request, uid): + """``uid`` is one request id or a sequence of them (an ``n > 1`` fan-out streams + several generations into one response): a disconnect aborts every one. While the + generator is silent (a long prefill), send an SSE comment every + ``STREAM_KEEPALIVE_SECONDS`` and check whether the client has gone.""" + uids = list(uid) if isinstance(uid, (list, tuple)) else [uid] + finished = False + chunks = generator.__aiter__() + pending: asyncio.Future | None = None try: - async for chunk in generator: - # detect if the client has disconnected + while True: + if pending is None: + pending = asyncio.ensure_future(chunks.__anext__()) + done, _ = await asyncio.wait({pending}, timeout=STREAM_KEEPALIVE_SECONDS) if await request.is_disconnected(): - logger.info("Client disconnected for user %s", uid) + logger.info("Client disconnected for user %s", uids) raise asyncio.CancelledError + if not done: + yield _STREAM_KEEPALIVE + continue + try: + chunk = pending.result() + except StopAsyncIteration: + finished = True + return + pending = None yield chunk - except asyncio.CancelledError: - asyncio.create_task(self.abort_user(uid)) - raise + finally: + if pending is not None and not pending.done(): + pending.cancel() + # finally, not `except CancelledError`: the abort must go out however the + # stream dies -- the server cancelling the response task, a late close + # delivering GeneratorExit, or an exception out of the generator. Only a + # stream that ran to completion leaves the engine with nothing to stop. + if not finished: + for one in uids: + try: + await asyncio.shield(self.abort_user(one)) + except Exception: # noqa: BLE001 + logger.exception("Failed to deliver abort for user %s", one) async def abort_user(self, uid: int): + if uid in self.aborted_uids: + return + self.aborted_uids[uid] = None + while len(self.aborted_uids) > 65_536: + self.aborted_uids.pop(next(iter(self.aborted_uids))) + self.event_map.pop(uid, None) + self.ack_map.pop(uid, None) await asyncio.sleep(0.1) - if uid in self.ack_map: - del self.ack_map[uid] - if uid in self.event_map: - del self.event_map[uid] self.stats.on_abort(uid) logger.warning("Aborting request for user %s", uid) await self.send_one(AbortMsg(uid=uid)) @@ -417,6 +463,39 @@ def install_cors(app: FastAPI, origins_csv: str) -> None: ) +class _ApiKeyMiddleware: + def __init__(self, app, api_key: str) -> None: + self.app = app + self._digest = hashlib.sha256(api_key.encode()).digest() + + def _matches(self, value: str | None) -> bool: + # compare digests so the time taken does not depend on the key length + return value is not None and hmac.compare_digest( + hashlib.sha256(value.encode()).digest(), self._digest + ) + + async def __call__(self, scope, receive, send) -> None: + if scope["type"] not in ("http", "websocket") or scope["path"] == "/health": + return await self.app(scope, receive, send) + headers = Headers(scope=scope) + scheme, _, token = headers.get("authorization", "").partition(" ") + if (scheme.lower() == "bearer" and self._matches(token)) or self._matches( + headers.get("x-api-key") + ): + return await self.app(scope, receive, send) + await JSONResponse({"error": "Unauthorized"}, status_code=401)(scope, receive, send) + + +def install_api_key(app: FastAPI, api_key: str | None) -> None: + """Require ``api_key`` on every route but /health. + + Run it before install_cors, so CORS answers preflights and adds its headers to 401s.""" + if api_key is None: + return + app.add_middleware(_ApiKeyMiddleware, api_key=api_key) + logger.info("API key required on every route except /health") + + app = FastAPI(title="FreeToken API Server", version=__version__, lifespan=lifespan) register_openai_routes(app, get_global_state, lambda: _MODEL_SAMPLING) register_anthropic_routes(app, get_global_state, lambda: _MODEL_SAMPLING) @@ -446,38 +525,57 @@ def _served_model_name() -> str | None: return getattr(cfg, "served_model_name", None) -@app.middleware("http") -async def _record_request_middleware(request: Request, call_next): +class _RecordRequestMiddleware: """Time every generation request into the ring for /v1/requests + /v1/stats p95. Single- model server, so model = served_model_name; stream is inferred from the response media - type. Token counts are P3 (SSE usage arrives after the handler returns) — kept as None.""" - path = request.url.path - if path.startswith(_UNTRACKED_REQUEST_PREFIXES) or not path.startswith( - _TRACKED_REQUEST_PREFIXES - ): - return await call_next(request) - import time as _time - - start = _time.monotonic() - response = await call_next(request) - duration_ms = int((_time.monotonic() - start) * 1000) - ctype = response.headers.get("content-type", "") - request_ring.record_request( - request_ring.RequestRecord( - ts=_time.strftime("%Y-%m-%dT%H:%M:%SZ", _time.gmtime()), - method=request.method, - path=path, - status=response.status_code, - model=_served_model_name(), - duration_ms=duration_ms, - ttft_ms=None, - prompt_tokens=None, - completion_tokens=None, - stream=ctype.startswith("text/event-stream"), - error=None, - ) - ) - return response + type. Token counts are P3 (SSE usage arrives after the handler returns) — kept as None. + + Pure ASGI on purpose: Starlette's ``@app.middleware("http")`` (BaseHTTPMiddleware) hands the + handler a wrapped receive channel that never yields the client's ``http.disconnect`` while a + non-streaming request is still computing, so ``request.is_disconnected()`` stays False and an + abandoned request runs to max_tokens. The recorder is the only such middleware here.""" + + def __init__(self, app): + self.app = app + + async def __call__(self, scope, receive, send): + path = scope.get("path", "") if scope["type"] == "http" else "" + if not path.startswith(_TRACKED_REQUEST_PREFIXES) or path.startswith(_UNTRACKED_REQUEST_PREFIXES): + await self.app(scope, receive, send) + return + import time as _time + + start = _time.monotonic() + + async def send_recording(message): + # the base-class version recorded once the handler returned its response object, + # i.e. at response start, before a streaming body -- keep that timing + if message["type"] == "http.response.start": + ctype = "" + for name, value in message.get("headers", ()): + if name.lower() == b"content-type": + ctype = value.decode("latin-1") + request_ring.record_request( + request_ring.RequestRecord( + ts=_time.strftime("%Y-%m-%dT%H:%M:%SZ", _time.gmtime()), + method=scope["method"], + path=path, + status=message["status"], + model=_served_model_name(), + duration_ms=int((_time.monotonic() - start) * 1000), + ttft_ms=None, + prompt_tokens=None, + completion_tokens=None, + stream=ctype.startswith("text/event-stream"), + error=None, + ) + ) + await send(message) + + await self.app(scope, receive, send_recording) + + +app.add_middleware(_RecordRequestMiddleware) class CacheRebuildRequest(BaseModel): @@ -497,7 +595,7 @@ class CacheRebuildRequest(BaseModel): # is deferred (needs the drain-gate machinery); constraining the Literal makes an # unsupported value fail fast with a 422 at the API layer instead of a generic 503. mode: Literal["if_idle"] = "if_idle" - timeout: float = 300.0 + timeout: float = Field(default=300.0, gt=0, le=3600) async def dispatch_rebuild( @@ -563,8 +661,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 +706,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 +816,30 @@ def _reasoning_geometry(state: Any) -> dict | None: return {"gears": list(gears), "default": default, "kwargs": kwargs} +def kv_pool_geometry(state: Any) -> tuple[int, int]: + """The allocated KV pool as ``(num_pages, page_size)``, most-recent-truth first: the last + rebuild's result, else the running UserReply snapshot, else the load-time ("meta", …) ack. + ``num_pages`` is 0 until one of those lands (still loading / older engine build). + + Split out of cache_geometry so /v1/models can clamp the advertised context to the pool + without recomputing the whole panel payload -- both callers must agree on the resolution + order or they would report different capacities for the same server. + """ + stats = getattr(state, "stats", None) + last = getattr(state, "last_rebuild", None) or {} + pools = getattr(state, "cache_pools", None) or {} + num_pages = int( + last.get("num_pages") + or getattr(stats, "kv_total_pages", 0) + or pools.get("num_pages", 0) + or 0 + ) + page_size = int( + pools.get("page_size", 0) or getattr(getattr(state, "config", None), "page_size", 1) or 1 + ) + return num_pages, page_size + + def cache_geometry(state: Any) -> dict: """Current cache geometry for the desktop cache panel. Each pool size resolves most-recent-truth first: the last rebuild's result, else the running UserReply snapshot @@ -725,11 +855,10 @@ def cache_geometry(state: Any) -> dict: config = state.config last = getattr(state, "last_rebuild", None) or {} pools = getattr(state, "cache_pools", None) or {} - num_pages = int(last.get("num_pages") or tr.kv_total_pages or pools.get("num_pages", 0) or 0) + num_pages, page_size = kv_pool_geometry(state) num_mamba_slots = int( last.get("mamba_slots") or tr.mamba_total_slots or pools.get("num_mamba_slots", 0) or 0 ) - page_size = int(pools.get("page_size", 0) or getattr(config, "page_size", 1) or 1) moe_cache_size = last.get("moe_cache_size") if moe_cache_size is None: moe_cache_size = int(pools.get("moe_cache_size", 0) or 0) or configured_moe_cache_size( @@ -739,11 +868,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 +914,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), @@ -876,7 +1004,7 @@ def _flag_shutdown(signum, frame) -> None: signal.signal(sig, _flag_shutdown) -def _serve_and_run_shell(host: str, port: int) -> None: +def _serve_and_run_shell(host: str, port: int, api_key: str | None) -> None: """Shell mode: serve the API here, and attach the terminal client to it over the loopback. The shell is an ordinary API client (see ``freetoken.shell``), so shell mode is just @@ -909,7 +1037,7 @@ def _serve_and_run_shell(host: str, port: int) -> None: # /health and echoes the same load progress the desktop app polls for. A ^C during that # wait is a stop, not a crash -- exit through the teardown below, not a traceback. with contextlib.suppress(KeyboardInterrupt): - asyncio.run(run_shell(origin, connect_grace=30.0)) + asyncio.run(run_shell(origin, connect_grace=30.0, api_key=api_key)) finally: server.should_exit = True thread.join(timeout=15) @@ -951,6 +1079,7 @@ def run_api_server(config: ServerArgs, start_backend: Callable[[], "Any"], run_s host = config.server_host port = config.server_port + install_api_key(app, config.api_key) # Create/validate FREETOKEN_API_LOG_DIR and start the writer thread up front, so a # bad path is reported at boot rather than silently on the first request. install_cors(app, config.cors_origins) @@ -1037,7 +1166,7 @@ def _on_meta(meta: dict) -> None: ).start() if run_shell: - _serve_and_run_shell(host, port) + _serve_and_run_shell(host, port, config.api_key) return # uvicorn stays on the main thread (signal handling unchanged); ^C reaches the worker group. uvicorn.run(app, host=host, port=port) diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py index ab2fb9b748..6ed30ff073 100644 --- a/python/freetoken/server/args.py +++ b/python/freetoken/server/args.py @@ -3,7 +3,7 @@ import argparse import json import os -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import List, Tuple import torch @@ -38,6 +38,8 @@ def _nvfp4_entry(value: str) -> str: class ServerArgs(SchedulerConfig): server_host: str = "127.0.0.1" server_port: int = 1919 + # kept out of repr: parse_args logs the whole ServerArgs + api_key: str | None = field(default=None, repr=False) num_tokenizer: int = 0 silent_output: bool = False # The terminal shell is attached to this server (ft shell --model / ft serve --shell-mode). @@ -49,6 +51,12 @@ class ServerArgs(SchedulerConfig): # Reasoning parser that splits reasoning from content for OpenAI # responses. None disables it (default for models without a reasoning protocol). reasoning_parser: str | None = None + # Server-wide default thinking mode for reasoning-capable models: "auto" keeps the + # current per-request behavior; "chat" turns thinking off for every request that does + # not choose itself (template kwargs or a protocol-level effort/thinking field), for + # OpenAI-compatible clients that never send template kwargs, like Vercel AI SDK or + # llama-swap; "thinking" turns it on the same way. + default_thinking_mode: str = "auto" # "model": fill unspecified request sampling params from generation_config.json # (temperature/top_k/top_p), like sglang. "none": use framework defaults only. sampling_defaults: str = "model" @@ -59,6 +67,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 +112,6 @@ def backend_create_detokenizer_link(self) -> bool: def frontend_create_tokenizer_link(self) -> bool: return not self.share_tokenizer - @property - def distributed_addr(self) -> str: - return f"tcp://127.0.0.1:{self.server_port + 1}" - def _json_object(text: str) -> dict: try: @@ -131,6 +140,7 @@ def parse_args( from freetoken.attention import validate_attn_backend from freetoken.kvcache import SUPPORTED_CACHE_MANAGER from freetoken.moe import MOE_STRATEGIES + from freetoken.moe.pool_policy import SCOPES def _parse_quant_backend(value: str) -> str: from freetoken.layers.quantization import QuantBackend @@ -150,6 +160,27 @@ def _parse_moe_cache_rate(value: str) -> float: raise argparse.ArgumentTypeError("must be in [0, 1]") return rate + def _parse_mem_size(value: str) -> int: + units = { + "b": 1, "k": 1 << 10, "kb": 1 << 10, "kib": 1 << 10, "m": 1 << 20, "mb": 1 << 20, "mib": 1 << 20, + "g": 1 << 30, "gb": 1 << 30, "gib": 1 << 30, + } + text = value.strip().lower() + number = text.rstrip("abgikm") + unit = text[len(number):] + try: + size = int(float(number) * units[unit]) + except (KeyError, ValueError, OverflowError) as exc: + raise argparse.ArgumentTypeError("must be a size with a unit, e.g. 12GiB or 800MiB") from exc + if size <= 0: + raise argparse.ArgumentTypeError("must be positive") + return size + + def _parse_cache_split(value: str) -> str: + if value == "equal" or (value.startswith("profile:") and len(value) > len("profile:")): + return value + raise argparse.ArgumentTypeError("must be equal or profile:") + def _positive_int(value: str) -> int: try: n = int(value) @@ -159,6 +190,15 @@ def _positive_int(value: str) -> int: raise argparse.ArgumentTypeError("must be >= 1") return n + def _valid_port(value: str) -> int: + try: + n = int(value) + except ValueError as exc: + raise argparse.ArgumentTypeError("must be an integer") from exc + if not 0 <= n <= 65535: + raise argparse.ArgumentTypeError("must be between 0 and 65535") + return n + def _lazy_gpu_arg(value: str) -> tuple[str, ...]: from freetoken.gpu_select import gpu_arg @@ -203,6 +243,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): @@ -272,6 +314,17 @@ def _infer_reasoning_parser(model_path: str) -> str | None: help="Data type for model weights and activations. 'auto' will use FP16 for FP32/FP16 models and BF16 for BF16 models.", ) + parser.add_argument( + "--hf-overrides", + type=_json_object, + default=None, + metavar="JSON", + help="JSON object applied to the checkpoint config the model is built from, as vLLM's " + "--hf-overrides: a nested config section is updated key by key, any other value is " + "replaced whole. A YaRN rope_parameters override extends the servable context to " + "original_max_position_embeddings * factor.", + ) + parser.add_argument( "--tensor-parallel-size", "--tp-size", @@ -354,6 +407,27 @@ def _infer_reasoning_parser(model_path: str) -> str | None: help="The port number for the server to listen on.", ) + parser.add_argument( + "--api-key", + type=str, + default=ServerArgs.api_key, + help="Require this key on every route except /health, sent as " + "'Authorization: Bearer ' or 'x-api-key: '. Defaults to $FREETOKEN_API_KEY; " + "unset serves without authentication.", + ) + + parser.add_argument( + "--dist-port", + type=_valid_port, + dest="distributed_port", + default=None, + help=( + "Port for the internal TP rendezvous store, loopback-only regardless of --host. " + "Defaults to --port + 1; override when that collides with another instance or " + "service on the same host." + ), + ) + parser.add_argument( "--cuda-graph-max-bs", "--graph", @@ -407,6 +481,20 @@ def _infer_reasoning_parser(model_path: str) -> str | None: ), ) + parser.add_argument( + "--decode-interleave-every", + type=int, + default=ServerArgs.decode_interleave_every, + help=( + "Force one decode step after this many consecutive prefill steps. A long " + "prompt is prefilled in chunks, so it occupies that many consecutive " + "scheduler steps; measured on an 8xRTX4090 DSV4 deployment, a 300k-token " + "context left an already-decoding request unscheduled for 267-317 s. " + "Unset keeps the historical prefill-first order (the scheduler's own TODO " + "names this policy)." + ), + ) + parser.add_argument( "--page-size", type=int, @@ -414,6 +502,23 @@ def _infer_reasoning_parser(model_path: str) -> str | None: help="Set the page size for system management.", ) + parser.add_argument( + "--kv-cache-dtype", + dest="kv_quant", + type=str, + default=ServerArgs.kv_quant, + choices=["auto", "bf16", "fp8", "nvfp4"], + help=( + "KV-cache storage format. 'bf16' (default) stores the compute dtype; 'fp8'" + " stores e4m3 codes plus one fp32 scale per (token, kv head), roughly " + "doubling the tokens that fit in the same VRAM. Requires the triton" + " attention backend and a plain paged, hybrid-SWA or QSA sparse KV pool" + " (not MLA/DSA, DSV4, or MiniMax-M3 block-sparse models)." + " 'nvfp4' stores packed E2M1 with block/row scales; supports only" + " paged FULL, hybrid-SWA, or QSA sparse attention with head_dim divisible by 16." + ), + ) + parser.add_argument( "--attention-backend", "--attn", @@ -440,6 +545,27 @@ def _infer_reasoning_parser(model_path: str) -> str | None: "as a GDN-aware radix (cross-request GDN-state prefix reuse); pass 'naive' to opt out.", ) + parser.add_argument( + "--logprob-read-snapshot-donate", + type=str, + default=ServerArgs.logprob_read_snapshot_donate, + choices=["all", "boundary", "none"], + help="Hybrid (GDN) models: which state snapshots a finished logprob read (max_tokens=1 " + "with logprobs, e.g. AnyJev label scoring) donates to the prefix cache. 'all' (default): " + "the x64 prefill boundary and the final state; 'boundary': only the former; 'none': " + "neither, so a burst of reads cannot evict other conversations' snapshots.", + ) + + parser.add_argument( + "--logprob-read-isolate", + action="store_true", + default=ServerArgs.logprob_read_isolate, + help="Make a logprob read's (max_tokens=1 with logprobs, e.g. AnyJev label scoring) " + "logprobs independent of other traffic: it prefills in a pass of its own and neither " + "resumes from nor stores prefix-cache entries. Costs read throughput under concurrency " + "and all prefix reuse for reads; other requests are unchanged.", + ) + parser.add_argument( "--text-model-only", action="store_true", @@ -524,6 +650,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, @@ -556,6 +706,7 @@ def _infer_reasoning_parser(model_path: str) -> str | None: "qwen3_coder", "mistral", "deepseekv32", + "deepseekv41", "gemma4", "glm47", "minimax", @@ -583,14 +734,28 @@ 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." + "How the routed experts are served. 'auto' resolves a MoE model to 'offload', and to " + "resident 'fused' experts on unified-memory GPUs (GB10 / DGX Spark)." ), ) @@ -625,6 +790,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, @@ -660,6 +846,12 @@ def _infer_reasoning_parser(model_path: str) -> str | None: default=ServerArgs.moe_cache_rate, help="The fraction of all MoE experts to keep in GPU cache.", ) + moe_cache_group.add_argument( + "--moe-cache-mem", + type=_parse_mem_size, + default=ServerArgs.moe_cache_mem, + help="GPU memory for the MoE expert slot cache, e.g. 12GiB or 800MiB (G, GB and GiB all mean GiB); the slot count is this over the bytes per slot.", + ) moe_cache_group.add_argument( "--moe-cache-auto", action="store_true", @@ -674,14 +866,67 @@ def _infer_reasoning_parser(model_path: str) -> str | None: "--kv-reserve-tokens", type=int, default=ServerArgs.kv_reserve_tokens, - help="KV-cache token floor reserved before --moe-cache-auto fills experts.", + help=( + "Usable KV-cache token floor reserved before --moe-cache-auto fills experts " + "(the internal dummy page is additional)." + ), + ) + + parser.add_argument( + "--moe-cache-scope", + choices=SCOPES, + default=ServerArgs.moe_cache_scope, + help=( + "How MoE layers share the GPU expert slot cache: unified = one LRU pool per expert quant " + "(all layers when they share one quant format); per-layer = each layer has its own pool. " + "Pools are sized by --moe-cache-split." + ), + ) + parser.add_argument( + "--moe-cache-split", + type=_parse_cache_split, + default=ServerArgs.moe_cache_split, + help=( + "How pools share the slots: equal (by layer count), or profile: = by expert use from a " + "profile file. With --moe-strategy offload, a pool smaller than the expert count limits the " + "decode batch to its slots / top_k." + ), + ) + + parser.add_argument( + "--kv-host-pages", + type=int, + default=ServerArgs.kv_host_pages, + help=( + "Extra KV pages mirrored to pinned host RAM (QSA models): the GPU pool becomes " + "an LRU cache over the logical page space, extending context past VRAM capacity. " + "0 = off." + ), ) parser.add_argument( "--moe-cache-policy", default=ServerArgs.moe_cache_policy, - choices=["lru"], - help="The unified MoE cache eviction policy.", + choices=["lru", "kd", "kdfb", "rule"], + help=( + "The unified MoE cache eviction policy. lru (default); kd: evict the expert passed " + "over for the most tokens, then the one whose layer comes up last; kdfb: kd plus a " + "decayed per-expert use count; rule: kdfb plus router near misses counted as recent " + "(margin FREETOKEN_MOE_NEAR_MISS_THR logits). GPU decode only." + ), + ) + + parser.add_argument( + "--moe-collect-stats", + action="store_true", + default=ServerArgs.moe_collect_stats, + help=( + "Log MoE expert-cache miss rate and routing skew during decode. The counters " + "are captured into the decode CUDA graph, so this can only be chosen at startup. " + "Measured cost is below noise (46.1 vs 45.8 tok/s median on Ornith-35B-A3B " + "IQ3_S), but it stays off by default since it is a diagnostic and the readout " + "costs a host sync." + ), ) parser.add_argument( @@ -706,7 +951,7 @@ def _infer_reasoning_parser(model_path: str) -> str | None: "only, where CUDA pinned memory is capped: it locks just enough head+tail layers " "for the banks over the pin budget. Any value, 'auto' included, commits to CPU " "decode before the model is built, so the expert format must have a CPU executor " - "path (bf16, nvfp4, mxfp4); do not pass it on Linux. Unset = every layer on the " + "path (bf16, fp8_block, nvfp4, mxfp4); do not pass it on Linux. Unset = every layer on the " "GPU; a boot whose banks exceed a known pin budget stops and asks for this flag." ), ) @@ -719,11 +964,23 @@ def _infer_reasoning_parser(model_path: str) -> str | None: "For --moe-strategy hybrid: max experts fetched over PCIe per (layer, decode " "step); the rest of that step's misses are computed on the CPU, overlapped. " "-1 (default) = auto: fetch the benched pcie/cpu bandwidth fraction of each " - "step's misses (perfect overlap; needs an `ft bench bw` profile, else 1). " + "step's misses (perfect overlap; needs the `ft bench bw` row of the layer's kind, " + "from `ft bench bw --model `, else 1). " "0 = never fetch (all misses on CPU); large = behaves like plain offload." ), ) + 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", @@ -736,6 +993,17 @@ def _infer_reasoning_parser(model_path: str) -> str | None: ), ) + parser.add_argument( + "--moe-prefill-demand-tokens", + type=int, + default=ServerArgs.moe_prefill_demand_tokens, + help=( + "A prefill chunk of at most this many tokens loads only the experts it routes to, " + "through the expert cache as decode does, instead of every expert of each layer. " + "0 streams whole layers for every chunk." + ), + ) + parser.add_argument( "--enable-special-token-ckpt", action="store_true", @@ -808,6 +1076,16 @@ def _infer_reasoning_parser(model_path: str) -> str | None: if entry: kwargs["quant_backend"] = entry + # an empty flag usually means an unset shell variable; refuse it rather than serve open + if kwargs["api_key"] == "": + parser.error("--api-key is empty; omit it to serve without authentication") + from freetoken.launch import resolve_api_key + + kwargs["api_key"] = resolve_api_key(kwargs["api_key"]) + + if kwargs["distributed_port"] is None: + kwargs["distributed_port"] = kwargs["server_port"] + 1 + if kwargs["model_path"].startswith("~"): kwargs["model_path"] = os.path.expanduser(kwargs["model_path"]) @@ -840,21 +1118,32 @@ def _infer_reasoning_parser(model_path: str) -> str | None: _no_cache_flag = ( kwargs["moe_cache_size"] == 0 and not kwargs["moe_cache_auto"] + and kwargs["moe_cache_mem"] is None and (kwargs["moe_cache_rate"] is None or kwargs["moe_cache_rate"] == 0) ) if is_offload_moe_strategy(kwargs["moe_strategy"]) and _no_cache_flag: kwargs["moe_cache_auto"] = True - if kwargs["model_source"] == "modelscope": - model_path = kwargs["model_path"] - if not os.path.isdir(model_path): + # Resolve a hub repo id ("org/model") to a local checkpoint directory. This has to + # happen here, after served_model_name and the parser cascade have read the repo id + # (they want "DeepSeek-V4-Flash-0731", not a snapshot hash) and before anything opens + # the checkpoint: config parsing reads model-specific files straight off disk, so a + # bare repo id reaches `open()` as a relative path and dies with a FileNotFoundError. + model_path = kwargs["model_path"] + if not os.path.isdir(model_path): + if kwargs["model_source"] == "modelscope": from modelscope import snapshot_download ignore_patterns = [] if kwargs["use_dummy_weight"]: ignore_patterns = ["*.bin", "*.safetensors", "*.pt", "*.ckpt"] - model_path = snapshot_download(model_path, ignore_patterns=ignore_patterns) - kwargs["model_path"] = model_path + kwargs["model_path"] = snapshot_download(model_path, ignore_patterns=ignore_patterns) + else: + from freetoken.utils import download_hf_checkpoint + + kwargs["model_path"] = download_hf_checkpoint( + model_path, dummy_weight=kwargs["use_dummy_weight"] + ) del kwargs["model_source"] # "auto" (or an unspecified dtype) resolves to the checkpoint's dtype. Multimodal / @@ -892,6 +1181,7 @@ def _infer_reasoning_parser(model_path: str) -> str | None: image_max_tokens=image_max_tokens, processor_kwargs=kwargs.pop("mm_processor_kwargs") or {}, ) + kwargs["hf_overrides"] = kwargs["hf_overrides"] or {} result = ServerArgs(**kwargs) logger.info(f"Parsed arguments:\n{result}") return result, run_shell 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..0d2906c84c 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", "", ] @@ -259,6 +261,11 @@ class BaseFormatDetector(ABC): # bot_token is a parse trigger, not a uniqueness claim. toolcall_opener: str | None = None + # True when detect_and_parse replays the streaming state machine: its normal_text + # is then already every piece of text the stream emits (before, between and after + # calls), so parse_non_stream must return it as is. + one_shot_replays_stream = False + def __init__(self): # Streaming state management # Buffer for accumulating incomplete patterns that arrive across multiple streaming chunks @@ -347,17 +354,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 +518,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: @@ -655,6 +783,12 @@ def parse_streaming_increment(self, new_text: str, tools: List[Tool]) -> Streami +# A bare invoke opener's name: bounded so a literal "]{1,128})>\s*") +_PS_BARE_NAME_PARTIAL_RE = re.compile(r"[^\s<>]{0,128}") + + class InvokeParamStreamMixin: """Value-level streaming for invoke/parameter block formats (qwen3_coder, minimax — the same shape vLLM's ParserEngine covers): text outside blocks @@ -670,6 +804,13 @@ class InvokeParamStreamMixin: _ps_trim_single True: at most ONE leading/trailing trim char (qwen3_coder) _ps_missing_type schema type assumed for undeclared params ("string" streams, "loose" buffers and loose-parses — per-family legacy typing) + _ps_trim_after_block True: whitespace after a closed block is held, then + dropped if another block or the end follows, kept if + text does -- the same for any chunking (qwen3_coder) + _ps_bare_invoke True: an invoke without the wrapper block is a call when it + names a declared tool and its body starts (qwen3_coder) + _ps_literal_open True: a block opener whose next non-whitespace text cannot + start an invoke or the block close is text (qwen3_coder) """ _ps_outer_open: str = "" @@ -677,9 +818,15 @@ class InvokeParamStreamMixin: _ps_trim: str = "\n" _ps_trim_single: bool = False _ps_missing_type: str = "string" + _ps_trim_after_block: bool = False + _ps_bare_invoke: bool = False + _ps_literal_open: bool = False def _ps_reset(self) -> None: self._ps_mode = "idle" + self._ps_after_block = False + self._ps_bare = False + self._ps_after_bare = False self._ps_key = "" self._ps_lead = "{" self._ps_emitted_any = False @@ -706,6 +853,58 @@ def _ps_trim_trailing(self, text: str) -> str: return text[:-1] if text[-1:] and text[-1:] in self._ps_trim else text return text.rstrip(self._ps_trim) + def _ps_bare_status(self, text: str) -> str: + """Classify the bare invoke opener ``text`` starts with: "call", "partial" + (may still become one) or "literal". A declared name is required -- not + the unknown-tool forwarding a wrapped block gets -- and the body must + start, so the format quoted in prose or code stays text.""" + rest = text[len(self._ps_invoke_open_prefix):] + m = _PS_BARE_NAME_RE.match(rest) + if m is None: + return "partial" if _PS_BARE_NAME_PARTIAL_RE.fullmatch(rest) else "literal" + if m.group(1) not in self._tool_indices: + return "literal" + body = rest[m.end():] + for tok in (self._ps_param_open_prefix, self._ps_invoke_close): + if body.startswith(tok): + return "call" + if tok.startswith(body): + return "partial" + return "literal" + + def _ps_find_bare(self, buf: str, limit: int) -> Tuple[int, str]: + """First bare invoke opener before ``limit`` that is not literal.""" + i = buf.find(self._ps_invoke_open_prefix) + while i != -1 and i < limit: + status = self._ps_bare_status(buf[i:]) + if status != "literal": + return i, status + i = buf.find(self._ps_invoke_open_prefix, i + 1) + return -1, "" + + def _ps_opens_block(self, text: str) -> str: + """Whether ``text`` starts a block: "yes", "no" or "maybe" (too short to tell).""" + if text.startswith(self._ps_outer_open): + if not self._ps_literal_open: + return "yes" + body = text[len(self._ps_outer_open):].lstrip() + # A repeated opener is a stutter: the block ignores it, as before. + for tok in (self._ps_invoke_open_prefix, self._ps_outer_close, self._ps_outer_open): + if body.startswith(tok): + return "yes" + if tok.startswith(body): + return "maybe" + return "no" + if self._ps_outer_open.startswith(text): + return "maybe" + if self._ps_bare_invoke: + prefix = self._ps_invoke_open_prefix + if text.startswith(prefix): + return {"call": "yes", "partial": "maybe"}.get(self._ps_bare_status(text), "no") + if prefix.startswith(text): + return "maybe" + return "no" + def parse_streaming_increment(self, new_text: str, tools: List[Tool]) -> StreamingParseResult: self._buffer += new_text if not hasattr(self, "_tool_indices"): @@ -743,9 +942,42 @@ def _update_prev() -> None: if mode == "idle": if calls: break # text after a call defers to the next step (wire order) + if self._ps_after_block: + rest = buf.lstrip() + if self._ps_after_bare: + if rest.startswith(self._ps_outer_close): + # The model dropped the opener but kept the closer. + self._ps_after_bare = False + self._buffer = rest[len(self._ps_outer_close):] + continue + if self._ps_outer_close.startswith(rest): + break + self._ps_after_bare = False + opens = self._ps_opens_block(rest) + if opens == "maybe": + break # a separator or a text gap: undecided until more arrives + self._ps_after_block = False + if opens == "yes": + buf = self._buffer = rest pos = buf.find(self._ps_outer_open) + if self._ps_bare_invoke: + bare, status = self._ps_find_bare(buf, len(buf) if pos == -1 else pos) + if bare > 0: + normal_parts.append(buf[:bare]) + self._buffer = buf[bare:] + continue + if bare == 0: + if status == "partial": + break + self._ps_bare = True + self._ps_mode = "block" # parses the invoke at the buffer head + continue if pos == -1: hold = self._ends_with_partial_token(buf, self._ps_outer_open) + if self._ps_bare_invoke: + hold = max( + hold, self._ends_with_partial_token(buf, self._ps_invoke_open_prefix) + ) release = buf[: len(buf) - hold] if hold else buf if release: normal_parts.append(release) @@ -755,7 +987,13 @@ def _update_prev() -> None: normal_parts.append(buf[:pos]) self._buffer = buf[pos:] continue + opens = self._ps_opens_block(buf) + if opens == "maybe": + break self._buffer = buf[len(self._ps_outer_open):] + if opens == "no": + normal_parts.append(self._ps_outer_open) # the tag named in prose + continue self._ps_mode = "block" continue @@ -765,6 +1003,7 @@ def _update_prev() -> None: if close != -1 and (inv == -1 or close < inv): self._buffer = buf[close + len(self._ps_outer_close):] self._ps_mode = "idle" + self._ps_after_block = self._ps_trim_after_block continue if inv != -1: m = self._ps_invoke_open_re.search(buf, inv) @@ -816,7 +1055,14 @@ def _update_prev() -> None: while len(self.streamed_args_for_tool) <= self.current_tool_id: self.streamed_args_for_tool.append("") self._buffer = buf[e + len(self._ps_invoke_close):] - self._ps_mode = "block" + if self._ps_bare: + # A bare invoke is its own block: back to text after it. + self._ps_bare = False + self._ps_mode = "idle" + self._ps_after_block = self._ps_trim_after_block + self._ps_after_bare = True + else: + self._ps_mode = "block" continue if p != -1: m = self._ps_param_open_re.search(buf, p) @@ -897,9 +1143,52 @@ def _update_prev() -> None: return StreamingParseResult(normal_text="".join(normal_parts), calls=calls) + def _ps_finalize_open_call(self) -> List[ToolCallItem]: + """The generation ended (max_tokens) inside an open call: close its arguments the + way the missing closers would have, so the streamed fragments end as valid JSON and + the one-shot replay returns the same arguments. A held suffix that could only grow + into the parameter closer is markup and is dropped, the value before it is kept. + Outside an open call (a cut header, a skipped unknown tool) nothing changes.""" + mode = getattr(self, "_ps_mode", "idle") + if mode not in ("invoke", "pstr", "pbuf"): + return [] + buf = self._buffer + ledger = self.streamed_args_for_tool[self.current_tool_id] + held = self._ends_with_partial_token(buf, self._ps_param_close) if mode != "invoke" else 0 + value = buf[: len(buf) - held] + if mode == "pstr": + if not self._ps_lead_trimmed: + value = self._ps_trim_leading(value) + frag = self._json_escape_chunk(self._ps_trim_trailing(value)) + '"}' + elif mode == "pbuf" and held: + # Only the closer was cut: the typed value itself is complete. + raw = self._ps_trim_trailing(self._ps_trim_leading(value)) + converted = self._ps_convert_value(self._ps_key, raw) + frag = ( + self._ps_lead + json.dumps(self._ps_key, ensure_ascii=False) + ":" + + json.dumps(converted, ensure_ascii=False) + "}" + ) + else: + frag = "}" if ledger else "{}" + index = self.current_tool_id + self.streamed_args_for_tool[index] += frag + try: + parsed = json.loads(self.streamed_args_for_tool[index]) + except (json.JSONDecodeError, ValueError): + parsed = {} + self.prev_tool_call_arr[index]["arguments"] = parsed if isinstance(parsed, dict) else {} + self.current_tool_id += 1 + while len(self.streamed_args_for_tool) <= self.current_tool_id: + self.streamed_args_for_tool.append("") + self._buffer = "" + self._ps_reset() + return [ToolCallItem(tool_index=index, name=None, parameters=frag)] + def finish_streaming(self) -> str: residual, self._buffer = self._buffer, "" mode = getattr(self, "_ps_mode", "idle") + if getattr(self, "_ps_after_block", False) and not residual.strip(): + residual = "" self._ps_reset() if mode != "idle" or (self._ps_outer_open and self._ps_outer_open in residual): return "" @@ -1563,37 +1852,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 +1984,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: @@ -1885,9 +2182,13 @@ def finish_streaming(self) -> str: class Qwen3CoderDetector(InvokeParamStreamMixin, BaseFormatDetector): toolcall_opener = "" + one_shot_replays_stream = True _ps_trim = "\n" _ps_trim_single = True _ps_missing_type = "string" + _ps_trim_after_block = True + _ps_bare_invoke = True + _ps_literal_open = True """ Detector for Qwen3-Coder XML-style function call format. @@ -1921,8 +2222,6 @@ def __init__(self): self.tool_call_separator = "\n" # Regex patterns - self.tool_call_block_regex = re.compile(r"(.*?)", re.DOTALL) - self.function_regex = re.compile(r"||(?=)|$)", re.DOTALL ) @@ -1932,7 +2231,8 @@ def __init__(self): self._ps_outer_open = "" self._ps_outer_close = "" self._ps_invoke_open_prefix = "]*)>") + # No "<" in a name: an opener cut off before its ">" must not run into . + self._ps_invoke_open_re = re.compile(r"]*)>") self._ps_invoke_close = "" self._ps_param_open_prefix = "]*)>") @@ -1942,44 +2242,10 @@ def __init__(self): def has_tool_call(self, text: str) -> bool: return " Optional[ToolCallItem]: - """Parse a single ... block into a ToolCallItem.""" - try: - end_index = function_str.index(">") - except ValueError: - return None - - func_name = function_str[:end_index].strip() - tool_indices = self._get_tool_indices(tools) - if func_name not in tool_indices and not _should_forward_unknown_tool(func_name): - logger.warning(f"Model attempted to call undefined function: {func_name}") - return None - - parameters_text = function_str[end_index + 1 :] - param_config = self._get_param_config(func_name, tools) - param_dict = {} - - for match in self.parameter_regex.findall(parameters_text): - try: - idx = match.index(">") - except ValueError: - continue - param_name = match[:idx].strip() - param_value = match[idx + 1 :] - # Strip leading/trailing newlines from value - if param_value.startswith("\n"): - param_value = param_value[1:] - if param_value.endswith("\n"): - param_value = param_value[:-1] - - param_dict[param_name] = self._convert_param_value(param_value, param_name, param_config, func_name) - - return ToolCallItem( - tool_index=tool_indices.get(func_name, 0), - name=func_name, - parameters=json.dumps(param_dict, ensure_ascii=False), - ) + def finalize_streaming(self) -> List[ToolCallItem]: + """End-of-stream hook (FunctionCallParser.finalize_stream): close a call that + max_tokens cut off mid-arguments.""" + return self._ps_finalize_open_call() def _build_partial_arguments_json(self, func_name: str, partial_body: str, tools: List[Tool]) -> Optional[str]: """Build the current argument JSON from a partial XML tool-call body.""" @@ -2024,28 +2290,41 @@ def _build_partial_arguments_json(self, func_name: str, partial_body: str, tools return json.dumps(param_dict, ensure_ascii=False) def detect_and_parse(self, text: str, tools: List[Tool]) -> StreamingParseResult: - idx = text.find(self.bot_token) - normal_text = text[:idx].strip() if idx != -1 else text - - if " 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 +2896,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 +2918,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 +2929,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 +3801,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 +3822,7 @@ class FunctionCallParser: ToolCallParserEnum: Dict[str, Type[BaseFormatDetector]] = { "deepseekv32": DeepSeekV32Detector, + "deepseekv41": DeepSeekV41Detector, "gemma4": Gemma4Detector, "gpt-oss": GptOssDetector, "gpt_oss": GptOssDetector, @@ -3581,6 +3881,8 @@ def parse_non_stream(self, full_text: str) -> StreamingParseResult: if not self.tools: return StreamingParseResult(normal_text=full_text, calls=[]) parsed_result = self.detector.detect_and_parse(full_text, self.tools) + if self.detector.one_shot_replays_stream: + return parsed_result tool_call_list = parsed_result.calls if tool_call_list: # Keep text after the LAST tool block: the streaming path emits it, so diff --git a/python/freetoken/server/generation.py b/python/freetoken/server/generation.py index 91056e8d1f..4fe7d13580 100644 --- a/python/freetoken/server/generation.py +++ b/python/freetoken/server/generation.py @@ -14,7 +14,9 @@ from __future__ import annotations import asyncio +import dataclasses import json +import math import time from collections.abc import AsyncIterator from dataclasses import dataclass, field @@ -24,7 +26,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 @@ -35,14 +38,266 @@ _TemplateError = () from .function_call_parser import FunctionCallParser, TOOLS_TAG_LIST, ToolCallItem +from .logprobs import ( + LogprobMarkers, + LogprobSegmenter, + resolve_logprob_markers, + untrimmed, +) from .reasoning_parser import ( DSV4_SPECIAL_TOKENS, + THINK_END_TOKEN, + THINK_START_TOKEN, ReasoningParser, build_reasoning_parser, strip_special_tokens, ) +# 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 @@ -65,6 +320,16 @@ def __init__(self, message: str, code: str | None = None): # so generation/parsing logic lives in exactly one place and no adapter has to # re-parse a serialized OpenAI stream. # --------------------------------------------------------------------------- # +@dataclass +class PromptAdmitted: + """The scheduler admitted the prompt: its token count and prefix-cache hit are known + before any output. Sent once, ahead of the first delta, so a protocol that reports + input usage up front (Anthropic's message_start) need not wait for GenDone.""" + + prompt_tokens: int + cached_tokens: int = 0 + + @dataclass class ReasoningDelta: text: str @@ -73,6 +338,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 +376,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,9 +403,14 @@ class GenDone: completion_tokens: int matched_stop: str | None = None cached_tokens: int = 0 + reasoning_tokens: int = 0 + timings: GenTimings = field(default_factory=GenTimings) + # Content logprob entries no content delta carried (text held to the end, or + # dropped as a trailing separator); None when there are none. + logprobs: list[dict] | None = None -GenEvent = ReasoningDelta | ContentDelta | ToolCallStart | ToolCallArgsDelta | ToolCallsDelta | GenDone +GenEvent = PromptAdmitted | ReasoningDelta | ContentDelta | ToolCallStart | ToolCallArgsDelta | ToolCallsDelta | GenDone @dataclass @@ -129,6 +423,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 logprob entries of the message-content tokens (empty when the request did + # not ask). + logprobs: list[dict] = field(default_factory=list) @dataclass @@ -142,6 +441,10 @@ 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 + # Prompt ids counted by the frontend pre-check, handed to the worker so it does not + # render and tokenize again (also reused by every choice of an n > 1 request). + prompt_ids: Any = field(default=None, repr=False, compare=False) @property def parse_tools(self) -> bool: @@ -155,6 +458,11 @@ def parse_tools(self) -> bool: # --max-output-tokens (the Responses adapter passes that through); clamped to the remaining # context by the scheduler regardless. DEFAULT_MAX_OUTPUT_TOKENS = 32768 +# vLLM's MAX_NUM_ALLOWED_TOKEN_IDS: bounds the per-row mask upload +MAX_ALLOWED_TOKEN_IDS = 1024 +# The int32 the engine holds ids in; also keeps them within what msgpack carries. The API +# narrows this to len(tokenizer) when it has a frontend tokenizer, the tokenizer worker always. +MAX_TOKEN_ID = 2**31 - 1 def resolve_sampling( @@ -167,9 +475,20 @@ 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, + allowed_token_ids: list[int] | 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 +500,101 @@ def pick(value, key, framework): # non-positive value is a client error. if max_tokens is not None and max_tokens < 1: raise ValueError(f"max_tokens must be at least 1, got {max_tokens}") - return SamplingParams( + resolved_temperature = pick(temperature, "temperature", 0.0) + if not math.isfinite(resolved_temperature) or resolved_temperature < 0: + raise ValueError(f"temperature must be a finite number >= 0, got {resolved_temperature}") + resolved_top_p = pick(top_p, "top_p", 1.0) + if not math.isfinite(resolved_top_p) or not 0 < resolved_top_p <= 1: + raise ValueError(f"top_p must be in (0, 1], got {resolved_top_p}") + resolved_top_k = pick(top_k, "top_k", -1) + # Any top_k <= 0 means "no top-k filter", as in vLLM and the other OpenAI-compatible + # servers: Ollama-style clients (Open WebUI) and some SDKs send top_k=0 for "off", and a + # checkpoint's generation_config.json can carry it too. Map it to the engine's -1. + if resolved_top_k <= 0: + resolved_top_k = -1 + resolved_min_p = float(pick(min_p, "min_p", 0.0)) + resolved_rep = float(pick(repetition_penalty, "repetition_penalty", 1.0)) + if not 0.0 <= resolved_min_p <= 1.0: + raise ValueError(f"min_p must be in [0, 1], got {resolved_min_p}") + for name, value in (("presence_penalty", presence_penalty), ("frequency_penalty", frequency_penalty)): + if not -2.0 <= float(value) <= 2.0: + raise ValueError(f"{name} must be in [-2, 2], got {value}") + # not `<= 0`: NaN compares False and would poison every seen token's logit + if not (math.isfinite(resolved_rep) and resolved_rep > 0.0): + raise ValueError(f"repetition_penalty must be a positive finite number, got {resolved_rep}") + if min_tokens < 0: + raise ValueError(f"min_tokens must be >= 0, got {min_tokens}") + bias: list[list[float]] | None = None + if logit_bias: + bias = [] + for key, value in logit_bias.items(): + try: + tid = int(key) + except (TypeError, ValueError): + raise ValueError(f"logit_bias keys must be token ids, got {key!r}") from None + if not 0 <= tid <= MAX_TOKEN_ID: + raise ValueError(f"logit_bias token id must be in [0, {MAX_TOKEN_ID}], 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 not all(0 <= t <= MAX_TOKEN_ID for t in ids): + raise ValueError(f"stop_token_ids must be token ids in [0, {MAX_TOKEN_ID}]") + allowed = None + if allowed_token_ids is not None: + if not 1 <= len(allowed_token_ids) <= MAX_ALLOWED_TOKEN_IDS: + raise ValueError( + f"allowed_token_ids must hold 1 to {MAX_ALLOWED_TOKEN_IDS} ids, got {len(allowed_token_ids)}" + ) + allowed = sorted({int(t) for t in allowed_token_ids}) + if not 0 <= allowed[0] <= allowed[-1] <= MAX_TOKEN_ID: + raise ValueError(f"allowed_token_ids must be token ids in [0, {MAX_TOKEN_ID}]") + params = SamplingParams( ignore_eos=ignore_eos, max_tokens=default_max_tokens if max_tokens is None else max_tokens, - temperature=pick(temperature, "temperature", 0.0), - top_k=pick(top_k, "top_k", -1), - top_p=pick(top_p, "top_p", 1.0), + temperature=resolved_temperature, + top_k=resolved_top_k, + top_p=resolved_top_p, stop_strs=[s for s in stop_list if s], # drop empty strings (would match everything) + 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), + allowed_token_ids=allowed, ) + 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]: @@ -290,21 +689,112 @@ def split_tool_lists( # --------------------------------------------------------------------------- # async def submit_generation(spec: GenSpec, state: Any) -> int: """Enqueue one generation from a GenSpec; return its uid. Every protocol adapter - calls this — it takes the neutral spec, not a wire request type.""" + calls this — it takes the neutral spec, not a wire request type. Raises + GenerationError (a 400 in every adapter) when the prompt cannot be rendered or does + not fit the context, before anything is committed to the client.""" refs = collect_image_refs(spec.messages) images = await _resolve_images(refs, state) if refs else None - uid = state.new_user() - await state.send_one( - TokenizeMsg( - uid=uid, - text=spec.messages, - sampling_params=spec.sampling_params, - chat_template_kwargs=spec.chat_template_kwargs, - tools=spec.template_tools, - images=images, - ) + msg = TokenizeMsg( + uid=0, + text=spec.messages, + sampling_params=spec.sampling_params, + chat_template_kwargs=spec.chat_template_kwargs, + tools=spec.template_tools, + images=images, + inline_system_policy=spec.inline_system_policy, ) - return uid + if spec.prompt_ids is None: + spec.prompt_ids = await frontend_prompt_ids(msg, state) + msg.uid = state.new_user() + msg.input_ids = spec.prompt_ids + await state.send_one(msg) + return msg.uid + + +def model_context_length(state: Any) -> int | None: + """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 + if value <= 0: + return None + try: + # Local import: api_server imports the protocol modules, which import this one. + 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 + + +def check_context_fits(prompt_tokens: int, state: Any) -> None: + """Raise the scheduler's context_length_exceeded error up front when the prompt leaves + no room for one output token. A max_tokens that does not fit is not an error: the + scheduler clamps it to the room left, as agents that always send a large one expect.""" + limit = model_context_length(state) + if limit is not None and prompt_tokens >= limit: + # "prompt is too long: N tokens > M" is the phrasing Claude Code and OpenClaw match + # on (the scheduler sends the same text). + raise GenerationError( + f"prompt is too long: {prompt_tokens} tokens > {limit} maximum (prompt + generation); " + f"shorten the prompt or increase the KV cache budget", + "context_length_exceeded", + ) + + +async def frontend_prompt_ids(msg: TokenizeMsg, state: Any) -> Any | None: + """Render and tokenize ``msg`` frontend-side and check it against the context, so a + request that cannot run is a 400 *before* an adapter commits an SSE stream (once headers + go out the failure can only ride in-stream, where some agents show nothing but "empty + response", and clients that compact on "prompt is too long" never see it). + + Returns the prompt ids for the worker to serve as they are (no second tokenization), or + None when there is nothing to hand over: no frontend tokenizer (minimal embeddings), one + that fails to *initialize* (a server fault, not this request's), or an image prompt (the + worker's multimodal processor expands it; only the render is checked here and the + scheduler keeps the length check). Raises GenerationError for a render or encode failure + (the worker's own classification) and for a prompt that does not fit.""" + build = getattr(state, "frontend_tokenizer", None) + if build is None: + return None + try: + manager = await asyncio.to_thread(build) + except Exception: # noqa: BLE001 -- server fault, not this request's problem + return None + if manager is None: + return None + + if msg.images: + def work() -> Any: + manager.render_prompt(dataclasses.replace(msg, images=None)) + return None + else: + def work() -> Any: + return manager.encode_prompt(msg, manager.render_prompt(msg)) + + try: + ids = await asyncio.to_thread(work) + except Exception as exc: # noqa: BLE001 -- mirror the worker's classification + raise GenerationError(f"could not encode request: {exc}") from exc + if ids is None: + return None + check_context_fits(int(ids.numel()), state) + return ids async def count_prompt_tokens( @@ -312,6 +802,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,55 +826,38 @@ 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()) -async def prerender_error(spec: GenSpec, state: Any) -> GenerationError | None: - """Render ``spec``'s prompt frontend-side, returning the failure a streaming - adapter should surface as an HTTP 400 *before* committing an SSE stream — - once headers go out, a template rejection can only ride in-stream, where - some agents show nothing but "empty response". Render only; the worker - still renders and encodes authoritatively. Best-effort: a state without a - frontend tokenizer, or one that fails to *initialize*, skips validation - rather than blocking the generation path. - """ - build = getattr(state, "frontend_tokenizer", None) - if build is None: - return None - msg = TokenizeMsg( - uid=0, - text=spec.messages, - sampling_params=SamplingParams(), - chat_template_kwargs=spec.chat_template_kwargs, - tools=spec.template_tools, - ) - try: - manager = await asyncio.to_thread(build) - except Exception: # noqa: BLE001 -- server fault, not this request's problem - return None - try: - await asyncio.to_thread(manager.render_prompt, msg) - except Exception as exc: # noqa: BLE001 -- mirror the worker's classification - return GenerationError(f"could not encode request: {exc}") - return None - - def _make_reasoning_parser(spec: GenSpec, state: Any) -> ReasoningParser | None: """Build a reasoning parser for this generation, or None if the server has no - reasoning parser configured. ``force_reasoning`` matches the encode-side - thinking mode so chat-mode content is never mislabeled as reasoning.""" + reasoning parser configured.""" + return build_reasoning_parser(state.config, _force_reasoning(spec, state)) + + +def _force_reasoning(spec: GenSpec, state: Any) -> bool: + """Whether the generation starts inside reasoning: matches the encode-side thinking + mode so chat-mode content is never mislabeled as reasoning.""" parser_name = getattr(state.config, "reasoning_parser", None) if parser_name == "qwen3": # 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. @@ -406,7 +881,7 @@ def _make_reasoning_parser(spec: GenSpec, state: Any) -> ReasoningParser | None: force_reasoning = ( resolve_thinking_mode(spec.chat_template_kwargs, spec.template_tools) == "thinking" ) - return build_reasoning_parser(state.config, force_reasoning) + return force_reasoning def _split_reasoning(text: str, spec: GenSpec, state: Any) -> tuple[str, str]: @@ -447,11 +922,75 @@ def _parse_tool_response( return None parser = _make_tool_parser(spec, state) result = parser.parse_non_stream(text) - if not result.calls: + # A detector that replays its stream can drop markup without producing a call + # (an empty block): the one-shot content must drop it too. + if not result.calls and result.normal_text == text: return None return result.normal_text, result.calls +_UNRESOLVED = object() + + +async def chat_logprob_markers(state: Any) -> LogprobMarkers | None: + """The token ids chat logprobs are segmented on, resolved once per server from the + frontend tokenizer; None when this server has none.""" + cached = getattr(state, "_logprob_markers", _UNRESOLVED) + if cached is not _UNRESOLVED: + return cached + markers = None + build = getattr(state, "frontend_tokenizer", None) + if build is not None: + try: + manager = await asyncio.to_thread(build) + hidden = _QWEN_SEMANTIC_HIDDEN_TOKENS if _uses_qwen_semantic_protocol(state) else () + markers = await asyncio.to_thread( + resolve_logprob_markers, + manager.tokenizer, THINK_START_TOKEN, THINK_END_TOKEN, hidden, + ) + except Exception: # noqa: BLE001 -- server fault: retried on the next request + return None + state._logprob_markers = markers + return markers + + +def chat_logprobs_supported(state: Any, parse_tools: bool, markers: LogprobMarkers | None) -> bool: + """Whether chat logprobs can be confined to message content for this server's parsers: + the qwen3 reasoning split needs single-id think markers, and tool markup needs a + detector whose one-shot parse replays its stream (the segmenter shadows that stream).""" + config = state.config + reasoning = getattr(config, "reasoning_parser", None) + if reasoning not in (None, "qwen3"): + return False + if reasoning and (markers is None or markers.think_start is None or markers.think_end is None): + return False + if parse_tools: + name = getattr(config, "tool_call_parser", None) + detector = FunctionCallParser.ToolCallParserEnum.get(name) + if detector is None or not (detector.one_shot_replays_stream and detector.supports_streaming): + return False + return True + + +def _wants_logprobs(spec: GenSpec) -> bool: + return bool(getattr(spec.sampling_params, "logprobs", False)) + + +async def _make_logprob_segmenter(spec: GenSpec, state: Any) -> LogprobSegmenter | None: + """The segmenter for a request that asked for logprobs, or None when its entries + cannot be confined to message content (they are then all dropped).""" + markers = await chat_logprob_markers(state) + if not chat_logprobs_supported(state, spec.parse_tools, markers): + return None + reasoning = getattr(state.config, "reasoning_parser", None) is not None + return LogprobSegmenter( + markers or LogprobMarkers(None, None, frozenset()), + reasoning=reasoning, + force_reasoning=reasoning and _force_reasoning(spec, state), + tool_parser=_make_tool_parser(spec, state) if spec.parse_tools else None, + ) + + def _valid_json(text: str) -> bool: if not text: return False @@ -491,6 +1030,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 +1150,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 +1176,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 +1213,30 @@ 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) + + # Content entries ride the next content delta; reasoning and tool-call entries are + # never collected, so a hidden token's entry cannot surface on a visible chunk. + segmenter = await _make_logprob_segmenter(spec, state) if _wants_logprobs(spec) else None tool_parser: FunctionCallParser | None = None if parse_tools: @@ -608,6 +1247,9 @@ async def _generate_events_impl(uid: int, spec: GenSpec, state: Any) -> AsyncIte if candidate is not None and candidate.supports_streaming(): tool_parser = candidate frag_stable = tool_parser.args_fragments_prefix_stable() if tool_parser else True + # A detector that replays its stream for the one-shot parse already decides which + # whitespace after a call is a separator; dropping more here would split the modes. + detector_owns_ws = tool_parser is not None and tool_parser.detector.one_shot_replays_stream # Streaming tool-call assembly: detectors emit fragments (name first, then # argument diffs); they accumulate here and the call is emitted complete when @@ -619,6 +1261,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 = not detector_owns_ws + 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 +1298,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 +1311,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 +1321,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,29 +1329,39 @@ 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 engine_matched_stop: str | None = None + admitted = False 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 ack.prompt_tokens_delta and not admitted: + admitted = True + yield PromptAdmitted(prompt_tokens, cached_tokens) + if segmenter is not None: + pending_logprobs.extend(segmenter.feed(untrimmed(getattr(ack, "logprobs", None)))) 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,18 +1371,23 @@ 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) break + if segmenter is not None: + pending_logprobs.extend(segmenter.finish()) + # Drain residue held in the reasoning parser (a deferred tool block, or a # trailing partial token) so it is not silently dropped. 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 +1397,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,63 +1413,104 @@ 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 ToolCallsDelta(tool_calls) - if finish_reason != "length": - finish_reason = "tool_calls" + yield _content_delta(normal_text) + if tool_calls: + 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) yield GenDone( finish_reason, prompt_tokens, completion_tokens, matched_stop=engine_matched_stop, cached_tokens=cached_tokens, + reasoning_tokens=reasoning_tokens, + timings=timer.finish(), + logprobs=pending_logprobs or None, ) -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] = [] + wants_logprobs = _wants_logprobs(spec) 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 wants_logprobs: + logprob_entries.extend(untrimmed(getattr(ack, "logprobs", None))) if ack.finished: engine_finish_reason = getattr(ack, "finish_reason", None) engine_matched_stop = getattr(ack, "matched_stop", None) break + if wants_logprobs: + segmenter = await _make_logprob_segmenter(spec, state) + logprob_entries = ( + segmenter.feed(logprob_entries) + segmenter.finish() if segmenter is not None else [] + ) reasoning_text, content_text = _split_reasoning(full_content, spec, state) # Engine reason ("stop"/"length"); a tool call overrides it, but a truncation (length) wins. finish_reason = engine_finish_reason or "stop" @@ -802,17 +1518,26 @@ async def _generate_full_impl(uid: int, spec: GenSpec, state: Any) -> GenResult: parsed = _parse_tool_response(content_text, spec, state) if parsed is not None: content_text, tool_calls = parsed - if finish_reason != "length": + if tool_calls and finish_reason != "length": finish_reason = "tool_calls" specials = _leaked_special_tokens(state) + reasoning_special_filter = _make_semantic_special_token_filter(state) + content_special_filter = _make_semantic_special_token_filter(state) return GenResult( - reasoning=strip_special_tokens(reasoning_text, specials).strip(), - content=strip_special_tokens(content_text, specials), + reasoning=reasoning_special_filter.feed( + strip_special_tokens(reasoning_text, specials), final=True + ).strip(), + content=content_special_filter.feed( + strip_special_tokens(content_text, specials), final=True + ), tool_calls=tool_calls, finish_reason=finish_reason, prompt_tokens=prompt_tokens, completion_tokens=completion_tokens, matched_stop=engine_matched_stop, cached_tokens=cached_tokens, + reasoning_tokens=reasoning_tokens, + timings=timer.finish(), + logprobs=logprob_entries, ) diff --git a/python/freetoken/server/launch.py b/python/freetoken/server/launch.py index bb0c256f46..b02876c332 100644 --- a/python/freetoken/server/launch.py +++ b/python/freetoken/server/launch.py @@ -52,13 +52,32 @@ def _run_tokenize_worker(detach: bool, **kwargs) -> None: _detach_process_group() from freetoken.tokenizer import tokenize_worker - tokenize_worker(**kwargs) + ack_queue = kwargs.get("ack_queue") + try: + tokenize_worker(**kwargs) + except Exception as exc: # noqa: BLE001 -- surface the reason, then let it propagate + # Same contract as _run_scheduler: without this, a tokenizer/detokenizer that dies + # before its ready ack reaches the parent only as a dead process, so the supervisor + # reports the generic "backend worker freetoken-detokenizer-0 exited during load" and + # the real cause (an unreadable tokenizer.json, a missing checkpoint file, a ZMQ bind + # clash) is lost. The traceback still prints and the process still exits non-zero. + if ack_queue is not None: + _report_startup_error(ack_queue, exc) + raise def _run_scheduler(args: ServerArgs, ack_queue: mp.Queue[str]) -> None: if args.shell_mode: _detach_process_group() + import threading + + from tqdm import tqdm + + # Only rank 0 draws a bar, so no bar lock spans processes; tqdm's default one is a multiprocessing + # semaphore, which this worker leaves to resource_tracker as leaked when a signal ends it. + tqdm.set_lock(threading.RLock()) + # published (not bound) here: the engine binds it after the allocator setup from freetoken.gpu_select import set_assigned_gpu diff --git a/python/freetoken/server/logprobs.py b/python/freetoken/server/logprobs.py new file mode 100644 index 0000000000..1dce7cc413 --- /dev/null +++ b/python/freetoken/server/logprobs.py @@ -0,0 +1,269 @@ +"""OpenAI logprobs formatting: neutral engine entries -> wire shapes.""" + +from __future__ import annotations + +from collections import deque +from dataclasses import dataclass +from typing import Any + +# AnyJev reads up to 26 letter labels; its vLLM recipe runs --max-logprobs 32 +MAX_COMPLETION_LOGPROBS = 32 +# vLLM's stand-in for -inf: Starlette's JSON encoder refuses non-finite floats +LOGPROB_FLOOR = -9999.0 + + +def _clamped(logprob: float) -> float: + return logprob if logprob >= LOGPROB_FLOOR else LOGPROB_FLOOR # NaN too + + +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 <= MAX_COMPLETION_LOGPROBS: + return f"logprobs must be between 0 and {MAX_COMPLETION_LOGPROBS}" + if req.echo and req.logprobs is not None: + return "echo with logprobs is not supported" + return None + + +def untrimmed(entries: list[dict] | None) -> list[dict]: + """A reply's entries minus those of tokens a matched stop string cut from the text.""" + return [entry for entry in entries or () if not entry.get("stop_trimmed")] + + +def chat_content_entry(entry: dict) -> dict: + return { + "token": entry["token"], + "logprob": _clamped(entry["logprob"]), + "bytes": entry["bytes"], + "top_logprobs": [ + { + "token": candidate["token"], + "logprob": _clamped(candidate["logprob"]), + "bytes": candidate["bytes"], + } + for candidate in entry["top"] + ], + } + + +def completions_logprobs( + entries: list[dict], text_len: int | None = None, as_token_ids: bool = False +) -> dict: + """Each entry's detokenizer-stamped ``text_offset`` is where its token starts in the + generated text (the pieces of a split multi-byte character share its start). Given the + finished text's ``text_len``, offsets are clipped to it, since a matched stop string + trims tokens off the end. ``as_token_ids`` (vLLM return_tokens_as_token_ids) names + every token ``token_id:N`` instead of by its text.""" + + def name(item: dict) -> str: + return f"token_id:{item['token_id']}" if as_token_ids else item["token"] + + tokens: list[str] = [] + token_logprobs: list[float] = [] + top_logprobs: list[dict[str, float]] = [] + + for entry in entries: + tokens.append(name(entry)) + token_logprobs.append(_clamped(entry["logprob"])) + top: dict[str, float] = {} + for candidate in entry["top"]: + # The legacy shape keys alternatives by text, and partial-UTF-8 candidates + # all decode to U+FFFD: keep the most likely one per key. + key, logprob = name(candidate), _clamped(candidate["logprob"]) + if key not in top or logprob > top[key]: + top[key] = logprob + top_logprobs.append(top) + + return { + "tokens": tokens, + "token_logprobs": token_logprobs, + "top_logprobs": top_logprobs, + "text_offset": [ + entry["text_offset"] if text_len is None else min(entry["text_offset"], text_len) + for entry in entries + ], + } + + +def chat_logprobs_object(entries: list[dict]) -> dict: + return {"content": [chat_content_entry(entry) for entry in entries], "refusal": None} + + +@dataclass(frozen=True) +class LogprobMarkers: + """Token ids the chat segmenter keys on, resolved once per server. A think marker is + None unless the tokenizer encodes it as exactly one id that decodes back to it.""" + + think_start: int | None + think_end: int | None + # Special and protocol-hidden tokens: never message content. + hidden: frozenset[int] + + +def resolve_logprob_markers( + tokenizer: Any, think_start: str, think_end: str, hidden_tokens: tuple[str, ...] +) -> LogprobMarkers: + def single_id(marker: str) -> int | None: + ids = tokenizer.encode(marker, add_special_tokens=False) + if len(ids) == 1 and tokenizer.decode(ids) == marker: + return int(ids[0]) + return None + + hidden = {int(tid) for tid in getattr(tokenizer, "all_special_ids", None) or ()} + hidden.update(tid for tid in map(single_id, hidden_tokens) if tid is not None) + return LogprobMarkers(single_id(think_start), single_id(think_end), frozenset(hidden)) + + +def _visible_text(entry: dict) -> str: + """The part of an entry's token that reached the output text.""" + return entry["token"][: entry.get("stop_visible")] + + +class LogprobSegmenter: + """Picks a chat generation's message-content logprob entries out of the sampled ones. + + A pure function of the entry sequence (token ids and single-token strings), so the + streaming and non-streaming paths select the same entries for any ack chunking. + The reasoning split is keyed on the think marker ids and mirrors the qwen3 parser: + a leading ```` (after whitespace) opens reasoning once, the first ```` + ends it for good, and content after it is whitespace-trimmed at both ends. Tool + markup is decided by a shadow tool-call parser fed one token string at a time: its + held buffer is always a suffix of what it was fed, so every resolved character is + either emitted as text or consumed as markup, and a token is content when any of + its characters is emitted. + """ + + def __init__( + self, + markers: LogprobMarkers, + *, + reasoning: bool, + force_reasoning: bool, + tool_parser: Any | None = None, + ) -> None: + self._markers = markers + self._state = "head" if reasoning else "content" + self._force = force_reasoning + self._trim = False # content follows a real : trim it at both ends + self._content_started = False + self._held: list[dict] = [] # whitespace-only tokens awaiting the next text + self._tools = tool_parser + self._queue: deque[list] = deque() # [entry, start, end, is_content] being resolved + self._fed = 0 + self._out: list[dict] = [] + + def feed(self, entries: list[dict]) -> list[dict]: + """Content entries that became final, in generation order.""" + for entry in entries: + # Hidden ids skip the parsers whole: the final EOS adds no text to the output. + if entry["token_id"] not in self._markers.hidden: + self._split_reasoning(entry) + out, self._out = self._out, [] + return out + + def finish(self) -> list[dict]: + held, self._held = self._held, [] + if self._state == "head" and not self._force: + for entry in held: + self._content(entry) + if self._tools is not None: + det = self._tools.detector + base = self._fed - len(det._buffer) + region = det._buffer + if self._tools.recover_truncated_call(): + self._resolve(base, region, "") + base, region = self._fed, "" + self._resolve(base, region, self._tools.finish_stream()) + out, self._out = self._out, [] + return out + + def _split_reasoning(self, entry: dict) -> None: + # A token a stop string cut inside is only its visible text, never a marker. + tid = -1 if "stop_visible" in entry else entry["token_id"] + text = _visible_text(entry) + blank = not text.strip() + if self._state == "head": + if tid == self._markers.think_start: + self._held = [] + self._state = "reasoning" + elif blank: + self._held.append(entry) + elif self._force: + self._held = [] + self._state = "reasoning" + self._split_reasoning(entry) + else: + self._state = "content" + held, self._held = self._held, [] + for e in held + [entry]: + self._content(e) + elif self._state == "reasoning": + if tid == self._markers.think_end: + self._state = "content" + self._trim = True + elif not self._trim: + self._content(entry) + elif blank: + if self._content_started: + self._held.append(entry) + elif not self._content_started: + # The reasoning parser strips this leading whitespace before the tool + # parser sees it: " <" opening a bare call is markup, not a space. + self._content_started = True + self._content(entry, text.lstrip()) + else: + held, self._held = self._held, [] + for e in held + [entry]: + self._content(e) + + def _content(self, entry: dict, text: str | None = None) -> None: + if self._tools is None: + self._out.append(entry) + return + text = _visible_text(entry) if text is None else text + held = self._tools.detector._buffer + base = self._fed - len(held) + self._queue.append([entry, self._fed, self._fed + len(text), False]) + self._fed += len(text) + normal = "".join( + payload for kind, payload in self._tools.parse_stream_events(text) if kind == "text" + ) + resolved = self._fed - len(self._tools.detector._buffer) - base + region = (held + text)[:resolved] if normal else "" + self._resolve(base, region, normal, resolved) + + def _resolve(self, base: int, region: str, normal: str, length: int | None = None) -> None: + """``region`` (``length`` chars from offset ``base``) left the parser's buffer, + emitting ``normal``: mark the tokens it emitted from, then release resolved ones. + ``region`` itself is only needed when something was emitted.""" + if normal: + if region.startswith(normal): + emitted = [(0, len(normal))] + elif region.endswith(normal): + # Text after markup in one token, as when a bare call's ">" merges with + # the next character (">n"): a leftmost match could land inside the markup. + emitted = [(len(region) - len(normal), len(region))] + else: # not seen in practice: fall back to the leftmost alignment + emitted, i = [], 0 + for ch in normal: + j = region.find(ch, i) + if j == -1: + break + emitted.append((j, j + 1)) + i = j + 1 + for item in self._queue: + start, end = item[1] - base, item[2] - base + if any(lo < end and start < hi for lo, hi in emitted): + item[3] = True + resolved = base + (len(region) if length is None else length) + while self._queue and self._queue[0][2] <= resolved: + entry, _, _, is_content = self._queue.popleft() + if is_content: + self._out.append(entry) 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..d447f2a3b2 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,10 @@ from freetoken.tokenizer.effort import EFFORT_SCALE, KNOWN_REASONING_EFFORTS from .api_models import ( + MAX_N, + MAX_PROMPTS, + DetokenizeRequest, + TokenizeRequest, ChatCompletionRequest, CompletionRequest, ModelCard, @@ -22,24 +27,41 @@ ) from .function_call_parser import ToolCallItem from .request_logger import log_request +from .logprobs import ( + chat_logprobs_error, + chat_logprobs_object, + completion_logprobs_error, + completions_logprobs, + untrimmed, +) from .generation import ( DEFAULT_MAX_OUTPUT_TOKENS, + MAX_TOKEN_ID, ContentDelta, GenDone, GenerationError, GenSpec, + GenTimings, ReasoningDelta, ToolCallArgsDelta, ToolCallsDelta, ToolCallStart, + build_metrics, + chat_logprob_markers, + chat_logprobs_supported, + check_context_fits, + frontend_prompt_ids, generate_events, generate_full, - prerender_error, + metrics_enabled, + model_context_length, render_messages, resolve_sampling, submit_generation, ) +logger = logging.getLogger(__name__) + #: The wire superset plus "off", DeepSeek's disable synonym that #: effort_toggle_kwargs has always honored. _ACCEPTED_EFFORTS = (*KNOWN_REASONING_EFFORTS, "off") @@ -53,33 +75,28 @@ def _thinking_type(req: Any) -> str | None: if value in ("enabled", "disabled"): return value return None - - - def chat_request_to_genspec( req: ChatCompletionRequest, model_sampling: dict[str, Any], default_max_tokens: int = DEFAULT_MAX_OUTPUT_TOKENS, + default_thinking_mode: str | None = None, ) -> GenSpec: """OpenAI ChatCompletionRequest -> GenSpec (the OpenAI 'to_sampling_params').""" - from .model_meta import effort_toggle_kwargs + from .model_meta import apply_default_thinking_mode, effort_toggle_kwargs ctk = req.chat_template_kwargs thinking_type = _thinking_type(req) if req.reasoning_effort or thinking_type: ctk = effort_toggle_kwargs(req.reasoning_effort, ctk, thinking_type=thinking_type) + ctk = apply_default_thinking_mode(ctk, default_thinking_mode) + if req.continue_final_message: + ctk = {**ctk, "continue_final_message": True} + sampling_params = _resolve_sampling(req, model_sampling, default_max_tokens=default_max_tokens) + sampling_params.logprobs = bool(req.logprobs) + sampling_params.top_logprobs = req.top_logprobs or 0 return GenSpec( messages=render_messages([m.model_dump(exclude_none=True) for m in req.messages]), - sampling_params=resolve_sampling( - temperature=req.temperature, - top_k=req.top_k, - top_p=req.top_p, - max_tokens=req.max_tokens, - ignore_eos=req.ignore_eos, - model_sampling=model_sampling, - stop=req.stop, - default_max_tokens=default_max_tokens, - ), + 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,11 +148,21 @@ 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() model_id = _served_model_name(state) - ctx = _model_context_length(state) + ctx = model_context_length(state) efforts, default_effort = await _effort_fields(state) return ModelList(data=[ModelCard( id=model_id, @@ -146,6 +173,124 @@ 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", + ) + + +async def _chat_logprobs_conflict(req: ChatCompletionRequest, state: Any) -> str | None: + """Refuse chat logprobs when the server's parsers can hide or reclassify generated + tokens in a way the logprob segmenter cannot follow: entries must cover message + content only, never reasoning or tool-call tokens. The raw /v1/completions path has + no semantic layer and keeps full logprobs support.""" + if not req.logprobs: + return None + markers = await chat_logprob_markers(state) + if not chat_logprobs_supported(state, False, markers): + return ( + "logprobs on /v1/chat/completions are not supported with the " + f"{state.config.reasoning_parser!r} reasoning parser: reasoning tokens cannot " + "be told apart from message content. Use /v1/completions for raw token logprobs." + ) + if not chat_logprobs_supported(state, _should_parse_tools(req), markers): + return ( + "logprobs with tool parsing on /v1/chat/completions are not supported for the " + f"{state.config.tool_call_parser!r} tool-call parser: tool-call tokens cannot be " + "told apart from 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 +300,24 @@ 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") + unimplemented = _unimplemented_field(req, _UNIMPLEMENTED_CHAT_FIELDS) + if unimplemented is not None: + return create_error_response(unimplemented) + logprobs_error = chat_logprobs_error(req) + if logprobs_error is not None: + return create_error_response(logprobs_error, param="top_logprobs") + logprobs_conflict = await _chat_logprobs_conflict(req, state) + if logprobs_conflict is not None: + return create_error_response(logprobs_conflict, param="logprobs") if _response_format_unsupported(req.response_format): + rtype = (req.response_format or {}).get("type") return create_error_response( - "response_format json_object/json_schema is not supported (no constrained decoding)", + f"response_format.type {rtype!r} is not supported " + "(supported: text, json_object, json_schema)", param="response_format", ) - if req.n != 1: - return create_error_response("Only n=1 is supported", param="n") + if not 1 <= req.n <= MAX_N: + return create_error_response(f"n must be between 1 and {MAX_N}", param="n") # Case/whitespace and the "off" disable synonym stay accepted here because # effort_toggle_kwargs normalizes and honors them downstream. effort = req.reasoning_effort.strip().lower() if isinstance(req.reasoning_effort, str) else None @@ -185,56 +339,100 @@ 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)) + if (bad := await _token_id_error(state, spec.sampling_params.client_token_ids())) is not None: + return bad - if req.stream: - # Non-stream requests already surface render failures as a clean 400 - # through GenerationError; only the stream path needs the pre-check. - err = await prerender_error(spec, state) - 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) + choices.append({ + "index": index, + "message": message, + # OpenAI always sends the key, null when logprobs were not requested + "logprobs": chat_logprobs_object(result.logprobs) if req.logprobs else None, + "finish_reason": result.finish_reason, + }) + + 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 +440,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 +487,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"] = chat_logprobs_object(ev.logprobs) + yield _sse(_chat_chunk(req, uid, [choice])) elif isinstance(ev, ToolCallStart): open_tool = { "index": tool_calls_sent, @@ -305,7 +512,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 +529,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 +553,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 +565,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 +573,30 @@ 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 + done: dict[str, Any] = {"delta": {}, "index": index, "finish_reason": ev.finish_reason} + if ev.logprobs: + done["logprobs"] = chat_logprobs_object(ev.logprobs) + yield _sse(_chat_chunk(req, uid, [done])) + 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,128 +610,395 @@ 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") + prompts = _prompt_list(req.prompt) + num_prompts = len(prompts) + if num_prompts * req.n > MAX_PROMPTS: + return create_error_response( + f"prompts x n must be at most {MAX_PROMPTS}, got {num_prompts * req.n}", param="prompt" + ) + 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)) + token_prompts = _is_token_prompt(req.prompt) + id_fields = sampling.client_token_ids() + if token_prompts: + id_fields += [("prompt", ids) for ids in prompts] + if (bad := await _token_id_error(state, id_fields)) is not None: + return bad - prompts = [req.prompt] if isinstance(req.prompt, str) else req.prompt - assert isinstance(prompts, list) - 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 + 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) + prompt_ids = [] + try: + for prompt in prompts: + if token_prompts: + check_context_fits(len(prompt), state) + prompt_ids.append(None) + continue + prompt_ids.append(await frontend_prompt_ids( + TokenizeMsg(uid=0, text=prompt, sampling_params=sampling), state )) - ) - chunks = stream_completion_chunks(uid, req, state) + except GenerationError as exc: + return create_error_response(str(exc), code=exc.code) + uids: list[int] = [] + for prompt, ids in zip(prompts, prompt_ids): + for _ in range(req.n): + uid = state.new_user() + msg = ( + TokenizeMsg(uid=uid, text="", input_ids=prompt, sampling_params=sampling) + if token_prompts + else TokenizeMsg(uid=uid, text=prompt, sampling_params=sampling, input_ids=ids) + ) + await state.send_one(msg) + uids.append(uid) + + if req.stream: + 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) + return create_error_response(ack.error, code=getattr(ack, "error_code", None)) prompt_tokens += ack.prompt_tokens_delta completion_tokens += ack.completion_tokens_delta cached_tokens += ack.cached_tokens text += ack.incremental_output + if req.logprobs is not None: + entries.extend(untrimmed(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, len(text), req.return_tokens_as_token_ids) + 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 + sent_len = 0 + # Logprob entries of acks with no visible text yet (a held partial UTF-8 character, an + # EOS), carried to the next chunk that has text; a reply can bring several at once. + pending: list[dict] = [] 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" + error = {"message": ack.error, "type": "invalid_request_error", "code": getattr(ack, "error_code", None)} + yield _sse({"error": error}) + 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 req.logprobs is not None: + pending.extend(untrimmed(ack.logprobs)) if ack.incremental_output: + sent_len += len(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, + # Mid-stream, a stamp past sent_len is held-back text still to come. + "logprobs": ( + completions_logprobs( + pending, + sent_len if ack.finished else None, + req.return_tokens_as_token_ids, + ) + if pending + else None + ), + }], + ) ) + pending = [] if ack.finished: finish_reason = getattr(ack, "finish_reason", None) or "stop" break + # Entries whose text never showed (EOS, a stop string, a trailing partial character) + # ride the finish chunk, so the joined stream matches the non-stream list. 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}], - } + _completion_chunk( + req, uid, + [{ + "text": "", + "index": index, + "finish_reason": finish_reason, + "logprobs": ( + completions_logprobs(pending, sent_len, req.return_tokens_as_token_ids) + if pending + else 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 _token_id_error(state: Any, fields: list[tuple[str, list[int]]]) -> JSONResponse | None: + """A 400 for the first client token id outside the vocabulary, before anything is + queued. Without a frontend tokenizer the bound is the engine's int32; the tokenizer + worker checks against its own tokenizer either way.""" + if not fields: + return None + bound = await _vocab_len(state) or MAX_TOKEN_ID + 1 + for name, ids in fields: + if not all(0 <= t < bound for t in ids): + what = "prompt token ids" if name == "prompt" else name + return create_error_response(f"{what} must be in [0, {bound})", param=name) + return None + + +async def _vocab_len(state: Any) -> int | None: + """``len(tokenizer)`` (added tokens included), or None without a frontend tokenizer.""" + build = getattr(state, "frontend_tokenizer", None) + if build is None: + return None + try: + manager = await asyncio.to_thread(build) + return len(manager.tokenizer) + except Exception: # noqa: BLE001 -- the tokenizer worker still range-checks + return None + + +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 +1024,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 +1033,44 @@ 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, + allowed_token_ids=req.allowed_token_ids if isinstance(req, CompletionRequest) else None, ) + if isinstance(req, CompletionRequest) and req.logprobs is not None: + sampling_params.logprobs = True + sampling_params.top_logprobs = req.logprobs + if sampling_params.allowed_token_ids: + _check_restricted_logprobs(sampling_params) + return sampling_params + + +def _check_restricted_logprobs(sp: SamplingParams) -> None: + """Logprobs restricted to allowed_token_ids are vLLM's processed_logprobs at temperature 0: + log_softmax over the allowed ids of the unprocessed logits. Anything vLLM would also fold + into them is refused rather than silently left out.""" + if sp.temperature > 0: + raise ValueError("allowed_token_ids with logprobs needs temperature=0") + if sp.presence_penalty or sp.frequency_penalty or sp.repetition_penalty != 1.0: + raise ValueError("allowed_token_ids with logprobs does not support penalties") + if sp.min_tokens > 0: + raise ValueError("allowed_token_ids with logprobs does not support min_tokens") + if sp.logit_bias: + raise ValueError("allowed_token_ids with logprobs does not support logit_bias") + + +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 +1133,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 +1151,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,30 +1163,102 @@ def _usage(prompt_tokens: int, completion_tokens: int, cached_tokens: int = 0) - # disabled report and a 0-token hit serialize identically. if cached_tokens > 0: usage["prompt_tokens_details"] = {"cached_tokens": cached_tokens} + # OpenAI's completion_tokens_details.reasoning_tokens: tokens up to and including the + # reasoning end tag, only when the model emitted one. + if reasoning_tokens > 0: + usage["completion_tokens_details"] = {"reasoning_tokens": reasoning_tokens} return usage def _response_format_unsupported(response_format: dict[str, Any] | None) -> bool: - # We have no constrained/guided decoding; only plain text ('text' or unset) is honored. - return response_format is not None and response_format.get("type") not in (None, "text") + # Soft-accept OpenAI json_object/json_schema. FreeToken has no constrained + # decoding; clients that already ask for JSON in the prompt still need the + # wire field accepted (HTTP 200) instead of a hard 400. Unknown types 400. + if response_format is None: + return False + return response_format.get("type") not in (None, "text", "json_object", "json_schema") + + +# vLLM /v1/completions fields that change the result and are not implemented here: a set +# value is a 400 rather than silently dropped by extra="allow". +_UNIMPLEMENTED_COMPLETION_FIELDS = ( + "logprob_token_ids", + "prompt_logprobs", + "guided_json", + "guided_regex", + "guided_choice", + "guided_grammar", + "structured_outputs", + "bad_words", +) + + +# chat has neither of the completions-only ones +_UNIMPLEMENTED_CHAT_FIELDS = ( + "allowed_token_ids", + "return_tokens_as_token_ids", + *_UNIMPLEMENTED_COMPLETION_FIELDS, +) + + +def _unimplemented_field(req: Any, names: tuple[str, ...]) -> str | None: + for name in names: + if _is_set(getattr(req, name, None)): + return f"{name} is not supported by this server" + return None + + +def _is_set(value: Any) -> bool: + # not truthiness: prompt_logprobs=0 is a real request, and 0 == False + if value is None or value is False: + return False + return not (isinstance(value, (str, list, dict)) and not value) + + +def _empty_prompt(req: CompletionRequest) -> bool: + prompt = req.prompt + if isinstance(prompt, str): + # suffix-only fill-in-the-middle is a real request + return not prompt and req.suffix is None + if not prompt: + return True + return any( + (isinstance(item, list) and not item) + or (isinstance(item, str) and not item and req.suffix is None) + for item in prompt + ) def _completion_unsupported_reason(req: CompletionRequest) -> str | None: + unimplemented = _unimplemented_field(req, _UNIMPLEMENTED_COMPLETION_FIELDS) + if unimplemented is not None: + return unimplemented + if _empty_prompt(req): + return "prompt must not be empty" 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" + if req.echo: + return "echo is not supported with token-id prompts" + if req.suffix is not None: + return "suffix is not supported with token-id prompts" + logprobs_error = completion_logprobs_error(req) + if logprobs_error is not None: + return logprobs_error if _response_format_unsupported(req.response_format): - return "response_format json_object/json_schema is not supported (no constrained decoding)" + rtype = (req.response_format or {}).get("type") + return ( + f"response_format.type {rtype!r} is not supported " + "(supported: text, json_object, json_schema)" + ) return None +def _prompt_list(prompt: Any) -> list: + """One entry per prompt: a string or a token-id list.""" + if isinstance(prompt, str) or (prompt and all(isinstance(t, int) for t in prompt)): + return [prompt] + return list(prompt) + + def _is_token_prompt(prompt: Any) -> bool: return ( isinstance(prompt, list) @@ -700,11 +1295,3 @@ def _served_model_name(state: Any) -> str: return getattr(state.config, "served_model_name", None) or state.config.model_path -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.""" - 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 diff --git a/python/freetoken/server/reasoning_parser.py b/python/freetoken/server/reasoning_parser.py index 6788675ca9..c571b97a79 100644 --- a/python/freetoken/server/reasoning_parser.py +++ b/python/freetoken/server/reasoning_parser.py @@ -393,7 +393,7 @@ def _scan(self, text: str, *, hold_partial: bool) -> tuple[str, str]: class ThinkReasoningParser(BaseReasoningParser): """Generic ````/```` reasoning parser for models that wrap their - chain-of-thought in think tags (Qwen3/3.5, GLM-4.x, MiniMax-M2). No DSML tool + chain-of-thought in think tags (GLM-4.x, MiniMax-M2). No DSML tool marker — reasoning ends at ````. Assumes the model closes ```` before any tool call. These families do @@ -412,6 +412,113 @@ def __init__(self, force_reasoning: bool = False, stream_reasoning: bool = True) ) +class Qwen3ReasoningParser(BaseReasoningParser): + """Qwen3 ```` parser with one positional rule for both paths. + + Reasoning is open from the start when the template pre-opened it + (``force_reasoning``) or when the generation itself opens with ```` + (stripped once). It ends at the FIRST ```` and is never re-entered: + every later ````/```` -- quoted in prose, code or tool-call + arguments -- is literal content. Reasoning is whitespace-trimmed at both + ends, and so is content that follows a real ````; streaming holds + whitespace-only tails so its deltas join to exactly the one-shot result. + The one-shot parse replays the streaming machine, so the two cannot drift. + """ + + def __init__(self, force_reasoning: bool = False, stream_reasoning: bool = True) -> None: + super().__init__( + think_start_token=THINK_START_TOKEN, + think_end_token=THINK_END_TOKEN, + force_reasoning=force_reasoning, + stream_reasoning=stream_reasoning, + ) + # The generation start, held while it can still grow into ws* + "". + self._head: Optional[str] = "" + self._trim_content = False + self._reasoning_started = False + self._content_started = False + self._unstreamed_reasoning = "" + + def detect_and_parse(self, text: str) -> ReasoningParseResult: + clone = type(self)( + force_reasoning=self.force_reasoning, stream_reasoning=self.stream_reasoning + ) + first = clone.parse_streaming_increment(text) + rest = clone.flush() + return ReasoningParseResult( + reasoning_text=first.reasoning_text + rest.reasoning_text, + normal_text=first.normal_text + rest.normal_text, + ) + + def parse_streaming_increment(self, new_text: str) -> ReasoningParseResult: + if self._head is not None: + self._head += new_text + head = self._head.lstrip() + if head.startswith(self.think_start_token): + self._head = None + self._in_reasoning = True + return self._consume(head[len(self.think_start_token):], final=False) + if not head or self.think_start_token.startswith(head): + return ReasoningParseResult() + new_text, self._head = self._head, None + return self._consume(new_text, final=False) + + def flush(self) -> ReasoningParseResult: + text, self._head = self._head or "", None + return self._consume(text, final=True) + + def _reasoning_piece(self, piece: str, closing: bool) -> str: + if not self._reasoning_started: + piece = piece.lstrip() + self._reasoning_started = bool(piece) + if self.stream_reasoning: + return piece + self._unstreamed_reasoning += piece + if not closing: + return "" + piece, self._unstreamed_reasoning = self._unstreamed_reasoning, "" + return piece + + def _consume(self, text: str, final: bool) -> ReasoningParseResult: + buf = self._buffer + text + self._buffer = "" + reasoning = "" + if self._in_reasoning: + end = buf.find(self.think_end_token) + if end == -1: + if final: + return ReasoningParseResult( + reasoning_text=self._reasoning_piece(buf.rstrip(), closing=True) + ) + safe, _ = self._split_trailing_partial(buf) + body = safe.rstrip() + self._buffer = buf[len(body):] + return ReasoningParseResult( + reasoning_text=self._reasoning_piece(body, closing=False) + ) + reasoning = self._reasoning_piece(buf[:end].rstrip(), closing=True) + self._in_reasoning = False + self._trim_content = True + buf = buf[end + len(self.think_end_token):] + if not self._trim_content: + return ReasoningParseResult(reasoning_text=reasoning, normal_text=buf) + if not self._content_started: + buf = buf.lstrip() + self._content_started = bool(buf) + body = buf.rstrip() + if not final: + self._buffer = buf[len(body):] + return ReasoningParseResult(reasoning_text=reasoning, normal_text=body) + + def _split_trailing_partial(self, text: str) -> tuple[str, str]: + # Only is tracked: a later is literal, never a marker. + token = self.think_end_token + for k in range(min(len(token) - 1, len(text)), 0, -1): + if text.endswith(token[:k]): + return text[:-k], text[-k:] + return text, "" + + class MiniMaxM3ReasoningParser(BaseReasoningParser): """Reasoning parser for MiniMax-M3's ``...`` protocol. @@ -876,7 +983,7 @@ class ReasoningParser: ReasoningParserEnum: Dict[str, Type[BaseReasoningParser]] = { "deepseekv32": DeepSeekV32ReasoningParser, "gpt_oss": GptOssHarmonyReasoningParser, - "qwen3": ThinkReasoningParser, + "qwen3": Qwen3ReasoningParser, "glm": ThinkReasoningParser, "minimax": ThinkReasoningParser, "minimax_m3": MiniMaxM3ReasoningParser, diff --git a/python/freetoken/server/responses_api.py b/python/freetoken/server/responses_api.py index 71d6eed873..4240080510 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" @@ -452,17 +467,24 @@ def build_responses_response( usage=_usage( result.prompt_tokens, result.completion_tokens, result.cached_tokens if cache_report else 0, + reasoning_tokens=result.reasoning_tokens, ), 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,9 +498,16 @@ 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: +def _usage( + prompt_tokens: int, completion_tokens: int, cached_tokens: int = 0, reasoning_tokens: int = 0 +) -> ResponseUsage: # input_tokens stays inclusive of the cached prefix (OpenAI semantics); # cached_tokens is nonzero only under --enable-cache-report. return ResponseUsage( @@ -486,7 +515,7 @@ def _usage(prompt_tokens: int, completion_tokens: int, cached_tokens: int = 0) - output_tokens=completion_tokens, total_tokens=prompt_tokens + completion_tokens, input_tokens_details=InputTokensDetails(cached_tokens=cached_tokens, cache_write_tokens=0), - output_tokens_details=OutputTokensDetails(reasoning_tokens=0), + output_tokens_details=OutputTokensDetails(reasoning_tokens=reasoning_tokens), ) @@ -499,13 +528,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( @@ -522,6 +553,7 @@ def snapshot(status, output, usage=None, incomplete_reason=None): usage_pt = 0 usage_ct = 0 usage_cached = 0 + usage_reasoning = 0 def close_current() -> list[str]: nonlocal current, output_index @@ -707,6 +739,13 @@ 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 + usage_reasoning = ev.reasoning_tokens + 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": @@ -714,8 +753,9 @@ def args_delta_frame(fragment: str) -> str: type="response.incomplete", sequence_number=seq.next(), response=snapshot( "incomplete", output_items, - usage=_usage(usage_pt, usage_ct, usage_cached), + usage=_usage(usage_pt, usage_ct, usage_cached, usage_reasoning), incomplete_reason="max_output_tokens", + request_metrics=request_metrics, ), )) else: @@ -723,7 +763,8 @@ def args_delta_frame(fragment: str) -> str: type="response.completed", sequence_number=seq.next(), response=snapshot( "completed", output_items, - usage=_usage(usage_pt, usage_ct, usage_cached), + usage=_usage(usage_pt, usage_ct, usage_cached, usage_reasoning), + request_metrics=request_metrics, ), )) except GenerationError as exc: @@ -736,7 +777,7 @@ def args_delta_frame(fragment: str) -> str: type="response.failed", sequence_number=seq.next(), response=_response_obj( response_id, created, req.model, output_items, status="failed", - usage=_usage(usage_pt, usage_ct, usage_cached), + usage=_usage(usage_pt, usage_ct, usage_cached, usage_reasoning), # codex reads this code to tell a blown context window from a generic failure. # model_construct because ResponseError.code is a closed Literal in the SDK. error=ResponseError.model_construct( @@ -751,7 +792,7 @@ def args_delta_frame(fragment: str) -> str: type="response.failed", sequence_number=seq.next(), response=_response_obj( response_id, created, req.model, output_items, status="failed", - usage=_usage(usage_pt, usage_ct, usage_cached), + usage=_usage(usage_pt, usage_ct, usage_cached, usage_reasoning), error=ResponseError(code="server_error", message=str(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/server/supervisor.py b/python/freetoken/server/supervisor.py index 322419a65f..20deb7d132 100644 --- a/python/freetoken/server/supervisor.py +++ b/python/freetoken/server/supervisor.py @@ -55,6 +55,12 @@ class BackendHandle: processes: List[Any] = field(default_factory=list) expected_acks: int = 0 + def release_acks(self) -> None: + """Drop the ack queue once readiness has read it for the last time. A multiprocessing queue's + semaphores are unlinked only when it is collected, and a main process that a signal ends never + collects it, so resource_tracker reports them as leaked.""" + self.ack_queue = None + class WorkerDied(Exception): def __init__(self, message: str) -> None: @@ -172,6 +178,7 @@ def _shutting_down() -> bool: on_failure(exc.message) return + handle.release_acks() on_ready() while True: diff --git a/python/freetoken/shell/__init__.py b/python/freetoken/shell/__init__.py index e429f487a0..d31eed4cd5 100644 --- a/python/freetoken/shell/__init__.py +++ b/python/freetoken/shell/__init__.py @@ -45,6 +45,10 @@ def _build_parser(prog: str) -> argparse.ArgumentParser: default=None, help="FreeToken server URL (default: $FREETOKEN_HOST, else http://127.0.0.1:1919)", ) + parser.add_argument( + "--api-key", + help="API key of a server started with --api-key (default: $FREETOKEN_API_KEY)", + ) return parser @@ -63,7 +67,7 @@ def main(argv: Sequence[str] | None = None, *, prog: str = "ft shell") -> int: except SystemExit as exc: return int(exc.code) if isinstance(exc.code, int) else 2 - from freetoken.launch import resolve_server_url + from freetoken.launch import resolve_api_key, resolve_server_url try: server = resolve_server_url(parsed.server) @@ -74,7 +78,7 @@ def main(argv: Sequence[str] | None = None, *, prog: str = "ft shell") -> int: from .tui import run_shell try: - return asyncio.run(run_shell(server.origin)) + return asyncio.run(run_shell(server.origin, api_key=resolve_api_key(parsed.api_key))) except KeyboardInterrupt: return 130 diff --git a/python/freetoken/shell/client.py b/python/freetoken/shell/client.py index 183fa4dff2..d111536294 100644 --- a/python/freetoken/shell/client.py +++ b/python/freetoken/shell/client.py @@ -27,7 +27,7 @@ import openai from openai import AsyncOpenAI -# Any string is accepted by the server; sent so the SDK doesn't refuse to build a request. +# Sent when no key is given: a server without --api-key accepts any string, and the SDK needs one. LOCAL_API_KEY = "freetoken-local" CONTROL_TIMEOUT = 10.0 # Per-read timeout on the chat stream. Generous: it has to cover the gap before the first @@ -97,14 +97,15 @@ def __init__( self, origin: str, *, - api_key: str = LOCAL_API_KEY, + api_key: str | None = None, timeout: float = CONTROL_TIMEOUT, ) -> None: self.origin = origin.rstrip("/") self.timeout = timeout + self._api_key = api_key or LOCAL_API_KEY self._openai = AsyncOpenAI( base_url=f"{self.origin}/v1", - api_key=api_key, + api_key=self._api_key, max_retries=0, # a retried chat request would generate the turn twice timeout=openai.Timeout(CHAT_READ_TIMEOUT, connect=10.0), ) @@ -118,7 +119,7 @@ def _request_json_blocking( self, method: str, path: str, body: dict[str, Any] | None, timeout: float ) -> dict[str, Any]: data = None - headers = {"Accept": "application/json"} + headers = {"Accept": "application/json", "Authorization": f"Bearer {self._api_key}"} if body is not None: data = json.dumps(body).encode("utf-8") headers["Content-Type"] = "application/json" diff --git a/python/freetoken/shell/tui.py b/python/freetoken/shell/tui.py index 93bd18cc33..72d7ebe0e0 100644 --- a/python/freetoken/shell/tui.py +++ b/python/freetoken/shell/tui.py @@ -515,13 +515,15 @@ def _format_load_progress(doc: dict) -> str: return f"loading ({phase})..." -async def run_shell(origin: str, *, connect_grace: float = 0.0) -> int: +async def run_shell( + origin: str, *, connect_grace: float = 0.0, api_key: str | None = None +) -> int: """Attach to the FreeToken server at ``origin`` and run the terminal chat. ``connect_grace`` is how long to keep retrying a refused connection before giving up -- left at 0 when attaching to a server the user says is already running, raised when the caller just started one in this process (see ``server/api_server.py``).""" - client = ShellClient(origin) + client = ShellClient(origin, api_key=api_key) try: return await _run_shell(client, origin, connect_grace=connect_grace) finally: diff --git a/python/freetoken/third_party/__init__.py b/python/freetoken/third_party/__init__.py new file mode 100644 index 0000000000..633cc8a798 --- /dev/null +++ b/python/freetoken/third_party/__init__.py @@ -0,0 +1 @@ +"""Vendored third-party Python code; each package carries its LICENSE and VERSION.""" diff --git a/python/freetoken/third_party/gguf/LICENSE b/python/freetoken/third_party/gguf/LICENSE new file mode 100644 index 0000000000..76f67efdc6 --- /dev/null +++ b/python/freetoken/third_party/gguf/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2023 Georgi Gerganov + +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. diff --git a/python/freetoken/third_party/gguf/VERSION b/python/freetoken/third_party/gguf/VERSION new file mode 100644 index 0000000000..f3ec79b970 --- /dev/null +++ b/python/freetoken/third_party/gguf/VERSION @@ -0,0 +1,10 @@ +llama.cpp 8c2134a8755d8a9792af513dab20282affe903e4 (2026-10-14), the same pin as kernel/csrc/ggml + +Only the gguf-py modules that the GGUF reader and the ggml type tables need: gguf-py/gguf/{constants,gguf_reader,quants,lazy}.py. +The PyPI gguf release lags the ggml types in this tree (gguf 0.19.0 cannot read Q2_0 tensors), so FreeToken reads GGUF files with this copy. + +Local changes, kept in freetoken.patch (already applied): +- __init__.py: imports only the vendored modules (upstream also pulls in the writer, vocab and metadata modules and their dependencies). +- gguf_reader.py: imports constants relatively, so the installed gguf package is never picked up. + +To update: copy the same files from the llama.cpp commit that kernel/csrc/ggml pins, apply freetoken.patch, and update this file. diff --git a/python/freetoken/third_party/gguf/__init__.py b/python/freetoken/third_party/gguf/__init__.py new file mode 100644 index 0000000000..f4c74ef9e5 --- /dev/null +++ b/python/freetoken/third_party/gguf/__init__.py @@ -0,0 +1,4 @@ +from .constants import * +from .lazy import * +from .gguf_reader import * +from .quants import * diff --git a/python/freetoken/third_party/gguf/constants.py b/python/freetoken/third_party/gguf/constants.py new file mode 100644 index 0000000000..c31ab4a42e --- /dev/null +++ b/python/freetoken/third_party/gguf/constants.py @@ -0,0 +1,6332 @@ +from __future__ import annotations + +from enum import Enum, IntEnum, auto +from typing import Any + +# +# constants +# + +GGUF_MAGIC = 0x46554747 # "GGUF" +GGUF_VERSION = 3 +GGUF_DEFAULT_ALIGNMENT = 32 +GGML_QUANT_VERSION = 2 # GGML_QNT_VERSION from ggml.h +GGML_MAX_DIMS = 4 # GGML_MAX_DIMS from ggml.h + +# +# metadata keys +# + + +class Keys: + class General: + TYPE = "general.type" + ARCHITECTURE = "general.architecture" + QUANTIZATION_VERSION = "general.quantization_version" + ALIGNMENT = "general.alignment" + FILE_TYPE = "general.file_type" + + # Per-tensor extra options (tensor name array + parallel option arrays, e.g. prec_a4). + TENSOR_EXTRA_NAME = "general.tensor_extra.name" + TENSOR_EXTRA_PREC_A4 = "general.tensor_extra.prec_a4" + + # Recommended Sampler Parameters + SAMPLING_SEQUENCE = "general.sampling.sequence" + SAMPLING_TOP_K = "general.sampling.top_k" + SAMPLING_TOP_P = "general.sampling.top_p" + SAMPLING_MIN_P = "general.sampling.min_p" + SAMPLING_XTC_PROBABILITY = "general.sampling.xtc_probability" + SAMPLING_XTC_THRESHOLD = "general.sampling.xtc_threshold" + SAMPLING_TEMP = "general.sampling.temp" + SAMPLING_PENALTY_LAST_N = "general.sampling.penalty_last_n" + SAMPLING_PENALTY_REPEAT = "general.sampling.penalty_repeat" + SAMPLING_MIROSTAT = "general.sampling.mirostat" + SAMPLING_MIROSTAT_TAU = "general.sampling.mirostat_tau" + SAMPLING_MIROSTAT_ETA = "general.sampling.mirostat_eta" + + # Authorship Metadata + NAME = "general.name" + AUTHOR = "general.author" + VERSION = "general.version" + ORGANIZATION = "general.organization" + + FINETUNE = "general.finetune" + BASENAME = "general.basename" + + DESCRIPTION = "general.description" + QUANTIZED_BY = "general.quantized_by" + + SIZE_LABEL = "general.size_label" + + # Licensing details + LICENSE = "general.license" + LICENSE_NAME = "general.license.name" + LICENSE_LINK = "general.license.link" + + # Typically represents the converted GGUF repo (Unless native) + URL = "general.url" # Model Website/Paper + DOI = "general.doi" + UUID = "general.uuid" + REPO_URL = "general.repo_url" # Model Source Repository (git/svn/etc...) + + # Model Source during conversion + SOURCE_URL = "general.source.url" # Model Website/Paper + SOURCE_DOI = "general.source.doi" + SOURCE_UUID = "general.source.uuid" + SOURCE_REPO_URL = "general.source.repo_url" # Model Source Repository (git/svn/etc...) + + # Base Model Source. There can be more than one source if it's a merged + # model like with 'Mistral-7B-Merge-14-v0.1'. This will assist in + # tracing linage of models as it is finetuned or merged over time. + BASE_MODEL_COUNT = "general.base_model.count" + BASE_MODEL_NAME = "general.base_model.{id}.name" + BASE_MODEL_AUTHOR = "general.base_model.{id}.author" + BASE_MODEL_VERSION = "general.base_model.{id}.version" + BASE_MODEL_ORGANIZATION = "general.base_model.{id}.organization" + BASE_MODEL_DESCRIPTION = "general.base_model.{id}.description" + BASE_MODEL_URL = "general.base_model.{id}.url" # Model Website/Paper + BASE_MODEL_DOI = "general.base_model.{id}.doi" + BASE_MODEL_UUID = "general.base_model.{id}.uuid" + BASE_MODEL_REPO_URL = "general.base_model.{id}.repo_url" # Model Source Repository (git/svn/etc...) + + # Dataset Source + DATASET_COUNT = "general.dataset.count" + DATASET_NAME = "general.dataset.{id}.name" + DATASET_AUTHOR = "general.dataset.{id}.author" + DATASET_VERSION = "general.dataset.{id}.version" + DATASET_ORGANIZATION = "general.dataset.{id}.organization" + DATASET_DESCRIPTION = "general.dataset.{id}.description" + DATASET_URL = "general.dataset.{id}.url" # Model Website/Paper + DATASET_DOI = "general.dataset.{id}.doi" + DATASET_UUID = "general.dataset.{id}.uuid" + DATASET_REPO_URL = "general.dataset.{id}.repo_url" # Model Source Repository (git/svn/etc...) + + # Array based KV stores + TAGS = "general.tags" + LANGUAGES = "general.languages" + + class LLM: + VOCAB_SIZE = "{arch}.vocab_size" + CONTEXT_LENGTH = "{arch}.context_length" + EMBEDDING_LENGTH = "{arch}.embedding_length" + EMBEDDING_LENGTH_OUT = "{arch}.embedding_length_out" + FEATURES_LENGTH = "{arch}.features_length" + BLOCK_COUNT = "{arch}.block_count" + LEADING_DENSE_BLOCK_COUNT = "{arch}.leading_dense_block_count" + FEED_FORWARD_LENGTH = "{arch}.feed_forward_length" + EXPERT_FEED_FORWARD_LENGTH = "{arch}.expert_feed_forward_length" + EXPERT_SHARED_FEED_FORWARD_LENGTH = "{arch}.expert_shared_feed_forward_length" + EXPERT_CHUNK_FEED_FORWARD_LENGTH = "{arch}.expert_chunk_feed_forward_length" + USE_PARALLEL_RESIDUAL = "{arch}.use_parallel_residual" + TENSOR_DATA_LAYOUT = "{arch}.tensor_data_layout" + EXPERT_COUNT = "{arch}.expert_count" + EXPERT_USED_COUNT = "{arch}.expert_used_count" + EXPERT_SHARED_COUNT = "{arch}.expert_shared_count" + EXPERT_GROUP_COUNT = "{arch}.expert_group_count" + EXPERT_GROUP_USED_COUNT = "{arch}.expert_group_used_count" + EXPERT_WEIGHTS_SCALE = "{arch}.expert_weights_scale" + EXPERT_WEIGHTS_NORM = "{arch}.expert_weights_norm" + EXPERT_GATING_FUNC = "{arch}.expert_gating_func" + EXPERT_GROUP_SCALE = "{arch}.expert_group_scale" + EXPERT_LATENT_LENGTH = "{arch}.expert_latent_length" + EXPERTS_PER_GROUP = "{arch}.experts_per_group" + MOE_EVERY_N_LAYERS = "{arch}.moe_every_n_layers" + MOE_LATENT_SIZE = "{arch}.moe_latent_size" + NEXTN_PREDICT_LAYERS = "{arch}.nextn_predict_layers" + NUM_DEEPSTACK_LAYERS = "{arch}.n_deepstack_layers" + DEEPSTACK_MAPPING = "{arch}.deepstack_mapping" + POOLING_TYPE = "{arch}.pooling_type" + LOGIT_SCALE = "{arch}.logit_scale" + DECODER_START_TOKEN_ID = "{arch}.decoder_start_token_id" + DECODER_BLOCK_COUNT = "{arch}.decoder_block_count" + ATTN_LOGIT_SOFTCAPPING = "{arch}.attn_logit_softcapping" + ROUTER_LOGIT_SOFTCAPPING = "{arch}.router_logit_softcapping" + FINAL_LOGIT_SOFTCAPPING = "{arch}.final_logit_softcapping" + SWIN_NORM = "{arch}.swin_norm" + RESCALE_EVERY_N_LAYERS = "{arch}.rescale_every_n_layers" + TIME_MIX_EXTRA_DIM = "{arch}.time_mix_extra_dim" + TIME_DECAY_EXTRA_DIM = "{arch}.time_decay_extra_dim" + RESIDUAL_SCALE = "{arch}.residual_scale" + EMBEDDING_SCALE = "{arch}.embedding_scale" + TOKEN_SHIFT_COUNT = "{arch}.token_shift_count" + INTERLEAVE_MOE_LAYER_STEP = "{arch}.interleave_moe_layer_step" + FULL_ATTENTION_INTERVAL = "{arch}.full_attention_interval" + NUM_LOOPS = "{arch}.num_loops" + SKIP_LOOP_FINAL_NORM = "{arch}.skip_loop_final_norm" + HASH_LAYER_COUNT = "{arch}.hash_layer_count" + ACTIVATION_SPARSITY_SCALE = "{arch}.activation_sparsity_scale" + ALTUP_ACTIVE_IDX = "{arch}.altup.active_idx" + ALTUP_NUM_INPUTS = "{arch}.altup.num_inputs" + EMBD_LENGTH_PER_LAYER_INP = "{arch}.embedding_length_per_layer_input" + SWIGLU_CLAMP_EXP = "{arch}.swiglu_clamp_exp" + SWIGLU_CLAMP_SHEXP = "{arch}.swiglu_clamp_shexp" + HIDDEN_ACT = "{arch}.hidden_activation" + DENSE_FEAT_IN_SIZE = "{arch}.{dense}_feat_in" + DENSE_FEAT_OUT_SIZE = "{arch}.{dense}_feat_out" + TARGET_LAYERS = "{arch}.target_layers" + TARGET_HIDDEN_SIZE = "{arch}.target_hidden_size" + BLOCK_SIZE = "{arch}.block_size" + CONV_KERNEL_SIZE = "{arch}.conv_kernel_size" + CONV_GROUP_SIZE = "{arch}.conv_group_size" + SELECTOR_RANK = "{arch}.selector_rank" + SELECTOR_TOP_K = "{arch}.selector_top_k" + SAMPLE_FROM_ANCHOR = "{arch}.sample_from_anchor" + HAS_CONFIDENCE_HEAD = "{arch}.has_confidence_head" + NORM_BEFORE_RESIDUAL = "{arch}.norm_before_residual" + NORM_BEFORE_FC = "{arch}.norm_before_fc" + + class Adapters: + COUNT = "{arch}.adapters.count" + TOKEN_IDS_ACTIVATE = "{arch}.adapters.token_ids_activate" + TOKEN_IDS_SUBSTITUTE = "{arch}.adapters.token_ids_substitute" + LORA_RANK = "{arch}.adapters.lora_rank" + ROUTER_GAIN = "{arch}.adapters.router_gain" + + class Attention: + HEAD_COUNT = "{arch}.attention.head_count" + HEAD_COUNT_KV = "{arch}.attention.head_count_kv" + MAX_ALIBI_BIAS = "{arch}.attention.max_alibi_bias" + CLAMP_KQV = "{arch}.attention.clamp_kqv" + KEY_LENGTH = "{arch}.attention.key_length" + VALUE_LENGTH = "{arch}.attention.value_length" + LAYERNORM_EPS = "{arch}.attention.layer_norm_epsilon" + LAYERNORM_RMS_EPS = "{arch}.attention.layer_norm_rms_epsilon" + GROUPNORM_EPS = "{arch}.attention.group_norm_epsilon" + GROUPNORM_GROUPS = "{arch}.attention.group_norm_groups" + CAUSAL = "{arch}.attention.causal" + Q_LORA_RANK = "{arch}.attention.q_lora_rank" + KV_LORA_RANK = "{arch}.attention.kv_lora_rank" + DECAY_LORA_RANK = "{arch}.attention.decay_lora_rank" + ICLR_LORA_RANK = "{arch}.attention.iclr_lora_rank" + VALUE_RESIDUAL_MIX_LORA_RANK = "{arch}.attention.value_residual_mix_lora_rank" + GATE_LORA_RANK = "{arch}.attention.gate_lora_rank" + REL_BUCKETS_COUNT = "{arch}.attention.relative_buckets_count" + SLIDING_WINDOW = "{arch}.attention.sliding_window" + SCALE = "{arch}.attention.scale" + OUTPUT_GROUP_COUNT = "{arch}.attention.output_group_count" + OUTPUT_LORA_RANK = "{arch}.attention.output_lora_rank" + OUTPUT_SCALE = "{arch}.attention.output_scale" + VALUE_SCALE = "{arch}.attention.value_scale" + COMPRESS_RATIOS = "{arch}.attention.compress_ratios" + COMPRESS_ROPE_FREQ_BASE = "{arch}.attention.compress_rope_freq_base" + TEMPERATURE_LENGTH = "{arch}.attention.temperature_length" + KEY_LENGTH_MLA = "{arch}.attention.key_length_mla" + VALUE_LENGTH_MLA = "{arch}.attention.value_length_mla" + KEY_LENGTH_SWA = "{arch}.attention.key_length_swa" + VALUE_LENGTH_SWA = "{arch}.attention.value_length_swa" + KEY_LENGTH_MLA_SWA = "{arch}.attention.key_length_mla_swa" + VALUE_LENGTH_MLA_SWA = "{arch}.attention.value_length_mla_swa" + KV_LORA_RANK_SWA = "{arch}.attention.kv_lora_rank_swa" + SHARED_KV_LAYERS = "{arch}.attention.shared_kv_layers" + SLIDING_WINDOW_PATTERN = "{arch}.attention.sliding_window_pattern" + RECURRENT_LAYERS = "{arch}.attention.recurrent_layers" + TEMPERATURE_SCALE = "{arch}.attention.temperature_scale" + ROPE_PATTERN = "{arch}.attention.rope_pattern" + VALUE_EXPERT_COUNT = "{arch}.attention.value_expert_count" + VALUE_EXPERT_USED_COUNT = "{arch}.attention.value_expert_used_count" + + class Indexer: + HEAD_COUNT = "{arch}.attention.indexer.head_count" + KEY_LENGTH = "{arch}.attention.indexer.key_length" + TOP_K = "{arch}.attention.indexer.top_k" + BLOCK_SIZE = "{arch}.attention.indexer.block_size" # MSA + LOCAL_BLOCKS = "{arch}.attention.indexer.local_blocks" # MSA + TYPES = "{arch}.attention.indexer.types" + KPOOL = "{arch}.attention.indexer.kpool" # GLM5-Next + KPOOL_SELECT_TAIL = "{arch}.attention.indexer.kpool_select_tail" # GLM5-Next + + class HyperConnection: + COUNT = "{arch}.hyper_connection.count" + SINKHORN_ITERATIONS = "{arch}.hyper_connection.sinkhorn_iterations" + EPSILON = "{arch}.hyper_connection.epsilon" + # scale of the post gate (DeepSeek-V4 hardcodes 2.0) + MAGNITUDE = "{arch}.hyper_connection.magnitude" + # absent means the mix projection is full rank (DeepSeek-V4 behaviour) + LOW_RANK = "{arch}.hyper_connection.low_rank" + + class PerLayerEmbedding: + LAYERS = "{arch}.ple.layers" + NGRAM_SIZE = "{arch}.ple.ngram_size" + HEADS_PER_NGRAM = "{arch}.ple.heads_per_ngram" + CONV_KERNEL = "{arch}.ple.conv_kernel" + LAYER_MULTIPLIERS = "{arch}.ple.layer_multipliers" + HEAD_OFFSETS = "{arch}.ple.head_offsets" + HEAD_VOCAB_SIZES = "{arch}.ple.head_vocab_sizes" + EOS_TOKEN_ID = "{arch}.ple.eos_token_id" + IMAGE_TOKEN_ID = "{arch}.ple.image_token_id" + + class Rope: + DIMENSION_COUNT = "{arch}.rope.dimension_count" + DIMENSION_COUNT_SWA = "{arch}.rope.dimension_count_swa" + DIMENSION_SECTIONS = "{arch}.rope.dimension_sections" + SECTION_ORDER = "{arch}.rope.section_order" + FREQ_BASE = "{arch}.rope.freq_base" + FREQ_BASE_SWA = "{arch}.rope.freq_base_swa" + SCALING_TYPE = "{arch}.rope.scaling.type" + SCALING_FACTOR = "{arch}.rope.scaling.factor" + SCALING_ALPHA = "{arch}.rope.scaling.alpha" + SCALING_ATTN_FACTOR = "{arch}.rope.scaling.attn_factor" + SCALING_ORIG_CTX_LEN = "{arch}.rope.scaling.original_context_length" + SCALING_FINETUNED = "{arch}.rope.scaling.finetuned" + SCALING_YARN_TRUNCATE = "{arch}.rope.scaling.yarn_truncate" + SCALING_YARN_LOG_MUL = "{arch}.rope.scaling.yarn_log_multiplier" + SCALING_YARN_EXT_FACTOR = "{arch}.rope.scaling.yarn_ext_factor" + SCALING_YARN_ATTN_FACTOR = "{arch}.rope.scaling.yarn_attn_factor" + SCALING_YARN_BETA_FAST = "{arch}.rope.scaling.yarn_beta_fast" + SCALING_YARN_BETA_SLOW = "{arch}.rope.scaling.yarn_beta_slow" + + class Activation: + SITU_BETA = "{arch}.activation.situ_beta" + SITU_LINEAR_BETA = "{arch}.activation.situ_linear_beta" + + class AttnRes: + BLOCK_SIZE = "{arch}.attn_res.block_size" + + class Split: + LLM_KV_SPLIT_NO = "split.no" + LLM_KV_SPLIT_COUNT = "split.count" + LLM_KV_SPLIT_TENSORS_COUNT = "split.tensors.count" + + class HRM: + LAYERS_PER_STACK = "{arch}.hrm.layers_per_stack" + H_CYCLES = "{arch}.hrm.h_cycles" + L_CYCLES = "{arch}.hrm.l_cycles" + PREFIX_LM = "{arch}.hrm.prefix_lm" + + class SSM: + CONV_KERNEL = "{arch}.ssm.conv_kernel" + INNER_SIZE = "{arch}.ssm.inner_size" + STATE_SIZE = "{arch}.ssm.state_size" + TIME_STEP_RANK = "{arch}.ssm.time_step_rank" + GROUP_COUNT = "{arch}.ssm.group_count" + DT_B_C_RMS = "{arch}.ssm.dt_b_c_rms" + + class KDA: + HEAD_DIM = "{arch}.kda.head_dim" + SAFE_GATE = "{arch}.kda.safe_gate" + GATE_LOWER_BOUND = "{arch}.kda.gate_lower_bound" + + class WKV: + HEAD_SIZE = "{arch}.wkv.head_size" + + class PosNet: + EMBEDDING_LENGTH = "{arch}.posnet.embedding_length" + BLOCK_COUNT = "{arch}.posnet.block_count" + + class ConvNext: + EMBEDDING_LENGTH = "{arch}.convnext.embedding_length" + BLOCK_COUNT = "{arch}.convnext.block_count" + + class Classifier: + OUTPUT_LABELS = "{arch}.classifier.output_labels" + POOLING_TYPE = "{arch}.classifier.pooling_type" + ACTIVATION = "{arch}.classifier.activation" + + class ShortConv: + L_CACHE = "{arch}.shortconv.l_cache" + + class Decision: + TYPE = "{arch}.decision.type" + # note: single-use-case keys can be hard-coded in cpp code + BLOCK_COUNT = "{arch}.decision.block_count" + ROUTING_BLOCK_COUNT = "{arch}.decision.routing_block_count" + HEAD_COUNT = "{arch}.decision.head_count" + MAX_HEAD_TOKENS = "{arch}.decision.max_head_tokens" + TEMPERATURE = "{arch}.decision.temperature.{name}" # name: "" or "." + SCORE_BIAS = "{arch}.decision.score_bias.{n}.{i}" # added to the logit of level i of a score question with n levels + + class Tokenizer: + MODEL = "tokenizer.ggml.model" + PRE = "tokenizer.ggml.pre" + LIST = "tokenizer.ggml.tokens" + TOKEN_TYPE = "tokenizer.ggml.token_type" + TOKEN_TYPE_COUNT = "tokenizer.ggml.token_type_count" # for BERT-style token types + SCORES = "tokenizer.ggml.scores" + MERGES = "tokenizer.ggml.merges" + BOS_ID = "tokenizer.ggml.bos_token_id" + EOS_ID = "tokenizer.ggml.eos_token_id" + EOT_ID = "tokenizer.ggml.eot_token_id" + EOM_ID = "tokenizer.ggml.eom_token_id" + UNK_ID = "tokenizer.ggml.unknown_token_id" + SEP_ID = "tokenizer.ggml.seperator_token_id" + PAD_ID = "tokenizer.ggml.padding_token_id" + MASK_ID = "tokenizer.ggml.mask_token_id" + ADD_BOS = "tokenizer.ggml.add_bos_token" + ADD_EOS = "tokenizer.ggml.add_eos_token" + ADD_SEP = "tokenizer.ggml.add_sep_token" + ADD_PREFIX = "tokenizer.ggml.add_space_prefix" + REMOVE_EXTRA_WS = "tokenizer.ggml.remove_extra_whitespaces" + PRECOMPILED_CHARSMAP = "tokenizer.ggml.precompiled_charsmap" + SUPPRESS_TOKENS = "tokenizer.ggml.suppress_tokens" + HF_JSON = "tokenizer.huggingface.json" + RWKV = "tokenizer.rwkv.world" + CHAT_TEMPLATE = "tokenizer.chat_template" + CHAT_TEMPLATE_N = "tokenizer.chat_template.{name}" + CHAT_TEMPLATES = "tokenizer.chat_templates" + # Normalizer constants + NORMALIZER_NFC = "tokenizer.ggml.normalizer.nfc" + NORMALIZER_LOWERCASE = "tokenizer.ggml.normalizer.lowercase" + NORMALIZER_STRIP_ACCENTS = "tokenizer.ggml.normalizer.strip_accents" + # FIM/Infill special tokens constants + FIM_PRE_ID = "tokenizer.ggml.fim_pre_token_id" + FIM_SUF_ID = "tokenizer.ggml.fim_suf_token_id" + FIM_MID_ID = "tokenizer.ggml.fim_mid_token_id" + FIM_PAD_ID = "tokenizer.ggml.fim_pad_token_id" + FIM_REP_ID = "tokenizer.ggml.fim_rep_token_id" + FIM_SEP_ID = "tokenizer.ggml.fim_sep_token_id" + # deprecated: + PREFIX_ID = "tokenizer.ggml.prefix_token_id" + SUFFIX_ID = "tokenizer.ggml.suffix_token_id" + MIDDLE_ID = "tokenizer.ggml.middle_token_id" + + class Adapter: + TYPE = "adapter.type" + LORA_ALPHA = "adapter.lora.alpha" + LORA_TASK_NAME = "adapter.lora.task_name" + LORA_PROMPT_PREFIX = "adapter.lora.prompt_prefix" + ALORA_INVOCATION_TOKENS = "adapter.alora.invocation_tokens" + + class IMatrix: + CHUNK_COUNT = "imatrix.chunk_count" + CHUNK_SIZE = "imatrix.chunk_size" + DATASETS = "imatrix.datasets" + + class Clip: + PROJECTOR_TYPE = "clip.projector_type" + HAS_VISION_ENCODER = "clip.has_vision_encoder" + HAS_AUDIO_ENCODER = "clip.has_audio_encoder" + HAS_GEN_AUDIO_ENCODER = "clip.has_gen_audio_encoder" + HAS_LLAVA_PROJECTOR = "clip.has_llava_projector" + + class ClipVision: + PROJECTOR_TYPE = "clip.vision.projector_type" # for mixed modality models + IMAGE_SIZE = "clip.vision.image_size" + IMAGE_MIN_PIXELS = "clip.vision.image_min_pixels" + IMAGE_MAX_PIXELS = "clip.vision.image_max_pixels" + PREPROC_MIN_TILES = "clip.vision.preproc_min_tiles" + PREPROC_MAX_TILES = "clip.vision.preproc_max_tiles" + PREPROC_IMAGE_SIZE = "clip.vision.preproc_image_size" + PATCH_SIZE = "clip.vision.patch_size" + EMBEDDING_LENGTH = "clip.vision.embedding_length" + FEED_FORWARD_LENGTH = "clip.vision.feed_forward_length" + PROJECTION_DIM = "clip.vision.projection_dim" + BLOCK_COUNT = "clip.vision.block_count" + IMAGE_MEAN = "clip.vision.image_mean" + IMAGE_STD = "clip.vision.image_std" + MAX_SLICE_NUMS = "clip.vision.max_slice_nums" + IMAGE_RESIZE_ALGO = "clip.vision.image_resize_algo" + SPATIAL_MERGE_SIZE = "clip.vision.spatial_merge_size" + SWIGLU_CLAMP = "clip.vision.swiglu_clamp" + EXPERT_COUNT_PER_LAYER = "clip.vision.expert_count_per_layer" # dots3note pyramid MoE, 0 = dense layer + EXPERT_USED_COUNT = "clip.vision.expert_used_count" + USE_GELU = "clip.use_gelu" + USE_SILU = "clip.use_silu" + N_WA_PATTERN = "clip.vision.n_wa_pattern" # used by qwen2.5vl + WA_LAYER_INDEXES = "clip.vision.wa_layer_indexes" # used by youtuvl + WA_PATTERN_MODE = "clip.vision.wa_pattern_mode" # used by mimovl, per-layer -1/0/1 + IS_DEEPSTACK_LAYERS = "clip.vision.is_deepstack_layers" + WINDOW_SIZE = "clip.vision.window_size" + FEATURE_LAYERS = "clip.vision.feature_layer" # Granite4 Vision + IMAGE_GRID_PINPOINTS = "clip.vision.image_grid_pinpoints" # Granite4 Vision + + class Attention: + HEAD_COUNT = "clip.vision.attention.head_count" + HEAD_COUNT_KV = "clip.vision.attention.head_count_kv" # used by mimovl (GQA) + HEAD_DIM = "clip.vision.attention.head_dim" # set when qkv width != n_embd + LAYERNORM_EPS = "clip.vision.attention.layer_norm_epsilon" + + class Projector: + SCALE_FACTOR = "clip.vision.projector.scale_factor" + QUERY_SIDE = "clip.vision.projector.query_side" + WINDOW_SIDE = "clip.vision.projector.window_side" + SPATIAL_OFFSETS = "clip.vision.projector.spatial_offsets" + + class SAM: + BLOCK_COUNT = "clip.vision.sam.block_count" + EMBEDDING_LENGTH = "clip.vision.sam.embedding_length" + HEAD_COUNT = "clip.vision.sam.head_count" + + class ClipAudio: + PROJECTOR_TYPE = "clip.audio.projector_type" # for mixed modality models + NUM_MEL_BINS = "clip.audio.num_mel_bins" + EMBEDDING_LENGTH = "clip.audio.embedding_length" + FEED_FORWARD_LENGTH = "clip.audio.feed_forward_length" + PROJECTION_DIM = "clip.audio.projection_dim" + BLOCK_COUNT = "clip.audio.block_count" + SUBSAMPLING_FACTOR = "clip.audio.subsampling_factor" + CHUNK_SIZE = "clip.audio.chunk_size" + CONV_KERNEL_SIZE = "clip.audio.conv_kernel_size" + MAX_POS_EMB = "clip.audio.max_pos_emb" + FEATURE_LAYERS = "clip.audio.feature_layer" # Granite Speech Plus + RVQ_NUM_QUANTIZERS = "clip.audio.rvq.num_quantizers" + RVQ_CODEBOOK_SIZE = "clip.audio.rvq.codebook_size" + WA_PATTERN_MODE = "clip.audio.wa_pattern_mode" # per-layer -1 (full) / 0 (windowed) + WINDOW_SIZE = "clip.audio.window_size" + LOCAL_BLOCK_COUNT = "clip.audio.local_block_count" # mimo-v2.5: input_local_transformer layer count + LOCAL_GROUP_SIZE = "clip.audio.local_group_size" # mimo-v2.5: input_local_transformer grouping size + + class Attention: + HEAD_COUNT = "clip.audio.attention.head_count" + LAYERNORM_EPS = "clip.audio.attention.layer_norm_epsilon" + + class Projector: + STACK_FACTOR = "clip.audio.projector.stack_factor" + WINDOW_SIZE = "clip.audio.projector.window_size" + DOWNSAMPLE_RATE = "clip.audio.projector.downsample_rate" + HEAD_COUNT = "clip.audio.projector.head_count" + + class ClipGenAudio: + PROJECTOR_TYPE = "clip.gen.audio.projector_type" # for mixed modality models + # name of the weight variant, for settings that are not in the checkpoint + MODEL_VARIANT = "clip.gen.audio.model_variant" + EMBEDDING_LENGTH = "clip.gen.audio.embedding_length" + FEED_FORWARD_LENGTH = "clip.gen.audio.feed_forward_length" + BLOCK_COUNT = "clip.gen.audio.block_count" + PROJECTION_DIM = "clip.gen.audio.projection_dim" + + class Attention: + HEAD_COUNT = "clip.gen.audio.attention.head_count" + HEAD_COUNT_KV = "clip.gen.audio.attention.head_count_kv" + LAYERNORM_EPS = "clip.gen.audio.attention.layer_norm_epsilon" + + class Diffusion: + SHIFT_LOGITS = "diffusion.shift_logits" + + class xIELU: + ALPHA_P = "xielu.alpha_p" + ALPHA_N = "xielu.alpha_n" + BETA = "xielu.beta" + EPS = "xielu.eps" + + class PrismHadamard: + VERSION = "prism.hadamard.version" + TIED_OUTPUT = "prism.hadamard.tied_output" + BLOCK_SIZE = "prism.hadamard.block_size" + TRANSFORM = "prism.hadamard.transform" + AXIS = "prism.hadamard.axis" + SIGN_MODE = "prism.hadamard.sign_mode" + SIGN_WIDTHS = "prism.hadamard.sign_widths" + SIGN_VALUES = "prism.hadamard.sign_values" + WEIGHT_NAMES = "prism.hadamard.weight_names" + INVERSE_WEIGHT_NAMES = "prism.hadamard.inverse_weight_names" + GDN_V_GROUPED = "prism.hadamard.gdn_v_grouped" + + +# +# recommended mapping of model tensor names for storage in gguf +# + + +class GGUFType: + MODEL = "model" + ADAPTER = "adapter" + IMATRIX = "imatrix" + MMPROJ = "mmproj" # dummy, unused for now + + +class MODEL_ARCH(IntEnum): + MMPROJ = auto() # dummy arch for clip.cpp + LLAMA = auto() + LLAMA4 = auto() + DECI = auto() + FALCON = auto() + FALCON_H1 = auto() + BAICHUAN = auto() + GROK = auto() + GPT2 = auto() + GPTJ = auto() + GPTNEOX = auto() + MPT = auto() + STARCODER = auto() + REFACT = auto() + BERT = auto() + MODERN_BERT = auto() + NOMIC_BERT = auto() + NOMIC_BERT_MOE = auto() + NEO_BERT = auto() + JINA_BERT_V2 = auto() + JINA_BERT_V3 = auto() + EUROBERT = auto() + BLOOM = auto() + STABLELM = auto() + QWEN = auto() + QWEN2 = auto() + QWEN2MOE = auto() + QWEN2VL = auto() + QWEN3 = auto() + QWEN3MOE = auto() + QWEN3NEXT = auto() + HRM_TEXT = auto() + QWEN3VL = auto() + QWEN3VLMOE = auto() + QWEN35 = auto() + QWEN35MOE = auto() + CLEF = auto() + VLLM_DECISION2 = auto() + QWEN4EXP = auto() + PHI2 = auto() + PHI3 = auto() + PHIMOE = auto() + PLAMO = auto() + PLAMO2 = auto() + PLAMO3 = auto() + CODESHELL = auto() + ORION = auto() + INTERNLM2 = auto() + MINICPM = auto() + MINICPM3 = auto() + GEMMA = auto() + GEMMA2 = auto() + GEMMA3 = auto() + GEMMA3N = auto() + GEMMA4 = auto() + GEMMA4_ASSISTANT = auto() + GEMMA_EMBEDDING = auto() + GEMMA_EMBEDDING2 = auto() + STARCODER2 = auto() + RWKV6 = auto() + RWKV6QWEN2 = auto() + RWKV7 = auto() + ARWKV7 = auto() + MAMBA = auto() + MAMBA2 = auto() + MAPLE = auto() + JAMBA = auto() + XVERSE = auto() + COMMAND_R = auto() + COHERE2 = auto() + COHERE2MOE = auto() + DBRX = auto() + OLMO = auto() + OLMO2 = auto() + OLMOE = auto() + MUSE_GLIMMER = auto() + OPENELM = auto() + ARCTIC = auto() + DEEPSEEK = auto() + DEEPSEEK2 = auto() + DEEPSEEK2OCR = auto() + DEEPSEEK32 = auto() + DEEPSEEK4 = auto() + CHATGLM = auto() + GLM4 = auto() + GLM4_MOE = auto() + GLM_DSA = auto() + GLM5_NEXT = auto() + BITNET = auto() + T5 = auto() + T5ENCODER = auto() + JAIS = auto() + JAIS2 = auto() + NEMOTRON = auto() + NEMOTRON_H = auto() + NEMOTRON_H_MOE = auto() + EXAONE = auto() + EXAONE4 = auto() + EXAONE_MOE = auto() + GRANITE = auto() + GRANITE_MOE = auto() + GRANITE_HYBRID = auto() + GRANITE_SWITCH = auto() + GRANITE_SWA = auto() + CHAMELEON = auto() + WAVTOKENIZER_DEC = auto() + PLM = auto() + BAILINGMOE = auto() + BAILINGMOE2 = auto() + BAILINGMOE3 = auto() + DOTS1 = auto() + DOTS3NOTE = auto() + ARCEE = auto() + AFMOE = auto() + LAGUNA = auto() + ERNIE4_5 = auto() + ERNIE4_5_MOE = auto() + HUNYUAN_MOE = auto() + HUNYUAN_DENSE = auto() + HUNYUAN_VL = auto() + HY_V3 = auto() + HY_V4 = auto() + SMOLLM3 = auto() + GPT_OSS = auto() + LFM2 = auto() + LFM2MOE = auto() + DREAM = auto() + SMALLTHINKER = auto() + LLADA = auto() + LLADA_MOE = auto() + SEED_OSS = auto() + GROVEMOE = auto() + APERTUS = auto() + COGVLM = auto() + MINIMAX01 = auto() + MINIMAXM2 = auto() + MINIMAXM3 = auto() + RND1 = auto() + PANGU_EMBED = auto() + MISTRAL3 = auto() + EAGLE3 = auto() + DFLASH = auto() + MISTRAL4 = auto() + PADDLEOCR = auto() + MIMO2 = auto() + STEP35 = auto() + SPARK2_5 = auto() + LLAMA_EMBED = auto() + MAINCODER = auto() + KIMI_LINEAR = auto() + KIMI_K3 = auto() + TALKIE = auto() + MELLUM = auto() + NANBEIGE = auto() + QWEN3TTS = auto() + POCKETTTS = auto() + K2HORIZON = auto() + + +class VISION_PROJECTOR_TYPE(IntEnum): + MLP = auto() + LDP = auto() + LDPV2 = auto() + RESAMPLER = auto() + GLM_EDGE = auto() + MERGER = auto() + GEMMA3N = auto() + GEMMA3 = auto() + QWEN3VL = auto() + LING3VL = auto() + STEP3VL = auto() + COGVLM = auto() + + +class MODEL_TENSOR(IntEnum): + TOKEN_EMBD = auto() + TOKEN_EMBD_NORM = auto() + MASKED_EMBD_CENTROIDS= auto() + MASKED_EMBD_ORDERING = auto() + TOKEN_TYPES = auto() + POS_EMBD = auto() + OUTPUT = auto() + HRM_Z_L_INIT = auto() + DENSE_2_OUT = auto() # embeddinggemma 2_Dense + DENSE_3_OUT = auto() # embeddinggemma 3_Dense + OUTPUT_NORM = auto() + HC_HEAD_FN = auto() + HC_HEAD_BASE = auto() + HC_HEAD_SCALE = auto() + HC_HEAD_NORM = auto() # qwen4exp + HC_HEAD_DOWN = auto() # qwen4exp + HC_HEAD_UP = auto() # qwen4exp + ROPE_FREQS = auto() + ROPE_FACTORS_LONG = auto() + ROPE_FACTORS_SHORT = auto() + ATTN_Q = auto() + ATTN_K = auto() + ATTN_V = auto() + ATTN_QKV = auto() + ATTN_OUT = auto() + ATTN_NORM = auto() + ATTN_NORM_2 = auto() + ATTN_OUT_NORM = auto() + ATTN_POST_NORM = auto() + ATTN_ROT_EMBD = auto() + ATTN_SINKS = auto() + ATTN_GATE = auto() + FFN_GATE_INP = auto() + FFN_GATE_INP_SHEXP = auto() + FFN_NORM = auto() + FFN_PRE_NORM = auto() # alias of FFN_NORM + FFN_PRE_NORM_2 = auto() # gemma4 + FFN_POST_NORM = auto() + FFN_POST_NORM_1 = auto() # gemma4 + FFN_POST_NORM_2 = auto() # gemma4 + FFN_GATE = auto() + FFN_DOWN = auto() + FFN_UP = auto() + FFN_ACT = auto() + FFN_NORM_EXP = auto() + FFN_GATE_EXP = auto() + FFN_DOWN_EXP = auto() + FFN_UP_EXP = auto() + FFN_GATE_UP_EXP = auto() + FFN_GATE_SHEXP = auto() + FFN_DOWN_SHEXP = auto() + FFN_UP_SHEXP = auto() + FFN_GATE_CHEXP = auto() + FFN_DOWN_CHEXP = auto() + FFN_UP_CHEXP = auto() + FFN_EXP_PROBS_B = auto() + FFN_EXP_PROBS_B_VL = auto() # deepseek4 vision (bias for image tokens) + FFN_GATE_TID2EID = auto() + MOE_LATENT_DOWN = auto() # nemotron 3 super + MOE_LATENT_UP = auto() # nemotron 3 super + ATTN_Q_NORM = auto() + ATTN_K_NORM = auto() + LAYER_OUT_NORM = auto() + LAYER_OUT_SCALE = auto() + PER_LAYER_TOKEN_EMBD = auto() # gemma3n + PER_LAYER_MODEL_PROJ = auto() # gemma3n + PER_LAYER_INP_GATE = auto() # gemma3n + PER_LAYER_PROJ = auto() # gemma3n + PER_LAYER_PROJ_NORM = auto() # gemma3n + PER_LAYER_POST_NORM = auto() # gemma3n + ALTUP_PROJ = auto() # gemma3n + ALTUP_UNEMBD_PROJ = auto() # gemma3n + ALTUP_CORRECT_COEF = auto() # gemma3n + ALTUP_CORRECT_SCALE = auto() # gemma3n + ALTUP_PREDICT_COEF = auto() # gemma3n + ALTUP_ROUTER = auto() # gemma3n + ALTUP_ROUTER_NORM = auto() # gemma3n + LAUREL_L = auto() # gemma3n + LAUREL_R = auto() # gemma3n + LAUREL_POST_NORM = auto() # gemma3n + SSM_IN = auto() + SSM_CONV1D = auto() + SSM_X = auto() + SSM_DT = auto() + SSM_DT_NORM = auto() + SSM_A = auto() + SSM_B_NORM = auto() + SSM_C_NORM = auto() + SSM_D = auto() + SSM_NORM = auto() + SSM_OUT = auto() + SSM_ALPHA = auto() # qwen3.5 + SSM_BETA_ALPHA = auto() # qwen3next + SSM_CONV1D_Q = auto() # Kimi Linear + SSM_CONV1D_K = auto() # Kimi Linear + SSM_CONV1D_V = auto() # Kimi Linear + SSM_F_A = auto() # Kimi Linear + SSM_F_B = auto() # Kimi Linear + SSM_BETA = auto() # Kimi Linear qwen3.5 + SSM_G_A = auto() # Kimi Linear + SSM_G_B = auto() # Kimi Linear + SSM_G = auto() # Kimi K3 (full-rank KDA gate, replaces SSM_G_A/SSM_G_B) + ATTN_RES_SCORE = auto() # Kimi K3 (fused res_norm * res_proj, pre-attention) + FFN_RES_SCORE = auto() # Kimi K3 (fused res_norm * res_proj, pre-FFN) + OUTPUT_RES_SCORE = auto() # Kimi K3 (fused res_norm * res_proj, final) + FFN_ROUTED_DOWN = auto() # Kimi K3 (latent MoE: hidden -> latent) + FFN_ROUTED_UP = auto() # Kimi K3 (latent MoE: latent -> hidden) + FFN_ROUTED_NORM = auto() # Kimi K3 (latent MoE: norm on expert output) + TIME_MIX_W0 = auto() + TIME_MIX_W1 = auto() + TIME_MIX_W2 = auto() + TIME_MIX_A0 = auto() + TIME_MIX_A1 = auto() + TIME_MIX_A2 = auto() + TIME_MIX_V0 = auto() + TIME_MIX_V1 = auto() + TIME_MIX_V2 = auto() + TIME_MIX_G1 = auto() + TIME_MIX_G2 = auto() + TIME_MIX_K_K = auto() + TIME_MIX_K_A = auto() + TIME_MIX_R_K = auto() + TIME_MIX_LERP_X = auto() + TIME_MIX_LERP_K = auto() + TIME_MIX_LERP_V = auto() + TIME_MIX_LERP_R = auto() + TIME_MIX_LERP_G = auto() + TIME_MIX_LERP_FUSED = auto() + TIME_MIX_LERP_W = auto() + TIME_MIX_FIRST = auto() + TIME_MIX_DECAY = auto() + TIME_MIX_DECAY_W1 = auto() + TIME_MIX_DECAY_W2 = auto() + TIME_MIX_KEY = auto() + TIME_MIX_VALUE = auto() + TIME_MIX_RECEPTANCE = auto() + TIME_MIX_GATE = auto() + TIME_MIX_LN = auto() + TIME_MIX_OUTPUT = auto() + CHANNEL_MIX_LERP_K = auto() + CHANNEL_MIX_LERP_R = auto() + CHANNEL_MIX_KEY = auto() + CHANNEL_MIX_RECEPTANCE = auto() + CHANNEL_MIX_VALUE = auto() + ATTN_Q_A = auto() + ATTN_Q_B = auto() + ATTN_KV_A_MQA = auto() + ATTN_KV_B = auto() + ATTN_K_B = auto() + ATTN_V_B = auto() + ATTN_Q_A_NORM = auto() + ATTN_KV_A_NORM = auto() + ATTN_KV = auto() + ATTN_KV_NORM = auto() + ATTN_OUT_A = auto() + ATTN_OUT_B = auto() + HC_ATTN_FN = auto() + HC_ATTN_BASE = auto() + HC_ATTN_SCALE = auto() + HC_FFN_FN = auto() + HC_FFN_BASE = auto() + HC_FFN_SCALE = auto() + HC_ATTN_NORM = auto() # qwen4exp + HC_ATTN_DOWN = auto() # qwen4exp + HC_ATTN_UP = auto() # qwen4exp + HC_ATTN_INJECT = auto() # qwen4exp + HC_FFN_NORM = auto() # qwen4exp + HC_FFN_DOWN = auto() # qwen4exp + HC_FFN_UP = auto() # qwen4exp + HC_FFN_INJECT = auto() # qwen4exp + PLE_KEY = auto() # qwen4exp + PLE_VALUE = auto() # qwen4exp + PLE_NORM_KEY = auto() # qwen4exp + PLE_NORM_QUERY = auto() # qwen4exp + PLE_NORM_CONV = auto() # qwen4exp + PLE_CONV1D = auto() # qwen4exp + ATTN_COMPRESSOR_WKV = auto() + ATTN_COMPRESSOR_WGATE = auto() + ATTN_COMPRESSOR_APE = auto() + ATTN_COMPRESSOR_NORM = auto() + FFN_SUB_NORM = auto() + ATTN_SUB_NORM = auto() + DEC_ATTN_NORM = auto() + DEC_ATTN_Q = auto() + DEC_ATTN_K = auto() + DEC_ATTN_V = auto() + DEC_ATTN_OUT = auto() + DEC_ATTN_REL_B = auto() + DEC_CROSS_ATTN_NORM = auto() + DEC_CROSS_ATTN_NORM_KV = auto() + DEC_CROSS_ATTN_Q = auto() + DEC_CROSS_ATTN_K = auto() + DEC_CROSS_ATTN_V = auto() + DEC_CROSS_ATTN_OUT = auto() + DEC_CROSS_ATTN_REL_B = auto() + DEC_FFN_NORM = auto() + DEC_FFN_GATE = auto() + DEC_FFN_DOWN = auto() + DEC_FFN_UP = auto() + DEC_OUTPUT_NORM = auto() + ENC_ATTN_NORM = auto() + ENC_ATTN_Q = auto() + ENC_ATTN_K = auto() + ENC_ATTN_V = auto() + ENC_ATTN_OUT = auto() + ENC_ATTN_REL_B = auto() + ENC_FFN_NORM = auto() + ENC_FFN_GATE = auto() + ENC_FFN_DOWN = auto() + ENC_FFN_UP = auto() + ENC_OUTPUT_NORM = auto() + CLS = auto() # classifier + CLS_OUT = auto() # classifier output projection + CLS_NORM = auto() + DECISION_HIDDEN_NORM = auto() + DECISION_PROJ_MEMORY = auto() + DECISION_PROJ_QUESTION = auto() + DECISION_PROJ_OPTION_QUESTION = auto() + DECISION_PROJ_GLOBAL = auto() + DECISION_PROJ_OPTION_CONTEXT = auto() + DECISION_PROJ_OPTION_LEXICAL = auto() + DECISION_OPTION_SUMMARY_NORM = auto() + DECISION_FIELD_NORM = auto() + DECISION_OPTION_NORM = auto() + DECISION_SCALES = auto() + DECISION_SCORER = auto() + DECISION_SCORER_OUT = auto() + DECISION_QUERY_NORM = auto() + DECISION_KEY = auto() + DECISION_QUERY = auto() + DECISION_OPTION_MLP = auto() + DECISION_QUERY_MLP = auto() + CONV1D = auto() + CONVNEXT_DW = auto() + CONVNEXT_NORM = auto() + CONVNEXT_PW1 = auto() + CONVNEXT_PW2 = auto() + CONVNEXT_GAMMA = auto() + POSNET_CONV1 = auto() + POSNET_CONV2 = auto() + POSNET_NORM = auto() + POSNET_NORM1 = auto() + POSNET_NORM2 = auto() + POSNET_ATTN_NORM = auto() + POSNET_ATTN_Q = auto() + POSNET_ATTN_K = auto() + POSNET_ATTN_V = auto() + POSNET_ATTN_OUT = auto() + SHORTCONV_CONV = auto() + SHORTCONV_INPROJ = auto() + SHORTCONV_OUTPROJ = auto() + VISEXP_ATTN_QKV = auto() + VISEXP_ATTN_OUT = auto() + VISEXP_GATE = auto() + VISEXP_DOWN = auto() + VISEXP_UP = auto() + INDEXER_K_NORM = auto() + INDEXER_PROJ = auto() + INDEXER_ATTN_K = auto() + INDEXER_ATTN_Q_B = auto() + INDEXER_Q_PROJ = auto() + INDEXER_K_PROJ = auto() + INDEXER_Q_NORM = auto() + INDEXER_COMPRESSOR_WKV = auto() + INDEXER_COMPRESSOR_WGATE = auto() + INDEXER_COMPRESSOR_APE = auto() + INDEXER_COMPRESSOR_NORM = auto() + INDEXER_KPOOL_GATE = auto() + INDEXER_KPOOL_APE = auto() + ATTN_V_GATE = auto() # k2-horizon MoVA router + ATTN_V_EXP = auto() # k2-horizon MoVA value experts + # vision + V_MMPROJ = auto() + V_MMPROJ_FC = auto() + V_MMPROJ_MLP = auto() + V_MMPROJ_PEG = auto() + V_ENC_EMBD_CLS = auto() + V_ENC_EMBD_PATCH = auto() + V_ENC_EMBD_NORM = auto() + V_ENC_EMBD_PATCH_NORM = auto() # allow multiple norms in the same embd, e.g. for gemma4u + V_ENC_EMBD_POS = auto() + V_ENC_INPUT_NORM = auto() + V_ENC_ATTN_QKV = auto() + V_ENC_ATTN_Q = auto() + V_ENC_ATTN_Q_NORM = auto() + V_ENC_ATTN_K = auto() + V_ENC_ATTN_K_NORM = auto() + V_ENC_ATTN_V = auto() + V_ENC_ATTN_O = auto() + V_ENC_ATTN_O_NORM = auto() + V_ENC_ATTN_SINKS = auto() # mimovl + V_ENC_POST_ATTN_NORM = auto() + V_ENC_FFN_UP = auto() + V_ENC_FFN_GATE = auto() + V_ENC_FFN_DOWN = auto() + V_ENC_FFN_GATE_INP = auto() # dots3note vision MoE router + V_ENC_FFN_GATE_EXPS = auto() + V_ENC_FFN_UP_EXPS = auto() + V_ENC_FFN_DOWN_EXPS = auto() + V_ENC_FFN_EXP_PROBS_B = auto() + V_ENC_ATTN_POST_NORM = auto() # gemma4 + V_ENC_FFN_POST_NORM = auto() + V_LAYER_SCALE_1 = auto() + V_LAYER_SCALE_2 = auto() + V_LAYER_OUT_SCALE = auto() + V_PRE_NORM = auto() + V_POST_NORM = auto() + V_MM_PRE_NORM = auto() # hunyuanvl + V_MM_POST_NORM = auto() + V_MM_INP_NORM = auto() + V_MM_INP_PROJ = auto() # gemma3 + V_MM_SOFT_EMB_NORM = auto() # gemma3 + V_MM_EMBEDDING = auto() # gemma3n + V_MM_HARD_EMB_NORM = auto() # gemma3n + V_ENC_CONV_STEM = auto() # gemma3n + V_ENC_CONV_STEM_NORM = auto() # gemma3n + V_ENC_MSFA_EXP = auto() # gemma3n + V_ENC_MSFA_EXP_NORM = auto() # gemma3n + V_ENC_MSFA_PROJ = auto() # gemma3n + V_ENC_MSFA_PROJ_NORM = auto() # gemma3n + V_ENC_MSFA_NORM = auto() # gemma3n + V_RESMPL_POS_EMBD_K = auto() # minicpmv + V_RESMPL_ATTN_Q = auto() # minicpmv + V_RESMPL_ATTN_K = auto() # minicpmv + V_RESMPL_ATTN_V = auto() # minicpmv + V_RESMPL_ATTN_OUT = auto() # minicpmv + V_RESMPL_KV = auto() # minicpmv + V_RESMPL_KV_NORM = auto() # minicpmv + V_RESMPL_POST_NORM = auto() # minicpmv + V_RESMPL_Q_NORM = auto() # minicpmv + V_RESMPL_PROJ = auto() # minicpmv + V_RESMPL_QUERY = auto() # minicpmv + V_TOK_EMBD_IMG_BREAK = auto() # pixtral + V_TOK_EMBD_IMG_START = auto() # deepseek4v + V_TOK_EMBD_IMG_END = auto() # deepseek4v + V_TOK_EMBD_IMG_PAD = auto() # deepseek4v + V_MM_PATCH_MERGER = auto() # mistral small 3.1 + V_DS_NORM = auto() # qwen3vl + V_DS_FC1 = auto() # qwen3vl + V_DS_FC2 = auto() # qwen3vl + V_MERGER_LN1 = auto() # minicpmv4_6 + V_MERGER_ATTN_Q = auto() # minicpmv4_6 + V_MERGER_ATTN_K = auto() # minicpmv4_6 + V_MERGER_ATTN_V = auto() # minicpmv4_6 + V_MERGER_ATTN_O = auto() # minicpmv4_6 + V_MERGER_DS_LN = auto() # minicpmv4_6 + V_MERGER_DS_UP = auto() # minicpmv4_6 + V_MERGER_DS_DOWN = auto() # minicpmv4_6 + V_MM_POST_FC_NORM = auto() # cogvlm + V_MM_UP = auto() # cogvlm + V_MM_DOWN = auto() # cogvlm + V_MM_GATE = auto() # cogvlm + V_MM_MERGER_FC1 = auto() # minimax-m3 (patch-merge MLP) + V_MM_MERGER_FC2 = auto() # minimax-m3 (patch-merge MLP) + V_TOK_BOI = auto() # cogvlm + V_TOK_EOI = auto() # cogvlm + V_TOK_IMG_BEGIN = auto() # hunyuanvl + V_TOK_IMG_END = auto() # hunyuanvl + V_STD_BIAS = auto() # gemma4 + V_STD_SCALE = auto() # gemma4 + V_SAM_POS_EMBD = auto() # Deepseek-OCR + V_SAM_PATCH_EMBD = auto() # Deepseek-OCR + V_SAM_PRE_NORM = auto() # Deepseek-OCR + V_SAM_POST_NORM = auto() # Deepseek-OCR + V_SAM_ATTN_POS_H = auto() # Deepseek-OCR + V_SAM_ATTN_POS_W = auto() # Deepseek-OCR + V_SAM_ATTN_QKV = auto() # Deepseek-OCR + V_SAM_ATTN_OUT = auto() # Deepseek-OCR + V_SAM_MLP_LIN_1 = auto() # Deepseek-OCR + V_SAM_MLP_LIN_2 = auto() # Deepseek-OCR + V_SAM_NECK = auto() # Deepseek-OCR + V_SAM_NET_2 = auto() # Deepseek-OCR + V_SAM_NET_3 = auto() # Deepseek-OCR + V_ENC_EMBD_IMGNL = auto() # Deepseek-OCR + V_ENC_EMBD_VSEP = auto() # Deepseek-OCR + V_TOK_EMBD_SEP = auto() # MiniCPM-V 4.7 + V_RESMPL_QUERY_768 = auto() # Deepseek-OCR-2 + V_RESMPL_QUERY_1024 = auto() # Deepseek-OCR-2 + + # qformer projector (vision) - Granite4 Vision + V_QF_PROJ_QUERY = auto() + V_QF_PROJ_NORM = auto() + V_QF_PROJ_LINEAR = auto() + V_QF_SELF_ATTN_Q = auto() + V_QF_SELF_ATTN_K = auto() + V_QF_SELF_ATTN_V = auto() + V_QF_SELF_ATTN_O = auto() + V_QF_SELF_ATTN_NORM = auto() + V_QF_CROSS_ATTN_Q = auto() + V_QF_CROSS_ATTN_K = auto() + V_QF_CROSS_ATTN_V = auto() + V_QF_CROSS_ATTN_O = auto() + V_QF_CROSS_ATTN_NORM = auto() + V_QF_FFN_UP = auto() + V_QF_FFN_DOWN = auto() + V_QF_FFN_NORM = auto() + V_PROJ_NORM = auto() + # multi-projector (bid => projector id) - Granite4 vision + V_MULTI_PROJ_IMG_POS = auto() + V_MULTI_PROJ_QUERY = auto() + V_MULTI_PROJ_NORM = auto() + V_MULTI_PROJ_LINEAR = auto() + V_MULTI_PROJ_POST_NORM = auto() + + # audio (mtmd) + A_ENC_EMBD_POS = auto() + A_ENC_EMBD_NORM = auto() + A_ENC_EMBD_TO_LOGITS = auto() # lfm2 + A_ENC_INP_PROJ = auto() # gemma4 + A_ENC_CONV1D = auto() + A_ENC_CONV1D_NORM = auto() # gemma3n + A_ENC_CONV2D = auto() + A_ENC_CONV_OUT = auto() + A_PRE_NORM = auto() + A_POST_NORM = auto() + A_ENC_LAYER_PRE_NORM = auto() # gemma3n + A_ENC_ATTN_Q = auto() + A_ENC_ATTN_K = auto() + A_ENC_ATTN_V = auto() + A_ENC_ATTN_POST_NORM = auto() + A_ENC_ATTN_PRE_NORM = auto() + A_ENC_ATTN_K_REL = auto() # gemma4 + A_ENC_PER_DIM_SCALE = auto() # gemma3n + A_ENC_INPUT_NORM = auto() + A_ENC_OUTPUT = auto() # TODO @ngxson: rename to ATTN_OUT + A_ENC_OUTPUT_NORM = auto() # TODO @ngxson: rename to ATTN_OUT + A_ENC_FFN_UP = auto() + A_ENC_FFN_NORM = auto() + A_ENC_FFN_POST_NORM = auto() # gemma3n + A_ENC_FFN_SCALE = auto() # gemma3n + A_ENC_FFN_GATE = auto() + A_ENC_FFN_DOWN = auto() + A_ENC_FFN_UP_1 = auto() # lfm2, gemma3n + A_ENC_FFN_NORM_1 = auto() # lfm2, gemma3n (pre-norm) + A_ENC_FFN_POST_NORM_1 = auto() # gemma3n + A_ENC_FFN_SCALE_1 = auto() # gemma3n + A_ENC_FFN_GATE_1 = auto() # lfm2, gemma3n + A_ENC_FFN_DOWN_1 = auto() # lfm2, gemma3n + A_ENC_DOWNSAMPLE_CONV = auto() # mimo-audio-tokenizer: post-transformer downsample conv + A_ENC_DOWNSAMPLE_NORM = auto() # mimo-audio-tokenizer: post-transformer downsample norm + A_ENC_RVQ_CODEBOOK = auto() # mimo-audio-tokenizer: residual vector quantizer codebook, per quantizer index + A_ENC_CONV_RES2 = auto() # qwen3tts + A_ENC_SE_CONV1 = auto() # qwen3tts + A_ENC_SE_CONV2 = auto() # qwen3tts + A_ENC_ASP_ATTN = auto() # qwen3tts + A_ENC_ASP_TDNN = auto() # qwen3tts + # qwen3tts code_predictor: predicts the remaining RVQ codebooks + A_GEN_CODE_PROJ_IN = auto() # small_to_mtp_projection + A_GEN_CODE_EMBD = auto() # per-codebook embedding table, merged 3D [n_codebooks, vocab, dim] + A_GEN_CODE_HEAD = auto() # per-codebook output head, merged 3D [n_codebooks, vocab, dim] + A_GEN_CODE_OUT_EMBD = auto() # codebook-0 embedding, re-fed into the talker backbone (talker.model.codec_embedding) + A_GEN_CODE_ATTN_NORM = auto() + A_GEN_CODE_ATTN_Q = auto() + A_GEN_CODE_ATTN_Q_NORM = auto() + A_GEN_CODE_ATTN_K = auto() + A_GEN_CODE_ATTN_K_NORM = auto() + A_GEN_CODE_ATTN_V = auto() + A_GEN_CODE_ATTN_OUT = auto() + A_GEN_CODE_FFN_NORM = auto() + A_GEN_CODE_FFN_GATE = auto() + A_GEN_CODE_FFN_UP = auto() + A_GEN_CODE_FFN_DOWN = auto() + A_GEN_CODE_OUTPUT_NORM = auto() + # qwen3tts code2wav: RVQ codes -> raw PCM + A_GEN_WAV_QUANT_FIRST_IN = auto() # semantic RVQ, in_proj (1x1 conv, loaded as 2D) + A_GEN_WAV_QUANT_FIRST_OUT = auto() # semantic RVQ, out_proj + A_GEN_WAV_QUANT_FIRST_CB = auto() # semantic RVQ codebook (1 layer), folded from embedding_sum/cluster_usage + A_GEN_WAV_QUANT_REST_IN = auto() # acoustic RVQ, in_proj + A_GEN_WAV_QUANT_REST_OUT = auto() # acoustic RVQ, out_proj + A_GEN_WAV_QUANT_REST_CB = auto() # acoustic RVQ codebooks, merged 3D [15, vocab, dim] + A_GEN_WAV_PRE_CONV = auto() + A_GEN_WAV_TFM_IN_PROJ = auto() + A_GEN_WAV_TFM_OUT_PROJ = auto() + A_GEN_WAV_TFM_OUTPUT_NORM = auto() + A_GEN_WAV_TFM_ATTN_NORM = auto() + A_GEN_WAV_TFM_ATTN_Q = auto() + A_GEN_WAV_TFM_ATTN_K = auto() + A_GEN_WAV_TFM_ATTN_V = auto() + A_GEN_WAV_TFM_ATTN_OUT = auto() + A_GEN_WAV_TFM_ATTN_SCALE = auto() # layer scale (gamma) on the attn output + A_GEN_WAV_TFM_FFN_NORM = auto() + A_GEN_WAV_TFM_FFN_GATE = auto() + A_GEN_WAV_TFM_FFN_UP = auto() + A_GEN_WAV_TFM_FFN_DOWN = auto() + A_GEN_WAV_TFM_FFN_SCALE = auto() # layer scale (gamma) on the FFN output + A_GEN_WAV_UP_CONV = auto() # causal ConvTranspose1d, 2x upsample + A_GEN_WAV_UP_DWCONV = auto() # ConvNeXt depthwise conv + A_GEN_WAV_UP_NORM = auto() # ConvNeXt LayerNorm + A_GEN_WAV_UP_PW1 = auto() # ConvNeXt pointwise conv 1 (expand) + A_GEN_WAV_UP_PW2 = auto() # ConvNeXt pointwise conv 2 (project) + A_GEN_WAV_UP_GAMMA = auto() # ConvNeXt layer scale + A_GEN_WAV_DAC_ENTRY = auto() # DAC conv_pre + A_GEN_WAV_DAC_UP_SNAKE = auto() # DAC per-block SnakeBeta before the upsample conv + A_GEN_WAV_DAC_UP_CONV = auto() # DAC per-block causal ConvTranspose1d + A_GEN_WAV_DAC_RES_ACT1 = auto() # DAC residual unit, SnakeBeta before conv1 + A_GEN_WAV_DAC_RES_CONV1 = auto() # DAC residual unit, dilated causal conv + A_GEN_WAV_DAC_RES_ACT2 = auto() # DAC residual unit, SnakeBeta before conv2 + A_GEN_WAV_DAC_RES_CONV2 = auto() # DAC residual unit, pointwise causal conv + A_GEN_WAV_DAC_POST_SNAKE = auto() # DAC final SnakeBeta + A_GEN_WAV_DAC_POST_CONV = auto() # DAC conv_post -> 1-channel PCM + # pocket-tts: SEANet encoder (speaker path) and decoder (a.gen.wav path) + A_ENC_SEANET_CONV_IN = auto() + A_ENC_SEANET_CONV_OUT = auto() + A_ENC_SEANET_RES_CONV1 = auto() # residual unit, dilated conv + A_ENC_SEANET_RES_CONV2 = auto() # residual unit, pointwise conv + A_ENC_SEANET_SCALE_CONV = auto() # strided downsample conv + A_ENC_ATTN_SCALE = auto() # layer scale (gamma) on the attn output + A_ENC_FFN_SCALE_LS = auto() # layer scale (gamma) on the FFN output + A_ENC_SPEAKER_PROJ = auto() # voice latent -> backbone embd + A_GEN_FLOW_INPUT_PROJ = auto() + A_GEN_FLOW_COND_EMBD = auto() + A_GEN_FLOW_TIME_FREQS = auto() # timestep embedder, stored cos/sin frequencies + A_GEN_FLOW_TIME_UP = auto() + A_GEN_FLOW_TIME_DOWN = auto() + A_GEN_FLOW_TIME_NORM = auto() # RMSNorm alpha + A_GEN_FLOW_BLK_NORM = auto() # AdaLN res block, in_ln + A_GEN_FLOW_BLK_UP = auto() + A_GEN_FLOW_BLK_DOWN = auto() + A_GEN_FLOW_BLK_ADA = auto() # AdaLN modulation, -> shift/scale/gate + A_GEN_FLOW_FINAL_ADA = auto() # final layer AdaLN modulation, -> shift/scale + A_GEN_FLOW_FINAL_PROJ = auto() + A_GEN_OUT_EOS = auto() # end-of-speech head on the backbone hidden state + A_GEN_INPUT_LINEAR = auto() # generated latent -> backbone embd + A_GEN_EMB_MEAN = auto() # latent denormalization stats + A_GEN_EMB_STD = auto() + A_GEN_WAV_QUANT_OUT = auto() # DummyQuantizer output_proj, latent -> decoder dim + A_GEN_WAV_UPSAMPLE = auto() # frame rate -> encoder frame rate, depthwise convtr + A_GEN_WAV_SEANET_CONV_IN = auto() + A_GEN_WAV_SEANET_CONV_OUT = auto() # -> 1-channel PCM + A_GEN_WAV_SEANET_RES_CONV1 = auto() + A_GEN_WAV_SEANET_RES_CONV2 = auto() + A_GEN_WAV_SEANET_SCALE_CONV = auto() # strided upsample convtr + A_MMPROJ = auto() + A_MMPROJ_FC = auto() + A_MM_NORM_PRE = auto() + A_MM_NORM_MID = auto() + A_MM_EMBEDDING = auto() # gemma3n + A_MM_HARD_EMB_NORM = auto() # gemma3n + A_MM_SOFT_EMB_NORM = auto() # gemma3n + A_MM_INP_PROJ = auto() # gemma3n + A_MM_CODE_EMBD = auto() # mimo: text-side RVQ code embedding table ("text codebook"), merged 3D [n_channels, vocab, dim] + A_MM_LOCAL_ATTN_Q = auto() # mimo: input_local_transformer (LLM-side connector) + A_MM_LOCAL_ATTN_K = auto() + A_MM_LOCAL_ATTN_V = auto() + A_MM_LOCAL_ATTN_OUT = auto() + A_MM_LOCAL_FFN_GATE = auto() + A_MM_LOCAL_FFN_UP = auto() + A_MM_LOCAL_FFN_DOWN = auto() + A_MM_LOCAL_LN1 = auto() + A_MM_LOCAL_LN2 = auto() + A_MM_LOCAL_NORM = auto() # final norm after all input_local_transformer layers + A_PER_DIM_K_SCALE = auto() # gemma4 + A_PER_DIM_SCALE = auto() # gemma4 + # nextn/mtp + NEXTN_PROJ_PRE = auto() + NEXTN_PROJ_POST = auto() + NEXTN_EH_PROJ = auto() + NEXTN_EMBED_TOKENS = auto() + NEXTN_ENORM = auto() + NEXTN_HNORM = auto() + NEXTN_SHARED_HEAD_HEAD = auto() + NEXTN_SHARED_HEAD_NORM = auto() + NEXTN_HC_HEAD_NORM = auto() + NEXTN_HC_HEAD_DOWN = auto() + NEXTN_HC_HEAD_UP = auto() + # eagle3 + FC = auto() # feature fusion layer + D2T = auto() # draft to target vocabulary mapping + # dspark + DSPARK_MARKOV_W1 = auto() # markov head: prev-token embed + DSPARK_MARKOV_W2 = auto() # markov head: bias projection + DSPARK_CONF_PROJ = auto() # confidence head + DFLASH_ATTN_CONV_BASE = auto() + DFLASH_ATTN_CONV_PROJ = auto() + DFLASH_FFN_CONV_BASE = auto() + DFLASH_FFN_CONV_PROJ = auto() + DFLASH_SELECTOR_PREV = auto() + DFLASH_SELECTOR_NEXT = auto() + DFLASH_SELECTOR_HIDDEN = auto() + # lfm2 audio + A_ENC_NORM_CONV = auto() + A_ENC_LINEAR_POS = auto() + A_ENC_POS_BIAS_U = auto() + A_ENC_POS_BIAS_V = auto() + A_ENC_OUT = auto() + A_ENC_CONV_DW = auto() # SSM conv + A_ENC_CONV_NORM = auto() # SSM conv + A_ENC_CONV_PW1 = auto() + A_ENC_CONV_PW2 = auto() + A_ENC_CONV_NORM_MEAN = auto() # parakeet + A_ENC_CONV_NORM_VAR = auto() # parakeet + A_ENC_MEL_FILTERS = auto() # parakeet + A_ENC_WINDOW = auto() # parakeet + A_CTC_OUT = auto() + A_CTC_OUT_MID = auto() + A_ENC_ATTN_REL_POS_EMB = auto() + # audio qformer projector + A_QF_PROJ_QUERY = auto() + A_QF_PROJ_NORM = auto() + A_QF_PROJ_LINEAR = auto() + A_QF_SELF_ATTN_Q = auto() + A_QF_SELF_ATTN_K = auto() + A_QF_SELF_ATTN_V = auto() + A_QF_SELF_ATTN_O = auto() + A_QF_SELF_ATTN_NORM = auto() + A_QF_CROSS_ATTN_Q = auto() + A_QF_CROSS_ATTN_K = auto() + A_QF_CROSS_ATTN_V = auto() + A_QF_CROSS_ATTN_O = auto() + A_QF_CROSS_ATTN_NORM = auto() + A_QF_FFN_UP = auto() + A_QF_FFN_DOWN = auto() + A_QF_FFN_NORM = auto() + + +MODEL_ARCH_NAMES: dict[MODEL_ARCH, str] = { + MODEL_ARCH.MMPROJ: "clip", # dummy arch for clip.cpp + MODEL_ARCH.LLAMA: "llama", + MODEL_ARCH.LLAMA4: "llama4", + MODEL_ARCH.DECI: "deci", + MODEL_ARCH.FALCON: "falcon", + MODEL_ARCH.BAICHUAN: "baichuan", + MODEL_ARCH.GROK: "grok", + MODEL_ARCH.GPT2: "gpt2", + MODEL_ARCH.GPTJ: "gptj", + MODEL_ARCH.GPTNEOX: "gptneox", + MODEL_ARCH.MPT: "mpt", + MODEL_ARCH.STARCODER: "starcoder", + MODEL_ARCH.REFACT: "refact", + MODEL_ARCH.BERT: "bert", + MODEL_ARCH.MODERN_BERT: "modern-bert", + MODEL_ARCH.NOMIC_BERT: "nomic-bert", + MODEL_ARCH.NOMIC_BERT_MOE: "nomic-bert-moe", + MODEL_ARCH.NEO_BERT: "neo-bert", + MODEL_ARCH.JINA_BERT_V2: "jina-bert-v2", + MODEL_ARCH.JINA_BERT_V3: "jina-bert-v3", + MODEL_ARCH.EUROBERT: "eurobert", + MODEL_ARCH.BLOOM: "bloom", + MODEL_ARCH.STABLELM: "stablelm", + MODEL_ARCH.QWEN: "qwen", + MODEL_ARCH.QWEN2: "qwen2", + MODEL_ARCH.QWEN2MOE: "qwen2moe", + MODEL_ARCH.QWEN2VL: "qwen2vl", + MODEL_ARCH.QWEN3: "qwen3", + MODEL_ARCH.QWEN3MOE: "qwen3moe", + MODEL_ARCH.QWEN3NEXT: "qwen3next", + MODEL_ARCH.HRM_TEXT: "hrm_text", + MODEL_ARCH.QWEN3VL: "qwen3vl", + MODEL_ARCH.QWEN3VLMOE: "qwen3vlmoe", + MODEL_ARCH.QWEN35: "qwen35", + MODEL_ARCH.QWEN35MOE: "qwen35moe", + MODEL_ARCH.CLEF: "clef", + MODEL_ARCH.VLLM_DECISION2: "vllm-decision2", + MODEL_ARCH.QWEN4EXP: "qwen4exp", + MODEL_ARCH.PHI2: "phi2", + MODEL_ARCH.PHI3: "phi3", + MODEL_ARCH.PHIMOE: "phimoe", + MODEL_ARCH.PLAMO: "plamo", + MODEL_ARCH.PLAMO2: "plamo2", + MODEL_ARCH.PLAMO3: "plamo3", + MODEL_ARCH.CODESHELL: "codeshell", + MODEL_ARCH.ORION: "orion", + MODEL_ARCH.INTERNLM2: "internlm2", + MODEL_ARCH.MINICPM: "minicpm", + MODEL_ARCH.MINICPM3: "minicpm3", + MODEL_ARCH.GEMMA: "gemma", + MODEL_ARCH.GEMMA2: "gemma2", + MODEL_ARCH.GEMMA3: "gemma3", + MODEL_ARCH.GEMMA3N: "gemma3n", + MODEL_ARCH.GEMMA4: "gemma4", + MODEL_ARCH.GEMMA4_ASSISTANT: "gemma4-assistant", + MODEL_ARCH.GEMMA_EMBEDDING: "gemma-embedding", + MODEL_ARCH.GEMMA_EMBEDDING2: "gemma-embedding2", + MODEL_ARCH.STARCODER2: "starcoder2", + MODEL_ARCH.RWKV6: "rwkv6", + MODEL_ARCH.RWKV6QWEN2: "rwkv6qwen2", + MODEL_ARCH.RWKV7: "rwkv7", + MODEL_ARCH.ARWKV7: "arwkv7", + MODEL_ARCH.MAMBA: "mamba", + MODEL_ARCH.MAMBA2: "mamba2", + MODEL_ARCH.MAPLE: "maple", + MODEL_ARCH.JAMBA: "jamba", + MODEL_ARCH.XVERSE: "xverse", + MODEL_ARCH.COMMAND_R: "command-r", + MODEL_ARCH.COHERE2: "cohere2", + MODEL_ARCH.COHERE2MOE: "cohere2moe", + MODEL_ARCH.DBRX: "dbrx", + MODEL_ARCH.OLMO: "olmo", + MODEL_ARCH.OLMO2: "olmo2", + MODEL_ARCH.OLMOE: "olmoe", + MODEL_ARCH.MUSE_GLIMMER: "muse-glimmer", + MODEL_ARCH.OPENELM: "openelm", + MODEL_ARCH.ARCTIC: "arctic", + MODEL_ARCH.DEEPSEEK: "deepseek", + MODEL_ARCH.DEEPSEEK2: "deepseek2", + MODEL_ARCH.DEEPSEEK2OCR: "deepseek2-ocr", + MODEL_ARCH.DEEPSEEK32: "deepseek32", + MODEL_ARCH.DEEPSEEK4: "deepseek4", + MODEL_ARCH.CHATGLM: "chatglm", + MODEL_ARCH.GLM4: "glm4", + MODEL_ARCH.GLM4_MOE: "glm4moe", + MODEL_ARCH.GLM_DSA: "glm-dsa", + MODEL_ARCH.GLM5_NEXT: "glm5-next", + MODEL_ARCH.BITNET: "bitnet", + MODEL_ARCH.T5: "t5", + MODEL_ARCH.T5ENCODER: "t5encoder", + MODEL_ARCH.JAIS: "jais", + MODEL_ARCH.JAIS2: "jais2", + MODEL_ARCH.NEMOTRON: "nemotron", + MODEL_ARCH.NEMOTRON_H: "nemotron_h", + MODEL_ARCH.NEMOTRON_H_MOE: "nemotron_h_moe", + MODEL_ARCH.EXAONE: "exaone", + MODEL_ARCH.EXAONE4: "exaone4", + MODEL_ARCH.EXAONE_MOE: "exaone-moe", + MODEL_ARCH.GRANITE: "granite", + MODEL_ARCH.GRANITE_MOE: "granitemoe", + MODEL_ARCH.GRANITE_HYBRID: "granitehybrid", + MODEL_ARCH.GRANITE_SWITCH: "graniteswitch", + MODEL_ARCH.GRANITE_SWA: "granite_swa", + MODEL_ARCH.CHAMELEON: "chameleon", + MODEL_ARCH.WAVTOKENIZER_DEC: "wavtokenizer-dec", + MODEL_ARCH.PLM: "plm", + MODEL_ARCH.BAILINGMOE: "bailingmoe", + MODEL_ARCH.BAILINGMOE2: "bailingmoe2", + MODEL_ARCH.BAILINGMOE3: "bailingmoe3", + MODEL_ARCH.DOTS1: "dots1", + MODEL_ARCH.DOTS3NOTE: "dots3note", + MODEL_ARCH.ARCEE: "arcee", + MODEL_ARCH.AFMOE: "afmoe", + MODEL_ARCH.LAGUNA: "laguna", + MODEL_ARCH.ERNIE4_5: "ernie4_5", + MODEL_ARCH.ERNIE4_5_MOE: "ernie4_5-moe", + MODEL_ARCH.FALCON_H1: "falcon-h1", + MODEL_ARCH.HUNYUAN_MOE: "hunyuan-moe", + MODEL_ARCH.HUNYUAN_DENSE: "hunyuan-dense", + MODEL_ARCH.HUNYUAN_VL: "hunyuan_vl", + MODEL_ARCH.HY_V3: "hy_v3", + MODEL_ARCH.HY_V4: "hy_v4", + MODEL_ARCH.SMOLLM3: "smollm3", + MODEL_ARCH.GPT_OSS: "gpt-oss", + MODEL_ARCH.LFM2: "lfm2", + MODEL_ARCH.LFM2MOE: "lfm2moe", + MODEL_ARCH.DREAM: "dream", + MODEL_ARCH.SMALLTHINKER: "smallthinker", + MODEL_ARCH.LLADA: "llada", + MODEL_ARCH.LLADA_MOE: "llada-moe", + MODEL_ARCH.SEED_OSS: "seed_oss", + MODEL_ARCH.GROVEMOE: "grovemoe", + MODEL_ARCH.APERTUS: "apertus", + MODEL_ARCH.MINIMAX01: "minimax-01", + MODEL_ARCH.MINIMAXM2: "minimax-m2", + MODEL_ARCH.MINIMAXM3: "minimax-m3", + MODEL_ARCH.COGVLM: "cogvlm", + MODEL_ARCH.RND1: "rnd1", + MODEL_ARCH.PANGU_EMBED: "pangu-embedded", + MODEL_ARCH.MISTRAL3: "mistral3", + MODEL_ARCH.EAGLE3: "eagle3", + MODEL_ARCH.DFLASH: "dflash", + MODEL_ARCH.MISTRAL4: "mistral4", + MODEL_ARCH.PADDLEOCR: "paddleocr", + MODEL_ARCH.MIMO2: "mimo2", + MODEL_ARCH.STEP35: "step35", + MODEL_ARCH.SPARK2_5: "spark2_5", + MODEL_ARCH.LLAMA_EMBED: "llama-embed", + MODEL_ARCH.MAINCODER: "maincoder", + MODEL_ARCH.KIMI_LINEAR: "kimi-linear", + MODEL_ARCH.KIMI_K3: "kimi-k3", + MODEL_ARCH.TALKIE: "talkie", + MODEL_ARCH.MELLUM: "mellum", + MODEL_ARCH.NANBEIGE: "nanbeige", + MODEL_ARCH.QWEN3TTS: "qwen3tts", + MODEL_ARCH.POCKETTTS: "pockettts", + MODEL_ARCH.K2HORIZON: "k2-horizon", +} + +VISION_PROJECTOR_TYPE_NAMES: dict[VISION_PROJECTOR_TYPE, str] = { + VISION_PROJECTOR_TYPE.MLP: "mlp", + VISION_PROJECTOR_TYPE.LDP: "ldp", + VISION_PROJECTOR_TYPE.LDPV2: "ldpv2", + VISION_PROJECTOR_TYPE.RESAMPLER: "resampler", + VISION_PROJECTOR_TYPE.GLM_EDGE: "adapter", + VISION_PROJECTOR_TYPE.MERGER: "qwen2vl_merger", + VISION_PROJECTOR_TYPE.GEMMA3: "gemma3", + VISION_PROJECTOR_TYPE.QWEN3VL: "qwen3vl_merger", + VISION_PROJECTOR_TYPE.LING3VL: "ling3vl", + VISION_PROJECTOR_TYPE.STEP3VL: "step3vl", +} + +TENSOR_NAMES: dict[MODEL_TENSOR, str] = { + MODEL_TENSOR.TOKEN_EMBD: "token_embd", + MODEL_TENSOR.TOKEN_EMBD_NORM: "token_embd_norm", + MODEL_TENSOR.TOKEN_TYPES: "token_types", + MODEL_TENSOR.MASKED_EMBD_CENTROIDS: "masked_embd_centroids", + MODEL_TENSOR.MASKED_EMBD_ORDERING: "masked_embd_ordering", + MODEL_TENSOR.POS_EMBD: "position_embd", + MODEL_TENSOR.OUTPUT_NORM: "output_norm", + MODEL_TENSOR.OUTPUT: "output", + MODEL_TENSOR.HRM_Z_L_INIT: "hrm.z_l_init", + MODEL_TENSOR.DENSE_2_OUT: "dense_2", # embeddinggemma 2_Dense + MODEL_TENSOR.DENSE_3_OUT: "dense_3", # embeddinggemma 2_Dense + MODEL_TENSOR.HC_HEAD_FN: "output_hc_fn", + MODEL_TENSOR.HC_HEAD_BASE: "output_hc_base", + MODEL_TENSOR.HC_HEAD_SCALE: "output_hc_scale", + MODEL_TENSOR.HC_HEAD_NORM: "output_hc_norm", # qwen4exp + MODEL_TENSOR.HC_HEAD_DOWN: "output_hc_down", # qwen4exp + MODEL_TENSOR.HC_HEAD_UP: "output_hc_up", # qwen4exp + MODEL_TENSOR.ROPE_FREQS: "rope_freqs", + MODEL_TENSOR.ROPE_FACTORS_LONG: "rope_factors_long", + MODEL_TENSOR.ROPE_FACTORS_SHORT: "rope_factors_short", + MODEL_TENSOR.ATTN_NORM: "blk.{bid}.attn_norm", + MODEL_TENSOR.ATTN_NORM_2: "blk.{bid}.attn_norm_2", + MODEL_TENSOR.ATTN_QKV: "blk.{bid}.attn_qkv", + MODEL_TENSOR.ATTN_Q: "blk.{bid}.attn_q", + MODEL_TENSOR.ATTN_K: "blk.{bid}.attn_k", + MODEL_TENSOR.ATTN_V: "blk.{bid}.attn_v", + MODEL_TENSOR.ATTN_OUT: "blk.{bid}.attn_output", + MODEL_TENSOR.ATTN_ROT_EMBD: "blk.{bid}.attn_rot_embd", + MODEL_TENSOR.ATTN_SINKS: "blk.{bid}.attn_sinks", + MODEL_TENSOR.ATTN_GATE: "blk.{bid}.attn_gate", + MODEL_TENSOR.ATTN_Q_NORM: "blk.{bid}.attn_q_norm", + MODEL_TENSOR.ATTN_K_NORM: "blk.{bid}.attn_k_norm", + MODEL_TENSOR.ATTN_OUT_NORM: "blk.{bid}.attn_output_norm", + MODEL_TENSOR.ATTN_POST_NORM: "blk.{bid}.post_attention_norm", + MODEL_TENSOR.FFN_GATE_INP: "blk.{bid}.ffn_gate_inp", + MODEL_TENSOR.FFN_GATE_INP_SHEXP: "blk.{bid}.ffn_gate_inp_shexp", + MODEL_TENSOR.FFN_NORM: "blk.{bid}.ffn_norm", + MODEL_TENSOR.FFN_PRE_NORM: "blk.{bid}.ffn_norm", + MODEL_TENSOR.FFN_POST_NORM: "blk.{bid}.post_ffw_norm", + MODEL_TENSOR.FFN_PRE_NORM_2: "blk.{bid}.pre_ffw_norm_2", # gemma4 + MODEL_TENSOR.FFN_POST_NORM_1: "blk.{bid}.post_ffw_norm_1", # gemma4 + MODEL_TENSOR.FFN_POST_NORM_2: "blk.{bid}.post_ffw_norm_2", # gemma4 + MODEL_TENSOR.FFN_GATE: "blk.{bid}.ffn_gate", + MODEL_TENSOR.FFN_DOWN: "blk.{bid}.ffn_down", + MODEL_TENSOR.FFN_UP: "blk.{bid}.ffn_up", + MODEL_TENSOR.FFN_GATE_SHEXP: "blk.{bid}.ffn_gate_shexp", + MODEL_TENSOR.FFN_DOWN_SHEXP: "blk.{bid}.ffn_down_shexp", + MODEL_TENSOR.FFN_UP_SHEXP: "blk.{bid}.ffn_up_shexp", + MODEL_TENSOR.FFN_GATE_CHEXP: "blk.{bid}.ffn_gate_chexps", + MODEL_TENSOR.FFN_DOWN_CHEXP: "blk.{bid}.ffn_down_chexps", + MODEL_TENSOR.FFN_UP_CHEXP: "blk.{bid}.ffn_up_chexps", + MODEL_TENSOR.FFN_ACT: "blk.{bid}.ffn", + MODEL_TENSOR.FFN_NORM_EXP: "blk.{bid}.ffn_norm_exps", + MODEL_TENSOR.FFN_GATE_EXP: "blk.{bid}.ffn_gate_exps", + MODEL_TENSOR.FFN_DOWN_EXP: "blk.{bid}.ffn_down_exps", + MODEL_TENSOR.FFN_UP_EXP: "blk.{bid}.ffn_up_exps", + MODEL_TENSOR.FFN_GATE_UP_EXP: "blk.{bid}.ffn_gate_up_exps", + MODEL_TENSOR.FFN_EXP_PROBS_B: "blk.{bid}.exp_probs_b", + MODEL_TENSOR.FFN_EXP_PROBS_B_VL: "blk.{bid}.exp_probs_b_vl", + MODEL_TENSOR.FFN_GATE_TID2EID: "blk.{bid}.ffn_gate_tid2eid", + MODEL_TENSOR.MOE_LATENT_DOWN: "blk.{bid}.ffn_latent_down", # nemotron 3 super + MODEL_TENSOR.MOE_LATENT_UP: "blk.{bid}.ffn_latent_up", # nemotron 3 super + MODEL_TENSOR.LAYER_OUT_NORM: "blk.{bid}.layer_output_norm", + MODEL_TENSOR.LAYER_OUT_SCALE: "blk.{bid}.layer_output_scale", + MODEL_TENSOR.PER_LAYER_TOKEN_EMBD: "per_layer_token_embd", # gemma3n + MODEL_TENSOR.PER_LAYER_MODEL_PROJ: "per_layer_model_proj", # gemma3n + MODEL_TENSOR.PER_LAYER_PROJ_NORM: "per_layer_proj_norm", # gemma3n + MODEL_TENSOR.ALTUP_UNEMBD_PROJ: "altup_unembd_proj", # gemma3n + MODEL_TENSOR.ALTUP_PROJ: "altup_proj", # gemma3n + MODEL_TENSOR.PER_LAYER_INP_GATE: "blk.{bid}.inp_gate", # gemma3n + MODEL_TENSOR.PER_LAYER_PROJ: "blk.{bid}.proj", # gemma3n + MODEL_TENSOR.PER_LAYER_POST_NORM: "blk.{bid}.post_norm", # gemma3n + MODEL_TENSOR.ALTUP_CORRECT_COEF: "blk.{bid}.altup_correct_coef", # gemma3n + MODEL_TENSOR.ALTUP_CORRECT_SCALE: "blk.{bid}.altup_correct_scale", # gemma3n + MODEL_TENSOR.ALTUP_PREDICT_COEF: "blk.{bid}.altup_predict_coef", # gemma3n + MODEL_TENSOR.ALTUP_ROUTER: "blk.{bid}.altup_router", # gemma3n + MODEL_TENSOR.ALTUP_ROUTER_NORM: "blk.{bid}.altup_router_norm", # gemma3n + MODEL_TENSOR.LAUREL_L: "blk.{bid}.laurel_l", # gemma3n + MODEL_TENSOR.LAUREL_R: "blk.{bid}.laurel_r", # gemma3n + MODEL_TENSOR.LAUREL_POST_NORM: "blk.{bid}.laurel_post_norm", # gemma3n + MODEL_TENSOR.SSM_IN: "blk.{bid}.ssm_in", + MODEL_TENSOR.SSM_CONV1D: "blk.{bid}.ssm_conv1d", + MODEL_TENSOR.SSM_X: "blk.{bid}.ssm_x", + MODEL_TENSOR.SSM_DT: "blk.{bid}.ssm_dt", + MODEL_TENSOR.SSM_DT_NORM: "blk.{bid}.ssm_dt_norm", + MODEL_TENSOR.SSM_A: "blk.{bid}.ssm_a", + MODEL_TENSOR.SSM_B_NORM: "blk.{bid}.ssm_b_norm", + MODEL_TENSOR.SSM_C_NORM: "blk.{bid}.ssm_c_norm", + MODEL_TENSOR.SSM_D: "blk.{bid}.ssm_d", + MODEL_TENSOR.SSM_NORM: "blk.{bid}.ssm_norm", + MODEL_TENSOR.SSM_OUT: "blk.{bid}.ssm_out", + MODEL_TENSOR.SSM_ALPHA: "blk.{bid}.ssm_alpha", # qwen3.5 + MODEL_TENSOR.SSM_BETA_ALPHA: "blk.{bid}.ssm_ba", + MODEL_TENSOR.SSM_CONV1D_Q: "blk.{bid}.ssm_conv1d_q", # Kimi Linear + MODEL_TENSOR.SSM_CONV1D_K: "blk.{bid}.ssm_conv1d_k", # Kimi Linear + MODEL_TENSOR.SSM_CONV1D_V: "blk.{bid}.ssm_conv1d_v", # Kimi Linear + MODEL_TENSOR.SSM_F_A: "blk.{bid}.ssm_f_a", # Kimi Linear + MODEL_TENSOR.SSM_F_B: "blk.{bid}.ssm_f_b", # Kimi Linear + MODEL_TENSOR.SSM_BETA: "blk.{bid}.ssm_beta", # Kimi Linear qwen3.5 + MODEL_TENSOR.SSM_G_A: "blk.{bid}.ssm_g_a", # Kimi Linear + MODEL_TENSOR.SSM_G_B: "blk.{bid}.ssm_g_b", # Kimi Linear + MODEL_TENSOR.SSM_G: "blk.{bid}.ssm_g", # Kimi K3 + MODEL_TENSOR.ATTN_RES_SCORE: "blk.{bid}.attn_res_score", # Kimi K3 + MODEL_TENSOR.FFN_RES_SCORE: "blk.{bid}.ffn_res_score", # Kimi K3 + MODEL_TENSOR.OUTPUT_RES_SCORE: "output_res_score", # Kimi K3 + MODEL_TENSOR.FFN_ROUTED_DOWN: "blk.{bid}.ffn_routed_down", # Kimi K3 + MODEL_TENSOR.FFN_ROUTED_UP: "blk.{bid}.ffn_routed_up", # Kimi K3 + MODEL_TENSOR.FFN_ROUTED_NORM: "blk.{bid}.ffn_routed_norm", # Kimi K3 + MODEL_TENSOR.TIME_MIX_W0: "blk.{bid}.time_mix_w0", + MODEL_TENSOR.TIME_MIX_W1: "blk.{bid}.time_mix_w1", + MODEL_TENSOR.TIME_MIX_W2: "blk.{bid}.time_mix_w2", + MODEL_TENSOR.TIME_MIX_A0: "blk.{bid}.time_mix_a0", + MODEL_TENSOR.TIME_MIX_A1: "blk.{bid}.time_mix_a1", + MODEL_TENSOR.TIME_MIX_A2: "blk.{bid}.time_mix_a2", + MODEL_TENSOR.TIME_MIX_V0: "blk.{bid}.time_mix_v0", + MODEL_TENSOR.TIME_MIX_V1: "blk.{bid}.time_mix_v1", + MODEL_TENSOR.TIME_MIX_V2: "blk.{bid}.time_mix_v2", + MODEL_TENSOR.TIME_MIX_G1: "blk.{bid}.time_mix_g1", + MODEL_TENSOR.TIME_MIX_G2: "blk.{bid}.time_mix_g2", + MODEL_TENSOR.TIME_MIX_K_K: "blk.{bid}.time_mix_k_k", + MODEL_TENSOR.TIME_MIX_K_A: "blk.{bid}.time_mix_k_a", + MODEL_TENSOR.TIME_MIX_R_K: "blk.{bid}.time_mix_r_k", + MODEL_TENSOR.TIME_MIX_LERP_X: "blk.{bid}.time_mix_lerp_x", + MODEL_TENSOR.TIME_MIX_LERP_K: "blk.{bid}.time_mix_lerp_k", + MODEL_TENSOR.TIME_MIX_LERP_V: "blk.{bid}.time_mix_lerp_v", + MODEL_TENSOR.TIME_MIX_LERP_R: "blk.{bid}.time_mix_lerp_r", + MODEL_TENSOR.TIME_MIX_LERP_G: "blk.{bid}.time_mix_lerp_g", + MODEL_TENSOR.TIME_MIX_LERP_FUSED: "blk.{bid}.time_mix_lerp_fused", + MODEL_TENSOR.TIME_MIX_LERP_W: "blk.{bid}.time_mix_lerp_w", + MODEL_TENSOR.TIME_MIX_FIRST: "blk.{bid}.time_mix_first", + MODEL_TENSOR.TIME_MIX_DECAY: "blk.{bid}.time_mix_decay", + MODEL_TENSOR.TIME_MIX_DECAY_W1: "blk.{bid}.time_mix_decay_w1", + MODEL_TENSOR.TIME_MIX_DECAY_W2: "blk.{bid}.time_mix_decay_w2", + MODEL_TENSOR.TIME_MIX_KEY: "blk.{bid}.time_mix_key", + MODEL_TENSOR.TIME_MIX_VALUE: "blk.{bid}.time_mix_value", + MODEL_TENSOR.TIME_MIX_RECEPTANCE: "blk.{bid}.time_mix_receptance", + MODEL_TENSOR.TIME_MIX_GATE: "blk.{bid}.time_mix_gate", + MODEL_TENSOR.TIME_MIX_LN: "blk.{bid}.time_mix_ln", + MODEL_TENSOR.TIME_MIX_OUTPUT: "blk.{bid}.time_mix_output", + MODEL_TENSOR.CHANNEL_MIX_LERP_K: "blk.{bid}.channel_mix_lerp_k", + MODEL_TENSOR.CHANNEL_MIX_LERP_R: "blk.{bid}.channel_mix_lerp_r", + MODEL_TENSOR.CHANNEL_MIX_KEY: "blk.{bid}.channel_mix_key", + MODEL_TENSOR.CHANNEL_MIX_RECEPTANCE: "blk.{bid}.channel_mix_receptance", + MODEL_TENSOR.CHANNEL_MIX_VALUE: "blk.{bid}.channel_mix_value", + MODEL_TENSOR.ATTN_Q_A: "blk.{bid}.attn_q_a", + MODEL_TENSOR.ATTN_Q_B: "blk.{bid}.attn_q_b", + MODEL_TENSOR.ATTN_KV_A_MQA: "blk.{bid}.attn_kv_a_mqa", + MODEL_TENSOR.ATTN_KV_B: "blk.{bid}.attn_kv_b", + MODEL_TENSOR.ATTN_K_B: "blk.{bid}.attn_k_b", + MODEL_TENSOR.ATTN_V_B: "blk.{bid}.attn_v_b", + MODEL_TENSOR.ATTN_Q_A_NORM: "blk.{bid}.attn_q_a_norm", + MODEL_TENSOR.ATTN_KV_A_NORM: "blk.{bid}.attn_kv_a_norm", + MODEL_TENSOR.ATTN_KV: "blk.{bid}.attn_kv", + MODEL_TENSOR.ATTN_KV_NORM: "blk.{bid}.attn_kv_a_norm", + MODEL_TENSOR.ATTN_OUT_A: "blk.{bid}.attn_output_a", + MODEL_TENSOR.ATTN_OUT_B: "blk.{bid}.attn_output_b", + MODEL_TENSOR.HC_ATTN_FN: "blk.{bid}.hc_attn_fn", + MODEL_TENSOR.HC_ATTN_BASE: "blk.{bid}.hc_attn_base", + MODEL_TENSOR.HC_ATTN_SCALE: "blk.{bid}.hc_attn_scale", + MODEL_TENSOR.HC_FFN_FN: "blk.{bid}.hc_ffn_fn", + MODEL_TENSOR.HC_FFN_BASE: "blk.{bid}.hc_ffn_base", + MODEL_TENSOR.HC_FFN_SCALE: "blk.{bid}.hc_ffn_scale", + MODEL_TENSOR.HC_ATTN_NORM: "blk.{bid}.hc_attn_norm", # qwen4exp + MODEL_TENSOR.HC_ATTN_DOWN: "blk.{bid}.hc_attn_down", # qwen4exp + MODEL_TENSOR.HC_ATTN_UP: "blk.{bid}.hc_attn_up", # qwen4exp + MODEL_TENSOR.HC_ATTN_INJECT: "blk.{bid}.hc_attn_inject", # qwen4exp + MODEL_TENSOR.HC_FFN_NORM: "blk.{bid}.hc_ffn_norm", # qwen4exp + MODEL_TENSOR.HC_FFN_DOWN: "blk.{bid}.hc_ffn_down", # qwen4exp + MODEL_TENSOR.HC_FFN_UP: "blk.{bid}.hc_ffn_up", # qwen4exp + MODEL_TENSOR.HC_FFN_INJECT: "blk.{bid}.hc_ffn_inject", # qwen4exp + MODEL_TENSOR.PLE_KEY: "blk.{bid}.ple_key", # qwen4exp + MODEL_TENSOR.PLE_VALUE: "blk.{bid}.ple_value", # qwen4exp + MODEL_TENSOR.PLE_NORM_KEY: "blk.{bid}.ple_norm_key", # qwen4exp + MODEL_TENSOR.PLE_NORM_QUERY: "blk.{bid}.ple_norm_query", # qwen4exp + MODEL_TENSOR.PLE_NORM_CONV: "blk.{bid}.ple_norm_conv", # qwen4exp + MODEL_TENSOR.PLE_CONV1D: "blk.{bid}.ple_conv1d", # qwen4exp + MODEL_TENSOR.ATTN_COMPRESSOR_WKV: "blk.{bid}.attn_compressor_kv", + MODEL_TENSOR.ATTN_COMPRESSOR_WGATE: "blk.{bid}.attn_compressor_gate", + MODEL_TENSOR.ATTN_COMPRESSOR_APE: "blk.{bid}.attn_compressor_ape", + MODEL_TENSOR.ATTN_COMPRESSOR_NORM: "blk.{bid}.attn_compressor_norm", + MODEL_TENSOR.ATTN_SUB_NORM: "blk.{bid}.attn_sub_norm", + MODEL_TENSOR.FFN_SUB_NORM: "blk.{bid}.ffn_sub_norm", + MODEL_TENSOR.DEC_ATTN_NORM: "dec.blk.{bid}.attn_norm", + MODEL_TENSOR.DEC_ATTN_Q: "dec.blk.{bid}.attn_q", + MODEL_TENSOR.DEC_ATTN_K: "dec.blk.{bid}.attn_k", + MODEL_TENSOR.DEC_ATTN_V: "dec.blk.{bid}.attn_v", + MODEL_TENSOR.DEC_ATTN_OUT: "dec.blk.{bid}.attn_o", + MODEL_TENSOR.DEC_ATTN_REL_B: "dec.blk.{bid}.attn_rel_b", + MODEL_TENSOR.DEC_CROSS_ATTN_NORM: "dec.blk.{bid}.cross_attn_norm", + MODEL_TENSOR.DEC_CROSS_ATTN_NORM_KV: "dec.blk.{bid}.cross_attn_norm_kv", + MODEL_TENSOR.DEC_CROSS_ATTN_Q: "dec.blk.{bid}.cross_attn_q", + MODEL_TENSOR.DEC_CROSS_ATTN_K: "dec.blk.{bid}.cross_attn_k", + MODEL_TENSOR.DEC_CROSS_ATTN_V: "dec.blk.{bid}.cross_attn_v", + MODEL_TENSOR.DEC_CROSS_ATTN_OUT: "dec.blk.{bid}.cross_attn_o", + MODEL_TENSOR.DEC_CROSS_ATTN_REL_B: "dec.blk.{bid}.cross_attn_rel_b", + MODEL_TENSOR.DEC_FFN_NORM: "dec.blk.{bid}.ffn_norm", + MODEL_TENSOR.DEC_FFN_GATE: "dec.blk.{bid}.ffn_gate", + MODEL_TENSOR.DEC_FFN_DOWN: "dec.blk.{bid}.ffn_down", + MODEL_TENSOR.DEC_FFN_UP: "dec.blk.{bid}.ffn_up", + MODEL_TENSOR.DEC_OUTPUT_NORM: "dec.output_norm", + MODEL_TENSOR.ENC_ATTN_NORM: "enc.blk.{bid}.attn_norm", + MODEL_TENSOR.ENC_ATTN_Q: "enc.blk.{bid}.attn_q", + MODEL_TENSOR.ENC_ATTN_K: "enc.blk.{bid}.attn_k", + MODEL_TENSOR.ENC_ATTN_V: "enc.blk.{bid}.attn_v", + MODEL_TENSOR.ENC_ATTN_OUT: "enc.blk.{bid}.attn_o", + MODEL_TENSOR.ENC_ATTN_REL_B: "enc.blk.{bid}.attn_rel_b", + MODEL_TENSOR.ENC_FFN_NORM: "enc.blk.{bid}.ffn_norm", + MODEL_TENSOR.ENC_FFN_GATE: "enc.blk.{bid}.ffn_gate", + MODEL_TENSOR.ENC_FFN_DOWN: "enc.blk.{bid}.ffn_down", + MODEL_TENSOR.ENC_FFN_UP: "enc.blk.{bid}.ffn_up", + MODEL_TENSOR.ENC_OUTPUT_NORM: "enc.output_norm", + MODEL_TENSOR.CLS: "cls", + MODEL_TENSOR.CLS_OUT: "cls.output", + MODEL_TENSOR.CLS_NORM: "cls.norm", + MODEL_TENSOR.DECISION_HIDDEN_NORM: "decision.hidden_norm", + MODEL_TENSOR.DECISION_PROJ_MEMORY: "decision.proj_memory", + MODEL_TENSOR.DECISION_PROJ_QUESTION: "decision.proj_question", + MODEL_TENSOR.DECISION_PROJ_OPTION_QUESTION: "decision.proj_option_question", + MODEL_TENSOR.DECISION_PROJ_GLOBAL: "decision.proj_global", + MODEL_TENSOR.DECISION_PROJ_OPTION_CONTEXT: "decision.proj_option_context", + MODEL_TENSOR.DECISION_PROJ_OPTION_LEXICAL: "decision.proj_option_lexical", + MODEL_TENSOR.DECISION_OPTION_SUMMARY_NORM: "decision.option_summary_norm", + MODEL_TENSOR.DECISION_FIELD_NORM: "decision.field_norm", + MODEL_TENSOR.DECISION_OPTION_NORM: "decision.option_norm", + MODEL_TENSOR.DECISION_SCALES: "decision.scales", + MODEL_TENSOR.DECISION_SCORER: "decision.scorer", + MODEL_TENSOR.DECISION_SCORER_OUT: "decision.scorer_out", + MODEL_TENSOR.DECISION_QUERY_NORM: "decision.query_norm", + MODEL_TENSOR.DECISION_KEY: "decision.key", + MODEL_TENSOR.DECISION_QUERY: "decision.query", + MODEL_TENSOR.DECISION_OPTION_MLP: "decision.option_mlp", + MODEL_TENSOR.DECISION_QUERY_MLP: "decision.query_mlp", + MODEL_TENSOR.CONV1D: "conv1d", + MODEL_TENSOR.CONVNEXT_DW: "convnext.{bid}.dw", + MODEL_TENSOR.CONVNEXT_NORM: "convnext.{bid}.norm", + MODEL_TENSOR.CONVNEXT_PW1: "convnext.{bid}.pw1", + MODEL_TENSOR.CONVNEXT_PW2: "convnext.{bid}.pw2", + MODEL_TENSOR.CONVNEXT_GAMMA: "convnext.{bid}.gamma", + MODEL_TENSOR.POSNET_CONV1: "posnet.{bid}.conv1", + MODEL_TENSOR.POSNET_CONV2: "posnet.{bid}.conv2", + MODEL_TENSOR.POSNET_NORM: "posnet.{bid}.norm", + MODEL_TENSOR.POSNET_NORM1: "posnet.{bid}.norm1", + MODEL_TENSOR.POSNET_NORM2: "posnet.{bid}.norm2", + MODEL_TENSOR.POSNET_ATTN_NORM: "posnet.{bid}.attn_norm", + MODEL_TENSOR.POSNET_ATTN_Q: "posnet.{bid}.attn_q", + MODEL_TENSOR.POSNET_ATTN_K: "posnet.{bid}.attn_k", + MODEL_TENSOR.POSNET_ATTN_V: "posnet.{bid}.attn_v", + MODEL_TENSOR.POSNET_ATTN_OUT: "posnet.{bid}.attn_output", + MODEL_TENSOR.SHORTCONV_CONV: "blk.{bid}.shortconv.conv", + MODEL_TENSOR.SHORTCONV_INPROJ: "blk.{bid}.shortconv.in_proj", + MODEL_TENSOR.SHORTCONV_OUTPROJ: "blk.{bid}.shortconv.out_proj", + MODEL_TENSOR.VISEXP_ATTN_QKV: "blk.{bid}.vis_attn_qkv", + MODEL_TENSOR.VISEXP_ATTN_OUT: "blk.{bid}.vis_attn_output", + MODEL_TENSOR.VISEXP_GATE: "blk.{bid}.vis_gate", + MODEL_TENSOR.VISEXP_DOWN: "blk.{bid}.vis_down", + MODEL_TENSOR.VISEXP_UP: "blk.{bid}.vis_up", + MODEL_TENSOR.INDEXER_K_NORM: "blk.{bid}.indexer.k_norm", + MODEL_TENSOR.INDEXER_PROJ: "blk.{bid}.indexer.proj", + MODEL_TENSOR.INDEXER_ATTN_K: "blk.{bid}.indexer.attn_k", + MODEL_TENSOR.INDEXER_ATTN_Q_B: "blk.{bid}.indexer.attn_q_b", + MODEL_TENSOR.INDEXER_Q_PROJ: "blk.{bid}.indexer.q_proj", + MODEL_TENSOR.INDEXER_K_PROJ: "blk.{bid}.indexer.k_proj", + MODEL_TENSOR.INDEXER_Q_NORM: "blk.{bid}.indexer.q_norm", + MODEL_TENSOR.INDEXER_COMPRESSOR_WKV: "blk.{bid}.indexer_compressor_kv", + MODEL_TENSOR.INDEXER_COMPRESSOR_WGATE: "blk.{bid}.indexer_compressor_gate", + MODEL_TENSOR.INDEXER_COMPRESSOR_APE: "blk.{bid}.indexer_compressor_ape", + MODEL_TENSOR.INDEXER_COMPRESSOR_NORM: "blk.{bid}.indexer_compressor_norm", + MODEL_TENSOR.INDEXER_KPOOL_GATE: "blk.{bid}.indexer_compressor_gate", + MODEL_TENSOR.INDEXER_KPOOL_APE: "blk.{bid}.indexer_compressor_ape", + # vision + MODEL_TENSOR.V_MMPROJ: "mm.{bid}", + MODEL_TENSOR.V_MMPROJ_FC: "mm.model.fc", + MODEL_TENSOR.V_MMPROJ_MLP: "mm.model.mlp.{bid}", + MODEL_TENSOR.V_MMPROJ_PEG: "mm.model.peg.{bid}", + MODEL_TENSOR.V_ENC_EMBD_CLS: "v.class_embd", + MODEL_TENSOR.V_ENC_EMBD_PATCH: "v.patch_embd", + MODEL_TENSOR.V_ENC_EMBD_NORM: "v.norm_embd", + MODEL_TENSOR.V_ENC_EMBD_PATCH_NORM: "v.patch_norm.{bid}", + MODEL_TENSOR.V_ENC_EMBD_POS: "v.position_embd", + MODEL_TENSOR.V_ENC_ATTN_QKV: "v.blk.{bid}.attn_qkv", + MODEL_TENSOR.V_ENC_ATTN_Q: "v.blk.{bid}.attn_q", + MODEL_TENSOR.V_ENC_ATTN_Q_NORM: "v.blk.{bid}.attn_q_norm", + MODEL_TENSOR.V_ENC_ATTN_K: "v.blk.{bid}.attn_k", + MODEL_TENSOR.V_ENC_ATTN_K_NORM: "v.blk.{bid}.attn_k_norm", + MODEL_TENSOR.V_ENC_ATTN_V: "v.blk.{bid}.attn_v", + MODEL_TENSOR.V_ENC_INPUT_NORM: "v.blk.{bid}.ln1", + MODEL_TENSOR.V_ENC_ATTN_O: "v.blk.{bid}.attn_out", + MODEL_TENSOR.V_ENC_ATTN_O_NORM: "v.blk.{bid}.attn_out_norm", + MODEL_TENSOR.V_ENC_ATTN_SINKS: "v.blk.{bid}.attn_sinks", + MODEL_TENSOR.V_ENC_POST_ATTN_NORM: "v.blk.{bid}.ln2", + MODEL_TENSOR.V_ENC_FFN_UP: "v.blk.{bid}.ffn_up", + MODEL_TENSOR.V_ENC_FFN_GATE: "v.blk.{bid}.ffn_gate", + MODEL_TENSOR.V_ENC_FFN_DOWN: "v.blk.{bid}.ffn_down", + MODEL_TENSOR.V_ENC_FFN_GATE_INP: "v.blk.{bid}.ffn_gate_inp", + MODEL_TENSOR.V_ENC_FFN_GATE_EXPS: "v.blk.{bid}.ffn_gate_exps", + MODEL_TENSOR.V_ENC_FFN_UP_EXPS: "v.blk.{bid}.ffn_up_exps", + MODEL_TENSOR.V_ENC_FFN_DOWN_EXPS: "v.blk.{bid}.ffn_down_exps", + MODEL_TENSOR.V_ENC_FFN_EXP_PROBS_B: "v.blk.{bid}.exp_probs_b", + MODEL_TENSOR.V_ENC_ATTN_POST_NORM: "v.blk.{bid}.attn_post_norm", + MODEL_TENSOR.V_ENC_FFN_POST_NORM: "v.blk.{bid}.ffn_post_norm", + MODEL_TENSOR.V_LAYER_SCALE_1: "v.blk.{bid}.ls1", + MODEL_TENSOR.V_LAYER_SCALE_2: "v.blk.{bid}.ls2", + MODEL_TENSOR.V_LAYER_OUT_SCALE: "v.blk.{bid}.out_scale", + MODEL_TENSOR.V_PRE_NORM: "v.pre_ln", + MODEL_TENSOR.V_POST_NORM: "v.post_ln", + MODEL_TENSOR.V_MM_POST_NORM: "mm.post_norm", + MODEL_TENSOR.V_MM_INP_PROJ: "mm.input_projection", + MODEL_TENSOR.V_MM_INP_NORM: "mm.input_norm", + MODEL_TENSOR.V_MM_SOFT_EMB_NORM: "mm.soft_emb_norm", # gemma3n + MODEL_TENSOR.V_MM_EMBEDDING: "mm.embedding", # gemma3n + MODEL_TENSOR.V_MM_HARD_EMB_NORM: "mm.hard_emb_norm", # gemma3n + MODEL_TENSOR.V_ENC_CONV_STEM: "v.conv_stem.conv", # gemma3n + MODEL_TENSOR.V_ENC_CONV_STEM_NORM: "v.conv_stem.bn", # gemma3n + MODEL_TENSOR.V_ENC_MSFA_EXP: "v.msfa.ffn.pw_exp.conv", # gemma3n + MODEL_TENSOR.V_ENC_MSFA_EXP_NORM: "v.msfa.ffn.pw_exp.bn", # gemma3n + MODEL_TENSOR.V_ENC_MSFA_PROJ: "v.msfa.ffn.pw_proj.conv", # gemma3n + MODEL_TENSOR.V_ENC_MSFA_PROJ_NORM: "v.msfa.ffn.pw_proj.bn", # gemma3n + MODEL_TENSOR.V_ENC_MSFA_NORM: "v.msfa.norm", # gemma3n + MODEL_TENSOR.V_RESMPL_POS_EMBD_K: "resampler.pos_embd_k", + MODEL_TENSOR.V_RESMPL_ATTN_Q: "resampler.attn.q", + MODEL_TENSOR.V_RESMPL_ATTN_K: "resampler.attn.k", + MODEL_TENSOR.V_RESMPL_ATTN_V: "resampler.attn.v", + MODEL_TENSOR.V_RESMPL_ATTN_OUT: "resampler.attn.out", + MODEL_TENSOR.V_RESMPL_KV: "resampler.kv", + MODEL_TENSOR.V_RESMPL_KV_NORM: "resampler.ln_kv", + MODEL_TENSOR.V_RESMPL_POST_NORM: "resampler.ln_post", + MODEL_TENSOR.V_RESMPL_Q_NORM: "resampler.ln_q", + MODEL_TENSOR.V_RESMPL_PROJ: "resampler.proj", + MODEL_TENSOR.V_RESMPL_QUERY: "resampler.query", + MODEL_TENSOR.V_TOK_EMBD_IMG_BREAK: "v.token_embd.img_break", # pixtral + MODEL_TENSOR.V_TOK_EMBD_IMG_START: "v.token_embd.img_start", # deepseek4v + MODEL_TENSOR.V_TOK_EMBD_IMG_END: "v.token_embd.img_end", # deepseek4v + MODEL_TENSOR.V_TOK_EMBD_IMG_PAD: "v.token_embd.img_pad", # deepseek4v + MODEL_TENSOR.V_MM_PATCH_MERGER: "mm.patch_merger", # mistral small 3.1 + MODEL_TENSOR.V_DS_NORM: "v.deepstack.{bid}.norm", + MODEL_TENSOR.V_DS_FC1: "v.deepstack.{bid}.fc1", + MODEL_TENSOR.V_DS_FC2: "v.deepstack.{bid}.fc2", + MODEL_TENSOR.V_MERGER_LN1: "v.vit_merger.ln1", + MODEL_TENSOR.V_MERGER_ATTN_Q: "v.vit_merger.attn_q", + MODEL_TENSOR.V_MERGER_ATTN_K: "v.vit_merger.attn_k", + MODEL_TENSOR.V_MERGER_ATTN_V: "v.vit_merger.attn_v", + MODEL_TENSOR.V_MERGER_ATTN_O: "v.vit_merger.attn_out", + MODEL_TENSOR.V_MERGER_DS_LN: "v.vit_merger.ds_ln", + MODEL_TENSOR.V_MERGER_DS_UP: "v.vit_merger.ds_ffn_up", + MODEL_TENSOR.V_MERGER_DS_DOWN: "v.vit_merger.ds_ffn_down", + MODEL_TENSOR.V_MM_POST_FC_NORM: "mm.post_fc_norm", # cogvlm + MODEL_TENSOR.V_MM_UP: "mm.up", + MODEL_TENSOR.V_MM_DOWN: "mm.down", + MODEL_TENSOR.V_MM_GATE: "mm.gate", + MODEL_TENSOR.V_MM_MERGER_FC1: "mm.merger.fc1", + MODEL_TENSOR.V_MM_MERGER_FC2: "mm.merger.fc2", + MODEL_TENSOR.V_TOK_BOI: "v.boi", + MODEL_TENSOR.V_TOK_EOI: "v.eoi", + MODEL_TENSOR.V_MM_PRE_NORM: "mm.pre_norm", + MODEL_TENSOR.V_TOK_IMG_BEGIN: "mm.image_begin", + MODEL_TENSOR.V_TOK_IMG_END: "mm.image_end", + MODEL_TENSOR.V_STD_BIAS: "v.std_bias", # gemma4 + MODEL_TENSOR.V_STD_SCALE: "v.std_scale", # gemma4 + # DeepSeek-OCR SAM + MODEL_TENSOR.V_SAM_POS_EMBD: "v.sam.pos_embd", + MODEL_TENSOR.V_SAM_PATCH_EMBD: "v.sam.patch_embd", + MODEL_TENSOR.V_SAM_PRE_NORM: "v.sam.blk.{bid}.pre_ln", + MODEL_TENSOR.V_SAM_POST_NORM: "v.sam.blk.{bid}.post_ln", + MODEL_TENSOR.V_SAM_ATTN_POS_H: "v.sam.blk.{bid}.attn.pos_h", + MODEL_TENSOR.V_SAM_ATTN_POS_W: "v.sam.blk.{bid}.attn.pos_w", + MODEL_TENSOR.V_SAM_ATTN_QKV: "v.sam.blk.{bid}.attn.qkv", + MODEL_TENSOR.V_SAM_ATTN_OUT: "v.sam.blk.{bid}.attn.out", + MODEL_TENSOR.V_SAM_MLP_LIN_1: "v.sam.blk.{bid}.mlp.lin1", + MODEL_TENSOR.V_SAM_MLP_LIN_2: "v.sam.blk.{bid}.mlp.lin2", + MODEL_TENSOR.V_SAM_NECK: "v.sam.neck.{bid}", + MODEL_TENSOR.V_SAM_NET_2: "v.sam.net_2", + MODEL_TENSOR.V_SAM_NET_3: "v.sam.net_3", + MODEL_TENSOR.V_ENC_EMBD_IMGNL: "v.image_newline", # Deepseek-OCR, Granite4Vision + MODEL_TENSOR.V_ENC_EMBD_VSEP: "v.view_seperator", # Deepseek-OCR + MODEL_TENSOR.V_TOK_EMBD_SEP: "v.tok_embd_sep", # MiniCPM-V 4.7 + MODEL_TENSOR.V_RESMPL_QUERY_768: "v.resample_query_768", # Deepseek-OCR-2 qwen2 + MODEL_TENSOR.V_RESMPL_QUERY_1024: "v.resample_query_1024", # Deepseek-OCR-2 qwen2 + # Granite4 Vision + # qformer layers (bid => proj_id) + # NOTE: Names align with A_QF_* + MODEL_TENSOR.V_QF_SELF_ATTN_Q: "v.proj_blk.{bid}.self_attn_q", + MODEL_TENSOR.V_QF_SELF_ATTN_K: "v.proj_blk.{bid}.self_attn_k", + MODEL_TENSOR.V_QF_SELF_ATTN_V: "v.proj_blk.{bid}.self_attn_v", + MODEL_TENSOR.V_QF_SELF_ATTN_O: "v.proj_blk.{bid}.self_attn_out", + MODEL_TENSOR.V_QF_SELF_ATTN_NORM: "v.proj_blk.{bid}.self_attn_norm", + MODEL_TENSOR.V_QF_CROSS_ATTN_Q: "v.proj_blk.{bid}.cross_attn_q", + MODEL_TENSOR.V_QF_CROSS_ATTN_K: "v.proj_blk.{bid}.cross_attn_k", + MODEL_TENSOR.V_QF_CROSS_ATTN_V: "v.proj_blk.{bid}.cross_attn_v", + MODEL_TENSOR.V_QF_CROSS_ATTN_O: "v.proj_blk.{bid}.cross_attn_out", + MODEL_TENSOR.V_QF_CROSS_ATTN_NORM: "v.proj_blk.{bid}.cross_attn_norm", + MODEL_TENSOR.V_QF_FFN_UP: "v.proj_blk.{bid}.ffn_up", + MODEL_TENSOR.V_QF_FFN_DOWN: "v.proj_blk.{bid}.ffn_down", + MODEL_TENSOR.V_QF_FFN_NORM: "v.proj_blk.{bid}.ffn_norm", + # multi-projector (bid => projector ID) + MODEL_TENSOR.V_MULTI_PROJ_IMG_POS: "v.proj_blk.{bid}.img_pos", + MODEL_TENSOR.V_MULTI_PROJ_QUERY: "v.proj_blk.{bid}.query", + MODEL_TENSOR.V_MULTI_PROJ_NORM: "v.proj_blk.{bid}.norm", + MODEL_TENSOR.V_MULTI_PROJ_LINEAR: "v.proj_blk.{bid}.linear", + MODEL_TENSOR.V_MULTI_PROJ_POST_NORM: "v.proj_blk.{bid}.post_norm", + + # audio (mtmd) + # note: all audio tensor names must use prefix "a." or "mm.a." + MODEL_TENSOR.A_ENC_EMBD_POS: "a.position_embd", + MODEL_TENSOR.A_ENC_EMBD_NORM: "a.position_embd_norm", + MODEL_TENSOR.A_ENC_EMBD_TO_LOGITS: "a.embd_to_logits", + MODEL_TENSOR.A_ENC_INP_PROJ: "a.input_projection", + MODEL_TENSOR.A_ENC_CONV1D: "a.conv1d.{bid}", + MODEL_TENSOR.A_ENC_CONV2D: "a.conv2d.{bid}", + MODEL_TENSOR.A_ENC_CONV_OUT: "a.conv_out", + MODEL_TENSOR.A_ENC_CONV1D_NORM: "a.conv1d.{bid}.norm", + MODEL_TENSOR.A_PRE_NORM: "a.pre_ln", + MODEL_TENSOR.A_POST_NORM: "a.post_ln", + MODEL_TENSOR.A_ENC_LAYER_PRE_NORM: "a.blk.{bid}.layer_pre_norm", + MODEL_TENSOR.A_ENC_ATTN_Q: "a.blk.{bid}.attn_q", + MODEL_TENSOR.A_ENC_ATTN_K: "a.blk.{bid}.attn_k", + MODEL_TENSOR.A_ENC_ATTN_V: "a.blk.{bid}.attn_v", + MODEL_TENSOR.A_ENC_ATTN_POST_NORM: "a.blk.{bid}.attn_post_norm", + MODEL_TENSOR.A_ENC_ATTN_PRE_NORM: "a.blk.{bid}.attn_pre_norm", + MODEL_TENSOR.A_ENC_ATTN_K_REL: "a.blk.{bid}.attn_k_rel", + MODEL_TENSOR.A_ENC_PER_DIM_SCALE: "a.blk.{bid}.per_dim_scale", + MODEL_TENSOR.A_ENC_INPUT_NORM: "a.blk.{bid}.ln1", + MODEL_TENSOR.A_ENC_OUTPUT: "a.blk.{bid}.attn_out", + MODEL_TENSOR.A_ENC_OUTPUT_NORM: "a.blk.{bid}.ln2", + MODEL_TENSOR.A_ENC_FFN_NORM: "a.blk.{bid}.ffn_norm", + MODEL_TENSOR.A_ENC_FFN_POST_NORM: "a.blk.{bid}.ffn_post_norm", + MODEL_TENSOR.A_ENC_FFN_SCALE: "a.blk.{bid}.ffn_scale", + MODEL_TENSOR.A_ENC_FFN_UP: "a.blk.{bid}.ffn_up", + MODEL_TENSOR.A_ENC_FFN_GATE: "a.blk.{bid}.ffn_gate", + MODEL_TENSOR.A_ENC_FFN_DOWN: "a.blk.{bid}.ffn_down", + MODEL_TENSOR.A_ENC_FFN_NORM_1: "a.blk.{bid}.ffn_norm_1", + MODEL_TENSOR.A_ENC_FFN_POST_NORM_1: "a.blk.{bid}.ffn_post_norm_1", + MODEL_TENSOR.A_ENC_FFN_SCALE_1: "a.blk.{bid}.ffn_scale_1", + MODEL_TENSOR.A_ENC_FFN_UP_1: "a.blk.{bid}.ffn_up_1", + MODEL_TENSOR.A_ENC_FFN_GATE_1: "a.blk.{bid}.ffn_gate_1", + MODEL_TENSOR.A_ENC_FFN_DOWN_1: "a.blk.{bid}.ffn_down_1", + MODEL_TENSOR.A_ENC_DOWNSAMPLE_CONV: "a.downsample.conv", + MODEL_TENSOR.A_ENC_DOWNSAMPLE_NORM: "a.downsample.norm", + MODEL_TENSOR.A_ENC_RVQ_CODEBOOK: "a.rvq.codebook", + MODEL_TENSOR.A_ENC_CONV_RES2: "a.blk.{bid}.res2.{xid}", + MODEL_TENSOR.A_ENC_SE_CONV1: "a.blk.{bid}.se_conv1", + MODEL_TENSOR.A_ENC_SE_CONV2: "a.blk.{bid}.se_conv2", + MODEL_TENSOR.A_ENC_ASP_ATTN: "a.asp_attn", + MODEL_TENSOR.A_ENC_ASP_TDNN: "a.asp_tdnn", + MODEL_TENSOR.A_GEN_CODE_PROJ_IN: "a.gen.code.proj_in", + MODEL_TENSOR.A_GEN_CODE_EMBD: "a.gen.code.embd", + MODEL_TENSOR.A_GEN_CODE_HEAD: "a.gen.code.head", + MODEL_TENSOR.A_GEN_CODE_OUT_EMBD: "a.gen.code.out_embd", + MODEL_TENSOR.A_GEN_CODE_ATTN_NORM: "a.gen.code.blk.{bid}.ln1", # reuses the generic clip.cpp block loader (TN_LN_1) + MODEL_TENSOR.A_GEN_CODE_ATTN_Q: "a.gen.code.blk.{bid}.attn_q", + MODEL_TENSOR.A_GEN_CODE_ATTN_Q_NORM: "a.gen.code.blk.{bid}.attn_q_norm", + MODEL_TENSOR.A_GEN_CODE_ATTN_K: "a.gen.code.blk.{bid}.attn_k", + MODEL_TENSOR.A_GEN_CODE_ATTN_K_NORM: "a.gen.code.blk.{bid}.attn_k_norm", + MODEL_TENSOR.A_GEN_CODE_ATTN_V: "a.gen.code.blk.{bid}.attn_v", + MODEL_TENSOR.A_GEN_CODE_ATTN_OUT: "a.gen.code.blk.{bid}.attn_out", + MODEL_TENSOR.A_GEN_CODE_FFN_NORM: "a.gen.code.blk.{bid}.ln2", # reuses the generic clip.cpp block loader (TN_LN_2) + MODEL_TENSOR.A_GEN_CODE_FFN_GATE: "a.gen.code.blk.{bid}.ffn_gate", + MODEL_TENSOR.A_GEN_CODE_FFN_UP: "a.gen.code.blk.{bid}.ffn_up", + MODEL_TENSOR.A_GEN_CODE_FFN_DOWN: "a.gen.code.blk.{bid}.ffn_down", + MODEL_TENSOR.A_GEN_CODE_OUTPUT_NORM: "a.gen.code.output_norm", + MODEL_TENSOR.A_GEN_WAV_QUANT_FIRST_IN: "a.gen.wav.quant.first.in_proj", + MODEL_TENSOR.A_GEN_WAV_QUANT_FIRST_OUT: "a.gen.wav.quant.first.out_proj", + MODEL_TENSOR.A_GEN_WAV_QUANT_FIRST_CB: "a.gen.wav.quant.first.codebook", + MODEL_TENSOR.A_GEN_WAV_QUANT_REST_IN: "a.gen.wav.quant.rest.in_proj", + MODEL_TENSOR.A_GEN_WAV_QUANT_REST_OUT: "a.gen.wav.quant.rest.out_proj", + MODEL_TENSOR.A_GEN_WAV_QUANT_REST_CB: "a.gen.wav.quant.rest.codebook", + MODEL_TENSOR.A_GEN_WAV_PRE_CONV: "a.gen.wav.pre_conv", + MODEL_TENSOR.A_GEN_WAV_TFM_IN_PROJ: "a.gen.wav.tfm.in_proj", + MODEL_TENSOR.A_GEN_WAV_TFM_OUT_PROJ: "a.gen.wav.tfm.out_proj", + MODEL_TENSOR.A_GEN_WAV_TFM_OUTPUT_NORM: "a.gen.wav.tfm.output_norm", + MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_NORM: "a.gen.wav.tfm.blk.{bid}.ln1", + MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_Q: "a.gen.wav.tfm.blk.{bid}.attn_q", + MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_K: "a.gen.wav.tfm.blk.{bid}.attn_k", + MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_V: "a.gen.wav.tfm.blk.{bid}.attn_v", + MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_OUT: "a.gen.wav.tfm.blk.{bid}.attn_out", + MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_SCALE: "a.gen.wav.tfm.blk.{bid}.ls1", + MODEL_TENSOR.A_GEN_WAV_TFM_FFN_NORM: "a.gen.wav.tfm.blk.{bid}.ln2", + MODEL_TENSOR.A_GEN_WAV_TFM_FFN_GATE: "a.gen.wav.tfm.blk.{bid}.ffn_gate", + MODEL_TENSOR.A_GEN_WAV_TFM_FFN_UP: "a.gen.wav.tfm.blk.{bid}.ffn_up", + MODEL_TENSOR.A_GEN_WAV_TFM_FFN_DOWN: "a.gen.wav.tfm.blk.{bid}.ffn_down", + MODEL_TENSOR.A_GEN_WAV_TFM_FFN_SCALE: "a.gen.wav.tfm.blk.{bid}.ls2", + MODEL_TENSOR.A_GEN_WAV_UP_CONV: "a.gen.wav.up.blk.{bid}.conv", + MODEL_TENSOR.A_GEN_WAV_UP_DWCONV: "a.gen.wav.up.blk.{bid}.dwconv", + MODEL_TENSOR.A_GEN_WAV_UP_NORM: "a.gen.wav.up.blk.{bid}.norm", + MODEL_TENSOR.A_GEN_WAV_UP_PW1: "a.gen.wav.up.blk.{bid}.pw1", + MODEL_TENSOR.A_GEN_WAV_UP_PW2: "a.gen.wav.up.blk.{bid}.pw2", + MODEL_TENSOR.A_GEN_WAV_UP_GAMMA: "a.gen.wav.up.blk.{bid}.gamma", + MODEL_TENSOR.A_GEN_WAV_DAC_ENTRY: "a.gen.wav.dac.entry", + MODEL_TENSOR.A_GEN_WAV_DAC_UP_SNAKE: "a.gen.wav.dac.blk.{bid}.snake", + MODEL_TENSOR.A_GEN_WAV_DAC_UP_CONV: "a.gen.wav.dac.blk.{bid}.conv", + MODEL_TENSOR.A_GEN_WAV_DAC_RES_ACT1: "a.gen.wav.dac.blk.{bid}.res.{xid}.act1", + MODEL_TENSOR.A_GEN_WAV_DAC_RES_CONV1: "a.gen.wav.dac.blk.{bid}.res.{xid}.conv1", + MODEL_TENSOR.A_GEN_WAV_DAC_RES_ACT2: "a.gen.wav.dac.blk.{bid}.res.{xid}.act2", + MODEL_TENSOR.A_GEN_WAV_DAC_RES_CONV2: "a.gen.wav.dac.blk.{bid}.res.{xid}.conv2", + MODEL_TENSOR.A_GEN_WAV_DAC_POST_SNAKE: "a.gen.wav.dac.post_snake", + MODEL_TENSOR.A_GEN_WAV_DAC_POST_CONV: "a.gen.wav.dac.post_conv", + MODEL_TENSOR.A_ENC_SEANET_CONV_IN: "a.seanet.conv_in", + MODEL_TENSOR.A_ENC_SEANET_CONV_OUT: "a.seanet.conv_out", + MODEL_TENSOR.A_ENC_SEANET_RES_CONV1: "a.seanet.blk.{bid}.res_conv1", + MODEL_TENSOR.A_ENC_SEANET_RES_CONV2: "a.seanet.blk.{bid}.res_conv2", + MODEL_TENSOR.A_ENC_SEANET_SCALE_CONV: "a.seanet.blk.{bid}.scale_conv", + MODEL_TENSOR.A_ENC_ATTN_SCALE: "a.blk.{bid}.ls1", + MODEL_TENSOR.A_ENC_FFN_SCALE_LS: "a.blk.{bid}.ls2", + MODEL_TENSOR.A_ENC_SPEAKER_PROJ: "a.speaker_proj", + MODEL_TENSOR.A_GEN_FLOW_INPUT_PROJ: "a.gen.flow.input_proj", + MODEL_TENSOR.A_GEN_FLOW_COND_EMBD: "a.gen.flow.cond_embd", + MODEL_TENSOR.A_GEN_FLOW_TIME_FREQS: "a.gen.flow.time.{bid}.freqs", + MODEL_TENSOR.A_GEN_FLOW_TIME_UP: "a.gen.flow.time.{bid}.up", + MODEL_TENSOR.A_GEN_FLOW_TIME_DOWN: "a.gen.flow.time.{bid}.down", + MODEL_TENSOR.A_GEN_FLOW_TIME_NORM: "a.gen.flow.time.{bid}.norm", + MODEL_TENSOR.A_GEN_FLOW_BLK_NORM: "a.gen.flow.blk.{bid}.norm", + MODEL_TENSOR.A_GEN_FLOW_BLK_UP: "a.gen.flow.blk.{bid}.up", + MODEL_TENSOR.A_GEN_FLOW_BLK_DOWN: "a.gen.flow.blk.{bid}.down", + MODEL_TENSOR.A_GEN_FLOW_BLK_ADA: "a.gen.flow.blk.{bid}.ada", + MODEL_TENSOR.A_GEN_FLOW_FINAL_ADA: "a.gen.flow.final.ada", + MODEL_TENSOR.A_GEN_FLOW_FINAL_PROJ: "a.gen.flow.final.proj", + MODEL_TENSOR.A_GEN_OUT_EOS: "a.gen.out_eos", + MODEL_TENSOR.A_GEN_INPUT_LINEAR: "a.gen.input_linear", + MODEL_TENSOR.A_GEN_EMB_MEAN: "a.gen.emb_mean", + MODEL_TENSOR.A_GEN_EMB_STD: "a.gen.emb_std", + MODEL_TENSOR.A_GEN_WAV_QUANT_OUT: "a.gen.wav.quant_out", + MODEL_TENSOR.A_GEN_WAV_UPSAMPLE: "a.gen.wav.upsample", + MODEL_TENSOR.A_GEN_WAV_SEANET_CONV_IN: "a.gen.wav.seanet.conv_in", + MODEL_TENSOR.A_GEN_WAV_SEANET_CONV_OUT: "a.gen.wav.seanet.conv_out", + MODEL_TENSOR.A_GEN_WAV_SEANET_RES_CONV1: "a.gen.wav.seanet.blk.{bid}.res_conv1", + MODEL_TENSOR.A_GEN_WAV_SEANET_RES_CONV2: "a.gen.wav.seanet.blk.{bid}.res_conv2", + MODEL_TENSOR.A_GEN_WAV_SEANET_SCALE_CONV: "a.gen.wav.seanet.blk.{bid}.scale_conv", + MODEL_TENSOR.A_MMPROJ: "mm.a.mlp.{bid}", + MODEL_TENSOR.A_MMPROJ_FC: "mm.a.fc", + MODEL_TENSOR.A_MM_NORM_PRE: "mm.a.norm_pre", + MODEL_TENSOR.A_MM_NORM_MID: "mm.a.norm_mid", + MODEL_TENSOR.A_MM_INP_PROJ: "mm.a.input_projection", # gemma3n + MODEL_TENSOR.A_MM_SOFT_EMB_NORM: "mm.a.soft_emb_norm", # gemma3n + MODEL_TENSOR.A_MM_EMBEDDING: "mm.a.embedding", # gemma3n + MODEL_TENSOR.A_MM_HARD_EMB_NORM: "mm.a.hard_emb_norm", # gemma3n + MODEL_TENSOR.A_MM_CODE_EMBD: "mm.a.code_embd", + MODEL_TENSOR.A_MM_LOCAL_ATTN_Q: "mm.a.local_blk.{bid}.attn_q", + MODEL_TENSOR.A_MM_LOCAL_ATTN_K: "mm.a.local_blk.{bid}.attn_k", + MODEL_TENSOR.A_MM_LOCAL_ATTN_V: "mm.a.local_blk.{bid}.attn_v", + MODEL_TENSOR.A_MM_LOCAL_ATTN_OUT: "mm.a.local_blk.{bid}.attn_out", + MODEL_TENSOR.A_MM_LOCAL_FFN_GATE: "mm.a.local_blk.{bid}.ffn_gate", + MODEL_TENSOR.A_MM_LOCAL_FFN_UP: "mm.a.local_blk.{bid}.ffn_up", + MODEL_TENSOR.A_MM_LOCAL_FFN_DOWN: "mm.a.local_blk.{bid}.ffn_down", + MODEL_TENSOR.A_MM_LOCAL_LN1: "mm.a.local_blk.{bid}.ln1", + MODEL_TENSOR.A_MM_LOCAL_LN2: "mm.a.local_blk.{bid}.ln2", + MODEL_TENSOR.A_MM_LOCAL_NORM: "mm.a.local_norm", + MODEL_TENSOR.A_PER_DIM_K_SCALE: "a.blk.{bid}.per_dim_k_scale", # gemma4 + MODEL_TENSOR.A_PER_DIM_SCALE: "a.blk.{bid}.per_dim_scale", # gemma4 + # lfm2 audio + MODEL_TENSOR.A_ENC_NORM_CONV: "a.blk.{bid}.norm_conv", + MODEL_TENSOR.A_ENC_LINEAR_POS: "a.blk.{bid}.linear_pos", + MODEL_TENSOR.A_ENC_POS_BIAS_U: "a.blk.{bid}.pos_bias_u", + MODEL_TENSOR.A_ENC_POS_BIAS_V: "a.blk.{bid}.pos_bias_v", + MODEL_TENSOR.A_ENC_OUT: "a.pre_encode.out", + MODEL_TENSOR.A_ENC_CONV_DW: "a.blk.{bid}.conv_dw", + MODEL_TENSOR.A_ENC_CONV_NORM: "a.blk.{bid}.conv_norm", + MODEL_TENSOR.A_ENC_CONV_PW1: "a.blk.{bid}.conv_pw1", + MODEL_TENSOR.A_ENC_CONV_PW2: "a.blk.{bid}.conv_pw2", + MODEL_TENSOR.A_ENC_CONV_NORM_MEAN: "a.blk.{bid}.conv_norm_mean", + MODEL_TENSOR.A_ENC_CONV_NORM_VAR: "a.blk.{bid}.conv_norm_var", + MODEL_TENSOR.A_ENC_MEL_FILTERS: "a.mel_filters", + MODEL_TENSOR.A_ENC_WINDOW: "a.window", + MODEL_TENSOR.A_CTC_OUT: "a.enc_ctc_out", + MODEL_TENSOR.A_CTC_OUT_MID: "a.enc_ctc_out_mid", + MODEL_TENSOR.A_ENC_ATTN_REL_POS_EMB: "a.blk.{bid}.attn_rel_pos_emb", + # qformer projector + MODEL_TENSOR.A_QF_PROJ_QUERY: "a.proj_query", + MODEL_TENSOR.A_QF_PROJ_NORM: "a.proj_norm", + MODEL_TENSOR.A_QF_PROJ_LINEAR: "a.proj_linear", + MODEL_TENSOR.A_QF_SELF_ATTN_Q: "a.proj_blk.{bid}.self_attn_q", + MODEL_TENSOR.A_QF_SELF_ATTN_K: "a.proj_blk.{bid}.self_attn_k", + MODEL_TENSOR.A_QF_SELF_ATTN_V: "a.proj_blk.{bid}.self_attn_v", + MODEL_TENSOR.A_QF_SELF_ATTN_O: "a.proj_blk.{bid}.self_attn_out", + MODEL_TENSOR.A_QF_SELF_ATTN_NORM: "a.proj_blk.{bid}.self_attn_norm", + MODEL_TENSOR.A_QF_CROSS_ATTN_Q: "a.proj_blk.{bid}.cross_attn_q", + MODEL_TENSOR.A_QF_CROSS_ATTN_K: "a.proj_blk.{bid}.cross_attn_k", + MODEL_TENSOR.A_QF_CROSS_ATTN_V: "a.proj_blk.{bid}.cross_attn_v", + MODEL_TENSOR.A_QF_CROSS_ATTN_O: "a.proj_blk.{bid}.cross_attn_out", + MODEL_TENSOR.A_QF_CROSS_ATTN_NORM: "a.proj_blk.{bid}.cross_attn_norm", + MODEL_TENSOR.A_QF_FFN_UP: "a.proj_blk.{bid}.ffn_up", + MODEL_TENSOR.A_QF_FFN_DOWN: "a.proj_blk.{bid}.ffn_down", + MODEL_TENSOR.A_QF_FFN_NORM: "a.proj_blk.{bid}.ffn_norm", + # NextN/MTP + MODEL_TENSOR.NEXTN_PROJ_PRE: "nextn.pre_projection", + MODEL_TENSOR.NEXTN_PROJ_POST: "nextn.post_projection", + MODEL_TENSOR.NEXTN_EH_PROJ: "blk.{bid}.nextn.eh_proj", + MODEL_TENSOR.NEXTN_EMBED_TOKENS: "blk.{bid}.nextn.embed_tokens", + MODEL_TENSOR.NEXTN_ENORM: "blk.{bid}.nextn.enorm", + MODEL_TENSOR.NEXTN_HNORM: "blk.{bid}.nextn.hnorm", + MODEL_TENSOR.NEXTN_SHARED_HEAD_HEAD: "blk.{bid}.nextn.shared_head_head", + MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM: "blk.{bid}.nextn.shared_head_norm", + MODEL_TENSOR.NEXTN_HC_HEAD_NORM: "blk.{bid}.nextn.hc_head_norm", + MODEL_TENSOR.NEXTN_HC_HEAD_DOWN: "blk.{bid}.nextn.hc_head_down", + MODEL_TENSOR.NEXTN_HC_HEAD_UP: "blk.{bid}.nextn.hc_head_up", + MODEL_TENSOR.FC: "fc", + MODEL_TENSOR.DSPARK_MARKOV_W1: "markov_w1", + MODEL_TENSOR.DSPARK_MARKOV_W2: "markov_w2", + MODEL_TENSOR.DSPARK_CONF_PROJ: "conf_proj", + MODEL_TENSOR.DFLASH_ATTN_CONV_BASE: "blk.{bid}.attn_conv_base", + MODEL_TENSOR.DFLASH_ATTN_CONV_PROJ: "blk.{bid}.attn_conv_proj", + MODEL_TENSOR.DFLASH_FFN_CONV_BASE: "blk.{bid}.ffn_conv_base", + MODEL_TENSOR.DFLASH_FFN_CONV_PROJ: "blk.{bid}.ffn_conv_proj", + MODEL_TENSOR.DFLASH_SELECTOR_PREV: "selector_predecessor", + MODEL_TENSOR.DFLASH_SELECTOR_NEXT: "selector_successor", + MODEL_TENSOR.DFLASH_SELECTOR_HIDDEN: "selector_hidden", + MODEL_TENSOR.D2T: "d2t", + MODEL_TENSOR.ATTN_V_GATE: "blk.{bid}.attn_v_gate", + MODEL_TENSOR.ATTN_V_EXP: "blk.{bid}.attn_v_exps", +} + +MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { + MODEL_ARCH.MMPROJ: [ + MODEL_TENSOR.V_MMPROJ, + MODEL_TENSOR.V_MMPROJ_FC, + MODEL_TENSOR.V_MMPROJ_MLP, + MODEL_TENSOR.V_MMPROJ_PEG, + MODEL_TENSOR.V_ENC_EMBD_CLS, + MODEL_TENSOR.V_ENC_EMBD_PATCH, + MODEL_TENSOR.V_ENC_EMBD_NORM, + MODEL_TENSOR.V_ENC_EMBD_PATCH_NORM, + MODEL_TENSOR.V_ENC_EMBD_POS, + MODEL_TENSOR.V_ENC_EMBD_IMGNL, + MODEL_TENSOR.V_ENC_EMBD_VSEP, + MODEL_TENSOR.V_TOK_EMBD_SEP, + MODEL_TENSOR.V_ENC_INPUT_NORM, + MODEL_TENSOR.V_ENC_ATTN_QKV, + MODEL_TENSOR.V_ENC_ATTN_Q, + MODEL_TENSOR.V_ENC_ATTN_Q_NORM, + MODEL_TENSOR.V_ENC_ATTN_K, + MODEL_TENSOR.V_ENC_ATTN_K_NORM, + MODEL_TENSOR.V_ENC_ATTN_V, + MODEL_TENSOR.V_ENC_ATTN_O, + MODEL_TENSOR.V_ENC_ATTN_O_NORM, + MODEL_TENSOR.V_ENC_ATTN_SINKS, + MODEL_TENSOR.V_ENC_POST_ATTN_NORM, + MODEL_TENSOR.V_ENC_FFN_UP, + MODEL_TENSOR.V_ENC_FFN_GATE, + MODEL_TENSOR.V_ENC_FFN_DOWN, + MODEL_TENSOR.V_ENC_FFN_GATE_INP, + MODEL_TENSOR.V_ENC_FFN_GATE_EXPS, + MODEL_TENSOR.V_ENC_FFN_UP_EXPS, + MODEL_TENSOR.V_ENC_FFN_DOWN_EXPS, + MODEL_TENSOR.V_ENC_FFN_EXP_PROBS_B, + MODEL_TENSOR.V_ENC_ATTN_POST_NORM, + MODEL_TENSOR.V_ENC_FFN_POST_NORM, + MODEL_TENSOR.V_LAYER_SCALE_1, + MODEL_TENSOR.V_LAYER_SCALE_2, + MODEL_TENSOR.V_LAYER_OUT_SCALE, + MODEL_TENSOR.V_PRE_NORM, + MODEL_TENSOR.V_POST_NORM, + MODEL_TENSOR.V_MM_POST_NORM, + MODEL_TENSOR.V_MM_INP_PROJ, + MODEL_TENSOR.V_MM_INP_NORM, + MODEL_TENSOR.V_MM_SOFT_EMB_NORM, + MODEL_TENSOR.V_MM_EMBEDDING, + MODEL_TENSOR.V_MM_HARD_EMB_NORM, + MODEL_TENSOR.V_ENC_CONV_STEM, + MODEL_TENSOR.V_ENC_CONV_STEM_NORM, + MODEL_TENSOR.V_ENC_MSFA_EXP, + MODEL_TENSOR.V_ENC_MSFA_EXP_NORM, + MODEL_TENSOR.V_ENC_MSFA_PROJ, + MODEL_TENSOR.V_ENC_MSFA_PROJ_NORM, + MODEL_TENSOR.V_ENC_MSFA_NORM, + MODEL_TENSOR.V_RESMPL_POS_EMBD_K, + MODEL_TENSOR.V_RESMPL_ATTN_Q, + MODEL_TENSOR.V_RESMPL_ATTN_K, + MODEL_TENSOR.V_RESMPL_ATTN_V, + MODEL_TENSOR.V_RESMPL_ATTN_OUT, + MODEL_TENSOR.V_RESMPL_KV, + MODEL_TENSOR.V_RESMPL_KV_NORM, + MODEL_TENSOR.V_RESMPL_POST_NORM, + MODEL_TENSOR.V_RESMPL_Q_NORM, + MODEL_TENSOR.V_RESMPL_PROJ, + MODEL_TENSOR.V_RESMPL_QUERY, + MODEL_TENSOR.V_TOK_EMBD_IMG_BREAK, + MODEL_TENSOR.V_TOK_EMBD_IMG_START, + MODEL_TENSOR.V_TOK_EMBD_IMG_END, + MODEL_TENSOR.V_TOK_EMBD_IMG_PAD, + MODEL_TENSOR.V_MM_PATCH_MERGER, + MODEL_TENSOR.V_MM_MERGER_FC1, + MODEL_TENSOR.V_MM_MERGER_FC2, + MODEL_TENSOR.V_DS_NORM, + MODEL_TENSOR.V_DS_FC1, + MODEL_TENSOR.V_DS_FC2, + MODEL_TENSOR.V_MERGER_LN1, + MODEL_TENSOR.V_MERGER_ATTN_Q, + MODEL_TENSOR.V_MERGER_ATTN_K, + MODEL_TENSOR.V_MERGER_ATTN_V, + MODEL_TENSOR.V_MERGER_ATTN_O, + MODEL_TENSOR.V_MERGER_DS_LN, + MODEL_TENSOR.V_MERGER_DS_UP, + MODEL_TENSOR.V_MERGER_DS_DOWN, + MODEL_TENSOR.V_MM_POST_FC_NORM, + MODEL_TENSOR.V_MM_UP, + MODEL_TENSOR.V_MM_DOWN, + MODEL_TENSOR.V_MM_GATE, + MODEL_TENSOR.V_TOK_BOI, + MODEL_TENSOR.V_TOK_EOI, + MODEL_TENSOR.V_MM_PRE_NORM, + MODEL_TENSOR.V_TOK_IMG_BEGIN, + MODEL_TENSOR.V_TOK_IMG_END, + MODEL_TENSOR.V_STD_BIAS, + MODEL_TENSOR.V_STD_SCALE, + MODEL_TENSOR.V_SAM_POS_EMBD, + MODEL_TENSOR.V_SAM_PATCH_EMBD, + MODEL_TENSOR.V_SAM_PRE_NORM, + MODEL_TENSOR.V_SAM_POST_NORM, + MODEL_TENSOR.V_SAM_ATTN_POS_H, + MODEL_TENSOR.V_SAM_ATTN_POS_W, + MODEL_TENSOR.V_SAM_ATTN_QKV, + MODEL_TENSOR.V_SAM_ATTN_OUT, + MODEL_TENSOR.V_SAM_MLP_LIN_1, + MODEL_TENSOR.V_SAM_MLP_LIN_2, + MODEL_TENSOR.V_SAM_NECK, + MODEL_TENSOR.V_SAM_NET_2, + MODEL_TENSOR.V_SAM_NET_3, + MODEL_TENSOR.V_RESMPL_QUERY_768, + MODEL_TENSOR.V_RESMPL_QUERY_1024, + MODEL_TENSOR.V_PROJ_NORM, + MODEL_TENSOR.V_QF_PROJ_QUERY, + MODEL_TENSOR.V_QF_PROJ_NORM, + MODEL_TENSOR.V_QF_PROJ_LINEAR, + MODEL_TENSOR.V_QF_SELF_ATTN_Q, + MODEL_TENSOR.V_QF_SELF_ATTN_K, + MODEL_TENSOR.V_QF_SELF_ATTN_V, + MODEL_TENSOR.V_QF_SELF_ATTN_O, + MODEL_TENSOR.V_QF_SELF_ATTN_NORM, + MODEL_TENSOR.V_QF_CROSS_ATTN_Q, + MODEL_TENSOR.V_QF_CROSS_ATTN_K, + MODEL_TENSOR.V_QF_CROSS_ATTN_V, + MODEL_TENSOR.V_QF_CROSS_ATTN_O, + MODEL_TENSOR.V_QF_CROSS_ATTN_NORM, + MODEL_TENSOR.V_QF_FFN_UP, + MODEL_TENSOR.V_QF_FFN_DOWN, + MODEL_TENSOR.V_QF_FFN_NORM, + MODEL_TENSOR.V_QF_PROJ_NORM, + MODEL_TENSOR.V_MULTI_PROJ_IMG_POS, + MODEL_TENSOR.V_MULTI_PROJ_QUERY, + MODEL_TENSOR.V_MULTI_PROJ_LINEAR, + MODEL_TENSOR.V_MULTI_PROJ_NORM, + MODEL_TENSOR.V_MULTI_PROJ_POST_NORM, + # audio + MODEL_TENSOR.A_ENC_EMBD_POS, + MODEL_TENSOR.A_ENC_EMBD_NORM, + MODEL_TENSOR.A_ENC_EMBD_TO_LOGITS, + MODEL_TENSOR.A_ENC_INP_PROJ, + MODEL_TENSOR.A_ENC_CONV1D, + MODEL_TENSOR.A_ENC_CONV2D, + MODEL_TENSOR.A_ENC_CONV_OUT, + MODEL_TENSOR.A_ENC_CONV1D_NORM, + MODEL_TENSOR.A_PRE_NORM, + MODEL_TENSOR.A_POST_NORM, + MODEL_TENSOR.A_ENC_LAYER_PRE_NORM, + MODEL_TENSOR.A_ENC_ATTN_Q, + MODEL_TENSOR.A_ENC_ATTN_K, + MODEL_TENSOR.A_ENC_ATTN_V, + MODEL_TENSOR.A_ENC_ATTN_POST_NORM, + MODEL_TENSOR.A_ENC_ATTN_PRE_NORM, + MODEL_TENSOR.A_ENC_ATTN_K_REL, + MODEL_TENSOR.A_ENC_PER_DIM_SCALE, + MODEL_TENSOR.A_ENC_INPUT_NORM, + MODEL_TENSOR.A_ENC_OUTPUT, + MODEL_TENSOR.A_ENC_OUTPUT_NORM, + MODEL_TENSOR.A_ENC_FFN_NORM, + MODEL_TENSOR.A_ENC_FFN_POST_NORM, + MODEL_TENSOR.A_ENC_FFN_SCALE, + MODEL_TENSOR.A_ENC_FFN_UP, + MODEL_TENSOR.A_ENC_FFN_GATE, + MODEL_TENSOR.A_ENC_FFN_DOWN, + MODEL_TENSOR.A_ENC_FFN_NORM_1, + MODEL_TENSOR.A_ENC_FFN_POST_NORM_1, + MODEL_TENSOR.A_ENC_FFN_SCALE_1, + MODEL_TENSOR.A_ENC_FFN_UP_1, + MODEL_TENSOR.A_ENC_FFN_GATE_1, + MODEL_TENSOR.A_ENC_FFN_DOWN_1, + MODEL_TENSOR.A_ENC_DOWNSAMPLE_CONV, + MODEL_TENSOR.A_ENC_DOWNSAMPLE_NORM, + MODEL_TENSOR.A_ENC_RVQ_CODEBOOK, + MODEL_TENSOR.A_MMPROJ, + MODEL_TENSOR.A_MMPROJ_FC, + MODEL_TENSOR.A_MM_NORM_PRE, + MODEL_TENSOR.A_MM_NORM_MID, + MODEL_TENSOR.A_MM_CODE_EMBD, + MODEL_TENSOR.A_MM_LOCAL_ATTN_Q, + MODEL_TENSOR.A_MM_LOCAL_ATTN_K, + MODEL_TENSOR.A_MM_LOCAL_ATTN_V, + MODEL_TENSOR.A_MM_LOCAL_ATTN_OUT, + MODEL_TENSOR.A_MM_LOCAL_FFN_GATE, + MODEL_TENSOR.A_MM_LOCAL_FFN_UP, + MODEL_TENSOR.A_MM_LOCAL_FFN_DOWN, + MODEL_TENSOR.A_MM_LOCAL_LN1, + MODEL_TENSOR.A_MM_LOCAL_LN2, + MODEL_TENSOR.A_MM_LOCAL_NORM, + MODEL_TENSOR.A_ENC_NORM_CONV, + MODEL_TENSOR.A_ENC_LINEAR_POS, + MODEL_TENSOR.A_ENC_POS_BIAS_U, + MODEL_TENSOR.A_ENC_POS_BIAS_V, + MODEL_TENSOR.A_ENC_OUT, + MODEL_TENSOR.A_ENC_CONV_DW, + MODEL_TENSOR.A_ENC_CONV_NORM, + MODEL_TENSOR.A_ENC_CONV_PW1, + MODEL_TENSOR.A_ENC_CONV_PW2, + MODEL_TENSOR.A_ENC_CONV_RES2, + MODEL_TENSOR.A_ENC_SE_CONV1, + MODEL_TENSOR.A_ENC_SE_CONV2, + MODEL_TENSOR.A_ENC_ASP_ATTN, + MODEL_TENSOR.A_ENC_ASP_TDNN, + MODEL_TENSOR.A_GEN_CODE_PROJ_IN, + MODEL_TENSOR.A_GEN_CODE_EMBD, + MODEL_TENSOR.A_GEN_CODE_HEAD, + MODEL_TENSOR.A_GEN_CODE_OUT_EMBD, + MODEL_TENSOR.A_GEN_CODE_ATTN_NORM, + MODEL_TENSOR.A_GEN_CODE_ATTN_Q, + MODEL_TENSOR.A_GEN_CODE_ATTN_Q_NORM, + MODEL_TENSOR.A_GEN_CODE_ATTN_K, + MODEL_TENSOR.A_GEN_CODE_ATTN_K_NORM, + MODEL_TENSOR.A_GEN_CODE_ATTN_V, + MODEL_TENSOR.A_GEN_CODE_ATTN_OUT, + MODEL_TENSOR.A_GEN_CODE_FFN_NORM, + MODEL_TENSOR.A_GEN_CODE_FFN_GATE, + MODEL_TENSOR.A_GEN_CODE_FFN_UP, + MODEL_TENSOR.A_GEN_CODE_FFN_DOWN, + MODEL_TENSOR.A_GEN_CODE_OUTPUT_NORM, + MODEL_TENSOR.A_GEN_WAV_QUANT_FIRST_IN, + MODEL_TENSOR.A_GEN_WAV_QUANT_FIRST_OUT, + MODEL_TENSOR.A_GEN_WAV_QUANT_FIRST_CB, + MODEL_TENSOR.A_GEN_WAV_QUANT_REST_IN, + MODEL_TENSOR.A_GEN_WAV_QUANT_REST_OUT, + MODEL_TENSOR.A_GEN_WAV_QUANT_REST_CB, + MODEL_TENSOR.A_GEN_WAV_PRE_CONV, + MODEL_TENSOR.A_GEN_WAV_TFM_IN_PROJ, + MODEL_TENSOR.A_GEN_WAV_TFM_OUT_PROJ, + MODEL_TENSOR.A_GEN_WAV_TFM_OUTPUT_NORM, + MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_NORM, + MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_Q, + MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_K, + MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_V, + MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_OUT, + MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_SCALE, + MODEL_TENSOR.A_GEN_WAV_TFM_FFN_NORM, + MODEL_TENSOR.A_GEN_WAV_TFM_FFN_GATE, + MODEL_TENSOR.A_GEN_WAV_TFM_FFN_UP, + MODEL_TENSOR.A_GEN_WAV_TFM_FFN_DOWN, + MODEL_TENSOR.A_GEN_WAV_TFM_FFN_SCALE, + MODEL_TENSOR.A_GEN_WAV_UP_CONV, + MODEL_TENSOR.A_GEN_WAV_UP_DWCONV, + MODEL_TENSOR.A_GEN_WAV_UP_NORM, + MODEL_TENSOR.A_GEN_WAV_UP_PW1, + MODEL_TENSOR.A_GEN_WAV_UP_PW2, + MODEL_TENSOR.A_GEN_WAV_UP_GAMMA, + MODEL_TENSOR.A_GEN_WAV_DAC_ENTRY, + MODEL_TENSOR.A_GEN_WAV_DAC_UP_SNAKE, + MODEL_TENSOR.A_GEN_WAV_DAC_UP_CONV, + MODEL_TENSOR.A_GEN_WAV_DAC_RES_ACT1, + MODEL_TENSOR.A_GEN_WAV_DAC_RES_CONV1, + MODEL_TENSOR.A_GEN_WAV_DAC_RES_ACT2, + MODEL_TENSOR.A_GEN_WAV_DAC_RES_CONV2, + MODEL_TENSOR.A_GEN_WAV_DAC_POST_SNAKE, + MODEL_TENSOR.A_GEN_WAV_DAC_POST_CONV, + MODEL_TENSOR.A_ENC_SEANET_CONV_IN, + MODEL_TENSOR.A_ENC_SEANET_CONV_OUT, + MODEL_TENSOR.A_ENC_SEANET_RES_CONV1, + MODEL_TENSOR.A_ENC_SEANET_RES_CONV2, + MODEL_TENSOR.A_ENC_SEANET_SCALE_CONV, + MODEL_TENSOR.A_ENC_ATTN_SCALE, + MODEL_TENSOR.A_ENC_FFN_SCALE_LS, + MODEL_TENSOR.A_ENC_SPEAKER_PROJ, + MODEL_TENSOR.A_GEN_FLOW_INPUT_PROJ, + MODEL_TENSOR.A_GEN_FLOW_COND_EMBD, + MODEL_TENSOR.A_GEN_FLOW_TIME_FREQS, + MODEL_TENSOR.A_GEN_FLOW_TIME_UP, + MODEL_TENSOR.A_GEN_FLOW_TIME_DOWN, + MODEL_TENSOR.A_GEN_FLOW_TIME_NORM, + MODEL_TENSOR.A_GEN_FLOW_BLK_NORM, + MODEL_TENSOR.A_GEN_FLOW_BLK_UP, + MODEL_TENSOR.A_GEN_FLOW_BLK_DOWN, + MODEL_TENSOR.A_GEN_FLOW_BLK_ADA, + MODEL_TENSOR.A_GEN_FLOW_FINAL_ADA, + MODEL_TENSOR.A_GEN_FLOW_FINAL_PROJ, + MODEL_TENSOR.A_GEN_OUT_EOS, + MODEL_TENSOR.A_GEN_INPUT_LINEAR, + MODEL_TENSOR.A_GEN_EMB_MEAN, + MODEL_TENSOR.A_GEN_EMB_STD, + MODEL_TENSOR.A_GEN_WAV_QUANT_OUT, + MODEL_TENSOR.A_GEN_WAV_UPSAMPLE, + MODEL_TENSOR.A_GEN_WAV_SEANET_CONV_IN, + MODEL_TENSOR.A_GEN_WAV_SEANET_CONV_OUT, + MODEL_TENSOR.A_GEN_WAV_SEANET_RES_CONV1, + MODEL_TENSOR.A_GEN_WAV_SEANET_RES_CONV2, + MODEL_TENSOR.A_GEN_WAV_SEANET_SCALE_CONV, + MODEL_TENSOR.A_ENC_CONV_NORM_MEAN, + MODEL_TENSOR.A_ENC_CONV_NORM_VAR, + MODEL_TENSOR.A_ENC_MEL_FILTERS, + MODEL_TENSOR.A_ENC_WINDOW, + MODEL_TENSOR.A_MM_INP_PROJ, + MODEL_TENSOR.A_MM_SOFT_EMB_NORM, + MODEL_TENSOR.A_MM_EMBEDDING, + MODEL_TENSOR.A_MM_HARD_EMB_NORM, + MODEL_TENSOR.A_PER_DIM_K_SCALE, + MODEL_TENSOR.A_PER_DIM_SCALE, + MODEL_TENSOR.A_CTC_OUT, + MODEL_TENSOR.A_CTC_OUT_MID, + MODEL_TENSOR.A_ENC_ATTN_REL_POS_EMB, + # qformer projector + MODEL_TENSOR.A_QF_PROJ_QUERY, + MODEL_TENSOR.A_QF_PROJ_NORM, + MODEL_TENSOR.A_QF_PROJ_LINEAR, + MODEL_TENSOR.A_QF_SELF_ATTN_Q, + MODEL_TENSOR.A_QF_SELF_ATTN_K, + MODEL_TENSOR.A_QF_SELF_ATTN_V, + MODEL_TENSOR.A_QF_SELF_ATTN_O, + MODEL_TENSOR.A_QF_SELF_ATTN_NORM, + MODEL_TENSOR.A_QF_CROSS_ATTN_Q, + MODEL_TENSOR.A_QF_CROSS_ATTN_K, + MODEL_TENSOR.A_QF_CROSS_ATTN_V, + MODEL_TENSOR.A_QF_CROSS_ATTN_O, + MODEL_TENSOR.A_QF_CROSS_ATTN_NORM, + MODEL_TENSOR.A_QF_FFN_UP, + MODEL_TENSOR.A_QF_FFN_DOWN, + MODEL_TENSOR.A_QF_FFN_NORM, + ], + MODEL_ARCH.LLAMA: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ROPE_FREQS, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.ATTN_ROT_EMBD, + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.FFN_GATE_EXP, + MODEL_TENSOR.FFN_DOWN_EXP, + MODEL_TENSOR.FFN_UP_EXP, + ], + MODEL_ARCH.LLAMA4: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ROPE_FREQS, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.ATTN_ROT_EMBD, + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.FFN_GATE_EXP, + MODEL_TENSOR.FFN_DOWN_EXP, + MODEL_TENSOR.FFN_UP_EXP, + MODEL_TENSOR.FFN_GATE_SHEXP, + MODEL_TENSOR.FFN_DOWN_SHEXP, + MODEL_TENSOR.FFN_UP_SHEXP, + ], + MODEL_ARCH.DECI: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ROPE_FREQS, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.ATTN_ROT_EMBD, + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.FFN_GATE_EXP, + MODEL_TENSOR.FFN_DOWN_EXP, + MODEL_TENSOR.FFN_UP_EXP, + ], + MODEL_ARCH.GROK: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ROPE_FREQS, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.ATTN_ROT_EMBD, + MODEL_TENSOR.ATTN_OUT_NORM, + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.FFN_GATE_EXP, + MODEL_TENSOR.FFN_DOWN_EXP, + MODEL_TENSOR.FFN_UP_EXP, + MODEL_TENSOR.FFN_POST_NORM, + MODEL_TENSOR.LAYER_OUT_NORM, + ], + MODEL_ARCH.GPTNEOX: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + ], + MODEL_ARCH.FALCON: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_NORM_2, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + ], + MODEL_ARCH.BAICHUAN: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ROPE_FREQS, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.ATTN_ROT_EMBD, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + ], + MODEL_ARCH.STARCODER: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.POS_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + ], + MODEL_ARCH.BERT: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.TOKEN_EMBD_NORM, + MODEL_TENSOR.TOKEN_TYPES, + MODEL_TENSOR.POS_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.ATTN_OUT_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.LAYER_OUT_NORM, + MODEL_TENSOR.CLS, + MODEL_TENSOR.CLS_OUT, + ], + MODEL_ARCH.MODERN_BERT: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.TOKEN_EMBD_NORM, + MODEL_TENSOR.TOKEN_TYPES, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.CLS, + MODEL_TENSOR.CLS_OUT, + MODEL_TENSOR.CLS_NORM, + ], + MODEL_ARCH.NOMIC_BERT: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.TOKEN_EMBD_NORM, + MODEL_TENSOR.TOKEN_TYPES, + MODEL_TENSOR.POS_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.ATTN_OUT_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.LAYER_OUT_NORM, + ], + MODEL_ARCH.NOMIC_BERT_MOE: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.TOKEN_EMBD_NORM, + MODEL_TENSOR.TOKEN_TYPES, + MODEL_TENSOR.POS_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.ATTN_OUT_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_DOWN_EXP, + MODEL_TENSOR.FFN_UP_EXP, + MODEL_TENSOR.LAYER_OUT_NORM, + ], + MODEL_ARCH.NEO_BERT: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.ENC_OUTPUT_NORM, + MODEL_TENSOR.CLS, + MODEL_TENSOR.CLS_OUT, + ], + MODEL_ARCH.JINA_BERT_V2: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.TOKEN_EMBD_NORM, + MODEL_TENSOR.TOKEN_TYPES, + MODEL_TENSOR.ATTN_NORM_2, + MODEL_TENSOR.ATTN_OUT_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_Q_NORM, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_K_NORM, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.LAYER_OUT_NORM, + MODEL_TENSOR.CLS, + ], + MODEL_ARCH.JINA_BERT_V3: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.TOKEN_EMBD_NORM, + MODEL_TENSOR.TOKEN_TYPES, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.ATTN_OUT_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.LAYER_OUT_NORM, + ], + MODEL_ARCH.EUROBERT: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.FFN_DOWN, + ], + MODEL_ARCH.MPT: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.FFN_ACT, + MODEL_TENSOR.ATTN_Q_NORM, + MODEL_TENSOR.ATTN_K_NORM, + MODEL_TENSOR.POS_EMBD, + ], + MODEL_ARCH.GPTJ: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + ], + MODEL_ARCH.REFACT: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + ], + MODEL_ARCH.BLOOM: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.TOKEN_EMBD_NORM, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + ], + MODEL_ARCH.STABLELM: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ROPE_FREQS, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.ATTN_Q_NORM, + MODEL_TENSOR.ATTN_K_NORM, + ], + MODEL_ARCH.QWEN: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ROPE_FREQS, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.ATTN_ROT_EMBD, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + ], + MODEL_ARCH.QWEN2: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ROPE_FREQS, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + ], + MODEL_ARCH.DREAM: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ROPE_FREQS, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + ], + MODEL_ARCH.LLADA: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ROPE_FREQS, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + ], + MODEL_ARCH.QWEN2VL: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + ], + MODEL_ARCH.QWEN2MOE: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_GATE_EXP, + MODEL_TENSOR.FFN_DOWN_EXP, + MODEL_TENSOR.FFN_UP_EXP, + MODEL_TENSOR.FFN_GATE_INP_SHEXP, + MODEL_TENSOR.FFN_GATE_SHEXP, + MODEL_TENSOR.FFN_DOWN_SHEXP, + MODEL_TENSOR.FFN_UP_SHEXP, + ], + MODEL_ARCH.QWEN3: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ROPE_FREQS, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_Q_NORM, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_K_NORM, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + ], + MODEL_ARCH.QWEN3MOE: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_Q_NORM, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_K_NORM, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_GATE_EXP, + MODEL_TENSOR.FFN_DOWN_EXP, + MODEL_TENSOR.FFN_UP_EXP, + ], + MODEL_ARCH.QWEN3NEXT: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_Q_NORM, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_K_NORM, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.ATTN_POST_NORM, + MODEL_TENSOR.ATTN_GATE, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_GATE_INP_SHEXP, + MODEL_TENSOR.FFN_UP_SHEXP, + MODEL_TENSOR.FFN_DOWN_SHEXP, + MODEL_TENSOR.FFN_GATE_SHEXP, + MODEL_TENSOR.FFN_DOWN_EXP, + MODEL_TENSOR.FFN_UP_EXP, + MODEL_TENSOR.FFN_GATE_EXP, + MODEL_TENSOR.FFN_GATE_UP_EXP, + MODEL_TENSOR.SSM_A, + MODEL_TENSOR.SSM_CONV1D, + MODEL_TENSOR.SSM_DT, + MODEL_TENSOR.SSM_NORM, + MODEL_TENSOR.SSM_IN, + MODEL_TENSOR.SSM_BETA_ALPHA, + MODEL_TENSOR.SSM_OUT, + MODEL_TENSOR.NEXTN_EH_PROJ, + MODEL_TENSOR.NEXTN_EMBED_TOKENS, + MODEL_TENSOR.NEXTN_ENORM, + MODEL_TENSOR.NEXTN_HNORM, + MODEL_TENSOR.NEXTN_SHARED_HEAD_HEAD, + MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM, + ], + MODEL_ARCH.HRM_TEXT: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.HRM_Z_L_INIT, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_GATE, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + ], + MODEL_ARCH.QWEN3VL: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ROPE_FREQS, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_Q_NORM, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_K_NORM, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + ], + MODEL_ARCH.QWEN3VLMOE: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_Q_NORM, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_K_NORM, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_GATE_EXP, + MODEL_TENSOR.FFN_DOWN_EXP, + MODEL_TENSOR.FFN_UP_EXP, + ], + MODEL_ARCH.QWEN35: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.CLS_OUT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_Q_NORM, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_K_NORM, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.ATTN_POST_NORM, + MODEL_TENSOR.ATTN_GATE, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.SSM_A, + MODEL_TENSOR.SSM_CONV1D, + MODEL_TENSOR.SSM_DT, + MODEL_TENSOR.SSM_NORM, + MODEL_TENSOR.SSM_BETA, + MODEL_TENSOR.SSM_ALPHA, + MODEL_TENSOR.SSM_OUT, + # NextN/MTP tensors - preserved but unused + MODEL_TENSOR.NEXTN_EH_PROJ, + MODEL_TENSOR.NEXTN_EMBED_TOKENS, + MODEL_TENSOR.NEXTN_ENORM, + MODEL_TENSOR.NEXTN_HNORM, + MODEL_TENSOR.NEXTN_SHARED_HEAD_HEAD, + MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM, + ], + MODEL_ARCH.CLEF: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_Q_NORM, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_K_NORM, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.ATTN_POST_NORM, + MODEL_TENSOR.ATTN_GATE, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.SSM_A, + MODEL_TENSOR.SSM_CONV1D, + MODEL_TENSOR.SSM_DT, + MODEL_TENSOR.SSM_NORM, + MODEL_TENSOR.SSM_BETA, + MODEL_TENSOR.SSM_ALPHA, + MODEL_TENSOR.SSM_OUT, + # decision head + MODEL_TENSOR.TOKEN_TYPES, + MODEL_TENSOR.DEC_ATTN_NORM, + MODEL_TENSOR.DEC_ATTN_Q, + MODEL_TENSOR.DEC_ATTN_K, + MODEL_TENSOR.DEC_ATTN_V, + MODEL_TENSOR.DEC_ATTN_OUT, + MODEL_TENSOR.DEC_CROSS_ATTN_NORM, + MODEL_TENSOR.DEC_CROSS_ATTN_NORM_KV, + MODEL_TENSOR.DEC_CROSS_ATTN_Q, + MODEL_TENSOR.DEC_CROSS_ATTN_K, + MODEL_TENSOR.DEC_CROSS_ATTN_V, + MODEL_TENSOR.DEC_CROSS_ATTN_OUT, + MODEL_TENSOR.DEC_FFN_NORM, + MODEL_TENSOR.DEC_FFN_DOWN, + MODEL_TENSOR.DEC_FFN_UP, + MODEL_TENSOR.DECISION_HIDDEN_NORM, + MODEL_TENSOR.DECISION_PROJ_MEMORY, + MODEL_TENSOR.DECISION_PROJ_QUESTION, + MODEL_TENSOR.DECISION_PROJ_OPTION_QUESTION, + MODEL_TENSOR.DECISION_PROJ_GLOBAL, + MODEL_TENSOR.DECISION_PROJ_OPTION_CONTEXT, + MODEL_TENSOR.DECISION_PROJ_OPTION_LEXICAL, + MODEL_TENSOR.DECISION_OPTION_SUMMARY_NORM, + MODEL_TENSOR.DECISION_FIELD_NORM, + MODEL_TENSOR.DECISION_OPTION_NORM, + MODEL_TENSOR.DECISION_SCALES, + MODEL_TENSOR.DECISION_SCORER, + MODEL_TENSOR.DECISION_SCORER_OUT, + ], + MODEL_ARCH.VLLM_DECISION2: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_Q_NORM, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_K_NORM, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + # decision head + MODEL_TENSOR.DECISION_OPTION_NORM, + MODEL_TENSOR.DECISION_QUERY_NORM, + MODEL_TENSOR.DECISION_KEY, + MODEL_TENSOR.DECISION_QUERY, + MODEL_TENSOR.DECISION_OPTION_MLP, + MODEL_TENSOR.DECISION_QUERY_MLP, + MODEL_TENSOR.DECISION_SCORER_OUT, + ], + MODEL_ARCH.QWEN35MOE: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_Q_NORM, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_K_NORM, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.ATTN_POST_NORM, + MODEL_TENSOR.ATTN_GATE, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_GATE_INP_SHEXP, + MODEL_TENSOR.FFN_UP_SHEXP, + MODEL_TENSOR.FFN_DOWN_SHEXP, + MODEL_TENSOR.FFN_GATE_SHEXP, + MODEL_TENSOR.FFN_DOWN_EXP, + MODEL_TENSOR.FFN_UP_EXP, + MODEL_TENSOR.FFN_GATE_EXP, + MODEL_TENSOR.FFN_GATE_UP_EXP, + MODEL_TENSOR.SSM_A, + MODEL_TENSOR.SSM_CONV1D, + MODEL_TENSOR.SSM_DT, + MODEL_TENSOR.SSM_NORM, + MODEL_TENSOR.SSM_BETA, + MODEL_TENSOR.SSM_ALPHA, + MODEL_TENSOR.SSM_OUT, + # NextN/MTP tensors - preserved but unused + MODEL_TENSOR.NEXTN_EH_PROJ, + MODEL_TENSOR.NEXTN_EMBED_TOKENS, + MODEL_TENSOR.NEXTN_ENORM, + MODEL_TENSOR.NEXTN_HNORM, + MODEL_TENSOR.NEXTN_SHARED_HEAD_HEAD, + MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM, + ], + MODEL_ARCH.QWEN4EXP: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT, + # no OUTPUT_NORM / ATTN_NORM / ATTN_POST_NORM: hyper-connections replace every layer norm + MODEL_TENSOR.HC_HEAD_NORM, + MODEL_TENSOR.HC_HEAD_DOWN, + MODEL_TENSOR.HC_HEAD_UP, + MODEL_TENSOR.HC_ATTN_NORM, + MODEL_TENSOR.HC_ATTN_DOWN, + MODEL_TENSOR.HC_ATTN_UP, + MODEL_TENSOR.HC_ATTN_INJECT, + MODEL_TENSOR.HC_FFN_NORM, + MODEL_TENSOR.HC_FFN_DOWN, + MODEL_TENSOR.HC_FFN_UP, + MODEL_TENSOR.HC_FFN_INJECT, + # full attention layers: ATTN_Q holds [q|gate] interleaved per head + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_Q_NORM, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_K_NORM, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.INDEXER_Q_PROJ, + MODEL_TENSOR.INDEXER_K_PROJ, + MODEL_TENSOR.INDEXER_Q_NORM, + MODEL_TENSOR.INDEXER_K_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_GATE, + MODEL_TENSOR.SSM_A, + MODEL_TENSOR.SSM_CONV1D, + MODEL_TENSOR.SSM_DT, + MODEL_TENSOR.SSM_NORM, + MODEL_TENSOR.SSM_BETA, + MODEL_TENSOR.SSM_ALPHA, + MODEL_TENSOR.SSM_OUT, + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_GATE_INP_SHEXP, + MODEL_TENSOR.FFN_UP_SHEXP, + MODEL_TENSOR.FFN_DOWN_SHEXP, + MODEL_TENSOR.FFN_GATE_SHEXP, + MODEL_TENSOR.FFN_DOWN_EXP, + MODEL_TENSOR.FFN_UP_EXP, + MODEL_TENSOR.FFN_GATE_EXP, + MODEL_TENSOR.FFN_GATE_UP_EXP, + MODEL_TENSOR.PER_LAYER_TOKEN_EMBD, + MODEL_TENSOR.PLE_KEY, + MODEL_TENSOR.PLE_VALUE, + MODEL_TENSOR.PLE_NORM_KEY, + MODEL_TENSOR.PLE_NORM_QUERY, + MODEL_TENSOR.PLE_NORM_CONV, + MODEL_TENSOR.PLE_CONV1D, + # MTP block: [fc_embedding | fc_hidden] as eh_proj, its own hyper-connection mixer as the head + MODEL_TENSOR.NEXTN_EH_PROJ, + MODEL_TENSOR.NEXTN_ENORM, + MODEL_TENSOR.NEXTN_HNORM, + MODEL_TENSOR.NEXTN_HC_HEAD_NORM, + MODEL_TENSOR.NEXTN_HC_HEAD_DOWN, + MODEL_TENSOR.NEXTN_HC_HEAD_UP, + ], + MODEL_ARCH.PLAMO: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ROPE_FREQS, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.ATTN_ROT_EMBD, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + ], + MODEL_ARCH.PLAMO2: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ROPE_FREQS, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.ATTN_ROT_EMBD, + MODEL_TENSOR.ATTN_Q_NORM, + MODEL_TENSOR.ATTN_K_NORM, + MODEL_TENSOR.ATTN_POST_NORM, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.FFN_POST_NORM, + MODEL_TENSOR.SSM_IN, + MODEL_TENSOR.SSM_CONV1D, + MODEL_TENSOR.SSM_X, + MODEL_TENSOR.SSM_DT, + MODEL_TENSOR.SSM_A, + MODEL_TENSOR.SSM_D, + MODEL_TENSOR.SSM_OUT, + MODEL_TENSOR.SSM_DT_NORM, + MODEL_TENSOR.SSM_B_NORM, + MODEL_TENSOR.SSM_C_NORM, + ], + MODEL_ARCH.PLAMO3: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_Q_NORM, + MODEL_TENSOR.ATTN_K_NORM, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.ATTN_POST_NORM, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.FFN_POST_NORM, + ], + MODEL_ARCH.GPT2: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.POS_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + ], + MODEL_ARCH.PHI2: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + ], + MODEL_ARCH.PHI3: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ROPE_FACTORS_LONG, + MODEL_TENSOR.ROPE_FACTORS_SHORT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + ], + MODEL_ARCH.PHIMOE: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ROPE_FACTORS_LONG, + MODEL_TENSOR.ROPE_FACTORS_SHORT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_GATE_EXP, + MODEL_TENSOR.FFN_DOWN_EXP, + MODEL_TENSOR.FFN_UP_EXP, + ], + MODEL_ARCH.CODESHELL: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.POS_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.ATTN_ROT_EMBD, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + ], + MODEL_ARCH.ORION: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ROPE_FREQS, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.ATTN_ROT_EMBD, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + ], + MODEL_ARCH.INTERNLM2: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.ATTN_ROT_EMBD, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + ], + MODEL_ARCH.MINICPM: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.ROPE_FREQS, + MODEL_TENSOR.ROPE_FACTORS_LONG, + MODEL_TENSOR.ROPE_FACTORS_SHORT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.ATTN_ROT_EMBD, + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.FFN_GATE_EXP, + MODEL_TENSOR.FFN_DOWN_EXP, + MODEL_TENSOR.FFN_UP_EXP, + ], + MODEL_ARCH.MINICPM3: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ROPE_FACTORS_LONG, + MODEL_TENSOR.ROPE_FACTORS_SHORT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_Q_A, + MODEL_TENSOR.ATTN_Q_B, + MODEL_TENSOR.ATTN_KV_A_MQA, + MODEL_TENSOR.ATTN_KV_B, + MODEL_TENSOR.ATTN_Q_A_NORM, + MODEL_TENSOR.ATTN_KV_A_NORM, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + ], + MODEL_ARCH.GEMMA: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.FFN_NORM, + ], + MODEL_ARCH.GEMMA2: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_POST_NORM, + MODEL_TENSOR.FFN_PRE_NORM, + MODEL_TENSOR.FFN_POST_NORM, + ], + MODEL_ARCH.GEMMA3: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_Q_NORM, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_K_NORM, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_POST_NORM, + MODEL_TENSOR.FFN_PRE_NORM, + MODEL_TENSOR.FFN_POST_NORM, + ], + MODEL_ARCH.GEMMA3N: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_Q_NORM, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_K_NORM, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_POST_NORM, + MODEL_TENSOR.FFN_PRE_NORM, + MODEL_TENSOR.FFN_POST_NORM, + # altup / laurel + MODEL_TENSOR.PER_LAYER_TOKEN_EMBD, + MODEL_TENSOR.PER_LAYER_MODEL_PROJ, + MODEL_TENSOR.PER_LAYER_INP_GATE, + MODEL_TENSOR.PER_LAYER_PROJ, + MODEL_TENSOR.PER_LAYER_PROJ_NORM, + MODEL_TENSOR.PER_LAYER_POST_NORM, + MODEL_TENSOR.ALTUP_PROJ, + MODEL_TENSOR.ALTUP_UNEMBD_PROJ, + MODEL_TENSOR.ALTUP_CORRECT_COEF, + MODEL_TENSOR.ALTUP_CORRECT_SCALE, + MODEL_TENSOR.ALTUP_PREDICT_COEF, + MODEL_TENSOR.ALTUP_ROUTER, + MODEL_TENSOR.ALTUP_ROUTER_NORM, + MODEL_TENSOR.LAUREL_L, + MODEL_TENSOR.LAUREL_R, + MODEL_TENSOR.LAUREL_POST_NORM, + ], + MODEL_ARCH.GEMMA4: [ + MODEL_TENSOR.ROPE_FREQS, + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_Q_NORM, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_K_NORM, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.FFN_GATE_UP_EXP, + MODEL_TENSOR.FFN_GATE_EXP, + MODEL_TENSOR.FFN_UP_EXP, + MODEL_TENSOR.FFN_DOWN_EXP, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_POST_NORM, + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_PRE_NORM, + MODEL_TENSOR.FFN_PRE_NORM_2, + MODEL_TENSOR.FFN_POST_NORM, + MODEL_TENSOR.FFN_POST_NORM_1, + MODEL_TENSOR.FFN_POST_NORM_2, + MODEL_TENSOR.LAYER_OUT_SCALE, + MODEL_TENSOR.PER_LAYER_TOKEN_EMBD, + MODEL_TENSOR.PER_LAYER_MODEL_PROJ, + MODEL_TENSOR.PER_LAYER_INP_GATE, + MODEL_TENSOR.PER_LAYER_PROJ, + MODEL_TENSOR.PER_LAYER_PROJ_NORM, + MODEL_TENSOR.PER_LAYER_POST_NORM, + ], + MODEL_ARCH.GEMMA4_ASSISTANT: [ + MODEL_TENSOR.ROPE_FREQS, + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.MASKED_EMBD_CENTROIDS, + MODEL_TENSOR.MASKED_EMBD_ORDERING, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.NEXTN_PROJ_PRE, + MODEL_TENSOR.NEXTN_PROJ_POST, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_Q_NORM, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_POST_NORM, + MODEL_TENSOR.FFN_PRE_NORM, + MODEL_TENSOR.FFN_POST_NORM, + MODEL_TENSOR.LAYER_OUT_SCALE, + ], + MODEL_ARCH.GEMMA_EMBEDDING: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.DENSE_2_OUT, + MODEL_TENSOR.DENSE_3_OUT, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_Q_NORM, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_K_NORM, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_POST_NORM, + MODEL_TENSOR.FFN_PRE_NORM, + MODEL_TENSOR.FFN_POST_NORM, + ], + MODEL_ARCH.GEMMA_EMBEDDING2: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_Q_NORM, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_K_NORM, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_POST_NORM, + MODEL_TENSOR.FFN_PRE_NORM, + MODEL_TENSOR.FFN_POST_NORM, + MODEL_TENSOR.LAYER_OUT_SCALE, + MODEL_TENSOR.PER_LAYER_MODEL_PROJ, + MODEL_TENSOR.PER_LAYER_INP_GATE, + MODEL_TENSOR.PER_LAYER_PROJ, + MODEL_TENSOR.PER_LAYER_PROJ_NORM, + MODEL_TENSOR.PER_LAYER_POST_NORM, + ], + MODEL_ARCH.STARCODER2: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ROPE_FREQS, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.ATTN_ROT_EMBD, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + ], + MODEL_ARCH.RWKV6: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.TOKEN_EMBD_NORM, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_NORM_2, + MODEL_TENSOR.TIME_MIX_W1, + MODEL_TENSOR.TIME_MIX_W2, + MODEL_TENSOR.TIME_MIX_LERP_X, + MODEL_TENSOR.TIME_MIX_LERP_K, + MODEL_TENSOR.TIME_MIX_LERP_V, + MODEL_TENSOR.TIME_MIX_LERP_R, + MODEL_TENSOR.TIME_MIX_LERP_G, + MODEL_TENSOR.TIME_MIX_LERP_W, + MODEL_TENSOR.TIME_MIX_LERP_FUSED, + MODEL_TENSOR.TIME_MIX_FIRST, + MODEL_TENSOR.TIME_MIX_DECAY, + MODEL_TENSOR.TIME_MIX_DECAY_W1, + MODEL_TENSOR.TIME_MIX_DECAY_W2, + MODEL_TENSOR.TIME_MIX_KEY, + MODEL_TENSOR.TIME_MIX_VALUE, + MODEL_TENSOR.TIME_MIX_RECEPTANCE, + MODEL_TENSOR.TIME_MIX_GATE, + MODEL_TENSOR.TIME_MIX_LN, + MODEL_TENSOR.TIME_MIX_OUTPUT, + MODEL_TENSOR.CHANNEL_MIX_LERP_K, + MODEL_TENSOR.CHANNEL_MIX_LERP_R, + MODEL_TENSOR.CHANNEL_MIX_KEY, + MODEL_TENSOR.CHANNEL_MIX_RECEPTANCE, + MODEL_TENSOR.CHANNEL_MIX_VALUE, + ], + MODEL_ARCH.RWKV6QWEN2: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.TIME_MIX_W1, + MODEL_TENSOR.TIME_MIX_W2, + MODEL_TENSOR.TIME_MIX_LERP_X, + MODEL_TENSOR.TIME_MIX_LERP_K, + MODEL_TENSOR.TIME_MIX_LERP_V, + MODEL_TENSOR.TIME_MIX_LERP_R, + MODEL_TENSOR.TIME_MIX_LERP_G, + MODEL_TENSOR.TIME_MIX_LERP_W, + MODEL_TENSOR.TIME_MIX_LERP_FUSED, + MODEL_TENSOR.TIME_MIX_FIRST, + MODEL_TENSOR.TIME_MIX_DECAY, + MODEL_TENSOR.TIME_MIX_DECAY_W1, + MODEL_TENSOR.TIME_MIX_DECAY_W2, + MODEL_TENSOR.TIME_MIX_KEY, + MODEL_TENSOR.TIME_MIX_VALUE, + MODEL_TENSOR.TIME_MIX_RECEPTANCE, + MODEL_TENSOR.TIME_MIX_GATE, + MODEL_TENSOR.TIME_MIX_LN, + MODEL_TENSOR.TIME_MIX_OUTPUT, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + ], + MODEL_ARCH.RWKV7: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.TOKEN_EMBD_NORM, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_NORM_2, + MODEL_TENSOR.TIME_MIX_LERP_FUSED, + MODEL_TENSOR.TIME_MIX_W0, + MODEL_TENSOR.TIME_MIX_W1, + MODEL_TENSOR.TIME_MIX_W2, + MODEL_TENSOR.TIME_MIX_A0, + MODEL_TENSOR.TIME_MIX_A1, + MODEL_TENSOR.TIME_MIX_A2, + MODEL_TENSOR.TIME_MIX_V0, + MODEL_TENSOR.TIME_MIX_V1, + MODEL_TENSOR.TIME_MIX_V2, + MODEL_TENSOR.TIME_MIX_G1, + MODEL_TENSOR.TIME_MIX_G2, + MODEL_TENSOR.TIME_MIX_K_K, + MODEL_TENSOR.TIME_MIX_K_A, + MODEL_TENSOR.TIME_MIX_R_K, + MODEL_TENSOR.TIME_MIX_KEY, + MODEL_TENSOR.TIME_MIX_VALUE, + MODEL_TENSOR.TIME_MIX_RECEPTANCE, + MODEL_TENSOR.TIME_MIX_LN, + MODEL_TENSOR.TIME_MIX_OUTPUT, + MODEL_TENSOR.CHANNEL_MIX_LERP_K, + MODEL_TENSOR.CHANNEL_MIX_KEY, + MODEL_TENSOR.CHANNEL_MIX_VALUE, + ], + MODEL_ARCH.ARWKV7: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.TOKEN_EMBD_NORM, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.TIME_MIX_LERP_FUSED, + MODEL_TENSOR.TIME_MIX_W0, + MODEL_TENSOR.TIME_MIX_W1, + MODEL_TENSOR.TIME_MIX_W2, + MODEL_TENSOR.TIME_MIX_A0, + MODEL_TENSOR.TIME_MIX_A1, + MODEL_TENSOR.TIME_MIX_A2, + MODEL_TENSOR.TIME_MIX_V0, + MODEL_TENSOR.TIME_MIX_V1, + MODEL_TENSOR.TIME_MIX_V2, + MODEL_TENSOR.TIME_MIX_G1, + MODEL_TENSOR.TIME_MIX_G2, + MODEL_TENSOR.TIME_MIX_K_K, + MODEL_TENSOR.TIME_MIX_K_A, + MODEL_TENSOR.TIME_MIX_R_K, + MODEL_TENSOR.TIME_MIX_KEY, + MODEL_TENSOR.TIME_MIX_VALUE, + MODEL_TENSOR.TIME_MIX_RECEPTANCE, + MODEL_TENSOR.TIME_MIX_LN, + MODEL_TENSOR.TIME_MIX_OUTPUT, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + ], + MODEL_ARCH.MAMBA: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.SSM_IN, + MODEL_TENSOR.SSM_CONV1D, + MODEL_TENSOR.SSM_X, + MODEL_TENSOR.SSM_DT, + MODEL_TENSOR.SSM_A, + MODEL_TENSOR.SSM_D, + MODEL_TENSOR.SSM_OUT, + ], + MODEL_ARCH.MAMBA2: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.SSM_IN, + MODEL_TENSOR.SSM_CONV1D, + MODEL_TENSOR.SSM_DT, + MODEL_TENSOR.SSM_A, + MODEL_TENSOR.SSM_D, + MODEL_TENSOR.SSM_NORM, + MODEL_TENSOR.SSM_OUT, + ], + MODEL_ARCH.MAPLE: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_Q_NORM, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_K_NORM, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_GATE_EXP, + MODEL_TENSOR.FFN_DOWN_EXP, + MODEL_TENSOR.FFN_UP_EXP, + ], + MODEL_ARCH.JAMBA: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.SSM_IN, + MODEL_TENSOR.SSM_CONV1D, + MODEL_TENSOR.SSM_X, + MODEL_TENSOR.SSM_DT, + MODEL_TENSOR.SSM_DT_NORM, + MODEL_TENSOR.SSM_A, + MODEL_TENSOR.SSM_B_NORM, + MODEL_TENSOR.SSM_C_NORM, + MODEL_TENSOR.SSM_D, + MODEL_TENSOR.SSM_OUT, + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.FFN_GATE_EXP, + MODEL_TENSOR.FFN_DOWN_EXP, + MODEL_TENSOR.FFN_UP_EXP, + ], + MODEL_ARCH.XVERSE: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ROPE_FREQS, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.ATTN_ROT_EMBD, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + ], + MODEL_ARCH.COMMAND_R: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.ATTN_K_NORM, + MODEL_TENSOR.ATTN_Q_NORM, + ], + MODEL_ARCH.COHERE2: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + ], + MODEL_ARCH.COHERE2MOE: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_GATE_EXP, + MODEL_TENSOR.FFN_GATE_UP_EXP, + MODEL_TENSOR.FFN_DOWN_EXP, + MODEL_TENSOR.FFN_UP_EXP, + MODEL_TENSOR.FFN_GATE_SHEXP, + MODEL_TENSOR.FFN_DOWN_SHEXP, + MODEL_TENSOR.FFN_UP_SHEXP, + MODEL_TENSOR.NEXTN_EH_PROJ, + MODEL_TENSOR.NEXTN_EMBED_TOKENS, + MODEL_TENSOR.NEXTN_ENORM, + MODEL_TENSOR.NEXTN_HNORM, + MODEL_TENSOR.NEXTN_SHARED_HEAD_HEAD, + MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM, + ], + MODEL_ARCH.DBRX: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.ATTN_OUT_NORM, + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_GATE_EXP, + MODEL_TENSOR.FFN_DOWN_EXP, + MODEL_TENSOR.FFN_UP_EXP, + ], + MODEL_ARCH.OLMO: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + ], + MODEL_ARCH.OLMO2: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.ATTN_POST_NORM, + MODEL_TENSOR.ATTN_Q_NORM, + MODEL_TENSOR.ATTN_K_NORM, + MODEL_TENSOR.FFN_POST_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + ], + MODEL_ARCH.SEED_OSS: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.ATTN_POST_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + ], + MODEL_ARCH.OLMOE: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_Q_NORM, + MODEL_TENSOR.ATTN_K_NORM, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_GATE_EXP, + MODEL_TENSOR.FFN_UP_EXP, + MODEL_TENSOR.FFN_DOWN_EXP, + ], + MODEL_ARCH.MUSE_GLIMMER: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_Q_NORM, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_K_NORM, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.ATTN_GATE, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_POST_NORM, + MODEL_TENSOR.FFN_PRE_NORM, + MODEL_TENSOR.FFN_POST_NORM, + ], + MODEL_ARCH.OPENELM: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_Q_NORM, + MODEL_TENSOR.ATTN_K_NORM, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + ], + MODEL_ARCH.ARCTIC: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ROPE_FREQS, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.ATTN_ROT_EMBD, + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.FFN_NORM_EXP, + MODEL_TENSOR.FFN_GATE_EXP, + MODEL_TENSOR.FFN_DOWN_EXP, + MODEL_TENSOR.FFN_UP_EXP, + ], + MODEL_ARCH.DEEPSEEK: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ROPE_FREQS, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.ATTN_ROT_EMBD, + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.FFN_GATE_EXP, + MODEL_TENSOR.FFN_DOWN_EXP, + MODEL_TENSOR.FFN_UP_EXP, + MODEL_TENSOR.FFN_GATE_SHEXP, + MODEL_TENSOR.FFN_DOWN_SHEXP, + MODEL_TENSOR.FFN_UP_SHEXP, + ], + MODEL_ARCH.DEEPSEEK2: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ROPE_FREQS, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_Q_A, + MODEL_TENSOR.ATTN_Q_B, + MODEL_TENSOR.ATTN_KV_A_MQA, + MODEL_TENSOR.ATTN_KV_B, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_K_B, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_V_B, + MODEL_TENSOR.ATTN_Q_A_NORM, + MODEL_TENSOR.ATTN_KV_A_NORM, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.ATTN_ROT_EMBD, + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.FFN_GATE_EXP, + MODEL_TENSOR.FFN_DOWN_EXP, + MODEL_TENSOR.FFN_UP_EXP, + MODEL_TENSOR.FFN_GATE_SHEXP, + MODEL_TENSOR.FFN_DOWN_SHEXP, + MODEL_TENSOR.FFN_UP_SHEXP, + MODEL_TENSOR.FFN_EXP_PROBS_B, + # NextN/MTP tensors + MODEL_TENSOR.NEXTN_EH_PROJ, + MODEL_TENSOR.NEXTN_EMBED_TOKENS, + MODEL_TENSOR.NEXTN_ENORM, + MODEL_TENSOR.NEXTN_HNORM, + MODEL_TENSOR.NEXTN_SHARED_HEAD_HEAD, + MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM, + ], + MODEL_ARCH.DEEPSEEK2OCR: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ROPE_FREQS, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_Q_A, + MODEL_TENSOR.ATTN_Q_B, + MODEL_TENSOR.ATTN_KV_A_MQA, + MODEL_TENSOR.ATTN_KV_B, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_K_B, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_V_B, + MODEL_TENSOR.ATTN_Q_A_NORM, + MODEL_TENSOR.ATTN_KV_A_NORM, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.ATTN_ROT_EMBD, + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.FFN_GATE_EXP, + MODEL_TENSOR.FFN_DOWN_EXP, + MODEL_TENSOR.FFN_UP_EXP, + MODEL_TENSOR.FFN_GATE_UP_EXP, + MODEL_TENSOR.FFN_GATE_SHEXP, + MODEL_TENSOR.FFN_DOWN_SHEXP, + MODEL_TENSOR.FFN_UP_SHEXP, + MODEL_TENSOR.FFN_EXP_PROBS_B, + ], + MODEL_ARCH.DEEPSEEK32: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ROPE_FREQS, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_Q_A, + MODEL_TENSOR.ATTN_Q_B, + MODEL_TENSOR.ATTN_KV_A_MQA, + MODEL_TENSOR.ATTN_K_B, + MODEL_TENSOR.ATTN_V_B, + MODEL_TENSOR.ATTN_Q_A_NORM, + MODEL_TENSOR.ATTN_KV_A_NORM, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.ATTN_ROT_EMBD, + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.FFN_GATE_EXP, + MODEL_TENSOR.FFN_DOWN_EXP, + MODEL_TENSOR.FFN_UP_EXP, + MODEL_TENSOR.FFN_GATE_SHEXP, + MODEL_TENSOR.FFN_DOWN_SHEXP, + MODEL_TENSOR.FFN_UP_SHEXP, + MODEL_TENSOR.FFN_EXP_PROBS_B, + MODEL_TENSOR.INDEXER_K_NORM, + MODEL_TENSOR.INDEXER_PROJ, + MODEL_TENSOR.INDEXER_ATTN_K, + MODEL_TENSOR.INDEXER_ATTN_Q_B, + # NextN/MTP tensors - preserved but unused + MODEL_TENSOR.NEXTN_EH_PROJ, + MODEL_TENSOR.NEXTN_EMBED_TOKENS, + MODEL_TENSOR.NEXTN_ENORM, + MODEL_TENSOR.NEXTN_HNORM, + MODEL_TENSOR.NEXTN_SHARED_HEAD_HEAD, + MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM, + ], + MODEL_ARCH.DEEPSEEK4: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.HC_HEAD_FN, + MODEL_TENSOR.HC_HEAD_BASE, + MODEL_TENSOR.HC_HEAD_SCALE, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_SINKS, + MODEL_TENSOR.ATTN_Q_A, + MODEL_TENSOR.ATTN_Q_B, + MODEL_TENSOR.ATTN_Q_A_NORM, + MODEL_TENSOR.ATTN_KV, + MODEL_TENSOR.ATTN_KV_NORM, + MODEL_TENSOR.ATTN_OUT_A, + MODEL_TENSOR.ATTN_OUT_B, + MODEL_TENSOR.HC_ATTN_FN, + MODEL_TENSOR.HC_ATTN_BASE, + MODEL_TENSOR.HC_ATTN_SCALE, + MODEL_TENSOR.HC_FFN_FN, + MODEL_TENSOR.HC_FFN_BASE, + MODEL_TENSOR.HC_FFN_SCALE, + MODEL_TENSOR.ATTN_COMPRESSOR_WKV, + MODEL_TENSOR.ATTN_COMPRESSOR_WGATE, + MODEL_TENSOR.ATTN_COMPRESSOR_APE, + MODEL_TENSOR.ATTN_COMPRESSOR_NORM, + MODEL_TENSOR.INDEXER_PROJ, + MODEL_TENSOR.INDEXER_ATTN_Q_B, + MODEL_TENSOR.INDEXER_COMPRESSOR_WKV, + MODEL_TENSOR.INDEXER_COMPRESSOR_WGATE, + MODEL_TENSOR.INDEXER_COMPRESSOR_APE, + MODEL_TENSOR.INDEXER_COMPRESSOR_NORM, + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_GATE_TID2EID, + MODEL_TENSOR.FFN_EXP_PROBS_B, + MODEL_TENSOR.FFN_EXP_PROBS_B_VL, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE_EXP, + MODEL_TENSOR.FFN_DOWN_EXP, + MODEL_TENSOR.FFN_UP_EXP, + MODEL_TENSOR.FFN_GATE_SHEXP, + MODEL_TENSOR.FFN_DOWN_SHEXP, + MODEL_TENSOR.FFN_UP_SHEXP, + MODEL_TENSOR.NEXTN_EH_PROJ, + MODEL_TENSOR.NEXTN_EMBED_TOKENS, + MODEL_TENSOR.NEXTN_ENORM, + MODEL_TENSOR.NEXTN_HNORM, + MODEL_TENSOR.NEXTN_SHARED_HEAD_HEAD, + MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM, + ], + MODEL_ARCH.ERNIE4_5_MOE: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_GATE_EXP, + MODEL_TENSOR.FFN_DOWN_EXP, + MODEL_TENSOR.FFN_UP_EXP, + MODEL_TENSOR.FFN_GATE_SHEXP, + MODEL_TENSOR.FFN_DOWN_SHEXP, + MODEL_TENSOR.FFN_UP_SHEXP, + MODEL_TENSOR.FFN_EXP_PROBS_B, + ], + MODEL_ARCH.PLM: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_KV_A_MQA, + MODEL_TENSOR.ATTN_KV_A_NORM, + MODEL_TENSOR.ATTN_KV_B, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.FFN_DOWN, + ], + MODEL_ARCH.CHATGLM : [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.ROPE_FREQS, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + ], + MODEL_ARCH.GLM4 : [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.ROPE_FREQS, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.ATTN_POST_NORM, + MODEL_TENSOR.FFN_POST_NORM, + # NextN/MTP tensors - preserved but unused + MODEL_TENSOR.NEXTN_EH_PROJ, + MODEL_TENSOR.NEXTN_EMBED_TOKENS, + MODEL_TENSOR.NEXTN_ENORM, + MODEL_TENSOR.NEXTN_HNORM, + MODEL_TENSOR.NEXTN_SHARED_HEAD_HEAD, + MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM, + ], + MODEL_ARCH.GLM4_MOE: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_POST_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.ATTN_Q_NORM, + MODEL_TENSOR.ATTN_K_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_GATE_EXP, + MODEL_TENSOR.FFN_DOWN_EXP, + MODEL_TENSOR.FFN_UP_EXP, + MODEL_TENSOR.FFN_GATE_SHEXP, + MODEL_TENSOR.FFN_DOWN_SHEXP, + MODEL_TENSOR.FFN_UP_SHEXP, + MODEL_TENSOR.FFN_EXP_PROBS_B, + # NextN/MTP tensors + MODEL_TENSOR.NEXTN_EH_PROJ, + MODEL_TENSOR.NEXTN_EMBED_TOKENS, + MODEL_TENSOR.NEXTN_ENORM, + MODEL_TENSOR.NEXTN_HNORM, + MODEL_TENSOR.NEXTN_SHARED_HEAD_HEAD, + MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM, + ], + MODEL_ARCH.GLM_DSA: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ROPE_FREQS, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_Q_A, + MODEL_TENSOR.ATTN_Q_B, + MODEL_TENSOR.ATTN_KV_A_MQA, + MODEL_TENSOR.ATTN_KV_B, + MODEL_TENSOR.ATTN_K_B, + MODEL_TENSOR.ATTN_V_B, + MODEL_TENSOR.ATTN_Q_A_NORM, + MODEL_TENSOR.ATTN_KV_A_NORM, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.ATTN_ROT_EMBD, + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.FFN_GATE_EXP, + MODEL_TENSOR.FFN_DOWN_EXP, + MODEL_TENSOR.FFN_UP_EXP, + MODEL_TENSOR.FFN_GATE_SHEXP, + MODEL_TENSOR.FFN_DOWN_SHEXP, + MODEL_TENSOR.FFN_UP_SHEXP, + MODEL_TENSOR.FFN_EXP_PROBS_B, + MODEL_TENSOR.INDEXER_K_NORM, + MODEL_TENSOR.INDEXER_PROJ, + MODEL_TENSOR.INDEXER_ATTN_K, + MODEL_TENSOR.INDEXER_ATTN_Q_B, + # NextN/MTP tensors - preserved but unused + MODEL_TENSOR.NEXTN_EH_PROJ, + MODEL_TENSOR.NEXTN_EMBED_TOKENS, + MODEL_TENSOR.NEXTN_ENORM, + MODEL_TENSOR.NEXTN_HNORM, + MODEL_TENSOR.NEXTN_SHARED_HEAD_HEAD, + MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM, + ], + MODEL_ARCH.GLM5_NEXT: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + # mHC + MODEL_TENSOR.HC_ATTN_FN, + MODEL_TENSOR.HC_ATTN_BASE, + MODEL_TENSOR.HC_ATTN_SCALE, + MODEL_TENSOR.HC_FFN_FN, + MODEL_TENSOR.HC_FFN_BASE, + MODEL_TENSOR.HC_FFN_SCALE, + # KDA (linear attention) layers + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.SSM_CONV1D_Q, + MODEL_TENSOR.SSM_CONV1D_K, + MODEL_TENSOR.SSM_CONV1D_V, + MODEL_TENSOR.SSM_F_A, + MODEL_TENSOR.SSM_F_B, + MODEL_TENSOR.SSM_BETA, + MODEL_TENSOR.SSM_A, + MODEL_TENSOR.SSM_G_A, + MODEL_TENSOR.SSM_G_B, + MODEL_TENSOR.SSM_DT, + MODEL_TENSOR.SSM_NORM, + # MLA (nope) + DSA layers + MODEL_TENSOR.ATTN_Q_A, + MODEL_TENSOR.ATTN_Q_B, + MODEL_TENSOR.ATTN_KV_A_MQA, + MODEL_TENSOR.ATTN_KV_B, + MODEL_TENSOR.ATTN_K_B, + MODEL_TENSOR.ATTN_V_B, + MODEL_TENSOR.ATTN_Q_A_NORM, + MODEL_TENSOR.ATTN_KV_A_NORM, + MODEL_TENSOR.INDEXER_K_NORM, + MODEL_TENSOR.INDEXER_PROJ, + MODEL_TENSOR.INDEXER_ATTN_K, + MODEL_TENSOR.INDEXER_ATTN_Q_B, + MODEL_TENSOR.INDEXER_KPOOL_GATE, + MODEL_TENSOR.INDEXER_KPOOL_APE, + # FFN + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_GATE_EXP, + MODEL_TENSOR.FFN_DOWN_EXP, + MODEL_TENSOR.FFN_UP_EXP, + MODEL_TENSOR.FFN_GATE_SHEXP, + MODEL_TENSOR.FFN_DOWN_SHEXP, + MODEL_TENSOR.FFN_UP_SHEXP, + MODEL_TENSOR.FFN_EXP_PROBS_B, + # NextN/MTP tensors - preserved but unused + MODEL_TENSOR.NEXTN_EH_PROJ, + MODEL_TENSOR.NEXTN_EMBED_TOKENS, + MODEL_TENSOR.NEXTN_ENORM, + MODEL_TENSOR.NEXTN_HNORM, + MODEL_TENSOR.NEXTN_SHARED_HEAD_HEAD, + MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM, + ], + MODEL_ARCH.BITNET: [ + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.ATTN_SUB_NORM, + MODEL_TENSOR.FFN_SUB_NORM, + ], + MODEL_ARCH.T5: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.DEC_ATTN_NORM, + MODEL_TENSOR.DEC_ATTN_Q, + MODEL_TENSOR.DEC_ATTN_K, + MODEL_TENSOR.DEC_ATTN_V, + MODEL_TENSOR.DEC_ATTN_OUT, + MODEL_TENSOR.DEC_ATTN_REL_B, + MODEL_TENSOR.DEC_CROSS_ATTN_NORM, + MODEL_TENSOR.DEC_CROSS_ATTN_Q, + MODEL_TENSOR.DEC_CROSS_ATTN_K, + MODEL_TENSOR.DEC_CROSS_ATTN_V, + MODEL_TENSOR.DEC_CROSS_ATTN_OUT, + MODEL_TENSOR.DEC_CROSS_ATTN_REL_B, + MODEL_TENSOR.DEC_FFN_NORM, + MODEL_TENSOR.DEC_FFN_GATE, + MODEL_TENSOR.DEC_FFN_DOWN, + MODEL_TENSOR.DEC_FFN_UP, + MODEL_TENSOR.DEC_OUTPUT_NORM, + MODEL_TENSOR.ENC_ATTN_NORM, + MODEL_TENSOR.ENC_ATTN_Q, + MODEL_TENSOR.ENC_ATTN_K, + MODEL_TENSOR.ENC_ATTN_V, + MODEL_TENSOR.ENC_ATTN_OUT, + MODEL_TENSOR.ENC_ATTN_REL_B, + MODEL_TENSOR.ENC_FFN_NORM, + MODEL_TENSOR.ENC_FFN_GATE, + MODEL_TENSOR.ENC_FFN_DOWN, + MODEL_TENSOR.ENC_FFN_UP, + MODEL_TENSOR.ENC_OUTPUT_NORM, + ], + MODEL_ARCH.T5ENCODER: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ENC_ATTN_NORM, + MODEL_TENSOR.ENC_ATTN_Q, + MODEL_TENSOR.ENC_ATTN_K, + MODEL_TENSOR.ENC_ATTN_V, + MODEL_TENSOR.ENC_ATTN_OUT, + MODEL_TENSOR.ENC_ATTN_REL_B, + MODEL_TENSOR.ENC_FFN_NORM, + MODEL_TENSOR.ENC_FFN_GATE, + MODEL_TENSOR.ENC_FFN_DOWN, + MODEL_TENSOR.ENC_FFN_UP, + MODEL_TENSOR.ENC_OUTPUT_NORM, + ], + MODEL_ARCH.JAIS: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_UP, + ], + MODEL_ARCH.JAIS2: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + ], + MODEL_ARCH.NEMOTRON: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ROPE_FREQS, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.ATTN_ROT_EMBD, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + ], + MODEL_ARCH.NEMOTRON_H: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.SSM_IN, + MODEL_TENSOR.SSM_CONV1D, + MODEL_TENSOR.SSM_DT, + MODEL_TENSOR.SSM_A, + MODEL_TENSOR.SSM_D, + MODEL_TENSOR.SSM_NORM, + MODEL_TENSOR.SSM_OUT, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + ], + MODEL_ARCH.NEMOTRON_H_MOE: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.SSM_IN, + MODEL_TENSOR.SSM_CONV1D, + MODEL_TENSOR.SSM_DT, + MODEL_TENSOR.SSM_A, + MODEL_TENSOR.SSM_D, + MODEL_TENSOR.SSM_NORM, + MODEL_TENSOR.SSM_OUT, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + # experts + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_UP_EXP, + MODEL_TENSOR.FFN_DOWN_EXP, + # expert latent + MODEL_TENSOR.MOE_LATENT_DOWN, + MODEL_TENSOR.MOE_LATENT_UP, + # shared expert + MODEL_TENSOR.FFN_DOWN_SHEXP, + MODEL_TENSOR.FFN_UP_SHEXP, + MODEL_TENSOR.FFN_EXP_PROBS_B, + # NextN/MTP (draft head) + MODEL_TENSOR.ATTN_POST_NORM, + MODEL_TENSOR.NEXTN_EH_PROJ, + MODEL_TENSOR.NEXTN_ENORM, + MODEL_TENSOR.NEXTN_HNORM, + MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM, + ], + MODEL_ARCH.EXAONE: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ROPE_FREQS, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.ATTN_ROT_EMBD, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + ], + MODEL_ARCH.EXAONE4: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ROPE_FREQS, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_Q_NORM, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_K_NORM, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.ATTN_POST_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.FFN_POST_NORM, + # NextN/MTP tensors - preserved but unused + MODEL_TENSOR.NEXTN_EH_PROJ, + MODEL_TENSOR.NEXTN_EMBED_TOKENS, + MODEL_TENSOR.NEXTN_ENORM, + MODEL_TENSOR.NEXTN_HNORM, + MODEL_TENSOR.NEXTN_SHARED_HEAD_HEAD, + MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM, + ], + MODEL_ARCH.EXAONE_MOE: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ROPE_FREQS, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_Q_NORM, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_K_NORM, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_GATE_EXP, + MODEL_TENSOR.FFN_DOWN_EXP, + MODEL_TENSOR.FFN_UP_EXP, + MODEL_TENSOR.FFN_GATE_SHEXP, + MODEL_TENSOR.FFN_DOWN_SHEXP, + MODEL_TENSOR.FFN_UP_SHEXP, + MODEL_TENSOR.FFN_EXP_PROBS_B, + # NextN/MTP tensors - preserved but unused + MODEL_TENSOR.NEXTN_EH_PROJ, + MODEL_TENSOR.NEXTN_EMBED_TOKENS, + MODEL_TENSOR.NEXTN_ENORM, + MODEL_TENSOR.NEXTN_HNORM, + MODEL_TENSOR.NEXTN_SHARED_HEAD_HEAD, + MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM, + ], + MODEL_ARCH.GRANITE: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + ], + MODEL_ARCH.GRANITE_MOE: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_GATE_EXP, + MODEL_TENSOR.FFN_DOWN_EXP, + MODEL_TENSOR.FFN_UP_EXP, + MODEL_TENSOR.FFN_GATE_SHEXP, + MODEL_TENSOR.FFN_UP_SHEXP, + MODEL_TENSOR.FFN_DOWN_SHEXP, + ], + MODEL_ARCH.GRANITE_HYBRID: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.SSM_IN, + MODEL_TENSOR.SSM_CONV1D, + MODEL_TENSOR.SSM_DT, + MODEL_TENSOR.SSM_A, + MODEL_TENSOR.SSM_D, + MODEL_TENSOR.SSM_NORM, + MODEL_TENSOR.SSM_OUT, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_NORM, + # MoE + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_GATE_EXP, + MODEL_TENSOR.FFN_DOWN_EXP, + MODEL_TENSOR.FFN_UP_EXP, + MODEL_TENSOR.FFN_GATE_SHEXP, + MODEL_TENSOR.FFN_UP_SHEXP, + MODEL_TENSOR.FFN_DOWN_SHEXP, + # Dense + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + ], + MODEL_ARCH.GRANITE_SWITCH: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + ], + MODEL_ARCH.GRANITE_SWA: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.ATTN_SINKS, + MODEL_TENSOR.ROPE_FREQS, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + # MoE (GraniteMoeSWA) + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_GATE_EXP, + MODEL_TENSOR.FFN_GATE_UP_EXP, + MODEL_TENSOR.FFN_DOWN_EXP, + MODEL_TENSOR.FFN_UP_EXP, + # Shared expert - gate+up kept fused in FFN_UP_SHEXP (LLM_FFN_SWIGLU) + MODEL_TENSOR.FFN_UP_SHEXP, + MODEL_TENSOR.FFN_DOWN_SHEXP, + ], + MODEL_ARCH.CHAMELEON: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_Q_NORM, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_K_NORM, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + ], + MODEL_ARCH.WAVTOKENIZER_DEC: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.TOKEN_EMBD_NORM, + MODEL_TENSOR.CONV1D, + MODEL_TENSOR.CONVNEXT_DW, + MODEL_TENSOR.CONVNEXT_NORM, + MODEL_TENSOR.CONVNEXT_PW1, + MODEL_TENSOR.CONVNEXT_PW2, + MODEL_TENSOR.CONVNEXT_GAMMA, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.POSNET_CONV1, + MODEL_TENSOR.POSNET_CONV2, + MODEL_TENSOR.POSNET_NORM, + MODEL_TENSOR.POSNET_NORM1, + MODEL_TENSOR.POSNET_NORM2, + MODEL_TENSOR.POSNET_ATTN_NORM, + MODEL_TENSOR.POSNET_ATTN_Q, + MODEL_TENSOR.POSNET_ATTN_K, + MODEL_TENSOR.POSNET_ATTN_V, + MODEL_TENSOR.POSNET_ATTN_OUT, + ], + MODEL_ARCH.BAILINGMOE: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ROPE_FREQS, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE_EXP, + MODEL_TENSOR.FFN_DOWN_EXP, + MODEL_TENSOR.FFN_UP_EXP, + MODEL_TENSOR.FFN_GATE_SHEXP, + MODEL_TENSOR.FFN_DOWN_SHEXP, + MODEL_TENSOR.FFN_UP_SHEXP, + ], + MODEL_ARCH.BAILINGMOE2: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_Q_NORM, + MODEL_TENSOR.ATTN_K_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_EXP_PROBS_B, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.FFN_GATE_EXP, + MODEL_TENSOR.FFN_DOWN_EXP, + MODEL_TENSOR.FFN_UP_EXP, + MODEL_TENSOR.FFN_GATE_SHEXP, + MODEL_TENSOR.FFN_DOWN_SHEXP, + MODEL_TENSOR.FFN_UP_SHEXP, + MODEL_TENSOR.NEXTN_EH_PROJ, + MODEL_TENSOR.NEXTN_EMBED_TOKENS, + MODEL_TENSOR.NEXTN_ENORM, + MODEL_TENSOR.NEXTN_HNORM, + MODEL_TENSOR.NEXTN_SHARED_HEAD_HEAD, + MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM, + MODEL_TENSOR.LAYER_OUT_NORM, + ], + MODEL_ARCH.BAILINGMOE3: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_Q_A, + MODEL_TENSOR.ATTN_Q_B, + MODEL_TENSOR.ATTN_Q_A_NORM, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.ATTN_GATE, + MODEL_TENSOR.ATTN_KV_A_MQA, + MODEL_TENSOR.ATTN_KV_B, + MODEL_TENSOR.ATTN_K_B, + MODEL_TENSOR.ATTN_V_B, + MODEL_TENSOR.ATTN_KV_A_NORM, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_GATE_EXP, + MODEL_TENSOR.FFN_DOWN_EXP, + MODEL_TENSOR.FFN_UP_EXP, + MODEL_TENSOR.FFN_GATE_SHEXP, + MODEL_TENSOR.FFN_DOWN_SHEXP, + MODEL_TENSOR.FFN_UP_SHEXP, + MODEL_TENSOR.FFN_EXP_PROBS_B, + MODEL_TENSOR.SSM_CONV1D_Q, + MODEL_TENSOR.SSM_CONV1D_K, + MODEL_TENSOR.SSM_CONV1D_V, + MODEL_TENSOR.SSM_F_A, + MODEL_TENSOR.SSM_BETA, + MODEL_TENSOR.SSM_A, + MODEL_TENSOR.SSM_G_A, + MODEL_TENSOR.SSM_DT, + MODEL_TENSOR.SSM_NORM, + MODEL_TENSOR.NEXTN_EH_PROJ, + MODEL_TENSOR.NEXTN_ENORM, + MODEL_TENSOR.NEXTN_HNORM, + MODEL_TENSOR.LAYER_OUT_NORM, + ], + MODEL_ARCH.DOTS1: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_Q_NORM, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_K_NORM, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_EXP_PROBS_B, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_GATE_EXP, + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_GATE_SHEXP, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_DOWN_EXP, + MODEL_TENSOR.FFN_DOWN_SHEXP, + MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.FFN_UP_EXP, + MODEL_TENSOR.FFN_UP_SHEXP, + ], + MODEL_ARCH.DOTS3NOTE: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_Q_A, + MODEL_TENSOR.ATTN_Q_B, + MODEL_TENSOR.ATTN_KV_A_MQA, + MODEL_TENSOR.ATTN_K_B, + MODEL_TENSOR.ATTN_V_B, + MODEL_TENSOR.ATTN_Q_A_NORM, + MODEL_TENSOR.ATTN_KV_A_NORM, + MODEL_TENSOR.ATTN_K_NORM, + MODEL_TENSOR.ATTN_GATE, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_EXP_PROBS_B, + MODEL_TENSOR.FFN_GATE_EXP, + MODEL_TENSOR.FFN_DOWN_EXP, + MODEL_TENSOR.FFN_UP_EXP, + MODEL_TENSOR.FFN_GATE_SHEXP, + MODEL_TENSOR.FFN_DOWN_SHEXP, + MODEL_TENSOR.FFN_UP_SHEXP, + MODEL_TENSOR.INDEXER_K_NORM, + MODEL_TENSOR.INDEXER_PROJ, + MODEL_TENSOR.INDEXER_ATTN_K, + MODEL_TENSOR.INDEXER_ATTN_Q_B, + # NextN/MTP tensors - preserved but unused + MODEL_TENSOR.NEXTN_EH_PROJ, + MODEL_TENSOR.NEXTN_EMBED_TOKENS, + MODEL_TENSOR.NEXTN_ENORM, + MODEL_TENSOR.NEXTN_HNORM, + MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM, + ], + MODEL_ARCH.ARCEE: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ROPE_FREQS, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.ATTN_ROT_EMBD, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + ], + MODEL_ARCH.AFMOE: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_POST_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.ATTN_Q_NORM, + MODEL_TENSOR.ATTN_K_NORM, + MODEL_TENSOR.ATTN_GATE, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_GATE_EXP, + MODEL_TENSOR.FFN_DOWN_EXP, + MODEL_TENSOR.FFN_UP_EXP, + MODEL_TENSOR.FFN_GATE_SHEXP, + MODEL_TENSOR.FFN_UP_SHEXP, + MODEL_TENSOR.FFN_DOWN_SHEXP, + MODEL_TENSOR.FFN_PRE_NORM, + MODEL_TENSOR.FFN_POST_NORM, + MODEL_TENSOR.FFN_EXP_PROBS_B, + ], + MODEL_ARCH.LAGUNA: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.ATTN_Q_NORM, + MODEL_TENSOR.ATTN_K_NORM, + MODEL_TENSOR.ATTN_GATE, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_EXP_PROBS_B, + MODEL_TENSOR.FFN_GATE_EXP, + MODEL_TENSOR.FFN_DOWN_EXP, + MODEL_TENSOR.FFN_UP_EXP, + MODEL_TENSOR.FFN_GATE_SHEXP, + MODEL_TENSOR.FFN_UP_SHEXP, + MODEL_TENSOR.FFN_DOWN_SHEXP, + ], + MODEL_ARCH.ERNIE4_5: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + ], + MODEL_ARCH.PADDLEOCR: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + ], + MODEL_ARCH.FALCON_H1: [ + # Token embedding + MODEL_TENSOR.TOKEN_EMBD, + + # Input layernorm + MODEL_TENSOR.ATTN_NORM, + + # Attention components + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_Q, # Query projection + MODEL_TENSOR.ATTN_K, # Key projection + MODEL_TENSOR.ATTN_V, # Value projection + MODEL_TENSOR.ATTN_OUT, # Output projection + + # SSM components (Mamba2 specific) + MODEL_TENSOR.SSM_IN, # Input projection for SSM + MODEL_TENSOR.SSM_CONV1D, # Convolution layer + MODEL_TENSOR.SSM_DT, # Delta time projection + MODEL_TENSOR.SSM_A, # A parameter (log form) + MODEL_TENSOR.SSM_D, # D parameter + MODEL_TENSOR.SSM_NORM, # Normalization in SSM + MODEL_TENSOR.SSM_OUT, # Output projection + + # Pre-feedforward layernorm + MODEL_TENSOR.FFN_PRE_NORM, + + # Feed-forward network components + MODEL_TENSOR.FFN_GATE, # Gate projection (SwiGLU) + MODEL_TENSOR.FFN_DOWN, # Down projection + MODEL_TENSOR.FFN_UP, # Up projection + + # Post-feedforward layernorm + MODEL_TENSOR.OUTPUT_NORM, # Final layer norm + MODEL_TENSOR.OUTPUT, # Output projection (lm_head) + ], + MODEL_ARCH.HUNYUAN_MOE: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ROPE_FREQS, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_Q_NORM, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_K_NORM, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE_EXP, + MODEL_TENSOR.FFN_DOWN_EXP, + MODEL_TENSOR.FFN_UP_EXP, + MODEL_TENSOR.FFN_GATE_SHEXP, + MODEL_TENSOR.FFN_DOWN_SHEXP, + MODEL_TENSOR.FFN_UP_SHEXP, + ], + MODEL_ARCH.HUNYUAN_DENSE: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_Q_NORM, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_K_NORM, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + ], + MODEL_ARCH.HUNYUAN_VL: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_Q_NORM, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_K_NORM, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + ], + MODEL_ARCH.HY_V3: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_Q_NORM, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_K_NORM, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_EXP_PROBS_B, + MODEL_TENSOR.FFN_GATE_EXP, + MODEL_TENSOR.FFN_DOWN_EXP, + MODEL_TENSOR.FFN_UP_EXP, + MODEL_TENSOR.FFN_GATE_SHEXP, + MODEL_TENSOR.FFN_DOWN_SHEXP, + MODEL_TENSOR.FFN_UP_SHEXP, + # NextN/MTP tensors (draft head) + MODEL_TENSOR.NEXTN_EH_PROJ, + MODEL_TENSOR.NEXTN_EMBED_TOKENS, + MODEL_TENSOR.NEXTN_ENORM, + MODEL_TENSOR.NEXTN_HNORM, + MODEL_TENSOR.NEXTN_SHARED_HEAD_HEAD, + MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM, + ], + MODEL_ARCH.HY_V4: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ROPE_FREQS, + MODEL_TENSOR.HC_HEAD_FN, + MODEL_TENSOR.HC_HEAD_BASE, + MODEL_TENSOR.HC_HEAD_SCALE, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_SINKS, + MODEL_TENSOR.ATTN_Q_A, + MODEL_TENSOR.ATTN_Q_A_NORM, + MODEL_TENSOR.ATTN_Q_B, + MODEL_TENSOR.ATTN_KV_A_MQA, + MODEL_TENSOR.ATTN_KV_A_NORM, + MODEL_TENSOR.ATTN_K_B, + MODEL_TENSOR.ATTN_V_B, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.ATTN_GATE, + MODEL_TENSOR.INDEXER_K_NORM, + MODEL_TENSOR.INDEXER_PROJ, + MODEL_TENSOR.INDEXER_ATTN_K, + MODEL_TENSOR.INDEXER_ATTN_Q_B, + MODEL_TENSOR.HC_ATTN_FN, + MODEL_TENSOR.HC_ATTN_BASE, + MODEL_TENSOR.HC_ATTN_SCALE, + MODEL_TENSOR.HC_FFN_FN, + MODEL_TENSOR.HC_FFN_BASE, + MODEL_TENSOR.HC_FFN_SCALE, + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_EXP_PROBS_B, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.FFN_GATE_EXP, + MODEL_TENSOR.FFN_DOWN_EXP, + MODEL_TENSOR.FFN_UP_EXP, + MODEL_TENSOR.FFN_GATE_SHEXP, + MODEL_TENSOR.FFN_DOWN_SHEXP, + MODEL_TENSOR.FFN_UP_SHEXP, + ], + MODEL_ARCH.SMOLLM3: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ROPE_FREQS, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.ATTN_ROT_EMBD, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + ], + MODEL_ARCH.GPT_OSS: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_POST_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.ATTN_SINKS, + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_GATE_EXP, + MODEL_TENSOR.FFN_DOWN_EXP, + MODEL_TENSOR.FFN_UP_EXP, + ], + MODEL_ARCH.LFM2: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.TOKEN_EMBD_NORM, + MODEL_TENSOR.SHORTCONV_CONV, + MODEL_TENSOR.SHORTCONV_INPROJ, + MODEL_TENSOR.SHORTCONV_OUTPROJ, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.ATTN_NORM, # operator_norm + MODEL_TENSOR.ATTN_Q_NORM, + MODEL_TENSOR.ATTN_K_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.DENSE_2_OUT, # LFM2-ColBert-350M + MODEL_TENSOR.TOKEN_TYPES, # decision head + MODEL_TENSOR.CLS, + MODEL_TENSOR.CLS_NORM, + MODEL_TENSOR.CLS_OUT, + ], + MODEL_ARCH.LFM2MOE: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.TOKEN_EMBD_NORM, + MODEL_TENSOR.SHORTCONV_CONV, + MODEL_TENSOR.SHORTCONV_INPROJ, + MODEL_TENSOR.SHORTCONV_OUTPROJ, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.ATTN_NORM, # operator_norm + MODEL_TENSOR.ATTN_Q_NORM, + MODEL_TENSOR.ATTN_K_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_GATE_EXP, + MODEL_TENSOR.FFN_DOWN_EXP, + MODEL_TENSOR.FFN_UP_EXP, + MODEL_TENSOR.FFN_EXP_PROBS_B, + ], + MODEL_ARCH.SMALLTHINKER: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_GATE_EXP, + MODEL_TENSOR.FFN_DOWN_EXP, + MODEL_TENSOR.FFN_UP_EXP, + ], + MODEL_ARCH.APERTUS: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ROPE_FREQS, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.ATTN_ROT_EMBD, + MODEL_TENSOR.ATTN_Q_NORM, + MODEL_TENSOR.ATTN_K_NORM, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + ], + MODEL_ARCH.LLADA_MOE: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_Q_NORM, + MODEL_TENSOR.ATTN_K_NORM, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_GATE_EXP, + MODEL_TENSOR.FFN_UP_EXP, + MODEL_TENSOR.FFN_DOWN_EXP, + ], + MODEL_ARCH.GROVEMOE: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_Q_NORM, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_K_NORM, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_GATE_EXP, + MODEL_TENSOR.FFN_DOWN_EXP, + MODEL_TENSOR.FFN_UP_EXP, + MODEL_TENSOR.FFN_GATE_CHEXP, + MODEL_TENSOR.FFN_DOWN_CHEXP, + MODEL_TENSOR.FFN_UP_CHEXP, + ], + MODEL_ARCH.MINIMAX01: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_NORM_2, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.ATTN_GATE, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_GATE_EXP, + MODEL_TENSOR.FFN_DOWN_EXP, + MODEL_TENSOR.FFN_UP_EXP, + ], + MODEL_ARCH.MINIMAXM2: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_Q_NORM, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_K_NORM, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_GATE_EXP, + MODEL_TENSOR.FFN_DOWN_EXP, + MODEL_TENSOR.FFN_UP_EXP, + MODEL_TENSOR.FFN_EXP_PROBS_B, + ], + MODEL_ARCH.MINIMAXM3: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_Q_NORM, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_K_NORM, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_EXP_PROBS_B, + MODEL_TENSOR.FFN_GATE_EXP, + MODEL_TENSOR.FFN_DOWN_EXP, + MODEL_TENSOR.FFN_UP_EXP, + MODEL_TENSOR.FFN_GATE_SHEXP, + MODEL_TENSOR.FFN_DOWN_SHEXP, + MODEL_TENSOR.FFN_UP_SHEXP, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.INDEXER_Q_PROJ, + MODEL_TENSOR.INDEXER_K_PROJ, + MODEL_TENSOR.INDEXER_Q_NORM, + MODEL_TENSOR.INDEXER_K_NORM, + ], + MODEL_ARCH.COGVLM: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.VISEXP_ATTN_QKV, + MODEL_TENSOR.VISEXP_ATTN_OUT, + MODEL_TENSOR.VISEXP_GATE, + MODEL_TENSOR.VISEXP_UP, + MODEL_TENSOR.VISEXP_DOWN, + ], + MODEL_ARCH.RND1: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_Q_NORM, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_K_NORM, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_GATE_EXP, + MODEL_TENSOR.FFN_DOWN_EXP, + MODEL_TENSOR.FFN_UP_EXP, + ], + MODEL_ARCH.PANGU_EMBED: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + ], + MODEL_ARCH.MISTRAL3: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ROPE_FREQS, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.ATTN_ROT_EMBD, + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.FFN_GATE_EXP, + MODEL_TENSOR.FFN_DOWN_EXP, + MODEL_TENSOR.FFN_UP_EXP, + ], + MODEL_ARCH.EAGLE3: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ROPE_FREQS, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_NORM_2, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.FC, + MODEL_TENSOR.ENC_OUTPUT_NORM, + MODEL_TENSOR.D2T, + ], + MODEL_ARCH.DFLASH: [ + MODEL_TENSOR.ATTN_POST_NORM, + MODEL_TENSOR.FFN_POST_NORM, + MODEL_TENSOR.LAYER_OUT_SCALE, + MODEL_TENSOR.ROPE_FREQS, + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.ATTN_Q_NORM, + MODEL_TENSOR.ATTN_K_NORM, + MODEL_TENSOR.ATTN_SINKS, + MODEL_TENSOR.ATTN_Q_A, + MODEL_TENSOR.ATTN_Q_B, + MODEL_TENSOR.ATTN_Q_A_NORM, + MODEL_TENSOR.ATTN_KV, + MODEL_TENSOR.ATTN_KV_NORM, + MODEL_TENSOR.ATTN_OUT_A, + MODEL_TENSOR.ATTN_OUT_B, + MODEL_TENSOR.HC_ATTN_FN, + MODEL_TENSOR.HC_ATTN_BASE, + MODEL_TENSOR.HC_ATTN_SCALE, + MODEL_TENSOR.HC_FFN_FN, + MODEL_TENSOR.HC_FFN_BASE, + MODEL_TENSOR.HC_FFN_SCALE, + MODEL_TENSOR.HC_HEAD_FN, + MODEL_TENSOR.HC_HEAD_BASE, + MODEL_TENSOR.HC_HEAD_SCALE, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_EXP_PROBS_B, + MODEL_TENSOR.FFN_GATE_EXP, + MODEL_TENSOR.FFN_DOWN_EXP, + MODEL_TENSOR.FFN_UP_EXP, + MODEL_TENSOR.FFN_GATE_SHEXP, + MODEL_TENSOR.FFN_DOWN_SHEXP, + MODEL_TENSOR.FFN_UP_SHEXP, + MODEL_TENSOR.FC, + MODEL_TENSOR.ENC_OUTPUT_NORM, + MODEL_TENSOR.D2T, + # optional DSpark heads + MODEL_TENSOR.DSPARK_MARKOV_W1, + MODEL_TENSOR.DSPARK_MARKOV_W2, + MODEL_TENSOR.DSPARK_CONF_PROJ, + MODEL_TENSOR.DFLASH_ATTN_CONV_BASE, + MODEL_TENSOR.DFLASH_ATTN_CONV_PROJ, + MODEL_TENSOR.DFLASH_FFN_CONV_BASE, + MODEL_TENSOR.DFLASH_FFN_CONV_PROJ, + MODEL_TENSOR.DFLASH_SELECTOR_PREV, + MODEL_TENSOR.DFLASH_SELECTOR_NEXT, + MODEL_TENSOR.DFLASH_SELECTOR_HIDDEN, + ], + MODEL_ARCH.MISTRAL4: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ROPE_FREQS, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_Q_A, + MODEL_TENSOR.ATTN_Q_B, + MODEL_TENSOR.ATTN_KV_A_MQA, + MODEL_TENSOR.ATTN_KV_B, + MODEL_TENSOR.ATTN_K_B, + MODEL_TENSOR.ATTN_V_B, + MODEL_TENSOR.ATTN_Q_A_NORM, + MODEL_TENSOR.ATTN_KV_A_NORM, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.ATTN_ROT_EMBD, + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.FFN_GATE_EXP, + MODEL_TENSOR.FFN_DOWN_EXP, + MODEL_TENSOR.FFN_UP_EXP, + MODEL_TENSOR.FFN_GATE_UP_EXP, + MODEL_TENSOR.FFN_GATE_SHEXP, + MODEL_TENSOR.FFN_DOWN_SHEXP, + MODEL_TENSOR.FFN_UP_SHEXP, + MODEL_TENSOR.FFN_EXP_PROBS_B, + ], + MODEL_ARCH.MIMO2: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_SINKS, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_GATE_EXP, + MODEL_TENSOR.FFN_DOWN_EXP, + MODEL_TENSOR.FFN_UP_EXP, + MODEL_TENSOR.FFN_EXP_PROBS_B, + MODEL_TENSOR.LAYER_OUT_NORM, + MODEL_TENSOR.NEXTN_EH_PROJ, + MODEL_TENSOR.NEXTN_EMBED_TOKENS, + MODEL_TENSOR.NEXTN_ENORM, + MODEL_TENSOR.NEXTN_HNORM, + MODEL_TENSOR.NEXTN_SHARED_HEAD_HEAD, + MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM, + ], + MODEL_ARCH.STEP35: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ROPE_FREQS, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_Q_NORM, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_K_NORM, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_GATE, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_GATE_EXP, + MODEL_TENSOR.FFN_DOWN_EXP, + MODEL_TENSOR.FFN_UP_EXP, + MODEL_TENSOR.FFN_UP_SHEXP, + MODEL_TENSOR.FFN_GATE_SHEXP, + MODEL_TENSOR.FFN_DOWN_SHEXP, + MODEL_TENSOR.FFN_EXP_PROBS_B, + # NextN/MTP tensors (Step3p5 draft head) + MODEL_TENSOR.NEXTN_EH_PROJ, + MODEL_TENSOR.NEXTN_EMBED_TOKENS, + MODEL_TENSOR.NEXTN_ENORM, + MODEL_TENSOR.NEXTN_HNORM, + MODEL_TENSOR.NEXTN_SHARED_HEAD_HEAD, + MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM, + ], + MODEL_ARCH.SPARK2_5: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_GATE, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + ], + MODEL_ARCH.LLAMA_EMBED: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ROPE_FREQS, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.ATTN_ROT_EMBD, + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.FFN_GATE_EXP, + MODEL_TENSOR.FFN_DOWN_EXP, + MODEL_TENSOR.FFN_UP_EXP, + ], + MODEL_ARCH.MAINCODER: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_Q_NORM, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_K_NORM, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + ], + MODEL_ARCH.KIMI_LINEAR: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.ATTN_Q_A, + MODEL_TENSOR.ATTN_Q_B, + MODEL_TENSOR.ATTN_KV_A_MQA, + MODEL_TENSOR.ATTN_KV_B, + MODEL_TENSOR.ATTN_K_B, + MODEL_TENSOR.ATTN_V_B, + MODEL_TENSOR.ATTN_Q_A_NORM, + MODEL_TENSOR.ATTN_KV_A_NORM, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_GATE_EXP, + MODEL_TENSOR.FFN_DOWN_EXP, + MODEL_TENSOR.FFN_UP_EXP, + MODEL_TENSOR.SSM_CONV1D_Q, + MODEL_TENSOR.SSM_CONV1D_K, + MODEL_TENSOR.SSM_CONV1D_V, + MODEL_TENSOR.SSM_F_A, + MODEL_TENSOR.SSM_F_B, + MODEL_TENSOR.SSM_BETA, + MODEL_TENSOR.SSM_A, + MODEL_TENSOR.SSM_G_A, + MODEL_TENSOR.SSM_G_B, + MODEL_TENSOR.SSM_DT, + MODEL_TENSOR.SSM_NORM, + MODEL_TENSOR.FFN_EXP_PROBS_B, + MODEL_TENSOR.FFN_GATE_SHEXP, + MODEL_TENSOR.FFN_DOWN_SHEXP, + MODEL_TENSOR.FFN_UP_SHEXP, + ], + MODEL_ARCH.KIMI_K3: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.OUTPUT_RES_SCORE, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_RES_SCORE, + MODEL_TENSOR.FFN_RES_SCORE, + # MLA (full-attention layers) + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.ATTN_GATE, + MODEL_TENSOR.ATTN_Q_A, + MODEL_TENSOR.ATTN_Q_B, + MODEL_TENSOR.ATTN_KV_A_MQA, + MODEL_TENSOR.ATTN_KV_B, + MODEL_TENSOR.ATTN_K_B, + MODEL_TENSOR.ATTN_V_B, + MODEL_TENSOR.ATTN_Q_A_NORM, + MODEL_TENSOR.ATTN_KV_A_NORM, + # KDA (linear-attention layers) + MODEL_TENSOR.SSM_CONV1D_Q, + MODEL_TENSOR.SSM_CONV1D_K, + MODEL_TENSOR.SSM_CONV1D_V, + MODEL_TENSOR.SSM_F_A, + MODEL_TENSOR.SSM_F_B, + MODEL_TENSOR.SSM_BETA, + MODEL_TENSOR.SSM_A, + MODEL_TENSOR.SSM_G, + MODEL_TENSOR.SSM_DT, + MODEL_TENSOR.SSM_NORM, + # FFN + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_EXP_PROBS_B, + MODEL_TENSOR.FFN_GATE_EXP, + MODEL_TENSOR.FFN_DOWN_EXP, + MODEL_TENSOR.FFN_UP_EXP, + MODEL_TENSOR.FFN_GATE_SHEXP, + MODEL_TENSOR.FFN_DOWN_SHEXP, + MODEL_TENSOR.FFN_UP_SHEXP, + MODEL_TENSOR.FFN_ROUTED_DOWN, + MODEL_TENSOR.FFN_ROUTED_UP, + MODEL_TENSOR.FFN_ROUTED_NORM, + ], + MODEL_ARCH.TALKIE: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_Q_NORM, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.LAYER_OUT_SCALE, + ], + MODEL_ARCH.MELLUM: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_Q_NORM, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_K_NORM, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_GATE_EXP, + MODEL_TENSOR.FFN_DOWN_EXP, + MODEL_TENSOR.FFN_UP_EXP, + ], + MODEL_ARCH.NANBEIGE: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ROPE_FREQS, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.ATTN_ROT_EMBD, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + ], + MODEL_ARCH.QWEN3TTS: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_Q_NORM, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_K_NORM, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + ], + MODEL_ARCH.POCKETTTS: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + ], + MODEL_ARCH.K2HORIZON: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_Q_NORM, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_K_NORM, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_V_GATE, + MODEL_TENSOR.ATTN_V_EXP, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.ATTN_GATE, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_EXP_PROBS_B, + MODEL_TENSOR.FFN_GATE_EXP, + MODEL_TENSOR.FFN_UP_EXP, + MODEL_TENSOR.FFN_DOWN_EXP, + MODEL_TENSOR.FFN_GATE_SHEXP, + MODEL_TENSOR.FFN_UP_SHEXP, + MODEL_TENSOR.FFN_DOWN_SHEXP, + ], +} + +# tensors that will not be serialized +MODEL_TENSOR_SKIP: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { + MODEL_ARCH.LLAMA: [ + MODEL_TENSOR.ROPE_FREQS, + MODEL_TENSOR.ATTN_ROT_EMBD, + ], + MODEL_ARCH.DECI: [ + MODEL_TENSOR.ROPE_FREQS, + MODEL_TENSOR.ATTN_ROT_EMBD, + ], + MODEL_ARCH.BAICHUAN: [ + MODEL_TENSOR.ROPE_FREQS, + MODEL_TENSOR.ATTN_ROT_EMBD, + ], + MODEL_ARCH.QWEN: [ + MODEL_TENSOR.ROPE_FREQS, + MODEL_TENSOR.ATTN_ROT_EMBD, + ], + MODEL_ARCH.CODESHELL: [ + MODEL_TENSOR.ROPE_FREQS, + MODEL_TENSOR.ATTN_ROT_EMBD, + ], + MODEL_ARCH.ORION: [ + MODEL_TENSOR.ROPE_FREQS, + MODEL_TENSOR.ATTN_ROT_EMBD, + ], + MODEL_ARCH.STARCODER2: [ + MODEL_TENSOR.ROPE_FREQS, + MODEL_TENSOR.ATTN_ROT_EMBD, + ], + MODEL_ARCH.XVERSE: [ + MODEL_TENSOR.ROPE_FREQS, + MODEL_TENSOR.ATTN_ROT_EMBD, + ], + MODEL_ARCH.DEEPSEEK: [ + MODEL_TENSOR.ROPE_FREQS, + MODEL_TENSOR.ATTN_ROT_EMBD, + ], + MODEL_ARCH.DEEPSEEK2: [ + MODEL_TENSOR.ROPE_FREQS, + MODEL_TENSOR.ATTN_ROT_EMBD, + ], + MODEL_ARCH.DEEPSEEK2OCR: [ + MODEL_TENSOR.ROPE_FREQS, + MODEL_TENSOR.ATTN_ROT_EMBD, + ], + MODEL_ARCH.DEEPSEEK32: [ + MODEL_TENSOR.ROPE_FREQS, + MODEL_TENSOR.ATTN_ROT_EMBD, + ], + MODEL_ARCH.HY_V4: [ + MODEL_TENSOR.ROPE_FREQS, + MODEL_TENSOR.ATTN_ROT_EMBD, + ], + MODEL_ARCH.CHATGLM: [ + MODEL_TENSOR.ROPE_FREQS, + ], + MODEL_ARCH.NEMOTRON: [ + MODEL_TENSOR.ROPE_FREQS, + MODEL_TENSOR.ATTN_ROT_EMBD, + ], + MODEL_ARCH.BAILINGMOE: [ + MODEL_TENSOR.ROPE_FREQS, + ], + MODEL_ARCH.PANGU_EMBED: [ + MODEL_TENSOR.ROPE_FREQS, + MODEL_TENSOR.ATTN_ROT_EMBD, + ], + MODEL_ARCH.NANBEIGE: [ + MODEL_TENSOR.ROPE_FREQS, + MODEL_TENSOR.ATTN_ROT_EMBD, + ], +} + +# +# types +# + + +class TokenType(IntEnum): + NORMAL = 1 + UNKNOWN = 2 + CONTROL = 3 + USER_DEFINED = 4 + UNUSED = 5 + BYTE = 6 + + +class RopeScalingType(Enum): + NONE = 'none' + LINEAR = 'linear' + YARN = 'yarn' + LONGROPE = 'longrope' + + +# M-RoPE: input position slot (t, y, x, z) that feeds each RoPE section, in section order +class RopeSectionOrder(Enum): + TYXZ = 'tyxz' # default + ZYXT = 'zyxt' + + +class PoolingType(IntEnum): + NONE = 0 + MEAN = 1 + CLS = 2 + LAST = 3 + RANK = 4 + + +class GGMLQuantizationType(IntEnum): + F32 = 0 + F16 = 1 + Q4_0 = 2 + Q4_1 = 3 + Q5_0 = 6 + Q5_1 = 7 + Q8_0 = 8 + Q8_1 = 9 + Q2_K = 10 + Q3_K = 11 + Q4_K = 12 + Q5_K = 13 + Q6_K = 14 + Q8_K = 15 + IQ2_XXS = 16 + IQ2_XS = 17 + IQ3_XXS = 18 + IQ1_S = 19 + IQ4_NL = 20 + IQ3_S = 21 + IQ2_S = 22 + IQ4_XS = 23 + I8 = 24 + I16 = 25 + I32 = 26 + I64 = 27 + F64 = 28 + IQ1_M = 29 + BF16 = 30 + TQ1_0 = 34 + TQ2_0 = 35 + MXFP4 = 39 + NVFP4 = 40 + Q1_0 = 41 + Q2_0 = 42 + + +class ExpertGatingFuncType(IntEnum): + SOFTMAX = 1 + SIGMOID = 2 + SQRTSOFTPLUS = 4 + + +# TODO: add GGMLFileType from ggml_ftype in ggml.h + + +# from llama_ftype in llama.h +# ALL VALUES SHOULD BE THE SAME HERE AS THEY ARE OVER THERE. +class LlamaFileType(IntEnum): + ALL_F32 = 0 + MOSTLY_F16 = 1 # except 1d tensors + MOSTLY_Q4_0 = 2 # except 1d tensors + MOSTLY_Q4_1 = 3 # except 1d tensors + # MOSTLY_Q4_1_SOME_F16 = 4 # tok_embeddings.weight and output.weight are F16 + # MOSTLY_Q4_2 = 5 # support has been removed + # MOSTLY_Q4_3 = 6 # support has been removed + MOSTLY_Q8_0 = 7 # except 1d tensors + MOSTLY_Q5_0 = 8 # except 1d tensors + MOSTLY_Q5_1 = 9 # except 1d tensors + MOSTLY_Q2_K = 10 # except 1d tensors + MOSTLY_Q3_K_S = 11 # except 1d tensors + MOSTLY_Q3_K_M = 12 # except 1d tensors + MOSTLY_Q3_K_L = 13 # except 1d tensors + MOSTLY_Q4_K_S = 14 # except 1d tensors + MOSTLY_Q4_K_M = 15 # except 1d tensors + MOSTLY_Q5_K_S = 16 # except 1d tensors + MOSTLY_Q5_K_M = 17 # except 1d tensors + MOSTLY_Q6_K = 18 # except 1d tensors + MOSTLY_IQ2_XXS = 19 # except 1d tensors + MOSTLY_IQ2_XS = 20 # except 1d tensors + MOSTLY_Q2_K_S = 21 # except 1d tensors + MOSTLY_IQ3_XS = 22 # except 1d tensors + MOSTLY_IQ3_XXS = 23 # except 1d tensors + MOSTLY_IQ1_S = 24 # except 1d tensors + MOSTLY_IQ4_NL = 25 # except 1d tensors + MOSTLY_IQ3_S = 26 # except 1d tensors + MOSTLY_IQ3_M = 27 # except 1d tensors + MOSTLY_IQ2_S = 28 # except 1d tensors + MOSTLY_IQ2_M = 29 # except 1d tensors + MOSTLY_IQ4_XS = 30 # except 1d tensors + MOSTLY_IQ1_M = 31 # except 1d tensors + MOSTLY_BF16 = 32 # except 1d tensors + # MOSTLY_Q4_0_4_4 = 33 # removed from gguf files, use Q4_0 and runtime repack + # MOSTLY_Q4_0_4_8 = 34 # removed from gguf files, use Q4_0 and runtime repack + # MOSTLY_Q4_0_8_8 = 35 # removed from gguf files, use Q4_0 and runtime repack + MOSTLY_TQ1_0 = 36 # except 1d tensors + MOSTLY_TQ2_0 = 37 # except 1d tensors + MOSTLY_MXFP4_MOE = 38 # except 1d tensors + MOSTLY_NVFP4 = 39 # except 1d tensors + MOSTLY_Q1_0 = 40 # except 1d tensors + MOSTLY_Q2_0 = 41 # except 1d tensors + + GUESSED = 1024 # not specified in the model file + + +class GGUFEndian(IntEnum): + LITTLE = 0 + BIG = 1 + + +class GGUFValueType(IntEnum): + UINT8 = 0 + INT8 = 1 + UINT16 = 2 + INT16 = 3 + UINT32 = 4 + INT32 = 5 + FLOAT32 = 6 + BOOL = 7 + STRING = 8 + ARRAY = 9 + UINT64 = 10 + INT64 = 11 + FLOAT64 = 12 + + @staticmethod + def get_type(val: Any) -> GGUFValueType: + if isinstance(val, (str, bytes, bytearray)): + return GGUFValueType.STRING + elif isinstance(val, list): + return GGUFValueType.ARRAY + elif isinstance(val, float): + return GGUFValueType.FLOAT32 + elif isinstance(val, bool): + return GGUFValueType.BOOL + elif isinstance(val, int): + return GGUFValueType.INT32 + # TODO: need help with 64-bit types in Python + else: + raise ValueError(f"Unknown type: {type(val)}") + + +class DecisionType: + LAYA = "laya" # head blocks + scorer on the hidden state of one marker token per option + OPENJEV = "openjev" # logits of one label token per option + LEV = "lev" # same as openjev, noul is read from a rating scale + KEV = "kev" # dot product of the hidden states of the last token and of one end token per option + NIMBLE = "nimble" # same as openjev, the prompt lists all the questions of the request + CLEF = "clef" # joint head over all questions, one score per option + PPLX_DECIDER = "pplx-decider" # same as openjev, label codes of 1 or 2 letters + VLLM_DECISION2 = "vllm-decision2" # head over the end token of each option and the last token, one score per option + VLLM_D3 = "vllm-d3" # same as pplx-decider, non-causal attention + LFM2_D1 = "lfm2-d1" # same as openjev, the labels depend on the question type + LFM2_D1_OMNI = "lfm2-d1-omni" # same head as laya on a bidirectional LFM2 trunk, other prompt layout + + +class VisionProjectorType: + GEMMA3 = "gemma3" + GEMMA3NV = "gemma3nv" + GEMMA3NA = "gemma3na" + GEMMA4V = "gemma4v" + GEMMA4A = "gemma4a" + GEMMA4UV = "gemma4uv" # "unified" variant + GEMMA4UA = "gemma4ua" # "unified" variant + PHI4 = "phi4" + IDEFICS3 = "idefics3" + PIXTRAL = "pixtral" + LLAMA4 = "llama4" + QWEN2VL = "qwen2vl_merger" + QWEN25VL = "qwen2.5vl_merger" + EXAONE4_5 = "exaone4_5" + QWEN3VL = "qwen3vl_merger" + LING3VL = "ling3vl" + STEP3VL = "step3vl" + ULTRAVOX = "ultravox" + INTERNVL = "internvl" + QWEN2A = "qwen2a" # audio + QWEN3A = "qwen3a" # audio + GLMA = "glma" # audio + QWEN25O = "qwen2.5o" # omni + VOXTRAL = "voxtral" + MERALION = "meralion" # audio: Whisper + gated MLP adaptor + LFM2 = "lfm2" + KIMIVL = "kimivl" + PADDLEOCR = "paddleocr" + KIMIK25 = "kimik25" + LIGHTONOCR = "lightonocr" + COGVLM = "cogvlm" + JANUS_PRO = "janus_pro" + DOTSOCR = "dots_ocr" + DOTS3NOTE_V = "dots3note_v" + DOTS3NOTE_A = "dots3note_a" # audio + DEEPSEEKOCR = "deepseekocr" + DEEPSEEKOCR2 = "deepseekocr2" + DEEPSEEK4V = "deepseek4v" + LFM2A = "lfm2a" # audio + MUSIC_FLAMINGO = "musicflamingo" # audio + GLM4V = "glm4v" + GLM5V = "glm5v" + YOUTUVL = "youtuvl" + NEMOTRON_V2_VL = "nemotron_v2_vl" + QWEN3TTS_SPKENC = "qwen3tts_spkenc" # audio: ECAPA-TDNN speaker encoder + QWEN3TTS_GEN = "qwen3tts_gen" # audio generation: code_predictor + POCKETTTS_SPKENC = "pockettts_spkenc" # audio: mimi encoder as voice-prompt encoder + POCKETTTS_GEN = "pockettts_gen" # audio generation: flow-matching decoder + mimi decoder + HUNYUANVL = "hunyuanvl" + PARAKEET = "parakeet" # audio + MINIMAXM3 = "minimax_m3" + MINICPMV4_6 = "minicpmv4_6" + MINICPMV4_7 = "minicpmv4_7" + GRANITE_SPEECH = "granite_speech" # audio + MIMOVL = "mimovl" + MIMO_AUDIO = "mimo_audio" + GRANITE4_VISION = "granite4_vision" + MUSE_GLIMMER = "muse-glimmer" + COHERE2V = "cohere2v" + D1OMNI_V = "d1omni_v" # lfm2 vision, without separator tokens + D1OMNI_A = "d1omni_a" # lfm2a audio, with a residual block after the projector + + +# Items here are (block size, type size) +QK_K = 256 +GGML_QUANT_SIZES: dict[GGMLQuantizationType, tuple[int, int]] = { + GGMLQuantizationType.F32: (1, 4), + GGMLQuantizationType.F16: (1, 2), + GGMLQuantizationType.Q4_0: (32, 2 + 16), + GGMLQuantizationType.Q4_1: (32, 2 + 2 + 16), + GGMLQuantizationType.Q5_0: (32, 2 + 4 + 16), + GGMLQuantizationType.Q5_1: (32, 2 + 2 + 4 + 16), + GGMLQuantizationType.Q8_0: (32, 2 + 32), + GGMLQuantizationType.Q8_1: (32, 2 + 2 + 32), + GGMLQuantizationType.Q2_K: (256, 2 + 2 + QK_K // 16 + QK_K // 4), + GGMLQuantizationType.Q3_K: (256, 2 + QK_K // 4 + QK_K // 8 + 12), + GGMLQuantizationType.Q4_K: (256, 2 + 2 + QK_K // 2 + 12), + GGMLQuantizationType.Q5_K: (256, 2 + 2 + QK_K // 2 + QK_K // 8 + 12), + GGMLQuantizationType.Q6_K: (256, 2 + QK_K // 2 + QK_K // 4 + QK_K // 16), + GGMLQuantizationType.Q8_K: (256, 4 + QK_K + QK_K // 8), + GGMLQuantizationType.IQ2_XXS: (256, 2 + QK_K // 4), + GGMLQuantizationType.IQ2_XS: (256, 2 + QK_K // 4 + QK_K // 32), + GGMLQuantizationType.IQ3_XXS: (256, 2 + QK_K // 4 + QK_K // 8), + GGMLQuantizationType.IQ1_S: (256, 2 + QK_K // 8 + QK_K // 16), + GGMLQuantizationType.IQ4_NL: (32, 2 + 16), + GGMLQuantizationType.IQ3_S: (256, 2 + QK_K // 4 + QK_K // 8 + QK_K // 32 + 4), + GGMLQuantizationType.IQ2_S: (256, 2 + QK_K // 4 + QK_K // 16), + GGMLQuantizationType.IQ4_XS: (256, 2 + 2 + QK_K // 2 + QK_K // 64), + GGMLQuantizationType.I8: (1, 1), + GGMLQuantizationType.I16: (1, 2), + GGMLQuantizationType.I32: (1, 4), + GGMLQuantizationType.I64: (1, 8), + GGMLQuantizationType.F64: (1, 8), + GGMLQuantizationType.IQ1_M: (256, QK_K // 8 + QK_K // 16 + QK_K // 32), + GGMLQuantizationType.BF16: (1, 2), + GGMLQuantizationType.TQ1_0: (256, 2 + 4 * 13), + GGMLQuantizationType.TQ2_0: (256, 2 + 64), + GGMLQuantizationType.MXFP4: (32, 1 + 16), + GGMLQuantizationType.NVFP4: (64, 4 + 32), + GGMLQuantizationType.Q1_0: (128, 2 + 16), + GGMLQuantizationType.Q2_0: (64, 2 + 16), +} + + +# Aliases for backward compatibility. + +# general +KEY_GENERAL_ARCHITECTURE = Keys.General.ARCHITECTURE +KEY_GENERAL_QUANTIZATION_VERSION = Keys.General.QUANTIZATION_VERSION +KEY_GENERAL_ALIGNMENT = Keys.General.ALIGNMENT +KEY_GENERAL_NAME = Keys.General.NAME +KEY_GENERAL_AUTHOR = Keys.General.AUTHOR +KEY_GENERAL_URL = Keys.General.URL +KEY_GENERAL_DESCRIPTION = Keys.General.DESCRIPTION +KEY_GENERAL_LICENSE = Keys.General.LICENSE +KEY_GENERAL_SOURCE_URL = Keys.General.SOURCE_URL +KEY_GENERAL_FILE_TYPE = Keys.General.FILE_TYPE + +# LLM +KEY_VOCAB_SIZE = Keys.LLM.VOCAB_SIZE +KEY_CONTEXT_LENGTH = Keys.LLM.CONTEXT_LENGTH +KEY_EMBEDDING_LENGTH = Keys.LLM.EMBEDDING_LENGTH +KEY_BLOCK_COUNT = Keys.LLM.BLOCK_COUNT +KEY_FEED_FORWARD_LENGTH = Keys.LLM.FEED_FORWARD_LENGTH +KEY_USE_PARALLEL_RESIDUAL = Keys.LLM.USE_PARALLEL_RESIDUAL +KEY_TENSOR_DATA_LAYOUT = Keys.LLM.TENSOR_DATA_LAYOUT + +# attention +KEY_ATTENTION_HEAD_COUNT = Keys.Attention.HEAD_COUNT +KEY_ATTENTION_HEAD_COUNT_KV = Keys.Attention.HEAD_COUNT_KV +KEY_ATTENTION_MAX_ALIBI_BIAS = Keys.Attention.MAX_ALIBI_BIAS +KEY_ATTENTION_CLAMP_KQV = Keys.Attention.CLAMP_KQV +KEY_ATTENTION_LAYERNORM_EPS = Keys.Attention.LAYERNORM_EPS +KEY_ATTENTION_LAYERNORM_RMS_EPS = Keys.Attention.LAYERNORM_RMS_EPS + +# RoPE +KEY_ROPE_DIMENSION_COUNT = Keys.Rope.DIMENSION_COUNT +KEY_ROPE_FREQ_BASE = Keys.Rope.FREQ_BASE +KEY_ROPE_SCALING_TYPE = Keys.Rope.SCALING_TYPE +KEY_ROPE_SCALING_FACTOR = Keys.Rope.SCALING_FACTOR +KEY_ROPE_SCALING_ORIG_CTX_LEN = Keys.Rope.SCALING_ORIG_CTX_LEN +KEY_ROPE_SCALING_FINETUNED = Keys.Rope.SCALING_FINETUNED + +# SSM +KEY_SSM_CONV_KERNEL = Keys.SSM.CONV_KERNEL +KEY_SSM_INNER_SIZE = Keys.SSM.INNER_SIZE +KEY_SSM_STATE_SIZE = Keys.SSM.STATE_SIZE +KEY_SSM_TIME_STEP_RANK = Keys.SSM.TIME_STEP_RANK +KEY_SSM_GROUP_COUNT = Keys.SSM.GROUP_COUNT +KEY_SSM_DT_B_C_RMS = Keys.SSM.DT_B_C_RMS + +# KDA +KEY_KDA_HEAD_DIM = Keys.KDA.HEAD_DIM +KEY_KDA_SAFE_GATE = Keys.KDA.SAFE_GATE +KEY_KDA_GATE_LOWER_BOUND = Keys.KDA.GATE_LOWER_BOUND + +# tokenization +KEY_TOKENIZER_MODEL = Keys.Tokenizer.MODEL +KEY_TOKENIZER_PRE = Keys.Tokenizer.PRE +KEY_TOKENIZER_LIST = Keys.Tokenizer.LIST +KEY_TOKENIZER_TOKEN_TYPE = Keys.Tokenizer.TOKEN_TYPE +KEY_TOKENIZER_SCORES = Keys.Tokenizer.SCORES +KEY_TOKENIZER_MERGES = Keys.Tokenizer.MERGES +KEY_TOKENIZER_BOS_ID = Keys.Tokenizer.BOS_ID +KEY_TOKENIZER_EOS_ID = Keys.Tokenizer.EOS_ID +KEY_TOKENIZER_EOT_ID = Keys.Tokenizer.EOT_ID +KEY_TOKENIZER_EOM_ID = Keys.Tokenizer.EOM_ID +KEY_TOKENIZER_UNK_ID = Keys.Tokenizer.UNK_ID +KEY_TOKENIZER_SEP_ID = Keys.Tokenizer.SEP_ID +KEY_TOKENIZER_PAD_ID = Keys.Tokenizer.PAD_ID +KEY_TOKENIZER_MASK_ID = Keys.Tokenizer.MASK_ID +KEY_TOKENIZER_HF_JSON = Keys.Tokenizer.HF_JSON +KEY_TOKENIZER_RWKV = Keys.Tokenizer.RWKV + +KEY_TOKENIZER_FIM_PRE_ID = Keys.Tokenizer.FIM_PRE_ID +KEY_TOKENIZER_FIM_SUF_ID = Keys.Tokenizer.FIM_SUF_ID +KEY_TOKENIZER_FIM_MID_ID = Keys.Tokenizer.FIM_MID_ID +KEY_TOKENIZER_FIM_PAD_ID = Keys.Tokenizer.FIM_PAD_ID +KEY_TOKENIZER_FIM_REP_ID = Keys.Tokenizer.FIM_REP_ID +KEY_TOKENIZER_FIM_SEP_ID = Keys.Tokenizer.FIM_SEP_ID + +# deprecated +KEY_TOKENIZER_PREFIX_ID = Keys.Tokenizer.PREFIX_ID +KEY_TOKENIZER_SUFFIX_ID = Keys.Tokenizer.SUFFIX_ID +KEY_TOKENIZER_MIDDLE_ID = Keys.Tokenizer.MIDDLE_ID diff --git a/python/freetoken/third_party/gguf/freetoken.patch b/python/freetoken/third_party/gguf/freetoken.patch new file mode 100644 index 0000000000..b8016ff554 --- /dev/null +++ b/python/freetoken/third_party/gguf/freetoken.patch @@ -0,0 +1,23 @@ +--- a/gguf-py/gguf/__init__.py ++++ b/gguf-py/gguf/__init__.py +@@ -1,9 +1,4 @@ + from .constants import * + from .lazy import * + from .gguf_reader import * +-from .gguf_writer import * + from .quants import * +-from .tensor_mapping import * +-from .vocab import * +-from .utility import * +-from .metadata import * +--- a/gguf-py/gguf/gguf_reader.py ++++ b/gguf-py/gguf/gguf_reader.py +@@ -21,7 +21,7 @@ + # Allow running file in package as a script. + sys.path.insert(0, str(Path(__file__).parent.parent)) + +-from gguf.constants import ( ++from .constants import ( + GGML_MAX_DIMS, + GGML_QUANT_SIZES, + GGUF_DEFAULT_ALIGNMENT, diff --git a/python/freetoken/third_party/gguf/gguf_reader.py b/python/freetoken/third_party/gguf/gguf_reader.py new file mode 100644 index 0000000000..41eb6521ea --- /dev/null +++ b/python/freetoken/third_party/gguf/gguf_reader.py @@ -0,0 +1,390 @@ +# +# GGUF file reading/modification support. For API usage information, +# please see the files scripts/ for some fairly simple examples. +# +from __future__ import annotations + +import logging +import os +import sys +from collections import OrderedDict +from typing import Any, Literal, NamedTuple, TypeVar, Union + +import numpy as np +import numpy.typing as npt + +from .quants import quant_shape_to_byte_shape + +if __name__ == "__main__": + from pathlib import Path + + # Allow running file in package as a script. + sys.path.insert(0, str(Path(__file__).parent.parent)) + +from .constants import ( + GGML_MAX_DIMS, + GGML_QUANT_SIZES, + GGUF_DEFAULT_ALIGNMENT, + GGUF_MAGIC, + GGUF_VERSION, + GGMLQuantizationType, + GGUFValueType, + GGUFEndian, +) + +# limits mirroring ggml/src/gguf.cpp (not part of gguf.h) +GGUF_MAX_STRING_LENGTH = 1024 * 1024 * 1024 +GGUF_MAX_ARRAY_ELEMENTS = 1024 * 1024 * 1024 + +logger = logging.getLogger(__name__) + +READER_SUPPORTED_VERSIONS = [2, GGUF_VERSION] + + +class ReaderField(NamedTuple): + # Offset to start of this field. + offset: int + + # Name of the field (not necessarily from file data). + name: str + + # Data parts. Some types have multiple components, such as strings + # that consist of a length followed by the string data. + parts: list[npt.NDArray[Any]] = [] + + # Indexes into parts that we can call the actual data. For example + # an array of strings will be populated with indexes to the actual + # string data. + data: list[int] = [-1] + + types: list[GGUFValueType] = [] + + def contents(self, index_or_slice: int | slice = slice(None)) -> Any: + if self.types: + to_string = lambda x: str(x.tobytes(), encoding='utf-8') # noqa: E731 + main_type = self.types[0] + + if main_type == GGUFValueType.ARRAY: + sub_type = self.types[-1] + + if sub_type == GGUFValueType.STRING: + indices = self.data[index_or_slice] + + if isinstance(index_or_slice, int): + return to_string(self.parts[indices]) # type: ignore + else: + return [to_string(self.parts[idx]) for idx in indices] # type: ignore + else: + # FIXME: When/if _get_field_parts() support multi-dimensional arrays, this must do so too + + # Check if it's unsafe to perform slice optimization on data + # if any(True for idx in self.data if len(self.parts[idx]) != 1): + # optim_slice = slice(None) + # else: + # optim_slice = index_or_slice + # index_or_slice = slice(None) + + # if isinstance(optim_slice, int): + # return self.parts[self.data[optim_slice]].tolist()[0] + # else: + # return [pv for idx in self.data[optim_slice] for pv in self.parts[idx].tolist()][index_or_slice] + + if isinstance(index_or_slice, int): + return self.parts[self.data[index_or_slice]].tolist()[0] + else: + return [pv for idx in self.data[index_or_slice] for pv in self.parts[idx].tolist()] + + if main_type == GGUFValueType.STRING: + return to_string(self.parts[-1]) + else: + return self.parts[-1].tolist()[0] + + return None + + +class ReaderTensor(NamedTuple): + name: str + tensor_type: GGMLQuantizationType + shape: npt.NDArray[np.uint32] + n_elements: int + n_bytes: int + data_offset: int + data: npt.NDArray[Any] + field: ReaderField + + +class GGUFReader: + # I - same as host, S - swapped + byte_order: Literal['I', 'S'] = 'I' + alignment: int = GGUF_DEFAULT_ALIGNMENT + data_offset: int + + # Note: Internal helper, API may change. + gguf_scalar_to_np: dict[GGUFValueType, type[np.generic]] = { + GGUFValueType.UINT8: np.uint8, + GGUFValueType.INT8: np.int8, + GGUFValueType.UINT16: np.uint16, + GGUFValueType.INT16: np.int16, + GGUFValueType.UINT32: np.uint32, + GGUFValueType.INT32: np.int32, + GGUFValueType.FLOAT32: np.float32, + GGUFValueType.UINT64: np.uint64, + GGUFValueType.INT64: np.int64, + GGUFValueType.FLOAT64: np.float64, + GGUFValueType.BOOL: np.bool_, + } + + def __init__(self, path: os.PathLike[str] | str, mode: Literal['r', 'r+', 'c'] = 'r'): + self.data = np.memmap(path, mode = mode) + offs = 0 + + # Check for GGUF magic + if self._get(offs, np.uint32, override_order = '<')[0] != GGUF_MAGIC: + raise ValueError('GGUF magic invalid') + offs += 4 + + # Check GGUF version + temp_version = self._get(offs, np.uint32) + if temp_version[0] & 65535 == 0: + # If we get 0 here that means it's (probably) a GGUF file created for + # the opposite byte order of the machine this script is running on. + self.byte_order = 'S' + temp_version = temp_version.view(temp_version.dtype.newbyteorder(self.byte_order)) + version = temp_version[0] + if version not in READER_SUPPORTED_VERSIONS: + raise ValueError(f'Sorry, file appears to be version {version} which we cannot handle') + if sys.byteorder == "little": + # Host is little endian + host_endian = GGUFEndian.LITTLE + swapped_endian = GGUFEndian.BIG + else: + # Sorry PDP or other weird systems that don't use BE or LE. + host_endian = GGUFEndian.BIG + swapped_endian = GGUFEndian.LITTLE + self.endianess = swapped_endian if self.byte_order == "S" else host_endian + self.fields: OrderedDict[str, ReaderField] = OrderedDict() + self.tensors: list[ReaderTensor] = [] + offs += self._push_field(ReaderField(offs, 'GGUF.version', [temp_version], [0], [GGUFValueType.UINT32])) + + # Check tensor count and kv count + temp_counts = self._get(offs, np.uint64, 2) + offs += self._push_field(ReaderField(offs, 'GGUF.tensor_count', [temp_counts[:1]], [0], [GGUFValueType.UINT64])) + offs += self._push_field(ReaderField(offs, 'GGUF.kv_count', [temp_counts[1:]], [0], [GGUFValueType.UINT64])) + tensor_count, kv_count = temp_counts + if tensor_count > GGUF_MAX_ARRAY_ELEMENTS: + raise ValueError(f'Tensor count {tensor_count} exceeds maximum {GGUF_MAX_ARRAY_ELEMENTS}') + if kv_count > GGUF_MAX_ARRAY_ELEMENTS: + raise ValueError(f'KV count {kv_count} exceeds maximum {GGUF_MAX_ARRAY_ELEMENTS}') + offs = self._build_fields(offs, kv_count) + + # Build Tensor Info Fields + offs, tensors_fields = self._build_tensor_info(offs, tensor_count) + new_align = self.fields.get('general.alignment') + if new_align is not None: + if new_align.types != [GGUFValueType.UINT32]: + raise ValueError('Bad type for general.alignment field') + self.alignment = new_align.parts[-1][0] + # Ensure alignment is a non-zero power of two + if self.alignment == 0 or (self.alignment & (self.alignment - 1)) != 0: + raise ValueError('Invalid alignment: must be a non-zero power of two') + padding = offs % self.alignment + if padding != 0: + offs += self.alignment - padding + self.data_offset = offs + self._build_tensors(offs, tensors_fields) + + _DT = TypeVar('_DT', bound = npt.DTypeLike) + + # Fetch a key/value metadata field by key. + def get_field(self, key: str) -> Union[ReaderField, None]: + return self.fields.get(key, None) + + # Fetch a tensor from the list by index. + def get_tensor(self, idx: int) -> ReaderTensor: + return self.tensors[idx] + + def _get( + self, offset: int, dtype: npt.DTypeLike, count: int = 1, override_order: None | Literal['I', 'S', '<'] = None, + ) -> npt.NDArray[Any]: + count = int(count) + itemsize = int(np.empty([], dtype = dtype).itemsize) + end_offs = offset + itemsize * count + arr = self.data[offset:end_offs].view(dtype=dtype)[:count] + return arr.view(arr.dtype.newbyteorder(self.byte_order if override_order is None else override_order)) + + def _push_field(self, field: ReaderField, skip_sum: bool = False) -> int: + if field.name in self.fields: + # TODO: add option to make this a warning and accept duplicate keys like below + raise KeyError(f'Duplicate {field.name} already in list at offset {field.offset}') + + # logger.warning(f'Duplicate key {field.name} at offset {field.offset}') + # self.fields[field.name + '_{}'.format(field.offset)] = field + else: + self.fields[field.name] = field + return 0 if skip_sum else sum(int(part.nbytes) for part in field.parts) + + def _get_str(self, offset: int) -> tuple[npt.NDArray[np.uint64], npt.NDArray[np.uint8]]: + slen = self._get(offset, np.uint64) + if int(slen[0]) > GGUF_MAX_STRING_LENGTH: + raise ValueError(f'String length {int(slen[0])} exceeds maximum {GGUF_MAX_STRING_LENGTH}') + if offset + 8 + int(slen[0]) > self.data.nbytes: + raise ValueError(f'String length {int(slen[0])} exceeds remaining file size {self.data.nbytes - offset - 8}') + return slen, self._get(offset + 8, np.uint8, slen[0]) + + def _get_field_parts( + self, orig_offs: int, raw_type: int, + ) -> tuple[int, list[npt.NDArray[Any]], list[int], list[GGUFValueType]]: + offs = orig_offs + types: list[GGUFValueType] = [] + gtype = GGUFValueType(raw_type) + types.append(gtype) + # Handle strings. + if gtype == GGUFValueType.STRING: + sparts: list[npt.NDArray[Any]] = list(self._get_str(offs)) + size = sum(int(part.nbytes) for part in sparts) + return size, sparts, [1], types + # Check if it's a simple scalar type. + nptype = self.gguf_scalar_to_np.get(gtype) + if nptype is not None: + val = self._get(offs, nptype) + return int(val.nbytes), [val], [0], types + # Handle arrays. + if gtype == GGUFValueType.ARRAY: + raw_itype = self._get(offs, np.uint32) + offs += int(raw_itype.nbytes) + alen = self._get(offs, np.uint64) + if int(alen[0]) > GGUF_MAX_ARRAY_ELEMENTS: + raise ValueError(f'Array length {int(alen[0])} exceeds maximum {GGUF_MAX_ARRAY_ELEMENTS}') + offs += int(alen.nbytes) + aparts: list[npt.NDArray[Any]] = [raw_itype, alen] + data_idxs: list[int] = [] + # FIXME: Handle multi-dimensional arrays properly instead of flattening + for idx in range(alen[0]): + curr_size, curr_parts, curr_idxs, curr_types = self._get_field_parts(offs, raw_itype[0]) + if idx == 0: + types += curr_types + idxs_offs = len(aparts) + aparts += curr_parts + data_idxs += (idx + idxs_offs for idx in curr_idxs) + offs += curr_size + return offs - orig_offs, aparts, data_idxs, types + # We can't deal with this one. + raise ValueError(f'Unknown/unhandled field type {gtype}') + + def _get_tensor_info_field(self, orig_offs: int) -> ReaderField: + offs = orig_offs + + # Get Tensor Name + name_len, name_data = self._get_str(offs) + offs += int(name_len.nbytes + name_data.nbytes) + + # Get Tensor Dimensions Count + n_dims = self._get(offs, np.uint32) + offs += int(n_dims.nbytes) + if n_dims[0] > GGML_MAX_DIMS: + raise ValueError(f'Tensor dimensions count {n_dims[0]} exceeds GGML_MAX_DIMS ({GGML_MAX_DIMS})') + + # Get Tensor Dimension Array + dims = self._get(offs, np.uint64, n_dims[0]) + offs += int(dims.nbytes) + + # Get Tensor Encoding Scheme Type + raw_dtype = self._get(offs, np.uint32) + offs += int(raw_dtype.nbytes) + + # Get Tensor Offset + offset_tensor = self._get(offs, np.uint64) + offs += int(offset_tensor.nbytes) + + return ReaderField( + orig_offs, + str(bytes(name_data), encoding = 'utf-8'), + [name_len, name_data, n_dims, dims, raw_dtype, offset_tensor], + [1, 3, 4, 5], + ) + + def _build_fields(self, offs: int, count: int) -> int: + for _ in range(count): + orig_offs = offs + kv_klen, kv_kdata = self._get_str(offs) + offs += int(kv_klen.nbytes + kv_kdata.nbytes) + raw_kv_type = self._get(offs, np.uint32) + offs += int(raw_kv_type.nbytes) + parts: list[npt.NDArray[Any]] = [kv_klen, kv_kdata, raw_kv_type] + idxs_offs = len(parts) + field_size, field_parts, field_idxs, field_types = self._get_field_parts(offs, raw_kv_type[0]) + parts += field_parts + self._push_field(ReaderField( + orig_offs, + str(bytes(kv_kdata), encoding = 'utf-8'), + parts, + [idx + idxs_offs for idx in field_idxs], + field_types, + ), skip_sum = True) + offs += field_size + return offs + + def _build_tensor_info(self, offs: int, count: int) -> tuple[int, list[ReaderField]]: + tensor_fields = [] + for _ in range(count): + field = self._get_tensor_info_field(offs) + offs += sum(int(part.nbytes) for part in field.parts) + tensor_fields.append(field) + return offs, tensor_fields + + def _build_tensors(self, start_offs: int, fields: list[ReaderField]) -> None: + tensors = [] + tensor_names = set() # keep track of name to prevent duplicated tensors + for field in fields: + _name_len, name_data, _n_dims, dims, raw_dtype, offset_tensor = field.parts + # check if there's any tensor having same name already in the list + tensor_name = str(bytes(name_data), encoding = 'utf-8') + if tensor_name in tensor_names: + raise ValueError(f'Found duplicated tensor with name {tensor_name}') + tensor_names.add(tensor_name) + ggml_type = GGMLQuantizationType(raw_dtype[0]) + # use Python ints: np.prod on uint64 wraps silently on overflow + n_elems = 1 + for dim in dims.tolist(): + n_elems *= int(dim) + np_dims = tuple(reversed(dims.tolist())) + block_size, type_size = GGML_QUANT_SIZES[ggml_type] + n_bytes = n_elems * type_size // block_size + data_offs = int(start_offs + offset_tensor[0]) + item_type: npt.DTypeLike + if ggml_type == GGMLQuantizationType.F16: + item_count = n_elems + item_type = np.float16 + elif ggml_type == GGMLQuantizationType.F32: + item_count = n_elems + item_type = np.float32 + elif ggml_type == GGMLQuantizationType.F64: + item_count = n_elems + item_type = np.float64 + elif ggml_type == GGMLQuantizationType.I8: + item_count = n_elems + item_type = np.int8 + elif ggml_type == GGMLQuantizationType.I16: + item_count = n_elems + item_type = np.int16 + elif ggml_type == GGMLQuantizationType.I32: + item_count = n_elems + item_type = np.int32 + elif ggml_type == GGMLQuantizationType.I64: + item_count = n_elems + item_type = np.int64 + else: + item_count = n_bytes + item_type = np.uint8 + np_dims = quant_shape_to_byte_shape(np_dims, ggml_type) + tensors.append(ReaderTensor( + name = tensor_name, + tensor_type = ggml_type, + shape = dims, + n_elements = n_elems, + n_bytes = n_bytes, + data_offset = data_offs, + data = self._get(data_offs, item_type, item_count).reshape(np_dims), + field = field, + )) + self.tensors = tensors diff --git a/python/freetoken/third_party/gguf/lazy.py b/python/freetoken/third_party/gguf/lazy.py new file mode 100644 index 0000000000..a39f223215 --- /dev/null +++ b/python/freetoken/third_party/gguf/lazy.py @@ -0,0 +1,293 @@ +from __future__ import annotations +from abc import ABC, ABCMeta, abstractmethod + +import logging +from typing import Any, Callable + +import numpy as np +from numpy.typing import DTypeLike + + +logger = logging.getLogger(__name__) + + +class LazyMeta(ABCMeta): + + def __new__(cls, name: str, bases: tuple[type, ...], namespace: dict[str, Any], **kwargs): + def __getattr__(self, name: str) -> Any: + meta_attr = getattr(self._meta, name) + if callable(meta_attr): + return type(self)._wrap_fn( + (lambda s, *args, **kwargs: getattr(s, name)(*args, **kwargs)), + use_self=self, + ) + elif isinstance(meta_attr, self._tensor_type): + # e.g. self.T with torch.Tensor should still be wrapped + return type(self)._wrap_fn(lambda s: getattr(s, name))(self) + else: + # no need to wrap non-tensor properties, + # and they likely don't depend on the actual contents of the tensor + return meta_attr + + namespace["__getattr__"] = __getattr__ + + # need to make a builder for the wrapped wrapper to copy the name, + # or else it fails with very cryptic error messages, + # because somehow the same string would end up in every closures + def mk_wrap(op_name: str, *, meta_noop: bool = False): + # need to wrap the wrapper to get self + def wrapped_special_op(self, *args, **kwargs): + return type(self)._wrap_fn( + getattr(type(self)._tensor_type, op_name), + meta_noop=meta_noop, + )(self, *args, **kwargs) + return wrapped_special_op + + # special methods bypass __getattr__, so they need to be added manually + # ref: https://docs.python.org/3/reference/datamodel.html#special-lookup + # NOTE: doing this from a metaclass is very convenient + # TODO: make this even more comprehensive + for binary_op in ( + "lt", "le", "eq", "ne", "ge", "gt", + "add", "and", "floordiv", "lshift", "mod", "mul", "matmul", + "or", "pow", "rshift", "sub", "truediv", "xor", + "iadd", "iand", "ifloordiv", "ilshift", "imod", "imul", "ior", "irshift", "isub", "ixor", + "radd", "rand", "rfloordiv", "rmul", "ror", "rpow", "rsub", "rtruediv", "rxor", + ): + attr_name = f"__{binary_op}__" + # evaluation on the meta tensor is needed in case there's broadcasting + namespace[attr_name] = mk_wrap(attr_name, meta_noop=False) + + for unary_op in ("not", "abs", "invert", "neg", "pos"): + attr_name = f"__{unary_op}__" + # the result of these operators usually has the same shape and dtype as the input, + # so evaluation on the meta tensor can be skipped. + namespace[attr_name] = mk_wrap(attr_name, meta_noop=True) + + for special_op in ( + "getitem", "setitem", "len", + ): + attr_name = f"__{special_op}__" + namespace[attr_name] = mk_wrap(attr_name, meta_noop=False) + + return super().__new__(cls, name, bases, namespace, **kwargs) + + +# Tree of lazy tensors +class LazyBase(ABC, metaclass=LazyMeta): + _tensor_type: type + _meta: Any + _data: Any | None + _args: tuple + _kwargs: dict[str, Any] + _func: Callable[[Any], Any] | None + + def __init__(self, *, meta: Any, data: Any | None = None, args: tuple = (), kwargs: dict[str, Any] | None = None, func: Callable[[Any], Any] | None = None): + super().__init__() + self._meta = meta + self._data = data + self._args = args + self._kwargs = kwargs if kwargs is not None else {} + self._func = func + assert self._func is not None or self._data is not None + + def __init_subclass__(cls) -> None: + if "_tensor_type" not in cls.__dict__: + raise TypeError(f"property '_tensor_type' must be defined for {cls!r}") + return super().__init_subclass__() + + @staticmethod + def _recurse_apply(o: Any, fn: Callable[[Any], Any]) -> Any: + # TODO: dict and set + if isinstance(o, (list, tuple)): + L = [] + for item in o: + L.append(LazyBase._recurse_apply(item, fn)) + if isinstance(o, tuple): + L = tuple(L) + return L + elif isinstance(o, LazyBase): + return fn(o) + else: + return o + + @classmethod + def _wrap_fn(cls, fn: Callable, *, use_self: LazyBase | None = None, meta_noop: bool | DTypeLike | tuple[DTypeLike, Callable[[tuple[int, ...]], tuple[int, ...]]] = False) -> Callable[[Any], Any]: + def wrapped_fn(*args, **kwargs): + if kwargs is None: + kwargs = {} + args = ((use_self,) if use_self is not None else ()) + args + + meta_args = LazyBase._recurse_apply(args, lambda t: t._meta) + # TODO: maybe handle tensors in kwargs too + + if isinstance(meta_noop, bool) and not meta_noop: + try: + res = fn(*meta_args, **kwargs) + except NotImplementedError: + # running some operations on PyTorch's Meta tensors can cause this exception + res = None + else: + # some operators don't need to actually run on the meta tensors + assert len(args) > 0 + res = args[0] + assert isinstance(res, cls) + res = res._meta + # allow operations to override the dtype and shape + if meta_noop is not True: + if isinstance(meta_noop, tuple): + dtype, shape = meta_noop + assert callable(shape) + res = cls.meta_with_dtype_and_shape(dtype, shape(res.shape)) # ty: ignore[call-top-callable] + else: + res = cls.meta_with_dtype_and_shape(meta_noop, res.shape) + + if isinstance(res, cls._tensor_type): + return cls(meta=cls.eager_to_meta(res), args=args, kwargs=kwargs, func=fn) + elif isinstance(res, tuple) and all(isinstance(t, cls._tensor_type) for t in res): + # share the evaluation between lazy tuple elements + shared_args: list = [args, None] + + def eager_tuple_element(a: list[Any], i: int = 0, /, **kw) -> LazyBase: + assert len(a) == 2 + if a[1] is None: + a[1] = fn(*a[0], **kw) + return a[1][i] + return tuple(cls(meta=cls.eager_to_meta(res[i]), args=(shared_args, i), kwargs=kwargs, func=eager_tuple_element) for i in range(len(res))) + else: + del res # not needed + # non-tensor return likely relies on the contents of the args + # (e.g. the result of torch.equal) + eager_args = cls.to_eager(args) + return fn(*eager_args, **kwargs) + return wrapped_fn + + @classmethod + def to_eager(cls, t: Any) -> Any: + def simple_to_eager(_t: LazyBase) -> Any: + if _t._data is not None: + return _t._data + + # NOTE: there's a recursion limit in Python (usually 1000) + + assert _t._func is not None + _t._args = cls._recurse_apply(_t._args, simple_to_eager) + _t._data = _t._func(*_t._args, **_t._kwargs) + # sanity check + assert _t._data is not None + assert _t._data.dtype == _t._meta.dtype + assert _t._data.shape == _t._meta.shape + + return _t._data + + # recurse into lists and/or tuples, keeping their structure + return cls._recurse_apply(t, simple_to_eager) + + @classmethod + def eager_to_meta(cls, t: Any) -> Any: + return cls.meta_with_dtype_and_shape(t.dtype, t.shape) + + # must be overridden, meta tensor init is backend-specific + @classmethod + @abstractmethod + def meta_with_dtype_and_shape(cls, dtype: Any, shape: Any) -> Any: pass + + @classmethod + def from_eager(cls, t: Any) -> Any: + if type(t) is cls: + # already lazy + return t + elif isinstance(t, cls._tensor_type): + return cls(meta=cls.eager_to_meta(t), data=t) + else: + return TypeError(f"{type(t)!r} is not compatible with {cls._tensor_type!r}") + + +class LazyNumpyTensor(LazyBase): + _tensor_type = np.ndarray + + shape: tuple[int, ...] # Makes the type checker happy in quants.py + + @classmethod + def meta_with_dtype_and_shape(cls, dtype: DTypeLike, shape: tuple[int, ...]) -> np.ndarray[Any, Any]: + # The initial idea was to use np.nan as the fill value, + # but non-float types like np.int16 can't use that. + # So zero it is. + cheat = np.zeros(1, dtype) + return np.lib.stride_tricks.as_strided(cheat, shape, (0 for _ in shape)) + + def astype(self, dtype, *args, **kwargs): + meta = type(self).meta_with_dtype_and_shape(dtype, self._meta.shape) + full_args = (self, dtype,) + args + return type(self)(meta=meta, args=full_args, kwargs=kwargs, func=(lambda a, *args, **kwargs: a.astype(*args, **kwargs))) + + def tofile(self, *args, **kwargs): + eager = LazyNumpyTensor.to_eager(self) + return eager.tofile(*args, **kwargs) + + # TODO: __array_function__ + + +# Tensor written to file one row-chunk at a time +class LazyChunkedTensor: + + def __init__( + self, chunks: list[Callable[[], np.ndarray]], shape: tuple[int, ...], dtype: DTypeLike, + qtype: Any = None, byteswap: bool = False, + ): + self._chunks = chunks + self._qtype = qtype + self._byteswap = byteswap + self.shape = tuple(shape) + self.dtype = np.dtype(dtype) + + @property + def nbytes(self) -> int: + n = self.dtype.itemsize + for d in self.shape: + n *= d + return n + + def numpy(self) -> LazyChunkedTensor: + return self + + def __array__(self, *args, **kwargs): + # numpy would otherwise make a 1-element object array of self, and write 8 bytes + raise TypeError("LazyChunkedTensor cannot become an ndarray, it is written in chunks") + + def quantize(self, qtype: Any) -> LazyChunkedTensor: + from .constants import GGMLQuantizationType + from .quants import QuantError, quant_shape_to_byte_shape + + if qtype == GGMLQuantizationType.F32: + shape, dtype = self.shape, np.dtype(np.float32) + elif qtype == GGMLQuantizationType.F16: + shape, dtype = self.shape, np.dtype(np.float16) + else: + try: + shape, dtype = quant_shape_to_byte_shape(self.shape, qtype), np.dtype(np.uint8) + except ValueError as e: + # raised here and not per chunk, so callers can still fall back to F16 + raise QuantError(str(e)) from e + return LazyChunkedTensor(self._chunks, shape, dtype, qtype, self._byteswap) + + def byteswap(self, inplace: bool = False) -> LazyChunkedTensor: + if inplace: + raise NotImplementedError("a chunked tensor cannot be byteswapped in place") + return LazyChunkedTensor(self._chunks, self.shape, self.dtype, self._qtype, not self._byteswap) + + def tofile(self, *args, **kwargs) -> None: + from .quants import quantize + + written = 0 + for load_chunk in self._chunks: + chunk = load_chunk() + if self._qtype is not None: + # exact only because chunks split on rows, and blocks never cross one + chunk = quantize(chunk, self._qtype) + if self._byteswap: + chunk = chunk.byteswap(inplace=False) + chunk.tofile(*args, **kwargs) + written += chunk.nbytes + del chunk + assert written == self.nbytes, f"chunked tensor wrote {written} bytes, expected {self.nbytes}" diff --git a/python/freetoken/third_party/gguf/quants.py b/python/freetoken/third_party/gguf/quants.py new file mode 100644 index 0000000000..80966b6ef1 --- /dev/null +++ b/python/freetoken/third_party/gguf/quants.py @@ -0,0 +1,1378 @@ +from __future__ import annotations +from abc import ABC, abstractmethod +from typing import Any, Callable, Sequence +from math import log2, ceil + +from numpy.typing import DTypeLike + +from .constants import GGML_QUANT_SIZES, GGMLQuantizationType, QK_K +from .lazy import LazyNumpyTensor + +import numpy as np + + +def quant_shape_to_byte_shape(shape: Sequence[int], quant_type: GGMLQuantizationType) -> tuple[int, ...]: + block_size, type_size = GGML_QUANT_SIZES[quant_type] + if shape[-1] % block_size != 0: + raise ValueError(f"Quantized tensor row size ({shape[-1]}) is not a multiple of {quant_type.name} block size ({block_size})") + return (*shape[:-1], shape[-1] // block_size * type_size) + + +def quant_shape_from_byte_shape(shape: Sequence[int], quant_type: GGMLQuantizationType) -> tuple[int, ...]: + block_size, type_size = GGML_QUANT_SIZES[quant_type] + if shape[-1] % type_size != 0: + raise ValueError(f"Quantized tensor bytes per row ({shape[-1]}) is not a multiple of {quant_type.name} type size ({type_size})") + return (*shape[:-1], shape[-1] // type_size * block_size) + + +# This is faster than np.vectorize and np.apply_along_axis because it works on more than one row at a time +def _apply_over_grouped_rows(func: Callable[[np.ndarray], np.ndarray], arr: np.ndarray, otype: DTypeLike, oshape: tuple[int, ...]) -> np.ndarray: + rows = arr.reshape((-1, arr.shape[-1])) + assert len(rows.shape) + osize = 1 + for dim in oshape: + osize *= dim + out = np.empty(shape=osize, dtype=otype) + # compute over groups of 16 rows (arbitrary, but seems good for performance) + n_groups = (rows.shape[0] // 16) or 1 + np.concatenate([func(group).ravel() for group in np.array_split(rows, n_groups)], axis=0, out=out) + return out.reshape(oshape) + + +# round away from zero +# ref: https://stackoverflow.com/a/59143326/22827863 +def np_roundf(n: np.ndarray) -> np.ndarray: + a = abs(n) + floored = np.floor(a) + b = floored + np.floor(2 * (a - floored)) + return np.sign(n) * b + + +class QuantError(Exception): ... + + +_type_traits: dict[GGMLQuantizationType, type[__Quant]] = {} + + +def quantize(data: np.ndarray, qtype: GGMLQuantizationType) -> np.ndarray: + if qtype == GGMLQuantizationType.F32: + return data.astype(np.float32, copy=False) + elif qtype == GGMLQuantizationType.F16: + return data.astype(np.float16, copy=False) + elif (q := _type_traits.get(qtype)) is not None: + return q.quantize(data) + else: + raise NotImplementedError(f"Quantization for {qtype.name} is not yet implemented") + + +def dequantize(data: np.ndarray, qtype: GGMLQuantizationType) -> np.ndarray: + if qtype == GGMLQuantizationType.F32: + return data.view(np.float32) + elif qtype == GGMLQuantizationType.F16: + return data.view(np.float16).astype(np.float32) + elif (q := _type_traits.get(qtype)) is not None: + return q.dequantize(data) + else: + raise NotImplementedError(f"Dequantization for {qtype.name} is not yet implemented") + + +class __Quant(ABC): + qtype: GGMLQuantizationType + block_size: int + type_size: int + + grid: np.ndarray[Any, np.dtype[np.float32]] | None = None + grid_shape: tuple[int, int] = (0, 0) + grid_map: tuple[int | float, ...] = () + grid_hex: bytes | None = None + + def __init__(self): + return TypeError("Quant conversion classes can't have instances") + + def __init_subclass__(cls, qtype: GGMLQuantizationType) -> None: + cls.qtype = qtype + cls.block_size, cls.type_size = GGML_QUANT_SIZES[qtype] + cls.__quantize_lazy: Any = LazyNumpyTensor._wrap_fn( + cls.__quantize_array, + meta_noop=(np.uint8, cls.__shape_to_bytes) + ) + cls.__dequantize_lazy: Any = LazyNumpyTensor._wrap_fn( + cls.__dequantize_array, + meta_noop=(np.float32, cls.__shape_from_bytes) + ) + assert qtype not in _type_traits + _type_traits[qtype] = cls + + @classmethod + def init_grid(cls): + if cls.grid is not None or cls.grid_hex is None: + return + + bits_per_elem = ceil(log2(len(cls.grid_map))) + assert bits_per_elem != 0, cls.qtype.name + elems_per_byte = 8 // bits_per_elem + + grid = np.frombuffer(cls.grid_hex, dtype=np.uint8) + # decode hexadecimal chars from grid + grid = grid.reshape((-1, 2)) + grid = (np.where(grid > 0x40, grid + 9, grid) & 0x0F) << np.array([4, 0], dtype=np.uint8).reshape((1, 2)) + grid = grid[..., 0] | grid[..., 1] + # unpack the grid values + grid = grid.reshape((-1, 1)) >> np.array([i for i in range(0, 8, 8 // elems_per_byte)], dtype=np.uint8).reshape((1, elems_per_byte)) + grid = (grid & ((1 << bits_per_elem) - 1)).reshape((-1, 1)) + grid_map = np.array(cls.grid_map, dtype=np.float32).reshape((1, -1)) + grid = np.take_along_axis(grid_map, grid, axis=-1) + cls.grid = grid.reshape((1, 1, *cls.grid_shape)) + + @classmethod + @abstractmethod + def quantize_blocks(cls, blocks: np.ndarray) -> np.ndarray: + raise NotImplementedError + + @classmethod + @abstractmethod + def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray: + raise NotImplementedError + + @classmethod + def quantize_rows(cls, rows: np.ndarray) -> np.ndarray: + rows = rows.astype(np.float32, copy=False) + shape = rows.shape + n_blocks = rows.size // cls.block_size + blocks = rows.reshape((n_blocks, cls.block_size)) + blocks = cls.quantize_blocks(blocks) + assert blocks.dtype == np.uint8 + assert blocks.shape[-1] == cls.type_size + return blocks.reshape(cls.__shape_to_bytes(shape)) + + @classmethod + def dequantize_rows(cls, rows: np.ndarray) -> np.ndarray: + rows = rows.view(np.uint8) + shape = rows.shape + n_blocks = rows.size // cls.type_size + blocks = rows.reshape((n_blocks, cls.type_size)) + blocks = cls.dequantize_blocks(blocks) + assert blocks.dtype == np.float32 + assert blocks.shape[-1] == cls.block_size + return blocks.reshape(cls.__shape_from_bytes(shape)) + + @classmethod + def __shape_to_bytes(cls, shape: Sequence[int]): + return quant_shape_to_byte_shape(shape, cls.qtype) + + @classmethod + def __shape_from_bytes(cls, shape: Sequence[int]): + return quant_shape_from_byte_shape(shape, cls.qtype) + + @classmethod + def __quantize_array(cls, array: np.ndarray) -> np.ndarray: + return _apply_over_grouped_rows(cls.quantize_rows, arr=array, otype=np.uint8, oshape=cls.__shape_to_bytes(array.shape)) + + @classmethod + def __dequantize_array(cls, array: np.ndarray) -> np.ndarray: + cls.init_grid() + return _apply_over_grouped_rows(cls.dequantize_rows, arr=array, otype=np.float32, oshape=cls.__shape_from_bytes(array.shape)) + + @classmethod + def __quantize_lazy(cls, lazy_tensor: LazyNumpyTensor, /) -> Any: + pass + + @classmethod + def __dequantize_lazy(cls, lazy_tensor: LazyNumpyTensor, /) -> Any: + pass + + @classmethod + def can_quantize(cls, tensor: np.ndarray | LazyNumpyTensor) -> bool: + return tensor.shape[-1] % cls.block_size == 0 + + @classmethod + def quantize(cls, tensor: np.ndarray | LazyNumpyTensor) -> np.ndarray: + if not cls.can_quantize(tensor): + raise QuantError(f"Can't quantize tensor with shape {tensor.shape} to {cls.qtype.name}") + if isinstance(tensor, LazyNumpyTensor): + return cls.__quantize_lazy(tensor) + else: + return cls.__quantize_array(tensor) + + @classmethod + def dequantize(cls, tensor: np.ndarray | LazyNumpyTensor) -> np.ndarray: + if isinstance(tensor, LazyNumpyTensor): + return cls.__dequantize_lazy(tensor) + else: + return cls.__dequantize_array(tensor) + + +class BF16(__Quant, qtype=GGMLQuantizationType.BF16): + @classmethod + # same as ggml_compute_fp32_to_bf16 in ggml-impl.h + def quantize_blocks(cls, blocks: np.ndarray) -> np.ndarray: + n = blocks.view(np.uint32) + # force nan to quiet + n = np.where((n & 0x7fffffff) > 0x7f800000, (n & np.uint32(0xffff0000)) | np.uint32(64 << 16), n) + # round to nearest even + n = (np.uint64(n) + (0x7fff + ((n >> 16) & 1))) >> 16 + return n.astype(np.uint16).view(np.uint8) + + @classmethod + def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray: + return (blocks.view(np.int16).astype(np.int32) << 16).view(np.float32) + + +class Q4_0(__Quant, qtype=GGMLQuantizationType.Q4_0): + @classmethod + def quantize_blocks(cls, blocks: np.ndarray) -> np.ndarray: + n_blocks = blocks.shape[0] + + imax = abs(blocks).argmax(axis=-1, keepdims=True) + max = np.take_along_axis(blocks, imax, axis=-1) + + d = max / -8 + with np.errstate(divide="ignore"): + id = np.where(d == 0, 0, 1 / d) + qs = np.trunc((blocks * id) + np.float32(8.5), dtype=np.float32).astype(np.uint8).clip(0, 15) + + qs = qs.reshape((n_blocks, 2, cls.block_size // 2)) + qs = qs[..., 0, :] | (qs[..., 1, :] << np.uint8(4)) + + d = d.astype(np.float16).view(np.uint8) + + return np.concatenate([d, qs], axis=-1) + + @classmethod + def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray: + n_blocks = blocks.shape[0] + + d, qs = np.hsplit(blocks, [2]) + + d = d.view(np.float16).astype(np.float32) + + qs = qs.reshape((n_blocks, -1, 1, cls.block_size // 2)) >> np.array([0, 4], dtype=np.uint8).reshape((1, 1, 2, 1)) + qs = (qs & np.uint8(0x0F)).reshape((n_blocks, -1)).astype(np.int8) - np.int8(8) + + return (d * qs.astype(np.float32)) + + +class Q4_1(__Quant, qtype=GGMLQuantizationType.Q4_1): + @classmethod + def quantize_blocks(cls, blocks: np.ndarray) -> np.ndarray: + n_blocks = blocks.shape[0] + + max = blocks.max(axis=-1, keepdims=True) + min = blocks.min(axis=-1, keepdims=True) + + d = (max - min) / 15 + with np.errstate(divide="ignore"): + id = np.where(d == 0, 0, 1 / d) + qs = np.trunc((blocks - min) * id + np.float32(0.5), dtype=np.float32).astype(np.uint8).clip(0, 15) + + qs = qs.reshape((n_blocks, 2, cls.block_size // 2)) + qs = qs[..., 0, :] | (qs[..., 1, :] << np.uint8(4)) + + d = d.astype(np.float16).view(np.uint8) + m = min.astype(np.float16).view(np.uint8) + + return np.concatenate([d, m, qs], axis=-1) + + @classmethod + def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray: + n_blocks = blocks.shape[0] + + d, rest = np.hsplit(blocks, [2]) + m, qs = np.hsplit(rest, [2]) + + d = d.view(np.float16).astype(np.float32) + m = m.view(np.float16).astype(np.float32) + + qs = qs.reshape((n_blocks, -1, 1, cls.block_size // 2)) >> np.array([0, 4], dtype=np.uint8).reshape((1, 1, 2, 1)) + qs = (qs & np.uint8(0x0F)).reshape((n_blocks, -1)).astype(np.float32) + + return (d * qs) + m + + +class Q5_0(__Quant, qtype=GGMLQuantizationType.Q5_0): + @classmethod + def quantize_blocks(cls, blocks: np.ndarray) -> np.ndarray: + n_blocks = blocks.shape[0] + + imax = abs(blocks).argmax(axis=-1, keepdims=True) + max = np.take_along_axis(blocks, imax, axis=-1) + + d = max / -16 + with np.errstate(divide="ignore"): + id = np.where(d == 0, 0, 1 / d) + q = np.trunc((blocks * id) + np.float32(16.5), dtype=np.float32).astype(np.uint8).clip(0, 31) + + qs = q.reshape((n_blocks, 2, cls.block_size // 2)) + qs = (qs[..., 0, :] & np.uint8(0x0F)) | (qs[..., 1, :] << np.uint8(4)) + + qh = np.packbits(q.reshape((n_blocks, 1, 32)) >> np.uint8(4), axis=-1, bitorder="little").reshape(n_blocks, 4) + + d = d.astype(np.float16).view(np.uint8) + + return np.concatenate([d, qh, qs], axis=-1) + + @classmethod + def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray: + n_blocks = blocks.shape[0] + + d, rest = np.hsplit(blocks, [2]) + qh, qs = np.hsplit(rest, [4]) + + d = d.view(np.float16).astype(np.float32) + qh = qh.view(np.uint32) + + qh = qh.reshape((n_blocks, 1)) >> np.array([i for i in range(32)], dtype=np.uint32).reshape((1, 32)) + ql = qs.reshape((n_blocks, -1, 1, cls.block_size // 2)) >> np.array([0, 4], dtype=np.uint8).reshape((1, 1, 2, 1)) + qh = (qh & np.uint32(0x01)).astype(np.uint8) + ql = (ql & np.uint8(0x0F)).reshape((n_blocks, -1)) + + qs = (ql | (qh << np.uint8(4))).astype(np.int8) - np.int8(16) + + return (d * qs.astype(np.float32)) + + +class Q5_1(__Quant, qtype=GGMLQuantizationType.Q5_1): + @classmethod + def quantize_blocks(cls, blocks: np.ndarray) -> np.ndarray: + n_blocks = blocks.shape[0] + + max = blocks.max(axis=-1, keepdims=True) + min = blocks.min(axis=-1, keepdims=True) + + d = (max - min) / 31 + with np.errstate(divide="ignore"): + id = np.where(d == 0, 0, 1 / d) + q = np.trunc((blocks - min) * id + np.float32(0.5), dtype=np.float32).astype(np.uint8).clip(0, 31) + + qs = q.reshape((n_blocks, 2, cls.block_size // 2)) + qs = (qs[..., 0, :] & np.uint8(0x0F)) | (qs[..., 1, :] << np.uint8(4)) + + qh = np.packbits(q.reshape((n_blocks, 1, 32)) >> np.uint8(4), axis=-1, bitorder="little").reshape(n_blocks, 4) + + d = d.astype(np.float16).view(np.uint8) + m = min.astype(np.float16).view(np.uint8) + + return np.concatenate([d, m, qh, qs], axis=-1) + + @classmethod + def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray: + n_blocks = blocks.shape[0] + + d, rest = np.hsplit(blocks, [2]) + m, rest = np.hsplit(rest, [2]) + qh, qs = np.hsplit(rest, [4]) + + d = d.view(np.float16).astype(np.float32) + m = m.view(np.float16).astype(np.float32) + qh = qh.view(np.uint32) + + qh = qh.reshape((n_blocks, 1)) >> np.array([i for i in range(32)], dtype=np.uint32).reshape((1, 32)) + ql = qs.reshape((n_blocks, -1, 1, cls.block_size // 2)) >> np.array([0, 4], dtype=np.uint8).reshape((1, 1, 2, 1)) + qh = (qh & np.uint32(0x01)).astype(np.uint8) + ql = (ql & np.uint8(0x0F)).reshape((n_blocks, -1)) + + qs = (ql | (qh << np.uint8(4))).astype(np.float32) + + return (d * qs) + m + + +class Q8_0(__Quant, qtype=GGMLQuantizationType.Q8_0): + @classmethod + # Implementation of Q8_0 with bit-exact same results as reference implementation in ggml-quants.c + def quantize_blocks(cls, blocks: np.ndarray) -> np.ndarray: + + d = abs(blocks).max(axis=1, keepdims=True) / 127 + with np.errstate(divide="ignore"): + id = np.where(d == 0, 0, 1 / d) + qs = np_roundf(blocks * id) + + # (n_blocks, 2) + d = d.astype(np.float16).view(np.uint8) + # (n_blocks, block_size) + qs = qs.astype(np.int8).view(np.uint8) + + return np.concatenate([d, qs], axis=1) + + @classmethod + def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray: + d, x = np.split(blocks, [2], axis=1) + d = d.view(np.float16).astype(np.float32) + x = x.view(np.int8).astype(np.float32) + + return (x * d) + + +class Q2_K(__Quant, qtype=GGMLQuantizationType.Q2_K): + @classmethod + def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray: + n_blocks = blocks.shape[0] + + scales, rest = np.hsplit(blocks, [QK_K // 16]) + qs, rest = np.hsplit(rest, [QK_K // 4]) + d, dmin = np.hsplit(rest, [2]) + + d = d.view(np.float16).astype(np.float32) + dmin = dmin.view(np.float16).astype(np.float32) + + # (n_blocks, 16, 1) + dl = (d * (scales & 0xF).astype(np.float32)).reshape((n_blocks, QK_K // 16, 1)) + ml = (dmin * (scales >> 4).astype(np.float32)).reshape((n_blocks, QK_K // 16, 1)) + + shift = np.array([0, 2, 4, 6], dtype=np.uint8).reshape((1, 1, 4, 1)) + + qs = (qs.reshape((n_blocks, -1, 1, 32)) >> shift) & np.uint8(3) + + qs = qs.reshape((n_blocks, QK_K // 16, 16)).astype(np.float32) + + qs = dl * qs - ml + + return qs.reshape((n_blocks, -1)) + + +class Q3_K(__Quant, qtype=GGMLQuantizationType.Q3_K): + @classmethod + def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray: + n_blocks = blocks.shape[0] + + hmask, rest = np.hsplit(blocks, [QK_K // 8]) + qs, rest = np.hsplit(rest, [QK_K // 4]) + scales, d = np.hsplit(rest, [12]) + + d = d.view(np.float16).astype(np.float32) + + # The scales are packed at 6-bit each in this pattern: + # 0: IIIIAAAA + # 1: JJJJBBBB + # 2: KKKKCCCC + # 3: LLLLDDDD + # 4: MMMMEEEE + # 5: NNNNFFFF + # 6: OOOOGGGG + # 7: PPPPHHHH + # 8: MMIIEEAA + # 9: NNJJFFBB + # 10: OOKKGGCC + # 11: PPLLHHDD + lscales, hscales = np.hsplit(scales, [8]) + lscales = lscales.reshape((n_blocks, 1, 8)) >> np.array([0, 4], dtype=np.uint8).reshape((1, 2, 1)) + lscales = lscales.reshape((n_blocks, 16)) + hscales = hscales.reshape((n_blocks, 1, 4)) >> np.array([0, 2, 4, 6], dtype=np.uint8).reshape((1, 4, 1)) + hscales = hscales.reshape((n_blocks, 16)) + scales = (lscales & np.uint8(0x0F)) | ((hscales & np.uint8(0x03)) << np.uint8(4)) + scales = (scales.astype(np.int8) - np.int8(32)).astype(np.float32) + + dl = (d * scales).reshape((n_blocks, 16, 1)) + + ql = qs.reshape((n_blocks, -1, 1, 32)) >> np.array([0, 2, 4, 6], dtype=np.uint8).reshape((1, 1, 4, 1)) + qh = hmask.reshape(n_blocks, -1, 1, 32) >> np.array([i for i in range(8)], dtype=np.uint8).reshape((1, 1, 8, 1)) + ql = ql.reshape((n_blocks, 16, QK_K // 16)) & np.uint8(3) + qh = (qh.reshape((n_blocks, 16, QK_K // 16)) & np.uint8(1)) + qh = qh ^ np.uint8(1) # strangely, the offset is zero when the bitmask is 1 + q = (ql.astype(np.int8) - (qh << np.uint8(2)).astype(np.int8)).astype(np.float32) + + return (dl * q).reshape((n_blocks, QK_K)) + + +class Q4_K(__Quant, qtype=GGMLQuantizationType.Q4_K): + K_SCALE_SIZE = 12 + + @staticmethod + def get_scale_min(scales: np.ndarray) -> tuple[np.ndarray, np.ndarray]: + n_blocks = scales.shape[0] + scales = scales.view(np.uint8) + ### Unpacking the following: ### + # 0 EEAAAAAA + # 1 FFBBBBBB + # 2 GGCCCCCC + # 3 HHDDDDDD + # 4 eeaaaaaa + # 5 ffbbbbbb + # 6 ggcccccc + # 7 hhdddddd + # 8 eeeeEEEE + # 9 ffffFFFF + # 10 ggggGGGG + # 11 hhhhHHHH + scales = scales.reshape((n_blocks, 3, 4)) + d, m, m_d = np.split(scales, 3, axis=-2) + + sc = np.concatenate([d & 0x3F, (m_d & 0x0F) | ((d >> 2) & 0x30)], axis=-1) + min = np.concatenate([m & 0x3F, (m_d >> 4) | ((m >> 2) & 0x30)], axis=-1) + + return (sc.reshape((n_blocks, 8)), min.reshape((n_blocks, 8))) + + @classmethod + def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray: + n_blocks = blocks.shape[0] + + d, rest = np.hsplit(blocks, [2]) + dmin, rest = np.hsplit(rest, [2]) + scales, qs = np.hsplit(rest, [cls.K_SCALE_SIZE]) + + d = d.view(np.float16).astype(np.float32) + dmin = dmin.view(np.float16).astype(np.float32) + + sc, m = Q4_K.get_scale_min(scales) + + d = (d * sc.astype(np.float32)).reshape((n_blocks, -1, 1)) + dm = (dmin * m.astype(np.float32)).reshape((n_blocks, -1, 1)) + + qs = qs.reshape((n_blocks, -1, 1, 32)) >> np.array([0, 4], dtype=np.uint8).reshape((1, 1, 2, 1)) + qs = (qs & np.uint8(0x0F)).reshape((n_blocks, -1, 32)).astype(np.float32) + + return (d * qs - dm).reshape((n_blocks, QK_K)) + + +class Q5_K(__Quant, qtype=GGMLQuantizationType.Q5_K): + @classmethod + def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray: + n_blocks = blocks.shape[0] + + d, rest = np.hsplit(blocks, [2]) + dmin, rest = np.hsplit(rest, [2]) + scales, rest = np.hsplit(rest, [Q4_K.K_SCALE_SIZE]) + qh, qs = np.hsplit(rest, [QK_K // 8]) + + d = d.view(np.float16).astype(np.float32) + dmin = dmin.view(np.float16).astype(np.float32) + + sc, m = Q4_K.get_scale_min(scales) + + d = (d * sc.astype(np.float32)).reshape((n_blocks, -1, 1)) + dm = (dmin * m.astype(np.float32)).reshape((n_blocks, -1, 1)) + + ql = qs.reshape((n_blocks, -1, 1, 32)) >> np.array([0, 4], dtype=np.uint8).reshape((1, 1, 2, 1)) + qh = qh.reshape((n_blocks, -1, 1, 32)) >> np.array([i for i in range(8)], dtype=np.uint8).reshape((1, 1, 8, 1)) + ql = (ql & np.uint8(0x0F)).reshape((n_blocks, -1, 32)) + qh = (qh & np.uint8(0x01)).reshape((n_blocks, -1, 32)) + q = (ql | (qh << np.uint8(4))).astype(np.float32) + + return (d * q - dm).reshape((n_blocks, QK_K)) + + +class Q6_K(__Quant, qtype=GGMLQuantizationType.Q6_K): + @classmethod + def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray: + n_blocks = blocks.shape[0] + + ql, rest = np.hsplit(blocks, [QK_K // 2]) + qh, rest = np.hsplit(rest, [QK_K // 4]) + scales, d = np.hsplit(rest, [QK_K // 16]) + + scales = scales.view(np.int8).astype(np.float32) + d = d.view(np.float16).astype(np.float32) + d = (d * scales).reshape((n_blocks, QK_K // 16, 1)) + + ql = ql.reshape((n_blocks, -1, 1, 64)) >> np.array([0, 4], dtype=np.uint8).reshape((1, 1, 2, 1)) + ql = (ql & np.uint8(0x0F)).reshape((n_blocks, -1, 32)) + qh = qh.reshape((n_blocks, -1, 1, 32)) >> np.array([0, 2, 4, 6], dtype=np.uint8).reshape((1, 1, 4, 1)) + qh = (qh & np.uint8(0x03)).reshape((n_blocks, -1, 32)) + q = (ql | (qh << np.uint8(4))).astype(np.int8) - np.int8(32) + q = q.reshape((n_blocks, QK_K // 16, -1)).astype(np.float32) + + return (d * q).reshape((n_blocks, QK_K)) + + +class TQ1_0(__Quant, qtype=GGMLQuantizationType.TQ1_0): + @classmethod + def quantize_blocks(cls, blocks: np.ndarray) -> np.ndarray: + n_blocks = blocks.shape[0] + + d = abs(blocks).max(axis=-1, keepdims=True) + with np.errstate(divide="ignore"): + id = np.where(d == 0, 0, 1 / d) + qs = np_roundf(blocks * id) + qs = (qs.astype(np.int8) + np.int8(1)).astype(np.uint8) + + qs0, qs1, qh = qs[..., :(32 * 5)], qs[..., (32 * 5):(48 * 5)], qs[..., (48 * 5):] + qs0 = qs0.reshape((n_blocks, -1, 5, 32)) * np.array([81, 27, 9, 3, 1], dtype=np.uint8).reshape((1, 1, 5, 1)) + qs0 = np.sum(qs0, axis=-2).reshape((n_blocks, -1)) + qs1 = qs1.reshape((n_blocks, -1, 5, 16)) * np.array([81, 27, 9, 3, 1], dtype=np.uint8).reshape((1, 1, 5, 1)) + qs1 = np.sum(qs1, axis=-2).reshape((n_blocks, -1)) + qh = qh.reshape((n_blocks, -1, 4, 4)) * np.array([81, 27, 9, 3], dtype=np.uint8).reshape((1, 1, 4, 1)) + qh = np.sum(qh, axis=-2).reshape((n_blocks, -1)) + qs = np.concatenate([qs0, qs1, qh], axis=-1) + qs = (qs.astype(np.uint16) * 256 + (243 - 1)) // 243 + + qs = qs.astype(np.uint8) + d = d.astype(np.float16).view(np.uint8) + + return np.concatenate([qs, d], axis=-1) + + @classmethod + def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray: + n_blocks = blocks.shape[0] + + qs, rest = np.hsplit(blocks, [(QK_K - 4 * QK_K // 64) // 5]) + qh, d = np.hsplit(rest, [QK_K // 64]) + + d = d.view(np.float16).astype(np.float32) + + qs0, qs1 = qs[..., :32], qs[..., 32:] + qs0 = qs0.reshape((n_blocks, -1, 1, 32)) * np.array([1, 3, 9, 27, 81], dtype=np.uint8).reshape((1, 1, 5, 1)) + qs0 = qs0.reshape((n_blocks, -1)) + qs1 = qs1.reshape((n_blocks, -1, 1, 16)) * np.array([1, 3, 9, 27, 81], dtype=np.uint8).reshape((1, 1, 5, 1)) + qs1 = qs1.reshape((n_blocks, -1)) + qh = qh.reshape((n_blocks, -1, 1, 4)) * np.array([1, 3, 9, 27], dtype=np.uint8).reshape((1, 1, 4, 1)) + qh = qh.reshape((n_blocks, -1)) + qs = np.concatenate([qs0, qs1, qh], axis=-1) + qs = ((qs.astype(np.uint16) * 3) >> 8).astype(np.int8) - np.int8(1) + + return (d * qs.astype(np.float32)) + + +class TQ2_0(__Quant, qtype=GGMLQuantizationType.TQ2_0): + @classmethod + def quantize_blocks(cls, blocks: np.ndarray) -> np.ndarray: + n_blocks = blocks.shape[0] + + d = abs(blocks).max(axis=-1, keepdims=True) + with np.errstate(divide="ignore"): + id = np.where(d == 0, 0, 1 / d) + qs = np_roundf(blocks * id) + qs = (qs.astype(np.int8) + np.int8(1)).astype(np.uint8) + + qs = qs.reshape((n_blocks, -1, 4, 32)) << np.array([0, 2, 4, 6], dtype=np.uint8).reshape((1, 1, 4, 1)) + qs = qs[..., 0, :] | qs[..., 1, :] | qs[..., 2, :] | qs[..., 3, :] + qs = qs.reshape((n_blocks, -1)) + + d = d.astype(np.float16).view(np.uint8) + + return np.concatenate([qs, d], axis=-1) + + @classmethod + def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray: + n_blocks = blocks.shape[0] + + qs, d = np.hsplit(blocks, [QK_K // 4]) + + d = d.view(np.float16).astype(np.float32) + + qs = qs.reshape((n_blocks, -1, 1, 32)) >> np.array([0, 2, 4, 6], dtype=np.uint8).reshape((1, 1, 4, 1)) + qs = (qs & 0x03).reshape((n_blocks, -1)).astype(np.int8) - np.int8(1) + + return (d * qs.astype(np.float32)) + + +class MXFP4(__Quant, qtype=GGMLQuantizationType.MXFP4): + # e2m1 values (doubled) + # ref: https://www.opencompute.org/documents/ocp-microscaling-formats-mx-v1-0-spec-final-pdf + kvalues = (0, 1, 2, 3, 4, 6, 8, 12, 0, -1, -2, -3, -4, -6, -8, -12) + + @staticmethod + # see ggml_e8m0_to_fp32_half in ggml-impl.h + def e8m0_to_fp32_half(x: np.ndarray) -> np.ndarray: + bits = np.where(x < 2, np.uint32(0x00200000) << np.uint32(x), np.uint32(x - 1) << np.uint32(23)) + return bits.view(np.float32) + + @classmethod + def quantize_blocks(cls, blocks: np.ndarray) -> np.ndarray: + n_blocks = blocks.shape[0] + + d = abs(blocks).max(axis=-1, keepdims=True) + + with np.errstate(divide="ignore"): + e = np.where(d > 0, np.floor(np.log2(d)) - 2 + 127, 0).astype(np.uint8) + + d = cls.e8m0_to_fp32_half(e) + + kvalues = np.array(cls.kvalues, dtype=np.int8).reshape((1, 1, 16)) + + errs = np.abs(d.reshape((n_blocks, 1, 1)) * kvalues.astype(np.float32) - blocks.reshape((n_blocks, cls.block_size, 1))) + best = np.argmin(errs, axis=-1, keepdims=True) + + qs = best.reshape(n_blocks, 2, cls.block_size // 2).astype(np.uint8) + qs = qs[:, 0] | (qs[:, 1] << np.uint8(4)) + + qs = qs.reshape((n_blocks, cls.block_size // 2)) + + return np.concatenate([e, qs], axis=-1) + + @classmethod + def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray: + n_blocks = blocks.shape[0] + + e, qs = np.hsplit(blocks, [1]) + + d = cls.e8m0_to_fp32_half(e) + + qs = qs.reshape((n_blocks, 1, cls.block_size // 2)) >> np.array([0, 4], dtype=np.uint8).reshape((1, 2, 1)) + qs = (qs & np.uint8(0x0F)).view(np.int8) + + kvalues = np.array(cls.kvalues, dtype=np.int8).reshape(1, 1, 16) + qs = np.take_along_axis(kvalues, qs, axis=-1).reshape((n_blocks, cls.block_size)) + + return (d * qs.astype(np.float32)) + + +class NVFP4(__Quant, qtype=GGMLQuantizationType.NVFP4): + # E2M1 values doubled (kvalues_mxfp4 convention) + kvalues = (0, 1, 2, 3, 4, 6, 8, 12, 0, -1, -2, -3, -4, -6, -8, -12) + + @staticmethod + def ue4m3_to_fp32(x: np.ndarray) -> np.ndarray: + """Decode unsigned E4M3 (bias=7) to float, with 0.5 factor for kvalues convention.""" + exp = (x >> 3).astype(np.int32) & 0xF + man = (x & 0x7).astype(np.float32) + raw = np.where( + exp == 0, + man * 2**-9, + (1.0 + man / 8.0) * (2.0 ** (exp.astype(np.float32) - 7))) + return np.where((x == 0) | (x == 0x7F), 0.0, raw * 0.5) + + @staticmethod + def fp32_to_ue4m3(x: np.ndarray) -> np.ndarray: + """Vectorized float32 to unsigned E4M3, matching ggml_fp32_to_ue4m3 in C.""" + x = np.clip(x, 0.0, 448.0).astype(np.float32) + bits = x.view(np.uint32) + fp32_exp = ((bits >> 23) & 0xFF).astype(np.int32) - 127 + fp32_man = ((bits >> 20) & 0x7).astype(np.int32) + ue4m3_exp = fp32_exp + 7 + + # Subnormal + sub_man = np.clip((x * 512.0 + 0.5).astype(np.int32), 0, 7) + sub_result = np.where(sub_man >= 1, sub_man, 0).astype(np.uint8) + + # Normal with rounding + round_bit = ((bits >> 19) & 1).astype(np.int32) + man = fp32_man + round_bit + exp = ue4m3_exp.copy() + overflow = man > 7 + man = np.where(overflow, 0, man) + exp = np.where(overflow, exp + 1, exp) + normal_result = np.where(exp >= 15, np.uint8(0x7E), ((exp << 3) | man).astype(np.uint8)) + + return np.where(x <= 0.0, np.uint8(0), + np.where(ue4m3_exp <= 0, sub_result, + np.where(ue4m3_exp >= 15, np.uint8(0x7E), normal_result))) + + @classmethod + def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray: + n_super = blocks.shape[0] + + d_bytes, qs = np.hsplit(blocks, [4]) + d = cls.ue4m3_to_fp32(d_bytes).reshape(n_super, 4, 1) # (n_super, 4, 1) + + qs = qs.reshape(n_super, 4, 8) + lo = (qs & np.uint8(0x0F)).view(np.int8) + hi = (qs >> np.uint8(4)).view(np.int8) + vals = np.concatenate([lo, hi], axis=-1) # (n_super, 4, 16) + + kvalues = np.array(cls.kvalues, dtype=np.int8).reshape(1, 1, 16) + vals = np.take_along_axis(kvalues, vals, axis=-1) + + return (d * vals.astype(np.float32)).reshape(n_super, 64) + + +class IQ2_XXS(__Quant, qtype=GGMLQuantizationType.IQ2_XXS): + ksigns: bytes = ( + b"\x00\x81\x82\x03\x84\x05\x06\x87\x88\x09\x0a\x8b\x0c\x8d\x8e\x0f" + b"\x90\x11\x12\x93\x14\x95\x96\x17\x18\x99\x9a\x1b\x9c\x1d\x1e\x9f" + b"\xa0\x21\x22\xa3\x24\xa5\xa6\x27\x28\xa9\xaa\x2b\xac\x2d\x2e\xaf" + b"\x30\xb1\xb2\x33\xb4\x35\x36\xb7\xb8\x39\x3a\xbb\x3c\xbd\xbe\x3f" + b"\xc0\x41\x42\xc3\x44\xc5\xc6\x47\x48\xc9\xca\x4b\xcc\x4d\x4e\xcf" + b"\x50\xd1\xd2\x53\xd4\x55\x56\xd7\xd8\x59\x5a\xdb\x5c\xdd\xde\x5f" + b"\x60\xe1\xe2\x63\xe4\x65\x66\xe7\xe8\x69\x6a\xeb\x6c\xed\xee\x6f" + b"\xf0\x71\x72\xf3\x74\xf5\xf6\x77\x78\xf9\xfa\x7b\xfc\x7d\x7e\xff" + ) + + # iq2xxs_grid, but with each byte of the original packed in 2 bits, + # by mapping 0x08 to 0, 0x19 to 1, and 0x2b to 2. + grid_shape = (256, 8) + grid_map = (0x08, 0x19, 0x2b) + grid_hex = ( + b"00000200050008000a00110014002000220028002a0041004400500058006100" + b"6400800082008a00a20001010401100115014001840198010002020222028202" + b"010404041004210424044004420448046004810484049004a404000502050805" + b"200546056905800591050906100640068406a406000805080808140828084108" + b"440850085208880804094009020a140a01100410101021104010601084109010" + b"951000110811201150115a118011241245120014081420142514491480141815" + b"6215001616160118041810184018811800190519a019511a002002200a204420" + b"6120802082202921482100220222012404241024402456240025412564259026" + b"082820289428442a014004401040184021402440404048405640604081408440" + b"9040004120416141804185410142104248425642684200440844204480449944" + b"124524450046014804481048404845480049584961498249454a904a00500850" + b"1150195020508050885004514251a4519152905492540a550156545600581158" + b"195864584059085a046010604060686000615561186260620064056410651265" + b"84654268008002800a8041808280048118814081118201840484108415844084" + b"608400854685948509864086608602880489118a0490109024904090a1901691" + b"8091459200942294449451958198209902a050a085a009a100a218a450a804a9" + ) + + @classmethod + def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray: + n_blocks = blocks.shape[0] + + d, qs = np.hsplit(blocks, [2]) + + d = d.view(np.float16).astype(np.float32) + + qs = qs.view(np.uint32).reshape(n_blocks, -1, 2) + + db = d * (np.float32(0.5) + (qs[..., 1] >> 28).astype(np.float32)) * np.float32(0.25) + db = db.reshape((n_blocks, -1, 1, 1)) + + # get the sign indices and unpack the bits + signs = qs[..., 1].reshape((n_blocks, -1, 1)) >> np.array([0, 7, 14, 21], dtype=np.uint32).reshape((1, 1, 4)) + ksigns = np.frombuffer(cls.ksigns, dtype=np.uint8).reshape((1, 1, 1, 128)) + signs = (signs & np.uint32(0x7F)).reshape((n_blocks, -1, 4, 1)) + signs = np.take_along_axis(ksigns, signs, axis=-1) + signs = signs.reshape((n_blocks, -1, 4, 1)) >> np.array([i for i in range(8)], dtype=np.uint8).reshape((1, 1, 1, 8)) + signs = signs & np.uint8(0x01) + signs = np.where(signs == 0, np.float32(1), np.float32(-1)) + signs = signs.reshape((n_blocks, -1, 4, 8)) + + assert cls.grid is not None + grid = np.take_along_axis(cls.grid, qs[..., 0].copy().view(np.uint8).reshape((n_blocks, -1, 1, 1)), axis=-2) + grid = grid.reshape((n_blocks, -1, 4, 8)) + + return (db * grid * signs).reshape((n_blocks, -1)) + + +class IQ2_XS(__Quant, qtype=GGMLQuantizationType.IQ2_XS): + # iq2xs_grid, but with each byte of the original packed in 2 bits, + # by mapping 0x08 to 0, 0x19 to 1, and 0x2b to 2. + grid_shape = (512, 8) + grid_map = (0x08, 0x19, 0x2b) + grid_hex = ( + b"00000200050008000a0011001400160019002000220025002800410044004600" + b"49005000520055005800610064008000820085008800910094009900a0000101" + b"04010601090110011201150118011a0121012401400142014501480151015401" + b"6001680181018401900100020202050208021102140220024102440250025502" + b"80028a0201040404060409041004120415041804210424044004420445044804" + b"5104540456046004810484049004000502050505080511051405200541054405" + b"500561058005010604061006260640064206840600080208050808080a081108" + b"14082008250841084408500858088008a008aa08010904091009400981098909" + b"000a200a280a960aa00a01100410061009101010121015101810211024104010" + b"4210451048105110541060106a10811084109010001102110511081111111411" + b"2011411144115011801194119611011204120612101240126012001402140514" + b"0814111414142014411444144914501464148014011504151015401500161416" + b"49160118041810181218401854188618001905196619511aa91a002002200520" + b"08200a201120142020204120442050208020a020012104211021402148216521" + b"002222228022a82201240424102429244024002541255225992501261a26a626" + b"002808280a28202855288828a22868299029082a202a822a882a8a2a01400440" + b"0640094010401240154018402140244040404240454048404a40514054406040" + b"6540814084409040004102410541084111411441204141414441504180418541" + b"a241014204421042124229424042004402440544084411441444194420444144" + b"4444504480449444014504451045244540459a4500460a464446504601480448" + b"1048404845485448624800491149444950496949044a00500250055008501150" + b"145020502850415044505050805001510451105115514051425100524452aa52" + b"0154045410542154405460548154a154005508558055885521566856a1560058" + b"14584158505899581a5940594259855a0160046010604060546062608660a960" + b"006124624a62926200641664106540654565a46501686a682569066a546a626a" + b"00800280058008801180148020802a8041804480508080808280a880aa800181" + b"0481068110814081518159810082208280828282a082a8820184048410841284" + b"158440846084898400854485a58518866a860088088825885a8880888288a888" + b"0689228a808a888a968aa88a0190049010904090569084900091229164915692" + b"89920094059444945094589429959095929541965198a6984999159a609a00a0" + b"02a008a00aa020a02aa0a0a051a159a1a6a100a202a208a22aa280a2a0a240a4" + b"95a465a698a60aa820a822a828a8a0a8a8a804a984a986a928aa2aaa91aaaaaa" + ) + + @classmethod + def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray: + n_blocks = blocks.shape[0] + + d, rest = np.hsplit(blocks, [2]) + qs, scales = np.hsplit(rest, [2 * QK_K // 8]) + + d = d.view(np.float16).astype(np.float32) + qs = qs.view(np.uint16) + + scales = scales.reshape((n_blocks, -1, 1)) >> np.array([0, 4], dtype=np.uint8).reshape((1, 1, 2)) + scales = (scales & 0x0F).reshape((n_blocks, -1)) + db = d * (np.float32(0.5) + scales) * np.float32(0.25) + db = db.reshape((n_blocks, -1, 1, 1)) + + # get the sign indices and unpack the bits + signs = np.frombuffer(IQ2_XXS.ksigns, dtype=np.uint8).reshape(1, 1, 128) + signs = np.take_along_axis(signs, (qs >> 9).reshape((n_blocks, -1, 1)), axis=-1) + signs = signs.reshape((n_blocks, -1, 1)) >> np.array([i for i in range(8)], dtype=np.uint8).reshape((1, 1, 8)) + signs = signs & np.uint8(0x01) + signs = np.where(signs == 0, np.float32(1), np.float32(-1)) + signs = signs.reshape((n_blocks, -1, 2, 8)) + + assert cls.grid is not None + grid = np.take_along_axis(cls.grid, (qs & np.uint16(511)).reshape((n_blocks, -1, 1, 1)), axis=-2) + grid = grid.reshape((n_blocks, -1, 2, 8)) + + return (db * grid * signs).reshape((n_blocks, -1)) + + +class IQ2_S(__Quant, qtype=GGMLQuantizationType.IQ2_S): + # iq2s_grid, but with each byte of the original packed in 2 bits, + # by mapping 0x08 to 0, 0x19 to 1, and 0x2b to 2. + grid_shape = (1024, 8) + grid_map = (0x08, 0x19, 0x2b) + grid_hex = ( + b"00000200050008000a0011001400160019002000220025002800410044004600" + b"490050005200550058006100640066006900800082008500880091009400a000" + b"a500aa0001010401060109011001120115011801210124014001420145014801" + b"510154015601590160016501680181018401900192019501a101a40100020202" + b"050208021102140220022a02410244024602490250025502800285028a029402" + b"a202010404040604090410041204150418042104240426042904400442044504" + b"48044a0451045404560459046004620465048104840486048904900495049804" + b"a104a40400050205050508050a05110514051605190520052505280541054405" + b"46054905500552055505580561056405800582058505880591059405a0050106" + b"0406060609061006150640064506480651065406600681068406900600080208" + b"050808081108140816081908200825082a084108440846084908500852085508" + b"580861086408800885089408aa08010904091009120915091809210940094509" + b"480951095409600981099009000a110a140a220a280a2a0a500a990a01100410" + b"0610091010101210151018102110241026104010421045104810511054105610" + b"59106010621065106810811084108610901095109810a110a410001102110511" + b"08110a1111111411161119112011221125112811411144114611491150115211" + b"5511581161116411801182118511881191119411011204120912101215122112" + b"2412401245125112541281128412901200140214051408141114141416141914" + b"2014251428144114441446144914501452145514581461146414801482148514" + b"881491149414a014011504150615091510151215151518152115241540154215" + b"4515481551155415601581158415901500160516081611161416201641164416" + b"50168016aa160118041806180918101815181818211840184218451848185118" + b"541860188118841800190219051908191119141920194119441950196919a219" + b"041a101a401a561a00200220052008201120142016201920202025202a204120" + b"4420502052205520642080208a209420aa200121042110211221152121214021" + b"4221452151215421602181218421902100220a22222228222a22442250228822" + b"8a22a82201240424062409241024152418242124242440244224452448245124" + b"5424602481248424902400250525082511251425202541254425502566258025" + b"0126042610264026592600280528112814284128442850288a28aa2801290429" + b"102995290a2a222a642a882a8a2a014004400640094010401240154018401a40" + b"21402440264040404240454048404a4051405440564059406040624065408140" + b"8440904095409840a140a4400041024105410841114114411641194120412241" + b"2541414144414641494150415241554158416141644180418241854188419141" + b"9441a04101420442104212421542184224424042454248425142544260428142" + b"844200440244054408440a441144144416441944204422442544284441444444" + b"46444944504452445544584461446444804482448544884491449444a0440145" + b"0445064509451045124515451845214524454045424545454845514554456045" + b"6a4581458445904500460246054608461146144620464146444650468046a546" + b"0148044809481048124815481848214824484048424845484848514854486048" + b"84489048004902490549084911491449204941494449504980499649014a044a" + b"104a404a00500250055008501150145016501950205022502550285041504450" + b"4650495050505250555058506150645080508250855088509150945001510451" + b"0651095110511251155118512151245140514251455148515151545160518151" + b"8451905100520552085211521452205241524452505269528052015404540654" + b"0954105412541554185421542454405442544554485451545454605481548454" + b"9054005502550555085511551455205541554455505580550156045610562656" + b"405600580258055808581158145820584158445850585a588058015904591059" + b"4059005a195a855aa85a01600460066010601260156018602160246040604560" + b"4860516054606060846090600061026105610861116114612061416144615061" + b"806199610462106240625662a162006405640864116414642064416444645064" + b"806401650465106540654a656865926500669466016804681068656898680069" + b"2a69426aa16a0080028005800880118014801980208025804180448050805280" + b"5580588061808080858091809480018104810981108112811581188121812481" + b"408142814581488151815481818184819081a981008205820a82118214824182" + b"4482508201840484068409841084128415841884218440844284458448845184" + b"5484608481848484908400850285058508851185148520854185448550858085" + b"8a85018604861086298640860088058811881488418844885088a28801890489" + b"40896589228a588a5a8a828aa28a019004900990109012901590189024904090" + b"4290459048905190549060908190849090900091059111911491419144915091" + b"5a910192049210924092a6920094029405940894119414942094419444945094" + b"8094969401950495109540959895a19500964696649601980498109826984098" + b"a998009949995299909a00a005a00aa014a022a02aa041a044a050a0a2a0aaa0" + b"40a165a102a20aa222a228a22aa282a288a28aa2a8a201a404a410a440a489a4" + b"a4a400a519a551a60aa828a8a2a854a986a908aa0aaa20aa22aa28aa88aaaaaa" + ) + + @classmethod + def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray: + n_blocks = blocks.shape[0] + + d, rest = np.hsplit(blocks, [2]) + qs, rest = np.hsplit(rest, [QK_K // 8]) + signs, rest = np.hsplit(rest, [QK_K // 8]) + qh, scales = np.hsplit(rest, [QK_K // 32]) + + d = d.view(np.float16).astype(np.float32) + + scales = scales.reshape((n_blocks, -1, 1)) >> np.array([0, 4], dtype=np.uint8).reshape((1, 1, 2)) + scales = (scales & 0x0F).reshape((n_blocks, -1)) + db = d * (np.float32(0.5) + scales) * np.float32(0.25) + db = db.reshape((n_blocks, -1, 1, 1)) + + # unpack the sign bits + signs = signs.reshape((n_blocks, -1, 1)) >> np.array([i for i in range(8)], dtype=np.uint8).reshape((1, 1, 8)) + signs = signs & np.uint8(0x01) + signs = np.where(signs == 0, np.float32(1), np.float32(-1)) + signs = signs.reshape((n_blocks, -1, 2, 8)) + + qh = qh.reshape((n_blocks, -1, 1)) >> np.array([0, 2, 4, 6], dtype=np.uint8).reshape((1, 1, 4)) + qs = qs.astype(np.uint16) | ((qh & 0x03).astype(np.uint16) << 8).reshape((n_blocks, -1)) + + assert cls.grid is not None + grid = np.take_along_axis(cls.grid, qs.reshape((n_blocks, -1, 1, 1)), axis=-2) + grid = grid.reshape((n_blocks, -1, 2, 8)) + + return (db * grid * signs).reshape((n_blocks, -1)) + + +class IQ3_XXS(__Quant, qtype=GGMLQuantizationType.IQ3_XXS): + grid_shape = (256, 4) + grid_map = (0x04, 0x0c, 0x14, 0x1c, 0x24, 0x2c, 0x34, 0x3e) + grid_hex = ( + b"0000020004001100130017002000220031004200730075000101030110011201" + b"2101250130013201410154017001000202020402110220022202310233023702" + b"5102570275020103070310031203250370031304370444045704730475040105" + b"0705320552053506640610071407160743076107011003101010121021102310" + b"3010321034104710501000110211111120112211011203121012121221123012" + b"7212001302132013311346136613011405145014201524154615711505162217" + b"4017002002201120132020202220262031204220012103210521102112212121" + b"3021632167217021002202221122172220222222372240225522012310231423" + b"7023742335245324032527254125742501270327162745270130103012302130" + b"2330503065307230003102312031313144314631013203321032253252327232" + b"1133333330344734723400350635223555351436363663363337603704401740" + b"3540374053405740744120423742404260426642074345430444514464442545" + b"4345704505471047124730471250415070500051065126515551145232527252" + b"0253535310542354275472540255315550562457425724604460466064602161" + b"6161176264623063366344640565526533660367216703700570077010703270" + b"5270267140711272457252720073157333736073217441740075027524753076" + ) + + @classmethod + def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray: + n_blocks = blocks.shape[0] + + d, rest = np.hsplit(blocks, [2]) + qs, scales = np.hsplit(rest, [QK_K // 4]) + + d = d.view(np.float16).astype(np.float32) + scales = scales.view(np.uint32) + + db = d * (np.float32(0.5) + (scales >> 28).astype(np.float32)) * np.float32(0.5) + db = db.reshape((n_blocks, -1, 1, 1)) + + # get the sign indices and unpack the bits + signs = scales.reshape((n_blocks, -1, 1)) >> np.array([0, 7, 14, 21], dtype=np.uint32).reshape((1, 1, 4)) + ksigns = np.frombuffer(IQ2_XXS.ksigns, dtype=np.uint8).reshape((1, 1, 1, 128)) + signs = (signs & np.uint32(0x7F)).reshape((n_blocks, -1, 4, 1)) + signs = np.take_along_axis(ksigns, signs, axis=-1) + signs = signs.reshape((n_blocks, -1, 4, 1)) >> np.array([i for i in range(8)], dtype=np.uint8).reshape((1, 1, 1, 8)) + signs = signs & np.uint8(0x01) + signs = np.where(signs == 0, np.float32(1), np.float32(-1)) + signs = signs.reshape((n_blocks, -1, 4, 8)) + + assert cls.grid is not None + grid = np.take_along_axis(cls.grid, qs.reshape((n_blocks, -1, 1, 1)), axis=-2) + grid = grid.reshape((n_blocks, -1, 4, 8)) + + return (db * grid * signs).reshape((n_blocks, -1)) + + +class IQ3_S(__Quant, qtype=GGMLQuantizationType.IQ3_S): + grid_shape = (512, 4) + grid_map = (0x01, 0x03, 0x05, 0x07, 0x09, 0x0b, 0x0d, 0x0f) + grid_hex = ( + b"0000010002000500070010001100120014001600200021002500330040004200" + b"4500470051005300600062007100740077000001010102010401100111011501" + b"2001230127013101350144016101650172010002010205020702100213021602" + b"2102250230023402420245024702510253027002730203031103150320032203" + b"3103330336034403500352036703710375030004130417042104240432044004" + b"4304510470040205040520052205260533054105450547056605730506061106" + b"1306310652067106000702070407200722072607330750075407001001100210" + b"0410101011101310151017102010221031103410361054105610611072100011" + b"0111031106111011141121113011331141115011521170117611001212121512" + b"1712201224123212401243125512601272120113041307131013131321132713" + b"3013341341136213701303140514121414143114331442144614501454140115" + b"1015131521153015321551152016241627164416461601170317101712172117" + b"3517411762177017002001200320052007201020122014201620212023202720" + b"3020322041204320452050205220672070207320752000210221102113211721" + b"2221252131213421422151210122042207222122232230223722412253225722" + b"7122742200230223052311232223242331233323422350236623012407242024" + b"2324322435244124722475240425112522253725402553257025002602260726" + b"2126552661260527112726273027432750270230113013301530173022303130" + b"3330353042304430473051306330713001310331053114312131233140316031" + b"7231763100321232203232323432503201331033143321332333273330334133" + b"4333473355337333033411341634223431345234603464340135103512352535" + b"3235443556357335163641360137033720372237353700400440124020402440" + b"2740324041405040704002410741114113412241304135414341514155410142" + b"0342104215422142334240425742624270420443114313432043224331433543" + b"0044024424443744404471440545074521456245134634466046104715473047" + b"4347514702501050145022504050445047505250665074500151035105511251" + b"2151325172510052115223523052365253520253075310532753445351536553" + b"7353015404542054325446541255265551555355425602570457225711601360" + b"1560316033606060006120612761646112623462426255626262706200631463" + b"2163406325644364626400650365346560650566406611671367007004700770" + b"2070227036704070547062700271117124714371457101720472107216722172" + b"3072517202733273357353730174057413742074507422754275027631760077" + ) + + @classmethod + def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray: + n_blocks = blocks.shape[0] + + d, rest = np.hsplit(blocks, [2]) + qs, rest = np.hsplit(rest, [QK_K // 4]) + qh, rest = np.hsplit(rest, [QK_K // 32]) + signs, scales = np.hsplit(rest, [QK_K // 8]) + + d = d.view(np.float16).astype(np.float32) + + scales = scales.reshape((n_blocks, -1, 1)) >> np.array([0, 4], dtype=np.uint8).reshape((1, 1, 2)) + scales = (scales & 0x0F).reshape((n_blocks, -1)) + db = d * (1 + 2 * scales) + db = db.reshape((n_blocks, -1, 1, 1)) + + # unpack the sign bits + signs = signs.reshape((n_blocks, -1, 1)) >> np.array([i for i in range(8)], dtype=np.uint8).reshape((1, 1, 8)) + signs = signs & np.uint8(0x01) + signs = np.where(signs == 0, np.float32(1), np.float32(-1)) + signs = signs.reshape((n_blocks, -1, 4, 8)) + + qh = qh.reshape((n_blocks, -1, 1)) >> np.array([i for i in range(8)], dtype=np.uint8) + qh = (qh & 0x01).astype(np.uint16).reshape((n_blocks, -1)) + qs = qs.astype(np.uint16) | (qh << 8) + + assert cls.grid is not None + grid = np.take_along_axis(cls.grid, qs.reshape((n_blocks, -1, 1, 1)), axis=-2) + grid = grid.reshape((n_blocks, -1, 4, 8)) + + return (db * grid * signs).reshape((n_blocks, -1)) + + +class IQ1_S(__Quant, qtype=GGMLQuantizationType.IQ1_S): + # iq1s_grid, with each byte packed into 2 bits + # -1, 0, 1 <=> 0, 1, 2 + grid_shape = (2048, 8) + grid_map = (-1, 0, 1) + grid_hex = ( + b"00000200050008000a00110015002000220028002a0045005100540056006500" + b"8000820088008a009500a000a200a800aa000401050111011401160119011a01" + b"2501410146014901520155015a0161016401660168018501910194019601a501" + b"0002020208020a0215022002220228022a024502510259026402690280028202" + b"88028a02910295029902a002a202a802aa021104140416042504410449045504" + b"5a046404650491049904a5040105040505050605150518051a05290540054505" + b"4a0550055105540555055605590560056205650568056a058105910595059805" + b"9a05a105a405a505a605a9051406190641064406500652065506580660066106" + b"6606690685069106940699060008020808080a0815082008220828082a084508" + b"5108560865088008820888088a089508a008a208a808aa080509110914091909" + b"2409250941095009510955096109640969099109940996099909a509000a020a" + b"080a0a0a150a200a220a280a2a0a450a510a590a610a650a800a820a850a880a" + b"8a0a950aa00aa20aa80aaa0a1010111014101910241025104110441050105510" + b"58106110641065106910911094109610a110a510011104110611091110111211" + b"1511181121112411291145114a11501151115211541155115611591160116511" + b"841192119511a111a41111121412161225124012461249125212551258125a12" + b"641266128512911294129612a512011406140914141415141814191421142614" + b"41144514461448144a1451145414551456145914621465146814841489149014" + b"94149514981499149a14a114a414a514a914021505150a151115141515151615" + b"191520152215251528152a154115441545154615511552155415551556155915" + b"5a1561156415651566156915801582158415851588158a159015911594159515" + b"961599159a15a015a215a51501160416051606161516161618161a1621162616" + b"401642164416451648164a165116551656165816591661166416651668166916" + b"6a1686168a1692169516a416a916111816182518411844184618491850185518" + b"58185a1860186118641866186918851891189418a5181019121915191a192119" + b"25194219441945194819511954195519561959195a19601965196a1989199119" + b"921995199819a119a619a919091a161a241a261a441a461a491a501a521a551a" + b"581a611a661a691a851a911a961a9a1a0020022008200a201520202022202520" + b"28202a20452051205920612065208020822088208a209520a020a220a520a820" + b"aa2005211121142119212521422144214921552158215a216121642165216621" + b"8521902196219921a521012208220a22112215222022222228222a2245225122" + b"562259226522812288228a2291229522a022a222a822aa220524142416241924" + b"252444244524462449245224552458245a2466248524912494249924a124a524" + b"0925152521252925402545254825512554255525592562256525682589259025" + b"9425952598259a25a125a425a625a92505261026122619262526412649265526" + b"6026612669268426862690269a260028022808280a2815282028222828282a28" + b"45285128542865288028822888288a28a028a228a828aa280929112914291929" + b"2529462949295229552961296429662969298529902996299929a429a529002a" + b"022a082a0a2a202a222a282a2a2a452a512a562a592a652a802a822a882a8a2a" + b"952aa02aa22aa82aaa2a054011401640254049405240554058405a4061406440" + b"664094409940a140a6400041014104410641094112411541164118411a412141" + b"26412941454148414a41514154415541564159415a41654168416a4181418441" + b"8641904192419541a041a141a241054211421442164225424142524255425a42" + b"6442694289429442a5420144154419442944454448444a445144544455445644" + b"61446244654468446a44814486448944904492449544a044a144a94401450245" + b"05450a4511451445154516451945204525452a45414544454545464549455045" + b"5145544555455645584559456145644565456645694582458445854588459145" + b"94459545964599459a45a545a845aa450146054609461446154618461a462146" + b"2446294640464246454648465046514652465546564659466246654668468146" + b"85468a4694469546a146a446a6460548114815481a4825484248494850485548" + b"5848614864486648694885489148944896489948a5480149054906490a491049" + b"144915491849214924492649404945494a495149524954495549564959496049" + b"6249654966496a49864989499249954996499849a149a449a649a949164a444a" + b"464a494a554a584a5a4a644a694a944aa54a0150045005500650095012501550" + b"1a50215024502950405045504850515054505550565059506550685086508950" + b"95509850a050a150a650a9500551085109510a51115114511551165118511951" + b"20512551265128512a5141514451455146514951505151515251545155515651" + b"585159515a51615164516551665169518251855191519451955196519951a051" + b"a551aa5101520652125215521a5221522452425245524a525152545255525652" + b"595262526552855290529252955299529a52a452045405541154145415541654" + b"185419542154255428542a54415444544554465449544a545054515454545554" + b"5654585459545a54615462546454655466546954805488548a54915494549554" + b"96549954a154a454a554aa540155025504550555065509551055115512551455" + b"1555165519551a55215524552555265529554055415542554455455546554855" + b"4955505551555255545555555655585559555a55605561556455655566556855" + b"69556a5581558455855589558a559055915594559555965598559955a155a455" + b"a555a655a9550056015602560456065608560956115614561556185619562056" + b"2156225624562556265628562956415645564656485649564a56505651565256" + b"545655565656585659565a566156645665566956825685568656885689568a56" + b"915695569a56a256a556a656a856a95604580558065809581058155818582158" + b"2a58455848584a58515854585558565858585958605862586458655882588958" + b"9058925895589858a158a9580159025905590a59115914591559165919592559" + b"41594459455946594959505951595259545955595659585959595a5961596459" + b"655966596959815985598959915994599559965998599959a559045a085a155a" + b"1a5a205a255a265a295a455a485a495a515a555a565a585a595a625a655a685a" + b"6a5a815a8a5a925a955a965a985a9a5aa15a0560146016601960256044605060" + b"5560566058605a60616064606660696081609660a56001610461066109611261" + b"15612161226126612961456149615161556156615961656166616a6184618a61" + b"92619561a161a661a96111621662196240624162466255625662586260628562" + b"91629662a56211641264156416641a6421642664296440644264456448644a64" + b"516454645564566459645a646064626465648464856489649064926494649564" + b"966498649a64a164a464a964056508650a651165156516651965446545654665" + b"496550655165546555655665596561656465656566656965866589658a659165" + b"9565966599659a65a265a565a665a86502660966156620662666286629664066" + b"456648664a66516654665566566658665a666066656668668066826685668a66" + b"9466966698669966a066a466a666aa661668196825684168526855685a686168" + b"6968856891689868a66801690469106915692169246926692969406941694569" + b"4669486951695469556956695969606965696a69826984698a699569a169a469" + b"a569a969116a166a186a416a446a496a506a556a586a5a6a646a656a696a866a" + b"946a986a9a6aa66a0080028008800a802080228028802a804580508051805480" + b"5680598065808080828088808a809580a080a280a880aa800581118114811681" + b"1981258141814481498150815281558156815881598164816681698185818981" + b"948196819981a5810082028208820a8215822082228228822a82518254825982" + b"65828082828288828a829582a082a282a882aa82148419844184448451845584" + b"5a846184648469849484998401850985128515851a8526852985408541854585" + b"4885518554855585568559855a856585668568856a8581858485868589859085" + b"928595859885a68511861686198625864186448649864a865086558659865a86" + b"618666866a86858691869a86a4860088028808880a8815882088228828882a88" + b"41884588518854885988658869888088828888888a889588a088a288a888aa88" + b"05890689118914891689258941894489468949895089528955895a8961896489" + b"858996899989a589008a028a088a0a8a158a208a228a288a2a8a458a518a548a" + b"568a808a828a888a8a8a958aa08aa28aa88aaa8a059011901690189019902590" + b"419046904990559058905a9069906a9085909190949096909990a59001910491" + b"069109911091159118911a912191249126912991409145915091519154915591" + b"569159916291659184918691929195919891a191a491a691a991059211921492" + b"19922592449246924992509252925592589266926992859294929692a9920194" + b"04940694109415941894269440944a9451945494559456945894599460946194" + b"62946594849486949294949495949894a194a9940095059508950a9510951195" + b"14951595169519952195259529952a9541954495459546954995509551955295" + b"549555955695589559955a956195649565956695699581958595889591959295" + b"94959595969599959a95a095a295a595a895aa95019604961096159619962096" + b"2696299645964896499651965296559656965996659668968296849689968a96" + b"929694969596a496a696a9960598169819982598419846985098529855985698" + b"5a98649865988598919896989998a59804990699099910991299159918991a99" + b"209921992499269940994299459948994a995199549955995699599962996599" + b"66996a99819984999099929995999a99a199a699059a159a259a449a469a499a" + b"509a559a589a619a859a919a949a959a969a00a002a008a00aa015a020a022a0" + b"28a02aa045a051a054a056a059a080a082a088a08aa095a0a0a0a2a0a8a0aaa0" + b"05a109a111a114a116a119a11aa146a149a151a155a158a15aa161a164a185a1" + b"90a192a196a199a102a208a20aa210a219a222a228a22aa245a251a256a259a2" + b"65a280a282a288a28aa295a2a0a2a2a2a8a2aaa219a425a441a444a450a454a4" + b"55a458a45aa461a465a466a468a469a485a406a509a510a512a515a518a526a5" + b"29a542a545a551a554a555a556a559a565a56aa581a584a585a586a589a592a5" + b"95a598a505a611a616a61aa621a625a644a646a64aa652a655a656a658a660a6" + b"62a686a690a695a696a699a6a1a6a4a6a6a600a802a808a80aa820a822a828a8" + b"2aa851a854a856a859a880a882a888a88aa895a8a0a8a2a8a8a8aaa805a914a9" + b"19a921a925a941a950a955a95aa961a966a969a990a996a900aa02aa08aa0aaa" + b"20aa22aa28aa2aaa51aa54aa56aa80aa82aa88aa8aaa95aaa0aaa2aaa8aaaaaa" + ) + + delta = np.float32(0.125) + + @classmethod + def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray: + n_blocks = blocks.shape[0] + + d, rest = np.hsplit(blocks, [2]) + qs, qh = np.hsplit(rest, [QK_K // 8]) + + d = d.view(np.float16).astype(np.float32) + qh = qh.view(np.uint16) + + dl = d * (2 * ((qh >> 12) & 7) + 1) + dl = dl.reshape((n_blocks, -1, 1, 1)) + delta = np.where((qh & np.uint16(0x8000)) == 0, cls.delta, -cls.delta) + delta = delta.reshape((n_blocks, -1, 1, 1)) + + qh = qh.reshape((n_blocks, -1, 1)) >> np.array([0, 3, 6, 9], dtype=np.uint16).reshape((1, 1, 4)) + qs = qs.astype(np.uint16) | ((qh & 7) << 8).reshape((n_blocks, -1)) + + assert cls.grid is not None + grid = np.take_along_axis(cls.grid, qs.reshape((n_blocks, -1, 1, 1)), axis=-2) + grid = grid.reshape((n_blocks, -1, 4, 8)) + + return (dl * (grid + delta)).reshape((n_blocks, -1)) + + +class IQ1_M(__Quant, qtype=GGMLQuantizationType.IQ1_M): + grid_shape = IQ1_S.grid_shape + grid_map = IQ1_S.grid_map + grid_hex = IQ1_S.grid_hex + + delta = IQ1_S.delta + + # Okay *this* type is weird. It's the only one which stores the f16 scales in multiple parts. + @classmethod + def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray: + n_blocks = blocks.shape[0] + + qs, rest = np.hsplit(blocks, [QK_K // 8]) + qh, scales = np.hsplit(rest, [QK_K // 16]) + + # The f16 scale is packed across multiple bytes + scales = scales.view(np.uint16) + d = (scales.reshape((n_blocks, 4)) & np.uint16(0xF000)) >> np.array([12, 8, 4, 0], dtype=np.uint16).reshape((1, 4)) + d = d[..., 0] | d[..., 1] | d[..., 2] | d[..., 3] + d = d.view(np.float16).astype(np.float32).reshape((n_blocks, 1)) + + scales = scales.reshape(n_blocks, -1, 1) >> np.array([0, 3, 6, 9], dtype=np.uint16).reshape((1, 1, 4)) + scales = (scales & 0x07).reshape((n_blocks, -1)) + dl = d * (2 * scales + 1) + dl = dl.reshape((n_blocks, -1, 2, 1, 1)) + + qh = qh.reshape((n_blocks, -1, 1)) >> np.array([0, 4], dtype=np.uint8).reshape((1, 1, 2)) + qs = qs.astype(np.uint16) | ((qh & 0x07).astype(np.uint16) << 8).reshape((n_blocks, -1)) + + delta = np.where(qh & 0x08 == 0, cls.delta, -cls.delta) + delta = delta.reshape((n_blocks, -1, 2, 2, 1)) + + assert cls.grid is not None + grid = np.take_along_axis(cls.grid, qs.reshape((n_blocks, -1, 1, 1)), axis=-2) + grid = grid.reshape((n_blocks, -1, 2, 2, 8)) + + return (dl * (grid + delta)).reshape((n_blocks, -1)) + + +class IQ4_NL(__Quant, qtype=GGMLQuantizationType.IQ4_NL): + kvalues = (-127, -104, -83, -65, -49, -35, -22, -10, 1, 13, 25, 38, 53, 69, 89, 113) + + @classmethod + def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray: + n_blocks = blocks.shape[0] + + d, qs = np.hsplit(blocks, [2]) + + d = d.view(np.float16).astype(np.float32) + + qs = qs.reshape((n_blocks, -1, 1, cls.block_size // 2)) >> np.array([0, 4], dtype=np.uint8).reshape((1, 1, 2, 1)) + + qs = (qs & np.uint8(0x0F)).reshape((n_blocks, -1, 1)) + + kvalues = np.array(cls.kvalues, dtype=np.int8).reshape(1, 1, 16) + qs = np.take_along_axis(kvalues, qs, axis=-1).astype(np.float32).reshape((n_blocks, -1)) + + return (d * qs) + + +class IQ4_XS(__Quant, qtype=GGMLQuantizationType.IQ4_XS): + @classmethod + def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray: + n_blocks = blocks.shape[0] + + d, rest = np.hsplit(blocks, [2]) + scales_h, rest = np.hsplit(rest, [2]) + scales_l, qs = np.hsplit(rest, [QK_K // 64]) + + d = d.view(np.float16).astype(np.float32) + scales_h = scales_h.view(np.uint16) + + scales_l = scales_l.reshape((n_blocks, -1, 1)) >> np.array([0, 4], dtype=np.uint8).reshape((1, 1, 2)) + scales_h = scales_h.reshape((n_blocks, 1, -1)) >> np.array([2 * i for i in range(QK_K // 32)], dtype=np.uint16).reshape((1, -1, 1)) + scales_l = scales_l.reshape((n_blocks, -1)) & np.uint8(0x0F) + scales_h = scales_h.reshape((n_blocks, -1)).astype(np.uint8) & np.uint8(0x03) + + scales = (scales_l | (scales_h << np.uint8(4))).astype(np.int8) - np.int8(32) + dl = (d * scales.astype(np.float32)).reshape((n_blocks, -1, 1)) + + qs = qs.reshape((n_blocks, -1, 1, 16)) >> np.array([0, 4], dtype=np.uint8).reshape((1, 1, 2, 1)) + qs = qs.reshape((n_blocks, -1, 32, 1)) & np.uint8(0x0F) + + kvalues = np.array(IQ4_NL.kvalues, dtype=np.int8).reshape((1, 1, 1, -1)) + qs = np.take_along_axis(kvalues, qs, axis=-1).astype(np.float32).reshape((n_blocks, -1, 32)) + + return (dl * qs).reshape((n_blocks, -1)) diff --git a/python/freetoken/tokenizer/detokenize.py b/python/freetoken/tokenizer/detokenize.py index d138ac1a68..44c3149948 100644 --- a/python/freetoken/tokenizer/detokenize.py +++ b/python/freetoken/tokenizer/detokenize.py @@ -1,5 +1,7 @@ -from dataclasses import dataclass -from typing import Dict, FrozenSet, List +import json +import weakref +from dataclasses import dataclass, field +from typing import Any, Dict, FrozenSet, List from freetoken.message import DetokenizeMsg from transformers import PreTrainedTokenizerBase @@ -71,11 +73,74 @@ 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 + in_reasoning: bool = False # the prompt or the output opened a reasoning block + generated: int = 0 # sampled tokens, EOS included + text_end: int = 0 # decoded text so far, before holdback: where the next token starts + # (text start, entry) of logprob entries whose token text is still held for a stop string + held_logprobs: List[tuple] = field(default_factory=list) + + def reasoning_tokens(self) -> int: + if self.think_end_at is not None: + return self.think_end_at + # stopped inside the block (max_tokens, or EOS before the end tag): all of it + return self.generated if self.in_reasoning else 0 + + def release_logprobs( + self, entry: dict | None, token_start: int, emit_end: int, text_len: int, flush: bool + ) -> List[dict]: + """Hold this token's logprob entry and release those whose text is now sent.""" + if entry is not None: + self.held_logprobs.append((token_start, {**entry, "text_offset": token_start})) + out: List[dict] = [] + # A held token's text ends where the next one's starts. + starts = [start for start, _ in self.held_logprobs] + ends = starts[1:] + [text_len] if starts else [] + if flush: + trimmed = emit_end < text_len + for (start, held), end in zip(self.held_logprobs, ends, strict=True): + if trimmed and start >= emit_end: + held = {**held, "stop_trimmed": True} + elif trimmed and emit_end < end: + held = {**held, "stop_visible": emit_end - start} + out.append(held) + self.held_logprobs = [] + elif self.held_logprobs: + # Wait for all of a token's text, so a stop match inside it can still mark it. + n = 0 + while n < len(ends) and starts[n] < emit_end and ends[n] <= emit_end: + n += 1 + out = [held for _, held in self.held_logprobs[:n]] + self.held_logprobs = self.held_logprobs[n:] + return out + + 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, + think_start_id: int | None = None, ) -> None: # uid -> DecodeStatus self.decode_map: Dict[int, DecodeStatus] = {} @@ -83,6 +148,10 @@ 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) + self.think_start_id = think_start_id if think_start_id is None else int(think_start_id) def discard(self, uid: int) -> None: """Drop a uid's decode state without a finished DetokenizeMsg. An aborted or errored @@ -91,8 +160,59 @@ 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], List[int]]: + """Incremental text per message; for a finished message, its reasoning token count + (0 otherwise or when the generation never was in a reasoning block); and where each + message's token starts in the generated text. That offset ignores stop-string + holdback, which delays when text is sent but not where it lands.""" + texts, counts, _, offsets = self._detokenize(msgs, [None] * len(msgs)) + return texts, counts, offsets + + def detokenize_with_logprobs( + self, msgs: List[DetokenizeMsg], entries: List[dict | None] + ) -> tuple[List[str], List[int], List[List[dict]]]: + """``detokenize_with_meta`` plus the logprob entries each reply carries, each stamped + with its token's ``text_offset``. With stop strings, an entry waits for the reply that + releases all its token's text, so it never precedes text that may still be trimmed; + entries of tokens a matched stop string trimmed come on the finished reply marked + ``stop_trimmed``, and that of a token it cut inside with ``stop_visible``, the length + of the token's text left visible.""" + texts, counts, released, _ = self._detokenize(msgs, entries) + return texts, counts, released + + def _detokenize( + self, msgs: List[DetokenizeMsg], entries: List[dict | None] + ) -> tuple[List[str], List[int], List[List[dict]], List[int]]: + # 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([msg], [entry]) for msg, entry in zip(msgs, entries, strict=True) + ] + return ( + [p[0][0] for p in parts], + [p[1][0] for p in parts], + [p[2][0] for p in parts], + [p[3][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 +221,55 @@ 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)), + in_reasoning=bool(getattr(msg, "starts_in_reasoning", False)), ) s = self.decode_map[msg.uid] - if not (msg.finished and msg.next_token in self.eos_token_ids): + s.generated += 1 + 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) + if s.think_end_at is None and msg.next_token == self.think_start_id: + s.in_reasoning = True 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] = [] - for msg, read_str, surr_str in zip(msgs, read_texts, surr_texts, strict=True): + reasoning_tokens: List[int] = [] + released: List[List[dict]] = [] + text_offsets: List[int] = [] + for msg, entry, read_str, surr_str in zip( + msgs, entries, 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 + token_start = s.text_end + text_offsets.append(token_start) + output_str = s.update(read_str, surr_str) + s.text_end = max(s.text_end, len(output_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 +280,115 @@ 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.reasoning_tokens() if msg.finished else 0) + if entry is not None or s.held_logprobs: + flush = msg.finished or not msg.stop_strs + out = s.release_logprobs(entry, token_start, emit_end, len(output_str), flush) + else: + out = [] # no logprobs asked for: skip the hold bookkeeping + released.append(out) if msg.finished: del self.decode_map[msg.uid] - return incremental_strs + return incremental_strs, reasoning_tokens, released, text_offsets + + +def _gpt2_byte_decoder() -> dict[str, int]: + """Inverse of GPT-2's ``bytes_to_unicode``: the printable stand-in char -> its byte.""" + bs = [*range(ord("!"), ord("~") + 1), *range(ord("¡"), ord("¬") + 1), *range(ord("®"), ord("ÿ") + 1)] + cs = bs[:] + n = 0 + for b in range(256): + if b not in bs: + bs.append(b) + cs.append(256 + n) + n += 1 + return {chr(c): b for b, c in zip(bs, cs, strict=True)} + + +class _TokenBytes: + """A token id's own bytes. A single-id decode cannot give them for a piece of a + multi-byte character (it decodes to U+FFFD), so read the vocab piece instead.""" + + def __init__(self, tokenizer: Any) -> None: + decoder: dict = {} + backend = getattr(tokenizer, "backend_tokenizer", None) + try: + decoder = json.loads(backend.decoder.__getstate__()) + except Exception: # noqa: BLE001 -- no fast backend: fall back to the decoded text + pass + kinds = {d.get("type") for d in decoder.get("decoders", [decoder])} + self.byte_level = _gpt2_byte_decoder() if "ByteLevel" in kinds else None + self.byte_fallback = "ByteFallback" in kinds + added = getattr(tokenizer, "added_tokens_decoder", None) or {} + self.added = {int(tid): str(getattr(tok, "content", tok)) for tid, tok in added.items()} + self.convert = getattr(tokenizer, "convert_ids_to_tokens", None) + + def __call__(self, tid: int, text: str) -> bytes: + if tid in self.added: + # the GPT-2 table cannot invert an added token's raw (e.g. full-width) chars + return self.added[tid].encode("utf-8") + try: + piece = self.convert(tid) if self.convert is not None else None + except Exception: # noqa: BLE001 -- out-of-vocab id + piece = None + if isinstance(piece, str): + if self.byte_level is not None: + try: + return bytes(self.byte_level[c] for c in piece) + except KeyError: + pass + elif self.byte_fallback and len(piece) == 6 and piece.startswith("<0x") and piece.endswith(">"): + try: + return bytes([int(piece[3:5], 16)]) + except ValueError: + pass + return text.encode("utf-8") + + +_TOKEN_BYTES: "weakref.WeakKeyDictionary[Any, _TokenBytes]" = weakref.WeakKeyDictionary() + + +def _token_bytes(tokenizer: Any) -> _TokenBytes: + try: + cached = _TOKEN_BYTES.get(tokenizer) + if cached is None: + cached = _TOKEN_BYTES[tokenizer] = _TokenBytes(tokenizer) + except TypeError: # not weak-referenceable or unhashable + return _TokenBytes(tokenizer) + return cached + + +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, and ``bytes`` are the token's own bytes, so clients can reassemble a + character split across tokens. Ids past the real vocab (the padded tail of the + logits) report an empty token.""" + token_bytes = _token_bytes(tokenizer) + + def _fields(tid: int) -> tuple[str, list[int]]: + try: + text = tokenizer.decode([tid]) + except Exception: # noqa: BLE001 -- a padded id some tokenizers refuse to decode + text = "" + return text, list(token_bytes(tid, text)) + + token, data = _fields(token_id) + top = [] + for tid, logprob in zip(top_ids or [], top_logprobs or []): + text, cand = _fields(int(tid)) + top.append({"token_id": int(tid), "token": text, "bytes": cand, "logprob": float(logprob)}) + return { + "token_id": token_id, + "token": token, + "bytes": data, + "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..43c5515761 100644 --- a/python/freetoken/tokenizer/server.py +++ b/python/freetoken/tokenizer/server.py @@ -81,10 +81,37 @@ 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.""" + return _vocab_tag_id(tokenizer, ("", "<|end_of_thinking|>")) + + +def _vocab_tag_id(tokenizer, tags: tuple[str, ...]) -> int | None: + for tag in tags: + 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 _prompt_opens_reasoning(input_ids: torch.Tensor, start_id: int, end_id: int) -> bool: + """The prompt's last reasoning tag is an opening one: the template pre-opened the + block (qwen3 with thinking on), so generation starts inside it. A raw prompt that just + mentions a literal trips this too; harmless, since only the chat and responses + endpoints, whose prompts are always templated, report reasoning_tokens.""" + tags = torch.nonzero((input_ids == start_id) | (input_ids == end_id)).flatten() + return tags.numel() > 0 and int(input_ids[tags[-1]]) == start_id + + def _tokenize_requests( tokenize_manager: Any, messages: List[TokenizeMsg], logger: Any, + reasoning_tags: tuple[int, int] | None = None, ) -> tuple[List[UserMsg], List[UserReply]]: """Tokenize independently, returning backend work plus terminal frontend errors. @@ -119,6 +146,10 @@ def _tokenize_requests( ) ) continue + if reasoning_tags is not None: + user_msg.sampling_params.starts_in_reasoning = _prompt_opens_reasoning( + user_msg.input_ids, *reasoning_tags + ) backend.append(user_msg) return backend, errors @@ -146,12 +177,22 @@ 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)) + think_start_id = _vocab_tag_id(tokenizer, ("",)) + think_end_id = _reasoning_end_id(tokenizer) + reasoning_tags = ( + (think_start_id, think_end_id) + if think_start_id is not None and think_end_id is not None + else None + ) detokenize_manager = DetokenizeManager( - tokenizer, load_eos_token_ids(tokenizer_path, tokenizer) + tokenizer, + load_eos_token_ids(tokenizer_path, tokenizer), + think_end_id=think_end_id, + think_start_id=think_start_id, ) if ack_queue is not None: @@ -209,7 +250,21 @@ def tokenize_worker( ) sampled_replies: List[UserReply] = [] if len(detokenize_msg) > 0: - replies = detokenize_manager.detokenize(detokenize_msg) + entries = [ + 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 in detokenize_msg + ] + replies, reasoning_counts, released = detokenize_manager.detokenize_with_logprobs( + detokenize_msg, entries + ) sampled_replies = [ UserReply( uid=msg.uid, @@ -217,6 +272,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 +281,12 @@ 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, + logprobs=reply_logprobs or None, + ) + for msg, reply, reasoning_tokens, reply_logprobs in zip( + detokenize_msg, replies, reasoning_counts, released, 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 @@ -247,7 +307,9 @@ def tokenize_worker( # Tokenize per-message so a single un-renderable request (e.g. a chat template # that rejects the message layout) becomes a terminal error reply for THAT uid # instead of an uncaught exception that kills the worker and bricks the server. - backend, errors = _tokenize_requests(tokenize_manager, tokenize_msg, logger) + backend, errors = _tokenize_requests( + tokenize_manager, tokenize_msg, logger, reasoning_tags + ) if errors: send_frontend.put( errors[0] if len(errors) == 1 else BatchFrontendMsg(data=errors) diff --git a/python/freetoken/tokenizer/tokenize.py b/python/freetoken/tokenizer/tokenize.py index 758ffd4dcf..dc7144d574 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,40 @@ 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 + + +def _ids_below(ids: List[int], bound: int) -> bool: + return min(ids) >= 0 and max(ids) < bound + + class TokenizeManager: def __init__(self, tokenizer: PreTrainedTokenizerBase, mm_processor: MMProcessor | None = None) -> None: self.tokenizer = tokenizer @@ -57,25 +99,30 @@ 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] = [] # TODO: batch tokenization for msg in msgs: - prompt = self.render_prompt(msg) - # A jinja chat template owns every special token (HF's apply_chat_template - # tokenizes with add_special_tokens=False for the same reason): tokenizers - # that auto-add bos (muse-glimmer's, llama's) would otherwise double it -- - # the template already rendered one. Raw-string prompts and the dsv4 - # encoder path keep the default. - templated = isinstance(msg.text, list) and self._dsv4_encoder is None - input_ids: torch.Tensor = ( # type: ignore - self.tokenizer.encode( - prompt, return_tensors="pt", add_special_tokens=not templated - ) - ) - input_ids = input_ids.view(-1).to(torch.int32) + # past the real vocabulary an id indexes the padded logits tail or beyond them + for name, ids in msg.sampling_params.client_token_ids(): + if not _ids_below(ids, len(self.tokenizer)): + raise ValueError(f"{name} must be in [0, {len(self.tokenizer)})") + if isinstance(msg.input_ids, torch.Tensor) and not msg.images: + # the frontend's own encoding of this prompt (frontend_prompt_ids) + input_ids = msg.input_ids.view(-1).to(torch.int32) + elif msg.input_ids is not None: + # an id past the embedding rows would fault the engine, not just this request + if msg.input_ids and not _ids_below(msg.input_ids, len(self.tokenizer)): + raise ValueError(f"prompt token ids must be in [0, {len(self.tokenizer)})") + ids = torch.tensor(msg.input_ids, dtype=torch.int32) + results.append(UserMsg(uid=msg.uid, input_ids=ids, sampling_params=msg.sampling_params)) + continue + else: + input_ids = self.encode_prompt(msg, self.render_prompt(msg)) if msg.images: if self.mm_processor is None: raise ValueError("image input is not supported for this model") @@ -96,6 +143,20 @@ def tokenize(self, msgs: List[TokenizeMsg]) -> List[UserMsg]: ) return results + def encode_prompt(self, msg: TokenizeMsg, prompt: str) -> torch.Tensor: + """The tokenizer half of ``tokenize``: ``prompt`` (from ``render_prompt``) as 1-D + int32 ids, exposed so the frontend can count a request before committing a stream.""" + # A jinja chat template owns every special token (HF's apply_chat_template + # tokenizes with add_special_tokens=False for the same reason): tokenizers + # that auto-add bos (muse-glimmer's, llama's) would otherwise double it -- + # the template already rendered one. Raw-string prompts and the dsv4 + # encoder path keep the default. + templated = isinstance(msg.text, list) and self._dsv4_encoder is None + input_ids: torch.Tensor = self.tokenizer.encode( # type: ignore + prompt, return_tensors="pt", add_special_tokens=not templated + ) + return input_ids.view(-1).to(torch.int32) + def render_prompt(self, msg: TokenizeMsg) -> str: """The template/encoder half of ``tokenize``, exposed so the frontend can validate a request before committing an SSE stream. Sanitizes @@ -103,21 +164,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 +213,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 +287,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 +309,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 2e4ad15f2f..a1dd10b7c3 100644 --- a/python/freetoken/utils/__init__.py +++ b/python/freetoken/utils/__init__.py @@ -1,5 +1,7 @@ from .arch import ( is_arch_supported, + is_rocm, + get_rocm_gfx_arch, is_sm90_family, is_sm90_supported, is_sm100_family, @@ -7,6 +9,7 @@ ) from .hf import ( cached_load_hf_config, + download_hf_checkpoint, download_hf_weight, load_eos_token_ids, load_generation_sampling, @@ -28,6 +31,7 @@ __all__ = [ "cached_load_hf_config", + "download_hf_checkpoint", "download_hf_weight", "load_eos_token_ids", "load_generation_sampling", @@ -35,6 +39,8 @@ "load_toolcall_anchor_id", "init_logger", "is_arch_supported", + "is_rocm", + "get_rocm_gfx_arch", "is_sm90_family", "is_sm90_supported", "is_sm100_family", diff --git a/python/freetoken/utils/arch.py b/python/freetoken/utils/arch.py index 8c1c6c3d56..304505966c 100644 --- a/python/freetoken/utils/arch.py +++ b/python/freetoken/utils/arch.py @@ -1,14 +1,66 @@ from __future__ import annotations import functools +import os +import re from typing import Tuple +_GFX_ARCH_RE = re.compile(r"gfx\d+[a-z]?") + + +def _gfx_arch_from(value: object) -> str | None: + match = _GFX_ARCH_RE.search(str(value).lower()) + return match.group(0) if match else None + + +@functools.cache +def is_rocm() -> bool: + """True when torch is built for ROCm (AMD GPU) instead of CUDA.""" + import torch + return getattr(torch.version, "hip", None) is not None + + +@functools.cache +def get_rocm_gfx_arch() -> str | None: + """Return the current AMD GPU target (for example ``gfx1201``). + + Prefer the runtime device because build variables may contain multiple + semicolon-separated targets. Environment variables remain useful for + cross-compilation and systems where no GPU is currently visible; in that + fallback mode the first target is returned. The result is process-cached + for FreeToken's one-process-per-GPU execution model, so callers must select + the intended device before the first call. + """ + if not is_rocm(): + return None + + import torch + + if torch.cuda.is_available(): + try: + props = torch.cuda.get_device_properties(torch.cuda.current_device()) + for attr in ("gcnArchName", "arch"): + arch = _gfx_arch_from(getattr(props, attr, "")) + if arch: + return arch + except (AttributeError, RuntimeError): + pass + + for env_var in ("FREETOKEN_ROCM_ARCH", "PYTORCH_ROCM_ARCH", "HCC_AMDGPU_TARGET"): + arch = _gfx_arch_from(os.getenv(env_var, "")) + if arch: + return arch + return None + + @functools.cache def _get_torch_cuda_version() -> Tuple[int, int] | None: import torch import torch.version + if is_rocm(): + return None if not torch.cuda.is_available() or not torch.version.cuda: return None return torch.cuda.get_device_capability() diff --git a/python/freetoken/utils/hf.py b/python/freetoken/utils/hf.py index b93d3b1203..ce2b574a26 100644 --- a/python/freetoken/utils/hf.py +++ b/python/freetoken/utils/hf.py @@ -1,9 +1,9 @@ +import copy import functools import json import os -from typing import Any - -from typing import FrozenSet +import re +from typing import Any, FrozenSet, Mapping from huggingface_hub import hf_hub_download, snapshot_download from huggingface_hub.utils import EntryNotFoundError @@ -63,7 +63,7 @@ def load_eos_token_ids( if (gguf_src := gguf_config_source(model_path)) is not None: from freetoken.models.gguf.tokenizer import gguf_eos_token_ids - return frozenset(gguf_eos_token_ids(gguf_src, tokenizer)) + return frozenset(gguf_eos_token_ids(gguf_src)) ids: set[int] = set() if tokenizer.eos_token_id is not None: @@ -235,20 +235,40 @@ def _load_hf_config(model_path: str) -> Any: return config -def cached_load_hf_config(model_path: str) -> PretrainedConfig: +def _with_overrides(config: Any, data: dict, overrides: Mapping[str, Any]) -> dict: + merged = dict(data) + for key, value in overrides.items(): + section = getattr(config, key, None) + if isinstance(value, Mapping) and isinstance(section, (PretrainedConfig, RawConfigShim)): + merged[key] = _with_overrides(section, merged.get(key) or {}, value) + else: + merged[key] = copy.deepcopy(value) + return merged + + +def cached_load_hf_config( + model_path: str, overrides: Mapping[str, Any] | None = None +) -> PretrainedConfig: + """A fresh copy of the checkpoint's config with ``overrides`` applied as vLLM's --hf-overrides + does: nested config sections update key by key, every other value is replaced whole.""" # A .gguf file (or an FTW dir converted from one) carries its own metadata (no HF # config.json); return a shim the model registry dispatches on instead of a # PretrainedConfig. from freetoken.models.gguf.reader import gguf_config_source if (gguf_src := gguf_config_source(model_path)) is not None: + if overrides: + raise ValueError("--hf-overrides applies to a HF config.json; a GGUF carries its own metadata") from freetoken.models.gguf.config import build_gguf_shim return build_gguf_shim(gguf_src) config = _load_hf_config(model_path) + data = config.to_dict() + if overrides: + data = _with_overrides(config, data, overrides) if isinstance(config, RawConfigShim): - return RawConfigShim(config.to_dict()) - return type(config)(**config.to_dict()) + return RawConfigShim(data) + return type(config)(**data) def _weight_allow_patterns(repo_id: str) -> list[str]: @@ -277,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/checkpoint_diagnose/README.md b/scripts/checkpoint_diagnose/README.md new file mode 100644 index 0000000000..578952889e --- /dev/null +++ b/scripts/checkpoint_diagnose/README.md @@ -0,0 +1,177 @@ +# checkpoint_diagnose + +Find out why FreeToken does not load a checkpoint, sort the failure into one kind, and repair the checkpoints whose own files are non-standard. +Run everything from the repo root with `PYTHONPATH=python`. Nothing here needs a GPU or the checkpoint's weights. + +> [!WARNING] +> The diagnosis is heuristic and can be wrong: the failure kind, the findings and the suggested fixes are a best guess from the headers and configs, not a proof. Check them against the problems it prints before acting on them. + +> [!IMPORTANT] +> For agents: explain the diagnosis to the user in plain terms first: what fails, which kind it is, and what the fix would change. Ask the user before you run a fix script or open an issue. Do not change the user's checkpoint on your own. + +| Script | Does | +|---|---| +| `diagnose.py` | checks a checkpoint dir or HF repo id, prints the failure kind and the next step | +| `fix_config.py` | level 1 fix: edits only `config.json` / `hf_quant_config.json` | +| `rename_tensors.py` | level 2 fix: renames tensors (shard headers, index, quant skip lists); tensor data is kept byte for byte | +| `fetch_headers.py` | header-only stub of an HF repo; `diagnose.py` calls it for a repo id | +| `ckpt_edit.py` | the edit primitives both fix scripts use, for a case they do not cover | + +## Step 1: which format + +`diagnose.py` routes by itself, but know where each format goes: + +| The checkpoint has | Go to | +|---|---| +| `*.gguf` files | `scripts/gguf_check/` (its README), then the GGUF section below | +| `freetoken_weight.json` (FTW, made by `ft checkpoint`) | diagnose the HF checkpoint it was converted from; FTW repairs are in `docs/ftw-hotfix.md` | +| root `*.safetensors` shards | this directory, from step 2 | +| only `*.bin` / `*.pt`, or MLX weights | not served: FreeToken reads safetensors and GGUF | + +### GGUF: three kinds of failure + +`scripts/gguf_check/check.py` prints the problems of each file; the first one tells the kind. + +| Kind | check.py says | Do | +|---|---|---| +| a ggml type has no kernel | ` : no kernel` (needs `--kernels` from `probe_kernels.py`) | the build or GPU lacks kernels for that quant type: pick a file with another type, or report it with the `probe_kernels.py` table | +| the model's loader does not read the file | `unexpected tensor ...`, `N tensors never loaded`, `tensor ... yielded twice`, `: got ..., model has ...`, `MoE layer N ... bank is ..., the model needs ...`, or a `NotImplementedError` from the family's config parse | a FreeToken gap for that layout: report it with the check.py record | +| the GGUF file itself | `is not a GGUF file`, `duplicate GGUF key / tensor`, `runs past the end of the file`, `bad general.alignment`, `tokenizer: no EOS token`, `tokenizer: no chat template`, `chat template: ` | the file is broken or incomplete: download it again, or use another file of the repo | + +A file whose architecture FreeToken does not serve shows as `skip`, not as a failure. + +## Step 2: diagnose a safetensors checkpoint + +```bash +python scripts/checkpoint_diagnose/diagnose.py --out report.jsonl +``` + +- A local dir is read as is: the headers and JSON files, so a 400 GB checkpoint takes seconds. A weight reader that needs tensor values gets the real bytes it asks for. +- A repo id is fetched as a header-only stub into `~/.cache/freetoken/checkpoint-stubs/` (a few MB). Nothing is downloaded beyond that, and no `.py` file: a repo whose config is its own code (`auto_map`) is skipped, diagnose a local download of it instead. +- A remote repo is only diagnosed: its record has no fixes, and the fix scripts refuse a stub. When it says `fixable`, download the checkpoint and diagnose the local directory to get the fix commands. +- Add `--text-model-only` when the checkpoint will be served with `--text-model-only`. +- `--out` appends one JSON record per checkpoint. The fix scripts read their fixes from it with `--from-report`. + +It first checks the files themselves (data offsets cover each shard exactly, the index agrees with the shard headers), then runs the server's own load path on the CPU with meta tensors: config parse, model build, dense weight load through the model's `load_state_dict`, and expert-bank fill for MoE models with the reader the engine picks. +When that fails, it also checks the tensor names against transformers' own model class and its checkpoint conversion rules, and lints the files for the known non-standard forms. + +Output, per checkpoint: `PASS`, or `FAIL at : [] ...` with the problems, what transformers makes of the names, the findings, and the next step. + +## Step 3: the kind of failure + +Each failure gets one kind. `diagnose.py` tries, in this order: broken files (`checkpoint`), `arch`, `quant_method`, a broken prefix (`checkpoint`), `quant_model`, names transformers cannot map or another non-standard finding (`checkpoint`), `packing`; anything left is `loader`. + +| Kind | Means | Do | +|---|---|---| +| `arch` | the architecture is not registered in `python/freetoken/models/register.py` | not served; a new architecture starts as a feature request | +| `quant_method` | FreeToken has no reader for the quantization format (GPTQ, AWQ, MLX, bitsandbytes, an unsupported compressed-tensors scheme, ...) | not served; use a release in a format `python/freetoken/layers/quantization/configs/` reads | +| `quant_model` | the format is read, but this model or layer kind has no method for it (`no quant method for moe.`), or the reader does not produce the scales the method needs | a FreeToken gap: report it | +| `loader` | the model and format are served, and transformers loads these names, but FreeToken's reader does not | a FreeToken gap: report it. A finding marked `loader` may offer a fix that works around it | +| `packing` | the expert banks do not fill, and transformers repacks these expert tensors (for example one tensor per expert where the family's release stacks them) | a FreeToken gap: report it. Repacking tensor data is not offered here | +| `checkpoint` | the checkpoint's own files are broken or non-standard: an incomplete shard or index, names transformers would not load either, or a config that contradicts what the shards store | fix it, step 4 | + +To report a FreeToken gap, use the "model checkpoint" issue template and paste the `diagnose.py` output. + +How to tell `loader` / `packing` from `checkpoint`: the line `transformers ... would not load it either` means transformers' model class and conversion rules cannot map the names either, so the checkpoint is at fault. `transformers loads these names` means the names are a form transformers accepts, so FreeToken should read them too. + +## Step 4: fix a non-standard checkpoint + +Fixes apply to a local copy of the checkpoint only, never to a remote repo or its stub. + +Rules: + +1. Fix into a new directory (`--out `). Unchanged files there are symlinks to the original, so it costs only the rewritten shards. Use `--in-place` only when the user asks for it. It keeps a backup, and `--restore` undoes the latest in-place edit. +2. One level per round. If the record has a `rename` fix, apply that first, then run `diagnose.py` on the new directory, then apply the config fixes of that new record. A rename changes the names the config findings talk about. +3. Read the plan before applying it. Every run prints the plan first and changes nothing until it is confirmed. Check that: + - each rule renames the tensors you expect, and the example `old -> new` lines are right; + - nothing is renamed that should not be. A `warning: ... matches are inside the name` line means the pattern is not anchored; + - the config edits name the right modules; + - the shard costs are what you expect. `rewritten, copies N GiB` needs that much free disk. + + If anything in the plan is not what you expected, do not apply it: change the rule or the edit and plan again. +4. Confirming: in a terminal the script asks `[y/N]`. Without a terminal (an agent), pass `--plan `. The first run saves the plan there and changes nothing. Read it, show it to the user, and run the exact same command again to apply it only after the user agrees. If the command or the checkpoint changed in between, the second run saves the new plan instead of applying anything. +5. After every fix, run `diagnose.py` on the result, until it passes or nothing fixable is left. + +The usual round, for an agent: + +```bash +D=scripts/checkpoint_diagnose +python $D/diagnose.py /models/x --out r0.jsonl +python $D/rename_tensors.py /models/x --out /models/x-fixed --from-report r0.jsonl --plan plan.json # prints and saves the plan +python $D/rename_tensors.py /models/x --out /models/x-fixed --from-report r0.jsonl --plan plan.json # after the user agrees to the plan: applies it +python $D/diagnose.py /models/x-fixed --out r1.jsonl +python $D/fix_config.py /models/x-fixed --in-place --from-report r1.jsonl --plan plan.json # twice, as above +python $D/diagnose.py /models/x-fixed +``` + +`--from-report` applies the fixes `diagnose.py` found for that checkpoint; anything else is a manual edit, below. + +### Fix levels + +| Level | Script | Changes | Typical use | +|---|---|---|---| +| 1 | `fix_config.py` | `config.json`, `hf_quant_config.json` | the quantization skip lists, compressed-tensors targets, a config key | +| 2 | `rename_tensors.py` | shard headers, `model.safetensors.index.json`, and the skip-list entries / targets that are module names or globs (a `re:` entry is not rewritten; the plan warns about every entry left naming the old modules) | broken prefixes, legacy tensor names | +| - | not offered | tensor data | repacking experts, dropping tensors, dequantizing: re-export the checkpoint, or report the FreeToken gap | + +### Manual edits + +When no automatic fix fits, write the edit yourself. It goes through the same plan and confirmation. + +```bash +# rename: re.sub(PATTERN, REPL, name), rules in order; anchor with ^ and escape dots +python $D/rename_tensors.py /models/x --out /models/x-fixed --rule '^model\.language_model\.visual\.' 'model.visual.' + +# config key: dotted path, JSON value +python $D/fix_config.py /models/x --out /models/x-fixed --set text_config.tie_word_embeddings false + +# quantization skip lists, in the dialect's syntax: +# ModelOpt: glob (model.language_model.layers.*.mlp.gate) +# compressed-tensors: module name or re: (re:^model\.visual\..*) +# fp8: module name (model.language_model.layers.0.linear_attn.in_proj_a) +python $D/fix_config.py /models/x --out /models/x-fixed --skip-add 'model.language_model.layers.*.mlp.gate' +python $D/fix_config.py /models/x --out /models/x-fixed --skip-remove 'lm_head' +``` + +A rule that renames nothing, and a rename that gives two tensors the same name, are refused. So is an `--out` directory that already has files in it. + +### Known non-standard forms + +The findings `diagnose.py` prints, and what they fix. `kind` is the failure kind the finding points to. + +| Finding | The checkpoint | Level | Fix | +|---|---|---|---| +| `prefix` | nests the text tower more than once (`model.language_model.language_model.`) | 2 | rename to `model.language_model.` | +| `prefix` | stores the vision tower under the text tower (`model.language_model.visual.`) | 2 | rename to `model.visual.` | +| `prefix` | uses vLLM's internal names (`language_model.model.`, `language_model.lm_head.`) | 2 | rename to `model.language_model.` / `lm_head.` | +| `quant_declared_but_dense` | stores modules in bf16 that its quantization_config does not skip (vision tower, lm_head) | 1 | add them to the skip lists | +| `quant_skipped_but_quantized` | stores modules quantized that its skip list covers, often through a container entry (`...linear_attn`) | 1 | remove the entry, and list its dense children instead | +| `ct_target_vllm_names` | writes compressed-tensors targets with vLLM's names, so they match no tensor | 1 | rewrite the targets to `model.language_model.` | +| `quant_unlisted_dense_router` (kind `loader`) | stores in bf16, without listing them as skipped, modules the family serves in bf16 anyway (the Qwen3.5 routers) | 1 | list them; works around the reader | +| `kv_output_scale` (kind `loader`) | stores fp8 KV scales as `k_proj.output_scale` / `v_proj.output_scale` (only reported for the families whose reader drops `k_scale`) | 2 | rename to `k_scale` / `v_scale`, which the reader drops | +| `expert_tensors_incomplete` | lacks some expert's tensors in some layer | - | download it again, or ask its uploader | +| (stage `files`) | has a shard whose tensors do not cover the file, or an index that disagrees with the shards | - | download it again | +| `no_vision_tower` | is text-only under a multimodal architecture | runtime | serve with `--text-model-only` | +| `no_processor_config` | has a vision tower but no `preprocessor_config.json` (image requests fail with 400) | file | copy that file from the base model | +| `fp8_without_config` | stores fp8 weights with no quantization_config | - | not fixable here; the scales and block size are unknown | +| `mlx` (kind `quant_method`) | is MLX-quantized (`.scales` / `.biases`) | - | not served; use a non-MLX release | + +## Step 5: check the fixed checkpoint + +`diagnose.py` passing means: the names, shapes, quantization methods and expert-bank layout all line up with what the model builds. +It does not read weight values. On a real checkpoint, finish with `ft serve --model-path ` and a short request. For a checkpoint with a vision tower, also send an image request. + +## Using the primitives + +```python +import sys; sys.path.insert(0, "scripts/checkpoint_diagnose") +from ckpt_edit import Checkpoint, confirm, plan_lines_for_save + +ck = Checkpoint("/models/x") +ck.rename_regex([(r"^old\.", "new.")]) +ck.quant_skip_add(["model.visual.*"]) +print("\n".join(ck.log + plan_lines_for_save(ck, "/models/x-fixed"))) +ck.save(out="/models/x-fixed") # ck.save() edits in place, with a backup +``` + +`ck.tensors()` gives every tensor's shard, dtype and shape from the headers. `ck.quant_configs()` gives every quantization config the checkpoint carries. diff --git a/scripts/checkpoint_diagnose/ckpt_edit.py b/scripts/checkpoint_diagnose/ckpt_edit.py new file mode 100644 index 0000000000..5ec5215f55 --- /dev/null +++ b/scripts/checkpoint_diagnose/ckpt_edit.py @@ -0,0 +1,512 @@ +"""Edit primitives for HF safetensors checkpoints: config keys, quantization skip lists, tensor names. Tensor data is never changed. + + from ckpt_edit import Checkpoint + ck = Checkpoint("/path/to/ckpt") + ck.set_config("text_config.tie_word_embeddings", False) + ck.quant_skip_add(["model.language_model.layers.*.mlp.gate"]) + ck.rename(lambda name: name.replace("model.language_model.language_model.", "model.language_model.")) + ck.save(out="/path/to/fixed") # or ck.save() to edit in place + +``save(out=...)`` writes a new directory: changed files are written there, every other file is a symlink to the original. +``save()`` edits in place and first copies each changed JSON file and shard header to ``.ckpt_edit_backup/